From 6d48907167b6e335b177d91999935dfae552aeca Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 31 Dec 2023 17:51:34 -0500 Subject: [PATCH 01/12] remove gala license - blessed by @adrn --- src/galax/units.py | 29 +---------------------------- 1 file changed, 1 insertion(+), 28 deletions(-) diff --git a/src/galax/units.py b/src/galax/units.py index 849c6702..0482ce35 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -1,31 +1,4 @@ -"""Paired down UnitSystem class from gala. - -See gala's license below. - -``` -The MIT License (MIT) - -Copyright (c) 2012-2023 Adrian M. Price-Whelan - -Permission is hereby granted, free of charge, to any person obtaining a copy -of this software and associated documentation files (the "Software"), to deal -in the Software without restriction, including without limitation the rights -to use, copy, modify, merge, publish, distribute, sublicense, and/or sell -copies of the Software, and to permit persons to whom the Software is -furnished to do so, subject to the following conditions: - -The above copyright notice and this permission notice shall be included in all -copies or substantial portions of the Software. - -THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR -IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, -FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE -AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER -LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, -OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE -SOFTWARE. -``` -""" +"""Tools for representing systems of units using ``astropy.units``.""" __all__ = [ "UnitSystem", From db6c8bcbd38962be1088f9f3e9615eb0c81c2f32 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Sun, 31 Dec 2023 17:52:02 -0500 Subject: [PATCH 02/12] add docstring and type hints for input --- src/galax/units.py | 57 ++++++++++++++++++++++++++++++++++++++++++---- 1 file changed, 53 insertions(+), 4 deletions(-) diff --git a/src/galax/units.py b/src/galax/units.py index 0482ce35..2cf07ce3 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -9,7 +9,7 @@ ] from collections.abc import Iterator -from typing import Any, ClassVar, no_type_check +from typing import ClassVar, Union, no_type_check import astropy.units as u from astropy.units.physical import _physical_unit_mapping @@ -17,7 +17,55 @@ @no_type_check # TODO: get beartype working with this class UnitSystem: - """Represents a system of units.""" + """Represents a system of units. + + At minimum, this consists of a set of length, time, mass, and angle units, but may + also contain preferred representations for composite units. For example, the base + unit system could be ``{kpc, Myr, Msun, radian}``, but you can also specify a + preferred velocity unit, such as ``km/s``. + + This class behaves like a dictionary with keys set by physical types (i.e. "length", + "velocity", "energy", etc.). If a unit for a particular physical type is not + specified on creation, a composite unit will be created with the base units. See the + examples below for some demonstrations. + + Parameters + ---------- + **units, *units + The units that define the unit system. At minimum, this must contain length, + time, mass, and angle units. + + Examples + -------- + If only base units are specified, any physical type specified as a key + to this object will be composed out of the base units:: + + >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian) + >>> usys["velocity"] + Unit("m / s") + + However, preferred representations for composite units can also be specified when + initializing, either as a positional argument:: + + >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, u.erg) + >>> usys.preferred("energy") + Unit("erg") + + Or as a keyword argument:: + + >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, energy=u.erg) + >>> usys.preferred("energy") + Unit("erg") + + This is useful for Galactic dynamics where lengths and times are usually given in + terms of ``kpc`` and ``Myr``, but velocities are often specified in ``km/s``:: + + >>> usys = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, velocity=u.km/u.s) + >>> usys["velocity"] + Unit("kpc / Myr") + >>> usys.preferred("velocity") + Unit("km / s") + """ _core_units: list[u.UnitBase] _registry: dict[u.PhysicalType, u.UnitBase] @@ -29,8 +77,9 @@ class UnitSystem: u.get_physical_type("angle"), ] - # TODO: type hint `units` - def __init__(self, units: Any, *args: u.UnitBase) -> None: + def __init__( + self, units: Union[dict[str, u.UnitBase], "UnitSystem"], *args: u.UnitBase + ) -> None: if isinstance(units, UnitSystem): if len(args) > 0: msg = "If passing in a UnitSystem, cannot pass in additional units." From c3b4b734ffde3e7f567e1543f39fc24faa2ec58e Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 12:29:11 -0500 Subject: [PATCH 03/12] docstring fixes --- src/galax/units.py | 10 +++++++--- 1 file changed, 7 insertions(+), 3 deletions(-) diff --git a/src/galax/units.py b/src/galax/units.py index 2cf07ce3..fbb5b458 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -31,9 +31,10 @@ class UnitSystem: Parameters ---------- - **units, *units + *units, **units The units that define the unit system. At minimum, this must contain length, - time, mass, and angle units. + time, mass, and angle units. If passing in keyword arguments, the keys must be + valid :mod:`astropy.units` physical types. Examples -------- @@ -78,7 +79,9 @@ class UnitSystem: ] def __init__( - self, units: Union[dict[str, u.UnitBase], "UnitSystem"], *args: u.UnitBase + self, + units: Union[u.UnitBase, dict[str, u.UnitBase], "UnitSystem"], + *args: u.UnitBase, ) -> None: if isinstance(units, UnitSystem): if len(args) > 0: @@ -127,6 +130,7 @@ def __getitem__(self, key: str | u.PhysicalType) -> u.UnitBase: return unit def __len__(self) -> int: + # Note: This is required for q.decompose(usys) to work, where q is a Quantity return len(self._core_units) def __iter__(self) -> Iterator[u.UnitBase]: From 0b8df309896029d6a438bd6ea49ebad6f1f2b3b6 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 12:29:35 -0500 Subject: [PATCH 04/12] Add nox command to run doctests --- noxfile.py | 15 +++++++++++++++ 1 file changed, 15 insertions(+) diff --git a/noxfile.py b/noxfile.py index 227ff1e0..c7f17d2a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -38,6 +38,21 @@ def tests(session: nox.Session) -> None: session.run("pytest", *session.posargs) +@nox.session +def doctests(session: nox.Session) -> None: + """Run the regular tests and doctests.""" + session.install(".[test]") + session.run( + "pytest", + "--doctest-modules", + '--doctest-glob="*.rst"', + '--doctest-glob="*.md"', + "docs", + "src/galax", + *session.posargs, + ) + + @nox.session(reuse_venv=True) def docs(session: nox.Session) -> None: """Build the docs. Pass "--serve" to serve. Pass "-b linkcheck" to check links.""" From 1e2330202902c42f4b46ad751a104c915cee26ec Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 12:29:53 -0500 Subject: [PATCH 05/12] started implementation of preferred() --- src/galax/units.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/src/galax/units.py b/src/galax/units.py index fbb5b458..bd7342c4 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -112,7 +112,8 @@ def __init__( self._core_units.append(self._registry[phys_type]) def __getitem__(self, key: str | u.PhysicalType) -> u.UnitBase: - if key in self._registry: + key = u.get_physical_type(key) + if key in self._required_dimensions: return self._registry[key] unit = None @@ -151,6 +152,13 @@ def __hash__(self) -> int: """Hash the unit system.""" return hash(tuple(self._core_units) + tuple(self._required_dimensions)) + def preferred(self, key: str | u.PhysicalType) -> u.UnitBase: + """Return the preferred unit for a given physical type.""" + key = u.get_physical_type(key) + if key in self._registry: + return self._registry[key] + return self[key] + class DimensionlessUnitSystem(UnitSystem): """A unit system with only dimensionless units.""" From e546772bdeca2187c1dc26cea952e8f2b4268f5c Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 12:52:50 -0500 Subject: [PATCH 06/12] oops fix wrongness --- src/galax/units.py | 13 ++++--------- 1 file changed, 4 insertions(+), 9 deletions(-) diff --git a/src/galax/units.py b/src/galax/units.py index bd7342c4..ef75abf1 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -45,23 +45,18 @@ class UnitSystem: >>> usys["velocity"] Unit("m / s") - However, preferred representations for composite units can also be specified when - initializing, either as a positional argument:: + However, preferred representations for composite units can also be specified:: >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, u.erg) - >>> usys.preferred("energy") - Unit("erg") - - Or as a keyword argument:: - - >>> usys = UnitSystem(u.m, u.s, u.kg, u.radian, energy=u.erg) + >>> usys["energy"] + Unit("m2 kg / s2") >>> usys.preferred("energy") Unit("erg") This is useful for Galactic dynamics where lengths and times are usually given in terms of ``kpc`` and ``Myr``, but velocities are often specified in ``km/s``:: - >>> usys = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, velocity=u.km/u.s) + >>> usys = UnitSystem(u.kpc, u.Myr, u.Msun, u.radian, u.km/u.s) >>> usys["velocity"] Unit("kpc / Myr") >>> usys.preferred("velocity") From 62333694202148d546f87085ca32b18612e896b0 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 13:19:23 -0500 Subject: [PATCH 07/12] fix type annotation --- src/galax/units.py | 10 ++++++++-- 1 file changed, 8 insertions(+), 2 deletions(-) diff --git a/src/galax/units.py b/src/galax/units.py index ef75abf1..9525988b 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -75,8 +75,14 @@ class UnitSystem: def __init__( self, - units: Union[u.UnitBase, dict[str, u.UnitBase], "UnitSystem"], - *args: u.UnitBase, + units: Union[ + u.UnitBase, + u.Quantity, + dict[str, u.UnitBase], + dict[str, u.Quantity], + "UnitSystem", + ], + *args: u.UnitBase | u.Quantity, ) -> None: if isinstance(units, UnitSystem): if len(args) > 0: From 86162c096db995e18894f8b814c40462e3843788 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 13:19:37 -0500 Subject: [PATCH 08/12] ignore ruff telling me pickle is unsafe --- pyproject.toml | 1 + 1 file changed, 1 insertion(+) diff --git a/pyproject.toml b/pyproject.toml index 09f77296..398c4604 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -156,6 +156,7 @@ ignore = [ "TD002", # Missing author in TODO "TD003", # Missing issue link on the line following this TODO "UP037", # Remove quote from type annotation <- jaxtyping + "S301", # Pickle unsafe # TODO: fix these "ARG001", "ARG002", From 07d0cfea038719e7f9df26c9cf782a3b12d8f7b4 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 13:20:02 -0500 Subject: [PATCH 09/12] add as_preferred and unit tests of units...yep --- src/galax/units.py | 4 ++ tests/unit/test_units.py | 91 ++++++++++++++++++++++++++++++++++++++++ 2 files changed, 95 insertions(+) create mode 100644 tests/unit/test_units.py diff --git a/src/galax/units.py b/src/galax/units.py index 9525988b..bc6ee217 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -160,6 +160,10 @@ def preferred(self, key: str | u.PhysicalType) -> u.UnitBase: return self._registry[key] return self[key] + def as_preferred(self, quantity: u.Quantity) -> u.Quantity: + """Convert a quantity to the preferred unit for this unit system.""" + return quantity.to(self.preferred(quantity.unit.physical_type)) + class DimensionlessUnitSystem(UnitSystem): """A unit system with only dimensionless units.""" diff --git a/tests/unit/test_units.py b/tests/unit/test_units.py new file mode 100644 index 00000000..3b940e87 --- /dev/null +++ b/tests/unit/test_units.py @@ -0,0 +1,91 @@ +# Standard library +import pickle + +# Third party +import astropy.units as u +import numpy as np +import pytest + +# This package +from galax.units import UnitSystem, dimensionless + + +def test_init(): + usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun) + + with pytest.raises( + ValueError, match="must specify a unit for the physical type .*mass" + ): + UnitSystem(u.kpc, u.Myr, u.radian) # no mass + + with pytest.raises( + ValueError, match="must specify a unit for the physical type .*angle" + ): + UnitSystem(u.kpc, u.Myr, u.Msun) + + with pytest.raises( + ValueError, match="must specify a unit for the physical type .*time" + ): + UnitSystem(u.kpc, u.radian, u.Msun) + + with pytest.raises( + ValueError, match="must specify a unit for the physical type .*length" + ): + UnitSystem(u.Myr, u.radian, u.Msun) + + usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun) + usys = UnitSystem(usys) + + +def test_quantity_init(): + usys = UnitSystem(5 * u.kpc, 50 * u.Myr, 1e5 * u.Msun, u.rad) + assert np.isclose((8 * u.Myr).decompose(usys).value, 8 / 50) + + +def test_preferred(): + usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.km / u.s) + q = 15.0 * u.km / u.s + assert usys.preferred("velocity") == u.km / u.s + assert q.decompose(usys).unit == u.kpc / u.Myr + assert usys.as_preferred(q).unit == u.km / u.s + + +def test_dimensionless(): + assert dimensionless["dimensionless"] == u.one + assert dimensionless["length"] == u.one + + with pytest.raises(ValueError, match="can not be decomposed into"): + (15 * u.kpc).decompose(dimensionless) + + with pytest.raises(ValueError, match="are not convertible"): + dimensionless.as_preferred(15 * u.kpc) + + +def test_compare(): + usys1 = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.mas / u.yr) + usys1_clone = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.mas / u.yr) + + usys2 = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun, u.kiloarcsecond / u.yr) + usys3 = UnitSystem(u.kpc, u.Myr, u.radian, u.kg, u.mas / u.yr) + + assert usys1 == usys1_clone + assert usys1_clone == usys1 + + assert usys1 != usys2 + assert usys2 != usys1 + + assert usys1 != usys3 + assert usys3 != usys1 + + +def test_pickle(tmpdir): + usys = UnitSystem(u.kpc, u.Myr, u.radian, u.Msun) + + path = tmpdir / "test.pkl" + with path.open(mode="wb") as f: + pickle.dump(usys, f) + + with path.open(mode="rb") as f: + usys2 = pickle.load(f) + + assert usys == usys2 From 7e97205465f55ae3c3e4db0aee49ccf1463b25fe Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Mon, 1 Jan 2024 13:31:43 -0500 Subject: [PATCH 10/12] remove no type check --- src/galax/units.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/src/galax/units.py b/src/galax/units.py index bc6ee217..da415629 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -9,13 +9,12 @@ ] from collections.abc import Iterator -from typing import ClassVar, Union, no_type_check +from typing import ClassVar, Union import astropy.units as u from astropy.units.physical import _physical_unit_mapping -@no_type_check # TODO: get beartype working with this class UnitSystem: """Represents a system of units. From 7841d3d6c5ee805a5e0366c519d4e99e2ae45512 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 2 Jan 2024 16:37:14 -0500 Subject: [PATCH 11/12] address review comments --- pyproject.toml | 3 +-- src/galax/units.py | 2 -- 2 files changed, 1 insertion(+), 4 deletions(-) diff --git a/pyproject.toml b/pyproject.toml index 398c4604..1a0f2ae5 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -156,7 +156,6 @@ ignore = [ "TD002", # Missing author in TODO "TD003", # Missing issue link on the line following this TODO "UP037", # Remove quote from type annotation <- jaxtyping - "S301", # Pickle unsafe # TODO: fix these "ARG001", "ARG002", @@ -169,7 +168,7 @@ ignore = [ [tool.ruff.lint.per-file-ignores] "__init__.py" = ["F403"] "tests/**" = [ - "ANN", "D10", "E731", "INP001", "S101", "SLF001", "T20", + "ANN", "D10", "E731", "INP001", "S101", "S301", "SLF001", "T20", "TID252", # Relative imports from parent modules are banned ] "noxfile.py" = ["ERA001", "T20"] diff --git a/src/galax/units.py b/src/galax/units.py index da415629..09809901 100644 --- a/src/galax/units.py +++ b/src/galax/units.py @@ -77,8 +77,6 @@ def __init__( units: Union[ u.UnitBase, u.Quantity, - dict[str, u.UnitBase], - dict[str, u.Quantity], "UnitSystem", ], *args: u.UnitBase | u.Quantity, From 1d3dc08284f747b5d9df946f1d357a3bc24f2e26 Mon Sep 17 00:00:00 2001 From: Adrian Price-Whelan Date: Tue, 2 Jan 2024 18:19:51 -0500 Subject: [PATCH 12/12] add doctests to nox sessions list --- noxfile.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/noxfile.py b/noxfile.py index c7f17d2a..4be50f7a 100644 --- a/noxfile.py +++ b/noxfile.py @@ -8,7 +8,7 @@ DIR = Path(__file__).parent.resolve() -nox.options.sessions = ["lint", "tests"] +nox.options.sessions = ["lint", "tests", "doctests"] @nox.session