def cf_1gm_geos_b_a():
return xr.Dataset(
{
"Rad":
xr.DataArray(
da.zeros((OTHER_DIM_SIZE, Y_DIM_SIZE, X_DIM_SIZE)),
dims=("other", "b", "a"),
attrs={"grid_mapping": "goes_imager_projection"},
)
},
coords={
"b":
xr.DataArray(
da.linspace(0.1265, 0.04257, Y_DIM_SIZE),
dims=("b", ),
attrs={"units": "rad"},
),
"a":
xr.DataArray(
da.linspace(-0.07503, 0.06495, X_DIM_SIZE),
dims=("a", ),
attrs={"units": "rad"},
),
"t":
np.array("2017-09-02T18:03:34", dtype=np.datetime64),
"band_id":
xr.DataArray(np.array([1], dtype=np.uint8),
dims=("band", ),
attrs={"units": "1"}),
},
)
def raw_coords_lats1d_lons1d():
return xr.DataArray(
da.empty((Y_DIM_SIZE, X_DIM_SIZE)),
dims=("lats", "lons"),
coords={
"lons": da.linspace(25, 35, X_DIM_SIZE),
"lats": da.linspace(25, 35, Y_DIM_SIZE),
},
)
def test_linspace(endpoint):
darr = da.linspace(6, 49, endpoint=endpoint, chunks=5)
nparr = np.linspace(6, 49, endpoint=endpoint)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13)
nparr = np.linspace(1.4, 4.9, endpoint=endpoint, num=13)
assert_eq(darr, nparr)
darr = da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float)
nparr = np.linspace(6, 49, endpoint=endpoint, dtype=float)
assert_eq(darr, nparr)
darr, dstep = da.linspace(6, 49, endpoint=endpoint, chunks=5, retstep=True)
nparr, npstep = np.linspace(6, 49, endpoint=endpoint, retstep=True)
assert np.allclose(dstep, npstep)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13, dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, endpoint=endpoint, dtype=int)
assert_eq(darr, nparr)
assert (sorted(da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13).dask) ==
sorted(da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13).dask))
assert (sorted(da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float).dask) ==
sorted(da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float).dask))
def test_linspace(endpoint):
darr = da.linspace(6, 49, endpoint=endpoint, chunks=5)
nparr = np.linspace(6, 49, endpoint=endpoint)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13)
nparr = np.linspace(1.4, 4.9, endpoint=endpoint, num=13)
assert_eq(darr, nparr)
darr = da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float)
nparr = np.linspace(6, 49, endpoint=endpoint, dtype=float)
assert_eq(darr, nparr)
darr, dstep = da.linspace(6, 49, endpoint=endpoint, chunks=5, retstep=True)
nparr, npstep = np.linspace(6, 49, endpoint=endpoint, retstep=True)
assert np.allclose(dstep, npstep)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13, dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, endpoint=endpoint, dtype=int)
assert_eq(darr, nparr)
assert sorted(
da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13).dask
) == sorted(da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13).dask)
assert sorted(
da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float).dask
) == sorted(da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float).dask)
def _perlin_dask_numpy(data: da.Array, freq: tuple, seed: int) -> da.Array:
np.random.seed(seed)
p = np.random.permutation(2**20)
p = np.append(p, p)
height, width = data.shape
linx = da.linspace(0, freq[0], width, endpoint=False, dtype=np.float32)
liny = da.linspace(0, freq[1], height, endpoint=False, dtype=np.float32)
x, y = da.meshgrid(linx, liny)
_func = partial(_perlin, p)
data = da.map_blocks(_func, x, y, meta=np.array((), dtype=np.float32))
data = (data - da.min(data)) / da.ptp(data)
return data
def test_1d_tiles_from_coords_chunks(self):
gm = GridMapping.from_coords(
x_coords=xr.DataArray(da.linspace(177.5, 184.5, 8, chunks=4),
dims='lon'),
y_coords=xr.DataArray(da.linspace(4.5, -4.5, 10, chunks=5),
dims='lat'),
crs=GEO_CRS)
self.assertEqual((8, 10), gm.size)
self.assertEqual((4, 5), gm.tile_size)
self.assertEqual((1, 1), gm.xy_res)
self.assertEqual((177, -5, 185, 5), gm.xy_bbox)
self.assertEqual(GEO_CRS, gm.crs)
self.assertEqual(True, gm.is_regular)
self.assertEqual(False, gm.is_j_axis_up)
self.assertEqual(True, gm.is_lon_360)
def generate(self):
"""
Sub-classable method for generating a factorial design of specified 'levels' in the given domain.
The number of generated points is levels^d.
Returns
-------
dask.delayed
"""
if hasattr(self, 'random_idx'):
del self.random_idx
# Get grid coordinates
grid_coords = [
da.linspace(lb, ub, num=self.levels)
for lb, ub in zip(self.xmin, self.xmax)
]
# Generate the full grid
x = da.meshgrid(*grid_coords)
dim_idx = [item.ravel() for item in x]
x = da.vstack(dim_idx).T
x = x.rechunk(('auto', x.shape[1]))
if self.use_logger:
self.logger.info(
"Factorial design: generated {0} points in {1} dimensions".
format(len(x), len(self.xmin)))
self.generated = x
return x
def discretize(x_data, n_state, min_=0.0, max_=1.0, chunks=()):
"""Primitive discretization of a microstructure.
Args:
x_data: the data to discrtize
n_state: the number of local states
min_: the minimum local state
max_: the maximum local state
Returns:
the discretized microstructure
>>> discretize(da.random.random((12, 9), chunks=(3, 9)),
... 3,
... chunks=(1,)).chunks
((3, 3, 3, 3), (9,), (1, 1, 1))
>>> discretize(np.array([[0, 1], [0.5, 0.5]]), 3, chunks=(1,)).chunks
((2,), (2,), (1, 1, 1))
>>> discretize(np.array([[0, 1], [0.5, 0.5]]), 3, chunks=(1,)).compute()
array([[[ 1., 0., 0.],
[ 0., 0., 1.]],
<BLANKLINE>
[[ 0., 1., 0.],
[ 0., 1., 0.]]])
"""
return da.maximum(
discretize_nomax(
da.clip(x_data, min_, max_),
da.linspace(min_, max_, n_state, chunks=chunks or (n_state, ))), 0)
def plot_hyperplane(clf, min_x, max_x, linestyle, label):
# get the separating hyperplane
w = clf.coef_[0]
a = -w[0] / w[1]
xx = da.linspace(min_x - 5, max_x + 5) # make sure the line is long enough
yy = a * xx - (clf.intercept_[0]) / w[1]
plt.plot(xx, yy, linestyle, label=label, linewidth=2)
def test_linspace():
darr = da.linspace(6, 49, chunks=5)
nparr = np.linspace(6, 49)
eq(darr, nparr)
darr = da.linspace(1.4, 4.9, chunks=5, num=13)
nparr = np.linspace(1.4, 4.9, num=13)
eq(darr, nparr)
darr = da.linspace(6, 49, chunks=5, dtype=float)
nparr = np.linspace(6, 49, dtype=float)
eq(darr, nparr)
darr = da.linspace(1.4, 4.9, chunks=5, num=13, dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, dtype=int)
eq(darr, nparr)
def test_calibrate(self, *mocks):
"""Test calibrate."""
lut = np.linspace(1e6, 1.6e6, num=1024).astype(np.int32)
lut = np.tile(lut, (10, 1))
fh = HRITGOMSFileHandler()
fh.prologue = {'ImageCalibration': lut}
fh.chid = 1
# Set up test input data
counts = DataArray(
da.linspace(1, 1023, 25, chunks=5, dtype=np.uint16).reshape(5, 5))
# Test that calibration fails if given a silly mode
self.assertRaises(NotImplementedError, fh.calibrate, counts,
'nonsense')
# Test that 'counts' calibration returns identical values to input
out = fh.calibrate(counts, 'counts')
self.assertTrue(np.all(out.values == counts.values))
# Test that 'radiance' calibrates successfully
out = fh.calibrate(counts, 'radiance')
self.assertTrue(np.allclose(out.values, lut[0, counts] / 1000.))
# Test that 'brightness_temperature' calibrates successfully
out = fh.calibrate(counts, 'brightness_temperature')
self.assertTrue(np.allclose(out.values, lut[0, counts] / 1000.))
def test_linspace():
darr = da.linspace(6, 49, blocksize=5)
nparr = np.linspace(6, 49)
eq(darr, nparr)
darr = da.linspace(1.4, 4.9, blocksize=5, num=13)
nparr = np.linspace(1.4, 4.9, num=13)
eq(darr, nparr)
darr = da.linspace(6, 49, blocksize=5, dtype=float)
nparr = np.linspace(6, 49, dtype=float)
eq(darr, nparr)
darr = da.linspace(1.4, 4.9, blocksize=5, num=13, dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, dtype=int)
eq(darr, nparr)
def atm_variables_finder(mus,
muv,
phi,
height,
tau,
tO3,
tH2O,
taustep4sphalb,
tO2=1.0):
tau_step = da.linspace(taustep4sphalb,
MAXNUMSPHALBVALUES * taustep4sphalb,
MAXNUMSPHALBVALUES,
chunks=int(MAXNUMSPHALBVALUES / 2))
sphalb0 = csalbr(tau_step)
taur = tau * da.exp(-height / SCALEHEIGHT)
rhoray, trdown, trup = chand(phi, muv, mus, taur)
if isinstance(height, xr.DataArray):
def _sphalb_index(index_arr, sphalb0):
# FIXME: if/when dask can support lazy index arrays then remove this
return sphalb0[index_arr]
sphalb = da.map_blocks(_sphalb_index, (taur / taustep4sphalb +
0.5).astype(np.int32).data,
sphalb0.compute(),
dtype=sphalb0.dtype)
else:
sphalb = sphalb0[(taur / taustep4sphalb + 0.5).astype(np.int32)]
Ttotrayu = ((2 / 3. + muv) + (2 / 3. - muv) * trup) / (4 / 3. + taur)
Ttotrayd = ((2 / 3. + mus) + (2 / 3. - mus) * trdown) / (4 / 3. + taur)
TtotraytH2O = Ttotrayu * Ttotrayd * tH2O
tOG = tO3 * tO2
return sphalb, rhoray, TtotraytH2O, tOG
def test_linspace(endpoint):
darr = da.linspace(6, 49, endpoint=endpoint, chunks=5)
nparr = np.linspace(6, 49, endpoint=endpoint)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5, num=13)
nparr = np.linspace(1.4, 4.9, endpoint=endpoint, num=13)
assert_eq(darr, nparr)
darr = da.linspace(6, 49, endpoint=endpoint, chunks=5, dtype=float)
nparr = np.linspace(6, 49, endpoint=endpoint, dtype=float)
assert_eq(darr, nparr)
darr, dstep = da.linspace(6, 49, endpoint=endpoint, chunks=5, retstep=True)
nparr, npstep = np.linspace(6, 49, endpoint=endpoint, retstep=True)
assert np.allclose(dstep, npstep)
assert_eq(darr, nparr)
darr = da.linspace(1.4,
4.9,
endpoint=endpoint,
chunks=5,
num=13,
dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, endpoint=endpoint, dtype=int)
assert_eq(darr, nparr)
assert sorted(
da.linspace(1.4, 4.9, endpoint=endpoint, chunks=5,
num=13).dask) == sorted(
da.linspace(1.4,
4.9,
endpoint=endpoint,
chunks=5,
num=13).dask)
assert sorted(
da.linspace(6, 49, endpoint=endpoint, chunks=5,
dtype=float).dask) == sorted(
da.linspace(6,
49,
endpoint=endpoint,
chunks=5,
dtype=float).dask)
x = da.array([0.2, 6.4, 3.0, 1.6])
nparr = np.linspace(0, 2, 8, endpoint=endpoint)
darr = da.linspace(da.argmin(x), da.argmax(x) + 1, 8, endpoint=endpoint)
assert_eq(darr, nparr)
def _terrain_dask_numpy(data: da.Array,
seed: int,
x_range_scaled: tuple,
y_range_scaled: tuple,
zfactor: int) -> da.Array:
data = data * 0
height, width = data.shape
linx = da.linspace(
x_range_scaled[0], x_range_scaled[1], width, endpoint=False,
dtype=np.float32
)
liny = da.linspace(
y_range_scaled[0], y_range_scaled[1], height, endpoint=False,
dtype=np.float32
)
x, y = da.meshgrid(linx, liny)
nrange = np.arange(2 ** 20, dtype=int)
# multiplier, (xfreq, yfreq)
NOISE_LAYERS = ((1 / 2 ** i, (2 ** i, 2 ** i)) for i in range(16))
for i, (m, (xfreq, yfreq)) in enumerate(NOISE_LAYERS):
np.random.seed(seed + i)
p = np.random.permutation(nrange)
p = np.append(p, p)
_func = partial(_perlin, p)
noise = da.map_blocks(
_func,
x * xfreq,
y * yfreq,
meta=np.array((), dtype=np.float32)
)
data += noise * m
data /= (1.00 + 0.50 + 0.25 + 0.13 + 0.06 + 0.03)
data = data ** 3
data = (data - np.min(data)) / np.ptp(data)
data[data < 0.3] = 0 # create water
data *= zfactor
return data
def gx_y_x():
crs = CRS.from_epsg(4326)
return xr.DataArray(
da.empty((Y_DIM_SIZE, X_DIM_SIZE)),
dims=("y", "x"),
attrs={
"grid_mapping": "spatial_ref",
},
coords={
"spatial_ref": xr.DataArray(
0,
attrs={
"crs_wkt": crs.to_wkt(),
"spatial_ref": crs.to_wkt(),
},
),
"y": da.linspace(0, 15000, X_DIM_SIZE),
"x": da.linspace(-15000, 10000, Y_DIM_SIZE),
},
)
def _new_xy_coords(self) -> xr.DataArray:
self._assert_regular()
x_res_05, y_res_05 = self.x_res / 2, self.y_res / 2
x1, x2 = self.x_min + x_res_05, self.x_max - x_res_05
y1, y2 = self.y_min + y_res_05, self.y_max - y_res_05
if not self.is_j_axis_up:
y1, y2 = y2, y1
x_name, y_name = self.xy_dim_names
x_coords_1d = xr.DataArray(da.linspace(x1, x2, self.width,
chunks=self.tile_width),
dims=x_name)
y_coords_1d = xr.DataArray(da.linspace(y1, y2, self.height,
chunks=self.tile_height),
dims=y_name)
y_coords_2d, x_coords_2d = xr.broadcast(y_coords_1d, x_coords_1d)
xy_coords = xr.concat([x_coords_2d, y_coords_2d],
dim='coord').chunk((2,
self.tile_height,
self.tile_width))
xy_coords.name = 'xy_coords'
return xy_coords
def test_calibrate(self):
"""Test calibration"""
# Generate test data
counts = DataArray(da.linspace(0, 1200, 25, chunks=5).reshape(5, 5))
refl = np.array(
[[np.nan, 4.79247312, 9.68494624, 14.57741935, 19.46989247],
[24.36236559, 29.25483871, 34.14731183, 39.03978495, 43.93225806],
[48.82473118, 53.7172043, 58.60967742, 63.50215054, 68.39462366],
[73.28709677, 78.17956989, 83.07204301, 87.96451613, 92.85698925],
[97.74946237, 100., 100., 100., 100.]])
bt = np.array(
[[np.nan, 320.20678397, 310.43356794, 300.66035191, 290.88713587],
[
281.11391984, 271.34070381, 261.56748778, 251.79427175,
242.02105572
],
[
232.24783969, 222.47462366, 212.70140762, 202.92819159,
193.15497556
],
[
183.38175953, 173.6085435, 163.83532747, 154.06211144,
144.28889541
], [134.51567937, 130.02, 130.02, 130.02, 130.02]])
# Choose an area near the subsatellite point to avoid masking
# of space pixels
mda = self._get_mda(nlines=5,
ncols=5,
loff=1375.0,
coff=1375.0,
segno=0)
reader = self._get_reader(mda=mda)
# 1. Counts
res = reader.calibrate(data=counts, calibration='counts')
self.assertTrue(np.all(counts.values == res.values))
# 2. Reflectance
res = reader.calibrate(data=counts, calibration='reflectance')
np.testing.assert_allclose(refl, res.values) # also compares NaN
# 3. Brightness temperature
mda_bt = self._get_mda(nlines=5,
ncols=5,
loff=1375.0,
coff=1375.0,
segno=0,
vis=False)
reader_bt = self._get_reader(mda=mda_bt)
res = reader_bt.calibrate(data=counts,
calibration='brightness_temperature')
np.testing.assert_allclose(bt, res.values) # also compares NaN
def test_dataset_create_table(tmp_path, dataset_chunks, dtype):
datasets = []
names = []
datas = []
row_sum = 0
for chunks in dataset_chunks:
shapes = {k: sum(c) for k, c in chunks.items()}
row_sum += shapes['row']
# Make some visibilities
dims = ("row", "chan", "corr")
shape = tuple(shapes[d] for d in dims)
data_chunks = tuple(chunks[d] for d in dims)
data = da.random.random(shape, chunks=data_chunks).astype(dtype)
data_var = Variable(dims, data, {})
# Make some string names
dims = ("row", )
shape = tuple(shapes[d] for d in dims)
str_chunks = tuple(chunks[d] for d in dims)
np_str_array = np.asarray(["BOB"] * shape[0], dtype=np.object)
da_str_array = da.from_array(np_str_array, chunks=str_chunks)
str_array_var = Variable(dims, da_str_array, {})
datasets.append(Dataset({"DATA": data_var, "NAMES": str_array_var}))
datas.append(data)
names.extend(np_str_array.tolist())
freq = da.linspace(.856e9, 2 * .856e9, 64, chunks=16)
sub_datasets = [Dataset({"FREQ": (("row", "chan"), freq[None, :])})]
# Write the data to new tables
table_name = os.path.join(str(tmp_path), 'test.table')
writes = write_datasets(table_name, datasets, ["DATA", "NAMES"])
subt_writes = write_datasets(table_name + "::SPW", sub_datasets, ["FREQ"])
dask.compute(writes, subt_writes)
# Check written data
with pt.table(table_name, readonly=True, lockoptions='auto',
ack=False) as T:
assert row_sum == T.nrows()
assert_array_equal(T.getcol("DATA"), np.concatenate(datas))
assert_array_equal(T.getcol("NAMES"), names)
# Sub-table correctly linked and populated
with pt.table(table_name + "::SPW",
readonly=True,
lockoptions='auto',
ack=False) as T:
assert T.nrows() == 1
assert_array_equal(T.getcol("FREQ")[0], freq)
def test_linspace():
darr = da.linspace(6, 49, chunks=5)
nparr = np.linspace(6, 49)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, chunks=5, num=13)
nparr = np.linspace(1.4, 4.9, num=13)
assert_eq(darr, nparr)
darr = da.linspace(6, 49, chunks=5, dtype=float)
nparr = np.linspace(6, 49, dtype=float)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, chunks=5, num=13, dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, dtype=int)
assert_eq(darr, nparr)
assert (sorted(da.linspace(1.4, 4.9, chunks=5, num=13).dask) ==
sorted(da.linspace(1.4, 4.9, chunks=5, num=13).dask))
assert (sorted(da.linspace(6, 49, chunks=5, dtype=float).dask) ==
sorted(da.linspace(6, 49, chunks=5, dtype=float).dask))
def test_linspace():
darr = da.linspace(6, 49, chunks=5)
nparr = np.linspace(6, 49)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, chunks=5, num=13)
nparr = np.linspace(1.4, 4.9, num=13)
assert_eq(darr, nparr)
darr = da.linspace(6, 49, chunks=5, dtype=float)
nparr = np.linspace(6, 49, dtype=float)
assert_eq(darr, nparr)
darr = da.linspace(1.4, 4.9, chunks=5, num=13, dtype=int)
nparr = np.linspace(1.4, 4.9, num=13, dtype=int)
assert_eq(darr, nparr)
assert (sorted(da.linspace(1.4, 4.9, chunks=5, num=13).dask) == sorted(
da.linspace(1.4, 4.9, chunks=5, num=13).dask))
assert (sorted(da.linspace(6, 49, chunks=5, dtype=float).dask) == sorted(
da.linspace(6, 49, chunks=5, dtype=float).dask))
def test_temperature_difference(self, tmpdir, abi_l1b_c01_data_array):
new_data_arr = abi_l1b_c01_data_array.copy()
data = da.linspace(-10, 10,
new_data_arr.size).reshape(new_data_arr.shape)
new_data_arr.data = data
new_data_arr.attrs["name"] = "test_temperature_difference"
scn = Scene()
scn["test_temperature_difference"] = new_data_arr
out_fn = str(tmpdir + "test_temperature_difference.tif")
scn.save_datasets(filename=out_fn)
with rasterio.open(out_fn, "r") as out_ds:
assert out_ds.count == 2
l_data = out_ds.read(1)
# see polar2grid/tests/etc/enhancements/generic.yaml
flat_l_data = l_data.ravel()
data = data.ravel().compute()
exp_out = np.round(np.linspace(5.0, 205.0,
data.size)).astype(np.uint8)
np.testing.assert_allclose(flat_l_data, exp_out)
def test_p2g_palettize(self, keep_palette, ds_name, tmpdir,
abi_l1b_c01_data_array):
if ds_name == "test_p2g_palettize3":
shutil.copy(os.path.join(TEST_ETC_DIR, f"{ds_name}.npy"), tmpdir)
new_data_arr = abi_l1b_c01_data_array.copy()
data = da.linspace(180, 280,
new_data_arr.size).reshape(new_data_arr.shape)
new_data_arr.data = data
new_data_arr.attrs["name"] = ds_name
scn = Scene()
scn[ds_name] = new_data_arr
out_fn = str(tmpdir + f"{ds_name}_{keep_palette}.tif")
with easy_cwd(tmpdir):
scn.save_datasets(filename=out_fn, keep_palette=keep_palette)
with rasterio.open(out_fn, "r") as out_ds:
is_palette = keep_palette and "palettize" in ds_name
num_bands = 1 if is_palette else 4
assert out_ds.count == num_bands
if is_palette:
assert out_ds.colormap(1) is not None
def discretize(x_data, n_state=2, min_=0.0, max_=1.0, chunks=None):
"""Primitive discretization of a microstructure.
Args:
x_data: the data to discretize
n_state: the number of local states
min_: the minimum local state
max_: the maximum local state
chunks: chunks size for state axis
Returns:
the discretized microstructure
>>> discretize(da.random.random((12, 9), chunks=(3, 9)),
... 3,
... chunks=1).chunks
((3, 3, 3, 3), (9,), (1, 1, 1))
>>> discretize(np.array([[0, 1], [0.5, 0.5]]), 3, chunks=1).chunks
((2,), (2,), (1, 1, 1))
>>> assert np.allclose(
... discretize(
... np.array([[0, 1], [0.5, 0.5]]),
... 3,
... chunks=1
... ).compute(),
... [[[1, 0, 0], [0, 0, 1]], [[0, 1, 0], [0, 1, 0]]]
... )
"""
return da.maximum(
discretize_nomax(
da.clip(x_data, min_, max_),
da.linspace(min_, max_, n_state, chunks=(chunks or n_state,)),
),
0,
)
def test_dask_degridder_gridder():
from africanus.filters import convolution_filter
from africanus.gridding.simple.dask import grid, degrid
import dask.array as da
row = 100
chan = 16
corr = (2, 2)
nx = ny = 1024
row_chunk = 25
chan_chunk = 4
corr_chunk = corr
vis_shape = (row, chan) + corr
vis_chunks = (row_chunk, chan_chunk) + corr_chunk
vis = (da.random.random(vis_shape, chunks=vis_chunks) +
1j * da.random.random(vis_shape, chunks=vis_chunks))
uvw = da.random.random((row, 3), chunks=(row_chunk, 3))
# 4 channels of MeerKAT L band
ref_wave = lightspeed / da.linspace(
.856e9, .856e9 * 2, chan, chunks=chan_chunk)
flags = da.random.randint(0, 1, size=vis_shape, chunks=vis_chunks)
weights = da.random.random(vis_shape, chunks=vis_chunks)
conv_filter = convolution_filter(3, 63, "sinc")
vis_grid = grid(vis, uvw, flags, weights, ref_wave, conv_filter, ny, nx)
degrid_vis = degrid(vis_grid, uvw, weights, ref_wave, conv_filter)
np_vis_grid, np_degrid_vis = da.compute(vis_grid, degrid_vis)
assert np_vis_grid.shape == (ny, nx) + corr
assert np_degrid_vis.shape == (row, chan) + corr
def test_ms_create(Dataset, tmp_path, chunks, num_chans, corr_types, sources):
# Set up
rs = np.random.RandomState(42)
ms_path = tmp_path / "create.ms"
ms_table_name = str(ms_path)
ant_table_name = "::".join((ms_table_name, "ANTENNA"))
ddid_table_name = "::".join((ms_table_name, "DATA_DESCRIPTION"))
pol_table_name = "::".join((ms_table_name, "POLARIZATION"))
spw_table_name = "::".join((ms_table_name, "SPECTRAL_WINDOW"))
# SOURCE is an optional MS sub-table
src_table_name = "::".join((ms_table_name, "SOURCE"))
ms_datasets = []
ant_datasets = []
ddid_datasets = []
pol_datasets = []
spw_datasets = []
src_datasets = []
# For comparison
all_data_desc_id = []
all_data = []
# Create ANTENNA dataset of 64 antennas
# Each column in the ANTENNA has a fixed shape so we
# can represent all rows with one dataset
na = 64
position = da.random.random((na, 3)) * 10000
offset = da.random.random((na, 3))
names = np.array(['ANTENNA-%d' % i for i in range(na)], dtype=np.object)
ds = Dataset({
'POSITION': (("row", "xyz"), position),
'OFFSET': (("row", "xyz"), offset),
'NAME': (("row", ), da.from_array(names, chunks=na)),
})
ant_datasets.append(ds)
# Create SOURCE datasets
for s, (name, direction, rest_freq) in enumerate(sources):
dask_num_lines = da.full((1, ), len(rest_freq), dtype=np.int32)
dask_direction = da.asarray(direction)[None, :]
dask_rest_freq = da.asarray(rest_freq)[None, :]
dask_name = da.asarray(np.asarray([name], dtype=np.object))
ds = Dataset({
"NUM_LINES": (("row", ), dask_num_lines),
"NAME": (("row", ), dask_name),
"REST_FREQUENCY": (("row", "line"), dask_rest_freq),
"DIRECTION": (("row", "dir"), dask_direction),
})
src_datasets.append(ds)
# Create POLARISATION datasets.
# Dataset per output row required because column shapes are variable
for r, corr_type in enumerate(corr_types):
dask_num_corr = da.full((1, ), len(corr_type), dtype=np.int32)
dask_corr_type = da.from_array(corr_type,
chunks=len(corr_type))[None, :]
ds = Dataset({
"NUM_CORR": (("row", ), dask_num_corr),
"CORR_TYPE": (("row", "corr"), dask_corr_type),
})
pol_datasets.append(ds)
# Create multiple MeerKAT L-band SPECTRAL_WINDOW datasets
# Dataset per output row required because column shapes are variable
for num_chan in num_chans:
dask_num_chan = da.full((1, ), num_chan, dtype=np.int32)
dask_chan_freq = da.linspace(.856e9,
2 * .856e9,
num_chan,
chunks=num_chan)[None, :]
dask_chan_width = da.full((1, num_chan), .856e9 / num_chan)
ds = Dataset({
"NUM_CHAN": (("row", ), dask_num_chan),
"CHAN_FREQ": (("row", "chan"), dask_chan_freq),
"CHAN_WIDTH": (("row", "chan"), dask_chan_width),
})
spw_datasets.append(ds)
# For each cartesian product of SPECTRAL_WINDOW and POLARIZATION
# create a corresponding DATA_DESCRIPTION.
# Each column has fixed shape so we handle all rows at once
spw_ids, pol_ids = zip(
*product(range(len(num_chans)), range(len(corr_types))))
dask_spw_ids = da.asarray(np.asarray(spw_ids, dtype=np.int32))
dask_pol_ids = da.asarray(np.asarray(pol_ids, dtype=np.int32))
ddid_datasets.append(
Dataset({
"SPECTRAL_WINDOW_ID": (("row", ), dask_spw_ids),
"POLARIZATION_ID": (("row", ), dask_pol_ids),
}))
# Now create the associated MS dataset
for ddid, (spw_id, pol_id) in enumerate(zip(spw_ids, pol_ids)):
# Infer row, chan and correlation shape
row = sum(chunks['row'])
chan = spw_datasets[spw_id].CHAN_FREQ.shape[1]
corr = pol_datasets[pol_id].CORR_TYPE.shape[1]
# Create some dask vis data
dims = ("row", "chan", "corr")
np_data = (rs.normal(size=(row, chan, corr)) +
1j * rs.normal(size=(row, chan, corr))).astype(np.complex64)
data_chunks = tuple((chunks['row'], chan, corr))
dask_data = da.from_array(np_data, chunks=data_chunks)
# Create dask ddid column
dask_ddid = da.full(row, ddid, chunks=chunks['row'], dtype=np.int32)
dataset = Dataset({
'DATA': (dims, dask_data),
'DATA_DESC_ID': (("row", ), dask_ddid)
})
ms_datasets.append(dataset)
all_data.append(dask_data)
all_data_desc_id.append(dask_ddid)
ms_writes = xds_to_table(ms_datasets, ms_table_name, columns="ALL")
ant_writes = xds_to_table(ant_datasets, ant_table_name, columns="ALL")
pol_writes = xds_to_table(pol_datasets, pol_table_name, columns="ALL")
spw_writes = xds_to_table(spw_datasets, spw_table_name, columns="ALL")
ddid_writes = xds_to_table(ddid_datasets, ddid_table_name, columns="ALL")
source_writes = xds_to_table(src_datasets, src_table_name, columns="ALL")
dask.compute(ms_writes, ant_writes, pol_writes, spw_writes, ddid_writes,
source_writes)
# Check ANTENNA table correctly created
with pt.table(ant_table_name, ack=False) as A:
assert_array_equal(A.getcol("NAME"), names)
assert_array_equal(A.getcol("POSITION"), position)
assert_array_equal(A.getcol("OFFSET"), offset)
required_desc = pt.required_ms_desc("ANTENNA")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(A.colnames()) == set(required_columns)
# Check POLARIZATION table correctly created
with pt.table(pol_table_name, ack=False) as P:
for r, corr_type in enumerate(corr_types):
assert_array_equal(P.getcol("CORR_TYPE", startrow=r, nrow=1),
[corr_type])
assert_array_equal(P.getcol("NUM_CORR", startrow=r, nrow=1),
[len(corr_type)])
required_desc = pt.required_ms_desc("POLARIZATION")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(P.colnames()) == set(required_columns)
# Check SPECTRAL_WINDOW table correctly created
with pt.table(spw_table_name, ack=False) as S:
for r, num_chan in enumerate(num_chans):
assert_array_equal(
S.getcol("NUM_CHAN", startrow=r, nrow=1)[0], num_chan)
assert_array_equal(
S.getcol("CHAN_FREQ", startrow=r, nrow=1)[0],
np.linspace(.856e9, 2 * .856e9, num_chan))
assert_array_equal(
S.getcol("CHAN_WIDTH", startrow=r, nrow=1)[0],
np.full(num_chan, .856e9 / num_chan))
required_desc = pt.required_ms_desc("SPECTRAL_WINDOW")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(S.colnames()) == set(required_columns)
# We should get a cartesian product out
with pt.table(ddid_table_name, ack=False) as D:
spw_id, pol_id = zip(
*product(range(len(num_chans)), range(len(corr_types))))
assert_array_equal(pol_id, D.getcol("POLARIZATION_ID"))
assert_array_equal(spw_id, D.getcol("SPECTRAL_WINDOW_ID"))
required_desc = pt.required_ms_desc("DATA_DESCRIPTION")
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
assert set(D.colnames()) == set(required_columns)
with pt.table(src_table_name, ack=False) as S:
for r, (name, direction, rest_freq) in enumerate(sources):
assert_array_equal(S.getcol("NAME", startrow=r, nrow=1)[0], [name])
assert_array_equal(S.getcol("REST_FREQUENCY", startrow=r, nrow=1),
[rest_freq])
assert_array_equal(S.getcol("DIRECTION", startrow=r, nrow=1),
[direction])
with pt.table(ms_table_name, ack=False) as T:
# DATA_DESC_ID's are all the same shape
assert_array_equal(T.getcol("DATA_DESC_ID"),
da.concatenate(all_data_desc_id))
# DATA is variably shaped (on DATA_DESC_ID) so we
# compared each one separately.
for ddid, data in enumerate(all_data):
ms_data = T.getcol("DATA", startrow=ddid * row, nrow=row)
assert_array_equal(ms_data, data)
required_desc = pt.required_ms_desc()
required_columns = set(k for k in required_desc.keys()
if not k.startswith("_"))
# Check we have the required columns
assert set(T.colnames()) == required_columns.union(
["DATA", "DATA_DESC_ID"])
def genGridFile(fname="renoGrid.hdf5", iniFname="grid_file.ini"):
#load the ini construction file
with open(iniFname, 'r') as fileObject:
ini = [
line.strip() for line in fileObject if ("#" not in line.strip())
]
ini = json.loads(r''.join(ini))
#create hdf5 grid file
#note that flipping x and y IS NOT A MISTAKE!
y = da.linspace(0,
ini["NX"] * ini["DX"],
ini["NX"],
chunks=(ini["CHUNKX"], )).astype(np.int32)
x = da.linspace(0,
ini["NY"] * ini["DY"],
ini["NY"],
chunks=(ini["CHUNKY"], )).astype(np.int32)
z = da.linspace(0,
ini["NZ"] * ini["DZ"],
ini["NZ"],
chunks=(ini["CHUNKZ"], )).astype(np.int32)
topo = da.zeros(shape=(x.shape[0], y.shape[0]),
chunks=(ini["CHUNKX"], ini["CHUNKY"]))
basinSurface = da.zeros(shape=(x.shape[0], y.shape[0]),
chunks=(ini["CHUNKX"], ini["CHUNKY"]))
#compute a dask shape
#note that I chose to use cartesian indexing
gridCoords = da.meshgrid(x, y, z, sparse=False, indexing='xy')
#try flipping z -- Note that an rFile scans down from a maximum elevation to the bottom (not from bottome to top as the mesh grid does, this corrects that orientation)
#gridCoords[2] = da.flip(gridCoords[2],axis=-1)
#now write to the hdf5 file
da.to_hdf5(fname, '/grid/x', gridCoords[0])
da.to_hdf5(fname, '/grid/y', gridCoords[1])
da.to_hdf5(fname, '/grid/z', gridCoords[2])
#save an empty (bullshit) topography
da.to_hdf5(fname, '/grid/topo', topo)
#create all of the empty coordinate spaces
#-666 is the not interpolated, empty space flag (not to be confuesd with -999 the no value flag for sw4
da.to_hdf5(
fname, '/grid/vp',
da.full(gridCoords[0].shape,
-666,
chunks=gridCoords[0].chunksize,
dtype=np.float32).flatten())
da.to_hdf5(
fname, '/grid/vs',
da.full(gridCoords[0].shape,
-666,
chunks=gridCoords[0].chunksize,
dtype=np.float32).flatten())
da.to_hdf5(
fname, '/grid/p',
da.full(gridCoords[0].shape,
-666,
chunks=gridCoords[0].chunksize,
dtype=np.float32).flatten())
da.to_hdf5(
fname, '/grid/qp',
da.full(gridCoords[0].shape,
-666,
chunks=gridCoords[0].chunksize,
dtype=np.float32).flatten())
da.to_hdf5(
fname, '/grid/qs',
da.full(gridCoords[0].shape,
-666,
chunks=gridCoords[0].chunksize,
dtype=np.float32).flatten())
#build a unit descriptor array
da.to_hdf5(
fname, '/grid/unit',
da.full(gridCoords[0].shape,
-1,
chunks=gridCoords[0].chunksize,
dtype=np.int8).flatten())
#now write the config file to the hdf5 file header
result = h5py.File(fname, 'r+')
#there really must be a better way of doing this! This reencodes the json as ascii to remove any unicode which might be in it
ini = [i.encode("ascii", "ignore") for i in json.dumps(ini)]
result.create_dataset('/grid/ini', (len(ini), 1), 'S10', ini)
print(
"\nPressure fields for the following cases are computed (case, U, a, p0, x0)"
)
with open(f"{current_dir}/parameters_pressure.txt", "w") as file:
for i in range(num_cases):
line = f"{cases[i]:02.0f},{U_list[i]:0.0f},{a_list[i]:0.0f},{p0_list[i]:0.0f},{x0_list[i]:0.0f}"
print(line.replace(",", "\t"))
file.write(line + "\n")
del line, i
### Grid
x_num = int((x_max - x_min) / x_step + 1)
y_num = int((y_max - y_min) / y_step + 1)
x = da.linspace(x_min, x_max, x_num, chunks=256)
y = da.linspace(y_min, y_max, y_num, chunks=1024)
t = da.arange(t_min, t_max + 1, t_step, chunks=128)
t_num = t.size
tt, yy, xx = da.meshgrid(t, y, x, indexing="ij", sparse=True)
print("Grid parameters:")
print(f"{x.size=}\t\t{y.size=}\t\t{t.size=}")
print(f"{x.chunksize=}\t{y.chunksize=}\t{t.chunksize=}")
### Function
def pressure(x, y, t, t0=10000., U=50., a=200000., p0=2000., x0=0.):
return (p0 * (1. - da.exp(-t / t0)) * da.exp(-((x - x0)**2. +
def _lin_axis(min_: Union[float, int], max_: Union[float, int],
points: int) -> da.core.Array:
"""Generates linearly spaced array."""
return da.linspace(min_, max_, points)
def fixture_fake_dataset():
"""Create fake dataset."""
count_ir = da.linspace(0, 255, 4, dtype=np.uint8).reshape(2, 2)
count_wv = da.linspace(0, 255, 4, dtype=np.uint8).reshape(2, 2)
count_vis = da.linspace(0, 255, 16, dtype=np.uint8).reshape(4, 4)
sza = da.from_array(np.array([[45, 90], [0, 45]], dtype=np.float32))
mask = da.from_array(
np.array(
[
[0, 0, 0, 0],
[0, 0, 0, 0],
[0, 0, 1, 0], # 1 = "invalid"
[0, 0, 0, 0]
],
dtype=np.uint8))
time = np.arange(4).astype('datetime64[h]').reshape(2, 2)
ds = xr.Dataset(data_vars={
'count_vis': (('y', 'x'), count_vis),
'count_wv': (('y_ir_wv', 'x_ir_wv'), count_wv),
'count_ir': (('y_ir_wv', 'x_ir_wv'), count_ir),
'toa_bidirectional_reflectance_vis':
vis_refl_exp / 100,
'u_independent_toa_bidirectional_reflectance':
u_vis_refl_exp / 100,
'quality_pixel_bitmask': (('y', 'x'), mask),
'solar_zenith_angle': (('y_tie', 'x_tie'), sza),
'time_ir_wv': (('y_ir_wv', 'x_ir_wv'), time),
'a_ir':
-5.0,
'b_ir':
1.0,
'bt_a_ir':
10.0,
'bt_b_ir':
-1000.0,
'a_wv':
-0.5,
'b_wv':
0.05,
'bt_a_wv':
10.0,
'bt_b_wv':
-2000.0,
'years_since_launch':
20.0,
'a0_vis':
1.0,
'a1_vis':
0.01,
'a2_vis':
-0.0001,
'mean_count_space_vis':
1.0,
'distance_sun_earth':
1.0,
'solar_irradiance_vis':
650.0,
'sub_satellite_longitude_start':
57.1,
'sub_satellite_longitude_end':
np.nan,
'sub_satellite_latitude_start':
np.nan,
'sub_satellite_latitude_end':
0.1,
},
coords={
'y': [1, 2, 3, 4],
'x': [1, 2, 3, 4],
'y_ir_wv': [1, 2],
'x_ir_wv': [1, 2],
'y_tie': [1, 2],
'x_tie': [1, 2]
},
attrs={'foo': 'bar'})
ds['count_ir'].attrs['ancillary_variables'] = 'a_ir b_ir'
ds['count_wv'].attrs['ancillary_variables'] = 'a_wv b_wv'
return ds
def _log_axis(min_: Union[float, int], max_: Union[float, int],
points: int) -> da.core.Array:
"""Generates logarithmically spaced array."""
min_ = np.log10(min_)
max_ = np.log10(max_)
return 10.0**da.linspace(min_, max_, points)
def setUp(self):
"""Create fake data for the tests."""
data = da.linspace(0, 1, 16).reshape((4, 4))
self.da = xr.DataArray(data, dims=('y', 'x'), attrs={'test': 'test'})
