some changes

sivasathyaseeelan · sivasathyaseeelan · commit 3750a8ad0e45 · 2025-06-06T18:55:11.000+05:30
diff --git a/benchmarks/Jumps/VR_Aggregator_Benchmark.jmd b/benchmarks/Jumps/VR_Aggregator_Benchmark.jmd
@@ -29,11 +29,11 @@ The test cases are:
 7. **Matrix ODE with Variable Rate Jump**: Solved with `Tsit5`.
 8. **Complex ODE with Variable Rate Jump**: Solved with `Tsit5`.
 
-For visualization, we solve one trajectory per test case with 2 jumps (2x2 matrix for Test 7). For benchmarking, we vary jumps from 1 to 20 (2x2 to 10x10 for Test 7), running 100 trajectories per configuration.
+For visualization, we solve one trajectory per test case with 2 jumps. For benchmarking, we vary jumps from 1 to 20, running 100 trajectories per configuration.
 
 # Benchmark and Visualization Setup
 
-We define factories for each test case to create problems with a variable number of jumps (or matrix size for Test 7).
+We define factories for each test case to create problems with a variable number of jumps.
 
 ```julia
 algorithms = Tuple{Any, Any, String, String}[
@@ -45,18 +45,14 @@ algorithms = Tuple{Any, Any, String, String}[
     (VR_FRM(), Rosenbrock23(), "VR_FRM", "Test 1 Rosenbrock23 (autodiff, VR_FRM)"),
     (VR_Direct(), SRIW1(), "VR_Direct", "Test 2 SRIW1 (VR_Direct)"),
     (VR_FRM(), SRIW1(), "VR_FRM", "Test 2 SRIW1 (VR_FRM)"),
-    (VR_Direct(), SRA1(), "VR_Direct", "Test 3 SRA1 (VR_Direct)"),
-    (VR_FRM(), SRA1(), "VR_FRM", "Test 3 SRA1 (VR_FRM)"),
-    (VR_Direct(), Tsit5(), "VR_Direct", "Test 4 Tsit5 (VR_Direct, ConstantRateJump)"),
-    (VR_FRM(), Tsit5(), "VR_FRM", "Test 4 Tsit5 (VR_FRM, ConstantRateJump)"),
-    (VR_Direct(), Tsit5(), "VR_Direct", "Test 5 Tsit5 (VR_Direct)"),
-    (VR_FRM(), Tsit5(), "VR_FRM", "Test 5 Tsit5 (VR_FRM)"),
-    (VR_Direct(), SRIW1(), "VR_Direct", "Test 6 SRIW1 (VR_Direct)"),
-    (VR_FRM(), SRIW1(), "VR_FRM", "Test 6 SRIW1 (VR_FRM)"),
-    (VR_Direct(), Tsit5(), "VR_Direct", "Test 7 Tsit5 (VR_Direct)"),
-    (VR_FRM(), Tsit5(), "VR_FRM", "Test 7 Tsit5 (VR_FRM)"),
-    (VR_Direct(), Tsit5(), "VR_Direct", "Test 8 Tsit5 (VR_Direct)"),
-    (VR_FRM(), Tsit5(), "VR_FRM", "Test 8 Tsit5 (VR_FRM)"),
+    (VR_Direct(), Tsit5(), "VR_Direct", "Test 3 Tsit5 (VR_Direct, ConstantRateJump)"),
+    (VR_FRM(), Tsit5(), "VR_FRM", "Test 3 Tsit5 (VR_FRM, ConstantRateJump)"),
+    (VR_Direct(), Tsit5(), "VR_Direct", "Test 4 Tsit5 (VR_Direct)"),
+    (VR_FRM(), Tsit5(), "VR_FRM", "Test 4 Tsit5 (VR_FRM)"),
+    (VR_Direct(), SRIW1(), "VR_Direct", "Test 5 SRIW1 (VR_Direct)"),
+    (VR_FRM(), SRIW1(), "VR_FRM", "Test 5 SRIW1 (VR_FRM)"),
+    (VR_Direct(), Tsit5(), "VR_Direct", "Test 6 Tsit5 (VR_Direct)"),
+    (VR_FRM(), Tsit5(), "VR_FRM", "Test 6 Tsit5 (VR_FRM)"),
 ]
 
 function create_test1_problem(num_jumps, vr_aggregator, solver)
@@ -79,19 +75,6 @@ function create_test2_problem(num_jumps, vr_aggregator, solver)
 end
 
 function create_test3_problem(num_jumps, vr_aggregator, solver)
-    ff = (du, u, p, t) -> (du .= p == 0 ? 1.01u : 2.01u)
-    gg = (du, u, p, t) -> begin
-        du[1, 1] = 0.3u[1]; du[1, 2] = 0.6u[1]
-        du[2, 1] = 1.2u[1]; du[2, 2] = 0.2u[2]
-    end
-    prob = SDEProblem(ff, gg, ones(2), (0.0, 1.0), 0, noise_rate_prototype=zeros(2, 2))
-    jumps = [VariableRateJump((u, p, t) -> u[1] * 1.0, (integrator) -> (integrator.p = 1)) for _ in 1:num_jumps]
-    jump_prob = JumpProblem(prob, Direct(), jumps...; vr_aggregator=vr_aggregator, rng=rng)
-    ensemble_prob = EnsembleProblem(prob)
-    return ensemble_prob, jump_prob
-end
-
-function create_test4_problem(num_jumps, vr_aggregator, solver)
     f2 = (du, u, p, t) -> (du[1] = u[1])
     prob = ODEProblem(f2, [0.2], (0.0, 10.0))
     jumps = [ConstantRateJump((u, p, t) -> 2, (integrator) -> (integrator.u[1] = integrator.u[1] / 2)) for _ in 1:num_jumps]
@@ -100,7 +83,7 @@ function create_test4_problem(num_jumps, vr_aggregator, solver)
     return ensemble_prob, jump_prob
 end
 
-function create_test5_problem(num_jumps, vr_aggregator, solver)
+function create_test4_problem(num_jumps, vr_aggregator, solver)
     f2 = (du, u, p, t) -> (du[1] = u[1])
     prob = ODEProblem(f2, [0.2], (0.0, 10.0))
     jumps = [VariableRateJump((u, p, t) -> u[1], (integrator) -> (integrator.u[1] = integrator.u[1] / 2); interp_points=100) for _ in 1:num_jumps]
@@ -109,7 +92,7 @@ function create_test5_problem(num_jumps, vr_aggregator, solver)
     return ensemble_prob, jump_prob
 end
 
-function create_test6_problem(num_jumps, vr_aggregator, solver)
+function create_test5_problem(num_jumps, vr_aggregator, solver)
     f2 = (du, u, p, t) -> (du[1] = u[1])
     g2 = (du, u, p, t) -> (du[1] = u[1])
     prob = SDEProblem(f2, g2, [0.2], (0.0, 10.0))
@@ -119,19 +102,7 @@ function create_test6_problem(num_jumps, vr_aggregator, solver)
     return ensemble_prob, jump_prob
 end
 
-function create_test7_problem(num_jumps, vr_aggregator, solver, matrix_size=2)
-    f3 = (du, u, p, t) -> (du .= u)
-    u0 = ones(matrix_size, matrix_size)
-    prob = ODEProblem(f3, u0, (0.0, 1.0))
-    rate3 = (u, p, t) -> sum(u[1, :])
-    affect3! = (integrator) -> (integrator.u .= range(0.25, 1.0, length=matrix_size^2))
-    jumps = [VariableRateJump(rate3, affect3!) for _ in 1:num_jumps]
-    jump_prob = JumpProblem(prob, Direct(), jumps...; vr_aggregator=vr_aggregator, rng=rng)
-    ensemble_prob = EnsembleProblem(prob)
-    return ensemble_prob, jump_prob
-end
-
-function create_test8_problem(num_jumps, vr_aggregator, solver)
+function create_test6_problem(num_jumps, vr_aggregator, solver)
     f4 = (dx, x, p, t) -> (dx[1] = x[1])
     rate4 = (x, p, t) -> t
     affect4! = (integrator) -> (integrator.u[1] = integrator.u[1] * 0.5)
@@ -145,19 +116,15 @@ end
 
 # Solution Visualization
 
-We solve one trajectory for each test case with 2 jumps (2x2 matrix for Test 7) using `VR_Direct` and plot the state variables vs. time.
+We solve one trajectory for each test case with 2 jumps using `VR_Direct` and plot the state variables vs. time.
 
 ```julia
 let figs = []
