← Back to Scimba Jax Module · scimba_jax

Collocation-based methods

No mesh, no cells: the solution is a combination of basis functions spread over the whole domain, and the PDE is enforced by requiring the residual to vanish AT a finite set of sampled points — the collocation points, which double as the basis functions' own centers. Two families live here, both with their centers and shapes drawn once and kept fixed. The solve is written in JAX and is differentiable, exactly like the mesh-based solvers — the same linear solve can sit inside an inverse problem.

1. Kernel methods

Every basis function is a translate of the same radial kernel, centred on one sampled point \(x_j\):

\(\varphi_j(x) = k(x, x_j)\)

and the approximation is the linear combination \(u_h(x) = \sum_j c_j\, \varphi_j(x)\). The linear system is dense (every point sees every other), but for a modest number of points this is often the fastest way to a very accurate answer, with nothing to train beyond the coefficients \(c_j\). A few of the kernels available, with \(r = \lVert x - x_j \rVert\):

\(k_{\text{Gaussian}}(x, x_j) = \exp\!\left(-\dfrac{r^2}{2\sigma^2}\right)\)

\(k_{\text{Matérn-3/2}}(x, x_j) = \left(1 + \dfrac{\sqrt{3}\,r}{\sigma}\right)\exp\!\left(-\dfrac{\sqrt{3}\,r}{\sigma}\right)\)

\(k_{\text{Matérn-5/2}}(x, x_j) = \left(1 + \dfrac{\sqrt{5}\,r}{\sigma} + \dfrac{5 r^2}{3\sigma^2}\right)\exp\!\left(-\dfrac{\sqrt{5}\,r}{\sigma}\right)\)

The Matérn family trades the Gaussian's infinite smoothness for a fixed, finite one — useful when the solution itself is not perfectly smooth. A fifth option, FundamentalSolutionKernel, is not a shape choice at all: it plugs in the PDE's own Green's function (\(\log r\) in 2D for the Laplacian), so the basis already solves the equation away from its center. 

from scimba_jax.linear_approximation.basis.kernel_basis import KernelBasis
from scimba_jax.linear_approximation.basis.kernel_function import GaussianKernel
from scimba_jax.linear_approximation.collocation.collocation_elliptic import (
    EllipticCollocationScheme,
)
from scimba_jax.linear_approximation.variables.collocation_variables import (
    CollocationVariables,
)

# No mesh, no cells: ONE basis function per sampled point -- the centers
# ARE the collocation points, and the PDE is enforced AT those same points.
basis = KernelBasis(
    dim=2,
    output_dim=1,
    kernel_function=GaussianKernel(sigma=0.5),
    centers=interior_points,       # (n_points, 2), drawn by a DomainSampler
    learnable_center_bool=False,   # frozen: only the linear weights are fit
    basis_type="scalar",
)
variables = CollocationVariables(basis=basis, nb_variables=1)
scheme = EllipticCollocationScheme(
    pde=pde,
    variables=variables,
    collocation_points=interior_points,
    bc_collocation_points=boundary_points,
)
scheme = EllipticCollocationScheme.solve(scheme)
# 600 interior + 80 boundary points, sigma=0.5: 3.0e-06 relative L2 error
# against sin(pi x) sin(pi y) on the unit square.

# Swapping GaussianKernel for another kernel_function is the whole change --
# same basis, same scheme. At the same pair of points 0.5 apart:
#   GaussianKernel(sigma=0.5)            -> 0.6065
#   Matern32Kernel(sigma=0.5)            -> 0.4834
#   Matern52Kernel(sigma=0.5)            -> 0.5240
#   ExponentialKernel(sigma=0.5)         -> 0.3679  (C0 only, sharp at r=0)
#   FundamentalSolutionKernel(dim=2)     -> -0.1103 (the Laplacian's own
#                                           Green's function, log(r)/2pi)

2. Random features

Same collocation scheme, a nonlinear basis instead of a radial one:

\(\varphi_j(x) = \tanh\!\left(\dfrac{W_j \cdot (x - x_j)}{s} + b_j\right)\)

with \((W_j, b_j)\) drawn from a fixed random distribution once and never touched again — the extreme-learning-machine convention. Only the downstream coefficients \(c_j\) are fit, exactly as for a kernel. 

from scimba_jax.linear_approximation.basis.random_network_basis import (
    RandomFeaturesBasis,
)

# Same collocation scheme, a nonlinear basis instead: (W_j, b_j) drawn ONCE
# and frozen -- the extreme-learning-machine convention.
basis = RandomFeaturesBasis(
    dim=2,
    output_dim=1,
    centers=interior_points,
    scale=0.5,
    activation="tanh",
    learnable_center_bool=False,    # the defaults -- frozen is the point
    learnable_features_bool=False,  # of a RANDOM feature
)
variables = CollocationVariables(basis=basis, nb_variables=1)
scheme = EllipticCollocationScheme(
    pde=pde, variables=variables,
    collocation_points=interior_points, bc_collocation_points=boundary_points,
)
scheme = EllipticCollocationScheme.solve(scheme)
# Same problem as above: 1.4e-08 relative L2 error -- a random nonlinear
# feature can outperform a fixed kernel shape, at the cost of a PRNG key to
# keep track of.

3. Localized random features (RFM)

The Random Feature Method wraps each patch's random features in a smooth window that vanishes outside it, so the assembled system is sparse the way a mesh-based one is — without a mesh. It is the bridge between a purely meshless method and the locality a finite element or finite volume scheme gets for free. 

from scimba_jax.linear_approximation.basis.random_network_basis import (
    RandomFeaturePOUBasis,
)

# Random Feature Method (Chen, Chi, E & Yang, 2022): features LOCALISED to
# a patch by a partition-of-unity window instead of the whole domain --
# sparse the way a mesh-based basis is, without building a mesh.
basis = RandomFeaturePOUBasis(
    dim=2,
    output_dim=1,
    centers=patch_centers,       # (n_patches, 2)
    radii=0.6,                   # shared, or (n_patches, 2) for anisotropic patches
    n_features_per_patch=8,
    activation="tanh",
)

Every basis on this page is frozen by construction — centers and shapes are drawn once, not learned. Letting them move under gradient descent instead turns the same basis into a small neural network, which is the territory of PINNs rather than this page.