update aims, librpa, and kpath linearizer

mushroom/aims/stdout.py

@@ -41,13 +41,13 @@ def __init__(self, pstdout: Path, lazy_load: bool = False):
         with open(pstdout, 'r', encoding='utf-8') as h:
             lines = h.readlines()
         _logger.info("Reading standard output from: %s", pstdout)
-        self._not_converged = True
+        self._converged = False
         self._finished = False
         if len(lines) > 1:
             testline = lines[-2].strip()
             self._finished = testline == 'Have a nice day.' or testline == '*** scf_solver: SCF cycle not converged.'
             if testline == 'Have a nice day.':
-                self._not_converged = False
+                self._converged = True
 
         self._aims_version = None
 
@@ -163,7 +163,7 @@ def debug(msg):
                 if self._scf_lines is not None:
                     self._scf_lines = self._scf_lines[:self._scf_lines.index(l) + 1]
                 self._postscf_lines = lines[i:]
-            if l.startswith("          Leaving FHI-aims.") and not self._not_converged:
+            if l.startswith("          Leaving FHI-aims.") and self._converged:
                 # in case that SCF is not converged, it will leave aims without starting postscf
                 if self._postscf_lines is not None:
                     debug("End of post-SCF")
@@ -344,7 +344,7 @@ def _handle_timing_statistics(self):
     def is_finished(self, require_converge: bool = True):
         """check if the calculation is finished successfully"""
         if require_converge:
-            return self._finished and not self._not_converged
+            return self._finished and self._converged
         return self._finished
 
     @property
@@ -495,7 +495,7 @@ def get_QP_sigx(self):
 
     get_exx = get_QP_sigx
 
-    def get_QP_bandstructure(self, kind="eqp"):
+    def get_QP_bandstructure(self, kind="qp", **kwargs):
         """get the QP band structure
 
         Args:
@@ -508,10 +508,12 @@ def get_QP_bandstructure(self, kind="eqp"):
         """
         d, kpts = self.get_QP_result()
         if kind in ["eqp", "eps"]:
-            return BandStructure(d[kind], d["occ"], unit='ev'), kpts
+            return BandStructure(d[kind], d["occ"], unit='ev', **kwargs), kpts
+        elif kind in ["qp"]:
+            return BandStructure(d["eqp"], d["occ"], unit='ev', **kwargs), kpts
         elif kind in ["exx", "hf"]:
             eigen = d["eps"] - d["vxc"] + d["exx"]
-            return BandStructure(eigen, d["occ"], unit='ev'), kpts
+            return BandStructure(eigen, d["occ"], unit='ev', **kwargs), kpts
         # TODO: which case does the occupation number refer to when
         #       there is a band reordering?
         raise ValueError("Use eqp/eps for QP/KS and exx/hf for EXX/HF band structure")

mushroom/aims/workflow.py

@@ -1,7 +1,39 @@
 # -*- coding: utf-8 -*-
 """workflow objects for FHI-aims"""
-import os
+import pathlib
 
+from mushroom.aims.stdout import StdOut
 
-class Aims:
-    """class to define a FHI-aims calculation"""
+
+class AimsWorkdir:
+    """Class to handle a FHI-aims calculation working folder"""
+
+    def __init__(self, dirpath: str, aimsout: str = "aims.out", aimserr: str = None):
+        dirpath = pathlib.Path(dirpath).absolute()
+        self.dirpath = dirpath
+        self.path_control = dirpath / "control.in"
+        self.path_geometry = dirpath / "geometry.in"
+        self.path_aimsout = dirpath / aimsout
+        self.path_aimserr = None
+        if aimserr is not None:
+            self.path_aimserr = dirpath / aimserr
+
+        self.band_output_glob = "band[0-9][0-9][0-9][0-9].out"
+
+        self.efermi = None
+        self.nspin = None
+        self.nelect = None
+
+        # Objects to handle inputs and outputs
+        self._control = None
+        self._geometry = None
+        self._aimsout = None
+
+    def load_output(self):
+        self._aimsout = StdOut(self.path_aimsout, lazy_load=False)
+
+    def is_finished(self):
+        return self._aimsout.is_finished(False)
+
+    def is_converged(self):
+        return self._aimsout.is_finished(True)

mushroom/core/bs.py

@@ -6,7 +6,7 @@
 from itertools import permutations, product
 from typing import Sequence, Union
 from copy import deepcopy
-from numbers import Real
+from numbers import Real, Number
 
 import numpy as np
 
