FastLSQ 0.1.5

Solving PDEs in one shot via Fourier features with exact analytical derivatives

FastLSQ

BerkeleyLab ATAP Talk

Solving PDEs in one shot via Fourier features with exact analytical derivatives.

FastLSQ is a lightweight PDE solver built around SinusoidalBasis, an analytical derivative engine for random Fourier features. For sinusoidal features phi_j(x) = sin(W_j . x + b_j), every derivative of every order admits an exact closed-form expression -- no automatic differentiation needed.

Linear PDEs are solved in a single least-squares step; nonlinear PDEs are solved via Newton-Raphson iteration with Tikhonov regularisation, 1/sqrt(N) feature normalisation, and continuation/homotopy.

Installation

pip install fastlsq

For development (includes testing and build tools):

git clone https://github.com/asulc/FastLSQ.git
cd FastLSQ
pip install -e ".[dev]"

Quick start

Solve a linear PDE in one line

from fastlsq import solve_linear
from fastlsq.problems.linear import PoissonND

problem = PoissonND()
result = solve_linear(problem, scale=5.0)

u_fn = result["u_fn"]
print(f"Value error: {result['metrics']['val_err']:.2e}")

Solve a nonlinear PDE

from fastlsq import solve_nonlinear
from fastlsq.problems.nonlinear import NLPoisson2D

problem = NLPoisson2D()
result = solve_nonlinear(problem, max_iter=30)

print(f"Converged in {result['n_iters']} iterations")
print(f"Value error: {result['metrics']['val_err']:.2e}")

Use the basis directly

import torch
from fastlsq.basis import SinusoidalBasis

basis = SinusoidalBasis.random(input_dim=2, n_features=1500, sigma=5.0)
x = torch.rand(5000, 2)

# Arbitrary mixed partial via multi-index
d2_dxdy = basis.derivative(x, alpha=(1, 1))

# Or use fast-path methods
H     = basis.evaluate(x)            # (5000, 1500)
dH    = basis.gradient(x)            # (5000, 2, 1500)
lap_H = basis.laplacian(x)           # (5000, 1500)

Compose PDE operators symbolically

import torch
from fastlsq.basis import SinusoidalBasis, Op

basis = SinusoidalBasis.random(input_dim=2, n_features=1500, sigma=5.0)
x = torch.rand(5000, 2)

# Coefficients can be scalars or nn.Parameter (for AdamW optimisation)
k, c = 10.0, 2.0
helmholtz = Op.laplacian(d=2) + k**2 * Op.identity(d=2)
A_pde = helmholtz.apply(basis, x)    # (5000, 1500)

wave = Op.partial(dim=2, order=2, d=3) - c**2 * Op.laplacian(d=3, dims=[0, 1])

Vector-valued solutions

solve_linear / solve_nonlinear support vector-valued u: ℝᵈ → ℝᵏ for coupled systems (elasticity, Stokes, Maxwell vector potential, …) and for decoupled multi-output problems sharing one basis. The math is unchanged; the solver just allocates beta with shape (N, k) so that solver.predict(x) returns shape (M, k) directly.

A problem opts in by setting self.n_outputs = k and assembling its operator in block-stacked form A ∈ ℝ^{Mk × Nk}, b ∈ ℝ^{Mk × 1}. The helper block_concat removes the manual torch.cat bookkeeping:

import torch
from fastlsq import solve_linear, block_concat

class Stokes2D:
    n_outputs = 3        # (u, v, p)
    dim = 2
    name = "Stokes 2D"
    # ... exact, exact_grad, get_train_data, get_test_points ...

    def build(self, slv, x, bcs, f):
        basis = slv.basis
        cache = basis.cache(x)
        dx = basis.derivative(x, (1, 0), cache=cache)
        dy = basis.derivative(x, (0, 1), cache=cache)
        lap = basis.laplacian(x, cache=cache)

        # Rows = equations (mom_x, mom_y, continuity);
        # columns = coefficient blocks (u, v, p)
        A = block_concat([
            [-lap,  None,  dx  ],   # -Δu + ∂p/∂x = f_x
            [ None, -lap,  dy  ],   # -Δv + ∂p/∂y = f_y
            [ dx,   dy,    None],   #  ∂u/∂x + ∂v/∂y = 0
        ])
        b = block_concat([[f[:, 0:1]], [f[:, 1:2]], [torch.zeros_like(f[:, 0:1])]])
        # ... add BC blocks the same way ...
        return A, b