-    for test_num in 1:8
+    for test_num in 1:6
         # Select a representative solver for each test
-        algo, stepper = if test_num == 1
-            VR_Direct(), Tsit5()
-        elseif test_num == 2 || test_num == 6
+        algo, stepper = if  test_num == 2 || test_num == 5
             VR_Direct(), SRIW1()
-        elseif test_num == 3
-            VR_Direct(), SRA1()
-        elseif test_num in [4, 5, 7, 8]
+        else
             VR_Direct(), Tsit5()
         end
         label = "Test $test_num"
@@ -175,29 +142,16 @@ let figs = []
             create_test5_problem(2, algo, stepper)
         elseif test_num == 6
             create_test6_problem(2, algo, stepper)
-        elseif test_num == 7
-            create_test7_problem(2, algo, stepper, 2)
-        elseif test_num == 8
-            create_test8_problem(2, algo, stepper)
         end
-        
-        # Solve one trajectory
-        solver_kwargs = test_num == 3 ? (dt=1.0,) : ()
-        try
-            sol = solve(jump_prob, stepper; saveat=0.01, solver_kwargs...)
             
+        try
+            sol = solve(jump_prob, stepper; saveat=0.01)
+
             # Plot solution
             fig = plot(title="Test $test_num: Solution Trajectory", xlabel="Time", ylabel="State")
