speed up getMasses by looking up each distinct element once

prody/atomic/atomgroup.py

@@ -1968,10 +1968,7 @@ def setData(self, array, var=fname, dtype=field.dtype,
 
                 if not np.isscalar(array):
                     if var == 'chain':
-                        max_len = 0
-                        for val in array:
-                            if len(val) > max_len:
-                                max_len = len(val)
+                        max_len = max(map(len, array), default=0)
 
                         if max_len > int(dtype[1:]):
                             dtype = dtype[0] + str(max_len)

prody/tests/utilities/test_misctools.py

@@ -1,6 +1,10 @@
+from numpy import array
+from numpy.testing import assert_allclose
+
 from prody.tests import TestCase
 
 from prody.utilities import rangeString
+from prody.utilities.misctools import getMasses
 
 
 class TestRangeString(TestCase):
@@ -26,3 +30,48 @@ def testRepeated(self):
         self.assertEqual(rangeString(list(range(10, 20)) +
                                      list(range(15, 20)) +
                                      list(range(30))), '0 to 29')
+
+
+class TestGetMasses(TestCase):
+
+    def testKnownElements(self):
+
+        assert_allclose(getMasses(['C', 'N', 'O', 'S']),
+                        [12.0107, 14.0067, 15.9994, 32.065])
+
+    def testCaseInsensitive(self):
+        """Element symbols are matched regardless of case."""
+
+        assert_allclose(getMasses(['c', 'n', 'FE']),
+                        getMasses(['C', 'N', 'Fe']))
+
+    def testUnknownElementIsZero(self):
+        """An unrecognised symbol contributes zero mass, it does not raise."""
+
+        assert_allclose(getMasses(['C', 'Xx', 'O']), [12.0107, 0., 15.9994])
+
+    def testRepeatedElements(self):
+        """Repeated symbols must all map back to their own mass.
+
+        The lookup is done once per distinct symbol and mapped back onto the
+        atoms, so a wrong mapping would show up here as masses landing on the
+        wrong atoms.
+        """
+
+        elements = ['O', 'C', 'C', 'N', 'O', 'S', 'C', 'N']
+        expected = [15.9994, 12.0107, 12.0107, 14.0067,
+                    15.9994, 32.065, 12.0107, 14.0067]
+        assert_allclose(getMasses(elements), expected)
+
+    def testEmpty(self):
+
+        self.assertEqual(len(getMasses([])), 0)
+
+    def testString(self):
+        """A single symbol still returns a scalar."""
+
+        self.assertAlmostEqual(getMasses('c'), 12.0107)
+
+    def testArrayInput(self):
+
+        assert_allclose(getMasses(array(['C', 'O'])), [12.0107, 15.9994])

prody/utilities/misctools.py

@@ -370,13 +370,15 @@ def getMasses(elements):
     if isinstance(elements, str):
         return mass_dict[elements.capitalize()]
     else:
-        masses = zeros(len(elements))
-        for i,element in enumerate(elements):
-            if element.capitalize() in mass_dict:
-                masses[i] = mass_dict[element.capitalize()]
-            else:
-                masses[i] = 0.
-        return masses
+        elements = asarray(elements)
+        if elements.size == 0:
+            return zeros(0)
+        # a structure has only a handful of distinct element symbols, so the
+        # lookup is done once per distinct symbol rather than once per atom
+        unique_elements, inverse = unique(elements, return_inverse=True)
+        unique_masses = array([mass_dict.get(str(element).capitalize(), 0.)
+                               for element in unique_elements])
+        return unique_masses[inverse.reshape(-1)]
 
 def count(L, a=None):
     return len([b for b in L if b is a])