# Add coordinate transformation before first layer in NeuralPDE.jl

**URL:** <https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769>\
**Category:** Machine Learning\
**Tags:** pde, sciml, neural-network, lux\
**Created:** [August 29, 2024, 1:20pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769 "2024-08-29T13:20:05Z")\
**Posts on this page:** 6\
**Page:** 1

<div class="post-metadata">

**Author:** ![nico](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nico/32/211667_2.png) [@nico](https://discourse.julialang.org/u/nico)\
**Post date:** [August 29, 2024, 1:20pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769/1 "2024-08-29T13:20:05Z")

</div>

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(B_x);sin(B_x)], 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](https://global.discourse-cdn.com/julialang/original/3X/9/4/94c09360778426913731265327a88dcdf83e399b.png)  
Thank you for any help!

---

<div class="post-metadata">

**Author:** ![ChrisRackauckas](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/chrisrackauckas/32/77_2.png) [@ChrisRackauckas](https://discourse.julialang.org/u/ChrisRackauckas)\
**Post date:** [August 29, 2024, 2:54pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769/2 "2024-08-29T14:54:47Z")

</div>

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

---

<div class="post-metadata">

**Author:** ![nico](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nico/32/211667_2.png) [@nico](https://discourse.julialang.org/u/nico)\
**Post date:** [August 29, 2024, 4:44pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769/3 "2024-08-29T16:44:04Z")

</div>

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.

---

<div class="post-metadata">

**Author:** ![bertschi](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/bertschi/32/33462_2.png) [@bertschi](https://discourse.julialang.org/u/bertschi)\
**Post date:** [August 29, 2024, 7:39pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769/4 "2024-08-29T19:39:21Z")

</div>

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

```

---

<div class="post-metadata">

**Author:** ![nico](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nico/32/211667_2.png) [@nico](https://discourse.julialang.org/u/nico)\
**Post date:** [August 29, 2024, 8:15pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769/5 "2024-08-29T20:15:45Z")

</div>

thanks!

---

<div class="post-metadata">

**Author:** ![nico](https://sea2.discourse-cdn.com/julialang/user_avatar/discourse.julialang.org/nico/32/211667_2.png) [@nico](https://discourse.julialang.org/u/nico)\
**Post date:** [August 29, 2024, 8:16pm UTC](https://discourse.julialang.org/t/add-coordinate-transformation-before-first-layer-in-neuralpde-jl/118769/6 "2024-08-29T20:16:54Z")

</div>

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

```julia
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") 

```