def test_length_bounds_are_satisfied(ndim, data):
min_size = data.draw(st.integers(0, ndim), label="min_size")
max_size = data.draw(st.integers(min_size, ndim), label="max_size")
axes = data.draw(
nps.valid_tuple_axes(ndim, min_size=min_size, max_size=max_size), label="axes"
)
assert min_size <= len(axes) <= max_size
def test_mapped_positive_axes_are_unique(ndim, data):
min_size = data.draw(st.integers(0, ndim), label="min_size")
max_size = data.draw(st.integers(min_size, ndim), label="max_size")
axes = data.draw(
nps.valid_tuple_axes(ndim, min_size=min_size, max_size=max_size), label="axes"
)
assert len(set(axes)) == len({i if 0 < i else ndim + i for i in axes})
def test_minimize_negative_tuple_axes(ndim, data):
min_size = data.draw(st.integers(0, ndim), label="min_size")
max_size = data.draw(st.integers(min_size, ndim), label="max_size")
smallest = minimal(
nps.valid_tuple_axes(ndim, min_size, max_size), lambda x: all(i < 0 for i in x)
)
assert len(smallest) == min_size
def gen_roll_args(draw, arr):
shift = draw(st.integers() | st.tuples(*(st.integers()
for i in arr.shape)))
ax_strat = hnp.valid_tuple_axes(
arr.ndim,
**(dict(min_size=len(shift), max_size=len(shift)) if isinstance(
shift, tuple) else {}))
axis = draw(st.none() | st.integers(-arr.ndim, arr.ndim - 1) | ax_strat)
return dict(shift=shift, axis=axis)
def test_minimize_tuple_axes(ndim, data):
min_size = data.draw(st.integers(0, ndim), label="min_size")
max_size = data.draw(st.integers(min_size, ndim), label="max_size")
smallest = minimal(nps.valid_tuple_axes(ndim, min_size=min_size, max_size=max_size))
assert len(smallest) == min_size and all(k > -1 for k in smallest)
def test_axes_are_valid_inputs_to_sum(shape, data):
x = np.zeros(shape, dtype="uint8")
axes = data.draw(nps.valid_tuple_axes(ndim=len(shape)), label="axes")
np.sum(x, axes)
def valid_axes(
ndim: int,
pos_only: bool = False,
single_axis_only: bool = False,
permit_none: bool = True,
permit_int: bool = True,
min_dim: int = 0,
max_dim: Optional[int] = None,
) -> st.SearchStrategy[Union[None, int, Tuple[int, ...]]]:
""" Hypothesis search strategy: Given array dimensionality, generate valid
`axis` arguments (including `None`) for numpy's sequential functions.
Examples from this strategy shrink towards an empty tuple of axes.
If `single_axis_only=True`, then it shrinks towards 0.
Parameters
----------
ndim : int
The dimensionality of the array.
pos_only : bool, optional (default=False)
If True, the returned value(s) will be positive.
single_axis_only : bool, optional (default=False)
If True, a single integer axis or `None` (assuming `permit_none=True`)
will be returned.
permit_none : bool, optional (default=True)
If True, `None` may be returned instead of a tuple of all of the
available axes.
permit_int: bool, optional (default=True)
If True, the returned value may be an integer
min_dim : int, optional (default=0)
The smallest number of entries permitted in the returned tuple of axes
max_dim : Optional[int]
The largest number of entries permitted in the returned tuple of axes.
The defaults is ``ndim``.
Returns
-------
st.SearchStrategy[Union[None, int, Tuple[int, ...]]]
Examples
--------
>>> valid_axes(4).example()
(0, 1)
"""
if isinstance(ndim, (tuple, list)):
ndim = len(ndim)
single_axis_strat = st.integers(-ndim, ndim - 1) if ndim else st.just(0)
strats = []
if permit_none:
strats.append(st.none())
if permit_int and min_dim <= 1 and (max_dim is None or 1 <= max_dim):
strats.append(single_axis_strat)
if not single_axis_only:
strats.append(
hnp.valid_tuple_axes(ndim, min_size=min_dim, max_size=max_dim))
strat = st.one_of(*strats)
if pos_only:
strat = strat.map(lambda x: x if x is None else _to_positive(x, ndim))
return strat
