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 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..0ea852738 100644 --- a/trident/ion_balance.py +++ b/trident/ion_balance.py @@ -40,7 +40,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 +135,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 +183,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 +202,15 @@ 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 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 Determines whether or not to append a suffix to the field name that @@ -277,7 +289,9 @@ 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, @@ -371,9 +385,9 @@ 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: + 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 +403,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 +451,15 @@ 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 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 Determines whether or not to append a suffix to the field @@ -464,6 +488,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 = solar_abundance + else: + abundance_store = update_abundances(abundance_dict) + atom = atom.capitalize() field = "%s_p%d_number_density" % (atom, ion-1) @@ -480,6 +511,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 +559,15 @@ 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 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 Determines whether or not to append a suffix to the field @@ -555,6 +596,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 +605,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 +614,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 +663,15 @@ 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 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 Determines whether or not to append a suffix to the field @@ -646,6 +700,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 +709,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 +826,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[ftype, 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 +846,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, @@ -940,16 +998,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) @@ -971,9 +1029,23 @@ 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() -# Taken from Cloudy documentation. + # 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: + raise RuntimeError(f"Unrecognized element {key} provided to abundance_dict. Only elements up through Zn supported.") + + return abundances + +# 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, 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..8c0fa3fa5 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, @@ -400,7 +411,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), @@ -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)