Merge main

Author: ryan-williams

CONTRIBUTING.md

@@ -170,9 +170,10 @@ You can run `uv run -m ocean_emulators.train --help` to see all the options avai
 To learn more about other datasets used during training, please see the _Data Engineering_ section below.
 
 To run a remote training job with Skypilot, use the following command:
+
 ```shell
 # export WANDB_API_KEY=<my-key>  # Get your key at https://wandb.ai/authorize
-# sky launch -c fomo-cluster train.sky.yaml  --env WANDB_API_KEY --env-file <my-vars>.env
+uv run sky launch train.sky.yaml  --env WANDB_API_KEY --env-file <my-vars>.env
 ```
 
 Please read the docstring in the `train.sky.yaml` for more information.
@@ -186,10 +187,26 @@ uv run scripts/clone_data.py $DATA_PATH --compact_variables
 uv run -m ocean_emulators.eval configs/eval_om4.yaml --ckpt_path path/to/checkpoint --experiment.data_root $DATA_PATH
 ```
 
+This produces a `predictions.zarr` file in the output directory with the rollout of the model.
+
 You can run `uv run -m ocean_emulators.eval --help` to see all the options available.
 
 To learn more about other datasets used during evaluation, please see the _Data Engineering_ section below.
 
+### Visualizing outputs from the model
+
+```bash
+uv run -m ocean_emulators.viz configs/viz_om4.yaml --data_root path/to/data --name my_experiment_viz --runs='[{"name": "my_experiment", "location": "path/to/eval/output/predictions.zarr"}]'
+```
+
+You can run `uv run -m ocean_emulators.viz --help` to see all the options available.
+
+After making changes to the visualization code, you can run the following command to compare old and new plots:
+
+```bash
+uv run -m ocean_emulators.utils.compare path/to/old/viz path/to/new/viz
+```
+
 ## Configuration
 
 ### Configuration files

configs/data/om4_sky_local.yaml

@@ -3,4 +3,5 @@
 data_location: OM4.zarr
 data_means_location: OM4_means.zarr
 data_stds_location: OM4_stds.zarr
-static_data_vars: []
+static_data_vars: []
+concurrent_compute: true

configs/data/public.yaml

@@ -0,0 +1,48 @@
+# yaml-language-server: $schema=schemas/EvalConfig.json
+
+debug: false
+save_zarr: true
+disk_mode: true
+inference_time:
+  start: "2014-10-10"
+  end: "2022-12-24"
+num_model_steps_forward: -1
+
+experiment:
+  name: local_om4_samudra
+  rand_seed: 15
+  base_output_dir: .LOCAL
+  wandb:
+    mode: online
+    project: ocean-emulators
+    entity: m2lines
+    group: samudra-eval-om4
+  network: Samudra
+  prognostic_vars_key: thermo_dynamic_all
+  boundary_vars_key: tau_hfds_hfds_anom
+
+data: !include data/public.yaml
+
+samudra:
+  checkpointing: all
+  ch_width: [157, 280, 380, 480, 520]
+  n_out: 154
+  dilation: [1, 2, 4, 8]
+  n_layers: [1, 1, 1, 1]
+  pred_residuals: false
+  last_kernel_size: 3
+  pad: "circular"
+
+  core_block:
+    block_type: "conv_next_block"
+    kernel_size: 3
+    activation: "capped_gelu"
+    upscale_factor: 2
+    norm: "batch"
+
+  down_sampling_block: "avg_pool"
+  up_sampling_block: "bilinear_upsample"
+
+  corrector:
+    non_negative_corrector_names: null
+    ocean_heat_corrector: false

configs/eval_om4_halfdeg.yaml

@@ -5,8 +5,8 @@ disk_mode: true
 pin_mem: true
 save_freq: 5
 epochs: 70
-batch_size: 2
-learning_rate: 0.0002
+batch_size: 3
+learning_rate: 0.0006
 scheduler: { type: cosine }
 loss: mse
 finetune: false
@@ -44,7 +44,7 @@ preemptible: false
 backend: auto
 
 experiment:
-  name: om4_samudra_halfdeg
+  name: om4_samudra_halfdeg_wide
   rand_seed: 15
   base_output_dir: .LOCAL
   wandb:
@@ -55,11 +55,12 @@ experiment:
   prognostic_vars_key: thermo_dynamic_all
   boundary_vars_key: tau_hfds_hfds_anom
 data:
-  !include data/om4_halfdeg_remote.yaml
+  !include data/public.yaml
 
 samudra:
-  ch_width: [157, 200, 250, 300, 400]
-  n_out: 77
+  checkpointing: all
+  ch_width: [157, 280, 380, 480, 520]
+  n_out: 154
   dilation: [1, 2, 4, 8]
   n_layers: [1, 1, 1, 1]
   pred_residuals: false
@@ -70,7 +71,7 @@ samudra:
     block_type: "conv_next_block"
     kernel_size: 3
     activation: "capped_gelu"
-    upscale_factor: 4
+    upscale_factor: 2
     norm: "batch"
 
   down_sampling_block: "avg_pool"

configs/train_om4_halfdeg.yaml

@@ -0,0 +1,17 @@
+# yaml-language-server: $schema=schemas/VizConfig.json
+
+base_output_dir: ".LOCAL"
+name: "viz_om4_halfdeg"
+dataset_name: "OM4"
+groundtruth_location: "OM4.zarr"
+basins_location:
+  type: s3
+  bucket: "emulators"
+  path: "jr7309/basins/basin_masks_original.zarr"
+  endpoint_url: "https://nyu1.osn.mghpcc.org"
+groundtruth_time_range:
+  start: "2014-10-20"
+  end: "2022-12-24"
+runs:
+  - name: "local_om4_samudra"
+    location: ".LOCAL/local_om4_samudra/predictions.zarr"

configs/viz_om4.yaml

@@ -0,0 +1,17 @@
+# yaml-language-server: $schema=schemas/VizConfig.json
+
+base_output_dir: ".LOCAL"
+name: "viz_om4_halfdeg"
+dataset_name: "OM4"
+groundtruth_location: "OM4.zarr"
+basins_location:
+  type: s3
+  bucket: "emulators"
+  path: "jr7309/basins/basin_masks_regridded.zarr"
+  endpoint_url: "https://nyu1.osn.mghpcc.org"
+groundtruth_time_range:
+  start: "2014-10-20"
+  end: "2022-12-24"
+runs:
+  - name: "local_om4_samudra"
+    location: ".LOCAL/local_om4_samudra/predictions.zarr"