ENH: add support for molecular fields (AMRex)

Co-authored-by: Clément Robert <cr52@protonmail.com>

nose_ignores.txt

@@ -43,3 +43,4 @@
 --ignore-file=test_alt_ray_tracers\.py
 --ignore-file=test_minimal_representation\.py
 --ignore-file=test_set_log_level\.py
+--ignore-file=test_field_parsing\.py

tests/tests.yaml

@@ -187,6 +187,7 @@ other_tests:
      - "--ignore-file=test_ewah_write_load\\.py"
      - "--ignore-file=test_external_frontends\\.py"
      - "--ignore-file=test_field_access_pytest\\.py"
+     - "--ignore-file=test_field_parsing\\.py"
      - "--ignore-file=test_file_sanitizer\\.py"
      - "--ignore-file=test_firefly\\.py"
      - "--ignore-file=test_geometries\\.py"

yt/frontends/amrex/fields.py

@@ -1,4 +1,5 @@
 import re
+import sys
 
 import numpy as np
 
@@ -7,12 +8,15 @@
 from yt.units import YTQuantity
 from yt.utilities.physical_constants import amu_cgs, boltzmann_constant_cgs, c
 
+if sys.version_info >= (3, 10):
+    from typing import TypeAlias
+else:
+    from typing_extensions import TypeAlias
+
 rho_units = "code_mass / code_length**3"
 mom_units = "code_mass / (code_time * code_length**2)"
 eden_units = "code_mass / (code_time**2 * code_length)"  # erg / cm^3
 
-spec_finder = re.compile(r".*\((\D*)(\d*)\).*")
-
 
 def _thermal_energy_density(field, data):
     # What we've got here is UEINT:
@@ -357,16 +361,22 @@ def setup_fluid_fields(self):
         for _, field in self.ds.field_list:
             if field.startswith("X("):
                 # We have a fraction
