This issue is used to trigger TagBot; feel free to unsubscribe.

If you haven't already, you should update your TagBot.yml to include issue comment triggers. Please see this post on Discourse for instructions and more details.

If you'd like for me to do this for you, comment TagBot fix on this issue. I'll open a PR within a few hours, please be patient!

Seems like the 2D Poisson equation with Neural adapter test is broken. I tested it in master and it failed. Seems to be related to ChainRulesCore.

ERROR: LoadError: MethodError: no method matching *(::Tuple{Int64, Int64})
Closest candidates are:
  *(::Any, ::ChainRulesCore.Tangent) at /Users/gabrielbirnbaum/.julia/packages/ChainRulesCore/1L9My/src/tangent_arithmetic.jl:151
  *(::Any, ::ChainRulesCore.AbstractThunk) at /Users/gabrielbirnbaum/.julia/packages/ChainRulesCore/1L9My/src/tangent_arithmetic.jl:125
  *(::Any, ::ChainRulesCore.ZeroTangent) at /Users/gabrielbirnbaum/.julia/packages/ChainRulesCore/1L9My/src/tangent_arithmetic.jl:104

using Flux
using DiffEqFlux
using ModelingToolkit
using Test, NeuralPDE
using GalacticOptim
using SciMLBase
import ModelingToolkit: Interval

## Example, 2D Poisson equation with Neural adapter
println("Example, 2D Poisson equation with Neural adapter")
@parameters x y
@variables u(..)
Dxx = Differential(x)^2
Dyy = Differential(y)^2

# 2D PDE
eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

# Initial and boundary conditions
bcs = [u(0,y) ~ 0.0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
       u(x,0) ~ 0.0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
# Space and time domains
domains = [x ∈ Interval(0.0,1.0),
           y ∈ Interval(0.0,1.0)]
quadrature_strategy = NeuralPDE.QuadratureTraining(reltol=1e-2,abstol=1e-2,
                                                   maxiters =50, batch=100)
inner = 8
af = Flux.tanh
chain1 = Chain(Dense(2,inner,af),
               Dense(inner,inner,af),
               Dense(inner,1))
initθ = Float64.(DiffEqFlux.initial_params(chain1))
discretization = NeuralPDE.PhysicsInformedNN(chain1,
                                             quadrature_strategy;
                                             init_params = initθ)

@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x,y)])
prob = NeuralPDE.discretize(pde_system,discretization)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)
res = GalacticOptim.solve(prob, BFGS();  maxiters=2000) #  LoadError: MethodError: no method matching *(::Tuple{Int64, Int64})

Trying to run PINN fp doc example, the MTK5 update of

@parameters x @variables p(..) Dx = Differential(x) Dxx = Differential(x)^2

yields the following error:

MethodError: no method matching ^(::Differential, ::Int64) Closest candidates are: ^(!Matched::Float32, ::Integer) at math.jl:907 ^(!Matched::Irrational{:ℯ}, ::Integer) at mathconstants.jl:91 ^(!Matched::Irrational{:ℯ}, ::Number) at mathconstants.jl:91 ...

Stacktrace: [1] macro expansion at .\none:0 [inlined] [2] literal_pow(::typeof(^), ::Differential, ::Val{2}) at .\none:0 [3] top-level scope at In[3]:4 [4] include_string(::Function, ::Module, ::String, ::String) at .\loading.jl:1091

Issue

Looking at the 1D wave equation example, I am not sure if the presented solution is correct. I discussed this on a julia discourse thread. To briefly summarize the solution in that NeuralPDE.jl example isn't what I would expect to see, and it doesn't match the results from a matlab wave-solver when I try to replicate the results.

Then when I went to convert the NeuralPDE example to a "purely" ModelingToolkit version I got a solution that looked like it was from a diffusion problem. I did have a similar problem when I wrote a custom wave-solver using a spectral method, which ended up being due to improperly defined initial conditions (i.e. du(0,x)/dx = 0 instead of the derivative of the u(0,x)). So I exchanged Dt(u(0,x)) ~ 0. for Dx(u(0,x)) ~ 1-2x in the bcs without it having any effect.

