diff --git a/pyfv3/stencils/copy_corners.py b/pyfv3/stencils/copy_corners.py new file mode 100644 index 00000000..266627df --- /dev/null +++ b/pyfv3/stencils/copy_corners.py @@ -0,0 +1,169 @@ +from ndsl import StencilFactory, orchestrate +from ndsl.dsl.typing import FloatField + + +def corner_copy_x(field_to_copy): + """Equivalent to the copy_corners_x functions in fortran. + + This is written to operate on plain ndarrarys and not use the GT4Py framework. + This choice was made because we've seen a lot of performance left on the table using + orchestration without explicitly describing the operations but rather have full 3d- + sweeps with conditionals. + Since DaCe can handle (simple) operations on ndarrays directly this gives us a more + explicit entrypoint to the language and more optimization-potential. + + Args: + field_to_copy (ndarray): field to apply the corner copy on. + This is explicitly not type-hinted for orchestration + """ + field_to_copy[0, 0] = field_to_copy[0, 5] + field_to_copy[0, 1] = field_to_copy[1, 5] + field_to_copy[0, 2] = field_to_copy[2, 5] + + field_to_copy[1, 0] = field_to_copy[0, 4] + field_to_copy[1, 1] = field_to_copy[1, 4] + field_to_copy[1, 2] = field_to_copy[2, 4] + + field_to_copy[2, 0] = field_to_copy[0, 3] + field_to_copy[2, 1] = field_to_copy[1, 3] + field_to_copy[2, 2] = field_to_copy[2, 3] + + field_to_copy[0, -4] = field_to_copy[2, -7] + field_to_copy[0, -3] = field_to_copy[1, -7] + field_to_copy[0, -2] = field_to_copy[0, -7] + + field_to_copy[1, -4] = field_to_copy[2, -6] + field_to_copy[1, -3] = field_to_copy[1, -6] + field_to_copy[1, -2] = field_to_copy[0, -6] + + field_to_copy[2, -4] = field_to_copy[2, -5] + field_to_copy[2, -3] = field_to_copy[1, -5] + field_to_copy[2, -2] = field_to_copy[0, -5] + + field_to_copy[-4, 0] = field_to_copy[-2, 3] + field_to_copy[-4, 1] = field_to_copy[-3, 3] + field_to_copy[-4, 2] = field_to_copy[-4, 3] + + field_to_copy[-3, 0] = field_to_copy[-2, 4] + field_to_copy[-3, 1] = field_to_copy[-3, 4] + field_to_copy[-3, 2] = field_to_copy[-4, 4] + + field_to_copy[-2, 0] = field_to_copy[-2, 5] + field_to_copy[-2, 1] = field_to_copy[-3, 5] + field_to_copy[-2, 2] = field_to_copy[-4, 5] + + field_to_copy[-4, -2] = field_to_copy[-2, -5] + field_to_copy[-4, -3] = field_to_copy[-3, -5] + field_to_copy[-4, -4] = field_to_copy[-4, -5] + + field_to_copy[-3, -2] = field_to_copy[-2, -6] + field_to_copy[-3, -3] = field_to_copy[-3, -6] + field_to_copy[-3, -4] = field_to_copy[-4, -6] + + field_to_copy[-2, -2] = field_to_copy[-2, -7] + field_to_copy[-2, -3] = field_to_copy[-3, -7] + field_to_copy[-2, -4] = field_to_copy[-4, -7] + + +def corner_copy_y(field_to_copy): + """Equivalent to the copy_corners_y functions in fortran. + + This is written to operate on plain ndarrarys and not use the GT4Py framework. + This choice was made because we've seen a lot of performance left on the table using + orchestration without explicitly describing the operations but rather have full 3d- + sweeps with conditionals. + Since DaCe can handle (simple) operations on ndarrays directly this gives us a more + explicit entrypoint to the language and more optimization-potential. + + Args: + field_to_copy (ndarray): field to apply the corner copy on. + This is explicitly not type-hinted for orchestration + """ + field_to_copy[0, 0] = field_to_copy[5, 0] + field_to_copy[1, 0] = field_to_copy[5, 1] + field_to_copy[2, 0] = field_to_copy[5, 2] + + field_to_copy[0, 1] = field_to_copy[4, 0] + field_to_copy[1, 1] = field_to_copy[4, 1] + field_to_copy[2, 1] = field_to_copy[4, 2] + + field_to_copy[0, 2] = field_to_copy[3, 0] + field_to_copy[1, 2] = field_to_copy[3, 1] + field_to_copy[2, 2] = field_to_copy[3, 2] + + field_to_copy[-4, 0] = field_to_copy[-7, 2] + field_to_copy[-3, 0] = field_to_copy[-7, 1] + field_to_copy[-2, 0] = field_to_copy[-7, 0] + + field_to_copy[-4, 1] = field_to_copy[-6, 2] + field_to_copy[-3, 1] = field_to_copy[-6, 1] + field_to_copy[-2, 1] = field_to_copy[-6, 0] + + field_to_copy[-4, 2] = field_to_copy[-5, 2] + field_to_copy[-3, 2] = field_to_copy[-5, 1] + field_to_copy[-2, 2] = field_to_copy[-5, 0] + + field_to_copy[0, -2] = field_to_copy[5, -2] + field_to_copy[0, -3] = field_to_copy[4, -2] + field_to_copy[0, -4] = field_to_copy[3, -2] + + field_to_copy[1, -2] = field_to_copy[5, -3] + field_to_copy[1, -3] = field_to_copy[4, -3] + field_to_copy[1, -4] = field_to_copy[3, -3] + + field_to_copy[2, -2] = field_to_copy[5, -4] + field_to_copy[2, -3] = field_to_copy[4, -4] + field_to_copy[2, -4] = field_to_copy[3, -4] + + field_to_copy[-2, -4] = field_to_copy[-5, -2] + field_to_copy[-2, -3] = field_to_copy[-6, -2] + field_to_copy[-2, -2] = field_to_copy[-7, -2] + + field_to_copy[-3, -4] = field_to_copy[-5, -3] + field_to_copy[-3, -3] = field_to_copy[-6, -3] + field_to_copy[-3, -2] = field_to_copy[-7, -3] + + field_to_copy[-4, -4] = field_to_copy[-5, -4] + field_to_copy[-4, -3] = field_to_copy[-6, -4] + field_to_copy[-4, -2] = field_to_copy[-7, -4] + + +class CopyCornersX: + """ + Helper-class to copy corners corresponding to the fortran function copy_corners_x + """ + + def __init__(self, stencil_factory: StencilFactory) -> None: + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + ) + + if stencil_factory.grid_indexing.n_halo != 3: + raise NotImplementedError( + "Corner-Copy only implemented for exactly 3 Halo-Points" + ) + + def __call__(self, field: FloatField): + corner_copy_x(field) + + +class CopyCornersY: + """ + Helper-class to copy corners corresponding to the fortran function + copy_corners_y + """ + + def __init__(self, stencil_factory: StencilFactory) -> None: + orchestrate( + obj=self, + config=stencil_factory.config.dace_config, + ) + + if stencil_factory.grid_indexing.n_halo != 3: + raise NotImplementedError( + "Corner-Copy only implemented for exactly 3 Halo-Points" + ) + + def __call__(self, field: FloatField): + corner_copy_y(field) diff --git a/pyfv3/stencils/delnflux.py b/pyfv3/stencils/delnflux.py index f80d210b..6384f9ce 100644 --- a/pyfv3/stencils/delnflux.py +++ b/pyfv3/stencils/delnflux.py @@ -1,5 +1,7 @@ from typing import Optional +import dace + from ndsl import Quantity, QuantityFactory, StencilFactory, orchestrate from ndsl.constants import X_DIM, X_INTERFACE_DIM, Y_DIM, Y_INTERFACE_DIM, Z_DIM from ndsl.dsl.gt4py import PARALLEL, computation @@ -8,6 +10,7 @@ from ndsl.dsl.stencil import get_stencils_with_varied_bounds from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ, FloatFieldK from ndsl.grid import DampingCoefficients +from pyfv3.stencils.copy_corners import corner_copy_x, corner_copy_y def calc_damp(damp_c: Quantity, da_min: Float, nord: Quantity) -> Quantity: @@ -120,7 +123,7 @@ def d2_highorder_stencil( ): with computation(PARALLEL), interval(...): if nord > current_nord: - d2 = ((fx - fx[1, 0, 0]) + (fy - fy[0, 1, 0])) * rarea + d2 = (fx - fx[1, 0, 0] + fy - fy[0, 1, 0]) * rarea def d2_damp_interval( @@ -182,176 +185,16 @@ def diffusive_damp( fy = fy + 0.5 * damp * (mass[0, -1, 0] + mass) * fy2 -def copy_corners_y_nord( - q_in: FloatField, q_out: FloatField, nord: FloatFieldK, current_nord: int -): - """ - Args: - q_in (in): - q_out (out): - """ - from __externals__ import i_end, i_start, j_end, j_start - - with computation(PARALLEL), interval(...): - if nord > current_nord: - with horizontal( - region[i_start - 3, j_start - 3], region[i_start - 3, j_end + 3] - ): - q_out = q_in[5, 0, 0] - with horizontal( - region[i_start - 2, j_start - 3], region[i_start - 3, j_end + 2] - ): - q_out = q_in[4, 1, 0] - with horizontal( - region[i_start - 1, j_start - 3], region[i_start - 3, j_end + 1] - ): - q_out = q_in[3, 2, 0] - with horizontal( - region[i_start - 3, j_start - 2], region[i_start - 2, j_end + 3] - ): - q_out = q_in[4, -1, 0] - with horizontal( - region[i_start - 2, j_start - 2], region[i_start - 2, j_end + 2] - ): - q_out = q_in[3, 0, 0] - with horizontal( - region[i_start - 1, j_start - 2], region[i_start - 2, j_end + 1] - ): - q_out = q_in[2, 1, 0] - with horizontal( - region[i_start - 3, j_start - 1], region[i_start - 1, j_end + 3] - ): - q_out = q_in[3, -2, 0] - with horizontal( - region[i_start - 2, j_start - 1], region[i_start - 1, j_end + 2] - ): - q_out = q_in[2, -1, 0] - with horizontal( - region[i_start - 1, j_start - 1], region[i_start - 1, j_end + 1] - ): - q_out = q_in[1, 0, 0] - with horizontal( - region[i_end + 1, j_start - 3], region[i_end + 3, j_end + 1] - ): - q_out = q_in[-3, 2, 0] - with horizontal( - region[i_end + 2, j_start - 3], region[i_end + 3, j_end + 2] - ): - q_out = q_in[-4, 1, 0] - with horizontal( - region[i_end + 3, j_start - 3], region[i_end + 3, j_end + 3] - ): - q_out = q_in[-5, 0, 0] - with horizontal( - region[i_end + 1, j_start - 2], region[i_end + 2, j_end + 1] - ): - q_out = q_in[-2, 1, 0] - with horizontal( - region[i_end + 2, j_start - 2], region[i_end + 2, j_end + 2] - ): - q_out = q_in[-3, 0, 0] - with horizontal( - region[i_end + 3, j_start - 2], region[i_end + 2, j_end + 3] - ): - q_out = q_in[-4, -1, 0] - with horizontal( - region[i_end + 1, j_start - 1], region[i_end + 1, j_end + 1] - ): - q_out = q_in[-1, 0, 0] - with horizontal( - region[i_end + 2, j_start - 1], region[i_end + 1, j_end + 2] - ): - q_out = q_in[-2, -1, 0] - with horizontal( - region[i_end + 3, j_start - 1], region[i_end + 1, j_end + 3] - ): - q_out = q_in[-3, -2, 0] - +def copy_corners_y_nord(field_to_copy, nord): + for k in dace.map[0 : nord.data.shape[0]]: + if nord.data[k] > 0: + corner_copy_y(field_to_copy[:, :, k]) -def copy_corners_x_nord( - q_in: FloatField, q_out: FloatField, nord: FloatFieldK, current_nord: int -): - """ - Args: - q_in (in): - q_out (out): - """ - from __externals__ import i_end, i_start, j_end, j_start - with computation(PARALLEL), interval(...): - if nord > current_nord: - with horizontal( - region[i_start - 3, j_start - 3], region[i_end + 3, j_start - 3] - ): - q_out = q_in[0, 5, 0] - with horizontal( - region[i_start - 2, j_start - 3], region[i_end + 3, j_start - 2] - ): - q_out = q_in[-1, 4, 0] - with horizontal( - region[i_start - 1, j_start - 3], region[i_end + 3, j_start - 1] - ): - q_out = q_in[-2, 3, 0] - with horizontal( - region[i_start - 3, j_start - 2], region[i_end + 2, j_start - 3] - ): - q_out = q_in[1, 4, 0] - with horizontal( - region[i_start - 2, j_start - 2], region[i_end + 2, j_start - 2] - ): - q_out = q_in[0, 3, 0] - with horizontal( - region[i_start - 1, j_start - 2], region[i_end + 2, j_start - 1] - ): - q_out = q_in[-1, 2, 0] - with horizontal( - region[i_start - 3, j_start - 1], region[i_end + 1, j_start - 3] - ): - q_out = q_in[2, 3, 0] - with horizontal( - region[i_start - 2, j_start - 1], region[i_end + 1, j_start - 2] - ): - q_out = q_in[1, 2, 0] - with horizontal( - region[i_start - 1, j_start - 1], region[i_end + 1, j_start - 1] - ): - q_out = q_in[0, 1, 0] - with horizontal( - region[i_start - 3, j_end + 1], region[i_end + 1, j_end + 3] - ): - q_out = q_in[2, -3, 0] - with horizontal( - region[i_start - 2, j_end + 1], region[i_end + 1, j_end + 2] - ): - q_out = q_in[1, -2, 0] - with horizontal( - region[i_start - 1, j_end + 1], region[i_end + 1, j_end + 1] - ): - q_out = q_in[0, -1, 0] - with horizontal( - region[i_start - 3, j_end + 2], region[i_end + 2, j_end + 3] - ): - q_out = q_in[1, -4, 0] - with horizontal( - region[i_start - 2, j_end + 2], region[i_end + 2, j_end + 2] - ): - q_out = q_in[0, -3, 0] - with horizontal( - region[i_start - 1, j_end + 2], region[i_end + 2, j_end + 1] - ): - q_out = q_in[-1, -2, 0] - with horizontal( - region[i_start - 3, j_end + 3], region[i_end + 3, j_end + 3] - ): - q_out = q_in[0, -5, 0] - with horizontal( - region[i_start - 2, j_end + 3], region[i_end + 3, j_end + 2] - ): - q_out = q_in[-1, -4, 0] - with horizontal( - region[i_start - 1, j_end + 3], region[i_end + 3, j_end + 1] - ): - q_out = q_in[-2, -3, 0] +def copy_corners_x_nord(field_to_copy, nord): + for k in dace.map[0 : nord.data.shape[0]]: + if nord.data[k] > 0: + corner_copy_x(field_to_copy[:, :, k]) class DelnFlux: @@ -599,27 +442,6 @@ def __init__( origin=fx_origin, domain=(f1_nx - 1, f1_ny + 1, nk), ) - corner_origin, corner_domain = grid_indexing.get_origin_domain( - dims=[X_DIM, Y_DIM, Z_DIM], - halos=(grid_indexing.n_halo, grid_indexing.n_halo), - ) - corner_domain = corner_domain[:2] + (nk,) - corner_axis_offsets = grid_indexing.axis_offsets(corner_origin, corner_domain) - - self._copy_corners_x_nord = stencil_factory.from_origin_domain( - copy_corners_x_nord, - externals={**corner_axis_offsets}, - origin=corner_origin, - domain=corner_domain, - skip_passes=("UnreachableStmtPruning",), - ) - self._copy_corners_y_nord = stencil_factory.from_origin_domain( - copy_corners_y_nord, - externals={**corner_axis_offsets}, - origin=corner_origin, - domain=corner_domain, - skip_passes=("UnreachableStmtPruning",), - ) def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): """ @@ -643,11 +465,11 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): else: self._copy_stencil_interval(q_in=q, q_out=d2, nord=self._nord) - self._copy_corners_x_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + copy_corners_x_nord(d2.data, self._nord) self._fx_calc_stencil(q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord) - self._copy_corners_y_nord(q_in=d2, q_out=d2, nord=self._nord, current_nord=0) + copy_corners_y_nord(d2.data, self._nord) self._fy_calc_stencil(q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord) @@ -661,17 +483,13 @@ def __call__(self, q, fx2, fy2, damp_c, d2, mass=None): current_nord=n, ) - self._copy_corners_x_nord( - q_in=d2, q_out=d2, nord=self._nord, current_nord=n - ) + copy_corners_x_nord(d2.data, self._nord) self._column_conditional_fx_calculation[n]( q=d2, del6_v=self._del6_v, fx=fx2, nord=self._nord, current_nord=n ) - self._copy_corners_y_nord( - q_in=d2, q_out=d2, nord=self._nord, current_nord=n - ) + copy_corners_y_nord(d2.data, self._nord) self._column_conditional_fy_calculation[n]( q=d2, del6_u=self._del6_u, fy=fy2, nord=self._nord, current_nord=n diff --git a/pyfv3/stencils/fvtp2d.py b/pyfv3/stencils/fvtp2d.py index 9c20ff49..83e2e64d 100644 --- a/pyfv3/stencils/fvtp2d.py +++ b/pyfv3/stencils/fvtp2d.py @@ -7,7 +7,7 @@ from ndsl.dsl.gt4py import horizontal, interval, region from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from ndsl.grid import DampingCoefficients, GridData -from ndsl.stencils import corners +from pyfv3.stencils.copy_corners import CopyCornersX, CopyCornersY from pyfv3.stencils.delnflux import DelnFlux from pyfv3.stencils.xppm import XPiecewiseParabolic from pyfv3.stencils.yppm import YPiecewiseParabolic @@ -178,7 +178,7 @@ def make_quantity(): # self.delnflux = None self._do_delnflux = False - self._copy_corners_y = corners.CopyCorners("y", stencil_factory) + self._copy_corners_y = CopyCornersY(stencil_factory) self.y_piecewise_parabolic_inner = YPiecewiseParabolic( stencil_factory=stencil_factory, dya=grid_data.dya, @@ -201,7 +201,7 @@ def make_quantity(): domain=idx.domain_compute(add=(1, 1, 1)), ) - self._copy_corners_x = corners.CopyCorners("x", stencil_factory) + self._copy_corners_x = CopyCornersX(stencil_factory) self.x_piecewise_parabolic_inner = XPiecewiseParabolic( stencil_factory=stencil_factory, dxa=grid_data.dxa, @@ -313,7 +313,7 @@ def __call__( # y_area_flux as an input (flux = area_flux * advected_mean), since a flux is # easier to understand than the current output. This would be like merging # yppm with q_i_stencil and xppm with q_j_stencil. - self._copy_corners_y(q) + self._copy_corners_y(q.data) self.y_piecewise_parabolic_inner(q, cry, self._q_y_advected_mean) # q_y_advected_mean is 1/Delta_area * curly-F, where curly-F is defined in # equation 4.3 of the FV3 documentation and Delta_area is the advected area @@ -330,7 +330,7 @@ def __call__( ) # q_advected_y_x_advected_mean is now rho^n + F(rho^y) in PL07 eq 16 - self._copy_corners_x(q) + self._copy_corners_x(q.data) # similarly below for x<->y self.x_piecewise_parabolic_inner(q, crx, self._q_x_advected_mean) self.q_j_stencil( diff --git a/tests/savepoint/translate/translate_corners.py b/tests/savepoint/translate/translate_corners.py index 57ce464a..74916a6a 100644 --- a/tests/savepoint/translate/translate_corners.py +++ b/tests/savepoint/translate/translate_corners.py @@ -3,6 +3,7 @@ import ndsl.dsl.gt4py_utils as utils from ndsl import Namelist, StencilFactory from ndsl.stencils import corners +from pyfv3.stencils.copy_corners import CopyCornersX, CopyCornersY from pyfv3.testing import TranslateDycoreFortranData2Py @@ -100,8 +101,8 @@ def __init__( self.in_vars["data_vars"] = {"q": {}} self.in_vars["parameters"] = ["dir"] self.out_vars: Dict[str, Any] = {"q": {}} - self._copy_corners_x = corners.CopyCorners("x", stencil_factory=stencil_factory) - self._copy_corners_y = corners.CopyCorners("y", stencil_factory=stencil_factory) + self._copy_corners_x = CopyCornersX(stencil_factory) + self._copy_corners_y = CopyCornersY(stencil_factory) self.stencil_factory = stencil_factory def compute_from_storage(self, inputs): diff --git a/tests/savepoint/translate/translate_d_sw.py b/tests/savepoint/translate/translate_d_sw.py index c30798b5..63e705f7 100644 --- a/tests/savepoint/translate/translate_d_sw.py +++ b/tests/savepoint/translate/translate_d_sw.py @@ -3,7 +3,8 @@ import pyfv3 import pyfv3.stencils.d_sw as d_sw from ndsl import Namelist, StencilFactory -from ndsl.dsl.typing import FloatField, FloatFieldIJ +from ndsl.constants import X_DIM, Y_DIM, Z_INTERFACE_DIM +from ndsl.dsl.typing import Float, FloatField, FloatFieldIJ from pyfv3.testing import TranslateDycoreFortranData2Py @@ -63,6 +64,34 @@ def __init__( self.out_vars = self.in_vars["data_vars"].copy() del self.out_vars["zh"] + def compute(self, inputs): + self.make_storage_data_input_vars(inputs) + # Convert relevant inputs to quantities: + delp = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], units="unknown", dtype=Float + ) + delp.data[:] = delp.np.asarray(inputs.pop("delp")) + inputs["delp"] = delp + w = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], units="unknown", dtype=Float + ) + w.data[:] = delp.np.asarray(inputs.pop("w")) + inputs["w"] = w + q_con = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], units="unknown", dtype=Float + ) + q_con.data[:] = delp.np.asarray(inputs.pop("q_con")) + inputs["q_con"] = q_con + + pt = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], units="unknown", dtype=Float + ) + pt.data[:] = delp.np.asarray(inputs.pop("pt")) + inputs["pt"] = pt + + self.compute_func(**inputs) + return self.slice_output(inputs) + def ubke( uc: FloatField, diff --git a/tests/savepoint/translate/translate_del6vtflux.py b/tests/savepoint/translate/translate_del6vtflux.py index f8706317..dd922820 100644 --- a/tests/savepoint/translate/translate_del6vtflux.py +++ b/tests/savepoint/translate/translate_del6vtflux.py @@ -1,6 +1,7 @@ import pyfv3.stencils.delnflux as delnflux from ndsl import Namelist, StencilFactory -from ndsl.constants import Z_DIM +from ndsl.constants import X_DIM, Y_DIM, Z_DIM, Z_INTERFACE_DIM +from ndsl.dsl.typing import Float from pyfv3.testing import TranslateDycoreFortranData2Py @@ -44,5 +45,13 @@ def compute(self, inputs): self.grid.rarea, nord_col, ) + + # Convert relevant inputs to quantities: + d2 = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_INTERFACE_DIM], units="unknown", dtype=Float + ) + d2.data[:] = d2.np.asarray(inputs.pop("d2")) + inputs["d2"] = d2 + self.compute_func(**inputs) return self.slice_output(inputs) diff --git a/tests/savepoint/translate/translate_fvtp2d.py b/tests/savepoint/translate/translate_fvtp2d.py index ef392fa5..98d0e8ae 100644 --- a/tests/savepoint/translate/translate_fvtp2d.py +++ b/tests/savepoint/translate/translate_fvtp2d.py @@ -1,6 +1,6 @@ import ndsl.dsl.gt4py_utils as utils from ndsl import Namelist, StencilFactory -from ndsl.constants import Z_DIM +from ndsl.constants import X_DIM, Y_DIM, Z_DIM from ndsl.dsl.typing import Float from pyfv3.stencils import FiniteVolumeTransport from pyfv3.testing import TranslateDycoreFortranData2Py @@ -59,6 +59,12 @@ def compute_from_storage(self, inputs): dims=[Z_DIM], units="unknown", dtype=Float ) damp_c.data[:] = damp_c.np.asarray(inputs.pop("damp_c")) + + q = self.grid.quantity_factory.zeros( + dims=[X_DIM, Y_DIM, Z_DIM], units="unknown", dtype=Float + ) + q.data[:] = q.np.asarray(inputs.pop("q")) + inputs["q"] = q for optional_arg in ["mass"]: if optional_arg not in inputs: inputs[optional_arg] = None