WIP: Replacing unit management system

autochem/rate/_03rate.py

@@ -3,9 +3,9 @@
 import numpy
 import pint
 import pydantic
-from numpy.typing import ArrayLike
+from numpy.typing import ArrayLike, NDArray
 
-from ..util import unit_old
+from ..unit_ import UnitsData
 from ._00read import parse_results_from_chemkin_string
 from ._02const import RateConstant_, extract_rate_constant_from_chemkin_parse_results
 
@@ -24,8 +24,8 @@ def unit(self) -> pint.Unit:
         return self.rate_constant.unit
 
     def __call__(
-        self, t: ArrayLike, p: ArrayLike, units: unit_old.UnitsData | None = None
-    ) -> float | numpy.ndarray:
+        self, t: ArrayLike, p: ArrayLike = 1, units: UnitsData | None = None
+    ) -> NDArray[numpy.float64]:
         """Evaluate rate constant.
 
         Uses:
@@ -43,19 +43,9 @@ def __call__(
         """
         return self.rate_constant(t=t, p=p, units=units)
 
-    def convert_units(self, units: unit_old.UnitsData | None = None):
-        """Convert units.
-
-        :param units: Units
-        :return: Object
-        """
-        rate = self.model_copy()
-        rate.rate_constant = self.rate_constant.convert_units(units)
-        return rate
-
 
 def from_chemkin_string(
-    rate_str: str, units: unit_old.UnitsData | None = None, strict: bool = True
+    rate_str: str, units: UnitsData | None = None, strict: bool = True
 ) -> Rate:
     """Read rate from Chemkin string.
 

autochem/rate/__init__.py

@@ -8,13 +8,13 @@
     TroeBlendingFunction,
 )
 from ._02const import (
-    ActivatedRateConstant,
+    # ActivatedRateConstant,
     ArrheniusRateConstant,
-    BlendedRateConstant,
-    ChebRateConstant,
-    FalloffRateConstant,
+    # BlendedRateConstant,
+    # ChebRateConstant,
+    # FalloffRateConstant,
     ParamRateConstant,
-    PlogRateConstant,
+    # PlogRateConstant,
     RateConstant,
     RawRateConstant,
 )

autochem/unit_/__init__.py

@@ -0,0 +1,15 @@
+"""Units functions."""
+
+from . import system
+from ._manager import UnitManager, manage_units
+from .system import UNITS, Dimension, Units, UnitsData
+
+__all__ = [
+    "UNITS",
+    "Dimension",
+    "UnitManager",
+    "Units",
+    "UnitsData",
+    "manage_units",
+    "system",
+]

autochem/unit_/_manager.py

@@ -0,0 +1,76 @@
+"""Defines abstract base class for automatically handling model units."""
+
+import abc
+import functools
+from collections.abc import Callable, Sequence
+from typing import ClassVar, Protocol, Self
+
+import numpy
+from numpy.typing import NDArray
+
+from ..util.type_ import Frozen
+from . import system
+from .system import UNITS, Dimension, Units, UnitsData
+
+
+class UnitManager(Frozen, abc.ABC):
+
+    _dimensions: ClassVar[dict[str, Dimension]]
+
+    def __init__(self, units: UnitsData | None = None, **kwargs):
+        if units is not None:
+            units0 = Units.model_validate(units)
+            for key, dim in self.__class__._dimensions.items():
+                val0 = kwargs.get(key)
+                kwargs[key] = system.convert_dimension_value(
+                    units0, UNITS, dim, val0, **kwargs
+                )
+
+        super().__init__(**kwargs)
+
+
+def manage_units(arg_dims: Sequence[Dimension], ret_dim: Dimension):
+    """Transform function into a unit managing function.
+
+    Converts arguments to internal units, calls the function, then converts the return
+    value back to desired units.
+
+    :param arg_dims: Argument dimensions
+    :param ret_dim: Return dimensions
+    :return: Function decorator
+    """
+
+    # def manage_units_(func0: UnitMethod) -> UnitMethod:
+    def manage_units_(func0):
+
+        @functools.wraps(func0)
+        def func(
+            self, *args, units: UnitsData | None = None, **kwargs
+        ) -> float | NDArray[numpy.float64]:
+            # If no units were specified, return as-is
+            if units is None:
+                return func0(self, *args, units=units, **kwargs)
+
+            # Process units
+            units0 = Units.model_validate(units)
+
+            # Convert arguments
+            args_ = tuple(
+                system.convert_dimension_value(
+                    units0, UNITS, dim, arg, **self.model_dump()
+                )
+                for dim, arg in zip(arg_dims, args, strict=True)
+            )
+
+            # Call function
+            ret = func0(self, *args_, units=units, **kwargs)
+
+            # Convert return
+            ret_ = system.convert_dimension_value(
+                UNITS, units0, ret_dim, ret, **self.model_dump()
+            )
+            return ret_
+
+        return func
+
+    return manage_units_

autochem/util/unit_/system.py → autochem/unit_/system.py

@@ -9,7 +9,7 @@
 import pint
 from numpy.typing import NDArray
 
-from ..type_ import Frozen, Unit_
+from ..util.type_ import Frozen, Unit_
 
 
 # Model for specifying units
@@ -27,30 +27,30 @@ class Units(Frozen):
     @functools.cached_property
     def energy_per_substance(self) -> pint.Unit:
         """Energy per substance unit."""
-        return self.energy / self.substance
+        return pint.Unit(self.energy / self.substance)
 
     @functools.cached_property
     def concentration(self) -> pint.Unit:
         """Concentration unit."""
-        return self.substance / self.length**3
+        return pint.Unit(self.substance / self.length**3)
 
     def rate_constant(self, order: int) -> pint.Unit:
         """Rate constant unit.
 
         :param order: Reaction order
         """
         if order == 1:
-            return 1 / self.time
+            return pint.Unit(1 / self.time)
 
-        return self.concentration ** (1 - order) / self.time
+        return pint.Unit(self.concentration ** (1 - order) / self.time)
 
 
 # Alias for unit-convertible data
 UnitData = str | pint.Unit
 UnitsData = Mapping[str, UnitData] | Units
 
 
-# Instance representing internal unit system
+# Internal unit system (defaults of above Units class)
 UNITS = Units()
 
 
@@ -76,6 +76,17 @@ class Dimension(enum.StrEnum):
 
 
 # Functions
+def gas_constant_value(units: Units) -> float:
+    """Determine gas constant value in unit system.
+
+    :param units: Unit system
+    :return: Gas constant value
+    """
+    return pint.Quantity("molar_gas_constant").m_as(
+        units.energy_per_substance / units.temperature
+    )
+
+
 def dimension_unit(units: Units, dim: str | Dimension, **kwargs) -> pint.Unit:
     """Determine the unit for a dimension.
 
@@ -92,9 +103,9 @@ def dimension_unit(units: Units, dim: str | Dimension, **kwargs) -> pint.Unit:
     return getattr(units, dim)
 
 
-def convert_dimension(
+def convert_dimension_value(
     units: Units, new_units: Units, dim: str | Dimension, val: Any, **kwargs
-) -> float | NDArray[numpy.float64]:
+) -> NDArray[numpy.float64]:
     """Convert dimension value to new units.
 
     :param units: Units sytem
@@ -107,4 +118,4 @@ def convert_dimension(
     unit = dimension_unit(units, dim, **kwargs)
     new_unit = dimension_unit(new_units, dim, **kwargs)
     new_val = pint.Quantity(val, unit).m_as(new_unit)
-    return new_val
+    return numpy.array(new_val, dtype=numpy.float64)

autochem/util/plot.py

@@ -61,6 +61,6 @@ def arrhenius_plot(
     return (
         altair.Chart(data)
         .mark_line()
-        .transform_fold(fold=kTs.keys())
+        .transform_fold(fold=list(kTs.keys()))
         .encode(x=x, y=y, color=color)
     )