Support array arguments in Lattice dist, angle.

Author: pavoljuhas

diffpy/Structure/lattice.py

@@ -330,9 +330,15 @@ def rnorm(self, hkl):
 
 
     def dist(self, u, v):
-        """Return distance of 2 points in lattice coordinates.
+        """Return distance between 2 points in lattice coordinates.
+
+        u    -- vector or N-by-3 matrix of fractional coordinates.
+        v    -- vector or N-by-3 matrix of fractional coordinates.
+                Sizes of u, v must match when both of them are matrices.
+
+        Return float or an array of the same length as the matrix.
         """
-        duv = numpy.array(u) - v
+        duv = numpy.asarray(u) - v
         return self.norm(duv)
 
 
@@ -341,8 +347,14 @@ def angle(self, u, v):
         """
         ca = self.dot(u, v)/( self.norm(u)*self.norm(v) )
         # avoid round-off errors that would make abs(ca) greater than 1
-        ca = max(min(ca, 1), -1)
-        return math.degrees(math.acos(ca))
+        if numpy.isscalar(ca):
+            ca = max(min(ca, 1), -1)
+            rv = math.degrees(math.acos(ca))
+        else:
+            ca[ca < -1] = -1
+            ca[ca > +1] = +1
+            rv = numpy.degrees(numpy.arccos(ca))
+        return rv
 
 
     def isanisotropic(self, umx):

diffpy/Structure/tests/TestLattice.py

@@ -18,6 +18,7 @@
 
 import unittest
 import numpy
+import numpy.linalg as numalg
 
 from diffpy.Structure import Lattice, LatticeError
 
@@ -168,8 +169,6 @@ def test_readonly_properties(self):
 
     def test_setLatBase(self):
         """check calculation of unit cell rotation"""
-        import numpy
-        import numpy.linalg as numalg
         base = numpy.array([[ 1.0,  1.0,  0.0],
                           [ 0.0,  1.0,  1.0],
                           [ 1.0,  0.0,  1.0]])
@@ -214,7 +213,7 @@ def test_dot(self):
         '''check dot product of lattice vectors.'''
         L = self.lattice
         L.setLatPar(gamma=120)
-        self.assertAlmostEqual(-0.5, L.dot([1, 0, 0], [0, 1, 0]))
+        self.assertAlmostEqual(-0.5, L.dot([1, 0, 0], [0, 1, 0]), self.places)
         va5 = numpy.tile([1.0, 0.0, 0.0], (5, 1))
         vb5 = numpy.tile([0.0, 1.0, 0.0], (5, 1))
         self.assertTrue(numpy.array_equal(5 * [-0.5], L.dot(va5, vb5)))
@@ -229,7 +228,7 @@ def test_norm(self):
         u = numpy.array([[3, 4, 0], [1, 1, 1]])
         self.assertListAlmostEqual([5, 3**0.5], self.lattice.norm(u))
         self.lattice.setLatPar(gamma=120)
-        self.assertAlmostEqual(1, self.lattice.norm([1, 1, 0]))
+        self.assertAlmostEqual(1, self.lattice.norm([1, 1, 0]), self.places)
         return
 
 
@@ -239,12 +238,51 @@ def test_rnorm(self):
         L.setLatPar(1, 1.5, 2.3, 80, 95, 115)
         r = L.reciprocal()
         hkl = [0.5, 0.3, 0.2]
-        self.assertAlmostEqual(r.norm(hkl), L.rnorm(hkl))
+        self.assertAlmostEqual(r.norm(hkl), L.rnorm(hkl), self.places)
         hkl5 = numpy.tile(hkl, (5, 1))
         self.assertListAlmostEqual(5 * [r.norm(hkl)], L.rnorm(hkl5))
         return
 
 
+    def test_dist(self):
+        '''check dist function for distance between lattice points.'''
+        L = self.lattice
+        L.setLatPar(1, 1.5, 2.3, 80, 95, 115)
+        u = [0.1, 0.3, 0.7]
+        v = [0.3, 0.7, 0.7]
+        d0 = numalg.norm(L.cartesian(numpy.array(u) - v))
+        self.assertAlmostEqual(d0, L.dist(u, v), self.places)
+        self.assertAlmostEqual(d0, L.dist(v, u), self.places)
+        u5 = numpy.tile(u, (5, 1))
+        v5 = numpy.tile(v, (5, 1))
+        self.assertListAlmostEqual(5 * [d0], L.dist(u, v5))
+        self.assertListAlmostEqual(5 * [d0], L.dist(u5, v))
+        self.assertListAlmostEqual(5 * [d0], L.dist(v5, u5))
+        return
+
+
+    def test_angle(self):
+        '''check angle calculation between lattice vectors.'''
+        from math import degrees, acos
+        L = self.lattice
+        L.setLatPar(1, 1.5, 2.3, 80, 95, 115)
+        u = [0.1, 0.3, 0.7]
+        v = [0.3, 0.7, 0.7]
+        uc = L.cartesian(u)
+        vc = L.cartesian(v)
+        a0 = degrees(acos(numpy.dot(uc, vc) /
+                          (numalg.norm(uc) * numalg.norm(vc))))
+        self.assertAlmostEqual(a0, L.angle(u, v), self.places)
+        self.assertAlmostEqual(a0, L.angle(v, u), self.places)
+        u5 = numpy.tile(u, (5, 1))
+        v5 = numpy.tile(v, (5, 1))
+        self.assertListAlmostEqual(5 * [a0], L.angle(u, v5))
+        self.assertListAlmostEqual(5 * [a0], L.angle(u5, v))
+        self.assertListAlmostEqual(5 * [a0], L.angle(v5, u5))
+        return
+
+
+
     def test_repr(self):
         """check string representation of this lattice"""
         r = repr(self.lattice)