From fb43eed4a3f76d69a4706ecdda8977efd9663bb3 Mon Sep 17 00:00:00 2001
From: Zachary Hafen <zachary.h.hafen@gmail.com>
Date: Fri, 26 Feb 2021 16:20:25 -0800
Subject: [PATCH] Added basic integration with linetools.

---
 tests/test_line_database.py |  8 ++++++
 trident/ion_balance.py      |  3 ++-
 trident/line_database.py    | 50 +++++++++++++++++++++++++++++++++++--
 3 files changed, 58 insertions(+), 3 deletions(-)

diff --git a/tests/test_line_database.py b/tests/test_line_database.py
index 3561dfd51d..72b42105ab 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 b7afa03f7a..3238d3923a 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 af1be99e3b..1cad1e3e15 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 \
@@ -146,8 +147,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`.
 
@@ -187,6 +188,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**
 
@@ -198,6 +203,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))
 
@@ -450,6 +481,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: