Skip to content

Commit 02ba34d

Browse files
mcflugenmcflugen
committed
add tests from regrid with mask; tidy imports
1 parent c9e54f6 commit 02ba34d

File tree

2 files changed

+266
-33
lines changed

2 files changed

+266
-33
lines changed

‎src/addon/esmpy/src/esmpy/test/test_api/test_mesh.py‎

Lines changed: 30 additions & 24 deletions
Original file line numberDiff line numberDiff line change
@@ -1,63 +1,69 @@
1-
21
"""
32
mesh unit test file
43
"""
5-
6-
import pytest
7-
84
import os
95

106
import numpy as np
11-
from numpy.testing import assert_array_equal
12-
13-
from esmpy import *
7+
import pytest
8+
from esmpy.api.constants import _ESMF_NETCDF
9+
from esmpy.api.constants import _ESMF_PIO
10+
from esmpy.api.constants import CoordSys
11+
from esmpy.api.esmpymanager import local_pet
12+
from esmpy.api.esmpymanager import Manager
13+
from esmpy.api.esmpymanager import pet_count
14+
from esmpy.api.field import element
15+
from esmpy.api.field import node
16+
from esmpy.api.mesh import Mesh
1417
from esmpy.test.base import TestBase
15-
from esmpy.api.constants import _ESMF_NETCDF, _ESMF_PIO
16-
from esmpy.util.mesh_utilities import *
1718
from esmpy.util.cache_data import DATA_DIR
19+
from esmpy.util.mesh_utilities import *
20+
from numpy.testing import assert_array_equal
1821

1922

2023
@pytest.mark.parametrize(
21-
"mask",
24+
"mask_value",
2225
(
2326
None,
24-
[0] * 12,
25-
[True] * 12,
26-
[1] * 12,
27-
[False] * 12,
28-
[0, 0, 0, 0, 1, 1, 1, 1, 0, 0, 0, 0],
29-
[[0, 0, 0, 0], [-1, -1, -1, -1], [0, 0, 0, 0]],
27+
0,
28+
True,
29+
1,
30+
False,
31+
[[0, 0, 0, 0], [-1, -1, -1, -1], [1, 1, 1, 1]],
3032
),
3133
)
32-
def test_add_nodes_with_mask(mask):
33-
src = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=esmpy.CoordSys.CART)
34+
def test_add_nodes_with_mask(mask_value):
35+
if mask_value is None:
36+
mask = None
37+
else:
38+
mask = np.broadcast_to(mask_value, (3, 4))
39+
mesh = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=CoordSys.CART)
3440

3541
x_of_points, y_of_points = np.meshgrid([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0])
3642
xy_of_points = np.c_[x_of_points.flat, y_of_points.flat]
3743

38-
src.add_nodes(
44+
mesh.add_nodes(
3945
len(xy_of_points),
4046
np.arange(len(xy_of_points)),
4147
xy_of_points,
4248
np.zeros(len(xy_of_points), dtype=int),
4349
node_mask=mask,
4450
)
4551

46-
assert src.mask[1] is None
52+
assert mesh.mask[1] is None
4753
if mask is None:
48-
assert src.mask[0] is None
54+
assert mesh.mask[0] is None
4955
else:
50-
assert_array_equal(src.mask[0], np.asarray(mask).flat)
56+
assert_array_equal(mesh.mask[0], np.asarray(mask).flat)
5157

5258

5359
def test_add_nodes_with_mask_bad():
54-
src = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=esmpy.CoordSys.CART)
60+
mesh = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=CoordSys.CART)
5561

5662
x_of_points, y_of_points = np.meshgrid([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0])
5763
xy_of_points = np.c_[x_of_points.flat, y_of_points.flat]
5864

