Can I use this approach with PDE in 2D? #46

kjrathore · 2025-02-14T22:17:07Z

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 · 2025-02-17T14:00:02Z

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

kjrathore · 2025-02-24T23:22:34Z

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 · 2025-03-01T18:58:37Z

What did you try?

kjrathore · 2025-03-03T16:31:50Z

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 · 2025-03-03T17:50:11Z

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

kjrathore · 2025-03-06T14:55:29Z

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

SebastianM-C · 2025-03-06T14:59:09Z

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 · 2025-03-06T15:08:01Z

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 · 2025-03-10T19:02:22Z

@SebastianM-C ?

SebastianM-C · 2025-03-10T19:41:33Z

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.

kjrathore added the question Further information is requested label Feb 14, 2025

ChrisRackauckas assigned SebastianM-C Mar 3, 2025

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Can I use this approach with PDE in 2D? #46

Can I use this approach with PDE in 2D? #46

kjrathore commented Feb 14, 2025

ChrisRackauckas commented Feb 17, 2025

kjrathore commented Feb 24, 2025

ChrisRackauckas commented Mar 1, 2025

kjrathore commented Mar 3, 2025 •

edited

Loading

ChrisRackauckas commented Mar 3, 2025

kjrathore commented Mar 6, 2025

SebastianM-C commented Mar 6, 2025

kjrathore commented Mar 6, 2025

kjrathore commented Mar 10, 2025

SebastianM-C commented Mar 10, 2025 •

edited

Loading

Can I use this approach with PDE in 2D? #46

Can I use this approach with PDE in 2D? #46

Comments

kjrathore commented Feb 14, 2025

ChrisRackauckas commented Feb 17, 2025

kjrathore commented Feb 24, 2025

ChrisRackauckas commented Mar 1, 2025

kjrathore commented Mar 3, 2025 • edited Loading

ChrisRackauckas commented Mar 3, 2025

kjrathore commented Mar 6, 2025

SebastianM-C commented Mar 6, 2025

kjrathore commented Mar 6, 2025

kjrathore commented Mar 10, 2025

SebastianM-C commented Mar 10, 2025 • edited Loading

kjrathore commented Mar 3, 2025 •

edited

Loading

SebastianM-C commented Mar 10, 2025 •

edited

Loading