Matlab Code

clearvars;

% =========================================================================
% SIMULATION
% =========================================================================

% create the computational grid
Lx =1;           
dx = 0.1;                 % grid point spacing in the x direction [m]    
Nx = round(Lx/dx); % number of grid points in the x (row) direction

kgrid = kWaveGrid(Nx, dx);

% define the properties of the propagation medium
medium.sound_speed = 1;  % [m/s] 

kgrid.makeTime(medium.sound_speed, [], 5);

% create initial pressure distribution
x = 0:dx:Lx-dx;
p0 = x.*(1-x);
source.p0=p0';

% define a sensor
sensor.mask = zeros(1,Nx);

% run the simulation
args = {'PMLInside', false, 'RecordMovie', true};
sensor_data = kspaceFirstOrder1D(kgrid, medium, source, sensor,args{:});

Matlab Results

Keep in mind the solution below has a larger spatial dim than the NeuralPDE example, due to the fact that this wave solver is using a spectral method and has to have a perfectly-matched-layer to avoid "wrap around" effects.

Converted NeuralPDE Example

using Plots, DifferentialEquations, ModelingToolkit, DiffEqOperators

@parameters t, x
@variables u(..)
Dxx = Differential(x)^2
Dtt = Differential(t)^2
Dt = Differential(t)
Dx = Differential(x)

#2D PDE
C=1
eq  = Dtt(u(t,x)) ~ C^2*Dxx(u(t,x))

# Initial and boundary conditions
bcs = [u(t,0) ~ 0.,# for all t > 0
       u(t,1) ~ 0.,# for all t > 0
       u(0,x) ~ x*(1. - x), #for all 0 < x < 1
       Dt(u(0,x)) ~ 0.0, #for all  0 < x < 1
]

# Space and time domains
domains = [t ∈ (0.0,1.0),
           x ∈ (0.0,1.0)]
# Method of lines discretization
dx = 0.1
order = 2
discretization = MOLFiniteDifference([x=>dx],t)

# PDE system
pdesys = PDESystem(eq,bcs,domains,[t,x],[u(t,x)])


# Convert the PDE problem into an ODE problem
prob = discretize(pdesys,discretization)

# Solve ODE problem
sol = solve(prob)

# Plot results
anim = @animate for i ∈ 1:length(sol.t)
plot(sol.u[i], label = "wave", ylims =[-0.25, 0.25])
end every 5

gif(anim, "1Dwave.gif", fps = 10)

Converted NeuralPDE Results

I tried to adapt the https://neuralpde.sciml.ai/dev/pinn/2D/ GPU tutorial to a system of PDEs and unfortunately failed. I need to turn the initθ into a CuArray but I get a warning that Scalar indexing is disallowed. What is the performant/correct way to do the mapping I am doing here with CUDA?

using Flux, CUDA, DiffEqFlux
chain = [FastChain(FastDense(3, 16, Flux.σ), FastDense(16,16,Flux.σ), FastDense(16, 1)),
         FastChain(FastDense(2, 16, Flux.σ), FastDense(16,16,Flux.σ), FastDense(16, 1))]

initθ = map(c -> CuArray(Float64.(c)), DiffEqFlux.initial_params.(chain))

I started working on @zoemcc's PR last weekend because 1. I really want us to support heterogeneous inputs and 2. I thought it would be a good learning opportunity given that Zoe has done such a great job thinking about the architecture of this. My goal was to bring this PR up to date with the code base and get rid of remaining bugs so we could merge.

At this point this is almost done and you can check out Zoe's PR for an explanation of how this works https://github.com/SciML/NeuralPDE.jl/pull/298. This is mostly her work and I just opened a new branch because it was more pleasant to merge it bit-by-bit. I hope that is ok!

The PR is up to date with master, and I fixed a few bugs. The main bug that remains is that I still get a LoadError: DimensionMismatch("A has dimensions (3,1) but B has dimensions (2,30)") when I try to run a heterogeneous system. I am still not sure where this bug is coming from, but it probably has something to do with the fact that you need a array of chains of different input sizes to get heterogeneous systems to work and we didn't handle that correctly somewhere (probably in build_symbolic_function or discretize).

