def _lazywhere(cond, arrays, f, fillvalue=None, f2=None): """ np.where(cond, x, fillvalue) always evaluates x even where cond is False. This one only evaluates f(arr1[cond], arr2[cond], ...). For example, >>> a, b = np.array([1, 2, 3, 4]), np.array([5, 6, 7, 8]) >>> def f(a, b): return a*b >>> _lazywhere(a > 2, (a, b), f, np.nan) array([ nan, nan, 21., 32.]) Notice it assumes that all `arrays` are of the same shape, or can be broadcasted together. """ if fillvalue is None: if f2 is None: raise ValueError("One of (fillvalue, f2) must be given.") else: fillvalue = np.nan else: if f2 is not None: raise ValueError("Only one of (fillvalue, f2) can be given.") arrays = np.broadcast_arrays(*arrays) temp = tuple(np.extract(cond, arr) for arr in arrays) tcode = np.mintypecode([a.dtype.char for a in arrays]) out = _valarray(np.shape(arrays[0]), value=fillvalue, typecode=tcode) np.place(out, cond, f(*temp)) if f2 is not None: temp = tuple(np.extract(~cond, arr) for arr in arrays) np.place(out, ~cond, f2(*temp)) return out
def _lazyselect(condlist, choicelist, arrays, default=0): """ Mimic `np.select(condlist, choicelist)`. Notice it assumes that all `arrays` are of the same shape, or can be broadcasted together. All functions in `choicelist` must accept array arguments in the order given in `arrays` and must return an array of the same shape as broadcasted `arrays`. Examples -------- >>> x = np.arange(6) >>> np.select([x <3, x > 3], [x**2, x**3], default=0) array([ 0, 1, 4, 0, 64, 125]) >>> _lazyselect([x < 3, x > 3], [lambda x: x**2, lambda x: x**3], (x,)) array([ 0., 1., 4., 0., 64., 125.]) >>> a = -np.ones_like(x) >>> _lazyselect([x < 3, x > 3], ... [lambda x, a: x**2, lambda x, a: a * x**3], ... (x, a), default=np.nan) array([ 0., 1., 4., nan, -64., -125.]) """ arrays = np.broadcast_arrays(*arrays) tcode = np.mintypecode([a.dtype.char for a in arrays]) out = _valarray(np.shape(arrays[0]), value=default, typecode=tcode) for index in range(len(condlist)): func, cond = choicelist[index], condlist[index] if np.all(cond is False): continue cond, _ = np.broadcast_arrays(cond, arrays[0]) temp = tuple(np.extract(cond, arr) for arr in arrays) np.place(out, cond, func(*temp)) return out
def _lazyselect(condlist, choicelist, arrays, default=0): """ Mimic `np.select(condlist, choicelist)`. Notice, it assumes that all `arrays` are of the same shape or can be broadcasted together. All functions in `choicelist` must accept array arguments in the order given in `arrays` and must return an array of the same shape as broadcasted `arrays`. Examples -------- >>> x = np.arange(6) >>> np.select([x <3, x > 3], [x**2, x**3], default=0) array([ 0, 1, 4, 0, 64, 125]) >>> _lazyselect([x < 3, x > 3], [lambda x: x**2, lambda x: x**3], (x,)) array([ 0., 1., 4., 0., 64., 125.]) >>> a = -np.ones_like(x) >>> _lazyselect([x < 3, x > 3], ... [lambda x, a: x**2, lambda x, a: a * x**3], ... (x, a), default=np.nan) array([ 0., 1., 4., nan, -64., -125.]) """ arrays = np.broadcast_arrays(*arrays) tcode = np.mintypecode([a.dtype.char for a in arrays]) out = np.full(np.shape(arrays[0]), fill_value=default, dtype=tcode) for index in range(len(condlist)): func, cond = choicelist[index], condlist[index] if np.all(cond is False): continue cond, _ = np.broadcast_arrays(cond, arrays[0]) temp = tuple(np.extract(cond, arr) for arr in arrays) np.place(out, cond, func(*temp)) return out
import numpy as np np.mintypecode(['d', 'f', 'S']) x = np.array([1.1, 2 - 3.j]) np.mintypecode(x) np.mintypecode('abceh', default='G')
AR_f16: npt.NDArray[np.floating[npt._128Bit]] AR_c8: npt.NDArray[np.complex64] AR_c16: npt.NDArray[np.complex128] AR_LIKE_f: List[float] class RealObj: real: slice class ImagObj: imag: slice reveal_type(np.mintypecode(["f8"], typeset="qfQF")) reveal_type( np.asfarray(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] reveal_type( np.asfarray(AR_LIKE_f)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] reveal_type( np.asfarray(AR_f8, dtype="c16") ) # E: numpy.ndarray[Any, numpy.dtype[numpy.complexfloating[Any, Any]]] reveal_type(np.asfarray( AR_f8, dtype="i8")) # E: numpy.ndarray[Any, numpy.dtype[numpy.floating[Any]]] reveal_type(np.real(RealObj())) # E: slice reveal_type(np.real(AR_f8)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]] reveal_type(np.real(AR_c16)) # E: numpy.ndarray[Any, numpy.dtype[{float64}]]
import numpy as np import numpy.typing as npt DTYPE_i8: np.dtype[np.int64] np.mintypecode(DTYPE_i8) # E: incompatible type np.iscomplexobj(DTYPE_i8) # E: incompatible type np.isrealobj(DTYPE_i8) # E: incompatible type np.typename(DTYPE_i8) # E: No overload variant np.typename("invalid") # E: No overload variant np.common_type(np.timedelta64()) # E: incompatible type