Skip to content
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.

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

Refactor advection stencils to use StencilTest #336

Merged
merged 27 commits into from
Jan 17, 2024
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
Show all changes
27 commits
Select commit Hold shift + click to select a range
0dbb5ea
Refactor a few advection stencils
samkellerhals Dec 13, 2023
0fd8ea4
Disable C2V offset provider due to missing connectivity
samkellerhals Dec 14, 2023
3a5f2e5
Update with main
samkellerhals Jan 8, 2024
a8a622f
Allow comparing subset of output field
samkellerhals Jan 8, 2024
c983e7c
port btraj_dreg_03
samkellerhals Jan 8, 2024
7a028fc
Add face_val_ppm_stencil_02
samkellerhals Jan 8, 2024
00bc933
Add face_val_ppm_02a
samkellerhals Jan 8, 2024
db284fe
Add Output helper class
samkellerhals Jan 9, 2024
ebcee6a
Add more stencils and reshape helper
samkellerhals Jan 9, 2024
d89a007
Fix test
samkellerhals Jan 9, 2024
d0b14dc
Add hflx limiter 2
samkellerhals Jan 10, 2024
afaba3f
Add hflx limiter 4
samkellerhals Jan 10, 2024
a810ab7
Add hflx limiter pd 1
samkellerhals Jan 10, 2024
62b543f
add horadv stencil 1
samkellerhals Jan 11, 2024
e786e21
rbf intp edge stencil 1
samkellerhals Jan 11, 2024
37b1c4d
add zero stencils
samkellerhals Jan 11, 2024
f150408
add upwind_vflux_ppm 1
samkellerhals Jan 11, 2024
7c817b7
TestVertAdvStencil01
samkellerhals Jan 11, 2024
b82d9c7
Add vlimit prbl sm stencils
samkellerhals Jan 11, 2024
90db4d8
Add missing assert in test
samkellerhals Jan 11, 2024
057d9c2
Fix error in numpy test_hflx_limiter_mo_stencil_03
Jan 12, 2024
b154bc1
Fix missing numpy ref. stencils calls
Jan 12, 2024
ffb6b79
refactor hflx limiter mo stencil 03
samkellerhals Jan 16, 2024
d971c05
hflx limiter pd stencil 02
samkellerhals Jan 16, 2024
fab8722
Fix test
samkellerhals Jan 16, 2024
c6cb54d
Use GridType
samkellerhals Jan 17, 2024
de88c3f
Slice p_face
samkellerhals Jan 17, 2024
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
Prev Previous commit
Next Next commit
Add Output helper class
  • Loading branch information
samkellerhals committed Jan 9, 2024
commit db284feefdd44f7efd7d858c95d6d8ceb0201a8b
Original file line number Diff line number Diff line change
Expand Up @@ -10,80 +10,69 @@
# distribution for a copy of the license or check <https://www.gnu.org/licenses/>.
#
# SPDX-License-Identifier: GPL-3.0-or-later

import numpy as np
import pytest
from gt4py.next import as_field
from gt4py.next.ffront.fbuiltins import int32

from icon4py.model.atmosphere.advection.face_val_ppm_stencil_01 import face_val_ppm_stencil_01
from icon4py.model.common.dimension import CellDim, KDim
from icon4py.model.common.grid.simple import SimpleGrid
from icon4py.model.common.test_utils.helpers import _shape, random_field


def face_val_ppm_stencil_01_numpy(
p_cc: np.array,
p_cellhgt_mc_now: np.array,
k: np.array,
elev: int32,
):
# this is a comment: k = np.broadcast_to(k, p_cc.shape)
from icon4py.model.common.test_utils.helpers import (
Output,
StencilTest,
_shape,
random_field,
zero_field,
)

# 01a
zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / (
p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2]
)
zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1])
z_slope_a = (
p_cellhgt_mc_now[:, 1:-1]
/ (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:])
) * (
(2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac
+ (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1
)

# 01b
zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / (
p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2]
)
zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1])
z_slope_b = (
p_cellhgt_mc_now[:, 1:-1]
/ (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1])
) * (
(2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac
+ (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1
class TestFaceValPpmStencil01(StencilTest):
PROGRAM = face_val_ppm_stencil_01
OUTPUTS = (
Output(
"z_slope", refslice=(slice(None), slice(None, -1)), gtslice=(slice(None), slice(1, -1))
Copy link
Contributor

Choose a reason for hiding this comment

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

This is useful thanks. If I understood correctly gtslice is the slicing of the GridTools output and in parenthesis are the dimensions.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

Yes so refslice is the slice for the reference output, and gtslice is the slice for the gt4py output field. By default both are None which means the whole field.

Copy link
Contributor Author

Choose a reason for hiding this comment

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

You can also still use the normal syntax: OUTPUTS = ("a", "b") but then the fields are not sliced.

),
)

z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b)
@staticmethod
def reference(
grid, p_cc: np.array, p_cellhgt_mc_now: np.array, k: np.array, elev: int32, **kwargs
):
zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / (
p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2]
)
zfac = (p_cc[:, 2:] - p_cc[:, 1:-1]) / (p_cellhgt_mc_now[:, 2:] + p_cellhgt_mc_now[:, 1:-1])
z_slope_a = (
p_cellhgt_mc_now[:, 1:-1]
/ (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 2:])
) * (
(2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac
+ (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 2:]) * zfac_m1
)

