Skip to content

Add LLC CPU data canonicalization #670

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
fomo-bot merged 22 commits into from
Jul 24, 2026
Merged

Add LLC CPU data canonicalization #670

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
22 commits
Select commit Hold shift + click to select a range
7d07dd1
Add LLC CPU data canonicalization
oa-jder-bot May 7, 2026
e1affc6
Add LLC data loader tests
jder Jul 9, 2026
3cfe346
Handle extra LLC grid variables
jder Jul 9, 2026
38bec3e
Strengthen LLC root canonicalization tests
jder Jul 9, 2026
301a7b2
Make dataset time splits source-specific
jder Jul 15, 2026
8f18a92
Make inference times source-specific
jder Jul 15, 2026
ae0a4a3
Fix Empire AI smoke time overrides
jder Jul 15, 2026
d541437
Support indexed CLI config overrides
jder Jul 16, 2026
f2d95dd
Build data splits from source configs
jder Jul 16, 2026
9c8bcb0
Pr comments, todo
jder Jul 16, 2026
7f269ee
assert to avoid growing pr further
jder Jul 16, 2026
37cdab3
Make source time configs immutable
jder Jul 16, 2026
c6a6e88
Test LLC canonical output positively
jder Jul 16, 2026
40b4b26
Address LLC config review feedback
jder Jul 16, 2026
50466b3
comment
jder Jul 16, 2026
16a2638
Reject LLC static data vars
fomo-bot Jul 16, 2026
7649a17
Use LLC staggered masks
fomo-bot Jul 16, 2026
b6df79a
Format LLC mask helper
fomo-bot Jul 16, 2026
ee56e6c
Update src/samudra/config.py
jder Jul 22, 2026
7799e54
PR feedback
jder Jul 22, 2026
56d2d13
move
jder Jul 22, 2026
1c26183
Merge remote-tracking branch 'origin/main' into llc-stack-03-llc-cpu-…
jder Jul 23, 2026
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
27 changes: 27 additions & 0 deletions configs/data/llc.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,27 @@
# SPDX-FileCopyrightText: 2026 Ocean Emulator Authors
#
# SPDX-License-Identifier: Apache-2.0

# yaml-language-server: $schema=../schemas/DataConfig.json

sources:
- type: llc
face: 1
i_start: 0
i_end: 720
j_start: 0
j_end: 720
Comment thread
alxmrs marked this conversation as resolved.
train_time:
start: "2011-09-10T12:00:00Z"
end: "2012-09-01T12:00:00Z"
val_time:
start: "2012-09-01T12:00:00Z"
end: "2012-11-15T12:00:00Z"
data_location: LLC.zarr
data_means_location: LLC_means.zarr
data_stds_location: LLC_stds.zarr
static_data_vars: []
concurrent_compute: false
loading:
type: cpu
num_workers: 4
19 changes: 13 additions & 6 deletions configs/data/om4.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,13 +4,20 @@

# yaml-language-server: $schema=../schemas/DataConfig.json

dataset:
type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds_hfds_anom

sources:
- data_location: OM4.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds_hfds_anom
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
Comment thread
jder marked this conversation as resolved.
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "2014-10-10"
end: "2022-12-24"
data_location: OM4.zarr
data_means_location: OM4_means.zarr
data_stds_location: OM4_stds.zarr
static_data_vars: []
Expand Down
15 changes: 0 additions & 15 deletions configs/samudra_mini_om4/train.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -18,21 +18,6 @@ loss:
finetune: false
resume_ckpt_path: null
inference_epochs: []
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "1975-01-03"
end: "1976-01-03"
- start: "1985-01-03"
end: "1986-01-03"
- start: "1995-01-03"
end: "1996-01-03"
- start: "2004-10-05"
end: "2005-10-05"
data_stride: [1]
steps: [4]
step_transition: []
Expand Down
38 changes: 28 additions & 10 deletions configs/samudra_multi_om4/eval.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,9 +20,6 @@
debug: false
save_zarr: true
disk_mode: true
inference_time:
start: "2014-10-05"
end: "2015-10-05"
num_model_steps_forward: 25

experiment:
Expand All @@ -34,20 +31,41 @@ experiment:
project: default