result = solve_linear(Stokes2D(), scale=5.0)
u = result["u_fn"](x_test)        # shape (M, 3): columns are (u, v, p)

Partial derivatives for a vector u

The basis-level operators (basis.derivative, basis.gradient, basis.laplacian, DiffOperator.apply) all return shape (M, N) regardless of how many components u has — vector-ness only enters when you contract with beta:

# Full Jacobian, then slice (M, d, k) -> per (component, dim)
u, J = solver.predict_with_grad(x)   # J shape (M, d, k); J[:, j, c] = ∂u_c/∂x_j

# Single operator on a single component
D_y = solver.basis.derivative(x, alpha=(0, 1))   # (M, N): ∂φ/∂y
du0_dy = D_y @ solver.beta[:, 0:1]               # ∂u_0/∂y

# Symbolic operator, all components at once
from fastlsq import Op
yy = Op.partial(dim=1, order=2, d=2)
A  = yy.apply(solver.basis, x)                   # (M, N)
u_yy = A @ solver.beta                           # (M, k): ∂²u/∂y² per component

Scalar problems are untouched: n_outputs defaults to 1, solver.beta keeps shape (N, 1), and predict_with_grad returns gradient shape (M, d) for backward compatibility (the trailing component axis is squeezed when k=1). ElasticWave2D in fastlsq/problems/linear.py is the canonical coupled vector example.

Plot solutions

from fastlsq.plotting import plot_solution_2d_contour, plot_convergence

plot_solution_2d_contour(result["solver"], problem, save_path="solution.png")
plot_convergence(result["history"], problem_name=problem.name, save_path="convergence.png")

Benchmarks

# Linear PDE benchmark (Fast-LSQ vs PIELM)
python examples/run_linear.py

# Nonlinear PDE benchmark (Newton-Raphson)
python examples/run_nonlinear.py

# Learnable Helmholtz wavenumber (nn.Parameter + AdamW)
python examples/learnable_helmholtz.py

Inverse problems

The analytical derivatives enable gradients through the pre-factored solve, making inverse problems tractable. Example: recovering 4 anisotropic Gaussian heat sources (24 parameters) from 4 sparse sensors. The heat equation is solved in space-time; L-BFGS-B optimises source positions and shapes to match sensor time-series. (Click image for animation.)

python examples/inverse_heat_source.py

Core architecture

The framework is built around SinusoidalBasis -- the analytical derivative engine:

Class	Purpose
SinusoidalBasis	Evaluates basis functions and arbitrary-order derivatives in O(1) via the cyclic identity
BasisCache	Pre-computes sin(Z)/cos(Z) once, reuses across multiple derivative evaluations
DiffOperator / Op	Symbolic linear differential operators that compose via +, -, scalar *; coefficients can be nn.Parameter for learnable PDEs
FeatureBasis	Adapter for non-sinusoidal solvers (e.g. PIELM with tanh)
FastLSQSolver	Manages feature blocks; exposes .basis for all derivative computations
LearnableFastLSQ	Differentiable solver with learnable bandwidth via reparameterisation trick
block_concat, pack_beta, unpack_beta	Block-structured assembly helpers for vector-valued u (coupled systems). solver.beta has shape (N, k); scalar problems are the k=1 case

How it works

1. Basis construction. Given collocation points x, construct a SinusoidalBasis with random weights W and biases b.

2. Analytical derivatives. Exploit the cyclic derivative identity: the n-th derivative of sin(z) cycles through {sin, cos, -sin, -cos} with monomial weight prefactors. Any mixed partial D^alpha phi_j(x) is computed in O(1) -- no computational graph, no automatic differentiation.

3. PDE assembly. Define the differential operator symbolically with Op (e.g. Op.laplacian(d=2)) and apply it to the basis to get the system matrix A.

