Corr diff - testing CI/CD - DO NOT SUBMIT

.gitignore

@@ -175,6 +175,8 @@ pyrightconfig.json
 *.torch
 plots/*
 *.npz
+outputs/*
+logs/*
 
 # conda
 .conda/*
@@ -183,7 +185,6 @@ plots/*
 temp.*
 
 # local script
-interpolate.sh
 cosmo-grid.npz
 *.out
 .conda/* 
@@ -194,7 +195,9 @@ temp.zarr.sync*
 src/hirad/eval/__pycache__/*
 interpolate_basic.log
 interpolated.torch
+mlruns/
+.secrets.env
 out
 core
-*.png
-*.nc
+*.nc
+*.err

README.md

@@ -2,121 +2,211 @@
 
 HiRAD-Gen is short for high-resolution atmospheric downscaling using generative models. This repository contains the code and configuration required to train and use the model.
 
-## Installation (Alps)
+[Showcase](#Showcase)  
+[Setup - clariden/santis](#setup-claridensantis)   
+[Inference - clariden/santis](#running-inference-on-alps)  
+[Regression training - clariden/santis](#run-regression-model-training-alps)  
+[Diffusion training - clariden/santis](#run-diffusion-model-training-alps)  
+
+## Showcase
+
+<table>
+  <tr>
+    <th width="5%"></th>
+    <th width="31%">Input ERA5</th>
+    <th width="31%">Prediction</th>
+    <th width="31%">Target REAL-CH1</th>
+  </tr>
+  <tr>
+    <td><b>2t</b></td>
+    <td><img src="docs/images/showcase/2t-input.png" width="100%"></td>
+    <td><img src="docs/images/showcase/2t-pred.png" width="100%"></td>
+    <td><img src="docs/images/showcase/2t-target.png" width="100%"></td>
+  </tr>
+  <tr>
+    <td><b>10u</b></td>
+    <td><img src="docs/images/showcase/10-input.png" width="100%"></td>
+    <td><img src="docs/images/showcase/10u-pred.png" width="100%"></td>
+    <td><img src="docs/images/showcase/10-target.png" width="100%"></td>
+  </tr>
+  <tr>
+    <td><b>10v</b></td>
+    <td><img src="docs/images/showcase/10v-input.png" width="100%"></td>
+    <td><img src="docs/images/showcase/10v-pred.png" width="100%"></td>
+    <td><img src="docs/images/showcase/10v-target.png" width="100%"></td>
+  </tr>
+  <tr>
+    <td><b>tp</b></td>
+    <td><img src="docs/images/showcase/tp-input.png" width="100%"></td>
+    <td><img src="docs/images/showcase/tp-pred1.png" width="100%"></td>
+    <td><img src="docs/images/showcase/tp-target.png" width="100%"></td>
+  </tr>
+</table>
+
+### Ensemble Total Preceipitatin 1h
+
+<table>
+  <tr>
+    <td width="50%"><img src="docs/images/showcase/tp-pred1.png" width="100%"></td>
+    <td width="50%"><img src="docs/images/showcase/tp-pred2.png" width="100%"></td>
+  </tr>
+  <tr>
+    <td><img src="docs/images/showcase/tp-pred3.png" width="100%"></td>
+    <td><img src="docs/images/showcase/tp-pred4.png" width="100%"></td>
+  </tr>
+</table>
+
+## Setup clariden/santis container environment
+Container environment setup needed to run training and inference experiments on clariden/santis is contained in this repository under `ci/edf/modulus_env.toml`. Image squash is on clariden/alps under `/capstor/scratch/cscs/pstamenk/corr_diff.sqsh`. All the jobs can be run using this environment without additional installations and setup.
 
-To set up the environment for **HiRAD-Gen** on Alps supercomputer, follow these steps:
-
-1. **Start the PyTorch user environment**:
-   ```bash
-   uenv start pytorch/v2.6.0:v1 --view=default
-   ```
-
-2. **Create a Python virtual environment** (replace `{env_name}` with your desired environment name):
-   ```bash
-   python -m venv ./{env_name}
-   ```
-
-3. **Activate the virtual environment**:
-   ```bash
-   source ./{env_name}/bin/activate
-   ```
-
-4. **Install project dependencies**:
-   ```bash
-   pip install -e .
-   ```
-
-This will set up the necessary environment to run HiRAD-Gen within the Alps infrastructure.
-
-## Training
+## Inference
 
-### Run regression model training (Alps)
+### Running inference on Alps
 
-1. Script for running the training of regression model is in `src/hirad/train_regression.sh`. 
+1. Script for running the inference is in `src/hirad/generate.sh`. 
 Inside this script set the following:
 ```bash
 ### OUTPUT ###
-#SBATCH --output=your_path_to_output_log
-#SBATCH --error=your_path_to_output_error
+#SBATCH --output=path_to_output_log
+#SBATCH --error=path_to_output_error
 ```
 ```bash
-#SBATCH -A your_compute_group
+#SBATCH -A compute_group
 ```
 ```bash
-srun bash -c "
-    . ./{your_env_name}/bin/activate
-    python src/hirad/training/train.py --config-name=training_era_cosmo_regression.yaml
+srun --mpi=pmix --network=disable_rdzv_get --environment=./ci/edf/modulus_env.toml bash -c "
+    pip install -e .
+    python src/hirad/inference/generate.py --config-name=main-config-file-in-src/hirad/conf.yaml
 "
 ```
 
 2. Set up the following config files in `src/hirad/conf`:
 
-- In `training_era_cosmo_regression.yaml` set:
+- In main config file (by default `generate_era_real.yaml`) set:
 ```
 hydra:
   run:
-    dir: your_path_to_save_training_output
+    dir: your_path_to_save_inference_output
 ```
-- In `training/era_cosmo_regression.yaml` set:
+- In generation config file (by default `generation/era_real.yaml`):
+Choose the inference mode:
 ```
-hp:
-    training_duration: number of samples to train for (set to 4 for debugging, 512 fits into 30 minutes on 1 gpu with total_batch_size: 4)
+inference_mode: all/regression/diffusion
+```
+by default `all` does both regression and diffusion. Depending on mode, regression and/or diffusion model pretrained weights should be provided:
 ```
-- In `dataset/era_cosmo.yaml` set the `dataset_path` if different from default.
+io:
+  res_ckpt_path: path_to_directory_containing_diffusion_training_model_checkpoints
+  reg_ckpt_path: path_to_directory_containing_regression_training_model_checkpoints
+```
+Finally, from the dataset, subset of time steps can be chosen to do inference for.
+
+One way is to list steps under `times:` in format `%Y%m%d-%H%M` for era5_cosmo dataset.
+
+The other way is to specify `times_range:` with three items: first time step (`%Y%m%d-%H%M`), last time step (`%Y%m%d-%H%M`), hour shift (int). Hour shift specifies distance in hours between closest time steps for specific dataset.
 
 3. Submit the job with:
 ```bash
-sbatch src/hirad/train_regression.sh
+sbatch src/hirad/generate.sh
 ```
 
-### Run diffusion model training (Alps)
-Before training diffusion model, checkpoint for regression model has to exist.
+### Visualizing results
 
-1. Script for running the training of diffusion model is in `src/hirad/train_diffusion.sh`. 
+After generation is finished, visualization of results can be done using `src/hirad/snapshots.sh`. Set:
+```bash
+### OUTPUT ###
+#SBATCH --output=path_to_output_log
+```
+```bash
+### ENVIRONMENT ####
+#SBATCH -A compute_group
+```
+```bash
+srun --mpi=pmix --network=disable_rdzv_get --environment=./ci/edf/modulus_env.toml bash -c "
+    pip install -e .
+    python src/hirad/eval/snapshots.py --config-name=src/hirad/conf/config-file-in-src/hirad/conf.yaml
+"
+```
+In config file (by default `eval_real.yaml`) set:
+```bash
+# Path to the inference output directory
+inference_output_dir: '/path/to/generated/results/directory'
+results_dir_name: 'name_of_directory_to_save_output_plots'
+```
+If you want to generate plots for subset of times from inference set (follow same convection as in generate config):
+```
+times: list of times to visualize
+times_range: [start time, end time, time step] to visualize
+```
+
+Other setting can be changed according to output grid.
+
+Submit the job with:
+```bash
+sbatch src/hirad/snapshots.sh
+```
+
+### Evaluation of generated data
+
+Evaluation of generated samples can be done using `src/hirad/eval_precip.sh` and `src/hirad/eval_wind.sh`. Set:
+```bash
+### OUTPUT ###
+#SBATCH --output=path_to_output_log
+
+### ENVIRONMENT ####
+#SBATCH -A compute_group
+
+### CONFIG ###
+CONFIG_NAME="src/hirad/conf/config_file.yaml"
+```
+Default config file is the same as for visualization `eval_real.yaml`, and requires to set the same fileds. In both `eval_precip.sh` and `eval_wind.sh` there are several python scripts called. They are all commented out by default. Comment out the ones you want to run.
+
+Submit jobs with:
+```bash
+sbatch src/hirad/eval_precip.sh
+sbatch src/hirad/eval_wind.sh
+```
+
+## Training
+
+### Run regression model training (Alps)
+
+1. Script for running the training of regression model is in `src/hirad/train_regression.sh`. Here, you can change the sbatch settings.
 Inside this script set the following:
 ```bash
 ### OUTPUT ###
-#SBATCH --output=your_path_to_output_log
-#SBATCH --error=your_path_to_output_error
+#SBATCH --output=path_to_output_log
+#SBATCH --error=path_to_output_error
 ```
 ```bash
-#SBATCH -A your_compute_group
+#SBATCH -A compute_group
 ```
 ```bash
-srun bash -c "
-    . ./{your_env_name}/bin/activate
-    python src/hirad/training/train.py --config-name=training_era_cosmo_diffusion.yaml
+srun --mpi=pmix --network=disable_rdzv_get --environment=./ci/edf/modulus_env.toml bash -c "
+    pip install -e .
+    python src/hirad/training/train.py --config-name=main-config-file-in-src/hirad/conf.yaml
 "
 ```
 
 2. Set up the following config files in `src/hirad/conf`:
 
-- In `training_era_cosmo_diffusion.yaml` set:
+- In main config file (by default `training_era_real_regression.yaml`) set:
 ```
 hydra:
   run:
-    dir: your_path_to_save_training_output
-```
-- In `training/era_cosmo_regression.yaml` set:
-```
-hp:
-    training_duration: number of samples to train for (set to 4 for debugging, 512 fits into 30 minutes on 1 gpu with total_batch_size: 4)
-io:
-    regression_checkpoint_path: path_to_directory_containing_regression_training_model_checkpoints
+    dir: your_path_to_save_training_outputs
 ```
-- In `dataset/era_cosmo.yaml` set the `dataset_path` if different from default.
+- All other parameters for training regression can be changed in the main config file and config files the main config is referencing (default values are working for debugging purposes).
 
 3. Submit the job with:
 ```bash
-sbatch src/hirad/train_diffusion.sh
+sbatch src/hirad/train_regression.sh
 ```
 
-## Inference
-
-### Running inference on Alps
+### Run diffusion model training (Alps)
+Before training diffusion model, checkpoint for regression model has to exist.
 
-1. Script for running the inference is in `src/hirad/generate.sh`. 
-Inside this script set the following:
+1. Script for running the training of diffusion model is in `src/hirad/train_diffusion.sh`. Here, you can change the sbatch settings. Inside this script set the following:
 ```bash
 ### OUTPUT ###
 #SBATCH --output=your_path_to_output_log
@@ -126,42 +216,33 @@ Inside this script set the following:
 #SBATCH -A your_compute_group
 ```
 ```bash
-srun bash -c "
-    . ./{your_env_name}/bin/activate
-    python src/hirad/inference/generate.py --config-name=generate_era_cosmo.yaml
+srun --mpi=pmix --network=disable_rdzv_get --environment=./ci/edf/modulus_env.toml bash -c "
+    pip install -e .
+    python src/hirad/training/train.py --config-name=main-config-file-in-src/hirad/conf.yaml
 "
 ```
 
 2. Set up the following config files in `src/hirad/conf`:
 
-- In `generate_era_cosmo.yaml` set:
+- In main config file (by default `training_era_real_diffusion_patched.yaml`) set:
 ```
 hydra:
   run:
-    dir: your_path_to_save_inference_output
-```
-- In `generation/era_cosmo.yaml`:
-Choose the inference mode:
-```
-inference_mode: all/regression/diffusion
+    dir: your_path_to_save_training_output
 ```
-by default `all` does both regression and diffusion. Depending on mode, regression and/or diffusion model pretrained weights should be provided:
+- In training config file (by default `training/era_real_diffusion_patched.yaml`) set:
 ```
 io:
-  res_ckpt_path: path_to_directory_containing_diffusion_training_model_checkpoints
-  reg_ckpt_path: path_to_directory_containing_regression_training_model_checkpoints
+    regression_checkpoint_path: path_to_directory_containing_regression_training_model_checkpoints
 ```
-Finally, from the dataset, subset of time steps can be chosen to do inference for.
-
-One way is to list steps under `times:` in format `%Y%m%d-%H%M` for era5_cosmo dataset.
-
-The other way is to specify `times_range:` with three items: first time step (`%Y%m%d-%H%M`), last time step (`%Y%m%d-%H%M`), hour shift (int). Hour shift specifies distance in hours between closest time steps for specific dataset (6 for era_cosmo).
-
-By default, inference is done for one time step `20160101-0000`
-
-- In `dataset/era_cosmo.yaml` set the `dataset_path` if different from default.
+- All other parameters for training regression can be changed in the main config file and config files the main config is referencing (default values are working for debugging purposes).
 
 3. Submit the job with:
 ```bash
-sbatch src/hirad/generate.sh
-```
+sbatch src/hirad/train_diffusion.sh
+```
+
+## MLflow logging
+
+During training MLflow can be used to log metrics.
+Logging config files for regression and diffusion are located in `src/hirad/conf/logging/`. Set `method` to `mlflow` and specify `uri` if you want to log on remote server, otherwise run will be logged locally in output directory. Other options can also be modified here.

ci/cscs.yml

@@ -12,14 +12,20 @@ build_job:
   stage: build
   extends: .container-builder-cscs-gh200
   variables:
-    DOCKERFILE: ci/docker/Dockerfile
+    DOCKERFILE: ci/docker/Dockerfile.ci
+    KUBERNETES_MEMORY_REQUEST: '64Gi'
+    KUBERNETES_MEMORY_LIMIT: '64Gi'
 
-#test_job:
-#  stage: test
-#  extends: .container-runner-clariden-gh200
-#  image: $PERSIST_IMAGE_NAME
-#  script:
-#    - /opt/helloworld/bin/hello
-#  variables:
-#    SLURM_JOB_NUM_NODES: 2
-#    SLURM_NTASKS: 2
+test_job:
+  stage: test
+  extends: .container-runner-santis-gh200
+  image: $PERSIST_IMAGE_NAME
+  script:
+    - pytest /opt/hirad-gen/tests -v
+  variables:
+    USE_MPI: NO
+    SLURM_MPI_TYPE: pmix
+    SLURM_NETWORK: disable_rdzv_get
+    PMIX_MCA_psec: native
+    SLURM_JOB_NUM_NODES: 1
+    SLURM_NTASKS: 1