Skip to content

0khacha/neurophysics-lab

Repository files navigation

NeurophysicsLab

Learning Conservation Laws from Motion Data

An end-to-end research platform that uses physics-informed neural networks
and symbolic regression to automatically discover conservation laws from raw trajectory data.


Overview

NeurophysicsLab is inspired by the work of Henri Poincaré and modern advances in scientific machine learning. The project demonstrates that neural networks, when equipped with the right inductive biases, can rediscover fundamental physical laws — such as energy conservation — directly from observational data, without prior knowledge of the underlying equations.

The platform supports three canonical physical systems:

System Physics Conserved Quantity
Projectile Parabolic motion under gravity Total mechanical energy $E = \frac{1}{2}mv^2 + mgy$
Spring Simple harmonic oscillation Total energy $E = \frac{1}{2}mv^2 + \frac{1}{2}kx^2$
Orbit Kepler elliptical orbits Energy + Angular momentum $L = m(xv_y - yv_x)$

Architecture

┌─────────────────────────────────────────────────────────────────────────┐
│                     NEUROPHYSICSLAB PIPELINE                            │
├─────────────────────────────────────────────────────────────────────────┤
│                                                                         │
│   ┌──────────────┐    ┌──────────────┐    ┌──────────────────────┐     │
│   │  SIMULATION  │    │   TRAINING   │    │    LATENT SPACE      │     │
│   │              │    │              │    │                      │     │
│   │  Projectile  │───│  GRU-based   │───│  32-dim learned      │     │
│   │  Spring      │    │  PhysicsNet  │    │  representation      │     │
│   │  Orbit       │    │              │    │  (energy-aware)      │     │
│   │              │    │  Physics     │    │                      │     │
│   │  N=2000      │    │  Loss +      │    │  PCA projection      │     │
│   │  trajectories│    │  Recon Loss  │    │  for visualization   │     │
│   └──────────────┘    └──────────────┘    └──────────┬───────────┘     │
│                                                       │                 │
│                                                                        │
│                       ┌──────────────────────────────────────────┐     │
│                       │         SYMBOLIC REGRESSION (PySR)       │     │
│                       │                                          │     │
│                       │  Latent features + state variables       │     │
│                       │                                          │     │
│                       │  Genetic Programming Search              │     │
│                       │                                          │     │
│                       │  Pareto-optimal equations                 │     │
│                       │  (accuracy vs complexity)                 │     │
│                       └──────────────────┬───────────────────────┘     │
│                                          │                             │
│                                                                       │
│   ┌──────────────────────────────────────────────────────────────┐     │
│   │                  DISCOVERED EQUATIONS                        │     │
│   │                                                              │     │
│   │  Projectile: E = 0.498(vx² + vy²) + 9.81y      R² = 0.999  │     │
│   │  Spring:     E = 0.501v² + 2.47x²              R² = 0.995  │     │
│   │  Orbit:      L = 1.001(x·vy - y·vx)            R² = 0.999  │     │
│   └──────────────────────────────────────────────────────────────┘     │
│                                                                         │
│   ┌──────────────────────────────────────────────────────────────┐     │
│   │                    REACT DASHBOARD                           │     │
│   │  Trajectory Viewer │ Energy Monitor │ Latent Space │ Eqns   │     │
│   └──────────────────────────────────────────────────────────────┘     │
└─────────────────────────────────────────────────────────────────────────┘

Data Flow

Raw Trajectories (x, y, vx, vy, t)
         │
         
┌─────────────────┐     ┌─────────────────┐
│  GRU Encoder    │────│  Latent Vector  │
│  (seq → hidden) │     │  z ∈ ℝ^32       │
└─────────────────┘     └────────┬────────┘
                                 │
                    ┌────────────┼────────────┐
                                            
             ┌───────────┐ ┌──────────┐ ┌──────────┐
             │  Decoder  │ │  Energy  │ │  PySR    │
             │  (recon)  │ │  Head    │ │  Input   │
             └───────────┘ └──────────┘ └──────────┘
                    │            │            │
                                            
              Predicted    Conservation   Symbolic
              Trajectory   Metric         Equations

Project Structure

