From 6a34c842d92b219dd5a410ce2bcc23618fb3ee3a Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 12 Dec 2024 14:13:12 +0000 Subject: [PATCH 001/125] 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 002/125] 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 003/125] 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 004/125] 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 005/125] 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 006/125] 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 007/125] 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 008/125] 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 009/125] 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 010/125] 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 011/125] 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 012/125] 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 013/125] 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 014/125] 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 015/125] 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 016/125] 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 017/125] 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 018/125] 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 019/125] 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 020/125] 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 021/125] 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 022/125] 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 023/125] 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 024/125] 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 025/125] 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 026/125] 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 027/125] 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 028/125] 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 029/125] 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) From c86b9c91b07f6c5eac1670ee5a1a0c83537c6946 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 28 Jan 2025 17:23:10 +0000 Subject: [PATCH 030/125] Replace some internal uses of unyt_array with cosmo_array. --- swiftsimio/masks.py | 55 +++++++++++++++++++++++++++++++++++---------- tests/test_mask.py | 3 ++- 2 files changed, 45 insertions(+), 13 deletions(-) diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index eb05a248..434fe9d8 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -5,17 +5,21 @@ import warnings -import unyt import h5py import numpy as np from swiftsimio.metadata.objects import SWIFTMetadata -from swiftsimio.objects import InvalidSnapshot +from swiftsimio.objects import ( + InvalidSnapshot, + cosmo_array, + cosmo_quantity, + cosmo_factor, + a, +) from swiftsimio.accelerated import ranges_from_array -from typing import Dict class SWIFTMask(object): @@ -229,11 +233,19 @@ def _unpack_cell_metadata(self): self.counts[key] = counts[sort] # Also need to sort centers in the same way - self.centers = unyt.unyt_array(centers_handle[:][sort], units=self.units.length) + self.centers = cosmo_array( + centers_handle[:][sort], + units=self.units.length, + comoving=True, + cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), + ) # Note that we cannot assume that these are cubic, unfortunately. - self.cell_size = unyt.unyt_array( - metadata_handle.attrs["size"], units=self.units.length + self.cell_size = cosmo_array( + metadata_handle.attrs["size"], + units=self.units.length, + comoving=True, + cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), ) return @@ -242,8 +254,8 @@ def constrain_mask( self, group_name: str, quantity: str, - lower: unyt.array.unyt_quantity, - upper: unyt.array.unyt_quantity, + lower: cosmo_quantity, + upper: cosmo_quantity, ): """ Constrains the mask further for a given particle type, and bounds a @@ -263,10 +275,10 @@ def constrain_mask( quantity : str quantity being constrained - lower : unyt.array.unyt_quantity + lower : ~swiftsimio.objects.cosmo_quantity constraint lower bound - upper : unyt.array.unyt_quantity + upper : ~swiftsimio.objects.cosmo_quantity constraint upper bound See Also @@ -300,13 +312,32 @@ def constrain_mask( handle = handle_dict[quantity] + physical_dict = { + k: v + for k, v in zip(group_metadata.field_names, group_metadata.field_physicals) + } + + physical = physical_dict[quantity] + + cosmologies_dict = { + k: v + for k, v in zip( + group_metadata.field_names, group_metadata.field_cosmologies + ) + } + + cosmology_factor = cosmologies_dict[quantity] + # Load in the relevant data. with h5py.File(self.metadata.filename, "r") as file: # Surprisingly this is faster than just using the boolean # indexing because h5py has slow indexing routines. - data = unyt.unyt_array( - np.take(file[handle], np.where(current_mask)[0], axis=0), units=unit + data = cosmo_array( + np.take(file[handle], np.where(current_mask)[0], axis=0), + units=unit, + comoving=not physical, + cosmo_factor=cosmology_factor, ) new_mask = np.logical_and.reduce([data > lower, data <= upper]) diff --git a/tests/test_mask.py b/tests/test_mask.py index 1a75b498..f6af42e1 100644 --- a/tests/test_mask.py +++ b/tests/test_mask.py @@ -6,7 +6,8 @@ from swiftsimio import load, mask import numpy as np -from unyt import unyt_array as array, dimensionless +from unyt import dimensionless +from swiftsimio import cosmo_array as array @requires("cosmological_volume.hdf5") From 59594ff80f29e028593bf48271338d92ca8bba59 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 29 Jan 2025 12:46:52 +0000 Subject: [PATCH 031/125] Replace unyt array/quantity with cosmo array/quantity in tests. --- tests/test_soap.py | 22 +++++++++++++++------- 1 file changed, 15 insertions(+), 7 deletions(-) diff --git a/tests/test_soap.py b/tests/test_soap.py index 66ec93fa..3469ebf8 100644 --- a/tests/test_soap.py +++ b/tests/test_soap.py @@ -4,8 +4,8 @@ from tests.helper import requires -from swiftsimio import load, mask -import unyt +from swiftsimio import load, mask, cosmo_quantity +from swiftsimio.objects import cosmo_quantity, cosmo_factor, a @requires("soap_example.hdf5") @@ -43,8 +43,18 @@ def test_soap_can_mask_non_spatial(filename): def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): this_mask = mask(filename, spatial_only=False) - lower = unyt.unyt_quantity(1e5, "Msun") - upper = unyt.unyt_quantity(1e13, "Msun") + lower = cosmo_quantity( + 1e5, + "Msun", + comoving=True, + cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), + ) + upper = cosmo_quantity( + 1e13, + "Msun", + comoving=True, + cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), + ) this_mask.constrain_mask( "spherical_overdensity_200_mean", "total_mass", lower, upper ) @@ -60,9 +70,7 @@ def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): masses2 = data2.spherical_overdensity_200_mean.total_mass # Manually mask - custom_mask = ( - unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) >= lower - ) & (unyt.unyt_array(masses2.to_value(masses2.units), masses2.units) <= upper) + custom_mask = (masses2 >= lower) & (masses2 <= upper) assert len(masses2[custom_mask]) == len(masses) From 791b24a116d9734794b49b0a89f7e115c748d574 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 29 Jan 2025 12:48:05 +0000 Subject: [PATCH 032/125] Cleanup some unused imports and unused variables. --- tests/test_visualisation.py | 18 +++++++----------- 1 file changed, 7 insertions(+), 11 deletions(-) diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 99a91790..fabdcd9d 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -7,7 +7,6 @@ project_pixel_grid, ) from swiftsimio.visualisation.slice import ( - slice_scatter, slice_scatter_parallel, slice_gas, ) @@ -23,13 +22,11 @@ from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths from swiftsimio.visualisation.projection_backends import ( backends as projection_backends, - backends_parallel as projection_backends_parallel, ) from swiftsimio.visualisation.slice_backends import ( backends as slice_backends, backends_parallel as slice_backends_parallel, ) -from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths from swiftsimio.optional_packages import CudaSupportError, CUDA_AVAILABLE from swiftsimio.objects import cosmo_array, a from unyt.array import unyt_array @@ -87,7 +84,7 @@ def test_scatter_mass_conservation(): for resolution in resolutions: image = scatter(x, y, m, h, resolution, 1.0, 1.0) - mass_in_image = image.sum() / (resolution ** 2) + mass_in_image = image.sum() / (resolution**2) # Check mass conservation to 5% assert np.isclose(mass_in_image, total_mass, 0.05) @@ -364,7 +361,7 @@ def test_comoving_versus_physical(filename): img = func(data, resolution=256, project="masses") # check that we get a physical result assert not img.comoving - assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 + assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 # densities are still compatible with physical img = func(data, resolution=256, project="densities") assert not img.comoving @@ -426,14 +423,14 @@ def test_panel_rendering(filename): plt.imsave( "panels_added.png", - plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), + plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1))), ) projected = project_gas(data, res, "masses", backend="renormalised") plt.imsave( "projected.png", - plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), + plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)), ) fullstack = np.zeros((res, res)) @@ -606,7 +603,6 @@ def test_folding_deposit(): y = np.array([100, 200]) z = np.array([100, 200]) m = np.array([1, 1]) - h = np.array([1e-10, 1e-10]) res = 256 boxsize = 1.0 * res @@ -634,8 +630,8 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v + mean_density_deposit = (np.sum(deposition) / npix**3).to("Msun / kpc**3").v + mean_density_volume = (np.sum(volume) / npix**3).to("Msun / kpc**3").v mean_density_calculated = ( (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) .to("Msun / kpc**3") @@ -697,7 +693,7 @@ def test_dark_matter_power_spectrum(filename, save=False): folds[folding] = deposition - folding_output[2 ** folding] = (k, power_spectrum, scatter) + folding_output[2**folding] = (k, power_spectrum, scatter) # Now try doing them all together at once. From 6501b3f8c1ab12c1a09554c4cd4c881773c609d0 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 29 Jan 2025 16:13:03 +0000 Subject: [PATCH 033/125] Remove duplicate comoving=... in __repr__. --- swiftsimio/objects.py | 9 +-------- 1 file changed, 1 insertion(+), 8 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index ce970514..dc403af7 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -978,14 +978,7 @@ def __str__(self): return super().__str__() + " " + comoving_str def __repr__(self): - if self.comoving: - comoving_str = ", comoving=True)" - elif self.comoving is None: - comoving_str = ")" - else: - comoving_str = ", comoving=False)" - # Remove final parenthesis and append comoving flag - return super().__repr__()[:-1] + comoving_str + return super().__repr__() def __reduce__(self): """ From efa66d5ab28f3a04a11313f7b8a24e529c819809 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 29 Jan 2025 17:03:30 +0000 Subject: [PATCH 034/125] Prepare to handle np functions that 'just work' in unyt but don't for us. --- swiftsimio/_array_functions.py | 32 ++++++- swiftsimio/objects.py | 4 +- tests/test_cosmo_array.py | 149 ++++++++++++++++++++++++++++++++- 3 files changed, 176 insertions(+), 9 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 0b9850ec..a9158a35 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -15,19 +15,22 @@ _return_without_cosmo_factor, ) -_HANDLED_FUNCTIONS = dict() +_HANDLED_FUNCTIONS = {} def _return_helper(res, helper_result, ret_cf, out=None): if out is None: - if isinstance(res, unyt_quantity): + if isinstance(res, cosmo_array) and res.shape == (): + # happens when handling a function that unyt didn't handle explicitly + return cosmo_quantity(res) + elif isinstance(res, unyt_quantity) and not isinstance(res, cosmo_quantity): return cosmo_quantity( res, comoving=helper_result["comoving"], cosmo_factor=ret_cf, compression=helper_result["compression"], ) - elif isinstance(res, unyt_array): + elif isinstance(res, unyt_array) and not isinstance(res, cosmo_array): return cosmo_array( res, comoving=helper_result["comoving"], @@ -43,7 +46,16 @@ def _return_helper(res, helper_result, ret_cf, out=None): out.cosmo_factor = ret_cf if hasattr(out, "compression"): out.compression = helper_result["compression"] - return cosmo_array( # confused, do we set out, or return? + if res.shape == (): + return cosmo_quantity( + res.to_value(res.units), + res.units, + bypass_validation=True, + comoving=helper_result["comoving"], + cosmo_factor=ret_cf, + compression=helper_result["compression"], + ) + return cosmo_array( res.to_value(res.units), res.units, bypass_validation=True, @@ -1836,3 +1848,15 @@ def take(a, indices, axis=None, out=None, mode="raise"): 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) + + +@implements(np.amax) +def amax( + a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, where=np._NoValue +): + helper_result = _prepare_array_func_args( + a, axis=axis, out=out, keepdims=keepdims, initial=initial, where=where + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = np.amax._implementation(*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 dc403af7..1ead7f73 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -1308,7 +1308,9 @@ def __array_function__(self, func, types, args, kwargs): # leaving other arguments a chance to take the lead return NotImplemented - if func not in _HANDLED_FUNCTIONS and func in _UNYT_HANDLED_FUNCTIONS: + if func in _HANDLED_FUNCTIONS: + return _HANDLED_FUNCTIONS[func](*args, **kwargs) + elif 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: diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 302737e6..1a7e6b3e 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -105,14 +105,16 @@ def test_init_from_list_of_cosmo_arrays(self): class TestNumpyFunctions: - """ - Functions specially handled by unyt risk silently casting to unyt_array/unyt_quantity. - """ - def test_handled_funcs(self): + def test_explicitly_handled_funcs(self): + """ + Make sure we at least handle everything that unyt does, and anything that + 'just worked' for unyt but that we need to handle by hand. + """ from unyt._array_functions import _HANDLED_FUNCTIONS functions_to_check = { + # FUNCTIONS UNYT HANDLES EXPLICITLY: "array2string": (ca(np.arange(3)),), "dot": (ca(np.arange(3)), ca(np.arange(3))), "vdot": (ca(np.arange(3)), ca(np.arange(3))), @@ -220,6 +222,145 @@ def test_handled_funcs(self): "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 THAT UNYT DOESN'T HANDLE EXPLICITLY (THEY "JUST WORK"): + "all": (ca(np.arange(3)),), + "amax": (ca(np.arange(3)),), + "amin": (ca(np.arange(3)),), + # np.angle, # expects complex numbers + # np.any, # works out of the box (tested) + # np.append, # we get it for free with np.concatenate (tested) + # np.apply_along_axis, # works out of the box (tested) + # np.argmax, # returns pure numbers + # np.argmin, # returns pure numbers + # np.argpartition, # returns pure numbers + # np.argsort, # returns pure numbers + # np.argwhere, # returns pure numbers + # np.array_str, # hooks into __str__ + # np.atleast_1d, # works out of the box (tested) + # np.atleast_2d, # works out of the box (tested) + # np.atleast_3d, # works out of the box (tested) + # np.average, # works out of the box (tested) + # np.can_cast, # works out of the box (tested) + # np.common_type, # works out of the box (tested) + # np.result_type, # works out of the box (tested) + # np.iscomplex, # works out of the box (tested) + # np.iscomplexobj, # works out of the box (tested) + # np.isreal, # works out of the box (tested) + # np.isrealobj, # works out of the box (tested) + # np.nan_to_num, # works out of the box (tested) + # np.nanargmax, # return pure numbers + # np.nanargmin, # return pure numbers + # np.nanmax, # works out of the box (tested) + # np.nanmean, # works out of the box (tested) + # np.nanmedian, # works out of the box (tested) + # np.nanmin, # works out of the box (tested) + # np.trim_zeros, # works out of the box (tested) + # np.max, # works out of the box (tested) + # np.mean, # works out of the box (tested) + # np.median, # works out of the box (tested) + # np.min, # works out of the box (tested) + # np.ndim, # return pure numbers + # np.shape, # returns pure numbers + # np.size, # returns pure numbers + # np.sort, # works out of the box (tested) + # np.sum, # works out of the box (tested) + # np.repeat, # works out of the box (tested) + # np.tile, # works out of the box (tested) + # np.shares_memory, # works out of the box (tested) + # np.nonzero, # works out of the box (tested) + # np.count_nonzero, # returns pure numbers + # np.flatnonzero, # works out of the box (tested) + # np.isneginf, # works out of the box (tested) + # np.isposinf, # works out of the box (tested) + # np.empty_like, # works out of the box (tested) + # np.full_like, # works out of the box (tested) + # np.ones_like, # works out of the box (tested) + # np.zeros_like, # works out of the box (tested) + # np.copy, # works out of the box (tested) + # np.meshgrid, # works out of the box (tested) + # np.transpose, # works out of the box (tested) + # np.reshape, # works out of the box (tested) + # np.resize, # works out of the box (tested) + # np.roll, # works out of the box (tested) + # np.rollaxis, # works out of the box (tested) + # np.rot90, # works out of the box (tested) + # np.expand_dims, # works out of the box (tested) + # np.squeeze, # works out of the box (tested) + # np.flip, # works out of the box (tested) + # np.fliplr, # works out of the box (tested) + # np.flipud, # works out of the box (tested) + # np.delete, # works out of the box (tested) + # np.partition, # works out of the box (tested) + # np.broadcast_to, # works out of the box (tested) + # np.broadcast_arrays, # works out of the box (tested) + # np.split, # works out of the box (tested) + # np.array_split, # works out of the box (tested) + # np.dsplit, # works out of the box (tested) + # np.hsplit, # works out of the box (tested) + # np.vsplit, # works out of the box (tested) + # np.swapaxes, # works out of the box (tested) + # np.moveaxis, # works out of the box (tested) + # np.nansum, # works out of the box (tested) + # np.std, # works out of the box (tested) + # np.nanstd, # works out of the box (tested) + # np.nanvar, # works out of the box (tested) + # np.nanprod, # works out of the box (tested) + # np.diag, # works out of the box (tested) + # np.diag_indices_from, # returns pure numbers + # np.diagflat, # works out of the box (tested) + # np.diagonal, # works out of the box (tested) + # np.ravel, # returns pure numbers + # np.ravel_multi_index, # returns pure numbers + # np.unravel_index, # returns pure numbers + # np.fix, # works out of the box (tested) + # np.round, # is implemented via np.around + # np.may_share_memory, # returns pure numbers (booleans) + # np.linalg.matrix_power, # works out of the box (tested) + # np.linalg.cholesky, # works out of the box (tested) + # np.linalg.multi_dot, # works out of the box (tested) + # np.linalg.matrix_rank, # returns pure numbers + # np.linalg.qr, # works out of the box (tested) + # np.linalg.slogdet, # undefined units + # np.linalg.cond, # works out of the box (tested) + # np.gradient, # works out of the box (tested) + # np.cumsum, # works out of the box (tested) + # np.nancumsum, # works out of the box (tested) + # np.nancumprod, # we get it for free with np.cumprod (tested) + # np.bincount, # works out of the box (tested) + # np.unique, # works out of the box (tested) + # np.min_scalar_type, # returns dtypes + # np.extract, # works out of the box (tested) + # np.setxor1d, # we get it for free with previously implemented functions (tested) + # np.lexsort, # returns pure numbers + # np.digitize, # returns pure numbers + # np.tril_indices_from, # returns pure numbers + # np.triu_indices_from, # returns pure numbers + # np.imag, # works out of the box (tested) + # np.real, # works out of the box (tested) + # np.real_if_close, # works out of the box (tested) + # np.einsum_path, # returns pure numbers + # np.cov, # returns pure numbers + # np.corrcoef, # returns pure numbers + # np.compress, # works out of the box (tested) + # np.take_along_axis, # works out of the box (tested) + # # the following all work out of the box (tested): + # np.linalg.cross, + # np.linalg.diagonal, + # np.linalg.matmul, + # np.linalg.matrix_norm, + # np.linalg.matrix_transpose, + # np.linalg.svdvals, + # np.linalg.tensordot, + # np.linalg.trace, + # np.linalg.vecdot, + # np.linalg.vector_norm, + # np.astype, + # np.matrix_transpose, + # np.unique_all, + # np.unique_counts, + # np.unique_inverse, + # np.unique_values, + # np.vecdot, } functions_checked = list() bad_funcs = dict() From 75078d2d7f07a01a753e6d68986b3b1c5c818e84 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 29 Jan 2025 17:04:45 +0000 Subject: [PATCH 035/125] Make use of cosmo_array being supported in numpy in power_spectrum functions in visualisation (partial, test failing pending completing np support). --- swiftsimio/visualisation/power_spectrum.py | 58 ++++++++++------------ tests/test_visualisation.py | 20 +++++--- 2 files changed, 41 insertions(+), 37 deletions(-) diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index cdd56f17..2940f36e 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -199,7 +199,7 @@ def render_to_deposit( """ # Get the positions and masses - folding = 2.0 ** folding + folding = 2.0**folding positions = data.coordinates quantity = getattr(data, project) @@ -244,10 +244,10 @@ def render_to_deposit( units = 1.0 / ( data.metadata.boxsize[0] * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) units *= quantity.units - new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor ** 3) + new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor**3) return cosmo_array( deposition, comoving=comoving, cosmo_factor=new_cosmo_factor, units=units @@ -399,7 +399,7 @@ def folded_depositions_to_power_spectrum( if folding != final_folding: cutoff_wavenumber = ( - 2.0 ** folding * np.min(depositions[folding].shape) / np.min(box_size) + 2.0**folding * np.min(depositions[folding].shape) / np.min(box_size) ) if cutoff_above_wavenumber_fraction is not None: @@ -424,7 +424,7 @@ def folded_depositions_to_power_spectrum( corrected_wavenumber_centers[prefer_bins] = folded_wavenumber_centers[ prefer_bins ].to(corrected_wavenumber_centers.units) - folding_tracker[prefer_bins] = 2.0 ** folding + folding_tracker[prefer_bins] = 2.0**folding contributed_counts[prefer_bins] = folded_counts[prefer_bins] elif transition == "average": @@ -457,7 +457,7 @@ def folded_depositions_to_power_spectrum( # For debugging, we calculate an effective fold number. folding_tracker[use_bins] = ( - (folding_tracker * existing_weight + (2.0 ** folding) * new_weight) + (folding_tracker * existing_weight + (2.0**folding) * new_weight) / transition_norm )[use_bins] @@ -474,31 +474,31 @@ def folded_depositions_to_power_spectrum( def deposition_to_power_spectrum( - deposition: unyt.unyt_array, + deposition: cosmo_array, box_size: cosmo_array, folding: int = 0, - cross_deposition: Optional[unyt.unyt_array] = None, - wavenumber_bins: Optional[unyt.unyt_array] = None, + cross_deposition: Optional[cosmo_array] = None, + wavenumber_bins: Optional[cosmo_array] = None, workers: Optional[int] = None, shot_noise_norm: Optional[float] = None, -) -> Tuple[unyt.unyt_array]: +) -> Tuple[cosmo_array]: """ Convert a deposition to a power spectrum, by default using a linear binning strategy. Parameters ---------- - deposition: unyt.unyt_array[float32, float32, float32] + deposition: ~swiftsimio.objects.cosmo_array[float32, float32, float32] The deposition to convert to a power spectrum. - box_size: cosmo_array + box_size: ~swiftsimio.objects.cosmo_array The box size of the deposition, from the dataset. folding: int The folding number (i.e. box-size is divided by 2^folding) that was used here. - cross_deposition: unyt.unyt_array[float32, float32, float32] + cross_deposition: ~swiftsimio.objects.cosmo_array[float32, float32, float32] An optional second deposition to cross-correlate with the first. If not provided, we assume you want an auto-spectrum. - wavenumber_bins: unyt.unyt_array[float32], optional + wavenumber_bins: ~swiftsimio.objects.cosmo_array[float32], optional Optionally you can provide the specific bins that you would like to use. workers: Optional[int] The number of threads to use. @@ -509,12 +509,12 @@ def deposition_to_power_spectrum( Returns ------- - wavenumber_centers: unyt.unyt_array[float32] + wavenumber_centers: ~swiftsimio.objects.cosmo_array[float32] The k-values of the power spectrum, with units. These are the real bin centers, calculated from the mean value of k that was used in the binning process. - power_spectrum: unyt.unyt_array[float32] + power_spectrum: ~swiftsimio.objects.cosmo_array[float32] The power spectrum, with units. binned_counts: np.array[int32] @@ -538,20 +538,20 @@ def deposition_to_power_spectrum( deposition.shape == cross_deposition.shape ), "Depositions must have the same shape" - folding = 2.0 ** folding + folding = 2.0**folding box_size_folded = box_size[0] / folding npix = deposition.shape[0] - mean_deposition = np.mean(deposition.v) - overdensity = (deposition.v - mean_deposition) / mean_deposition + mean_deposition = np.mean(deposition) + overdensity = (deposition - mean_deposition) / mean_deposition fft = scipy.fft.fftn(overdensity / np.prod(deposition.shape), workers=workers) if cross_deposition is not None: - mean_cross_deposition = np.mean(cross_deposition.v) + mean_cross_deposition = np.mean(cross_deposition) cross_overdensity = ( - cross_deposition.v - mean_cross_deposition + cross_deposition - mean_cross_deposition ) / mean_cross_deposition conj_fft = scipy.fft.fftn( @@ -560,7 +560,7 @@ def deposition_to_power_spectrum( else: conj_fft = fft.conj() - fourier_amplitudes = (fft * conj_fft).real * box_size_folded ** 3 + fourier_amplitudes = (fft * conj_fft).real * box_size_folded**3 # Calculate k-value spacing (centered FFT) dk = 2 * np.pi / (box_size_folded) @@ -578,7 +578,7 @@ def deposition_to_power_spectrum( else: kbins = wavenumber_bins - binned_amplitudes = np.histogram(knrm, bins=kbins, weights=fourier_amplitudes.v)[0] + binned_amplitudes = np.histogram(knrm, bins=kbins, weights=fourier_amplitudes)[0] binned_counts = np.histogram(knrm, bins=kbins)[0] # Also compute the 'real' average wavenumber point contributing to this bin. binned_wavenumbers = np.histogram(knrm, bins=kbins, weights=knrm)[0] @@ -592,21 +592,17 @@ def deposition_to_power_spectrum( divisor[zero_mask] = 1 # Correct for folding - binned_amplitudes *= folding ** 3 + binned_amplitudes *= folding**3 # Correct units and names - wavenumbers = unyt.unyt_array( - binned_wavenumbers / divisor, units=knrm.units, name="Wavenumber $k$" - ) + wavenumbers = binned_wavenumbers / divisor + wavenumbers.name = "Wavenumber $k$" shot_noise = ( (box_size[0] ** 3 / shot_noise_norm) if shot_noise_norm is not None else 0.0 ) - power_spectrum = ( - unyt.unyt_array((binned_amplitudes) / divisor, units=fourier_amplitudes.units) - - shot_noise - ) + power_spectrum = (binned_amplitudes / divisor) - shot_noise power_spectrum.name = "Power Spectrum $P(k)$" diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index fabdcd9d..4951d3b1 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -28,7 +28,7 @@ backends_parallel as slice_backends_parallel, ) from swiftsimio.optional_packages import CudaSupportError, CUDA_AVAILABLE -from swiftsimio.objects import cosmo_array, a +from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a from unyt.array import unyt_array import unyt @@ -654,13 +654,21 @@ def test_dark_matter_power_spectrum(filename, save=False): # Collate a bunch of raw depositions folds = {} - min_k = 1e-2 / unyt.Mpc - max_k = 1e2 / unyt.Mpc - - bins = unyt.unyt_array( - np.logspace(np.log10(min_k.v), np.log10(max_k.v), 32), units=min_k.units + min_k = cosmo_quantity( + 1e-2, + unyt.Mpc**-1, + comoving=True, + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + ) + max_k = cosmo_quantity( + 1e2, + unyt.Mpc**-1, + comoving=True, + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), ) + bins = np.geomspace(min_k, max_k, 32) + output = {} for npix in [32, 128]: # Deposit the particles From cb08f9d1d866a58de3b4b6eff1c5c59824b12bf1 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 30 Jan 2025 11:53:29 +0000 Subject: [PATCH 036/125] Make checks on numpy functions more stringent and fix revealed bugs. --- swiftsimio/_array_functions.py | 26 ++- swiftsimio/objects.py | 26 ++- tests/test_cosmo_array.py | 285 +++++++++++++++++---------------- 3 files changed, 194 insertions(+), 143 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index a9158a35..ac941c52 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1,7 +1,6 @@ import warnings import numpy as np from unyt import unyt_quantity, unyt_array -from unyt._array_functions import implements from .objects import ( cosmo_array, cosmo_quantity, @@ -18,6 +17,17 @@ _HANDLED_FUNCTIONS = {} +def implements(numpy_function): + """Register an __array_function__ implementation for unyt_array objects.""" + + # See NEP 18 https://numpy.org/neps/nep-0018-array-function-protocol.html + def decorator(func): + _HANDLED_FUNCTIONS[numpy_function] = func + return func + + return decorator + + def _return_helper(res, helper_result, ret_cf, out=None): if out is None: if isinstance(res, cosmo_array) and res.shape == (): @@ -1707,7 +1717,7 @@ def einsum( *helper_result_operands["args"], **helper_result["kwargs"], ) - return _return_helper(res, helper_result, ret_cf, out=out) + return _return_helper(res, helper_result_operands, ret_cf, out=out) @implements(np.convolve) @@ -1860,3 +1870,15 @@ def amax( ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) res = np.amax._implementation(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf, out=out) + + +@implements(np.amin) +def amin( + a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, where=np._NoValue +): + helper_result = _prepare_array_func_args( + a, axis=axis, out=out, keepdims=keepdims, initial=initial, where=where + ) + ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + res = np.amin._implementation(*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 1ead7f73..4f1083b9 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -1408,16 +1408,32 @@ def __new__( Description of the compression filters that were applied to that array in the hdf5 file. """ - 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 + + units = getattr(input_scalar, "units", None) if units is None else units + name = getattr(input_scalar, "name", None) if name is None else name + cosmo_factor = ( + getattr(input_scalar, "cosmo_factor", None) + if cosmo_factor is None + else cosmo_factor + ) + comoving = ( + getattr(input_scalar, "comoving", None) if comoving is None else comoving + ) + valid_transform = ( + getattr(input_scalar, "valid_transform", None) + if valid_transform is not None + else valid_transform + ) + compression = ( + getattr(input_scalar, "compression", None) + if compression is None + else compression + ) ret = super().__new__( cls, np.asarray(input_scalar), diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 1a7e6b3e..413c0e27 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -23,6 +23,13 @@ def ca(x, unit=u.Mpc): return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) +def arg_to_ua(arg): + if type(arg) in (list, tuple): + return type(arg)([arg_to_ua(a) for a in arg]) + else: + return to_ua(arg) + + def to_ua(x): return u.unyt_array(x.to_value(x.units), x.units) if hasattr(x, "comoving") else x @@ -42,6 +49,10 @@ def check_result(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)) + if isinstance(x_c, cosmo_array): # includes cosmo_quantity + assert x_c.comoving is False + if x_c.units != u.dimensionless: + assert x_c.cosmo_factor is not None return @@ -226,146 +237,148 @@ def test_explicitly_handled_funcs(self): "all": (ca(np.arange(3)),), "amax": (ca(np.arange(3)),), "amin": (ca(np.arange(3)),), - # np.angle, # expects complex numbers - # np.any, # works out of the box (tested) - # np.append, # we get it for free with np.concatenate (tested) - # np.apply_along_axis, # works out of the box (tested) - # np.argmax, # returns pure numbers - # np.argmin, # returns pure numbers - # np.argpartition, # returns pure numbers - # np.argsort, # returns pure numbers - # np.argwhere, # returns pure numbers - # np.array_str, # hooks into __str__ - # np.atleast_1d, # works out of the box (tested) - # np.atleast_2d, # works out of the box (tested) - # np.atleast_3d, # works out of the box (tested) - # np.average, # works out of the box (tested) - # np.can_cast, # works out of the box (tested) - # np.common_type, # works out of the box (tested) - # np.result_type, # works out of the box (tested) - # np.iscomplex, # works out of the box (tested) - # np.iscomplexobj, # works out of the box (tested) - # np.isreal, # works out of the box (tested) - # np.isrealobj, # works out of the box (tested) - # np.nan_to_num, # works out of the box (tested) - # np.nanargmax, # return pure numbers - # np.nanargmin, # return pure numbers - # np.nanmax, # works out of the box (tested) - # np.nanmean, # works out of the box (tested) - # np.nanmedian, # works out of the box (tested) - # np.nanmin, # works out of the box (tested) - # np.trim_zeros, # works out of the box (tested) - # np.max, # works out of the box (tested) - # np.mean, # works out of the box (tested) - # np.median, # works out of the box (tested) - # np.min, # works out of the box (tested) - # np.ndim, # return pure numbers - # np.shape, # returns pure numbers - # np.size, # returns pure numbers - # np.sort, # works out of the box (tested) - # np.sum, # works out of the box (tested) - # np.repeat, # works out of the box (tested) - # np.tile, # works out of the box (tested) - # np.shares_memory, # works out of the box (tested) - # np.nonzero, # works out of the box (tested) - # np.count_nonzero, # returns pure numbers - # np.flatnonzero, # works out of the box (tested) - # np.isneginf, # works out of the box (tested) - # np.isposinf, # works out of the box (tested) - # np.empty_like, # works out of the box (tested) - # np.full_like, # works out of the box (tested) - # np.ones_like, # works out of the box (tested) - # np.zeros_like, # works out of the box (tested) - # np.copy, # works out of the box (tested) - # np.meshgrid, # works out of the box (tested) - # np.transpose, # works out of the box (tested) - # np.reshape, # works out of the box (tested) - # np.resize, # works out of the box (tested) - # np.roll, # works out of the box (tested) - # np.rollaxis, # works out of the box (tested) - # np.rot90, # works out of the box (tested) - # np.expand_dims, # works out of the box (tested) - # np.squeeze, # works out of the box (tested) - # np.flip, # works out of the box (tested) - # np.fliplr, # works out of the box (tested) - # np.flipud, # works out of the box (tested) - # np.delete, # works out of the box (tested) - # np.partition, # works out of the box (tested) - # np.broadcast_to, # works out of the box (tested) - # np.broadcast_arrays, # works out of the box (tested) - # np.split, # works out of the box (tested) - # np.array_split, # works out of the box (tested) - # np.dsplit, # works out of the box (tested) - # np.hsplit, # works out of the box (tested) - # np.vsplit, # works out of the box (tested) - # np.swapaxes, # works out of the box (tested) - # np.moveaxis, # works out of the box (tested) - # np.nansum, # works out of the box (tested) - # np.std, # works out of the box (tested) - # np.nanstd, # works out of the box (tested) - # np.nanvar, # works out of the box (tested) - # np.nanprod, # works out of the box (tested) - # np.diag, # works out of the box (tested) - # np.diag_indices_from, # returns pure numbers - # np.diagflat, # works out of the box (tested) - # np.diagonal, # works out of the box (tested) - # np.ravel, # returns pure numbers - # np.ravel_multi_index, # returns pure numbers - # np.unravel_index, # returns pure numbers - # np.fix, # works out of the box (tested) - # np.round, # is implemented via np.around - # np.may_share_memory, # returns pure numbers (booleans) - # np.linalg.matrix_power, # works out of the box (tested) - # np.linalg.cholesky, # works out of the box (tested) - # np.linalg.multi_dot, # works out of the box (tested) - # np.linalg.matrix_rank, # returns pure numbers - # np.linalg.qr, # works out of the box (tested) - # np.linalg.slogdet, # undefined units - # np.linalg.cond, # works out of the box (tested) - # np.gradient, # works out of the box (tested) - # np.cumsum, # works out of the box (tested) - # np.nancumsum, # works out of the box (tested) - # np.nancumprod, # we get it for free with np.cumprod (tested) - # np.bincount, # works out of the box (tested) - # np.unique, # works out of the box (tested) - # np.min_scalar_type, # returns dtypes - # np.extract, # works out of the box (tested) - # np.setxor1d, # we get it for free with previously implemented functions (tested) - # np.lexsort, # returns pure numbers - # np.digitize, # returns pure numbers - # np.tril_indices_from, # returns pure numbers - # np.triu_indices_from, # returns pure numbers - # np.imag, # works out of the box (tested) - # np.real, # works out of the box (tested) - # np.real_if_close, # works out of the box (tested) - # np.einsum_path, # returns pure numbers - # np.cov, # returns pure numbers - # np.corrcoef, # returns pure numbers - # np.compress, # works out of the box (tested) - # np.take_along_axis, # works out of the box (tested) - # # the following all work out of the box (tested): - # np.linalg.cross, - # np.linalg.diagonal, - # np.linalg.matmul, - # np.linalg.matrix_norm, - # np.linalg.matrix_transpose, - # np.linalg.svdvals, - # np.linalg.tensordot, - # np.linalg.trace, - # np.linalg.vecdot, - # np.linalg.vector_norm, - # np.astype, - # np.matrix_transpose, - # np.unique_all, - # np.unique_counts, - # np.unique_inverse, - # np.unique_values, - # np.vecdot, + # angle, # expects complex numbers + # any, # works out of the box (tested) + # append, # we get it for free with np.concatenate (tested) + # apply_along_axis, # works out of the box (tested) + # argmax, # returns pure numbers + # argmin, # returns pure numbers + # argpartition, # returns pure numbers + # argsort, # returns pure numbers + # argwhere, # returns pure numbers + # array_str, # hooks into __str__ + # atleast_1d, # works out of the box (tested) + # atleast_2d, # works out of the box (tested) + # atleast_3d, # works out of the box (tested) + # average, # works out of the box (tested) + # can_cast, # works out of the box (tested) + # common_type, # works out of the box (tested) + # result_type, # works out of the box (tested) + # iscomplex, # works out of the box (tested) + # iscomplexobj, # works out of the box (tested) + # isreal, # works out of the box (tested) + # isrealobj, # works out of the box (tested) + # nan_to_num, # works out of the box (tested) + # nanargmax, # return pure numbers + # nanargmin, # return pure numbers + # nanmax, # works out of the box (tested) + # nanmean, # works out of the box (tested) + # nanmedian, # works out of the box (tested) + # nanmin, # works out of the box (tested) + # trim_zeros, # works out of the box (tested) + # max, # works out of the box (tested) + # mean, # works out of the box (tested) + # median, # works out of the box (tested) + # min, # works out of the box (tested) + # ndim, # return pure numbers + # shape, # returns pure numbers + # size, # returns pure numbers + # sort, # works out of the box (tested) + # sum, # works out of the box (tested) + # repeat, # works out of the box (tested) + # tile, # works out of the box (tested) + # shares_memory, # works out of the box (tested) + # nonzero, # works out of the box (tested) + # count_nonzero, # returns pure numbers + # flatnonzero, # works out of the box (tested) + # isneginf, # works out of the box (tested) + # isposinf, # works out of the box (tested) + # empty_like, # works out of the box (tested) + # full_like, # works out of the box (tested) + # ones_like, # works out of the box (tested) + # zeros_like, # works out of the box (tested) + # copy, # works out of the box (tested) + "meshgrid": (ca(np.arange(3)), ca(np.arange(3))), + # transpose, # works out of the box (tested) + # reshape, # works out of the box (tested) + # resize, # works out of the box (tested) + # roll, # works out of the box (tested) + # rollaxis, # works out of the box (tested) + # rot90, # works out of the box (tested) + # expand_dims, # works out of the box (tested) + # squeeze, # works out of the box (tested) + # flip, # works out of the box (tested) + # fliplr, # works out of the box (tested) + # flipud, # works out of the box (tested) + # delete, # works out of the box (tested) + # partition, # works out of the box (tested) + # broadcast_to, # works out of the box (tested) + # broadcast_arrays, # works out of the box (tested) + # split, # works out of the box (tested) + # array_split, # works out of the box (tested) + # dsplit, # works out of the box (tested) + # hsplit, # works out of the box (tested) + # vsplit, # works out of the box (tested) + # swapaxes, # works out of the box (tested) + # moveaxis, # works out of the box (tested) + # nansum, # works out of the box (tested) + # std, # works out of the box (tested) + # nanstd, # works out of the box (tested) + # nanvar, # works out of the box (tested) + # nanprod, # works out of the box (tested) + # diag, # works out of the box (tested) + # diag_indices_from, # returns pure numbers + # diagflat, # works out of the box (tested) + # diagonal, # works out of the box (tested) + # ravel, # returns pure numbers + # ravel_multi_index, # returns pure numbers + # unravel_index, # returns pure numbers + # fix, # works out of the box (tested) + # round, # is implemented via np.around + # may_share_memory, # returns pure numbers (booleans) + # linalg.matrix_power, # works out of the box (tested) + # linalg.cholesky, # works out of the box (tested) + # linalg.multi_dot, # works out of the box (tested) + # linalg.matrix_rank, # returns pure numbers + # linalg.qr, # works out of the box (tested) + # linalg.slogdet, # undefined units + # linalg.cond, # works out of the box (tested) + # gradient, # works out of the box (tested) + # cumsum, # works out of the box (tested) + # nancumsum, # works out of the box (tested) + # nancumprod, # we get it for free with np.cumprod (tested) + # bincount, # works out of the box (tested) + # unique, # works out of the box (tested) + # min_scalar_type, # returns dtypes + # extract, # works out of the box (tested) + # setxor1d, # we get it for free with previously implemented functions (tested) + # lexsort, # returns pure numbers + # digitize, # returns pure numbers + # tril_indices_from, # returns pure numbers + # triu_indices_from, # returns pure numbers + # imag, # works out of the box (tested) + # real, # works out of the box (tested) + # real_if_close, # works out of the box (tested) + # einsum_path, # returns pure numbers + # cov, # returns pure numbers + # corrcoef, # returns pure numbers + # compress, # works out of the box (tested) + # take_along_axis, # works out of the box (tested) + # he following all work out of the box (tested): + # linalg.cross, + # linalg.diagonal, + # linalg.matmul, + # linalg.matrix_norm, + # linalg.matrix_transpose, + # linalg.svdvals, + # linalg.tensordot, + # linalg.trace, + # linalg.vecdot, + # linalg.vector_norm, + # astype, + # matrix_transpose, + # unique_all, + # unique_counts, + # unique_inverse, + # unique_values, + # vecdot, } functions_checked = list() bad_funcs = dict() for fname, args in functions_to_check.items(): - ua_args = tuple(to_ua(arg) for arg in args) + ua_args = list() + for arg in args: + ua_args.append(arg_to_ua(arg)) func = getfunc(fname) try: with warnings.catch_warnings(): From 7de44a7e37510d93037382e930575f90c7da0eca Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 30 Jan 2025 15:04:02 +0000 Subject: [PATCH 037/125] Simplify implementation of amin, amax, and implement meshgrid. --- swiftsimio/_array_functions.py | 33 +++++++++++++-------------------- 1 file changed, 13 insertions(+), 20 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index ac941c52..d7ee0c5f 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -12,6 +12,8 @@ _comparison_cosmo_factor, _power_cosmo_factor, _return_without_cosmo_factor, + _propagate_cosmo_array_attributes, + _copy_cosmo_array_attributes, ) _HANDLED_FUNCTIONS = {} @@ -1860,25 +1862,16 @@ def take(a, indices, axis=None, out=None, mode="raise"): return _return_helper(res, helper_result, ret_cf, out=out) -@implements(np.amax) -def amax( - a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, where=np._NoValue -): - helper_result = _prepare_array_func_args( - a, axis=axis, out=out, keepdims=keepdims, initial=initial, where=where - ) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) - res = np.amax._implementation(*helper_result["args"], **helper_result["kwargs"]) - return _return_helper(res, helper_result, ret_cf, out=out) +amax = implements(np.amax)(_propagate_cosmo_array_attributes(np.amax._implementation)) +amin = implements(np.amin)(_propagate_cosmo_array_attributes(np.amin._implementation)) -@implements(np.amin) -def amin( - a, axis=None, out=None, keepdims=np._NoValue, initial=np._NoValue, where=np._NoValue -): - helper_result = _prepare_array_func_args( - a, axis=axis, out=out, keepdims=keepdims, initial=initial, where=where - ) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) - res = np.amin._implementation(*helper_result["args"], **helper_result["kwargs"]) - return _return_helper(res, helper_result, ret_cf, out=out) +@implements(np.meshgrid) +def meshgrid(*xi, **kwargs): + # meshgrid is a unique case: arguments never interact with each other, so we don't + # want to use our _prepare_array_func_args helper (that will try to coerce to + # compatible comoving, cosmo_factor). + # However we can't just use _propagate_cosmo_array_attributes because we need to + # iterate over arguments. + res = np.meshgrid._implementation(*xi, **kwargs) + return tuple(_copy_cosmo_array_attributes(x, r) for (x, r) in zip(xi, res)) From 63f4557a7f0f63328672049dddf4f51f38de1add Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 30 Jan 2025 15:18:32 +0000 Subject: [PATCH 038/125] Split out cosmo attribute copying into separate function & fix some bugs. --- swiftsimio/objects.py | 54 ++++++++++++++++++++++++------------------- 1 file changed, 30 insertions(+), 24 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 4f1083b9..d1c75204 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -113,18 +113,29 @@ def __init__(self, message="Could not convert to comoving coordinates"): self.message = message +def _copy_cosmo_array_attributes(from_ca, to_ca): + if not isinstance(to_ca, cosmo_array): + return to_ca + if hasattr(from_ca, "cosmo_factor"): + to_ca.cosmo_factor = from_ca.cosmo_factor + if hasattr(from_ca, "comoving"): + to_ca.comoving = from_ca.comoving + if hasattr(from_ca, "valid_transform"): + to_ca.valid_transform = from_ca.valid_transform + return to_ca + + def _propagate_cosmo_array_attributes(func): - def wrapped(self, *args, **kwargs): - ret = func(self, *args, **kwargs) + # can work on methods (obj is self) and functions (obj is first argument) + def wrapped(obj, *args, **kwargs): + ret = func(obj, *args, **kwargs) if not isinstance(ret, cosmo_array): return ret - if hasattr(self, "cosmo_factor"): - ret.cosmo_factor = self.cosmo_factor - if hasattr(self, "comoving"): - ret.comoving = self.comoving - if hasattr(self, "valid_transform"): - ret.valid_transform = self.valid_transform - return ret + ret = _copy_cosmo_array_attributes(obj, ret) + if ret.shape == (): + return cosmo_quantity(ret) + else: + return ret return wrapped @@ -1028,9 +1039,9 @@ def take(self, indices, **kwargs): return cosmo_array(taken) @_propagate_cosmo_array_attributes - def reshape(self, shape, **kwargs): - reshaped = unyt_array.reshape(self, shape, **kwargs) - if shape == (): + def reshape(self, shape, /, *, order="C"): + reshaped = unyt_array.reshape(self, shape, order=order) + if shape == () or shape is None: return cosmo_quantity(reshaped) else: return reshaped @@ -1426,7 +1437,7 @@ def __new__( ) valid_transform = ( getattr(input_scalar, "valid_transform", None) - if valid_transform is not None + if valid_transform is None else valid_transform ) compression = ( @@ -1448,18 +1459,13 @@ def __new__( compression=compression, ) if ret.size > 1: - raise RuntimeError("unyt_quantity instances must be scalars") + raise RuntimeError("cosmo_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] + @_propagate_cosmo_array_attributes + def reshape(self, shape, /, *, order="C"): + reshaped = unyt_array.reshape(self, shape, order=order) if shape == () or shape is None: - return super().reshape(shape, order=order) + return reshaped else: - return cosmo_array(self).reshape(shape, order=order) + return cosmo_array(reshaped) From 38c49d8a8ddbf49c75aa79cc8f3cb42f9b0956cf Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 30 Jan 2025 15:18:53 +0000 Subject: [PATCH 039/125] Check against unyt warning when running unyt version of function. --- tests/test_cosmo_array.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 413c0e27..673df1c9 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -386,7 +386,7 @@ def test_explicitly_handled_funcs(self): warnings.filterwarnings( action="ignore", category=UserWarning, - message="numpy.savetxt does not preserve units or cosmo", + message="numpy.savetxt does not preserve units", ) ua_result = func(*ua_args) except u.exceptions.UnytError: From 4452fa73ef5a91929694500a29136f4f51aa228e Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 30 Jan 2025 15:34:48 +0000 Subject: [PATCH 040/125] Add a few more functions and a check that we're covering all functions mentioned by unyt. --- tests/test_cosmo_array.py | 26 ++++++++++++++++++-------- 1 file changed, 18 insertions(+), 8 deletions(-) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 673df1c9..aeebffbf 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -123,6 +123,7 @@ def test_explicitly_handled_funcs(self): 'just worked' for unyt but that we need to handle by hand. """ from unyt._array_functions import _HANDLED_FUNCTIONS + from unyt.tests.test_array_functions import NOOP_FUNCTIONS functions_to_check = { # FUNCTIONS UNYT HANDLES EXPLICITLY: @@ -238,11 +239,14 @@ def test_explicitly_handled_funcs(self): "amax": (ca(np.arange(3)),), "amin": (ca(np.arange(3)),), # angle, # expects complex numbers - # any, # works out of the box (tested) - # append, # we get it for free with np.concatenate (tested) - # apply_along_axis, # works out of the box (tested) - # argmax, # returns pure numbers - # argmin, # returns pure numbers + "any": (ca(np.arange(3)),), + "append": ( + ca(np.arange(3)), + ca(1), + ), + "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))), + "argmax": (ca(np.arange(3)),), + "argmin": (ca(np.arange(3)),), # argpartition, # returns pure numbers # argsort, # returns pure numbers # argwhere, # returns pure numbers @@ -371,7 +375,7 @@ def test_explicitly_handled_funcs(self): # unique_counts, # unique_inverse, # unique_values, - # vecdot, + # vecdot, # actually a ufunc! } functions_checked = list() bad_funcs = dict() @@ -443,7 +447,9 @@ def test_explicitly_handled_funcs(self): "(obtained, obtained with unyt input): " + str(bad_funcs) ) unchecked_functions = [ - f for f in _HANDLED_FUNCTIONS if f not in functions_checked + f + for f in set(_HANDLED_FUNCTIONS) | NOOP_FUNCTIONS + if f not in functions_checked ] try: assert len(unchecked_functions) == 0 @@ -451,7 +457,11 @@ def test_explicitly_handled_funcs(self): raise AssertionError( "Did not check functions", [ - ".".join((f.__module__, f.__name__)).replace("numpy", "np") + ( + ".".join((f.__module__, f.__name__)).replace("numpy", "np") + if type(f) is not np.ufunc + else f"{f.__name__} is a ufunc!" + ) for f in unchecked_functions ], ) From 453ef89a9116e1274d3a270529649f4d735c0ad1 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 30 Jan 2025 17:56:43 +0000 Subject: [PATCH 041/125] Check many more array funcs. --- swiftsimio/_array_functions.py | 18 ++- tests/test_cosmo_array.py | 200 ++++++++++++++++----------------- 2 files changed, 113 insertions(+), 105 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index d7ee0c5f..6a5ce268 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1862,8 +1862,20 @@ def take(a, indices, axis=None, out=None, mode="raise"): return _return_helper(res, helper_result, ret_cf, out=out) -amax = implements(np.amax)(_propagate_cosmo_array_attributes(np.amax._implementation)) -amin = implements(np.amin)(_propagate_cosmo_array_attributes(np.amin._implementation)) +implements(np.amax)(_propagate_cosmo_array_attributes(np.amax._implementation)) +implements(np.amin)(_propagate_cosmo_array_attributes(np.amin._implementation)) +implements(np.average)(_propagate_cosmo_array_attributes(np.average._implementation)) +implements(np.nanmax)(_propagate_cosmo_array_attributes(np.nanmax._implementation)) +implements(np.nanmean)(_propagate_cosmo_array_attributes(np.nanmean._implementation)) +implements(np.nanmedian)(_propagate_cosmo_array_attributes(np.nanmedian._implementation)) +implements(np.nanmin)(_propagate_cosmo_array_attributes(np.nanmin._implementation)) +implements(np.max)(_propagate_cosmo_array_attributes(np.max._implementation)) +implements(np.min)(_propagate_cosmo_array_attributes(np.min._implementation)) +implements(np.mean)(_propagate_cosmo_array_attributes(np.mean._implementation)) +implements(np.median)(_propagate_cosmo_array_attributes(np.median._implementation)) +implements(np.sort)(_propagate_cosmo_array_attributes(np.sort._implementation)) +implements(np.sum)(_propagate_cosmo_array_attributes(np.sum._implementation)) +implements(np.partition)(_propagate_cosmo_array_attributes(np.partition._implementation)) @implements(np.meshgrid) @@ -1875,3 +1887,5 @@ def meshgrid(*xi, **kwargs): # iterate over arguments. res = np.meshgrid._implementation(*xi, **kwargs) return tuple(_copy_cosmo_array_attributes(x, r) for (x, r) in zip(xi, res)) + + diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index aeebffbf..829cddb5 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 arg_to_ua(arg): @@ -38,6 +38,9 @@ def check_result(x_c, x_u): if isinstance(x_u, str): assert isinstance(x_c, str) return + if isinstance(x_u, type) or isinstance(x_u, np.dtype): + assert x_u == x_c + return # careful, unyt_quantity is a subclass of unyt_array if isinstance(x_u, u.unyt_quantity): assert isinstance(x_c, cosmo_quantity) @@ -61,7 +64,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") @@ -72,7 +75,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") @@ -82,7 +85,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") @@ -92,7 +95,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") @@ -103,20 +106,19 @@ 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 class TestNumpyFunctions: - def test_explicitly_handled_funcs(self): """ Make sure we at least handle everything that unyt does, and anything that @@ -240,75 +242,72 @@ def test_explicitly_handled_funcs(self): "amin": (ca(np.arange(3)),), # angle, # expects complex numbers "any": (ca(np.arange(3)),), - "append": ( - ca(np.arange(3)), - ca(1), - ), + "append": (ca(np.arange(3)), ca(1)), "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))), "argmax": (ca(np.arange(3)),), "argmin": (ca(np.arange(3)),), # argpartition, # returns pure numbers - # argsort, # returns pure numbers - # argwhere, # returns pure numbers - # array_str, # hooks into __str__ - # atleast_1d, # works out of the box (tested) - # atleast_2d, # works out of the box (tested) - # atleast_3d, # works out of the box (tested) - # average, # works out of the box (tested) - # can_cast, # works out of the box (tested) - # common_type, # works out of the box (tested) - # result_type, # works out of the box (tested) - # iscomplex, # works out of the box (tested) - # iscomplexobj, # works out of the box (tested) - # isreal, # works out of the box (tested) - # isrealobj, # works out of the box (tested) + "argsort": (ca(np.arange(3)),), + "argwhere": (ca(np.arange(3)),), + "array_str": (ca(np.arange(3)),), + "atleast_1d": (ca(np.arange(3)),), + "atleast_2d": (ca(np.arange(3)),), + "atleast_3d": (ca(np.arange(3)),), + "average": (ca(np.arange(3)),), + "can_cast": (ca(np.arange(3)), np.float64), + "common_type": (ca(np.arange(3)), ca(np.arange(3))), + "result_type": (ca(np.ones(3)), ca(np.ones(3))), + "iscomplex": (ca(np.arange(3)),), + "iscomplexobj": (ca(np.arange(3)),), + "isreal": (ca(np.arange(3)),), + "isrealobj": (ca(np.arange(3)),), # nan_to_num, # works out of the box (tested) - # nanargmax, # return pure numbers - # nanargmin, # return pure numbers - # nanmax, # works out of the box (tested) - # nanmean, # works out of the box (tested) - # nanmedian, # works out of the box (tested) - # nanmin, # works out of the box (tested) - # trim_zeros, # works out of the box (tested) - # max, # works out of the box (tested) - # mean, # works out of the box (tested) - # median, # works out of the box (tested) - # min, # works out of the box (tested) - # ndim, # return pure numbers - # shape, # returns pure numbers - # size, # returns pure numbers - # sort, # works out of the box (tested) - # sum, # works out of the box (tested) - # repeat, # works out of the box (tested) - # tile, # works out of the box (tested) - # shares_memory, # works out of the box (tested) - # nonzero, # works out of the box (tested) - # count_nonzero, # returns pure numbers - # flatnonzero, # works out of the box (tested) - # isneginf, # works out of the box (tested) - # isposinf, # works out of the box (tested) - # empty_like, # works out of the box (tested) - # full_like, # works out of the box (tested) - # ones_like, # works out of the box (tested) - # zeros_like, # works out of the box (tested) - # copy, # works out of the box (tested) + "nanargmax": (ca(np.arange(3)),), + "nanargmin": (ca(np.arange(3)),), + "nanmax": (ca(np.arange(3)),), + "nanmean": (ca(np.arange(3)),), + "nanmedian": (ca(np.arange(3)),), + "nanmin": (ca(np.arange(3)),), + "trim_zeros": (ca(np.arange(3)),), + "max": (ca(np.arange(3)),), + "mean": (ca(np.arange(3)),), + "median": (ca(np.arange(3)),), + "min": (ca(np.arange(3)),), + "ndim": (ca(np.arange(3)),), + "shape": (ca(np.arange(3)),), + "size": (ca(np.arange(3)),), + "sort": (ca(np.arange(3)),), + "sum": (ca(np.arange(3)),), + "repeat": (ca(np.arange(3)), 2), + # "tile": (ca(np.arange(3)), 2), # reshape broken? + "shares_memory": (ca(np.arange(3)), ca(np.arange(3))), + "nonzero": (ca(np.arange(3)),), + "count_nonzero": (ca(np.arange(3)),), + "flatnonzero": (ca(np.arange(3)),), + "isneginf": (ca(np.arange(3)),), + "isposinf": (ca(np.arange(3)),), + # "empty_like": (ca(np.arange(3)),), # unyt result wrong? + "full_like": (ca(np.arange(3)), ca(1)), + "ones_like": (ca(np.arange(3)),), + "zeros_like": (ca(np.arange(3)),), + "copy": (ca(np.arange(3)),), "meshgrid": (ca(np.arange(3)), ca(np.arange(3))), - # transpose, # works out of the box (tested) - # reshape, # works out of the box (tested) - # resize, # works out of the box (tested) - # roll, # works out of the box (tested) - # rollaxis, # works out of the box (tested) - # rot90, # works out of the box (tested) - # expand_dims, # works out of the box (tested) - # squeeze, # works out of the box (tested) - # flip, # works out of the box (tested) - # fliplr, # works out of the box (tested) - # flipud, # works out of the box (tested) - # delete, # works out of the box (tested) - # partition, # works out of the box (tested) - # broadcast_to, # works out of the box (tested) - # broadcast_arrays, # works out of the box (tested) - # split, # works out of the box (tested) + "transpose": (ca(np.eye(3)),), + "reshape": (ca(np.arange(3)), (3,)), + "resize": (ca(np.arange(3)), 6), + "roll": (ca(np.arange(3)), 1), + "rollaxis": (ca(np.arange(3)), 0), + "rot90": (ca(np.eye(3)),), + "expand_dims": (ca(np.arange(3)), 0), + "squeeze": (ca(np.arange(3)),), + "flip": (ca(np.eye(3)),), + "fliplr": (ca(np.eye(3)),), + "flipud": (ca(np.eye(3)),), + "delete": (ca(np.arange(3)), 0), + "partition": (ca(np.arange(3)), 1), + "broadcast_to": (ca(np.arange(3)), 3), + "broadcast_arrays": (ca(np.arange(3)),), + "split": (ca(np.arange(3)), 1), # array_split, # works out of the box (tested) # dsplit, # works out of the box (tested) # hsplit, # works out of the box (tested) @@ -375,7 +374,6 @@ def test_explicitly_handled_funcs(self): # unique_counts, # unique_inverse, # unique_values, - # vecdot, # actually a ufunc! } functions_checked = list() bad_funcs = dict() @@ -457,11 +455,7 @@ def test_explicitly_handled_funcs(self): raise AssertionError( "Did not check functions", [ - ( - ".".join((f.__module__, f.__name__)).replace("numpy", "np") - if type(f) is not np.ufunc - else f"{f.__name__} is a ufunc!" - ) + (".".join((f.__module__, f.__name__)).replace("numpy", "np")) for f in unchecked_functions ], ) @@ -477,7 +471,7 @@ def test_explicitly_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), ), ), ), @@ -488,13 +482,13 @@ def test_explicitly_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), ), ), ), @@ -506,19 +500,19 @@ def test_explicitly_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), ), ], ), @@ -530,7 +524,7 @@ def test_explicitly_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]), ), @@ -547,19 +541,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] @@ -606,9 +600,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): @@ -616,9 +610,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): @@ -626,9 +620,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 = { @@ -640,9 +634,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), ] ), ): @@ -682,12 +676,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"]) @@ -696,10 +690,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 f21f1b7732a7ee07ee610e43269eff88460971fd Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 10:28:01 +0000 Subject: [PATCH 042/125] Keep cosmo_array to cosmo_quantity conversion in one location. --- swiftsimio/_array_functions.py | 15 ++++---- swiftsimio/objects.py | 40 ++++++++++++-------- tests/test_cosmo_array.py | 68 +++++++++++++++++----------------- 3 files changed, 66 insertions(+), 57 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 6a5ce268..ce6fd5bb 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -32,10 +32,7 @@ def decorator(func): def _return_helper(res, helper_result, ret_cf, out=None): if out is None: - if isinstance(res, cosmo_array) and res.shape == (): - # happens when handling a function that unyt didn't handle explicitly - return cosmo_quantity(res) - elif isinstance(res, unyt_quantity) and not isinstance(res, cosmo_quantity): + if isinstance(res, unyt_quantity) and not isinstance(res, cosmo_quantity): return cosmo_quantity( res, comoving=helper_result["comoving"], @@ -1867,7 +1864,9 @@ def take(a, indices, axis=None, out=None, mode="raise"): implements(np.average)(_propagate_cosmo_array_attributes(np.average._implementation)) implements(np.nanmax)(_propagate_cosmo_array_attributes(np.nanmax._implementation)) implements(np.nanmean)(_propagate_cosmo_array_attributes(np.nanmean._implementation)) -implements(np.nanmedian)(_propagate_cosmo_array_attributes(np.nanmedian._implementation)) +implements(np.nanmedian)( + _propagate_cosmo_array_attributes(np.nanmedian._implementation) +) implements(np.nanmin)(_propagate_cosmo_array_attributes(np.nanmin._implementation)) implements(np.max)(_propagate_cosmo_array_attributes(np.max._implementation)) implements(np.min)(_propagate_cosmo_array_attributes(np.min._implementation)) @@ -1875,7 +1874,9 @@ def take(a, indices, axis=None, out=None, mode="raise"): implements(np.median)(_propagate_cosmo_array_attributes(np.median._implementation)) implements(np.sort)(_propagate_cosmo_array_attributes(np.sort._implementation)) implements(np.sum)(_propagate_cosmo_array_attributes(np.sum._implementation)) -implements(np.partition)(_propagate_cosmo_array_attributes(np.partition._implementation)) +implements(np.partition)( + _propagate_cosmo_array_attributes(np.partition._implementation) +) @implements(np.meshgrid) @@ -1887,5 +1888,3 @@ def meshgrid(*xi, **kwargs): # iterate over arguments. res = np.meshgrid._implementation(*xi, **kwargs) return tuple(_copy_cosmo_array_attributes(x, r) for (x, r) in zip(xi, res)) - - diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index d1c75204..0b4111a5 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -132,10 +132,11 @@ def wrapped(obj, *args, **kwargs): if not isinstance(ret, cosmo_array): return ret ret = _copy_cosmo_array_attributes(obj, ret) - if ret.shape == (): - return cosmo_quantity(ret) - else: - return ret + # if ret.shape == (): + # return cosmo_quantity(ret) + # else: + # return ret + return ret return wrapped @@ -1319,19 +1320,28 @@ def __array_function__(self, func, types, args, kwargs): # leaving other arguments a chance to take the lead return NotImplemented - if func in _HANDLED_FUNCTIONS: - return _HANDLED_FUNCTIONS[func](*args, **kwargs) - elif 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 cosmo_array objects if not all(issubclass(t, cosmo_array) or t is np.ndarray for t in types): + # Note: this allows subclasses that don't override + # __array_function__ to handle cosmo_array objects return NotImplemented - return _HANDLED_FUNCTIONS[func](*args, **kwargs) + + if func in _HANDLED_FUNCTIONS: + ret = _HANDLED_FUNCTIONS[func](*args, **kwargs) + elif func in _UNYT_HANDLED_FUNCTIONS: + ret = _UNYT_HANDLED_FUNCTIONS[func](*args, **kwargs) + else: + # default to numpy's private implementation + ret = func._implementation(*args, **kwargs) + if ( + isinstance(ret, cosmo_array) + and ret.shape == () + and not isinstance(ret, cosmo_quantity) + ): + return cosmo_quantity(ret) + elif isinstance(ret, cosmo_quantity) and ret.shape != (): + return cosmo_array(ret) + else: + return ret class cosmo_quantity(cosmo_array, unyt_quantity): diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 829cddb5..7cbe2d9c 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 arg_to_ua(arg): @@ -64,7 +64,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") @@ -75,7 +75,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") @@ -85,7 +85,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") @@ -95,7 +95,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") @@ -106,14 +106,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 @@ -240,7 +240,7 @@ def test_explicitly_handled_funcs(self): "all": (ca(np.arange(3)),), "amax": (ca(np.arange(3)),), "amin": (ca(np.arange(3)),), - # angle, # expects complex numbers + "angle": (ca(complex(1, 1)),), "any": (ca(np.arange(3)),), "append": (ca(np.arange(3)), ca(1)), "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))), @@ -471,7 +471,7 @@ def test_explicitly_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), ), ), ), @@ -482,13 +482,13 @@ def test_explicitly_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), ), ), ), @@ -500,19 +500,19 @@ def test_explicitly_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), ), ], ), @@ -524,7 +524,7 @@ def test_explicitly_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]), ), @@ -541,19 +541,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] @@ -600,9 +600,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): @@ -610,9 +610,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): @@ -620,9 +620,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 = { @@ -634,9 +634,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), ] ), ): @@ -676,12 +676,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"]) @@ -690,10 +690,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 5dfcb02b7bc01edf3981acf5ecb1ae4053a506a9 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 15:05:52 +0000 Subject: [PATCH 043/125] Simplify wrapping some methods and properties. --- swiftsimio/objects.py | 74 ++++++++++++++++++++----------------------- 1 file changed, 35 insertions(+), 39 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 0b4111a5..b9dcdf95 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -96,6 +96,7 @@ isnat, heaviside, matmul, + vecdot, ) from numpy._core.umath import _ones_like @@ -132,10 +133,27 @@ def wrapped(obj, *args, **kwargs): if not isinstance(ret, cosmo_array): return ret ret = _copy_cosmo_array_attributes(obj, ret) - # if ret.shape == (): - # return cosmo_quantity(ret) - # else: - # return ret + return ret + + return wrapped + + +def _ensure_cosmo_array_or_quantity(func): + # can work on methods (obj is self) and functions (obj is first argument) + def wrapped(obj, *args, **kwargs): + ret = func(obj, *args, **kwargs) + if isinstance(ret, unyt_quantity) and not isinstance(ret, cosmo_quantity): + ret = cosmo_quantity(ret) + elif isinstance(ret, unyt_array) and not isinstance(ret, cosmo_array): + ret = cosmo_array(ret) + if ( + isinstance(ret, cosmo_array) + and not isinstance(ret, cosmo_quantity) + and ret.shape == () + ): + ret = cosmo_quantity(ret) + elif isinstance(ret, cosmo_quantity) and ret.shape != (): + ret = cosmo_array(ret) return ret return wrapped @@ -852,6 +870,7 @@ class cosmo_array(unyt_array): _ones_like: _preserve_cosmo_factor, matmul: _multiply_cosmo_factor, clip: _passthrough_cosmo_factor, + vecdot: _multiply_cosmo_factor, } def __new__( @@ -1031,43 +1050,20 @@ def __setstate__(self, state): 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, /, *, order="C"): - reshaped = unyt_array.reshape(self, shape, order=order) - if shape == () or shape is None: - 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 + take = _propagate_cosmo_array_attributes( + _ensure_cosmo_array_or_quantity(unyt_array.take) + ) + reshape = _propagate_cosmo_array_attributes( + _ensure_cosmo_array_or_quantity(unyt_array.reshape) + ) + __getitem__ = _propagate_cosmo_array_attributes( + _ensure_cosmo_array_or_quantity(unyt_array.__getitem__) + ) # Also wrap some array "properties": - - @property - def T(self): - return self.transpose() # transpose is wrapped above. - - @property - def ua(self): - return _propagate_cosmo_array_attributes(np.ones_like)(self) - - @property - def unit_array(self): - return _propagate_cosmo_array_attributes(np.ones_like)(self) + T = property(_propagate_cosmo_array_attributes(unyt_array.transpose)) + ua = property(_propagate_cosmo_array_attributes(np.ones_like)) + unit_array = property(_propagate_cosmo_array_attributes(np.ones_like)) def convert_to_comoving(self) -> None: """ From 9e46bf7593741dccb43564bc4910cce656976dbe Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 15:06:28 +0000 Subject: [PATCH 044/125] Simplify array function checks and finish checking all unyt-supported numpy functions. --- tests/test_cosmo_array.py | 206 +++++++++++++++++++------------------- 1 file changed, 104 insertions(+), 102 deletions(-) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 7cbe2d9c..bd66d3e1 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -31,27 +31,38 @@ def arg_to_ua(arg): def to_ua(x): - return u.unyt_array(x.to_value(x.units), x.units) if hasattr(x, "comoving") else x + return u.unyt_array(x) if hasattr(x, "comoving") else x -def check_result(x_c, x_u): +def check_result(x_c, x_u, ignore_values=False): + if x_u is None: + assert x_c is None + return if isinstance(x_u, str): assert isinstance(x_c, str) return if isinstance(x_u, type) or isinstance(x_u, np.dtype): assert x_u == x_c return - # careful, unyt_quantity is a subclass of unyt_array + if type(x_u) in (list, tuple): + assert type(x_u) is type(x_c) + assert len(x_u) == len(x_c) + for x_c_i, x_u_i in zip(x_c, x_u): + check_result(x_c_i, x_u_i) + 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 np.allclose(x_c, x_u) + if not ignore_values: + 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)) + if not ignore_values: + assert np.allclose(x_c.to_value(x_c.units), x_u.to_value(x_u.units)) if isinstance(x_c, cosmo_array): # includes cosmo_quantity assert x_c.comoving is False if x_c.units != u.dimensionless: @@ -238,16 +249,16 @@ def test_explicitly_handled_funcs(self): "take": (ca(np.arange(3)), np.arange(3)), # FUNCTIONS THAT UNYT DOESN'T HANDLE EXPLICITLY (THEY "JUST WORK"): "all": (ca(np.arange(3)),), - "amax": (ca(np.arange(3)),), - "amin": (ca(np.arange(3)),), + "amax": (ca(np.arange(3)),), # implemented via max + "amin": (ca(np.arange(3)),), # implemented via min "angle": (ca(complex(1, 1)),), "any": (ca(np.arange(3)),), "append": (ca(np.arange(3)), ca(1)), "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))), - "argmax": (ca(np.arange(3)),), - "argmin": (ca(np.arange(3)),), - # argpartition, # returns pure numbers - "argsort": (ca(np.arange(3)),), + "argmax": (ca(np.arange(3)),), # implemented via max + "argmin": (ca(np.arange(3)),), # implemented via min + "argpartition": (ca(np.arange(3)), 1), # implemented via partition + "argsort": (ca(np.arange(3)),), # implemented via sort "argwhere": (ca(np.arange(3)),), "array_str": (ca(np.arange(3)),), "atleast_1d": (ca(np.arange(3)),), @@ -261,13 +272,13 @@ def test_explicitly_handled_funcs(self): "iscomplexobj": (ca(np.arange(3)),), "isreal": (ca(np.arange(3)),), "isrealobj": (ca(np.arange(3)),), - # nan_to_num, # works out of the box (tested) - "nanargmax": (ca(np.arange(3)),), - "nanargmin": (ca(np.arange(3)),), - "nanmax": (ca(np.arange(3)),), - "nanmean": (ca(np.arange(3)),), - "nanmedian": (ca(np.arange(3)),), - "nanmin": (ca(np.arange(3)),), + "nan_to_num": (ca(np.arange(3)),), + "nanargmax": (ca(np.arange(3)),), # implemented via max + "nanargmin": (ca(np.arange(3)),), # implemented via min + "nanmax": (ca(np.arange(3)),), # implemented via max + "nanmean": (ca(np.arange(3)),), # implemented via mean + "nanmedian": (ca(np.arange(3)),), # implemented via median + "nanmin": (ca(np.arange(3)),), # implemented via min "trim_zeros": (ca(np.arange(3)),), "max": (ca(np.arange(3)),), "mean": (ca(np.arange(3)),), @@ -279,14 +290,14 @@ def test_explicitly_handled_funcs(self): "sort": (ca(np.arange(3)),), "sum": (ca(np.arange(3)),), "repeat": (ca(np.arange(3)), 2), - # "tile": (ca(np.arange(3)), 2), # reshape broken? + "tile": (ca(np.arange(3)), 2), "shares_memory": (ca(np.arange(3)), ca(np.arange(3))), "nonzero": (ca(np.arange(3)),), "count_nonzero": (ca(np.arange(3)),), "flatnonzero": (ca(np.arange(3)),), "isneginf": (ca(np.arange(3)),), "isposinf": (ca(np.arange(3)),), - # "empty_like": (ca(np.arange(3)),), # unyt result wrong? + "empty_like": (ca(np.arange(3)),), "full_like": (ca(np.arange(3)), ca(1)), "ones_like": (ca(np.arange(3)),), "zeros_like": (ca(np.arange(3)),), @@ -308,72 +319,74 @@ def test_explicitly_handled_funcs(self): "broadcast_to": (ca(np.arange(3)), 3), "broadcast_arrays": (ca(np.arange(3)),), "split": (ca(np.arange(3)), 1), - # array_split, # works out of the box (tested) - # dsplit, # works out of the box (tested) - # hsplit, # works out of the box (tested) - # vsplit, # works out of the box (tested) - # swapaxes, # works out of the box (tested) - # moveaxis, # works out of the box (tested) - # nansum, # works out of the box (tested) - # std, # works out of the box (tested) - # nanstd, # works out of the box (tested) - # nanvar, # works out of the box (tested) - # nanprod, # works out of the box (tested) - # diag, # works out of the box (tested) - # diag_indices_from, # returns pure numbers - # diagflat, # works out of the box (tested) - # diagonal, # works out of the box (tested) - # ravel, # returns pure numbers - # ravel_multi_index, # returns pure numbers - # unravel_index, # returns pure numbers - # fix, # works out of the box (tested) - # round, # is implemented via np.around - # may_share_memory, # returns pure numbers (booleans) - # linalg.matrix_power, # works out of the box (tested) - # linalg.cholesky, # works out of the box (tested) - # linalg.multi_dot, # works out of the box (tested) - # linalg.matrix_rank, # returns pure numbers - # linalg.qr, # works out of the box (tested) - # linalg.slogdet, # undefined units - # linalg.cond, # works out of the box (tested) - # gradient, # works out of the box (tested) - # cumsum, # works out of the box (tested) - # nancumsum, # works out of the box (tested) - # nancumprod, # we get it for free with np.cumprod (tested) - # bincount, # works out of the box (tested) - # unique, # works out of the box (tested) - # min_scalar_type, # returns dtypes - # extract, # works out of the box (tested) - # setxor1d, # we get it for free with previously implemented functions (tested) - # lexsort, # returns pure numbers - # digitize, # returns pure numbers - # tril_indices_from, # returns pure numbers - # triu_indices_from, # returns pure numbers - # imag, # works out of the box (tested) - # real, # works out of the box (tested) - # real_if_close, # works out of the box (tested) - # einsum_path, # returns pure numbers - # cov, # returns pure numbers - # corrcoef, # returns pure numbers - # compress, # works out of the box (tested) - # take_along_axis, # works out of the box (tested) - # he following all work out of the box (tested): - # linalg.cross, - # linalg.diagonal, - # linalg.matmul, - # linalg.matrix_norm, - # linalg.matrix_transpose, - # linalg.svdvals, - # linalg.tensordot, - # linalg.trace, - # linalg.vecdot, - # linalg.vector_norm, - # astype, - # matrix_transpose, - # unique_all, - # unique_counts, - # unique_inverse, - # unique_values, + "array_split": (ca(np.arange(3)), 1), + "dsplit": (ca(np.arange(27)).reshape(3, 3, 3), 1), + "hsplit": (ca(np.arange(3)), 1), + "vsplit": (ca(np.eye(3)), 1), + "swapaxes": (ca(np.eye(3)), 0, 1), + "moveaxis": (ca(np.eye(3)), 0, 1), + "nansum": (ca(np.arange(3)),), # implemented via sum + "std": (ca(np.arange(3)),), + "nanstd": (ca(np.arange(3)),), + "nanvar": (ca(np.arange(3)),), + "nanprod": (ca(np.arange(3)),), + "diag": (ca(np.eye(3)),), + "diag_indices_from": (ca(np.eye(3)),), + "diagflat": (ca(np.eye(3)),), + "diagonal": (ca(np.eye(3)),), + "ravel": (ca(np.arange(3)),), + "ravel_multi_index": (np.eye(2, dtype=int), (2, 2)), + "unravel_index": (np.arange(3), (3,)), + "fix": (ca(np.arange(3)),), + "round": (ca(np.arange(3)),), # implemented via around + "may_share_memory": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.matrix_power": (ca(np.eye(3)), 2), + "linalg.cholesky": (ca(np.eye(3)),), + "linalg.multi_dot": ((ca(np.eye(3)), ca(np.eye(3))),), + "linalg.matrix_rank": (ca(np.eye(3)),), + "linalg.qr": (ca(np.eye(3)),), + "linalg.slogdet": (ca(np.eye(3)),), + "linalg.cond": (ca(np.eye(3)),), + "gradient": (ca(np.arange(3)),), + "cumsum": (ca(np.arange(3)),), + "nancumsum": (ca(np.arange(3)),), + "nancumprod": (ca(np.arange(3)),), + "bincount": (ca(np.arange(3)),), + "unique": (ca(np.arange(3)),), + "min_scalar_type": (ca(np.arange(3)),), + "extract": (0, ca(np.arange(3))), + "setxor1d": (ca(np.arange(3)), ca(np.arange(3))), + "lexsort": (ca(np.arange(3)),), + "digitize": (ca(np.arange(3)), ca(np.arange(3))), + "tril_indices_from": (ca(np.eye(3)),), + "triu_indices_from": (ca(np.eye(3)),), + "imag": (ca(np.arange(3)),), + "real": (ca(np.arange(3)),), + "real_if_close": (ca(np.arange(3)),), + "einsum_path": ("ij,jk->ik", ca(np.eye(3)), ca(np.eye(3))), + "cov": (ca(np.arange(3)),), + "corrcoef": (ca(np.arange(3)),), + "compress": (np.zeros(3), ca(np.arange(3))), + "take_along_axis": (ca(np.arange(3)), np.ones(3, dtype=int), 0), + "linalg.cross": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.diagonal": (ca(np.eye(3)),), + "linalg.matmul": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.matrix_norm": (ca(np.eye(3)),), + "linalg.matrix_transpose": (ca(np.eye(3)),), + "linalg.svdvals": (ca(np.eye(3)),), + "linalg.tensordot": (ca(np.eye(3)), ca(np.eye(3))), + "linalg.trace": (ca(np.eye(3)),), + "linalg.vecdot": (ca(np.arange(3)), ca(np.arange(3))), + "linalg.vector_norm": (ca(np.arange(3)),), + "astype": (ca(np.arange(3)), float), + "matrix_transpose": (ca(np.eye(3)),), + "unique_all": (ca(np.arange(3)),), + "unique_counts": (ca(np.arange(3)),), + "unique_inverse": (ca(np.arange(3)),), + "unique_values": (ca(np.arange(3)),), + "cumulative_sum": (ca(np.arange(3)),), + "cumulative_prod": (ca(np.arange(3)),), + "unstack": (ca(np.arange(3)),), } functions_checked = list() bad_funcs = dict() @@ -423,22 +436,11 @@ def test_explicitly_handled_funcs(self): ua_result = ua_args[0] 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 - 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 + ignore_values = fname in {"empty_like"} # empty_like has arbitrary data + try: + check_result(result, ua_result, ignore_values=ignore_values) + except AssertionError: + bad_funcs["np." + fname] = result, ua_result if len(bad_funcs) > 0: raise AssertionError( "Some functions did not return expected types " From 583665034695aba590ed6bd30f5081d6a62acdcc Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 15:06:43 +0000 Subject: [PATCH 045/125] Remove some redundant wrappers. --- swiftsimio/_array_functions.py | 8 -------- 1 file changed, 8 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index ce6fd5bb..a2be59a6 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1859,15 +1859,7 @@ def take(a, indices, axis=None, out=None, mode="raise"): return _return_helper(res, helper_result, ret_cf, out=out) -implements(np.amax)(_propagate_cosmo_array_attributes(np.amax._implementation)) -implements(np.amin)(_propagate_cosmo_array_attributes(np.amin._implementation)) implements(np.average)(_propagate_cosmo_array_attributes(np.average._implementation)) -implements(np.nanmax)(_propagate_cosmo_array_attributes(np.nanmax._implementation)) -implements(np.nanmean)(_propagate_cosmo_array_attributes(np.nanmean._implementation)) -implements(np.nanmedian)( - _propagate_cosmo_array_attributes(np.nanmedian._implementation) -) -implements(np.nanmin)(_propagate_cosmo_array_attributes(np.nanmin._implementation)) implements(np.max)(_propagate_cosmo_array_attributes(np.max._implementation)) implements(np.min)(_propagate_cosmo_array_attributes(np.min._implementation)) implements(np.mean)(_propagate_cosmo_array_attributes(np.mean._implementation)) From e687143e95c02898b01b5db038333ddc4d5c9f88 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 16:47:35 +0000 Subject: [PATCH 046/125] Greatly simplify array function handling with wrapper functions. --- swiftsimio/_array_functions.py | 1182 ++++++++------------------------ 1 file changed, 302 insertions(+), 880 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index a2be59a6..fea96f12 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -15,12 +15,115 @@ _propagate_cosmo_array_attributes, _copy_cosmo_array_attributes, ) +from unyt._array_functions import ( + dot as unyt_dot, + vdot as unyt_vdot, + inner as unyt_inner, + outer as unyt_outer, + kron as unyt_kron, + histogram_bin_edges as unyt_histogram_bin_edges, + linalg_inv as unyt_linalg_inv, + linalg_tensorinv as unyt_linalg_tensorinv, + linalg_pinv as unyt_linalg_pinv, + linalg_svd as unyt_linalg_svd, + histogram as unyt_histogram, + histogram2d as unyt_histogram2d, + histogramdd as unyt_histogramdd, + concatenate as unyt_concatenate, + intersect1d as unyt_intersect1d, + union1d as unyt_union1d, + norm as unyt_linalg_norm, # not linalg_norm, doesn't follow usual pattern + vstack as unyt_vstack, + hstack as unyt_hstack, + dstack as unyt_dstack, + column_stack as unyt_column_stack, + stack as unyt_stack, + around as unyt_around, + block as unyt_block, + ftt_fft as unyt_fft_fft, # unyt has a copy-pasted typo fft -> ftt + ftt_fft2 as unyt_fft_fft2, + ftt_fftn as unyt_fft_fftn, + ftt_hfft as unyt_fft_hfft, + ftt_rfft as unyt_fft_rfft, + ftt_rfft2 as unyt_fft_rfft2, + ftt_rfftn as unyt_fft_rfftn, + ftt_ifft as unyt_fft_ifft, + ftt_ifft2 as unyt_fft_ifft2, + ftt_ifftn as unyt_fft_ifftn, + ftt_ihfft as unyt_fft_ihfft, + ftt_irfft as unyt_fft_irfft, + ftt_irfft2 as unyt_fft_irfft2, + ftt_irfftn as unyt_fft_irfftn, + fft_fftshift as unyt_fft_fftshift, + fft_ifftshift as unyt_fft_ifftshift, + sort_complex as unyt_sort_complex, + isclose as unyt_isclose, + allclose as unyt_allclose, + array2string as unyt_array2string, + cross as unyt_cross, + array_equal as unyt_array_equal, + array_equiv as unyt_array_equiv, + linspace as unyt_linspace, + logspace as unyt_logspace, + geomspace as unyt_geomspace, + copyto as unyt_copyto, + prod as unyt_prod, + var as unyt_var, + trace as unyt_trace, + percentile as unyt_percentile, + quantile as unyt_quantile, + nanpercentile as unyt_nanpercentile, + nanquantile as unyt_nanquantile, + linalg_det as unyt_linalg_det, + diff as unyt_diff, + ediff1d as unyt_ediff1d, + ptp as unyt_ptp, + pad as unyt_pad, + choose as unyt_choose, + insert as unyt_insert, + linalg_lstsq as unyt_linalg_lstsq, + linalg_solve as unyt_linalg_solve, + linalg_tensorsolve as unyt_linalg_tensorsolve, + linalg_eig as unyt_linalg_eig, + linalg_eigh as unyt_linalg_eigh, + linalg_eigvals as unyt_linalg_eigvals, + linalg_eigvalsh as unyt_linalg_eigvalsh, + savetxt as unyt_savetxt, + fill_diagonal as unyt_fill_diagonal, + isin as unyt_isin, + place as unyt_place, + put as unyt_put, + put_along_axis as unyt_put_along_axis, + putmask as unyt_putmask, + searchsorted as unyt_searchsorted, + select as unyt_select, + setdiff1d as unyt_setdiff1d, + sinc as unyt_sinc, + clip as unyt_clip, + where as unyt_where, + triu as unyt_triu, + tril as unyt_tril, + einsum as unyt_einsum, + convolve as unyt_convolve, + correlate as unyt_correlate, + tensordot as unyt_tensordot, + unwrap as unyt_unwrap, + interp as unyt_interp, + array_repr as unyt_array_repr, + linalg_outer as unyt_linalg_outer, + trapezoid as unyt_trapezoid, + isin as unyt_in1d, + take as unyt_take, +) _HANDLED_FUNCTIONS = {} +# first we define helper functions to handle repetitive operations in wrapping unyt & +# numpy functions (we will actually wrap the functions below): + def implements(numpy_function): - """Register an __array_function__ implementation for unyt_array objects.""" + """Register an __array_function__ implementation for cosmo_array objects.""" # See NEP 18 https://numpy.org/neps/nep-0018-array-function-protocol.html def decorator(func): @@ -74,6 +177,80 @@ def _return_helper(res, helper_result, ret_cf, out=None): ) +def _default_unary_wrapper(unyt_func, cosmo_factor_wrapper): + + # assumes that we have one primary argument that will be handled + # by the cosmo_factor_wrapper + def wrapper(*args, **kwargs): + helper_result = _prepare_array_func_args(*args, **kwargs) + ret_cf = cosmo_factor_wrapper(helper_result["ca_cfs"][0]) + res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) + if "out" in kwargs: + return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) + else: + return _return_helper(res, helper_result, ret_cf) + + return wrapper + + +def _default_binary_wrapper(unyt_func, cosmo_factor_wrapper): + + # assumes we have two primary arguments that will be handled + # by the cosmo_factor_wrapper + def wrapper(*args, **kwargs): + helper_result = _prepare_array_func_args(*args, **kwargs) + ret_cf = cosmo_factor_wrapper( + helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] + ) + res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) + if "out" in kwargs: + return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) + else: + return _return_helper(res, helper_result, ret_cf) + + return wrapper + + +def _default_comparison_wrapper(unyt_func): + + # assumes we have two primary arguments that will be handled with + # _comparison_cosmo_factor with them as the inputs + def wrapper(*args, **kwargs): + helper_result = _prepare_array_func_args(*args, **kwargs) + ret_cf = _comparison_cosmo_factor( + helper_result["ca_cfs"][0], + helper_result["ca_cfs"][1], + inputs=args[:2], + ) + res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) + return _return_helper(res, helper_result, ret_cf) + + return wrapper + + +def _default_oplist_wrapper(unyt_func): + + # assumes first argument is a list of operands + # assumes that we always preserve the cosmo factor of the first + # element in the list of operands + def wrapper(*args, **kwargs): + helper_result = _prepare_array_func_args(*args, **kwargs) + helper_result_oplist = _prepare_array_func_args(*args[0]) + ret_cf = _preserve_cosmo_factor(helper_result_oplist["ca_cfs"][0]) + res = unyt_func( + helper_result_oplist["args"], + *helper_result["args"][1:], + **helper_result["kwargs"], + ) + return _return_helper(res, helper_result_oplist, ret_cf) + + return wrapper + + +# Now we wrap functions that unyt handles explicitly (below that will be those not handled +# explicitly): + + @implements(np.array2string) def array2string( a, @@ -92,7 +269,6 @@ def array2string( *, legacy=None, ): - from unyt._array_functions import array2string as unyt_array2string res = unyt_array2string( a, @@ -119,69 +295,15 @@ def array2string( 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.dot)(_default_binary_wrapper(unyt_dot, _multiply_cosmo_factor)) +implements(np.vdot)(_default_binary_wrapper(unyt_vdot, _multiply_cosmo_factor)) +implements(np.inner)(_default_binary_wrapper(unyt_inner, _multiply_cosmo_factor)) +implements(np.outer)(_default_binary_wrapper(unyt_outer, _multiply_cosmo_factor)) +implements(np.kron)(_default_binary_wrapper(unyt_kron, _multiply_cosmo_factor)) @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: @@ -196,56 +318,22 @@ def histogram_bin_edges(a, bins=10, range=None, weights=None): 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(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( - 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(np.linalg.inv)( + _default_unary_wrapper(unyt_linalg_inv, _reciprocal_cosmo_factor) +) +implements(np.linalg.tensorinv)( + _default_unary_wrapper(unyt_linalg_tensorinv, _reciprocal_cosmo_factor) +) +implements(np.linalg.pinv)( + _default_unary_wrapper(unyt_linalg_pinv, _reciprocal_cosmo_factor) +) +implements(np.linalg.svd)( + _default_unary_wrapper(unyt_linalg_svd, _preserve_cosmo_factor) +) @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( a, bins=bins, range=range, density=density, weights=weights @@ -277,7 +365,6 @@ def histogram(a, bins=10, range=None, density=None, weights=None): @implements(np.histogram2d) def histogram2d(x, y, bins=10, range=None, density=None, weights=None): - from unyt._array_functions import histogram2d as unyt_histogram2d if range is not None: xrange, yrange = range @@ -394,7 +481,6 @@ 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 D = len(sample) if range is not None: @@ -500,157 +586,21 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): ) -@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( - 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(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 = _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 - - 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 - - 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): - # 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 = _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) -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(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 = _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) +implements(np.concatenate)(_default_oplist_wrapper(unyt_concatenate)) +implements(np.cross)(_default_binary_wrapper(unyt_cross, _multiply_cosmo_factor)) +implements(np.intersect1d)( + _default_binary_wrapper(unyt_intersect1d, _preserve_cosmo_factor) +) +implements(np.union1d)(_default_binary_wrapper(unyt_union1d, _preserve_cosmo_factor)) +implements(np.linalg.norm)( + _default_unary_wrapper(unyt_linalg_norm, _preserve_cosmo_factor) +) +implements(np.vstack)(_default_oplist_wrapper(unyt_vstack)) +implements(np.hstack)(_default_oplist_wrapper(unyt_hstack)) +implements(np.dstack)(_default_oplist_wrapper(unyt_dstack)) +implements(np.column_stack)(_default_oplist_wrapper(unyt_column_stack)) +implements(np.stack)(_default_oplist_wrapper(unyt_stack)) +implements(np.around)(_default_unary_wrapper(unyt_around, _preserve_cosmo_factor)) def _recursive_to_comoving(lst): @@ -718,249 +668,58 @@ def _prepare_array_block_args(lst, recursing=False): @implements(np.block) def block(arrays): - from unyt._array_functions import block as unyt_block - + # block is a special case since we need to recurse more than one level + # down the list of arrays. 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 - - -@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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(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(x, axes=None): - from unyt._array_functions import fft_fftshift as unyt_fft_fftshift - - 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(x, axes=None): - from unyt._array_functions import fft_ifftshift as unyt_fft_ifftshift - - 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) -def sort_complex(a): - from unyt._array_functions import sort_complex as unyt_sort_complex - - 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) -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( - 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(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 - ) - 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(a1, a2, equal_nan=False): - from unyt._array_functions import array_equal as unyt_array_equal - - 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(a1, a2): - from unyt._array_functions import array_equiv as unyt_array_equiv +implements(np.fft.fft)(_default_unary_wrapper(unyt_fft_fft, _reciprocal_cosmo_factor)) +implements(np.fft.fft2)(_default_unary_wrapper(unyt_fft_fft2, _reciprocal_cosmo_factor)) +implements(np.fft.fftn)(_default_unary_wrapper(unyt_fft_fftn, _reciprocal_cosmo_factor)) +implements(np.fft.hfft)(_default_unary_wrapper(unyt_fft_hfft, _reciprocal_cosmo_factor)) +implements(np.fft.rfft)(_default_unary_wrapper(unyt_fft_rfft, _reciprocal_cosmo_factor)) +implements(np.fft.rfft2)( + _default_unary_wrapper(unyt_fft_rfft2, _reciprocal_cosmo_factor) +) +implements(np.fft.rfftn)( + _default_unary_wrapper(unyt_fft_rfftn, _reciprocal_cosmo_factor) +) +implements(np.fft.ifft)(_default_unary_wrapper(unyt_fft_ifft, _reciprocal_cosmo_factor)) +implements(np.fft.ifft2)( + _default_unary_wrapper(unyt_fft_ifft2, _reciprocal_cosmo_factor) +) +implements(np.fft.ifftn)( + _default_unary_wrapper(unyt_fft_ifftn, _reciprocal_cosmo_factor) +) +implements(np.fft.ihfft)( + _default_unary_wrapper(unyt_fft_ihfft, _reciprocal_cosmo_factor) +) +implements(np.fft.irfft)( + _default_unary_wrapper(unyt_fft_irfft, _reciprocal_cosmo_factor) +) +implements(np.fft.irfft2)( + _default_unary_wrapper(unyt_fft_irfft2, _reciprocal_cosmo_factor) +) +implements(np.fft.irfftn)( + _default_unary_wrapper(unyt_fft_irfftn, _reciprocal_cosmo_factor) +) +implements(np.fft.fftshift)( + _default_unary_wrapper(unyt_fft_fftshift, _preserve_cosmo_factor) +) +implements(np.fft.ifftshift)( + _default_unary_wrapper(unyt_fft_ifftshift, _preserve_cosmo_factor) +) - 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.sort_complex)( + _default_unary_wrapper(unyt_sort_complex, _preserve_cosmo_factor) +) +implements(np.isclose)(_default_comparison_wrapper(unyt_isclose)) +implements(np.allclose)(_default_comparison_wrapper(unyt_allclose)) +implements(np.array_equal)(_default_comparison_wrapper(unyt_array_equal)) +implements(np.array_equiv)(_default_comparison_wrapper(unyt_array_equiv)) @implements(np.linspace) @@ -975,7 +734,6 @@ def linspace( *, device=None, ): - from unyt._array_functions import linspace as unyt_linspace helper_result = _prepare_array_func_args( start, @@ -999,7 +757,6 @@ def linspace( @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( start, stop, num=num, endpoint=endpoint, base=base, dtype=dtype, axis=axis @@ -1009,23 +766,13 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0): return _return_helper(res, helper_result, ret_cf) -@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( - 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.geomspace)( + _default_binary_wrapper(unyt_geomspace, _preserve_cosmo_factor) +) @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(dst, src, casting=casting, where=where) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][1]) @@ -1049,7 +796,6 @@ def prod( initial=np._NoValue, where=np._NoValue, ): - from unyt._array_functions import prod as unyt_prod helper_result = _prepare_array_func_args( a, @@ -1069,181 +815,22 @@ def prod( return _return_helper(res, helper_result, ret_cf, out=out) -@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( - 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) -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( - 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) -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( - 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( - 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( - 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( - 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( - 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( - 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( - 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.var)(_default_unary_wrapper(unyt_var, _preserve_cosmo_factor)) +implements(np.trace)(_default_unary_wrapper(unyt_trace, _preserve_cosmo_factor)) +implements(np.percentile)( + _default_unary_wrapper(unyt_percentile, _preserve_cosmo_factor) +) +implements(np.quantile)(_default_unary_wrapper(unyt_quantile, _preserve_cosmo_factor)) +implements(np.nanpercentile)( + _default_unary_wrapper(unyt_nanpercentile, _preserve_cosmo_factor) +) +implements(np.nanquantile)( + _default_unary_wrapper(unyt_nanquantile, _preserve_cosmo_factor) +) @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(a) ret_cf = _power_cosmo_factor( @@ -1255,46 +842,14 @@ 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.diff)(_default_unary_wrapper(unyt_diff, _preserve_cosmo_factor)) +implements(np.ediff1d)(_default_unary_wrapper(unyt_ediff1d, _preserve_cosmo_factor)) +implements(np.ptp)(_default_unary_wrapper(unyt_ptp, _preserve_cosmo_factor)) +# implements(np.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 @@ -1306,7 +861,6 @@ def pad(array, pad_width, mode="constant", **kwargs): @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) @@ -1317,7 +871,6 @@ def choose(a, choices, out=None, mode="raise"): @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( @@ -1329,7 +882,6 @@ def insert(arr, obj, values, axis=None): @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(a, b, rcond=rcond) ret_cf = _divide_cosmo_factor( @@ -1348,7 +900,6 @@ def linalg_lstsq(a, b, rcond=None): @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(a, b) ret_cf = _divide_cosmo_factor( @@ -1360,7 +911,6 @@ def linalg_solve(a, b): @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(a, b, axes=axes) ret_cf = _divide_cosmo_factor( @@ -1372,7 +922,6 @@ def linalg_tensorsolve(a, b, axes=None): @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(a) ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) @@ -1385,7 +934,6 @@ def linalg_eig(a): @implements(np.linalg.eigh) def linalg_eigh(a, UPLO="L"): - from unyt._array_functions import linalg_eigh as unyt_linalg_eigh helper_result = _prepare_array_func_args(a, UPLO=UPLO) ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) @@ -1396,24 +944,12 @@ def linalg_eigh(a, UPLO="L"): ) -@implements(np.linalg.eigvals) -def linalg_eigvals(a): - from unyt._array_functions import linalg_eigvals as unyt_linalg_eigvals - - 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) - - -@implements(np.linalg.eigvalsh) -def linalg_eigvalsh(a, UPLO="L"): - from unyt._array_functions import linalg_eigvalsh as unyt_linalg_eigvalsh - - 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) +implements(np.linalg.eigvals)( + _default_unary_wrapper(unyt_linalg_eigvals, _preserve_cosmo_factor) +) +implements(np.linalg.eigvalsh)( + _default_unary_wrapper(unyt_linalg_eigvalsh, _preserve_cosmo_factor) +) @implements(np.savetxt) @@ -1428,7 +964,6 @@ def savetxt( comments="# ", encoding=None, ): - from unyt._array_functions import savetxt as unyt_savetxt warnings.warn( "numpy.savetxt does not preserve units or cosmo_array information, " @@ -1471,7 +1006,6 @@ def apply_over_axes(func, a, axes): @implements(np.fill_diagonal) def fill_diagonal(a, val, wrap=False): - from unyt._array_functions import fill_diagonal as unyt_fill_diagonal helper_result = _prepare_array_func_args(a, val, wrap=wrap) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][1]) @@ -1484,25 +1018,11 @@ def fill_diagonal(a, val, wrap=False): unyt_fill_diagonal(a, val, **helper_result["kwargs"]) -@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( - 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.isin)(_default_comparison_wrapper(unyt_isin)) @implements(np.place) def place(arr, mask, vals): - from unyt._array_functions import place as unyt_place helper_result = _prepare_array_func_args(arr, mask, vals) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) @@ -1518,7 +1038,6 @@ def place(arr, mask, vals): @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(a, ind, v, mode=mode) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) @@ -1534,7 +1053,6 @@ def put(a, ind, v, mode="raise"): @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(arr, indices, values, axis) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) @@ -1550,7 +1068,6 @@ def put_along_axis(arr, indices, values, axis): @implements(np.putmask) def putmask(a, mask, values): - from unyt._array_functions import putmask as unyt_putmask helper_result = _prepare_array_func_args(a, mask, values) _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) @@ -1564,21 +1081,13 @@ def putmask(a, mask, values): unyt_putmask(a, mask, values) -@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(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.searchsorted)( + _default_binary_wrapper(unyt_searchsorted, _return_without_cosmo_factor) +) @implements(np.select) def select(condlist, choicelist, default=0): - from unyt._array_functions import select as unyt_select helper_result = _prepare_array_func_args(condlist, choicelist, default=default) helper_result_choicelist = _prepare_array_func_args(*choicelist) @@ -1591,22 +1100,13 @@ def select(condlist, choicelist, default=0): return _return_helper(res, helper_result, ret_cf) -@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(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.setdiff1d)( + _default_binary_wrapper(unyt_setdiff1d, _preserve_cosmo_factor) +) @implements(np.sinc) def sinc(x): - from unyt._array_functions import sinc as unyt_sinc # unyt just casts to array and calls the numpy implementation # so let's just hand off to them @@ -1624,7 +1124,6 @@ def clip( max=np._NoValue, **kwargs, ): - from unyt._array_functions import clip as unyt_clip # can't work out how to properly handle min and max, # just leave them in kwargs I guess (might be a numpy version conflict?) @@ -1652,7 +1151,6 @@ def clip( @implements(np.where) def where(condition, *args): - from unyt._array_functions import where as unyt_where helper_result = _prepare_array_func_args(condition, *args) if len(args) == 0: # just condition @@ -1668,24 +1166,8 @@ def where(condition, *args): return _return_helper(res, helper_result, ret_cf) -@implements(np.triu) -def triu(m, k=0): - from unyt._array_functions import triu as unyt_triu - - 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(m, k=0): - from unyt._array_functions import tril as unyt_tril - - 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.triu)(_default_unary_wrapper(unyt_triu, _preserve_cosmo_factor)) +implements(np.tril)(_default_unary_wrapper(unyt_tril, _preserve_cosmo_factor)) @implements(np.einsum) @@ -1698,7 +1180,6 @@ def einsum( casting="safe", optimize=False, ): - from unyt._array_functions import einsum as unyt_einsum helper_result = _prepare_array_func_args( subscripts, @@ -1719,45 +1200,17 @@ def einsum( return _return_helper(res, helper_result_operands, ret_cf, out=out) -@implements(np.convolve) -def convolve(a, v, mode="full"): - from unyt._array_functions import convolve as unyt_convolve - - 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(a, v, mode="valid"): - from unyt._array_functions import correlate as unyt_correlate - - 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(a, b, axes=2): - from unyt._array_functions import tensordot as unyt_tensordot - - 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.convolve)(_default_binary_wrapper(unyt_convolve, _multiply_cosmo_factor)) +implements(np.correlate)( + _default_binary_wrapper(unyt_correlate, _multiply_cosmo_factor) +) +implements(np.tensordot)( + _default_binary_wrapper(unyt_tensordot, _multiply_cosmo_factor) +) @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( p, discont=discont, axis=axis, period=period @@ -1773,7 +1226,6 @@ def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586): @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( x, xp, fp, left=left, right=right, period=period @@ -1785,7 +1237,6 @@ def interp(x, xp, fp, left=None, right=None, period=None): @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( arr, @@ -1804,21 +1255,13 @@ def array_repr(arr, max_line_width=None, precision=None, suppress_small=None): return rep -@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(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.linalg.outer)( + _default_binary_wrapper(unyt_linalg_outer, _multiply_cosmo_factor) +) @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(y, x=x, dx=dx, axis=axis) if x is None: @@ -1833,31 +1276,10 @@ def trapezoid(y, x=None, dx=1.0, axis=-1): return _return_helper(res, helper_result, ret_cf) -@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( - 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) - - -@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) +implements(np.in1d)(_default_comparison_wrapper(unyt_in1d)) +implements(np.take)(_default_unary_wrapper(unyt_take, _preserve_cosmo_factor)) +# Now we wrap functions that unyt does not handle explicitly: implements(np.average)(_propagate_cosmo_array_attributes(np.average._implementation)) implements(np.max)(_propagate_cosmo_array_attributes(np.max._implementation)) From 37870ce4d87eac09d37dacbe088a5dffe1ca5337 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 16:48:41 +0000 Subject: [PATCH 047/125] Remove version guard since we're going to pin a higher unyt version as dependency. --- swiftsimio/objects.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index b9dcdf95..afefd207 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -7,14 +7,9 @@ import unyt from unyt import unyt_array, unyt_quantity -from unyt.array import multiple_output_operators, _iterable +from unyt.array import multiple_output_operators, _iterable, POWER_MAPPING from numbers import Number as numeric_type -try: - from unyt.array import POWER_MAPPING -except ImportError: - raise ImportError("unyt >=2.9.0 required") - import sympy import numpy as np from numpy import ( From 218b156f392779afcc6311dbc53f806742573db7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 16:50:29 +0000 Subject: [PATCH 048/125] Cleanup another import guard since we'll require newer numpy now. --- swiftsimio/objects.py | 7 +------ 1 file changed, 1 insertion(+), 6 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index afefd207..5679d711 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -93,12 +93,7 @@ matmul, vecdot, ) -from numpy._core.umath import _ones_like - -try: - from numpy._core.umath import clip -except ImportError: - clip = None +from numpy._core.umath import _ones_like, clip # The scale factor! a = sympy.symbols("a") From 5bb02fb8e99825024cd60f39489f93ca60c289eb Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 17:48:10 +0000 Subject: [PATCH 049/125] Move array_func helper functions to _array_functions.py. --- swiftsimio/_array_functions.py | 457 ++++++++++++++++++++++++++++++--- swiftsimio/objects.py | 422 ++---------------------------- 2 files changed, 442 insertions(+), 437 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index fea96f12..12ffa392 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1,20 +1,8 @@ import warnings import numpy as np +import unyt from unyt import unyt_quantity, unyt_array -from .objects import ( - cosmo_array, - cosmo_quantity, - _prepare_array_func_args, - _multiply_cosmo_factor, - _divide_cosmo_factor, - _preserve_cosmo_factor, - _reciprocal_cosmo_factor, - _comparison_cosmo_factor, - _power_cosmo_factor, - _return_without_cosmo_factor, - _propagate_cosmo_array_attributes, - _copy_cosmo_array_attributes, -) +from swiftsimio import objects from unyt._array_functions import ( dot as unyt_dot, vdot as unyt_vdot, @@ -122,6 +110,409 @@ # numpy functions (we will actually wrap the functions below): +def _copy_cosmo_array_attributes(from_ca, to_ca): + if not isinstance(to_ca, objects.cosmo_array): + return to_ca + if hasattr(from_ca, "cosmo_factor"): + to_ca.cosmo_factor = from_ca.cosmo_factor + if hasattr(from_ca, "comoving"): + to_ca.comoving = from_ca.comoving + if hasattr(from_ca, "valid_transform"): + to_ca.valid_transform = from_ca.valid_transform + return to_ca + + +def _propagate_cosmo_array_attributes(func): + # can work on methods (obj is self) and functions (obj is first argument) + def wrapped(obj, *args, **kwargs): + ret = func(obj, *args, **kwargs) + if not isinstance(ret, objects.cosmo_array): + return ret + ret = _copy_cosmo_array_attributes(obj, ret) + return ret + + return wrapped + + +def _ensure_cosmo_array_or_quantity(func): + # can work on methods (obj is self) and functions (obj is first argument) + def wrapped(obj, *args, **kwargs): + ret = func(obj, *args, **kwargs) + if isinstance(ret, unyt_quantity) and not isinstance( + ret, objects.cosmo_quantity + ): + ret = objects.cosmo_quantity(ret) + elif isinstance(ret, unyt_array) and not isinstance(ret, objects.cosmo_array): + ret = objects.cosmo_array(ret) + if ( + isinstance(ret, objects.cosmo_array) + and not isinstance(ret, objects.cosmo_quantity) + and ret.shape == () + ): + ret = objects.cosmo_quantity(ret) + elif isinstance(ret, objects.cosmo_quantity) and ret.shape != (): + ret = objects.cosmo_array(ret) + return ret + + return wrapped + + +def _sqrt_cosmo_factor(ca_cf, **kwargs): + return _power_cosmo_factor( + ca_cf, (False, None), power=0.5 + ) # ufunc sqrt not supported + + +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): + # neither has cosmo_factor information: + return None + elif not ca1 and ca2: + # one is not a cosmo_array, allow e.g. multiplication by constants: + return cf2 + elif ca1 and not ca2: + # two is not a cosmo_array, allow e.g. multiplication by constants: + return cf1 + elif (ca1 and ca2) and ((cf1 is None) or (cf2 is None)): + # both cosmo_array but not both with cosmo_factor + # (both without shortcircuited above already): + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors ({cf1} and {cf2})," + f" discarding cosmo_factor in return value.", + RuntimeWarning, + ) + return None + elif (ca1 and ca2) and ((cf1 is not None) and (cf2 is not None)): + # both cosmo_array and both with cosmo_factor: + return cf1 * cf2 # cosmo_factor.__mul__ raises if scale factors differ + else: + raise RuntimeError("Unexpected state, please report this error on github.") + + +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: + # single argument, return promptly + return cf1 + elif (cf1 is None) and (cf2 is None): + # neither has cosmo_factor information: + return None + elif ca1 and not ca2: + # only one is cosmo_array + return cf1 + elif ca2 and not ca1: + # only one is cosmo_array + return cf2 + elif (ca1 and ca2) and (cf1 is None and cf2 is not None): + # both cosmo_array, but not both with cosmo_factor + # (both without shortcircuited above already): + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + return cf2 + elif (ca1 and ca2) and (cf1 is not None and cf2 is None): + # both cosmo_array, but not both with cosmo_factor + # (both without shortcircuited above already): + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf1}) for all arguments.", + RuntimeWarning, + ) + return cf1 + elif (ca1 and ca2) and (cf1 != cf2): + raise ValueError( + f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." + ) + elif (ca1 and ca2) and (cf1 == cf2): + return cf1 # or cf2, they're equal + else: + # not dealing with cosmo_arrays at all + return None + + +def _power_cosmo_factor(ca_cf1, ca_cf2, inputs=None, power=None): + if inputs is not None and power is not None: + raise ValueError + ca1, cf1 = ca_cf1 + ca2, cf2 = ca_cf2 + power = inputs[1] if inputs else power + if hasattr(power, "units"): + if not power.units.is_dimensionless: + raise ValueError("Exponent must be dimensionless.") + elif power.units is not unyt.dimensionless: + power = power.to_value(unyt.dimensionless) + # else power.units is unyt.dimensionless, do nothing + if ca2 and cf2.a_factor != 1.0: + raise ValueError("Exponent has scaling with scale factor != 1.") + if cf1 is None: + return None + return np.power(cf1, power) + + +def _square_cosmo_factor(ca_cf, **kwargs): + return _power_cosmo_factor(ca_cf, (False, None), power=2) + + +def _cbrt_cosmo_factor(ca_cf, **kwargs): + return _power_cosmo_factor(ca_cf, (False, None), power=1.0 / 3.0) + + +def _divide_cosmo_factor(ca_cf1, ca_cf2, **kwargs): + ca1, cf1 = ca_cf1 + ca2, cf2 = ca_cf2 + return _multiply_cosmo_factor( + (ca1, cf1), (ca2, _reciprocal_cosmo_factor((ca2, cf2))) + ) + + +def _reciprocal_cosmo_factor(ca_cf, **kwargs): + return _power_cosmo_factor(ca_cf, (False, None), power=-1) + + +def _passthrough_cosmo_factor(ca_cf, ca_cf2=None, **kwargs): + ca, cf = ca_cf + ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) + if ca_cf2 is None: + # no second argument, return promptly + return cf + elif (cf2 is not None) and cf != cf2: + # if both have cosmo_factor information and it differs this is an error + raise ValueError( + f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." + ) + else: + # passthrough is for e.g. ufuncs with a second dimensionless argument, + # so ok if cf2 is None and cf1 is not + return cf + + +def _return_without_cosmo_factor(ca_cf, ca_cf2=None, inputs=None, zero_comparison=None): + ca, cf = ca_cf + ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) + if ca_cf2 is None: + # no second argument + pass + elif ca and not ca2: + # one is not a cosmo_array, warn on e.g. comparison to constants: + if not zero_comparison: + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing" + f" assuming provided cosmo_factor ({cf}) for all arguments.", + RuntimeWarning, + ) + elif not ca and ca2: + # two is not a cosmo_array, warn on e.g. comparison to constants: + if not zero_comparison: + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing" + f" assuming provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + elif (ca and ca2) and (cf is not None and cf2 is None): + # one has no cosmo_factor information, warn: + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf}) for all arguments.", + RuntimeWarning, + ) + elif (ca and ca2) and (cf is None and cf2 is not None): + # two has no cosmo_factor information, warn: + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + elif (cf is not None) and (cf2 is not None) and (cf != cf2): + # both have cosmo_factor, don't match: + raise ValueError( + f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." + ) + elif (cf is not None) and (cf2 is not None) and (cf == cf2): + # both have cosmo_factor, and they match: + pass + else: + # not dealing with cosmo_arrays at all + pass + # return without cosmo_factor + return None + + +def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): + ca1, cf1 = ca_cf1 + ca2, cf2 = ca_cf2 + if (cf1 is None) and (cf2 is None): + return None + if cf1 is None and cf2 is not None: + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf2}) for all arguments.", + RuntimeWarning, + ) + if cf1 is not None and cf2 is None: + warnings.warn( + f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf1}) for all arguments.", + RuntimeWarning, + ) + if (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): + raise ValueError( + f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." + ) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + + +def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): + ca1, cf1 = ca_cf1 + ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) + try: + iter(inputs[0]) + except TypeError: + if ca1: + input1_iszero = not inputs[0].value and inputs[0] is not False + else: + input1_iszero = not inputs[0] and inputs[0] is not False + else: + if ca1: + input1_iszero = not inputs[0].value.any() + else: + input1_iszero = not inputs[0].any() + try: + iter(inputs[1]) + except IndexError: + input2_iszero = None + except TypeError: + if ca2: + input2_iszero = not inputs[1].value and inputs[1] is not False + else: + input2_iszero = not inputs[1] and inputs[1] is not False + else: + if ca2: + input2_iszero = not inputs[1].value.any() + else: + input2_iszero = not inputs[1].any() + zero_comparison = input1_iszero or input2_iszero + return _return_without_cosmo_factor( + ca_cf1, ca_cf2=ca_cf2, inputs=inputs, zero_comparison=zero_comparison + ) + + +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: + # >>> 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)) + 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 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]]): + # 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 + # 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] + 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, + ) + + def implements(numpy_function): """Register an __array_function__ implementation for cosmo_array objects.""" @@ -135,15 +526,17 @@ def decorator(func): def _return_helper(res, helper_result, ret_cf, out=None): if out is None: - if isinstance(res, unyt_quantity) and not isinstance(res, cosmo_quantity): - return cosmo_quantity( + if isinstance(res, unyt_quantity) and not isinstance( + res, objects.cosmo_quantity + ): + return objects.cosmo_quantity( res, comoving=helper_result["comoving"], cosmo_factor=ret_cf, compression=helper_result["compression"], ) - elif isinstance(res, unyt_array) and not isinstance(res, cosmo_array): - return cosmo_array( + elif isinstance(res, unyt_array) and not isinstance(res, objects.cosmo_array): + return objects.cosmo_array( res, comoving=helper_result["comoving"], cosmo_factor=ret_cf, @@ -159,7 +552,7 @@ def _return_helper(res, helper_result, ret_cf, out=None): if hasattr(out, "compression"): out.compression = helper_result["compression"] if res.shape == (): - return cosmo_quantity( + return objects.cosmo_quantity( res.to_value(res.units), res.units, bypass_validation=True, @@ -167,7 +560,7 @@ def _return_helper(res, helper_result, ret_cf, out=None): cosmo_factor=ret_cf, compression=helper_result["compression"], ) - return cosmo_array( + return objects.cosmo_array( res.to_value(res.units), res.units, bypass_validation=True, @@ -353,7 +746,7 @@ def histogram(a, bins=10, range=None, density=None, weights=None): else: ret_cf_counts = ret_cf_dens if density else None if isinstance(counts, unyt_array): - counts = cosmo_array( + counts = objects.cosmo_array( counts.to_value(counts.units), counts.units, comoving=helper_result["comoving"], @@ -413,7 +806,7 @@ 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_w["kw_ca_cfs"]["weights"]) if isinstance(counts, unyt_array): - counts = cosmo_array( + counts = objects.cosmo_array( counts.to_value(counts.units), counts.units, comoving=helper_result_w["comoving"], @@ -465,7 +858,7 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): else: ret_cf_counts = _reciprocal_cosmo_factor((ret_cf_xy is not None, ret_cf_xy)) if isinstance(counts, unyt_array): - counts = cosmo_array( + counts = objects.cosmo_array( counts.to_value(counts.units), counts.units, comoving=helper_result["comoving"], @@ -528,7 +921,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) if isinstance(counts, unyt_array): - counts = cosmo_array( + counts = objects.cosmo_array( counts.to_value(counts.units), counts.units, comoving=helper_result_w["comoving"], @@ -570,7 +963,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): (ret_cf_sample is not None, ret_cf_sample) ) if isinstance(counts, unyt_array): - counts = cosmo_array( + counts = objects.cosmo_array( counts.to_value(counts.units), counts.units, comoving=helper_result["comoving"], @@ -579,10 +972,10 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): ) return ( counts, - [ + tuple( _return_helper(b, helper_result, ret_cf) for b, helper_result, ret_cf in zip(bins, helper_results, ret_cfs) - ], + ), ) @@ -777,7 +1170,7 @@ def copyto(dst, src, casting="same_kind", where=True): 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): + if isinstance(src, objects.cosmo_array): comoving = getattr(dst, "comoving", None) if comoving: src.convert_to_comoving() @@ -1027,7 +1420,7 @@ def place(arr, mask, vals): 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): + if isinstance(vals, objects.cosmo_array): comoving = getattr(arr, "comoving", None) if comoving: vals.convert_to_comoving() @@ -1042,7 +1435,7 @@ def put(a, ind, v, mode="raise"): 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): + if isinstance(v, objects.cosmo_array): comoving = getattr(a, "comoving", None) if comoving: v.convert_to_comoving() @@ -1057,7 +1450,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): + if isinstance(values, objects.cosmo_array): comoving = getattr(arr, "comoving", None) if comoving: values.convert_to_comoving() @@ -1072,7 +1465,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): + if isinstance(values, objects.cosmo_array): comoving = getattr(a, "comoving", None) if comoving: values.convert_to_comoving() diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 5679d711..4a2ad274 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -3,9 +3,6 @@ unyt_array that we use, called cosmo_array. """ -import warnings - -import unyt from unyt import unyt_array, unyt_quantity from unyt.array import multiple_output_operators, _iterable, POWER_MAPPING from numbers import Number as numeric_type @@ -94,6 +91,23 @@ vecdot, ) from numpy._core.umath import _ones_like, clip +from ._array_functions import ( + _propagate_cosmo_array_attributes, + _ensure_cosmo_array_or_quantity, + _sqrt_cosmo_factor, + _multiply_cosmo_factor, + _preserve_cosmo_factor, + _power_cosmo_factor, + _square_cosmo_factor, + _cbrt_cosmo_factor, + _divide_cosmo_factor, + _reciprocal_cosmo_factor, + _passthrough_cosmo_factor, + _return_without_cosmo_factor, + _arctan2_cosmo_factor, + _comparison_cosmo_factor, + _prepare_array_func_args, +) # The scale factor! a = sympy.symbols("a") @@ -104,407 +118,6 @@ def __init__(self, message="Could not convert to comoving coordinates"): self.message = message -def _copy_cosmo_array_attributes(from_ca, to_ca): - if not isinstance(to_ca, cosmo_array): - return to_ca - if hasattr(from_ca, "cosmo_factor"): - to_ca.cosmo_factor = from_ca.cosmo_factor - if hasattr(from_ca, "comoving"): - to_ca.comoving = from_ca.comoving - if hasattr(from_ca, "valid_transform"): - to_ca.valid_transform = from_ca.valid_transform - return to_ca - - -def _propagate_cosmo_array_attributes(func): - # can work on methods (obj is self) and functions (obj is first argument) - def wrapped(obj, *args, **kwargs): - ret = func(obj, *args, **kwargs) - if not isinstance(ret, cosmo_array): - return ret - ret = _copy_cosmo_array_attributes(obj, ret) - return ret - - return wrapped - - -def _ensure_cosmo_array_or_quantity(func): - # can work on methods (obj is self) and functions (obj is first argument) - def wrapped(obj, *args, **kwargs): - ret = func(obj, *args, **kwargs) - if isinstance(ret, unyt_quantity) and not isinstance(ret, cosmo_quantity): - ret = cosmo_quantity(ret) - elif isinstance(ret, unyt_array) and not isinstance(ret, cosmo_array): - ret = cosmo_array(ret) - if ( - isinstance(ret, cosmo_array) - and not isinstance(ret, cosmo_quantity) - and ret.shape == () - ): - ret = cosmo_quantity(ret) - elif isinstance(ret, cosmo_quantity) and ret.shape != (): - ret = cosmo_array(ret) - return ret - - return wrapped - - -def _sqrt_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor( - ca_cf, (False, None), power=0.5 - ) # ufunc sqrt not supported - - -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): - # neither has cosmo_factor information: - return None - elif not ca1 and ca2: - # one is not a cosmo_array, allow e.g. multiplication by constants: - return cf2 - elif ca1 and not ca2: - # two is not a cosmo_array, allow e.g. multiplication by constants: - return cf1 - elif (ca1 and ca2) and ((cf1 is None) or (cf2 is None)): - # both cosmo_array but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors ({cf1} and {cf2})," - f" discarding cosmo_factor in return value.", - RuntimeWarning, - ) - return None - elif (ca1 and ca2) and ((cf1 is not None) and (cf2 is not None)): - # both cosmo_array and both with cosmo_factor: - return cf1 * cf2 # cosmo_factor.__mul__ raises if scale factors differ - else: - raise RuntimeError("Unexpected state, please report this error on github.") - - -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: - # single argument, return promptly - return cf1 - elif (cf1 is None) and (cf2 is None): - # neither has cosmo_factor information: - return None - elif ca1 and not ca2: - # only one is cosmo_array - return cf1 - elif ca2 and not ca1: - # only one is cosmo_array - return cf2 - elif (ca1 and ca2) and (cf1 is None and cf2 is not None): - # both cosmo_array, but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - return cf2 - elif (ca1 and ca2) and (cf1 is not None and cf2 is None): - # both cosmo_array, but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf1}) for all arguments.", - RuntimeWarning, - ) - return cf1 - elif (ca1 and ca2) and (cf1 != cf2): - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." - ) - elif (ca1 and ca2) and (cf1 == cf2): - return cf1 # or cf2, they're equal - else: - # not dealing with cosmo_arrays at all - return None - - -def _power_cosmo_factor(ca_cf1, ca_cf2, inputs=None, power=None): - if inputs is not None and power is not None: - raise ValueError - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - power = inputs[1] if inputs else power - if hasattr(power, "units"): - if not power.units.is_dimensionless: - raise ValueError("Exponent must be dimensionless.") - elif power.units is not unyt.dimensionless: - power = power.to_value(unyt.dimensionless) - # else power.units is unyt.dimensionless, do nothing - if ca2 and cf2.a_factor != 1.0: - raise ValueError("Exponent has scaling with scale factor != 1.") - if cf1 is None: - return None - return np.power(cf1, power) - - -def _square_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=2) - - -def _cbrt_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=1.0 / 3.0) - - -def _divide_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - return _multiply_cosmo_factor( - (ca1, cf1), (ca2, _reciprocal_cosmo_factor((ca2, cf2))) - ) - - -def _reciprocal_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=-1) - - -def _passthrough_cosmo_factor(ca_cf, ca_cf2=None, **kwargs): - ca, cf = ca_cf - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # no second argument, return promptly - return cf - elif (cf2 is not None) and cf != cf2: - # if both have cosmo_factor information and it differs this is an error - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." - ) - else: - # passthrough is for e.g. ufuncs with a second dimensionless argument, - # so ok if cf2 is None and cf1 is not - return cf - - -def _return_without_cosmo_factor(ca_cf, ca_cf2=None, inputs=None, zero_comparison=None): - ca, cf = ca_cf - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # no second argument - pass - elif ca and not ca2: - # one is not a cosmo_array, warn on e.g. comparison to constants: - if not zero_comparison: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing" - f" assuming provided cosmo_factor ({cf}) for all arguments.", - RuntimeWarning, - ) - elif not ca and ca2: - # two is not a cosmo_array, warn on e.g. comparison to constants: - if not zero_comparison: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing" - f" assuming provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - elif (ca and ca2) and (cf is not None and cf2 is None): - # one has no cosmo_factor information, warn: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf}) for all arguments.", - RuntimeWarning, - ) - elif (ca and ca2) and (cf is None and cf2 is not None): - # two has no cosmo_factor information, warn: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - elif (cf is not None) and (cf2 is not None) and (cf != cf2): - # both have cosmo_factor, don't match: - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." - ) - elif (cf is not None) and (cf2 is not None) and (cf == cf2): - # both have cosmo_factor, and they match: - pass - else: - # not dealing with cosmo_arrays at all - pass - # return without cosmo_factor - return None - - -def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - if (cf1 is None) and (cf2 is None): - return None - if cf1 is None and cf2 is not None: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) - if cf1 is not None and cf2 is None: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf1}) for all arguments.", - RuntimeWarning, - ) - if (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." - ) - return cosmo_factor(a**0, scale_factor=cf1.scale_factor) - - -def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - try: - iter(inputs[0]) - except TypeError: - if ca1: - input1_iszero = not inputs[0].value and inputs[0] is not False - else: - input1_iszero = not inputs[0] and inputs[0] is not False - else: - if ca1: - input1_iszero = not inputs[0].value.any() - else: - input1_iszero = not inputs[0].any() - try: - iter(inputs[1]) - except IndexError: - input2_iszero = None - except TypeError: - if ca2: - input2_iszero = not inputs[1].value and inputs[1] is not False - else: - input2_iszero = not inputs[1] and inputs[1] is not False - else: - if ca2: - input2_iszero = not inputs[1].value.any() - else: - input2_iszero = not inputs[1].any() - zero_comparison = input1_iszero or input2_iszero - return _return_without_cosmo_factor( - ca_cf1, ca_cf2=ca_cf2, inputs=inputs, zero_comparison=zero_comparison - ) - - -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: - # >>> 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)) - 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 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]]): - # 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 - # 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] - 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 @@ -771,7 +384,6 @@ class cosmo_array(unyt_array): """ - # TODO: _cosmo_factor_ufunc_registry = { add: _preserve_cosmo_factor, subtract: _preserve_cosmo_factor, From f891b85ffea316f67889f4af08f7ef84a7c3f2eb Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 17:48:55 +0000 Subject: [PATCH 050/125] Replace awkward unyt usage with new numpy supported cosmo versions. --- swiftsimio/visualisation/power_spectrum.py | 54 +++++++++++----------- 1 file changed, 26 insertions(+), 28 deletions(-) diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index 2940f36e..e602f221 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -5,12 +5,12 @@ from numpy import float32, float64, int32, zeros, ndarray import numpy as np import scipy.fft -import unyt from swiftsimio.optional_packages import tqdm from swiftsimio.accelerated import jit, NUM_THREADS, prange -from swiftsimio import cosmo_array +from swiftsimio import cosmo_array, cosmo_quantity from swiftsimio.reader import __SWIFTGroupDataset +from swiftsimio._array_functions import _reciprocal_cosmo_factor from typing import Optional, Dict, Tuple @@ -206,12 +206,14 @@ def render_to_deposit( if positions.comoving: if not quantity.compatible_with_comoving(): raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' + f'Physical quantity "{project}" is not compatible with comoving ' + "coordinates!" ) else: if not quantity.compatible_with_physical(): raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' + f'Comoving quantity "{project}" is not compatible with physical ' + "coordinates!" ) # Get the box size @@ -259,7 +261,7 @@ def folded_depositions_to_power_spectrum( box_size: cosmo_array, number_of_wavenumber_bins: int, cross_depositions: Optional[Dict[int, cosmo_array]] = None, - wavenumber_range: Optional[Tuple[unyt.unyt_quantity]] = None, + wavenumber_range: Optional[Tuple[cosmo_quantity]] = None, log_wavenumber_bins: bool = True, workers: Optional[int] = None, minimal_sample_modes: Optional[int] = 0, @@ -267,7 +269,7 @@ def folded_depositions_to_power_spectrum( track_progress: bool = False, transition: str = "simple", shot_noise_norm: Optional[float] = None, -) -> Tuple[unyt.unyt_array]: +) -> Tuple[cosmo_array]: """ Convert some folded depositions to power spectra. @@ -284,7 +286,7 @@ def folded_depositions_to_power_spectrum( The number of bins to use in the power spectrum. cross_depositions: Optional[dict[int, cosmo_array]] An optional dictionary of cross-depositions, where the key is the folding. - wavenumber_range: Optional[tuple[unyt.unyt_quantity]] + wavenumber_range: Optional[tuple[cosmo_quantity]] The range of wavenumbers to use. Officially optional, but is required for now. log_wavenumber_bins: bool @@ -313,13 +315,13 @@ def folded_depositions_to_power_spectrum( Returns ------- - wavenumber_bins: unyt.unyt_array[float32] + wavenumber_bins: cosmo_array[float32] The wavenumber bins. - wavenumber_centers: unyt.unyt_array[float32] + wavenumber_centers: cosmo_array[float32] The centers of the wavenumber bins. - power_spectrum: unyt.unyt_array[float32] + power_spectrum: cosmo_array[float32] The power spectrum. folding_tracker: np.array @@ -336,32 +338,28 @@ def folded_depositions_to_power_spectrum( raise NotImplementedError if log_wavenumber_bins: - wavenumber_bins = unyt.unyt_array( - np.logspace( - np.log10(min(wavenumber_range).v), - np.log10(max(wavenumber_range).v), - number_of_wavenumber_bins + 1, - ), - wavenumber_range[0].units, - name="Wavenumber bins", + wavenumber_bins = np.geomspace( + np.min(cosmo_array(wavenumber_range)), + np.max(cosmo_array(wavenumber_range)), + number_of_wavenumber_bins + 1, ) + else: - wavenumber_bins = unyt.unyt_array( - np.linspace( - min(wavenumber_range).v, - max(wavenumber_range).v, - number_of_wavenumber_bins + 1, - ), - wavenumber_range[0].units, - name="Wavenumber bins", + wavenumber_bins = np.linspace( + np.min(cosmo_array(wavenumber_range)), + np.max(cosmo_array(wavenumber_range)), + number_of_wavenumber_bins + 1, ) + wavenumber_bins.name = "Wavenumber bins" wavenumber_centers = 0.5 * (wavenumber_bins[1:] + wavenumber_bins[:-1]) wavenumber_centers.name = r"Wavenumbers $k$" box_volume = np.prod(box_size) - power_spectrum = unyt.unyt_array( + power_spectrum = cosmo_array( np.zeros(number_of_wavenumber_bins), units=box_volume.units, + comoving=box_volume.comoving, + cosmo_factor=box_volume.cosmo_factor**-1, name="Power spectrum $P(k)$", ) folding_tracker = np.ones(number_of_wavenumber_bins, dtype=float) @@ -406,7 +404,7 @@ def folded_depositions_to_power_spectrum( maximally_sampled_wavenumber = np.max( folded_wavenumber_centers[use_bins] ) - cutoff_wavenumber = min( + cutoff_wavenumber = np.min( cutoff_above_wavenumber_fraction * maximally_sampled_wavenumber, cutoff_above_wavenumber_fraction * cutoff_wavenumber, ) From d120fdbdedc56a44be66ab3d286c2f05ad130eea Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 31 Jan 2025 22:19:33 +0000 Subject: [PATCH 051/125] Point requirements to unyt github main branch (temporarily). --- requirements.txt | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/requirements.txt b/requirements.txt index 7e7a9182..2d58369a 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,4 @@ astropy numpy h5py numba -unyt +git+https://github.com/yt-project/unyt.git From 44477fb80361d6863f4deec7483c205b954d2916 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 1 Feb 2025 09:09:07 +0000 Subject: [PATCH 052/125] Handle comoving is None in str representation --- swiftsimio/objects.py | 2 ++ 1 file changed, 2 insertions(+) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 4a2ad274..2017cce8 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -605,6 +605,8 @@ def __array_finalize__(self, obj): def __str__(self): if self.comoving: comoving_str = "(Comoving)" + elif self.comoving is None: + comoving_str = "(Physical/comoving not set)" else: comoving_str = "(Physical)" From 04e522abc3463303c3e65257b339ee3cd7036336 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 5 Feb 2025 13:39:54 +0000 Subject: [PATCH 053/125] Refactor to keep code for casting return types in one place. --- swiftsimio/_array_functions.py | 164 ++++++++++++++------------------- swiftsimio/objects.py | 45 +++------ 2 files changed, 82 insertions(+), 127 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 12ffa392..d4206548 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -134,24 +134,46 @@ def wrapped(obj, *args, **kwargs): return wrapped +def _promote_unyt_to_cosmo(input_object): + # converts a unyt_quantity to a cosmo_quantity, or a unyt_array to a cosmo_array + if isinstance(input_object, unyt_quantity) and not isinstance( + input_object, objects.cosmo_quantity + ): + ret = objects.cosmo_quantity(input_object) + elif isinstance(input_object, unyt_array) and not isinstance( + input_object, objects.cosmo_array + ): + ret = objects.cosmo_array(input_object) + else: + ret = input_object + return ret + + +def _ensure_array_or_quantity_matches_shape(input_object): + if ( + isinstance(input_object, objects.cosmo_array) + and not isinstance(input_object, objects.cosmo_quantity) + and input_object.shape == () + ): + ret = objects.cosmo_quantity(input_object) + elif isinstance(input_object, objects.cosmo_quantity) and input_object.shape != (): + ret = objects.cosmo_array(input_object) + else: + ret = input_object + return ret + + def _ensure_cosmo_array_or_quantity(func): # can work on methods (obj is self) and functions (obj is first argument) def wrapped(obj, *args, **kwargs): ret = func(obj, *args, **kwargs) - if isinstance(ret, unyt_quantity) and not isinstance( - ret, objects.cosmo_quantity - ): - ret = objects.cosmo_quantity(ret) - elif isinstance(ret, unyt_array) and not isinstance(ret, objects.cosmo_array): - ret = objects.cosmo_array(ret) - if ( - isinstance(ret, objects.cosmo_array) - and not isinstance(ret, objects.cosmo_quantity) - and ret.shape == () - ): - ret = objects.cosmo_quantity(ret) - elif isinstance(ret, objects.cosmo_quantity) and ret.shape != (): - ret = objects.cosmo_array(ret) + if isinstance(ret, tuple): + ret = tuple( + _ensure_array_or_quantity_matches_shape(_promote_unyt_to_cosmo(item)) + for item in ret + ) + else: + ret = _ensure_array_or_quantity_matches_shape(_promote_unyt_to_cosmo(ret)) return ret return wrapped @@ -525,49 +547,16 @@ def decorator(func): def _return_helper(res, helper_result, ret_cf, out=None): - if out is None: - if isinstance(res, unyt_quantity) and not isinstance( - res, objects.cosmo_quantity - ): - return objects.cosmo_quantity( - res, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf, - compression=helper_result["compression"], - ) - elif isinstance(res, unyt_array) and not isinstance(res, objects.cosmo_array): - return objects.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"): + res = _promote_unyt_to_cosmo(res) + if isinstance(res, objects.cosmo_array): # also recognizes cosmo_quantity + res.comoving = helper_result["comoving"] + res.cosmo_factor = ret_cf + res.compression = helper_result["compression"] + if isinstance(out, objects.cosmo_array): # also recognizes cosmo_quantity out.comoving = helper_result["comoving"] - if hasattr(out, "cosmo_factor"): out.cosmo_factor = ret_cf - if hasattr(out, "compression"): out.compression = helper_result["compression"] - if res.shape == (): - return objects.cosmo_quantity( - res.to_value(res.units), - res.units, - bypass_validation=True, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf, - compression=helper_result["compression"], - ) - return objects.cosmo_array( - res.to_value(res.units), - res.units, - bypass_validation=True, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf, - compression=helper_result["compression"], - ) + return res def _default_unary_wrapper(unyt_func, cosmo_factor_wrapper): @@ -745,14 +734,11 @@ def histogram(a, bins=10, range=None, density=None, weights=None): ) else: ret_cf_counts = ret_cf_dens if density else None - if isinstance(counts, unyt_array): - counts = objects.cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf_counts, - compression=helper_result["compression"], - ) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity + counts.comoving = helper_result["comoving"] + counts.cosmo_factor = ret_cf_counts + counts.compression = helper_result["compression"] return counts, _return_helper(bins, helper_result, ret_cf_bins) @@ -805,14 +791,13 @@ 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_w["kw_ca_cfs"]["weights"]) - if isinstance(counts, unyt_array): - counts = objects.cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result_w["comoving"], - cosmo_factor=ret_cf_w, - compression=helper_result_w["compression"], - ) + counts = _promote_unyt_to_cosmo(counts) + if isinstance( + counts, objects.cosmo_array + ): # also recognizes cosmo_quantity + counts.comoving = helper_result_w["comoving"] + counts.cosmo_factor = ret_cf_w + counts.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 @@ -857,14 +842,11 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): ) else: ret_cf_counts = _reciprocal_cosmo_factor((ret_cf_xy is not None, ret_cf_xy)) - if isinstance(counts, unyt_array): - counts = objects.cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf_counts, - compression=helper_result["compression"], - ) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity + counts.comoving = helper_result["comoving"] + counts.cosmo_factor = ret_cf_counts + counts.compression = helper_result["compression"] return ( counts, _return_helper(xbins, helper_result_x, ret_cf_x), @@ -920,14 +902,11 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): ) if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) - if isinstance(counts, unyt_array): - counts = objects.cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result_w["comoving"], - cosmo_factor=ret_cf_w, - compression=helper_result_w["compression"], - ) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): + counts.comoving = helper_result_w["comoving"] + counts.cosmo_factor = ret_cf_w + counts.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 @@ -962,14 +941,11 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): ret_cf_counts = _reciprocal_cosmo_factor( (ret_cf_sample is not None, ret_cf_sample) ) - if isinstance(counts, unyt_array): - counts = objects.cosmo_array( - counts.to_value(counts.units), - counts.units, - comoving=helper_result["comoving"], - cosmo_factor=ret_cf_counts, - compression=helper_result["compression"], - ) + counts = _promote_unyt_to_cosmo(counts) + if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity + counts.comoving = helper_result["comoving"] + counts.cosmo_factor = ret_cf_counts + counts.compression = helper_result["compression"] return ( counts, tuple( diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 2017cce8..48ef9341 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -865,21 +865,17 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): else: ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*ca_cfs, inputs=inputs) - ret = super().__array_ufunc__(ufunc, method, *inputs, **kwargs) + ret = _ensure_cosmo_array_or_quantity(super().__array_ufunc__)( + ufunc, method, *inputs, **kwargs + ) # if we get a tuple we have multiple return values to deal with - # if unyt returns a bare ndarray, do the same - # otherwise we create a view and attach our attributes if isinstance(ret, tuple): - ret = tuple( - r.view(type(self)) if isinstance(r, unyt_array) else r for r in ret - ) for r in ret: - if isinstance(r, type(self)): + if isinstance(r, cosmo_array): # also recognizes cosmo_quantity r.comoving = helper_result["comoving"] r.cosmo_factor = ret_cf r.compression = helper_result["compression"] - if isinstance(ret, unyt_array): - ret = ret.view(type(self)) + elif isinstance(ret, cosmo_array): # also recognizes cosmo_quantity ret.comoving = helper_result["comoving"] ret.cosmo_factor = ret_cf ret.compression = helper_result["compression"] @@ -887,13 +883,13 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): out = kwargs.pop("out") if ufunc not in multiple_output_operators: out = out[0] - if isinstance(out, cosmo_array): + if isinstance(out, cosmo_array): # also recognizes cosmo_quantity out.comoving = helper_result["comoving"] out.cosmo_factor = ret_cf out.compression = helper_result["compression"] else: for o in out: - if isinstance(o, type(self)): + if isinstance(o, cosmo_array): # also recognizes cosmo_quantity o.comoving = helper_result["comoving"] o.cosmo_factor = ret_cf o.compression = helper_result["compression"] @@ -926,22 +922,13 @@ def __array_function__(self, func, types, args, kwargs): return NotImplemented if func in _HANDLED_FUNCTIONS: - ret = _HANDLED_FUNCTIONS[func](*args, **kwargs) + function_to_invoke = _HANDLED_FUNCTIONS[func] elif func in _UNYT_HANDLED_FUNCTIONS: - ret = _UNYT_HANDLED_FUNCTIONS[func](*args, **kwargs) + function_to_invoke = _UNYT_HANDLED_FUNCTIONS[func] else: # default to numpy's private implementation - ret = func._implementation(*args, **kwargs) - if ( - isinstance(ret, cosmo_array) - and ret.shape == () - and not isinstance(ret, cosmo_quantity) - ): - return cosmo_quantity(ret) - elif isinstance(ret, cosmo_quantity) and ret.shape != (): - return cosmo_array(ret) - else: - return ret + function_to_invoke = func._implementation + return function_to_invoke(*args, **kwargs) class cosmo_quantity(cosmo_array, unyt_quantity): @@ -1033,7 +1020,7 @@ def __new__( bypass_validation or isinstance(input_scalar, (numeric_type, np.number, np.ndarray)) ): - raise RuntimeError("unyt_quantity values must be numeric") + raise RuntimeError("cosmo_quantity values must be numeric") units = getattr(input_scalar, "units", None) if units is None else units name = getattr(input_scalar, "name", None) if name is None else name @@ -1071,11 +1058,3 @@ def __new__( if ret.size > 1: raise RuntimeError("cosmo_quantity instances must be scalars") return ret - - @_propagate_cosmo_array_attributes - def reshape(self, shape, /, *, order="C"): - reshaped = unyt_array.reshape(self, shape, order=order) - if shape == () or shape is None: - return reshaped - else: - return cosmo_array(reshaped) From 65cd3fa3e3dfc320aa4ad4db0fd0000d233872d2 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 5 Feb 2025 13:41:23 +0000 Subject: [PATCH 054/125] Cleanup while checking unyt usage. --- swiftsimio/conversions.py | 10 +++--- swiftsimio/metadata/objects.py | 59 +++++++++++++++------------------- swiftsimio/swiftsnap.py | 1 - tests/test_read_ic.py | 7 ++-- 4 files changed, 35 insertions(+), 42 deletions(-) diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index d246c1a1..4fa415fe 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -4,12 +4,11 @@ """ from swiftsimio.optional_packages import ASTROPY_AVAILABLE -import unyt +from swiftsimio.objects import cosmo_quantity if ASTROPY_AVAILABLE: from astropy.cosmology import w0waCDM from astropy.cosmology.core import Cosmology - import astropy.version import astropy.constants as const import astropy.units as astropy_units import numpy as np @@ -74,7 +73,7 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: correct parameters. """ - H0 = unyt.unyt_quantity(cosmo["H0 [internal units]"][0], units=1.0 / units.time) + H0 = cosmo_quantity(cosmo["H0 [internal units]"][0], units=1.0 / units.time) Omega_b = cosmo["Omega_b"][0] Omega_lambda = cosmo["Omega_lambda"][0] @@ -97,13 +96,13 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: # expressions taken directly from astropy, since they do no longer # allow access to these attributes (since version 5.1+) critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value - a_B_c2 = (4.0 * const.sigma_sb / const.c ** 3).cgs.value + a_B_c2 = (4.0 * const.sigma_sb / const.c**3).cgs.value # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy. # This is an exact inversion of the procedure performed in astropy. critical_density_0 = astropy_units.Quantity( critdens_const * H0.to("1/s").value ** 2, - astropy_units.g / astropy_units.cm ** 3, + astropy_units.g / astropy_units.cm**3, ) Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0) @@ -142,7 +141,6 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: m_nu=ap_m_nu, ) - else: def swift_cosmology_to_astropy(cosmo: dict, units) -> dict: diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index 411deeb7..b0cefcf6 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -113,12 +113,14 @@ def postprocess_header(self): # items including the scale factor. # These must be unpacked as they are stored as length-1 arrays - header_unpack_float_units = metadata.metadata_fields.generate_units_header_unpack_single_float( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, + header_unpack_float_units = ( + metadata.metadata_fields.generate_units_header_unpack_single_float( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, + ) ) for field, names in metadata.metadata_fields.header_unpack_single_float.items(): try: @@ -164,15 +166,19 @@ def postprocess_header(self): self.scale_factor = 1.0 # These are just read straight in to variables - header_unpack_arrays_units = metadata.metadata_fields.generate_units_header_unpack_arrays( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, + header_unpack_arrays_units = ( + metadata.metadata_fields.generate_units_header_unpack_arrays( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, + ) ) - header_unpack_arrays_cosmo_args = metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( - self.scale_factor + header_unpack_arrays_cosmo_args = ( + metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( + self.scale_factor + ) ) for field, name in metadata.metadata_fields.header_unpack_arrays.items(): @@ -604,18 +610,14 @@ def get_units(unit_attribute): # Need to check if the exponent is 0 manually because of float precision unit_exponent = unit_attribute[f"U_{exponent} exponent"][0] if unit_exponent != 0.0: - units *= unit ** unit_exponent + units *= unit**unit_exponent except KeyError: # Can't load that data! # We should probably warn the user here... pass - # Deal with case where we _really_ have a dimensionless quantity. Comparing with - # 1.0 doesn't work, beacause in these cases unyt reverts to a floating point - # comparison. - try: - units.units - except AttributeError: + # Deal with case where we _really_ have a dimensionless quantity. + if not hasattr(units, "units"): units = None return units @@ -702,7 +704,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a ** cosmo_exponent + a_factor_this_dataset = a**cosmo_exponent return cosmo_factor(a_factor_this_dataset, current_scale_factor) @@ -847,17 +849,8 @@ def handle(self): Property that gets the file handle, which can be shared with other objects for efficiency reasons. """ - if isinstance(self._handle, h5py.File): - # Can be open or closed, let's test. - try: - file = self._handle.file - - return self._handle - except ValueError: - # This will be the case if there is no active file handle - pass - - self._handle = h5py.File(self.filename, "r") + if not self._handle: # if self._handle is None, or if file closed (h5py #1363) + self._handle = h5py.File(self.filename, "r") return self._handle diff --git a/swiftsimio/swiftsnap.py b/swiftsimio/swiftsnap.py index 3c586460..ac74b601 100755 --- a/swiftsimio/swiftsnap.py +++ b/swiftsimio/swiftsnap.py @@ -64,7 +64,6 @@ def decode(bytestring: bytes) -> str: def swiftsnap(): import swiftsimio as sw from swiftsimio.metadata.objects import metadata_discriminator - import unyt from swiftsimio.metadata.particle import particle_name_underscores from textwrap import wrap diff --git a/tests/test_read_ic.py b/tests/test_read_ic.py index 4db743ac..112d4ffb 100644 --- a/tests/test_read_ic.py +++ b/tests/test_read_ic.py @@ -1,6 +1,7 @@ from swiftsimio import load from swiftsimio import Writer from swiftsimio.units import cosmo_units +from swiftsimio import cosmo_array import unyt import numpy as np @@ -33,7 +34,7 @@ def simple_snapshot_data(): x = Writer(cosmo_units, boxsize) # 32^3 particles. - n_p = 32 ** 3 + n_p = 32**3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) @@ -70,7 +71,9 @@ def test_reading_ic_units(simple_snapshot_data, field): data = load(test_filename) - assert unyt.array.allclose_units( + assert isinstance(getattr(data.gas, field), cosmo_array) + # np.allclose checks unit consistency + assert np.allclose( getattr(data.gas, field), getattr(simple_snapshot_data.gas, field), rtol=1.0e-4 ) return From 0d9a1c8d8f1f7a4de3383c52be38d7f6af44e98e Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 5 Feb 2025 13:52:32 +0000 Subject: [PATCH 055/125] Add a test for cosmo_quantity conversion. --- tests/test_cosmo_array.py | 15 ++++ tests/test_cosmo_array_attrs.py | 146 +++++++++++++++++--------------- 2 files changed, 93 insertions(+), 68 deletions(-) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index bd66d3e1..2bac216a 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -686,6 +686,21 @@ def test_propagation_func(self, func, args): assert res.cosmo_factor == cosmo_factor(a**1, 1.0) assert res.valid_transform is True + def test_scalar_return_func(self): + """ + Make sure that default-wrapped functions that take a cosmo_array and return a + scalar convert to a cosmo_quantity. + """ + ca = cosmo_array( + np.arange(3), + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a**1, 1.0), + valid_transform=True, + ) + res = np.min(ca) + assert isinstance(res, cosmo_quantity) + @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"]) def test_propagation_props(self, prop): cq = cosmo_quantity( diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_attrs.py index bfa625d2..c67bab4c 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_attrs.py @@ -105,7 +105,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 0.5), + cosmo_factor=cosmo_factor(a**1, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -114,7 +114,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 0.5), + cosmo_factor=cosmo_factor(a**1, 0.5), comoving=False, ) assert not arr.compatible_with_comoving() @@ -123,7 +123,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 0, 0.5), + cosmo_factor=cosmo_factor(a**0, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -132,7 +132,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 0, 0.5), + cosmo_factor=cosmo_factor(a**0, 0.5), comoving=False, ) assert arr.compatible_with_comoving() @@ -141,7 +141,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), comoving=True, ) assert arr.compatible_with_comoving() @@ -150,7 +150,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), comoving=False, ) assert arr.compatible_with_comoving() @@ -176,7 +176,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.ones_like(inp) assert res.to_value(u.kpc) == 1 @@ -207,7 +207,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): @@ -221,7 +221,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): res = inp1 + inp2 @@ -233,13 +233,13 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Ufunc arguments have cosmo_factors that differ" @@ -250,7 +250,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp + inp assert res.to_value(u.kpc) == 4 @@ -261,7 +261,7 @@ def test_multiplying_ufunc(self): # no cosmo_factors inp = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) res = inp * inp - assert res.to_value(u.kpc ** 2) == 4 + assert res.to_value(u.kpc**2) == 4 assert res.comoving is False assert res.cosmo_factor is None # one is not cosmo_array @@ -283,12 +283,12 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): res = inp1 * inp2 - assert res.to_value(u.kpc ** 2) == 4 + assert res.to_value(u.kpc**2) == 4 assert res.comoving is False assert res.cosmo_factor is None # only two has cosmo_factor @@ -297,11 +297,11 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): res = inp1 * inp2 - assert res.to_value(u.kpc ** 2) == 4 + assert res.to_value(u.kpc**2) == 4 assert res.comoving is False assert res.cosmo_factor is None # cosmo_factors both present @@ -309,24 +309,24 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp * inp - assert res.to_value(u.kpc ** 2) == 4 + assert res.to_value(u.kpc**2) == 4 assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 2 + assert res.cosmo_factor == inp.cosmo_factor**2 def test_dividing_ufunc(self): inp = cosmo_array( [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp / inp assert res.to_value(u.dimensionless) == 1 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 0 + assert res.cosmo_factor == inp.cosmo_factor**0 def test_return_without_ufunc(self): # 1 argument @@ -334,7 +334,7 @@ def test_return_without_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.logical_not(inp) assert res == np.logical_not(1) @@ -344,7 +344,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.logaddexp(inp, inp) assert res == np.logaddexp(2, 2) @@ -354,13 +354,13 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Ufunc arguments have cosmo_factors that differ" @@ -372,7 +372,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): res = np.logaddexp(inp1, inp2) @@ -383,7 +383,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): @@ -396,7 +396,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): res = np.logaddexp(inp1, inp2) @@ -407,7 +407,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = u.unyt_array([2], u.kpc) with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): @@ -420,31 +420,31 @@ def test_sqrt_ufunc(self): [4], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.sqrt(inp) - assert res.to_value(u.kpc ** 0.5) == 2 # also ensures units ok + assert res.to_value(u.kpc**0.5) == 2 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 0.5 + assert res.cosmo_factor == inp.cosmo_factor**0.5 def test_square_ufunc(self): inp = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.square(inp) - assert res.to_value(u.kpc ** 2) == 4 # also ensures units ok + assert res.to_value(u.kpc**2) == 4 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 2 + assert res.cosmo_factor == inp.cosmo_factor**2 def test_cbrt_ufunc(self): inp = cosmo_array( [8], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.cbrt(inp) assert res.to_value(u.kpc ** (1.0 / 3.0)) == 2 # also ensures units ok @@ -456,12 +456,12 @@ def test_reciprocal_ufunc(self): [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.reciprocal(inp) - assert res.to_value(u.kpc ** -1) == 0.5 # also ensures units ok + assert res.to_value(u.kpc**-1) == 0.5 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** -1 + assert res.cosmo_factor == inp.cosmo_factor**-1 def test_passthrough_ufunc(self): # 1 argument @@ -469,7 +469,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.negative(inp) assert res.to_value(u.kpc) == -2 @@ -480,7 +480,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.copysign(inp, inp) assert res.to_value(u.kpc) == inp.to_value(u.kpc) @@ -491,13 +491,13 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises(ValueError): res = np.copysign(inp1, inp2) @@ -507,7 +507,7 @@ def test_arctan2_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.arctan2(inp, inp) assert res.to_value(u.dimensionless) == np.arctan2(2, 2) @@ -519,13 +519,13 @@ def test_comparison_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp1 < inp2 assert res.all() @@ -536,44 +536,54 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) out = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) np.abs(inp, out=out) assert out.to_value(u.kpc) == np.abs(inp.to_value(u.kpc)) assert out.comoving is False assert out.cosmo_factor == inp.cosmo_factor + inp = cosmo_array( + [1], + u.kpc, + comoving=False, + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + ) + # make sure we can also pass a non-cosmo type for out without crashing + out = np.array([np.nan]) + np.abs(inp, out=out) + assert out == np.abs(inp.to_value(u.kpc)) def test_reduce_multiply(self): inp = cosmo_array( [[1, 2], [3, 4]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) res = np.multiply.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc ** 2), np.array([3.0, 8.0])) + np.testing.assert_allclose(res.to_value(u.kpc**2), np.array([3.0, 8.0])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 2 + assert res.cosmo_factor == inp.cosmo_factor**2 def test_reduce_divide(self): inp = cosmo_array( [[1.0, 2.0], [1.0, 4.0], [1.0, 1.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) res = np.divide.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc ** -1), np.array([1.0, 0.5])) + np.testing.assert_allclose(res.to_value(u.kpc**-1), np.array([1.0, 0.5])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** -1 + assert res.cosmo_factor == inp.cosmo_factor**-1 def test_reduce_other(self): inp = cosmo_array( [[1.0, 2.0], [1.0, 2.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.add.reduce(inp, axis=0) np.testing.assert_allclose(res.to_value(u.kpc), np.array([2.0, 4.0])) @@ -586,7 +596,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res1, res2 = np.modf(inp) assert res1.to_value(u.kpc) == 0.5 @@ -600,7 +610,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res1, res2 = np.frexp(inp) assert res1 == 0.625 @@ -614,7 +624,7 @@ def test_multi_output_with_out_arg(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) out1 = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) out2 = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) @@ -631,7 +641,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = 0 res = inp1 > inp2 @@ -640,7 +650,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = 0.5 with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): @@ -650,20 +660,20 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [0, 0, 0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) assert (inp1 > inp2).all() inp1 = cosmo_array( [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = np.ones(3) * u.kpc with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): @@ -672,7 +682,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = np.zeros(3) * u.kpc assert (inp1 > inp2).all() @@ -680,13 +690,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( 1, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp1 == inp2 assert res.all() @@ -694,13 +704,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( 0, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp1 > inp2 assert res.all() From 0657ba931dd2ce029ab79942583da7447933d59a Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 5 Feb 2025 14:42:30 +0000 Subject: [PATCH 056/125] Cleanup while checking unyt usage. --- swiftsimio/_array_functions.py | 4 +- swiftsimio/snapshot_writer.py | 13 ++-- swiftsimio/visualisation/projection.py | 61 +++++++------------ swiftsimio/visualisation/ray_trace.py | 20 +++--- swiftsimio/visualisation/slice.py | 35 +++++------ .../smoothing_length/generate.py | 36 +++++------ swiftsimio/visualisation/volume_render.py | 39 ++++++------ tests/basic_test.py | 7 +-- 8 files changed, 92 insertions(+), 123 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index d4206548..b4826716 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -111,7 +111,9 @@ def _copy_cosmo_array_attributes(from_ca, to_ca): - if not isinstance(to_ca, objects.cosmo_array): + if not isinstance(to_ca, objects.cosmo_array) and isinstance( + from_ca, objects.cosmo_array + ): return to_ca if hasattr(from_ca, "cosmo_factor"): to_ca.cosmo_factor = from_ca.cosmo_factor diff --git a/swiftsimio/snapshot_writer.py b/swiftsimio/snapshot_writer.py index 2d9a4206..fbb6362d 100644 --- a/swiftsimio/snapshot_writer.py +++ b/swiftsimio/snapshot_writer.py @@ -14,6 +14,7 @@ from functools import reduce from swiftsimio import metadata +from swiftsimio.objects import cosmo_array from swiftsimio.metadata.cosmology.cosmology_fields import a_exponents @@ -51,7 +52,7 @@ class __SWIFTWriterParticleDataset(object): checks if all required datasets are empty. check_consistent(self) performs consistency checks on dataset - generate_smoothing_lengths(self, boxsize: Union[List[unyt.unyt_quantity], unyt.unyt_quantity], dimension: int) + generate_smoothing_lengths(self, boxsize: cosmo_array, dimension: int) automatically generates the smoothing lengths write_particle_group(self, file_handle: h5py.File, compress: bool) writes the particle group's required properties to file. @@ -164,7 +165,7 @@ def check_consistent(self) -> bool: def generate_smoothing_lengths( self, - boxsize: Union[List[unyt.unyt_quantity], unyt.unyt_quantity], + boxsize: cosmo_array, dimension: int, ): """ @@ -175,7 +176,7 @@ def generate_smoothing_lengths( Parameters ---------- - boxsize : unyt.unyt_quantity or list of unyt.unyt_quantity + boxsize : cosmo_array or cosmo_quantity size of SWIFT computational box dimension : int number of box dimensions @@ -271,7 +272,7 @@ def get_attributes(self, scale_factor: float) -> dict: # Find the scale factor associated quantities a_exp = a_exponents.get(name, 0) - a_factor = scale_factor ** a_exp + a_factor = scale_factor**a_exp attributes_dict[output_handle] = { "Conversion factor to CGS (not including cosmological corrections)": [ @@ -506,7 +507,7 @@ class SWIFTSnapshotWriter(object): def __init__( self, unit_system: Union[unyt.UnitSystem, str], - box_size: Union[list, unyt.unyt_quantity], + box_size: cosmo_array, dimension=3, compress=True, extra_header: Union[None, dict] = None, @@ -522,7 +523,7 @@ def __init__( ---------- unit_system : unyt.UnitSystem or str unit system for dataset - boxsize : list or unyt.unyt_quantity + boxsize : cosmo_array size of simulation box and associated units dimension : int, optional dimensions of simulation diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index fa0738b3..8d375257 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -3,53 +3,36 @@ """ from typing import Union -from math import sqrt, ceil from numpy import ( - float64, float32, - int32, - zeros, array, - arange, - ndarray, ones, - isclose, matmul, empty_like, logical_and, s_, + ceil, ) -from unyt import unyt_array, unyt_quantity, exceptions +from unyt import unyt_quantity, exceptions from swiftsimio import SWIFTDataset, cosmo_array from swiftsimio.reader import __SWIFTGroupDataset -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends import backends, backends_parallel -# Backwards compatability - -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_gamma, - kernel_constant, -) -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_single_precision as kernel, -) - scatter = backends["fast"] scatter_parallel = backends_parallel["fast"] def project_pixel_grid( data: __SWIFTGroupDataset, - boxsize: unyt_array, + boxsize: cosmo_array, resolution: int, project: Union[str, None] = "masses", - region: Union[None, unyt_array] = None, + region: Union[None, cosmo_array] = None, mask: Union[None, array] = None, rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, parallel: bool = False, backend: str = "fast", periodic: bool = True, @@ -68,7 +51,7 @@ def project_pixel_grid( data: __SWIFTGroupDataset The SWIFT dataset that you wish to visualise (get this from ``load``) - boxsize: unyt_array + boxsize: cosmo_array The box-size of the simulation. resolution: int @@ -81,7 +64,7 @@ def project_pixel_grid( always create it as ``data.gas.my_variable = data.gas.other_variable * data.gas.masses``. - region: unyt_array, optional + region: cosmo_array, optional Region, determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges) if it is not None. It should have a length of four or six, and take the form: @@ -116,7 +99,7 @@ def project_pixel_grid( Returns ------- - image: unyt_array + image: cosmo_array Projected image with units of project / length^2, of size ``res`` x ``res``. @@ -140,7 +123,7 @@ def project_pixel_grid( ) except AttributeError: raise exceptions.InvalidUnitOperation( - "Ensure that rotation_center is a unyt array with the same units as coordinates" + "Ensure that rotation_center is a cosmo_array with the same units as coordinates" ) number_of_particles = data.coordinates.shape[0] @@ -202,12 +185,12 @@ def project_pixel_grid( if data.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( - f"Physical smoothing length is not compatible with comoving coordinates!" + "Physical smoothing length is not compatible with comoving coordinates!" ) else: if not hsml.compatible_with_physical(): raise AttributeError( - f"Comoving smoothing length is not compatible with physical coordinates!" + "Comoving smoothing length is not compatible with physical coordinates!" ) except AttributeError: # No hsml present. If they are using the 'histogram' backend, we @@ -269,10 +252,10 @@ def project_gas_pixel_grid( data: SWIFTDataset, resolution: int, project: Union[str, None] = "masses", - region: Union[None, unyt_array] = None, + region: Union[None, cosmo_array] = None, mask: Union[None, array] = None, rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, parallel: bool = False, backend: str = "fast", periodic: bool = True, @@ -303,7 +286,7 @@ def project_gas_pixel_grid( always create it as ``data.gas.my_variable = data.gas.other_variable * data.gas.masses``. - region: unyt_array, optional + region: cosmo_array, optional Region, determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges) if it is not None. It should have a length of four or six, and take the form: @@ -374,9 +357,9 @@ def project_gas( data: SWIFTDataset, resolution: int, project: Union[str, None] = "masses", - region: Union[None, unyt_array] = None, + region: Union[None, cosmo_array] = None, mask: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, rotation_matrix: Union[None, array] = None, parallel: bool = False, backend: str = "fast", @@ -406,7 +389,7 @@ def project_gas( always create it as ``data.gas.my_variable = data.gas.other_variable * data.gas.masses``. - region: unyt_array, optional + region: cosmo_array, optional Region, determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges) if it is not None. It should have a length of four or six, and take the form: @@ -442,7 +425,7 @@ def project_gas( Returns ------- - image: unyt_array + image: cosmo_array Projected image with units of project / length^2, of size ``res`` x ``res``. @@ -474,23 +457,23 @@ def project_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 61960109..3ef77ff3 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -217,15 +217,15 @@ def core_panels_parallel( def panel_pixel_grid( data: __SWIFTGroupDataset, - boxsize: unyt.unyt_array, + boxsize: cosmo_array, resolution: int, panels: int, project: Union[str, None] = "masses", - region: Union[None, unyt.unyt_array] = None, + region: Union[None, cosmo_array] = None, mask: Union[None, np.array] = None, rotation_matrix: Union[None, np.array] = None, - rotation_center: Union[None, unyt.unyt_array] = None, -) -> unyt.unyt_array: + rotation_center: Union[None, cosmo_array] = None, +) -> cosmo_array: if rotation_center is not None: try: if rotation_center.units == data.coordinates.units: @@ -332,10 +332,10 @@ def panel_gas( resolution: int, panels: int, project: Union[str, None] = "masses", - region: Union[None, unyt.unyt_array] = None, + region: Union[None, cosmo_array] = None, mask: Union[None, np.array] = None, rotation_matrix: Union[None, np.array] = None, - rotation_center: Union[None, unyt.unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, ) -> cosmo_array: image = panel_pixel_grid( data=data.gas, @@ -353,23 +353,23 @@ def panel_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index 0240311e..bc5b67e0 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -4,17 +4,10 @@ from typing import Union, Optional from numpy import float32, array, ones, matmul -from unyt import unyt_array, unyt_quantity -from swiftsimio import SWIFTDataset, cosmo_array +from swiftsimio import SWIFTDataset, cosmo_array, cosmo_quantity from swiftsimio.visualisation.slice_backends import backends, backends_parallel from swiftsimio.visualisation.smoothing_length import backends_get_hsml -from swiftsimio.visualisation.slice_backends.sph import ( - kernel, - kernel_constant, - kernel_gamma, -) - slice_scatter = backends["sph"] slice_scatter_parallel = backends_parallel["sph"] @@ -22,12 +15,12 @@ def slice_gas_pixel_grid( data: SWIFTDataset, resolution: int, - z_slice: Optional[unyt_quantity] = None, + z_slice: Optional[cosmo_quantity] = None, project: Union[str, None] = "masses", parallel: bool = False, rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, + region: Union[None, cosmo_array] = None, backend: str = "sph", periodic: bool = True, ): @@ -43,7 +36,7 @@ def slice_gas_pixel_grid( resolution : int Specifies size of return array - z_slice : unyt_quantity + z_slice : cosmo_quantity Specifies the location along the z-axis where the slice is to be extracted, relative to the rotation center or the origin of the box if no rotation center is provided. If the perspective is rotated @@ -67,7 +60,7 @@ def slice_gas_pixel_grid( Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. - region : unyt_array, optional + region : cosmo_array, optional determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges) if it is not None. It should have a length of four, and take the form: @@ -205,12 +198,12 @@ def slice_gas_pixel_grid( def slice_gas( data: SWIFTDataset, resolution: int, - z_slice: Optional[unyt_quantity] = None, + z_slice: Optional[cosmo_quantity] = None, project: Union[str, None] = "masses", parallel: bool = False, rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, + region: Union[None, cosmo_array] = None, backend: str = "sph", periodic: bool = True, ): @@ -225,7 +218,7 @@ def slice_gas( resolution : int Specifies size of return array - z_slice : unyt_quantity + z_slice : cosmo_quantity Specifies the location along the z-axis where the slice is to be extracted, relative to the rotation center or the origin of the box if no rotation center is provided. If the perspective is rotated @@ -304,7 +297,7 @@ def slice_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 3) + units = 1.0 / (max_range**3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units( @@ -312,17 +305,17 @@ def slice_gas( ) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 3) + units = 1.0 / (max_range**3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) diff --git a/swiftsimio/visualisation/smoothing_length/generate.py b/swiftsimio/visualisation/smoothing_length/generate.py index dd271f60..cb2ae76b 100644 --- a/swiftsimio/visualisation/smoothing_length/generate.py +++ b/swiftsimio/visualisation/smoothing_length/generate.py @@ -3,18 +3,17 @@ that do not usually carry a smoothing length field (e.g. dark matter). """ -from typing import Union -from unyt import unyt_array from numpy import empty, float32 -from swiftsimio import SWIFTDataset, cosmo_array -from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma +from swiftsimio import cosmo_array from swiftsimio.optional_packages import KDTree, TREE_AVAILABLE +from swiftsimio._array_functions import _propagate_cosmo_array_attributes +@_propagate_cosmo_array_attributes # copies attrs of first arg to result, if present def generate_smoothing_lengths( - coordinates: Union[unyt_array, cosmo_array], - boxsize: Union[unyt_array, cosmo_array], + coordinates: cosmo_array, + boxsize: cosmo_array, kernel_gamma: float32, neighbours=32, speedup_fac=2, @@ -25,9 +24,9 @@ def generate_smoothing_lengths( Parameters ---------- - coordinates : unyt_array or cosmo_array + coordinates : cosmo_array a cosmo_array that gives the co-ordinates of all particles - boxsize : unyt_array or cosmo_array + boxsize : cosmo_array the size of the box (3D) kernel_gamma : float32 the kernel gamma of the kernel being used @@ -46,8 +45,8 @@ def generate_smoothing_lengths( Returns ------- - smoothing lengths : unyt_array - an unyt array of smoothing lengths. + smoothing lengths : cosmo_array + a cosmo_array of smoothing lengths. """ if not TREE_AVAILABLE: @@ -103,15 +102,8 @@ def generate_smoothing_lengths( smoothing_lengths[starting_index:ending_index] = d[:, -1] - if isinstance(coordinates, cosmo_array): - return cosmo_array( - smoothing_lengths * (hsml_correction_fac_speedup / kernel_gamma), - units=coordinates.units, - comoving=coordinates.comoving, - cosmo_factor=coordinates.cosmo_factor, - ) - else: - return unyt_array( - smoothing_lengths * (hsml_correction_fac_speedup / kernel_gamma), - units=coordinates.units, - ) + return type(coordinates)( + smoothing_lengths + * (hsml_correction_fac_speedup / kernel_gamma) + * coordinates.units + ) diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index fbb9e591..724ed87c 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -18,13 +18,12 @@ matmul, max as np_max, ) -from unyt import unyt_array from swiftsimio import SWIFTDataset, cosmo_array from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.optional_packages import plt -from .slice import kernel, kernel_gamma +from swiftsimio.visualisation.slice_backends.sph import kernel, kernel_gamma @jit(nopython=True, fastmath=True) @@ -588,8 +587,8 @@ def render_gas_voxel_grid( project: Union[str, None] = "masses", parallel: bool = False, rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, + region: Union[None, cosmo_array] = None, periodic: bool = True, ): """ @@ -618,11 +617,11 @@ def render_gas_voxel_grid( ``rotation_center``. In the default case, this provides a volume render viewed along the z axis. - rotation_center: np.array, optional + rotation_center: cosmo_array, optional Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. - region : unyt_array, optional + region : cosmo_array, optional determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges, and front and back edges) if it is not None. It should have a @@ -712,12 +711,12 @@ def render_gas_voxel_grid( if data.gas.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( - f"Physical smoothing length is not compatible with comoving coordinates!" + "Physical smoothing length is not compatible with comoving coordinates!" ) else: if not hsml.compatible_with_physical(): raise AttributeError( - f"Comoving smoothing length is not compatible with physical coordinates!" + "Comoving smoothing length is not compatible with physical coordinates!" ) if periodic: @@ -755,8 +754,8 @@ def render_gas( project: Union[str, None] = "masses", parallel: bool = False, rotation_matrix: Union[None, array] = None, - rotation_center: Union[None, unyt_array] = None, - region: Union[None, unyt_array] = None, + rotation_center: Union[None, cosmo_array] = None, + region: Union[None, cosmo_array] = None, periodic: bool = True, ): """ @@ -785,11 +784,11 @@ def render_gas( ``rotation_center``. In the default case, this provides a volume render viewed along the z axis. - rotation_center: np.array, optional + rotation_center: cosmo_array, optional Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. - region : unyt_array, optional + region : cosmo_array, optional determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges, and front and back edges) if it is not None. It should have a @@ -845,14 +844,14 @@ def render_gas( * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) @@ -896,8 +895,8 @@ def visualise_render( render : np.array The render to visualise. You should scale this appropriately before using this function (e.g. use a logarithmic transform!) - and pass in the 'value' array, not the original cosmo array or - unyt array. + and pass in the 'value' array, not the original cosmo_array or + unyt_array. centers : list[float] The centers of your rendering functions @@ -907,8 +906,8 @@ def visualise_render( will have the same width. cmap : str - The colormap to use for the rendering functions. - + The colormap to use for the rendering functions. + return_type : Literal["all", "lighten", "add"] The type of return. If "all", all images are returned. If "lighten", the maximum of all images is returned. If "add", the sum of all images @@ -974,11 +973,11 @@ def visualise_render_options( centers : list[float] The centers of your rendering functions - + widths : list[float] | float The widths of your rendering functions. If a single float, all functions will have the same width. - + cmap : str The colormap to use for the rendering functions. diff --git a/tests/basic_test.py b/tests/basic_test.py index e125da1f..24fb9295 100644 --- a/tests/basic_test.py +++ b/tests/basic_test.py @@ -4,7 +4,6 @@ from swiftsimio import load from swiftsimio import Writer -from swiftsimio.units import cosmo_units import unyt import numpy as np @@ -24,7 +23,7 @@ def test_write(): x = Writer("galactic", boxsize) # 32^3 particles. - n_p = 32 ** 3 + n_p = 32**3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) @@ -53,7 +52,7 @@ def test_load(): """ x = load("test.hdf5") - density = x.gas.internal_energy - coordinates = x.gas.coordinates + x.gas.internal_energy + x.gas.coordinates os.remove("test.hdf5") From 92894809cc81797740796a9bc5c9e890e6b3300c Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 5 Feb 2025 22:47:18 +0000 Subject: [PATCH 057/125] Refactor ca_cfs to just use cf directly. --- swiftsimio/_array_functions.py | 418 +++++++++++++------------------- swiftsimio/objects.py | 63 +++-- tests/test_cosmo_array_attrs.py | 120 ++++----- tests/test_data.py | 9 +- 4 files changed, 276 insertions(+), 334 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index b4826716..e941954b 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -181,44 +181,31 @@ def wrapped(obj, *args, **kwargs): return wrapped -def _sqrt_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor( - ca_cf, (False, None), power=0.5 - ) # ufunc sqrt not supported +def _sqrt_cosmo_factor(cf, **kwargs): + return _power_cosmo_factor(cf, None, power=0.5) # ufunc sqrt not supported 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) + cfs = args + if len(cfs) == 1: + return __multiply_cosmo_factor(cfs[0]) + retval = __multiply_cosmo_factor(cfs[0], cfs[1]) + for cf in cfs[2:]: + retval = __multiply_cosmo_factor(retval, cf) return retval -def __multiply_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 +def __multiply_cosmo_factor(cf1, cf2, **kwargs): if (cf1 is None) and (cf2 is None): # neither has cosmo_factor information: return None - elif not ca1 and ca2: - # one is not a cosmo_array, allow e.g. multiplication by constants: + elif (cf1 is None) and (cf2 is not None): + # first has no cosmo information, allow e.g. multiplication by constants: return cf2 - elif ca1 and not ca2: - # two is not a cosmo_array, allow e.g. multiplication by constants: + elif (cf1 is not None) and (cf2 is None): + # second has no cosmo information, allow e.g. multiplication by constants: return cf1 - elif (ca1 and ca2) and ((cf1 is None) or (cf2 is None)): - # both cosmo_array but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors ({cf1} and {cf2})," - f" discarding cosmo_factor in return value.", - RuntimeWarning, - ) - return None - elif (ca1 and ca2) and ((cf1 is not None) and (cf2 is not None)): + elif (cf1 is not None) and (cf2 is not None): # both cosmo_array and both with cosmo_factor: return cf1 * cf2 # cosmo_factor.__mul__ raises if scale factors differ else: @@ -226,64 +213,46 @@ def __multiply_cosmo_factor(ca_cf1, ca_cf2, **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) + cfs = args + if len(cfs) == 1: + return cfs[0] + retval = __preserve_cosmo_factor(cfs[0], cfs[1]) + for cf in cfs[2:]: + retval = __preserve_cosmo_factor(retval, 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: - # single argument, return promptly - return cf1 - elif (cf1 is None) and (cf2 is None): +def __preserve_cosmo_factor(cf1, cf2, **kwargs): + if (cf1 is None) and (cf2 is None): # neither has cosmo_factor information: return None - elif ca1 and not ca2: + elif (cf1 is not None) and (cf2 is None): # only one is cosmo_array + warnings.warn( + f"Mixing arguments with and without cosmo_factors, continuing assuming" + f" provided cosmo_factor ({cf1}) for all arguments.", + RuntimeWarning, + ) return cf1 - elif ca2 and not ca1: + elif (cf1 is None) and (cf2 is not None): # only one is cosmo_array - return cf2 - elif (ca1 and ca2) and (cf1 is None and cf2 is not None): - # both cosmo_array, but not both with cosmo_factor - # (both without shortcircuited above already): warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) return cf2 - elif (ca1 and ca2) and (cf1 is not None and cf2 is None): - # both cosmo_array, but not both with cosmo_factor - # (both without shortcircuited above already): - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf1}) for all arguments.", - RuntimeWarning, - ) - return cf1 - elif (ca1 and ca2) and (cf1 != cf2): - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." - ) - elif (ca1 and ca2) and (cf1 == cf2): + elif cf1 != cf2: + raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") + elif cf1 == cf2: return cf1 # or cf2, they're equal else: - # not dealing with cosmo_arrays at all - return None + raise RuntimeError("Unexpected state, please report this error on github.") -def _power_cosmo_factor(ca_cf1, ca_cf2, inputs=None, power=None): +def _power_cosmo_factor(cf1, cf2, inputs=None, power=None): if inputs is not None and power is not None: raise ValueError - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 power = inputs[1] if inputs else power if hasattr(power, "units"): if not power.units.is_dimensionless: @@ -291,158 +260,121 @@ def _power_cosmo_factor(ca_cf1, ca_cf2, inputs=None, power=None): elif power.units is not unyt.dimensionless: power = power.to_value(unyt.dimensionless) # else power.units is unyt.dimensionless, do nothing - if ca2 and cf2.a_factor != 1.0: + exp_afactor = getattr(cf2, "a_factor", None) + if (exp_afactor is not None) and (exp_afactor != 1.0): raise ValueError("Exponent has scaling with scale factor != 1.") if cf1 is None: return None return np.power(cf1, power) -def _square_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=2) +def _square_cosmo_factor(cf, **kwargs): + return _power_cosmo_factor(cf, None, power=2) -def _cbrt_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=1.0 / 3.0) +def _cbrt_cosmo_factor(cf, **kwargs): + return _power_cosmo_factor(cf, None, power=1.0 / 3.0) -def _divide_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 - return _multiply_cosmo_factor( - (ca1, cf1), (ca2, _reciprocal_cosmo_factor((ca2, cf2))) - ) +def _divide_cosmo_factor(cf1, cf2, **kwargs): + return _multiply_cosmo_factor(cf1, _reciprocal_cosmo_factor(cf2)) -def _reciprocal_cosmo_factor(ca_cf, **kwargs): - return _power_cosmo_factor(ca_cf, (False, None), power=-1) +def _reciprocal_cosmo_factor(cf, **kwargs): + return _power_cosmo_factor(cf, None, power=-1) -def _passthrough_cosmo_factor(ca_cf, ca_cf2=None, **kwargs): - ca, cf = ca_cf - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # no second argument, return promptly +def _passthrough_cosmo_factor(cf, cf2="__not_provided__", **kwargs): + if isinstance(cf2, str) and cf2 == "__not_provided__": return cf - elif (cf2 is not None) and cf != cf2: + if (cf2 is not None) and cf != cf2: # if both have cosmo_factor information and it differs this is an error - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." - ) + raise ValueError(f"Arguments have cosmo_factors that differ: {cf} and {cf2}.") else: # passthrough is for e.g. ufuncs with a second dimensionless argument, # so ok if cf2 is None and cf1 is not return cf -def _return_without_cosmo_factor(ca_cf, ca_cf2=None, inputs=None, zero_comparison=None): - ca, cf = ca_cf - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) - if ca_cf2 is None: - # no second argument - pass - elif ca and not ca2: +def _return_without_cosmo_factor( + cf, cf2="__not_provided__", zero_comparison=None, **kwargs +): + if isinstance(cf2, str) and cf2 == "__not_provided__": + return None + if (cf is not None) and (cf2 is None): # one is not a cosmo_array, warn on e.g. comparison to constants: if not zero_comparison: warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing" + f"Mixing arguments with and without cosmo_factors, continuing" f" assuming provided cosmo_factor ({cf}) for all arguments.", RuntimeWarning, ) - elif not ca and ca2: + elif (cf is None) and (cf2 is not None): # two is not a cosmo_array, warn on e.g. comparison to constants: if not zero_comparison: warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing" + f"Mixing arguments with and without cosmo_factors, continuing" f" assuming provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - elif (ca and ca2) and (cf is not None and cf2 is None): - # one has no cosmo_factor information, warn: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf}) for all arguments.", - RuntimeWarning, - ) - elif (ca and ca2) and (cf is None and cf2 is not None): - # two has no cosmo_factor information, warn: - warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" - f" provided cosmo_factor ({cf2}) for all arguments.", - RuntimeWarning, - ) elif (cf is not None) and (cf2 is not None) and (cf != cf2): # both have cosmo_factor, don't match: - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf} and {cf2}." - ) + raise ValueError(f"Arguments have cosmo_factors that differ: {cf} and {cf2}.") elif (cf is not None) and (cf2 is not None) and (cf == cf2): # both have cosmo_factor, and they match: pass else: - # not dealing with cosmo_arrays at all - pass + raise RuntimeError("Unexpected state, please report this error on github.") # return without cosmo_factor return None -def _arctan2_cosmo_factor(ca_cf1, ca_cf2, **kwargs): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 +def _arctan2_cosmo_factor(cf1, cf2, **kwargs): if (cf1 is None) and (cf2 is None): return None - if cf1 is None and cf2 is not None: + elif (cf1 is None) and (cf2 is not None): warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - if cf1 is not None and cf2 is None: + return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) + elif (cf1 is not None) and (cf2 is None): warnings.warn( - f"Mixing ufunc arguments with and without cosmo_factors, continuing assuming" + f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - if (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): - raise ValueError( - f"Ufunc arguments have cosmo_factors that differ: {cf1} and {cf2}." - ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): + raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") + elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + else: + raise RuntimeError("Unexpected state, please report this error on github.") -def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): - ca1, cf1 = ca_cf1 - ca2, cf2 = ca_cf2 if ca_cf2 is not None else (None, None) +def _comparison_cosmo_factor(cf1, cf2, inputs=None): try: iter(inputs[0]) except TypeError: - if ca1: - input1_iszero = not inputs[0].value and inputs[0] is not False - else: - input1_iszero = not inputs[0] and inputs[0] is not False + input1_iszero = ( + not getattr(inputs[0], "value", inputs[0]) and inputs[0] is not False + ) else: - if ca1: - input1_iszero = not inputs[0].value.any() - else: - input1_iszero = not inputs[0].any() + input1_iszero = not getattr(inputs[0], "value", inputs[0]).any() try: iter(inputs[1]) except IndexError: input2_iszero = None except TypeError: - if ca2: - input2_iszero = not inputs[1].value and inputs[1] is not False - else: - input2_iszero = not inputs[1] and inputs[1] is not False + input2_iszero = ( + not getattr(inputs[1], "value", inputs[1]) and inputs[1] is not False + ) else: - if ca2: - input2_iszero = not inputs[1].value.any() - else: - input2_iszero = not inputs[1].any() + input2_iszero = not getattr(inputs[1], "value", inputs[1]).any() zero_comparison = input1_iszero or input2_iszero - return _return_without_cosmo_factor( - ca_cf1, ca_cf2=ca_cf2, inputs=inputs, zero_comparison=zero_comparison - ) + return _return_without_cosmo_factor(cf1, cf2=cf2, zero_comparison=zero_comparison) def _prepare_array_func_args(*args, _default_cm=True, **kwargs): @@ -461,10 +393,7 @@ def _prepare_array_func_args(*args, _default_cm=True, **kwargs): # 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)) - for arg in args - ] + cfs = [getattr(arg, "cosmo_factor", None) for arg in args] comps = [ (hasattr(arg, "compression"), getattr(arg, "compression", None)) for arg in args ] @@ -472,10 +401,7 @@ def _prepare_array_func_args(*args, _default_cm=True, **kwargs): 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_cfs = {k: 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() @@ -530,8 +456,8 @@ def _prepare_array_func_args(*args, _default_cm=True, **kwargs): return dict( args=args, kwargs=kwargs, - ca_cfs=ca_cfs, - kw_ca_cfs=kw_ca_cfs, + cfs=cfs, + kw_cfs=kw_cfs, comoving=ret_cm, compression=ret_comp, ) @@ -567,7 +493,7 @@ def _default_unary_wrapper(unyt_func, cosmo_factor_wrapper): # by the cosmo_factor_wrapper def wrapper(*args, **kwargs): helper_result = _prepare_array_func_args(*args, **kwargs) - ret_cf = cosmo_factor_wrapper(helper_result["ca_cfs"][0]) + ret_cf = cosmo_factor_wrapper(helper_result["cfs"][0]) res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) if "out" in kwargs: return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) @@ -583,9 +509,7 @@ def _default_binary_wrapper(unyt_func, cosmo_factor_wrapper): # by the cosmo_factor_wrapper def wrapper(*args, **kwargs): helper_result = _prepare_array_func_args(*args, **kwargs) - ret_cf = cosmo_factor_wrapper( - helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] - ) + ret_cf = cosmo_factor_wrapper(helper_result["cfs"][0], helper_result["cfs"][1]) res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) if "out" in kwargs: return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) @@ -602,8 +526,8 @@ def _default_comparison_wrapper(unyt_func): def wrapper(*args, **kwargs): helper_result = _prepare_array_func_args(*args, **kwargs) ret_cf = _comparison_cosmo_factor( - helper_result["ca_cfs"][0], - helper_result["ca_cfs"][1], + helper_result["cfs"][0], + helper_result["cfs"][1], inputs=args[:2], ) res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) @@ -620,7 +544,7 @@ def _default_oplist_wrapper(unyt_func): def wrapper(*args, **kwargs): helper_result = _prepare_array_func_args(*args, **kwargs) helper_result_oplist = _prepare_array_func_args(*args[0]) - ret_cf = _preserve_cosmo_factor(helper_result_oplist["ca_cfs"][0]) + ret_cf = _preserve_cosmo_factor(helper_result_oplist["cfs"][0]) res = unyt_func( helper_result_oplist["args"], *helper_result["args"][1:], @@ -692,10 +616,10 @@ def histogram_bin_edges(a, bins=10, range=None, weights=None): 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"]) + ret_cf = _preserve_cosmo_factor(helper_result["kw_cfs"]["bins"]) else: # bins based on values in a - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) res = unyt_histogram_bin_edges( *helper_result["args"], **helper_result["kwargs"] ) @@ -722,17 +646,13 @@ 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_bins = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) - ret_cf_dens = _reciprocal_cosmo_factor(helper_result["ca_cfs"][0]) + ret_cf_bins = _preserve_cosmo_factor(helper_result["cfs"][0]) + ret_cf_dens = _reciprocal_cosmo_factor(helper_result["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_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) 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 + _multiply_cosmo_factor(ret_cf_w, ret_cf_dens) if density else ret_cf_w ) else: ret_cf_counts = ret_cf_dens if density else None @@ -772,8 +692,8 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): ) 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]) + ret_cf_x = _preserve_cosmo_factor(helper_result_x["cfs"][0]) + ret_cf_y = _preserve_cosmo_factor(helper_result_y["cfs"][0]) if (helper_result_x["kwargs"]["range"] is None) and ( helper_result_y["kwargs"]["range"] is None ): @@ -792,7 +712,7 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): weights=helper_result_w["kwargs"]["weights"], ) if weights is not None: - ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_cfs"]["weights"]) counts = _promote_unyt_to_cosmo(counts) if isinstance( counts, objects.cosmo_array @@ -812,8 +732,8 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): 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]) + ret_cf_x = _preserve_cosmo_factor(helper_result_x["cfs"][0]) + ret_cf_y = _preserve_cosmo_factor(helper_result_y["cfs"][0]) if (helper_result["kwargs"]["xrange"] is None) and ( helper_result["kwargs"]["yrange"] is None ): @@ -832,18 +752,18 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): weights=helper_result["kwargs"]["weights"], ) ret_cf_xy = _multiply_cosmo_factor( - helper_result["ca_cfs"][0], - helper_result["ca_cfs"][1], + helper_result["cfs"][0], + helper_result["cfs"][1], ) if weights is not None: - ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_cfs"]["weights"]) - inv_ret_cf_xy = _reciprocal_cosmo_factor((ret_cf_xy is not None, ret_cf_xy)) + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) + inv_ret_cf_xy = _reciprocal_cosmo_factor(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), + ret_cf_w, + inv_ret_cf_xy, ) else: - ret_cf_counts = _reciprocal_cosmo_factor((ret_cf_xy is not None, ret_cf_xy)) + ret_cf_counts = _reciprocal_cosmo_factor(ret_cf_xy) counts = _promote_unyt_to_cosmo(counts) if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity counts.comoving = helper_result["comoving"] @@ -881,7 +801,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): if not density: helper_result_w = _prepare_array_func_args(weights=weights) ret_cfs = [ - _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + _preserve_cosmo_factor(helper_result["cfs"][0]) for helper_result in helper_results ] if all( @@ -903,7 +823,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): weights=helper_result_w["kwargs"]["weights"], ) if weights is not None: - ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_cfs"]["weights"]) counts = _promote_unyt_to_cosmo(counts) if isinstance(counts, objects.cosmo_array): counts.comoving = helper_result_w["comoving"] @@ -918,7 +838,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): range=range, weights=weights, ) - ret_cfs = D * [_preserve_cosmo_factor(helper_result["ca_cfs"][0])] + ret_cfs = D * [_preserve_cosmo_factor(helper_result["cfs"][0])] counts, bins = unyt_histogramdd( helper_result["args"], bins=helper_result["kwargs"]["bins"], @@ -926,23 +846,19 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): 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]) + if len(helper_result["cfs"]) == 1: + ret_cf_sample = _preserve_cosmo_factor(helper_result["cfs"][0]) else: - ret_cf_sample = _multiply_cosmo_factor(*helper_result["ca_cfs"]) + ret_cf_sample = _multiply_cosmo_factor(*helper_result["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( - (ret_cf_sample is not None, ret_cf_sample) - ) + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) + inv_ret_cf_sample = _reciprocal_cosmo_factor(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), + ret_cf_w, + inv_ret_cf_sample, ) else: - ret_cf_counts = _reciprocal_cosmo_factor( - (ret_cf_sample is not None, ret_cf_sample) - ) + ret_cf_counts = _reciprocal_cosmo_factor(ret_cf_sample) counts = _promote_unyt_to_cosmo(counts) if isinstance(counts, objects.cosmo_array): # also recognizes cosmo_quantity counts.comoving = helper_result["comoving"] @@ -999,7 +915,7 @@ def _prepare_array_block_args(lst, recursing=False): 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] + cfs = [hr["cfs"] for hr in helper_results] convert_to_cm = False if all(cms): ret_cm = True @@ -1020,9 +936,9 @@ def _prepare_array_block_args(lst, recursing=False): 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: + ret_cf = cfs[0] + for cf in cfs[1:]: + if cf != ret_cf: raise ValueError("Mixed cosmo_factor values in input.") if convert_to_cm: ret_lst = _recursive_to_comoving(lst) @@ -1116,9 +1032,7 @@ def linspace( axis=axis, device=device, ) - ret_cf = _preserve_cosmo_factor( - helper_result["ca_cfs"][0], helper_result["ca_cfs"][1] - ) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["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) @@ -1132,7 +1046,7 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0): 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"]) + ret_cf = _preserve_cosmo_factor(helper_result["kw_cfs"]["base"]) res = unyt_logspace(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1146,7 +1060,9 @@ def logspace(start, stop, num=50, endpoint=True, base=10.0, dtype=None, axis=0): def copyto(dst, src, casting="same_kind", where=True): 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]) + if isinstance(src, objects.cosmo_array) and isinstance(dst, objects.cosmo_array): + # if we're copyting across two + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][1]) # must pass dst directly here because it's modified in-place if isinstance(src, objects.cosmo_array): comoving = getattr(dst, "comoving", None) @@ -1179,8 +1095,8 @@ def prod( ) res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) ret_cf = _power_cosmo_factor( - helper_result["ca_cfs"][0], - (False, None), + helper_result["cfs"][0], + None, power=a.size // res.size, ) return _return_helper(res, helper_result, ret_cf, out=out) @@ -1205,8 +1121,8 @@ def linalg_det(a): helper_result = _prepare_array_func_args(a) ret_cf = _power_cosmo_factor( - helper_result["ca_cfs"][0], - (False, None), + helper_result["cfs"][0], + None, power=a.shape[0], ) res = unyt_linalg_det(*helper_result["args"], **helper_result["kwargs"]) @@ -1225,7 +1141,7 @@ def pad(array, pad_width, mode="constant", **kwargs): 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]) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) res = unyt_pad(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1235,7 +1151,7 @@ def choose(a, choices, out=None, mode="raise"): 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"]) + ret_cf = _preserve_cosmo_factor(*helper_result_choices["cfs"]) res = unyt_choose(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf, out=out) @@ -1244,9 +1160,7 @@ def choose(a, choices, out=None, mode="raise"): def insert(arr, obj, values, axis=None): 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] - ) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) res = unyt_insert(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1255,11 +1169,9 @@ def insert(arr, obj, values, axis=None): def linalg_lstsq(a, b, rcond=None): 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]) + ret_cf = _divide_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][0]) + resid_cf = _power_cosmo_factor(helper_result["cfs"][1], None, power=2) + sing_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) ress = unyt_linalg_lstsq(*helper_result["args"], **helper_result["kwargs"]) return ( _return_helper(ress[0], helper_result, ret_cf), @@ -1273,9 +1185,7 @@ def linalg_lstsq(a, b, rcond=None): def linalg_solve(a, b): helper_result = _prepare_array_func_args(a, b) - ret_cf = _divide_cosmo_factor( - helper_result["ca_cfs"][1], helper_result["ca_cfs"][0] - ) + ret_cf = _divide_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][0]) res = unyt_linalg_solve(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1284,9 +1194,7 @@ def linalg_solve(a, b): def linalg_tensorsolve(a, b, axes=None): 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["cfs"][1], helper_result["cfs"][0]) res = unyt_linalg_tensorsolve(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1295,7 +1203,7 @@ def linalg_tensorsolve(a, b, axes=None): def linalg_eig(a): helper_result = _prepare_array_func_args(a) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) ress = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) return ( _return_helper(ress[0], helper_result, ret_cf), @@ -1307,7 +1215,7 @@ def linalg_eig(a): def linalg_eigh(a, UPLO="L"): helper_result = _prepare_array_func_args(a, UPLO=UPLO) - ret_cf = _preserve_cosmo_factor(helper_result["ca_cfs"][0]) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) ress = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) return ( _return_helper(ress[0], helper_result, ret_cf), @@ -1379,7 +1287,7 @@ def apply_over_axes(func, a, axes): def fill_diagonal(a, val, wrap=False): helper_result = _prepare_array_func_args(a, val, wrap=wrap) - _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][1]) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][1]) # must pass a directly here because it's modified in-place comoving = getattr(a, "comoving", None) if comoving: @@ -1396,7 +1304,7 @@ def fill_diagonal(a, val, wrap=False): def place(arr, mask, vals): helper_result = _prepare_array_func_args(arr, mask, vals) - _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) # must pass arr directly here because it's modified in-place if isinstance(vals, objects.cosmo_array): comoving = getattr(arr, "comoving", None) @@ -1411,7 +1319,7 @@ def place(arr, mask, vals): def put(a, ind, v, mode="raise"): helper_result = _prepare_array_func_args(a, ind, v, mode=mode) - _preserve_cosmo_factor(helper_result["ca_cfs"][0], helper_result["ca_cfs"][2]) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) # must pass arr directly here because it's modified in-place if isinstance(v, objects.cosmo_array): comoving = getattr(a, "comoving", None) @@ -1426,7 +1334,7 @@ def put(a, ind, v, mode="raise"): 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]) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) # must pass arr directly here because it's modified in-place if isinstance(values, objects.cosmo_array): comoving = getattr(arr, "comoving", None) @@ -1441,7 +1349,7 @@ def put_along_axis(arr, indices, values, axis): 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]) + _preserve_cosmo_factor(helper_result["cfs"][0], helper_result["cfs"][2]) # must pass arr directly here because it's modified in-place if isinstance(values, objects.cosmo_array): comoving = getattr(a, "comoving", None) @@ -1462,7 +1370,7 @@ def select(condlist, choicelist, default=0): 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"]) + ret_cf = _preserve_cosmo_factor(*helper_result_choicelist["cfs"]) res = unyt_select( helper_result["args"][0], helper_result_choicelist["args"], @@ -1506,9 +1414,9 @@ def clip( **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"], + helper_result["cfs"][0], + helper_result["kw_cfs"]["a_min"], + helper_result["kw_cfs"]["a_max"], ) res = unyt_clip( helper_result["args"][0], @@ -1525,14 +1433,12 @@ def where(condition, *args): 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]) + ret_cf = _return_without_cosmo_factor(helper_result["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] - ) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][1], helper_result["cfs"][2]) res = unyt_where(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1562,7 +1468,7 @@ def einsum( optimize=optimize, ) helper_result_operands = _prepare_array_func_args(*operands) - ret_cf = _preserve_cosmo_factor(*helper_result_operands["ca_cfs"]) + ret_cf = _preserve_cosmo_factor(*helper_result_operands["cfs"]) res = unyt_einsum( helper_result["args"][0], *helper_result_operands["args"], @@ -1587,9 +1493,9 @@ def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586): 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"], + helper_result["cfs"][0], + helper_result["kw_cfs"]["discont"], + helper_result["kw_cfs"]["period"], ) res = unyt_unwrap(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1601,7 +1507,7 @@ def interp(x, xp, fp, left=None, right=None, period=None): 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]) + ret_cf = _preserve_cosmo_factor(helper_result["cfs"][2]) res = unyt_interp(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1637,11 +1543,11 @@ def trapezoid(y, x=None, dx=1.0, axis=-1): 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"] + helper_result["cfs"][0], helper_result["kw_cfs"]["dx"] ) else: ret_cf = _multiply_cosmo_factor( - helper_result["ca_cfs"][0], helper_result["kw_ca_cfs"]["x"] + helper_result["cfs"][0], helper_result["kw_cfs"]["x"] ) res = unyt_trapezoid(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 48ef9341..35ea6476 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -233,6 +233,8 @@ def a_factor(self): float the a-factor for given unit """ + if (self.expr is None) or (self.scale_factor is None): + return None return float(self.expr.subs(a, self.scale_factor)) @property @@ -253,6 +255,8 @@ def redshift(self): ..math:: z = \\frac{1}{a} - 1, where :math: `a` is the scale factor """ + if self.scale_factor is None: + return None return (1.0 / self.scale_factor) - 1.0 def __add__(self, b): @@ -292,6 +296,17 @@ def __mul__(self, b): f"{self.scale_factor} and {b.scale_factor}" ) + if ((self.expr is None) and (b.expr is not None)) or ( + (self.expr is not None) and (b.expr is None) + ): + raise InvalidScaleFactor( + "Attempting to multiply an initialized cosmo_factor with an " + f"uninitialized cosmo_factor {self} and {b}." + ) + if (self.expr is None) and (b.expr is None): + # let's be permissive and allow two uninitialized cosmo_factors through + return cosmo_factor(expr=None, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr * b.expr, scale_factor=self.scale_factor) def __truediv__(self, b): @@ -301,6 +316,17 @@ def __truediv__(self, b): f"{self.scale_factor} and {b.scale_factor}" ) + if ((self.expr is None) and (b.expr is not None)) or ( + (self.expr is not None) and (b.expr is None) + ): + raise InvalidScaleFactor( + "Attempting to divide an initialized cosmo_factor with an " + f"uninitialized cosmo_factor {self} and {b}." + ) + if (self.expr is None) and (b.expr is None): + # let's be permissive and allow two uninitialized cosmo_factors through + return cosmo_factor(expr=None, scale_factor=self.scale_factor) + return cosmo_factor(expr=self.expr / b.expr, scale_factor=self.scale_factor) def __radd__(self, b): @@ -316,6 +342,8 @@ def __rtruediv__(self, b): return b.__truediv__(self) def __pow__(self, p): + if self.expr is None: + return cosmo_factor(expr=None, scale_factor=self.scale_factor) return cosmo_factor(expr=self.expr**p, scale_factor=self.scale_factor) def __lt__(self, b): @@ -349,7 +377,10 @@ def __repr__(self): str string to print exponent and current scale factor """ - return self.__str__() + return f"cosmo_factor(expr={self.expr}, scale_factor={self.scale_factor})" + + +NULL_CF = cosmo_factor(None, None) class cosmo_array(unyt_array): @@ -484,7 +515,7 @@ def __new__( bypass_validation=False, input_units=None, name=None, - cosmo_factor=None, + cosmo_factor=cosmo_factor(None, None), comoving=None, valid_transform=True, compression=None, @@ -539,8 +570,8 @@ def __new__( 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), + cosmo_factor, + getattr(input_array, "cosmo_factor", NULL_CF), ) if not valid_transform: input_array.convert_to_physical() @@ -558,7 +589,7 @@ def __new__( 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"] + cosmo_factor, *helper_result["cfs"] ) if not valid_transform: input_array.convert_to_physical() @@ -597,7 +628,7 @@ def __array_finalize__(self, obj): super().__array_finalize__(obj) if obj is None: return - self.cosmo_factor = getattr(obj, "cosmo_factor", None) + self.cosmo_factor = getattr(obj, "cosmo_factor", NULL_CF) self.comoving = getattr(obj, "comoving", None) self.compression = getattr(obj, "compression", None) self.valid_transform = getattr(obj, "valid_transform", True) @@ -752,7 +783,7 @@ def from_astropy( arr, unit_registry=None, comoving=None, - cosmo_factor=None, + cosmo_factor=cosmo_factor(None, None), compression=None, valid_transform=True, ): @@ -798,7 +829,7 @@ def from_pint( arr, unit_registry=None, comoving=None, - cosmo_factor=None, + cosmo_factor=cosmo_factor(None, None), compression=None, valid_transform=True, ): @@ -846,7 +877,7 @@ def from_pint( def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): helper_result = _prepare_array_func_args(*inputs, **kwargs) - ca_cfs = helper_result["ca_cfs"] + 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 @@ -854,16 +885,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( - ca_cfs[0], - (False, None), + cfs[0], + None, power=power_map(inputs[0].shape[kwargs["axis"]]), ) else: ret_cf = _power_cosmo_factor( - ca_cfs[0], (False, None), power=power_map(inputs[0].size) + cfs[0], None, power=power_map(inputs[0].size) ) else: - ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*ca_cfs, inputs=inputs) + ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs) ret = _ensure_cosmo_array_or_quantity(super().__array_ufunc__)( ufunc, method, *inputs, **kwargs @@ -974,7 +1005,7 @@ def __new__( dtype=None, bypass_validation=False, name=None, - cosmo_factor=None, + cosmo_factor=cosmo_factor(None, None), comoving=None, valid_transform=True, compression=None, @@ -1025,8 +1056,8 @@ def __new__( units = getattr(input_scalar, "units", None) if units is None else units name = getattr(input_scalar, "name", None) if name is None else name cosmo_factor = ( - getattr(input_scalar, "cosmo_factor", None) - if cosmo_factor is None + getattr(input_scalar, "cosmo_factor", NULL_CF) + if cosmo_factor == NULL_CF else cosmo_factor ) comoving = ( diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_attrs.py index c67bab4c..7b94443b 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_attrs.py @@ -6,7 +6,13 @@ import pytest import numpy as np import unyt as u -from swiftsimio.objects import cosmo_array, cosmo_factor, a, multiple_output_operators +from swiftsimio.objects import ( + cosmo_array, + cosmo_factor, + a, + multiple_output_operators, + InvalidScaleFactor, +) class TestCopyFuncs: @@ -183,25 +189,27 @@ def test_preserving_ufunc(self): assert res.comoving is False assert res.cosmo_factor == inp.cosmo_factor # 2 argument, no cosmo_factors - inp = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp = cosmo_array([2], u.kpc, comoving=False) res = inp + inp assert res.to_value(u.kpc) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # 2 argument, one is not cosmo_array inp1 = u.unyt_array([2], u.kpc) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - res = inp1 + inp2 + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.warns(RuntimeWarning, match="Mixing arguments"): + res = inp1 + inp2 assert res.to_value(u.kpc) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # 2 argument, two is not cosmo_array - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = u.unyt_array([2], u.kpc) - res = inp1 + inp2 + with pytest.warns(RuntimeWarning, match="Mixing arguments"): + res = inp1 + inp2 assert res.to_value(u.kpc) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # 2 argument, only one has cosmo_factor inp1 = cosmo_array( [2], @@ -209,25 +217,23 @@ def test_preserving_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 + inp2 - assert res.to_value(u.kpc) == 4 - assert res.comoving is False - assert res.cosmo_factor == inp1.cosmo_factor + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + inp1 + inp2 # 2 argument, only two has cosmo_factor - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = cosmo_array( [2], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 + inp2 - assert res.to_value(u.kpc) == 4 - assert res.comoving is False - assert res.cosmo_factor == inp2.cosmo_factor + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + inp1 + inp2 # 2 argument, mismatched cosmo_factors inp1 = cosmo_array( [2], @@ -242,9 +248,9 @@ def test_preserving_ufunc(self): cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( - ValueError, match="Ufunc arguments have cosmo_factors that differ" + ValueError, match="Arguments have cosmo_factors that differ" ): - res = inp1 + inp2 + inp1 + inp2 # 2 argument, matched cosmo_factors inp = cosmo_array( [2], @@ -259,20 +265,20 @@ def test_preserving_ufunc(self): def test_multiplying_ufunc(self): # no cosmo_factors - inp = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp = cosmo_array([2], u.kpc, comoving=False) res = inp * inp assert res.to_value(u.kpc**2) == 4 assert res.comoving is False - assert res.cosmo_factor is None + assert res.cosmo_factor == cosmo_factor(None, None) # one is not cosmo_array inp1 = 2 - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp2 = cosmo_array([2], u.kpc, comoving=False) res = inp1 * inp2 assert res.to_value(u.kpc) == 4 assert res.comoving is False assert res.cosmo_factor == inp2.cosmo_factor # two is not cosmo_array - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = 2 res = inp1 * inp2 assert res.to_value(u.kpc) == 4 @@ -285,25 +291,19 @@ def test_multiplying_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 * inp2 - assert res.to_value(u.kpc**2) == 4 - assert res.comoving is False - assert res.cosmo_factor is None + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): + inp1 * inp2 # only two has cosmo_factor - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = cosmo_array( [2], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = inp1 * inp2 - assert res.to_value(u.kpc**2) == 4 - assert res.comoving is False - assert res.cosmo_factor is None + with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): + inp1 * inp2 # cosmo_factors both present inp = cosmo_array( [2], @@ -363,21 +363,21 @@ def test_return_without_ufunc(self): cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( - ValueError, match="Ufunc arguments have cosmo_factors that differ" + ValueError, match="Arguments have cosmo_factors that differ" ): - res = np.logaddexp(inp1, inp2) + np.logaddexp(inp1, inp2) # 2 arguments, one missing comso_factor - inp1 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) + inp1 = cosmo_array([2], u.kpc, comoving=False) inp2 = cosmo_array( [2], u.kpc, comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = np.logaddexp(inp1, inp2) - assert res == np.logaddexp(2, 2) - assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + np.logaddexp(inp1, inp2) # 2 arguments, two missing comso_factor inp1 = cosmo_array( [2], @@ -385,9 +385,11 @@ def test_return_without_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - inp2 = cosmo_array([2], u.kpc, comoving=False, cosmo_factor=None) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): - res = np.logaddexp(inp1, inp2) + inp2 = cosmo_array([2], u.kpc, comoving=False) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + np.logaddexp(inp1, inp2) assert res == np.logaddexp(2, 2) assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) # 2 arguments, one not cosmo_array @@ -398,7 +400,7 @@ def test_return_without_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) assert res == np.logaddexp(2, 2) assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) @@ -410,7 +412,7 @@ def test_return_without_ufunc(self): cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = u.unyt_array([2], u.kpc) - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) assert res == np.logaddexp(2, 2) assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) @@ -499,8 +501,10 @@ def test_passthrough_ufunc(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) - with pytest.raises(ValueError): - res = np.copysign(inp1, inp2) + with pytest.raises( + ValueError, match="Arguments have cosmo_factors that differ" + ): + np.copysign(inp1, inp2) def test_arctan2_ufunc(self): inp = cosmo_array( @@ -538,7 +542,7 @@ def test_out_arg(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - out = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) + out = cosmo_array([np.nan], u.dimensionless, comoving=True) np.abs(inp, out=out) assert out.to_value(u.kpc) == np.abs(inp.to_value(u.kpc)) assert out.comoving is False @@ -626,8 +630,8 @@ def test_multi_output_with_out_arg(self): comoving=False, cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) - out1 = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) - out2 = cosmo_array([np.nan], u.dimensionless, comoving=True, cosmo_factor=None) + out1 = cosmo_array([np.nan], u.dimensionless, comoving=True) + out2 = cosmo_array([np.nan], u.dimensionless, comoving=True) np.modf(inp, out=(out1, out2)) assert out1.to_value(u.kpc) == 0.5 assert out2.to_value(u.kpc) == 2.0 @@ -653,7 +657,7 @@ def test_comparison_with_zero(self): cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = 0.5 - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = inp1 > inp2 assert res.all() inp1 = cosmo_array( @@ -676,7 +680,7 @@ def test_comparison_with_zero(self): cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = np.ones(3) * u.kpc - with pytest.warns(RuntimeWarning, match="Mixing ufunc arguments"): + with pytest.warns(RuntimeWarning, match="Mixing arguments"): assert (inp1 == inp2).all() inp1 = cosmo_array( [1, 1, 1], diff --git a/tests/test_data.py b/tests/test_data.py index e6d61b08..25de5eaa 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -165,7 +165,7 @@ def test_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), ) start_offset = offsets @@ -212,7 +212,7 @@ def test_dithered_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), ) offsets = mask_region.offsets["dark_matter"] counts = mask_region.counts["dark_matter"] @@ -256,7 +256,7 @@ def test_reading_select_region_metadata(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.2, boxsize * 0.8]).T @@ -267,6 +267,7 @@ def test_reading_select_region_metadata(filename): selected_coordinates = selected_data.gas.coordinates # Now need to repeat teh selection by hand: + subset_mask = logical_and.reduce( [ logical_and(x > y_lower, x < y_upper) @@ -307,7 +308,7 @@ def test_reading_select_region_metadata_not_spatial_only(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.26, boxsize * 0.74]).T From 303668d1453799ec4a5cec9b55b2defe1210c5a4 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 6 Feb 2025 17:55:09 +0000 Subject: [PATCH 058/125] Troubleshooting after refactor. --- swiftsimio/_array_functions.py | 85 ++++------------ swiftsimio/masks.py | 4 +- swiftsimio/objects.py | 84 ++++++++++------ swiftsimio/visualisation/projection.py | 17 ++-- tests/helper.py | 6 +- tests/test_cosmo_array.py | 102 +++++++++++-------- tests/test_cosmo_array_attrs.py | 134 ++++++++++++------------- 7 files changed, 214 insertions(+), 218 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index e941954b..f2cbbf96 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -2,6 +2,7 @@ import numpy as np import unyt from unyt import unyt_quantity, unyt_array +from unyt.array import _iterable from swiftsimio import objects from unyt._array_functions import ( dot as unyt_dot, @@ -284,9 +285,7 @@ def _reciprocal_cosmo_factor(cf, **kwargs): return _power_cosmo_factor(cf, None, power=-1) -def _passthrough_cosmo_factor(cf, cf2="__not_provided__", **kwargs): - if isinstance(cf2, str) and cf2 == "__not_provided__": - return cf +def _passthrough_cosmo_factor(cf, cf2=None, **kwargs): if (cf2 is not None) and cf != cf2: # if both have cosmo_factor information and it differs this is an error raise ValueError(f"Arguments have cosmo_factors that differ: {cf} and {cf2}.") @@ -296,10 +295,8 @@ def _passthrough_cosmo_factor(cf, cf2="__not_provided__", **kwargs): return cf -def _return_without_cosmo_factor( - cf, cf2="__not_provided__", zero_comparison=None, **kwargs -): - if isinstance(cf2, str) and cf2 == "__not_provided__": +def _return_without_cosmo_factor(cf, cf2=np._NoValue, zero_comparison=None, **kwargs): + if cf2 is np._NoValue: return None if (cf is not None) and (cf2 is None): # one is not a cosmo_array, warn on e.g. comparison to constants: @@ -338,18 +335,18 @@ def _arctan2_cosmo_factor(cf1, cf2, **kwargs): f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) else: raise RuntimeError("Unexpected state, please report this error on github.") @@ -390,7 +387,7 @@ def _prepare_array_func_args(*args, _default_cm=True, **kwargs): # 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 pass the first argument (of np.concatenate - an inp) 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] cfs = [getattr(arg, "cosmo_factor", None) for arg in args] @@ -526,9 +523,7 @@ def _default_comparison_wrapper(unyt_func): def wrapper(*args, **kwargs): helper_result = _prepare_array_func_args(*args, **kwargs) ret_cf = _comparison_cosmo_factor( - helper_result["cfs"][0], - helper_result["cfs"][1], - inputs=args[:2], + helper_result["cfs"][0], helper_result["cfs"][1], inputs=args[:2] ) res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -680,16 +675,8 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): 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, - ) + 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["cfs"][0]) @@ -752,16 +739,12 @@ def histogram2d(x, y, bins=10, range=None, density=None, weights=None): weights=helper_result["kwargs"]["weights"], ) ret_cf_xy = _multiply_cosmo_factor( - helper_result["cfs"][0], - helper_result["cfs"][1], + helper_result["cfs"][0], helper_result["cfs"][1] ) if weights is not None: ret_cf_w = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"]) inv_ret_cf_xy = _reciprocal_cosmo_factor(ret_cf_xy) - ret_cf_counts = _multiply_cosmo_factor( - ret_cf_w, - inv_ret_cf_xy, - ) + ret_cf_counts = _multiply_cosmo_factor(ret_cf_w, inv_ret_cf_xy) else: ret_cf_counts = _reciprocal_cosmo_factor(ret_cf_xy) counts = _promote_unyt_to_cosmo(counts) @@ -791,11 +774,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): # bins is an integer bins = D * [bins] helper_results = [ - _prepare_array_func_args( - s, - bins=b, - range=r, - ) + _prepare_array_func_args(s, bins=b, range=r) for s, b, r in zip(sample, bins, ranges) ] if not density: @@ -833,10 +812,7 @@ def histogramdd(sample, bins=10, range=None, density=None, weights=None): # 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, + *sample, bins=bins, range=range, weights=weights ) ret_cfs = D * [_preserve_cosmo_factor(helper_result["cfs"][0])] counts, bins = unyt_histogramdd( @@ -853,10 +829,7 @@ 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_cfs"]["weights"]) inv_ret_cf_sample = _reciprocal_cosmo_factor(ret_cf_sample) - ret_cf_counts = _multiply_cosmo_factor( - ret_cf_w, - inv_ret_cf_sample, - ) + ret_cf_counts = _multiply_cosmo_factor(ret_cf_w, inv_ret_cf_sample) else: ret_cf_counts = _reciprocal_cosmo_factor(ret_cf_sample) counts = _promote_unyt_to_cosmo(counts) @@ -1095,9 +1068,7 @@ def prod( ) res = unyt_prod(*helper_result["args"], **helper_result["kwargs"]) ret_cf = _power_cosmo_factor( - helper_result["cfs"][0], - None, - power=a.size // res.size, + helper_result["cfs"][0], None, power=a.size // res.size ) return _return_helper(res, helper_result, ret_cf, out=out) @@ -1120,11 +1091,7 @@ def prod( def linalg_det(a): helper_result = _prepare_array_func_args(a) - ret_cf = _power_cosmo_factor( - helper_result["cfs"][0], - None, - power=a.shape[0], - ) + ret_cf = _power_cosmo_factor(helper_result["cfs"][0], None, power=a.shape[0]) res = unyt_linalg_det(*helper_result["args"], **helper_result["kwargs"]) return _return_helper(res, helper_result, ret_cf) @@ -1205,10 +1172,7 @@ def linalg_eig(a): helper_result = _prepare_array_func_args(a) ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) ress = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) - return ( - _return_helper(ress[0], helper_result, ret_cf), - ress[1], - ) + return (_return_helper(ress[0], helper_result, ret_cf), ress[1]) @implements(np.linalg.eigh) @@ -1217,10 +1181,7 @@ def linalg_eigh(a, UPLO="L"): helper_result = _prepare_array_func_args(a, UPLO=UPLO) ret_cf = _preserve_cosmo_factor(helper_result["cfs"][0]) ress = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) - return ( - _return_helper(ress[0], helper_result, ret_cf), - ress[1], - ) + return (_return_helper(ress[0], helper_result, ret_cf), ress[1]) implements(np.linalg.eigvals)( @@ -1407,11 +1368,7 @@ def clip( # 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, + a, a_min=a_min, a_max=a_max, out=out, **kwargs ) ret_cf = _preserve_cosmo_factor( helper_result["cfs"][0], diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index 434fe9d8..189cb1d0 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -237,7 +237,7 @@ def _unpack_cell_metadata(self): centers_handle[:][sort], units=self.units.length, comoving=True, - cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 1, self.metadata.scale_factor), ) # Note that we cannot assume that these are cubic, unfortunately. @@ -245,7 +245,7 @@ def _unpack_cell_metadata(self): metadata_handle.attrs["size"], units=self.units.length, comoving=True, - cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 1, self.metadata.scale_factor), ) return diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 35ea6476..f5b1ccf7 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -344,7 +344,7 @@ def __rtruediv__(self, b): def __pow__(self, p): if self.expr is None: return cosmo_factor(expr=None, scale_factor=self.scale_factor) - 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 @@ -515,7 +515,7 @@ def __new__( bypass_validation=False, input_units=None, name=None, - cosmo_factor=cosmo_factor(None, None), + cosmo_factor=None, comoving=None, valid_transform=True, compression=None, @@ -569,32 +569,43 @@ def __new__( input_array.convert_to_physical() else: comoving = input_array.comoving - cosmo_factor = _preserve_cosmo_factor( - cosmo_factor, - getattr(input_array, "cosmo_factor", NULL_CF), + cosmo_factor = ( + getattr(input_array, "cosmo_factor") + if cosmo_factor is None + else _preserve_cosmo_factor( + cosmo_factor, getattr(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"] + elif isinstance(input_array, np.ndarray) and input_array.dtype != object: + # guard np.ndarray so it doesn't get caught by _iterable in next case + if cosmo_factor is None: + cosmo_factor = NULL_CF + elif _iterable(input_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"] + if cosmo_factor is None: + cosmo_factor = _preserve_cosmo_factor(*helper_result["cfs"]) + if cosmo_factor is None: + cosmo_factor = NULL_CF + elif all([cf is None for cf in helper_result["cfs"]]): + cosmo_factor = cosmo_factor + else: cosmo_factor = _preserve_cosmo_factor( cosmo_factor, *helper_result["cfs"] ) - if not valid_transform: - input_array.convert_to_physical() - if compression != helper_result["compression"]: - compression = None # just drop it + if not valid_transform: + input_array.convert_to_physical() + if compression != helper_result["compression"]: + compression = None # just drop it obj = super().__new__( cls, @@ -609,7 +620,10 @@ def __new__( if isinstance(obj, unyt_array) and not isinstance(obj, cls): obj = obj.view(cls) - obj.cosmo_factor = cosmo_factor + # unyt allows creating a unyt_array from e.g. arrays with heterogenous units + # (it probably shouldn't...), so we don't bother recursing deeply and therefore + # can't guarantee that cosmo_factor isn't None at this point, guard with default + obj.cosmo_factor = cosmo_factor if cosmo_factor is not None else NULL_CF obj.comoving = comoving obj.compression = compression obj.valid_transform = valid_transform @@ -620,7 +634,7 @@ def __new__( if obj.comoving: assert ( obj.valid_transform - ), "Comoving Cosmo arrays must be able to be transformed to physical" + ), "Comoving cosmo_arrays must be able to be transformed to physical" return obj @@ -885,14 +899,17 @@ 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], - None, - power=power_map(inputs[0].shape[kwargs["axis"]]), + cfs[0], None, power=power_map(inputs[0].shape[kwargs["axis"]]) ) else: ret_cf = _power_cosmo_factor( cfs[0], None, power=power_map(inputs[0].size) ) + elif ( + ufunc in (logical_and, logical_or, logical_xor, logical_not) + and method == "reduce" + ): + ret_cf = _return_without_cosmo_factor(cfs[0]) else: ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs) @@ -1005,7 +1022,7 @@ def __new__( dtype=None, bypass_validation=False, name=None, - cosmo_factor=cosmo_factor(None, None), + cosmo_factor=None, comoving=None, valid_transform=True, compression=None, @@ -1055,11 +1072,16 @@ def __new__( units = getattr(input_scalar, "units", None) if units is None else units name = getattr(input_scalar, "name", None) if name is None else name - cosmo_factor = ( - getattr(input_scalar, "cosmo_factor", NULL_CF) - if cosmo_factor == NULL_CF - else cosmo_factor - ) + if hasattr(input_scalar, "cosmo_factor") and (cosmo_factor is not None): + cosmo_factor = _preserve_cosmo_factor( + cosmo_factor, getattr(input_scalar, "cosmo_factor") + ) + elif cosmo_factor is None and hasattr(input_scalar, "cosmo_factor"): + cosmo_factor = getattr(input_scalar, "cosmo_factor") + elif cosmo_factor is not None and not hasattr(input_scalar, "cosmo_factor"): + pass + else: + cosmo_factor = NULL_CF comoving = ( getattr(input_scalar, "comoving", None) if comoving is None else comoving ) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 8d375257..59545f50 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -12,8 +12,9 @@ logical_and, s_, ceil, + zeros_like, ) -from unyt import unyt_quantity, exceptions +from unyt import exceptions from swiftsimio import SWIFTDataset, cosmo_array from swiftsimio.reader import __SWIFTGroupDataset @@ -160,12 +161,12 @@ def project_pixel_grid( z_slice_included = True z_min, z_max = region[4:] else: - z_min = unyt_quantity(0.0, units=box_z.units) + z_min = zeros_like(box_z) z_max = box_z else: - x_min = unyt_quantity(0.0, units=box_x.units) + x_min = zeros_like(box_x) x_max = box_x - y_min = unyt_quantity(0.0, units=box_y.units) + y_min = zeros_like(box_y) y_max = box_y x_range = x_max - x_min @@ -457,23 +458,23 @@ def project_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/tests/helper.py b/tests/helper.py index f2d6907d..58b0afd8 100644 --- a/tests/helper.py +++ b/tests/helper.py @@ -6,8 +6,8 @@ import os import h5py from swiftsimio.subset_writer import find_links, write_metadata -from swiftsimio import mask, cosmo_array, load -from numpy import mean +from swiftsimio import mask, cosmo_array +from numpy import mean, zeros_like webstorage_location = "http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/" test_data_location = "test_data/" @@ -82,7 +82,7 @@ def create_single_particle_dataset(filename: str, output_name: str): # Create a dummy mask in order to write metadata data_mask = mask(filename) boxsize = data_mask.metadata.boxsize - region = [[0, b] for b in boxsize] + region = [[zeros_like(b), b] for b in boxsize] data_mask.constrain_spatial(region) # Write the metadata diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 2bac216a..c1a0d4ee 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -20,7 +20,13 @@ 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 cq(x, unit=u.Mpc): + return cosmo_quantity( + x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5) + ) def arg_to_ua(arg): @@ -75,7 +81,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") @@ -86,7 +92,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") @@ -96,7 +102,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") @@ -106,7 +112,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") @@ -117,14 +123,17 @@ 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 @@ -187,9 +196,9 @@ def test_explicitly_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)), - "logspace": (ca(1, unit=u.dimensionless), ca(2, unit=u.dimensionless)), - "geomspace": (ca(1), ca(1)), + "linspace": (cq(1), cq(2)), + "logspace": (cq(1, unit=u.dimensionless), cq(2, unit=u.dimensionless)), + "geomspace": (cq(1), cq(1)), "copyto": (ca(np.arange(3)), ca(np.arange(3))), "prod": (ca(np.arange(3)),), "var": (ca(np.arange(3)),), @@ -205,7 +214,7 @@ def test_explicitly_handled_funcs(self): "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)), + "insert": (ca(np.arange(3)), 1, cq(1)), "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))), "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))), "linalg.tensorsolve": ( @@ -228,11 +237,11 @@ def test_explicitly_handled_funcs(self): "select": ( [np.arange(3) < 1, np.arange(3) > 1], [ca(np.arange(3)), ca(np.arange(3))], - ca(1), + cq(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)), + "clip": (ca(np.arange(3)), cq(1), cq(2)), "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))),), @@ -251,9 +260,9 @@ def test_explicitly_handled_funcs(self): "all": (ca(np.arange(3)),), "amax": (ca(np.arange(3)),), # implemented via max "amin": (ca(np.arange(3)),), # implemented via min - "angle": (ca(complex(1, 1)),), + "angle": (cq(complex(1, 1)),), "any": (ca(np.arange(3)),), - "append": (ca(np.arange(3)), ca(1)), + "append": (ca(np.arange(3)), cq(1)), "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))), "argmax": (ca(np.arange(3)),), # implemented via max "argmin": (ca(np.arange(3)),), # implemented via min @@ -298,7 +307,7 @@ def test_explicitly_handled_funcs(self): "isneginf": (ca(np.arange(3)),), "isposinf": (ca(np.arange(3)),), "empty_like": (ca(np.arange(3)),), - "full_like": (ca(np.arange(3)), ca(1)), + "full_like": (ca(np.arange(3)), cq(1)), "ones_like": (ca(np.arange(3)),), "zeros_like": (ca(np.arange(3)),), "copy": (ca(np.arange(3)),), @@ -422,6 +431,13 @@ def test_explicitly_handled_funcs(self): category=UserWarning, message="numpy.savetxt does not preserve units or cosmo", ) + if "unwrap" in fname: + # haven't bothered to pass a cosmo_quantity for period + warnings.filterwarnings( + action="ignore", + category=RuntimeWarning, + message="Mixing arguments with and without cosmo_factors", + ) result = func(*args) if fname.split(".")[-1] in ( "fill_diagonal", @@ -473,7 +489,7 @@ def test_explicitly_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), ), ), ), @@ -484,13 +500,13 @@ def test_explicitly_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), ), ), ), @@ -502,19 +518,19 @@ def test_explicitly_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), ), ], ), @@ -526,7 +542,7 @@ def test_explicitly_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]), ), @@ -543,19 +559,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] @@ -602,9 +618,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): @@ -612,9 +628,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): @@ -622,9 +638,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 = { @@ -636,9 +652,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), ] ), ): @@ -678,12 +694,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 def test_scalar_return_func(self): @@ -695,7 +711,7 @@ def test_scalar_return_func(self): np.arange(3), u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), valid_transform=True, ) res = np.min(ca) @@ -707,10 +723,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 diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_attrs.py index 7b94443b..1ec07d8b 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_attrs.py @@ -111,7 +111,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 0.5), + cosmo_factor=cosmo_factor(a ** 1, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -120,7 +120,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 0.5), + cosmo_factor=cosmo_factor(a ** 1, 0.5), comoving=False, ) assert not arr.compatible_with_comoving() @@ -129,7 +129,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**0, 0.5), + cosmo_factor=cosmo_factor(a ** 0, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -138,7 +138,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**0, 0.5), + cosmo_factor=cosmo_factor(a ** 0, 0.5), comoving=False, ) assert arr.compatible_with_comoving() @@ -147,7 +147,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=True, ) assert arr.compatible_with_comoving() @@ -156,7 +156,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=False, ) assert arr.compatible_with_comoving() @@ -182,7 +182,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.ones_like(inp) assert res.to_value(u.kpc) == 1 @@ -215,7 +215,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises( @@ -228,7 +228,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -239,13 +239,13 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -256,7 +256,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp + inp assert res.to_value(u.kpc) == 4 @@ -267,7 +267,7 @@ def test_multiplying_ufunc(self): # no cosmo_factors inp = cosmo_array([2], u.kpc, comoving=False) res = inp * inp - assert res.to_value(u.kpc**2) == 4 + assert res.to_value(u.kpc ** 2) == 4 assert res.comoving is False assert res.cosmo_factor == cosmo_factor(None, None) # one is not cosmo_array @@ -289,7 +289,7 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): @@ -300,7 +300,7 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): inp1 * inp2 @@ -309,24 +309,24 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp * inp - assert res.to_value(u.kpc**2) == 4 + assert res.to_value(u.kpc ** 2) == 4 assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**2 + assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_dividing_ufunc(self): inp = cosmo_array( [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp / inp assert res.to_value(u.dimensionless) == 1 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**0 + assert res.cosmo_factor == inp.cosmo_factor ** 0 def test_return_without_ufunc(self): # 1 argument @@ -334,7 +334,7 @@ def test_return_without_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.logical_not(inp) assert res == np.logical_not(1) @@ -344,7 +344,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.logaddexp(inp, inp) assert res == np.logaddexp(2, 2) @@ -354,13 +354,13 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -372,7 +372,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -383,7 +383,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises( @@ -398,7 +398,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) @@ -409,7 +409,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = u.unyt_array([2], u.kpc) with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -422,31 +422,31 @@ def test_sqrt_ufunc(self): [4], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.sqrt(inp) - assert res.to_value(u.kpc**0.5) == 2 # also ensures units ok + assert res.to_value(u.kpc ** 0.5) == 2 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**0.5 + assert res.cosmo_factor == inp.cosmo_factor ** 0.5 def test_square_ufunc(self): inp = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.square(inp) - assert res.to_value(u.kpc**2) == 4 # also ensures units ok + assert res.to_value(u.kpc ** 2) == 4 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**2 + assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_cbrt_ufunc(self): inp = cosmo_array( [8], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.cbrt(inp) assert res.to_value(u.kpc ** (1.0 / 3.0)) == 2 # also ensures units ok @@ -458,12 +458,12 @@ def test_reciprocal_ufunc(self): [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.reciprocal(inp) - assert res.to_value(u.kpc**-1) == 0.5 # also ensures units ok + assert res.to_value(u.kpc ** -1) == 0.5 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**-1 + assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_passthrough_ufunc(self): # 1 argument @@ -471,7 +471,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.negative(inp) assert res.to_value(u.kpc) == -2 @@ -482,7 +482,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.copysign(inp, inp) assert res.to_value(u.kpc) == inp.to_value(u.kpc) @@ -493,13 +493,13 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -511,7 +511,7 @@ def test_arctan2_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.arctan2(inp, inp) assert res.to_value(u.dimensionless) == np.arctan2(2, 2) @@ -523,13 +523,13 @@ def test_comparison_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp1 < inp2 assert res.all() @@ -540,7 +540,7 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) out = cosmo_array([np.nan], u.dimensionless, comoving=True) np.abs(inp, out=out) @@ -551,7 +551,7 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) # make sure we can also pass a non-cosmo type for out without crashing out = np.array([np.nan]) @@ -563,31 +563,31 @@ def test_reduce_multiply(self): [[1, 2], [3, 4]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) res = np.multiply.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc**2), np.array([3.0, 8.0])) + np.testing.assert_allclose(res.to_value(u.kpc ** 2), np.array([3.0, 8.0])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**2 + assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_reduce_divide(self): inp = cosmo_array( [[1.0, 2.0], [1.0, 4.0], [1.0, 1.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) res = np.divide.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc**-1), np.array([1.0, 0.5])) + np.testing.assert_allclose(res.to_value(u.kpc ** -1), np.array([1.0, 0.5])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**-1 + assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_reduce_other(self): inp = cosmo_array( [[1.0, 2.0], [1.0, 2.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.add.reduce(inp, axis=0) np.testing.assert_allclose(res.to_value(u.kpc), np.array([2.0, 4.0])) @@ -600,7 +600,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res1, res2 = np.modf(inp) assert res1.to_value(u.kpc) == 0.5 @@ -614,7 +614,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res1, res2 = np.frexp(inp) assert res1 == 0.625 @@ -628,7 +628,7 @@ def test_multi_output_with_out_arg(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) out1 = cosmo_array([np.nan], u.dimensionless, comoving=True) out2 = cosmo_array([np.nan], u.dimensionless, comoving=True) @@ -645,7 +645,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = 0 res = inp1 > inp2 @@ -654,7 +654,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = 0.5 with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -664,20 +664,20 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [0, 0, 0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) assert (inp1 > inp2).all() inp1 = cosmo_array( [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = np.ones(3) * u.kpc with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -686,7 +686,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = np.zeros(3) * u.kpc assert (inp1 > inp2).all() @@ -694,13 +694,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( 1, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp1 == inp2 assert res.all() @@ -708,13 +708,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( 0, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp1 > inp2 assert res.all() From 448ba7dc3b81f7bcc8b03b698e7e2dad66c118ef Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 6 Feb 2025 17:58:40 +0000 Subject: [PATCH 059/125] Remove unused import. --- swiftsimio/visualisation/power_spectrum.py | 21 ++++++++++----------- 1 file changed, 10 insertions(+), 11 deletions(-) diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index e602f221..6411de5b 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -10,7 +10,6 @@ from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio import cosmo_array, cosmo_quantity from swiftsimio.reader import __SWIFTGroupDataset -from swiftsimio._array_functions import _reciprocal_cosmo_factor from typing import Optional, Dict, Tuple @@ -199,7 +198,7 @@ def render_to_deposit( """ # Get the positions and masses - folding = 2.0**folding + folding = 2.0 ** folding positions = data.coordinates quantity = getattr(data, project) @@ -246,10 +245,10 @@ def render_to_deposit( units = 1.0 / ( data.metadata.boxsize[0] * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) units *= quantity.units - new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor**3) + new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor ** 3) return cosmo_array( deposition, comoving=comoving, cosmo_factor=new_cosmo_factor, units=units @@ -359,7 +358,7 @@ def folded_depositions_to_power_spectrum( np.zeros(number_of_wavenumber_bins), units=box_volume.units, comoving=box_volume.comoving, - cosmo_factor=box_volume.cosmo_factor**-1, + cosmo_factor=box_volume.cosmo_factor ** -1, name="Power spectrum $P(k)$", ) folding_tracker = np.ones(number_of_wavenumber_bins, dtype=float) @@ -397,7 +396,7 @@ def folded_depositions_to_power_spectrum( if folding != final_folding: cutoff_wavenumber = ( - 2.0**folding * np.min(depositions[folding].shape) / np.min(box_size) + 2.0 ** folding * np.min(depositions[folding].shape) / np.min(box_size) ) if cutoff_above_wavenumber_fraction is not None: @@ -422,7 +421,7 @@ def folded_depositions_to_power_spectrum( corrected_wavenumber_centers[prefer_bins] = folded_wavenumber_centers[ prefer_bins ].to(corrected_wavenumber_centers.units) - folding_tracker[prefer_bins] = 2.0**folding + folding_tracker[prefer_bins] = 2.0 ** folding contributed_counts[prefer_bins] = folded_counts[prefer_bins] elif transition == "average": @@ -455,7 +454,7 @@ def folded_depositions_to_power_spectrum( # For debugging, we calculate an effective fold number. folding_tracker[use_bins] = ( - (folding_tracker * existing_weight + (2.0**folding) * new_weight) + (folding_tracker * existing_weight + (2.0 ** folding) * new_weight) / transition_norm )[use_bins] @@ -536,7 +535,7 @@ def deposition_to_power_spectrum( deposition.shape == cross_deposition.shape ), "Depositions must have the same shape" - folding = 2.0**folding + folding = 2.0 ** folding box_size_folded = box_size[0] / folding npix = deposition.shape[0] @@ -558,7 +557,7 @@ def deposition_to_power_spectrum( else: conj_fft = fft.conj() - fourier_amplitudes = (fft * conj_fft).real * box_size_folded**3 + fourier_amplitudes = (fft * conj_fft).real * box_size_folded ** 3 # Calculate k-value spacing (centered FFT) dk = 2 * np.pi / (box_size_folded) @@ -590,7 +589,7 @@ def deposition_to_power_spectrum( divisor[zero_mask] = 1 # Correct for folding - binned_amplitudes *= folding**3 + binned_amplitudes *= folding ** 3 # Correct units and names wavenumbers = binned_wavenumbers / divisor From 85f3900df68c496b72009b3768e904fd46a38286 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 6 Feb 2025 20:33:50 +0000 Subject: [PATCH 060/125] Replace some unyt usage with cosmo versions. --- swiftsimio/visualisation/power_spectrum.py | 7 ++++--- swiftsimio/visualisation/ray_trace.py | 14 +++++++------- swiftsimio/visualisation/slice.py | 22 +++++++++++----------- 3 files changed, 22 insertions(+), 21 deletions(-) diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index 6411de5b..fb090937 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -2,7 +2,7 @@ Tools for creating power spectra from SWIFT data. """ -from numpy import float32, float64, int32, zeros, ndarray +from numpy import float32, float64, int32, zeros, ndarray, zeros_like import numpy as np import scipy.fft @@ -596,9 +596,10 @@ def deposition_to_power_spectrum( wavenumbers.name = "Wavenumber $k$" shot_noise = ( - (box_size[0] ** 3 / shot_noise_norm) if shot_noise_norm is not None else 0.0 + (box_size[0] ** 3 / shot_noise_norm) + if shot_noise_norm is not None + else zeros_like(box_size[0] ** 3) # copy cosmo properties ) - power_spectrum = (binned_amplitudes / divisor) - shot_noise power_spectrum.name = "Power Spectrum $P(k)$" diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 3ef77ff3..c5d65cb1 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -273,11 +273,11 @@ def panel_pixel_grid( z_min = unyt.unyt_quantity(0.0, units=box_z.units) z_max = box_z else: - x_min = unyt.unyt_quantity(0.0, units=box_x.units) + x_min = np.zeros_like(box_x) x_max = box_x - y_min = unyt.unyt_quantity(0.0, units=box_y.units) + y_min = np.zeros_like(box_y) y_max = box_y - z_min = unyt.unyt_quantity(0.0, units=box_z.units) + z_min = np.zeros_like(box_z) z_max = box_z x_range = x_max - x_min @@ -353,23 +353,23 @@ def panel_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index bc5b67e0..f9782186 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -3,7 +3,7 @@ """ from typing import Union, Optional -from numpy import float32, array, ones, matmul +from numpy import float32, array, ones, matmul, zeros_like from swiftsimio import SWIFTDataset, cosmo_array, cosmo_quantity from swiftsimio.visualisation.slice_backends import backends, backends_parallel from swiftsimio.visualisation.smoothing_length import backends_get_hsml @@ -91,7 +91,7 @@ def slice_gas_pixel_grid( """ if z_slice is None: - z_slice = 0.0 * data.gas.coordinates.units + z_slice = zeros_like(data.metadata.boxsize[0]) number_of_gas_particles = data.gas.coordinates.shape[0] @@ -113,16 +113,16 @@ def slice_gas_pixel_grid( box_x, box_y, box_z = data.metadata.boxsize - if z_slice > box_z or z_slice < (0 * box_z): + if z_slice > box_z or z_slice < zeros_like(box_z): raise ValueError("Please enter a slice value inside the box.") # Set the limits of the image. if region is not None: x_min, x_max, y_min, y_max = region else: - x_min = (0 * box_x).to(box_x.units) + x_min = zeros_like(box_x) x_max = box_x - y_min = (0 * box_y).to(box_y.units) + y_min = zeros_like(box_y) y_max = box_y x_range = x_max - x_min @@ -146,7 +146,7 @@ def slice_gas_pixel_grid( else: x, y, z = data.gas.coordinates.T - z_center = 0 * box_z + z_center = zeros_like(box_z) hsml = backends_get_hsml[backend](data) if data.gas.coordinates.comoving: @@ -278,7 +278,7 @@ def slice_gas( """ if z_slice is None: - z_slice = 0.0 * data.gas.coordinates.units + z_slice = zeros_like(data.metadata.boxsize[0]) image = slice_gas_pixel_grid( data, @@ -297,7 +297,7 @@ def slice_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range**3) + units = 1.0 / (max_range ** 3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units( @@ -305,17 +305,17 @@ def slice_gas( ) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**3) + units = 1.0 / (max_range ** 3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) From 0f80416149b6ab5bef6a5d328123300cd5494d37 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 7 Feb 2025 12:53:39 +0000 Subject: [PATCH 061/125] Some cleanup of array constructors. --- swiftsimio/objects.py | 68 +++++++++++++++++++++++++++++-------------- 1 file changed, 46 insertions(+), 22 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index f5b1ccf7..5df6075c 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -344,7 +344,7 @@ def __rtruediv__(self, b): def __pow__(self, p): if self.expr is None: return cosmo_factor(expr=None, scale_factor=self.scale_factor) - 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 @@ -539,9 +539,9 @@ def __new__( 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. + corrupted or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. Defaults to False. input_units : str, optional deprecated in favour of units option name : str, optional @@ -562,6 +562,22 @@ def __new__( cosmo_factor: cosmo_factor + if bypass_validation is True: + obj = super().__new__( + cls, + input_array, + units=units, + registry=registry, + dtype=dtype, + bypass_validation=bypass_validation, + name=name, + ) + # dtype, units, registry & name handled by unyt + obj.comoving = comoving + obj.cosmo_factor = cosmo_factor if cosmo_factor is not None else NULL_CF + obj.valid_transform = valid_transform + obj.compression = compression + return obj if isinstance(input_array, cosmo_array): if comoving: input_array.convert_to_comoving() @@ -594,8 +610,6 @@ def __new__( input_array = helper_result["args"] if cosmo_factor is None: cosmo_factor = _preserve_cosmo_factor(*helper_result["cfs"]) - if cosmo_factor is None: - cosmo_factor = NULL_CF elif all([cf is None for cf in helper_result["cfs"]]): cosmo_factor = cosmo_factor else: @@ -1046,9 +1060,9 @@ def __new__( 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. + corrupted or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. 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 @@ -1064,23 +1078,33 @@ def __new__( Description of the compression filters that were applied to that array in the hdf5 file. """ - if not ( - bypass_validation - or isinstance(input_scalar, (numeric_type, np.number, np.ndarray)) - ): + if bypass_validation is True: + ret = super().__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 not isinstance(input_scalar, (numeric_type, np.number, np.ndarray)): raise RuntimeError("cosmo_quantity values must be numeric") + # Use values from kwargs, if None use values from input_scalar units = getattr(input_scalar, "units", None) if units is None else units name = getattr(input_scalar, "name", None) if name is None else name - if hasattr(input_scalar, "cosmo_factor") and (cosmo_factor is not None): - cosmo_factor = _preserve_cosmo_factor( - cosmo_factor, getattr(input_scalar, "cosmo_factor") - ) - elif cosmo_factor is None and hasattr(input_scalar, "cosmo_factor"): - cosmo_factor = getattr(input_scalar, "cosmo_factor") - elif cosmo_factor is not None and not hasattr(input_scalar, "cosmo_factor"): - pass - else: + cosmo_factor = ( + getattr(input_scalar, "cosmo_factor", None) + if cosmo_factor is None + else cosmo_factor + ) + if cosmo_factor is None: cosmo_factor = NULL_CF comoving = ( getattr(input_scalar, "comoving", None) if comoving is None else comoving From b2ce2df2f7c44abc6f0948cc3a4419654330e028 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 7 Feb 2025 17:51:35 +0000 Subject: [PATCH 062/125] Cleanup and document _array_functions.py --- swiftsimio/_array_functions.py | 862 +++++++++++++++--- swiftsimio/objects.py | 160 ++-- .../smoothing_length/generate.py | 4 +- 3 files changed, 843 insertions(+), 183 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index f2cbbf96..5371989f 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -1,8 +1,19 @@ +""" +Overloaded implementations of unyt and numpy functions to correctly handle +:class:`~swiftsimio.objects.cosmo_array` input. + +This module also defines wrappers and helper functions to facilitate overloading +functions and handling the processing of our custom array attributes. + +Nothing in this module is intended to be user-facing, but the helpers and wrappers +are documented to assist in maintenance and development of swiftsimio. +""" + import warnings import numpy as np +from typing import Callable, Tuple import unyt from unyt import unyt_quantity, unyt_array -from unyt.array import _iterable from swiftsimio import objects from unyt._array_functions import ( dot as unyt_dot, @@ -111,92 +122,239 @@ # numpy functions (we will actually wrap the functions below): -def _copy_cosmo_array_attributes(from_ca, to_ca): - if not isinstance(to_ca, objects.cosmo_array) and isinstance( - from_ca, objects.cosmo_array +def _copy_cosmo_array_attributes(from_ca: object, to_ca: object) -> object: + """ + Copy :class:`~swiftsimio.objects.cosmo_array` attributes across two objects. + + Copies the ``cosmo_factor``, ``comoving``, ``valid_transform`` and ``compression`` + attributes across if both the source and destination objects are + :class:`~swiftsimio.objects.cosmo_array` instances (else returns input). + + Parameters + ---------- + from_ca : :obj:`object` + The source object. + + to_ca : :obj:`object` + The destination object. + + Returns + ------- + out : :obj:`object` + The destination object (with attributes copied if copy occurred). + """ + if not ( + isinstance(to_ca, objects.cosmo_array) + and isinstance(from_ca, objects.cosmo_array) ): return to_ca - if hasattr(from_ca, "cosmo_factor"): - to_ca.cosmo_factor = from_ca.cosmo_factor - if hasattr(from_ca, "comoving"): - to_ca.comoving = from_ca.comoving - if hasattr(from_ca, "valid_transform"): - to_ca.valid_transform = from_ca.valid_transform + to_ca.cosmo_factor = from_ca.cosmo_factor + to_ca.comoving = from_ca.comoving + to_ca.valid_transform = from_ca.valid_transform + to_ca.compression = from_ca.compression return to_ca -def _propagate_cosmo_array_attributes(func): - # can work on methods (obj is self) and functions (obj is first argument) +def _propagate_cosmo_array_attributes_to_result(func: Callable) -> Callable: + """ + Wrapper that copies :class:`~swiftsimio.objects.cosmo_array` attributes from first + input argument to first output. + + Many functions take one input (or have a first input that has a close correspondance + to the output) and one output. This helper copies the ``cosmo_factor``, ``comoving``, + ``valid_transform`` and ``compression`` attributes from the first input argument to + the output. Can be used as a decorator on functions (the first argument is then the + first argument of the function) or methods (the first argument is then ``self``). + If the output is not a :class:`~swiftsimio.objects.cosmo_array` it is not promoted + (and then no attributes are copied). + + Parameters + ---------- + func : callable + The function whose argument attributes will be copied to its result. + + Returns + ------- + out : callable + The wrapped function. + """ + def wrapped(obj, *args, **kwargs): - ret = func(obj, *args, **kwargs) - if not isinstance(ret, objects.cosmo_array): - return ret - ret = _copy_cosmo_array_attributes(obj, ret) - return ret + """ + Call the function and copy the attributes. + """ + return _copy_cosmo_array_attributes(obj, func(obj, *args, **kwargs)) return wrapped -def _promote_unyt_to_cosmo(input_object): - # converts a unyt_quantity to a cosmo_quantity, or a unyt_array to a cosmo_array +def _promote_unyt_to_cosmo(input_object: object) -> object: + """ + Upgrades the input unyt instance to its cosmo equivalent. + + In many cases we can obtain a unyt class instance and want to promote it to its cosmo + equivalent to attach our cosmo attributes. This helper promotes an input + :class:`~unyt.array.unyt_array` to a :class:`~swiftsimio.objects.cosmo_array` or an + input :class:`~unyt.array.unyt_quantity` to a + :class:`~swiftsimio.objects.cosmo_quantity`. If the input is neither type, it is just + returned. + + Parameters + ---------- + input_object : :obj:`object` + Object to consider for promotion from unyt instance to cosmo instance. + """ if isinstance(input_object, unyt_quantity) and not isinstance( input_object, objects.cosmo_quantity ): - ret = objects.cosmo_quantity(input_object) + return input_object.view(objects.cosmo_quantity) elif isinstance(input_object, unyt_array) and not isinstance( input_object, objects.cosmo_array ): - ret = objects.cosmo_array(input_object) + return input_object.view(objects.cosmo_array) else: - ret = input_object - return ret + return input_object -def _ensure_array_or_quantity_matches_shape(input_object): +def _ensure_array_or_quantity_matches_shape(input_object: object) -> object: + """ + Convert scalars to :class:`~swiftsimio.objects.cosmo_quantity` and arrays to + :class:`~swiftsimio.objects.cosmo_array`. + + Scalar quantities are meant to be contained in + :class:`~swiftsimio.objects.cosmo_quantity` and arrays in + :class:`~swiftsimio.objects.cosmo_array`. Many functions (e.g. from numpy) can change + the data contents without changing the containing class. This helper checks the input + to make sure the data match the container type and converts if not. + + Parameters + ---------- + input_object : :obj:`object` + The object whose data is to be checked against its type. + + Returns + ------- + out : :obj:`object` + A version of the input with container type matching data contents. + """ if ( isinstance(input_object, objects.cosmo_array) and not isinstance(input_object, objects.cosmo_quantity) and input_object.shape == () ): - ret = objects.cosmo_quantity(input_object) + return input_object.view(objects.cosmo_quantity) elif isinstance(input_object, objects.cosmo_quantity) and input_object.shape != (): - ret = objects.cosmo_array(input_object) + return input_object.view(objects.cosmo_array) else: - ret = input_object - return ret + return input_object -def _ensure_cosmo_array_or_quantity(func): - # can work on methods (obj is self) and functions (obj is first argument) - def wrapped(obj, *args, **kwargs): - ret = func(obj, *args, **kwargs) - if isinstance(ret, tuple): - ret = tuple( +def _ensure_result_is_cosmo_array_or_quantity(func: Callable) -> Callable: + """ + Wrapper that converts any :class:`~unyt.array.unyt_array` or + :class:`~unyt.array.unyt_quantity` instances in function output to cosmo equivalents. + + If the wrapped function returns a :obj:`tuple` (as many numpy functions do) it is + iterated over (but not recursively) and each element with a unyt class type is + upgraded to its cosmo equivalent. If anything but a :obj:`tuple` is returned, that + object is promoted to the cosmo equivalent if it is of a unyt class type. + + Parameters + ---------- + func : Callable + The function whose result(s) will be upgraded to + :class:`~swiftsimio.objects.cosmo_array` or + :class:`~swifsimio.objects.cosmo_quantity`. + Returns + ------- + out : Callable + The wrapped function. + """ + + def wrapped(*args, **kwargs) -> object: + """ + Promote unyt types in function output to cosmo types. + """ + result = func(*args, **kwargs) + if isinstance(result, tuple): + return tuple( _ensure_array_or_quantity_matches_shape(_promote_unyt_to_cosmo(item)) - for item in ret + for item in result ) else: - ret = _ensure_array_or_quantity_matches_shape(_promote_unyt_to_cosmo(ret)) - return ret + return _ensure_array_or_quantity_matches_shape( + _promote_unyt_to_cosmo(result) + ) return wrapped -def _sqrt_cosmo_factor(cf, **kwargs): - return _power_cosmo_factor(cf, None, power=0.5) # ufunc sqrt not supported +def _sqrt_cosmo_factor(cf: "objects.cosmo_factor", **kwargs) -> "objects.cosmo_factor": + """ + Take the square root of a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` whose square root should be taken. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The square root of the input :class:`~swiftsimio.objects.cosmo_factor`. + """ + return _power_cosmo_factor(cf, None, power=0.5) + + +def _multiply_cosmo_factor( + *cfs: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Recursively multiply :class:`~swiftsimio.objects.cosmo_factor`s. + All arguments are expected to be of type :class:`~swiftsimio.objects.cosmo_factor`. + They are cumumatively multipled together and the result returned. -def _multiply_cosmo_factor(*args, **kwargs): - cfs = args + Parameters + ---------- + cfs : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor`s to be multiplied. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the input :class:`~swiftsimio.objects.cosmo_factor`s. + """ if len(cfs) == 1: - return __multiply_cosmo_factor(cfs[0]) - retval = __multiply_cosmo_factor(cfs[0], cfs[1]) + return __binary_multiply_cosmo_factor(cfs[0]) + retval = __binary_multiply_cosmo_factor(cfs[0], cfs[1]) for cf in cfs[2:]: - retval = __multiply_cosmo_factor(retval, cf) + retval = __binary_multiply_cosmo_factor(retval, cf) return retval -def __multiply_cosmo_factor(cf1, cf2, **kwargs): +def __binary_multiply_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Multiply two :class:`~swiftsimio.objects.cosmo_factor`s. + + Not intended for direct use but only as a helper for + :func:`~swiftsimio._array_functions._multiply_cosmo_factor`. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + The first :class:`~swiftsimio.objects.cosmo_factor`. + + cf2 : swiftsimio.objects.cosmo_factor + The second :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the :class:`~swiftsimio.objects.cosmo_factor`s. + """ if (cf1 is None) and (cf2 is None): # neither has cosmo_factor information: return None @@ -209,21 +367,67 @@ def __multiply_cosmo_factor(cf1, cf2, **kwargs): elif (cf1 is not None) and (cf2 is not None): # both cosmo_array and both with cosmo_factor: return cf1 * cf2 # cosmo_factor.__mul__ raises if scale factors differ - else: - raise RuntimeError("Unexpected state, please report this error on github.") -def _preserve_cosmo_factor(*args, **kwargs): - cfs = args +def _preserve_cosmo_factor( + *cfs: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Helper to preserve the :class:`~swiftsimio.objects.cosmo_factor` of input. + + If there is a single argument, return its ``cosmo_factor``. If there are multiple + arguments, check that they all have matching ``cosmo_factor``. Any arguments that + are not :class:`~swiftsimio.objects.cosmo_array`s are ignored for this purpose. + + Parameters + ---------- + cfs : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor`s to be preserved. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The preserved :class:`~swiftsimio.objects.cosmo_factor`. + """ if len(cfs) == 1: return cfs[0] - retval = __preserve_cosmo_factor(cfs[0], cfs[1]) + retval = __binary_preserve_cosmo_factor(cfs[0], cfs[1]) for cf in cfs[2:]: - retval = __preserve_cosmo_factor(retval, cf) + retval = __binary_preserve_cosmo_factor(retval, cf) return retval -def __preserve_cosmo_factor(cf1, cf2, **kwargs): +def __binary_preserve_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Given two :class:`~swiftsimio.objects.cosmo_factor`s, get it if they match. + + Not intended for direct use but only as a helper for + :func:`~swiftsimio._array_functions._preserve_cosmo_factor`. If the two inputs + are compatible, return the compatible :class:`~swiftsimio.objects.cosmo_factor`. + If one of them is ``None``, produce a warning. If they are incompatible, raise. + + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + The first :class:`~swiftsimio.objects.cosmo_factor`. + + cf2 : swiftsimio.objects.cosmo_factor + The second :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The preserved :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + Raised if the two arguments are :class:`~swiftsimio.objects.cosmo_factor`s + that do not have matching attributes. + """ if (cf1 is None) and (cf2 is None): # neither has cosmo_factor information: return None @@ -247,11 +451,44 @@ def __preserve_cosmo_factor(cf1, cf2, **kwargs): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif cf1 == cf2: return cf1 # or cf2, they're equal - else: - raise RuntimeError("Unexpected state, please report this error on github.") -def _power_cosmo_factor(cf1, cf2, inputs=None, power=None): +def _power_cosmo_factor( + cf1: "objects.cosmo_factor", + cf2: "objects.cosmo_factor", + inputs: "Tuple[objects.cosmo_array]" = None, + power: float = None, +) -> "objects.cosmo_factor": + """ + Raise a :class:`~swiftsimio.objects.cosmo_factor` to a power of another + :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` that is the base. + + cf2 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` that is the exponent. + + inputs : :obj:`tuple` + The objects that ``cf1`` and ``cf2`` are attached to. Give ``inputs`` or + ``power``, not both. + + power : float + The power to raise ``cf1`` to. Give ``inputs`` or ``power``, not both. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The exponentiated :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the exponent is not a dimensionless quantity, or the exponent has a + ``cosmo_factor`` whose scaling with scale factor is not ``1.0``. + """ if inputs is not None and power is not None: raise ValueError power = inputs[1] if inputs else power @@ -269,34 +506,164 @@ def _power_cosmo_factor(cf1, cf2, inputs=None, power=None): return np.power(cf1, power) -def _square_cosmo_factor(cf, **kwargs): +def _square_cosmo_factor( + cf: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Square a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` to square. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The squared :class:`~swiftsimio.objects.cosmo_factor`. + """ return _power_cosmo_factor(cf, None, power=2) -def _cbrt_cosmo_factor(cf, **kwargs): +def _cbrt_cosmo_factor(cf: "objects.cosmo_factor", **kwargs) -> "objects.cosmo_factor": + """ + Take the cube root of a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` whose cube root should be taken. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The cube root of the input :class:`~swiftsimio.objects.cosmo_factor`. + """ return _power_cosmo_factor(cf, None, power=1.0 / 3.0) -def _divide_cosmo_factor(cf1, cf2, **kwargs): +def _divide_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Divide two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + Numerator :class:`~swiftsimio.objects.cosmo_factor`. + + cf1 : swiftsimio.objects.cosmo_factor + Denominator :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The ratio of the input :class:`~swiftsimio.objects.cosmo_factor`s. + """ return _multiply_cosmo_factor(cf1, _reciprocal_cosmo_factor(cf2)) -def _reciprocal_cosmo_factor(cf, **kwargs): +def _reciprocal_cosmo_factor( + cf: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Take the inverse of a :class:`~swiftsimio.objects.cosmo_factor`. + + Parameters + ---------- + cf : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` to be inverted. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The inverted :class:`~swiftsimio.objects.cosmo_factor`. + """ return _power_cosmo_factor(cf, None, power=-1) -def _passthrough_cosmo_factor(cf, cf2=None, **kwargs): +def _passthrough_cosmo_factor( + cf: "objects.cosmo_factor", cf2: "objects.cosmo_factor" = None, **kwargs +) -> "objects.cosmo_factor": + """ + Preserve a :class:`~swiftsimio.objects.cosmo_factor`, optionally checking that it + matches a second :class:`~swiftsimio.objects.cosmo_factor`. + + This helper is intended for e.g. numpy ufuncs with a second dimensionless argument + so it's ok if ``cf2`` is ``None`` and ``cf1`` is not. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` to pass through. + + cf2 : swiftsimio.objects.cosmo_factor + Optional second :class:`~swiftsimio.objects.cosmo_factor` to check matches. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The input :class:`~swiftsimio.objects.cosmo_factor`. + + Raises + ------ + ValueError + If ``cf2`` is provided, is not ``None`` and does not match ``cf1``. + """ if (cf2 is not None) and cf != cf2: # if both have cosmo_factor information and it differs this is an error raise ValueError(f"Arguments have cosmo_factors that differ: {cf} and {cf2}.") else: - # passthrough is for e.g. ufuncs with a second dimensionless argument, - # so ok if cf2 is None and cf1 is not return cf -def _return_without_cosmo_factor(cf, cf2=np._NoValue, zero_comparison=None, **kwargs): +def _return_without_cosmo_factor( + cf: "objects.cosmo_factor", + cf2: "objects.cosmo_factor" = np._NoValue, + zero_comparison: bool = None, + **kwargs, +) -> None: + """ + Return ``None``, but first check that argument + :class:`~swiftsimio.objects.cosmo_factor`s match, raising or warning if not. + + Comparisons are a special case that wraps around this wrapper, see + :func:`~swiftsimio._array_functions._comparison_cosmo_factor`. + + We borrow ``np._NoValue`` as a default for ``cf2`` because ``None`` here + represents the absence of a :class:`~swiftsimio.objects.cosmo_factor`, and + we need to handle that case. + + Warnings are produced when one argument has no cosmo factor, relevant e.g. + when comparing to constants. We handle comparison with zero, that is unambiguous, + as a special case. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to discard. + + cf2 : swiftsimio.objects.cosmo_factor + Optional second :class:`~swiftsimio.objects.cosmo_factor` to check for a + match with ``cf1``. + + zero_comparison : bool + If ``True``, silences warnings when exactly one of ``cf1`` and ``cf2`` is + ``None``. Enables comparing with zero without warning. + + Returns + ------- + out : None + The :class:`~swiftsimio.objects.cosmo_factor` is discarded. + + Raises + ------ + ValueError + If ``cf2`` is provided, is not ``None`` and does not match ``cf1``. + """ if cf2 is np._NoValue: + # there was no second argument, return promptly return None if (cf is not None) and (cf2 is None): # one is not a cosmo_array, warn on e.g. comparison to constants: @@ -320,13 +687,37 @@ def _return_without_cosmo_factor(cf, cf2=np._NoValue, zero_comparison=None, **kw elif (cf is not None) and (cf2 is not None) and (cf == cf2): # both have cosmo_factor, and they match: pass - else: - raise RuntimeError("Unexpected state, please report this error on github.") # return without cosmo_factor return None -def _arctan2_cosmo_factor(cf1, cf2, **kwargs): +def _arctan2_cosmo_factor( + cf1: "objects.cosmo_factor", cf2: "objects.cosmo_factor", **kwargs +) -> "objects.cosmo_factor": + """ + Helper specifically to handle the :class:`~swiftsimio.objects.cosmo_factor`s for the + ``arctan2`` ufunc from numpy. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` for the first ``arctan2`` argument. + + cf2 : swiftsimio.objects.cosmo_factor + :class:`~swiftsimio.objects.cosmo_factor` for the second ``arctan2`` argument. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` for the ``arctan2`` result. + + Raises + ------ + ValueError + If the input :class:`~swiftsimio.objects.cosmo_factor`s differ, they will + not cancel out and this is an error. + + """ if (cf1 is None) and (cf2 is None): return None elif (cf1 is None) and (cf2 is not None): @@ -335,23 +726,51 @@ def _arctan2_cosmo_factor(cf1, cf2, **kwargs): f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) - else: - raise RuntimeError("Unexpected state, please report this error on github.") + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + + +def _comparison_cosmo_factor( + cf1: "objects.cosmo_factor", + cf2: "objects.cosmo_factor", + inputs: "Tuple[objects.cosmo_array]" = None, +) -> None: + """ + Helper to enable comparisons involving :class:`~swiftsimio.objects.cosmo_factor`s. + + Warnings are emitted when the comparison is ambiguous, for instance if comparing to a + bare :obj:`float` or similar. Comparison to zero is a special case where we suppress + warnings. + + See also :func:`~swiftsimio._array_functions._return_without_cosmo_factor` that is + used in implementing this function. + + Parameters + ---------- + cf1 : swiftsimio.objects.cosmo_factor + First :class:`~swiftsimio.objects.cosmo_factor` to compare. + + cf2 : swiftsimio.objects.cosmo_factor + Second :class:`~swiftsimio.objects.cosmo_factor` to compare. + inputs : :obj:`tuple` + The objects that ``cf1`` and ``cf2`` are attached to. -def _comparison_cosmo_factor(cf1, cf2, inputs=None): + Returns + ------- + out : None + The :class:`~swiftsimio.objects.cosmo_factor` is discarded. + """ try: iter(inputs[0]) except TypeError: @@ -374,21 +793,52 @@ def _comparison_cosmo_factor(cf1, cf2, inputs=None): return _return_without_cosmo_factor(cf1, cf2=cf2, zero_comparison=zero_comparison) -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: - # >>> 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 inp) to this function - # to check consistency and attempt to coerce to comoving if needed. +def _prepare_array_func_args(*args, _default_cm: bool = True, **kwargs) -> dict: + """ + Coerce args and kwargs to a common ``comoving`` and collect ``cosmo_factor``s. + + This helper function is mostly intended for writing wrappers for unyt and numpy + functions. It checks for consistency for all args and kwargs, but is not recursive + so mixed cosmo attributes could be passed in the first argument to + :func:`numpy.concatenate`, for instance. This function can be used "recursively" in a + limited way manually: in functions like :func:`numpy.concatenate` where a list of + arrays is expected, it makes sense to pass the first argument to this function + to check consistency and attempt to coerce to comoving if needed. + + Note that unyt allows creating a :class:`~unyt.array.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 us to try to cover all possible invalid user input without + unyt being stricter. + + The best way to understand the usage of this helper function is to look at the many + wrapped unyt and numpy functions in :mod:`~swiftsimio._array_functions`. + + Parameters + ---------- + _default_cm: bool + If mixed ``comoving`` attributes are found, their data are converted such that + their ``comoving`` has the value of this argument. (Default: ``True``) + + Returns + ------- + out : dict + A dictionary containing the input `args`` and ``kwargs`` coerced to a common + state, and lists of their ``cosmo_factor`` attributes, and ``comoving`` and + ``compression`` values that can be used in return values for wrapped functions, + when relevant. + + Raises + ------ + ValueError + If the input arrays cannot be coerced to a consistent state of ``comoving``. + """ cms = [(hasattr(arg, "comoving"), getattr(arg, "comoving", None)) for arg in args] cfs = [getattr(arg, "cosmo_factor", None) for arg in args] comps = [ @@ -460,18 +910,75 @@ def _prepare_array_func_args(*args, _default_cm=True, **kwargs): ) -def implements(numpy_function): - """Register an __array_function__ implementation for cosmo_array objects.""" +def implements(numpy_function: Callable) -> Callable: + """ + Register an __array_function__ implementation for cosmo_array objects. + + Intended for use as a decorator. + + Parameters + ---------- + numpy_function : Callable + A function handle from numpy (not ufuncs). + + Returns + ------- + out : Callable + The wrapped function. + """ # See NEP 18 https://numpy.org/neps/nep-0018-array-function-protocol.html - def decorator(func): + def decorator(func: Callable) -> Callable: + """ + Actually register the specified function. + + Parameters + ---------- + func: Callable + The function wrapping the numpy equivalent. + + Returns + ------- + out : Callable + The input ``func``. + """ _HANDLED_FUNCTIONS[numpy_function] = func return func return decorator -def _return_helper(res, helper_result, ret_cf, out=None): +def _return_helper( + res: np.ndarray, + helper_result: dict, + ret_cf: "objects.cosmo_factor", + out: np.ndarray = None, +) -> "objects.cosmo_array": + """ + Helper function to attach our cosmo attributes to return values of wrapped functions. + + The return value is first promoted to be a :class:`~swiftsimio.objects.cosmo_array` + (or quantity) if necessary. If the return value is still not one of our cosmo + types, we don't attach attributes. + + Parameters + ---------- + res : numpy.ndarray + The output array of a function to attach our attributes to. + helper_result : dict + A helper :obj:`dict` returned by + :func:`~swiftsimio._array_functions._prepare_array_func_args`. + ret_cf : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to attach to the result. + out : numpy.ndarray + For functions that can place output in an ``out`` argument, a reference to + the output array (optional). + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The input return value of a wrapped function with our cosmo attributes applied. + """ res = _promote_unyt_to_cosmo(res) if isinstance(res, objects.cosmo_array): # also recognizes cosmo_quantity res.comoving = helper_result["comoving"] @@ -484,13 +991,44 @@ def _return_helper(res, helper_result, ret_cf, out=None): return res -def _default_unary_wrapper(unyt_func, cosmo_factor_wrapper): +def _default_unary_wrapper( + unyt_func: Callable, cosmo_factor_handler: Callable +) -> Callable: + """ + Wrapper helper for unary functions with typical behaviour. + + For many numpy and unyt functions with one (main) input argument, the wrapping + code that we need to apply is repetitive. Just prepare the arguments, apply + the chosen processing to the ``cosmo_factor``, and attach our cosmo attributes + to the return value. This function facilitates writing these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) function to be wrapped. + cosmo_factor_handler : Callable + The function that handles the ``cosmo_factor``s, chosen from those defined + in :mod:`~swiftsimio._array_functions`. + + Returns + ------- + out : Callable + The wrapped function. + """ - # assumes that we have one primary argument that will be handled - # by the cosmo_factor_wrapper def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + + Returns + ------- + out : Callable + The wrapped function. + """ helper_result = _prepare_array_func_args(*args, **kwargs) - ret_cf = cosmo_factor_wrapper(helper_result["cfs"][0]) + ret_cf = cosmo_factor_handler(helper_result["cfs"][0]) res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) if "out" in kwargs: return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) @@ -500,13 +1038,44 @@ def wrapper(*args, **kwargs): return wrapper -def _default_binary_wrapper(unyt_func, cosmo_factor_wrapper): +def _default_binary_wrapper( + unyt_func: Callable, cosmo_factor_handler: Callable +) -> Callable: + """ + Wrapper helper for binary functions with typical behaviour. + + For many numpy and unyt functions with two (main) input arguments, the wrapping + code that we need to apply is repetitive. Just prepare the arguments, apply + the chosen processing to the cosmo_factors, and attach our cosmo attributes + to the return value. This function facilitates writing these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) function to be wrapped. + cosmo_factor_handler : Callable + The function that handles the ``cosmo_factor``s, chosen from those defined + in :mod:`~swiftsimio._array_functions`. + + Returns + ------- + out : Callable + The wrapped function. + """ - # assumes we have two primary arguments that will be handled - # by the cosmo_factor_wrapper def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + + Returns + ------- + out : Callable + The wrapped function. + """ helper_result = _prepare_array_func_args(*args, **kwargs) - ret_cf = cosmo_factor_wrapper(helper_result["cfs"][0], helper_result["cfs"][1]) + ret_cf = cosmo_factor_handler(helper_result["cfs"][0], helper_result["cfs"][1]) res = unyt_func(*helper_result["args"], **helper_result["kwargs"]) if "out" in kwargs: return _return_helper(res, helper_result, ret_cf, out=kwargs["out"]) @@ -516,11 +1085,39 @@ def wrapper(*args, **kwargs): return wrapper -def _default_comparison_wrapper(unyt_func): +def _default_comparison_wrapper(unyt_func: Callable) -> Callable: + """ + Wrapper helper for binary comparison functions with typical behaviour. + + For many numpy and unyt comparison functions with two (main) input arguments, the + wrapping code that we need to apply is repetitive. Just prepare the arguments, + process the ``cosmo_factors``, and attach our cosmo attributes + to the return value. This function facilitates writing these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) comparison function to be wrapped. + + Returns + ------- + out : Callable + The wrapped function. + """ # assumes we have two primary arguments that will be handled with # _comparison_cosmo_factor with them as the inputs def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + + Returns + ------- + out : Callable + The wrapped function. + """ helper_result = _prepare_array_func_args(*args, **kwargs) ret_cf = _comparison_cosmo_factor( helper_result["cfs"][0], helper_result["cfs"][1], inputs=args[:2] @@ -531,12 +1128,38 @@ def wrapper(*args, **kwargs): return wrapper -def _default_oplist_wrapper(unyt_func): +def _default_oplist_wrapper(unyt_func: Callable) -> Callable: + """ + Wrapper helper for functions accepting a list of operands with typical behaviour. + + For many numpy and unyt functions taking a list of operands as an argument, the + wrapping code that we need to apply is repetitive. Just prepare the arguments, + preserve the ``cosmo_factor`` (after checking that it's common across the list), and + attach our cosmo attributes to the return value. This function facilitates writing + these wrappers. + + Can be used as a decorator. + + Parameters + ---------- + unyt_func : Callable + The unyt (or numpy) function to be wrapped. + + Returns + ------- + out : Callable + The wrapped function. + """ - # assumes first argument is a list of operands - # assumes that we always preserve the cosmo factor of the first - # element in the list of operands def wrapper(*args, **kwargs): + """ + Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + + Returns + ------- + out : Callable + The wrapped function. + """ helper_result = _prepare_array_func_args(*args, **kwargs) helper_result_oplist = _prepare_array_func_args(*args[0]) ret_cf = _preserve_cosmo_factor(helper_result_oplist["cfs"][0]) @@ -550,6 +1173,9 @@ def wrapper(*args, **kwargs): return wrapper +# Next we wrap functions from unyt and numpy. There's not much point in writing docstrings +# or type hints for all of these. + # Now we wrap functions that unyt handles explicitly (below that will be those not handled # explicitly): @@ -1515,15 +2141,23 @@ def trapezoid(y, x=None, dx=1.0, axis=-1): # Now we wrap functions that unyt does not handle explicitly: -implements(np.average)(_propagate_cosmo_array_attributes(np.average._implementation)) -implements(np.max)(_propagate_cosmo_array_attributes(np.max._implementation)) -implements(np.min)(_propagate_cosmo_array_attributes(np.min._implementation)) -implements(np.mean)(_propagate_cosmo_array_attributes(np.mean._implementation)) -implements(np.median)(_propagate_cosmo_array_attributes(np.median._implementation)) -implements(np.sort)(_propagate_cosmo_array_attributes(np.sort._implementation)) -implements(np.sum)(_propagate_cosmo_array_attributes(np.sum._implementation)) +implements(np.average)( + _propagate_cosmo_array_attributes_to_result(np.average._implementation) +) +implements(np.max)(_propagate_cosmo_array_attributes_to_result(np.max._implementation)) +implements(np.min)(_propagate_cosmo_array_attributes_to_result(np.min._implementation)) +implements(np.mean)( + _propagate_cosmo_array_attributes_to_result(np.mean._implementation) +) +implements(np.median)( + _propagate_cosmo_array_attributes_to_result(np.median._implementation) +) +implements(np.sort)( + _propagate_cosmo_array_attributes_to_result(np.sort._implementation) +) +implements(np.sum)(_propagate_cosmo_array_attributes_to_result(np.sum._implementation)) implements(np.partition)( - _propagate_cosmo_array_attributes(np.partition._implementation) + _propagate_cosmo_array_attributes_to_result(np.partition._implementation) ) @@ -1532,7 +2166,7 @@ def meshgrid(*xi, **kwargs): # meshgrid is a unique case: arguments never interact with each other, so we don't # want to use our _prepare_array_func_args helper (that will try to coerce to # compatible comoving, cosmo_factor). - # However we can't just use _propagate_cosmo_array_attributes because we need to - # iterate over arguments. + # However we can't just use _propagate_cosmo_array_attributes_to_result because we + # need to iterate over arguments. res = np.meshgrid._implementation(*xi, **kwargs) return tuple(_copy_cosmo_array_attributes(x, r) for (x, r) in zip(xi, res)) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 5df6075c..213c0e2f 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -92,8 +92,8 @@ ) from numpy._core.umath import _ones_like, clip from ._array_functions import ( - _propagate_cosmo_array_attributes, - _ensure_cosmo_array_or_quantity, + _propagate_cosmo_array_attributes_to_result, + _ensure_result_is_cosmo_array_or_quantity, _sqrt_cosmo_factor, _multiply_cosmo_factor, _preserve_cosmo_factor, @@ -113,6 +113,17 @@ a = sympy.symbols("a") +def _verify_valid_transform_validity(obj): + if not obj.valid_transform: + assert ( + not obj.comoving + ), "Cosmo arrays without a valid transform to comoving units must be physical" + if obj.comoving: + assert ( + obj.valid_transform + ), "Comoving cosmo_arrays must be able to be transformed to physical" + + class InvalidConversionError(Exception): def __init__(self, message="Could not convert to comoving coordinates"): self.message = message @@ -563,6 +574,7 @@ def __new__( cosmo_factor: cosmo_factor if bypass_validation is True: + obj = super().__new__( cls, input_array, @@ -572,54 +584,74 @@ def __new__( bypass_validation=bypass_validation, name=name, ) - # dtype, units, registry & name handled by unyt + + # dtype, units, registry & name are handled by unyt obj.comoving = comoving obj.cosmo_factor = cosmo_factor if cosmo_factor is not None else NULL_CF obj.valid_transform = valid_transform obj.compression = compression + return obj + if isinstance(input_array, cosmo_array): - if comoving: - input_array.convert_to_comoving() + + obj = input_array.view(cls) + + # do cosmo_factor first since it can be used in comoving/physical conversion: + if cosmo_factor is not None: + obj.cosmo_factor = cosmo_factor + # else is already copied from input_array + + if comoving is True: + obj.convert_to_comoving() elif comoving is False: - input_array.convert_to_physical() - else: - comoving = input_array.comoving - cosmo_factor = ( - getattr(input_array, "cosmo_factor") - if cosmo_factor is None - else _preserve_cosmo_factor( - cosmo_factor, getattr(input_array, "cosmo_factor") - ) + obj.convert_to_physical() + # else is already copied from input_array + + # only overwrite valid_transform after transforming so that invalid + # transformations raise: + obj.valid_transform = valid_transform + _verify_valid_transform_validity(obj) + + obj.compression = ( + compression if compression is not None else obj.compression ) - if not valid_transform: - input_array.convert_to_physical() - if compression != input_array.compression: - compression = None # just drop it + + return obj + elif isinstance(input_array, np.ndarray) and input_array.dtype != object: + # guard np.ndarray so it doesn't get caught by _iterable in next case - if cosmo_factor is None: - cosmo_factor = NULL_CF - elif _iterable(input_array): + + # ndarray with object dtype goes to next case to properly handle e.g. + # ndarrays containing cosmo_quantities + + pass + + elif _iterable(input_array) and input_array: + + # if _prepare_array_func_args finds cosmo_array input it will convert to: default_cm = comoving if comoving is not None else True + + # coerce any cosmo_array inputs to consistency: 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"] - if cosmo_factor is None: - cosmo_factor = _preserve_cosmo_factor(*helper_result["cfs"]) - elif all([cf is None for cf in helper_result["cfs"]]): - cosmo_factor = cosmo_factor - else: - cosmo_factor = _preserve_cosmo_factor( - cosmo_factor, *helper_result["cfs"] - ) - if not valid_transform: - input_array.convert_to_physical() - if compression != helper_result["compression"]: - compression = None # just drop it + + # default to comoving, cosmo_factor and compression given as kwargs + comoving = helper_result["comoving"] if comoving is None else comoving + cosmo_factor = ( + _preserve_cosmo_factor(*helper_result["cfs"]) + if cosmo_factor is None + else cosmo_factor + ) + compression = ( + helper_result["compression"] if compression is None else compression + ) + # valid_transform has a non-None default, so we have to decide to always + # respect it obj = super().__new__( cls, @@ -634,21 +666,15 @@ def __new__( if isinstance(obj, unyt_array) and not isinstance(obj, cls): obj = obj.view(cls) + # attach our attributes: + obj.comoving = comoving # unyt allows creating a unyt_array from e.g. arrays with heterogenous units - # (it probably shouldn't...), so we don't bother recursing deeply and therefore + # (it probably shouldn't...), so we don't recurse deeply and therefore # can't guarantee that cosmo_factor isn't None at this point, guard with default obj.cosmo_factor = cosmo_factor if cosmo_factor is not None else NULL_CF - obj.comoving = comoving obj.compression = compression obj.valid_transform = valid_transform - if not obj.valid_transform: - assert ( - not obj.comoving - ), "Cosmo arrays without a valid transform to comoving units must be physical" - if obj.comoving: - assert ( - obj.valid_transform - ), "Comoving cosmo_arrays must be able to be transformed to physical" + _verify_valid_transform_validity(obj) return obj @@ -701,32 +727,32 @@ def __setstate__(self, state): # Wrap functions that return copies of cosmo_arrays so that our # attributes get passed through: - 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) - compress = _propagate_cosmo_array_attributes(unyt_array.compress) - diagonal = _propagate_cosmo_array_attributes(unyt_array.diagonal) - flatten = _propagate_cosmo_array_attributes(unyt_array.flatten) - ravel = _propagate_cosmo_array_attributes(unyt_array.ravel) - repeat = _propagate_cosmo_array_attributes(unyt_array.repeat) - swapaxes = _propagate_cosmo_array_attributes(unyt_array.swapaxes) - transpose = _propagate_cosmo_array_attributes(unyt_array.transpose) - view = _propagate_cosmo_array_attributes(unyt_array.view) - - take = _propagate_cosmo_array_attributes( - _ensure_cosmo_array_or_quantity(unyt_array.take) + astype = _propagate_cosmo_array_attributes_to_result(unyt_array.astype) + in_units = _propagate_cosmo_array_attributes_to_result(unyt_array.in_units) + byteswap = _propagate_cosmo_array_attributes_to_result(unyt_array.byteswap) + compress = _propagate_cosmo_array_attributes_to_result(unyt_array.compress) + diagonal = _propagate_cosmo_array_attributes_to_result(unyt_array.diagonal) + flatten = _propagate_cosmo_array_attributes_to_result(unyt_array.flatten) + ravel = _propagate_cosmo_array_attributes_to_result(unyt_array.ravel) + repeat = _propagate_cosmo_array_attributes_to_result(unyt_array.repeat) + swapaxes = _propagate_cosmo_array_attributes_to_result(unyt_array.swapaxes) + transpose = _propagate_cosmo_array_attributes_to_result(unyt_array.transpose) + view = _propagate_cosmo_array_attributes_to_result(unyt_array.view) + + take = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_array.take) ) - reshape = _propagate_cosmo_array_attributes( - _ensure_cosmo_array_or_quantity(unyt_array.reshape) + reshape = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_array.reshape) ) - __getitem__ = _propagate_cosmo_array_attributes( - _ensure_cosmo_array_or_quantity(unyt_array.__getitem__) + __getitem__ = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_array.__getitem__) ) # Also wrap some array "properties": - T = property(_propagate_cosmo_array_attributes(unyt_array.transpose)) - ua = property(_propagate_cosmo_array_attributes(np.ones_like)) - unit_array = property(_propagate_cosmo_array_attributes(np.ones_like)) + T = property(_propagate_cosmo_array_attributes_to_result(unyt_array.transpose)) + ua = property(_propagate_cosmo_array_attributes_to_result(np.ones_like)) + unit_array = property(_propagate_cosmo_array_attributes_to_result(np.ones_like)) def convert_to_comoving(self) -> None: """ @@ -927,7 +953,7 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): else: ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs) - ret = _ensure_cosmo_array_or_quantity(super().__array_ufunc__)( + ret = _ensure_result_is_cosmo_array_or_quantity(super().__array_ufunc__)( ufunc, method, *inputs, **kwargs ) # if we get a tuple we have multiple return values to deal with diff --git a/swiftsimio/visualisation/smoothing_length/generate.py b/swiftsimio/visualisation/smoothing_length/generate.py index cb2ae76b..82eab392 100644 --- a/swiftsimio/visualisation/smoothing_length/generate.py +++ b/swiftsimio/visualisation/smoothing_length/generate.py @@ -7,10 +7,10 @@ from swiftsimio import cosmo_array from swiftsimio.optional_packages import KDTree, TREE_AVAILABLE -from swiftsimio._array_functions import _propagate_cosmo_array_attributes +from swiftsimio._array_functions import _propagate_cosmo_array_attributes_to_result -@_propagate_cosmo_array_attributes # copies attrs of first arg to result, if present +@_propagate_cosmo_array_attributes_to_result # copies attrs of first arg to result def generate_smoothing_lengths( coordinates: cosmo_array, boxsize: cosmo_array, From 3b0fc84996ff86471556796bd0f3d0a69b90f949 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 8 Feb 2025 09:58:57 +0000 Subject: [PATCH 063/125] Start reviewing objects.py docstrings. --- swiftsimio/objects.py | 30 ++++++++++++++++++++++++++---- 1 file changed, 26 insertions(+), 4 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 213c0e2f..b28b6847 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -1,6 +1,12 @@ """ -Contains global objects, e.g. the superclass version of the -unyt_array that we use, called cosmo_array. +Contains classes for our custom :class:`~swiftsimio.objects.cosmo_array`, +:class:`~swiftsimio.objects.cosmo_quantity` and +:class:`~swiftsimio.objects.cosmo_factor` objects for cosmology-aware +arrays, extending the functionality of the :class:`~unyt.array.unyt_array`. + +For developers, see also :mod:`~swiftsimio._array_functions` containing +helpers, wrappers and implementations that enable most :mod:`numpy` and +:mod:`unyt` functions to work with our cosmology-aware arrays. """ from unyt import unyt_array, unyt_quantity @@ -113,7 +119,23 @@ a = sympy.symbols("a") -def _verify_valid_transform_validity(obj): +def _verify_valid_transform_validity(obj: "cosmo_array") -> None: + """ + Checks that ``comoving`` and ``valid_transform`` attributes are compatible. + + Comoving arrays must be able to transform, while arrays that don't transform must + be physical. This function raises if this is not the case. + + Parameters + ---------- + obj : swiftsimio.objects.cosmo_array + The array whose validity is to be checked. + + Raises + ------ + AssertionError + When an invalid combination of ``comoving`` and ``valid_transform`` is found. + """ if not obj.valid_transform: assert ( not obj.comoving @@ -125,7 +147,7 @@ def _verify_valid_transform_validity(obj): class InvalidConversionError(Exception): - def __init__(self, message="Could not convert to comoving coordinates"): + def __init__(self, message="Could not convert to comoving coordinates."): self.message = message From a396554f5136516ca57d4d62b951a765d7901f2a Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 8 Feb 2025 15:55:29 +0000 Subject: [PATCH 064/125] Run black. --- swiftsimio/_array_functions.py | 6 ++-- swiftsimio/conversions.py | 5 +-- swiftsimio/metadata/objects.py | 38 ++++++++++------------- swiftsimio/snapshot_writer.py | 8 ++--- swiftsimio/visualisation/volume_render.py | 4 +-- tests/basic_test.py | 2 +- tests/test_data.py | 8 ++--- tests/test_read_ic.py | 2 +- tests/test_soap.py | 4 +-- tests/test_visualisation.py | 31 ++++++++---------- 10 files changed, 47 insertions(+), 61 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 5371989f..181ba372 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -726,18 +726,18 @@ def _arctan2_cosmo_factor( f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor( diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index 4fa415fe..7eef7108 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -96,13 +96,13 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: # expressions taken directly from astropy, since they do no longer # allow access to these attributes (since version 5.1+) critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value - a_B_c2 = (4.0 * const.sigma_sb / const.c**3).cgs.value + a_B_c2 = (4.0 * const.sigma_sb / const.c ** 3).cgs.value # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy. # This is an exact inversion of the procedure performed in astropy. critical_density_0 = astropy_units.Quantity( critdens_const * H0.to("1/s").value ** 2, - astropy_units.g / astropy_units.cm**3, + astropy_units.g / astropy_units.cm ** 3, ) Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0) @@ -141,6 +141,7 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: m_nu=ap_m_nu, ) + else: def swift_cosmology_to_astropy(cosmo: dict, units) -> dict: diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index aba8b5ca..e3731620 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -113,14 +113,12 @@ def postprocess_header(self): # items including the scale factor. # These must be unpacked as they are stored as length-1 arrays - header_unpack_float_units = ( - metadata.metadata_fields.generate_units_header_unpack_single_float( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) + header_unpack_float_units = metadata.metadata_fields.generate_units_header_unpack_single_float( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, ) for field, names in metadata.metadata_fields.header_unpack_single_float.items(): try: @@ -166,19 +164,15 @@ def postprocess_header(self): self.scale_factor = 1.0 # These are just read straight in to variables - header_unpack_arrays_units = ( - metadata.metadata_fields.generate_units_header_unpack_arrays( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) + header_unpack_arrays_units = metadata.metadata_fields.generate_units_header_unpack_arrays( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, ) - header_unpack_arrays_cosmo_args = ( - metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( - self.scale_factor - ) + header_unpack_arrays_cosmo_args = metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( + self.scale_factor ) for field, name in metadata.metadata_fields.header_unpack_arrays.items(): @@ -610,7 +604,7 @@ def get_units(unit_attribute): # Need to check if the exponent is 0 manually because of float precision unit_exponent = unit_attribute[f"U_{exponent} exponent"][0] if unit_exponent != 0.0: - units *= unit**unit_exponent + units *= unit ** unit_exponent except KeyError: # Can't load that data! # We should probably warn the user here... @@ -704,7 +698,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a**cosmo_exponent + a_factor_this_dataset = a ** cosmo_exponent return cosmo_factor(a_factor_this_dataset, current_scale_factor) diff --git a/swiftsimio/snapshot_writer.py b/swiftsimio/snapshot_writer.py index fbb6362d..388ead1b 100644 --- a/swiftsimio/snapshot_writer.py +++ b/swiftsimio/snapshot_writer.py @@ -163,11 +163,7 @@ def check_consistent(self) -> bool: return True - def generate_smoothing_lengths( - self, - boxsize: cosmo_array, - dimension: int, - ): + def generate_smoothing_lengths(self, boxsize: cosmo_array, dimension: int): """ Automatically generates the smoothing lengths as 2 * the mean interparticle separation. @@ -272,7 +268,7 @@ def get_attributes(self, scale_factor: float) -> dict: # Find the scale factor associated quantities a_exp = a_exponents.get(name, 0) - a_factor = scale_factor**a_exp + a_factor = scale_factor ** a_exp attributes_dict[output_handle] = { "Conversion factor to CGS (not including cosmological corrections)": [ diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 724ed87c..3430cb0d 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -844,14 +844,14 @@ def render_gas( * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) diff --git a/tests/basic_test.py b/tests/basic_test.py index 24fb9295..adb4e7fc 100644 --- a/tests/basic_test.py +++ b/tests/basic_test.py @@ -23,7 +23,7 @@ def test_write(): x = Writer("galactic", boxsize) # 32^3 particles. - n_p = 32**3 + n_p = 32 ** 3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) diff --git a/tests/test_data.py b/tests/test_data.py index 25de5eaa..44a69f1b 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -165,7 +165,7 @@ def test_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), ) start_offset = offsets @@ -212,7 +212,7 @@ def test_dithered_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), ) offsets = mask_region.offsets["dark_matter"] counts = mask_region.counts["dark_matter"] @@ -256,7 +256,7 @@ def test_reading_select_region_metadata(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.2, boxsize * 0.8]).T @@ -308,7 +308,7 @@ def test_reading_select_region_metadata_not_spatial_only(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.26, boxsize * 0.74]).T diff --git a/tests/test_read_ic.py b/tests/test_read_ic.py index 112d4ffb..d1f39f5c 100644 --- a/tests/test_read_ic.py +++ b/tests/test_read_ic.py @@ -34,7 +34,7 @@ def simple_snapshot_data(): x = Writer(cosmo_units, boxsize) # 32^3 particles. - n_p = 32**3 + n_p = 32 ** 3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) diff --git a/tests/test_soap.py b/tests/test_soap.py index 3469ebf8..1d59768c 100644 --- a/tests/test_soap.py +++ b/tests/test_soap.py @@ -47,13 +47,13 @@ def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): 1e5, "Msun", comoving=True, - cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 0, this_mask.metadata.scale_factor), ) upper = cosmo_quantity( 1e13, "Msun", comoving=True, - cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 0, this_mask.metadata.scale_factor), ) this_mask.constrain_mask( "spherical_overdensity_200_mean", "total_mass", lower, upper diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 4951d3b1..e35eda9d 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -6,10 +6,7 @@ project_gas, project_pixel_grid, ) -from swiftsimio.visualisation.slice import ( - slice_scatter_parallel, - slice_gas, -) +from swiftsimio.visualisation.slice import slice_scatter_parallel, slice_gas from swiftsimio.visualisation.volume_render import render_gas from swiftsimio.visualisation.volume_render import scatter as volume_scatter from swiftsimio.visualisation.power_spectrum import ( @@ -20,9 +17,7 @@ ) from swiftsimio.visualisation.ray_trace import panel_gas from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths -from swiftsimio.visualisation.projection_backends import ( - backends as projection_backends, -) +from swiftsimio.visualisation.projection_backends import backends as projection_backends from swiftsimio.visualisation.slice_backends import ( backends as slice_backends, backends_parallel as slice_backends_parallel, @@ -84,7 +79,7 @@ def test_scatter_mass_conservation(): for resolution in resolutions: image = scatter(x, y, m, h, resolution, 1.0, 1.0) - mass_in_image = image.sum() / (resolution**2) + mass_in_image = image.sum() / (resolution ** 2) # Check mass conservation to 5% assert np.isclose(mass_in_image, total_mass, 0.05) @@ -361,7 +356,7 @@ def test_comoving_versus_physical(filename): img = func(data, resolution=256, project="masses") # check that we get a physical result assert not img.comoving - assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 + assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 # densities are still compatible with physical img = func(data, resolution=256, project="densities") assert not img.comoving @@ -423,14 +418,14 @@ def test_panel_rendering(filename): plt.imsave( "panels_added.png", - plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1))), + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), ) projected = project_gas(data, res, "masses", backend="renormalised") plt.imsave( "projected.png", - plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)), + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), ) fullstack = np.zeros((res, res)) @@ -630,8 +625,8 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix**3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix**3).to("Msun / kpc**3").v + mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v + mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v mean_density_calculated = ( (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) .to("Msun / kpc**3") @@ -656,15 +651,15 @@ def test_dark_matter_power_spectrum(filename, save=False): min_k = cosmo_quantity( 1e-2, - unyt.Mpc**-1, + unyt.Mpc ** -1, comoving=True, - cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), ) max_k = cosmo_quantity( 1e2, - unyt.Mpc**-1, + unyt.Mpc ** -1, comoving=True, - cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), ) bins = np.geomspace(min_k, max_k, 32) @@ -701,7 +696,7 @@ def test_dark_matter_power_spectrum(filename, save=False): folds[folding] = deposition - folding_output[2**folding] = (k, power_spectrum, scatter) + folding_output[2 ** folding] = (k, power_spectrum, scatter) # Now try doing them all together at once. From 1b89f1cde9eef92c1b576c7b8a22e0e846eb7c54 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 8 Feb 2025 16:03:31 +0000 Subject: [PATCH 065/125] Edit a few more docstrings. --- swiftsimio/objects.py | 53 +++++++++++++++++++++++++++++++++---------- 1 file changed, 41 insertions(+), 12 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 43a0cf67..774231cb 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -147,7 +147,26 @@ def _verify_valid_transform_validity(obj: "cosmo_array") -> None: class InvalidConversionError(Exception): - def __init__(self, message="Could not convert to comoving coordinates."): + """ + Raised when converting from comoving from physical to comoving is not allowed. + + Parameters + ---------- + message : str, optional + Message to print in case of invalid conversion. + """ + + def __init__( + self, message: str = "Could not convert to comoving coordinates." + ) -> None: + """ + Constructor for warning of invalid conversion. + + Parameters + ---------- + message : str, optional + Message to print in case of invalid conversion. + """ self.message = message @@ -155,23 +174,32 @@ class InvalidScaleFactor(Exception): """ Raised when a scale factor is invalid, such as when adding two cosmo_factors with inconsistent scale factors. + + Parameters + ---------- + message : str, optional + Message to print in case of invalid scale factor. """ - def __init__(self, message=None, *args): + def __init__(self, message: str = None, *args) -> None: """ - Constructor for warning of invalid scale factor + Constructor for warning of invalid scale factor. Parameters ---------- - message : str, optional - Message to print in case of invalid scale factor + Message to print in case of invalid scale factor. """ self.message = message def __str__(self): """ - Print warning message of invalid scale factor + Print warning message for invalid scale factor. + + Returns + ------- + out : str + The error message. """ return f"InvalidScaleFactor: {self.message}" @@ -180,6 +208,11 @@ class InvalidSnapshot(Exception): """ Generated when a snapshot is invalid (e.g. you are trying to partially load a sub-snapshot). + + Parameters + ---------- + message : str, optional + Message to print in case of invalid snapshot. """ def __init__(self, message=None, *args): @@ -188,7 +221,6 @@ def __init__(self, message=None, *args): Parameters ---------- - message : str, optional Message to print in case of invalid snapshot """ @@ -207,7 +239,7 @@ class cosmo_factor: comoving and physical coordinates. This takes the expected exponent of the array that can be parsed - by sympy, and the current value of the cosmological scale factor a. + by sympy, and the current value of the cosmological scale factor ``a``. This should be given as the conversion from comoving to physical, i.e. @@ -215,7 +247,6 @@ class cosmo_factor: Examples -------- - Typically this would make cosmo_factor = a for the conversion between comoving positions r' and physical co-ordinates r. @@ -227,11 +258,10 @@ class cosmo_factor: def __init__(self, expr, scale_factor): """ - Constructor for cosmology factor class + Constructor for cosmology factor class. Parameters ---------- - expr : sympy.expr expression used to convert between comoving and physical coordinates scale_factor : float @@ -247,7 +277,6 @@ def __str__(self): Returns ------- - str string to print exponent and current scale factor """ From d9fe3c9c125e1d1027122d45e5e3909872d0afe3 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 8 Feb 2025 16:14:25 +0000 Subject: [PATCH 066/125] Fix warnings from neutrino cosmology. --- swiftsimio/conversions.py | 9 ++++++--- 1 file changed, 6 insertions(+), 3 deletions(-) diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index 7eef7108..356a964f 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -45,12 +45,15 @@ def swift_neutrinos_to_astropy(N_eff, N_ur, M_nu_eV, deg_nu): raise AttributeError( "SWIFTsimIO uses astropy, which cannot handle this cosmological model." ) - if not int(N_eff) == deg_nu.astype(int).sum() + int(N_ur): + print(N_eff, deg_nu, N_ur) + if not int(N_eff) == deg_nu.astype(int).sum() + int(N_ur.squeeze()): raise AttributeError( - "SWIFTsimIO uses astropy, which cannot handle this cosmological model." + "SWIFTSimIO uses astropy, which cannot handle this cosmological model." ) ap_m_nu = [[m] * int(d) for m, d in zip(M_nu_eV, deg_nu)] # replicate - ap_m_nu = sum(ap_m_nu, []) + [0.0] * int(N_ur) # flatten + add massless + ap_m_nu = sum(ap_m_nu, []) + [0.0] * int( + N_ur.squeeze() + ) # flatten + add massless ap_m_nu = np.array(ap_m_nu) * astropy_units.eV return ap_m_nu From 489d60c86a4e91af50f9c6401971efb5fd90e1c3 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 9 Feb 2025 12:24:15 +0000 Subject: [PATCH 067/125] Use squeeze safely. --- swiftsimio/conversions.py | 4 +- swiftsimio/objects.py | 77 ++++++++++++++++++++++++--------------- 2 files changed, 50 insertions(+), 31 deletions(-) diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index 356a964f..85917d91 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -46,13 +46,13 @@ def swift_neutrinos_to_astropy(N_eff, N_ur, M_nu_eV, deg_nu): "SWIFTsimIO uses astropy, which cannot handle this cosmological model." ) print(N_eff, deg_nu, N_ur) - if not int(N_eff) == deg_nu.astype(int).sum() + int(N_ur.squeeze()): + if not int(N_eff) == deg_nu.astype(int).sum() + int(np.squeeze(N_ur)): raise AttributeError( "SWIFTSimIO uses astropy, which cannot handle this cosmological model." ) ap_m_nu = [[m] * int(d) for m, d in zip(M_nu_eV, deg_nu)] # replicate ap_m_nu = sum(ap_m_nu, []) + [0.0] * int( - N_ur.squeeze() + np.squeeze(N_ur) ) # flatten + add massless ap_m_nu = np.array(ap_m_nu) * astropy_units.eV return ap_m_nu diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 774231cb..570851e3 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -6,7 +6,7 @@ For developers, see also :mod:`~swiftsimio._array_functions` containing helpers, wrappers and implementations that enable most :mod:`numpy` and -:mod:`unyt` functions to work with our cosmology-aware arrays. +:mod:`unyt` functions to work with our cosmology-aware arrays. """ from unyt import unyt_array, unyt_quantity @@ -161,7 +161,7 @@ def __init__( ) -> None: """ Constructor for warning of invalid conversion. - + Parameters ---------- message : str, optional @@ -239,23 +239,48 @@ class cosmo_factor: comoving and physical coordinates. This takes the expected exponent of the array that can be parsed - by sympy, and the current value of the cosmological scale factor ``a``. + by :mod:`sympy`, and the current value of the cosmological scale factor ``a``. This should be given as the conversion from comoving to physical, i.e. + :math:`r = a^f \times r` where :math:`a` is the scale factor, + :math:`r` is a physical quantity and :math`r'` a comoving quantity. - r = cosmo_factor * r' with r in physical and r' comoving + Parameters + ---------- + expr : sympy.expr + Expression used to convert between comoving and physical coordinates. + scale_factor : float + The scale factor (a). + Attributes + ---------- + expr : sympy.expr + Expression used to convert between comoving and physical coordinates. + scale_factor : float + The scale factor (a). + Examples -------- - Typically this would make cosmo_factor = a for the conversion between - comoving positions r' and physical co-ordinates r. + Mass density transforms as :math:`a^3`. To set up a ``cosmo_factor``, supposing + a current ``scale_factor=0.97``, we import the scale factor ``a`` and initialize + as: - To do this, use the a imported from objects multiplied as you'd like: + :: - ``density_cosmo_factor = cosmo_factor(a**3, scale_factor=0.97)`` + from swiftsimio.objects import a # the scale factor (a sympy symbol object) + density_cosmo_factor = cosmo_factor(a**3, scale_factor=0.97) + :class:`~swiftsimio.objects.cosmo_factor` support arithmetic, for example: + + :: + + >>> cosmo_factor(a**2, scale_factor=0.5) * cosmo_factor(a**-1, scale_factor=0.5) + cosmo_factor(expr=a, scale_factor=0.5) """ + expr: sympy.expr + scale_factor: float + def __init__(self, expr, scale_factor): """ Constructor for cosmology factor class. @@ -263,9 +288,9 @@ def __init__(self, expr, scale_factor): Parameters ---------- expr : sympy.expr - expression used to convert between comoving and physical coordinates + Expression used to convert between comoving and physical coordinates. scale_factor : float - the scale factor of the simulation data + The scale factor (a). """ self.expr = expr self.scale_factor = scale_factor @@ -273,27 +298,26 @@ def __init__(self, expr, scale_factor): def __str__(self): """ - Print exponent and current scale factor + Print exponent and current scale factor. Returns ------- - str - string to print exponent and current scale factor + out : str + String with exponent and current scale factor. """ return str(self.expr) + f" at a={self.scale_factor}" @property def a_factor(self): """ - The a-factor for the unit. + The multiplicative factor for conversion from comoving to physical. - e.g. for density this is 1 / a**3. + For example, for density this is :math:`a^{-3}`. Returns ------- - - float - the a-factor for given unit + out : float + The multiplicative factor for conversion from comoving to physical. """ if (self.expr is None) or (self.scale_factor is None): return None @@ -302,20 +326,15 @@ def a_factor(self): @property def redshift(self): """ - Compute the redshift from the scale factor. + The redshift computed from the scale factor. + + Returns the redshift :math:`z = \frac{1}{a} - 1`, where :math:`a` is the scale + factor. Returns ------- - - float - redshift from the given scale factor - - Notes - ----- - - Returns the redshift - ..math:: z = \\frac{1}{a} - 1, - where :math: `a` is the scale factor + out : float + The redshift. """ if self.scale_factor is None: return None From 9c8eb0a6846db0427892b2b9944033917a9b7883 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 9 Feb 2025 14:27:37 +0000 Subject: [PATCH 068/125] Review type hints and docstrings in objects.py. --- swiftsimio/objects.py | 962 +++++++++++++++++++++++++++++++++--------- 1 file changed, 765 insertions(+), 197 deletions(-) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 570851e3..8edf68d5 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -9,9 +9,12 @@ :mod:`unyt` functions to work with our cosmology-aware arrays. """ +import unyt from unyt import unyt_array, unyt_quantity from unyt.array import multiple_output_operators, _iterable, POWER_MAPPING from numbers import Number as numeric_type +from typing import Iterable, Union, Tuple, Callable +from collections.abc import Collection import sympy import numpy as np @@ -115,6 +118,16 @@ _prepare_array_func_args, ) +try: + import pint +except ImportError: + pass # only for type hinting + +try: + import astropy.units +except ImportError: + pass # only for type hinting + # The scale factor! a = sympy.symbols("a") @@ -215,7 +228,7 @@ class InvalidSnapshot(Exception): Message to print in case of invalid snapshot. """ - def __init__(self, message=None, *args): + def __init__(self, message: str = None, *args) -> None: """ Constructor for warning of invalid snapshot @@ -226,14 +239,14 @@ def __init__(self, message=None, *args): """ self.message = message - def __str__(self): + def __str__(self) -> str: """ Print warning message of invalid snapshot """ return f"InvalidSnapshot: {self.message}" -class cosmo_factor: +class cosmo_factor(object): """ Cosmology factor class for storing and computing conversion between comoving and physical coordinates. @@ -247,7 +260,7 @@ class cosmo_factor: Parameters ---------- - expr : sympy.expr + expr : sympy.Expr Expression used to convert between comoving and physical coordinates. scale_factor : float The scale factor (a). @@ -258,7 +271,7 @@ class cosmo_factor: Expression used to convert between comoving and physical coordinates. scale_factor : float The scale factor (a). - + Examples -------- Mass density transforms as :math:`a^3`. To set up a ``cosmo_factor``, supposing @@ -278,10 +291,10 @@ class cosmo_factor: cosmo_factor(expr=a, scale_factor=0.5) """ - expr: sympy.expr + expr: sympy.Expr scale_factor: float - def __init__(self, expr, scale_factor): + def __init__(self, expr: sympy.Expr, scale_factor: float) -> None: """ Constructor for cosmology factor class. @@ -296,7 +309,7 @@ def __init__(self, expr, scale_factor): self.scale_factor = scale_factor pass - def __str__(self): + def __str__(self) -> str: """ Print exponent and current scale factor. @@ -308,7 +321,7 @@ def __str__(self): return str(self.expr) + f" at a={self.scale_factor}" @property - def a_factor(self): + def a_factor(self) -> float: """ The multiplicative factor for conversion from comoving to physical. @@ -324,7 +337,7 @@ def a_factor(self): return float(self.expr.subs(a, self.scale_factor)) @property - def redshift(self): + def redshift(self) -> float: """ The redshift computed from the scale factor. @@ -340,7 +353,31 @@ def redshift(self): return None return (1.0 / self.scale_factor) - 1.0 - def __add__(self, b): + def __add__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Add two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to add to this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The sum of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be summed is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only add cosmo_factor to another cosmo_factor.") if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to add two cosmo_factors with different scale factors " @@ -355,7 +392,34 @@ def __add__(self, b): return cosmo_factor(expr=self.expr, scale_factor=self.scale_factor) - def __sub__(self, b): + def __sub__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Subtract two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to subtract from this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The difference of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be subtracted is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError( + "Can only subtract cosmo_factor from another cosmo_factor." + ) if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to subtract two cosmo_factors with different scale factors " @@ -370,7 +434,34 @@ def __sub__(self, b): return cosmo_factor(expr=self.expr, scale_factor=self.scale_factor) - def __mul__(self, b): + def __mul__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Multiply two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to multiply this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be multiplied is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError( + "Can only multiply cosmo_factor with another cosmo_factor." + ) if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to multiply two cosmo_factors with different scale factors " @@ -390,7 +481,31 @@ def __mul__(self, b): return cosmo_factor(expr=self.expr * b.expr, scale_factor=self.scale_factor) - def __truediv__(self, b): + def __truediv__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Divide two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to divide this one by. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The quotient of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to divide by is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only divide cosmo_factor with another cosmo_factor.") if not self.scale_factor == b.scale_factor: raise InvalidScaleFactor( "Attempting to divide two cosmo_factors with different scale factors " @@ -410,90 +525,410 @@ def __truediv__(self, b): return cosmo_factor(expr=self.expr / b.expr, scale_factor=self.scale_factor) - def __radd__(self, b): + def __radd__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Add two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to add to this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The sum of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be summed is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return self.__add__(b) - def __rsub__(self, b): + def __rsub__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Subtract two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to subtract from this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The difference of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be subtracted is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return self.__sub__(b) - def __rmul__(self, b): + def __rmul__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Multiply two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to multiply this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The product of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to be multiplied is not a + :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return self.__mul__(b) - def __rtruediv__(self, b): + def __rtruediv__(self, b: "cosmo_factor") -> "cosmo_factor": + """ + Divide two :class:`~swiftsimio.objects.cosmo_factor`s. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to divide this one by. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The quotient of the two :class:`~swiftsimio.objects.cosmo_factor`s. + + Raises + ------ + ValueError + If the object to divide by is not a :class:`~swiftsimio.objects.cosmo_factor`. + + swiftsimio.objects.InvalidScaleFactor + If the :class:`~swiftsimio.objects.cosmo_factor` has a ``scale_factor`` that + does not match this one's. + """ return b.__truediv__(self) - def __pow__(self, p): + def __pow__(self, p: float) -> "cosmo_factor": + """ + Raise this :class:`~swiftsimio.objects.cosmo_factor` to an exponent. + + Parameters + ---------- + p : float + The exponent by which to raise this :class:`~swiftsimio.objects.cosmo_factor`. + + Returns + ------- + out : swiftsimio.objects.cosmo_factor + The exponentiated :class:`~swiftsimio.objects.cosmo_factor`s. + """ if self.expr is None: return cosmo_factor(expr=None, scale_factor=self.scale_factor) return cosmo_factor(expr=self.expr ** p, scale_factor=self.scale_factor) - def __lt__(self, b): + def __lt__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") return self.a_factor < b.a_factor - def __gt__(self, b): + def __gt__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") return self.a_factor > b.a_factor - def __le__(self, b): + def __le__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") return self.a_factor <= b.a_factor - def __ge__(self, b): + def __ge__(self, b: "cosmo_factor") -> bool: + """ + Compare the values of two :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") return self.a_factor >= b.a_factor - def __eq__(self, b): - # Doesn't handle some corner cases, e.g. cosmo_factor(a ** 1, scale_factor=1) - # is considered equal to cosmo_factor(a ** 2, scale_factor=1) because - # 1 ** 1 == 1 ** 2. Should check self.expr vs b.expr with sympy? + def __eq__(self, b: "cosmo_factor") -> bool: + """ + Compare the expressions and values of two + :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. Notice that unlike ``__gt__``, + ``__ge__``, ``__lt__`` and ``__le__``, (in)equality comparisons check that the + expression is (un)equal as well as the value. This is so that e.g. ``a**1`` and + ``a**2``, both with ``scale_factor=1.0`` are not equal (both have + ``a_factor==1``). + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ + if not isinstance(b, cosmo_factor): + raise ValueError("Can only compare cosmo_factor with another cosmo_factor.") return (self.scale_factor == b.scale_factor) and (self.a_factor == b.a_factor) - def __ne__(self, b): + def __ne__(self, b: "cosmo_factor") -> bool: + """ + Compare the expressions and values of two + :meth:`~swiftsimio.objects.cosmo_factor.a_factor`s. + + The :meth:`~swiftsimio.objects.cosmo_factor.a_factor` is the ``expr`` attribute + evaluated given the ``scale_factor`` attribute. Notice that unlike ``__gt__``, + ``__ge__``, ``__lt__`` and ``__le__``, (in)equality comparisons check that the + expression is (un)equal as well as the value. This is so that e.g. ``a**1`` and + ``a**2``, both with ``scale_factor=1.0`` are not equal (both have + ``a_factor==1``). + + Parameters + ---------- + b : swiftsimio.objects.cosmo_factor + The :class:`~swiftsimio.objects.cosmo_factor` to compare with this one. + + Returns + ------- + out : bool + The result of the comparison. + + Raises + ------ + ValueError + If the object to compare is not a :class:`~swiftsimio.objects.cosmo_factor`. + """ return not self.__eq__(b) - def __repr__(self): + def __repr__(self) -> str: """ - Print exponent and current scale factor + Get a string representation of the scaling with the scale factor. Returns ------- - - str - string to print exponent and current scale factor + out : str + String representation of the scaling with the scale factor. """ return f"cosmo_factor(expr={self.expr}, scale_factor={self.scale_factor})" -NULL_CF = cosmo_factor(None, None) +NULL_CF = cosmo_factor(None, None) # helps avoid name collisions with kwargs below class cosmo_array(unyt_array): """ Cosmology array class. - This inherits from the unyt.unyt_array, and adds - four variables: compression, cosmo_factor, comoving, and valid_transform. + This inherits from the :class:`~unyt.array.unyt_array`, and adds + four attributes: ``compression``, ``cosmo_factor``, ``comoving``, and + ``valid_transform``. + + .. note:: + + :class:`~swiftsimio.objects.cosmo_array` and the related + :class:`~swiftsimio.objects.cosmo_quantity` are now intended to support all + :mod:`numpy` functions, propagating units (thanks to :mod:`unyt`) and + cosmology information. There are a large number of functions, and a very + large number of possible parameter combinations, so some corner cases may + have been missed in testing. Please report any issues on github, they are + usually easy to fix for future use! Currently :mod:`scipy` functions are + not supported (although some might "just work"). Requests to fully support + specific functions can also be submitted as github issues. Parameters ---------- - - unyt_array : unyt.unyt_array - the inherited unyt_array + input_array : np.ndarray, unyt.array.unyt_array or iterable + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, 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 ``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 or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. 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 : swiftsimio.objects.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. If ``False``, + then ``comoving`` must be ``False``. + compression : string + Description of the compression filters that were applied to that array in the + hdf5 file. Attributes ---------- - comoving : bool - if True then the array is in comoving co-ordinates, and if - False then it is in physical units. + If ``True`` then the array is in comoving coordinates, if``False`` then it is in + physical units. - cosmo_factor : float - Object to store conversion data between comoving and physical coordinates + cosmo_factor : swiftsimio.objects.cosmo_factor + 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 + If ``True`` then the array can be converted from physical to comoving units. + + Notes + ----- + This class will generally try to make sense of input and initialize an array-like + object consistent with the input, and warn or raise if this cannot be done + consistently. However, the way that :class:`~unyt.array.unyt_array` handles input + imposes some limits to this. In particular, nested non-numpy containers given in + input are not traversed recursively, but only one level deep. This means that + while with this input the attributes are detected by the new array correctly: + + :: + + >>> from swiftsimio.objects import cosmo_array, cosmo_factor, a + >>> x = cosmo_array( + ... np.arange(3), + ... u.kpc, + ... comoving=True, + ... cosmo_factor=cosmo_factor(a**1, scale_factor=1.0) + ... ) + >>> cosmo_array([x, x]) + cosmo_array([[0, 1, 2], + [0, 1, 2]], 'kpc', comoving='True', cosmo_factor='a at a=1.0', + valid_transform='True') + + with this input they are lost: + + :: + + >>> cosmo_array([[x, x],[x, x]]) + cosmo_array([[[0, 1, 2],[0, 1, 2]],[[0, 1, 2],[0, 1, 2]]], + '(dimensionless)', comoving='None', cosmo_factor='None at a=None', + valid_transform='True') + See Also + -------- + swiftsimio.objects.cosmo_quantity """ _cosmo_factor_ufunc_registry = { @@ -589,56 +1024,53 @@ class cosmo_array(unyt_array): def __new__( cls, - input_array, - units=None, - registry=None, - dtype=None, - bypass_validation=False, - input_units=None, - name=None, - cosmo_factor=None, - comoving=None, - valid_transform=True, - compression=None, - ): + input_array: Iterable, + units: Union[str, unyt.unit_object.Unit, "astropy.units.core.Unit"] = None, + registry: unyt.unit_registry.UnitRegistry = None, + dtype: Union[np.dtype, str] = None, + bypass_validation: bool = False, + name: str = None, + cosmo_factor: cosmo_factor = None, + comoving: bool = None, + valid_transform: bool = True, + compression: str = None, + ) -> "cosmo_array": """ - Essentially a copy of the __new__ constructor. + Closely inspired by the :meth:`unyt.array.unyt_array.__new__` constructor. Parameters ---------- - input_array : iterable - A tuple, list, or array to attach units 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). + input_array : np.ndarray, unyt.array.unyt_array or iterable + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, 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. + The registry to create units from. If ``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 + none is found, uses ``np.float64``. bypass_validation : bool, optional - If True, all input validation is skipped. Using this option may produce + If ``True``, all input validation is skipped. Using this option may produce corrupted or invalid data, but can lead to significant speedups in the input validation logic adds significant overhead. If set, minimally - pass valid values for units, comoving and cosmo_factor. Defaults to False. - input_units : str, optional - deprecated in favour of units option + pass valid values for units, comoving and cosmo_factor. Defaults to ``False``. name : str, optional - The name of the array. Defaults to None. This attribute does not propagate + 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 + cosmo_factor : swiftsimio.objects.cosmo_factor + Object to store conversion data between comoving and physical coordinates. comoving : bool - flag to indicate whether using comoving coordinates + Flag to indicate whether using comoving coordinates. valid_transform : bool - flag to indicate whether this array can be converted to comoving + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. compression : string - description of the compression filters that were applied to that array in the - hdf5 file + Description of the compression filters that were applied to that array in the + hdf5 file. """ cosmo_factor: cosmo_factor @@ -748,7 +1180,7 @@ def __new__( return obj - def __array_finalize__(self, obj): + def __array_finalize__(self, obj: "cosmo_array") -> None: super().__array_finalize__(obj) if obj is None: return @@ -757,7 +1189,7 @@ def __array_finalize__(self, obj): self.compression = getattr(obj, "compression", None) self.valid_transform = getattr(obj, "valid_transform", True) - def __str__(self): + def __str__(self) -> str: if self.comoving: comoving_str = "(Comoving)" elif self.comoving is None: @@ -767,15 +1199,20 @@ def __str__(self): return super().__str__() + " " + comoving_str - def __repr__(self): + def __repr__(self) -> str: return super().__repr__() - def __reduce__(self): + def __reduce__(self) -> tuple: """ - Pickle reduction method + Pickle reduction method. - Here we add an extra element at the start of the unyt_array state - tuple to store the cosmology info. + Here we add an extra element at the start of the :class:`~unyt.array.unyt_array` + state tuple to store the cosmology info. + + Returns + ------- + out : tuple + The state ready for pickling. """ np_ret = super(cosmo_array, self).__reduce__() obj_state = np_ret[2] @@ -785,12 +1222,17 @@ def __reduce__(self): new_ret = np_ret[:2] + cosmo_state + np_ret[3:] return new_ret - def __setstate__(self, state): + def __setstate__(self, state: Tuple) -> None: """ - Pickle setstate method + Pickle setstate method. Here we extract the extra cosmology info we added to the object - state and pass the rest to unyt_array.__setstate__. + state and pass the rest to :meth:`unyt.array.unyt_array.__setstate__`. + + Parameters + ---------- + state : tuple + A :obj:`tuple` containing the extra state information. """ super(cosmo_array, self).__setstate__(state[1:]) self.cosmo_factor, self.comoving, self.valid_transform = state[0] @@ -828,7 +1270,7 @@ def __setstate__(self, state): def convert_to_comoving(self) -> None: """ - Convert the internal data to be in comoving units. + Convert the internal data in-place to be in comoving units. """ if self.comoving: return @@ -842,7 +1284,7 @@ def convert_to_comoving(self) -> None: def convert_to_physical(self) -> None: """ - Convert the internal data to be in physical units. + Convert the internal data in-place to be in physical units. """ if self.comoving is None: raise InvalidConversionError @@ -855,28 +1297,28 @@ def convert_to_physical(self) -> None: values *= self.cosmo_factor.a_factor self.comoving = False - def to_physical(self): + def to_physical(self) -> "cosmo_array": """ Creates a copy of the data in physical units. Returns ------- - cosmo_array - copy of cosmo_array in physical units + out : swiftsimio.objects.cosmo_array + Copy of this array in physical units. """ copied_data = self.in_units(self.units, cosmo_factor=self.cosmo_factor) copied_data.convert_to_physical() return copied_data - def to_comoving(self): + def to_comoving(self) -> "cosmo_array": """ Creates a copy of the data in comoving units. Returns ------- - cosmo_array - copy of cosmo_array in comoving units + out : swiftsimio.objects.cosmo_array + Copy of this array in comoving units """ if not self.valid_transform: raise InvalidConversionError @@ -885,60 +1327,82 @@ def to_comoving(self): return copied_data - def compatible_with_comoving(self): + def compatible_with_comoving(self) -> bool: """ - Is this cosmo_array compatible with a comoving cosmo_array? + Is this :class:`~swiftsimio.objects.cosmo_array` compatible with a comoving + :class:`~swiftsimio.objects.cosmo_array`? + + This is the case if the :class:`~swiftsimio.objects.cosmo_array` is comoving, or + if the scale factor exponent is 0, or the scale factor is 1 + (either case satisfies ``cosmo_factor.a_factor() == 1``). - This is the case if the cosmo_array is comoving, or if the scale factor - exponent is 0 (cosmo_factor.a_factor() == 1) + Returns + ------- + out : bool + ``True`` if compatible, ``False`` otherwise. """ return self.comoving or (self.cosmo_factor.a_factor == 1.0) - def compatible_with_physical(self): + def compatible_with_physical(self) -> bool: """ - Is this cosmo_array compatible with a physical cosmo_array? + Is this :class:`~swiftsimio.objects.cosmo_array` compatible with a physical + :class:`~swiftsimio.objects.cosmo_array`? - This is the case if the cosmo_array is physical, or if the scale factor - exponent is 0 (cosmo_factor.a_factor == 1) + This is the case if the :class:`~swiftsimio.objects.cosmo_array` is physical, or + if the scale factor exponent is 0, or the scale factor is 1 + (either case satisfies ``cosmo_factor.a_factor() == 1``). + + Returns + ------- + out : bool + ``True`` if compatible, ``False`` otherwise. """ return (not self.comoving) or (self.cosmo_factor.a_factor == 1.0) @classmethod def from_astropy( cls, - arr, - unit_registry=None, - comoving=None, - cosmo_factor=cosmo_factor(None, None), - compression=None, - valid_transform=True, - ): + arr: "astropy.units.quantity.Quantity", + unit_registry: unyt.unit_registry.UnitRegistry = None, + comoving: bool = None, + cosmo_factor: cosmo_factor = cosmo_factor(None, None), + compression: str = None, + valid_transform: bool = True, + ) -> "cosmo_array": """ - Convert an AstroPy "Quantity" to a cosmo_array. + Convert an :class:`astropy.units.quantity.Quantity` to a + :class:`~swiftsimio.objects.cosmo_array`. Parameters ---------- - arr: AstroPy Quantity - The Quantity to convert from. - unit_registry: yt UnitRegistry, optional - A yt unit registry to use in the conversion. If one is not supplied, the + arr: astropy.units.quantity.Quantity + The quantity to convert from. + unit_registry : unyt.unit_registry.UnitRegistry, optional + A unyt registry to use in the conversion. If one is not supplied, the default one will be used. 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 + Flag to indicate whether using comoving coordinates. + cosmo_factor : swiftsimio.objects.cosmo_factor + 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. + Description of the compression filters that were applied to that array in the + hdf5 file. valid_transform : bool - flag to indicate whether this array can be converted to comoving + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + A cosmology-aware array. Example ------- - >>> from astropy.units import kpc - >>> cosmo_array.from_astropy([1, 2, 3] * kpc) - cosmo_array([1., 2., 3.], 'kpc') + :: + + >>> from astropy.units import kpc + >>> cosmo_array.from_astropy([1, 2, 3] * kpc) + cosmo_array([1., 2., 3.], 'kpc') """ obj = super().from_astropy(arr, unit_registry=unit_registry).view(cls) @@ -952,46 +1416,54 @@ def from_astropy( @classmethod def from_pint( cls, - arr, - unit_registry=None, - comoving=None, - cosmo_factor=cosmo_factor(None, None), - compression=None, - valid_transform=True, - ): + arr: "pint.registry.Quantity", + unit_registry: unyt.unit_registry.UnitRegistry = None, + comoving: bool = None, + cosmo_factor: cosmo_factor = cosmo_factor(None, None), + compression: str = None, + valid_transform: bool = True, + ) -> "cosmo_array": """ - Convert a Pint "Quantity" to a cosmo_array. + Convert a :class:`pint.registry.Quantity` to a + :class:`~swiftsimio.objects.cosmo_array`. Parameters ---------- - arr : Pint Quantity - The Quantity to convert from. - unit_registry : yt UnitRegistry, optional - A yt unit registry to use in the conversion. If one is not - supplied, the default one will be used. + arr: pint.registry.Quantity + The quantity to convert from. + unit_registry : unyt.unit_registry.UnitRegistry, optional + A unyt registry to use in the conversion. If one is not supplied, the + default one will be used. 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 + Flag to indicate whether using comoving coordinates. + cosmo_factor : swiftsimio.objects.cosmo_factor + 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. + Description of the compression filters that were applied to that array in the + hdf5 file. valid_transform : bool - flag to indicate whether this array can be converted to comoving + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + A cosmology-aware array. Examples -------- - >>> from pint import UnitRegistry - >>> import numpy as np - >>> ureg = UnitRegistry() - >>> a = np.arange(4) - >>> b = ureg.Quantity(a, "erg/cm**3") - >>> b - - >>> c = cosmo_array.from_pint(b) - >>> c - cosmo_array([0, 1, 2, 3], 'erg/cm**3') + :: + + >>> from pint import UnitRegistry + >>> import numpy as np + >>> ureg = UnitRegistry() + >>> a = np.arange(4) + >>> b = ureg.Quantity(a, "erg/cm**3") + >>> b + + >>> c = cosmo_array.from_pint(b) + >>> c + cosmo_array([0, 1, 2, 3], 'erg/cm**3') """ obj = super().from_pint(arr, unit_registry=unit_registry).view(cls) obj.comoving = comoving @@ -1001,7 +1473,39 @@ def from_pint( return obj - def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): + def __array_ufunc__( + self, ufunc: np.ufunc, method: str, *inputs, **kwargs + ) -> object: + """ + Handles :mod:`numpy` ufunc calls on :class:`~swiftsimio.objects.cosmo_array` + input. + + :mod:`numpy` facilitates wrapping array classes by handing off to this function + when a function of :class:`numpy.ufunc` type is called with arguments from an + inheriting array class. Since we inherit from :class:`~unyt.array.unyt_array`, + we let :mod:`unyt` handle what to do with the units and take care of processing + the cosmology information via our helper functions. + + Parameters + ---------- + ufunc : numpy.ufunc + The numpy function being called. + + method : str, optional + Some ufuncs have methods accessed as attributes, such as ``"reduce"``. + If using such a method, this argument receives its name. + + inputs : tuple + Arguments to the ufunc. + + kwargs : dict + Keyword arguments to the ufunc. + + Returns + ------- + out : object + The result of the ufunc call, with our cosmology attribute processing applied. + """ helper_result = _prepare_array_func_args(*inputs, **kwargs) cfs = helper_result["cfs"] @@ -1056,7 +1560,39 @@ def __array_ufunc__(self, ufunc, method, *inputs, **kwargs): return ret - def __array_function__(self, func, types, args, kwargs): + def __array_function__( + self, func: Callable, types: Collection, args: tuple, kwargs: dict + ): + """ + Handles :mod:`numpy` function calls on :class:`~swiftsimio.objects.cosmo_array` + input. + + :mod:`numpy` facilitates wrapping array classes by handing off to this function + when a numpy-defined function is called with arguments from an + inheriting array class. Since we inherit from :class:`~unyt.array.unyt_array`, + we let :mod:`unyt` handle what to do with the units and take care of processing + the cosmology information via our helper functions. + + Parameters + ---------- + func : Callable + The numpy function being called. + + types : collections.abc.Collection + A collection of unique argument types from the original :mod:`numpy` function + call that implement ``__array_function__``. + + args : tuple + Arguments to the functions. + + kwargs : dict + Keyword arguments to the function. + + Returns + ------- + out : object + The result of the ufunc call, with our cosmology attribute processing applied. + """ # Follow NEP 18 guidelines # https://numpy.org/neps/nep-0018-array-function-protocol.html from ._array_functions import _HANDLED_FUNCTIONS @@ -1095,21 +1631,54 @@ 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. + This inherits from both the :class:`~swiftsimio.objects.cosmo_array` and the + :class:`~unyt.array.unyt_quantity`, and has the same four attributes as + :class:`~swiftsimio.objects.cosmo_array`: ``compression``, ``cosmo_factor``, + ``comoving``, and ``valid_transform``. - Parameters - ---------- + Like :class:`unyt.array.unyt_quantity`, it is intended to hold a scalar value. + Values of this type will be returned by :mod:`numpy` functions that return + scalar values. - cosmo_array : cosmo_array - the inherited cosmo_array + Other than containing a scalar, functionality is identical to + :class:`~swiftsimio.objects.cosmo_array`. Refer to that class's documentation. - unyt_quantity : unyt.unyt_quantity - the inherited unyt_quantity + Parameters + ---------- + input_scalar : float or unyt.array.unyt_quantity + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, 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 ``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 or invalid data, but can lead to significant speedups + in the input validation logic adds significant overhead. If set, minimally + pass valid values for units, comoving and cosmo_factor. 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 : swiftsimio.objects.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. If + ``False``, then ``comoving`` must be ``False``. + compression : string + Description of the compression filters that were applied to that array in the + hdf5 file. Attributes ---------- - comoving : bool if True then the array is in comoving co-ordinates, and if False then it is in physical units. @@ -1123,55 +1692,54 @@ class cosmo_quantity(cosmo_array, unyt_quantity): 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_scalar: numeric_type, + units: Union[str, unyt.unit_object.Unit, "astropy.units.core.Unit"] = None, + registry: unyt.unit_registry.UnitRegistry = None, + dtype: Union[np.dtype, str] = None, + bypass_validation: bool = False, + name: str = None, + cosmo_factor: cosmo_factor = None, + comoving: bool = None, + valid_transform: bool = True, + compression: str = None, + ) -> "cosmo_quantity": """ - Essentially a copy of the unyt_quantity.__new__ constructor. + Closely inspired by the :meth:`unyt.array.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). + input_scalar : float or unyt.array.unyt_quantity + A tuple, list, or array to attach units and cosmology information to. + units : str, unyt.unit_object.Unit or astropy.units.core.Unit, optional + The units of the array. When using strings, 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. + The registry to create units from. If ``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 + none is found, uses ``np.float64``. bypass_validation : bool, optional - If True, all input validation is skipped. Using this option may produce + If ``True``, all input validation is skipped. Using this option may produce corrupted or invalid data, but can lead to significant speedups in the input validation logic adds significant overhead. If set, minimally - pass valid values for units, comoving and cosmo_factor. Defaults to False. + pass valid values for units, comoving and cosmo_factor. Defaults to ``False``. name : str, optional - The name of the array. Defaults to None. This attribute does not propagate + 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. + cosmo_factor : swiftsimio.objects.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. + Flag to indicate whether this array can be converted to comoving. If + ``False``, then ``comoving`` must be ``False``. compression : string Description of the compression filters that were applied to that array in the hdf5 file. From a1689ae46738d43c8dde2eaade3d5b980f11a516 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 9 Feb 2025 14:44:07 +0000 Subject: [PATCH 069/125] Fix some sphinx issues. --- swiftsimio/_array_functions.py | 8 ++------ swiftsimio/objects.py | 16 +++++++--------- 2 files changed, 9 insertions(+), 15 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 181ba372..fe72a05d 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -180,9 +180,7 @@ def _propagate_cosmo_array_attributes_to_result(func: Callable) -> Callable: """ def wrapped(obj, *args, **kwargs): - """ - Call the function and copy the attributes. - """ + # omit docstring so that sphinx picks up docstring of wrapped function return _copy_cosmo_array_attributes(obj, func(obj, *args, **kwargs)) return wrapped @@ -272,9 +270,7 @@ def _ensure_result_is_cosmo_array_or_quantity(func: Callable) -> Callable: """ def wrapped(*args, **kwargs) -> object: - """ - Promote unyt types in function output to cosmo types. - """ + # omit docstring so that sphinx picks up docstring of wrapped function result = func(*args, **kwargs) if isinstance(result, tuple): return tuple( diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 8edf68d5..0516d64a 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -4,7 +4,7 @@ :class:`~swiftsimio.objects.cosmo_factor` objects for cosmology-aware arrays, extending the functionality of the :class:`~unyt.array.unyt_array`. -For developers, see also :mod:`~swiftsimio._array_functions` containing +For developers, see also :mod:`swiftsimio._array_functions` containing helpers, wrappers and implementations that enable most :mod:`numpy` and :mod:`unyt` functions to work with our cosmology-aware arrays. """ @@ -267,8 +267,9 @@ class cosmo_factor(object): Attributes ---------- - expr : sympy.expr + expr : sympy.Expr Expression used to convert between comoving and physical coordinates. + scale_factor : float The scale factor (a). @@ -283,16 +284,14 @@ class cosmo_factor(object): from swiftsimio.objects import a # the scale factor (a sympy symbol object) density_cosmo_factor = cosmo_factor(a**3, scale_factor=0.97) - :class:`~swiftsimio.objects.cosmo_factor` support arithmetic, for example: + :class:`~swiftsimio.objects.cosmo_factor` supports arithmetic, for example: :: >>> cosmo_factor(a**2, scale_factor=0.5) * cosmo_factor(a**-1, scale_factor=0.5) cosmo_factor(expr=a, scale_factor=0.5) - """ - expr: sympy.Expr - scale_factor: float + """ def __init__(self, expr: sympy.Expr, scale_factor: float) -> None: """ @@ -302,6 +301,7 @@ def __init__(self, expr: sympy.Expr, scale_factor: float) -> None: ---------- expr : sympy.expr Expression used to convert between comoving and physical coordinates. + scale_factor : float The scale factor (a). """ @@ -341,7 +341,7 @@ def redshift(self) -> float: """ The redshift computed from the scale factor. - Returns the redshift :math:`z = \frac{1}{a} - 1`, where :math:`a` is the scale + Returns the redshift :math:`z = \\frac{1}{a} - 1`, where :math:`a` is the scale factor. Returns @@ -1073,8 +1073,6 @@ def __new__( hdf5 file. """ - cosmo_factor: cosmo_factor - if bypass_validation is True: obj = super().__new__( From 2b46c57116119e070f0d2213f2bbc69979e7ec59 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 9 Feb 2025 14:45:29 +0000 Subject: [PATCH 070/125] Remove a and cosmo_factor from global namespace (was temporarily added for dev). --- swiftsimio/__init__.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/swiftsimio/__init__.py b/swiftsimio/__init__.py index ceb146f0..90efe849 100644 --- a/swiftsimio/__init__.py +++ b/swiftsimio/__init__.py @@ -8,7 +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 +from swiftsimio.objects import cosmo_array, cosmo_quantity import swiftsimio.visualisation as visualisation import swiftsimio.units as units import swiftsimio.subset_writer as subset_writer From 95704e0d021e06152e4940bbc4d91426360b1413 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sun, 9 Feb 2025 15:23:21 +0000 Subject: [PATCH 071/125] Fill in descriptions of all cosmo_array tests. --- tests/test_cosmo_array.py | 118 +++++++++++++++++++++++++++++++- tests/test_cosmo_array_attrs.py | 107 ++++++++++++++++++++++++++++- 2 files changed, 221 insertions(+), 4 deletions(-) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index eb9b6c52..ba112464 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -14,6 +14,10 @@ def getfunc(fname): + """ + Helper for our tests: get the function handle from a name (possibly with attribute + access). + """ func = np for attr in fname.split("."): func = getattr(func, attr) @@ -21,16 +25,26 @@ def getfunc(fname): def ca(x, unit=u.Mpc): + """ + Helper for our tests: turn an array into a cosmo_array. + """ return cosmo_array(x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5)) def cq(x, unit=u.Mpc): + """ + Helper for our tests: turn a scalar into a cosmo_quantity. + """ return cosmo_quantity( x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5) ) def arg_to_ua(arg): + """ + Helper for our tests: recursively convert cosmo_* in an argument (possibly an + iterable) to their unyt_* equivalents. + """ if type(arg) in (list, tuple): return type(arg)([arg_to_ua(a) for a in arg]) else: @@ -38,10 +52,22 @@ def arg_to_ua(arg): def to_ua(x): + """ + Helper for our tests: turn a cosmo_* object into its unyt_* equivalent. + """ return u.unyt_array(x) if hasattr(x, "comoving") else x def check_result(x_c, x_u, ignore_values=False): + """ + Helper for our tests: check that a result with cosmo input matches what we + expected based on the result with unyt input. + + We check: + - that the type of the result makes sense, recursing if needed. + - that the value of the result matches (unless ignore_values=False). + - that the units match. + """ if x_u is None: assert x_c is None return @@ -78,7 +104,14 @@ def check_result(x_c, x_u, ignore_values=False): class TestCosmoArrayInit: + """ + Test different ways of initializing a cosmo_array. + """ + def test_init_from_ndarray(self): + """ + Check initializing from a bare numpy array. + """ arr = cosmo_array( np.ones(5), units=u.Mpc, @@ -90,6 +123,9 @@ def test_init_from_ndarray(self): assert isinstance(arr, cosmo_array) def test_init_from_list(self): + """ + Check initializing from a list of values. + """ arr = cosmo_array( [1, 1, 1, 1, 1], units=u.Mpc, @@ -101,6 +137,9 @@ def test_init_from_list(self): assert isinstance(arr, cosmo_array) def test_init_from_unyt_array(self): + """ + Check initializing from a unyt_array. + """ arr = cosmo_array( u.unyt_array(np.ones(5), units=u.Mpc), cosmo_factor=cosmo_factor(a ** 1, 1), @@ -111,6 +150,12 @@ def test_init_from_unyt_array(self): assert isinstance(arr, cosmo_array) def test_init_from_list_of_unyt_arrays(self): + """ + Check initializing from a list of unyt_array's. + + Note that unyt won't recurse deeper than one level on inputs, so we don't test + deeper than one level of lists. This behaviour is documented in cosmo_array. + """ arr = cosmo_array( [u.unyt_array(1, units=u.Mpc) for _ in range(5)], cosmo_factor=cosmo_factor(a ** 1, 1), @@ -121,6 +166,12 @@ def test_init_from_list_of_unyt_arrays(self): assert isinstance(arr, cosmo_array) def test_init_from_list_of_cosmo_arrays(self): + """ + Check initializing from a list of cosmo_array's. + + Note that unyt won't recurse deeper than one level on inputs, so we don't test + deeper than one level of lists. This behaviour is documented in cosmo_array. + """ arr = cosmo_array( [ cosmo_array( @@ -140,10 +191,34 @@ def test_init_from_list_of_cosmo_arrays(self): class TestNumpyFunctions: + """ + Check that numpy functions recognize our cosmo classes as input and handle them + correctly. + """ + def test_explicitly_handled_funcs(self): """ Make sure we at least handle everything that unyt does, and anything that 'just worked' for unyt but that we need to handle by hand. + + We don't try to be exhaustive here, but at give some basic input to every function + that we expect to be able to take cosmo input. We then use our helpers defined + above to convert the inputs to unyt equivalents and call the numpy function on + both cosmo and unyt input. Then we use our helpers to check the results for + consistency. For instnace, if with unyt input we got back a unyt_array, we + should expect a cosmo_array. + + We are not currently explicitly testing that the results of any specific function + are numerically what we expected them to be (seems like overkill), nor that the + cosmo_factor's make sense given the input. The latter would be a useful addition, + but I can't think of a sensible way to implement this besides writing in the + expectation for every output value of every function by hand. + + As long as no functions outright crash, the test will report the list of functions + that we should have covered that we didn't cover in tests, and/or the list of + functions whose output values were not what we expected based on running them with + unyt input. Otherwise we just get a stack trace of the first function that + crashed. """ from unyt._array_functions import _HANDLED_FUNCTIONS from unyt.tests.test_array_functions import NOOP_FUNCTIONS @@ -413,7 +488,11 @@ def test_explicitly_handled_funcs(self): category=UserWarning, message="numpy.savetxt does not preserve units", ) - ua_result = func(*ua_args) + try: + ua_result = func(*ua_args) + except: + print(f"Crashed in {fname} with unyt input.") + raise except u.exceptions.UnytError: raises_unyt_error = True else: @@ -439,7 +518,11 @@ def test_explicitly_handled_funcs(self): category=RuntimeWarning, message="Mixing arguments with and without cosmo_factors", ) - result = func(*args) + try: + result = func(*args) + except: + print(f"Crashed in {fname} with cosmo input.") + raise if fname.split(".")[-1] in ( "fill_diagonal", "copyto", @@ -551,6 +634,13 @@ def test_explicitly_handled_funcs(self): @pytest.mark.parametrize("bins_type", ("int", "np", "ca")) @pytest.mark.parametrize("density", (None, True)) def test_histograms(self, func_args, weights, bins_type, density): + """ + Test that histograms give sensible output. + + Histograms are tricky with possible density and weights arguments, and the way + that attributes need validation and propagation between the bins and values. + They are also commonly used. They therefore need a bespoke test. + """ func, args = func_args bins = { "int": 10, @@ -663,17 +753,31 @@ def test_histograms(self, func_args, weights, bins_type, density): assert b.cosmo_factor == expt_cf def test_getitem(self): + """ + Make sure that we don't degrade to an ndarray on slicing. + """ assert isinstance(ca(np.arange(3))[0], cosmo_quantity) def test_reshape_to_scalar(self): + """ + Make sure that we convert to a cosmo_quantity when we reshape to a scalar. + """ assert isinstance(ca(np.ones(1)).reshape(tuple()), cosmo_quantity) def test_iter(self): + """ + Make sure that we get cosmo_quantity's when iterating over a cosmo_array. + """ for cq in ca(np.arange(3)): assert isinstance(cq, cosmo_quantity) class TestCosmoQuantity: + """ + Test that the cosmo_quantity class works as desired, mostly around issues converting + back and forth with cosmo_array. + """ + @pytest.mark.parametrize( "func, args", [ @@ -691,6 +795,9 @@ class TestCosmoQuantity: ], ) def test_propagation_func(self, func, args): + """ + Test that functions that are supposed to propagate our attributes do so. + """ cq = cosmo_quantity( 1, u.m, @@ -720,6 +827,9 @@ def test_scalar_return_func(self): @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"]) def test_propagation_props(self, prop): + """ + Test that properties propagate our attributes as intended. + """ cq = cosmo_quantity( 1, u.m, @@ -734,6 +844,10 @@ def test_propagation_props(self, prop): class TestCosmoArrayCopy: + """ + Tests of explicit (deep)copying of cosmo_array. + """ + def test_copy(self): """ Check that when we copy a cosmo_array it preserves its values and attributes. diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_attrs.py index 1ec07d8b..d4b1562c 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_attrs.py @@ -1,6 +1,5 @@ """ -Tests that functions returning copies of cosmo_array - preserve the comoving and cosmo_factor attributes. +Tests that ufuncs handling cosmo_array's properly handle our extra attributes. """ import pytest @@ -16,10 +15,18 @@ class TestCopyFuncs: + """ + Test ufuncs that copy arrays. + """ + @pytest.mark.parametrize( ("func"), ["byteswap", "diagonal", "flatten", "ravel", "transpose", "view"] ) def test_argless_copyfuncs(self, func): + """ + Make sure that our attributes are preserved through copying functions that + take no arguments. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -30,6 +37,9 @@ def test_argless_copyfuncs(self, func): assert hasattr(getattr(arr, func)(), "comoving") def test_astype(self): + """ + Make sure that our attributes are preserved through astype. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -41,6 +51,9 @@ def test_astype(self): assert hasattr(res, "comoving") def test_in_units(self): + """ + Make sure that our attributes are preserved through in_units (from unyt). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -52,6 +65,9 @@ def test_in_units(self): assert hasattr(res, "comoving") def test_compress(self): + """ + Make sure that our attributes are preserved through compress. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -63,6 +79,9 @@ def test_compress(self): assert hasattr(res, "comoving") def test_repeat(self): + """ + Make sure that our attributes are preserved through repeat. + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -74,6 +93,9 @@ def test_repeat(self): assert hasattr(res, "comoving") def test_T(self): + """ + Make sure that our attributes are preserved through transpose (T). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -85,6 +107,9 @@ def test_T(self): assert hasattr(res, "comoving") def test_ua(self): + """ + Make sure that our attributes are preserved through ua (from unyt). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -96,6 +121,9 @@ def test_ua(self): assert hasattr(res, "comoving") def test_unit_array(self): + """ + Make sure that our attributes are preserved through unit_array (from unyt). + """ arr = cosmo_array( np.ones((10, 10)), units="Mpc", @@ -107,6 +135,10 @@ def test_unit_array(self): assert hasattr(res, "comoving") def test_compatibility(self): + """ + Check that the compatible_with_comoving and compatible_with_physical functions + give correct compatibility checks. + """ # comoving array at high redshift arr = cosmo_array( np.ones((10, 10)), @@ -164,19 +196,37 @@ def test_compatibility(self): class TestCheckUfuncCoverage: + """ + Check that we've wrapped all functions that unyt wraps. + """ + def test_multi_output_coverage(self): + """ + Compare our list of multi_output_operators with unyt's to make sure we cover + everything. + """ assert set(multiple_output_operators.keys()) == set( (np.modf, np.frexp, np.divmod) ) def test_ufunc_coverage(self): + """ + Compare our list of ufuncs with unyt's to make sure we cover everything. + """ assert set(u.unyt_array._ufunc_registry.keys()) == set( cosmo_array._ufunc_registry.keys() ) class TestCosmoArrayUfuncs: + """ + Test some example functions using each of our wrappers for correct output. + """ + def test_preserving_ufunc(self): + """ + Tests of the _preserve_cosmo_factor wrapper. + """ # 1 argument inp = cosmo_array( [2], @@ -264,6 +314,9 @@ def test_preserving_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor def test_multiplying_ufunc(self): + """ + Tests of the _multiply_cosmo_factor wrapper. + """ # no cosmo_factors inp = cosmo_array([2], u.kpc, comoving=False) res = inp * inp @@ -317,6 +370,9 @@ def test_multiplying_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_dividing_ufunc(self): + """ + Tests of the _divide_cosmo_factor wrapper. + """ inp = cosmo_array( [2.0], u.kpc, @@ -329,6 +385,9 @@ def test_dividing_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 0 def test_return_without_ufunc(self): + """ + Tests of the _return_without_cosmo_factor wrapper. + """ # 1 argument inp = cosmo_array( [1], @@ -418,6 +477,9 @@ def test_return_without_ufunc(self): assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) def test_sqrt_ufunc(self): + """ + Tests of the _sqrt_cosmo_factor wrapper. + """ inp = cosmo_array( [4], u.kpc, @@ -430,6 +492,9 @@ def test_sqrt_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 0.5 def test_square_ufunc(self): + """ + Tests of the _square_cosmo_factor wrapper. + """ inp = cosmo_array( [2], u.kpc, @@ -442,6 +507,9 @@ def test_square_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_cbrt_ufunc(self): + """ + Tests of the _cbrt_cosmo_factor wrapper. + """ inp = cosmo_array( [8], u.kpc, @@ -454,6 +522,9 @@ def test_cbrt_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** (1.0 / 3.0) def test_reciprocal_ufunc(self): + """ + Tests of the _reciprocal_cosmo_factor wrapper. + """ inp = cosmo_array( [2.0], u.kpc, @@ -466,6 +537,9 @@ def test_reciprocal_ufunc(self): assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_passthrough_ufunc(self): + """ + Tests of the _passthrough_cosmo_factor wrapper. + """ # 1 argument inp = cosmo_array( [2], @@ -507,6 +581,9 @@ def test_passthrough_ufunc(self): np.copysign(inp1, inp2) def test_arctan2_ufunc(self): + """ + Tests of the _arctan2_cosmo_factor wrapper. + """ inp = cosmo_array( [2], u.kpc, @@ -519,6 +596,9 @@ def test_arctan2_ufunc(self): assert res.cosmo_factor.a_factor == 1 # also ensures cosmo_factor present def test_comparison_ufunc(self): + """ + Tests of the _comparison_cosmo_factor wrapper. + """ inp1 = cosmo_array( [1], u.kpc, @@ -536,6 +616,9 @@ def test_comparison_ufunc(self): assert isinstance(res, np.ndarray) and not isinstance(res, u.unyt_array) def test_out_arg(self): + """ + Test that our helpers can handle functions with an ``out`` kwarg. + """ inp = cosmo_array( [1], u.kpc, @@ -559,6 +642,9 @@ def test_out_arg(self): assert out == np.abs(inp.to_value(u.kpc)) def test_reduce_multiply(self): + """ + Test that we can handle the reduce method for the multiply ufunc. + """ inp = cosmo_array( [[1, 2], [3, 4]], u.kpc, @@ -571,6 +657,9 @@ def test_reduce_multiply(self): assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_reduce_divide(self): + """ + Test that we can handle the reduce method for the divide ufunc. + """ inp = cosmo_array( [[1.0, 2.0], [1.0, 4.0], [1.0, 1.0]], u.kpc, @@ -583,6 +672,9 @@ def test_reduce_divide(self): assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_reduce_other(self): + """ + Test that we can handle other ufuncs with a reduce method. + """ inp = cosmo_array( [[1.0, 2.0], [1.0, 2.0]], u.kpc, @@ -595,6 +687,9 @@ def test_reduce_other(self): assert res.cosmo_factor == inp.cosmo_factor def test_multi_output(self): + """ + Test that we can handle functions with multiple return values. + """ # with passthrough inp = cosmo_array( [2.5], @@ -623,6 +718,10 @@ def test_multi_output(self): assert isinstance(res2, np.ndarray) and not isinstance(res2, u.unyt_array) def test_multi_output_with_out_arg(self): + """ + Test that we can handle multiple return values in conjunction with an ``out`` + kwarg. + """ # with two out arrays inp = cosmo_array( [2.5], @@ -641,6 +740,10 @@ def test_multi_output_with_out_arg(self): assert out2.cosmo_factor == inp.cosmo_factor def test_comparison_with_zero(self): + """ + Test that we don't produce warnings for dangerous comparisons on comparison with + zero. + """ inp1 = cosmo_array( [1, 1, 1], u.kpc, From 7cc732e00f375d4e7467cb2185506545a347dbb1 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 16:34:46 +0000 Subject: [PATCH 072/125] Remove some more unused imports. --- swiftsimio/visualisation/rotation.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/swiftsimio/visualisation/rotation.py b/swiftsimio/visualisation/rotation.py index c9d01c2f..5fab663d 100644 --- a/swiftsimio/visualisation/rotation.py +++ b/swiftsimio/visualisation/rotation.py @@ -2,10 +2,9 @@ Rotation matrix calculation routines. """ -from swiftsimio.accelerated import jit -from numpy import float64, array, matrix, cross, identity, dot, matmul +from numpy import float64, array, cross, identity, dot, matmul from numpy.linalg import norm, inv -from math import sin, cos, sqrt, acos +from math import sin, acos def rotation_matrix_from_vector(vector: float64, axis: str = "z") -> array: From de7ffad8ce545972bc977836b10dcaf6657b3471 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 16:41:02 +0000 Subject: [PATCH 073/125] Add now-redundant checks on rotation_center argument. --- swiftsimio/visualisation/projection.py | 21 ++++----------------- swiftsimio/visualisation/ray_trace.py | 20 ++++---------------- 2 files changed, 8 insertions(+), 33 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 59545f50..753de175 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -114,19 +114,6 @@ def project_pixel_grid( array if you want it to be visualised the 'right way up'. """ - if rotation_center is not None: - try: - if rotation_center.units == data.coordinates.units: - pass - else: - raise exceptions.InvalidUnitOperation( - "Units of coordinates and rotation center must agree" - ) - except AttributeError: - raise exceptions.InvalidUnitOperation( - "Ensure that rotation_center is a cosmo_array with the same units as coordinates" - ) - number_of_particles = data.coordinates.shape[0] if project is None: @@ -458,23 +445,23 @@ def project_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index c5d65cb1..6b68ae3c 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -226,18 +226,6 @@ def panel_pixel_grid( rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, ) -> cosmo_array: - if rotation_center is not None: - try: - if rotation_center.units == data.coordinates.units: - pass - else: - raise unyt.exceptions.InvalidUnitOperation( - "Units of coordinates and rotation center must agree" - ) - except AttributeError: - raise unyt.exceptions.InvalidUnitOperation( - "Ensure that rotation_center is a unyt array with the same units as coordinates" - ) number_of_particles = data.coordinates.shape[0] @@ -353,23 +341,23 @@ def panel_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) From d3dcdff7aa817cce920062391c5309e127235c34 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 16:47:57 +0000 Subject: [PATCH 074/125] Cleanup more imports. --- swiftsimio/visualisation/projection.py | 49 ++++++++++---------------- 1 file changed, 19 insertions(+), 30 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 753de175..97936a6a 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -3,18 +3,7 @@ """ from typing import Union -from numpy import ( - float32, - array, - ones, - matmul, - empty_like, - logical_and, - s_, - ceil, - zeros_like, -) -from unyt import exceptions +import numpy as np from swiftsimio import SWIFTDataset, cosmo_array from swiftsimio.reader import __SWIFTGroupDataset @@ -31,8 +20,8 @@ def project_pixel_grid( resolution: int, project: Union[str, None] = "masses", region: Union[None, cosmo_array] = None, - mask: Union[None, array] = None, - rotation_matrix: Union[None, array] = None, + mask: Union[None, np.array] = None, + rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, parallel: bool = False, backend: str = "fast", @@ -117,7 +106,7 @@ def project_pixel_grid( number_of_particles = data.coordinates.shape[0] if project is None: - m = ones(number_of_particles, dtype=float32) + m = np.ones(number_of_particles, dtype=np.float32) else: m = getattr(data, project) if data.coordinates.comoving: @@ -134,7 +123,7 @@ def project_pixel_grid( # This provides a default 'slice it all' mask. if mask is None: - mask = s_[:] + mask = np.s_[:] box_x, box_y, box_z = boxsize @@ -148,12 +137,12 @@ def project_pixel_grid( z_slice_included = True z_min, z_max = region[4:] else: - z_min = zeros_like(box_z) + z_min = np.zeros_like(box_z) z_max = box_z else: - x_min = zeros_like(box_x) + x_min = np.zeros_like(box_x) x_max = box_x - y_min = zeros_like(box_y) + y_min = np.zeros_like(box_y) y_max = box_y x_range = x_max - x_min @@ -184,13 +173,13 @@ def project_pixel_grid( # No hsml present. If they are using the 'histogram' backend, we # should just mock them to be anything as it doesn't matter. if backend == "histogram": - hsml = empty_like(m) + hsml = np.empty_like(m) else: raise AttributeError if rotation_center is not None: # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.coordinates - rotation_center).T) + x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) x += rotation_center[0] y += rotation_center[1] @@ -200,9 +189,9 @@ def project_pixel_grid( x, y, z = data.coordinates.T if z_slice_included: - combined_mask = logical_and(mask, logical_and(z <= z_max, z >= z_min)).astype( - bool - ) + combined_mask = np.logical_and( + mask, np.logical_and(z <= z_max, z >= z_min) + ).astype(bool) else: combined_mask = mask @@ -229,8 +218,8 @@ def project_pixel_grid( image = backends[backend](**common_arguments) # determine the effective number of pixels for each dimension - xres = int(ceil(resolution * (x_range / max_range))) - yres = int(ceil(resolution * (y_range / max_range))) + xres = int(np.ceil(resolution * (x_range / max_range))) + yres = int(np.ceil(resolution * (y_range / max_range))) # trim the image to remove empty pixels return image[:xres, :yres] @@ -241,8 +230,8 @@ def project_gas_pixel_grid( resolution: int, project: Union[str, None] = "masses", region: Union[None, cosmo_array] = None, - mask: Union[None, array] = None, - rotation_matrix: Union[None, array] = None, + mask: Union[None, np.array] = None, + rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, parallel: bool = False, backend: str = "fast", @@ -346,9 +335,9 @@ def project_gas( resolution: int, project: Union[str, None] = "masses", region: Union[None, cosmo_array] = None, - mask: Union[None, array] = None, + mask: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, - rotation_matrix: Union[None, array] = None, + rotation_matrix: Union[None, np.array] = None, parallel: bool = False, backend: str = "fast", periodic: bool = True, From 8e9052b1976651062d6f4d7ebcf31bb3383ae5ff Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:02:49 +0000 Subject: [PATCH 075/125] Refactor try/except into getattr for fetching hsml. --- swiftsimio/visualisation/projection.py | 8 +++----- swiftsimio/visualisation/ray_trace.py | 8 +++----- swiftsimio/visualisation/volume_render.py | 12 +++++------- 3 files changed, 11 insertions(+), 17 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 97936a6a..34ed95cd 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -154,11 +154,9 @@ def project_pixel_grid( max_range = max(x_range, y_range) try: - try: - hsml = data.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.smoothing_length + hsml = getattr( + data, "smoothing_lengths", data.smoothing_length + ) # backwards compatible if data.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 6b68ae3c..80b83db6 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -276,11 +276,9 @@ def panel_pixel_grid( # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) - try: - hsml = data.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.smoothing_length + hsml = getattr( + data, "smoothing_lengths", data.smoothing_length + ) # backwards compatible if data.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 3430cb0d..adff7991 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -703,11 +703,9 @@ def render_gas_voxel_grid( else: x, y, z = data.gas.coordinates.T - try: - hsml = data.gas.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.gas.smoothing_length + hsml = getattr( + data, "smoothing_lengths", data.smoothing_length + ) # backwards compatible if data.gas.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( @@ -844,14 +842,14 @@ def render_gas( * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) From 36a4ba2d624ceb182199e7f068f01b14df9d0733 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:16:33 +0000 Subject: [PATCH 076/125] Remove stray print statement from debugging. --- swiftsimio/conversions.py | 6 ++---- 1 file changed, 2 insertions(+), 4 deletions(-) diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index 85917d91..2e6d65c0 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -45,7 +45,6 @@ def swift_neutrinos_to_astropy(N_eff, N_ur, M_nu_eV, deg_nu): raise AttributeError( "SWIFTsimIO uses astropy, which cannot handle this cosmological model." ) - print(N_eff, deg_nu, N_ur) if not int(N_eff) == deg_nu.astype(int).sum() + int(np.squeeze(N_ur)): raise AttributeError( "SWIFTSimIO uses astropy, which cannot handle this cosmological model." @@ -99,13 +98,13 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: # expressions taken directly from astropy, since they do no longer # allow access to these attributes (since version 5.1+) critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value - a_B_c2 = (4.0 * const.sigma_sb / const.c ** 3).cgs.value + a_B_c2 = (4.0 * const.sigma_sb / const.c**3).cgs.value # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy. # This is an exact inversion of the procedure performed in astropy. critical_density_0 = astropy_units.Quantity( critdens_const * H0.to("1/s").value ** 2, - astropy_units.g / astropy_units.cm ** 3, + astropy_units.g / astropy_units.cm**3, ) Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0) @@ -144,7 +143,6 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: m_nu=ap_m_nu, ) - else: def swift_cosmology_to_astropy(cosmo: dict, units) -> dict: From 5871b1ee0085a1a3dde938ad5d56cb59408224f8 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:28:17 +0000 Subject: [PATCH 077/125] Gradually coercing vis functions to share code. --- swiftsimio/visualisation/projection.py | 39 +++++++++++------------ swiftsimio/visualisation/ray_trace.py | 8 +++-- swiftsimio/visualisation/volume_render.py | 21 ++++++------ 3 files changed, 35 insertions(+), 33 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 34ed95cd..02f45bc6 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -153,27 +153,24 @@ def project_pixel_grid( # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) - try: - hsml = getattr( - data, "smoothing_lengths", data.smoothing_length - ) # backwards compatible - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - except AttributeError: - # No hsml present. If they are using the 'histogram' backend, we - # should just mock them to be anything as it doesn't matter. - if backend == "histogram": - hsml = np.empty_like(m) - else: - raise AttributeError + if backend == "histogram": + hsml = np.empty_like(m) # not used anyway for this backend + else: + hsml = ( + data.smoothing_lengths + if hasattr(data, "smoothing_lengths") + else data.smoothing_length # backwards compatibility + ) + if data.coordinates.comoving: + if not hsml.compatible_with_comoving(): + raise AttributeError( + "Physical smoothing length is not compatible with comoving coordinates!" + ) + else: + if not hsml.compatible_with_physical(): + raise AttributeError( + "Comoving smoothing length is not compatible with physical coordinates!" + ) if rotation_center is not None: # Rotate co-ordinates as required diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 80b83db6..7e038496 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -276,9 +276,11 @@ def panel_pixel_grid( # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) - hsml = getattr( - data, "smoothing_lengths", data.smoothing_length - ) # backwards compatible + hsml = ( + data.smoothing_lengths + if hasattr(data, "smoothing_lengths") + else data.smoothing_length # backwards compatibility + ) if data.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index adff7991..0fd26594 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -647,14 +647,15 @@ def render_gas_voxel_grid( slice_gas_pixel_grid : Creates a 2D slice of a SWIFT dataset """ + data = data.gas # coerce rest of this function to share with other vis functions - number_of_gas_particles = data.gas.particle_ids.size + number_of_gas_particles = data.particle_ids.size if project is None: m = ones(number_of_gas_particles, dtype=float32) else: - m = getattr(data.gas, project) - if data.gas.coordinates.comoving: + m = getattr(data, project) + if data.coordinates.comoving: if not m.compatible_with_comoving(): raise AttributeError( f'Physical quantity "{project}" is not compatible with comoving coordinates!' @@ -694,19 +695,21 @@ def render_gas_voxel_grid( # Let's just hope that the box is square otherwise we're probably SOL if rotation_center is not None: # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.gas.coordinates - rotation_center).T) + x, y, z = matmul(rotation_matrix, (data.coordinates - rotation_center).T) x += rotation_center[0] y += rotation_center[1] z += rotation_center[2] else: - x, y, z = data.gas.coordinates.T + x, y, z = data.coordinates.T - hsml = getattr( - data, "smoothing_lengths", data.smoothing_length - ) # backwards compatible - if data.gas.coordinates.comoving: + hsml = ( + data.smoothing_lengths + if hasattr(data, "smoothing_lengths") + else data.smoothing_length # backwards compatibility + ) + if data.coordinates.comoving: if not hsml.compatible_with_comoving(): raise AttributeError( "Physical smoothing length is not compatible with comoving coordinates!" From 42747980418bc3d8fc6c9aecfb37a7ed2e87c653 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:31:09 +0000 Subject: [PATCH 078/125] Remove redundant boxsize argument, get from metadata. --- swiftsimio/visualisation/projection.py | 7 +------ swiftsimio/visualisation/ray_trace.py | 4 +--- 2 files changed, 2 insertions(+), 9 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 02f45bc6..fbb40bf7 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -16,7 +16,6 @@ def project_pixel_grid( data: __SWIFTGroupDataset, - boxsize: cosmo_array, resolution: int, project: Union[str, None] = "masses", region: Union[None, cosmo_array] = None, @@ -41,9 +40,6 @@ def project_pixel_grid( data: __SWIFTGroupDataset The SWIFT dataset that you wish to visualise (get this from ``load``) - boxsize: cosmo_array - The box-size of the simulation. - resolution: int The resolution of the image. All images returned are square, ``res`` by ``res``, pixel grids. @@ -125,7 +121,7 @@ def project_pixel_grid( if mask is None: mask = np.s_[:] - box_x, box_y, box_z = boxsize + box_x, box_y, box_z = data.metadata.boxsize # Set the limits of the image. z_slice_included = False @@ -310,7 +306,6 @@ def project_gas_pixel_grid( image = project_pixel_grid( data=data.gas, - boxsize=data.metadata.boxsize, resolution=resolution, project=project, mask=mask, diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 7e038496..bbf5ce70 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -217,7 +217,6 @@ def core_panels_parallel( def panel_pixel_grid( data: __SWIFTGroupDataset, - boxsize: cosmo_array, resolution: int, panels: int, project: Union[str, None] = "masses", @@ -249,7 +248,7 @@ def panel_pixel_grid( if mask is None: mask = np.s_[:] - box_x, box_y, box_z = boxsize + box_x, box_y, box_z = data.metadata.boxsize # Set the limits of the image. if region is not None: @@ -327,7 +326,6 @@ def panel_gas( ) -> cosmo_array: image = panel_pixel_grid( data=data.gas, - boxsize=data.metadata.boxsize, resolution=resolution, panels=panels, project=project, From c80267e05b501ce81b47254ae3d870f414996ec3 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:35:22 +0000 Subject: [PATCH 079/125] Remove argument from tests, too. --- tests/test_visualisation.py | 23 +++++++++++------------ 1 file changed, 11 insertions(+), 12 deletions(-) diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index e35eda9d..543cf473 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -79,7 +79,7 @@ def test_scatter_mass_conservation(): for resolution in resolutions: image = scatter(x, y, m, h, resolution, 1.0, 1.0) - mass_in_image = image.sum() / (resolution ** 2) + mass_in_image = image.sum() / (resolution**2) # Check mass conservation to 5% assert np.isclose(mass_in_image, total_mass, 0.05) @@ -356,7 +356,7 @@ def test_comoving_versus_physical(filename): img = func(data, resolution=256, project="masses") # check that we get a physical result assert not img.comoving - assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 + assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 # densities are still compatible with physical img = func(data, resolution=256, project="densities") assert not img.comoving @@ -380,7 +380,6 @@ def test_nongas_smoothing_lengths(filename): ) project_pixel_grid( data.dark_matter, - boxsize=data.metadata.boxsize, resolution=256, project="masses", ) @@ -418,14 +417,14 @@ def test_panel_rendering(filename): plt.imsave( "panels_added.png", - plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), + plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1))), ) projected = project_gas(data, res, "masses", backend="renormalised") plt.imsave( "projected.png", - plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), + plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)), ) fullstack = np.zeros((res, res)) @@ -625,8 +624,8 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v + mean_density_deposit = (np.sum(deposition) / npix**3).to("Msun / kpc**3").v + mean_density_volume = (np.sum(volume) / npix**3).to("Msun / kpc**3").v mean_density_calculated = ( (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) .to("Msun / kpc**3") @@ -651,15 +650,15 @@ def test_dark_matter_power_spectrum(filename, save=False): min_k = cosmo_quantity( 1e-2, - unyt.Mpc ** -1, + unyt.Mpc**-1, comoving=True, - cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), ) max_k = cosmo_quantity( 1e2, - unyt.Mpc ** -1, + unyt.Mpc**-1, comoving=True, - cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), ) bins = np.geomspace(min_k, max_k, 32) @@ -696,7 +695,7 @@ def test_dark_matter_power_spectrum(filename, save=False): folds[folding] = deposition - folding_output[2 ** folding] = (k, power_spectrum, scatter) + folding_output[2**folding] = (k, power_spectrum, scatter) # Now try doing them all together at once. From fc2177dd0a9162eb71f9304ab24feab8f762c116 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:35:34 +0000 Subject: [PATCH 080/125] Cleanup more imports. --- swiftsimio/visualisation/volume_render.py | 283 +++++++++++----------- 1 file changed, 136 insertions(+), 147 deletions(-) diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 0fd26594..3a74d3cb 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -5,19 +5,7 @@ from typing import List, Literal, Tuple, Union from math import sqrt, exp, pi -from numpy import ( - float64, - float32, - int32, - zeros, - array, - ndarray, - ones, - isclose, - linspace, - matmul, - max as np_max, -) +import numpy as np from swiftsimio import SWIFTDataset, cosmo_array from swiftsimio.accelerated import jit, NUM_THREADS, prange @@ -28,16 +16,16 @@ @jit(nopython=True, fastmath=True) def scatter( - x: float64, - y: float64, - z: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, res: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, - box_z: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: """ Creates a weighted voxel grid @@ -48,41 +36,41 @@ def scatter( Parameters ---------- - x : np.array[float64] - array of x-positions of the particles. Must be bounded by [0, 1]. + x : np.np.array[np.float64] + np.array of x-positions of the particles. Must be bounded by [0, 1]. - y : np.array[float64] - array of y-positions of the particles. Must be bounded by [0, 1]. + y : np.np.array[np.float64] + np.array of y-positions of the particles. Must be bounded by [0, 1]. - z : np.array[float64] - array of z-positions of the particles. Must be bounded by [0, 1]. + z : np.np.array[np.float64] + np.array of z-positions of the particles. Must be bounded by [0, 1]. - m : np.array[float32] - array of masses (or otherwise weights) of the particles + m : np.np.array[np.float32] + np.array of masses (or otherwise weights) of the particles - h : np.array[float32] - array of smoothing lengths of the particles + h : np.np.array[np.float32] + np.array of smoothing lengths of the particles res : int the number of voxels along one axis, i.e. this returns a cube of res * res * res. - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x, y and z. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x, y and z. Used for periodic wrapping. - box_z: float64 + box_z: np.float64 box size in z, in the same rescaled length units as x, y and z. Used for periodic wrapping Returns ------- - np.array[float32, float32, float32] + np.np.array[np.float32, np.float32, np.float32] voxel grid of quantity See Also @@ -97,19 +85,19 @@ def scatter( Explicitly defining the types in this function allows for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba ones. + floats and integers is also an improvement over using the numba np.ones. """ - # Output array for our image - image = zeros((res, res, res), dtype=float32) - maximal_array_index = int32(res) - 1 + # Output np.array for our image + image = np.zeros((res, res, res), dtype=np.float32) + maximal_array_index = np.int32(res) - 1 # Change that integer to a float, we know that our x, y are bounded # by [0, 1]. - float_res = float32(res) + float_res = np.float32(res) pixel_width = 1.0 / float_res # We need this for combining with the x_pos and y_pos variables. - float_res_64 = float64(res) + float_res_64 = np.float64(res) # If the kernel width is smaller than this, we drop to just PIC method drop_to_single_cell = pixel_width * 0.5 @@ -149,15 +137,15 @@ def scatter( # Calculate the cell that this particle; use the 64 bit version of the # resolution as this is the same type as the positions - particle_cell_x = int32(float_res_64 * x_pos) - particle_cell_y = int32(float_res_64 * y_pos) - particle_cell_z = int32(float_res_64 * z_pos) + particle_cell_x = np.int32(float_res_64 * x_pos) + particle_cell_y = np.int32(float_res_64 * y_pos) + particle_cell_z = np.int32(float_res_64 * z_pos) # SWIFT stores hsml as the FWHM. kernel_width = kernel_gamma * hsml # The number of cells that this kernel spans - cells_spanned = int32(1.0 + kernel_width * float_res) + cells_spanned = np.int32(1.0 + kernel_width * float_res) if ( particle_cell_x + cells_spanned < 0 @@ -188,7 +176,7 @@ def scatter( for cell_x in range( # Ensure that the lowest x value is 0, otherwise we'll segfault max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the array bounds, + # Ensure that the highest x value lies within the np.array bounds, # otherwise we'll segfault (oops). min( particle_cell_x + cells_spanned, maximal_array_index + 1 @@ -197,8 +185,8 @@ def scatter( # The distance in x to our new favourite cell -- remember that our x, y # are all in a box of [0, 1]; calculate the distance to the cell centre distance_x = ( - float32(cell_x) + 0.5 - ) * pixel_width - float32(x_pos) + np.float32(cell_x) + 0.5 + ) * pixel_width - np.float32(x_pos) distance_x_2 = distance_x * distance_x for cell_y in range( max(0, particle_cell_y - cells_spanned), @@ -208,8 +196,8 @@ def scatter( ), ): distance_y = ( - float32(cell_y) + 0.5 - ) * pixel_width - float32(y_pos) + np.float32(cell_y) + 0.5 + ) * pixel_width - np.float32(y_pos) distance_y_2 = distance_y * distance_y for cell_z in range( max(0, particle_cell_z - cells_spanned), @@ -219,8 +207,8 @@ def scatter( ), ): distance_z = ( - float32(cell_z) + 0.5 - ) * pixel_width - float32(z_pos) + np.float32(cell_z) + 0.5 + ) * pixel_width - np.float32(z_pos) distance_z_2 = distance_z * distance_z r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) @@ -234,17 +222,17 @@ def scatter( @jit(nopython=True, fastmath=True) def scatter_limited_z( - x: float64, - y: float64, - z: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, res: int, res_ratio_z: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, - box_z: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: """ Creates a weighted voxel grid @@ -255,20 +243,20 @@ def scatter_limited_z( Parameters ---------- - x : np.array[float64] - array of x-positions of the particles. Must be bounded by [0, 1]. + x : np.np.array[np.float64] + np.array of x-positions of the particles. Must be bounded by [0, 1]. - y : np.array[float64] - array of y-positions of the particles. Must be bounded by [0, 1]. + y : np.np.array[np.float64] + np.array of y-positions of the particles. Must be bounded by [0, 1]. - z : np.array[float64] - array of z-positions of the particles. Must be bounded by [0, 1]. + z : np.np.array[np.float64] + np.array of z-positions of the particles. Must be bounded by [0, 1]. - m : np.array[float32] - array of masses (or otherwise weights) of the particles + m : np.np.array[np.float32] + np.array of masses (or otherwise weights) of the particles - h : np.array[float32] - array of smoothing lengths of the particles + h : np.np.array[np.float32] + np.array of smoothing lengths of the particles res : int the number of voxels along one axis, i.e. this returns a cube @@ -277,24 +265,24 @@ def scatter_limited_z( res_ratio_z: int the number of voxels along the x and y axes relative to the z axis. If this is, for instance, 8, and the res is 128, then the - output array will be 128 x 128 x 16. + output np.array will be 128 x 128 x 16. - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x, y and z. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x, y and z. Used for periodic wrapping. - box_z: float64 + box_z: np.float64 box size in z, in the same rescaled length units as x, y and z. Used for periodic wrapping Returns ------- - np.array[float32, float32, float32] + np.np.array[np.float32, np.float32, np.float32] voxel grid of quantity See Also @@ -309,24 +297,24 @@ def scatter_limited_z( Explicitly defining the types in this function allows for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba ones. + floats and integers is also an improvement over using the numba np.ones. """ - # Output array for our image + # Output np.array for our image res_z = res // res_ratio_z - image = zeros((res, res, res_z), dtype=float32) - maximal_array_index = int32(res) - 1 - maximal_array_index_z = int32(res_z) - 1 + image = np.zeros((res, res, res_z), dtype=np.float32) + maximal_array_index = np.int32(res) - 1 + maximal_array_index_z = np.int32(res_z) - 1 # Change that integer to a float, we know that our x, y are bounded # by [0, 1]. - float_res = float32(res) - float_res_z = float32(res_z) + float_res = np.float32(res) + float_res_z = np.float32(res_z) pixel_width = 1.0 / float_res pixel_width_z = 1.0 / float_res_z # We need this for combining with the x_pos and y_pos variables. - float_res_64 = float64(res) - float_res_z_64 = float64(res_z) + float_res_64 = np.float64(res) + float_res_z_64 = np.float64(res_z) # If the kernel width is smaller than this, we drop to just PIC method drop_to_single_cell = pixel_width * 0.5 @@ -367,16 +355,16 @@ def scatter_limited_z( # Calculate the cell that this particle; use the 64 bit version of the # resolution as this is the same type as the positions - particle_cell_x = int32(float_res_64 * x_pos) - particle_cell_y = int32(float_res_64 * y_pos) - particle_cell_z = int32(float_res_z_64 * z_pos) + particle_cell_x = np.int32(float_res_64 * x_pos) + particle_cell_y = np.int32(float_res_64 * y_pos) + particle_cell_z = np.int32(float_res_z_64 * z_pos) # SWIFT stores hsml as the FWHM. kernel_width = kernel_gamma * hsml # The number of cells that this kernel spans - cells_spanned = int32(1.0 + kernel_width * float_res) - cells_spanned_z = int32(1.0 + kernel_width * float_res_z) + cells_spanned = np.int32(1.0 + kernel_width * float_res) + cells_spanned_z = np.int32(1.0 + kernel_width * float_res_z) if ( particle_cell_x + cells_spanned < 0 @@ -410,7 +398,7 @@ def scatter_limited_z( for cell_x in range( # Ensure that the lowest x value is 0, otherwise we'll segfault max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the array bounds, + # Ensure that the highest x value lies within the np.array bounds, # otherwise we'll segfault (oops). min( particle_cell_x + cells_spanned, maximal_array_index + 1 @@ -419,8 +407,8 @@ def scatter_limited_z( # The distance in x to our new favourite cell -- remember that our x, y # are all in a box of [0, 1]; calculate the distance to the cell centre distance_x = ( - float32(cell_x) + 0.5 - ) * pixel_width - float32(x_pos) + np.float32(cell_x) + 0.5 + ) * pixel_width - np.float32(x_pos) distance_x_2 = distance_x * distance_x for cell_y in range( max(0, particle_cell_y - cells_spanned), @@ -430,8 +418,8 @@ def scatter_limited_z( ), ): distance_y = ( - float32(cell_y) + 0.5 - ) * pixel_width - float32(y_pos) + np.float32(cell_y) + 0.5 + ) * pixel_width - np.float32(y_pos) distance_y_2 = distance_y * distance_y for cell_z in range( max(0, particle_cell_z - cells_spanned_z), @@ -441,8 +429,8 @@ def scatter_limited_z( ), ): distance_z = ( - float32(cell_z) + 0.5 - ) * pixel_width_z - float32(z_pos) + np.float32(cell_z) + 0.5 + ) * pixel_width_z - np.float32(z_pos) distance_z_2 = distance_z * distance_z r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) @@ -456,17 +444,17 @@ def scatter_limited_z( @jit(nopython=True, fastmath=True, parallel=True) def scatter_parallel( - x: float64, - y: float64, - z: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, res: int, res_ratio_z: int = 1, - box_x: float64 = 0.0, - box_y: float64 = 0.0, - box_z: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: """ Parallel implementation of scatter @@ -476,20 +464,20 @@ def scatter_parallel( Parameters ---------- - x : array of float64 - array of x-positions of the particles. Must be bounded by [0, 1]. + x : np.array of np.float64 + np.array of x-positions of the particles. Must be bounded by [0, 1]. - y : array of float64 - array of y-positions of the particles. Must be bounded by [0, 1]. + y : np.array of np.float64 + np.array of y-positions of the particles. Must be bounded by [0, 1]. - z : array of float64 - array of z-positions of the particles. Must be bounded by [0, 1]. + z : np.array of np.float64 + np.array of z-positions of the particles. Must be bounded by [0, 1]. - m : array of float32 - array of masses (or otherwise weights) of the particles + m : np.array of np.float32 + np.array of masses (or otherwise weights) of the particles - h : array of float32 - array of smoothing lengths of the particles + h : np.array of np.float32 + np.array of smoothing lengths of the particles res : int the number of voxels along one axis, i.e. this returns a cube @@ -498,22 +486,22 @@ def scatter_parallel( res_ratio_z: int the number of voxels along the x and y axes relative to the z - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x, y and z. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x, y and z. Used for periodic wrapping. - box_z: float64 + box_z: np.float64 box size in z, in the same rescaled length units as x, y and z. Used for periodic wrapping Returns ------- - ndarray of float32 + np.ndarray of np.float32 voxel grid of quantity See Also @@ -528,7 +516,7 @@ def scatter_parallel( Explicitly defining the types in this function allows for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba ones. + floats and integers is also an improvement over using the numba np.ones. """ # Same as scatter, but executes in parallel! This is actually trivial, @@ -537,7 +525,7 @@ def scatter_parallel( number_of_particles = x.size core_particles = number_of_particles // NUM_THREADS - output = zeros((res, res, res), dtype=float32) + output = np.zeros((res, res, res), dtype=np.float32) for thread in prange(NUM_THREADS): # Left edge is easy, just start at 0 and go to 'final' @@ -586,7 +574,7 @@ def render_gas_voxel_grid( resolution: int, project: Union[str, None] = "masses", parallel: bool = False, - rotation_matrix: Union[None, array] = None, + rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, region: Union[None, cosmo_array] = None, periodic: bool = True, @@ -601,7 +589,7 @@ def render_gas_voxel_grid( Dataset from which slice is extracted resolution : int - Specifies size of return array + Specifies size of return np.array project : str, optional Data field to be projected. Default is mass. If None then simply @@ -612,7 +600,7 @@ def render_gas_voxel_grid( defaults to False, but can speed up the creation of large images significantly at the cost of increased memory usage. - rotation_matrix: np.array, optional + rotation_matrix: np.np.array, optional Rotation matrix (3x3) that describes the rotation of the box around ``rotation_center``. In the default case, this provides a volume render viewed along the z axis. @@ -638,8 +626,8 @@ def render_gas_voxel_grid( Returns ------- - ndarray of float32 - Creates a `resolution` x `resolution` x `resolution` array and + np.ndarray of np.float32 + Creates a `resolution` x `resolution` x `resolution` np.array and returns it, without appropriate units. See Also @@ -652,7 +640,7 @@ def render_gas_voxel_grid( number_of_gas_particles = data.particle_ids.size if project is None: - m = ones(number_of_gas_particles, dtype=float32) + m = np.ones(number_of_gas_particles, dtype=np.float32) else: m = getattr(data, project) if data.coordinates.comoving: @@ -686,7 +674,8 @@ def render_gas_voxel_grid( # Test that we've got a cubic box if not ( - isclose(x_range.value, y_range.value) and isclose(x_range.value, z_range.value) + np.isclose(x_range.value, y_range.value) + and np.isclose(x_range.value, z_range.value) ): raise AttributeError( "Projection code is currently not able to handle non-cubic images" @@ -695,7 +684,7 @@ def render_gas_voxel_grid( # Let's just hope that the box is square otherwise we're probably SOL if rotation_center is not None: # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.coordinates - rotation_center).T) + x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) x += rotation_center[0] y += rotation_center[1] @@ -754,7 +743,7 @@ def render_gas( resolution: int, project: Union[str, None] = "masses", parallel: bool = False, - rotation_matrix: Union[None, array] = None, + rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, region: Union[None, cosmo_array] = None, periodic: bool = True, @@ -769,7 +758,7 @@ def render_gas( Dataset from which slice is extracted resolution : int - Specifies size of return array + Specifies size of return np.array project : str, optional Data field to be projected. Default is mass. If None then simply @@ -780,7 +769,7 @@ def render_gas( defaults to False, but can speed up the creation of large images significantly at the cost of increased memory usage. - rotation_matrix: np.array, optional + rotation_matrix: np.np.array, optional Rotation matrix (3x3) that describes the rotation of the box around ``rotation_center``. In the default case, this provides a volume render viewed along the z axis. @@ -805,8 +794,8 @@ def render_gas( Returns ------- - ndarray of float32 - a `resolution` x `resolution` x `resolution` array of the contribution + np.ndarray of np.float32 + a `resolution` x `resolution` x `resolution` np.array of the contribution of the projected data field to the voxel grid from all of the particles See Also @@ -863,7 +852,7 @@ def render_gas( @jit(nopython=True, fastmath=True) def render_voxel_to_array(data, center, width): - output = zeros((data.shape[0], data.shape[1])) + output = np.zeros((data.shape[0], data.shape[1])) for i in range(data.shape[0]): for j in range(data.shape[1]): @@ -880,23 +869,23 @@ def render_voxel_to_array(data, center, width): def visualise_render( - render: ndarray, + render: np.ndarray, centers: List[float], widths: Union[List[float], float], cmap: str = "viridis", return_type: Literal["all", "lighten", "add"] = "lighten", norm: Union[List["plt.Normalize"], "plt.Normalize", None] = None, -) -> Tuple[Union[List[ndarray], ndarray], List["plt.Normalize"]]: +) -> Tuple[Union[List[np.ndarray], np.ndarray], List["plt.Normalize"]]: """ Visualises a render with multiple centers and widths. Parameters ---------- - render : np.array + render : np.np.array The render to visualise. You should scale this appropriately before using this function (e.g. use a logarithmic transform!) - and pass in the 'value' array, not the original cosmo_array or + and pass in the 'value' np.array, not the original cosmo_array or unyt_array. centers : list[float] @@ -921,7 +910,7 @@ def visualise_render( Returns ------- - list[np.array] | np.array + list[np.np.array] | np.np.array The images of the rendering functions. If return_type is "all", this will be a list of images. If return_type is "lighten" or "add", this will be a single image. @@ -938,7 +927,7 @@ def visualise_render( elif not isinstance(norm, list): norm = [norm] * len(centers) - colors = plt.get_cmap(cmap)(linspace(0, 1, len(centers)))[:, :3] + colors = plt.get_cmap(cmap)(np.linspace(0, 1, len(centers)))[:, :3] images = [ n(render_voxel_to_array(render, center, width)) @@ -946,7 +935,7 @@ def visualise_render( ] images = [ - array([color[0] * x, color[1] * x, color[2] * x]).T + np.array([color[0] * x, color[1] * x, color[2] * x]).T for color, x in zip(colors, images) ] @@ -954,7 +943,7 @@ def visualise_render( return images, norm if return_type == "lighten": - return np_max(images, axis=0), norm + return np.max(images, axis=0), norm if return_type == "add": return sum(images), norm @@ -993,10 +982,10 @@ def visualise_render_options( if isinstance(widths, float): widths = [widths] * len(centers) - colors = plt.get_cmap(cmap)(linspace(0, 1, len(centers)))[:, :3] + colors = plt.get_cmap(cmap)(np.linspace(0, 1, len(centers)))[:, :3] for center, width, color in zip(centers, widths, colors): - xs = linspace(center - 5.0 * width, center + 5.0 * width, 100) + xs = np.linspace(center - 5.0 * width, center + 5.0 * width, 100) ys = [ exp(-0.5 * ((center - x) / width) ** 2) / (width * sqrt(2.0 * pi)) for x in xs From 20534b4902e7d19903cdc91ab28da95910a039dd Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 17:54:54 +0000 Subject: [PATCH 081/125] Prepared to create helper function in vis functions. --- swiftsimio/visualisation/projection.py | 7 +- swiftsimio/visualisation/ray_trace.py | 15 ++- swiftsimio/visualisation/slice.py | 113 ++++++++++++++-------- swiftsimio/visualisation/volume_render.py | 50 ++++++---- 4 files changed, 123 insertions(+), 62 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index fbb40bf7..bab5ec47 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -140,15 +140,19 @@ def project_pixel_grid( x_max = box_x y_min = np.zeros_like(box_y) y_max = box_y + # + # x_range = x_max - x_min y_range = y_max - y_min + # # Deal with non-cubic boxes: # we always use the maximum of x_range and y_range to normalise the coordinates # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) + # if backend == "histogram": hsml = np.empty_like(m) # not used anyway for this backend else: @@ -175,10 +179,11 @@ def project_pixel_grid( x += rotation_center[0] y += rotation_center[1] z += rotation_center[2] - else: x, y, z = data.coordinates.T + # ------------------------------ + if z_slice_included: combined_mask = np.logical_and( mask, np.logical_and(z <= z_max, z >= z_min) diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index bbf5ce70..cabf8635 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -250,14 +250,20 @@ def panel_pixel_grid( box_x, box_y, box_z = data.metadata.boxsize + # + # + # # Set the limits of the image. + # + # if region is not None: x_min, x_max, y_min, y_max = region[:4] if len(region) == 6: + # z_min, z_max = region[4:] else: - z_min = unyt.unyt_quantity(0.0, units=box_z.units) + z_min = np.zeros_like(box_z) z_max = box_z else: x_min = np.zeros_like(box_x) @@ -269,12 +275,17 @@ def panel_pixel_grid( x_range = x_max - x_min y_range = y_max - y_min + # # Deal with non-cubic boxes: # we always use the maximum of x_range and y_range to normalise the coordinates # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) + # + # + # + # hsml = ( data.smoothing_lengths if hasattr(data, "smoothing_lengths") @@ -301,6 +312,8 @@ def panel_pixel_grid( else: x, y, z = data.coordinates.T + # ------------------------------ + return core_panels( x=x[mask] / max_range, y=y[mask] / max_range, diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index f9782186..42517986 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -3,7 +3,7 @@ """ from typing import Union, Optional -from numpy import float32, array, ones, matmul, zeros_like +import numpy as np from swiftsimio import SWIFTDataset, cosmo_array, cosmo_quantity from swiftsimio.visualisation.slice_backends import backends, backends_parallel from swiftsimio.visualisation.smoothing_length import backends_get_hsml @@ -18,7 +18,7 @@ def slice_gas_pixel_grid( z_slice: Optional[cosmo_quantity] = None, project: Union[str, None] = "masses", parallel: bool = False, - rotation_matrix: Union[None, array] = None, + rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, region: Union[None, cosmo_array] = None, backend: str = "sph", @@ -34,7 +34,7 @@ def slice_gas_pixel_grid( Dataset from which slice is extracted resolution : int - Specifies size of return array + Specifies size of return np.array z_slice : cosmo_quantity Specifies the location along the z-axis where the slice is to be @@ -51,12 +51,12 @@ def slice_gas_pixel_grid( defaults to False, but can speed up the creation of large images significantly at the cost of increased memory usage. - rotation_matrix: np.array, optional + rotation_matrix: np.np.array, optional Rotation matrix (3x3) that describes the rotation of the box around ``rotation_center``. In the default case, this provides a slice perpendicular to the z axis. - rotation_center: np.array, optional + rotation_center: np.np.array, optional Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. @@ -80,8 +80,8 @@ def slice_gas_pixel_grid( Returns ------- - ndarray of float32 - Creates a `resolution` x `resolution` array and returns it, + ndarray of np.float32 + Creates a `resolution` x `resolution` np.array and returns it, without appropriate units. See Also @@ -89,17 +89,18 @@ def slice_gas_pixel_grid( render_gas_voxel_grid : Creates a 3D voxel grid from a SWIFT dataset """ + data = data.gas if z_slice is None: - z_slice = zeros_like(data.metadata.boxsize[0]) + z_slice = np.zeros_like(data.metadata.boxsize[0]) - number_of_gas_particles = data.gas.coordinates.shape[0] + number_of_particles = data.coordinates.shape[0] if project is None: - m = ones(number_of_gas_particles, dtype=float32) + m = np.ones(number_of_particles, dtype=np.float32) else: - m = getattr(data.gas, project) - if data.gas.coordinates.comoving: + m = getattr(data, project) + if data.coordinates.comoving: if not m.compatible_with_comoving(): raise AttributeError( f'Physical quantity "{project}" is not compatible with comoving coordinates!' @@ -111,54 +112,80 @@ def slice_gas_pixel_grid( ) m = m.value + # + # + # + # box_x, box_y, box_z = data.metadata.boxsize - if z_slice > box_z or z_slice < zeros_like(box_z): + if z_slice > box_z or z_slice < np.zeros_like(box_z): raise ValueError("Please enter a slice value inside the box.") # Set the limits of the image. + # + # if region is not None: x_min, x_max, y_min, y_max = region + # + # + # + # + # + # + # else: - x_min = zeros_like(box_x) + x_min = np.zeros_like(box_x) x_max = box_x - y_min = zeros_like(box_y) + y_min = np.zeros_like(box_y) y_max = box_y + # + # x_range = x_max - x_min y_range = y_max - y_min + # # Deal with non-cubic boxes: # we always use the maximum of x_range and y_range to normalise the coordinates # empty pixels in the resulting square image are trimmed afterwards max_range = max(x_range, y_range) + # + # + # + # + hsml = backends_get_hsml[backend](data) + # + # + # + # + if data.coordinates.comoving: + if not hsml.compatible_with_comoving(): + raise AttributeError( + "Physical smoothing length is not compatible with comoving coordinates!" + ) + else: + if not hsml.compatible_with_physical(): + raise AttributeError( + "Comoving smoothing length is not compatible with physical coordinates!" + ) + if rotation_center is not None: # Rotate co-ordinates as required - x, y, z = matmul(rotation_matrix, (data.gas.coordinates - rotation_center).T) + x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) x += rotation_center[0] y += rotation_center[1] z += rotation_center[2] - - z_center = rotation_center[2] - else: - x, y, z = data.gas.coordinates.T + x, y, z = data.coordinates.T - z_center = zeros_like(box_z) + # ------------------------------ - hsml = backends_get_hsml[backend](data) - if data.gas.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - f"Physical smoothing length is not compatible with comoving coordinates!" - ) + if rotation_center is not None: + z_center = rotation_center[2] else: - if not hsml.compatible_with_physical(): - raise AttributeError( - f"Comoving smoothing length is not compatible with physical coordinates!" - ) + z_center = np.zeros_like(box_z) if periodic: periodic_box_x = box_x / max_range @@ -201,7 +228,7 @@ def slice_gas( z_slice: Optional[cosmo_quantity] = None, project: Union[str, None] = "masses", parallel: bool = False, - rotation_matrix: Union[None, array] = None, + rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, region: Union[None, cosmo_array] = None, backend: str = "sph", @@ -216,7 +243,7 @@ def slice_gas( Dataset from which slice is extracted resolution : int - Specifies size of return array + Specifies size of return np.array z_slice : cosmo_quantity Specifies the location along the z-axis where the slice is to be @@ -233,16 +260,16 @@ def slice_gas( defaults to False, but can speed up the creation of large images significantly at the cost of increased memory usage. - rotation_matrix: np.array, optional + rotation_matrix: np.np.array, optional Rotation matrix (3x3) that describes the rotation of the box around ``rotation_center``. In the default case, this provides a slice perpendicular to the z axis. - rotation_center: np.array, optional + rotation_center: np.np.array, optional Center of the rotation. If you are trying to rotate around a galaxy, this should be the most bound particle. - region : array, optional + region : np.array, optional determines where the image will be created (this corresponds to the left and right-hand edges, and top and bottom edges) if it is not None. It should have a length of four, and take the form: @@ -262,8 +289,8 @@ def slice_gas( Returns ------- - ndarray of float32 - a `resolution` x `resolution` array of the contribution + ndarray of np.float32 + a `resolution` x `resolution` np.array of the contribution of the projected data field to the voxel grid from all of the particles See Also @@ -278,7 +305,7 @@ def slice_gas( """ if z_slice is None: - z_slice = zeros_like(data.metadata.boxsize[0]) + z_slice = np.zeros_like(data.metadata.boxsize[0]) image = slice_gas_pixel_grid( data, @@ -297,7 +324,7 @@ def slice_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 3) + units = 1.0 / (max_range**3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units( @@ -305,17 +332,17 @@ def slice_gas( ) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 3) + units = 1.0 / (max_range**3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 3a74d3cb..b6952ef8 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -637,10 +637,10 @@ def render_gas_voxel_grid( """ data = data.gas # coerce rest of this function to share with other vis functions - number_of_gas_particles = data.particle_ids.size + number_of_particles = data.particle_ids.size if project is None: - m = np.ones(number_of_gas_particles, dtype=np.float32) + m = np.ones(number_of_particles, dtype=np.float32) else: m = getattr(data, project) if data.coordinates.comoving: @@ -655,17 +655,33 @@ def render_gas_voxel_grid( ) m = m.value + # + # + # + # box_x, box_y, box_z = data.metadata.boxsize + # + # + # # Set the limits of the image. + # + # if region is not None: x_min, x_max, y_min, y_max, z_min, z_max = region + # + # + # + # + # + # + # else: - x_min = (0 * box_x).to(box_x.units) + x_min = np.zeros_like(box_x) x_max = box_x - y_min = (0 * box_y).to(box_y.units) + y_min = np.zeros_like(box_y) y_max = box_y - z_min = (0 * box_z).to(box_z.units) + z_min = np.zeros_like(box_z) z_max = box_z x_range = x_max - x_min @@ -681,18 +697,6 @@ def render_gas_voxel_grid( "Projection code is currently not able to handle non-cubic images" ) - # Let's just hope that the box is square otherwise we're probably SOL - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - - else: - x, y, z = data.coordinates.T - hsml = ( data.smoothing_lengths if hasattr(data, "smoothing_lengths") @@ -709,6 +713,18 @@ def render_gas_voxel_grid( "Comoving smoothing length is not compatible with physical coordinates!" ) + if rotation_center is not None: + # Rotate co-ordinates as required + x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) + + x += rotation_center[0] + y += rotation_center[1] + z += rotation_center[2] + else: + x, y, z = data.coordinates.T + + # ------------------------------ + if periodic: periodic_box_x = box_x / x_range periodic_box_y = box_y / y_range From 9530290803b5c2efed13ed0301a60c1b0a979fa5 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Mon, 10 Feb 2025 23:14:57 +0000 Subject: [PATCH 082/125] Start refactoring visualisation functions with helpers. --- swiftsimio/visualisation/_vistools.py | 81 ++++++++++++ swiftsimio/visualisation/projection.py | 104 ++++------------ swiftsimio/visualisation/ray_trace.py | 100 +++------------ swiftsimio/visualisation/slice.py | 116 +++--------------- .../visualisation/smoothing_length/sph.py | 20 ++- swiftsimio/visualisation/volume_render.py | 110 +++-------------- 6 files changed, 170 insertions(+), 361 deletions(-) create mode 100644 swiftsimio/visualisation/_vistools.py diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py new file mode 100644 index 00000000..ed8dde32 --- /dev/null +++ b/swiftsimio/visualisation/_vistools.py @@ -0,0 +1,81 @@ +import numpy as np + + +def _get_projection_field(data, field_name): + if field_name is None: + m = np.ones(data.particle_ids.size) + else: + m = getattr(data, field_name) + if data.coordinates.comoving: + if not m.compatible_with_comoving(): + raise AttributeError( + f'Physical quantity "{field_name}" is not compatible with comoving ' + "coordinates!" + ) + else: + if not m.compatible_with_physical(): + raise AttributeError( + f'Comoving quantity "{field_name}" is not compatible with physical ' + "coordinates!" + ) + m = m.value # slated for removal + return m + + +def _get_region_limits(data, region, z_slice=None, require_cubic=False): + z_slice_included = z_slice is not None + if not z_slice_included: + z_slice = np.zeros_like(data.metadata.boxsize[0]) + box_x, box_y, box_z = data.metadata.boxsize + if region is not None: + x_min, x_max, y_min, y_max = region[:4] + if len(region) == 6: + z_slice_included = True + z_min, z_max = region[4:] + else: + z_min, z_max = np.zeros_like(box_z), box_z + else: + x_min, x_max = np.zeros_like(box_x), box_x + y_min, y_max = np.zeros_like(box_y), box_y + z_min, z_max = np.zeros_like(box_z), box_z + + if z_slice_included and (z_slice > box_z) or (z_slice < np.zeros_like(box_z)): + raise ValueError("Please enter a slice value inside the box.") + + x_range = x_max - x_min + y_range = y_max - y_min + z_range = z_max - z_min + max_range = np.max([x_range, y_range]) + + if require_cubic and not ( + np.isclose(x_range, y_range) and np.isclose(x_range, z_range) + ): + raise AttributeError( + "Projection code is currently not able to handle non-cubic images" + ) + + return ( + x_min, + x_max, + y_min, + y_max, + z_min, + z_max, + x_range, + y_range, + z_range, + max_range, + ) + + +def _get_rotated_coordinates(data, rotation_matrix, rotation_center): + if rotation_center is not None: + # Rotate co-ordinates as required + x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) + + x += rotation_center[0] + y += rotation_center[1] + z += rotation_center[2] + else: + x, y, z = data.coordinates.T + return x, y, z diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index bab5ec47..7601a429 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -9,6 +9,12 @@ from swiftsimio.reader import __SWIFTGroupDataset from swiftsimio.visualisation.projection_backends import backends, backends_parallel +from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_limits, + _get_rotated_coordinates, +) scatter = backends["fast"] scatter_parallel = backends_parallel["fast"] @@ -99,92 +105,28 @@ def project_pixel_grid( array if you want it to be visualised the 'right way up'. """ - number_of_particles = data.coordinates.shape[0] - - if project is None: - m = np.ones(number_of_particles, dtype=np.float32) - else: - m = getattr(data, project) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value + m = _get_projection_field(data, project) # This provides a default 'slice it all' mask. if mask is None: mask = np.s_[:] - box_x, box_y, box_z = data.metadata.boxsize - - # Set the limits of the image. - z_slice_included = False + x_min, x_max, y_min, y_max, z_min, z_max, x_range, y_range, _, max_range = _get_region_limits( + data, region + ) - if region is not None: - x_min, x_max, y_min, y_max = region[:4] - - if len(region) == 6: - z_slice_included = True - z_min, z_max = region[4:] - else: - z_min = np.zeros_like(box_z) - z_max = box_z - else: - x_min = np.zeros_like(box_x) - x_max = box_x - y_min = np.zeros_like(box_y) - y_max = box_y - # - # - - x_range = x_max - x_min - y_range = y_max - y_min - # - - # Deal with non-cubic boxes: - # we always use the maximum of x_range and y_range to normalise the coordinates - # empty pixels in the resulting square image are trimmed afterwards - max_range = max(x_range, y_range) - - # if backend == "histogram": hsml = np.empty_like(m) # not used anyway for this backend else: - hsml = ( - data.smoothing_lengths - if hasattr(data, "smoothing_lengths") - else data.smoothing_length # backwards compatibility - ) - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - else: - x, y, z = data.coordinates.T + hsml = backends_get_hsml["sph"](data) + + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) # ------------------------------ - if z_slice_included: + if (region is not None) and len( + region + ) != 6: # if not z_slice_included: ...should refactor combined_mask = np.logical_and( mask, np.logical_and(z <= z_max, z >= z_min) ).astype(bool) @@ -192,11 +134,9 @@ def project_pixel_grid( combined_mask = mask if periodic: - periodic_box_x = box_x / max_range - periodic_box_y = box_y / max_range + periodic_box_x, periodic_box_y, _ = data.metadata.boxsize / max_range else: - periodic_box_x = 0.0 - periodic_box_y = 0.0 + periodic_box_x, periodic_box_y = 0.0, 0.0 common_arguments = dict( x=(x[combined_mask] - x_min) / max_range, @@ -429,23 +369,23 @@ def project_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index cabf8635..68b564a7 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -20,7 +20,12 @@ ) from swiftsimio.accelerated import jit, prange, NUM_THREADS -import unyt +from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_limits, + _get_rotated_coordinates, +) @jit(nopython=True, fastmath=True) @@ -226,93 +231,18 @@ def panel_pixel_grid( rotation_center: Union[None, cosmo_array] = None, ) -> cosmo_array: - number_of_particles = data.coordinates.shape[0] - - if project is None: - m = np.ones(number_of_particles, dtype=np.float32) - else: - m = getattr(data, project) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value + m = _get_projection_field(data, project) # This provides a default 'slice it all' mask. if mask is None: mask = np.s_[:] - box_x, box_y, box_z = data.metadata.boxsize - - # - # - # - # Set the limits of the image. - # - # - if region is not None: - x_min, x_max, y_min, y_max = region[:4] - - if len(region) == 6: - # - z_min, z_max = region[4:] - else: - z_min = np.zeros_like(box_z) - z_max = box_z - else: - x_min = np.zeros_like(box_x) - x_max = box_x - y_min = np.zeros_like(box_y) - y_max = box_y - z_min = np.zeros_like(box_z) - z_max = box_z - - x_range = x_max - x_min - y_range = y_max - y_min - # - - # Deal with non-cubic boxes: - # we always use the maximum of x_range and y_range to normalise the coordinates - # empty pixels in the resulting square image are trimmed afterwards - max_range = max(x_range, y_range) - - # - # - # - # - hsml = ( - data.smoothing_lengths - if hasattr(data, "smoothing_lengths") - else data.smoothing_length # backwards compatibility + x_min, x_max, y_min, y_max, z_min, z_max, _, _, _, max_range = _get_region_limits( + data, region ) - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) - - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - else: - x, y, z = data.coordinates.T + hsml = backends_get_hsml["sph"](data) - # ------------------------------ + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) return core_panels( x=x[mask] / max_range, @@ -352,23 +282,23 @@ def panel_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**2) + units = 1.0 / (max_range ** 2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index 42517986..1197dfb4 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -7,6 +7,11 @@ from swiftsimio import SWIFTDataset, cosmo_array, cosmo_quantity from swiftsimio.visualisation.slice_backends import backends, backends_parallel from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_limits, + _get_rotated_coordinates, +) slice_scatter = backends["sph"] slice_scatter_parallel = backends_parallel["sph"] @@ -91,110 +96,27 @@ def slice_gas_pixel_grid( """ data = data.gas - if z_slice is None: - z_slice = np.zeros_like(data.metadata.boxsize[0]) + m = _get_projection_field(data, project) - number_of_particles = data.coordinates.shape[0] + x_min, x_max, y_min, y_max, _, _, x_range, y_range, _, max_range = _get_region_limits( + data, region, z_slice=z_slice + ) - if project is None: - m = np.ones(number_of_particles, dtype=np.float32) - else: - m = getattr(data, project) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value - - # - # - # - # - box_x, box_y, box_z = data.metadata.boxsize - - if z_slice > box_z or z_slice < np.zeros_like(box_z): - raise ValueError("Please enter a slice value inside the box.") - - # Set the limits of the image. - # - # - if region is not None: - x_min, x_max, y_min, y_max = region - # - # - # - # - # - # - # - else: - x_min = np.zeros_like(box_x) - x_max = box_x - y_min = np.zeros_like(box_y) - y_max = box_y - # - # - - x_range = x_max - x_min - y_range = y_max - y_min - # - - # Deal with non-cubic boxes: - # we always use the maximum of x_range and y_range to normalise the coordinates - # empty pixels in the resulting square image are trimmed afterwards - max_range = max(x_range, y_range) - - # - # - # - # hsml = backends_get_hsml[backend](data) - # - # - # - # - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - else: - x, y, z = data.coordinates.T - - # ------------------------------ + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) if rotation_center is not None: z_center = rotation_center[2] else: - z_center = np.zeros_like(box_z) + z_center = np.zeros_like(data.metadata.boxsize[2]) if periodic: - periodic_box_x = box_x / max_range - periodic_box_y = box_y / max_range - periodic_box_z = box_z / max_range + periodic_box_x, periodic_box_y, periodic_box_z = ( + data.metadata.boxsize / max_range + ) else: - periodic_box_x = 0.0 - periodic_box_y = 0.0 - periodic_box_z = 0.0 + periodic_box_x, periodic_box_y, periodic_box_z = 0.0, 0.0, 0.0 # determine the effective number of pixels for each dimension xres = int(resolution * x_range / max_range) @@ -324,7 +246,7 @@ def slice_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range**3) + units = 1.0 / (max_range ** 3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units( @@ -332,17 +254,17 @@ def slice_gas( ) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**3) + units = 1.0 / (max_range ** 3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) diff --git a/swiftsimio/visualisation/smoothing_length/sph.py b/swiftsimio/visualisation/smoothing_length/sph.py index 29d09446..a38f1035 100644 --- a/swiftsimio/visualisation/smoothing_length/sph.py +++ b/swiftsimio/visualisation/smoothing_length/sph.py @@ -14,9 +14,19 @@ def get_hsml(data: SWIFTDataset) -> cosmo_array: ------- The extracted smoothing lengths. """ - try: - hsml = data.gas.smoothing_lengths - except AttributeError: - # Backwards compatibility - hsml = data.gas.smoothing_length + hsml = ( + data.smoothing_lengths + if hasattr(data, "smoothing_lengths") + else data.smoothing_length # backwards compatibility + ) + if data.coordinates.comoving: + if not hsml.compatible_with_comoving(): + raise AttributeError( + "Physical smoothing length is not compatible with comoving coordinates!" + ) + else: + if not hsml.compatible_with_physical(): + raise AttributeError( + "Comoving smoothing length is not compatible with physical coordinates!" + ) return hsml diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index b6952ef8..97905c75 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -12,6 +12,12 @@ from swiftsimio.optional_packages import plt from swiftsimio.visualisation.slice_backends.sph import kernel, kernel_gamma +from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation._vistools import ( + _get_projection_field, + _get_region_limits, + _get_rotated_coordinates, +) @jit(nopython=True, fastmath=True) @@ -635,104 +641,24 @@ def render_gas_voxel_grid( slice_gas_pixel_grid : Creates a 2D slice of a SWIFT dataset """ - data = data.gas # coerce rest of this function to share with other vis functions + data = data.gas - number_of_particles = data.particle_ids.size + m = _get_projection_field(data, project) - if project is None: - m = np.ones(number_of_particles, dtype=np.float32) - else: - m = getattr(data, project) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{project}" is not compatible with comoving coordinates!' - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{project}" is not compatible with physical coordinates!' - ) - m = m.value - - # - # - # - # - box_x, box_y, box_z = data.metadata.boxsize - - # - # - # - # Set the limits of the image. - # - # - if region is not None: - x_min, x_max, y_min, y_max, z_min, z_max = region - # - # - # - # - # - # - # - else: - x_min = np.zeros_like(box_x) - x_max = box_x - y_min = np.zeros_like(box_y) - y_max = box_y - z_min = np.zeros_like(box_z) - z_max = box_z - - x_range = x_max - x_min - y_range = y_max - y_min - z_range = z_max - z_min - - # Test that we've got a cubic box - if not ( - np.isclose(x_range.value, y_range.value) - and np.isclose(x_range.value, z_range.value) - ): - raise AttributeError( - "Projection code is currently not able to handle non-cubic images" - ) - - hsml = ( - data.smoothing_lengths - if hasattr(data, "smoothing_lengths") - else data.smoothing_length # backwards compatibility + x_min, x_max, y_min, y_max, z_min, z_max, x_range, y_range, z_range, _ = _get_region_limits( + data, region, require_cubic=True ) - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - if rotation_center is not None: - # Rotate co-ordinates as required - x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) + hsml = backends_get_hsml["sph"](data) - x += rotation_center[0] - y += rotation_center[1] - z += rotation_center[2] - else: - x, y, z = data.coordinates.T - - # ------------------------------ + x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) if periodic: - periodic_box_x = box_x / x_range - periodic_box_y = box_y / y_range - periodic_box_z = box_z / z_range + periodic_box_x, periodic_box_y, periodic_box_z = ( + data.metadata.boxsize / cosmo_array([x_range, y_range, z_range]) + ) else: - periodic_box_x = 0.0 - periodic_box_y = 0.0 - periodic_box_z = 0.0 + periodic_box_x, periodic_box_y, periodic_box_z = 0.0, 0.0, 0.0 arguments = dict( x=(x - x_min) / x_range, @@ -850,14 +776,14 @@ def render_gas( * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) From 63bfb25e99b6f3965ec6ea134ab25866829d340d Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 11 Feb 2025 13:13:40 +0000 Subject: [PATCH 083/125] Move more to helper functions and cleanup. --- swiftsimio/visualisation/_vistools.py | 34 ++++---- swiftsimio/visualisation/projection.py | 78 +++++++------------ swiftsimio/visualisation/ray_trace.py | 41 +++++----- swiftsimio/visualisation/slice.py | 57 ++++++-------- .../smoothing_length/__init__.py | 2 + swiftsimio/visualisation/volume_render.py | 62 ++++++--------- 6 files changed, 119 insertions(+), 155 deletions(-) diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py index ed8dde32..f12b84eb 100644 --- a/swiftsimio/visualisation/_vistools.py +++ b/swiftsimio/visualisation/_vistools.py @@ -22,7 +22,7 @@ def _get_projection_field(data, field_name): return m -def _get_region_limits(data, region, z_slice=None, require_cubic=False): +def _get_region_info(data, region, z_slice=None, require_cubic=False, periodic=True): z_slice_included = z_slice is not None if not z_slice_included: z_slice = np.zeros_like(data.metadata.boxsize[0]) @@ -51,22 +51,30 @@ def _get_region_limits(data, region, z_slice=None, require_cubic=False): np.isclose(x_range, y_range) and np.isclose(x_range, z_range) ): raise AttributeError( - "Projection code is currently not able to handle non-cubic images" + "Projection code is currently not able to handle non-cubic images." ) - return ( - x_min, - x_max, - y_min, - y_max, - z_min, - z_max, - x_range, - y_range, - z_range, - max_range, + periodic_box_x, periodic_box_y, periodic_box_z = ( + data.metadata.boxsize / max_range if periodic else np.zeros(3) ) + return { + "x_min": x_min, + "x_max": x_max, + "y_min": y_min, + "y_max": y_max, + "z_min": z_min, + "z_max": z_max, + "x_range": x_range, + "y_range": y_range, + "z_range": z_range, + "max_range": max_range, + "z_slice_included": z_slice_included, + "periodic_box_x": periodic_box_x, + "periodic_box_y": periodic_box_y, + "periodic_box_z": periodic_box_z, + } + def _get_rotated_coordinates(data, rotation_matrix, rotation_center): if rotation_center is not None: diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 7601a429..c0ee401e 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -12,7 +12,7 @@ from swiftsimio.visualisation.smoothing_length import backends_get_hsml from swiftsimio.visualisation._vistools import ( _get_projection_field, - _get_region_limits, + _get_region_info, _get_rotated_coordinates, ) @@ -106,56 +106,38 @@ def project_pixel_grid( """ m = _get_projection_field(data, project) - - # This provides a default 'slice it all' mask. - if mask is None: - mask = np.s_[:] - - x_min, x_max, y_min, y_max, z_min, z_max, x_range, y_range, _, max_range = _get_region_limits( - data, region - ) - - if backend == "histogram": - hsml = np.empty_like(m) # not used anyway for this backend - else: - hsml = backends_get_hsml["sph"](data) - + region_info = _get_region_info(data, region) + hsml = backends_get_hsml["sph" if backend != "histogram" else "histogram"](data) x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) - - # ------------------------------ - - if (region is not None) and len( - region - ) != 6: # if not z_slice_included: ...should refactor - combined_mask = np.logical_and( - mask, np.logical_and(z <= z_max, z >= z_min) + mask = mask if mask is not None else np.s_[...] + if not region_info["z_slice_included"]: + mask = np.logical_and( + mask, + np.logical_and(z <= region_info["z_max"], z >= region_info["z_min"]), ).astype(bool) - else: - combined_mask = mask - - if periodic: - periodic_box_x, periodic_box_y, _ = data.metadata.boxsize / max_range - else: - periodic_box_x, periodic_box_y = 0.0, 0.0 - common_arguments = dict( - x=(x[combined_mask] - x_min) / max_range, - y=(y[combined_mask] - y_min) / max_range, - m=m[combined_mask], - h=hsml[combined_mask] / max_range, + arguments = dict( + x=(x[mask] - region_info["x_min"]) / region_info["max_range"], + y=(y[mask] - region_info["y_min"]) / region_info["max_range"], + m=m[mask], + h=hsml[mask] / region_info["max_range"], res=resolution, - box_x=periodic_box_x, - box_y=periodic_box_y, + box_x=region_info["periodic_box_x"], + box_y=region_info["periodic_box_y"], + ) + image = ( + backends_parallel[backend](**arguments) + if parallel + else backends[backend](**arguments) ) - - if parallel: - image = backends_parallel[backend](**common_arguments) - else: - image = backends[backend](**common_arguments) # determine the effective number of pixels for each dimension - xres = int(np.ceil(resolution * (x_range / max_range))) - yres = int(np.ceil(resolution * (y_range / max_range))) + xres = int( + np.ceil(resolution * (region_info["x_range"] / region_info["max_range"])) + ) + yres = int( + np.ceil(resolution * (region_info["y_range"] / region_info["max_range"])) + ) # trim the image to remove empty pixels return image[:xres, :yres] @@ -369,23 +351,23 @@ def project_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 68b564a7..480fb516 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -23,7 +23,7 @@ from swiftsimio.visualisation.smoothing_length import backends_get_hsml from swiftsimio.visualisation._vistools import ( _get_projection_field, - _get_region_limits, + _get_region_info, _get_rotated_coordinates, ) @@ -232,28 +232,21 @@ def panel_pixel_grid( ) -> cosmo_array: m = _get_projection_field(data, project) - - # This provides a default 'slice it all' mask. - if mask is None: - mask = np.s_[:] - - x_min, x_max, y_min, y_max, z_min, z_max, _, _, _, max_range = _get_region_limits( - data, region - ) + region_info = _get_region_info(data, region) hsml = backends_get_hsml["sph"](data) - x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) + mask = np.s_[...] if mask is None else mask return core_panels( - x=x[mask] / max_range, - y=y[mask] / max_range, + x=x[mask] / region_info["max_range"], + y=y[mask] / region_info["max_range"], z=z[mask], - h=hsml[mask] / max_range, + h=hsml[mask] / region_info["max_range"], m=m[mask], res=resolution, panels=panels, - min_z=z_min, - max_z=z_max, + min_z=region_info["z_min"], + max_z=region_info["z_max"], ) @@ -282,23 +275,23 @@ def panel_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units(1.0 / (x_range.units * y_range.units)) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 2) + units = 1.0 / (max_range**2) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 2) + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 2 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 else: new_cosmo_factor = coord_cosmo_factor ** (-2) @@ -388,7 +381,11 @@ def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): # if factor > 0.5: # factor = 1.0 - factor -# value += 1 / (width * np.sqrt(2.0 * np.pi)) * np.exp(-0.5 * ((i - center) / width) ** 2) +# value += ( +# 1 +# / (width * np.sqrt(2.0 * np.pi)) +# * np.exp(-0.5 * ((i - center) / width) ** 2) +# ) # output[x, y] = value @@ -403,7 +400,9 @@ def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): # width = 0.05 # centers = [np.mean(log_data) + x * std for x in [0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]] # # %% -# colors = plt.get_cmap("swift.nineteen_eighty_nine")((np.linspace(0, 1, len(centers))))[:, :3] +# colors = plt.get_cmap("swift.nineteen_eighty_nine")((np.linspace(0, 1, len(centers))))[ +# :, :3 +# ] # grids = [ # make_grid(color, center, width) for color, center in zip(colors, centers) diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index 1197dfb4..7820b436 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -9,7 +9,7 @@ from swiftsimio.visualisation.smoothing_length import backends_get_hsml from swiftsimio.visualisation._vistools import ( _get_projection_field, - _get_region_limits, + _get_region_info, _get_rotated_coordinates, ) @@ -97,43 +97,31 @@ def slice_gas_pixel_grid( data = data.gas m = _get_projection_field(data, project) - - x_min, x_max, y_min, y_max, _, _, x_range, y_range, _, max_range = _get_region_limits( - data, region, z_slice=z_slice - ) - + region_info = _get_region_info(data, region, z_slice=z_slice, periodic=periodic) hsml = backends_get_hsml[backend](data) - x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) - - if rotation_center is not None: - z_center = rotation_center[2] - else: - z_center = np.zeros_like(data.metadata.boxsize[2]) - - if periodic: - periodic_box_x, periodic_box_y, periodic_box_z = ( - data.metadata.boxsize / max_range - ) - else: - periodic_box_x, periodic_box_y, periodic_box_z = 0.0, 0.0, 0.0 + z_center = ( + rotation_center[2] + if rotation_center is not None + else np.zeros_like(data.metadata.boxsize[2]) + ) # determine the effective number of pixels for each dimension - xres = int(resolution * x_range / max_range) - yres = int(resolution * y_range / max_range) + xres = int(resolution * region_info["x_range"] / region_info["max_range"]) + yres = int(resolution * region_info["y_range"] / region_info["max_range"]) common_parameters = dict( - x=(x - x_min) / max_range, - y=(y - y_min) / max_range, - z=z / max_range, + x=(x - region_info["x_min"]) / region_info["max_range"], + y=(y - region_info["y_min"]) / region_info["max_range"], + z=z / region_info["max_range"], m=m, - h=hsml / max_range, - z_slice=(z_center + z_slice) / max_range, + h=hsml / region_info["max_range"], + z_slice=(z_center + z_slice) / region_info["max_range"], xres=xres, yres=yres, - box_x=periodic_box_x, - box_y=periodic_box_y, - box_z=periodic_box_z, + box_x=region_info["periodic_box_x"], + box_y=region_info["periodic_box_y"], + box_z=region_info["periodic_box_z"], ) if parallel: @@ -226,8 +214,7 @@ def slice_gas( appropriate """ - if z_slice is None: - z_slice = np.zeros_like(data.metadata.boxsize[0]) + z_slice = np.zeros_like(data.metadata.boxsize[0]) if z_slice is None else z_slice image = slice_gas_pixel_grid( data, @@ -246,7 +233,7 @@ def slice_gas( x_range = region[1] - region[0] y_range = region[3] - region[2] max_range = max(x_range, y_range) - units = 1.0 / (max_range ** 3) + units = 1.0 / (max_range**3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... units.convert_to_units( @@ -254,17 +241,17 @@ def slice_gas( ) else: max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range ** 3) + units = 1.0 / (max_range**3) # Unfortunately this is required to prevent us from {over,under}flowing # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) diff --git a/swiftsimio/visualisation/smoothing_length/__init__.py b/swiftsimio/visualisation/smoothing_length/__init__.py index 7cccb0fe..2417ab74 100644 --- a/swiftsimio/visualisation/smoothing_length/__init__.py +++ b/swiftsimio/visualisation/smoothing_length/__init__.py @@ -1,8 +1,10 @@ +import numpy as np from .sph import get_hsml as sph_get_hsml from .nearest_neighbours import get_hsml as nearest_neighbours_get_hsml from .generate import generate_smoothing_lengths backends_get_hsml = { + "histogram": lambda m: np.empty_like(m), "sph": sph_get_hsml, "nearest_neighbours": nearest_neighbours_get_hsml, } diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 97905c75..11f26878 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -15,7 +15,7 @@ from swiftsimio.visualisation.smoothing_length import backends_get_hsml from swiftsimio.visualisation._vistools import ( _get_projection_field, - _get_region_limits, + _get_region_info, _get_rotated_coordinates, ) @@ -180,16 +180,17 @@ def scatter( else: # Now we loop over the square of cells that the kernel lives in for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we'll segfault + # Ensure that the lowest x value is 0, otherwise we segfault max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the np.array bounds, - # otherwise we'll segfault (oops). + # Ensure that the highest x value lies within the np.array + # bounds, otherwise we'll segfault (oops). min( particle_cell_x + cells_spanned, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that + # our x, y are all in a box of [0, 1]; calculate the distance + # to the cell centre distance_x = ( np.float32(cell_x) + 0.5 ) * pixel_width - np.float32(x_pos) @@ -402,16 +403,17 @@ def scatter_limited_z( else: # Now we loop over the square of cells that the kernel lives in for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we'll segfault + # Ensure that the lowest x value is 0, otherwise we segfault max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the np.array bounds, - # otherwise we'll segfault (oops). + # Ensure that the highest x value lies within the np.array + # bounds, otherwise we'll segfault (oops). min( particle_cell_x + cells_spanned, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that + # our x, y are all in a box of [0, 1]; calculate the distance + # to the cell centre distance_x = ( np.float32(cell_x) + 0.5 ) * pixel_width - np.float32(x_pos) @@ -644,38 +646,22 @@ def render_gas_voxel_grid( data = data.gas m = _get_projection_field(data, project) - - x_min, x_max, y_min, y_max, z_min, z_max, x_range, y_range, z_range, _ = _get_region_limits( - data, region, require_cubic=True - ) - + region_info = _get_region_info(data, region, require_cubic=True) hsml = backends_get_hsml["sph"](data) - x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) - if periodic: - periodic_box_x, periodic_box_y, periodic_box_z = ( - data.metadata.boxsize / cosmo_array([x_range, y_range, z_range]) - ) - else: - periodic_box_x, periodic_box_y, periodic_box_z = 0.0, 0.0, 0.0 - arguments = dict( - x=(x - x_min) / x_range, - y=(y - y_min) / y_range, - z=(z - z_min) / z_range, + x=(x - region_info["x_min"]) / region_info["x_range"], + y=(y - region_info["y_min"]) / region_info["y_range"], + z=(z - region_info["z_min"]) / region_info["z_range"], m=m, - h=hsml / x_range, + h=hsml / region_info["x_range"], # cubic so x_range == y_range == z_range res=resolution, - box_x=periodic_box_x, - box_y=periodic_box_y, - box_z=periodic_box_z, + box_x=region_info["periodic_box_x"], + box_y=region_info["periodic_box_y"], + box_z=region_info["periodic_box_z"], ) - - if parallel: - image = scatter_parallel(**arguments) - else: - image = scatter(**arguments) + image = scatter_parallel(**arguments) if parallel else scatter(**arguments) return image @@ -776,14 +762,14 @@ def render_gas( * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor if project is not None: units *= getattr(data.gas, project).units project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor ** 3 + new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 else: new_cosmo_factor = coord_cosmo_factor ** (-3) From 530c78ffc526b4a540f025f1f8d2114f3a9d2fd7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 11 Feb 2025 21:00:34 +0000 Subject: [PATCH 084/125] Continue cleanup of visualisation functions. --- swiftsimio/visualisation/_vistools.py | 5 ++--- swiftsimio/visualisation/projection.py | 20 ++++++------------- .../visualisation/projection_backends/fast.py | 10 +--------- 3 files changed, 9 insertions(+), 26 deletions(-) diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py index f12b84eb..00b6c7d9 100644 --- a/swiftsimio/visualisation/_vistools.py +++ b/swiftsimio/visualisation/_vistools.py @@ -3,7 +3,7 @@ def _get_projection_field(data, field_name): if field_name is None: - m = np.ones(data.particle_ids.size) + m = np.ones_like(data.particle_ids) else: m = getattr(data, field_name) if data.coordinates.comoving: @@ -18,7 +18,6 @@ def _get_projection_field(data, field_name): f'Comoving quantity "{field_name}" is not compatible with physical ' "coordinates!" ) - m = m.value # slated for removal return m @@ -45,7 +44,7 @@ def _get_region_info(data, region, z_slice=None, require_cubic=False, periodic=T x_range = x_max - x_min y_range = y_max - y_min z_range = z_max - z_min - max_range = np.max([x_range, y_range]) + max_range = np.r_[x_range, y_range].max() if require_cubic and not ( np.isclose(x_range, y_range) and np.isclose(x_range, z_range) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index c0ee401e..99ccf828 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -130,6 +130,7 @@ def project_pixel_grid( if parallel else backends[backend](**arguments) ) + raise RuntimeError # determine the effective number of pixels for each dimension xres = int( @@ -334,6 +335,7 @@ def project_gas( array if you want it to be visualised the 'right way up'. """ + # REFACTOR? can we get this image to come out as a cosmo_array already? image = project_gas_pixel_grid( data=data, resolution=resolution, @@ -347,20 +349,10 @@ def project_gas( periodic=periodic, ) - if region is not None: - x_range = region[1] - region[0] - y_range = region[3] - region[2] - max_range = max(x_range, y_range) - units = 1.0 / (max_range**2) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units(1.0 / (x_range.units * y_range.units)) - else: - max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**2) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) + region_info = _get_region_info(data.gas, region) + units = 1.0 / (region_info["max_range"] ** 2) + # Unfortunately this is required to prevent us from {over,under}flowing the units: + units.convert_to_units(1.0 / data.metadata.boxsize.units**2) comoving = data.gas.coordinates.comoving coord_cosmo_factor = data.gas.coordinates.cosmo_factor diff --git a/swiftsimio/visualisation/projection_backends/fast.py b/swiftsimio/visualisation/projection_backends/fast.py index 4555e2cb..7f0f715d 100644 --- a/swiftsimio/visualisation/projection_backends/fast.py +++ b/swiftsimio/visualisation/projection_backends/fast.py @@ -2,14 +2,9 @@ Fast backend. This uses float32 precision and no special cases. -""" - -""" The original smoothing code. This provides no renormalisation. """ - -from typing import Union from math import sqrt from numpy import float64, float32, int32, zeros, ndarray @@ -17,10 +12,7 @@ from swiftsimio.visualisation.projection_backends.kernels import ( kernel_single_precision as kernel, ) -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_constant, - kernel_gamma, -) +from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma @jit(nopython=True, fastmath=True) From 79780f97bdd2c0fcf6f0bb3a49478508a116cea0 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 12 Feb 2025 01:39:16 +0000 Subject: [PATCH 085/125] Consolidate handling cosmo_arrays in visualisation and big namespace cleanup. --- swiftsimio/_array_functions.py | 15 +- swiftsimio/visualisation/__init__.py | 7 +- swiftsimio/visualisation/_vistools.py | 93 ++- swiftsimio/visualisation/projection.py | 139 +--- .../projection_backends/__init__.py | 30 +- swiftsimio/visualisation/ray_trace.py | 361 +-------- .../ray_trace_backends/__init__.py | 18 + .../ray_trace_backends/core_panels.py | 332 +++++++++ swiftsimio/visualisation/slice.py | 141 +--- .../visualisation/slice_backends/__init__.py | 11 +- .../visualisation/smoothing_length/sph.py | 13 +- swiftsimio/visualisation/volume_render.py | 682 +----------------- .../volume_render_backends/__init__.py | 18 + .../volume_render_backends/scatter.py | 568 +++++++++++++++ tests/test_visualisation.py | 384 ++++++---- 15 files changed, 1328 insertions(+), 1484 deletions(-) create mode 100644 swiftsimio/visualisation/ray_trace_backends/__init__.py create mode 100644 swiftsimio/visualisation/ray_trace_backends/core_panels.py create mode 100644 swiftsimio/visualisation/volume_render_backends/__init__.py create mode 100644 swiftsimio/visualisation/volume_render_backends/scatter.py diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index fe72a05d..5ee38806 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -122,7 +122,9 @@ # numpy functions (we will actually wrap the functions below): -def _copy_cosmo_array_attributes(from_ca: object, to_ca: object) -> object: +def _copy_cosmo_array_attributes( + from_ca: object, to_ca: object, copy_units=False +) -> object: """ Copy :class:`~swiftsimio.objects.cosmo_array` attributes across two objects. @@ -138,6 +140,9 @@ def _copy_cosmo_array_attributes(from_ca: object, to_ca: object) -> object: to_ca : :obj:`object` The destination object. + copy_units : bool + If ``True`` also copy ``units`` attribute (usually let :mod:`unyt` handle this). + Returns ------- out : :obj:`object` @@ -148,6 +153,8 @@ def _copy_cosmo_array_attributes(from_ca: object, to_ca: object) -> object: and isinstance(from_ca, objects.cosmo_array) ): return to_ca + if copy_units: + to_ca.units = from_ca.units to_ca.cosmo_factor = from_ca.cosmo_factor to_ca.comoving = from_ca.comoving to_ca.valid_transform = from_ca.valid_transform @@ -722,18 +729,18 @@ def _arctan2_cosmo_factor( f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor( diff --git a/swiftsimio/visualisation/__init__.py b/swiftsimio/visualisation/__init__.py index 05928f1b..b496dede 100644 --- a/swiftsimio/visualisation/__init__.py +++ b/swiftsimio/visualisation/__init__.py @@ -2,7 +2,8 @@ Visualisation sub-module for swiftismio. """ -from .projection import scatter, project_gas, project_gas_pixel_grid -from .slice import slice_scatter as slice -from .slice import slice_gas, slice_gas_pixel_grid +from .projection import project_gas +from .slice import slice_gas +from .volume_render import render_gas +from .ray_trace import panel_gas from .smoothing_length import generate_smoothing_lengths diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py index 00b6c7d9..1836b2fd 100644 --- a/swiftsimio/visualisation/_vistools.py +++ b/swiftsimio/visualisation/_vistools.py @@ -1,31 +1,33 @@ import numpy as np +from warnings import warn +from swiftsimio.objects import cosmo_array +from swiftsimio._array_functions import _copy_cosmo_array_attributes def _get_projection_field(data, field_name): - if field_name is None: - m = np.ones_like(data.particle_ids) - else: - m = getattr(data, field_name) - if data.coordinates.comoving: - if not m.compatible_with_comoving(): - raise AttributeError( - f'Physical quantity "{field_name}" is not compatible with comoving ' - "coordinates!" - ) - else: - if not m.compatible_with_physical(): - raise AttributeError( - f'Comoving quantity "{field_name}" is not compatible with physical ' - "coordinates!" - ) - return m + return ( + getattr(data, field_name) + if field_name is not None + else np.ones_like(data.particle_ids) + ) def _get_region_info(data, region, z_slice=None, require_cubic=False, periodic=True): + boxsize = data.metadata.boxsize + if region is not None: + region = cosmo_array(region) + if data.coordinates.comoving: + boxsize.convert_to_comoving() + if region is not None: + region.convert_to_comoving() + elif data.coordinates.comoving is False: # compare to False in case None + boxsize.convert_to_physical() + if region is not None: + region.convert_to_physical() z_slice_included = z_slice is not None if not z_slice_included: - z_slice = np.zeros_like(data.metadata.boxsize[0]) - box_x, box_y, box_z = data.metadata.boxsize + z_slice = np.zeros_like(boxsize[0]) + box_x, box_y, box_z = boxsize if region is not None: x_min, x_max, y_min, y_max = region[:4] if len(region) == 6: @@ -54,7 +56,7 @@ def _get_region_info(data, region, z_slice=None, require_cubic=False, periodic=T ) periodic_box_x, periodic_box_y, periodic_box_z = ( - data.metadata.boxsize / max_range if periodic else np.zeros(3) + boxsize / max_range if periodic else np.zeros(3) ) return { @@ -77,6 +79,10 @@ def _get_region_info(data, region, z_slice=None, require_cubic=False, periodic=T def _get_rotated_coordinates(data, rotation_matrix, rotation_center): if rotation_center is not None: + if data.coordinates.comoving: + rotation_center = rotation_center.to_comoving() + elif data.coordinates.comoving is False: + rotation_center = rotation_center.to_physical() # Rotate co-ordinates as required x, y, z = np.matmul(rotation_matrix, (data.coordinates - rotation_center).T) @@ -86,3 +92,50 @@ def _get_rotated_coordinates(data, rotation_matrix, rotation_center): else: x, y, z = data.coordinates.T return x, y, z + + +def backends_restore_cosmo_and_units(backend_func): + + def wrapper(*args, **kwargs): + norm = kwargs.pop("norm") + comoving = getattr(kwargs["m"], "comoving", None) + if comoving is True: + if kwargs["x"].comoving is False or kwargs["y"].comoving is False: + warn( + "Projecting a comoving quantity with physical input for coordinates. " + "Converting coordinate grid to comoving." + ) + kwargs["x"].convert_to_comoving() + kwargs["y"].convert_to_comoving() + if kwargs["h"].comoving is False: + warn( + "Projecting a comoving quantity with physical input for smoothing " + "lengths. Converting smoothing lengths to comoving." + ) + kwargs["h"].convert_to_comoving() + norm.convert_to_comoving() + elif comoving is False: # don't use else in case None + if kwargs["x"].comoving or kwargs["y"].comoving: + warn( + "Projecting a physical quantity with comoving input for coordinates. " + "Converting coordinate grid to physical." + ) + kwargs["x"].convert_to_physical() + kwargs["y"].convert_to_physical() + if kwargs["h"].comoving: + warn( + "Projecting a physical quantity with comoving input for smoothing " + "lengths. Converting smoothing lengths to physical." + ) + kwargs["h"].convert_to_physical() + norm.convert_to_physical() + return ( + _copy_cosmo_array_attributes( + kwargs["m"], + backend_func(*args, **kwargs).view(cosmo_array), + copy_units=True, + ) + / norm + ) + + return wrapper diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 99ccf828..707a6dd3 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -7,7 +7,6 @@ from swiftsimio import SWIFTDataset, cosmo_array from swiftsimio.reader import __SWIFTGroupDataset - from swiftsimio.visualisation.projection_backends import backends, backends_parallel from swiftsimio.visualisation.smoothing_length import backends_get_hsml from swiftsimio.visualisation._vistools import ( @@ -16,9 +15,6 @@ _get_rotated_coordinates, ) -scatter = backends["fast"] -scatter_parallel = backends_parallel["fast"] - def project_pixel_grid( data: __SWIFTGroupDataset, @@ -116,7 +112,7 @@ def project_pixel_grid( np.logical_and(z <= region_info["z_max"], z >= region_info["z_min"]), ).astype(bool) - arguments = dict( + kwargs = dict( x=(x[mask] - region_info["x_min"]) / region_info["max_range"], y=(y[mask] - region_info["y_min"]) / region_info["max_range"], m=m[mask], @@ -124,13 +120,13 @@ def project_pixel_grid( res=resolution, box_x=region_info["periodic_box_x"], box_y=region_info["periodic_box_y"], + norm=(region_info["x_range"] * region_info["y_range"]), ) image = ( - backends_parallel[backend](**arguments) + backends_parallel[backend](**kwargs) if parallel - else backends[backend](**arguments) + else backends[backend](**kwargs) ) - raise RuntimeError # determine the effective number of pixels for each dimension xres = int( @@ -144,110 +140,6 @@ def project_pixel_grid( return image[:xres, :yres] -def project_gas_pixel_grid( - data: SWIFTDataset, - resolution: int, - project: Union[str, None] = "masses", - region: Union[None, cosmo_array] = None, - mask: Union[None, np.array] = None, - rotation_matrix: Union[None, np.array] = None, - rotation_center: Union[None, cosmo_array] = None, - parallel: bool = False, - backend: str = "fast", - periodic: bool = True, -): - r""" - Creates a 2D projection of a SWIFT dataset, projected by the "project" - variable (e.g. if project is Temperature, we return: \bar{T} = \sum_j T_j - W_{ij}). - - This function is the same as ``project_gas`` but does not include units. - - Default projection variable is mass. If it is None, then we don't - weight with anything, providing a number density image. - - Parameters - ---------- - - data: SWIFTDataset - The SWIFT dataset that you wish to visualise (get this from ``load``) - - resolution: int - The resolution of the image. All images returned are square, ``res`` - by ``res``, pixel grids. - - project: str, optional - Variable to project to get the weighted density of. By default, this - is mass. If you would like to mass-weight any other variable, you can - always create it as ``data.gas.my_variable = data.gas.other_variable - * data.gas.masses``. - - region: cosmo_array, optional - Region, determines where the image will be created (this corresponds - to the left and right-hand edges, and top and bottom edges) if it is - not None. It should have a length of four or six, and take the form: - ``[x_min, x_max, y_min, y_max, {z_min, z_max}]`` - - mask: np.array, optional - Allows only a sub-set of the particles in data to be visualised. Useful - in cases where you have read data out of a ``velociraptor`` catalogue, - or if you only want to visualise e.g. star forming particles. This boolean - mask is applied just before visualisation. - - rotation_center: np.array, optional - Center of the rotation. If you are trying to rotate around a galaxy, this - should be the most bound particle. - - rotation_matrix: np.array, optional - Rotation matrix (3x3) that describes the rotation of the box around - ``rotation_center``. In the default case, this provides a projection - along the z axis. - - parallel: bool, optional - Defaults to ``False``, whether or not to create the image in parallel. - The parallel version of this function uses significantly more memory. - - backend: str, optional - Backend to use. See documentation for details. Defaults to 'fast'. - - periodic: bool, optional - Account for periodic boundary conditions for the simulation box? - Defaults to ``True``. - - Returns - ------- - - image: np.array - Projected image with dimensions of project / length^2, of size - ``res`` x ``res``. - - - Notes - ----- - - + Particles outside of this range are still considered if their smoothing - lengths overlap with the range. - + The returned array has x as the first component and y as the second component, - which is the opposite to what ``imshow`` requires. You should transpose the - array if you want it to be visualised the 'right way up'. - """ - - image = project_pixel_grid( - data=data.gas, - resolution=resolution, - project=project, - mask=mask, - parallel=parallel, - region=region, - rotation_matrix=rotation_matrix, - rotation_center=rotation_center, - backend=backend, - periodic=periodic, - ) - - return image - - def project_gas( data: SWIFTDataset, resolution: int, @@ -335,9 +227,8 @@ def project_gas( array if you want it to be visualised the 'right way up'. """ - # REFACTOR? can we get this image to come out as a cosmo_array already? - image = project_gas_pixel_grid( - data=data, + return project_pixel_grid( + data=data.gas, resolution=resolution, project=project, mask=mask, @@ -348,21 +239,3 @@ def project_gas( backend=backend, periodic=periodic, ) - - region_info = _get_region_info(data.gas, region) - units = 1.0 / (region_info["max_range"] ** 2) - # Unfortunately this is required to prevent us from {over,under}flowing the units: - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) - - comoving = data.gas.coordinates.comoving - coord_cosmo_factor = data.gas.coordinates.cosmo_factor - if project is not None: - units *= getattr(data.gas, project).units - project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 - else: - new_cosmo_factor = coord_cosmo_factor ** (-2) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) diff --git a/swiftsimio/visualisation/projection_backends/__init__.py b/swiftsimio/visualisation/projection_backends/__init__.py index 3715abf7..14a620ef 100644 --- a/swiftsimio/visualisation/projection_backends/__init__.py +++ b/swiftsimio/visualisation/projection_backends/__init__.py @@ -6,6 +6,8 @@ being a developer-only indended feature. """ +from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units + from swiftsimio.visualisation.projection_backends.fast import scatter as fast from swiftsimio.visualisation.projection_backends.fast import ( scatter_parallel as fast_parallel, @@ -48,21 +50,21 @@ ) backends = { - "histogram": histogram, - "fast": fast, - "renormalised": renormalised, - "subsampled": subsampled, - "subsampled_extreme": subsampled_extreme, - "reference": reference, - "gpu": gpu, + "histogram": backends_restore_cosmo_and_units(histogram), + "fast": backends_restore_cosmo_and_units(fast), + "renormalised": backends_restore_cosmo_and_units(renormalised), + "subsampled": backends_restore_cosmo_and_units(subsampled), + "subsampled_extreme": backends_restore_cosmo_and_units(subsampled_extreme), + "reference": backends_restore_cosmo_and_units(reference), + "gpu": backends_restore_cosmo_and_units(gpu), } backends_parallel = { - "histogram": histogram_parallel, - "fast": fast_parallel, - "renormalised": renormalised_parallel, - "subsampled": subsampled_parallel, - "subsampled_extreme": subsampled_extreme_parallel, - "reference": reference_parallel, - "gpu": gpu_parallel, + "histogram": backends_restore_cosmo_and_units(histogram_parallel), + "fast": backends_restore_cosmo_and_units(fast_parallel), + "renormalised": backends_restore_cosmo_and_units(renormalised_parallel), + "subsampled": backends_restore_cosmo_and_units(subsampled_parallel), + "subsampled_extreme": backends_restore_cosmo_and_units(subsampled_extreme_parallel), + "reference": backends_restore_cosmo_and_units(reference_parallel), + "gpu": backends_restore_cosmo_and_units(gpu_parallel), } diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index 480fb516..cb101543 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -2,24 +2,13 @@ Ray tracing module for visualisation. """ -# There should be three implementations here: -# - An example case that uses 'screening' with multiple panes. -# - An example case that builds a 3D mesh and uses that -# - An example case that uses a 'real' algorithm (even if it has to -# be single-threaded) - from typing import Union import numpy as np -import math from swiftsimio.objects import cosmo_array from swiftsimio.reader import __SWIFTGroupDataset, SWIFTDataset -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_gamma, - kernel_double_precision as kernel, -) -from swiftsimio.accelerated import jit, prange, NUM_THREADS +from swiftsimio.visualisation.ray_trace_backends import backends from swiftsimio.visualisation.smoothing_length import backends_get_hsml from swiftsimio.visualisation._vistools import ( _get_projection_field, @@ -28,198 +17,6 @@ ) -@jit(nopython=True, fastmath=True) -def core_panels( - x: np.float64, - y: np.float64, - z: np.float64, - h: np.float32, - m: np.float32, - res: int, - panels: int, - min_z: np.float64, - max_z: np.float64, -) -> np.array: - """ - Creates a 2D array of the projected density of particles in a 3D volume using the - 'renormalised' strategy, with multiple panels across the z-range. - - Parameters - ---------- - - x: np.array[np.float64] - The x-coordinates of the particles. - - y: np.array[np.float64] - The y-coordinates of the particles. - - z: np.array[np.float64] - The z-coordinates of the particles. - - h: np.array[np.float32] - The smoothing lengths of the particles. - - m: np.array[np.float32] - The masses of the particles. - - res: int - The resolution of the output array. - - panels: int - The number of panels to use in the z-direction. - - min_z: np.float64 - The minimum z-coordinate of the volume. - - max_z: np.float64 - The maximum z-coordinate of the volume. - - Returns - ------- - - A 3D array of shape (res, res, panels) containing the projected density in each pixel. - """ - output = np.zeros((res, res, panels)) - maximal_array_index = res - 1 - - number_of_particles = len(x) - float_res = float(res) - pixel_width = 1.0 / float_res - - assert len(y) == number_of_particles - assert len(z) == number_of_particles - assert len(h) == number_of_particles - assert len(m) == number_of_particles - - z_per_panel = (max_z - min_z) / panels - - inverse_cell_area = float_res * float_res - - for i in range(number_of_particles): - panel = int(z[i] / z_per_panel) - - if panel < 0 or panel >= panels: - continue - - particle_cell_x = int(float_res * x[i]) - particle_cell_y = int(float_res * y[i]) - - kernel_width = kernel_gamma * h[i] - cells_spanned = int(1.0 + kernel_width * float_res) - - if ( - particle_cell_x + cells_spanned < 0 - or particle_cell_x - cells_spanned > maximal_array_index - or particle_cell_y + cells_spanned < 0 - or particle_cell_y - cells_spanned > maximal_array_index - ): - # Can happily skip this particle - continue - - if cells_spanned <= 1: - if ( - particle_cell_x >= 0 - and particle_cell_x <= maximal_array_index - and particle_cell_y >= 0 - and particle_cell_y <= maximal_array_index - ): - output[particle_cell_x, particle_cell_y, panel] += ( - m[i] * inverse_cell_area - ) - continue - - normalisation = 0.0 - - for cell_x in range( - particle_cell_x - cells_spanned, particle_cell_x + cells_spanned + 1 - ): - distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] - distance_x_2 = distance_x * distance_x - - for cell_y in range( - particle_cell_y - cells_spanned, particle_cell_y + cells_spanned + 1 - ): - distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] - distance_y_2 = distance_y * distance_y - - r = math.sqrt(distance_x_2 + distance_y_2) - - normalisation += kernel(r, kernel_width) - - # Now have the normalisation - normalisation = m[i] * inverse_cell_area / normalisation - - for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we'll segfault - max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the array bounds, - # otherwise we'll segfault (oops). - min(particle_cell_x + cells_spanned + 1, maximal_array_index + 1), - ): - distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] - distance_x_2 = distance_x * distance_x - - for cell_y in range( - max(0, particle_cell_y - cells_spanned), - min(particle_cell_y + cells_spanned + 1, maximal_array_index + 1), - ): - distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] - distance_y_2 = distance_y * distance_y - - r = math.sqrt(distance_x_2 + distance_y_2) - - output[cell_x, cell_y, panel] += kernel(r, kernel_width) * normalisation - - return output - - -@jit(nopython=True, fastmath=True) -def core_panels_parallel( - x: np.float64, - y: np.float64, - z: np.float64, - h: np.float32, - m: np.float32, - res: int, - panels: int, - min_z: np.float64, - max_z: np.float64, -): - # Same as scatter, but executes in parallel! This is actually trivial, - # we just make NUM_THREADS images and add them together at the end. - - number_of_particles = x.size - core_particles = number_of_particles // NUM_THREADS - - output = np.zeros((res, res, panels), dtype=np.float32) - - for thread in prange(NUM_THREADS): - # Left edge is easy, just start at 0 and go to 'final' - left_edge = thread * core_particles - - # Right edge is harder in case of left over particles... - right_edge = thread + 1 - - if right_edge == NUM_THREADS: - right_edge = number_of_particles - else: - right_edge *= core_particles - - output += core_panels( - x[left_edge:right_edge], - y[left_edge:right_edge], - z[left_edge:right_edge], - h[left_edge:right_edge], - m[left_edge:right_edge], - res, - panels, - min_z, - max_z, - ) - - return output - - def panel_pixel_grid( data: __SWIFTGroupDataset, resolution: int, @@ -237,7 +34,9 @@ def panel_pixel_grid( x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) mask = np.s_[...] if mask is None else mask - return core_panels( + # There's a parallel version of core_panels but it seems + # that it's never used anywhere. + return backends["core_panels"]( x=x[mask] / region_info["max_range"], y=y[mask] / region_info["max_range"], z=z[mask], @@ -247,6 +46,7 @@ def panel_pixel_grid( panels=panels, min_z=region_info["z_min"], max_z=region_info["z_max"], + norm=(region_info["x_range"] * region_info["y_range"]), ) @@ -260,7 +60,7 @@ def panel_gas( rotation_matrix: Union[None, np.array] = None, rotation_center: Union[None, cosmo_array] = None, ) -> cosmo_array: - image = panel_pixel_grid( + return panel_pixel_grid( data=data.gas, resolution=resolution, panels=panels, @@ -270,152 +70,3 @@ def panel_gas( rotation_matrix=rotation_matrix, rotation_center=rotation_center, ) - - if region is not None: - x_range = region[1] - region[0] - y_range = region[3] - region[2] - max_range = max(x_range, y_range) - units = 1.0 / (max_range**2) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units(1.0 / (x_range.units * y_range.units)) - else: - max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**2) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**2) - - comoving = data.gas.coordinates.comoving - coord_cosmo_factor = data.gas.coordinates.cosmo_factor - if project is not None: - units *= getattr(data.gas, project).units - project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**2 - else: - new_cosmo_factor = coord_cosmo_factor ** (-2) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) - - -# --- Functions that actually perform the 'ray tracing'. - - -def transfer_function(value, width, center): - """ - A simple gaussian transfer function centered around a specific value. - """ - return ( - 1 - / (width * np.sqrt(2.0 * np.pi)) - * np.exp(-0.5 * ((value - center) / width) ** 2) - ) - - -@jit(fastmath=True, nopython=True) -def integrate_ray_numba_specific( - input: np.array, red: float, green: float, blue: float, center: float, width: float -): - """ - Given a ray, integrate the transfer function along it - """ - - value = np.array([0.0, 0.0, 0.0], dtype=np.float32) - color = np.array([red, green, blue], dtype=np.float32) - - for i in input: - value += ( - color - * 1 - / (width * np.sqrt(2.0 * np.pi)) - * np.exp(-0.5 * ((i - center) / width) ** 2) - ) - - return value / len(input) - - -@jit(fastmath=True, nopython=True) -def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): - """ - Given a ray, integrate the transfer function along it - """ - - value = np.float32(0.0) - - for i in input: - value *= 0.99 - value += ( - 1 - / (width * np.sqrt(2.0 * np.pi)) - * np.exp(-0.5 * ((i - center) / width) ** 2) - ) - - return np.float32(value / len(input)) - - -# #%% -# data = np.load("voxel_1024.npy") -# # %% -# log_data = np.log10(data) -# # %% -# transfer = lambda x: transfer_function(x, np.mean(log_data), np.std(log_data) * 0.5) -# # %% -# color = np.array([1.0, 0.0, 0.0], dtype=np.float32) -# #%% -# from tqdm import tqdm -# # %% -# @numba.njit(fastmath=True) -# def make_grid(color, center, width): -# output = np.zeros((len(log_data), len(log_data[0])), dtype=np.float32) -# for x in numba.prange(len(log_data)): -# for y in range(len(log_data)): -# data = log_data[x, y] - -# value = np.float32(0.0) - -# for index, i in enumerate(data): -# factor = index / len(data) - -# if factor > 0.5: -# factor = 1.0 - factor - -# value += ( -# 1 -# / (width * np.sqrt(2.0 * np.pi)) -# * np.exp(-0.5 * ((i - center) / width) ** 2) -# ) - -# output[x, y] = value - -# return output - - -# # %% -# import matplotlib.pyplot as plt -# import swiftascmaps -# # %% -# std = np.std(log_data) -# width = 0.05 -# centers = [np.mean(log_data) + x * std for x in [0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]] -# # %% -# colors = plt.get_cmap("swift.nineteen_eighty_nine")((np.linspace(0, 1, len(centers))))[ -# :, :3 -# ] - -# grids = [ -# make_grid(color, center, width) for color, center in zip(colors, centers) -# ] -# #%% - -# #%% -# make_image = lambda x, y: np.array([x * y[0], x * y[1], x * y[2]]).T -# images = [make_image(grid / np.max(grid), color) for color, grid in zip(colors, grids)] -# # %% -# combined_image = sum(images) -# plt.imsave("test.png", combined_image / np.max(combined_image)) -# # %% -# for id, image in zip(centers, images): -# plt.imsave(f"test{id}.png", image) -# # %% diff --git a/swiftsimio/visualisation/ray_trace_backends/__init__.py b/swiftsimio/visualisation/ray_trace_backends/__init__.py new file mode 100644 index 00000000..9c25b3db --- /dev/null +++ b/swiftsimio/visualisation/ray_trace_backends/__init__.py @@ -0,0 +1,18 @@ +""" +Backends for ray tracing +""" + +from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units + +from swiftsimio.visualisation.ray_trace_backends.core_panels import ( + core_panels, + core_panels_parallel, +) + +backends = { + "core_panels": backends_restore_cosmo_and_units(core_panels), +} + +backends_parallel = { + "core_panels": backends_restore_cosmo_and_units(core_panels_parallel), +} diff --git a/swiftsimio/visualisation/ray_trace_backends/core_panels.py b/swiftsimio/visualisation/ray_trace_backends/core_panels.py new file mode 100644 index 00000000..9919aa18 --- /dev/null +++ b/swiftsimio/visualisation/ray_trace_backends/core_panels.py @@ -0,0 +1,332 @@ +""" +Ray tracing module for visualisation. +""" + +# There should be three implementations here: +# - An example case that uses 'screening' with multiple panes. +# - An example case that builds a 3D mesh and uses that +# - An example case that uses a 'real' algorithm (even if it has to +# be single-threaded) + +import numpy as np +import math + +from swiftsimio.visualisation.projection_backends.kernels import ( + kernel_gamma, + kernel_double_precision as kernel, +) + +from swiftsimio.accelerated import jit, prange, NUM_THREADS + + +@jit(nopython=True, fastmath=True) +def core_panels( + x: np.float64, + y: np.float64, + z: np.float64, + h: np.float32, + m: np.float32, + res: int, + panels: int, + min_z: np.float64, + max_z: np.float64, +) -> np.array: + """ + Creates a 2D array of the projected density of particles in a 3D volume using the + 'renormalised' strategy, with multiple panels across the z-range. + + Parameters + ---------- + + x: np.array[np.float64] + The x-coordinates of the particles. + + y: np.array[np.float64] + The y-coordinates of the particles. + + z: np.array[np.float64] + The z-coordinates of the particles. + + h: np.array[np.float32] + The smoothing lengths of the particles. + + m: np.array[np.float32] + The masses of the particles. + + res: int + The resolution of the output array. + + panels: int + The number of panels to use in the z-direction. + + min_z: np.float64 + The minimum z-coordinate of the volume. + + max_z: np.float64 + The maximum z-coordinate of the volume. + + Returns + ------- + + A 3D array of shape (res, res, panels) containing the projected density in each pixel. + """ + output = np.zeros((res, res, panels)) + maximal_array_index = res - 1 + + number_of_particles = len(x) + float_res = float(res) + pixel_width = 1.0 / float_res + + assert len(y) == number_of_particles + assert len(z) == number_of_particles + assert len(h) == number_of_particles + assert len(m) == number_of_particles + + z_per_panel = (max_z - min_z) / panels + + inverse_cell_area = float_res * float_res + + for i in range(number_of_particles): + panel = int(z[i] / z_per_panel) + + if panel < 0 or panel >= panels: + continue + + particle_cell_x = int(float_res * x[i]) + particle_cell_y = int(float_res * y[i]) + + kernel_width = kernel_gamma * h[i] + cells_spanned = int(1.0 + kernel_width * float_res) + + if ( + particle_cell_x + cells_spanned < 0 + or particle_cell_x - cells_spanned > maximal_array_index + or particle_cell_y + cells_spanned < 0 + or particle_cell_y - cells_spanned > maximal_array_index + ): + # Can happily skip this particle + continue + + if cells_spanned <= 1: + if ( + particle_cell_x >= 0 + and particle_cell_x <= maximal_array_index + and particle_cell_y >= 0 + and particle_cell_y <= maximal_array_index + ): + output[particle_cell_x, particle_cell_y, panel] += ( + m[i] * inverse_cell_area + ) + continue + + normalisation = 0.0 + + for cell_x in range( + particle_cell_x - cells_spanned, particle_cell_x + cells_spanned + 1 + ): + distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] + distance_x_2 = distance_x * distance_x + + for cell_y in range( + particle_cell_y - cells_spanned, particle_cell_y + cells_spanned + 1 + ): + distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] + distance_y_2 = distance_y * distance_y + + r = math.sqrt(distance_x_2 + distance_y_2) + + normalisation += kernel(r, kernel_width) + + # Now have the normalisation + normalisation = m[i] * inverse_cell_area / normalisation + + for cell_x in range( + # Ensure that the lowest x value is 0, otherwise we'll segfault + max(0, particle_cell_x - cells_spanned), + # Ensure that the highest x value lies within the array bounds, + # otherwise we'll segfault (oops). + min(particle_cell_x + cells_spanned + 1, maximal_array_index + 1), + ): + distance_x = (float(cell_x) + 0.5) * pixel_width - x[i] + distance_x_2 = distance_x * distance_x + + for cell_y in range( + max(0, particle_cell_y - cells_spanned), + min(particle_cell_y + cells_spanned + 1, maximal_array_index + 1), + ): + distance_y = (float(cell_y) + 0.5) * pixel_width - y[i] + distance_y_2 = distance_y * distance_y + + r = math.sqrt(distance_x_2 + distance_y_2) + + output[cell_x, cell_y, panel] += kernel(r, kernel_width) * normalisation + + return output + + +@jit(nopython=True, fastmath=True) +def core_panels_parallel( + x: np.float64, + y: np.float64, + z: np.float64, + h: np.float32, + m: np.float32, + res: int, + panels: int, + min_z: np.float64, + max_z: np.float64, +): + # Same as scatter, but executes in parallel! This is actually trivial, + # we just make NUM_THREADS images and add them together at the end. + + number_of_particles = x.size + core_particles = number_of_particles // NUM_THREADS + + output = np.zeros((res, res, panels), dtype=np.float32) + + for thread in prange(NUM_THREADS): + # Left edge is easy, just start at 0 and go to 'final' + left_edge = thread * core_particles + + # Right edge is harder in case of left over particles... + right_edge = thread + 1 + + if right_edge == NUM_THREADS: + right_edge = number_of_particles + else: + right_edge *= core_particles + + output += core_panels( + x[left_edge:right_edge], + y[left_edge:right_edge], + z[left_edge:right_edge], + h[left_edge:right_edge], + m[left_edge:right_edge], + res, + panels, + min_z, + max_z, + ) + + return output + + +# --- Functions that actually perform the 'ray tracing'. + + +def transfer_function(value, width, center): + """ + A simple gaussian transfer function centered around a specific value. + """ + return ( + 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((value - center) / width) ** 2) + ) + + +@jit(fastmath=True, nopython=True) +def integrate_ray_numba_specific( + input: np.array, red: float, green: float, blue: float, center: float, width: float +): + """ + Given a ray, integrate the transfer function along it + """ + + value = np.array([0.0, 0.0, 0.0], dtype=np.float32) + color = np.array([red, green, blue], dtype=np.float32) + + for i in input: + value += ( + color + * 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((i - center) / width) ** 2) + ) + + return value / len(input) + + +@jit(fastmath=True, nopython=True) +def integrate_ray_numba_nocolor(input: np.array, center: float, width: float): + """ + Given a ray, integrate the transfer function along it + """ + + value = np.float32(0.0) + + for i in input: + value *= 0.99 + value += ( + 1 + / (width * np.sqrt(2.0 * np.pi)) + * np.exp(-0.5 * ((i - center) / width) ** 2) + ) + + return np.float32(value / len(input)) + + +# #%% +# data = np.load("voxel_1024.npy") +# # %% +# log_data = np.log10(data) +# # %% +# transfer = lambda x: transfer_function(x, np.mean(log_data), np.std(log_data) * 0.5) +# # %% +# color = np.array([1.0, 0.0, 0.0], dtype=np.float32) +# #%% +# from tqdm import tqdm +# # %% +# @numba.njit(fastmath=True) +# def make_grid(color, center, width): +# output = np.zeros((len(log_data), len(log_data[0])), dtype=np.float32) +# for x in numba.prange(len(log_data)): +# for y in range(len(log_data)): +# data = log_data[x, y] + +# value = np.float32(0.0) + +# for index, i in enumerate(data): +# factor = index / len(data) + +# if factor > 0.5: +# factor = 1.0 - factor + +# value += ( +# 1 +# / (width * np.sqrt(2.0 * np.pi)) +# * np.exp(-0.5 * ((i - center) / width) ** 2) +# ) + +# output[x, y] = value + +# return output + + +# # %% +# import matplotlib.pyplot as plt +# import swiftascmaps +# # %% +# std = np.std(log_data) +# width = 0.05 +# centers = [np.mean(log_data) + x * std for x in [0, 1.5, 2.0, 2.5, 3.0, 3.5, 4.0, 4.5]] +# # %% +# colors = plt.get_cmap("swift.nineteen_eighty_nine")((np.linspace(0, 1, len(centers))))[ +# :, :3 +# ] + +# grids = [ +# make_grid(color, center, width) for color, center in zip(colors, centers) +# ] +# #%% + +# #%% +# make_image = lambda x, y: np.array([x * y[0], x * y[1], x * y[2]]).T +# images = [make_image(grid / np.max(grid), color) for color, grid in zip(colors, grids)] +# # %% +# combined_image = sum(images) +# plt.imsave("test.png", combined_image / np.max(combined_image)) +# # %% +# for id, image in zip(centers, images): +# plt.imsave(f"test{id}.png", image) +# # %% diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index 7820b436..fe57b260 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -13,11 +13,8 @@ _get_rotated_coordinates, ) -slice_scatter = backends["sph"] -slice_scatter_parallel = backends_parallel["sph"] - -def slice_gas_pixel_grid( +def slice_gas( data: SWIFTDataset, resolution: int, z_slice: Optional[cosmo_quantity] = None, @@ -95,6 +92,7 @@ def slice_gas_pixel_grid( """ data = data.gas + z_slice = np.zeros_like(data.metadata.boxsize[0]) if z_slice is None else z_slice m = _get_projection_field(data, project) region_info = _get_region_info(data, region, z_slice=z_slice, periodic=periodic) @@ -110,7 +108,7 @@ def slice_gas_pixel_grid( xres = int(resolution * region_info["x_range"] / region_info["max_range"]) yres = int(resolution * region_info["y_range"] / region_info["max_range"]) - common_parameters = dict( + kwargs = dict( x=(x - region_info["x_min"]) / region_info["max_range"], y=(y - region_info["y_min"]) / region_info["max_range"], z=z / region_info["max_range"], @@ -122,139 +120,12 @@ def slice_gas_pixel_grid( box_x=region_info["periodic_box_x"], box_y=region_info["periodic_box_y"], box_z=region_info["periodic_box_z"], + norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]), ) if parallel: - image = backends_parallel[backend](**common_parameters) + image = backends_parallel[backend](**kwargs) else: - image = backends[backend](**common_parameters) + image = backends[backend](**kwargs) return image - - -def slice_gas( - data: SWIFTDataset, - resolution: int, - z_slice: Optional[cosmo_quantity] = None, - project: Union[str, None] = "masses", - parallel: bool = False, - rotation_matrix: Union[None, np.array] = None, - rotation_center: Union[None, cosmo_array] = None, - region: Union[None, cosmo_array] = None, - backend: str = "sph", - periodic: bool = True, -): - """ - Creates a 2D slice of a SWIFT dataset, weighted by data field - - Parameters - ---------- - data : SWIFTDataset - Dataset from which slice is extracted - - resolution : int - Specifies size of return np.array - - z_slice : cosmo_quantity - Specifies the location along the z-axis where the slice is to be - extracted, relative to the rotation center or the origin of the box - if no rotation center is provided. If the perspective is rotated - this value refers to the location along the rotated z-axis. - - project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles - - parallel : bool, optional - used to determine if we will create the image in parallel. This - defaults to False, but can speed up the creation of large images - significantly at the cost of increased memory usage. - - rotation_matrix: np.np.array, optional - Rotation matrix (3x3) that describes the rotation of the box around - ``rotation_center``. In the default case, this provides a slice - perpendicular to the z axis. - - rotation_center: np.np.array, optional - Center of the rotation. If you are trying to rotate around a galaxy, this - should be the most bound particle. - - region : np.array, optional - determines where the image will be created - (this corresponds to the left and right-hand edges, and top and bottom edges) - if it is not None. It should have a length of four, and take the form: - - [x_min, x_max, y_min, y_max] - - Particles outside of this range are still considered if their - smoothing lengths overlap with the range. - - backend : str, optional - Backend to use. Choices are "sph" for interpolation using kernel weights or - "nearest_neighbours" for nearest neighbour interpolation. - - periodic : bool, optional - Account for periodic boundaries for the simulation box? - Default is ``True``. - - Returns - ------- - ndarray of np.float32 - a `resolution` x `resolution` np.array of the contribution - of the projected data field to the voxel grid from all of the particles - - See Also - -------- - slice_gas_pixel grid : Creates a 2D slice of a SWIFT dataset - render_gas : Creates a 3D voxel grid of a SWIFT dataset with appropriate units - - Notes - ----- - This is a wrapper function for slice_gas_pixel_grid ensuring that output units are - appropriate - """ - - z_slice = np.zeros_like(data.metadata.boxsize[0]) if z_slice is None else z_slice - - image = slice_gas_pixel_grid( - data, - resolution, - z_slice, - project, - parallel, - rotation_matrix, - rotation_center, - region, - backend, - periodic, - ) - - if region is not None: - x_range = region[1] - region[0] - y_range = region[3] - region[2] - max_range = max(x_range, y_range) - units = 1.0 / (max_range**3) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units( - 1.0 / (x_range.units * y_range.units * data.metadata.boxsize.units) - ) - else: - max_range = max(data.metadata.boxsize[0], data.metadata.boxsize[1]) - units = 1.0 / (max_range**3) - # Unfortunately this is required to prevent us from {over,under}flowing - # the units... - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) - - comoving = data.gas.coordinates.comoving - coord_cosmo_factor = data.gas.coordinates.cosmo_factor - if project is not None: - units *= getattr(data.gas, project).units - project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 - else: - new_cosmo_factor = coord_cosmo_factor ** (-3) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) diff --git a/swiftsimio/visualisation/slice_backends/__init__.py b/swiftsimio/visualisation/slice_backends/__init__.py index 0fd2e2ba..099d8ada 100644 --- a/swiftsimio/visualisation/slice_backends/__init__.py +++ b/swiftsimio/visualisation/slice_backends/__init__.py @@ -2,6 +2,8 @@ Backends for density slicing. """ +from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units + from swiftsimio.visualisation.slice_backends.sph import ( slice_scatter as sph, slice_scatter_parallel as sph_parallel, @@ -11,9 +13,12 @@ slice_scatter_parallel as nearest_neighbours_parallel, ) -backends = {"sph": sph, "nearest_neighbours": nearest_neighbours} +backends = { + "sph": backends_restore_cosmo_and_units(sph), + "nearest_neighbours": backends_restore_cosmo_and_units(nearest_neighbours), +} backends_parallel = { - "sph": sph_parallel, - "nearest_neighbours": nearest_neighbours_parallel, + "sph": backends_restore_cosmo_and_units(sph_parallel), + "nearest_neighbours": backends_restore_cosmo_and_units(nearest_neighbours_parallel), } diff --git a/swiftsimio/visualisation/smoothing_length/sph.py b/swiftsimio/visualisation/smoothing_length/sph.py index a38f1035..d761ad46 100644 --- a/swiftsimio/visualisation/smoothing_length/sph.py +++ b/swiftsimio/visualisation/smoothing_length/sph.py @@ -14,19 +14,8 @@ def get_hsml(data: SWIFTDataset) -> cosmo_array: ------- The extracted smoothing lengths. """ - hsml = ( + return ( data.smoothing_lengths if hasattr(data, "smoothing_lengths") else data.smoothing_length # backwards compatibility ) - if data.coordinates.comoving: - if not hsml.compatible_with_comoving(): - raise AttributeError( - "Physical smoothing length is not compatible with comoving coordinates!" - ) - else: - if not hsml.compatible_with_physical(): - raise AttributeError( - "Comoving smoothing length is not compatible with physical coordinates!" - ) - return hsml diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 11f26878..b0dda8dc 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -7,12 +7,12 @@ from math import sqrt, exp, pi import numpy as np from swiftsimio import SWIFTDataset, cosmo_array +from swiftsimio.accelerated import jit -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.optional_packages import plt -from swiftsimio.visualisation.slice_backends.sph import kernel, kernel_gamma from swiftsimio.visualisation.smoothing_length import backends_get_hsml +from swiftsimio.visualisation.volume_render_backends import backends, backends_parallel from swiftsimio.visualisation._vistools import ( _get_projection_field, _get_region_info, @@ -20,564 +20,7 @@ ) -@jit(nopython=True, fastmath=True) -def scatter( - x: np.float64, - y: np.float64, - z: np.float64, - m: np.float32, - h: np.float32, - res: int, - box_x: np.float64 = 0.0, - box_y: np.float64 = 0.0, - box_z: np.float64 = 0.0, -) -> np.ndarray: - """ - Creates a weighted voxel grid - - Computes contributions to a voxel grid from particles with positions - (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. - This includes periodic boundary effects. - - Parameters - ---------- - - x : np.np.array[np.float64] - np.array of x-positions of the particles. Must be bounded by [0, 1]. - - y : np.np.array[np.float64] - np.array of y-positions of the particles. Must be bounded by [0, 1]. - - z : np.np.array[np.float64] - np.array of z-positions of the particles. Must be bounded by [0, 1]. - - m : np.np.array[np.float32] - np.array of masses (or otherwise weights) of the particles - - h : np.np.array[np.float32] - np.array of smoothing lengths of the particles - - res : int - the number of voxels along one axis, i.e. this returns a cube - of res * res * res. - - box_x: np.float64 - box size in x, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_y: np.float64 - box size in y, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_z: np.float64 - box size in z, in the same rescaled length units as x, y and z. - Used for periodic wrapping - - Returns - ------- - - np.np.array[np.float32, np.float32, np.float32] - voxel grid of quantity - - See Also - -------- - - scatter_parallel : Parallel implementation of this function - slice_scatter : Create scatter plot of a slice of data - slice_scatter_parallel : Create scatter plot of a slice of data in parallel - - Notes - ----- - - Explicitly defining the types in this function allows - for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba np.ones. - """ - # Output np.array for our image - image = np.zeros((res, res, res), dtype=np.float32) - maximal_array_index = np.int32(res) - 1 - - # Change that integer to a float, we know that our x, y are bounded - # by [0, 1]. - float_res = np.float32(res) - pixel_width = 1.0 / float_res - - # We need this for combining with the x_pos and y_pos variables. - float_res_64 = np.float64(res) - - # If the kernel width is smaller than this, we drop to just PIC method - drop_to_single_cell = pixel_width * 0.5 - - # Pre-calculate this constant for use with the above - inverse_cell_volume = float_res * float_res * float_res - - if box_x == 0.0: - xshift_min = 0 - xshift_max = 1 - else: - xshift_min = -1 - xshift_max = 2 - if box_y == 0.0: - yshift_min = 0 - yshift_max = 1 - else: - yshift_min = -1 - yshift_max = 2 - if box_z == 0.0: - zshift_min = 0 - zshift_max = 1 - else: - zshift_min = -1 - zshift_max = 2 - - for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( - x, y, z, m, h - ): - # loop over periodic copies of the particle - for xshift in range(xshift_min, xshift_max): - for yshift in range(yshift_min, yshift_max): - for zshift in range(zshift_min, zshift_max): - x_pos = x_pos_original + xshift * box_x - y_pos = y_pos_original + yshift * box_y - z_pos = z_pos_original + zshift * box_z - - # Calculate the cell that this particle; use the 64 bit version of the - # resolution as this is the same type as the positions - particle_cell_x = np.int32(float_res_64 * x_pos) - particle_cell_y = np.int32(float_res_64 * y_pos) - particle_cell_z = np.int32(float_res_64 * z_pos) - - # SWIFT stores hsml as the FWHM. - kernel_width = kernel_gamma * hsml - - # The number of cells that this kernel spans - cells_spanned = np.int32(1.0 + kernel_width * float_res) - - if ( - particle_cell_x + cells_spanned < 0 - or particle_cell_x - cells_spanned > maximal_array_index - or particle_cell_y + cells_spanned < 0 - or particle_cell_y - cells_spanned > maximal_array_index - or particle_cell_z + cells_spanned < 0 - or particle_cell_z - cells_spanned > maximal_array_index - ): - # Can happily skip this particle - continue - - if kernel_width < drop_to_single_cell: - # Easygame, gg - if ( - particle_cell_x >= 0 - and particle_cell_x <= maximal_array_index - and particle_cell_y >= 0 - and particle_cell_y <= maximal_array_index - and particle_cell_z >= 0 - and particle_cell_z <= maximal_array_index - ): - image[ - particle_cell_x, particle_cell_y, particle_cell_z - ] += (mass * inverse_cell_volume) - else: - # Now we loop over the square of cells that the kernel lives in - for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we segfault - max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the np.array - # bounds, otherwise we'll segfault (oops). - min( - particle_cell_x + cells_spanned, maximal_array_index + 1 - ), - ): - # The distance in x to our new favourite cell -- remember that - # our x, y are all in a box of [0, 1]; calculate the distance - # to the cell centre - distance_x = ( - np.float32(cell_x) + 0.5 - ) * pixel_width - np.float32(x_pos) - distance_x_2 = distance_x * distance_x - for cell_y in range( - max(0, particle_cell_y - cells_spanned), - min( - particle_cell_y + cells_spanned, - maximal_array_index + 1, - ), - ): - distance_y = ( - np.float32(cell_y) + 0.5 - ) * pixel_width - np.float32(y_pos) - distance_y_2 = distance_y * distance_y - for cell_z in range( - max(0, particle_cell_z - cells_spanned), - min( - particle_cell_z + cells_spanned, - maximal_array_index + 1, - ), - ): - distance_z = ( - np.float32(cell_z) + 0.5 - ) * pixel_width - np.float32(z_pos) - distance_z_2 = distance_z * distance_z - - r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) - - kernel_eval = kernel(r, kernel_width) - - image[cell_x, cell_y, cell_z] += mass * kernel_eval - - return image - - -@jit(nopython=True, fastmath=True) -def scatter_limited_z( - x: np.float64, - y: np.float64, - z: np.float64, - m: np.float32, - h: np.float32, - res: int, - res_ratio_z: int, - box_x: np.float64 = 0.0, - box_y: np.float64 = 0.0, - box_z: np.float64 = 0.0, -) -> np.ndarray: - """ - Creates a weighted voxel grid - - Computes contributions to a voxel grid from particles with positions - (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. - This includes periodic boundary effects. - - Parameters - ---------- - - x : np.np.array[np.float64] - np.array of x-positions of the particles. Must be bounded by [0, 1]. - - y : np.np.array[np.float64] - np.array of y-positions of the particles. Must be bounded by [0, 1]. - - z : np.np.array[np.float64] - np.array of z-positions of the particles. Must be bounded by [0, 1]. - - m : np.np.array[np.float32] - np.array of masses (or otherwise weights) of the particles - - h : np.np.array[np.float32] - np.array of smoothing lengths of the particles - - res : int - the number of voxels along one axis, i.e. this returns a cube - of res * res * res. - - res_ratio_z: int - the number of voxels along the x and y axes relative to the z - axis. If this is, for instance, 8, and the res is 128, then the - output np.array will be 128 x 128 x 16. - - box_x: np.float64 - box size in x, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_y: np.float64 - box size in y, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_z: np.float64 - box size in z, in the same rescaled length units as x, y and z. - Used for periodic wrapping - - Returns - ------- - - np.np.array[np.float32, np.float32, np.float32] - voxel grid of quantity - - See Also - -------- - - scatter_parallel : Parallel implementation of this function - slice_scatter : Create scatter plot of a slice of data - slice_scatter_parallel : Create scatter plot of a slice of data in parallel - - Notes - ----- - - Explicitly defining the types in this function allows - for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba np.ones. - """ - # Output np.array for our image - res_z = res // res_ratio_z - image = np.zeros((res, res, res_z), dtype=np.float32) - maximal_array_index = np.int32(res) - 1 - maximal_array_index_z = np.int32(res_z) - 1 - - # Change that integer to a float, we know that our x, y are bounded - # by [0, 1]. - float_res = np.float32(res) - float_res_z = np.float32(res_z) - pixel_width = 1.0 / float_res - pixel_width_z = 1.0 / float_res_z - - # We need this for combining with the x_pos and y_pos variables. - float_res_64 = np.float64(res) - float_res_z_64 = np.float64(res_z) - - # If the kernel width is smaller than this, we drop to just PIC method - drop_to_single_cell = pixel_width * 0.5 - drop_to_single_cell_z = pixel_width_z * 0.5 - - # Pre-calculate this constant for use with the above - inverse_cell_volume = float_res * float_res * float_res_z - - if box_x == 0.0: - xshift_min = 0 - xshift_max = 1 - else: - xshift_min = -1 - xshift_max = 2 - if box_y == 0.0: - yshift_min = 0 - yshift_max = 1 - else: - yshift_min = -1 - yshift_max = 2 - if box_z == 0.0: - zshift_min = 0 - zshift_max = 1 - else: - zshift_min = -1 - zshift_max = 2 - - for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( - x, y, z, m, h - ): - # loop over periodic copies of the particle - for xshift in range(xshift_min, xshift_max): - for yshift in range(yshift_min, yshift_max): - for zshift in range(zshift_min, zshift_max): - x_pos = x_pos_original + xshift * box_x - y_pos = y_pos_original + yshift * box_y - z_pos = z_pos_original + zshift * box_z - - # Calculate the cell that this particle; use the 64 bit version of the - # resolution as this is the same type as the positions - particle_cell_x = np.int32(float_res_64 * x_pos) - particle_cell_y = np.int32(float_res_64 * y_pos) - particle_cell_z = np.int32(float_res_z_64 * z_pos) - - # SWIFT stores hsml as the FWHM. - kernel_width = kernel_gamma * hsml - - # The number of cells that this kernel spans - cells_spanned = np.int32(1.0 + kernel_width * float_res) - cells_spanned_z = np.int32(1.0 + kernel_width * float_res_z) - - if ( - particle_cell_x + cells_spanned < 0 - or particle_cell_x - cells_spanned > maximal_array_index - or particle_cell_y + cells_spanned < 0 - or particle_cell_y - cells_spanned > maximal_array_index - or particle_cell_z + cells_spanned_z < 0 - or particle_cell_z - cells_spanned_z > maximal_array_index_z - ): - # Can happily skip this particle - continue - - if ( - kernel_width < drop_to_single_cell - or kernel_width < drop_to_single_cell_z - ): - # Easygame, gg - if ( - particle_cell_x >= 0 - and particle_cell_x <= maximal_array_index - and particle_cell_y >= 0 - and particle_cell_y <= maximal_array_index - and particle_cell_z >= 0 - and particle_cell_z <= maximal_array_index_z - ): - image[ - particle_cell_x, particle_cell_y, particle_cell_z - ] += (mass * inverse_cell_volume) - else: - # Now we loop over the square of cells that the kernel lives in - for cell_x in range( - # Ensure that the lowest x value is 0, otherwise we segfault - max(0, particle_cell_x - cells_spanned), - # Ensure that the highest x value lies within the np.array - # bounds, otherwise we'll segfault (oops). - min( - particle_cell_x + cells_spanned, maximal_array_index + 1 - ), - ): - # The distance in x to our new favourite cell -- remember that - # our x, y are all in a box of [0, 1]; calculate the distance - # to the cell centre - distance_x = ( - np.float32(cell_x) + 0.5 - ) * pixel_width - np.float32(x_pos) - distance_x_2 = distance_x * distance_x - for cell_y in range( - max(0, particle_cell_y - cells_spanned), - min( - particle_cell_y + cells_spanned, - maximal_array_index + 1, - ), - ): - distance_y = ( - np.float32(cell_y) + 0.5 - ) * pixel_width - np.float32(y_pos) - distance_y_2 = distance_y * distance_y - for cell_z in range( - max(0, particle_cell_z - cells_spanned_z), - min( - particle_cell_z + cells_spanned_z, - maximal_array_index_z + 1, - ), - ): - distance_z = ( - np.float32(cell_z) + 0.5 - ) * pixel_width_z - np.float32(z_pos) - distance_z_2 = distance_z * distance_z - - r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) - - kernel_eval = kernel(r, kernel_width) - - image[cell_x, cell_y, cell_z] += mass * kernel_eval - - return image - - -@jit(nopython=True, fastmath=True, parallel=True) -def scatter_parallel( - x: np.float64, - y: np.float64, - z: np.float64, - m: np.float32, - h: np.float32, - res: int, - res_ratio_z: int = 1, - box_x: np.float64 = 0.0, - box_y: np.float64 = 0.0, - box_z: np.float64 = 0.0, -) -> np.ndarray: - """ - Parallel implementation of scatter - - Compute contributions to a voxel grid from particles with positions - (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. - This ignores boundary effects. - - Parameters - ---------- - x : np.array of np.float64 - np.array of x-positions of the particles. Must be bounded by [0, 1]. - - y : np.array of np.float64 - np.array of y-positions of the particles. Must be bounded by [0, 1]. - - z : np.array of np.float64 - np.array of z-positions of the particles. Must be bounded by [0, 1]. - - m : np.array of np.float32 - np.array of masses (or otherwise weights) of the particles - - h : np.array of np.float32 - np.array of smoothing lengths of the particles - - res : int - the number of voxels along one axis, i.e. this returns a cube - of res * res * res. - - res_ratio_z: int - the number of voxels along the x and y axes relative to the z - - box_x: np.float64 - box size in x, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_y: np.float64 - box size in y, in the same rescaled length units as x, y and z. - Used for periodic wrapping. - - box_z: np.float64 - box size in z, in the same rescaled length units as x, y and z. - Used for periodic wrapping - - Returns - ------- - - np.ndarray of np.float32 - voxel grid of quantity - - See Also - -------- - - scatter : Create voxel grid of quantity - slice_scatter : Create scatter plot of a slice of data - slice_scatter_parallel : Create scatter plot of a slice of data in parallel - - Notes - ----- - - Explicitly defining the types in this function allows - for a 25-50% performance improvement. In our testing, using numpy - floats and integers is also an improvement over using the numba np.ones. - - """ - # Same as scatter, but executes in parallel! This is actually trivial, - # we just make NUM_THREADS images and add them together at the end. - - number_of_particles = x.size - core_particles = number_of_particles // NUM_THREADS - - output = np.zeros((res, res, res), dtype=np.float32) - - for thread in prange(NUM_THREADS): - # Left edge is easy, just start at 0 and go to 'final' - left_edge = thread * core_particles - - # Right edge is harder in case of left over particles... - right_edge = thread + 1 - - if right_edge == NUM_THREADS: - right_edge = number_of_particles - else: - right_edge *= core_particles - - # using kwargs is unsupported in numba - if res_ratio_z == 1: - output += scatter( - x=x[left_edge:right_edge], - y=y[left_edge:right_edge], - z=z[left_edge:right_edge], - m=m[left_edge:right_edge], - h=h[left_edge:right_edge], - res=res, - box_x=box_x, - box_y=box_y, - box_z=box_z, - ) - else: - output += scatter_limited_z( - x=x[left_edge:right_edge], - y=y[left_edge:right_edge], - z=z[left_edge:right_edge], - m=m[left_edge:right_edge], - h=h[left_edge:right_edge], - res=res, - res_ratio_z=res_ratio_z, - box_x=box_x, - box_y=box_y, - box_z=box_z, - ) - - return output - - -def render_gas_voxel_grid( +def render_gas( data: SWIFTDataset, resolution: int, project: Union[str, None] = "masses", @@ -660,122 +103,15 @@ def render_gas_voxel_grid( box_x=region_info["periodic_box_x"], box_y=region_info["periodic_box_y"], box_z=region_info["periodic_box_z"], + norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]), ) - image = scatter_parallel(**arguments) if parallel else scatter(**arguments) - - return image - - -def render_gas( - data: SWIFTDataset, - resolution: int, - project: Union[str, None] = "masses", - parallel: bool = False, - rotation_matrix: Union[None, np.array] = None, - rotation_center: Union[None, cosmo_array] = None, - region: Union[None, cosmo_array] = None, - periodic: bool = True, -): - """ - Creates a 3D voxel grid of a SWIFT dataset, weighted by data field - - Parameters - ---------- - - data : SWIFTDataset - Dataset from which slice is extracted - - resolution : int - Specifies size of return np.array - - project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles - - parallel : bool - used to determine if we will create the image in parallel. This - defaults to False, but can speed up the creation of large images - significantly at the cost of increased memory usage. - - rotation_matrix: np.np.array, optional - Rotation matrix (3x3) that describes the rotation of the box around - ``rotation_center``. In the default case, this provides a volume render - viewed along the z axis. - - rotation_center: cosmo_array, optional - Center of the rotation. If you are trying to rotate around a galaxy, this - should be the most bound particle. - - region : cosmo_array, optional - determines where the image will be created - (this corresponds to the left and right-hand edges, and top and bottom - edges, and front and back edges) if it is not None. It should have a - length of six, and take the form: - ``[x_min, x_max, y_min, y_max, z_min, z_max]`` - Particles outside of this range are still considered if their - smoothing lengths overlap with the range. - - periodic : bool, optional - Account for periodic boundaries for the simulation box? - Default is ``True``. - - Returns - ------- - - np.ndarray of np.float32 - a `resolution` x `resolution` x `resolution` np.array of the contribution - of the projected data field to the voxel grid from all of the particles - - See Also - -------- - - slice_gas : Creates a 2D slice of a SWIFT dataset with appropriate units - render_gas_voxel_grid : Creates a 3D voxel grid of a SWIFT dataset - - Notes - ----- - - This is a wrapper function for slice_gas_pixel_grid ensuring that output - units are appropriate - """ - - image = render_gas_voxel_grid( - data, - resolution, - project, - parallel, - rotation_matrix, - rotation_center, - region=region, - periodic=periodic, + image = ( + backends_parallel["scatter"](**arguments) + if parallel + else backends["scatter"](**arguments) ) - if region is not None: - x_range = region[1] - region[0] - y_range = region[3] - region[2] - z_range = region[5] - region[4] - units = 1.0 / (x_range * y_range * z_range) - units.convert_to_units(1.0 / (x_range.units * y_range.units * z_range.units)) - else: - units = 1.0 / ( - data.metadata.boxsize[0] - * data.metadata.boxsize[1] - * data.metadata.boxsize[2] - ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) - - comoving = data.gas.coordinates.comoving - coord_cosmo_factor = data.gas.coordinates.cosmo_factor - if project is not None: - units *= getattr(data.gas, project).units - project_cosmo_factor = getattr(data.gas, project).cosmo_factor - new_cosmo_factor = project_cosmo_factor / coord_cosmo_factor**3 - else: - new_cosmo_factor = coord_cosmo_factor ** (-3) - - return cosmo_array( - image, units=units, cosmo_factor=new_cosmo_factor, comoving=comoving - ) + return image @jit(nopython=True, fastmath=True) diff --git a/swiftsimio/visualisation/volume_render_backends/__init__.py b/swiftsimio/visualisation/volume_render_backends/__init__.py new file mode 100644 index 00000000..0752a1f7 --- /dev/null +++ b/swiftsimio/visualisation/volume_render_backends/__init__.py @@ -0,0 +1,18 @@ +""" +Backends for volume rendering +""" + +from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units + +from swiftsimio.visualisation.volume_render_backends.scatter import ( + scatter, + scatter_parallel, +) + +backends = { + "scatter": backends_restore_cosmo_and_units(scatter), +} + +backends_parallel = { + "scatter": backends_restore_cosmo_and_units(scatter_parallel), +} diff --git a/swiftsimio/visualisation/volume_render_backends/scatter.py b/swiftsimio/visualisation/volume_render_backends/scatter.py new file mode 100644 index 00000000..d52218bd --- /dev/null +++ b/swiftsimio/visualisation/volume_render_backends/scatter.py @@ -0,0 +1,568 @@ +""" +Basic volume render for SPH data. This takes the 3D positions +of the particles and projects them onto a grid. +""" + +from math import sqrt +import numpy as np + +from swiftsimio.accelerated import jit, NUM_THREADS, prange + +from swiftsimio.visualisation.slice_backends.sph import kernel, kernel_gamma + + +@jit(nopython=True, fastmath=True) +def scatter( + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, + res: int, + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: + """ + Creates a weighted voxel grid + + Computes contributions to a voxel grid from particles with positions + (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. + This includes periodic boundary effects. + + Parameters + ---------- + + x : np.np.array[np.float64] + np.array of x-positions of the particles. Must be bounded by [0, 1]. + + y : np.np.array[np.float64] + np.array of y-positions of the particles. Must be bounded by [0, 1]. + + z : np.np.array[np.float64] + np.array of z-positions of the particles. Must be bounded by [0, 1]. + + m : np.np.array[np.float32] + np.array of masses (or otherwise weights) of the particles + + h : np.np.array[np.float32] + np.array of smoothing lengths of the particles + + res : int + the number of voxels along one axis, i.e. this returns a cube + of res * res * res. + + box_x: np.float64 + box size in x, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_y: np.float64 + box size in y, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_z: np.float64 + box size in z, in the same rescaled length units as x, y and z. + Used for periodic wrapping + + Returns + ------- + + np.np.array[np.float32, np.float32, np.float32] + voxel grid of quantity + + See Also + -------- + + scatter_parallel : Parallel implementation of this function + slice_scatter : Create scatter plot of a slice of data + slice_scatter_parallel : Create scatter plot of a slice of data in parallel + + Notes + ----- + + Explicitly defining the types in this function allows + for a 25-50% performance improvement. In our testing, using numpy + floats and integers is also an improvement over using the numba np.ones. + """ + # Output np.array for our image + image = np.zeros((res, res, res), dtype=np.float32) + maximal_array_index = np.int32(res) - 1 + + # Change that integer to a float, we know that our x, y are bounded + # by [0, 1]. + float_res = np.float32(res) + pixel_width = 1.0 / float_res + + # We need this for combining with the x_pos and y_pos variables. + float_res_64 = np.float64(res) + + # If the kernel width is smaller than this, we drop to just PIC method + drop_to_single_cell = pixel_width * 0.5 + + # Pre-calculate this constant for use with the above + inverse_cell_volume = float_res * float_res * float_res + + if box_x == 0.0: + xshift_min = 0 + xshift_max = 1 + else: + xshift_min = -1 + xshift_max = 2 + if box_y == 0.0: + yshift_min = 0 + yshift_max = 1 + else: + yshift_min = -1 + yshift_max = 2 + if box_z == 0.0: + zshift_min = 0 + zshift_max = 1 + else: + zshift_min = -1 + zshift_max = 2 + + for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( + x, y, z, m, h + ): + # loop over periodic copies of the particle + for xshift in range(xshift_min, xshift_max): + for yshift in range(yshift_min, yshift_max): + for zshift in range(zshift_min, zshift_max): + x_pos = x_pos_original + xshift * box_x + y_pos = y_pos_original + yshift * box_y + z_pos = z_pos_original + zshift * box_z + + # Calculate the cell that this particle; use the 64 bit version of the + # resolution as this is the same type as the positions + particle_cell_x = np.int32(float_res_64 * x_pos) + particle_cell_y = np.int32(float_res_64 * y_pos) + particle_cell_z = np.int32(float_res_64 * z_pos) + + # SWIFT stores hsml as the FWHM. + kernel_width = kernel_gamma * hsml + + # The number of cells that this kernel spans + cells_spanned = np.int32(1.0 + kernel_width * float_res) + + if ( + particle_cell_x + cells_spanned < 0 + or particle_cell_x - cells_spanned > maximal_array_index + or particle_cell_y + cells_spanned < 0 + or particle_cell_y - cells_spanned > maximal_array_index + or particle_cell_z + cells_spanned < 0 + or particle_cell_z - cells_spanned > maximal_array_index + ): + # Can happily skip this particle + continue + + if kernel_width < drop_to_single_cell: + # Easygame, gg + if ( + particle_cell_x >= 0 + and particle_cell_x <= maximal_array_index + and particle_cell_y >= 0 + and particle_cell_y <= maximal_array_index + and particle_cell_z >= 0 + and particle_cell_z <= maximal_array_index + ): + image[ + particle_cell_x, particle_cell_y, particle_cell_z + ] += (mass * inverse_cell_volume) + else: + # Now we loop over the square of cells that the kernel lives in + for cell_x in range( + # Ensure that the lowest x value is 0, otherwise we segfault + max(0, particle_cell_x - cells_spanned), + # Ensure that the highest x value lies within the np.array + # bounds, otherwise we'll segfault (oops). + min( + particle_cell_x + cells_spanned, maximal_array_index + 1 + ), + ): + # The distance in x to our new favourite cell -- remember that + # our x, y are all in a box of [0, 1]; calculate the distance + # to the cell centre + distance_x = ( + np.float32(cell_x) + 0.5 + ) * pixel_width - np.float32(x_pos) + distance_x_2 = distance_x * distance_x + for cell_y in range( + max(0, particle_cell_y - cells_spanned), + min( + particle_cell_y + cells_spanned, + maximal_array_index + 1, + ), + ): + distance_y = ( + np.float32(cell_y) + 0.5 + ) * pixel_width - np.float32(y_pos) + distance_y_2 = distance_y * distance_y + for cell_z in range( + max(0, particle_cell_z - cells_spanned), + min( + particle_cell_z + cells_spanned, + maximal_array_index + 1, + ), + ): + distance_z = ( + np.float32(cell_z) + 0.5 + ) * pixel_width - np.float32(z_pos) + distance_z_2 = distance_z * distance_z + + r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) + + kernel_eval = kernel(r, kernel_width) + + image[cell_x, cell_y, cell_z] += mass * kernel_eval + + return image + + +@jit(nopython=True, fastmath=True) +def scatter_limited_z( + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, + res: int, + res_ratio_z: int, + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: + """ + Creates a weighted voxel grid + + Computes contributions to a voxel grid from particles with positions + (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. + This includes periodic boundary effects. + + Parameters + ---------- + + x : np.np.array[np.float64] + np.array of x-positions of the particles. Must be bounded by [0, 1]. + + y : np.np.array[np.float64] + np.array of y-positions of the particles. Must be bounded by [0, 1]. + + z : np.np.array[np.float64] + np.array of z-positions of the particles. Must be bounded by [0, 1]. + + m : np.np.array[np.float32] + np.array of masses (or otherwise weights) of the particles + + h : np.np.array[np.float32] + np.array of smoothing lengths of the particles + + res : int + the number of voxels along one axis, i.e. this returns a cube + of res * res * res. + + res_ratio_z: int + the number of voxels along the x and y axes relative to the z + axis. If this is, for instance, 8, and the res is 128, then the + output np.array will be 128 x 128 x 16. + + box_x: np.float64 + box size in x, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_y: np.float64 + box size in y, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_z: np.float64 + box size in z, in the same rescaled length units as x, y and z. + Used for periodic wrapping + + Returns + ------- + + np.np.array[np.float32, np.float32, np.float32] + voxel grid of quantity + + See Also + -------- + + scatter_parallel : Parallel implementation of this function + slice_scatter : Create scatter plot of a slice of data + slice_scatter_parallel : Create scatter plot of a slice of data in parallel + + Notes + ----- + + Explicitly defining the types in this function allows + for a 25-50% performance improvement. In our testing, using numpy + floats and integers is also an improvement over using the numba np.ones. + """ + # Output np.array for our image + res_z = res // res_ratio_z + image = np.zeros((res, res, res_z), dtype=np.float32) + maximal_array_index = np.int32(res) - 1 + maximal_array_index_z = np.int32(res_z) - 1 + + # Change that integer to a float, we know that our x, y are bounded + # by [0, 1]. + float_res = np.float32(res) + float_res_z = np.float32(res_z) + pixel_width = 1.0 / float_res + pixel_width_z = 1.0 / float_res_z + + # We need this for combining with the x_pos and y_pos variables. + float_res_64 = np.float64(res) + float_res_z_64 = np.float64(res_z) + + # If the kernel width is smaller than this, we drop to just PIC method + drop_to_single_cell = pixel_width * 0.5 + drop_to_single_cell_z = pixel_width_z * 0.5 + + # Pre-calculate this constant for use with the above + inverse_cell_volume = float_res * float_res * float_res_z + + if box_x == 0.0: + xshift_min = 0 + xshift_max = 1 + else: + xshift_min = -1 + xshift_max = 2 + if box_y == 0.0: + yshift_min = 0 + yshift_max = 1 + else: + yshift_min = -1 + yshift_max = 2 + if box_z == 0.0: + zshift_min = 0 + zshift_max = 1 + else: + zshift_min = -1 + zshift_max = 2 + + for x_pos_original, y_pos_original, z_pos_original, mass, hsml in zip( + x, y, z, m, h + ): + # loop over periodic copies of the particle + for xshift in range(xshift_min, xshift_max): + for yshift in range(yshift_min, yshift_max): + for zshift in range(zshift_min, zshift_max): + x_pos = x_pos_original + xshift * box_x + y_pos = y_pos_original + yshift * box_y + z_pos = z_pos_original + zshift * box_z + + # Calculate the cell that this particle; use the 64 bit version of the + # resolution as this is the same type as the positions + particle_cell_x = np.int32(float_res_64 * x_pos) + particle_cell_y = np.int32(float_res_64 * y_pos) + particle_cell_z = np.int32(float_res_z_64 * z_pos) + + # SWIFT stores hsml as the FWHM. + kernel_width = kernel_gamma * hsml + + # The number of cells that this kernel spans + cells_spanned = np.int32(1.0 + kernel_width * float_res) + cells_spanned_z = np.int32(1.0 + kernel_width * float_res_z) + + if ( + particle_cell_x + cells_spanned < 0 + or particle_cell_x - cells_spanned > maximal_array_index + or particle_cell_y + cells_spanned < 0 + or particle_cell_y - cells_spanned > maximal_array_index + or particle_cell_z + cells_spanned_z < 0 + or particle_cell_z - cells_spanned_z > maximal_array_index_z + ): + # Can happily skip this particle + continue + + if ( + kernel_width < drop_to_single_cell + or kernel_width < drop_to_single_cell_z + ): + # Easygame, gg + if ( + particle_cell_x >= 0 + and particle_cell_x <= maximal_array_index + and particle_cell_y >= 0 + and particle_cell_y <= maximal_array_index + and particle_cell_z >= 0 + and particle_cell_z <= maximal_array_index_z + ): + image[ + particle_cell_x, particle_cell_y, particle_cell_z + ] += (mass * inverse_cell_volume) + else: + # Now we loop over the square of cells that the kernel lives in + for cell_x in range( + # Ensure that the lowest x value is 0, otherwise we segfault + max(0, particle_cell_x - cells_spanned), + # Ensure that the highest x value lies within the np.array + # bounds, otherwise we'll segfault (oops). + min( + particle_cell_x + cells_spanned, maximal_array_index + 1 + ), + ): + # The distance in x to our new favourite cell -- remember that + # our x, y are all in a box of [0, 1]; calculate the distance + # to the cell centre + distance_x = ( + np.float32(cell_x) + 0.5 + ) * pixel_width - np.float32(x_pos) + distance_x_2 = distance_x * distance_x + for cell_y in range( + max(0, particle_cell_y - cells_spanned), + min( + particle_cell_y + cells_spanned, + maximal_array_index + 1, + ), + ): + distance_y = ( + np.float32(cell_y) + 0.5 + ) * pixel_width - np.float32(y_pos) + distance_y_2 = distance_y * distance_y + for cell_z in range( + max(0, particle_cell_z - cells_spanned_z), + min( + particle_cell_z + cells_spanned_z, + maximal_array_index_z + 1, + ), + ): + distance_z = ( + np.float32(cell_z) + 0.5 + ) * pixel_width_z - np.float32(z_pos) + distance_z_2 = distance_z * distance_z + + r = sqrt(distance_x_2 + distance_y_2 + distance_z_2) + + kernel_eval = kernel(r, kernel_width) + + image[cell_x, cell_y, cell_z] += mass * kernel_eval + + return image + + +@jit(nopython=True, fastmath=True, parallel=True) +def scatter_parallel( + x: np.float64, + y: np.float64, + z: np.float64, + m: np.float32, + h: np.float32, + res: int, + res_ratio_z: int = 1, + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, + box_z: np.float64 = 0.0, +) -> np.ndarray: + """ + Parallel implementation of scatter + + Compute contributions to a voxel grid from particles with positions + (`x`,`y`,`z`) with smoothing lengths `h` weighted by quantities `m`. + This ignores boundary effects. + + Parameters + ---------- + x : np.array of np.float64 + np.array of x-positions of the particles. Must be bounded by [0, 1]. + + y : np.array of np.float64 + np.array of y-positions of the particles. Must be bounded by [0, 1]. + + z : np.array of np.float64 + np.array of z-positions of the particles. Must be bounded by [0, 1]. + + m : np.array of np.float32 + np.array of masses (or otherwise weights) of the particles + + h : np.array of np.float32 + np.array of smoothing lengths of the particles + + res : int + the number of voxels along one axis, i.e. this returns a cube + of res * res * res. + + res_ratio_z: int + the number of voxels along the x and y axes relative to the z + + box_x: np.float64 + box size in x, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_y: np.float64 + box size in y, in the same rescaled length units as x, y and z. + Used for periodic wrapping. + + box_z: np.float64 + box size in z, in the same rescaled length units as x, y and z. + Used for periodic wrapping + + Returns + ------- + + np.ndarray of np.float32 + voxel grid of quantity + + See Also + -------- + + scatter : Create voxel grid of quantity + slice_scatter : Create scatter plot of a slice of data + slice_scatter_parallel : Create scatter plot of a slice of data in parallel + + Notes + ----- + + Explicitly defining the types in this function allows + for a 25-50% performance improvement. In our testing, using numpy + floats and integers is also an improvement over using the numba np.ones. + + """ + # Same as scatter, but executes in parallel! This is actually trivial, + # we just make NUM_THREADS images and add them together at the end. + + number_of_particles = x.size + core_particles = number_of_particles // NUM_THREADS + + output = np.zeros((res, res, res), dtype=np.float32) + + for thread in prange(NUM_THREADS): + # Left edge is easy, just start at 0 and go to 'final' + left_edge = thread * core_particles + + # Right edge is harder in case of left over particles... + right_edge = thread + 1 + + if right_edge == NUM_THREADS: + right_edge = number_of_particles + else: + right_edge *= core_particles + + # using kwargs is unsupported in numba + if res_ratio_z == 1: + output += scatter( + x=x[left_edge:right_edge], + y=y[left_edge:right_edge], + z=z[left_edge:right_edge], + m=m[left_edge:right_edge], + h=h[left_edge:right_edge], + res=res, + box_x=box_x, + box_y=box_y, + box_z=box_z, + ) + else: + output += scatter_limited_z( + x=x[left_edge:right_edge], + y=y[left_edge:right_edge], + z=z[left_edge:right_edge], + m=m[left_edge:right_edge], + h=h[left_edge:right_edge], + res=res, + res_ratio_z=res_ratio_z, + box_x=box_x, + box_y=box_y, + box_z=box_z, + ) + + return output diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 543cf473..c1bdffd7 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -1,28 +1,32 @@ import pytest -from swiftsimio import load -from swiftsimio.visualisation import scatter, slice, volume_render -from swiftsimio.visualisation.projection import ( - scatter_parallel, - project_gas, - project_pixel_grid, -) -from swiftsimio.visualisation.slice import slice_scatter_parallel, slice_gas +from swiftsimio import load, mask +from swiftsimio.visualisation.projection import project_gas, project_pixel_grid +from swiftsimio.visualisation.slice import slice_gas from swiftsimio.visualisation.volume_render import render_gas -from swiftsimio.visualisation.volume_render import scatter as volume_scatter +from swiftsimio.visualisation.ray_trace import panel_gas + +from swiftsimio.visualisation.slice_backends import ( + backends as slice_backends, + backends_parallel as slice_backends_parallel, +) +from swiftsimio.visualisation.volume_render_backends import ( + backends as volume_render_backends, + backends_parallel as volume_render_backends_parallel, +) +from swiftsimio.visualisation.projection_backends import ( + backends as projection_backends, + backends_parallel as projection_backends_parallel, +) + from swiftsimio.visualisation.power_spectrum import ( deposit, deposition_to_power_spectrum, render_to_deposit, folded_depositions_to_power_spectrum, ) -from swiftsimio.visualisation.ray_trace import panel_gas from swiftsimio.visualisation.smoothing_length import generate_smoothing_lengths -from swiftsimio.visualisation.projection_backends import backends as projection_backends -from swiftsimio.visualisation.slice_backends import ( - backends as slice_backends, - backends_parallel as slice_backends_parallel, -) from swiftsimio.optional_packages import CudaSupportError, CUDA_AVAILABLE + from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a from unyt.array import unyt_array import unyt @@ -46,13 +50,14 @@ def test_scatter(save=False): for backend in projection_backends.keys(): try: image = projection_backends[backend]( - np.array([0.0, 1.0, 1.0, -0.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([1.0, 1.0, 1.0, 1.0]), - np.array([0.2, 0.2, 0.2, 0.000002]), - 256, - 1.0, - 1.0, + x=np.array([0.0, 1.0, 1.0, -0.000001]), + y=np.array([0.0, 0.0, 1.0, 1.000001]), + m=np.array([1.0, 1.0, 1.0, 1.0]), + h=np.array([0.2, 0.2, 0.2, 0.000002]), + res=256, + box_x=1.0, + box_y=1.0, + norm=1.0, ) except CudaSupportError: if CUDA_AVAILABLE: @@ -78,11 +83,21 @@ def test_scatter_mass_conservation(): total_mass = np.sum(m) for resolution in resolutions: - image = scatter(x, y, m, h, resolution, 1.0, 1.0) + scatter = projection_backends["fast"] + image = scatter( + x=x, + y=y, + m=m, + h=h, + res=resolution, + box_x=1.0, + box_y=1.0, + norm=1.0, + ) mass_in_image = image.sum() / (resolution**2) # Check mass conservation to 5% - assert np.isclose(mass_in_image, total_mass, 0.05) + assert np.isclose(mass_in_image.view(np.ndarray), total_mass, 0.05) return @@ -105,9 +120,27 @@ def test_scatter_parallel(save=False): hsml = np.random.rand(number_of_parts).astype(np.float32) * h_max masses = np.ones(number_of_parts, dtype=np.float32) - image = scatter(coordinates[0], coordinates[1], masses, hsml, resolution, 1.0, 1.0) + scatter = projection_backends["fast"] + scatter_parallel = projection_backends_parallel["fast"] + image = scatter( + x=coordinates[0], + y=coordinates[1], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + norm=1.0, + ) image_par = scatter_parallel( - coordinates[0], coordinates[1], masses, hsml, resolution, 1.0, 1.0 + x=coordinates[0], + y=coordinates[1], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + norm=1.0, ) if save: @@ -119,17 +152,20 @@ def test_scatter_parallel(save=False): def test_slice(save=False): + slice = slice_backends["sph"] image = slice( - np.array([0.0, 1.0, 1.0, -0.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([1.0, 1.0, 1.0, 1.0]), - np.array([0.2, 0.2, 0.2, 0.000002]), - 0.99, - 256, - 1.0, - 1.0, - 1.0, + x=np.array([0.0, 1.0, 1.0, -0.000001]), + y=np.array([0.0, 0.0, 1.0, 1.000001]), + z=np.array([0.0, 0.0, 1.0, 1.000001]), + m=np.array([1.0, 1.0, 1.0, 1.0]), + h=np.array([0.2, 0.2, 0.2, 0.000002]), + z_slice=0.99, + xres=256, + yres=256, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, ) if save: @@ -159,30 +195,32 @@ def test_slice_parallel(save=False): for backend in slice_backends.keys(): image = slice_backends[backend]( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - z_slice, - resolution, - resolution, - 1.0, - 1.0, - 1.0, + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + z_slice=z_slice, + xres=resolution, + yres=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, ) image_par = slice_backends_parallel[backend]( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - z_slice, - resolution, - resolution, - 1.0, - 1.0, - 1.0, + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + z_slice=z_slice, + xres=resolution, + yres=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, ) assert np.isclose(image, image_par).all() @@ -195,16 +233,18 @@ def test_slice_parallel(save=False): def test_volume_render(): # render image - volume_render.scatter( - np.array([0.0, 1.0, 1.0, -0.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([0.0, 0.0, 1.0, 1.000001]), - np.array([1.0, 1.0, 1.0, 1.0]), - np.array([0.2, 0.2, 0.2, 0.000002]), - 64, - 1.0, - 1.0, - 1.0, + scatter = volume_render_backends["scatter"] + scatter( + x=np.array([0.0, 1.0, 1.0, -0.000001]), + y=np.array([0.0, 0.0, 1.0, 1.000001]), + z=np.array([0.0, 0.0, 1.0, 1.000001]), + m=np.array([1.0, 1.0, 1.0, 1.0]), + h=np.array([0.2, 0.2, 0.2, 0.000002]), + res=64, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, ) return @@ -223,28 +263,31 @@ def test_volume_parallel(): hsml = np.random.rand(number_of_parts).astype(np.float32) * h_max masses = np.ones(number_of_parts, dtype=np.float32) - image = volume_render.scatter( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - resolution, - 1.0, - 1.0, - 1.0, + scatter = volume_render_backends["scatter"] + image = scatter( + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, ) - image_par = volume_render.scatter_parallel( - coordinates[0], - coordinates[1], - coordinates[2], - masses, - hsml, - resolution, - 1, - 1.0, - 1.0, - 1.0, + scatter_parallel = volume_render_backends_parallel["scatter"] + image_par = scatter_parallel( + x=coordinates[0], + y=coordinates[1], + z=coordinates[2], + m=masses, + h=hsml, + res=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, ) assert np.isclose(image, image_par).all() @@ -266,7 +309,7 @@ def test_selection_render(filename): project_gas(data, 256, parallel=True, region=[0 * bs, 0.001 * bs] * 2) # render non-square project_gas( - data, 256, parallel=True, region=[0 * bs, 0.00 * bs, 0.25 * bs, 0.75 * bs] + data, 256, parallel=True, region=[0 * bs, 0.50 * bs, 0.25 * bs, 0.75 * bs] ) # Slicing @@ -309,20 +352,48 @@ def test_render_outside_region(): h = 10 ** np.random.rand(number_of_parts) - 1.0 h[h > 0.5] = 0.05 m = np.ones_like(h) - projection_backends["histogram"](x, y, m, h, resolution, 1.0, 1.0) + projection_backends["histogram"]( + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, norm=1.0 + ) for backend in projection_backends.keys(): try: - projection_backends[backend](x, y, m, h, resolution, 1.0, 1.0) + projection_backends[backend]( + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, norm=1.0 + ) except CudaSupportError: if CUDA_AVAILABLE: raise ImportError("Optional loading of the CUDA module is broken") else: continue - slice_scatter_parallel(x, y, z, m, h, 0.2, resolution, 1.0, 1.0, 1.0) + slice_backends_parallel["sph"]( + x=x, + y=y, + z=z, + m=m, + h=h, + z_slice=0.2, + xres=resolution, + yres=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, + ) - volume_render.scatter_parallel(x, y, z, m, h, resolution, 1, 1.0, 1.0, 1.0) + volume_render_backends_parallel["scatter"]( + x=x, + y=y, + z=z, + m=m, + h=h, + res=resolution, + box_x=1.0, + box_y=1.0, + box_z=1.0, + norm=1.0, + ) @requires("cosmological_volume.hdf5") @@ -331,40 +402,81 @@ def test_comoving_versus_physical(filename): Test what happens if you try to mix up physical and comoving quantities. """ - for func, aexp in [(project_gas, -2.0), (slice_gas, -3.0), (render_gas, -3.0)]: + # this test is pretty slow if we don't mask out some particles + m = mask(filename) + boxsize = m.metadata.boxsize + m.constrain_spatial([[0.0 * b, 0.2 * b] for b in boxsize]) + region = [ + 0.0 * boxsize[0], + 0.2 * boxsize[0], + 0.0 * boxsize[1], + 0.2 * boxsize[1], + 0.0 * boxsize[2], + 0.2 * boxsize[2], + ] + for func, aexp in [ + (project_gas, -2.0), + (slice_gas, -3.0), + (render_gas, -3.0), + ]: # normal case: everything comoving - data = load(filename) + data = load(filename, mask=m) # we force the default (project="masses") to check the cosmo_factor # conversion in this case - img = func(data, resolution=256, project=None) - assert img.comoving + img = func(data, resolution=64, project=None, region=region) + assert data.gas.masses.comoving and img.comoving assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0 - img = func(data, resolution=256, project="densities") - assert img.comoving + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving and img.comoving assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 - # try to mix comoving coordinates with a physical variable + # try to mix comoving coordinates with a physical variable: + # the coordinates should convert to physical internally and warn data.gas.densities.convert_to_physical() - with pytest.raises(AttributeError, match="not compatible with comoving"): - img = func(data, resolution=256, project="densities") - # convert coordinates to physical (but not smoothing lengths) + with pytest.warns( + UserWarning, match="Converting smoothing lengths to physical." + ): + with pytest.warns( + UserWarning, match="Converting coordinate grid to physical." + ): + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving is False and img.comoving is False + assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 + # convert coordinates to physical (but not smoothing lengths): + # the coordinates (copy) should convert back to comoving to match the masses data.gas.coordinates.convert_to_physical() - with pytest.raises(AttributeError, match=""): - img = func(data, resolution=256, project="masses") + with pytest.warns(UserWarning, match="Converting coordinate grid to comoving."): + img = func(data, resolution=64, project="masses", region=region) + assert data.gas.masses.comoving and img.comoving + assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0 # also convert smoothing lengths to physical + # everything should still convert back to comoving to match masses data.gas.smoothing_lengths.convert_to_physical() - # masses are always compatible with either - img = func(data, resolution=256, project="masses") - # check that we get a physical result - assert not img.comoving + with pytest.warns( + UserWarning, match="Converting smoothing lengths to comoving." + ): + with pytest.warns( + UserWarning, match="Converting coordinate grid to comoving." + ): + img = func(data, resolution=64, project="masses", region=region) + assert data.gas.masses.comoving and img.comoving assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 - # densities are still compatible with physical - img = func(data, resolution=256, project="densities") - assert not img.comoving + # densities are physical, make sure this works with physical coordinates and + # smoothing lengths + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving is False and img.comoving is False assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 - # now try again with comoving densities + # now try again with comoving densities, should work and give a comoving img + # with internal conversions to comoving data.gas.densities.convert_to_comoving() - with pytest.raises(AttributeError, match="not compatible with physical"): - img = func(data, resolution=256, project="densities") + with pytest.warns( + UserWarning, match="Converting smoothing lengths to comoving." + ): + with pytest.warns( + UserWarning, match="Converting coordinate grid to comoving." + ): + img = func(data, resolution=64, project="densities", region=region) + assert data.gas.densities.comoving and img.comoving + assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0 @requires("cosmological_volume.hdf5") @@ -373,7 +485,7 @@ def test_nongas_smoothing_lengths(filename): Test that the visualisation tools to calculate smoothing lengths give usable results. """ - # If project_pixel_grid runs without error the smoothing lengths seem usable. + # If project_gas runs without error the smoothing lengths seem usable. data = load(filename) data.dark_matter.smoothing_length = generate_smoothing_lengths( data.dark_matter.coordinates, data.metadata.boxsize, kernel_gamma=1.8 @@ -478,6 +590,7 @@ def test_periodic_boundary_wrapping(): res=pixel_resolution, box_x=boxsize, box_y=boxsize, + norm=boxsize**2, ) image2 = projection_backends[backend]( x=coordinates_non_periodic[:, 0], @@ -487,6 +600,7 @@ def test_periodic_boundary_wrapping(): res=pixel_resolution, box_x=0.0, box_y=0.0, + norm=boxsize**2, ) assert (image1 == image2).all() except CudaSupportError: @@ -509,6 +623,7 @@ def test_periodic_boundary_wrapping(): box_x=boxsize, box_y=boxsize, box_z=boxsize, + norm=boxsize**3, ) image2 = slice_backends[backend]( x=coordinates_non_periodic[:, 0], @@ -522,12 +637,14 @@ def test_periodic_boundary_wrapping(): box_x=0.0, box_y=0.0, box_z=0.0, + norm=boxsize**3, ) assert (image1 == image2).all() # test the volume rendering scatter function - image1 = volume_render.scatter( + scatter = volume_render_backends["scatter"] + image1 = scatter( x=coordinates_periodic[:, 0], y=coordinates_periodic[:, 1], z=coordinates_periodic[:, 2], @@ -537,9 +654,9 @@ def test_periodic_boundary_wrapping(): box_x=boxsize, box_y=boxsize, box_z=boxsize, + norm=boxsize**3, ) - - image2 = volume_render.scatter( + image2 = scatter( x=coordinates_non_periodic[:, 0], y=coordinates_non_periodic[:, 1], z=coordinates_non_periodic[:, 2], @@ -549,6 +666,7 @@ def test_periodic_boundary_wrapping(): box_x=0.0, box_y=0.0, box_z=0.0, + norm=boxsize**3, ) assert (image1 == image2).all() @@ -572,20 +690,22 @@ def test_volume_render_and_unfolded_deposit(): # 1.0 implies no folding deposition = deposit(x, y, z, m, res, 1.0, boxsize, boxsize, boxsize) - # Need to norm for the volume render - volume = volume_scatter( - x / boxsize, - y / boxsize, - z / boxsize, - m, - h / boxsize, - res, - boxsize, - boxsize, - boxsize, + # Need to norm coords and box for the volume render + volume = volume_render_backends["scatter"]( + x=x / boxsize, + y=y / boxsize, + z=z / boxsize, + m=m, + h=h / boxsize, + res=res, + box_x=boxsize, + box_y=boxsize, + box_z=boxsize, + norm=boxsize**3, ) - assert np.allclose(deposition, volume) + # need to divide out the box volume for the deposition + assert np.allclose(deposition / boxsize**3, volume.view(np.ndarray)) def test_folding_deposit(): From e85741cd7ff237c32b8e2bf4024494c96994800d Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 12 Feb 2025 01:42:46 +0000 Subject: [PATCH 086/125] Run black. --- swiftsimio/accelerated.py | 2 +- swiftsimio/masks.py | 4 +- .../metadata/metadata/metadata_fields.py | 2 +- swiftsimio/metadata/objects.py | 38 ++--- swiftsimio/metadata/soap/soap_types.py | 1 + swiftsimio/metadata/unit/unit_fields.py | 8 +- swiftsimio/objects.py | 2 +- swiftsimio/snapshot_writer.py | 2 +- swiftsimio/visualisation/power_spectrum.py | 20 +-- .../projection_backends/histogram.py | 1 - .../projection_backends/reference.py | 1 - .../projection_backends/renormalised.py | 1 - .../projection_backends/subsampled.py | 1 - .../projection_backends/subsampled_extreme.py | 1 - .../slice_backends/nearest_neighbours.py | 2 +- tests/basic_test.py | 2 +- tests/test_cosmo_array.py | 78 +++++----- tests/test_cosmo_array_attrs.py | 134 +++++++++--------- tests/test_data.py | 8 +- tests/test_extraparts.py | 7 +- tests/test_read_ic.py | 2 +- tests/test_soap.py | 4 +- tests/test_vistools.py | 2 +- 23 files changed, 162 insertions(+), 161 deletions(-) diff --git a/swiftsimio/accelerated.py b/swiftsimio/accelerated.py index a8cd8ad8..7382f6d1 100644 --- a/swiftsimio/accelerated.py +++ b/swiftsimio/accelerated.py @@ -123,7 +123,7 @@ def read_ranges_from_file_unchunked( already_read = 0 handle_multidim = handle.ndim > 1 - for (read_start, read_end) in ranges: + for read_start, read_end in ranges: if read_end == read_start: continue diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index 189cb1d0..434fe9d8 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -237,7 +237,7 @@ def _unpack_cell_metadata(self): centers_handle[:][sort], units=self.units.length, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, self.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), ) # Note that we cannot assume that these are cubic, unfortunately. @@ -245,7 +245,7 @@ def _unpack_cell_metadata(self): metadata_handle.attrs["size"], units=self.units.length, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, self.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), ) return diff --git a/swiftsimio/metadata/metadata/metadata_fields.py b/swiftsimio/metadata/metadata/metadata_fields.py index a032e039..75e12a5e 100644 --- a/swiftsimio/metadata/metadata/metadata_fields.py +++ b/swiftsimio/metadata/metadata/metadata_fields.py @@ -77,7 +77,7 @@ def generate_cosmo_args_header_unpack_arrays(scale_factor) -> dict: # should not be cosmo_array'd). cosmo_args = { "boxsize": dict( - cosmo_factor=cosmo_factor(a ** 1, scale_factor=scale_factor), + cosmo_factor=cosmo_factor(a**1, scale_factor=scale_factor), comoving=True, # if it's not, then a=1 and it doesn't matter valid_transform=True, ) diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index e3731620..aba8b5ca 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -113,12 +113,14 @@ def postprocess_header(self): # items including the scale factor. # These must be unpacked as they are stored as length-1 arrays - header_unpack_float_units = metadata.metadata_fields.generate_units_header_unpack_single_float( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, + header_unpack_float_units = ( + metadata.metadata_fields.generate_units_header_unpack_single_float( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, + ) ) for field, names in metadata.metadata_fields.header_unpack_single_float.items(): try: @@ -164,15 +166,19 @@ def postprocess_header(self): self.scale_factor = 1.0 # These are just read straight in to variables - header_unpack_arrays_units = metadata.metadata_fields.generate_units_header_unpack_arrays( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, + header_unpack_arrays_units = ( + metadata.metadata_fields.generate_units_header_unpack_arrays( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, + ) ) - header_unpack_arrays_cosmo_args = metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( - self.scale_factor + header_unpack_arrays_cosmo_args = ( + metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( + self.scale_factor + ) ) for field, name in metadata.metadata_fields.header_unpack_arrays.items(): @@ -604,7 +610,7 @@ def get_units(unit_attribute): # Need to check if the exponent is 0 manually because of float precision unit_exponent = unit_attribute[f"U_{exponent} exponent"][0] if unit_exponent != 0.0: - units *= unit ** unit_exponent + units *= unit**unit_exponent except KeyError: # Can't load that data! # We should probably warn the user here... @@ -698,7 +704,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a ** cosmo_exponent + a_factor_this_dataset = a**cosmo_exponent return cosmo_factor(a_factor_this_dataset, current_scale_factor) diff --git a/swiftsimio/metadata/soap/soap_types.py b/swiftsimio/metadata/soap/soap_types.py index cecb2730..5caa12c2 100644 --- a/swiftsimio/metadata/soap/soap_types.py +++ b/swiftsimio/metadata/soap/soap_types.py @@ -2,6 +2,7 @@ Includes the fancy names. """ + # Describes the conversion of hdf5 groups to names def get_soap_name_underscore(group: str) -> str: soap_name_underscores = { diff --git a/swiftsimio/metadata/unit/unit_fields.py b/swiftsimio/metadata/unit/unit_fields.py index ae5756bd..ed4f6915 100644 --- a/swiftsimio/metadata/unit/unit_fields.py +++ b/swiftsimio/metadata/unit/unit_fields.py @@ -54,11 +54,11 @@ def generate_units(m, l, t, I, T): } gas = { - "density": m / (l ** 3), - "entropy": m * l ** 2 / (t ** 2 * T), + "density": m / (l**3), + "entropy": m * l**2 / (t**2 * T), "internal_energy": (l / t) ** 2, "smoothing_length": l, - "pressure": m / (l * t ** 2), + "pressure": m / (l * t**2), "diffusion": None, "sfr": m / t, "temperature": T, @@ -79,7 +79,7 @@ def generate_units(m, l, t, I, T): sinks = {**shared} stars = { - "birth_density": m / (l ** 3), + "birth_density": m / (l**3), "birth_time": t, "initial_masses": m, "smoothing_length": l, diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 0516d64a..f60c65a7 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -643,7 +643,7 @@ def __pow__(self, p: float) -> "cosmo_factor": """ if self.expr is None: return cosmo_factor(expr=None, scale_factor=self.scale_factor) - 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: "cosmo_factor") -> bool: """ diff --git a/swiftsimio/snapshot_writer.py b/swiftsimio/snapshot_writer.py index 388ead1b..0d05495f 100644 --- a/swiftsimio/snapshot_writer.py +++ b/swiftsimio/snapshot_writer.py @@ -268,7 +268,7 @@ def get_attributes(self, scale_factor: float) -> dict: # Find the scale factor associated quantities a_exp = a_exponents.get(name, 0) - a_factor = scale_factor ** a_exp + a_factor = scale_factor**a_exp attributes_dict[output_handle] = { "Conversion factor to CGS (not including cosmological corrections)": [ diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index fb090937..a7fd0680 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -198,7 +198,7 @@ def render_to_deposit( """ # Get the positions and masses - folding = 2.0 ** folding + folding = 2.0**folding positions = data.coordinates quantity = getattr(data, project) @@ -245,10 +245,10 @@ def render_to_deposit( units = 1.0 / ( data.metadata.boxsize[0] * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) + units.convert_to_units(1.0 / data.metadata.boxsize.units**3) units *= quantity.units - new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor ** 3) + new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor**3) return cosmo_array( deposition, comoving=comoving, cosmo_factor=new_cosmo_factor, units=units @@ -358,7 +358,7 @@ def folded_depositions_to_power_spectrum( np.zeros(number_of_wavenumber_bins), units=box_volume.units, comoving=box_volume.comoving, - cosmo_factor=box_volume.cosmo_factor ** -1, + cosmo_factor=box_volume.cosmo_factor**-1, name="Power spectrum $P(k)$", ) folding_tracker = np.ones(number_of_wavenumber_bins, dtype=float) @@ -396,7 +396,7 @@ def folded_depositions_to_power_spectrum( if folding != final_folding: cutoff_wavenumber = ( - 2.0 ** folding * np.min(depositions[folding].shape) / np.min(box_size) + 2.0**folding * np.min(depositions[folding].shape) / np.min(box_size) ) if cutoff_above_wavenumber_fraction is not None: @@ -421,7 +421,7 @@ def folded_depositions_to_power_spectrum( corrected_wavenumber_centers[prefer_bins] = folded_wavenumber_centers[ prefer_bins ].to(corrected_wavenumber_centers.units) - folding_tracker[prefer_bins] = 2.0 ** folding + folding_tracker[prefer_bins] = 2.0**folding contributed_counts[prefer_bins] = folded_counts[prefer_bins] elif transition == "average": @@ -454,7 +454,7 @@ def folded_depositions_to_power_spectrum( # For debugging, we calculate an effective fold number. folding_tracker[use_bins] = ( - (folding_tracker * existing_weight + (2.0 ** folding) * new_weight) + (folding_tracker * existing_weight + (2.0**folding) * new_weight) / transition_norm )[use_bins] @@ -535,7 +535,7 @@ def deposition_to_power_spectrum( deposition.shape == cross_deposition.shape ), "Depositions must have the same shape" - folding = 2.0 ** folding + folding = 2.0**folding box_size_folded = box_size[0] / folding npix = deposition.shape[0] @@ -557,7 +557,7 @@ def deposition_to_power_spectrum( else: conj_fft = fft.conj() - fourier_amplitudes = (fft * conj_fft).real * box_size_folded ** 3 + fourier_amplitudes = (fft * conj_fft).real * box_size_folded**3 # Calculate k-value spacing (centered FFT) dk = 2 * np.pi / (box_size_folded) @@ -589,7 +589,7 @@ def deposition_to_power_spectrum( divisor[zero_mask] = 1 # Correct for folding - binned_amplitudes *= folding ** 3 + binned_amplitudes *= folding**3 # Correct units and names wavenumbers = binned_wavenumbers / divisor diff --git a/swiftsimio/visualisation/projection_backends/histogram.py b/swiftsimio/visualisation/projection_backends/histogram.py index cb82a540..0fa392f1 100644 --- a/swiftsimio/visualisation/projection_backends/histogram.py +++ b/swiftsimio/visualisation/projection_backends/histogram.py @@ -4,7 +4,6 @@ Uses double precision. """ - from typing import Union from math import sqrt, ceil from numpy import float32, float64, int32, zeros, ndarray diff --git a/swiftsimio/visualisation/projection_backends/reference.py b/swiftsimio/visualisation/projection_backends/reference.py index dfdd59a0..942f3a10 100644 --- a/swiftsimio/visualisation/projection_backends/reference.py +++ b/swiftsimio/visualisation/projection_backends/reference.py @@ -5,7 +5,6 @@ Uses double precision. """ - from typing import Union from math import sqrt, ceil from numpy import float32, float64, int32, zeros, ndarray diff --git a/swiftsimio/visualisation/projection_backends/renormalised.py b/swiftsimio/visualisation/projection_backends/renormalised.py index 295cab0c..06fe392f 100644 --- a/swiftsimio/visualisation/projection_backends/renormalised.py +++ b/swiftsimio/visualisation/projection_backends/renormalised.py @@ -6,7 +6,6 @@ conserved up to floating point precision. """ - """ The original smoothing code. This provides basic renormalisation of the kernel on each call. diff --git a/swiftsimio/visualisation/projection_backends/subsampled.py b/swiftsimio/visualisation/projection_backends/subsampled.py index 15ee5432..7f859c86 100644 --- a/swiftsimio/visualisation/projection_backends/subsampled.py +++ b/swiftsimio/visualisation/projection_backends/subsampled.py @@ -7,7 +7,6 @@ Uses double precision. """ - """ The original smoothing code. This provides a paranoid supersampling of the kernel. diff --git a/swiftsimio/visualisation/projection_backends/subsampled_extreme.py b/swiftsimio/visualisation/projection_backends/subsampled_extreme.py index cfd0a2f3..8e2cab5c 100644 --- a/swiftsimio/visualisation/projection_backends/subsampled_extreme.py +++ b/swiftsimio/visualisation/projection_backends/subsampled_extreme.py @@ -7,7 +7,6 @@ Uses double precision. """ - """ The original smoothing code. This provides a paranoid supersampling of the kernel. diff --git a/swiftsimio/visualisation/slice_backends/nearest_neighbours.py b/swiftsimio/visualisation/slice_backends/nearest_neighbours.py index 37656bb1..bb5e9504 100644 --- a/swiftsimio/visualisation/slice_backends/nearest_neighbours.py +++ b/swiftsimio/visualisation/slice_backends/nearest_neighbours.py @@ -130,7 +130,7 @@ def slice_scatter( tree = build_tree(x, y, z, box_x, box_y, box_z) _, i = tree.query(pixel_coordinates, workers=workers) - values = m / h ** 3 + values = m / h**3 return values[i].reshape(xres, yres) diff --git a/tests/basic_test.py b/tests/basic_test.py index adb4e7fc..24fb9295 100644 --- a/tests/basic_test.py +++ b/tests/basic_test.py @@ -23,7 +23,7 @@ def test_write(): x = Writer("galactic", boxsize) # 32^3 particles. - n_p = 32 ** 3 + n_p = 32**3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index ba112464..65fa3081 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -28,16 +28,14 @@ def ca(x, unit=u.Mpc): """ Helper for our tests: turn an array into a cosmo_array. """ - 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 cq(x, unit=u.Mpc): """ Helper for our tests: turn a scalar into a cosmo_quantity. """ - return cosmo_quantity( - x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5) - ) + return cosmo_quantity(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) def arg_to_ua(arg): @@ -115,7 +113,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") @@ -129,7 +127,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") @@ -142,7 +140,7 @@ 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") @@ -158,7 +156,7 @@ 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") @@ -178,14 +176,14 @@ def test_init_from_list_of_cosmo_arrays(self): [1], units=u.Mpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1), + 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 @@ -573,7 +571,7 @@ def test_explicitly_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), ), ), ), @@ -584,13 +582,13 @@ def test_explicitly_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), ), ), ), @@ -602,19 +600,19 @@ def test_explicitly_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), ), ], ), @@ -626,7 +624,7 @@ def test_explicitly_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]), ), @@ -650,19 +648,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] @@ -709,9 +707,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): @@ -719,9 +717,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): @@ -729,9 +727,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 = { @@ -743,9 +741,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), ] ), ): @@ -802,12 +800,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 def test_scalar_return_func(self): @@ -819,7 +817,7 @@ def test_scalar_return_func(self): np.arange(3), u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), valid_transform=True, ) res = np.min(ca) @@ -834,12 +832,12 @@ 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 @@ -855,7 +853,7 @@ def test_copy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) copy_arr = copy(arr) @@ -871,7 +869,7 @@ def test_deepcopy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) copy_arr = deepcopy(arr) @@ -887,7 +885,7 @@ def test_to_cgs(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) cgs_arr = arr.in_cgs() diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_attrs.py index d4b1562c..9b38cfbc 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_attrs.py @@ -143,7 +143,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 0.5), + cosmo_factor=cosmo_factor(a**1, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -152,7 +152,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 0.5), + cosmo_factor=cosmo_factor(a**1, 0.5), comoving=False, ) assert not arr.compatible_with_comoving() @@ -161,7 +161,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 0, 0.5), + cosmo_factor=cosmo_factor(a**0, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -170,7 +170,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 0, 0.5), + cosmo_factor=cosmo_factor(a**0, 0.5), comoving=False, ) assert arr.compatible_with_comoving() @@ -179,7 +179,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), comoving=True, ) assert arr.compatible_with_comoving() @@ -188,7 +188,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), comoving=False, ) assert arr.compatible_with_comoving() @@ -232,7 +232,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.ones_like(inp) assert res.to_value(u.kpc) == 1 @@ -265,7 +265,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises( @@ -278,7 +278,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -289,13 +289,13 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -306,7 +306,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp + inp assert res.to_value(u.kpc) == 4 @@ -320,7 +320,7 @@ def test_multiplying_ufunc(self): # no cosmo_factors inp = cosmo_array([2], u.kpc, comoving=False) res = inp * inp - assert res.to_value(u.kpc ** 2) == 4 + assert res.to_value(u.kpc**2) == 4 assert res.comoving is False assert res.cosmo_factor == cosmo_factor(None, None) # one is not cosmo_array @@ -342,7 +342,7 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): @@ -353,7 +353,7 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): inp1 * inp2 @@ -362,12 +362,12 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp * inp - assert res.to_value(u.kpc ** 2) == 4 + assert res.to_value(u.kpc**2) == 4 assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 2 + assert res.cosmo_factor == inp.cosmo_factor**2 def test_dividing_ufunc(self): """ @@ -377,12 +377,12 @@ def test_dividing_ufunc(self): [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp / inp assert res.to_value(u.dimensionless) == 1 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 0 + assert res.cosmo_factor == inp.cosmo_factor**0 def test_return_without_ufunc(self): """ @@ -393,7 +393,7 @@ def test_return_without_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.logical_not(inp) assert res == np.logical_not(1) @@ -403,7 +403,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.logaddexp(inp, inp) assert res == np.logaddexp(2, 2) @@ -413,13 +413,13 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -431,7 +431,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -442,7 +442,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises( @@ -457,7 +457,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) @@ -468,7 +468,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = u.unyt_array([2], u.kpc) with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -484,12 +484,12 @@ def test_sqrt_ufunc(self): [4], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.sqrt(inp) - assert res.to_value(u.kpc ** 0.5) == 2 # also ensures units ok + assert res.to_value(u.kpc**0.5) == 2 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 0.5 + assert res.cosmo_factor == inp.cosmo_factor**0.5 def test_square_ufunc(self): """ @@ -499,12 +499,12 @@ def test_square_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.square(inp) - assert res.to_value(u.kpc ** 2) == 4 # also ensures units ok + assert res.to_value(u.kpc**2) == 4 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 2 + assert res.cosmo_factor == inp.cosmo_factor**2 def test_cbrt_ufunc(self): """ @@ -514,7 +514,7 @@ def test_cbrt_ufunc(self): [8], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.cbrt(inp) assert res.to_value(u.kpc ** (1.0 / 3.0)) == 2 # also ensures units ok @@ -529,12 +529,12 @@ def test_reciprocal_ufunc(self): [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.reciprocal(inp) - assert res.to_value(u.kpc ** -1) == 0.5 # also ensures units ok + assert res.to_value(u.kpc**-1) == 0.5 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** -1 + assert res.cosmo_factor == inp.cosmo_factor**-1 def test_passthrough_ufunc(self): """ @@ -545,7 +545,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.negative(inp) assert res.to_value(u.kpc) == -2 @@ -556,7 +556,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.copysign(inp, inp) assert res.to_value(u.kpc) == inp.to_value(u.kpc) @@ -567,13 +567,13 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -588,7 +588,7 @@ def test_arctan2_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.arctan2(inp, inp) assert res.to_value(u.dimensionless) == np.arctan2(2, 2) @@ -603,13 +603,13 @@ def test_comparison_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp1 < inp2 assert res.all() @@ -623,7 +623,7 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) out = cosmo_array([np.nan], u.dimensionless, comoving=True) np.abs(inp, out=out) @@ -634,7 +634,7 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) # make sure we can also pass a non-cosmo type for out without crashing out = np.array([np.nan]) @@ -649,12 +649,12 @@ def test_reduce_multiply(self): [[1, 2], [3, 4]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) res = np.multiply.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc ** 2), np.array([3.0, 8.0])) + np.testing.assert_allclose(res.to_value(u.kpc**2), np.array([3.0, 8.0])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** 2 + assert res.cosmo_factor == inp.cosmo_factor**2 def test_reduce_divide(self): """ @@ -664,12 +664,12 @@ def test_reduce_divide(self): [[1.0, 2.0], [1.0, 4.0], [1.0, 1.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), ) res = np.divide.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc ** -1), np.array([1.0, 0.5])) + np.testing.assert_allclose(res.to_value(u.kpc**-1), np.array([1.0, 0.5])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor ** -1 + assert res.cosmo_factor == inp.cosmo_factor**-1 def test_reduce_other(self): """ @@ -679,7 +679,7 @@ def test_reduce_other(self): [[1.0, 2.0], [1.0, 2.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = np.add.reduce(inp, axis=0) np.testing.assert_allclose(res.to_value(u.kpc), np.array([2.0, 4.0])) @@ -695,7 +695,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res1, res2 = np.modf(inp) assert res1.to_value(u.kpc) == 0.5 @@ -709,7 +709,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res1, res2 = np.frexp(inp) assert res1 == 0.625 @@ -727,7 +727,7 @@ def test_multi_output_with_out_arg(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) out1 = cosmo_array([np.nan], u.dimensionless, comoving=True) out2 = cosmo_array([np.nan], u.dimensionless, comoving=True) @@ -748,7 +748,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = 0 res = inp1 > inp2 @@ -757,7 +757,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = 0.5 with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -767,20 +767,20 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( [0, 0, 0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) assert (inp1 > inp2).all() inp1 = cosmo_array( [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = np.ones(3) * u.kpc with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -789,7 +789,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = np.zeros(3) * u.kpc assert (inp1 > inp2).all() @@ -797,13 +797,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( 1, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp1 == inp2 assert res.all() @@ -811,13 +811,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) inp2 = cosmo_array( 0, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), ) res = inp1 > inp2 assert res.all() diff --git a/tests/test_data.py b/tests/test_data.py index 44a69f1b..25de5eaa 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -165,7 +165,7 @@ def test_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), ) start_offset = offsets @@ -212,7 +212,7 @@ def test_dithered_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), ) offsets = mask_region.offsets["dark_matter"] counts = mask_region.counts["dark_matter"] @@ -256,7 +256,7 @@ def test_reading_select_region_metadata(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.2, boxsize * 0.8]).T @@ -308,7 +308,7 @@ def test_reading_select_region_metadata_not_spatial_only(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.26, boxsize * 0.74]).T diff --git a/tests/test_extraparts.py b/tests/test_extraparts.py index 796d4cc3..bfb99e04 100644 --- a/tests/test_extraparts.py +++ b/tests/test_extraparts.py @@ -1,6 +1,7 @@ """ Test for extra particle types. """ + from swiftsimio import load, metadata from swiftsimio import Writer from swiftsimio.units import cosmo_units @@ -33,11 +34,11 @@ def generate_units(m, l, t, I, T): "particle_ids": None, "velocities": l / t, "potential": l * l / (t * t), - "density": m / (l ** 3), - "entropy": m * l ** 2 / (t ** 2 * T), + "density": m / (l**3), + "entropy": m * l**2 / (t**2 * T), "internal_energy": (l / t) ** 2, "smoothing_length": l, - "pressure": m / (l * t ** 2), + "pressure": m / (l * t**2), "diffusion": None, "sfr": m / t, "temperature": T, diff --git a/tests/test_read_ic.py b/tests/test_read_ic.py index d1f39f5c..112d4ffb 100644 --- a/tests/test_read_ic.py +++ b/tests/test_read_ic.py @@ -34,7 +34,7 @@ def simple_snapshot_data(): x = Writer(cosmo_units, boxsize) # 32^3 particles. - n_p = 32 ** 3 + n_p = 32**3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) diff --git a/tests/test_soap.py b/tests/test_soap.py index 1d59768c..3469ebf8 100644 --- a/tests/test_soap.py +++ b/tests/test_soap.py @@ -47,13 +47,13 @@ def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): 1e5, "Msun", comoving=True, - cosmo_factor=cosmo_factor(a ** 0, this_mask.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), ) upper = cosmo_quantity( 1e13, "Msun", comoving=True, - cosmo_factor=cosmo_factor(a ** 0, this_mask.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), ) this_mask.constrain_mask( "spherical_overdensity_200_mean", "total_mass", lower, upper diff --git a/tests/test_vistools.py b/tests/test_vistools.py index e4a7fae6..83da6ca4 100644 --- a/tests/test_vistools.py +++ b/tests/test_vistools.py @@ -35,7 +35,7 @@ def vertical_func(x): return abs(1.0 - 2.0 * x) def horizontal_func(y): - return y ** 2 + return y**2 raster_at = np.linspace(0, 1, 1024) From 399bec604e91cf08c1570974acd3ae95b1d145c1 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 12 Feb 2025 01:43:47 +0000 Subject: [PATCH 087/125] Run CI black version. --- swiftsimio/_array_functions.py | 6 +- swiftsimio/conversions.py | 5 +- swiftsimio/masks.py | 4 +- .../metadata/metadata/metadata_fields.py | 2 +- swiftsimio/metadata/objects.py | 38 +++-- swiftsimio/metadata/unit/unit_fields.py | 8 +- swiftsimio/objects.py | 2 +- swiftsimio/snapshot_writer.py | 2 +- swiftsimio/visualisation/_vistools.py | 1 - swiftsimio/visualisation/power_spectrum.py | 20 +-- swiftsimio/visualisation/projection.py | 3 +- .../ray_trace_backends/__init__.py | 6 +- .../slice_backends/nearest_neighbours.py | 2 +- .../volume_render_backends/__init__.py | 8 +- tests/basic_test.py | 2 +- tests/test_cosmo_array.py | 78 +++++----- tests/test_cosmo_array_attrs.py | 134 +++++++++--------- tests/test_data.py | 8 +- tests/test_extraparts.py | 6 +- tests/test_read_ic.py | 2 +- tests/test_soap.py | 4 +- tests/test_vistools.py | 2 +- tests/test_visualisation.py | 59 +++----- 23 files changed, 188 insertions(+), 214 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 5ee38806..1557f840 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -729,18 +729,18 @@ def _arctan2_cosmo_factor( f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor( diff --git a/swiftsimio/conversions.py b/swiftsimio/conversions.py index 2e6d65c0..0ff3878f 100644 --- a/swiftsimio/conversions.py +++ b/swiftsimio/conversions.py @@ -98,13 +98,13 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: # expressions taken directly from astropy, since they do no longer # allow access to these attributes (since version 5.1+) critdens_const = (3.0 / (8.0 * np.pi * const.G)).cgs.value - a_B_c2 = (4.0 * const.sigma_sb / const.c**3).cgs.value + a_B_c2 = (4.0 * const.sigma_sb / const.c ** 3).cgs.value # SWIFT provides Omega_r, but we need a consistent Tcmb0 for astropy. # This is an exact inversion of the procedure performed in astropy. critical_density_0 = astropy_units.Quantity( critdens_const * H0.to("1/s").value ** 2, - astropy_units.g / astropy_units.cm**3, + astropy_units.g / astropy_units.cm ** 3, ) Tcmb0 = (Omega_r * critical_density_0.value / a_B_c2) ** (1.0 / 4.0) @@ -143,6 +143,7 @@ def swift_cosmology_to_astropy(cosmo: dict, units) -> Cosmology: m_nu=ap_m_nu, ) + else: def swift_cosmology_to_astropy(cosmo: dict, units) -> dict: diff --git a/swiftsimio/masks.py b/swiftsimio/masks.py index 434fe9d8..189cb1d0 100644 --- a/swiftsimio/masks.py +++ b/swiftsimio/masks.py @@ -237,7 +237,7 @@ def _unpack_cell_metadata(self): centers_handle[:][sort], units=self.units.length, comoving=True, - cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 1, self.metadata.scale_factor), ) # Note that we cannot assume that these are cubic, unfortunately. @@ -245,7 +245,7 @@ def _unpack_cell_metadata(self): metadata_handle.attrs["size"], units=self.units.length, comoving=True, - cosmo_factor=cosmo_factor(a**1, self.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 1, self.metadata.scale_factor), ) return diff --git a/swiftsimio/metadata/metadata/metadata_fields.py b/swiftsimio/metadata/metadata/metadata_fields.py index 75e12a5e..a032e039 100644 --- a/swiftsimio/metadata/metadata/metadata_fields.py +++ b/swiftsimio/metadata/metadata/metadata_fields.py @@ -77,7 +77,7 @@ def generate_cosmo_args_header_unpack_arrays(scale_factor) -> dict: # should not be cosmo_array'd). cosmo_args = { "boxsize": dict( - cosmo_factor=cosmo_factor(a**1, scale_factor=scale_factor), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=scale_factor), comoving=True, # if it's not, then a=1 and it doesn't matter valid_transform=True, ) diff --git a/swiftsimio/metadata/objects.py b/swiftsimio/metadata/objects.py index aba8b5ca..e3731620 100644 --- a/swiftsimio/metadata/objects.py +++ b/swiftsimio/metadata/objects.py @@ -113,14 +113,12 @@ def postprocess_header(self): # items including the scale factor. # These must be unpacked as they are stored as length-1 arrays - header_unpack_float_units = ( - metadata.metadata_fields.generate_units_header_unpack_single_float( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) + header_unpack_float_units = metadata.metadata_fields.generate_units_header_unpack_single_float( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, ) for field, names in metadata.metadata_fields.header_unpack_single_float.items(): try: @@ -166,19 +164,15 @@ def postprocess_header(self): self.scale_factor = 1.0 # These are just read straight in to variables - header_unpack_arrays_units = ( - metadata.metadata_fields.generate_units_header_unpack_arrays( - m=self.units.mass, - l=self.units.length, - t=self.units.time, - I=self.units.current, - T=self.units.temperature, - ) + header_unpack_arrays_units = metadata.metadata_fields.generate_units_header_unpack_arrays( + m=self.units.mass, + l=self.units.length, + t=self.units.time, + I=self.units.current, + T=self.units.temperature, ) - header_unpack_arrays_cosmo_args = ( - metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( - self.scale_factor - ) + header_unpack_arrays_cosmo_args = metadata.metadata_fields.generate_cosmo_args_header_unpack_arrays( + self.scale_factor ) for field, name in metadata.metadata_fields.header_unpack_arrays.items(): @@ -610,7 +604,7 @@ def get_units(unit_attribute): # Need to check if the exponent is 0 manually because of float precision unit_exponent = unit_attribute[f"U_{exponent} exponent"][0] if unit_exponent != 0.0: - units *= unit**unit_exponent + units *= unit ** unit_exponent except KeyError: # Can't load that data! # We should probably warn the user here... @@ -704,7 +698,7 @@ def get_cosmo(dataset): # Can't load, 'graceful' fallback. cosmo_exponent = 0.0 - a_factor_this_dataset = a**cosmo_exponent + a_factor_this_dataset = a ** cosmo_exponent return cosmo_factor(a_factor_this_dataset, current_scale_factor) diff --git a/swiftsimio/metadata/unit/unit_fields.py b/swiftsimio/metadata/unit/unit_fields.py index ed4f6915..ae5756bd 100644 --- a/swiftsimio/metadata/unit/unit_fields.py +++ b/swiftsimio/metadata/unit/unit_fields.py @@ -54,11 +54,11 @@ def generate_units(m, l, t, I, T): } gas = { - "density": m / (l**3), - "entropy": m * l**2 / (t**2 * T), + "density": m / (l ** 3), + "entropy": m * l ** 2 / (t ** 2 * T), "internal_energy": (l / t) ** 2, "smoothing_length": l, - "pressure": m / (l * t**2), + "pressure": m / (l * t ** 2), "diffusion": None, "sfr": m / t, "temperature": T, @@ -79,7 +79,7 @@ def generate_units(m, l, t, I, T): sinks = {**shared} stars = { - "birth_density": m / (l**3), + "birth_density": m / (l ** 3), "birth_time": t, "initial_masses": m, "smoothing_length": l, diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index f60c65a7..0516d64a 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -643,7 +643,7 @@ def __pow__(self, p: float) -> "cosmo_factor": """ if self.expr is None: return cosmo_factor(expr=None, scale_factor=self.scale_factor) - 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: "cosmo_factor") -> bool: """ diff --git a/swiftsimio/snapshot_writer.py b/swiftsimio/snapshot_writer.py index 0d05495f..388ead1b 100644 --- a/swiftsimio/snapshot_writer.py +++ b/swiftsimio/snapshot_writer.py @@ -268,7 +268,7 @@ def get_attributes(self, scale_factor: float) -> dict: # Find the scale factor associated quantities a_exp = a_exponents.get(name, 0) - a_factor = scale_factor**a_exp + a_factor = scale_factor ** a_exp attributes_dict[output_handle] = { "Conversion factor to CGS (not including cosmological corrections)": [ diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py index 1836b2fd..72dbb298 100644 --- a/swiftsimio/visualisation/_vistools.py +++ b/swiftsimio/visualisation/_vistools.py @@ -95,7 +95,6 @@ def _get_rotated_coordinates(data, rotation_matrix, rotation_center): def backends_restore_cosmo_and_units(backend_func): - def wrapper(*args, **kwargs): norm = kwargs.pop("norm") comoving = getattr(kwargs["m"], "comoving", None) diff --git a/swiftsimio/visualisation/power_spectrum.py b/swiftsimio/visualisation/power_spectrum.py index a7fd0680..fb090937 100644 --- a/swiftsimio/visualisation/power_spectrum.py +++ b/swiftsimio/visualisation/power_spectrum.py @@ -198,7 +198,7 @@ def render_to_deposit( """ # Get the positions and masses - folding = 2.0**folding + folding = 2.0 ** folding positions = data.coordinates quantity = getattr(data, project) @@ -245,10 +245,10 @@ def render_to_deposit( units = 1.0 / ( data.metadata.boxsize[0] * data.metadata.boxsize[1] * data.metadata.boxsize[2] ) - units.convert_to_units(1.0 / data.metadata.boxsize.units**3) + units.convert_to_units(1.0 / data.metadata.boxsize.units ** 3) units *= quantity.units - new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor**3) + new_cosmo_factor = quantity.cosmo_factor / (coord_cosmo_factor ** 3) return cosmo_array( deposition, comoving=comoving, cosmo_factor=new_cosmo_factor, units=units @@ -358,7 +358,7 @@ def folded_depositions_to_power_spectrum( np.zeros(number_of_wavenumber_bins), units=box_volume.units, comoving=box_volume.comoving, - cosmo_factor=box_volume.cosmo_factor**-1, + cosmo_factor=box_volume.cosmo_factor ** -1, name="Power spectrum $P(k)$", ) folding_tracker = np.ones(number_of_wavenumber_bins, dtype=float) @@ -396,7 +396,7 @@ def folded_depositions_to_power_spectrum( if folding != final_folding: cutoff_wavenumber = ( - 2.0**folding * np.min(depositions[folding].shape) / np.min(box_size) + 2.0 ** folding * np.min(depositions[folding].shape) / np.min(box_size) ) if cutoff_above_wavenumber_fraction is not None: @@ -421,7 +421,7 @@ def folded_depositions_to_power_spectrum( corrected_wavenumber_centers[prefer_bins] = folded_wavenumber_centers[ prefer_bins ].to(corrected_wavenumber_centers.units) - folding_tracker[prefer_bins] = 2.0**folding + folding_tracker[prefer_bins] = 2.0 ** folding contributed_counts[prefer_bins] = folded_counts[prefer_bins] elif transition == "average": @@ -454,7 +454,7 @@ def folded_depositions_to_power_spectrum( # For debugging, we calculate an effective fold number. folding_tracker[use_bins] = ( - (folding_tracker * existing_weight + (2.0**folding) * new_weight) + (folding_tracker * existing_weight + (2.0 ** folding) * new_weight) / transition_norm )[use_bins] @@ -535,7 +535,7 @@ def deposition_to_power_spectrum( deposition.shape == cross_deposition.shape ), "Depositions must have the same shape" - folding = 2.0**folding + folding = 2.0 ** folding box_size_folded = box_size[0] / folding npix = deposition.shape[0] @@ -557,7 +557,7 @@ def deposition_to_power_spectrum( else: conj_fft = fft.conj() - fourier_amplitudes = (fft * conj_fft).real * box_size_folded**3 + fourier_amplitudes = (fft * conj_fft).real * box_size_folded ** 3 # Calculate k-value spacing (centered FFT) dk = 2 * np.pi / (box_size_folded) @@ -589,7 +589,7 @@ def deposition_to_power_spectrum( divisor[zero_mask] = 1 # Correct for folding - binned_amplitudes *= folding**3 + binned_amplitudes *= folding ** 3 # Correct units and names wavenumbers = binned_wavenumbers / divisor diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 707a6dd3..aced21e6 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -108,8 +108,7 @@ def project_pixel_grid( mask = mask if mask is not None else np.s_[...] if not region_info["z_slice_included"]: mask = np.logical_and( - mask, - np.logical_and(z <= region_info["z_max"], z >= region_info["z_min"]), + mask, np.logical_and(z <= region_info["z_max"], z >= region_info["z_min"]) ).astype(bool) kwargs = dict( diff --git a/swiftsimio/visualisation/ray_trace_backends/__init__.py b/swiftsimio/visualisation/ray_trace_backends/__init__.py index 9c25b3db..be85d05c 100644 --- a/swiftsimio/visualisation/ray_trace_backends/__init__.py +++ b/swiftsimio/visualisation/ray_trace_backends/__init__.py @@ -9,10 +9,8 @@ core_panels_parallel, ) -backends = { - "core_panels": backends_restore_cosmo_and_units(core_panels), -} +backends = {"core_panels": backends_restore_cosmo_and_units(core_panels)} backends_parallel = { - "core_panels": backends_restore_cosmo_and_units(core_panels_parallel), + "core_panels": backends_restore_cosmo_and_units(core_panels_parallel) } diff --git a/swiftsimio/visualisation/slice_backends/nearest_neighbours.py b/swiftsimio/visualisation/slice_backends/nearest_neighbours.py index bb5e9504..37656bb1 100644 --- a/swiftsimio/visualisation/slice_backends/nearest_neighbours.py +++ b/swiftsimio/visualisation/slice_backends/nearest_neighbours.py @@ -130,7 +130,7 @@ def slice_scatter( tree = build_tree(x, y, z, box_x, box_y, box_z) _, i = tree.query(pixel_coordinates, workers=workers) - values = m / h**3 + values = m / h ** 3 return values[i].reshape(xres, yres) diff --git a/swiftsimio/visualisation/volume_render_backends/__init__.py b/swiftsimio/visualisation/volume_render_backends/__init__.py index 0752a1f7..485cf175 100644 --- a/swiftsimio/visualisation/volume_render_backends/__init__.py +++ b/swiftsimio/visualisation/volume_render_backends/__init__.py @@ -9,10 +9,6 @@ scatter_parallel, ) -backends = { - "scatter": backends_restore_cosmo_and_units(scatter), -} +backends = {"scatter": backends_restore_cosmo_and_units(scatter)} -backends_parallel = { - "scatter": backends_restore_cosmo_and_units(scatter_parallel), -} +backends_parallel = {"scatter": backends_restore_cosmo_and_units(scatter_parallel)} diff --git a/tests/basic_test.py b/tests/basic_test.py index 24fb9295..adb4e7fc 100644 --- a/tests/basic_test.py +++ b/tests/basic_test.py @@ -23,7 +23,7 @@ def test_write(): x = Writer("galactic", boxsize) # 32^3 particles. - n_p = 32**3 + n_p = 32 ** 3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 65fa3081..ba112464 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -28,14 +28,16 @@ def ca(x, unit=u.Mpc): """ Helper for our tests: turn an array into a cosmo_array. """ - 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 cq(x, unit=u.Mpc): """ Helper for our tests: turn a scalar into a cosmo_quantity. """ - return cosmo_quantity(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) + return cosmo_quantity( + x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5) + ) def arg_to_ua(arg): @@ -113,7 +115,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") @@ -127,7 +129,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") @@ -140,7 +142,7 @@ 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") @@ -156,7 +158,7 @@ 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") @@ -176,14 +178,14 @@ def test_init_from_list_of_cosmo_arrays(self): [1], units=u.Mpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1), + 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 @@ -571,7 +573,7 @@ def test_explicitly_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), ), ), ), @@ -582,13 +584,13 @@ def test_explicitly_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), ), ), ), @@ -600,19 +602,19 @@ def test_explicitly_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), ), ], ), @@ -624,7 +626,7 @@ def test_explicitly_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]), ), @@ -648,19 +650,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] @@ -707,9 +709,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): @@ -717,9 +719,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): @@ -727,9 +729,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 = { @@ -741,9 +743,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), ] ), ): @@ -800,12 +802,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 def test_scalar_return_func(self): @@ -817,7 +819,7 @@ def test_scalar_return_func(self): np.arange(3), u.m, comoving=False, - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), valid_transform=True, ) res = np.min(ca) @@ -832,12 +834,12 @@ 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 @@ -853,7 +855,7 @@ def test_copy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) copy_arr = copy(arr) @@ -869,7 +871,7 @@ def test_deepcopy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) copy_arr = deepcopy(arr) @@ -885,7 +887,7 @@ def test_to_cgs(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a**1, 1), + cosmo_factor=cosmo_factor(a ** 1, 1), comoving=False, ) cgs_arr = arr.in_cgs() diff --git a/tests/test_cosmo_array_attrs.py b/tests/test_cosmo_array_attrs.py index 9b38cfbc..d4b1562c 100644 --- a/tests/test_cosmo_array_attrs.py +++ b/tests/test_cosmo_array_attrs.py @@ -143,7 +143,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 0.5), + cosmo_factor=cosmo_factor(a ** 1, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -152,7 +152,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 0.5), + cosmo_factor=cosmo_factor(a ** 1, 0.5), comoving=False, ) assert not arr.compatible_with_comoving() @@ -161,7 +161,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**0, 0.5), + cosmo_factor=cosmo_factor(a ** 0, 0.5), comoving=True, ) assert arr.compatible_with_comoving() @@ -170,7 +170,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**0, 0.5), + cosmo_factor=cosmo_factor(a ** 0, 0.5), comoving=False, ) assert arr.compatible_with_comoving() @@ -179,7 +179,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=True, ) assert arr.compatible_with_comoving() @@ -188,7 +188,7 @@ def test_compatibility(self): arr = cosmo_array( np.ones((10, 10)), units="Mpc", - cosmo_factor=cosmo_factor(a**1, 1.0), + cosmo_factor=cosmo_factor(a ** 1, 1.0), comoving=False, ) assert arr.compatible_with_comoving() @@ -232,7 +232,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.ones_like(inp) assert res.to_value(u.kpc) == 1 @@ -265,7 +265,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises( @@ -278,7 +278,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -289,13 +289,13 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -306,7 +306,7 @@ def test_preserving_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp + inp assert res.to_value(u.kpc) == 4 @@ -320,7 +320,7 @@ def test_multiplying_ufunc(self): # no cosmo_factors inp = cosmo_array([2], u.kpc, comoving=False) res = inp * inp - assert res.to_value(u.kpc**2) == 4 + assert res.to_value(u.kpc ** 2) == 4 assert res.comoving is False assert res.cosmo_factor == cosmo_factor(None, None) # one is not cosmo_array @@ -342,7 +342,7 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): @@ -353,7 +353,7 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.raises(InvalidScaleFactor, match="Attempting to multiply"): inp1 * inp2 @@ -362,12 +362,12 @@ def test_multiplying_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp * inp - assert res.to_value(u.kpc**2) == 4 + assert res.to_value(u.kpc ** 2) == 4 assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**2 + assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_dividing_ufunc(self): """ @@ -377,12 +377,12 @@ def test_dividing_ufunc(self): [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp / inp assert res.to_value(u.dimensionless) == 1 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**0 + assert res.cosmo_factor == inp.cosmo_factor ** 0 def test_return_without_ufunc(self): """ @@ -393,7 +393,7 @@ def test_return_without_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.logical_not(inp) assert res == np.logical_not(1) @@ -403,7 +403,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.logaddexp(inp, inp) assert res == np.logaddexp(2, 2) @@ -413,13 +413,13 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -431,7 +431,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -442,7 +442,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array([2], u.kpc, comoving=False) with pytest.raises( @@ -457,7 +457,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) with pytest.warns(RuntimeWarning, match="Mixing arguments"): res = np.logaddexp(inp1, inp2) @@ -468,7 +468,7 @@ def test_return_without_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = u.unyt_array([2], u.kpc) with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -484,12 +484,12 @@ def test_sqrt_ufunc(self): [4], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.sqrt(inp) - assert res.to_value(u.kpc**0.5) == 2 # also ensures units ok + assert res.to_value(u.kpc ** 0.5) == 2 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**0.5 + assert res.cosmo_factor == inp.cosmo_factor ** 0.5 def test_square_ufunc(self): """ @@ -499,12 +499,12 @@ def test_square_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.square(inp) - assert res.to_value(u.kpc**2) == 4 # also ensures units ok + assert res.to_value(u.kpc ** 2) == 4 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**2 + assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_cbrt_ufunc(self): """ @@ -514,7 +514,7 @@ def test_cbrt_ufunc(self): [8], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.cbrt(inp) assert res.to_value(u.kpc ** (1.0 / 3.0)) == 2 # also ensures units ok @@ -529,12 +529,12 @@ def test_reciprocal_ufunc(self): [2.0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.reciprocal(inp) - assert res.to_value(u.kpc**-1) == 0.5 # also ensures units ok + assert res.to_value(u.kpc ** -1) == 0.5 # also ensures units ok assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**-1 + assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_passthrough_ufunc(self): """ @@ -545,7 +545,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.negative(inp) assert res.to_value(u.kpc) == -2 @@ -556,7 +556,7 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.copysign(inp, inp) assert res.to_value(u.kpc) == inp.to_value(u.kpc) @@ -567,13 +567,13 @@ def test_passthrough_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) with pytest.raises( ValueError, match="Arguments have cosmo_factors that differ" @@ -588,7 +588,7 @@ def test_arctan2_ufunc(self): [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.arctan2(inp, inp) assert res.to_value(u.dimensionless) == np.arctan2(2, 2) @@ -603,13 +603,13 @@ def test_comparison_ufunc(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [2], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp1 < inp2 assert res.all() @@ -623,7 +623,7 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) out = cosmo_array([np.nan], u.dimensionless, comoving=True) np.abs(inp, out=out) @@ -634,7 +634,7 @@ def test_out_arg(self): [1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) # make sure we can also pass a non-cosmo type for out without crashing out = np.array([np.nan]) @@ -649,12 +649,12 @@ def test_reduce_multiply(self): [[1, 2], [3, 4]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) res = np.multiply.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc**2), np.array([3.0, 8.0])) + np.testing.assert_allclose(res.to_value(u.kpc ** 2), np.array([3.0, 8.0])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**2 + assert res.cosmo_factor == inp.cosmo_factor ** 2 def test_reduce_divide(self): """ @@ -664,12 +664,12 @@ def test_reduce_divide(self): [[1.0, 2.0], [1.0, 4.0], [1.0, 1.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=0.5), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=0.5), ) res = np.divide.reduce(inp, axis=0) - np.testing.assert_allclose(res.to_value(u.kpc**-1), np.array([1.0, 0.5])) + np.testing.assert_allclose(res.to_value(u.kpc ** -1), np.array([1.0, 0.5])) assert res.comoving is False - assert res.cosmo_factor == inp.cosmo_factor**-1 + assert res.cosmo_factor == inp.cosmo_factor ** -1 def test_reduce_other(self): """ @@ -679,7 +679,7 @@ def test_reduce_other(self): [[1.0, 2.0], [1.0, 2.0]], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = np.add.reduce(inp, axis=0) np.testing.assert_allclose(res.to_value(u.kpc), np.array([2.0, 4.0])) @@ -695,7 +695,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res1, res2 = np.modf(inp) assert res1.to_value(u.kpc) == 0.5 @@ -709,7 +709,7 @@ def test_multi_output(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res1, res2 = np.frexp(inp) assert res1 == 0.625 @@ -727,7 +727,7 @@ def test_multi_output_with_out_arg(self): [2.5], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) out1 = cosmo_array([np.nan], u.dimensionless, comoving=True) out2 = cosmo_array([np.nan], u.dimensionless, comoving=True) @@ -748,7 +748,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = 0 res = inp1 > inp2 @@ -757,7 +757,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = 0.5 with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -767,20 +767,20 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( [0, 0, 0], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) assert (inp1 > inp2).all() inp1 = cosmo_array( [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = np.ones(3) * u.kpc with pytest.warns(RuntimeWarning, match="Mixing arguments"): @@ -789,7 +789,7 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = np.zeros(3) * u.kpc assert (inp1 > inp2).all() @@ -797,13 +797,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( 1, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp1 == inp2 assert res.all() @@ -811,13 +811,13 @@ def test_comparison_with_zero(self): [1, 1, 1], u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) inp2 = cosmo_array( 0, u.kpc, comoving=False, - cosmo_factor=cosmo_factor(a**1, scale_factor=1.0), + cosmo_factor=cosmo_factor(a ** 1, scale_factor=1.0), ) res = inp1 > inp2 assert res.all() diff --git a/tests/test_data.py b/tests/test_data.py index 25de5eaa..44a69f1b 100644 --- a/tests/test_data.py +++ b/tests/test_data.py @@ -165,7 +165,7 @@ def test_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), ) start_offset = offsets @@ -212,7 +212,7 @@ def test_dithered_cell_metadata_is_valid(filename): boxsize = cosmo_array( boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, mask_region.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, mask_region.metadata.a), ) offsets = mask_region.offsets["dark_matter"] counts = mask_region.counts["dark_matter"] @@ -256,7 +256,7 @@ def test_reading_select_region_metadata(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.2, boxsize * 0.8]).T @@ -308,7 +308,7 @@ def test_reading_select_region_metadata_not_spatial_only(filename): boxsize = cosmo_array( full_data.metadata.boxsize, comoving=True, - cosmo_factor=cosmo_factor(a**1, full_data.metadata.a), + cosmo_factor=cosmo_factor(a ** 1, full_data.metadata.a), ) restrict = cosmo_array([boxsize * 0.26, boxsize * 0.74]).T diff --git a/tests/test_extraparts.py b/tests/test_extraparts.py index bfb99e04..1a2f9db0 100644 --- a/tests/test_extraparts.py +++ b/tests/test_extraparts.py @@ -34,11 +34,11 @@ def generate_units(m, l, t, I, T): "particle_ids": None, "velocities": l / t, "potential": l * l / (t * t), - "density": m / (l**3), - "entropy": m * l**2 / (t**2 * T), + "density": m / (l ** 3), + "entropy": m * l ** 2 / (t ** 2 * T), "internal_energy": (l / t) ** 2, "smoothing_length": l, - "pressure": m / (l * t**2), + "pressure": m / (l * t ** 2), "diffusion": None, "sfr": m / t, "temperature": T, diff --git a/tests/test_read_ic.py b/tests/test_read_ic.py index 112d4ffb..d1f39f5c 100644 --- a/tests/test_read_ic.py +++ b/tests/test_read_ic.py @@ -34,7 +34,7 @@ def simple_snapshot_data(): x = Writer(cosmo_units, boxsize) # 32^3 particles. - n_p = 32**3 + n_p = 32 ** 3 # Randomly spaced coordinates from 0, 100 Mpc in each direction x.gas.coordinates = np.random.rand(n_p, 3) * (100 * unyt.Mpc) diff --git a/tests/test_soap.py b/tests/test_soap.py index 3469ebf8..1d59768c 100644 --- a/tests/test_soap.py +++ b/tests/test_soap.py @@ -47,13 +47,13 @@ def test_soap_can_mask_spatial_and_non_spatial_actually_use(filename): 1e5, "Msun", comoving=True, - cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 0, this_mask.metadata.scale_factor), ) upper = cosmo_quantity( 1e13, "Msun", comoving=True, - cosmo_factor=cosmo_factor(a**0, this_mask.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** 0, this_mask.metadata.scale_factor), ) this_mask.constrain_mask( "spherical_overdensity_200_mean", "total_mass", lower, upper diff --git a/tests/test_vistools.py b/tests/test_vistools.py index 83da6ca4..e4a7fae6 100644 --- a/tests/test_vistools.py +++ b/tests/test_vistools.py @@ -35,7 +35,7 @@ def vertical_func(x): return abs(1.0 - 2.0 * x) def horizontal_func(y): - return y**2 + return y ** 2 raster_at = np.linspace(0, 1, 1024) diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index c1bdffd7..429e7730 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -85,16 +85,9 @@ def test_scatter_mass_conservation(): for resolution in resolutions: scatter = projection_backends["fast"] image = scatter( - x=x, - y=y, - m=m, - h=h, - res=resolution, - box_x=1.0, - box_y=1.0, - norm=1.0, + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, norm=1.0 ) - mass_in_image = image.sum() / (resolution**2) + mass_in_image = image.sum() / (resolution ** 2) # Check mass conservation to 5% assert np.isclose(mass_in_image.view(np.ndarray), total_mass, 0.05) @@ -414,11 +407,7 @@ def test_comoving_versus_physical(filename): 0.0 * boxsize[2], 0.2 * boxsize[2], ] - for func, aexp in [ - (project_gas, -2.0), - (slice_gas, -3.0), - (render_gas, -3.0), - ]: + for func, aexp in [(project_gas, -2.0), (slice_gas, -3.0), (render_gas, -3.0)]: # normal case: everything comoving data = load(filename, mask=m) # we force the default (project="masses") to check the cosmo_factor @@ -459,7 +448,7 @@ def test_comoving_versus_physical(filename): ): img = func(data, resolution=64, project="masses", region=region) assert data.gas.masses.comoving and img.comoving - assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 + assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 # densities are physical, make sure this works with physical coordinates and # smoothing lengths img = func(data, resolution=64, project="densities", region=region) @@ -490,11 +479,7 @@ def test_nongas_smoothing_lengths(filename): data.dark_matter.smoothing_length = generate_smoothing_lengths( data.dark_matter.coordinates, data.metadata.boxsize, kernel_gamma=1.8 ) - project_pixel_grid( - data.dark_matter, - resolution=256, - project="masses", - ) + project_pixel_grid(data.dark_matter, resolution=256, project="masses") assert isinstance(data.dark_matter.smoothing_length, cosmo_array) # We should also be able to use a unyt_array (rather than cosmo_array) as input, @@ -529,14 +514,14 @@ def test_panel_rendering(filename): plt.imsave( "panels_added.png", - plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1))), + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), ) projected = project_gas(data, res, "masses", backend="renormalised") plt.imsave( "projected.png", - plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)), + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), ) fullstack = np.zeros((res, res)) @@ -590,7 +575,7 @@ def test_periodic_boundary_wrapping(): res=pixel_resolution, box_x=boxsize, box_y=boxsize, - norm=boxsize**2, + norm=boxsize ** 2, ) image2 = projection_backends[backend]( x=coordinates_non_periodic[:, 0], @@ -600,7 +585,7 @@ def test_periodic_boundary_wrapping(): res=pixel_resolution, box_x=0.0, box_y=0.0, - norm=boxsize**2, + norm=boxsize ** 2, ) assert (image1 == image2).all() except CudaSupportError: @@ -623,7 +608,7 @@ def test_periodic_boundary_wrapping(): box_x=boxsize, box_y=boxsize, box_z=boxsize, - norm=boxsize**3, + norm=boxsize ** 3, ) image2 = slice_backends[backend]( x=coordinates_non_periodic[:, 0], @@ -637,7 +622,7 @@ def test_periodic_boundary_wrapping(): box_x=0.0, box_y=0.0, box_z=0.0, - norm=boxsize**3, + norm=boxsize ** 3, ) assert (image1 == image2).all() @@ -654,7 +639,7 @@ def test_periodic_boundary_wrapping(): box_x=boxsize, box_y=boxsize, box_z=boxsize, - norm=boxsize**3, + norm=boxsize ** 3, ) image2 = scatter( x=coordinates_non_periodic[:, 0], @@ -666,7 +651,7 @@ def test_periodic_boundary_wrapping(): box_x=0.0, box_y=0.0, box_z=0.0, - norm=boxsize**3, + norm=boxsize ** 3, ) assert (image1 == image2).all() @@ -701,11 +686,11 @@ def test_volume_render_and_unfolded_deposit(): box_x=boxsize, box_y=boxsize, box_z=boxsize, - norm=boxsize**3, + norm=boxsize ** 3, ) # need to divide out the box volume for the deposition - assert np.allclose(deposition / boxsize**3, volume.view(np.ndarray)) + assert np.allclose(deposition / boxsize ** 3, volume.view(np.ndarray)) def test_folding_deposit(): @@ -744,8 +729,8 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix**3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix**3).to("Msun / kpc**3").v + mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v + mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v mean_density_calculated = ( (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) .to("Msun / kpc**3") @@ -770,15 +755,15 @@ def test_dark_matter_power_spectrum(filename, save=False): min_k = cosmo_quantity( 1e-2, - unyt.Mpc**-1, + unyt.Mpc ** -1, comoving=True, - cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), ) max_k = cosmo_quantity( 1e2, - unyt.Mpc**-1, + unyt.Mpc ** -1, comoving=True, - cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), ) bins = np.geomspace(min_k, max_k, 32) @@ -815,7 +800,7 @@ def test_dark_matter_power_spectrum(filename, save=False): folds[folding] = deposition - folding_output[2**folding] = (k, power_spectrum, scatter) + folding_output[2 ** folding] = (k, power_spectrum, scatter) # Now try doing them all together at once. From c4b5795cf8a3a54d050cf9c15411904e6363ee63 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 12 Feb 2025 16:37:45 +0000 Subject: [PATCH 088/125] Move wrapper from backends to be called in frontend functions. --- swiftsimio/visualisation/_vistools.py | 2 +- swiftsimio/visualisation/projection.py | 8 ++--- .../projection_backends/__init__.py | 30 +++++++++---------- swiftsimio/visualisation/ray_trace.py | 3 +- .../ray_trace_backends/__init__.py | 8 ++--- swiftsimio/visualisation/slice.py | 8 ++--- .../visualisation/slice_backends/__init__.py | 11 ++----- swiftsimio/visualisation/volume_render.py | 10 +++---- .../volume_render_backends/__init__.py | 6 ++-- 9 files changed, 34 insertions(+), 52 deletions(-) diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py index 72dbb298..cdbf4fb2 100644 --- a/swiftsimio/visualisation/_vistools.py +++ b/swiftsimio/visualisation/_vistools.py @@ -94,7 +94,7 @@ def _get_rotated_coordinates(data, rotation_matrix, rotation_center): return x, y, z -def backends_restore_cosmo_and_units(backend_func): +def backend_restore_cosmo_and_units(backend_func): def wrapper(*args, **kwargs): norm = kwargs.pop("norm") comoving = getattr(kwargs["m"], "comoving", None) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index aced21e6..f5f60346 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -13,6 +13,7 @@ _get_projection_field, _get_region_info, _get_rotated_coordinates, + backend_restore_cosmo_and_units, ) @@ -121,11 +122,8 @@ def project_pixel_grid( box_y=region_info["periodic_box_y"], norm=(region_info["x_range"] * region_info["y_range"]), ) - image = ( - backends_parallel[backend](**kwargs) - if parallel - else backends[backend](**kwargs) - ) + backend_func = (backends_parallel if parallel else backends)[backend] + image = backend_restore_cosmo_and_units(backend_func)(**kwargs) # determine the effective number of pixels for each dimension xres = int( diff --git a/swiftsimio/visualisation/projection_backends/__init__.py b/swiftsimio/visualisation/projection_backends/__init__.py index 14a620ef..3715abf7 100644 --- a/swiftsimio/visualisation/projection_backends/__init__.py +++ b/swiftsimio/visualisation/projection_backends/__init__.py @@ -6,8 +6,6 @@ being a developer-only indended feature. """ -from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units - from swiftsimio.visualisation.projection_backends.fast import scatter as fast from swiftsimio.visualisation.projection_backends.fast import ( scatter_parallel as fast_parallel, @@ -50,21 +48,21 @@ ) backends = { - "histogram": backends_restore_cosmo_and_units(histogram), - "fast": backends_restore_cosmo_and_units(fast), - "renormalised": backends_restore_cosmo_and_units(renormalised), - "subsampled": backends_restore_cosmo_and_units(subsampled), - "subsampled_extreme": backends_restore_cosmo_and_units(subsampled_extreme), - "reference": backends_restore_cosmo_and_units(reference), - "gpu": backends_restore_cosmo_and_units(gpu), + "histogram": histogram, + "fast": fast, + "renormalised": renormalised, + "subsampled": subsampled, + "subsampled_extreme": subsampled_extreme, + "reference": reference, + "gpu": gpu, } backends_parallel = { - "histogram": backends_restore_cosmo_and_units(histogram_parallel), - "fast": backends_restore_cosmo_and_units(fast_parallel), - "renormalised": backends_restore_cosmo_and_units(renormalised_parallel), - "subsampled": backends_restore_cosmo_and_units(subsampled_parallel), - "subsampled_extreme": backends_restore_cosmo_and_units(subsampled_extreme_parallel), - "reference": backends_restore_cosmo_and_units(reference_parallel), - "gpu": backends_restore_cosmo_and_units(gpu_parallel), + "histogram": histogram_parallel, + "fast": fast_parallel, + "renormalised": renormalised_parallel, + "subsampled": subsampled_parallel, + "subsampled_extreme": subsampled_extreme_parallel, + "reference": reference_parallel, + "gpu": gpu_parallel, } diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index cb101543..f5a4274f 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -14,6 +14,7 @@ _get_projection_field, _get_region_info, _get_rotated_coordinates, + backend_restore_cosmo_and_units, ) @@ -36,7 +37,7 @@ def panel_pixel_grid( # There's a parallel version of core_panels but it seems # that it's never used anywhere. - return backends["core_panels"]( + return backend_restore_cosmo_and_units(backends["core_panels"])( x=x[mask] / region_info["max_range"], y=y[mask] / region_info["max_range"], z=z[mask], diff --git a/swiftsimio/visualisation/ray_trace_backends/__init__.py b/swiftsimio/visualisation/ray_trace_backends/__init__.py index be85d05c..d14af439 100644 --- a/swiftsimio/visualisation/ray_trace_backends/__init__.py +++ b/swiftsimio/visualisation/ray_trace_backends/__init__.py @@ -2,15 +2,11 @@ Backends for ray tracing """ -from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units - from swiftsimio.visualisation.ray_trace_backends.core_panels import ( core_panels, core_panels_parallel, ) -backends = {"core_panels": backends_restore_cosmo_and_units(core_panels)} +backends = {"core_panels": core_panels} -backends_parallel = { - "core_panels": backends_restore_cosmo_and_units(core_panels_parallel) -} +backends_parallel = {"core_panels": core_panels_parallel} diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index fe57b260..60110658 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -11,6 +11,7 @@ _get_projection_field, _get_region_info, _get_rotated_coordinates, + backend_restore_cosmo_and_units, ) @@ -122,10 +123,7 @@ def slice_gas( box_z=region_info["periodic_box_z"], norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]), ) - - if parallel: - image = backends_parallel[backend](**kwargs) - else: - image = backends[backend](**kwargs) + backend_func = (backends_parallel if parallel else backends)[backend] + image = backend_restore_cosmo_and_units(backend_func)(**kwargs) return image diff --git a/swiftsimio/visualisation/slice_backends/__init__.py b/swiftsimio/visualisation/slice_backends/__init__.py index 099d8ada..0fd2e2ba 100644 --- a/swiftsimio/visualisation/slice_backends/__init__.py +++ b/swiftsimio/visualisation/slice_backends/__init__.py @@ -2,8 +2,6 @@ Backends for density slicing. """ -from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units - from swiftsimio.visualisation.slice_backends.sph import ( slice_scatter as sph, slice_scatter_parallel as sph_parallel, @@ -13,12 +11,9 @@ slice_scatter_parallel as nearest_neighbours_parallel, ) -backends = { - "sph": backends_restore_cosmo_and_units(sph), - "nearest_neighbours": backends_restore_cosmo_and_units(nearest_neighbours), -} +backends = {"sph": sph, "nearest_neighbours": nearest_neighbours} backends_parallel = { - "sph": backends_restore_cosmo_and_units(sph_parallel), - "nearest_neighbours": backends_restore_cosmo_and_units(nearest_neighbours_parallel), + "sph": sph_parallel, + "nearest_neighbours": nearest_neighbours_parallel, } diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index b0dda8dc..ceb9b593 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -17,6 +17,7 @@ _get_projection_field, _get_region_info, _get_rotated_coordinates, + backend_restore_cosmo_and_units, ) @@ -93,7 +94,7 @@ def render_gas( hsml = backends_get_hsml["sph"](data) x, y, z = _get_rotated_coordinates(data, rotation_matrix, rotation_center) - arguments = dict( + kwargs = dict( x=(x - region_info["x_min"]) / region_info["x_range"], y=(y - region_info["y_min"]) / region_info["y_range"], z=(z - region_info["z_min"]) / region_info["z_range"], @@ -105,11 +106,8 @@ def render_gas( box_z=region_info["periodic_box_z"], norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]), ) - image = ( - backends_parallel["scatter"](**arguments) - if parallel - else backends["scatter"](**arguments) - ) + backend_func = (backends_parallel if parallel else backends)["scatter"] + image = backend_restore_cosmo_and_units(backend_func)(**kwargs) return image diff --git a/swiftsimio/visualisation/volume_render_backends/__init__.py b/swiftsimio/visualisation/volume_render_backends/__init__.py index 485cf175..f8108b5f 100644 --- a/swiftsimio/visualisation/volume_render_backends/__init__.py +++ b/swiftsimio/visualisation/volume_render_backends/__init__.py @@ -2,13 +2,11 @@ Backends for volume rendering """ -from swiftsimio.visualisation._vistools import backends_restore_cosmo_and_units - from swiftsimio.visualisation.volume_render_backends.scatter import ( scatter, scatter_parallel, ) -backends = {"scatter": backends_restore_cosmo_and_units(scatter)} +backends = {"scatter": scatter} -backends_parallel = {"scatter": backends_restore_cosmo_and_units(scatter_parallel)} +backends_parallel = {"scatter": scatter_parallel} From 06cc3ca8e7c5b56f8b1a6321a8026d2e8654d4d6 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Wed, 12 Feb 2025 16:42:44 +0000 Subject: [PATCH 089/125] Move norm argument onto wrapper function. --- swiftsimio/visualisation/_vistools.py | 3 +-- swiftsimio/visualisation/projection.py | 4 ++-- swiftsimio/visualisation/ray_trace.py | 4 ++-- swiftsimio/visualisation/slice.py | 4 ++-- swiftsimio/visualisation/volume_render.py | 4 ++-- 5 files changed, 9 insertions(+), 10 deletions(-) diff --git a/swiftsimio/visualisation/_vistools.py b/swiftsimio/visualisation/_vistools.py index cdbf4fb2..b3d8479f 100644 --- a/swiftsimio/visualisation/_vistools.py +++ b/swiftsimio/visualisation/_vistools.py @@ -94,9 +94,8 @@ def _get_rotated_coordinates(data, rotation_matrix, rotation_center): return x, y, z -def backend_restore_cosmo_and_units(backend_func): +def backend_restore_cosmo_and_units(backend_func, norm=1.0): def wrapper(*args, **kwargs): - norm = kwargs.pop("norm") comoving = getattr(kwargs["m"], "comoving", None) if comoving is True: if kwargs["x"].comoving is False or kwargs["y"].comoving is False: diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index f5f60346..24da3638 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -120,10 +120,10 @@ def project_pixel_grid( res=resolution, box_x=region_info["periodic_box_x"], box_y=region_info["periodic_box_y"], - norm=(region_info["x_range"] * region_info["y_range"]), ) + norm = region_info["x_range"] * region_info["y_range"] backend_func = (backends_parallel if parallel else backends)[backend] - image = backend_restore_cosmo_and_units(backend_func)(**kwargs) + image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs) # determine the effective number of pixels for each dimension xres = int( diff --git a/swiftsimio/visualisation/ray_trace.py b/swiftsimio/visualisation/ray_trace.py index f5a4274f..55bc2a1c 100644 --- a/swiftsimio/visualisation/ray_trace.py +++ b/swiftsimio/visualisation/ray_trace.py @@ -37,7 +37,8 @@ def panel_pixel_grid( # There's a parallel version of core_panels but it seems # that it's never used anywhere. - return backend_restore_cosmo_and_units(backends["core_panels"])( + norm = region_info["x_range"] * region_info["y_range"] + return backend_restore_cosmo_and_units(backends["core_panels"], norm=norm)( x=x[mask] / region_info["max_range"], y=y[mask] / region_info["max_range"], z=z[mask], @@ -47,7 +48,6 @@ def panel_pixel_grid( panels=panels, min_z=region_info["z_min"], max_z=region_info["z_max"], - norm=(region_info["x_range"] * region_info["y_range"]), ) diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index 60110658..fa5e850f 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -121,9 +121,9 @@ def slice_gas( box_x=region_info["periodic_box_x"], box_y=region_info["periodic_box_y"], box_z=region_info["periodic_box_z"], - norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]), ) + norm = region_info["x_range"] * region_info["y_range"] * region_info["z_range"] backend_func = (backends_parallel if parallel else backends)[backend] - image = backend_restore_cosmo_and_units(backend_func)(**kwargs) + image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs) return image diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index ceb9b593..1e3296a6 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -104,10 +104,10 @@ def render_gas( box_x=region_info["periodic_box_x"], box_y=region_info["periodic_box_y"], box_z=region_info["periodic_box_z"], - norm=(region_info["x_range"] * region_info["y_range"] * region_info["z_range"]), ) + norm = region_info["x_range"] * region_info["y_range"] * region_info["z_range"] backend_func = (backends_parallel if parallel else backends)["scatter"] - image = backend_restore_cosmo_and_units(backend_func)(**kwargs) + image = backend_restore_cosmo_and_units(backend_func, norm=norm)(**kwargs) return image From d784b29b2f435a2ff19334978ea376ea45636e99 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 13 Feb 2025 09:19:19 +0000 Subject: [PATCH 090/125] Update tests with norm arg removed from backend functions. --- tests/test_visualisation.py | 51 +++++++++++-------------------------- 1 file changed, 15 insertions(+), 36 deletions(-) diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 429e7730..cae91ddf 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -57,7 +57,6 @@ def test_scatter(save=False): res=256, box_x=1.0, box_y=1.0, - norm=1.0, ) except CudaSupportError: if CUDA_AVAILABLE: @@ -84,10 +83,8 @@ def test_scatter_mass_conservation(): for resolution in resolutions: scatter = projection_backends["fast"] - image = scatter( - x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, norm=1.0 - ) - mass_in_image = image.sum() / (resolution ** 2) + image = scatter(x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0) + mass_in_image = image.sum() / (resolution**2) # Check mass conservation to 5% assert np.isclose(mass_in_image.view(np.ndarray), total_mass, 0.05) @@ -123,7 +120,6 @@ def test_scatter_parallel(save=False): res=resolution, box_x=1.0, box_y=1.0, - norm=1.0, ) image_par = scatter_parallel( x=coordinates[0], @@ -133,7 +129,6 @@ def test_scatter_parallel(save=False): res=resolution, box_x=1.0, box_y=1.0, - norm=1.0, ) if save: @@ -158,7 +153,6 @@ def test_slice(save=False): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) if save: @@ -199,7 +193,6 @@ def test_slice_parallel(save=False): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) image_par = slice_backends_parallel[backend]( x=coordinates[0], @@ -213,7 +206,6 @@ def test_slice_parallel(save=False): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) assert np.isclose(image, image_par).all() @@ -237,7 +229,6 @@ def test_volume_render(): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) return @@ -267,7 +258,6 @@ def test_volume_parallel(): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) scatter_parallel = volume_render_backends_parallel["scatter"] image_par = scatter_parallel( @@ -280,7 +270,6 @@ def test_volume_parallel(): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) assert np.isclose(image, image_par).all() @@ -346,13 +335,13 @@ def test_render_outside_region(): h[h > 0.5] = 0.05 m = np.ones_like(h) projection_backends["histogram"]( - x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, norm=1.0 + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0 ) for backend in projection_backends.keys(): try: projection_backends[backend]( - x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, norm=1.0 + x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0 ) except CudaSupportError: if CUDA_AVAILABLE: @@ -372,7 +361,6 @@ def test_render_outside_region(): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) volume_render_backends_parallel["scatter"]( @@ -385,7 +373,6 @@ def test_render_outside_region(): box_x=1.0, box_y=1.0, box_z=1.0, - norm=1.0, ) @@ -448,7 +435,7 @@ def test_comoving_versus_physical(filename): ): img = func(data, resolution=64, project="masses", region=region) assert data.gas.masses.comoving and img.comoving - assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 + assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 # densities are physical, make sure this works with physical coordinates and # smoothing lengths img = func(data, resolution=64, project="densities", region=region) @@ -514,14 +501,14 @@ def test_panel_rendering(filename): plt.imsave( "panels_added.png", - plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), + plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1))), ) projected = project_gas(data, res, "masses", backend="renormalised") plt.imsave( "projected.png", - plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), + plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)), ) fullstack = np.zeros((res, res)) @@ -575,7 +562,6 @@ def test_periodic_boundary_wrapping(): res=pixel_resolution, box_x=boxsize, box_y=boxsize, - norm=boxsize ** 2, ) image2 = projection_backends[backend]( x=coordinates_non_periodic[:, 0], @@ -585,7 +571,6 @@ def test_periodic_boundary_wrapping(): res=pixel_resolution, box_x=0.0, box_y=0.0, - norm=boxsize ** 2, ) assert (image1 == image2).all() except CudaSupportError: @@ -608,7 +593,6 @@ def test_periodic_boundary_wrapping(): box_x=boxsize, box_y=boxsize, box_z=boxsize, - norm=boxsize ** 3, ) image2 = slice_backends[backend]( x=coordinates_non_periodic[:, 0], @@ -622,7 +606,6 @@ def test_periodic_boundary_wrapping(): box_x=0.0, box_y=0.0, box_z=0.0, - norm=boxsize ** 3, ) assert (image1 == image2).all() @@ -639,7 +622,6 @@ def test_periodic_boundary_wrapping(): box_x=boxsize, box_y=boxsize, box_z=boxsize, - norm=boxsize ** 3, ) image2 = scatter( x=coordinates_non_periodic[:, 0], @@ -651,7 +633,6 @@ def test_periodic_boundary_wrapping(): box_x=0.0, box_y=0.0, box_z=0.0, - norm=boxsize ** 3, ) assert (image1 == image2).all() @@ -686,11 +667,9 @@ def test_volume_render_and_unfolded_deposit(): box_x=boxsize, box_y=boxsize, box_z=boxsize, - norm=boxsize ** 3, ) - # need to divide out the box volume for the deposition - assert np.allclose(deposition / boxsize ** 3, volume.view(np.ndarray)) + assert np.allclose(deposition, volume.view(np.ndarray)) def test_folding_deposit(): @@ -729,8 +708,8 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v + mean_density_deposit = (np.sum(deposition) / npix**3).to("Msun / kpc**3").v + mean_density_volume = (np.sum(volume) / npix**3).to("Msun / kpc**3").v mean_density_calculated = ( (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) .to("Msun / kpc**3") @@ -755,15 +734,15 @@ def test_dark_matter_power_spectrum(filename, save=False): min_k = cosmo_quantity( 1e-2, - unyt.Mpc ** -1, + unyt.Mpc**-1, comoving=True, - cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), ) max_k = cosmo_quantity( 1e2, - unyt.Mpc ** -1, + unyt.Mpc**-1, comoving=True, - cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), ) bins = np.geomspace(min_k, max_k, 32) @@ -800,7 +779,7 @@ def test_dark_matter_power_spectrum(filename, save=False): folds[folding] = deposition - folding_output[2 ** folding] = (k, power_spectrum, scatter) + folding_output[2**folding] = (k, power_spectrum, scatter) # Now try doing them all together at once. From bccf1ca2f6b3d1eb4a73511d801501399e371d24 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Thu, 13 Feb 2025 09:33:42 +0000 Subject: [PATCH 091/125] Run black. --- swiftsimio/_array_functions.py | 2 +- tests/test_visualisation.py | 56 +++++++++++++++------------------- 2 files changed, 25 insertions(+), 33 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 1557f840..334a7d52 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -2073,7 +2073,7 @@ def einsum( @implements(np.unwrap) -def unwrap(p, discont=None, axis=-1, *, period=6.283185307179586): +def unwrap(p, discont=None, axis=-1, *, period=6.283_185_307_179_586): helper_result = _prepare_array_func_args( p, discont=discont, axis=axis, period=period diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index cae91ddf..f371da81 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -50,10 +50,10 @@ def test_scatter(save=False): for backend in projection_backends.keys(): try: image = projection_backends[backend]( - x=np.array([0.0, 1.0, 1.0, -0.000001]), - y=np.array([0.0, 0.0, 1.0, 1.000001]), + x=np.array([0.0, 1.0, 1.0, -0.000_001]), + y=np.array([0.0, 0.0, 1.0, 1.000_001]), m=np.array([1.0, 1.0, 1.0, 1.0]), - h=np.array([0.2, 0.2, 0.2, 0.000002]), + h=np.array([0.2, 0.2, 0.2, 0.000_002]), res=256, box_x=1.0, box_y=1.0, @@ -71,7 +71,7 @@ def test_scatter(save=False): def test_scatter_mass_conservation(): - np.random.seed(971263) + np.random.seed(971_263) # Width of 0.8 centered on 0.5, 0.5. x = 0.8 * np.random.rand(100) + 0.1 y = 0.8 * np.random.rand(100) + 0.1 @@ -84,7 +84,7 @@ def test_scatter_mass_conservation(): for resolution in resolutions: scatter = projection_backends["fast"] image = scatter(x=x, y=y, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0) - mass_in_image = image.sum() / (resolution**2) + mass_in_image = image.sum() / (resolution ** 2) # Check mass conservation to 5% assert np.isclose(mass_in_image.view(np.ndarray), total_mass, 0.05) @@ -142,11 +142,11 @@ def test_scatter_parallel(save=False): def test_slice(save=False): slice = slice_backends["sph"] image = slice( - x=np.array([0.0, 1.0, 1.0, -0.000001]), - y=np.array([0.0, 0.0, 1.0, 1.000001]), - z=np.array([0.0, 0.0, 1.0, 1.000001]), + x=np.array([0.0, 1.0, 1.0, -0.000_001]), + y=np.array([0.0, 0.0, 1.0, 1.000_001]), + z=np.array([0.0, 0.0, 1.0, 1.000_001]), m=np.array([1.0, 1.0, 1.0, 1.0]), - h=np.array([0.2, 0.2, 0.2, 0.000002]), + h=np.array([0.2, 0.2, 0.2, 0.000_002]), z_slice=0.99, xres=256, yres=256, @@ -220,11 +220,11 @@ def test_volume_render(): # render image scatter = volume_render_backends["scatter"] scatter( - x=np.array([0.0, 1.0, 1.0, -0.000001]), - y=np.array([0.0, 0.0, 1.0, 1.000001]), - z=np.array([0.0, 0.0, 1.0, 1.000001]), + x=np.array([0.0, 1.0, 1.0, -0.000_001]), + y=np.array([0.0, 0.0, 1.0, 1.000_001]), + z=np.array([0.0, 0.0, 1.0, 1.000_001]), m=np.array([1.0, 1.0, 1.0, 1.0]), - h=np.array([0.2, 0.2, 0.2, 0.000002]), + h=np.array([0.2, 0.2, 0.2, 0.000_002]), res=64, box_x=1.0, box_y=1.0, @@ -364,15 +364,7 @@ def test_render_outside_region(): ) volume_render_backends_parallel["scatter"]( - x=x, - y=y, - z=z, - m=m, - h=h, - res=resolution, - box_x=1.0, - box_y=1.0, - box_z=1.0, + x=x, y=y, z=z, m=m, h=h, res=resolution, box_x=1.0, box_y=1.0, box_z=1.0 ) @@ -435,7 +427,7 @@ def test_comoving_versus_physical(filename): ): img = func(data, resolution=64, project="masses", region=region) assert data.gas.masses.comoving and img.comoving - assert (img.cosmo_factor.expr - a**aexp).simplify() == 0 + assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0 # densities are physical, make sure this works with physical coordinates and # smoothing lengths img = func(data, resolution=64, project="densities", region=region) @@ -501,14 +493,14 @@ def test_panel_rendering(filename): plt.imsave( "panels_added.png", - plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(np.sum(panel, axis=-1))), + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(np.sum(panel, axis=-1))), ) projected = project_gas(data, res, "masses", backend="renormalised") plt.imsave( "projected.png", - plt.get_cmap()(LogNorm(vmin=10**6, vmax=10**6.5)(projected)), + plt.get_cmap()(LogNorm(vmin=10 ** 6, vmax=10 ** 6.5)(projected)), ) fullstack = np.zeros((res, res)) @@ -708,8 +700,8 @@ def test_volume_render_and_unfolded_deposit_with_units(filename): # Volume render the particles volume = render_gas(data, npix, parallel=False).to_physical() - mean_density_deposit = (np.sum(deposition) / npix**3).to("Msun / kpc**3").v - mean_density_volume = (np.sum(volume) / npix**3).to("Msun / kpc**3").v + mean_density_deposit = (np.sum(deposition) / npix ** 3).to("Msun / kpc**3").v + mean_density_volume = (np.sum(volume) / npix ** 3).to("Msun / kpc**3").v mean_density_calculated = ( (np.sum(data.gas.masses) / (data.metadata.boxsize[0] * data.metadata.a) ** 3) .to("Msun / kpc**3") @@ -734,15 +726,15 @@ def test_dark_matter_power_spectrum(filename, save=False): min_k = cosmo_quantity( 1e-2, - unyt.Mpc**-1, + unyt.Mpc ** -1, comoving=True, - cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), ) max_k = cosmo_quantity( 1e2, - unyt.Mpc**-1, + unyt.Mpc ** -1, comoving=True, - cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor), + cosmo_factor=cosmo_factor(a ** -1, data.metadata.scale_factor), ) bins = np.geomspace(min_k, max_k, 32) @@ -779,7 +771,7 @@ def test_dark_matter_power_spectrum(filename, save=False): folds[folding] = deposition - folding_output[2**folding] = (k, power_spectrum, scatter) + folding_output[2 ** folding] = (k, power_spectrum, scatter) # Now try doing them all together at once. From 1b3feae703b4bf08d3c933a746effbaae1ded033 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 14 Feb 2025 08:55:03 +0000 Subject: [PATCH 092/125] Update some vis docstrings. --- swiftsimio/visualisation/projection.py | 11 +++++++---- swiftsimio/visualisation/slice.py | 11 ++++++----- swiftsimio/visualisation/volume_render.py | 12 +++++++----- 3 files changed, 20 insertions(+), 14 deletions(-) diff --git a/swiftsimio/visualisation/projection.py b/swiftsimio/visualisation/projection.py index 24da3638..a615e035 100644 --- a/swiftsimio/visualisation/projection.py +++ b/swiftsimio/visualisation/projection.py @@ -51,7 +51,8 @@ def project_pixel_grid( Variable to project to get the weighted density of. By default, this is mass. If you would like to mass-weight any other variable, you can always create it as ``data.gas.my_variable = data.gas.other_variable - * data.gas.masses``. + * data.gas.masses``. The result is comoving if this is comoving, else + it is physical. region: cosmo_array, optional Region, determines where the image will be created (this corresponds @@ -90,6 +91,7 @@ def project_pixel_grid( image: cosmo_array Projected image with units of project / length^2, of size ``res`` x ``res``. + Comoving if ``project`` data are comoving, else physical. Notes @@ -171,7 +173,8 @@ def project_gas( Variable to project to get the weighted density of. By default, this is mass. If you would like to mass-weight any other variable, you can always create it as ``data.gas.my_variable = data.gas.other_variable - * data.gas.masses``. + * data.gas.masses``. The result is comoving if this is comoving, else + it is physical. region: cosmo_array, optional Region, determines where the image will be created (this corresponds @@ -210,8 +213,8 @@ def project_gas( ------- image: cosmo_array - Projected image with units of project / length^2, of size ``res`` x - ``res``. + Projected image with units of project / length^2, of size ``res`` x ``res``. + Comoving if ``project`` data are comoving, else physical. Notes diff --git a/swiftsimio/visualisation/slice.py b/swiftsimio/visualisation/slice.py index fa5e850f..25ad1dd0 100644 --- a/swiftsimio/visualisation/slice.py +++ b/swiftsimio/visualisation/slice.py @@ -46,8 +46,9 @@ def slice_gas( this value refers to the location along the rotated z-axis. project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles + Data field to be projected. Default is mass. If ``None`` then simply + count number of particles. The result is comoving if this is comoving, + else it is physical. parallel : bool used to determine if we will create the image in parallel. This @@ -83,9 +84,9 @@ def slice_gas( Returns ------- - ndarray of np.float32 - Creates a `resolution` x `resolution` np.array and returns it, - without appropriate units. + image : cosmo_array + Slice image with units of project / length^2, of size ``res`` x ``res``. + Comoving if ``project`` data are comoving, else physical. See Also -------- diff --git a/swiftsimio/visualisation/volume_render.py b/swiftsimio/visualisation/volume_render.py index 1e3296a6..8662acdb 100644 --- a/swiftsimio/visualisation/volume_render.py +++ b/swiftsimio/visualisation/volume_render.py @@ -44,8 +44,9 @@ def render_gas( Specifies size of return np.array project : str, optional - Data field to be projected. Default is mass. If None then simply - count number of particles + Data field to be projected. Default is ``"mass"``. If ``None`` then simply + count number of particles. The result is comoving if this is comoving, else + it is physical. parallel : bool used to determine if we will create the image in parallel. This @@ -78,9 +79,10 @@ def render_gas( Returns ------- - np.ndarray of np.float32 - Creates a `resolution` x `resolution` x `resolution` np.array and - returns it, without appropriate units. + cosmo_array + Voxel grid with units of project / length^3, of size ``resolution`` x + ``resolution`` x ``resolution``. Comoving if ``project`` data are + comoving, else physical. See Also -------- From 71c9212ff5474f0bf491d9e71b8a8fc04d1935ee Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 14 Feb 2025 13:24:49 +0000 Subject: [PATCH 093/125] Support __round__ for cosmo_quantity and dot for cosmo_array. --- swiftsimio/objects.py | 6 ++++++ tests/test_cosmo_array.py | 27 +++++++++++++++++++++++++++ 2 files changed, 33 insertions(+) diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 0516d64a..6f5a5d04 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -116,6 +116,7 @@ _arctan2_cosmo_factor, _comparison_cosmo_factor, _prepare_array_func_args, + _default_binary_wrapper, ) try: @@ -1260,6 +1261,7 @@ def __setstate__(self, state: Tuple) -> None: __getitem__ = _propagate_cosmo_array_attributes_to_result( _ensure_result_is_cosmo_array_or_quantity(unyt_array.__getitem__) ) + dot = _default_binary_wrapper(unyt_array.dot, _multiply_cosmo_factor) # Also wrap some array "properties": T = property(_propagate_cosmo_array_attributes_to_result(unyt_array.transpose)) @@ -1799,3 +1801,7 @@ def __new__( if ret.size > 1: raise RuntimeError("cosmo_quantity instances must be scalars") return ret + + __round__ = _propagate_cosmo_array_attributes_to_result( + _ensure_result_is_cosmo_array_or_quantity(unyt_quantity.__round__) + ) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index ba112464..008638e3 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -771,6 +771,16 @@ def test_iter(self): for cq in ca(np.arange(3)): assert isinstance(cq, cosmo_quantity) + def test_dot(self): + """ + Make sure that we get a cosmo_array when we use array attribute dot. + """ + res = ca(np.arange(3)).dot(ca(np.arange(3))) + assert isinstance(res, cosmo_quantity) + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5) + assert res.valid_transform is True + class TestCosmoQuantity: """ @@ -810,6 +820,23 @@ def test_propagation_func(self, func, args): assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) assert res.valid_transform is True + def test_round(self): + """ + Test that attributes propagate through the round builtin. + """ + cq = cosmo_quantity( + 1.03, + u.m, + comoving=False, + cosmo_factor=cosmo_factor(a ** 1, 1.0), + valid_transform=True, + ) + res = round(cq) + assert res.value == 1.0 + assert res.comoving is False + assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0) + assert res.valid_transform is True + def test_scalar_return_func(self): """ Make sure that default-wrapped functions that take a cosmo_array and return a From 7b765b7ddb6c0e000194c41949e9ac1b6ad9b6c0 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 14 Feb 2025 23:06:59 +0000 Subject: [PATCH 094/125] Make a pass over the visualisation docs. --- docs/source/visualisation/index.rst | 11 +-- docs/source/visualisation/power_spectra.rst | 52 ++++++++----- docs/source/visualisation/projection.rst | 39 +++++----- docs/source/visualisation/ray_trace.rst | 4 + docs/source/visualisation/slice.rst | 26 ++++--- docs/source/visualisation/tools.rst | 8 +- docs/source/visualisation/volume_render.rst | 31 ++++---- swiftsimio/_array_functions.py | 18 +++-- .../visualisation/projection_backends/fast.py | 5 +- .../projection_backends/histogram.py | 75 +++++++++---------- .../projection_backends/kernels.py | 31 ++++---- .../projection_backends/reference.py | 9 +-- .../projection_backends/renormalised.py | 21 ++---- .../projection_backends/subsampled.py | 12 +-- .../projection_backends/subsampled_extreme.py | 12 +-- .../smoothing_length/generate.py | 9 +-- 16 files changed, 185 insertions(+), 178 deletions(-) create mode 100644 docs/source/visualisation/ray_trace.rst diff --git a/docs/source/visualisation/index.rst b/docs/source/visualisation/index.rst index 2b6ad60c..036e7be2 100644 --- a/docs/source/visualisation/index.rst +++ b/docs/source/visualisation/index.rst @@ -1,12 +1,12 @@ Visualisation ============= -:mod:`swiftsimio` provides visualisation routines accelerated with the -:mod:`numba` module. They work without this module, but we strongly recommend -installing it for the best performance (1000x+ speedups). These are provided -in the :mod:`swiftismio.visualisation` sub-modules. +:mod:`swiftsimio` provides visualisation routines in the +:mod:`swiftsimio.visualisation` sub-module. They are accelerated with the +:mod:`numba` module. They can work without :mod:`numba`, but we strongly recommend +installing it for the best performance (1000x+ speedups). -The three built-in rendering types (described below) have the following +The four built-in rendering types (described below) have the following common interface: .. code-block:: python @@ -30,6 +30,7 @@ additional functionality. projection slice volume_render + ray_trace power_spectra tools diff --git a/docs/source/visualisation/power_spectra.rst b/docs/source/visualisation/power_spectra.rst index 7a892356..2f7211be 100644 --- a/docs/source/visualisation/power_spectra.rst +++ b/docs/source/visualisation/power_spectra.rst @@ -6,15 +6,15 @@ runs on-the-fly, and as such after a run has completed you may wish to create a number of more non-standard power spectra. These tools are available as part of the :mod:`swiftsimio.visualisation.power_spectrum` -package. Making a power spectrum consists of two major steps: depositing the particles +module. Making a power spectrum consists of two major steps: depositing the particles on grid(s), and then binning their fourier transform to get the one-dimensional power. -Depositing on a Grid +Depositing on a grid -------------------- Depositing your particles on a grid is performed using -:meth:`swiftsimio.visualisation.power_spectrum.render_to_deposit`. This function +:func:`swiftsimio.visualisation.power_spectrum.render_to_deposit`. This function performs a nearest-grid-point (NGP) of all particles in the provided particle dataset. For example: @@ -35,17 +35,17 @@ dataset. For example: The specific field being depositied can be controlled with the ``project`` keyword argument. The ``resolution``` argument gives the one-dimensional resolution of the 3D grid, so in this case you would recieve a ``512x512x512`` -grid. Note that the ``gas_mass_deposit`` is a :obj:`swiftsimio.cosmo_array`, +grid. Note that the ``gas_mass_deposit`` is a :obj:`~swiftsimio.objects.cosmo_array`, and as such includes cosmological and unit information that is used later in the process. -Generating a Power Spectrum +Generating a power spectrum --------------------------- Once you have your grid deposited, you can easily generate a power spectrum using the -:meth:`swiftsimio.visualisation.power_spectrum.deposition_to_power_spectrum` +:func:`~swiftsimio.visualisation.power_spectrum.deposition_to_power_spectrum` function. For example, using the above deposit: .. code-block:: python @@ -64,7 +64,7 @@ Wavenumbers are calculated to be at the weighted mean of the k-values in each bin, rather than representing the center of the bin. -More Complex Scenarios +More complex scenarios ---------------------- In a realistic simualted power spectrum, you will need to perform 'folding' @@ -89,13 +89,14 @@ The ``folding`` parameter is available for both ``render_to_deposit`` and ``deposition_to_power_spectrum``, but it may be easier to use the utility functions provided for automatically stitching together the folded spectra. The function -:meth:`swiftsimio.visualsation.power_spectrum.folded_depositions_to_power_spectrum` +:func:`~swiftsimio.visualsation.power_spectrum.folded_depositions_to_power_spectrum` allows you to do this easily: .. code-block:: python + import unyt as u from swiftsimio.visualisation.power_spectrum import folded_depositions_to_power_spectrum - import unyt + from swiftsimio.objects import cosmo_quantity, cosmo_factor, a folded_depositions = {} @@ -112,7 +113,20 @@ allows you to do this easily: depositions=folded_depositions, box_size=data.metadata.box_size, number_of_wavenumber_bins=128, - wavenumber_range=[1e-2 / unyt.Mpc, 1e2 / unyt.Mpc], + wavenumber_range=[ + cosmo_quantity( + 1e-2, + u.Mpc**-1, + comoving=True, + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor, + ) + cosmo_quantity( + 1e2, + u.Mpc**-1, + comoving=True, + cosmo_factor=cosmo_factor(a**-1, data.metadata.scale_factor, + ), + ], log_wavenumber_bins=True, workers=4, minimal_sample_modes=8192, @@ -129,19 +143,19 @@ visualisation. There are a few crucial parameters to this function: 1. ``workers`` is the number of threads to use for the calculation of - the fourier transforms. + the Fourier transforms. 2. ``minimal_sample_modes`` is the minimum number of modes that must be - present in a bin for it to be included in the final power spectrum. - Generally for a big simulation you want to set this to around 10'000, - and this number is ignored for the lowest wavenumber bin. + present in a bin for it to be included in the final power spectrum. + Generally for a big simulation you want to set this to around 10,000, + and this number is ignored for the lowest wavenumber bin. 3. ``cutoff_above_wavenumber_fraction`` is the fraction of the individual fold's (as represented by the FFT itself) maximally sampled wavenumber. Ignored for the last fold, and we always cap the maximal - wavenumber to the nyquist frequency. + wavenumber to the Nyquist frequency. 4. ``shot_noise_norm`` is the number of particles in the simulation - that contribute to the power spectrum. This is used to normalise - the power spectrum to the shot noise level. This is very - important in this case because of the use of NGP deposition. + that contribute to the power spectrum. This is used to normalise + the power spectrum to the shot noise level. This is very + important in this case because of the use of NGP deposition. Foldings are stitched using a simple method where the 'better sampled' -foldings are used preferentially, up to the cutoff value. \ No newline at end of file +foldings are used preferentially, up to the cutoff value. diff --git a/docs/source/visualisation/projection.rst b/docs/source/visualisation/projection.rst index e92428e1..8898130d 100644 --- a/docs/source/visualisation/projection.rst +++ b/docs/source/visualisation/projection.rst @@ -15,7 +15,7 @@ with :math:`\tilde{A}_i` the smoothed quantity in pixel :math:`i`, and Here we use the Wendland-C2 kernel. The primary function here is -:meth:`swiftsimio.visualisation.projection.project_gas`, which allows you to +:func:`swiftsimio.visualisation.projection.project_gas`, which allows you to create a gas projection of any field. See the example below. Example @@ -29,7 +29,7 @@ Example data = load("cosmo_volume_example.hdf5") # This creates a grid that has units msun / Mpc^2, and can be transformed like - # any other unyt quantity + # any other cosmo_array mass_map = project_gas( data, resolution=1024, @@ -52,7 +52,7 @@ Example This basic demonstration creates a mass surface density map. To create, for example, a projected temperature map, we need to remove the -surface density dependence (i.e. :meth:`project_gas` returns a surface +surface density dependence (i.e. :func:`~swiftsimio.visualisation.projection.project_gas` returns a surface temperature in units of K / kpc^2 and we just want K) by dividing out by this: @@ -114,11 +114,11 @@ backends are as follows: + ``fast``: The default backend - this is extremely fast, and provides very basic smoothing, with a return type of single precision floating point numbers. + ``histogram``: This backend provides zero smoothing, and acts in a similar way - to the ``np.hist2d`` function but with the same arguments as ``scatter``. -+ ``reference``: The same backend as ``fast`` but with two distinguishing features; + to the :func:`~numpy.histogram2d` function but with the same arguments as ``scatter``. ++ ``reference``: The same backend as ``fast`` but with two distinguishing features: all calculations are performed in double precision, and it will return early with a warning message if there are not enough pixels to fully resolve each kernel. - Regular users should not use this mode. + Intended for developer usage, regular users should not use this mode. + ``renormalised``: The same as ``fast``, but each kernel is evaluated twice and renormalised to ensure mass conservation within floating point precision. Returns single precision arrays. @@ -166,7 +166,7 @@ All visualisation functions by default assume a periodic box. Rather than simply projecting each individual particle once, four additional periodic copies of each particle are also projected. Most copies will project outside the valid pixel range, but the copies that do not ensure that pixels close to the edge -receive all necessary contributions. Thanks to Numba optimisations, the overhead +receive all necessary contributions. Thanks to :mod:`numba` optimisations, the overhead of these additional copies is relatively small. There are some caveats with this approach. If you try to visualise a subset of @@ -183,8 +183,8 @@ Rotations Sometimes you will need to visualise a galaxy from a different perspective. The :mod:`swiftsimio.visualisation.rotation` sub-module provides routines to generate rotation matrices corresponding to vectors, which can then be -provided to the ``rotation_matrix`` argument of :meth:`project_gas` (and -:meth:`project_gas_pixel_grid`). You will also need to supply the +provided to the ``rotation_matrix`` argument of :func:`~swiftsimio.visualisation.projection.project_gas` (and +:func:`~swiftsimio.visualisation.projection.project_gas_pixel_grid`). You will also need to supply the ``rotation_center`` argument, as the rotation takes place around this given point. The example code below loads a snapshot, and a halo catalogue, and creates an edge-on and face-on projection using the integration in @@ -287,15 +287,12 @@ Other particle types -------------------- Other particle types are able to be visualised through the use of the -:meth:`swiftsimio.visualisation.projection.project_pixel_grid` function. This -does not attach correct symbolic units, so you will have to work those out -yourself, but it does perform the smoothing. We aim to introduce the feature -of correctly applied units to these projections soon. +:func:`swiftsimio.visualisation.projection.project_pixel_grid` function. To use this feature for particle types that do not have smoothing lengths, you will need to generate them, as in the example below where we create a mass density map for dark matter. We provide a utility to do this through -:meth:`swiftsimio.visualisation.smoothing_length.generate_smoothing_lengths`. +:func:`~swiftsimio.visualisation.smoothing_length.generate.generate_smoothing_lengths`. .. code-block:: python @@ -348,15 +345,15 @@ smoothing lengths, and smoothed quantities, to generate a pixel grid that represents the smoothed version of the data. This API is available through -:meth:`swiftsimio.visualisation.projection.scatter` and -:meth:`swiftsimio.visualisation.projection.scatter_parallel` for the parallel +:obj:`swiftsimio.visualisation.projection_backends.backends["scatter"]` and +:obj:`swiftsimio.visualisation.projection_backends.backends_parallel["scatter"]` for the parallel version. The parallel version uses significantly more memory as it allocates a thread-local image array for each thread, summing them in the end. Here we will only describe the ``scatter`` variant, but they behave in the exact same way. By default this uses the "fast" backend. To use the others, you can select them manually from the module, or by using the ``backends`` and ``backends_parallel`` -dictionaries in :mod:`swiftsimio.visualisation.projection`. +dictionaries in :mod:`swiftsimio.visualisation.projection_backends`. To use this function, you will need: @@ -374,7 +371,9 @@ The key here is that only particles in the domain [0, 1] in x, and [0, 1] in y will be visible in the image. You may have particles outside of this range; they will not crash the code, and may even contribute to the image if their smoothing lengths overlap with [0, 1]. You will need to re-scale your data -such that it lives within this range. Then you may use the function as follows: +such that it lives within this range. You should also pass raw numpy arrays (not +:class:`~swiftsimio.objects.cosmo_array` or :class:`~unyt.array.unyt_array`, the +inputs are dimensionless). Then you may use the function as follows: .. code-block:: python @@ -383,12 +382,12 @@ such that it lives within this range. Then you may use the function as follows: # Using the variable names from above out = scatter(x=x, y=y, h=h, m=m, res=res) -``out`` will be a 2D :mod:`numpy` grid of shape ``[res, res]``. You will need +``out`` will be a 2D :class:`~numpy.ndarray` grid of shape ``[res, res]``. You will need to re-scale this back to your original dimensions to get it in the correct units, and do not forget that it now represents the smoothed quantity per surface area. If the optional arguments ``box_x`` and ``box_y`` are provided, they should contain the simulation box size in the same re-scaled coordinates as ``x`` and ``y``. The projection backend will then correctly apply periodic boundary -wrapping. If ``box_x`` and ``box_y`` are not provided or set to 0, no +wrapping. If ``box_x`` and ``box_y`` are not provided or set to ``0``, no periodic boundaries are applied. diff --git a/docs/source/visualisation/ray_trace.rst b/docs/source/visualisation/ray_trace.rst new file mode 100644 index 00000000..9830b113 --- /dev/null +++ b/docs/source/visualisation/ray_trace.rst @@ -0,0 +1,4 @@ +Ray tracing +=========== + +Documentation to be completed... diff --git a/docs/source/visualisation/slice.rst b/docs/source/visualisation/slice.rst index 7fb7ca8c..4345a396 100644 --- a/docs/source/visualisation/slice.rst +++ b/docs/source/visualisation/slice.rst @@ -5,7 +5,7 @@ The :mod:`swiftsimio.visualisation.slice` sub-module provides an interface to render SWIFT data onto a slice. This takes your 3D data and finds the 3D density at fixed z-position, slicing through the box. -The default :code:`"sph"` backend effectively solves the equation: +The default ``"sph"`` backend effectively solves the equation: :math:`\tilde{A}_i = \sum_j A_j W_{ij, 3D}` @@ -19,7 +19,7 @@ nearest-neighbour interpolation to compute the densities at each pixel. This backend is more suited for use with moving-mesh hydrodynamics schemes. The primary function here is -:meth:`swiftsimio.visualisation.slice.slice_gas`, which allows you to +:func:`swiftsimio.visualisation.slice.slice_gas`, which allows you to create a gas slice of any field. See the example below. Example @@ -58,8 +58,9 @@ Example This basic demonstration creates a mass density map. To create, for example, a projected temperature map, we need to remove the -density dependence (i.e. :meth:`slice_gas` returns a volumetric temperature -in units of K / kpc^3 and we just want K) by dividing out by this: +density dependence (i.e. :func:`~swiftsimio.visualisation.slice.slice_gas` +returns a volumetric temperature in units of K / kpc^3 and we just want K) +by dividing out by this: .. code-block:: python @@ -121,7 +122,7 @@ All visualisation functions by default assume a periodic box. Rather than simply summing each individual particle once, eight additional periodic copies of each particle are also accounted for. Most copies will contribute outside the valid pixel range, but the copies that do not ensure that pixels close to the -edge receive all necessary contributions. Thanks to Numba optimisations, the +edge receive all necessary contributions. Thanks to :mod:`numba` optimisations, the overhead of these additional copies is relatively small. There are some caveats with this approach. If you try to visualise a subset of @@ -139,11 +140,12 @@ Rotations Rotations of the box prior to slicing are provided in a similar fashion to the :mod:`swiftsimio.visualisation.projection` sub-module, by using the :mod:`swiftsimio.visualisation.rotation` sub-module. To rotate the perspective -prior to slicing a ``rotation_center`` argument in :meth:`slice_gas` needs +prior to slicing a ``rotation_center`` argument in +:func:`~swiftsimio.visualisation.slice.slice_gas` needs to be provided, specifying the point around which the rotation takes place. The angle of rotation is specified with a matrix, supplied by ``rotation_matrix`` -in :meth:`slice_gas`. The rotation matrix may be computed with -:meth:`rotation_matrix_from_vector`. This will result in the perspective being +in :func:`~swiftsimio.visualisation.slice.slice_gas`. The rotation matrix may be computed with +:func:`~swiftsimio.visualisation.rotation.rotation_matrix_from_vector`. This will result in the perspective being rotated to be along the provided vector. This approach to rotations applied to the above example is shown below. @@ -209,8 +211,8 @@ smoothing lengths, and smoothed quantities, to generate a pixel grid that represents the smoothed, sliced, version of the data. This API is available through -:meth:`swiftsimio.visualisation.slice.slice_scatter` and -:meth:`swiftsimio.visualisation.slice.slice_scatter_parallel` for the parallel +:func:`swiftsimio.visualisation.slice_backends.backends["sph"]` and +:func:`swiftsimio.visualisation.slice_backends.backends_parallel["sph"]` for the parallel version. The parallel version uses significantly more memory as it allocates a thread-local image array for each thread, summing them in the end. Here we will only describe the ``scatter`` variant, but they behave in the exact same way. @@ -235,7 +237,9 @@ not crash the code, and may even contribute to the image if their smoothing lengths overlap with [0, 1]. You will need to re-scale your data such that it lives within this range. Smoothing lengths and z coordinates need to be re-scaled in the same way (using the same scaling factor), but z coordinates do -not need to lie in the domain [0, 1]. Then you may use the function as follows: +not need to lie in the domain [0, 1]. You should provide inputs as raw numpy arrays +(not :class:`~swiftsimio.objects.cosmo_array` or :class:`~unyt.array.unyt_array`). +Then you may use the function as follows: .. code-block:: python diff --git a/docs/source/visualisation/tools.rst b/docs/source/visualisation/tools.rst index d4244927..6f92a8d4 100644 --- a/docs/source/visualisation/tools.rst +++ b/docs/source/visualisation/tools.rst @@ -9,8 +9,8 @@ Tools The :mod:`swiftsimio.visualisation.tools.cmaps` module includes three objects that can be used to deploy two dimensional colour maps. The first, -:class:`swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2D`, and second -:class:`swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2DHSV`, allow +:class:`~swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2D`, and second +:class:`~swiftsimio.visualisation.tools.cmaps.LinearSegmentedCmap2DHSV`, allow you to generate new color maps from sets of colors and coordinates. .. code-block:: python @@ -22,7 +22,7 @@ you to generate new color maps from sets of colors and coordinates. ) This generates a color map that is a quasi-linear interpolation between all -of the points. The map can be displayed using the ``plot`` method, +of the points. The map can be displayed using the :func:`~matplotlib.pyplot.plot` function, .. code-block:: python @@ -30,7 +30,7 @@ of the points. The map can be displayed using the ``plot`` method, bower.plot(ax) -Which generates: +which generates: .. image:: bower_cmap.png diff --git a/docs/source/visualisation/volume_render.rst b/docs/source/visualisation/volume_render.rst index 2c3d8ee5..a2bcd6db 100644 --- a/docs/source/visualisation/volume_render.rst +++ b/docs/source/visualisation/volume_render.rst @@ -15,7 +15,7 @@ with :math:`\tilde{A}_i` the smoothed quantity in pixel :math:`i`, and Here we use the Wendland-C2 kernel. The primary function here is -:meth:`swiftsimio.visualisation.volume_render.render_gas`, which allows you +:func:`swiftsimio.visualisation.volume_render.render_gas`, which allows you to create a gas density grid of any field, see the example below. Example @@ -41,7 +41,7 @@ Example This basic demonstration creates a mass density cube. To create, for example, a projected temperature cube, we need to remove the -density dependence (i.e. :meth:`render_gas` returns a volumetric +density dependence (i.e. :func:`~swiftsimio.visualisation.volume_render.render_gas` returns a volumetric temperature in units of K / kpc^3 and we just want K) by dividing out by this: @@ -89,7 +89,7 @@ All visualisation functions by default assume a periodic box. Rather than simply summing each individual particle once, eight additional periodic copies of each particle are also taken into account. Most copies will contribute outside the valid voxel range, but the copies that do not ensure that voxels -close to the edge receive all necessary contributions. Thanks to Numba +close to the edge receive all necessary contributions. Thanks to :mod:`numba` optimisations, the overhead of these additional copies is relatively small. There are some caveats with this approach. If you try to visualise a subset of @@ -106,11 +106,11 @@ Rotations Rotations of the box prior to volume rendering are provided in a similar fashion to the :mod:`swiftsimio.visualisation.projection` sub-module, by using the :mod:`swiftsimio.visualisation.rotation` sub-module. To rotate the perspective -prior to slicing a ``rotation_center`` argument in :meth:`render_gas` needs +prior to slicing a ``rotation_center`` argument in :func:`~swiftsimio.visualisation.volume_render.render_gas` needs to be provided, specifying the point around which the rotation takes place. The angle of rotation is specified with a matrix, supplied by ``rotation_matrix`` -in :meth:`render_gas`. The rotation matrix may be computed with -:meth:`rotation_matrix_from_vector`. This will result in the perspective being +in :func:`~swiftsimio.visualisation.volume_render.render_gas`. The rotation matrix may be computed with +:func:`~swiftsimio.visualisation.rotation.rotation_matrix_from_vector`. This will result in the perspective being rotated to be along the provided vector. This approach to rotations applied to the above example is shown below. @@ -160,8 +160,8 @@ Rendering --------- We provide a volume rendering function that can be used to make images highlighting -specific density contours. The notable function here is -:meth:``swiftsimio.visualisation.volume_render.visualise_render``. This takes +specific density contours. The key function here is +:func:`swiftsimio.visualisation.volume_render.visualise_render`. This takes in your volume rendering, along with a colour map and centers, to create these highlights. The example below shows how to use this. @@ -238,7 +238,7 @@ Here we can see the quick view of this image. It's just a regular density projec plt.savefig("volume_render_options.png") -This function :meth:`swiftsimio.visualisation.volume_render.visualise_render_options` allows +This function :func:`swiftsimio.visualisation.volume_render.visualise_render_options` allows you to see what densities your rendering is picking out: .. image:: volume_render_options.png @@ -280,8 +280,8 @@ smoothing lengths, and smoothed quantities, to generate a pixel grid that represents the smoothed, volume rendered, version of the data. This API is available through -:meth:`swiftsimio.visualisation.volume_render.scatter` and -:meth:`swiftsimio.visualisation.volume_render.scatter_parallel` for the parallel +:func:`swiftsimio.visualisation.volume_render_backends.backends["scatter"]` and +:func:`swiftsimio.visualisation.volume_render_backends.backends_parallel["scatter"]` for the parallel version. The parallel version uses significantly more memory as it allocates a thread-local image array for each thread, summing them in the end. Here we will only describe the ``scatter`` variant, but they behave in the exact same way. @@ -303,8 +303,9 @@ The key here is that only particles in the domain [0, 1] in x, [0, 1] in y, and [0, 1] in z. will be visible in the cube. You may have particles outside of this range; they will not crash the code, and may even contribute to the image if their smoothing lengths overlap with [0, 1]. You will need to -re-scale your data such that it lives within this range. Then you may use the -function as follows: +re-scale your data such that it lives within this range. You should pass in +raw numpy array (not :class:`~swiftsimio.objects.cosmo_array` or :class:`~unyt.array.unyt_array`). +Then you may use the function as follows: .. code-block:: python @@ -313,7 +314,7 @@ function as follows: # Using the variable names from above out = scatter(x=x, y=y, z=z, h=h, m=m, res=res) -``out`` will be a 3D :mod:`numpy` grid of shape ``[res, res, res]``. You will +``out`` will be a 3D :class:`~numpy.ndarray` grid of shape ``[res, res, res]``. You will need to re-scale this back to your original dimensions to get it in the correct units, and do not forget that it now represents the smoothed quantity per volume. @@ -322,4 +323,4 @@ If the optional arguments ``box_x``, ``box_y`` and ``box_z`` are provided, they should contain the simulation box size in the same re-scaled coordinates as ``x``, ``y`` and ``z``. The rendering function will then correctly apply periodic boundary wrapping. If ``box_x``, ``box_y`` and ``box_z`` are not -provided or set to 0, no periodic boundaries are applied +provided or set to 0, no periodic boundaries are applied. diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 334a7d52..6132f926 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -729,18 +729,18 @@ def _arctan2_cosmo_factor( f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor( @@ -1023,7 +1023,8 @@ def _default_unary_wrapper( def wrapper(*args, **kwargs): """ - Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + Prepare arguments, handle ``cosmo_factor`` attriubtes, and attach attributes to + output. Returns ------- @@ -1070,7 +1071,8 @@ def _default_binary_wrapper( def wrapper(*args, **kwargs): """ - Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + Prepare arguments, handle ``cosmo_factor`` attributes, and attach attributes to + output. Returns ------- @@ -1114,7 +1116,8 @@ def _default_comparison_wrapper(unyt_func: Callable) -> Callable: # _comparison_cosmo_factor with them as the inputs def wrapper(*args, **kwargs): """ - Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + Prepare arguments, handle ``cosmo_factor`` attributes, and attach attributes to + output. Returns ------- @@ -1156,7 +1159,8 @@ def _default_oplist_wrapper(unyt_func: Callable) -> Callable: def wrapper(*args, **kwargs): """ - Prepare arguments, handle ``cosmo_factor``s, and attach attributes to output. + Prepare arguments, handle ``cosmo_factor`` attributes, and attach attributes to + output. Returns ------- diff --git a/swiftsimio/visualisation/projection_backends/fast.py b/swiftsimio/visualisation/projection_backends/fast.py index 7f0f715d..96121146 100644 --- a/swiftsimio/visualisation/projection_backends/fast.py +++ b/swiftsimio/visualisation/projection_backends/fast.py @@ -154,8 +154,9 @@ def scatter( particle_cell_x + cells_spanned + 1, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float32(cell_x) + 0.5) * pixel_width - float32( x_pos ) diff --git a/swiftsimio/visualisation/projection_backends/histogram.py b/swiftsimio/visualisation/projection_backends/histogram.py index 0fa392f1..871da53d 100644 --- a/swiftsimio/visualisation/projection_backends/histogram.py +++ b/swiftsimio/visualisation/projection_backends/histogram.py @@ -4,23 +4,20 @@ Uses double precision. """ -from typing import Union -from math import sqrt, ceil -from numpy import float32, float64, int32, zeros, ndarray - +import numpy as np from swiftsimio.accelerated import jit, NUM_THREADS, prange @jit(nopython=True, fastmath=True) def scatter( - x: float64, - y: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + m: np.float32, + h: np.float32, res: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, +) -> np.ndarray: """ Creates a weighted scatter plot @@ -31,34 +28,34 @@ def scatter( Parameters ---------- - x : np.array[float64] + x : np.array[np.float64] array of x-positions of the particles. Must be bounded by [0, 1]. - y : np.array[float64] + y : np.array[np.float64] array of y-positions of the particles. Must be bounded by [0, 1]. - m : np.array[float32] + m : np.array[np.float32] array of masses (or otherwise weights) of the particles - h : np.array[float32] + h : np.array[np.float32] array of smoothing lengths of the particles res : int the number of pixels along one axis, i.e. this returns a square of res * res. - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x and y. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x and y. Used for periodic wrapping. Returns ------- - np.array[float32, float32, float32] + np.array[np.float32, np.float32, np.float32] pixel grid of quantity See Also @@ -74,13 +71,13 @@ def scatter( floats and integers is also an improvement over using the numba ones. """ # Output array for our image - image = zeros((res, res), dtype=float64) - maximal_array_index = int32(res) - 1 + image = np.zeros((res, res), dtype=np.float64) + maximal_array_index = np.int32(res) - 1 # Change that integer to a float, we know that our x, y are bounded # by [0, 1]. # We need this for combining with the x_pos and y_pos variables. - float_res = float64(res) + float_res = np.float64(res) # Pre-calculate this constant for use with the above inverse_cell_area = float_res * float_res @@ -106,8 +103,8 @@ def scatter( # Calculate the cell that this particle; use the 64 bit version of the # resolution as this is the same type as the positions - particle_cell_x = int32(float_res * x_pos) - particle_cell_y = int32(float_res * y_pos) + particle_cell_x = np.int32(float_res * x_pos) + particle_cell_y = np.int32(float_res * y_pos) if not ( particle_cell_x < 0 @@ -116,7 +113,7 @@ def scatter( or particle_cell_y >= maximal_array_index ): image[particle_cell_x, particle_cell_y] += ( - float64(mass) * inverse_cell_area + np.float64(mass) * inverse_cell_area ) return image @@ -124,14 +121,14 @@ def scatter( @jit(nopython=True, fastmath=True, parallel=True) def scatter_parallel( - x: float64, - y: float64, - m: float32, - h: float32, + x: np.float64, + y: np.float64, + m: np.float32, + h: np.float32, res: int, - box_x: float64 = 0.0, - box_y: float64 = 0.0, -) -> ndarray: + box_x: np.float64 = 0.0, + box_y: np.float64 = 0.0, +) -> np.ndarray: """ Parallel implementation of scatter @@ -142,34 +139,34 @@ def scatter_parallel( Parameters ---------- - x : np.array[float64] + x : np.array[np.float64] array of x-positions of the particles. Must be bounded by [0, 1]. - y : np.array[float64] + y : np.array[np.float64] array of y-positions of the particles. Must be bounded by [0, 1]. - m : np.array[float32] + m : np.array[np.float32] array of masses (or otherwise weights) of the particles - h : np.array[float32] + h : np.array[np.float32] array of smoothing lengths of the particles res : int the number of pixels along one axis, i.e. this returns a square of res * res. - box_x: float64 + box_x: np.float64 box size in x, in the same rescaled length units as x and y. Used for periodic wrapping. - box_y: float64 + box_y: np.float64 box size in y, in the same rescaled length units as x and y. Used for periodic wrapping. Returns ------- - np.array[float32, float32, float32] + np.array[np.float32, np.float32, np.float32] pixel grid of quantity See Also @@ -191,7 +188,7 @@ def scatter_parallel( number_of_particles = x.size core_particles = number_of_particles // NUM_THREADS - output = zeros((res, res), dtype=float64) + output = np.zeros((res, res), dtype=np.float64) for thread in prange(NUM_THREADS): # Left edge is easy, just start at 0 and go to 'final' diff --git a/swiftsimio/visualisation/projection_backends/kernels.py b/swiftsimio/visualisation/projection_backends/kernels.py index dc62cfb1..bc69990d 100644 --- a/swiftsimio/visualisation/projection_backends/kernels.py +++ b/swiftsimio/visualisation/projection_backends/kernels.py @@ -2,19 +2,16 @@ Projection kernels. """ -from typing import Union -from math import sqrt -from numpy import float64, float32, int32 - -from swiftsimio.accelerated import jit, NUM_THREADS, prange +import numpy as np +from swiftsimio.accelerated import jit # Taken from Dehnen & Aly 2012 -kernel_gamma = float32(1.897367) -kernel_constant = float32(7.0 / 3.14159) +kernel_gamma = np.float32(1.897367) +kernel_constant = np.float32(7.0 / 3.14159) @jit("float32(float32, float32)", nopython=True, fastmath=True) -def kernel_single_precision(r: float32, H: float32): +def kernel_single_precision(r: np.float32, H: np.float32): """ Single precision kernel implementation for swiftsimio. @@ -23,16 +20,16 @@ def kernel_single_precision(r: float32, H: float32): Parameters ---------- - r : float32 + r : np.float32 radius used in kernel computation - H : float32 + H : np.float32 kernel width (i.e. radius of compact support for the kernel) Returns ------- - float32 + np.float32 Contribution to the density by the particle See Also @@ -45,7 +42,7 @@ def kernel_single_precision(r: float32, H: float32): .. [1] Dehnen W., Aly H., 2012, MNRAS, 425, 1068 """ - kernel_constant = float32(2.22817109) + kernel_constant = np.float32(2.22817109) inverse_H = 1.0 / H ratio = r * inverse_H @@ -65,7 +62,7 @@ def kernel_single_precision(r: float32, H: float32): @jit("float64(float64, float64)", nopython=True, fastmath=True) -def kernel_double_precision(r: float64, H: float64): +def kernel_double_precision(r: np.float64, H: np.float64): """ Single precision kernel implementation for swiftsimio. @@ -74,15 +71,15 @@ def kernel_double_precision(r: float64, H: float64): Parameters ---------- - r : float32 + r : np.float32 radius used in kernel computation - H : float32 + H : np.float32 kernel width (i.e. radius of compact support for the kernel) Returns ------- - float32 + np.float32 Contribution to the density by the particle See Also @@ -95,7 +92,7 @@ def kernel_double_precision(r: float64, H: float64): .. [2] Dehnen W., Aly H., 2012, MNRAS, 425, 1068 """ - kernel_constant = float64(2.22817109) + kernel_constant = np.float64(2.22817109) inverse_H = 1.0 / H ratio = r * inverse_H diff --git a/swiftsimio/visualisation/projection_backends/reference.py b/swiftsimio/visualisation/projection_backends/reference.py index 942f3a10..affe96c5 100644 --- a/swiftsimio/visualisation/projection_backends/reference.py +++ b/swiftsimio/visualisation/projection_backends/reference.py @@ -5,11 +5,9 @@ Uses double precision. """ -from typing import Union -from math import sqrt, ceil +from math import sqrt from numpy import float32, float64, int32, zeros, ndarray -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends.kernels import ( kernel_double_precision as kernel, @@ -151,8 +149,9 @@ def scatter( particle_cell_x + cells_spanned + 1, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float64(cell_x) + 0.5) * pixel_width - float64( x_pos ) diff --git a/swiftsimio/visualisation/projection_backends/renormalised.py b/swiftsimio/visualisation/projection_backends/renormalised.py index 06fe392f..4ae6bb60 100644 --- a/swiftsimio/visualisation/projection_backends/renormalised.py +++ b/swiftsimio/visualisation/projection_backends/renormalised.py @@ -4,15 +4,11 @@ This version of the function is the same as `fast` but provides an explicit renormalisation of each kernel such that the mass is conserved up to floating point precision. -""" -""" The original smoothing code. This provides basic renormalisation -of the kernel on each call. +of the kernel on each call. """ - -from typing import Union from math import sqrt from numpy import float64, float32, int32, zeros, ndarray @@ -21,10 +17,7 @@ from swiftsimio.visualisation.projection_backends.kernels import ( kernel_single_precision as kernel, ) -from swiftsimio.visualisation.projection_backends.kernels import ( - kernel_constant, - kernel_gamma, -) +from swiftsimio.visualisation.projection_backends.kernels import kernel_gamma @jit(nopython=True, fastmath=True) @@ -166,8 +159,9 @@ def scatter( particle_cell_x - cells_spanned + 1, particle_cell_x + cells_spanned, ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float32(cell_x) + 0.5) * pixel_width - float32( x_pos ) @@ -198,8 +192,9 @@ def scatter( particle_cell_x + cells_spanned + 1, maximal_array_index + 1 ), ): - # The distance in x to our new favourite cell -- remember that our x, y - # are all in a box of [0, 1]; calculate the distance to the cell centre + # The distance in x to our new favourite cell -- remember that our + # x, y are all in a box of [0, 1]; calculate the distance to the + # cell centre distance_x = (float32(cell_x) + 0.5) * pixel_width - float32( x_pos ) diff --git a/swiftsimio/visualisation/projection_backends/subsampled.py b/swiftsimio/visualisation/projection_backends/subsampled.py index 7f859c86..fc95f08a 100644 --- a/swiftsimio/visualisation/projection_backends/subsampled.py +++ b/swiftsimio/visualisation/projection_backends/subsampled.py @@ -5,19 +5,14 @@ scales uses subsampling. Uses double precision. -""" -""" The original smoothing code. This provides a paranoid supersampling of the kernel. """ - -from typing import Union from math import sqrt, ceil from numpy import float32, float64, int32, zeros, ndarray -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends.kernels import ( kernel_double_precision as kernel, @@ -281,9 +276,10 @@ def scatter( ), ): float_cell_y = float64(cell_y) - # Now we subsample the pixels to get a more accurate determination - # of the kernel weight. We take the mean of the kernel evaluations - # within a given pixel and apply this as the true 'kernel evaluation'. + # Now we subsample the pixels to get a more accurate + # determination of the kernel weight. We take the mean of the + # kernel evaluations within a given pixel and apply this as + # the true 'kernel evaluation'. kernel_eval = float64(0.0) for subsample_x in range(0, subsample_factor): diff --git a/swiftsimio/visualisation/projection_backends/subsampled_extreme.py b/swiftsimio/visualisation/projection_backends/subsampled_extreme.py index 8e2cab5c..56d0c375 100644 --- a/swiftsimio/visualisation/projection_backends/subsampled_extreme.py +++ b/swiftsimio/visualisation/projection_backends/subsampled_extreme.py @@ -5,19 +5,14 @@ scales uses subsampling. Uses double precision. -""" -""" The original smoothing code. This provides a paranoid supersampling of the kernel. """ - -from typing import Union from math import sqrt, ceil from numpy import float32, float64, int32, zeros, ndarray -from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.accelerated import jit, NUM_THREADS, prange from swiftsimio.visualisation.projection_backends.kernels import ( kernel_double_precision as kernel, @@ -283,9 +278,10 @@ def scatter( ), ): float_cell_y = float64(cell_y) - # Now we subsample the pixels to get a more accurate determination - # of the kernel weight. We take the mean of the kernel evaluations - # within a given pixel and apply this as the true 'kernel evaluation'. + # Now we subsample the pixels to get a more accurate + # determination of the kernel weight. We take the mean of the + # kernel evaluations within a given pixel and apply this as + # the true 'kernel evaluation'. kernel_eval = float64(0.0) for subsample_x in range(0, subsample_factor): diff --git a/swiftsimio/visualisation/smoothing_length/generate.py b/swiftsimio/visualisation/smoothing_length/generate.py index 82eab392..e151fc0b 100644 --- a/swiftsimio/visualisation/smoothing_length/generate.py +++ b/swiftsimio/visualisation/smoothing_length/generate.py @@ -3,8 +3,7 @@ that do not usually carry a smoothing length field (e.g. dark matter). """ -from numpy import empty, float32 - +import numpy as np from swiftsimio import cosmo_array from swiftsimio.optional_packages import KDTree, TREE_AVAILABLE from swiftsimio._array_functions import _propagate_cosmo_array_attributes_to_result @@ -14,7 +13,7 @@ def generate_smoothing_lengths( coordinates: cosmo_array, boxsize: cosmo_array, - kernel_gamma: float32, + kernel_gamma: np.float32, neighbours=32, speedup_fac=2, dimension=3, @@ -28,7 +27,7 @@ def generate_smoothing_lengths( a cosmo_array that gives the co-ordinates of all particles boxsize : cosmo_array the size of the box (3D) - kernel_gamma : float32 + kernel_gamma : np.float32 the kernel gamma of the kernel being used neighbours : int, optional the number of neighbours to encompass @@ -58,7 +57,7 @@ def generate_smoothing_lengths( tree = KDTree(coordinates.value, boxsize=boxsize.to(coordinates.units).value) - smoothing_lengths = empty(number_of_parts, dtype=float32) + smoothing_lengths = np.empty(number_of_parts, dtype=np.float32) smoothing_lengths[-1] = -0.1 # Include speedup_fac stuff here: From 4f7d3d7564e85d554144f4f0cbb4d3e847119ca6 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 14 Feb 2025 23:16:45 +0000 Subject: [PATCH 095/125] Switch on intersphinx, why not. --- docs/source/conf.py | 11 +++++++++++ 1 file changed, 11 insertions(+) diff --git a/docs/source/conf.py b/docs/source/conf.py index 3088ca6d..d0f61711 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -40,6 +40,7 @@ "sphinx.ext.napoleon", "sphinx.ext.mathjax", "sphinx.ext.autosummary", + "sphinx.ext.intersphinx", ] # Add any paths that contain templates here, relative to this directory. @@ -98,3 +99,13 @@ def run_apidoc(_): def setup(app): app.connect("builder-inited", run_apidoc) + + +intersphinx_mapping = dict( + numpy=("https://numpy.org/doc/stable/", None), + numba=("https://numba.readthedocs.io/en/stable/", None), + unyt=("https://unyt.readthedocs.io/en/stable/", None), + scipy=("https://docs.scipy.org/doc/scipy/", None), + swiftgalaxy=("https://swiftsimio.readthedocs.io/en/latest/", None), + velociraptor=("https://velociraptor-python.readthedocs.io/en/latest/", None), +) From bc3ff27697b5e6b964bea47a92fe2ddc37474586 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 15 Feb 2025 00:11:30 +0000 Subject: [PATCH 096/125] Reorganize docs a little bit and flesh out the cosmo_array description. --- docs/source/cosmo_array/index.rst | 82 +++++++++++++++++++++++++++ docs/source/getting_started/index.rst | 4 +- docs/source/index.rst | 2 + docs/source/loading_data/index.rst | 76 ------------------------- docs/source/soap/index.rst | 41 ++++++++++++++ 5 files changed, 127 insertions(+), 78 deletions(-) create mode 100644 docs/source/cosmo_array/index.rst create mode 100644 docs/source/soap/index.rst diff --git a/docs/source/cosmo_array/index.rst b/docs/source/cosmo_array/index.rst new file mode 100644 index 00000000..5a73723a --- /dev/null +++ b/docs/source/cosmo_array/index.rst @@ -0,0 +1,82 @@ +The ``cosmo_array`` +=================== + +:mod:`swiftsimio` uses a customized class based on the :class:`~unyt.array.unyt_array` +to store data arrays. The :class:`~swiftsimio.objects.cosmo_array` has all of the same +functionality as the :class:`~unyt.array.unyt_array`, but also adds information about +data transformation between physical and comoving coordinates, and descriptive metadata. + +For instance, should you ever need to know what a dataset represents, you can +ask for a description by accessing the ``name`` attribute: + +.. code-block:: python + + print(rho_gas.name) + +which will output ``Co-moving mass densities of the particles``. + +The cosmology information is stored in three attributes: + + + ``comoving`` + + ``cosmo_factor`` + + ``valid_transform`` + +The ``comoving`` attribute specifies whether the array is a physical (``True``) or +comoving (``False``) quantity, while the ``cosmo_factor`` stores the expression needed +to convert back and forth between comoving and physical quantities and the value of +the scale factor. The conversion factors can be accessed like this: + +.. code-block:: python + + # Conversion factor to make the densities a physical quantity + print(rho_gas.cosmo_factor.a_factor) + physical_rho_gas = rho_gas.cosmo_factor.a_factor * rho_gas + + # Symbolic scale-factor expression + print(rho_gas.cosmo_factor.expr) + +which will output ``132651.002785671`` and ``a**(-3.0)``. Converting an array to/from physical/comoving +is done with the :meth:`~swiftsimio.objects.cosmo_array.to_physical`, :meth:`~swiftsimio.objects.cosmo_array.to_comoving`, :meth:`~swiftsimio.objects.cosmo_array.convert_to_physical` and :meth:`~swiftsimio.objects.cosmo_array.to_comoving` methods, for instance: + +.. code-block:: python + + physical_rho_gas = rho_gas.to_physical() + + # Convert in-place + rho_gas.convert_to_physical() + +The ``valid_transform`` is a boolean flag that is set to ``False`` for some arrays that don't make sense to convert to comoving. + +:class:`~swiftsimio.objects.cosmo_array` supports array arithmetic and the entire :mod:`numpy` range of functions. Attempting to combine arrays (e.g. by addition) will validate the cosmology information first. The implementation is designed to be permissive: it will only raise exceptions when a genuinely invalid combination is encountered, but is tolerant of missing cosmology information. When one argument in a relevant operation (like addition) is not a :class:`~swiftsimio.objects.cosmo_array` the attributes of the :class:`~swiftsimio.objects.cosmo_array` will be assumed for both arguments. In such cases a warning is produced stating that this assumption has been made. + +.. note:: + + :class:`~swiftsimio.objects.cosmo_array` and the related :class:`~swiftsimio.objects.cosmo_quantity` are now intended to support all :mod:`numpy` functions, propagating units and cosmology information correctly through mathematical operations. Try making a histogram with weights and ``density=True`` with :func:`numpy.histogram`! There are a large number of functions and a very large number of possible parameter combinations, so some corner cases may have been missed in development. Please report any errors or unexpected results using github issues or other channels so that they can be fixed. Currently :mod:`scipy` functions are not supported (although some might "just work"). Requests to support specific functions can be accommodated. + +To make the most of the utility offered by the :class:`~swiftsimio.objects.cosmo_array` class, it is helpful to know how to create your own. A good template for this looks like: + +.. code-block:: python + + import unyt as u + from swiftsimio.objects import cosmo_array, cosmo_factor + from swiftsimio.objects import a # the scale factor represented by a sympy expression + + my_cosmo_array = cosmo_array( + [1, 2, 3], + u.Mpc, + comoving=True, + cosmo_factor=cosmo_factor(a**1, 1.0) # 1.0 is the current scale factor + ) + # consider getting the scale factor from metadata when applicable, like: + # data.metadata.a + +There is also a very similar :class:`~swiftsimio.objects.cosmo_quantity` class designed for scalar values, +analogous to the :class:`~unyt.array.unyt_quantity`. You may encounter this being returned by :mod:`numpy` functions. Cosmology-aware scalar values can be initialized similarly: + +.. code-block:: python + + import unyt as u + from swiftsimio.objects import cosmo_quantity, cosmo_factor, a + + my_cosmo_quantity = cosmo_quantity(2, u.Mpc, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) + diff --git a/docs/source/getting_started/index.rst b/docs/source/getting_started/index.rst index 9715fdaa..a4366936 100644 --- a/docs/source/getting_started/index.rst +++ b/docs/source/getting_started/index.rst @@ -109,8 +109,8 @@ In the above it's important to note the following: + Only the density and temperatures (corresponding to the ``PartType0/Densities`` and ``PartType0/Temperatures``) datasets are read in. + That data is only read in once the - :meth:`swiftsimio.objects.cosmo_array.convert_to_cgs` method is called. -+ :meth:`swiftsimio.objects.cosmo_array.convert_to_cgs` converts data in-place; + :meth:`~swiftsimio.objects.cosmo_array.convert_to_cgs` method is called. ++ :meth:`~swiftsimio.objects.cosmo_array.convert_to_cgs` converts data in-place; i.e. it returns `None`. + The data is cached and not re-read in when ``plt.scatter`` is called. diff --git a/docs/source/index.rst b/docs/source/index.rst index f7ecc438..9470b8ea 100644 --- a/docs/source/index.rst +++ b/docs/source/index.rst @@ -17,8 +17,10 @@ snapshots to enable partial reading. getting_started/index loading_data/index + cosmo_array/index masking/index visualisation/index + soap/index velociraptor/index creating_initial_conditions/index statistics/index diff --git a/docs/source/loading_data/index.rst b/docs/source/loading_data/index.rst index e62f7b30..2a901069 100644 --- a/docs/source/loading_data/index.rst +++ b/docs/source/loading_data/index.rst @@ -247,47 +247,6 @@ Then, to access individual columns (in this case element abundances): data.gas.element_mass_fractions.silicon -Non-unyt properties -------------------- - -Each data array has some custom properties that are not present within the base -:obj:`unyt.unyt_array` class. We create our own version of this in -:obj:`swiftsimio.objects.cosmo_array`, which allows each dataset to contain -its own cosmology and name properties. - -For instance, should you ever need to know what a dataset represents, you can -ask for a description: - -.. code-block:: python - - print(rho_gas.name) - -which will output ``Co-moving mass densities of the particles``. They include -scale-factor information, too, through the ``cosmo_factor`` object, - -.. code-block:: python - - # Conversion factor to make the densities a physical quantity - print(rho_gas.cosmo_factor.a_factor) - physical_rho_gas = rho_gas.cosmo_factor.a_factor * rho_gas - - # Symbolic scale-factor expression - print(rho_gas.cosmo_factor.expr) - -which will output ``132651.002785671`` and ``a**(-3.0)``. This is an easy way -to convert your co-moving values to physical ones. - -An even easier way to convert your properties to physical is to use the -built-in ``to_physical`` and ``convert_to_physical`` methods, as follows: - -.. code-block:: python - - physical_rho_gas = rho_gas.to_physical() - - # Convert in-place - rho_gas.convert_to_physical() - - User-defined particle types --------------------------- @@ -310,38 +269,3 @@ in SWIFT will be automatically read. "extra_test.hdf5", ) - -Halo Catalogues ---------------- - -SWIFT-compatible halo catalogues, such as those written with SOAP, can be -loaded entirely transparently with ``swiftsimio``. It is generally possible -to use all of the functionality (masking, visualisation, etc.) that is used -with snapshots with these files, assuming the files conform to the -correct metadata standard. - -An example SOAP file is available at -``http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/soap_example.hdf5`` - -You can load SOAP files as follows: - -.. code-block:: python - - from swiftsimio import load - - catalogue = load("soap_example.hdf5") - - print(catalogue.spherical_overdensity_200_mean.total_mass) - - # >>> [ 591. 328.5 361. 553. 530. 507. 795. - # 574. 489.5 233.75 0. 1406. 367.5 2308. - # ... - # 0. 534. 0. 191.75 1450. 600. 290. ] 10000000000.0*Msun (Physical) - -What's going on here? Under the hood, ``swiftsimio`` has a discrimination function -between different metadata types, based upon a property stored in the HDF5 file, -``Header/OutputType``. If this is set to ``FullVolume``, we have a snapshot, -and use the :obj:`swiftsimio.metadata.objects.SWIFTSnapshotMetadata` -class. If it is ``SOAP``, we use -:obj:`swiftsimio.metadata.objects.SWIFTSOAPMetadata`, which instructs -``swiftsimio`` to read slightly different properties from the HDF5 file. diff --git a/docs/source/soap/index.rst b/docs/source/soap/index.rst new file mode 100644 index 00000000..c143c9ec --- /dev/null +++ b/docs/source/soap/index.rst @@ -0,0 +1,41 @@ +Halo Catalogues & SOAP integration +================================== + +SWIFT-compatible halo catalogues, such as those written with SOAP, can be +loaded entirely transparently with ``swiftsimio``. It is generally possible +to use all of the functionality (masking, visualisation, etc.) that is used +with snapshots with these files, assuming the files conform to the +correct metadata standard. + +An example SOAP file is available at +``http://virgodb.cosma.dur.ac.uk/swift-webstorage/IOExamples/soap_example.hdf5`` + +You can load SOAP files as follows: + +.. code-block:: python + + from swiftsimio import load + + catalogue = load("soap_example.hdf5") + + print(catalogue.spherical_overdensity_200_mean.total_mass) + + # >>> [ 591. 328.5 361. 553. 530. 507. 795. + # 574. 489.5 233.75 0. 1406. 367.5 2308. + # ... + # 0. 534. 0. 191.75 1450. 600. 290. ] 10000000000.0*Msun (Physical) + +What's going on here? Under the hood, ``swiftsimio`` has a discrimination function +between different metadata types, based upon a property stored in the HDF5 file, +``Header/OutputType``. If this is set to ``FullVolume``, we have a snapshot, +and use the :obj:`swiftsimio.metadata.objects.SWIFTSnapshotMetadata` +class. If it is ``SOAP``, we use +:obj:`swiftsimio.metadata.objects.SWIFTSOAPMetadata`, which instructs +``swiftsimio`` to read slightly different properties from the HDF5 file. + +swiftgalaxy +----------- + +The :mod:`swiftgalaxy` companion package to :mod:`swiftsimio` offers further integration with halo catalogues in SOAP, Caesar and Velociraptor formats (so far). It greatly simplifies efficient loading of particles belonging to an object from a catalogue, and additional tools that are useful when working with a galaxy or other localized collection of particles. Refer to the `swiftgalaxy documentation`_ for details. + +.. _swiftgalaxy documentation: https://swiftgalaxy.readthedocs.io/en/latest/ From 53d6fe1093ceefc8ed1d0b81dc437194d219a71b Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Sat, 15 Feb 2025 00:12:04 +0000 Subject: [PATCH 097/125] Run black. --- swiftsimio/_array_functions.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 6132f926..5b113d7d 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -729,18 +729,18 @@ def _arctan2_cosmo_factor( f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor( From 34d9497c37cee52d5f1c8982a27539765c8d38e8 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 18 Feb 2025 09:45:14 +0000 Subject: [PATCH 098/125] Add new vis backend submodules (from reorganising) to installed modules. --- pyproject.toml | 2 ++ 1 file changed, 2 insertions(+) diff --git a/pyproject.toml b/pyproject.toml index 83482749..b5b44589 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -16,6 +16,8 @@ packages = [ "swiftsimio.visualisation", "swiftsimio.visualisation.projection_backends", "swiftsimio.visualisation.slice_backends", + "swiftsimio.visualisation.volume_render_backends", + "swiftsimio.visualisation.ray_trace_backends", "swiftsimio.visualisation.tools", "swiftsimio.visualisation.smoothing_length", ] From 14a0f7301fac75898ffbdfd9463a0d5196dd83c6 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 11 Mar 2025 17:18:47 +0000 Subject: [PATCH 099/125] Unyt fixed ftt->fft typo, update to match. --- swiftsimio/_array_functions.py | 28 ++++++++++++++-------------- 1 file changed, 14 insertions(+), 14 deletions(-) diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 5b113d7d..fdaabde0 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -40,20 +40,20 @@ stack as unyt_stack, around as unyt_around, block as unyt_block, - ftt_fft as unyt_fft_fft, # unyt has a copy-pasted typo fft -> ftt - ftt_fft2 as unyt_fft_fft2, - ftt_fftn as unyt_fft_fftn, - ftt_hfft as unyt_fft_hfft, - ftt_rfft as unyt_fft_rfft, - ftt_rfft2 as unyt_fft_rfft2, - ftt_rfftn as unyt_fft_rfftn, - ftt_ifft as unyt_fft_ifft, - ftt_ifft2 as unyt_fft_ifft2, - ftt_ifftn as unyt_fft_ifftn, - ftt_ihfft as unyt_fft_ihfft, - ftt_irfft as unyt_fft_irfft, - ftt_irfft2 as unyt_fft_irfft2, - ftt_irfftn as unyt_fft_irfftn, + fft_fft as unyt_fft_fft, + fft_fft2 as unyt_fft_fft2, + fft_fftn as unyt_fft_fftn, + fft_hfft as unyt_fft_hfft, + fft_rfft as unyt_fft_rfft, + fft_rfft2 as unyt_fft_rfft2, + fft_rfftn as unyt_fft_rfftn, + fft_ifft as unyt_fft_ifft, + fft_ifft2 as unyt_fft_ifft2, + fft_ifftn as unyt_fft_ifftn, + fft_ihfft as unyt_fft_ihfft, + fft_irfft as unyt_fft_irfft, + fft_irfft2 as unyt_fft_irfft2, + fft_irfftn as unyt_fft_irfftn, fft_fftshift as unyt_fft_fftshift, fft_ifftshift as unyt_fft_ifftshift, sort_complex as unyt_sort_complex, From feae3670bb13ee7560e5c93fb63646172f17608d Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 11 Mar 2025 17:24:24 +0000 Subject: [PATCH 100/125] cache@v2 in workflows deprecated, update so that we can run CI. --- .github/workflows/pytest.yml | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/.github/workflows/pytest.yml b/.github/workflows/pytest.yml index aab5bc05..309cfa2a 100644 --- a/.github/workflows/pytest.yml +++ b/.github/workflows/pytest.yml @@ -16,7 +16,7 @@ jobs: with: python-version: ${{ matrix.python-version }} - name: Cache pip - uses: actions/cache@v2 + uses: actions/cache@v4 with: # This path is specific to Ubuntu path: ~/.cache/pip From fdd0c76742676fae013f8e58f147850323e8e443 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Fri, 14 Mar 2025 14:00:27 +0000 Subject: [PATCH 101/125] Point to unyt on pypi now that 3.0.4 is released. --- pyproject.toml | 4 ++-- requirements.txt | 2 +- 2 files changed, 3 insertions(+), 3 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index b5b44589..d357269c 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -39,9 +39,9 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "numpy", + "numpy>=2.1.0", "h5py", - "unyt>=3.0.2", + "unyt>=3.0.4", "numba>=0.50.0", ] diff --git a/requirements.txt b/requirements.txt index 2d58369a..7e7a9182 100644 --- a/requirements.txt +++ b/requirements.txt @@ -2,4 +2,4 @@ astropy numpy h5py numba -git+https://github.com/yt-project/unyt.git +unyt From db688b07c4aee4f65fb7a30fcdcd3ad6557546d7 Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 18 Mar 2025 18:11:26 +0000 Subject: [PATCH 102/125] Add support for multiplying a cosmo_array by a bare unit object from unyt. --- pyproject.toml | 2 +- swiftsimio/_array_functions.py | 6 +- swiftsimio/objects.py | 116 ++++++++++++++++++++++++- tests/test_cosmo_array.py | 154 ++++++++++++++++++++++++--------- 4 files changed, 230 insertions(+), 48 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index d357269c..940aac0f 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -41,7 +41,7 @@ classifiers = [ dependencies = [ "numpy>=2.1.0", "h5py", - "unyt>=3.0.4", + "unyt@git+https://github.com/kyleaoman/unyt.git#egg=fix_block_wrapper", "numba>=0.50.0", ] diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index fdaabde0..3995d5d4 100644 --- a/swiftsimio/_array_functions.py +++ b/swiftsimio/_array_functions.py @@ -729,18 +729,18 @@ def _arctan2_cosmo_factor( f" provided cosmo_factor ({cf2}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf2.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf2.scale_factor) elif (cf1 is not None) and (cf2 is None): warnings.warn( f"Mixing arguments with and without cosmo_factors, continuing assuming" f" provided cosmo_factor ({cf1}) for all arguments.", RuntimeWarning, ) - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) elif (cf1 is not None) and (cf2 is not None) and (cf1 != cf2): raise ValueError(f"Arguments have cosmo_factors that differ: {cf1} and {cf2}.") elif (cf1 is not None) and (cf2 is not None) and (cf1 == cf2): - return objects.cosmo_factor(objects.a ** 0, scale_factor=cf1.scale_factor) + return objects.cosmo_factor(objects.a**0, scale_factor=cf1.scale_factor) def _comparison_cosmo_factor( diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index 6f5a5d04..be02f02f 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -644,7 +644,7 @@ def __pow__(self, p: float) -> "cosmo_factor": """ if self.expr is None: return cosmo_factor(expr=None, scale_factor=self.scale_factor) - 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: "cosmo_factor") -> bool: """ @@ -1626,6 +1626,120 @@ def __array_function__( function_to_invoke = func._implementation return function_to_invoke(*args, **kwargs) + def __mul__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Multiply this :class:`~swiftsimio.objects.cosmo_array`. + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to multiply with this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the multiplication. + """ + if isinstance(b, unyt.unit_object.Unit): + retval = self.__copy__() + retval.units = retval.units * b + return retval + else: + return super().__mul__(b) + + def __rmul__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Multiply this :class:`~swiftsimio.objects.cosmo_array` (as the right argument). + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + .. note:: + + This function is never called when `b` is a :class:`unyt.unit_object.Unit` + because :mod:`unyt` handles the operation. This results in a silent demotion + to a :class:`unyt.array.unyt_array`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to multiply with this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the multiplication. + """ + if isinstance(b, unyt.unit_object.Unit): + return self.__mul__(b) + else: + return super().__rmul__(b) + + def __truediv__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Divide this :class:`~swiftsimio.objects.cosmo_array`. + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to divide this one by. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the division. + """ + if isinstance(b, unyt.unit_object.Unit): + return self.__mul__(1 / b) + else: + return super().__truediv__(b) + + def __rtruediv__( + self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit] + ) -> "cosmo_array": + """ + Divide this :class:`~swiftsimio.objects.cosmo_array` (as the right argument). + + We delegate most cases to :mod:`unyt`, but we need to handle the case where the + second argument is a :class:`~unyt.unit_object.Unit`. + + .. note:: + + This function is never called when `b` is a :class:`unyt.unit_object.Unit` + because :mod:`unyt` handles the operation. This results in a silent demotion + to a :class:`unyt.array.unyt_array`. + + Parameters + ---------- + b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \ + :class:`~unyt.unit_object.Unit` + The object to divide by this one. + + Returns + ------- + out : swiftsimio.objects.cosmo_array + The result of the division. + """ + if isinstance(b, unyt.unit_object.Unit): + return (1 / self).__mul__(b) + else: + return super().__rtruediv__(b) + class cosmo_quantity(cosmo_array, unyt_quantity): """ diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 008638e3..8998b724 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -28,16 +28,14 @@ def ca(x, unit=u.Mpc): """ Helper for our tests: turn an array into a cosmo_array. """ - 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 cq(x, unit=u.Mpc): """ Helper for our tests: turn a scalar into a cosmo_quantity. """ - return cosmo_quantity( - x, unit, comoving=False, cosmo_factor=cosmo_factor(a ** 1, 0.5) - ) + return cosmo_quantity(x, unit, comoving=False, cosmo_factor=cosmo_factor(a**1, 0.5)) def arg_to_ua(arg): @@ -115,7 +113,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") @@ -129,7 +127,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") @@ -142,7 +140,7 @@ 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") @@ -158,7 +156,7 @@ 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") @@ -178,14 +176,14 @@ def test_init_from_list_of_cosmo_arrays(self): [1], units=u.Mpc, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1), + 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 @@ -573,7 +571,7 @@ def test_explicitly_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), ), ), ), @@ -584,13 +582,13 @@ def test_explicitly_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), ), ), ), @@ -602,19 +600,19 @@ def test_explicitly_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), ), ], ), @@ -626,7 +624,7 @@ def test_explicitly_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]), ), @@ -650,19 +648,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] @@ -709,9 +707,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): @@ -719,9 +717,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): @@ -729,9 +727,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 = { @@ -743,9 +741,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), ] ), ): @@ -778,7 +776,7 @@ def test_dot(self): res = ca(np.arange(3)).dot(ca(np.arange(3))) assert isinstance(res, cosmo_quantity) assert res.comoving is False - assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5) + assert res.cosmo_factor == cosmo_factor(a**2, 0.5) assert res.valid_transform is True @@ -812,12 +810,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 def test_round(self): @@ -828,13 +826,13 @@ def test_round(self): 1.03, u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), valid_transform=True, ) res = round(cq) assert res.value == 1.0 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 def test_scalar_return_func(self): @@ -846,7 +844,7 @@ def test_scalar_return_func(self): np.arange(3), u.m, comoving=False, - cosmo_factor=cosmo_factor(a ** 1, 1.0), + cosmo_factor=cosmo_factor(a**1, 1.0), valid_transform=True, ) res = np.min(ca) @@ -861,12 +859,12 @@ 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 @@ -882,7 +880,7 @@ def test_copy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) copy_arr = copy(arr) @@ -898,7 +896,7 @@ def test_deepcopy(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) copy_arr = deepcopy(arr) @@ -914,7 +912,7 @@ def test_to_cgs(self): units = u.Mpc arr = cosmo_array( u.unyt_array(np.ones(5), units=units), - cosmo_factor=cosmo_factor(a ** 1, 1), + cosmo_factor=cosmo_factor(a**1, 1), comoving=False, ) cgs_arr = arr.in_cgs() @@ -922,3 +920,73 @@ def test_to_cgs(self): assert cgs_arr.units == u.cm assert cgs_arr.cosmo_factor == arr.cosmo_factor assert cgs_arr.comoving == arr.comoving + + +class TestMultiplicationByUnyt: + + def test_multiplication_by_unyt(self): + """ + We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as + for `cosmo_array(...) * u.Mpc`. + + Right-sided multiplication & division can't be supported without upstream + changes in unyt, see `test_rmultiplication_by_unyt`. + """ + ca = cosmo_array( + np.ones(3), u.Mpc, comoving=True, cosmo_factor=cosmo_factor(a**1, 1.0) + ) + # required so that can test right-sided division with the same assertions: + assert np.allclose(ca.to_value(ca.units), 1) + # the reference result: + multiplied_by_quantity = ca * (1 * u.Mpc) # parentheses very important here + # get the same result twice through left-sided multiplication and division: + lmultiplied_by_unyt = ca * u.Mpc + ldivided_by_unyt = ca / u.Mpc**-1 + + for multiplied_by_unyt in ( + lmultiplied_by_unyt, + ldivided_by_unyt, + ): + assert isinstance(multiplied_by_quantity, cosmo_array) + assert isinstance(multiplied_by_unyt, cosmo_array) + assert np.allclose( + multiplied_by_unyt.to_value(multiplied_by_quantity.units), + multiplied_by_quantity.to_value(multiplied_by_quantity.units), + ) + + @pytest.mark.xfail + def test_rmultiplication_by_unyt(self): + """ + We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as + for `cosmo_array(...) * u.Mpc`. + + But unyt will call it's own __mul__ before we get a chance to use our __rmul__ + when the cosmo_array is the right-hand argument. + + We can't handle this case without upstream changes in unyt, so this test is marked + to xfail. + + If this is fixed in the future this test will pass and can be merged with + `test_multiplication_by_unyt` to tidy up. + """ + ca = cosmo_array( + np.ones(3), u.Mpc, comoving=True, cosmo_factor=cosmo_factor(a**1, 1.0) + ) + # required so that can test right-sided division with the same assertions: + assert np.allclose(ca.to_value(ca.units), 1) + # the reference result: + multiplied_by_quantity = ca * (1 * u.Mpc) # parentheses very important here + # get 2x the same result through right-sided multiplication and division: + rmultiplied_by_unyt = u.Mpc * ca + rdivided_by_unyt = u.Mpc**2 / ca + + for multiplied_by_unyt in ( + rmultiplied_by_unyt, + rdivided_by_unyt, + ): + assert isinstance(multiplied_by_quantity, cosmo_array) + assert isinstance(multiplied_by_unyt, cosmo_array) + assert np.allclose( + multiplied_by_unyt.to_value(multiplied_by_quantity.units), + multiplied_by_quantity.to_value(multiplied_by_quantity.units), + ) From ab722156bfa959fe32e8abe8abd3ddd3c45617fd Mon Sep 17 00:00:00 2001 From: Kyle Oman Date: Tue, 18 Mar 2025 18:18:07 +0000 Subject: [PATCH 103/125] Re-point unyt to pypi (tests will fail). --- panels_added.png | Bin 0 -> 6582 bytes projected.png | Bin 0 -> 6582 bytes pyproject.toml | 2 +- stacked.png | Bin 0 -> 649355 bytes swiftsimio/_array_functions.py | 6 +- swiftsimio/objects.py | 2 +- tests/test_cosmo_array.py | 103 ++++++++++++++++----------------- 7 files changed, 54 insertions(+), 59 deletions(-) create mode 100644 panels_added.png create mode 100644 projected.png create mode 100644 stacked.png diff --git a/panels_added.png b/panels_added.png new file mode 100644 index 0000000000000000000000000000000000000000..add23b354d24ebf3bfc04f1d7101b01615f60410 GIT binary patch literal 6582 zcmeAS@N?(olHy`uVBq!ia0y~yU;#2&7&zE~RK2WrGmx?@ag8Vm&QB{TPb^Ah@J%c! z$jL9s$xKoxOD!tS%+E8{v(z)xQOGDMDJZtm*UyD3(917M*S{?JP7!DbXMsm#F#`jC zA_z0CuwS0Zz#ysY>EaktG3V_KN8Sblh64u-bwB5Gi?*zc@OUw|XEaktG3V_KN8Sblh64u-bwB5Gi?*zc@OUw|X=2.1.0", "h5py", - "unyt@git+https://github.com/kyleaoman/unyt.git#egg=fix_block_wrapper", + "unyt>=3.0.4", "numba>=0.50.0", ] diff --git a/stacked.png b/stacked.png new file mode 100644 index 0000000000000000000000000000000000000000..e2697fe14096d4e59988f4b4626e86a10b6c674f GIT binary patch literal 649355 zcmV)LK)Jt(P)V>bE;%kSEFfrfbZ~PzFE4FjbZ~5MbZlv2E^l&Y zFVh+BApigX32;bRa{vGjVE_ORVF9Q=r)dBI002ouK~#90)ctMNB)h6Lih&|)?Q=Wf z_PQ1BrhRv12>*~klPMILPt`iFhc!l3<`WtEmVgNe>Hp*Z_5U>SWFpdae44IjA|mRa zCelQt*?6WVsv_ojyo?TiV;NkIuIj}O`GoQ}oBu?#*($?ywWCuXih5Ng=k<2q8f9qJ z7VL+YKl`m+L~jwDW#MbyMNtu(<7Zfwcctt2q&jW(eSEGmayw?;o`*K-ZmS%jy0H(@ zmZ}Yk-`DuX`?gPOwD%IR%hnuTHs?@8)j#L_h={f>*JW~DU9Q&Ub=~sM>y}@w%dhK} zm-dm*>zZ@kY2=%jsG2nRCnC~KrJKs#o4l>b@B5N_KC}5LBGPo_*ZH@nZYuXI(&~qBk!B%>$DD2-E+-1XvZOcv<25*TlvjtP@{(0UVqdx-#4T^m|t@K)VQ|jj?OjP znu)e9s`cDP_DO5ne7Ch8Ye>~%o%OjEdY|in+mQxq4Q!vZ%jVh{^jC9VJkN=C&2^?h zyPvE~ds|}jjGZDJdD<*()u@^ zd+!s^{b~KTOiPaI__T(k*h8M2_Vzg!o!WtE6AWfk`q~4D4=34Sw!Z-_#Uf90r@0Jq zTCOlei4Kz-0=y&g;hsj>gYzr21J_SQ3wZu;%oO`>nWj|qXOmxuhT!^R*^$eMJ`}V= zWZ#dykmZQ~B8f%%s{%6I#s_y^<1YS7NEJeu0#GW*1%VM>|5;ANbYr+ z!P|j!2oRcEAzX5n&>^g4;j-YpRxG8y-nJG zyR_e3^xdVsyU4qX_C5e&cltYoe7jVnwYi3o&MNd*#hQw}RQjdTf3@-bmr8%B^p{A# z6zB=cdvgezj54nI{aXKCv9WQoc+lf}W*P`Sb>_1^DXBjY-3bLVo%9FF2HYig9IE$LJF4Db# zChXD+80T%flm6`d&njOZ$d1q#lO?G8F-3kCdXVnNgB zQzkLD@o6Fp_%i!BCt8~ z5KG(TVJRPnExu!2Wt~u0B@8Kv-xbTkZ^cKk&8PmQ zzS8I4tXr+5o*xOnyxOJmj4w zo!a~pnJU_AD%}=6{3fZA$TWFm80cgNzY6=>`W&M+AM_=U8IEVAh!xm-_%HH!*$ez8 zxmTQ_jCQE=eq13H_zkjSFvlQ6@_rC%Qz<_^^v6WBTiWYtCOVj5{_H*e9?BdW_^?gc zw!6HqV%-7+>zpj#Bfh6}*6j*5S1^X$jN=@<7&@hEENr&_wUA1oai|82g_ksVu2q|5 zx?uH_#a{^__;)EdFG_4r?Z+%Of$L8>{;|NHXK|+VwLt!5TTx#AQ(m&_q;V$#!c07b zFUYejcb3+WP9=)EMc#{^1lkFjf+a;_kN-RK_EWccYrCOP5sMfsEB z4Ax?LuJszQDY%4!*M*`cU58qGDJ=vP2y?D_n@YOzZYb<}OCA6>c%1~e76;kI`ecA1 zdzZ@Y|1D2|RN>5jtzdESp(vQ;j|%iu`~Fkk%VqKB9jkP9KS@bfg}N{h;&O$MkJd!c zPg(R+JAfmk!6(bLjNh!;$lD6=?Q8Tz!$pTM%=?gKKWr<`=~P0gE1y*M{aSrb&y)(r zHWzKjd6@h-DFk86^khLV#v%PU>qRo2>i*MxyU||yUho(8ctiQiQ`0U@JN#{O+jxG7 z%KM>7l_eNGp=9HuHDASKl}u2M#||Jv)Wpd_Cq>)>{0&D$VTUI!tPwr|fVR^} zds)?F=$YW42G)ObO?(U{xea{Vuh?v#UrqHZxD*sjsL~I_gaa7Z1PCh#C+6`n$IVXN z{F>pikRZ=vmbJBonGrVb()2I)$so4P$PF=DePt5jW&K^Cw-uf1rkj=#U7)*YvI3{~ zoalm!&T06}bcnFCF8?$r=O-zK{Hy)<#LU0Uj)D*tNN$h0aa6ZW;iUs17sy))=n40# zXrC^0_*)@(rRB!;&v%!S0Qw=+7wj_%PF{JUJH=eMkd$c4)X!R>sg(?n*E5~Z=lTKK z`8SG~V=9#=KUpyFozvdsRsNIKIr9Gi8DKk1_Qn-LH@iQEMjZrbp*>KpU$3sOg>DXj zRixid%z9YvMFfEGJRYPf7H0Y`BA4{mL%4MMarY+gyG>X47H;jUZ`m5Pj_2fKOF%}; zl&p74?D{^j^57F)ilLi(x5~roANKFY|3!BG^G7_7)}YnpY_9yLHe&oi4KcR0km+X<@(V{V+f&hPc%z)ixe~%w;H??Ab6`A+_$NpW>c_rkSG#DDYqYK zIfzVa9E`>$jbUNAokowC$4cfY_ZQ9CEKh~%02kV z{Fr-uk)ajuG}&daTiP_Aj?`_qEw8^dKG*kwOkD1I_u&TZO22!$FXG<{tKH2lv zmgAHHf&x7|Lj?LAxIl&j~Z_p9p{e z2pKT4FytxR`JH+DD=28;`aP^C1i@L5Ih;KVVPk&n5<+`u@+u($g+Hu2gpwh^m68CD zCD9y=Gtk?lQHVN%5of8#$I$7dyU0ef!GqHrHF=Wj6P^ULtB&wvYr*T<+{ zL0fAWEgE^wq3tf3gVol&&FCsG1lni}(c0y59`+3_Mq7RtOgrYaRfcl=u@4p&ozVwZ z0C@Wi+CclC=|BhCZ)?9sT|O=x3jxbw>mi+gPG@-wdKO0pc{vr}S8OS%ytzve0fK%~ z<^XW3)_(d|^OHV~>2C=^>;STp2sp@=jZKbe$fSbwza{KAqb8mvqJBw$lb-XiQEE;gJJ1y6L*@~IOiMS zceYWN;R;4CK=5n){pBp+rgw%9S#>U|Me$DukSG-lBMCAhnb`PPxj(tdMvPLqn1{9g zIPU!bf-50{o@K&+2GQv}f(svs;h*bsxi5ndr=*(mXIU&^-jpwm`2+aIG5oY#kT7uI ziLm@4Az9#8HxP{HCikiw*y7((Sb0~+$r>yx2o?se;mO)#t>7X2@h9i%eE4qSCuaYZ9LZYr?Aa1LNu)+1c^eToeSJB%sWDsbXPi zG=_=*0mie&-rj+q{K2;?^e!3uO7 z6Zc_0GnMcjQa_F`*Ulr|C=~t+1WR56CL8$O_5!XwtkVkcdkSO~GEOcdgN-!~uB1|N z@245>Ve-SWO)~STt@VRrVt#DUM&PJ+ufUuj>tI^)Rg-RAep{1Y^A~@9!;y||8T>Nv zb%TWiUk@hQ`;gKQ{8k`%b#z&eL#9Gyb&S1jAUMw<%8FD6?wvfybz&l03D>NAK-PHR z|Br1?biYsp!EbD|2rF(ugOX(!t-sR^T5p?{Vh6@x{q3N!>9ZkMTOhbW+!6o0C5))r z+5*tG4bSv8Ezc&>Uu!*w3%j5FZlgZ5`uborQ+bJRo|EidSbr7Ha;6Qe5^l-bhMS}` z)bv(D8ckqSyAv1{W?Z%e(kCX!-xCO?cTm`>i0bxT%f%wVRJh3uf{){Y&w|Nh3(f+w z`~FhP-#7Y*jG?l?9f?iR@M`^R;CE2fIQtu#-wb{gTtpPzqak(EH9u50y4-k5F4DTc zhyO$t7{!1NxQUQCEQ@h5R0N_=zaI&EMTmKyFp-C>n-J@~-snA{xtBH1D17B+EdlRn zN#^z_0bzJ=6P~DzhM#?b=#v0N50~|$FH54T)@v8mc%nJet;d25O(GrVcUD`5me$)m z3JgBr0w6I!R*GcnL%n`d9%O>tA1T%ouAQem2pmG#`exks@npKQ+Oad~RDiEXJy?CV zawvav$AgRI@z}sK2~IjKEknDkD~fLO;9*lp5t$n5Y2BI@V(TWoX&mjiLD}bm-hFF; zU%Mv=hUm}Xp4vfs@82sIwu%qHw1!Ze#-ZRtya8<;WvED0O8A1&`OhaSo>0L=FG@kQ z0M47&1FnIp+y^fL%wGwX6%SP@#uLl(_ZUFq?}xcVH{vu7{25{r%v%RAOLT0y_~^=J833tdMOJc!_TT&vqv{HR z^`3!dMWwxMMR)cyz-9r2z0Dg4Hk%gU<4hv>b%AlNp|%)4k`@)P%KD`Co%nndmiW`b zV>-j}z>dWzvSRLs?LyzZY8Hr+xtjJ8QGfMMPcDeWl-imj9u>IJ`>F`8|5R zU02a$g?^ubuT(9f4>Ir-S!aG5xkC$GLxiaUzxX!-;p8?=)?Wee>j4+~JZm7O1voQ_ zk5lmA?y55ico1K*JpEMyVJJNlB6jp!2wbSVL8fc@-ACcEAO|I}<$ReGKwX73dF%Gf zmG<%n+7Wdi@N=SS%OyPAY1e)VMZ=s!2f<_P_5olZi<;}lvrjaBB~5t}-k#{)Q>ETl zvOImADg|c(v2=}ACBT~I#ef0;t%Xq51Rp!)g}QMZbk0eeXUNfoGL$hO*-k*G{yDFI z80Y7+#|zGPH+=@b9$W_mm!z@B>pMjDPr`zUcc1densd(^av|;cZ5cR^Lh=QBJuBk& zHNh#lLKpPssKj^z!Mxuq%LykYM7G9#v;HD&-1cA2e^RE?aK?ABnv`uiL4kQ_6GUP_ zo}xW#JpggF{O4H-L3uHM0lrTz^xUif!6^=fZewsV2C(N_WoU`Hplhh&_j0@dR0x3I zPGbRU_t@s@7!k^PxCi}bEyfV?WN-xgXYw75@0L#&LSiucdwR8a7$OPF@NWHm+ZAVh z4ZyH_-XZu6pxxI$wD6V*0HDSDG)M7$K~JVR$zmJ=$0vsPMfz06w^dB#g_7)^;vaoe zF4HMenpX5g9rzNB5rZY7E^!IOvKv>@< z2F4li?t_W$$pM0O{e1!_Kb?XtdS8IwV!eYD64%xg(ug?_`-Ui1j2=6K?>P@;0k8nS z34}ZiN&$W|fLY~xzm8UHPP+^SR)L3EO65b~EQs?mkS*>L2zutMwJBp zn{(fW+x75@QgE5hy;n z8nN;GpAcEB%bN^XWR$ zmp;}OJ+FWq??Z+B4!#?)wufQfj*ihdY-4cX10U(W<+;>q9TZvPeH_ES#sbG>;VWKC zn#$VtExZ@5&^QJx?xuO2sslKNu=fU zzXKfKRDS}$rHr=vvfh^j?$$*4-yec7R!ao&&@mrth~TG^e_}r#cq!5C>L^s6)1oZi zm(=W9CaC+W0dgOjtfXFGH|Hv`?>&HK^&p}q3q(6v@36j`Ta*hXgx5T(fCU%p!^34m5<^LgO8X>~;h9<-Svs$)3* za#OFLEY9X?e(fYNorJP)678dHxv>y~TJNVeobqBtCm?=EA^2><_J_g^c> zmwkZtXL_%zBeM5zE9b;_`0wLf){8HV5&z@yfIOyknU6|f<^0q)94}B5_#cUfZf)G< zln$P4>A08!Tb;q8kyi==!&cw!#kF%?>%w((abA04R2~E+Vaq^l8ye=T$w@2(lKTb&sIXwF$w0_Y+_YUCq8*%JgiBJ(~U$LbSs{ohQ*I*>1yz^2fEcK_Qt!<35g z3GmC{9@=Ig!&j5RH)VLbP34-M%x8cL>o0+6lWchw$zFG;xc%cJgAa-2qNQXx3?hNw z&A>rgm_&;}My~GhRXKYUKD9mCi@f)!9I?3lV*{^h%W4aB-QwSnHT?)hywiq4SXHO> zRYfkHXU*HN7H<}f34Gi7``E#6#{-qR&MUp6aNK5q=<)3JC*b$GcFUb76OrDVs16O0 z^f3s_PaA-+NNN}g35l7p@KP~Z{LTxa94+X#KS&5x7NC$~Awk{%juB)h3pQJP3RqaV z4ZuWXdld6x4f1Q&f`(?XAkZfOof@YQpo9e#ot6v)zHpZ>uAgx6gY&-a>H&PO*5^OC za2-G|jzMGaE91M+S>SFS@ldB#*3y+oWt#8H)JUpYD6CqM7q4VW)*Gl7Ex}KzPl2xY zCarZbYfH$UR$q^J8T$nG(GZxYi>27IOgS;_I?f&<^W)7b`}$OVm-d$l{dN$Tj~JyY zmxQq8kpO#4b6(s_TOyia)Q6P55WGMBhOhlV&`xV8{$4>!o~llK4c+mT=IADAZp z9rUVDmR7AIV7=8DuNLk7hIT)6=dTlhfXxa#pWEb%+ZjY^M1bN5@T98_<^c$51MHr` zZ`s#sU2i;V|GP32lqt)jiRN#6yDE(MPukNI1S_(*FpPZbLp>Yc=1D{vg7d}e_9u9` zpmj+HRv4$46#Fe2crJG%+1USrwrSc&-?zQ$!=I#0&Sen%Z=?fXO5a&#d@AP$Dd`4$ ziT%NtJ;{C_={BZMW05MY8gZ_~!6^b&Om5M2!vJ0j8Se-HUdce1-^C`T;E=}_LVtfc zR(uaHlo3Z=*i->D4*@I!(d_-jM4--x{luGeiZEGwKEHa2`91W?`DP*wq7il+u;E47de?NoYq%hKTRx|XN$lekU*w@Et z*~H$NvO-qG*6S>qsyV+k=OI!bp{e*iHtIMC^J`oDcgr8+z{+FtfZ)DI8N8OsoWg6F zX&aXb0l{2(?t!AVrWCX{l}W>d5B7PtTt;-`OSD>`-E_Bn@p{nhW0^nP{{;_7N&qT* zGa77dp+0=k9<~*%wM#?9@}7@>{Jf7J?qhtG+J5ZIZ2k|FAyG&0a{{bj*+mIyfnT@} z7(pz%QA0qhf90r}vIL}czk?o=yecq1Q!$VNO~E@FK@APF!9JQt+Jx!E;?bOUG*xR5 zjKG}Q;p}=)%ACn~GgCLija~AF7GitBF|`d5v|^rHuep z-LHvWi#ja72xc+uLeOmI@@}+Cr$ySIYe^5c|owQyJzgaHNOl*kK<|qzz}-5WGimdDiiTAfQ=AW;R@){z*6IYoQk=yXS-@9| zDAekjwuq{zO9Ke`x6wA0yXyqJl0fjLG>UZr>vPkttlx)DAV=;j(_Eur(W-bp3Cu?O z&Xq0N59ewG_KoDV=S^0XB>2UTwBK0}$FeuYkF33(+71e&EeKB0?zLx_hSaW52a=H@ zoY$7(B`v!ZcrSY~PYpeYW`Wj#0al1+Lx&H{vZ9-fgvgjgMhgdFGZD7gUp27!&DTE( z!@B6)lkk1ai|yq3x7>UI`W)hkyMl3dyb2v}$C+V|AnivQm^pyu*RMB9;P|N@aP%RK z4ZM?;sIVXWH2@DM8j!$mgC%yNF%{s7J!zTHDuKzeyp!T8BdiEkwzVaLCmgHY>#UF= zIPZLJUDwvrOruLV&x`PPNO>O0l?sgOK6lYP+qe7%pwe9M7%U-t+k?DCilO-`!8mS$ zYq5{2`j5P5zVx(&;)U`hXGYXs216;w3UZhgx|uF*;FtKIGWEsf2k)o8%bCnqU6aY{ zD(j*U_{F|5s6D01{TZ#lpAz1V@pMh(qXpHZrX}Q^D1^b#id)KWXxP!bImcdE%4ngs z^Q}3NKW>Sb^VoF*!QOVSh!ySOwyv^b#`qfcH8X#Y0FJ1+y=MJ=81f9QzSjuiD1OzD z$ppW6EvQYt+IAcXxhWU4C7+{Bzy%YF&P{I{+{zbX|7P>mc|Q-xpGH zVmt1Yvjl{z^wOTJ#nMII(vrm?pW9!+w-w-}&2n^JZP-LMM#(PT&Jc9o8T!7jv|r4Y z(kNde;sFTenL?O9yCfnOqY(9!)Rwz6g5dlT5a_uMci{)TAPS)`7y~kK$nSM;AlTmr zrBSj%3G7u}ZAopXe6fd$eH=^zpA!=e#!2o_+puZ14)Q0C(4&Chz!!z6xzE!TJUAXO ze%N-blL)A{B$tIkSOcYLst*(}!L;}GCmL?c<#*Ah&%$pmhu$yr1wO2a<~|78si_2L z83g+aX_`@Z@>kjtZ{u1ml7N@xcYztdANHg#F`80?7$h4DP5U5&AqC}8xeDi=!M#+0 zu^O-;Xr97UBhND&e8bE%An{KPz9;9AE=Wlqt*D5qK*TN@ysO`6!33p2{zNq|S5 zj?XV-I%f|up0{cY7zp#&B^jh>+7r*%0CXC*gCgl^yem0GU|cmn-qu`2(=cK+wOdmD zrNG*(|4dhy-lc60?(4hX34%Dl>`v6Vq7()g1t%2W+%?M5kNy)w&KK+(qop7T|uVG_y&b zbd3dgF^%>#QSH*SiM(ynefB1GoHsMN4A;FMbL=tVtmnBV;B6}XHP0N!UPcd~^(^Xg zTT*x+c@*gA%pwpW)J+#w2Qv|rsP(Yog=?s}C9BT`-^J#6;I&Zv9xLyt#)N0V#5J~aqZFE7utFaLaLi}c;Yu3aHMm*(UbrxO{|wIIM)|JY0j+r@$N_*g z#xbLcE|U~ zXI*|puUaT`a0n3my3%LWuVFowDQICo@LaQhzuaqL zN+V`HEWv&sU-uo8y!EF7=_(lDBwTpez1s7>ad(|35L|-aP56&n&i4e z?!|48T%QXsiYRMtfnwkJI!^m<8xTp>$dr(2g%)_%B9}??KEY2s18ltSfo`D%IT|$I ziGHISwB-C-pVb*J2!BXVNN~e(OhQOl6fkoE1gr6maz4;}eC5OgysBsUZm(a0+~=|i zeo}*Asu$_-q=yzl+5~0J2e`fz{EEOY({uPcgzh<(8@o&U&xHXxisQ6E;OK;`-vc-` z_zrS|oYsHvJ>Y-jd6tqZQ~0@l$P@85G6?2&#AGGj>GMTN1rox=_RUk0d zvd-YCyFjO?$PXT52iYM3sA3L`=AY#UA96!i=qGv0KVgrSV%A}M<1oi7w64=<8izs# zqOnVgooy_#cS_oq)e!<3M{G;tea63K>v5$tk~|HW93J*lDD^n={wf3j0v)BeEHNqA zTpTO<@;urFo}I4%v^_xXD9b(Pvh_EvI_g6v^J-l-a$e`hCu7{4ePBn5HI$Y=9S0I{ zAx5x?;I}DRaXlcHk_+_A>+iR=U)yPE?AD#oDZU38i7@r9%L`HSUS9ck7itIDaojD* zSR@GHdN(|0tjR-ZSbe9(*Q2%1o327kp196CvSgDsk42hjw_(v;rs8FxRiZ5>5#?EL zPzreA?t2WVkT|)iUzV%(x{k(@Tk6ZTKXU~7eWmEV7L&dx zrAIEq1#%>`QV3>(lx=D8va0`~oF^b%yBp9d3f%U9AH(v+Lz!jGZ z%S?3c3x#C_w%V5l9iyjL;Mbak1}}kWvRrNZ4h23GpkOSLFj#2((|28T^4o;2{kt3V zL!?W53)fZL$Elmw-g>E$ib$eEpw09+^%s~hf;Y7lu{