PDE solving example
In this section we solve the following PDE problem with a spline grid:
\[\begin{align*} \begin{cases} u : \Omega \rightarrow \mathbb{R} \\ \Delta u = u^3 &\text{ for }& \mathbf{x} \in \Omega \\ u = g(\mathbf{x}) \; &\text{ for }& \mathbf{x} \in \partial\Omega \end{cases} \end{align*},\]
where the domain is given by the open unit cube: $\Omega = (0,1)^3$. We assume that $g(x,y,z) = 0$ for $z \in (0,1)$ and $g(x,y,z) = f(x,y)$ for $z = 0, 1$.
We solve this problem by sampling the domain and enforcing the PDE on the interior points and the boundary condition on the boundary points.
The residual kernel
We first define the kernel which calculates the residual of a given approximation to the solution of the problem above. Note that this kernel is agnostic of the fact that the solution will come from a spline grid.
using KernelAbstractions
@kernel function pde_residual_kernel(
residual,
@Const(f),
@Const(u),
@Const(∂₁²u),
@Const(∂₂²u),
@Const(∂₃²u)
)
I = @index(Global, Cartesian)
is_boundary = false
for (i, i_max) in zip(Tuple(I)[1:(end - 1)], size(residual)[1:(end - 1)])
if (i == 1) || (i == i_max)
is_boundary = true
break
end
end
residual[I] = if is_boundary
if I[3] == 1 || I[3] == size(residual)[3]
u[I] - f[I[1], I[2]]
else
u[I]
end
else
∂₁²u[I] + ∂₂²u[I] + ∂₃²u[I] - u[I]^3
end
endpde_residual_kernel (generic function with 4 methods)The residual as a function of the spline grid control points
Whe now define a function which generates the input of the above residual kernel from a spline grid and computes the residual in place. It assumes a parameter object with the spline grid, f (an array that specifies the boundary condition for the bottom and top of the domain, the boundary condition is 0 elsewhere) and caches for the spline grid evaluation.
function pde_residual!(residual, control_points_flat, p)::Nothing
(; spline_grid, f, u, ∂₁²u, ∂₂²u, ∂₃²u) = p
control_points = reshape(
control_points_flat,
size(spline_grid.control_points)
)
evaluate!(spline_grid; control_points, eval = u)
evaluate!(spline_grid; control_points, eval = ∂₁²u, derivative_order = (2, 0, 0))
evaluate!(spline_grid; control_points, eval = ∂₂²u, derivative_order = (0, 2, 0))
evaluate!(spline_grid; control_points, eval = ∂₃²u, derivative_order = (0, 0, 2))
backend = get_backend(u)
pde_residual_kernel(backend)(
reshape(residual, size(u)),
f,
u,
∂₁²u,
∂₂²u,
∂₃²u,
ndrange = size(u)
)
synchronize(backend)
return nothing
endpde_residual! (generic function with 1 method)Defining the spline grid
We would like to define a residual function $\mathbb{R}^N \rightarrow \mathbb{R}^N$, i.e. the number of sample points must equal the number of control points. For each dimension the maximum derivative order is 2, to evaluate the result of the Laplacian operator on the spline.
using SplineGrids
n_control_points = (25, 25, 25)
degree = (2, 2, 2)
n_sample_points = n_control_points
dim_out = 1
spline_dimensions = SplineDimension.(
n_control_points, degree, n_sample_points; max_derivative_order = 2)
spline_grid = SplineGrid(spline_dimensions, dim_out)
spline_gridSplineGrid volume with outputs in ℝ¹ (Float32)
----------------------------------------------
* Properties per dimension:
┌─────────────────┬────────┬───────────────────┬─────────────────┐
│ input dimension │ degree │ # basis functions │ # sample points │
├─────────────────┼────────┼───────────────────┼─────────────────┤
│ 1 │ 2 │ 25 │ 25 │
│ 2 │ 2 │ 25 │ 25 │
│ 3 │ 2 │ 25 │ 25 │
└─────────────────┴────────┴───────────────────┴─────────────────┘
* Control points:
DefaultControlPoints for grid of size (25, 25, 25) in ℝ¹ (Float32).
Defining the boundary condition
We define a boundary condition inspired by the Julia logo.
using GLMakie
function hill(x, y, x0, y0, β)
r = sqrt((x - x0)^2 + (y - y0)^2)
return exp(-(r / β)^2)
end
function hills(x, y)
R = 0.25
out = 0.0
for θ in range(0, 2π, length = 4)[1:(end - 1)]
x0 = 0.5 + R * cos(θ)
y0 = 0.5 + R * sin(θ)
out += hill(x, y, x0, y0, R / 2)
end
return 100 * out
end
f = [hills.(x, y)
for
x in spline_dimensions[1].sample_points,
y in spline_dimensions[2].sample_points]
heatmap(f)
Defining the problem
We will solve the problem with a Jacobian free Newton-Krylov method. To do this, we need to provide a Jacobian-vector product (JVP) function.
using NonlinearSolve
using Enzyme
p = (;
spline_grid,
f,
u = similar(spline_grid.eval),
∂₁²u = similar(spline_grid.eval),
∂₂²u = similar(spline_grid.eval),
∂₃²u = similar(spline_grid.eval)
)
meta_p = (;
residual = similar(spline_grid.eval),
p_duplicated = DuplicatedNoNeed(p, make_zero(p))
)
function pde_residual_jvp!(Jv, v, control_points_flat, meta_p)::Nothing
for val in values(meta_p.p_duplicated.dval)
make_zero!(val)
end
autodiff(
Forward,
pde_residual!,
DuplicatedNoNeed(vec(meta_p.residual), Jv),
Duplicated(control_points_flat, v),
meta_p.p_duplicated
)
return nothing
end
nonlinear_function = NonlinearFunction(
(residual, control_points_flat, p_) -> pde_residual!(residual, control_points_flat, p);
jvp = pde_residual_jvp!
)
# Initial guess
x0 = zeros(Float32, length(spline_grid.control_points))
problem = NonlinearProblem(
nonlinear_function,
x0,
meta_p
)NonlinearProblem with uType Vector{Float32}. In-place: true
u0: 15625-element Vector{Float32}:
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
⋮
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0
0.0Solving the problem
sol = solve(problem, NewtonRaphson(linsolve = KrylovJL_GMRES()))retcode: MaxIters
u: 15625-element Vector{Float32}:
7.203033f-5
0.00012896715
0.00088540546
0.004786132
0.020288704
0.06699376
0.17124008
0.33741593
0.512363
0.6007665
⋮
0.0013418678
0.00020978265
3.4804158f-5
1.1672109f-5
6.2266195f-6
2.9949101f-6
1.1428368f-6
3.4744426f-7
2.0612349f-7Viewing the solution
spline_grid.control_points .= reshape(sol.u, size(spline_grid.control_points))
evaluate!(spline_grid)
fig = Figure()
ax, plt = volume(fig[1, 1], log.(spline_grid.eval[:, :, :, 1] .+ 1))
Colorbar(fig[1, 2], plt; label = L"\log(u + 1)")
fig