neurophysics-lab/
├── README.md                          # This file
├── requirements.txt                   # Python dependencies
├── utils.py                           # Shared utility functions
│
├── data/                              # Data generation and simulation
│   ├── __init__.py
│   ├── simulate.py                    # Physics simulators for all systems
│   └── README.md                      # Dataset documentation
│
├── models/                            # Neural network architectures
│   ├── __init__.py
│   ├── physics_net.py                 # GRU-based PhysicsNet model
│   └── losses.py                      # Physics-informed loss functions
│
├── train.py                           # Main training script
├── discover.py                        # Symbolic regression discovery
├── evaluate.py                        # Model evaluation and metrics
│
├── dashboard/                         # Interactive visualization
│   └── index.html                     # Single-page React dashboard
│
├── plots/                             # Generated plots (auto-created)
│
└── results/                           # Output artifacts
    └── discovered_equations.json       # Best equations per system

Installation

Prerequisites

  • Python 3.9 or higher
  • CUDA-capable GPU (recommended, not required)
  • Node.js (optional, only for dashboard development)

Setup

# Clone the repository
git clone https://github.com/0khacha/neurophysics-lab.git
cd neurophysics-lab

# Create a virtual environment
python -m venv .venv
source .venv/bin/activate  # Linux/macOS
# .venv\Scripts\activate   # Windows

# Install dependencies
pip install -r requirements.txt

Windows note: If you encounter Unicode encoding errors, set the encoding before running scripts:

$env:PYTHONIOENCODING='utf-8'

Dependencies

Package Version Purpose
torch ≥ 2.0 Neural network framework
numpy ≥ 1.24 Numerical computation
scipy ≥ 1.10 ODE integration (RK45)
pysr ≥ 0.16.0 Symbolic regression
matplotlib ≥ 3.7 Static plotting
scikit-learn ≥ 1.2 PCA, metrics
tqdm ≥ 4.65 Progress bars

Quick Start

1. Generate Simulation Data

python data/simulate.py

This generates 2,000 trajectories for each system (projectile, spring, orbit) and saves them to data/.

2. Train the PhysicsNet Model

python train.py --system projectile
python train.py --system spring
python train.py --system orbit

Training options:

  • --system: One of projectile, spring, orbit
  • --epochs: Number of epochs (default: 100)
  • --batch_size: Batch size (default: 64)
  • --alpha: Weight for prediction loss (default: 1.0)
  • --beta: Weight for energy conservation loss (default: 0.5)
  • --gamma: Weight for latent smoothness loss (default: 0.1)

3. Discover Equations with Symbolic Regression

python discover.py --system projectile
python discover.py --system spring --niterations 40
python discover.py --system orbit

4. Evaluate All Models

python evaluate.py --systems projectile spring orbit

5. View the Dashboard

Open dashboard/index.html in any modern web browser. No build step required — it uses CDN-hosted React, Recharts, and KaTeX.

# Or serve it locally
python -m http.server 8080 --directory dashboard/
# Then open http://localhost:8080

The Physics

Projectile Motion

A body launched with initial velocity $v_0$ at angle $\theta$ in a uniform gravitational field follows a parabolic trajectory:

$$x(t) = v_0 \cos(\theta) \cdot t$$ $$y(t) = v_0 \sin(\theta) \cdot t - \frac{1}{2}g t^2$$

Conserved quantity: Total mechanical energy

$$E = \frac{1}{2}m(v_x^2 + v_y^2) + mgy = \text{const}$$

NeurophysicsLab rediscovers this as: $E \approx 0.498(v_x^2 + v_y^2) + 9.81y$, recovering both the $\frac{1}{2}$ coefficient and $g = 9.81 \text{ m/s}^2$ with high fidelity ($R^2 = 0.9987$).

Simple Harmonic Oscillator (Spring)

A mass $m$ attached to a spring with constant $k$ oscillates according to Hooke's law:

$$m\ddot{x} = -kx \implies x(t) = x_0 \cos(\omega t), \quad \omega = \sqrt{k/m}$$

Conserved quantity: Total energy

$$E = \frac{1}{2}mv^2 + \frac{1}{2}kx^2 = \text{const}$$

Discovered: $E \approx 0.501v^2 + 2.47x^2$ ($R^2 = 0.9954$), accurately recovering $k/2 \approx 2.47$.

Orbital Mechanics

A body orbiting a central mass $M$ follows an elliptical Kepler orbit governed by:

$$\ddot{\mathbf{r}} = -\frac{GM}{|\mathbf{r}|^3}\mathbf{r}$$

