Add coordinate transformation before first layer in NeuralPDE.jl

Specific Domains Machine Learning
, , ,
1

Hi, I woud like to add a transformation to the coordinates x and y before they are passed to the first layer of a neural net defined with Lux using NerualPDE.jl. I am implementing physics informed neural networks.
The idea is to implement Fourier features as described in:
Wang, S., Sankaran, S., Wang, H., & Perdikaris, P. (2023). An expert’s guide to training physics-informed neural networks. arXiv preprint arXiv:2308.08468 .
That is, gamma(x) = [cos(Bx);sin(Bx)], where B is a matrix with randomly and fixed coefficients. gamma(x) is then passed to the neural network, which then has as input dimentions two times the number of rows of B.
image
Thank you for any help!

2

You just add that function as the first one in the chain. Should be one line of code. What did you try?

3

I was not sure how to formulate this. I mean how to “interact” with the Chain of Lux. Didnt expect it to be that flexible.
So you say that a function from x,y and that returns a n-vector can work as first element of the chain?
Thank you for your reply Chris by the way.

4

Yes, any which can handle single or batched vectors, i.e., \mathbb{R}^{n (\times b)} \to \mathbb{R}^{m (\times b)}, can be used as a layer.
Here is an example:

julia> using Lux; import Random

julia> n, m = 4, 3;

julia> rng = Random.default_rng(); Random.seed!(rng, 0);

julia> blayer = let B = randn(Float32, m, n); x -> vcat(cos.(B * x), sin.(B * x)) end
#5 (generic function with 1 method)

julia> model = Chain(blayer, Dense(2*m => 5));

julia> ps, st = Lux.setup(rng, model);

julia> x = rand(rng, Float32, n, 8);

julia> size(x)
(4, 8)

julia> y, st = Lux.apply(model, x, ps, st);

julia> size(y)
(5, 8)
5

thanks!

6

Solved!
Here is the code I tested it with. It is based on the Poisson tutorial from NeuralPDE.jl (It works very well BTW!! I encourage you to try it.):

using NeuralPDE, Lux, ModelingToolkit, Optimization, OptimizationOptimJL
using Random
using LineSearches
import ModelingToolkit: Interval, infimum, supremum
using Plots

@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)

# Boundary conditions
bcs = [u(0, y) ~ 0.0, u(1, y) ~ 0,
    u(x, 0) ~ 0.0, u(x, 1) ~ 0]
# Space and time domains
domains = [x ∈ Interval(0.0, 1.0),
    y ∈ Interval(0.0, 1.0)]

dim = 2 # number of dimensions
# Fourier features
nFourier = 10
B = π*rand(nFourier,dim)
gammaFunc(x) = [cos.(B*x); sin.(B*x)]
# Neural network
nn = 10
chain = Lux.Chain(gammaFunc, Dense(2*nFourier, nn, Lux.σ), Dense(nn, nn, Lux.σ), Dense(nn, 1))

#strategy = QuadratureTraining(; batch = 200, abstol = 1e-6, reltol = 1e-6)
dx = 0.05
strategy = GridTraining(dx)
discretization = PhysicsInformedNN(chain, strategy)

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

lossHist = []
iter = 0
callback = function (p, l)
    global iter += 1
    println("Iter: $iter, Loss: $l")
    push!(lossHist, l)
    return false
end

# Train
opt = OptimizationOptimJL.LBFGS(linesearch = BackTracking())
res = Optimization.solve(prob, opt; callback = callback, maxiters = 1000)

# Extract
phi = discretization.phi

# Plot
dx = 0.05
xs, ys = [infimum(d.domain):(dx / 10):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)

p1 = plot(xs, ys, u_real, linetype = :contourf, title = "analytic");
p2 = plot(xs, ys, u_predict, linetype = :contourf, title = "predict");
p3 = plot(xs, ys, diff_u, linetype = :contourf, title = "error");
plot(p1, p2, p3)
savefig("poisson_PINN_u.png")

plot(lossHist, yscale=:log10)
savefig("poisson_PINN_loss.png")
1 Like