Merge pull request #65 from ahxbcn/fix_vibration

Fix bug in vibration.py

Author: pxlxingliang

src/abacusagent/modules/submodules/vibration.py

@@ -93,17 +93,65 @@ def collect_force(abacusjob_dir):
 
     return metrics['force']
 
-def filter_force(forces, selected_atoms):
+def prepare_abacus_vibration_analysis(input_params: Dict[str, Any],
+                                      stru: AbacusStru,
+                                      work_path: Path,
+                                      abacus_inputs_dir: Path,
+                                      selected_atoms: Optional[List[int]] = None,
+                                      stepsize: float = 0.01,
+                                      vib_cache_dir: str = "vib"):
     """
-    Select forces belong to selected atoms
+    Prepare ABACUS input files for vibrational analysis.
     """
-    selected_atoms_force_idx = []
-    for selected_atom in selected_atoms:
-        selected_atoms_force_idx += [selected_atom*3, selected_atom*3+1, selected_atom*3+2]
+    stru_file = input_params.get('stru_file', "STRU")
+    displaced_stru = copy.deepcopy(stru)
+    original_stru_coord = np.array(stru.get_coord(bohr=False, direct=False))
+
+    if selected_atoms is None:
+        selected_atoms = [i for i in range(stru.get_natoms())]
+
+    selected_atoms.sort()
 
-    filtered_force = np.array(forces)[selected_atoms_force_idx]
+    DIRECTION_MAP = ['x', 'y', 'z']
+    STEP_MAP = {'+': 1, '-': -1}
+    disped_stru_job_paths = []
+    disped_stru_cache_labels = []
+
+    # Prepare ABACUS input files for the given structure
+    abacus_scf_work_path = os.path.join(work_path, "SCF")
+    original_stru_job_path = os.path.join(abacus_scf_work_path, "eq")
+    os.makedirs(original_stru_job_path)
 
