ZDM / GB heterogeneous input

[killah-t-cell]

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 #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).

[killah-t-cell]

I am now running into problems with the way I define the cord (we changed it to call it input_cord). We define it like this input_cord = :(adapt(DiffEqBase.parameterless_type($θ), vcat($(interior_u_args...)))) and this is giving me a dimension error LoadError: DimensionMismatch("matrix A has dimensions (3,1), matrix B has dimensions (94,30)") . I tried to debug line by line but I don't understand why this is happening yet. Do you have any idea why? @KirillZubov

Will take a break and look at it later. Hopefully with fresh eyes.

[killah-t-cell]

@KirillZubov as requested, an MWE for input_cord

## MWE
#chain = [[FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]; [FastChain(FastDense(2, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1)) for i in 1:2]]
chain = FastChain(FastDense(2, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1))
ex2 = :(s(y, x))
depvar = ex2.args[1]
indvars = ex2.args[2:end]
num_depvar = dict_depvars[depvar]
interior_u_args = map(indvars) do arg
    if arg isa Symbol
        arg
    else
        :(fill($arg, (1, $:batch_size)))
    end
end
input_cord = :(adapt(DiffEqBase.parameterless_type($θ), vcat($(interior_u_args...))))
ex.args = if !(typeof(chain) <: AbstractVector)
    [:($(Expr(:$, :u))), input_cord, :($θ), :phi]
else
    [:($(Expr(:$, :u))), input_cord, Symbol(:($θ), num_depvar), Symbol(:phi, num_depvar)]
end

# 1 var
chain = FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1))
ex2 = :(r(1))
depvar = ex2.args[1]
indvars = ex2.args[2:end]
num_depvar = dict_depvars[depvar]
interior_u_args = map(indvars) do arg
    if arg isa Symbol
        arg
    else
        :(fill($arg, (1, $:batch_size)))
    end
end
input_cord = :(adapt(DiffEqBase.parameterless_type($θ), vcat($(interior_u_args...))))
ex.args = if !(typeof(chain) <: AbstractVector)
    [:($(Expr(:$, :u))), input_cord, :($θ), :phi]
else
    [:($(Expr(:$, :u))), input_cord, Symbol(:($θ), num_depvar), Symbol(:phi, num_depvar)]
end

# bc
chain = FastChain(FastDense(1, numhid, Flux.σ), FastDense(numhid, numhid, Flux.σ), FastDense(numhid, 1))
ex2 = :(p(1, x))
depvar = ex2.args[1]
indvars = ex2.args[2:end]
num_depvar = dict_depvars[depvar]
interior_u_args = map(indvars) do arg
    if arg isa Symbol
        arg
    else
        :(fill($arg, (1, $:batch_size)))
    end
end
input_cord = :(adapt(DiffEqBase.parameterless_type($θ), vcat($(interior_u_args...))))
ex.args = if !(typeof(chain) <: AbstractVector)
    [:($(Expr(:$, :u))), input_cord, :($θ), :phi]
else
    [:($(Expr(:$, :u))), input_cord, Symbol(:($θ), num_depvar), Symbol(:phi, num_depvar)]
end

[KirillZubov]

It does not need to adapt here. cord is already adapted in the computational part.

θ = gensym("θ")
interior_u_args = [:y, :x]
julia> input_cord = :(adapt((DiffEqBase.parameterless_type($θ)), vcat($(interior_u_args...))))
:(adapt(DiffEqBase.parameterless_type(var"##θ#347"), vcat(y, x)))

julia> input_cord =  :(vcat($(interior_u_args...)))
:(vcat(y, x))

also, I think vcat should be done not in arguments of variables. it confuses and makes the inner representation hard to read

I suggest this way:

ex.args = if !(typeof(chain) <: AbstractVector)
    [:($(Expr(:$, :u))), Symbol(:cord,num_depvar), :($θ), :phi]
else
    [:($(Expr(:$, :u))), Symbol(:cord,num_depvar), Symbol(:($θ), num_depvar), Symbol(:phi, num_depvar)]
