diff --git a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py index b6c10acf85..a4f80bc943 100644 --- a/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py +++ b/model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py @@ -40,6 +40,32 @@ def _compute_nbidx( return nbidx[jc, ind, :] +def _compute_z_vintcoeff( + k_start: list, + k_end: list, + z_mc: np.array, + z_mc_off: np.array, + z_vintcoeff: np.array, + jc: int, + ind: int, + nlev: int, +) -> np.array: + jk_start = nlev - 2 + for jk in reversed(range(k_start[jc], k_end[jc])): + for jk1 in reversed(range(jk_start + 1)): + if ( + z_mc[jc, jk] <= z_mc_off[jc, ind, jk1] + and z_mc[jc, jk] >= z_mc_off[jc, ind, jk1 + 1] + ): + z_vintcoeff[jc, ind, jk] = (z_mc[jc, jk] - z_mc_off[jc, ind, jk1 + 1]) / ( + z_mc_off[jc, ind, jk1] - z_mc_off[jc, ind, jk1 + 1] + ) + jk_start = jk1 + break + + return z_vintcoeff[jc, ind, :] + + def _compute_i_params( k_start: list, k_end: list, @@ -109,15 +135,23 @@ def _compute_k_start_end( return k_start, k_end -def _compute_zd_vertoffset_dsl( - k_start: list, - k_end: list, +def compute_diffusion_metrics( z_mc: np.array, z_mc_off: np.array, - nbidx: np.array, + k_start: list, + k_end: list, i_indlist: list, i_listdim: int, + nbidx: np.array, + z_vintcoeff: np.array, + z_maxslp_avg: np.array, + z_maxhgtd_avg: np.array, + mask_hdiff: np.array, + zd_diffcoef_dsl: np.array, + zd_intcoef_dsl: np.array, zd_vertoffset_dsl: np.array, + thslp_zdiffu: float, + thhgtd_zdiffu: float, nlev: int, ) -> np.array: for ji in range(i_listdim): @@ -126,41 +160,25 @@ def _compute_zd_vertoffset_dsl( nbidx[jc, 0, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 0, nlev) nbidx[jc, 1, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 1, nlev) nbidx[jc, 2, :] = _compute_nbidx(k_start, k_end, z_mc, z_mc_off, nbidx, jc, 2, nlev) + + z_vintcoeff[jc, 0, :] = _compute_z_vintcoeff( + k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 0, nlev + ) + z_vintcoeff[jc, 1, :] = _compute_z_vintcoeff( + k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 1, nlev + ) + z_vintcoeff[jc, 2, :] = _compute_z_vintcoeff( + k_start, k_end, z_mc, z_mc_off, z_vintcoeff, jc, 2, nlev + ) for jk in range(k_start[jc], k_end[jc]): + zd_intcoef_dsl[jc, 0, jk] = z_vintcoeff[jc, 0, jk] + zd_intcoef_dsl[jc, 1, jk] = z_vintcoeff[jc, 1, jk] + zd_intcoef_dsl[jc, 2, jk] = z_vintcoeff[jc, 2, jk] + zd_vertoffset_dsl[jc, 0, jk] = nbidx[jc, 0, jk] - jk zd_vertoffset_dsl[jc, 1, jk] = nbidx[jc, 1, jk] - jk zd_vertoffset_dsl[jc, 2, jk] = nbidx[jc, 2, jk] - jk - return zd_vertoffset_dsl - - -def _compute_mask_hdiff( - mask_hdiff: np.array, k_start: list, k_end: list, i_indlist: list, i_listdim: int -) -> np.array: - for ji in range(i_listdim): - jc = i_indlist[ji] - if k_start[jc] is not None and k_end[jc] is not None: - for jk in range(k_start[jc], k_end[jc]): - mask_hdiff[jc, jk] = True - - return mask_hdiff - - -def _compute_zd_diffcoef_dsl( - z_maxslp_avg: np.array, - z_maxhgtd_avg: np.array, - k_start: list, - k_end: list, - i_indlist: list, - i_listdim: int, - zd_diffcoef_dsl: np.array, - thslp_zdiffu: float, - thhgtd_zdiffu: float, -) -> np.array: - for ji in range(i_listdim): - jc = i_indlist[ji] - if k_start[jc] is not None and k_end[jc] is not None: - for jk in range(k_start[jc], k_end[jc]): zd_diffcoef_dsl_var = max( 0.0, math.sqrt(max(0.0, z_maxslp_avg[jc, jk] - thslp_zdiffu)) / 250.0, @@ -169,4 +187,6 @@ def _compute_zd_diffcoef_dsl( zd_diffcoef_dsl_var = min(0.002, zd_diffcoef_dsl_var) zd_diffcoef_dsl[jc, jk] = zd_diffcoef_dsl_var - return zd_diffcoef_dsl + mask_hdiff[jc, jk] = True + + return mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl diff --git a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py index d992a448b1..21f292bf02 100644 --- a/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py +++ b/model/common/src/icon4py/model/common/test_utils/serialbox_utils.py @@ -683,7 +683,7 @@ def wgtfacq_e_dsl( def zd_diffcoef(self): return self._get_field("zd_diffcoef", CellDim, KDim) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) def zd_intcoef(self): return self._read_and_reorder_sparse_field("vcoef") @@ -692,6 +692,7 @@ def geopot(self): def _read_and_reorder_sparse_field(self, name: str, sparse_size=3): ser_input = np.squeeze(self.serializer.read(name, self.savepoint))[:, :, :] + ser_input = self._reduce_to_dim_size(ser_input, (CellDim, C2E2CDim, KDim)) if ser_input.shape[1] != sparse_size: ser_input = np.moveaxis(ser_input, 1, -1) @@ -706,10 +707,9 @@ def _linearize_first_2dims( assert old_shape[1] == sparse_size return as_field(target_dims, data.reshape(old_shape[0] * old_shape[1], old_shape[2])) - @IconSavepoint.optionally_registered() + @IconSavepoint.optionally_registered(CellDim, C2E2CDim, KDim) def zd_vertoffset(self): - # return self._read_and_reorder_sparse_field("zd_vertoffset") - return self._get_field("zd_vertoffset", CellDim, C2E2CDim, KDim) + return self._read_and_reorder_sparse_field("zd_vertoffset") def zd_vertidx(self): return np.squeeze(self.serializer.read("zd_vertidx", self.savepoint)) diff --git a/model/common/tests/metric_tests/test_metric_fields.py b/model/common/tests/metric_tests/test_metric_fields.py index 960e56aff1..70bd4e4d99 100644 --- a/model/common/tests/metric_tests/test_metric_fields.py +++ b/model/common/tests/metric_tests/test_metric_fields.py @@ -21,6 +21,7 @@ from icon4py.model.common import constants from icon4py.model.common.dimension import ( C2E2CDim, + CECDim, CellDim, E2CDim, ECDim, @@ -68,9 +69,7 @@ from icon4py.model.common.metrics.stencils.compute_diffusion_metrics import ( _compute_i_params, _compute_k_start_end, - _compute_mask_hdiff, - _compute_zd_diffcoef_dsl, - _compute_zd_vertoffset_dsl, + compute_diffusion_metrics, ) from icon4py.model.common.metrics.stencils.compute_zdiff_gradp_dsl import compute_zdiff_gradp_dsl from icon4py.model.common.test_utils.datatest_utils import ( @@ -910,32 +909,26 @@ def test_compute_diffusion_metrics( backend = None mask_hdiff = zero_field(icon_grid, CellDim, KDim, dtype=bool).asnumpy() zd_vertoffset_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() - nlev = icon_grid.num_levels - mask_hdiff_ref = metrics_savepoint.mask_hdiff() - zd_vertoffset_ref = metrics_savepoint.zd_vertoffset() - # TODO: inquire to Magda as to why 3316 works instead of cell_nudging - # cell_nudging = icon_grid.get_end_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) - cell_nudging = 3316 - cell_lateral = icon_grid.get_start_index( - CellDim, - HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, - ) + z_vintcoeff = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() + zd_intcoef_dsl = zero_field(icon_grid, CellDim, C2E2CDim, KDim).asnumpy() z_maxslp_avg = zero_field(icon_grid, CellDim, KDim) z_maxhgtd_avg = zero_field(icon_grid, CellDim, KDim) zd_diffcoef_dsl = zero_field(icon_grid, CellDim, KDim).asnumpy() - c_bln_avg = interpolation_savepoint.c_bln_avg() - thslp_zdiffu = ( - 0.02 # TODO: import from Fortran, note: same variable in diffusion has different value - ) - thhgtd_zdiffu = ( - 125 # TODO: import from Fortran, note: same variable in diffusion has different value - ) - c_owner_mask = grid_savepoint.c_owner_mask().asnumpy() maxslp = zero_field(icon_grid, CellDim, KDim) maxhgtd = zero_field(icon_grid, CellDim, KDim) max_nbhgt = zero_field(icon_grid, CellDim) + c2e2c = icon_grid.connectivities[C2E2CDim] nbidx = constant_field(icon_grid, 1, CellDim, C2E2CDim, KDim, dtype=int).asnumpy() + c_bln_avg = interpolation_savepoint.c_bln_avg() + thslp_zdiffu = 0.02 + thhgtd_zdiffu = 125 + cell_nudging = icon_grid.get_start_index(CellDim, HorizontalMarkerIndex.nudging(CellDim)) + cell_lateral = icon_grid.get_start_index( + CellDim, + HorizontalMarkerIndex.lateral_boundary(CellDim) + 1, + ) + nlev = icon_grid.num_levels _compute_maxslp_maxhgtd( ddxn_z_full=metrics_savepoint.ddxn_z_full(), @@ -945,10 +938,9 @@ def test_compute_diffusion_metrics( offset_provider={"C2E": icon_grid.get_offset_provider("C2E")}, ) - z_ifc = metrics_savepoint.z_ifc() z_mc = zero_field(icon_grid, CellDim, KDim, extend={KDim: 1}) compute_z_mc.with_backend(backend)( - z_ifc, + metrics_savepoint.z_ifc(), z_mc, horizontal_start=int32(0), horizontal_end=icon_grid.num_cells, @@ -956,7 +948,6 @@ def test_compute_diffusion_metrics( vertical_end=nlev, offset_provider={"Koff": icon_grid.get_offset_provider("Koff")}, ) - z_mc_off = z_mc.asnumpy()[c2e2c] compute_z_maxslp_avg_z_maxhgtd_avg.with_backend(backend)( maxslp=maxslp, @@ -982,56 +973,60 @@ def test_compute_diffusion_metrics( ) k_start, k_end = _compute_k_start_end( - z_mc.asnumpy(), - max_nbhgt.asnumpy(), - z_maxslp_avg.asnumpy(), - z_maxhgtd_avg.asnumpy(), - c_owner_mask, - thslp_zdiffu, - thhgtd_zdiffu, - cell_nudging, - icon_grid.num_cells, - nlev, + z_mc=z_mc.asnumpy(), + max_nbhgt=max_nbhgt.asnumpy(), + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=icon_grid.num_cells, + nlev=nlev, ) i_indlist, i_listreduce, ji = _compute_i_params( - k_start, - k_end, - z_maxslp_avg.asnumpy(), - z_maxhgtd_avg.asnumpy(), - c_owner_mask, - thslp_zdiffu, - thhgtd_zdiffu, - cell_nudging, - icon_grid.num_cells, - nlev, + k_start=k_start, + k_end=k_end, + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + c_owner_mask=grid_savepoint.c_owner_mask().asnumpy(), + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + cell_nudging=cell_nudging, + n_cells=icon_grid.num_cells, + nlev=nlev, ) i_listdim = ji - i_listreduce - mask_hdiff = _compute_mask_hdiff(mask_hdiff, k_start, k_end, i_indlist, i_listdim) - zd_diffcoef_dsl = _compute_zd_diffcoef_dsl( - z_maxslp_avg.asnumpy(), - z_maxhgtd_avg.asnumpy(), - k_start, - k_end, - i_indlist, - i_listdim, - zd_diffcoef_dsl, - thslp_zdiffu, - thhgtd_zdiffu, - ) - zd_vertoffset_dsl = _compute_zd_vertoffset_dsl( - k_start, - k_end, - z_mc.asnumpy(), - z_mc_off, - nbidx, - i_indlist, - i_listdim, - zd_vertoffset_dsl, - nlev, - ) - - assert dallclose(mask_hdiff, mask_hdiff_ref.asnumpy()) + + mask_hdiff, zd_diffcoef_dsl, zd_intcoef_dsl, zd_vertoffset_dsl = compute_diffusion_metrics( + z_mc=z_mc.asnumpy(), + z_mc_off=z_mc.asnumpy()[c2e2c], + k_start=k_start, + k_end=k_end, + i_indlist=i_indlist, + i_listdim=i_listdim, + nbidx=nbidx, + z_vintcoeff=z_vintcoeff, + z_maxslp_avg=z_maxslp_avg.asnumpy(), + z_maxhgtd_avg=z_maxhgtd_avg.asnumpy(), + mask_hdiff=mask_hdiff, + zd_diffcoef_dsl=zd_diffcoef_dsl, + zd_intcoef_dsl=zd_intcoef_dsl, + zd_vertoffset_dsl=zd_vertoffset_dsl, + thslp_zdiffu=thslp_zdiffu, + thhgtd_zdiffu=thhgtd_zdiffu, + nlev=nlev, + ) + zd_intcoef_dsl = flatten_first_two_dims( + CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_intcoef_dsl) + ) + zd_vertoffset_dsl = flatten_first_two_dims( + CECDim, KDim, field=as_field((CellDim, C2E2CDim, KDim), zd_vertoffset_dsl) + ) + + assert dallclose(mask_hdiff, metrics_savepoint.mask_hdiff().asnumpy()) assert dallclose(zd_diffcoef_dsl, metrics_savepoint.zd_diffcoef().asnumpy(), rtol=1.0e-11) - assert dallclose(zd_vertoffset_dsl, zd_vertoffset_ref.asnumpy()) + assert dallclose(zd_vertoffset_dsl.asnumpy(), metrics_savepoint.zd_vertoffset().asnumpy()) + assert dallclose(zd_intcoef_dsl.asnumpy(), metrics_savepoint.zd_intcoef())