From 6a34c842d92b219dd5a410ce2bcc23618fb3ee3a Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 12 Dec 2024 14:13:12 +0000 Subject: [PATCH 01/29] Implement tests for cosmo_quantity. --- tests/test_cosmo_array.py | 225 +++++++++++++++++++++++++++++++++++++- 1 file changed, 220 insertions(+), 5 deletions(-) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 4cb83535..5cc6aebc 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -2,15 +2,52 @@ Tests the initialisation of a cosmo_array. """ +import pytest +import os import numpy as np import unyt as u -from swiftsimio.objects import cosmo_array, cosmo_factor +from swiftsimio.objects import cosmo_array, cosmo_factor, a + +savetxt_file = "saved_array.txt" + + +# dummy until actually implemented: +class cosmo_quantity: + pass + + +def getfunc(fname): + func = np + for attr in fname.split("."): + func = getattr(func, attr) + return func + + +def ca(x, unit=u.Mpc): + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) + + +def to_ua(x): + return u.unyt_array(x.to_value(x.units), x.units) if hasattr(x, "comoving") else x + + +def check_result(x_c, x_u): + # careful, unyt_quantity subclass of unyt_array + if isinstance(x_u, u.unyt_quantity): + assert isinstance(x_c, cosmo_quantity) + elif isinstance(x_u, u.unyt_array): + assert isinstance(x_c, cosmo_array) and not isinstance(x_c, cosmo_quantity) + else: + assert not isinstance(x_c, cosmo_array) class TestCosmoArrayInit: def test_init_from_ndarray(self): arr = cosmo_array( - np.ones(5), units=u.Mpc, cosmo_factor=cosmo_factor("a^1", 1), comoving=False + np.ones(5), + units=u.Mpc, + cosmo_factor=cosmo_factor(a ** 1, 1), + comoving=False, ) assert hasattr(arr, "cosmo_factor") assert hasattr(arr, "comoving") @@ -20,7 +57,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor("a^1", 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -30,7 +67,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor("a^1", 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -40,9 +77,187 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor("a^1", 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") assert hasattr(arr, "comoving") assert isinstance(arr, cosmo_array) + + +class TestNumpyFunctions: + """ + Functions specially handled by unyt risk silently casting to unyt_array/unyt_quantity. + """ + + def test_handled_funcs(self): + from unyt._array_functions import _HANDLED_FUNCTIONS + + functions_to_check = { + "dot": (ca(np.arange(3)), ca(np.arange(3))), + "vdot": (ca(np.arange(3)), ca(np.arange(3))), + "inner": (ca(np.arange(3)), ca(np.arange(3))), + "outer": (ca(np.arange(3)), ca(np.arange(3))), + "kron": (ca(np.arange(3)), ca(np.arange(3))), + "histogram_bin_edges": (ca(np.arange(3)),), + "concatenate": (ca(np.eye(3)),), + "cross": (ca(np.arange(3)), ca(np.arange(3))), + "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), + "union1d": (ca(np.arange(3)), ca(np.arange(3))), + "vstack": (ca(np.arange(3)),), + "hstack": (ca(np.arange(3)),), + "dstack": (ca(np.arange(3)),), + "column_stack": (ca(np.arange(3)),), + "stack": (ca(np.arange(3)),), + "around": (ca(np.arange(3)),), + "block": (ca(np.arange(3)),), + "sort_complex": (ca(np.arange(3)),), + "linspace": (ca(1), ca(2), 10), + "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), + "geomspace": (ca(1), ca(1), 10), + "prod": (ca(np.arange(3)),), + "var": (ca(np.arange(3)),), + "trace": (ca(np.eye(3)),), + "percentile": (ca(np.arange(3)), 30), + "quantile": (ca(np.arange(3)), 0.3), + "nanpercentile": (ca(np.arange(3)), 30), + "nanquantile": (ca(np.arange(3)), 0.3), + "diff": (ca(np.arange(3)),), + "ediff1d": (ca(np.arange(3)),), + "ptp": (ca(np.arange(3)),), + "cumprod": (ca(np.arange(3)),), + "pad": (ca(np.arange(3)), 3), + "choose": (np.arange(3), ca(np.eye(3))), + "insert": (ca(np.arange(3)), 1, ca(1)), + "isin": (ca(np.arange(3)), ca(np.arange(3))), + "in1d": (ca(np.arange(3)), ca(np.arange(3))), + "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "isclose": (ca(np.arange(3)), ca(np.arange(3))), + "allclose": (ca(np.arange(3)), ca(np.arange(3))), + "array_equal": (ca(np.arange(3)), ca(np.arange(3))), + "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), + "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), + "select": ( + [np.arange(3) < 1, np.arange(3) > 1], + [ca(np.arange(3)), ca(np.arange(3))], + ca(1), + ), + "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), + "sinc": (ca(np.arange(3)),), + "clip": (ca(np.arange(3)), ca(1), ca(2)), + "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "triu": (ca(np.arange(3)),), + "tril": (ca(np.arange(3)),), + "einsum": ("ii->i", ca(np.eye(3))), + "convolve": (ca(np.arange(3)), ca(np.arange(3))), + "correlate": (ca(np.arange(3)), ca(np.arange(3))), + "tensordot": (ca(np.eye(3)), ca(np.eye(3))), + "unwrap": (ca(np.arange(3)),), + "histogram": (ca(np.arange(3)),), + "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), + "histogramdd": (ca(np.arange(3)),), + "linalg.norm": (ca(np.arange(3)),), + "linalg.inv": (ca(np.eye(3)),), + "linalg.tensorinv": (ca(np.eye(9).reshape((3, 3, 3, 3))),), + "linalg.pinv": (ca(np.eye(3)),), + "linalg.det": (ca(np.eye(3)),), + "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.tensorsolve": ( + ca(np.eye(24).reshape((6, 4, 2, 3, 4))), + ca(np.ones((6, 4))), + ), + "linalg.eigvals": (ca(np.eye(3)),), + "linalg.eigvalsh": (ca(np.eye(3)),), + "linalg.svd": (ca(np.eye(3)),), + "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.eig": (ca(np.eye(3)),), + "linalg.eigh": (ca(np.eye(3)),), + "fft.fft": (ca(np.arange(3)),), + "fft.fft2": (ca(np.eye(3)),), + "fft.fftn": (ca(np.arange(3)),), + "fft.hfft": (ca(np.arange(3)),), + "fft.rfft": (ca(np.arange(3)),), + "fft.rfft2": (ca(np.eye(3)),), + "fft.rfftn": (ca(np.arange(3)),), + "fft.ifft": (ca(np.arange(3)),), + "fft.ifft2": (ca(np.eye(3)),), + "fft.ifftn": (ca(np.arange(3)),), + "fft.ihfft": (ca(np.arange(3)),), + "fft.irfft": (ca(np.arange(3)),), + "fft.irfft2": (ca(np.eye(3)),), + "fft.irfftn": (ca(np.arange(3)),), + "fft.fftshift": (ca(np.arange(3)),), + "fft.ifftshift": (ca(np.arange(3)),), + "array2string": (ca(np.arange(3)),), + "copyto": (ca(np.arange(3)), ca(np.arange(3))), + "savetxt": (savetxt_file, ca(np.arange(3))), + "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), + "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), + "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), + "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), + "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + "array_repr": (ca(np.arange(3)),), + "trapezoid": (ca(np.arange(3)),), + } + functions_checked = list() + bad_funcs = dict() + for fname, args in functions_to_check.items(): + ua_args = tuple(to_ua(arg) for arg in args) + func = getfunc(fname) + try: + ua_result = func(*ua_args) + except u.exceptions.UnytError: + raises_unyt_error = True + else: + raises_unyt_error = False + if "savetxt" in fname and os.path.isfile(savetxt_file): + os.remove(savetxt_file) + functions_checked.append(func) + if raises_unyt_error: + with pytest.raises(u.exceptions.UnytError): + result = func(*args) + continue + result = func(*args) + if "savetxt" in fname and os.path.isfile(savetxt_file): + os.remove(savetxt_file) + if ua_result is None: + try: + assert result is None + except AssertionError: + bad_funcs["np." + fname] = result, ua_result + try: + if isinstance(ua_result, tuple): + assert isinstance(result, tuple) + assert len(result) == len(ua_result) + for r, ua_r in zip(result, ua_result): + check_result(r, ua_r) + else: + check_result(result, ua_result) + except AssertionError: + bad_funcs["np." + fname] = result, ua_result + if len(bad_funcs) > 0: + raise AssertionError( + "Some functions did not return expected types " + "(obtained, obtained with unyt input): " + str(bad_funcs) + ) + unchecked_functions = [ + f for f in _HANDLED_FUNCTIONS if f not in functions_checked + ] + try: + assert len(unchecked_functions) == 0 + except AssertionError: + raise AssertionError( + "Did not check functions", [f.__name__ for f in unchecked_functions] + ) + + def test_getitem(self): + assert isinstance(ca(np.arange(3))[0], cosmo_quantity) + + def test_reshape_to_scalar(self): + assert isinstance(ca(np.ones(1)).reshape(tuple()), cosmo_quantity) + + def test_iter(self): + for cq in ca(np.arange(3)): + assert isinstance(cq, cosmo_quantity) From 07aac86928650bf4957a529374a7e17ab04d8e44 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 12 Dec 2024 14:40:29 +0000 Subject: [PATCH 02/29] Initial cosmo_quantity implementation. --- swiftsimio/objects.py | 81 +++++++++++++++++++++++++++++++++++++++++-- 1 file changed, 78 insertions(+), 3 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 7ddaf140..bd5d0d6e 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -6,8 +6,9 @@ import warnings import unyt -from unyt import unyt_array +from unyt import unyt_array, unyt_quantity from unyt.array import multiple_output_operators +from numbers import Number as numeric_type try: from unyt.array import POWER_MAPPING @@ -599,8 +600,6 @@ class cosmo_array(unyt_array): This inherits from the unyt.unyt_array, and adds four variables: compression, cosmo_factor, comoving, and valid_transform. - Data is assumed to be comoving when passed to the object but you - can override this by setting the latter flag to be False. Parameters ---------- @@ -1168,3 +1167,79 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): o.compression = ret_comp return ret + + +class cosmo_quantity(cosmo_array, unyt_quantity): + """ + Cosmology scalar class. + + This inherits from both the cosmo_array and the unyt.unyt_array, and has the same four + attributes as cosmo_array: compression, cosmo_factor, comoving, and valid_transform. + + Parameters + ---------- + + cosmo_array : cosmo_array + the inherited cosmo_array + + unyt_quantity : unyt.unyt_quantity + the inherited unyt_quantity + + Attributes + ---------- + + comoving : bool + if True then the array is in comoving co-ordinates, and if + False then it is in physical units. + + cosmo_factor : float + Object to store conversion data between comoving and physical coordinates + + compression : string + String describing any compression that was applied to this array in the + hdf5 file. + + valid_transform: bool + if True then the array can be converted from physical to comoving units + + """ + + def __new__( + cls, + input_scalar, + units=None, + registry=None, + dtype=None, + bypass_validation=False, + name=None, + cosmo_factor=None, + comoving=None, + valid_transform=True, + compression=None, + ): + input_units = units + if not ( + bypass_validation + or isinstance(input_scalar, (numeric_type, np.number, np.ndarray)) + ): + raise RuntimeError("unyt_quantity values must be numeric") + if input_units is None: + units = getattr(input_scalar, "units", None) + else: + units = input_units + ret = cosmo_array.__new__( + cls, + np.asarray(input_scalar), + units, + registry, + dtype=dtype, + bypass_validation=bypass_validation, + name=name, + cosmo_factor=cosmo_factor, + comoving=comoving, + valid_transform=valid_transform, + compression=compression, + ) + if ret.size > 1: + raise RuntimeError("unyt_quantity instances must be scalars") + return ret From 6eede3374fd929d19ccf929c816812931c01bfb7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 12 Dec 2024 19:54:25 +0000 Subject: [PATCH 03/29] First rough working implementation of a numpy func overload. --- swiftsimio/__init__.py | 1 + swiftsimio/_array_functions.py | 95 ++++++++++++++++++++++++++++++++++ swiftsimio/objects.py | 79 +++++++++++++++++++++++++++- 3 files changed, 174 insertions(+), 1 deletion(-) create mode 100644 swiftsimio/_array_functions.py diff --git a/swiftsimio/__init__.py b/swiftsimio/__init__.py index f9fa573e..ceb146f0 100644 --- a/swiftsimio/__init__.py +++ b/swiftsimio/__init__.py @@ -8,6 +8,7 @@ import swiftsimio.metadata as metadata import swiftsimio.accelerated as accelerated import swiftsimio.objects as objects +from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a import swiftsimio.visualisation as visualisation import swiftsimio.units as units import swiftsimio.subset_writer as subset_writer diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py new file mode 100644 index 00000000..f02195b5 --- /dev/null +++ b/swiftsimio/_array_functions.py @@ -0,0 +1,95 @@ +import numpy as np +from unyt import unyt_array, unyt_quantity +from unyt._array_functions import implements +from unyt._array_functions import _HANDLED_FUNCTIONS as _UNYT_HANDLED_FUNCTIONS +from .objects import _multiply_cosmo_factor, cosmo_array, cosmo_quantity + +_HANDLED_FUNCTIONS = dict() + + +def arg_helper(*args, **kwargs): + cms = [(hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args] + cfs = [ + (hasattr(arg, "cosmo_factor"), getattr(arg, "cosmo_factor", None)) + for arg in args + ] + comps = [ + (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args + ] + kw_cms = [ + (hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args + ] + kw_cfs = [ + (hasattr(arg, "cosmo_factor"), getattr(arg, "cosmo_factor", None)) + for arg in args + ] + kw_comps = [ + (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args + ] + if all([cm[1] for cm in cms + kw_cms if cm[0]]): + # all cosmo inputs are comoving + ret_cm = True + elif all([cm[1] is None for cm in cms + kw_cms if cm[0]]): + # all cosmo inputs have comoving=None + ret_cm = None + elif any([cm[1] is None for cm in cms + kw_cms if cm[0]]): + # only some cosmo inputs have comoving=None + raise ValueError( + "Some arguments have comoving=None and others have comoving=True|False. " + "Result is undefined!" + ) + elif not any([cm[1] for cm in cms + kw_cms if cm[0]]): + # all cosmo_array inputs are physical + ret_cm = False + else: + # mix of comoving and physical inputs + args = [ + arg.to_comoving() if cm[0] and not cm[1] else arg + for arg, cm in zip(args, cms) + ] + kwargs = [ + kwarg.to_comoving() if cm[0] and not cm[1] else kwarg + for kwarg, cm in zip(kwargs, cms) + ] + ret_cm = True + if len(set(comps + kw_comps)) == 1: + # all compressions identical, preserve it + ret_comp = comps[0] + else: + # mixed compressions, strip it off + ret_comp = None + return args, cfs, kwargs, kw_cfs, ret_cm, ret_comp + + +@implements(np.dot) +def dot(a, b, out=None): + from unyt._array_functions import dot as unyt_dot + + args, cfs, kwargs, kw_cfs, ret_cm, ret_comp = arg_helper(a, b, out=out) + ret_cf = ( + cfs[0][1] * cfs[1][1] + ) # this needs helper function handling different cases + res = unyt_dot(*args, **kwargs) + if out is None: + # should we be using __array_wrap__ somehow? + if isinstance(res, unyt_quantity): + return cosmo_quantity( + res, comoving=ret_cm, cosmo_factor=ret_cf, compression=ret_comp + ) + return cosmo_array( + res, comoving=ret_cm, cosmo_factor=ret_cf, compression=ret_comp + ) + if hasattr(out, "comoving"): + out.comoving = ret_cm + if hasattr(out, "cosmo_factor"): + out.cosmo_factor = ret_cf + if hasattr(out, "compression"): + out.compression = ret_comp + return cosmo_array( # confused, do we set out, or return? + res.to_value(res.units), + res.units, + bypass_validation=True, + comoving=ret_cm, + cosmo_factor=ret_cf, + compression=ret_comp, + ) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 8492b2b4..921dca30 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -454,7 +454,6 @@ def __init__(self, expr, scale_factor): expr : sympy.expr expression used to convert between comoving and physical coordinates - scale_factor : float the scale factor of the simulation data """ @@ -1101,6 +1100,15 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): if all([cm[1] for cm in cms if cm[0]]): # all cosmo_array inputs are comoving ret_cm = True + elif all([cm[1] is None for cm in cms if cm[0]]): + # all cosmo inputs have comoving=None + ret_cm = None + elif any([cm[1] is None for cm in cms if cm[0]]): + # only some cosmo inputs have comoving=None + raise ValueError( + "Some arguments have comoving=None and others have comoving=True|False. " + "Result is undefined!" + ) elif not any([cm[1] for cm in cms if cm[0]]): # all cosmo_array inputs are physical ret_cm = False @@ -1171,6 +1179,38 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): return ret + def __array_function__(self, func, types, args, kwargs): + # Follow NEP 18 guidelines + # https://numpy.org/neps/nep-0018-array-function-protocol.html + from ._array_functions import _HANDLED_FUNCTIONS + from unyt._array_functions import ( + _HANDLED_FUNCTIONS as _UNYT_HANDLED_FUNCTIONS, + _UNSUPPORTED_FUNCTIONS as _UNYT_UNSUPPORTED_FUNCTIONS, + ) + + # Let's claim to support everything supported by unyt. + # If we can't do this in future, follow their pattern of + # defining out own _UNSUPPORTED_FUNCTIONS in a _array_functions.py file + _UNSUPPORTED_FUNCTIONS = _UNYT_UNSUPPORTED_FUNCTIONS + + if func in _UNSUPPORTED_FUNCTIONS: + # following NEP 18, return NotImplemented as a sentinel value + # which will lead to raising a TypeError, while + # leaving other arguments a chance to take the lead + return NotImplemented + + if func not in _HANDLED_FUNCTIONS and func in _UNYT_HANDLED_FUNCTIONS: + # first look for unyt's implementation + return _UNYT_HANDLED_FUNCTIONS[func](*args, **kwargs) + elif func not in _UNYT_HANDLED_FUNCTIONS: + # otherwise default to numpy's private implementation + return func._implementation(*args, **kwargs) + # Note: this allows subclasses that don't override + # __array_function__ to handle unyt_array objects + if not all(issubclass(t, cosmo_array) or t is np.ndarray for t in types): + return NotImplemented + return _HANDLED_FUNCTIONS[func](*args, **kwargs) + class cosmo_quantity(cosmo_array, unyt_quantity): """ @@ -1220,6 +1260,43 @@ def __new__( valid_transform=True, compression=None, ): + """ + Essentially a copy of the unyt_quantity.__new__ constructor. + + Parameters + ---------- + input_scalar : an integer of floating point scalar + A scalar to attach units and cosmological transofrmations to. + units : str, unyt.unit_symbols or astropy.unit, optional + The units of the array. Powers must be specified using python syntax + (cm**3, not cm^3). + registry : unyt.unit_registry.UnitRegistry, optional + The registry to create units from. If input_units is already associated with a + unit registry and this is specified, this will be used instead of the registry + associated with the unit object. + dtype : np.dtype or str, optional + The dtype of the array data. Defaults to the dtype of the input data, or, if + none is found, uses np.float64 + bypass_validation : bool, optional + If True, all input validation is skipped. Using this option may produce + corrupted, invalid units or array data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, input_units + must be a valid unit object. Defaults to False. + name : str, optional + The name of the array. Defaults to None. This attribute does not propagate + through mathematical operations, but is preserved under indexing and unit + conversions. + cosmo_factor : cosmo_factor + cosmo_factor object to store conversion data between comoving and physical + coordinates. + comoving : bool + Flag to indicate whether using comoving coordinates. + valid_transform : bool + Flag to indicate whether this array can be converted to comoving. + compression : string + Description of the compression filters that were applied to that array in the + hdf5 file. + """ input_units = units if not ( bypass_validation From 32f9a2e73580ea3255bfb6cb3218ec796ce5a885 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 13 Dec 2024 00:26:00 +0000 Subject: [PATCH 04/29] Start trudging through list of functions to wrap. --- swiftsimio/_array_functions.py | 340 ++++++++++++++++++++++++++------- swiftsimio/objects.py | 126 +++++++----- tests/test_cosmo_array.py | 37 ++-- 3 files changed, 362 insertions(+), 141 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index f02195b5..27b8bc7c 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1,95 +1,289 @@ import numpy as np -from unyt import unyt_array, unyt_quantity +from unyt import unyt_quantity from unyt._array_functions import implements -from unyt._array_functions import _HANDLED_FUNCTIONS as _UNYT_HANDLED_FUNCTIONS -from .objects import _multiply_cosmo_factor, cosmo_array, cosmo_quantity +from .objects import ( + cosmo_array, + cosmo_quantity, + _prepare_array_func_args, + _multiply_cosmo_factor, + _preserve_cosmo_factor, + _reciprocal_cosmo_factor, +) _HANDLED_FUNCTIONS = dict() -def arg_helper(*args, **kwargs): - cms = [(hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args] - cfs = [ - (hasattr(arg, "cosmo_factor"), getattr(arg, "cosmo_factor", None)) - for arg in args - ] - comps = [ - (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args - ] - kw_cms = [ - (hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args - ] - kw_cfs = [ - (hasattr(arg, "cosmo_factor"), getattr(arg, "cosmo_factor", None)) - for arg in args - ] - kw_comps = [ - (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args - ] - if all([cm[1] for cm in cms + kw_cms if cm[0]]): - # all cosmo inputs are comoving - ret_cm = True - elif all([cm[1] is None for cm in cms + kw_cms if cm[0]]): - # all cosmo inputs have comoving=None - ret_cm = None - elif any([cm[1] is None for cm in cms + kw_cms if cm[0]]): - # only some cosmo inputs have comoving=None - raise ValueError( - "Some arguments have comoving=None and others have comoving=True|False. " - "Result is undefined!" - ) - elif not any([cm[1] for cm in cms + kw_cms if cm[0]]): - # all cosmo_array inputs are physical - ret_cm = False - else: - # mix of comoving and physical inputs - args = [ - arg.to_comoving() if cm[0] and not cm[1] else arg - for arg, cm in zip(args, cms) - ] - kwargs = [ - kwarg.to_comoving() if cm[0] and not cm[1] else kwarg - for kwarg, cm in zip(kwargs, cms) - ] - ret_cm = True - if len(set(comps + kw_comps)) == 1: - # all compressions identical, preserve it - ret_comp = comps[0] - else: - # mixed compressions, strip it off - ret_comp = None - return args, cfs, kwargs, kw_cfs, ret_cm, ret_comp - - -@implements(np.dot) -def dot(a, b, out=None): - from unyt._array_functions import dot as unyt_dot - - args, cfs, kwargs, kw_cfs, ret_cm, ret_comp = arg_helper(a, b, out=out) - ret_cf = ( - cfs[0][1] * cfs[1][1] - ) # this needs helper function handling different cases - res = unyt_dot(*args, **kwargs) +def _return_helper(res, helper_result, ret_cf, out=None): if out is None: - # should we be using __array_wrap__ somehow? if isinstance(res, unyt_quantity): return cosmo_quantity( - res, comoving=ret_cm, cosmo_factor=ret_cf, compression=ret_comp + res, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf, + compression=helper_result["compression"], ) return cosmo_array( - res, comoving=ret_cm, cosmo_factor=ret_cf, compression=ret_comp + res, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf, + compression=helper_result["compression"], ) if hasattr(out, "comoving"): - out.comoving = ret_cm + out.comoving = helper_result["comoving"] if hasattr(out, "cosmo_factor"): out.cosmo_factor = ret_cf if hasattr(out, "compression"): - out.compression = ret_comp + out.compression = helper_result["compression"] return cosmo_array( # confused, do we set out, or return? res.to_value(res.units), res.units, bypass_validation=True, - comoving=ret_cm, + comoving=helper_result["comoving"], cosmo_factor=ret_cf, - compression=ret_comp, + compression=helper_result["compression"], + ) + + +@implements(np.array2string) +def array2string( + a, + max_line_width=None, + precision=None, + suppress_small=None, + separator=" ", + prefix="", + style=np._NoValue, + formatter=None, + threshold=None, + edgeitems=None, + sign=None, + floatmode=None, + suffix="", + *, + legacy=None, +): + from unyt._array_functions import array2string as unyt_array2string + + res = unyt_array2string( + a, + max_line_width=max_line_width, + precision=precision, + suppress_small=suppress_small, + separator=separator, + prefix=prefix, + style=style, + formatter=formatter, + threshold=threshold, + edgeitems=edgeitems, + sign=sign, + floatmode=floatmode, + suffix=suffix, + legacy=legacy, + ) + if a.comoving: + append = " (comoving)" + elif a.comoving is False: + append = " (physical)" + elif a.comoving is None: + append = "" + return res + append + + +@implements(np.dot) +def dot(a, b, out=None): + from unyt._array_functions import dot as unyt_dot + + helper_result = _prepare_array_func_args(a, b, out=out) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] ) + res = unyt_dot(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) + + +@implements(np.vdot) +def vdot(a, b, /): + from unyt._array_functions import vdot as unyt_vdot + + helper_result = _prepare_array_func_args(a, b) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_vdot(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.inner) +def inner(a, b, /): + from unyt._array_functions import inner as unyt_inner + + helper_result = _prepare_array_func_args(a, b) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_inner(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.outer) +def outer(a, b, out=None): + from unyt._array_functions import outer as unyt_outer + + helper_result = _prepare_array_func_args(a, b, out=out) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_outer(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) + + +@implements(np.kron) +def kron(a, b): + from unyt._array_functions import kron as unyt_kron + + helper_result = _prepare_array_func_args(a, b) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_kron(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.histogram_bin_edges) +def histogram_bin_edges(a, bins=10, range=None, weights=None): + from unyt._array_functions import histogram_bin_edges as unyt_histogram_bin_edges + + helper_result = _prepare_array_func_args(a, bins=bins, range=range, weights=weights) + if not isinstance(bins, str) and np.ndim(bins) == 1: + # we got bin edges as input + ret_cf = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["bins"]) + else: + # bins based on values in a + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_histogram_bin_edges(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.linalg.inv) +def linalg_inv(a): + from unyt._array_functions import linalg_inv as unyt_linalg_inv + + helper_result = _prepare_array_func_args(a) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_linalg_inv(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.linalg.tensorinv) +def linalg_tensorinv(a, ind=2): + from unyt._array_functions import linalg_tensorinv as unyt_linalg_tensorinv + + helper_result = _prepare_array_func_args(a, ind=ind) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_linalg_tensorinv(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.linalg.pinv) +def linalg_pinv(a, rcond=None, hermitian=False, *, rtol=np._NoValue): + from unyt._array_functions import linalg_pinv as unyt_linalg_pinv + + helper_result = _prepare_array_func_args( + a, + rcond=rcond, + hermitian=hermitian, + rtol=rtol + ) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_linalg_pinv(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(...) +# def ...(...): +# from unyt._array_functions import ... as unyt_... + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 921dca30..89196ced 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -371,6 +371,69 @@ def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): ) +def _prepare_array_func_args(*args, **kwargs): + cms = [(hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args] + ca_cfs = [ + (hasattr(arg, "cosmo_factor"), getattr(arg, "cosmo_factor", None)) + for arg in args + ] + comps = [ + (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args + ] + kw_cms = { + k: (hasattr(kwarg, "comoving"), getattr(kwarg, "comoving", None)) + for k, kwarg in kwargs.items() + } + kw_ca_cfs = { + k: (hasattr(kwarg, "cosmo_factor"), getattr(kwarg, "cosmo_factor", None)) + for k, kwarg in kwargs.items() + } + kw_comps = { + k: (hasattr(kwarg, "compression"), getattr(kwarg, "compression", None)) + for k, kwarg in kwargs.items() + } + if all([cm[1] for cm in cms + list(kw_cms.values()) if cm[0]]): + # all cosmo inputs are comoving + ret_cm = True + elif all([cm[1] is None for cm in cms + list(kw_cms.values()) if cm[0]]): + # all cosmo inputs have comoving=None + ret_cm = None + elif any([cm[1] is None for cm in cms + list(kw_cms.values()) if cm[0]]): + # only some cosmo inputs have comoving=None + raise ValueError( + "Some arguments have comoving=None and others have comoving=True|False. " + "Result is undefined!" + ) + elif all([cm[1] is False for cm in cms + list(kw_cms.values()) if cm[0]]): + # all cosmo_array inputs are physical + ret_cm = False + else: + # mix of comoving and physical inputs + args = [ + arg.to_comoving() if cm[0] and not cm[1] else arg + for arg, cm in zip(args, cms) + ] + kwargs = { + k: kwarg.to_comoving() if kw_cms[k][0] and not kw_cms[k][1] else kwarg + for k, kwarg in kwargs.items() + } + ret_cm = True + if len(set(comps + list(kw_comps.values()))) == 1: + # all compressions identical, preserve it + ret_comp = comps[0] + else: + # mixed compressions, strip it off + ret_comp = None + return dict( + args=args, + kwargs=kwargs, + ca_cfs=ca_cfs, + kw_ca_cfs=kw_ca_cfs, + comoving=ret_cm, + compression=ret_comp, + ) + + class InvalidScaleFactor(Exception): """ Raised when a scale factor is invalid, such as when adding @@ -1082,50 +1145,9 @@ def from_pint( return obj - # TODO: def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): - cms = [ - (hasattr(inp, "comoving"), getattr(inp, "comoving", None)) for inp in inputs - ] - cfs = [ - (hasattr(inp, "cosmo_factor"), getattr(inp, "cosmo_factor", None)) - for inp in inputs - ] - comps = [ - (hasattr(inp, "compression"), getattr(inp, "compression", None)) - for inp in inputs - ] - - # if we're here at least one input must be a cosmo_array - if all([cm[1] for cm in cms if cm[0]]): - # all cosmo_array inputs are comoving - ret_cm = True - elif all([cm[1] is None for cm in cms if cm[0]]): - # all cosmo inputs have comoving=None - ret_cm = None - elif any([cm[1] is None for cm in cms if cm[0]]): - # only some cosmo inputs have comoving=None - raise ValueError( - "Some arguments have comoving=None and others have comoving=True|False. " - "Result is undefined!" - ) - elif not any([cm[1] for cm in cms if cm[0]]): - # all cosmo_array inputs are physical - ret_cm = False - else: - # mix of comoving and physical inputs - inputs = [ - inp.to_comoving() if cm[0] and not cm[1] else inp - for inp, cm in zip(inputs, cms) - ] - ret_cm = True - - if len(set(comps)) == 1: - # all compressions identical, preserve it - ret_comp = comps[0] - else: - # mixed compressions, strip it off - ret_comp = None + helper_result = _prepare_array_func_args(*inputs, **kwargs) + cfs = helper_result["cfs"] # make sure we evaluate the cosmo_factor_ufunc_registry function: # might raise/warn even if we're not returning a cosmo_array @@ -1154,28 +1176,28 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): ) for r in ret: if isinstance(r, type(self)): - r.comoving = ret_cm + r.comoving = helper_result["comoving"] r.cosmo_factor = ret_cf - r.compression = ret_comp + r.compression = helper_result["compression"] if isinstance(ret, unyt_array): ret = ret.view(type(self)) - ret.comoving = ret_cm + ret.comoving = helper_result["comoving"] ret.cosmo_factor = ret_cf - ret.compression = ret_comp + ret.compression = helper_result["compression"] if "out" in kwargs: out = kwargs.pop("out") if ufunc not in multiple_output_operators: out = out[0] if isinstance(out, cosmo_array): - out.comoving = ret_cm + out.comoving = helper_result["comoving"] out.cosmo_factor = ret_cf - out.compression = ret_comp + out.compression = helper_result["compression"] else: for o in out: if isinstance(o, type(self)): - o.comoving = ret_cm + o.comoving = helper_result["comoving"] o.cosmo_factor = ret_cf - o.compression = ret_comp + o.compression = helper_result["compression"] return ret @@ -1206,7 +1228,7 @@ def __array_function__(self, func, types, args, kwargs): # otherwise default to numpy's private implementation return func._implementation(*args, **kwargs) # Note: this allows subclasses that don't override - # __array_function__ to handle unyt_array objects + # __array_function__ to handle cosmo_array objects if not all(issubclass(t, cosmo_array) or t is np.ndarray for t in types): return NotImplemented return _HANDLED_FUNCTIONS[func](*args, **kwargs) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 5cc6aebc..ff64360d 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -6,16 +6,11 @@ import os import numpy as np import unyt as u -from swiftsimio.objects import cosmo_array, cosmo_factor, a +from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a savetxt_file = "saved_array.txt" -# dummy until actually implemented: -class cosmo_quantity: - pass - - def getfunc(fname): func = np for attr in fname.split("."): @@ -32,13 +27,19 @@ def to_ua(x): def check_result(x_c, x_u): - # careful, unyt_quantity subclass of unyt_array + if isinstance(x_u, str): + assert isinstance(x_c, str) + return + # careful, unyt_quantity is a subclass of unyt_array if isinstance(x_u, u.unyt_quantity): assert isinstance(x_c, cosmo_quantity) elif isinstance(x_u, u.unyt_array): assert isinstance(x_c, cosmo_array) and not isinstance(x_c, cosmo_quantity) else: assert not isinstance(x_c, cosmo_array) + assert x_c.units == x_u.units + assert np.allclose(x_c.to_value(x_c.units), x_u.to_value(x_u.units)) + return class TestCosmoArrayInit: @@ -94,12 +95,20 @@ def test_handled_funcs(self): from unyt._array_functions import _HANDLED_FUNCTIONS functions_to_check = { + "array2string": (ca(np.arange(3)),), "dot": (ca(np.arange(3)), ca(np.arange(3))), "vdot": (ca(np.arange(3)), ca(np.arange(3))), "inner": (ca(np.arange(3)), ca(np.arange(3))), "outer": (ca(np.arange(3)), ca(np.arange(3))), "kron": (ca(np.arange(3)), ca(np.arange(3))), "histogram_bin_edges": (ca(np.arange(3)),), + "linalg.inv": (ca(np.eye(3)),), + "linalg.tensorinv": (ca(np.eye(9).reshape((3, 3, 3, 3))),), + "linalg.pinv": (ca(np.eye(3)),), + "linalg.svd": (ca(np.eye(3)),), + "histogram": (ca(np.arange(3)),), + "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), + "histogramdd": (ca(np.arange(3)),), "concatenate": (ca(np.eye(3)),), "cross": (ca(np.arange(3)), ca(np.arange(3))), "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), @@ -153,13 +162,7 @@ def test_handled_funcs(self): "correlate": (ca(np.arange(3)), ca(np.arange(3))), "tensordot": (ca(np.eye(3)), ca(np.eye(3))), "unwrap": (ca(np.arange(3)),), - "histogram": (ca(np.arange(3)),), - "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), - "histogramdd": (ca(np.arange(3)),), "linalg.norm": (ca(np.arange(3)),), - "linalg.inv": (ca(np.eye(3)),), - "linalg.tensorinv": (ca(np.eye(9).reshape((3, 3, 3, 3))),), - "linalg.pinv": (ca(np.eye(3)),), "linalg.det": (ca(np.eye(3)),), "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), @@ -169,7 +172,6 @@ def test_handled_funcs(self): ), "linalg.eigvals": (ca(np.eye(3)),), "linalg.eigvalsh": (ca(np.eye(3)),), - "linalg.svd": (ca(np.eye(3)),), "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), "linalg.eig": (ca(np.eye(3)),), "linalg.eigh": (ca(np.eye(3)),), @@ -189,7 +191,6 @@ def test_handled_funcs(self): "fft.irfftn": (ca(np.arange(3)),), "fft.fftshift": (ca(np.arange(3)),), "fft.ifftshift": (ca(np.arange(3)),), - "array2string": (ca(np.arange(3)),), "copyto": (ca(np.arange(3)), ca(np.arange(3))), "savetxt": (savetxt_file, ca(np.arange(3))), "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), @@ -249,7 +250,11 @@ def test_handled_funcs(self): assert len(unchecked_functions) == 0 except AssertionError: raise AssertionError( - "Did not check functions", [f.__name__ for f in unchecked_functions] + "Did not check functions", + [ + ".".join((f.__module__, f.__name__)).replace("numpy", "np") + for f in unchecked_functions + ], ) def test_getitem(self): From c9b8b6b9ceab21da1a3e393a24b131f48dd80403 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 13 Dec 2024 17:58:51 +0000 Subject: [PATCH 05/29] Work out an approach for concatenate-like functions. --- swiftsimio/_array_functions.py | 60 +++++++++++++++++++--------------- swiftsimio/objects.py | 18 ++++++++-- tests/test_cosmo_array.py | 10 +++--- 3 files changed, 54 insertions(+), 34 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 27b8bc7c..c5de68e7 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -46,21 +46,21 @@ def _return_helper(res, helper_result, ret_cf, out=None): @implements(np.array2string) def array2string( - a, - max_line_width=None, - precision=None, - suppress_small=None, - separator=" ", - prefix="", - style=np._NoValue, - formatter=None, - threshold=None, - edgeitems=None, - sign=None, - floatmode=None, - suffix="", - *, - legacy=None, + a, + max_line_width=None, + precision=None, + suppress_small=None, + separator=" ", + prefix="", + style=np._NoValue, + formatter=None, + threshold=None, + edgeitems=None, + sign=None, + floatmode=None, + suffix="", + *, + legacy=None, ): from unyt._array_functions import array2string as unyt_array2string @@ -160,7 +160,9 @@ def histogram_bin_edges(a, bins=10, range=None, weights=None): else: # bins based on values in a ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) - res = unyt_histogram_bin_edges(*helper_result["args"], **helper_result["kwargs"]) + res = unyt_histogram_bin_edges( + *helper_result["args"], **helper_result["kwargs"] + ) return _return_helper(res, helper_result, ret_cf) @@ -189,10 +191,7 @@ def linalg_pinv(a, rcond=None, hermitian=False, *, rtol=np._NoValue): from unyt._array_functions import linalg_pinv as unyt_linalg_pinv helper_result = _prepare_array_func_args( - a, - rcond=rcond, - hermitian=hermitian, - rtol=rtol + a, rcond=rcond, hermitian=hermitian, rtol=rtol ) ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) res = unyt_linalg_pinv(*helper_result["args"], **helper_result["kwargs"]) @@ -229,14 +228,21 @@ def linalg_pinv(a, rcond=None, hermitian=False, *, rtol=np._NoValue): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(...) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +@implements(np.concatenate) +def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): + from unyt._array_functions import concatenate as unyt_concatenate -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + tup, axis=axis, out=out, dtype=dtype, casting=casting + ) + helper_result_concat_items = _prepare_array_func_args(*tup) + ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) + res = unyt_concatenate( + helper_result_concat_items["args"], + *helper_result["args"][1:], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_concat_items, ret_cf, out=out) # @implements(...) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 89196ced..17f59d83 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -333,7 +333,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) + return cosmo_factor(a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -372,6 +372,20 @@ def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): def _prepare_array_func_args(*args, **kwargs): + # unyt allows creating a unyt_array from e.g. arrays with heterogenous units + # (it probably shouldn't...). + # Example: + # >>> u.unyt_array([np.arange(3), np.arange(3) * u.m]) + # unyt_array([[0, 1, 2], + # [0, 1, 2]], '(dimensionless)') + # It's impractical for cosmo_array to try to cover + # all possible invalid user input without unyt being stricter. + # This function checks for consistency for all args and kwargs, but is not recursive + # so mixed cosmo attributes could be passed in the first argument to np.concatenate, + # for instance. This function can be used "recursively" in a limited way manually: + # in functions like np.concatenate where a list of arrays is expected, it makes sense + # to pass the first argument (of np.concatenate - an iterable) to this function + # to check consistency and attempt to coerce to comoving if needed. cms = [(hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args] ca_cfs = [ (hasattr(arg, "cosmo_factor"), getattr(arg, "cosmo_factor", None)) @@ -632,7 +646,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index ff64360d..ef078c48 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -19,7 +19,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) def to_ua(x): @@ -47,7 +47,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -58,7 +58,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -68,7 +68,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -78,7 +78,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") From 37dd6afee2882d00b8e7a5ac7c2b45b536f631b7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 13 Dec 2024 22:29:33 +0000 Subject: [PATCH 06/29] Implement a few more funcs, rough in all the rest. --- swiftsimio/_array_functions.py | 933 +++++++++++++++++++++++++++++++-- tests/test_cosmo_array.py | 12 +- 2 files changed, 905 insertions(+), 40 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index c5de68e7..1d3b6bd8 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1,5 +1,5 @@ import numpy as np -from unyt import unyt_quantity +from unyt import unyt_quantity, unyt_array from unyt._array_functions import implements from .objects import ( cosmo_array, @@ -22,12 +22,16 @@ def _return_helper(res, helper_result, ret_cf, out=None): cosmo_factor=ret_cf, compression=helper_result["compression"], ) - return cosmo_array( - res, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf, - compression=helper_result["compression"], - ) + elif isinstance(res, unyt_array): + return cosmo_array( + res, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf, + compression=helper_result["compression"], + ) + else: + # unyt returned a bare array + return res if hasattr(out, "comoving"): out.comoving = helper_result["comoving"] if hasattr(out, "cosmo_factor"): @@ -198,33 +202,79 @@ def linalg_pinv(a, rcond=None, hermitian=False, *, rtol=np._NoValue): return _return_helper(res, helper_result, ret_cf) -# @implements(...) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +@implements(np.linalg.svd) +def linalg_svd(a, full_matrices=True, compute_uv=True, hermitian=False): + from unyt._array_functions import linalg_svd as unyt_linalg_svd -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + full_matrices=full_matrices, + compute_uv=compute_uv, + hermitian=hermitian + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ress = unyt_linalg_svd(*helper_result["args"], **helper_result["kwargs"]) + if compute_uv: + return tuple(_return_helper(res, helper_result, ret_cf) for res in ress) + else: + return _return_helper(ress, helper_result, ret_cf) -# @implements(...) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +@implements(np.histogram) +def histogram(a, bins=10, range=None, density=None, weights=None): + from unyt._array_functions import histogram as unyt_histogram -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + bins=bins, + range=range, + density=density, + weights=weights + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + counts, bins = unyt_histogram(*helper_result["args"], **helper_result["kwargs"]) + return counts, _return_helper(bins, helper_result, ret_cf) -# @implements(...) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +# ND HISTOGRAMS ARE TRICKY - EACH AXIS CAN HAVE DIFFERENT COSMO FACTORS -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# @implements(np.histogram2d) +# def histogram2d(x, y, bins=10, range=None, density=None, weights=None): +# from unyt._array_functions import histogram2d as unyt_histogram2d + +# helper_result = _prepare_array_func_args( +# x, +# y, +# bins=bins, +# range=range, +# density=density, +# weights=weights +# ) +# ret_cf_x = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) +# ret_cf_y = _preserve_cosmo_factor(helper_result["ca_cfs"][1]) +# counts, xbins, ybins = unyt_histogram2d( +# *helper_result["args"], **helper_result["kwargs"] +# ) +# return ( +# counts, +# _return_helper(xbins, helper_result, ret_cf), +# _return_helper(ybins, helper_result, ret_cf), +# ) + + +# @implements(np.histogramdd) +# def histogramdd(sample, bins=10, range=None, density=None, weights=None): +# from unyt._array_functions import histogramdd as unyt_histogramdd + +# helper_result = _prepare_array_func_args( +# sample, +# bins=bins, +# range=range, +# density=density, +# weights=weights +# ) +# ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) +# counts, bins = unyt_histogramdd(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) @@ -245,7 +295,7 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): return _return_helper(res, helper_result_concat_items, ret_cf, out=out) -# @implements(...) +# @implements(np.cross) # def ...(...): # from unyt._array_functions import ... as unyt_... @@ -255,7 +305,7 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(...) +# @implements(np.intersect1d) # def ...(...): # from unyt._array_functions import ... as unyt_... @@ -265,7 +315,7 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(...) +# @implements(np.union1d) # def ...(...): # from unyt._array_functions import ... as unyt_... @@ -275,7 +325,7 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(...) +# @implements(np.linalg.norm) # def ...(...): # from unyt._array_functions import ... as unyt_... @@ -285,11 +335,824 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(...) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +@implements(np.vstack) +def vstack(tup, *, dtype=None, casting="same_kind"): + from unyt._array_functions import vstack as unyt_vstack + + helper_result = _prepare_array_func_args( + tup, dtype=dtype, casting=casting + ) + helper_result_concat_items = _prepare_array_func_args(*tup) + ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) + res = unyt_vstack( + helper_result_concat_items["args"], + *helper_result["args"][1:], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_concat_items, ret_cf) + + +@implements(np.hstack) +def hstack(tup, *, dtype=None, casting="same_kind"): + from unyt._array_functions import hstack as unyt_hstack + + helper_result = _prepare_array_func_args( + tup, dtype=dtype, casting=casting + ) + helper_result_concat_items = _prepare_array_func_args(*tup) + ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) + res = unyt_hstack( + helper_result_concat_items["args"], + *helper_result["args"][1:], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_concat_items, ret_cf) + + +@implements(np.dstack) +def dstack(tup): + from unyt._array_functions import dstack as unyt_dstack + + helper_result_concat_items = _prepare_array_func_args(*tup) + ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) + res = unyt_dstack(helper_result_concat_items["args"]) + return _return_helper(res, helper_result_concat_items, ret_cf) + + +@implements(np.column_stack) +def column_stack(tup): + from unyt._array_functions import column_stack as unyt_column_stack + + helper_result_concat_items = _prepare_array_func_args(*tup) + ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) + res = unyt_column_stack(helper_result_concat_items["args"]) + return _return_helper(res, helper_result_concat_items, ret_cf) + + +@implements(np.stack) +def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"): + from unyt._array_functions import stack as unyt_stack + + helper_result = _prepare_array_func_args( + arrays, axis=axis, out=out, dtype=dtype, casting=casting + ) + helper_result_concat_items = _prepare_array_func_args(*arrays) + ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) + res = unyt_stack( + helper_result_concat_items["args"], + *helper_result["args"][1:], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_concat_items, ret_cf, out=out) + + +# @implements(np.around) +# def around(...): +# from unyt._array_functions import around as unyt_around # helper_result = _prepare_array_func_args(...) # ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# res = unyt_around(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.block) +# def block(...): +# from unyt._array_functions import block as unyt_block + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_block(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# UNYT HAS A COPY-PASTED TYPO fft -> ftt + +# @implements(np.fft.fft) +# def ftt_fft(...): +# from unyt._array_functions import ftt_fft as unyt_fft_fft + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_fft(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.fft2) +# def ftt_fft2(...): +# from unyt._array_functions import ftt_fft2 as unyt_fft_fft2 + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_fft2(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.fftn) +# def ftt_fftn(...): +# from unyt._array_functions import ftt_fftn as unyt_fft_fftn + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_fftn(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.hfft) +# def ftt_hfft(...): +# from unyt._array_functions import ftt_hfft as unyt_fft_hfft + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_hfft(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.rfft) +# def ftt_rfft(...): +# from unyt._array_functions import ftt_rfft as unyt_fft_rfft + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_rfft(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.rfft2) +# def fft_rfft2(...): +# from unyt._array_functions import ftt_rfft2 as unyt_fft_rfft2 + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_rfft2(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.rfftn) +# def fft_rfftn(...): +# from unyt._array_functions import ftt_rfftn as unyt_fft_rfftn + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_rfftn(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.ifft) +# def fft_ifft(...): +# from unyt._array_functions import ftt_ifft as unyt_fft_ifft + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_ifft(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.ifft2) +# def fft_ifft2(...): +# from unyt._array_functions import ftt_ifft2 as unyt_fft_ifft2 + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_ifft2(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.ifftn) +# def fft_ifftn(...): +# from unyt._array_functions import ftt_ifftn as unyt_fft_ifftn + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_ifftn(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.ihfft) +# def fft_ihfft(...): +# from unyt._array_functions import ftt_ihfft as unyt_fft_ihfft + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_ihfft(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.irfft) +# def fft_irfft(...): +# from unyt._array_functions import ftt_irfft as unyt_fft_irfft + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_irfft(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.irfft2) +# def fft_irfft2(...): +# from unyt._array_functions import ftt_irfft2 as unyt_fft_irfft2 + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_irfft2(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.irfftn) +# def fft_irfftn(...): +# from unyt._array_functions import ftt_irfftn as unyt_fft_irfftn + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_irfftn(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.fftshift) +# def fft_fftshift(...): +# from unyt._array_functions import fft_fftshift as unyt_fft_fftshift + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_fftshift(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fft.ifftshift) +# def fft_ifftshift(...): +# from unyt._array_functions import fft_ifftshift as unyt_fft_ifftshift + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fft_ifftshift(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.sort_complex) +# def sort_complex(...): +# from unyt._array_functions import sort_complex as unyt_sort_complex + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_sort_complex(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.isclose) +# def isclose(...): +# from unyt._array_functions import isclose as unyt_isclose + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_isclose(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.allclose) +# def allclose(...): +# from unyt._array_functions import allclose as unyt_allclose + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_allclose(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.array_equal) +# def array_equal(...): +# from unyt._array_functions import array_equal as unyt_array_equal + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_array_equal(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.array_equiv) +# def array_equiv(...): +# from unyt._array_functions import array_equiv as unyt_array_equiv + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_array_equiv(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linspace) +# def linspace(...): +# from unyt._array_functions import linspace as unyt_linspace + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linspace(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.logspace) +# def logspace(...): +# from unyt._array_functions import logspace as unyt_logspace + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_logspace(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.geomspace) +# def geomspace(...): +# from unyt._array_functions import geomspace as unyt_geomspace + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_geomspace(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.copyto) +# def copyto(...): +# from unyt._array_functions import copyto as unyt_copyto + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_copyto(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.prod) +# def prod(...): +# from unyt._array_functions import prod as unyt_prod + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.var) +# def var(...): +# from unyt._array_functions import var as unyt_var + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_var(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.trace) +# def trace(...): +# from unyt._array_functions import trace as unyt_trace + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_trace(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.percentile) +# def percentile(...): +# from unyt._array_functions import percentile as unyt_percentile + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_percentile(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.quantile) +# def quantile(...): +# from unyt._array_functions import quantile as unyt_quantile + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_quantile(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.nanpercentile) +# def nanpercentile(...): +# from unyt._array_functions import nanpercentile as unyt_nanpercentile + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_nanpercentile(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.nanquantile) +# def nanquantile(...): +# from unyt._array_functions import nanquantile as unyt_nanquantile + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_nanquantile(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.det) +# def linalg_det(...): +# from unyt._array_functions import linalg_det as unyt_linalg_det + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_det(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.lstsq) +# def linalg_lstsq(...): +# from unyt._array_functions import linalg_lstsq as unyt_linalg_lstsq + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_lstsq(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.solve) +# def linalg_solve(...): +# from unyt._array_functions import linalg_solve as unyt_linalg_solve + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_solve(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.tensorsolve) +# def linalg_tensorsolve(...): +# from unyt._array_functions import linalg_tensorsolve as unyt_linalg_tensorsolve + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_tensorsolve(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.eig) +# def linalg_eig(...): +# from unyt._array_functions import linalg_eig as unyt_linalg_eig + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.eigh) +# def linalg_eigh(...): +# from unyt._array_functions import linalg_eigh as unyt_linalg_eigh + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.eigvals) +# def linalg_eigvals(...): +# from unyt._array_functions import linalg_eigvals as unyt_linalg_eigvals + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_eigvals(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.eigvalsh) +# def linalg_eigvalsh(...): +# from unyt._array_functions import linalg_eigvalsh as unyt_linalg_eigvalsh + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_eigvalsh(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.savetxt) +# def savetxt(...): +# from unyt._array_functions import savetxt as unyt_savetxt + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_savetxt(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.apply_over_axes) +# def apply_over_axes(...): +# from unyt._array_functions import apply_over_axes as unyt_apply_over_axes + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_apply_over_axes(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.diff) +# def diff(...): +# from unyt._array_functions import diff as unyt_diff + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_diff(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.ediff1d) +# def ediff1d(...): +# from unyt._array_functions import ediff1d as unyt_ediff1d + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_ediff1d(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.ptp) +# def ptp(...): +# from unyt._array_functions import ptp as unyt_ptp + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_ptp(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.cumprod) +# def cumprod(...): +# from unyt._array_functions import cumprod as unyt_cumprod + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_cumprod(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.pad) +# def pad(...): +# from unyt._array_functions import pad as unyt_pad + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_pad(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.choose) +# def choose(...): +# from unyt._array_functions import choose as unyt_choose + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_choose(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.fill_diagonal) +# def fill_diagonal(...): +# from unyt._array_functions import fill_diagonal as unyt_fill_diagonal + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_fill_diagonal(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.insert) +# def insert(...): +# from unyt._array_functions import insert as unyt_insert + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_insert(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.isin) +# def isin(...): +# from unyt._array_functions import isin as unyt_isin + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.place) +# def place(...): +# from unyt._array_functions import place as unyt_place + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_place(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.put) +# def ...(...): +# from unyt._array_functions import put as unyt_put + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_put(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.put_along_axis) +# def put_along_axis(...): +# from unyt._array_functions import put_along_axis as unyt_put_along_axis + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_put_along_axis(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.putmask) +# def putmask(...): +# from unyt._array_functions import putmask as unyt_putmask + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_putmask(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.searchsorted) +# def searchsorted(...): +# from unyt._array_functions import searchsorted as unyt_searchsorted + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_searchsorted(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.select) +# def select(...): +# from unyt._array_functions import select as unyt_select + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_select(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.setdiff1d) +# def setdiff1d(...): +# from unyt._array_functions import setdiff1d as unyt_setdiff1d + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_setdiff1d(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.sinc) +# def sinc(...): +# from unyt._array_functions import sinc as unyt_sinc + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_sinc(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.clip) +# def clip(...): +# from unyt._array_functions import clip as unyt_clip + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_clip(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.where) +# def where(...): +# from unyt._array_functions import where as unyt_where + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.triu) +# def triu(...): +# from unyt._array_functions import triu as unyt_triu + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_triu(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.tril) +# def tril(...): +# from unyt._array_functions import tril as unyt_tril + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_tril(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.einsum) +# def einsum(...): +# from unyt._array_functions import einsum as unyt_einsum + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_einsum(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.convolve) +# def convolve(...): +# from unyt._array_functions import convolve as unyt_convolve + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_convolve(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.correlate) +# def correlate(...): +# from unyt._array_functions import correlate as unyt_correlate + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_correlate(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.tensordot) +# def tensordot(...): +# from unyt._array_functions import tensordot as unyt_tensordot + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_tensordot(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.unwrap) +# def unwrap(...): +# from unyt._array_functions import unwrap as unyt_unwrap + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_unwrap(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.interp) +# def interp(...): +# from unyt._array_functions import interp as unyt_interp + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_interp(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.array_repr) +# def array_repr(...): +# from unyt._array_functions import array_repr as unyt_array_repr + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_array_repr(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.linalg.outer) +# def linalg_outer(...): +# from unyt._array_functions import linalg_outer as unyt_linalg_outer + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_linalg_outer(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.trapezoid) +# def trapezoid(...): +# from unyt._array_functions import trapezoid as unyt_trapezoid + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_trapezoid(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.in1d) +# def in1d(...): +# from unyt._array_functions import in1d as unyt_in1d + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_in1d(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index ef078c48..44fcd126 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -19,7 +19,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) def to_ua(x): @@ -37,6 +37,8 @@ def check_result(x_c, x_u): assert isinstance(x_c, cosmo_array) and not isinstance(x_c, cosmo_quantity) else: assert not isinstance(x_c, cosmo_array) + assert np.allclose(x_c, x_u) + return assert x_c.units == x_u.units assert np.allclose(x_c.to_value(x_c.units), x_u.to_value(x_u.units)) return @@ -47,7 +49,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -58,7 +60,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -68,7 +70,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -78,7 +80,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") From fd255dc22ed0ab0c30b97a6706cfcd303c74b856 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 14 Dec 2024 09:04:41 +0000 Subject: [PATCH 07/29] Fill in FFT functions, continue roughing in. --- swiftsimio/_array_functions.py | 270 ++++++++++++++++----------------- tests/test_cosmo_array.py | 144 +++++++++--------- 2 files changed, 207 insertions(+), 207 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 1d3b6bd8..91054c18 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -296,42 +296,42 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): # @implements(np.cross) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +# def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None): +# from unyt._array_functions import cross as unyt_cross -# helper_result = _prepare_array_func_args(...) +# helper_result = _prepare_array_func_args(a, b, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis) # ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# res = unyt_cross(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.intersect1d) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +# def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): +# from unyt._array_functions import intersect1d as unyt_intersect1d -# helper_result = _prepare_array_func_args(...) +# helper_result = _prepare_array_func_args(ar1, ar2, assume_unique=assume_unique, return_indices=return_indices) # ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# res = unyt_intersect1d(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.union1d) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +# def union1d(ar1, ar2): +# from unyt._array_functions import union1d as unyt_union1d -# helper_result = _prepare_array_func_args(...) +# helper_result = _prepare_array_func_args(ar1, ar2) # ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# res = unyt_union1d(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.linalg.norm) -# def ...(...): -# from unyt._array_functions import ... as unyt_... +# def linalg_norm(x, ord=None, axis=None, keepdims=False): +# from unyt._array_functions import linalg_norm as unyt_linalg_norm -# helper_result = _prepare_array_func_args(...) +# helper_result = _prepare_array_func_args(x, ord=ord, axis=axis, keepdims=keepdims) # ret_cf = ...() -# res = unyt_...(*helper_result["args"], **helper_result["kwargs"]) +# res = unyt_linalg_norm(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) @@ -407,185 +407,185 @@ def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"): # @implements(np.around) -# def around(...): +# def around(a, decimals=0, out=None): # from unyt._array_functions import around as unyt_around -# helper_result = _prepare_array_func_args(...) +# helper_result = _prepare_array_func_args(a, decimals=decimals, out=out) # ret_cf = ...() # res = unyt_around(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.block) -# def block(...): +# def block(arrays): # from unyt._array_functions import block as unyt_block -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_block(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) +# helper_result_concat_items = _prepare_array_func_args(*arrays) +# ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) +# res = unyt_block(helper_result_concat_items["args"]) +# return _return_helper(res, helper_result_concat_items, ret_cf) # UNYT HAS A COPY-PASTED TYPO fft -> ftt -# @implements(np.fft.fft) -# def ftt_fft(...): -# from unyt._array_functions import ftt_fft as unyt_fft_fft +@implements(np.fft.fft) +def ftt_fft(a, n=None, axis=-1, norm=None, out=None): + from unyt._array_functions import ftt_fft as unyt_fft_fft -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_fft(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, n=n, axis=axis, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_fft(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.fft2) -# def ftt_fft2(...): -# from unyt._array_functions import ftt_fft2 as unyt_fft_fft2 +@implements(np.fft.fft2) +def ftt_fft2(a, s=None, axes=(-2, -1), norm=None, out=None): + from unyt._array_functions import ftt_fft2 as unyt_fft_fft2 -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_fft2(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_fft2(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.fftn) -# def ftt_fftn(...): -# from unyt._array_functions import ftt_fftn as unyt_fft_fftn +@implements(np.fft.fftn) +def ftt_fftn(a, s=None, axes=None, norm=None, out=None): + from unyt._array_functions import ftt_fftn as unyt_fft_fftn -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_fftn(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_fftn(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.hfft) -# def ftt_hfft(...): -# from unyt._array_functions import ftt_hfft as unyt_fft_hfft +@implements(np.fft.hfft) +def ftt_hfft(a, n=None, axis=-1, norm=None, out=None): + from unyt._array_functions import ftt_hfft as unyt_fft_hfft -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_hfft(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, n=n, axis=axis, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_hfft(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.rfft) -# def ftt_rfft(...): -# from unyt._array_functions import ftt_rfft as unyt_fft_rfft +@implements(np.fft.rfft) +def ftt_rfft(a, n=None, axis=-1, norm=None, out=None): + from unyt._array_functions import ftt_rfft as unyt_fft_rfft -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_rfft(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, n=n, axis=axis, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_rfft(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.rfft2) -# def fft_rfft2(...): -# from unyt._array_functions import ftt_rfft2 as unyt_fft_rfft2 +@implements(np.fft.rfft2) +def fft_rfft2(a, s=None, axes=(-2, -1), norm=None, out=None): + from unyt._array_functions import ftt_rfft2 as unyt_fft_rfft2 -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_rfft2(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_rfft2(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.rfftn) -# def fft_rfftn(...): -# from unyt._array_functions import ftt_rfftn as unyt_fft_rfftn +@implements(np.fft.rfftn) +def fft_rfftn(a, s=None, axes=None, norm=None, out=None): + from unyt._array_functions import ftt_rfftn as unyt_fft_rfftn -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_rfftn(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_rfftn(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.ifft) -# def fft_ifft(...): -# from unyt._array_functions import ftt_ifft as unyt_fft_ifft +@implements(np.fft.ifft) +def fft_ifft(a, n=None, axis=-1, norm=None, out=None): + from unyt._array_functions import ftt_ifft as unyt_fft_ifft -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_ifft(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, n=n, axis=axis, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_ifft(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.ifft2) -# def fft_ifft2(...): -# from unyt._array_functions import ftt_ifft2 as unyt_fft_ifft2 +@implements(np.fft.ifft2) +def fft_ifft2(a, s=None, axes=(-2, -1), norm=None, out=None): + from unyt._array_functions import ftt_ifft2 as unyt_fft_ifft2 -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_ifft2(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_ifft2(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.ifftn) -# def fft_ifftn(...): -# from unyt._array_functions import ftt_ifftn as unyt_fft_ifftn +@implements(np.fft.ifftn) +def fft_ifftn(a, s=None, axes=None, norm=None, out=None): + from unyt._array_functions import ftt_ifftn as unyt_fft_ifftn -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_ifftn(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_ifftn(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.ihfft) -# def fft_ihfft(...): -# from unyt._array_functions import ftt_ihfft as unyt_fft_ihfft +@implements(np.fft.ihfft) +def fft_ihfft(a, n=None, axis=-1, norm=None, out=None): + from unyt._array_functions import ftt_ihfft as unyt_fft_ihfft -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_ihfft(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, n=n, axis=axis, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_ihfft(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.irfft) -# def fft_irfft(...): -# from unyt._array_functions import ftt_irfft as unyt_fft_irfft +@implements(np.fft.irfft) +def fft_irfft(a, n=None, axis=-1, norm=None, out=None): + from unyt._array_functions import ftt_irfft as unyt_fft_irfft -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_irfft(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, n=n, axis=axis, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_irfft(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.irfft2) -# def fft_irfft2(...): -# from unyt._array_functions import ftt_irfft2 as unyt_fft_irfft2 +@implements(np.fft.irfft2) +def fft_irfft2(a, s=None, axes=(-2, -1), norm=None, out=None): + from unyt._array_functions import ftt_irfft2 as unyt_fft_irfft2 -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_irfft2(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_irfft2(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.irfftn) -# def fft_irfftn(...): -# from unyt._array_functions import ftt_irfftn as unyt_fft_irfftn +@implements(np.fft.irfftn) +def fft_irfftn(a, s=None, axes=None, norm=None, out=None): + from unyt._array_functions import ftt_irfftn as unyt_fft_irfftn -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_irfftn(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, s=s, axes=axes, norm=norm, out=out) + ret_cf = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_irfftn(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.fft.fftshift) -# def fft_fftshift(...): -# from unyt._array_functions import fft_fftshift as unyt_fft_fftshift +@implements(np.fft.fftshift) +def fft_fftshift(x, axes=None): + from unyt._array_functions import fft_fftshift as unyt_fft_fftshift -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_fftshift(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(x, axes=axes) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_fftshift(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.fft.ifftshift) -# def fft_ifftshift(...): -# from unyt._array_functions import fft_ifftshift as unyt_fft_ifftshift +@implements(np.fft.ifftshift) +def fft_ifftshift(x, axes=None): + from unyt._array_functions import fft_ifftshift as unyt_fft_ifftshift -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fft_ifftshift(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(x, axes=axes) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_fft_ifftshift(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.sort_complex) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 44fcd126..5240ae3f 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -109,74 +109,19 @@ def test_handled_funcs(self): "linalg.pinv": (ca(np.eye(3)),), "linalg.svd": (ca(np.eye(3)),), "histogram": (ca(np.arange(3)),), - "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), - "histogramdd": (ca(np.arange(3)),), + # "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), + # "histogramdd": (ca(np.arange(3)),), "concatenate": (ca(np.eye(3)),), - "cross": (ca(np.arange(3)), ca(np.arange(3))), - "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), - "union1d": (ca(np.arange(3)), ca(np.arange(3))), + # "cross": (ca(np.arange(3)), ca(np.arange(3))), + # "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), + # "union1d": (ca(np.arange(3)), ca(np.arange(3))), "vstack": (ca(np.arange(3)),), "hstack": (ca(np.arange(3)),), "dstack": (ca(np.arange(3)),), "column_stack": (ca(np.arange(3)),), "stack": (ca(np.arange(3)),), - "around": (ca(np.arange(3)),), - "block": (ca(np.arange(3)),), - "sort_complex": (ca(np.arange(3)),), - "linspace": (ca(1), ca(2), 10), - "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), - "geomspace": (ca(1), ca(1), 10), - "prod": (ca(np.arange(3)),), - "var": (ca(np.arange(3)),), - "trace": (ca(np.eye(3)),), - "percentile": (ca(np.arange(3)), 30), - "quantile": (ca(np.arange(3)), 0.3), - "nanpercentile": (ca(np.arange(3)), 30), - "nanquantile": (ca(np.arange(3)), 0.3), - "diff": (ca(np.arange(3)),), - "ediff1d": (ca(np.arange(3)),), - "ptp": (ca(np.arange(3)),), - "cumprod": (ca(np.arange(3)),), - "pad": (ca(np.arange(3)), 3), - "choose": (np.arange(3), ca(np.eye(3))), - "insert": (ca(np.arange(3)), 1, ca(1)), - "isin": (ca(np.arange(3)), ca(np.arange(3))), - "in1d": (ca(np.arange(3)), ca(np.arange(3))), - "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), - "isclose": (ca(np.arange(3)), ca(np.arange(3))), - "allclose": (ca(np.arange(3)), ca(np.arange(3))), - "array_equal": (ca(np.arange(3)), ca(np.arange(3))), - "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), - "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), - "select": ( - [np.arange(3) < 1, np.arange(3) > 1], - [ca(np.arange(3)), ca(np.arange(3))], - ca(1), - ), - "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), - "sinc": (ca(np.arange(3)),), - "clip": (ca(np.arange(3)), ca(1), ca(2)), - "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), - "triu": (ca(np.arange(3)),), - "tril": (ca(np.arange(3)),), - "einsum": ("ii->i", ca(np.eye(3))), - "convolve": (ca(np.arange(3)), ca(np.arange(3))), - "correlate": (ca(np.arange(3)), ca(np.arange(3))), - "tensordot": (ca(np.eye(3)), ca(np.eye(3))), - "unwrap": (ca(np.arange(3)),), - "linalg.norm": (ca(np.arange(3)),), - "linalg.det": (ca(np.eye(3)),), - "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), - "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), - "linalg.tensorsolve": ( - ca(np.eye(24).reshape((6, 4, 2, 3, 4))), - ca(np.ones((6, 4))), - ), - "linalg.eigvals": (ca(np.eye(3)),), - "linalg.eigvalsh": (ca(np.eye(3)),), - "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), - "linalg.eig": (ca(np.eye(3)),), - "linalg.eigh": (ca(np.eye(3)),), + # "around": (ca(np.arange(3)),), + # "block": (ca(np.arange(3)),), "fft.fft": (ca(np.arange(3)),), "fft.fft2": (ca(np.eye(3)),), "fft.fftn": (ca(np.arange(3)),), @@ -193,16 +138,71 @@ def test_handled_funcs(self): "fft.irfftn": (ca(np.arange(3)),), "fft.fftshift": (ca(np.arange(3)),), "fft.ifftshift": (ca(np.arange(3)),), - "copyto": (ca(np.arange(3)), ca(np.arange(3))), - "savetxt": (savetxt_file, ca(np.arange(3))), - "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), - "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), - "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), - "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), - "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), - "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), - "array_repr": (ca(np.arange(3)),), - "trapezoid": (ca(np.arange(3)),), + # "sort_complex": (ca(np.arange(3)),), + # "linspace": (ca(1), ca(2), 10), + # "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), + # "geomspace": (ca(1), ca(1), 10), + # "prod": (ca(np.arange(3)),), + # "var": (ca(np.arange(3)),), + # "trace": (ca(np.eye(3)),), + # "percentile": (ca(np.arange(3)), 30), + # "quantile": (ca(np.arange(3)), 0.3), + # "nanpercentile": (ca(np.arange(3)), 30), + # "nanquantile": (ca(np.arange(3)), 0.3), + # "diff": (ca(np.arange(3)),), + # "ediff1d": (ca(np.arange(3)),), + # "ptp": (ca(np.arange(3)),), + # "cumprod": (ca(np.arange(3)),), + # "pad": (ca(np.arange(3)), 3), + # "choose": (np.arange(3), ca(np.eye(3))), + # "insert": (ca(np.arange(3)), 1, ca(1)), + # "isin": (ca(np.arange(3)), ca(np.arange(3))), + # "in1d": (ca(np.arange(3)), ca(np.arange(3))), + # "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + # "isclose": (ca(np.arange(3)), ca(np.arange(3))), + # "allclose": (ca(np.arange(3)), ca(np.arange(3))), + # "array_equal": (ca(np.arange(3)), ca(np.arange(3))), + # "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), + # "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), + # "select": ( + # [np.arange(3) < 1, np.arange(3) > 1], + # [ca(np.arange(3)), ca(np.arange(3))], + # ca(1), + # ), + # "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), + # "sinc": (ca(np.arange(3)),), + # "clip": (ca(np.arange(3)), ca(1), ca(2)), + # "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + # "triu": (ca(np.arange(3)),), + # "tril": (ca(np.arange(3)),), + # "einsum": ("ii->i", ca(np.eye(3))), + # "convolve": (ca(np.arange(3)), ca(np.arange(3))), + # "correlate": (ca(np.arange(3)), ca(np.arange(3))), + # "tensordot": (ca(np.eye(3)), ca(np.eye(3))), + # "unwrap": (ca(np.arange(3)),), + # "linalg.norm": (ca(np.arange(3)),), + # "linalg.det": (ca(np.eye(3)),), + # "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), + # "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), + # "linalg.tensorsolve": ( + # ca(np.eye(24).reshape((6, 4, 2, 3, 4))), + # ca(np.ones((6, 4))), + # ), + # "linalg.eigvals": (ca(np.eye(3)),), + # "linalg.eigvalsh": (ca(np.eye(3)),), + # "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), + # "linalg.eig": (ca(np.eye(3)),), + # "linalg.eigh": (ca(np.eye(3)),), + # "copyto": (ca(np.arange(3)), ca(np.arange(3))), + # "savetxt": (savetxt_file, ca(np.arange(3))), + # "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), + # "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), + # "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), + # "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + # "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), + # "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + # "array_repr": (ca(np.arange(3)),), + # "trapezoid": (ca(np.arange(3)),), } functions_checked = list() bad_funcs = dict() From 960bdbbc7ee4ab4310e99bbff4aa0f2dbb223834 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 14 Dec 2024 11:14:46 +0000 Subject: [PATCH 08/29] Implement block, around. --- swiftsimio/_array_functions.py | 89 ++++++++++++++++++++++++++++------ swiftsimio/objects.py | 4 +- tests/test_cosmo_array.py | 4 +- 3 files changed, 79 insertions(+), 18 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 91054c18..5ec412ee 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -406,24 +406,85 @@ def stack(arrays, axis=0, out=None, *, dtype=None, casting="same_kind"): return _return_helper(res, helper_result_concat_items, ret_cf, out=out) -# @implements(np.around) -# def around(a, decimals=0, out=None): -# from unyt._array_functions import around as unyt_around +@implements(np.around) +def around(a, decimals=0, out=None): + from unyt._array_functions import around as unyt_around -# helper_result = _prepare_array_func_args(a, decimals=decimals, out=out) -# ret_cf = ...() -# res = unyt_around(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, decimals=decimals, out=out) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_around(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) + + +def _recursive_to_comoving(lst): + ret_lst = list() + for item in lst: + if isinstance(item, list): + ret_lst.append(_recursive_to_comoving(item)) + else: + ret_lst.append(item.to_comoving()) + return ret_lst + + +def _prepare_array_block_args(lst, recursing=False): + """ + Block accepts only a nested list of array "blocks". We need to recurse on this. + """ + helper_results = list() + if isinstance(lst, list): + for item in lst: + if isinstance(item, list): + helper_results += _prepare_array_block_args(item, recursing=True) + else: + helper_results.append(_prepare_array_func_args(item)) + if recursing: + return helper_results + cms = [hr["comoving"] for hr in helper_results] + comps = [hr["compression"] for hr in helper_results] + ca_cfs = [hr["ca_cfs"] for hr in helper_results] + convert_to_cm = False + if all(cms): + ret_cm = True + elif all([cm is None for cm in cms]): + ret_cm = None + elif any([cm is None for cm in cms]) and not all([cm is None for cm in cms]): + raise ValueError("Some input has comoving=None and others have " + "comoving=True|False. Result is undefined!") + elif all([cm is False for cm in cms]): + ret_cm = False + else: + # mix of True and False only + ret_cm = True + convert_to_cm = True + if len(set(comps)) == 1: + ret_comp = comps[0] + else: + ret_comp = None + ret_cf = ca_cfs[0] + for ca_cf in ca_cfs[1:]: + if ca_cf != ret_cf: + raise ValueError("Mixed cosmo_factor values in input.") + if convert_to_cm: + ret_lst = _recursive_to_comoving(lst) + else: + ret_lst = lst + return dict( + args=ret_lst, + kwargs=dict(), + comoving=ret_cm, + cosmo_factor=ret_cf, + compression=ret_comp, + ) -# @implements(np.block) -# def block(arrays): -# from unyt._array_functions import block as unyt_block +@implements(np.block) +def block(arrays): + from unyt._array_functions import block as unyt_block -# helper_result_concat_items = _prepare_array_func_args(*arrays) -# ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) -# res = unyt_block(helper_result_concat_items["args"]) -# return _return_helper(res, helper_result_concat_items, ret_cf) + helper_result_block = _prepare_array_block_args(arrays) + ret_cf = helper_result_block["cosmo_factor"] + res = unyt_block(helper_result_block["args"]) + return _return_helper(res, helper_result_block, ret_cf) # UNYT HAS A COPY-PASTED TYPO fft -> ftt diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 17f59d83..7b9f404f 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -333,7 +333,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a**0, scale_factor=cf1.scale_factor) + return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -646,7 +646,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 5240ae3f..6aac44bf 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -120,8 +120,8 @@ def test_handled_funcs(self): "dstack": (ca(np.arange(3)),), "column_stack": (ca(np.arange(3)),), "stack": (ca(np.arange(3)),), - # "around": (ca(np.arange(3)),), - # "block": (ca(np.arange(3)),), + "around": (ca(np.arange(3)),), + "block": ([[ca(np.arange(3))], [ca(np.arange(3))]],), "fft.fft": (ca(np.arange(3)),), "fft.fft2": (ca(np.eye(3)),), "fft.fftn": (ca(np.arange(3)),), From 77f1d275e106b6e8420146189fe22cb3de540a0d Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 14 Dec 2024 13:01:53 +0000 Subject: [PATCH 09/29] Implement a few more funcs. --- swiftsimio/_array_functions.py | 92 +++++++++++++++++++++------------- tests/test_cosmo_array.py | 10 ++-- 2 files changed, 62 insertions(+), 40 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 5ec412ee..28ee5440 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -295,44 +295,66 @@ def concatenate(tup, axis=0, out=None, dtype=None, casting="same_kind"): return _return_helper(res, helper_result_concat_items, ret_cf, out=out) -# @implements(np.cross) -# def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None): -# from unyt._array_functions import cross as unyt_cross +@implements(np.cross) +def cross(a, b, axisa=-1, axisb=-1, axisc=-1, axis=None): + from unyt._array_functions import cross as unyt_cross -# helper_result = _prepare_array_func_args(a, b, axisa=axisa, axisb=axisb, axisc=axisc, axis=axis) -# ret_cf = ...() -# res = unyt_cross(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + b, + axisa=axisa, + axisb=axisb, + axisc=axisc, + axis=axis, + ) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_cross(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.intersect1d) -# def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): -# from unyt._array_functions import intersect1d as unyt_intersect1d +@implements(np.intersect1d) +def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): + from unyt._array_functions import intersect1d as unyt_intersect1d -# helper_result = _prepare_array_func_args(ar1, ar2, assume_unique=assume_unique, return_indices=return_indices) -# ret_cf = ...() -# res = unyt_intersect1d(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + ar1, + ar2, + assume_unique=assume_unique, + return_indices=return_indices + ) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_intersect1d(*helper_result["args"], **helper_result["kwargs"]) + if return_indices: + return res + else: + return _return_helper(res, helper_result, ret_cf) -# @implements(np.union1d) -# def union1d(ar1, ar2): -# from unyt._array_functions import union1d as unyt_union1d +@implements(np.union1d) +def union1d(ar1, ar2): + from unyt._array_functions import union1d as unyt_union1d -# helper_result = _prepare_array_func_args(ar1, ar2) -# ret_cf = ...() -# res = unyt_union1d(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(ar1, ar2) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_union1d(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.linalg.norm) -# def linalg_norm(x, ord=None, axis=None, keepdims=False): -# from unyt._array_functions import linalg_norm as unyt_linalg_norm +@implements(np.linalg.norm) +def linalg_norm(x, ord=None, axis=None, keepdims=False): + # they didn't use linalg_norm, doesn't follow usual pattern: + from unyt._array_functions import norm as unyt_linalg_norm -# helper_result = _prepare_array_func_args(x, ord=ord, axis=axis, keepdims=keepdims) -# ret_cf = ...() -# res = unyt_linalg_norm(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(x, ord=ord, axis=axis, keepdims=keepdims) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_linalg_norm(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) @implements(np.vstack) @@ -649,14 +671,14 @@ def fft_ifftshift(x, axes=None): return _return_helper(res, helper_result, ret_cf) -# @implements(np.sort_complex) -# def sort_complex(...): -# from unyt._array_functions import sort_complex as unyt_sort_complex +@implements(np.sort_complex) +def sort_complex(a): + from unyt._array_functions import sort_complex as unyt_sort_complex -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_sort_complex(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_sort_complex(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.isclose) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 6aac44bf..8040908d 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -112,9 +112,10 @@ def test_handled_funcs(self): # "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), # "histogramdd": (ca(np.arange(3)),), "concatenate": (ca(np.eye(3)),), - # "cross": (ca(np.arange(3)), ca(np.arange(3))), - # "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), - # "union1d": (ca(np.arange(3)), ca(np.arange(3))), + "cross": (ca(np.arange(3)), ca(np.arange(3))), + "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), + "union1d": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.norm": (ca(np.arange(3)),), "vstack": (ca(np.arange(3)),), "hstack": (ca(np.arange(3)),), "dstack": (ca(np.arange(3)),), @@ -138,7 +139,7 @@ def test_handled_funcs(self): "fft.irfftn": (ca(np.arange(3)),), "fft.fftshift": (ca(np.arange(3)),), "fft.ifftshift": (ca(np.arange(3)),), - # "sort_complex": (ca(np.arange(3)),), + "sort_complex": (ca(np.arange(3)),), # "linspace": (ca(1), ca(2), 10), # "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), # "geomspace": (ca(1), ca(1), 10), @@ -180,7 +181,6 @@ def test_handled_funcs(self): # "correlate": (ca(np.arange(3)), ca(np.arange(3))), # "tensordot": (ca(np.eye(3)), ca(np.eye(3))), # "unwrap": (ca(np.arange(3)),), - # "linalg.norm": (ca(np.arange(3)),), # "linalg.det": (ca(np.eye(3)),), # "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), # "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), From 914910a9c3a87e42826115f92ff304ebc7f9286d Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 14 Dec 2024 20:39:12 +0000 Subject: [PATCH 10/29] Implement a few more numpy funcs. --- swiftsimio/_array_functions.py | 106 +++++++++++++++++++++++---------- swiftsimio/objects.py | 6 +- tests/test_cosmo_array.py | 8 +-- 3 files changed, 84 insertions(+), 36 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 28ee5440..c2098596 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -8,6 +8,7 @@ _multiply_cosmo_factor, _preserve_cosmo_factor, _reciprocal_cosmo_factor, + _comparison_cosmo_factor, ) _HANDLED_FUNCTIONS = dict() @@ -681,44 +682,72 @@ def sort_complex(a): return _return_helper(res, helper_result, ret_cf) -# @implements(np.isclose) -# def isclose(...): -# from unyt._array_functions import isclose as unyt_isclose +@implements(np.isclose) +def isclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False): + from unyt._array_functions import isclose as unyt_isclose -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_isclose(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + b, + rtol=rtol, + atol=atol, + equal_nan=equal_nan, + ) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=(a, b), + ) + res = unyt_isclose(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.allclose) -# def allclose(...): -# from unyt._array_functions import allclose as unyt_allclose +@implements(np.allclose) +def allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False): + from unyt._array_functions import allclose as unyt_allclose -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_allclose(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + b, + rtol=rtol, + atol=atol, + equal_nan=equal_nan + ) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=(a, b), + ) + res = unyt_allclose(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.array_equal) -# def array_equal(...): -# from unyt._array_functions import array_equal as unyt_array_equal +@implements(np.array_equal) +def array_equal(a1, a2, equal_nan=False): + from unyt._array_functions import array_equal as unyt_array_equal -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_array_equal(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a1, a2, equal_nan=equal_nan) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=(a1, a2), + ) + res = unyt_array_equal(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.array_equiv) -# def array_equiv(...): -# from unyt._array_functions import array_equiv as unyt_array_equiv +@implements(np.array_equiv) +def array_equiv(a1, a2): + from unyt._array_functions import array_equiv as unyt_array_equiv -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_array_equiv(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a1, a2) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=(a1, a2), + ) + res = unyt_array_equiv(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.linspace) @@ -761,11 +790,28 @@ def sort_complex(a): # return _return_helper(res, helper_result, ret_cf, out=out) +# UNYT.PROD HAS A BUG - IF AXIS IS USED DIMENSIONS ARE WRONG # @implements(np.prod) -# def prod(...): +# def prod( +# a, +# axis=None, +# dtype=None, +# out=None, +# keepdims=np._NoValue, +# initial=np._NoValue, +# where=np._NoValue +# ): # from unyt._array_functions import prod as unyt_prod -# helper_result = _prepare_array_func_args(...) +# helper_result = _prepare_array_func_args( +# a, +# axis=axis, +# dtype=dtype, +# out=out, +# keepdims=keepdims, +# initial=initial, +# where=where, +# ) # ret_cf = ...() # res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) # return _return_helper(res, helper_result, ret_cf, out=out) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 7b9f404f..003f351a 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -203,7 +203,8 @@ def _preserve_cosmo_factor(ca_cf1, ca_cf2=None, **kwargs): elif (ca1 and ca2) and (cf1 == cf2): return cf1 # or cf2, they're equal else: - raise RuntimeError("Unexpected state, please report this error on github.") + # not dealing with cosmo_arrays at all + return None def _power_cosmo_factor(ca_cf1, ca_cf2, inputs=None, power=None): @@ -307,7 +308,8 @@ def _return_without_cosmo_factor(ca_cf, ca_cf2=None, inputs=None, zero_compariso # both have cosmo_factor, and they match: pass else: - raise RuntimeError("Unexpected state, please report this error on github.") + # not dealing with cosmo_arrays at all + pass # return without cosmo_factor return None diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 8040908d..f63b3b3b 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -140,6 +140,10 @@ def test_handled_funcs(self): "fft.fftshift": (ca(np.arange(3)),), "fft.ifftshift": (ca(np.arange(3)),), "sort_complex": (ca(np.arange(3)),), + "isclose": (ca(np.arange(3)), ca(np.arange(3))), + "allclose": (ca(np.arange(3)), ca(np.arange(3))), + "array_equal": (ca(np.arange(3)), ca(np.arange(3))), + "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), # "linspace": (ca(1), ca(2), 10), # "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), # "geomspace": (ca(1), ca(1), 10), @@ -160,10 +164,6 @@ def test_handled_funcs(self): # "isin": (ca(np.arange(3)), ca(np.arange(3))), # "in1d": (ca(np.arange(3)), ca(np.arange(3))), # "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), - # "isclose": (ca(np.arange(3)), ca(np.arange(3))), - # "allclose": (ca(np.arange(3)), ca(np.arange(3))), - # "array_equal": (ca(np.arange(3)), ca(np.arange(3))), - # "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), # "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), # "select": ( # [np.arange(3) < 1, np.arange(3) > 1], From a30350a20ee19d7b9e3d47c7f08dde8d1ddc53fd Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 15 Dec 2024 02:22:41 +0000 Subject: [PATCH 11/29] Start implementing histograms. --- swiftsimio/_array_functions.py | 147 +++++++++++++++++++++++++++------ swiftsimio/objects.py | 2 +- tests/test_cosmo_array.py | 2 +- 3 files changed, 126 insertions(+), 25 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index c2098596..f1c309cd 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -233,34 +233,135 @@ def histogram(a, bins=10, range=None, density=None, weights=None): weights=weights ) ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ret_cf_dens = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) counts, bins = unyt_histogram(*helper_result["args"], **helper_result["kwargs"]) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) + ret_cf_counts = ret_cf_w * ret_cf_dens if density else ret_cf_w + else: + ret_cf_counts = ret_cf_dens if density else None + if isinstance(counts, unyt_array): + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf_counts, + compression=helper_result["compression"], + ) return counts, _return_helper(bins, helper_result, ret_cf) -# ND HISTOGRAMS ARE TRICKY - EACH AXIS CAN HAVE DIFFERENT COSMO FACTORS +@implements(np.histogram2d) +def histogram2d(x, y, bins=10, range=None, density=None, weights=None): + from unyt._array_functions import histogram2d as unyt_histogram2d -# @implements(np.histogram2d) -# def histogram2d(x, y, bins=10, range=None, density=None, weights=None): -# from unyt._array_functions import histogram2d as unyt_histogram2d - -# helper_result = _prepare_array_func_args( -# x, -# y, -# bins=bins, -# range=range, -# density=density, -# weights=weights -# ) -# ret_cf_x = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) -# ret_cf_y = _preserve_cosmo_factor(helper_result["ca_cfs"][1]) -# counts, xbins, ybins = unyt_histogram2d( -# *helper_result["args"], **helper_result["kwargs"] -# ) -# return ( -# counts, -# _return_helper(xbins, helper_result, ret_cf), -# _return_helper(ybins, helper_result, ret_cf), -# ) + if range is not None: + xrange, yrange = range + else: + xrange, yrange = None, None + + try: + N = len(bins) + except TypeError: + N = 1 + if N != 2: + xbins = ybins = bins + elif N == 2: + xbins, ybins = bins + helper_result_x = _prepare_array_func_args( + x, + bins=xbins, + range=xrange, + ) + helper_result_y = _prepare_array_func_args( + y, + bins=ybins, + range=yrange, + ) + if not density: + helper_result_w = _prepare_array_func_args(weights=weights) + ret_cf_x = _preserve_cosmo_factor(helper_result_x["ca_cfs"][0]) + ret_cf_y = _preserve_cosmo_factor(helper_result_y["ca_cfs"][0]) + if ( + (helper_result_x["kwargs"]["range"] is None) + and (helper_result_y["kwargs"]["range"] is None) + ): + safe_range = None + else: + safe_range = ( + helper_result_x["kwargs"]["range"], + helper_result_y["kwargs"]["range"], + ) + counts, xbins, ybins = unyt_histogram2d( + helper_result_x["args"][0], + helper_result_y["args"][0], + bins=(helper_result_x["kwargs"]["bins"], helper_result_y["kwargs"]["bins"]), + range=safe_range, + density=density, + weights=helper_result_w["kwargs"]["weights"], + ) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result_w["comoving"], + cosmo_factor=ret_cf_w, + compression=helper_result_w["compression"], + ) + else: # density=True + # now x, y and weights must be compatible because they will combine + # we unpack input to the helper to get everything checked for compatibility + helper_result = _prepare_array_func_args( + x, + y, + xbins=xbins, + ybins=ybins, + xrange=xrange, + yrange=yrange, + weights=weights, + ) + ret_cf_x = _preserve_cosmo_factor(helper_result_x["ca_cfs"][0]) + ret_cf_y = _preserve_cosmo_factor(helper_result_y["ca_cfs"][0]) + if ( + (helper_result["kwargs"]["xrange"] is None) + and (helper_result["kwargs"]["yrange"] is None) + ): + safe_range = None + else: + safe_range = ( + helper_result["kwargs"]["xrange"], + helper_result["kwargs"]["yrange"], + ) + counts, xbins, ybins = unyt_histogram2d( + helper_result["args"][0], + helper_result["args"][1], + bins=(helper_result["kwargs"]["xbins"], helper_result["kwargs"]["ybins"]), + range=safe_range, + density=density, + weights=helper_result["kwargs"]["weights"], + ) + ret_cf_xy = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + ) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) + ret_cf_counts = ret_cf_w / ret_cf_xy + else: + ret_cf_counts = ret_cf_xy ** -1 + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf_counts, + compression=helper_result["compression"], + ) + return ( + counts, + _return_helper(xbins, helper_result_x, ret_cf_x), + _return_helper(ybins, helper_result_y, ret_cf_y), + ) # @implements(np.histogramdd) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 003f351a..59b464fe 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -436,7 +436,7 @@ def _prepare_array_func_args(*args, **kwargs): ret_cm = True if len(set(comps + list(kw_comps.values()))) == 1: # all compressions identical, preserve it - ret_comp = comps[0] + ret_comp = (comps + list(kw_comps.values()))[0] else: # mixed compressions, strip it off ret_comp = None diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index f63b3b3b..e2824cb1 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -109,7 +109,7 @@ def test_handled_funcs(self): "linalg.pinv": (ca(np.eye(3)),), "linalg.svd": (ca(np.eye(3)),), "histogram": (ca(np.arange(3)),), - # "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), + "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), # "histogramdd": (ca(np.arange(3)),), "concatenate": (ca(np.eye(3)),), "cross": (ca(np.arange(3)), ca(np.arange(3))), From c51cd0b056b374d1335ba003de610ed595c30277 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 15 Dec 2024 02:49:23 +0000 Subject: [PATCH 12/29] Implement a few more numpy funcs. --- swiftsimio/_array_functions.py | 81 ++++++++++++++++++---------------- tests/test_cosmo_array.py | 6 +-- 2 files changed, 46 insertions(+), 41 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index f1c309cd..1fedd224 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -9,6 +9,7 @@ _preserve_cosmo_factor, _reciprocal_cosmo_factor, _comparison_cosmo_factor, + _power_cosmo_factor, ) _HANDLED_FUNCTIONS = dict() @@ -891,31 +892,34 @@ def array_equiv(a1, a2): # return _return_helper(res, helper_result, ret_cf, out=out) -# UNYT.PROD HAS A BUG - IF AXIS IS USED DIMENSIONS ARE WRONG -# @implements(np.prod) -# def prod( -# a, -# axis=None, -# dtype=None, -# out=None, -# keepdims=np._NoValue, -# initial=np._NoValue, -# where=np._NoValue -# ): -# from unyt._array_functions import prod as unyt_prod +@implements(np.prod) +def prod( + a, + axis=None, + dtype=None, + out=None, + keepdims=np._NoValue, + initial=np._NoValue, + where=np._NoValue +): + from unyt._array_functions import prod as unyt_prod -# helper_result = _prepare_array_func_args( -# a, -# axis=axis, -# dtype=dtype, -# out=out, -# keepdims=keepdims, -# initial=initial, -# where=where, -# ) -# ret_cf = ...() -# res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + axis=axis, + dtype=dtype, + out=out, + keepdims=keepdims, + initial=initial, + where=where, + ) + res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) + ret_cf = _power_cosmo_factor( + helper_result["ca_cfs"][0], + (False, None), + power=a.size // res.size, + ) + return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.var) @@ -928,14 +932,21 @@ def array_equiv(a1, a2): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.trace) -# def trace(...): -# from unyt._array_functions import trace as unyt_trace +@implements(np.trace) +def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None): + from unyt._array_functions import trace as unyt_trace -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_trace(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + offset=offset, + axis1=axis1, + axis2=axis2, + dtype=dtype, + out=out, + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_trace(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.percentile) @@ -1110,13 +1121,7 @@ def array_equiv(a1, a2): # @implements(np.cumprod) # def cumprod(...): -# from unyt._array_functions import cumprod as unyt_cumprod - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_cumprod(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - +# Omitted because unyt just raises if called. # @implements(np.pad) # def pad(...): diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index e2824cb1..cefe77df 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -147,9 +147,9 @@ def test_handled_funcs(self): # "linspace": (ca(1), ca(2), 10), # "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), # "geomspace": (ca(1), ca(1), 10), - # "prod": (ca(np.arange(3)),), + "prod": (ca(np.arange(3)),), # "var": (ca(np.arange(3)),), - # "trace": (ca(np.eye(3)),), + "trace": (ca(np.eye(3)),), # "percentile": (ca(np.arange(3)), 30), # "quantile": (ca(np.arange(3)), 0.3), # "nanpercentile": (ca(np.arange(3)), 30), @@ -157,7 +157,7 @@ def test_handled_funcs(self): # "diff": (ca(np.arange(3)),), # "ediff1d": (ca(np.arange(3)),), # "ptp": (ca(np.arange(3)),), - # "cumprod": (ca(np.arange(3)),), + "cumprod": (ca(np.arange(3)),), # "pad": (ca(np.arange(3)), 3), # "choose": (np.arange(3), ca(np.eye(3))), # "insert": (ca(np.arange(3)), 1, ca(1)), From f27c5277b2c3220b1be05b7015bd0d5b172d9c2a Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 16 Dec 2024 17:09:27 +0000 Subject: [PATCH 13/29] Implement linalg.det. --- swiftsimio/_array_functions.py | 18 +++++++++++------- tests/test_cosmo_array.py | 2 +- 2 files changed, 12 insertions(+), 8 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 1fedd224..927db375 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -989,14 +989,18 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.linalg.det) -# def linalg_det(...): -# from unyt._array_functions import linalg_det as unyt_linalg_det +@implements(np.linalg.det) +def linalg_det(a): + from unyt._array_functions import linalg_det as unyt_linalg_det -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_det(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a) + ret_cf = _power_cosmo_factor( + helper_result["ca_cfs"][0], + (False, None), + power=a.shape[0], + ) + res = unyt_linalg_det(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.linalg.lstsq) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index cefe77df..64627898 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -154,6 +154,7 @@ def test_handled_funcs(self): # "quantile": (ca(np.arange(3)), 0.3), # "nanpercentile": (ca(np.arange(3)), 30), # "nanquantile": (ca(np.arange(3)), 0.3), + "linalg.det": (ca(np.eye(3)),), # "diff": (ca(np.arange(3)),), # "ediff1d": (ca(np.arange(3)),), # "ptp": (ca(np.arange(3)),), @@ -181,7 +182,6 @@ def test_handled_funcs(self): # "correlate": (ca(np.arange(3)), ca(np.arange(3))), # "tensordot": (ca(np.eye(3)), ca(np.eye(3))), # "unwrap": (ca(np.arange(3)),), - # "linalg.det": (ca(np.eye(3)),), # "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), # "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), # "linalg.tensorsolve": ( From 21585c55e7c21b30216d227956afda645223c362 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 16 Dec 2024 17:43:04 +0000 Subject: [PATCH 14/29] Attempt histogramdd implementation. --- swiftsimio/_array_functions.py | 186 +++++++++++++++++++++++---------- tests/test_cosmo_array.py | 12 +-- 2 files changed, 139 insertions(+), 59 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 927db375..09b03452 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -209,10 +209,7 @@ def linalg_svd(a, full_matrices=True, compute_uv=True, hermitian=False): from unyt._array_functions import linalg_svd as unyt_linalg_svd helper_result = _prepare_array_func_args( - a, - full_matrices=full_matrices, - compute_uv=compute_uv, - hermitian=hermitian + a, full_matrices=full_matrices, compute_uv=compute_uv, hermitian=hermitian ) ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) ress = unyt_linalg_svd(*helper_result["args"], **helper_result["kwargs"]) @@ -227,11 +224,7 @@ def histogram(a, bins=10, range=None, density=None, weights=None): from unyt._array_functions import histogram as unyt_histogram helper_result = _prepare_array_func_args( - a, - bins=bins, - range=range, - density=density, - weights=weights + a, bins=bins, range=range, density=density, weights=weights ) ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) ret_cf_dens = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) @@ -283,9 +276,8 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): helper_result_w = _prepare_array_func_args(weights=weights) ret_cf_x = _preserve_cosmo_factor(helper_result_x["ca_cfs"][0]) ret_cf_y = _preserve_cosmo_factor(helper_result_y["ca_cfs"][0]) - if ( - (helper_result_x["kwargs"]["range"] is None) - and (helper_result_y["kwargs"]["range"] is None) + if (helper_result_x["kwargs"]["range"] is None) and ( + helper_result_y["kwargs"]["range"] is None ): safe_range = None else: @@ -302,6 +294,9 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): weights=helper_result_w["kwargs"]["weights"], ) if weights is not None: + raise NotImplementedError( + "Need to handle numpy array weights here and stress-test histogram2d." + ) ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) counts = cosmo_array( counts.to_value(counts.units), @@ -324,9 +319,8 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): ) ret_cf_x = _preserve_cosmo_factor(helper_result_x["ca_cfs"][0]) ret_cf_y = _preserve_cosmo_factor(helper_result_y["ca_cfs"][0]) - if ( - (helper_result["kwargs"]["xrange"] is None) - and (helper_result["kwargs"]["yrange"] is None) + if (helper_result["kwargs"]["xrange"] is None) and ( + helper_result["kwargs"]["yrange"] is None ): safe_range = None else: @@ -350,7 +344,7 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) ret_cf_counts = ret_cf_w / ret_cf_xy else: - ret_cf_counts = ret_cf_xy ** -1 + ret_cf_counts = ret_cf_xy**-1 counts = cosmo_array( counts.to_value(counts.units), counts.units, @@ -365,20 +359,114 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): ) -# @implements(np.histogramdd) -# def histogramdd(sample, bins=10, range=None, density=None, weights=None): -# from unyt._array_functions import histogramdd as unyt_histogramdd +@implements(np.histogramdd) +def histogramdd(sample, bins=10, range=None, density=None, weights=None): + from unyt._array_functions import histogramdd as unyt_histogramdd -# helper_result = _prepare_array_func_args( -# sample, -# bins=bins, -# range=range, -# density=density, -# weights=weights -# ) -# ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) -# counts, bins = unyt_histogramdd(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + N, D = sample.shape + if range is not None: + ranges = range + else: + ranges = D * [None] + + try: + len(bins) + except TypeError: + # bins is an integer + bins = D * [bins] + helper_results = [ + _prepare_array_func_args( + s, + bins=b, + range=r, + ) + for s, b, r in zip(sample, bins, ranges) + ] + if not density: + helper_result_w = _prepare_array_func_args(weights=weights) + ret_cfs = [ + _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + for helper_result in helper_results + ] + if all( + [ + helper_result["kwargs"]["range"] is None + for helper_result in helper_results + ] + ): + safe_range = None + else: + safe_range = [ + helper_result["kwargs"]["range"] for helper_result in helper_results + ] + counts, bins = unyt_histogramdd( + [helper_result["args"][0] for helper_result in helper_results], + bins=[helper_result["kwargs"]["bins"] for helper_result in helper_results], + range=safe_range, + density=density, + weights=helper_result_w["kwargs"]["weights"], + ) + if weights is not None: + raise NotImplementedError( + "Need to handle numpy array weights here and stress-test histogramdd." + ) + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result_w["comoving"], + cosmo_factor=ret_cf_w, + compression=helper_result_w["compression"], + ) + else: # density=True + # now sample and weights must be compatible because they will combine + # we unpack input to the helper to get everything checked for compatibility + helper_result = _prepare_array_func_args( + sample, + bins=bins, + range=range, + weights=weights, + ) + helper_results = D * [helper_result] + ret_cfs = D * [_preserve_cosmo_factor(helper_result["ca_cfs"][0])] + counts, bins = unyt_histogramdd( + helper_result["args"][0], + bins=helper_result["kwargs"]["bins"], + range=helper_result["kwargs"]["range"], + density=density, + weights=helper_result["kwargs"]["weights"], + ) + if len(helper_result["ca_cfs"]) == 1: + ret_cf_sample = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + else: + ret_cf_sample = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + ) + if len(helper_result["ca_cfs"]) > 2: + for hr_ca_cfs_i in helper_result["ca_cfs"][2:]: + ret_cf_sample = _multiply_cosmo_factor( + (True, ret_cf_sample), hr_ca_cfs_i + ) + if weights is not None: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) + ret_cf_counts = ret_cf_w / ret_cf_sample + else: + ret_cf_counts = ret_cf_sample**-1 + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf_counts, + compression=helper_result["compression"], + ) + return ( + counts, + [ + _return_helper(b, helper_result, ret_cf) + for b, helper_result, ret_cf in zip(bins, helper_results, ret_cfs) + ], + ) @implements(np.concatenate) @@ -422,10 +510,7 @@ def intersect1d(ar1, ar2, assume_unique=False, return_indices=False): from unyt._array_functions import intersect1d as unyt_intersect1d helper_result = _prepare_array_func_args( - ar1, - ar2, - assume_unique=assume_unique, - return_indices=return_indices + ar1, ar2, assume_unique=assume_unique, return_indices=return_indices ) ret_cf = _preserve_cosmo_factor( helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] @@ -464,9 +549,7 @@ def linalg_norm(x, ord=None, axis=None, keepdims=False): def vstack(tup, *, dtype=None, casting="same_kind"): from unyt._array_functions import vstack as unyt_vstack - helper_result = _prepare_array_func_args( - tup, dtype=dtype, casting=casting - ) + helper_result = _prepare_array_func_args(tup, dtype=dtype, casting=casting) helper_result_concat_items = _prepare_array_func_args(*tup) ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) res = unyt_vstack( @@ -481,9 +564,7 @@ def vstack(tup, *, dtype=None, casting="same_kind"): def hstack(tup, *, dtype=None, casting="same_kind"): from unyt._array_functions import hstack as unyt_hstack - helper_result = _prepare_array_func_args( - tup, dtype=dtype, casting=casting - ) + helper_result = _prepare_array_func_args(tup, dtype=dtype, casting=casting) helper_result_concat_items = _prepare_array_func_args(*tup) ret_cf = _preserve_cosmo_factor(helper_result_concat_items["ca_cfs"][0]) res = unyt_hstack( @@ -573,8 +654,10 @@ def _prepare_array_block_args(lst, recursing=False): elif all([cm is None for cm in cms]): ret_cm = None elif any([cm is None for cm in cms]) and not all([cm is None for cm in cms]): - raise ValueError("Some input has comoving=None and others have " - "comoving=True|False. Result is undefined!") + raise ValueError( + "Some input has comoving=None and others have " + "comoving=True|False. Result is undefined!" + ) elif all([cm is False for cm in cms]): ret_cm = False else: @@ -614,6 +697,7 @@ def block(arrays): # UNYT HAS A COPY-PASTED TYPO fft -> ftt + @implements(np.fft.fft) def ftt_fft(a, n=None, axis=-1, norm=None, out=None): from unyt._array_functions import ftt_fft as unyt_fft_fft @@ -809,11 +893,7 @@ def allclose(a, b, rtol=1e-05, atol=1e-08, equal_nan=False): from unyt._array_functions import allclose as unyt_allclose helper_result = _prepare_array_func_args( - a, - b, - rtol=rtol, - atol=atol, - equal_nan=equal_nan + a, b, rtol=rtol, atol=atol, equal_nan=equal_nan ) ret_cf = _comparison_cosmo_factor( helper_result["ca_cfs"][0], @@ -894,13 +974,13 @@ def array_equiv(a1, a2): @implements(np.prod) def prod( - a, - axis=None, - dtype=None, - out=None, - keepdims=np._NoValue, - initial=np._NoValue, - where=np._NoValue + a, + axis=None, + dtype=None, + out=None, + keepdims=np._NoValue, + initial=np._NoValue, + where=np._NoValue, ): from unyt._array_functions import prod as unyt_prod diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 64627898..f30efca6 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -19,7 +19,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) def to_ua(x): @@ -49,7 +49,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -60,7 +60,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -70,7 +70,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -80,7 +80,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -110,7 +110,7 @@ def test_handled_funcs(self): "linalg.svd": (ca(np.eye(3)),), "histogram": (ca(np.arange(3)),), "histogram2d": (ca(np.arange(3)), ca(np.arange(3))), - # "histogramdd": (ca(np.arange(3)),), + "histogramdd": (ca(np.arange(3)).reshape((1, 3)),), "concatenate": (ca(np.eye(3)),), "cross": (ca(np.arange(3)), ca(np.arange(3))), "intersect1d": (ca(np.arange(3)), ca(np.arange(3))), From daa2196a1a9800737d2a0288babdec6c7e748f24 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 17 Dec 2024 17:54:25 +0000 Subject: [PATCH 15/29] Wrap up histogram testing. --- swiftsimio/_array_functions.py | 101 ++++++++++-------- swiftsimio/objects.py | 21 ++-- tests/test_cosmo_array.py | 183 +++++++++++++++++++++++++++++++++ 3 files changed, 255 insertions(+), 50 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 09b03452..dac6a642 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -231,7 +231,13 @@ def histogram(a, bins=10, range=None, density=None, weights=None): counts, bins = unyt_histogram(*helper_result["args"], **helper_result["kwargs"]) if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) - ret_cf_counts = ret_cf_w * ret_cf_dens if density else ret_cf_w + ret_cf_counts = ( + _multiply_cosmo_factor( + (ret_cf_w is not None, ret_cf_w), (ret_cf_dens is not None, ret_cf_dens) + ) + if density + else ret_cf_w + ) else: ret_cf_counts = ret_cf_dens if density else None if isinstance(counts, unyt_array): @@ -294,17 +300,15 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): weights=helper_result_w["kwargs"]["weights"], ) if weights is not None: - raise NotImplementedError( - "Need to handle numpy array weights here and stress-test histogram2d." - ) ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) - counts = cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result_w["comoving"], - cosmo_factor=ret_cf_w, - compression=helper_result_w["compression"], - ) + if isinstance(counts, unyt_array): + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result_w["comoving"], + cosmo_factor=ret_cf_w, + compression=helper_result_w["compression"], + ) else: # density=True # now x, y and weights must be compatible because they will combine # we unpack input to the helper to get everything checked for compatibility @@ -342,16 +346,21 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): ) if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) - ret_cf_counts = ret_cf_w / ret_cf_xy + inv_ret_cf_xy = _reciprocal_cosmo_factor((ret_cf_xy is not None, ret_cf_xy)) + ret_cf_counts = _multiply_cosmo_factor( + (ret_cf_w is not None, ret_cf_w), + (inv_ret_cf_xy is not None, inv_ret_cf_xy), + ) else: - ret_cf_counts = ret_cf_xy**-1 - counts = cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf_counts, - compression=helper_result["compression"], - ) + ret_cf_counts = _reciprocal_cosmo_factor((ret_cf_xy is not None, ret_cf_xy)) + if isinstance(counts, unyt_array): + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf_counts, + compression=helper_result["compression"], + ) return ( counts, _return_helper(xbins, helper_result_x, ret_cf_x), @@ -363,7 +372,7 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): def histogramdd(sample, bins=10, range=None, density=None, weights=None): from unyt._array_functions import histogramdd as unyt_histogramdd - N, D = sample.shape + D = len(sample) if range is not None: ranges = range else: @@ -407,30 +416,27 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): weights=helper_result_w["kwargs"]["weights"], ) if weights is not None: - raise NotImplementedError( - "Need to handle numpy array weights here and stress-test histogramdd." - ) ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) - counts = cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result_w["comoving"], - cosmo_factor=ret_cf_w, - compression=helper_result_w["compression"], - ) + if isinstance(counts, unyt_array): + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result_w["comoving"], + cosmo_factor=ret_cf_w, + compression=helper_result_w["compression"], + ) else: # density=True # now sample and weights must be compatible because they will combine # we unpack input to the helper to get everything checked for compatibility helper_result = _prepare_array_func_args( - sample, + *sample, bins=bins, range=range, weights=weights, ) - helper_results = D * [helper_result] ret_cfs = D * [_preserve_cosmo_factor(helper_result["ca_cfs"][0])] counts, bins = unyt_histogramdd( - helper_result["args"][0], + helper_result["args"], bins=helper_result["kwargs"]["bins"], range=helper_result["kwargs"]["range"], density=density, @@ -450,16 +456,25 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): ) if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) - ret_cf_counts = ret_cf_w / ret_cf_sample + inv_ret_cf_sample = _reciprocal_cosmo_factor( + (ret_cf_sample is not None, ret_cf_sample) + ) + ret_cf_counts = _multiply_cosmo_factor( + (ret_cf_w is not None, ret_cf_w), + (inv_ret_cf_sample is not None, inv_ret_cf_sample), + ) else: - ret_cf_counts = ret_cf_sample**-1 - counts = cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf_counts, - compression=helper_result["compression"], - ) + ret_cf_counts = _reciprocal_cosmo_factor( + (ret_cf_sample is not None, ret_cf_sample) + ) + if isinstance(counts, unyt_array): + counts = cosmo_array( + counts.to_value(counts.units), + counts.units, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf_counts, + compression=helper_result["compression"], + ) return ( counts, [ diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 59b464fe..f4ce1861 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -335,7 +335,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) + return cosmo_factor(a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -408,7 +408,10 @@ def _prepare_array_func_args(*args, **kwargs): k: (hasattr(kwarg, "compression"), getattr(kwarg, "compression", None)) for k, kwarg in kwargs.items() } - if all([cm[1] for cm in cms + list(kw_cms.values()) if cm[0]]): + if len([cm[1] for cm in cms + list(kw_cms.values()) if cm[0]]) == 0: + # no cosmo inputs + ret_cm = None + elif all([cm[1] for cm in cms + list(kw_cms.values()) if cm[0]]): # all cosmo inputs are comoving ret_cm = True elif all([cm[1] is None for cm in cms + list(kw_cms.values()) if cm[0]]): @@ -440,6 +443,10 @@ def _prepare_array_func_args(*args, **kwargs): else: # mixed compressions, strip it off ret_comp = None + args = [unyt_array(arg) if isinstance(arg, cosmo_array) else arg for arg in args] + kwargs = { + k: unyt_array(v) if isinstance(v, cosmo_array) else v for k, v in kwargs.items() + } return dict( args=args, kwargs=kwargs, @@ -648,7 +655,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor @@ -1163,7 +1170,7 @@ def from_pint( def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): helper_result = _prepare_array_func_args(*inputs, **kwargs) - cfs = helper_result["cfs"] + ca_cfs = helper_result["ca_cfs"] # make sure we evaluate the cosmo_factor_ufunc_registry function: # might raise/warn even if we're not returning a cosmo_array @@ -1171,16 +1178,16 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): power_map = POWER_MAPPING[ufunc] if "axis" in kwargs and kwargs["axis"] is not None: ret_cf = _power_cosmo_factor( - cfs[0], + ca_cfs[0], (False, None), power=power_map(inputs[0].shape[kwargs["axis"]]), ) else: ret_cf = _power_cosmo_factor( - cfs[0], (False, None), power=power_map(inputs[0].size) + ca_cfs[0], (False, None), power=power_map(inputs[0].size) ) else: - ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs) + ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*ca_cfs, inputs=inputs) ret = super().__array_ufunc__(ufunc, method, *inputs, **kwargs) # if we get a tuple we have multiple return values to deal with diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index f30efca6..a16a3dbc 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -259,6 +259,189 @@ def test_handled_funcs(self): ], ) + # the combinations of units and cosmo_factors is nonsense but it's just for testing... + @pytest.mark.parametrize( + "func_args", + ( + ( + np.histogram, + ( + cosmo_array( + [1, 2, 3], + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a**1, 1.0), + ), + ), + ), + ( + np.histogram2d, + ( + cosmo_array( + [1, 2, 3], + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a**1, 1.0), + ), + cosmo_array( + [1, 2, 3], + u.K, + comoving=False, + cosmo_factor=cosmo_factor(a**2, 1.0), + ), + ), + ), + ( + np.histogramdd, + ( + [ + cosmo_array( + [1, 2, 3], + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a**1, 1.0), + ), + cosmo_array( + [1, 2, 3], + u.K, + comoving=False, + cosmo_factor=cosmo_factor(a**2, 1.0), + ), + cosmo_array( + [1, 2, 3], + u.kg, + comoving=False, + cosmo_factor=cosmo_factor(a**3, 1.0), + ), + ], + ), + ), + ), + ) + @pytest.mark.parametrize( + "weights", + ( + None, + cosmo_array( + [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a**1, 1.0) + ), + np.array([1, 2, 3]), + ), + ) + @pytest.mark.parametrize("bins", ("int", "np", "ca")) + @pytest.mark.parametrize("density", (None, True)) + def test_histograms(self, func_args, weights, bins, density): + func, args = func_args + _bins = { + "int": 10, + "np": [np.linspace(0, 5, 11)] * 3, + "ca": [ + cosmo_array( + np.linspace(0, 5, 11), + u.kpc, + comoving=False, + cosmo_factor=cosmo_factor(a**1, 1.0), + ), + cosmo_array( + np.linspace(0, 5, 11), + u.K, + comoving=False, + cosmo_factor=cosmo_factor(a**2, 1.0), + ), + cosmo_array( + np.linspace(0, 5, 11), + u.Msun, + comoving=False, + cosmo_factor=cosmo_factor(a**3, 1.0), + ), + ], + }[bins] + bins = ( + _bins[ + { + np.histogram: np.s_[0], + np.histogram2d: np.s_[:2], + np.histogramdd: np.s_[:], + }[func] + ] + if bins in ("np", "ca") + else _bins + ) + result = func(*args, bins=bins, density=density, weights=weights) + ua_args = tuple( + ( + to_ua(arg) + if not isinstance(arg, tuple) + else tuple(to_ua(item) for item in arg) + ) + for arg in args + ) + ua_bins = ( + to_ua(bins) + if not isinstance(bins, tuple) + else tuple(to_ua(item) for item in bins) + ) + ua_result = func( + *ua_args, bins=ua_bins, density=density, weights=to_ua(weights) + ) + if isinstance(ua_result, tuple): + assert isinstance(result, tuple) + assert len(result) == len(ua_result) + for r, ua_r in zip(result, ua_result): + check_result(r, ua_r) + else: + check_result(result, ua_result) + if density is None and not isinstance(weights, cosmo_array): + assert not hasattr(result[0], "comoving") + else: + assert result[0].comoving is False + if density and not isinstance(weights, cosmo_array): + assert ( + result[0].cosmo_factor + == { + np.histogram: cosmo_factor(a**-1, 1.0), + np.histogram2d: cosmo_factor(a**-3, 1.0), + np.histogramdd: cosmo_factor(a**-6, 1.0), + }[func] + ) + elif density and isinstance(weights, cosmo_array): + assert result[0].comoving is False + assert ( + result[0].cosmo_factor + == { + np.histogram: cosmo_factor(a**0, 1.0), + np.histogram2d: cosmo_factor(a**-2, 1.0), + np.histogramdd: cosmo_factor(a**-5, 1.0), + }[func] + ) + elif density is None and isinstance(weights, cosmo_array): + assert result[0].comoving is False + assert ( + result[0].cosmo_factor + == { + np.histogram: cosmo_factor(a**1, 1.0), + np.histogram2d: cosmo_factor(a**1, 1.0), + np.histogramdd: cosmo_factor(a**1, 1.0), + }[func] + ) + ret_bins = { + np.histogram: [result[1]], + np.histogram2d: result[1:], + np.histogramdd: result[1], + }[func] + for b, expt_cf in zip( + ret_bins, + ( + [ + cosmo_factor(a**1, 1.0), + cosmo_factor(a**2, 1.0), + cosmo_factor(a**3, 1.0), + ] + ), + ): + assert b.comoving is False + assert b.cosmo_factor == expt_cf + def test_getitem(self): assert isinstance(ca(np.arange(3))[0], cosmo_quantity) From 9f5f6e544084a86421024214cd5270fae2959cef Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 17 Dec 2024 17:56:41 +0000 Subject: [PATCH 16/29] Amend physical conversion test because array_equal is now cosmo_array-aware. --- tests/test_physical_conversion.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/tests/test_physical_conversion.py b/tests/test_physical_conversion.py index 26b72d29..2ade790f 100644 --- a/tests/test_physical_conversion.py +++ b/tests/test_physical_conversion.py @@ -10,7 +10,12 @@ def test_convert(filename): """ data = load(filename) coords = data.gas.coordinates + units = coords.units coords_physical = coords.to_physical() - assert array_equal(coords * data.metadata.a, coords_physical) + # array_equal applied to cosmo_array's is aware of physical & comoving + # make sure to compare bare arrays: + assert array_equal( + coords.to_value(units) * data.metadata.a, coords_physical.to_value(units) + ) return From 16e40cce9e8e5567b8b10c5b8fd75756de182e32 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 17 Dec 2024 18:14:59 +0000 Subject: [PATCH 17/29] Implement a couple more numpy functions. --- swiftsimio/_array_functions.py | 30 ++++++++++++++++-------------- swiftsimio/objects.py | 5 +---- tests/test_cosmo_array.py | 4 ++-- 3 files changed, 19 insertions(+), 20 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index dac6a642..1a5219b3 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1188,24 +1188,26 @@ def linalg_det(a): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.diff) -# def diff(...): -# from unyt._array_functions import diff as unyt_diff +@implements(np.diff) +def diff(a, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue): + from unyt._array_functions import diff as unyt_diff -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_diff(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, n=n, axis=axis, prepend=prepend, append=append + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_diff(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.ediff1d) -# def ediff1d(...): -# from unyt._array_functions import ediff1d as unyt_ediff1d +@implements(np.ediff1d) +def ediff1d(ary, to_end=None, to_begin=None): + from unyt._array_functions import ediff1d as unyt_ediff1d -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_ediff1d(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(ary, to_end=to_end, to_begin=to_begin) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_ediff1d(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.ptp) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index f4ce1861..d6c60d29 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -443,10 +443,7 @@ def _prepare_array_func_args(*args, **kwargs): else: # mixed compressions, strip it off ret_comp = None - args = [unyt_array(arg) if isinstance(arg, cosmo_array) else arg for arg in args] - kwargs = { - k: unyt_array(v) if isinstance(v, cosmo_array) else v for k, v in kwargs.items() - } + # WE SHOULD COMPLAIN HERE IF WE HAVE DIFFERENT SCALE FACTORS IN COSMO_FACTOR'S?? return dict( args=args, kwargs=kwargs, diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index a16a3dbc..112876be 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -155,8 +155,8 @@ def test_handled_funcs(self): # "nanpercentile": (ca(np.arange(3)), 30), # "nanquantile": (ca(np.arange(3)), 0.3), "linalg.det": (ca(np.eye(3)),), - # "diff": (ca(np.arange(3)),), - # "ediff1d": (ca(np.arange(3)),), + "diff": (ca(np.arange(3)),), + "ediff1d": (ca(np.arange(3)),), # "ptp": (ca(np.arange(3)),), "cumprod": (ca(np.arange(3)),), # "pad": (ca(np.arange(3)), 3), From 240b99172cf378b6d87f49be95491f39418459f7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 18 Dec 2024 08:59:12 +0000 Subject: [PATCH 18/29] Check cosmo_factor for consistent scale_factor when preparing arguments. --- swiftsimio/objects.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index d6c60d29..d77430d4 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -335,7 +335,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a**0, scale_factor=cf1.scale_factor) + return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -443,7 +443,14 @@ def _prepare_array_func_args(*args, **kwargs): else: # mixed compressions, strip it off ret_comp = None - # WE SHOULD COMPLAIN HERE IF WE HAVE DIFFERENT SCALE FACTORS IN COSMO_FACTOR'S?? + expected_scale_factor = None + for ca_cf in ca_cfs + list(kw_ca_cfs.values()): + if ca_cf[0]: + if expected_scale_factor is None: + expected_scale_factor = ca_cf[1].scale_factor + else: + if ca_cf[1].scale_factor != expected_scale_factor: + raise ValueError("Mismatched scale factors in cosmo_array's.") return dict( args=args, kwargs=kwargs, @@ -652,7 +659,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor From 8880a13e482e7de93bea41d62bfe14fa6a61e825 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 18 Dec 2024 09:31:13 +0000 Subject: [PATCH 19/29] Implement a few more numpy functions. --- swiftsimio/_array_functions.py | 213 ++++++++++++++++++++++++++------- tests/test_cosmo_array.py | 66 +++++----- 2 files changed, 204 insertions(+), 75 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 1a5219b3..6c844a73 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -947,14 +947,38 @@ def array_equiv(a1, a2): return _return_helper(res, helper_result, ret_cf) -# @implements(np.linspace) -# def linspace(...): -# from unyt._array_functions import linspace as unyt_linspace +@implements(np.linspace) +def linspace( + start, + stop, + num=50, + endpoint=True, + retstep=False, + dtype=None, + axis=0, + *, + device=None, +): + from unyt._array_functions import linspace as unyt_linspace -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linspace(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + start, + stop, + num=num, + endpoint=endpoint, + retstep=retstep, + dtype=dtype, + axis=axis, + device=device, + ) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + ress = unyt_linspace(*helper_result["args"], **helper_result["kwargs"]) + if retstep: + return tuple(_return_helper(res, helper_result, ret_cf) for res in ress) + else: + return _return_helper(ress, helper_result, ret_cf) # @implements(np.logspace) @@ -1017,14 +1041,35 @@ def prod( return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.var) -# def var(...): -# from unyt._array_functions import var as unyt_var +@implements(np.var) +def var( + a, + axis=None, + dtype=None, + out=None, + ddof=0, + keepdims=np._NoValue, + *, + where=np._NoValue, + mean=np._NoValue, + correction=np._NoValue +): + from unyt._array_functions import var as unyt_var -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_var(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + axis=axis, + dtype=dtype, + out=out, + ddof=ddof, + keepdims=keepdims, + where=where, + mean=mean, + correction=correction, + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_var(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) @implements(np.trace) @@ -1044,44 +1089,128 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None): return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.percentile) -# def percentile(...): -# from unyt._array_functions import percentile as unyt_percentile +@implements(np.percentile) +def percentile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None +): + from unyt._array_functions import percentile as unyt_percentile -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_percentile(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + q, + axis=axis, + out=out, + overwrite_input=overwrite_input, + method=method, + keepdims=keepdims, + weights=weights, + interpolation=interpolation, + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_percentile(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.quantile) -# def quantile(...): -# from unyt._array_functions import quantile as unyt_quantile +@implements(np.quantile) +def quantile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method='linear', + keepdims=False, + *, + weights=None, + interpolation=None +): + from unyt._array_functions import quantile as unyt_quantile -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_quantile(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + q, + axis=axis, + out=out, + overwrite_input=overwrite_input, + method=method, + keepdims=keepdims, + weights=weights, + interpolation=interpolation, + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_quantile(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.nanpercentile) -# def nanpercentile(...): -# from unyt._array_functions import nanpercentile as unyt_nanpercentile +@implements(np.nanpercentile) +def percentile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None +): + from unyt._array_functions import nanpercentile as unyt_nanpercentile -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_nanpercentile(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + q, + axis=axis, + out=out, + overwrite_input=overwrite_input, + method=method, + keepdims=keepdims, + weights=weights, + interpolation=interpolation, + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_nanpercentile(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.nanquantile) -# def nanquantile(...): -# from unyt._array_functions import nanquantile as unyt_nanquantile +@implements(np.nanquantile) +def nanquantile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method='linear', + keepdims=False, + *, + weights=None, + interpolation=None +): + from unyt._array_functions import nanquantile as unyt_nanquantile -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_nanquantile(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + a, + q, + axis=axis, + out=out, + overwrite_input=overwrite_input, + method=method, + keepdims=keepdims, + weights=weights, + interpolation=interpolation, + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_nanquantile(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) @implements(np.linalg.det) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 112876be..42644c0a 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -19,7 +19,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) def to_ua(x): @@ -49,7 +49,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -60,7 +60,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -70,7 +70,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -80,7 +80,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -144,16 +144,16 @@ def test_handled_funcs(self): "allclose": (ca(np.arange(3)), ca(np.arange(3))), "array_equal": (ca(np.arange(3)), ca(np.arange(3))), "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), - # "linspace": (ca(1), ca(2), 10), + "linspace": (ca(1), ca(2), 10), # "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), # "geomspace": (ca(1), ca(1), 10), "prod": (ca(np.arange(3)),), - # "var": (ca(np.arange(3)),), + "var": (ca(np.arange(3)),), "trace": (ca(np.eye(3)),), - # "percentile": (ca(np.arange(3)), 30), - # "quantile": (ca(np.arange(3)), 0.3), - # "nanpercentile": (ca(np.arange(3)), 30), - # "nanquantile": (ca(np.arange(3)), 0.3), + "percentile": (ca(np.arange(3)), 30), + "quantile": (ca(np.arange(3)), 0.3), + "nanpercentile": (ca(np.arange(3)), 30), + "nanquantile": (ca(np.arange(3)), 0.3), "linalg.det": (ca(np.eye(3)),), "diff": (ca(np.arange(3)),), "ediff1d": (ca(np.arange(3)),), @@ -270,7 +270,7 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), ), ), @@ -281,13 +281,13 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a**2, 1.0), + cosmo_factor=cosmo_factor(a ** 2, 1.0), ), ), ), @@ -299,19 +299,19 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a**2, 1.0), + cosmo_factor=cosmo_factor(a ** 2, 1.0), ), cosmo_array( [1, 2, 3], u.kg, comoving=False, - cosmo_factor=cosmo_factor(a**3, 1.0), + cosmo_factor=cosmo_factor(a ** 3, 1.0), ), ], ), @@ -323,7 +323,7 @@ def test_handled_funcs(self): ( None, cosmo_array( - [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a**1, 1.0) + [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 1.0) ), np.array([1, 2, 3]), ), @@ -340,19 +340,19 @@ def test_histograms(self, func_args, weights, bins, density): np.linspace(0, 5, 11), u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.K, comoving=False, - cosmo_factor=cosmo_factor(a**2, 1.0), + cosmo_factor=cosmo_factor(a ** 2, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.Msun, comoving=False, - cosmo_factor=cosmo_factor(a**3, 1.0), + cosmo_factor=cosmo_factor(a ** 3, 1.0), ), ], }[bins] @@ -399,9 +399,9 @@ def test_histograms(self, func_args, weights, bins, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a**-1, 1.0), - np.histogram2d: cosmo_factor(a**-3, 1.0), - np.histogramdd: cosmo_factor(a**-6, 1.0), + np.histogram: cosmo_factor(a ** -1, 1.0), + np.histogram2d: cosmo_factor(a ** -3, 1.0), + np.histogramdd: cosmo_factor(a ** -6, 1.0), }[func] ) elif density and isinstance(weights, cosmo_array): @@ -409,9 +409,9 @@ def test_histograms(self, func_args, weights, bins, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a**0, 1.0), - np.histogram2d: cosmo_factor(a**-2, 1.0), - np.histogramdd: cosmo_factor(a**-5, 1.0), + np.histogram: cosmo_factor(a ** 0, 1.0), + np.histogram2d: cosmo_factor(a ** -2, 1.0), + np.histogramdd: cosmo_factor(a ** -5, 1.0), }[func] ) elif density is None and isinstance(weights, cosmo_array): @@ -419,9 +419,9 @@ def test_histograms(self, func_args, weights, bins, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a**1, 1.0), - np.histogram2d: cosmo_factor(a**1, 1.0), - np.histogramdd: cosmo_factor(a**1, 1.0), + np.histogram: cosmo_factor(a ** 1, 1.0), + np.histogram2d: cosmo_factor(a ** 1, 1.0), + np.histogramdd: cosmo_factor(a ** 1, 1.0), }[func] ) ret_bins = { @@ -433,9 +433,9 @@ def test_histograms(self, func_args, weights, bins, density): ret_bins, ( [ - cosmo_factor(a**1, 1.0), - cosmo_factor(a**2, 1.0), - cosmo_factor(a**3, 1.0), + cosmo_factor(a ** 1, 1.0), + cosmo_factor(a ** 2, 1.0), + cosmo_factor(a ** 3, 1.0), ] ), ): From 30f4e35c91412135a0a064e08d4114bbf5152063 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 18 Dec 2024 12:43:06 +0000 Subject: [PATCH 20/29] Implement a few more numpy functions. --- swiftsimio/_array_functions.py | 327 +++++++++++++++++---------------- swiftsimio/objects.py | 36 ++-- tests/test_cosmo_array.py | 94 +++++----- 3 files changed, 241 insertions(+), 216 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 6c844a73..f345fe1c 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -226,7 +226,7 @@ def histogram(a, bins=10, range=None, density=None, weights=None): helper_result = _prepare_array_func_args( a, bins=bins, range=range, density=density, weights=weights ) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ret_cf_bins = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) ret_cf_dens = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) counts, bins = unyt_histogram(*helper_result["args"], **helper_result["kwargs"]) if weights is not None: @@ -248,7 +248,7 @@ def histogram(a, bins=10, range=None, density=None, weights=None): cosmo_factor=ret_cf_counts, compression=helper_result["compression"], ) - return counts, _return_helper(bins, helper_result, ret_cf) + return counts, _return_helper(bins, helper_result, ret_cf_bins) @implements(np.histogram2d) @@ -445,15 +445,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): if len(helper_result["ca_cfs"]) == 1: ret_cf_sample = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) else: - ret_cf_sample = _multiply_cosmo_factor( - helper_result["ca_cfs"][0], - helper_result["ca_cfs"][1], - ) - if len(helper_result["ca_cfs"]) > 2: - for hr_ca_cfs_i in helper_result["ca_cfs"][2:]: - ret_cf_sample = _multiply_cosmo_factor( - (True, ret_cf_sample), hr_ca_cfs_i - ) + ret_cf_sample = _multiply_cosmo_factor(*helper_result["ca_cfs"]) if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) inv_ret_cf_sample = _reciprocal_cosmo_factor( @@ -981,24 +973,30 @@ def linspace( return _return_helper(ress, helper_result, ret_cf) -# @implements(np.logspace) -# def logspace(...): -# from unyt._array_functions import logspace as unyt_logspace +@implements(np.logspace) +def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0): + from unyt._array_functions import logspace as unyt_logspace -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_logspace(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + start, stop, num=num, endpoint=endpoint, base=base, dtype=dtype, axis=axis + ) + ret_cf = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["base"]) + res = unyt_logspace(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.geomspace) -# def geomspace(...): -# from unyt._array_functions import geomspace as unyt_geomspace +@implements(np.geomspace) +def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0): + from unyt._array_functions import geomspace as unyt_geomspace -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_geomspace(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + start, stop, num=num, endpoint=endpoint, dtype=dtype, axis=axis + ) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_geomspace(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.copyto) @@ -1043,16 +1041,16 @@ def prod( @implements(np.var) def var( - a, - axis=None, - dtype=None, - out=None, - ddof=0, - keepdims=np._NoValue, - *, - where=np._NoValue, - mean=np._NoValue, - correction=np._NoValue + a, + axis=None, + dtype=None, + out=None, + ddof=0, + keepdims=np._NoValue, + *, + where=np._NoValue, + mean=np._NoValue, + correction=np._NoValue, ): from unyt._array_functions import var as unyt_var @@ -1091,16 +1089,16 @@ def trace(a, offset=0, axis1=0, axis2=1, dtype=None, out=None): @implements(np.percentile) def percentile( - a, - q, - axis=None, - out=None, - overwrite_input=False, - method="linear", - keepdims=False, - *, - weights=None, - interpolation=None + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None, ): from unyt._array_functions import percentile as unyt_percentile @@ -1122,16 +1120,16 @@ def percentile( @implements(np.quantile) def quantile( - a, - q, - axis=None, - out=None, - overwrite_input=False, - method='linear', - keepdims=False, - *, - weights=None, - interpolation=None + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None, ): from unyt._array_functions import quantile as unyt_quantile @@ -1152,17 +1150,17 @@ def quantile( @implements(np.nanpercentile) -def percentile( - a, - q, - axis=None, - out=None, - overwrite_input=False, - method="linear", - keepdims=False, - *, - weights=None, - interpolation=None +def nanpercentile( + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None, ): from unyt._array_functions import nanpercentile as unyt_nanpercentile @@ -1184,16 +1182,16 @@ def percentile( @implements(np.nanquantile) def nanquantile( - a, - q, - axis=None, - out=None, - overwrite_input=False, - method='linear', - keepdims=False, - *, - weights=None, - interpolation=None + a, + q, + axis=None, + out=None, + overwrite_input=False, + method="linear", + keepdims=False, + *, + weights=None, + interpolation=None, ): from unyt._array_functions import nanquantile as unyt_nanquantile @@ -1227,6 +1225,88 @@ def linalg_det(a): return _return_helper(res, helper_result, ret_cf) +@implements(np.diff) +def diff(a, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue): + from unyt._array_functions import diff as unyt_diff + + helper_result = _prepare_array_func_args( + a, n=n, axis=axis, prepend=prepend, append=append + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_diff(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.ediff1d) +def ediff1d(ary, to_end=None, to_begin=None): + from unyt._array_functions import ediff1d as unyt_ediff1d + + helper_result = _prepare_array_func_args(ary, to_end=to_end, to_begin=to_begin) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_ediff1d(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.ptp) +def ptp(a, axis=None, out=None, keepdims=np._NoValue): + from unyt._array_functions import ptp as unyt_ptp + + helper_result = _prepare_array_func_args(a, axis=axis, out=out, keepdims=keepdims) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_ptp(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) + + +# @implements(np.cumprod) +# def cumprod(...): +# Omitted because unyt just raises if called. + + +@implements(np.pad) +def pad(array, pad_width, mode="constant", **kwargs): + from unyt._array_functions import pad as unyt_pad + + helper_result = _prepare_array_func_args(array, pad_width, mode=mode, **kwargs) + # the number of options is huge, including user defined functions to handle data + # let's just preserve the cosmo_factor of the input `array` and trust the user... + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_pad(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +@implements(np.choose) +def choose(a, choices, out=None, mode="raise"): + from unyt._array_functions import choose as unyt_choose + + helper_result = _prepare_array_func_args(a, choices, out=out, mode=mode) + helper_result_choices = _prepare_array_func_args(*choices) + ret_cf = _preserve_cosmo_factor(*helper_result_choices["ca_cfs"]) + res = unyt_choose(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) + + +@implements(np.insert) +def insert(arr, obj, values, axis=None): + from unyt._array_functions import insert as unyt_insert + + helper_result = _prepare_array_func_args(arr, obj, values, axis=axis) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][2] + ) + res = unyt_insert(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + +# @implements(np.isin) +# def isin(...): +# from unyt._array_functions import isin as unyt_isin + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + # @implements(np.linalg.lstsq) # def linalg_lstsq(...): # from unyt._array_functions import linalg_lstsq as unyt_linalg_lstsq @@ -1317,90 +1397,19 @@ def linalg_det(a): # return _return_helper(res, helper_result, ret_cf, out=out) -@implements(np.diff) -def diff(a, n=1, axis=-1, prepend=np._NoValue, append=np._NoValue): - from unyt._array_functions import diff as unyt_diff - - helper_result = _prepare_array_func_args( - a, n=n, axis=axis, prepend=prepend, append=append - ) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) - res = unyt_diff(*helper_result["args"], **helper_result["kwargs"]) - return _return_helper(res, helper_result, ret_cf) - - -@implements(np.ediff1d) -def ediff1d(ary, to_end=None, to_begin=None): - from unyt._array_functions import ediff1d as unyt_ediff1d - - helper_result = _prepare_array_func_args(ary, to_end=to_end, to_begin=to_begin) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) - res = unyt_ediff1d(*helper_result["args"], **helper_result["kwargs"]) - return _return_helper(res, helper_result, ret_cf) - +@implements(np.fill_diagonal) +def fill_diagonal(a, val, wrap=False): + from unyt._array_functions import fill_diagonal as unyt_fill_diagonal -# @implements(np.ptp) -# def ptp(...): -# from unyt._array_functions import ptp as unyt_ptp - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_ptp(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.cumprod) -# def cumprod(...): -# Omitted because unyt just raises if called. - -# @implements(np.pad) -# def pad(...): -# from unyt._array_functions import pad as unyt_pad - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_pad(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.choose) -# def choose(...): -# from unyt._array_functions import choose as unyt_choose - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_choose(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.fill_diagonal) -# def fill_diagonal(...): -# from unyt._array_functions import fill_diagonal as unyt_fill_diagonal - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_fill_diagonal(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.insert) -# def insert(...): -# from unyt._array_functions import insert as unyt_insert - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_insert(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.isin) -# def isin(...): -# from unyt._array_functions import isin as unyt_isin - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, val, wrap=wrap) + _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][1]) + # must pass a directly here because it's modified in-place + comoving = getattr(a, "comoving", None) + if comoving: + val = val.to_comoving() + elif comoving is False: + val = val.to_physical() + unyt_fill_diagonal(a, val, **helper_result["kwargs"]) # @implements(np.place) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index d77430d4..db2be603 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -135,7 +135,17 @@ def _sqrt_cosmo_factor(ca_cf, **kwargs): ) # ufunc sqrt not supported -def _multiply_cosmo_factor(ca_cf1, ca_cf2, **kwargs): +def _multiply_cosmo_factor(*args, **kwargs): + ca_cfs = args + if len(ca_cfs) == 1: + return __multiply_cosmo_factor(ca_cfs[0]) + retval = __multiply_cosmo_factor(ca_cfs[0], ca_cfs[1]) + for ca_cf in ca_cfs[2:]: + retval = __multiply_cosmo_factor((retval is not None, retval), ca_cf) + return retval + + +def __multiply_cosmo_factor(ca_cf1, ca_cf2, **kwargs): ca1, cf1 = ca_cf1 ca2, cf2 = ca_cf2 if (cf1 is None) and (cf2 is None): @@ -163,7 +173,17 @@ def _multiply_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise RuntimeError("Unexpected state, please report this error on github.") -def _preserve_cosmo_factor(ca_cf1, ca_cf2=None, **kwargs): +def _preserve_cosmo_factor(*args, **kwargs): + ca_cfs = args + if len(ca_cfs) == 1: + return __preserve_cosmo_factor(ca_cfs[0]) + retval = __preserve_cosmo_factor(ca_cfs[0], ca_cfs[1]) + for ca_cf in ca_cfs[2:]: + retval = __preserve_cosmo_factor((retval is not None, retval), ca_cf) + return retval + + +def __preserve_cosmo_factor(ca_cf1, ca_cf2=None, **kwargs): ca1, cf1 = ca_cf1 ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) if ca_cf2 is None: @@ -335,7 +355,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) + return cosmo_factor(a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -443,14 +463,6 @@ def _prepare_array_func_args(*args, **kwargs): else: # mixed compressions, strip it off ret_comp = None - expected_scale_factor = None - for ca_cf in ca_cfs + list(kw_ca_cfs.values()): - if ca_cf[0]: - if expected_scale_factor is None: - expected_scale_factor = ca_cf[1].scale_factor - else: - if ca_cf[1].scale_factor != expected_scale_factor: - raise ValueError("Mismatched scale factors in cosmo_array's.") return dict( args=args, kwargs=kwargs, @@ -659,7 +671,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 42644c0a..166561ef 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -19,7 +19,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) def to_ua(x): @@ -49,7 +49,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -60,7 +60,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -70,7 +70,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -80,7 +80,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -144,9 +144,9 @@ def test_handled_funcs(self): "allclose": (ca(np.arange(3)), ca(np.arange(3))), "array_equal": (ca(np.arange(3)), ca(np.arange(3))), "array_equiv": (ca(np.arange(3)), ca(np.arange(3))), - "linspace": (ca(1), ca(2), 10), - # "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless), 10), - # "geomspace": (ca(1), ca(1), 10), + "linspace": (ca(1), ca(2)), + "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless)), + "geomspace": (ca(1), ca(1)), "prod": (ca(np.arange(3)),), "var": (ca(np.arange(3)),), "trace": (ca(np.eye(3)),), @@ -157,11 +157,11 @@ def test_handled_funcs(self): "linalg.det": (ca(np.eye(3)),), "diff": (ca(np.arange(3)),), "ediff1d": (ca(np.arange(3)),), - # "ptp": (ca(np.arange(3)),), + "ptp": (ca(np.arange(3)),), "cumprod": (ca(np.arange(3)),), - # "pad": (ca(np.arange(3)), 3), - # "choose": (np.arange(3), ca(np.eye(3))), - # "insert": (ca(np.arange(3)), 1, ca(1)), + "pad": (ca(np.arange(3)), 3), + "choose": (np.arange(3), ca(np.eye(3))), + "insert": (ca(np.arange(3)), 1, ca(1)), # "isin": (ca(np.arange(3)), ca(np.arange(3))), # "in1d": (ca(np.arange(3)), ca(np.arange(3))), # "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), @@ -195,7 +195,7 @@ def test_handled_funcs(self): # "linalg.eigh": (ca(np.eye(3)),), # "copyto": (ca(np.arange(3)), ca(np.arange(3))), # "savetxt": (savetxt_file, ca(np.arange(3))), - # "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), + "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), # "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), # "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), # "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), @@ -223,6 +223,10 @@ def test_handled_funcs(self): result = func(*args) continue result = func(*args) + if "fill_diagonal" in fname: + # treat inplace modified values for relevant functions as result + result = args[0] + ua_result = ua_args[0] if "savetxt" in fname and os.path.isfile(savetxt_file): os.remove(savetxt_file) if ua_result is None: @@ -270,7 +274,7 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), ), ), @@ -281,13 +285,13 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a ** 2, 1.0), + cosmo_factor=cosmo_factor(a**2, 1.0), ), ), ), @@ -299,19 +303,19 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a ** 2, 1.0), + cosmo_factor=cosmo_factor(a**2, 1.0), ), cosmo_array( [1, 2, 3], u.kg, comoving=False, - cosmo_factor=cosmo_factor(a ** 3, 1.0), + cosmo_factor=cosmo_factor(a**3, 1.0), ), ], ), @@ -323,16 +327,16 @@ def test_handled_funcs(self): ( None, cosmo_array( - [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 1.0) + [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a**1, 1.0) ), np.array([1, 2, 3]), ), ) - @pytest.mark.parametrize("bins", ("int", "np", "ca")) + @pytest.mark.parametrize("bins_type", ("int", "np", "ca")) @pytest.mark.parametrize("density", (None, True)) - def test_histograms(self, func_args, weights, bins, density): + def test_histograms(self, func_args, weights, bins_type, density): func, args = func_args - _bins = { + bins = { "int": 10, "np": [np.linspace(0, 5, 11)] * 3, "ca": [ @@ -340,32 +344,32 @@ def test_histograms(self, func_args, weights, bins, density): np.linspace(0, 5, 11), u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.K, comoving=False, - cosmo_factor=cosmo_factor(a ** 2, 1.0), + cosmo_factor=cosmo_factor(a**2, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.Msun, comoving=False, - cosmo_factor=cosmo_factor(a ** 3, 1.0), + cosmo_factor=cosmo_factor(a**3, 1.0), ), ], - }[bins] + }[bins_type] bins = ( - _bins[ + bins[ { np.histogram: np.s_[0], np.histogram2d: np.s_[:2], np.histogramdd: np.s_[:], }[func] ] - if bins in ("np", "ca") - else _bins + if bins_type in ("np", "ca") + else bins ) result = func(*args, bins=bins, density=density, weights=weights) ua_args = tuple( @@ -391,17 +395,17 @@ def test_histograms(self, func_args, weights, bins, density): check_result(r, ua_r) else: check_result(result, ua_result) - if density is None and not isinstance(weights, cosmo_array): - assert not hasattr(result[0], "comoving") + if not density and not isinstance(weights, cosmo_array): + assert not isinstance(result[0], cosmo_array) else: assert result[0].comoving is False if density and not isinstance(weights, cosmo_array): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a ** -1, 1.0), - np.histogram2d: cosmo_factor(a ** -3, 1.0), - np.histogramdd: cosmo_factor(a ** -6, 1.0), + np.histogram: cosmo_factor(a**-1, 1.0), + np.histogram2d: cosmo_factor(a**-3, 1.0), + np.histogramdd: cosmo_factor(a**-6, 1.0), }[func] ) elif density and isinstance(weights, cosmo_array): @@ -409,19 +413,19 @@ def test_histograms(self, func_args, weights, bins, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a ** 0, 1.0), - np.histogram2d: cosmo_factor(a ** -2, 1.0), - np.histogramdd: cosmo_factor(a ** -5, 1.0), + np.histogram: cosmo_factor(a**0, 1.0), + np.histogram2d: cosmo_factor(a**-2, 1.0), + np.histogramdd: cosmo_factor(a**-5, 1.0), }[func] ) - elif density is None and isinstance(weights, cosmo_array): + elif not density and isinstance(weights, cosmo_array): assert result[0].comoving is False assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a ** 1, 1.0), - np.histogram2d: cosmo_factor(a ** 1, 1.0), - np.histogramdd: cosmo_factor(a ** 1, 1.0), + np.histogram: cosmo_factor(a**1, 1.0), + np.histogram2d: cosmo_factor(a**1, 1.0), + np.histogramdd: cosmo_factor(a**1, 1.0), }[func] ) ret_bins = { @@ -433,9 +437,9 @@ def test_histograms(self, func_args, weights, bins, density): ret_bins, ( [ - cosmo_factor(a ** 1, 1.0), - cosmo_factor(a ** 2, 1.0), - cosmo_factor(a ** 3, 1.0), + cosmo_factor(a**1, 1.0), + cosmo_factor(a**2, 1.0), + cosmo_factor(a**3, 1.0), ] ), ): From cf11cfdf6ced5c09461044ba77c0900e90609903 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 18 Dec 2024 15:45:55 +0000 Subject: [PATCH 21/29] Implement a few more numpy functions. --- swiftsimio/_array_functions.py | 100 ++++++++++++++++++++------------- tests/test_cosmo_array.py | 44 +++++++-------- 2 files changed, 82 insertions(+), 62 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index f345fe1c..905853ce 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -6,6 +6,7 @@ cosmo_quantity, _prepare_array_func_args, _multiply_cosmo_factor, + _divide_cosmo_factor, _preserve_cosmo_factor, _reciprocal_cosmo_factor, _comparison_cosmo_factor, @@ -999,14 +1000,20 @@ def geomspace(start, stop, num=50, endpoint=True, dtype=None, axis=0): return _return_helper(res, helper_result, ret_cf) -# @implements(np.copyto) -# def copyto(...): -# from unyt._array_functions import copyto as unyt_copyto +@implements(np.copyto) +def copyto(dst, src, casting="same_kind", where=True): + from unyt._array_functions import copyto as unyt_copyto -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_copyto(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(dst, src, casting=casting, where=where) + _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][1]) + # must pass dst directly here because it's modified in-place + if isinstance(src, cosmo_array): + comoving = getattr(dst, "comoving", None) + if comoving: + src.convert_to_comoving() + elif comoving is False: + src.convert_to_physical() + unyt_copyto(dst, src, **helper_result["kwargs"]) @implements(np.prod) @@ -1297,34 +1304,35 @@ def insert(arr, obj, values, axis=None): return _return_helper(res, helper_result, ret_cf) -# @implements(np.isin) -# def isin(...): -# from unyt._array_functions import isin as unyt_isin - -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.linalg.lstsq) -# def linalg_lstsq(...): -# from unyt._array_functions import linalg_lstsq as unyt_linalg_lstsq +@implements(np.linalg.lstsq) +def linalg_lstsq(a, b, rcond=None): + from unyt._array_functions import linalg_lstsq as unyt_linalg_lstsq -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_lstsq(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, b, rcond=rcond) + ret_cf = _divide_cosmo_factor( + helper_result["ca_cfs"][1], helper_result["ca_cfs"][0] + ) + resid_cf = _power_cosmo_factor(helper_result["ca_cfs"][1], (False, None), power=2) + sing_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ress = unyt_linalg_lstsq(*helper_result["args"], **helper_result["kwargs"]) + return ( + _return_helper(ress[0], helper_result, ret_cf), + _return_helper(ress[1], helper_result, resid_cf), + ress[2], + _return_helper(ress[3], helper_result, sing_cf), + ) -# @implements(np.linalg.solve) -# def linalg_solve(...): -# from unyt._array_functions import linalg_solve as unyt_linalg_solve +@implements(np.linalg.solve) +def linalg_solve(a, b): + from unyt._array_functions import linalg_solve as unyt_linalg_solve -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_solve(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, b) + ret_cf = _divide_cosmo_factor( + helper_result["ca_cfs"][1], helper_result["ca_cfs"][0] + ) + res = unyt_linalg_solve(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.linalg.tensorsolve) @@ -1406,12 +1414,22 @@ def fill_diagonal(a, val, wrap=False): # must pass a directly here because it's modified in-place comoving = getattr(a, "comoving", None) if comoving: - val = val.to_comoving() + val.convert_to_comoving() elif comoving is False: - val = val.to_physical() + val.convert_to_physical() unyt_fill_diagonal(a, val, **helper_result["kwargs"]) +# @implements(np.isin) +# def isin(...): +# from unyt._array_functions import isin as unyt_isin + +# helper_result = _prepare_array_func_args(...) +# ret_cf = ...() +# res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) +# return _return_helper(res, helper_result, ret_cf, out=out) + + # @implements(np.place) # def place(...): # from unyt._array_functions import place as unyt_place @@ -1602,14 +1620,16 @@ def fill_diagonal(a, val, wrap=False): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.linalg.outer) -# def linalg_outer(...): -# from unyt._array_functions import linalg_outer as unyt_linalg_outer +@implements(np.linalg.outer) +def linalg_outer(x1, x2, /): + from unyt._array_functions import linalg_outer as unyt_linalg_outer -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_outer(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(x1, x2) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_linalg_outer(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.trapezoid) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 166561ef..b81444f0 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -147,6 +147,7 @@ def test_handled_funcs(self): "linspace": (ca(1), ca(2)), "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless)), "geomspace": (ca(1), ca(1)), + "copyto": (ca(np.arange(3)), ca(np.arange(3))), "prod": (ca(np.arange(3)),), "var": (ca(np.arange(3)),), "trace": (ca(np.eye(3)),), @@ -162,9 +163,24 @@ def test_handled_funcs(self): "pad": (ca(np.arange(3)), 3), "choose": (np.arange(3), ca(np.eye(3))), "insert": (ca(np.arange(3)), 1, ca(1)), + "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), + # "linalg.tensorsolve": ( + # ca(np.eye(24).reshape((6, 4, 2, 3, 4))), + # ca(np.ones((6, 4))), + # ), + # "linalg.eig": (ca(np.eye(3)),), + # "linalg.eigh": (ca(np.eye(3)),), + # "linalg.eigvals": (ca(np.eye(3)),), + # "linalg.eigvalsh": (ca(np.eye(3)),), + # "savetxt": (savetxt_file, ca(np.arange(3))), + "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), + # "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), # "isin": (ca(np.arange(3)), ca(np.arange(3))), - # "in1d": (ca(np.arange(3)), ca(np.arange(3))), - # "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + # "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), + # "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + # "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), + # "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), # "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), # "select": ( # [np.arange(3) < 1, np.arange(3) > 1], @@ -182,27 +198,11 @@ def test_handled_funcs(self): # "correlate": (ca(np.arange(3)), ca(np.arange(3))), # "tensordot": (ca(np.eye(3)), ca(np.eye(3))), # "unwrap": (ca(np.arange(3)),), - # "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), - # "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), - # "linalg.tensorsolve": ( - # ca(np.eye(24).reshape((6, 4, 2, 3, 4))), - # ca(np.ones((6, 4))), - # ), - # "linalg.eigvals": (ca(np.eye(3)),), - # "linalg.eigvalsh": (ca(np.eye(3)),), - # "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), - # "linalg.eig": (ca(np.eye(3)),), - # "linalg.eigh": (ca(np.eye(3)),), - # "copyto": (ca(np.arange(3)), ca(np.arange(3))), - # "savetxt": (savetxt_file, ca(np.arange(3))), - "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), - # "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), - # "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), - # "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), - # "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), - # "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + # "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), # "array_repr": (ca(np.arange(3)),), + "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), # "trapezoid": (ca(np.arange(3)),), + # "in1d": (ca(np.arange(3)), ca(np.arange(3))), } functions_checked = list() bad_funcs = dict() @@ -223,7 +223,7 @@ def test_handled_funcs(self): result = func(*args) continue result = func(*args) - if "fill_diagonal" in fname: + if ("fill_diagonal" in fname) or ("copyto" in fname): # treat inplace modified values for relevant functions as result result = args[0] ua_result = ua_args[0] From dd106fd92d4171e7ec216c9245cb866b903a8fd7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Dec 2024 12:33:15 +0000 Subject: [PATCH 22/29] Implement more numpy functions. --- swiftsimio/_array_functions.py | 311 +++++++++++++++++++++------------ swiftsimio/objects.py | 4 +- tests/test_cosmo_array.py | 120 +++++++------ 3 files changed, 270 insertions(+), 165 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 905853ce..d8292c02 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1,3 +1,4 @@ +import warnings import numpy as np from unyt import unyt_quantity, unyt_array from unyt._array_functions import implements @@ -11,6 +12,7 @@ _reciprocal_cosmo_factor, _comparison_cosmo_factor, _power_cosmo_factor, + _return_without_cosmo_factor, ) _HANDLED_FUNCTIONS = dict() @@ -1335,74 +1337,113 @@ def linalg_solve(a, b): return _return_helper(res, helper_result, ret_cf) -# @implements(np.linalg.tensorsolve) -# def linalg_tensorsolve(...): -# from unyt._array_functions import linalg_tensorsolve as unyt_linalg_tensorsolve +@implements(np.linalg.tensorsolve) +def linalg_tensorsolve(a, b, axes=None): + from unyt._array_functions import linalg_tensorsolve as unyt_linalg_tensorsolve -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_tensorsolve(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, b, axes=axes) + ret_cf = _divide_cosmo_factor(helper_result["ca_cfs"][1], helper_result["ca_cfs"][0]) + res = unyt_linalg_tensorsolve(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.linalg.eig) -# def linalg_eig(...): -# from unyt._array_functions import linalg_eig as unyt_linalg_eig +@implements(np.linalg.eig) +def linalg_eig(a): + from unyt._array_functions import linalg_eig as unyt_linalg_eig -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) - - -# @implements(np.linalg.eigh) -# def linalg_eigh(...): -# from unyt._array_functions import linalg_eigh as unyt_linalg_eigh + helper_result = _prepare_array_func_args(a) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ress = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) + return ( + _return_helper(ress[0], helper_result, ret_cf), + ress[1], + ) -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) +@implements(np.linalg.eigh) +def linalg_eigh(a, UPLO="L"): + from unyt._array_functions import linalg_eigh as unyt_linalg_eigh -# @implements(np.linalg.eigvals) -# def linalg_eigvals(...): -# from unyt._array_functions import linalg_eigvals as unyt_linalg_eigvals + helper_result = _prepare_array_func_args(a, UPLO=UPLO) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ress = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) + return ( + _return_helper(ress[0], helper_result, ret_cf), + ress[1], + ) -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_eigvals(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) +@implements(np.linalg.eigvals) +def linalg_eigvals(a): + from unyt._array_functions import linalg_eigvals as unyt_linalg_eigvals -# @implements(np.linalg.eigvalsh) -# def linalg_eigvalsh(...): -# from unyt._array_functions import linalg_eigvalsh as unyt_linalg_eigvalsh + helper_result = _prepare_array_func_args(a) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_linalg_eigvals(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_linalg_eigvalsh(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) +@implements(np.linalg.eigvalsh) +def linalg_eigvalsh(a, UPLO="L"): + from unyt._array_functions import linalg_eigvalsh as unyt_linalg_eigvalsh -# @implements(np.savetxt) -# def savetxt(...): -# from unyt._array_functions import savetxt as unyt_savetxt + helper_result = _prepare_array_func_args(a, UPLO=UPLO) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_linalg_eigvalsh(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_savetxt(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) +@implements(np.savetxt) +def savetxt( + fname, + X, + fmt="%.18e", + delimiter=" ", + newline="\n", + header="", + footer="", + comments="# ", + encoding=None +): + from unyt._array_functions import savetxt as unyt_savetxt + + warnings.warn( + "numpy.savetxt does not preserve units or cosmo_array information, " + "and will only save the raw numerical data from the cosmo_array object.\n" + "If this is the intended behaviour, call `numpy.savetxt(file, arr.d)` " + "to silence this warning.\n", + stacklevel=4, + ) + helper_result = _prepare_array_func_args( + fname, + X, + fmt=fmt, + delimiter=delimiter, + newline=newline, + header=header, + footer=footer, + comments=comments, + encoding=encoding, + ) + with warnings.catch_warnings(): + warnings.filterwarnings( + action="ignore", + category=UserWarning, + message="numpy.savetxt does not preserve units" + ) + unyt_savetxt(*helper_result["args"], **helper_result["kwargs"]) + return -# @implements(np.apply_over_axes) -# def apply_over_axes(...): -# from unyt._array_functions import apply_over_axes as unyt_apply_over_axes -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_apply_over_axes(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) +@implements(np.apply_over_axes) +def apply_over_axes(func, a, axes): + res = func(a, axes[0]) + if len(axes) > 1: + # this function is recursive by nature, + # here we intentionally do not call the base _implementation + return np.apply_over_axes(func, res, axes[1:]) + else: + return res @implements(np.fill_diagonal) @@ -1420,64 +1461,94 @@ def fill_diagonal(a, val, wrap=False): unyt_fill_diagonal(a, val, **helper_result["kwargs"]) -# @implements(np.isin) -# def isin(...): -# from unyt._array_functions import isin as unyt_isin +@implements(np.isin) +def isin(element, test_elements, assume_unique=False, invert=False, *, kind=None): + from unyt._array_functions import isin as unyt_isin -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + element, test_elements, assume_unique=assume_unique, invert=invert, kind=kind + ) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=(element, test_elements), + ) + res = unyt_isin(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.place) -# def place(...): -# from unyt._array_functions import place as unyt_place +@implements(np.place) +def place(arr, mask, vals): + from unyt._array_functions import place as unyt_place -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_place(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(arr, mask, vals) + _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + # must pass arr directly here because it's modified in-place + if isinstance(vals, cosmo_array): + comoving = getattr(arr, "comoving", None) + if comoving: + vals.convert_to_comoving() + elif comoving is False: + vals.convert_to_physical() + unyt_place(arr, mask, vals) -# @implements(np.put) -# def ...(...): -# from unyt._array_functions import put as unyt_put +@implements(np.put) +def put(a, ind, v, mode="raise"): + from unyt._array_functions import put as unyt_put -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_put(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, ind, v, mode=mode) + _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + # must pass arr directly here because it's modified in-place + if isinstance(v, cosmo_array): + comoving = getattr(a, "comoving", None) + if comoving: + v.convert_to_comoving() + elif comoving is False: + v.convert_to_physical() + unyt_put(a, ind, v, **helper_result["kwargs"]) -# @implements(np.put_along_axis) -# def put_along_axis(...): -# from unyt._array_functions import put_along_axis as unyt_put_along_axis +@implements(np.put_along_axis) +def put_along_axis(arr, indices, values, axis): + from unyt._array_functions import put_along_axis as unyt_put_along_axis -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_put_along_axis(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(arr, indices, values, axis) + _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + if isinstance(values, cosmo_array): + comoving = getattr(arr, "comoving", None) + if comoving: + values.convert_to_comoving() + elif comoving is False: + values.convert_to_physical() + unyt_put_along_axis(arr, indices, values, axis) -# @implements(np.putmask) -# def putmask(...): -# from unyt._array_functions import putmask as unyt_putmask +@implements(np.putmask) +def putmask(a, mask, values): + from unyt._array_functions import putmask as unyt_putmask -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_putmask(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, mask, values) + _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + if isinstance(values, cosmo_array): + comoving = getattr(a, "comoving", None) + if comoving: + values.convert_to_comoving() + elif comoving is False: + values.convert_to_physical() + unyt_putmask(a, mask, values) -# @implements(np.searchsorted) -# def searchsorted(...): -# from unyt._array_functions import searchsorted as unyt_searchsorted +@implements(np.searchsorted) +def searchsorted(a, v, side="left", sorter=None): + from unyt._array_functions import searchsorted as unyt_searchsorted -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_searchsorted(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, v, side=side, sorter=sorter) + ret_cf = _return_without_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_searchsorted(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) # @implements(np.select) @@ -1500,24 +1571,50 @@ def fill_diagonal(a, val, wrap=False): # return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.sinc) -# def sinc(...): -# from unyt._array_functions import sinc as unyt_sinc +@implements(np.sinc) +def sinc(x): + from unyt._array_functions import sinc as unyt_sinc -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_sinc(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + # unyt just casts to array and calls the numpy implementation + # so let's just hand off to them + return unyt_sinc(x) -# @implements(np.clip) -# def clip(...): -# from unyt._array_functions import clip as unyt_clip +@implements(np.clip) +def clip( + a, + a_min=np._NoValue, + a_max=np._NoValue, + out=None, + *, + min=np._NoValue, + max=np._NoValue, + **kwargs, +): + from unyt._array_functions import clip as unyt_clip -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_clip(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + # can't work out how to properly handle min and max, + # just leave them in kwargs I guess (might be a numpy version conflict?) + helper_result = _prepare_array_func_args( + a, + a_min=a_min, + a_max=a_max, + out=out, + **kwargs, + ) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["kw_ca_cfs"]["a_min"], + helper_result["kw_ca_cfs"]["a_max"], + ) + res = unyt_clip( + helper_result["args"][0], + helper_result["kwargs"]["a_min"], + helper_result["kwargs"]["a_max"], + out=helper_result["kwargs"]["out"], + **kwargs, + ) + return _return_helper(res, helper_result, ret_cf, out=out) # @implements(np.where) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index db2be603..58a0fc2e 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -355,7 +355,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a**0, scale_factor=cf1.scale_factor) + return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -671,7 +671,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index b81444f0..0c84887f 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -19,7 +19,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) def to_ua(x): @@ -49,7 +49,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -60,7 +60,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -70,7 +70,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -80,7 +80,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -165,31 +165,31 @@ def test_handled_funcs(self): "insert": (ca(np.arange(3)), 1, ca(1)), "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), - # "linalg.tensorsolve": ( - # ca(np.eye(24).reshape((6, 4, 2, 3, 4))), - # ca(np.ones((6, 4))), - # ), - # "linalg.eig": (ca(np.eye(3)),), - # "linalg.eigh": (ca(np.eye(3)),), - # "linalg.eigvals": (ca(np.eye(3)),), - # "linalg.eigvalsh": (ca(np.eye(3)),), - # "savetxt": (savetxt_file, ca(np.arange(3))), + "linalg.tensorsolve": ( + ca(np.eye(24).reshape((6, 4, 2, 3, 4))), + ca(np.ones((6, 4))), + ), + "linalg.eig": (ca(np.eye(3)),), + "linalg.eigh": (ca(np.eye(3)),), + "linalg.eigvals": (ca(np.eye(3)),), + "linalg.eigvalsh": (ca(np.eye(3)),), + "savetxt": (savetxt_file, ca(np.arange(3))), "fill_diagonal": (ca(np.eye(3)), ca(np.arange(3))), - # "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), - # "isin": (ca(np.arange(3)), ca(np.arange(3))), - # "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), - # "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), - # "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), - # "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), - # "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), + "apply_over_axes": (lambda x, axis: x, ca(np.eye(3)), (0, 1)), + "isin": (ca(np.arange(3)), ca(np.arange(3))), + "place": (ca(np.arange(3)), np.arange(3) > 0, ca(np.arange(3))), + "put": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), + "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), + "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), # "select": ( # [np.arange(3) < 1, np.arange(3) > 1], # [ca(np.arange(3)), ca(np.arange(3))], # ca(1), # ), # "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), - # "sinc": (ca(np.arange(3)),), - # "clip": (ca(np.arange(3)), ca(1), ca(2)), + "sinc": (ca(np.arange(3)),), + "clip": (ca(np.arange(3)), ca(1), ca(2)), # "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), # "triu": (ca(np.arange(3)),), # "tril": (ca(np.arange(3)),), @@ -223,7 +223,14 @@ def test_handled_funcs(self): result = func(*args) continue result = func(*args) - if ("fill_diagonal" in fname) or ("copyto" in fname): + if fname.split(".")[-1] in ( + "fill_diagonal", + "copyto", + "place", + "put", + "put_along_axis", + "putmask", + ): # treat inplace modified values for relevant functions as result result = args[0] ua_result = ua_args[0] @@ -234,16 +241,17 @@ def test_handled_funcs(self): assert result is None except AssertionError: bad_funcs["np." + fname] = result, ua_result - try: - if isinstance(ua_result, tuple): - assert isinstance(result, tuple) - assert len(result) == len(ua_result) - for r, ua_r in zip(result, ua_result): - check_result(r, ua_r) - else: - check_result(result, ua_result) - except AssertionError: - bad_funcs["np." + fname] = result, ua_result + else: + try: + if isinstance(ua_result, tuple): + assert isinstance(result, tuple) + assert len(result) == len(ua_result) + for r, ua_r in zip(result, ua_result): + check_result(r, ua_r) + else: + check_result(result, ua_result) + except AssertionError: + bad_funcs["np." + fname] = result, ua_result if len(bad_funcs) > 0: raise AssertionError( "Some functions did not return expected types " @@ -274,7 +282,7 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), ), ), @@ -285,13 +293,13 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a**2, 1.0), + cosmo_factor=cosmo_factor(a ** 2, 1.0), ), ), ), @@ -303,19 +311,19 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a**2, 1.0), + cosmo_factor=cosmo_factor(a ** 2, 1.0), ), cosmo_array( [1, 2, 3], u.kg, comoving=False, - cosmo_factor=cosmo_factor(a**3, 1.0), + cosmo_factor=cosmo_factor(a ** 3, 1.0), ), ], ), @@ -327,7 +335,7 @@ def test_handled_funcs(self): ( None, cosmo_array( - [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a**1, 1.0) + [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 1.0) ), np.array([1, 2, 3]), ), @@ -344,19 +352,19 @@ def test_histograms(self, func_args, weights, bins_type, density): np.linspace(0, 5, 11), u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.K, comoving=False, - cosmo_factor=cosmo_factor(a**2, 1.0), + cosmo_factor=cosmo_factor(a ** 2, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.Msun, comoving=False, - cosmo_factor=cosmo_factor(a**3, 1.0), + cosmo_factor=cosmo_factor(a ** 3, 1.0), ), ], }[bins_type] @@ -403,9 +411,9 @@ def test_histograms(self, func_args, weights, bins_type, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a**-1, 1.0), - np.histogram2d: cosmo_factor(a**-3, 1.0), - np.histogramdd: cosmo_factor(a**-6, 1.0), + np.histogram: cosmo_factor(a ** -1, 1.0), + np.histogram2d: cosmo_factor(a ** -3, 1.0), + np.histogramdd: cosmo_factor(a ** -6, 1.0), }[func] ) elif density and isinstance(weights, cosmo_array): @@ -413,9 +421,9 @@ def test_histograms(self, func_args, weights, bins_type, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a**0, 1.0), - np.histogram2d: cosmo_factor(a**-2, 1.0), - np.histogramdd: cosmo_factor(a**-5, 1.0), + np.histogram: cosmo_factor(a ** 0, 1.0), + np.histogram2d: cosmo_factor(a ** -2, 1.0), + np.histogramdd: cosmo_factor(a ** -5, 1.0), }[func] ) elif not density and isinstance(weights, cosmo_array): @@ -423,9 +431,9 @@ def test_histograms(self, func_args, weights, bins_type, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a**1, 1.0), - np.histogram2d: cosmo_factor(a**1, 1.0), - np.histogramdd: cosmo_factor(a**1, 1.0), + np.histogram: cosmo_factor(a ** 1, 1.0), + np.histogram2d: cosmo_factor(a ** 1, 1.0), + np.histogramdd: cosmo_factor(a ** 1, 1.0), }[func] ) ret_bins = { @@ -437,9 +445,9 @@ def test_histograms(self, func_args, weights, bins_type, density): ret_bins, ( [ - cosmo_factor(a**1, 1.0), - cosmo_factor(a**2, 1.0), - cosmo_factor(a**3, 1.0), + cosmo_factor(a ** 1, 1.0), + cosmo_factor(a ** 2, 1.0), + cosmo_factor(a ** 3, 1.0), ] ), ): From 77db41e339077b99c33c08dcd4d245efc2b4bd90 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Dec 2024 16:44:48 +0000 Subject: [PATCH 23/29] Finish implementing numpy functions, immediate tests pass. --- swiftsimio/_array_functions.py | 265 +++++++++++++++++++++------------ tests/test_cosmo_array.py | 115 +++++++++++--- 2 files changed, 262 insertions(+), 118 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index d8292c02..5d71caf4 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1515,6 +1515,7 @@ def put_along_axis(arr, indices, values, axis): helper_result = _prepare_array_func_args(arr, indices, values, axis) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + # must pass arr directly here because it's modified in-place if isinstance(values, cosmo_array): comoving = getattr(arr, "comoving", None) if comoving: @@ -1530,6 +1531,7 @@ def putmask(a, mask, values): helper_result = _prepare_array_func_args(a, mask, values) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + # must pass arr directly here because it's modified in-place if isinstance(values, cosmo_array): comoving = getattr(a, "comoving", None) if comoving: @@ -1551,24 +1553,32 @@ def searchsorted(a, v, side="left", sorter=None): return _return_helper(res, helper_result, ret_cf) -# @implements(np.select) -# def select(...): -# from unyt._array_functions import select as unyt_select +@implements(np.select) +def select(condlist, choicelist, default=0): + from unyt._array_functions import select as unyt_select -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_select(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(condlist, choicelist, default=default) + helper_result_choicelist = _prepare_array_func_args(*choicelist) + ret_cf = _preserve_cosmo_factor(*helper_result_choicelist["ca_cfs"]) + res = unyt_select( + helper_result["args"][0], + helper_result_choicelist["args"], + **helper_result["kwargs"] + ) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.setdiff1d) -# def setdiff1d(...): -# from unyt._array_functions import setdiff1d as unyt_setdiff1d +@implements(np.setdiff1d) +def setdiff1d(ar1, ar2, assume_unique=False): + from unyt._array_functions import setdiff1d as unyt_setdiff1d -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_setdiff1d(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(ar1, ar2, assume_unique=assume_unique) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + ) + res = unyt_setdiff1d(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) @implements(np.sinc) @@ -1617,104 +1627,150 @@ def clip( return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.where) -# def where(...): -# from unyt._array_functions import where as unyt_where +@implements(np.where) +def where(condition, *args): + from unyt._array_functions import where as unyt_where -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(condition, *args) + if len(args) == 0: # just condition + ret_cf = _return_without_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) + elif len(args) < 2: + # error message borrowed from numpy 1.24.1 + raise ValueError("either both or neither of x and y should be given") + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][1], helper_result["ca_cfs"][2] + ) + res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.triu) -# def triu(...): -# from unyt._array_functions import triu as unyt_triu +@implements(np.triu) +def triu(m, k=0): + from unyt._array_functions import triu as unyt_triu -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_triu(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(m, k=0) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_triu(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.tril) -# def tril(...): -# from unyt._array_functions import tril as unyt_tril +@implements(np.tril) +def tril(m, k=0): + from unyt._array_functions import tril as unyt_tril -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_tril(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(m, k=0) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_tril(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.einsum) -# def einsum(...): -# from unyt._array_functions import einsum as unyt_einsum +@implements(np.einsum) +def einsum( + subscripts, *operands, out=None, dtype=None, order="K", casting="safe", optimize=False +): + from unyt._array_functions import einsum as unyt_einsum -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_einsum(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + subscripts, + operands, + out=out, + dtype=dtype, + order=order, + casting=casting, + optimize=optimize, + ) + helper_result_operands = _prepare_array_func_args(*operands) + ret_cf = _preserve_cosmo_factor(*helper_result_operands["ca_cfs"]) + res = unyt_einsum( + helper_result["args"][0], + *helper_result_operands["args"], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result, ret_cf, out=out) -# @implements(np.convolve) -# def convolve(...): -# from unyt._array_functions import convolve as unyt_convolve +@implements(np.convolve) +def convolve(a, v, mode="full"): + from unyt._array_functions import convolve as unyt_convolve -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_convolve(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, v, mode=mode) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_convolve(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.correlate) -# def correlate(...): -# from unyt._array_functions import correlate as unyt_correlate +@implements(np.correlate) +def correlate(a, v, mode="valid"): + from unyt._array_functions import correlate as unyt_correlate -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_correlate(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, v, mode=mode) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_correlate(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.tensordot) -# def tensordot(...): -# from unyt._array_functions import tensordot as unyt_tensordot +@implements(np.tensordot) +def tensordot(a, b, axes=2): + from unyt._array_functions import tensordot as unyt_tensordot -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_tensordot(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(a, b, axes=axes) + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_tensordot(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.unwrap) -# def unwrap(...): -# from unyt._array_functions import unwrap as unyt_unwrap +@implements(np.unwrap) +def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586): + from unyt._array_functions import unwrap as unyt_unwrap -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_unwrap(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(p, discont=discont, axis=axis, period=period) + ret_cf = _preserve_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["kw_ca_cfs"]["discont"], + helper_result["kw_ca_cfs"]["period"], + ) + res = unyt_unwrap(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.interp) -# def interp(...): -# from unyt._array_functions import interp as unyt_interp +@implements(np.interp) +def interp(x, xp, fp, left=None, right=None, period=None): + from unyt._array_functions import interp as unyt_interp -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_interp(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + x, xp, fp, left=left, right=right, period=period + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][2]) + res = unyt_interp(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.array_repr) -# def array_repr(...): -# from unyt._array_functions import array_repr as unyt_array_repr +@implements(np.array_repr) +def array_repr(arr, max_line_width=None, precision=None, suppress_small=None): + from unyt._array_functions import array_repr as unyt_array_repr -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_array_repr(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + arr, + max_line_width=max_line_width, + precision=precision, + suppress_small=suppress_small + ) + rep = unyt_array_repr(*helper_result["args"], **helper_result["kwargs"])[:-1] + if hasattr(arr, "comoving"): + rep += f", comoving='{arr.comoving}'" + if hasattr(arr, "cosmo_factor"): + rep += f", cosmo_factor='{arr.cosmo_factor}'" + if hasattr(arr, "valid_transform"): + rep += f", valid_transform='{arr.valid_transform}'" + rep += ")" + return rep @implements(np.linalg.outer) @@ -1729,21 +1785,34 @@ def linalg_outer(x1, x2, /): return _return_helper(res, helper_result, ret_cf) -# @implements(np.trapezoid) -# def trapezoid(...): -# from unyt._array_functions import trapezoid as unyt_trapezoid +@implements(np.trapezoid) +def trapezoid(y, x=None, dx=1.0, axis=-1): + from unyt._array_functions import trapezoid as unyt_trapezoid -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_trapezoid(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args(y, x=x, dx=dx, axis=axis) + if x is None: + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["kw_ca_cfs"]["dx"] + ) + else: + ret_cf = _multiply_cosmo_factor( + helper_result["ca_cfs"][0], helper_result["kw_ca_cfs"]["x"] + ) + res = unyt_trapezoid(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) -# @implements(np.in1d) -# def in1d(...): -# from unyt._array_functions import in1d as unyt_in1d +@implements(np.in1d) +def in1d(ar1, ar2, assume_unique=False, invert=False, *, kind=None): + from unyt._array_functions import isin as unyt_in1d -# helper_result = _prepare_array_func_args(...) -# ret_cf = ...() -# res = unyt_in1d(*helper_result["args"], **helper_result["kwargs"]) -# return _return_helper(res, helper_result, ret_cf, out=out) + helper_result = _prepare_array_func_args( + ar1, ar2, assume_unique=assume_unique, invert=invert, kind=kind + ) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=(ar1, ar2), + ) + res = unyt_in1d(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 0c84887f..1309410d 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -4,6 +4,7 @@ import pytest import os +import warnings import numpy as np import unyt as u from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a @@ -87,6 +88,21 @@ def test_init_from_list_of_unyt_arrays(self): assert hasattr(arr, "comoving") assert isinstance(arr, cosmo_array) + def test_init_from_list_of_cosmo_arrays(self): + arr = cosmo_array( + [ + cosmo_array( + 1, units=u.Mpc, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 1) + ) + for _ in range(5) + ] + ) + assert isinstance(arr, cosmo_array) + assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor( + a ** 1, 1 + ) + assert hasattr(arr, "comoving") and arr.comoving is False + class TestNumpyFunctions: """ @@ -182,27 +198,27 @@ def test_handled_funcs(self): "put_along_axis": (ca(np.arange(3)), np.arange(3), ca(np.arange(3)), 0), "putmask": (ca(np.arange(3)), np.arange(3), ca(np.arange(3))), "searchsorted": (ca(np.arange(3)), ca(np.arange(3))), - # "select": ( - # [np.arange(3) < 1, np.arange(3) > 1], - # [ca(np.arange(3)), ca(np.arange(3))], - # ca(1), - # ), - # "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), + "select": ( + [np.arange(3) < 1, np.arange(3) > 1], + [ca(np.arange(3)), ca(np.arange(3))], + ca(1), + ), + "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))), "sinc": (ca(np.arange(3)),), "clip": (ca(np.arange(3)), ca(1), ca(2)), - # "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), - # "triu": (ca(np.arange(3)),), - # "tril": (ca(np.arange(3)),), - # "einsum": ("ii->i", ca(np.eye(3))), - # "convolve": (ca(np.arange(3)), ca(np.arange(3))), - # "correlate": (ca(np.arange(3)), ca(np.arange(3))), - # "tensordot": (ca(np.eye(3)), ca(np.eye(3))), - # "unwrap": (ca(np.arange(3)),), - # "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), - # "array_repr": (ca(np.arange(3)),), + "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "triu": (ca(np.ones((3, 3))),), + "tril": (ca(np.ones((3, 3))),), + "einsum": ("ii->i", ca(np.eye(3))), + "convolve": (ca(np.arange(3)), ca(np.arange(3))), + "correlate": (ca(np.arange(3)), ca(np.arange(3))), + "tensordot": (ca(np.eye(3)), ca(np.eye(3))), + "unwrap": (ca(np.arange(3)),), + "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "array_repr": (ca(np.arange(3)),), "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), - # "trapezoid": (ca(np.arange(3)),), - # "in1d": (ca(np.arange(3)), ca(np.arange(3))), + "trapezoid": (ca(np.arange(3)),), + "in1d": (ca(np.arange(3)), ca(np.arange(3))), # np deprecated } functions_checked = list() bad_funcs = dict() @@ -210,7 +226,14 @@ def test_handled_funcs(self): ua_args = tuple(to_ua(arg) for arg in args) func = getfunc(fname) try: - ua_result = func(*ua_args) + with warnings.catch_warnings(): + if "savetxt" in fname: + warnings.filterwarnings( + action="ignore", + category=UserWarning, + message="numpy.savetxt does not preserve units or cosmo", + ) + ua_result = func(*ua_args) except u.exceptions.UnytError: raises_unyt_error = True else: @@ -222,7 +245,14 @@ def test_handled_funcs(self): with pytest.raises(u.exceptions.UnytError): result = func(*args) continue - result = func(*args) + with warnings.catch_warnings(): + if "savetxt" in fname: + warnings.filterwarnings( + action="ignore", + category=UserWarning, + message="numpy.savetxt does not preserve units or cosmo", + ) + result = func(*args) if fname.split(".")[-1] in ( "fill_diagonal", "copyto", @@ -463,3 +493,48 @@ def test_reshape_to_scalar(self): def test_iter(self): for cq in ca(np.arange(3)): assert isinstance(cq, cosmo_quantity) + + +class TestCosmoQuantity: + @pytest.mark.parametrize( + "func, args", + [ + ("astype", (float,)), + ("in_units", (u.m,)), + ("byteswap", tuple()), + ("compress", ([True],)), + ("flatten", tuple()), + ("ravel", tuple()), + ("repeat", (1,)), + ("reshape", (1,)), + ("take", ([0],)), + ("transpose", tuple()), + ("view", tuple()), + ], + ) + def test_propagation_func(self, func, args): + cq = cosmo_quantity( + 1, + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a ** 1, 1.0), + valid_transform=True, + ) + res = getattr(cq, func)(*args) + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.valid_transform is True + + @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"]) + def test_propagation_props(self, prop): + cq = cosmo_quantity( + 1, + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a ** 1, 1.0), + valid_transform=True, + ) + res = getattr(cq, prop) + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.valid_transform is True From 922de1b0eaaba6ce5152a23df6edb9fa3c12bd01 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Dec 2024 16:45:13 +0000 Subject: [PATCH 24/29] Overhaul initialisation to allow list of cosmo_arrays as input. --- swiftsimio/objects.py | 142 ++++++++++++++++++++++++++++-------------- 1 file changed, 95 insertions(+), 47 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 58a0fc2e..9ddc283a 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -7,7 +7,7 @@ import unyt from unyt import unyt_array, unyt_quantity -from unyt.array import multiple_output_operators +from unyt.array import multiple_output_operators, _iterable from numbers import Number as numeric_type try: @@ -116,7 +116,7 @@ def __init__(self, message="Could not convert to comoving coordinates"): def _propagate_cosmo_array_attributes(func): def wrapped(self, *args, **kwargs): ret = func(self, *args, **kwargs) - if not type(ret) is cosmo_array: + if not isinstance(ret, cosmo_array): return ret if hasattr(self, "cosmo_factor"): ret.cosmo_factor = self.cosmo_factor @@ -393,7 +393,7 @@ def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): ) -def _prepare_array_func_args(*args, **kwargs): +def _prepare_array_func_args(*args, _default_cm=True, **kwargs): # unyt allows creating a unyt_array from e.g. arrays with heterogenous units # (it probably shouldn't...). # Example: @@ -448,15 +448,27 @@ def _prepare_array_func_args(*args, **kwargs): ret_cm = False else: # mix of comoving and physical inputs - args = [ - arg.to_comoving() if cm[0] and not cm[1] else arg - for arg, cm in zip(args, cms) - ] - kwargs = { - k: kwarg.to_comoving() if kw_cms[k][0] and not kw_cms[k][1] else kwarg - for k, kwarg in kwargs.items() - } - ret_cm = True + # better to modify inplace (convert_to_comoving)? + if _default_cm: + args = [ + arg.to_comoving() if cm[0] and not cm[1] else arg + for arg, cm in zip(args, cms) + ] + kwargs = { + k: kwarg.to_comoving() if kw_cms[k][0] and not kw_cms[k][1] else kwarg + for k, kwarg in kwargs.items() + } + ret_cm = True + else: + args = [ + arg.to_physical() if cm[0] and not cm[1] else arg + for arg, cm in zip(args, cms) + ] + kwargs = { + k: kwarg.to_physical() if kw_cms[k][0] and not kw_cms[k][1] else kwarg + for k, kwarg in kwargs.items() + } + ret_cm = False if len(set(comps + list(kw_comps.values()))) == 1: # all compressions identical, preserve it ret_comp = (comps + list(kw_comps.values()))[0] @@ -694,6 +706,18 @@ def __eq__(self, b): def __ne__(self, b): return not self.__eq__(b) + def __repr__(self): + """ + Print exponent and current scale factor + + Returns + ------- + + str + string to print exponent and current scale factor + """ + return self.__str__() + class cosmo_array(unyt_array): """ @@ -874,39 +898,49 @@ def __new__( cosmo_factor: cosmo_factor - try: - obj = super().__new__( - cls, - input_array, - units=units, - registry=registry, - dtype=dtype, - bypass_validation=bypass_validation, - name=name, - ) - except TypeError: - # Older versions of unyt (before input_units was deprecated) - obj = super().__new__( - cls, - input_array, - units=units, - registry=registry, - dtype=dtype, - bypass_validation=bypass_validation, - input_units=input_units, - name=name, - ) - except TypeError: - # Even older versions of unyt (before name was added) - obj = super().__new__( - cls, - input_array, - units=units, - registry=registry, - dtype=dtype, - bypass_validation=bypass_validation, - input_units=input_units, + if isinstance(input_array, cosmo_array): + if comoving: + input_array.convert_to_comoving() + elif comoving is False: + input_array.convert_to_physical() + else: + comoving = input_array.comoving + cosmo_factor = _preserve_cosmo_factor( + (cosmo_factor is not None, cosmo_factor), + (input_array.cosmo_factor is not None, input_array.cosmo_factor), ) + if not valid_transform: + input_array.convert_to_physical() + if compression != input_array.compression: + compression = None # just drop it + elif isinstance(input_array, np.ndarray): + pass # guard np.ndarray so it doesn't get caught by _iterable in next case + elif _iterable(input_array) and input_array: + if isinstance(input_array[0], cosmo_array): + default_cm = comoving if comoving is not None else True + helper_result = _prepare_array_func_args( + *input_array, _default_cm=default_cm + ) + if comoving is None: + comoving = helper_result["comoving"] + input_array = helper_result["args"] + cosmo_factor = _preserve_cosmo_factor( + (cosmo_factor is not None, cosmo_factor), *helper_result["ca_cfs"] + ) + if not valid_transform: + input_array.convert_to_physical() + if compression != helper_result["compression"]: + compression = None # just drop it + + obj = super().__new__( + cls, + input_array, + units=units, + registry=registry, + dtype=dtype, + bypass_validation=bypass_validation, + name=name, + ) if isinstance(obj, unyt_array) and not isinstance(obj, cls): obj = obj.view(cls) @@ -980,7 +1014,6 @@ def __setstate__(self, state): # Wrap functions that return copies of cosmo_arrays so that our # attributes get passed through: - __getitem__ = _propagate_cosmo_array_attributes(unyt_array.__getitem__) astype = _propagate_cosmo_array_attributes(unyt_array.astype) in_units = _propagate_cosmo_array_attributes(unyt_array.in_units) byteswap = _propagate_cosmo_array_attributes(unyt_array.byteswap) @@ -989,13 +1022,28 @@ def __setstate__(self, state): flatten = _propagate_cosmo_array_attributes(unyt_array.flatten) ravel = _propagate_cosmo_array_attributes(unyt_array.ravel) repeat = _propagate_cosmo_array_attributes(unyt_array.repeat) - reshape = _propagate_cosmo_array_attributes(unyt_array.reshape) swapaxes = _propagate_cosmo_array_attributes(unyt_array.swapaxes) take = _propagate_cosmo_array_attributes(unyt_array.take) transpose = _propagate_cosmo_array_attributes(unyt_array.transpose) view = _propagate_cosmo_array_attributes(unyt_array.view) - # Also wrap some array "attributes": + @_propagate_cosmo_array_attributes + def reshape(self, shape, **kwargs): + reshaped = unyt_array.reshape(self, shape, **kwargs) + if shape == (): + return cosmo_quantity(reshaped) + else: + return reshaped + + @_propagate_cosmo_array_attributes + def __getitem__(self, *args, **kwargs): + item = unyt_array.__getitem__(self, *args, *kwargs) + if item.shape == (): + return cosmo_quantity(item) + else: + return item + + # Also wrap some array "properties": @property def T(self): From ec9e932984c255dad135bb75c8986586de824397 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Dec 2024 16:56:33 +0000 Subject: [PATCH 25/29] Cleanup some caught warnings that no longer happen :) --- tests/test_data.py | 75 +++++++++++++++++----------------------------- 1 file changed, 28 insertions(+), 47 deletions(-) diff --git a/tests/test_data.py b/tests/test_data.py index 232864d2..e6d61b08 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -187,11 +187,8 @@ def test_cell_metadata_is_valid(filename): continue # Give it a little wiggle room. - # Mask_region provides unyt_array, not cosmo_array, anticipate warnings. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert max <= upper * 1.05 - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert min > lower * 0.95 + assert max <= upper * 1.05 + assert min > lower * 0.95 @requires("cosmological_volume_dithered.hdf5") @@ -240,12 +237,8 @@ def test_dithered_cell_metadata_is_valid(filename): continue # Give it a little wiggle room - # Mask_region provides unyt_array, not cosmo_array, anticipate warnings. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert max <= upper * 1.05 - # Mask_region provides unyt_array, not cosmo_array, anticipate warnings. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - assert min > lower * 0.95 + assert max <= upper * 1.05 + assert min > lower * 0.95 @requires("cosmological_volume.hdf5") @@ -274,27 +267,21 @@ def test_reading_select_region_metadata(filename): selected_coordinates = selected_data.gas.coordinates # Now need to repeat teh selection by hand: - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) - ] - ) + subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) + ] + ) # We also need to repeat for the thing we just selected; the cells only give # us an _approximate_ selection! - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - selected_subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip( - selected_data.gas.coordinates.T, restrict - ) - ] - ) + selected_subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(selected_data.gas.coordinates.T, restrict) + ] + ) hand_selected_coordinates = full_data.gas.coordinates[subset_mask] @@ -331,27 +318,21 @@ def test_reading_select_region_metadata_not_spatial_only(filename): selected_coordinates = selected_data.gas.coordinates # Now need to repeat the selection by hand: - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) - ] - ) + subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(full_data.gas.coordinates.T, restrict) + ] + ) # We also need to repeat for the thing we just selected; the cells only give # us an _approximate_ selection! - # Iterating a cosmo_array gives unyt_quantities, anticipate the warning for comparing to cosmo_array. - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - selected_subset_mask = logical_and.reduce( - [ - logical_and(x > y_lower, x < y_upper) - for x, (y_lower, y_upper) in zip( - selected_data.gas.coordinates.T, restrict - ) - ] - ) + selected_subset_mask = logical_and.reduce( + [ + logical_and(x > y_lower, x < y_upper) + for x, (y_lower, y_upper) in zip(selected_data.gas.coordinates.T, restrict) + ] + ) hand_selected_coordinates = full_data.gas.coordinates[subset_mask] From 44a6f59d42d7d21ca0dc8917911b0a63be6a1229 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Dec 2024 16:56:52 +0000 Subject: [PATCH 26/29] Remove some unused imports and variables. --- tests/test_rotate_visualisations.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/tests/test_rotate_visualisations.py b/tests/test_rotate_visualisations.py index 694f6213..1e73aec5 100644 --- a/tests/test_rotate_visualisations.py +++ b/tests/test_rotate_visualisations.py @@ -6,7 +6,6 @@ from swiftsimio.visualisation.rotation import rotation_matrix_from_vector from numpy import array_equal from os import remove -import pytest @requires("cosmological_volume.hdf5") @@ -29,7 +28,6 @@ def test_project(filename): centre = data.gas.coordinates[0] rotate_vec = [0.5, 0.5, 0.5] matrix = rotation_matrix_from_vector(rotate_vec, axis="z") - boxsize = data.metadata.boxsize unrotated = project_gas(data, resolution=1024, project="masses", parallel=True) @@ -67,7 +65,6 @@ def test_slice(filename): centre = data.gas.coordinates[0] rotate_vec = [0.5, 0.5, 0.5] matrix = rotation_matrix_from_vector(rotate_vec, axis="z") - boxsize = data.metadata.boxsize slice_z = centre[2] @@ -114,7 +111,6 @@ def test_render(filename): centre = data.gas.coordinates[0] rotate_vec = [0.5, 0.5, 0.5] matrix = rotation_matrix_from_vector(rotate_vec, axis="z") - boxsize = data.metadata.boxsize unrotated = render_gas(data, resolution=256, project="masses", parallel=True) From 3d8770818c10139115504cfcaa85143e98c5770f Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 19 Dec 2024 17:32:00 +0000 Subject: [PATCH 27/29] Add reshape to cosmo_quantity to enable some functions like meshgrid. --- swiftsimio/objects.py | 15 ++++++++++++++- 1 file changed, 14 insertions(+), 1 deletion(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 9ddc283a..f7361aff 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -1416,7 +1416,7 @@ def __new__( units = getattr(input_scalar, "units", None) else: units = input_units - ret = cosmo_array.__new__( + ret = super().__new__( cls, np.asarray(input_scalar), units, @@ -1432,3 +1432,16 @@ def __new__( if ret.size > 1: raise RuntimeError("unyt_quantity instances must be scalars") return ret + + def reshape(self, *shape, order="C"): + # this is necessary to support some numpy operations + # natively, like numpy.meshgrid, which internally performs + # reshaping, e.g., arr.reshape(1, -1), which doesn't affect the size, + # but does change the object's internal representation to a >0D array + # see https://github.com/yt-project/unyt/issues/224 + if len(shape) == 1: + shape = shape[0] + if shape == () or shape is None: + return super().reshape(shape, order=order) + else: + return cosmo_array(self).reshape(shape, order=order) From 165f42239ab943ac24340a0c781e93d9b6f6e9a9 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 10 Jan 2025 17:10:57 +0000 Subject: [PATCH 28/29] Implement numpy take function. --- swiftsimio/_array_functions.py | 66 +++++++++++++++++++++------------ tests/test_cosmo_array.py | 67 +++++++++++++++++----------------- 2 files changed, 77 insertions(+), 56 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 5d71caf4..0b9850ec 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1342,7 +1342,9 @@ def linalg_tensorsolve(a, b, axes=None): from unyt._array_functions import linalg_tensorsolve as unyt_linalg_tensorsolve helper_result = _prepare_array_func_args(a, b, axes=axes) - ret_cf = _divide_cosmo_factor(helper_result["ca_cfs"][1], helper_result["ca_cfs"][0]) + ret_cf = _divide_cosmo_factor( + helper_result["ca_cfs"][1], helper_result["ca_cfs"][0] + ) res = unyt_linalg_tensorsolve(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1395,15 +1397,15 @@ def linalg_eigvalsh(a, UPLO="L"): @implements(np.savetxt) def savetxt( - fname, - X, - fmt="%.18e", - delimiter=" ", - newline="\n", - header="", - footer="", - comments="# ", - encoding=None + fname, + X, + fmt="%.18e", + delimiter=" ", + newline="\n", + header="", + footer="", + comments="# ", + encoding=None, ): from unyt._array_functions import savetxt as unyt_savetxt @@ -1429,7 +1431,7 @@ def savetxt( warnings.filterwarnings( action="ignore", category=UserWarning, - message="numpy.savetxt does not preserve units" + message="numpy.savetxt does not preserve units", ) unyt_savetxt(*helper_result["args"], **helper_result["kwargs"]) return @@ -1563,7 +1565,7 @@ def select(condlist, choicelist, default=0): res = unyt_select( helper_result["args"][0], helper_result_choicelist["args"], - **helper_result["kwargs"] + **helper_result["kwargs"], ) return _return_helper(res, helper_result, ret_cf) @@ -1592,14 +1594,14 @@ def sinc(x): @implements(np.clip) def clip( - a, - a_min=np._NoValue, - a_max=np._NoValue, - out=None, - *, - min=np._NoValue, - max=np._NoValue, - **kwargs, + a, + a_min=np._NoValue, + a_max=np._NoValue, + out=None, + *, + min=np._NoValue, + max=np._NoValue, + **kwargs, ): from unyt._array_functions import clip as unyt_clip @@ -1667,7 +1669,13 @@ def tril(m, k=0): @implements(np.einsum) def einsum( - subscripts, *operands, out=None, dtype=None, order="K", casting="safe", optimize=False + subscripts, + *operands, + out=None, + dtype=None, + order="K", + casting="safe", + optimize=False, ): from unyt._array_functions import einsum as unyt_einsum @@ -1730,7 +1738,9 @@ def tensordot(a, b, axes=2): def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586): from unyt._array_functions import unwrap as unyt_unwrap - helper_result = _prepare_array_func_args(p, discont=discont, axis=axis, period=period) + helper_result = _prepare_array_func_args( + p, discont=discont, axis=axis, period=period + ) ret_cf = _preserve_cosmo_factor( helper_result["ca_cfs"][0], helper_result["kw_ca_cfs"]["discont"], @@ -1760,7 +1770,7 @@ def array_repr(arr, max_line_width=None, precision=None, suppress_small=None): arr, max_line_width=max_line_width, precision=precision, - suppress_small=suppress_small + suppress_small=suppress_small, ) rep = unyt_array_repr(*helper_result["args"], **helper_result["kwargs"])[:-1] if hasattr(arr, "comoving"): @@ -1816,3 +1826,13 @@ def in1d(ar1, ar2, assume_unique=False, invert=False, *, kind=None): ) res = unyt_in1d(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) + + +@implements(np.take) +def take(a, indices, axis=None, out=None, mode="raise"): + from unyt._array_functions import take as unyt_take + + helper_result = _prepare_array_func_args(a, indices, axis=axis, out=out, mode=mode) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = unyt_take(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf, out=out) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 1309410d..302737e6 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -20,7 +20,7 @@ def getfunc(fname): def ca(x, unit=u.Mpc): - return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) + return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) def to_ua(x): @@ -50,7 +50,7 @@ def test_init_from_ndarray(self): arr = cosmo_array( np.ones(5), units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -61,7 +61,7 @@ def test_init_from_list(self): arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -71,7 +71,7 @@ def test_init_from_list(self): def test_init_from_unyt_array(self): arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -81,7 +81,7 @@ def test_init_from_unyt_array(self): def test_init_from_list_of_unyt_arrays(self): arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) assert hasattr(arr, "cosmo_factor") @@ -92,14 +92,14 @@ def test_init_from_list_of_cosmo_arrays(self): arr = cosmo_array( [ cosmo_array( - 1, units=u.Mpc, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 1) + 1, units=u.Mpc, comoving=False, cosmo_factor=cosmo_factor(a**1, 1) ) for _ in range(5) ] ) assert isinstance(arr, cosmo_array) assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor( - a ** 1, 1 + a**1, 1 ) assert hasattr(arr, "comoving") and arr.comoving is False @@ -219,6 +219,7 @@ def test_handled_funcs(self): "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), "trapezoid": (ca(np.arange(3)),), "in1d": (ca(np.arange(3)), ca(np.arange(3))), # np deprecated + "take": (ca(np.arange(3)), np.arange(3)), } functions_checked = list() bad_funcs = dict() @@ -312,7 +313,7 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), ), ), @@ -323,13 +324,13 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a ** 2, 1.0), + cosmo_factor=cosmo_factor(a**2, 1.0), ), ), ), @@ -341,19 +342,19 @@ def test_handled_funcs(self): [1, 2, 3], u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), cosmo_array( [1, 2, 3], u.K, comoving=False, - cosmo_factor=cosmo_factor(a ** 2, 1.0), + cosmo_factor=cosmo_factor(a**2, 1.0), ), cosmo_array( [1, 2, 3], u.kg, comoving=False, - cosmo_factor=cosmo_factor(a ** 3, 1.0), + cosmo_factor=cosmo_factor(a**3, 1.0), ), ], ), @@ -365,7 +366,7 @@ def test_handled_funcs(self): ( None, cosmo_array( - [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 1.0) + [1, 2, 3], u.s, comoving=False, cosmo_factor=cosmo_factor(a**1, 1.0) ), np.array([1, 2, 3]), ), @@ -382,19 +383,19 @@ def test_histograms(self, func_args, weights, bins_type, density): np.linspace(0, 5, 11), u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.K, comoving=False, - cosmo_factor=cosmo_factor(a ** 2, 1.0), + cosmo_factor=cosmo_factor(a**2, 1.0), ), cosmo_array( np.linspace(0, 5, 11), u.Msun, comoving=False, - cosmo_factor=cosmo_factor(a ** 3, 1.0), + cosmo_factor=cosmo_factor(a**3, 1.0), ), ], }[bins_type] @@ -441,9 +442,9 @@ def test_histograms(self, func_args, weights, bins_type, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a ** -1, 1.0), - np.histogram2d: cosmo_factor(a ** -3, 1.0), - np.histogramdd: cosmo_factor(a ** -6, 1.0), + np.histogram: cosmo_factor(a**-1, 1.0), + np.histogram2d: cosmo_factor(a**-3, 1.0), + np.histogramdd: cosmo_factor(a**-6, 1.0), }[func] ) elif density and isinstance(weights, cosmo_array): @@ -451,9 +452,9 @@ def test_histograms(self, func_args, weights, bins_type, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a ** 0, 1.0), - np.histogram2d: cosmo_factor(a ** -2, 1.0), - np.histogramdd: cosmo_factor(a ** -5, 1.0), + np.histogram: cosmo_factor(a**0, 1.0), + np.histogram2d: cosmo_factor(a**-2, 1.0), + np.histogramdd: cosmo_factor(a**-5, 1.0), }[func] ) elif not density and isinstance(weights, cosmo_array): @@ -461,9 +462,9 @@ def test_histograms(self, func_args, weights, bins_type, density): assert ( result[0].cosmo_factor == { - np.histogram: cosmo_factor(a ** 1, 1.0), - np.histogram2d: cosmo_factor(a ** 1, 1.0), - np.histogramdd: cosmo_factor(a ** 1, 1.0), + np.histogram: cosmo_factor(a**1, 1.0), + np.histogram2d: cosmo_factor(a**1, 1.0), + np.histogramdd: cosmo_factor(a**1, 1.0), }[func] ) ret_bins = { @@ -475,9 +476,9 @@ def test_histograms(self, func_args, weights, bins_type, density): ret_bins, ( [ - cosmo_factor(a ** 1, 1.0), - cosmo_factor(a ** 2, 1.0), - cosmo_factor(a ** 3, 1.0), + cosmo_factor(a**1, 1.0), + cosmo_factor(a**2, 1.0), + cosmo_factor(a**3, 1.0), ] ), ): @@ -517,12 +518,12 @@ def test_propagation_func(self, func, args): 1, u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), valid_transform=True, ) res = getattr(cq, func)(*args) assert res.comoving is False - assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.cosmo_factor == cosmo_factor(a**1, 1.0) assert res.valid_transform is True @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"]) @@ -531,10 +532,10 @@ def test_propagation_props(self, prop): 1, u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), valid_transform=True, ) res = getattr(cq, prop) assert res.comoving is False - assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.cosmo_factor == cosmo_factor(a**1, 1.0) assert res.valid_transform is True From ef43675e9e672dc2309cf8e6229a607191b887cb Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 10 Jan 2025 17:37:33 +0000 Subject: [PATCH 29/29] Wrap np.ndarray.take properly. --- swiftsimio/objects.py | 13 ++++++++++--- 1 file changed, 10 insertions(+), 3 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index f7361aff..ce970514 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -355,7 +355,7 @@ def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): raise ValueError( f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." ) - return cosmo_factor(a ** 0, scale_factor=cf1.scale_factor) + return cosmo_factor(a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): @@ -683,7 +683,7 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): - return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) def __lt__(self, b): return self.a_factor < b.a_factor @@ -1023,10 +1023,17 @@ def __setstate__(self, state): ravel = _propagate_cosmo_array_attributes(unyt_array.ravel) repeat = _propagate_cosmo_array_attributes(unyt_array.repeat) swapaxes = _propagate_cosmo_array_attributes(unyt_array.swapaxes) - take = _propagate_cosmo_array_attributes(unyt_array.take) transpose = _propagate_cosmo_array_attributes(unyt_array.transpose) view = _propagate_cosmo_array_attributes(unyt_array.view) + @_propagate_cosmo_array_attributes + def take(self, indices, **kwargs): + taken = unyt_array.take(self, indices, **kwargs) + if np.ndim(indices) == 0: + return cosmo_quantity(taken) + else: + return cosmo_array(taken) + @_propagate_cosmo_array_attributes def reshape(self, shape, **kwargs): reshaped = unyt_array.reshape(self, shape, **kwargs)