PDFT Chkfile Support (#61)

* fix text_mcpdft for ci

* fix test_grad_mcpdft

* wait

* update and seems to work

* add chkfile tests

* fix sideffect to test

* fix flake

* add e_mcscf to pdft h5 chk

* maybe working for lpdft

* not working still

* add hard links for pdft

* add lpdft chkfile support

* format

* fix tests for current PySCF version

pyscf/grad/lpdft.py

@@ -652,6 +652,8 @@ def get_Aop_Adiag(
 
         hop, Adiag = newton_casscf.gen_g_hop(fcasscf, mo, ci, eris, verbose)[2:]
 
+        # TODO: (MR Hennefarth) This is not safe way to check if something is a
+        # mixed solver, can probably do better than this...
         if hasattr(self.base, "_irrep_slices"):
             for ham_slice in ham_od:
                 ham_slice[np.diag_indices_from(ham_slice)] = 0.0

pyscf/mcpdft/chkfile.py

@@ -0,0 +1,130 @@
+#!/usr/bin/env python
+# Copyright 2014-2024 The PySCF Developers. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+#
+# Author: Matthew Hennefath <[email protected]>
+
+import h5py
+from pyscf.lib.chkfile import load_mol, load, dump
+from pyscf.lib import H5FileWrap
+from pyscf.mcscf.addons import StateAverageMixFCISolver
+
+
+def load_pdft(chkfile):
+    return load_mol(chkfile), load(chkfile, "pdft")
+
+
+def dump_mcpdft(
+    mc,
+    chkfile=None,
+    key="pdft",
+    e_tot=None,
+    e_ot=None,
+    e_states=None,
+    e_mcscf=None,
+    mcscf_key="mcscf",
+):
+    """Save MC-PDFT calculation results in chkfile"""
+    if chkfile is None:
+        chkfile = mc.chkfile
+    if e_tot is None:
+        e_tot = mc.e_tot
+    if e_ot is None:
+        e_ot = mc.e_ot
+
+    if h5py.is_hdf5(chkfile):
+        mode = "a"
+    else:
+        mode = "w"
+
+    with H5FileWrap(chkfile, mode) as fh5:
+        if mode == "a" and key in fh5:
+            del fh5[key]
+
+        def store(subkey, val):
+            if val is not None:
+
+                fh5[key + "/" + subkey] = val
+
+        store("e_tot", e_tot)
+        store("e_ot", e_ot)
+        store("e_states", e_states)
+        store("e_mcscf", e_mcscf)
+
+        # Now lets link
+        if mcscf_key in fh5:
+            hard_links = {
+                "mo_coeff": "mo_coeff",
+                "ci__from_list__": "ci__from_list__",
+                "ci": "ci",
+            }
+
+            for s, d in hard_links.items():
+                if s in fh5[mcscf_key]:
+                    fh5[key + "/" + d] = fh5[mcscf_key + "/" + s]
+
+
+def dump_lpdft(
+    mc,
+    chkfile=None,
+    key="pdft",
+    e_tot=None,
+    e_states=None,
+    e_mcscf=None,
+    ci=None,
+    mcscf_key="mcscf",
+):
+    """Save L-PDFT calculation results in chkfile"""
+    if chkfile is None:
+        chkfile = mc.chkfile
+
+    if e_tot is None:
+        e_tot = mc.e_tot
+
+    if e_states is None:
+        e_states = mc.e_states
+
+    if e_mcscf is None:
+        e_mcscf = mc.e_mcscf
+
+    if h5py.is_hdf5(chkfile):
+        mode = "a"
+
+    else:
+        mode = "w"
+
+    mixed_ci = isinstance(mc.fcisolver, StateAverageMixFCISolver) and ci is not None
+
+    with H5FileWrap(chkfile, mode) as fh5:
+        if mode == "a" and key in fh5:
+            del fh5[key]
+
+        def store(subkey, val):
+            if val is not None:
+                fh5[key + "/" + subkey] = val
+
+        store("e_tot", e_tot)
+        store("e_states", e_states)
+        store("e_mcscf", e_mcscf)
+        if not mixed_ci:
+            store("ci", ci)
+
+        if mcscf_key in fh5:
+            hard_links = {"mo_coeff": "mo_coeff"}
+            for s, d in hard_links.items():
+                if s in fh5[mcscf_key]:
+                    fh5[key + "/" + d] = fh5[mcscf_key + "/" + s]
+
+    if mixed_ci:
+        dump(chkfile, "pdft/ci", ci)