return z_slope
zfac_m1 = (p_cc[:, 1:-1] - p_cc[:, :-2]) / (
p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, :-2]
)
zfac = (p_cc[:, 1:-1] - p_cc[:, 1:-1]) / (
p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1]
)
z_slope_b = (
p_cellhgt_mc_now[:, 1:-1]
/ (p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1] + p_cellhgt_mc_now[:, 1:-1])
) * (
(2.0 * p_cellhgt_mc_now[:, :-2] + p_cellhgt_mc_now[:, 1:-1]) * zfac
+ (p_cellhgt_mc_now[:, 1:-1] + 2.0 * p_cellhgt_mc_now[:, 1:-1]) * zfac_m1
)

z_slope = np.where(k[1:-1] < elev, z_slope_a, z_slope_b)
return dict(z_slope=z_slope)

def test_face_val_ppm_stencil_01(backend):
grid = SimpleGrid()
p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1})
p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1})

k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32))
elev = k[-2]

z_slope = random_field(grid, CellDim, KDim)

ref = face_val_ppm_stencil_01_numpy(
p_cc.asnumpy(),
p_cellhgt_mc_now.asnumpy(),
k.asnumpy(),
elev,
)

face_val_ppm_stencil_01.with_backend(backend)(
p_cc,
p_cellhgt_mc_now,
k,
elev,
z_slope,
offset_provider={"Koff": KDim},
)
@pytest.fixture
def input_data(self, grid):
z_slope = zero_field(grid, CellDim, KDim)
p_cc = random_field(grid, CellDim, KDim, extend={KDim: 1})
p_cellhgt_mc_now = random_field(grid, CellDim, KDim, extend={KDim: 1})
k = as_field((KDim,), np.arange(0, _shape(grid, KDim, extend={KDim: 1})[0], dtype=int32))
elev = k[-2]

assert np.allclose(ref[:, :-1], z_slope.asnumpy()[:, 1:-1])
return dict(p_cc=p_cc, p_cellhgt_mc_now=p_cellhgt_mc_now, k=k, elev=elev, z_slope=z_slope)
Original file line number Diff line number Diff line change
Expand Up @@ -15,12 +15,15 @@

from icon4py.model.atmosphere.advection.face_val_ppm_stencil_02a import face_val_ppm_stencil_02a
from icon4py.model.common.dimension import CellDim, KDim
from icon4py.model.common.test_utils.helpers import StencilTest, random_field
from icon4py.model.common.test_utils.helpers import Output, StencilTest, random_field


outslice = (slice(None), slice(1, None))


class TestFaceValPpmStencil02a(StencilTest):
PROGRAM = face_val_ppm_stencil_02a
OUTPUTS = (("p_face", (slice(None), slice(1, None))),)
OUTPUTS = (Output("p_face", refslice=outslice, gtslice=outslice),)

@staticmethod
def reference(grid, p_cc: np.array, p_cellhgt_mc_now: np.array, **kwargs):
Expand Down
20 changes: 14 additions & 6 deletions model/common/src/icon4py/model/common/test_utils/helpers.py
Original file line number Diff line number Diff line change
Expand Up @@ -11,6 +11,7 @@
#
# SPDX-License-Identifier: GPL-3.0-or-later

from dataclasses import dataclass, field
from typing import ClassVar, Optional

import numpy as np
Expand Down Expand Up @@ -148,6 +149,13 @@ def allocate_data(backend, input_data):
return input_data


@dataclass(frozen=True)
class Output:
name: str
refslice: tuple[slice, ...] = field(default_factory=lambda: (slice(None),))
gtslice: tuple[slice, ...] = field(default_factory=lambda: (slice(None),))


def _test_validation(self, grid, backend, input_data):
reference_outputs = self.reference(
grid,
Expand All @@ -164,14 +172,14 @@ def _test_validation(self, grid, backend, input_data):
offset_provider=grid.get_all_offset_providers(),
)
for out in self.OUTPUTS:
if isinstance(out, tuple):
name, subset = out
else:
name = out
subset = slice(None)
name, refslice, gtslice = (
(out.name, out.refslice, out.gtslice)
if isinstance(out, Output)
else (out, (slice(None),), (slice(None),))
)

assert np.allclose(
input_data[name].asnumpy()[subset], reference_outputs[name][subset], equal_nan=True
input_data[name].asnumpy()[gtslice], reference_outputs[name][refslice], equal_nan=True
), f"Validation failed for '{name}'"


Expand Down
Loading