end

so inner representation will look like this:

(Expr[:((cord, var"##θ#338", phi, derivative, u, p)->begin
          begin
              (var"##θ#3381", var"##θ#3382") = (var"##θ#338"[1:61], var"##θ#338"[62:122])
              (phi1, phi2,phi3, phi4) = (phi[1], phi[2], phi[3], phi[4])
              let (x, y) = (cord[[1], :], cord[[2], :])
                  cord1 = vcat(x)
                  cord2 = vcat(y)
                  cord3 = vcat(x, y)
                  cord4 = vcat(y, x)
                  expression_of_equation(phi1(cord1,...),   phi2(cord2,...), ... derivative(phi1, u, cord1,...))
              end
          end
      end, ...

it's already done something similar for Quadrature
check here

NeuralPDE.jl/src/pinns_pde_solve.jl

Line 455 in 5a881ff

vcat_expr = :(cord = vcat($(indvars...)))

and sym_prob output of this example

NeuralPDE.jl/test/NNPDE_tests.jl

Lines 223 to 257 in 5a881ff

	## Example 4, system of pde
	println("Example 4, system of pde")
	@parameters x, y
	@variables u1(..), u2(..)
	Dx = Differential(x)
	Dy = Differential(y)
	# System of pde
	eqs = [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
	Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0]
	# 3*u1(x,0) ~ 2*u2(x,0)]
	# Initial and boundary conditions
	bcs = [u1(x,0) ~ 2*x, u2(x,0) ~ 3*x]
	# Space and time domains
	domains = [x ∈ Interval(0.0,1.0), y ∈ Interval(0.0,1.0)]
	# Neural network
	chain1 = FastChain(FastDense(2,15,Flux.tanh),FastDense(15,1))
	chain2 = FastChain(FastDense(2,15,Flux.tanh),FastDense(15,1))
	quadrature_strategy = NeuralPDE.QuadratureTraining(quadrature_alg=CubatureJLh(),
	reltol=1e-3,abstol=1e-3,
	maxiters =50, batch=100)
	chain = [chain1,chain2]
	initθ = map(c -> Float64.(c), DiffEqFlux.initial_params.(chain))
	discretization = NeuralPDE.PhysicsInformedNN(chain,quadrature_strategy; init_params = initθ)
	@named pde_system = PDESystem(eqs,bcs,domains,[x,y],[u1,u2])
	prob = NeuralPDE.discretize(pde_system,discretization)
	sym_prob = NeuralPDE.symbolic_discretize(pde_system,discretization)

[KirillZubov]

in general, it is no longer necessary to store a cord in one multidimensional array, it is outdated.
Now I think better - one coordinate in one one-dimensional vector.
but it is already another issue

cord = rand(2,10)
(x, y) = (cord[[1], :],cord[[2], :])

julia> @benchmark (x, y) = (cord[[1], :],cord[[2], :])
BenchmarkTools.Trial: 10000 samples with 178 evaluations.
 Range (min … max):  591.865 ns … 165.837 μs  ┊ GC (min … max): 0.00% … 99.26%
 Time  (median):     626.309 ns               ┊ GC (median):    0.00%
 Time  (mean ± σ):   795.487 ns ±   3.504 μs  ┊ GC (mean ± σ):  9.79% ±  2.22%

cord = [rand(1,10),rand(1,10)]

@benchmark (x, y) = (cord[1],cord[2])
BenchmarkTools.Trial: 10000 samples with 973 evaluations.
Range (min … max):  72.491 ns …  29.747 μs  ┊ GC (min … max): 0.00% … 99.58%
Time  (median):     77.020 ns               ┊ GC (median):    0.00%
Time  (mean ± σ):   89.026 ns ± 297.478 ns  ┊ GC (mean ± σ):  3.33% ±  1.00%

[killah-t-cell]

Thanks! This explains a lot and cleans up a lot of the code. I have a question:

in build_symbolic_loss_function, in order to get the right cords to vcat for each equation in a system:

...
  cord1 = vcat(x)
  cord2 = vcat(y)
  cord3 = vcat(x, y)
  cord4 = vcat(y, x)
...

I need to know which dependent variables are being used in each equation. How do you recommend I do that inside of build_symbolic_loss_function?

I tried using this_eq_depvars = filter(depvar -> dict_depvar_input[depvar] ⊆ bc_indvars, depvars) but that does not work so well.

[KirillZubov]

draft

function pair(eq, depvars,dict_depvars, dict_depvar_input)
    expr = toexpr(eq)
    pair_ = map(depvars) do depvar
        if !isempty(find_thing_in_expr(expr,  depvar))
            dict_depvars[depvar] => dict_depvar_input[depvar]
        end
    end
    Dict(filter(p -> p !==nothing , pair_))
end

julia> eq_pair = pair(eq, depvars,dict_depvars, dict_depvar_input)
Dict{Int64, Vector{Symbol}} with 4 entries:
  4 => [:y, :x]
  2 => [:y]
  3 => [:x, :y]
  1 => [:x]

 julia> bcs_pair = map(bc ->pair(bc, depvars,dict_depvars, dict_depvar_input),bcs)
 6-element Vector{Dict{Int64, Vector{Symbol}}}:
  Dict(1 => [:x])
  Dict(2 => [:y])
  Dict(3 => [:x, :y])
  Dict(3 => [:x, :y])
  Dict(4 => [:y, :x])
  Dict(4 => [:y, :x])

some_eq = Dx(p(x, y)) + Dy(s(y, x)) - Dy(s(0, x)) ~ p(0,0)

julia>  some_eq_pair =pair(some_eq, depvars,dict_depvars, dict_depvar_input)
Dict{Int64, Vector{Symbol}} with 2 entries:
  4 => [:y, :x]
  1 => [:x]


 

[KirillZubov]

        eq_pair_expr  =Expr[]
        for i in keys(some_eq_pair)
            push!(eq_pair_expr, :( $(Symbol(:cord, :($i))) = vcat($(some_eq_pair[i]...))))
        end
        julia> eq_pair_expr
        2-element Vector{Expr}:
        :(cord4 = vcat(y, x))
        :(cord1 = vcat(x))

[killah-t-cell]

Thanks!

I now have restructured build_symbolic_loss so it outputs this for a heterogeneous system

eqs = u1(x, y) + Dx(u2(x)) ~ 0

bcs = [u1(x,0) ~ cos(x), u2(x) ~ 0]

(Expr[:((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[[1], :], cord[[2], :])
                  begin
                      cord2 = vcat(x)
                      cord1 = vcat(x, y)
                  end
                  (+).(derivative(phi2, u, cord2, [[6.0554544523933395e-6]], 1, var"##θ#3542"), u(cord1, var"##θ#3541", phi1)) .- 0
              end
          end
      end)], Expr[:((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[[1], :], fill(0, size(cord[[1], :])))
                  begin
                      cord1 = vcat(x, y)
                  end
                  u(cord1, var"##θ#3541", phi1) .- cos.(x)
              end
          end
      end), :((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x,) = (cord[[1], :],)
                  begin
                      cord2 = vcat(x)
                  end
                  u(cord2, var"##θ#3542", phi2) .- 0
              end
          end
      end)])

And for a normal system

# System of pde
eqs = [Dx(u1(x,y)) + 4*Dy(u2(x,y)) ~ 0,
      Dx(u2(x,y)) + 9*Dy(u1(x,y)) ~ 0]

# Initial and boundary conditions
bcs = [u1(x,0) ~ 2*x, u2(x,0) ~ 3*x]

(Expr[:((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[[1], :], cord[[2], :])
                  begin
                      cord2 = vcat(x, y)
                      cord1 = vcat(x, y)
                  end
                  (+).(derivative(phi1, u, cord1, [[6.0554544523933395e-6, 0.0]], 1, var"##θ#3541"), (*).(4, derivative(phi2, u, cord2, [[0.0, 6.0554544523933395e-6]], 1, var"##θ#3542"))) .- 0
              end
          end
      end), :((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[[1], :], cord[[2], :])
                  begin
                      cord2 = vcat(x, y)
                      cord1 = vcat(x, y)
                  end
                  (+).(derivative(phi2, u, cord2, [[6.0554544523933395e-6, 0.0]], 1, var"##θ#3542"), (*).(9, derivative(phi1, u, cord1, [[0.0, 6.0554544523933395e-6]], 1, var"##θ#3541"))) .- 0
              end
          end
      end)], Expr[:((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[[1], :], fill(0, size(cord[[1], :])))
                  begin
                      cord1 = vcat(x, y)
                  end
                  u(cord1, var"##θ#3541", phi1) .- (*).(2, x)
              end
          end
      end), :((cord, var"##θ#354", phi, derivative, u, p)->begin
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          #= /Users/gabrielbirnbaum/.julia/dev/NeuralPDE/src/pinns_pde_solve.jl:514 =#
          begin
              (var"##θ#3541", var"##θ#3542") = (var"##θ#354"[1:61], var"##θ#354"[62:122])
              (phi1, phi2) = (phi[1], phi[2])
              let (x, y) = (cord[[1], :], fill(0, size(cord[[1], :])))
                  begin
                      cord2 = vcat(x, y)
                  end
                  u(cord2, var"##θ#3542", phi2) .- (*).(3, x)
              end
          end
      end)])

The normal systems and tests are working now which is good. But the heterogeneous system is still throwing out an LoadError: DimensionMismatch("A has dimensions (15,2) but B has dimensions (1,2)") error.

I think this might have something to do with discretize or with build_loss_function but right now I don't yet know.

EDIT: I had a typo. Both systems are working now! 😄

[killah-t-cell]

Ok so here is the current to do list before opening the PR:

• Fix heterogeneous systems when using strategies that are not QuadratureTraning
• Clean up pinns_pde_solve.jl (remove @shows, unnecessary spacing changes, mwe)
• Write up tests

I plan to open the PR tomorrow.

[killah-t-cell]

@KirillZubov this is ready for a code review. Are there any good examples of heterogeneous systems we could use for the tests and the documentation?

[killah-t-cell]

An update:

I cleaned up the white spaces. There are only 2 problems remaining now:

Problem 1: I am getting a BoundsError: attempt to access 1×1 Matrix{Float32} at index [[2], 1:1] in the Approximation of function 2D test. The symbolic discretization looks like this (see below). I am not yet sure how to fix this properly (all other things are working so I am not sure why this test in particular is broken).

(Expr[:((cord, var"##θ#332", phi, derivative, integral, u, p)->begin
          begin
              let (x, y) = (cord[[1], :], cord[[2], :])
                  begin
                      cord1 = vcat(x, y)
                  end
                  u(cord1, var"##θ#332", phi) .- (*).(-1, (cos).(x), (cos).(y), (exp).((+).((*).(-1, (^).((+).(-3.141592653589793, x), 2)), (*).(-1, (^).((+).(-3.141592653589793, y), 2)))))
              end
          end
      end)], Expr[:((cord, var"##θ#332", phi, derivative, integral, u, p)->begin
          begin
              let (x, y) = (cord[[1], :], cord[[2], :])
                  begin
                      cord1 = vcat(x, y)
                  end
                  u(cord1, var"##θ#332", phi) .- u(cord1, var"##θ#332", phi)
              end
          end
      end)])

Problem 2: The way we discretize a Integro-DE is very different. From what I gathered from looking at master, this is supposed to look a bit like this (see below). In this world I am not sure how to adapt the way we load cords to this yet. _transform_expression is calling build_symbolic_loss_function with eqs=nothing which breaks the way I find out which dependent variables are being used in each equation.

        # if eqs = nothing, this_eq_pair can't get the information it needs.
        this_eq_pair = pair(eqs, depvars, dict_depvars, dict_depvar_input)
        this_eq_indvars = sort(collect(Set(vcat(values(this_eq_pair)...))))
        ...
        eq_pair_expr = Expr[]
        for i in keys(this_eq_pair)
            push!(eq_pair_expr, :( $(Symbol(:cord, :($i))) = vcat($(this_eq_pair[i]...))))
        end
        vcat_expr = Expr(:block, :($(eq_pair_expr...)))
        vcat_expr_loss_functions = Expr(:block, vcat_expr, loss_function)
        ...

(Expr[:((cord, var"##θ#331", phi, derivative, integral, u, p)->begin
          begin
              let (x,) = (cord[[1], :],)
                  integral(u, cord, phi, [1], RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#331"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0x1a7e437d, 0x633f4d2b, 0x0c12a9a0, 0xdba3679e, 0x4f5c1eae)}(quote
    begin
        let (x,) = (cord[[1], :],)
            u(cord, var"##θ#331", phi)
        end
    end
end), Any[0], Any[RuntimeGeneratedFunctions.RuntimeGeneratedFunction{(:cord, Symbol("##θ#331"), :phi, :derivative, :integral, :u, :p), NeuralPDE.var"#_RGF_ModTag", NeuralPDE.var"#_RGF_ModTag", (0xece15f21, 0xdd268b8a, 0x3a6073e5, 0x93a3d03a, 0x3d656f4c)}(quote
    begin
        let (x,) = (cord[[1], :],)
            x
        end
    end
end)], var"##θ#331") .- (*).(1//3, (^).(x, 3))
              end
          end
      end)], Expr[:((cord, var"##θ#331", phi, derivative, integral, u, p)->begin
          begin
              let (x,) = (cord[[1], :],)
                  u(cord, var"##θ#331", phi) .- 0.0
              end
          end
      end)])

I will be back at it tomorrow to try to figure this out.

[killah-t-cell]

Ok I now fixed Problem 2 and IDEs work with the heterogeneous system. I am still getting the BoundsError: attempt to access 1×1 Matrix{Float32} at index [[2], 1:1] in the Approximation of function 2D test. The symbolic_loss_function output looks right to me (see my previous comment) so I am not sure where to look. I wonder if @KirillZubov will find something during code review or if he has any tips.

[killah-t-cell]

I know I said this before but – this is almost done now.

I am training the Kermack–McKendrick model right now. This system is big and takes a while to train. It would be maybe a good idea to use another system to do a convergence test on just to improve developer experience. I may also just be training it wrong right now (it currently stops converging when the loss hits about 80). I did not spend too much time playing with this yet, so I will try to focus on it now as it seems like this is the current bottleneck.

This PR is now 2 weeks old, so I am trying to wrap this up asap. I started writing some documentation for the feature. We can probably use the Kermack–McKendrick model there at least (as it is a pretty cool example).

I would love to get some feedback on some questions I have regarding order:

1. I define this_eq_indvars = unique(vcat(values(this_eq_pair)...)). I wonder if this is ok or if I should be more order-aware.
2. Heterogeneous system are currently order-dependent. That means that I need to declare the dependent variables in PDESystem in the same order that I declared the chain. Is that ok? Is there a quick fix? (which we can do in another PR maybe for the sake of de-risking)

[ChrisRackauckas]

Yes, let's wrap it off. It seems like this is getting side tracked. Get to a version you think is ready to merge, and then let's merge and follow up with the other issues.

That means that I need to declare the dependent variables in PDESystem in the same order that I declared the chain. Is that ok?

I think that is expected. If it's u(x,y), it better always be u(x,y).

[killah-t-cell]

Alright @ChrisRackauckas I pushed a version I think is ready to merge. It doesn't have a convergence test but I am still gonna try to get either this SIR model or another heterogeneous system to work. I will make another PR for that though.

[ChrisRackauckas]

lgtm

[ChrisRackauckas]

Note this will need a breaking release.

[ChrisRackauckas]

NNPDE should be split. It doesn't make sense to have one test set be over an hour.