Symbolic ODE/PDE in Vector form with Symbolics/ModelingToolkit

I’ve been having a hard time trying to use vector-valued equations (as opposed to vectors of equations) in Symbolics/ModelingToolkit. I remember reading somewhere in either docs that his feature was somewhat experimental, yet I’ve been unable to find this piece of information again, as well as any tutorial or guidance on how to do this.

Should I just leave this alone? My work involves large ODEs in vector/matrix form and my solver (a PINN) gives me solutions also in vector form; It just feels wrong to have to split both the equation and the solution into their individual lines and link them one by one.

An example of the libraries not working how one (I?) would expect:

using ModelingToolkit

@variables t, u(..), v(..), U(..)[1:2]

#Simple Harmonic Oscillator
Dt = Differential(t) 
eqs = [Dt(u(t)) ~ v(t), Dt(v(t)) ~ -u(t)] #Works
vec_eqs = Dt(U(t)) ~ [0 1;-1 0]*U(t)      #ERROR: axes of Differential(t)((U(t))[1:2]) not known

It follows Julia’s vector semantics, so you need to broadcast that:

@variables t, u(..), v(..), U(..)[1:2]
Dt = Differential(t) 
vec_eqs = Dt.(U(t)) .~ [0 1;-1 0]*U(t)

That could probably be simplified though.

Thanks Chris. I had tried broadcasting, but it seems to mess up the types. In particular, vec_eqs is not an Equation, so it cannot be used with DifferentialEquations/NeuralPDE. Or am I missing something?

using ModelingToolkit

@parameters σ ρ
@variables t x(t) y(t) X(t)[1:2]
D = Differential(t)

eqs = [D(x) ~ σ*y,
       D(y) ~ ρ*x]
vec_eqs = D.(X) .~ [0 σ;ρ 0]*X

typeof(eqs)#Vector{Equation}
typeof(vec_eqs)#Symbolics.Arr{Any,1}

@named sys = ODESystem(eqs,t,[x,y],[σ,ρ],tspan=(0, 1000.0)) #Works
@named vec_sys = ODESystem(vec_eqs,t,[X],[σ,ρ],tspan=(0, 1000.0))#ERROR

@xtalax were you working on vector derivatives in NeuralPDE?

Please correct me if I’m saying something very silly, but it seems like NeuralPDE itself doesn’t really care if it’s dealing with vectors or scalars? As long as the equations themselves are understood by Symbolics things seem to work.

For example, take this (trivial) system of two conservation laws with a single PINN:

#Vector_Neural_PDE
using NeuralPDE, Lux, Optimization, OptimizationOptimJL
import ModelingToolkit: Interval

@parameters x t
@variables u(..) #Declare u as if it were a scalar
Dx = Differential(x)
Dt  = Differential(t)

#Set PDE normally
c = 1
eq  = Dt(u(x,t)) ~ c*Dx(u(x,t))

#Give a vector disguised as a scalar for boundary conditions
bcs = [u(x,0) .~ exp(-x^2)] 

# Space and time domains
domains = [x ∈ Interval(-5,5),
           t ∈ Interval(0.0,10.0)]

# PINN with TWO outputs
dim_in = 2
dim_out = 2
chain = Lux.Chain(Dense(dim_in,12,Lux.tanh),Dense(12,12,Lux.tanh),Dense(12,12,Lux.tanh),Dense(12,dim_out))

# Discretization
dx = 0.05
discretization = PhysicsInformedNN(chain,GridTraining(dx))
@named pde_system = PDESystem(eq,bcs,domains,[x,t],[u(x, t)])
prob = discretize(pde_system,discretization)

#Optimizer
opt = OptimizationOptimJL.BFGS()

res = Optimization.solve(prob, opt, maxtime=20)

All of this seems to work just fine, with both outputs of our PINN being trained accordingly:

using Plots

phi = discretization.phi

xs,ts = [infimum(d.domain):dx/10:supremum(d.domain) for d in domains]

U(x,t) = phi([x,t],res.u)[1]
V(x,t) = phi([x,t],res.u)[2]

plot(xs,U.(xs,0))
plot!(xs,U.(xs,1))
plot!(xs,U.(xs,2))

plot(xs,V.(xs,0))
plot!(xs,V.(xs,1))
plot!(xs,V.(xs,2))

Yes, that works. But the training process could specialize on vector calculus to improve its operations.

Yes, see here Add ArrayDifferentialOperators for Vector calculus by xtalax · Pull Request #942 · JuliaSymbolics/Symbolics.jl · GitHub