Can I use this approach with PDE in 2D?

[kjrathore]

Question❓

Can implement the symbolic UDE approach in 2D PDE?
Similar to https://github.com/SciML/MethodOfLines.jl/blob/master/test/pde_systems/MOL_2D_Diffusion.jl
since ModelingToolKit provides a nice way of solving the equations and discretization methods. Would be happy to know other way around to apply missing physics UDE in 2D PDEs.

[ChrisRackauckas]

This should be fine. You'd just use the same apply function from Lux

[kjrathore]

Hi @ChrisRackauckas
I am following the friction example, implemented it for my 1D system. it works!
However, for PDE and discretization:
systems = [nn_in, nn_out] arguments were not recognized by (PDESystem · ModelingToolkit.jl).

Similar to https://discourse.julialang.org/t/methodoflines-pdesystem-with-internal-systems/104409
ERROR: LoadError: type PDESystem has no field nn_in Stacktrace: [1] getproperty(x::PDESystem, sym::Symbol)

[ChrisRackauckas]

What did you try?

[kjrathore]

I tried this:

using ModelingToolkitNeuralNets
using ModelingToolkit
import ModelingToolkit.t_nounits as t
import ModelingToolkit.D_nounits as Dt
using ModelingToolkitStandardLibrary.Blocks
using OrdinaryDiffEq
using Optimization
using OptimizationOptimisers: Adam
using SciMLStructures
using SciMLStructures: Tunable
using SymbolicIndexingInterface
using StableRNGs
using Lux
using Plots, Random, Statistics
using MethodOfLines, LinearAlgebra, DomainSets
using ModelingToolkit: Differential
# using DiffEqFlux, Lux, ComponentArrays, Zygote, StableRNGs
# using Optimization, OptimizationOptimisers
rng = StableRNG(1111)


@parameters x y
Dxx = Differential(x)^2
Dyy = Differential(y)^2

# Time and space domains
t_min, t_max = 0.0, 0.1
x_min, x_max = 0.0, 1.0
y_min, y_max = 0.0, 1.0
dx, dy = 0.1, 0.1
order = 2  # Approximation order

# define variables and constants
@variables N(..) P(..)  # Define as callable functions
@constants a=0.3 b=0.5 c=0.5 d=0.05 e=0.08 μ=0.18 τ=0.3
@constants DN=0.5 DP=0.01

# Boundry conditions
bcs = [
    N(t_min,x,y) ~ 0,
    # Dirichlet-like approximation of Neumann for N
    N(t, x_min, y) ~ 0.9, N(t, x_max, y) ~ 0.9,
    N(t, x, y_min) ~ 0.9, N(t, x, y_max) ~ 0.9,

    P(t_min, x, y) ~ 0,
    # Dirichlet-like approximation of Neumann for P
    P(t, x_min, y) ~ 0.25, P(t, x_max, y) ~ 0.25,
    P(t, x, y_min) ~ 0.25, P(t, x, y_max) ~ 0.25
]

# Space and time domains
domains = [t ∈ Interval(t_min, t_max),
            x ∈ Interval(x_min, x_max),
            y ∈ Interval(y_min, y_max)]


function hab_true()

    eqs = [
        Dt(N(t,x,y)) ~ a - b * N(t,x,y) * P(t,x,y) - e * N(t,x,y) + 
                        DN * (Dxx(N(t,x,y)) + Dyy(N(t,x,y))),
        Dt(P(t,x,y)) ~ c * N(t,x,y) * P(t,x,y) - d * P(t,x,y) 
            - τ * (P(t,x,y)^2 / (μ^2 + P(t,x,y)^2)) 
            + DP * (Dxx(P(t,x,y)) + Dyy(P(t,x,y)))
    ]
    # return ODESystem(eqs, ModelingToolkit.t_nounits, name = :hab_true)
    return PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)], name=:hab_true)
end

model_true = hab_true()
# Method of lines discretization
discretization = MOLFiniteDifference([x => dx, y => dy], t) # Discretize space
prob_true = discretize(model_true, discretization) # Convert PDE → ODE system
# simplified_sys = structural_simplify(discretized_sys)
println("--------------------------------")
println("Discretized the System.")

# Solution of the ODE system
sol_ref = solve(prob_true, Tsit5(); saveat = 0.001)
println("--------------------------------")
println("Solved PDE..")
# println(sol_ref.u)