data:
dataset:
type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
static_data_vars: []
concurrent_compute: true
sources:
- data_location: om4_halfdeg_v4/OM4.zarr
# First source is used for inference. Reorder to eval a different scale.
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
train_time: &train_time
Comment thread
alxmrs marked this conversation as resolved.
start: "1975-01-03"
end: "2013-10-05"
val_time: &val_time
start: "2013-10-05"
end: "2014-10-05"
inference_times: &inference_times
- start: "2014-10-05"
end: "2015-10-05"
data_location: om4_halfdeg_v4/OM4.zarr
data_means_location: om4_halfdeg_v4/OM4_means.zarr
data_stds_location: om4_halfdeg_v4/OM4_stds.zarr
- data_location: om4_onedeg_v3/OM4.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
train_time: *train_time
val_time: *val_time
inference_times: *inference_times
data_location: om4_onedeg_v3/OM4.zarr
data_means_location: om4_onedeg_v3/OM4_means.zarr
data_stds_location: om4_onedeg_v3/OM4_stds.zarr
- data_location: om4_quarterdeg_v2/OM4.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
train_time: *train_time
val_time: *val_time
inference_times: *inference_times
data_location: om4_quarterdeg_v2/OM4.zarr
data_means_location: om4_quarterdeg_v2/OM4_means.zarr
data_stds_location: om4_quarterdeg_v2/OM4_stds.zarr

Expand Down
49 changes: 27 additions & 22 deletions configs/samudra_multi_om4/train.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,21 +19,6 @@ loss:
finetune: false
resume_ckpt_path: null
inference_epochs: []
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "1975-01-03"
end: "1976-01-03" # "1985-01-03"
- start: "1985-01-03"
end: "1986-01-03" # "1995-01-03"
- start: "1995-01-03"
end: "1996-01-03" # "2005-01-03"
- start: "2004-10-05"
end: "2005-10-05" # "2014-10-05"
data_stride: [1]
steps: [4]
step_transition: []
Expand All @@ -49,23 +34,43 @@ experiment:
mode: disabled
project: default
data:
dataset:
type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
static_data_vars: []
concurrent_compute: true
loading:
type: cpu
num_workers: 4
sources:
- data_location: om4_quarterdeg_v2/OM4.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
train_time: &train_time
start: "1975-01-03"
end: "2013-10-05"
val_time: &val_time
start: "2013-10-05"
end: "2014-10-05"
inference_times: &inference_times
- start: "2014-10-05"
end: "2015-10-05"
data_location: om4_quarterdeg_v2/OM4.zarr
data_means_location: om4_quarterdeg_v2/OM4_means.zarr
data_stds_location: om4_quarterdeg_v2/OM4_stds.zarr
- data_location: om4_halfdeg_v4/OM4.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
train_time: *train_time
val_time: *val_time
inference_times: *inference_times
data_location: om4_halfdeg_v4/OM4.zarr
data_means_location: om4_halfdeg_v4/OM4_means.zarr
data_stds_location: om4_halfdeg_v4/OM4_stds.zarr
- data_location: om4_onedeg_v3/OM4.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_all
boundary_vars_key: tau_hfds
train_time: *train_time
val_time: *val_time
inference_times: *inference_times
data_location: om4_onedeg_v3/OM4.zarr
data_means_location: om4_onedeg_v3/OM4_means.zarr
data_stds_location: om4_onedeg_v3/OM4_stds.zarr
model: !include model.yaml
3 changes: 0 additions & 3 deletions configs/samudra_om4/eval.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,6 @@
debug: false
save_zarr: true
disk_mode: true
inference_time:
start: "2014-10-10"
end: "2022-12-24"
num_model_steps_forward: 25

experiment:
Expand Down
15 changes: 0 additions & 15 deletions configs/samudra_om4/train.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,21 +19,6 @@ loss:
finetune: false
resume_ckpt_path: null
inference_epochs: []
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "1975-01-03"
end: "1976-01-03" # "1985-01-03"
- start: "1985-01-03"
end: "1986-01-03" # "1995-01-03"
- start: "1995-01-03"
end: "1996-01-03" # "2005-01-03"
- start: "2004-10-05"
end: "2005-10-05" # "2014-10-05"
data_stride: [1]
steps: [4]
step_transition: []
Expand Down
3 changes: 0 additions & 3 deletions configs/samudra_om4_v1/eval.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,6 @@
debug: false
save_zarr: true
disk_mode: true
inference_time:
start: "2014-10-10"
end: "2022-12-24"
num_model_steps_forward: 25

