Refactor species field definition (#175)

Co-authored-by: google-labs-jules[bot] <161369871+google-labs-jules[bot]@users.noreply.github.com>

Author: henry2004y

src/flekspy/yt/yt.py

@@ -24,8 +24,7 @@ class FLEKSFieldInfo(FieldInfoContainer):
     rho_units = "code_density"
     mass_units = "code_mass"
 
-    # TODO: find a way to avoid repeating s0, s1...
-    known_other_fields = (
+    _base_fields = (
         ("Bx", (b_units, ["magnetic_field_x"], r"B_x")),
         ("By", (b_units, ["magnetic_field_y"], r"B_y")),
         ("Bz", (b_units, ["magnetic_field_z"], r"B_z")),
@@ -35,28 +34,23 @@ class FLEKSFieldInfo(FieldInfoContainer):
         ("X", (l_units, [], r"X")),
         ("Y", (l_units, [], r"Y")),
         ("Z", (l_units, [], r"Z")),
-        ("rhos0", (rho_units, [], r"\rho")),
-        ("uxs0", (v_units, [], r"u_x")),
-        ("uys0", (v_units, [], r"u_y")),
-        ("uzs0", (v_units, [], r"u_z")),
-        ("pxxs0", (p_units, [], r"P_{xx}")),
-        ("pyys0", (p_units, [], r"P_{yy}")),
-        ("pzzs0", (p_units, [], r"P_{zz}")),
-        ("pxys0", (p_units, [], r"P_{xy}")),
-        ("pxzs0", (p_units, [], r"P_{xz}")),
-        ("pyzs0", (p_units, [], r"P_{yz}")),
-        ("rhos1", (rho_units, [], r"\rho")),
-        ("uxs1", (v_units, [], r"u_x")),
-        ("uys1", (v_units, [], r"u_y")),
-        ("uzs1", (v_units, [], r"u_z")),
-        ("pxxs1", (p_units, [], r"P_{xx}")),
-        ("pyys1", (p_units, [], r"P_{yy}")),
-        ("pzzs1", (p_units, [], r"P_{zz}")),
-        ("pxys1", (p_units, [], r"P_{xy}")),
-        ("pxzs1", (p_units, [], r"P_{xz}")),
-        ("pyzs1", (p_units, [], r"P_{yz}")),
     )
 
+    _species_fields_template = (
+        ("rho", (rho_units, [], r"\rho")),
+        ("ux", (v_units, [], r"u_x")),
+        ("uy", (v_units, [], r"u_y")),
+        ("uz", (v_units, [], r"u_z")),
+        ("pxx", (p_units, [], r"P_{xx}")),
+        ("pyy", (p_units, [], r"P_{yy}")),
+        ("pzz", (p_units, [], r"P_{zz}")),
+        ("pxy", (p_units, [], r"P_{xy}")),
+        ("pxz", (p_units, [], r"P_{xz}")),
+        ("pyz", (p_units, [], r"P_{yz}")),
+    )
+
+    known_other_fields = _base_fields
+
     known_particle_fields = (
         ("particle_weight", (mass_units, ["p_w"], r"weight")),
         ("particle_position_x", (l_units, ["p_x"], "x")),
@@ -78,11 +72,57 @@ def __init__(self, ds, field_list):
             finfo.nodal_flag = ds.nodal_flags[field]
 
     def setup_fluid_fields(self):
+        import re
+
         from yt.fields.magnetic_field import setup_magnetic_field_aliases
 
         for field in self.known_other_fields:
             fname = field[0]
-            self.alias(("mesh", fname), ("boxlib", fname))
+            # Try to alias to boxlib first, then raw
+            original_name = ("boxlib", fname)
+            if original_name not in self:
+                original_name = ("raw", fname)
+
+            self.alias(("mesh", fname), original_name)
+
+        # Dynamically alias species fields from the input dataset
+        # Use field_list instead of index.raw_fields to avoid forcing index construction
+        species_regex = re.compile(r"^(?P<base>.*)s(?P<idx>\d+)$")
+
+        # field_list typically contains tuples like ('boxlib', 'varname') or ('raw', 'varname')
+        # We need to scan for on-disk fields.
+        # Use self.field_list instead of self.ds.field_list to be robust against incomplete ds mocks in tests
+        for ftype, fname in self.field_list:
+            # We are interested in fields that might be raw fluid fields
+            if ftype not in ("boxlib", "raw"):
+                continue
+
+            match = species_regex.match(fname)
+            if match:
+                base = match.group("base")
+                for template_field, props in self._species_fields_template:
+                    if base == template_field:
+                        # Construct proper LaTeX display name merging subscripts
+                        base_display = props[2]
+                        suffix = f"s{match.group('idx')}"
+                        if "_" in base_display:
+                            name_part, sub_part = base_display.split("_", 1)
+                            if sub_part.startswith("{") and sub_part.endswith("}"):
+                                sub_content = sub_part[1:-1]
+                                display_name = f"{name_part}_{{{sub_content},{suffix}}}"
+                            else:
+                                display_name = f"{name_part}_{{{sub_part},{suffix}}}"
+                        else:
+                            display_name = f"{base_display}_{{{suffix}}}"
+
+                        # Alias ("mesh", fname) to the found field
+                        self.alias(
+                            ("mesh", fname), (ftype, fname), units=props[0]
+                        )
+
+                        if ("mesh", fname) in self:
+                            self[("mesh", fname)].display_name = display_name
+                        break
 
         # This function is required by yt to correctly handle magnetic field
         # units and set up aliases. See:

tests/test_dynamic_species.py

@@ -0,0 +1,60 @@
+
+import unittest
+from unittest.mock import MagicMock
+from flekspy.yt.yt import FLEKSFieldInfo
+
+class TestDynamicFields(unittest.TestCase):
+    def test_dynamic_species_registration(self):
+        import yt.geometry.geometry_enum as geometry_enum
+
+        # Mock the dataset
+        mock_ds = MagicMock()
+        # Set geometry to CARTESIAN to avoid assert_never(geometry) in yt
+        mock_ds.geometry = geometry_enum.Geometry.CARTESIAN
+
+        # Create a fake list of raw fields that includes a high index species
+        mock_ds.index.raw_fields = ["Bx", "By", "Bz", "rhos0", "rhos99", "uxs99"]
+        mock_ds.nodal_flags = {f: 0 for f in mock_ds.index.raw_fields}
+
+        # Configure field_units on the mock dataset to return string units
+        mock_ds.field_units.get.return_value = "code_length" # Default for unknown
+
+        # Initialize FieldInfo
+        # Pass the raw fields as field_list to simulate what yt does
+        field_list = [("raw", f) for f in mock_ds.index.raw_fields]
+        fi = FLEKSFieldInfo(mock_ds, field_list)
+
+        # Setup fluid fields
+        fi.setup_fluid_fields()
+
+        # Verify that dynamic aliasing logic was executed and created expected fields
+        # Note: in mocked environment without a real index hierarchy,
+        # 'alias' might not behave exactly as in real yt unless perfectly mocked.
+        # But we can check if self.alias was called or if the fields are in 'fi'.
+
+        # If alias failed because original_name not in self (which alias() checks),
+        # we need to make sure original_name IS in self.
+
+        # Since we passed field_list to __init__, ('raw', 'rhos0') etc should be in fi.
+        self.assertIn(("raw", "rhos0"), fi)
+        self.assertIn(("raw", "rhos99"), fi)
+
+        # Check if mesh aliases exist
+        self.assertIn(("mesh", "rhos0"), fi)
+        self.assertEqual(fi[("mesh", "rhos0")].units, "code_density")
+
+        self.assertIn(("mesh", "rhos99"), fi)
+        self.assertEqual(fi[("mesh", "rhos99")].units, "code_density")
+        self.assertIn(r"\rho_{s99}", fi[("mesh", "rhos99")].display_name)
+
+        self.assertIn(("mesh", "uxs99"), fi)
+        self.assertEqual(fi[("mesh", "uxs99")].units, "code_velocity")
+        # u_x + s99 -> u_{x,s99}
+        self.assertIn(r"u_{x,s99}", fi[("mesh", "uxs99")].display_name)
+
+        # Check that unknown patterns are not registered blindly (though regular fields loop handles aliases for known_other_fields)
+        # Bx is in known_other_fields
+        self.assertIn(("mesh", "Bx"), fi)
+
+if __name__ == "__main__":
+    unittest.main()