Move the with_stack hack inside of the DCM code (#526)

* move the `with_stack` hack inside of the DCM code

* don't use `with_stack` on the DCM test

---------

Co-authored-by: haris organtzidis <organtzh@gmail.com>

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

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
@@ -239,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