add: 'get_stored_wavefunction()' method for ORCA and XTB

See Issue #260 on Github.

added parsing of XTB wavefunction w/ proper amount of electrons
and contraction coefficient fixing. wtf XTB.

mod: deprecated json_dump in favor of wavefunction_dump in ORCA

pysisyphus/calculators/Calculator.py

@@ -17,6 +17,7 @@
 from pysisyphus.constants import BOHR2ANG
 from pysisyphus.helpers import geom_loader
 from pysisyphus.linalg import finite_difference_hessian
+from pysisyphus.wavefunction import Wavefunction
 
 
 KeepKind = Enum("KeepKind", ["ALL", "LATEST", "NONE"])
@@ -37,7 +38,6 @@ def __post_init__(self):
 
 
 class Calculator:
-
     conf_key = None
     _set_plans = []
 
@@ -238,6 +238,12 @@ def restore_org_hessian(self):
         self.get_hessian = self._org_get_hessian
         self.hessian_kind = HessKind["ORG"]
 
+    def load_wavefunction_from_file(self, fn, **kwargs):
+        return Wavefunction.from_file(fn, **kwargs)
+
+    def get_stored_wavefunction(self, **kwargs):
+        raise Exception("Not implemented!")
+
     def make_fn(self, name, counter=None, return_str=False):
         """Make a full filename.
 

pysisyphus/calculators/ORCA.py

@@ -367,14 +367,14 @@ def get_exc_ens_fosc(wf_fn, cis_fn, log_fn):
 
 
 class ORCA(OverlapCalculator):
-
     conf_key = "orca"
     _set_plans = (
         "gbw",
         "out",
         "cis",
         "densities",
         ("molden", "mwfn_wf"),
+        "json",
     )
 
     def __init__(
@@ -384,7 +384,8 @@ def __init__(
         gbw=None,
         do_stable=False,
         numfreq=False,
-        json_dump=True,
+        json_dump=None,
+        wavefunction_dump=True,
         **kwargs,
     ):
         """ORCA calculator.
@@ -412,7 +413,9 @@ def __init__(
             before every calculation.
         numfreq : bool, optional
             Use numerical frequencies instead of analytical ones.
-        json_dump : bool, optional
+        json_dump : bool, optional, deprecated
+            Use 'wavefunction_dump' instead.
+        wavefunction_dump : bool, optional
             Whether to dump the wavefunction to JSON via orca_2json. The JSON can become
             very large in calculations comprising many basis functions.
         """
@@ -423,7 +426,13 @@ def __init__(
         self.gbw = gbw
         self.do_stable = bool(do_stable)
         self.freq_keyword = "numfreq" if numfreq else "freq"
-        self.json_dump = bool(json_dump)
+        if json_dump is not None:
+            warnings.warn(
+                "Use of 'json_dump' is deprecated! Use 'wavefunction_dump' instead!",
+                DeprecationWarning,
+            )
+            wavefunction_dump = json_dump
+        self.wavefunction_dump = bool(wavefunction_dump)
 
         assert ("pal" not in keywords) and ("nprocs" not in blocks), (
             "PALn/nprocs not " "allowed! Use 'pal: n' in the 'calc' section instead."
@@ -610,11 +619,14 @@ def get_hessian(self, atoms, coords, **prepare_kwargs):
 
         inp = self.prepare_input(atoms, coords, calc_type, **prepare_kwargs)
         results = self.run(inp, calc="hessian")
-        # results = self.store_and_track(
-        # results, self.get_hessian, atoms, coords, **prepare_kwargs
-        # )
+        results = self.store_and_track(
+            results, self.get_hessian, atoms, coords, **prepare_kwargs
+        )
         return results
 
+    def get_stored_wavefunction(self, **kwargs):
+        return self.load_wavefunction_from_file(self.json, **kwargs)
+
     def run_calculation(self, atoms, coords, **prepare_kwargs):
         """Basically some kind of dummy method that can be called
         to execute ORCA with the stored cmd of this calculator."""
@@ -634,7 +646,7 @@ def run_after(self, path):
             self.popen(cmd, cwd=path)
             shutil.copy(path / "orca.molden.input", path / "orca.molden")
 
-        if self.json_dump: 
+        if self.wavefunction_dump:
             # Will silently fail with ECPs
             cmd = "orca_2json orca"
             proc = self.popen(cmd, cwd=path)

pysisyphus/calculators/XTB.py

@@ -21,13 +21,8 @@
 
 
 class XTB(Calculator):
-
     conf_key = "xtb"
-    _set_plans = (
-        "charges",
-        "json",
-        "xtbrestart",
-    )
+    _set_plans = ("charges", "json", "xtbrestart", "molden")
 
     def __init__(
         self,
@@ -42,6 +37,7 @@ def __init__(
         topo=None,
         topo_update=None,
         quiet=False,
+        wavefunction_dump=False,
         **kwargs,
     ):
         """XTB calculator.
@@ -74,6 +70,8 @@ def __init__(
             Mememory per core in MB.
         quiet : bool, optional
             Suppress creation of log files.
+        wavefunction_dump : bool
+            Whether to dump a molden file.
         """
         super().__init__(**kwargs)
 
@@ -92,6 +90,7 @@ def __init__(
         self.topo = topo
         self.topo_update = topo_update
         self.quiet = quiet
+        self.wavefunction_dump = wavefunction_dump
 
         self.topo_used = 0
         self.xtbrestart = None
@@ -109,12 +108,12 @@ def __init__(
             "xtbopt.xyz",
             "g98.out",
             "xtb.trj",
-            # "json:xtbout.json",
             "charges:charges",
             "xcontrol",
+            "molden:molden.input",
         )
         if self.restart:
-            self.to_keep += ("xtbrestart", )
+            self.to_keep += ("xtbrestart",)
         if self.quiet:
             self.to_keep = ()
 
@@ -193,6 +192,8 @@ def prepare_add_args(self, xcontrol=None):
         elif self.alpb:
             alpb = f"--alpb {self.alpb}".split()
             add_args = add_args + alpb
+        if self.wavefunction_dump:
+            add_args = add_args + ["--molden"]
         # Select parametrization
         gfn = ["--gfnff"] if self.gfn == "ff" else f"--gfn {self.gfn}".split()
         add_args = add_args + gfn
@@ -238,6 +239,9 @@ def get_hessian(self, atoms, coords, **prepare_kwargs):
         results = self.run(inp, **kwargs)
         return results
 
+    def get_stored_wavefunction(self, **kwargs):
+        return self.load_wavefunction_from_file(self.molden, xtb_nuc_charges=True)
+
     def run_calculation(self, atoms, coords, **prepare_kwargs):
         self.prepare_input(atoms, coords, "calculation", **prepare_kwargs)
         inp = self.prepare_coords(atoms, coords)
@@ -377,7 +381,7 @@ def parse_hessian(self, path):
         with open(path / "hessian") as handle:
             text = handle.read()
         hessian = np.array(text.split()[1:], dtype=float)
-        coord_num = int(hessian.size ** 0.5)
+        coord_num = int(hessian.size**0.5)
         hessian = hessian.reshape(coord_num, coord_num)
         energy = self.parse_energy(path)
         results = {