mod: root argument for OverlapCalculator

- all classes deriving from OverlapCalculator now support the 'root'
keyword
- ORCA sets appropriate iroot, even when keyword is not present in the
input
- Turbomole and Gaussian are not yet implemented

pysisyphus/calculators/DFTBp.py

@@ -1,3 +1,4 @@
+import itertools as it
 from math import ceil
 from pathlib import Path
 import re
@@ -35,13 +36,15 @@ def parse_xplusy(text):
     for i in range(states):
         block = lines[i * block_size : (i + 1) * block_size]
         _, *rest = block
-        xpy = np.array([line.split() for line in rest], dtype=float)
+        xpy = list()
+        for line in block[1:]:
+            xpy.extend(line.split())
+        xpy = np.array(xpy, dtype=float)
         xpys.append(xpy)
     return size, states, np.array(xpys)
 
 
 class DFTBp(OverlapCalculator):
-
     conf_key = "dftbp"
     _set_plans = (
         "out",
@@ -95,7 +98,7 @@ class DFTBp(OverlapCalculator):
         },
     }
 
-    def __init__(self, parameter, *args, slakos=None, root=None, **kwargs):
+    def __init__(self, parameter, *args, slakos=None, **kwargs):
         super().__init__(*args, **kwargs)
 
         assert self.mult == 1, "Open-shell not yet supported!"
@@ -107,8 +110,6 @@ def __init__(self, parameter, *args, slakos=None, root=None, **kwargs):
             f"Expected '{self.parameter}' sub-directory in '{self.slakos_prefix}' "
             "but could not find it!"
         )
-        self.root = root
-
         self.base_cmd = self.get_cmd()
         self.gen_geom_fn = "geometry.gen"
         self.inp_fn = "dftb_in.hsd"
@@ -203,8 +204,8 @@ def get_gen_str(atoms, coords):
         return gen_str
 
     @staticmethod
-    def get_excited_state_str(root, forces=False):
-        if root is None:
+    def get_excited_state_str(track, root, nroots, forces=False):
+        if root is None and (track == False):
             return ""
 
         casida_tpl = jinja2.Template(
@@ -213,7 +214,7 @@ def get_excited_state_str(root, forces=False):
             Casida {
                 NrOfExcitations = {{ nstates }}
                 Symmetry = Singlet
-                StateOfInterest = {{ root }}
+                {% if root %}StateOfInterest = {{ root }}{% endif %}
                 WriteXplusY = Yes
                 {{ es_forces }}
             }
@@ -222,7 +223,7 @@ def get_excited_state_str(root, forces=False):
         )
         es_forces = "ExcitedStateForces = Yes" if forces else ""
         es_str = casida_tpl.render(
-            nstates=root + 5,
+            nstates=nroots if nroots else root + 5,
             root=root,
             es_forces=es_forces,
         )
@@ -236,7 +237,7 @@ def prepare_input(self, atoms, coords, calc_type):
         analysis = list()
         if calc_type == "forces":
             analysis.append("CalculateForces = Yes")
-        if self.root:
+        if self.track or self.root:
             analysis.extend(("WriteEigenvectors = Yes", "EigenvectorsAsText = Yes"))
         ang_moms = self.max_ang_moms[self.parameter]
         unique_atoms = set(atoms)
@@ -265,7 +266,9 @@ def prepare_input(self, atoms, coords, calc_type):
             parameter=self.parameter,
             max_ang_moms=max_ang_moms,
             hubbard_derivs=hubbard_derivs,
-            excited_state_str=self.get_excited_state_str(self.root, es_forces),
+            excited_state_str=self.get_excited_state_str(
+                self.track, self.root, self.nroots, es_forces
+            ),
             analysis=analysis,
         )
         return inp, path

pysisyphus/calculators/Gaussian16.py

@@ -53,7 +53,6 @@ def __init__(
             for kw, option in [self.parse_keyword(kw) for kw in self.route.split()]
         }
         exc_keyword = [key for key in "td tda cis".split() if key in keywords]
