def test_interpolation__masked(self):
levels = np.array([0.5, 1.5])
new_data = np.empty([len(levels)] + list(self.shape[1:]), dtype=float)
new_data[:, 0, :] = np.nan
new_data_mask = np.isnan(new_data)
scheme = 'linear'
mask = [[[False], [True]], [[True], [False]], [[False], [False]]]
masked = ma.empty(self.shape)
masked.mask = mask
cube = _make_cube(masked, dtype=self.dtype)
# save cube to test the lazy data interpolation too
iris.save(cube, self.filename)
with mock.patch('stratify.interpolate',
return_value=new_data) as mocker:
# first test lazy
loaded_cube = iris.load_cube(self.filename)
result_from_lazy = extract_levels(loaded_cube, levels, scheme)
self.assertEqual(result_from_lazy, self.created_cube)
# then test realized
result = extract_levels(cube, levels, scheme)
self.assertEqual(result, self.created_cube)
args, kwargs = mocker.call_args
# Check the stratify.interpolate args ...
self.assertEqual(len(args), 3)
self.assert_array_equal(args[0], levels)
pts = cube.coord(axis='z', dim_coords=True).points
src_levels_broadcast = np.broadcast_to(pts.reshape(self.z, 1, 1),
cube.shape)
self.assert_array_equal(args[1], src_levels_broadcast)
self.assert_array_equal(args[2], cube.data)
# Check the stratify.interpolate kwargs ...
self.assertEqual(
kwargs, dict(axis=0, interpolation=scheme,
extrapolation='nan'))
args, kwargs = self.mock_create_cube.call_args
# in-place for new extract_levels with nan's
new_data[np.isnan(new_data)] = _MDI
# Check the _create_cube args ...
self.assertEqual(len(args), 4)
self.assertEqual(args[0].metadata, cube.metadata)
coord_comparison = iris.analysis.coord_comparison(args[0], cube)
self.assertFalse(coord_comparison['not_equal']
or coord_comparison['non_equal_data_dimension'])
self.assert_array_equal(args[0].data, cube.data)
new_data_mask = np.zeros(new_data.shape, bool)
new_data_mask[new_data == _MDI] = True
new_data = np.ma.array(new_data, mask=new_data_mask)
self.assert_array_equal(args[1], new_data)
self.assertTrue(ma.isMaskedArray(args[1]))
self.assert_array_equal(args[1].mask, new_data_mask)
self.assert_array_equal(args[2],
self.cube.coord(axis='z', dim_coords=True))
self.assert_array_equal(args[3], levels)
# Check the _create_cube kwargs ...
self.assertEqual(kwargs, dict())
def test_nop__levels_match(self):
vcoord = _make_vcoord(self.z)
self.assertEqual(self.cube.coord(axis='z', dim_coords=True), vcoord)
levels = vcoord.points
result = extract_levels(self.cube, levels, 'linear')
self.assertEqual(result, self.cube)
self.assertEqual(id(result), id(self.cube))
def test_interpolation(self):
new_data = np.array(True)
levels = np.array([0.5, 1.5])
scheme = 'linear'
with mock.patch('stratify.interpolate',
return_value=new_data) as mocker:
result = extract_levels(self.cube, levels, scheme)
self.assertEqual(result, self.created_cube)
args, kwargs = mocker.call_args
# Check the stratify.interpolate args ...
self.assertEqual(len(args), 3)
self.assert_array_equal(args[0], levels)
pts = self.cube.coord(axis='z', dim_coords=True).points
src_levels_broadcast = np.broadcast_to(pts.reshape(self.z, 1, 1),
self.cube.shape)
self.assert_array_equal(args[1], src_levels_broadcast)
self.assert_array_equal(args[2], self.cube.data)
# Check the stratify.interpolate kwargs ...
self.assertEqual(
kwargs, dict(axis=0, interpolation=scheme,
extrapolation='nan'))
args, kwargs = self.mock_create_cube.call_args
# Check the _create_cube args ...
self.assertEqual(len(args), 4)
self.assertEqual(args[0], self.cube)
self.assert_array_equal(args[1], new_data)
self.assert_array_equal(args[2],
self.cube.coord(axis='z', dim_coords=True))
self.assert_array_equal(args[3], levels)
# Check the _create_cube kwargs ...
self.assertEqual(kwargs, dict())
def test_interpolation__extrapolated_nan_filling(self):
new_data = np.array([0, np.nan])
levels = [0.5, 1.5]
scheme = 'nearest'
with mock.patch('stratify.interpolate',
return_value=new_data) as mocker:
result = extract_levels(self.cube, levels, scheme)
self.assertEqual(result, self.created_cube)
args, kwargs = mocker.call_args
# Check the stratify.interpolate args ...
self.assertEqual(len(args), 3)
self.assert_array_equal(args[0], levels)
pts = self.cube.coord(axis='z', dim_coords=True).points
src_levels_broadcast = np.broadcast_to(pts.reshape(self.z, 1, 1),
self.cube.shape)
self.assert_array_equal(args[1], src_levels_broadcast)
self.assert_array_equal(args[2], self.cube.data)
# Check the stratify.interpolate kwargs ...
self.assertEqual(
kwargs, dict(axis=0, interpolation=scheme,
extrapolation='nan'))
args, kwargs = self.mock_create_cube.call_args
# Check the _create_cube args ...
self.assertEqual(len(args), 4)
self.assert_array_equal(args[0], self.cube)
new_data[np.isnan(new_data)] = _MDI
new_data_mask = np.zeros(new_data.shape, bool)
new_data_mask[new_data == _MDI] = True
new_data = np.ma.array(new_data, mask=new_data_mask)
self.assert_array_equal(args[1], new_data)
self.assert_array_equal(args[2],
self.cube.coord(axis='z', dim_coords=True))
self.assert_array_equal(args[3], levels)
# Check the _create_cube kwargs ...
self.assertEqual(kwargs, dict())
def test_interpolation__scalar_collapse(self):
level = 1
scheme = 'nearest'
result = extract_levels(self.cube, level, scheme)
expected = np.array([[[4., 5.], [6., 7.]], [[16., 17.], [18., 19.]]])
self.assertArrayEqual(result.data, expected)
del self.shape[self.z_dim]
self.assertEqual(result.shape, tuple(self.shape))
def test_interpolation__nearest(self):
levels = [0.49, 1.51]
scheme = 'nearest'
result = extract_levels(self.cube, levels, scheme)
expected = np.array([[[[0., 1.], [2., 3.]], [[8., 9.], [10., 11.]]],
[[[12., 13.], [14., 15.]], [[20., 21.],
[22., 23.]]]])
self.assertArrayEqual(result.data, expected)
self.shape[self.z_dim] = len(levels)
self.assertEqual(result.shape, tuple(self.shape))
def test_interpolation__linear(self):
levels = [0.5, 1.5]
scheme = 'linear'
result = extract_levels(self.cube, levels, scheme)
expected = np.array([[[[2., 3.], [4., 5.]], [[6., 7.], [8., 9.]]],
[[[14., 15.], [16., 17.]], [[18., 19.],
[20., 21.]]]])
self.assertArrayEqual(result.data, expected)
self.shape[self.z_dim] = len(levels)
self.assertEqual(result.shape, tuple(self.shape))
def test_add_alt_coord(self):
assert self.cube.coords('air_pressure')
assert not self.cube.coords('altitude')
result = extract_levels(self.cube, [1, 2],
'linear_horizontal_extrapolate_vertical',
coordinate='altitude')
assert not result.coords('air_pressure')
assert result.coords('altitude')
assert result.shape == (2, 2, 2, 2)
np.testing.assert_allclose(result.coord('altitude').points, [1.0, 2.0])
def test_extraction(self):
levels = [0, 2]
result = extract_levels(self.cube, levels, 'linear')
data = np.array([0, 1, 4, 5], dtype=self.dtype).reshape(2, 2, 1)
expected = _make_cube(
data, aux_coord=False, dim_coord=False, dtype=self.dtype)
coord = self.cube.coord('Pressure Slice').copy()
expected.add_aux_coord(coord[levels], (0, 1))
coord = self.cube.coord('air_pressure').copy()
expected.add_dim_coord(coord[levels], 0)
self.assertEqual(result, expected)
def test_interpolation__extrapolated_nan_filling(self):
levels = [-10, 1, 2, 10]
scheme = 'nearest'
result = extract_levels(self.cube, levels, scheme)
expected = np.array([[[[_MDI, _MDI], [_MDI, _MDI]], [[4., 5.],
[6., 7.]],
[[8., 9.], [10., 11.]],
[[_MDI, _MDI], [_MDI, _MDI]]],
[[[_MDI, _MDI], [_MDI, _MDI]],
[[16., 17.], [18., 19.]], [[20., 21.],
[22., 23.]],
[[_MDI, _MDI], [_MDI, _MDI]]]])
self.assertArrayEqual(result.data, expected)
self.shape[self.z_dim] = len(levels)
self.assertEqual(result.shape, tuple(self.shape))
def test_add_plev_coord(self):
self.cube.coord('air_pressure').standard_name = 'altitude'
self.cube.coord('altitude').var_name = 'alt'
self.cube.coord('altitude').long_name = 'altitude'
self.cube.coord('altitude').units = 'm'
assert not self.cube.coords('air_pressure')
assert self.cube.coords('altitude')
result = extract_levels(self.cube, [1, 2],
'linear_horizontal_extrapolate_vertical',
coordinate='air_pressure')
assert result.coords('air_pressure')
assert not result.coords('altitude')
assert result.shape == (2, 2, 2, 2)
np.testing.assert_allclose(
result.coord('air_pressure').points, [1.0, 2.0])
def test_extraction__failure(self):
levels = [0, 2]
with mock.patch('iris.cube.Cube.extract', return_value=None):
emsg = 'Failed to extract levels'
with self.assertRaisesRegex(ValueError, emsg):
extract_levels(self.cube, levels, 'linear')
def test_invalid_scheme__unknown(self):
levels = mock.sentinel.levels
scheme = mock.sentinel.scheme
emsg = 'Unknown vertical interpolation scheme'
with self.assertRaisesRegex(ValueError, emsg):
extract_levels(self.cube, levels, scheme)
