fix(dftbplus): keep cell and stress of periodic geometries

dpdata/formats/dftbplus/output.py

@@ -1,6 +1,6 @@
 from __future__ import annotations
 
-from typing import TYPE_CHECKING
+from typing import TYPE_CHECKING, Any
 
 import numpy as np
 
@@ -33,12 +33,38 @@ def read_dftb_plus(
     np.ndarray
         atomic forces
 
+    """
+    frame = parse_dftb_plus(fn_1, fn_2)
+    return frame["symbols"], frame["coords"], frame["energy"], frame["forces"]
+
+
+def parse_dftb_plus(fn_1: FileType, fn_2: FileType) -> dict[str, Any]:
+    """Read one labeled frame, including the cell and stress, from DFTB+ files.
+
+    Parameters
+    ----------
+    fn_1 : str
+        DFTB+ input file name
+    fn_2 : str
+        DFTB+ output file name
+
+    Returns
+    -------
+    dict
+        ``symbols``, ``coords``, ``energy`` and ``forces`` as returned by
+        :func:`read_dftb_plus`; ``cell``, the lattice vectors of a periodic
+        (``S`` or ``F``) GenFormat geometry; and ``stress``, the ``Total stress
+        tensor`` in Hartree/Bohr^3. ``cell`` and ``stress`` are ``None`` when
+        the files do not contain them.
+
     """
     coord = None
     symbols = None
     forces = None
     energy = None
     natoms = None
+    cell = None
+    stress = None
     with open_file(fn_1) as f:
         lines = iter(f)
         for line in lines:
@@ -58,15 +84,18 @@ def read_dftb_plus(
                 s = next(lines).split()
                 symbols.append(components[int(s[1]) - 1])
                 coord.append([float(s[2]), float(s[3]), float(s[4])])
-            if coordinate_mode == "F":
-                # Fractional GenFormat coordinates are expressed in the
-                # lattice-vector basis that follows the atom records.
+            if coordinate_mode in {"S", "F"}:
+                # Periodic GenFormat geometries end with the origin and the
+                # three lattice vectors.
                 origin = np.array([float(value) for value in next(lines).split()])
-                lattice = np.array(
+                cell = np.array(
                     [[float(value) for value in next(lines).split()] for _ in range(3)]
                 )
-                coord = (np.asarray(coord) @ lattice + origin).tolist()
-            elif coordinate_mode not in {"C", "S"}:
+                if coordinate_mode == "F":
+                    # Fractional GenFormat coordinates are expressed in the
+                    # lattice-vector basis.
+                    coord = (np.asarray(coord) @ cell + origin).tolist()
+            elif coordinate_mode != "C":
                 raise ValueError(
                     f"unsupported GenFormat coordinate mode: {coordinate_mode}"
                 )
@@ -84,10 +113,21 @@ def read_dftb_plus(
             elif line.startswith("Total energy:"):
                 s = line.split()
                 energy = float(s[2])
+            elif line.startswith("Total stress tensor"):
+                stress = np.array(
+                    [[float(value) for value in next(lines).split()] for _ in range(3)]
+                )
 
     symbols = np.array(symbols)
     forces = np.array(forces)
     coord = np.array(coord)
     assert coord.shape == forces.shape
 
-    return symbols, coord, energy, forces
+    return {
+        "symbols": symbols,
+        "coords": coord,
+        "energy": energy,
+        "forces": forces,
+        "cell": cell,
+        "stress": stress,
+    }

dpdata/plugins/dftbplus.py

@@ -3,20 +3,23 @@
 import numpy as np
 
 from dpdata.format import Format
-from dpdata.formats.dftbplus.output import read_dftb_plus
-from dpdata.unit import EnergyConversion, ForceConversion
+from dpdata.formats.dftbplus.output import parse_dftb_plus
+from dpdata.unit import EnergyConversion, ForceConversion, PressureConversion
 
 energy_convert = EnergyConversion("hartree", "eV").value()
 force_convert = ForceConversion("hartree/bohr", "eV/angstrom").value()
+pressure_convert = PressureConversion("hartree/bohr^3", "eV/angstrom^3").value()
 
 
 @Format.register("dftbplus")
 class DFTBplusFormat(Format):
-    """DFTB+ input/output pair for one labeled molecular configuration.
+    """DFTB+ input/output pair for one labeled configuration.
 
     Pass a tuple containing the DFTB+ input geometry file and output result
     file. The reader combines symbols and coordinates from the input with the
-    energy and forces from the output. See the `DFTB+ documentation
+    energy and forces from the output. A periodic GenFormat geometry (``S`` or
+    ``F``) keeps its cell, and the ``Total stress tensor`` in the output, when
+    present, is converted to the virial. See the `DFTB+ documentation
     <https://dftbplus.org/documentation>`_ for the underlying file formats.
     """
 
@@ -27,28 +30,41 @@ def from_labeled_system(self, file_paths, **kwargs):
         ----------
         file_paths : tuple
             A tuple containing the input and output file paths.
-            - Input file (file_in): Contains information about symbols and coord.
-            - Output file (file_out): Contains information about energy and force.
+            - Input file (file_in): Contains information about symbols, coord and cell.
+            - Output file (file_out): Contains information about energy, force and stress.
         **kwargs : dict
             other parameters
 
         """
         file_in, file_out = file_paths
-        symbols, coord, energy, forces = read_dftb_plus(file_in, file_out)
+        frame = parse_dftb_plus(file_in, file_out)
+        symbols = frame["symbols"]
+        coord = frame["coords"]
         last_occurrence = {v: i for i, v in enumerate(symbols)}
         atom_names = np.array(sorted(np.unique(symbols), key=last_occurrence.get))
         atom_types = np.array([np.where(atom_names == s)[0][0] for s in symbols])
         atom_numbs = np.array([np.sum(atom_types == i) for i in range(len(atom_names))])
         natoms = coord.shape[0]
 
-        return {
+        data = {
             "atom_types": atom_types,
             "atom_names": list(atom_names),
             "atom_numbs": list(atom_numbs),
             "coords": coord.reshape((1, natoms, 3)),
-            "energies": np.array([energy * energy_convert]),
-            "forces": (forces * force_convert).reshape((1, natoms, 3)),
-            "cells": np.zeros((1, 3, 3)),
+            "energies": np.array([frame["energy"] * energy_convert]),
+            "forces": (frame["forces"] * force_convert).reshape((1, natoms, 3)),
             "orig": np.zeros(3),
-            "nopbc": True,
         }
+        cell = frame["cell"]
+        if cell is None:
+            data["cells"] = np.zeros((1, 3, 3))
+            data["nopbc"] = True
+        else:
+            data["cells"] = cell.reshape((1, 3, 3))
+            if frame["stress"] is not None:
+                # DFTB+ reports the stress with the opposite sign to ASE's
+                # (see ase.calculators.dftb), so the virial is +volume * stress.
+                volume = abs(np.linalg.det(cell))
+                virial = frame["stress"] * volume * pressure_convert
+                data["virials"] = virial.reshape((1, 3, 3))
+        return data

tests/test_dftbplus.py

@@ -121,5 +121,73 @@ def test_parser_rejects_input_without_genformat_geometry(self):
             )
 
 
+class TestDftbplusPeriodic(unittest.TestCase):
+    forces_text = """Total energy: -1.0 H
+Total Forces
+1 0.0 0.0 0.0
+2 0.0 0.0 0.0
+"""
+    stress_text = """
+Total stress tensor
+1.0e-6 0.0 0.0
+0.0 2.0e-6 0.0
+0.0 0.0 3.0e-6
+"""
+
+    @staticmethod
+    def load(input_text, output_text):
+        return dpdata.LabeledSystem(
+            (StringIO(input_text), StringIO(output_text)), fmt="dftbplus"
+        )
+
+    def test_supercell_keeps_cell_and_converts_stress_to_virial(self):
+        input_text = """Geometry = GenFormat {
+2 S
+Si
+1 1 0.0 0.0 0.0
+2 1 1.0 1.0 1.0
+0.0 0.0 0.0
+5.0 0.0 0.0
+0.0 5.0 0.0
+0.0 0.0 5.0
+}
+"""
+        system = self.load(input_text, self.forces_text + self.stress_text)
+
+        self.assertFalse(system.nopbc)
+        np.testing.assert_allclose(system["cells"][0], np.eye(3) * 5.0)
+        # DFTB+ prints the stress in Hartree/Bohr^3 with the opposite sign to
+        # ASE's, so the virial is +volume * stress.
+        pressure_convert = dpdata.unit.PressureConversion(
+            "hartree/bohr^3", "eV/angstrom^3"
+        ).value()
+        np.testing.assert_allclose(
+            system["virials"][0],
+            np.diag([1.0e-6, 2.0e-6, 3.0e-6]) * 125.0 * pressure_convert,
+        )
+
+    def test_fractional_geometry_keeps_cell_without_virial(self):
+        input_text = """Geometry = GenFormat {
+2 F
+H O
+1 1 0.5 0.0 0.0
+2 2 0.0 0.5 0.0
+0.0 0.0 0.0
+2.0 0.0 0.0
+0.0 4.0 0.0
+0.0 0.0 6.0
+}
+"""
+        system = self.load(input_text, self.forces_text)
+
+        self.assertFalse(system.nopbc)
+        np.testing.assert_allclose(system["cells"][0], np.diag([2.0, 4.0, 6.0]))
+        np.testing.assert_allclose(
+            system["coords"][0], [[1.0, 0.0, 0.0], [0.0, 2.0, 0.0]]
+        )
+        # Without a stress tensor in the output there is no virial label.
+        self.assertNotIn("virials", system.data)
+
+
 if __name__ == "__main__":
     unittest.main()