From 00d63cbd40c2e6715cb44db134c094b62bd3347e Mon Sep 17 00:00:00 2001
From: Paul Shen <pxshen@alumni.stanford.edu>
Date: Mon, 30 Dec 2024 20:55:53 -0500
Subject: [PATCH]

---
 src/pic/run.jl | 8 +++++---
 src/pictest.jl | 6 +++---
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/src/pic/run.jl b/src/pic/run.jl
index 65d3567..75edda4 100644
--- a/src/pic/run.jl
+++ b/src/pic/run.jl
@@ -203,7 +203,7 @@ function picrun(path; gpuarray=nothing, kw...)
             end
         end
         model = models[1]
-        opt = AreaChangeOptimiser(model)
+        opt = AreaChangeOptimiser(model; minchange=0.001)
         opt_state = Flux.setup(opt, model)
         println("starting optimization... first iter will be slow due to adjoint compilation.")
         img = nothing
@@ -232,7 +232,7 @@ function picrun(path; gpuarray=nothing, kw...)
                     # T * dϕ / π
                 end
             else
-                @time global l, (dldm,) = Flux.withgradient(model) do model
+                function f(model)
                     models = [model]
                     res = make_pic_sim_prob(runs, run_probs, lb, dl,
                         designs, design_config, models, ;
@@ -282,6 +282,9 @@ function picrun(path; gpuarray=nothing, kw...)
                     println("    weighted total loss $l")
                     l
                 end
+
+                @time global l, (dldm,) = Flux.withgradient(f, model)
+
             end
             @assert !isnothing(dldm)
             if !isnothing(stoploss) && l < stoploss
@@ -325,7 +328,6 @@ function picrun(path; gpuarray=nothing, kw...)
             opt.maxchange = 0.001 + relu.(l - [0.1, 0.3, 0.7]) ⋅ [0.01, 0.01, 0.01]
             Jello.update_loss!(opt, l)
             Flux.update!(opt_state, model, dldm)# |> gpu)
-            repair!(model)
         end
         if framerate > 0
             make_pic_sim_prob(runs, run_probs, lb, dl,
diff --git a/src/pictest.jl b/src/pictest.jl
index 34ff7f5..0ac5bd1 100644
--- a/src/pictest.jl
+++ b/src/pictest.jl
@@ -8,7 +8,7 @@ ENV["JULIA_PKG_PRECOMPILE_AUTO"] = 0
 # picrun(joinpath("runs", "bend_R5"))
 # picrun(joinpath("runs", "mode_converter"))
 # picrun(joinpath("runs", "demux"))
-# picrun(joinpath("runs", "splitter"))
+picrun(joinpath("runs", "splitter"))
 
-using CUDA
-picrun(joinpath("runs", "tiny"); gpuarray=cu)
\ No newline at end of file
+# using CUDA
+# picrun(joinpath("runs", "tiny"); gpuarray=cu)
\ No newline at end of file