diff --git a/tests/test_line_database.py b/tests/test_line_database.py index ff742e874..c09ad4ded 100644 --- a/tests/test_line_database.py +++ b/tests/test_line_database.py @@ -62,6 +62,14 @@ def test_line_database_from_input(): assert ld.lines_all[0].identifier == HI.identifier print(ld) +def test_line_database_add_line_linetools(): + ld = LineDatabase() + HI = Line('H', 'I', 1216, 626500000.0, 2.3, identifier='Ly a') + ld.add_line('H', 'I', 1216, use_linetools=True, identifier='Ly a') + np.testing.assert_allclose( ld.lines_all[0].gamma, HI.gamma ) + assert ld.lines_all[0].identifier == HI.identifier + print(ld) + def test_select_lines_from_line_database(): ld = LineDatabase('lines.txt') assert len(ld.select_lines('Ne')) == 8 # 8 listed Ne lines diff --git a/trident/ion_balance.py b/trident/ion_balance.py index 39a9c062d..355f187bb 100644 --- a/trident/ion_balance.py +++ b/trident/ion_balance.py @@ -245,7 +245,8 @@ def add_ion_fields(ds, ions, ftype='gas', # If line_database is set, then use the underlying file as the line list # to select ions from. if line_database is not None: - line_database = LineDatabase(line_database) + if not isinstance( line_database, LineDatabase ): + line_database = LineDatabase(line_database) ion_list = line_database.parse_subset_to_ions(ions) # Otherwise, any ion can be selected (not just ones in the line list). diff --git a/trident/line_database.py b/trident/line_database.py index 1f3f45e04..08fd584ca 100644 --- a/trident/line_database.py +++ b/trident/line_database.py @@ -12,6 +12,7 @@ #----------------------------------------------------------------------------- import os +import warnings from yt.funcs import \ mylog from trident.config import \ @@ -147,8 +148,8 @@ def __init__(self, input_file=None): else: self.input_file = 'Manually Entered' - def add_line(self, element, ion_state, wavelength, gamma, - f_value, field=None, identifier=None): + def add_line(self, element, ion_state, wavelength, gamma=None, + f_value=None, field=None, identifier=None, use_linetools=False): """ Manually add a line to the :class:`~trident.LineDatabase`. @@ -188,6 +189,10 @@ def add_line(self, element, ion_state, wavelength, gamma, An optional identifier for the transition Example: 'Ly a' for Lyman alpha + :use_linetools: bool + + If True retrieve the values for wavelength, gamma, and f_value + from the linetools python package (linetools.readthedocs.io). **Example** @@ -199,6 +204,32 @@ def add_line(self, element, ion_state, wavelength, gamma, >>> ldb.add_line('H', 'I', 1215.67, 469860000, 0.41641, 'Ly a') >>> print(ldb.lines_all) """ + + if use_linetools: + # Retrieve + def linetools_id( wrest ): + return '%s%s %d' % (element, ion_state, round(float(wrest), 0)) + line = self.linetools_linelist[ linetools_id(wavelength) ] + + # Check for rounding differences + if line is None: + warnings.warn( 'Line ' + linetools_id(wavelength) + \ + ' not found in linetools database. Checking adjacent wavelengths.') + line = self.linetools_linelist[ linetools_id(wavelength-1) ] + if line is None: + line = self.linetools_linelist[ linetools_id(wavelength+1) ] + if line is None: + raise KeyError( 'No line found in linetools database.' ) + + # Format + wavelength = line['wrest'].value + gamma = line['gamma'].value + f_value = line['f'] + + else: + assert gamma is not None + assert f_value is not None + self.lines_all.append(Line(element, ion_state, wavelength, gamma, f_value, field, identifier)) @@ -451,6 +482,21 @@ def parse_subset_to_ions(self, subsets=None): ions = uniquify(ions) return ions + @property + def linetools_linelist(self): + '''Handler for a linetools LineList class. + ''' + + if not hasattr( self, '_linetools_linelist' ): + + # Only doing the import here so rest of the class doesn't + # break if the feature isn't enabled. + from linetools.lists.linelist import LineList + + self._linetools_linelist = LineList('ISM') + + return self._linetools_linelist + def __repr__(self): disp = "" for line in self.lines_all: