Skip to content

Commit

Permalink
Add collocation for time-dependent PDEs (#43)
Browse files Browse the repository at this point in the history
* add collocation method for time-dependent PDEs

* format

* relax tolerance for CI

* increase coverage

* format
  • Loading branch information
JoshuaLampert committed May 3, 2024
1 parent 1610525 commit f188609
Show file tree
Hide file tree
Showing 17 changed files with 575 additions and 141 deletions.
4 changes: 4 additions & 0 deletions Project.toml
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,8 @@ ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
ReadVTK = "dc215faf-f008-4882-a9f7-a79a826fadc3"
RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
SciMLBase = "0bca4576-84f4-4d90-8ffe-ffa030f20462"
SimpleUnPack = "ce78b400-467f-4804-87d8-8f486da07d0a"
SpecialFunctions = "276daf66-3868-5448-9aa4-cd146d93841b"
StaticArrays = "90137ffa-7385-5640-81b9-e52037218182"
TypedPolynomials = "afbbf031-7a57-5f58-a1b9-b774a0fad08d"
Expand All @@ -18,6 +20,8 @@ ForwardDiff = "0.10.36"
LinearAlgebra = "1"
ReadVTK = "0.2"
RecipesBase = "1.1"
SciMLBase = "1.93, 2"
SimpleUnPack = "1.1"
SpecialFunctions = "2"
StaticArrays = "1"
TypedPolynomials = "0.4.1"
Expand Down
3 changes: 2 additions & 1 deletion docs/src/pde.md
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

TODO:

* Explain basic setup
* Explain basic setup and basics of collocation for stationary and time-dependent PDEs
* Define custom differential operators and PDEs and solve them
* Stationary and time-dependent PDEs
* AD vs analytic derivatives
28 changes: 28 additions & 0 deletions examples/PDEs/heat_2d_basic.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,28 @@
using KernelInterpolation
using OrdinaryDiffEq
using Plots

# right-hand-side of Heat equation
f(t, x, equations) = 1.0
pde = HeatEquation(1.0, f)

# initial condition
u(t, x, equations) = 0.0

n = 10
nodeset_inner = homogeneous_hypercube(n, (0.1, 0.1), (0.9, 0.9); dim = 2)
n_boundary = 3
nodeset_boundary = homogeneous_hypercube_boundary(n_boundary; dim = 2)
# Dirichlet boundary condition
g(t, x) = 0.0

kernel = WendlandKernel{2}(3, shape_parameter = 0.3)
sd = Semidiscretization(pde, nodeset_inner, g, nodeset_boundary, u, kernel)
tspan = (0.0, 1.0)
ode = semidiscretize(sd, tspan)
sol = solve(ode, Rosenbrock23(), saveat = 0.01)
titp = TemporalInterpolation(sol)

many_nodes = homogeneous_hypercube(20; dim = 2)

plot(many_nodes, titp(last(tspan)))
32 changes: 32 additions & 0 deletions examples/PDEs/heat_2d_manufactured.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,32 @@
using KernelInterpolation
using OrdinaryDiffEq
using Plots

# right-hand-side of heat equation
function f(t, x, equations)
exp(t) * (x[1] * (x[1] - 1) * x[2] * (x[2] - 1) -
2 * equations.diffusivity * (x[1] * (x[1] - 1) + x[2] * (x[2] - 1)))
end
pde = HeatEquation(2.0, f)

# analytical solution
u(t, x, equations) = exp(t) * x[1] * (x[1] - 1) * x[2] * (x[2] - 1)

n = 10
nodeset_inner = homogeneous_hypercube(n, (0.1, 0.1), (0.9, 0.9); dim = 2)
n_boundary = 3
nodeset_boundary = homogeneous_hypercube_boundary(n_boundary; dim = 2)
# Dirichlet boundary condition
g(t, x) = 0.0

kernel = WendlandKernel{2}(3, shape_parameter = 0.3)
sd = Semidiscretization(pde, nodeset_inner, g, nodeset_boundary, u, kernel)
tspan = (0.0, 1.0)
ode = semidiscretize(sd, tspan)
sol = solve(ode, Rosenbrock23(), saveat = 0.01)
titp = TemporalInterpolation(sol)

many_nodes = homogeneous_hypercube(20; dim = 2)

plot(many_nodes, titp(last(tspan)))
plot!(many_nodes, x -> u(last(tspan), x, pde))
Original file line number Diff line number Diff line change
Expand Up @@ -2,25 +2,24 @@ using KernelInterpolation
using Plots

# right-hand-side of Poisson equation
f(x) = 5 / 4 * pi^2 * sin(pi * x[1]) * cos(pi * x[2] / 2)
f(x, equations) = 5 / 4 * pi^2 * sin(pi * x[1]) * cos(pi * x[2] / 2)
pde = PoissonEquation(f)

# analytical solution of equation
u(x) = sin(pi * x[1]) * cos(pi * x[2] / 2)
u(x, equations) = sin(pi * x[1]) * cos(pi * x[2] / 2)

n = 10
nodeset_inner = homogeneous_hypercube(n, (0.1, 0.1), (0.9, 0.9); dim = 2)
n_boundary = 3
nodeset_boundary = homogeneous_hypercube_boundary(n_boundary; dim = 2)
values = f.(nodeset_inner)
# Dirichlet boundary condition (here taken from analytical solution)
g(x) = u(x)
values_boundary = g.(nodeset_boundary)
g(x) = u(x, pde)

kernel = WendlandKernel{2}(3, shape_parameter = 0.3)
itp = solve(pde, nodeset_inner, nodeset_boundary, values_boundary, kernel)
sd = SpatialDiscretization(pde, nodeset_inner, g, nodeset_boundary, kernel)
itp = solve_stationary(sd)

many_nodes = homogeneous_hypercube(20; dim = 2)

plot(many_nodes, itp)
plot!(many_nodes, u)
plot!(many_nodes, x -> u(x, pde))
14 changes: 10 additions & 4 deletions src/KernelInterpolation.jl
Original file line number Diff line number Diff line change
@@ -1,9 +1,11 @@
module KernelInterpolation

using ForwardDiff: ForwardDiff
using LinearAlgebra: norm, Symmetric, tr
using LinearAlgebra: Symmetric, norm, tr, muladd
using ReadVTK: VTKFile, get_points, get_point_data, get_data
using RecipesBase: RecipesBase, @recipe, @series
using SciMLBase: ODEFunction, ODEProblem, ODESolution
using SimpleUnPack: @unpack
using SpecialFunctions: besselk, loggamma
using StaticArrays: StaticArrays, MVector
using TypedPolynomials: Variable, monomials, degree
Expand All @@ -12,8 +14,10 @@ using WriteVTK: vtk_grid, MeshCell, VTKCellTypes
include("kernels/kernels.jl")
include("nodes.jl")
include("differential_operators.jl")
include("pdes.jl")
include("equations.jl")
include("kernel_matrices.jl")
include("interpolation.jl")
include("discretization.jl")
include("visualization.jl")
include("io.jl")
include("util.jl")
Expand All @@ -26,13 +30,15 @@ export GaussKernel, MultiquadricKernel, InverseMultiquadricKernel,
TransformationKernel, ProductKernel, SumKernel
export phi, Phi, order
export Laplacian
export PoissonEquation
export PoissonEquation, HeatEquation
export SpatialDiscretization, Semidiscretization, semidiscretize
export NodeSet, separation_distance, dim, eachdim, values_along_dim, distance_matrix,
random_hypercube, random_hypercube_boundary, homogeneous_hypercube,
homogeneous_hypercube_boundary, random_hypersphere, random_hypersphere_boundary
export interpolation_kernel, nodeset, coefficients, kernel_coefficients,
polynomial_coefficients, polynomial_basis, polyvars, system_matrix,
interpolate, solve, kernel_inner_product, kernel_norm
interpolate, solve_stationary, kernel_inner_product, kernel_norm,
TemporalInterpolation
export vtk_save, vtk_read
export examples_dir, get_examples, default_example, include_example

Expand Down
149 changes: 149 additions & 0 deletions src/discretization.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,149 @@
"""
SpatialDiscretization(equations, nodeset_inner, boundary_condition, nodeset_boundary)
Spatial discretization of a partial differential equation with Dirichlet boundary conditions.
The `nodeset_inner` are the nodes in the domain and `nodeset_boundary` are the nodes on the boundary. The `boundary_condition`
is a function describing the Dirichlet boundary conditions.
See also [`Semidiscretization`](@ref), [`solve_stationary`](@ref).
"""
struct SpatialDiscretization{Dim, RealT, Equations, BoundaryCondition,
Kernel <: AbstractKernel{Dim}}
equations::Equations
nodeset_inner::NodeSet{Dim, RealT}
boundary_condition::BoundaryCondition
nodeset_boundary::NodeSet{Dim, RealT}
kernel::Kernel

function SpatialDiscretization(equations, nodeset_inner, boundary_condition,
nodeset_boundary,
kernel = GaussKernel{dim(nodeset_inner)}())
new{dim(nodeset_inner), eltype(nodeset_inner), typeof(equations),
typeof(boundary_condition), typeof(kernel)}(equations,
nodeset_inner, boundary_condition,
nodeset_boundary, kernel)
end
end

function Base.show(io::IO, sd::SpatialDiscretization)
print(io,
"SpatialDiscretization with $(dim(sd)) dimensions, $(length(sd.nodeset_inner)) inner nodes, $(length(sd.nodeset_boundary)) boundary nodes, and kernel $(sd.kernel)")
end

dim(::SpatialDiscretization{Dim}) where {Dim} = Dim
Base.eltype(::SpatialDiscretization{Dim, RealT}) where {Dim, RealT} = RealT

"""
solve_stationary(spatial_discretization)
Solve a stationary partial differential equation discretized as `spatial_discretization` with Dirichlet boundary
conditions by non-symmetric collocation (Kansa method).
Returns an [`Interpolation`](@ref) object.
"""
function solve_stationary(spatial_discretization::SpatialDiscretization{Dim, RealT}) where {
Dim,
RealT
}
@unpack equations, nodeset_inner, boundary_condition, nodeset_boundary, kernel = spatial_discretization
nodeset = merge(nodeset_inner, nodeset_boundary)

pd_matrix = pde_matrix(equations, nodeset_inner, nodeset, kernel)
b_matrix = kernel_matrix(nodeset_boundary, nodeset, kernel)
system_matrix = [pd_matrix
b_matrix]
b = [rhs(nodeset_inner, equations); boundary_condition.(nodeset_boundary)]
c = system_matrix \ b

# Do not support additional polynomial basis for now
xx = polyvars(Dim)
ps = monomials(xx, 0:-1)
return Interpolation(kernel, nodeset, c, system_matrix, ps, xx)
end

"""
Semidiscretization(spatial_discretization, initial_condition)
Semidiscretization(equations, nodeset_inner, boundary_condition, nodeset_boundary, initial_condition, kernel = GaussKernel{dim(nodeset_inner)}())
Semidiscretization of a partial differential equation with Dirichlet boundary conditions and initial condition `initial_condition`. The `boundary_condition` function
can be time- and space-dependent. The `initial_condition` function is time- and space-dependent to be able to reuse it as analytical solution if available. If no
analytical solution is available, the time variable can be ignored in the `initial_condition` function.
See also [`SpatialDiscretization`](@ref), [`semidiscretize`](@ref).
"""
struct Semidiscretization{InitialCondition, Cache}
spatial_discretization::SpatialDiscretization
initial_condition::InitialCondition
cache::Cache
end

function Semidiscretization(spatial_discretization, initial_condition)
@unpack equations, nodeset_inner, boundary_condition, nodeset_boundary, kernel = spatial_discretization
nodeset = merge(nodeset_inner, nodeset_boundary)
k_matrix_inner = kernel_matrix(nodeset_inner, nodeset, kernel)
k_matrix_boundary = kernel_matrix(nodeset_boundary, nodeset, kernel)
# whole kernel matrix is not needed for rhs, but for initial condition
k_matrix = [k_matrix_inner
k_matrix_boundary]
pd_matrix = pde_matrix(equations, nodeset_inner, nodeset, kernel)
b_matrix = kernel_matrix(nodeset_boundary, nodeset, kernel)
pde_boundary_matrix = [pd_matrix
b_matrix]
m_matrix = [k_matrix_inner
zeros(eltype(k_matrix_inner), size(k_matrix_boundary)...)]
cache = (; kernel_matrix = k_matrix, mass_matrix = m_matrix,
pde_boundary_matrix = pde_boundary_matrix)
return Semidiscretization{typeof(initial_condition), typeof(cache)}(spatial_discretization,
initial_condition,
cache)
end

function Semidiscretization(equations, nodeset_inner, boundary_condition, nodeset_boundary,
initial_condition, kernel = GaussKernel{dim(nodeset_inner)}())
return Semidiscretization(SpatialDiscretization(equations, nodeset_inner,
boundary_condition, nodeset_boundary,
kernel), initial_condition)
end

function Base.show(io::IO, semi::Semidiscretization)
print(io,
"Semidiscretization with $(dim(semi)) dimensions, $(length(semi.spatial_discretization.nodeset_inner)) inner nodes, $(length(semi.spatial_discretization.nodeset_boundary)) boundary nodes, and kernel $(semi.spatial_discretization.kernel)")
end

dim(semi::Semidiscretization) = dim(semi.spatial_discretization)
Base.eltype(semi::Semidiscretization) = eltype(semi.spatial_discretization)

# right-hand side of the ODE
# M c' = -A c + b
# where M is the (singular) mass matrix
# M = (A_I; 0)∈R^{N x N}, A_I∈R^{N_I x N}
# A_{LB} = (A_L; A_B)∈R^{N x N}, A_L∈R^{N_I x N}, A_B∈R^{N_B x N}
# b = (g; g)∈R^N, f∈R^{N_I}, g∈R^{N_B}

# We can get u from c by
# u = A c
# A = (A_I; A_B)∈R^{N x N}
function rhs!(dc, c, semi, t)
@unpack pde_boundary_matrix = semi.cache
@unpack equations, nodeset_inner, boundary_condition, nodeset_boundary = semi.spatial_discretization
rhs_vector = [rhs(t, nodeset_inner, equations);
boundary_condition.(Ref(t), nodeset_boundary)]
# dc = -pde_boundary_matrix * c + rhs_vector
dc[:] = muladd(pde_boundary_matrix, -c, rhs_vector)
return nothing
end

"""
semidiscetize(semi::Semidiscretization, tspan)
Wrap a `Semidiscretization` object into an `ODEProblem` object with time span `tspan`.
"""
function semidiscretize(semi::Semidiscretization, tspan)
nodeset = merge(semi.spatial_discretization.nodeset_inner,
semi.spatial_discretization.nodeset_boundary)
u0 = semi.initial_condition.(Ref(first(tspan)), nodeset,
Ref(semi.spatial_discretization.equations))
c0 = semi.cache.kernel_matrix \ u0
iip = true # is-inplace, i.e., we modify a vector when calling rhs!
f = ODEFunction{iip}(rhs!, mass_matrix = semi.cache.mass_matrix)
return ODEProblem(f, c0, tspan, semi)
end
73 changes: 73 additions & 0 deletions src/equations.jl
Original file line number Diff line number Diff line change
@@ -0,0 +1,73 @@
abstract type AbstractEquation end

abstract type AbstractStationaryEquation <: AbstractEquation end

function rhs(nodeset::NodeSet, equations::AbstractStationaryEquation)
if equations.f isa AbstractVector
return equations.f
else
return equations.f.(nodeset, Ref(equations))
end
end

@doc raw"""
PoissonEquation(f)
Poisson equation with right-hand side `f`, which can be a space-dependent function or a vector. The Poisson equation is
defined as
```math
-Δu = f
```
"""
struct PoissonEquation{F} <: AbstractStationaryEquation where {F}
f::F

# accept either a function or a vector
function PoissonEquation(f)
return new{typeof(f)}(f)
end
end

function Base.show(io::IO, ::PoissonEquation)
print(io, "-Δu = f")
end

function (::PoissonEquation)(kernel::RadialSymmetricKernel, x, y)
return -Laplacian()(kernel, x, y)
end

abstract type AbstractTimeDependentEquation <: AbstractEquation end

function rhs(t, nodeset::NodeSet, equations::AbstractTimeDependentEquation)
if equations.f isa AbstractVector
return equations.f
else
return equations.f.(Ref(t), nodeset, Ref(equations))
end
end

@doc raw"""
HeatEquation(diffusivity, f)
Heat equation with thermal diffusivity `diffusivity`. The heat equation is defined as
```math
∂_t u = κΔu + f,
```
where `κ` is the thermal diffusivity and `f` is the right-hand side, which can be a time- and space-dependent function or a vector.
"""
struct HeatEquation{RealT, F} <: AbstractTimeDependentEquation
diffusivity::RealT
f::F

function HeatEquation(diffusivity, f)
return new{typeof(diffusivity), typeof(f)}(diffusivity, f)
end
end

function Base.show(io::IO, ::HeatEquation)
print(io, "∂_t u = κΔu + f")
end

function (equations::HeatEquation)(kernel::RadialSymmetricKernel, x, y)
return -equations.diffusivity * Laplacian()(kernel, x, y)
end
Loading

0 comments on commit f188609

Please sign in to comment.