Skip to content

Commit 1ee1716

Browse files
valeriupredoivaleriupredoi
authored
Fix da.broadcast_to call when the fx cube has different shape than target data cube (#1350)
* fix the broadcasting issue in mask * make 4d data get fx masked in tests * rename fx coordinate specifications
1 parent e506259 commit 1ee1716

File tree

2 files changed

+58
-33
lines changed

2 files changed

+58
-33
lines changed

‎esmvalcore/preprocessor/_mask.py

+4-4
Original file line numberDiff line numberDiff line change
@@ -114,8 +114,8 @@ def mask_landsea(cube, mask_out, always_use_ne_mask=False):
114114
'not found in cube. Check fx_file availability.')
115115

116116
if fx_cube:
117-
fx_cube.data = da.broadcast_to(fx_cube.core_data(), cube.shape)
118-
landsea_mask = _get_fx_mask(fx_cube.data, mask_out,
117+
fx_cube_data = da.broadcast_to(fx_cube.core_data(), cube.shape)
118+
landsea_mask = _get_fx_mask(fx_cube_data, mask_out,
119119
fx_cube.var_name)
120120
cube.data = _apply_fx_mask(landsea_mask, cube.data)
121121
logger.debug("Applying land-sea mask: %s", fx_cube.var_name)
@@ -181,8 +181,8 @@ def mask_landseaice(cube, mask_out):
181181
logger.debug('Ancillary variable land ice area fraction '
182182
'not found in cube. Check fx_file availability.')
183183
if fx_cube:
184-
fx_cube.data = da.broadcast_to(fx_cube.core_data(), cube.shape)
185-
landice_mask = _get_fx_mask(fx_cube.data, mask_out, fx_cube.var_name)
184+
fx_cube_data = da.broadcast_to(fx_cube.core_data(), cube.shape)
185+
landice_mask = _get_fx_mask(fx_cube_data, mask_out, fx_cube.var_name)
186186
cube.data = _apply_fx_mask(landice_mask, cube.data)
187187
logger.debug("Applying landsea-ice mask: sftgif")
188188
else:

‎tests/integration/preprocessor/_mask/test_mask.py

+54-29
Original file line numberDiff line numberDiff line change
@@ -25,7 +25,7 @@ def setUp(self):
2525
fx_data = np.empty((3, 3))
2626
fx_data[:] = 60.
2727
fx_data[1, 2] = 30.
28-
self.new_cube_data = np.empty((3, 3))
28+
self.new_cube_data = np.empty((2, 3, 3))
2929
self.new_cube_data[:] = 200.
3030
crd_sys = iris.coord_systems.GeogCS(iris.fileformats.pp.EARTH_RADIUS)
3131
self.lons = iris.coords.DimCoord([0, 1.5, 3],
@@ -38,6 +38,11 @@ def setUp(self):
3838
bounds=[[0, 1], [1, 2], [2, 3]],
3939
units='degrees_north',
4040
coord_system=crd_sys)
41+
self.zcoord = iris.coords.DimCoord([0.5, 5.],
42+
long_name='zcoord',
43+
bounds=[[0., 2.5], [2.5, 25.]],
44+
units='m',
45+
attributes={'positive': 'down'})
4146
self.times = iris.coords.DimCoord([0, 1.5, 2.5, 3.5],
4247
standard_name='time',
4348
bounds=[[0, 1], [1, 2], [2, 3],
@@ -47,9 +52,11 @@ def setUp(self):
4752
standard_name='time',
4853
bounds=[[0, 1], [1, 2], [2, 3]],
4954
units='hours')
50-
self.coords_spec = [(self.lats, 0), (self.lons, 1)]
55+
self.fx_coords_spec = [(self.lats, 0), (self.lons, 1)]
56+
self.cube_coords_spec = [(self.zcoord, 0),
57+
(self.lats, 1), (self.lons, 2)]
5158
self.fx_mask = iris.cube.Cube(fx_data,
52-
dim_coords_and_dims=self.coords_spec,
59+
dim_coords_and_dims=self.fx_coords_spec,
5360
units='%')
5461
self.mock_data = np.ma.empty((4, 3, 3))
5562
self.mock_data[:] = 10.
@@ -69,8 +76,10 @@ def test_components_fx_var(self, tmp_path):
6976
'frequency': 'fx',
7077
'filename': sftlf_file}
7178
}
72-
new_cube_land = iris.cube.Cube(self.new_cube_data,
73-
dim_coords_and_dims=self.coords_spec)
79+
new_cube_land = iris.cube.Cube(
80+
self.new_cube_data,
81+
dim_coords_and_dims=self.cube_coords_spec
82+
)
7483
new_cube_land = add_fx_variables(new_cube_land, fx_vars,
7584
CheckLevels.IGNORE)
7685
result_land = mask_landsea(
@@ -92,8 +101,10 @@ def test_components_fx_var(self, tmp_path):
92101
'frequency': 'fx',
93102
'filename': sftlf_file}
94103
}
95-
new_cube_ice = iris.cube.Cube(self.new_cube_data,
96-
dim_coords_and_dims=self.coords_spec)
104+
new_cube_ice = iris.cube.Cube(
105+
self.new_cube_data,
106+
dim_coords_and_dims=self.cube_coords_spec
107+
)
97108
new_cube_ice = add_fx_variables(new_cube_ice, fx_vars,
98109
CheckLevels.IGNORE)
99110
result_ice = mask_landseaice(
@@ -117,12 +128,16 @@ def test_mask_landsea(self, tmp_path):
117128
'frequency': 'fx',
118129
'filename': sftlf_file}
119130
}
120-
new_cube_land = iris.cube.Cube(self.new_cube_data,
121-
dim_coords_and_dims=self.coords_spec)
131+
new_cube_land = iris.cube.Cube(
132+
self.new_cube_data,
133+
dim_coords_and_dims=self.cube_coords_spec
134+
)
122135
new_cube_land = add_fx_variables(new_cube_land, fx_vars,
123136
CheckLevels.IGNORE)
124-
new_cube_sea = iris.cube.Cube(self.new_cube_data,
125-
dim_coords_and_dims=self.coords_spec)
137+
new_cube_sea = iris.cube.Cube(
138+
self.new_cube_data,
139+
dim_coords_and_dims=self.cube_coords_spec
140+
)
126141
new_cube_sea = add_fx_variables(new_cube_sea, fx_vars,
127142
CheckLevels.IGNORE)
128143

