Skip to content

Commit 19ef0fb

Browse files
[DSL] Update model code to NDSL v2025.10.00-dirty (#1177)
* SatTables: move `compute_tables` in flag Fix compilation for layout above 1x1 * Fix call to get_tables * Update QuantityFactory data dimensions change Update remove of numpy module in quantity factory * Remove unused code * Add strict editable `build` directories to gitignore * Move `Namelist` import to `f90nml` Lint * Clean up unused API * Add strides to early pyMoist wrapper report. Log only on rank 0. Make `black` consistently happy * Clean up interface prints * Proper strides into bytes calculation * Update data dim API on the QuantityFactory * Update base NDSL version
1 parent 922902b commit 19ef0fb

34 files changed

+302
-262
lines changed

GEOSagcm_GridComp/GEOSphysics_GridComp/.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -4,3 +4,4 @@
44
/@GEOSradiation_GridComp
55
/GEOSradiation_GridComp
66
/GEOSradiation_GridComp@
7+
**/build/*

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/.gitignore

Lines changed: 1 addition & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -13,3 +13,4 @@ sandbox/
1313
.gt4py_cache/
1414
.gt_cache_*/
1515
.translate*/
16+
**/build/

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/GFDL_1M/driver/driver.py

Lines changed: 4 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -84,7 +84,10 @@ def __init__(
8484
# generate saturation specific humidity tables
8585
# -----------------------------------------------------------------------
8686

87-
self.sat_tables = get_tables(stencil_factory.backend)
87+
self.sat_tables = get_tables(
88+
stencil_factory.backend,
89+
stencil_factory.config.dace_config,
90+
)
8891

8992
# -----------------------------------------------------------------------
9093
# initialize stencils

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/GFDL_1M/driver/sat_tables.py

Lines changed: 6 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -2,6 +2,7 @@
22

33
from ndsl.boilerplate import get_factories_single_tile
44
from ndsl.constants import X_DIM, Y_DIM, Z_DIM
5+
from ndsl.dsl.dace.dace_config import DaceConfig
56
from ndsl.dsl.gt4py import FORWARD, PARALLEL, GlobalTable, K, computation, exp, interval, log, log10
67
from ndsl.dsl.typing import Float, FloatField, Int
78
from pyMoist.GFDL_1M.driver.constants import constants
@@ -189,7 +190,7 @@ class GFDL_driver_tables:
189190
Reference Fortran: gfdl_cloud_microphys.F90: qsmith_init.py
190191
"""
191192

192-
def __init__(self, backend):
193+
def __init__(self, backend, dace_config: DaceConfig):
193194
table_compute_domain = (1, 1, constants.LENGTH)
194195

195196
stencil_factory, quantity_factory = get_factories_single_tile(
@@ -213,7 +214,7 @@ def __init__(self, backend):
213214
# Cancel multi-node compile for tables
214215
# TODO: this should come for free with the rewrite of the gt:X stencils
215216
# compilation mode
216-
if MPI.COMM_WORLD.Get_rank() != 0:
217+
if not dace_config.do_compile:
217218
MPI.COMM_WORLD.Barrier()
218219

219220
compute_qs_table_1 = stencil_factory.from_origin_domain(
@@ -258,7 +259,7 @@ def __init__(self, backend):
258259
self._table4,
259260
)
260261

261-
if MPI.COMM_WORLD.Get_rank() == 0:
262+
if dace_config.do_compile:
262263
MPI.COMM_WORLD.Barrier()
263264

264265
self.table1 = self._table1.view[0, 0, :]
@@ -277,8 +278,8 @@ def __init__(self, backend):
277278
}
278279

279280

280-
def get_tables(backend):
281+
def get_tables(backend, dace_config):
281282
if _cached_table["driver_qsat"] is None:
282-
_cached_table["driver_qsat"] = GFDL_driver_tables(backend)
283+
_cached_table["driver_qsat"] = GFDL_driver_tables(backend, dace_config)
283284

284285
return _cached_table["driver_qsat"]

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/UW/compute_uwshcu.py

Lines changed: 22 additions & 28 deletions
Original file line numberDiff line numberDiff line change
@@ -7512,37 +7512,31 @@ def __init__(
75127512
self.stop_cin = self.quantity_factory.zeros([X_DIM, Y_DIM], "n/a", dtype=bool)
75137513
self.stop_buoyancy_sort = self.quantity_factory.zeros([X_DIM, Y_DIM], "n/a", dtype=bool)
75147514

7515-
self.ntracers_quantity_factory = self.make_ntracers_quantity_factory(
7516-
self.quantity_factory,
7515+
self.quantity_factory.add_data_dimensions(
7516+
{
7517+
"ntracers": constants.NCNST,
7518+
}
75177519
)
75187520

75197521
# Create tracer fields
7520-
self.trsrc = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7521-
self.trsrc_o = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7522-
self.tre = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7523-
self.trmin = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7524-
self.sstr0 = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7525-
self.tr0 = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7526-
self.tr0_o = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7527-
self.sstr0_o = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7528-
self.trten = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7529-
self.tr0_s = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7530-
self.tr0_inoutvar = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7531-
self.tr0_inout = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7532-
self.trflx = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7533-
self.tru = self.ntracers_quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7534-
self.tru_emf = self.ntracers_quantity_factory.zeros(
7535-
[X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a"
7536-
)
7537-
self.xflx_ndim = self.ntracers_quantity_factory.zeros(
7538-
[X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a"
7539-
)
7540-
self.trflx_d = self.ntracers_quantity_factory.zeros(
7541-
[X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a"
7542-
)
7543-
self.trflx_u = self.ntracers_quantity_factory.zeros(
7544-
[X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a"
7545-
)
7522+
self.trsrc = self.quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7523+
self.trsrc_o = self.quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7524+
self.tre = self.quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7525+
self.trmin = self.quantity_factory.zeros([X_DIM, Y_DIM, "ntracers"], "n/a")
7526+
self.sstr0 = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7527+
self.tr0 = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7528+
self.tr0_o = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7529+
self.sstr0_o = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7530+
self.trten = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7531+
self.tr0_s = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7532+
self.tr0_inoutvar = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7533+
self.tr0_inout = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_DIM, "ntracers"], "n/a")
7534+
self.trflx = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7535+
self.tru = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7536+
self.tru_emf = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7537+
self.xflx_ndim = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7538+
self.trflx_d = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
7539+
self.trflx_u = self.quantity_factory.zeros([X_DIM, Y_DIM, Z_INTERFACE_DIM, "ntracers"], "n/a")
75467540

75477541
self.saturation_vapor_pressure_table = get_saturation_vapor_pressure_table(
75487542
self.stencil_factory.backend

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/aer_activation.py

Lines changed: 11 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -1,5 +1,3 @@
1-
import copy
2-
31
from gt4py.cartesian.gtscript import PARALLEL, computation, exp, float64, interval, log, sqrt
42

53
import pyMoist.constants as constants
@@ -326,28 +324,33 @@ def __init__(
326324
raise NotImplementedError("Non NN Aerosol not implemented")
327325

328326
# Temporary buffers
329-
nmodes_quantity_factory = AerActivation.make_nmodes_quantity_factory(quantity_factory)
330-
self._nact = nmodes_quantity_factory.zeros(
327+
quantity_factory.add_data_dimensions(
328+
**{
329+
"n_modes": constants.N_MODES,
330+
}
331+
)
332+
333+
self._nact = quantity_factory.zeros(
331334
[X_DIM, Y_DIM, Z_DIM, "n_modes"],
332335
units="n/a",
333336
dtype=Float,
334337
)
335-
self._ni = nmodes_quantity_factory.zeros(
338+
self._ni = quantity_factory.zeros(
336339
[X_DIM, Y_DIM, Z_DIM, "n_modes"],
337340
units="n/a",
338341
dtype=Float,
339342
)
340-
self._rg = nmodes_quantity_factory.zeros(
343+
self._rg = quantity_factory.zeros(
341344
[X_DIM, Y_DIM, Z_DIM, "n_modes"],
342345
units="n/a",
343346
dtype=Float,
344347
)
345-
self._sig0 = nmodes_quantity_factory.zeros(
348+
self._sig0 = quantity_factory.zeros(
346349
[X_DIM, Y_DIM, Z_DIM, "n_modes"],
347350
units="n/a",
348351
dtype=Float,
349352
)
350-
self._bibar = nmodes_quantity_factory.zeros(
353+
self._bibar = quantity_factory.zeros(
351354
[X_DIM, Y_DIM, Z_DIM, "n_modes"],
352355
units="n/a",
353356
dtype=Float,
@@ -359,16 +362,6 @@ def __init__(
359362
compute_dims=[X_DIM, Y_DIM, Z_DIM],
360363
)
361364

362-
@staticmethod
363-
def make_nmodes_quantity_factory(ijk_quantity_factory: QuantityFactory):
364-
nmodes_quantity_factory = copy.deepcopy(ijk_quantity_factory)
365-
nmodes_quantity_factory.set_extra_dim_lengths(
366-
**{
367-
"n_modes": constants.N_MODES,
368-
}
369-
)
370-
return nmodes_quantity_factory
371-
372365
def __call__(
373366
self,
374367
aero_dgn: FloatField_NModes,

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/mapl/memory_factory.py

Lines changed: 8 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,9 +8,15 @@
88
from MAPLpyish import MAPLBridge, MAPLState
99

1010
from ndsl import QuantityFactory, ndsl_log
11+
from ndsl.dsl.gt4py_utils import is_gpu_backend
12+
from ndsl.optional_imports import cupy as cp
1113
from pyMoist.interface.f_py_conversion import FortranPythonConversion
1214

1315

16+
if cp is None:
17+
cp = np
18+
19+
1420
class MAPLMemoryRepository:
1521
"""A factory capable of accessing memory handlded by MAPL and formatting it
1622
for direct use in DSL application."""
@@ -32,11 +38,12 @@ def __init__(self, state: MAPLState, quantity_factory: QuantityFactory) -> None:
3238
self._quantity_factory = quantity_factory
3339
self._state = state
3440
self._bridge = MAPLBridge()
41+
self._quantity_factory.backend
3542
self._f_py_converter = FortranPythonConversion(
3643
self._quantity_factory.sizer.nx,
3744
self._quantity_factory.sizer.ny,
3845
self._quantity_factory.sizer.nz,
39-
self._quantity_factory._numpy,
46+
cp if is_gpu_backend(self._quantity_factory.backend) else np,
4047
)
4148
self._fortran_pointers: dict[str, MAPLMemoryRepository.FortranMemory] = {}
4249

GEOSagcm_GridComp/GEOSphysics_GridComp/GEOSmoist_GridComp/pyMoist/pyMoist/interface/python_bridge.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -29,7 +29,6 @@ def init(
2929
self.rank = MPI.COMM_WORLD.Get_rank()
3030
self.backend = backend
3131
self.flags = moist_flags_f_to_python(pyMoist_flags)
32-
print(f"Moist Flags:\n{self.flags}")
3332
# For Fortran<->NumPy conversion
3433
if is_gpu_backend(self.backend):
3534
numpy_module = cp

0 commit comments

Comments
 (0)