From b221953b9be3b9735655249b402d6686faf53637 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sat, 2 Aug 2025 10:03:21 +0530 Subject: [PATCH 1/2] 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 (96%) 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 96% rename from src/sophia.jl rename to lib/OptimizationSophia/src/OptimizationSophia.jl index 6abef831a..5cce1e3f7 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 + struct Sophia η::Float64 βs::Tuple{Float64, Float64} @@ -119,3 +125,5 @@ function SciMLBase.__solve(cache::OptimizationCache{ θ, x) 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 8d0257dd1..9908e677b 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") include("auglag.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 0b581954d7c636a1209a7c0e28a9982fcc743d14 Mon Sep 17 00:00:00 2001 From: paramthakkar123 Date: Sat, 2 Aug 2025 10:09:05 +0530 Subject: [PATCH 2/2] 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