From a9f5c705080d78b0984edb70e5ec7e8ccb8b4941 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 21 Mar 2025 14:02:14 +0000
Subject: [PATCH 01/23] Messing around trying to find a workable approach.

---
 swiftsimio/objects.py     | 131 ++++++++++++++++++++------------------
 tests/test_cosmo_array.py |  69 ++++++++++++--------
 2 files changed, 110 insertions(+), 90 deletions(-)

diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py
index 3538012e..3acf87ed 100644
--- a/swiftsimio/objects.py
+++ b/swiftsimio/objects.py
@@ -103,6 +103,7 @@
 from ._array_functions import (
     _propagate_cosmo_array_attributes_to_result,
     _ensure_result_is_cosmo_array_or_quantity,
+    _copy_cosmo_array_attributes_if_present,
     _sqrt_cosmo_factor,
     _multiply_cosmo_factor,
     _preserve_cosmo_factor,
@@ -1600,6 +1601,62 @@ def from_pint(
 
         return obj
 
+    @classmethod
+    def __unyt_ufunc_prepare__(cls, ufunc: np.ufunc, method: str, *inputs, **kwargs):
+        helper_result = _prepare_array_func_args(*inputs, **kwargs)
+        return ufunc, method, helper_result["args"], helper_result["kwargs"]
+
+    def __unyt_ufunc_finalize__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
+        helper_result = _prepare_array_func_args(*inputs, **kwargs)
+        cfs = helper_result["cfs"]
+        ret = self
+        # make sure we evaluate the cosmo_factor_ufunc_registry function:
+        # might raise/warn even if we're not returning a cosmo_array
+        if ufunc in (multiply, divide) and method == "reduce":
+            power_map = POWER_MAPPING[ufunc]
+            if "axis" in kwargs and kwargs["axis"] is not None:
+                ret_cf = _power_cosmo_factor(
+                    cfs[0], None, power=power_map(inputs[0].shape[kwargs["axis"]])
+                )
+            else:
+                ret_cf = _power_cosmo_factor(
+                    cfs[0], None, power=power_map(inputs[0].size)
+                )
+        elif (
+            ufunc in (logical_and, logical_or, logical_xor, logical_not)
+            and method == "reduce"
+        ):
+            ret_cf = _return_without_cosmo_factor(cfs[0])
+        else:
+            ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs)
+        # if we get a tuple we have multiple return values to deal with
+        if isinstance(ret, tuple):
+            for r in ret:
+                if isinstance(r, cosmo_array):  # also recognizes cosmo_quantity
+                    r.comoving = helper_result["comoving"]
+                    r.cosmo_factor = ret_cf
+                    r.compression = helper_result["compression"]
+        elif isinstance(ret, cosmo_array):  # also recognizes cosmo_quantity
+            ret.comoving = helper_result["comoving"]
+            ret.cosmo_factor = ret_cf
+            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):  # also recognizes cosmo_quantity
+                    out.comoving = helper_result["comoving"]
+                    out.cosmo_factor = ret_cf
+                    out.compression = helper_result["compression"]
+            else:
+                for o in out:
+                    if isinstance(o, cosmo_array):  # also recognizes cosmo_quantity
+                        o.comoving = helper_result["comoving"]
+                        o.cosmo_factor = ret_cf
+                        o.compression = helper_result["compression"]
+
+        return ret
+
     def __array_ufunc__(
         self, ufunc: np.ufunc, method: str, *inputs, **kwargs
     ) -> object:
@@ -1753,6 +1810,8 @@ def __array_function__(
             function_to_invoke = func._implementation
         return function_to_invoke(*args, **kwargs)
 
+    @_propagate_cosmo_array_attributes_to_result
+    @_ensure_result_is_cosmo_array_or_quantity
     def __mul__(
         self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit]
     ) -> "cosmo_array":
@@ -1773,13 +1832,15 @@ def __mul__(
         out : swiftsimio.objects.cosmo_array
             The result of the multiplication.
         """
-        if isinstance(b, unyt.unit_object.Unit):
-            retval = self.__copy__()
-            retval.units = retval.units * b
-            return retval
+        if getattr(b, "is_Unit", False):
+            return b.__mul__(
+                self.view(unyt_quantity) if self.shape == () else self.view(unyt_array)
+            )
         else:
             return super().__mul__(b)
 
+    @_propagate_cosmo_array_attributes_to_result
+    @_ensure_result_is_cosmo_array_or_quantity
     def __rmul__(
         self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit]
     ) -> "cosmo_array":
@@ -1806,67 +1867,13 @@ def __rmul__(
         out : swiftsimio.objects.cosmo_array
             The result of the multiplication.
         """
-        if isinstance(b, unyt.unit_object.Unit):
-            return self.__mul__(b)
+        if getattr(b, "is_Unit", False):
+            return b.__rmul__(
+                self.view(unyt_quantity) if self.shape == () else self.view(unyt_array)
+            )
         else:
             return super().__rmul__(b)
 
-    def __truediv__(
-        self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit]
-    ) -> "cosmo_array":
-        """
-        Divide this :class:`~swiftsimio.objects.cosmo_array`.
-
-        We delegate most cases to :mod:`unyt`, but we need to handle the case where the
-        second argument is a :class:`~unyt.unit_object.Unit`.
-
-        Parameters
-        ----------
-        b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \
-        :class:`~unyt.unit_object.Unit`
-            The object to divide this one by.
-
-        Returns
-        -------
-        out : swiftsimio.objects.cosmo_array
-            The result of the division.
-        """
-        if isinstance(b, unyt.unit_object.Unit):
-            return self.__mul__(1 / b)
-        else:
-            return super().__truediv__(b)
-
-    def __rtruediv__(
-        self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit]
-    ) -> "cosmo_array":
-        """
-        Divide this :class:`~swiftsimio.objects.cosmo_array` (as the right argument).
-
-        We delegate most cases to :mod:`unyt`, but we need to handle the case where the
-        second argument is a :class:`~unyt.unit_object.Unit`.
-
-        .. note::
-
-            This function is never called when `b` is a :class:`unyt.unit_object.Unit`
-            because :mod:`unyt` handles the operation. This results in a silent demotion
-            to a :class:`unyt.array.unyt_array`.
-
-        Parameters
-        ----------
-        b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \
-        :class:`~unyt.unit_object.Unit`
-            The object to divide by this one.
-
-        Returns
-        -------
-        out : swiftsimio.objects.cosmo_array
-            The result of the division.
-        """
-        if isinstance(b, unyt.unit_object.Unit):
-            return (1 / self).__mul__(b)
-        else:
-            return super().__rtruediv__(b)
-
 
 class cosmo_quantity(cosmo_array, unyt_quantity):
     """
diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 3f87702d..219c9ea4 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -1078,23 +1078,28 @@ def test_multiplication_by_unyt(self):
         ca = cosmo_array(
             np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
         )
-        # required so that can test right-sided division with the same assertions:
-        assert np.allclose(ca.to_value(ca.units), 1)
-        # the reference result:
-        multiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
-        # get the same result twice through left-sided multiplication and division:
+
+        lmultiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
         lmultiplied_by_unyt = ca * u.Mpc
-        ldivided_by_unyt = ca / u.Mpc ** -1
+        assert isinstance(lmultiplied_by_quantity, cosmo_array)
+        assert isinstance(lmultiplied_by_unyt, cosmo_array)
+        assert lmultiplied_by_unyt.comoving == lmultiplied_by_quantity.comoving
+        assert np.allclose(
+            lmultiplied_by_unyt.to_value(lmultiplied_by_quantity.units),
+            lmultiplied_by_quantity.to_value(lmultiplied_by_quantity.units),
+        )
 
-        for multiplied_by_unyt in (lmultiplied_by_unyt, ldivided_by_unyt):
-            assert isinstance(multiplied_by_quantity, cosmo_array)
-            assert isinstance(multiplied_by_unyt, cosmo_array)
-            assert np.allclose(
-                multiplied_by_unyt.to_value(multiplied_by_quantity.units),
-                multiplied_by_quantity.to_value(multiplied_by_quantity.units),
-            )
+        ldivided_by_quantity = ca / (1 * u.Mpc)
+        ldivided_by_unyt = ca / u.Mpc
+        assert isinstance(ldivided_by_quantity, cosmo_array)
+        assert isinstance(ldivided_by_unyt, cosmo_array)
+        assert ldivided_by_unyt.comoving == ldivided_by_quantity.comoving
+        assert np.allclose(
+            ldivided_by_unyt.to_value(ldivided_by_quantity.units),
+            ldivided_by_quantity.to_value(ldivided_by_quantity.units),
+        )
 
-    @pytest.mark.xfail
+    # @pytest.mark.xfail
     def test_rmultiplication_by_unyt(self):
         """
         We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as
@@ -1107,25 +1112,33 @@ def test_rmultiplication_by_unyt(self):
         to xfail.
 
         If this is fixed in the future this test will pass and can be merged with
-        `test_multiplication_by_unyt` to tidy up.
+        `test_multiplication_by_unyt` to tidy up. Also clean up docstrings of
+        `cosmo_array.__mul__`, `cosmo_array.__rmul__`, `cosmo_array.__truediv__`,
+        `cosmo_array.__rtruediv__`.
 
         See https://github.com/yt-project/unyt/pull/572
         """
         ca = cosmo_array(
             np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
         )
-        # required so that can test right-sided division with the same assertions:
-        assert np.allclose(ca.to_value(ca.units), 1)
-        # the reference result:
-        multiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
-        # get 2x the same result through right-sided multiplication and division:
+
+        rmultiplied_by_quantity = (1 * u.Mpc) * ca  # parentheses very important here
+        assert rmultiplied_by_quantity.comoving
         rmultiplied_by_unyt = u.Mpc * ca
-        rdivided_by_unyt = u.Mpc ** 2 / ca
+        assert isinstance(rmultiplied_by_quantity, cosmo_array)
+        assert isinstance(rmultiplied_by_unyt, cosmo_array)
+        assert rmultiplied_by_unyt.comoving == rmultiplied_by_quantity.comoving
+        assert np.allclose(
+            rmultiplied_by_unyt.to_value(rmultiplied_by_quantity.units),
+            rmultiplied_by_quantity.to_value(rmultiplied_by_quantity.units),
+        )
 
-        for multiplied_by_unyt in (rmultiplied_by_unyt, rdivided_by_unyt):
-            assert isinstance(multiplied_by_quantity, cosmo_array)
-            assert isinstance(multiplied_by_unyt, cosmo_array)
-            assert np.allclose(
-                multiplied_by_unyt.to_value(multiplied_by_quantity.units),
-                multiplied_by_quantity.to_value(multiplied_by_quantity.units),
-            )
+        rdivided_by_quantity = (1 * u.Mpc) / ca  # parentheses very important here
+        rdivided_by_unyt = u.Mpc / ca
+        assert isinstance(rdivided_by_quantity, cosmo_array)
+        assert isinstance(rdivided_by_unyt, cosmo_array)
+        assert rdivided_by_unyt.comoving == rdivided_by_quantity.comoving
+        assert np.allclose(
+            rdivided_by_unyt.to_value(rdivided_by_quantity.units),
+            rdivided_by_quantity.to_value(rdivided_by_quantity.units),
+        )

From 11517045627d5e218ac8c3baf7c1c20911cc7089 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 21 Mar 2025 16:29:13 +0000
Subject: [PATCH 02/23] Improvements to a test.

---
 tests/test_visualisation.py | 31 +++++++++++++++----------------
 1 file changed, 15 insertions(+), 16 deletions(-)

diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py
index 3059fcc4..ca1cabec 100644
--- a/tests/test_visualisation.py
+++ b/tests/test_visualisation.py
@@ -377,24 +377,17 @@ def test_comoving_versus_physical(filename):
     # this test is pretty slow if we don't mask out some particles
     m = mask(filename)
     boxsize = m.metadata.boxsize
-    m.constrain_spatial([[0.0 * b, 0.2 * b] for b in boxsize])
-    region = [
-        0.0 * boxsize[0],
-        0.2 * boxsize[0],
-        0.0 * boxsize[1],
-        0.2 * boxsize[1],
-        0.0 * boxsize[2],
-        0.2 * boxsize[2],
-    ]
+    region = cosmo_array([np.zeros_like(boxsize), 0.2 * boxsize]).T
+    m.constrain_spatial(region)
     for func, aexp in [(project_gas, -2.0), (slice_gas, -3.0), (render_gas, -3.0)]:
         # normal case: everything comoving
         data = load(filename, mask=m)
         # we force the default (project="masses") to check the cosmo_factor
         # conversion in this case
-        img = func(data, resolution=64, project=None, region=region)
+        img = func(data, resolution=64, project=None, region=region.flatten())
         assert data.gas.masses.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0
-        img = func(data, resolution=64, project="densities", region=region)
+        img = func(data, resolution=64, project="densities", region=region.flatten())
         assert data.gas.densities.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0
         # try to mix comoving coordinates with a physical variable:
@@ -406,14 +399,16 @@ def test_comoving_versus_physical(filename):
             with pytest.warns(
                 UserWarning, match="Converting coordinate grid to physical."
             ):
-                img = func(data, resolution=64, project="densities", region=region)
+                img = func(
+                    data, resolution=64, project="densities", region=region.flatten()
+                )
         assert data.gas.densities.comoving is False and img.comoving is False
         assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0
         # convert coordinates to physical (but not smoothing lengths):
         # the coordinates (copy) should convert back to comoving to match the masses
         data.gas.coordinates.convert_to_physical()
         with pytest.warns(UserWarning, match="Converting coordinate grid to comoving."):
-            img = func(data, resolution=64, project="masses", region=region)
+            img = func(data, resolution=64, project="masses", region=region.flatten())
         assert data.gas.masses.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0
         # also convert smoothing lengths to physical
@@ -425,12 +420,14 @@ def test_comoving_versus_physical(filename):
             with pytest.warns(
                 UserWarning, match="Converting coordinate grid to comoving."
             ):
-                img = func(data, resolution=64, project="masses", region=region)
+                img = func(
+                    data, resolution=64, project="masses", region=region.flatten()
+                )
         assert data.gas.masses.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0
         # densities are physical, make sure this works with physical coordinates and
         # smoothing lengths
-        img = func(data, resolution=64, project="densities", region=region)
+        img = func(data, resolution=64, project="densities", region=region.flatten())
         assert data.gas.densities.comoving is False and img.comoving is False
         assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0
         # now try again with comoving densities, should work and give a comoving img
@@ -442,7 +439,9 @@ def test_comoving_versus_physical(filename):
             with pytest.warns(
                 UserWarning, match="Converting coordinate grid to comoving."
             ):
-                img = func(data, resolution=64, project="densities", region=region)
+                img = func(
+                    data, resolution=64, project="densities", region=region.flatten()
+                )
         assert data.gas.densities.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0
 

From d268816836273a3f5f7da7f2cc1e58d4827b3312 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Sat, 22 Mar 2025 13:16:10 +0000
Subject: [PATCH 03/23] Tidy up and ensure tests pass.

---
 swiftsimio/objects.py     | 106 ++++---
 tests/test_cosmo_array.py | 603 ++++++++++++++++++++------------------
 2 files changed, 388 insertions(+), 321 deletions(-)

diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py
index 3acf87ed..f0f47ac9 100644
--- a/swiftsimio/objects.py
+++ b/swiftsimio/objects.py
@@ -1606,10 +1606,12 @@ def __unyt_ufunc_prepare__(cls, ufunc: np.ufunc, method: str, *inputs, **kwargs)
         helper_result = _prepare_array_func_args(*inputs, **kwargs)
         return ufunc, method, helper_result["args"], helper_result["kwargs"]
 
-    def __unyt_ufunc_finalize__(self, ufunc: np.ufunc, method: str, *inputs, **kwargs):
+    @classmethod
+    def __unyt_ufunc_finalize__(
+        cls, result, ufunc: np.ufunc, method: str, *inputs, **kwargs
+    ):
         helper_result = _prepare_array_func_args(*inputs, **kwargs)
         cfs = helper_result["cfs"]
-        ret = self
         # make sure we evaluate the cosmo_factor_ufunc_registry function:
         # might raise/warn even if we're not returning a cosmo_array
         if ufunc in (multiply, divide) and method == "reduce":
@@ -1628,34 +1630,69 @@ def __unyt_ufunc_finalize__(self, ufunc: np.ufunc, method: str, *inputs, **kwarg
         ):
             ret_cf = _return_without_cosmo_factor(cfs[0])
         else:
-            ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs)
+            ret_cf = cls._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs)
         # if we get a tuple we have multiple return values to deal with
-        if isinstance(ret, tuple):
-            for r in ret:
+        if isinstance(result, tuple):
+            r = tuple(
+                (
+                    r.view(cosmo_quantity)
+                    if r.shape == ()
+                    else (
+                        r.view(cosmo_array)
+                        if isinstance(r, unyt_array) and not isinstance(r, cosmo_array)
+                        else r
+                    )
+                )
+                for r in result
+            )
+            for r in result:
                 if isinstance(r, cosmo_array):  # also recognizes cosmo_quantity
                     r.comoving = helper_result["comoving"]
                     r.cosmo_factor = ret_cf
                     r.compression = helper_result["compression"]
-        elif isinstance(ret, cosmo_array):  # also recognizes cosmo_quantity
-            ret.comoving = helper_result["comoving"]
-            ret.cosmo_factor = ret_cf
-            ret.compression = helper_result["compression"]
+        elif isinstance(result, unyt_array):  # also recognizes cosmo_quantity
+            if not isinstance(result, cosmo_array):
+                result = (
+                    result.view(cosmo_quantity)
+                    if result.shape == ()
+                    else result.view(cosmo_array)
+                )
+            result.comoving = helper_result["comoving"]
+            result.cosmo_factor = ret_cf
+            result.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, unyt_array) and not isinstance(out, cosmo_array):
+                    out = (
+                        out.view(cosmo_quantity)
+                        if out.shape == ()
+                        else out.view(cosmo_array)
+                    )
                 if isinstance(out, cosmo_array):  # also recognizes cosmo_quantity
                     out.comoving = helper_result["comoving"]
                     out.cosmo_factor = ret_cf
                     out.compression = helper_result["compression"]
             else:
+                out = tuple(
+                    (
+                        (
+                            o.view(cosmo_quantity)
+                            if o.shape == ()
+                            else o.view(cosmo_array)
+                        )
+                        if isinstance(o, unyt_array) and not isinstance(o, cosmo_array)
+                        else o
+                    )
+                    for o in out
+                )
                 for o in out:
                     if isinstance(o, cosmo_array):  # also recognizes cosmo_quantity
                         o.comoving = helper_result["comoving"]
                         o.cosmo_factor = ret_cf
                         o.compression = helper_result["compression"]
-
-        return ret
+        return result
 
     def __array_ufunc__(
         self, ufunc: np.ufunc, method: str, *inputs, **kwargs
@@ -1713,20 +1750,20 @@ def __array_ufunc__(
         else:
             ret_cf = self._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs)
 
-        ret = _ensure_result_is_cosmo_array_or_quantity(super().__array_ufunc__)(
+        result = _ensure_result_is_cosmo_array_or_quantity(super().__array_ufunc__)(
             ufunc, method, *inputs, **kwargs
         )
         # if we get a tuple we have multiple return values to deal with
-        if isinstance(ret, tuple):
-            for r in ret:
+        if isinstance(result, tuple):
+            for r in result:
                 if isinstance(r, cosmo_array):  # also recognizes cosmo_quantity
                     r.comoving = helper_result["comoving"]
                     r.cosmo_factor = ret_cf
                     r.compression = helper_result["compression"]
-        elif isinstance(ret, cosmo_array):  # also recognizes cosmo_quantity
-            ret.comoving = helper_result["comoving"]
-            ret.cosmo_factor = ret_cf
-            ret.compression = helper_result["compression"]
+        elif isinstance(result, cosmo_array):  # also recognizes cosmo_quantity
+            result.comoving = helper_result["comoving"]
+            result.cosmo_factor = ret_cf
+            result.compression = helper_result["compression"]
         if "out" in kwargs:
             out = kwargs.pop("out")
             if ufunc not in multiple_output_operators:
@@ -1742,7 +1779,7 @@ def __array_ufunc__(
                         o.cosmo_factor = ret_cf
                         o.compression = helper_result["compression"]
 
-        return ret
+        return result
 
     def __array_function__(
         self, func: Callable, types: Collection, args: tuple, kwargs: dict
@@ -1810,8 +1847,6 @@ def __array_function__(
             function_to_invoke = func._implementation
         return function_to_invoke(*args, **kwargs)
 
-    @_propagate_cosmo_array_attributes_to_result
-    @_ensure_result_is_cosmo_array_or_quantity
     def __mul__(
         self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit]
     ) -> "cosmo_array":
@@ -1833,14 +1868,17 @@ def __mul__(
             The result of the multiplication.
         """
         if getattr(b, "is_Unit", False):
