PhysAI: Unified Operator Synthesis & Cross-Validation

Overview

PhysAI is a research library for approximating solutions to partial and ordinary differential equations with neural networks, built around the physics rather than around any single deep learning framework.

• 57 governing equations — elliptic/static problems, diffusion–reaction systems, hyperbolic/wave equations, nonlinear transport, fluid dynamics, quantum and dispersive systems, kinetic/probabilistic equations, excitable media, and relativistic/quantum-field theory (general relativity, the Dirac equation, quantum-gas statistics, phonon transport).
• Three model architectures — a Fourier-feature PINN (physai.models.pinn), a Fourier Neural Operator (physai.models.fno, Li et al. 2020), and a Chebyshev-basis USENO (physai.models.spectral_element / spectralpinn) with C⁰/C¹ interface-continuity losses for stiff multiphysics problems.
• AutoOptimizer — reads equation order, nonlinearity, and stiffness to emit a frozen RuntimeConfig: learning rate, optimizer choice, collocation-point density, loss weighting, warm-up schedule, network size, and precision.
• Arbitrary geometry via SDFs — box, ball, cylinder, half-space, ellipsoid, torus, capsule, 2-D polygon primitives; CSG combinators; user-defined SDFs; mesh import (.stl/.obj/.ply/.off); Dirichlet/Neumann/Robin/periodic boundary conditions on arbitrary regions.
• Numerical cross-validation — Trainer.cross_validate(...) compares the model with an independent Dedalus box solve, embedded-boundary finite differences, or optional native solver adapters for FiPy, FEniCS, FEniCSx, and Meep.
• LaTeX-to-residual AutoSolve — validate supported equation syntax with build_latex_residual(...), then pass the differentiable residual to autosolve(residual, geometry, ...); or register it with register_latex_pde(...) and solve by its PDE name.
• Physics-focused visualization — loss-history, residual-field, spectrum, and 1-D/2-D plots, plus an N-D toolkit for slicing, projection, isosurfaces, and animation.
• Multi-backend by construction — AbstractBackend fixes the tensor/autodiff/optimizer surface; TorchBackend, JAXBackend, TensorFlowBackend, and PaddleBackend each implement it, so residuals and losses are written once and run correctly on all four.
• Optional, consent-gated extras — a live terminal dashboard (rich) with an optional local-LLM chat panel, and a bundled Conda installer for Dedalus, FiPy, FEniCS, FEniCSx, Meep, and CuPy. Nothing installs on pip install/import physai; setup is user-invoked and inert in CI.

Installation

A base install stays light — torch + numpy + matplotlib only. Backends and features are opt-in extras.

pip install physai

pip install "physai[jax]"          # JAX backend (jax, flax, optax)
pip install "physai[tensorflow]"   # TensorFlow backend
pip install "physai[paddle]"       # PaddlePaddle backend
pip install "physai[dashboard]"    # live terminal training dashboard (rich)
pip install "physai[chat]"         # dashboard + local GGUF chat side panel (llama-cpp-python)
pip install "physai[all]"          # every backend and feature extra above

Or, from source:

git clone https://github.com/MS-AGI/PhysAI.git
cd PhysAI
pip install -e ".[jax,tensorflow,dashboard]"

Optional classical solvers use a separate Conda environment. Run python -m physai.solver_setup, activate physai-solvers, then install PhysAI in that environment with python -m pip install physai. The installer includes Dedalus, FiPy, classic FEniCS, FEniCSx, Meep, and CuPy.

Repository Structure

src/physai/
├── core/
│   ├── pde_residual.py      # PDE_REGISTRY: 57 governing-equation residuals
│   ├── latex_pde.py         # Validated LaTeX equations -> backend residuals
│   ├── auto_optimizer.py    # ProblemSpec -> AutoOptimizer -> RuntimeConfig
│   └── losses.py            # residual / dirichlet / neumann / robin / periodic losses
├── backends/                 # AbstractBackend + torch / jax / tensorflow / paddle
├── models/
│   ├── pinn.py               # Fourier-feature PINN, hard-constraint support
│   ├── fno.py                 # Fourier Neural Operator
│   └── spectral_element.py, spectralpinn.py   # USENO (Chebyshev spectral element)
├── solvers/
│   └── solver.py             # Dedalus, embedded-boundary FD, and optional native solver adapters
├── geometry.py                # SDF primitives, CSG, mesh import, BoundaryConditionSet
├── trainer.py                 # Trainer: training loop, callbacks, cross_validate, per-backend step logic
├── visualization.py            # 1-D/2-D plots, loss curves, spectra, animations
├── visualization_nd.py          # slicing / projection / isosurfaces / animation for N-D fields
├── dashboard/live.py            # optional live terminal dashboard (Callback)
├── chat_setup.py, solver_setup.py   # consent-gated optional solver setup
└── utils.py                    # sampling (LHS/Sobol), metrics, seeding, dtype helpers
tests/
└── test_pde_everything.py     # Tier A/B/C suite — see Testing below
examples/
README.md
pyproject.toml
requirements.txt

Quick Start

These three runnable examples demonstrate a user-registered Maxwell residual with animation, a wave PINN compared with AutoSolve or the optional custom solve-box path, and the Einstein vacuum residual with Schwarzschild exterior metric data. Run the commands from the repository root after installing PhysAI and the dependencies required by each example.

