def concat(key):
return pipe(
lambda x: func(np.array(x)),
delayed,
lambda x: fmap(lambda y: (y.shape, x(y)), data.blocks),
fmap(lambda x: (x[0], get(key, x[1]))),
fmap(lambda x: from_delayed(key, x[0], x[1])),
list,
lambda x: da.concatenate(x, axis=0),
)
def fit_fourier(fx_data, fy_data, redundancy_func):
"""Fit the data in Fourier space.
Fit the data after it has been discretized and transformed into
Fourier space.
Args:
fx_data: microstructure in k-space
fy_data: response in k-space
redundancy_func: helps remove redundancies in the coefficient matrix
Returns:
the coefficient matrix (unchunked)
>>> make_data = lambda s: da.from_array(np.arange(np.prod(s),
... dtype=float).reshape(s),
... chunks=s)
>>> matrix = fit_fourier(make_data((5, 4, 4, 3)),
... make_data((5, 4, 4)),
... lambda _: (slice(None),))
>>> print(matrix.shape)
(4, 4, 3)
>>> test_matrix = np.resize([5. / 18., 1. / 9., -1. / 18.], (4, 4, 3))
>>> assert np.allclose(matrix, test_matrix)
"""
lstsq_mks_ = lstsq_mks(fx_data, fy_data, redundancy_func)
return pipe(
fmap(lstsq_mks_, np.ndindex(fx_data.shape[1:-1])),
list,
array_from_tuple(shape=fx_data.shape[1:], dtype=np.complex)
)
def zero_pad(arr, shape, chunks):
"""Zero pad an array with zeros
Args:
arr: the array to pad
shape: the shape of the new array
chunks: how to rechunk the new array
Returns:
the new padded version of the array
>>> print(
... zero_pad(
... np.arange(4).reshape([1, 2, 2, 1]),
... (1, 4, 5, 1),
... None
... )[0,...,0].compute()
... )
[[0 0 0 0 0]
[0 0 0 1 0]
[0 0 2 3 0]
[0 0 0 0 0]]
>>> print(zero_pad(np.arange(4).reshape([2, 2]), (4, 5), None).compute())
[[0 0 0 0 0]
[0 0 0 1 0]
[0 0 2 3 0]
[0 0 0 0 0]]
>>> zero_pad(zero_pad(np.arange(4).reshape([2, 2]), (4, 5, 1), None))
Traceback (most recent call last):
...
RuntimeError: length of shape is incorrect
>>> zero_pad(zero_pad(np.arange(4).reshape([2, 2]), (1, 2), None))
Traceback (most recent call last):
...
RuntimeError: resize shape is too small
>>> arr = da.from_array(np.arange(4).reshape((2, 2)), chunks=(2, 1))
>>> out = zero_pad(arr, (4, 3), (-1, 1))
>>> out.shape
(4, 3)
>>> out.chunks
((4,), (1, 1, 1))
"""
if len(shape) != len(arr.shape):
raise RuntimeError("length of shape is incorrect")
if not np.all(shape >= arr.shape):
raise RuntimeError("resize shape is too small")
return pipe(
np.array(shape) - np.array(arr.shape),
lambda x: np.concatenate(
((x - (x // 2))[..., None], (x // 2)[..., None]), axis=1
),
fmap(tuple),
tuple,
lambda x: da.pad(arr, x, "constant", constant_values=0),
lambda x: da.rechunk(x, chunks=chunks or x.shape),
)
def apply_dict_func(func, data, shape_dict):
"""Apply a function that returns a dictionary of arrays to a Dask
array in parallel.
e.g.
>>> data = da.from_array(np.arange(36).reshape(4, 3, 3), chunks=(2, 3, 3))
>>> def func(arr):
... return dict(
... a=np.resize(
... arr,
... (arr.shape[0],) + (arr.shape[1] + 1,) + (arr.shape[2] + 1,)
... ),
... b=np.resize(arr, (arr.shape[:3]) + (1,))
... )
>>> out = apply_dict_func(func, data, dict(a=(4, 4, 4), b=(4, 3, 3, 1)))
>>> print(out['a'].chunks)
((2, 2), (4,), (4,))
>>> print(out['b'].shape)
(4, 3, 3, 1)
>>> print(out['a'].compute().shape)
(4, 4, 4)
Args:
func: the function to apply to the Dask array
data: the Dask array to call the function with
shape_dict: a dictionary of shapes, the keys are the keys
returned by the funcion and the shapes correspond to the
shapes that are output from the function
"""
@curry
def from_delayed(key, shape, delayed_func):
return da.from_delayed(
delayed_func, dtype=float, shape=(shape[0],) + shape_dict[key][1:]
)
def concat(key):
return pipe(
lambda x: func(np.array(x)),
delayed,
lambda x: fmap(lambda y: (y.shape, x(y)), data.blocks),
fmap(lambda x: (x[0], get(key, x[1]))),
fmap(lambda x: from_delayed(key, x[0], x[1])),
list,
lambda x: da.concatenate(x, axis=0),
)
return pipe(shape_dict.keys(), fmap(lambda x: (x, concat(x))), dict)
def apply_dataframe_func(func, data):
"""Daskerize a function that takes an array and returns a dataframe
>>> import pandas
>>> def func(x):
... return pandas.DataFrame(x)
>>> x = np.random.random((10, 4))
>>> df = apply_dataframe_func(func, da.from_array(x, chunks=(2, 4)))
>>> df.get_partition(0).compute().shape
(2, 4)
"""
return pipe(data.blocks, fmap(delayed(func)), list, ddf.from_delayed)
def fit_disc(x_data, y_data, redundancy_func):
"""Fit the discretized data.
Fit the data after the data has already been discretized.
Args:
x_data: the discretized mircrostructure field
y_data: the discretized response field
redundancy_func: helps remove redundancies in the coefficient matrix
Returns:
the chunked coefficient matrix based on the chunking of local
state space from the discretized mircrostructure field
>>> make_data = lambda s, c: da.from_array(
... np.arange(np.prod(s),
... dtype=float).reshape(s),
... chunks=c
... )
>>> matrix = fit_disc(make_data((6, 4, 4, 3), (2, 4, 4, 1)),
... make_data((6, 4, 4), (2, 4, 4)),
... lambda _: (slice(None),))
>>> print(matrix.shape)
(4, 4, 3)
>>> print(matrix.chunks)
((4,), (4,), (1, 1, 1))
>>> assert np.allclose(matrix.compute()[0, 0, 0], 5. / 18.)
"""
chunks = lambda x: (None, ) * (len(x.shape) - 1) + (x_data.chunks[-1], )
return pipe(
[x_data, y_data],
fmap(dafftn(axes=faxes(x_data))),
list,
lambda x: fit_fourier(*x, redundancy_func),
lambda x: da.from_array(x, chunks=chunks(x)),
)
def extend(shape, arr):
"""Extend an array by adding new axes with shape of shape argument.
The values from the existing axes are repeated in the new
axes. The is achieved using repeated uses of np.repeat followed by
np.reshape.
Args:
shape: the new shape to extend by
arr: the array to extend
Returns:
a new extended array
>>> a = np.arange(6).reshape((2, 3))
>>> extend((4,), a).shape
(2, 3, 4)
>>> print(extend((2, 3), a))
[[[[0 0 0]
[0 0 0]]
<BLANKLINE>
[[1 1 1]
[1 1 1]]
<BLANKLINE>
[[2 2 2]
[2 2 2]]]
<BLANKLINE>
<BLANKLINE>
[[[3 3 3]
[3 3 3]]
<BLANKLINE>
[[4 4 4]
[4 4 4]]
<BLANKLINE>
[[5 5 5]
[5 5 5]]]]
"""
extend_ = curry(lambda s, x: np.repeat(x, s, axis=-1).reshape(x.shape +
(s, )))
fextend = sequence(fmap(extend_), list, lambda x: sequence(*x))
return fextend(shape)(arr)
def quantile_ndim(arr, quantiles):
"""Use np.quantile across varying quantiles
Args:
arr: array to quantize (at least 2D)
quantiles: same dimensions as arr
Returns:
quantized array
>>> a = np.array([np.arange(10), np.arange(10)])
>>> q = np.array(((0.2, 0.8), (0.41, 0.59)))
>>> print(quantile_ndim(a, q))
[[2 7]
[4 5]]
"""
return pipe(
zip(arr, quantiles),
fmap(lambda x: np.quantile(*x, axis=-1, interpolation="nearest").T),
np.stack,
)
def dist_mesh(shape):
"""Calculate a mesh of distances
Assume dx=1 and the center pixel is shifted to the left for odd
shaped axes.
Args:
shape: the shape of the domain
Returns:
an array of distances from the center
>>> assert np.allclose(
... dist_mesh((3, 2, 4)),
... np.sqrt(np.array(
... [[[6, 3, 2, 3],
... [5, 2, 1, 2]],
... [[5, 2, 1, 2],
... [4, 1, 0, 1]],
... [[6, 3, 2, 3],
... [5, 2, 1, 2]]]
... ))
... )
>>> dist_mesh((3, 5, 4)).shape
(3, 5, 4)
"""
center = lambda x: np.reshape(
np.array(x) // 2, (len(x), ) + (1, ) * len(x))
return sequence(
fmap(lambda x: np.linspace(0, x - 1, x)),
lambda x: np.array(np.meshgrid(*x, indexing="ij")),
lambda x: x - center(x.shape[1:]),
lambda x: np.linalg.norm(x, axis=0),
)(shape)
