Skip to content

Commit

Permalink
finishes
Browse files Browse the repository at this point in the history
  • Loading branch information
@nfarabullini
nfarabullini committed May 29, 2024
1 parent cb990b8 commit c1a6343
Show file tree
Hide file tree
Showing 3 changed files with 124 additions and 109 deletions.
90 changes: 55 additions & 35 deletions model/common/src/icon4py/model/common/metrics/stencils/compute_diffusion_metrics.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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):
Expand All @@ -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,
Expand All @@ -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
8 changes: 4 additions & 4 deletions model/common/src/icon4py/model/common/test_utils/serialbox_utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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")

Expand All @@ -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)

Expand All @@ -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))
Expand Down
135 changes: 65 additions & 70 deletions model/common/tests/metric_tests/test_metric_fields.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -21,6 +21,7 @@
from icon4py.model.common import constants
from icon4py.model.common.dimension import (
C2E2CDim,
CECDim,
CellDim,
E2CDim,
ECDim,
Expand Down Expand Up @@ -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 (
Expand Down Expand Up @@ -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(),
Expand All @@ -945,18 +938,16 @@ 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,
vertical_start=int32(0),
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,
Expand All @@ -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())

0 comments on commit c1a6343

Please sign in to comment.