I gave this issue another shot with fresh eyes and made some real progress. The symbolic transformation now looks right.

For this system

@parameters v x t
@variables f(..)
Iv = Integral((t,x) in DomainSets.ProductDomain(ClosedInterval(-Inf, Inf),ClosedInterval(-Inf, Inf)))
Dx = Differential(x)
eqs_ = Iv(f(t, x, v)*x) + Dx(f(t,x,v)) ~ π

bcs = [f(0,x,v) ~ 2]

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

# Neural Network
chain = [FastChain(FastDense(3, 16, Flux.σ), FastDense(16,16,Flux.σ), FastDense(16, 1))]
initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))

discretization = NeuralPDE.PhysicsInformedNN(chain, QuadratureTraining(), init_params= initθ)
@named pde_system = PDESystem(eqs_, bcs, domains, [t,x,v], [f(t,x,v)])
prob = SciMLBase.symbolic_discretize(pde_system, discretization)
prob = SciMLBase.discretize(pde_system, discretization)

cb = function (p,l)
    println("Current loss is: $l")
    return false
end

res = GalacticOptim.solve(prob, BFGS(); cb=cb, maxiters=100)

The symbolic transformation is

(Expr[:((cord, var"##θ#333", phi, derivative, integral, u, p)->begin
          begin
              (var"##θ#3331",) = (var"##θ#333"[1:353],)
              (phi1,) = (phi[1],)
              let (t, x, v) = (cord[[1], :], cord[[2], :], cord[[3], :])
                  begin
                      cord1 = vcat(t, x, v)
                  end
                  (+).(integral(u, cord1, phi, [1, 2], RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#333"), :phi, :derivative, :integral, :u, :p), Main.NeuralPDE.var"#_RGF_ModTag", Main.NeuralPDE.var"#_RGF_ModTag", (0x1892a7ea, 0xc1142249, 0x44da67b6, 0x4db8ecb4, 0x0fa6b1a4)}(quote
    begin
        (var"##θ#3331",) = (var"##θ#333"[1:353],)
        (phi1,) = (phi[1],)
        let (t, x, v) = (fill(t ./ (1 .- t .^ 2), size(cord[[1], :])), fill(x ./ (1 .- x .^ 2), size(cord[[1], :])), cord[[1], :])
            begin
                cord1 = vcat(t, x, v)
            end
            (*).((*).(x, u(cord1, var"##θ#3331", phi1)), (/).((+).(1, (^).(t, 2)), (^).((-).(1, (^).(t, 2)), 2)), (/).((+).(1, (^).(x, 2)), (^).((-).(1, (^).(x, 2)), 2)))
        end
    end
end), Any[-1.0, -1.0], Any[1.0, 1.0], var"##θ#333"), derivative(phi1, u, cord1, [[0.0, 6.0554544523933395e-6, 0.0]], 1, var"##θ#3331")) .- π
              end
          end
      end)], Expr[:((cord, var"##θ#333", phi, derivative, integral, u, p)->begin
          begin
              (var"##θ#3331",) = (var"##θ#333"[1:353],)
              (phi1,) = (phi[1],)
              let (t, x, v) = (fill(0, size(cord[[1], :])), cord[[1], :], cord[[2], :])
                  begin
                      cord1 = vcat(t, x, v)
                  end
                  u(cord1, var"##θ#3331", phi1) .- 2
              end
          end
      end)])

I am opening this as a draft PR, because I am still getting a LoadError: TypeError: non-boolean (Symbolics.Num) used in boolean context in solve.

In Example1 2D Poisson and many other examples, there is a bug in the syntax for Discretization:

dx = 0.05 discretization = PhysicsInformedNN(chain, GridTraining(dx))

Returns:

MethodError: no method matching GridTraining() Closest candidates are: GridTraining(!Matched::Any) at C:\Users\Denis.julia\dev\NeuralPDE.jl-master\src\pinns_pde_solve.jl:64