-            return b.__mul__(
-                self.view(unyt_quantity) if self.shape == () else self.view(unyt_array)
+            return _copy_cosmo_array_attributes_if_present(
+                self,
+                _ensure_result_is_cosmo_array_or_quantity(b.__mul__)(
+                    self.view(unyt_quantity)
+                    if self.shape == ()
+                    else self.view(unyt_array)
+                ),
             )
         else:
             return super().__mul__(b)
 
-    @_propagate_cosmo_array_attributes_to_result
-    @_ensure_result_is_cosmo_array_or_quantity
     def __rmul__(
         self, b: Union[int, float, np.ndarray, unyt.unit_object.Unit]
     ) -> "cosmo_array":
@@ -1850,12 +1888,6 @@ def __rmul__(
         We delegate most cases to :mod:`unyt`, but we need to handle the case where the
         second argument is a :class:`~unyt.unit_object.Unit`.
 
-        .. note::
-
-            This function is never called when `b` is a :class:`unyt.unit_object.Unit`
-            because :mod:`unyt` handles the operation. This results in a silent demotion
-            to a :class:`unyt.array.unyt_array`.
-
         Parameters
         ----------
         b : :class:`~numpy.ndarray`, :obj:`int`, :obj:`float` or \
@@ -1868,9 +1900,7 @@ def __rmul__(
             The result of the multiplication.
         """
         if getattr(b, "is_Unit", False):
-            return b.__rmul__(
-                self.view(unyt_quantity) if self.shape == () else self.view(unyt_array)
-            )
+            return self.__mul__(b)
         else:
             return super().__rmul__(b)
 
@@ -2006,7 +2036,7 @@ def __new__(
             hdf5 file.
         """
         if bypass_validation is True:
-            ret = super().__new__(
+            result = super().__new__(
                 cls,
                 np.asarray(input_scalar),
                 units=units,
@@ -2046,7 +2076,7 @@ def __new__(
             if compression is None
             else compression
         )
-        ret = super().__new__(
+        result = super().__new__(
             cls,
             np.asarray(input_scalar),
             units=units,
@@ -2061,9 +2091,9 @@ def __new__(
             valid_transform=valid_transform,
             compression=compression,
         )
-        if ret.size > 1:
+        if result.size > 1:
             raise RuntimeError("cosmo_quantity instances must be scalars")
-        return ret
+        return result
 
     __round__ = _propagate_cosmo_array_attributes_to_result(
         _ensure_result_is_cosmo_array_or_quantity(unyt_quantity.__round__)
diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 219c9ea4..c06a56a3 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -24,7 +24,7 @@ def getfunc(fname):
     return func
 
 
-def ca(x, unit=u.Mpc):
+def cosmo_obect(x, unit=u.Mpc):
     """
     Helper for our tests: turn an array into a cosmo_array.
     """
@@ -329,253 +329,281 @@ def test_explicitly_handled_funcs(self):
 
         functions_to_check = {
             # FUNCTIONS UNYT HANDLES EXPLICITLY:
-            "array2string": (ca(np.arange(3)),),
-            "dot": (ca(np.arange(3)), ca(np.arange(3))),
-            "vdot": (ca(np.arange(3)), ca(np.arange(3))),
-            "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))),
+            "array2string": (cosmo_obect(np.arange(3)),),
+            "dot": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "vdot": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "inner": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "outer": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "kron": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "histogram_bin_edges": (cosmo_obect(np.arange(3)),),
+            "linalg.inv": (cosmo_obect(np.eye(3)),),
+            "linalg.tensorinv": (cosmo_obect(np.eye(9).reshape((3, 3, 3, 3))),),
+            "linalg.pinv": (cosmo_obect(np.eye(3)),),
+            "linalg.svd": (cosmo_obect(np.eye(3)),),
+            "histogram": (cosmo_obect(np.arange(3)),),
+            "histogram2d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "histogramdd": (cosmo_obect(np.arange(3)).reshape((1, 3)),),
+            "concatenate": (cosmo_obect(np.eye(3)),),
+            "cross": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "intersect1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "union1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "linalg.norm": (cosmo_obect(np.arange(3)),),
+            "vstack": (cosmo_obect(np.arange(3)),),
+            "hstack": (cosmo_obect(np.arange(3)),),
+            "dstack": (cosmo_obect(np.arange(3)),),
+            "column_stack": (cosmo_obect(np.arange(3)),),
+            "stack": (cosmo_obect(np.arange(3)),),
+            "around": (cosmo_obect(np.arange(3)),),
+            "block": ([[cosmo_obect(np.arange(3))], [cosmo_obect(np.arange(3))]],),
+            "fft.fft": (cosmo_obect(np.arange(3)),),
+            "fft.fft2": (cosmo_obect(np.eye(3)),),
+            "fft.fftn": (cosmo_obect(np.arange(3)),),
+            "fft.hfft": (cosmo_obect(np.arange(3)),),
+            "fft.rfft": (cosmo_obect(np.arange(3)),),
+            "fft.rfft2": (cosmo_obect(np.eye(3)),),
+            "fft.rfftn": (cosmo_obect(np.arange(3)),),
+            "fft.ifft": (cosmo_obect(np.arange(3)),),
+            "fft.ifft2": (cosmo_obect(np.eye(3)),),
+            "fft.ifftn": (cosmo_obect(np.arange(3)),),
+            "fft.ihfft": (cosmo_obect(np.arange(3)),),
+            "fft.irfft": (cosmo_obect(np.arange(3)),),
+            "fft.irfft2": (cosmo_obect(np.eye(3)),),
+            "fft.irfftn": (cosmo_obect(np.arange(3)),),
+            "fft.fftshift": (cosmo_obect(np.arange(3)),),
+            "fft.ifftshift": (cosmo_obect(np.arange(3)),),
+            "sort_complex": (cosmo_obect(np.arange(3)),),
+            "isclose": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "allclose": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "array_equal": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "array_equiv": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
             "linspace": (cq(1), cq(2)),
             "logspace": (cq(1, unit=u.dimensionless), cq(2, unit=u.dimensionless)),
             "geomspace": (cq(1), cq(1)),
-            "copyto": (ca(np.arange(3)), ca(np.arange(3))),
-            "prod": (ca(np.arange(3)),),
-            "var": (ca(np.arange(3)),),
-            "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, cq(1)),
-            "linalg.lstsq": (ca(np.eye(3)), ca(np.eye(3))),
-            "linalg.solve": (ca(np.eye(3)), ca(np.eye(3))),
+            "copyto": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "prod": (cosmo_obect(np.arange(3)),),
+            "var": (cosmo_obect(np.arange(3)),),
+            "trace": (cosmo_obect(np.eye(3)),),
+            "percentile": (cosmo_obect(np.arange(3)), 30),
+            "quantile": (cosmo_obect(np.arange(3)), 0.3),
+            "nanpercentile": (cosmo_obect(np.arange(3)), 30),
+            "nanquantile": (cosmo_obect(np.arange(3)), 0.3),
+            "linalg.det": (cosmo_obect(np.eye(3)),),
+            "diff": (cosmo_obect(np.arange(3)),),
+            "ediff1d": (cosmo_obect(np.arange(3)),),
+            "ptp": (cosmo_obect(np.arange(3)),),
+            "cumprod": (cosmo_obect(np.arange(3)),),
+            "pad": (cosmo_obect(np.arange(3)), 3),
+            "choose": (np.arange(3), cosmo_obect(np.eye(3))),
+            "insert": (cosmo_obect(np.arange(3)), 1, cq(1)),
+            "linalg.lstsq": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
+            "linalg.solve": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
             "linalg.tensorsolve": (
-                ca(np.eye(24).reshape((6, 4, 2, 3, 4))),
-                ca(np.ones((6, 4))),
+                cosmo_obect(np.eye(24).reshape((6, 4, 2, 3, 4))),
+                cosmo_obect(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))),
+            "linalg.eig": (cosmo_obect(np.eye(3)),),
+            "linalg.eigh": (cosmo_obect(np.eye(3)),),
+            "linalg.eigvals": (cosmo_obect(np.eye(3)),),
+            "linalg.eigvalsh": (cosmo_obect(np.eye(3)),),
+            "savetxt": (savetxt_file, cosmo_obect(np.arange(3))),
+            "fill_diagonal": (cosmo_obect(np.eye(3)), cosmo_obect(np.arange(3))),
+            "apply_over_axes": (lambda x, axis: x, cosmo_obect(np.eye(3)), (0, 1)),
+            "isin": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "place": (
+                cosmo_obect(np.arange(3)),
+                np.arange(3) > 0,
+                cosmo_obect(np.arange(3)),
+            ),
+            "put": (cosmo_obect(np.arange(3)), np.arange(3), cosmo_obect(np.arange(3))),
+            "put_along_axis": (
+                cosmo_obect(np.arange(3)),
+                np.arange(3),
+                cosmo_obect(np.arange(3)),
+                0,
+            ),
+            "putmask": (
+                cosmo_obect(np.arange(3)),
+                np.arange(3),
+                cosmo_obect(np.arange(3)),
+            ),
+            "searchsorted": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
             "select": (
                 [np.arange(3) < 1, np.arange(3) > 1],
-                [ca(np.arange(3)), ca(np.arange(3))],
+                [cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))],
                 cq(1),
             ),
-            "setdiff1d": (ca(np.arange(3)), ca(np.arange(3, 6))),
-            "sinc": (ca(np.arange(3)),),
-            "clip": (ca(np.arange(3)), cq(1), cq(2)),
-            "where": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))),
-            "triu": (ca(np.ones((3, 3))),),
-            "tril": (ca(np.ones((3, 3))),),
-            "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)),
+            "setdiff1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3, 6))),
+            "sinc": (cosmo_obect(np.arange(3)),),
+            "clip": (cosmo_obect(np.arange(3)), cq(1), cq(2)),
+            "where": (
+                cosmo_obect(np.arange(3)),
+                cosmo_obect(np.arange(3)),
+                cosmo_obect(np.arange(3)),
+            ),
+            "triu": (cosmo_obect(np.ones((3, 3))),),
+            "tril": (cosmo_obect(np.ones((3, 3))),),
+            "einsum": ("ii->i", cosmo_obect(np.eye(3))),
+            "convolve": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "correlate": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "tensordot": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
+            "unwrap": (cosmo_obect(np.arange(3)),),
+            "interp": (
+                cosmo_obect(np.arange(3)),
+                cosmo_obect(np.arange(3)),
+                cosmo_obect(np.arange(3)),
+            ),
+            "array_repr": (cosmo_obect(np.arange(3)),),
+            "linalg.outer": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "trapezoid": (cosmo_obect(np.arange(3)),),
+            "in1d": (
+                cosmo_obect(np.arange(3)),
+                cosmo_obect(np.arange(3)),
+            ),  # np deprecated
+            "take": (cosmo_obect(np.arange(3)), np.arange(3)),
             # FUNCTIONS THAT UNYT DOESN'T HANDLE EXPLICITLY (THEY "JUST WORK"):
-            "all": (ca(np.arange(3)),),
-            "amax": (ca(np.arange(3)),),  # implemented via max
-            "amin": (ca(np.arange(3)),),  # implemented via min
+            "all": (cosmo_obect(np.arange(3)),),
+            "amax": (cosmo_obect(np.arange(3)),),  # implemented via max
+            "amin": (cosmo_obect(np.arange(3)),),  # implemented via min
             "angle": (cq(complex(1, 1)),),
-            "any": (ca(np.arange(3)),),
-            "append": (ca(np.arange(3)), cq(1)),
-            "apply_along_axis": (lambda x: x, 0, ca(np.eye(3))),
-            "argmax": (ca(np.arange(3)),),  # implemented via max
-            "argmin": (ca(np.arange(3)),),  # implemented via min
-            "argpartition": (ca(np.arange(3)), 1),  # implemented via partition
-            "argsort": (ca(np.arange(3)),),  # implemented via sort
-            "argwhere": (ca(np.arange(3)),),
-            "array_str": (ca(np.arange(3)),),
-            "atleast_1d": (ca(np.arange(3)),),
-            "atleast_2d": (ca(np.arange(3)),),
-            "atleast_3d": (ca(np.arange(3)),),
-            "average": (ca(np.arange(3)),),
-            "can_cast": (ca(np.arange(3)), np.float64),
-            "common_type": (ca(np.arange(3)), ca(np.arange(3))),
-            "result_type": (ca(np.ones(3)), ca(np.ones(3))),
-            "iscomplex": (ca(np.arange(3)),),
-            "iscomplexobj": (ca(np.arange(3)),),
-            "isreal": (ca(np.arange(3)),),
-            "isrealobj": (ca(np.arange(3)),),
-            "nan_to_num": (ca(np.arange(3)),),
-            "nanargmax": (ca(np.arange(3)),),  # implemented via max
-            "nanargmin": (ca(np.arange(3)),),  # implemented via min
-            "nanmax": (ca(np.arange(3)),),  # implemented via max
-            "nanmean": (ca(np.arange(3)),),  # implemented via mean
-            "nanmedian": (ca(np.arange(3)),),  # implemented via median
-            "nanmin": (ca(np.arange(3)),),  # implemented via min
-            "trim_zeros": (ca(np.arange(3)),),
-            "max": (ca(np.arange(3)),),
-            "mean": (ca(np.arange(3)),),
-            "median": (ca(np.arange(3)),),
-            "min": (ca(np.arange(3)),),
-            "ndim": (ca(np.arange(3)),),
-            "shape": (ca(np.arange(3)),),
-            "size": (ca(np.arange(3)),),
-            "sort": (ca(np.arange(3)),),
-            "sum": (ca(np.arange(3)),),
-            "repeat": (ca(np.arange(3)), 2),
-            "tile": (ca(np.arange(3)), 2),
-            "shares_memory": (ca(np.arange(3)), ca(np.arange(3))),
-            "nonzero": (ca(np.arange(3)),),
-            "count_nonzero": (ca(np.arange(3)),),
-            "flatnonzero": (ca(np.arange(3)),),
-            "isneginf": (ca(np.arange(3)),),
-            "isposinf": (ca(np.arange(3)),),
-            "empty_like": (ca(np.arange(3)),),
-            "full_like": (ca(np.arange(3)), cq(1)),
-            "ones_like": (ca(np.arange(3)),),
-            "zeros_like": (ca(np.arange(3)),),
-            "copy": (ca(np.arange(3)),),
-            "meshgrid": (ca(np.arange(3)), ca(np.arange(3))),
-            "transpose": (ca(np.eye(3)),),
-            "reshape": (ca(np.arange(3)), (3,)),
-            "resize": (ca(np.arange(3)), 6),
-            "roll": (ca(np.arange(3)), 1),
-            "rollaxis": (ca(np.arange(3)), 0),
-            "rot90": (ca(np.eye(3)),),
-            "expand_dims": (ca(np.arange(3)), 0),
-            "squeeze": (ca(np.arange(3)),),
-            "flip": (ca(np.eye(3)),),
-            "fliplr": (ca(np.eye(3)),),
-            "flipud": (ca(np.eye(3)),),
-            "delete": (ca(np.arange(3)), 0),
-            "partition": (ca(np.arange(3)), 1),
-            "broadcast_to": (ca(np.arange(3)), 3),
-            "broadcast_arrays": (ca(np.arange(3)),),
-            "split": (ca(np.arange(3)), 1),
-            "array_split": (ca(np.arange(3)), 1),
-            "dsplit": (ca(np.arange(27)).reshape(3, 3, 3), 1),
-            "hsplit": (ca(np.arange(3)), 1),
-            "vsplit": (ca(np.eye(3)), 1),
-            "swapaxes": (ca(np.eye(3)), 0, 1),
-            "moveaxis": (ca(np.eye(3)), 0, 1),
-            "nansum": (ca(np.arange(3)),),  # implemented via sum
-            "std": (ca(np.arange(3)),),
-            "nanstd": (ca(np.arange(3)),),
-            "nanvar": (ca(np.arange(3)),),
-            "nanprod": (ca(np.arange(3)),),
-            "diag": (ca(np.eye(3)),),
-            "diag_indices_from": (ca(np.eye(3)),),
-            "diagflat": (ca(np.eye(3)),),
-            "diagonal": (ca(np.eye(3)),),
-            "ravel": (ca(np.arange(3)),),
+            "any": (cosmo_obect(np.arange(3)),),
+            "append": (cosmo_obect(np.arange(3)), cq(1)),
+            "apply_along_axis": (lambda x: x, 0, cosmo_obect(np.eye(3))),
+            "argmax": (cosmo_obect(np.arange(3)),),  # implemented via max
+            "argmin": (cosmo_obect(np.arange(3)),),  # implemented via min
+            "argpartition": (cosmo_obect(np.arange(3)), 1),  # implemented via partition
+            "argsort": (cosmo_obect(np.arange(3)),),  # implemented via sort
+            "argwhere": (cosmo_obect(np.arange(3)),),
+            "array_str": (cosmo_obect(np.arange(3)),),
+            "atleast_1d": (cosmo_obect(np.arange(3)),),
+            "atleast_2d": (cosmo_obect(np.arange(3)),),
+            "atleast_3d": (cosmo_obect(np.arange(3)),),
+            "average": (cosmo_obect(np.arange(3)),),
+            "can_cast": (cosmo_obect(np.arange(3)), np.float64),
+            "common_type": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "result_type": (cosmo_obect(np.ones(3)), cosmo_obect(np.ones(3))),
+            "iscomplex": (cosmo_obect(np.arange(3)),),
+            "iscomplexobj": (cosmo_obect(np.arange(3)),),
+            "isreal": (cosmo_obect(np.arange(3)),),
+            "isrealobj": (cosmo_obect(np.arange(3)),),
+            "nan_to_num": (cosmo_obect(np.arange(3)),),
+            "nanargmax": (cosmo_obect(np.arange(3)),),  # implemented via max
+            "nanargmin": (cosmo_obect(np.arange(3)),),  # implemented via min
+            "nanmax": (cosmo_obect(np.arange(3)),),  # implemented via max
+            "nanmean": (cosmo_obect(np.arange(3)),),  # implemented via mean
+            "nanmedian": (cosmo_obect(np.arange(3)),),  # implemented via median
+            "nanmin": (cosmo_obect(np.arange(3)),),  # implemented via min
+            "trim_zeros": (cosmo_obect(np.arange(3)),),
+            "max": (cosmo_obect(np.arange(3)),),
+            "mean": (cosmo_obect(np.arange(3)),),
+            "median": (cosmo_obect(np.arange(3)),),
+            "min": (cosmo_obect(np.arange(3)),),
+            "ndim": (cosmo_obect(np.arange(3)),),
+            "shape": (cosmo_obect(np.arange(3)),),
+            "size": (cosmo_obect(np.arange(3)),),
+            "sort": (cosmo_obect(np.arange(3)),),
+            "sum": (cosmo_obect(np.arange(3)),),
+            "repeat": (cosmo_obect(np.arange(3)), 2),
+            "tile": (cosmo_obect(np.arange(3)), 2),
+            "shares_memory": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "nonzero": (cosmo_obect(np.arange(3)),),
+            "count_nonzero": (cosmo_obect(np.arange(3)),),
+            "flatnonzero": (cosmo_obect(np.arange(3)),),
+            "isneginf": (cosmo_obect(np.arange(3)),),
+            "isposinf": (cosmo_obect(np.arange(3)),),
+            "empty_like": (cosmo_obect(np.arange(3)),),
+            "full_like": (cosmo_obect(np.arange(3)), cq(1)),
+            "ones_like": (cosmo_obect(np.arange(3)),),
+            "zeros_like": (cosmo_obect(np.arange(3)),),
+            "copy": (cosmo_obect(np.arange(3)),),
+            "meshgrid": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "transpose": (cosmo_obect(np.eye(3)),),
+            "reshape": (cosmo_obect(np.arange(3)), (3,)),
+            "resize": (cosmo_obect(np.arange(3)), 6),
+            "roll": (cosmo_obect(np.arange(3)), 1),
+            "rollaxis": (cosmo_obect(np.arange(3)), 0),
+            "rot90": (cosmo_obect(np.eye(3)),),
+            "expand_dims": (cosmo_obect(np.arange(3)), 0),
+            "squeeze": (cosmo_obect(np.arange(3)),),
+            "flip": (cosmo_obect(np.eye(3)),),
+            "fliplr": (cosmo_obect(np.eye(3)),),
+            "flipud": (cosmo_obect(np.eye(3)),),
+            "delete": (cosmo_obect(np.arange(3)), 0),
+            "partition": (cosmo_obect(np.arange(3)), 1),
+            "broadcast_to": (cosmo_obect(np.arange(3)), 3),
+            "broadcast_arrays": (cosmo_obect(np.arange(3)),),
+            "split": (cosmo_obect(np.arange(3)), 1),
+            "array_split": (cosmo_obect(np.arange(3)), 1),
+            "dsplit": (cosmo_obect(np.arange(27)).reshape(3, 3, 3), 1),
+            "hsplit": (cosmo_obect(np.arange(3)), 1),
+            "vsplit": (cosmo_obect(np.eye(3)), 1),
+            "swapaxes": (cosmo_obect(np.eye(3)), 0, 1),
+            "moveaxis": (cosmo_obect(np.eye(3)), 0, 1),
+            "nansum": (cosmo_obect(np.arange(3)),),  # implemented via sum
+            "std": (cosmo_obect(np.arange(3)),),
+            "nanstd": (cosmo_obect(np.arange(3)),),
+            "nanvar": (cosmo_obect(np.arange(3)),),
+            "nanprod": (cosmo_obect(np.arange(3)),),
+            "diag": (cosmo_obect(np.eye(3)),),
+            "diag_indices_from": (cosmo_obect(np.eye(3)),),
+            "diagflat": (cosmo_obect(np.eye(3)),),
+            "diagonal": (cosmo_obect(np.eye(3)),),
+            "ravel": (cosmo_obect(np.arange(3)),),
             "ravel_multi_index": (np.eye(2, dtype=int), (2, 2)),
             "unravel_index": (np.arange(3), (3,)),
-            "fix": (ca(np.arange(3)),),
-            "round": (ca(np.arange(3)),),  # implemented via around
-            "may_share_memory": (ca(np.arange(3)), ca(np.arange(3))),
-            "linalg.matrix_power": (ca(np.eye(3)), 2),
-            "linalg.cholesky": (ca(np.eye(3)),),
-            "linalg.multi_dot": ((ca(np.eye(3)), ca(np.eye(3))),),
-            "linalg.matrix_rank": (ca(np.eye(3)),),
-            "linalg.qr": (ca(np.eye(3)),),
-            "linalg.slogdet": (ca(np.eye(3)),),
-            "linalg.cond": (ca(np.eye(3)),),
-            "gradient": (ca(np.arange(3)),),
-            "cumsum": (ca(np.arange(3)),),
-            "nancumsum": (ca(np.arange(3)),),
-            "nancumprod": (ca(np.arange(3)),),
-            "bincount": (ca(np.arange(3)),),
-            "unique": (ca(np.arange(3)),),
-            "min_scalar_type": (ca(np.arange(3)),),
-            "extract": (0, ca(np.arange(3))),
-            "setxor1d": (ca(np.arange(3)), ca(np.arange(3))),
-            "lexsort": (ca(np.arange(3)),),
-            "digitize": (ca(np.arange(3)), ca(np.arange(3))),
-            "tril_indices_from": (ca(np.eye(3)),),
-            "triu_indices_from": (ca(np.eye(3)),),
-            "imag": (ca(np.arange(3)),),
-            "real": (ca(np.arange(3)),),
-            "real_if_close": (ca(np.arange(3)),),
-            "einsum_path": ("ij,jk->ik", ca(np.eye(3)), ca(np.eye(3))),
-            "cov": (ca(np.arange(3)),),
-            "corrcoef": (ca(np.arange(3)),),
-            "compress": (np.zeros(3), ca(np.arange(3))),
-            "take_along_axis": (ca(np.arange(3)), np.ones(3, dtype=int), 0),
-            "linalg.cross": (ca(np.arange(3)), ca(np.arange(3))),
-            "linalg.diagonal": (ca(np.eye(3)),),
-            "linalg.matmul": (ca(np.eye(3)), ca(np.eye(3))),
-            "linalg.matrix_norm": (ca(np.eye(3)),),
-            "linalg.matrix_transpose": (ca(np.eye(3)),),
-            "linalg.svdvals": (ca(np.eye(3)),),
-            "linalg.tensordot": (ca(np.eye(3)), ca(np.eye(3))),
-            "linalg.trace": (ca(np.eye(3)),),
-            "linalg.vecdot": (ca(np.arange(3)), ca(np.arange(3))),
-            "linalg.vector_norm": (ca(np.arange(3)),),
-            "astype": (ca(np.arange(3)), float),
-            "matrix_transpose": (ca(np.eye(3)),),
-            "unique_all": (ca(np.arange(3)),),
-            "unique_counts": (ca(np.arange(3)),),
-            "unique_inverse": (ca(np.arange(3)),),
-            "unique_values": (ca(np.arange(3)),),
-            "cumulative_sum": (ca(np.arange(3)),),
-            "cumulative_prod": (ca(np.arange(3)),),
-            "unstack": (ca(np.arange(3)),),
+            "fix": (cosmo_obect(np.arange(3)),),
+            "round": (cosmo_obect(np.arange(3)),),  # implemented via around
+            "may_share_memory": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "linalg.matrix_power": (cosmo_obect(np.eye(3)), 2),
+            "linalg.cholesky": (cosmo_obect(np.eye(3)),),
+            "linalg.multi_dot": ((cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),),
+            "linalg.matrix_rank": (cosmo_obect(np.eye(3)),),
+            "linalg.qr": (cosmo_obect(np.eye(3)),),
+            "linalg.slogdet": (cosmo_obect(np.eye(3)),),
+            "linalg.cond": (cosmo_obect(np.eye(3)),),
+            "gradient": (cosmo_obect(np.arange(3)),),
+            "cumsum": (cosmo_obect(np.arange(3)),),
+            "nancumsum": (cosmo_obect(np.arange(3)),),
+            "nancumprod": (cosmo_obect(np.arange(3)),),
+            "bincount": (cosmo_obect(np.arange(3)),),
+            "unique": (cosmo_obect(np.arange(3)),),
+            "min_scalar_type": (cosmo_obect(np.arange(3)),),
+            "extract": (0, cosmo_obect(np.arange(3))),
+            "setxor1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "lexsort": (cosmo_obect(np.arange(3)),),
+            "digitize": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "tril_indices_from": (cosmo_obect(np.eye(3)),),
+            "triu_indices_from": (cosmo_obect(np.eye(3)),),
+            "imag": (cosmo_obect(np.arange(3)),),
+            "real": (cosmo_obect(np.arange(3)),),
+            "real_if_close": (cosmo_obect(np.arange(3)),),
+            "einsum_path": (
+                "ij,jk->ik",
+                cosmo_obect(np.eye(3)),
+                cosmo_obect(np.eye(3)),
+            ),
+            "cov": (cosmo_obect(np.arange(3)),),
+            "corrcoef": (cosmo_obect(np.arange(3)),),
+            "compress": (np.zeros(3), cosmo_obect(np.arange(3))),
+            "take_along_axis": (cosmo_obect(np.arange(3)), np.ones(3, dtype=int), 0),
+            "linalg.cross": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "linalg.diagonal": (cosmo_obect(np.eye(3)),),
+            "linalg.matmul": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
+            "linalg.matrix_norm": (cosmo_obect(np.eye(3)),),
+            "linalg.matrix_transpose": (cosmo_obect(np.eye(3)),),
+            "linalg.svdvals": (cosmo_obect(np.eye(3)),),
+            "linalg.tensordot": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
+            "linalg.trace": (cosmo_obect(np.eye(3)),),
+            "linalg.vecdot": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "linalg.vector_norm": (cosmo_obect(np.arange(3)),),
+            "astype": (cosmo_obect(np.arange(3)), float),
+            "matrix_transpose": (cosmo_obect(np.eye(3)),),
+            "unique_all": (cosmo_obect(np.arange(3)),),
+            "unique_counts": (cosmo_obect(np.arange(3)),),
+            "unique_inverse": (cosmo_obect(np.arange(3)),),
+            "unique_values": (cosmo_obect(np.arange(3)),),
+            "cumulative_sum": (cosmo_obect(np.arange(3)),),
+            "cumulative_prod": (cosmo_obect(np.arange(3)),),
+            "unstack": (cosmo_obect(np.arange(3)),),
         }
         functions_checked = list()
         bad_funcs = dict()
@@ -685,7 +713,8 @@ def test_block_is_broken(self):
         - Remove this test.
         """
         assert isinstance(
-            np.block([[ca(np.arange(3))], [ca(np.arange(3))]]), cosmo_array
+            np.block([[cosmo_obect(np.arange(3))], [cosmo_obect(np.arange(3))]]),
+            cosmo_array,
         )
 
     # the combinations of units and cosmo_factors is nonsense but it's just for testing...
@@ -763,7 +792,7 @@ def test_block_is_broken(self):
             np.array([1, 2, 3]),
         ),
     )
-    @pytest.mark.parametrize("bins_type", ("int", "np", "ca"))
+    @pytest.mark.parametrize("bins_type", ("int", "np", "cosmo_obect"))
     @pytest.mark.parametrize("density", (None, True))
     def test_histograms(self, func_args, weights, bins_type, density):
         """
@@ -777,7 +806,7 @@ def test_histograms(self, func_args, weights, bins_type, density):
         bins = {
             "int": 10,
             "np": [np.linspace(0, 5, 11)] * 3,
-            "ca": [
+            "cosmo_obect": [
                 cosmo_array(
                     np.linspace(0, 5, 11),
                     u.kpc,
@@ -809,7 +838,7 @@ def test_histograms(self, func_args, weights, bins_type, density):
                     np.histogramdd: np.s_[:],
                 }[func]
             ]
-            if bins_type in ("np", "ca")
+            if bins_type in ("np", "cosmo_obect")
             else bins
         )
         result = func(*args, bins=bins, density=density, weights=weights)
@@ -891,26 +920,26 @@ def test_getitem(self):
         """
         Make sure that we don't degrade to an ndarray on slicing.
         """
-        assert isinstance(ca(np.arange(3))[0], cosmo_quantity)
+        assert isinstance(cosmo_obect(np.arange(3))[0], cosmo_quantity)
 
     def test_reshape_to_scalar(self):
         """
         Make sure that we convert to a cosmo_quantity when we reshape to a scalar.
         """
-        assert isinstance(ca(np.ones(1)).reshape(tuple()), cosmo_quantity)
+        assert isinstance(cosmo_obect(np.ones(1)).reshape(tuple()), cosmo_quantity)
 
     def test_iter(self):
         """
         Make sure that we get cosmo_quantity's when iterating over a cosmo_array.
         """
-        for cq in ca(np.arange(3)):
+        for cq in cosmo_obect(np.arange(3)):
             assert isinstance(cq, cosmo_quantity)
 
     def test_dot(self):
         """
         Make sure that we get a cosmo_array when we use array attribute dot.
         """
-        res = ca(np.arange(3)).dot(ca(np.arange(3)))
+        res = cosmo_obect(np.arange(3)).dot(cosmo_obect(np.arange(3)))
         assert isinstance(res, cosmo_quantity)
         assert res.comoving is False
         assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5)
@@ -979,7 +1008,7 @@ def test_scalar_return_func(self):
         Make sure that default-wrapped functions that take a cosmo_array and return a
         scalar convert to a cosmo_quantity.
         """
-        ca = cosmo_array(
+        cosmo_obect = cosmo_array(
             np.arange(3),
             u.m,
             comoving=False,
@@ -987,7 +1016,7 @@ def test_scalar_return_func(self):
             scale_exponent=1,
             valid_transform=True,
         )
-        res = np.min(ca)
+        res = np.min(cosmo_obect)
         assert isinstance(res, cosmo_quantity)
 
     @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"])
@@ -1008,6 +1037,24 @@ def test_propagation_props(self, prop):
         assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
+    def test_multiply_quantities(self):
+        """
+        Test multiplying two quantities.
+        """
+        cq = cosmo_quantity(
+            2,
+            u.m,
+            comoving=False,
+            scale_factor=0.5,
+            scale_exponent=1,
+            valid_transform=True,
+        )
+        multiplied = cq * cq
+        assert type(multiplied) is cosmo_quantity
+        assert multiplied.comoving is False
+        assert multiplied.cosmo_factor == cosmo_factor(a ** 2, 0.5)
+        assert multiplied.to_value(u.m ** 2) == 4
+
 
 class TestCosmoArrayCopy:
     """
@@ -1067,20 +1114,27 @@ def test_to_cgs(self):
 
 
 class TestMultiplicationByUnyt:
-    def test_multiplication_by_unyt(self):
+    @pytest.mark.parametrize(
+        "cosmo_object",
+        [
+            cosmo_array(
+                np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
+            ),
+            cosmo_quantity(
+                np.ones(1), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
+            ),
+        ],
+    )
+    def test_multiplication_by_unyt(self, cosmo_object):
         """
         We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as
         for `cosmo_array(...) * u.Mpc`.
-
-        Right-sided multiplication & division can't be supported without upstream
-        changes in unyt, see `test_rmultiplication_by_unyt`.
         """
-        ca = cosmo_array(
-            np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
-        )
 
-        lmultiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
-        lmultiplied_by_unyt = ca * u.Mpc
+        lmultiplied_by_quantity = cosmo_object * (
+            1 * u.Mpc
+        )  # parentheses very important here
+        lmultiplied_by_unyt = cosmo_object * u.Mpc
         assert isinstance(lmultiplied_by_quantity, cosmo_array)
         assert isinstance(lmultiplied_by_unyt, cosmo_array)
         assert lmultiplied_by_unyt.comoving == lmultiplied_by_quantity.comoving
@@ -1089,8 +1143,10 @@ def test_multiplication_by_unyt(self):
             lmultiplied_by_quantity.to_value(lmultiplied_by_quantity.units),
         )
 
-        ldivided_by_quantity = ca / (1 * u.Mpc)
-        ldivided_by_unyt = ca / u.Mpc
+        ldivided_by_quantity = cosmo_object / (
+            1 * u.Mpc
+        )  # parentheses very important here
+        ldivided_by_unyt = cosmo_object / u.Mpc
         assert isinstance(ldivided_by_quantity, cosmo_array)
         assert isinstance(ldivided_by_unyt, cosmo_array)
         assert ldivided_by_unyt.comoving == ldivided_by_quantity.comoving
@@ -1099,32 +1155,11 @@ def test_multiplication_by_unyt(self):
             ldivided_by_quantity.to_value(ldivided_by_quantity.units),
         )
 
