|
5 | 5 |
|
6 | 6 | import math |
7 | 7 | import warnings |
8 | | -from collections.abc import Sequence |
| 8 | +from collections.abc import Callable, Sequence |
9 | 9 | from types import ModuleType |
10 | | -from typing import cast |
| 10 | +from typing import cast, overload |
11 | 11 |
|
12 | 12 | from ._at import at |
13 | 13 | from ._utils import _compat, _helpers |
14 | | -from ._utils._compat import array_namespace, is_jax_array |
15 | | -from ._utils._helpers import asarrays |
| 14 | +from ._utils._compat import ( |
| 15 | + array_namespace, |
| 16 | + is_dask_namespace, |
| 17 | + is_jax_array, |
| 18 | + is_jax_namespace, |
| 19 | +) |
| 20 | +from ._utils._helpers import asarrays, meta_namespace |
16 | 21 | from ._utils._typing import Array |
17 | 22 |
|
18 | 23 | __all__ = [ |
| 24 | + "apply_where", |
19 | 25 | "atleast_nd", |
20 | 26 | "broadcast_shapes", |
21 | 27 | "cov", |
|
29 | 35 | ] |
30 | 36 |
|
31 | 37 |
|
| 38 | +@overload |
| 39 | +def apply_where( # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08 |
| 40 | + cond: Array, |
| 41 | + f1: Callable[..., Array], |
| 42 | + f2: Callable[..., Array], |
| 43 | + args: Array | tuple[Array, ...], |
| 44 | + /, |
| 45 | + *, |
| 46 | + xp: ModuleType | None = None, |
| 47 | +) -> Array: ... |
| 48 | + |
| 49 | + |
| 50 | +@overload |
| 51 | +def apply_where( # type: ignore[no-any-explicit,no-any-decorated] # numpydoc ignore=GL08 |
| 52 | + cond: Array, |
| 53 | + f1: Callable[..., Array], |
| 54 | + args: Array | tuple[Array, ...], |
| 55 | + /, |
| 56 | + *, |
| 57 | + fill_value: Array | int | float | complex | bool, |
| 58 | + xp: ModuleType | None = None, |
| 59 | +) -> Array: ... |
| 60 | + |
| 61 | + |
| 62 | +def apply_where( # type: ignore[no-any-explicit] # numpydoc ignore=PR01,PR02 |
| 63 | + cond: Array, |
| 64 | + f1: Callable[..., Array], |
| 65 | + f2: Callable[..., Array] | Array | tuple[Array], # optional positional argument |
| 66 | + args: Array | tuple[Array, ...] | None = None, |
| 67 | + /, |
| 68 | + *, |
| 69 | + fill_value: Array | int | float | complex | bool | None = None, |
| 70 | + xp: ModuleType | None = None, |
| 71 | +) -> Array: |
| 72 | + """ |
| 73 | + Run one of two elementwise functions depending on a condition. |
| 74 | +
|
| 75 | + Equivalent to ``f1(*args) if cond else fill_value`` performed elementwise |
| 76 | + when `fill_value` is defined, otherwise to ``f1(*args) if cond else f2(*args)``. |
| 77 | +
|
| 78 | + Parameters |
| 79 | + ---------- |
| 80 | + cond : array |
| 81 | + The condition, expressed as a boolean array. |
| 82 | + f1 : callable |
| 83 | + Elementwise function of `args`, returning a single array. |
| 84 | + Where `cond` is True, output will be ``f1(arg0[cond], arg1[cond], ...)``. |
| 85 | + f2 : callable, optional |
| 86 | + Elementwise function of `args`, returning a single array. |
| 87 | + Where `cond` is False, output will be ``f2(arg0[cond], arg1[cond], ...)``. |
| 88 | + Mutually exclusive with `fill_value`. |
| 89 | + args : Array or tuple of Arrays |
| 90 | + Argument(s) to `f1` (and `f2`). Must be broadcastable with `cond`. |
| 91 | + fill_value : Array or scalar, optional |
| 92 | + If provided, value with which to fill output array where `cond` is False. |
| 93 | + It does not need to be scalar. |
| 94 | + Mutually exclusive with `f2`. You must provide one or the other. |
| 95 | + xp : array_namespace, optional |
| 96 | + The standard-compatible namespace for `cond` and `args`. Default: infer. |
| 97 | +
|
| 98 | + Returns |
| 99 | + ------- |
| 100 | + Array |
| 101 | + An array with elements from the output of `f1` where `cond` is True and either |
| 102 | + the output of `f2` or `fill_value` where `cond` is False. The returned array has |
| 103 | + data type determined by type promotion rules between the output of `f1` and |
| 104 | + either `fill_value` or the output of `f2`. |
| 105 | +
|
| 106 | + Notes |
| 107 | + ----- |
| 108 | + ``xp.where(cond, f1(*args), f2(*args))`` requires explicitly evaluating `f1` even |
| 109 | + when `cond` is False, and `f2` when cond is True. This function evaluates each |
| 110 | + function only for their matching condition, if the backend allows for it. |
| 111 | +
|
| 112 | + On Dask, `f1` and `f2` are applied to the individual chunks and should use functions |
| 113 | + from the namespace of the chunks. |
| 114 | +
|
| 115 | + Examples |
| 116 | + -------- |
| 117 | + >>> a = xp.asarray([5, 4, 3]) |
| 118 | + >>> b = xp.asarray([0, 2, 2]) |
| 119 | + >>> def f(a, b): |
| 120 | + ... return a // b |
| 121 | + >>> apply_where(b != 0, f, (a, b), fill_value=xp.nan) |
| 122 | + array([ nan, 2., 1.]) |
| 123 | + """ |
| 124 | + # Parse and normalize arguments |
| 125 | + mutually_exc_msg = "Exactly one of `fill_value` or `f2` must be given." |
| 126 | + if args is None: |
| 127 | + f2, args = None, f2 |
| 128 | + if fill_value is None: |
| 129 | + raise TypeError(mutually_exc_msg) |
| 130 | + else: |
| 131 | + if not callable(f2): |
| 132 | + msg = "Third parameter must be a callable, Array, or tuple of Arrays." |
| 133 | + raise TypeError(msg) |
| 134 | + if fill_value is not None: |
| 135 | + raise TypeError(mutually_exc_msg) |
| 136 | + |
| 137 | + if not isinstance(args, tuple): |
| 138 | + args = (args,) |
| 139 | + f2 = cast(Callable[..., Array] | None, f2) # type: ignore[no-any-explicit] |
| 140 | + args = cast(tuple[Array, ...], args) |
| 141 | + |
| 142 | + xp = array_namespace(cond, *args) if xp is None else xp |
| 143 | + |
| 144 | + if getattr(fill_value, "ndim", 0): |
| 145 | + cond, fill_value, *args = xp.broadcast_arrays(cond, fill_value, *args) |
| 146 | + else: |
| 147 | + cond, *args = xp.broadcast_arrays(cond, *args) |
| 148 | + |
| 149 | + if is_dask_namespace(xp): |
| 150 | + meta_xp = meta_namespace(cond, fill_value, *args, xp=xp) |
| 151 | + return xp.map_blocks(_apply_where, cond, f1, f2, fill_value, *args, xp=meta_xp) |
| 152 | + return _apply_where(cond, f1, f2, fill_value, *args, xp=xp) |
| 153 | + |
| 154 | + |
| 155 | +def _apply_where( # type: ignore[no-any-explicit] # numpydoc ignore=PR01,RT01 |
| 156 | + cond: Array, |
| 157 | + f1: Callable[..., Array], |
| 158 | + f2: Callable[..., Array] | None, |
| 159 | + fill_value: Array | int | float | complex | bool | None, |
| 160 | + *args: Array, |
| 161 | + xp: ModuleType, |
| 162 | +) -> Array: |
| 163 | + """Helper of `apply_where`. On Dask, this runs on a single chunk.""" |
| 164 | + |
| 165 | + if is_jax_namespace(xp): |
| 166 | + # jax.jit does not support assignment by boolean mask |
| 167 | + return xp.where(cond, f1(*args), f2(*args) if f2 is not None else fill_value) |
| 168 | + |
| 169 | + device = _compat.device(cond) |
| 170 | + temp1 = f1(*(arr[cond] for arr in args)) |
| 171 | + |
| 172 | + if f2 is None: |
| 173 | + # TODO remove asarrays once all backends support Array API 2024.12 |
| 174 | + dtype = xp.result_type(*asarrays(temp1, fill_value, xp=xp)) |
| 175 | + if getattr(fill_value, "ndim", 0): |
| 176 | + fill_value = xp.astype(fill_value, dtype) |
| 177 | + return at(fill_value, cond).set(temp1, copy=True) |
| 178 | + out = xp.full(cond.shape, fill_value=fill_value, dtype=dtype, device=device) |
| 179 | + else: |
| 180 | + ncond = ~cond |
| 181 | + temp2 = f2(*(arr[ncond] for arr in args)) |
| 182 | + dtype = xp.result_type(temp1, temp2) |
| 183 | + out = xp.empty(cond.shape, dtype=dtype, device=device) |
| 184 | + out = at(out, ncond).set(temp2) |
| 185 | + |
| 186 | + return at(out, cond).set(temp1) |
| 187 | + |
| 188 | + |
32 | 189 | def atleast_nd(x: Array, /, *, ndim: int, xp: ModuleType | None = None) -> Array: |
33 | 190 | """ |
34 | 191 | Recursively expand the dimension of an array to at least `ndim`. |
@@ -385,12 +542,15 @@ def isclose( |
385 | 542 | a_inexact = xp.isdtype(a.dtype, ("real floating", "complex floating")) |
386 | 543 | b_inexact = xp.isdtype(b.dtype, ("real floating", "complex floating")) |
387 | 544 | if a_inexact or b_inexact: |
388 | | - # FIXME: use scipy's lazywhere to suppress warnings on inf |
389 | | - out = xp.where( |
| 545 | + # prevent warnings on numpy and dask on inf - inf |
| 546 | + mxp = meta_namespace(a, b, xp=xp) |
| 547 | + out = apply_where( |
390 | 548 | xp.isinf(a) | xp.isinf(b), |
391 | | - xp.isinf(a) & xp.isinf(b) & (xp.sign(a) == xp.sign(b)), |
| 549 | + lambda a, b: mxp.isinf(a) & mxp.isinf(b) & (mxp.sign(a) == mxp.sign(b)), # pyright: ignore[reportUnknownArgumentType] |
392 | 550 | # Note: inf <= inf is True! |
393 | | - xp.abs(a - b) <= (atol + rtol * xp.abs(b)), |
| 551 | + lambda a, b: mxp.abs(a - b) <= (atol + rtol * mxp.abs(b)), # pyright: ignore[reportUnknownArgumentType] |
| 552 | + (a, b), |
| 553 | + xp=xp, |
394 | 554 | ) |
395 | 555 | if equal_nan: |
396 | 556 | out = xp.where(xp.isnan(a) & xp.isnan(b), xp.asarray(True), out) |
|
0 commit comments