jax-pinn: Physics-Informed Neural Network for Structural Analysis

JAX/Flax implementation of a physics-informed neural network (PINN) for real-time structural beam analysis, with production export to TensorFlow via jax2tf.

Overview

This project demonstrates a complete ML engineering pipeline from research to production:

1. Analytical solvers generate ground-truth training data using Euler-Bernoulli beam theory
2. Flax PI-ResMLP model learns the physics with heteroscedastic uncertainty quantification
3. jax2tf bridge exports the trained JAX model to TensorFlow SavedModel for production serving
4. Gradio demo provides an interactive web interface for beam analysis

Why JAX + jax2tf?

• JAX provides XLA compilation, automatic differentiation, and hardware-agnostic acceleration (CPU/GPU/TPU)
• jax2tf enables deployment to TensorFlow Serving, TF Lite, and TF.js without rewriting the model
• Flax offers a clean functional API for neural networks with explicit parameter management

Architecture

Input Features (7)
    |
    v
Dense(7, 256) -> LayerNorm -> SiLU
    |
    v
ResidualBlock(256) x 4
    |  [LayerNorm -> Dense -> SiLU -> Dropout -> Dense -> + residual]
    v
Dense(256, 128) -> LayerNorm -> SiLU
    |
    +---> Dense(128, 2) --> stress    [mean, log_var]
    +---> Dense(128, 2) --> deflection [mean, log_var]
    +---> Dense(128, 3) --> safety     [safe, marginal, failure]

Physics-Informed Loss

L_total = L_regression + 0.3 * L_classification + 0.1 * L_physics

L_regression:      Heteroscedastic NLL (mean + log-variance)
L_classification:  Cross-entropy for safety category
L_physics:         Energy bounds + safety consistency penalty

Export Pipeline

Flax Model + Params --> jax2tf.convert() --> tf.Module --> tf.saved_model.save()
                                                              |
                                            TF Serving / TF Lite / TF.js

Beam Configurations

Config	Support	Loading	Deflection Formula	Stress Formula
beam_ss_point	Simply Supported	Central Point Load	PL^3 / (48EI)	PL / (4S)
beam_cantilever_point	Cantilever	Tip Point Load	PL^3 / (3EI)	PL / S
beam_fixed_udl	Fixed-Fixed	Uniform Distributed	wL^4 / (384EI)	wL^2 / (12S)

Benchmark Results

Backend	Batch Size	Mean Latency (ms)	Throughput (samples/s)
JAX (CPU)	1	TBD	TBD
JAX (CPU)	32	TBD	TBD
JAX (CPU)	128	TBD	TBD
TF SavedModel	1	TBD	TBD
TF SavedModel	32	TBD	TBD
TF SavedModel	128	TBD	TBD

Installation

# Clone repository
git clone https://github.com/wolfwdavid/jax-pinn.git
cd jax-pinn

# Install dependencies
pip install -e ".[dev]"

# For GPU support
pip install -e ".[gpu]"

Usage

Generate Training Data

from src.data.generate_dataset import generate_full_dataset

samples = generate_full_dataset(n_samples_per_config=10000, seed=42)
print(f"Generated {len(samples)} total samples")

Train Model

import jax
from src.data.dataset import BeamDataset
from src.data.generate_dataset import generate_full_dataset
from src.training.train_state import create_train_state
from src.training.train import train_epoch

# Generate data
samples = generate_full_dataset(n_samples_per_config=5000)
dataset = BeamDataset(samples)
train_ds, val_ds = dataset.train_val_split(val_fraction=0.2)

# Create model and optimizer
key = jax.random.PRNGKey(42)
state = create_train_state(key=key, input_dim=dataset.input_dim)

# Training loop
for epoch in range(100):
    key, epoch_key = jax.random.split(key)
    batches = train_ds.get_batches(batch_size=256, key=epoch_key)
    state, metrics = train_epoch(state, batches, epoch_key)
    print(f"Epoch {epoch}: loss={metrics['total']:.4f}")

Export to TensorFlow

from src.export.jax2tf_export import export_to_saved_model

export_path = export_to_saved_model(
    apply_fn=state.apply_fn,
    params=state.params,
    input_dim=7,
    export_dir="exported_models/beam_pinn",
)
print(f"SavedModel exported to {export_path}")

Run Benchmarks

from src.benchmark.runner import BenchmarkRunner
from src.benchmark.report import generate_report

runner = BenchmarkRunner(num_runs=100, warmup_runs=10)
results = runner.run_all(
    apply_fn=state.apply_fn,
    params=state.params,
    input_dim=7,
    saved_model_dir="exported_models/beam_pinn",
)
print(generate_report(results))

Launch Gradio Demo

python app.py

Testing

# Run all tests
pytest tests/ -v

# Run specific test suite
pytest tests/test_solvers/ -v
pytest tests/test_model/ -v
pytest tests/test_training/ -v

# With coverage
pytest tests/ --cov=src --cov-report=html

Project Structure

jax-pinn/
├── src/
│   ├── solvers/          # Analytical beam solvers (ground truth)
│   │   ├── base.py       # ABC for solvers
│   │   ├── beam.py       # 3 beam configurations
│   │   └── schema.py     # Pydantic data models
│   ├── data/             # Data generation and loading
│   │   ├── generate_dataset.py  # LHS sampling
│   │   └── dataset.py    # JAX-native data loader
│   ├── model/            # Neural network components
│   │   ├── architecture.py      # Flax PI-ResMLP
│   │   ├── physics_loss.py      # JAX physics-informed loss
│   │   └── normalization.py     # Feature normalization
│   ├── training/         # Training loop
│   │   ├── train.py      # JIT-compiled train/eval steps
│   │   ├── train_state.py # Flax TrainState factory
│   │   └── evaluate.py   # Evaluation metrics
│   ├── export/           # Model export
│   │   └── jax2tf_export.py     # JAX -> TF SavedModel
│   ├── benchmark/        # Inference benchmarking
│   │   ├── runner.py     # JAX vs TF latency measurement
│   │   └── report.py     # Results formatting
│   └── utils/            # Utilities
│       └── device.py     # JAX device detection
├── tests/                # Comprehensive test suite
├── configs/              # YAML configuration files
├── app.py                # Gradio web demo
├── Dockerfile            # Multi-stage Docker build
└── pyproject.toml        # Project configuration

Technical Decisions

Decision	Rationale
MLP over Transformer	Tabular regression on 7 numeric features gains nothing from attention. See Grinsztajn et al. (2022).
Residual connections	Prevents vanishing gradients; allows identity mappings where input features directly predict output.
Heteroscedastic outputs	Predicting mean + log-variance provides calibrated uncertainty without ensembles.
Log-space targets	Stress/deflection span 6+ orders of magnitude (Pa to GPa); log-space gives balanced gradients.
jax2tf over ONNX	Direct TF integration; no intermediate format; preserves JAX semantics exactly.
LHS over random	Latin Hypercube Sampling provides better coverage of the parameter space.

References

• Timoshenko, S.P. & Gere, J.M. "Mechanics of Materials"
• Grinsztajn, L. et al. "Why do tree-based models still outperform deep learning on tabular data?" (2022)
• Raissi, M. et al. "Physics-informed neural networks" (2019)
• JAX documentation: https://jax.readthedocs.io/
• Flax documentation: https://flax.readthedocs.io/

Author

David White Wolf (@wolfwdavid)

License

Apache-2.0