# # Extract the 3D solution arrays for N and P
N_sol = Array(sol_ref[N(t, x, y)])
P_sol = Array(sol_ref[P(t, x, y)])
println("N sol:",size(N_sol), size(P_sol))
# t = sol.t
N_m = mean(N_sol, dims = 1)
P_m = mean(P_sol, dims = 1)
noise_magnitude = 5e-2
Nₙ = N_sol .+ (noise_magnitude * N_m) .* randn(rng, eltype(N_sol), size(N_sol))
Pₙ = P_sol .+ (noise_magnitude * P_m) .* randn(rng, eltype(P_sol), size(P_sol))
println("--------------------------------")
println("Prepared obs data..")

# anim = @animate for t_idx in 1:length(sol_ref.t)
#     p1 = heatmap(N_sol[t_idx, :, :], title="Nutrient (N) at t=$(round(sol_ref.t[t_idx], digits=2))",
#                  xlabel="x", ylabel="y", color=:viridis)
#     p2 = heatmap(P_sol[t_idx, :, :], title="Phytoplankton (P) at t=$(round(sol_ref.t[t_idx], digits=2))",
#                  xlabel="x", ylabel="y", color=:plasma)
#     plot(p1, p2, layout=(1,2))  # Side-by-side plots
# end

# gif(anim, "N_P_solution.gif", fps=10)
# println("========================================\n")

#discretization space
nx = floor(Int64, (x_max - x_min) / dx) + 1
ny = floor(Int64, (y_max - y_min) / dy) + 1
println("Discretization space: $nx, $ny")

#xy space should be the input to Neural Network.
# we may need to flatten the features before passing to NN.
# 2*nx*ny
L = nx*ny
nnf = 2*nx*ny  #number of feature to NN


function hab_ude()    
    @named nn_in = RealInputArray(nin = nnf)
    @named nn_out = RealOutputArray(nout = nnf)

    println("Shape of nn_out.u:", size(nn_out.u))
        # Dt(N(t,x,y)) ~ a - b * N(t,x,y) * P(t,x,y) - e * N(t,x,y) + 
        #                 DN * (Dxx(N(t,x,y)) + Dyy(N(t,x,y))),
        # Dt(P(t,x,y)) ~ c * N(t,x,y) * P(t,x,y) - d * P(t,x,y) 
        #     - τ * (P(t,x,y)^2 / (μ^2 + P(t,x,y)^2)) + DP * (Dxx(P(t,x,y)) + Dyy(P(t,x,y)))
    eqs = [Dt(N(t,x,y)) ~ a .- reshape(nn_in.u[1:L], nx, ny) .- e .* N(t,x,y),
                N(t,x,y) ~ reshape(nn_out.u[1:L], nx, ny),
            Dt(P(t,x,y)) ~ reshape(nn_in.u[L+1:end], nx, ny) .- d .* P(t,x,y) .- τ .* (P(t,x,y)^2 ./(μ^2 .+ P(t,x,y)^2)),
                P(t,x,y) ~ reshape(nn_out.u[L+1:end], nx, ny)
        ]

    # return ODESystem(eqs, t, name = :hab, systems = [nn_in, nn_out])
    return PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)], 
                name = :hab, systems = [nn_in, nn_out])
end


model = hab_ude()

chain = Lux.Chain(
    Lux.Dense(nnf => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => nnf, use_bias = false)
)
@named nn = NeuralNetworkBlock(nnf, nnf; chain = chain, 
                                rng = rng)
# println("model systems", model.nn_in)

println()
# Combine equations from the model and the neural network connections

eqs = [connect(model.nn_in, nn.output)
       connect(model.nn_out, nn.input)]


ude_sys = complete(PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)], name=:ude_sys))
discretization = MOLFiniteDifference([x => dx, y => dy], t) # Discretize space
sys = discretize(ude_sys, discretization) # Convert PDE → ODE system
println("--------------------------------")
println("Discretized the UDE.")


function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
    new_p = set_x(prob, x)
    new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0))
    ts = sol_ref.t
    new_sol = solve(new_prob, Rodas4(), saveat = ts, abstol = 1e-8, reltol = 1e-8)
    loss = zero(eltype(x))
    for i in eachindex(new_sol.u)
        loss += sum(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i)))
    end
    return SciMLBase.successful_retcode(new_sol) ? loss : Inf
