Test ScatteringFactorTable.lookup with arrays argument.

Author: pavoljuhas

diffpy/srreal/tests/testscatteringfactortable.py

@@ -4,9 +4,9 @@
 """
 
 
-import os
 import unittest
 import cPickle
+import numpy
 
 from diffpy.srreal.scatteringfactortable import ScatteringFactorTable
 
@@ -132,6 +132,23 @@ def test_derived_clone(self):
         self.assertEqual(set(['Xy']), lsft2.getCustomSymbols())
         return
 
+    def test_lookup(self):
+        """Check ScatteringFactorTable.lookup handling of array arguments.
+        """
+        qa = numpy.linspace(0, 50)
+        sftx = self.sftx
+        fmn0 = numpy.array([sftx.lookup('Mn', x) for x in qa])
+        fmn1 = sftx.lookup('Mn', qa)
+        self.assertTrue(numpy.array_equal(fmn0, fmn1))
+        self.assertTrue(numpy.array_equal(
+            fmn0.reshape(5, 10), sftx.lookup('Mn', qa.reshape(5, 10))))
+        self.assertTrue(numpy.array_equal(
+            fmn0.reshape(5, 2, 5), sftx.lookup('Mn', qa.reshape(5, 2, 5))))
+        self.assertTrue(numpy.array_equal(fmn0, sftx.lookup('Mn', list(qa))))
+        self.assertRaises(TypeError, sftx.lookup, 'Na', 'asdf')
+        self.assertRaises(TypeError, sftx.lookup, 'Na', {})
+        return
+
 # End of class TestScatteringFactorTable
 
 if __name__ == '__main__':