Merge branch 'abinit:develop' into embedding_ifc

abipy/dynamics/cpx.py

@@ -23,7 +23,9 @@
 import numpy as np
 import pandas as pd
 import dataclasses
+import abipy.core.abinit_units as abu
 
+from pathlib import Path
 from monty.functools import lazy_property
 from monty.bisect import find_le
 from abipy.core.mixins import TextFile, NotebookWriter
@@ -33,8 +35,7 @@
     get_fig_plotly, add_plotly_fig_kwargs, PlotlyRowColDesc, plotly_klabels, plotly_set_lims)
 
 
-
-def parse_file_with_blocks(filepath: PathLike, block_size: int) -> tuple[int, np.ndarray, list]:
+def parse_file_with_blocks(filepath: PathLike, block_len: int) -> tuple[int, np.ndarray, list]:
     """
     Parse a QE file whose format is:
 
@@ -47,19 +48,19 @@ def parse_file_with_blocks(filepath: PathLike, block_size: int) -> tuple[int, np
          0.78345550025202E+01     0.67881458009071E+01     0.12273445304341E+01
          0.20762325213327E+01     0.59955571558384E+01     0.47335647385293E+01
 
-    i.e. a header followed by block_size lines
+    i.e. a header followed by block_len lines
 
     Args:
         filepath: Filename
-        block_size: Number of lines in block.
+        block_len: Number of lines in block.
 
     Return: (number_of_steps, array_to_be_reshaped, list_of_headers)
     """
     nsteps, data, headers = 0, [], []
     with open(filepath, "rt") as fh:
         for il, line in enumerate(fh):
             toks = line.split()
-            if il % (block_size + 1) == 0:
+            if il % (block_len + 1) == 0:
                 #assert len(toks) == 2
                 headers.append([int(toks[0]), float(toks[1])])
                 nsteps += 1
@@ -112,7 +113,7 @@ class EvpFile(TextFile, NotebookWriter):
     for each time step of a CP simulation.
     The evp file has the following format:
 
-    # nfi tps(ps) ekinc T cell(K) Tion(K) etot enthal econs
+    # nfi tps(ps) ekinc Tcell(K) Tion(K) etot enthal econs
           0            0          300   -2099.0164    7.4066063   -2091.6098    19363.353    2188.2406    5.0978969
           10         0.01    557.43004   -2105.3429    13.762216   -2091.5807    16917.626    2188.2406    5.0978969
     ...
@@ -250,3 +251,147 @@ def traj_to_qepos(traj_filepath: PathLike, pos_filepath: PathLike) -> None:
             fh.write(str(it) + '\n')
             for ia in range(natoms):
                 fh.write(str(pos_tac[it,ia,0]) + ' ' + str(pos_tac[it,ia,1]) + ' ' + str(pos_tac[it,ia,2]) + '\n')