experiment:
Expand Down
23 changes: 0 additions & 23 deletions configs/samudra_om4_v1/train.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -16,29 +16,6 @@ loss: mse
finetune: false
resume_ckpt_path: null
inference_epochs: []
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "1975-01-03"
end: "1980-01-03"
- start: "1980-01-03"
end: "1985-01-03"
- start: "1985-01-03"
end: "1990-01-03"
- start: "1990-01-03"
end: "1995-01-03"
- start: "1995-01-03"
end: "2000-01-03"
- start: "2000-01-03"
end: "2005-01-03"
- start: "2005-01-03"
end: "2010-01-03"
- start: "2010-01-03"
end: "2014-10-05"
data_stride: [1]
steps: [4]
step_transition: []
Expand Down
3 changes: 0 additions & 3 deletions configs/samudra_om4_v2/eval.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,6 @@
debug: false
save_zarr: true
disk_mode: true
inference_time:
start: "2014-10-10"
end: "2022-12-24"
num_model_steps_forward: 25

experiment:
Expand Down
15 changes: 0 additions & 15 deletions configs/samudra_om4_v2/train.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,21 +20,6 @@ loss:
finetune: false
resume_ckpt_path: null
inference_epochs: []
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "1975-01-03"
end: "1976-01-03" # "1985-01-03"
- start: "1985-01-03"
end: "1986-01-03" # "1995-01-03"
- start: "1995-01-03"
end: "1996-01-03" # "2005-01-03"
- start: "2004-10-05"
end: "2005-10-05" # "2014-10-05"
data_stride: [1]
steps: [4]
step_transition: []
Expand Down
3 changes: 0 additions & 3 deletions configs/samudra_om4_v2_highres/eval.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,9 +7,6 @@
debug: false
save_zarr: true
disk_mode: true
inference_time:
start: "2014-10-10"
end: "2022-12-24"
num_model_steps_forward: 25

experiment:
Expand Down
15 changes: 0 additions & 15 deletions configs/samudra_om4_v2_highres/train.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -20,21 +20,6 @@ loss:
finetune: false
resume_ckpt_path: null
inference_epochs: []
train_time:
start: "1975-01-03"
end: "2013-10-05"
val_time:
start: "2013-10-05"
end: "2014-10-05"
inference_times:
- start: "1975-01-03"
end: "1976-01-03" # "1985-01-03"
- start: "1985-01-03"
end: "1986-01-03" # "1995-01-03"
- start: "1995-01-03"
end: "1996-01-03" # "2005-01-03"
- start: "2004-10-05"
end: "2005-10-05" # "2014-10-05"
data_stride: [1]
steps: [4]
step_transition: []
Expand Down
27 changes: 13 additions & 14 deletions configs/test/train_default.yaml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -19,15 +19,6 @@ loss: mse
finetune: false
resume_ckpt_path: null
inference_epochs: [-1]
train_time:
start: "1975-08-15"
end: "1975-10-25"
val_time:
start: "1975-10-25"
end: "1975-11-25"
inference_times:
- start: "1975-11-25"
end: "1975-12-25"
data_stride: [1]
steps: [1]
step_transition: []
Expand All @@ -44,12 +35,20 @@ experiment:
group: samudra-train-om4
data_root: .data_cache/
data:
dataset:
type: om4
prognostic_vars_key: thermo_dynamic_5
boundary_vars_key: tau_hfds
sources:
- data_location: data.zarr
- type: om4
prognostic_vars_key: thermo_dynamic_5
boundary_vars_key: tau_hfds
train_time:
start: "1975-08-15"
end: "1975-10-25"
val_time:
start: "1975-10-25"
end: "1975-11-25"
inference_times:
- start: "1975-11-25"
end: "1975-12-25"
data_location: data.zarr
data_means_location: means.nc
data_stds_location: stds.nc
loading:
Expand Down
Loading
Loading