Enable total energy extractor from ABACUS output

.gitignore

@@ -157,3 +157,10 @@ cython_debug/
 #  and can be added to the global gitignore or merged into this file.  For a more nuclear
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
+
+
+# test files
+test/data/siesta/siesta_out/*
+example/siesta_io/siesta_io.ipynb
+CLAUDE.md
+playground/*

dftio/__main__.py

@@ -153,6 +153,12 @@ def main_parser() -> argparse.ArgumentParser:
         default=0,
         help="The initial band index for eigenvalues to save.(0-band_index_min) bands will be ignored!"
     )
+    parser_parse.add_argument(
+        "-energy",
+        "--energy",
+        action="store_true",
+        help="Whether to parse the total energy (Etot) from DFT output",
+    )
 
     parser_band = subparsers.add_parser(
         "band",

dftio/data/_keys.py

@@ -103,6 +103,7 @@
 
 PER_ATOM_ENERGY_KEY: Final[str] = "atomic_energy"
 TOTAL_ENERGY_KEY: Final[str] = "total_energy"
+UNCONVERGED_FRAME_INDICES_KEY: Final[str] = "unconverged_frames"
 FORCE_KEY: Final[str] = "forces"
 PARTIAL_FORCE_KEY: Final[str] = "partial_forces"
 STRESS_KEY: Final[str] = "stress"

dftio/io/abacus/abacus_parser.py

@@ -13,6 +13,9 @@
 from dftio.data import _keys
 from dftio.register import Register
 import lmdb
+import logging
+
+log = logging.getLogger(__name__)
 import pickle
 import shutil
 
@@ -414,6 +417,160 @@ def transform(self, mat, l_lefts, l_rights):
 
         return block_lefts @ mat @ block_rights.T
 
+    @staticmethod
+    def _extract_energy_from_log(loglines, mode, dump_freq=1):
+        """
+        Extract total energy from ABACUS log file.
+
+        Note that this extractor is only validated for ABACUS versions 3.9.0.
+
+        For SCF/NSCF, extracts the final total energy.
+        For MD, extracts energies at each dump interval, filtering out unconverged frames.
+        For RELAX, extracts energies for the final converged structure.
+
+        Parameters
+        ----------
+        loglines : list of str
+            Lines from the log file
+        mode : str
+            Calculation mode (scf, nscf, md, relax)
+        nframes : int, optional
+            Expected number of frames (for MD/relax validation)
+        dump_freq : int, optional
+            Dump frequency for MD calculations (default: 1)
+
+        Returns
+        -------
+        tuple or None
+            Tuple of (energy_array, unconverged_indices) where:
+            - energy_array: np.ndarray of energies in eV (filtered, converged only for MD)
+            - unconverged_indices: list of frame indices that did not converge
+            Returns None if extraction fails
+        """
+        energy = []
+
+        if mode in ["scf", "nscf"]:
+            # For SCF/NSCF, search from the end for the final energy
+            for line in reversed(loglines):
+                if "final etot is" in line:
+                    # LTS version format: "final etot is <value> eV"
+                    Etot = float(line.split()[-2])
+                    return (np.array([Etot], dtype=np.float64), [])
+                elif "TOTAL ENERGY" in line:
+                    # Develop version format
+                    Etot = float(line.split()[-2])
+                    return (np.array([Etot], dtype=np.float64), [])
+                elif "convergence has NOT been achieved!" in line or \
+                     "convergence has not been achieved" in line:
+                    # SCF did not converge
+                    return None
+            return None
+
+        elif mode == "md":
+            # For MD, extract all energies at dump intervals
+            nenergy = 0
+            for line in loglines:
+                if "final etot is" in line:
+                    if nenergy % dump_freq == 0:
+                        energy.append(float(line.split()[-2]))
+                    nenergy += 1
+                elif "!! convergence has not been achieved" in line:
+                    if nenergy % dump_freq == 0:
+                        energy.append(np.nan)
+                    nenergy += 1
+
+            if len(energy) == 0:
+                return None
+
+            # Filter out unconverged frames (NaN values) and track their indices
+            energy = np.array(energy, dtype=np.float64)
+            valid_mask = ~np.isnan(energy)
+            if not valid_mask.any():
+                return None
+
+            # Get indices of unconverged frames
+            unconverged_indices = np.where(~valid_mask)[0].tolist()
+
+            # Filter to keep only converged frames
+            energy = energy[valid_mask]
+            return (energy, unconverged_indices)
+
+        elif mode == "relax":
+            relax_success = False
+            # For RELAX, extract energy for the converged structures
+            for line in loglines:
+                if "Relaxation is converged!" in line:
+                    relax_success = True
+                if "!FINAL_ETOT_IS" in line:
+                    energy.append(float(line.split()[-2]))
+
+            if relax_success and len(energy) > 0:
+                return (np.array(energy, dtype=np.float64), [])
+            else:
+                # raise error when relaxation did not converge or no energies found
+                raise ValueError("Relaxation did not converge or no energies found.")
+
+        else:
+            return None
+
+    def get_etot(self, idx):
+        """
+        Extract total energy (Etot) from ABACUS output.
+
+        Parameters
+        ----------
+        idx : int
+            Index of the structure/trajectory
+
+        Returns
+        -------
+        dict
+            Dictionary with:
+            - _keys.TOTAL_ENERGY_KEY: energy array (shape: [1,] for SCF/NSCF, [nframes,] for MD/RELAX)
+            - _keys.UNCONVERGED_FRAME_INDICES_KEY: list of unconverged frame indices (empty if all converged)
+            Energy values are in eV.
+            Returns None if energy extraction fails or structures are unconverged.
+        """
+        mode = self.get_mode(idx)
+        logfile = "running_" + mode + ".log"
+        logpath = os.path.join(self.raw_datas[idx], "OUT.ABACUS", logfile)
+
+        # Check if log file exists
+        if not os.path.exists(logpath):
+            raise FileNotFoundError(f"Log file {logpath} does not exist.")
+
+        # Read log file
+        with open(logpath, 'r') as f:
+            loglines = f.readlines()
+
+        # Get dump frequency for MD mode
+        dump_freq = 1
+        if mode == "md":
+            input_path = os.path.join(self.raw_datas[idx], "INPUT")
+            if os.path.exists(input_path):
+                with open(input_path, 'r') as f:
+                    for line in f:
+                        if len(line) > 0 and "md_dumpfreq" in line and "md_dumpfreq" == line.split()[0]:
+                            dump_freq = int(line.split()[1])
+                            break
+
+        # Extract energy
+        result = self._extract_energy_from_log(loglines, mode, dump_freq)
+
+        if result is None:
+            return None
+
+        energy, unconverged_indices = result
+
+        # Log warning if there are unconverged frames
+        if len(unconverged_indices) > 0:
+            log.warning(f"Energy extraction: frames {unconverged_indices} did not converge")
+
+        return {
+            _keys.TOTAL_ENERGY_KEY: energy,
+            _keys.UNCONVERGED_FRAME_INDICES_KEY: unconverged_indices
+        }
+
     def get_abs_h0_folders(self, h0_root):
         # Build a map of all directory names to their full paths to avoid repeated os.walk calls
         folder_path_map = {}

dftio/io/parse.py

@@ -211,13 +211,13 @@ def check_blocks(self, idx, hamiltonian: bool=False, overlap: bool=False, densit
 
         return True
 
-    def write(self, idx, outroot, format, eigenvalue, hamiltonian, overlap, density_matrix, band_index_min, **kwargs):
+    def write(self, idx, outroot, format, eigenvalue, hamiltonian, overlap, density_matrix, band_index_min, energy=False, **kwargs):
         if format == "hdf5":
-            self.write_hdf5(idx=idx, outroot=outroot, eigenvalue=eigenvalue, hamiltonian=hamiltonian, overlap=overlap, density_matrix=density_matrix,band_index_min=band_index_min)
+            self.write_hdf5(idx=idx, outroot=outroot, eigenvalue=eigenvalue, hamiltonian=hamiltonian, overlap=overlap, density_matrix=density_matrix,band_index_min=band_index_min, energy=energy)
         elif format in ["dat", "ase"]:
-            self.write_dat(idx=idx, outroot=outroot, fmt=format, eigenvalue=eigenvalue, hamiltonian=hamiltonian, overlap=overlap, density_matrix=density_matrix,band_index_min=band_index_min)
+            self.write_dat(idx=idx, outroot=outroot, fmt=format, eigenvalue=eigenvalue, hamiltonian=hamiltonian, overlap=overlap, density_matrix=density_matrix,band_index_min=band_index_min, energy=energy)
         elif format == "lmdb":
-            self.write_lmdb(idx=idx, outroot=outroot, eigenvalue=eigenvalue, hamiltonian=hamiltonian, overlap=overlap, density_matrix=density_matrix,band_index_min=band_index_min)
+            self.write_lmdb(idx=idx, outroot=outroot, eigenvalue=eigenvalue, hamiltonian=hamiltonian, overlap=overlap, density_matrix=density_matrix,band_index_min=band_index_min, energy=energy)
         else:
             raise NotImplementedError(f"Format: {format} is not implemented!")
 
@@ -242,10 +242,10 @@ def write_struct(self, structure, out_dir, fmt='dat'):
         else:
             raise NotImplementedError(f"Format: {fmt} is not implemented!")
 
-    def write_dat(self, idx, outroot, fmt='dat', eigenvalue=False, hamiltonian=False, overlap=False, density_matrix=False, band_index_min=0):
+    def write_dat(self, idx, outroot, fmt='dat', eigenvalue=False, hamiltonian=False, overlap=False, density_matrix=False, band_index_min=0, energy=False):
         # write structure
         os.makedirs(outroot, exist_ok=True)
-       
+
         structure = self.get_structure(idx)
 
         out_dir = os.path.join(outroot, self.formula(idx=idx)+".{}".format(idx))
@@ -255,7 +255,7 @@ def write_dat(self, idx, outroot, fmt='dat', eigenvalue=False, hamiltonian=False
         # np.savetxt(os.path.join(out_dir, "positions.dat"), structure[_keys.POSITIONS_KEY].reshape(-1, 3))
         # np.savetxt(os.path.join(out_dir, "atomic_numbers.dat"), structure[_keys.ATOMIC_NUMBERS_KEY], fmt='%d')
         # np.savetxt(os.path.join(out_dir, "pbc.dat"), structure[_keys.PBC_KEY])
-        
+
         # write structure
         self.write_struct(structure, out_dir, fmt=fmt)
 
@@ -266,6 +266,26 @@ def write_dat(self, idx, outroot, fmt='dat', eigenvalue=False, hamiltonian=False
             np.save(os.path.join(out_dir, "kpoints.npy"), eigstatus[_keys.KPOINT_KEY])
             np.save(os.path.join(out_dir, "eigenvalues.npy"), eigstatus[_keys.ENERGY_EIGENVALUE_KEY])
 
+        # write energy
+        if energy:
+            if hasattr(self, 'get_etot'):
+                energy_data = self.get_etot(idx)
+                if energy_data is not None:
+                    np.savetxt(os.path.join(out_dir, "total_energy.dat"), energy_data[_keys.TOTAL_ENERGY_KEY])
+
+                    # Write unconverged frame indices if present
+                    if _keys.UNCONVERGED_FRAME_INDICES_KEY in energy_data:
+                        unconverged_indices = energy_data[_keys.UNCONVERGED_FRAME_INDICES_KEY]
+                        if len(unconverged_indices) > 0:
+                            with open(os.path.join(out_dir, "unconverged_frames.dat"), 'w') as f:
+                                f.write("# Frame indices that did not converge during MD/RELAX\n")
+                                for idx_frame in unconverged_indices:
+                                    f.write(f"{idx_frame}\n")
+                else:
+                    log.warning(f"Failed to extract energy for structure {idx}")
+            else:
+                log.warning(f"Parser does not implement get_etot method")
+
         # write blocks
         if any([hamiltonian is not None, overlap is not None, density_matrix is not None]) and any([hamiltonian, overlap, density_matrix]):
             with open(os.path.join(out_dir, "basis.dat"), 'w') as f:
@@ -279,34 +299,55 @@ def write_dat(self, idx, outroot, fmt='dat', eigenvalue=False, hamiltonian=False
                         for key_str, value in ham[i].items():
                             default_group.create_dataset(key_str, data=value)
             del ham
-            
+
             if overlap:
                 with h5py.File(os.path.join(out_dir, "overlaps.h5"), 'w') as fid:
                     for i in range(len(ovp)):
                         default_group = fid.create_group(str(i))
                         for key_str, value in ovp[i].items():
                             default_group.create_dataset(key_str, data=value)
             del ovp
-            
+
             if density_matrix:
                 with h5py.File(os.path.join(out_dir, "density_matrices.h5"), 'w') as fid:
                     for i in range(len(dm)):
                         default_group = fid.create_group(str(i))
                         for key_str, value in dm[i].items():
                             default_group.create_dataset(key_str, data=value)
-            
+
             del dm
 
         return True
 
-    def write_lmdb(self, idx, outroot, eigenvalue: bool=False, hamiltonian: bool=False, overlap: bool=False, density_matrix: bool=False,band_index_min=0):
+    def write_lmdb(self, idx, outroot, eigenvalue: bool=False, hamiltonian: bool=False, overlap: bool=False, density_matrix: bool=False,band_index_min=0, energy: bool=False):
         os.makedirs(outroot, exist_ok=True)
         out_dir = os.path.join(outroot, "data.{}.lmdb".format(os.getpid()))
         structure = self.get_structure(idx)
         if any([hamiltonian, overlap, density_matrix]):
             ham, ovp, dm = self.get_blocks(idx, hamiltonian, overlap, density_matrix)
         if eigenvalue:
             eigstatus = self.get_eigenvalue(idx=idx, band_index_min=band_index_min)
+        if energy:
+            if hasattr(self, 'get_etot'):
+                energy_data = self.get_etot(idx)
+            else:
+                energy_data = None
+                log.warning(f"Parser does not implement get_etot method")
+
+        # Build frame index mapping for energy data
+        # If there are unconverged frames, energy array will be shorter than n_frames
+        energy_frame_mapping = None
+        if energy and energy_data is not None:
+            unconverged_indices = energy_data.get(_keys.UNCONVERGED_FRAME_INDICES_KEY, [])
+            if len(unconverged_indices) > 0:
+                # Build mapping: structure_frame_idx -> energy_array_idx
+                energy_frame_mapping = {}
+                energy_idx = 0
+                n_frames_total = structure[_keys.POSITIONS_KEY].shape[0]
+                for frame_idx in range(n_frames_total):
+                    if frame_idx not in unconverged_indices:
+                        energy_frame_mapping[frame_idx] = energy_idx
+                        energy_idx += 1
 
         n_frames = structure[_keys.POSITIONS_KEY].shape[0]
         lmdb_env = lmdb.open(out_dir, map_size=1048576000000, lock=True)
@@ -321,6 +362,23 @@ def write_lmdb(self, idx, outroot, eigenvalue: bool=False, hamiltonian: bool=Fal
                 data_dict[_keys.ENERGY_EIGENVALUE_KEY] = eigstatus[_keys.ENERGY_EIGENVALUE_KEY][nf]
                 data_dict[_keys.KPOINT_KEY] = eigstatus[_keys.KPOINT_KEY]
 
+            if energy and energy_data is not None:
+                # For single structure (SCF/NSCF), energy_data has shape [1,]
+                # For trajectories (MD/RELAX), energy_data has shape [nframes,] or less if unconverged
+                if energy_data[_keys.TOTAL_ENERGY_KEY].shape[0] == 1:
+                    # Single structure case
+                    data_dict[_keys.TOTAL_ENERGY_KEY] = energy_data[_keys.TOTAL_ENERGY_KEY][0]
+                else:
+                    # Trajectory case - use mapping if unconverged frames exist
+                    if energy_frame_mapping is not None:
+                        if nf in energy_frame_mapping:
+                            energy_idx = energy_frame_mapping[nf]
+                            data_dict[_keys.TOTAL_ENERGY_KEY] = energy_data[_keys.TOTAL_ENERGY_KEY][energy_idx]
+                        # else: skip energy for unconverged frames (don't add to data_dict)
+                    else:
+                        # No unconverged frames, direct indexing
+                        data_dict[_keys.TOTAL_ENERGY_KEY] = energy_data[_keys.TOTAL_ENERGY_KEY][nf]
+
             if hamiltonian:
                 data_dict["hamiltonian"] = ham[nf]
             if overlap: