Final GSOC 2025 commit.

Author: AstitvaAggarwal

src/NN_SDE_solve.jl

@@ -363,7 +363,8 @@ function generate_DataMoments_loss(
 
     # moment matching (MSE across time for 1st, 2nd moments) - assumes diffusion is a Gaussian at each timepoint
     # uses sample variance
-    return (θ, _) -> begin
+    return (θ,
+        _) -> begin
         sum(abs2,
             mean(process, dims = 2) .-
             mean.(Base.Fix2(phi, θ).(sdephi_inputs))) / length(ts) +
@@ -406,7 +407,8 @@ function generate_EM_L2loss(dataset::Vector{<:Vector}, f, g)
     X_increments = reduce(hcat, get_increments.(dataset[1]))
     n, n_samples = size(X_increments)
 
-    loss_fn = (θ, _) -> begin
+    loss_fn = (θ,
+        _) -> begin
         gx = reduce(hcat,
             [[g(process[i, j], θ.p, dataset[2][i])^2 * Δt[i] for i in 1:n]
              for j in 1:n_samples])
@@ -658,7 +660,8 @@ function SciMLBase.__solve(
     # For weak training: higher sub_batch corresponds with a narrower confidence band/ increased certainty in the Weak solution.
     # For strong training: it means more strong paths to train over.
     # weak loss-> weak training is default solve mode.
-    (; param_estim, sub_batch, strong_loss, moment_loss, chain, opt, autodiff, init_params, batch,
+    (; param_estim, sub_batch, strong_loss, moment_loss,
+        chain, opt, autodiff, init_params, batch,
         additional_loss, dataset, numensemble, data_sub_batch) = alg
     n_z = chain[1].in_dims - 1
     sde_phi, init_params = generate_phi(chain, t0, u0, init_params)
@@ -709,7 +712,8 @@ function SciMLBase.__solve(
         if moment_loss
             # min batch for L2 mean is sub samples of the dataset
             data_sub_batch = max(data_sub_batch, length(dataset[1]))
-            DataMoments_loss, dataset_training_sets = generate_DataMoments_loss(
+            DataMoments_loss,
+            dataset_training_sets = generate_DataMoments_loss(
                 dataset, n_z, sde_phi, f, g,
                 autodiff, p, param_estim, data_sub_batch, train_type)
         end

test/NN_SDE_tests.jl

@@ -76,8 +76,8 @@ end
          z3 * sin((3 - 1 / 2) * π * t) / ((3 - 1 / 2) * π))
     end
     truncated_sol(
-    u0, t, z1, z2, z3) = u0 *
-                         exp((α - β^2 / 2) * t + β * W_kkl(t, z1, z2, z3))
+        u0, t, z1, z2, z3) = u0 *
+                             exp((α - β^2 / 2) * t + β * W_kkl(t, z1, z2, z3))
 
     num_samples = 2000
     num_time_steps = dt
@@ -446,8 +446,8 @@ end
 
     # estimated sde parameter tests (we trained with 15 observed solution paths).
     # absolute value taken for 2nd estimated parameter as loss for variance is independent of this parameter's direction.
-    @test sol_1.estimated_params[1].≈ideal_p[1] rtol=2e-1
-    @test abs(sol_1.estimated_params[2]).≈ideal_p[2] rtol=8e-2
-    @test sol_2.estimated_params[1].≈ideal_p[1] rtol=2e-1
-    @test abs(sol_2.estimated_params[2]).≈ideal_p[2] rtol=8e-2
+    @test sol_1.estimated_params[1] .≈ ideal_p[1] rtol=2e-1
+    @test abs(sol_1.estimated_params[2]) .≈ ideal_p[2] rtol=8e-2
+    @test sol_2.estimated_params[1] .≈ ideal_p[1] rtol=2e-1
+    @test abs(sol_2.estimated_params[2]) .≈ ideal_p[2] rtol=8e-2
 end