diff --git a/arkouda/numpy/numeric.py b/arkouda/numpy/numeric.py index ac0d214ae50..cf26a607341 100644 --- a/arkouda/numpy/numeric.py +++ b/arkouda/numpy/numeric.py @@ -30,8 +30,11 @@ ARKOUDA_SUPPORTED_INTS, _datatype_check, bigint, + bool_scalars, int_scalars, + isSupportedBool, isSupportedNumber, + numeric_and_bool_scalars, numeric_scalars, resolve_scalar_dtype, str_, @@ -52,7 +55,7 @@ ) from arkouda.numpy.pdarrayclass import all as ak_all from arkouda.numpy.pdarrayclass import any as ak_any -from arkouda.numpy.pdarraycreation import array, linspace, scalar_array +from arkouda.numpy.pdarraycreation import array, linspace from arkouda.numpy.sorting import sort from arkouda.numpy.strings import Strings @@ -154,7 +157,7 @@ class ErrorMode(Enum): # TODO: standardize error checking in python interface -# merge_where comes in handy in arctan2 and some other functions. +# _merge_where comes in handy in various functions. def _merge_where(new_pda, where, ret): @@ -163,6 +166,15 @@ def _merge_where(new_pda, where, ret): return new_pda +# _normalize_scalar is needed in arctan2 to handle broadcasts of np.bool_ scalars + + +def _normalize_scalar(x): + if isinstance(x, np.generic): + return x.item() + return x + + @overload def cast( pda: pdarray, @@ -1397,10 +1409,13 @@ def arctan(pda: pdarray, where: Union[bool, pdarray] = True) -> pdarray: @typechecked def arctan2( - num: Union[pdarray, numeric_scalars], - denom: Union[pdarray, numeric_scalars], - where: Union[bool, pdarray] = True, -) -> pdarray: + x1: Union[pdarray, numeric_and_bool_scalars], + x2: Union[pdarray, numeric_and_bool_scalars], + /, + out: Optional[pdarray] = None, + *, + where: Optional[Union[bool_scalars, pdarray]] = None, +) -> Union[pdarray, numeric_scalars]: """ Return the element-wise inverse tangent of the array pair. The result chosen is the signed angle in radians between the ray ending at the origin and passing through the @@ -1409,11 +1424,13 @@ def arctan2( Parameters ---------- - num : pdarray or numeric_scalars + x1 : pdarray or numeric_scalars Numerator of the arctan2 argument. - denom : pdarray or numeric_scalars + x2 : pdarray or numeric_scalars Denominator of the arctan2 argument. - where : bool or pdarray, default=True + out : Optional, pdarray + A pdarray in which to store the result, or to use as a source when where is False. + where : Optional, bool_scalars or pdarray, default=None This condition is broadcast over the input. At locations where the condition is True, the inverse tangent will be applied to the corresponding values. Elsewhere, it will retain its original value. Default set to True. @@ -1432,6 +1449,8 @@ def arctan2( | Raised if any element of pdarrays num and denom is not a supported type | Raised if both num and denom are scalars | Raised if where is neither boolean nor a pdarray of boolean + ValueError + Raised if broadcasting of the given parameters isn't possible. Examples -------- @@ -1441,72 +1460,210 @@ def arctan2( >>> ak.arctan2(y,x) array([0.78539816... 2.35619449... -2.35619449... -0.78539816...]) """ - from arkouda.client import generic_msg + # The line below is needed in order to enable use of arkouda's "where" function without + # a name conflict, since "where" is also a parameter name. + + from arkouda.numpy.numeric import where as ak_where + + # And we may need these imports for broadcasting. + from arkouda.numpy.util import broadcast_shapes as bcast_shapes + from arkouda.numpy.util import broadcast_to as bcast_to + + def _isSupported(arg): + return isSupportedNumber(arg) or isSupportedBool(arg) + + def _bool_case(x1, x2): + return type(x1) in (bool, np.bool_) and type(x2) in (bool, np.bool) + + # The function below is needed for the boolean scalar case. + + # np.arctan2 returns float16 if x1 and x2 are both bool. This helper converts it to float64 + # because subsequent functions such as where can't accomodate float16. + + def nparctan2(x1, x2): + return np.float64(np.arctan2(x1, x2)) if _bool_case(x1, x2) else np.arctan2(x1, x2) + + # Begin with various checks. - if not all(isSupportedNumber(arg) or isinstance(arg, pdarray) for arg in [num, denom]): + # First the scalar case. + + if np.isscalar(x1) and np.isscalar(x2): + if out is None: + return nparctan2(x1, x2) if where is None or where is True else np.float64(1.0) + + if out is not None: + if not isinstance(out, pdarray): + raise TypeError("return arrays must be of type pdarray") # matches numpy's error + else: + if where is None or where is True: + out[:] = nparctan2(x1, x2) + return out + else: + try: + _where = where if isinstance(where, pdarray) else _normalize_scalar(where) + new_where = bcast_to(_where, out.shape) + out[:] = ak_where(new_where, nparctan2(x1, x2), out) + return out + except Exception as e: + raise ValueError("Cannot broadcast inputs to common shape in arctan2.") from e + + # That covers the scalar case. From here onward, at least one of x1, x2 is a pdarray. + + if out is None and where is not None: + raise ValueError("In arctan2, 'out' must be specified if 'where' is used.") + + if out is not None and out.dtype != ak_float64: + raise TypeError(f"Cannot return arctan2 result as type {out.dtype}") + + if where is False: + if out is not None: + return out # This is the one instance where you can get away with "where" but not "out" + else: + raise ValueError("In arctan2, 'out' must be specified if 'where' is used.") + + if isinstance(where, pdarray) and where.dtype != bool: + raise TypeError(f"where must have dtype bool, got {where.dtype} instead") + + if np.isscalar(where) and type(where) is not bool: + raise TypeError(f"where must have dtype bool, got {type(where)} instead") + + # At this point, we know we have both out and where. Check for valid types. + + if not all(_isSupported(arg) or isinstance(arg, pdarray) for arg in [x1, x2]): raise TypeError( - f"Unsupported types {type(num)} and/or {type(denom)}. Supported " - "types are numeric scalars and pdarrays. At least one argument must be a pdarray." + f"Unsupported types {type(x1)} and/or {type(x2)}. Supported " + "types are numeric scalars and pdarrays." ) - if isSupportedNumber(num) and isSupportedNumber(denom): + if _isSupported(x1) and _isSupported(x2): raise TypeError( - f"Unsupported types {type(num)} and/or {type(denom)}. Supported " - "types are numeric scalars and pdarrays. At least one argument must be a pdarray." + f"Unsupported types {type(x1)} and/or {type(x2)}. Supported " + "types are numeric scalars and pdarrays." ) - # TODO: handle shape broadcasting for multidimensional arrays - if where is True: - pass - elif where is False: - return num / denom # type: ignore - elif where.dtype != bool: - raise TypeError(f"where must have dtype bool, got {where.dtype} instead") + # Now broadcast as needed. Because each of (x1, x2, where) may be a + # scalar or pdarray, we use assign some temporary variables below so + # that the broadcast will work in any allowed case. + # Note that (1,) is sort of a "dummy shape," broadcastable to anything. - if isinstance(num, pdarray) or isinstance(denom, pdarray): - ndim = num.ndim if isinstance(num, pdarray) else denom.ndim # type: ignore[union-attr] - - # The code below will create the command string for arctan2vv, arctan2vs or arctan2sv, based - # on a and b. - - if isinstance(num, pdarray) and isinstance(denom, pdarray): - cmdstring = f"arctan2vv<{num.dtype},{ndim},{denom.dtype}>" - if where is True: - argdict = { - "a": num, - "b": denom, - } - elif where is False: - return num / denom # type: ignore - else: - argdict = { - "a": num[where], - "b": denom[where], - } - elif not isinstance(denom, pdarray): - ts = resolve_scalar_dtype(denom) + ws = where.shape if isinstance(where, pdarray) else (1,) + x1s = x1.shape if isinstance(x1, pdarray) else (1,) + x2s = x2.shape if isinstance(x2, pdarray) else (1,) + + # the "normalize scalar" call is needed since we support np.bool_, + # which unlike bool, can't be broadcast. + + _where = where if isinstance(where, pdarray) else _normalize_scalar(where) + _x1 = x1 if isinstance(x1, pdarray) else _normalize_scalar(x1) + _x2 = x2 if isinstance(x2, pdarray) else _normalize_scalar(x2) + + # If there is no out parameter, we try to find a common shape. + + common_shape = () + + if out is None: + if where is None: + try: + common_shape = bcast_shapes(x1s, x2s) + _x1 = bcast_to(_x1, common_shape) + _x2 = bcast_to(_x2, common_shape) + except Exception as e: + raise ValueError("Cannot broadcast inputs to common shape in arctan2.") from e + if where is not None: + try: + common_shape = bcast_shapes(x1s, x2s, ws) + _x1 = bcast_to(_x1, common_shape) + _x2 = bcast_to(_x2, common_shape) + _where = bcast_to(where, common_shape) + except Exception as e: + raise ValueError("Cannot broadcast inputs to common shape in arctan2.") from e + + # But if there is an out parameter, we use its shape as the output shape. + + if out is not None: + common_shape = out.shape + try: + _x1 = bcast_to(_x1, common_shape) + _x2 = bcast_to(_x2, common_shape) + if where is not None: + _where = bcast_to(where, common_shape) + except Exception as e: + raise ValueError(f"bcast_to(x2) failed: {type(e).__name__}: {e}") from e + + # Do the computation. I'm keeping this in a separate function for readability. + + tmp = _arctan2_(_x1, _x2) + + # Now handle the "where" parameter as needed. + + if _where is None or _where is True: + if out is not None: + out[:] = tmp + return tmp + else: + if out is None: + raise ValueError("In arctan2, 'out' must be specified if 'where' is used.") + + out[:] = ak_where(_where, tmp, out) + return out + + +def handle_bools(x): + if x.dtype in (bool, np.bool_, ak_bool): + return x.astype(ak_float64) + else: + return x + + +@typechecked +def _arctan2_( + x1: Union[pdarray, numeric_and_bool_scalars], + x2: Union[pdarray, numeric_and_bool_scalars], +) -> pdarray: + from arkouda.client import generic_msg + + if isinstance(x1, pdarray) or isinstance(x2, pdarray): + # These four ifs look awkward, since one of x1, x2 MUST be a pdarray, but since mypy + # doesn't know that, this is how we set ndim and handle bools without causing a mypy error. + + if np.isscalar(x1): + if isinstance(x1, (bool, np.bool_)): + x1 = int(x1) + if np.isscalar(x2): + if isinstance(x2, (bool, np.bool_)): + x1 = int(x2) + + if isinstance(x1, pdarray): + ndim = x1.ndim + x1 = handle_bools(x1) + if isinstance(x2, pdarray): + ndim = x2.ndim + x2 = handle_bools(x2) + + # The code below will create the command string for arctan2vv, arctan2vs + # or arctan2sv, based on x1 and x2. + + argdict = {"a": x1, "b": x2} + if isinstance(x1, pdarray) and isinstance(x2, pdarray): + cmdstring = f"arctan2vv<{x1.dtype},{ndim},{x2.dtype}>" + + elif isinstance(x1, pdarray) and not isinstance(x2, pdarray): + ts = resolve_scalar_dtype(x2) if ts in ["float64", "int64", "uint64", "bool"]: - cmdstring = "arctan2vs_" + ts + f"<{num.dtype},{ndim}>" # type: ignore[union-attr] + cmdstring = "arctan2vs_" + ts + f"<{x1.dtype},{ndim}>" else: - raise TypeError(f"{ts} is not an allowed denom type for arctan2") - argdict = {"a": num if where is True else num[where], "b": denom} # type: ignore - elif not isinstance(num, pdarray): - ts = resolve_scalar_dtype(num) + raise TypeError(f"{ts} is not an allowed x2 type for arctan2") + + elif isinstance(x2, pdarray) and not isinstance(x1, pdarray): + ts = resolve_scalar_dtype(x1) if ts in ["float64", "int64", "uint64", "bool"]: - cmdstring = "arctan2sv_" + ts + f"<{denom.dtype},{ndim}>" + cmdstring = "arctan2sv_" + ts + f"<{x2.dtype},{ndim}>" else: - raise TypeError(f"{ts} is not an allowed num type for arctan2") - argdict = {"a": num, "b": denom if where is True else denom[where]} # type: ignore + raise TypeError(f"{ts} is not an allowed x1 type for arctan2") repMsg = generic_msg(cmd=cmdstring, args=argdict) - ret = create_pdarray(repMsg) - if where is True: - return ret - else: - new_pda = num / denom # type : ignore - return _merge_where(new_pda, where, ret) + return create_pdarray(repMsg) - else: - return scalar_array(arctan2(num, denom) if where else num / denom) + raise TypeError("_arctan2_ helper function called with no pdarray arguments.") @typechecked diff --git a/arkouda/numpy/pdarrayclass.py b/arkouda/numpy/pdarrayclass.py index 6c4b7750ff9..d2947b14345 100644 --- a/arkouda/numpy/pdarrayclass.py +++ b/arkouda/numpy/pdarrayclass.py @@ -2,6 +2,7 @@ import builtins import json +import warnings from functools import reduce from math import ceil @@ -36,6 +37,7 @@ from arkouda.numpy.dtypes import int64 as akint64 from arkouda.numpy.dtypes import str_ as akstr_ from arkouda.numpy.dtypes import uint64 as akuint64 +from arkouda.numpy.err import geterr as ak_geterr module = modules[__name__] @@ -222,6 +224,28 @@ def stats_reduce( SUPPORTED_STATS_REDUCTION_OPS = ["var", "std"] +def _dbz_check(op, flag): + if op not in ["/", "//"] or not flag: + return + + mode = ak_geterr()["divide"] + + msg = ( + "divide by zero encountered in divide" + if op == "/" + else "divide by zero encountered in floor_divide" + ) + + if mode == "ignore": + return + + elif mode == "warn": + warnings.warn(msg, RuntimeWarning) + + elif mode == "raise": + raise FloatingPointError(msg) + + def _axis_parser(axis): if axis is None: return None @@ -855,6 +879,10 @@ def _binop(self, other: Union[pdarray, numeric_scalars], op: str) -> pdarray: cmd=f"binopvv<{self.dtype},{other.dtype},{x1.ndim}>", args={"op": op, "a": x1, "b": x2}, ) + + # Create a runtime warning if this caused divide-by-zero. + _dbz_check(op, any(x2 == 0)) + if tmp_x1: del x1 if tmp_x2: @@ -878,6 +906,10 @@ def _binop(self, other: Union[pdarray, numeric_scalars], op: str) -> pdarray: cmd=f"binopvs<{self.dtype},{dt},{self.ndim}>", args={"op": op, "a": self, "value": other}, ) + + # Create a runtime warning if this caused divide-by-zero. + _dbz_check(op, other == 0) + return create_pdarray(repMsg) # reverse binary operators @@ -929,6 +961,10 @@ def _r_binop(self, other: Union[pdarray, numeric_scalars], op: str) -> pdarray: cmd=f"binopsv<{self.dtype},{dt},{self.ndim}>", args={"op": op, "a": self, "value": other}, ) + + # Create a runtime warning if this caused divide-by-zero. + _dbz_check(op, any(self == 0)) + return create_pdarray(repMsg) def transfer(self, hostname: str, port: int_scalars): diff --git a/tests/numpy/datetime_test.py b/tests/numpy/datetime_test.py index c7d1a7c06b4..6d36f95bce9 100644 --- a/tests/numpy/datetime_test.py +++ b/tests/numpy/datetime_test.py @@ -4,6 +4,7 @@ import arkouda as ak +from arkouda.numpy import err as akerr from arkouda.numpy import timeclass @@ -148,10 +149,20 @@ def test_op_types(self, verbose=pytest.verbose): eval(f"fcvec {op} scsca") metrics["ak_not_supported"] += 1 else: + oldmodes = akerr.geterr() + akerr.seterr(divide="ignore") compare_flag = True ret = eval(f"fcvec {op} scvec") assert isinstance(ret, return_type) metrics["ak_supported"] += 1 + # try: + # compare_flag = True + # ret = eval(f"fcvec {op} scvec") + # assert isinstance(ret, return_type) + # metrics["ak_supported"] += 1 + # except RuntimeWarning: + # continue # this test can cause divide-by-zero, which would + # # also cause a following test to fail, so skip it try: pdret = eval(f"pdfcvec {op} pdscvec") except TypeError: @@ -161,6 +172,7 @@ def test_op_types(self, verbose=pytest.verbose): ) metrics["ak_yes_pd_no"] += 1 compare_flag = False + akerr.seterr(**oldmodes) if compare_flag: # Arkouda currently does not handle NaT, so replace with zero if pdret.dtype.kind == "m": diff --git a/tests/numpy/numeric_test.py b/tests/numpy/numeric_test.py index adf2a8ed4eb..272c6565679 100644 --- a/tests/numpy/numeric_test.py +++ b/tests/numpy/numeric_test.py @@ -54,6 +54,15 @@ (ak.bool_, ak.bool_), ] +# WHERE_OUT is used in test_arctan2 + +WHERE_OUT = [ + pytest.param(False, False, id="whereF_outF"), + pytest.param(False, True, id="whereF_outT"), + pytest.param(True, True, id="whereT_outT"), +] + + # There are many ways to create a vector of alternating values. # This is a fairly fast and fairly straightforward approach. @@ -64,6 +73,19 @@ def alternate(L, R, n): return v +def random_scalar(dtype): + if dtype is np.float64: + return np.float64(np.random.random()) + elif dtype is np.int64: + return np.random.randint(-100, 100, dtype=np.int64) + elif dtype is np.uint64: + return np.random.randint(0, 100, dtype=np.uint64) + elif dtype is np.bool_: + return np.random.randint(2, dtype=np.bool_) + else: + raise TypeError(f"Unsupported dtype: {dtype}") + + # The following tuples support a simplification of the trigonometric # and hyperbolic testing. @@ -118,6 +140,7 @@ def alternate(L, R, n): np.array([np.nan, -np.inf, -1.0, 1.0, np.inf]), ] ), + ak.bool_: alternate(True, False, 10), ak.uint64: np.arange(2**64 - 10, 2**64, dtype=np.uint64), } @@ -506,110 +529,223 @@ def test_trig_and_hyp(self, num_type): with pytest.raises(TypeError): akfunc(np.array([range(0, 10)]).astype(num_type)) - @pytest.mark.parametrize("num_type", NO_BOOL) - @pytest.mark.parametrize("denom_type", NO_BOOL) - def test_arctan2(self, num_type, denom_type): - np.random.seed(pytest.seed) - na_num = np.random.permutation(NP_TRIG_ARRAYS[num_type]) - na_denom = np.random.permutation(DENOM_ARCTAN2_ARRAYS[denom_type]) + @pytest.mark.parametrize("num_type", NUMERIC_TYPES) + @pytest.mark.parametrize("denom_type", NUMERIC_TYPES) + @pytest.mark.parametrize("num_scalar", [True, False]) + @pytest.mark.parametrize("denom_scalar", [True, False]) + @pytest.mark.parametrize("uses_where, uses_out", WHERE_OUT) + def test_arctan2(self, num_type, denom_type, num_scalar, denom_scalar, uses_out, uses_where): + from arkouda.numpy import cast as akcast - pda_num = ak.array(na_num, dtype=num_type) - pda_denom = ak.array(na_denom, dtype=denom_type) + np.random.seed(pytest.seed) - truth_np = alternate(True, False, len(na_num)) - truth_ak = ak.array(truth_np) + # Create numerator and denominator - assert np.allclose( - np.arctan2(na_num, na_denom, where=True), - ak.arctan2(pda_num, pda_denom, where=True).to_ndarray(), - equal_nan=True, + na_num = ( + random_scalar(num_type) if num_scalar else np.random.permutation(NP_TRIG_ARRAYS[num_type]) ) - - assert np.allclose( - np.arctan2(na_num[0], na_denom, where=True), - ak.arctan2(pda_num[0], pda_denom, where=True).to_ndarray(), - equal_nan=True, - ) - assert np.allclose( - np.arctan2(na_num, na_denom[0], where=True), - ak.arctan2(pda_num, pda_denom[0], where=True).to_ndarray(), - equal_nan=True, + na_denom = ( + random_scalar(denom_type) + if denom_scalar + else np.random.permutation(DENOM_ARCTAN2_ARRAYS[denom_type]) ) + pda_num = na_num if num_scalar else ak.array(na_num, dtype=num_type) + pda_denom = na_denom if denom_scalar else ak.array(na_denom, dtype=denom_type) - assert np.allclose( - na_num / na_denom, - ak.arctan2(pda_num, pda_denom, where=False).tolist(), - equal_nan=True, - ) - assert np.allclose( - na_num[0] / na_denom, - ak.arctan2(pda_num[0], pda_denom, where=False).to_ndarray(), - equal_nan=True, - ) - assert np.allclose( - na_num / na_denom[0], - ak.arctan2(pda_num, pda_denom[0], where=False).to_ndarray(), - equal_nan=True, - ) + # Handle the bool-bool case. - assert np.allclose( - [ - (np.arctan2(na_num[i], na_denom[i]) if truth_np[i] else na_num[i] / na_denom[i]) - for i in range(len(na_num)) - ], - ak.arctan2(pda_num, pda_denom, where=truth_ak).to_ndarray(), - equal_nan=True, - ) - assert np.allclose( - [ - (np.arctan2(na_num[0], na_denom[i]) if truth_np[i] else na_num[0] / na_denom[i]) - for i in range(len(na_denom)) - ], - ak.arctan2(pda_num[0], pda_denom, where=truth_ak).to_ndarray(), - equal_nan=True, - ) - assert np.allclose( - [ - (np.arctan2(na_num[i], na_denom[0]) if truth_np[i] else na_num[i] / na_denom[0]) - for i in range(len(na_num)) - ], - ak.arctan2(pda_num, pda_denom[0], where=truth_ak).to_ndarray(), - equal_nan=True, - ) + # This step is needed because of the comparison to numpy. Numpy's + # arctan2, when operating on bool types, returns np.float16. This + # causes a unit test issue: + # - because the accuracy/precision is much less than float64, + # the tolerances have to be adjusted for comparisons. - # Edge cases: infinities and zeros. Doesn't use _infinity_edge_case_helper - # because arctan2 needs two numbers (numerator and denominator) rather than one. + bool_case = (num_type in (bool, np.bool_)) and (denom_type in (bool, np.bool_)) - na1 = np.array([np.inf, -np.inf]) - pda1 = ak.array(na1) - na2 = np.array([1, 10]) - pda2 = ak.array(na2) + (rtol, atol) = (1.0e-3, 1.0e-3) if bool_case else (1.0e-5, 1.0e-08) - assert np.allclose( - np.arctan2(na1, na2), - ak.arctan2(pda1, pda2).to_ndarray(), - equal_nan=True, - ) + def recast_if_bool_case(x): + if np.isscalar(x): + return np.float64(x) + elif isinstance(x, np.ndarray): + return x.astype(np.float64) + elif isinstance(x, ak.pdarray): + return akcast(x, ak.float64) + else: + raise ValueError("recast_if_bool_case only handles scalars, ndarrays, and pdarrays") - assert np.allclose( - np.arctan2(na2, na1), - ak.arctan2(pda2, pda1).to_ndarray(), - equal_nan=True, - ) - assert np.allclose(np.arctan2(na1, 5), ak.arctan2(pda1, 5).to_ndarray(), equal_nan=True) - assert np.allclose(np.arctan2(5, na1), ak.arctan2(5, pda1).to_ndarray(), equal_nan=True) - assert np.allclose(np.arctan2(na1, 0), ak.arctan2(pda1, 0).to_ndarray(), equal_nan=True) - assert np.allclose(np.arctan2(0, na1), ak.arctan2(0, pda1).to_ndarray(), equal_nan=True) + # Many cases to test. Begin with both scalar. + # This test also covers broadcasting, because we supply "out" as a 1D pdarray. - with pytest.raises(TypeError): - ak.arctan2( - np.array([range(0, 10)]).astype(num_type), - np.array([range(10, 20)]).astype(num_type), - ) - with pytest.raises(TypeError): - ak.arctan2(pda_num[0], np.array([range(10, 20)]).astype(num_type)) - with pytest.raises(TypeError): - ak.arctan2(np.array([range(0, 10)]).astype(num_type), pda_denom[0]) + if num_scalar and denom_scalar: + if not uses_out: + assert np.arctan2(na_num, na_denom) == ak.arctan2(pda_num, pda_denom) + else: + na_out = np.array([1.0, 1.0, 1.0, 1.0]) + na_outc = na_out.copy() + pda_out = ak.array(na_out) + pda_outc = pda_out.copy() + if not uses_where: + na_res = np.arctan2(na_num, na_denom, na_outc) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + else: + na_res = np.arctan2(na_num, na_denom, na_outc, where=True) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=True) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + na_outc = na_out.copy() + pda_outc = pda_out.copy() + na_res = np.arctan2(na_num, na_denom, na_outc, where=False) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=False) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + + # Now the cases where one is scalar and the other is a vector. + + if num_scalar ^ denom_scalar: + if not uses_out: + na_res = np.arctan2(na_num, na_denom) + na_res = recast_if_bool_case(na_res) + ak_assert_almost_equivalent(na_res, ak.arctan2(pda_num, pda_denom), rtol, atol) + else: + na_out = np.ones(len(na_denom)) if num_scalar else np.ones(len(na_num)) + na_outc = na_out.copy() + pda_out = ak.array(na_out) + pda_outc = pda_out.copy() + if not uses_where: + na_res = np.arctan2(na_num, na_denom, na_outc) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + else: + na_where = ( + alternate(True, False, len(na_denom)) + if num_scalar + else alternate(True, False, len(na_num)) + ) + pda_where = ak.array(na_where) + na_res = np.arctan2(na_num, na_denom, na_outc, where=True) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=True) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + na_outc = na_out.copy() + pda_outc = pda_out.copy() + na_res = np.arctan2(na_num, na_denom, na_outc, where=False) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=False) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + na_outc = na_out.copy() + pda_outc = pda_out.copy() + na_res = np.arctan2(na_num, na_denom, na_outc, where=na_where) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=pda_where) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + + # Finally the case where both are vectors. + # This test covers broadcasting as well, provided 2 is in the compiled ranks. + # Because in that case, we create out as a 2D array of size 2 x len(the 1D data vectors). + + if not (num_scalar or denom_scalar): + if not uses_out: + na_res = np.arctan2(na_num, na_denom) + na_res = recast_if_bool_case(na_res) + ak_assert_almost_equivalent(na_res, ak.arctan2(pda_num, pda_denom), rtol, atol) + else: + na_out = ( + np.ones(len(na_denom)) + if not (2 in get_array_ranks()) + else np.ones((2, len(na_denom))) + ) + na_outc = na_out.copy() + pda_out = ak.array(na_out) + pda_outc = pda_out.copy() + if not uses_where: + na_res = np.arctan2(na_num, na_denom, na_outc) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + else: + na_where = alternate(True, False, len(na_denom)) + pda_where = ak.array(na_where) + na_res = np.arctan2(na_num, na_denom, na_outc, where=True) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=True) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + na_outc = na_out.copy() + pda_outc = pda_out.copy() + na_res = np.arctan2(na_num, na_denom, na_outc, where=False) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=False) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + na_outc = na_out.copy() + pda_outc = pda_out.copy() + na_res = np.arctan2(na_num, na_denom, na_outc, where=na_where) + na_res = recast_if_bool_case(na_res) + na_outc = recast_if_bool_case(na_outc) + pda_res = ak.arctan2(pda_num, pda_denom, pda_outc, where=pda_where) + ak.assert_almost_equivalent(na_res, pda_res, rtol, atol) + ak.assert_almost_equivalent(na_outc, pda_outc, rtol, atol) + + # Edge cases: infinities and zeros. Doesn't use _infinity_edge_case_helper + # because arctan2 needs two numbers (numerator and denominator) rather than one. + + # This really doesn't need to be done for every case of the inputs. May pull + # it into a separate test. + + if ( + num_type is ak.float64 + and denom_type is ak.float64 + and num_scalar is False + and denom_scalar is False + and uses_where is False + and uses_out is False + ): + na1 = np.array([np.inf, -np.inf]) + pda1 = ak.array(na1) + na2 = np.array([1, 10]) + pda2 = ak.array(na2) + pda_out = ak.ones(len(pda1)) + + ak_assert_almost_equivalent(np.arctan2(na1, na2), ak.arctan2(pda1, pda2)) + ak_assert_almost_equivalent(np.arctan2(na2, na1), ak.arctan2(pda2, pda1)) + ak_assert_almost_equivalent(np.arctan2(na1, 5), ak.arctan2(pda1, 5)) + ak_assert_almost_equivalent(np.arctan2(5, na1), ak.arctan2(5, pda1)) + ak_assert_almost_equivalent(np.arctan2(na1, 0), ak.arctan2(pda1, 0)) + ak_assert_almost_equivalent(np.arctan2(0, na1), ak.arctan2(0, pda1)) + + with pytest.raises(TypeError): + ak.arctan2( + np.array([range(0, 10)]).astype(num_type), + np.array([range(10, 20)]).astype(num_type), + ) + with pytest.raises(TypeError): + ak.arctan2(pda_num[0], np.array([range(10, 20)]).astype(num_type)) + with pytest.raises(TypeError): + ak.arctan2(np.array([range(0, 10)]).astype(num_type), pda_denom[0]) + with pytest.raises(TypeError): + ak.arctan2(pda1, pda2, pda_out, where=pda1) @pytest.mark.parametrize("num_type", NO_BOOL) def test_rad2deg(self, num_type):