Merge pull request #621 from avcopan/dev

Refactor: Stick with functional rate interface
Auto-Mech · Feb 13, 2025 · 7259273 · 7259273
2 parents 0e4bd8d + 37564b3
commit 7259273
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 44 deletions.
diff --git a/autochem/rate/_01const.py b/autochem/rate/_01const.py
@@ -1,7 +1,7 @@
 """Rate constant models."""
 
 import abc
-from collections.abc import Mapping
+from collections.abc import Mapping, Sequence
 from typing import Annotated, ClassVar
 
 import altair
@@ -91,7 +91,8 @@ def process_output(self, ktp: ArrayLike) -> NDArray[numpy.float64]:
 
     def display(
         self,
-        others: "Mapping[str, RateConstant] | None" = None,
+        others: "Sequence[RateConstant]" = (),
+        labels: Sequence[str] = (),
         t_range: tuple[Number, Number] = (400, 1250),
         p: Number = 1,
         units: UnitsData | None = None,
@@ -101,7 +102,8 @@ def display(
     ) -> altair.Chart:
         """Display as an Arrhenius plot.
 
-        :param others: Other rate constants by label
+        :param others: Other rate constants
+        :param others_labels: Labels for other rate constants
         :param t_range: Temperature range
         :param p: Pressure
         :param units: Units
@@ -118,13 +120,14 @@ def display(
         x_label = f"{x_label} ({x_unit})"
         y_label = f"{y_label} ({y_unit})"
 
-        # Gather functions
-        others = {} if others is None else others
-        funcs = {label: self, **others}
+        # Gather rate constants and labels
+        assert len(others) == len(labels), f"{labels} !~ {others}"
+        all_ks = [self, *others]
+        all_labels = [label, *labels]
 
         # Gether data from functons
         t = numpy.linspace(*t_range, num=500)
-        data_dct = {lab: func(t, p) for lab, func in funcs.items()}
+        data_dct = {lab: k(t, p) for lab, k in zip(all_labels, all_ks, strict=True)}
         data = pandas.DataFrame({"x": numpy.divide(1000, t), **data_dct})
 
         # Determine exponent range

diff --git a/autochem/rate/_02rate.py b/autochem/rate/_02rate.py
@@ -78,39 +78,6 @@ def __call__(
         """
         return self.rate_constant(t=t, p=p, units=units)
 
-    def display(
-        self,
-        others: "Mapping[str, Rate] | None" = None,
-        t_range: tuple[Number, Number] = (400, 1250),
-        p: Number = 1,
-        units: UnitsData | None = None,
-        label: str = "This work",
-        x_label: str = "1000/T",
-        y_label: str = "k",
-    ) -> altair.Chart:
-        """Display as an Arrhenius plot.
-
-        :param others: Other rate constants by label
-        :param t_range: Temperature range
-        :param p: Pressure
-        :param units: Units
-        :param x_label: X-axis label
-        :param y_label: Y-axis label
-        :return: Chart
-        """
-        if others is not None:
-            others = {lab: obj.rate_constant for lab, obj in others.items()}
-
-        return self.rate_constant.display(
-            others=others,
-            t_range=t_range,
-            p=p,
-            units=units,
-            label=label,
-            x_label=x_label,
-            y_label=y_label,
-        )
-
 
 # Constructors
 def from_chemkin_string(
@@ -252,3 +219,39 @@ def chemkin_string(rate: Rate, eq_width: int = 55, dup: bool = False) -> str:
     rate_str = rate_constant_chemkin_string(rate.rate_constant, eq_width=eq_width)
     reac_str = f"{eq:<{eq_width}} {rate_str}"
     return chemkin.write_with_dup(reac_str, dup=dup)
+
+
+# Display
+def display(
+    rate: Rate,
+    comp_rates: Sequence[Rate] = (),
+    comp_labels: Sequence[str] = (),
+    t_range: tuple[Number, Number] = (400, 1250),
+    p: Number = 1,
+    units: UnitsData | None = None,
+    label: str = "This work",
+    x_label: str = "1000/T",
+    y_label: str = "k",
+) -> altair.Chart:
+    """Display as an Arrhenius plot, optionally comparing to other rates.
+
+    :param rate: Rate
+    :param comp_rates: Rates for comparison
+    :param comp_labels: Labels for comparison
+    :param t_range: Temperature range
+    :param p: Pressure
+    :param units: Units
+    :param x_label: X-axis label
+    :param y_label: Y-axis label
+    :return: Chart
+    """
+    return rate.rate_constant.display(
+        others=[r.rate_constant for r in comp_rates],
+        labels=comp_labels,
+        t_range=t_range,
+        p=p,
+        units=units,
+        label=label,
+        x_label=x_label,
+        y_label=y_label,
+    )
diff --git a/autochem/rate/__init__.py b/autochem/rate/__init__.py
@@ -21,6 +21,7 @@
     Rate,
     chemkin_equation,
     chemkin_string,
+    display,
     expand_lumped,
     from_chemkin_string,
 )
@@ -42,6 +43,7 @@
     "TroeBlendingFunction",
     "chemkin_equation",
     "chemkin_string",
+    "display",
     "expand_lumped",
     "from_chemkin_string",
 ]
diff --git a/autochem/tests/test__rate.py b/autochem/tests/test__rate.py
@@ -123,10 +123,10 @@ def test__from_chemkin_string(name, data, check_roundtrip: bool):
         assert k == k_, f"\n   {k}\n!= {k_}"
 
     # Plot against another rate
-    other_units = SIMPLE.get("units")
-    other_chem_str = SIMPLE.get("chemkin")
-    other_k = rate.from_chemkin_string(other_chem_str, units=other_units)
-    k.display({"Other": other_k})
+    copm_units = SIMPLE.get("units")
+    copm_chem_str = SIMPLE.get("chemkin")
+    comp_k = rate.from_chemkin_string(copm_chem_str, units=copm_units)
+    rate.display(k, comp_rates=[comp_k], comp_labels=["comp"])
 
 
 if __name__ == "__main__":