frontends/swift: add SpeciesFractions-derived fields and fix scalar conversion

yt/frontends/swift/data_structures.py

@@ -58,17 +58,25 @@ def _set_code_unit_attributes(self):
             msg = "Assuming length units are in comoving centimetres"
             only_on_root(mylog.info, msg)
             self.length_unit = self.quan(
-                float(units["Unit length in cgs (U_L)"]), "cmcm"
+                float(np.asarray(units["Unit length in cgs (U_L)"]).item()),
+                "cmcm",  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
             )
         else:
             msg = "Assuming length units are in physical centimetres"
             only_on_root(mylog.info, msg)
-            self.length_unit = self.quan(float(units["Unit length in cgs (U_L)"]), "cm")
-
-        self.mass_unit = self.quan(float(units["Unit mass in cgs (U_M)"]), "g")
-        self.time_unit = self.quan(float(units["Unit time in cgs (U_t)"]), "s")
+            self.length_unit = self.quan(
+                float(np.asarray(units["Unit length in cgs (U_L)"]).item()), "cm"
+            )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
+
+        self.mass_unit = self.quan(
+            float(np.asarray(units["Unit mass in cgs (U_M)"]).item()), "g"
+        )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
+        self.time_unit = self.quan(
+            float(np.asarray(units["Unit time in cgs (U_t)"]).item()), "s"
+        )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
         self.temperature_unit = self.quan(
-            float(units["Unit temperature in cgs (U_T)"]), "K"
+            float(np.asarray(units["Unit temperature in cgs (U_T)"]).item()),
+            "K",  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
         )
 
         return
@@ -120,40 +128,60 @@ def _parse_parameter_file(self):
         self.domain_right_edge = header["BoxSize"]
         self.domain_left_edge = np.zeros_like(self.domain_right_edge)
 
-        self.dimensionality = int(header["Dimension"])
+        self.dimensionality = int(
+            np.asarray(header["Dimension"]).item()
+        )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
 
         # SWIFT is either all periodic, or not periodic at all
         if has_runtime_pars:
-            periodic = int(runtime_parameters["PeriodicBoundariesOn"])
+            periodic = int(
+                np.asarray(runtime_parameters["PeriodicBoundariesOn"]).item()
+            )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
         else:
-            periodic = int(parameters["InitialConditions:periodic"])
+            periodic = int(
+                np.asarray(parameters["InitialConditions:periodic"]).item()
+            )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
 
         if periodic:
             self._periodicity = [True] * self.dimensionality
         else:
             self._periodicity = [False] * self.dimensionality
 
         # Units get attached to this
-        self.current_time = float(header["Time"])
+        self.current_time = float(
+            np.asarray(header["Time"]).item()
+        )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
 
         # Now cosmology enters the fray, as a runtime parameter.
-        self.cosmological_simulation = int(policy["cosmological integration"])
+        self.cosmological_simulation = int(
+            np.asarray(policy["cosmological integration"]).item()
+        )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
 
         if self.cosmological_simulation:
             try:
-                self.current_redshift = float(header["Redshift"])
+                self.current_redshift = float(
+                    np.asarray(header["Redshift"]).item()
+                )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
                 # These won't be present if self.cosmological_simulation is false
-                self.omega_lambda = float(parameters["Cosmology:Omega_lambda"])
+                self.omega_lambda = float(
+                    np.asarray(parameters["Cosmology:Omega_lambda"]).item()
+                )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
                 # Cosmology:Omega_m parameter deprecated at SWIFT commit d2783c2
                 # Between SWIFT versions 0.9.0 and 1.0.0
                 if "Cosmology:Omega_cdm" in parameters:
-                    self.omega_matter = float(parameters["Cosmology:Omega_b"]) + float(
-                        parameters["Cosmology:Omega_cdm"]
-                    )
+                    self.omega_matter = float(
+                        np.asarray(parameters["Cosmology:Omega_b"]).item()
+                    ) + float(
+                        np.asarray(parameters["Cosmology:Omega_cdm"]).item()
+                    )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
                 else:
-                    self.omega_matter = float(parameters["Cosmology:Omega_m"])
+                    self.omega_matter = float(
+                        np.asarray(parameters["Cosmology:Omega_m"]).item()
+                    )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
                 # This is "little h"
