DlWP HEALpix unifying original and coupled model config (#678)

* Merge dlwp_healpix and dlwp_healpix_coupled

Signed-off-by: root <iauyeung@nvidia.com>

---------

Signed-off-by: root <iauyeung@nvidia.com>
Co-authored-by: David Pruitt <dpruitt@nvidia.com>

examples/weather/dlwp_healpix/README.md

@@ -5,12 +5,37 @@ This example is an implementation of the
 model. The DLWP model can be used to predict the state of the atmosphere given a previous
 atmospheric state.  You can infer a 320-member ensemble set of six-week forecasts at 1.4°
 resolution within a couple of minutes, demonstrating the potential of AI in developing
-near real-time digital twins for weather prediction
+near real-time digital twins for weather prediction. This example also contains an
+implementation of the coupled Ocean-Atmosphere DLWP model.
 
 ## Problem overview
 
 The goal is to train an AI model that can emulate the state of the atmosphere and predict
 global weather over a certain time span. The Deep Learning Weather Prediction (DLWP) model
 uses deep CNNs for globally gridded weather prediction. DLWP CNNs directly map u(t) to
 its future state u(t+Δt) by learning from historical observations of the weather,
-with Δt set to 6 hr
+with Δt set to 6 hr. The Deep Learning Ocean Model (DLOM) that is designed to couple with
+deep learning weather prediction (DLWP) model. The DLOM forecasts sea surface
+temperature (SST). DLOMs use deep learning techniques as in DLWP models but are
+configured with different architectures and slower time stepping. DLOMs and DLWP models
+are trained to learn atmosphere-ocean coupling.
+
+## Getting Started
+
+To train the DLWP HEALPix model, run
+
+```bash
+python train.py
+```
+
+To train the coupled DLWP model, run
+
+```bash
+python train.py --config-name config_hpx32_coupled_dlwp
+```
+
+To train the coupled DLOM model, run
+
+```bash
+python train.py --config-name config_hpx32_coupled_dlom
+```

examples/weather/dlwp_healpix_coupled/configs/model/hpx_rec_unet.yaml → examples/weather/dlwp_healpix/config_hpx32_coupled_dlom.yaml

@@ -15,22 +15,28 @@
 # limitations under the License.
 
 defaults:
-  - modules/encoder@encoder: rec_unet_enc
-  - modules/decoder@decoder: rec_unet_dec
+  - _self_
+  - data: era5_hpx32_dlom_sst-z1000-ws_48H-dt
+  - model: coupled_hpx_unet_dlom
+  - trainer: dlom
 
-_target_: modulus.models.dlwp_healpix.HEALPixRecUNet
-_recursive_: false
-presteps: 1
-input_time_dim: ${data.input_time_dim}
-output_time_dim: ${data.output_time_dim}
-delta_time: ${data.time_step}
+experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S}
+output_dir: outputs/${experiment_name}
+# checkpoints names are in the form training-state-<name>.mdlus
+checkpoint_name: last
+load_weights_only: false
+seed: 0
 
-# Parameters automatically overridden in train code
-input_channels: 7
-output_channels: 7
-n_constants: 2
-decoder_input_channels: 1
+# Training specifications
+batch_size: 64
+learning_rate: 1e-4
+num_workers: 8
 
-# some perf parameters
-enable_nhwc: false
-enable_healpixpad: false
+# Distributed setup (multi GPU)
+port: 29450
+master_address: localhost
+
+hydra:
+  verbose: true
+  run:
+    dir: ${output_dir}

examples/weather/dlwp_healpix_coupled/configs/config.yaml → examples/weather/dlwp_healpix/configs/config_hpx32_coupled_dlwp.yaml

@@ -16,9 +16,9 @@
 
 defaults:
   - _self_
-  - data: era5_hpx64_7var_6h_24h
-  - model: hpx_rec_unet
-  - trainer: default
+  - data: era5_hpx32_8var-coupled_6h_24h
+  - model: coupled_hpx_rec_unet
+  - trainer: dlwp
 
 experiment_name: ${now:%Y-%m-%d}/${now:%H-%M-%S}
 output_dir: outputs/${experiment_name}

examples/weather/dlwp_healpix/configs/data/scaling/hpx32.yaml

@@ -49,4 +49,19 @@ sst:
 tp6:
   mean: 0.
   std: 1.
+z1000-24H:
+  mean: 936.7376098632812
+  std: 883.0859375
+ws10-24H:
+  mean: 6.15248966217041
+  std: 3.321399688720703
+ws10:
+  mean: 6.1497307
+  std: 3.583117 
+ws10-48H:
+  mean: 6.081215
+  std: 3.1224248
+adt:
+  mean: 0.40309871564035454
+  std: 0.6165622319307328
 

examples/weather/dlwp_healpix_coupled/configs/model/coupled_hpx_rec_unet_dlom.yaml → examples/weather/dlwp_healpix/configs/model/coupled_hpx_unet_dlom.yaml

@@ -15,10 +15,10 @@
 # limitations under the License.
 
 defaults:
-  - modules/encoder@encoder: rec_unet_enc
-  - modules/decoder@decoder: rec_unet_dec
+  - modules/encoder@encoder: unet_enc
+  - modules/decoder@decoder: unet_dec
 
-_target_: modulus.models.dlwp_healpix.HEALPixRecUNet
+_target_: modulus.models.dlwp_healpix.HEALPixUNet
 _recursive_: false
 presteps: 0
 input_time_dim: ${data.input_time_dim}

examples/weather/dlwp_healpix/configs/trainer/criterion/ocean_mse.yaml

@@ -16,5 +16,5 @@
 
 _target_: modulus.metrics.climate.healpix_loss.OceanMSE
 
-lsm_file: /invariants/hpx32_1950-2022_3h_sst-only.zarr
+lsm_file: /datasets/healpix/HPX32/hpx32_1950-2022_3h_sst_coupled.zarr