diff --git a/yt/frontends/swift/data_structures.py b/yt/frontends/swift/data_structures.py index 5740edef679..c3d065d51a9 100644 --- a/yt/frontends/swift/data_structures.py +++ b/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,13 +128,19 @@ 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 @@ -134,26 +148,40 @@ def _parse_parameter_file(self): 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, " diff --git a/yt/frontends/swift/fields.py b/yt/frontends/swift/fields.py index 19ec84822c4..fecc9c6be73 100644 --- a/yt/frontends/swift/fields.py +++ b/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: + + - "_fraction_raw": the unmodified column value. + - "_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 "_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")) diff --git a/yt/frontends/swift/tests/test_outputs.py b/yt/frontends/swift/tests/test_outputs.py index d6dccb41e20..c78d450d781 100644 --- a/yt/frontends/swift/tests/test_outputs.py +++ b/yt/frontends/swift/tests/test_outputs.py @@ -8,6 +8,7 @@ keplerian_ring = "KeplerianRing/keplerian_ring_0020.hdf5" EAGLE_6 = "EAGLE_6/eagle_0005.hdf5" +stromgen_sphere = "Stromgren/output_7bin_HHeZdust_real_0060.hdf5" # Combined the tests for loading a file and ensuring the units have been @@ -111,3 +112,49 @@ def test_cosmo_dataset_selection(): ds = load(EAGLE_6) psc = ParticleSelectionComparison(ds) psc.run_defaults() + + +@requires_module("h5py") +@requires_file(stromgen_sphere) +def test_species_fractions(): + """ + Chemistry-network SWIFT snapshots expose a packed "SpeciesFractions" + field; SwiftFieldInfo unpack it into individual per-ion and + per-molecule "_fraction" (normalized) and + "_fraction_raw" (on-disk, unnormalized) fields. + """ + ds = load(stromgen_sphere) + if ("PartType0", "SpeciesFractions") not in ds.field_list: + return # this sample dataset has no chemistry network fields + + ad = ds.all_data() + + # normalized ion fraction always lies in [0, 1] + for name in ("HI", "HII"): + frac = ad["PartType0", f"{name}_fraction"] + assert frac.min() >= 0.0 + assert frac.max() <= 1.0 + + # the raw column value need not be bounded the same way, and for + # H (single-ion-state indices aside) should be present and finite + raw = ad["PartType0", "HII_fraction_raw"] + assert np.isfinite(raw).all() + + # HI + HII + Hm should renormalize to 1 (mask out zero-total particles) + hi = ad["PartType0", "HI_fraction"] + hii = ad["PartType0", "HII_fraction"] + hm = ad["PartType0", "Hm_fraction"] + total_raw = ( + ad["PartType0", "HI_fraction_raw"] + + ad["PartType0", "HII_fraction_raw"] + + ad["PartType0", "Hm_fraction_raw"] + ) + nonzero = total_raw > 0 + assert_almost_equal( + (hi + hii + hm)[nonzero], np.ones(nonzero.sum(), dtype="float64") + ) + + # a molecule's normalized fraction is just an alias for its raw value + assert_almost_equal( + ad["PartType0", "H2_fraction"].d, ad["PartType0", "H2_fraction_raw"].d + )