5965
with pytest.raises(ValueError):
60-
src.add_nodes(
66+
mesh.add_nodes(
6167
len(xy_of_points),
6268
np.arange(len(xy_of_points)),
6369
xy_of_points,

‎src/addon/esmpy/src/esmpy/test/test_api/test_regrid.py‎

Lines changed: 236 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,18 +1,245 @@
11
"""
22
regrid unit test file
33
"""
4-
5-
import pytest
6-
74
import os
85

9-
from esmpy import *
6+
import esmpy
7+
import numpy as np
8+
import pytest
9+
from esmpy.api.constants import _ESMF_NETCDF
10+
from esmpy.api.constants import _ESMF_PIO
11+
from esmpy.api.constants import _ESMF_USE_INMEM_FACTORS
12+
from esmpy.api.constants import CoordSys
13+
from esmpy.api.constants import LineType
14+
from esmpy.api.constants import Region
15+
from esmpy.api.constants import StaggerLoc
16+
from esmpy.api.constants import UnmappedAction
17+
from esmpy.api.esmpymanager import local_pet
18+
from esmpy.api.esmpymanager import Manager
19+
from esmpy.api.esmpymanager import pet_count
20+
from esmpy.api.field import Field
21+
from esmpy.api.grid import Grid
22+
from esmpy.api.mesh import Mesh
23+
from esmpy.api.mesh import MeshElemType
24+
from esmpy.api.mesh import MeshLoc
25+
from esmpy.api.regrid import Regrid
26+
from esmpy.api.regrid import RegridMethod
1027
from esmpy.test.base import TestBase
11-
from esmpy.api.constants import _ESMF_NETCDF, _ESMF_PIO
12-
from esmpy.util.field_utilities import compare_fields
13-
from esmpy.util.grid_utilities import *
14-
from esmpy.util.mesh_utilities import *
1528
from esmpy.util.cache_data import DATA_DIR
29+
from esmpy.util.field_utilities import compare_fields
30+
from esmpy.util.grid_utilities import compute_mass_grid
31+
from esmpy.util.grid_utilities import grid_create_from_bounds
32+
from esmpy.util.grid_utilities import grid_create_from_bounds_3d
33+
from esmpy.util.grid_utilities import grid_create_from_bounds_periodic
34+
from esmpy.util.grid_utilities import grid_create_from_bounds_periodic_3d
35+
from esmpy.util.grid_utilities import initialize_field_grid
36+
from esmpy.util.grid_utilities import initialize_field_grid_3d
37+
from esmpy.util.grid_utilities import initialize_field_grid_periodic
38+
from esmpy.util.grid_utilities import initialize_field_grid_periodic_3d
39+
from esmpy.util.mesh_utilities import compute_mass_mesh
40+
from esmpy.util.mesh_utilities import initialize_field_mesh
41+
from esmpy.util.mesh_utilities import mesh_create_10
42+
from esmpy.util.mesh_utilities import mesh_create_50
43+
from esmpy.util.mesh_utilities import mesh_create_50_ngons
44+
from numpy.testing import assert_array_almost_equal
45+
46+
47+
NON_CONSERVATIVE_METHODS = (
48+
RegridMethod.BILINEAR,
49+
RegridMethod.NEAREST_DTOS,
50+
RegridMethod.NEAREST_STOD,
51+
RegridMethod.PATCH,
52+
)
53+
54+
55+
def create_raster_field(x_of_col, y_of_row, node_mask=None):
56+
"""Create a mesh that is a structured grid of squares."""
57+
x_of_col, y_of_row = np.asarray(x_of_col), np.asarray(y_of_row)
58+
if node_mask is not None:
59+
node_mask = np.asarray(node_mask).flatten()
60+
61+
mesh = Mesh(parametric_dim=2, spatial_dim=2, coord_sys=CoordSys.CART)
62+
63+
x_of_point, y_of_point = np.meshgrid(x_of_col, y_of_row)
64+
xy_of_point = np.c_[x_of_point.flat, y_of_point.flat]
65+
n_rows, n_cols = len(y_of_row), len(x_of_col)
66+
n_points = n_rows * n_cols
67+
id_of_point = np.arange(n_points).reshape((n_rows, n_cols))
68+
69+
x_of_element, y_of_element = np.meshgrid(
70+
0.5 * (x_of_col[:-1] + x_of_col[1:]),
71+
0.5 * (y_of_row[:-1] + y_of_row[1:]),
72+
)
73+
xy_of_element = np.c_[x_of_element.flat, y_of_element.flat]
74+
n_elements = (n_rows - 1) * (n_cols - 1)
75+
id_of_element = np.arange(n_elements)
76+
points_at_element = np.c_[
77+
id_of_point[:-1,:-1].flat,
78+
id_of_point[1:,:-1].flat,
79+
id_of_point[1:,1:].flat,
80+
id_of_point[:-1,1:].flat,
81+
]
82+
83+
mesh.add_nodes(
84+
n_points,
85+
id_of_point.flatten(),
86+
xy_of_point.flatten(),
87+
np.zeros(n_points, dtype=int),
88+
node_mask=node_mask,
89+
)
90+
91+
mesh.add_elements(
92+
n_elements,
93+
np.arange(n_elements, dtype=int),
94+
np.full(n_elements, MeshElemType.QUAD, dtype=int),
95+
points_at_element,
96+
xy_of_element,
97+
)
98+
99+
return Field(mesh, meshloc=MeshLoc.NODE)
100+
101+
102+
@pytest.mark.parametrize("mask_value", (None, 0, True, 1, False))
103+
@pytest.mark.parametrize("method", NON_CONSERVATIVE_METHODS)
104+
def test_regrid_with_const_node_mask(mask_value, method):
105+
"""Check regriding with an all-or-nothing mask."""
106+
if mask_value is None:
107+
mask = None
108+
else:
109+
mask = np.full(12, mask_value)
110+
111+
src = create_raster_field([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0], node_mask=mask)
112+
dst = create_raster_field([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0], node_mask=mask)
113+
src.data[:] = 100.0
114+
dst.data[:] = -999.0
115+
116+
regrid = Regrid(
117+
src,
118+
dst,
119+
regrid_method=method,
120+
unmapped_action=UnmappedAction.IGNORE,
121+
src_mask_values=(1,),
122+
dst_mask_values=(1,),
123+
)
124+
125+
if mask is None:
126+
mask = np.full(src.data.size, False, dtype=bool)
127+
else:
128+
mask = np.asarray(mask, dtype=bool).reshape(-1)
129+
130+
assert_array_almost_equal(dst.data[~mask], src.data[~mask])
131+
assert_array_almost_equal(dst.data[mask], -999)
132+
133+
134+
@pytest.mark.parametrize(
135+
"mask",
136+
(
137+
[[1, 1, 1, 1], [0, 0, 0, 0], [0, 0, 0, 0]],
138+
[[0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1]],
139+
[[1, 0, 0, 0], [1, 0, 0, 0], [1, 0, 0, 0]],
140+
[[1, 0, 0, 1], [1, 0, 0, 1], [1, 0, 0, 1]],
141+
[[0, 0, 0, 1], [0, 0, 0, 1], [0, 0, 0, 1]],
142+
),
143+
)
144+
@pytest.mark.parametrize("method", NON_CONSERVATIVE_METHODS)
145+
def test_regrid_with_node_mask(mask, method):
146+
src = create_raster_field([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0], node_mask=mask)
147+
dst = create_raster_field([0.0, 1.0, 2.0, 3.0], [0.0, 1.0, 2.0], node_mask=mask)
148+
src.data[:] = 100.0
149+
dst.data[:] = -999.0
150+
151+
regrid = Regrid(
152+
src,
153+
dst,
154+
regrid_method=method,
155+
unmapped_action=UnmappedAction.IGNORE,
156+
src_mask_values=(1,),
157+
dst_mask_values=(1,),
158+
)
159+
160+
if mask is None:
161+
mask = np.full(src.data.size, False, dtype=bool)
162+
else:
163+
mask = np.asarray(mask, dtype=bool).reshape(-1)
164+
165+
assert_array_almost_equal(dst.data[~mask], src.data[~mask])
166+
assert_array_almost_equal(dst.data[mask], -999)
167+
168+
169+
@pytest.mark.parametrize("method", NON_CONSERVATIVE_METHODS)
170+
def test_regrid_with_random_node_mask(method):
171+
n_rows, n_cols = 100, 200
172+
173+
src_data = np.random.uniform(-1.0, 1.0, size=n_rows * n_cols)
174+
mask = src_data > 0.0
175+
176+
src = create_raster_field(
177+
np.linspace(0.0, 200.0, num=n_cols, endpoint=True),
178+
np.linspace(0.0, 100.0, num=n_rows, endpoint=True),
179+
node_mask=mask,
180+
)
181+
dst = create_raster_field(
182+
np.linspace(0.25, 199.25, num=n_cols - 1, endpoint=True),
183+
np.linspace(0.25, 99.25, num=n_rows - 1, endpoint=True),
184+
node_mask=None,
185+
)
186+
187+
src.data[:] = src_data
188+
dst.data[:] = 999.0
189+
190+
regrid = Regrid(
191+
src,
192+
dst,
193+
regrid_method=method,
194+
unmapped_action=UnmappedAction.IGNORE,
195+
src_mask_values=(True,),
196+
dst_mask_values=(True,),
197+
)
198+
199+
assert np.all((dst.data <= 0.0) | (dst.data == 999.0))
200+
201+
202+
@pytest.mark.parametrize("method", NON_CONSERVATIVE_METHODS)
203+
def test_regrid_with_multivalued_node_mask(method):
204+
n_rows, n_cols = 100, 200
205+
206+
src_data = np.random.uniform(-1.0, 1.0, size=n_rows * n_cols)
207+
208+
mask = np.zeros(src_data.size, dtype=int)
209+
mask[src_data < -0.5] = 1
210+
mask[(src_data >= -0.5) & (src_data < 0.5)] = 2
211+
mask[src_data >= 0.5] = 3
212+
213+
src = create_raster_field(
214+
np.linspace(0.0, 200.0, num=n_cols, endpoint=True),
215+
np.linspace(0.0, 100.0, num=n_rows, endpoint=True),
216+
node_mask=mask,
217+
)
218+
dst = create_raster_field(
219+
np.linspace(0.25, 199.25, num=n_cols - 1, endpoint=True),
220+
np.linspace(0.25, 99.25, num=n_rows - 1, endpoint=True),
221+
node_mask=None,
222+
)
223+
224+
src.data[:] = src_data
225+
dst.data[:] = 999.0
226+
227+
regrid = Regrid(
228+
src,
229+
dst,
230+
regrid_method=method,
231+
unmapped_action=UnmappedAction.IGNORE,
232+
src_mask_values=(1, 3),
233+
)
234+
235+
actual = np.asarray(dst.data).reshape((n_rows - 1, n_cols - 1))
236+
if method == RegridMethod.NEAREST_DTOS:
237+
# For this method and grid layout, the last row and column take
238+
# input from multiple source nodes, and thus we can't be sure about
239+
# the values at the destination nodes.
240+
actual = actual[:-1, :-1]
241+
242+
assert np.all(((actual >= -0.5) & (actual < 0.5)) | (actual == 999.0))
16243

17244

18245
class TestRegrid(TestBase):
@@ -83,7 +310,7 @@ def test_field_regrid(self):
83310
line_type=LineType.CART, factors=False)
84311
_ = rh(srcfield, dstfield)
85312

86-
@pytest.mark.skipif(not constants._ESMF_USE_INMEM_FACTORS, reason="compiler does not support in-memory weights")
313+
@pytest.mark.skipif(not _ESMF_USE_INMEM_FACTORS, reason="compiler does not support in-memory weights")
87314
@pytest.mark.skipif(pet_count()!=1, reason="test must be run in serial")
88315
def test_field_regrid_factor_retrieval(self):
89316
# Test retrieving factors from a route handle.

0 commit comments

Comments
 (0)