From a337e73e2a7451f39df1a904991171fbc6e58903 Mon Sep 17 00:00:00 2001 From: Evelyn Date: Tue, 25 May 2021 15:58:54 -0400 Subject: [PATCH 1/5] Allow non-solar abundances to be specified reading_func_args accessed by SpectrumGenerator changed some variable names bc I thought trident was confused but alas it is me who is confused (cue billy joel frog) Change how we pass abundances heyo Propogated changes to spec gen Added abundance table args to ray generators No method has_key :( abundance_store not being set; attempt fix abundance_store not being set; attempt fix GLOBAL VARIABLES ARE DUMB Let's be explicit with all our globals beginning of final modifications got ri of print statements and commented out parts of reading function I probably do not need anymore some_rad_science is only commented out right now I think she is finally done --- trident/__init__.py | 3 +- trident/ion_balance.py | 126 ++++++++++++++++++++++++++++------ trident/light_ray.py | 6 +- trident/ray_generator.py | 21 +++++- trident/spectrum_generator.py | 32 +++++++-- 5 files changed, 157 insertions(+), 31 deletions(-) diff --git a/trident/__init__.py b/trident/__init__.py index c4a998956..df77d2ea3 100644 --- a/trident/__init__.py +++ b/trident/__init__.py @@ -37,7 +37,8 @@ add_ion_mass_field, \ solar_abundance, \ atomic_mass, \ - calculate_ion_fraction + calculate_ion_fraction, \ + update_abundances from trident.instrument import \ Instrument diff --git a/trident/ion_balance.py b/trident/ion_balance.py index 44b760aaf..28149e715 100644 --- a/trident/ion_balance.py +++ b/trident/ion_balance.py @@ -11,6 +11,7 @@ # The full license is in the file LICENSE, distributed with this software. #----------------------------------------------------------------------------- +from sympy import abundance from yt.fields.field_detector import \ FieldDetector from yt.utilities.linear_interpolators import \ @@ -40,7 +41,9 @@ fraction_zero_point = 1.e-9 zero_out_value = -30. -table_store = {} +# Global variables for altering elemental abundance & ionization fractions +abundance_store = {} +ion_table_store = {} class IonBalanceTable(object): def __init__(self, filename=None, atom=None): @@ -133,6 +136,7 @@ def _log_T(field, data): def add_ion_fields(ds, ions, ftype='gas', ionization_table=None, + abundance_dict=None, field_suffix=False, line_database=None, sampling_type='local', @@ -180,16 +184,16 @@ def add_ion_fields(ds, ions, ftype='gas', :ions: list of strings - List of strings matching possible lines. Strings can be of the - form: - * Atom - Examples: "H", "C", "Mg" - * Ion - Examples: "H I", "H II", "C IV", "Mg II" - * Line - Examples: "H I 1216", "C II 1336", "Mg II 1240" + List of strings matching possible lines. Strings can be of the + form: + * Atom - Examples: "H", "C", "Mg" + * Ion - Examples: "H I", "H II", "C IV", "Mg II" + * Line - Examples: "H I 1216", "C II 1336", "Mg II 1240" - If set to 'all', creates **all** ions for the first 30 elements: - (ie hydrogen to zinc). If set to 'all' with ``line_database`` - keyword set, then creates **all** ions associated with the lines - specified in the equivalent :class:`~trident.LineDatabase`. + If set to 'all', creates **all** ions for the first 30 elements: + (ie hydrogen to zinc). If set to 'all' with ``line_database`` + keyword set, then creates **all** ions associated with the lines + specified in the equivalent :class:`~trident.LineDatabase`. :ionization_table: string, optional @@ -199,6 +203,14 @@ def add_ion_fields(ds, ions, ftype='gas', specified in ~/.trident/config Default: None + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + :field_suffix: boolean, optional Determines whether or not to append a suffix to the field name that @@ -277,10 +289,13 @@ def add_ion_fields(ds, ions, ftype='gas', # - X_P#_density for (atom, ion) in ion_list: add_ion_mass_field(atom, ion, ds, ftype, ionization_table, - field_suffix=field_suffix, sampling_type=sampling_type) + abundance_dict=abundance_dict, + field_suffix=field_suffix, + sampling_type=sampling_type) def add_ion_fraction_field(atom, ion, ds, ftype="gas", ionization_table=None, + abundance_dict=None, field_suffix=False, sampling_type='local', particle_type=None): @@ -324,6 +339,14 @@ def add_ion_fraction_field(atom, ion, ds, ftype="gas", metallicity, and redshift. By default, it uses the table specified in ~/.trident/config + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + :field_suffix: boolean, optional Determines whether or not to append a suffix to the field name that indicates what ionization table was used @@ -371,9 +394,10 @@ def add_ion_fraction_field(atom, ion, ds, ftype="gas", if field_suffix: field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] - if field not in table_store: + global ion_table_store + if field not in ion_table_store: ionTable = IonBalanceTable(ionization_table, atom) - table_store[field] = {'fraction': copy.deepcopy(ionTable.ion_fraction[ion-1]), + ion_table_store[field] = {'fraction': copy.deepcopy(ionTable.ion_fraction[ion-1]), 'parameters': copy.deepcopy(ionTable.parameters)} del ionTable @@ -389,6 +413,7 @@ def add_ion_fraction_field(atom, ion, ds, ftype="gas", def add_ion_number_density_field(atom, ion, ds, ftype="gas", ionization_table=None, + abundance_dict=None, field_suffix=False, sampling_type='local', particle_type=None): @@ -436,6 +461,14 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", metallicity, and redshift. By default, it uses the table specified in ~/.trident/config + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + :field_suffix: boolean, optional Determines whether or not to append a suffix to the field @@ -464,6 +497,13 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", if ionization_table is None: ionization_table = ion_table_filepath + + global abundance_store + if abundance_dict is None: + abundance_store = copy.copy(solar_abundance) + else: + abundance_store = update_abundances(abundance_dict) + atom = atom.capitalize() field = "%s_p%d_number_density" % (atom, ion-1) @@ -472,6 +512,7 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] add_ion_fraction_field(atom, ion, ds, ftype, ionization_table, + abundance_dict=abundance_dict, field_suffix=field_suffix, sampling_type=sampling_type) @@ -480,6 +521,7 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", def add_ion_density_field(atom, ion, ds, ftype="gas", ionization_table=None, + abundance_dict=None, field_suffix=False, sampling_type='local', particle_type=None): @@ -527,6 +569,14 @@ def add_ion_density_field(atom, ion, ds, ftype="gas", metallicity, and redshift. By default, it uses the table specified in ~/.trident/config + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + :field_suffix: boolean, optional Determines whether or not to append a suffix to the field @@ -555,6 +605,7 @@ def add_ion_density_field(atom, ion, ds, ftype="gas", if ionization_table is None: ionization_table = ion_table_filepath + atom = atom.capitalize() field = "%s_p%d_density" % (atom, ion-1) @@ -563,6 +614,8 @@ def add_ion_density_field(atom, ion, ds, ftype="gas", field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] add_ion_number_density_field(atom, ion, ds, ftype, ionization_table, + abundance_dict=abundance_dict, + field_suffix=field_suffix, sampling_type=sampling_type) _add_field(ds, ("gas", field), function=_ion_density, @@ -570,6 +623,7 @@ def add_ion_density_field(atom, ion, ds, ftype="gas", def add_ion_mass_field(atom, ion, ds, ftype="gas", ionization_table=None, + abundance_dict=None, field_suffix=False, sampling_type='local', particle_type=None): @@ -618,6 +672,14 @@ def add_ion_mass_field(atom, ion, ds, ftype="gas", metallicity, and redshift. By default, it uses the table specified in ~/.trident/config + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + :field_suffix: boolean, optional Determines whether or not to append a suffix to the field @@ -646,6 +708,7 @@ def add_ion_mass_field(atom, ion, ds, ftype="gas", if ionization_table is None: ionization_table = ion_table_filepath + atom = atom.capitalize() field = "%s_p%s_mass" % (atom, ion-1) @@ -654,6 +717,7 @@ def add_ion_mass_field(atom, ion, ds, ftype="gas", field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] add_ion_density_field(atom, ion, ds, ftype, ionization_table, + abundance_dict=abundance_dict, field_suffix=field_suffix, sampling_type=sampling_type) @@ -770,11 +834,12 @@ def _ion_number_density(field, data): atomic_mass[atom] / mh if atom == 'H' or atom == 'He': - number_density = solar_abundance[atom] * data[ftype, fraction_field_name] + number_density = abundance_store[atom] * data[fraction_field_name] else: - number_density = data.ds.quan(solar_abundance[atom], "1.0/Zsun") * \ + number_density = data.ds.quan(abundance_store[atom], "1.0/Zsun") * \ data[ftype, fraction_field_name] * \ data[ftype, "metallicity"] + # convert to number density # use the on disk hydrogen number density if possible if (ftype, "H_nuclei_density") in data.ds.derived_field_list: @@ -789,26 +854,27 @@ def _ion_fraction_field(field, data): of an ion over a dataset by plugging in the density, temperature, metallicity and redshift of the output into the ionization table. """ + if isinstance(field.name, tuple): ftype = field.name[0] field_name = field.name[1] else: ftype = "gas" field_name = field.name - n_parameters = len(table_store[field_name]['parameters']) + n_parameters = len(ion_table_store[field_name]['parameters']) if n_parameters == 1: - ionFraction = table_store[field_name]['fraction'] - t_param = table_store[field_name]['parameters'][0] + ionFraction = ion_table_store[field_name]['fraction'] + t_param = ion_table_store[field_name]['parameters'][0] bds = t_param.astype("=f8") interp = UnilinearFieldInterpolator(ionFraction, bds, 'log_T', truncate=True) elif n_parameters == 3: - ionFraction = table_store[field_name]['fraction'] - n_param = table_store[field_name]['parameters'][0] - z_param = table_store[field_name]['parameters'][1] - t_param = table_store[field_name]['parameters'][2] + ionFraction = ion_table_store[field_name]['fraction'] + n_param = ion_table_store[field_name]['parameters'][0] + z_param = ion_table_store[field_name]['parameters'][1] + t_param = ion_table_store[field_name]['parameters'][2] bds = [n_param.astype("=f8"), z_param.astype("=f8"), t_param.astype("=f8")] interp = TrilinearFieldInterpolator(ionFraction, bds, @@ -971,6 +1037,21 @@ def calculate_ion_fraction(ion, density, temperature, redshift, ionization_table fraction = np.clip(fraction, 0.0, 1.0) return fraction +def update_abundances(abundance_replacements): + + # Start with solar abundances as a "base" + abundances = solar_abundance.copy() + + # Modify the provided elements. Could be all of them! + # Validate element keys using existing solar_abundance dict + for key, val in abundance_replacements.items(): + if key in solar_abundance: + abundances[key] = val + else: + print(f"PROBLEM KEY: {key}") + raise RuntimeError(f"Unrecognized element {key} provided to abundance_dict. Only elements up through Zn supported.") + + return abundances # Taken from Cloudy documentation. @@ -986,6 +1067,7 @@ def calculate_ion_fraction(ion, density, temperature, redshift, ionization_table 'Mn': 2.88e-07, 'Fe': 2.82e-05, 'Co': 8.32e-08, 'Ni': 1.78e-06, 'Cu': 1.62e-08, 'Zn': 3.98e-08} + atomic_mass = { 'H' : 1.00794, 'He': 4.002602, 'Li': 6.941, 'Be': 9.012182, 'B' : 10.811, 'C' : 12.0107, diff --git a/trident/light_ray.py b/trident/light_ray.py index f6ab5d696..a5db3ee87 100644 --- a/trident/light_ray.py +++ b/trident/light_ray.py @@ -1,3 +1,4 @@ + """ LightRay class and member functions. @@ -50,7 +51,7 @@ class LightRay(CosmologySplice): For compound rays, the LightRay stacks together multiple datasets in a time series in order to approximate a LightRay's path through a volume and redshift interval larger than a single simulation data output. - The outcome is something akin to a synthetic QSO line of sight. + The outcome is thing akin to a synthetic QSO line of sight. Once the LightRay object is set up, use LightRay.make_light_ray to begin making rays. Different randomizations can be created with a @@ -521,9 +522,10 @@ def make_light_ray(self, seed=None, periodic=True, if field_parameters is None: field_parameters = {} + else: # Initialize data structures. - self._data = {} + self._data = {} # temperature field is automatically added to fields if fields is None: fields = [] if ('gas', 'temperature') not in fields: diff --git a/trident/ray_generator.py b/trident/ray_generator.py index c838658a9..c9632cfe9 100644 --- a/trident/ray_generator.py +++ b/trident/ray_generator.py @@ -32,7 +32,7 @@ def make_simple_ray(dataset_file, start_position, end_position, solution_filename=None, data_filename=None, trajectory=None, redshift=None, field_parameters=None, setup_function=None, load_kwargs=None, - line_database=None, ionization_table=None, + line_database=None, ionization_table=None, abundance_dict=None, fail_empty=True): """ Create a yt LightRay object for a single dataset (eg CGM). This is a @@ -170,6 +170,15 @@ def make_simple_ray(dataset_file, start_position, end_position, it uses the table specified in ~/.trident/config Default: None + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + Default: None + :fail_empty: optional, bool If True, Trident will fail when it tries to create an empty Ray @@ -251,6 +260,7 @@ def make_compound_ray(parameter_filename, simulation_type, find_outputs=False, seed=None, setup_function=None, load_kwargs=None, line_database=None, ionization_table=None, + abundance_dict=None, field_parameters = None, fail_empty=True): """ @@ -442,6 +452,15 @@ def make_compound_ray(parameter_filename, simulation_type, it uses the table specified in ~/.trident/config Default: None + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + Default: None + :field_parameters: optional, dict Used to set field parameters in light rays. For example, if the 'bulk_velocity' field parameter is set, the relative diff --git a/trident/spectrum_generator.py b/trident/spectrum_generator.py index 90a9e6a42..1034ea1bc 100644 --- a/trident/spectrum_generator.py +++ b/trident/spectrum_generator.py @@ -163,6 +163,15 @@ class SpectrumGenerator(AbsorptionSpectrum): file. Default: None + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + Default: None + **Example** Create a one-zone ray, and generate a COS spectrum from that ray. @@ -192,7 +201,7 @@ class SpectrumGenerator(AbsorptionSpectrum): def __init__(self, instrument=None, lambda_min=None, lambda_max=None, n_lambda=None, dlambda=None, lsf_kernel=None, line_database='lines.txt', ionization_table=None, - bin_space='wavelength'): + abundance_dict=None, bin_space='wavelength'): if instrument is None and \ ((lambda_min is None or lambda_max is None) or \ (dlambda is None and n_lambda is None)): @@ -241,6 +250,8 @@ def __init__(self, instrument=None, lambda_min=None, lambda_max=None, else: self.ionization_table = None + self.abundance_dict = abundance_dict + def make_spectrum(self, ray, lines='all', output_file=None, output_absorbers_file=None, @@ -360,6 +371,7 @@ def make_spectrum(self, ray, lines='all', >>> sg.plot_spectrum('spec_raw.png') """ self.observing_redshift = observing_redshift + if isinstance(ray, str): ray = load(ray) @@ -400,7 +412,8 @@ def make_spectrum(self, ray, lines='all', my_lev = int(on_ion[1][1:]) + 1 mylog.info("Creating %s from ray's fields." % (line.field[1])) add_ion_number_density_field(on_ion[0], my_lev, ray, - ionization_table=self.ionization_table) + ionization_table=self.ionization_table, + abundance_dict=self.abundance_dict) self.add_line(line.identifier, line.field, float(line.wavelength), @@ -419,7 +432,6 @@ def make_spectrum(self, ray, lines='all', element='H', ion_state='I') if len(H_lines) > 0 and ly_continuum: self.add_continuum('Ly C', H_lines[0].field, 912.32336, 1.6e17, 3.0) - AbsorptionSpectrum.make_spectrum(self, ad, output_file=None, line_list_file=None, @@ -1044,7 +1056,7 @@ def __repr__(self): return disp def load_spectrum(filename, format='auto', instrument=None, lsf_kernel=None, - line_database='lines.txt', ionization_table=None): + line_database='lines.txt', ionization_table=None, abundance_dict=None): """ Load a previously saved spectrum from disk. @@ -1085,6 +1097,15 @@ def load_spectrum(filename, format='auto', instrument=None, lsf_kernel=None, based on its density, temperature, metallicity, and redshift. Default: None + :abundance_dict: dictionary, optional + + Dictionary of elemental abundances normalized to hydrogen. Keys should + be elemental symbols, e.g., 'He'. By default, Trident assumes the solar + abundances of REF. Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, specify + the dictionary should include all elements up through zinc. + Default: None + **Example** Create a simple spectrum, save it to disk, and load it back as a new @@ -1132,7 +1153,8 @@ def load_spectrum(filename, format='auto', instrument=None, lsf_kernel=None, sg = SpectrumGenerator(instrument=instrument, lambda_min=lambda_min, lambda_max=lambda_max, n_lambda=n_lambda, lsf_kernel=lsf_kernel, line_database=line_database, - ionization_table=ionization_table) + ionization_table=ionization_table, + abundance_dict=abundance_dict) if tau_field is not None: sg.load_spectrum(lambda_field=lambda_field, tau_field=tau_field, flux_field=flux_field) From 26af0a06d599d8456ccf0add807f68eda0b7723c Mon Sep 17 00:00:00 2001 From: Claire Kopenhafer Date: Wed, 3 Jun 2026 11:21:40 -0400 Subject: [PATCH 2/5] Cleanup rebase of abundance table additions --- trident/ion_balance.py | 19 ++++++++----------- trident/light_ray.py | 6 ++---- trident/spectrum_generator.py | 2 +- 3 files changed, 11 insertions(+), 16 deletions(-) diff --git a/trident/ion_balance.py b/trident/ion_balance.py index 28149e715..2fb8e5e67 100644 --- a/trident/ion_balance.py +++ b/trident/ion_balance.py @@ -11,7 +11,6 @@ # The full license is in the file LICENSE, distributed with this software. #----------------------------------------------------------------------------- -from sympy import abundance from yt.fields.field_detector import \ FieldDetector from yt.utilities.linear_interpolators import \ @@ -500,7 +499,7 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", global abundance_store if abundance_dict is None: - abundance_store = copy.copy(solar_abundance) + abundance_store = solar_abundance else: abundance_store = update_abundances(abundance_dict) @@ -834,7 +833,7 @@ def _ion_number_density(field, data): atomic_mass[atom] / mh if atom == 'H' or atom == 'He': - number_density = abundance_store[atom] * data[fraction_field_name] + number_density = abundance_store[atom] * data[ftype, fraction_field_name] else: number_density = data.ds.quan(abundance_store[atom], "1.0/Zsun") * \ data[ftype, fraction_field_name] * \ @@ -1006,16 +1005,16 @@ def calculate_ion_fraction(ion, density, temperature, redshift, ionization_table field = "%s_p%d_ion_fraction" % (atom, ion_state-1) field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] - if field not in table_store: + if field not in ion_table_store: ionTable = IonBalanceTable(ionization_table, atom) - table_store[field] = {'fraction': copy.deepcopy(ionTable.ion_fraction[ion_state-1]), + ion_table_store[field] = {'fraction': copy.deepcopy(ionTable.ion_fraction[ion_state-1]), 'parameters': copy.deepcopy(ionTable.parameters)} del ionTable - ion_fraction = table_store[field]['fraction'] - n_param = table_store[field]['parameters'][0] - z_param = table_store[field]['parameters'][1] - t_param = table_store[field]['parameters'][2] + ion_fraction = ion_table_store[field]['fraction'] + n_param = ion_table_store[field]['parameters'][0] + z_param = ion_table_store[field]['parameters'][1] + t_param = ion_table_store[field]['parameters'][2] # x,y,z coordinates for all ion_fractions from table coords = (n_param, z_param, t_param) @@ -1048,7 +1047,6 @@ def update_abundances(abundance_replacements): if key in solar_abundance: abundances[key] = val else: - print(f"PROBLEM KEY: {key}") raise RuntimeError(f"Unrecognized element {key} provided to abundance_dict. Only elements up through Zn supported.") return abundances @@ -1067,7 +1065,6 @@ def update_abundances(abundance_replacements): 'Mn': 2.88e-07, 'Fe': 2.82e-05, 'Co': 8.32e-08, 'Ni': 1.78e-06, 'Cu': 1.62e-08, 'Zn': 3.98e-08} - atomic_mass = { 'H' : 1.00794, 'He': 4.002602, 'Li': 6.941, 'Be': 9.012182, 'B' : 10.811, 'C' : 12.0107, diff --git a/trident/light_ray.py b/trident/light_ray.py index a5db3ee87..f6ab5d696 100644 --- a/trident/light_ray.py +++ b/trident/light_ray.py @@ -1,4 +1,3 @@ - """ LightRay class and member functions. @@ -51,7 +50,7 @@ class LightRay(CosmologySplice): For compound rays, the LightRay stacks together multiple datasets in a time series in order to approximate a LightRay's path through a volume and redshift interval larger than a single simulation data output. - The outcome is thing akin to a synthetic QSO line of sight. + The outcome is something akin to a synthetic QSO line of sight. Once the LightRay object is set up, use LightRay.make_light_ray to begin making rays. Different randomizations can be created with a @@ -522,10 +521,9 @@ def make_light_ray(self, seed=None, periodic=True, if field_parameters is None: field_parameters = {} - else: # Initialize data structures. - self._data = {} + self._data = {} # temperature field is automatically added to fields if fields is None: fields = [] if ('gas', 'temperature') not in fields: diff --git a/trident/spectrum_generator.py b/trident/spectrum_generator.py index 1034ea1bc..8c0fa3fa5 100644 --- a/trident/spectrum_generator.py +++ b/trident/spectrum_generator.py @@ -371,7 +371,6 @@ def make_spectrum(self, ray, lines='all', >>> sg.plot_spectrum('spec_raw.png') """ self.observing_redshift = observing_redshift - if isinstance(ray, str): ray = load(ray) @@ -432,6 +431,7 @@ def make_spectrum(self, ray, lines='all', element='H', ion_state='I') if len(H_lines) > 0 and ly_continuum: self.add_continuum('Ly C', H_lines[0].field, 912.32336, 1.6e17, 3.0) + AbsorptionSpectrum.make_spectrum(self, ad, output_file=None, line_list_file=None, From b617daa0734fce9f0a7ae89503c9317999e897a5 Mon Sep 17 00:00:00 2001 From: Claire Kopenhafer Date: Thu, 4 Jun 2026 15:36:41 -0400 Subject: [PATCH 3/5] Clean up abundance docstrings/API --- trident/ion_balance.py | 30 ++++++++++++------------------ 1 file changed, 12 insertions(+), 18 deletions(-) diff --git a/trident/ion_balance.py b/trident/ion_balance.py index 2fb8e5e67..950909876 100644 --- a/trident/ion_balance.py +++ b/trident/ion_balance.py @@ -206,8 +206,9 @@ def add_ion_fields(ds, ions, ftype='gas', Dictionary of elemental abundances normalized to hydrogen. Keys should be elemental symbols, e.g., 'He'. By default, Trident assumes the solar - abundances of REF. Entries in this dictionary will replace the default - solar values. To completely replace the default solar abundances, specify + abundances from CLOUDY (Ferland et al. 2017). + Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, the dictionary should include all elements up through zinc. :field_suffix: boolean, optional @@ -294,7 +295,6 @@ def add_ion_fields(ds, ions, ftype='gas', def add_ion_fraction_field(atom, ion, ds, ftype="gas", ionization_table=None, - abundance_dict=None, field_suffix=False, sampling_type='local', particle_type=None): @@ -338,14 +338,6 @@ def add_ion_fraction_field(atom, ion, ds, ftype="gas", metallicity, and redshift. By default, it uses the table specified in ~/.trident/config - :abundance_dict: dictionary, optional - - Dictionary of elemental abundances normalized to hydrogen. Keys should - be elemental symbols, e.g., 'He'. By default, Trident assumes the solar - abundances of REF. Entries in this dictionary will replace the default - solar values. To completely replace the default solar abundances, specify - the dictionary should include all elements up through zinc. - :field_suffix: boolean, optional Determines whether or not to append a suffix to the field name that indicates what ionization table was used @@ -464,8 +456,9 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", Dictionary of elemental abundances normalized to hydrogen. Keys should be elemental symbols, e.g., 'He'. By default, Trident assumes the solar - abundances of REF. Entries in this dictionary will replace the default - solar values. To completely replace the default solar abundances, specify + abundances from CLOUDY (Ferland et al. 2017). + Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, the dictionary should include all elements up through zinc. :field_suffix: boolean, optional @@ -511,7 +504,6 @@ def add_ion_number_density_field(atom, ion, ds, ftype="gas", field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] add_ion_fraction_field(atom, ion, ds, ftype, ionization_table, - abundance_dict=abundance_dict, field_suffix=field_suffix, sampling_type=sampling_type) @@ -572,8 +564,9 @@ def add_ion_density_field(atom, ion, ds, ftype="gas", Dictionary of elemental abundances normalized to hydrogen. Keys should be elemental symbols, e.g., 'He'. By default, Trident assumes the solar - abundances of REF. Entries in this dictionary will replace the default - solar values. To completely replace the default solar abundances, specify + abundances from CLOUDY (Ferland et al. 2017). + Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, the dictionary should include all elements up through zinc. :field_suffix: boolean, optional @@ -675,8 +668,9 @@ def add_ion_mass_field(atom, ion, ds, ftype="gas", Dictionary of elemental abundances normalized to hydrogen. Keys should be elemental symbols, e.g., 'He'. By default, Trident assumes the solar - abundances of REF. Entries in this dictionary will replace the default - solar values. To completely replace the default solar abundances, specify + abundances from CLOUDY (Ferland et al. 2017). + Entries in this dictionary will replace the default + solar values. To completely replace the default solar abundances, the dictionary should include all elements up through zinc. :field_suffix: boolean, optional From ba77a5cce19b7598bccc2ab4109534867d8c0e4d Mon Sep 17 00:00:00 2001 From: Claire Kopenhafer Date: Thu, 4 Jun 2026 15:36:59 -0400 Subject: [PATCH 4/5] Add tests for abundance & ion tables --- .circleci/config.yml | 5 ++++ tests/test_ion_balance.py | 57 +++++++++++++++++++++++++++++++++++++++ 2 files changed, 62 insertions(+) diff --git a/.circleci/config.yml b/.circleci/config.yml index a988692e5..22d9ca180 100644 --- a/.circleci/config.yml +++ b/.circleci/config.yml @@ -113,6 +113,11 @@ commands: gunzip hm2012_lr.h5.gz mv hm2012_lr.h5 $TRIDENT_ION_DATA fi + if [ ! -f $TRIDENT_ION_DATA/fg2009_lr.h5 ]; then + wget http://trident-project.org/data/ion_table/fg2009_lr.h5.gz + gunzip fg2009_lr.h5.gz + mv fg2009_lr.h5 $TRIDENT_ION_DATA + fi # download answer test data if [ ! -f $TRIDENT_ANSWER_DATA/enzo_small/AMRCosmology.enzo ]; then pushd tests diff --git a/tests/test_ion_balance.py b/tests/test_ion_balance.py index 66206e601..ba7ce9d49 100644 --- a/tests/test_ion_balance.py +++ b/tests/test_ion_balance.py @@ -26,12 +26,14 @@ from yt.testing import \ fake_random_ds, \ fake_amr_ds +from yt.utilities.physical_constants import mh import tempfile import shutil from trident.testing import \ answer_test_data_dir, \ assert_array_rel_equal import os +import os.path as path import numpy as np @@ -282,6 +284,26 @@ def test_add_ion_fields_to_enzo(): SlicePlot(ds, 'x', field).save(dirpath) shutil.rmtree(dirpath) +def test_add_ion_fields_to_enzo_with_nonsolar_abundances(): + """ + Test adding a non-tracked metal to an Enzo dataset with nonsolar abundance + """ + abun = {"O": 5e-3} + + ds = load(ISO_GALAXY) + ad = ds.all_data() + add_ion_number_density_field('O', 6, ds, abundance_dict=abun) + field = ('gas', 'O_p5_number_density') + assert field in ds.derived_field_list + assert isinstance(ad[field], np.ndarray) + + # Test values of added field + num_dens = ds.quan(abun["O"], "1/Zsun") \ + * ad[("gas","metallicity")] \ + * ad[("gas","O_p5_ion_fraction")] \ + * ad[("gas","H_nuclei_density")] + assert np.allclose(num_dens, ad[field]) + def test_add_ion_fields_to_gizmo(): """ Test to add various ion fields to gizmo dataset and slice on them @@ -299,6 +321,26 @@ def test_add_ion_fields_to_gizmo(): SlicePlot(ds, 'x', field).save(dirpath) shutil.rmtree(dirpath) +def test_add_ion_fields_to_gizmo_with_nonsolar_abundances(): + """ + Test adding a non-tracked metal to an Enzo dataset with nonsolar abundance + """ + abun = {"Na": 2e-5} + + ds = load(FIRE_SIM) + ad = ds.all_data() + add_ion_number_density_field('Na', 2, ds, abundance_dict=abun) + field = ('gas', 'Na_p1_number_density') + assert field in ds.derived_field_list + assert isinstance(ad[field], np.ndarray) + + # Test values of added field + num_dens = ds.quan(abun["Na"], "1.0/Zsun") \ + * ad[("gas","metallicity")] \ + * ad[("gas","Na_p1_ion_fraction")] \ + * ad[("gas","H_nuclei_density")] + assert np.allclose(num_dens, ad[field]) + def test_ion_fraction_field_is_from_on_disk_fields(): """ Test to add various ion fields to Enzo dataset and slice on them @@ -353,6 +395,21 @@ def test_calculate_ion_fraction(): # Does it return all hydrogen being ionized at 1e7 K? assert calculate_ion_fraction('H II', 1e-2, 1e7, 0) == 1 + # Can it swap ionization tables? + ion_filepath = os.environ.get("TRIDENT_ION_DATA", + path.join(path.expanduser("~"), ".trident")) + ionfile_hm12 = path.join(ion_filepath, "hm2012_lr.h5") + ionfile_fg09 = path.join(ion_filepath, "fg2009_lr.h5") + + # O VI shows the biggest difference between "model families"; + # Taira et al. 2025 https://ui.adsabs.harvard.edu/abs/2025ApJ...991..221T/abstract + frac_hm12 = calculate_ion_fraction("O VI", dens, temp, reds, + ionfile_hm12) + frac_fg09 = calculate_ion_fraction("O VI", dens, temp, reds, + ionfile_fg09) + + assert not np.allclose(frac_hm12, frac_fg09) + def test_species_fraction_field_is_used_for_ion_mass_and_number_density(): """ Ensure Gadget/AREPO-style X_fraction fields are used in preference to the From ad82b09dc6859c998c7e505f64d589711ad09742 Mon Sep 17 00:00:00 2001 From: Claire Kopenhafer Date: Fri, 5 Jun 2026 09:30:39 -0400 Subject: [PATCH 5/5] Linting --- trident/ion_balance.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/trident/ion_balance.py b/trident/ion_balance.py index 950909876..0ea852738 100644 --- a/trident/ion_balance.py +++ b/trident/ion_balance.py @@ -385,7 +385,6 @@ def add_ion_fraction_field(atom, ion, ds, ftype="gas", if field_suffix: field += "_%s" % ionization_table.split(os.sep)[-1].split(".h5")[0] - global ion_table_store if field not in ion_table_store: ionTable = IonBalanceTable(ionization_table, atom) ion_table_store[field] = {'fraction': copy.deepcopy(ionTable.ion_fraction[ion-1]), @@ -1043,10 +1042,10 @@ def update_abundances(abundance_replacements): else: raise RuntimeError(f"Unrecognized element {key} provided to abundance_dict. Only elements up through Zn supported.") - return abundances + return abundances -# Taken from Cloudy documentation. +# Taken from Cloudy documentation. solar_abundance = { 'H' : 1.00e+00, 'He': 1.00e-01, 'Li': 2.04e-09, 'Be': 2.63e-11, 'B' : 6.17e-10, 'C' : 2.45e-04,