diff --git a/prody/atomic/atomgroup.py b/prody/atomic/atomgroup.py index 31bf14ebb..56a7cd182 100644 --- a/prody/atomic/atomgroup.py +++ b/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) diff --git a/prody/tests/utilities/test_misctools.py b/prody/tests/utilities/test_misctools.py index 27d6bfa6a..90509ab7b 100644 --- a/prody/tests/utilities/test_misctools.py +++ b/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]) diff --git a/prody/utilities/misctools.py b/prody/utilities/misctools.py index f9d8754dc..a4f9c22f6 100644 --- a/prody/utilities/misctools.py +++ b/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])