diff --git a/.github/workflows/docs.yml b/.github/workflows/docs.yml index 73ef4193..0ffef95e 100644 --- a/.github/workflows/docs.yml +++ b/.github/workflows/docs.yml @@ -57,7 +57,7 @@ jobs: # https://github.com/JamesIves/github-pages-deploy-action/tree/dev#using-an-ssh-deploy-key- - name: Deploy uses: JamesIves/github-pages-deploy-action@v4 - if: ${{ github.ref == 'refs/heads/master' }} + if: ${{ github.ref == 'refs/heads/main' }} with: folder: docs/_build/html ssh-key: ${{ secrets.DEPLOY_KEY }} diff --git a/docs/api.rst b/docs/api.rst index 9ff9e894..97162042 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -62,6 +62,14 @@ the index objects. .. autofunction:: ndindex.broadcast_shapes +Exceptions +========== + +These are some custom exceptions that are raised by a few functions in +ndindex. Note that most functions in ndindex will raise `IndexError` +(if the index would be invalid), or `TypeError` or `ValueError` (if the input +types or values are incorrect). + .. autoexception:: ndindex.BroadcastError .. autoexception:: ndindex.AxisError @@ -106,3 +114,5 @@ relied on as they may be removed or changed. .. autofunction:: ndindex.shapetools.remove_indices .. autofunction:: ndindex.shapetools.unremove_indices + +.. autofunction:: ndindex.shapetools.normalize_skip_axes diff --git a/docs/changelog.md b/docs/changelog.md index 93d4a555..11f3c2af 100644 --- a/docs/changelog.md +++ b/docs/changelog.md @@ -110,9 +110,9 @@ ### Minor Changes - Added - [CODE_OF_CONDUCT.md](https://github.com/Quansight-Labs/ndindex/blob/master/CODE_OF_CONDUCT.md) + [CODE_OF_CONDUCT.md](https://github.com/Quansight-Labs/ndindex/blob/main/CODE_OF_CONDUCT.md) to the ndindex repository. ndindex follows the [Quansight Code of - Conduct](https://github.com/Quansight/.github/blob/master/CODE_OF_CONDUCT.md). + Conduct](https://github.com/Quansight/.github/blob/main/CODE_OF_CONDUCT.md). - Avoid precomputing all iterated values for slices with large steps in {any}`ChunkSize.as_subchunks()`. diff --git a/docs/index.md b/docs/index.md index 26ec156a..16806dab 100644 --- a/docs/index.md +++ b/docs/index.md @@ -277,7 +277,7 @@ There are two primary types of tests that we employ to verify this: - Hypothesis tests. Hypothesis is a library that can intelligently check a combinatorial search space of inputs. This requires writing hypothesis strategies that can generate all the relevant types of indices (see - [ndindex/tests/helpers.py](https://github.com/Quansight-Labs/ndindex/blob/master/ndindex/tests/helpers.py)). + [ndindex/tests/helpers.py](https://github.com/Quansight-Labs/ndindex/blob/main/ndindex/tests/helpers.py)). For more information on hypothesis, see . All tests have hypothesis tests, even if they are also tested exhaustively. @@ -307,7 +307,7 @@ Benchmarks for ndindex are published ## License -[MIT License](https://github.com/Quansight-Labs/ndindex/blob/master/LICENSE) +[MIT License](https://github.com/Quansight-Labs/ndindex/blob/main/LICENSE) (acknowledgments)= ## Acknowledgments diff --git a/docs/slices.md b/docs/slices.md index 3cb8832d..0c76d625 100644 --- a/docs/slices.md +++ b/docs/slices.md @@ -1693,7 +1693,7 @@ hard to write slice arithmetic. The arithmetic is already hard enough due to the modular nature of `step`, but the discontinuous aspect of `start` and `stop` increases this tenfold. If you are unconvinced of this, take a look at the [source -code](https://github.com/Quansight-labs/ndindex/blob/master/ndindex/slice.py) for +code](https://github.com/Quansight-labs/ndindex/blob/main/ndindex/slice.py) for `ndindex.Slice()`. You will see lots of nested `if` blocks.[^source-footnote] This is because slices have *fundamentally* different definitions if the `start` or `stop` are `None`, negative, or nonnegative. Furthermore, `None` is diff --git a/ndindex/booleanarray.py b/ndindex/booleanarray.py index 1b99c26a..a4620a9c 100644 --- a/ndindex/booleanarray.py +++ b/ndindex/booleanarray.py @@ -108,8 +108,7 @@ def _raise_indexerror(self, shape, axis=0): for i in range(axis, axis+self.ndim): if self.shape[i-axis] != 0 and shape[i] != self.shape[i-axis]: - - raise IndexError(f"boolean index did not match indexed array along dimension {i}; dimension is {shape[i]} but corresponding boolean dimension is {self.shape[i-axis]}") + raise IndexError(f'boolean index did not match indexed array along axis {i}; size of axis is {shape[i]} but size of corresponding boolean axis is {self.shape[i-axis]}') def reduce(self, shape=None, *, axis=0, negative_int=False): """ @@ -125,7 +124,7 @@ def reduce(self, shape=None, *, axis=0, negative_int=False): >>> idx.reduce((3,)) Traceback (most recent call last): ... - IndexError: boolean index did not match indexed array along dimension 0; dimension is 3 but corresponding boolean dimension is 2 + IndexError: boolean index did not match indexed array along axis 0; size of axis is 3 but size of corresponding boolean axis is 2 >>> idx.reduce((2,)) BooleanArray([True, False]) diff --git a/ndindex/integer.py b/ndindex/integer.py index 10a1a548..c33498f4 100644 --- a/ndindex/integer.py +++ b/ndindex/integer.py @@ -1,5 +1,5 @@ from .ndindex import NDIndex, operator_index -from .shapetools import asshape +from .shapetools import AxisError, asshape class Integer(NDIndex): """ @@ -62,7 +62,7 @@ def _raise_indexerror(self, shape, axis=0): size = shape[axis] raise IndexError(f"index {self.raw} is out of bounds for axis {axis} with size {size}") - def reduce(self, shape=None, *, axis=0, negative_int=False): + def reduce(self, shape=None, *, axis=0, negative_int=False, axiserror=False): """ Reduce an Integer index on an array of shape `shape`. @@ -96,7 +96,14 @@ def reduce(self, shape=None, *, axis=0, negative_int=False): if shape is None: return self + if axiserror: + if not isinstance(shape, int): # pragma: no cover + raise TypeError("axiserror=True requires shape to be an integer") + if not self.isvalid(shape): + raise AxisError(self.raw, shape) + shape = asshape(shape, axis=axis) + self._raise_indexerror(shape, axis) if self.raw < 0 and not negative_int: diff --git a/ndindex/shapetools.py b/ndindex/shapetools.py index e3c5854b..567c9d91 100644 --- a/ndindex/shapetools.py +++ b/ndindex/shapetools.py @@ -1,14 +1,16 @@ import numbers import itertools -from collections import defaultdict from collections.abc import Sequence +from ._crt import prod from .ndindex import ndindex, operator_index class BroadcastError(ValueError): """ - Exception raised by :func:`iter_indices()` when the input shapes are not - broadcast compatible. + Exception raised by :func:`iter_indices()` and + :func:`broadcast_shapes()` when the input shapes are not broadcast + compatible. + """ __slots__ = ("arg1", "shape1", "arg2", "shape2") @@ -24,8 +26,8 @@ def __str__(self): class AxisError(ValueError, IndexError): """ - Exception raised by :func:`iter_indices()` when the `skip_axes` argument - is out of bounds. + Exception raised by :func:`iter_indices()` and + :func:`broadcast_shapes()` when the `skip_axes` argument is out-of-bounds. This is used instead of the NumPy exception of the same name so that `iter_indices` does not need to depend on NumPy. @@ -34,6 +36,9 @@ class AxisError(ValueError, IndexError): __slots__ = ("axis", "ndim") def __init__(self, axis, ndim): + # NumPy allows axis=-1 for 0-d arrays + if (ndim < 0 or -ndim <= axis < ndim) and not (ndim == 0 and axis == -1): + raise ValueError(f"Invalid AxisError ({axis}, {ndim})") self.axis = axis self.ndim = ndim @@ -49,7 +54,7 @@ def broadcast_shapes(*shapes, skip_axes=()): shape with `skip_axes`. If the shapes are not broadcast compatible (excluding `skip_axes`), - `BroadcastError` is raised. + :class:`BroadcastError` is raised. >>> from ndindex import broadcast_shapes >>> broadcast_shapes((2, 3), (3,), (4, 2, 1)) @@ -62,52 +67,41 @@ def broadcast_shapes(*shapes, skip_axes=()): Axes in `skip_axes` apply to each shape *before* being broadcasted. Each shape will be broadcasted together with these axes removed. The dimensions in skip_axes do not need to be equal or broadcast compatible with one - another. The final broadcasted shape will have `None` in each `skip_axes` - location, and the broadcasted remaining `shapes` axes elsewhere. + another. The final broadcasted shape be the result of broadcasting all the + non-skip axes. >>> broadcast_shapes((10, 3, 2), (20, 2), skip_axes=(0,)) - (None, 3, 2) + (3, 2) """ - skip_axes = asshape(skip_axes, allow_negative=True) shapes = [asshape(shape, allow_int=False) for shape in shapes] - - if any(i >= 0 for i in skip_axes) and any(i < 0 for i in skip_axes): - # See the comments in remove_indices and iter_indices - raise NotImplementedError("Mixing both negative and nonnegative skip_axes is not yet supported") + skip_axes = normalize_skip_axes(shapes, skip_axes) if not shapes: - if skip_axes: - # Raise IndexError - ndindex(skip_axes[0]).reduce(0) return () - dims = [len(shape) for shape in shapes] - shape_skip_axes = [[ndindex(i).reduce(n, negative_int=True) for i in skip_axes] for n in dims] + non_skip_shapes = [remove_indices(shape, skip_axis) for shape, skip_axis in zip(shapes, skip_axes)] + dims = [len(shape) for shape in non_skip_shapes] N = max(dims) - broadcasted_skip_axes = [ndindex(i).reduce(N) for i in skip_axes] - broadcasted_shape = [None if i in broadcasted_skip_axes else 1 for i in range(N)] + broadcasted_shape = [1]*N arg = None for i in range(-1, -N-1, -1): for j in range(len(shapes)): if dims[j] < -i: continue - shape = shapes[j] - idx = associated_axis(shape, broadcasted_shape, i, skip_axes) - broadcasted_side = broadcasted_shape[idx] + shape = non_skip_shapes[j] + broadcasted_side = broadcasted_shape[i] shape_side = shape[i] - if i in shape_skip_axes[j]: - continue - elif shape_side == 1: + if shape_side == 1: continue elif broadcasted_side == 1: broadcasted_side = shape_side arg = j elif shape_side != broadcasted_side: raise BroadcastError(arg, shapes[arg], j, shapes[j]) - broadcasted_shape[idx] = broadcasted_side + broadcasted_shape[i] = broadcasted_side return tuple(broadcasted_shape) @@ -216,73 +210,54 @@ def iter_indices(*shapes, skip_axes=(), _debug=False): Tuple(1, 2) """ - skip_axes = asshape(skip_axes, allow_negative=True) + skip_axes = normalize_skip_axes(shapes, skip_axes) shapes = [asshape(shape, allow_int=False) for shape in shapes] - if any(i >= 0 for i in skip_axes) and any(i < 0 for i in skip_axes): - # Mixing positive and negative skip_axes is too difficult to deal with - # (see the comment in unremove_indices). It's a bit of an unusual - # thing to support, at least in the general case, because a positive - # and negative index can index the same element, but only for shapes - # that are a specific size. So while, in principle something like - # iter_indices((2, 10, 20, 4), (2, 30, 4), skip_axes=(1, -2)) could - # make sense, it's a bit odd to do so. Of course, there's no reason we - # couldn't support cases like that, but they complicate the - # implementation and, especially, complicate the test generation in - # the hypothesis strategies. Given that I'm not completely sure how to - # implement it correctly, and I don't actually need support for it, - # I'm leaving it as not implemented for now. - raise NotImplementedError("Mixing both negative and nonnegative skip_axes is not yet supported") - - n = len(skip_axes) - if len(set(skip_axes)) != n: - raise ValueError("skip_axes should not contain duplicate axes") - if not shapes: - if skip_axes: - raise AxisError(skip_axes[0], 0) yield () return shapes = [asshape(shape) for shape in shapes] - ndim = len(max(shapes, key=len)) - min_ndim = len(min(shapes, key=len)) - - _skip_axes = defaultdict(list) - for shape in shapes: - for a in skip_axes: - try: - a = ndindex(a).reduce(len(shape)).raw - except IndexError: - raise AxisError(a, min_ndim) - _skip_axes[shape].append(a) - _skip_axes[shape].sort() + S = len(shapes) - iters = [[] for i in range(len(shapes))] + iters = [[] for i in range(S)] broadcasted_shape = broadcast_shapes(*shapes, skip_axes=skip_axes) - non_skip_broadcasted_shape = remove_indices(broadcasted_shape, skip_axes) - for i in range(-1, -ndim-1, -1): - for it, shape in zip(iters, shapes): + idxes = [-1]*S + + while any(i is not None for i in idxes): + for s, it, shape, sk in zip(range(S), iters, shapes, skip_axes): + i = idxes[s] + if i is None: + continue if -i > len(shape): - # for every dimension prepended by broadcasting, repeat the - # indices that many times - for j in range(len(it)): - if non_skip_broadcasted_shape[i+n] not in [0, 1]: - it[j] = ncycles(it[j], non_skip_broadcasted_shape[i+n]) - break - elif len(shape) + i in _skip_axes[shape]: + if not shape: + pass + elif len(shape) == len(sk): + # The whole shape is skipped. Just repeat the most recent slice + it[0] = ncycles(it[0], prod(broadcasted_shape)) + else: + # Find the first non-skipped axis and repeat by however + # many implicit axes are left in the broadcasted shape + for j in range(-len(shape), 0): + if j not in sk: + break + it[j] = ncycles(it[j], prod(broadcasted_shape[:len(sk)-len(shape)+len(broadcasted_shape)])) + + idxes[s] = None + continue + + val = associated_axis(broadcasted_shape, i, sk) + if i in sk: it.insert(0, [slice(None)]) else: - idx = associated_axis(shape, broadcasted_shape, i, skip_axes) - val = broadcasted_shape[idx] - assert val is not None if val == 0: return elif val != 1 and shape[i] == 1: it.insert(0, ncycles(range(shape[i]), val)) else: it.insert(0, range(shape[i])) + idxes[s] -= 1 if _debug: # pragma: no cover print(f"{iters = }") @@ -351,39 +326,31 @@ def asshape(shape, axis=None, *, allow_int=True, allow_negative=False): return tuple(newshape) -def associated_axis(shape, broadcasted_shape, i, skip_axes): +def associated_axis(broadcasted_shape, i, skip_axes): """ - Return the associated index into `broadcast_shape` corresponding to - `shape[i]` given `skip_axes`. + Return the associated element of `broadcasted_shape` corresponding to + `shape[i]` given `skip_axes`. If there is not such element (i.e., it's out + of bounds), returns None. This function makes implicit assumptions about its input and is only designed for internal use. """ - n = len(shape) - N = len(broadcasted_shape) skip_axes = sorted(skip_axes, reverse=True) if i >= 0: raise NotImplementedError - if not skip_axes: - return i - # We assume skip_axes are either all negative or all nonnegative - if skip_axes[0] < 0: - return i - elif skip_axes[0] >= 0: - invmapping = [None]*N - for s in skip_axes: - invmapping[s] = s - - for j in range(-1, i-1, -1): - if j + n in skip_axes: - k = j + n #- N - continue - for k in range(-1, -N-1, -1): - if invmapping[k] is None: - invmapping[k] = j - break - return k + if i in skip_axes: + return None + # We assume skip_axes are all negative and sorted + j = i + for sk in skip_axes: + if sk >= i: + j += 1 + else: + break + if ndindex(j).isvalid(len(broadcasted_shape)): + return broadcasted_shape[j] + return None def remove_indices(x, idxes): """ @@ -410,8 +377,8 @@ def unremove_indices(x, idxes, *, val=None): This function is only intended for internal usage. """ if any(i >= 0 for i in idxes) and any(i < 0 for i in idxes): - # A mix of positive and negative indices provides a fundamental - # problem. Sometimes, the result is not unique: for example, x = [0]; + # A mix of positive and negative indices presents a fundamental + # problem: sometimes the result is not unique. For example, x = [0]; # idxes = [1, -1] could be satisfied by both [0, None] or [0, None, # None], depending on whether each index refers to a separate None or # not (note that both cases are supported by remove_indices(), because @@ -470,3 +437,54 @@ def __repr__(self): def __iter__(self): return itertools.chain.from_iterable(itertools.repeat(tuple(self.iterable), self.n)) + +def normalize_skip_axes(shapes, skip_axes): + """ + Return a canonical form of `skip_axes` corresponding to `shapes`. + + A canonical form of `skip_axes` is a list of tuples of integers, one for + each shape in `shapes`, which are a unique set of axes for each + corresponding shape. + + If `skip_axes` is an integer, this is basically `[(skip_axes,) for s + in shapes]`. If `skip_axes` is a tuple, it is like `[skip_axes for s in + shapes]`. + + The `skip_axes` must always refer to unique axes in each shape. + + The returned `skip_axes` will always be negative integers and will be + sorted. + + This function is only intended for internal usage. + + """ + # Note: we assume asshape has already been called on the shapes in shapes + if isinstance(skip_axes, Sequence): + if skip_axes and all(isinstance(i, Sequence) for i in skip_axes): + if len(skip_axes) != len(shapes): + raise ValueError(f"Expected {len(shapes)} skip_axes") + return [normalize_skip_axes([shape], skip_axis)[0] for shape, skip_axis in zip(shapes, skip_axes)] + else: + try: + [operator_index(i) for i in skip_axes] + except TypeError: + raise TypeError("skip_axes must be an integer, a tuple of integers, or a list of tuples of integers") + + skip_axes = asshape(skip_axes, allow_negative=True) + + # From here, skip_axes is a single tuple of integers + + if not shapes and skip_axes: + raise ValueError("skip_axes must be empty if there are no shapes") + + new_skip_axes = [] + for shape in shapes: + s = tuple(sorted(ndindex(i).reduce(len(shape), negative_int=True, axiserror=True).raw for i in skip_axes)) + if len(s) != len(set(s)): + err = ValueError(f"skip_axes {skip_axes} are not unique for shape {shape}") + # For testing + err.skip_axes = skip_axes + err.shape = shape + raise err + new_skip_axes.append(s) + return new_skip_axes diff --git a/ndindex/tests/helpers.py b/ndindex/tests/helpers.py index 8f2366cf..1cceaff9 100644 --- a/ndindex/tests/helpers.py +++ b/ndindex/tests/helpers.py @@ -11,9 +11,9 @@ from hypothesis import assume, note from hypothesis.strategies import (integers, none, one_of, lists, just, builds, shared, composite, sampled_from, - booleans) + nothing, tuples as hypothesis_tuples) from hypothesis.extra.numpy import (arrays, mutually_broadcastable_shapes as - mbs, BroadcastableShapes) + mbs, BroadcastableShapes, valid_tuple_axes) from ..ndindex import ndindex from ..shapetools import remove_indices, unremove_indices @@ -154,89 +154,221 @@ def _mutually_broadcastable_shapes(draw, *, shapes=short_shapes, min_shapes=0, m mutually_broadcastable_shapes = shared(_mutually_broadcastable_shapes()) -@composite -def _skip_axes_st(draw, - mutually_broadcastable_shapes=mutually_broadcastable_shapes, - num_skip_axes=None): - shapes, result_shape = draw(mutually_broadcastable_shapes) - if result_shape == (): - assume(num_skip_axes is None) - return () - negative = draw(booleans(), label='skip_axes < 0') - N = len(min(shapes, key=len)) - if num_skip_axes is not None: - min_size = max_size = num_skip_axes - assume(N >= num_skip_axes) - else: - min_size = 0 - max_size = None - if N == 0: - return () - if negative: - axes = draw(lists(integers(-N, -1), min_size=min_size, max_size=max_size, unique=True)) - else: - axes = draw(lists(integers(0, N-1), min_size=min_size, max_size=max_size, unique=True)) - axes = tuple(axes) - # Sometimes return an integer - if num_skip_axes is None and len(axes) == 1 and draw(booleans(), label='skip_axes integer'): # pragma: no cover - return axes[0] - return axes +def _fill_shape(draw, + *, + result_shape, + skip_axes, + skip_axes_values): + max_n = max([i + 1 if i >= 0 else -i for i in skip_axes], default=0) + assume(max_n <= len(skip_axes) + len(result_shape)) + dim = draw(integers(min_value=max_n, max_value=len(skip_axes) + len(result_shape))) + new_shape = ['placeholder']*dim + for i in skip_axes: + assume(new_shape[i] is not None) # skip_axes must be unique + new_shape[i] = None + j = -1 + for i in range(-1, -dim - 1, -1): + if new_shape[i] is None: + new_shape[i] = draw(skip_axes_values) + else: + new_shape[i] = draw(sampled_from([result_shape[j], 1])) + j -= 1 + while new_shape and new_shape[0] == 'placeholder': # pragma: no cover + # Can happen if positive and negative skip_axes refer to the same + # entry + new_shape.pop(0) + + # This will happen if the skip axes are too large + assume('placeholder' not in new_shape) + + if prod([i for i in new_shape if i]) >= SHORT_MAX_ARRAY_SIZE: + note(f"Filtering the shape {new_shape} (too many elements)") + assume(False) + + return tuple(new_shape) -skip_axes_st = shared(_skip_axes_st()) +skip_axes_with_broadcasted_shape_type = shared(sampled_from([int, tuple, list])) @composite -def mutually_broadcastable_shapes_with_skipped_axes(draw, skip_axes_st=skip_axes_st, mutually_broadcastable_shapes=mutually_broadcastable_shapes, -skip_axes_values=integers(0)): +def _mbs_and_skip_axes( + draw, + shapes=short_shapes, + min_shapes=0, + max_shapes=32, + skip_axes_type_st=skip_axes_with_broadcasted_shape_type, + skip_axes_values=integers(0, 20), + num_skip_axes=None, +): """ mutually_broadcastable_shapes except skip_axes() axes might not be broadcastable The result_shape will be None in the position of skip_axes. """ - skip_axes_ = draw(skip_axes_st) - shapes, result_shape = draw(mutually_broadcastable_shapes) - if isinstance(skip_axes_, int): - skip_axes_ = (skip_axes_,) - - # Randomize the shape values in the skipped axes - shapes_ = [] - for shape in shapes: - _shape = list(unremove_indices(shape, skip_axes_)) - # sanity check - assert remove_indices(_shape, skip_axes_) == shape, (_shape, skip_axes_, shape) - - # Replace None values with random values - for j in range(len(_shape)): - if _shape[j] is None: - _shape[j] = draw(skip_axes_values) - shapes_.append(tuple(_shape)) - - result_shape_ = unremove_indices(result_shape, skip_axes_) - # sanity check - assert remove_indices(result_shape_, skip_axes_) == result_shape - - for shape in shapes_: - if prod([i for i in shape if i]) >= SHORT_MAX_ARRAY_SIZE: - note(f"Filtering the shape {shape} (too many elements)") - assume(False) - return BroadcastableShapes(shapes_, result_shape_) - -two_mutually_broadcastable_shapes_1 = shared(_mutually_broadcastable_shapes( - shapes=_short_shapes(1), - min_shapes=2, - max_shapes=2, - min_side=1)) -one_skip_axes = shared(_skip_axes_st( - mutually_broadcastable_shapes=two_mutually_broadcastable_shapes_1, - num_skip_axes=1)) -two_mutually_broadcastable_shapes_2 = shared(_mutually_broadcastable_shapes( - shapes=_short_shapes(2), - min_shapes=2, - max_shapes=2, - min_side=2)) -two_skip_axes = shared(_skip_axes_st( - mutually_broadcastable_shapes=two_mutually_broadcastable_shapes_2, - num_skip_axes=2)) + skip_axes_type = draw(skip_axes_type_st) + _result_shape = draw(shapes) + if _result_shape == (): + assume(num_skip_axes is None) + + ndim = len(_result_shape) + num_shapes = draw(integers(min_value=min_shapes, max_value=max_shapes)) + if not num_shapes: + assume(num_skip_axes is None) + num_skip_axes = 0 + if not ndim: + return BroadcastableShapes([()]*num_shapes, ()), () + + if num_skip_axes is not None: + min_skip_axes = max_skip_axes = num_skip_axes + else: + min_skip_axes = 0 + max_skip_axes = None + + # int and single tuple cases must be limited to N to ensure that they are + # correct for all shapes + if skip_axes_type == int: + assume(num_skip_axes in [None, 1]) + skip_axes = draw(valid_tuple_axes(ndim, min_size=1, max_size=1))[0] + _skip_axes = [(skip_axes,)]*num_shapes + elif skip_axes_type == tuple: + skip_axes = draw(tuples(integers(-ndim, ndim-1), min_size=min_skip_axes, + max_size=max_skip_axes, unique=True)) + _skip_axes = [skip_axes]*num_shapes + elif skip_axes_type == list: + skip_axes = [] + for i in range(num_shapes): + skip_axes.append(draw(tuples(integers(-ndim, ndim+1), min_size=min_skip_axes, + max_size=max_skip_axes, unique=True))) + _skip_axes = skip_axes + + shapes = [] + for i in range(num_shapes): + shapes.append(_fill_shape(draw, result_shape=_result_shape, skip_axes=_skip_axes[i], + skip_axes_values=skip_axes_values)) + + non_skip_shapes = [remove_indices(shape, sk) for shape, sk in + zip(shapes, _skip_axes)] + # Broadcasting the result _fill_shape may produce a shape different from + # _result_shape because it might not have filled all dimensions, or it + # might have chosen 1 for a dimension every time. Ideally we would just be + # using shapes from mutually_broadcastable_shapes, but I don't know how to + # reverse inject skip axes into shapes in general (see the comment in + # unremove_indices). So for now, we just use the actual broadcast of the + # non-skip shapes. Note that we use np.broadcast_shapes here instead of + # ndindex.broadcast_shapes because test_broadcast_shapes itself uses this + # strategy. + broadcasted_shape = broadcast_shapes(*non_skip_shapes) + + return BroadcastableShapes(shapes, broadcasted_shape), skip_axes + +mbs_and_skip_axes = shared(_mbs_and_skip_axes()) + +mutually_broadcastable_shapes_with_skipped_axes = mbs_and_skip_axes.map( + lambda i: i[0]) +skip_axes_st = mbs_and_skip_axes.map(lambda i: i[1]) + +@composite +def _cross_shapes_and_skip_axes(draw): + (shapes, _broadcasted_shape), skip_axes = draw(_mbs_and_skip_axes( + shapes=_short_shapes(2), + min_shapes=2, + max_shapes=2, + num_skip_axes=1, + # TODO: Test other skip axes types + skip_axes_type_st=just(list), + skip_axes_values=just(3), + )) + + broadcasted_skip_axis = draw(integers(-len(_broadcasted_shape)-1, len(_broadcasted_shape))) + broadcasted_shape = unremove_indices(_broadcasted_shape, + [broadcasted_skip_axis], val=3) + skip_axes.append((broadcasted_skip_axis,)) + + return BroadcastableShapes(shapes, broadcasted_shape), skip_axes + +cross_shapes_and_skip_axes = shared(_cross_shapes_and_skip_axes()) +cross_shapes = cross_shapes_and_skip_axes.map(lambda i: i[0]) +cross_skip_axes = cross_shapes_and_skip_axes.map(lambda i: i[1]) + +@composite +def cross_arrays_st(draw): + broadcastable_shapes = draw(cross_shapes) + shapes, broadcasted_shape = broadcastable_shapes + + # Sanity check + assert len(shapes) == 2 + # We need to generate fairly random arrays. Otherwise, if they are too + # similar to each other, like two arange arrays would be, the cross + # product will be 0. We also disable the fill feature in arrays() for the + # same reason, as it would otherwise generate too many vectors that are + # colinear. + a = draw(arrays(dtype=int, shape=shapes[0], elements=integers(-100, 100), fill=nothing())) + b = draw(arrays(dtype=int, shape=shapes[1], elements=integers(-100, 100), fill=nothing())) + + return a, b + +@composite +def _matmul_shapes_and_skip_axes(draw): + (shapes, _broadcasted_shape), skip_axes = draw(_mbs_and_skip_axes( + shapes=_short_shapes(2), + min_shapes=2, + max_shapes=2, + num_skip_axes=2, + # TODO: Test other skip axes types + skip_axes_type_st=just(list), + skip_axes_values=just(None), + )) + + broadcasted_skip_axes = draw(hypothesis_tuples(*[ + integers(-len(_broadcasted_shape)-1, len(_broadcasted_shape)) + ]*2)) + + try: + broadcasted_shape = unremove_indices(_broadcasted_shape, + broadcasted_skip_axes) + except NotImplementedError: + # TODO: unremove_indices only works with both positive or both negative + assume(False) + # Make sure the indices are unique + assume(len(set(broadcasted_skip_axes)) == len(broadcasted_skip_axes)) + + skip_axes.append(broadcasted_skip_axes) + + # (n, m) @ (m, k) -> (n, k) + n, m, k = draw(hypothesis_tuples(integers(0, 10), integers(0, 10), + integers(0, 10))) + shape1, shape2 = map(list, shapes) + ax1, ax2 = skip_axes[0] + shape1[ax1] = n + shape1[ax2] = m + ax1, ax2 = skip_axes[1] + shape2[ax1] = m + shape2[ax2] = k + broadcasted_shape = list(broadcasted_shape) + ax1, ax2 = skip_axes[2] + broadcasted_shape[ax1] = n + broadcasted_shape[ax2] = k + + shapes = (tuple(shape1), tuple(shape2)) + broadcasted_shape = tuple(broadcasted_shape) + + return BroadcastableShapes(shapes, broadcasted_shape), skip_axes + +matmul_shapes_and_skip_axes = shared(_matmul_shapes_and_skip_axes()) +matmul_shapes = matmul_shapes_and_skip_axes.map(lambda i: i[0]) +matmul_skip_axes = matmul_shapes_and_skip_axes.map(lambda i: i[1]) + +@composite +def matmul_arrays_st(draw): + broadcastable_shapes = draw(matmul_shapes) + shapes, broadcasted_shape = broadcastable_shapes + + # Sanity check + assert len(shapes) == 2 + a = draw(arrays(dtype=int, shape=shapes[0], elements=integers(-100, 100))) + b = draw(arrays(dtype=int, shape=shapes[1], elements=integers(-100, 100))) + + return a, b reduce_kwargs = sampled_from([{}, {'negative_int': False}, {'negative_int': True}]) diff --git a/ndindex/tests/test_booleanarray.py b/ndindex/tests/test_booleanarray.py index 0ee19e1d..171140b6 100644 --- a/ndindex/tests/test_booleanarray.py +++ b/ndindex/tests/test_booleanarray.py @@ -1,4 +1,6 @@ -from numpy import prod, arange, array, bool_, empty, full +from numpy import prod, arange, array, bool_, empty, full, __version__ as np_version + +NP1 = np_version.startswith('1') from hypothesis import given, example from hypothesis.strategies import one_of, integers @@ -60,7 +62,9 @@ def test_booleanarray_reduce_hypothesis(idx, shape, kwargs): index = BooleanArray(idx) - check_same(a, index.raw, ndindex_func=lambda a, x: a[x.reduce(shape, **kwargs).raw]) + same_exception = not NP1 + check_same(a, index.raw, ndindex_func=lambda a, x: a[x.reduce(shape, **kwargs).raw], + same_exception=same_exception) try: reduced = index.reduce(shape, **kwargs) diff --git a/ndindex/tests/test_chunking.py b/ndindex/tests/test_chunking.py index 4439ab7c..6e6eaccc 100644 --- a/ndindex/tests/test_chunking.py +++ b/ndindex/tests/test_chunking.py @@ -50,7 +50,7 @@ def test_ChunkSize_args(chunk_size_tuple, idx): try: ndindex(idx) - except ValueError: + except ValueError: # pragma: no cover # Filter out invalid slices (TODO: do this in the strategy) assume(False) diff --git a/ndindex/tests/test_isvalid.py b/ndindex/tests/test_isvalid.py index e307c345..13995b81 100644 --- a/ndindex/tests/test_isvalid.py +++ b/ndindex/tests/test_isvalid.py @@ -5,6 +5,7 @@ from .helpers import ndindices, shapes, MAX_ARRAY_SIZE, check_same, prod +@example([0], (1,)) @example(..., (1, 2, 3)) @example(slice(0, 1), ()) @example(slice(0, 1), (1,)) diff --git a/ndindex/tests/test_shapetools.py b/ndindex/tests/test_shapetools.py index 17531deb..ba8333e3 100644 --- a/ndindex/tests/test_shapetools.py +++ b/ndindex/tests/test_shapetools.py @@ -7,27 +7,27 @@ from hypothesis import assume, given, example from hypothesis.strategies import (one_of, integers, tuples as hypothesis_tuples, just, lists, shared, - composite, nothing) -from hypothesis.extra.numpy import arrays + ) from pytest import raises from ..ndindex import ndindex from ..shapetools import (asshape, iter_indices, ncycles, BroadcastError, AxisError, broadcast_shapes, remove_indices, - unremove_indices, associated_axis) + unremove_indices, associated_axis, + normalize_skip_axes) from ..integer import Integer from ..tuple import Tuple from .helpers import (prod, mutually_broadcastable_shapes_with_skipped_axes, skip_axes_st, mutually_broadcastable_shapes, tuples, - shapes, two_mutually_broadcastable_shapes_1, - two_mutually_broadcastable_shapes_2, one_skip_axes, - two_skip_axes, assert_equal) - -@example([((1, 1), (1, 1)), (None, 1)], (0,)) -@example([((0,), (0,)), (None,)], (0,)) -@example([((1, 2), (2, 1)), (2, None)], 1) -@given(mutually_broadcastable_shapes_with_skipped_axes(), skip_axes_st) + shapes, assert_equal, cross_shapes, cross_skip_axes, + cross_arrays_st, matmul_shapes, matmul_skip_axes, + matmul_arrays_st) + +@example([[(1, 1), (1, 1)], (1,)], (0,)) +@example([[(0,), (0,)], ()], (0,)) +@example([[(1, 2), (2, 1)], (2,)], 1) +@given(mutually_broadcastable_shapes_with_skipped_axes, skip_axes_st) def test_iter_indices(broadcastable_shapes, skip_axes): # broadcasted_shape will contain None on the skip_axes, as those axes # might not be broadcast compatible @@ -37,94 +37,76 @@ def test_iter_indices(broadcastable_shapes, skip_axes): assume(len(broadcasted_shape) < 32) # 1. Normalize inputs - _skip_axes = (skip_axes,) if isinstance(skip_axes, int) else skip_axes - - # Double check the mutually_broadcastable_shapes_with_skipped_axes - # strategy - for i in _skip_axes: - assert broadcasted_shape[i] is None + _skip_axes = normalize_skip_axes(shapes, skip_axes) + _skip_axes_kwarg_default = [()]*len(shapes) # Skipped axes may not be broadcast compatible. Since the index for a # skipped axis should always be a slice(None), the result should be the # same if the skipped axes are all moved to the end of the shape. canonical_shapes = [] - for s in shapes: - c = remove_indices(s, _skip_axes) - c = c + tuple(1 for i in _skip_axes) + for s, sk in zip(shapes, _skip_axes): + c = remove_indices(s, sk) canonical_shapes.append(c) - canonical_skip_axes = list(range(-1, -len(_skip_axes) - 1, -1)) - broadcasted_canonical_shape = list(broadcast_shapes(*canonical_shapes, - skip_axes=canonical_skip_axes)) - for i in range(len(broadcasted_canonical_shape)): - if broadcasted_canonical_shape[i] is None: - broadcasted_canonical_shape[i] = 1 - - skip_shapes = [tuple(1 for i in _skip_axes) for shape in shapes] - non_skip_shapes = [remove_indices(shape, skip_axes) for shape in shapes] - broadcasted_non_skip_shape = remove_indices(broadcasted_shape, skip_axes) - assert None not in broadcasted_non_skip_shape - assert broadcast_shapes(*non_skip_shapes) == broadcasted_non_skip_shape - - nitems = prod(broadcasted_non_skip_shape) - - if _skip_axes == (): + + non_skip_shapes = [remove_indices(shape, sk) for shape, sk in zip(shapes, _skip_axes)] + assert np.broadcast_shapes(*non_skip_shapes) == broadcasted_shape + + nitems = prod(broadcasted_shape) + + if skip_axes == (): # kwarg default res = iter_indices(*shapes) - broadcasted_res = iter_indices(broadcast_shapes(*shapes)) else: res = iter_indices(*shapes, skip_axes=skip_axes) - broadcasted_res = iter_indices(broadcasted_canonical_shape, - skip_axes=canonical_skip_axes) + broadcasted_res = iter_indices(broadcasted_shape) sizes = [prod(shape) for shape in shapes] arrays = [np.arange(size).reshape(shape) for size, shape in zip(sizes, shapes)] canonical_sizes = [prod(shape) for shape in canonical_shapes] canonical_arrays = [np.arange(size).reshape(shape) for size, shape in zip(canonical_sizes, canonical_shapes)] + canonical_broadcasted_array = np.arange(nitems).reshape(broadcasted_shape) # 2. Check that iter_indices is the same whether or not the shapes are # broadcasted together first. Also check that every iterated index is the # expected type and there are as many as expected. vals = [] + bvals = [] n = -1 - def _move_slices_to_end(idx): + def _remove_slices(idx): assert isinstance(idx, Tuple) - idx2 = list(idx.args) - slices = [i for i in range(len(idx2)) if idx2[i] == slice(None)] - idx2 = remove_indices(idx2, slices) - idx2 = idx2 + (slice(None),)*len(slices) + idx2 = [i for i in idx.args if i != slice(None)] return Tuple(*idx2) for n, (idxes, bidxes) in enumerate(zip(res, broadcasted_res)): assert len(idxes) == len(shapes) assert len(bidxes) == 1 - for idx, shape in zip(idxes, shapes): + for idx, shape, sk in zip(idxes, shapes, _skip_axes): assert isinstance(idx, Tuple) assert len(idx.args) == len(shape) - normalized_skip_axes = sorted(ndindex(i).reduce(len(shape)).raw for i in _skip_axes) - for i in range(len(idx.args)): - if i in normalized_skip_axes: + for i in range(-1, -len(idx.args) - 1, -1): + if i in sk: assert idx.args[i] == slice(None) else: assert isinstance(idx.args[i], Integer) - canonical_idxes = [_move_slices_to_end(idx) for idx in idxes] + canonical_idxes = [_remove_slices(idx) for idx in idxes] a_indexed = tuple([a[idx.raw] for a, idx in zip(arrays, idxes)]) canonical_a_indexed = tuple([a[idx.raw] for a, idx in zip(canonical_arrays, canonical_idxes)]) + canonical_b_indexed = canonical_broadcasted_array[bidxes[0].raw] - for c_indexed, skip_shape in zip(canonical_a_indexed, skip_shapes): - assert c_indexed.shape == skip_shape + for c_indexed in canonical_a_indexed: + assert c_indexed.shape == () + assert canonical_b_indexed.shape == () - if _skip_axes: - # If there are skipped axes, recursively call iter_indices to - # get each individual element of the resulting subarrays. - for subidxes in iter_indices(*[x.shape for x in canonical_a_indexed]): - items = [x[i.raw] for x, i in zip(canonical_a_indexed, subidxes)] - vals.append(tuple(items)) + if _skip_axes != _skip_axes_kwarg_default: + vals.append(tuple(canonical_a_indexed)) else: vals.append(a_indexed) + bvals.append(canonical_b_indexed) + # assert both iterators have the same length raises(StopIteration, lambda: next(res)) raises(StopIteration, lambda: next(broadcasted_res)) @@ -141,41 +123,20 @@ def _move_slices_to_end(idx): if not arrays: assert vals == [()] else: - correct_vals = [tuple(i) for i in np.stack(np.broadcast_arrays(*canonical_arrays), axis=-1).reshape((nitems, len(arrays)))] + correct_vals = list(zip(*[x.flat for x in np.broadcast_arrays(*canonical_arrays)])) # Also test that the indices are produced in a lexicographic order # (even though this isn't strictly guaranteed by the iter_indices # docstring) in the case when there are no skip axes. The order when # there are skip axes is more complicated because the skipped axes are # iterated together. - if not _skip_axes: + if _skip_axes == _skip_axes_kwarg_default: assert vals == correct_vals else: assert set(vals) == set(correct_vals) + assert bvals == list(canonical_broadcasted_array.flat) -cross_shapes = mutually_broadcastable_shapes_with_skipped_axes( - mutually_broadcastable_shapes=two_mutually_broadcastable_shapes_1, - skip_axes_st=one_skip_axes, - skip_axes_values=integers(3, 3)) - -@composite -def cross_arrays_st(draw): - broadcastable_shapes = draw(cross_shapes) - shapes, broadcasted_shape = broadcastable_shapes - - # Sanity check - assert len(shapes) == 2 - # We need to generate fairly random arrays. Otherwise, if they are too - # similar to each other, like two arange arrays would be, the cross - # product will be 0. We also disable the fill feature in arrays() for the - # same reason, as it would otherwise generate too many vectors that are - # colinear. - a = draw(arrays(dtype=int, shape=shapes[0], elements=integers(-100, 100), fill=nothing())) - b = draw(arrays(dtype=int, shape=shapes[1], elements=integers(-100, 100), fill=nothing())) - - return a, b - -@given(cross_arrays_st(), cross_shapes, one_skip_axes) -def test_iter_indices_cross(cross_arrays, broadcastable_shapes, skip_axes): +@given(cross_arrays_st(), cross_shapes, cross_skip_axes) +def test_iter_indices_cross(cross_arrays, broadcastable_shapes, _skip_axes): # Test iter_indices behavior against np.cross, which effectively skips the # crossed axis. Note that we don't test against cross products of size 2 # because a 2 x 2 cross product just returns the z-axis (i.e., it doesn't @@ -183,17 +144,17 @@ def test_iter_indices_cross(cross_arrays, broadcastable_shapes, skip_axes): # going to be removed in NumPy 2.0. a, b = cross_arrays shapes, broadcasted_shape = broadcastable_shapes - skip_axis = skip_axes[0] - broadcasted_shape = list(broadcasted_shape) - # Remove None from the shape for iter_indices - broadcasted_shape[skip_axis] = 3 - broadcasted_shape = tuple(broadcasted_shape) + # Sanity check + skip_axes = normalize_skip_axes([*shapes, broadcasted_shape], _skip_axes) + for sh, sk in zip([*shapes, broadcasted_shape], skip_axes): + assert len(sk) == 1 + assert sh[sk[0]] == 3 - res = np.cross(a, b, axisa=skip_axis, axisb=skip_axis, axisc=skip_axis) + res = np.cross(a, b, axisa=skip_axes[0][0], axisb=skip_axes[1][0], axisc=skip_axes[2][0]) assert res.shape == broadcasted_shape - for idx1, idx2, idx3 in iter_indices(*shapes, broadcasted_shape, skip_axes=skip_axes): + for idx1, idx2, idx3 in iter_indices(*shapes, broadcasted_shape, skip_axes=_skip_axes): assert a[idx1.raw].shape == (3,) assert b[idx2.raw].shape == (3,) assert_equal(np.cross( @@ -201,46 +162,7 @@ def test_iter_indices_cross(cross_arrays, broadcastable_shapes, skip_axes): b[idx2.raw]), res[idx3.raw]) - -@composite -def _matmul_shapes(draw): - broadcastable_shapes = draw(mutually_broadcastable_shapes_with_skipped_axes( - mutually_broadcastable_shapes=two_mutually_broadcastable_shapes_2, - skip_axes_st=two_skip_axes, - skip_axes_values=just(None), - )) - shapes, broadcasted_shape = broadcastable_shapes - skip_axes = draw(two_skip_axes) - # (n, m) @ (m, k) -> (n, k) - n, m, k = draw(hypothesis_tuples(integers(0, 10), integers(0, 10), - integers(0, 10))) - - shape1, shape2 = map(list, shapes) - ax1, ax2 = skip_axes - shape1[ax1] = n - shape1[ax2] = m - shape2[ax1] = m - shape2[ax2] = k - broadcasted_shape = list(broadcasted_shape) - broadcasted_shape[ax1] = n - broadcasted_shape[ax2] = k - return [tuple(shape1), tuple(shape2)], tuple(broadcasted_shape) - -matmul_shapes = shared(_matmul_shapes()) - -@composite -def matmul_arrays_st(draw): - broadcastable_shapes = draw(matmul_shapes) - shapes, broadcasted_shape = broadcastable_shapes - - # Sanity check - assert len(shapes) == 2 - a = draw(arrays(dtype=int, shape=shapes[0], elements=integers(-100, 100))) - b = draw(arrays(dtype=int, shape=shapes[1], elements=integers(-100, 100))) - - return a, b - -@given(matmul_arrays_st(), matmul_shapes, two_skip_axes) +@given(matmul_arrays_st(), matmul_shapes, matmul_skip_axes) def test_iter_indices_matmul(matmul_arrays, broadcastable_shapes, skip_axes): # Test iter_indices behavior against np.matmul, which effectively skips the # contracted axis (they aren't broadcasted together, even when they are @@ -248,20 +170,44 @@ def test_iter_indices_matmul(matmul_arrays, broadcastable_shapes, skip_axes): a, b = matmul_arrays shapes, broadcasted_shape = broadcastable_shapes - ax1, ax2 = skip_axes - n, m, k = shapes[0][ax1], shapes[0][ax2], shapes[1][ax2] + # Note, we don't use normalize_skip_axes here because it sorts the skip + # axes - res = np.matmul(a, b, axes=[skip_axes, skip_axes, skip_axes]) + ax1, ax2 = skip_axes[0] + ax3 = skip_axes[1][1] + n, m, k = shapes[0][ax1], shapes[0][ax2], shapes[1][ax3] + + # Sanity check + sk0, sk1, sk2 = skip_axes + shape1, shape2 = shapes + assert a.shape == shape1 + assert b.shape == shape2 + assert shape1[sk0[0]] == n + assert shape1[sk0[1]] == m + assert shape2[sk1[0]] == m + assert shape2[sk1[1]] == k + assert broadcasted_shape[sk2[0]] == n + assert broadcasted_shape[sk2[1]] == k + + res = np.matmul(a, b, axes=skip_axes) assert res.shape == broadcasted_shape + is_ordered = lambda sk, shape: (Integer(sk[0]).reduce(len(shape)).raw <= Integer(sk[1]).reduce(len(shape)).raw) + orders = [ + is_ordered(sk0, shapes[0]), + is_ordered(sk1, shapes[1]), + is_ordered(sk2, broadcasted_shape), + ] + for idx1, idx2, idx3 in iter_indices(*shapes, broadcasted_shape, skip_axes=skip_axes): - assert a[idx1.raw].shape == (n, m) if ax1 <= ax2 else (m, n) - assert b[idx2.raw].shape == (m, k) if ax1 <= ax2 else (k, m) - if ax1 <= ax2: - sub_res = np.matmul(a[idx1.raw], b[idx2.raw]) - else: - sub_res = np.matmul(a[idx1.raw], b[idx2.raw], - axes=[(1, 0), (1, 0), (1, 0)]) + assert a[idx1.raw].shape == (n, m) if orders[0] else (m, n) + assert b[idx2.raw].shape == (m, k) if orders[1] else (k, m) + sub_res_axes = [ + (0, 1) if orders[0] else (1, 0), + (0, 1) if orders[1] else (1, 0), + (0, 1) if orders[2] else (1, 0), + ] + sub_res = np.matmul(a[idx1.raw], b[idx2.raw], axes=sub_res_axes) assert_equal(sub_res, res[idx3.raw]) def test_iter_indices_errors(): @@ -301,12 +247,11 @@ def test_iter_indices_errors(): # Older versions of NumPy do not have the more helpful error message assert ndindex_msg == np_msg - raises(NotImplementedError, lambda: list(iter_indices((1, 2), skip_axes=(0, -1)))) - - with raises(ValueError, match=r"duplicate axes"): + with raises(ValueError, match=r"not unique"): list(iter_indices((1, 2), skip_axes=(0, 1, 0))) - raises(AxisError, lambda: list(iter_indices(skip_axes=(0,)))) + raises(AxisError, lambda: list(iter_indices((0,), skip_axes=(3,)))) + raises(ValueError, lambda: list(iter_indices(skip_axes=(0,)))) raises(TypeError, lambda: list(iter_indices(1, 2))) raises(TypeError, lambda: list(iter_indices(1, 2, (2, 2)))) raises(TypeError, lambda: list(iter_indices([(1, 2), (2, 2)]))) @@ -403,7 +348,7 @@ def test_broadcast_shapes_errors(shapes): raises(TypeError, lambda: broadcast_shapes(1, 2, (2, 2))) raises(TypeError, lambda: broadcast_shapes([(1, 2), (2, 2)])) -@given(mutually_broadcastable_shapes_with_skipped_axes(), skip_axes_st) +@given(mutually_broadcastable_shapes_with_skipped_axes, skip_axes_st) def test_broadcast_shapes_skip_axes(broadcastable_shapes, skip_axes): shapes, broadcasted_shape = broadcastable_shapes assert broadcast_shapes(*shapes, skip_axes=skip_axes) == broadcasted_shape @@ -412,38 +357,32 @@ def test_broadcast_shapes_skip_axes(broadcastable_shapes, skip_axes): @example([[(0, 1)], (0, 1)], (2,)) @example([[(0, 1)], (0, 1)], (0, -1)) @example([[(0, 1, 0, 0, 0), (2, 0, 0, 0)], (0, 2, 0, 0, 0)], [1]) -@given(mutually_broadcastable_shapes, lists(integers(-20, 20), max_size=20)) +@given(mutually_broadcastable_shapes, + one_of( + integers(-20, 20), + tuples(integers(-20, 20), max_size=20), + lists(tuples(integers(-20, 20), max_size=20), max_size=32))) def test_broadcast_shapes_skip_axes_errors(broadcastable_shapes, skip_axes): shapes, broadcasted_shape = broadcastable_shapes - if any(i < 0 for i in skip_axes) and any(i >= 0 for i in skip_axes): - raises(NotImplementedError, lambda: broadcast_shapes(*shapes, skip_axes=skip_axes)) - return + # All errors should come from normalize_skip_axes, which is tested + # separately below. try: - if not shapes and skip_axes: - raise IndexError - for shape in shapes: - for i in skip_axes: - shape[i] - except IndexError: - error = True - else: - error = False + normalize_skip_axes(shapes, skip_axes) + except (TypeError, ValueError, IndexError) as e: + raises(type(e), lambda: broadcast_shapes(*shapes, + skip_axes=skip_axes)) + return try: broadcast_shapes(*shapes, skip_axes=skip_axes) except IndexError: - if not error: # pragma: no cover - raise RuntimeError("broadcast_shapes raised but should not have") - return + raise RuntimeError("broadcast_shapes raised but should not have") # pragma: no cover except BroadcastError: # Broadcastable shapes can become unbroadcastable after skipping axes # (see the @example above). pass - if error: # pragma: no cover - raise RuntimeError("broadcast_shapes did not raise but should have") - remove_indices_n = shared(integers(0, 100)) @given(remove_indices_n, @@ -454,6 +393,8 @@ def test_remove_indices(n, idxes): assume(min(idxes) >= -n) a = tuple(range(n)) b = remove_indices(a, idxes) + if len(idxes) == 1: + assert remove_indices(a, idxes[0]) == b A = list(a) for i in idxes: @@ -479,53 +420,49 @@ def test_remove_indices(n, idxes): raises(NotImplementedError, lambda: unremove_indices(b, idxes)) # Meta-test for the hypothesis strategy -@given(mutually_broadcastable_shapes_with_skipped_axes(), skip_axes_st) +@given(mutually_broadcastable_shapes_with_skipped_axes, skip_axes_st) def test_mutually_broadcastable_shapes_with_skipped_axes(broadcastable_shapes, skip_axes): # pragma: no cover shapes, broadcasted_shape = broadcastable_shapes - _skip_axes = (skip_axes,) if isinstance(skip_axes, int) else skip_axes + _skip_axes = normalize_skip_axes(shapes, skip_axes) + + assert len(_skip_axes) == len(shapes) for shape in shapes: assert None not in shape - for i in _skip_axes: - assert broadcasted_shape[i] is None + assert None not in broadcasted_shape - _shapes = [remove_indices(shape, skip_axes) for shape in shapes] - _broadcasted_shape = remove_indices(broadcasted_shape, skip_axes) + _shapes = [remove_indices(shape, sk) for shape, sk in zip(shapes, _skip_axes)] - assert None not in _broadcasted_shape - assert broadcast_shapes(*_shapes) == _broadcasted_shape + assert broadcast_shapes(*_shapes) == broadcasted_shape -@example([[(2, 10, 3, 4), (10, 3, 4)], (2, None, 3, 4)], (-3,)) +@example([[(2, 10, 3, 4), (10, 3, 4)], (2, 3, 4)], (-3,)) @example([[(0, 10, 2, 3, 10, 4), (1, 10, 1, 0, 10, 2, 3, 4)], - (1, None, 1, 0, None, 2, 3, 4)], (1, 4)) -@example([[(2, 0, 3, 4)], (2, None, 3, 4)], (1,)) -@example([[(0, 0, 0, 0, 0), (0, 0, 0, 0, 0, 0)], (0, None, None, 0, 0, 0)], (1, 2)) -@given(mutually_broadcastable_shapes_with_skipped_axes(), skip_axes_st) + (1, 1, 0, 2, 3, 4)], (1, 4)) +@example([[(2, 0, 3, 4)], (2, 3, 4)], (1,)) +@example([[(0, 0, 0, 0, 0), (0, 0, 0, 0, 0, 0)], (0, 0, 0, 0)], (1, 2)) +@given(mutually_broadcastable_shapes_with_skipped_axes, skip_axes_st) def test_associated_axis(broadcastable_shapes, skip_axes): - _skip_axes = (skip_axes,) if isinstance(skip_axes, int) else skip_axes - shapes, broadcasted_shape = broadcastable_shapes - ndim = len(broadcasted_shape) + _skip_axes = normalize_skip_axes(shapes, skip_axes) - normalized_skip_axes = [ndindex(i).reduce(ndim) for i in _skip_axes] - - for shape in shapes: + for shape, sk in zip(shapes, _skip_axes): n = len(shape) for i in range(-len(shape), 0): val = shape[i] - idx = associated_axis(shape, broadcasted_shape, i, _skip_axes) - bval = broadcasted_shape[idx] + bval = associated_axis(broadcasted_shape, i, sk) if bval is None: - if _skip_axes[0] >= 0: - assert ndindex(i).reduce(n) == ndindex(idx).reduce(ndim) in normalized_skip_axes - else: - assert ndindex(i).reduce(n, negative_int=True) == \ - ndindex(idx).reduce(ndim, negative_int=True) in _skip_axes + assert ndindex(i).reduce(n, negative_int=True) in sk, (shape, i) else: - assert val == 1 or bval == val + assert val == 1 or bval == val, (shape, i) + + sk = max(_skip_axes, key=len, default=()) + for i in range(-len(broadcasted_shape)-len(sk)-10, -len(broadcasted_shape)-len(sk)): + assert associated_axis(broadcasted_shape, i, sk) is None + +# TODO: add a hypothesis test for asshape def test_asshape(): assert asshape(1) == (1,) assert asshape(np.int64(2)) == (2,) @@ -557,3 +494,118 @@ def test_asshape(): raises(TypeError, lambda: asshape(-1, allow_int=False)) raises(TypeError, lambda: asshape(-1, allow_negative=True, allow_int=False)) raises(TypeError, lambda: asshape(np.int64(1), allow_int=False)) + raises(IndexError, lambda: asshape((2, 3), 3)) + +@example([], []) +@example([()], []) +@example([(0, 1)], 0) +@example([(2, 3), (2, 3, 4)], [(3,), (0,)]) +@example([(0, 1)], 0) +@example([(2, 3)], (0, -2)) +@example([(2, 4), (2, 3, 4)], [(0,), (-3,)]) +@given(lists(tuples(integers(0))), + one_of(integers(), tuples(integers()), lists(tuples(integers())))) +def test_normalize_skip_axes(shapes, skip_axes): + if not shapes: + if skip_axes in [(), []]: + assert normalize_skip_axes(shapes, skip_axes) == [] + else: + raises(ValueError, lambda: normalize_skip_axes(shapes, skip_axes)) + return + + min_dim = min(len(shape) for shape in shapes) + + if isinstance(skip_axes, int): + if not (-min_dim <= skip_axes < min_dim): + raises(AxisError, lambda: normalize_skip_axes(shapes, skip_axes)) + return + _skip_axes = [(skip_axes,)]*len(shapes) + skip_len = 1 + elif isinstance(skip_axes, tuple): + if not all(-min_dim <= s < min_dim for s in skip_axes): + raises(AxisError, lambda: normalize_skip_axes(shapes, skip_axes)) + return + _skip_axes = [skip_axes]*len(shapes) + skip_len = len(skip_axes) + elif not skip_axes: + # empty list will be interpreted as a single skip_axes tuple + assert normalize_skip_axes(shapes, skip_axes) == [()]*len(shapes) + return + else: + if len(shapes) != len(skip_axes): + raises(ValueError, lambda: normalize_skip_axes(shapes, skip_axes)) + return + _skip_axes = skip_axes + skip_len = len(skip_axes[0]) + + try: + res = normalize_skip_axes(shapes, skip_axes) + except AxisError as e: + axis, ndim = e.args + assert any(axis in s for s in _skip_axes) + assert any(ndim == len(shape) for shape in shapes) + assert axis < -ndim or axis >= ndim + return + except ValueError as e: + if 'not unique' in str(e): + bad_skip_axes, bad_shape = e.skip_axes, e.shape + assert str(bad_skip_axes) in str(e) + assert str(bad_shape) in str(e) + assert bad_skip_axes in _skip_axes + assert bad_shape in shapes + indexed = [bad_shape[i] for i in bad_skip_axes] + assert len(indexed) != len(set(indexed)) + return + else: # pragma: no cover + raise + + assert isinstance(res, list) + assert all(isinstance(x, tuple) for x in res) + assert all(isinstance(i, int) for x in res for i in x) + + assert len(res) == len(shapes) + for shape, new_skip_axes in zip(shapes, res): + assert len(new_skip_axes) == len(set(new_skip_axes)) == skip_len + assert new_skip_axes == tuple(sorted(new_skip_axes)) + for i in new_skip_axes: + assert i < 0 + assert ndindex(i).reduce(len(shape), negative_int=True) == i + + # TODO: Assert the order is maintained (doesn't actually matter for now + # but could for future applications) + +def test_normalize_skip_axes_errors(): + raises(TypeError, lambda: normalize_skip_axes([(1,)], {0: 1})) + raises(TypeError, lambda: normalize_skip_axes([(1,)], {0})) + raises(TypeError, lambda: normalize_skip_axes([(1,)], [(0,), 0])) + raises(TypeError, lambda: normalize_skip_axes([(1,)], [0, (0,)])) + +@example(10, 5) +@given(integers(), integers()) +def test_axiserror(axis, ndim): + if ndim == 0 and axis in [0, -1]: + # NumPy allows axis=0 or -1 for 0-d arrays + AxisError(axis, ndim) + return + + try: + if ndim >= 0: + range(ndim)[axis] + except IndexError: + e = AxisError(axis, ndim) + else: + raises(ValueError, lambda: AxisError(axis, ndim)) + return + + try: + raise e + except AxisError as e2: + assert e2.args == (axis, ndim) + if ndim <= 32 and -1000 < axis < 1000: + a = np.empty((0,)*ndim) + try: + np.sum(a, axis=axis) + except np_AxisError as e3: + assert str(e2) == str(e3) + else: + raise RuntimeError("numpy didn't raise AxisError") # pragma: no cover diff --git a/requirements-dev.txt b/requirements-dev.txt new file mode 100644 index 00000000..7d960665 --- /dev/null +++ b/requirements-dev.txt @@ -0,0 +1,8 @@ +hypothesis +numpy +packaging +pyflakes +pytest +pytest-cov +pytest-doctestplus +pytest-flakes