4. Linear solve. Solve A beta = b via least squares (optionally Tikhonov-regularised).

5. Newton iteration (nonlinear). Linearise the PDE residual, solve J delta_beta = -R with backtracking line search, and repeat.

Adding your own PDE

Define a problem class and use solver.basis to build the linear system:

import torch, numpy as np
from fastlsq import solve_linear, Op
from fastlsq.geometry import sample_box, sample_boundary_box

class MyPoisson2D:
    def __init__(self):
        self.name = "My Poisson"
        self.dim = 2
        self.pde_op = -Op.laplacian(d=2)

    def exact(self, x):
        return torch.sin(np.pi * x[:, 0:1]) * torch.sin(np.pi * x[:, 1:2])

    def exact_grad(self, x):
        sx, cx = torch.sin(np.pi * x[:, 0:1]), torch.cos(np.pi * x[:, 0:1])
        sy, cy = torch.sin(np.pi * x[:, 1:2]), torch.cos(np.pi * x[:, 1:2])
        return torch.cat([np.pi * cx * sy, np.pi * sx * cy], dim=1)

    def source(self, x):
        return 2 * np.pi**2 * self.exact(x)

    def get_train_data(self, n_pde=5000, n_bc=1000):
        x_pde = sample_box(n_pde, self.dim)
        f_pde = self.source(x_pde)
        x_bc = sample_boundary_box(n_bc, self.dim)
        u_bc = self.exact(x_bc)
        return x_pde, [(x_bc, u_bc)], f_pde

    def build(self, solver, x_pde, bcs, f_pde):
        basis = solver.basis
        cache = basis.cache(x_pde)
        A_pde = self.pde_op.apply(basis, x_pde, cache=cache)
        As, bs = [A_pde], [f_pde]
        for (x_bc, u_bc) in bcs:
            As.append(100.0 * basis.evaluate(x_bc))
            bs.append(100.0 * u_bc)
        return torch.cat(As), torch.cat(bs)

    def get_test_points(self, n=5000):
        return sample_box(n, self.dim)

result = solve_linear(MyPoisson2D(), scale=5.0)

See examples/add_your_own_pde.py for the complete tutorial.

Features

• Analytical derivative engine: SinusoidalBasis computes arbitrary-order derivatives exactly in O(1) -- the foundation of the entire framework
• Symbolic PDE operators: Compose differential operators with Op (Laplacian, wave, Helmholtz, biharmonic, custom) via intuitive arithmetic; coefficients can be nn.Parameter for AdamW optimisation
• Vector-valued solutions: First-class support for u: ℝᵈ → ℝᵏ (elasticity, Stokes, Maxwell). Problems declare n_outputs = k; block_concat assembles coupled block systems; solver.predict(x) returns shape (M, k). Scalar problems are the k=1 case
• High-level API: Solve PDEs in one line with solve_linear() and solve_nonlinear()
• Learnable bandwidth: LearnableFastLSQ optimises the bandwidth (scalar or anisotropic) via reparameterisation
• Learnable PDE coefficients: Plug nn.Parameter into Op (e.g. Helmholtz wavenumber k) and optimise via AdamW; gradients flow through the prebuilt linear solve
• Auto-tuning: Automatic scale selection via grid search
• Built-in plotting: Solution visualization, convergence plots, spectral sensitivity
• Geometry samplers: Box, ball, sphere, interval, custom samplers
• Diagnostics: Problem validation, conditioning checks, error detection
• Export utilities: NumPy conversion, checkpoint saving/loading
• PyTorch Lightning: Integration for training loops
• 20+ benchmark problems: Linear, nonlinear, and regression-mode PDEs

Paper

The full preprint is available on arXiv

Citing this work

If you use FastLSQ in your research, please cite:

@misc{sulc2026fastlsqframeworkoneshotpde,
      title={FastLSQ: A Framework for One-Shot PDE Solving}, 
      author={Antonin Sulc},
      year={2026},
      eprint={2602.10541},
      archivePrefix={arXiv},
      primaryClass={math.NA},
      url={https://arxiv.org/abs/2602.10541}, 
}

License

This project is licensed under the MIT License -- see LICENSE for details.