add unittest for convert cell

Author: felix5572

dpdata/cp2k/output.py

@@ -17,6 +17,59 @@
 
 avail_patterns.append(re.compile(r'^ INITIAL POTENTIAL ENERGY'))
 avail_patterns.append(re.compile(r'^ ENSEMBLE TYPE'))
+
+def cell_to_low_triangle(A,B,C,alpha,beta,gamma):
+    """
+        Convert cell to low triangle matrix.
+
+        Parameters
+        ----------
+        A : float
+            cell length A
+        B : float
+            cell length B
+        C : float
+            cell length C
+        alpha : float
+            radian. The angle between vector B and  vector C.
+        beta : float
+            radian. The angle between vector A and  vector C.
+        gamma : float
+            radian. The angle between vector B and  vector C.
+        
+        Returns
+        -------
+        cell : list
+            The cell matrix used by dpdata in low triangle form.
+    """
+    if not np.pi*5/180<alpha< np.pi*175/180:
+        raise RuntimeError("alpha=={}: must be a radian, and \
+            must be in np.pi*5/180 < alpha < np.pi*175/180".format(alpha))
+    if not np.pi*5/180<beta< np.pi*175/180:
+        raise RuntimeError("beta=={}: must be a radian, and \
+            must be in np.pi*5/180 < beta < np.pi*175/180".format(beta))
+    if not np.pi*5/180<gamma< np.pi*175/180:
+            raise RuntimeError("gamma=={}: must be a radian, and \
+                must be in np.pi*5/180 < gamma < np.pi*175/180".format(gamma))
+    if not A > 0.2:
+        raise RuntimeError("A=={}, must be greater than 0.2".format(A))
+    if not B > 0.2:
+        raise RuntimeError("B=={}, must be greater than 0.2".format(B))
+    if not C > 0.2:
+        raise RuntimeError("C=={}, must be greater than 0.2".format(C))
+
+    lx = A
+    xy = B * np.cos(gamma)
+    xz = C * np.cos(beta)
+    ly = B* np.sin(gamma)
+    if not ly > 0.1:
+        raise RuntimeError("ly:=B* np.sin(gamma)=={}, must be greater than 0.1",format(ly))
+    yz = (B*C*np.cos(alpha)-xy*xz)/ly
+    lz = np.sqrt(C**2-xz**2-yz**2)
+    cell = np.asarray([[lx, 0 , 0],
+            [xy, ly, 0 ],
+            [xz, yz, lz]]).astype('float32')
+    return cell
 class Cp2kSystems(object):
     """
     deal with cp2k outputfile 
@@ -123,15 +176,17 @@ def handle_single_log_frame(self, lines):
                 cell_gamma = np.deg2rad(float(cell_angle_pattern.match(line).groupdict()['gamma']))
                 cell_flag+=1
         if cell_flag == 2:
-            lx = cell_A
-            xy = cell_B * np.cos(cell_gamma)
-            xz = cell_C * np.cos(cell_beta)
-            ly = cell_B* np.sin(cell_gamma)
-            yz = (cell_B*cell_C*np.cos(cell_alpha)-xy*xz)/ly
-            lz = np.sqrt(cell_C**2-xz**2-yz**2)
-            self.cell = [[lx, 0 , 0],
-                    [xy, ly, 0 ],
-                    [xz, yz, lz]]
+            self.cell = cell_to_low_triangle(cell_A,cell_B,cell_C,
+                cell_alpha,cell_beta,cell_gamma)
+            # lx = cell_A
+            # xy = cell_B * np.cos(cell_gamma)
+            # xz = cell_C * np.cos(cell_beta)
+            # ly = cell_B* np.sin(cell_gamma)
+            # yz = (cell_B*cell_C*np.cos(cell_alpha)-xy*xz)/ly
+            # lz = np.sqrt(cell_C**2-xz**2-yz**2)
+            # self.cell = [[lx, 0 , 0],
+            #         [xy, ly, 0 ],
+            #         [xz, yz, lz]]
 
         element_index = -1
         element_dict = OrderedDict()

tests/test_cell_to_low_triangle.py

@@ -0,0 +1,35 @@
+import os
+import numpy as np
+import unittest
+from dpdata.cp2k.output import cell_to_low_triangle
+
+class CellToLowTriangle(unittest.TestCase):
+    def test_func1(self):
+        cell_1 = cell_to_low_triangle(6,6,6,np.pi*1/2, np.pi*1/2, np.pi*1/2)
+        cell_2 = np.asarray([[6,0,0],[0,6,0],[0,0,6]])
+        for ii in range(3):
+            for jj in range(3):
+                self.assertAlmostEqual(cell_1[ii,jj], cell_2[ii,jj], places=6)
+
+    def test_func2(self):
+        cell_1 = cell_to_low_triangle(6,6,6,np.pi*1/3, np.pi*1/3, np.pi*1/3)
+        cell_2 = np.asarray([
+            [6,0,0],
+            [3,3*np.sqrt(3),0],
+            [3,np.sqrt(3),2*np.sqrt(6)]])
+        for ii in range(3):
+            for jj in range(3):
+                self.assertAlmostEqual(cell_1[ii,jj], cell_2[ii,jj], places=6)
+
+    def test_func3(self):
+        with self.assertRaises(Exception) as c:
+            cell_to_low_triangle(0.1,6,6,np.pi*1/2,np.pi*1/2,np.pi*1/2)
+        self.assertTrue("A==0.1" in str(c.exception))
+
+    def test_func4(self):
+        with self.assertRaises(Exception) as c:
+            cell_to_low_triangle(6,6,6,np.pi*3/180,np.pi*1/2,np.pi*1/2)
+        self.assertTrue("alpha" in str(c.exception))
+
+if __name__ == '__main__':
+    unittest.main()