Merge pull request #2524 from cta-observatory/unit_trait

Add a trait for astropy quantities

docs/changes/2524.feature.rst

@@ -0,0 +1 @@
+Add an ``AstroQuantity`` trait which can hold any ``astropy.units.Quantity``.

src/ctapipe/core/tests/test_traits.py

@@ -2,13 +2,15 @@
 import pathlib
 import tempfile
 from abc import ABCMeta, abstractmethod
+from subprocess import CalledProcessError
 from unittest import mock
 
 import pytest
 from traitlets import CaselessStrEnum, HasTraits, Int
 
-from ctapipe.core import Component
+from ctapipe.core import Component, Tool, run_tool
 from ctapipe.core.traits import (
+    AstroQuantity,
     AstroTime,
     List,
     Path,
@@ -278,6 +280,99 @@ class NoNone(Component):
         c.time = None
 
 
+def test_quantity():
+    import astropy.units as u
+
+    class SomeComponentWithQuantityTrait(Component):
+        quantity = AstroQuantity()
+
+    c = SomeComponentWithQuantityTrait()
+    c.quantity = -1.754 * u.m / (u.s * u.deg)
+    assert isinstance(c.quantity, u.Quantity)
+    assert c.quantity.value == -1.754
+    assert c.quantity.unit == u.Unit("m / (deg s)")
+
+    c.quantity = "1337 erg / centimeter**2 second"
+    assert isinstance(c.quantity, u.Quantity)
+    assert c.quantity.value == 1337
+    assert c.quantity.unit == u.Unit("erg / (s cm2)")
+
+    with pytest.raises(TraitError):
+        c.quantity = "No quantity"
+
+    # Try misspelled/ non-existent unit
+    with pytest.raises(TraitError):
+        c.quantity = "5 meters"
+
+    # Test definition of physical type
+    class SomeComponentWithEnergyTrait(Component):
+        energy = AstroQuantity(physical_type=u.physical.energy)
+
+    c = SomeComponentWithEnergyTrait()
+
+    class AnotherComponentWithEnergyTrait(Component):
+        energy = AstroQuantity(physical_type=u.TeV)
+
+    c = AnotherComponentWithEnergyTrait()
+
+    with pytest.raises(
+        TraitError,
+        match="Given physical type must be either of type"
+        + " astropy.units.PhysicalType or a subclass of"
+        + f" astropy.units.UnitBase, was {type(5 * u.TeV)}.",
+    ):
+
+        class SomeBadComponentWithEnergyTrait(Component):
+            energy = AstroQuantity(physical_type=5 * u.TeV)
+
+    with pytest.raises(
+        TraitError,
+        match=f"Given physical type {u.physical.energy} does not match"
+        + f" physical type of the default value, {u.get_physical_type(5 * u.m)}.",
+    ):
+
+        class AnotherBadComponentWithEnergyTrait(Component):
+            energy = AstroQuantity(
+                default_value=5 * u.m, physical_type=u.physical.energy
+            )
+
+
+def test_quantity_tool(capsys):
+    import astropy.units as u
+
+    class MyTool(Tool):
+        energy = AstroQuantity(physical_type=u.physical.energy).tag(config=True)
+
+    tool = MyTool()
+    run_tool(tool, ["--MyTool.energy=5 TeV"])
+    assert tool.energy == 5 * u.TeV
+
+    with pytest.raises(CalledProcessError):
+        run_tool(tool, ["--MyTool.energy=5 m"], raises=True)
+
+    captured = capsys.readouterr()
+    assert (
+        captured.err.split(":")[-1]
+        == f" Given quantity is of physical type {u.get_physical_type(5 * u.m)}."
+        + f" Expected {u.physical.energy}.\n"
+    )
+
+
+def test_quantity_none():
+    class AllowNone(Component):
+        quantity = AstroQuantity(default_value=None, allow_none=True)
+
+    c = AllowNone()
+    assert c.quantity is None
+
+    class NoNone(Component):
+        quantity = AstroQuantity(default_value="5 meter", allow_none=False)
+
+    c = NoNone()
+    with pytest.raises(TraitError):
+        c.quantity = None
+
+
 def test_component_name():
     from ctapipe.core.traits import ComponentName, ComponentNameList
 

src/ctapipe/core/traits.py

@@ -5,6 +5,7 @@
 import pathlib
 from urllib.parse import urlparse
 
+import astropy.units as u
 import traitlets
 import traitlets.config
 from astropy.time import Time
@@ -17,6 +18,7 @@
 
 __all__ = [
     # Implemented here
+    "AstroQuantity",
     "AstroTime",
     "BoolTelescopeParameter",
     "IntTelescopeParameter",
@@ -72,8 +74,60 @@
 flag = traitlets.config.boolean_flag
 
 
+class AstroQuantity(TraitType):
+    """A trait containing an ``astropy.units`` quantity."""
+
+    def __init__(self, physical_type=None, **kwargs):
+        super().__init__(**kwargs)
+        if physical_type is not None:
+            if isinstance(physical_type, u.PhysicalType):
+                self.physical_type = physical_type
+            elif isinstance(physical_type, u.UnitBase):
+                self.physical_type = u.get_physical_type(physical_type)
+            else:
+                raise TraitError(
+                    "Given physical type must be either of type"
+                    " astropy.units.PhysicalType or a subclass of"
+                    f" astropy.units.UnitBase, was {type(physical_type)}."
+                )
+        else:
+            self.physical_type = physical_type
+
+        if self.default_value is not Undefined and self.physical_type is not None:
+            default_type = u.get_physical_type(self.default_value)
+            if default_type != self.physical_type:
+                raise TraitError(
+                    f"Given physical type {self.physical_type} does not match"
+                    f" physical type of the default value, {default_type}."
+                )
+
+    def info(self):
+        info = "An ``astropy.units.Quantity`` instance"
+        if self.allow_none:
+            info += "or None"
+        return info
+
+    def validate(self, obj, value):
+        try:
+            quantity = u.Quantity(value)
+        except TypeError:
+            return self.error(obj, value)
+        except ValueError:
+            return self.error(obj, value)
+
+        if self.physical_type is not None:
+            given_type = u.get_physical_type(quantity)
+            if given_type != self.physical_type:
+                raise TraitError(
+                    f"Given quantity is of physical type {given_type}."
+                    f" Expected {self.physical_type}."
+                )
+
+        return quantity
+
+
 class AstroTime(TraitType):
-    """A trait representing a point in Time, as understood by `astropy.time`"""
+    """A trait representing a point in Time, as understood by ``astropy.time``."""
 
     def validate(self, obj, value):
         """try to parse and return an ISO time string"""