diff --git a/src/samudra/config.py b/src/samudra/config.py index 598d1469f..a06c6405d 100644 --- a/src/samudra/config.py +++ b/src/samudra/config.py @@ -458,7 +458,7 @@ def build( BlockType = Literal["conv_next_block", "conv_block"] ActivationType = Literal["relu", "gelu", "capped_gelu"] -NormType = Literal["batch", "instance", "layer"] +NormType = Literal["batch", "instance", "group", "nonorm", "layer"] class BlockConfig(BaseConfig): @@ -467,6 +467,7 @@ class BlockConfig(BaseConfig): activation: ActivationType = "capped_gelu" upscale_factor: int = 4 norm: NormType = "batch" + group_norm_groups: int = Field(default=32, ge=1) pointwise_linear: bool = False def build(self) -> CoreBlockBuilder: @@ -511,6 +512,7 @@ def create_block( kernel_size=self.kernel_size, upscale_factor=self.upscale_factor, norm=self.norm, + group_norm_groups=self.group_norm_groups, activation=activation, pointwise_linear=self.pointwise_linear, ) @@ -1227,6 +1229,7 @@ class TrainConfig(TopLevelConfig): # Data parameters at root level data_stride: list[int] = [1] + temporal_stride: int = Field(default=1, ge=1) steps: list[int] = [4] step_transition: list[int] = [] inference_epochs: list[int] = [-1] diff --git a/src/samudra/configs/data/llc.yaml b/src/samudra/configs/data/llc.yaml index 62bb5ad0d..433b7197f 100644 --- a/src/samudra/configs/data/llc.yaml +++ b/src/samudra/configs/data/llc.yaml @@ -1,4 +1,4 @@ -# SPDX-FileCopyrightText: 2026 Ocean Emulator Authors +# SPDX-FileCopyrightText: 2026 Samudra Authors # # SPDX-License-Identifier: Apache-2.0 @@ -6,6 +6,8 @@ sources: - type: llc + prognostic_vars_key: single_1 + boundary_vars_key: single_1 face: 1 i_start: 0 i_end: 720 @@ -17,6 +19,9 @@ sources: val_time: start: "2012-09-01T12:00:00Z" end: "2012-11-15T12:00:00Z" + inference_times: + - start: "2012-10-16T12:00:00Z" + end: "2012-11-14T12:00:00Z" data_location: LLC.zarr data_means_location: LLC_means.zarr data_stds_location: LLC_stds.zarr diff --git a/src/samudra/configs/samudra_llc/eval.yaml b/src/samudra/configs/samudra_llc/eval.yaml new file mode 100644 index 000000000..a79fa1617 --- /dev/null +++ b/src/samudra/configs/samudra_llc/eval.yaml @@ -0,0 +1,20 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/EvalConfig.json + +debug: false +save_zarr: false +disk_mode: true +num_model_steps_forward: 25 + +experiment: + name: samudra_llc_eval + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/llc.yaml +model: !include model.yaml diff --git a/src/samudra/configs/samudra_llc/model.yaml b/src/samudra/configs/samudra_llc/model.yaml new file mode 100644 index 000000000..5ee2c5a81 --- /dev/null +++ b/src/samudra/configs/samudra_llc/model.yaml @@ -0,0 +1,27 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/SamudraConfig.json + +checkpointing: all +pred_residuals: false +last_kernel_size: 3 +pad: "constant" +use_bfloat16: false + +unet: + ch_width: [256, 384, 512, 512] + dilation: [1, 2, 4, 8] + n_layers: [1, 1, 1, 1] + + core_block: + block_type: "conv_next_block" + kernel_size: 3 + activation: "capped_gelu" + upscale_factor: 2 + norm: "group" + group_norm_groups: 32 + + down_sampling_block: "avg_pool" + up_sampling_block: "bilinear_upsample" diff --git a/src/samudra/configs/samudra_llc/train.yaml b/src/samudra/configs/samudra_llc/train.yaml new file mode 100644 index 000000000..0e6f6900a --- /dev/null +++ b/src/samudra/configs/samudra_llc/train.yaml @@ -0,0 +1,38 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +# yaml-language-server: $schema=../schemas/TrainConfig.json + +debug: false +disk_mode: true +pin_mem: true +save_freq: 5 +epochs: 70 +batch_size: 1 +learning_rate: 0.0006 +gradient_accumulation_steps: 4 +scheduler: { type: cosine } +loss: + type: dynamic + metric: mse +finetune: false +resume_ckpt_path: null +inference_epochs: [] +data_stride: [1] +temporal_stride: 24 +steps: [1] +step_transition: [] +preemptible: false + +backend: auto + +experiment: + name: samudra_llc + rand_seed: 15 + base_output_dir: .LOCAL + wandb: + mode: disabled + project: default +data: !include ../data/llc.yaml +model: !include model.yaml diff --git a/src/samudra/datasets.py b/src/samudra/datasets.py index 5955ca41e..a29a7a456 100644 --- a/src/samudra/datasets.py +++ b/src/samudra/datasets.py @@ -467,6 +467,7 @@ def __init__( masked_fill_value: float, stride: int = 1, concurrent_compute_: bool = False, + temporal_stride: int = 1, ): super().__init__() self.id = f"{self.__class__.__name__}_{str(id(self))}" @@ -474,6 +475,9 @@ def __init__( self.hist: int = hist self.steps: int = steps self.stride: int = stride + if temporal_stride < 1: + raise ValueError("temporal_stride must be >= 1") + self.temporal_stride: int = temporal_stride self.normalize_before_mask: bool = normalize_before_mask self.masked_fill_value: float = masked_fill_value self._concurrent_compute = concurrent_compute_ @@ -512,11 +516,15 @@ def __init__( output_resolution_cpu=self.prognostic_src.resolution, ) - self.size: int = ( + base_size = ( time_.size - self.steps * (self.hist + 1) * self.stride - self.hist * self.stride ) + self.size: int = max( + 0, + (base_size + self.temporal_stride - 1) // self.temporal_stride, + ) def __len__(self) -> int: return self.size @@ -644,7 +652,7 @@ def _get_x_index(self, idx: int, step: int) -> xr.DataArray: if idx >= len(self): raise IndexError("Index out of range") - window_index = idx + step * (self.hist + 1) * self.stride + window_index = idx * self.temporal_stride + step * (self.hist + 1) * self.stride return self.rolling_indices.isel(window=window_index, drop=True) diff --git a/src/samudra/models/modules/blocks.py b/src/samudra/models/modules/blocks.py index 8c9435ada..0a7fa2d9f 100644 --- a/src/samudra/models/modules/blocks.py +++ b/src/samudra/models/modules/blocks.py @@ -280,6 +280,7 @@ def __init__( pad="circular", upscale_factor: int = 4, norm="batch", + group_norm_groups: int = 32, checkpoint_simple: bool = False, pointwise_linear: bool = False, ): @@ -302,16 +303,13 @@ def __init__( dilation=dilation, ) ) - # BatchNorm - if norm == "batch": - convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor)) - # Instance Norm - elif norm == "instance": - convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor)) - elif norm == "nonorm": - pass - else: - raise NotImplementedError + norm_layer = self._build_norm_layer( + norm=norm, + channels=int(in_channels * upscale_factor), + group_norm_groups=group_norm_groups, + ) + if norm_layer is not None: + convblock.append(norm_layer) if activation is not None: convblock.append(activation()) convblock.append( @@ -322,16 +320,13 @@ def __init__( dilation=dilation, ) ) - # BatchNorm - if norm == "batch": - convblock.append(torch.nn.BatchNorm2d(in_channels * upscale_factor)) - # Instance Norm - elif norm == "instance": - convblock.append(torch.nn.InstanceNorm2d(in_channels * upscale_factor)) - elif norm == "nonorm": - pass - else: - raise NotImplementedError + norm_layer = self._build_norm_layer( + norm=norm, + channels=int(in_channels * upscale_factor), + group_norm_groups=group_norm_groups, + ) + if norm_layer is not None: + convblock.append(norm_layer) if activation is not None: convblock.append(activation()) # Linear postprocessing @@ -343,6 +338,29 @@ def __init__( self.convblock = torch.nn.Sequential(*convblock) self.checkpoint_simple = checkpoint_simple + @staticmethod + def _build_norm_layer( + norm: str, + channels: int, + group_norm_groups: int, + ) -> torch.nn.Module | None: + if norm == "batch": + return torch.nn.BatchNorm2d(channels) + if norm == "instance": + return torch.nn.InstanceNorm2d(channels) + if norm == "group": + if group_norm_groups < 1: + raise ValueError("group_norm_groups must be >= 1") + num_groups = min(group_norm_groups, channels) + while channels % num_groups != 0: + num_groups -= 1 + return torch.nn.GroupNorm(num_groups=num_groups, num_channels=channels) + if norm == "layer": + return torch.nn.GroupNorm(num_groups=1, num_channels=channels) + if norm == "nonorm": + return None + raise NotImplementedError(f"Unsupported normalization mode {norm!r}") + def forward(self, x: torch.Tensor) -> torch.Tensor: # return self.skip_module(x) + self.convblock(x) skip = self.skip_module(x) diff --git a/src/samudra/train.py b/src/samudra/train.py index f43c4a4d4..4e9fc4a80 100644 --- a/src/samudra/train.py +++ b/src/samudra/train.py @@ -323,6 +323,7 @@ def __init__(self, cfg: TrainConfig) -> None: self.search_run = cfg.experiment.search self.debug = cfg.debug self.data_stride: list[int] = cfg.data_stride + self.temporal_stride: int = cfg.temporal_stride self.batch_size: int = cfg.batch_size self.gradient_accumulation_steps: int = cfg.gradient_accumulation_steps self.num_workers: int = data_num_workers @@ -1090,6 +1091,7 @@ def init_data_loaders(self, cur_step: int) -> None: masked_fill_value=self.normalize_fill_value, stride=stride, concurrent_compute_=self.concurrent_compute, + temporal_stride=self.temporal_stride, ) for stride in self.data_stride for src in self.data_container.train_sources @@ -1109,6 +1111,7 @@ def init_data_loaders(self, cur_step: int) -> None: masked_fill_value=self.normalize_fill_value, stride=stride, concurrent_compute_=self.concurrent_compute, + temporal_stride=self.temporal_stride, ) for stride in self.data_stride ] diff --git a/tests/test_blocks.py b/tests/test_blocks.py index ce69f6d5d..bced532b1 100644 --- a/tests/test_blocks.py +++ b/tests/test_blocks.py @@ -92,3 +92,70 @@ def test_gradient_flows_through_trunk_when_skip_dropped(self): torch.testing.assert_close(trunk.grad, torch.ones_like(trunk)) # Skip gradient should be all zeros (dropped) torch.testing.assert_close(skip.grad, torch.zeros_like(skip)) + + +def test_convnext_block_group_norm_uses_divisible_group_count(): + block = ConvNeXtBlock( + in_channels=10, + out_channels=10, + kernel_size=3, + dilation=1, + n_layers=1, + norm="group", + group_norm_groups=6, + ) + + norm_layers = [ + layer for layer in block.convblock if isinstance(layer, nn.GroupNorm) + ] + + assert len(norm_layers) == 2 + assert all(layer.num_channels == 40 for layer in norm_layers) + assert all(layer.num_groups == 5 for layer in norm_layers) + + +def test_convnext_block_layer_norm_uses_single_group(): + block = ConvNeXtBlock( + in_channels=8, + out_channels=8, + kernel_size=3, + dilation=1, + n_layers=1, + norm="layer", + ) + + norm_layers = [ + layer for layer in block.convblock if isinstance(layer, nn.GroupNorm) + ] + + assert len(norm_layers) == 2 + assert all(layer.num_groups == 1 for layer in norm_layers) + + +def test_convnext_block_nonorm_inserts_no_normalization_layers(): + block = ConvNeXtBlock( + in_channels=8, + out_channels=8, + kernel_size=3, + dilation=1, + n_layers=1, + norm="nonorm", + ) + + assert not any( + isinstance(layer, (nn.BatchNorm2d, nn.InstanceNorm2d, nn.GroupNorm)) + for layer in block.convblock + ) + + +def test_convnext_block_group_norm_rejects_nonpositive_group_count(): + with pytest.raises(ValueError, match="group_norm_groups must be >= 1"): + ConvNeXtBlock( + in_channels=8, + out_channels=8, + kernel_size=3, + dilation=1, + n_layers=1, + norm="group", + group_norm_groups=0, + ) diff --git a/tests/test_config.py b/tests/test_config.py index 199374146..f9b611616 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -22,9 +22,11 @@ Om4DataSourceConfig, Om4TimeConfig, RolloutValidationConfig, + SamudraConfig, TrainConfig, ) from samudra.config_schema import get_pydantic_models +from samudra.models.modules.blocks import ConvNeXtBlock from samudra.utils.location import LocalLocation, UnresolvedLocation from tests.conftest import DEFAULT_CONFIG, TEST_CONFIGS_DIR from tests.llc_fixtures import write_raw_llc_datasets @@ -392,6 +394,40 @@ def test_get_pydantic_models_collects_loading_variants(): assert models["GpuDataLoadingConfig"] is GpuDataLoadingConfig +def test_llc_train_config_uses_group_norm_and_temporal_stride(tmp_path): + cfg = TrainConfig.from_yaml_and_cli( + [ + "samudra_llc/train.yaml", + "--experiment.base_output_dir", + str(tmp_path / "outputs"), + ] + ) + + assert cfg.temporal_stride == 24 + assert cfg.data.sources[0].type == "llc" + assert cfg.data.sources[0].prognostic_vars_key == "single_1" + assert cfg.data.sources[0].boundary_vars_key == "single_1" + assert isinstance(cfg.model, SamudraConfig) + assert cfg.model.unet.core_block.norm == "group" + assert cfg.model.unet.core_block.group_norm_groups == 32 + block = cfg.model.unet.core_block.build()(1, 1, 1, 1, "constant", False) + assert isinstance(block, ConvNeXtBlock) + + +def test_llc_train_config_allows_cli_override_for_temporal_stride(tmp_path): + cfg = TrainConfig.from_yaml_and_cli( + [ + "samudra_llc/train.yaml", + "--experiment.base_output_dir", + str(tmp_path / "outputs"), + "--temporal_stride", + "12", + ] + ) + + assert cfg.temporal_stride == 12 + + @pytest.mark.parametrize( "preset", sorted( diff --git a/tests/test_temporal_stride.py b/tests/test_temporal_stride.py new file mode 100644 index 000000000..1357adfeb --- /dev/null +++ b/tests/test_temporal_stride.py @@ -0,0 +1,63 @@ +# SPDX-FileCopyrightText: 2026 Samudra Authors +# +# SPDX-License-Identifier: Apache-2.0 + +import numpy as np +import torch +import xarray as xr + +from samudra.constants import build_om4_spec +from samudra.datasets import TorchTrainDataset +from samudra.utils.data import DataSource, Masks + + +def test_torch_train_dataset_temporal_stride_subsamples_windows(): + coords = {"time": range(10), "lat": range(1), "lon": range(1)} + data = xr.Dataset( + { + "thetao_0": xr.DataArray( + np.arange(10, dtype=np.float32).reshape(10, 1, 1), + dims=["time", "lat", "lon"], + coords=coords, + ), + "hfds": xr.DataArray( + (100 + np.arange(10, dtype=np.float32)).reshape(10, 1, 1), + dims=["time", "lat", "lon"], + coords=coords, + ), + } + ) + means = xr.Dataset({"thetao_0": 0.0, "hfds": 0.0}) + stds = xr.Dataset({"thetao_0": 1.0, "hfds": 1.0}) + masks = Masks( + prognostic=torch.ones((1, 1, 1), dtype=torch.bool), + boundary=torch.ones((1, 1), dtype=torch.bool), + ) + dataset_spec = build_om4_spec( + prognostic_vars_key="thetao_1", boundary_vars_key="hfds" + ) + source = DataSource( + "test", + data, + means, + stds, + masks=masks, + dataset_spec=dataset_spec, + ) + + dataset = TorchTrainDataset( + src=source, + prognostic_var_names=["thetao_0"], + boundary_var_names=["hfds"], + hist=1, + steps=2, + normalize_before_mask=True, + masked_fill_value=0.0, + stride=1, + temporal_stride=2, + ) + + assert len(dataset) == 3 + assert dataset[0].raw_data[0][0][:, 0, 0, 0].tolist() == [0.0, 1.0] + assert dataset[1].raw_data[0][0][:, 0, 0, 0].tolist() == [2.0, 3.0] + assert dataset[2].raw_data[0][0][:, 0, 0, 0].tolist() == [4.0, 5.0]