Feat/online training

.gitignore

@@ -2,6 +2,8 @@ data/3dmeshes/*.pk
 data/ground_truths/**/*
 data/output/*.pt
 data/output/**/*
+data/train_datasets/*.pt
 *.egg-info
 tensorboard_logs
-**/__pycache__
+**/__pycache__
+wandb

README.md

@@ -17,62 +17,82 @@ conda env create -f environment.yml
 conda activate masconcube
 ```
 
-Or, if you want to install also development dependencies:
+## Project structure
 
-```bash
-conda env create -f environment_dev.yml
-conda activate masconcube
+```
+masconCube/
+├── data/                           # Data folder
+│   ├── 3dmeshes/                   # 3D meshes (to be downloaded from darioizzo/geodesyNets)
+│   ├── ground_truths/              # Ground-truth mascon models (to be generated with scripts/generate_ground_truth.py)
+│   ├── output/                     # Output folder for trained models and results
+|   ├── train_configs/              # Training configurations for MasconCube (yaml files)
+|   ├── train_datasets/             # Training datasets (to be generated with scripts/generate_train_datasets.py)
+|   ├── val_datasets/               # Validation datasets (to be generated with scripts/generate_val_datasets.py)
+|   └── test_datasets/              # Test datasets (to be generated with scripts/generate_val_datasets.py)
+├── mascon_cube/                    # Main package
+│   ├── data/                       # Data loading and processing
+|   ├── geodesynet/                 # GeodesyNets original implementation, mostly unmodified
+|   ├── pinn_gm/                    # PINN-GM III original implementation, mostly unmodified
+│   ├── contstants.py               
+│   ├── losses.py                   
+│   ├── metrics.py
+│   ├── models.py
+│   ├── training.py
+│   ├── trajectory.py
+│   ├── utils.py
+│   └── visualization.py
+├── notebooks/                      # Jupyter notebooks                 
+├── scripts/                        # Scripts for data generation, training, and evaluation
+└── environment.yml                 # Conda environment file
 ```
 
-Development dependencies include packages for linting and contributing to the project.
+## Data
 
