Scimba Jax
scimba_jax is a JAX-based, functional rewrite of Scimba. Its
building blocks for domains, approximation spaces and PDE solvers are
designed to be composed and to work with jit and
vmap. It is released as an experimental module: APIs may
change, and its test coverage is lower than that of scimba_torch.
Functional core
Pure functions and pytrees throughout, designed around jax.jit and jax.vmap.
Composable modules
Domains, samplers, approximation spaces and physical models compose the
same way as in scimba_torch.
Active development
New solvers (DG schemes, optimal transport, kinetic models) are developed here first, so the API changes frequently.
Tools / methods map
An overview of the main tools and methods available in Scimba Jax.
Domains & meshes
Meshless and mesh-based domains, mappings, signed-distance functions, and
structured/DG meshes (domains/, linear_approximation/meshes/).
Samplers
Collocation-point sampling strategies (uniform, adaptive and residual-based) over domains, boundaries and time.
Physical models
Elliptic PDEs, Monge-Ampère / optimal transport, Grad-Shafranov and
more (physical_models/).
Mesh-based solvers
FEM, DG and finite volumes, plus multigrid, basis functions and
nonlinear solvers (linear_approximation/).
Collocation-based methods
Kernel and random-feature approximation spaces, sampled once and kept fixed. They require neither a mesh nor any training.
Nonlinear approximation space
Networks (MLP, ResNet, ICNN), Fourier & periodic embeddings,
activation functions, approximation spaces with pre/post-processing
and create_variables().
PINNs & projector
Physics-informed networks and least-squares projectors for function approximation and PDE residual minimization.
Neural operator
Data-driven and physics-informed operator learning, with FNO, GINO, U-Net and DeepONet architectures for parametric solution maps.
ROMs
Reduced-order models built from snapshots of the solution manifold. Under construction.
ODE discovery and learning flows
Learn governing ODEs and flow maps from data, including neural ODE and ODE-flow based approximation strategies.
Particles & Lagrangian approaches
Particle-in-cell and Lagrangian schemes for kinetic and transport problems. Under construction.
Hybrid numerical solver
Combine network-based and classical linear approximation spaces within the same least-squares training loop.
Scimba Torch
scimba_torch is the original PyTorch backend. It covers
function projection as well as elliptic, time-dependent and kinetic
parametric PDEs, with a wide range of training strategies and optimizers.
It no longer receives new developments, but it is fully tested and stable.
Fully tested
The primary backend, with the test suite and coverage reports tracked in CI.
Rich domain support
Meshless and mesh-based domains in 2D/3D, parametric hypersurfaces, level sets.
Many solvers
PINNs, Deep Ritz, neural Galerkin, neural semi-Lagrangian, ODE flows, structure-preserving networks.
Tools / methods map
The packages listed under
Available Packages in the API
documentation, which is generated from the source code with autosummary.
Domains & geometry
domain/ and geometry/: meshless and mesh-based domains, parametric hypersurfaces.
Approximation & neural nets
approximation_space/ and neural_nets/, including structure-preserving networks.
Physical models & flows
physical_models/ and flows/: elliptic PDEs, ODEs and time-dependent PDEs.
Numerical solvers & optimizers
numerical_solvers/ and optimizers/: PINN training loops, natural gradient, SS-Broyden.