-            if test_num == 7
-                # For matrix ODE, plot sum of elements
-                plot!(sol.t, [sum(sol.u[i]) for i in 1:length(sol.u)], label="Sum of Matrix Elements")
-            elseif test_num == 8
+            if test_num == 6
                 # For complex ODE, plot real part
                 plot!(sol.t, real.(sol[1,:]), label="Real Part")
-            elseif test_num == 3
-                # For 2D SDE, plot both components
-                plot!(sol.t, sol[1,:], label="u[1]")
-                plot!(sol.t, sol[2,:], label="u[2]")
             else
                 # For scalar problems, plot state
                 plot!(sol.t, sol[1,:], label="u[1]")
@@ -213,11 +167,10 @@ end
 
 # Benchmark Execution
 
-We benchmark each test case for 1 to 20 jumps (2x2 to 10x10 for Test 7), running 100 trajectories. Errors are logged to diagnose failures.
+We benchmark each test case for 1 to 20 jumps, running 100 trajectories. Errors are logged to diagnose failures.
 
 ```julia
 num_jumps_range = append!([1], 5:5:20)
-matrix_sizes = [2, 4, 6, 8, 10]
 bs = Vector{Vector{BenchmarkTools.Trial}}()
 errors = Dict{String, Vector{String}}()
 
@@ -227,38 +180,36 @@ for (algo, stepper, agg_name, label) in algorithms
     errors[label] = String[]
     _bs = bs[end]
     test_num = parse(Int, match(r"Test (\d+)", label).captures[1])
-    is_matrix_test = test_num == 7
-    range_var = is_matrix_test ? matrix_sizes : num_jumps_range
+    range_var = num_jumps_range
     for (i, var) in enumerate(range_var)
         if test_num == 1
-            ensemble_prob, jump_prob = create_test1_problem(is_matrix_test ? 2 : var, algo, stepper)
+            ensemble_prob, jump_prob = create_test1_problem(var, algo, stepper)
         elseif test_num == 2
-            ensemble_prob, jump_prob = create_test2_problem(is_matrix_test ? 2 : var, algo, stepper)
+            ensemble_prob, jump_prob = create_test2_problem(var, algo, stepper)
         elseif test_num == 3
-            ensemble_prob, jump_prob = create_test3_problem(is_matrix_test ? 2 : var, algo, stepper)
+            ensemble_prob, jump_prob = create_test3_problem(var, algo, stepper)
         elseif test_num == 4
-            ensemble_prob, jump_prob = create_test4_problem(is_matrix_test ? 2 : var, algo, stepper)
+            ensemble_prob, jump_prob = create_test4_problem(var, algo, stepper)
         elseif test_num == 5
-            ensemble_prob, jump_prob = create_test5_problem(is_matrix_test ? 2 : var, algo, stepper)
+            ensemble_prob, jump_prob = create_test5_problem(var, algo, stepper)
         elseif test_num == 6
-            ensemble_prob, jump_prob = create_test6_problem(is_matrix_test ? 2 : var, algo, stepper)
-        elseif test_num == 7
-            ensemble_prob, jump_prob = create_test7_problem(2, algo, stepper, var)
-        elseif test_num == 8
-            ensemble_prob, jump_prob = create_test8_problem(is_matrix_test ? 2 : var, algo, stepper)
+            ensemble_prob, jump_prob = create_test6_problem(var, algo, stepper)
         end
-        solver_kwargs = test_num == 3 ? (dt=1.0,) : ""
         trial = try
-            @benchmark solve($ensemble_prob, $stepper, EnsembleSerial(), trajectories=100, jump_prob=$jump_prob; $solver_kwargs...) samples=50 evals=1 seconds=10
+            @benchmark(
+                solve($jump_prob, $stepper), 
+                samples=50, 
+                evals=1, 
+                seconds=10
+            )
         catch e
-            push!(errors[label], "Error at $(is_matrix_test ? "Matrix Size" : "Num Jumps") = $var: $(sprint(showerror, e))")
+            push!(errors[label], "Error at Num Jumps = $var: $(sprint(showerror, e))")
             BenchmarkTools.Trial(BenchmarkTools.Parameters(samples=50, evals=1, seconds=10))
         end
         push!(_bs, trial)
-        if (var == 1 || var % (is_matrix_test ? 2 : 5) == 0)
-            median_time = length(trial) > 0 ? "$(BenchmarkTools.prettytime(median(trial.times)))" : "nan"
-            println("algo=$label, $(is_matrix_test ? "Matrix Size" : "Num Jumps") = $var, length = $(length(trial.times)), median time = $median_time")
-        end
+
+        median_time = length(trial) > 0 ? "$(BenchmarkTools.prettytime(median(trial.times)))" : "nan"
+        println("algo=$label, Num Jumps = $var, length = $(length(trial.times)), median time = $median_time")
     end
 end
 
@@ -279,7 +230,7 @@ We plot the median execution times for each test case, comparing `VR_Direct` and
 
 ```julia
 let figs = []
-    for test_num in 1:8
+    for test_num in 1:6
         test_algorithms = filter(a -> parse(Int, match(r"Test (\d+)", a[4]).captures[1]) == test_num, algorithms)
         is_matrix_test = test_num == 7
         range_var = is_matrix_test ? matrix_sizes : num_jumps_range
@@ -307,6 +258,6 @@ let figs = []
         end
         push!(figs, fig)
     end
-    plot(figs..., layout=(4, 2), format=fmt, size=(width_px, 4*height_px/2))
+    plot(figs..., layout=(6, 1), format=fmt, size=(width_px, 8*height_px/2))
 end
 ```
