jaxdiffmodal

Fast, differentiable, and GPU-accelerated simulation framework for modelling the dynamics of strings, membranes, and plates using modal methods implemented in JAX.

Features

• Differentiable implementation using JAX
• Simulates linear and nonlinear models:
  • Tension-modulated string (Kirchhoff–Carrier)
  • Tension-modulated membrane (Berger model)
  • von Kármán nonlinear plate
• Fast GPU-accelerated time integration
• Designed for real-time synthesis, inverse modelling, and dataset generation
• Includes example notebooks to reproduce results from the paper

Installation for Development

It is recommended to use the uv package manager to install the environment and dependencies.

uv sync --extra dev --extra benchmark

Otherwise you can create a virtual environment and install the dependencies manually:

python3 -m venv .venv
source .venv/bin/activate
pip install -e .

Testing

The project includes comprehensive tests to validate the Python implementation against MATLAB reference code. Tests follow the MATLAB-Python cross-validation pattern documented in CLAUDE.md.

Running Tests

Prerequisites: The tests use pytest and require scipy for eigenvalue computations.

Run all tests from the project root directory:

python -m pytest tests/ -v

Run specific test modules:

# Test K and M matrix assembly with MATLAB comparison
python -m pytest tests/test_K_M_matlab_comparison.py -v

# Test integral functions (int1, int2, int4)
python -m pytest tests/test_intx_matlab_comparison.py -v

# Test matrix integration functions (i_mat)
python -m pytest tests/test_imat_matlab_comparison.py -v

Generating MATLAB Reference Data

To regenerate MATLAB reference data (requires MATLAB installation):

cd tests
matlab -batch test_K_M_matlab_reference
matlab -batch test_intx_matlab_reference_fixed
matlab -batch test_imat_matlab_reference

Test Coverage

• K, M matrix assembly: Perfect match with MATLAB (1e-8 tolerance)
• Airy stress coefficients: Validated using MATLAB eigenvalues/eigenvectors as input (1e-6 tolerance)
• Integral functions: Comprehensive parameter space validation
• Basic functionality: Matrix properties, dimensions, and mathematical consistency

See CLAUDE.md for detailed testing methodology and patterns.

Examples

The nbs/examples directory includes:

• Synthetic and real-world inverse modelling experiments for strings and plates
• Optimisation of nonlinear parameters and coupling tensors
• Scripts to reproduce figures from the paper

The benchmark directory includes comparisons against:

• An optimised C++ implementation using Eigen and BLAS
• A JIT-compiled PyTorch implementation (GPU)
• A MATLAB baseline

We plan to add more benchmarks, examples, and real-time synthesis demos in future updates.

Acknowledgements

• Mode processing adapted from VKPlate
• Plate mode computation using magpie-python
• Coupling coefficient implementation based on VKGong