# greydanus/hamiltonian-nn architecture

> Architecture diagram of the GitHub repository greydanus/hamiltonian-nn, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-28). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/greydanus/hamiltonian-nn
- Repository: https://github.com/greydanus/hamiltonian-nn

## Overview

This repository is a research implementation of Hamiltonian Neural Networks (HNNs), with training and analysis workflows across six physical-system tasks. Task-specific data modules generate or acquire states and derivatives; training scripts fit either an HNN or baseline, and save model weights and metrics. The shared HNN computes structured vector fields, while the pixel task additionally learns an autoencoder representation. Analysis notebooks examine task results. The README documents these capabilities; some task-specific source was not sampled, so links below show only documented or directly supported wiring.

Main entry points: [`hnn.py`](https://github.com/greydanus/hamiltonian-nn/blob/master/hnn.py) (HNN dynamics), [`experiment-pixels/train.py`](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pixels/train.py) (Pixel trainer), [`experiment-spring/train.py`](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-spring/train.py) (Spring trainer), [`experiment-pend/train.py`](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pend/train.py) (Pendulum trainer)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_data["Task data"]
  node_spring_data["Spring data<br/>[data.py]"]
  node_pend_data["Pendulum data<br/>[data.py]"]
  node_real_data["Real pendulum data"]
  node_two_data["Two-body data<br/>[data.py]"]
  node_three_data["Three-body data<br/>[data.py]"]
  node_pixels_data["Pixel data<br/>[data.py]"]
end

subgraph group_training["Training"]
  node_spring_train["Spring trainer<br/>[train.py]"]
  node_pend_train["Pendulum trainer<br/>[train.py]"]
  node_real_train["Real trainer<br/>[train.py]"]
  node_two_train["Two-body trainer<br/>[train.py]"]
  node_three_train["Three-body trainer<br/>[train.py]"]
  node_pixels_train["Pixel trainer<br/>[train.py]"]
  node_weights["Saved weights"]
end

subgraph group_core["Model runtime"]
  node_hnn["HNN dynamics<br/>[hnn.py]"]
  node_nn["Neural networks<br/>[nn_models.py]"]
  node_utils["Integration utilities<br/>[utils.py]"]
  node_pixel_model["Pixel HNN<br/>[hnn.py]"]
end

subgraph group_analysis["Results analysis"]
  node_notebooks["Task analyses"]
end

node_researcher(("Researcher"))

node_researcher -->|"starts training"| node_spring_train
node_researcher -->|"starts training"| node_pend_train
node_researcher -->|"starts training"| node_real_train
node_researcher -->|"starts training"| node_two_train
node_researcher -->|"starts training"| node_three_train
node_researcher -->|"starts training"| node_pixels_train
node_spring_train -->|"loads data"| node_spring_data
node_pend_train -->|"loads data"| node_pend_data
node_real_train -->|"loads data"| node_real_data
node_two_train -->|"loads data"| node_two_data
node_three_train -->|"loads data"| node_three_data
node_pixels_train -->|"loads data"| node_pixels_data
node_spring_train -->|"trains model"| node_hnn
node_pend_train -->|"trains model"| node_hnn
node_real_train -->|"trains model"| node_hnn
node_two_train -->|"trains model"| node_hnn
node_three_train -->|"trains model"| node_hnn
node_hnn -->|"uses network"| node_nn
node_hnn -->|"integrates with"| node_utils
node_pixels_train -->|"trains model"| node_pixel_model
node_pixel_model -->|"delegates dynamics"| node_hnn
node_pixel_model -->|"uses network"| node_nn
node_spring_train -->|"saves weights"| node_weights
node_pend_train -->|"saves weights"| node_weights
node_real_train -->|"saves weights"| node_weights
node_two_train -->|"saves weights"| node_weights
node_three_train -->|"saves weights"| node_weights
node_pixels_train -->|"saves weights"| node_weights
node_researcher -->|"inspects results"| node_notebooks
node_pixels_data -.->|"samples Gym"| node_pixels_train

click node_spring_data "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-spring/data.py"
click node_pend_data "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pend/data.py"
click node_real_data "https://github.com/greydanus/hamiltonian-nn/tree/master/experiment-real"
click node_two_data "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-2body/data.py"
click node_three_data "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-3body/data.py"
click node_pixels_data "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pixels/data.py"
click node_spring_train "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-spring/train.py"
click node_pend_train "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pend/train.py"
click node_real_train "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-real/train.py"
click node_two_train "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-2body/train.py"
click node_three_train "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-3body/train.py"
click node_pixels_train "https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pixels/train.py"
click node_hnn "https://github.com/greydanus/hamiltonian-nn/blob/master/hnn.py"
click node_nn "https://github.com/greydanus/hamiltonian-nn/blob/master/nn_models.py"
click node_utils "https://github.com/greydanus/hamiltonian-nn/blob/master/utils.py"
click node_pixel_model "https://github.com/greydanus/hamiltonian-nn/blob/master/hnn.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_spring_data,node_pend_data,node_real_data,node_two_data,node_three_data,node_pixels_data toneBlue
class node_spring_train,node_pend_train,node_real_train,node_two_train,node_three_train,node_pixels_train,node_weights toneAmber
class node_hnn,node_nn,node_utils,node_pixel_model toneMint
class node_notebooks toneRose
class node_researcher toneIndigo
```

## Components

### Task data

- [Spring data](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-spring/data.py): `experiment-spring/data.py`
- [Pendulum data](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pend/data.py): `experiment-pend/data.py`
- [Real pendulum data](https://github.com/greydanus/hamiltonian-nn/tree/master/experiment-real): `experiment-real`
- [Two-body data](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-2body/data.py): `experiment-2body/data.py`
- [Three-body data](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-3body/data.py): `experiment-3body/data.py`
- [Pixel data](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pixels/data.py): `experiment-pixels/data.py`

### Training

- [Spring trainer](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-spring/train.py): `experiment-spring/train.py`
- [Pendulum trainer](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pend/train.py): `experiment-pend/train.py`
- [Real trainer](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-real/train.py): `experiment-real/train.py`
- [Two-body trainer](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-2body/train.py): `experiment-2body/train.py`
- [Three-body trainer](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-3body/train.py): `experiment-3body/train.py`
- [Pixel trainer](https://github.com/greydanus/hamiltonian-nn/blob/master/experiment-pixels/train.py): `experiment-pixels/train.py`
- **Saved weights**

### Model runtime

- [HNN dynamics](https://github.com/greydanus/hamiltonian-nn/blob/master/hnn.py): `hnn.py`
- [Neural networks](https://github.com/greydanus/hamiltonian-nn/blob/master/nn_models.py): `nn_models.py`
- [Integration utilities](https://github.com/greydanus/hamiltonian-nn/blob/master/utils.py): `utils.py`
- [Pixel HNN](https://github.com/greydanus/hamiltonian-nn/blob/master/hnn.py): `hnn.py`

### Results analysis

- **Task analyses**

### Other components

- **Researcher**

## Connections

- Researcher → Spring trainer: starts training
- Researcher → Pendulum trainer: starts training
- Researcher → Real trainer: starts training
- Researcher → Two-body trainer: starts training
- Researcher → Three-body trainer: starts training
- Researcher → Pixel trainer: starts training
- Spring trainer → Spring data: loads data
- Pendulum trainer → Pendulum data: loads data
- Real trainer → Real pendulum data: loads data
- Two-body trainer → Two-body data: loads data
- Three-body trainer → Three-body data: loads data
- Pixel trainer → Pixel data: loads data
- Spring trainer → HNN dynamics: trains model
- Pendulum trainer → HNN dynamics: trains model
- Real trainer → HNN dynamics: trains model
- Two-body trainer → HNN dynamics: trains model
- Three-body trainer → HNN dynamics: trains model
- HNN dynamics → Neural networks: uses network
- HNN dynamics → Integration utilities: integrates with
- Pixel trainer → Pixel HNN: trains model
- Pixel HNN → HNN dynamics: delegates dynamics
- Pixel HNN → Neural networks: uses network
- Spring trainer → Saved weights: saves weights
- Pendulum trainer → Saved weights: saves weights
- Real trainer → Saved weights: saves weights
- Two-body trainer → Saved weights: saves weights
- Three-body trainer → Saved weights: saves weights
- Pixel trainer → Saved weights: saves weights
- Researcher → Task analyses: inspects results
- Pixel data → Pixel trainer: samples Gym