Adding Training Strategies to dae_solvers.jl

src/dae_solve.jl

@@ -72,6 +72,18 @@ function inner_loss(phi::ODEPhi{C, T, U}, f, autodiff::Bool, t::AbstractVector,
     sum(abs2, loss) / length(t)
 end
 
+function generate_loss(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p, differential_vars::AbstractVector)
+    integrand(t::Number, θ) = abs2(inner_loss(phi, f, autodiff, t, θ, p, differential_vars))
+    integrand(ts, θ) = [abs2(inner_loss(phi, f, autodiff, t, θ, p, differential_vars)) for t in ts] WHAT IS TS HERE?? - DO WE NEED THIS LINE
+    function loss(θ, _)
+        intf = BatchIntegralFunction(integrand, max_batch = strategy.batch)
+        intprob = IntegralProblem(intf, (tspan[1], tspan[2]), θ)
+        sol = solve(intprob, strategy.quadrature_alg; abstol = strategy.abstol, reltol = strategy.reltol, maxiters = strategy.maxiters)
+        sol.u
+    end
+    return loss
+end
+
 function generate_loss(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p,
         differential_vars::AbstractVector)
     ts = tspan[1]:(strategy.dx):tspan[2]
@@ -82,6 +94,42 @@ function generate_loss(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p,
     return loss
 end
 
+function generate_loss(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p, differential_vars::AbstractVector)
+    autodiff && throw(ArgumentError("autodiff not supported for StochasticTraining."))
+    function loss(θ, _)
+        ts = adapt(parameterless_type(θ),
+            [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)])
+        sum(abs2, inner_loss(phi, f, autodiff, ts, θ, p, differential_vars))
+    end
+    return loss
+end
+
+
+function generate_loss(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p,differential_vars::AbstractVector)
+    autodiff && throw(ArgumentError("autodiff not supported for WeightedIntervalTraining."))
+    minT = tspan[1]
+    maxT = tspan[2]
+
+    weights = strategy.weights ./ sum(strategy.weights)
+
+    N = length(weights)
+    points = strategy.points
+
+    difference = (maxT - minT) / N
+
+    data = Float64[]
+    for (index, item) in enumerate(weights)
+        temp_data = rand(1, trunc(Int, points * item)) .* difference .+ minT .+ ((index - 1) * difference)
+        data = append!(data, temp_data)
+    end
+
+    ts = data
+    function loss(θ, _)
+        sum(abs2, inner_loss(phi, f, autodiff, ts, θ, p, differential_vars))
+    end
+    return loss
+end
+
 function DiffEqBase.__solve(prob::DiffEqBase.AbstractDAEProblem,
         alg::NNDAE,
         args...;

src/ode_solve.jl

@@ -219,79 +219,87 @@ end
 
 Representation of the loss function, parametric on the training strategy `strategy`.
 """
-function generate_loss(strategy::QuadratureTraining, phi, f, autodiff::Bool, tspan, p,
-        batch, param_estim::Bool)
-    integrand(t::Number, θ) = abs2(inner_loss(phi, f, autodiff, t, θ, p, param_estim))
-
-    integrand(ts, θ) = [abs2(inner_loss(phi, f, autodiff, t, θ, p, param_estim)) for t in ts]
 
-    function loss(θ, _)
-        intf = BatchIntegralFunction(integrand, max_batch = strategy.batch)
-        intprob = IntegralProblem(intf, (tspan[1], tspan[2]), θ)
-        sol = solve(intprob, strategy.quadrature_alg; abstol = strategy.abstol, reltol = strategy.reltol, maxiters = strategy.maxiters)
-        sol.u
-    end
-    return loss
-end
+function generate_loss(strategy, phi, f, autodiff::Bool, tspan, p,
+        batch, param_estim::Bool)
 
-function generate_loss(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p, batch, param_estim::Bool)
-    ts = tspan[1]:(strategy.dx):tspan[2]
-    autodiff && throw(ArgumentError("autodiff not supported for GridTraining."))
-    function loss(θ, _)
-        if batch
-            inner_loss(phi, f, autodiff, ts, θ, p, param_estim)
-        else
-            sum([inner_loss(phi, f, autodiff, t, θ, p, param_estim) for t in ts])
+    if strategy isa QuadratureTraining
+
+        integrand(t::Number, θ) = abs2(inner_loss(phi, f, autodiff, t, θ, p, param_estim))
+        function integrand(ts, θ)
+            [abs2(inner_loss(phi, f, autodiff, t, θ, p, param_estim)) for t in ts]
         end
-    end
-    return loss
-end
-
-function generate_loss(strategy::StochasticTraining, phi, f, autodiff::Bool, tspan, p,
-        batch, param_estim::Bool)
-    autodiff && throw(ArgumentError("autodiff not supported for StochasticTraining."))
-    function loss(θ, _)
-        ts = adapt(parameterless_type(θ),
-            [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)])
-        if batch
-            inner_loss(phi, f, autodiff, ts, θ, p, param_estim)
-        else
-            sum([inner_loss(phi, f, autodiff, t, θ, p, param_estim) for t in ts])
+
+        function loss(θ, _)
+            intf = BatchIntegralFunction(integrand, max_batch = strategy.batch)
+            intprob = IntegralProblem(intf, (tspan[1], tspan[2]), θ)
+            sol = solve(intprob, strategy.quadrature_alg; abstol = strategy.abstol,
+                reltol = strategy.reltol, maxiters = strategy.maxiters)
+            sol.u
         end
-    end
-    return loss
-end
-
-function generate_loss(strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p,
-        batch, param_estim::Bool)
-    autodiff && throw(ArgumentError("autodiff not supported for WeightedIntervalTraining."))
-    minT = tspan[1]
-    maxT = tspan[2]
-
-    weights = strategy.weights ./ sum(strategy.weights)
-
-    N = length(weights)
-    points = strategy.points
-
-    difference = (maxT - minT) / N
+        return loss
+
+    elseif strategy isa GridTraining
+        ts = tspan[1]:(strategy.dx):tspan[2]
+        autodiff && throw(ArgumentError("autodiff not supported for GridTraining."))
+        function loss(θ, _)
+            if batch
+                inner_loss(phi, f, autodiff, ts, θ, p, param_estim)
+            else
+                sum([inner_loss(phi, f, autodiff, t, θ, p, param_estim) for t in ts])
+            end
+        end
+        return loss
+
+    elseif strategy isa StochasticTraining
+        autodiff && throw(ArgumentError("autodiff not supported for StochasticTraining."))
+        function loss(θ, _)
+            ts = adapt(parameterless_type(θ),
+                [(tspan[2] - tspan[1]) * rand() + tspan[1] for i in 1:(strategy.points)])
+            if batch
+                inner_loss(phi, f, autodiff, ts, θ, p, param_estim)
+            else
+                sum([inner_loss(phi, f, autodiff, t, θ, p, param_estim) for t in ts])
+            end
+        end
+        return loss
+
+    elseif strategy isa WeightedIntervalTraining
+
+        autodiff && throw(ArgumentError("autodiff not supported for WeightedIntervalTraining."))
+        minT = tspan[1]
+        maxT = tspan[2]
+
+        weights = strategy.weights ./ sum(strategy.weights)
+
+        N = length(weights)
+        points = strategy.points
+
+        difference = (maxT - minT) / N
+
+        data = Float64[]
+        for (index, item) in enumerate(weights)
+            temp_data = rand(1, trunc(Int, points * item)) .* difference .+ minT .+
+            ((index - 1) * difference)
+            data = append!(data, temp_data)
+        end
+
+        ts = data
+        function loss(θ, _)
+            if batch
+                inner_loss(phi, f, autodiff, ts, θ, p, param_estim)
+            else
+                sum([inner_loss(phi, f, autodiff, t, θ, p, param_estim) for t in ts])
+            end
+        end
+        return loss
 
-    data = Float64[]
-    for (index, item) in enumerate(weights)
-        temp_data = rand(1, trunc(Int, points * item)) .* difference .+ minT .+
-                    ((index - 1) * difference)
-        data = append!(data, temp_data)
+    else strategy isa QuasiRandomTraining
+        error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.")
     end
+end  
+
 
-    ts = data
-    function loss(θ, _)
-        if batch
-            inner_loss(phi, f, autodiff, ts, θ, p, param_estim)
-        else
-            sum([inner_loss(phi, f, autodiff, t, θ, p, param_estim) for t in ts])
-        end
-    end
-    return loss
-end
 
 function evaluate_tstops_loss(phi, f, autodiff::Bool, tstops, p, batch, param_estim::Bool)
     function loss(θ, _)
@@ -304,9 +312,6 @@ function evaluate_tstops_loss(phi, f, autodiff::Bool, tstops, p, batch, param_es
     return loss
 end
 
-function generate_loss(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, tspan)
-    error("QuasiRandomTraining is not supported by NNODE since it's for high dimensional spaces only. Use StochasticTraining instead.")
-end
 
 struct NNODEInterpolation{T <: ODEPhi, T2}
     phi::T

test/NNDAE_tests.jl

@@ -62,3 +62,40 @@ end
 
     @test ground_sol(0:(1 / 100):(pi / 2))≈sol atol=0.4
 end
+
+@testset "WeightedIntervalTraining" begin
+        println("WeightedIntervalTraining")
+        function f(out, du, u, p, t)
+            out[1] = du[1] -u[2]
+            out[2] = u[1] + u[2]
+        end
+        p = []
+        u0 = [1.0/4.0, 1.0/4.0]
+        tspan = (0.0, 100000.0)
+        differential_vars = [true, false]
+        prob = DAEProblem(f, u0, tspan, differential_vars = differential_vars)
+        true_out_1(t) = exp(-t)/4.0
+        true_out_2(t) = -1.0 * exp(-t)/4.0
+        func = Lux.σ
+        N = 12
+        chain = Lux.Chain(Lux.Dense(1, N, func), Lux.Dense(N, N, func), Lux.Dense(N, N, func),
+                        Lux.Dense(N, N, func), Lux.Dense(N, length(u0)))
+        opt = OptimizationOptimisers.Adam(0.01)
+        weights = [0.7, 0.2, 0.1]
+        points = 200
+        alg = NNDAE(chain, opt,init_params = nothing;  autodiff = false,
+                            strategy = NeuralPDE.WeightedIntervalTraining(weights, points))
+        sol = solve(prob, alg, verbose = false, maxiters = 5000, saveat = 0.01)
+        #@test abs(mean(sol) - mean(true_sol)) < 0.2
+        """Sol would have 2 outputs: one for u[1] and the other for u[2] so just I need compute the total error for all t in tspan """
+        total_error = 0
+        for i in tspan:
+            total_error  = total_error + abs(sol(i) - [true_out_1(i) true_out_2(i)])
+        end
+        if total_error < 0.01:
+            print("It works!")
+        else:
+            print("Total error exceeds bound")
+        end
+    end
+