Stacktrace: [1] PhysicsInformedNN(::Function, ::GridTraining) at C:\Users\Denis.julia\dev\NeuralPDE.jl-master\src\pinns_pde_solve.jl:28 [2] top-level scope at In[13]:3 [3] include_string(::Function, ::Module, ::String, ::String) at .\loading.jl:1091

MWE:

using DiffEqFlux, Flux, NeuralPDE, ModelingToolkit, DomainSets, Optimization, OptimizationFlux, Test

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

# 2D PDE
eq  = Dxx(u(x,y)) + Dyy(u(x,y)) ~ -sin(pi*x)*sin(pi*y)

# Initial and boundary conditions
bcs = [u(0,y) ~ 0.0, u(1,y) ~ -sin(pi*1)*sin(pi*y),
       u(x,0) ~ 0.0, u(x,1) ~ -sin(pi*x)*sin(pi*1)]
# Space and time domains
domains = [x ∈ Interval(0.0,1.0),
           y ∈ Interval(0.0,1.0)]

@named pde_system = PDESystem(eq,bcs,domains,[x,y],[u(x, y)])

fastchain = FastChain(FastDense(2,12,Flux.σ),FastDense(12,12,Flux.σ),FastDense(12,1))
fluxchain = Chain(Dense(2,12,Flux.σ),Dense(12,12,Flux.σ),Dense(12,1))
initθ = Float64.(DiffEqFlux.initial_params(fastchain))
grid_strategy = NeuralPDE.GridTraining(0.1)

p,re = Flux.destructure(fluxchain)

discretization1 = NeuralPDE.PhysicsInformedNN(fastchain,
                                             grid_strategy;
                                             init_params = initθ)

discretization2 = NeuralPDE.PhysicsInformedNN(fluxchain,
                                             grid_strategy;
                                             init_params = initθ)


prob1 = NeuralPDE.discretize(pde_system,discretization1)
prob2 = NeuralPDE.discretize(pde_system,discretization2)
sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization1)

Zygote.gradient((x)->prob1.f(x,nothing),initθ)
Zygote.gradient((x)->prob2.f(x,nothing),initθ) # Very very different???

function callback(p,l)
    @show l
    false
end
res = Optimization.solve(prob1, ADAM(0.1); callback=callback,maxiters=1000)
phi = discretization1.phi

xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2)

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@show maximum(abs2,u_predict - u_real)
@test u_predict ≈ u_real atol = 2.0

res = Optimization.solve(prob2, ADAM(0.1); callback=callback,maxiters=1000)
phi = discretization2.phi

xs,ys = [infimum(d.domain):0.01:supremum(d.domain) for d in domains]
analytic_sol_func(x,y) = (sin(pi*x)*sin(pi*y))/(2pi^2)

u_predict = reshape([first(phi([x,y],res.minimizer)) for x in xs for y in ys],(length(xs),length(ys)))
u_real = reshape([analytic_sol_func(x,y) for x in xs for y in ys], (length(xs),length(ys)))
diff_u = abs.(u_predict .- u_real)

@show maximum(abs2,u_predict - u_real)
@test u_predict ≈ u_real atol = 2.0

See fluxchain fails and the gradient is off.

I tried the 2D Poisson's equation Tutorial () and got the following error executing the solver:

`julia> res = GalacticOptim.solve(prob, BFGS(), progress = false; cb = cb, maxiters=1000)

ERROR: MethodError: no method matching Optim.Options(; extended_trace=true, callback=GalacticOptim.var"#_cb#25"{var"#1#2",BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Nothing,Flat},Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}}(var"#1#2"(), BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Nothing,Flat}(LineSearches.InitialStatic{Float64} alpha: Float64 1.0 scaled: Bool false , LineSearches.HagerZhang{Float64,Base.RefValue{Bool}} delta: Float64 0.1 sigma: Float64 0.9 alphamax: Float64 Inf rho: Float64 5.0 epsilon: Float64 1.0e-6 gamma: Float64 0.66 linesearchmax: Int64 50 psi3: Float64 0.1 display: Int64 0 mayterminate: Base.RefValue{Bool} , nothing, nothing, Flat()), Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}((GalacticOptim.NullData(),)), Core.Box(2), Core.Box(GalacticOptim.NullData()), Core.Box(#undef)), iterations=1000, progress=false) Closest candidates are: Optim.Options(; x_tol, f_tol, g_tol, x_abstol, x_reltol, f_abstol, f_reltol, g_abstol, g_reltol, outer_x_tol, outer_f_tol, outer_g_tol, outer_x_abstol, outer_x_reltol, outer_f_abstol, outer_f_reltol, outer_g_abstol, outer_g_reltol, f_calls_limit, g_calls_limit, h_calls_limit, allow_f_increases, allow_outer_f_increases, successive_f_tol, iterations, outer_iterations, store_trace, trace_simplex, show_trace, extended_trace, show_every, callback, time_limit) at /home/ah/.julia/packages/Optim/auGGa/src/types.jl:73 got unsupported keyword argument "progress" Optim.Options(::T, ::T, ::T, ::T, ::T, ::T, ::T, ::T, ::T, ::T, ::T, ::T, ::Int64, ::Int64, ::Int64, ::Bool, ::Bool, ::Int64, ::Int64, ::Int64, ::Bool, ::Bool, ::Bool, ::Bool, ::Int64, ::TCallback, ::Float64) where {T, TCallback} at /home/ah/.julia/packages/Optim/auGGa/src/types.jl:44 got unsupported keyword arguments "extended_trace", "callback", "iterations", "progress" Stacktrace: [1] kwerr(::NamedTuple{(:extended_trace, :callback, :iterations, :progress),Tuple{Bool,GalacticOptim.var"#_cb#25"{var"#1#2",BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Nothing,Flat},Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}},Int64,Bool}}, ::Type{T} where T) at ./error.jl:157 [2] __solve(::OptimizationProblem{true,OptimizationFunction{true,GalacticOptim.AutoZygote,NeuralPDE.var"#loss_function#191"{NeuralPDE.var"#177#183"{Int64,NeuralPDE.var"#175#181"{NeuralPDE.var"#168#170"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#73"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#73"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#73"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#172#173"{Float32},NeuralPDE.var"#inner_loss#179"}},NeuralPDE.var"#177#183"{Int64,NeuralPDE.var"#175#181"{NeuralPDE.var"#168#170"{FastChain{Tuple{FastDense{typeof(σ),DiffEqFlux.var"#initial_params#73"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(σ),DiffEqFlux.var"#initial_params#73"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}},FastDense{typeof(identity),DiffEqFlux.var"#initial_params#73"{typeof(Flux.glorot_uniform),typeof(Flux.zeros),Int64,Int64}}}}},NeuralPDE.var"#172#173"{Float32},NeuralPDE.var"#inner_loss#179"}}},Nothing,Nothing,Nothing,Nothing,Nothing,Nothing},Array{Float32,1},DiffEqBase.NullParameters,Nothing,Nothing,Nothing,Base.Iterators.Pairs{Union{},Union{},Tuple{},NamedTuple{(),Tuple{}}}}, ::BFGS{LineSearches.InitialStatic{Float64},LineSearches.HagerZhang{Float64,Base.RefValue{Bool}},Nothing,Nothing,Flat}, ::Base.Iterators.Cycle{Tuple{GalacticOptim.NullData}}; cb::Function, maxiters::Int64, kwargs::Base.Iterators.Pairs{Symbol,Bool,Tuple{Symbol},NamedTuple{(:progress,),Tuple{Bool}}}) at /home/ah/.julia/packages/GalacticOptim/TfGcK/src/solve.jl:208 [3] #solve#1 at /home/ah/.julia/packages/GalacticOptim/TfGcK/src/solve.jl:12 [inlined] [4] top-level scope at REPL[16]:1 ` Should the tutorial code be updated or is there another problem?

I ran: https://github.com/JuliaDiffEq/NeuralNetDiffEq.jl/blob/master/test/NNPDENS_tests.jl

For "Black-Scholes-Barenblatt equation"

ans = solve(prob, pdealg, verbose=true, maxiters=250, trajectories=m,
alg=EM(), dt=dt, pabstol = 1f-6)