+
+
+class Qe2Extxyz:
+    """
+    Convert QE/CP output files into ASE extended xyz format.
+
+    Example:
+
+        from abipy.dynamics.cpx import Qe2Extxyz
+        converter = Qe2Extxyz.from_input("cp.in", code="cp")
+        converter.write_xyz("extended.xyz", take_every=1)
+    """
+
+    @classmethod
+    def from_input(cls, qe_input, code, prefix=None):
+        """
+        Build object from QE/CP input assuming all the other output files
+        are located in the same directory with the given prefix.
+        """
+        qe_input = Path(str(qe_input)).absolute()
+        directory = qe_input.cwd()
+        from ase.io.espresso import read_fortran_namelist
+        with open(qe_input, "rt") as fh:
+            sections, card_lines = read_fortran_namelist(fh)
+
+        # https://www.quantum-espresso.org/Doc/INPUT_CP.html
+        control = sections["control"]
+        calculation = control["calculation"]
+        #outdir = control["outdir"]
+
+        default_prefix = "cp" if code == "cp" else "pwscf"
+        if prefix is None:
+            prefix = control.get("prefix", default_prefix)
+
+        pos_filepath = directory / f"{prefix}.pos"
+        for_filepath = directory / f"{prefix}.for"
+        cell_filepath = directory / f"{prefix}.cel"
+        evp_filepath = directory / f"{prefix}.evp"
+        str_filepath = directory / f"{prefix}.str"
+        if not str_filepath.exists():
+            # File with stresses is optional.
+            str_filepath = None
+
+        return cls(code, qe_input, pos_filepath, for_filepath, cell_filepath, evp_filepath, str_filepath=str_filepath)
+
+    def __init__(self, code, qe_input, pos_filepath, for_filepath, cell_filepath, evp_filepath, str_filepath=None):
+        """
+        Args:
+            code: Either "pwscf" or "cp".
+            qe_input: QE/CP input file
+            pos_filepath: File with cartesian positions.
+            for_filepath: File with cartesian forces.
+            cell_filepath: File with lattice vectors.
+            evp_filepath: File with energies
+        """
+        print(f"""
+Reading symbols from: {qe_input=}
+Reading positions from: {pos_filepath=}
+Reading forces from: {for_filepath=}
+Reading cell from: {cell_filepath=}
+Reading energies from: {evp_filepath=}
+Reading stresses from: {str_filepath=}
+        """)
+
+        # Parse input file to get initial_atoms and symbols
+        from ase.io.espresso import read_espresso_in
+        with open(qe_input, "rt") as fh:
+            self.initial_atoms = read_espresso_in(fh)
+            natom = len(self.initial_atoms)
+            print("initial_atoms:", self.initial_atoms)
+
+        # FIXME: Clarify units for positions, forces and stresses.
+        # NB: in QE/CP lengths are in Bohr
+        # but CP uses Hartree for energies whereas PW uses Rydberg
+        e_fact = 1.0 if code == "cp" else 0.5
+
+        # Parse lattice vectors NB: in the cel file, lattice vectors are stored as column-vectors.
+        # so he have to transpose the data as ASE uses row-vectors.
+        cell_nsteps, cells, cell_headers = parse_file_with_blocks(cell_filepath, 3)
+        cells = np.reshape(cells, (cell_nsteps, 3, 3)) * abu.Bohr_Ang
+        self.cells = np.transpose(cells, (0, 2, 1)).copy()
+
+        # Get energies from the evp file and convert from Ha to eV.
+        with EvpFile(evp_filepath) as evp:
+            self.energies = evp.df["etot"].values * (e_fact * abu.Ha_to_eV)
+            if len(self.energies) != cell_nsteps:
+                raise RuntimeError(f"{len(energies)=} != {cell_nsteps=}")
+
+        # Parse Cartesian positions (Bohr --> Ang)
+        pos_nsteps, pos_cart, pos_headers = parse_file_with_blocks(pos_filepath, natom)
+        self.pos_cart = np.reshape(pos_cart, (pos_nsteps, natom, 3)) * abu.Bohr_Ang
+        if pos_nsteps != cell_nsteps:
+            raise RuntimeError(f"{pos_nsteps=} != {cell_nsteps=}")
+
+        # Parse Cartesian forces (Ha/Bohr --> eV/Ang)
+        for_nsteps, forces_list, for_headers = parse_file_with_blocks(for_filepath, natom)
+        self.forces_step = np.reshape(forces_list, (for_nsteps, natom, 3)) * (e_fact * abu.Ha_to_eV / abu.Bohr_Ang)
+        if for_nsteps != cell_nsteps:
+            raise RuntimeError(f"{for_nsteps=} != {cell_nsteps=}")
+
+        # Optionally, parse stress. In ASE, stress is eV/Ang^3
+        self.stresses = None
+        if str_filepath is not None:
+            str_nsteps, stresses, str_headers = parse_file_with_blocks(str_filepath, 3)
+            self.stresses = np.reshape(stresses, (str_nsteps, 3, 3)) * (e_fact * abu.Ha_to_eV / abu.Bohr_Ang**3)
+            if str_nsteps != cell_nsteps:
+                raise RuntimeError(f"{str_nsteps=} != {cell_nsteps=}")
+
+        self.nsteps = cell_nsteps
+
+    def write_xyz(self, xyz_filepath, take_every=1, pbc=True) -> None:
+        """
+        Write results in ASE extended xyz format.
+
+        Args:
+            xyz_filepath: Name of the XYZ file.
+            take_every: Used to downsample the trajectory.
+        """
+        print(f"Writing results in extended xyz format in: {xyz_filepath}")
+        from ase.io import write
+        with open(xyz_filepath, "wt") as fh:
+            for istep, atoms in enumerate(self.yield_atoms(pbc)):
+                if istep % take_every != 0: continue
+                write(fh, atoms, format='extxyz', append=True)
+
+    def yield_atoms(self, pbc):
+        """Yields ASE atoms along the trajectory."""
+        from ase.atoms import Atoms
+        from ase.calculators.singlepoint import SinglePointCalculator
+        for istep in range(self.nsteps):
+            atoms = Atoms(symbols=self.initial_atoms.symbols,
+                          positions=self.pos_cart[istep],
+                          cell=self.cells[istep],
+                          pbc=pbc,
+                          )
+
+            # Attach calculator with results.
+            atoms.calc = SinglePointCalculator(atoms,
+                                               energy=self.energies[istep],
+                                               free_energy=self.energies[istep],
+                                               forces=self.forces_step[istep],
+                                               stress=self.stresses[istep] if self.stresses is not None else None,
+                                               )
+            yield atoms

abipy/examples/plot/plot_lruj.py

@@ -4,6 +4,8 @@
 =====================
 
 This example shows how to parse the output file produced by lruj and plot the results
+
+
 See also https://docs.abinit.org/tutorial/lruj/#43-execution-of-the-lruj-post-processinng-utility
 """
 from abipy.electrons.lruj import LrujAnalyzer, LrujResults
@@ -17,8 +19,7 @@
 #%%
 # Plot the fits.
 
-lr.plot(degrees="all", insetdegree=4, ptcolor0='blue',
-         ptitle="Hello World", fontsize=9)
+lr.plot(degrees="all", insetdegree=4, ptcolor0='blue', ptitle="Hello World", fontsize=9)
 
 #filepaths = [
 #    "tlruj_2.o_DS1_LRUJ.nc",

abipy/flowtk/tasks.py

@@ -653,6 +653,7 @@ def from_user_config(cls) -> TaskManager:
         _USER_CONFIG_TASKMANAGER = cls.from_file(path)
         return _USER_CONFIG_TASKMANAGER
 
+
     @classmethod
     def from_file(cls, filepath: str) -> TaskManager:
         """Read the configuration parameters from the Yaml file filepath."""