I am trying to train a neural network which is a part of a PDE, I am solving the PDE with MOL and ODE solve. Below is the forward pass i.e getting the solution with ODE solve which runs fine.
import ModelingToolkit: Interval, infimum, supremum using DiffEqSensitivity using OrdinaryDiffEq, ModelingToolkit, MethodOfLines, DomainSets using Plots, Flux using Interpolations x_dom = 0:0.1:2 dim = 1 net = Flux.Chain(Dense(dim,16,Flux.σ),Dense(16,16,Flux.σ),Dense(16,1)) p = Flux.params(net) # p, re = Flux.destructure(net) A = [net([x1]) for x1 in x_dom] itp = interpolate(A, BSpline(Cubic(Line(OnGrid())))) itps = scale(itp, x_dom) net_func(x) = itps(x) @parameters t x @variables u(..) @register_symbolic net_func(x) Dxx = Differential(x)^2 Dt = Differential(t) t_min= 0. t_max = 2.0 x_min = 0. x_max = 2. y_min = 0. y_max = 2. dim = 2 eq = Dt(u(t,x)) ~ Dxx(u(t,x)) + net_func(x) bcs = [u(t_min,x) ~ sin(2*pi*x), u(t,x_min) ~ 0.0, u(t,x_max) ~ 1.0] domains = [t ∈ Interval(t_min,t_max), x ∈ Interval(x_min,x_max)] dx = 0.1 order = 2 discretization = MOLFiniteDifference([x => dx], t, approx_order=order) @named pde_system = PDESystem(eq,bcs,domains,[t,x],[u(t, x)]) prob = discretize(pde_system,discretization) sol = solve(prob, p = p, Tsit5(), saveat=0.1)
However, I am not able to train this network, below is the code for the loss function I am using, which is chosen just to see if a backward pass is possible or not.
function predict_sol() Array(solve(prob,Tsit5(),p = p, progress=true)) end function loss_func() _sol = predict_sol() loss = sum(abs2,_sol) @show loss loss end data = Iterators.repeated((), 2) opt = ADAM(0.025) Flux.train!(loss_func, p, data, opt)
I am curious if it is possible to train a network within MOL framework if so, then can someone help me figure out where I might be wrong?