-                nice_name, tex_label = _nice_species_name(field)
-                self.alias(
-                    ("gas", f"{nice_name}_fraction"), ("boxlib", field), units=""
+                sub = Substance(field)
+                # Overwrite field to use nicer tex_label display_name
+                self.add_output_field(
+                    ("boxlib", field),
+                    sampling_type="cell",
+                    units="",
+                    display_name=rf"X\left({sub.to_tex()}\right)",
                 )
-                func = _create_density_func(("gas", f"{nice_name}_fraction"))
+                self.alias(("gas", f"{sub}_fraction"), ("boxlib", field), units="")
+                func = _create_density_func(("gas", f"{sub}_fraction"))
                 self.add_field(
-                    name=("gas", f"{nice_name}_density"),
+                    name=("gas", f"{sub}_density"),
                     sampling_type="cell",
                     function=func,
                     units=self.ds.unit_system["density"],
+                    display_name=rf"\rho {sub.to_tex()}",
                 )
 
 
@@ -462,29 +472,27 @@ def setup_fluid_fields(self):
         for _, field in self.ds.field_list:
             if field.startswith("X("):
                 # We have a mass fraction
-                nice_name, tex_label = _nice_species_name(field)
+                sub = Substance(field)
                 # Overwrite field to use nicer tex_label display_name
                 self.add_output_field(
                     ("boxlib", field),
                     sampling_type="cell",
                     units="",
-                    display_name=tex_label,
-                )
-                self.alias(
-                    ("gas", f"{nice_name}_fraction"), ("boxlib", field), units=""
+                    display_name=rf"X\left({sub.to_tex()}\right)",
                 )
-                func = _create_density_func(("gas", f"{nice_name}_fraction"))
+                self.alias(("gas", f"{sub}_fraction"), ("boxlib", field), units="")
+                func = _create_density_func(("gas", f"{sub}_fraction"))
                 self.add_field(
-                    name=("gas", f"{nice_name}_density"),
+                    name=("gas", f"{sub}_density"),
                     sampling_type="cell",
                     function=func,
                     units=unit_system["density"],
-                    display_name=rf"\rho {tex_label}",
+                    display_name=rf"\rho {sub.to_tex()}",
                 )
 
             elif field.startswith("omegadot("):
-                nice_name, tex_label = _nice_species_name(field)
-                display_name = rf"\dot{{\omega}}\left[{tex_label}\right]"
+                sub = Substance(field)
+                display_name = rf"\dot{{\omega}}\left[{sub.to_tex()}\right]"
                 # Overwrite field to use nicer tex_label'ed display_name
                 self.add_output_field(
                     ("boxlib", field),
@@ -493,23 +501,70 @@ def setup_fluid_fields(self):
                     display_name=display_name,
                 )
                 self.alias(
-                    ("gas", f"{nice_name}_creation_rate"),
+                    ("gas", f"{sub}_creation_rate"),
                     ("boxlib", field),
                     units=unit_system["frequency"],
                 )
 
 
-def _nice_species_name(field):
-    spec_match = spec_finder.search(field)
-    nice_name = "".join(spec_match.groups())
-    # if the species field is a descriptive name, then the match
-    # on the integer will be blank
-    # modify the tex string in this case to remove spurious tex spacing
-    lab = r"X\left(^{%s}%s\right)"
-    if spec_match.groups()[-1] == "":
-        lab = r"X\left(%s%s\right)"
-    tex_label = lab % spec_match.groups()[::-1]
-    return nice_name, tex_label
+substance_expr_re = re.compile(r"\(([a-zA-Z][a-zA-Z0-9]*)\)")
+substance_elements_re = re.compile(r"(?P<element>[a-zA-Z]+)(?P<digits>\d*)")
+SubstanceSpec: TypeAlias = list[tuple[str, int]]
+
+
+class Substance:
+    def __init__(self, data: str) -> None:
+        if (m := substance_expr_re.search(data)) is None:
+            raise ValueError(f"{data!r} doesn't match expected regular expression")
+        sub_str = m.group()
+        constituents = substance_elements_re.findall(sub_str)
+
+        # 0 is used as a sentinel value to mark descriptive names
+        default_value = 1 if len(constituents) > 1 else 0
+        self._spec: SubstanceSpec = [
+            (name, int(count or default_value)) for (name, count) in constituents
+        ]
+
+    def get_spec(self) -> SubstanceSpec:
+        return self._spec.copy()
+
+    def is_isotope(self) -> bool:
+        return len(self._spec) == 1 and self._spec[0][1] > 0
+
+    def is_molecule(self) -> bool:
+        return len(self._spec) != 1
+
+    def is_descriptive_name(self) -> bool:
+        return len(self._spec) == 1 and self._spec[0][1] == 0
+
+    def __str__(self) -> str:
+        return "".join(
+            f"{element}{count if count > 1 else ''}" for element, count in self._spec
+        )
+
+    def _to_tex_isotope(self) -> str:
+        element, count = self._spec[0]
+        return rf"^{{{count}}}{element}"
+
+    def _to_tex_molecule(self) -> str:
+        return "".join(
+            rf"{element}_{{{count if count>1 else ''}}}"
+            for element, count in self._spec
+        )
+
+    def _to_tex_descriptive(self) -> str:
+        return str(self)
+
+    def to_tex(self) -> str:
+        if self.is_isotope():
+            return self._to_tex_isotope()
+        elif self.is_molecule():
+            return self._to_tex_molecule()
+        elif self.is_descriptive_name():
+            return self._to_tex_descriptive()
+        else:
+            # should only be reachable in case of a regular expression defect
+            raise RuntimeError
 
 
 def _create_density_func(field_name):

yt/frontends/amrex/tests/test_field_parsing.py

@@ -0,0 +1,51 @@
+import pytest
+
+from yt.frontends.amrex.fields import Substance
+
+
+@pytest.mark.parametrize(
+    "data, expected",
+    [
+        pytest.param("X(He5)", [("He", 5)], id="isotope_1"),
+        pytest.param("X(C12)", [("C", 12)], id="isotope_2"),
+        pytest.param("X(A1B2C3)", [("A", 1), ("B", 2), ("C", 3)], id="molecule_1"),
+        pytest.param("X(C12H24)", [("C", 12), ("H", 24)], id="molecule_2"),
+        pytest.param("X(H2O)", [("H", 2), ("O", 1)], id="molecule_3"),
+        pytest.param("X(ash)", [("ash", 0)], id="descriptive_name"),
+    ],
+)
+def test_Substance_spec(data, expected):
+    assert Substance(data)._spec == expected
+
+
+@pytest.mark.parametrize(
+    "data, expected_type",
+    [
+        pytest.param("X(He5)", "isotope", id="isotope_1"),
+        pytest.param("X(C12)", "isotope", id="isotope_2"),
+        pytest.param("X(A1B2C3)", "molecule", id="molecule_1"),
+        pytest.param("X(C12H24)", "molecule", id="molecule_2"),
+        pytest.param("X(H2O)", "molecule", id="molecule_3"),
+        pytest.param("X(ash)", "descriptive_name", id="descriptive_name"),
+    ],
+)
+def test_Substance_type(data, expected_type):
+    sub = Substance(data)
+    assert getattr(sub, f"is_{expected_type}")()
+
+
+@pytest.mark.parametrize(
+    "data, expected_str, expected_tex",
+    [
+        pytest.param("X(He5)", "He5", "^{5}He", id="isotope_1"),
+        pytest.param("X(C12)", "C12", "^{12}C", id="isotope_2"),
+        pytest.param("X(A1B2C3)", "AB2C3", "A_{}B_{2}C_{3}", id="molecule_1"),
+        pytest.param("X(C12H24)", "C12H24", "C_{12}H_{24}", id="molecule_2"),
+        pytest.param("X(H2O)", "H2O", "H_{2}O_{}", id="molecule_2"),
+        pytest.param("X(ash)", "ash", "ash", id="descriptive_name"),
+    ],
+)
+def test_Substance_to_str(data, expected_str, expected_tex):
+    sub = Substance(data)
+    assert str(sub) == expected_str
+    assert sub.to_tex() == expected_tex