Scimba Jax
scimba_jax is a JAX-based, functional rewrite of Scimba. It
favors composable, jit/vmap-friendly building
blocks for domains, approximation spaces and PDE solvers. It is released as
an experimental module: APIs may change, and test coverage does not yet
match 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) land here first — expect rapid iteration.
Tools / methods map
A quick tour of the main tools and methods proposed by 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 held fixed — no mesh, no 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 — FNO, GINO, U-Net and DeepONet architectures for parametric solution maps.
ROMs
Reduced-order models built from snapshots of the solution manifold. In 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. In 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, fully-tested 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.
This is the recommended starting point for production-grade SciML workflows.
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 docs —
generated automatically from the source via autosummary.
Domains & geometry
domain/, geometry/ — meshless/mesh-based domains, parametric hypersurfaces.
Approximation & neural nets
approximation_space/, neural_nets/ — including structure-preserving networks.
Physical models & flows
physical_models/, flows/ — elliptic, ODE, and time-dependent PDEs.
Numerical solvers & optimizers
numerical_solvers/, optimizers/ — PINN training loops, natural gradient, SS-Broyden.