Merge pull request #620 from avcopan/dev

New: Simplifies plotting

autochem/rate/_01const.py

@@ -4,8 +4,10 @@
 from collections.abc import Mapping
 from typing import Annotated, ClassVar
 
+import altair
 import more_itertools as mit
 import numpy
+import pandas
 import pint
 import pydantic
 import xarray
@@ -16,7 +18,7 @@
 from .. import unit_
 from ..unit_ import UNITS, Dimension, UnitManager, Units, UnitsData
 from ..util import chemkin
-from ..util.type_ import Frozen, NDArray_, Scalable, Scalers, SubclassTyped
+from ..util.type_ import Frozen, NDArray_, Number, Scalable, Scalers, SubclassTyped
 from ._00func import (
     BlendingFunction_,
     extract_blending_function_from_chemkin_parse_results,
@@ -79,6 +81,79 @@ def process_input(
 
         return t, p
 
+    def process_output(self, ktp: ArrayLike) -> NDArray[numpy.float64]:
+        """Normalize rate constant output, clipping unphyiscal negative values.
+
+        :param ktp: Rate constant values
+        :return: Rate constant values
+        """
+        return numpy.where(numpy.less_equal(ktp, 0), numpy.nan, ktp)
+
+    def display(
+        self,
+        others: "Mapping[str, RateConstant] | 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
+        """
+        # Process units
+        units = UNITS if units is None else Units.model_validate(units)
+        x_unit = unit_.pretty_string(units.temperature**-1)
+        y_unit = unit_.pretty_string(units.rate_constant(self.order))
+
+        # Add units to labels
+        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}
+
+        # Gether data from functons
+        t = numpy.linspace(*t_range, num=500)
+        data_dct = {lab: func(t, p) for lab, func in funcs.items()}
+        data = pandas.DataFrame({"x": numpy.divide(1000, t), **data_dct})
+
+        # Determine exponent range
+        vals_arr = numpy.array(list(data_dct.values()))
+        is_nan = numpy.isnan(vals_arr)
+        exp_arr = numpy.log10(vals_arr, where=~is_nan)
+        exp_arr[is_nan] = 0.0
+        exp_arr = numpy.rint(exp_arr).astype(int)
+        exp_max = numpy.max(exp_arr).item()
+        exp_min = numpy.min(exp_arr).item()
+        y_vals = [10**x for x in range(exp_min, exp_max + 2)]
+
+        # Prepare encoding parameters
+        x = altair.X("x", title=x_label)
+        y = (
+            altair.Y("value:Q", title=y_label)
+            .scale(type="log")
+            .axis(format=".1e", values=y_vals)
+        )
+        color = "key:N" if others else altair.Undefined
+
+        # Create chart
+        return (
+            altair.Chart(data)
+            .mark_line()
+            .transform_fold(fold=list(data_dct.keys()))
+            .encode(x=x, y=y, color=color)
+        )
+
 
 class RawRateConstant(RateConstant):
     ts: list[float]
@@ -111,7 +186,7 @@ def __call__(
         """Evaluate rate constant."""
         t, p = self.process_input(t, p)
         ktp: NDArray[numpy.float64] = self.ktp.sel(t=t, p=p, method="ffill").data
-        return ktp
+        return self.process_output(ktp)
 
 
 class ParamRateConstant(RateConstant):
@@ -160,7 +235,7 @@ def __call__(
         t, p = self.process_input(t, p)
         r = unit_.system.gas_constant_value(UNITS)
         ktp = self.A * (t**self.b) * numpy.exp(-self.E / (r * t))
-        return ktp
+        return self.process_output(ktp)
 
 
 class BlendedRateConstant(ParamRateConstant, abc.ABC):  # type: ignore[misc]
@@ -245,7 +320,7 @@ def __call__(
         _, k_high = self.arrhenius_functions
         p_r = self.effective_reduced_pressure(t, p)
         ktp = k_high(t) * p_r / (1 + p_r) * self.function(t, p_r)
-        return ktp
+        return self.process_output(ktp)
 
 
 class ActivatedRateConstant(BlendedRateConstant):
@@ -264,7 +339,7 @@ def __call__(
         k_low, _ = self.arrhenius_functions
         p_r = self.effective_reduced_pressure(t, p)
         ktp = k_low(t) / (1 + p_r) * self.function(t, p_r)
-        return ktp
+        return self.process_output(ktp)
 
 
 class PlogRateConstant(ParamRateConstant):
@@ -305,7 +380,7 @@ def __call__(
         elif p == p0:
             ktp = kt0
 
-        return ktp
+        return self.process_output(ktp)
 
     @property
     def arrhenius_functions(self) -> list[ArrheniusRateConstant]:
@@ -406,7 +481,7 @@ def __call__(
 
         # AVC: I don't understand why I need to transpose the coefficient matrix here
         ktp = chebyshev.chebgrid2d(t_, p_, self.coeffs.T)
-        return ktp
+        return self.process_output(ktp)
 
 
 RateConstant_ = Annotated[

autochem/rate/_02rate.py

@@ -5,14 +5,15 @@
 from collections.abc import Mapping, Sequence
 from typing import ClassVar
 
+import altair
 import numpy
 import pint
 import pydantic
 from numpy.typing import ArrayLike, NDArray
 
 from ..unit_ import UnitsData
 from ..util import chemkin
-from ..util.type_ import Scalable, Scalers
+from ..util.type_ import Number, Scalable, Scalers
 from ._01const import (
     ArrheniusRateConstant,
     ParamRateConstant,
@@ -77,6 +78,39 @@ 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(

autochem/tests/test__rate.py

@@ -1,10 +1,8 @@
 """Test autochem.rate."""
 
-import numpy
 import pytest
 
 from autochem import rate
-from autochem.util import plot
 
 # Elementary
 SIMPLE = {
@@ -115,19 +113,21 @@ def test__from_chemkin_string(name, data, check_roundtrip: bool):
     # Read
     k = rate.from_chemkin_string(chem_str, units=units)
 
-    # Plot
-    T = numpy.linspace(400, 1250, 500)
-    P = 10
-    k(T, P)
-    plot.arrhenius_plot(T, k(T, P), y_unit=k.unit)
-
     # Write
     chem_str_ = rate.chemkin_string(k)
     print(chem_str_)
     k_ = rate.from_chemkin_string(chem_str_)
+
+    # Check roundtrip
     if check_roundtrip:
         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})
+
 
 if __name__ == "__main__":
     # test__from_chemkin_string("FALLOFF", FALLOFF)

autochem/unit_/__init__.py

@@ -2,7 +2,7 @@
 
 from . import system
 from ._manager import UnitManager, manage_units
-from ._unit import string
+from ._unit import pretty_string, string
 from .system import UNITS, Dimension, Units, UnitsData
 
 __all__ = [
@@ -12,6 +12,7 @@
     "Units",
     "UnitsData",
     "manage_units",
+    "pretty_string",
     "string",
     "system",
 ]

autochem/unit_/_unit.py

@@ -4,4 +4,8 @@
 
 
 def string(unit: pint.Unit) -> str:
-    return format(unit, "~")
+    return format(unit, "~C")
+
+
+def pretty_string(unit: pint.Unit) -> str:
+    return format(unit, "~P")

autochem/util/__init__.py

@@ -1,5 +1,5 @@
 """Utilities."""
 
-from . import chemkin, plot, type_
+from . import chemkin, type_
 
-__all__ = ["chemkin", "plot", "type_"]
+__all__ = ["chemkin", "type_"]