From 3ed5da05aabb7dfa68a2b82a845cd20c2442885e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 22 Oct 2025 15:19:17 -0400 Subject: [PATCH 1/8] Add `add|update_data_dimensions` Deprecate `set_extra_dim_lengths` + utest --- ndsl/initialization/allocator.py | 38 ++++++++++++++++++++++++++++++++ tests/test_dimension_sizer.py | 14 ++++++++++++ 2 files changed, 52 insertions(+) diff --git a/ndsl/initialization/allocator.py b/ndsl/initialization/allocator.py index e466cd9e..cf5b4662 100644 --- a/ndsl/initialization/allocator.py +++ b/ndsl/initialization/allocator.py @@ -10,6 +10,7 @@ from ndsl.constants import SPATIAL_DIMS from ndsl.dsl.typing import Float from ndsl.initialization import GridSizer +from ndsl.logging import ndsl_log_on_rank_0 from ndsl.quantity import Quantity, QuantityHaloSpec @@ -52,8 +53,45 @@ def set_extra_dim_lengths(self, **kwargs: Any) -> None: """ Set the length of extra (non-x/y/z) dimensions. """ + ndsl_log_on_rank_0.warning( + "`QuantityFactory.set_extra_dim_lengths` is deprecated. " + "Use `add_data_dimensions` or `update_data_dimensions`.", + ) self.sizer.extra_dim_lengths.update(kwargs) + def update_data_dimensions( + self, + data_dimension_descriptions: dict[str, int], + ) -> None: + """ + Update the length of extra (non-x/y/z) dimensions, unknown data dimensions + will be added, existing updated. + + Args: + data_dimension_descriptions: Dict of name/length pairs + """ + self.sizer.extra_dim_lengths.update(data_dimension_descriptions) + + def add_data_dimensions( + self, + data_dimension_descriptions: dict[str, int], + ) -> None: + """ + Add new data (non-x/y/z) dimensions via a key-length pair. If the dimension + already exists, it will error out. + + Args: + data_dimension_descriptions: Dict of name/length pairs + """ + for name in data_dimension_descriptions.keys(): + if name in self.sizer.extra_dim_lengths.keys(): + raise ValueError( + f"[NDSL] Data dimension {name} already exists! " + "Use `update_data_dimensions` if you meant to update the length." + ) + + self.sizer.extra_dim_lengths.update(data_dimension_descriptions) + @classmethod def from_backend(cls, sizer: GridSizer, backend: str) -> QuantityFactory: """Initialize a QuantityFactory to use a specific gt4py backend. diff --git a/tests/test_dimension_sizer.py b/tests/test_dimension_sizer.py index 7ccc578d..d939f5db 100644 --- a/tests/test_dimension_sizer.py +++ b/tests/test_dimension_sizer.py @@ -220,3 +220,17 @@ def test_allocator_empty(sizer, dim_case, units, dtype): assert quantity.origin == dim_case.origin assert quantity.extent == dim_case.extent assert quantity.data.shape == dim_case.shape + + +def test_allocator_data_dimensions_operations(sizer): + quantity_factory = QuantityFactory.from_backend(sizer, "numpy") + quantity_factory.add_data_dimensions({"D0": 11}) + assert "D0" in quantity_factory.sizer.extra_dim_lengths.keys() + assert quantity_factory.sizer.extra_dim_lengths["D0"] == 11 + quantity_factory.update_data_dimensions({"D0": 22}) + assert quantity_factory.sizer.extra_dim_lengths["D0"] == 22 + with pytest.raises( + ValueError, + match="Use `update_data_dimensions` if you meant to update the length.", + ): + quantity_factory.add_data_dimensions({"D0": 33}) From 52ac3fa008376921faf6ff121be1bb94ab7cb196 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Wed, 22 Oct 2025 15:41:54 -0400 Subject: [PATCH 2/8] Update usage of `set_extra_dim_lenghts` --- tests/test_4d_fields.py | 8 ++++---- 1 file changed, 4 insertions(+), 4 deletions(-) diff --git a/tests/test_4d_fields.py b/tests/test_4d_fields.py index d08f1552..f8e8bac1 100644 --- a/tests/test_4d_fields.py +++ b/tests/test_4d_fields.py @@ -38,8 +38,8 @@ def __call__(self, q_in: FloatFieldTracer, q_out: FloatField): def test_non_orchestrated_call() -> None: stencil_factory, quantity_factory = get_factories_single_tile(24, 24, 91, 3) - quantity_factory.set_extra_dim_lengths( - **{ + quantity_factory.add_data_dimensions( + { TRACER_DIM: ntracers, } ) @@ -63,8 +63,8 @@ def test_orchestrated_call() -> None: stencil_factory, quantity_factory = get_factories_single_tile_orchestrated( 24, 24, 91, 3 ) - quantity_factory.set_extra_dim_lengths( - **{ + quantity_factory.add_data_dimensions( + { TRACER_DIM: ntracers, } ) From a132ace698087636c6b72a0ba334117930d63532 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 23 Oct 2025 11:49:17 -0400 Subject: [PATCH 3/8] `GridSizer` and `SubtileGridSizer` rename `extra_dim_lengths` attempt (dependencies might die) --- examples/mpi/zarr_monitor.py | 2 +- ndsl/boilerplate.py | 2 +- ndsl/dsl/stencil.py | 2 +- ndsl/grid/generation.py | 6 ++-- ndsl/initialization/allocator.py | 13 ++++---- ndsl/initialization/grid_sizer.py | 14 +++++++-- ndsl/initialization/subtile_grid_sizer.py | 37 +++++++++++++++-------- ndsl/quantity/field_bundle.py | 4 +-- ndsl/quantity/state.py | 6 ++-- ndsl/stencils/testing/grid.py | 2 +- tests/test_dimension_sizer.py | 6 ++-- 11 files changed, 58 insertions(+), 36 deletions(-) diff --git a/examples/mpi/zarr_monitor.py b/examples/mpi/zarr_monitor.py index 0c089af4..b81136b2 100644 --- a/examples/mpi/zarr_monitor.py +++ b/examples/mpi/zarr_monitor.py @@ -19,7 +19,7 @@ def get_example_state(time): - sizer = SubtileGridSizer(nx=48, ny=48, nz=70, n_halo=3, extra_dim_lengths={}) + sizer = SubtileGridSizer(nx=48, ny=48, nz=70, n_halo=3, data_dimensions={}) allocator = QuantityFactory(sizer, np) air_temperature = allocator.zeros([X_DIM, Y_DIM, Z_DIM], units="degK") air_temperature.view[:] = np.random.randn(*air_temperature.extent) diff --git a/ndsl/boilerplate.py b/ndsl/boilerplate.py index 7b8522c1..1c7b9997 100644 --- a/ndsl/boilerplate.py +++ b/ndsl/boilerplate.py @@ -60,7 +60,7 @@ def _get_factories( ny_tile=ny, nz=nz, n_halo=nhalo, - extra_dim_lengths={}, + data_dimensions={}, layout=partitioner.layout, tile_partitioner=partitioner, ) diff --git a/ndsl/dsl/stencil.py b/ndsl/dsl/stencil.py index a12c91bf..e0bfd9c4 100644 --- a/ndsl/dsl/stencil.py +++ b/ndsl/dsl/stencil.py @@ -593,7 +593,7 @@ def domain(self, domain: Index3D) -> None: ny=domain[1], nz=domain[2], n_halo=self.n_halo, - extra_dim_lengths={}, + data_dimensions={}, ) @classmethod diff --git a/ndsl/grid/generation.py b/ndsl/grid/generation.py index d78920ab..12b20288 100644 --- a/ndsl/grid/generation.py +++ b/ndsl/grid/generation.py @@ -242,8 +242,8 @@ def __init__( self._tile_partitioner = self._comm.tile.partitioner self._rank = self._comm.rank self.quantity_factory = quantity_factory - self.quantity_factory.set_extra_dim_lengths( - **{ + self.quantity_factory.add_data_dimensions( + { self.LON_OR_LAT_DIM: 2, self.TILE_DIM: 6, self.CARTESIAN_DIM: 3, @@ -500,7 +500,7 @@ def from_tile_sizing( ny_tile=npy - 1, nz=npz, n_halo=N_HALO_DEFAULT, - extra_dim_lengths={ + data_dimensions={ cls.LON_OR_LAT_DIM: 2, cls.TILE_DIM: 6, cls.CARTESIAN_DIM: 3, diff --git a/ndsl/initialization/allocator.py b/ndsl/initialization/allocator.py index 1e0531bb..ffbe6ee4 100644 --- a/ndsl/initialization/allocator.py +++ b/ndsl/initialization/allocator.py @@ -10,7 +10,6 @@ from ndsl.constants import SPATIAL_DIMS from ndsl.dsl.typing import Float from ndsl.initialization import GridSizer -from ndsl.logging import ndsl_log_on_rank_0 from ndsl.quantity import Quantity, QuantityHaloSpec @@ -53,11 +52,13 @@ def set_extra_dim_lengths(self, **kwargs: Any) -> None: """ Set the length of extra (non-x/y/z) dimensions. """ - ndsl_log_on_rank_0.warning( + warnings.warn( "`QuantityFactory.set_extra_dim_lengths` is deprecated. " "Use `add_data_dimensions` or `update_data_dimensions`.", + DeprecationWarning, + 2, ) - self.sizer.extra_dim_lengths.update(kwargs) + self.sizer.data_dimensions.update(kwargs) def update_data_dimensions( self, @@ -70,7 +71,7 @@ def update_data_dimensions( Args: data_dimension_descriptions: Dict of name/length pairs """ - self.sizer.extra_dim_lengths.update(data_dimension_descriptions) + self.sizer.data_dimensions.update(data_dimension_descriptions) def add_data_dimensions( self, @@ -84,13 +85,13 @@ def add_data_dimensions( data_dimension_descriptions: Dict of name/length pairs """ for name in data_dimension_descriptions.keys(): - if name in self.sizer.extra_dim_lengths.keys(): + if name in self.sizer.data_dimensions.keys(): raise ValueError( f"[NDSL] Data dimension {name} already exists! " "Use `update_data_dimensions` if you meant to update the length." ) - self.sizer.extra_dim_lengths.update(data_dimension_descriptions) + self.sizer.data_dimensions.update(data_dimension_descriptions) @classmethod def from_backend(cls, sizer: GridSizer, backend: str) -> QuantityFactory: diff --git a/ndsl/initialization/grid_sizer.py b/ndsl/initialization/grid_sizer.py index 8dee7965..961ab793 100644 --- a/ndsl/initialization/grid_sizer.py +++ b/ndsl/initialization/grid_sizer.py @@ -1,3 +1,4 @@ +import warnings from collections.abc import Sequence from dataclasses import dataclass @@ -12,8 +13,17 @@ class GridSizer: """Length of the z compute dimension for produced arrays.""" n_halo: int """Number of horizontal halo points for produced arrays.""" - extra_dim_lengths: dict[str, int] - """Lengths of any non-x/y/z dimensions, such as land or radiation dimensions.""" + data_dimensions: dict[str, int] + """Name/Lengths pair of any non-x/y/z dimensions, such as land or radiation dimensions.""" + + @property + def extra_dim_lengths(self) -> dict[str, int]: + warnings.warn( + "`GridSizer.extra_dim_lengths` is a deprecated API, use `GridSizer.data_dimensions`.", + DeprecationWarning, + 2, + ) + return self.data_dimensions def get_origin(self, dims: Sequence[str]) -> tuple[int, ...]: raise NotImplementedError() diff --git a/ndsl/initialization/subtile_grid_sizer.py b/ndsl/initialization/subtile_grid_sizer.py index aadf755b..923247f5 100644 --- a/ndsl/initialization/subtile_grid_sizer.py +++ b/ndsl/initialization/subtile_grid_sizer.py @@ -1,3 +1,4 @@ +import warnings from collections.abc import Iterable from typing import Self @@ -15,10 +16,12 @@ def from_tile_params( ny_tile: int, nz: int, n_halo: int, - extra_dim_lengths: dict[str, int], layout: tuple[int, int], + *, + data_dimensions: dict[str, int] = {}, tile_partitioner: TilePartitioner | None = None, tile_rank: int = 0, + extra_dim_lengths: dict[str, int] | None = None, ) -> Self: """Create a SubtileGridSizer from parameters about the full tile. @@ -27,13 +30,21 @@ def from_tile_params( ny_tile: number of y cell centers on the tile nz: number of vertical levels n_halo: number of halo points - extra_dim_lengths: lengths of any non-x/y/z dimensions, + data_dimensions: lengths of any non-x/y/z dimensions, such as land or radiation dimensions layout: (y, x) number of ranks along tile edges tile_partitioner (optional): partitioner object for the tile. By default, a TilePartitioner is created with the given layout tile_rank (optional): rank of this subtile. + extra_dim_lengths: DEPRECATED API - use `data_dimensions` """ + if extra_dim_lengths is not None: + warnings.warn( + "`extra_dim_lengths` is a deprecated name, please `data_dimensions`.", + DeprecationWarning, + 2, + ) + data_dimensions = extra_dim_lengths if tile_partitioner is None: tile_partitioner = TilePartitioner(layout) y_slice, x_slice = tile_partitioner.subtile_slice( @@ -55,7 +66,7 @@ def from_tile_params( "SubtileGridSizer::from_tile_params: Compute domain extent must be greater than halo size" ) - return cls(nx, ny, nz, n_halo, extra_dim_lengths) + return cls(nx, ny, nz, n_halo, data_dimensions) @classmethod def from_namelist( @@ -92,19 +103,19 @@ def from_namelist( "expected to find nx_tile or fv_core_nml" ) return cls.from_tile_params( - nx_tile, - ny_tile, - nz, - N_HALO_DEFAULT, - {}, - layout, - tile_partitioner, - tile_rank, + nx_tile=nx_tile, + ny_tile=ny_tile, + nz=nz, + n_halo=N_HALO_DEFAULT, + data_dimensions={}, + layout=layout, + tile_partitioner=tile_partitioner, + tile_rank=tile_rank, ) @property def dim_extents(self) -> dict[str, int]: - return_dict = self.extra_dim_lengths.copy() + return_dict = self.data_dimensions.copy() return_dict.update( { constants.X_DIM: self.nx, @@ -128,7 +139,7 @@ def get_extent(self, dims: Iterable[str]) -> tuple[int, ...]: return tuple(extents[dim] for dim in dims) def get_shape(self, dims: Iterable[str]) -> tuple[int, ...]: - shape_dict = self.extra_dim_lengths.copy() + shape_dict = self.data_dimensions.copy() # must pad non-interface variables to have the same shape as interface variables shape_dict.update( { diff --git a/ndsl/quantity/field_bundle.py b/ndsl/quantity/field_bundle.py index fc6a5dd7..f3d22968 100644 --- a/ndsl/quantity/field_bundle.py +++ b/ndsl/quantity/field_bundle.py @@ -118,8 +118,8 @@ def extend_3D_quantity_factory( extra_dims: dict of [name, size] of the data dimensions to add. """ new_factory = copy.copy(quantity_factory) - new_factory.set_extra_dim_lengths( - **{ + new_factory.add_data_dimensions( + { **extra_dims, } ) diff --git a/ndsl/quantity/state.py b/ndsl/quantity/state.py index b1c776a1..64137458 100644 --- a/ndsl/quantity/state.py +++ b/ndsl/quantity/state.py @@ -84,8 +84,8 @@ def __init__(self, factory: QuantityFactory, ddims: dict[str, int]): self._factory = factory def __enter__(self) -> None: - self._original_dims = self._factory.sizer.extra_dim_lengths - self._factory.sizer.extra_dim_lengths = self._ddims + self._original_dims = self._factory.sizer.data_dimensions + self._factory.sizer.data_dimensions = self._ddims def __exit__( self, @@ -93,7 +93,7 @@ def __exit__( value: BaseException | None, traceback: TracebackType | None, ) -> None: - self._factory.sizer.extra_dim_lengths = self._original_dims + self._factory.sizer.data_dimensions = self._original_dims @classmethod def empty( diff --git a/ndsl/stencils/testing/grid.py b/ndsl/stencils/testing/grid.py index 06a90704..c6c6d16d 100644 --- a/ndsl/stencils/testing/grid.py +++ b/ndsl/stencils/testing/grid.py @@ -143,7 +143,7 @@ def sizer(self): ny_tile=self.npy - 1, nz=self.npz, n_halo=self.halo, - extra_dim_lengths={ + data_dimensions={ MetricTerms.LON_OR_LAT_DIM: 2, MetricTerms.TILE_DIM: 6, MetricTerms.CARTESIAN_DIM: 3, diff --git a/tests/test_dimension_sizer.py b/tests/test_dimension_sizer.py index d939f5db..6181de44 100644 --- a/tests/test_dimension_sizer.py +++ b/tests/test_dimension_sizer.py @@ -225,10 +225,10 @@ def test_allocator_empty(sizer, dim_case, units, dtype): def test_allocator_data_dimensions_operations(sizer): quantity_factory = QuantityFactory.from_backend(sizer, "numpy") quantity_factory.add_data_dimensions({"D0": 11}) - assert "D0" in quantity_factory.sizer.extra_dim_lengths.keys() - assert quantity_factory.sizer.extra_dim_lengths["D0"] == 11 + assert "D0" in quantity_factory.sizer.data_dimensions.keys() + assert quantity_factory.sizer.data_dimensions["D0"] == 11 quantity_factory.update_data_dimensions({"D0": 22}) - assert quantity_factory.sizer.extra_dim_lengths["D0"] == 22 + assert quantity_factory.sizer.data_dimensions["D0"] == 22 with pytest.raises( ValueError, match="Use `update_data_dimensions` if you meant to update the length.", From 5251e37cf0ee0e51057f7acf53bc7fea1cde3e5a Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 23 Oct 2025 11:51:55 -0400 Subject: [PATCH 4/8] Verbose comment --- ndsl/initialization/allocator.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/ndsl/initialization/allocator.py b/ndsl/initialization/allocator.py index ffbe6ee4..e9d8857e 100644 --- a/ndsl/initialization/allocator.py +++ b/ndsl/initialization/allocator.py @@ -65,8 +65,8 @@ def update_data_dimensions( data_dimension_descriptions: dict[str, int], ) -> None: """ - Update the length of extra (non-x/y/z) dimensions, unknown data dimensions - will be added, existing updated. + Update the length of data (non-x/y/z) dimensions, unknown data dimensions + will be added, existing ones updated. Args: data_dimension_descriptions: Dict of name/length pairs From b8b9d3a005ac262cd427b26ec78e188372c2fb31 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 23 Oct 2025 12:03:45 -0400 Subject: [PATCH 5/8] Missed pytest.fixture --- tests/test_dimension_sizer.py | 12 ++++++------ 1 file changed, 6 insertions(+), 6 deletions(-) diff --git a/tests/test_dimension_sizer.py b/tests/test_dimension_sizer.py index 6181de44..ebad25ef 100644 --- a/tests/test_dimension_sizer.py +++ b/tests/test_dimension_sizer.py @@ -70,12 +70,12 @@ def namelist(nx_tile, ny_tile, nz, layout): def sizer(request, nx_tile, ny_tile, nz, layout, namelist, extra_dimension_lengths): if request.param == "from_tile_params": sizer = SubtileGridSizer.from_tile_params( - nx_tile, - ny_tile, - nz, - N_HALO_DEFAULT, - extra_dimension_lengths, - layout, + nx_tile=nx_tile, + ny_tile=ny_tile, + nz=nz, + n_halo=N_HALO_DEFAULT, + layout=layout, + data_dimensions=extra_dimension_lengths, ) elif request.param == "from_namelist": sizer = SubtileGridSizer.from_namelist(namelist) From ca3993fc29442d893468815caecc4620e899ebc6 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 23 Oct 2025 13:03:32 -0400 Subject: [PATCH 6/8] Better doc --- ndsl/initialization/subtile_grid_sizer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/ndsl/initialization/subtile_grid_sizer.py b/ndsl/initialization/subtile_grid_sizer.py index 923247f5..835c2cd6 100644 --- a/ndsl/initialization/subtile_grid_sizer.py +++ b/ndsl/initialization/subtile_grid_sizer.py @@ -40,7 +40,7 @@ def from_tile_params( """ if extra_dim_lengths is not None: warnings.warn( - "`extra_dim_lengths` is a deprecated name, please `data_dimensions`.", + "`extra_dim_lengths` is a deprecated name, please use `data_dimensions` instead.", DeprecationWarning, 2, ) From aa1bebeed5e9fe7195075ad3311c734d29158fd9 Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 23 Oct 2025 15:02:14 -0400 Subject: [PATCH 7/8] MetricTerms: default constructor take care of data dimensions --- ndsl/grid/generation.py | 5 ----- 1 file changed, 5 deletions(-) diff --git a/ndsl/grid/generation.py b/ndsl/grid/generation.py index 12b20288..1bc37bc5 100644 --- a/ndsl/grid/generation.py +++ b/ndsl/grid/generation.py @@ -500,11 +500,6 @@ def from_tile_sizing( ny_tile=npy - 1, nz=npz, n_halo=N_HALO_DEFAULT, - data_dimensions={ - cls.LON_OR_LAT_DIM: 2, - cls.TILE_DIM: 6, - cls.CARTESIAN_DIM: 3, - }, layout=communicator.partitioner.tile.layout, ) quantity_factory = QuantityFactory.from_backend(sizer, backend=backend) From bcea590565899feabc3b84f241363e2d54b1856e Mon Sep 17 00:00:00 2001 From: Florian Deconinck Date: Thu, 23 Oct 2025 15:43:35 -0400 Subject: [PATCH 8/8] Make `MetricTerms` reentrant --- ndsl/grid/generation.py | 1 + tests/grid/test_eta.py | 6 +++--- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/ndsl/grid/generation.py b/ndsl/grid/generation.py index 1bc37bc5..fbf98936 100644 --- a/ndsl/grid/generation.py +++ b/ndsl/grid/generation.py @@ -501,6 +501,7 @@ def from_tile_sizing( nz=npz, n_halo=N_HALO_DEFAULT, layout=communicator.partitioner.tile.layout, + data_dimensions={}, ) quantity_factory = QuantityFactory.from_backend(sizer, backend=backend) return cls( diff --git a/tests/grid/test_eta.py b/tests/grid/test_eta.py index 9e9c7e92..d08600fa 100755 --- a/tests/grid/test_eta.py +++ b/tests/grid/test_eta.py @@ -53,7 +53,7 @@ def test_set_hybrid_pressure_coefficients_correct(levels): ny_tile=ny, nz=nz, n_halo=nhalo, - extra_dim_lengths={}, + data_dimensions={}, layout=layout, tile_partitioner=partitioner.tile, tile_rank=communicator.tile.rank, @@ -102,7 +102,7 @@ def test_set_hybrid_pressure_coefficients_nofile(): ny_tile=ny, nz=nz, n_halo=nhalo, - extra_dim_lengths={}, + data_dimensions={}, layout=layout, tile_partitioner=partitioner.tile, tile_rank=communicator.tile.rank, @@ -146,7 +146,7 @@ def test_set_hybrid_pressure_coefficients_not_mono(): ny_tile=ny, nz=nz, n_halo=nhalo, - extra_dim_lengths={}, + data_dimensions={}, layout=layout, tile_partitioner=partitioner.tile, tile_rank=communicator.tile.rank,