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..0b9850ec --- /dev/null +++ b/swiftsimio/_array_functions.py @@ -0,0 +1,1838 @@ +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, + _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, +) + +_HANDLED_FUNCTIONS = dict() + + +def _return_helper(res, helper_result, ret_cf, out=None): + if out is None: + if isinstance(res, unyt_quantity): + return cosmo_quantity( + 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"): + out.cosmo_factor = ret_cf + if hasattr(out, "compression"): + 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=helper_result["comoving"], + cosmo_factor=ret_cf, + 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(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.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 + ) + 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: + ret_cf_w = _preserve_cosmo_factor(helper_result["kw_ca_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 + ) + 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_bins) + + +@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 + 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"]) + 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 + 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"]) + 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 = _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), + _return_helper(ybins, helper_result_y, ret_cf_y), + ) + + +@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: + 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: + ret_cf_w = _preserve_cosmo_factor(helper_result_w["kw_ca_cfs"]["weights"]) + 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, + bins=bins, + range=range, + weights=weights, + ) + ret_cfs = D * [_preserve_cosmo_factor(helper_result["ca_cfs"][0])] + counts, bins = unyt_histogramdd( + helper_result["args"], + 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"]) + 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_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 = _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, + [ + _return_helper(b, helper_result, ret_cf) + for b, helper_result, ret_cf in zip(bins, helper_results, ret_cfs) + ], + ) + + +@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) + + +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 + + 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 + + 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) +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( + 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) +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 + ) + 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(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.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]) + # 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) +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, + ) + 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) +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.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( + 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.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.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( + 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(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( + 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) +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] + ) + 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(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]) + ress = unyt_linalg_eig(*helper_result["args"], **helper_result["kwargs"]) + return ( + _return_helper(ress[0], helper_result, ret_cf), + ress[1], + ) + + +@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]) + ress = unyt_linalg_eigh(*helper_result["args"], **helper_result["kwargs"]) + return ( + _return_helper(ress[0], helper_result, ret_cf), + ress[1], + ) + + +@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.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(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) +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]) + # must pass a directly here because it's modified in-place + comoving = getattr(a, "comoving", None) + if comoving: + val.convert_to_comoving() + elif comoving is False: + val.convert_to_physical() + 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.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]) + # 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 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]) + # 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(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]) + # must pass arr directly here because it's modified in-place + 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(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]) + # must pass arr directly here because it's modified in-place + 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(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.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) + 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(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.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 + return unyt_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, +): + 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?) + 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) +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 + 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(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.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( + 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(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.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 + ) + 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(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 + ) + 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(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, + 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) +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.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: + 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(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) 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/swiftsimio/objects.py b/swiftsimio/objects.py index dad748ee..ce970514 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -6,8 +6,9 @@ import warnings import unyt -from unyt import unyt_array -from unyt.array import multiple_output_operators +from unyt import unyt_array, unyt_quantity +from unyt.array import multiple_output_operators, _iterable +from numbers import Number as numeric_type try: from unyt.array import POWER_MAPPING @@ -115,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 @@ -134,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): @@ -162,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: @@ -202,7 +223,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): @@ -306,7 +328,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 @@ -332,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): @@ -370,6 +393,98 @@ def _comparison_cosmo_factor(ca_cf1, ca_cf2=None, inputs=None): ) +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 @@ -453,7 +568,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 """ @@ -569,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 @@ -592,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): """ @@ -599,8 +725,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 ---------- @@ -774,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) @@ -880,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) @@ -889,13 +1022,35 @@ 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 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) + 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): @@ -1084,41 +1239,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 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) + 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 @@ -1126,16 +1249,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 @@ -1147,27 +1270,185 @@ 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 + + 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 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) + + +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, + ): + """ + 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 + 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 = 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 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) diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py index 4cb83535..302737e6 100644 --- a/tests/test_cosmo_array.py +++ b/tests/test_cosmo_array.py @@ -2,15 +2,56 @@ Tests the initialisation of a cosmo_array. """ +import pytest +import os +import warnings import numpy as np import unyt as u -from swiftsimio.objects import cosmo_array, cosmo_factor +from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a + +savetxt_file = "saved_array.txt" + + +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): + 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 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 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 +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") @@ -30,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") @@ -40,9 +81,461 @@ 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) + + 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: + """ + 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 = { + "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)).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))), + "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)),), + "column_stack": (ca(np.arange(3)),), + "stack": (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)),), + "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)),), + "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)), + "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)),), + "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)),), + "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)), + "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))), + "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)), + "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))), # np deprecated + "take": (ca(np.arange(3)), 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: + 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: + 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 + 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", + "place", + "put", + "put_along_axis", + "putmask", + ): + # 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: + 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 + 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", + [ + ".".join((f.__module__, f.__name__)).replace("numpy", "np") + for f in unchecked_functions + ], + ) + + # 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_type", ("int", "np", "ca")) + @pytest.mark.parametrize("density", (None, True)) + def test_histograms(self, func_args, weights, bins_type, 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_type] + bins = ( + bins[ + { + np.histogram: np.s_[0], + np.histogram2d: np.s_[:2], + np.histogramdd: np.s_[:], + }[func] + ] + if bins_type 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 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), + }[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 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), + }[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) + + 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) + + +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 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] 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") 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 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)