Merge pull request #248 from willwheelera/black

Black files

pyqmc/__init__.py

@@ -36,7 +36,9 @@ def gradient_generator(mol, wf, to_opt=None, **ewald_kwargs):
     )
 
 
-def default_slater(mol, mf, optimize_orbitals=False, twist=None, optimize_zeros=True, epsilon=1e-8):
+def default_slater(
+    mol, mf, optimize_orbitals=False, twist=None, optimize_zeros=True, epsilon=1e-8
+):
     """
     Construct a Slater determinant
     Args:
@@ -53,11 +55,13 @@ def default_slater(mol, mf, optimize_orbitals=False, twist=None, optimize_zeros=
             to_opt[k] = np.ones(wf.parameters[k].shape).astype(bool)
             if not optimize_zeros:
                 to_opt[k][np.abs(wf.parameters[k]) < epsilon] = False
-        
+
     return wf, to_opt
 
 
-def default_multislater(mol, mf, mc, tol=None, optimize_orbitals=False,optimize_zeros=True, epsilon=1e-8):
+def default_multislater(
+    mol, mf, mc, tol=None, optimize_orbitals=False, optimize_zeros=True, epsilon=1e-8
+):
     import numpy as np
 
     wf = MultiSlater(mol, mf, mc, tol)
@@ -229,53 +233,55 @@ def recover_pyscf(chkfile, ci_checkfile=None, cancel_outputs=True):
     import h5py
     import json
 
-    with h5py.File(chkfile,'r') as f:
-        periodic = 'a' in json.loads(f['mol'][()]).keys()
+    with h5py.File(chkfile, "r") as f:
+        periodic = "a" in json.loads(f["mol"][()]).keys()
 
     if not periodic:
         mol = pyscf.lib.chkfile.load_mol(chkfile)
-        with h5py.File(chkfile,'r') as f:
-            mo_occ_shape = f['scf/mo_occ'].shape
+        with h5py.File(chkfile, "r") as f:
+            mo_occ_shape = f["scf/mo_occ"].shape
         if cancel_outputs:
             mol.output = None
             mol.stdout = None
-        if len(mo_occ_shape)==2:
+        if len(mo_occ_shape) == 2:
             mf = pyscf.scf.UHF(mol)
-        elif len(mo_occ_shape)==1:
-            mf = pyscf.scf.ROHF(mol) if mol.spin !=0 else pyscf.scf.RHF(mol)
+        elif len(mo_occ_shape) == 1:
+            mf = pyscf.scf.ROHF(mol) if mol.spin != 0 else pyscf.scf.RHF(mol)
         else:
             raise Exception("Couldn't determine type from chkfile")
-    else: 
+    else:
         import pyscf.pbc
+
         mol = pyscf.pbc.lib.chkfile.load_cell(chkfile)
-        with h5py.File(chkfile,'r') as f:
-            has_kpts = 'mo_occ__from_list__' in f['/scf'].keys()
+        with h5py.File(chkfile, "r") as f:
+            has_kpts = "mo_occ__from_list__" in f["/scf"].keys()
             if has_kpts:
-                rhf = '000000' in f['/scf/mo_occ__from_list__/'].keys()
+                rhf = "000000" in f["/scf/mo_occ__from_list__/"].keys()
             else:
-                rhf = len(f['/scf/mo_occ'].shape)==1
+                rhf = len(f["/scf/mo_occ"].shape) == 1
         if cancel_outputs:
             mol.output = None
             mol.stdout = None
         if not rhf and has_kpts:
             mf = pyscf.pbc.scf.KUHF(mol)
         elif has_kpts:
-            mf = pyscf.pbc.scf.KROHF(mol) if mol.spin !=0 else pyscf.pbc.scf.KRHF(mol)
+            mf = pyscf.pbc.scf.KROHF(mol) if mol.spin != 0 else pyscf.pbc.scf.KRHF(mol)
         elif rhf:
-            mf = pyscf.pbc.scf.ROHF(mol) if mol.spin !=0 else pyscf.pbc.scf.RHF(mol)
+            mf = pyscf.pbc.scf.ROHF(mol) if mol.spin != 0 else pyscf.pbc.scf.RHF(mol)
         else:
             mf = pyscf.pbc.scf.UHF(mol)
     mf.__dict__.update(pyscf.scf.chkfile.load(chkfile, "scf"))
 
     if ci_checkfile is not None:
-        with h5py.File(ci_checkfile,'r') as f:
-            hci='ci/_strs' in f.keys()
-        if hci:            
+        with h5py.File(ci_checkfile, "r") as f:
+            hci = "ci/_strs" in f.keys()
+        if hci:
             mc = pyscf.hci.SCI(mol)
         else:
             import pyscf.casci
+
             mc = pyscf.casci.CASCI(mol)
-        mc.__dict__.update(pyscf.lib.chkfile.load(ci_checkfile,'ci'))
+        mc.__dict__.update(pyscf.lib.chkfile.load(ci_checkfile, "ci"))
 
         return mol, mf, mc
     return mol, mf
@@ -304,8 +310,10 @@ def read_wf(wf, wf_file):
             grp = hdf["wf"]
             for k in grp.keys():
                 new_parms = np.array(grp[k])
-                if wf.parameters[k].shape!=new_parms.shape:
-                    raise Exception(f"For wave function parameter {k}, shape in {wf_file} is {new_parms.shape}, while current shape is {wf.parameters[k].shape}")
+                if wf.parameters[k].shape != new_parms.shape:
+                    raise Exception(
+                        f"For wave function parameter {k}, shape in {wf_file} is {new_parms.shape}, while current shape is {wf.parameters[k].shape}"
+                    )
                 wf.parameters[k] = new_parms
         else:
             raise Exception("Did not find wf in hdf file")

pyqmc/dmc.py

@@ -166,13 +166,16 @@ def dmc_propagate_parallel(wf, configs, weights, client, npartitions, *args, **k
     confweight_avg = confweight / (np.mean(confweight) * npartitions)
     weight = np.array([w["weight"] for w in allresults[0]])
     weight_avg = weight / np.mean(weight)
-    block_avg = {k: np.sum(
+    block_avg = {
+        k: np.sum(
             [
                 res[k] * ww * cw
                 for res, cw, ww in zip(allresults[0], confweight_avg, weight_avg)
             ],
             axis=0,
-        ) for k in allresults[0][0].keys()}
+        )
+        for k in allresults[0][0].keys()
+    }
     block_avg["weight"] = np.mean(weight)
     return block_avg, configs, weights
 
@@ -235,7 +238,9 @@ def branch(configs, weights):
     probability = np.cumsum(weights)
     wtot = probability[-1]
     base = np.random.rand()
-    newinds = np.searchsorted(probability, (base + np.linspace(0,wtot,nconfig)) % wtot)
+    newinds = np.searchsorted(
+        probability, (base + np.linspace(0, wtot, nconfig)) % wtot
+    )
     configs.resample(newinds)
     weights.fill(wtot / nconfig)
     return configs, weights

pyqmc/ewald.py

@@ -190,7 +190,9 @@ def ei_const(self, ne):
         return -ne * self.i_sum * self.ijconst
 
     def e_single(self, ne):
-        return 0.5 * (ne - 1) * self.ijconst - self.i_sum * self.ijconst + self.squareconst
+        return (
+            0.5 * (ne - 1) * self.ijconst - self.i_sum * self.ijconst + self.squareconst
+        )
 
     def ewald_ion(self):
         r"""

pyqmc/func3d.py

@@ -242,7 +242,7 @@ def gradient_laplacian(self, rvec, r):
         Returns:
           grad, lap: (nconfig,...,3) vectors (components of laplacian d^2/dx_i^2 separately)
         """
-        grad=np.zeros(rvec.shape)
+        grad = np.zeros(rvec.shape)
         lap = np.zeros(rvec.shape)
         mask = r < self.parameters["rcut"]
         r = r[mask, np.newaxis]
@@ -258,7 +258,9 @@ def gradient_laplacian(self, rvec, r):
         d2pdz2_over_dpdz = (3 * z - 1) / (z * (z - 1))
         d2bdp2_over_dbdp = -2 * beta / (1 + beta * p)
         d2zdx2 = (1 - (rvec / r) ** 2) / (r * self.parameters["rcut"])
-        lap[mask] = grad[mask] * (d2bdp2_over_dbdp * dpdz * dzdx + d2pdz2_over_dpdz * dzdx)
+        lap[mask] = grad[mask] * (
+            d2bdp2_over_dbdp * dpdz * dzdx + d2pdz2_over_dpdz * dzdx
+        )
         lap[mask] += dbdp * dpdz * d2zdx2
         return grad, lap
 

pyqmc/optimize_ortho.py

@@ -525,7 +525,7 @@ def optimize_orthogonal(
     # warm up to equilibrate the configurations before running optimization
     if warmup_options is None:
         warmup_options = {}
- 
+
     data, coords = pyqmc.mc.vmc(
         wfs[-1],
         coords,
@@ -535,7 +535,6 @@ def optimize_orthogonal(
         **warmup_options
     )
 
-
     # One set of configurations for every wave function
     allcoords = [coords.copy() for _ in wfs[:-1]]
     dtype = np.complex if wfs[-1].iscomplex else np.float

pyqmc/recipes.py

@@ -6,11 +6,12 @@
 import scipy.stats
 import pandas as pd
 
+
 def OPTIMIZE(
     dft_checkfile,
     output,
     nconfig=1000,
-    ci_checkfile = None,
+    ci_checkfile=None,
     start_from=None,
     S=None,
     client=None,
@@ -22,19 +23,19 @@ def OPTIMIZE(
     if linemin_kws is None:
         linemin_kws = {}
 
-    target_root=0
+    target_root = 0
     if ci_checkfile is None:
         mol, mf = pyqmc.recover_pyscf(dft_checkfile)
         mc = None
     else:
         mol, mf, mc = pyqmc.recover_pyscf(dft_checkfile, ci_checkfile=ci_checkfile)
-        mc.ci=mc.ci[target_root]
+        mc.ci = mc.ci[target_root]
 
     if S is not None:
         mol = pyqmc.get_supercell(mol, np.asarray(S))
 
     wf, to_opt = pyqmc.generate_wf(
-        mol, mf, mc = mc, jastrow_kws=jastrow_kws, slater_kws=slater_kws
+        mol, mf, mc=mc, jastrow_kws=jastrow_kws, slater_kws=slater_kws
     )
     if start_from is not None:
         pyqmc.read_wf(wf, start_from)
@@ -56,20 +57,26 @@ def OPTIMIZE(
 def generate_accumulators(mol, mf, energy=True, rdm1=False, extra_accumulators=None):
     acc = {} if extra_accumulators is None else extra_accumulators
 
-    if len(mf.mo_coeff.shape)==2:
+    if len(mf.mo_coeff.shape) == 2:
         mo_coeff = [mf.mo_coeff, mf.mo_coeff]
     else:
-        mo_coeff=mf.mo_coeff
+        mo_coeff = mf.mo_coeff
 
     if energy:
-        if 'energy' in acc.keys():
+        if "energy" in acc.keys():
             raise Exception("Found energy in extra_accumulators and energy is True")
         acc["energy"] = pyqmc.EnergyAccumulator(mol)
     if rdm1:
-        if 'rdm1_up' in acc.keys() or 'rdm1_down' in acc.keys():
-            raise Exception("Found rdm1_up or rdm1_down in extra_accumulators and rdm1 is True")
-        acc['rdm1_up']=pyqmc.obdm.OBDMAccumulator(mol, orb_coeff=mo_coeff[0], spin=0),
-        acc['rdm1_down']=pyqmc.obdm.OBDMAccumulator(mol, orb_coeff=mo_coeff[1], spin=1),
+        if "rdm1_up" in acc.keys() or "rdm1_down" in acc.keys():
+            raise Exception(
+                "Found rdm1_up or rdm1_down in extra_accumulators and rdm1 is True"
+            )
+        acc["rdm1_up"] = (
+            pyqmc.obdm.OBDMAccumulator(mol, orb_coeff=mo_coeff[0], spin=0),
+        )
+        acc["rdm1_down"] = (
+            pyqmc.obdm.OBDMAccumulator(mol, orb_coeff=mo_coeff[1], spin=1),
+        )
 
     return acc
 
@@ -91,21 +98,23 @@ def VMC(
     if vmc_kws is None:
         vmc_kws = {}
 
-    target_root=0
+    target_root = 0
     if ci_checkfile is None:
         mol, mf = pyqmc.recover_pyscf(dft_checkfile)
         mc = None
     else:
         mol, mf, mc = pyqmc.recover_pyscf(dft_checkfile, ci_checkfile=ci_checkfile)
-        mc.ci=mc.ci[target_root]
+        mc.ci = mc.ci[target_root]
 
     if S is not None:
         mol = pyqmc.get_supercell(mol, np.asarray(S))
 
     if accumulators is None:
         accumulators = {}
 
-    wf, _ = pyqmc.generate_wf(mol, mf, mc = mc, jastrow_kws=jastrow_kws, slater_kws=slater_kws)
+    wf, _ = pyqmc.generate_wf(
+        mol, mf, mc=mc, jastrow_kws=jastrow_kws, slater_kws=slater_kws
+    )
 
     if start_from is not None:
         pyqmc.read_wf(wf, start_from)
@@ -141,13 +150,13 @@ def DMC(
     if dmc_kws is None:
         dmc_kws = {}
 
-    target_root=0
+    target_root = 0
     if ci_checkfile is None:
         mol, mf = pyqmc.recover_pyscf(dft_checkfile)
         mc = None
     else:
         mol, mf, mc = pyqmc.recover_pyscf(dft_checkfile, ci_checkfile=ci_checkfile)
-        mc.ci=mc.ci[target_root]
+        mc.ci = mc.ci[target_root]
 
     if S is not None:
         mol = pyqmc.get_supercell(mol, np.asarray(S))
@@ -173,26 +182,24 @@ def DMC(
     )
 
 
-
-
 def read_opt(fname):
     with h5py.File(fname) as f:
-        return pd.DataFrame({
-            'energy':f['energy'][...],
-            'iteration':f['iteration'][...],
-            'error':f['energy_error'][...],
-            'fname':[fname]*len(f['energy'])
-        })
-
+        return pd.DataFrame(
+            {
+                "energy": f["energy"][...],
+                "iteration": f["iteration"][...],
+                "error": f["energy_error"][...],
+                "fname": [fname] * len(f["energy"]),
+            }
+        )
 
 
 def read_mc_output(fname, warmup=5, reblock=16):
-    ret = {'fname':fname, 'warmup':warmup,'reblock':reblock}
+    ret = {"fname": fname, "warmup": warmup, "reblock": reblock}
     with h5py.File(fname) as f:
         for k in f.keys():
-            if 'energy' in k:
-                vals = pyqmc.reblock.reblock(f[k][warmup:],reblock)
+            if "energy" in k:
+                vals = pyqmc.reblock.reblock(f[k][warmup:], reblock)
                 ret[k] = np.mean(vals)
-                ret[k+"_err"] = scipy.stats.sem(vals)
+                ret[k + "_err"] = scipy.stats.sem(vals)
     return ret
-

tests/integration/test_shci_wf.py

@@ -49,7 +49,7 @@ def run_test():
     )
     en, err = avg(data["energytotal"][1:])
     nsigma = 4
-    assert len(wf.parameters['det_coeff']) == len(cisolver.ci)
+    assert len(wf.parameters["det_coeff"]) == len(cisolver.ci)
     assert en - nsigma * err < e_hf
     assert en + nsigma * err > ci_energy
 

tests/unit/test_accumulators.py

@@ -60,11 +60,13 @@ def test_info_functions_pbc():
     from pyscf.pbc import gto, scf
     from pyqmc.supercell import get_supercell
 
-    mol = gto.Cell(atom="He 0.00 0.00 0.00", basis="gth-szv", pseudo='gth-pade', unit="B")
+    mol = gto.Cell(
+        atom="He 0.00 0.00 0.00", basis="gth-szv", pseudo="gth-pade", unit="B"
+    )
     mol.a = 5.61 * np.eye(3)
     mol.build()
 
-    mf = scf.KRKS(mol, kpts=mol.make_kpts([2, 2, 2])) #.density_fit()
+    mf = scf.KRKS(mol, kpts=mol.make_kpts([2, 2, 2]))  # .density_fit()
     ehf = mf.kernel()
 
     supercell = get_supercell(mol, 2 * np.eye(3))