end

of = OptimizationFunction{true}(loss, AutoForwardDiff())

prob = ODEProblem(sys, [], (t_min, t_max), [])
get_vars = getu(sys, [sys.hab.N, sys.hab.P])
get_refs = getu(model_true, [model_true.N, model_true.P])
set_x = setp_oop(sys, sys.nn.p)
x0 = default_values(sys)[sys.nn.p]

cb = (opt_state, loss) -> begin
    @info "step $(opt_state.iter), loss: $loss"
    return false
end

op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs, set_x))
res = solve(op, Adam(5e-3); maxiters = 10, callback = cb)

Outputs:
ERROR: LoadError: type PDESystem has no field nn_in

[ChrisRackauckas]

@SebastianM-C is going to try to get an example together.

[kjrathore]

Hi @SebastianM-C,
I'm happy and eager to get your help on this.

[SebastianM-C]

There are several places where a NN could go, I think that starting with a boundary condition is an option. I'm looking at an example with the heat equation. Are you looking for something else?

[kjrathore]

I couldn't think of boundary condition options. And Heat equation example has 2 parameters (t,x) whereas I have 3 (t, x, y) in my case, although that is manageable. I am trying this:

I tried this:

using ModelingToolkitNeuralNets
using ModelingToolkit
import ModelingToolkit.t_nounits as t
import ModelingToolkit.D_nounits as Dt
using ModelingToolkitStandardLibrary.Blocks
using OrdinaryDiffEq
using Optimization
using OptimizationOptimisers: Adam
using SciMLStructures
using SciMLStructures: Tunable
using SymbolicIndexingInterface
using StableRNGs
using Lux
using Plots, Random, Statistics
using MethodOfLines, LinearAlgebra, DomainSets
using ModelingToolkit: Differential
# using DiffEqFlux, Lux, ComponentArrays, Zygote, StableRNGs
# using Optimization, OptimizationOptimisers
rng = StableRNG(1111)


@parameters x y
Dxx = Differential(x)^2
Dyy = Differential(y)^2

# Time and space domains
t_min, t_max = 0.0, 0.1
x_min, x_max = 0.0, 1.0
y_min, y_max = 0.0, 1.0
dx, dy = 0.1, 0.1
order = 2  # Approximation order

# define variables and constants
@variables N(..) P(..)  # Define as callable functions
@constants a=0.3 b=0.5 c=0.5 d=0.05 e=0.08 μ=0.18 τ=0.3
@constants DN=0.5 DP=0.01

# Boundry conditions
bcs = [
    N(t_min,x,y) ~ 0,
    # Dirichlet-like approximation of Neumann for N
    N(t, x_min, y) ~ 0.9, N(t, x_max, y) ~ 0.9,
    N(t, x, y_min) ~ 0.9, N(t, x, y_max) ~ 0.9,

    P(t_min, x, y) ~ 0,
    # Dirichlet-like approximation of Neumann for P
    P(t, x_min, y) ~ 0.25, P(t, x_max, y) ~ 0.25,
    P(t, x, y_min) ~ 0.25, P(t, x, y_max) ~ 0.25
]

# Space and time domains
domains = [t ∈ Interval(t_min, t_max),
            x ∈ Interval(x_min, x_max),
            y ∈ Interval(y_min, y_max)]


function hab_true()

    eqs = [
        Dt(N(t,x,y)) ~ a - b * N(t,x,y) * P(t,x,y) - e * N(t,x,y) + 
                        DN * (Dxx(N(t,x,y)) + Dyy(N(t,x,y))),
        Dt(P(t,x,y)) ~ c * N(t,x,y) * P(t,x,y) - d * P(t,x,y) 
            - τ * (P(t,x,y)^2 / (μ^2 + P(t,x,y)^2)) 
            + DP * (Dxx(P(t,x,y)) + Dyy(P(t,x,y)))
    ]
    # return ODESystem(eqs, ModelingToolkit.t_nounits, name = :hab_true)
    return PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)], name=:hab_true)
end

model_true = hab_true()
# Method of lines discretization
discretization = MOLFiniteDifference([x => dx, y => dy], t) # Discretize space
prob_true = discretize(model_true, discretization) # Convert PDE → ODE system
# simplified_sys = structural_simplify(discretized_sys)
println("--------------------------------")
println("Discretized the System.")