@@ -135,26 +150,30 @@ def test_mask_landsea(self, tmp_path):
135150
new_cube_sea,
136151
'sea',
137152
)
138-
expected = np.ma.empty((3, 3))
153+
expected = np.ma.empty((2, 3, 3))
139154
expected.data[:] = 200.
140-
expected.mask = np.ones((3, 3), bool)
141-
expected.mask[1, 2] = False
155+
expected.mask = np.ones((2, 3, 3), bool)
156+
expected.mask[:, 1, 2] = False
142157
# set fillvalues so we are sure they are equal
143158
np.ma.set_fill_value(result_land.data, 1e+20)
144159
np.ma.set_fill_value(result_sea.data, 1e+20)
145160
np.ma.set_fill_value(expected, 1e+20)
146161
assert_array_equal(result_land.data, expected)
147-
expected.mask = np.zeros((3, 3), bool)
148-
expected.mask[1, 2] = True
162+
expected.mask = np.zeros((2, 3, 3), bool)
163+
expected.mask[:, 1, 2] = True
149164
assert_array_equal(result_sea.data, expected)
150165

151166
# Mask with shp files although sftlf is available
152-
new_cube_land = iris.cube.Cube(self.new_cube_data,
153-
dim_coords_and_dims=self.coords_spec)
167+
new_cube_land = iris.cube.Cube(
168+
self.new_cube_data,
169+
dim_coords_and_dims=self.cube_coords_spec
170+
)
154171
new_cube_land = add_fx_variables(new_cube_land, fx_vars,
155172
CheckLevels.IGNORE)
156-
new_cube_sea = iris.cube.Cube(self.new_cube_data,
157-
dim_coords_and_dims=self.coords_spec)
173+
new_cube_sea = iris.cube.Cube(
174+
self.new_cube_data,
175+
dim_coords_and_dims=self.cube_coords_spec
176+
)
158177
new_cube_sea = add_fx_variables(new_cube_sea, fx_vars,
159178
CheckLevels.IGNORE)
160179
result_land = mask_landsea(
@@ -177,10 +196,14 @@ def test_mask_landsea(self, tmp_path):
177196
assert_array_equal(result_sea.data, expected)
178197

179198
# mask with shp files
180-
new_cube_land = iris.cube.Cube(self.new_cube_data,
181-
dim_coords_and_dims=self.coords_spec)
182-
new_cube_sea = iris.cube.Cube(self.new_cube_data,
183-
dim_coords_and_dims=self.coords_spec)
199+
new_cube_land = iris.cube.Cube(
200+
self.new_cube_data,
201+
dim_coords_and_dims=self.cube_coords_spec
202+
)
203+
new_cube_sea = iris.cube.Cube(
204+
self.new_cube_data,
205+
dim_coords_and_dims=self.cube_coords_spec
206+
)
184207
result_land = mask_landsea(new_cube_land, 'land')
185208
result_sea = mask_landsea(new_cube_sea, 'sea')
186209

@@ -207,15 +230,17 @@ def test_mask_landseaice(self, tmp_path):
207230
'frequency': 'fx',
208231
'filename': sftgif_file}
209232
}
210-
new_cube_ice = iris.cube.Cube(self.new_cube_data,
211-
dim_coords_and_dims=self.coords_spec)
233+
new_cube_ice = iris.cube.Cube(
234+
self.new_cube_data,
235+
dim_coords_and_dims=self.cube_coords_spec
236+
)
212237
new_cube_ice = add_fx_variables(new_cube_ice, fx_vars,
213238
CheckLevels.IGNORE)
214239
result_ice = mask_landseaice(new_cube_ice, 'ice')
215-
expected = np.ma.empty((3, 3))
240+
expected = np.ma.empty((2, 3, 3))
216241
expected.data[:] = 200.
217-
expected.mask = np.ones((3, 3), bool)
218-
expected.mask[1, 2] = False
242+
expected.mask = np.ones((2, 3, 3), bool)
243+
expected.mask[:, 1, 2] = False
219244
np.ma.set_fill_value(result_ice.data, 1e+20)
220245
np.ma.set_fill_value(expected, 1e+20)
221246
assert_array_equal(result_ice.data, expected)

0 commit comments

Comments
 (0)