def prop1FT(u, step, L1, wavel, z, fft_object=None):
M, N = np.shape(u)
x = np.linspace(-L1 / 2.0, L1 / 2.0 - step, M)
y = np.linspace(-L1 / 2.0, L1 / 2.0 - step, N)
'''
#Kenan's approach
fx = np.fft.fftfreq(M,d=step)
fy = np.fft.fftfreq(N,d=step)
fx = pyfftw.interfaces.numpy_fft.fftshift(fx)
fy = pyfftw.interfaces.numpy_fft.fftshift(fy)
FX,FY = da.meshgrid((fx),(fy))
c = np.exp((-1j*z*2*np.pi/wavel)*np.sqrt(1+wavel**2*(FX**2+FY**2)))
'''
L_out = wavel * z / step
X, Y = da.meshgrid(x, y)
u = ne.evaluate('exp((-1j*2*pi/wavel)*sqrt(X**2+Y**2+z**2))*u')
del X, Y
if fft_object is not None:
fft_object.run_fft2(u)
else:
u = np.fft.fft2(u)
u = np.fft.fftshift(u)
x2 = np.linspace(-L_out / 2.0, L_out / 2.0, M)
y2 = np.linspace(-L_out / 2.0, L_out / 2.0, N)
X2, Y2 = da.meshgrid(x2, y2)
u = ne.evaluate('exp((-1j*2*pi/wavel)*sqrt(X2**2+Y2**2+z**2))*u')
del X2, Y2
u = ne.evaluate('u*(1j/(wavel*z))*step*step')
return u, L_out
def test_meshgrid(shapes, chunks, indexing, sparse):
xi_a = []
xi_d = []
xi_dc = []
for each_shape, each_chunk in zip(shapes, chunks):
xi_a.append(np.random.random(each_shape))
xi_d_e = da.from_array(xi_a[-1], chunks=each_chunk)
xi_d.append(xi_d_e)
xi_d_ef = xi_d_e.flatten()
xi_dc.append(xi_d_ef.chunks[0])
do = list(range(len(xi_dc)))
if indexing == "xy" and len(xi_dc) > 1:
do[0], do[1] = do[1], do[0]
xi_dc[0], xi_dc[1] = xi_dc[1], xi_dc[0]
xi_dc = tuple(xi_dc)
r_a = np.meshgrid(*xi_a, indexing=indexing, sparse=sparse)
r_d = da.meshgrid(*xi_d, indexing=indexing, sparse=sparse)
assert isinstance(r_d, list)
assert len(r_a) == len(r_d)
for e_r_a, e_r_d, i in zip(r_a, r_d, do):
assert_eq(e_r_a, e_r_d)
if sparse:
assert e_r_d.chunks[i] == xi_dc[i]
else:
assert e_r_d.chunks == xi_dc
def expand_array(data,
scans,
c_align,
c_exp,
scan_size=16,
tpz_size=16,
nties=200,
track_offset=0.5,
scan_offset=0.5):
"""Expand *data* according to alignment and expansion."""
nties = np.asscalar(nties)
tpz_size = np.asscalar(tpz_size)
s_scan, s_track = da.meshgrid(np.arange(nties * tpz_size),
np.arange(scans * scan_size))
s_track = (s_track.reshape(scans, scan_size, nties, tpz_size) % scan_size +
track_offset) / scan_size
s_scan = (s_scan.reshape(scans, scan_size, nties, tpz_size) % tpz_size +
scan_offset) / tpz_size
a_scan = s_scan + s_scan * (1 - s_scan) * c_exp + s_track * (
1 - s_track) * c_align
a_track = s_track
data_a = data[:scans * 2:2, np.newaxis, :-1, np.newaxis]
data_b = data[:scans * 2:2, np.newaxis, 1:, np.newaxis]
data_c = data[1:scans * 2:2, np.newaxis, 1:, np.newaxis]
data_d = data[1:scans * 2:2, np.newaxis, :-1, np.newaxis]
fdata = ((1 - a_track) *
((1 - a_scan) * data_a + a_scan * data_b) + a_track * (
(1 - a_scan) * data_d + a_scan * data_c))
return fdata.reshape(scans * scan_size, nties * tpz_size)
def expand_array(data,
scans,
c_align,
c_exp,
scan_size=16,
tpz_size=16,
nties=200,
track_offset=0.5,
scan_offset=0.5):
"""Expand *data* according to alignment and expansion."""
nties = np.asscalar(nties)
tpz_size = np.asscalar(tpz_size)
s_scan, s_track = da.meshgrid(np.arange(nties * tpz_size),
np.arange(scans * scan_size))
s_track = (s_track.reshape(scans, scan_size, nties, tpz_size) % scan_size +
track_offset) / scan_size
s_scan = (s_scan.reshape(scans, scan_size, nties, tpz_size) % tpz_size +
scan_offset) / tpz_size
a_scan = s_scan + s_scan * (1 - s_scan) * c_exp + s_track * (
1 - s_track) * c_align
a_track = s_track
data_a = data[:scans * 2:2, np.newaxis, :-1, np.newaxis]
data_b = data[:scans * 2:2, np.newaxis, 1:, np.newaxis]
data_c = data[1:scans * 2:2, np.newaxis, 1:, np.newaxis]
data_d = data[1:scans * 2:2, np.newaxis, :-1, np.newaxis]
fdata = ((1 - a_track) * ((1 - a_scan) * data_a + a_scan * data_b) +
a_track * ((1 - a_scan) * data_d + a_scan * data_c))
return fdata.reshape(scans * scan_size, nties * tpz_size)
def interpolate_angles(self, angles, resolution):
"""Interpolate the angles."""
# FIXME: interpolate in cartesian coordinates if the lons or lats are
# problematic
from geotiepoints.multilinear import MultilinearInterpolator
geocoding = self.root.find('.//Tile_Geocoding')
rows = int(
geocoding.find('Size[@resolution="' + str(resolution) +
'"]/NROWS').text)
cols = int(
geocoding.find('Size[@resolution="' + str(resolution) +
'"]/NCOLS').text)
smin = [0, 0]
smax = np.array(angles.shape) - 1
orders = angles.shape
minterp = MultilinearInterpolator(smin, smax, orders)
minterp.set_values(da.atleast_2d(angles.ravel()))
x = da.arange(rows, dtype=angles.dtype,
chunks=CHUNK_SIZE) / (rows - 1) * (angles.shape[0] - 1)
y = da.arange(cols, dtype=angles.dtype,
chunks=CHUNK_SIZE) / (cols - 1) * (angles.shape[1] - 1)
xcoord, ycoord = da.meshgrid(x, y)
return da.map_blocks(self._do_interp,
minterp,
xcoord,
ycoord,
dtype=angles.dtype,
chunks=xcoord.chunks)
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 interpolate_xarray_linear(xpoints,
ypoints,
values,
shape,
chunks=CHUNK_SIZE):
"""Interpolate linearly, generating a dask array."""
from scipy.interpolate.interpnd import (LinearNDInterpolator,
_ndim_coords_from_arrays)
if isinstance(chunks, (list, tuple)):
vchunks, hchunks = chunks
else:
vchunks, hchunks = chunks, chunks
points = _ndim_coords_from_arrays(
np.vstack((np.asarray(ypoints), np.asarray(xpoints))).T)
interpolator = LinearNDInterpolator(points, values)
grid_x, grid_y = da.meshgrid(da.arange(shape[1], chunks=hchunks),
da.arange(shape[0], chunks=vchunks))
# workaround for non-thread-safe first call of the interpolator:
interpolator((0, 0))
res = da.map_blocks(intp, grid_x, grid_y, interpolator=interpolator)
return DataArray(res, dims=('y', 'x'))
def test_meshgrid(shapes, chunks, indexing, sparse):
xi_a = []
xi_d = []
xi_dc = []
for each_shape, each_chunk in zip(shapes, chunks):
xi_a.append(np.random.random(each_shape))
xi_d_e = da.from_array(xi_a[-1], chunks=each_chunk)
xi_d.append(xi_d_e)
xi_d_ef = xi_d_e.flatten()
xi_dc.append(xi_d_ef.chunks[0])
do = list(range(len(xi_dc)))
if indexing == "xy" and len(xi_dc) > 1:
do[0], do[1] = do[1], do[0]
xi_dc[0], xi_dc[1] = xi_dc[1], xi_dc[0]
xi_dc = tuple(xi_dc)
r_a = np.meshgrid(*xi_a, indexing=indexing, sparse=sparse)
r_d = da.meshgrid(*xi_d, indexing=indexing, sparse=sparse)
assert isinstance(r_d, list)
assert len(r_a) == len(r_d)
for e_r_a, e_r_d, i in zip(r_a, r_d, do):
assert_eq(e_r_a, e_r_d)
if sparse:
assert e_r_d.chunks[i] == xi_dc[i]
else:
assert e_r_d.chunks == xi_dc
def position_grid(shape, blocksize):
"""
"""
coords = da.meshgrid(*[range(x) for x in shape], indexing='ij')
coords = da.stack(coords, axis=-1).astype(np.int16)
return da.rechunk(coords, chunks=tuple(blocksize) + (3, ))
def interpolate_angles(self, angles, resolution):
# FIXME: interpolate in cartesian coordinates if the lons or lats are
# problematic
from geotiepoints.multilinear import MultilinearInterpolator
geocoding = self.root.find('.//Tile_Geocoding')
rows = int(geocoding.find('Size[@resolution="' + str(resolution) + '"]/NROWS').text)
cols = int(geocoding.find('Size[@resolution="' + str(resolution) + '"]/NCOLS').text)
smin = [0, 0]
smax = np.array(angles.shape) - 1
orders = angles.shape
minterp = MultilinearInterpolator(smin, smax, orders)
minterp.set_values(da.atleast_2d(angles.ravel()))
def _do_interp(minterp, xcoord, ycoord):
interp_points2 = np.vstack((xcoord.ravel(),
ycoord.ravel()))
res = minterp(interp_points2)
return res.reshape(xcoord.shape)
x = da.arange(rows, dtype=angles.dtype, chunks=CHUNK_SIZE) / (rows-1) * (angles.shape[0] - 1)
y = da.arange(cols, dtype=angles.dtype, chunks=CHUNK_SIZE) / (cols-1) * (angles.shape[1] - 1)
xcoord, ycoord = da.meshgrid(x, y)
return da.map_blocks(_do_interp, minterp, xcoord, ycoord, dtype=angles.dtype,
chunks=xcoord.chunks)
def expand_arrays(arrays,
scans,
c_align,
c_exp,
scan_size=16,
tpz_size=16,
nties=200,
track_offset=0.5,
scan_offset=0.5):
"""Expand *data* according to alignment and expansion."""
nties = nties.item()
tpz_size = tpz_size.item()
s_scan, s_track = da.meshgrid(da.arange(nties * tpz_size),
da.arange(scans * scan_size))
s_track = (s_track.reshape(scans, scan_size, nties, tpz_size) % scan_size
+ track_offset) / scan_size
s_scan = (s_scan.reshape(scans, scan_size, nties, tpz_size) % tpz_size
+ scan_offset) / tpz_size
a_scan = s_scan + s_scan * (1 - s_scan) * c_exp + s_track * (
1 - s_track) * c_align
a_track = s_track
expanded = []
coef_a = (1 - a_track) * (1 - a_scan)
coef_b = (1 - a_track) * a_scan
coef_d = a_track * (1 - a_scan)
coef_c = a_track * a_scan
corner_coefficients = (coef_a, coef_b, coef_c, coef_d)
for data in arrays:
fdata = _interpolate_data(data, corner_coefficients, scans)
expanded.append(fdata.reshape(scans * scan_size, nties * tpz_size))
return expanded
def test_map_iter(self,chunks):
iter_array, _ = da.meshgrid(range(11), range(10))
iter_array = iter_array.rechunk(chunks)
s_iter = hs.signals.BaseSignal(iter_array).T
s_iter = s_iter.as_lazy()
f = lambda a, b: a + b
s_out = self.s.map(function=f, b=s_iter, inplace=False)
np.testing.assert_array_equal(s_out.mean(axis=(2, 3)).data, iter_array)
def apply_distance_using_dask(t1,f1):
# Compute the grid for any configuration.
import numpy as np
import dask.array as da
# gps = da.stack(da.meshgrid(da.linspace(-0.5,0.5, t1.nx_total), da.linspace(-0.5,0.5, t1.ny_total),da.linspace(-0.5,0.5, t1.nz_total)), -1).reshape(-1, 3)
gps = da.stack(da.meshgrid(t1.x_grid, t1.y_grid,t1.z_grid, indexing='ij'), -1).reshape(-1, 3) # * Only the unit cell.
gps =gps.rechunk(10000,3)
grid = da.apply_along_axis(func1d=grid_point_distance, frameda=da.from_array(t1.coord), Ada=da.from_array(t1.A),sig = da.from_array(f1.sigma), sigda=da.from_array(f1.sigma_array), epsda=da.from_array(f1.epsilon_array), axis=1, arr=gps)
return grid
def _get_first_unmasked_data(array, axis):
"""Get first unmasked value of an array along an axis."""
mask = da.ma.getmaskarray(array)
numerical_mask = da.where(mask, -1.0, 1.0)
indices_first_positive = da.argmax(numerical_mask, axis=axis)
indices = da.meshgrid(
*[da.arange(array.shape[i]) for i in range(array.ndim) if i != axis],
indexing='ij')
indices.insert(axis, indices_first_positive)
first_unmasked_data = np.array(array)[tuple(indices)]
return first_unmasked_data
def test_meshgrid_inputcoercion():
a = [1, 2, 3]
b = np.array([4, 5, 6, 7])
x, y = np.meshgrid(a, b, indexing="ij")
z = x * y
x_d, y_d = da.meshgrid(a, b, indexing="ij")
z_d = x_d * y_d
assert z_d.shape == (len(a), len(b))
assert_eq(z, z_d)
def test_meshgrid_inputcoercion():
a = [1, 2, 3]
b = np.array([4, 5, 6, 7])
x, y = np.meshgrid(a, b, indexing='ij')
z = x * y
x_d, y_d = da.meshgrid(a, b, indexing='ij')
z_d = x_d * y_d
assert z_d.shape == (len(a), len(b))
assert_eq(z, z_d)
def interpolate(self, lon1, lat1, satz1):
cscan_len = self.cscan_len
cscan_full_width = self.cscan_full_width
fscan_width = self.fscan_width
fscan_len = self.fscan_len
scans = satz1.shape[0] // cscan_len
satz1 = satz1.data
satz1 = satz1.reshape((-1, cscan_len, cscan_full_width))
satz_a, satz_b, satz_c, satz_d = get_corners(da.deg2rad(satz1))
c_exp, c_ali = compute_expansion_alignment(satz_a, satz_b, satz_c, satz_d)
x, y = self.get_coords(scans)
i_rs, i_rt = da.meshgrid(x, y)
p_os = 0
p_ot = 0
s_s = (p_os + i_rs) * 1. / fscan_width
s_t = (p_ot + i_rt) * 1. / fscan_len
cols = fscan_width
lines = fscan_len
c_exp_full = self.expand_tiepoint_array(c_exp, lines, cols)
c_ali_full = self.expand_tiepoint_array(c_ali, lines, cols)
a_track = s_t
a_scan = (s_s + s_s * (1 - s_s) * c_exp_full + s_t * (1 - s_t) * c_ali_full)
res = []
datasets = lonlat2xyz(lon1, lat1)
for data in datasets:
data = data.data
data = data.reshape((-1, cscan_len, cscan_full_width))
data_a, data_b, data_c, data_d = get_corners(data)
data_a = self.expand_tiepoint_array(data_a, lines, cols)
data_b = self.expand_tiepoint_array(data_b, lines, cols)
data_c = self.expand_tiepoint_array(data_c, lines, cols)
data_d = self.expand_tiepoint_array(data_d, lines, cols)
data_1 = (1 - a_scan) * data_a + a_scan * data_b
data_2 = (1 - a_scan) * data_d + a_scan * data_c
data = (1 - a_track) * data_1 + a_track * data_2
res.append(data)
lon, lat = xyz2lonlat(*res)
return xr.DataArray(lon, dims=lon1.dims), xr.DataArray(lat, dims=lat1.dims)
def test_meshgrid(shapes, indexing, sparse):
xi_a = []
xi_d = []
for each_shape in shapes:
xi_a.append(np.random.random(each_shape))
xi_d.append(da.from_array(xi_a[-1], chunks=1))
r_a = np.meshgrid(*xi_a, indexing=indexing, sparse=sparse)
r_d = da.meshgrid(*xi_d, indexing=indexing, sparse=sparse)
assert isinstance(r_d, list)
assert len(r_a) == len(r_d)
for e_r_a, e_r_d in zip(r_a, r_d):
assert_eq(e_r_a, e_r_d)
def test_meshgrid(shapes, indexing, sparse):
xi_a = []
xi_d = []
for each_shape in shapes:
xi_a.append(np.random.random(each_shape))
xi_d.append(da.from_array(xi_a[-1], chunks=1))
r_a = np.meshgrid(*xi_a, indexing=indexing, sparse=sparse)
r_d = da.meshgrid(*xi_d, indexing=indexing, sparse=sparse)
assert isinstance(r_d, list)
assert len(r_a) == len(r_d)
for e_r_a, e_r_d in zip(r_a, r_d):
assert_eq(e_r_a, e_r_d)
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 propTF(u, step, L1, wavel, z, fft_object=None):
M, N = np.shape(u)
FX, FY = da.meshgrid(np.fft.fftfreq(M, step), np.fft.fftfreq(N, step))
if fft_object is not None:
fft_object.run_fft2(u)
else:
u = np.fft.fft2(u)
u = ne.evaluate('exp(-1j*(2*pi*z/wavel)*sqrt(1-wavel**2*(FX**2+FY**2)))*u')
if fft_object is not None:
fft_object.run_ifft2(u)
else:
u = np.fft.ifft2(u)
return u, L1
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 interpolate_xarray_linear(xpoints, ypoints, values, shape):
"""Interpolate linearly, generating a dask array."""
from scipy.interpolate.interpnd import (LinearNDInterpolator,
_ndim_coords_from_arrays)
points = _ndim_coords_from_arrays(np.vstack((np.asarray(ypoints),
np.asarray(xpoints))).T)
interpolator = LinearNDInterpolator(points, values)
def intp(grid_x, grid_y, interpolator):
return interpolator((grid_y, grid_x))
grid_x, grid_y = da.meshgrid(da.arange(shape[1], chunks=CHUNK_SIZE),
da.arange(shape[0], chunks=CHUNK_SIZE))
# workaround for non-thread-safe first call of the interpolator:
interpolator((0, 0))
res = da.map_blocks(intp, grid_x, grid_y, interpolator=interpolator)
return DataArray(res, dims=('y', 'x'))
def _get_freq_grid(shape, chunks, axis, n, dtype=float):
assert len(shape) == len(chunks)
shape = tuple(shape)
dtype = np.dtype(dtype).type
assert (issubclass(dtype, numbers.Real)
and not issubclass(dtype, numbers.Integral))
axis = axis % len(shape)
freq_grid = []
for ax, (s, c) in enumerate(zip(shape, chunks)):
if axis == ax and n > 0:
f = da.fft.rfftfreq(n, chunks=c).astype(dtype)
else:
f = da.fft.fftfreq(s, chunks=c).astype(dtype)
freq_grid.append(f)
freq_grid = da.meshgrid(*freq_grid, indexing="ij", sparse=True)
return freq_grid
def correlation_nxn(
data: da.Array,
columns: Optional[Sequence[str]] = None
) -> Tuple[da.Array, da.Array, Dict[CorrelationMethod, da.Array]]:
"""
Calculation of a n x n correlation matrix for n columns
"""
_, ncols = data.shape
cordx, cordy = da.meshgrid(range(ncols), range(ncols))
cordx, cordy = cordy.ravel(), cordx.ravel()
corrs = {
CorrelationMethod.Pearson: pearson_nxn(data).ravel(),
CorrelationMethod.Spearman: spearman_nxn(data).ravel(),
CorrelationMethod.KendallTau: kendall_tau_nxn(data).ravel(),
}
if columns is not None:
# The number of columns usually is not too large
cordx = da.from_array(columns[cordx.compute()], chunks=1)
cordy = da.from_array(columns[cordy.compute()], chunks=1)
return cordx, cordy, corrs
def expand_arrays(arrays,
scans,
c_align,
c_exp,
scan_size=16,
tpz_size=16,
nties=200,
track_offset=0.5,
scan_offset=0.5):
"""Expand *data* according to alignment and expansion."""
nties = nties.item()
tpz_size = tpz_size.item()
s_scan, s_track = da.meshgrid(da.arange(nties * tpz_size),
da.arange(scans * scan_size))
s_track = (s_track.reshape(scans, scan_size, nties, tpz_size) % scan_size +
track_offset) / scan_size
s_scan = (s_scan.reshape(scans, scan_size, nties, tpz_size) % tpz_size +
scan_offset) / tpz_size
a_scan = s_scan + s_scan * (1 - s_scan) * c_exp + s_track * (
1 - s_track) * c_align
a_track = s_track
expanded = []
coef_a = (1 - a_track) * (1 - a_scan)
coef_b = (1 - a_track) * a_scan
coef_d = a_track * (1 - a_scan)
coef_c = a_track * a_scan
for data in arrays:
data_a = data[:scans * 2:2, np.newaxis, :-1, np.newaxis]
data_b = data[:scans * 2:2, np.newaxis, 1:, np.newaxis]
data_c = data[1:scans * 2:2, np.newaxis, 1:, np.newaxis]
data_d = data[1:scans * 2:2, np.newaxis, :-1, np.newaxis]
fdata = (coef_a * data_a + coef_b * data_b + coef_d * data_d +
coef_c * data_c)
expanded.append(fdata.reshape(scans * scan_size, nties * tpz_size))
return expanded
def interpolate(self, lon1, lat1, satz1):
cscan_len = self.cscan_len
cscan_full_width = self.cscan_full_width
fscan_width = self.fscan_width
fscan_len = self.fscan_len
scans = satz1.shape[0] // cscan_len
satz1 = satz1.data
satz1 = satz1.reshape((-1, cscan_len, cscan_full_width))
satz_a, satz_b, satz_c, satz_d = get_corners(da.deg2rad(satz1))
c_exp, c_ali = compute_expansion_alignment(satz_a, satz_b, satz_c, satz_d)
x, y = self.get_coords(scans)
i_rs, i_rt = da.meshgrid(x, y)
p_os = 0
p_ot = 0
s_s = (p_os + i_rs) * 1. / fscan_width
s_t = (p_ot + i_rt) * 1. / fscan_len
cols = fscan_width
lines = fscan_len
c_exp_full = self.expand_tiepoint_array(c_exp, lines, cols)
c_ali_full = self.expand_tiepoint_array(c_ali, lines, cols)
a_track = s_t
a_scan = (s_s + s_s * (1 - s_s) * c_exp_full + s_t * (1 - s_t) * c_ali_full)
res = []
sublat = lat1[::16, ::16]
sublon = lon1[::16, ::16]
to_cart = abs(sublat).max() > 60 or (sublon.max() - sublon.min()) > 180
if to_cart:
datasets = lonlat2xyz(lon1, lat1)
else:
datasets = [lon1, lat1]
for data in datasets:
data_attrs = data.attrs
dims = data.dims
data = data.data
data = data.reshape((-1, cscan_len, cscan_full_width))
data_a, data_b, data_c, data_d = get_corners(data)
data_a = self.expand_tiepoint_array(data_a, lines, cols)
data_b = self.expand_tiepoint_array(data_b, lines, cols)
data_c = self.expand_tiepoint_array(data_c, lines, cols)
data_d = self.expand_tiepoint_array(data_d, lines, cols)
data_1 = (1 - a_scan) * data_a + a_scan * data_b
data_2 = (1 - a_scan) * data_d + a_scan * data_c
data = (1 - a_track) * data_1 + a_track * data_2
res.append(xr.DataArray(data, attrs=data_attrs, dims=dims))
if to_cart:
return xyz2lonlat(*res)
else:
return res
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)
def interpolate(self, lon1, lat1, satz1):
cscan_len = self.cscan_len
cscan_full_width = self.cscan_full_width
fscan_width = self.fscan_width
fscan_len = self.fscan_len
scans = lat1.shape[0] // cscan_len
latattrs = lat1.attrs
lonattrs = lon1.attrs
dims = lat1.dims
lat1 = lat1.data
lon1 = lon1.data
satz1 = satz1.data
lat1 = lat1.reshape((-1, cscan_len, cscan_full_width))
lon1 = lon1.reshape((-1, cscan_len, cscan_full_width))
satz1 = satz1.reshape((-1, cscan_len, cscan_full_width))
lats_a, lats_b, lats_c, lats_d = get_corners(lat1)
lons_a, lons_b, lons_c, lons_d = get_corners(lon1)
satz_a, satz_b, satz_c, satz_d = get_corners(da.deg2rad(satz1))
c_exp, c_ali = compute_expansion_alignment(satz_a, satz_b, satz_c,
satz_d)
x, y = self.get_coords(scans)
i_rs, i_rt = da.meshgrid(x, y)
p_os = 0
p_ot = 0
s_s = (p_os + i_rs) * 1. / fscan_width
s_t = (p_ot + i_rt) * 1. / fscan_len
cols = fscan_width
lines = fscan_len
c_exp_full = self.expand_tiepoint_array(c_exp, lines, cols)
c_ali_full = self.expand_tiepoint_array(c_ali, lines, cols)
a_track = s_t
a_scan = (s_s + s_s * (1 - s_s) * c_exp_full + s_t *
(1 - s_t) * c_ali_full)
lats_a = self.expand_tiepoint_array(lats_a, lines, cols)
lats_b = self.expand_tiepoint_array(lats_b, lines, cols)
lats_c = self.expand_tiepoint_array(lats_c, lines, cols)
lats_d = self.expand_tiepoint_array(lats_d, lines, cols)
lons_a = self.expand_tiepoint_array(lons_a, lines, cols)
lons_b = self.expand_tiepoint_array(lons_b, lines, cols)
lons_c = self.expand_tiepoint_array(lons_c, lines, cols)
lons_d = self.expand_tiepoint_array(lons_d, lines, cols)
lats_1 = (1 - a_scan) * lats_a + a_scan * lats_b
lats_2 = (1 - a_scan) * lats_d + a_scan * lats_c
lats = (1 - a_track) * lats_1 + a_track * lats_2
lons_1 = (1 - a_scan) * lons_a + a_scan * lons_b
lons_2 = (1 - a_scan) * lons_d + a_scan * lons_c
lons = (1 - a_track) * lons_1 + a_track * lons_2
return xr.DataArray(lons, attrs=lonattrs,
dims=dims), xr.DataArray(lats,
attrs=latattrs,
dims=dims)
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. +
(y - U * t)**2.) / a**2.))
### Compute fields
for case_number in range(num_cases):
## Set parameters
def compute_dgrid_gpu(data, spacing=0.2, chunk_size=5000):
from ase.io import read
import numpy as np
import dask.array as da
import cupy as cp
import pandas as pd
from tqdm import tqdm
import ase
# data = read(path_to_cif)
fpoints = cp.asarray(data.get_scaled_positions()-0.5)
cell = cp.asarray(ase.geometry.complete_cell(data.get_cell()) )
radius_series = pd.Series({'H': 0.31, 'He': 0.28, 'Li': 1.28, 'Be': 0.96,
'B': 0.84, 'C': 0.73, 'N': 0.71, 'O': 0.66,
'F': 0.57, 'Ne': 0.58, 'Na': 1.66, 'Mg': 1.41,
'Al': 1.21, 'Si': 1.11, 'P': 1.07, 'S': 1.05,
'Cl': 1.02, 'Ar': 1.06, 'K': 2.03, 'Ca': 1.76,
'Sc': 1.70, 'Ti': 1.60, 'V': 1.53, 'Cr': 1.39,
'Mn': 1.50, 'Fe': 1.42, 'Co': 1.38, 'Ni': 1.24,
'Cu': 1.32, 'Zn': 1.22, 'Ga': 1.22, 'Ge': 1.20,
'As': 1.19, 'Se': 1.20, 'Br': 1.20, 'Kr': 1.16,
'Rb': 2.20, 'Sr': 1.95, 'Y': 1.90, 'Zr': 1.75,
'Nb': 1.64, 'Mo': 1.54, 'Tc': 1.47, 'Ru': 1.46,
'Rh': 1.42, 'Pd': 1.39, 'Ag': 1.45, 'Cd': 1.44,
'In': 1.42, 'Sn': 1.39, 'Sb': 1.39, 'Te': 1.38,
'I': 1.39, 'Xe': 1.40, 'Cs': 2.44, 'Ba': 2.15,
'La': 2.07, 'Ce': 2.04, 'Pr': 2.03, 'Nd': 2.01,
'Pm': 1.99, 'Sm': 1.98, 'Eu': 1.98, 'Gd': 1.96,
'Tb': 1.94, 'Dy': 1.92, 'Ho': 1.92, 'Er': 1.89,
'Tm': 1.90, 'Yb': 1.87, 'Lu': 1.87, 'Hf': 1.75,
'Ta': 1.70, 'W': 1.62, 'Re': 1.51, 'Os': 1.44,
'Ir': 1.41, 'Pt': 1.36, 'Au': 1.36, 'Hg': 1.32,
'Tl': 1.45, 'Pb': 1.46, 'Bi': 1.48, 'Po': 1.40,
'At': 1.50, 'Rn': 1.50, 'Fr': 2.60, 'Ra': 2.21,
'Ac': 2.15, 'Th': 2.06, 'Pa': 2.00, 'U': 1.96,
'Np': 1.90, 'Pu': 1.87, 'Am': 1.80, 'Cm': 1.69})
radii = cp.asarray(radius_series[data.get_chemical_symbols()].values).reshape(-1,1)
[nx, ny, nz] = (data.get_cell_lengths_and_angles()[0:3]/spacing).astype(np.int)+1
# print([nx,ny,nz])
gpoints = da.stack(da.meshgrid(np.linspace(-0.5, 0.5, nx), np.linspace(-0.5, 0.5, ny),np.linspace(-0.5, 0.5, nz), indexing='ij'), -1).reshape(-1, 3).rechunk(chunk_size,3)
# def gpd(points, fpoints=fpoints, cell=cell, radii=radii):
# points = cp.asarray(points)
# return cp.min(cp.linalg.norm(cp.dot(cell, (cp.expand_dims(points, axis=1)-cp.expand_dims(fpoints, axis=0)-cp.around(cp.expand_dims(points, axis=1)-cp.expand_dims(fpoints, axis=0))).reshape(-1,3).T).T, axis=1).reshape(fpoints.shape[0],-1)-radii, axis=0)
def gpd(points, fpoints=fpoints, cell=cell, radii=radii):
points = cp.asarray(points)
diff = cp.expand_dims(points, axis=1)-cp.expand_dims(fpoints, axis=0)
diff = (diff - cp.around(diff)).reshape(-1,3)
diff = cp.dot(cell, diff.T).T
diff = cp.linalg.norm(diff, axis=1).reshape(-1,fpoints.shape[0])-radii.T
return cp.min(diff, axis=1)
gpoints_da = gpoints.map_blocks(gpd,chunks=(chunk_size,1))
distance_grid = []
for block in tqdm(gpoints_da.blocks):
# print("Working on block {} \n".format(i))
distance_grid.append(block.compute())
# Free the gpu memeory before returning the output
del data, cell, fpoints, radius_series, radii
cp._default_memory_pool.free_all_blocks()
cp._default_pinned_memory_pool.free_all_blocks()
return np.hstack(distance_grid).reshape(nx,ny,nz)
def test_map_nav_size_error(self):
iter_array, _ = da.meshgrid(range(12), range(10))
s_iter = hs.signals.BaseSignal(iter_array).T
f = lambda a, b: a + b
with pytest.raises(ValueError):
self.s.map(function=f, b=s_iter, inplace=False)
def interpolate(self, lon1, lat1, satz1):
cscan_len = self.cscan_len
cscan_full_width = self.cscan_full_width
fscan_width = self.fscan_width
fscan_len = self.fscan_len
scans = satz1.shape[0] // cscan_len
satz1 = satz1.data
satz1 = satz1.reshape((-1, cscan_len, cscan_full_width))
satz_a, satz_b, satz_c, satz_d = get_corners(da.deg2rad(satz1))
c_exp, c_ali = compute_expansion_alignment(satz_a, satz_b, satz_c, satz_d)
x, y = self.get_coords(scans)
i_rs, i_rt = da.meshgrid(x, y)
p_os = 0
p_ot = 0
s_s = (p_os + i_rs) * 1. / fscan_width
s_t = (p_ot + i_rt) * 1. / fscan_len
cols = fscan_width
lines = fscan_len
c_exp_full = self.expand_tiepoint_array(c_exp, lines, cols)
c_ali_full = self.expand_tiepoint_array(c_ali, lines, cols)
a_track = s_t
a_scan = (s_s + s_s * (1 - s_s) * c_exp_full + s_t*(1 - s_t) * c_ali_full)
res = []
sublat = lat1[::16, ::16]
sublon = lon1[::16, ::16]
to_cart = abs(sublat).max() > 60 or (sublon.max() - sublon.min()) > 180
if to_cart:
datasets = lonlat2xyz(lon1, lat1)
else:
datasets = [lon1, lat1]
for data in datasets:
data_attrs = data.attrs
dims = data.dims
data = data.data
data = data.reshape((-1, cscan_len, cscan_full_width))
data_a, data_b, data_c, data_d = get_corners(data)
data_a = self.expand_tiepoint_array(data_a, lines, cols)
data_b = self.expand_tiepoint_array(data_b, lines, cols)
data_c = self.expand_tiepoint_array(data_c, lines, cols)
data_d = self.expand_tiepoint_array(data_d, lines, cols)
data_1 = (1 - a_scan) * data_a + a_scan * data_b
data_2 = (1 - a_scan) * data_d + a_scan * data_c
data = (1 - a_track) * data_1 + a_track * data_2
res.append(xr.DataArray(data, attrs=data_attrs, dims=dims))
if to_cart:
return xyz2lonlat(*res)
else:
return res
