CNN Fixes & IRL Resimulation

.github/workflows/pre-commit.yml

@@ -1,6 +1,6 @@
 on:
   push:
-    branches: main
+    branches: master
     paths:
       - "**.py"
   pull_request:

.github/workflows/pytest.yml

@@ -1,6 +1,6 @@
 on:
   push:
-    branches: main
+    branches: master
     paths:
       - "**.py"
   pull_request:

README.md

@@ -6,21 +6,57 @@
 <h1 align="center" style="margin-top: 0px;">BCNF: Ballistic Conditional Normalizing Flows</h1>
 <h2 align="center" style="margin-top: 0px;">Generative Neural Networks for the Sciences: Final Project</h2>
 
-<!-- <div align="center">
+<div align="center">
 
-[![pytest](https://github.com/psaegert/elqm-INLPT-WS2023/actions/workflows/pytest.yml/badge.svg)](https://github.com/psaegert/elqm-INLPT-WS2023/actions/workflows/pytest.yml)
-[![quality checks](https://github.com/psaegert/elqm-INLPT-WS2023/actions/workflows/pre-commit.yml/badge.svg)](https://github.com/psaegert/elqm-INLPT-WS2023/actions/workflows/pre-commit.yml)
+[![pytest](https://github.com/MrWhatZitToYaa/IGNNS-final-project/actions/workflows/pytest.yml/badge.svg)](https://github.com/MrWhatZitToYaa/IGNNS-final-project/actions/workflows/pytest.yml)
+[![quality checks](https://github.com/MrWhatZitToYaa/IGNNS-final-project/actions/workflows/pre-commit.yml/badge.svg)](https://github.com/MrWhatZitToYaa/IGNNS-final-project/actions/workflows/pre-commit.yml)
 
-</div> -->
+</div>
 
 
 # Introduction
+Conditional Normlizing Flows (CNF) [[1]](https://arxiv.org/abs/1605.08803) [[2]](https://arxiv.org/abs/1410.8516) have achieved major advancements at solving inverse problems [[3]](https://arxiv.org/abs/1808.04730) and performing Amortized Simulation Based Inference [[4]](https://arxiv.org/abs/2312.05440) [[3]](https://arxiv.org/abs/1808.04730) [[5]](https://arxiv.org/abs/2003.06281). Motivated by the idea of exploring physical systems with relevant practical applications, we apply various variations of the CNF architecture with LSTM, Transformer-based and Convolutional feature networks on a family of simulated and real physical systems with ballistic properties, and evaluate our method with regards to log likelihood, calibration and resimulation. Our experiments show that various simulation parameters can be identified from 3D trajectory data with good precision and fair accuracy, but barely at all from rendered or real-world video data.
+
+# Results
+
+| Model                         | Train      | Validation | Test    |
+|-------------------------------|------------|------------|---------|
+| `trajectory_FC_large`         | **-51.52** | **-51.51** | **-53.19** |
+| `trajectory_LSTM_large`      | -45.82     | -34.70     | -35.32  |
+| `trajectory_TRF_large`       | -41.21     | -28.87     | -29.12  |
+| `trajectory_FC_small`         | -45.71     | -50.69     | -50.80  |
+| `trajectory_LSTM_small`      | -42.77     | -45.83     | -46.41  |
+| `trajectory_TRF_small`       | -41.71     | -45.06     | -46.11  |
+
+**Table**: **Negative Log Likelihood of the Training, Validation, and Test Set in the Trajectory Setting.** The trajectory data is input into a fully connected (FC) Long Short-Term Memory (LSTM) or Transformer (TRF) feature network.
+
+---
+
+![Metrics during Training of the Trajectory Models](results/logging/train_trajectory_comparison.png)
+
+**Figure**: **Metrics during Training of the Trajectory Models.** We show the training and validation loss (left), statistics of the mean (middle) and standard deviation (right) of the validation code distribution in each dimension during training of our method with a small (solid lines) and large (colored dashed lines) CNF. Every model except the two using a fully connected feature network (red solid line, blue dashed line) exhibit overfitting behavior. Ideally, all means in z space converge to 0 and standard deviations to 1 with no variation across the dimensions (black dashed lines). All metrics are logged with the W&B logging provided by Nikita.
+
+---
+
+![Resimulated Trajectories with Parameters Sampled by `trajectory_LSTM_large`](results/trajectory_LSTM_large/figures/resimulation.png)
+
+**Figure**: **Resimulated Trajectories with Parameters Sampled by `trajectory_LSTM_large`**. We show the ground-truth (red curve) and 500 resimulated (feint black curves) for six arbitrarily chosen instances from our test set. The point of impact is calculated by the (first) intersection of the simulated trajectory with the ground ($z$ = 0, green area). Units in Meters.
+
+---
+
+![Predicted Point of Impact Distribution for `trajectory_LSTM_large`](results/trajectory_LSTM_large/figures/resimulation_impact_heatmap.png)
+
+**Figure**: **Predicted Point of Impact Distribution for `trajectory_LSTM_large`**. Many impact distributions (heatmap) match well with the ground truth impact (center point of the white circle). Some predictions significantly deviate from the ground truth impact (Trajectory #5, #10). Units in Meters.
+
+
+For more results, please see the [results](results) folder.
+
 
 # Requirements
 
 ## Hardware
-- 8GB RAM
-- GPU (optional)
+- 32GB RAM
+- CUDA-enabled GPU with 16GB VRAM (recommended)
 
 ## Software
 - Python 3.11
@@ -60,17 +96,69 @@ pip install -e .
 
 # Usage
 
-**CLI**
+For pretrained models and data, please refer to our [bcnf-models](https://huggingface.co/BCNF/bcnf-models) and [bcnf-data-public](https://huggingface.co/datasets/BCNF/bcnf-data-public) repositories on Hugging Face.
 
+## Load a model from a configuration file:L
+[calibration.ipynb](./notebooks/calibration.ipynb)
+```python
+import os
+import json
+import torch
+from bcnf import CondRealNVP_v2
+
+MODEL_NAME = "trajectory_LSTM_large"
+
+# Load the config
+with open(os.path.join(get_dir('models', 'bcnf-models', MODEL_NAME), 'config.json'), 'r') as f:
+    config = load_config(json.load(f)['config_path'])
+
+# Load the model
+model = CondRealNVP_v2.from_config(config).to(device)
+cnf.load_state_dict(torch.load(os.path.join(get_dir('models', 'bcnf-models', MODEL_NAME), "state_dict.pt")))
+cnf.eval()
+```
+
+## Train your own model:
+
+From the command line:
 ```sh
-bcnf demo --dummy_option "Hello World"
+bcnf train -c configs/runs/trajectory_LSTM_large.yaml
 ```
 
-**Python API**
+In your own code:\
+[train.ipynb](./notebooks/train.ipynb)
 ```python
-from bcnf import ...
-```
+from bcnf.utils import get_dir, load_config, sub_root_path
+from bcnf.train import Trainer
+from bcnf import CondRealNVP_v2
+
+# Specify a path template for the config file
+config_path_pattern = os.path.join("{{BCNF_ROOT}}", "configs", "runs", "my_config.yaml")
 
+# Find the config file in the local filesystem
+config_path = sub_root_path(config_path_pattern)
+
+# Load the config and create a model
+config = load_config(config_path, verify=False)
+model = CondRealNVP_v2.from_config(config).to(device)
+
+# Create a Trainer instance and load the data specified in the config
+trainer = Trainer(
+    config={k.lower(): v for k, v in config.to_dict().items()},
+    project_name="bcnf-test",  # Name of the Weights & Biases project
+    parameter_index_mapping=model.parameter_index_mapping,
+    verbose=True,
+)
+
+# Train
+model = trainer.train(model)
+
+# Save
+torch.save(model.state_dict(), os.path.join(get_dir('models', 'bcnf-models', MODEL_NAME, create=True), f"state_dict.pt"))
+
+with open(os.path.join(get_dir('models', 'bcnf-models', MODEL_NAME, create=True), 'config.json'), 'w') as f:
+    json.dump({'config_path': config_path_pattern}, f)
+```
 
 # Development
 
@@ -91,16 +179,14 @@ pytest tests --cov src
 ```
 
 # Citation
-If you use ELQM: Energy-Law Query-Master for your research, please cite it using the following
-
 ```bibtex
 @software{bcnf2024,
     author = {Christian Kleiber and Paul Saegert and Nikita Tatsch},
     title = {BCNF: Ballistic Conditional Normalizing Flows},
     month = mar,
     year = 2024,
     publisher = {GitHub},
-    version = {0.0.1},
+    version = {0.1.0},
     url = {https://github.com/MrWhatZitToYaa/IGNNS-final-project}
 }
 ```

configs/runs/trajectory_LSTM_FC_large.yaml

@@ -0,0 +1,86 @@
+global:
+  cnf_version: 2
+  parameter_selection: ['x0_x','x0_y','x0_z','v0_x','v0_y','v0_z','g','w_x','w_y','w_z','b','m','a_x','a_y','a_z','r','A','Cd','rho']
+  conditions :  # Consumed one by one in the ConcatenateContition layer of the feature network stack
+  - ["trajectories"]
+  - ["cam_radian", "cam_radius", "cam_angles", "cam_heights"]
+  dtype: "float32"
+
+data:
+  path: "{{BCNF_ROOT}}/data/bcnf-data/fixed_data_render_2s_15FPS/train"
+  config_file: "{{BCNF_ROOT}}/configs/data/config.yaml"
+  n_samples: 5000
+  output_type: "trajectories"
+  dt: 0.067
+  T: 2.0
+  ratio:
+    - 16
+    - 9
+  fov_horizontal: 70.0
+  cam1_pos: 0.0
+  num_cams: 2
+  break_on_impact: False
+  verbose: True
+  do_filter: False
+  device: "cuda:0"
+
+model:
+  kwargs:
+    size: 19
+    nested_sizes: [526, 526, 526, 526, 526]
+    n_conditions: 107
+    n_blocks: 26
+    dropout: 0.407
+    act_norm: True
+
+feature_networks:
+  - type: ConcatenateCondition
+    kwargs:
+      input_size: null
+      output_size: 3
+  - type: LSTM
+    kwargs:
+      input_size: 3
+      hidden_size: 16
+      output_size: 100
+      num_layers: 2
+      dropout: 0.5
+      bidirectional: True
+      pooling: "mean"
+  - type: ConcatenateCondition
+    kwargs:
+      input_size: 100
+      output_size: 107
+      dim: -1
+
+optimizer:
+  type: Adam
+  kwargs:
+    lr: 2e-4
+
+lr_scheduler:
+  type: ReduceLROnPlateau
+  kwargs:
+    mode: "min"
+    factor: 0.5
+    patience: 250
+    threshold_mode: "abs"
+    threshold: 1e-1
+
+training:
+  validation_split: 0.2
+  val_loss_window_size: 32
+  val_loss_patience: 500
+  val_loss_tolerance: 1e-1
+  val_loss_tolerance_mode: "abs"
+  learning_rate: 2e-4
+  random_state: 2024_03_25
+  batch_size: 256
+  num_workers: 0
+  pin_memory: False
+  n_epochs: 50_000
+  verbose: True
+  timeout: 10_000
+  wandb:
+    model_log_frequency: 10
+

configs/runs/trajectory_LSTM_large.yaml

@@ -40,7 +40,7 @@ feature_networks:
   - type: LSTM
     kwargs:
       input_size: 3
-      hidden_size: 212
+      hidden_size: 140
       output_size: 1360
       num_layers: 2
       dropout: 0.111

configs/runs/trajectory_LSTM_large_pretrained.yaml

@@ -0,0 +1,81 @@
+global:
+  parameter_selection: ['x0_x','x0_y','x0_z','v0_x','v0_y','v0_z','g','w_x','w_y','w_z','b','m','a_x','a_y','a_z','r','A','Cd','rho']
+  dtype: "float32"
+
+data:
+  path: "{{BCNF_ROOT}}/data/bcnf-data/fixed_data_render_2s_15FPS/train"
+  config_file: "{{BCNF_ROOT}}/configs/data/config.yaml"
+  n_samples: 5000
+  output_type: "trajectories"
+  dt: 0.067
+  T: 2.0
+  ratio:
+    - 16
+    - 9
+  fov_horizontal: 70.0
+  cam1_pos: 0.0
+  num_cams: 2
+  break_on_impact: False
+  verbose: True
+  do_filter: False
+  device: "cuda:0"
+
+model:
+  kwargs:
+    size: 19
+    nested_sizes: [526, 526, 526, 526, 526]
+    n_conditions: 1360
+    n_blocks: 26
+    dropout: 0.407
+    act_norm: True
+
+feature_network:
+  type: FullyConnected
+  kwargs:
+    sizes:
+      - 3
+      - 1000
+    dropout: 0.3
+
+time_series_network:
+  type: LSTM
+  kwargs:
+    input_size: 1000
+    hidden_size: 212
+    output_size: 1360
+    num_layers: 2
+    dropout: 0.111
+    bidirectional: True
+    pooling: "mean"
+
+optimizer:
+  type: Adam
+  kwargs:
+    lr: 2e-4
+
+lr_scheduler:
+  type: ReduceLROnPlateau
+  kwargs:
+    mode: "min"
+    factor: 0.5
+    patience: 250
+    threshold_mode: "abs"
+    threshold: 1e-1
+
+training:
+  validation_split: 0.2
+  val_loss_window_size: 32
+  val_loss_patience: 75
+  val_loss_tolerance: 1e-1
+  val_loss_tolerance_mode: "abs"
+  learning_rate: 2e-4
+  random_state: 2024_03_25
+  batch_size: 256
+  num_workers: 0
+  pin_memory: False
+  n_epochs: 50_000
+  verbose: True
+  timeout: 10_000
+  wandb:
+    model_log_frequency: 10
+

configs/runs/trajectory_LSTM_small.yaml

@@ -40,7 +40,7 @@ feature_networks:
   - type: LSTM
     kwargs:
       input_size: 3
-      hidden_size: 212
+      hidden_size: 12
       output_size: 80
       num_layers: 2
       dropout: 0.111

configs/runs/trajectory_TRF_small.yaml

@@ -40,10 +40,10 @@ feature_networks:
   - type: Transformer
     kwargs:
       input_size: 3
-      trf_size: 128
-      n_heads: 8
-      ff_size: 128
-      n_blocks: 8
+      trf_size: 16
+      n_heads: 4
+      ff_size: 8
+      n_blocks: 4
       output_size: 80
       dropout: 0.1