1. Register Maxwell's equations and animate the field

Registers a one-dimensional vacuum Maxwell system, trains six field outputs, and animates the transverse electric field Ey. Source: examples/maxwell_animation.py.

python examples/maxwell_animation.py

2. Solve the wave equation and cross-validate

Registers the two-field first-order wave system and trains (u, v). By default it compares against AutoSolve; use --reference solve-box to select the custom Dedalus box-solver path as an alternative. The solve-box option requires the Conda solver environment; see the Installation section. Source: examples/wave_cross_validation.py.

python examples/wave_cross_validation.py

3. Fit a Schwarzschild exterior with the Einstein field residual

Trains einstein_field against the isotropic-coordinate Schwarzschild vacuum metric outside the horizon, then prints a held-out relative L2 metric error. Source: examples/schwarzschild_einstein_residual.py.

python examples/schwarzschild_einstein_residual.py

Model Architectures

Three architectures share the same backend abstraction and residual definitions.

Numerical Cross-Validation

Trainer.cross_validate(...) evaluates a trained model at points from an independent classical solve and reports L2 errors. Box problems can use Dedalus spectral methods; arbitrary geometries use embedded-boundary finite-difference fast paths for Poisson, Helmholtz, and heat, or AutoSolve for other registered PDEs and caller-supplied residuals. Optional native adapters support FiPy, FEniCS, FEniCSx, and Meep. Native adapters accept solver-specific inputs through solver_kwargs; Meep and custom result types may need a result adapter.

The native solver packages are installed through the bundled Conda environment. Run python -m physai.solver_setup, activate physai-solvers, and install PhysAI in that environment before using these adapters.

Build a Residual from LaTeX

build_latex_residual parses the equation when called, checks its symbols and syntax, and returns a residual that uses the selected backend for automatic differentiation. Declare field names, coordinate order, and numeric parameters explicitly. The supported notation includes arithmetic, common scalar functions, partial derivative fractions and subscripts, and scalar Laplacians. Pass this residual directly to autosolve(residual, geometry, ...), or register it with register_latex_pde and ask AutoSolve to build it by name. In either case, provide the geometry and suitable boundary, initial, or data constraints; these are not inferred from the equation.

import numpy as np
from physai.backends import get_backend
from physai import BoundaryConditionSet, autosolve, box, register_latex_pde

backend = get_backend("torch")
equation = r"\frac{\partial u}{\partial t} + u\frac{\partial u}{\partial x} - \nu\frac{\partial^2 u}{\partial x^2} = 0"
register_latex_pde(
    "custom_burgers", equation,
    fields=("u",), coordinates=("x", "t"), parameters={"nu": 0.01},
)

geometry = box([(-1.0, 1.0)])
bcs = BoundaryConditionSet(geometry)
for endpoint in (-1.0, 1.0):
    bcs.add(
        "dirichlet", value=lambda x: np.zeros(len(x)),
        region=lambda x, endpoint=endpoint: np.isclose(x[:, 0], endpoint),
        name=f"x_{endpoint}",
    )

x0 = np.linspace(-1.0, 1.0, 32)
ic_points = np.column_stack((x0, np.zeros_like(x0)))
solution = autosolve(
    geometry=geometry, pde="custom_burgers", backend=backend,
    time_domain=(0.0, 1.0), boundary_conditions=bcs,
    ic_points=ic_points, ic_values=-np.sin(np.pi * x0),
)

For a one-off equation, pass the object returned by build_latex_residual(equation, backend, ...) directly as the first argument to autosolve(residual, geometry, ...). AutoSolve needs the problem constraints; they are not inferred from the PDE. Unsupported notation is rejected during parsing. Vector/tensor operators and arbitrary LaTeX macros are outside the current grammar.

Register the two Maxwell tensor laws

The current parser accepts scalar component equations rather than Einstein-index tensor notation. The runnable example expands the covariant laws ∂μFμν = μ₀Jν and ∂μF̃μν = 0 into Gauss, Ampère–Maxwell, no-magnetic-charge, and Faraday components, then registers all eight residual equations. Its fields include E, B, charge density ρ, and current density J; c and ε₀ are configurable parameters. It then trains a PINN on a periodic 3-D plane wave and animates the predicted E_y slice. The linked file contains the full registration, training, prediction, and animation code.

python examples/maxwell_latex_register.py

Source: examples/maxwell_latex_register.py.

Testing

Coverage and logs are checked in at tests/tests_log.txt. The suite runs across every backend whose framework is importable in the current environment — a missing framework is skipped, not a failure.

pytest tests/test_pde_everything.py -v

Backend Notes

Tensors, autodiff, and optimizers are abstracted behind AbstractBackend, implemented by TorchBackend, JAXBackend, TensorFlowBackend, and PaddleBackend.

Citation

If you use PhysAI in research, academic publication, or official work, citation is required.

@software{singh_physai_2026,
  author    = {Mankrit Singh},
  title     = {PhysAI: A Multi-Backend Physics-Informed Neural Network
               Framework for Solving, Cross-Validating, and Visualizing
               Ordinary and Partial Differential Equations},
  month     = sep,
  year      = 2026,
  publisher = {Zenodo},
  version   = {5.1.0},
  doi       = {10.5281/zenodo.17214724},
  url       = {https://doi.org/10.5281/zenodo.17214724},
  orcid     = {0009-0009-3913-6929}
}