Commit b41c932e authored by Kruyff,D.L.W. (Dylan)'s avatar Kruyff,D.L.W. (Dylan)
Browse files

Redo lost changes

parent 8d4ead19
......@@ -23,8 +23,8 @@ def register_cupyx():
from cupyx.scipy.sparse import spmatrix
try:
from cupy.sparse import hstack
from cupy.sparse import vstack
from cupyx.scipy.sparse import hstack
from cupyx.scipy.sparse import vstack
except ImportError as e:
raise ImportError(
"Stacking of sparse arrays requires at least CuPy version 8.0.0"
......
......@@ -23,7 +23,7 @@ def blockwise(
meta=None,
**kwargs
):
""" Tensor operation: Generalized inner and outer products
"""Tensor operation: Generalized inner and outer products
A broad class of blocked algorithms and patterns can be specified with a
concise multi-index notation. The ``blockwise`` function applies an in-memory
......
......@@ -87,7 +87,7 @@ with ignoring(AttributeError):
def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
""" Coarsen array by applying reduction to fixed size neighborhoods
"""Coarsen array by applying reduction to fixed size neighborhoods
Parameters
----------
......@@ -144,7 +144,7 @@ def coarsen(reduction, x, axes, trim_excess=False, **kwargs):
def trim(x, axes=None):
""" Trim boundaries off of array
"""Trim boundaries off of array
>>> x = np.arange(24).reshape((4, 6))
>>> trim(x, axes={0: 0, 1: 1})
......@@ -166,7 +166,7 @@ def trim(x, axes=None):
def topk(a, k, axis, keepdims):
""" Chunk and combine function of topk
"""Chunk and combine function of topk
Extract the k largest elements from a on the given axis.
If k is negative, extract the -k smallest elements instead.
......@@ -184,7 +184,7 @@ def topk(a, k, axis, keepdims):
def topk_aggregate(a, k, axis, keepdims):
""" Final aggregation function of topk
"""Final aggregation function of topk
Invoke topk one final time and then sort the results internally.
"""
......@@ -202,7 +202,7 @@ def topk_aggregate(a, k, axis, keepdims):
def argtopk_preprocess(a, idx):
""" Preparatory step for argtopk
"""Preparatory step for argtopk
Put data together with its original indices in a tuple.
"""
......@@ -210,7 +210,7 @@ def argtopk_preprocess(a, idx):
def argtopk(a_plus_idx, k, axis, keepdims):
""" Chunk and combine function of argtopk
"""Chunk and combine function of argtopk
Extract the indices of the k largest elements from a on the given axis.
If k is negative, extract the indices of the -k smallest elements instead.
......@@ -239,7 +239,7 @@ def argtopk(a_plus_idx, k, axis, keepdims):
def argtopk_aggregate(a_plus_idx, k, axis, keepdims):
""" Final aggregation function of argtopk
"""Final aggregation function of argtopk
Invoke argtopk one final time, sort the results internally, drop the data
and return the index only.
......@@ -278,7 +278,7 @@ def view(x, dtype, order="C"):
def slice_with_int_dask_array(x, idx, offset, x_size, axis):
""" Chunk function of `slice_with_int_dask_array_on_axis`.
"""Chunk function of `slice_with_int_dask_array_on_axis`.
Slice one chunk of x by one chunk of idx.
Parameters
......@@ -320,7 +320,7 @@ def slice_with_int_dask_array(x, idx, offset, x_size, axis):
def slice_with_int_dask_array_aggregate(idx, chunk_outputs, x_chunks, axis):
""" Final aggregation function of `slice_with_int_dask_array_on_axis`.
"""Final aggregation function of `slice_with_int_dask_array_on_axis`.
Aggregate all chunks of x by one chunk of idx, reordering the output of
`slice_with_int_dask_array`.
......
......@@ -105,7 +105,7 @@ def getter(a, b, asarray=True, lock=None):
def getter_nofancy(a, b, asarray=True, lock=None):
""" A simple wrapper around ``getter``.
"""A simple wrapper around ``getter``.
Used to indicate to the optimization passes that the backend doesn't
support fancy indexing.
......@@ -114,7 +114,7 @@ def getter_nofancy(a, b, asarray=True, lock=None):
def getter_inline(a, b, asarray=True, lock=None):
""" A getter function that optimizations feel comfortable inlining
"""A getter function that optimizations feel comfortable inlining
Slicing operations with this function may be inlined into a graph, such as
in the following rewrite
......@@ -166,7 +166,7 @@ def implements(*numpy_functions):
def slices_from_chunks(chunks):
""" Translate chunks tuple to a set of slices in product order
"""Translate chunks tuple to a set of slices in product order
>>> slices_from_chunks(((2, 2), (3, 3, 3))) # doctest: +NORMALIZE_WHITESPACE
[(slice(0, 2, None), slice(0, 3, None)),
......@@ -194,7 +194,7 @@ def getem(
asarray=True,
dtype=None,
):
""" Dask getting various chunks from an array-like
"""Dask getting various chunks from an array-like
>>> getem('X', chunks=(2, 3), shape=(4, 6)) # doctest: +SKIP
{('X', 0, 0): (getter, 'X', (slice(0, 2), slice(0, 3))),
......@@ -227,7 +227,7 @@ def getem(
def dotmany(A, B, leftfunc=None, rightfunc=None, **kwargs):
""" Dot product of many aligned chunks
"""Dot product of many aligned chunks
>>> x = np.array([[1, 2], [1, 2]])
>>> y = np.array([[10, 20], [10, 20]])
......@@ -249,7 +249,7 @@ def dotmany(A, B, leftfunc=None, rightfunc=None, **kwargs):
def _concatenate2(arrays, axes=[]):
""" Recursively Concatenate nested lists of arrays along axes
"""Recursively Concatenate nested lists of arrays along axes
Each entry in axes corresponds to each level of the nested list. The
length of axes should correspond to the level of nesting of arrays.
......@@ -369,7 +369,7 @@ def apply_infer_dtype(func, args, kwargs, funcname, suggest_dtype="dtype", nout=
def normalize_arg(x):
""" Normalize user provided arguments to blockwise or map_blocks
"""Normalize user provided arguments to blockwise or map_blocks
We do a few things:
......@@ -391,7 +391,7 @@ def normalize_arg(x):
def _pass_extra_kwargs(func, keys, *args, **kwargs):
""" Helper for :func:`map_blocks` to pass `block_info` or `block_id`.
"""Helper for :func:`map_blocks` to pass `block_info` or `block_id`.
For each element of `keys`, a corresponding element of args is changed
to a keyword argument with that key, before all arguments re passed on
......@@ -413,7 +413,7 @@ def map_blocks(
meta=None,
**kwargs,
):
""" Map a function across all blocks of a dask array.
"""Map a function across all blocks of a dask array.
Parameters
----------
......@@ -760,7 +760,7 @@ def map_blocks(
def broadcast_chunks(*chunkss):
""" Construct a chunks tuple that broadcasts many chunks tuples
"""Construct a chunks tuple that broadcasts many chunks tuples
>>> a = ((5, 5),)
>>> b = ((5, 5),)
......@@ -812,7 +812,7 @@ def store(
return_stored=False,
**kwargs,
):
""" Store dask arrays in array-like objects, overwrite data in target
"""Store dask arrays in array-like objects, overwrite data in target
This stores dask arrays into object that supports numpy-style setitem
indexing. It stores values chunk by chunk so that it does not have to
......@@ -999,7 +999,7 @@ See the following documentation page for details:
class Array(DaskMethodsMixin):
""" Parallel Dask Array
"""Parallel Dask Array
A parallel nd-array comprised of many numpy arrays arranged in a grid.
......@@ -1398,7 +1398,7 @@ class Array(DaskMethodsMixin):
return r
def to_svg(self, size=500):
""" Convert chunks from Dask Array into an SVG Image
"""Convert chunks from Dask Array into an SVG Image
Parameters
----------
......@@ -1419,7 +1419,7 @@ class Array(DaskMethodsMixin):
return svg(self.chunks, size=size)
def to_hdf5(self, filename, datapath, **kwargs):
""" Store array in HDF5 file
"""Store array in HDF5 file
>>> x.to_hdf5('myfile.hdf5', '/x') # doctest: +SKIP
......@@ -1435,7 +1435,7 @@ class Array(DaskMethodsMixin):
return to_hdf5(filename, datapath, self, **kwargs)
def to_dask_dataframe(self, columns=None, index=None, meta=None):
""" Convert dask Array to dask Dataframe
"""Convert dask Array to dask Dataframe
Parameters
----------
......@@ -1636,7 +1636,7 @@ class Array(DaskMethodsMixin):
@property
def blocks(self):
""" Slice an array by blocks
"""Slice an array by blocks
This allows blockwise slicing of a Dask array. You can perform normal
Numpy-style slicing but now rather than slice elements of the array you
......@@ -2118,7 +2118,7 @@ class Array(DaskMethodsMixin):
return map_blocks(func, self, *args, **kwargs)
def map_overlap(self, func, depth, boundary=None, trim=True, **kwargs):
""" Map a function over blocks of the array with some overlap
"""Map a function over blocks of the array with some overlap
We share neighboring zones between blocks of the array, then map a
function, then trim away the neighboring strips.
......@@ -2224,7 +2224,7 @@ class Array(DaskMethodsMixin):
return clip(self, min, max)
def view(self, dtype=None, order="C"):
""" Get a view of the array as a new data type
"""Get a view of the array as a new data type
Parameters
----------
......@@ -2355,7 +2355,7 @@ def ensure_int(f):
def normalize_chunks(chunks, shape=None, limit=None, dtype=None, previous_chunks=None):
""" Normalize chunks to tuple of tuples
"""Normalize chunks to tuple of tuples
This takes in a variety of input types and information and produces a full
tuple-of-tuples result for chunks, suitable to be passed to Array or
......@@ -2529,7 +2529,7 @@ def _compute_multiplier(limit: int, dtype, largest_block: int, result):
def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
""" Determine automatic chunks
"""Determine automatic chunks
This takes in a chunks value that contains ``"auto"`` values in certain
dimensions and replaces those values with concrete dimension sizes that try
......@@ -2654,7 +2654,7 @@ def auto_chunks(chunks, shape, limit, dtype, previous_chunks=None):
def round_to(c, s):
""" Return a chunk dimension that is close to an even multiple or factor
"""Return a chunk dimension that is close to an even multiple or factor
We want values for c that are nicely aligned with s.
......@@ -2690,7 +2690,7 @@ def from_array(
getitem=None,
meta=None,
):
""" Create dask array from something that looks like an array
"""Create dask array from something that looks like an array
Input must have a ``.shape``, ``.ndim``, ``.dtype`` and support numpy-style slicing.
......@@ -3029,7 +3029,7 @@ def _check_regular_chunks(chunkset):
def from_delayed(value, shape, dtype=None, meta=None, name=None):
""" Create a dask array from a dask delayed value
"""Create a dask array from a dask delayed value
This routine is useful for constructing dask arrays in an ad-hoc fashion
using dask delayed, particularly when combined with stack and concatenate.
......@@ -3062,7 +3062,7 @@ def from_delayed(value, shape, dtype=None, meta=None, name=None):
def from_func(func, shape, dtype=None, name=None, args=(), kwargs={}):
""" Create dask array in a single block by calling a function
"""Create dask array in a single block by calling a function
Calling the provided function with func(*args, **kwargs) should return a
NumPy array of the indicated shape and dtype.
......@@ -3090,7 +3090,7 @@ def from_func(func, shape, dtype=None, name=None, args=(), kwargs={}):
def common_blockdim(blockdims):
""" Find the common block dimensions from the list of block dimensions
"""Find the common block dimensions from the list of block dimensions
Currently only implements the simplest possible heuristic: the common
block-dimension is the only one that does not span fully span a dimension.
......@@ -3877,7 +3877,7 @@ def broadcast_shapes(*shapes):
def elemwise(op, *args, **kwargs):
""" Apply elementwise function across arguments
"""Apply elementwise function across arguments
Respects broadcasting rules
......@@ -3960,7 +3960,7 @@ def elemwise(op, *args, **kwargs):
def handle_out(out, result):
""" Handle out parameters
"""Handle out parameters
If out is a dask.array then this overwrites the contents of that array with
the result
......@@ -4127,7 +4127,7 @@ def broadcast_arrays(*args, **kwargs):
def offset_func(func, offset, *args):
""" Offsets inputs by offset
"""Offsets inputs by offset
>>> double = lambda x: x * 2
>>> f = offset_func(double, (10,))
......@@ -4148,7 +4148,7 @@ def offset_func(func, offset, *args):
def chunks_from_arrays(arrays):
""" Chunks tuple from nested list of arrays
"""Chunks tuple from nested list of arrays
>>> x = np.array([1, 2])
>>> chunks_from_arrays([x, x])
......@@ -4184,7 +4184,7 @@ def chunks_from_arrays(arrays):
def deepfirst(seq):
""" First element in a nested list
"""First element in a nested list
>>> deepfirst([[[1, 2], [3, 4]], [5, 6], [7, 8]])
1
......@@ -4204,7 +4204,7 @@ def shapelist(a):
def reshapelist(shape, seq):
""" Reshape iterator to nested shape
"""Reshape iterator to nested shape
>>> reshapelist((2, 3), range(6))
[[0, 1, 2], [3, 4, 5]]
......@@ -4217,7 +4217,7 @@ def reshapelist(shape, seq):
def transposelist(arrays, axes, extradims=0):
""" Permute axes of nested list
"""Permute axes of nested list
>>> transposelist([[1,1,1],[1,1,1]], [2,1])
[[[1, 1], [1, 1], [1, 1]]]
......@@ -4357,7 +4357,7 @@ def stack(seq, axis=0, allow_unknown_chunksizes=False):
def concatenate3(arrays):
""" Recursive np.concatenate
"""Recursive np.concatenate
Input should be a nested list of numpy arrays arranged in the order they
should appear in the array itself. Each array should have the same number
......@@ -4436,7 +4436,7 @@ def concatenate_axes(arrays, axes):
def to_hdf5(filename, *args, **kwargs):
""" Store arrays in HDF5 file
"""Store arrays in HDF5 file
This saves several dask arrays into several datapaths in an HDF5 file.
It creates the necessary datasets and handles clean file opening/closing.
......@@ -4670,7 +4670,7 @@ def _vindex_array(x, dict_indexes):
def _get_axis(indexes):
""" Get axis along which point-wise slicing results lie
"""Get axis along which point-wise slicing results lie
This is mostly a hack because I can't figure out NumPy's rule on this and
can't be bothered to go reading.
......@@ -4735,7 +4735,7 @@ def _vindex_merge(locations, values):
def to_npy_stack(dirname, x, axis=0):
""" Write dask array to a stack of .npy files
"""Write dask array to a stack of .npy files
This partitions the dask.array along one axis and stores each block along
that axis as a single .npy file in the specified directory
......@@ -4787,7 +4787,7 @@ def to_npy_stack(dirname, x, axis=0):
def from_npy_stack(dirname, mmap_mode="r"):
""" Load dask array from stack of npy files
"""Load dask array from stack of npy files
See ``da.to_npy_stack`` for docstring
......
......@@ -112,7 +112,7 @@ _out_chunk_fns = {
def fft_wrap(fft_func, kind=None, dtype=None):
""" Wrap 1D, 2D, and ND real and complex FFT functions
"""Wrap 1D, 2D, and ND real and complex FFT functions
Takes a function that behaves like ``numpy.fft`` functions and
a specified kind to match it to that are named after the functions
......
......@@ -12,7 +12,7 @@ from ..core import flatten
# Modified version of `numpy.lib.function_base._parse_gufunc_signature`
# Modifications:
# - Allow for zero input arguments
# See https://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html
# See https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html
_DIMENSION_NAME = r"\w+"
_CORE_DIMENSION_LIST = "(?:{0:}(?:,{0:})*,?)?".format(_DIMENSION_NAME)
_ARGUMENT = r"\({}\)".format(_CORE_DIMENSION_LIST)
......@@ -240,6 +240,9 @@ def apply_gufunc(func, signature, *args, **kwargs):
dimensions are to consist only of one chunk.
Warning: enabling this can increase memory usage significantly.
Defaults to ``False``.
meta: Optional, tuple, keyword only
tuple of empty ndarrays describing the shape and dtype of the output of the gufunc.
Defaults to ``None``.
**kwargs : dict
Extra keyword arguments to pass to `func`
......@@ -270,7 +273,7 @@ def apply_gufunc(func, signature, *args, **kwargs):
References
----------
.. [1] https://docs.scipy.org/doc/numpy/reference/ufuncs.html
.. [2] https://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html
.. [2] https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html
"""
axes = kwargs.pop("axes", None)
axis = kwargs.pop("axis", None)
......@@ -279,6 +282,7 @@ def apply_gufunc(func, signature, *args, **kwargs):
output_sizes = kwargs.pop("output_sizes", None)
vectorize = kwargs.pop("vectorize", None)
allow_rechunk = kwargs.pop("allow_rechunk", False)
meta = kwargs.pop("meta", None)
# Input processing:
## Signature
......@@ -417,27 +421,32 @@ significantly.".format(
### Use existing `blockwise` but only with loopdims to enforce
### concatenation for coredims that appear also at the output
### Modifying `blockwise` could improve things here.
try:
tmp = blockwise( # First try to compute meta
func, loop_output_dims, *arginds, concatenate=True, **kwargs
)
except ValueError:
# If computing meta doesn't work, provide it explicitly based on
# provided dtypes
sample = arginds[0]._meta
if isinstance(output_dtypes, tuple):
meta = tuple(
meta_from_array(sample, dtype=odt)
for ocd, odt in zip(output_coredimss, output_dtypes)
)
else:
meta = tuple(
meta_from_array(sample, dtype=odt)
for ocd, odt in zip((output_coredimss,), (output_dtypes,))
)
if meta is not None:
tmp = blockwise(
func, loop_output_dims, *arginds, concatenate=True, meta=meta, **kwargs
)
else:
try:
tmp = blockwise( # First try to compute meta
func, loop_output_dims, *arginds, concatenate=True, **kwargs
)
except ValueError:
# If computing meta doesn't work, provide it explicitly based on
# provided dtypes
sample = arginds[0]._meta
if isinstance(output_dtypes, tuple):
meta = tuple(
meta_from_array(sample, dtype=odt)
for ocd, odt in zip(output_coredimss, output_dtypes)
)
else:
meta = tuple(
meta_from_array(sample, dtype=odt)
for ocd, odt in zip((output_coredimss,), (output_dtypes,))
)
tmp = blockwise(
func, loop_output_dims, *arginds, concatenate=True, meta=meta, **kwargs
)
if isinstance(tmp._meta, tuple):
metas = tmp._meta
......@@ -580,7 +589,7 @@ class gufunc(object):
References
----------
.. [1] https://docs.scipy.org/doc/numpy/reference/ufuncs.html
.. [2] https://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html
.. [2] https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html
"""
def __init__(self, pyfunc, **kwargs):
......@@ -628,7 +637,7 @@ class gufunc(object):
output_sizes=self.output_sizes,
output_dtypes=self.output_dtypes,
allow_rechunk=self.allow_rechunk or kwargs.pop("allow_rechunk", False),
**kwargs
**kwargs,
)
......@@ -681,6 +690,9 @@ def as_gufunc(signature=None, **kwargs):
dimensions are to consist only of one chunk.
Warning: enabling this can increase memory usage significantly.
Defaults to ``False``.
meta: Optional, tuple, keyword only
tuple of empty ndarrays describing the shape and dtype of the output of the gufunc.
Defaults to ``None``.
Returns
-------
......@@ -710,7 +722,7 @@ def as_gufunc(signature=None, **kwargs):
References
----------
.. [1] https://docs.scipy.org/doc/numpy/reference/ufuncs.html
.. [2] https://docs.scipy.org/doc/numpy/reference/c-api.generalized-ufuncs.html
.. [2] https://docs.scipy.org/doc/numpy/reference/c-api/generalized-ufuncs.html
"""
_allowedkeys = {
"vectorize",
......@@ -720,6 +732,7 @@ def as_gufunc(signature=None, **kwargs):
"output_sizes",
"output_dtypes",
"allow_rechunk",
"meta",
}
if set(_allowedkeys).issubset(kwargs.keys()):
raise TypeError("Unsupported keyword argument(s) provided")
......
Supports Markdown
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment