fix optimized DF-CCSD(T) (#78)

Author: fishjojo

examples/cc/02-df-ccsd_t.py

@@ -0,0 +1,47 @@
+"""
+Nuclear gradient of density-fitting CCSD(T)
+
+Reference:
+CCSD(T) total energy: -76.30586034
+Nuclear gradient:
+[[ 5.81461842e-02  6.30418580e-17  8.76964385e-15]
+ [-2.90730921e-02 -3.65174463e-17 -1.27052571e-01]
+ [-2.90730921e-02 -2.65244117e-17  1.27052571e-01]]
+"""
+import jax
+from pyscfad import gto, scf
+from pyscfad.cc import dfccsd
+from pyscfad import config
+
+# Setting `pyscfad_moleintor_opt` to `True` will use the
+# efficient back-propagation CPU implementation. However, only
+# 1st order derivative is available.
+config.update("pyscfad_moleintor_opt", True)
+config.update("pyscfad_scf_implicit_diff", True)
+config.update("pyscfad_ccsd_implicit_diff", True)
+
+mol = gto.Mole()
+mol.atom = '''
+    O  0.000000  0.000000  0.000000
+    H  0.758602  0.000000  0.504284
+    H  0.758602  0.000000  -0.504284
+'''
+mol.basis = 'aug-ccpvtz'
+mol.verbose = 4
+mol.incore_anyway = True
+mol.max_memory = 16000
+mol.build(trace_exp=False, trace_ctr_coeff=False)
+
+def energy(mol):
+    mf = scf.RHF(mol).density_fit()
+    mf.kernel()
+    mycc = dfccsd.RCCSD(mf)
+    eris = mycc.ao2mo()
+    mycc.kernel(eris=eris)
+    et = mycc.ccsd_t(eris=eris)
+    return mycc.e_tot + et
+
+e, jac = jax.value_and_grad(energy)(mol)
+print(f"CCSD(T) total energy: {e:.8f}")
+print("Nuclear gradient:")
+print(jac.coords)

pyscfad/backend/_jax/pytree.py

@@ -110,11 +110,12 @@ class PytreeNodeMeta(type):
     def __new__(mcls, name, bases, dct, **kwargs):
         cls = super().__new__(mcls, name, bases, dct, **kwargs)
 
-        _dynamic_attr = set()
-        for base in cls.__mro__:
+        # preserve the order of the attributes
+        _dynamic_attr = []
+        for base in reversed(cls.__mro__):
             if hasattr(base, '_dynamic_attr'):
-                _dynamic_attr |= set(base._dynamic_attr)
-        _dynamic_attr = tuple(_dynamic_attr)
+                _dynamic_attr.extend(base._dynamic_attr)
+        _dynamic_attr = tuple(dict.fromkeys(_dynamic_attr))
 
         def _flatten(obj, keys=(), with_keys=False):
             if keys:

pyscfad/cc/ccsd.py

@@ -13,8 +13,8 @@
 from pyscfad.tools.linear_solver import gen_gmres
 
 # attributes explicitly appearing in :fun:`update_amps` are dynamic
-ERI_Tracers = {'fock', 'mo_energy',
-               'oooo', 'ovoo', 'ovov', 'oovv', 'ovvo', 'ovvv', 'vvvv'}
+ERI_Tracers = ('fock', 'mo_energy',
+               'oooo', 'ovoo', 'ovov', 'oovv', 'ovvo', 'ovvv', 'vvvv')
 
 def _converged_iter(amp, mycc, eris):
     t1, t2 = mycc.vector_to_amplitudes(amp)

pyscfad/cc/ccsd_t.py

@@ -1,16 +1,14 @@
 import ctypes
 import numpy
 from jax import custom_vjp
-from jax.tree_util import tree_flatten, tree_unflatten
+from jax.tree_util import tree_flatten_with_path, tree_unflatten
 from pyscf.lib import (
     prange_tril,
     num_threads,
     current_memory,
-#    load_library
 )
 from pyscf.cc import ccsd_t as pyscf_ccsd_t
 from pyscfad.lib import logger
-#libcc = load_library('libcc')
 from pyscfadlib import libcc_vjp as libcc
 
 def kernel(mycc, eris, t1=None, t2=None, verbose=logger.NOTE):
@@ -36,24 +34,32 @@ def _ccsd_t_kernel_fwd(eris, t1, t2):
     def _ccsd_t_kernel_bwd(res, ybar):
         eris, t1, t2 = res
 
-        leaves, tree = tree_flatten(eris)
-        assert len(leaves) == 9
-        shapes = [leaf.shape for leaf in leaves]
-        del leaves
+        # TODO clean up tree unflatten
+        path_vals, treedef = tree_flatten_with_path(eris)
+        keys = [item[0][0].name for item in path_vals]
+        shapes = [item[1].shape for item in path_vals]
+        path_vals = None
 
         t1_bar, t2_bar, fock_bar, mo_energy_bar,\
             ovoo_bar, ovov_bar, ovvv_bar = _ccsd_t_energy_vjp(eris, t1, t2, ybar, max_memory)
 
-        leaves = [fock_bar,
-                  mo_energy_bar,
-                  numpy.zeros(shapes[2]),
-                  ovoo_bar,
-                  ovov_bar,
-                  numpy.zeros(shapes[5]),
-                  numpy.zeros(shapes[6]),
-                  ovvv_bar,
-                  numpy.zeros(shapes[8])]
-        eris_bar = tree_unflatten(tree, leaves)
+        leaves = [None] * len(keys)
+        key_to_bar = {
+            'fock': fock_bar,
+            'mo_energy': mo_energy_bar,
+            'ovoo': ovoo_bar,
+            'ovov': ovov_bar,
+            'ovvv': ovvv_bar,
+        }
+
+        for k, val in key_to_bar.items():
+            leaves[keys.index(k)] = val
+
+        for i, leaf in enumerate(leaves):
+            if leaf is None:
+                leaves[i] = numpy.zeros(shapes[i])
+
+        eris_bar = tree_unflatten(treedef, leaves)
         return eris_bar, t1_bar, t2_bar
 
     _ccsd_t_kernel.defvjp(_ccsd_t_kernel_fwd, _ccsd_t_kernel_bwd)