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)