-                self.hubble_constant = float(parameters["Cosmology:h"])
+                self.hubble_constant = float(
+                    np.asarray(parameters["Cosmology:h"]).item()
+                )  ## Bug fixed for conversion of an array with ndim > 0 to a scalar
             except KeyError:
                 mylog.warning(
                     "Could not find cosmology information in Parameters, "

yt/frontends/swift/fields.py

@@ -1,5 +1,187 @@
+from functools import partial
+
 from yt.frontends.sph.fields import SPHFieldInfo
 
+# Column index of each species inside the packed SWIFT "SpeciesFractions"
+# particle field, as produced by the SWIFT chemistry/cooling network.
+# Elements map to a list of (ion_name, column_index); molecules map
+# straight to a column index.
+ELEMENT_MAP = {
+    "H": [("HI", 1), ("HII", 2), ("Hm", 3)],
+    "He": [("HeI", 4), ("HeII", 5), ("HeIII", 6)],
+    "C": [
+        ("CI", 7),
+        ("CII", 8),
+        ("CIII", 9),
+        ("CIV", 10),
+        ("CV", 11),
+        ("CVI", 12),
+        ("CVII", 13),
+        ("Cm", 14),
+    ],
+    "N": [
+        ("NI", 15),
+        ("NII", 16),
+        ("NIII", 17),
+        ("NIV", 18),
+        ("NV", 19),
+        ("NVI", 20),
+        ("NVII", 21),
+        ("NVIII", 22),
+    ],
+    "O": [
+        ("OI", 23),
+        ("OII", 24),
+        ("OIII", 25),
+        ("OIV", 26),
+        ("OV", 27),
+        ("OVI", 28),
+        ("OVII", 29),
+        ("OVIII", 30),
+        ("OIX", 31),
+        ("Om", 32),
+    ],
+    "Ne": [
+        ("NeI", 33),
+        ("NeII", 34),
+        ("NeIII", 35),
+        ("NeIV", 36),
+        ("NeV", 37),
+        ("NeVI", 38),
+        ("NeVII", 39),
+        ("NeVIII", 40),
+        ("NeIX", 41),
+        ("NeX", 42),
+        ("NeXI", 43),
+    ],
+    "Mg": [
+        ("MgI", 44),
+        ("MgII", 45),
+        ("MgIII", 46),
+        ("MgIV", 47),
+        ("MgV", 48),
+        ("MgVI", 49),
+        ("MgVII", 50),
+        ("MgVIII", 51),
+        ("MgIX", 52),
+        ("MgX", 53),
+        ("MgXI", 54),
+        ("MgXII", 55),
+        ("MgXIII", 56),
+    ],
+    "Si": [
+        ("SiI", 57),
+        ("SiII", 58),
+        ("SiIII", 59),
+        ("SiIV", 60),
+        ("SiV", 61),
+        ("SiVI", 62),
+        ("SiVII", 63),
+        ("SiVIII", 64),
+        ("SiIX", 65),
+        ("SiX", 66),
+        ("SiXI", 67),
+        ("SiXII", 68),
+        ("SiXIII", 69),
+        ("SiXIV", 70),
+        ("SiXV", 71),
+    ],
+    "S": [
+        ("SI", 72),
+        ("SII", 73),
+        ("SIII", 74),
+        ("SIV", 75),
+        ("SV", 76),
+        ("SVI", 77),
+        ("SVII", 78),
+        ("SVIII", 79),
+        ("SIX", 80),
+        ("SX", 81),
+        ("SXI", 82),
+        ("SXII", 83),
+        ("SXIII", 84),
+        ("SXIV", 85),
+        ("SXV", 86),
+        ("SXVI", 87),
+        ("SXVII", 88),
+    ],
+    "Ca": [
+        ("CaI", 89),
+        ("CaII", 90),
+        ("CaIII", 91),
+        ("CaIV", 92),
+        ("CaV", 93),
+        ("CaVI", 94),
+        ("CaVII", 95),
+        ("CaVIII", 96),
+        ("CaIX", 97),
+        ("CaX", 98),
+        ("CaXI", 99),
+        ("CaXII", 100),
+        ("CaXIII", 101),
+        ("CaXIV", 102),
+        ("CaXV", 103),
+        ("CaXVI", 104),
+        ("CaXVII", 105),
+        ("CaXVIII", 106),
+        ("CaXIX", 107),
+        ("CaXX", 108),
+        ("CaXXI", 109),
+    ],
+    "Fe": [
+        ("FeI", 110),
+        ("FeII", 111),
+        ("FeIII", 112),
+        ("FeIV", 113),
+        ("FeV", 114),
+        ("FeVI", 115),
+        ("FeVII", 116),
+        ("FeVIII", 117),
+        ("FeIX", 118),
+        ("FeX", 119),
+        ("FeXI", 120),
+        ("FeXII", 121),
+        ("FeXIII", 122),
+        ("FeXIV", 123),
+        ("FeXV", 124),
+        ("FeXVI", 125),
+        ("FeXVII", 126),
+        ("FeXVIII", 127),
+        ("FeXIX", 128),
+        ("FeXX", 129),
+        ("FeXXI", 130),
+        ("FeXXII", 131),
+        ("FeXXIII", 132),
+        ("FeXXIV", 133),
+        ("FeXXV", 134),
+        ("FeXXVI", 135),
+        ("FeXXVII", 136),
+    ],
+}
+
+MOLECULE_MAP = {
+    "H2": 137,
+    "H2p": 138,
+    "H3p": 139,
+    "OH": 140,
+    "H2O": 141,
+    "C2": 142,
+    "O2": 143,
+    "HCOp": 144,
+    "CH": 145,
+    "CH2": 146,
+    "CH3p": 147,
+    "CO": 148,
+    "CHp": 149,
+    "CH2p": 150,
+    "OHp": 151,
+    "H2Op": 152,
+    "H3Op": 153,
+    "COp": 154,
+    "HOCp": 155,
+    "O2p": 156,
+}
+
 
 class SwiftFieldInfo(SPHFieldInfo):
     def __init__(self, ds, field_list, slice_info=None):
@@ -18,6 +200,66 @@ def setup_particle_fields(self, ptype, *args, **kwargs):
 
         if ptype in ("PartType0", "Gas"):
             self.setup_gas_particle_fields(ptype)
+            if (ptype, "SpeciesFractions") in self.ds.field_list:
+                self.species_names = self._setup_species_fractions(ptype)
+
+    def _setup_species_fractions(self, ptype):
+        """
+        Splits the packed "SpeciesFractions" particle field (one column per
+        ionization state / molecule, as produced by SWIFT's chemistry
+        network) into two fields per ion/molecule:
+
+        - "<species>_fraction_raw": the unmodified column value.
+        - "<species>_fraction": for ions, this is renormalized so it sums
+          to 1 across all ionization states of the parent element (i.e.
+          "fraction of carbon atoms that are C IV"). For molecules, which
+          have no parent-element group to normalize against, this is
+          identical to "<species>_fraction_raw".
+        """
+
+        def _raw_value(field, data, idx: int):
+            return data[ptype, "SpeciesFractions"][:, idx]
+
+        def _ion_fraction(field, data, idx: int, indices: list):
+            all_species = data[ptype, "SpeciesFractions"]
+            num = all_species[:, idx]
+            total = all_species[:, indices].sum(axis=1)
+            out = num.copy()
+            mask = total > 0
+            out[mask] = num[mask] / total[mask]
+            out[~mask] = 0
+            return out
+
+        species_names = []
+
+        for _element, ions in ELEMENT_MAP.items():
+            indices = [idx for _, idx in ions]
+            for name, idx in ions:
+                self.add_field(
+                    (ptype, f"{name}_fraction_raw"),
+                    sampling_type="particle",
+                    function=partial(_raw_value, idx=idx),
+                    units="",
+                )
+                self.add_field(
+                    (ptype, f"{name}_fraction"),
+                    sampling_type="particle",
+                    function=partial(_ion_fraction, idx=idx, indices=indices),
+                    units="",
+                )
+                species_names.append(name)
+
+        for name, idx in MOLECULE_MAP.items():
+            self.add_field(
+                (ptype, f"{name}_fraction_raw"),
+                sampling_type="particle",
+                function=partial(_raw_value, idx=idx),
+                units="",
+            )
+            self.alias((ptype, f"{name}_fraction"), (ptype, f"{name}_fraction_raw"))
+            species_names.append(name)
+
+        return species_names
 
     def setup_gas_particle_fields(self, ptype):
         self.alias((ptype, "temperature"), (ptype, "Temperatures"))