From ec1b88ab212fb90ffbde1f8cfe443105ccc2b892 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sat, 2 Aug 2025 10:03:21 +0530 Subject: [PATCH 01/11] Added new subpackage for Sophia.jl --- lib/OptimizationSophia/Project.toml | 28 ++++++++ .../src/OptimizationSophia.jl | 8 +++ lib/OptimizationSophia/test/runtests.jl | 67 +++++++++++++++++++ src/Optimization.jl | 1 - test/minibatch.jl | 5 -- 5 files changed, 103 insertions(+), 6 deletions(-) create mode 100644 lib/OptimizationSophia/Project.toml rename src/sophia.jl => lib/OptimizationSophia/src/OptimizationSophia.jl (97%) create mode 100644 lib/OptimizationSophia/test/runtests.jl diff --git a/lib/OptimizationSophia/Project.toml b/lib/OptimizationSophia/Project.toml new file mode 100644 index 000000000..53bf80dd8 --- /dev/null +++ b/lib/OptimizationSophia/Project.toml @@ -0,0 +1,28 @@ +name = "OptimizationSophia" +uuid = "892fee11-dca1-40d6-b698-84ba0d87399a" +authors = ["paramthakkar123 "] +version = "0.1.0" + +[deps] +ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" +Lux = "b2108857-7c20-44ae-9111-449ecde12c47" +MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" +Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" +OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" +OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" +SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1" +Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" +Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" + +[compat] +ComponentArrays = "0.15.29" +Lux = "1.16.0" +MLUtils = "0.4.8" +Optimization = "4.5.0" +OptimizationBase = "2.10.0" +OrdinaryDiffEqTsit5 = "1.2.0" +Random = "1.11.0" +SciMLSensitivity = "7.88.0" +Test = "1.11.0" +Zygote = "0.7.10" diff --git a/src/sophia.jl b/lib/OptimizationSophia/src/OptimizationSophia.jl similarity index 97% rename from src/sophia.jl rename to lib/OptimizationSophia/src/OptimizationSophia.jl index 8516913a2..bd15be621 100644 --- a/src/sophia.jl +++ b/lib/OptimizationSophia/src/OptimizationSophia.jl @@ -1,3 +1,9 @@ +module OptimizationSophia + +using OptimizationBase.SciMLBase +using OptimizationBase: OptimizationCache +using Optimization + """ Sophia(; η = 1e-3, βs = (0.9, 0.999), ϵ = 1e-8, λ = 1e-1, k = 10, ρ = 0.04) @@ -171,3 +177,5 @@ function SciMLBase.__solve(cache::OptimizationCache{ θ, x, retcode = ReturnCode.Success) end + +end diff --git a/lib/OptimizationSophia/test/runtests.jl b/lib/OptimizationSophia/test/runtests.jl new file mode 100644 index 000000000..f5ba929a9 --- /dev/null +++ b/lib/OptimizationSophia/test/runtests.jl @@ -0,0 +1,67 @@ +using OptimizationBase, Optimization +using OptimizationBase.SciMLBase: solve, OptimizationFunction, OptimizationProblem +using OptimizationSophia +using Lux, MLUtils, Random, ComponentArrays +using SciMLSensitivity +using Test +using Zygote +using OrdinaryDiffEqTsit5 + +function dudt_(u, p, t) + ann(u, p, st)[1] .* u +end + +function newtons_cooling(du, u, p, t) + temp = u[1] + k, temp_m = p + du[1] = dT = -k * (temp - temp_m) +end + +function true_sol(du, u, p, t) + true_p = [log(2) / 8.0, 100.0] + newtons_cooling(du, u, true_p, t) +end + +function callback(state, l) #callback function to observe training + display(l) + return l < 1e-2 +end + +function predict_adjoint(fullp, time_batch) + Array(solve(prob, Tsit5(), p = fullp, saveat = time_batch)) +end + +function loss_adjoint(fullp, p) + (batch, time_batch) = p + pred = predict_adjoint(fullp, time_batch) + sum(abs2, batch .- pred) +end + +u0 = Float32[200.0] +datasize = 30 +tspan = (0.0f0, 1.5f0) +rng = Random.default_rng() + +ann = Lux.Chain(Lux.Dense(1, 8, tanh), Lux.Dense(8, 1, tanh)) +pp, st = Lux.setup(rng, ann) +pp = ComponentArray(pp) + +prob = ODEProblem{false}(dudt_, u0, tspan, pp) + +t = range(tspan[1], tspan[2], length = datasize) +true_prob = ODEProblem(true_sol, u0, tspan) +ode_data = Array(solve(true_prob, Tsit5(), saveat = t)) + +k = 10 +train_loader = MLUtils.DataLoader((ode_data, t), batchsize = k) + +l1 = loss_adjoint(pp, (train_loader.data[1], train_loader.data[2]))[1] + +optfun = OptimizationFunction(loss_adjoint, + OptimizationBase.AutoZygote()) +optprob = OptimizationProblem(optfun, pp, train_loader) + +res1 = solve(optprob, + OptimizationSophia.Sophia(), callback = callback, + maxiters = 2000) +@test 10res1.objective < l1 \ No newline at end of file diff --git a/src/Optimization.jl b/src/Optimization.jl index 3d2a540de..8b20ef818 100644 --- a/src/Optimization.jl +++ b/src/Optimization.jl @@ -23,7 +23,6 @@ export ObjSense, MaxSense, MinSense include("utils.jl") include("state.jl") include("lbfgsb.jl") -include("sophia.jl") export solve diff --git a/test/minibatch.jl b/test/minibatch.jl index 8185ec4bc..abd5a2610 100644 --- a/test/minibatch.jl +++ b/test/minibatch.jl @@ -58,11 +58,6 @@ optfun = OptimizationFunction(loss_adjoint, Optimization.AutoZygote()) optprob = OptimizationProblem(optfun, pp, train_loader) -res1 = Optimization.solve(optprob, - Optimization.Sophia(), callback = callback, - maxiters = 2000) -@test 10res1.objective < l1 - optfun = OptimizationFunction(loss_adjoint, Optimization.AutoForwardDiff()) optprob = OptimizationProblem(optfun, pp, train_loader) From 045dfe55205c6af6203492bc442f41f73cfa8eaa Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sat, 2 Aug 2025 10:09:05 +0530 Subject: [PATCH 02/11] Formatted --- lib/OptimizationSophia/test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OptimizationSophia/test/runtests.jl b/lib/OptimizationSophia/test/runtests.jl index f5ba929a9..3536c0512 100644 --- a/lib/OptimizationSophia/test/runtests.jl +++ b/lib/OptimizationSophia/test/runtests.jl @@ -64,4 +64,4 @@ optprob = OptimizationProblem(optfun, pp, train_loader) res1 = solve(optprob, OptimizationSophia.Sophia(), callback = callback, maxiters = 2000) -@test 10res1.objective < l1 \ No newline at end of file +@test 10res1.objective < l1 From 7ea04cacbddbf47be9506689565a8afab546432c Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 1 Oct 2025 19:06:51 +0200 Subject: [PATCH 03/11] add test runner --- .github/workflows/CI.yml | 1 + 1 file changed, 1 insertion(+) diff --git a/.github/workflows/CI.yml b/.github/workflows/CI.yml index 5b29379f7..fd9e26ead 100644 --- a/.github/workflows/CI.yml +++ b/.github/workflows/CI.yml @@ -38,6 +38,7 @@ jobs: - OptimizationPyCMA - OptimizationQuadDIRECT - OptimizationSciPy + - OptimizationSophia - OptimizationSpeedMapping - OptimizationPolyalgorithms - OptimizationNLPModels From 8a08402f1fcea7e91b673ec9719e889fd8b813bd Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 1 Oct 2025 13:30:04 -0400 Subject: [PATCH 04/11] Update Project.toml --- lib/OptimizationSophia/Project.toml | 9 +++++++-- 1 file changed, 7 insertions(+), 2 deletions(-) diff --git a/lib/OptimizationSophia/Project.toml b/lib/OptimizationSophia/Project.toml index 53bf80dd8..f27211167 100644 --- a/lib/OptimizationSophia/Project.toml +++ b/lib/OptimizationSophia/Project.toml @@ -4,11 +4,13 @@ authors = ["paramthakkar123 "] version = "0.1.0" [deps] +Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" +OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" + +[extras] ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" -Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" -OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1" @@ -26,3 +28,6 @@ Random = "1.11.0" SciMLSensitivity = "7.88.0" Test = "1.11.0" Zygote = "0.7.10" + +[targets] +test = ["Test", "ComponentArrays", "Lux", "MLUtils", "OrdinaryDiffEqTsit5", "Random", "SciMLSensitivity", "Zygote"] From 8cc2b4eb651888dc53b6ca1b89365cc58426fd32 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 1 Oct 2025 13:31:00 -0400 Subject: [PATCH 05/11] Update native.jl --- test/native.jl | 11 ----------- 1 file changed, 11 deletions(-) diff --git a/test/native.jl b/test/native.jl index e1df7c9a7..f7385fd0d 100644 --- a/test/native.jl +++ b/test/native.jl @@ -55,14 +55,3 @@ optf1 = OptimizationFunction(loss, AutoSparseForwardDiff()) prob1 = OptimizationProblem(optf1, rand(5), data) sol1 = solve(prob1, OptimizationOptimisers.Adam(), maxiters = 1000, callback = callback) @test sol1.objective < l0 - -# Test Sophia with ComponentArrays + Enzyme (shadow generation fix) -using ComponentArrays -x0_comp = ComponentVector(a = 0.0, b = 0.0) -rosenbrock_comp(x, p = nothing) = (1 - x.a)^2 + 100 * (x.b - x.a^2)^2 - -optf_sophia = OptimizationFunction(rosenbrock_comp, AutoEnzyme()) -prob_sophia = OptimizationProblem(optf_sophia, x0_comp) -res_sophia = solve(prob_sophia, Optimization.Sophia(η=0.01, k=5), maxiters = 50) -@test res_sophia.objective < rosenbrock_comp(x0_comp) # Test optimization progress -@test res_sophia.retcode == Optimization.SciMLBase.ReturnCode.Success From 36983a5320e51e3ce25366dabbd01ecf8c28ec73 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 1 Oct 2025 13:31:24 -0400 Subject: [PATCH 06/11] Update runtests.jl --- lib/OptimizationSophia/test/runtests.jl | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/lib/OptimizationSophia/test/runtests.jl b/lib/OptimizationSophia/test/runtests.jl index 3536c0512..c1c604cbb 100644 --- a/lib/OptimizationSophia/test/runtests.jl +++ b/lib/OptimizationSophia/test/runtests.jl @@ -65,3 +65,14 @@ res1 = solve(optprob, OptimizationSophia.Sophia(), callback = callback, maxiters = 2000) @test 10res1.objective < l1 + +# Test Sophia with ComponentArrays + Enzyme (shadow generation fix) +using ComponentArrays +x0_comp = ComponentVector(a = 0.0, b = 0.0) +rosenbrock_comp(x, p = nothing) = (1 - x.a)^2 + 100 * (x.b - x.a^2)^2 + +optf_sophia = OptimizationFunction(rosenbrock_comp, AutoEnzyme()) +prob_sophia = OptimizationProblem(optf_sophia, x0_comp) +res_sophia = solve(prob_sophia, Optimization.Sophia(η=0.01, k=5), maxiters = 50) +@test res_sophia.objective < rosenbrock_comp(x0_comp) # Test optimization progress +@test res_sophia.retcode == Optimization.SciMLBase.ReturnCode.Success From bdae56bc55e222e488405821ecbd0d4a05b29541 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Wed, 1 Oct 2025 14:05:10 -0400 Subject: [PATCH 07/11] Update lib/OptimizationSophia/Project.toml --- lib/OptimizationSophia/Project.toml | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/lib/OptimizationSophia/Project.toml b/lib/OptimizationSophia/Project.toml index f27211167..7f658d080 100644 --- a/lib/OptimizationSophia/Project.toml +++ b/lib/OptimizationSophia/Project.toml @@ -24,9 +24,9 @@ MLUtils = "0.4.8" Optimization = "4.5.0" OptimizationBase = "2.10.0" OrdinaryDiffEqTsit5 = "1.2.0" -Random = "1.11.0" +Random = "1.10.0" SciMLSensitivity = "7.88.0" -Test = "1.11.0" +Test = "1.10.0" Zygote = "0.7.10" [targets] From 220d7cfb80e5bf862e767f013d438d61f317d601 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Thu, 2 Oct 2025 09:20:01 +0200 Subject: [PATCH 08/11] Make sure to pull in Random stdlib --- Project.toml | 1 - lib/OptimizationSophia/Project.toml | 2 +- lib/OptimizationSophia/src/OptimizationSophia.jl | 1 + src/Optimization.jl | 2 +- 4 files changed, 3 insertions(+), 3 deletions(-) diff --git a/Project.toml b/Project.toml index 4c1a483df..511eb4f5e 100644 --- a/Project.toml +++ b/Project.toml @@ -14,7 +14,6 @@ LoggingExtras = "e6f89c97-d47a-5376-807f-9c37f3926c36" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" Printf = "de0858da-6303-5e67-8744-51eddeeeb8d7" ProgressLogging = "33c8b6b6-d38a-422a-b730-caa89a2f386c" -Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" Reexport = "189a3867-3050-52da-a836-e630ba90ab69" SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462" SparseArrays = "2f01184e-e22b-5df5-ae63-d93ebab69eaf" diff --git a/lib/OptimizationSophia/Project.toml b/lib/OptimizationSophia/Project.toml index 7f658d080..2bc8bc014 100644 --- a/lib/OptimizationSophia/Project.toml +++ b/lib/OptimizationSophia/Project.toml @@ -6,13 +6,13 @@ version = "0.1.0" [deps] Optimization = "7f7a1694-90dd-40f0-9382-eb1efda571ba" OptimizationBase = "bca83a33-5cc9-4baa-983d-23429ab6bcbb" +Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" [extras] ComponentArrays = "b0b7db55-cfe3-40fc-9ded-d10e2dbeff66" Lux = "b2108857-7c20-44ae-9111-449ecde12c47" MLUtils = "f1d291b0-491e-4a28-83b9-f70985020b54" OrdinaryDiffEqTsit5 = "b1df2697-797e-41e3-8120-5422d3b24e4a" -Random = "9a3f8284-a2c9-5f02-9a11-845980a1fd5c" SciMLSensitivity = "1ed8b502-d754-442c-8d5d-10ac956f44a1" Test = "8dfed614-e22c-5e08-85e1-65c5234f0b40" Zygote = "e88e6eb3-aa80-5325-afca-941959d7151f" diff --git a/lib/OptimizationSophia/src/OptimizationSophia.jl b/lib/OptimizationSophia/src/OptimizationSophia.jl index bd15be621..c5ce693b6 100644 --- a/lib/OptimizationSophia/src/OptimizationSophia.jl +++ b/lib/OptimizationSophia/src/OptimizationSophia.jl @@ -3,6 +3,7 @@ module OptimizationSophia using OptimizationBase.SciMLBase using OptimizationBase: OptimizationCache using Optimization +using Random """ Sophia(; η = 1e-3, βs = (0.9, 0.999), ϵ = 1e-8, λ = 1e-1, k = 10, ρ = 0.04) diff --git a/src/Optimization.jl b/src/Optimization.jl index 8b20ef818..138997ae5 100644 --- a/src/Optimization.jl +++ b/src/Optimization.jl @@ -12,7 +12,7 @@ if !isdefined(Base, :get_extension) end using Logging, ProgressLogging, ConsoleProgressMonitor, TerminalLoggers, LoggingExtras -using ArrayInterface, Base.Iterators, SparseArrays, LinearAlgebra, Random +using ArrayInterface, Base.Iterators, SparseArrays, LinearAlgebra import OptimizationBase: instantiate_function, OptimizationCache, ReInitCache import SciMLBase: OptimizationProblem, From b44069b7537a0856881cdfc6084ba7ee01099cef Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Thu, 2 Oct 2025 09:32:51 +0200 Subject: [PATCH 09/11] v4.8 for testing --- Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/Project.toml b/Project.toml index 511eb4f5e..806eb8707 100644 --- a/Project.toml +++ b/Project.toml @@ -1,6 +1,6 @@ name = "Optimization" uuid = "7f7a1694-90dd-40f0-9382-eb1efda571ba" -version = "5.0.0" +version = "4.8.0" [deps] ADTypes = "47edcb42-4c32-4615-8424-f2b9edc5f35b" From 5d5b16cb0e57a0a1c3f6c93f5d5aaedde52a1903 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Thu, 2 Oct 2025 10:16:27 +0200 Subject: [PATCH 10/11] fix test deps targets --- lib/OptimizationSophia/Project.toml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OptimizationSophia/Project.toml b/lib/OptimizationSophia/Project.toml index 2bc8bc014..41946ca19 100644 --- a/lib/OptimizationSophia/Project.toml +++ b/lib/OptimizationSophia/Project.toml @@ -30,4 +30,4 @@ Test = "1.10.0" Zygote = "0.7.10" [targets] -test = ["Test", "ComponentArrays", "Lux", "MLUtils", "OrdinaryDiffEqTsit5", "Random", "SciMLSensitivity", "Zygote"] +test = ["Test", "ComponentArrays", "Lux", "MLUtils", "OrdinaryDiffEqTsit5", "SciMLSensitivity", "Zygote"] From 0a6432157920897a8c8c6f9079a73bb8c4b246f1 Mon Sep 17 00:00:00 2001 From: Christopher Rackauckas Date: Thu, 2 Oct 2025 10:22:47 +0200 Subject: [PATCH 11/11] test typo --- lib/OptimizationSophia/test/runtests.jl | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/lib/OptimizationSophia/test/runtests.jl b/lib/OptimizationSophia/test/runtests.jl index c1c604cbb..63f17408d 100644 --- a/lib/OptimizationSophia/test/runtests.jl +++ b/lib/OptimizationSophia/test/runtests.jl @@ -73,6 +73,6 @@ rosenbrock_comp(x, p = nothing) = (1 - x.a)^2 + 100 * (x.b - x.a^2)^2 optf_sophia = OptimizationFunction(rosenbrock_comp, AutoEnzyme()) prob_sophia = OptimizationProblem(optf_sophia, x0_comp) -res_sophia = solve(prob_sophia, Optimization.Sophia(η=0.01, k=5), maxiters = 50) +res_sophia = solve(prob_sophia, OptimizationSophia.Sophia(η=0.01, k=5), maxiters = 50) @test res_sophia.objective < rosenbrock_comp(x0_comp) # Test optimization progress @test res_sophia.retcode == Optimization.SciMLBase.ReturnCode.Success