-    # @pytest.mark.xfail
-    def test_rmultiplication_by_unyt(self):
-        """
-        We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as
-        for `cosmo_array(...) * u.Mpc`.
-
-        But unyt will call it's own __mul__ before we get a chance to use our __rmul__
-        when the cosmo_array is the right-hand argument.
-
-        We can't handle this case without upstream changes in unyt, so this test is marked
-        to xfail.
-
-        If this is fixed in the future this test will pass and can be merged with
-        `test_multiplication_by_unyt` to tidy up. Also clean up docstrings of
-        `cosmo_array.__mul__`, `cosmo_array.__rmul__`, `cosmo_array.__truediv__`,
-        `cosmo_array.__rtruediv__`.
-
-        See https://github.com/yt-project/unyt/pull/572
-        """
-        ca = cosmo_array(
-            np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
-        )
-
-        rmultiplied_by_quantity = (1 * u.Mpc) * ca  # parentheses very important here
+        rmultiplied_by_quantity = (
+            1 * u.Mpc
+        ) * cosmo_object  # parentheses very important here
         assert rmultiplied_by_quantity.comoving
-        rmultiplied_by_unyt = u.Mpc * ca
+        rmultiplied_by_unyt = u.Mpc * cosmo_object
         assert isinstance(rmultiplied_by_quantity, cosmo_array)
         assert isinstance(rmultiplied_by_unyt, cosmo_array)
         assert rmultiplied_by_unyt.comoving == rmultiplied_by_quantity.comoving
@@ -1133,8 +1168,10 @@ def test_rmultiplication_by_unyt(self):
             rmultiplied_by_quantity.to_value(rmultiplied_by_quantity.units),
         )
 
-        rdivided_by_quantity = (1 * u.Mpc) / ca  # parentheses very important here
-        rdivided_by_unyt = u.Mpc / ca
+        rdivided_by_quantity = (
+            1 * u.Mpc
+        ) / cosmo_object  # parentheses very important here
+        rdivided_by_unyt = u.Mpc / cosmo_object
         assert isinstance(rdivided_by_quantity, cosmo_array)
         assert isinstance(rdivided_by_unyt, cosmo_array)
         assert rdivided_by_unyt.comoving == rdivided_by_quantity.comoving

From 3843127977a048cd6ace7a8a794a44ab424392c1 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 10 Oct 2025 17:30:25 +0100
Subject: [PATCH 04/23] Fix typo.

---
 tests/test_cosmo_array.py | 568 +++++++++++++++++++-------------------
 1 file changed, 289 insertions(+), 279 deletions(-)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index c06a56a3..67751d32 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -24,7 +24,7 @@ def getfunc(fname):
     return func
 
 
-def cosmo_obect(x, unit=u.Mpc):
+def cosmo_object(x, unit=u.Mpc):
     """
     Helper for our tests: turn an array into a cosmo_array.
     """
@@ -120,7 +120,7 @@ def test_init_from_ndarray(self):
         arr = cosmo_array(
             np.ones(5),
             units=u.Mpc,
-            cosmo_factor=cosmo_factor(a ** 1, 1.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -145,7 +145,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.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -168,7 +168,7 @@ def test_init_from_unyt_array(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             u.unyt_array(np.ones(5), units=u.Mpc),
-            cosmo_factor=cosmo_factor(a ** 1, 1.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -194,7 +194,7 @@ def test_init_from_list_of_unyt_arrays(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             [u.unyt_array(1, units=u.Mpc) for _ in range(5)],
-            cosmo_factor=cosmo_factor(a ** 1, 1.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -218,7 +218,7 @@ def test_init_from_list_of_cosmo_arrays(self):
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a ** 1, 1
+            a**1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
@@ -228,14 +228,14 @@ def test_init_from_list_of_cosmo_arrays(self):
                     [1],
                     units=u.Mpc,
                     comoving=False,
-                    cosmo_factor=cosmo_factor(a ** 1, 1.0),
+                    cosmo_factor=cosmo_factor(a**1, 1.0),
                 )
                 for _ in range(5)
             ]
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a ** 1, 1
+            a**1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
 
@@ -329,281 +329,291 @@ def test_explicitly_handled_funcs(self):
 
         functions_to_check = {
             # FUNCTIONS UNYT HANDLES EXPLICITLY:
-            "array2string": (cosmo_obect(np.arange(3)),),
-            "dot": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "vdot": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "inner": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "outer": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "kron": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "histogram_bin_edges": (cosmo_obect(np.arange(3)),),
-            "linalg.inv": (cosmo_obect(np.eye(3)),),
-            "linalg.tensorinv": (cosmo_obect(np.eye(9).reshape((3, 3, 3, 3))),),
-            "linalg.pinv": (cosmo_obect(np.eye(3)),),
-            "linalg.svd": (cosmo_obect(np.eye(3)),),
-            "histogram": (cosmo_obect(np.arange(3)),),
-            "histogram2d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "histogramdd": (cosmo_obect(np.arange(3)).reshape((1, 3)),),
-            "concatenate": (cosmo_obect(np.eye(3)),),
-            "cross": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "intersect1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "union1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "linalg.norm": (cosmo_obect(np.arange(3)),),
-            "vstack": (cosmo_obect(np.arange(3)),),
-            "hstack": (cosmo_obect(np.arange(3)),),
-            "dstack": (cosmo_obect(np.arange(3)),),
-            "column_stack": (cosmo_obect(np.arange(3)),),
-            "stack": (cosmo_obect(np.arange(3)),),
-            "around": (cosmo_obect(np.arange(3)),),
-            "block": ([[cosmo_obect(np.arange(3))], [cosmo_obect(np.arange(3))]],),
-            "fft.fft": (cosmo_obect(np.arange(3)),),
-            "fft.fft2": (cosmo_obect(np.eye(3)),),
-            "fft.fftn": (cosmo_obect(np.arange(3)),),
-            "fft.hfft": (cosmo_obect(np.arange(3)),),
-            "fft.rfft": (cosmo_obect(np.arange(3)),),
-            "fft.rfft2": (cosmo_obect(np.eye(3)),),
-            "fft.rfftn": (cosmo_obect(np.arange(3)),),
-            "fft.ifft": (cosmo_obect(np.arange(3)),),
-            "fft.ifft2": (cosmo_obect(np.eye(3)),),
-            "fft.ifftn": (cosmo_obect(np.arange(3)),),
-            "fft.ihfft": (cosmo_obect(np.arange(3)),),
-            "fft.irfft": (cosmo_obect(np.arange(3)),),
-            "fft.irfft2": (cosmo_obect(np.eye(3)),),
-            "fft.irfftn": (cosmo_obect(np.arange(3)),),
-            "fft.fftshift": (cosmo_obect(np.arange(3)),),
-            "fft.ifftshift": (cosmo_obect(np.arange(3)),),
-            "sort_complex": (cosmo_obect(np.arange(3)),),
-            "isclose": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "allclose": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "array_equal": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "array_equiv": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "array2string": (cosmo_object(np.arange(3)),),
+            "dot": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "vdot": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "inner": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "outer": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "kron": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "histogram_bin_edges": (cosmo_object(np.arange(3)),),
+            "linalg.inv": (cosmo_object(np.eye(3)),),
+            "linalg.tensorinv": (cosmo_object(np.eye(9).reshape((3, 3, 3, 3))),),
+            "linalg.pinv": (cosmo_object(np.eye(3)),),
+            "linalg.svd": (cosmo_object(np.eye(3)),),
+            "histogram": (cosmo_object(np.arange(3)),),
+            "histogram2d": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "histogramdd": (cosmo_object(np.arange(3)).reshape((1, 3)),),
+            "concatenate": (cosmo_object(np.eye(3)),),
+            "cross": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "intersect1d": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "union1d": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "linalg.norm": (cosmo_object(np.arange(3)),),
+            "vstack": (cosmo_object(np.arange(3)),),
+            "hstack": (cosmo_object(np.arange(3)),),
+            "dstack": (cosmo_object(np.arange(3)),),
+            "column_stack": (cosmo_object(np.arange(3)),),
+            "stack": (cosmo_object(np.arange(3)),),
+            "around": (cosmo_object(np.arange(3)),),
+            "block": ([[cosmo_object(np.arange(3))], [cosmo_object(np.arange(3))]],),
+            "fft.fft": (cosmo_object(np.arange(3)),),
+            "fft.fft2": (cosmo_object(np.eye(3)),),
+            "fft.fftn": (cosmo_object(np.arange(3)),),
+            "fft.hfft": (cosmo_object(np.arange(3)),),
+            "fft.rfft": (cosmo_object(np.arange(3)),),
+            "fft.rfft2": (cosmo_object(np.eye(3)),),
+            "fft.rfftn": (cosmo_object(np.arange(3)),),
+            "fft.ifft": (cosmo_object(np.arange(3)),),
+            "fft.ifft2": (cosmo_object(np.eye(3)),),
+            "fft.ifftn": (cosmo_object(np.arange(3)),),
+            "fft.ihfft": (cosmo_object(np.arange(3)),),
+            "fft.irfft": (cosmo_object(np.arange(3)),),
+            "fft.irfft2": (cosmo_object(np.eye(3)),),
+            "fft.irfftn": (cosmo_object(np.arange(3)),),
+            "fft.fftshift": (cosmo_object(np.arange(3)),),
+            "fft.ifftshift": (cosmo_object(np.arange(3)),),
+            "sort_complex": (cosmo_object(np.arange(3)),),
+            "isclose": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "allclose": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "array_equal": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "array_equiv": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
             "linspace": (cq(1), cq(2)),
             "logspace": (cq(1, unit=u.dimensionless), cq(2, unit=u.dimensionless)),
             "geomspace": (cq(1), cq(1)),
-            "copyto": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "prod": (cosmo_obect(np.arange(3)),),
-            "var": (cosmo_obect(np.arange(3)),),
-            "trace": (cosmo_obect(np.eye(3)),),
-            "percentile": (cosmo_obect(np.arange(3)), 30),
-            "quantile": (cosmo_obect(np.arange(3)), 0.3),
-            "nanpercentile": (cosmo_obect(np.arange(3)), 30),
-            "nanquantile": (cosmo_obect(np.arange(3)), 0.3),
-            "linalg.det": (cosmo_obect(np.eye(3)),),
-            "diff": (cosmo_obect(np.arange(3)),),
-            "ediff1d": (cosmo_obect(np.arange(3)),),
-            "ptp": (cosmo_obect(np.arange(3)),),
-            "cumprod": (cosmo_obect(np.arange(3)),),
-            "pad": (cosmo_obect(np.arange(3)), 3),
-            "choose": (np.arange(3), cosmo_obect(np.eye(3))),
-            "insert": (cosmo_obect(np.arange(3)), 1, cq(1)),
-            "linalg.lstsq": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
-            "linalg.solve": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
+            "copyto": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "prod": (cosmo_object(np.arange(3)),),
+            "var": (cosmo_object(np.arange(3)),),
+            "trace": (cosmo_object(np.eye(3)),),
+            "percentile": (cosmo_object(np.arange(3)), 30),
+            "quantile": (cosmo_object(np.arange(3)), 0.3),
+            "nanpercentile": (cosmo_object(np.arange(3)), 30),
+            "nanquantile": (cosmo_object(np.arange(3)), 0.3),
+            "linalg.det": (cosmo_object(np.eye(3)),),
+            "diff": (cosmo_object(np.arange(3)),),
+            "ediff1d": (cosmo_object(np.arange(3)),),
+            "ptp": (cosmo_object(np.arange(3)),),
+            "cumprod": (cosmo_object(np.arange(3)),),
+            "pad": (cosmo_object(np.arange(3)), 3),
+            "choose": (np.arange(3), cosmo_object(np.eye(3))),
+            "insert": (cosmo_object(np.arange(3)), 1, cq(1)),
+            "linalg.lstsq": (cosmo_object(np.eye(3)), cosmo_object(np.eye(3))),
+            "linalg.solve": (cosmo_object(np.eye(3)), cosmo_object(np.eye(3))),
             "linalg.tensorsolve": (
-                cosmo_obect(np.eye(24).reshape((6, 4, 2, 3, 4))),
-                cosmo_obect(np.ones((6, 4))),
+                cosmo_object(np.eye(24).reshape((6, 4, 2, 3, 4))),
+                cosmo_object(np.ones((6, 4))),
             ),
-            "linalg.eig": (cosmo_obect(np.eye(3)),),
-            "linalg.eigh": (cosmo_obect(np.eye(3)),),
-            "linalg.eigvals": (cosmo_obect(np.eye(3)),),
-            "linalg.eigvalsh": (cosmo_obect(np.eye(3)),),
-            "savetxt": (savetxt_file, cosmo_obect(np.arange(3))),
-            "fill_diagonal": (cosmo_obect(np.eye(3)), cosmo_obect(np.arange(3))),
-            "apply_over_axes": (lambda x, axis: x, cosmo_obect(np.eye(3)), (0, 1)),
-            "isin": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "linalg.eig": (cosmo_object(np.eye(3)),),
+            "linalg.eigh": (cosmo_object(np.eye(3)),),
+            "linalg.eigvals": (cosmo_object(np.eye(3)),),
+            "linalg.eigvalsh": (cosmo_object(np.eye(3)),),
+            "savetxt": (savetxt_file, cosmo_object(np.arange(3))),
+            "fill_diagonal": (cosmo_object(np.eye(3)), cosmo_object(np.arange(3))),
+            "apply_over_axes": (lambda x, axis: x, cosmo_object(np.eye(3)), (0, 1)),
+            "isin": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
             "place": (
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
                 np.arange(3) > 0,
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
+            ),
+            "put": (
+                cosmo_object(np.arange(3)),
+                np.arange(3),
+                cosmo_object(np.arange(3)),
             ),
-            "put": (cosmo_obect(np.arange(3)), np.arange(3), cosmo_obect(np.arange(3))),
             "put_along_axis": (
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
                 np.arange(3),
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
                 0,
             ),
             "putmask": (
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
                 np.arange(3),
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
             ),
-            "searchsorted": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
+            "searchsorted": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
             "select": (
                 [np.arange(3) < 1, np.arange(3) > 1],
-                [cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))],
+                [cosmo_object(np.arange(3)), cosmo_object(np.arange(3))],
                 cq(1),
             ),
-            "setdiff1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3, 6))),
-            "sinc": (cosmo_obect(np.arange(3)),),
-            "clip": (cosmo_obect(np.arange(3)), cq(1), cq(2)),
+            "setdiff1d": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3, 6))),
+            "sinc": (cosmo_object(np.arange(3)),),
+            "clip": (cosmo_object(np.arange(3)), cq(1), cq(2)),
             "where": (
-                cosmo_obect(np.arange(3)),
-                cosmo_obect(np.arange(3)),
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
+                cosmo_object(np.arange(3)),
+                cosmo_object(np.arange(3)),
             ),
-            "triu": (cosmo_obect(np.ones((3, 3))),),
-            "tril": (cosmo_obect(np.ones((3, 3))),),
-            "einsum": ("ii->i", cosmo_obect(np.eye(3))),
-            "convolve": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "correlate": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "tensordot": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
-            "unwrap": (cosmo_obect(np.arange(3)),),
+            "triu": (cosmo_object(np.ones((3, 3))),),
+            "tril": (cosmo_object(np.ones((3, 3))),),
+            "einsum": ("ii->i", cosmo_object(np.eye(3))),
+            "convolve": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "correlate": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "tensordot": (cosmo_object(np.eye(3)), cosmo_object(np.eye(3))),
+            "unwrap": (cosmo_object(np.arange(3)),),
             "interp": (
-                cosmo_obect(np.arange(3)),
-                cosmo_obect(np.arange(3)),
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
+                cosmo_object(np.arange(3)),
+                cosmo_object(np.arange(3)),
             ),
-            "array_repr": (cosmo_obect(np.arange(3)),),
-            "linalg.outer": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "trapezoid": (cosmo_obect(np.arange(3)),),
+            "array_repr": (cosmo_object(np.arange(3)),),
+            "linalg.outer": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "trapezoid": (cosmo_object(np.arange(3)),),
             "in1d": (
-                cosmo_obect(np.arange(3)),
-                cosmo_obect(np.arange(3)),
+                cosmo_object(np.arange(3)),
+                cosmo_object(np.arange(3)),
             ),  # np deprecated
-            "take": (cosmo_obect(np.arange(3)), np.arange(3)),
+            "take": (cosmo_object(np.arange(3)), np.arange(3)),
             # FUNCTIONS THAT UNYT DOESN'T HANDLE EXPLICITLY (THEY "JUST WORK"):
-            "all": (cosmo_obect(np.arange(3)),),
-            "amax": (cosmo_obect(np.arange(3)),),  # implemented via max
-            "amin": (cosmo_obect(np.arange(3)),),  # implemented via min
+            "all": (cosmo_object(np.arange(3)),),
+            "amax": (cosmo_object(np.arange(3)),),  # implemented via max
+            "amin": (cosmo_object(np.arange(3)),),  # implemented via min
             "angle": (cq(complex(1, 1)),),
-            "any": (cosmo_obect(np.arange(3)),),
-            "append": (cosmo_obect(np.arange(3)), cq(1)),
-            "apply_along_axis": (lambda x: x, 0, cosmo_obect(np.eye(3))),
-            "argmax": (cosmo_obect(np.arange(3)),),  # implemented via max
-            "argmin": (cosmo_obect(np.arange(3)),),  # implemented via min
-            "argpartition": (cosmo_obect(np.arange(3)), 1),  # implemented via partition
-            "argsort": (cosmo_obect(np.arange(3)),),  # implemented via sort
-            "argwhere": (cosmo_obect(np.arange(3)),),
-            "array_str": (cosmo_obect(np.arange(3)),),
-            "atleast_1d": (cosmo_obect(np.arange(3)),),
-            "atleast_2d": (cosmo_obect(np.arange(3)),),
-            "atleast_3d": (cosmo_obect(np.arange(3)),),
-            "average": (cosmo_obect(np.arange(3)),),
-            "can_cast": (cosmo_obect(np.arange(3)), np.float64),
-            "common_type": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "result_type": (cosmo_obect(np.ones(3)), cosmo_obect(np.ones(3))),
-            "iscomplex": (cosmo_obect(np.arange(3)),),
-            "iscomplexobj": (cosmo_obect(np.arange(3)),),
-            "isreal": (cosmo_obect(np.arange(3)),),
-            "isrealobj": (cosmo_obect(np.arange(3)),),
-            "nan_to_num": (cosmo_obect(np.arange(3)),),
-            "nanargmax": (cosmo_obect(np.arange(3)),),  # implemented via max
-            "nanargmin": (cosmo_obect(np.arange(3)),),  # implemented via min
-            "nanmax": (cosmo_obect(np.arange(3)),),  # implemented via max
-            "nanmean": (cosmo_obect(np.arange(3)),),  # implemented via mean
-            "nanmedian": (cosmo_obect(np.arange(3)),),  # implemented via median
-            "nanmin": (cosmo_obect(np.arange(3)),),  # implemented via min
-            "trim_zeros": (cosmo_obect(np.arange(3)),),
-            "max": (cosmo_obect(np.arange(3)),),
-            "mean": (cosmo_obect(np.arange(3)),),
-            "median": (cosmo_obect(np.arange(3)),),
-            "min": (cosmo_obect(np.arange(3)),),
-            "ndim": (cosmo_obect(np.arange(3)),),
-            "shape": (cosmo_obect(np.arange(3)),),
-            "size": (cosmo_obect(np.arange(3)),),
-            "sort": (cosmo_obect(np.arange(3)),),
-            "sum": (cosmo_obect(np.arange(3)),),
-            "repeat": (cosmo_obect(np.arange(3)), 2),
-            "tile": (cosmo_obect(np.arange(3)), 2),
-            "shares_memory": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "nonzero": (cosmo_obect(np.arange(3)),),
-            "count_nonzero": (cosmo_obect(np.arange(3)),),
-            "flatnonzero": (cosmo_obect(np.arange(3)),),
-            "isneginf": (cosmo_obect(np.arange(3)),),
-            "isposinf": (cosmo_obect(np.arange(3)),),
-            "empty_like": (cosmo_obect(np.arange(3)),),
-            "full_like": (cosmo_obect(np.arange(3)), cq(1)),
-            "ones_like": (cosmo_obect(np.arange(3)),),
-            "zeros_like": (cosmo_obect(np.arange(3)),),
-            "copy": (cosmo_obect(np.arange(3)),),
-            "meshgrid": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "transpose": (cosmo_obect(np.eye(3)),),
-            "reshape": (cosmo_obect(np.arange(3)), (3,)),
-            "resize": (cosmo_obect(np.arange(3)), 6),
-            "roll": (cosmo_obect(np.arange(3)), 1),
-            "rollaxis": (cosmo_obect(np.arange(3)), 0),
-            "rot90": (cosmo_obect(np.eye(3)),),
-            "expand_dims": (cosmo_obect(np.arange(3)), 0),
-            "squeeze": (cosmo_obect(np.arange(3)),),
-            "flip": (cosmo_obect(np.eye(3)),),
-            "fliplr": (cosmo_obect(np.eye(3)),),
-            "flipud": (cosmo_obect(np.eye(3)),),
-            "delete": (cosmo_obect(np.arange(3)), 0),
-            "partition": (cosmo_obect(np.arange(3)), 1),
-            "broadcast_to": (cosmo_obect(np.arange(3)), 3),
-            "broadcast_arrays": (cosmo_obect(np.arange(3)),),
-            "split": (cosmo_obect(np.arange(3)), 1),
-            "array_split": (cosmo_obect(np.arange(3)), 1),
-            "dsplit": (cosmo_obect(np.arange(27)).reshape(3, 3, 3), 1),
-            "hsplit": (cosmo_obect(np.arange(3)), 1),
-            "vsplit": (cosmo_obect(np.eye(3)), 1),
-            "swapaxes": (cosmo_obect(np.eye(3)), 0, 1),
-            "moveaxis": (cosmo_obect(np.eye(3)), 0, 1),
-            "nansum": (cosmo_obect(np.arange(3)),),  # implemented via sum
-            "std": (cosmo_obect(np.arange(3)),),
-            "nanstd": (cosmo_obect(np.arange(3)),),
-            "nanvar": (cosmo_obect(np.arange(3)),),
-            "nanprod": (cosmo_obect(np.arange(3)),),
-            "diag": (cosmo_obect(np.eye(3)),),
-            "diag_indices_from": (cosmo_obect(np.eye(3)),),
-            "diagflat": (cosmo_obect(np.eye(3)),),
-            "diagonal": (cosmo_obect(np.eye(3)),),
-            "ravel": (cosmo_obect(np.arange(3)),),
+            "any": (cosmo_object(np.arange(3)),),
+            "append": (cosmo_object(np.arange(3)), cq(1)),
+            "apply_along_axis": (lambda x: x, 0, cosmo_object(np.eye(3))),
+            "argmax": (cosmo_object(np.arange(3)),),  # implemented via max
+            "argmin": (cosmo_object(np.arange(3)),),  # implemented via min
+            "argpartition": (
+                cosmo_object(np.arange(3)),
+                1,
+            ),  # implemented via partition
+            "argsort": (cosmo_object(np.arange(3)),),  # implemented via sort
+            "argwhere": (cosmo_object(np.arange(3)),),
+            "array_str": (cosmo_object(np.arange(3)),),
+            "atleast_1d": (cosmo_object(np.arange(3)),),
+            "atleast_2d": (cosmo_object(np.arange(3)),),
+            "atleast_3d": (cosmo_object(np.arange(3)),),
+            "average": (cosmo_object(np.arange(3)),),
+            "can_cast": (cosmo_object(np.arange(3)), np.float64),
+            "common_type": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "result_type": (cosmo_object(np.ones(3)), cosmo_object(np.ones(3))),
+            "iscomplex": (cosmo_object(np.arange(3)),),
+            "iscomplexobj": (cosmo_object(np.arange(3)),),
+            "isreal": (cosmo_object(np.arange(3)),),
+            "isrealobj": (cosmo_object(np.arange(3)),),
+            "nan_to_num": (cosmo_object(np.arange(3)),),
+            "nanargmax": (cosmo_object(np.arange(3)),),  # implemented via max
+            "nanargmin": (cosmo_object(np.arange(3)),),  # implemented via min
+            "nanmax": (cosmo_object(np.arange(3)),),  # implemented via max
+            "nanmean": (cosmo_object(np.arange(3)),),  # implemented via mean
+            "nanmedian": (cosmo_object(np.arange(3)),),  # implemented via median
+            "nanmin": (cosmo_object(np.arange(3)),),  # implemented via min
+            "trim_zeros": (cosmo_object(np.arange(3)),),
+            "max": (cosmo_object(np.arange(3)),),
+            "mean": (cosmo_object(np.arange(3)),),
+            "median": (cosmo_object(np.arange(3)),),
+            "min": (cosmo_object(np.arange(3)),),
+            "ndim": (cosmo_object(np.arange(3)),),
+            "shape": (cosmo_object(np.arange(3)),),
+            "size": (cosmo_object(np.arange(3)),),
+            "sort": (cosmo_object(np.arange(3)),),
+            "sum": (cosmo_object(np.arange(3)),),
+            "repeat": (cosmo_object(np.arange(3)), 2),
+            "tile": (cosmo_object(np.arange(3)), 2),
+            "shares_memory": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "nonzero": (cosmo_object(np.arange(3)),),
+            "count_nonzero": (cosmo_object(np.arange(3)),),
+            "flatnonzero": (cosmo_object(np.arange(3)),),
+            "isneginf": (cosmo_object(np.arange(3)),),
+            "isposinf": (cosmo_object(np.arange(3)),),
+            "empty_like": (cosmo_object(np.arange(3)),),
+            "full_like": (cosmo_object(np.arange(3)), cq(1)),
+            "ones_like": (cosmo_object(np.arange(3)),),
+            "zeros_like": (cosmo_object(np.arange(3)),),
+            "copy": (cosmo_object(np.arange(3)),),
+            "meshgrid": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "transpose": (cosmo_object(np.eye(3)),),
+            "reshape": (cosmo_object(np.arange(3)), (3,)),
+            "resize": (cosmo_object(np.arange(3)), 6),
+            "roll": (cosmo_object(np.arange(3)), 1),
+            "rollaxis": (cosmo_object(np.arange(3)), 0),
+            "rot90": (cosmo_object(np.eye(3)),),
+            "expand_dims": (cosmo_object(np.arange(3)), 0),
+            "squeeze": (cosmo_object(np.arange(3)),),
+            "flip": (cosmo_object(np.eye(3)),),
+            "fliplr": (cosmo_object(np.eye(3)),),
+            "flipud": (cosmo_object(np.eye(3)),),
+            "delete": (cosmo_object(np.arange(3)), 0),
+            "partition": (cosmo_object(np.arange(3)), 1),
+            "broadcast_to": (cosmo_object(np.arange(3)), 3),
+            "broadcast_arrays": (cosmo_object(np.arange(3)),),
+            "split": (cosmo_object(np.arange(3)), 1),
+            "array_split": (cosmo_object(np.arange(3)), 1),
+            "dsplit": (cosmo_object(np.arange(27)).reshape(3, 3, 3), 1),
+            "hsplit": (cosmo_object(np.arange(3)), 1),
+            "vsplit": (cosmo_object(np.eye(3)), 1),
+            "swapaxes": (cosmo_object(np.eye(3)), 0, 1),
+            "moveaxis": (cosmo_object(np.eye(3)), 0, 1),
+            "nansum": (cosmo_object(np.arange(3)),),  # implemented via sum
+            "std": (cosmo_object(np.arange(3)),),
+            "nanstd": (cosmo_object(np.arange(3)),),
+            "nanvar": (cosmo_object(np.arange(3)),),
+            "nanprod": (cosmo_object(np.arange(3)),),
+            "diag": (cosmo_object(np.eye(3)),),
+            "diag_indices_from": (cosmo_object(np.eye(3)),),
+            "diagflat": (cosmo_object(np.eye(3)),),
+            "diagonal": (cosmo_object(np.eye(3)),),
+            "ravel": (cosmo_object(np.arange(3)),),
             "ravel_multi_index": (np.eye(2, dtype=int), (2, 2)),
             "unravel_index": (np.arange(3), (3,)),
-            "fix": (cosmo_obect(np.arange(3)),),
-            "round": (cosmo_obect(np.arange(3)),),  # implemented via around
-            "may_share_memory": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "linalg.matrix_power": (cosmo_obect(np.eye(3)), 2),
-            "linalg.cholesky": (cosmo_obect(np.eye(3)),),
-            "linalg.multi_dot": ((cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),),
-            "linalg.matrix_rank": (cosmo_obect(np.eye(3)),),
-            "linalg.qr": (cosmo_obect(np.eye(3)),),
-            "linalg.slogdet": (cosmo_obect(np.eye(3)),),
-            "linalg.cond": (cosmo_obect(np.eye(3)),),
-            "gradient": (cosmo_obect(np.arange(3)),),
-            "cumsum": (cosmo_obect(np.arange(3)),),
-            "nancumsum": (cosmo_obect(np.arange(3)),),
-            "nancumprod": (cosmo_obect(np.arange(3)),),
-            "bincount": (cosmo_obect(np.arange(3)),),
-            "unique": (cosmo_obect(np.arange(3)),),
-            "min_scalar_type": (cosmo_obect(np.arange(3)),),
-            "extract": (0, cosmo_obect(np.arange(3))),
-            "setxor1d": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "lexsort": (cosmo_obect(np.arange(3)),),
-            "digitize": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "tril_indices_from": (cosmo_obect(np.eye(3)),),
-            "triu_indices_from": (cosmo_obect(np.eye(3)),),
-            "imag": (cosmo_obect(np.arange(3)),),
-            "real": (cosmo_obect(np.arange(3)),),
-            "real_if_close": (cosmo_obect(np.arange(3)),),
+            "fix": (cosmo_object(np.arange(3)),),
+            "round": (cosmo_object(np.arange(3)),),  # implemented via around
+            "may_share_memory": (
+                cosmo_object(np.arange(3)),
+                cosmo_object(np.arange(3)),
+            ),
+            "linalg.matrix_power": (cosmo_object(np.eye(3)), 2),
+            "linalg.cholesky": (cosmo_object(np.eye(3)),),
+            "linalg.multi_dot": ((cosmo_object(np.eye(3)), cosmo_object(np.eye(3))),),
+            "linalg.matrix_rank": (cosmo_object(np.eye(3)),),
+            "linalg.qr": (cosmo_object(np.eye(3)),),
+            "linalg.slogdet": (cosmo_object(np.eye(3)),),
+            "linalg.cond": (cosmo_object(np.eye(3)),),
+            "gradient": (cosmo_object(np.arange(3)),),
+            "cumsum": (cosmo_object(np.arange(3)),),
+            "nancumsum": (cosmo_object(np.arange(3)),),
+            "nancumprod": (cosmo_object(np.arange(3)),),
+            "bincount": (cosmo_object(np.arange(3)),),
+            "unique": (cosmo_object(np.arange(3)),),
+            "min_scalar_type": (cosmo_object(np.arange(3)),),
+            "extract": (0, cosmo_object(np.arange(3))),
+            "setxor1d": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "lexsort": (cosmo_object(np.arange(3)),),
+            "digitize": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "tril_indices_from": (cosmo_object(np.eye(3)),),
+            "triu_indices_from": (cosmo_object(np.eye(3)),),
+            "imag": (cosmo_object(np.arange(3)),),
+            "real": (cosmo_object(np.arange(3)),),
+            "real_if_close": (cosmo_object(np.arange(3)),),
             "einsum_path": (
                 "ij,jk->ik",
-                cosmo_obect(np.eye(3)),
-                cosmo_obect(np.eye(3)),
+                cosmo_object(np.eye(3)),
+                cosmo_object(np.eye(3)),
             ),
-            "cov": (cosmo_obect(np.arange(3)),),
-            "corrcoef": (cosmo_obect(np.arange(3)),),
-            "compress": (np.zeros(3), cosmo_obect(np.arange(3))),
-            "take_along_axis": (cosmo_obect(np.arange(3)), np.ones(3, dtype=int), 0),
-            "linalg.cross": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "linalg.diagonal": (cosmo_obect(np.eye(3)),),
-            "linalg.matmul": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
-            "linalg.matrix_norm": (cosmo_obect(np.eye(3)),),
-            "linalg.matrix_transpose": (cosmo_obect(np.eye(3)),),
-            "linalg.svdvals": (cosmo_obect(np.eye(3)),),
-            "linalg.tensordot": (cosmo_obect(np.eye(3)), cosmo_obect(np.eye(3))),
-            "linalg.trace": (cosmo_obect(np.eye(3)),),
-            "linalg.vecdot": (cosmo_obect(np.arange(3)), cosmo_obect(np.arange(3))),
-            "linalg.vector_norm": (cosmo_obect(np.arange(3)),),
-            "astype": (cosmo_obect(np.arange(3)), float),
-            "matrix_transpose": (cosmo_obect(np.eye(3)),),
-            "unique_all": (cosmo_obect(np.arange(3)),),
-            "unique_counts": (cosmo_obect(np.arange(3)),),
-            "unique_inverse": (cosmo_obect(np.arange(3)),),
-            "unique_values": (cosmo_obect(np.arange(3)),),
-            "cumulative_sum": (cosmo_obect(np.arange(3)),),
-            "cumulative_prod": (cosmo_obect(np.arange(3)),),
-            "unstack": (cosmo_obect(np.arange(3)),),
+            "cov": (cosmo_object(np.arange(3)),),
+            "corrcoef": (cosmo_object(np.arange(3)),),
+            "compress": (np.zeros(3), cosmo_object(np.arange(3))),
+            "take_along_axis": (cosmo_object(np.arange(3)), np.ones(3, dtype=int), 0),
+            "linalg.cross": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "linalg.diagonal": (cosmo_object(np.eye(3)),),
+            "linalg.matmul": (cosmo_object(np.eye(3)), cosmo_object(np.eye(3))),
+            "linalg.matrix_norm": (cosmo_object(np.eye(3)),),
+            "linalg.matrix_transpose": (cosmo_object(np.eye(3)),),
+            "linalg.svdvals": (cosmo_object(np.eye(3)),),
+            "linalg.tensordot": (cosmo_object(np.eye(3)), cosmo_object(np.eye(3))),
+            "linalg.trace": (cosmo_object(np.eye(3)),),
+            "linalg.vecdot": (cosmo_object(np.arange(3)), cosmo_object(np.arange(3))),
+            "linalg.vector_norm": (cosmo_object(np.arange(3)),),
+            "astype": (cosmo_object(np.arange(3)), float),
+            "matrix_transpose": (cosmo_object(np.eye(3)),),
+            "unique_all": (cosmo_object(np.arange(3)),),
+            "unique_counts": (cosmo_object(np.arange(3)),),
+            "unique_inverse": (cosmo_object(np.arange(3)),),
+            "unique_values": (cosmo_object(np.arange(3)),),
+            "cumulative_sum": (cosmo_object(np.arange(3)),),
+            "cumulative_prod": (cosmo_object(np.arange(3)),),
+            "unstack": (cosmo_object(np.arange(3)),),
         }
         functions_checked = list()
         bad_funcs = dict()
@@ -713,7 +723,7 @@ def test_block_is_broken(self):
         - Remove this test.
         """
         assert isinstance(
-            np.block([[cosmo_obect(np.arange(3))], [cosmo_obect(np.arange(3))]]),
+            np.block([[cosmo_object(np.arange(3))], [cosmo_object(np.arange(3))]]),
             cosmo_array,
         )
 
@@ -792,7 +802,7 @@ def test_block_is_broken(self):
             np.array([1, 2, 3]),
         ),
     )
-    @pytest.mark.parametrize("bins_type", ("int", "np", "cosmo_obect"))
+    @pytest.mark.parametrize("bins_type", ("int", "np", "cosmo_object"))
     @pytest.mark.parametrize("density", (None, True))
     def test_histograms(self, func_args, weights, bins_type, density):
         """
@@ -806,7 +816,7 @@ def test_histograms(self, func_args, weights, bins_type, density):
         bins = {
             "int": 10,
             "np": [np.linspace(0, 5, 11)] * 3,
-            "cosmo_obect": [
+            "cosmo_object": [
                 cosmo_array(
                     np.linspace(0, 5, 11),
                     u.kpc,
@@ -838,7 +848,7 @@ def test_histograms(self, func_args, weights, bins_type, density):
                     np.histogramdd: np.s_[:],
                 }[func]
             ]
-            if bins_type in ("np", "cosmo_obect")
+            if bins_type in ("np", "cosmo_object")
             else bins
         )
         result = func(*args, bins=bins, density=density, weights=weights)
@@ -873,9 +883,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a ** -1, 1.0),
-                    np.histogram2d: cosmo_factor(a ** -3, 1.0),
-                    np.histogramdd: cosmo_factor(a ** -6, 1.0),
+                    np.histogram: cosmo_factor(a**-1, 1.0),
+                    np.histogram2d: cosmo_factor(a**-3, 1.0),
+                    np.histogramdd: cosmo_factor(a**-6, 1.0),
                 }[func]
             )
         elif density and isinstance(weights, cosmo_array):
@@ -883,9 +893,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a ** 0, 1.0),
-                    np.histogram2d: cosmo_factor(a ** -2, 1.0),
-                    np.histogramdd: cosmo_factor(a ** -5, 1.0),
+                    np.histogram: cosmo_factor(a**0, 1.0),
+                    np.histogram2d: cosmo_factor(a**-2, 1.0),
+                    np.histogramdd: cosmo_factor(a**-5, 1.0),
                 }[func]
             )
         elif not density and isinstance(weights, cosmo_array):
@@ -893,9 +903,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a ** 1, 1.0),
-                    np.histogram2d: cosmo_factor(a ** 1, 1.0),
-                    np.histogramdd: cosmo_factor(a ** 1, 1.0),
+                    np.histogram: cosmo_factor(a**1, 1.0),
+                    np.histogram2d: cosmo_factor(a**1, 1.0),
+                    np.histogramdd: cosmo_factor(a**1, 1.0),
                 }[func]
             )
         ret_bins = {
@@ -907,9 +917,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             ret_bins,
             (
                 [
-                    cosmo_factor(a ** 1, 1.0),
-                    cosmo_factor(a ** 2, 1.0),
-                    cosmo_factor(a ** 3, 1.0),
+                    cosmo_factor(a**1, 1.0),
+                    cosmo_factor(a**2, 1.0),
+                    cosmo_factor(a**3, 1.0),
                 ]
             ),
         ):
@@ -920,29 +930,29 @@ def test_getitem(self):
         """
         Make sure that we don't degrade to an ndarray on slicing.
         """
-        assert isinstance(cosmo_obect(np.arange(3))[0], cosmo_quantity)
+        assert isinstance(cosmo_object(np.arange(3))[0], cosmo_quantity)
 
     def test_reshape_to_scalar(self):
         """
         Make sure that we convert to a cosmo_quantity when we reshape to a scalar.
         """
-        assert isinstance(cosmo_obect(np.ones(1)).reshape(tuple()), cosmo_quantity)
+        assert isinstance(cosmo_object(np.ones(1)).reshape(tuple()), cosmo_quantity)
 
     def test_iter(self):
         """
         Make sure that we get cosmo_quantity's when iterating over a cosmo_array.
         """
-        for cq in cosmo_obect(np.arange(3)):
+        for cq in cosmo_object(np.arange(3)):
             assert isinstance(cq, cosmo_quantity)
 
     def test_dot(self):
         """
         Make sure that we get a cosmo_array when we use array attribute dot.
         """
-        res = cosmo_obect(np.arange(3)).dot(cosmo_obect(np.arange(3)))
+        res = cosmo_object(np.arange(3)).dot(cosmo_object(np.arange(3)))
         assert isinstance(res, cosmo_quantity)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5)
+        assert res.cosmo_factor == cosmo_factor(a**2, 0.5)
         assert res.valid_transform is True
 
 
@@ -982,7 +992,7 @@ def test_propagation_func(self, func, args):
         )
         res = getattr(cq, func)(*args)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
         assert res.valid_transform is True
 
     def test_round(self):
@@ -1000,7 +1010,7 @@ def test_round(self):
         res = round(cq)
         assert res.value == 1.0
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
         assert res.valid_transform is True
 
     def test_scalar_return_func(self):
@@ -1008,7 +1018,7 @@ def test_scalar_return_func(self):
         Make sure that default-wrapped functions that take a cosmo_array and return a
         scalar convert to a cosmo_quantity.
         """
-        cosmo_obect = cosmo_array(
+        cosmo_object = cosmo_array(
             np.arange(3),
             u.m,
             comoving=False,
@@ -1016,7 +1026,7 @@ def test_scalar_return_func(self):
             scale_exponent=1,
             valid_transform=True,
         )
-        res = np.min(cosmo_obect)
+        res = np.min(cosmo_object)
         assert isinstance(res, cosmo_quantity)
 
     @pytest.mark.parametrize("prop", ["T", "ua", "unit_array"])
@@ -1034,7 +1044,7 @@ def test_propagation_props(self, prop):
         )
         res = getattr(cq, prop)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
         assert res.valid_transform is True
 
     def test_multiply_quantities(self):
@@ -1052,8 +1062,8 @@ def test_multiply_quantities(self):
         multiplied = cq * cq
         assert type(multiplied) is cosmo_quantity
         assert multiplied.comoving is False
-        assert multiplied.cosmo_factor == cosmo_factor(a ** 2, 0.5)
-        assert multiplied.to_value(u.m ** 2) == 4
+        assert multiplied.cosmo_factor == cosmo_factor(a**2, 0.5)
+        assert multiplied.to_value(u.m**2) == 4
 
 
 class TestCosmoArrayCopy:

From d2acaa9d3bdf8f0d2b80034a801eeb1d951561a6 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 10 Oct 2025 17:46:18 +0100
Subject: [PATCH 05/23] Run formatter.

---
 tests/test_cosmo_array.py   | 50 ++++++++++++++++++-------------------
 tests/test_visualisation.py | 32 ++++++++++++++++++++----
 2 files changed, 52 insertions(+), 30 deletions(-)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 50228be0..9360eb46 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -121,7 +121,7 @@ def test_init_from_ndarray(self):
         arr = cosmo_array(
             np.ones(5),
             units=u.Mpc,
-            cosmo_factor=cosmo_factor(a**1, 1.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -146,7 +146,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.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -169,7 +169,7 @@ def test_init_from_unyt_array(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             u.unyt_array(np.ones(5), units=u.Mpc),
-            cosmo_factor=cosmo_factor(a**1, 1.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -195,7 +195,7 @@ def test_init_from_list_of_unyt_arrays(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             [u.unyt_array(1, units=u.Mpc) for _ in range(5)],
-            cosmo_factor=cosmo_factor(a**1, 1.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -219,7 +219,7 @@ def test_init_from_list_of_cosmo_arrays(self):
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a**1, 1
+            a ** 1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
@@ -229,14 +229,14 @@ def test_init_from_list_of_cosmo_arrays(self):
                     [1],
                     units=u.Mpc,
                     comoving=False,
-                    cosmo_factor=cosmo_factor(a**1, 1.0),
+                    cosmo_factor=cosmo_factor(a ** 1, 1.0),
                 )
                 for _ in range(5)
             ]
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a**1, 1
+            a ** 1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
 
@@ -893,9 +893,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a**-1, 1.0),
-                    np.histogram2d: cosmo_factor(a**-3, 1.0),
-                    np.histogramdd: cosmo_factor(a**-6, 1.0),
+                    np.histogram: cosmo_factor(a ** -1, 1.0),
+                    np.histogram2d: cosmo_factor(a ** -3, 1.0),
+                    np.histogramdd: cosmo_factor(a ** -6, 1.0),
                 }[func]
             )
         elif density and isinstance(weights, cosmo_array):
@@ -903,9 +903,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a**0, 1.0),
-                    np.histogram2d: cosmo_factor(a**-2, 1.0),
-                    np.histogramdd: cosmo_factor(a**-5, 1.0),
+                    np.histogram: cosmo_factor(a ** 0, 1.0),
+                    np.histogram2d: cosmo_factor(a ** -2, 1.0),
+                    np.histogramdd: cosmo_factor(a ** -5, 1.0),
                 }[func]
             )
         elif not density and isinstance(weights, cosmo_array):
@@ -913,9 +913,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a**1, 1.0),
-                    np.histogram2d: cosmo_factor(a**1, 1.0),
-                    np.histogramdd: cosmo_factor(a**1, 1.0),
+                    np.histogram: cosmo_factor(a ** 1, 1.0),
+                    np.histogram2d: cosmo_factor(a ** 1, 1.0),
+                    np.histogramdd: cosmo_factor(a ** 1, 1.0),
                 }[func]
             )
         ret_bins = {
@@ -927,9 +927,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             ret_bins,
             (
                 [
-                    cosmo_factor(a**1, 1.0),
-                    cosmo_factor(a**2, 1.0),
-                    cosmo_factor(a**3, 1.0),
+                    cosmo_factor(a ** 1, 1.0),
+                    cosmo_factor(a ** 2, 1.0),
+                    cosmo_factor(a ** 3, 1.0),
                 ]
             ),
         ):
@@ -962,7 +962,7 @@ def test_dot(self):
         res = cosmo_object(np.arange(3)).dot(cosmo_object(np.arange(3)))
         assert isinstance(res, cosmo_quantity)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**2, 0.5)
+        assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5)
         assert res.valid_transform is True
 
 
@@ -1002,7 +1002,7 @@ def test_propagation_func(self, func, args):
         )
         res = getattr(cq, func)(*args)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
     def test_round(self):
@@ -1020,7 +1020,7 @@ def test_round(self):
         res = round(cq)
         assert res.value == 1.0
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
     def test_scalar_return_func(self):
@@ -1054,7 +1054,7 @@ def test_propagation_props(self, prop):
         )
         res = getattr(cq, prop)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
     def test_multiply_quantities(self):
@@ -1072,8 +1072,8 @@ def test_multiply_quantities(self):
         multiplied = cq * cq
         assert type(multiplied) is cosmo_quantity
         assert multiplied.comoving is False
-        assert multiplied.cosmo_factor == cosmo_factor(a**2, 0.5)
-        assert multiplied.to_value(u.m**2) == 4
+        assert multiplied.cosmo_factor == cosmo_factor(a ** 2, 0.5)
+        assert multiplied.to_value(u.m ** 2) == 4
 
 
 class TestCosmoArrayCopy:
diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py
index bf84c40c..8052713d 100644
--- a/tests/test_visualisation.py
+++ b/tests/test_visualisation.py
@@ -1015,11 +1015,21 @@ def test_comoving_versus_physical(cosmological_volume_only_single):
         data = load(cosmological_volume_only_single, mask=m)
         # we force the default (project="masses") to check the cosmo_factor
         # conversion in this case
-        img = func(data, resolution=64, project="masses", region=region.flatten(), parallel=True)
+        img = func(
+            data,
+            resolution=64,
+            project="masses",
+            region=region.flatten(),
+            parallel=True,
+        )
         assert data.gas.masses.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0
         img = func(
-            data, resolution=64, project="densities", region=region.flatten(), parallel=True
+            data,
+            resolution=64,
+            project="densities",
+            region=region.flatten(),
+            parallel=True,
         )
         assert data.gas.densities.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0
@@ -1046,7 +1056,11 @@ def test_comoving_versus_physical(cosmological_volume_only_single):
         data.gas.coordinates.convert_to_physical()
         with pytest.warns(UserWarning, match="Converting coordinate grid to comoving."):
             img = func(
-                data, resolution=64, project="masses", region=region.flatten(), parallel=True
+                data,
+                resolution=64,
+                project="masses",
+                region=region.flatten(),
+                parallel=True,
             )
         assert data.gas.masses.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** (aexp)).simplify() == 0
@@ -1060,14 +1074,22 @@ def test_comoving_versus_physical(cosmological_volume_only_single):
                 UserWarning, match="Converting coordinate grid to comoving."
             ):
                 img = func(
-                    data, resolution=64, project="masses", region=region.flatten(), parallel=True
+                    data,
+                    resolution=64,
+                    project="masses",
+                    region=region.flatten(),
+                    parallel=True,
                 )
         assert data.gas.masses.comoving and img.comoving
         assert (img.cosmo_factor.expr - a ** aexp).simplify() == 0
         # densities are physical, make sure this works with physical coordinates and
         # smoothing lengths
         img = func(
-            data, resolution=64, project="densities", region=region.flatten(), parallel=True
+            data,
+            resolution=64,
+            project="densities",
+            region=region.flatten(),
+            parallel=True,
         )
         assert data.gas.densities.comoving is False and img.comoving is False
         assert (img.cosmo_factor.expr - a ** (aexp - 3.0)).simplify() == 0

From 6e44ef3ffcecf4fb0ecc4260606580d814c02c6e Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 10 Oct 2025 21:38:15 +0100
Subject: [PATCH 06/23] Depend on unyt main branch for testing.

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 38b3416a..0ee3c7f1 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,4 @@ astropy
 h5py
 numba
 numpy
-unyt
+git+https://github.com/yt-project/unyt.git

From ce72da30d8118e4b04c91556fab363a4d87f5fa6 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 10 Oct 2025 22:24:20 +0100
Subject: [PATCH 07/23] Fix botched merge and clean up tests based on notes
 from past self.

---
 tests/test_cosmo_array.py | 143 +++++++++++++-------------------------
 1 file changed, 49 insertions(+), 94 deletions(-)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 9360eb46..5c49c83d 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -121,7 +121,7 @@ def test_init_from_ndarray(self):
         arr = cosmo_array(
             np.ones(5),
             units=u.Mpc,
-            cosmo_factor=cosmo_factor(a ** 1, 1.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -146,7 +146,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.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -169,7 +169,7 @@ def test_init_from_unyt_array(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             u.unyt_array(np.ones(5), units=u.Mpc),
-            cosmo_factor=cosmo_factor(a ** 1, 1.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -195,7 +195,7 @@ def test_init_from_list_of_unyt_arrays(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             [u.unyt_array(1, units=u.Mpc) for _ in range(5)],
-            cosmo_factor=cosmo_factor(a ** 1, 1.0),
+            cosmo_factor=cosmo_factor(a**1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -219,7 +219,7 @@ def test_init_from_list_of_cosmo_arrays(self):
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a ** 1, 1
+            a**1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
@@ -229,14 +229,14 @@ def test_init_from_list_of_cosmo_arrays(self):
                     [1],
                     units=u.Mpc,
                     comoving=False,
-                    cosmo_factor=cosmo_factor(a ** 1, 1.0),
+                    cosmo_factor=cosmo_factor(a**1, 1.0),
                 )
                 for _ in range(5)
             ]
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a ** 1, 1
+            a**1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
 
@@ -628,13 +628,6 @@ def test_explicitly_handled_funcs(self):
         functions_checked = list()
         bad_funcs = dict()
         for fname, args in functions_to_check.items():
-            # ----- this is to be removed ------
-            # ---- see test_block_is_broken ----
-            if fname == "block":
-                # we skip this function due to issue in unyt with unreleased fix
-                functions_checked.append(np.block)
-                continue
-            # ----------------------------------
             ua_args = list()
             for arg in args:
                 ua_args.append(arg_to_ua(arg))
@@ -721,22 +714,6 @@ def test_explicitly_handled_funcs(self):
                 ],
             )
 
-    @pytest.mark.xfail
-    def test_block_is_broken(self):
-        """
-        There is an issue in unyt affecting np.block and fixed in
-        https://github.com/yt-project/unyt/pull/571
-
-        When this fix is released:
-        - This test will unexpectedly pass (instead of xfailing).
-        - Remove lines flagged with a comment in `test_explicitly_handled_funcs`.
-        - Remove this test.
-        """
-        assert isinstance(
-            np.block([[cosmo_object(np.arange(3))], [cosmo_object(np.arange(3))]]),
-            cosmo_array,
-        )
-
     # the combinations of units and cosmo_factors is nonsense but it's just for testing...
     @pytest.mark.parametrize(
         "func_args",
@@ -893,9 +870,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a ** -1, 1.0),
-                    np.histogram2d: cosmo_factor(a ** -3, 1.0),
-                    np.histogramdd: cosmo_factor(a ** -6, 1.0),
+                    np.histogram: cosmo_factor(a**-1, 1.0),
+                    np.histogram2d: cosmo_factor(a**-3, 1.0),
+                    np.histogramdd: cosmo_factor(a**-6, 1.0),
                 }[func]
             )
         elif density and isinstance(weights, cosmo_array):
@@ -903,9 +880,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a ** 0, 1.0),
-                    np.histogram2d: cosmo_factor(a ** -2, 1.0),
-                    np.histogramdd: cosmo_factor(a ** -5, 1.0),
+                    np.histogram: cosmo_factor(a**0, 1.0),
+                    np.histogram2d: cosmo_factor(a**-2, 1.0),
+                    np.histogramdd: cosmo_factor(a**-5, 1.0),
                 }[func]
             )
         elif not density and isinstance(weights, cosmo_array):
@@ -913,9 +890,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a ** 1, 1.0),
-                    np.histogram2d: cosmo_factor(a ** 1, 1.0),
-                    np.histogramdd: cosmo_factor(a ** 1, 1.0),
+                    np.histogram: cosmo_factor(a**1, 1.0),
+                    np.histogram2d: cosmo_factor(a**1, 1.0),
+                    np.histogramdd: cosmo_factor(a**1, 1.0),
                 }[func]
             )
         ret_bins = {
@@ -927,9 +904,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             ret_bins,
             (
                 [
-                    cosmo_factor(a ** 1, 1.0),
-                    cosmo_factor(a ** 2, 1.0),
-                    cosmo_factor(a ** 3, 1.0),
+                    cosmo_factor(a**1, 1.0),
+                    cosmo_factor(a**2, 1.0),
+                    cosmo_factor(a**3, 1.0),
                 ]
             ),
         ):
@@ -962,7 +939,7 @@ def test_dot(self):
         res = cosmo_object(np.arange(3)).dot(cosmo_object(np.arange(3)))
         assert isinstance(res, cosmo_quantity)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5)
+        assert res.cosmo_factor == cosmo_factor(a**2, 0.5)
         assert res.valid_transform is True
 
 
@@ -1002,7 +979,7 @@ def test_propagation_func(self, func, args):
         )
         res = getattr(cq, func)(*args)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
         assert res.valid_transform is True
 
     def test_round(self):
@@ -1020,7 +997,7 @@ def test_round(self):
         res = round(cq)
         assert res.value == 1.0
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
         assert res.valid_transform is True
 
     def test_scalar_return_func(self):
@@ -1054,7 +1031,7 @@ def test_propagation_props(self, prop):
         )
         res = getattr(cq, prop)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
         assert res.valid_transform is True
 
     def test_multiply_quantities(self):
@@ -1072,8 +1049,8 @@ def test_multiply_quantities(self):
         multiplied = cq * cq
         assert type(multiplied) is cosmo_quantity
         assert multiplied.comoving is False
-        assert multiplied.cosmo_factor == cosmo_factor(a ** 2, 0.5)
-        assert multiplied.to_value(u.m ** 2) == 4
+        assert multiplied.cosmo_factor == cosmo_factor(a**2, 0.5)
+        assert multiplied.to_value(u.m**2) == 4
 
 
 class TestCosmoArrayCopy:
@@ -1134,60 +1111,38 @@ def test_to_cgs(self):
 
 
 class TestMultiplicationByUnyt:
-    @pytest.mark.parametrize(
-        "cosmo_object",
-        [
-            cosmo_array(
-                np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
-            ),
-            cosmo_quantity(
-                np.ones(1), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
-            ),
-        ],
-    )
-    def test_multiplication_by_unyt(self, cosmo_object):
+    def test_multiplication_by_unyt(self):
         """
         We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as
         for `cosmo_array(...) * u.Mpc`.
         """
-
-        lmultiplied_by_quantity = cosmo_object * (
-            1 * u.Mpc
-        )  # parentheses very important here
-        lmultiplied_by_unyt = cosmo_object * u.Mpc
-        assert isinstance(lmultiplied_by_quantity, cosmo_array)
-        assert isinstance(lmultiplied_by_unyt, cosmo_array)
-        assert lmultiplied_by_unyt.comoving == lmultiplied_by_quantity.comoving
-        assert np.allclose(
-            lmultiplied_by_unyt.to_value(lmultiplied_by_quantity.units),
-            lmultiplied_by_quantity.to_value(lmultiplied_by_quantity.units),
+        ca = cosmo_array(
+            np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1
         )
-
-        ldivided_by_quantity = cosmo_object / (
-            1 * u.Mpc
-        )  # parentheses very important here
-        ldivided_by_unyt = cosmo_object / u.Mpc
-        assert isinstance(ldivided_by_quantity, cosmo_array)
-        assert isinstance(ldivided_by_unyt, cosmo_array)
-        assert ldivided_by_unyt.comoving == ldivided_by_quantity.comoving
-        assert np.allclose(
-            ldivided_by_unyt.to_value(ldivided_by_quantity.units),
-            ldivided_by_quantity.to_value(ldivided_by_quantity.units),
+        # required so that can test right-sided division with the same assertions:
+        assert np.allclose(ca.to_value(ca.units), 1)
+        # the reference result:
+        multiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
+        # get the same result twice through left-sided multiplication and division:
+        lmultiplied_by_unyt = ca * u.Mpc
+        ldivided_by_unyt = ca / u.Mpc**-1
+        rmultiplied_by_unyt = u.Mpc * ca
+        rdivided_by_unyt = u.Mpc**3 / ca
+        print(
+            lmultiplied_by_unyt,
+            ldivided_by_unyt,
+            rmultiplied_by_unyt,
+            rdivided_by_unyt,
         )
 
-        rmultiplied_by_quantity = (
-            1 * u.Mpc
-        ) * cosmo_object  # parentheses very important here
-        rmultiplied_by_unyt = u.Mpc * cosmo_object
-        rdivided_by_quantity = (
-            1 * u.Mpc
-        ) / cosmo_object  # parentheses very important here
-        rdivided_by_unyt = u.Mpc / cosmo_object
-        assert rmultiplied_by_quantity.comoving
-        for multiplied_by_unyt in (rmultiplied_by_unyt, rdivided_by_unyt):
+        for multiplied_by_unyt in (
+            lmultiplied_by_unyt,
+            ldivided_by_unyt,
+            rmultiplied_by_unyt,
+            rdivided_by_unyt,
+        ):
             assert isinstance(multiplied_by_quantity, cosmo_array)
             assert isinstance(multiplied_by_unyt, cosmo_array)
-            assert multiplied_by_unyt.comoving == multiplied_by_quantity.comoving
             assert np.allclose(
                 multiplied_by_unyt.to_value(multiplied_by_quantity.units),
                 multiplied_by_quantity.to_value(multiplied_by_quantity.units),
@@ -1202,7 +1157,7 @@ class TestPickle:
     def test_pickle(self):
         attrs = {
             "comoving": False,
-            "cosmo_factor": cosmo_factor(a ** 1, 0.5),
+            "cosmo_factor": cosmo_factor(a**1, 0.5),
             "compression": "FMantissa9",
             "valid_transform": True,
         }

From a4761a51e1d2156396f6f1935964d52dd41dc742 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 10 Oct 2025 22:25:48 +0100
Subject: [PATCH 08/23] Run black...

---
 tests/test_cosmo_array.py | 61 +++++++++++++++++++--------------------
 1 file changed, 29 insertions(+), 32 deletions(-)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 5c49c83d..adb109ee 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -121,7 +121,7 @@ def test_init_from_ndarray(self):
         arr = cosmo_array(
             np.ones(5),
             units=u.Mpc,
-            cosmo_factor=cosmo_factor(a**1, 1.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -146,7 +146,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.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -169,7 +169,7 @@ def test_init_from_unyt_array(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             u.unyt_array(np.ones(5), units=u.Mpc),
-            cosmo_factor=cosmo_factor(a**1, 1.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -195,7 +195,7 @@ def test_init_from_list_of_unyt_arrays(self):
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
         arr = cosmo_array(
             [u.unyt_array(1, units=u.Mpc) for _ in range(5)],
-            cosmo_factor=cosmo_factor(a**1, 1.0),
+            cosmo_factor=cosmo_factor(a ** 1, 1.0),
             comoving=False,
         )
         assert hasattr(arr, "cosmo_factor")
@@ -219,7 +219,7 @@ def test_init_from_list_of_cosmo_arrays(self):
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a**1, 1
+            a ** 1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
         # also with a cosmo_factor argument instead of scale_factor & scale_exponent
@@ -229,14 +229,14 @@ def test_init_from_list_of_cosmo_arrays(self):
                     [1],
                     units=u.Mpc,
                     comoving=False,
-                    cosmo_factor=cosmo_factor(a**1, 1.0),
+                    cosmo_factor=cosmo_factor(a ** 1, 1.0),
                 )
                 for _ in range(5)
             ]
         )
         assert isinstance(arr, cosmo_array)
         assert hasattr(arr, "cosmo_factor") and arr.cosmo_factor == cosmo_factor(
-            a**1, 1
+            a ** 1, 1
         )
         assert hasattr(arr, "comoving") and arr.comoving is False
 
@@ -870,9 +870,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a**-1, 1.0),
-                    np.histogram2d: cosmo_factor(a**-3, 1.0),
-                    np.histogramdd: cosmo_factor(a**-6, 1.0),
+                    np.histogram: cosmo_factor(a ** -1, 1.0),
+                    np.histogram2d: cosmo_factor(a ** -3, 1.0),
+                    np.histogramdd: cosmo_factor(a ** -6, 1.0),
                 }[func]
             )
         elif density and isinstance(weights, cosmo_array):
@@ -880,9 +880,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a**0, 1.0),
-                    np.histogram2d: cosmo_factor(a**-2, 1.0),
-                    np.histogramdd: cosmo_factor(a**-5, 1.0),
+                    np.histogram: cosmo_factor(a ** 0, 1.0),
+                    np.histogram2d: cosmo_factor(a ** -2, 1.0),
+                    np.histogramdd: cosmo_factor(a ** -5, 1.0),
                 }[func]
             )
         elif not density and isinstance(weights, cosmo_array):
@@ -890,9 +890,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             assert (
                 result[0].cosmo_factor
                 == {
-                    np.histogram: cosmo_factor(a**1, 1.0),
-                    np.histogram2d: cosmo_factor(a**1, 1.0),
-                    np.histogramdd: cosmo_factor(a**1, 1.0),
+                    np.histogram: cosmo_factor(a ** 1, 1.0),
+                    np.histogram2d: cosmo_factor(a ** 1, 1.0),
+                    np.histogramdd: cosmo_factor(a ** 1, 1.0),
                 }[func]
             )
         ret_bins = {
@@ -904,9 +904,9 @@ def test_histograms(self, func_args, weights, bins_type, density):
             ret_bins,
             (
                 [
-                    cosmo_factor(a**1, 1.0),
-                    cosmo_factor(a**2, 1.0),
-                    cosmo_factor(a**3, 1.0),
+                    cosmo_factor(a ** 1, 1.0),
+                    cosmo_factor(a ** 2, 1.0),
+                    cosmo_factor(a ** 3, 1.0),
                 ]
             ),
         ):
@@ -939,7 +939,7 @@ def test_dot(self):
         res = cosmo_object(np.arange(3)).dot(cosmo_object(np.arange(3)))
         assert isinstance(res, cosmo_quantity)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**2, 0.5)
+        assert res.cosmo_factor == cosmo_factor(a ** 2, 0.5)
         assert res.valid_transform is True
 
 
@@ -979,7 +979,7 @@ def test_propagation_func(self, func, args):
         )
         res = getattr(cq, func)(*args)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
     def test_round(self):
@@ -997,7 +997,7 @@ def test_round(self):
         res = round(cq)
         assert res.value == 1.0
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
     def test_scalar_return_func(self):
@@ -1031,7 +1031,7 @@ def test_propagation_props(self, prop):
         )
         res = getattr(cq, prop)
         assert res.comoving is False
-        assert res.cosmo_factor == cosmo_factor(a**1, 1.0)
+        assert res.cosmo_factor == cosmo_factor(a ** 1, 1.0)
         assert res.valid_transform is True
 
     def test_multiply_quantities(self):
@@ -1049,8 +1049,8 @@ def test_multiply_quantities(self):
         multiplied = cq * cq
         assert type(multiplied) is cosmo_quantity
         assert multiplied.comoving is False
-        assert multiplied.cosmo_factor == cosmo_factor(a**2, 0.5)
-        assert multiplied.to_value(u.m**2) == 4
+        assert multiplied.cosmo_factor == cosmo_factor(a ** 2, 0.5)
+        assert multiplied.to_value(u.m ** 2) == 4
 
 
 class TestCosmoArrayCopy:
@@ -1125,14 +1125,11 @@ def test_multiplication_by_unyt(self):
         multiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
         # get the same result twice through left-sided multiplication and division:
         lmultiplied_by_unyt = ca * u.Mpc
-        ldivided_by_unyt = ca / u.Mpc**-1
+        ldivided_by_unyt = ca / u.Mpc ** -1
         rmultiplied_by_unyt = u.Mpc * ca
-        rdivided_by_unyt = u.Mpc**3 / ca
+        rdivided_by_unyt = u.Mpc ** 3 / ca
         print(
-            lmultiplied_by_unyt,
-            ldivided_by_unyt,
-            rmultiplied_by_unyt,
-            rdivided_by_unyt,
+            lmultiplied_by_unyt, ldivided_by_unyt, rmultiplied_by_unyt, rdivided_by_unyt
         )
 
         for multiplied_by_unyt in (
@@ -1157,7 +1154,7 @@ class TestPickle:
     def test_pickle(self):
         attrs = {
             "comoving": False,
-            "cosmo_factor": cosmo_factor(a**1, 0.5),
+            "cosmo_factor": cosmo_factor(a ** 1, 0.5),
             "compression": "FMantissa9",
             "valid_transform": True,
         }

From 54851cdd73b834c1ee1e6b430219ba78eee25d21 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Sat, 25 Oct 2025 13:01:32 +0100
Subject: [PATCH 09/23] Run formatter.

---
 tests/test_cosmo_array.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 8092c2f1..1c414579 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -1049,8 +1049,8 @@ def test_multiply_quantities(self):
         multiplied = cq * cq
         assert type(multiplied) is cosmo_quantity
         assert multiplied.comoving is False
-        assert multiplied.cosmo_factor == cosmo_factor(a ** 2, 0.5)
-        assert multiplied.to_value(u.m ** 2) == 4
+        assert multiplied.cosmo_factor == cosmo_factor(a**2, 0.5)
+        assert multiplied.to_value(u.m**2) == 4
 
 
 class TestCosmoArrayCopy:
@@ -1125,10 +1125,10 @@ def test_multiplication_by_unyt(self):
         multiplied_by_quantity = ca * (1 * u.Mpc)  # parentheses very important here
         # get the same result twice through left-sided multiplication and division:
         lmultiplied_by_unyt = ca * u.Mpc
-        ldivided_by_unyt = ca / u.Mpc ** -1
+        ldivided_by_unyt = ca / u.Mpc**-1
         # and twice more through right-sided multiplication and division:
         rmultiplied_by_unyt = u.Mpc * ca
-        rdivided_by_unyt = u.Mpc ** 3 / ca
+        rdivided_by_unyt = u.Mpc**3 / ca
 
         for multiplied_by_unyt in (
             lmultiplied_by_unyt,

From 40cf579127ae143b583b5a5f40778449599173d1 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Mon, 27 Oct 2025 16:30:47 +0000
Subject: [PATCH 10/23] Bugfix.

---
 swiftsimio/objects.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py
index 467b93de..f6dcfd7b 100644
--- a/swiftsimio/objects.py
+++ b/swiftsimio/objects.py
@@ -1661,7 +1661,7 @@ def __unyt_ufunc_finalize__(
             ret_cf = cls._cosmo_factor_ufunc_registry[ufunc](*cfs, inputs=inputs)
         # if we get a tuple we have multiple return values to deal with
         if isinstance(result, tuple):
-            r = tuple(
+            result = tuple(
                 (
                     r.view(cosmo_quantity)
                     if r.shape == ()

From 8f5dab2abe8bca9232b9bf8f69810940efb64e3a Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Mon, 27 Oct 2025 16:55:57 +0000
Subject: [PATCH 11/23] Fill out docstrings and type hints.

---
 swiftsimio/objects.py     | 79 +++++++++++++++++++++++++++++++++++++--
 tests/test_cosmo_array.py |  4 +-
 2 files changed, 77 insertions(+), 6 deletions(-)

diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py
index f6dcfd7b..22f9d740 100644
--- a/swiftsimio/objects.py
+++ b/swiftsimio/objects.py
@@ -1630,14 +1630,87 @@ def from_pint(
         return obj
 
     @classmethod
-    def __unyt_ufunc_prepare__(cls, ufunc: np.ufunc, method: str, *inputs, **kwargs):
+    def __unyt_ufunc_prepare__(
+        cls, ufunc: np.ufunc, method: str, *inputs: tuple, **kwargs: dict
+    ) -> tuple[np.ufunc, str, tuple, dict]:
+        """
+        Prepare arguments for a ufunc call.
+
+        This function gives us the opportunity to pre-process arguments to a ufunc call
+        before handing control off to :mod:`unyt`. The arguments and kwargs are checked for
+        consistent ``cosmo_factor`` attributes and coerced to a common comoving/physical state.
+
+        Parameters
+        ----------
+        ufunc : np.ufunc
+            The ufunc that is about to be called.
+
+        method : str
+            The call method for the ufunc (for example `"call"` or `"reduce"`).
+
+        *inputs : tuple
+            The ufunc arguments.
+
+        **kwargs : dict
+            The ufunc kwargs.
+
+        Returns
+        -------
+        np.ufunc
+            The ufunc that is about to be called.
+
+        str
+            The call method for the ufunc.
+
+        tuple
+            The now prepared arguments for the ufunc.
+
+        dict
+            The now prepared kwargs for the ufunc.
+        """
         helper_result = _prepare_array_func_args(*inputs, **kwargs)
         return ufunc, method, helper_result["args"], helper_result["kwargs"]
 
     @classmethod
     def __unyt_ufunc_finalize__(
-        cls, result, ufunc: np.ufunc, method: str, *inputs, **kwargs
-    ):
+        cls,
+        result: tuple | unyt_array,
+        ufunc: np.ufunc,
+        method: str,
+        *inputs: tuple,
+        **kwargs: dict,
+    ) -> "tuple | cosmo_array":
+        """
+        Finalize results after a ufunc call.
+
+        This function gives us the opportunity to post-process return value(s) from a ufunc
+        when we get control back from :mod:`unyt`. We check that the return type is
+        consistent with its shape (i.e. a :class:`~swiftsimio.objects.cosmo_array` or
+        :class:`~swiftsimio.objects.cosmo_quantity`) and attach our cosmo attributes.
+
+        Parameters
+        ----------
+        result : :class:`~unyt.array.unyt_array` or tuple
+            The return value of the called ufunc.
+
+        ufunc : np.ufunc
+            The ufunc that was called.
+
+        method : str
+            The call method for the ufunc (for example `"call"` or `"reduce"`).
+
+        *inputs : tuple
+            The ufunc arguments.
+
+        **kwargs : dict
+            The ufunc kwargs.
+
+        Returns
+        -------
+        tuple or comso_array
+            The result of the ufunc call, with the appropriate type and cosmo attributes attached.
+        """
+        # wonder if we could cache helper_result during __unyt_ufunc_prepare__ to use here?
         helper_result = _prepare_array_func_args(*inputs, **kwargs)
         cfs = helper_result["cfs"]
         # make sure we evaluate the cosmo_factor_ufunc_registry function:
diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 86295325..5437d33a 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -1092,9 +1092,7 @@ def test_propagation_props(self, prop):
         assert res.valid_transform is True
 
     def test_multiply_quantities(self):
-        """
-        Test multiplying two quantities.
-        """
+        """Test multiplying two quantities."""
         cq = cosmo_quantity(
             2,
             u.m,

From c3723dcf206821e22ece0a169587a07f172aab8b Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 17:18:37 +0000
Subject: [PATCH 12/23] Temporarily point to unyt main branch as dependency
 while waiting for release.

---
 requirements.txt | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/requirements.txt b/requirements.txt
index 38b3416a..677cb0a2 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -2,4 +2,4 @@ astropy
 h5py
 numba
 numpy
-unyt
+git+https://github.com/yt-project/unyt.git@main

From cd7ecf63a82286dcd93338fa61a66a61f2e1fe3a Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 17:25:14 +0000
Subject: [PATCH 13/23] Add regression test.

---
 tests/test_cosmo_array.py | 11 +++++++++++
 1 file changed, 11 insertions(+)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 113735d1..3b35f5b7 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -988,6 +988,17 @@ def test_dot(self):
         assert res.cosmo_factor == cosmo_factor(a**2, 0.5)
         assert res.valid_transform is True
 
+    def test_average_with_returned(self):
+        """Make sure that sum of weights in numpy's average gets cosmo attributes."""
+        # regression test for https://github.com/SWIFTSIM/swiftsimio/issues/285
+        x = cosmo_array(
+            np.arange(9).reshape((3, 3)), u.kpc, scale_factor=1.0, scale_exponent=1.0
+        )
+        w = cosmo_array(np.arange(3), u.solMass, scale_factor=1.0, scale_exponent=0.0)
+        avg, wsum = np.average(x, weights=w, axis=-1, returned=True)
+        assert avg.cosmo_factor == x.cosmo_factor
+        assert wsum.cosmo_factor == w.cosmo_factor
+
 
 class TestCosmoQuantity:
     """

From 32a2e8e165d9f6690a72cfd592ca52c547dbe3a4 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 17:48:54 +0000
Subject: [PATCH 14/23] Patch works in numpy2.3.4

---
 swiftsimio/_array_functions.py | 40 +++++++++++++++++++++++++++++++---
 1 file changed, 37 insertions(+), 3 deletions(-)

diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py
index 9e0f2f22..f2fd6e9b 100644
--- a/swiftsimio/_array_functions.py
+++ b/swiftsimio/_array_functions.py
@@ -118,6 +118,10 @@
     isin as unyt_in1d,
     take as unyt_take,
 )
+from importlib.metadata import version
+from packaging.version import Version
+
+NUMPY_VERSION = Version(version("numpy"))
 
 _HANDLED_FUNCTIONS = {}
 
@@ -2267,9 +2271,39 @@ def trapezoid(y, x=None, dx=1.0, axis=-1):  # noqa numpydoc ignore=GL08
 
 # Now we wrap functions that unyt does not handle explicitly:
 
-implements(np.average)(
-    _propagate_cosmo_array_attributes_to_result(np.average._implementation)
-)
+
+if NUMPY_VERSION < Version("2.4.1"):
+
+    from unyt._array_functions import average as unyt_average
+
+    @implements(np.average)
+    def average(a, axis=None, weights=None, returned=False, *, keepdims=np._NoValue):
+        # Average suffered from a bug
+        # (https://github.com/SWIFTSIM/swiftsimio/issues/285)
+        # Correct results depend on unyt>=3.1.0
+        # (https://github.com/yt-project/unyt/pull/611)
+        # There is also a fix in numpy>=3.4.1
+        # (https://github.com/numpy/numpy/pull/30522)
+        # that means we no longer need any special handling here, but to support older
+        # versions we need a patch.
+        helper_result = _prepare_array_func_args(
+            a, axis=axis, weights=weights, returned=returned, keepdims=keepdims
+        )
+        res = unyt_average(
+            a, axis=axis, weights=weights, returned=returned, keepdims=keepdims
+        )
+        ret_cf_avg = _preserve_cosmo_factor(helper_result["cfs"][0])
+        if returned:
+            avg, wsum = res
+            ret_cf_wsum = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"])
+            return (
+                _return_helper(avg, helper_result, ret_cf_avg),
+                _return_helper(wsum, helper_result, ret_cf_wsum),
+            )
+        else:
+            return _return_helper(avg, helper_result, ret_cf_avg)
+
+
 implements(np.max)(_propagate_cosmo_array_attributes_to_result(np.max._implementation))
 implements(np.min)(_propagate_cosmo_array_attributes_to_result(np.min._implementation))
 implements(np.mean)(

From fe79d38d8ef4ab7e8d0497c4e732cb3316def11c Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 17:55:52 +0000
Subject: [PATCH 15/23] Patch works in numpy 2.4.1

---
 swiftsimio/_array_functions.py | 61 +++++++++++++++++-----------------
 1 file changed, 31 insertions(+), 30 deletions(-)

diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py
index f2fd6e9b..5d8bc78d 100644
--- a/swiftsimio/_array_functions.py
+++ b/swiftsimio/_array_functions.py
@@ -93,7 +93,6 @@
     linalg_eigvalsh as unyt_linalg_eigvalsh,
     savetxt as unyt_savetxt,
     fill_diagonal as unyt_fill_diagonal,
-    isin as unyt_isin,
     place as unyt_place,
     put as unyt_put,
     put_along_axis as unyt_put_along_axis,
@@ -115,7 +114,7 @@
     array_repr as unyt_array_repr,
     linalg_outer as unyt_linalg_outer,
     trapezoid as unyt_trapezoid,
-    isin as unyt_in1d,
+    isin as unyt_isin,
     take as unyt_take,
 )
 from importlib.metadata import version
@@ -2266,42 +2265,44 @@ def trapezoid(y, x=None, dx=1.0, axis=-1):  # noqa numpydoc ignore=GL08
     return _return_helper(res, helper_result, ret_cf)
 
 
-implements(np.in1d)(_default_comparison_wrapper(unyt_in1d))
+implements(np.isin)(_default_comparison_wrapper(unyt_isin))
 implements(np.take)(_default_unary_wrapper(unyt_take, _preserve_cosmo_factor))
 
 # Now we wrap functions that unyt does not handle explicitly:
 
 
-if NUMPY_VERSION < Version("2.4.1"):
+@implements(np.average)
+def average(a, axis=None, weights=None, returned=False, *, keepdims=np._NoValue):  # noqa numpydoc ignore=GL08
+    # Average suffered from a bug
+    # (https://github.com/SWIFTSIM/swiftsimio/issues/285)
+    # Correct results depend on unyt>=3.1.0
+    # (https://github.com/yt-project/unyt/pull/611)
+    # There is also a fix in numpy>=3.4.1
+    # (https://github.com/numpy/numpy/pull/30522)
+    # that means we no longer need any special handling here, but to support older
+    # versions we need a patch.
+    helper_result = _prepare_array_func_args(
+        a, axis=axis, weights=weights, returned=returned, keepdims=keepdims
+    )
+    if NUMPY_VERSION < Version("2.4.1"):
+        from unyt._array_functions import average as super_average
 
-    from unyt._array_functions import average as unyt_average
+    else:
+        super_average = np.average._implementation
 
-    @implements(np.average)
-    def average(a, axis=None, weights=None, returned=False, *, keepdims=np._NoValue):
-        # Average suffered from a bug
-        # (https://github.com/SWIFTSIM/swiftsimio/issues/285)
-        # Correct results depend on unyt>=3.1.0
-        # (https://github.com/yt-project/unyt/pull/611)
-        # There is also a fix in numpy>=3.4.1
-        # (https://github.com/numpy/numpy/pull/30522)
-        # that means we no longer need any special handling here, but to support older
-        # versions we need a patch.
-        helper_result = _prepare_array_func_args(
-            a, axis=axis, weights=weights, returned=returned, keepdims=keepdims
-        )
-        res = unyt_average(
-            a, axis=axis, weights=weights, returned=returned, keepdims=keepdims
+    res = super_average(
+        a, axis=axis, weights=weights, returned=returned, keepdims=keepdims
+    )
+    ret_cf_avg = _preserve_cosmo_factor(helper_result["cfs"][0])
+    if returned:
+        avg, wsum = res
+        ret_cf_wsum = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"])
+        return (
+            _return_helper(avg, helper_result, ret_cf_avg),
+            _return_helper(wsum, helper_result, ret_cf_wsum),
         )
-        ret_cf_avg = _preserve_cosmo_factor(helper_result["cfs"][0])
-        if returned:
-            avg, wsum = res
-            ret_cf_wsum = _preserve_cosmo_factor(helper_result["kw_cfs"]["weights"])
-            return (
-                _return_helper(avg, helper_result, ret_cf_avg),
-                _return_helper(wsum, helper_result, ret_cf_wsum),
-            )
-        else:
-            return _return_helper(avg, helper_result, ret_cf_avg)
+    else:
+        return _return_helper(avg, helper_result, ret_cf_avg)
 
 
 implements(np.max)(_propagate_cosmo_array_attributes_to_result(np.max._implementation))

From 949bb388de5d52474084a9647351fc1251cd3243 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 18:15:30 +0000
Subject: [PATCH 16/23] Compatibility fixes for unyt 3.1.0 and numpy 2.4.1

---
 swiftsimio/_array_functions.py | 37 +++-------------------------------
 tests/test_cosmo_array.py      | 32 +++--------------------------
 2 files changed, 6 insertions(+), 63 deletions(-)

diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py
index 5d8bc78d..82745fe7 100644
--- a/swiftsimio/_array_functions.py
+++ b/swiftsimio/_array_functions.py
@@ -1274,39 +1274,8 @@ def wrapper(*args: tuple[Any], **kwargs: dict[str, Any]) -> Callable:
 
 
 @implements(np.array2string)
-def array2string(  # noqa numpydoc ignore=GL08
-    a,  # noqa: ANN001
-    max_line_width=None,  # noqa: ANN001
-    precision=None,  # noqa: ANN001
-    suppress_small=None,  # noqa: ANN001
-    separator=" ",  # noqa: ANN001
-    prefix="",  # noqa: ANN001
-    style=np._NoValue,  # noqa: ANN001
-    formatter=None,  # noqa: ANN001
-    threshold=None,  # noqa: ANN001
-    edgeitems=None,  # noqa: ANN001
-    sign=None,  # noqa: ANN001
-    floatmode=None,  # noqa: ANN001
-    suffix="",  # noqa: ANN001
-    *,
-    legacy=None,  # noqa: ANN001
-):
-    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,
-    )
+def array2string(a, *args, **kwargs):  # noqa numpydoc ignore=GL08
+    res = unyt_array2string(a, *args, **kwargs)
     if a.comoving:
         append = " (comoving)"
     elif a.comoving is False:
@@ -2302,7 +2271,7 @@ def average(a, axis=None, weights=None, returned=False, *, keepdims=np._NoValue)
             _return_helper(wsum, helper_result, ret_cf_wsum),
         )
     else:
-        return _return_helper(avg, helper_result, ret_cf_avg)
+        return _return_helper(res, helper_result, ret_cf_avg)
 
 
 implements(np.max)(_propagate_cosmo_array_attributes_to_result(np.max._implementation))
diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 3b35f5b7..427d4356 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -486,7 +486,6 @@ def test_explicitly_handled_funcs(self):
             "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),
@@ -512,7 +511,7 @@ def test_explicitly_handled_funcs(self):
             "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
+            "isin": (ca(np.arange(3)), ca(np.arange(3))),
             "take": (ca(np.arange(3)), np.arange(3)),
             # FUNCTIONS THAT UNYT DOESN'T HANDLE EXPLICITLY (THEY "JUST WORK"):
             "all": (ca(np.arange(3)),),
@@ -658,13 +657,6 @@ def test_explicitly_handled_funcs(self):
         functions_checked = list()
         bad_funcs = dict()
         for fname, args in functions_to_check.items():
-            # ----- this is to be removed ------
-            # ---- see test_block_is_broken ----
-            if fname == "block":
-                # we skip this function due to issue in unyt with unreleased fix
-                functions_checked.append(np.block)
-                continue
-            # ----------------------------------
             ua_args = list()
             for arg in args:
                 ua_args.append(arg_to_ua(arg))
@@ -751,23 +743,6 @@ def test_explicitly_handled_funcs(self):
                 ],
             )
 
-    @pytest.mark.xfail
-    def test_block_is_broken(self):
-        """
-        Tracking upstream fix.
-
-        There is an issue in unyt affecting np.block and fixed in
-        https://github.com/yt-project/unyt/pull/571.
-
-        When this fix is released:
-        - This test will unexpectedly pass (instead of xfailing).
-        - Remove lines flagged with a comment in `test_explicitly_handled_funcs`.
-        - Remove this test.
-        """
-        assert isinstance(
-            np.block([[ca(np.arange(3))], [ca(np.arange(3))]]), cosmo_array
-        )
-
     # the combinations of units and cosmo_factors is nonsense but it's just for testing...
     @pytest.mark.parametrize(
         "func_args",
@@ -843,7 +818,7 @@ def test_block_is_broken(self):
             np.array([1, 2, 3]),
         ),
     )
-    @pytest.mark.parametrize("bins_type", ("int", "np", "ca"))
+    @pytest.mark.parametrize("bins_type", ("int", "ca"))
     @pytest.mark.parametrize("density", (None, True))
     def test_histograms(self, func_args, weights, bins_type, density):
         """
@@ -856,7 +831,6 @@ 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),
@@ -889,7 +863,7 @@ def test_histograms(self, func_args, weights, bins_type, density):
                     np.histogramdd: np.s_[:],
                 }[func]
             ]
-            if bins_type in ("np", "ca")
+            if bins_type == "ca"
             else bins
         )
         result = func(*args, bins=bins, density=density, weights=weights)

From 54ade28e90e49655d0c6a32da31c188afa6a9ba4 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 18:24:29 +0000
Subject: [PATCH 17/23] More version-dependent cases.

---
 tests/test_cosmo_array.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/tests/test_cosmo_array.py b/tests/test_cosmo_array.py
index 427d4356..a55edbc4 100644
--- a/tests/test_cosmo_array.py
+++ b/tests/test_cosmo_array.py
@@ -9,6 +9,10 @@
 from copy import copy, deepcopy
 import pickle
 from swiftsimio.objects import cosmo_array, cosmo_quantity, cosmo_factor, a
+from importlib.metadata import version
+from packaging.version import Version
+
+NUMPY_VERSION = Version(version("numpy"))
 
 savetxt_file = "saved_array.txt"
 
@@ -654,6 +658,8 @@ def test_explicitly_handled_funcs(self):
             "cumulative_prod": (ca(np.arange(3)),),
             "unstack": (ca(np.arange(3)),),
         }
+        if NUMPY_VERSION < Version("2.4.1"):
+            functions_to_check["in1d"] = (ca(np.arange(3)), ca(np.arange(3)))
         functions_checked = list()
         bad_funcs = dict()
         for fname, args in functions_to_check.items():

From 80ada0dae0e77673ff187498c65c04756b434e3a Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 22:49:03 +0000
Subject: [PATCH 18/23] Make a single source of truth for dependencies.

---
 .github/workflows/lint_and_test.yml   | 14 +-------------
 .github/workflows/python-publish.yml  |  1 -
 .readthedocs.yaml                     | 10 +---------
 CONTRIBUTING.md                       |  2 +-
 docs/requirements.txt                 |  4 ----
 docs/source/getting_started/index.rst | 11 ++---------
 optional_requirements.txt             |  7 -------
 pyproject.toml                        | 21 ++++++++++++++++++++-
 requirements.txt                      |  5 -----
 9 files changed, 25 insertions(+), 50 deletions(-)
 delete mode 100644 docs/requirements.txt
 delete mode 100644 optional_requirements.txt
 delete mode 100644 requirements.txt

diff --git a/.github/workflows/lint_and_test.yml b/.github/workflows/lint_and_test.yml
index 155be6c1..4874ff16 100644
--- a/.github/workflows/lint_and_test.yml
+++ b/.github/workflows/lint_and_test.yml
@@ -44,22 +44,10 @@ jobs:
       uses: actions/setup-python@v5
       with:
         python-version: ${{ matrix.python-version }}
-    - name: Cache pip
-      uses: actions/cache@v4
-      with:
-        # This path is specific to Ubuntu
-        path: ~/.cache/pip
-        # Look to see if there is a cache hit for the corresponding requirements file
-        key: ${{ runner.os }}-pip-${{ hashFiles('requirements.txt') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-${{ runner.os }}-
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install pytest
-        if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
-        if [ -f optional_requirements.txt ]; then pip install -r optional_requirements.txt; fi
-        pip install -e .
+        pip install -e .[testing]
     - name: Test with pytest
       run: |
         pytest
diff --git a/.github/workflows/python-publish.yml b/.github/workflows/python-publish.yml
index bbcd6037..0227c98e 100644
--- a/.github/workflows/python-publish.yml
+++ b/.github/workflows/python-publish.yml
@@ -30,7 +30,6 @@ jobs:
       run: |
         python -m pip install --upgrade pip
         pip install setuptools wheel build
-        pip install -r requirements.txt
         pip install .
     - name: Build
       run: |
diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index ec03d7c7..7fd18d6e 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -11,8 +11,7 @@ build:
     python: "3.12"
   jobs:
     post_create_environment:
-      - python3 -m pip install swiftsimio
-      - python3 -m pip install -r optional_requirements.txt
+      - python3 -m pip install swiftsimio[docs]
 
 # Build documentation in the "docs/" directory with Sphinx
 sphinx:
@@ -22,11 +21,4 @@ sphinx:
 formats:
   - pdf
   - epub
-
-# Optional but recommended, declare the Python requirements required
-# to build your documentation
-# See https://docs.readthedocs.io/en/stable/guides/reproducible-builds.html
-python:
-  install:
-    - requirements: docs/requirements.txt
   
diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 226f171a..81f903b6 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -33,7 +33,7 @@ Documentation
 
 The API documentation is built automatically from the docstrings of classes, functions, etc. in the source files. These follow the NumPy-style format. All public (i.e. not starting in `_`) modules, functions, classes, methods, etc. should have an appropriate docstring. Tests should also have descriptive docstrings, but full descriptions (e.g. of all parameters) are not required.
 
-In addition to this there is "narrative documentation" that should describe the features of the code. The docs are built with `sphinx` and use the "ReadTheDocs" theme. If you have the dependencies installed (check `/docs/requirements.txt`) you can build the documentation locally with `make html` in the `/docs` directory. Opening the `/docs/index.html` file with a browser will then allow you to browse the documentation and check your contributions.
+In addition to this there is "narrative documentation" that should describe the features of the code. The docs are built with `sphinx` and use the "ReadTheDocs" theme. If you have the dependencies installed (check `pyproject.toml` or use `pip install -e .[docs]` in the project root directory) you can build the documentation locally with `make html` in the `/docs` directory. Opening the `/docs/build/index.html` file with a browser will then allow you to browse the documentation and check your contributions.
 
 Docstrings
 ----------
diff --git a/docs/requirements.txt b/docs/requirements.txt
deleted file mode 100644
index e3707535..00000000
--- a/docs/requirements.txt
+++ /dev/null
@@ -1,4 +0,0 @@
-sphinx
-sphinx-rtd-theme
-recommonmark
-sphinx_design
diff --git a/docs/source/getting_started/index.rst b/docs/source/getting_started/index.rst
index 56ee8851..6c1d2d0d 100644
--- a/docs/source/getting_started/index.rst
+++ b/docs/source/getting_started/index.rst
@@ -30,19 +30,12 @@ To set up the code for development, first clone the latest master from GitHub:
    
    git clone https://github.com/SWIFTSIM/swiftsimio.git
 
-and install with ``pip`` using the ``-e`` ("editable") flag,
+and install with ``pip`` using the ``-e`` ("editable") flag, and specifying optional dependencies for development and building the documentation:
 
 .. code-block::
 
    cd swiftsimio
-   pip install -e .
-
-Then install the optioanl dependencies for the code and the documentation:
-
-.. code-block::
-
-   pip install -r optional_requirements.txt
-   pip install -r docs/requirements.txt
+   pip install -e .[dev,docs]
 
 .. include:: ../../../README.rst
    :start-after: COMMUNITY_START_LABEL
diff --git a/optional_requirements.txt b/optional_requirements.txt
deleted file mode 100644
index 1d25944b..00000000
--- a/optional_requirements.txt
+++ /dev/null
@@ -1,7 +0,0 @@
-# Development
-pytest
-ruff
-numpydoc
-matplotlib
-scipy
-wily
diff --git a/pyproject.toml b/pyproject.toml
index 03bcf6ae..15011599 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -41,7 +41,7 @@ dependencies = [
     "astropy>=5.0",
     "numpy>=2.1.0",
     "h5py",
-    "unyt>=3.0.4",
+    "unyt @ git+https://github.com/yt-project/unyt@main",
     "numba>=0.50.0",
 ]
 
@@ -53,6 +53,25 @@ dependencies = [
 [project.scripts]
 swiftsnap = "swiftsimio.swiftsnap:swiftsnap"
 
+[project.optional-dependencies]
+dev = [
+    "pytest",
+    "ruff",
+    "numpydoc",
+    "matplotlib",
+    "scipy",
+    "wily",
+]
+docs = [
+    "sphinx",
+    "sphinx-rtd-theme",
+    "recommonmark",
+    "sphinx_design",
+]
+all = [
+    "swiftsimio[dev,docs]"
+]
+
 [tool.ruff]
 exclude = ["docs/source/conf.py"]
 line-length = 88
diff --git a/requirements.txt b/requirements.txt
deleted file mode 100644
index 677cb0a2..00000000
--- a/requirements.txt
+++ /dev/null
@@ -1,5 +0,0 @@
-astropy
-h5py
-numba
-numpy
-git+https://github.com/yt-project/unyt.git@main

From 5bb3bf2019060848ad42a0871da3c3e55c5da2c8 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Tue, 13 Jan 2026 23:13:02 +0000
Subject: [PATCH 19/23] Fix mismatched extra deps name.

---
 .github/workflows/lint_and_test.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/lint_and_test.yml b/.github/workflows/lint_and_test.yml
index 4874ff16..e5895450 100644
--- a/.github/workflows/lint_and_test.yml
+++ b/.github/workflows/lint_and_test.yml
@@ -47,7 +47,7 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -e .[testing]
+        pip install -e .[dev]
     - name: Test with pytest
       run: |
         pytest

From c8a423be01bbff34c0427811e73be7fa71c6957a Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Wed, 14 Jan 2026 10:14:42 +0000
Subject: [PATCH 20/23] Install from repo, not pypi, on RTD.

---
 .readthedocs.yaml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.readthedocs.yaml b/.readthedocs.yaml
index 7fd18d6e..5caef67b 100644
--- a/.readthedocs.yaml
+++ b/.readthedocs.yaml
@@ -11,7 +11,7 @@ build:
     python: "3.12"
   jobs:
     post_create_environment:
-      - python3 -m pip install swiftsimio[docs]
+      - python3 -m pip install .[docs]
 
 # Build documentation in the "docs/" directory with Sphinx
 sphinx:

From 4be32cd470e5e40815ab60e2135428ec5c2d070e Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Thu, 15 Jan 2026 21:39:30 +0000
Subject: [PATCH 21/23] Pin to unyt 3.1.0, tentatively astropy 7.0

---
 pyproject.toml | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/pyproject.toml b/pyproject.toml
index 15011599..0c9ea99f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,10 +38,10 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "astropy>=5.0",
+    "astropy>=7.0",
     "numpy>=2.1.0",
     "h5py",
-    "unyt @ git+https://github.com/yt-project/unyt@main",
+    "unyt>=3.1.0",
     "numba>=0.50.0",
 ]
 

From 5ef628e745a7714907377b8510a3feee5a6fada9 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Thu, 15 Jan 2026 21:42:09 +0000
Subject: [PATCH 22/23] No astropy7 in py3.10,pin to >6?

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 0c9ea99f..ca2cd1a0 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -38,7 +38,7 @@ classifiers = [
     "Operating System :: OS Independent",
 ]
 dependencies = [
-    "astropy>=7.0",
+    "astropy>=6.0",
     "numpy>=2.1.0",
     "h5py",
     "unyt>=3.1.0",

From 545e6a8d01e8b4f55b378ff0a5b91b705e136762 Mon Sep 17 00:00:00 2001
From: Kyle Oman <kyle.a.oman@durham.ac.uk>
Date: Fri, 23 Jan 2026 15:56:50 +0000
Subject: [PATCH 23/23] Be deliberate about when latest & stable versions of
 docs are linked to.

---
 CONTRIBUTING.md            | 2 +-
 README.rst                 | 2 +-
 docs/source/conf.py        | 2 +-
 docs/source/soap/index.rst | 2 +-
 pyproject.toml             | 2 +-
 5 files changed, 5 insertions(+), 5 deletions(-)

diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md
index 81f903b6..b9ce812d 100644
--- a/CONTRIBUTING.md
+++ b/CONTRIBUTING.md
@@ -38,4 +38,4 @@ In addition to this there is "narrative documentation" that should describe the
 Docstrings
 ----------
 
-Ruff currently has limited support for [numpydoc](https://numpydoc.readthedocs.io/en/latest/index.html)-style docstrings. To run additional checks on docstrings use `numpydoc lint **/*.py` in the same directory as the `pyproject.toml` file. As more style rules become supported by `ruff` this will hopefully be phased out.
+Ruff currently has limited support for [numpydoc](https://numpydoc.readthedocs.io)-style docstrings. To run additional checks on docstrings use `numpydoc lint **/*.py` in the same directory as the `pyproject.toml` file. As more style rules become supported by `ruff` this will hopefully be phased out.
diff --git a/README.rst b/README.rst
index 1979ae53..a528c446 100644
--- a/README.rst
+++ b/README.rst
@@ -15,7 +15,7 @@ SWIFTsimIO
    :target: https://github.com/SWIFTSIM/swiftsimio/actions/workflows/lint_and_test.yml
    :alt: Build status
 .. |Documentation status| image:: https://readthedocs.org/projects/swiftsimio/badge/?version=latest
-   :target: https://swiftsimio.readthedocs.io/en/latest/?badge=latest
+   :target: https://swiftsimio.readthedocs.io
    :alt: Documentation status
 .. |JOSS| image:: https://joss.theoj.org/papers/e85c85f49b99389d98f9b6d81f090331/status.svg
    :target: https://joss.theoj.org/papers/e85c85f49b99389d98f9b6d81f090331
diff --git a/docs/source/conf.py b/docs/source/conf.py
index 74ce3a2c..564074c9 100644
--- a/docs/source/conf.py
+++ b/docs/source/conf.py
@@ -109,7 +109,7 @@ def setup(app):  # numpydoc ignore=GL08
     numba=("https://numba.readthedocs.io/en/stable/", None),
     unyt=("https://unyt.readthedocs.io/en/stable/", None),
     scipy=("https://docs.scipy.org/doc/scipy/", None),
-    swiftgalaxy=("https://swiftsimio.readthedocs.io/en/latest/", None),
+    swiftgalaxy=("https://swiftgalaxy.readthedocs.io/en/stable/", None),
     velociraptor=("https://velociraptor-python.readthedocs.io/en/latest/", None),
     astropy=("https://docs.astropy.org/en/stable/", None),
 )
diff --git a/docs/source/soap/index.rst b/docs/source/soap/index.rst
index c143c9ec..f8d644d0 100644
--- a/docs/source/soap/index.rst
+++ b/docs/source/soap/index.rst
@@ -38,4 +38,4 @@ swiftgalaxy
 
 The :mod:`swiftgalaxy` companion package to :mod:`swiftsimio` offers further integration with halo catalogues in SOAP, Caesar and Velociraptor formats (so far). It greatly simplifies efficient loading of particles belonging to an object from a catalogue, and additional tools that are useful when working with a galaxy or other localized collection of particles. Refer to the `swiftgalaxy documentation`_ for details.
 
-.. _swiftgalaxy documentation: https://swiftgalaxy.readthedocs.io/en/latest/
+.. _swiftgalaxy documentation: https://swiftgalaxy.readthedocs.io
diff --git a/pyproject.toml b/pyproject.toml
index ca2cd1a0..da27fa70 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -48,7 +48,7 @@ dependencies = [
 [project.urls]
 "Homepage" = "https://github.com/SWIFTSIM/swiftsimio"
 "Bug Tracker" = "https://github.com/SWIFTSIM/swiftsimio/issues"
-"Documentation" = "https://swiftsimio.readthedocs.io/en/latest"
+"Documentation" = "https://swiftsimio.readthedocs.io"
 
 [project.scripts]
 swiftsnap = "swiftsimio.swiftsnap:swiftsnap"