diff --git a/.github/workflows/lint_and_test.yml b/.github/workflows/lint_and_test.yml index 155be6c1..e5895450 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 .[dev] - 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..5caef67b 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 .[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..b9ce812d 100644 --- a/CONTRIBUTING.md +++ b/CONTRIBUTING.md @@ -33,9 +33,9 @@ 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 ---------- -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/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/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/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/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/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..da27fa70 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -38,21 +38,40 @@ classifiers = [ "Operating System :: OS Independent", ] dependencies = [ - "astropy>=5.0", + "astropy>=6.0", "numpy>=2.1.0", "h5py", - "unyt>=3.0.4", + "unyt>=3.1.0", "numba>=0.50.0", ] [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" +[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 38b3416a..00000000 --- a/requirements.txt +++ /dev/null @@ -1,5 +0,0 @@ -astropy -h5py -numba -numpy -unyt diff --git a/swiftsimio/_array_functions.py b/swiftsimio/_array_functions.py index 9e0f2f22..82745fe7 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,9 +114,13 @@ 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 +from packaging.version import Version + +NUMPY_VERSION = Version(version("numpy")) _HANDLED_FUNCTIONS = {} @@ -1271,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: @@ -2262,14 +2234,46 @@ 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: -implements(np.average)( - _propagate_cosmo_array_attributes_to_result(np.average._implementation) -) + +@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 + + else: + super_average = np.average._implementation + + 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), + ) + else: + return _return_helper(res, 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)( diff --git a/swiftsimio/objects.py b/swiftsimio/objects.py index a367b5a3..22f9d740 100644 --- a/swiftsimio/objects.py +++ b/swiftsimio/objects.py @@ -105,6 +105,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, @@ -1628,6 +1629,172 @@ def from_pint( return obj + @classmethod + 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: 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: + # 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 = 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): + result = 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(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 result + def __array_ufunc__( self, ufunc: np.ufunc, @@ -1687,20 +1854,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, *helper_result["args"], **helper_result["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: @@ -1716,7 +1883,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 @@ -1803,10 +1970,15 @@ def __mul__( ~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 _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) @@ -1819,12 +1991,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 \ @@ -1836,67 +2002,11 @@ def __rmul__( ~swiftsimio.objects.cosmo_array The result of the multiplication. """ - if isinstance(b, unyt.unit_object.Unit): + if getattr(b, "is_Unit", False): return self.__mul__(b) else: return super().__rmul__(b) - def __truediv__( - self, b: 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 - ------- - ~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: 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 - ------- - ~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): """ @@ -2007,7 +2117,7 @@ def __new__( The constructed object. """ if bypass_validation is True: - ret = super().__new__( + result = super().__new__( cls, np.asarray(input_scalar), units=units, @@ -2047,7 +2157,7 @@ def __new__( if compression is None else compression ) - ret = super().__new__( + result = super().__new__( cls, np.asarray(input_scalar), units=units, @@ -2062,9 +2172,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 113735d1..9af3344d 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" @@ -486,7 +490,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), @@ -500,7 +503,11 @@ def test_explicitly_handled_funcs(self): "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))), + "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))), @@ -508,11 +515,15 @@ def test_explicitly_handled_funcs(self): "correlate": (ca(np.arange(3)), ca(np.arange(3))), "tensordot": (ca(np.eye(3)), ca(np.eye(3))), "unwrap": (ca(np.arange(3)),), - "interp": (ca(np.arange(3)), ca(np.arange(3)), ca(np.arange(3))), + "interp": ( + ca(np.arange(3)), + ca(np.arange(3)), + ca(np.arange(3)), + ), "array_repr": (ca(np.arange(3)),), "linalg.outer": (ca(np.arange(3)), ca(np.arange(3))), "trapezoid": (ca(np.arange(3)),), - "in1d": (ca(np.arange(3)), ca(np.arange(3))), # 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)),), @@ -524,7 +535,10 @@ def test_explicitly_handled_funcs(self): "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 + "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)),), @@ -606,7 +620,10 @@ def test_explicitly_handled_funcs(self): "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))), + "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))),), @@ -630,7 +647,11 @@ def test_explicitly_handled_funcs(self): "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))), + "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))), @@ -655,16 +676,11 @@ 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(): - # ----- 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 +767,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 +842,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 +855,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 +887,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) @@ -988,6 +986,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: """ @@ -1077,6 +1086,22 @@ 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: """Tests of explicit (deep)copying of cosmo_array.""" @@ -1136,9 +1161,6 @@ def test_multiplication_by_unyt(self): We desire consistent behaviour for example for `cosmo_array(...) * (1 * u.Mpc)` as for `cosmo_array(...) * u.Mpc`. - - Right-sided multiplication & division can't be supported without upstream - changes in unyt, see `test_rmultiplication_by_unyt`. """ ca = cosmo_array( np.ones(3), u.Mpc, comoving=True, scale_factor=1.0, scale_exponent=1 @@ -1150,46 +1172,16 @@ def test_multiplication_by_unyt(self): # get the same result twice through left-sided multiplication and division: lmultiplied_by_unyt = ca * u.Mpc ldivided_by_unyt = ca / u.Mpc**-1 - - for multiplied_by_unyt in (lmultiplied_by_unyt, ldivided_by_unyt): - assert isinstance(multiplied_by_quantity, cosmo_array) - assert isinstance(multiplied_by_unyt, cosmo_array) - assert np.allclose( - multiplied_by_unyt.to_value(multiplied_by_quantity.units), - multiplied_by_quantity.to_value(multiplied_by_quantity.units), - ) - - @pytest.mark.xfail - def test_rmultiplication_by_unyt(self): - """ - Check that right-sided multiplication behaves itself. - - 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. - - 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: + # and twice more through right-sided multiplication and division: rmultiplied_by_unyt = u.Mpc * ca - rdivided_by_unyt = u.Mpc**2 / ca + rdivided_by_unyt = u.Mpc**3 / ca - 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 np.allclose( diff --git a/tests/test_visualisation.py b/tests/test_visualisation.py index 83cb588b..51337ca2 100644 --- a/tests/test_visualisation.py +++ b/tests/test_visualisation.py @@ -1055,25 +1055,28 @@ def test_comoving_versus_physical(cosmological_volume_only_single): # this test is pretty slow if we don't mask out some particles m = mask(cosmological_volume_only_single) boxsize = m.metadata.boxsize - m.constrain_spatial([[0.0 * b, 0.05 * b] for b in boxsize]) - region = [ - 0.0 * boxsize[0], - 0.05 * boxsize[0], - 0.0 * boxsize[1], - 0.05 * boxsize[1], - 0.0 * boxsize[2], - 0.05 * 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(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, 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, 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 @@ -1090,7 +1093,7 @@ def test_comoving_versus_physical(cosmological_volume_only_single): data, resolution=64, project="densities", - region=region, + region=region.flatten(), parallel=True, ) assert data.gas.densities.comoving is False and img.comoving is False @@ -1100,7 +1103,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, 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 @@ -1114,14 +1121,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, 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, 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 @@ -1138,7 +1153,7 @@ def test_comoving_versus_physical(cosmological_volume_only_single): data, resolution=64, project="densities", - region=region, + region=region.flatten(), parallel=True, ) assert data.gas.densities.comoving and img.comoving