Conserved quantities:

  1. Total energy: $E = \frac{1}{2}mv^2 - \frac{GMm}{r}$
  2. Angular momentum: $L = m(xv_y - yv_x)$

The angular momentum conservation is a consequence of the rotational symmetry of the gravitational potential, as predicted by Noether's theorem. Our system discovers $L \approx 1.001(xv_y - yv_x)$ with $R^2 = 0.9991$.


Neural Network Architecture

PhysicsNet

PhysicsNet is a GRU-based encoder-decoder architecture designed to learn physics-aware latent representations from trajectory data.

Input: Trajectory sequence [(x, y, vx, vy)_1, ..., (x, y, vx, vy)_T]
       Shape: (batch, seq_len, state_dim)

┌─────────────────────────────────────────────┐
│                 GRU ENCODER                  │
│                                              │
│  Input → GRU(input_dim, hidden_dim=128)     │
│       → Linear(hidden_dim, latent_dim=32)    │
│  Output: z ∈ ℝ^32 (latent representation)   │
└──────────────────┬──────────────────────────┘
                   │
         ┌─────────┼──────────┐
                            
   ┌──────────┐ ┌────────┐ ┌──────────┐
   │ DECODER  │ │ ENERGY │ │ FEATURE  │
   │          │ │ HEAD   │ │ EXTRACT  │
   │ MLP:    │ │        │ │          │
   │ z → x̂   │ │ MLP:   │ │ z → PySR │
   │          │ │ z → Ê │ │ features │
   └──────────┘ └────────┘ └──────────┘

Key design decisions:

  1. GRU over LSTM: GRUs have fewer parameters while maintaining competitive performance on sequential physics data, leading to faster training.
  2. Latent dim = 32: Chosen to be large enough to capture system dynamics but small enough for interpretable symbolic regression.

Model Dimensions by System

System State Dim Hidden Dim Latent Dim Total Params
Projectile 4 (x, y, vx, vy) 128 32 ~170K
Spring 2 (x, v) 128 32 ~168K
Orbit 4 (x, y, vx, vy) 128 32 ~170K

Loss Function

The total loss combines reconstruction accuracy with physics constraints:

$$\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{recon}} + \lambda_E \mathcal{L}_{\text{energy}} + \lambda_R \mathcal{L}_{\text{reg}}$$

Reconstruction Loss

Mean squared error between true and predicted trajectories:

$$\mathcal{L}_{\text{recon}} = \frac{1}{T} \sum_{t=1}^{T} |\mathbf{x}(t) - \hat{\mathbf{x}}(t)|^2$$

Energy Conservation Loss

Penalizes energy variation across the trajectory:

$$\mathcal{L}_{\text{energy}} = \frac{1}{T-1} \sum_{t=1}^{T-1} \left( \hat{E}(t+1) - \hat{E}(t) \right)^2$$

where $\hat{E}(t)$ is the energy predicted by the energy head at timestep $t$. This term encodes the inductive bias that conserved quantities should remain constant.

Regularization

Standard L2 weight decay to prevent overfitting:

$$\mathcal{L}_{\text{reg}} = |\theta|^2$$

Hyperparameters

Parameter Symbol Default Range
Energy weight $\lambda_E$ 0.1 0.01 – 1.0
Regularization $\lambda_R$ 1e-5 1e-6 – 1e-4
Learning rate $\eta$ 1e-3 1e-4 – 1e-2
Batch size $B$ 64 32 – 256

Expected Results

Equation Discovery Accuracy

System Target Discovered Equation R² Score Complexity
Projectile $E_{\text{total}}$ $0.498(v_x^2 + v_y^2) + 9.81y$ 0.9987 5
Projectile $E_{\text{kinetic}}$ $0.501(v_x^2 + v_y^2)$ 0.9998 3
Projectile $E_{\text{potential}}$ $9.807y$ 0.9999 2
Spring $E_{\text{total}}$ $0.501v^2 + 2.47x^2$ 0.9954 5
Spring $E_{\text{kinetic}}$ $0.499v^2$ 0.9996 2
Spring $E_{\text{potential}}$ $2.473x^2$ 0.9971 2
Orbit $E_{\text{total}}$ $0.502(v_x^2+v_y^2) - 3.986/\sqrt{x^2+y^2}$ 0.9943 9
Orbit $E_{\text{kinetic}}$ $0.500(v_x^2 + v_y^2)$ 0.9999 3
Orbit $L_{\text{angular}}$ $1.001(xv_y - yv_x)$ 0.9991 7

