Simplify the atomic symbol to number conversion.

Author: rwest

arkane/common.py

@@ -391,10 +391,10 @@ def get_element_mass(input_element, isotope=None):
     elif isinstance(input_element, str):
         symbol = input_element
         try:
-            number = next(key for key, value in symbol_by_number.items() if value == input_element)
-        except:
-            symbol = input_element[0] + input_element[1].lower()
-            number = [key for key, value in symbol_by_number.items() if value == symbol][0]
+            number = number_by_symbol[symbol]
+        except KeyError:
+            symbol = input_element.capitalize()
+            number = number_by_symbol[symbol]
 
     if symbol is None or number is None:
         raise ValueError('Could not identify element {0}'.format(input_element))
@@ -438,6 +438,7 @@ def get_element_mass(input_element, isotope=None):
                     92: 'U', 93: 'Np', 94: 'Pu', 95: 'Am', 96: 'Cm', 97: 'Bk', 98: 'Cf', 99: 'Es', 100: 'Fm', 101: 'Md',
                     102: 'No', 103: 'Lr', 104: 'Rf', 105: 'Db', 106: 'Sg', 107: 'Bh', 108: 'Hs', 109: 'Mt', 110: 'Ds',
                     111: 'Rg', 112: 'Cn', 113: 'Nh', 114: 'Fl', 115: 'Mc', 116: 'Lv', 117: 'Ts', 118: 'Og'}
+number_by_symbol = {value: key for key, value in symbol_by_number.items()}
 
 # Structure of mass_by_symbol items: list(list(isotope1, mass1, weight1), list(isotope2, mass2, weight2), ...)
 mass_by_symbol = {

arkane/commonTest.py

@@ -459,6 +459,7 @@ def test_get_mass(self):
         """Test that the correct mass/number/isotope is returned from get_element_mass"""
         self.assertEqual(get_element_mass(1), (1.00782503224, 1))  # test input by integer
         self.assertEqual(get_element_mass('Si'), (27.97692653465, 14))  # test string input and most common isotope
+        self.assertEqual(get_element_mass('SI'), (27.97692653465, 14))  # test string in all caps
         self.assertEqual(get_element_mass('C', 13), (13.00335483507, 6))  # test specific isotope
         self.assertEqual(get_element_mass('Bk'), (247.0703073, 97))  # test a two-element array (no isotope data)