Skip to content

refactor: prefer backend over gt4py_backend #170

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
oelbert merged 1 commit into from
Jan 13, 2026
Merged

refactor: prefer backend over gt4py_backend #170

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
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
8 changes: 4 additions & 4 deletions examples/notebooks/functions.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -106,7 +106,7 @@ def init_quantity(
units=units,
origin=(nhalo, nhalo, 0)[:skip_z],
extent=(nx, ny, nz)[:skip_z],
gt4py_backend=backend,
backend=backend,
)

if grid == VariableGrid.CellCorners:
Expand All @@ -116,7 +116,7 @@ def init_quantity(
units=units,
origin=(nhalo, nhalo, 0)[:skip_z],
extent=(nx + 1, ny + 1, nz)[:skip_z],
gt4py_backend=backend,
backend=backend,
)

elif grid == VariableGrid.StaggeredInX:
Expand All @@ -126,7 +126,7 @@ def init_quantity(
units=units,
origin=(nhalo, nhalo, 0)[:skip_z],
extent=(nx + 1, ny, nz)[:skip_z],
gt4py_backend=backend,
backend=backend,
)

elif grid == VariableGrid.StaggeredInY:
Expand All @@ -136,7 +136,7 @@ def init_quantity(
units=units,
origin=(nhalo, nhalo, 0)[:skip_z],
extent=(nx, ny + 1, nz)[:skip_z],
gt4py_backend=backend,
backend=backend,
)

return variable
Expand Down
6 changes: 3 additions & 3 deletions examples/notebooks/grid_generation.ipynb
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -267,15 +267,15 @@
" units=\"degrees\",\n",
" origin=(nhalo, nhalo),\n",
" extent=(nx + 1, ny + 1),\n",
" gt4py_backend=backend,\n",
" backend=backend,\n",
")\n",
"lat = Quantity(\n",
" metric_terms.lat.data * 180 / np.pi,\n",
" dims=(\"x_interface\", \"y_interface\"),\n",
" units=\"degrees\",\n",
" origin=(nhalo, nhalo),\n",
" extent=(nx + 1, ny + 1),\n",
" gt4py_backend=backend,\n",
" backend=backend,\n",
")\n",
"\n",
"# gather the distributed fields into a global field on the root rank\n",
Expand Down Expand Up @@ -357,7 +357,7 @@
" units=\"m2\",\n",
" origin=(nhalo, nhalo),\n",
" extent=(nx, ny),\n",
" gt4py_backend=backend,\n",
" backend=backend,\n",
")\n",
"\n",
"# rescale to 10^3 km2\n",
Expand Down
6 changes: 1 addition & 5 deletions pace/grid.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -181,12 +181,8 @@ def get_grid(
quantity_factory: QuantityFactory,
communicator: Communicator,
) -> Tuple[DampingCoefficients, DriverGridData, GridData]:
backend = quantity_factory.zeros(
dims=[X_DIM, Y_DIM], units="unknown"
).gt4py_backend

ndsl_log.info("Using serialized grid data")
grid = self._get_serialized_grid(communicator, backend)
grid = self._get_serialized_grid(communicator, quantity_factory.backend)
grid_data = grid.grid_data
driver_grid_data = grid.driver_grid_data
damping_coefficients = grid.damping_coefficients
Expand Down
8 changes: 3 additions & 5 deletions pace/initialization.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -286,11 +286,9 @@ def get_driver_state(
grid_data: GridData,
schemes: List[PHYSICS_PACKAGES],
) -> DriverState:
backend = quantity_factory.zeros(
dims=[X_DIM, Y_DIM], units="unknown"
).gt4py_backend

dycore_state = self._initialize_dycore_state(communicator, backend)
Comment on lines -289 to -293
Copy link
Collaborator Author

@romanc romanc Jan 6, 2026

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Since the refactor in NDSL, the backend is now known and saved on the QuantityFactory. No need to allocate something just to get the name of the backend.

dycore_state = self._initialize_dycore_state(
communicator, quantity_factory.backend
)
physics_state = PhysicsState.init_zeros(
quantity_factory=quantity_factory,
schemes=schemes,
Expand Down
4 changes: 2 additions & 2 deletions tests/main/driver/test_safety_checks.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -65,7 +65,7 @@ def test_check_state_domain_only():
"unknown",
origin=(1, 1, 0),
extent=(3, 3, 2),
gt4py_backend="numpy",
backend="numpy",
)
dycore_state = unittest.mock.MagicMock(u=u_quantity)
checker.check_state(dycore_state)
Expand All @@ -83,7 +83,7 @@ def test_check_nan_value():
"unknown",
origin=(0, 0, 0),
extent=(4, 4, 2),
gt4py_backend="numpy",
backend="numpy",
)
dycore_state = unittest.mock.MagicMock(u=u_quantity)
with pytest.raises(RuntimeError):
Expand Down