-    return filtered_force
+    link_abacusjob(src=abacus_inputs_dir,
+                   dst=original_stru_job_path,  
+                   copy_files=['INPUT', 'STRU', 'KPT', 'abacus.log'],
+                   exclude_directories=True)
+    WriteInput(input_params, os.path.join(original_stru_job_path, "INPUT"))
+    disped_stru_job_paths.append(original_stru_job_path)
+    disped_stru_cache_labels.append(os.path.join(vib_cache_dir, 'cache.eq.json'))
+
+    # Prepare ABACUS input files for each displaced structure. nfree is assumed to be 2.
+    for selected_atom in selected_atoms:
+        for direction in range(3): # x, y and z directions
+            displaced_stru_coord = copy.deepcopy(original_stru_coord)
+            for step in STEP_MAP.keys(): # Two steps along one direction
+                disped_stru_job_path = os.path.join(abacus_scf_work_path, f"disp-{selected_atom}-{DIRECTION_MAP[direction]}{step}")
+                os.makedirs(disped_stru_job_path)
+
+                link_abacusjob(src=abacus_inputs_dir,
+                               dst=disped_stru_job_path, 
+                               copy_files=['INPUT', 'STRU', 'KPT', 'abacus.log'],
+                               exclude_directories=True)
+                
+                displaced_stru_coord[selected_atom][direction] = original_stru_coord[selected_atom][direction] + stepsize * STEP_MAP[step]
+                displaced_stru.set_coord(displaced_stru_coord, bohr=False, direct=False)
+                WriteInput(input_params, os.path.join(disped_stru_job_path, "INPUT"))
+                displaced_stru.write(os.path.join(disped_stru_job_path, stru_file))
+
+                disped_stru_job_paths.append(disped_stru_job_path)
+                disped_stru_cache_labels.append(os.path.join(vib_cache_dir, f"cache.{selected_atom}{DIRECTION_MAP[direction]}{step}.json"))
+        
+    return disped_stru_job_paths, disped_stru_cache_labels
 
 def abacus_vibration_analysis(abacus_inputs_dir: Path,
                               selected_atoms: Optional[List[int]] = None,
@@ -142,7 +190,7 @@ def abacus_vibration_analysis(abacus_inputs_dir: Path,
             raise ValueError("stepsize should be positive.")
 
         work_path = Path(generate_work_path()).absolute()
-
+        
         input_params = ReadInput(os.path.join(abacus_inputs_dir, "INPUT"))
         stru_file = input_params.get('stru_file', "STRU")
         stru = AbacusStru.ReadStru(os.path.join(abacus_inputs_dir, stru_file))
@@ -157,66 +205,26 @@ def abacus_vibration_analysis(abacus_inputs_dir: Path,
                          indices=selected_atoms,
                          delta=stepsize,
                          nfree=nfree)
-        
-        displaced_stru = copy.deepcopy(stru)
-        original_stru_coord = np.array(stru.get_coord(bohr=False, direct=False))
-
-        if selected_atoms is None:
-            selected_atoms = [i for i in range(stru.get_natoms())]
-
-        selected_atoms.sort()
-        
-        DIRECTION_MAP = ['x', 'y', 'z']
-        STEP_MAP = {'+': 1, '-': -1}
-        disped_stru_job_paths = []
-        disped_stru_cache_labels = []
-
-        # Prepare ABACUS input files for the given structure
-        abacus_scf_work_path = os.path.join(work_path, "SCF")
-        original_stru_job_path = os.path.join(abacus_scf_work_path, "eq")
-        os.makedirs(original_stru_job_path)
-
-        link_abacusjob(src=abacus_inputs_dir,
-                       dst=original_stru_job_path,  
-                       copy_files=['INPUT', 'STRU', 'KPT', 'abacus.log'],
-                       exclude_directories=True)
-        WriteInput(input_params, os.path.join(original_stru_job_path, "INPUT"))
-        disped_stru_job_paths.append(original_stru_job_path)
-        disped_stru_cache_labels.append(os.path.join(vib_cache_dir, 'cache.eq.json'))
-
-        # Prepare ABACUS input files for each displaced structure. nfree is assumed to be 2.
-        for selected_atom in selected_atoms:
-            for direction in range(3): # x, y and z directions
-                displaced_stru_coord = copy.deepcopy(original_stru_coord)
-                for step in STEP_MAP.keys(): # Two steps along one direction
-                    disped_stru_job_path = os.path.join(abacus_scf_work_path, f"disp-{selected_atom}-{DIRECTION_MAP[direction]}{step}")
-                    os.makedirs(disped_stru_job_path)
-
-                    link_abacusjob(src=abacus_inputs_dir,
-                                   dst=disped_stru_job_path, 
-                                   copy_files=['INPUT', 'STRU', 'KPT', 'abacus.log'],
-                                   exclude_directories=True)
-                    
-                    displaced_stru_coord[selected_atom][direction] = original_stru_coord[selected_atom][direction] + stepsize * STEP_MAP[step]
-                    displaced_stru.set_coord(displaced_stru_coord, bohr=False, direct=False)
-                    WriteInput(input_params, os.path.join(disped_stru_job_path, "INPUT"))
-                    displaced_stru.write(os.path.join(disped_stru_job_path, stru_file))
 
-                    disped_stru_job_paths.append(disped_stru_job_path)
-                    disped_stru_cache_labels.append(os.path.join(vib_cache_dir, f"cache.{selected_atom}{DIRECTION_MAP[direction]}{step}.json"))
+        disped_stru_job_paths, disped_stru_cache_labels = prepare_abacus_vibration_analysis(input_params,
+                                                                                            stru, 
+                                                                                            work_path=work_path,
+                                                                                            abacus_inputs_dir=abacus_inputs_dir,
+                                                                                            selected_atoms=selected_atoms,
+                                                                                            stepsize=stepsize,
+                                                                                            vib_cache_dir=vib_cache_dir)
 
         # Run ABACUS calculation for all prepared jobs
         run_abacus(disped_stru_job_paths)
 
         # Collect needed forces for each ABACUS calculation and dump to cache directory used by ase.vibrations
         cache_forces_json = {"forces": {
-            "__ndarray__": [[len(selected_atoms), 3],
+            "__ndarray__": [[stru.get_natoms(), 3],
                             "float64",
                             []]
         }}
         for disped_stru_job_path, disped_stru_cache_label in list(zip(disped_stru_job_paths, disped_stru_cache_labels)):
-            force = collect_force(disped_stru_job_path)
-            cache_forces_json["forces"]["__ndarray__"][2] = list(filter_force(force, selected_atoms))
+            cache_forces_json["forces"]["__ndarray__"][2] = collect_force(disped_stru_job_path)
             with open(os.path.join(vib_cache_dir, disped_stru_cache_label), "w") as fin:
                 json.dump(cache_forces_json, fin)
 
@@ -253,6 +261,5 @@ def abacus_vibration_analysis(abacus_inputs_dir: Path,
                 'vib_free_energy': float(free_energy),
                 'vib_analysis_work_dir': Path(work_path).absolute()}
     except Exception as e:
-        print(e)
         return {'message': f"Doing vibration analysis failed: {e}"}
 

tests/integrate_test/data/ref_results.json

@@ -194,6 +194,15 @@
             "vib_free_energy": 0.1834753687243439
         }
     },
+    "test_abacus_vibration_analysis_h2_selected":
+    {
+        "result": {
+            "frequencies": [58.98644787490194, 58.98644787490194, 3035.9634063376648],
+            "zero_point_energy": 0.19551913052685363,
+            "vib_entropy": 0.0004398879187182639,
+            "vib_free_energy": 0.08145444835006974
+        }
+    },
     "test_abacus_run_md_nve":
     {
         "result": {}

tests/integrate_test/test_vibration.py

@@ -49,3 +49,31 @@ def test_abacus_vibration_analysis_h2(self):
         self.assertAlmostEqual(outputs['zero_point_energy'], ref_results['zero_point_energy'], places=4)
         self.assertAlmostEqual(outputs['vib_entropy'], ref_results['vib_entropy'], places=4)
         self.assertAlmostEqual(outputs['vib_free_energy'], ref_results['vib_free_energy'], places=4)
+
+    def test_abacus_vibration_analysis_h2_selected(self):
+        """
+        Test the abacus_vibration_analysis function for relaxed H2 molecule.
+        Only 1 atom are selected for vibration analysis.
+        """
+        test_func_name = inspect.currentframe().f_code.co_name
+        ref_results = load_test_ref_result(test_func_name)
+
+        test_work_dir = self.test_path / test_func_name
+        shutil.copytree(self.abacus_inputs_dir_h2, test_work_dir)
+        shutil.copy2(self.stru_h2_relaxed, test_work_dir / "STRU")
+
+        outputs = abacus_vibration_analysis(test_work_dir,
+                                            selected_atoms = [1],
+                                            stepsize = 0.01,
+                                            nfree = 2,
+                                            temperature=400)
+        print(outputs)
+        
+        self.assertIsInstance(outputs['vib_analysis_work_dir'], get_path_type())
+
+        for freq_output, freq_ref in zip(outputs['frequencies'], ref_results['frequencies']):
+            self.assertAlmostEqual(freq_output, freq_ref, places=2)
+
+        self.assertAlmostEqual(outputs['zero_point_energy'], ref_results['zero_point_energy'], places=4)
+        self.assertAlmostEqual(outputs['vib_entropy'], ref_results['vib_entropy'], places=4)
+        self.assertAlmostEqual(outputs['vib_free_energy'], ref_results['vib_free_energy'], places=4)