It says: MethodError: no method matching Float32(::Tracker.TrackedReal{Float32}) Closest candidates are: Float32(::Real, !Matched::RoundingMode) where T<:AbstractFloat at rounding.jl:200 Float32(::T) where T<:Number at boot.jl:718 Float32(!Matched::Int8) at float.jl:60

"Nonlinear Black-Scholes Equation with Default Risk"

@time ans = solve(prob, pdealg, verbose=true, maxiters=200, trajectories=m,
                            alg=EM(), dt=dt, pabstol = 1f-6)

Says: MethodError: vcat(::TrackedArray{…,Array{Float32,1}}, ::Array{Tracker.TrackedReal{Float32},1}) is ambiguous. Candidates:

Currently something like u(x) - u(0) ~ sin(x) gets parsed as u(x) - u(x) ~ sin(x) when generating the loss function, because all instances of u(anything) are assumed to have the same input. I've made changes to _transform_expression and dot that generate the correct loss function, and I made a minor edit to one 1D test case changing u(x) ~ xcos(...) to u(x) - u(0) ~ xcos(...), to test these changes and mimic someone trying to train a network that they new passed through (0,0) and so would represent by g(x) = u(x) - u(0).

This pull request changes the compat entry for the QuasiMonteCarlo package from 0.2 to 0.2, 0.3 for package docs. This keeps the compat entries for earlier versions.

Note: I have not tested your package with this new compat entry. It is your responsibility to make sure that your package tests pass before you merge this pull request.

This pull request changes the compat entry for the QuasiMonteCarlo package from 0.2.1 to 0.2.1, 0.3. This keeps the compat entries for earlier versions.

Note: I have not tested your package with this new compat entry. It is your responsibility to make sure that your package tests pass before you merge this pull request.

This pull request changes the compat entry for the DomainSets package from 0.5 to 0.5, 0.6 for package docs. This keeps the compat entries for earlier versions.

Note: I have not tested your package with this new compat entry. It is your responsibility to make sure that your package tests pass before you merge this pull request.

This pull request changes the compat entry for the DomainSets package from 0.5 to 0.5, 0.6. This keeps the compat entries for earlier versions.

Note: I have not tested your package with this new compat entry. It is your responsibility to make sure that your package tests pass before you merge this pull request.

Hi, I was about to solve a multi-dimensional Stochastic PDE like the one that used to be in the documentation. No wonder why you removed it, the old code does not work. Therefore, what modifications do I have to do to make it work now? I dig out the code from an old commit

using NeuralPDE
using Flux
using DifferentialEquations
using LinearAlgebra
d = 100 # number of dimensions
X0 = fill(0.0f0, d) # initial value of stochastic control process
tspan = (0.0f0, 1.0f0)
λ = 1.0f0

g(X) = log(0.5f0 + 0.5f0 * sum(X.^2))
f(X,u,σᵀ∇u,p,t) = -λ * sum(σᵀ∇u.^2)
μ_f(X,p,t) = zero(X)  # Vector d x 1 λ
σ_f(X,p,t) = Diagonal(sqrt(2.0f0) * ones(Float32, d)) # Matrix d x d
prob = TerminalPDEProblem(g, f, μ_f, σ_f, X0, tspan)
hls = 10 + d # hidden layer size
opt = Flux.ADAM(0.01)  # optimizer
# sub-neural network approximating solutions at the desired point
u0 = Flux.Chain(Dense(d, hls, relu),
                Dense(hls, hls, relu),
                Dense(hls, 1))
# sub-neural network approximating the spatial gradients at time point
σᵀ∇u = Flux.Chain(Dense(d + 1, hls, relu),
                  Dense(hls, hls, relu),
                  Dense(hls, hls, relu),
                  Dense(hls, d))
pdealg = NNPDENS(u0, σᵀ∇u, opt=opt)
@time ans = solve(prob, pdealg, verbose=true, maxiters=100, trajectories=100,
                            alg=EM(), dt=1.2, pabstol=1f-2)

I really appreciate if you could help me. Thanks.