Add add_atom method to PosCar.

Author: PytLab

tests/poscar_test.py

@@ -13,11 +13,11 @@ class PosCarTest(unittest.TestCase):
 
     def setUp(self):
         self.maxDiff = True
+        self.filename = path + "/POSCAR"
 
     def test_get_poscar_content(self):
         " Make sure we can get the correct poscar content. "
-        filename = path + "/POSCAR"
-        poscar = PosCar(filename)
+        poscar = PosCar(self.filename)
 
         ref_content = """Created by VASPy\n 1.000000000\n    7.29321435   -4.21073927    0.00000000\n    0.00000000    8.42147853    0.00000000\n   -0.00000000    0.00000000   16.87610843\nPt   \n   36\nSelective Dynamics\nDirect\n    0.244666665792    0.223999996980    0.135815443038F    F    F    \n    0.022444443570    0.112888885869    0.271630886077T    T    T    \n    0.133555554681    0.001777774758    0.000000000000F    F    F    \n    0.133555554681    0.001777774758    0.407446329115T    T    T    \n    0.577999999126    0.223999996980    0.135815443038F    F    F    \n    0.355777776904    0.112888885869    0.271630886077T    T    T    \n    0.466888888015    0.001777774758    0.000000000000F    F    F    \n    0.466888888015    0.001777774758    0.407446329115T    T    T    \n    0.911333332459    0.223999996980    0.135815443038F    F    F    \n    0.689111110237    0.112888885869    0.271630886077T    T    T    \n    0.800222221348    0.001777774758    0.000000000000F    F    F    \n    0.800222221348    0.001777774758    0.407446329115T    T    T    \n    0.244666665792    0.557333330313    0.135815443038F    F    F    \n    0.022444443570    0.446222219202    0.271630886077T    T    T    \n    0.133555554681    0.335111108091    0.000000000000F    F    F    \n    0.133555554681    0.335111108091    0.407446329115T    T    T    \n    0.577999999126    0.557333330313    0.135815443038F    F    F    \n    0.355777776904    0.446222219202    0.271630886077T    T    T    \n    0.466888888015    0.335111108091    0.000000000000F    F    F    \n    0.466888888015    0.335111108091    0.407446329115T    T    T    \n    0.911333332459    0.557333330313    0.135815443038F    F    F    \n    0.689111110237    0.446222219202    0.271630886077T    T    T    \n    0.800222221348    0.335111108091    0.000000000000F    F    F    \n    0.800222221348    0.335111108091    0.407446329115T    T    T    \n    0.244666665792    0.890666663647    0.135815443038F    F    F    \n    0.022444443570    0.779555552536    0.271630886077T    T    T    \n    0.133555554681    0.668444441424    0.000000000000F    F    F    \n    0.133555554681    0.668444441424    0.407446329115T    T    T    \n    0.577999999126    0.890666663647    0.135815443038F    F    F    \n    0.355777776904    0.779555552536    0.271630886077T    T    T    \n    0.466888888015    0.668444441424    0.000000000000F    F    F    \n    0.466888888015    0.668444441424    0.407446329115T    T    T    \n    0.911333332459    0.890666663647    0.135815443038F    F    F    \n    0.689111110237    0.779555552536    0.271630886077T    T    T    \n    0.800222221348    0.668444441424    0.000000000000F    F    F    \n    0.800222221348    0.668444441424    0.407446329115T    T    T    \n"""
 
@@ -27,15 +27,29 @@ def test_get_poscar_content(self):
 
     def test_get_xyz_content(self):
         " Make sure we can get correct xyz file content from poscar. "