@@ -993,17 +993,19 @@ def resolve(self, eres=0.01):
         """
         raise NotImplementedError
 
-    def __add__(self, y: Real):
-        if isinstance(y, Real):
-            if y == 0.0:
-                return self
-            newbs = deepcopy(self)
-            newbs._eigen = newbs._eigen + y
-            if newbs._efermi is not None:
-                newbs._efermi += y
-            newbs._bandedge_calculated = False
-            return newbs
-        raise TypeError("expected a Real, got {}".format(type(y)))
+    def __add__(self, y: Union[float, int]):
+        newbs = deepcopy(self)
+        newbs += y
+        return newbs
+
+    def __iadd__(self, y):
+        if y == 0.0:
+            return self
+        self._eigen += y
+        if self._efermi is not None:
+            self._efermi += y
+        self._bandedge_calculated = False
+        return self
 
     def __sub__(self, y):
         if isinstance(y, type(self)):

mushroom/core/cell.py

@@ -20,7 +20,7 @@
 except ImportError:
     symprec = 1.0E-5
 from mushroom.core.constants import PI, SQRT3
-from mushroom.core.elements import get_atomic_number, get_atomic_mass
+from mushroom.core.elements import get_atomic_number, get_atomic_mass, element_symbols
 from mushroom.core.unit import LengthUnit
 from mushroom.core.pkg import detect
 from mushroom.core.crystutils import (get_latt_consts_from_latt_vecs,
@@ -80,7 +80,7 @@ class Cell(LengthUnit):
 
     _err = CellError
     _dtype = 'float64'
-    avail_exporters = ['vasp', 'abi', 'json', 'qe', 'qe_alat', 'aims']
+    avail_exporters = ['vasp', 'abi', 'json', 'qe', 'qe_alat', 'aims', "abacus"]
     avail_readers = ['vasp', 'json', 'cif', 'aims']
 
     def __init__(self, latt: Latt3T3, atms: Sequence[str], posi: Sequence[RealVec3D],
@@ -142,6 +142,7 @@ def __init__(self, latt: Latt3T3, atms: Sequence[str], posi: Sequence[RealVec3D]
             'qe': self.export_qe,
             'qe_alat': self.export_qe_alat,
             'aims': self.export_aims,
+            'abacus': self.export_abacus,
         }
         assert all([x in self.exporters for x in self.avail_exporters])
 
@@ -955,6 +956,64 @@ def export(self, output_format: str, scale: float = 1.0) -> str:
             raise ValueError("Unsupported export:", output_format)
         return e(scale=scale)
 
+    def export_abacus(self, scale: float = 1.0):
+        """Export in ABACUS STRU format
+
+        The header for pseudopotential and numerical orbitals are exported
+        as a placeholder.
+
+        ATOMIC_SPECIES
+        C  12.011  C.upf
+
+        NUMERICAL_ORBITAL
+        C.orb
+
+        Caveats:
+        - Setting magnetic moment is not supported
+        - Relax flags are forced to 0, i.e. not relaxed
+        """
+        slist = []
+        slist.append("ATOMIC_SPECIES")
+        for symbol in self.atom_types:
+            slist.append("{s:<2s}  {m:f}  {s:s}.upf".format(s=symbol, m=get_atomic_mass(symbol)))
+        slist.append("")
+
+        slist.append("NUMERICAL_ORBITAL")
+        for symbol in self.atom_types:
+            slist.append("{s:s}.orb".format(s=element_symbols[get_atomic_number(symbol)]))
+        slist.append("")
+
+        slist.append("LATTICE_CONSTANT")
+        if self.unit == "ang":
+            slist.append("1.889726164")
+        else:
+            slist.append("1.000000000")
+        slist.append("")
+
+        slist.append("LATTICE_VECTORS")
+        for i in range(3):
+            slist.append("{:19.10f} {:19.10f} {:19.10f}".format(*(self.latt[i, :] * scale)))
+        slist.append("")
+
+        slist.append("ATOMIC_POSITIONS")
+        if self.coord_sys == "C":
+            slist.append("Cartesian")
+        else:
+            slist.append("Direct")
+        for symbol in self.atom_types:
+            indices = self.get_atm_indices(symbol)
+            slist.append(symbol)
+            slist.append("0.0")  # hard-coded magnetic moment
+            slist.append("{:d}".format(len(indices)))
+            for index in indices:
+                # dyn = self.sd_flag(ia=index)
+                # aflag = " ".join({True: "1", False: "0"}[d] for d in dyn)
+                aflag = "0 0 0"
+                slist.append("{:19.10f} {:19.10f} {:19.10f}  {:s}"
+                             .format(*self.posi[index, :], aflag))
+
+        return "\n".join(slist)
+
     def export_qe(self, scale: float = 1.0) -> str:
         """Export in Quantum Espresso format
         """

mushroom/core/kpoints.py

@@ -21,42 +21,63 @@
 _logger = loggers["kpoints"]
 
 
+def check_on_line_segment(x_0, x_st, x_ed, atol=1e-5):
+    """Check if point x_0 lies on the line segment specified by points x_st and x_ed
+
+    Args:
+        x_0, x_st, x_ed (1d-array-like)
+
+    Return:
+        True if x_0 lies between x_st and x_ed, otherwise False
+    """
+    x0st = np.subtract(x_0, x_st)
+    x0ed = np.subtract(x_0, x_ed)
+    dot = np.dot(x0st, x0ed)  # to check if dot is between the two points
+    distance = np.linalg.norm(np.cross(x0st, x0ed)) / np.linalg.norm(np.subtract(x_ed, x_st))
+    return np.isclose(distance, 0.0, atol=atol) and dot <= 0.0
+
+
 class KPathLinearizer:
     """object to manipulate path in reciprocal k space
 
     Special kpoints are recognized automatically.
 
     Args:
-        kpts (list): the coordinates of k points
+        kpts (list): the fractional coordinates of k points
         recp_latt (3x3 array): the reciprocal lattice vectors, [b1, b2, b3].
             If parsed, it will be used to convert the kpts to Cartisian coordiantes.
     """
 
     def __init__(self, kpts, recp_latt=None, unify_x: bool = False):
+        self._use_cartesian = False
+        self._recp_latt = None
         self._nkpts = len(kpts)
         if np.shape(kpts) != (self._nkpts, 3):
             raise ValueError("bad shape of parsed kpoints")
-        self.kpts = np.array(kpts)
+        self._kpts = np.array(kpts)
         if recp_latt is not None:
-            self.kpts = np.matmul(kpts, recp_latt)
+            self._kpts = np.matmul(kpts, recp_latt)
+            self._recp_latt = recp_latt
+            self._use_cartesian = True
         self._ksegs = None
         self._unify_x = unify_x
         self._x = None
         self._special_x = None
         self._index_special_x = None
+        self._xmax_non_unifty = None
         self._find_ksegs()
         _logger.debug("Found segments: %r", self._ksegs)
 
     def _find_ksegs(self):
-        self._ksegs = find_k_segments(self.kpts)
+        self._ksegs = find_k_segments(self._kpts)
 
     def _compute_x(self):
         """calculate 1d abscissa of kpoints"""
         xs = []
         ispks = []
         accumu_l = 0.0
         for i, (st, ed) in enumerate(self._ksegs):
-            l = np.linalg.norm(self.kpts[st, :] - self.kpts[ed, :])
+            l = np.linalg.norm(self._kpts[st, :] - self._kpts[ed, :])
             # remove duplicate
             if st not in ispks and st - 1 not in ispks:
                 ispks.append(st)
@@ -67,10 +88,11 @@ def _compute_x(self):
                 if st == self._ksegs[i - 1][1]:
                     skip = 1
             x = accumu_l + np.linalg.norm(
-                self.kpts[st:ed + 1, :] - self.kpts[st, :], axis=1)[skip:]
+                self._kpts[st:ed + 1, :] - self._kpts[st, :], axis=1)[skip:]
             xs.extend(x)
             accumu_l += l
         self._x = np.array(xs)
+        self._xmax_non_unifty = self._x[-1]
         if self._unify_x:
             self._x /= self._x[-1]
         self._special_x = self._x[ispks]
@@ -91,6 +113,42 @@ def special_x(self):
             self._compute_x()
         return self._special_x
 
+    @property
+    def X(self):
+        """alias to special_x"""
+        return self.special_x
+
+    def locate(self, kpts):
+        """Locate the 1d coordinates of k-points on the path
+
+        Args:
+            kpts (2d-array, shape N*3): k-points in fractional coordinates
+
+        Return:
+            a list of float lists or None
+        """
+        if self._x is None:
+            self._compute_x()
+        xs_all = []
+        if len(kpts) == 0:
+            return xs_all
+        if self._use_cartesian:
+            kpts = np.matmul(kpts, self._recp_latt)
+        kpts = np.array(kpts)
+        for k in kpts:
+            xs = None
+            for (st, ed) in self._ksegs:
+                if check_on_line_segment(k, self._kpts[st], self._kpts[ed]):
+                    x = np.linalg.norm(k - self._kpts[st])
+                    if self._unify_x:
+                        # print(x, self._xmax_non_unifty)
+                        x /= self._xmax_non_unifty
+                    if xs is None:
+                        xs = []
+                    xs.append(x + self.x[st])
+            xs_all.append(xs)
+        return xs_all
+
 
 class MPGrid:
     """Monkhorst-Pack kpoint mesh

mushroom/core/test/test_kpoints.py

@@ -122,9 +122,23 @@ def test_special_x(self):
         for kpts, spec_x in zip(self.valid_kpts, self.valid_spec_xs):
             kp = KPathLinearizer(kpts)
             self.assertEqual(len(spec_x), len(kp.special_x))
+            # alias of special_x
+            self.assertEqual(len(spec_x), len(kp.X))
             for x1, x2 in zip(spec_x, kp.special_x):
                 self.assertAlmostEqual(x1, x2)
 
+    def test_locate(self):
+        kpts_band = [
+            [0.0, 0.0, 0.0], [0.5, 0.5, 0.5]
+        ]
+        kp = KPathLinearizer(kpts_band, recp_latt=[[1, 0, 0], [0, 1, 0], [0, 0, 1]], unify_x=True)
+        xs = kp.locate([])
+        self.assertListEqual(xs, [])
+        xs = kp.locate([[0.25, 0.25, 0.25]])
+        xs_ref = [[0.5,]]
+        for x, x_ref in zip(xs, xs_ref):
+            self.assertListEqual(x, x_ref)
+
 
 class test_mpgrid(ut.TestCase):
     """test the Monkhorst-Pack grid generation"""