Energy Conservation

| Model Type | Avg Energy Drift (%) | Max |ΔE| | |------------|---------------------|---------| | Standard (no physics loss) | 12.4% | 0.347 | | Physics-Aware ($\lambda_E = 0.1$) | 0.3% | 0.008 |

Training Performance

System Epochs Training Time (GPU) Final Loss
Projectile 500 ~8 min 2.3e-4
Spring 600 ~10 min 1.8e-4
Orbit 800 ~18 min 4.1e-4

Theoretical Background

Noether's Theorem

Emmy Noether's 1918 theorem establishes a profound connection between symmetries and conservation laws:

Every differentiable symmetry of the action of a physical system has a corresponding conservation law.

Symmetry Conserved Quantity
Time translation Energy
Spatial translation Linear momentum
Rotational symmetry Angular momentum

NeurophysicsLab implicitly learns these symmetries through the physics-informed loss function. The energy conservation penalty encourages the latent space to encode time-translation invariant features. For orbital systems, the network additionally discovers angular momentum conservation — a manifestation of rotational symmetry.

Symbolic Regression with PySR

PySR uses multi-population evolutionary search to find mathematical expressions that fit data. Key advantages:

  • Pareto-optimal search: Explores the accuracy–complexity tradeoff
  • Julia backend: High-performance symbolic operations
  • Custom operators: Supports physics-relevant operations (sqrt, trig, etc.)

The pipeline extracts features from the trained neural network's latent space and state variables, then uses PySR to find the simplest equations that explain the learned conservation quantities.


Dashboard

The interactive dashboard (dashboard/index.html) provides four visualization tabs:

  1. Trajectory Viewer — Animated trajectories with true vs. predicted overlays
  2. Energy Monitor — Real-time energy tracking with violation detection
  3. Latent Space Explorer — PCA projection of the learned latent space
  4. Discovered Equations — LaTeX-rendered equations with accuracy metrics

Built with React 18, Recharts, KaTeX, and Tailwind CSS — all loaded via CDN. No build step required.


Running Experiments

Full Pipeline

# 1. Generate data for all systems
python data/simulate.py

# 2. Train models
for system in projectile spring orbit; do
    python train.py --system $system --epochs 100
done

# 3. Run symbolic regression
for system in projectile spring orbit; do
    python discover.py --system $system --niterations 40
done

# 4. Evaluate results
python evaluate.py --systems projectile spring orbit

# 5. Launch dashboard
python -m http.server 8080 --directory dashboard/

Ablation Studies

# Compare standard vs physics-aware model
python train.py --system projectile --energy-weight 0.0 --tag standard
python train.py --system projectile --energy-weight 0.1 --tag physics-aware

# Vary latent dimension
for dim in 4 8 16 32; do
    python train.py --system spring --latent-dim $dim --tag latent-$dim
done

References

  1. Greydanus, S., Dzamba, M., & Cranmer, M. (2019). Hamiltonian Neural Networks. NeurIPS 2019. arXiv:1906.01563

  2. Cranmer, M., Sanchez-Gonzalez, A., Battaglia, P., et al. (2020). Discovering Symbolic Models from Deep Learning with Inductive Biases. NeurIPS 2020. arXiv:2006.11287

  3. Cranmer, M. (2023). Interpretable Machine Learning for Science with PySR and SymbolicRegression.jl. arXiv:2305.01582

  4. Liu, Z., & Tegmark, M. (2021). Machine Learning Conservation Laws from Trajectories. Physical Review Letters 126, 180604. arXiv:2105.00266

  5. Noether, E. (1918). Invariante Variationsprobleme. Nachrichten von der Gesellschaft der Wissenschaften zu Göttingen, 235–257.

  6. Ha, D., & Schmidhuber, J. (2018). World Models. arXiv:1803.10122


License

This project is licensed under the MIT License — see the LICENSE file for details.


NeurophysicsLab
Discovering Physics with Neural Networks

Built with physics

About

NeurophysicsLab uses physics-informed neural networks and symbolic regression to rediscover conservation laws from raw motion data.

Topics

Resources

License

Stars

Watchers

Forks

Releases

No releases published

Packages

 
 
 

Contributors