-        filename = path + "/POSCAR"
-        poscar = PosCar(filename)
+        poscar = PosCar(self.filename)
 
         ref_content = """          36\nSTEP =        1\nPt   0.244666665792   0.22399999698  0.135815443038\nPt    0.02244444357  0.112888885869  0.271630886077\nPt   0.133555554681  0.001777774758             0.0\nPt   0.133555554681  0.001777774758  0.407446329115\nPt   0.577999999126   0.22399999698  0.135815443038\nPt   0.355777776904  0.112888885869  0.271630886077\nPt   0.466888888015  0.001777774758             0.0\nPt   0.466888888015  0.001777774758  0.407446329115\nPt   0.911333332459   0.22399999698  0.135815443038\nPt   0.689111110237  0.112888885869  0.271630886077\nPt   0.800222221348  0.001777774758             0.0\nPt   0.800222221348  0.001777774758  0.407446329115\nPt   0.244666665792  0.557333330313  0.135815443038\nPt    0.02244444357  0.446222219202  0.271630886077\nPt   0.133555554681  0.335111108091             0.0\nPt   0.133555554681  0.335111108091  0.407446329115\nPt   0.577999999126  0.557333330313  0.135815443038\nPt   0.355777776904  0.446222219202  0.271630886077\nPt   0.466888888015  0.335111108091             0.0\nPt   0.466888888015  0.335111108091  0.407446329115\nPt   0.911333332459  0.557333330313  0.135815443038\nPt   0.689111110237  0.446222219202  0.271630886077\nPt   0.800222221348  0.335111108091             0.0\nPt   0.800222221348  0.335111108091  0.407446329115\nPt   0.244666665792  0.890666663647  0.135815443038\nPt    0.02244444357  0.779555552536  0.271630886077\nPt   0.133555554681  0.668444441424             0.0\nPt   0.133555554681  0.668444441424  0.407446329115\nPt   0.577999999126  0.890666663647  0.135815443038\nPt   0.355777776904  0.779555552536  0.271630886077\nPt   0.466888888015  0.668444441424             0.0\nPt   0.466888888015  0.668444441424  0.407446329115\nPt   0.911333332459  0.890666663647  0.135815443038\nPt   0.689111110237  0.779555552536  0.271630886077\nPt   0.800222221348  0.668444441424             0.0\nPt   0.800222221348  0.668444441424  0.407446329115\n"""
 
         ret_content = poscar.get_xyz_content()
 
         self.assertEqual(ref_content, ret_content)
 
+    def test_add_atom(self):
+        "Make sure we can add a new atom to current poscar. "
+        poscar = PosCar(self.filename)
+        ori_data = poscar.data
+        ori_tf = poscar.tf
+
+        poscar.add_atom('O', [0.5, 0.5, 0.5])
+
+        self.assertListEqual(poscar.atom_types, ['Pt', 'O'])
+        self.assertListEqual(poscar.atom_numbers, [36, 1])
+        self.assertListEqual(ori_data.tolist() + [[0.5, 0.5, 0.5]],
+                             poscar.data.tolist())
+        self.assertListEqual(ori_tf.tolist() + [['T', 'T', 'T']],
+                             poscar.tf.tolist())
+
 if "__main__" == __name__: 
     suite = unittest.TestLoader().loadTestsFromTestCase(PosCarTest)
     unittest.TextTestRunner(verbosity=2).run(suite) 

vaspy/atomco.py

@@ -437,6 +437,46 @@ def get_content(self):
         content = self.get_poscar_content()
         return content
 
+    def add_atom(self, atom_type, coordinate, fix=['T', 'T', 'T']):
+        """
+        Add a new atom to coordinate file.
+
+        Parameters:
+        -----------
+        atom_type: element type of the atom, str.
+        coordinate: position of the added atom, list of float.
+        fix: flags for fixed atom in three directions, list of str.
+
+        Example:
+        --------
+        >>> poscar.add_atom('C', [0.5, 0.5, 0.3])
+        """
+        atomco_dict = self.atomco_dict
+        tf_dict = self.tf_dict
+
+        self.natom += 1
+        self.totline += 1
+
+        if atom_type in self.atom_types:
+            atomco_dict[atom_type].append(coordinate)
+            tf_dict[atom_type].append(fix)
+            idx = self.atom_types.index(atom_type)
+            self.atom_numbers[idx] += 1
+        else:
+            self.atom_types.append(atom_type)
+            atomco_dict[atom_type] = [coordinate]
+            tf_dict[atom_type] = [fix]
+            self.atom_numbers.append(1)
+
+        # New data and fix info.
+        data, tf = [], []
+        for atom_type in self.atom_types:
+            data += atomco_dict[atom_type]
+            tf += tf_dict[atom_type]
+
+        self.data = np.float64(np.array(data))
+        self.tf = np.array(tf)
+
     def tofile(self, filename='POSCAR_c'):
         "生成文件"
         "PosCar object to POSCAR or CONTCAR."