# Solution of the ODE system
sol_ref = solve(prob_true, Tsit5(); saveat = 0.001)
println("--------------------------------")
println("Solved PDE..")
# println(sol_ref.u)

# # Extract the 3D solution arrays for N and P
N_sol = Array(sol_ref[N(t, x, y)])
P_sol = Array(sol_ref[P(t, x, y)])
println("N sol:",size(N_sol), size(P_sol))
# t = sol.t
N_m = mean(N_sol, dims = 1)
P_m = mean(P_sol, dims = 1)
noise_magnitude = 5e-2
Nₙ = N_sol .+ (noise_magnitude * N_m) .* randn(rng, eltype(N_sol), size(N_sol))
Pₙ = P_sol .+ (noise_magnitude * P_m) .* randn(rng, eltype(P_sol), size(P_sol))
println("--------------------------------")
println("Prepared obs data..")

# anim = @animate for t_idx in 1:length(sol_ref.t)
#     p1 = heatmap(N_sol[t_idx, :, :], title="Nutrient (N) at t=$(round(sol_ref.t[t_idx], digits=2))",
#                  xlabel="x", ylabel="y", color=:viridis)
#     p2 = heatmap(P_sol[t_idx, :, :], title="Phytoplankton (P) at t=$(round(sol_ref.t[t_idx], digits=2))",
#                  xlabel="x", ylabel="y", color=:plasma)
#     plot(p1, p2, layout=(1,2))  # Side-by-side plots
# end

# gif(anim, "N_P_solution.gif", fps=10)
# println("========================================\n")

#discretization space
nx = floor(Int64, (x_max - x_min) / dx) + 1
ny = floor(Int64, (y_max - y_min) / dy) + 1
println("Discretization space: $nx, $ny")

#xy space should be the input to Neural Network.
# we may need to flatten the features before passing to NN.
# 2*nx*ny
L = nx*ny
nnf = 2*nx*ny  #number of feature to NN


function hab_ude()    
    @named nn_in = RealInputArray(nin = nnf)
    @named nn_out = RealOutputArray(nout = nnf)

    println("Shape of nn_out.u:", size(nn_out.u))
        # Dt(N(t,x,y)) ~ a - b * N(t,x,y) * P(t,x,y) - e * N(t,x,y) + 
        #                 DN * (Dxx(N(t,x,y)) + Dyy(N(t,x,y))),
        # Dt(P(t,x,y)) ~ c * N(t,x,y) * P(t,x,y) - d * P(t,x,y) 
        #     - τ * (P(t,x,y)^2 / (μ^2 + P(t,x,y)^2)) + DP * (Dxx(P(t,x,y)) + Dyy(P(t,x,y)))
    eqs = [Dt(N(t,x,y)) ~ a .- reshape(nn_in.u[1:L], nx, ny) .- e .* N(t,x,y),
                N(t,x,y) ~ reshape(nn_out.u[1:L], nx, ny),
            Dt(P(t,x,y)) ~ reshape(nn_in.u[L+1:end], nx, ny) .- d .* P(t,x,y) .- τ .* (P(t,x,y)^2 ./(μ^2 .+ P(t,x,y)^2)),
                P(t,x,y) ~ reshape(nn_out.u[L+1:end], nx, ny)
        ]

    # return ODESystem(eqs, t, name = :hab, systems = [nn_in, nn_out])
    return PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)], 
                name = :hab, systems = [nn_in, nn_out])
end


model = hab_ude()

chain = Lux.Chain(
    Lux.Dense(nnf => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => 10, Lux.mish, use_bias = false),
    Lux.Dense(10 => nnf, use_bias = false)
)
@named nn = NeuralNetworkBlock(nnf, nnf; chain = chain, 
                                rng = rng)
# println("model systems", model.nn_in)

println()
# Combine equations from the model and the neural network connections

eqs = [connect(model.nn_in, nn.output)
       connect(model.nn_out, nn.input)]


ude_sys = complete(PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)], name=:ude_sys))
discretization = MOLFiniteDifference([x => dx, y => dy], t) # Discretize space
sys = discretize(ude_sys, discretization) # Convert PDE → ODE system
println("--------------------------------")
println("Discretized the UDE.")


