From a040e7668dde3eee510bfe7cbffad3a6c252248f Mon Sep 17 00:00:00 2001 From: root Date: Wed, 25 Sep 2024 02:10:12 +0000 Subject: [PATCH 1/3] fix coupled dataloader batch inference Signed-off-by: root --- modulus/datapipes/healpix/couplers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modulus/datapipes/healpix/couplers.py b/modulus/datapipes/healpix/couplers.py index 5289b32778..19231d3746 100644 --- a/modulus/datapipes/healpix/couplers.py +++ b/modulus/datapipes/healpix/couplers.py @@ -173,14 +173,14 @@ def set_coupled_fields(self, coupled_fields): # create buffer for coupling coupled_fields = coupled_fields[ :, :, :, self.coupled_channel_indices, :, : - ].permute(0, 2, 3, 1, 4, 5) + ].permute(2, 0, 3, 1, 4, 5) self.preset_coupled_fields = th.empty( - [self.coupled_integration_dim, self.batch_size, self.timevar_dim] + [self.coupled_integration_dim, coupled_fields.shape[1], self.timevar_dim] + list(self.spatial_dims) ) for i in range(len(self.preset_coupled_fields)): self.preset_coupled_fields[i, :, :, :, :, :] = coupled_fields[ - 0, 0, -1, :, :, : + 0, :, -1:, :, :, : ] # flag for construct integrated coupling method to use this array self.coupled_mode = True From b60dad79b8a5a15c5c2ff1eef231c566ea3ad4a1 Mon Sep 17 00:00:00 2001 From: root Date: Sun, 29 Sep 2024 04:29:31 +0000 Subject: [PATCH 2/3] Add test to cover set_coupled_fields in dlwp-healpix coupled inference Signed-off-by: root --- test/datapipes/test_healpix_couple.py | 54 +++++++++++++++++++++++++++ 1 file changed, 54 insertions(+) diff --git a/test/datapipes/test_healpix_couple.py b/test/datapipes/test_healpix_couple.py index 833827c9ef..bd4bc13e8c 100644 --- a/test/datapipes/test_healpix_couple.py +++ b/test/datapipes/test_healpix_couple.py @@ -22,6 +22,7 @@ import numpy as np import pandas as pd import pytest +import torch as th import xarray as xr from omegaconf import DictConfig, OmegaConf from pytest_utils import nfsdata_or_fail @@ -138,6 +139,25 @@ def test_ConstantCoupler(data_dir, dataset_name, scaling_dict, pytestconfig): expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) assert np.array_equal(expected, coupler.coupled_scaling["std"]) + coupler.coupled_channel_indices = [0, 1] + coupled_fields_batch_size = 4 + coupled_fields_timedim = 2 + coupled_fields = th.rand( + coupled_fields_batch_size, + coupler.spatial_dims[0], + coupled_fields_timedim, + len(coupler.coupled_channel_indices), + coupler.spatial_dims[1], + coupler.spatial_dims[2], + ) + expected_shape = [ + coupler.coupled_integration_dim, + coupled_fields_batch_size, + coupler.timevar_dim, + ] + list(coupler.spatial_dims) + coupler.set_coupled_fields(coupled_fields) + assert list(coupler.preset_coupled_fields.shape) == expected_shape + DistributedManager.cleanup() @@ -196,6 +216,40 @@ def test_TrailingAverageCoupler(data_dir, dataset_name, scaling_dict, pytestconf expected = np.expand_dims(coupled_scaling["std"].to_numpy(), (0, 2, 3, 4)) assert np.array_equal(expected, coupler.coupled_scaling["std"]) + averaging_window_max_indices = [ + i // pd.Timedelta(data_time_step) for i in coupler.input_times + ] + di = averaging_window_max_indices[0] + averaging_slices = [] + for j in range(coupler.coupled_integration_dim): + averaging_slices.append([]) + for i, r in enumerate(averaging_window_max_indices): + averaging_slices[j].append( + slice( + coupler.input_time_dim * j * di + i * di, + coupler.input_time_dim * j * di + r, + ) + ) + coupler.averaging_slices = averaging_slices + coupler.coupled_channel_indices = [0, 1] + coupled_fields_batch_size = 4 + coupled_fields_timedim = 4 + coupled_fields = th.rand( + coupled_fields_batch_size, + coupler.spatial_dims[0], + coupled_fields_timedim, + len(coupler.coupled_channel_indices), + coupler.spatial_dims[1], + coupler.spatial_dims[2], + ) + expected_shape = [ + coupler.coupled_integration_dim, + coupled_fields_batch_size, + coupler.timevar_dim, + ] + list(coupler.spatial_dims) + coupler.set_coupled_fields(coupled_fields) + assert list(coupler.preset_coupled_fields.shape) == expected_shape + DistributedManager.cleanup() From be5e007d54127a432ad60a076d97970dd4b714a8 Mon Sep 17 00:00:00 2001 From: root Date: Thu, 19 Dec 2024 23:17:34 +0000 Subject: [PATCH 3/3] Fix indexing in constant coupler Signed-off-by: root --- modulus/datapipes/healpix/couplers.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/modulus/datapipes/healpix/couplers.py b/modulus/datapipes/healpix/couplers.py index d191a54703..42afb162ee 100644 --- a/modulus/datapipes/healpix/couplers.py +++ b/modulus/datapipes/healpix/couplers.py @@ -173,14 +173,14 @@ def set_coupled_fields(self, coupled_fields): # create buffer for coupling coupled_fields = coupled_fields[ :, :, :, self.coupled_channel_indices, :, : - ].permute(2, 0, 3, 1, 4, 5) + ].permute(0, 2, 3, 1, 4, 5) self.preset_coupled_fields = th.empty( - [self.coupled_integration_dim, coupled_fields.shape[1], self.timevar_dim] + [self.coupled_integration_dim, coupled_fields.shape[0], self.timevar_dim] + list(self.spatial_dims) ) for i in range(len(self.preset_coupled_fields)): self.preset_coupled_fields[i, :, :, :, :, :] = coupled_fields[ - 0, :, -1:, :, :, : + :, -1, :, :, :, : ] # flag for construct integrated coupling method to use this array self.coupled_mode = True