-## Reprouducing the results from the paper
+Data is not stored in this repository to keep the size of the repository small. You can populate the `data` folder by following these steps:
 
  1. Download the 3D meshes from [darioizzo/geodesyNets/3dmeshes]([darioizzo/geodesyNets/3dmeshes](https://github.com/darioizzo/geodesyNets/tree/1edbb64d1e8e355e124a41eac27a14d7c5c5d881/3dmeshes)) and copy them inside the `data/3dmeshes` folder. For more information, see the [data README](data/README.md).
  2. Generate the ground-truth mascon models in the `data/ground_truths` folder by running the following script:
     ```bash
     python scripts/generate_ground_truth.py
     ```
-3. Generate the validation datasets in the `data/val_datasets` and `data/test_datasets` folders by running the following script:
-    ```bash
-    python scripts/generate_val_datasets.py
-    ```
-4. Train MasconCubes with the following command:
-    ```bash
-    python scripts/train_cubes_all.py [--gpus <gpu1> <gpu2> ...]
-    ```
-5. Train GeodesyNets with the following command:
+ 3. Generate the training datasets using in the paper in the `data/train_datasets` folders by running the following script:
     ```bash
-    python scripts/train_geodesynet_all.py [--gpus <gpu1> <gpu2> ...]
+    python scripts/generate_train_datasets.py random
     ```
-6. Train PINN-GM III with the following command:
+ 4. Generate the validation datasets in the `data/val_datasets` and `data/test_datasets` folders by running the following script:
     ```bash
-    python scripts/train_pinn_all.py [--gpus <gpu1> <gpu2> ...]
+    python scripts/generate_val_datasets.py
     ```
-7. Evaluate the models and produce plots using the provided notebooks.
-
-Note that steps 5 and 6 are required only to compare the results with previous state-of-the-art methods, and they might take a long time to run. You can skip them if you are only interested in MasconCube. If you want to run them, multiple GPUs are recommended, so that you can run them in parallel. The `--gpus` argument allows you to specify which GPUs to use for training.
 
-### Training on single asteroids
-
-If you want to run the training on single asteroids, you can use the scripts `scripts/train.py`, `scripts/train_geodesynet.py`, and `scripts/train_pinn.py`. For example, to train MasconCube on `eros_uniform`, you can run:
+## Training MasconCube and other models (legacy method)
 
+MasconCube, GeodesyNet and PINN-GM III can be trained using their respective training scripts in the `scripts` folder: `train_cube.py`, `train_geodesynet.py`, and `train_pinn.py`.
+The cli interface for all the training scripts is similar:
 ```bash
-python scripts/train.py eros_uniform
+python scripts/train_<model>.py <asteroid_name>
 ```
 
-MasconCube trainings also support TensorBoard logging (development dependencies required). You can run the following command to start TensorBoard:
-
+It is also possible to train a model for every asteroid in `data/ground_truths` on multiple GPUs using the scripts `train_cubes_all.py`, `train_geodesynet_all.py`, and `train_pinn_all.py`.
+Again, the cli interface is similar for all the scripts:
 ```bash
-tensorboard --logdir runs
+python scripts/train_<model>_all.py [--gpus <gpu1> <gpu2> ...]
 ```
 
-And then enable logging in the training script by passing the `--tensorboard` argument:
+Default settings should replicate the results in the paper. However some small implemention details could lead to slightly different results. If you want to replicate exactly the results of the paper checkout [v1.0.0](https://github.com/esa/masconCube/tree/v.1.0.0).
 
+## Training MasconCube (new method)
+
+The recommended way to train MasconCube models is with yaml configuration files stored in the `data/train_configs` folder.
+The folder alredy contains a configuration file called `paper.yaml` that you can use as a starting point.
+It use similar settings to the one used in the paper, but for an exact comparison follow the instructions in the previous section.
+
+To train a MasconCube model using a configuration file, run the following command:
 ```bash
-python scripts/train.py eros_uniform --tensorboard
+python scripts/run_config.py data/train_configs/<config_file>.yaml [--use-wandb]
 ```
 
+The `--use-wandb` flag is optional, and it enables logging with [Weights & Biases](https://wandb.ai/).
+
 ## License
 The code is released under the [Apache 2.0 license](https://github.com/esa/masconCube?tab=Apache-2.0-1-ov-file).
 

data/train_configs/paper.yaml

@@ -0,0 +1,42 @@
+# Training configuration for the paper experiments
+
+asteroids:
+  bennu: 
+    train: bennu_100000_spherical_0.0_1.0_42.pt
+    val: bennu_1000_spherical_0_2.pt
+  eros_2: 
+    train: eros_2_100000_spherical_0.0_1.0_42.pt
+    val: eros_2_1000_spherical_0_2.pt
+  eros_3: 
+    train: eros_3_100000_spherical_0.0_1.0_42.pt
+    val: eros_3_1000_spherical_0_2.pt
+  eros_uniform: 
+    train: eros_uniform_100000_spherical_0.0_1.0_42.pt
+    val: eros_uniform_1000_spherical_0_2.pt
+  itokawa: 
+    train: itokawa_100000_spherical_0.0_1.0_42.pt
+    val: itokawa_1000_spherical_0_2.pt
+  itokawa_cos: 
+    train: itokawa_cos_100000_spherical_0.0_1.0_42.pt
+    val: itokawa_cos_1000_spherical_0_2.pt
+  planetesimal: 
+    train: planetesimal_100000_spherical_0.0_1.0_42.pt
+    val: planetesimal_1000_spherical_0_2.pt
+  planetesimal_decentered: 
+    train: planetesimal_decentered_100000_spherical_0.0_1.0_42.pt
+    val: planetesimal_decentered_1000_spherical_0_2.pt
+  planetesimal_uniform: 
+    train: planetesimal_uniform_100000_spherical_0.0_1.0_42.pt
+    val: planetesimal_uniform_1000_spherical_0_2.pt
+
+train:
+  n_epochs: 1
+  batch_size: 1000
+  loss_fn: normalized_l1_loss
+  lr: 1.0E-5
+  cube_side: 100
+  differential: False
+  normalize: True
+  activation_function: linear
+  data_from_trajectory: False
+  batch_persistency: 10

data/train_configs/test.yaml

@@ -0,0 +1,49 @@
+# Training configuration for the paper experiments
+
+asteroids:
+  # bennu: 
+  #   train: bennu_100000_spherical_0.0_1.0_42.pt
+  #   val: bennu_1000_spherical_0_2.pt
+  # eros_2: 
+  #   train: eros_2_100000_spherical_0.0_1.0_42.pt
+  #   val: eros_2_1000_spherical_0_2.pt
+  eros_3: 
+    train: eros_3_traj_1_10000_1.0.pt
+    val: eros_3_1000_spherical_0_2.pt
+  # eros_uniform: 
+  #   train: eros_uniform_100000_spherical_0.0_1.0_42.pt
+  #   val: eros_uniform_1000_spherical_0_2.pt
+  itokawa: 
+    train: itokawa_traj_1_10000_1.0.pt
+    val: itokawa_1000_spherical_0_2.pt
+  # itokawa_cos: 
+  #   train: itokawa_cos_100000_spherical_0.0_1.0_42.pt
+  #   val: itokawa_cos_1000_spherical_0_2.pt
+  # planetesimal: 
+  #   train: planetesimal_100000_spherical_0.0_1.0_42.pt
+  #   val: planetesimal_1000_spherical_0_2.pt
+  # planetesimal_decentered: 
+  #   train: planetesimal_decentered_100000_spherical_0.0_1.0_42.pt
+  #   val: planetesimal_decentered_1000_spherical_0_2.pt
+  # planetesimal_uniform: 
+  #   train: planetesimal_uniform_100000_spherical_0.0_1.0_42.pt
+  #   val: planetesimal_uniform_1000_spherical_0_2.pt
+
+train:
+  n_epochs: 100
+  batch_size: 1000
+  loss_fn: normalized_l1_loss
+  lr: 1.0E-5
+  scheduler_factor: 0.8
+  scheduler_patience: 2
+  scheduler_min_lr: 1.0E-7
+  cube_side: 100
+  differential: False
+  normalize: True
+  activation_function: linear
+  data_from_trajectory: False
+  batch_persistency: 3
+  starting_training_data: 2000
+  add_data_every_n_epochs: 5
+  data_to_add: 1000
+  warmup_epochs: 10

environment.yml

@@ -10,16 +10,25 @@ dependencies:
   - numpy ~= 1.26.4  # numpy 2 is not compatible with cascade
   - cascade ~= 0.1.9
   - lazy_import ~= 0.2.2
-  - scipy ~= 0.15.2
+  - scipy ~= 1.16.3
   - seaborn ~= 0.13.2
   - torchquad ~= 0.4.1
   - pyvista ~= 0.45.2
   - pykep ~= 2.6.4
+  - pydantic ~= 2.12.3
+  - tqdm
+    # dev
+  - pre-commit ~= 4.1.0
+  - ipykernel ~= 6.29.5
+  - nbformat ~= 5.10.4
 
   - pip:
     - tetgen == 0.6.4
     - -e .
     - torch
     - torchvision
     - torchaudio
+    - sobol_seq
+    - pyvistaqt
+    - wandb
     - --extra-index-url https://download.pytorch.org/whl/cu118

mascon_cube/constants.py

@@ -4,10 +4,11 @@
 DATA_DIR = ROOT_DIR / "data"
 GROUND_TRUTH_DIR = DATA_DIR / "ground_truths"
 MESH_DIR = DATA_DIR / "3dmeshes"
+TRAIN_DATASETS_DIR = DATA_DIR / "train_datasets"
 VAL_DATASETS_DIR = DATA_DIR / "val_datasets"
 TEST_DATASETS_DIR = DATA_DIR / "test_datasets"
 OUTPUT_DIR = DATA_DIR / "output"
-TENSORBOARD_DIR = ROOT_DIR / "tensorboard_logs"
+CONFIGS_DIR = DATA_DIR / "train_configs"
 
 
 # -------------------------------------------------------------------------------------------------------------------- #