function loss(x, (prob, sol_ref, get_vars, get_refs, set_x))
    new_p = set_x(prob, x)
    new_prob = remake(prob, p = new_p, u0 = eltype(x).(prob.u0))
    ts = sol_ref.t
    new_sol = solve(new_prob, Rodas4(), saveat = ts, abstol = 1e-8, reltol = 1e-8)
    loss = zero(eltype(x))
    for i in eachindex(new_sol.u)
        loss += sum(abs2.(get_vars(new_sol, i) .- get_refs(sol_ref, i)))
    end
    return SciMLBase.successful_retcode(new_sol) ? loss : Inf
end

of = OptimizationFunction{true}(loss, AutoForwardDiff())

prob = ODEProblem(sys, [], (t_min, t_max), [])
get_vars = getu(sys, [sys.hab.N, sys.hab.P])
get_refs = getu(model_true, [model_true.N, model_true.P])
set_x = setp_oop(sys, sys.nn.p)
x0 = default_values(sys)[sys.nn.p]

cb = (opt_state, loss) -> begin
    @info "step $(opt_state.iter), loss: $loss"
    return false
end

op = OptimizationProblem(of, x0, (prob, sol_ref, get_vars, get_refs, set_x))
res = solve(op, Adam(5e-3); maxiters = 10, callback = cb)

Outputs: ERROR: LoadError: type PDESystem has no field nn_in

[kjrathore]

@SebastianM-C ?

[SebastianM-C]

I'm working on a simpler version first, I think it would be good to start from a 1D PDE and then move to 2D. For the heat equation you can do it in more than 1D.

The main thing is that you can't currently use the block components approach with PDEs, so you would have to reformulate it to the direct calls to LuxCore.stateless_apply (see also #53).

In this case your equations would look something like

function hab_ude(chain)
    # @named nn_in = RealInputArray(nin = nnf)
    # @named nn_out = RealOutputArray(nout = nnf)

    # println("Shape of nn_out.u:", size(nn_out.u))
        # Dt(N(t,x,y)) ~ a - b * N(t,x,y) * P(t,x,y) - e * N(t,x,y) +
        #                 DN * (Dxx(N(t,x,y)) + Dyy(N(t,x,y))),
        # Dt(P(t,x,y)) ~ c * N(t,x,y) * P(t,x,y) - d * P(t,x,y)
        #     - τ * (P(t,x,y)^2 / (μ^2 + P(t,x,y)^2)) + DP * (Dxx(P(t,x,y)) + Dyy(P(t,x,y)))
    @variables nn_in(t)[1:nnf]
    @variables nn_out(t)[1:nnf]

    ps_nn_init, st_nn = Lux.setup(StableRNG(1234), chain)
    ca_nn = ComponentArray{Float64}(ps_nn_init)
    @parameters p_nn[1:length(ca_nn)] = Vector(ca_nn)
    @parameters T_nn::typeof(typeof(ca_nn))=typeof(ca_nn) [tunable = false]

    eqs = [Dt(N(t,x,y)) ~ a .- reshape(nn_in[1:L], nx, ny) .- e .* N(t,x,y),
                N(t,x,y) ~ reshape(nn_out[1:L], nx, ny),
            Dt(P(t,x,y)) ~ reshape(nn_in[L+1:end], nx, ny) .- d .* P(t,x,y) .- τ .* (P(t,x,y)^2 ./(μ^2 .+ P(t,x,y)^2)),
                P(t,x,y) ~ reshape(nn_out[L+1:end], nx, ny),
            nn_out ~ LuxCore.stateless_apply(chain, nn_in, lazyconvert(T_nn, p_nn))
        ]

    # return ODESystem(eqs, t, name = :hab, systems = [nn_in, nn_out])
    return PDESystem(eqs, bcs, domains, [t,x,y], [N(t,x,y), P(t,x,y)],
                name = :hab)
end

but currently there are some issues in how discretization interacts with vector quantities, so it doesn't work for now.

[ChrisRackauckas]

The new SymbolicNeuralNetwork is probably all that's needed here?

[SebastianM-C]

We need more robust vector handling in MOL. See SciML/PDEBase.jl#37 for example. Another thing that @xtalax mentioned is having the symbolic derivatives for the NN via nested AD.

[kjrathore]

@ChrisRackauckas Meanwhile can you suggest me any other way to implement UDE for 2D scenario?