def test_basedataobj_dim_noleadsample():
data = np.arange(24).reshape(2, 2, 3, 2)
coords = {
BDO.TIME: (2, [1, 2, 3]),
BDO.LEVEL: (1, [1000, 900]),
BDO.LAT: (0, [50, 55]),
BDO.LON: (3, [-80, -90])
}
obj = BDO(data, dim_coords=coords)
def test_basedataobj_grid_nomatch():
data = np.arange(24).reshape(4, 3, 2)
coords = {
BDO.TIME: (0, [1, 2, 3, 4]),
BDO.LAT: (1, [45, 50, 55, 60]),
BDO.LON: (2, [-80, -90])
}
obj = BDO(data, dim_coords=coords)
obj.get_coordinate_grids(BDO.LAT)
def test_basedataobj_dim_notime():
data = np.arange(24).reshape(2, 2, 3, 2)
coords = {
BDO.LEVEL: (1, [1000, 900]),
BDO.LAT: (2, [45, 50, 55]),
BDO.LON: (3, [-80, -90])
}
obj = BDO(data, dim_coords=coords)
assert len(list(obj._dim_idx.keys())) == 3
assert obj._full_shp == data.shape
def test_basedataobj_nanentry_noleadtime():
data = np.arange(24).reshape(3, 4, 2).astype(np.float32)
data[2, 3, 1] = np.nan
obj = BDO(data)
assert obj.is_masked
np.testing.assert_array_equal(data.flatten(), obj.orig_data)
assert np.array_equal(data.flatten()[:-1], obj.compressed_data)
assert np.array_equal(data.flatten()[:-1], obj.data)
np.testing.assert_array_equal(data,
obj.inflate_full_grid(reshape_orig=True),
err_msg='Inflation to full grid failed.')
def test_basedataobj_grid_flat():
data = np.arange(24).reshape(4, 3, 2)
coords = {
BDO.TIME: (0, [1, 2, 3, 4]),
BDO.LAT: (1, [45, 50, 55]),
BDO.LON: (2, [-80, -90])
}
obj = BDO(data, dim_coords=coords, force_flat=True)
tmp = obj.get_coordinate_grids([BDO.LAT, BDO.LON])
longrd, latgrd = np.meshgrid(coords[BDO.LON][1], coords[BDO.LAT][1])
assert np.array_equal(tmp[BDO.LON], longrd.flatten())
assert np.array_equal(tmp[BDO.LAT], latgrd.flatten())
def test_basedataobj_dim_mismatched():
data = np.arange(24).reshape(2, 2, 3, 2)
coords = {
BDO.TIME: (0, [1, 2]),
BDO.LEVEL: (1, [1000, 900, 800]),
BDO.LAT: (2, [45, 50, 55]),
BDO.LON: (3, [-80, -90, 100])
}
obj = BDO(data, dim_coords=coords)
assert len(list(obj._dim_idx.keys())) == 2
assert BDO.TIME in obj._dim_idx
assert BDO.LAT in obj._dim_idx
def test_basedataobj_compressed_noleadtime():
data = np.arange(24).reshape(3, 4, 2).astype(np.float32)
data[2, 3, 1] = np.nan
data[1, 0, 1] = np.nan
valid = np.isfinite(data)
comp = data[valid]
valid = valid.flatten()
obj = BDO(comp, valid_data=valid)
assert obj.is_masked
assert obj.data.shape == comp.shape
assert valid.shape == tuple(obj._full_shp)
np.testing.assert_array_equal(data.flatten(), obj.inflate_full_grid())
def test_basedataobj_masked_noleadtime():
data = np.arange(24).reshape(3, 4, 2).astype(np.float32)
mask = np.zeros_like(data, dtype=np.bool)
mask[2, 3, 1] = True
data = np.ma.array(data, mask=mask)
obj = BDO(data)
assert obj.is_masked
np.testing.assert_array_equal(data.flatten(), obj.orig_data)
assert np.array_equal(data.flatten()[:-1], obj.compressed_data)
assert np.array_equal(data.flatten()[:-1], obj.data)
np.testing.assert_array_equal(data,
obj.inflate_full_grid(reshape_orig=True),
err_msg='Inflation to full grid failed.')
def test_basedataobj_dim_matching():
data = np.arange(24).reshape(2, 2, 3, 2)
coords = {
BDO.TIME: (0, [1, 2]),
BDO.LEVEL: (1, [1000, 900]),
BDO.LAT: (2, [45, 50, 55]),
BDO.LON: (3, [-80, -90])
}
obj = BDO(data, dim_coords=coords)
assert obj._leading_time
assert obj._dim_idx == {
key: value[0]
for key, value in list(coords.items())
}
assert data.shape[0] == obj._time_shp[0]
assert data.shape[1:] == obj._spatial_shp
assert data.shape == obj.data.shape
def test_basedataobj_compressed_leadtime():
data = np.arange(24).reshape(3, 4, 2).astype(np.float32)
data[2, 3, 1] = np.nan
data[1, 0, 1] = np.nan
dim = {BDO.TIME: (0, [1, 2, 3])}
valid = np.isfinite(data[0])
for time in data:
valid &= np.isfinite(time)
comp = data[:, valid]
inflated = np.copy(data).reshape(3, 8)
inflated[:, ~valid.flatten()] = np.nan
obj = BDO(comp, dim_coords=dim, valid_data=valid.flatten())
assert obj._leading_time
assert obj.is_masked
np.testing.assert_array_equal(obj.inflate_full_grid(), inflated)
def test_basedataobj_grid_masked_compressed():
data = np.arange(24).reshape(4, 3, 2).astype(np.float16)
data[2, 2, 1] = np.nan
data[1, 0, 1] = np.nan
coords = {
BDO.TIME: (0, [1, 2, 3, 4]),
BDO.LAT: (1, [45, 50, 55]),
BDO.LON: (2, [-80, -90])
}
obj = BDO(data, dim_coords=coords, force_flat=True)
tmp = obj.get_coordinate_grids([BDO.LAT, BDO.LON])
longrd, latgrd = np.meshgrid(coords[BDO.LON][1], coords[BDO.LAT][1])
valid = np.isfinite(data[0])
for time in data:
valid &= np.isfinite(time)
assert np.array_equal(tmp[BDO.LON], longrd[valid].flatten())
assert np.array_equal(tmp[BDO.LAT], latgrd[valid].flatten())
def test_basedataobj_nanentry_leadtime():
data = np.arange(24).reshape(3, 2, 2, 2).astype(np.float32)
data[1, 0, 0, 1] = np.nan
data[2, 1, 1, 1] = np.nan
dim = {BDO.TIME: (0, [1, 2, 3])}
valid = np.isfinite(data[0])
for time in data:
valid &= np.isfinite(time)
full_valid = np.ones_like(data, dtype=np.bool) * valid
obj = BDO(data, dim_coords=dim)
assert obj._leading_time
assert obj.is_masked
assert obj.data.size == 18 # remove entire nan loc from entire sample
assert np.array_equal(obj.data, data[:, valid])
data[~full_valid] = np.nan
np.testing.assert_array_equal(data,
obj.inflate_full_grid(reshape_orig=True))
def test_basedataobj_masked_leadtime():
data = np.arange(24).reshape(3, 2, 2, 2).astype(np.float32)
mask = np.zeros_like(data, dtype=np.bool)
mask[1, 0, 0, 1] = True
mask[2, 1, 1, 1] = True
data = np.ma.array(data, mask=mask)
dim = {BDO.TIME: (0, [1, 2, 3])}
valid = np.logical_not(data.mask)
valid_composite = valid.sum(axis=0) == data.shape[0]
full_valid = np.ones_like(data, dtype=np.bool) * valid_composite
flat_data = data.data.reshape(3, -1)
compressed_data = np.compress(valid_composite.flatten(), flat_data, axis=1)
obj = BDO(data, dim_coords=dim)
assert obj._leading_time
assert obj.is_masked
assert obj.data.size == 18 # remove entire nan loc from entire sample
assert np.array_equal(obj.data, compressed_data)
data[~full_valid] = np.nan
np.testing.assert_array_equal(data.filled(np.nan),
obj.inflate_full_grid(reshape_orig=True))
def test_basedataobj_data_force_flat():
data = np.arange(20).reshape(2, 2, 5)
obj = BDO(data, force_flat=True)
assert np.array_equal(data.flatten(), obj.data[:])
assert obj.data.shape == data.flatten().shape
assert obj.orig_data.shape == data.flatten().shape
def test_basedataobj_data_nodim(shape):
data = np.arange(np.product(np.array(shape))).reshape(shape)
obj = BDO(data)
assert np.array_equal(obj.data, data)
assert obj.data.shape == data.shape
assert obj._full_shp == data.shape