From fbacadc0ad229ba17f79c3a526e0eea9bc88f35f Mon Sep 17 00:00:00 2001
From: Mason Protter <mason.protter@icloud.com>
Date: Fri, 10 Jan 2025 20:24:57 +0100
Subject: [PATCH 1/2] move the `with_stack` hack inside of the DCM code

---
 src/datafitting/spDCM_VL.jl |  8 ++++++++
 test/datafitting.jl         | 21 +++++++++------------
 2 files changed, 17 insertions(+), 12 deletions(-)

diff --git a/src/datafitting/spDCM_VL.jl b/src/datafitting/spDCM_VL.jl
index 3594d92e..7c98dcdd 100644
--- a/src/datafitting/spDCM_VL.jl
+++ b/src/datafitting/spDCM_VL.jl
@@ -460,7 +460,15 @@ function setup_sDCM(data, model, initcond, csdsetup, priors, hyperpriors, indice
     return (vlstate, vlsetup)
 end
 
+with_stack(f, n) = fetch(schedule(Task(f, n)));
+
 function run_sDCM_iteration!(state::VLState, setup::VLSetup)
+    with_stack(5_000_000) do
+        _run_sDCM_iteration!(state, setup)
+    end
+end
+
+function _run_sDCM_iteration!(state::VLState, setup::VLSetup)
     (;μθ_po, λ, v, ϵ_θ, dFdθ, dFdθθ) = state
 
     f = setup.model_at_x0
diff --git a/test/datafitting.jl b/test/datafitting.jl
index 8ff70b80..36273a08 100644
--- a/test/datafitting.jl
+++ b/test/datafitting.jl
@@ -84,19 +84,16 @@ using MAT
 
     # HACK: on machines with very small amounts of RAM, Julia can run out of stack space while compiling the code called in this loop
     # this should be rewritten to abuse the compiler less, but for now, an easy solution is just to run it with more allocated stack space.
-    with_stack(f, n) = fetch(schedule(Task(f,n)))
 
-    with_stack(5_000_000) do  # 5MB of stack space
-        for iter in 1:max_iter
-            state.iter = iter
-            run_sDCM_iteration!(state, setup)
-            print("iteration: ", iter, " - F:", state.F[end] - state.F[2], " - dF predicted:", state.dF[end], "\n")
-            if iter >= 4
-                criterion = state.dF[end-3:end] .< setup.tolerance
-                if all(criterion)
-                    print("convergence\n")
-                    break
-                end
+    for iter in 1:max_iter
+        state.iter = iter
+        run_sDCM_iteration!(state, setup)
+        print("iteration: ", iter, " - F:", state.F[end] - state.F[2], " - dF predicted:", state.dF[end], "\n")
+        if iter >= 4
+            criterion = state.dF[end-3:end] .< setup.tolerance
+            if all(criterion)
+                print("convergence\n")
+                break
             end
         end
     end

From fb3ee653a7f0ae47e7c547f2d00246bd5df5cffc Mon Sep 17 00:00:00 2001
From: Mason Protter <mason.protter@icloud.com>
Date: Fri, 10 Jan 2025 20:29:06 +0100
Subject: [PATCH 2/2] don't use `with_stack` on the DCM test

---
 test/datafitting.jl | 23 +++++++++--------------
 1 file changed, 9 insertions(+), 14 deletions(-)

diff --git a/test/datafitting.jl b/test/datafitting.jl
index 36273a08..584509db 100644
--- a/test/datafitting.jl
+++ b/test/datafitting.jl
@@ -236,21 +236,16 @@ end
     csdsetup = (mar_order = 8, freq = freq, dt = dt);
 
     (state, setup) = setup_sDCM(data, fullmodel, initcond, csdsetup, priors, hyperpriors, indices, modelparam, "LFP");
-    # HACK: on machines with very small amounts of RAM, Julia can run out of stack space while compiling the code called in this loop
-    # this should be rewritten to abuse the compiler less, but for now, an easy solution is just to run it with more allocated stack space.
-    with_stack(f, n) = fetch(schedule(Task(f,n)))
 
-    with_stack(5_000_000) do  # 5MB of stack space
-        for iter in 1:128
-            state.iter = iter
-            run_sDCM_iteration!(state, setup)
-            print("iteration: ", iter, " - F:", state.F[end] - state.F[2], " - dF predicted:", state.dF[end], "\n")
-            if iter >= 4
-                criterion = state.dF[end-3:end] .< setup.tolerance
-                if all(criterion)
-                    print("convergence\n")
-                    break
-                end
+    for iter in 1:128
+        state.iter = iter
+        run_sDCM_iteration!(state, setup)
+        print("iteration: ", iter, " - F:", state.F[end] - state.F[2], " - dF predicted:", state.dF[end], "\n")
+        if iter >= 4
+            criterion = state.dF[end-3:end] .< setup.tolerance
+            if all(criterion)
+                print("convergence\n")
+                break
             end
         end
     end