Adding Training Strategies to dae_solvers.jl

src/dae_solve.jl

@@ -79,6 +79,35 @@ function generate_loss(strategy::GridTraining, phi, f, autodiff::Bool, tspan, p,
     return loss
 end
 
+function generate_loss(
+        strategy::WeightedIntervalTraining, phi, f, autodiff::Bool, tspan, p,
+        differential_vars::AbstractVector)
+    autodiff && throw(ArgumentError("autodiff not supported for GridTraining."))
+    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(inner_loss(phi, f, autodiff, ts, θ, p, differential_vars))
+    end
+    return loss
+end
+
 function SciMLBase.__solve(prob::SciMLBase.AbstractDAEProblem,
         alg::NNDAE,
         args...;
@@ -138,6 +167,8 @@ function SciMLBase.__solve(prob::SciMLBase.AbstractDAEProblem,
         else
             error("dt is not defined")
         end
+    else
+        alg.strategy
     end
 
     inner_f = generate_loss(strategy, phi, f, autodiff, tspan, p, differential_vars)

src/ode_solve.jl

@@ -225,6 +225,7 @@ 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))
@@ -304,6 +305,8 @@ function generate_loss(
     return loss
 end
 
+
+
 function evaluate_tstops_loss(phi, f, autodiff::Bool, tstops, p, batch, param_estim::Bool)
     function loss(θ, _)
         if batch
@@ -319,6 +322,7 @@ function generate_loss(strategy::QuasiRandomTraining, phi, f, autodiff::Bool, ts
     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
     θ::T2

test/NNDAE_tests.jl

@@ -54,7 +54,39 @@ end
     prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
     chain = Lux.Chain(Lux.Dense(1, 15, Lux.σ), Lux.Dense(15, 2))
     opt = OptimizationOptimisers.Adam(0.1)
-    alg = NNDAE(chain, OptimizationOptimisers.Adam(0.1); autodiff = false)
+    alg = NeuralPDE.NNDAE(chain, OptimizationOptimisers.Adam(0.1); autodiff = false)
+
+    sol = solve(prob,
+        alg, verbose = false, dt = 1 / 100.0f0,
+        maxiters = 3000, abstol = 1.0f-10)
+
+    @test ground_sol(0:(1 / 100):(pi / 2))≈sol atol=0.4
+end
+
+@testset "WeightedIntervalTraining" begin
+    function example2(du, u, p, t)
+        du[1] = u[1] - t
+        du[2] = u[2] - t
+        nothing
+    end
+    M = [0.0 0
+         0 1]
+    u₀ = [0.0, 0.0]
+    du₀ = [0.0, 0.0]
+    tspan = (0.0f0, pi / 2.0f0)
+    f = ODEFunction(example2, mass_matrix = M)
+    prob_mm = ODEProblem(f, u₀, tspan)
+    ground_sol = solve(prob_mm, Rodas5(), reltol = 1e-8, abstol = 1e-8)
+
+    example = (du, u, p, t) -> [u[1] - t - du[1], u[2] - t - du[2]]
+    differential_vars = [false, true]
+    prob = DAEProblem(example, du₀, u₀, tspan; differential_vars = differential_vars)
+    chain = Lux.Chain(Lux.Dense(1, 15, Lux.σ), Lux.Dense(15, 2))
+    opt = OptimizationOptimisers.Adam(0.1)
+    weights = [0.7, 0.2, 0.1]
+    points = 200
+    alg = NeuralPDE.NNDAE(chain, OptimizationOptimisers.Adam(0.1),
+        strategy = NeuralPDE.WeightedIntervalTraining(weights, points); autodiff = false)
 
     sol = solve(prob,
         alg, verbose = false, dt = 1 / 100.0f0,