-        self.root = None
         self.nstates = None
         if exc_keyword:
             self.exc_key = exc_keyword[0]

pysisyphus/calculators/ORCA.py

@@ -599,10 +599,20 @@ def __init__(
         )
 
         self.do_tddft = False
-        if "tddft" in self.blocks:
+        self.es_block_header = [key for key in ("tddft", "cis") if key in self.blocks]
+        if self.es_block_header:
+            assert len(self.es_block_header) == 1
+            self.es_block_header = self.es_block_header[0]
+        if self.es_block_header:
             self.do_tddft = True
             try:
                 self.root = int(re.search(r"iroot\s*(\d+)", self.blocks).group(1))
+                warnings.warn(
+                    f"Using root {self.root}, as specified w/ 'iroot' keyword. Please "
+                    "use the designated 'root' keyword in the future, as the 'iroot' route "
+                    "will be deprecated.",
+                    DeprecationWarning,
+                )
             except AttributeError:
                 self.log("Doing TDA/TDDFT calculation without gradient.")
         self.triplets = bool(re.search(r"triplets\s+true", self.blocks))
@@ -694,9 +704,17 @@ def prepare_input(
 
     def get_block_str(self):
         block_str = self.blocks
-        # Use the correct root if we track it
-        if self.track:
-            block_str = re.sub(r"iroot\s+(\d+)", f"iroot {self.root}", self.blocks)
+        # Use the correct root if we track and a root is supplied
+        if self.track and (self.root is not None):
+            if "iroot" in self.blocks:
+                block_str = re.sub(r"iroot\s+(\d+)", f"iroot {self.root}", self.blocks)
+            # Insert appropriate iroot keyword if not already present
+            else:
+                block_str = re.sub(
+                    f"{self.es_block_header}",
+                    f"{self.es_block_header} iroot {self.root}",
+                    self.blocks,
+                )
             self.log(f"Using iroot '{self.root}' for excited state gradient.")
         return block_str
 

pysisyphus/calculators/OverlapCalculator.py

@@ -104,6 +104,8 @@ class OverlapCalculator(Calculator):
     def __init__(
         self,
         *args,
+        root=None,
+        nroots=None,
         track=False,
         ovlp_type="tden",
         double_mol=False,
@@ -126,7 +128,13 @@ def __init__(
     ):
         super().__init__(*args, **kwargs)
 
+        self.root = root
+        self.nroots = nroots
         self.track = track
+
+        # TODO: enable this, when all calculators implement self.root & self.nroots
+        # if self.track:
+        # assert self.root <= self.nroots, "'root' must be smaller " "than 'nroots'!"
         self.ovlp_type = ovlp_type
         assert (
             self.ovlp_type in self.OVLP_TYPE_VERBOSE.keys()
@@ -219,7 +227,6 @@ def __init__(
         self.ref_cycle = 0
         self.ref_cycles = list()
         self.atoms = None
-        self.root = None
 
         if track:
             self.log(

pysisyphus/calculators/PySCF.py

@@ -39,8 +39,6 @@ def __init__(
         basis,
         xc=None,
         method="scf",
-        root=None,
-        nstates=None,
         auxbasis=None,
         keep_chk=True,
         verbose=0,
@@ -58,14 +56,8 @@ def __init__(
             self.multisteps[self.method] = ("scf", self.method)
         if self.xc and self.method != "tddft":
             self.method = "dft"
-        self.root = root
-        self.nstates = nstates
         if self.method == "tddft":
-            assert self.nstates, "nstates must be set with method='tddft'!"
-        if self.track:
-            assert self.root <= self.nstates, (
-                "'root' must be smaller " "than 'nstates'!"
-            )
+            assert self.nroots, "nroots must be set with method='tddft'!"
         self.auxbasis = auxbasis
         self.keep_chk = keep_chk
         self.verbose = int(verbose)
@@ -121,10 +113,10 @@ def _get_driver():
             mf = mp2_mf(mf)
         elif mf and (step == "tddft"):
             mf = pyscf.tddft.TDDFT(mf)
-            mf.nstates = self.nstates
+            mf.nstates = self.nroots
         elif mf and (step == "tda"):
             mf = pyscf.tddft.TDA(mf)
-            mf.nstates = self.nstates
+            mf.nstates = self.nroots
         else:
             raise Exception("Unknown method '{step}'!")
         return mf
@@ -171,8 +163,15 @@ def get_energy(self, atoms, coords, **prepare_kwargs):
 
         mol = self.prepare_input(atoms, coords)
         mf = self.run(mol, point_charges=point_charges)
+        energy = mf.e_tot
+        root = 0 if self.root is None else self.root
+        try:
+            energy = energy[root]
+        except TypeError:
+            pass
+
         results = {
-            "energy": mf.e_tot,
+            "energy": energy,
         }
         results = self.store_and_track(
             results, self.get_energy, atoms, coords, **prepare_kwargs

pysisyphus/calculators/Turbomole.py

@@ -173,7 +173,6 @@ def __init__(
         self,
         control_path=None,
         simple_input=None,
-        root=None,
         double_mol_path=None,
         cosmo_kwargs=None,
         **kwargs,
@@ -209,7 +208,6 @@ def __init__(
         # Set provided control_path or use the one generated for simple_input
         self.control_path = Path(control_path).absolute()
 
-        self.root = root
         self.double_mol_path = double_mol_path
         if self.double_mol_path:
             self.double_mol_path = Path(self.double_mol_path)

tests/test_overlap_calculator/test_calculators.py

@@ -0,0 +1,60 @@
+import pytest
+
+from pysisyphus.calculators.PySCF import PySCF
+from pysisyphus.calculators import DFTBp, ORCA
+from pysisyphus.testing import using
+
+
+_ROOT_REF_ENERGIES = {
+    (ORCA, None): -74.4893,
+    (ORCA, 2): -74.3984,
+    (PySCF, None): -74.4893,
+    (PySCF, 2): -74.3714,
+    (DFTBp, None): -4.077751,
+    (DFTBp, 2): -3.33313,
+}
+
+
+@pytest.mark.parametrize(
+    "calc_cls, calc_kwargs",
+    (
+        pytest.param(
+            ORCA,
+            {
+                "keywords": f"hf sto-3g",
+                "blocks": "%tddft tda false nroots 3 end",
+            },
+            marks=(using("orca")),
+        ),
+        pytest.param(
+            PySCF,
+            {
+                "basis": "sto3g",
+                "method": "tddft",
+                "nroots": 3,
+            },
+            marks=(using("pyscf")),
+        ),
+        pytest.param(
+            DFTBp,
+            {
+                "parameter": "mio-ext",
+                "nroots": 4,
+            },
+            marks=(using("dftbp")),
+        ),
+    ),
+)
+@pytest.mark.parametrize("root", (None, 2))
+def test_root(calc_cls, calc_kwargs, root):
+    ref_energy = _ROOT_REF_ENERGIES[(calc_cls, root)]
+    calc_kwargs.update(
+        {
+            "root": root,
+            "base_name": f"root_{root}",
+            "track": True,
+        }
+    )
+    geom = calc_cls.geom_from_fn("lib:h2o.xyz", **calc_kwargs)
+    energy = geom.energy
+    assert energy == pytest.approx(ref_energy)

tests/test_overlap_calculator/test_overlap_calculator.py

@@ -6,7 +6,6 @@
 from pysisyphus.calculators.OverlapCalculator import OverlapCalculator
 from pysisyphus.calculators.PySCF import PySCF
 from pysisyphus.helpers import geom_loader
-from pysisyphus.helpers_pure import describe
 from pysisyphus.init_logging import init_logging
 from pysisyphus.testing import using
 from pysisyphus.wavefunction.excited_states import norm_ci_coeffs, tden_overlaps