Corr diff - testing CI/CD - DO NOT SUBMIT

.gitignore

@@ -175,6 +175,7 @@ pyrightconfig.json
 *.torch
 plots/*
 *.npz
+outputs/*
 
 # conda
 .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

README.md

@@ -2,37 +2,20 @@
 
 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)
+[Setup clariden/santis](#setup-claridensantis)   
+[Regression training - clariden/santis](#run-regression-model-training-alps)  
+[Diffusion training - clariden/santis](#run-diffusion-model-training-alps)  
+[Inference - clariden/santis](#running-inference-on-alps)  
+[Installation - uenv/venv - deprecated](#installation-alps-uenvvenv---deprecated)
 
-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.
+## 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/hirad.sqsh`. All the jobs can be run using this environment without additional installations and setup.  
 
 ## Training
 
 ### Run regression model training (Alps)
 
-1. Script for running the training of regression model is in `src/hirad/train_regression.sh`. 
+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 ###
@@ -55,14 +38,9 @@ srun bash -c "
 ```
 hydra:
   run:
-    dir: your_path_to_save_training_output
+    dir: your_path_to_save_training_outputs
 ```
-- 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)
-```
-- 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 `training_era_cosmo_regression.yaml` and config files the main config is referencing (default values are working for debugging purposes).
 
 3. Submit the job with:
 ```bash
@@ -72,8 +50,7 @@ sbatch src/hirad/train_regression.sh
 ### Run diffusion model training (Alps)
 Before training diffusion model, checkpoint for regression model has to exist.
 
-1. Script for running the training of diffusion model is in `src/hirad/train_diffusion.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
@@ -82,12 +59,6 @@ Inside this script set the following:
 ```bash
 #SBATCH -A your_compute_group
 ```
-```bash
-srun bash -c "
-    . ./{your_env_name}/bin/activate
-    python src/hirad/training/train.py --config-name=training_era_cosmo_diffusion.yaml
-"
-```
 
 2. Set up the following config files in `src/hirad/conf`:
 
@@ -97,14 +68,12 @@ hydra:
   run:
     dir: your_path_to_save_training_output
 ```
-- In `training/era_cosmo_regression.yaml` set:
+- In `training/era_cosmo_diffusion.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
 ```
-- 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 `training_era_cosmo_diffusion.yaml` and config files the main config is referencing (default values are working for debugging purposes).
 
 3. Submit the job with:
 ```bash
@@ -155,13 +124,42 @@ Finally, from the dataset, subset of time steps can be chosen to do inference fo
 
 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`
+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.
 
-- In `dataset/era_cosmo.yaml` set the `dataset_path` if different from default.
+- In `dataset/era_cosmo_inference.yaml` set the `dataset_path` if different from default. Make sure that specified times or times_range is contained in dataset_path.
 
 3. Submit the job with:
 ```bash
 sbatch src/hirad/generate.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.
+
+## Installation (Alps uenv/venv) - deprecated
+
+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.

ci/cscs.yml

@@ -12,14 +12,18 @@ build_job:
   stage: build
   extends: .container-builder-cscs-gh200
   variables:
-    DOCKERFILE: ci/docker/Dockerfile
+    DOCKERFILE: ci/docker/Dockerfile.ci
 
-#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-clariden-gh200
+  image: $PERSIST_IMAGE_NAME
+  script:
+     - echo 'hello world'
+  #  - pip install -e . --no-dependencies
+  #  - python src/hirad/training/train.py --config-name=training_era_cosmo_regression_test.yaml
+  variables:
+    SLURM_JOB_NUM_NODES: 2
+    SLURM_NTASKS: 2
+    #SLURM_ACCOUNT: a161
+    #SBATCH_ACCOUNT: a161

ci/docker/Dockerfile.ci

@@ -0,0 +1,15 @@
+FROM nvcr.io/nvidia/physicsnemo/physicsnemo:25.06
+
+# setup
+RUN apt-get update && apt-get install python3-pip python3-venv -y
+RUN pip install --upgrade pip
+
+# Install the rest of dependencies.
+RUN pip install \
+    Cartopy==0.22.0 \
+    xskillscore
+
+
+
+
+

ci/docker/Dockerfile.corrdiff

@@ -0,0 +1,14 @@
+FROM nvcr.io/nvidia/physicsnemo/physicsnemo:25.06
+
+# setup
+RUN apt-get update && apt-get install python3-pip python3-venv -y
+RUN pip install --upgrade pip
+
+# Install the rest of dependencies.
+RUN pip install \
+    anemoi.datasets \
+    Cartopy==0.22.0 \
+    xskillscore \
+    scoringrules \
+    mlflow \
+    meteodata-lab 

ci/docker/Dockerfile.python

@@ -0,0 +1,29 @@
+# Following some suggestions in https://meteoswiss.atlassian.net/wiki/spaces/APN/pages/719684202/Clariden+Alps+environment+setup
+
+FROM ubuntu:22.04 as builder
+#FROM nvcr.io/nvidia/pytorch:25.01-py3
+
+
+# setup
+RUN apt-get update && apt-get install python3-pip python3-venv -y
+RUN pip install --upgrade \
+        pip
+        #ninja
+        #wheel
+        #packaging
+        #setuptools
+
+# install the rest of dependencies
+# TODO: Factor pydeps into a separate file(s)
+# TODO: Add versions for things
+RUN pip install \
+        anemoi-datasets \
+        cartopy \
+        matplotlib \
+        numpy \
+        pandas \
+        scipy \
+        torch
+
+
+

ci/edf/modulus_env.toml

@@ -0,0 +1,10 @@
+image = "/capstor/scratch/cscs/pstamenk/hirad.sqsh"
+
+mounts = ["/capstor", "/iopsstor", "/users"]
+
+# The initial directory in the container. 
+workdir = "${PWD}"
+
+[annotations]
+com.hooks.aws_ofi_nccl.enabled = "true"
+com.hooks.aws_ofi_nccl.variant = "cuda12"

interpolate.sh

@@ -3,4 +3,4 @@
 #SBATCH --partition=postproc
 #SBATCH --time=23:59:00
 
-python src/input_data/interpolate_basic.py src/input_data/era-all.yaml src/input_data/cosmo-all.yaml /store_new/mch/msopr/hirad-gen/basic-torch/era5-cosmo-1h-all-channels/
+python src/hirad/input_data/interpolate_basic.py src/hirad/input_data/era-all.yaml src/hirad/input_data/cosmo-all.yaml /store_new/mch/msopr/hirad-gen/basic-torch/era5-cosmo-1h-all-channels/

pyproject.toml

@@ -10,19 +10,10 @@ authors = [
     { name="Petar Stamenkovic", email="petar.stamenkovic@meteoswiss.ch" }
 ]
 readme = "README.md"
-requires-python = ">=3.12"
+#requires-python = ">=3.12"
 license = {file = "LICENSE"}
 
 dependencies = [
-    "cartopy>=0.24.1",
-    "cftime>=1.6.4",
-    "hydra-core>=1.3.2",
-    "matplotlib>=3.10.1",
-    "omegaconf>=2.3.0",
-    "tensorboard>=2.19.0",
-    "termcolor>=3.1.0",
-    "torchinfo>=1.8.0",
-    "treelib>=1.7.1"
 ]
 
 [tool.setuptools]

src/hirad/conf/compute_eval.yaml

@@ -0,0 +1,12 @@
+hydra:
+  job:
+    chdir: true
+    name: diffusion_era5_cosmo_7500000_test
+  run:
+    dir: ./outputs/generation/${hydra:job.name}
+
+# Get defaults
+defaults:
+  - _self_
+  # Dataset
+  - dataset/era_cosmo_inference

src/hirad/conf/dataset/era_cosmo.yaml

@@ -1,3 +1,9 @@
 type: era5_cosmo
-dataset_path: /iopsstor/scratch/cscs/pstamenk/datasets/basic-torch/trim_19_full
-validation_path: /iopsstor/scratch/cscs/pstamenk/datasets/basic-torch/trim_19_full/validation
+dataset_path: /iopsstor/scratch/cscs/mmcgloho/run-1_2/train/
+# dataset_path: /iopsstor/scratch/cscs/pstamenk/datasets/basic-torch/era5-cosmo-1h-all-channels
+validation_path: /iopsstor/scratch/cscs/mmcgloho/run-1_2/validation/
+input_channel_names: [2t, 10u, 10v, tcw, t_850, z_850, u_850, v_850, t_500, z_500, u_500, v_500, tp]
+output_channel_names: [2t, 10u, 10v, tp]
+static_channel_names: ['hsurf']
+transform_channels: ['tp-box_cox_025']
+n_month_hour_channels: 4

src/hirad/conf/dataset/era_cosmo_inference.yaml

@@ -0,0 +1,8 @@
+type: era5_cosmo
+# dataset_path: /iopsstor/scratch/cscs/pstamenk/datasets/basic-torch/era5-cosmo-1h-linear-interpolation/validation
+dataset_path: /iopsstor/scratch/cscs/mmcgloho/run-1_2/validation
+input_channel_names: [2t, 10u, 10v, tcw, t_850, z_850, u_850, v_850, t_500, z_500, u_500, v_500, tp]
+output_channel_names: [2t, 10u, 10v, tp]
+static_channel_names: ['hsurf']
+transform_channels: ['tp-box_cox_025']
+n_month_hour_channels: 4

src/hirad/conf/generate_era_cosmo.yaml

@@ -1,15 +1,15 @@
 hydra:
   job:
     chdir: true
-    name: generation_full
+    name: diffusion_era5_cosmo_7500000_test
   run:
-    dir: /iopsstor/scratch/cscs/pstamenk/outputs/${hydra:job.name}
+    dir: ./outputs/generation/${hydra:job.name}
 
 # Get defaults
 defaults:
   - _self_
   # Dataset
-  - dataset/era_cosmo
+  - dataset/era_cosmo_inference
 
   # Sampler
   - sampler/stochastic

src/hirad/conf/generate_era_cosmo_test.yaml

@@ -0,0 +1,20 @@
+hydra:
+  job:
+    chdir: true
+    name: generation_era5_cosmo_test
+  run:
+    dir: ./outputs/generation/${hydra:job.name}
+
+# Get defaults
+defaults:
+  - _self_
+  # Dataset
+  - dataset/era_cosmo_inference
+
+  # Sampler
+  - sampler/stochastic
+  #- sampler/deterministic
+
+  # Generation
+  - generation/era_cosmo_test
+    #- generation/patched_based