universal DE. - possible?

[bgladman]

Question❓

In the Julia Con 2022 talk, it was suggested that is was possible to embed a NN as part of a system of PDEs. I am struggling with figuring out how to specify the weights (as some component vector, θ, which I want to optimise) as parameters. Something along the lines of using @parameters P, and @named sys = PDESystem(eqs, bcs, domain, [t, z, r], [C(t, z), q(t, z, r), V(t)], [P]; defaults=Dict(P => collect(θ)) ).
where eqs includes the NN([V,t], P) and then,

newprob = remake(prob, p =[collect(θ)])

Is this doable, if so what do I need to change to make this work?

Thanks in advance.

[ChrisRackauckas]

MTK supports it. I haven't tried specifically with the MethodOfLines discretizer, but if you use a registered symbolic function for the array I don't see why it wouldn't work given the MTK support. It would be good to test it better here, though for these applications I usually get a little bit more manual right now with the codegen since there's some tricks that the codegen wouldn't do here.

[bgladman]

Thanks Chris!. This could very well be a case of me running before learning to walk. I tried adapting the "adding Parameters" example from the Docs to optimise Dn and Dp as a test case, but hitting snags. Mwe as follows.

using ModelingToolkit, MethodOfLines, OrdinaryDiffEq, DomainSets
using SciMLStructures: replace!, Tunable
using Optimization, OptimizationPolyalgorithms, SciMLSensitivity,
      Zygote, Plots, ComponentArrays
      
@parameters t x
@parameters Dn, Dp [tunable = true]
@variables u(..) v(..)
Dt = Differential(t)
Dx = Differential(x)
Dxx = Differential(x)^2

eqs = [Dt(u(t, x)) ~ Dn * Dxx(u(t, x)) + u(t, x) * v(t, x),# + flow_rate_f(t),
    Dt(v(t, x)) ~ Dp * Dxx(v(t, x)) - u(t, x) * v(t, x)]
bcs = [u(0, x) ~ sin(pi * x / 2),
    v(0, x) ~ sin(pi * x / 2),
    u(t, 0) ~ 0.0, Dx(u(t, 1)) ~ 0.0, #sitp_f(t),
    v(t, 0) ~ 0.0, Dx(v(t, 1)) ~ 0.0]

domains = [t ∈ Interval(0.0, 1.0),
    x ∈ Interval(0.0, 1.0)]

@named pdesys = PDESystem(
    eqs, bcs, domains, [t, x], [u(t, x), v(t, x)],
    [Dn, Dp]; defaults = Dict(Dn => 0.5, Dp => 2.0))

discretization = MOLFiniteDifference([x => 0.1], t)

prob = discretize(pdesys, discretization) # This gives an ODEProblem since it's time-dependent
sol = solve(prob, Tsit5(), saveat=0.1)
solu = sol[u(t, x)]

function loss(p)
    newprob = remake(prob)
    replace!(Tunable(), newprob.p, [p.Dn, p.Dp]) #[p[1], p[2]])
    sol =  solve(newprob, Tsit5(), 
    sensealg = ForwardDiffSensitivity(),saveat = 0.1)
    sol_pr = sol[u(t, x)]
    loss = sum(abs2, sol_pr .- solu)
    return loss
end

p = ComponentArray(Dn=0.3, Dp=1.0)

callback = function (state, l)
    display(l)
    return false
end

adtype = Optimization.AutoZygote()
optf = Optimization.OptimizationFunction((x, p) -> loss(x), adtype)
optprob = Optimization.OptimizationProblem(optf, p)

result_ode = solve(optprob, PolyOpt(),
    callback = callback,
    maxiters = 100)

This gives the following Error.

ERROR: MethodError: no method matching SciMLBase.PDETimeSeriesSolution(::SciMLBase.PDETimeSeriesSolution{…}, ::MethodOfLines.MOLMetadata{…})

Closest candidates are:
  SciMLBase.PDETimeSeriesSolution(::SciMLBase.AbstractODESolution{T}, ::MethodOfLines.MOLMetadata) where T
   @ MethodOfLines ~/.julia/packages/MethodOfLines/EiZyI/src/interface/solution/timedep.jl:11

MOL) pkg> status
Status `~/Documents/Julia/MOL/Project.toml`
  [28f2ccd6] ApproxFun v0.13.27
  [2445eb08] DataDrivenDiffEq v1.5.0
  [5b588203] DataDrivenSparse v0.1.2
  [aae7a2af] DiffEqFlux v4.1.0
  [1130ab10] DiffEqParamEstim v2.2.0
  [0c46a032] DifferentialEquations v7.15.0
  [cf66c380] FastChebInterp v1.2.0
  [442a2c76] FastGaussQuadrature v1.0.2
  [587475ba] Flux v0.14.25
  [f6369f11] ForwardDiff v0.10.38
  [56d4f2e9] Gridap v0.18.7
  [a98d9a8b] Interpolations v0.15.1
  [2c470bb0] Kronecker v0.5.5
  [94925ecb] MethodOfLines v0.11.7
  [961ee093] ModelingToolkit v9.52.0
  [76087f3c] NLopt v1.1.1
  [315f7962] NeuralPDE v5.17.0
  [4e6fcdb7] OptimizationNLopt v0.3.2
  [500b13db] OptimizationPolyalgorithms v0.3.0
  [bac558e1] OrderedCollections v1.6.3
  [731186ca] RecursiveArrayTools v3.27.3
  [0bca4576] SciMLBase v2.62.0
  [1ed8b502] SciMLSensitivity v7.71.2
  [53ae85a6] SciMLStructures v1.6.1
  [860ef19b] StableRNGs v1.0.2
  [2913bbd2] StatsBase v0.34.3
  [c3572dad] Sundials v4.26.1
  [2efcf032] SymbolicIndexingInterface v0.3.35
  [0c5d862f] Symbolics v6.21.0
  [e88e6eb3] Zygote v0.6.73

Any ideas on where this is going wrong?

Thanks.