From 1f9f657b4ccd62699eccfab023139c1953c83623 Mon Sep 17 00:00:00 2001 From: Oskar Weser Date: Wed, 20 Nov 2024 13:02:32 -0500 Subject: [PATCH] apply `ruff format` purely cosmetic PEP8 conformance --- docs/source/conf.py | 14 +- example/kbe_polyacetylene.py | 30 +- example/molbe_dmrg_block2.py | 85 +- example/molbe_h8_chemical_potential.py | 37 +- example/molbe_h8_density_matching.py | 37 +- example/molbe_hexene_oneshot_uccsd.py | 20 +- example/molbe_io_fcidump.py | 9 +- example/molbe_octane.py | 21 +- example/molbe_octane_get_rdms.py | 12 +- example/molbe_oneshot_rbe_hcore.py | 54 +- example/molbe_oneshot_rbe_qmmm-fromchk.py | 48 +- example/molbe_oneshot_ube_qmmm.py | 48 +- example/molbe_ppp.py | 17 +- kbe/_opt.py | 71 +- kbe/autofrag.py | 2285 +++++++++++++++------ kbe/chain.py | 295 +-- kbe/fragment.py | 102 +- kbe/helper.py | 12 +- kbe/lo.py | 289 ++- kbe/lo_k.py | 97 +- kbe/misc.py | 77 +- kbe/pbe.py | 491 +++-- kbe/pfrag.py | 284 ++- kbe/solver.py | 15 +- molbe/_opt.py | 230 ++- molbe/autofrag.py | 261 ++- molbe/be_parallel.py | 415 ++-- molbe/be_var.py | 10 +- molbe/eri_onthefly.py | 102 +- molbe/external/ccsd_rdm.py | 52 +- molbe/external/cphf_utils.py | 347 ++-- molbe/external/cpmp2_utils.py | 301 +-- molbe/external/jac_utils.py | 97 +- molbe/external/lo_helper.py | 36 +- molbe/external/optqn.py | 313 +-- molbe/external/uccsd_eri.py | 38 +- molbe/external/unrestricted_utils.py | 68 +- molbe/fragment.py | 139 +- molbe/helper.py | 254 ++- molbe/lchain.py | 255 ++- molbe/lo.py | 314 +-- molbe/mbe.py | 381 ++-- molbe/misc.py | 153 +- molbe/pfrag.py | 235 ++- molbe/rdm.py | 225 +- molbe/solver.py | 644 +++--- molbe/ube.py | 220 +- setup.py | 25 +- tests/chem_dm_kBE_test.py | 94 +- tests/chempot_molBE_test.py | 61 +- tests/dm_molBE_test.py | 43 +- tests/dmrg_molBE_test.py | 60 +- tests/eri_onthefly_test.py | 30 +- tests/hf-in-hf_BE_test.py | 59 +- tests/ube-oneshot_test.py | 117 +- 55 files changed, 6314 insertions(+), 3715 deletions(-) diff --git a/docs/source/conf.py b/docs/source/conf.py index 7d9d865b..bd42be57 100644 --- a/docs/source/conf.py +++ b/docs/source/conf.py @@ -8,24 +8,24 @@ import os import sys -sys.path.insert(0, os.path.abspath('../..')) -project = 'QuEmb' -copyright = '2024, Oinam Romesh Meitei' -author = 'Oinam Romesh Meitei' +sys.path.insert(0, os.path.abspath("../..")) + +project = "QuEmb" +copyright = "2024, Oinam Romesh Meitei" +author = "Oinam Romesh Meitei" # -- General configuration --------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#general-configuration -extensions = ['sphinx.ext.autodoc', 'sphinx.ext.napoleon','sphinx_rtd_theme'] +extensions = ["sphinx.ext.autodoc", "sphinx.ext.napoleon", "sphinx_rtd_theme"] napoleon_google_docstring = False napoleon_include_init_with_doc = True napoleon_numpy_docstring = True exclude_patterns = [] - # -- Options for HTML output ------------------------------------------------- # https://www.sphinx-doc.org/en/master/usage/configuration.html#options-for-html-output -html_theme = 'sphinx_rtd_theme' +html_theme = "sphinx_rtd_theme" diff --git a/example/kbe_polyacetylene.py b/example/kbe_polyacetylene.py index 0532815f..95004dfc 100644 --- a/example/kbe_polyacetylene.py +++ b/example/kbe_polyacetylene.py @@ -9,18 +9,18 @@ kpt = [1, 1, 3] cell = gto.Cell() -a = 8. -b = 8. -c = 2.455 *2. +a = 8.0 +b = 8.0 +c = 2.455 * 2.0 lat = numpy.eye(3) -lat[0,0] = a -lat[1,1] = b -lat[2,2] = c +lat[0, 0] = a +lat[1, 1] = b +lat[2, 2] = c cell.a = lat -cell.atom=''' +cell.atom = """ H 1.4285621630072645 0.0 -0.586173422487319 C 0.3415633681566205 0.0 -0.5879921146011252 H -1.4285621630072645 0.0 0.586173422487319 @@ -29,11 +29,11 @@ C 0.3415633681566205 0.0 1.867007885398875 H -1.4285621630072645 0.0 3.041173422487319 C -0.3415633681566205 0.0 3.0429921146011254 -''' +""" -cell.unit='Angstrom' -cell.basis = 'sto-3g' -cell.verbose=0 +cell.unit = "Angstrom" +cell.basis = "sto-3g" +cell.verbose = 0 cell.build() kpts = cell.make_kpts(kpt, wrap_around=True) @@ -47,11 +47,9 @@ kpoint_energy = kmf.kernel() # Define fragment in the supercell -kfrag = fragpart(be_type='be2', mol=cell, - kpt=kpt, frozen_core=True) +kfrag = fragpart(be_type="be2", mol=cell, kpt=kpt, frozen_core=True) # Initialize BE -mykbe = BE(kmf, fobj, - kpts=kpts) +mykbe = BE(kmf, fobj, kpts=kpts) # Perform BE density matching -mykbe.optimize(solver='CCSD') +mykbe.optimize(solver="CCSD") diff --git a/example/molbe_dmrg_block2.py b/example/molbe_dmrg_block2.py index 09d6ae02..8f49dde9 100644 --- a/example/molbe_dmrg_block2.py +++ b/example/molbe_dmrg_block2.py @@ -18,8 +18,8 @@ for a in seps: # Hartree-Fock serves as the starting point for all BE calculations: mol = gto.M() - mol.atom = [['H', (0.,0.,i*a)] for i in range(8)] - mol.basis = 'sto-3g' + mol.atom = [["H", (0.0, 0.0, i * a)] for i in range(8)] + mol.basis = "sto-3g" mol.charge = 0 mol.spin = 0 mol.build() @@ -43,21 +43,21 @@ # any clear advantage to using any one scheme over another, # the Pipek-Mezey scheme continues to be the most popular. With # BE-DMRG, localization takes place prior to fragmentation: - fobj = fragpart(be_type='be1', mol=mol) + fobj = fragpart(be_type="be1", mol=mol) mybe = BE( mf, fobj, - lo_method='pipek-mezey', # or 'lowdin', 'iao', 'boys' - pop_method='lowdin' # or 'meta-lowdin', 'mulliken', 'iao', 'becke' - ) + lo_method="pipek-mezey", # or 'lowdin', 'iao', 'boys' + pop_method="lowdin", # or 'meta-lowdin', 'mulliken', 'iao', 'becke' + ) # Next, run BE-DMRG with default parameters and maxM=100. mybe.oneshot( - solver='block2', # or 'DMRG', 'DMRGSCF', 'DMRGCI' - scratch=scratch, # Scratch dir for fragment DMRG - maxM=100, # Max fragment bond dimension - force_cleanup=True, # Remove all fragment DMRG tmpfiles - ) + solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' + scratch=scratch, # Scratch dir for fragment DMRG + maxM=100, # Max fragment bond dimension + force_cleanup=True, # Remove all fragment DMRG tmpfiles + ) bedmrg_ecorr.append(mybe.ebe_tot - mf.e_tot) # Setting `force_cleanup=True` will clean the scratch directory after each @@ -70,31 +70,26 @@ # Finally, plot the resulting potential energy curves: fig, ax = plt.subplots() -ax.plot(seps, fci_ecorr, 'o-', linewidth=1, label='FCI') -ax.plot(seps, ccsd_ecorr, 'o-', linewidth=1, label='CCSD') -ax.plot(seps, ccsdt_ecorr, 'o-', linewidth=1, label='CCSD(T)') -ax.plot(seps, bedmrg_ecorr, 'o-', linewidth=1, label='BE1-DMRG') +ax.plot(seps, fci_ecorr, "o-", linewidth=1, label="FCI") +ax.plot(seps, ccsd_ecorr, "o-", linewidth=1, label="CCSD") +ax.plot(seps, ccsdt_ecorr, "o-", linewidth=1, label="CCSD(T)") +ax.plot(seps, bedmrg_ecorr, "o-", linewidth=1, label="BE1-DMRG") ax.legend() -plt.savefig(os.path.join(scratch, f'BEDMRG_H8_PES{num_points}.png')) +plt.savefig(os.path.join(scratch, f"BEDMRG_H8_PES{num_points}.png")) # (See ../quemb/example/figures/BEDMRG_H8_PES20.png for an example.) # For larger fragments, you'll want greater control over the fragment # DMRG calculations. Using the same setup as above for a single geometry: mol = gto.M() -mol.atom = [['H', (0.,0.,i*1.2)] for i in range(8)] -mol.basis = 'sto-3g' +mol.atom = [["H", (0.0, 0.0, i * 1.2)] for i in range(8)] +mol.basis = "sto-3g" mol.charge = 0 mol.spin = 0 mol.build() -fobj = fragpart(be_type='be2', mol=mol) -mybe = BE( - mf, - fobj, - lo_method='pipek-mezey', - pop_method='lowdin' - ) +fobj = fragpart(be_type="be2", mol=mol) +mybe = BE(mf, fobj, lo_method="pipek-mezey", pop_method="lowdin") # We automatically construct the fragment DMRG schedules based on user keywords. The following # input, for example, yields a 60 sweep schedule which uses the two-dot algorithm from sweeps 0-49, @@ -102,34 +97,34 @@ # vector procedure, along with a few other tweaks: mybe.optimize( - solver='block2', # or 'DMRG', 'DMRGSCF', 'DMRGCI' - scratch=scratch, # Scratch dir for fragment DMRG - startM=20, # Initial fragment bond dimension (1st sweep) - maxM=200, # Maximum fragment bond dimension - max_iter=60, # Max number of sweeps - twodot_to_onedot=50, # Sweep num to switch from two- to one-dot algo. - max_mem=40, # Max memory (in GB) allotted to fragment DMRG - max_noise=1e-3, # Max MPS noise introduced per sweep - min_tol=1e-8, # Tighest Davidson tolerance per sweep - block_extra_keyword=['fiedler'], # Specify orbital reordering algorithm - force_cleanup=True, # Remove all fragment DMRG tmpfiles + solver="block2", # or 'DMRG', 'DMRGSCF', 'DMRGCI' + scratch=scratch, # Scratch dir for fragment DMRG + startM=20, # Initial fragment bond dimension (1st sweep) + maxM=200, # Maximum fragment bond dimension + max_iter=60, # Max number of sweeps + twodot_to_onedot=50, # Sweep num to switch from two- to one-dot algo. + max_mem=40, # Max memory (in GB) allotted to fragment DMRG + max_noise=1e-3, # Max MPS noise introduced per sweep + min_tol=1e-8, # Tighest Davidson tolerance per sweep + block_extra_keyword=["fiedler"], # Specify orbital reordering algorithm + force_cleanup=True, # Remove all fragment DMRG tmpfiles only_chem=True, ) # Or, alternatively, we can construct a full schedule by hand: -schedule={ - 'scheduleSweeps': [0, 10, 20, 30, 40, 50], # Sweep indices - 'scheduleMaxMs': [25, 50, 100, 200, 500, 500], # Sweep maxMs - 'scheduleTols': [1e-5,1e-5, 1e-6, 1e-6, 1e-8, 1e-8], # Sweep Davidson tolerances - 'scheduleNoises': [0.01, 0.01, 0.001, 0.001, 1e-4, 0.0], # Sweep MPS noise +schedule = { + "scheduleSweeps": [0, 10, 20, 30, 40, 50], # Sweep indices + "scheduleMaxMs": [25, 50, 100, 200, 500, 500], # Sweep maxMs + "scheduleTols": [1e-5, 1e-5, 1e-6, 1e-6, 1e-8, 1e-8], # Sweep Davidson tolerances + "scheduleNoises": [0.01, 0.01, 0.001, 0.001, 1e-4, 0.0], # Sweep MPS noise } # and pass it to the fragment solver through `schedule_kwargs`: mybe.optimize( - solver='block2', + solver="block2", scratch=scratch, schedule_kwargs=schedule, - block_extra_keyword=['fiedler'], + block_extra_keyword=["fiedler"], force_cleanup=True, only_chem=True, ) @@ -138,5 +133,5 @@ # and `[scratch]/dmrg.out`, which are the fragment DMRG inputs and outputs, respectively, used # by `block2`. -#NOTE: Parameters in `schedule_kwargs` will overwrite any other DMRG kwargs. -#NOTE: The DMRG schedule kwargs and related syntax follows the standard notation used in block2. \ No newline at end of file +# NOTE: Parameters in `schedule_kwargs` will overwrite any other DMRG kwargs. +# NOTE: The DMRG schedule kwargs and related syntax follows the standard notation used in block2. diff --git a/example/molbe_h8_chemical_potential.py b/example/molbe_h8_chemical_potential.py index ab148dcf..f340a149 100644 --- a/example/molbe_h8_chemical_potential.py +++ b/example/molbe_h8_chemical_potential.py @@ -1,11 +1,12 @@ # Illustrates a simple molecular BE calculation with chemical # potential matching -from pyscf import gto,scf, fci +from pyscf import gto, scf, fci from molbe import BE, fragpart # PySCF HF generated mol & mf (molecular desciption & HF object) -mol = gto.M(atom=''' +mol = gto.M( + atom=""" H 0. 0. 0. H 0. 0. 1. H 0. 0. 2. @@ -14,7 +15,10 @@ H 0. 0. 5. H 0. 0. 6. H 0. 0. 7. -''',basis='sto-3g', charge=0) +""", + basis="sto-3g", + charge=0, +) mf = scf.RHF(mol) mf.conv_tol = 1e-12 @@ -23,39 +27,38 @@ # Perform PySCF FCI to get reference energy mc = fci.FCI(mf) fci_ecorr = mc.kernel()[0] - mf.e_tot -print(f'*** FCI Correlation Energy: {fci_ecorr:>14.8f} Ha', flush=True) +print(f"*** FCI Correlation Energy: {fci_ecorr:>14.8f} Ha", flush=True) # Perform BE calculations with different fragment schemes: # Define BE1 fragments -fobj = fragpart(be_type='be1', mol=mol) +fobj = fragpart(be_type="be1", mol=mol) # Initialize BE mybe = BE(mf, fobj) # Perform chemical potential optimization -mybe.optimize(solver='FCI', only_chem=True) +mybe.optimize(solver="FCI", only_chem=True) # Compute BE error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (fci_ecorr - be_ecorr)*100./fci_ecorr -print(f'*** BE1 Correlation Energy Error (%) : {err_:>8.4f} %') +err_ = (fci_ecorr - be_ecorr) * 100.0 / fci_ecorr +print(f"*** BE1 Correlation Energy Error (%) : {err_:>8.4f} %") # Define BE2 fragments -fobj = fragpart(be_type='be2', mol=mol) +fobj = fragpart(be_type="be2", mol=mol) mybe = BE(mf, fobj) -mybe.optimize(solver='FCI', only_chem=True) +mybe.optimize(solver="FCI", only_chem=True) # Compute BE error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (fci_ecorr - be_ecorr)*100./fci_ecorr -print(f'*** BE2 Correlation Energy Error (%) : {err_:>8.4f} %') +err_ = (fci_ecorr - be_ecorr) * 100.0 / fci_ecorr +print(f"*** BE2 Correlation Energy Error (%) : {err_:>8.4f} %") # Define BE3 fragments -fobj = fragpart(be_type='be3', mol=mol) +fobj = fragpart(be_type="be3", mol=mol) mybe = BE(mf, fobj) -mybe.optimize(solver='FCI', only_chem=True) +mybe.optimize(solver="FCI", only_chem=True) # Compute BE error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (fci_ecorr - be_ecorr)*100./fci_ecorr -print(f'*** BE3 Correlation Energy Error (%) : {err_:>8.4f} %') - +err_ = (fci_ecorr - be_ecorr) * 100.0 / fci_ecorr +print(f"*** BE3 Correlation Energy Error (%) : {err_:>8.4f} %") diff --git a/example/molbe_h8_density_matching.py b/example/molbe_h8_density_matching.py index ec529ac2..a647a7ea 100644 --- a/example/molbe_h8_density_matching.py +++ b/example/molbe_h8_density_matching.py @@ -1,11 +1,12 @@ # Illustrates a simple molecular BE calculation with BE # density matching between edge & centers of fragments. -from pyscf import gto,scf, fci +from pyscf import gto, scf, fci from molbe import BE, fragpart # PySCF HF generated mol & mf (molecular desciption & HF object) -mol = gto.M(atom=''' +mol = gto.M( + atom=""" H 0. 0. 0. H 0. 0. 1. H 0. 0. 2. @@ -14,7 +15,10 @@ H 0. 0. 5. H 0. 0. 6. H 0. 0. 7. -''',basis='sto-3g', charge=0) +""", + basis="sto-3g", + charge=0, +) mf = scf.RHF(mol) mf.conv_tol = 1e-12 @@ -23,39 +27,38 @@ # Perform PySCF FCI to get reference energy mc = fci.FCI(mf) fci_ecorr = mc.kernel()[0] - mf.e_tot -print(f'*** FCI Correlation Energy: {fci_ecorr:>14.8f} Ha', flush=True) +print(f"*** FCI Correlation Energy: {fci_ecorr:>14.8f} Ha", flush=True) # Perform BE calculations with different fragment schemes: # Define BE1 fragments -fobj = fragpart(be_type='be1', mol=mol) +fobj = fragpart(be_type="be1", mol=mol) # Initialize BE mybe = BE(mf, fobj) # Density matching in BE -mybe.optimize(solver='FCI') +mybe.optimize(solver="FCI") # Compute BE error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (fci_ecorr - be_ecorr)*100./fci_ecorr -print(f'*** BE1 Correlation Energy Error (%) : {err_:>8.4f} %') +err_ = (fci_ecorr - be_ecorr) * 100.0 / fci_ecorr +print(f"*** BE1 Correlation Energy Error (%) : {err_:>8.4f} %") # Define BE2 fragments -fobj = fragpart(be_type='be2', mol=mol) +fobj = fragpart(be_type="be2", mol=mol) mybe = BE(mf, fobj) -mybe.optimize(solver='FCI') +mybe.optimize(solver="FCI") # Compute BE error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (fci_ecorr - be_ecorr)*100./fci_ecorr -print(f'*** BE2 Correlation Energy Error (%) : {err_:>8.4f} %') +err_ = (fci_ecorr - be_ecorr) * 100.0 / fci_ecorr +print(f"*** BE2 Correlation Energy Error (%) : {err_:>8.4f} %") # Define BE3 fragments -fobj = fragpart(be_type='be3', mol=mol) +fobj = fragpart(be_type="be3", mol=mol) mybe = BE(mf, fobj) -mybe.optimize(solver='FCI') +mybe.optimize(solver="FCI") # Compute BE error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (fci_ecorr - be_ecorr)*100./fci_ecorr -print(f'*** BE3 Correlation Energy Error (%) : {err_:>8.4f} %') - +err_ = (fci_ecorr - be_ecorr) * 100.0 / fci_ecorr +print(f"*** BE3 Correlation Energy Error (%) : {err_:>8.4f} %") diff --git a/example/molbe_hexene_oneshot_uccsd.py b/example/molbe_hexene_oneshot_uccsd.py index 86fd96cd..af859970 100644 --- a/example/molbe_hexene_oneshot_uccsd.py +++ b/example/molbe_hexene_oneshot_uccsd.py @@ -8,27 +8,29 @@ import sys # Set up scratch directory settings -#be_var.SCRATCH='{scratch location}' -#be_var.CREATE_SCRATCH_DIR=True +# be_var.SCRATCH='{scratch location}' +# be_var.CREATE_SCRATCH_DIR=True # Give path to structure xyz file -structure = 'data/hexene.xyz' +structure = "data/hexene.xyz" -#Build PySCF molecule object +# Build PySCF molecule object mol = gto.M() mol.atom = structure -mol.basis = 'sto-3g' -mol.charge = -1; mol.spin = 1 +mol.basis = "sto-3g" +mol.charge = -1 +mol.spin = 1 mol.build() -#Run UHF with PySCF -mf = scf.UHF(mol); mf.kernel() +# Run UHF with PySCF +mf = scf.UHF(mol) +mf.kernel() # Specify number of processors nproc = 1 # Initialize fragments without frozen core approximation at BE2 level -fobj = fragpart(frag_type='autogen', be_type='be2', mol = mol, frozen_core=False) +fobj = fragpart(frag_type="autogen", be_type="be2", mol=mol, frozen_core=False) # Initialize UBE mybe = UBE(mf, fobj) diff --git a/example/molbe_io_fcidump.py b/example/molbe_io_fcidump.py index f0eddc2f..317061e1 100644 --- a/example/molbe_io_fcidump.py +++ b/example/molbe_io_fcidump.py @@ -6,15 +6,18 @@ from molbe import fragpart from molbe.misc import * from molbe import be_var -be_var.PRINT_LEVEL=3 + +be_var.PRINT_LEVEL = 3 # Read in molecular integrals expressed in libint basis ordering # numpy.loadtxt takes care of the input under the hood -mol, mf = libint2pyscf("data/octane.xyz", "data/hcore_libint_octane.dat", "STO-3G", hcore_skiprows=1) +mol, mf = libint2pyscf( + "data/octane.xyz", "data/hcore_libint_octane.dat", "STO-3G", hcore_skiprows=1 +) mf.kernel() # Construct fragments for BE -fobj = fragpart(be_type='be2', mol=mol) +fobj = fragpart(be_type="be2", mol=mol) oct_be = BE(mf, fobj) # Write out fcidump file for each fragment diff --git a/example/molbe_octane.py b/example/molbe_octane.py index 87fc53f4..67a7c35a 100644 --- a/example/molbe_octane.py +++ b/example/molbe_octane.py @@ -4,7 +4,8 @@ from molbe import fragpart, BE # Perform pyscf HF calculation to get mol & mf objects -mol = gto.M(atom=''' +mol = gto.M( + atom=""" C 0.4419364699 -0.6201930287 0.0000000000 C -0.4419364699 0.6201930287 0.0000000000 H -1.0972005331 0.5963340874 0.8754771384 @@ -31,7 +32,10 @@ H -0.9171145792 -4.5073104916 0.8797333088 H 0.3671153250 -5.3316378285 0.0000000000 H -0.3671153250 5.3316378285 0.0000000000 -''',basis='sto-3g', charge=0) +""", + basis="sto-3g", + charge=0, +) mf = scf.RHF(mol) @@ -40,22 +44,21 @@ # Perform CCSD calculation to get reference energy for comparison mc = cc.CCSD(mf, frozen=8) -mc.verbose=0 +mc.verbose = 0 ccsd_ecorr = mc.kernel()[0] -print(f'*** CCSD Correlation Energy: {ccsd_ecorr:>14.8f} Ha', flush=True) +print(f"*** CCSD Correlation Energy: {ccsd_ecorr:>14.8f} Ha", flush=True) # initialize fragments (use frozen core approximation) -fobj = fragpart(be_type='be2', mol=mol, frozen_core=True) +fobj = fragpart(be_type="be2", mol=mol, frozen_core=True) # Initialize BE mybe = BE(mf, fobj) # Perform BE density matching. # Uses 20 procs, each fragment calculation assigned OMP_NUM_THREADS to 4 # effectively running 5 fragment calculations in parallel -mybe.optimize(solver='CCSD', nproc=20, ompnum=4) +mybe.optimize(solver="CCSD", nproc=20, ompnum=4) # Compute error be_ecorr = mybe.ebe_tot - mybe.ebe_hf -err_ = (ccsd_ecorr - be_ecorr)*100./ccsd_ecorr -print(f'*** BE2 Correlation Energy Error (%) : {err_:>8.4f} %') - +err_ = (ccsd_ecorr - be_ecorr) * 100.0 / ccsd_ecorr +print(f"*** BE2 Correlation Energy Error (%) : {err_:>8.4f} %") diff --git a/example/molbe_octane_get_rdms.py b/example/molbe_octane_get_rdms.py index 1bc25a22..46da2c56 100644 --- a/example/molbe_octane_get_rdms.py +++ b/example/molbe_octane_get_rdms.py @@ -4,7 +4,8 @@ from molbe import fragpart, BE # Perform pyscf HF calculation to get mol & mf objects -mol = gto.M(atom=''' +mol = gto.M( + atom=""" C 0.4419364699 -0.6201930287 0.0000000000 C -0.4419364699 0.6201930287 0.0000000000 H -1.0972005331 0.5963340874 0.8754771384 @@ -31,7 +32,10 @@ H -0.9171145792 -4.5073104916 0.8797333088 H 0.3671153250 -5.3316378285 0.0000000000 H -0.3671153250 5.3316378285 0.0000000000 -''',basis='sto-3g', charge=0) +""", + basis="sto-3g", + charge=0, +) mf = scf.RHF(mol) @@ -39,13 +43,13 @@ mf.kernel() # initialize fragments (use frozen core approximation) -fobj = fragpart(be_type='be2', mol=mol, frozen_core=True) +fobj = fragpart(be_type="be2", mol=mol, frozen_core=True) # Initialize BE mybe = BE(mf, fobj) # Perform BE density matching. # Uses 20 procs, each fragment calculation assigned OMP_NUM_THREADS to 4 # effectively running 5 fragment calculations in parallel -mybe.optimize(solver='CCSD', nproc=20, ompnum=4) +mybe.optimize(solver="CCSD", nproc=20, ompnum=4) rdm1_ao, rdm2_ao = mybe.rdm1_fullbasis(return_ao=True) diff --git a/example/molbe_oneshot_rbe_hcore.py b/example/molbe_oneshot_rbe_hcore.py index 3ec7db48..3e1599ac 100644 --- a/example/molbe_oneshot_rbe_hcore.py +++ b/example/molbe_oneshot_rbe_hcore.py @@ -7,11 +7,12 @@ import numpy from pyscf import gto, scf, qmmm from molbe.misc import be2puffin -from molbe import be_var +from molbe import be_var # variables for scratch handling -#pbe_var.SCRATCH = '{}' -#pbe_var.CREATE_SCRATCH_DIR = True +# pbe_var.SCRATCH = '{}' +# pbe_var.CREATE_SCRATCH_DIR = True + # Convert PySCF integrals to libint format def pyscf2lint(mol, hcore_pyscf): @@ -31,20 +32,18 @@ def pyscf2lint(mol, hcore_pyscf): hcore_libint = hcore_pyscf[numpy.ix_(pyscf2libint_ind, pyscf2libint_ind)] return hcore_libint + # Set MM charges and their positions to use PySCF's QM/MM # functionality. Note that the units for the coordinates are # in Bohr and the units for the structure are in Angstrom # to match Q4Bio application. This can be changed in # misc/be2puffin -charges = [-.2, -.1, .15, .2] -coords = [(-3, -8, -2), - (-2, 6, 1), - (2, -5, 2), - (1, 8, 1.5)] +charges = [-0.2, -0.1, 0.15, 0.2] +coords = [(-3, -8, -2), (-2, 6, 1), (2, -5, 2), (1, 8, 1.5)] # Give structure XYZ, in Angstroms -structure = 'data/octane.xyz' +structure = "data/octane.xyz" """ Build hcore in the libint form for QM/MM @@ -55,7 +54,7 @@ def pyscf2lint(mol, hcore_pyscf): # Build QM/MM RHF Hamiltonian mf1 = scf.RHF(mol) -mf = qmmm.mm_charge(mf1, coords, charges, unit='bohr') +mf = qmmm.mm_charge(mf1, coords, charges, unit="bohr") hcore_pyscf = mf.get_hcore() # Note: can save hcore_pyscf with: numpy.savetxt("hcore_pyscf.dat", hcore_pyscf) @@ -67,21 +66,22 @@ def pyscf2lint(mol, hcore_pyscf): # returns BE energy with CCSD solver from RHF reference, # using checkfile from converged RHF -be_energy = be2puffin(structure, # the QM region XYZ geometry - 'sto-3g', # the chosen basis set - hcore = hcore_libint, # the loaded hamiltonian - libint_inp = True, # True if passing hcore in libint format, False for PySCF - use_df = False, # density fitting - charge = 0, # charge of QM region - spin = 0, # spin of QM region - nproc = 1, # number of processors to parallize across - ompnum = 2, - be_type = 'be2', # BE type: this sets the fragment size. - frozen_core = True, # Frozen core - unrestricted = False, # specify restricted calculation - from_chk = False, # can save the RHF as PySCF checkpoint. - # Set to true if running from converged UHF chk - checkfile = None) # if not None, will save RHF calculation to a checkfile. - # if rerunning from chk (from_chk=True), name the checkfile here +be_energy = be2puffin( + structure, # the QM region XYZ geometry + "sto-3g", # the chosen basis set + hcore=hcore_libint, # the loaded hamiltonian + libint_inp=True, # True if passing hcore in libint format, False for PySCF + use_df=False, # density fitting + charge=0, # charge of QM region + spin=0, # spin of QM region + nproc=1, # number of processors to parallize across + ompnum=2, + be_type="be2", # BE type: this sets the fragment size. + frozen_core=True, # Frozen core + unrestricted=False, # specify restricted calculation + from_chk=False, # can save the RHF as PySCF checkpoint. + # Set to true if running from converged UHF chk + checkfile=None, +) # if not None, will save RHF calculation to a checkfile. +# if rerunning from chk (from_chk=True), name the checkfile here # ecp = ecp) # can add ECP for heavy atoms as: {'Ru': 'def2-SVP'} - diff --git a/example/molbe_oneshot_rbe_qmmm-fromchk.py b/example/molbe_oneshot_rbe_qmmm-fromchk.py index 32c5d6cd..e30f00bd 100644 --- a/example/molbe_oneshot_rbe_qmmm-fromchk.py +++ b/example/molbe_oneshot_rbe_qmmm-fromchk.py @@ -5,11 +5,11 @@ import numpy from pyscf import gto, scf, qmmm from molbe.misc import be2puffin -from molbe import be_var +from molbe import be_var # variables for scratch handling -#pbe_var.SCRATCH = '{}' -#pbe_var.CREATE_SCRATCH_DIR = True +# pbe_var.SCRATCH = '{}' +# pbe_var.CREATE_SCRATCH_DIR = True # Set MM charges and their positions to use PySCF's QM/MM # functionality. Note that the units for the coordinates are @@ -17,32 +17,31 @@ # to match Q4Bio application. This can be changed in # misc/be2puffin -charges = [-.2, -.1, .15, .2] -coords = [(-3, -8, -2), - (-2, 6, 1), - (2, -5, 2), - (1, 8, 1.5)] +charges = [-0.2, -0.1, 0.15, 0.2] +coords = [(-3, -8, -2), (-2, 6, 1), (2, -5, 2), (1, 8, 1.5)] # Give structure XYZ, in Angstroms -structure = 'data/octane.xyz' +structure = "data/octane.xyz" # returns BE energy with CCSD solver from RHF reference, # using checkfile from converged RHF -be_energy = be2puffin(structure, # the QM region XYZ geometry - 'sto-3g', # the chosen basis set - pts_and_charges = [coords, charges], # the loaded hamiltonian - use_df = False, # density fitting - charge = 0, # charge of QM region - spin = 0, # spin of QM region - nproc = 1, # number of processors to parallize across - ompnum = 2, - be_type = 'be2', # BE type: this sets the fragment size. - frozen_core = False, # Frozen core - unrestricted = False, # specify restricted calculation - from_chk = True, # can save the RHF as PySCF checkpoint. - # Set to true if running from converged UHF chk - checkfile = 'data/oneshot_rbe_qmmm.chk') # if not None, will save RHF calculation to a checkfile. - # if rerunning from chk (from_chk=True), name the checkfile here +be_energy = be2puffin( + structure, # the QM region XYZ geometry + "sto-3g", # the chosen basis set + pts_and_charges=[coords, charges], # the loaded hamiltonian + use_df=False, # density fitting + charge=0, # charge of QM region + spin=0, # spin of QM region + nproc=1, # number of processors to parallize across + ompnum=2, + be_type="be2", # BE type: this sets the fragment size. + frozen_core=False, # Frozen core + unrestricted=False, # specify restricted calculation + from_chk=True, # can save the RHF as PySCF checkpoint. + # Set to true if running from converged UHF chk + checkfile="data/oneshot_rbe_qmmm.chk", +) # if not None, will save RHF calculation to a checkfile. +# if rerunning from chk (from_chk=True), name the checkfile here # ecp = ecp) # can add ECP for heavy atoms as: {'Ru': 'def2-SVP'} """ @@ -58,4 +57,3 @@ from_chk = True checkfile = {Name_of_checkfile} """ - diff --git a/example/molbe_oneshot_ube_qmmm.py b/example/molbe_oneshot_ube_qmmm.py index 564d5df7..eb82247f 100644 --- a/example/molbe_oneshot_ube_qmmm.py +++ b/example/molbe_oneshot_ube_qmmm.py @@ -7,11 +7,11 @@ from molbe import be_var # variables for scratch handling -#pbe_var.SCRATCH = '{}' -#pbe_var.CREATE_SCRATCH_DIR = True +# pbe_var.SCRATCH = '{}' +# pbe_var.CREATE_SCRATCH_DIR = True # Give structure XYZ, in Angstroms -structure = 'data/octane.xyz' +structure = "data/octane.xyz" # Set MM charges and their positions to use PySCF's QM/MM # functionality. Note that the units for the coordinates are @@ -19,29 +19,27 @@ # to match Q4Bio application. This can be changed in # misc/be2puffin -charges = [-.2, -.1, .15, .2] -coords = [(-3, -8, -2), - (-2, 6, 1), - (2, -5, 2), - (1, 8, 1.5)] +charges = [-0.2, -0.1, 0.15, 0.2] +coords = [(-3, -8, -2), (-2, 6, 1), (2, -5, 2), (1, 8, 1.5)] # returns UBE energy with UCCSD solver from UHF reference -be_energy = be2puffin(structure, # the QM region XYZ geometry - 'STO-3G', # the chosen basis set - pts_and_charges=[coords, charges], # the point coordinates and coordinates - use_df = False, # keep density fitting False for PySCF UHF - charge = -1, # charge of QM region - spin = 1, # spin of QM region - nproc = 1, # number of processors to parallize across - ompnum = 2, # number of nodes to parallelize across - be_type = 'be2', # BE type: this sets the fragment size. - frozen_core = False, # keep this to False for non-minimal basis: localization and - # numerical problems for ruthenium systems in non-minimal basis - unrestricted = True, # specify unrestricted calculation - from_chk = False, # can save the UHF as PySCF checkpoint. - # Set to true if running from converged UHF chk - checkfile = None) # if not None, will save UHF calculation to a checkfile. - # if rerunning from chk (from_chk=True), name the checkfile here +be_energy = be2puffin( + structure, # the QM region XYZ geometry + "STO-3G", # the chosen basis set + pts_and_charges=[coords, charges], # the point coordinates and coordinates + use_df=False, # keep density fitting False for PySCF UHF + charge=-1, # charge of QM region + spin=1, # spin of QM region + nproc=1, # number of processors to parallize across + ompnum=2, # number of nodes to parallelize across + be_type="be2", # BE type: this sets the fragment size. + frozen_core=False, # keep this to False for non-minimal basis: localization and + # numerical problems for ruthenium systems in non-minimal basis + unrestricted=True, # specify unrestricted calculation + from_chk=False, # can save the UHF as PySCF checkpoint. + # Set to true if running from converged UHF chk + checkfile=None, +) # if not None, will save UHF calculation to a checkfile. +# if rerunning from chk (from_chk=True), name the checkfile here # ecp = ecp) # can add ECP for heavy atoms as: {'Ru': 'def2-SVP'} - diff --git a/example/molbe_ppp.py b/example/molbe_ppp.py index a40cbd44..8ab6cde2 100644 --- a/example/molbe_ppp.py +++ b/example/molbe_ppp.py @@ -1,10 +1,11 @@ # Perform BE calculation with 6-31g basis set -from pyscf import gto,scf +from pyscf import gto, scf from molbe import fragpart, BE # Perform pyscf HF calculation to get mol & mf objects -mol = gto.M(atom=''' +mol = gto.M( + atom=""" C 3.74360 5.55710 7.14890 C 3.18510 4.41510 6.58860 C 3.18510 4.41510 5.17210 @@ -19,7 +20,10 @@ C 3.92620 5.88870 8.53990 H 4.87720 7.84630 8.50510 H 4.87720 7.84630 11.11180 -''',basis='6-31g', charge=0) +""", + basis="6-31g", + charge=0, +) mf = scf.RHF(mol) @@ -27,11 +31,10 @@ mf.kernel() # Define fragments; use IAO scheme with 'sto-3g' as the minimal basis set -fobj = fragpart(be_type='be2', mol=mol, - valence_basis='sto-3g', frozen_core=True) +fobj = fragpart(be_type="be2", mol=mol, valence_basis="sto-3g", frozen_core=True) # Initialize BE -mybe = BE(mg, fobj, lo_method='iao') +mybe = BE(mg, fobj, lo_method="iao") # Density matching with CCSD as local solver -mybe.optimize(solver='CCSD') +mybe.optimize(solver="CCSD") diff --git a/kbe/_opt.py b/kbe/_opt.py index 1be41d91..aa156754 100644 --- a/kbe/_opt.py +++ b/kbe/_opt.py @@ -3,9 +3,19 @@ from molbe._opt import BEOPT -def optimize(self, solver='MP2',method='QN', - only_chem=False, conv_tol = 1.e-6,relax_density=False, use_cumulant=True, - J0=None, nproc=1, ompnum=4, max_iter=500): +def optimize( + self, + solver="MP2", + method="QN", + only_chem=False, + conv_tol=1.0e-6, + relax_density=False, + use_cumulant=True, + J0=None, + nproc=1, + ompnum=4, + max_iter=500, +): """BE optimization function Interfaces BEOPT to perform bootstrap embedding optimization. @@ -41,35 +51,54 @@ def optimize(self, solver='MP2',method='QN', # Check if only chemical potential optimization is required if not only_chem: pot = self.pot - if self.be_type=='be1': - sys.exit('BE1 only works with chemical potential optimization. Set only_chem=True') + if self.be_type == "be1": + sys.exit( + "BE1 only works with chemical potential optimization. Set only_chem=True" + ) else: - pot = [0.] + pot = [0.0] # Initialize the BEOPT object - be_ = BEOPT(pot, self.Fobjs, self.Nocc, self.enuc, hf_veff = self.hf_veff, - nproc=nproc, ompnum=ompnum, - max_space=max_iter,conv_tol = conv_tol, - only_chem=only_chem, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff = self.ci_coeff_cutoff,relax_density=relax_density, - select_cutoff = self.select_cutoff, - solver=solver, ecore=self.E_core, ebe_hf=self.ebe_hf) + be_ = BEOPT( + pot, + self.Fobjs, + self.Nocc, + self.enuc, + hf_veff=self.hf_veff, + nproc=nproc, + ompnum=ompnum, + max_space=max_iter, + conv_tol=conv_tol, + only_chem=only_chem, + hci_cutoff=self.hci_cutoff, + ci_coeff_cutoff=self.ci_coeff_cutoff, + relax_density=relax_density, + select_cutoff=self.select_cutoff, + solver=solver, + ecore=self.E_core, + ebe_hf=self.ebe_hf, + ) - if method=='QN': + if method == "QN": # Prepare the initial Jacobian matrix if only_chem: - J0 = [[0.]] - J0 = self.get_be_error_jacobian(jac_solver='HF') - J0 = [[J0[-1,-1]]] + J0 = [[0.0]] + J0 = self.get_be_error_jacobian(jac_solver="HF") + J0 = [[J0[-1, -1]]] else: - J0 = self.get_be_error_jacobian(jac_solver='HF') + J0 = self.get_be_error_jacobian(jac_solver="HF") # Perform the optimization be_.optimize(method, J0=J0) self.ebe_tot = self.ebe_hf + be_.Ebe[0] # Print the energy components - print_energy(be_.Ebe[0], be_.Ebe[1][1], be_.Ebe[1][0]+be_.Ebe[1][2], self.ebe_hf, self.unitcell_nkpt) + print_energy( + be_.Ebe[0], + be_.Ebe[1][1], + be_.Ebe[1][0] + be_.Ebe[1][2], + self.ebe_hf, + self.unitcell_nkpt, + ) else: - print('This optimization method for BE is not supported') + print("This optimization method for BE is not supported") sys.exit() diff --git a/kbe/autofrag.py b/kbe/autofrag.py index 6b452c47..ef5a15da 100644 --- a/kbe/autofrag.py +++ b/kbe/autofrag.py @@ -5,11 +5,13 @@ from molbe.helper import get_core from itertools import compress + def warn_large_fragment(): - print('Fragments that spans more than 2 unit-cells are not supported') - print('Try with larger unit-cell(or super-cell)') + print("Fragments that spans more than 2 unit-cells are not supported") + print("Try with larger unit-cell(or super-cell)") sys.exit() + def add_check_k(min1, flist, sts, ksts, nk_): sts_ = sts.copy() for q_ in min1: @@ -26,8 +28,8 @@ def add_check_k(min1, flist, sts, ksts, nk_): sts.append(q_) ksts.append(nk_) -def nearestof2coord(coord1, coord2 , bond=2.6*1.88973): +def nearestof2coord(coord1, coord2, bond=2.6 * 1.88973): mind = 50000 lmin = () for idx, i in enumerate(coord1): @@ -36,7 +38,7 @@ def nearestof2coord(coord1, coord2 , bond=2.6*1.88973): continue dist = numpy.linalg.norm(i - j) - if dist < mind or dist-mind < 0.1: + if dist < mind or dist - mind < 0.1: if dist <= bond: lmin = (idx, jdx) mind = dist @@ -44,34 +46,45 @@ def nearestof2coord(coord1, coord2 , bond=2.6*1.88973): lunit_ = [lmin[0]] runit_ = [lmin[1]] else: - return([],[]) + return ([], []) for idx, i in enumerate(coord1): for jdx, j in enumerate(coord2): - if idx == jdx : + if idx == jdx: continue if idx == lmin[0] and jdx == lmin[1]: continue dist = numpy.linalg.norm(i - j) - if dist-mind<0.1 and dist <= bond: - + if dist - mind < 0.1 and dist <= bond: lunit_.append(idx) runit_.append(jdx) - return(lunit_, runit_) - - -def sidefunc(cell, Idx, unit1, unit2, main_list, sub_list, coord, - be_type, bond=2.6*1.88973, klist=[], - ext_list=[], NK=None, rlist=[]): - + return (lunit_, runit_) + + +def sidefunc( + cell, + Idx, + unit1, + unit2, + main_list, + sub_list, + coord, + be_type, + bond=2.6 * 1.88973, + klist=[], + ext_list=[], + NK=None, + rlist=[], +): if ext_list == []: - main_list.extend(unit2[numpy.where(unit1==Idx)[0]]) - sub_list.extend(unit2[numpy.where(unit1==Idx)[0]]) + main_list.extend(unit2[numpy.where(unit1 == Idx)[0]]) + sub_list.extend(unit2[numpy.where(unit1 == Idx)[0]]) else: - for sub_i in unit2[numpy.where(unit1==Idx)[0]]: - if sub_i in rlist: continue + for sub_i in unit2[numpy.where(unit1 == Idx)[0]]: + if sub_i in rlist: + continue if sub_i in ext_list: for tmp_jdx, tmpj in enumerate(ext_list): if sub_i == tmpj and klist[tmp_jdx] == NK: @@ -86,24 +99,30 @@ def sidefunc(cell, Idx, unit1, unit2, main_list, sub_list, coord, closest = sub_list.copy() close_be3 = [] - if be_type == 'be3' or be_type == 'be4': - for lmin1 in unit2[numpy.where(unit1==Idx)[0]]: + if be_type == "be3" or be_type == "be4": + for lmin1 in unit2[numpy.where(unit1 == Idx)[0]]: for jdx, j in enumerate(coord): - if not jdx in unit1 and not jdx in unit2 \ - and not cell.atom_pure_symbol(jdx) == 'H': + if ( + not jdx in unit1 + and not jdx in unit2 + and not cell.atom_pure_symbol(jdx) == "H" + ): dist = numpy.linalg.norm(coord[lmin1] - j) if dist <= bond: - if not jdx in sub_list: # avoid repeated occurence + if not jdx in sub_list: # avoid repeated occurence main_list.append(jdx) sub_list.append(jdx) close_be3.append(jdx) - if be_type == 'be4': + if be_type == "be4": for kdx, k in enumerate(coord): if kdx == jdx: continue - if not kdx in unit1 and not kdx in unit2 \ - and not cell.atom_pure_symbol(kdx) == 'H': + if ( + not kdx in unit1 + and not kdx in unit2 + and not cell.atom_pure_symbol(kdx) == "H" + ): dist = numpy.linalg.norm(coord[jdx] - k) if dist <= bond: main_list.append(kdx) @@ -111,17 +130,41 @@ def sidefunc(cell, Idx, unit1, unit2, main_list, sub_list, coord, return closest, close_be3 -def surround(cell, sidx, unit1, unit2, flist, coord, be_type, ext_list, - klist, NK, rlist=[], bond=1.8*1.88973): - +def surround( + cell, + sidx, + unit1, + unit2, + flist, + coord, + be_type, + ext_list, + klist, + NK, + rlist=[], + bond=1.8 * 1.88973, +): be_type_ = be_reduce(be_type) - if not rlist == [] and be_type_ == 'be3': - be_type_ = 'be2' + if not rlist == [] and be_type_ == "be3": + be_type_ = "be2" sublist_ = [] flist_ = flist.copy() if not be_type_ == 0: - sidefunc(cell, sidx, unit1, unit2, flist, sublist_, coord, be_type_, - bond=bond, klist=klist, ext_list=ext_list,NK=NK, rlist=rlist) + sidefunc( + cell, + sidx, + unit1, + unit2, + flist, + sublist_, + coord, + be_type_, + bond=bond, + klist=klist, + ext_list=ext_list, + NK=NK, + rlist=rlist, + ) sublist = [tmpi for tmpi in sublist_ if not tmpi in rlist] sublist = [] for tmpi in sublist_: @@ -137,8 +180,18 @@ def surround(cell, sidx, unit1, unit2, flist, coord, be_type, ext_list, klist.append(NK) -def kfrag_func(site_list, numk, nk1, uNs, Ns, nk2=None, debug=False, - debug1=False, shift=False, debug2=False): +def kfrag_func( + site_list, + numk, + nk1, + uNs, + Ns, + nk2=None, + debug=False, + debug1=False, + shift=False, + debug2=False, +): if nk2 is None: nk2 = nk1 frglist = [] @@ -146,50 +199,66 @@ def kfrag_func(site_list, numk, nk1, uNs, Ns, nk2=None, debug=False, for pq in site_list: if numk > nk1 * 2: if not uNs == Ns: - nk_ = numk - (nk1*2) - frglist.append(uNs*nk1*2 + (nk_*uNs) + pq) + nk_ = numk - (nk1 * 2) + frglist.append(uNs * nk1 * 2 + (nk_ * uNs) + pq) else: - frglist.append(uNs*numk + pq) - elif numk ==nk1: + frglist.append(uNs * numk + pq) + elif numk == nk1: if not uNs == Ns: - frglist.append(uNs*numk+pq) + frglist.append(uNs * numk + pq) else: if not debug2: - frglist.append(uNs*numk + pq) + frglist.append(uNs * numk + pq) else: - frglist.append(uNs*numk + pq) - elif numk >nk1: + frglist.append(uNs * numk + pq) + elif numk > nk1: if uNs == Ns: - frglist.append(uNs*numk + pq) + frglist.append(uNs * numk + pq) else: nk_ = numk - nk1 - frglist.append(uNs*nk1 + (nk_*uNs)+pq) + frglist.append(uNs * nk1 + (nk_ * uNs) + pq) else: - if not Ns==uNs: - frglist.append(uNs*numk + pq) + if not Ns == uNs: + frglist.append(uNs * numk + pq) else: - frglist.append(uNs*numk + pq) - if debug: print(uNs*numk + pq, end=' ') - if debug: print() + frglist.append(uNs * numk + pq) + if debug: + print(uNs * numk + pq, end=" ") + if debug: + print() return frglist + def be_reduce(be_type, N=1): - if N ==1: - if be_type == 'be2': + if N == 1: + if be_type == "be2": return 0 - elif be_type =='be3': - return 'be2' - elif be_type == 'be4': - return 'be3' - -def autogen(mol, kpt, frozen_core=True, be_type='be2', - write_geom=False, - unitcell=1, gamma_2d=False, gamma_1d=False, - long_bond=False, perpend_dist = 4.0, perpend_dist_tol = 1e-3, - nx=False, ny=False, nz=False, - valence_basis = None, interlayer = False, - print_frags=True): + elif be_type == "be3": + return "be2" + elif be_type == "be4": + return "be3" + + +def autogen( + mol, + kpt, + frozen_core=True, + be_type="be2", + write_geom=False, + unitcell=1, + gamma_2d=False, + gamma_1d=False, + long_bond=False, + perpend_dist=4.0, + perpend_dist_tol=1e-3, + nx=False, + ny=False, + nz=False, + valence_basis=None, + interlayer=False, + print_frags=True, +): """ Automatic cell partitioning @@ -245,7 +314,7 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', from .misc import sgeom if not float(unitcell).is_integer(): - print('Fractional unitcell is not supported!') + print("Fractional unitcell is not supported!") sys.exit() elif unitcell > 1: if not nx and not ny and not nz: @@ -262,14 +331,20 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', cell.build() print(flush=True) - print('No. of cells used in building fragments : {:>3}'.format(unitcell),flush=True) + print( + "No. of cells used in building fragments : {:>3}".format(unitcell), + flush=True, + ) for idx, i in enumerate(kpt): if not i == 1: - n_ = i/float(unitcell) + n_ = i / float(unitcell) if n_ > kpt[idx]: - print('Use a larger number of k-points; ',flush=True) - print('Fragment cell larger than all k-points combined is not supported',flush=True) + print("Use a larger number of k-points; ", flush=True) + print( + "Fragment cell larger than all k-points combined is not supported", + flush=True, + ) sys.exit() else: cell = mol.copy() @@ -283,7 +358,7 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', bond = 1.8 * ang2bohr if long_bond: normdist = 6.0 * ang2bohr - bond=2.6 *ang2bohr + bond = 2.6 * ang2bohr hbond = 1.2 * ang2bohr lunit = [] runit = [] @@ -292,7 +367,7 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', munit = [] tunit = [] - lnext = [i for i in kpt if i>1] + lnext = [i for i in kpt if i > 1] if not len(lnext) == 0: nk1 = lnext[0] @@ -300,24 +375,24 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', if len(lnext) > 1: nk2 = lnext[1] twoD = True - lkpt = [2 if i>1 else 1 for i in kpt] + lkpt = [2 if i > 1 else 1 for i in kpt] else: if gamma_1d: nk1 = 1 nk2 = 1 - lkpt = [1,1,2] - twoD=False + lkpt = [1, 1, 2] + twoD = False else: nk1 = 1 nk2 = 1 twoD = True - lkpt = [2,2,1] + lkpt = [2, 2, 1] if frozen_core: ncore__, no_core_idx__, core_list__ = get_core(mol) - Nsite = mol.aoslice_by_atom()[-1][3]-ncore__ -1 + Nsite = mol.aoslice_by_atom()[-1][3] - ncore__ - 1 else: - Nsite = mol.aoslice_by_atom()[-1][3]-1 + Nsite = mol.aoslice_by_atom()[-1][3] - 1 # kmesh is lkpt # Building neighbouring cells @@ -327,26 +402,26 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', lnk2 = 1 lattice_vector = cell.lattice_vectors() - Ts = lib.cartesian_prod((numpy.arange(lkpt[0]), - numpy.arange(lkpt[1]), - numpy.arange(lkpt[2]))) + Ts = lib.cartesian_prod( + (numpy.arange(lkpt[0]), numpy.arange(lkpt[1]), numpy.arange(lkpt[2])) + ) Ls = numpy.dot(Ts, lattice_vector) # 1-2-(1-2)-1-2 # * * - lcoord = Ls.reshape(-1,1,3)[lnk2]*-1 + coord - rcoord = Ls.reshape(-1,1,3)[lnk2] + coord + lcoord = Ls.reshape(-1, 1, 3)[lnk2] * -1 + coord + rcoord = Ls.reshape(-1, 1, 3)[lnk2] + coord lunit, runit = nearestof2coord(coord, lcoord, bond=bond) lunit_, runit_ = nearestof2coord(coord, rcoord, bond=bond) if not set(lunit) == set(runit_) or not set(runit) == set(lunit_): - print('Fragmentation error : wrong connection of unit cells ') + print("Fragmentation error : wrong connection of unit cells ") sys.exit() - if sum(i>1 for i in kpt) > 1 or gamma_2d: + if sum(i > 1 for i in kpt) > 1 or gamma_2d: # only 2D is supported # neighbours up-down or right-left @@ -355,32 +430,30 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', # (1-2) # | # *1-2 - lcoord2 = Ls.reshape(-1,1,3)[1]*-1 + coord - rcoord2 = Ls.reshape(-1,1,3)[1] + coord + lcoord2 = Ls.reshape(-1, 1, 3)[1] * -1 + coord + rcoord2 = Ls.reshape(-1, 1, 3)[1] + coord dunit, uunit = nearestof2coord(coord, lcoord2, bond=bond) dunit_, uunit_ = nearestof2coord(coord, rcoord2, bond=bond) if not set(uunit) == set(dunit_) or not set(dunit) == set(uunit_): - print('Fragmentation error : wrong connection of unit cells ') + print("Fragmentation error : wrong connection of unit cells ") sys.exit() - # diagonal # 1-2 1-2* # \ / # (1-2) # / \ # 1-2* 1-2 - lcoord3 = Ls.reshape(-1,1,3)[lnk2+1]*-1 + coord - rcoord3 = Ls.reshape(-1,1,3)[lnk2+1] + coord - + lcoord3 = Ls.reshape(-1, 1, 3)[lnk2 + 1] * -1 + coord + rcoord3 = Ls.reshape(-1, 1, 3)[lnk2 + 1] + coord munit, tunit = nearestof2coord(coord, lcoord3, bond=bond) munit_, tunit_ = nearestof2coord(coord, rcoord3, bond=bond) if not set(munit) == set(tunit_) or not set(tunit) == set(munit_): - print('Fragmentation error : wrong connection of unit cells ') + print("Fragmentation error : wrong connection of unit cells ") sys.exit() # kmesh lkpt ends @@ -412,7 +485,7 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', munit = numpy.asarray(munit) tunit = numpy.asarray(tunit) - inter_dist = 1000. + inter_dist = 1000.0 inter_idx = 0 if twoD and interlayer: inter_layer_axis = 2 @@ -421,9 +494,11 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', inter_dist = 1000 inter_idx = 0 for ajdx, aj in enumerate(coord): - if aidx==ajdx: continue - if ai[inter_layer_axis] == aj[inter_layer_axis]: continue - dist = numpy.linalg.norm(ai-aj) + if aidx == ajdx: + continue + if ai[inter_layer_axis] == aj[inter_layer_axis]: + continue + dist = numpy.linalg.norm(ai - aj) if dist > bond: if inter_dist > dist: inter_dist = dist @@ -431,23 +506,25 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', inter_dist_ = [] inter_idx_ = [] for ajdx, aj in enumerate(coord): - if aidx==ajdx: continue - if ai[inter_layer_axis] == aj[inter_layer_axis]: continue - dist = numpy.linalg.norm(ai-aj) - if abs(dist-inter_dist) < perpend_dist_tol: + if aidx == ajdx: + continue + if ai[inter_layer_axis] == aj[inter_layer_axis]: + continue + dist = numpy.linalg.norm(ai - aj) + if abs(dist - inter_dist) < perpend_dist_tol: inter_dist_.append(dist) inter_idx_.append(ajdx) inter_layer_dict.append([inter_idx_, inter_dist_]) # Assumes - minimum atom in a ring is 5 - if be_type == 'be4': + if be_type == "be4": if twoD: - print('*********************') - print('USE BE4 WITH CAUTION') - print('*********************') + print("*********************") + print("USE BE4 WITH CAUTION") + print("*********************") for idx, i in enumerate(normlist): - if cell.atom_pure_symbol(idx) == 'H': + if cell.atom_pure_symbol(idx) == "H": continue lsts = [] @@ -469,239 +546,759 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', clist = [] cout = 0 - for jdx,j in enumerate(tmplist): - if not idx==jdx and not cell.atom_pure_symbol(jdx) == 'H': - if abs(j)< normdist: + for jdx, j in enumerate(tmplist): + if not idx == jdx and not cell.atom_pure_symbol(jdx) == "H": + if abs(j) < normdist: clist.append(jdx) pedg = [] flist = [] - clist_check = [] # contains atoms from fragment that are - # part of l,r,d,u,m,t but are in the - # unit cell that is their k is 1 - # for example, klsts[i] = 1 + clist_check = [] # contains atoms from fragment that are + # part of l,r,d,u,m,t but are in the + # unit cell that is their k is 1 + # for example, klsts[i] = 1 ## WARNING !!! # For systems like Graphene, BN, SiC, hexagonal 2D sheets, # BE4 can give wrong fragmentations # the following code adds in redundant atoms if idx in lunit: - closest, close_be3 = sidefunc(cell, idx, lunit, runit, flist, - lsts, coord, be_type, bond=bond) + closest, close_be3 = sidefunc( + cell, idx, lunit, runit, flist, lsts, coord, be_type, bond=bond + ) for zdx in lsts: if twoD: - klsts.append(nk1*(nk2-1)+1) + klsts.append(nk1 * (nk2 - 1) + 1) else: klsts.append(nk1) for lsidx in closest: - if lsidx in lunit or lsidx in munit: warn_large_fragment() + if lsidx in lunit or lsidx in munit: + warn_large_fragment() if lsidx in runit: - surround(cell, lsidx, runit, lunit, flist, coord, be_type, - lsts, klsts, 1, rlist=[idx]+clist_check, bond=bond) #1 + surround( + cell, + lsidx, + runit, + lunit, + flist, + coord, + be_type, + lsts, + klsts, + 1, + rlist=[idx] + clist_check, + bond=bond, + ) # 1 if lsidx in uunit: - surround(cell, lsidx, uunit, dunit, flist, coord, be_type, - lsts, klsts, nk1*(nk2-1)+2, bond=bond) + surround( + cell, + lsidx, + uunit, + dunit, + flist, + coord, + be_type, + lsts, + klsts, + nk1 * (nk2 - 1) + 2, + bond=bond, + ) if lsidx in dunit: - surround(cell, lsidx, dunit, uunit, flist, coord, be_type, - lsts, klsts, nk1*nk2, bond=bond) + surround( + cell, + lsidx, + dunit, + uunit, + flist, + coord, + be_type, + lsts, + klsts, + nk1 * nk2, + bond=bond, + ) if lsidx in tunit: - surround(cell, lsidx, tunit, munit, flist, coord, be_type, - lsts, klsts, 2, bond=bond) - - if be_type == 'be4': + surround( + cell, + lsidx, + tunit, + munit, + flist, + coord, + be_type, + lsts, + klsts, + 2, + bond=bond, + ) + + if be_type == "be4": for lsidx in close_be3: - if lsidx in lunit or lsidx in munit: warn_large_fragment() + if lsidx in lunit or lsidx in munit: + warn_large_fragment() if lsidx in runit: - surround(cell, lsidx, runit, lunit, flist, coord, 'be3', - lsts, klsts, 1, rlist=[idx], bond=bond) #1 + surround( + cell, + lsidx, + runit, + lunit, + flist, + coord, + "be3", + lsts, + klsts, + 1, + rlist=[idx], + bond=bond, + ) # 1 if lsidx in uunit: - surround(cell, lsidx, uunit, dunit, flist, coord, 'be3', - lsts, klsts, nk1*(nk2-1)+2, bond=bond) + surround( + cell, + lsidx, + uunit, + dunit, + flist, + coord, + "be3", + lsts, + klsts, + nk1 * (nk2 - 1) + 2, + bond=bond, + ) if lsidx in dunit: - surround(cell, lsidx, dunit, uunit, flist, coord, 'be3', - lsts, klsts, nk1*nk2, bond=bond) + surround( + cell, + lsidx, + dunit, + uunit, + flist, + coord, + "be3", + lsts, + klsts, + nk1 * nk2, + bond=bond, + ) if lsidx in tunit: - surround(cell, lsidx, tunit, munit, flist, coord, 'be3', - lsts, klsts, 2, bond=bond) + surround( + cell, + lsidx, + tunit, + munit, + flist, + coord, + "be3", + lsts, + klsts, + 2, + bond=bond, + ) for i1dx, i1 in enumerate(klsts): - if i1 == 1: clist_check.append(lsts[i1dx]) + if i1 == 1: + clist_check.append(lsts[i1dx]) if idx in uunit: - closest, close_be3 = sidefunc(cell, idx, uunit, dunit, flist, usts, - coord, be_type, bond=bond) + closest, close_be3 = sidefunc( + cell, idx, uunit, dunit, flist, usts, coord, be_type, bond=bond + ) for kdx in usts: kusts.append(2) for usidx in closest: - if usidx in uunit or usidx in tunit: warn_large_fragment() + if usidx in uunit or usidx in tunit: + warn_large_fragment() if usidx in lunit: - surround(cell, usidx, lunit, runit, flist, coord, be_type, usts, - kusts, nk1*(nk2-1)+2, bond=bond) + surround( + cell, + usidx, + lunit, + runit, + flist, + coord, + be_type, + usts, + kusts, + nk1 * (nk2 - 1) + 2, + bond=bond, + ) if usidx in runit: - surround(cell, usidx, runit, lunit, flist, coord, be_type, usts, - kusts, nk1+2, bond=bond) + surround( + cell, + usidx, + runit, + lunit, + flist, + coord, + be_type, + usts, + kusts, + nk1 + 2, + bond=bond, + ) if usidx in munit: - surround(cell, usidx, munit, tunit, flist, coord, be_type, usts, - kusts, nk1*(nk2-1)+1, bond=bond) + surround( + cell, + usidx, + munit, + tunit, + flist, + coord, + be_type, + usts, + kusts, + nk1 * (nk2 - 1) + 1, + bond=bond, + ) if usidx in dunit: - surround(cell, usidx, dunit, uunit, flist, coord, be_type, usts, - kusts, 1, rlist=[idx]+clist_check, bond=bond) - if be_type == 'be4': + surround( + cell, + usidx, + dunit, + uunit, + flist, + coord, + be_type, + usts, + kusts, + 1, + rlist=[idx] + clist_check, + bond=bond, + ) + if be_type == "be4": for usidx in close_be3: - if usidx in uunit or usidx in tunit: warn_large_fragment() + if usidx in uunit or usidx in tunit: + warn_large_fragment() if usidx in lunit: - surround(cell, usidx, lunit, runit, flist, coord, 'be3', usts, - kusts, nk1*(nk2-1)+2, bond=bond) + surround( + cell, + usidx, + lunit, + runit, + flist, + coord, + "be3", + usts, + kusts, + nk1 * (nk2 - 1) + 2, + bond=bond, + ) if usidx in runit: - surround(cell, usidx, runit, lunit, flist, coord, 'be3', usts, - kusts, nk1+2, bond=bond) + surround( + cell, + usidx, + runit, + lunit, + flist, + coord, + "be3", + usts, + kusts, + nk1 + 2, + bond=bond, + ) if usidx in munit: - surround(cell, usidx, munit, tunit, flist, coord, 'be3', usts, - kusts, nk1*(nk2-1)+1, bond=bond) + surround( + cell, + usidx, + munit, + tunit, + flist, + coord, + "be3", + usts, + kusts, + nk1 * (nk2 - 1) + 1, + bond=bond, + ) if usidx in dunit: - surround(cell, usidx, dunit, uunit, flist, coord, 'be3', usts, - kusts, 1, rlist=[idx], bond=bond) + surround( + cell, + usidx, + dunit, + uunit, + flist, + coord, + "be3", + usts, + kusts, + 1, + rlist=[idx], + bond=bond, + ) for i1dx, i1 in enumerate(kusts): - if i1 == 1: clist_check.append(usts[i1dx]) + if i1 == 1: + clist_check.append(usts[i1dx]) if idx in munit: - closest, close_be3 = sidefunc(cell, idx, munit, tunit, flist, msts, coord, - be_type, bond=bond) + closest, close_be3 = sidefunc( + cell, idx, munit, tunit, flist, msts, coord, be_type, bond=bond + ) for kdx in msts: - kmsts.append(nk1*nk2) + kmsts.append(nk1 * nk2) for msidx in closest: - if msidx in lunit or msidx in munit or msidx in dunit: warn_large_fragment() + if msidx in lunit or msidx in munit or msidx in dunit: + warn_large_fragment() if msidx in runit: - surround(cell, msidx, runit, lunit, flist, coord, be_type, msts, kmsts, - nk1, bond=bond) + surround( + cell, + msidx, + runit, + lunit, + flist, + coord, + be_type, + msts, + kmsts, + nk1, + bond=bond, + ) if msidx in uunit: - surround(cell, msidx, uunit, dunit, flist, coord, be_type, msts, kmsts, - nk1*(nk2-1)+1, bond=bond) + surround( + cell, + msidx, + uunit, + dunit, + flist, + coord, + be_type, + msts, + kmsts, + nk1 * (nk2 - 1) + 1, + bond=bond, + ) if msidx in tunit: - surround(cell, msidx, tunit, munit, flist, coord, be_type, msts, kmsts, 1, - rlist=[idx]+clist_check, bond=bond) - if be_type == 'be4': + surround( + cell, + msidx, + tunit, + munit, + flist, + coord, + be_type, + msts, + kmsts, + 1, + rlist=[idx] + clist_check, + bond=bond, + ) + if be_type == "be4": for msidx in close_be3: - if msidx in lunit or msidx in munit or msidx in dunit: warn_large_fragment() + if msidx in lunit or msidx in munit or msidx in dunit: + warn_large_fragment() if msidx in runit: - surround(cell, msidx, runit, lunit, flist, coord, 'be3', msts, kmsts, - nk1, bond=bond) + surround( + cell, + msidx, + runit, + lunit, + flist, + coord, + "be3", + msts, + kmsts, + nk1, + bond=bond, + ) if msidx in uunit: - surround(cell, msidx, uunit, dunit, flist, coord, 'be3', msts, kmsts, - nk1*(nk2-1)+1, bond=bond) + surround( + cell, + msidx, + uunit, + dunit, + flist, + coord, + "be3", + msts, + kmsts, + nk1 * (nk2 - 1) + 1, + bond=bond, + ) if msidx in tunit: - surround(cell, msidx, tunit, munit, flist, coord, 'be3', msts, kmsts, - 1, rlist=[idx], bond=bond) + surround( + cell, + msidx, + tunit, + munit, + flist, + coord, + "be3", + msts, + kmsts, + 1, + rlist=[idx], + bond=bond, + ) for i1dx, i1 in enumerate(kmsts): - if i1 == 1: clist_check.append(msts[i1dx]) + if i1 == 1: + clist_check.append(msts[i1dx]) if idx in runit: - closest, close_be3 = sidefunc(cell, idx, runit, lunit, - flist, rsts, coord, be_type, bond=bond) + closest, close_be3 = sidefunc( + cell, idx, runit, lunit, flist, rsts, coord, be_type, bond=bond + ) for kdx in rsts: if twoD: - krsts.append(nk1+1) + krsts.append(nk1 + 1) else: krsts.append(2) for rsidx in closest: - if rsidx in runit or rsidx in tunit: warn_large_fragment() + if rsidx in runit or rsidx in tunit: + warn_large_fragment() if rsidx in lunit: - surround(cell, rsidx, lunit, runit, flist, coord, be_type, rsts, krsts, - 1, rlist=[idx]+clist_check, bond=bond) #1 + surround( + cell, + rsidx, + lunit, + runit, + flist, + coord, + be_type, + rsts, + krsts, + 1, + rlist=[idx] + clist_check, + bond=bond, + ) # 1 if rsidx in munit: - surround(cell, rsidx, munit, tunit, flist, coord, be_type, rsts, krsts, - nk1, bond=bond) + surround( + cell, + rsidx, + munit, + tunit, + flist, + coord, + be_type, + rsts, + krsts, + nk1, + bond=bond, + ) if rsidx in uunit: - surround(cell, rsidx, uunit, dunit, flist, coord, be_type, rsts, krsts, - nk1+2, bond=bond) + surround( + cell, + rsidx, + uunit, + dunit, + flist, + coord, + be_type, + rsts, + krsts, + nk1 + 2, + bond=bond, + ) if rsidx in dunit: - surround(cell, rsidx, dunit, uunit, flist, coord, be_type, rsts, krsts, - nk1*2, bond=bond) - if be_type == 'be4': + surround( + cell, + rsidx, + dunit, + uunit, + flist, + coord, + be_type, + rsts, + krsts, + nk1 * 2, + bond=bond, + ) + if be_type == "be4": for rsidx in close_be3: - if rsidx in runit or rsidx in tunit: warn_large_fragment() + if rsidx in runit or rsidx in tunit: + warn_large_fragment() if rsidx in lunit: - surround(cell, rsidx, lunit, runit, flist, coord, 'be3', rsts, krsts, - 1, rlist=[idx], bond=bond) #1 + surround( + cell, + rsidx, + lunit, + runit, + flist, + coord, + "be3", + rsts, + krsts, + 1, + rlist=[idx], + bond=bond, + ) # 1 if rsidx in munit: - surround(cell, rsidx, munit, tunit, flist, coord, 'be3', rsts, krsts, - nk1, bond=bond) + surround( + cell, + rsidx, + munit, + tunit, + flist, + coord, + "be3", + rsts, + krsts, + nk1, + bond=bond, + ) if rsidx in uunit: - surround(cell, rsidx, uunit, dunit, flist, coord, 'be3', rsts, krsts, - nk1+2, bond=bond) + surround( + cell, + rsidx, + uunit, + dunit, + flist, + coord, + "be3", + rsts, + krsts, + nk1 + 2, + bond=bond, + ) if rsidx in dunit: - surround(cell, rsidx, dunit, uunit, flist, coord, 'be3', rsts, krsts, - nk1*2, bond=bond) + surround( + cell, + rsidx, + dunit, + uunit, + flist, + coord, + "be3", + rsts, + krsts, + nk1 * 2, + bond=bond, + ) for i1dx, i1 in enumerate(krsts): - if i1 == 1: clist_check.append(rsts[i1dx]) + if i1 == 1: + clist_check.append(rsts[i1dx]) if idx in dunit: - closest, close_be3 = sidefunc(cell, idx, dunit, uunit, flist, dsts, - coord, be_type, bond=bond) + closest, close_be3 = sidefunc( + cell, idx, dunit, uunit, flist, dsts, coord, be_type, bond=bond + ) for kdx in dsts: kdsts.append(nk1) for dsidx in closest: - if dsidx in munit or dsidx in dunit: warn_large_fragment() + if dsidx in munit or dsidx in dunit: + warn_large_fragment() if dsidx in lunit: - surround(cell, dsidx, lunit, runit, flist, coord, be_type, dsts, kdsts, - nk1*nk2, bond=bond) + surround( + cell, + dsidx, + lunit, + runit, + flist, + coord, + be_type, + dsts, + kdsts, + nk1 * nk2, + bond=bond, + ) if dsidx in runit: - surround(cell, dsidx, runit, lunit, flist, coord, be_type, dsts, kdsts, - nk1*2, bond=bond) + surround( + cell, + dsidx, + runit, + lunit, + flist, + coord, + be_type, + dsts, + kdsts, + nk1 * 2, + bond=bond, + ) if dsidx in uunit: - surround(cell, dsidx, uunit, dunit, flist, coord, be_type, dsts, kdsts, - 1, rlist=[idx]+clist_check, bond=bond) + surround( + cell, + dsidx, + uunit, + dunit, + flist, + coord, + be_type, + dsts, + kdsts, + 1, + rlist=[idx] + clist_check, + bond=bond, + ) if dsidx in tunit: - surround(cell, dsidx, tunit, munit, flist, coord, be_type, dsts, kdsts, - nk1+1, bond=bond) - if be_type == 'be4': + surround( + cell, + dsidx, + tunit, + munit, + flist, + coord, + be_type, + dsts, + kdsts, + nk1 + 1, + bond=bond, + ) + if be_type == "be4": for dsidx in close_be3: - if dsidx in munit or dsidx in dunit: warn_large_fragment() + if dsidx in munit or dsidx in dunit: + warn_large_fragment() if dsidx in lunit: - surround(cell, dsidx, lunit, runit, flist, coord, 'be3', dsts, kdsts, - nk1*nk2, bond=bond) + surround( + cell, + dsidx, + lunit, + runit, + flist, + coord, + "be3", + dsts, + kdsts, + nk1 * nk2, + bond=bond, + ) if dsidx in runit: - surround(cell, dsidx, runit, lunit, flist, coord, 'be3', dsts, kdsts, - nk1*2, bond=bond) + surround( + cell, + dsidx, + runit, + lunit, + flist, + coord, + "be3", + dsts, + kdsts, + nk1 * 2, + bond=bond, + ) if dsidx in uunit: - surround(cell, dsidx, uunit, dunit, flist, coord, 'be3', dsts, kdsts, - 1, rlist=[idx], bond=bond) + surround( + cell, + dsidx, + uunit, + dunit, + flist, + coord, + "be3", + dsts, + kdsts, + 1, + rlist=[idx], + bond=bond, + ) if dsidx in tunit: - surround(cell, dsidx, tunit, munit, flist, coord, 'be3', dsts, kdsts, - nk1+1, bond=bond) + surround( + cell, + dsidx, + tunit, + munit, + flist, + coord, + "be3", + dsts, + kdsts, + nk1 + 1, + bond=bond, + ) for i1dx, i1 in enumerate(kdsts): - if i1 == 1: clist_check.append(dsts[i1dx]) + if i1 == 1: + clist_check.append(dsts[i1dx]) if idx in tunit: - closest, close_be3 = sidefunc(cell, idx, tunit, munit, flist, tsts, coord, - be_type, bond=bond) + closest, close_be3 = sidefunc( + cell, idx, tunit, munit, flist, tsts, coord, be_type, bond=bond + ) for kdx in tsts: - ktsts.append(nk1+2) + ktsts.append(nk1 + 2) for tsidx in closest: - if tsidx in runit or tsidx in uunit or tsidx in tunit: warn_large_fragment() + if tsidx in runit or tsidx in uunit or tsidx in tunit: + warn_large_fragment() if tsidx in lunit: - surround(cell, tsidx, lunit, runit, flist, coord, be_type, tsts, ktsts, 2, - bond=bond) + surround( + cell, + tsidx, + lunit, + runit, + flist, + coord, + be_type, + tsts, + ktsts, + 2, + bond=bond, + ) if tsidx in munit: - surround(cell, tsidx, munit, tunit, flist, coord, be_type, tsts, ktsts, 1, - rlist=[idx]+clist_check, bond=bond) + surround( + cell, + tsidx, + munit, + tunit, + flist, + coord, + be_type, + tsts, + ktsts, + 1, + rlist=[idx] + clist_check, + bond=bond, + ) if tsidx in dunit: - surround(cell, tsidx, dunit, uunit, flist, coord, be_type, tsts, ktsts, - nk1+1, bond=bond) - if be_type == 'be4': + surround( + cell, + tsidx, + dunit, + uunit, + flist, + coord, + be_type, + tsts, + ktsts, + nk1 + 1, + bond=bond, + ) + if be_type == "be4": for tsidx in close_be3: - if tsidx in runit or tsidx in uunit or tsidx in tunit: warn_large_fragment() + if tsidx in runit or tsidx in uunit or tsidx in tunit: + warn_large_fragment() if tsidx in lunit: - surround(cell, tsidx, lunit, runit, flist, coord, 'be3', tsts, ktsts, 2, - bond=bond) + surround( + cell, + tsidx, + lunit, + runit, + flist, + coord, + "be3", + tsts, + ktsts, + 2, + bond=bond, + ) if tsidx in munit: - surround(cell, tsidx, munit, tunit, flist, coord, 'be3', tsts, ktsts, 1, - rlist=[idx], bond=bond) + surround( + cell, + tsidx, + munit, + tunit, + flist, + coord, + "be3", + tsts, + ktsts, + 1, + rlist=[idx], + bond=bond, + ) if tsidx in dunit: - surround(cell, tsidx, dunit, uunit, flist, coord, 'be3', tsts, ktsts, - nk1+1, bond=bond) + surround( + cell, + tsidx, + dunit, + uunit, + flist, + coord, + "be3", + tsts, + ktsts, + nk1 + 1, + bond=bond, + ) for i1dx, i1 in enumerate(ktsts): - if i1 == 1: clist_check.append(tsts[i1dx]) + if i1 == 1: + clist_check.append(tsts[i1dx]) flist.append(idx) cen.append(idx) @@ -709,266 +1306,528 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', for jdx in clist: dist = numpy.linalg.norm(coord[idx] - coord[jdx]) - if (dist <= bond) or (interlayer and dist in inter_layer_dict[idx][1] \ - and jdx in inter_layer_dict[idx][0] and \ - dist < perpend_dist*ang2bohr and not jdx in pedg): - if not jdx in clist_check: - flist.append(jdx) - pedg.append(jdx) - if dist > bond: continue - if be_type=='be3' or be_type == 'be4': - if jdx in lunit: - lmin1 = runit[numpy.where(lunit==jdx)[0]] - if not twoD: - flist.extend(lmin1) - lsts.extend(lmin1) - for kdx in lmin1: - if twoD: - klsts.append(nk1*(nk2-1)+1) - else: - klsts.append(nk1) - else: - add_check_k(lmin1, flist, lsts, klsts, nk1*(nk2-1)+1) - if be_type == 'be4': - for kdx, k in enumerate(coord): - if kdx == jdx or kdx in lmin1 or cell.atom_pure_symbol(kdx) == 'H': - continue - dist = numpy.linalg.norm(coord[lmin1] - k) - if dist <= bond: - if kdx in lsts and klsts[lsts.index(kdx)] == nk1*(nk2-1)+1 \ - and twoD: continue - flist.append(kdx) - lsts.append(kdx) - if twoD: - klsts.append(nk1*(nk2-1)+1) - else: - klsts.append(nk1) - for lsidx in lmin1: - if lsidx in lunit or lsidx in munit: warn_large_fragment() - if lsidx in uunit: - surround(cell, lsidx, uunit, dunit, flist, coord, 'be3', - lsts, klsts, nk1*(nk2-1)+2, bond=bond) - if lsidx in dunit: - surround(cell, lsidx, dunit, uunit, flist, coord, 'be3', - lsts, klsts, nk1*nk2, bond=bond) - if lsidx in tunit: - surround(cell, lsidx, tunit, munit, flist, coord, 'be3', - lsts, klsts, 2, bond=bond) - if jdx in runit: - rmin1 = lunit[numpy.where(runit==jdx)[0]] - if not twoD: - flist.extend(rmin1) - rsts.extend(rmin1) - for kdx in rmin1: - if twoD: - krsts.append(nk1+1) - else: - krsts.append(2) - else: - add_check_k(rmin1, flist, rsts, krsts, nk1+1) - if be_type == 'be4': - for kdx, k in enumerate(coord): - if kdx == jdx or kdx in rmin1 or cell.atom_pure_symbol(kdx) == 'H': - continue - dist = numpy.linalg.norm(coord[rmin1] - k) - if dist <= bond: - if kdx in rsts and krsts[rsts.index(kdx)] == nk1+1 \ - and twoD: continue - flist.append(kdx) - rsts.append(kdx) - if twoD: - krsts.append(nk1+1) - else: - krsts.append(2) - for rsidx in rmin1: - if rsidx in runit or rsidx in tunit: warn_large_fragment() - if rsidx in munit: - surround(cell, rsidx, munit, tunit, flist, coord, 'be3', - rsts, krsts, nk1, bond=bond) - if rsidx in uunit: - surround(cell, rsidx, uunit, dunit, flist, coord, 'be3', - rsts, krsts, nk1+2, bond=bond) - if rsidx in dunit: - surround(cell, rsidx, dunit, uunit, flist, coord, 'be3', - rsts, krsts, nk1*2, bond=bond) - - if jdx in uunit: - umin1 = dunit[numpy.where(uunit==jdx)[0]] - add_check_k(umin1, flist, usts, kusts, 2) - if be_type == 'be4': - for kdx, k in enumerate(coord): - if kdx == jdx or kdx in umin1 or \ - cell.atom_pure_symbol(kdx) == 'H': - continue - dist = numpy.linalg.norm(coord[umin1] - k) - if dist <= bond: - if kdx in usts and kusts[usts.index(kdx)] == 2 \ - and twoD: continue - flist.append(kdx) - usts.append(kdx) - kusts.append(2) - for usidx in umin1: - if usidx in uunit or usidx in tunit: warn_large_fragment() - if usidx in lunit: - surround(cell, usidx, lunit, runit, flist, coord, 'be3', - usts, kusts, nk1*(nk2-1)+2, bond=bond) - if usidx in runit: - surround(cell, usidx, runit, lunit, flist, coord, 'be3', - usts, kusts, nk1+2, bond=bond) - if usidx in munit: - surround(cell, usidx, munit, tunit, flist, coord, 'be3', - usts, kusts, nk1*(nk2-1)+1, bond=bond) - if jdx in dunit: - dmin1 = uunit[numpy.where(dunit==jdx)[0]] - add_check_k(dmin1, flist, dsts, kdsts, nk1) - - if be_type == 'be4': - for kdx, k in enumerate(coord): - if kdx == jdx or kdx in dmin1 or \ - cell.atom_pure_symbol(kdx) == 'H': - continue - dist = numpy.linalg.norm(coord[dmin1] - k) - if dist <= bond: - if kdx in dsts and kdsts[dsts.index(kdx)] == nk1 \ - and twoD: continue - flist.append(kdx) - dsts.append(kdx) - kdsts.append(nk1) - for dsidx in dmin1: - if dsidx in munit or dsidx in dunit: warn_large_fragment() - if dsidx in lunit: - surround(cell, dsidx, lunit, runit, flist, coord, 'be3', - dsts, kdsts, nk1*nk2, bond=bond) - if dsidx in runit: - surround(cell, dsidx, runit, lunit, flist, coord, 'be3', - dsts, kdsts, nk1*2, bond=bond) - if dsidx in tunit: - surround(cell, dsidx, tunit, munit, flist, coord, 'be3', - dsts, kdsts, nk1+1, bond=bond) - if jdx in munit:# - mmin1 = tunit[numpy.where(munit==jdx)[0]] - add_check_k(mmin1, flist, msts, kmsts, nk1*nk2) - if be_type == 'be4': - for kdx, k in enumerate(coord): - if kdx == jdx or kdx in mmin1 or \ - cell.atom_pure_symbol(kdx) == 'H': - continue - dist = numpy.linalg.norm(coord[mmin1] - k) - if dist <= bond: - if kdx in msts and kmsts[msts.index(kdx)] == nk1*nk2 \ - and twoD: continue - flist.append(kdx) - msts.append(kdx) - kmsts.append(nk1*nk2) - for msidx in mmin1: - if msidx in lunit or msidx in munit or msidx in dunit: - warn_large_fragment() - if msidx in runit: - surround(cell, msidx, runit, lunit, flist, coord, 'be3', - msts, kmsts, nk1, bond=bond) - if msidx in uunit: - surround(cell, msidx, uunit, dunit, flist, coord, 'be3', - msts, kmsts, nk1*(nk2-1)+1, bond=bond) - - if jdx in tunit: - tmin1 = munit[numpy.where(tunit==jdx)[0]] - add_check_k(tmin1, flist, tsts, ktsts, nk1+2) - - if be_type == 'be4': - for kdx, k in enumerate(coord): - if kdx == jdx or kdx in tmin1 or \ - cell.atom_pure_symbol(kdx) == 'H': - continue - dist = numpy.linalg.norm(coord[tmin1] - k) - if dist <= bond: - if kdx in tsts and ktsts[tsts.index(kdx)] == nk1+2 and \ - twoD: continue - flist.append(kdx) - tsts.append(kdx) - ktsts.append(nk1+2) - for tsidx in tmin1: - if tsidx in runit or tsidx in uunit or tsidx in tunit: - warn_large_fragment() - if tsidx in lunit: - surround(cell, tsidx, lunit, runit, flist, coord, 'be3', - tsts, ktsts, 2, bond=bond) - if tsidx in dunit: - surround(cell, tsidx, dunit, uunit, flist, coord, 'be3', - tsts, ktsts, nk1+1, bond=bond) - - - for kdx in clist: - if not kdx == jdx: - dist = numpy.linalg.norm(coord[jdx] - coord[kdx]) - if (dist <= bond) or (interlayer and dist in inter_layer_dict[jdx][1] \ - and kdx in inter_layer_dict[jdx][0] and \ - dist < perpend_dist*ang2bohr): - if not kdx in pedg and not kdx in clist_check: - flist.append(kdx) - pedg.append(kdx) - if be_type=='be4': - if kdx in lunit: - lmin1 = runit[numpy.where(lunit==kdx)[0]] - for zdx in lmin1: - if zdx in lsts and klsts[lsts.index(zdx)] == nk1*(nk2-1)+1 \ - and twoD: continue - flist.append(zdx) - lsts.append(zdx) - if twoD: - klsts.append(nk1*(nk2-1)+1) - else: - klsts.append(nk1) - if kdx in runit: - rmin1 = lunit[numpy.where(runit==kdx)[0]] - for zdx in rmin1: - if zdx in rsts and krsts[rsts.index(zdx)] == nk1+1 and twoD: - continue - flist.append(zdx) - rsts.append(zdx) - if twoD: - krsts.append(nk1+1) - else: - krsts.append(2) - if kdx in uunit: - umin1 = dunit[numpy.where(uunit==kdx)[0]] - for zdx in umin1: - if zdx in usts and kusts[usts.index(zdx)] == 2 and twoD: - continue - flist.append(zdx) - usts.append(zdx) - kusts.append(2) - if kdx in dunit: - dmin1 = uunit[numpy.where(dunit==kdx)[0]] - for zdx in dmin1: - if zdx in dsts and kdsts[dsts.index(zdx)] == nk1 and twoD: - continue - flist.append(zdx) - dsts.append(zdx) - kdsts.append(nk1) - if kdx in munit: - mmin1 = tunit[numpy.where(munit==kdx)[0]] - for zdx in mmin1: - if zdx in msts and kmsts[msts.index(zdx)] == nk1*nk2 and twoD: - continue - flist.append(zdx) - msts.append(zdx) - kmsts.append(nk1*nk2) - if kdx in tunit: - tmin1 = munit[numpy.where(tunit==kdx)[0]] - for zdx in tmin1: - if zdx in tsts and ktsts[tsts.index(zdx)] == nk1+2 and twoD: - continue - flist.append(zdx) - tsts.append(zdx) - ktsts.append(nk1+2) - - for ldx, l in enumerate(coord): - if ldx==kdx or ldx==jdx or cell.atom_pure_symbol(ldx) == 'H'or ldx in pedg: - continue - dist = numpy.linalg.norm(coord[kdx] - l) - if dist <= bond: - flist.append(ldx) - pedg.append(ldx) + if (dist <= bond) or ( + interlayer + and dist in inter_layer_dict[idx][1] + and jdx in inter_layer_dict[idx][0] + and dist < perpend_dist * ang2bohr + and not jdx in pedg + ): + if not jdx in clist_check: + flist.append(jdx) + pedg.append(jdx) + if dist > bond: + continue + if be_type == "be3" or be_type == "be4": + if jdx in lunit: + lmin1 = runit[numpy.where(lunit == jdx)[0]] + if not twoD: + flist.extend(lmin1) + lsts.extend(lmin1) + for kdx in lmin1: + if twoD: + klsts.append(nk1 * (nk2 - 1) + 1) + else: + klsts.append(nk1) + else: + add_check_k(lmin1, flist, lsts, klsts, nk1 * (nk2 - 1) + 1) + if be_type == "be4": + for kdx, k in enumerate(coord): + if ( + kdx == jdx + or kdx in lmin1 + or cell.atom_pure_symbol(kdx) == "H" + ): + continue + dist = numpy.linalg.norm(coord[lmin1] - k) + if dist <= bond: + if ( + kdx in lsts + and klsts[lsts.index(kdx)] + == nk1 * (nk2 - 1) + 1 + and twoD + ): + continue + flist.append(kdx) + lsts.append(kdx) + if twoD: + klsts.append(nk1 * (nk2 - 1) + 1) + else: + klsts.append(nk1) + for lsidx in lmin1: + if lsidx in lunit or lsidx in munit: + warn_large_fragment() + if lsidx in uunit: + surround( + cell, + lsidx, + uunit, + dunit, + flist, + coord, + "be3", + lsts, + klsts, + nk1 * (nk2 - 1) + 2, + bond=bond, + ) + if lsidx in dunit: + surround( + cell, + lsidx, + dunit, + uunit, + flist, + coord, + "be3", + lsts, + klsts, + nk1 * nk2, + bond=bond, + ) + if lsidx in tunit: + surround( + cell, + lsidx, + tunit, + munit, + flist, + coord, + "be3", + lsts, + klsts, + 2, + bond=bond, + ) + if jdx in runit: + rmin1 = lunit[numpy.where(runit == jdx)[0]] + if not twoD: + flist.extend(rmin1) + rsts.extend(rmin1) + for kdx in rmin1: + if twoD: + krsts.append(nk1 + 1) + else: + krsts.append(2) + else: + add_check_k(rmin1, flist, rsts, krsts, nk1 + 1) + if be_type == "be4": + for kdx, k in enumerate(coord): + if ( + kdx == jdx + or kdx in rmin1 + or cell.atom_pure_symbol(kdx) == "H" + ): + continue + dist = numpy.linalg.norm(coord[rmin1] - k) + if dist <= bond: + if ( + kdx in rsts + and krsts[rsts.index(kdx)] == nk1 + 1 + and twoD + ): + continue + flist.append(kdx) + rsts.append(kdx) + if twoD: + krsts.append(nk1 + 1) + else: + krsts.append(2) + for rsidx in rmin1: + if rsidx in runit or rsidx in tunit: + warn_large_fragment() + if rsidx in munit: + surround( + cell, + rsidx, + munit, + tunit, + flist, + coord, + "be3", + rsts, + krsts, + nk1, + bond=bond, + ) + if rsidx in uunit: + surround( + cell, + rsidx, + uunit, + dunit, + flist, + coord, + "be3", + rsts, + krsts, + nk1 + 2, + bond=bond, + ) + if rsidx in dunit: + surround( + cell, + rsidx, + dunit, + uunit, + flist, + coord, + "be3", + rsts, + krsts, + nk1 * 2, + bond=bond, + ) + + if jdx in uunit: + umin1 = dunit[numpy.where(uunit == jdx)[0]] + add_check_k(umin1, flist, usts, kusts, 2) + if be_type == "be4": + for kdx, k in enumerate(coord): + if ( + kdx == jdx + or kdx in umin1 + or cell.atom_pure_symbol(kdx) == "H" + ): + continue + dist = numpy.linalg.norm(coord[umin1] - k) + if dist <= bond: + if ( + kdx in usts + and kusts[usts.index(kdx)] == 2 + and twoD + ): + continue + flist.append(kdx) + usts.append(kdx) + kusts.append(2) + for usidx in umin1: + if usidx in uunit or usidx in tunit: + warn_large_fragment() + if usidx in lunit: + surround( + cell, + usidx, + lunit, + runit, + flist, + coord, + "be3", + usts, + kusts, + nk1 * (nk2 - 1) + 2, + bond=bond, + ) + if usidx in runit: + surround( + cell, + usidx, + runit, + lunit, + flist, + coord, + "be3", + usts, + kusts, + nk1 + 2, + bond=bond, + ) + if usidx in munit: + surround( + cell, + usidx, + munit, + tunit, + flist, + coord, + "be3", + usts, + kusts, + nk1 * (nk2 - 1) + 1, + bond=bond, + ) + if jdx in dunit: + dmin1 = uunit[numpy.where(dunit == jdx)[0]] + add_check_k(dmin1, flist, dsts, kdsts, nk1) + + if be_type == "be4": + for kdx, k in enumerate(coord): + if ( + kdx == jdx + or kdx in dmin1 + or cell.atom_pure_symbol(kdx) == "H" + ): + continue + dist = numpy.linalg.norm(coord[dmin1] - k) + if dist <= bond: + if ( + kdx in dsts + and kdsts[dsts.index(kdx)] == nk1 + and twoD + ): + continue + flist.append(kdx) + dsts.append(kdx) + kdsts.append(nk1) + for dsidx in dmin1: + if dsidx in munit or dsidx in dunit: + warn_large_fragment() + if dsidx in lunit: + surround( + cell, + dsidx, + lunit, + runit, + flist, + coord, + "be3", + dsts, + kdsts, + nk1 * nk2, + bond=bond, + ) + if dsidx in runit: + surround( + cell, + dsidx, + runit, + lunit, + flist, + coord, + "be3", + dsts, + kdsts, + nk1 * 2, + bond=bond, + ) + if dsidx in tunit: + surround( + cell, + dsidx, + tunit, + munit, + flist, + coord, + "be3", + dsts, + kdsts, + nk1 + 1, + bond=bond, + ) + if jdx in munit: # + mmin1 = tunit[numpy.where(munit == jdx)[0]] + add_check_k(mmin1, flist, msts, kmsts, nk1 * nk2) + if be_type == "be4": + for kdx, k in enumerate(coord): + if ( + kdx == jdx + or kdx in mmin1 + or cell.atom_pure_symbol(kdx) == "H" + ): + continue + dist = numpy.linalg.norm(coord[mmin1] - k) + if dist <= bond: + if ( + kdx in msts + and kmsts[msts.index(kdx)] == nk1 * nk2 + and twoD + ): + continue + flist.append(kdx) + msts.append(kdx) + kmsts.append(nk1 * nk2) + for msidx in mmin1: + if msidx in lunit or msidx in munit or msidx in dunit: + warn_large_fragment() + if msidx in runit: + surround( + cell, + msidx, + runit, + lunit, + flist, + coord, + "be3", + msts, + kmsts, + nk1, + bond=bond, + ) + if msidx in uunit: + surround( + cell, + msidx, + uunit, + dunit, + flist, + coord, + "be3", + msts, + kmsts, + nk1 * (nk2 - 1) + 1, + bond=bond, + ) + + if jdx in tunit: + tmin1 = munit[numpy.where(tunit == jdx)[0]] + add_check_k(tmin1, flist, tsts, ktsts, nk1 + 2) + + if be_type == "be4": + for kdx, k in enumerate(coord): + if ( + kdx == jdx + or kdx in tmin1 + or cell.atom_pure_symbol(kdx) == "H" + ): + continue + dist = numpy.linalg.norm(coord[tmin1] - k) + if dist <= bond: + if ( + kdx in tsts + and ktsts[tsts.index(kdx)] == nk1 + 2 + and twoD + ): + continue + flist.append(kdx) + tsts.append(kdx) + ktsts.append(nk1 + 2) + for tsidx in tmin1: + if tsidx in runit or tsidx in uunit or tsidx in tunit: + warn_large_fragment() + if tsidx in lunit: + surround( + cell, + tsidx, + lunit, + runit, + flist, + coord, + "be3", + tsts, + ktsts, + 2, + bond=bond, + ) + if tsidx in dunit: + surround( + cell, + tsidx, + dunit, + uunit, + flist, + coord, + "be3", + tsts, + ktsts, + nk1 + 1, + bond=bond, + ) + + for kdx in clist: + if not kdx == jdx: + dist = numpy.linalg.norm(coord[jdx] - coord[kdx]) + if (dist <= bond) or ( + interlayer + and dist in inter_layer_dict[jdx][1] + and kdx in inter_layer_dict[jdx][0] + and dist < perpend_dist * ang2bohr + ): + if not kdx in pedg and not kdx in clist_check: + flist.append(kdx) + pedg.append(kdx) + if be_type == "be4": + if kdx in lunit: + lmin1 = runit[numpy.where(lunit == kdx)[0]] + for zdx in lmin1: + if ( + zdx in lsts + and klsts[lsts.index(zdx)] + == nk1 * (nk2 - 1) + 1 + and twoD + ): + continue + flist.append(zdx) + lsts.append(zdx) + if twoD: + klsts.append(nk1 * (nk2 - 1) + 1) + else: + klsts.append(nk1) + if kdx in runit: + rmin1 = lunit[numpy.where(runit == kdx)[0]] + for zdx in rmin1: + if ( + zdx in rsts + and krsts[rsts.index(zdx)] == nk1 + 1 + and twoD + ): + continue + flist.append(zdx) + rsts.append(zdx) + if twoD: + krsts.append(nk1 + 1) + else: + krsts.append(2) + if kdx in uunit: + umin1 = dunit[numpy.where(uunit == kdx)[0]] + for zdx in umin1: + if ( + zdx in usts + and kusts[usts.index(zdx)] == 2 + and twoD + ): + continue + flist.append(zdx) + usts.append(zdx) + kusts.append(2) + if kdx in dunit: + dmin1 = uunit[numpy.where(dunit == kdx)[0]] + for zdx in dmin1: + if ( + zdx in dsts + and kdsts[dsts.index(zdx)] == nk1 + and twoD + ): + continue + flist.append(zdx) + dsts.append(zdx) + kdsts.append(nk1) + if kdx in munit: + mmin1 = tunit[numpy.where(munit == kdx)[0]] + for zdx in mmin1: + if ( + zdx in msts + and kmsts[msts.index(zdx)] == nk1 * nk2 + and twoD + ): + continue + flist.append(zdx) + msts.append(zdx) + kmsts.append(nk1 * nk2) + if kdx in tunit: + tmin1 = munit[numpy.where(tunit == kdx)[0]] + for zdx in tmin1: + if ( + zdx in tsts + and ktsts[tsts.index(zdx)] == nk1 + 2 + and twoD + ): + continue + flist.append(zdx) + tsts.append(zdx) + ktsts.append(nk1 + 2) + + for ldx, l in enumerate(coord): + if ( + ldx == kdx + or ldx == jdx + or cell.atom_pure_symbol(ldx) == "H" + or ldx in pedg + ): + continue + dist = numpy.linalg.norm(coord[kdx] - l) + if dist <= bond: + flist.append(ldx) + pedg.append(ldx) lsites.append(lsts) rsites.append(rsts) @@ -988,67 +1847,94 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', hlist = [[] for i in coord] for idx, i in enumerate(normlist): - if cell.atom_pure_symbol(idx) == 'H': - + if cell.atom_pure_symbol(idx) == "H": tmplist = normlist - i tmplist = list(tmplist) clist = [] - for jdx,j in enumerate(tmplist): - if not idx==jdx and not cell.atom_pure_symbol(jdx) == 'H': - if abs(j)< normdist: + for jdx, j in enumerate(tmplist): + if not idx == jdx and not cell.atom_pure_symbol(jdx) == "H": + if abs(j) < normdist: clist.append(jdx) for jdx in clist: - dist = numpy.linalg.norm(coord[idx] - coord[jdx]) if dist <= hbond: hlist[jdx].append(idx) if print_frags: print(flush=True) - print('Fragment sites',flush=True) - print('--------------------------',flush=True) - print('Fragment | Center | Edges ',flush=True) - print('--------------------------',flush=True) - - for idx,i in enumerate(Frag): - print(' {:>4} | {:>5} |'.format(idx, - cell.atom_pure_symbol(cen[idx])+str(cen[idx]+1)), - end=' ', flush=True) + print("Fragment sites", flush=True) + print("--------------------------", flush=True) + print("Fragment | Center | Edges ", flush=True) + print("--------------------------", flush=True) + + for idx, i in enumerate(Frag): + print( + " {:>4} | {:>5} |".format( + idx, cell.atom_pure_symbol(cen[idx]) + str(cen[idx] + 1) + ), + end=" ", + flush=True, + ) for j in hlist[cen[idx]]: - print(' {:>5} '.format('*'+cell.atom_pure_symbol(j)+str(j+1)),end=' ', flush=True) + print( + " {:>5} ".format("*" + cell.atom_pure_symbol(j) + str(j + 1)), + end=" ", + flush=True, + ) for j in i: - if j == cen[idx]: continue - print(' {:>5} '.format(cell.atom_pure_symbol(j)+str(j+1)),end=' ', flush=True) + if j == cen[idx]: + continue + print( + " {:>5} ".format(cell.atom_pure_symbol(j) + str(j + 1)), + end=" ", + flush=True, + ) for k in hlist[j]: - print(' {:>5} '.format(cell.atom_pure_symbol(k)+str(k+1)),end=' ', flush=True) + print( + " {:>5} ".format(cell.atom_pure_symbol(k) + str(k + 1)), + end=" ", + flush=True, + ) print(flush=True) - print('--------------------------',flush=True) - print(' No. of fragments : ',len(Frag),flush=True) - print('*H : Center H atoms (printed as Edges above.)', flush=True) + print("--------------------------", flush=True) + print(" No. of fragments : ", len(Frag), flush=True) + print("*H : Center H atoms (printed as Edges above.)", flush=True) print(flush=True) if write_geom: - w = open('fragments.xyz','w') - for idx,i in enumerate(Frag): - w.write(str(len(i)+len(hlist[cen[idx]])+len(hlist[j]))+'\n') - w.write('Fragment - '+str(idx)+'\n') + w = open("fragments.xyz", "w") + for idx, i in enumerate(Frag): + w.write(str(len(i) + len(hlist[cen[idx]]) + len(hlist[j])) + "\n") + w.write("Fragment - " + str(idx) + "\n") for j in hlist[cen[idx]]: - w.write(' {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n'.format(cell.atom_pure_symbol(j), - coord[j][0]/ang2bohr, - coord[j][1]/ang2bohr, - coord[j][2]/ang2bohr)) + w.write( + " {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n".format( + cell.atom_pure_symbol(j), + coord[j][0] / ang2bohr, + coord[j][1] / ang2bohr, + coord[j][2] / ang2bohr, + ) + ) for j in i: - w.write(' {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n'.format(cell.atom_pure_symbol(j), - coord[j][0]/ang2bohr, - coord[j][1]/ang2bohr, - coord[j][2]/ang2bohr)) + w.write( + " {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n".format( + cell.atom_pure_symbol(j), + coord[j][0] / ang2bohr, + coord[j][1] / ang2bohr, + coord[j][2] / ang2bohr, + ) + ) for k in hlist[j]: - w.write(' {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n'.format(cell.atom_pure_symbol(k), - coord[k][0]/ang2bohr, - coord[k][1]/ang2bohr, - coord[k][2]/ang2bohr)) + w.write( + " {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n".format( + cell.atom_pure_symbol(k), + coord[k][0] / ang2bohr, + coord[k][1] / ang2bohr, + coord[k][2] / ang2bohr, + ) + ) w.close() if not valence_basis is None: @@ -1070,8 +1956,7 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', hshift = [0 for i in coord] for adx in range(cell.natm): - - if not cell.atom_pure_symbol(adx) == 'H': + if not cell.atom_pure_symbol(adx) == "H": bas = baslist[adx] start_ = bas[2] stop_ = bas[3] @@ -1082,8 +1967,9 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', if frozen_core: start_ -= coreshift ncore_ = core_list[adx] - stop_ -= coreshift+ncore_ - if pao: nbas2[adx] -= ncore_ + stop_ -= coreshift + ncore_ + if pao: + nbas2[adx] -= ncore_ coreshift += ncore_ b1list = [i for i in range(start_, stop_)] @@ -1094,9 +1980,7 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', hsites = [[] for i in coord] nbas2H = [0 for i in coord] for hdx, h in enumerate(hlist): - for hidx in h: - basH = baslist[hidx] startH = basH[2] stopH = basH[3] @@ -1126,15 +2010,22 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', mfrac = unitcell - int(unitcell) conmax = True nkcon = True - if gamma_2d or gamma_1d: conmax = False - Ns = Nsite+1 - uNs = Ns*unitcell + if gamma_2d or gamma_1d: + conmax = False + Ns = Nsite + 1 + uNs = Ns * unitcell Sites = sites__.copy() - for j_ in range(unitcell-1): + for j_ in range(unitcell - 1): for idx, i in enumerate(Sites): - if any (i>=unitcell*Ns + unitcell*Ns*j_): - sites__[idx] = [ (nk1*Ns)+j+(nk1*Ns*j_ - unitcell*Ns*j_) - (unitcell*Ns) for j in i] + if any(i >= unitcell * Ns + unitcell * Ns * j_): + sites__[idx] = [ + (nk1 * Ns) + + j + + (nk1 * Ns * j_ - unitcell * Ns * j_) + - (unitcell * Ns) + for j in i + ] for idx, i in enumerate(Frag): ftmp = [] @@ -1142,131 +2033,148 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', indix = 0 edind = [] edg = [] - for jdx_,jdx in enumerate(lsites[idx]): + for jdx_, jdx in enumerate(lsites[idx]): edg.append(jdx) if not conmax: frglist = sites__[jdx].copy() frglist.extend(hsites[jdx]) elif nkcon: - numk = klsites[idx][jdx_]-1 - frglist = kfrag_func(sites__[jdx]+hsites[jdx], numk, nk1, uNs, Ns) + numk = klsites[idx][jdx_] - 1 + frglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - frglist = [pq-max_site-1 for pq in sites__[jdx]] - frglist.extend([pq-max_site-1 for pq in hsites[jdx]]) + frglist = [pq - max_site - 1 for pq in sites__[jdx]] + frglist.extend([pq - max_site - 1 for pq in hsites[jdx]]) ftmp.extend(frglist) - ls = len(sites__[jdx])+len(hsites[jdx]) + ls = len(sites__[jdx]) + len(hsites[jdx]) if not pao: if not conmax: edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) elif nkcon: - edglist = kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + edglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - edglist = [pq-max_site-1 for pq in sites__[jdx]] - edglist.extend([pq-max_site-1 for pq in hsites[jdx]]) + edglist = [pq - max_site - 1 for pq in sites__[jdx]] + edglist.extend([pq - max_site - 1 for pq in hsites[jdx]]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: if not conmax: - edglist = sites__[jdx].copy()[:nbas2[jdx]] - edglist.extend( hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx].copy()[: nbas2[jdx]] + edglist.extend(hsites[jdx][: nbas2H[jdx]]) elif nkcon: - edglist = kfrag_func(sites__[jdx][:nbas2[jdx]]+hsites[jdx][:nbas2H[jdx]], - numk, nk1, uNs, Ns) + edglist = kfrag_func( + sites__[jdx][: nbas2[jdx]] + hsites[jdx][: nbas2H[jdx]], + numk, + nk1, + uNs, + Ns, + ) else: - edglist = [pq-max_site-1 for pq in sites__[jdx]][:nbas2[jdx]] - edglist.extend([pq-max_site-1 for pq in hsites[jdx]][:nbas2H[jdx]]) + edglist = [pq - max_site - 1 for pq in sites__[jdx]][: nbas2[jdx]] + edglist.extend( + [pq - max_site - 1 for pq in hsites[jdx]][: nbas2H[jdx]] + ) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls - for jdx_,jdx in enumerate(usites[idx]): + for jdx_, jdx in enumerate(usites[idx]): edg.append(jdx) if not conmax: frglist = sites__[jdx].copy() frglist.extend(hsites[jdx]) elif nkcon: - numk = kusites[idx][jdx_]-1 - frglist = kfrag_func(sites__[jdx]+hsites[jdx], numk, nk1, uNs, Ns) + numk = kusites[idx][jdx_] - 1 + frglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - frglist = [pq-max_site-1 for pq in sites__[jdx]] - frglist.extend([pq-max_site-1 for pq in hsites[jdx]]) + frglist = [pq - max_site - 1 for pq in sites__[jdx]] + frglist.extend([pq - max_site - 1 for pq in hsites[jdx]]) ftmp.extend(frglist) - ls = len(sites__[jdx])+len(hsites[jdx]) + ls = len(sites__[jdx]) + len(hsites[jdx]) if not pao: if not conmax: edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) elif nkcon: - edglist =kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + edglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - edglist = [pq-max_site-1 for pq in sites__[jdx]] - edglist.extend([pq-max_site-1 for pq in hsites[jdx]]) + edglist = [pq - max_site - 1 for pq in sites__[jdx]] + edglist.extend([pq - max_site - 1 for pq in hsites[jdx]]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: if not conmax: - edglist =sites__[jdx].copy()[:nbas2[jdx]] - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx].copy()[: nbas2[jdx]] + edglist.extend(hsites[jdx][: nbas2H[jdx]]) elif nkcon: - edglist =kfrag_func(sites__[jdx][:nbas2[jdx]]+hsites[jdx][:nbas2H[jdx]], - numk, nk1, uNs, Ns) + edglist = kfrag_func( + sites__[jdx][: nbas2[jdx]] + hsites[jdx][: nbas2H[jdx]], + numk, + nk1, + uNs, + Ns, + ) else: - edglist = [pq-max_site-1 for pq in sites__[jdx]][:nbas2[jdx]] - edglist.extend([pq-max_site-1 for pq in hsites[jdx]][:nbas2H[jdx]]) + edglist = [pq - max_site - 1 for pq in sites__[jdx]][: nbas2[jdx]] + edglist.extend( + [pq - max_site - 1 for pq in hsites[jdx]][: nbas2H[jdx]] + ) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls - for jdx_,jdx in enumerate(msites[idx]): + for jdx_, jdx in enumerate(msites[idx]): edg.append(jdx) if not conmax: - frglist =sites__[jdx].copy() + frglist = sites__[jdx].copy() frglist.extend(hsites[jdx]) elif nkcon: - numk = kmsites[idx][jdx_]-1 - frglist = kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + numk = kmsites[idx][jdx_] - 1 + frglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - frglist = [pq-max_site-1 for pq in sites__[jdx]] - frglist.extend([pq-max_site-1 for pq in hsites[jdx]]) + frglist = [pq - max_site - 1 for pq in sites__[jdx]] + frglist.extend([pq - max_site - 1 for pq in hsites[jdx]]) ftmp.extend(frglist) - ls = len(sites__[jdx])+len(hsites[jdx]) + ls = len(sites__[jdx]) + len(hsites[jdx]) if not pao: if not conmax: edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) elif nkcon: - edglist =kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + edglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - edglist = [pq-max_site-1 for pq in sites__[jdx]] - edglist.extend([pq-max_site-1 for pq in hsites[jdx]]) + edglist = [pq - max_site - 1 for pq in sites__[jdx]] + edglist.extend([pq - max_site - 1 for pq in hsites[jdx]]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: if not conmax: - edglist = sites__[jdx].copy()[:nbas2[jdx]] - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx].copy()[: nbas2[jdx]] + edglist.extend(hsites[jdx][: nbas2H[jdx]]) elif nkcon: - edglist =kfrag_func(sites__[jdx][:nbas2[jdx]]+hsites[jdx][:nbas2H[jdx]], - numk, nk1, uNs, Ns) + edglist = kfrag_func( + sites__[jdx][: nbas2[jdx]] + hsites[jdx][: nbas2H[jdx]], + numk, + nk1, + uNs, + Ns, + ) else: - edglist = [pq-max_site-1 for pq in sites__[jdx]][:nbas2[jdx]] - edglist.extend([pq-max_site-1 for pq in hsites[jdx]][:nbas2H[jdx]]) + edglist = [pq - max_site - 1 for pq in sites__[jdx]][: nbas2[jdx]] + edglist.extend( + [pq - max_site - 1 for pq in hsites[jdx]][: nbas2H[jdx]] + ) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls @@ -1274,13 +2182,13 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', frglist.extend(hsites[cen[idx]]) ftmp.extend(frglist) - ls = len(sites__[cen[idx]])+len(hsites[cen[idx]]) + ls = len(sites__[cen[idx]]) + len(hsites[cen[idx]]) if not pao: - centerf_idx.append([pq for pq in range(indix,indix+ls)]) + centerf_idx.append([pq for pq in range(indix, indix + ls)]) else: - cntlist = sites__[cen[idx]].copy()[:nbas2[cen[idx]]] - cntlist.extend(hsites[cen[idx]][:nbas2H[cen[idx]]]) - ind__ = [ indix+frglist.index(pq) for pq in cntlist] + cntlist = sites__[cen[idx]].copy()[: nbas2[cen[idx]]] + cntlist.extend(hsites[cen[idx]][: nbas2H[cen[idx]]]) + ind__ = [indix + frglist.index(pq) for pq in cntlist] centerf_idx.append(ind__) indix += ls @@ -1294,142 +2202,160 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: - edglist = sites__[jdx][:nbas2[jdx]].copy() - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx][: nbas2[jdx]].copy() + edglist.extend(hsites[jdx][: nbas2H[jdx]]) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls - - for jdx_,jdx in enumerate(rsites[idx]): + for jdx_, jdx in enumerate(rsites[idx]): edg.append(jdx) if not conmax: frglist = sites__[jdx].copy() frglist.extend(hsites[jdx]) elif nkcon: - numk = krsites[idx][jdx_]-1 - frglist = kfrag_func(sites__[jdx]+hsites[jdx], numk, - nk1, uNs, Ns, debug2=True) + numk = krsites[idx][jdx_] - 1 + frglist = kfrag_func( + sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns, debug2=True + ) else: - frglist = [pq+max_site+1 for pq in sites__[jdx]] - frglist.extend([pq+max_site+1 for pq in hsites[jdx]]) + frglist = [pq + max_site + 1 for pq in sites__[jdx]] + frglist.extend([pq + max_site + 1 for pq in hsites[jdx]]) ftmp.extend(frglist) - ls = len(sites__[jdx])+len(hsites[jdx]) + ls = len(sites__[jdx]) + len(hsites[jdx]) if not pao: if not conmax: edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) elif nkcon: - edglist = kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + edglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - edglist = [pq-max_site+1 for pq in sites__[jdx]] - edglist.extend([pq-max_site+1 for pq in hsites[jdx]]) + edglist = [pq - max_site + 1 for pq in sites__[jdx]] + edglist.extend([pq - max_site + 1 for pq in hsites[jdx]]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: if not conmax: - edglist =sites__[jdx].copy()[:nbas2[jdx]] - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx].copy()[: nbas2[jdx]] + edglist.extend(hsites[jdx][: nbas2H[jdx]]) elif nkcon: - edglist = kfrag_func(sites__[jdx][:nbas2[jdx]]+hsites[jdx][:nbas2H[jdx]], - numk, nk1, uNs, Ns) + edglist = kfrag_func( + sites__[jdx][: nbas2[jdx]] + hsites[jdx][: nbas2H[jdx]], + numk, + nk1, + uNs, + Ns, + ) else: - edglist = [pq+max_site+1 for pq in sites__[jdx]][:nbas2[jdx]] - edglist.extend([pq+max_site+1 for pq in hsites[jdx]][:nbas2H[jdx]]) + edglist = [pq + max_site + 1 for pq in sites__[jdx]][: nbas2[jdx]] + edglist.extend( + [pq + max_site + 1 for pq in hsites[jdx]][: nbas2H[jdx]] + ) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls - for jdx_,jdx in enumerate(dsites[idx]): + for jdx_, jdx in enumerate(dsites[idx]): edg.append(jdx) if not conmax: frglist = sites__[jdx].copy() frglist.extend(hsites[jdx]) elif nkcon: - numk = kdsites[idx][jdx_]-1 - frglist = kfrag_func(sites__[jdx]+hsites[jdx], numk, nk1, uNs, Ns) + numk = kdsites[idx][jdx_] - 1 + frglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - frglist = [pq+max_site+1 for pq in sites__[jdx]] - frglist.extend([pq+max_site+1 for pq in hsites[jdx]]) + frglist = [pq + max_site + 1 for pq in sites__[jdx]] + frglist.extend([pq + max_site + 1 for pq in hsites[jdx]]) ftmp.extend(frglist) - ls = len(sites__[jdx])+len(hsites[jdx]) + ls = len(sites__[jdx]) + len(hsites[jdx]) if not pao: if not conmax: edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) elif nkcon: - edglist = kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + edglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - edglist = [pq-max_site+1 for pq in sites__[jdx]] - edglist.extend([pq-max_site+1 for pq in hsites[jdx]]) + edglist = [pq - max_site + 1 for pq in sites__[jdx]] + edglist.extend([pq - max_site + 1 for pq in hsites[jdx]]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: if not conmax: - edglist = sites__[jdx].copy()[:nbas2[jdx]] - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx].copy()[: nbas2[jdx]] + edglist.extend(hsites[jdx][: nbas2H[jdx]]) elif nkcon: - edglist = kfrag_func(sites__[jdx][:nbas2[jdx]]+hsites[jdx][:nbas2H[jdx]], - numk, nk1, uNs, Ns) + edglist = kfrag_func( + sites__[jdx][: nbas2[jdx]] + hsites[jdx][: nbas2H[jdx]], + numk, + nk1, + uNs, + Ns, + ) else: - edglist = [pq+max_site+1 for pq in sites__[jdx]][:nbas2[jdx]] - edglist.extend([pq+max_site+1 for pq in hsites[jdx]][:nbas2H[jdx]]) + edglist = [pq + max_site + 1 for pq in sites__[jdx]][: nbas2[jdx]] + edglist.extend( + [pq + max_site + 1 for pq in hsites[jdx]][: nbas2H[jdx]] + ) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls - for jdx_,jdx in enumerate(tsites[idx]): + for jdx_, jdx in enumerate(tsites[idx]): edg.append(jdx) if not conmax: frglist = sites__[jdx].copy() frglist.extend(hsites[jdx]) elif nkcon: - numk = ktsites[idx][jdx_]-1 - frglist = kfrag_func(sites__[jdx]+hsites[jdx], numk, nk1, uNs, Ns) + numk = ktsites[idx][jdx_] - 1 + frglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - frglist = [pq+max_site+1 for pq in sites__[jdx]] - frglist.extend([pq+max_site+1 for pq in hsites[jdx]]) + frglist = [pq + max_site + 1 for pq in sites__[jdx]] + frglist.extend([pq + max_site + 1 for pq in hsites[jdx]]) ftmp.extend(frglist) - ls = len(sites__[jdx])+len(hsites[jdx]) + ls = len(sites__[jdx]) + len(hsites[jdx]) if not pao: if not conmax: edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) elif nkcon: - edglist = kfrag_func(sites__[jdx]+hsites[jdx], - numk, nk1, uNs, Ns) + edglist = kfrag_func(sites__[jdx] + hsites[jdx], numk, nk1, uNs, Ns) else: - edglist = [pq-max_site+1 for pq in sites__[jdx]] - edglist.extend([pq-max_site+1 for pq in hsites[jdx]]) + edglist = [pq - max_site + 1 for pq in sites__[jdx]] + edglist.extend([pq - max_site + 1 for pq in hsites[jdx]]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: if not conmax: - edglist = sites__[jdx].copy()[:nbas2[jdx]] - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx].copy()[: nbas2[jdx]] + edglist.extend(hsites[jdx][: nbas2H[jdx]]) elif nkcon: - edglist =kfrag_func(sites__[jdx][:nbas2[jdx]]+hsites[jdx][:nbas2H[jdx]], - numk, nk1, uNs, Ns) + edglist = kfrag_func( + sites__[jdx][: nbas2[jdx]] + hsites[jdx][: nbas2H[jdx]], + numk, + nk1, + uNs, + Ns, + ) else: - edglist = [pq+max_site+1 for pq in sites__[jdx]][:nbas2[jdx]] - edglist.extend([pq+max_site+1 for pq in hsites[jdx]][:nbas2H[jdx]]) + edglist = [pq + max_site + 1 for pq in sites__[jdx]][: nbas2[jdx]] + edglist.extend( + [pq + max_site + 1 for pq in hsites[jdx]][: nbas2H[jdx]] + ) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls @@ -1438,7 +2364,6 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', edgsites.append(ftmpe) edge_idx.append(edind) - center = [] for ix in edge: cen_ = [] @@ -1447,10 +2372,9 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', center.append(cen_) Nfrag = len(fsites) - ebe_weight=[] + ebe_weight = [] # Use IAO+PAO for computing energy for ix, i in enumerate(fsites): - tmp_ = [i.index(pq) for pq in sites__[cen[ix]]] tmp_.extend([i.index(pq) for pq in hsites[cen[ix]]]) ebe_weight.append([1.0, tmp_]) @@ -1467,13 +2391,10 @@ def autogen(mol, kpt, frozen_core=True, be_type='be2', cntlist.extend(hsites[cen[j]]) idx.append([fsites[j].index(k) for k in cntlist]) else: - cntlist = sites__[cen[j]].copy()[:nbas2[cen[j]]] - cntlist.extend(hsites[cen[j]][:nbas2H[cen[j]]]) + cntlist = sites__[cen[j]].copy()[: nbas2[cen[j]]] + cntlist.extend(hsites[cen[j]][: nbas2H[cen[j]]]) idx.append([fsites[j].index(k) for k in cntlist]) center_idx.append(idx) - return(fsites, edgsites, center, edge_idx, center_idx, centerf_idx, ebe_weight) - - - + return (fsites, edgsites, center, edge_idx, center_idx, centerf_idx, ebe_weight) diff --git a/kbe/chain.py b/kbe/chain.py index ad481b65..f7dada96 100644 --- a/kbe/chain.py +++ b/kbe/chain.py @@ -4,22 +4,27 @@ # Old fragmentation code - do not use other than debugging! -def findH(mol,nh, tmphlist = []): + +def findH(mol, nh, tmphlist=[]): import math + coord_ = mol.atom_coords() hidx = [] for idx in range(mol.natm): - if mol.atom_pure_symbol(idx) == 'H': - d = math.sqrt((nh[0]-coord_[idx][0])**2 + - (nh[1]-coord_[idx][1])**2 + - (nh[2]-coord_[idx][2])**2) - if d*0.52917721092< 1.6: + if mol.atom_pure_symbol(idx) == "H": + d = math.sqrt( + (nh[0] - coord_[idx][0]) ** 2 + + (nh[1] - coord_[idx][1]) ** 2 + + (nh[2] - coord_[idx][2]) ** 2 + ) + if d * 0.52917721092 < 1.6: if not idx in tmphlist: hidx.append(idx) return hidx + def polychain(self, mol, frozen_core=False, unitcell=1): """ Hard coded fragmentation for polymer chains. This is not recommended for any production level calculations. @@ -30,7 +35,7 @@ def polychain(self, mol, frozen_core=False, unitcell=1): baslist = mol.aoslice_by_atom() mol2 = mol.copy() - mol2.basis = 'sto-3g' + mol2.basis = "sto-3g" mol2.build() bas2list = mol2.aoslice_by_atom() @@ -39,9 +44,8 @@ def polychain(self, mol, frozen_core=False, unitcell=1): coreshift = 0 tmphlist = [] for adx in range(mol.natm): - - if not mol.atom_pure_symbol(adx) == 'H': - h1 = findH(mol, mol.atom_coord(adx),tmphlist = tmphlist ) + if not mol.atom_pure_symbol(adx) == "H": + h1 = findH(mol, mol.atom_coord(adx), tmphlist=tmphlist) tmphlist.extend(h1) bas = baslist[adx] @@ -56,7 +60,6 @@ def polychain(self, mol, frozen_core=False, unitcell=1): sites_ = [] sites2_ = [] for hidx in h1: - basH = baslist[hidx] basH2 = bas2list[hidx] @@ -76,21 +79,21 @@ def polychain(self, mol, frozen_core=False, unitcell=1): b1list = [i for i in range(startH_, stopH_)] b2list = [i for i in range(startH2_, stopH2_)] - sites2_.extend([ True if i in b2list else False for i in b1list]) + sites2_.extend([True if i in b2list else False for i in b1list]) sites_.extend(b1list) if frozen_core: start_ -= coreshift start2_ -= coreshift ncore_ = self.core_list[adx] - stop_ -= coreshift+ncore_ - stop2_ -= coreshift+ncore_ + stop_ -= coreshift + ncore_ + stop2_ -= coreshift + ncore_ coreshift += ncore_ b1list = [i for i in range(start_, stop_)] b2list = [i for i in range(start2_, stop2_)] sites_.extend(b1list) - sites2_.extend([ True if i in b2list else False for i in b1list]) + sites2_.extend([True if i in b2list else False for i in b1list]) sites__.append(sites_) sites2.append(sites2_) @@ -103,29 +106,29 @@ def polychain(self, mol, frozen_core=False, unitcell=1): ns_ = nsites[-1][-1] nsites = [] for p in sites: - nsites.append([q+ns_+1 for q in p]) + nsites.append([q + ns_ + 1 for q in p]) else: nsites = sites__ sites = [*sites, *nsites] else: int_sites = int(unitcell) - frac_sites = int(len(sites__)*(unitcell-int_sites)) + frac_sites = int(len(sites__) * (unitcell - int_sites)) for i in range(int_sites): if i: ns_ = nsites[-1][-1] nsites = [] for p in sites: - nsites.append([q+ns_+1 for q in p]) + nsites.append([q + ns_ + 1 for q in p]) else: nsites = sites__ sites = [*sites, *nsites] ns_ = sites[-1][-1] nsites = [] for i in range(frac_sites): - nsites.append([q+ns_+1 for q in sites__[i]]) + nsites.append([q + ns_ + 1 for q in sites__[i]]) sites = [*sites, *nsites] - elif unitcell <1: - frac_sites =int(len(sites__)*(unitcell)) + elif unitcell < 1: + frac_sites = int(len(sites__) * (unitcell)) ns_ = sites__[-1][-1] nsites = [] for i in range(frac_sites): @@ -134,90 +137,91 @@ def polychain(self, mol, frozen_core=False, unitcell=1): else: sites = sites__ - if self.be_type=='be2': - fs=[] + if self.be_type == "be2": + fs = [] if not self.self_match and not self.allcen: - - for i in range(len(sites)-2): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]) - fs.append([ sites[i], sites[i+1], sites[i+2]]) + for i in range(len(sites) - 2): + self.fsites.append(sites[i] + sites[i + 1] + sites[i + 2]) + fs.append([sites[i], sites[i + 1], sites[i + 2]]) self.Nfrag = len(self.fsites) self.edge.append([fs[0][2]]) for i in fs[1:-1]: - self.edge.append([i[0],i[-1]]) + self.edge.append([i[0], i[-1]]) self.edge.append([fs[-1][0]]) self.center.append([1]) - for i in range(self.Nfrag-2): - self.center.append([i,i+2]) - self.center.append([self.Nfrag-2]) + for i in range(self.Nfrag - 2): + self.center.append([i, i + 2]) + self.center.append([self.Nfrag - 2]) for ix, i in enumerate(self.fsites): tmp_ = [] - elist_ = [ xx for yy in self.edge[ix] for xx in yy] + elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) elif self.allcen: - sites_left = [-i-1 for i in sites[0]] + sites_left = [-i - 1 for i in sites[0]] ns = len(sites[-1]) - sites_right = [i+ns for i in sites[-1]] + sites_right = [i + ns for i in sites[-1]] - self.fsites.append(sites_left+sites[0]+sites[1]) + self.fsites.append(sites_left + sites[0] + sites[1]) fs.append([sites_left, sites[0], sites[1]]) - for i in range(len(sites)-2): - self.fsites.append(sites[i]+sites[i+1]+sites[i+2]) - fs.append([ sites[i], sites[i+1], sites[i+2]]) - self.fsites.append(sites[-2]+sites[-1]+sites_right) - fs.append([sites[-2],sites[-1],sites_right]) + for i in range(len(sites) - 2): + self.fsites.append(sites[i] + sites[i + 1] + sites[i + 2]) + fs.append([sites[i], sites[i + 1], sites[i + 2]]) + self.fsites.append(sites[-2] + sites[-1] + sites_right) + fs.append([sites[-2], sites[-1], sites_right]) self.Nfrag = len(self.fsites) - for i in fs: #[1:-1]: - self.edge.append([i[0],i[-1]]) + for i in fs: # [1:-1]: + self.edge.append([i[0], i[-1]]) - self.center.append([self.Nfrag-1, 1]) - for i in range(self.Nfrag-2): - self.center.append([i,i+2]) - self.center.append([self.Nfrag-2, 0]) + self.center.append([self.Nfrag - 1, 1]) + for i in range(self.Nfrag - 2): + self.center.append([i, i + 2]) + self.center.append([self.Nfrag - 2, 0]) for ix, i in enumerate(self.fsites): tmp_ = [] elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) else: - - for i in range(len(sites)-2): - self.fsites.append(sites[i]+sites[i+1]+sites[i+2]) - fs.append([ sites[i], sites[i+1], sites[i+2]]) + for i in range(len(sites) - 2): + self.fsites.append(sites[i] + sites[i + 1] + sites[i + 2]) + fs.append([sites[i], sites[i + 1], sites[i + 2]]) self.Nfrag = len(self.fsites) - #self.edge.append([fs[0][2]]) - for i in fs:#[1:-1]: - self.edge.append([i[0],i[-1]]) - #self.edge.append([fs[-1][0]]) + # self.edge.append([fs[0][2]]) + for i in fs: # [1:-1]: + self.edge.append([i[0], i[-1]]) + # self.edge.append([fs[-1][0]]) # Frag (0) outer edge is in Frag (-1) - self.center.append([self.Nfrag-1,1]) - for i in range(self.Nfrag-2): - self.center.append([i,i+2]) + self.center.append([self.Nfrag - 1, 1]) + for i in range(self.Nfrag - 2): + self.center.append([i, i + 2]) # Frag (-1) outer edge is in Frag (0) - self.center.append([self.Nfrag-2, 0]) + self.center.append([self.Nfrag - 2, 0]) for ix, i in enumerate(self.fsites): tmp_ = [] if ix == 0: tmp_.extend([i.index(k) for k in self.edge[ix][0]]) - elist_ = [ xx for yy in self.edge[ix] for xx in yy] + elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) - if ix == self.Nfrag-1: + if not j in elist_: + tmp_.append(i.index(j)) + if ix == self.Nfrag - 1: tmp_.extend([i.index(k) for k in self.edge[ix][1]]) self.ebe_weight.append([1.0, tmp_]) @@ -226,120 +230,134 @@ def polychain(self, mol, frozen_core=False, unitcell=1): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][1]]) print(self.fsites) - elif self.be_type=='be3': - fs=[] + elif self.be_type == "be3": + fs = [] if not self.self_match and not self.allcen: # change back to i,i+1,i+2 - for i in range(len(sites)-4): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]+ - sites[i+3]+ sites[i+4]) - fs.append([ sites[i], sites[i+1], sites[i+2], sites[i+3], - sites[i+4]]) + for i in range(len(sites) - 4): + self.fsites.append( + sites[i] + sites[i + 1] + sites[i + 2] + sites[i + 3] + sites[i + 4] + ) + fs.append( + [sites[i], sites[i + 1], sites[i + 2], sites[i + 3], sites[i + 4]] + ) self.Nfrag = len(self.fsites) self.edge.append([fs[0][3], fs[0][4]]) # change back 0->1 for i in fs[1:-1]: - self.edge.append([i[0],i[1], i[-2], i[-1]]) + self.edge.append([i[0], i[1], i[-2], i[-1]]) self.edge.append([fs[-1][0], fs[-1][1]]) - self.center.append([1,2]) - self.center.append([0,0,2,3]) - for i in range(self.Nfrag-4): - self.center.append([i, i+1, i+3, i+4]) - self.center.append([self.Nfrag-4, self.Nfrag-3, - self.Nfrag-1, self.Nfrag-1]) - self.center.append([self.Nfrag-3, self.Nfrag-2]) + self.center.append([1, 2]) + self.center.append([0, 0, 2, 3]) + for i in range(self.Nfrag - 4): + self.center.append([i, i + 1, i + 3, i + 4]) + self.center.append( + [self.Nfrag - 4, self.Nfrag - 3, self.Nfrag - 1, self.Nfrag - 1] + ) + self.center.append([self.Nfrag - 3, self.Nfrag - 2]) for i in range(self.Nfrag): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][2]]) for ix, i in enumerate(self.fsites): tmp_ = [] - elist_ = [ xx for yy in self.edge[ix] for xx in yy] + elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) elif self.allcen: ns = len(sites[-1]) - sites_left1 = [-i-1 for i in sites[0]] - sites_left0 = [i-ns for i in sites_left1] - sites_right0 = [i+ns for i in sites[-1]] - sites_right1 = [i+ns for i in sites_right0] - - self.fsites.append(sites_left0+sites_left1+sites[0]+sites[1]+sites[2]) - self.fsites.append(sites_left1+sites[0]+sites[1]+sites[2]+sites[3]) - fs.append([sites_left0,sites_left1,sites[0],sites[1],sites[2]]) - fs.append([sites_left1,sites[0],sites[1],sites[2],sites[3]]) - for i in range(len(sites)-4): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]+ - sites[i+3]+ sites[i+4]) - fs.append([ sites[i], sites[i+1], sites[i+2], sites[i+3], - sites[i+4]]) - self.fsites.append(sites[-4]+sites[-3]+sites[-2]+sites[-1]+sites_right0) - self.fsites.append(sites[-3]+sites[-2]+sites[-1]+sites_right0+sites_right1) - fs.append([sites[-4],sites[-3],sites[-2],sites[-1],sites_right0]) - fs.append([sites[-3],sites[-2],sites[-1],sites_right0,sites_right1]) + sites_left1 = [-i - 1 for i in sites[0]] + sites_left0 = [i - ns for i in sites_left1] + sites_right0 = [i + ns for i in sites[-1]] + sites_right1 = [i + ns for i in sites_right0] + + self.fsites.append( + sites_left0 + sites_left1 + sites[0] + sites[1] + sites[2] + ) + self.fsites.append(sites_left1 + sites[0] + sites[1] + sites[2] + sites[3]) + fs.append([sites_left0, sites_left1, sites[0], sites[1], sites[2]]) + fs.append([sites_left1, sites[0], sites[1], sites[2], sites[3]]) + for i in range(len(sites) - 4): + self.fsites.append( + sites[i] + sites[i + 1] + sites[i + 2] + sites[i + 3] + sites[i + 4] + ) + fs.append( + [sites[i], sites[i + 1], sites[i + 2], sites[i + 3], sites[i + 4]] + ) + self.fsites.append( + sites[-4] + sites[-3] + sites[-2] + sites[-1] + sites_right0 + ) + self.fsites.append( + sites[-3] + sites[-2] + sites[-1] + sites_right0 + sites_right1 + ) + fs.append([sites[-4], sites[-3], sites[-2], sites[-1], sites_right0]) + fs.append([sites[-3], sites[-2], sites[-1], sites_right0, sites_right1]) self.Nfrag = len(self.fsites) - #self.edge.append([fs[0][3], fs[0][4]]) - #self.edge.append([fs[1][1], fs[1][3], fs[1][4]]) - for i in fs: #[2:-2]: - self.edge.append([i[0],i[1], i[-2], i[-1]]) - #self.edge.append([fs[-2][0],fs[-2][1],fs[-2][3]]) - #self.edge.append([fs[-1][0],fs[-1][1]]) + # self.edge.append([fs[0][3], fs[0][4]]) + # self.edge.append([fs[1][1], fs[1][3], fs[1][4]]) + for i in fs: # [2:-2]: + self.edge.append([i[0], i[1], i[-2], i[-1]]) + # self.edge.append([fs[-2][0],fs[-2][1],fs[-2][3]]) + # self.edge.append([fs[-1][0],fs[-1][1]]) - self.center.append([self.Nfrag-2,self.Nfrag-1, 1,2]) - self.center.append([self.Nfrag-1,0,2,3]) - for i in range(self.Nfrag-4): - self.center.append([i, i+1, i+3, i+4]) - self.center.append([self.Nfrag-4, self.Nfrag-3, self.Nfrag-1, 0]) - self.center.append([self.Nfrag-3, self.Nfrag-2, 0, 1]) + self.center.append([self.Nfrag - 2, self.Nfrag - 1, 1, 2]) + self.center.append([self.Nfrag - 1, 0, 2, 3]) + for i in range(self.Nfrag - 4): + self.center.append([i, i + 1, i + 3, i + 4]) + self.center.append([self.Nfrag - 4, self.Nfrag - 3, self.Nfrag - 1, 0]) + self.center.append([self.Nfrag - 3, self.Nfrag - 2, 0, 1]) for ix, i in enumerate(self.fsites): tmp_ = [] elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) else: - - for i in range(len(sites)-4): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]+ - sites[i+3]+ sites[i+4]) - fs.append([ sites[i], sites[i+1], sites[i+2], sites[i+3], - sites[i+4]]) + for i in range(len(sites) - 4): + self.fsites.append( + sites[i] + sites[i + 1] + sites[i + 2] + sites[i + 3] + sites[i + 4] + ) + fs.append( + [sites[i], sites[i + 1], sites[i + 2], sites[i + 3], sites[i + 4]] + ) Nfrag = len(self.fsites) self.Nfrag = Nfrag for i in fs: - self.edge.append([i[0],i[1], i[-2], i[-1]]) + self.edge.append([i[0], i[1], i[-2], i[-1]]) - self.center.append([Nfrag-2,Nfrag-1,1,Nfrag-2]) - self.center.append([Nfrag-1,0,Nfrag-2,Nfrag-1]) - for i in range(self.Nfrag-4): - self.center.append([i, i+1, i+3, i+4]) + self.center.append([Nfrag - 2, Nfrag - 1, 1, Nfrag - 2]) + self.center.append([Nfrag - 1, 0, Nfrag - 2, Nfrag - 1]) + for i in range(self.Nfrag - 4): + self.center.append([i, i + 1, i + 3, i + 4]) if Nfrag > 2: - self.center.append([Nfrag-4, Nfrag-3, Nfrag-1, 0]) - self.center.append([Nfrag-3, Nfrag-2, 0, 1]) + self.center.append([Nfrag - 4, Nfrag - 3, Nfrag - 1, 0]) + self.center.append([Nfrag - 3, Nfrag - 2, 0, 1]) for ix, i in enumerate(self.fsites): tmp_ = [] - if ix==0: + if ix == 0: for edg_ix in range(2): tmp_.extend([i.index(k) for k in self.edge[ix][edg_ix]]) - elist_ = [ xx for yy in self.edge[ix] for xx in yy] + elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) - if ix == self.Nfrag-1: + if not j in elist_: + tmp_.append(i.index(j)) + if ix == self.Nfrag - 1: for edg_ix in range(2): - - tmp_.extend([i.index(k) for k in self.edge[ix][edg_ix+2]]) + tmp_.extend([i.index(k) for k in self.edge[ix][edg_ix + 2]]) self.ebe_weight.append([1.0, tmp_]) # center on each fragments ?? do we add more for PBC @@ -347,10 +365,10 @@ def polychain(self, mol, frozen_core=False, unitcell=1): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][2]]) print(flush=True) - print(' No. of fragments : ',self.Nfrag,flush=True) + print(" No. of fragments : ", self.Nfrag, flush=True) print(flush=True) - if not self.be_type=='be1': + if not self.be_type == "be1": for i in range(self.Nfrag): idx = [] for j in self.edge[i]: @@ -360,9 +378,12 @@ def polychain(self, mol, frozen_core=False, unitcell=1): for i in range(self.Nfrag): idx = [] for j in range(len(self.center[i])): - - idx.append([self.fsites[self.center[i][j]].index(k) - for k in self.edge[i][j]]) + idx.append( + [ + self.fsites[self.center[i][j]].index(k) + for k in self.edge[i][j] + ] + ) self.center_idx.append(idx) else: for i in range(self.Nfrag): @@ -375,7 +396,3 @@ def polychain(self, mol, frozen_core=False, unitcell=1): idx__.append(fidx) idx.append(idx__) self.center_idx.append(idx) - - - - diff --git a/kbe/fragment.py b/kbe/fragment.py index 8ab72eba..f89abdb3 100644 --- a/kbe/fragment.py +++ b/kbe/fragment.py @@ -4,21 +4,38 @@ from molbe.helper import get_core from .autofrag import autogen + def print_mol_missing(): - print('Provide pyscf gto.Cell object in fragpart() and restart!', - flush=True) - print('exiting',flush=True) + print("Provide pyscf gto.Cell object in fragpart() and restart!", flush=True) + print("exiting", flush=True) sys.exit() -class fragpart: - def __init__(self, natom=0, dim=1, frag_type='autogen', - unitcell=1, - gamma_2d = False, gamma_1d=False, interlayer=False, - long_bond=False, perpend_dist = 4.0, perpend_dist_tol = 1e-3, - nx=False, ny=False, nz=False,closed=False, - kpt = None, valence_basis=None, - be_type='be2', mol=None, frozen_core=False, self_match=False, allcen=True): +class fragpart: + def __init__( + self, + natom=0, + dim=1, + frag_type="autogen", + unitcell=1, + gamma_2d=False, + gamma_1d=False, + interlayer=False, + long_bond=False, + perpend_dist=4.0, + perpend_dist_tol=1e-3, + nx=False, + ny=False, + nz=False, + closed=False, + kpt=None, + valence_basis=None, + be_type="be2", + mol=None, + frozen_core=False, + self_match=False, + allcen=True, + ): """Fragment/partitioning definition Interfaces two main fragmentation functions (autogen & polychain) in MolBE. It defines edge & @@ -60,10 +77,10 @@ def __init__(self, natom=0, dim=1, frag_type='autogen', # No. of unitcells to use for fragment construction self.unitcell = unitcell self.mol = mol - self.frag_type=frag_type + self.frag_type = frag_type self.fsites = [] self.Nfrag = 0 - self.edge= [] + self.edge = [] self.center = [] self.ebe_weight = [] self.edge_idx = [] @@ -73,40 +90,59 @@ def __init__(self, natom=0, dim=1, frag_type='autogen', self.natom = natom self.frozen_core = frozen_core self.self_match = self_match - self.allcen=allcen - self.valence_basis=valence_basis + self.allcen = allcen + self.valence_basis = valence_basis self.kpt = kpt - self.molecule = False ### remove this + self.molecule = False ### remove this # Check for frozen core approximation if frozen_core: self.ncore, self.no_core_idx, self.core_list = get_core(mol) - if frag_type=='polychain': - if mol is None: print_mol_missing() + if frag_type == "polychain": + if mol is None: + print_mol_missing() self.polychain(mol, frozen_core=frozen_core, unitcell=unitcell) - elif frag_type=='autogen': - if mol is None: print_mol_missing() + elif frag_type == "autogen": + if mol is None: + print_mol_missing() if kpt is None: - print('Provide kpt mesh in fragpart() and restart!', - flush=True) - print('exiting',flush=True) + print("Provide kpt mesh in fragpart() and restart!", flush=True) + print("exiting", flush=True) sys.exit() - fgs = autogen(mol, kpt, be_type=be_type, frozen_core=frozen_core, - valence_basis=valence_basis, unitcell=unitcell, - nx=nx, ny=ny, nz=nz, - long_bond=long_bond, - perpend_dist = perpend_dist, perpend_dist_tol = perpend_dist_tol, - gamma_2d=gamma_2d, gamma_1d=gamma_1d,interlayer=interlayer) + fgs = autogen( + mol, + kpt, + be_type=be_type, + frozen_core=frozen_core, + valence_basis=valence_basis, + unitcell=unitcell, + nx=nx, + ny=ny, + nz=nz, + long_bond=long_bond, + perpend_dist=perpend_dist, + perpend_dist_tol=perpend_dist_tol, + gamma_2d=gamma_2d, + gamma_1d=gamma_1d, + interlayer=interlayer, + ) - self.fsites, self.edge, self.center, self.edge_idx, self.center_idx, self.centerf_idx, self.ebe_weight = fgs + ( + self.fsites, + self.edge, + self.center, + self.edge_idx, + self.center_idx, + self.centerf_idx, + self.ebe_weight, + ) = fgs self.Nfrag = len(self.fsites) else: - print('Fragmentation type = ',frag_type,' not implemented!', - flush=True) - print('exiting',flush=True) + print("Fragmentation type = ", frag_type, " not implemented!", flush=True) + print("exiting", flush=True) sys.exit() from .chain import polychain diff --git a/kbe/helper.py b/kbe/helper.py index 6a2a88a9..4d540cb3 100644 --- a/kbe/helper.py +++ b/kbe/helper.py @@ -2,9 +2,11 @@ import numpy, sys, functools + def get_veff(eri_, dm, S, TA, hf_veff, return_veff0=False): import functools from pyscf import scf + """ Calculate the effective HF potential (Veff) for a given density matrix and electron repulsion integrals. @@ -31,27 +33,25 @@ def get_veff(eri_, dm, S, TA, hf_veff, return_veff0=False): P_ = numpy.zeros((neo, neo), dtype=numpy.complex128) for k in range(nk): Cinv = numpy.dot(TA[k].conj().T, S[k]) - P_ += functools.reduce(numpy.dot, - (Cinv, dm[k], Cinv.conj().T)) + P_ += functools.reduce(numpy.dot, (Cinv, dm[k], Cinv.conj().T)) P_ /= float(nk) P_ = numpy.asarray(P_.real, dtype=numpy.double) eri_ = numpy.asarray(eri_, dtype=numpy.double) vj, vk = scf.hf.dot_eri_dm(eri_, P_, hermi=1, with_j=True, with_k=True) - Veff_ = vj - 0.5*vk + Veff_ = vj - 0.5 * vk # remove core contribution from hf_veff Veff0 = numpy.zeros((neo, neo), dtype=numpy.complex128) for k in range(nk): - Veff0 += functools.reduce(numpy.dot, - (TA[k].conj().T, hf_veff[k], TA[k])) + Veff0 += functools.reduce(numpy.dot, (TA[k].conj().T, hf_veff[k], TA[k])) Veff0 /= float(nk) Veff = Veff0 - Veff_ if return_veff0: - return(Veff0, Veff) + return (Veff0, Veff) return Veff diff --git a/kbe/lo.py b/kbe/lo.py index f51b7c3b..5754933b 100644 --- a/kbe/lo.py +++ b/kbe/lo.py @@ -2,26 +2,33 @@ # Henry Tran # from pyscf import lib -import numpy,sys +import numpy, sys from copy import deepcopy from functools import reduce from molbe.external.lo_helper import get_aoind_by_atom, reorder_by_atom_ # iao tmp - class KMF: - def __init__(self, cell, kpts = None, mo_coeff = None, mo_energy = None): - self. cell = cell + def __init__(self, cell, kpts=None, mo_coeff=None, mo_energy=None): + self.cell = cell self.kpts = kpts self.mo_coeff = mo_coeff.copy() self.mo_energy = mo_energy self.mo_energy_kpts = mo_energy self.mo_coeff_kpts = mo_coeff.copy() -def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True, - valence_only=False, iao_val_core=True): - """ Orbital localization + +def localize( + self, + lo_method, + mol=None, + valence_basis="sto-3g", + iao_wannier=True, + valence_only=False, + iao_val_core=True, +): + """Orbital localization Performs orbital localization computations for periodic systems. For large basis, IAO is recommended augmented with PAO orbitals. @@ -42,72 +49,72 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True """ from numpy.linalg import eigh from pyscf.lo.iao import iao - import scipy.linalg,functools + import scipy.linalg, functools from molbe.helper import ncore_ from .misc import get_phase, sgeom - if lo_method == 'iao': - if valence_basis == 'sto-3g': + if lo_method == "iao": + if valence_basis == "sto-3g": from .basis_sto3g_core_val import val_basis, core_basis - elif valence_basis == 'minao': + elif valence_basis == "minao": from .basis_minao_core_val import val_basis, core_basis elif iao_val_core: - sys.exit('valence_basis='+valence_basis+' not supported for iao_val_core=True') - - - if lo_method == 'lowdin': + sys.exit( + "valence_basis=" + + valence_basis + + " not supported for iao_val_core=True" + ) + if lo_method == "lowdin": # Lowdin orthogonalization with k-points W = numpy.zeros_like(self.S) nk, nao, nmo = self.C.shape if self.frozen_core: - W_nocore = numpy.zeros_like(self.S[:,:, self.ncore:]) - lmo_coeff = numpy.zeros_like(self.C[:,self.ncore:,self.ncore:]) - cinv_ = numpy.zeros((nk, nmo-self.ncore, nao), dtype=numpy.complex128) + W_nocore = numpy.zeros_like(self.S[:, :, self.ncore :]) + lmo_coeff = numpy.zeros_like(self.C[:, self.ncore :, self.ncore :]) + cinv_ = numpy.zeros((nk, nmo - self.ncore, nao), dtype=numpy.complex128) else: lmo_coeff = numpy.zeros_like(self.C) cinv_ = numpy.zeros((nk, nmo, nao), dtype=numpy.complex128) for k in range(self.nkpt): - es_, vs_ = scipy.linalg.eigh(self.S[k]) - edx = es_ > 1.e-14 + edx = es_ > 1.0e-14 - W[k] = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].conj().T) + W[k] = numpy.dot(vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].conj().T) for i in range(W[k].shape[1]): - if W[k][i,i] < 0: - W[:,i] *= -1 + if W[k][i, i] < 0: + W[:, i] *= -1 if self.frozen_core: - pcore = numpy.eye(W[k].shape[0]) - numpy.dot(self.P_core[k], self.S[k]) C_ = numpy.dot(pcore, W[k]) # PYSCF has basis in 1s2s3s2p2p2p3p3p3p format # fix no_core_idx - use population for now - #C_ = C_[:,self.no_core_idx] - Cpop = functools.reduce(numpy.dot, - (C_.conj().T, self.S[k], C_)) + # C_ = C_[:,self.no_core_idx] + Cpop = functools.reduce(numpy.dot, (C_.conj().T, self.S[k], C_)) Cpop = numpy.diag(Cpop.real) no_core_idx = numpy.where(Cpop > 0.7)[0] - C_ = C_[:,no_core_idx] + C_ = C_[:, no_core_idx] - S_ = functools.reduce(numpy.dot, - (C_.conj().T, self.S[k], C_)) + S_ = functools.reduce(numpy.dot, (C_.conj().T, self.S[k], C_)) es_, vs_ = scipy.linalg.eigh(S_) - edx = es_ > 1.e-14 - W_ = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].conj().T) - W_nocore[k] = numpy.dot(C_,W_) - - lmo_coeff[k] = functools.reduce(numpy.dot, - (W_nocore[k].conj().T,self.S[k], - self.C[k][:,self.ncore:])) + edx = es_ > 1.0e-14 + W_ = numpy.dot(vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].conj().T) + W_nocore[k] = numpy.dot(C_, W_) + + lmo_coeff[k] = functools.reduce( + numpy.dot, + (W_nocore[k].conj().T, self.S[k], self.C[k][:, self.ncore :]), + ) cinv_[k] = numpy.dot(W_nocore[k].conj().T, self.S[k]) else: - lmo_coeff[k] = functools.reduce(numpy.dot, - (W[k].conj().T,self.S[k], self.C[k])) + lmo_coeff[k] = functools.reduce( + numpy.dot, (W[k].conj().T, self.S[k], self.C[k]) + ) cinv_[k] = numpy.dot(W[k].conj().T, self.S[k]) if self.frozen_core: self.W = W_nocore @@ -116,34 +123,42 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True self.lmo_coeff = lmo_coeff self.cinv = cinv_ - elif lo_method=='iao': + elif lo_method == "iao": from libdmet.lo import pywannier90 from pyscf import lo import os, h5py - from .lo_k import get_xovlp_k, get_iao_k, get_pao_k, get_pao_native_k,symm_orth_k , remove_core_mo_k + from .lo_k import ( + get_xovlp_k, + get_iao_k, + get_pao_k, + get_pao_native_k, + symm_orth_k, + remove_core_mo_k, + ) from pyscf import gto, lo if not iao_val_core or not self.frozen_core: - Co = self.C[:,:,:self.Nocc].copy() + Co = self.C[:, :, : self.Nocc].copy() S12, S2 = get_xovlp_k(self.cell, self.kpts, basis=valence_basis) ciao_ = get_iao_k(Co, S12, self.S, S2=S2) arrange_by_atom = True # tmp - aos are not rearrange and so below is not necessary if arrange_by_atom: - nk, nao, nlo = ciao_.shape Ciao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) for k in range(self.nkpt): aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, iaoind_by_atom = reorder_by_atom_(ciao_[k], aoind_by_atom, self.S[k]) + ctmp, iaoind_by_atom = reorder_by_atom_( + ciao_[k], aoind_by_atom, self.S[k] + ) Ciao_[k] = ctmp else: Ciao_ = ciao_.copy() # get_pao_k returns canonical orthogonalized orbitals - #Cpao = get_pao_k(Ciao, self.S, S12, S2, self.cell) + # Cpao = get_pao_k(Ciao, self.S, S12, S2, self.cell) # get_pao_native_k returns symm orthogonalized orbitals cpao_ = get_pao_native_k(Ciao_, self.S, self.cell, valence_basis, self.kpts) @@ -152,7 +167,9 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True Cpao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) for k in range(self.nkpt): aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, paoind_by_atom = reorder_by_atom_(cpao_[k], aoind_by_atom, self.S[k]) + ctmp, paoind_by_atom = reorder_by_atom_( + cpao_[k], aoind_by_atom, self.S[k] + ) Cpao_[k] = ctmp else: Cpao_ = cpao_.copy() @@ -160,9 +177,11 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True nk, nao, nlo = Ciao_.shape if self.frozen_core: nk, nao, nlo = Ciao_.shape - Ciao_nocore = numpy.zeros((nk, nao, nlo-self.ncore), dtype=numpy.complex128) + Ciao_nocore = numpy.zeros( + (nk, nao, nlo - self.ncore), dtype=numpy.complex128 + ) for k in range(nk): - Ccore = self.C[k][:,:self.ncore] + Ccore = self.C[k][:, : self.ncore] Ciao_nocore[k] = remove_core_mo_k(Ciao_[k], Ccore, self.S[k]) Ciao_ = Ciao_nocore @@ -171,29 +190,36 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True # Begin core s12_core_, s2_core = get_xovlp_k(self.cell, self.kpts, basis=core_basis) - C_core_ = self.C[:,:,:self.ncore].copy() + C_core_ = self.C[:, :, : self.ncore].copy() nk_, nao_, nmo_ = C_core_.shape s1_core = numpy.zeros((nk_, nmo_, nmo_), dtype=self.S.dtype) - s12_core = numpy.zeros((nk_, nmo_, s12_core_.shape[-1]), dtype=s12_core_.dtype) + s12_core = numpy.zeros( + (nk_, nmo_, s12_core_.shape[-1]), dtype=s12_core_.dtype + ) C_core = numpy.zeros((nk_, self.ncore, self.ncore), dtype=C_core_.dtype) for k in range(nk_): C_core[k] = C_core_[k].conj().T @ self.S[k] @ C_core_[k] s1_core[k] = C_core_[k].conj().T @ self.S[k] @ C_core_[k] s12_core[k] = C_core_[k].conj().T @ s12_core_[k] ciao_core_ = get_iao_k(C_core, s12_core, s1_core, s2_core, ortho=False) - ciao_core = numpy.zeros((nk_, nao_, ciao_core_.shape[-1]), dtype=ciao_core_.dtype) + ciao_core = numpy.zeros( + (nk_, nao_, ciao_core_.shape[-1]), dtype=ciao_core_.dtype + ) for k in range(nk_): ciao_core[k] = C_core_[k] @ ciao_core_[k] - ciao_core[k] = symm_orth_k(ciao_core[k], ovlp = self.S[k]) + ciao_core[k] = symm_orth_k(ciao_core[k], ovlp=self.S[k]) # Begin valence s12_val_, s2_val = get_xovlp_k(self.cell, self.kpts, basis=val_basis) - C_nocore = self.C[:,:,self.ncore:].copy() - C_nocore_occ_ = C_nocore[:,:,:self.Nocc].copy() + C_nocore = self.C[:, :, self.ncore :].copy() + C_nocore_occ_ = C_nocore[:, :, : self.Nocc].copy() nk_, nao_, nmo_ = C_nocore.shape s1_val = numpy.zeros((nk_, nmo_, nmo_), dtype=self.S.dtype) s12_val = numpy.zeros((nk_, nmo_, s12_val_.shape[-1]), dtype=s12_val_.dtype) - C_nocore_occ = numpy.zeros((nk_, nao_-self.ncore, C_nocore_occ_.shape[-1]),dtype=C_nocore_occ_.dtype) + C_nocore_occ = numpy.zeros( + (nk_, nao_ - self.ncore, C_nocore_occ_.shape[-1]), + dtype=C_nocore_occ_.dtype, + ) for k in range(nk_): C_nocore_occ[k] = C_nocore[k].conj().T @ self.S[k] @ C_nocore_occ_[k] s1_val[k] = C_nocore[k].conj().T @ self.S[k] @ C_nocore[k] @@ -202,11 +228,13 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True Ciao_ = numpy.zeros((nk_, nao_, ciao_val_.shape[-1]), dtype=ciao_val_.dtype) for k in range(nk_): Ciao_[k] = C_nocore[k] @ ciao_val_[k] - Ciao_[k] = symm_orth_k(Ciao_[k], ovlp = self.S[k]) + Ciao_[k] = symm_orth_k(Ciao_[k], ovlp=self.S[k]) # stack core|val nao = self.S.shape[-1] - c_core_val = numpy.zeros((nk_, nao, Ciao_.shape[-1]+self.ncore), dtype=Ciao_.dtype) + c_core_val = numpy.zeros( + (nk_, nao, Ciao_.shape[-1] + self.ncore), dtype=Ciao_.dtype + ) for k in range(nk_): c_core_val[k] = numpy.hstack((ciao_core[k], Ciao_[k])) @@ -216,15 +244,21 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True nk, nao, nlo = c_core_val.shape for k in range(self.nkpt): aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, iaoind_by_atom = reorder_by_atom_(c_core_val[k], aoind_by_atom, self.S[k]) + ctmp, iaoind_by_atom = reorder_by_atom_( + c_core_val[k], aoind_by_atom, self.S[k] + ) - cpao_ = get_pao_native_k(c_core_val, self.S, self.cell, valence_basis, self.kpts, ortho=True) + cpao_ = get_pao_native_k( + c_core_val, self.S, self.cell, valence_basis, self.kpts, ortho=True + ) if arrange_by_atom: nk, nao, nlo = cpao_.shape Cpao_ = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) for k in range(self.nkpt): aoind_by_atom = get_aoind_by_atom(self.cell) - ctmp, paoind_by_atom = reorder_by_atom_(cpao_[k], aoind_by_atom, self.S[k]) + ctmp, paoind_by_atom = reorder_by_atom_( + cpao_[k], aoind_by_atom, self.S[k] + ) Cpao_[k] = ctmp Cpao = Cpao_.copy() @@ -233,14 +267,16 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True if iao_wannier: mo_energy_ = [] for k in range(nk): - fock_iao = reduce(numpy.dot, (Ciao_[k].conj().T, self.FOCK[k], Ciao_[k])) + fock_iao = reduce( + numpy.dot, (Ciao_[k].conj().T, self.FOCK[k], Ciao_[k]) + ) S_iao = reduce(numpy.dot, (Ciao_[k].conj().T, self.S[k], Ciao_[k])) e_iao, v_iao = scipy.linalg.eigh(fock_iao, S_iao) mo_energy_.append(e_iao) - iaomf = KMF(self.mol, kpts = self.kpts, mo_coeff = Ciao_, mo_energy = mo_energy_) + iaomf = KMF(self.mol, kpts=self.kpts, mo_coeff=Ciao_, mo_energy=mo_energy_) num_wann = numpy.asarray(iaomf.mo_coeff).shape[2] - keywords = ''' + keywords = """ num_iter = 5000 dis_num_iter = 0 conv_noise_amp = -2.0 @@ -248,13 +284,14 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True conv_tol = 1.0E-09 iprint = 3 kmesh_tol = 0.00001 - ''' + """ # set conv window # dis_num_iter=0 - w90 = pywannier90.W90(iaomf, self.kmesh, - num_wann, other_keywords=keywords) + w90 = pywannier90.W90(iaomf, self.kmesh, num_wann, other_keywords=keywords) - A_matrix = numpy.zeros((self.nkpt, num_wann, num_wann), dtype=numpy.complex128) + A_matrix = numpy.zeros( + (self.nkpt, num_wann, num_wann), dtype=numpy.complex128 + ) i_init = True for k in range(self.nkpt): @@ -263,14 +300,16 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True else: ovlp_ciao = uciao[k].conj().T @ self.S[k] @ Ciao[k] A_matrix[k] = ovlp_ciao - A_matrix = A_matrix.transpose(1,2,0) + A_matrix = A_matrix.transpose(1, 2, 0) w90.kernel(A_matrix=A_matrix) - u_mat = numpy.array(w90.U_matrix.transpose(2,0,1), order='C', dtype=numpy.complex128) + u_mat = numpy.array( + w90.U_matrix.transpose(2, 0, 1), order="C", dtype=numpy.complex128 + ) - os.system('cp wannier90.wout wannier90_iao.wout') - os.system('rm wannier90.*') + os.system("cp wannier90.wout wannier90_iao.wout") + os.system("rm wannier90.*") nk, nao, nlo = Ciao_.shape Ciao = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) @@ -279,8 +318,10 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True Ciao[k] = numpy.dot(Ciao_[k], u_mat[k]) # Stack Ciao|Cpao - Wstack = numpy.zeros((self.nkpt, Ciao.shape[1], Ciao.shape[2]+Cpao.shape[2]), - dtype=numpy.complex128) + Wstack = numpy.zeros( + (self.nkpt, Ciao.shape[1], Ciao.shape[2] + Cpao.shape[2]), + dtype=numpy.complex128, + ) if self.frozen_core: for k in range(self.nkpt): shift = 0 @@ -290,21 +331,21 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True ncore += nc niao = len(iaoind_by_atom[ix]) iaoind_ix = [i_ - ncore for i_ in iaoind_by_atom[ix][nc:]] - Wstack[k][:, shift:shift+niao-nc] = Ciao[k][:, iaoind_ix] - shift += niao-nc + Wstack[k][:, shift : shift + niao - nc] = Ciao[k][:, iaoind_ix] + shift += niao - nc npao = len(paoind_by_atom[ix]) - Wstack[k][:,shift:shift+npao] = Cpao[k][:, paoind_by_atom[ix]] + Wstack[k][:, shift : shift + npao] = Cpao[k][:, paoind_by_atom[ix]] shift += npao else: for k in range(self.nkpt): shift = 0 for ix in range(self.cell.natm): niao = len(iaoind_by_atom[ix]) - Wstack[k][:, shift:shift+niao] = Ciao[k][:, iaoind_by_atom[ix]] + Wstack[k][:, shift : shift + niao] = Ciao[k][:, iaoind_by_atom[ix]] shift += niao npao = len(paoind_by_atom[ix]) - Wstack[k][:,shift:shift+npao] = Cpao[k][:, paoind_by_atom[ix]] + Wstack[k][:, shift : shift + npao] = Cpao[k][:, paoind_by_atom[ix]] shift += npao self.W = Wstack @@ -316,81 +357,101 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True cinv_ = numpy.zeros((self.nkpt, nlo, nao), dtype=numpy.complex128) if nmo > nlo: - Co_nocore = self.C[:,:,self.ncore:self.Nocc] - Cv = self.C[:,:,self.Nocc:] + Co_nocore = self.C[:, :, self.ncore : self.Nocc] + Cv = self.C[:, :, self.Nocc :] # Ensure that the LOs span the occupied space for k in range(self.nkpt): - assert(numpy.allclose(numpy.sum((self.W[k].conj().T @ - self.S[k] @ Co_nocore[k])**2.), - self.Nocc - self.ncore)) + assert numpy.allclose( + numpy.sum((self.W[k].conj().T @ self.S[k] @ Co_nocore[k]) ** 2.0), + self.Nocc - self.ncore, + ) # Find virtual orbitals that lie in the span of LOs - u, l, vt = numpy.linalg.svd(self.W[k].conj().T @ self.S[k] @ Cv[k], - full_matrices=False) + u, l, vt = numpy.linalg.svd( + self.W[k].conj().T @ self.S[k] @ Cv[k], full_matrices=False + ) nvlo = nlo - self.Nocc - self.ncore - assert(numpy.allclose(numpy.sum(l[:nvlo]), nvlo)) + assert numpy.allclose(numpy.sum(l[:nvlo]), nvlo) C_ = numpy.hstack([Co_nocore[k], Cv[k] @ vt[:nvlo].conj().T]) lmo_ = self.W[k].conj().T @ self.S[k] @ C_ - assert(numpy.allclose(lmo_.conj().T @ lmo_, numpy.eye(lmo_.shape[1]))) + assert numpy.allclose(lmo_.conj().T @ lmo_, numpy.eye(lmo_.shape[1])) lmo_coeff.append(lmo_) else: for k in range(self.nkpt): - lmo_coeff[k] = reduce(numpy.dot, (self.W[k].conj().T, self.S[k], self.C[k][:,self.ncore:])) + lmo_coeff[k] = reduce( + numpy.dot, + (self.W[k].conj().T, self.S[k], self.C[k][:, self.ncore :]), + ) cinv_[k] = numpy.dot(self.W[k].conj().T, self.S[k]) - assert(numpy.allclose(lmo_coeff[k].conj().T @ lmo_coeff[k], numpy.eye(lmo_coeff[k].shape[1]))) + assert numpy.allclose( + lmo_coeff[k].conj().T @ lmo_coeff[k], + numpy.eye(lmo_coeff[k].shape[1]), + ) self.lmo_coeff = lmo_coeff self.cinv = cinv_ - elif lo_method == 'wannier': - #from pyscf.pbc.tools import pywannier90 + elif lo_method == "wannier": + # from pyscf.pbc.tools import pywannier90 from libdmet.lo import pywannier90 from .lo_k import get_symm_orth_mat_k, remove_core_mo_k + nk, nao, nmo = self.C.shape lorb = numpy.zeros((nk, nao, nmo), dtype=numpy.complex128) - lorb_nocore = numpy.zeros((nk, nao, nmo-self.ncore), dtype=numpy.complex128) + lorb_nocore = numpy.zeros((nk, nao, nmo - self.ncore), dtype=numpy.complex128) for k in range(nk): es_, vs_ = scipy.linalg.eigh(self.S[k]) - edx = es_ > 1.e-14 - lorb[k] = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].conj().T) + edx = es_ > 1.0e-14 + lorb[k] = numpy.dot( + vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].conj().T + ) if self.frozen_core: - Ccore = self.C[k][:,:self.ncore] + Ccore = self.C[k][:, : self.ncore] lorb_nocore[k] = remove_core_mo_k(lorb[k], Ccore, self.S[k]) if not self.frozen_core: - lmf = KMF(self.mol, kpts=self.kpts, mo_coeff = lorb, mo_energy = self.mo_energy) + lmf = KMF(self.mol, kpts=self.kpts, mo_coeff=lorb, mo_energy=self.mo_energy) else: mo_energy_nc = [] for k in range(nk): - fock_lnc = reduce(numpy.dot, (lorb_nocore[k].conj().T, self.FOCK[k], lorb_nocore[k])) - S_lnc = reduce(numpy.dot, (lorb_nocore[k].conj().T, self.S[k], lorb_nocore[k])) + fock_lnc = reduce( + numpy.dot, (lorb_nocore[k].conj().T, self.FOCK[k], lorb_nocore[k]) + ) + S_lnc = reduce( + numpy.dot, (lorb_nocore[k].conj().T, self.S[k], lorb_nocore[k]) + ) e__, v__ = scipy.linalg.eigh(fock_lnc, S_lnc) mo_energy_nc.append(e__) - lmf = KMF(self.mol, kpts=self.kpts, mo_coeff = lorb_nocore, mo_energy = mo_energy_nc) + lmf = KMF( + self.mol, kpts=self.kpts, mo_coeff=lorb_nocore, mo_energy=mo_energy_nc + ) num_wann = lmf.mo_coeff.shape[2] - keywords = ''' + keywords = """ num_iter = 10000 dis_num_iter = 0 conv_window = 10 conv_tol = 1.0E-09 iprint = 3 kmesh_tol = 0.00001 - ''' + """ - w90 = pywannier90.W90(lmf, self.kmesh, - num_wann, other_keywords=keywords) + w90 = pywannier90.W90(lmf, self.kmesh, num_wann, other_keywords=keywords) A_matrix = numpy.zeros((self.nkpt, num_wann, num_wann), dtype=numpy.complex128) - i_init = True # Using A=I + lowdin orbital and A= + |psi> is the same + i_init = ( + True # Using A=I + lowdin orbital and A= + |psi> is the same + ) for k in range(self.nkpt): if i_init: A_matrix[k] = numpy.eye(num_wann, dtype=numpy.complex128) - A_matrix = A_matrix.transpose(1,2,0) + A_matrix = A_matrix.transpose(1, 2, 0) w90.kernel(A_matrix=A_matrix) - u_mat = numpy.array(w90.U_matrix.transpose(2,0,1), order='C', dtype=numpy.complex128) + u_mat = numpy.array( + w90.U_matrix.transpose(2, 0, 1), order="C", dtype=numpy.complex128 + ) nk, nao, nlo = lmf.mo_coeff.shape W = numpy.zeros((nk, nao, nlo), dtype=numpy.complex128) @@ -398,17 +459,23 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', iao_wannier=True W[k] = numpy.dot(lmf.mo_coeff[k], u_mat[k]) self.W = W - lmo_coeff = numpy.zeros((self.nkpt, nlo, nmo-self.ncore), dtype=numpy.complex128) + lmo_coeff = numpy.zeros( + (self.nkpt, nlo, nmo - self.ncore), dtype=numpy.complex128 + ) cinv_ = numpy.zeros((self.nkpt, nlo, nao), dtype=numpy.complex128) for k in range(nk): - lmo_coeff[k] = reduce(numpy.dot, (self.W[k].conj().T, self.S[k], self.C[k][:,self.ncore:])) + lmo_coeff[k] = reduce( + numpy.dot, (self.W[k].conj().T, self.S[k], self.C[k][:, self.ncore :]) + ) cinv_[k] = numpy.dot(self.W[k].conj().T, self.S[k]) - assert(numpy.allclose(lmo_coeff[k].conj().T @ lmo_coeff[k], numpy.eye(lmo_coeff[k].shape[1]))) + assert numpy.allclose( + lmo_coeff[k].conj().T @ lmo_coeff[k], numpy.eye(lmo_coeff[k].shape[1]) + ) self.lmo_coeff = lmo_coeff self.cinv = cinv_ else: - print('lo_method = ',lo_method,' not implemented!',flush=True) - print('exiting',flush=True) + print("lo_method = ", lo_method, " not implemented!", flush=True) + print("exiting", flush=True) sys.exit() diff --git a/kbe/lo_k.py b/kbe/lo_k.py index 265f2a15..4481e85c 100644 --- a/kbe/lo_k.py +++ b/kbe/lo_k.py @@ -5,6 +5,7 @@ import numpy, sys, scipy from functools import reduce + def dot_gen(A, B, ovlp): if ovlp is None: Ad = numpy.dot(A.conj().T, B) @@ -12,43 +13,48 @@ def dot_gen(A, B, ovlp): Ad = reduce(numpy.dot, (A.conj().T, ovlp, B)) return Ad -def get_cano_orth_mat(A, thr=1.E-7, ovlp=None): - S = dot_gen(A,A,ovlp) + +def get_cano_orth_mat(A, thr=1.0e-7, ovlp=None): + S = dot_gen(A, A, ovlp) e, u = numpy.linalg.eigh(S) if thr > 0: - idx_keep = e/e[-1] > thr + idx_keep = e / e[-1] > thr else: idx_keep = list(range(e.shape[0])) - U = u[:,idx_keep] * e[idx_keep]**-0.5 + U = u[:, idx_keep] * e[idx_keep] ** -0.5 return U -def cano_orth(A, thr=1.E-7, ovlp=None): - """ Canonically orthogonalize columns of A - """ + +def cano_orth(A, thr=1.0e-7, ovlp=None): + """Canonically orthogonalize columns of A""" U = get_cano_orth_mat(A, thr, ovlp) return A @ U -def get_symm_orth_mat_k(A, thr=1.E-7, ovlp=None): - S = dot_gen(A,A,ovlp) + +def get_symm_orth_mat_k(A, thr=1.0e-7, ovlp=None): + S = dot_gen(A, A, ovlp) e, u = scipy.linalg.eigh(S) if int(numpy.sum(e < thr)) > 0: - raise ValueError("Linear dependence is detected in the column space of A: smallest eigenvalue (%.3E) is less than thr (%.3E). Please use 'cano_orth' instead." % (numpy.min(e), thr)) + raise ValueError( + "Linear dependence is detected in the column space of A: smallest eigenvalue (%.3E) is less than thr (%.3E). Please use 'cano_orth' instead." + % (numpy.min(e), thr) + ) U = reduce(numpy.dot, (u, numpy.diag(e**-0.5), u.conj().T)) - #U = reduce(numpy.dot, (u/numpy.sqrt(e), u.conj().T)) + # U = reduce(numpy.dot, (u/numpy.sqrt(e), u.conj().T)) return U -def symm_orth_k(A, thr=1.E-7, ovlp=None): - """ Symmetrically orthogonalize columns of A - """ + +def symm_orth_k(A, thr=1.0e-7, ovlp=None): + """Symmetrically orthogonalize columns of A""" U = get_symm_orth_mat_k(A, thr, ovlp) AU = numpy.dot(A, U) return AU -def get_xovlp_k(cell, kpts, basis='sto-3g'): +def get_xovlp_k(cell, kpts, basis="sto-3g"): """ Gets set of valence orbitals based on smaller (should be minimal) basis inumpy.t: @@ -59,34 +65,39 @@ def get_xovlp_k(cell, kpts, basis='sto-3g'): S22 - Overlap in new basis set """ - from pyscf.pbc import gto as pgto#intor_cross + from pyscf.pbc import gto as pgto # intor_cross cell_alt = cell.copy() cell_alt.basis = basis cell_alt.build() - S22 = numpy.array(cell_alt.pbc_intor('int1e_ovlp', hermi=1, kpts=kpts), dtype=numpy.complex128) - S12 = numpy.array(pgto.cell.intor_cross('int1e_ovlp', cell, cell_alt, kpts=kpts), dtype=numpy.complex128) + S22 = numpy.array( + cell_alt.pbc_intor("int1e_ovlp", hermi=1, kpts=kpts), dtype=numpy.complex128 + ) + S12 = numpy.array( + pgto.cell.intor_cross("int1e_ovlp", cell, cell_alt, kpts=kpts), + dtype=numpy.complex128, + ) + + return (S12, S22) - return(S12, S22) def remove_core_mo_k(Clo, Ccore, S, thr=0.5): - assert(numpy.allclose(Clo.conj().T@S@Clo,numpy.eye(Clo.shape[1]))) - assert(numpy.allclose(Ccore.conj().T@S@Ccore,numpy.eye(Ccore.shape[1]))) + assert numpy.allclose(Clo.conj().T @ S @ Clo, numpy.eye(Clo.shape[1])) + assert numpy.allclose(Ccore.conj().T @ S @ Ccore, numpy.eye(Ccore.shape[1])) - n,nlo = Clo.shape + n, nlo = Clo.shape ncore = Ccore.shape[1] - Pcore = Ccore@Ccore.conj().T @ S + Pcore = Ccore @ Ccore.conj().T @ S Clo1 = (numpy.eye(n) - Pcore) @ Clo pop = numpy.diag(Clo1.conj().T @ S @ Clo1) - idx_keep = numpy.where(pop>thr)[0] - assert(len(idx_keep) == nlo-ncore) - Clo2 = symm_orth_k(Clo1[:,idx_keep], ovlp=S) + idx_keep = numpy.where(pop > thr)[0] + assert len(idx_keep) == nlo - ncore + Clo2 = symm_orth_k(Clo1[:, idx_keep], ovlp=S) return Clo2 - def get_iao_k(Co, S12, S1, S2=None, ortho=True): """ Args: @@ -107,9 +118,9 @@ def get_iao_k(Co, S12, S1, S2=None, ortho=True): P1[k] = scipy.linalg.inv(S1[k]) P2[k] = scipy.linalg.inv(S2[k]) - Ciao = numpy.zeros((nk, nao, S12.shape[-1]), dtype= numpy.complex128) + Ciao = numpy.zeros((nk, nao, S12.shape[-1]), dtype=numpy.complex128) for k in range(nk): - #Cotil = P1[k] @ S12[k] @ P2[k] @ S12[k].conj().T @ Co[k] + # Cotil = P1[k] @ S12[k] @ P2[k] @ S12[k].conj().T @ Co[k] Cotil = reduce(numpy.dot, (P1[k], S12[k], P2[k], S12[k].conj().T, Co[k])) ptil = numpy.dot(P1[k], S12[k]) Stil = reduce(numpy.dot, (Cotil.conj().T, S1[k], Cotil)) @@ -120,13 +131,17 @@ def get_iao_k(Co, S12, S1, S2=None, ortho=True): Potil = reduce(numpy.dot, (Cotil, Stil_inv, Cotil.conj().T)) - Ciao[k] = (numpy.eye(nao, dtype=numpy.complex128) - \ - numpy.dot((Po + Potil - 2.* reduce(numpy.dot,(Po, S1[k], Potil))), S1[k])) @ ptil + Ciao[k] = ( + numpy.eye(nao, dtype=numpy.complex128) + - numpy.dot( + (Po + Potil - 2.0 * reduce(numpy.dot, (Po, S1[k], Potil))), S1[k] + ) + ) @ ptil if ortho: Ciao[k] = symm_orth_k(Ciao[k], ovlp=S1[k]) rep_err = numpy.linalg.norm(Ciao[k] @ Ciao[k].conj().T @ S1[k] @ Po - Po) - if rep_err > 1.E-10: + if rep_err > 1.0e-10: raise RuntimeError return Ciao @@ -151,15 +166,16 @@ def get_pao_k(Ciao, S, S12, S2): nonval = numpy.eye(nao) - s12 @ s12.conj().T Piao = Ciao[k] @ Ciao[k].conj().T @ S[k] - cpao_ = (numpy.eye(nao) - Piao)@ nonval + cpao_ = (numpy.eye(nao) - Piao) @ nonval numpy.o0 = cpao_.shape[-1] - Cpao.append(cano_orth(cpao_,ovlp=S[k])) + Cpao.append(cano_orth(cpao_, ovlp=S[k])) numpy.o1 = Cpao[k].shape[-1] Cpao = numpy.asarray(Cpao) return Cpao + def get_pao_native_k(Ciao, S, mol, valence_basis, kpts, ortho=True): """ Args: @@ -181,7 +197,11 @@ def get_pao_native_k(Ciao, S, mol, valence_basis, kpts, ortho=True): full_ao_labels = mol.ao_labels() valence_ao_labels = mol_alt.ao_labels() - vir_idx = [idx for idx, label in enumerate(full_ao_labels) if (not label in valence_ao_labels)] + vir_idx = [ + idx + for idx, label in enumerate(full_ao_labels) + if (not label in valence_ao_labels) + ] niao = len(vir_idx) Cpao = numpy.zeros((nk, nao, niao), dtype=numpy.complex128) @@ -196,12 +216,9 @@ def get_pao_native_k(Ciao, S, mol, valence_basis, kpts, ortho=True): npao0 = cpao_.shape[1] Cpao[k] = cano_orth(cpao_, ovlp=S[k]) npao1 = cpao_.shape[1] - print("# of PAO: %d --> %d" % (npao0,npao1), flush=True) + print("# of PAO: %d --> %d" % (npao0, npao1), flush=True) print("", flush=True) else: - Cpao[k] = cpao_.copy() + Cpao[k] = cpao_.copy() return Cpao - - - diff --git a/kbe/misc.py b/kbe/misc.py index 459ddaeb..c12e65a8 100644 --- a/kbe/misc.py +++ b/kbe/misc.py @@ -4,6 +4,7 @@ import numpy, functools import numpy as np + def sgeom(cell, kmesh=None): """ Get a supercell pyscf.pbc.gto.Cell object @@ -21,38 +22,53 @@ def sgeom(cell, kmesh=None): return scell -def get_phase(cell, kpts, kmesh): +def get_phase(cell, kpts, kmesh): a_vec = cell.lattice_vectors() - Ts = lib.cartesian_prod((numpy.arange(kmesh[0]), - numpy.arange(kmesh[1]), - numpy.arange(kmesh[2]))) + Ts = lib.cartesian_prod( + (numpy.arange(kmesh[0]), numpy.arange(kmesh[1]), numpy.arange(kmesh[2])) + ) Rs = numpy.dot(Ts, a_vec) tmp_ = numpy.dot(Rs, kpts.T) NRs = Rs.shape[0] - phase = 1/numpy.sqrt(NRs) * numpy.exp(1j* numpy.dot(Rs, kpts.T)) + phase = 1 / numpy.sqrt(NRs) * numpy.exp(1j * numpy.dot(Rs, kpts.T)) return phase -def get_phase1(cell, kpts, kmesh): +def get_phase1(cell, kpts, kmesh): a_vec = cell.lattice_vectors() - Ts = lib.cartesian_prod((numpy.arange(kmesh[0]), - numpy.arange(kmesh[1]), - numpy.arange(kmesh[2]))) + Ts = lib.cartesian_prod( + (numpy.arange(kmesh[0]), numpy.arange(kmesh[1]), numpy.arange(kmesh[2])) + ) Rs = numpy.dot(Ts, a_vec) NRs = Rs.shape[0] - phase = numpy.exp(-1.0j* numpy.dot(Rs, kpts.T)) + phase = numpy.exp(-1.0j * numpy.dot(Rs, kpts.T)) return phase + class storePBE: - def __init__(self, Nocc, hf_veff, hcore, - S, C, hf_dm, hf_etot, W, lmo_coeff, - enuc, ek, - E_core, C_core, P_core, core_veff): + def __init__( + self, + Nocc, + hf_veff, + hcore, + S, + C, + hf_dm, + hf_etot, + W, + lmo_coeff, + enuc, + ek, + E_core, + C_core, + P_core, + core_veff, + ): self.Nocc = Nocc self.hf_veff = hf_veff self.hcore = hcore @@ -71,22 +87,23 @@ def __init__(self, Nocc, hf_veff, hcore, def print_energy(ecorr, e_V_Kapprox, e_F_dg, e_hf, unitcell_nkpt): - - # Print energy results - print('-----------------------------------------------------', - flush=True) - print(' BE ENERGIES with cumulant-based expression', flush=True) - - print('-----------------------------------------------------', - flush=True) - print(' E_BE = E_HF + Tr(F del g) + Tr(V K_approx)', flush=True) - print(' E_HF : {:>14.8f} Ha'.format(e_hf/unitcell_nkpt), flush=True) - print(' Tr(F del g) : {:>14.8f} Ha'.format(e_F_dg/unitcell_nkpt), flush=True) - print(' Tr(V K_aprrox) : {:>14.8f} Ha'.format(e_V_Kapprox/unitcell_nkpt), flush=True) - print(' E_BE : {:>14.8f} Ha'.format((ecorr + e_hf)/unitcell_nkpt), flush=True) - print(' Ecorr BE : {:>14.8f} Ha'.format(ecorr/unitcell_nkpt), flush=True) - print('-----------------------------------------------------', - flush=True) + print("-----------------------------------------------------", flush=True) + print(" BE ENERGIES with cumulant-based expression", flush=True) + + print("-----------------------------------------------------", flush=True) + print(" E_BE = E_HF + Tr(F del g) + Tr(V K_approx)", flush=True) + print(" E_HF : {:>14.8f} Ha".format(e_hf / unitcell_nkpt), flush=True) + print(" Tr(F del g) : {:>14.8f} Ha".format(e_F_dg / unitcell_nkpt), flush=True) + print( + " Tr(V K_aprrox) : {:>14.8f} Ha".format(e_V_Kapprox / unitcell_nkpt), + flush=True, + ) + print( + " E_BE : {:>14.8f} Ha".format((ecorr + e_hf) / unitcell_nkpt), + flush=True, + ) + print(" Ecorr BE : {:>14.8f} Ha".format(ecorr / unitcell_nkpt), flush=True) + print("-----------------------------------------------------", flush=True) print(flush=True) diff --git a/kbe/pbe.py b/kbe/pbe.py index 0eb53711..87444b7a 100644 --- a/kbe/pbe.py +++ b/kbe/pbe.py @@ -3,12 +3,13 @@ from .pfrag import Frags from molbe.helper import get_core import molbe.be_var as be_var -import numpy,functools,sys, pickle +import numpy, functools, sys, pickle from pyscf import lib import h5py, os from .misc import storePBE + class BE: """ Class for handling periodic bootstrap embedding (BE) calculations. @@ -28,16 +29,31 @@ class BE: lo_method : str Method for orbital localization, default is 'lowdin'. """ - def __init__(self, mf, fobj, eri_file='eri_file.h5', - lo_method='lowdin',compute_hf=True, - restart=False, save=False, - restart_file='storebe.pk', - mo_energy = None, - save_file='storebe.pk',hci_pt=False, - nproc=1, ompnum=4, - hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, - iao_val_core=True, - exxdiv='ewald', kpts = None, cderi = None, iao_wannier = False): + + def __init__( + self, + mf, + fobj, + eri_file="eri_file.h5", + lo_method="lowdin", + compute_hf=True, + restart=False, + save=False, + restart_file="storebe.pk", + mo_energy=None, + save_file="storebe.pk", + hci_pt=False, + nproc=1, + ompnum=4, + hci_cutoff=0.001, + ci_coeff_cutoff=None, + select_cutoff=None, + iao_val_core=True, + exxdiv="ewald", + kpts=None, + cderi=None, + iao_wannier=False, + ): """ Constructor for BE object. @@ -74,7 +90,7 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', """ if restart: # Load previous calculation data from restart file - with open(restart_file, 'rb') as rfile: + with open(restart_file, "rb") as rfile: store_ = pickle.load(rfile) rfile.close() self.Nocc = store_.Nocc @@ -97,7 +113,7 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.ompnum = ompnum # Fragment information from fobj - self.frag_type=fobj.frag_type + self.frag_type = fobj.frag_type self.Nfrag = fobj.Nfrag self.fsites = fobj.fsites self.edge = fobj.edge @@ -114,13 +130,15 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', unitcell_nkpt = 1 for i in self.kmesh: - if i>1: unitcell_nkpt *= self.unitcell + if i > 1: + unitcell_nkpt *= self.unitcell self.unitcell_nkpt = unitcell_nkpt - self.ebe_hf = 0. + self.ebe_hf = 0.0 nkpts_ = 1 for i in self.kmesh: - if i>1: nkpts_ *= i + if i > 1: + nkpts_ *= i self.nkpt = nkpts_ self.kpts = kpts @@ -128,19 +146,21 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.hci_cutoff = hci_cutoff self.ci_coeff_cutoff = ci_coeff_cutoff self.select_cutoff = select_cutoff - self.hci_pt=hci_pt + self.hci_pt = hci_pt if not restart: self.mo_energy = mf.mo_energy mf.exxdiv = None self.mf = mf - self.Nocc = mf.cell.nelectron//2 + self.Nocc = mf.cell.nelectron // 2 self.enuc = mf.energy_nuc() self.hcore = mf.get_hcore() self.S = mf.get_ovlp() self.C = numpy.array(mf.mo_coeff) self.hf_dm = mf.make_rdm1() - self.hf_veff = mf.get_veff(self.cell, dm_kpts = self.hf_dm, hermi=1, kpts=self.kpts, kpts_band=None) + self.hf_veff = mf.get_veff( + self.cell, dm_kpts=self.hf_dm, hermi=1, kpts=self.kpts, kpts_band=None + ) self.hf_etot = mf.e_tot self.W = None self.lmo_coeff = None @@ -152,36 +172,40 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.cderi = cderi # Set scratch directory - jobid='' + jobid = "" if be_var.CREATE_SCRATCH_DIR: try: - jobid = str(os.environ['SLURM_JOB_ID']) + jobid = str(os.environ["SLURM_JOB_ID"]) except: - jobid = '' - if not be_var.SCRATCH=='': os.system('mkdir '+be_var.SCRATCH+str(jobid)) - if jobid == '': - self.eri_file = be_var.SCRATCH+eri_file + jobid = "" + if not be_var.SCRATCH == "": + os.system("mkdir " + be_var.SCRATCH + str(jobid)) + if jobid == "": + self.eri_file = be_var.SCRATCH + eri_file if cderi: - self.cderi = be_var.SCRATCH+cderi + self.cderi = be_var.SCRATCH + cderi else: - self.eri_file = be_var.SCRATCH+str(jobid)+'/'+eri_file + self.eri_file = be_var.SCRATCH + str(jobid) + "/" + eri_file if cderi: - self.cderi = be_var.SCRATCH+str(jobid)+'/'+cderi - + self.cderi = be_var.SCRATCH + str(jobid) + "/" + cderi - if exxdiv == 'ewald': + if exxdiv == "ewald": if not restart: self.ek = self.ewald_sum(kpts=self.kpts) - print('Energy contribution from Ewald summation : {:>12.8f} Ha'.format(self.ek), - flush=True) - print('Total HF Energy will contain this contribution. ') + print( + "Energy contribution from Ewald summation : {:>12.8f} Ha".format( + self.ek + ), + flush=True, + ) + print("Total HF Energy will contain this contribution. ") print(flush=True) elif exxdiv is None: - print('Setting exxdiv=None') - self.ek = 0. + print("Setting exxdiv=None") + self.ek = 0.0 else: - print('exxdiv = ',exxdiv,'not implemented!',flush=True) - print('Energy may diverse.',flush=True) + print("exxdiv = ", exxdiv, "not implemented!", flush=True) + print("Energy may diverse.", flush=True) print(flush=True) self.frozen_core = False if not fobj.frozen_core else True @@ -199,72 +223,103 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.core_list = fobj.core_list if not restart: - self.Nocc -=self.ncore + self.Nocc -= self.ncore nk, nao, nao = self.hf_dm.shape - dm_nocore = numpy.zeros((nk, nao, nao), dtype=numpy.result_type(self.C, self.C)) + dm_nocore = numpy.zeros( + (nk, nao, nao), dtype=numpy.result_type(self.C, self.C) + ) C_core = numpy.zeros((nk, nao, self.ncore), dtype=self.C.dtype) - P_core = numpy.zeros((nk, nao, nao), dtype=numpy.result_type(self.C, self.C)) + P_core = numpy.zeros( + (nk, nao, nao), dtype=numpy.result_type(self.C, self.C) + ) for k in range(nk): - dm_nocore[k]+= 2.*numpy.dot(self.C[k][:,self.ncore:self.ncore+self.Nocc], - self.C[k][:,self.ncore:self.ncore+self.Nocc].conj().T) - C_core[k] += self.C[k][:,:self.ncore] + dm_nocore[k] += 2.0 * numpy.dot( + self.C[k][:, self.ncore : self.ncore + self.Nocc], + self.C[k][:, self.ncore : self.ncore + self.Nocc].conj().T, + ) + C_core[k] += self.C[k][:, : self.ncore] P_core[k] += numpy.dot(C_core[k], C_core[k].conj().T) - self.C_core = C_core self.P_core = P_core self.hf_dm = dm_nocore - self.core_veff = mf.get_veff(self.cell, dm_kpts = self.P_core*2., hermi=1, kpts=self.kpts, kpts_band=None) - - ecore_h1 = 0. - ecore_veff = 0. + self.core_veff = mf.get_veff( + self.cell, + dm_kpts=self.P_core * 2.0, + hermi=1, + kpts=self.kpts, + kpts_band=None, + ) + + ecore_h1 = 0.0 + ecore_veff = 0.0 for k in range(nk): - ecore_h1 += numpy.einsum('ij,ji', self.hcore[k], 2.*self.P_core[k]) - ecore_veff += numpy.einsum('ij,ji', 2.*self.P_core[k], self.core_veff[k]) * .5 + ecore_h1 += numpy.einsum( + "ij,ji", self.hcore[k], 2.0 * self.P_core[k] + ) + ecore_veff += ( + numpy.einsum("ij,ji", 2.0 * self.P_core[k], self.core_veff[k]) + * 0.5 + ) ecore_h1 /= float(nk) ecore_veff /= float(nk) E_core = ecore_h1 + ecore_veff - if numpy.abs(E_core.imag).max() < 1.e-10: + if numpy.abs(E_core.imag).max() < 1.0e-10: self.E_core = E_core.real else: - - print('Imaginary density in E_core ', numpy.abs(E_core.imag).max()) + print("Imaginary density in E_core ", numpy.abs(E_core.imag).max()) sys.exit() for k in range(nk): self.hf_veff[k] -= self.core_veff[k] self.hcore[k] += self.core_veff[k] - # Needed for Wannier localization - if lo_method=='wannier' or iao_wannier: + if lo_method == "wannier" or iao_wannier: self.FOCK = self.mf.get_fock(self.hcore, self.S, self.hf_veff, self.hf_dm) - if not restart: # Localize orbitals - self.localize(lo_method, mol=self.cell, valence_basis=fobj.valence_basis, - iao_wannier=iao_wannier, iao_val_core=iao_val_core) + self.localize( + lo_method, + mol=self.cell, + valence_basis=fobj.valence_basis, + iao_wannier=iao_wannier, + iao_val_core=iao_val_core, + ) if save: # Save intermediate results for restart - store_ = storePBE(self.Nocc, self.hf_veff, self.hcore, - self.S, self.C, self.hf_dm, self.hf_etot, - self.W, self.lmo_coeff, self.enuc, self.ek, - self.E_core, self.C_core, self.P_core, self.core_veff) - with open(save_file, 'wb') as rfile: + store_ = storePBE( + self.Nocc, + self.hf_veff, + self.hcore, + self.S, + self.C, + self.hf_dm, + self.hf_etot, + self.W, + self.lmo_coeff, + self.enuc, + self.ek, + self.E_core, + self.C_core, + self.P_core, + self.core_veff, + ) + with open(save_file, "wb") as rfile: pickle.dump(store_, rfile, pickle.HIGHEST_PROTOCOL) rfile.close() - if not restart : - self.initialize(mf._eri,compute_hf) - + if not restart: + self.initialize(mf._eri, compute_hf) from ._opt import optimize + # this is a molbe method not BEOPT from molbe.external.optqn import get_be_error_jacobian from .lo import localize @@ -273,22 +328,20 @@ def print_ini(self): """ Print initialization banner for the kBE calculation. """ - print('-----------------------------------------------------------', - flush=True) - - print(' BBBBBBB EEEEEEE ',flush=True) - print(' BB B EE ',flush=True) - print(' PP PP BB B EE ',flush=True) - print(' PP PP BBBBBBB EEEEEEE ',flush=True) - print(' PPPP BB B EE ',flush=True) - print(' PP PP BB B EE ',flush=True) - print(' PP PP BBBBBBB EEEEEEE ',flush=True) + print("-----------------------------------------------------------", flush=True) + + print(" BBBBBBB EEEEEEE ", flush=True) + print(" BB B EE ", flush=True) + print(" PP PP BB B EE ", flush=True) + print(" PP PP BBBBBBB EEEEEEE ", flush=True) + print(" PPPP BB B EE ", flush=True) + print(" PP PP BB B EE ", flush=True) + print(" PP PP BBBBBBB EEEEEEE ", flush=True) print(flush=True) - print(' PERIODIC BOOTSTRAP EMBEDDING',flush=True) - print(' BEn = ',self.be_type,flush=True) - print('-----------------------------------------------------------', - flush=True) + print(" PERIODIC BOOTSTRAP EMBEDDING", flush=True) + print(" BEn = ", self.be_type, flush=True) + print("-----------------------------------------------------------", flush=True) print(flush=True) def ewald_sum(self, kpts=None): @@ -298,15 +351,19 @@ def ewald_sum(self, kpts=None): nk, nao, nao = dm_.shape vk_kpts = numpy.zeros(dm_.shape) * 1j - _ewald_exxdiv_for_G0(self.mf.cell, self.kpts, dm_.reshape(-1, nk, nao, nao), - vk_kpts.reshape(-1, nk, nao, nao), - self.kpts) - e_ = numpy.einsum("kij,kji->",vk_kpts,dm_)*0.25 + _ewald_exxdiv_for_G0( + self.mf.cell, + self.kpts, + dm_.reshape(-1, nk, nao, nao), + vk_kpts.reshape(-1, nk, nao, nao), + self.kpts, + ) + e_ = numpy.einsum("kij,kji->", vk_kpts, dm_) * 0.25 e_ /= float(nk) return e_.real - def initialize(self, eri_,compute_hf, restart=False): + def initialize(self, eri_, compute_hf, restart=False): """ Initialize the Bootstrap Embedding calculation. @@ -329,66 +386,104 @@ def initialize(self, eri_,compute_hf, restart=False): from pyscf.pbc import ao2mo as pao2mo from libdmet.basis_transform.eri_transform import get_emb_eri_fast_gdf - if compute_hf: E_hf = 0. - EH1 = 0. - ECOUL = 0. - EF = 0. + if compute_hf: + E_hf = 0.0 + EH1 = 0.0 + ECOUL = 0.0 + EF = 0.0 # Create a file to store ERIs if not restart: - file_eri = h5py.File(self.eri_file,'w') + file_eri = h5py.File(self.eri_file, "w") lentmp = len(self.edge_idx) - transform_parallel=False # hard set for now + transform_parallel = False # hard set for now for I in range(self.Nfrag): if lentmp: - fobjs_ = Frags(self.fsites[I], I, edge=self.edge[I], - eri_file=self.eri_file, - center=self.center[I], edge_idx=self.edge_idx[I], - center_idx=self.center_idx[I],efac=self.ebe_weight[I], - centerf_idx=self.centerf_idx[I], unitcell=self.unitcell, - unitcell_nkpt=self.unitcell_nkpt) + fobjs_ = Frags( + self.fsites[I], + I, + edge=self.edge[I], + eri_file=self.eri_file, + center=self.center[I], + edge_idx=self.edge_idx[I], + center_idx=self.center_idx[I], + efac=self.ebe_weight[I], + centerf_idx=self.centerf_idx[I], + unitcell=self.unitcell, + unitcell_nkpt=self.unitcell_nkpt, + ) else: - fobjs_ = Frags(self.fsites[I],I,edge=[],center=[], - eri_file=self.eri_file, - edge_idx=[],center_idx=[],centerf_idx=[], - efac=self.ebe_weight[I], unitcell=self.unitcell, - unitcell_nkpt=self.unitcell_nkpt) - - fobjs_.sd(self.W, self.lmo_coeff, self.Nocc, kmesh=self.kmesh, - cell=self.cell, frag_type=self.frag_type, kpts=self.kpts, h1=self.hcore) + fobjs_ = Frags( + self.fsites[I], + I, + edge=[], + center=[], + eri_file=self.eri_file, + edge_idx=[], + center_idx=[], + centerf_idx=[], + efac=self.ebe_weight[I], + unitcell=self.unitcell, + unitcell_nkpt=self.unitcell_nkpt, + ) + + fobjs_.sd( + self.W, + self.lmo_coeff, + self.Nocc, + kmesh=self.kmesh, + cell=self.cell, + frag_type=self.frag_type, + kpts=self.kpts, + h1=self.hcore, + ) fobjs_.cons_h1(self.hcore) fobjs_.heff = numpy.zeros_like(fobjs_.h1) - fobjs_.dm_init = fobjs_.get_nsocc(self.S, self.C, self.Nocc, ncore=self.ncore) + fobjs_.dm_init = fobjs_.get_nsocc( + self.S, self.C, self.Nocc, ncore=self.ncore + ) if self.cderi is None: if not restart: - - eri = get_emb_eri_fast_gdf(self.mf.cell, self.mf.with_df, - t_reversal_symm=True, - symmetry=4, - C_ao_eo=fobjs_.TA)[0] + eri = get_emb_eri_fast_gdf( + self.mf.cell, + self.mf.with_df, + t_reversal_symm=True, + symmetry=4, + C_ao_eo=fobjs_.TA, + )[0] file_eri.create_dataset(fobjs_.dname, data=eri) eri = ao2mo.restore(8, eri, fobjs_.nao) fobjs_.cons_fock(self.hf_veff, self.S, self.hf_dm, eri_=eri) else: - eri=None + eri = None self.Fobjs.append(fobjs_) # ERI & Fock parallelization for periodic calculations if self.cderi: if self.nproc == 1: - print('If cderi is set, try again with nproc > 1') + print("If cderi is set, try again with nproc > 1") sys.exit() - nprocs = int(self.nproc/self.ompnum) + nprocs = int(self.nproc / self.ompnum) pool_ = Pool(nprocs) - os.system('export OMP_NUM_THREADS='+str(self.ompnum)) + os.system("export OMP_NUM_THREADS=" + str(self.ompnum)) results = [] eris = [] for frg in range(self.Nfrag): - result = pool_.apply_async(eritransform_parallel, [self.mf.cell.a, self.mf.cell.atom, self.mf.cell.basis, self.kpts, self.Fobjs[frg].TA, self.cderi]) + result = pool_.apply_async( + eritransform_parallel, + [ + self.mf.cell.a, + self.mf.cell.atom, + self.mf.cell.basis, + self.kpts, + self.Fobjs[frg].TA, + self.cderi, + ], + ) results.append(result) [eris.append(result.get()) for result in results] pool_.close() @@ -398,37 +493,47 @@ def initialize(self, eri_,compute_hf, restart=False): eris = None file_eri.close() - nprocs = int(self.nproc/self.ompnum) + nprocs = int(self.nproc / self.ompnum) pool_ = Pool(nprocs) results = [] veffs = [] for frg in range(self.Nfrag): - result = pool_.apply_async(parallel_fock_wrapper, [self.Fobjs[frg].dname, self.Fobjs[frg].nao, self.hf_dm, self.S, self.Fobjs[frg].TA, self.hf_veff, self.eri_file]) + result = pool_.apply_async( + parallel_fock_wrapper, + [ + self.Fobjs[frg].dname, + self.Fobjs[frg].nao, + self.hf_dm, + self.S, + self.Fobjs[frg].TA, + self.hf_veff, + self.eri_file, + ], + ) results.append(result) [veffs.append(result.get()) for result in results] pool_.close() for frg in range(self.Nfrag): veff0, veff_ = veffs[frg] - if numpy.abs(veff_.imag).max() < 1.e-6: + if numpy.abs(veff_.imag).max() < 1.0e-6: self.Fobjs[frg].veff = veff_.real self.Fobjs[frg].veff0 = veff0.real else: - print('Imaginary Veff ', numpy.abs(veff_.imag).max()) + print("Imaginary Veff ", numpy.abs(veff_.imag).max()) sys.exit() self.Fobjs[frg].fock = self.Fobjs[frg].h1 + veff_.real veffs = None - # SCF parallelized if self.nproc == 1 and not transform_parallel: for frg in range(self.Nfrag): # SCF - self.Fobjs[frg].scf(fs=True, dm0 = self.Fobjs[frg].dm_init) + self.Fobjs[frg].scf(fs=True, dm0=self.Fobjs[frg].dm_init) else: - nprocs = int(self.nproc/self.ompnum) + nprocs = int(self.nproc / self.ompnum) pool_ = Pool(nprocs) - os.system('export OMP_NUM_THREADS='+str(self.ompnum)) + os.system("export OMP_NUM_THREADS=" + str(self.ompnum)) results = [] mo_coeffs = [] for frg in range(self.Nfrag): @@ -436,9 +541,10 @@ def initialize(self, eri_,compute_hf, restart=False): nocc = self.Fobjs[frg].nsocc dname = self.Fobjs[frg].dname h1 = self.Fobjs[frg].fock + self.Fobjs[frg].heff - result = pool_.apply_async(parallel_scf_wrapper, [dname, nao, nocc, h1, - self.Fobjs[frg].dm_init, - self.eri_file]) + result = pool_.apply_async( + parallel_scf_wrapper, + [dname, nao, nocc, h1, self.Fobjs[frg].dm_init, self.eri_file], + ) results.append(result) [mo_coeffs.append(result.get()) for result in results] pool_.close() @@ -446,8 +552,13 @@ def initialize(self, eri_,compute_hf, restart=False): self.Fobjs[frg]._mo_coeffs = mo_coeffs[frg] for frg in range(self.Nfrag): - self.Fobjs[frg].dm0 = numpy.dot( self.Fobjs[frg]._mo_coeffs[:,:self.Fobjs[frg].nsocc], - self.Fobjs[frg]._mo_coeffs[:,:self.Fobjs[frg].nsocc].conj().T) *2. + self.Fobjs[frg].dm0 = ( + numpy.dot( + self.Fobjs[frg]._mo_coeffs[:, : self.Fobjs[frg].nsocc], + self.Fobjs[frg]._mo_coeffs[:, : self.Fobjs[frg].nsocc].conj().T, + ) + * 2.0 + ) # energy if compute_hf: @@ -459,23 +570,26 @@ def initialize(self, eri_,compute_hf, restart=False): file_eri.close() if compute_hf: - E_hf /= self.unitcell_nkpt - hf_err = self.hf_etot-(E_hf+self.enuc+self.E_core) + hf_err = self.hf_etot - (E_hf + self.enuc + self.E_core) - self.ebe_hf = E_hf+self.enuc+self.E_core-self.ek - print('HF-in-HF error : {:>.4e} Ha'. - format(hf_err), flush=True) + self.ebe_hf = E_hf + self.enuc + self.E_core - self.ek + print( + "HF-in-HF error : {:>.4e} Ha".format(hf_err), + flush=True, + ) - if abs(hf_err)>1.e-5: - print('WARNING!!! Large HF-in-HF energy error') + if abs(hf_err) > 1.0e-5: + print("WARNING!!! Large HF-in-HF energy error") couti = 0 for fobj in self.Fobjs: fobj.udim = couti couti = fobj.set_udim(couti) - def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False): + def oneshot( + self, solver="MP2", nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False + ): """ Perform a one-shot bootstrap embedding calculation. @@ -497,31 +611,59 @@ def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean print("Calculating Energy by Fragment? ", calc_frag_energy) if nproc == 1: - rets = be_func(None, self.Fobjs, self.Nocc, solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff = self.ci_coeff_cutoff, - select_cutoff = self.select_cutoff, - nproc=ompnum, frag_energy=calc_frag_energy, - ereturn=True, eeval=True) + rets = be_func( + None, + self.Fobjs, + self.Nocc, + solver, + self.enuc, + hf_veff=self.hf_veff, + hci_cutoff=self.hci_cutoff, + ci_coeff_cutoff=self.ci_coeff_cutoff, + select_cutoff=self.select_cutoff, + nproc=ompnum, + frag_energy=calc_frag_energy, + ereturn=True, + eeval=True, + ) else: - rets = be_func_parallel(None, self.Fobjs, self.Nocc, solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff = self.ci_coeff_cutoff, - select_cutoff = self.select_cutoff, - ereturn=True, eeval=True, frag_energy=calc_frag_energy, - nproc=nproc, ompnum=ompnum) - - print('-----------------------------------------------------', - flush=True) - print(' One Shot BE ', flush=True) - print(' Solver : ',solver,flush=True) - print('-----------------------------------------------------', - flush=True) + rets = be_func_parallel( + None, + self.Fobjs, + self.Nocc, + solver, + self.enuc, + hf_veff=self.hf_veff, + hci_cutoff=self.hci_cutoff, + ci_coeff_cutoff=self.ci_coeff_cutoff, + select_cutoff=self.select_cutoff, + ereturn=True, + eeval=True, + frag_energy=calc_frag_energy, + nproc=nproc, + ompnum=ompnum, + ) + + print("-----------------------------------------------------", flush=True) + print(" One Shot BE ", flush=True) + print(" Solver : ", solver, flush=True) + print("-----------------------------------------------------", flush=True) print(flush=True) if calc_frag_energy: - print("Final Tr(F del g) is : {:>12.8f} Ha".format(rets[1][0]+rets[1][2]), flush=True) - print("Final Tr(V K_approx) is : {:>12.8f} Ha".format(rets[1][1]), flush=True) - print("Final e_corr is : {:>12.8f} Ha".format(rets[0]), flush=True) + print( + "Final Tr(F del g) is : {:>12.8f} Ha".format( + rets[1][0] + rets[1][2] + ), + flush=True, + ) + print( + "Final Tr(V K_approx) is : {:>12.8f} Ha".format(rets[1][1]), + flush=True, + ) + print( + "Final e_corr is : {:>12.8f} Ha".format(rets[0]), + flush=True, + ) self.ebe_tot = rets[0] @@ -551,7 +693,7 @@ def update_fock(self, heff=None): for idx, fobj in self.Fobjs: fobj.fock += heff[idx] - def write_heff(self, heff_file='bepotfile.h5'): + def write_heff(self, heff_file="bepotfile.h5"): """ Write the effective Hamiltonian to a file. @@ -560,13 +702,13 @@ def write_heff(self, heff_file='bepotfile.h5'): heff_file : str, optional Path to the file to store effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, 'w') + filepot = h5py.File(heff_file, "w") for fobj in self.Fobjs: print(fobj.heff.shape, fobj.dname, flush=True) filepot.create_dataset(fobj.dname, data=fobj.heff) filepot.close() - def read_heff(self, heff_file='bepotfile.h5'): + def read_heff(self, heff_file="bepotfile.h5"): """ Read the effective Hamiltonian from a file. @@ -575,13 +717,12 @@ def read_heff(self, heff_file='bepotfile.h5'): heff_file : str, optional Path to the file storing effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, 'r') + filepot = h5py.File(heff_file, "r") for fobj in self.Fobjs: fobj.heff = filepot.get(fobj.dname) filepot.close() - def initialize_pot(Nfrag, edge_idx): """ Initialize the potential array for bootstrap embedding. @@ -604,18 +745,18 @@ def initialize_pot(Nfrag, edge_idx): list of float Initialized potential array with zeros. """ - pot_=[] + pot_ = [] if not len(edge_idx) == 0: for I in range(Nfrag): for i in edge_idx[I]: for j in range(len(i)): for k in range(len(i)): - if j>k: + if j > k: continue pot_.append(0.0) - pot_.append(0.) + pot_.append(0.0) return pot_ @@ -630,18 +771,19 @@ def eritransform_parallel(a, atom, basis, kpts, C_ao_emb, cderi): cell.a = a cell.atom = atom cell.basis = basis - cell.charge=0 - cell.verbose=0 + cell.charge = 0 + cell.verbose = 0 cell.build() mydf = df.GDF(cell, kpts) mydf._cderi = cderi - eri = get_emb_eri_fast_gdf(cell, mydf, - t_reversal_symm=True, symmetry=4, - C_ao_eo=C_ao_emb) + eri = get_emb_eri_fast_gdf( + cell, mydf, t_reversal_symm=True, symmetry=4, C_ao_eo=C_ao_emb + ) return eri + def parallel_fock_wrapper(dname, nao, dm, S, TA, hf_veff, eri_file): """ Wrapper for parallel Fock transformation @@ -649,16 +791,17 @@ def parallel_fock_wrapper(dname, nao, dm, S, TA, hf_veff, eri_file): from .helper import get_veff, get_eri eri_ = get_eri(dname, nao, eri_file=eri_file, ignore_symm=True) - veff0, veff_ = get_veff(eri_, dm, S, TA, hf_veff, return_veff0 = True) + veff0, veff_ = get_veff(eri_, dm, S, TA, hf_veff, return_veff0=True) return veff0, veff_ -def parallel_scf_wrapper(dname, nao, nocc, h1, dm_init, eri_file): +def parallel_scf_wrapper(dname, nao, nocc, h1, dm_init, eri_file): """ Wrapper for performing fragment scf calculation """ from .helper import get_eri, get_scfObj + eri = get_eri(dname, nao, eri_file=eri_file) mf_ = get_scfObj(h1, eri, nocc, dm_init) diff --git a/kbe/pfrag.py b/kbe/pfrag.py index 145b90f7..87a41f73 100644 --- a/kbe/pfrag.py +++ b/kbe/pfrag.py @@ -4,10 +4,11 @@ from .helper import * from molbe.helper import get_eri, get_scfObj from .misc import * -import numpy,h5py -import functools,sys, math +import numpy, h5py +import functools, sys, math from pyscf import ao2mo + class Frags: """ Class for handling fragments in periodic bootstrap embedding. @@ -15,10 +16,22 @@ class Frags: This class contains various functionalities required for managing and manipulating fragments for periodic BE calculations. """ - def __init__(self, fsites, ifrag, edge=None, center=None, - edge_idx=None, center_idx=None, efac=None, - eri_file='eri_file.h5',unitcell_nkpt=1, - ewald_ek=None, centerf_idx=None, unitcell=1): + + def __init__( + self, + fsites, + ifrag, + edge=None, + center=None, + edge_idx=None, + center_idx=None, + efac=None, + eri_file="eri_file.h5", + unitcell_nkpt=1, + ewald_ek=None, + centerf_idx=None, + unitcell=1, + ): """Constructor function for `Frags` class. Parameters @@ -44,14 +57,14 @@ def __init__(self, fsites, ifrag, edge=None, center=None, """ self.fsites = fsites - self.unitcell=unitcell - self.unitcell_nkpt=unitcell_nkpt + self.unitcell = unitcell + self.unitcell_nkpt = unitcell_nkpt self.nfsites = len(fsites) self.TA = None self.TA_lo_eo = None self.h1 = None self.ifrag = ifrag - self.dname = 'f'+str(ifrag) + self.dname = "f" + str(ifrag) self.nao = None self.mo_coeffs = None self._mo_coeffs = None @@ -77,23 +90,31 @@ def __init__(self, fsites, ifrag, edge=None, center=None, self._del_rdm1 = None self.rdm1 = None self.genvs = None - self.ebe = 0. - self.ebe_hf = 0. + self.ebe = 0.0 + self.ebe_hf = 0.0 self.efac = efac - self.ewald_ek=ewald_ek + self.ewald_ek = ewald_ek self.fock = None self.veff = None self.veff0 = None self.dm_init = None self.dm0 = None self.eri_file = eri_file - self.pot=None - self.ebe_hf0 = 0. + self.pot = None + self.ebe_hf0 = 0.0 self.rdm1_lo_k = None - def sd(self, lao, lmo, nocc, - frag_type='autogen', - cell=None, kpts = None, kmesh=None, h1=None): + def sd( + self, + lao, + lmo, + nocc, + frag_type="autogen", + cell=None, + kpts=None, + kmesh=None, + h1=None, + ): """ Perform Schmidt decomposition for the fragment. @@ -116,23 +137,22 @@ def sd(self, lao, lmo, nocc, from .misc import get_phase nk, nao, nlo = lao.shape - rdm1_lo_k = numpy.zeros((nk, nlo, nlo), - dtype=numpy.result_type(lmo, lmo)) + rdm1_lo_k = numpy.zeros((nk, nlo, nlo), dtype=numpy.result_type(lmo, lmo)) for k in range(nk): - rdm1_lo_k[k] += numpy.dot(lmo[k][:,:nocc], lmo[k][:,:nocc].conj().T) + rdm1_lo_k[k] += numpy.dot(lmo[k][:, :nocc], lmo[k][:, :nocc].conj().T) self.rdm1_lo_k = rdm1_lo_k phase = get_phase(cell, kpts, kmesh) - supcell_rdm = numpy.einsum('Rk,kuv,Sk->RuSv', phase, rdm1_lo_k, phase.conj()) - supcell_rdm = supcell_rdm.reshape(nk*nlo, nk*nlo) + supcell_rdm = numpy.einsum("Rk,kuv,Sk->RuSv", phase, rdm1_lo_k, phase.conj()) + supcell_rdm = supcell_rdm.reshape(nk * nlo, nk * nlo) - if numpy.abs(supcell_rdm.imag).max() < 1.e-6: + if numpy.abs(supcell_rdm.imag).max() < 1.0e-6: supcell_rdm = supcell_rdm.real else: - print('Imaginary density in Full SD', numpy.abs(supcell_rdm.imag).max()) + print("Imaginary density in Full SD", numpy.abs(supcell_rdm.imag).max()) sys.exit() - Sites = [i+(nlo*0) for i in self.fsites] - if not frag_type == 'autogen': + Sites = [i + (nlo * 0) for i in self.fsites] + if not frag_type == "autogen": Sites.sort() TA_R = schmidt_decomp_svd(supcell_rdm, Sites) @@ -140,11 +160,12 @@ def sd(self, lao, lmo, nocc, TA_R = TA_R.reshape(nk, nlo, teo) phase1 = get_phase1(cell, kpts, kmesh) - TA_k = numpy.einsum('Rim, Rk -> kim', TA_R, phase1) + TA_k = numpy.einsum("Rim, Rk -> kim", TA_R, phase1) self.TA_lo_eo = TA_k - TA_ao_eo_k = numpy.zeros((nk, nao, teo), - dtype=numpy.result_type(lao.dtype, TA_k.dtype)) + TA_ao_eo_k = numpy.zeros( + (nk, nao, teo), dtype=numpy.result_type(lao.dtype, TA_k.dtype) + ) for k in range(nk): TA_ao_eo_k[k] = numpy.dot(lao[k], TA_k[k]) @@ -154,22 +175,23 @@ def sd(self, lao, lmo, nocc, # useful for debugging -- rdm1_eo = numpy.zeros((teo, teo), dtype=numpy.complex128) for k in range(nk): - rdm1_eo += functools.reduce(numpy.dot, - (TA_k[k].conj().T, rdm1_lo_k[k], - TA_k[k])) + rdm1_eo += functools.reduce( + numpy.dot, (TA_k[k].conj().T, rdm1_lo_k[k], TA_k[k]) + ) rdm1_eo /= float(nk) h1_eo = numpy.zeros((teo, teo), dtype=numpy.complex128) for k in range(nk): - h1_eo += functools.reduce(numpy.dot, - (self.TA[k].conj().T, h1[k], - self.TA[k])) + h1_eo += functools.reduce( + numpy.dot, (self.TA[k].conj().T, h1[k], self.TA[k]) + ) h1_eo /= float(nk) - e1 = 2.0 *numpy.einsum("ij,ij->i", h1_eo[:self.nfsites], - rdm1_eo[:self.nfsites]) - e_h1 = 0. + e1 = 2.0 * numpy.einsum( + "ij,ij->i", h1_eo[: self.nfsites], rdm1_eo[: self.nfsites] + ) + e_h1 = 0.0 for i in self.efac[1]: - e_h1 += self.efac[0]*e1[i] + e_h1 += self.efac[0] * e1[i] def cons_h1(self, h1): """ @@ -184,15 +206,15 @@ def cons_h1(self, h1): nk, nao, teo = self.TA.shape h1_eo = numpy.zeros((teo, teo), dtype=numpy.complex128) for k in range(nk): - h1_eo += functools.reduce(numpy.dot, - (self.TA[k].conj().T, h1[k], - self.TA[k])) + h1_eo += functools.reduce( + numpy.dot, (self.TA[k].conj().T, h1[k], self.TA[k]) + ) h1_eo /= float(nk) - if numpy.abs(h1_eo.imag).max() < 1.e-7: + if numpy.abs(h1_eo.imag).max() < 1.0e-7: self.h1 = h1_eo.real else: - print('Imaginary Hcore ', numpy.abs(h1_eo.imag).max()) + print("Imaginary Hcore ", numpy.abs(h1_eo.imag).max()) sys.exit() def cons_fock(self, hf_veff, S, dm, eri_=None): @@ -212,18 +234,20 @@ def cons_fock(self, hf_veff, S, dm, eri_=None): """ if eri_ is None: - eri_ = get_eri(self.dname, self.TA.shape[1], ignore_symm=True, eri_file=self.eri_file) + eri_ = get_eri( + self.dname, self.TA.shape[1], ignore_symm=True, eri_file=self.eri_file + ) veff0, veff_ = get_veff(eri_, dm, S, self.TA, hf_veff, return_veff0=True) - if numpy.abs(veff_.imag).max() < 1.e-6: + if numpy.abs(veff_.imag).max() < 1.0e-6: self.veff = veff_.real self.veff0 = veff0.real else: - print('Imaginary Veff ', numpy.abs(veff_.imag).max()) + print("Imaginary Veff ", numpy.abs(veff_.imag).max()) sys.exit() self.fock = self.h1 + veff_.real - def get_nsocc(self, S, C, nocc,ncore=0): + def get_nsocc(self, S, C, nocc, ncore=0): """ Get the number of occupied orbitals for the fragment. @@ -246,31 +270,38 @@ def get_nsocc(self, S, C, nocc,ncore=0): import scipy.linalg nk, nao, neo = self.TA.shape - dm_ = numpy.zeros((nk, nao, nao), dtype=numpy.result_type(C,C)) + dm_ = numpy.zeros((nk, nao, nao), dtype=numpy.result_type(C, C)) for k in range(nk): - dm_[k] = 2.* numpy.dot(C[k][:,ncore:ncore+nocc], C[k][:,ncore:ncore+nocc].conj().T) + dm_[k] = 2.0 * numpy.dot( + C[k][:, ncore : ncore + nocc], C[k][:, ncore : ncore + nocc].conj().T + ) P_ = numpy.zeros((neo, neo), dtype=numpy.complex128) for k in range(nk): Cinv = numpy.dot(self.TA[k].conj().T, S[k]) - P_ += functools.reduce(numpy.dot, - (Cinv, dm_[k], Cinv.conj().T)) + P_ += functools.reduce(numpy.dot, (Cinv, dm_[k], Cinv.conj().T)) P_ /= float(nk) - if numpy.abs(P_.imag).max() < 1.e-6: + if numpy.abs(P_.imag).max() < 1.0e-6: P_ = P_.real else: - print('Imaginary density in get_nsocc ', numpy.abs(P_.imag).max()) + print("Imaginary density in get_nsocc ", numpy.abs(P_.imag).max()) sys.exit() nsocc_ = numpy.trace(P_) - nsocc = int(numpy.round(nsocc_.real)/2) + nsocc = int(numpy.round(nsocc_.real) / 2) self.nsocc = nsocc return P_ - - def scf(self, heff=None, fs=False, eri=None, - pert_h=False,pert_list=None, save_chkfile=False, - dm0 = None): + def scf( + self, + heff=None, + fs=False, + eri=None, + pert_h=False, + pert_list=None, + save_chkfile=False, + dm0=None, + ): """ Perform self-consistent field (SCF) calculation for the fragment. @@ -286,22 +317,36 @@ def scf(self, heff=None, fs=False, eri=None, Initial density matrix, by default None. """ import copy - if self._mf is not None: self._mf = None - if self._mc is not None: self._mc = None - if heff is None: heff = self.heff + + if self._mf is not None: + self._mf = None + if self._mc is not None: + self._mc = None + if heff is None: + heff = self.heff if eri is None: eri = get_eri(self.dname, self.nao, eri_file=self.eri_file) if dm0 is None: - dm0 = numpy.dot( self._mo_coeffs[:,:self.nsocc], - self._mo_coeffs[:,:self.nsocc].conj().T) *2. - - mf_ = get_scfObj(self.fock + heff, eri, - self.nsocc, dm0 = dm0, - fname = self.dname, - pert_h=pert_h, pert_list=pert_list, - save_chkfile=save_chkfile) + dm0 = ( + numpy.dot( + self._mo_coeffs[:, : self.nsocc], + self._mo_coeffs[:, : self.nsocc].conj().T, + ) + * 2.0 + ) + + mf_ = get_scfObj( + self.fock + heff, + eri, + self.nsocc, + dm0=dm0, + fname=self.dname, + pert_h=pert_h, + pert_list=pert_list, + save_chkfile=save_chkfile, + ) if pert_h: return mf_ @@ -313,13 +358,18 @@ def scf(self, heff=None, fs=False, eri=None, self._mo_coeffs = mf_.mo_coeff.copy() dm0 = mf_.make_rdm1() - mf_= None - - def update_heff(self,u, cout = None, return_heff=False, - be_iter=None, - no_chempot=False, - tmp_add = False, - only_chem=False): + mf_ = None + + def update_heff( + self, + u, + cout=None, + return_heff=False, + be_iter=None, + no_chempot=False, + tmp_add=False, + only_chem=False, + ): """ Update the effective Hamiltonian for the fragment. """ @@ -332,9 +382,9 @@ def update_heff(self,u, cout = None, return_heff=False, else: cout = cout if not no_chempot: - for i,fi in enumerate(self.fsites): + for i, fi in enumerate(self.fsites): if not any(i in sublist for sublist in self.edge_idx): - heff_[i,i] -= u[-1] + heff_[i, i] -= u[-1] if only_chem: self.heff = heff_ @@ -342,13 +392,13 @@ def update_heff(self,u, cout = None, return_heff=False, if cout is None: return heff_ else: - return(cout, heff_) + return (cout, heff_) return cout - for idx,i in enumerate(self.edge_idx): + for idx, i in enumerate(self.edge_idx): for j in range(len(i)): for k in range(len(i)): - if j>k : + if j > k: continue heff_[i[j], i[k]] = u[cout] @@ -361,74 +411,82 @@ def update_heff(self,u, cout = None, return_heff=False, if cout is None: return heff_ else: - return(cout, heff_) + return (cout, heff_) return cout def set_udim(self, cout): for i in self.edge_idx: for j in range(len(i)): for k in range(len(i)): - if j>k : + if j > k: continue cout += 1 return cout - def energy_hf(self, rdm_hf=None, mo_coeffs = None, eri=None, return_e1=False, unrestricted = False): + def energy_hf( + self, rdm_hf=None, mo_coeffs=None, eri=None, return_e1=False, unrestricted=False + ): if mo_coeffs is None: mo_coeffs = self._mo_coeffs if rdm_hf is None: - rdm_hf = numpy.dot(mo_coeffs[:,:self.nsocc], - mo_coeffs[:,:self.nsocc].conj().T) + rdm_hf = numpy.dot( + mo_coeffs[:, : self.nsocc], mo_coeffs[:, : self.nsocc].conj().T + ) - unrestricted = 1. if unrestricted else 2. + unrestricted = 1.0 if unrestricted else 2.0 - e1 = unrestricted*numpy.einsum("ij,ij->i", self.h1[:self.nfsites], - rdm_hf[:self.nfsites]) + e1 = unrestricted * numpy.einsum( + "ij,ij->i", self.h1[: self.nfsites], rdm_hf[: self.nfsites] + ) - ec = 0.5 * unrestricted * numpy.einsum("ij,ij->i",self.veff[:self.nfsites], - rdm_hf[:self.nfsites]) + ec = ( + 0.5 + * unrestricted + * numpy.einsum( + "ij,ij->i", self.veff[: self.nfsites], rdm_hf[: self.nfsites] + ) + ) if self.TA.ndim == 3: jmax = self.TA[0].shape[1] else: jmax = self.TA.shape[1] if eri is None: - r = h5py.File(self.eri_file,'r') + r = h5py.File(self.eri_file, "r") eri = r[self.dname][()] r.close() - e2 = numpy.zeros_like(e1) for i in range(self.nfsites): for j in range(jmax): - ij = i*(i+1)//2+j if i > j else j*(j+1)//2+i - Gij = (2.*rdm_hf[i,j]*rdm_hf - - numpy.outer(rdm_hf[i], rdm_hf[j]))[:jmax,:jmax] + ij = i * (i + 1) // 2 + j if i > j else j * (j + 1) // 2 + i + Gij = (2.0 * rdm_hf[i, j] * rdm_hf - numpy.outer(rdm_hf[i], rdm_hf[j]))[ + :jmax, :jmax + ] Gij[numpy.diag_indices(jmax)] *= 0.5 Gij += Gij.T e2[i] += 0.5 * unrestricted * Gij[numpy.tril_indices(jmax)] @ eri[ij] - e_ = e1+e2+ec - etmp = 0. - e1_ = 0. - e2_ = 0. - ec_ = 0. + e_ = e1 + e2 + ec + etmp = 0.0 + e1_ = 0.0 + e2_ = 0.0 + ec_ = 0.0 for i in self.efac[1]: - etmp += self.efac[0]*e_[i] - e1_ += self.efac[0]*e1[i] - e2_ += self.efac[0]*e2[i] - ec_ += self.efac[0]*ec[i] + etmp += self.efac[0] * e_[i] + e1_ += self.efac[0] * e1[i] + e2_ += self.efac[0] * e2[i] + ec_ += self.efac[0] * ec[i] self.ebe_hf = etmp if return_e1: - e_h1 = 0. - e_coul = 0. + e_h1 = 0.0 + e_coul = 0.0 for i in self.efac[1]: + e_h1 += self.efac[0] * e1[i] + e_coul += self.efac[0] * (e2[i] + ec[i]) + return (e_h1, e_coul, e1 + e2 + ec) - e_h1 += self.efac[0]*e1[i] - e_coul += self.efac[0]*(e2[i]+ec[i]) - return(e_h1,e_coul, e1+e2+ec) - - return e1+e2+ec + return e1 + e2 + ec diff --git a/kbe/solver.py b/kbe/solver.py index 5eecce61..45a1484e 100644 --- a/kbe/solver.py +++ b/kbe/solver.py @@ -1,6 +1,7 @@ # Author(s): Oinam Romesh Meitei -import numpy,functools,sys, time,os, h5py +import numpy, functools, sys, time, os, h5py + def schmidt_decomp_svd(rdm, Frag_sites): """ @@ -28,18 +29,16 @@ def schmidt_decomp_svd(rdm, Frag_sites): thres = 1.0e-10 Tot_sites = rdm.shape[0] - Fragsites = [i if i>=0 else Tot_sites+i for i in Frag_sites] + Fragsites = [i if i >= 0 else Tot_sites + i for i in Frag_sites] - Env_sites1 = numpy.array([i for i in range(Tot_sites) - if not i in Fragsites]) + Env_sites1 = numpy.array([i for i in range(Tot_sites) if not i in Fragsites]) nfs = len(Frag_sites) Denv = rdm[Env_sites1][:, Fragsites] - U, sigma, V = scipy.linalg.svd(Denv, full_matrices=False, lapack_driver='gesvd') - nbath = ( sigma >= thres).sum() + U, sigma, V = scipy.linalg.svd(Denv, full_matrices=False, lapack_driver="gesvd") + nbath = (sigma >= thres).sum() TA = numpy.zeros((Tot_sites, nfs + nbath), dtype=numpy.complex128) TA[Fragsites, :nfs] = numpy.eye(nfs) - TA[Env_sites1, nfs:] = U[:,:nbath] + TA[Env_sites1, nfs:] = U[:, :nbath] return TA - diff --git a/molbe/_opt.py b/molbe/_opt.py index 8ba0c3f4..9455ddb1 100644 --- a/molbe/_opt.py +++ b/molbe/_opt.py @@ -2,7 +2,7 @@ from .solver import be_func from .be_parallel import be_func_parallel -import scipy,sys,numpy,time, h5py +import scipy, sys, numpy, time, h5py class BEOPT: @@ -40,29 +40,45 @@ class BEOPT: Hartree-Fock energy. Defaults to 0.0 """ - def __init__(self, pot, Fobjs, Nocc, enuc,solver='MP2', ecore=0., - nproc=1,ompnum=4, - only_chem=False, hf_veff = None, - hci_pt=False,hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, - max_space=500, conv_tol = 1.e-6,relax_density = False, - ebe_hf =0., scratch_dir=None, **solver_kwargs): - + def __init__( + self, + pot, + Fobjs, + Nocc, + enuc, + solver="MP2", + ecore=0.0, + nproc=1, + ompnum=4, + only_chem=False, + hf_veff=None, + hci_pt=False, + hci_cutoff=0.001, + ci_coeff_cutoff=None, + select_cutoff=None, + max_space=500, + conv_tol=1.0e-6, + relax_density=False, + ebe_hf=0.0, + scratch_dir=None, + **solver_kwargs, + ): # Initialize class attributes - self.ebe_hf=ebe_hf + self.ebe_hf = ebe_hf self.hf_veff = hf_veff self.pot = pot self.Fobjs = Fobjs self.Nocc = Nocc self.enuc = enuc - self.solver=solver + self.solver = solver self.ecore = ecore self.iter = 0 self.err = 0.0 self.Ebe = 0.0 - self.max_space=max_space + self.max_space = max_space self.nproc = nproc self.ompnum = ompnum - self.only_chem=only_chem + self.only_chem = only_chem self.conv_tol = conv_tol self.relax_density = relax_density # HCI parameters @@ -70,8 +86,8 @@ def __init__(self, pot, Fobjs, Nocc, enuc,solver='MP2', ecore=0., self.ci_coeff_cutoff = ci_coeff_cutoff self.select_cutoff = select_cutoff self.hci_pt = hci_pt - self.solver_kwargs=solver_kwargs - self.scratch_dir=scratch_dir + self.solver_kwargs = solver_kwargs + self.scratch_dir = scratch_dir def objfunc(self, xk): """ @@ -92,33 +108,58 @@ def objfunc(self, xk): # Choose the appropriate function based on the number of processors if self.nproc == 1: - err_, errvec_,ebe_ = be_func(xk, self.Fobjs, self.Nocc, self.solver, self.enuc, - eeval=True, return_vec=True, hf_veff = self.hf_veff, - only_chem=self.only_chem, - hci_cutoff=self.hci_cutoff, - nproc=self.ompnum, relax_density=self.relax_density, - ci_coeff_cutoff = self.ci_coeff_cutoff, - select_cutoff = self.select_cutoff, hci_pt=self.hci_pt, - ecore=self.ecore, ebe_hf=self.ebe_hf, be_iter=self.iter, - scratch_dir=self.scratch_dir, **self.solver_kwargs) + err_, errvec_, ebe_ = be_func( + xk, + self.Fobjs, + self.Nocc, + self.solver, + self.enuc, + eeval=True, + return_vec=True, + hf_veff=self.hf_veff, + only_chem=self.only_chem, + hci_cutoff=self.hci_cutoff, + nproc=self.ompnum, + relax_density=self.relax_density, + ci_coeff_cutoff=self.ci_coeff_cutoff, + select_cutoff=self.select_cutoff, + hci_pt=self.hci_pt, + ecore=self.ecore, + ebe_hf=self.ebe_hf, + be_iter=self.iter, + scratch_dir=self.scratch_dir, + **self.solver_kwargs, + ) else: - err_, errvec_,ebe_ = be_func_parallel(xk, self.Fobjs, self.Nocc, self.solver, self.enuc, - eeval=True, return_vec=True, hf_veff = self.hf_veff, - nproc=self.nproc, ompnum=self.ompnum, - only_chem=self.only_chem, - hci_cutoff=self.hci_cutoff,relax_density=self.relax_density, - ci_coeff_cutoff = self.ci_coeff_cutoff, - select_cutoff = self.select_cutoff, - ecore=self.ecore, ebe_hf=self.ebe_hf, be_iter=self.iter, - scratch_dir=self.scratch_dir, **self.solver_kwargs) + err_, errvec_, ebe_ = be_func_parallel( + xk, + self.Fobjs, + self.Nocc, + self.solver, + self.enuc, + eeval=True, + return_vec=True, + hf_veff=self.hf_veff, + nproc=self.nproc, + ompnum=self.ompnum, + only_chem=self.only_chem, + hci_cutoff=self.hci_cutoff, + relax_density=self.relax_density, + ci_coeff_cutoff=self.ci_coeff_cutoff, + select_cutoff=self.select_cutoff, + ecore=self.ecore, + ebe_hf=self.ebe_hf, + be_iter=self.iter, + scratch_dir=self.scratch_dir, + **self.solver_kwargs, + ) # Update error and BE energy self.err = err_ self.Ebe = ebe_ return errvec_ - - def optimize(self, method, J0 = None, trust_region=False): + def optimize(self, method, J0=None, trust_region=False): """Main kernel to perform BE optimization Parameters @@ -133,59 +174,72 @@ def optimize(self, method, J0 = None, trust_region=False): from molbe.external.optqn import FrankQN import sys - print('-----------------------------------------------------', - flush=True) - print(' Starting BE optimization ', flush=True) - print(' Solver : ',self.solver,flush=True) + print("-----------------------------------------------------", flush=True) + print(" Starting BE optimization ", flush=True) + print(" Solver : ", self.solver, flush=True) if self.only_chem: - print(' Chemical Potential Optimization', flush=True) - print('-----------------------------------------------------', - flush=True) + print(" Chemical Potential Optimization", flush=True) + print("-----------------------------------------------------", flush=True) print(flush=True) - if method=='QN': - - print('-- In iter ',self.iter, flush=True) + if method == "QN": + print("-- In iter ", self.iter, flush=True) # Initial step f0 = self.objfunc(self.pot) - print('Error in density matching : {:>2.4e}'.format(self.err), flush=True) + print( + "Error in density matching : {:>2.4e}".format(self.err), + flush=True, + ) print(flush=True) # Initialize the Quasi-Newton optimizer - optQN = FrankQN(self.objfunc, numpy.array(self.pot), - f0, J0, - max_space=self.max_space) + optQN = FrankQN( + self.objfunc, numpy.array(self.pot), f0, J0, max_space=self.max_space + ) if self.err < self.conv_tol: print(flush=True) - print('CONVERGED w/o Optimization Steps',flush=True) + print("CONVERGED w/o Optimization Steps", flush=True) print(flush=True) else: # Perform optimization steps for iter_ in range(self.max_space): - print('-- In iter ',self.iter, flush=True) + print("-- In iter ", self.iter, flush=True) optQN.next_step(trust_region=trust_region) self.iter += 1 - print('Error in density matching : {:>2.4e}'.format(self.err), flush=True) + print( + "Error in density matching : {:>2.4e}".format(self.err), + flush=True, + ) print(flush=True) if self.err < self.conv_tol: print(flush=True) - print('CONVERGED',flush=True) + print("CONVERGED", flush=True) print(flush=True) break else: - print('This optimization method for BE is not supported') + print("This optimization method for BE is not supported") sys.exit() - - -def optimize(self, solver='MP2',method='QN', - only_chem=False, conv_tol = 1.e-6, relax_density=False, use_cumulant=True, - J0=None, nproc=1, ompnum=4, max_iter=500, scratch_dir=None, trust_region=False, - **solver_kwargs): +def optimize( + self, + solver="MP2", + method="QN", + only_chem=False, + conv_tol=1.0e-6, + relax_density=False, + use_cumulant=True, + J0=None, + nproc=1, + ompnum=4, + max_iter=500, + scratch_dir=None, + trust_region=False, + **solver_kwargs, +): """BE optimization function Interfaces BEOPT to perform bootstrap embedding optimization. @@ -223,37 +277,53 @@ def optimize(self, solver='MP2',method='QN', # Check if only chemical potential optimization is required if not only_chem: pot = self.pot - if self.be_type=='be1': - sys.exit('BE1 only works with chemical potential optimization. Set only_chem=True') + if self.be_type == "be1": + sys.exit( + "BE1 only works with chemical potential optimization. Set only_chem=True" + ) else: - pot = [0.] + pot = [0.0] # Initialize the BEOPT object - be_ = BEOPT(pot, self.Fobjs, self.Nocc, self.enuc, hf_veff = self.hf_veff, - nproc=nproc, ompnum=ompnum, scratch_dir=scratch_dir, - max_space=max_iter,conv_tol = conv_tol, - only_chem=only_chem, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff = self.ci_coeff_cutoff,relax_density=relax_density, - select_cutoff = self.select_cutoff,hci_pt=self.hci_pt, - solver=solver, ecore=self.E_core, ebe_hf=self.ebe_hf, - **solver_kwargs) - - if method=='QN': + be_ = BEOPT( + pot, + self.Fobjs, + self.Nocc, + self.enuc, + hf_veff=self.hf_veff, + nproc=nproc, + ompnum=ompnum, + scratch_dir=scratch_dir, + max_space=max_iter, + conv_tol=conv_tol, + only_chem=only_chem, + hci_cutoff=self.hci_cutoff, + ci_coeff_cutoff=self.ci_coeff_cutoff, + relax_density=relax_density, + select_cutoff=self.select_cutoff, + hci_pt=self.hci_pt, + solver=solver, + ecore=self.E_core, + ebe_hf=self.ebe_hf, + **solver_kwargs, + ) + + if method == "QN": # Prepare the initial Jacobian matrix if only_chem: - J0 = [[0.]] - J0 = self.get_be_error_jacobian(jac_solver='HF') - J0 = [[J0[-1,-1]]] + J0 = [[0.0]] + J0 = self.get_be_error_jacobian(jac_solver="HF") + J0 = [[J0[-1, -1]]] else: - J0 = self.get_be_error_jacobian(jac_solver='HF') + J0 = self.get_be_error_jacobian(jac_solver="HF") # Perform the optimization be_.optimize(method, J0=J0, trust_region=trust_region) self.ebe_tot = self.ebe_hf + be_.Ebe[0] # Print the energy components - print_energy(be_.Ebe[0], be_.Ebe[1][1], be_.Ebe[1][0]+be_.Ebe[1][2], self.ebe_hf) + print_energy( + be_.Ebe[0], be_.Ebe[1][1], be_.Ebe[1][0] + be_.Ebe[1][2], self.ebe_hf + ) else: - print('This optimization method for BE is not supported') + print("This optimization method for BE is not supported") sys.exit() - diff --git a/molbe/autofrag.py b/molbe/autofrag.py index 28c9f815..4b007e6e 100644 --- a/molbe/autofrag.py +++ b/molbe/autofrag.py @@ -4,10 +4,16 @@ import numpy from .helper import get_core -def autogen(mol, frozen_core=True, be_type='be2', - write_geom=False, - valence_basis = None, valence_only = False, - print_frags=True): + +def autogen( + mol, + frozen_core=True, + be_type="be2", + write_geom=False, + valence_basis=None, + valence_only=False, + print_frags=True, +): """ Automatic molecular partitioning @@ -89,7 +95,7 @@ def autogen(mol, frozen_core=True, be_type='be2', # Check if the molecule is a hydrogen chain hchain = True for i in range(cell.natm): - if not cell.atom_pure_symbol(i) == 'H': + if not cell.atom_pure_symbol(i) == "H": hchain = False break @@ -98,7 +104,7 @@ def autogen(mol, frozen_core=True, be_type='be2', # Assumes that there can be only 5 member connected system for idx, i in enumerate(normlist): - if cell.atom_pure_symbol(idx) == 'H' and not hchain: + if cell.atom_pure_symbol(idx) == "H" and not hchain: continue tmplist = normlist - i @@ -106,39 +112,44 @@ def autogen(mol, frozen_core=True, be_type='be2', clist = [] cout = 0 - for jdx,j in enumerate(tmplist): - if not idx==jdx and (not cell.atom_pure_symbol(jdx) == 'H' or hchain): - if abs(j)< normdist: + for jdx, j in enumerate(tmplist): + if not idx == jdx and (not cell.atom_pure_symbol(jdx) == "H" or hchain): + if abs(j) < normdist: clist.append(jdx) pedg = [] flist = [] flist.append(idx) - if not be_type == 'be1': + if not be_type == "be1": for jdx in clist: dist = numpy.linalg.norm(coord[idx] - coord[jdx]) if dist <= bond: - flist.append(jdx) - pedg.append(jdx) - if be_type=='be3' or be_type == 'be4': - - for kdx in clist: - if not kdx == jdx: - dist = numpy.linalg.norm(coord[jdx] - coord[kdx]) - if dist <= bond: - if not kdx in pedg: - flist.append(kdx) - pedg.append(kdx) - if be_type=='be4': - for ldx, l in enumerate(coord): - if ldx==kdx or ldx==jdx or\ - (cell.atom_pure_symbol(ldx) == 'H' and not hchain)\ - or ldx in pedg: - continue - dist = numpy.linalg.norm(coord[kdx] - l) - if dist <= bond: - flist.append(ldx) - pedg.append(ldx) + flist.append(jdx) + pedg.append(jdx) + if be_type == "be3" or be_type == "be4": + for kdx in clist: + if not kdx == jdx: + dist = numpy.linalg.norm(coord[jdx] - coord[kdx]) + if dist <= bond: + if not kdx in pedg: + flist.append(kdx) + pedg.append(kdx) + if be_type == "be4": + for ldx, l in enumerate(coord): + if ( + ldx == kdx + or ldx == jdx + or ( + cell.atom_pure_symbol(ldx) == "H" + and not hchain + ) + or ldx in pedg + ): + continue + dist = numpy.linalg.norm(coord[kdx] - l) + if dist <= bond: + flist.append(ldx) + pedg.append(ldx) # Update fragment and edge lists based on current partitioning for pidx, frag_ in enumerate(Frag): @@ -147,8 +158,10 @@ def autogen(mol, frozen_core=True, be_type='be2', open_frag_cen.append(idx) break elif set(frag_).issubset(flist): - open_frag = [ oidx-1 if oidx > pidx else oidx for oidx in open_frag] - open_frag.append(len(Frag)-1) + open_frag = [ + oidx - 1 if oidx > pidx else oidx for oidx in open_frag + ] + open_frag.append(len(Frag) - 1) open_frag_cen.append(cen[pidx]) del cen[pidx] del Frag[pidx] @@ -164,13 +177,13 @@ def autogen(mol, frozen_core=True, be_type='be2', hlist = [[] for i in coord] if not hchain: for idx, i in enumerate(normlist): - if cell.atom_pure_symbol(idx) == 'H': + if cell.atom_pure_symbol(idx) == "H": tmplist = normlist - i tmplist = list(tmplist) clist = [] - for jdx,j in enumerate(tmplist): - if not idx==jdx and not cell.atom_pure_symbol(jdx) == 'H': - if abs(j)< normdist: + for jdx, j in enumerate(tmplist): + if not idx == jdx and not cell.atom_pure_symbol(jdx) == "H": + if abs(j) < normdist: clist.append(jdx) for jdx in clist: dist = numpy.linalg.norm(coord[idx] - coord[jdx]) @@ -180,47 +193,78 @@ def autogen(mol, frozen_core=True, be_type='be2', # Print fragments if requested if print_frags: print(flush=True) - print('Fragment sites',flush=True) - print('--------------------------',flush=True) - print('Fragment | Center | Edges ',flush=True) - print('--------------------------',flush=True) - - for idx,i in enumerate(Frag): - print(' {:>4} | {:>5} |'.format(idx, cell.atom_pure_symbol(cen[idx])+str(cen[idx]+1)),end=' ', flush=True) + print("Fragment sites", flush=True) + print("--------------------------", flush=True) + print("Fragment | Center | Edges ", flush=True) + print("--------------------------", flush=True) + + for idx, i in enumerate(Frag): + print( + " {:>4} | {:>5} |".format( + idx, cell.atom_pure_symbol(cen[idx]) + str(cen[idx] + 1) + ), + end=" ", + flush=True, + ) for j in hlist[cen[idx]]: - print(' {:>5} '.format('*'+cell.atom_pure_symbol(j)+str(j+1)),end=' ', flush=True) + print( + " {:>5} ".format("*" + cell.atom_pure_symbol(j) + str(j + 1)), + end=" ", + flush=True, + ) for j in i: - if j == cen[idx]: continue - print(' {:>5} '.format(cell.atom_pure_symbol(j)+str(j+1)),end=' ', flush=True) + if j == cen[idx]: + continue + print( + " {:>5} ".format(cell.atom_pure_symbol(j) + str(j + 1)), + end=" ", + flush=True, + ) for k in hlist[j]: - print(' {:>5} '.format(cell.atom_pure_symbol(k)+str(k+1)),end=' ', flush=True) + print( + " {:>5} ".format(cell.atom_pure_symbol(k) + str(k + 1)), + end=" ", + flush=True, + ) print(flush=True) - print('--------------------------',flush=True) - print(' No. of fragments : ',len(Frag),flush=True) - print('*H : Center H atoms (printed as Edges above.)', flush=True) + print("--------------------------", flush=True) + print(" No. of fragments : ", len(Frag), flush=True) + print("*H : Center H atoms (printed as Edges above.)", flush=True) print(flush=True) # Write fragment geometry to a file if requested if write_geom: - w = open('fragments.xyz','w') - for idx,i in enumerate(Frag): - w.write(str(len(i)+len(hlist[cen[idx]])+len(hlist[j]))+'\n') - w.write('Fragment - '+str(idx)+'\n') + w = open("fragments.xyz", "w") + for idx, i in enumerate(Frag): + w.write(str(len(i) + len(hlist[cen[idx]]) + len(hlist[j])) + "\n") + w.write("Fragment - " + str(idx) + "\n") for j in hlist[cen[idx]]: - w.write(' {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n'.format(cell.atom_pure_symbol(j), - coord[j][0]/ang2bohr, - coord[j][1]/ang2bohr, - coord[j][2]/ang2bohr)) + w.write( + " {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n".format( + cell.atom_pure_symbol(j), + coord[j][0] / ang2bohr, + coord[j][1] / ang2bohr, + coord[j][2] / ang2bohr, + ) + ) for j in i: - w.write(' {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n'.format(cell.atom_pure_symbol(j), - coord[j][0]/ang2bohr, - coord[j][1]/ang2bohr, - coord[j][2]/ang2bohr)) + w.write( + " {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n".format( + cell.atom_pure_symbol(j), + coord[j][0] / ang2bohr, + coord[j][1] / ang2bohr, + coord[j][2] / ang2bohr, + ) + ) for k in hlist[j]: - w.write(' {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n'.format(cell.atom_pure_symbol(k), - coord[k][0]/ang2bohr, - coord[k][1]/ang2bohr, - coord[k][2]/ang2bohr)) + w.write( + " {:>3} {:>10.7f} {:>10.7f} {:>10.7f} \n".format( + cell.atom_pure_symbol(k), + coord[k][0] / ang2bohr, + coord[k][1] / ang2bohr, + coord[k][2] / ang2bohr, + ) + ) w.close() # Prepare for PAO basis if requested @@ -252,7 +296,7 @@ def autogen(mol, frozen_core=True, be_type='be2', sites__[adx] = b1list continue - if not cell.atom_pure_symbol(adx) == 'H' and not hchain: + if not cell.atom_pure_symbol(adx) == "H" and not hchain: bas = baslist[adx] start_ = bas[2] stop_ = bas[3] @@ -263,8 +307,9 @@ def autogen(mol, frozen_core=True, be_type='be2', if frozen_core: start_ -= coreshift ncore_ = core_list[adx] - stop_ -= coreshift+ncore_ - if pao: nbas2[adx] -= ncore_ + stop_ -= coreshift + ncore_ + if pao: + nbas2[adx] -= ncore_ coreshift += ncore_ b1list = [i for i in range(start_, stop_)] @@ -308,32 +353,30 @@ def autogen(mol, frozen_core=True, be_type='be2', frglist = sites__[cen[idx]].copy() frglist.extend(hsites[cen[idx]]) - ls = len(sites__[cen[idx]])+len(hsites[cen[idx]]) + ls = len(sites__[cen[idx]]) + len(hsites[cen[idx]]) if idx in open_frag: - for pidx__,pid__ in enumerate(open_frag): + for pidx__, pid__ in enumerate(open_frag): if idx == pid__: - frglist.extend(sites__[open_frag_cen[pidx__]]) frglist.extend(hsites[open_frag_cen[pidx__]]) - ls += len(sites__[open_frag_cen[pidx__]]) +\ - len(hsites[open_frag_cen[pidx__]]) + ls += len(sites__[open_frag_cen[pidx__]]) + len( + hsites[open_frag_cen[pidx__]] + ) ftmp.extend(frglist) if not pao: - ls_ = len(sites__[cen[idx]])+len(hsites[cen[idx]]) - centerf_idx.append([pq for pq in range(indix,indix+ls_)]) + ls_ = len(sites__[cen[idx]]) + len(hsites[cen[idx]]) + centerf_idx.append([pq for pq in range(indix, indix + ls_)]) else: - cntlist = sites__[cen[idx]].copy()[:nbas2[cen[idx]]] - cntlist.extend(hsites[cen[idx]][:nbas2H[cen[idx]]]) - ind__ = [ indix+frglist.index(pq) for pq in cntlist] + cntlist = sites__[cen[idx]].copy()[: nbas2[cen[idx]]] + cntlist.extend(hsites[cen[idx]][: nbas2H[cen[idx]]]) + ind__ = [indix + frglist.index(pq) for pq in cntlist] centerf_idx.append(ind__) indix += ls - if not be_type=='be1': + if not be_type == "be1": for jdx in pedge[idx]: - if idx in open_frag: - if jdx == open_frag_cen[open_frag.index(idx)]: continue if jdx in open_frag_cen: @@ -348,13 +391,13 @@ def autogen(mol, frozen_core=True, be_type='be2', edglist = sites__[jdx].copy() edglist.extend(hsites[jdx]) ftmpe.append(edglist) - edind.append([pq for pq in range(indix,indix+ls)]) + edind.append([pq for pq in range(indix, indix + ls)]) else: - edglist = sites__[jdx][:nbas2[jdx]].copy() - edglist.extend(hsites[jdx][:nbas2H[jdx]]) + edglist = sites__[jdx][: nbas2[jdx]].copy() + edglist.extend(hsites[jdx][: nbas2H[jdx]]) ftmpe.append(edglist) - ind__ = [ indix+frglist.index(pq) for pq in edglist] + ind__ = [indix + frglist.index(pq) for pq in edglist] edind.append(ind__) indix += ls edge.append(edg) @@ -370,21 +413,21 @@ def autogen(mol, frozen_core=True, be_type='be2', elif jx in open_frag_cen: cen_.append(open_frag[open_frag_cen.index(jx)]) else: - print(' This is more complicated than I can handle') + print(" This is more complicated than I can handle") sys.exit() center.append(cen_) Nfrag = len(fsites) - add_centers=[[] for x in range(Nfrag)] # additional centers for mixed-basis - ebe_weight=[] + add_centers = [[] for x in range(Nfrag)] # additional centers for mixed-basis + ebe_weight = [] # Compute weights for each fragment for ix, i in enumerate(fsites): tmp_ = [i.index(pq) for pq in sites__[cen[ix]]] tmp_.extend([i.index(pq) for pq in hsites[cen[ix]]]) if ix in open_frag: - for pidx__,pid__ in enumerate(open_frag): + for pidx__, pid__ in enumerate(open_frag): if ix == pid__: add_centers[pid__].append(open_frag_cen[pidx__]) tmp_.extend([i.index(pq) for pq in sites__[open_frag_cen[pidx__]]]) @@ -392,10 +435,10 @@ def autogen(mol, frozen_core=True, be_type='be2', ebe_weight.append([1.0, tmp_]) center_idx = [] - if not be_type=='be1': + if not be_type == "be1": for i in range(Nfrag): idx = [] - for jdx,j in enumerate(center[i]): + for jdx, j in enumerate(center[i]): jdx_continue = False if j in open_frag: for kdx, k in enumerate(open_frag): @@ -406,25 +449,39 @@ def autogen(mol, frozen_core=True, be_type='be2', cntlist.extend(hsites[open_frag_cen[kdx]]) idx.append([fsites[j].index(k) for k in cntlist]) else: - cntlist = sites__[open_frag_cen[kdx]].copy()[:nbas2[cen[j]]] - cntlist.extend(hsites[open_frag_cen[kdx]][:nbas2H[cen[j]]]) + cntlist = sites__[open_frag_cen[kdx]].copy()[ + : nbas2[cen[j]] + ] + cntlist.extend( + hsites[open_frag_cen[kdx]][: nbas2H[cen[j]]] + ) idx.append([fsites[j].index(k) for k in cntlist]) jdx_continue = True break - if jdx_continue: continue + if jdx_continue: + continue if not pao: cntlist = sites__[cen[j]].copy() cntlist.extend(hsites[cen[j]]) idx.append([fsites[j].index(k) for k in cntlist]) else: - cntlist = sites__[cen[j]].copy()[:nbas2[cen[j]]] - cntlist.extend(hsites[cen[j]][:nbas2H[cen[j]]]) + cntlist = sites__[cen[j]].copy()[: nbas2[cen[j]]] + cntlist.extend(hsites[cen[j]][: nbas2H[cen[j]]]) idx.append([fsites[j].index(k) for k in cntlist]) center_idx.append(idx) - return(fsites, edgsites, center, edge_idx, center_idx, centerf_idx, ebe_weight, Frag, cen, hlist, add_centers) - - - + return ( + fsites, + edgsites, + center, + edge_idx, + center_idx, + centerf_idx, + ebe_weight, + Frag, + cen, + hlist, + add_centers, + ) diff --git a/molbe/be_parallel.py b/molbe/be_parallel.py index dae60dc5..93d9dfbf 100644 --- a/molbe/be_parallel.py +++ b/molbe/be_parallel.py @@ -1,7 +1,7 @@ # Author(s): Oinam Romesh Meitei, Leah Weisburn from .solver import solve_error -from .solver import solve_mp2, solve_ccsd,make_rdm1_ccsd_t1,solve_uccsd +from .solver import solve_mp2, solve_ccsd, make_rdm1_ccsd_t1, solve_uccsd from .solver import make_rdm2_urlx from .helper import get_frag_energy from molbe.external.unrestricted_utils import make_uhf_obj @@ -9,12 +9,32 @@ import functools, numpy, sys from .helper import * -def run_solver(h1, dm0, dname, nao, nocc, nfsites, - efac, TA, hf_veff, h1_e, - solver='MP2',eri_file='eri_file.h5', veff0=None, - hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, - ompnum=4, writeh1=False, - eeval=True, return_rdm_ao=True, use_cumulant=True, relax_density=False, frag_energy=False): + +def run_solver( + h1, + dm0, + dname, + nao, + nocc, + nfsites, + efac, + TA, + hf_veff, + h1_e, + solver="MP2", + eri_file="eri_file.h5", + veff0=None, + hci_cutoff=0.001, + ci_coeff_cutoff=None, + select_cutoff=None, + ompnum=4, + writeh1=False, + eeval=True, + return_rdm_ao=True, + use_cumulant=True, + relax_density=False, + frag_energy=False, +): """ Run a quantum chemistry solver to compute the reduced density matrices. @@ -76,34 +96,36 @@ def run_solver(h1, dm0, dname, nao, nocc, nfsites, rdm_return = True # Select solver - if solver=='MP2': + if solver == "MP2": mc_ = solve_mp2(mf_, mo_energy=mf_.mo_energy) rdm1_tmp = mc_.make_rdm1() - elif solver=='CCSD': + elif solver == "CCSD": if not rdm_return: - t1, t2 = solve_ccsd(mf_, - mo_energy=mf_.mo_energy, - rdm_return=False) + t1, t2 = solve_ccsd(mf_, mo_energy=mf_.mo_energy, rdm_return=False) rdm1_tmp = make_rdm1_ccsd_t1(t1) else: - t1, t2, rdm1_tmp, rdm2s = solve_ccsd(mf_, - mo_energy=mf_.mo_energy, - rdm_return=True, - rdm2_return = True, use_cumulant=use_cumulant, - relax=True) - elif solver=='FCI': + t1, t2, rdm1_tmp, rdm2s = solve_ccsd( + mf_, + mo_energy=mf_.mo_energy, + rdm_return=True, + rdm2_return=True, + use_cumulant=use_cumulant, + relax=True, + ) + elif solver == "FCI": from pyscf import fci mc_ = fci.FCI(mf_, mf_.mo_coeff) efci, civec = mc_.kernel() rdm1_tmp = mc_.make_rdm1(civec, mc_.norb, mc_.nelec) - elif solver=='HCI': - from pyscf import hci,ao2mo + elif solver == "HCI": + from pyscf import hci, ao2mo nao, nmo = mf_.mo_coeff.shape - eri = ao2mo.kernel(mf_._eri, mf_.mo_coeff, aosym='s4', - compact=False).reshape(4*((nmo),)) + eri = ao2mo.kernel(mf_._eri, mf_.mo_coeff, aosym="s4", compact=False).reshape( + 4 * ((nmo),) + ) ci_ = hci.SCI(mf_.mol) if select_cutoff is None and ci_coeff_cutoff is None: select_cutoff = hci_cutoff @@ -116,96 +138,123 @@ def run_solver(h1, dm0, dname, nao, nocc, nfsites, nelec = (nocc, nocc) h1_ = functools.reduce(numpy.dot, (mf_.mo_coeff.T, h1, mf_.mo_coeff)) - eci, civec = ci_.kernel(h1_, eri, nmo, nelec) + eci, civec = ci_.kernel(h1_, eri, nmo, nelec) civec = numpy.asarray(civec) (rdm1a_, rdm1b_), (rdm2aa, rdm2ab, rdm2bb) = ci_.make_rdm12s(civec, nmo, nelec) rdm1_tmp = rdm1a_ + rdm1b_ - rdm2s = rdm2aa + rdm2ab + rdm2ab.transpose(2,3,0,1) + rdm2bb + rdm2s = rdm2aa + rdm2ab + rdm2ab.transpose(2, 3, 0, 1) + rdm2bb - elif solver=='SHCI': + elif solver == "SHCI": from pyscf.shciscf import shci nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) mch = shci.SHCISCF(mf_, nmo, nelec, orbpath=dname) - mch.fcisolver.mpiprefix = 'mpirun -np '+str(ompnum) - mch.fcisolver.stochastic = True # this is for PT and doesnt add PT to rdm + mch.fcisolver.mpiprefix = "mpirun -np " + str(ompnum) + mch.fcisolver.stochastic = True # this is for PT and doesnt add PT to rdm mch.fcisolver.nPTiter = 0 mch.fcisolver.sweep_iter = [0] mch.fcisolver.DoRDM = True - mch.fcisolver.sweep_epsilon = [ hci_cutoff ] - mch.fcisolver.scratchDirectory='/scratch/oimeitei/'+dname + mch.fcisolver.sweep_epsilon = [hci_cutoff] + mch.fcisolver.scratchDirectory = "/scratch/oimeitei/" + dname if not writeh1: - mch.fcisolver.restart=True + mch.fcisolver.restart = True mch.mc1step() rdm1_tmp, rdm2s = mch.fcisolver.make_rdm12(0, nmo, nelec) - elif solver == 'SCI': + elif solver == "SCI": from pyscf import cornell_shci from pyscf import ao2mo, mcscf nao, nmo = mf_.mo_coeff.shape nelec = (nocc, nocc) - cas = mcscf.CASCI (mf_, nmo, nelec) + cas = mcscf.CASCI(mf_, nmo, nelec) h1, ecore = cas.get_h1eff(mo_coeff=mf_.mo_coeff) - eri = ao2mo.kernel(mf_._eri, mf_.mo_coeff, aosym='s4', compact=False).reshape(4*((nmo),)) + eri = ao2mo.kernel(mf_._eri, mf_.mo_coeff, aosym="s4", compact=False).reshape( + 4 * ((nmo),) + ) ci = cornell_shci.SHCI() - ci.runtimedir=dname - ci.restart=True - ci.config['var_only'] = True - ci.config['eps_vars'] = [hci_cutoff] - ci.config['get_1rdm_csv'] = True - ci.config['get_2rdm_csv'] = True + ci.runtimedir = dname + ci.restart = True + ci.config["var_only"] = True + ci.config["eps_vars"] = [hci_cutoff] + ci.config["get_1rdm_csv"] = True + ci.config["get_2rdm_csv"] = True ci.kernel(h1, eri, nmo, nelec) - rdm1_tmp, rdm2s = ci.make_rdm12(0,nmo,nelec) + rdm1_tmp, rdm2s = ci.make_rdm12(0, nmo, nelec) else: - print('Solver not implemented',flush=True) - print('exiting',flush=True) + print("Solver not implemented", flush=True) + print("exiting", flush=True) sys.exit() # Compute RDM1 - rdm1 = functools.reduce(numpy.dot, - (mf_.mo_coeff, - rdm1_tmp, - mf_.mo_coeff.T))*0.5 + rdm1 = functools.reduce(numpy.dot, (mf_.mo_coeff, rdm1_tmp, mf_.mo_coeff.T)) * 0.5 if eeval: - if solver =='CCSD' and not rdm_return: + if solver == "CCSD" and not rdm_return: with_dm1 = True - if use_cumulant: with_dm1 = False - rdm2s = make_rdm2_urlx(t1, t2, with_dm1 = with_dm1) + if use_cumulant: + with_dm1 = False + rdm2s = make_rdm2_urlx(t1, t2, with_dm1=with_dm1) - elif solver == 'MP2': + elif solver == "MP2": rdm2s = mc_.make_rdm2() - elif solver == 'FCI': + elif solver == "FCI": rdm2s = mc_.make_rdm2(civec, mc_.norb, mc_.nelec) if use_cumulant: hf_dm = numpy.zeros_like(rdm1_tmp) - hf_dm[numpy.diag_indices(nocc)] += 2. + hf_dm[numpy.diag_indices(nocc)] += 2.0 del_rdm1 = rdm1_tmp.copy() - del_rdm1[numpy.diag_indices(nocc)] -= 2. - nc = numpy.einsum('ij,kl->ijkl',hf_dm, hf_dm) + \ - numpy.einsum('ij,kl->ijkl',hf_dm, del_rdm1) + \ - numpy.einsum('ij,kl->ijkl',del_rdm1, hf_dm) - nc -= (numpy.einsum('ij,kl->iklj',hf_dm, hf_dm) + \ - numpy.einsum('ij,kl->iklj',hf_dm, del_rdm1) + \ - numpy.einsum('ij,kl->iklj',del_rdm1, hf_dm))*0.5 + del_rdm1[numpy.diag_indices(nocc)] -= 2.0 + nc = ( + numpy.einsum("ij,kl->ijkl", hf_dm, hf_dm) + + numpy.einsum("ij,kl->ijkl", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->ijkl", del_rdm1, hf_dm) + ) + nc -= ( + numpy.einsum("ij,kl->iklj", hf_dm, hf_dm) + + numpy.einsum("ij,kl->iklj", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->iklj", del_rdm1, hf_dm) + ) * 0.5 rdm2s -= nc - e_f = get_frag_energy(mf_.mo_coeff, nocc, nfsites, efac, TA, h1_e, hf_veff, rdm1_tmp, rdm2s, dname, eri_file, veff0) + e_f = get_frag_energy( + mf_.mo_coeff, + nocc, + nfsites, + efac, + TA, + h1_e, + hf_veff, + rdm1_tmp, + rdm2s, + dname, + eri_file, + veff0, + ) if frag_energy: return e_f if return_rdm_ao: - return(e_f, mf_.mo_coeff, rdm1, rdm2s, rdm1_tmp) + return (e_f, mf_.mo_coeff, rdm1, rdm2s, rdm1_tmp) return (e_f, mf_.mo_coeff, rdm1, rdm2s) -def run_solver_u(fobj_a, fobj_b, solver, enuc, hf_veff, - frag_energy=True, relax_density=False, frozen=False, - eri_file='eri_file.h5', use_cumulant=True, ereturn=True): +def run_solver_u( + fobj_a, + fobj_b, + solver, + enuc, + hf_veff, + frag_energy=True, + relax_density=False, + frozen=False, + eri_file="eri_file.h5", + use_cumulant=True, + ereturn=True, +): """ Run a quantum chemistry solver to compute the reduced density matrices. @@ -251,36 +300,47 @@ def run_solver_u(fobj_a, fobj_b, solver, enuc, hf_veff, if relax_density: rdm_return = True - if solver=='UCCSD': + if solver == "UCCSD": if rdm_return: - ucc, rdm1_tmp, rdm2s = solve_uccsd(full_uhf, eris, relax=relax_density, - rdm_return=True, rdm2_return=True, - frozen=frozen) + ucc, rdm1_tmp, rdm2s = solve_uccsd( + full_uhf, + eris, + relax=relax_density, + rdm_return=True, + rdm2_return=True, + frozen=frozen, + ) else: - ucc = solve_uccsd(full_uhf, eris, relax = relax_density, rdm_return=False, - frozen=frozen) - rdm1_tmp = make_rdm1_uccsd(ucc, relax=relax_density) + ucc = solve_uccsd( + full_uhf, eris, relax=relax_density, rdm_return=False, frozen=frozen + ) + rdm1_tmp = make_rdm1_uccsd(ucc, relax=relax_density) else: raise NotImplementedError("Only UCCSD Solver implemented") # Compute RDM1 fobj_a.__rdm1 = rdm1_tmp[0].copy() - fobj_a._rdm1 = functools.reduce(numpy.dot, - (fobj_a._mf.mo_coeff, - rdm1_tmp[0], - fobj_a._mf.mo_coeff.T))*0.5 + fobj_a._rdm1 = ( + functools.reduce( + numpy.dot, (fobj_a._mf.mo_coeff, rdm1_tmp[0], fobj_a._mf.mo_coeff.T) + ) + * 0.5 + ) fobj_b.__rdm1 = rdm1_tmp[1].copy() - fobj_b._rdm1 = functools.reduce(numpy.dot, - (fobj_b._mf.mo_coeff, - rdm1_tmp[1], - fobj_b._mf.mo_coeff.T))*0.5 + fobj_b._rdm1 = ( + functools.reduce( + numpy.dot, (fobj_b._mf.mo_coeff, rdm1_tmp[1], fobj_b._mf.mo_coeff.T) + ) + * 0.5 + ) # Calculate Energies if ereturn: - if solver =='UCCSD' and not rdm_return: + if solver == "UCCSD" and not rdm_return: with_dm1 = True - if use_cumulant: with_dm1=False + if use_cumulant: + with_dm1 = False rdm2s = make_rdm2_uccsd(ucc, with_dm1=with_dm1) else: raise NotImplementedError("RDM Return not Implemented") @@ -291,35 +351,56 @@ def run_solver_u(fobj_a, fobj_b, solver, enuc, hf_veff, # Calculate energy on a per-fragment basis if frag_energy: if frozen: - h1_ab = [full_uhf.h1[0]+full_uhf.full_gcore[0]+full_uhf.core_veffs[0], - full_uhf.h1[1]+full_uhf.full_gcore[1]+full_uhf.core_veffs[1]] + h1_ab = [ + full_uhf.h1[0] + full_uhf.full_gcore[0] + full_uhf.core_veffs[0], + full_uhf.h1[1] + full_uhf.full_gcore[1] + full_uhf.core_veffs[1], + ] else: h1_ab = [fobj_a.h1, fobj_b.h1] - e_f = get_frag_energy_u((fobj_a._mo_coeffs,fobj_b._mo_coeffs), - (fobj_a.nsocc,fobj_b.nsocc), - (fobj_a.nfsites, fobj_b.nfsites), - (fobj_a.efac, fobj_b.efac), - (fobj_a.TA, fobj_b.TA), - h1_ab, - hf_veff, - rdm1_tmp, - rdm2s, - fobj_a.dname, - eri_file=fobj_a.eri_file, - gcores=full_uhf.full_gcore, - frozen=frozen - ) + e_f = get_frag_energy_u( + (fobj_a._mo_coeffs, fobj_b._mo_coeffs), + (fobj_a.nsocc, fobj_b.nsocc), + (fobj_a.nfsites, fobj_b.nfsites), + (fobj_a.efac, fobj_b.efac), + (fobj_a.TA, fobj_b.TA), + h1_ab, + hf_veff, + rdm1_tmp, + rdm2s, + fobj_a.dname, + eri_file=fobj_a.eri_file, + gcores=full_uhf.full_gcore, + frozen=frozen, + ) return e_f else: return NotImplementedError("Energy only calculated on a per-fragment basis") -def be_func_parallel(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, - nproc=1, ompnum=4, - only_chem=False,relax_density=False,use_cumulant=True, - eeval=False, ereturn=False, frag_energy=False, - hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, - return_vec=False, ecore=0., ebe_hf=0., be_iter=None, writeh1=False): +def be_func_parallel( + pot, + Fobjs, + Nocc, + solver, + enuc, + hf_veff=None, + nproc=1, + ompnum=4, + only_chem=False, + relax_density=False, + use_cumulant=True, + eeval=False, + ereturn=False, + frag_energy=False, + hci_cutoff=0.001, + ci_coeff_cutoff=None, + select_cutoff=None, + return_vec=False, + ecore=0.0, + ebe_hf=0.0, + be_iter=None, + writeh1=False, +): """ Embarrassingly Parallel High-Level Computation @@ -376,14 +457,14 @@ def be_func_parallel(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, nfrag = len(Fobjs) # Create directories for fragments if required - if writeh1 and solver=='SCI': + if writeh1 and solver == "SCI": for nf in range(nfrag): dname = Fobjs[nf].dname - os.system('mkdir '+dname) + os.system("mkdir " + dname) # Set the number of OpenMP threads - os.system('export OMP_NUM_THREADS='+str(ompnum)) - nprocs = int(nproc/ompnum) + os.system("export OMP_NUM_THREADS=" + str(ompnum)) + nprocs = int(nproc / ompnum) # Update the effective Hamiltonian with potentials if not pot is None: @@ -407,11 +488,34 @@ def be_func_parallel(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, h1_e = Fobjs[nf].h1 veff0 = Fobjs[nf].veff0 - result = pool_.apply_async(run_solver, [h1, dm0, dname, nao, nocc, nfsites, - efac, TA, hf_veff, h1_e, - solver,Fobjs[nf].eri_file, veff0, - hci_cutoff, ci_coeff_cutoff,select_cutoff, - ompnum, writeh1, True, True, use_cumulant, relax_density, frag_energy]) + result = pool_.apply_async( + run_solver, + [ + h1, + dm0, + dname, + nao, + nocc, + nfsites, + efac, + TA, + hf_veff, + h1_e, + solver, + Fobjs[nf].eri_file, + veff0, + hci_cutoff, + ci_coeff_cutoff, + select_cutoff, + ompnum, + writeh1, + True, + True, + use_cumulant, + relax_density, + frag_energy, + ], + ) results.append(result) @@ -422,19 +526,19 @@ def be_func_parallel(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, if frag_energy: # Compute and return fragment energy # rdms are the returned energies, not density matrices! - e_1 = 0. - e_2 = 0. - e_c = 0. + e_1 = 0.0 + e_2 = 0.0 + e_c = 0.0 for i in range(len(rdms)): e_1 += rdms[i][0] e_2 += rdms[i][1] e_c += rdms[i][2] - return (e_1+e_2+e_c, (e_1, e_2, e_c)) + return (e_1 + e_2 + e_c, (e_1, e_2, e_c)) # Compute total energy - e_1 = 0. - e_2 = 0. - e_c = 0. + e_1 = 0.0 + e_2 = 0.0 + e_c = 0.0 for idx, fobj in enumerate(Fobjs): e_1 += rdms[idx][0][0] e_2 += rdms[idx][0][1] @@ -445,21 +549,34 @@ def be_func_parallel(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, fobj.__rdm1 = rdms[idx][4] del rdms - ernorm, ervec = solve_error(Fobjs,Nocc, only_chem=only_chem) + ernorm, ervec = solve_error(Fobjs, Nocc, only_chem=only_chem) if return_vec: - return (ernorm, ervec, [e_1+e_2+e_c, [e_1, e_2, e_c]]) + return (ernorm, ervec, [e_1 + e_2 + e_c, [e_1, e_2, e_c]]) if eeval: - print('Error in density matching : {:>2.4e}'.format(ernorm), flush=True) + print("Error in density matching : {:>2.4e}".format(ernorm), flush=True) return ernorm -def be_func_parallel_u(pot, Fobjs, solver, enuc, hf_veff=None, - nproc=1, ompnum=4, - relax_density=False,use_cumulant=True, - eeval=False, ereturn=False, frag_energy=False, - ecore=0., ebe_hf=0., frozen=False): + +def be_func_parallel_u( + pot, + Fobjs, + solver, + enuc, + hf_veff=None, + nproc=1, + ompnum=4, + relax_density=False, + use_cumulant=True, + eeval=False, + ereturn=False, + frag_energy=False, + ecore=0.0, + ebe_hf=0.0, + frozen=False, +): """ Embarrassingly Parallel High-Level Computation @@ -506,26 +623,30 @@ def be_func_parallel_u(pot, Fobjs, solver, enuc, hf_veff=None, import os # Set the number of OpenMP threads - os.system('export OMP_NUM_THREADS='+str(ompnum)) - nprocs = int(nproc/ompnum) + os.system("export OMP_NUM_THREADS=" + str(ompnum)) + nprocs = int(nproc / ompnum) pool_ = Pool(nprocs) results = [] energy_list = [] # Run solver in parallel for each fragment - for (fobj_a, fobj_b) in Fobjs: - - result = pool_.apply_async(run_solver_u, [fobj_a, fobj_b, - solver, - enuc, - hf_veff, - frag_energy, - relax_density, - frozen, - use_cumulant, - True - ]) + for fobj_a, fobj_b in Fobjs: + result = pool_.apply_async( + run_solver_u, + [ + fobj_a, + fobj_b, + solver, + enuc, + hf_veff, + frag_energy, + relax_density, + frozen, + use_cumulant, + True, + ], + ) results.append(result) # Collect results @@ -534,13 +655,15 @@ def be_func_parallel_u(pot, Fobjs, solver, enuc, hf_veff=None, if frag_energy: # Compute and return fragment energy - e_1 = 0. - e_2 = 0. - e_c = 0. + e_1 = 0.0 + e_2 = 0.0 + e_c = 0.0 for i in range(len(energy_list)): e_1 += energy_list[i][0] e_2 += energy_list[i][1] e_c += energy_list[i][2] - return (e_1+e_2+e_c, (e_1, e_2, e_c)) + return (e_1 + e_2 + e_c, (e_1, e_2, e_c)) else: - return NotImplementedError("Only fragment-wise energy return implemented, no RDM return") + return NotImplementedError( + "Only fragment-wise energy return implemented, no RDM return" + ) diff --git a/molbe/be_var.py b/molbe/be_var.py index fc97dc10..2ff58683 100644 --- a/molbe/be_var.py +++ b/molbe/be_var.py @@ -1,6 +1,6 @@ -PRINT_LEVEL=5 -WRITE_FILE=0 +PRINT_LEVEL = 5 +WRITE_FILE = 0 SOLVER_CALL = 0 -SCRATCH='' -CREATE_SCRATCH_DIR=False -INTEGRAL_TRANSFORM_MAX_MEMORY=50 # in GB +SCRATCH = "" +CREATE_SCRATCH_DIR = False +INTEGRAL_TRANSFORM_MAX_MEMORY = 50 # in GB diff --git a/molbe/eri_onthefly.py b/molbe/eri_onthefly.py index b9b2fd36..8a8f58c0 100644 --- a/molbe/eri_onthefly.py +++ b/molbe/eri_onthefly.py @@ -10,8 +10,9 @@ from scipy.linalg import cholesky, solve_triangular from . import be_var + def integral_direct_DF(mf, Fobjs, file_eri, auxbasis=None): - """ Calculate AO density-fitted 3-center integrals on-the-fly and transform to Schmidt space for given fragment objects + """Calculate AO density-fitted 3-center integrals on-the-fly and transform to Schmidt space for given fragment objects Parameters ---------- @@ -27,22 +28,42 @@ def integral_direct_DF(mf, Fobjs, file_eri, auxbasis=None): """ def calculate_pqL(aux_range): - """ Internal function to calculate the 3-center integrals for a given range of auxiliary indices + """Internal function to calculate the 3-center integrals for a given range of auxiliary indices Parameters ---------- aux_range : tuple of int (start index, end index) of the auxiliary basis functions to calculate the 3-center integrals, i.e. (pq|L) with L ∈ [start, end) is returned """ - if be_var.PRINT_LEVEL > 10: print('Start calculating (μν|P) for range', aux_range, flush=True) + if be_var.PRINT_LEVEL > 10: + print("Start calculating (μν|P) for range", aux_range, flush=True) p0, p1 = aux_range - shls_slice = (0, mf.mol.nbas, 0, mf.mol.nbas, mf.mol.nbas + p0, mf.mol.nbas + p1) - ints = getints3c(mf.mol._add_suffix('int3c2e'), atm, bas, env, shls_slice, 1, 's1', ao_loc, cintopt, out=None) - if be_var.PRINT_LEVEL > 10: print('Finish calculating (μν|P) for range', aux_range, flush=True) + shls_slice = ( + 0, + mf.mol.nbas, + 0, + mf.mol.nbas, + mf.mol.nbas + p0, + mf.mol.nbas + p1, + ) + ints = getints3c( + mf.mol._add_suffix("int3c2e"), + atm, + bas, + env, + shls_slice, + 1, + "s1", + ao_loc, + cintopt, + out=None, + ) + if be_var.PRINT_LEVEL > 10: + print("Finish calculating (μν|P) for range", aux_range, flush=True) return ints def block_step_size(nfrag, naux, nao): - """ Internal function to calculate the block step size for the 3-center integrals calculation + """Internal function to calculate the block step size for the 3-center integrals calculation Parameters ---------- @@ -53,40 +74,69 @@ def block_step_size(nfrag, naux, nao): nao : int Number of atomic orbitals """ - return max(1, int( - be_var.INTEGRAL_TRANSFORM_MAX_MEMORY * 1e9 / 8 / nao / nao / naux / nfrag - )) #max(int(500*.24e6/8/nao),1) + return max( + 1, + int( + be_var.INTEGRAL_TRANSFORM_MAX_MEMORY + * 1e9 + / 8 + / nao + / nao + / naux + / nfrag + ), + ) # max(int(500*.24e6/8/nao),1) if be_var.PRINT_LEVEL > 2: - print('Evaluating fragment ERIs on-the-fly using density fitting...', flush=True) - print('In this case, note that HF-in-HF error includes DF error on top of numerical error from embedding.', flush=True) - auxmol = make_auxmol(mf.mol, auxbasis = auxbasis) - j2c = auxmol.intor(mf.mol._add_suffix('int2c2e'), hermi=1) # (L|M) + print( + "Evaluating fragment ERIs on-the-fly using density fitting...", flush=True + ) + print( + "In this case, note that HF-in-HF error includes DF error on top of numerical error from embedding.", + flush=True, + ) + auxmol = make_auxmol(mf.mol, auxbasis=auxbasis) + j2c = auxmol.intor(mf.mol._add_suffix("int2c2e"), hermi=1) # (L|M) low = cholesky(j2c, lower=True) - pqL_frag = [numpy.zeros((auxmol.nao, fragobj.nao, fragobj.nao)) for fragobj in Fobjs] # place to store fragment (pq|L) + pqL_frag = [ + numpy.zeros((auxmol.nao, fragobj.nao, fragobj.nao)) for fragobj in Fobjs + ] # place to store fragment (pq|L) end = 0 - atm, bas, env = mole.conc_env(mf.mol._atm, mf.mol._bas, mf.mol._env, auxmol._atm, auxmol._bas, auxmol._env) - ao_loc = make_loc(bas, mf.mol._add_suffix('int3c2e')); cintopt = make_cintopt(atm, bas, env, mf.mol._add_suffix('int3c2e')) - blockranges = [(x, y) for x, y in lib.prange(0, auxmol.nbas, block_step_size(len(Fobjs), auxmol.nbas, mf.mol.nao))] + atm, bas, env = mole.conc_env( + mf.mol._atm, mf.mol._bas, mf.mol._env, auxmol._atm, auxmol._bas, auxmol._env + ) + ao_loc = make_loc(bas, mf.mol._add_suffix("int3c2e")) + cintopt = make_cintopt(atm, bas, env, mf.mol._add_suffix("int3c2e")) + blockranges = [ + (x, y) + for x, y in lib.prange( + 0, auxmol.nbas, block_step_size(len(Fobjs), auxmol.nbas, mf.mol.nao) + ) + ] if be_var.PRINT_LEVEL > 4: - print('Aux Basis Block Info: ', blockranges, flush=True) + print("Aux Basis Block Info: ", blockranges, flush=True) for idx, ints in enumerate(lib.map_with_prefetch(calculate_pqL, blockranges)): - if be_var.PRINT_LEVEL > 4: print('Calculating pq|L block #', idx, blockranges[idx], flush=True) + if be_var.PRINT_LEVEL > 4: + print("Calculating pq|L block #", idx, blockranges[idx], flush=True) # Transform pq (AO) to fragment space (ij) - start = end; end += ints.shape[2] + start = end + end += ints.shape[2] for fragidx in range(len(Fobjs)): - if be_var.PRINT_LEVEL > 10: print('(μν|P) -> (ij|P) for frag #', fragidx, flush=True) - Lqp = numpy.transpose(ints, axes=(2,1,0)) + if be_var.PRINT_LEVEL > 10: + print("(μν|P) -> (ij|P) for frag #", fragidx, flush=True) + Lqp = numpy.transpose(ints, axes=(2, 1, 0)) Lqi = Lqp @ Fobjs[fragidx].TA Liq = numpy.moveaxis(Lqi, 2, 1) pqL_frag[fragidx][start:end, :, :] = Liq @ Fobjs[fragidx].TA # Fit to get B_{ij}^{L} for fragidx in range(len(Fobjs)): - if be_var.PRINT_LEVEL > 10: print('Fitting B_{ij}^{L} for frag #', fragidx, flush=True) + if be_var.PRINT_LEVEL > 10: + print("Fitting B_{ij}^{L} for frag #", fragidx, flush=True) b = pqL_frag[fragidx].reshape(auxmol.nao, -1) bb = solve_triangular(low, b, lower=True, overwrite_b=True, check_finite=False) - if be_var.PRINT_LEVEL > 10: print('Finished obtaining B_{ij}^{L} for frag #', fragidx, flush=True) + if be_var.PRINT_LEVEL > 10: + print("Finished obtaining B_{ij}^{L} for frag #", fragidx, flush=True) eri_nosym = bb.T @ bb - eri = restore('4', eri_nosym, Fobjs[fragidx].nao) + eri = restore("4", eri_nosym, Fobjs[fragidx].nao) file_eri.create_dataset(Fobjs[fragidx].dname, data=eri) diff --git a/molbe/external/ccsd_rdm.py b/molbe/external/ccsd_rdm.py index a0305b4f..b10faa46 100644 --- a/molbe/external/ccsd_rdm.py +++ b/molbe/external/ccsd_rdm.py @@ -3,31 +3,35 @@ # The code has been slightly modified. # -import numpy,functools,sys, time,os +import numpy, functools, sys, time, os + def make_rdm1_ccsd_t1(t1): nocc, nvir = t1.shape nmo = nocc + nvir - dm = numpy.zeros((nmo,nmo), dtype=t1.dtype) - dm[:nocc,nocc:] = t1 - dm[nocc:,:nocc] = t1.T - dm[numpy.diag_indices(nocc)] += 2. + dm = numpy.zeros((nmo, nmo), dtype=t1.dtype) + dm[:nocc, nocc:] = t1 + dm[nocc:, :nocc] = t1.T + dm[numpy.diag_indices(nocc)] += 2.0 return dm + def make_rdm2_urlx(t1, t2, with_dm1=True): nocc, nvir = t1.shape nmo = nocc + nvir goovv = (numpy.einsum("ia,jb->ijab", t1, t1) + t2) * 0.5 - dovov = goovv.transpose(0,2,1,3) * 2 - goovv.transpose(1,2,0,3) + dovov = goovv.transpose(0, 2, 1, 3) * 2 - goovv.transpose(1, 2, 0, 3) - dm2 = numpy.zeros([nmo,nmo,nmo,nmo], dtype=t1.dtype) + dm2 = numpy.zeros([nmo, nmo, nmo, nmo], dtype=t1.dtype) dovov = numpy.asarray(dovov) - dm2[:nocc,nocc:,:nocc,nocc:] = dovov - dm2[:nocc,nocc:,:nocc,nocc:]+= dovov.transpose(2,3,0,1) - dm2[nocc:,:nocc,nocc:,:nocc] = dm2[:nocc,nocc:,:nocc,nocc:].transpose(1,0,3,2).conj() + dm2[:nocc, nocc:, :nocc, nocc:] = dovov + dm2[:nocc, nocc:, :nocc, nocc:] += dovov.transpose(2, 3, 0, 1) + dm2[nocc:, :nocc, nocc:, :nocc] = ( + dm2[:nocc, nocc:, :nocc, nocc:].transpose(1, 0, 3, 2).conj() + ) dovov = None if with_dm1: @@ -35,34 +39,38 @@ def make_rdm2_urlx(t1, t2, with_dm1=True): dm1[numpy.diag_indices(nocc)] -= 2 for i in range(nocc): - dm2[i,i,:,:] += dm1 * 2 - dm2[:,:,i,i] += dm1 * 2 - dm2[:,i,i,:] -= dm1 - dm2[i,:,:,i] -= dm1.T + dm2[i, i, :, :] += dm1 * 2 + dm2[:, :, i, i] += dm1 * 2 + dm2[:, i, i, :] -= dm1 + dm2[i, :, :, i] -= dm1.T for i in range(nocc): for j in range(nocc): - dm2[i,i,j,j] += 4 - dm2[i,j,j,i] -= 2 + dm2[i, i, j, j] += 4 + dm2[i, j, j, i] -= 2 return dm2 + def make_rdm1_uccsd(ucc, relax=False): from pyscf.cc.uccsd_rdm import make_rdm1 - if relax==True: + + if relax == True: rdm1 = make_rdm1(ucc, ucc.t1, ucc.t2, ucc.l1, ucc.l2) else: - l1 = [numpy.zeros_like(ucc.t1[s]) for s in [0,1]] - l2 = [numpy.zeros_like(ucc.t2[s]) for s in [0,1,2]] + l1 = [numpy.zeros_like(ucc.t1[s]) for s in [0, 1]] + l2 = [numpy.zeros_like(ucc.t2[s]) for s in [0, 1, 2]] rdm1 = make_rdm1(ucc, ucc.t1, ucc.t2, l1, l2) return rdm1 + def make_rdm2_uccsd(ucc, relax=False, with_dm1=True): from pyscf.cc.uccsd_rdm import make_rdm2 - if relax==True: + + if relax == True: rdm2 = make_rdm2(ucc, ucc.t1, ucc.t2, ucc.l1, ucc.l2, with_dm1=with_dm1) else: - l1 = [numpy.zeros_like(ucc.t1[s]) for s in [0,1]] - l2 = [numpy.zeros_like(ucc.t2[s]) for s in [0,1,2]] + l1 = [numpy.zeros_like(ucc.t1[s]) for s in [0, 1]] + l2 = [numpy.zeros_like(ucc.t2[s]) for s in [0, 1, 2]] rdm2 = make_rdm2(ucc, ucc.t1, ucc.t2, l1, l2, with_dm1=with_dm1) return rdm2 diff --git a/molbe/external/cphf_utils.py b/molbe/external/cphf_utils.py index f27d5b0e..d883373b 100644 --- a/molbe/external/cphf_utils.py +++ b/molbe/external/cphf_utils.py @@ -10,18 +10,21 @@ def get_cphf_A(C, moe, eri, no): nao = C.shape[0] nv = nao - no - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] nov = no * nv - Vovov = ao2mo.incore.general(eri, (Co,Cv,Co,Cv), - compact=False).reshape(no,nv,no,nv) - Voovv = ao2mo.incore.general(eri, (Co,Co,Cv,Cv), - compact=False).reshape(no,no,nv,nv) + Vovov = ao2mo.incore.general(eri, (Co, Cv, Co, Cv), compact=False).reshape( + no, nv, no, nv + ) + Voovv = ao2mo.incore.general(eri, (Co, Co, Cv, Cv), compact=False).reshape( + no, no, nv, nv + ) # 4*Viajb - Vibja - Vjiab - A = (4.*Vovov - Vovov.transpose(0,3,2,1) - \ - Voovv.transpose(0,2,1,3)).reshape(nov,nov) - denom = (moe[:no].reshape(-1,1) - moe[no:]).ravel() + A = ( + 4.0 * Vovov - Vovov.transpose(0, 3, 2, 1) - Voovv.transpose(0, 2, 1, 3) + ).reshape(nov, nov) + denom = (moe[:no].reshape(-1, 1) - moe[no:]).ravel() A -= np.diag(denom) return A @@ -30,8 +33,8 @@ def get_cphf_A(C, moe, eri, no): def get_cphf_rhs(C, no, v): nao = C.shape[0] nv = nao - no - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] return (Co.T @ v @ Cv).ravel() @@ -61,7 +64,7 @@ def cphf_kernel_batch(C, moe, eri, no, vs): # build RHS vectors B0s = np.zeros([nov, npot]) for i, v in enumerate(vs): - B0s[:,i] = get_cphf_rhs(C, no, v) + B0s[:, i] = get_cphf_rhs(C, no, v) # build A matrix A = get_cphf_A(C, moe, eri, no) @@ -75,110 +78,142 @@ def cphf_kernel_batch(C, moe, eri, no, vs): def get_rhf_dP_from_u(C, no, u): n = C.shape[0] nv = n - no - dP = -C[:,:no] @ u.reshape(no,nv) @ C[:,no:].T + dP = -C[:, :no] @ u.reshape(no, nv) @ C[:, no:].T dP += dP.T return dP -def get_full_u(C, moe, eri, no, v, u, thresh=1E8): +def get_full_u(C, moe, eri, no, v, u, thresh=1e8): nao = C.shape[0] nv = nao - no - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] moeo = moe[:no] moev = moe[no:] - Voovo = ao2mo.incore.general(eri, (Co,Co,Cv,Co), - compact=False).reshape(no,no,nv,no) - Vvvvo = ao2mo.incore.general(eri, (Cv,Cv,Cv,Co), - compact=False).reshape(nv,nv,nv,no) + Voovo = ao2mo.incore.general(eri, (Co, Co, Cv, Co), compact=False).reshape( + no, no, nv, no + ) + Vvvvo = ao2mo.incore.general(eri, (Cv, Cv, Cv, Co), compact=False).reshape( + nv, nv, nv, no + ) vmo = C.T @ v @ C uvo = u.reshape(nv, no) - U = np.zeros([nao,nao]) - U[no:,:no] = uvo - U[:no,no:] = -uvo.T + U = np.zeros([nao, nao]) + U[no:, :no] = uvo + U[:no, no:] = -uvo.T # occupied-occupied block # ( Qjk + sum_ai (4*Vjkai-Vjiak-Vikaj)*uai ) / (ek - ej) - denom_oo = moeo.reshape(no,1) - moeo - denom_oo[np.diag_indices(no)] = 1. + denom_oo = moeo.reshape(no, 1) - moeo + denom_oo[np.diag_indices(no)] = 1.0 if np.sum(np.abs(denom_oo**-1) > thresh) > 0: raise RuntimeError - U[:no,:no] = -(vmo[:no,:no] + \ - np.einsum("jkai,ai->jk", 4.*Voovo - \ - Voovo.transpose(0,3,2,1) - \ - Voovo.transpose(3,1,2,0), \ - uvo)) / denom_oo - U[np.diag_indices(no)] = 0. + U[:no, :no] = ( + -( + vmo[:no, :no] + + np.einsum( + "jkai,ai->jk", + 4.0 * Voovo - Voovo.transpose(0, 3, 2, 1) - Voovo.transpose(3, 1, 2, 0), + uvo, + ) + ) + / denom_oo + ) + U[np.diag_indices(no)] = 0.0 # virtual-virtual block # ( Qbc + sum_ai (4*Vbcai-Vacbi-Vbaci)*uai ) / (ec - eb) - denom_vv = moev.reshape(nv,1) - moev - denom_vv[np.diag_indices(nv)] = 1. + denom_vv = moev.reshape(nv, 1) - moev + denom_vv[np.diag_indices(nv)] = 1.0 if np.sum(np.abs(denom_vv**-1) > thresh) > 0: raise RuntimeError - uvv = -(vmo[no:,no:] + \ - np.einsum("bcai,ai->bc", 4.*Vvvvo - \ - Vvvvo.transpose(2,1,0,3) - \ - Vvvvo.transpose(0,2,1,3), \ - uvo)) / denom_vv - uvv[np.diag_indices(nv)] = 0. - U[no:,no:] = uvv + uvv = ( + -( + vmo[no:, no:] + + np.einsum( + "bcai,ai->bc", + 4.0 * Vvvvo - Vvvvo.transpose(2, 1, 0, 3) - Vvvvo.transpose(0, 2, 1, 3), + uvo, + ) + ) + / denom_vv + ) + uvv[np.diag_indices(nv)] = 0.0 + U[no:, no:] = uvv return U -def get_full_u_batch(C, moe, eri, no, vs, us, thresh=1E10, timing=False): +def get_full_u_batch(C, moe, eri, no, vs, us, thresh=1e10, timing=False): nao = C.shape[0] nv = nao - no npot = len(vs) - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] moeo = moe[:no] moev = moe[no:] - Voovo = ao2mo.incore.general(eri, (Co,Co,Cv,Co), - compact=False).reshape(no,no,nv,no) - Vvvvo = ao2mo.incore.general(eri, (Cv,Cv,Cv,Co), - compact=False).reshape(nv,nv,nv,no) + Voovo = ao2mo.incore.general(eri, (Co, Co, Cv, Co), compact=False).reshape( + no, no, nv, no + ) + Vvvvo = ao2mo.incore.general(eri, (Cv, Cv, Cv, Co), compact=False).reshape( + nv, nv, nv, no + ) Us = [None for i in range(npot)] for i in range(npot): vmo = C.T @ vs[i] @ C - uvo = us[:,i].reshape(nv, no) - U = np.zeros([nao,nao]) - U[no:,:no] = uvo - U[:no,no:] = -uvo.T + uvo = us[:, i].reshape(nv, no) + U = np.zeros([nao, nao]) + U[no:, :no] = uvo + U[:no, no:] = -uvo.T # occupied-occupied block # ( Qjk + sum_ai (4*Vjkai-Vjiak-Vikaj)*uai ) / (ek - ej) - denom_oo = moeo.reshape(no,1) - moeo - denom_oo[np.diag_indices(no)] = 1. + denom_oo = moeo.reshape(no, 1) - moeo + denom_oo[np.diag_indices(no)] = 1.0 if np.sum(np.abs(denom_oo**-1) > thresh) > 0: raise RuntimeError - U[:no,:no] = -(vmo[:no,:no] + \ - np.einsum("jkai,ai->jk", 4.*Voovo - \ - Voovo.transpose(0,3,2,1) - \ - Voovo.transpose(3,1,2,0), \ - uvo)) / denom_oo - U[np.diag_indices(no)] = 0. + U[:no, :no] = ( + -( + vmo[:no, :no] + + np.einsum( + "jkai,ai->jk", + 4.0 * Voovo + - Voovo.transpose(0, 3, 2, 1) + - Voovo.transpose(3, 1, 2, 0), + uvo, + ) + ) + / denom_oo + ) + U[np.diag_indices(no)] = 0.0 # virtual-virtual block # ( Qbc + sum_ai (4*Vbcai-Vacbi-Vbaci)*uai ) / (ec - eb) - denom_vv = moev.reshape(nv,1) - moev - denom_vv[np.diag_indices(nv)] = 1. + denom_vv = moev.reshape(nv, 1) - moev + denom_vv[np.diag_indices(nv)] = 1.0 if np.sum(np.abs(denom_vv**-1) > thresh) > 0: raise RuntimeError - uvv = -(vmo[no:,no:] + \ - np.einsum("bcai,ai->bc", 4.*Vvvvo - \ - Vvvvo.transpose(2,1,0,3) - \ - Vvvvo.transpose(0,2,1,3), \ - uvo)) / denom_vv - uvv[np.diag_indices(nv)] = 0. - U[no:,no:] = uvv + uvv = ( + -( + vmo[no:, no:] + + np.einsum( + "bcai,ai->bc", + 4.0 * Vvvvo + - Vvvvo.transpose(2, 1, 0, 3) + - Vvvvo.transpose(0, 2, 1, 3), + uvo, + ) + ) + / denom_vv + ) + uvv[np.diag_indices(nv)] = 0.0 + U[no:, no:] = uvv Us[i] = U @@ -190,10 +225,10 @@ def uvo_as_full_u_batch(nao, no, us): nv = nao - no Us = [None] * ncon for i in range(ncon): - uvo = us[:,i].reshape(nv, no) + uvo = us[:, i].reshape(nv, no) U = np.zeros([nao, nao]) - U[no:,:no] = uvo - U[:no,no:] = -uvo.T + U[no:, :no] = uvo + U[:no, no:] = -uvo.T Us[i] = U return Us @@ -203,16 +238,16 @@ def get_dP_lagrangian(C, no): nao = C.shape[0] nv = nao - no - L = np.zeros([nv*no,nao*(nao+1)//2]) + L = np.zeros([nv * no, nao * (nao + 1) // 2]) mn = -1 for mu in range(nao): - for nu in range(mu,nao): + for nu in range(mu, nao): mn += 1 ai = -1 for a in range(nv): for i in range(no): ai += 1 - L[ai,mn] = C[mu,a+no]*C[nu,i] + C[mu,i]*C[nu,a+no] + L[ai, mn] = C[mu, a + no] * C[nu, i] + C[mu, i] * C[nu, a + no] return L @@ -228,113 +263,131 @@ def get_zvec(C, moe, eri, no): L = get_dP_lagrangian(C, no) # solve for z vector - z = np.linalg.solve(np.eye(nv*no)-A.T, L) + z = np.linalg.solve(np.eye(nv * no) - A.T, L) return z """ For UHF """ + + def get_cpuhf_A_spinless_eri(C, moe, eri, no): n = C[0].shape[0] - nv = [n-no[s] for s in [0,1]] - nov = [no[s]*nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] - moeo = [moe[s][:no[s]] for s in [0,1]] - moev = [moe[s][no[s]:] for s in [0,1]] + nv = [n - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] + moeo = [moe[s][: no[s]] for s in [0, 1]] + moev = [moe[s][no[s] :] for s in [0, 1]] nA = sum(nov) - A = np.empty([nA,nA]) - for s in [0,1]: + A = np.empty([nA, nA]) + for s in [0, 1]: # same spin - Vss_ovov = ao2mo.incore.general(eri, (Co[s],Cv[s],Co[s],Cv[s]), - compact=False).reshape(no[s],nv[s],no[s],nv[s]) - Vss_oovv = ao2mo.incore.general(eri, (Co[s],Co[s],Cv[s],Cv[s]), - compact=False).reshape(no[s],no[s],nv[s],nv[s]) - Ass = (Vss_ovov*2 - - Vss_ovov.transpose(0,3,2,1)).reshape(nov[s],nov[s]) - \ - Vss_oovv.transpose(0,2,1,3).reshape(nov[s],nov[s]) - Dss = (moeo[s].reshape(-1,1) - moev[s]).ravel() + Vss_ovov = ao2mo.incore.general( + eri, (Co[s], Cv[s], Co[s], Cv[s]), compact=False + ).reshape(no[s], nv[s], no[s], nv[s]) + Vss_oovv = ao2mo.incore.general( + eri, (Co[s], Co[s], Cv[s], Cv[s]), compact=False + ).reshape(no[s], no[s], nv[s], nv[s]) + Ass = (Vss_ovov * 2 - Vss_ovov.transpose(0, 3, 2, 1)).reshape( + nov[s], nov[s] + ) - Vss_oovv.transpose(0, 2, 1, 3).reshape(nov[s], nov[s]) + Dss = (moeo[s].reshape(-1, 1) - moev[s]).ravel() Ass[np.diag_indices(nov[s])] -= Dss if s == 0: - A[:nov[s],:nov[s]] = Ass + A[: nov[s], : nov[s]] = Ass else: - A[nov[s]:,nov[s]:] = Ass + A[nov[s] :, nov[s] :] = Ass # opposite spin - Vos_ovov = ao2mo.incore.general(eri, (Co[s],Cv[s],Co[1-s],Cv[1-s]), - compact=False) * 2. + Vos_ovov = ( + ao2mo.incore.general( + eri, (Co[s], Cv[s], Co[1 - s], Cv[1 - s]), compact=False + ) + * 2.0 + ) if s == 0: - A[:nov[s],nov[s]:] = Vos_ovov + A[: nov[s], nov[s] :] = Vos_ovov else: - A[nov[s]:,:nov[s]] = Vos_ovov + A[nov[s] :, : nov[s]] = Vos_ovov return A def get_cpuhf_A_spin_eri(C, moe, eri, no): n = C[0].shape[0] - nv = [n-no[s] for s in [0,1]] - nov = [no[s]*nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] - moeo = [moe[s][:no[s]] for s in [0,1]] - moev = [moe[s][no[s]:] for s in [0,1]] + nv = [n - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] + moeo = [moe[s][: no[s]] for s in [0, 1]] + moev = [moe[s][no[s] :] for s in [0, 1]] nA = sum(nov) - A = np.empty([nA,nA]) - for s in [0,1]: + A = np.empty([nA, nA]) + for s in [0, 1]: # same spin - Vss_ovov = ao2mo.incore.general(eri[s], (Co[s],Cv[s],Co[s],Cv[s]), - compact=False) - Vss_oovv = ao2mo.incore.general(eri[s], (Co[s],Co[s],Cv[s],Cv[s]), - compact=False) - Ass = (Vss_ovov*2 - - Vss_ovov.transpose(0,3,2,1)).reshape(nov[s],nov[s]) - \ - Vss_oovv.transpose(0,2,1,3).reshape(nov[s],nov[s]) - Dss = (moeo[s].reshape(-1,1) - moev[s]).ravel() + Vss_ovov = ao2mo.incore.general( + eri[s], (Co[s], Cv[s], Co[s], Cv[s]), compact=False + ) + Vss_oovv = ao2mo.incore.general( + eri[s], (Co[s], Co[s], Cv[s], Cv[s]), compact=False + ) + Ass = (Vss_ovov * 2 - Vss_ovov.transpose(0, 3, 2, 1)).reshape( + nov[s], nov[s] + ) - Vss_oovv.transpose(0, 2, 1, 3).reshape(nov[s], nov[s]) + Dss = (moeo[s].reshape(-1, 1) - moev[s]).ravel() Ass[np.diag_indices(nov[s])] -= Dss if s == 0: - A[:nov[s],:nov[s]] = Ass + A[: nov[s], : nov[s]] = Ass else: - A[nov[s]:,nov[s]:] = Ass + A[nov[s] :, nov[s] :] = Ass # opposite spin - Vos_ovov = ao2mo.incore.general(eri[2], (Co[0],Cv[0],Co[1],Cv[1]), - compact=False).reshape(nov[0],nov[1]) * 2. - if s == 1: Vos_ovov = Vos_ovov.T + Vos_ovov = ( + ao2mo.incore.general( + eri[2], (Co[0], Cv[0], Co[1], Cv[1]), compact=False + ).reshape(nov[0], nov[1]) + * 2.0 + ) + if s == 1: + Vos_ovov = Vos_ovov.T if s == 0: - A[:nov[s],nov[s]:] = Vos_ovov + A[: nov[s], nov[s] :] = Vos_ovov else: - A[nov[s]:,:nov[s]] = Vos_ovov + A[nov[s] :, : nov[s]] = Vos_ovov return A def get_cpuhf_A(C, moe, eri, no): - """ eri could be a single numpy array (in spinless basis) or a tuple/list of 3 numpy arrays in the order [aa,bb,ab] - """ + """eri could be a single numpy array (in spinless basis) or a tuple/list of 3 numpy arrays in the order [aa,bb,ab]""" if isinstance(eri, np.ndarray): return get_cpuhf_A_spinless_eri(C, moe, eri, no) - elif isinstance(eri, (list,tuple)): + elif isinstance(eri, (list, tuple)): if len(eri) != 3: - raise ValueError("Input eri must be a list/tuple of 3 numpy arrays in the order [aa,bb,ab]") + raise ValueError( + "Input eri must be a list/tuple of 3 numpy arrays in the order [aa,bb,ab]" + ) return get_cpuhf_A_spin_eri(C, moe, eri, no) else: - raise ValueError("Input eri must be either a numpy array or a list/tuple of 3 numpy arrays in the order [aa,bb,ab].") + raise ValueError( + "Input eri must be either a numpy array or a list/tuple of 3 numpy arrays in the order [aa,bb,ab]." + ) def get_cpuhf_u(C, moe, eri, no, vpot): n = C[0].shape[0] - nv = [n-no[s] for s in [0,1]] - nov = [no[s]*nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] + nv = [n - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] # build lhs A = get_cpuhf_A(C, moe, eri, no) # build rhs - b = np.concatenate([(Co[s].T@vpot[s]@Cv[s]).ravel() for s in [0,1]]) + b = np.concatenate([(Co[s].T @ vpot[s] @ Cv[s]).ravel() for s in [0, 1]]) # solve u = np.linalg.solve(A, b) @@ -344,10 +397,10 @@ def get_cpuhf_u(C, moe, eri, no, vpot): def get_cpuhf_u_batch(C, moe, eri, no, vpots): n = C[0].shape[0] - nv = [n-no[s] for s in [0,1]] - nov = [no[s]*nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] + nv = [n - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] # build lhs A = get_cpuhf_A(C, moe, eri, no) @@ -355,10 +408,11 @@ def get_cpuhf_u_batch(C, moe, eri, no, vpots): # build rhs nA = A.shape[0] npot = len(vpots) - bs = np.zeros([nA,npot]) + bs = np.zeros([nA, npot]) for i in range(npot): - bs[:,i] = np.concatenate([(Co[s].T@vpots[i][s]@Cv[s]).ravel() - for s in [0,1]]) + bs[:, i] = np.concatenate( + [(Co[s].T @ vpots[i][s] @ Cv[s]).ravel() for s in [0, 1]] + ) # solve us = np.linalg.solve(A, bs) @@ -368,15 +422,18 @@ def get_cpuhf_u_batch(C, moe, eri, no, vpots): def get_uhf_dP_from_u(C, no, u): n = C[0].shape[0] - nv = [n-no[s] for s in [0,1]] - nov = [no[s]*nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] + nv = [n - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] dP = [None] * 2 - for s in [0,1]: - u_ = u[:nov[0]].reshape([no[s],nv[s]]) if s == 0 \ - else u[nov[0]:].reshape(no[s],nv[s]) - dP[s] = - Co[s] @ u_ @ Cv[s].T + for s in [0, 1]: + u_ = ( + u[: nov[0]].reshape([no[s], nv[s]]) + if s == 0 + else u[nov[0] :].reshape(no[s], nv[s]) + ) + dP[s] = -Co[s] @ u_ @ Cv[s].T dP[s] += dP[s].T return dP diff --git a/molbe/external/cpmp2_utils.py b/molbe/external/cpmp2_utils.py index b4938ef1..a6c276e3 100644 --- a/molbe/external/cpmp2_utils.py +++ b/molbe/external/cpmp2_utils.py @@ -11,13 +11,15 @@ """ RMP2 implementation """ + + def get_Diajb_r(moe, no): n = moe.size nv = n - no eo = moe[:no] ev = moe[no:] - Dia = (eo.reshape(-1,1) - ev).ravel() - Diajb = (Dia.reshape(-1,1) + Dia).reshape(no,nv,no,nv) + Dia = (eo.reshape(-1, 1) - ev).ravel() + Diajb = (Dia.reshape(-1, 1) + Dia).reshape(no, nv, no, nv) return Diajb @@ -27,19 +29,19 @@ def get_dF_r(no, V, C, Q, u): n = C.shape[0] nv = n - no - Co = C[:,:no] - Cv = C[:,no:] - uov = u.reshape(no,nv) + Co = C[:, :no] + Cv = C[:, no:] + uov = u.reshape(no, nv) dP = -Co @ uov @ Cv.T dP += dP.T - vj, vk = scf.hf.dot_eri_dm(V, dP*2., hermi=1) - dF = Q + vj-0.5*vk + vj, vk = scf.hf.dot_eri_dm(V, dP * 2.0, hermi=1) + dF = Q + vj - 0.5 * vk return dF def get_dmoe_F_r(C, dF): - de = np.einsum("pi,qi,pq->i",C,C,dF) + de = np.einsum("pi,qi,pq->i", C, C, dF) return de @@ -47,19 +49,19 @@ def get_dmoe_F_r(C, dF): def get_full_u_F_r(no, C, moe, dF, u): n = C.shape[0] nv = n - no - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] eo = moe[:no] ev = moe[no:] - uov = u.reshape(no,nv) - Dij = -eo.reshape(-1,1) + eo + uov = u.reshape(no, nv) + Dij = -eo.reshape(-1, 1) + eo np.fill_diagonal(Dij, 1) - dUoo = (Co.T@dF@Co) / Dij - np.fill_diagonal(dUoo, 0.) - Dab = -ev.reshape(-1,1) + ev + dUoo = (Co.T @ dF @ Co) / Dij + np.fill_diagonal(dUoo, 0.0) + Dab = -ev.reshape(-1, 1) + ev np.fill_diagonal(Dab, 1) - dUvv = (Cv.T@dF@Cv) / Dab - np.fill_diagonal(dUvv, 0.) + dUvv = (Cv.T @ dF @ Cv) / Dab + np.fill_diagonal(dUvv, 0.0) return np.block([[dUoo, uov], [-uov.T, dUvv]]) @@ -68,67 +70,69 @@ def get_dVovov_r(no, V, C, u): n = C.shape[0] nv = n - no nov = no * nv - Co = C[:,:no] - Cv = C[:,no:] - if u.size == no*nv: - uov = u.reshape(no,nv) + Co = C[:, :no] + Cv = C[:, no:] + if u.size == no * nv: + uov = u.reshape(no, nv) dCv = Co @ uov dCo = -Cv @ uov.T else: dC = C @ u - dCo = dC[:,:no] - dCv = dC[:,no:] - Vovov_ = ao2mo.incore.general(V,(Co,Cv,Co,dCv)).reshape(nov,nov) - Vovo_v = ao2mo.incore.general(V,(Co,Cv,dCo,Cv)).reshape(nov,nov) - dVovov = (Vovov_ + Vovov_.T + Vovo_v + Vovo_v.T).reshape(no,nv,no,nv) + dCo = dC[:, :no] + dCv = dC[:, no:] + Vovov_ = ao2mo.incore.general(V, (Co, Cv, Co, dCv)).reshape(nov, nov) + Vovo_v = ao2mo.incore.general(V, (Co, Cv, dCo, Cv)).reshape(nov, nov) + dVovov = (Vovov_ + Vovov_.T + Vovo_v + Vovo_v.T).reshape(no, nv, no, nv) return dVovov def get_Pmp2_r(t2l, t2r): - assert(t2l.ndim == t2r.ndim == 4) + assert t2l.ndim == t2r.ndim == 4 no = t2l.shape[0] - Poo = -np.einsum("iajb,majb->im", t2l, 2.*t2r-t2r.transpose(0,3,2,1), - optimize=True) - Pvv = np.einsum("iajb,icjb->ac", t2l, 2.*t2r-t2r.transpose(0,3,2,1), - optimize=True) + Poo = -np.einsum( + "iajb,majb->im", t2l, 2.0 * t2r - t2r.transpose(0, 3, 2, 1), optimize=True + ) + Pvv = np.einsum( + "iajb,icjb->ac", t2l, 2.0 * t2r - t2r.transpose(0, 3, 2, 1), optimize=True + ) - return slg.block_diag(Poo,Pvv) + return slg.block_diag(Poo, Pvv) def get_dPmp2_batch_r(C, moe, V, no, Qs, aorep=True): - """ Derivative of oo and vv block of the MP2 density - """ + """Derivative of oo and vv block of the MP2 density""" n = C.shape[0] nv = n - no - Co = C[:,:no] - Cv = C[:,no:] - Vovov = ao2mo.incore.general(V, (Co,Cv,Co,Cv)).reshape(no,nv,no,nv) + Co = C[:, :no] + Cv = C[:, no:] + Vovov = ao2mo.incore.general(V, (Co, Cv, Co, Cv)).reshape(no, nv, no, nv) Diajb = get_Diajb_r(moe, no) t2 = Vovov / Diajb - from .cphf_utils import (cphf_kernel_batch as cphf_kernel) + from .cphf_utils import cphf_kernel_batch as cphf_kernel + us = cphf_kernel(C, moe, V, no, Qs) nQ = len(Qs) dPs = [None] * nQ Phf = np.diag([1 if i < no else 0 for i in range(n)]) iQ = 0 - for u,Q in zip(us,Qs): + for u, Q in zip(us, Qs): dF = get_dF_r(no, V, C, Q, u) dmoe = get_dmoe_F_r(C, dF) dDiajb = get_Diajb_r(dmoe, no) U = get_full_u_F_r(no, C, moe, dF, u) dVovov = get_dVovov_r(no, V, C, U) - dt2 = (dVovov - t2*dDiajb) / Diajb + dt2 = (dVovov - t2 * dDiajb) / Diajb - P = get_Pmp2_r(t2,t2) + P = get_Pmp2_r(t2, t2) P += Phf dP = U @ P - P @ U - dP2 = get_Pmp2_r(dt2,t2) + dP2 = get_Pmp2_r(dt2, t2) dP2 += dP2.T - dPs[iQ] = (dP + dP2) * 2. + dPs[iQ] = (dP + dP2) * 2.0 iQ += 1 @@ -141,37 +145,42 @@ def get_dPmp2_batch_r(C, moe, V, no, Qs, aorep=True): """ UMP2 implementation """ + + def get_Diajb_u(moe, no): - n = [moe[s].size for s in [0,1]] - nv = [n[s] - no[s] for s in [0,1]] - eo = [moe[s][:no[s]] for s in [0,1]] - ev = [moe[s][no[s]:] for s in [0,1]] - Dia = [(eo[s].reshape(-1,1)-ev[s]).ravel() for s in [0,1]] + n = [moe[s].size for s in [0, 1]] + nv = [n[s] - no[s] for s in [0, 1]] + eo = [moe[s][: no[s]] for s in [0, 1]] + ev = [moe[s][no[s] :] for s in [0, 1]] + Dia = [(eo[s].reshape(-1, 1) - ev[s]).ravel() for s in [0, 1]] Diajb = [ - (Dia[0].reshape(-1,1)+Dia[0]).reshape(no[0],nv[0],no[0],nv[0]), - (Dia[1].reshape(-1,1)+Dia[1]).reshape(no[1],nv[1],no[1],nv[1]), - (Dia[0].reshape(-1,1)+Dia[1]).reshape(no[0],nv[0],no[1],nv[1]) + (Dia[0].reshape(-1, 1) + Dia[0]).reshape(no[0], nv[0], no[0], nv[0]), + (Dia[1].reshape(-1, 1) + Dia[1]).reshape(no[1], nv[1], no[1], nv[1]), + (Dia[0].reshape(-1, 1) + Dia[1]).reshape(no[0], nv[0], no[1], nv[1]), ] return Diajb def get_dF_u(no, V, C, Q, u): - n = [C[s].shape[0] for s in [0,1]] - nv = [n[s] - no[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] + n = [C[s].shape[0] for s in [0, 1]] + nv = [n[s] - no[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] dP = [None] * 2 - for s in [0,1]: - uov = u[s].reshape(no[s],nv[s]) + for s in [0, 1]: + uov = u[s].reshape(no[s], nv[s]) dP[s] = -Co[s] @ uov @ Cv[s].T dP[s] += dP[s].T dF = [None] * 2 - for s in [0,1]: + for s in [0, 1]: vj_ss = np.einsum("pqrs,sr->pq", V[s], dP[s]) - vj_os = np.einsum("pqrs,sr->pq", V[2], dP[1]) if s == 0 \ + vj_os = ( + np.einsum("pqrs,sr->pq", V[2], dP[1]) + if s == 0 else np.einsum("pqrs,qp->rs", V[2], dP[0]) + ) vk_ss = np.einsum("psrq,sr->pq", V[s], dP[s]) dF[s] = Q[s] + vj_ss + vj_os - vk_ss @@ -179,141 +188,163 @@ def get_dF_u(no, V, C, Q, u): def get_dmoe_F_u(C, dF): - de = [np.einsum("pi,qi,pq->i",C[s],C[s],dF[s]) for s in [0,1]] + de = [np.einsum("pi,qi,pq->i", C[s], C[s], dF[s]) for s in [0, 1]] return de def get_full_u_F_u(no, C, moe, dF, u): - n = [C[s].shape[0] for s in [0,1]] - nv = [n[s] - no[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] - eo = [moe[s][:no[s]] for s in [0,1]] - ev = [moe[s][no[s]:] for s in [0,1]] + n = [C[s].shape[0] for s in [0, 1]] + nv = [n[s] - no[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] + eo = [moe[s][: no[s]] for s in [0, 1]] + ev = [moe[s][no[s] :] for s in [0, 1]] U = [None] * 2 - for s in [0,1]: - uov = u[s].reshape(no[s],nv[s]) - Dij = -eo[s].reshape(-1,1) + eo[s] + for s in [0, 1]: + uov = u[s].reshape(no[s], nv[s]) + Dij = -eo[s].reshape(-1, 1) + eo[s] np.fill_diagonal(Dij, 1) - dUoo = (Co[s].T@dF[s]@Co[s]) / Dij - np.fill_diagonal(dUoo, 0.) - Dab = -ev[s].reshape(-1,1) + ev[s] + dUoo = (Co[s].T @ dF[s] @ Co[s]) / Dij + np.fill_diagonal(dUoo, 0.0) + Dab = -ev[s].reshape(-1, 1) + ev[s] np.fill_diagonal(Dab, 1) - dUvv = (Cv[s].T@dF[s]@Cv[s]) / Dab - np.fill_diagonal(dUvv, 0.) + dUvv = (Cv[s].T @ dF[s] @ Cv[s]) / Dab + np.fill_diagonal(dUvv, 0.0) U[s] = np.block([[dUoo, uov], [-uov.T, dUvv]]) return U def get_dVovov_u(no, V, C, u): - n = [C[s].shape[0] for s in [0,1]] - nv = [n[s] - no[s] for s in [0,1]] - nov = [no[s] * nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] + n = [C[s].shape[0] for s in [0, 1]] + nv = [n[s] - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] dCo = [None] * 2 dCv = [None] * 2 dVovov = [None] * 3 - for s in [0,1]: + for s in [0, 1]: dVovov[s] = get_dVovov_r(no[s], V[s], C[s], u[s]) - if u[s].size == no[s]*nv[s]: - uov = u[s].reshape(no[s],nv[s]) + if u[s].size == no[s] * nv[s]: + uov = u[s].reshape(no[s], nv[s]) dCv[s] = Co[s] @ uov dCo[s] = -Cv[s] @ uov.T else: dC = C[s] @ u[s] - dCo[s] = dC[:,:no[s]] - dCv[s] = dC[:,no[s]:] + dCo[s] = dC[:, : no[s]] + dCv[s] = dC[:, no[s] :] dVovov[2] = ( - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[2],dCo[0],Cv[0],Co[1],Cv[1], - optimize=True) + - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[2],Co[0],dCv[0],Co[1],Cv[1], - optimize=True) + - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[2],Co[0],Cv[0],dCo[1],Cv[1], - optimize=True) + - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[2],Co[0],Cv[0],Co[1],dCv[1], - optimize=True) + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[2], dCo[0], Cv[0], Co[1], Cv[1], optimize=True + ) + + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[2], Co[0], dCv[0], Co[1], Cv[1], optimize=True + ) + + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[2], Co[0], Cv[0], dCo[1], Cv[1], optimize=True + ) + + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[2], Co[0], Cv[0], Co[1], dCv[1], optimize=True + ) ) return dVovov def get_Pmp2_u(t2l, t2r): - assert(len(t2l) == len(t2r) == 3) # aa,bb,ab - assert(t2l[0].ndim == t2r[0].ndim == 4) # ovov - no = [t2l[s].shape[0] for s in [0,1]] - es_pattern = ["iajb,majb->im","iajb,iamb->jm"] - Poo = [- ( - np.einsum(es_pattern[0], t2l[s], t2r[s]-t2r[s].transpose(0,3,2,1), - optimize=True) + - np.einsum(es_pattern[s], t2l[2], t2r[2], optimize=True)) - for s in [0,1] + assert len(t2l) == len(t2r) == 3 # aa,bb,ab + assert t2l[0].ndim == t2r[0].ndim == 4 # ovov + no = [t2l[s].shape[0] for s in [0, 1]] + es_pattern = ["iajb,majb->im", "iajb,iamb->jm"] + Poo = [ + -( + np.einsum( + es_pattern[0], + t2l[s], + t2r[s] - t2r[s].transpose(0, 3, 2, 1), + optimize=True, + ) + + np.einsum(es_pattern[s], t2l[2], t2r[2], optimize=True) + ) + for s in [0, 1] ] - es_pattern = ["iajb,icjb->ac","iajb,iajc->bc"] - Pvv = [( - np.einsum(es_pattern[0], t2l[s], t2r[s]-t2r[s].transpose(0,3,2,1), - optimize=True) + - np.einsum(es_pattern[s], t2l[2], t2r[2], optimize=True)) - for s in [0,1] + es_pattern = ["iajb,icjb->ac", "iajb,iajc->bc"] + Pvv = [ + ( + np.einsum( + es_pattern[0], + t2l[s], + t2r[s] - t2r[s].transpose(0, 3, 2, 1), + optimize=True, + ) + + np.einsum(es_pattern[s], t2l[2], t2r[2], optimize=True) + ) + for s in [0, 1] ] - return [slg.block_diag(Poo[s],Pvv[s]) for s in [0,1]] + return [slg.block_diag(Poo[s], Pvv[s]) for s in [0, 1]] def get_dPmp2_batch_u(C, moe, V, no, Qs, aorep=True): - """ no = [noa, nob] - V = [Vaa, Vbb, Vab] - C = [Ca, Cb] - moe = [moea, moeb] + """no = [noa, nob] + V = [Vaa, Vbb, Vab] + C = [Ca, Cb] + moe = [moea, moeb] """ - n = [C[s].shape[0] for s in [0,1]] - nv = [n[s] - no[s] for s in [0,1]] - nov = [no[s]*nv[s] for s in [0,1]] - Co = [C[s][:,:no[s]] for s in [0,1]] - Cv = [C[s][:,no[s]:] for s in [0,1]] + n = [C[s].shape[0] for s in [0, 1]] + nv = [n[s] - no[s] for s in [0, 1]] + nov = [no[s] * nv[s] for s in [0, 1]] + Co = [C[s][:, : no[s]] for s in [0, 1]] + Cv = [C[s][:, no[s] :] for s in [0, 1]] Vovov = [ - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[0],Co[0],Cv[0],Co[0],Cv[0], - optimize=True), - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[1],Co[1],Cv[1],Co[1],Cv[1], - optimize=True), - np.einsum("pqrs,pi,qa,rj,sb->iajb",V[2],Co[0],Cv[0],Co[1],Cv[1], - optimize=True) + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[0], Co[0], Cv[0], Co[0], Cv[0], optimize=True + ), + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[1], Co[1], Cv[1], Co[1], Cv[1], optimize=True + ), + np.einsum( + "pqrs,pi,qa,rj,sb->iajb", V[2], Co[0], Cv[0], Co[1], Cv[1], optimize=True + ), ] Diajb = get_Diajb_u(moe, no) t2 = [Vovov[s] / Diajb[s] for s in range(3)] - from .cphf_utils import (get_cpuhf_u_batch as cphf_kernel) + from .cphf_utils import get_cpuhf_u_batch as cphf_kernel + us = cphf_kernel(C, moe, V, no, Qs) nQ = len(Qs) dPs = [None] * nQ - Phf = [np.diag([1 if i < no[s] else 0 for i in range(n[s])]) for s in [0,1]] + Phf = [np.diag([1 if i < no[s] else 0 for i in range(n[s])]) for s in [0, 1]] iQ = 0 - for u_,Q in zip(us,Qs): - u = [u_[:nov[0]], u_[nov[0]:]] + for u_, Q in zip(us, Qs): + u = [u_[: nov[0]], u_[nov[0] :]] dF = get_dF_u(no, V, C, Q, u) dmoe = get_dmoe_F_u(C, dF) dDiajb = get_Diajb_u(dmoe, no) U = get_full_u_F_u(no, C, moe, dF, u) dVovov = get_dVovov_u(no, V, C, U) - dt2 = [(dVovov[s] - t2[s]*dDiajb[s]) / Diajb[s] for s in range(3)] + dt2 = [(dVovov[s] - t2[s] * dDiajb[s]) / Diajb[s] for s in range(3)] - P = get_Pmp2_u(t2,t2) - for s in [0,1]: P[s] += Phf[s] - dP = [U[s] @ P[s] - P[s] @ U[s] for s in [0,1]] + P = get_Pmp2_u(t2, t2) + for s in [0, 1]: + P[s] += Phf[s] + dP = [U[s] @ P[s] - P[s] @ U[s] for s in [0, 1]] - dP2 = get_Pmp2_u(dt2,t2) - for s in [0,1]: dP2[s] += dP2[s].T + dP2 = get_Pmp2_u(dt2, t2) + for s in [0, 1]: + dP2[s] += dP2[s].T - dPs[iQ] = [dP[s] + dP2[s] for s in [0,1]] + dPs[iQ] = [dP[s] + dP2[s] for s in [0, 1]] iQ += 1 if aorep: for iQ in range(nQ): - for s in [0,1]: + for s in [0, 1]: dPs[iQ][s] = C[s] @ dPs[iQ][s] @ C[s].T return dPs diff --git a/molbe/external/jac_utils.py b/molbe/external/jac_utils.py index 8f3095a1..02f07e9b 100644 --- a/molbe/external/jac_utils.py +++ b/molbe/external/jac_utils.py @@ -12,31 +12,34 @@ t_ia = ((2*t2-t2)_ibjc g_cjba - g_ikbj (2*t2-t2)_jbka) / (e_i - e_a) This approximate t1 is by substituting the MP2 t2 amplitudes into the CCSD equation and run it for one cycle. """ + + def get_t1(no, nv, moe, Vovov, Voovo, Vvovv): - assert(moe.size == no+nv) - assert(Vovov.shape == (no,nv,no,nv)) - assert(Voovo.shape == (no,no,nv,no)) - assert(Vvovv.shape == (nv,no,nv,nv)) - eia = moe[:no].reshape(-1,1) - moe[no:] + assert moe.size == no + nv + assert Vovov.shape == (no, nv, no, nv) + assert Voovo.shape == (no, no, nv, no) + assert Vvovv.shape == (nv, no, nv, nv) + eia = moe[:no].reshape(-1, 1) - moe[no:] eia_ = eia.ravel() - eiajb = (eia_.reshape(-1,1) + eia_).reshape(no,nv,no,nv) + eiajb = (eia_.reshape(-1, 1) + eia_).reshape(no, nv, no, nv) t2 = Vovov / eiajb t1approx = ( - 2.*np.einsum("ibjc,cjba->ia", t2, Vvovv, optimize=True) - - np.einsum("jbic,cjba->ia", t2, Vvovv, optimize=True) - - 2.*np.einsum("ikbj,jbka->ia", Voovo, t2, optimize=True) + - np.einsum("ikbj,kbja->ia", Voovo, t2, optimize=True)) / eia + 2.0 * np.einsum("ibjc,cjba->ia", t2, Vvovv, optimize=True) + - np.einsum("jbic,cjba->ia", t2, Vvovv, optimize=True) + - 2.0 * np.einsum("ikbj,jbka->ia", Voovo, t2, optimize=True) + + np.einsum("ikbj,kbja->ia", Voovo, t2, optimize=True) + ) / eia return t1approx def get_Vmogen_r(no, V, C, pattern): - assert(set(pattern) == set("ov")) + assert set(pattern) == set("ov") n = C.shape[0] nv = n - no nov = {"o": no, "v": nv} - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] Cov = {"o": Co, "v": Cv} Cs = [] @@ -45,25 +48,25 @@ def get_Vmogen_r(no, V, C, pattern): Cs += [Cov[p]] shape += [nov[p]] - return ao2mo.incore.general(V,Cs,compact=False).reshape(*shape) + return ao2mo.incore.general(V, Cs, compact=False).reshape(*shape) def get_dVmogen_r(no, V, C, u, pattern): - assert(set(pattern) == set("ov")) + assert set(pattern) == set("ov") n = C.shape[0] nv = n - no nov = {"o": no, "v": nv} - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] Cov = {"o": Co, "v": Cv} - if u.size == no*nv: - uov = u.reshape(no,nv) + if u.size == no * nv: + uov = u.reshape(no, nv) dCv = Co @ uov dCo = -Cv @ uov.T else: dC = C @ u - dCo = dC[:,:no] - dCv = dC[:,no:] + dCo = dC[:, :no] + dCv = dC[:, no:] dCov = {"o": dCo, "v": dCv} # take care of symmetry p12 = pattern[:2] @@ -71,16 +74,16 @@ def get_dVmogen_r(no, V, C, u, pattern): def xform_1_index(ip0): Cs = [] - for ip,p in enumerate(pattern): + for ip, p in enumerate(pattern): Cs += [dCov[p]] if ip == ip0 else [Cov[p]] shape = tuple([nov[p] for p in pattern]) - return ao2mo.incore.general(V,Cs,compact=False).reshape(*shape) + return ao2mo.incore.general(V, Cs, compact=False).reshape(*shape) # xform ij dV = xform_1_index(0) + xform_1_index(1) # xform kl if p12 == p34: - dV += dV.transpose(2,3,0,1) + dV += dV.transpose(2, 3, 0, 1) else: dV += xform_1_index(2) + xform_1_index(3) @@ -90,14 +93,15 @@ def xform_1_index(ip0): def get_dt1ao_an(no, V, C, moe, Qs, us=None): n = C.shape[0] nv = n - no - Co = C[:,:no] - Cv = C[:,no:] + Co = C[:, :no] + Cv = C[:, no:] def get_Dia_r(moe_, no_): - return moe_[:no_].reshape(-1,1) - moe_[no_:] + return moe_[:no_].reshape(-1, 1) - moe_[no_:] # prepare integrals from .cpmp2_utils import get_Diajb_r + Vovov = get_Vmogen_r(no, V, C, "ovov") Vvovv = get_Vmogen_r(no, V, C, "vovv") Voovo = get_Vmogen_r(no, V, C, "oovo") @@ -109,21 +113,24 @@ def get_Dia_r(moe_, no_): # solve CPHF get u if us is None: from .cphf_utils import cphf_kernel_batch + us = cphf_kernel_batch(C, moe, V, no, Qs) dt1s = [None] * len(Qs) iu = 0 - for u,Q in zip(us,Qs): + for u, Q in zip(us, Qs): # get A from .cpmp2_utils import get_dF_r + A = -get_dF_r(no, V, C, Q, u) Aoo = Co.T @ A @ Co Avv = Cv.T @ A @ Cv # get tA - tA = np.einsum("lajb,li->iajb",t2,Aoo,optimize=True) -\ - np.einsum("idjb,da->iajb",t2,Avv,optimize=True) - tA += tA.transpose(2,3,0,1) + tA = np.einsum("lajb,li->iajb", t2, Aoo, optimize=True) - np.einsum( + "idjb,da->iajb", t2, Avv, optimize=True + ) + tA += tA.transpose(2, 3, 0, 1) # get VU dVovov = get_dVmogen_r(no, V, C, u, "ovov") @@ -132,19 +139,26 @@ def get_Dia_r(moe_, no_): # get dmoe and deov from .cpmp2_utils import get_dmoe_F_r + dmoe = get_dmoe_F_r(C, -A) deov = get_Dia_r(dmoe, no) # get dCov - uov = u.reshape(no,nv) + uov = u.reshape(no, nv) dCo = -Cv @ uov.T dCv = Co @ uov # show time - dt1 = Co @ (get_t1(no, nv, moe, tA, Voovo, Vvovv) + - get_t1(no, nv, moe, dVovov, Voovo, Vvovv) + - get_t1(no, nv, moe, Vovov, dVoovo, dVvovv) + - (Aoo @ t1 - t1 @ Avv) / eov) @ Cv.T + dt1 = ( + Co + @ ( + get_t1(no, nv, moe, tA, Voovo, Vvovv) + + get_t1(no, nv, moe, dVovov, Voovo, Vvovv) + + get_t1(no, nv, moe, Vovov, dVoovo, dVvovv) + + (Aoo @ t1 - t1 @ Avv) / eov + ) + @ Cv.T + ) dt1 += dCo @ t1 @ Cv.T + Co @ t1 @ dCv.T dt1 += dt1.T @@ -158,17 +172,18 @@ def get_Dia_r(moe_, no_): def get_dPccsdurlx_batch_u(C, moe, eri, no, vpots): # cphf from .cphf_utils import cphf_kernel_batch + us = cphf_kernel_batch(C, moe, eri, no, vpots) # get mp2 part dPt1aos = get_dt1ao_an(no, eri, C, moe, vpots, us=us) # HF contribution n = C.shape[0] nv = n - no - for iu,u in enumerate(us): - Co = C[:,:no] - Cv = C[:,no:] - dCo = -Cv @ u.reshape(no,nv).T - dPhfao = 2. * dCo @ Co.T + for iu, u in enumerate(us): + Co = C[:, :no] + Cv = C[:, no:] + dCo = -Cv @ u.reshape(no, nv).T + dPhfao = 2.0 * dCo @ Co.T dPhfao += dPhfao.T dPt1aos[iu] += dPhfao # total diff --git a/molbe/external/lo_helper.py b/molbe/external/lo_helper.py index e701f0b4..460373b8 100644 --- a/molbe/external/lo_helper.py +++ b/molbe/external/lo_helper.py @@ -5,11 +5,12 @@ # from pyscf import lib -import numpy,sys +import numpy, sys from copy import deepcopy from functools import reduce -def get_symm_mat_pow(A, p, check_symm=True, thresh=1.E-8): + +def get_symm_mat_pow(A, p, check_symm=True, thresh=1.0e-8): """A ** p where A is symmetric Note: @@ -19,7 +20,7 @@ def get_symm_mat_pow(A, p, check_symm=True, thresh=1.E-8): return numpy.linalg.matrix_power(A, int(p)) if check_symm: - assert(numpy.linalg.norm(A-A.conj().T) < thresh) + assert numpy.linalg.norm(A - A.conj().T) < thresh e, u = numpy.linalg.eigh(A) Ap = u @ numpy.diag(e**p) @ u.conj().T @@ -29,27 +30,24 @@ def get_symm_mat_pow(A, p, check_symm=True, thresh=1.E-8): def get_aoind_by_atom(mol, atomind_by_motif=None): import numpy as np + natom = mol.natm aoslice_by_atom = mol.aoslice_by_atom() - aoshift_by_atom = [0]+[aoslice_by_atom[ia][-1] - for ia in range(natom)] + aoshift_by_atom = [0] + [aoslice_by_atom[ia][-1] for ia in range(natom)] # if motif info is provided, group lo by motif if atomind_by_motif is None: - - aoind_by_atom = [list(range(*aoshift_by_atom[ia:ia+2])) - for ia in range(natom)] + aoind_by_atom = [ + list(range(*aoshift_by_atom[ia : ia + 2])) for ia in range(natom) + ] else: - nmotif = len(atomind_by_motif) - assert( - set([ia for im in range(nmotif) - for ia in atomind_by_motif[im]]) == set(range(natom)) + assert set([ia for im in range(nmotif) for ia in atomind_by_motif[im]]) == set( + range(natom) ) aoind_by_atom = [[] for im in range(nmotif)] for im in range(nmotif): for ia in atomind_by_motif[im]: - aoind_by_atom[im] += list( - range(*aoshift_by_atom[ia:ia+2])) + aoind_by_atom[im] += list(range(*aoshift_by_atom[ia : ia + 2])) return aoind_by_atom @@ -68,15 +66,15 @@ def reorder_by_atom_(Clo, aoind_by_atom, S, thr=0.5): loshift = 0 for ia in range(natom): ra = aoind_by_atom[ia] - poplo_by_atom = np.sum(Clo_soao[ra]**2., axis=0) - loind_a = np.where(poplo_by_atom>thr)[0].tolist() + poplo_by_atom = np.sum(Clo_soao[ra] ** 2.0, axis=0) + loind_a = np.where(poplo_by_atom > thr)[0].tolist() loind_reorder += loind_a nlo_a = len(loind_a) - loind_by_atom[ia] = list(range(loshift,loshift+nlo_a)) + loind_by_atom[ia] = list(range(loshift, loshift + nlo_a)) loshift += nlo_a if loind_reorder != list(range(nlo)): - print('REORDERD') - Clo_new = Clo[:,loind_reorder] + print("REORDERD") + Clo_new = Clo[:, loind_reorder] else: Clo_new = Clo return Clo_new, loind_by_atom diff --git a/molbe/external/optqn.py b/molbe/external/optqn.py index 56227170..46820b45 100644 --- a/molbe/external/optqn.py +++ b/molbe/external/optqn.py @@ -5,18 +5,17 @@ # The code has been slightly modified. # -import numpy,sys, h5py +import numpy, sys, h5py from .. import be_var def line_search_LF(func, xold, fold, dx, iter_): - """Adapted from D.-H. Li and M. Fukushima, Optimization Metheods and Software, 13, 181 (2000) - """ + """Adapted from D.-H. Li and M. Fukushima, Optimization Metheods and Software, 13, 181 (2000)""" beta = 0.1 rho = 0.9 - sigma1 = 1E-3 - sigma2 = 1E-3 - eta = (iter_+1)**-2. + sigma1 = 1e-3 + sigma2 = 1e-3 + eta = (iter_ + 1) ** -2.0 xk = xold + dx lcout = 0 @@ -27,15 +26,13 @@ def line_search_LF(func, xold, fold, dx, iter_): norm_dx = numpy.linalg.norm(dx) norm_fk = numpy.linalg.norm(fk) norm_fold = numpy.linalg.norm(fold) - alp = 1. - - if norm_fk > rho*norm_fold - sigma2*norm_dx**2.: - while norm_fk > (1.+eta)*norm_fold - sigma1*alp**2.*norm_dx**2.: + alp = 1.0 + if norm_fk > rho * norm_fold - sigma2 * norm_dx**2.0: + while norm_fk > (1.0 + eta) * norm_fold - sigma1 * alp**2.0 * norm_dx**2.0: alp *= beta xk = xold + alp * dx - fk = func(xk) lcout += 1 @@ -43,11 +40,12 @@ def line_search_LF(func, xold, fold, dx, iter_): if lcout == 20: break - print(' No. of line search steps in QN opt :', lcout, flush=True) + print(" No. of line search steps in QN opt :", lcout, flush=True) print(flush=True) return alp, xk, fk -def trustRegion(func, xold, fold, Binv, c = 0.5): + +def trustRegion(func, xold, fold, Binv, c=0.5): """Perform Trust Region Optimization. See "A Broyden Trust Region Quasi-Newton Method for Nonlinear Equations" (https://www.iaeng.org/IJCS/issues_v46/issue_3/IJCS_46_3_09.pdf) Algorithm 1 for more information @@ -71,55 +69,82 @@ def trustRegion(func, xold, fold, Binv, c = 0.5): x_{p+1} and f_{p+1}. These values are used to proceed with Broyden's Method. """ # c := initial trust radius (trust_radius = c^p) - microiter = 0 # p - rho = 0.001 # Threshold for trust region subproblem - ratio = 0 # Initial r - B = numpy.linalg.inv(Binv) # approx Jacobian - #dx_gn = - Binv@fold - dx_gn = -(Binv@Binv.T)@B.T@fold - dx_sd = - B.T@fold # Steepest Descent step - t = numpy.linalg.norm(dx_sd)**2 / numpy.linalg.norm(B@dx_sd)**2 + microiter = 0 # p + rho = 0.001 # Threshold for trust region subproblem + ratio = 0 # Initial r + B = numpy.linalg.inv(Binv) # approx Jacobian + # dx_gn = - Binv@fold + dx_gn = -(Binv @ Binv.T) @ B.T @ fold + dx_sd = -B.T @ fold # Steepest Descent step + t = numpy.linalg.norm(dx_sd) ** 2 / numpy.linalg.norm(B @ dx_sd) ** 2 prevdx = None - while (ratio < rho or ared < 0.): + while ratio < rho or ared < 0.0: # Trust Region subproblem # minimize (1/2) ||F_k + B_k d||^2 w.r.t. d, s.t. d w/i trust radius # to pick the optimal direction using dog leg method - if numpy.linalg.norm(dx_gn) < max(1.0, numpy.linalg.norm(xold)) * (c ** microiter): # Gauss-Newton step within the trust radius - print(' Trust Region Optimization Step ', microiter, ': Gauss-Newton', flush=True) + if numpy.linalg.norm(dx_gn) < max(1.0, numpy.linalg.norm(xold)) * ( + c**microiter + ): # Gauss-Newton step within the trust radius + print( + " Trust Region Optimization Step ", + microiter, + ": Gauss-Newton", + flush=True, + ) dx = dx_gn - elif t * numpy.linalg.norm(dx_sd) > max(1.0, numpy.linalg.norm(xold)) * (c ** microiter): # GN step outside, SD step also outside - print(' Trust Region Optimization Step ', microiter, ': Steepest Descent', flush=True) - dx = (c ** microiter) / numpy.linalg.norm(dx_sd) * dx_sd - else: # GN step outside, SD step inside (dog leg step) - # dx := t*dx_sd + s (dx_gn - t*dx_sd) s.t. ||dx|| = c^p - print(' Trust Region Optimization Step ', microiter, ': Dog Leg', flush=True) - tdx_sd = t*dx_sd + elif t * numpy.linalg.norm(dx_sd) > max(1.0, numpy.linalg.norm(xold)) * ( + c**microiter + ): # GN step outside, SD step also outside + print( + " Trust Region Optimization Step ", + microiter, + ": Steepest Descent", + flush=True, + ) + dx = (c**microiter) / numpy.linalg.norm(dx_sd) * dx_sd + else: # GN step outside, SD step inside (dog leg step) + # dx := t*dx_sd + s (dx_gn - t*dx_sd) s.t. ||dx|| = c^p + print( + " Trust Region Optimization Step ", microiter, ": Dog Leg", flush=True + ) + tdx_sd = t * dx_sd diff = dx_gn - tdx_sd - #s = (-dx_sd.T@diff + numpy.sqrt((dx_sd.T@diff)**2 - numpy.linalg.norm(diff)**2*(numpy.linalg.norm(dx_sd)**2-(c ** microiter)**2))) / (numpy.linalg.norm(dx_sd))**2 + # s = (-dx_sd.T@diff + numpy.sqrt((dx_sd.T@diff)**2 - numpy.linalg.norm(diff)**2*(numpy.linalg.norm(dx_sd)**2-(c ** microiter)**2))) / (numpy.linalg.norm(dx_sd))**2 # s is largest value in [0, 1] s.t. ||dx|| \le trust radius s = 1 - dx = tdx_sd + s*diff - while (numpy.linalg.norm(dx) > c ** microiter and s > 0): + dx = tdx_sd + s * diff + while numpy.linalg.norm(dx) > c**microiter and s > 0: s -= 0.001 - dx = tdx_sd + s*diff + dx = tdx_sd + s * diff if prevdx is None or not numpy.all(dx == prevdx): # Actual Reduction := f(x_k) - f(x_k + dx) fnew = func(xold + dx) - ared = 0.5 * (numpy.linalg.norm(fold)**2 - numpy.linalg.norm(fnew)**2) + ared = 0.5 * (numpy.linalg.norm(fold) ** 2 - numpy.linalg.norm(fnew) ** 2) # Predicted Reduction := q(0) - q(dx) where q = (1/2) ||F_k + B_k d||^2 - pred = 0.5 * (numpy.linalg.norm(fold)**2 - numpy.linalg.norm(fold + B@dx)**2) + pred = 0.5 * ( + numpy.linalg.norm(fold) ** 2 - numpy.linalg.norm(fold + B @ dx) ** 2 + ) # Trust Region convergence criteria # r = ared/pred \le rho ratio = ared / pred microiter += 1 if prevdx is None or not numpy.all(dx == prevdx) and be_var.PRINT_LEVEL > 2: - print(' ||δx||: ', numpy.linalg.norm(dx), flush=True) - print(' Reduction Ratio (Actual / Predicted): ', ared, '/', pred, '=', ratio, flush=True) + print(" ||δx||: ", numpy.linalg.norm(dx), flush=True) + print( + " Reduction Ratio (Actual / Predicted): ", + ared, + "/", + pred, + "=", + ratio, + flush=True, + ) prevdx = dx - return xold + dx, fnew # xnew + return xold + dx, fnew # xnew + class FrankQN: - """ Quasi Newton Optimization + """Quasi Newton Optimization Performs quasi newton optimization. Interfaces many functionalities of the frankestein code originaly written by Hong-Zhou Ye @@ -128,7 +153,6 @@ class FrankQN: """ def __init__(self, func, x0, f0, J0, trust=0.5, max_space=500): - self.x0 = x0 self.n = x0.size self.f0 = f0 @@ -138,21 +162,20 @@ def __init__(self, func, x0, f0, J0, trust=0.5, max_space=500): self.iter_ = 0 - self.tol_gmres = 1.e-6 - self.xnew = None # new errvec - self.xold = None # old errvec - self.fnew = None # new jacobian? - self.fold = None # old jacobian? + self.tol_gmres = 1.0e-6 + self.xnew = None # new errvec + self.xold = None # old errvec + self.fnew = None # new jacobian? + self.fold = None # old jacobian? self.max_subspace = max_space - self.dxs = numpy.empty([self.max_subspace,self.n]) - self.fs = numpy.empty([self.max_subspace,self.n]) - self.us = numpy.empty([self.max_subspace,self.n]) # u_m = B_m @ f_m - self.vs = numpy.empty([self.max_subspace,self.n]) # v_m = B_0 @ f_{m+1} + self.dxs = numpy.empty([self.max_subspace, self.n]) + self.fs = numpy.empty([self.max_subspace, self.n]) + self.us = numpy.empty([self.max_subspace, self.n]) # u_m = B_m @ f_m + self.vs = numpy.empty([self.max_subspace, self.n]) # v_m = B_0 @ f_{m+1} self.B = None self.trust = trust def next_step(self, trust_region=False): - if self.iter_ == 0: self.xnew = self.x0 self.fnew = self.func(self.xnew) if self.f0 is None else self.f0 @@ -168,75 +191,79 @@ def next_step(self, trust_region=False): self.xold = self.xnew.copy() self.fold = self.fnew.copy() - if not self.iter_ ==0: - - tmp__ = numpy.outer(dx_i - self.Binv@df_i, dx_i@self.Binv)/\ - (dx_i@self.Binv@df_i) + if not self.iter_ == 0: + tmp__ = numpy.outer(dx_i - self.Binv @ df_i, dx_i @ self.Binv) / ( + dx_i @ self.Binv @ df_i + ) self.Binv += tmp__ - us_tmp = self.Binv@self.fnew + us_tmp = self.Binv @ self.fnew if trust_region: - self.xnew, self.fnew = trustRegion(self.func, self.xold, self.fold, self.Binv, c = self.trust) + self.xnew, self.fnew = trustRegion( + self.func, self.xold, self.fold, self.Binv, c=self.trust + ) else: self.us[self.iter_] = self.get_Bnfn(self.iter_) alp, self.xnew, self.fnew = line_search_LF( - self.func, self.xold, self.fold, -self.us[self.iter_], self.iter_) + self.func, self.xold, self.fold, -self.us[self.iter_], self.iter_ + ) # udpate vs, dxs, and fs self.vs[self.iter_] = numpy.dot(self.B0, self.fnew) self.dxs[self.iter_] = self.xnew - self.xold - self.fs[self.iter_+1] = self.fnew.copy() + self.fs[self.iter_ + 1] = self.fnew.copy() self.iter_ += 1 def get_Bnfn(self, n): - # self.us; self.dxs; self.vs - if n == 0: return self.us[0] + if n == 0: + return self.us[0] vs = [None] * n - for i in range(n): vs[i] = self.vs[n-i-1] - for i in range(1,n+1): - un_ = self.us[i-1] - dxn_ = self.dxs[i-1] - vps = [None] * (n-i+1) - for j in range(n-i+1): + for i in range(n): + vs[i] = self.vs[n - i - 1] + for i in range(1, n + 1): + un_ = self.us[i - 1] + dxn_ = self.dxs[i - 1] + vps = [None] * (n - i + 1) + for j in range(n - i + 1): a = vs[j] - b = vs[n-i] - un_ + b = vs[n - i] - un_ - vps[j] = a + (dxn_@a)/(dxn_@b) * (dxn_ - b) + vps[j] = a + (dxn_ @ a) / (dxn_ @ b) * (dxn_ - b) vs = vps return vs[0] -def get_be_error_jacobian(self,jac_solver='HF'): - Jes = [None] * self.Nfrag - Jcs = [None] * self.Nfrag - xes = [None] * self.Nfrag - xcs = [None] * self.Nfrag - ys = [None] * self.Nfrag +def get_be_error_jacobian(self, jac_solver="HF"): + Jes = [None] * self.Nfrag + Jcs = [None] * self.Nfrag + xes = [None] * self.Nfrag + xcs = [None] * self.Nfrag + ys = [None] * self.Nfrag alphas = [None] * self.Nfrag - if jac_solver =='MP2': + if jac_solver == "MP2": res_func = mp2res_func - elif jac_solver =='CCSD': + elif jac_solver == "CCSD": res_func = ccsdres_func - elif jac_solver=='HF': + elif jac_solver == "HF": res_func = hfres_func Ncout = [None] * self.Nfrag for A in range(self.Nfrag): - - Jes[A], Jcs[A], xes[A], xcs[A], ys[A], alphas[A], Ncout[A] = \ + Jes[A], Jcs[A], xes[A], xcs[A], ys[A], alphas[A], Ncout[A] = ( get_atbe_Jblock_frag(self.Fobjs[A], res_func) + ) alpha = sum(alphas) # build Jacobian - ''' ignore! + """ ignore! F0-M1 F1-M2M2 F2-M3M3 F3-M4 M1 M2 M2 M3 M3 M4 M1 E0 C1-1 @@ -245,55 +272,57 @@ def get_be_error_jacobian(self,jac_solver='HF'): M3 C1-1 E2 E2 M3 E2 E2 C3-3 M4 C2-2 E3 - ''' + """ N_ = sum(Ncout) - J = numpy.zeros((N_+1, N_+1)) + J = numpy.zeros((N_ + 1, N_ + 1)) cout = 0 for findx, fobj in enumerate(self.Fobjs): - J[cout:Ncout[findx]+cout, cout:Ncout[findx]+cout] = Jes[findx] - J[cout:Ncout[findx]+cout, N_:] = numpy.array(xes[findx]).reshape(-1,1) - J[N_:, cout:Ncout[findx]+cout] = ys[findx] + J[cout : Ncout[findx] + cout, cout : Ncout[findx] + cout] = Jes[findx] + J[cout : Ncout[findx] + cout, N_:] = numpy.array(xes[findx]).reshape(-1, 1) + J[N_:, cout : Ncout[findx] + cout] = ys[findx] coutc = 0 coutc_ = 0 for cindx, cens in enumerate(fobj.center_idx): - coutc += Jcs[fobj.center[cindx]].shape[0] - start_ = sum(Ncout[:fobj.center[cindx]]) + coutc += Jcs[fobj.center[cindx]].shape[0] + start_ = sum(Ncout[: fobj.center[cindx]]) end_ = start_ + Ncout[fobj.center[cindx]] - J[cout+coutc_: cout+coutc, start_:end_] += Jcs[fobj.center[cindx]] - J[cout+coutc_: cout+coutc, N_:] += numpy.array( - xcs[fobj.center[cindx]]).reshape(-1,1) + J[cout + coutc_ : cout + coutc, start_:end_] += Jcs[fobj.center[cindx]] + J[cout + coutc_ : cout + coutc, N_:] += numpy.array( + xcs[fobj.center[cindx]] + ).reshape(-1, 1) coutc_ = coutc cout += Ncout[findx] - J[N_:,N_:] = alpha + J[N_:, N_:] = alpha return J + def get_atbe_Jblock_frag(fobj, res_func): from molbe.helper import get_scfObj, get_eri - vpots = get_vpots_frag(fobj.nao, fobj.edge_idx, - fobj.fsites) + vpots = get_vpots_frag(fobj.nao, fobj.edge_idx, fobj.fsites) eri_ = get_eri(fobj.dname, fobj.nao, eri_file=fobj.eri_file) - dm0 = numpy.dot( fobj._mo_coeffs[:,:fobj.nsocc], - fobj._mo_coeffs[:,:fobj.nsocc].T) *2. - mf_ = get_scfObj(fobj.fock+fobj.heff, eri_, fobj.nsocc, dm0=dm0) + dm0 = ( + numpy.dot(fobj._mo_coeffs[:, : fobj.nsocc], fobj._mo_coeffs[:, : fobj.nsocc].T) + * 2.0 + ) + mf_ = get_scfObj(fobj.fock + fobj.heff, eri_, fobj.nsocc, dm0=dm0) dPs, dP_mu = res_func(mf_, vpots, eri_, fobj.nsocc) Je = [] Jc = [] - y=[] + y = [] xe = [] xc = [] cout = 0 for edge in fobj.edge_idx: - for j_ in range(len(edge)): for k_ in range(len(edge)): - if j_>k_: + if j_ > k_: continue ## response w.r.t matching pot # edges @@ -304,27 +333,25 @@ def get_atbe_Jblock_frag(fobj, res_func): for j__ in range(lene): for k__ in range(lene): - if j__>k__: + if j__ > k__: continue - tmpje_.append(dPs[cout][edge_[j__],edge_[k__]]) - y_ = 0. + tmpje_.append(dPs[cout][edge_[j__], edge_[k__]]) + y_ = 0.0 for fidx, fval in enumerate(fobj.fsites): - - if not any(fidx in sublist - for sublist in fobj.edge_idx): + if not any(fidx in sublist for sublist in fobj.edge_idx): y_ += dPs[cout][fidx, fidx] y.append(y_) tmpjc_ = [] # center on the same fragment - #for cen in fobj.efac[1]: + # for cen in fobj.efac[1]: for j__ in fobj.centerf_idx: for k__ in fobj.centerf_idx: - if j__>k__: + if j__ > k__: continue - tmpjc_.append(-dPs[cout][j__,k__]) + tmpjc_.append(-dPs[cout][j__, k__]) Je.append(tmpje_) @@ -332,69 +359,70 @@ def get_atbe_Jblock_frag(fobj, res_func): ## response w.r.t. chem pot # edge - xe.append(dP_mu[edge[j_],edge[k_]]) + xe.append(dP_mu[edge[j_], edge[k_]]) cout += 1 Je = numpy.array(Je).T Jc = numpy.array(Jc).T - alpha = 0. + alpha = 0.0 for fidx, fval in enumerate(fobj.fsites): - if not any(fidx in sublist - for sublist in fobj.edge_idx): + if not any(fidx in sublist for sublist in fobj.edge_idx): alpha += dP_mu[fidx, fidx] for j__ in fobj.centerf_idx: for k__ in fobj.centerf_idx: - if j__>k__: + if j__ > k__: continue - xc.append(-dP_mu[j__,k__]) + xc.append(-dP_mu[j__, k__]) return Je, Jc, xe, xc, y, alpha, cout -def get_be_error_jacobian_selffrag(self,jac_solver='HF'): - - Jes = [None] * self.Nfrag - Jcs = [None] * self.Nfrag - xes = [None] * self.Nfrag - xcs = [None] * self.Nfrag - ys = [None] * self.Nfrag +def get_be_error_jacobian_selffrag(self, jac_solver="HF"): + Jes = [None] * self.Nfrag + Jcs = [None] * self.Nfrag + xes = [None] * self.Nfrag + xcs = [None] * self.Nfrag + ys = [None] * self.Nfrag alphas = [None] * self.Nfrag - if jac_solver =='MP2': + if jac_solver == "MP2": res_func = mp2res_func - elif jac_solver =='CCSD': + elif jac_solver == "CCSD": res_func = ccsdres_func - elif jac_solver=='HF': + elif jac_solver == "HF": res_func = hfres_func - Jes, Jcs, xes, xcs, ys, alphas, Ncout = \ - get_atbe_Jblock_frag(self.Fobjs[0], res_func) + Jes, Jcs, xes, xcs, ys, alphas, Ncout = get_atbe_Jblock_frag( + self.Fobjs[0], res_func + ) N_ = Ncout - J = numpy.zeros((N_+1, N_+1)) + J = numpy.zeros((N_ + 1, N_ + 1)) J[:Ncout, :Ncout] = Jes - J[:Ncout, N_:] = numpy.array(xes).reshape(-1,1) + J[:Ncout, N_:] = numpy.array(xes).reshape(-1, 1) J[N_:, :Ncout] = ys - J[:Ncout, N_:] += numpy.array([*xcs, *xcs]).reshape(-1,1) - J[N_:,N_:] = alphas + J[:Ncout, N_:] += numpy.array([*xcs, *xcs]).reshape(-1, 1) + J[N_:, N_:] = alphas return J def hfres_func(mf, vpots, eri, nsocc): - from molbe.external.cphf_utils import cphf_kernel_batch,get_rhf_dP_from_u + from molbe.external.cphf_utils import cphf_kernel_batch, get_rhf_dP_from_u + C = mf.mo_coeff moe = mf.mo_energy eri = mf._eri no = nsocc us = cphf_kernel_batch(C, moe, eri, no, vpots) - dPs = [get_rhf_dP_from_u(C, no, us[I]) for I in range(len(vpots)-1)] + dPs = [get_rhf_dP_from_u(C, no, us[I]) for I in range(len(vpots) - 1)] dP_mu = get_rhf_dP_from_u(C, no, us[-1]) return dPs, dP_mu + def mp2res_func(mf, vpots, eri, nsocc): from molbe.external.cpmp2_utils import get_dPmp2_batch_r @@ -404,7 +432,7 @@ def mp2res_func(mf, vpots, eri, nsocc): no = nsocc dPs_an = get_dPmp2_batch_r(C, moe, eri, no, vpots, aorep=True) - dPs_an = numpy.array([dp_*0.5 for dp_ in dPs_an]) + dPs_an = numpy.array([dp_ * 0.5 for dp_ in dPs_an]) dP_mu = dPs_an[-1] return dPs_an[:-1], dP_mu @@ -432,24 +460,19 @@ def get_vpots_frag(nao, edge_idx, fsites): lene = len(edge_) for j__ in range(lene): for k__ in range(lene): - if j__>k__: + if j__ > k__: continue - tmppot = numpy.zeros((nao,nao)) - tmppot[edge_[j__],edge_[k__]] = tmppot[edge_[k__],edge_[j__]] = 1 + tmppot = numpy.zeros((nao, nao)) + tmppot[edge_[j__], edge_[k__]] = tmppot[edge_[k__], edge_[j__]] = 1 vpots.append(tmppot) # only the centers # outer edges not included - tmppot = numpy.zeros((nao,nao)) + tmppot = numpy.zeros((nao, nao)) for fidx, fval in enumerate(fsites): - - if not any(fidx in sublist - for sublist in edge_idx): + if not any(fidx in sublist for sublist in edge_idx): tmppot[fidx, fidx] = -1 - vpots.append(tmppot) return vpots - - diff --git a/molbe/external/uccsd_eri.py b/molbe/external/uccsd_eri.py index 9c40ed5d..96547278 100644 --- a/molbe/external/uccsd_eri.py +++ b/molbe/external/uccsd_eri.py @@ -1,30 +1,42 @@ import numpy from pyscf.cc.uccsd import _make_eris_incore + def make_eris_incore(mycc, Vss, Vos, mo_coeff=None, ao2mofn=None, frozen=False): - vhf = frank_get_veff(mycc, mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ), Vss, Vos) + vhf = frank_get_veff( + mycc, mycc._scf.make_rdm1(mycc.mo_coeff, mycc.mo_occ), Vss, Vos + ) fockao = frank_get_fock(mycc, vhf, frozen) mycc._scf.get_veff = lambda *args, **kwargs: vhf mycc._scf.get_fock = lambda *args, **kwargs: fockao - return _make_eris_incore(mycc, mo_coeff=mo_coeff, ao2mofn=ao2mofn) #, frozen) + return _make_eris_incore(mycc, mo_coeff=mo_coeff, ao2mofn=ao2mofn) # , frozen) + def frank_get_veff(mycc, dm, Vss, Vos): - veffss = [numpy.einsum("pqrs,sr->pq", Vss[s], dm[s]) - - numpy.einsum("psrq,sr->pq", Vss[s], dm[s]) for s in [0,1]] - veffos = [numpy.einsum("pqrs,sr->pq", Vos, dm[1]), - numpy.einsum("pqrs,qp->rs", Vos, dm[0])] - veff = [veffss[s] + veffos[s] for s in [0,1]] + veffss = [ + numpy.einsum("pqrs,sr->pq", Vss[s], dm[s]) + - numpy.einsum("psrq,sr->pq", Vss[s], dm[s]) + for s in [0, 1] + ] + veffos = [ + numpy.einsum("pqrs,sr->pq", Vos, dm[1]), + numpy.einsum("pqrs,qp->rs", Vos, dm[0]), + ] + veff = [veffss[s] + veffos[s] for s in [0, 1]] return veff + def frank_get_fock(mycc, vhf, frozen): - if frozen==False: + if frozen == False: mycc._scf.full_gcore = None mycc._scf.full_hs = None - fock = [mycc._scf.h1[s]+mycc._scf.gcores_raw[s] for s in [0,1]] + fock = [mycc._scf.h1[s] + mycc._scf.gcores_raw[s] for s in [0, 1]] else: - mycc._scf.full_gcore = [mycc._scf.gcores_raw[s] - vhf[s] for s in [0,1]] - mycc._scf.full_hs = [mycc._scf.h1[s] + mycc._scf.full_gcore[s] + mycc._scf.core_veffs[s] for s in [0,1]] - fock = [mycc._scf.full_hs[s] + vhf[s] for s in [0,1]] + mycc._scf.full_gcore = [mycc._scf.gcores_raw[s] - vhf[s] for s in [0, 1]] + mycc._scf.full_hs = [ + mycc._scf.h1[s] + mycc._scf.full_gcore[s] + mycc._scf.core_veffs[s] + for s in [0, 1] + ] + fock = [mycc._scf.full_hs[s] + vhf[s] for s in [0, 1]] return fock - diff --git a/molbe/external/unrestricted_utils.py b/molbe/external/unrestricted_utils.py index 5a44e66a..ad995b44 100644 --- a/molbe/external/unrestricted_utils.py +++ b/molbe/external/unrestricted_utils.py @@ -4,8 +4,10 @@ import h5py import numpy, functools + def make_uhf_obj(fobj_a, fobj_b, frozen=False): from pyscf import scf + """ Constructs UHF object from the alpha and beta components """ @@ -24,36 +26,46 @@ def make_uhf_obj(fobj_a, fobj_b, frozen=False): full_uhf.h1 = (fobj_a.h1, fobj_b.h1) # Define the HF veff in alpha and beta spin Schmidt spaces - full_uhf.veff0_a = functools.reduce(numpy.dot,(fobj_a.TA.T,fobj_a.hf_veff,fobj_a.TA)) - full_uhf.veff0_b = functools.reduce(numpy.dot,(fobj_b.TA.T,fobj_b.hf_veff,fobj_b.TA)) + full_uhf.veff0_a = functools.reduce( + numpy.dot, (fobj_a.TA.T, fobj_a.hf_veff, fobj_a.TA) + ) + full_uhf.veff0_b = functools.reduce( + numpy.dot, (fobj_b.TA.T, fobj_b.hf_veff, fobj_b.TA) + ) # Build correct eris - Vs = uccsd_restore_eris((1,1,1), fobj_a, fobj_b) + Vs = uccsd_restore_eris((1, 1, 1), fobj_a, fobj_b) full_uhf.TA = [fobj_a.TA, fobj_b.TA] # Build core components if frozen: - full_uhf.gcores_raw = [fobj_a.TA.T @ (fobj_a.hf_veff-fobj_a.core_veff) @ fobj_a.TA, - fobj_b.TA.T @ (fobj_b.hf_veff-fobj_b.core_veff) @ fobj_b.TA] - full_uhf.core_veffs = [fobj_a.TA.T @ fobj_a.core_veff @ fobj_a.TA, - fobj_b.TA.T @ fobj_b.core_veff @ fobj_b.TA] + full_uhf.gcores_raw = [ + fobj_a.TA.T @ (fobj_a.hf_veff - fobj_a.core_veff) @ fobj_a.TA, + fobj_b.TA.T @ (fobj_b.hf_veff - fobj_b.core_veff) @ fobj_b.TA, + ] + full_uhf.core_veffs = [ + fobj_a.TA.T @ fobj_a.core_veff @ fobj_a.TA, + fobj_b.TA.T @ fobj_b.core_veff @ fobj_b.TA, + ] else: full_uhf.gcores_raw = [full_uhf.veff0_a, full_uhf.veff0_b] -# full_uhf.gcores_raw = [fobj_a.TA.T @ fobj_a.hf_veff @ fobj_a.TA, fobj_b.TA.T @ fobj_b.hf_veff @ fobj_b.TA] + # full_uhf.gcores_raw = [fobj_a.TA.T @ fobj_a.hf_veff @ fobj_a.TA, fobj_b.TA.T @ fobj_b.hf_veff @ fobj_b.TA] full_uhf.core_veffs = None - return full_uhf, Vs + return full_uhf, Vs + def uccsd_restore_eris(symm, fobj_a, fobj_b, pad0=True, skip_Vab=False): from pyscf import ao2mo + """ restore ERIs in the correct spin spaces """ Vsfile = fobj_a.eri_file Vsname = fobj_a.dname - nf = (fobj_a._mf.mo_coeff.shape[1],fobj_b._mf.mo_coeff.shape[1]) + nf = (fobj_a._mf.mo_coeff.shape[1], fobj_b._mf.mo_coeff.shape[1]) with h5py.File(Vsfile, "r") as fVs: Vs = [None] * 3 @@ -63,30 +75,33 @@ def uccsd_restore_eris(symm, fobj_a, fobj_b, pad0=True, skip_Vab=False): return Vs + def restore_eri_gen(targetsym, eri, norb1, norb2): """ An extension of PySCF's ao2mo.restore to Vaabb. """ - assert(targetsym in (1,4)) + assert targetsym in (1, 4) - npair1 = norb1*(norb1+1) // 2 - npair2 = norb2*(norb2+1) // 2 + npair1 = norb1 * (norb1 + 1) // 2 + npair2 = norb2 * (norb2 + 1) // 2 - if eri.size == norb1**2 * norb2**2: # s1 + if eri.size == norb1**2 * norb2**2: # s1 if targetsym == 1: - return eri.reshape(norb1,norb1,norb2,norb2) + return eri.reshape(norb1, norb1, norb2, norb2) else: return _convert_eri_gen(1, 4, eri, norb1, norb2) - elif eri.size == npair1 * npair2: # s4 + elif eri.size == npair1 * npair2: # s4 if targetsym == 1: return _convert_eri_gen(4, 1, eri, norb1, norb2) else: - return eri.reshape(npair1,npair2) + return eri.reshape(npair1, npair2) + def _convert_eri_gen(origsym, targetsym, eri, norb1, norb2): import ctypes from pyscf import lib - libao2mo = lib.load_library('libao2mo') + + libao2mo = lib.load_library("libao2mo") """ #NOTE: IF YOU GET AN ERROR ABOUT THIS ATTRIBUTE: This requires a custom PySCF compilation @@ -129,19 +144,20 @@ def _convert_eri_gen(origsym, targetsym, eri, norb1, norb2): } } } """ - fn = getattr(libao2mo, 'AO2MOrestore_nr%sto%s_gen'%(origsym,targetsym)) + fn = getattr(libao2mo, "AO2MOrestore_nr%sto%s_gen" % (origsym, targetsym)) if targetsym == 1: - eri_out = numpy.empty([norb1,norb1,norb2,norb2]) + eri_out = numpy.empty([norb1, norb1, norb2, norb2]) elif targetsym == 4: - npair1 = norb1*(norb1+1) // 2 - npair2 = norb2*(norb2+1) // 2 - eri_out = numpy.empty([npair1,npair2]) + npair1 = norb1 * (norb1 + 1) // 2 + npair2 = norb2 * (norb2 + 1) // 2 + eri_out = numpy.empty([npair1, npair2]) - fn(eri.ctypes.data_as(ctypes.c_void_p), + fn( + eri.ctypes.data_as(ctypes.c_void_p), eri_out.ctypes.data_as(ctypes.c_void_p), ctypes.c_int(norb1), - ctypes.c_int(norb2)) + ctypes.c_int(norb2), + ) return eri_out - diff --git a/molbe/fragment.py b/molbe/fragment.py index cd2ca22b..7c7b8074 100644 --- a/molbe/fragment.py +++ b/molbe/fragment.py @@ -4,6 +4,7 @@ from .helper import get_core from .autofrag import autogen + class fragpart: """Fragment/partitioning definition @@ -38,18 +39,24 @@ class fragpart: Whether to write 'fragment.xyz' file which contains all the fragments in cartesian coordinates. """ - def __init__(self, frag_type='autogen', - closed=False, - valence_basis=None,valence_only=False, - print_frags=True, write_geom=False, - be_type='be2', mol=None, frozen_core=False): - + def __init__( + self, + frag_type="autogen", + closed=False, + valence_basis=None, + valence_only=False, + print_frags=True, + write_geom=False, + be_type="be2", + mol=None, + frozen_core=False, + ): # Initialize class attributes self.mol = mol self.frag_type = frag_type self.fsites = [] self.Nfrag = 0 - self.edge= [] + self.edge = [] self.center = [] self.ebe_weight = [] self.edge_idx = [] @@ -71,101 +78,125 @@ def __init__(self, frag_type='autogen', self.ncore, self.no_core_idx, self.core_list = get_core(mol) # Check type of fragmentation function - if frag_type=='hchain_simple': + if frag_type == "hchain_simple": # This is an experimental feature. self.hchain_simple() - elif frag_type=='chain': + elif frag_type == "chain": if mol is None: - print('Provide pyscf gto.M object in fragpart() and restart!', - flush=True) - print('exiting',flush=True) + print( + "Provide pyscf gto.M object in fragpart() and restart!", flush=True + ) + print("exiting", flush=True) sys.exit() - self.chain(mol, frozen_core=frozen_core,closed=closed) - elif frag_type=='autogen': + self.chain(mol, frozen_core=frozen_core, closed=closed) + elif frag_type == "autogen": if mol is None: - print('Provide pyscf gto.M object in fragpart() and restart!', - flush=True) - print('exiting',flush=True) + print( + "Provide pyscf gto.M object in fragpart() and restart!", flush=True + ) + print("exiting", flush=True) sys.exit() - fgs = autogen(mol, be_type=be_type, frozen_core=frozen_core,write_geom=write_geom, - valence_basis=valence_basis, valence_only=valence_only, print_frags=print_frags) - - self.fsites, self.edge, self.center, self.edge_idx, self.center_idx, self.centerf_idx, self.ebe_weight, self.Frag_atom, self.center_atom, self.hlist_atom, self.add_center_atom = fgs + fgs = autogen( + mol, + be_type=be_type, + frozen_core=frozen_core, + write_geom=write_geom, + valence_basis=valence_basis, + valence_only=valence_only, + print_frags=print_frags, + ) + + ( + self.fsites, + self.edge, + self.center, + self.edge_idx, + self.center_idx, + self.centerf_idx, + self.ebe_weight, + self.Frag_atom, + self.center_atom, + self.hlist_atom, + self.add_center_atom, + ) = fgs self.Nfrag = len(self.fsites) else: - print('Fragmentation type = ',frag_type,' not implemented!', - flush=True) - print('exiting',flush=True) + print("Fragmentation type = ", frag_type, " not implemented!", flush=True) + print("exiting", flush=True) sys.exit() from .lchain import chain + def hchain_simple(self): - """Hard coded fragmentation feature - """ + """Hard coded fragmentation feature""" self.natom = self.mol.natm - if self.be_type=='be1': + if self.be_type == "be1": for i in range(self.natom): self.fsites.append([i]) self.edge.append([]) self.Nfrag = len(self.fsites) - elif self.be_type=='be2': - for i in range(self.natom-2): - self.fsites.append([i, i+1, i+2]) + elif self.be_type == "be2": + for i in range(self.natom - 2): + self.fsites.append([i, i + 1, i + 2]) self.centerf_idx.append([1]) self.Nfrag = len(self.fsites) self.edge.append([[2]]) for i in self.fsites[1:-1]: - self.edge.append([[i[0]],[i[-1]]]) + self.edge.append([[i[0]], [i[-1]]]) self.edge.append([[self.fsites[-1][0]]]) self.center.append([1]) - for i in range(self.Nfrag-2): - self.center.append([i,i+2]) - self.center.append([self.Nfrag-2]) + for i in range(self.Nfrag - 2): + self.center.append([i, i + 2]) + self.center.append([self.Nfrag - 2]) - elif self.be_type=='be3': - for i in range(self.natom-4): - self.fsites.append([i, i+1, i+2, i+3, i+4]) + elif self.be_type == "be3": + for i in range(self.natom - 4): + self.fsites.append([i, i + 1, i + 2, i + 3, i + 4]) self.centerf_idx.append([2]) self.Nfrag = len(self.fsites) - self.edge.append([[3],[4]]) + self.edge.append([[3], [4]]) for i in self.fsites[1:-1]: - self.edge.append([[i[0]],[i[1]],[i[-2]],[i[-1]]]) - self.edge.append([[self.fsites[-1][0]],[self.fsites[-1][1]]]) + self.edge.append([[i[0]], [i[1]], [i[-2]], [i[-1]]]) + self.edge.append([[self.fsites[-1][0]], [self.fsites[-1][1]]]) - self.center.append([1,2]) - self.center.append([0,0,2,3]) - for i in range(self.Nfrag-4): - self.center.append([i,i+1, i+3,i+4]) + self.center.append([1, 2]) + self.center.append([0, 0, 2, 3]) + for i in range(self.Nfrag - 4): + self.center.append([i, i + 1, i + 3, i + 4]) - self.center.append([self.Nfrag-4,self.Nfrag-3, - self.Nfrag-1,self.Nfrag-1]) - self.center.append([self.Nfrag-3,self.Nfrag-2]) + self.center.append( + [self.Nfrag - 4, self.Nfrag - 3, self.Nfrag - 1, self.Nfrag - 1] + ) + self.center.append([self.Nfrag - 3, self.Nfrag - 2]) for ix, i in enumerate(self.fsites): tmp_ = [] - elist_ = [ xx for yy in self.edge[ix] for xx in yy] + elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) - if not self.be_type=='be1': + if not self.be_type == "be1": for i in range(self.Nfrag): idx = [] for j in self.edge[i]: - idx.append([self.fsites[i].index(k) for k in j]) self.edge_idx.append(idx) for i in range(self.Nfrag): idx = [] for j in range(len(self.center[i])): - idx.append([self.fsites[self.center[i][j]].index(k) - for k in self.edge[i][j]]) + idx.append( + [ + self.fsites[self.center[i][j]].index(k) + for k in self.edge[i][j] + ] + ) self.center_idx.append(idx) - diff --git a/molbe/helper.py b/molbe/helper.py index e6be4208..4e8a88e2 100644 --- a/molbe/helper.py +++ b/molbe/helper.py @@ -4,6 +4,7 @@ import numpy, sys, functools, h5py from pyscf import ao2mo + def get_veff(eri_, dm, S, TA, hf_veff): """ Calculate the effective HF potential (Veff) for a given density matrix and electron repulsion integrals. @@ -35,7 +36,7 @@ def get_veff(eri_, dm, S, TA, hf_veff): # Transform the density matrix ST = numpy.dot(S, TA) - P_ = functools.reduce(numpy.dot,(ST.T, dm, ST)) + P_ = functools.reduce(numpy.dot, (ST.T, dm, ST)) # Ensure the transformed density matrix and ERI are real and double-precision P_ = numpy.asarray(P_.real, dtype=numpy.double) @@ -44,14 +45,23 @@ def get_veff(eri_, dm, S, TA, hf_veff): # Compute the Coulomb (J) and exchange (K) integrals vj, vk = scf.hf.dot_eri_dm(eri_, P_, hermi=1, with_j=True, with_k=True) Veff_ = vj - 0.5 * vk - Veff = functools.reduce(numpy.dot,(TA.T, hf_veff, TA)) - Veff_ + Veff = functools.reduce(numpy.dot, (TA.T, hf_veff, TA)) - Veff_ return Veff + # create pyscf pbc scf object -def get_scfObj(h1, Eri, nocc, dm0=None, enuc=0., - pert_h=False, pert_list=None, - save_chkfile=False, fname='f0'): +def get_scfObj( + h1, + Eri, + nocc, + dm0=None, + enuc=0.0, + pert_h=False, + pert_list=None, + save_chkfile=False, + fname="f0", +): """ Initialize and run a restricted Hartree-Fock (RHF) calculation. @@ -90,12 +100,12 @@ def get_scfObj(h1, Eri, nocc, dm0=None, enuc=0., # Initialize an RHF object mf_ = scf.RHF(mol) - mf_.get_hcore = lambda *args:h1 + mf_.get_hcore = lambda *args: h1 mf_.get_ovlp = lambda *args: S mf_._eri = Eri mf_.incore_anyway = True - mf_.max_cycle=50 - mf_.verbose=0 + mf_.max_cycle = 50 + mf_.verbose = 0 # Run the SCF calculation if dm0 is None: @@ -106,28 +116,31 @@ def get_scfObj(h1, Eri, nocc, dm0=None, enuc=0., # Check if the SCF calculation converged if not mf_.converged: print(flush=True) - print('WARNING!!! SCF not convereged - applying level_shift=0.2, diis_space=25 ',flush=True) + print( + "WARNING!!! SCF not convereged - applying level_shift=0.2, diis_space=25 ", + flush=True, + ) print(flush=True) - mf_.verbose=0 - mf_.level_shift=0.2 - mf_.diis_space=25 + mf_.verbose = 0 + mf_.level_shift = 0.2 + mf_.diis_space = 25 if dm0 is None: mf_.kernel() else: mf_.kernel(dm0=dm0) if not mf_.converged: print(flush=True) - print('WARNING!!! SCF still not convereged!',flush=True) + print("WARNING!!! SCF still not convereged!", flush=True) print(flush=True) else: print(flush=True) - print('SCF Converged!',flush=True) + print("SCF Converged!", flush=True) print(flush=True) return mf_ -def get_eri(i_frag, Nao, symm = 8, ignore_symm = False, eri_file='eri_file.h5'): +def get_eri(i_frag, Nao, symm=8, ignore_symm=False, eri_file="eri_file.h5"): """ Retrieve and optionally restore electron repulsion integrals (ERI) from an HDF5 file. @@ -155,12 +168,11 @@ def get_eri(i_frag, Nao, symm = 8, ignore_symm = False, eri_file='eri_file.h5'): import h5py # Open the HDF5 file and read the ERI for the specified fragment - r = h5py.File(eri_file,'r') + r = h5py.File(eri_file, "r") eri__ = numpy.array(r.get(i_frag)) # Optionally restore the symmetry of the ERI if not ignore_symm: - # Set the number of threads for the library to 1 lib.num_threads(1) eri__ = ao2mo.restore(symm, eri__, Nao) @@ -169,26 +181,25 @@ def get_eri(i_frag, Nao, symm = 8, ignore_symm = False, eri_file='eri_file.h5'): return eri__ -def ncore_(z): - if 1<= z<=2: +def ncore_(z): + if 1 <= z <= 2: nc = 0 - elif 2<=z<=5: - nc=1 - elif 5<=z<=12: - nc=1 - elif 12<=z<=30: - nc=5 - elif 31<=z<=38: - nc=9 - elif 39<=z<=48: - nc=14 - elif 49<=z<=56: - nc=18 + elif 2 <= z <= 5: + nc = 1 + elif 5 <= z <= 12: + nc = 1 + elif 12 <= z <= 30: + nc = 5 + elif 31 <= z <= 38: + nc = 9 + elif 39 <= z <= 48: + nc = 14 + elif 49 <= z <= 56: + nc = 18 else: - print('Ncore not computed in helper.ncore(), add it yourself!', - flush=True) - print('exiting',flush=True) + print("Ncore not computed in helper.ncore(), add it yourself!", flush=True) + print("exiting", flush=True) sys.exit() return nc @@ -215,11 +226,25 @@ def get_core(mol): ncore = ncore_(mol.atom_charge(ix)) corelist.append(ncore) Ncore += ncore - idx.extend([k for k in range(bas[2]+ncore, bas[3])]) - - return (Ncore,idx,corelist) - -def get_frag_energy(mo_coeffs, nsocc, nfsites, efac, TA, h1, hf_veff, rdm1, rdm2s, dname, eri_file='eri_file.h5', veff0=None): + idx.extend([k for k in range(bas[2] + ncore, bas[3])]) + + return (Ncore, idx, corelist) + + +def get_frag_energy( + mo_coeffs, + nsocc, + nfsites, + efac, + TA, + h1, + hf_veff, + rdm1, + rdm2s, + dname, + eri_file="eri_file.h5", + veff0=None, +): """ Compute the fragment energy. @@ -261,19 +286,18 @@ def get_frag_energy(mo_coeffs, nsocc, nfsites, efac, TA, h1, hf_veff, rdm1, rdm2 rdm1s_rot = mo_coeffs @ rdm1 @ mo_coeffs.T * 0.5 # Construct the Hartree-Fock 1-RDM - hf_1rdm = numpy.dot(mo_coeffs[:,:nsocc], - mo_coeffs[:,:nsocc].conj().T) + hf_1rdm = numpy.dot(mo_coeffs[:, :nsocc], mo_coeffs[:, :nsocc].conj().T) # Compute the difference between the rotated RDM1 and the Hartree-Fock 1-RDM delta_rdm1 = 2 * (rdm1s_rot - hf_1rdm) if veff0 is None: # Compute the effective potential in the transformed basis - veff0 = functools.reduce(numpy.dot,(TA.T,hf_veff,TA)) + veff0 = functools.reduce(numpy.dot, (TA.T, hf_veff, TA)) # Calculate the one-electron and effective potential energy contributions - e1 = numpy.einsum("ij,ij->i",h1[:nfsites], delta_rdm1[:nfsites]) - ec = numpy.einsum("ij,ij->i",veff0[:nfsites], delta_rdm1[:nfsites]) + e1 = numpy.einsum("ij,ij->i", h1[:nfsites], delta_rdm1[:nfsites]) + ec = numpy.einsum("ij,ij->i", veff0[:nfsites], delta_rdm1[:nfsites]) if TA.ndim == 3: jmax = TA[0].shape[1] @@ -281,13 +305,14 @@ def get_frag_energy(mo_coeffs, nsocc, nfsites, efac, TA, h1, hf_veff, rdm1, rdm2 jmax = TA.shape[1] # Load the electron repulsion integrals from the HDF5 file - r = h5py.File(eri_file,'r') + r = h5py.File(eri_file, "r") eri = r[dname][()] r.close() # Rotate the RDM2 into the MO basis - rdm2s = numpy.einsum("ijkl,pi,qj,rk,sl->pqrs", 0.5*rdm2s, - *([mo_coeffs]*4),optimize=True) + rdm2s = numpy.einsum( + "ijkl,pi,qj,rk,sl->pqrs", 0.5 * rdm2s, *([mo_coeffs] * 4), optimize=True + ) # Initialize the two-electron energy contribution e2 = numpy.zeros_like(e1) @@ -295,32 +320,47 @@ def get_frag_energy(mo_coeffs, nsocc, nfsites, efac, TA, h1, hf_veff, rdm1, rdm2 # Calculate the two-electron energy contribution for i in range(nfsites): for j in range(jmax): - ij = i*(i+1)//2+j if i > j else j*(j+1)//2+i - Gij = rdm2s[i,j,:jmax,:jmax].copy() + ij = i * (i + 1) // 2 + j if i > j else j * (j + 1) // 2 + i + Gij = rdm2s[i, j, :jmax, :jmax].copy() Gij[numpy.diag_indices(jmax)] *= 0.5 Gij += Gij.T e2[i] += Gij[numpy.tril_indices(jmax)] @ eri[ij] # Sum the energy contributions - e_ = e1+e2+ec + e_ = e1 + e2 + ec # Initialize temporary energy variables - etmp = 0. - e1_tmp = 0. - e2_tmp = 0. - ec_tmp = 0. + etmp = 0.0 + e1_tmp = 0.0 + e2_tmp = 0.0 + ec_tmp = 0.0 # Calculate the total energy contribution for the specified fragment indices for i in efac[1]: - etmp += efac[0]*e_[i] - e1_tmp += efac[0]*e1[i] - e2_tmp += efac[0]*e2[i] - ec_tmp += efac[0]*ec[i] - - return [e1_tmp,e2_tmp,ec_tmp] - -def get_frag_energy_u(mo_coeffs, nsocc, nfsites, efac, TA, h1, hf_veff, rdm1, rdm2s, dname, - eri_file='eri_file.h5', gcores=None, frozen=False, veff0=None): + etmp += efac[0] * e_[i] + e1_tmp += efac[0] * e1[i] + e2_tmp += efac[0] * e2[i] + ec_tmp += efac[0] * ec[i] + + return [e1_tmp, e2_tmp, ec_tmp] + + +def get_frag_energy_u( + mo_coeffs, + nsocc, + nfsites, + efac, + TA, + h1, + hf_veff, + rdm1, + rdm2s, + dname, + eri_file="eri_file.h5", + gcores=None, + frozen=False, + veff0=None, +): """ Compute the fragment energy for unrestricted calculations @@ -363,57 +403,73 @@ def get_frag_energy_u(mo_coeffs, nsocc, nfsites, efac, TA, h1, hf_veff, rdm1, rd """ # Rotate the RDM1 into the MO basis for both spins - rdm1s_rot = [mo_coeffs[s] @ rdm1[s] @ mo_coeffs[s].T for s in [0,1] ] + rdm1s_rot = [mo_coeffs[s] @ rdm1[s] @ mo_coeffs[s].T for s in [0, 1]] # Construct the Hartree-Fock RDM1 for both spin the the Schmidt space - hf_1rdm = [numpy.dot(mo_coeffs[s][:,:nsocc[s]], - mo_coeffs[s][:,:nsocc[s]].conj().T) for s in [0,1]] + hf_1rdm = [ + numpy.dot(mo_coeffs[s][:, : nsocc[s]], mo_coeffs[s][:, : nsocc[s]].conj().T) + for s in [0, 1] + ] # Compute the difference between the rotated RDM1 and the HF RDM1 - delta_rdm1 = [2 * (rdm1s_rot[s] - hf_1rdm[s]) for s in [0,1]] + delta_rdm1 = [2 * (rdm1s_rot[s] - hf_1rdm[s]) for s in [0, 1]] if veff0 is None: # Compute thte effective potential in the transformed basis - veff0 = [functools.reduce(numpy.dot,(TA[s].T,hf_veff[s],TA[s])) for s in [0,1]] + veff0 = [ + functools.reduce(numpy.dot, (TA[s].T, hf_veff[s], TA[s])) for s in [0, 1] + ] # For frozen care, remove core potential and Hamiltonian components if frozen: - for s in [0,1]: + for s in [0, 1]: veff0[s] -= gcores[s] h1[s] -= gcores[s] # Calculate the one-electron and effective potential energy contributions - e1 = [numpy.einsum("ij,ij->i",h1[s][:nfsites[s]], delta_rdm1[s][:nfsites[s]]) for s in [0,1]] - ec = [numpy.einsum("ij,ij->i",veff0[s][:nfsites[s]], delta_rdm1[s][:nfsites[s]]) for s in [0,1]] + e1 = [ + numpy.einsum("ij,ij->i", h1[s][: nfsites[s]], delta_rdm1[s][: nfsites[s]]) + for s in [0, 1] + ] + ec = [ + numpy.einsum("ij,ij->i", veff0[s][: nfsites[s]], delta_rdm1[s][: nfsites[s]]) + for s in [0, 1] + ] - jmax = [TA[0].shape[1],TA[1].shape[1]] + jmax = [TA[0].shape[1], TA[1].shape[1]] # Load ERIs from the HDF5 file - r = h5py.File(eri_file,'r') - Vs = [r[dname[0]][()],r[dname[1]][()],r[dname[2]][()]] + r = h5py.File(eri_file, "r") + Vs = [r[dname[0]][()], r[dname[1]][()], r[dname[2]][()]] r.close() # Rotate the RDM2 into the MO basis - rdm2s_k = [numpy.einsum("ijkl,pi,qj,rk,sl->pqrs", rdm2s[s], - *([mo_coeffs[s12[0]]]*2+[mo_coeffs[s12[1]]]*2), optimize=True) - for s,s12 in zip([0,1,2],[[0,0],[0,1],[1,1]])] + rdm2s_k = [ + numpy.einsum( + "ijkl,pi,qj,rk,sl->pqrs", + rdm2s[s], + *([mo_coeffs[s12[0]]] * 2 + [mo_coeffs[s12[1]]] * 2), + optimize=True, + ) + for s, s12 in zip([0, 1, 2], [[0, 0], [0, 1], [1, 1]]) + ] # Initialize the two-electron energy contribution - e2 = [numpy.zeros(h1[0].shape[0]),numpy.zeros(h1[1].shape[0])] + e2 = [numpy.zeros(h1[0].shape[0]), numpy.zeros(h1[1].shape[0])] # Calculate the two-electron energy contribution for alpha and beta def contract_2e(jmaxs, rdm2_, V_, s, sym): e2_ = numpy.zeros(nfsites[s]) - jmax1,jmax2 = [jmaxs]*2 if isinstance(jmaxs,int) else jmaxs + jmax1, jmax2 = [jmaxs] * 2 if isinstance(jmaxs, int) else jmaxs for i in range(nfsites[s]): for j in range(jmax1): - ij = i*(i+1)//2+j if i > j else j*(j+1)//2+i - if sym in [4,2]: - Gij = rdm2_[i,j,:jmax2,:jmax2].copy() + ij = i * (i + 1) // 2 + j if i > j else j * (j + 1) // 2 + i + if sym in [4, 2]: + Gij = rdm2_[i, j, :jmax2, :jmax2].copy() Vij = V_[ij] else: - Gij = rdm2_[:jmax2,:jmax2,i,j].copy() - Vij = V_[:,ij] + Gij = rdm2_[:jmax2, :jmax2, i, j].copy() + Vij = V_[:, ij] Gij[numpy.diag_indices(jmax2)] *= 0.5 Gij += Gij.T e2_[i] += Gij[numpy.tril_indices(jmax2)] @ Vij @@ -422,10 +478,10 @@ def contract_2e(jmaxs, rdm2_, V_, s, sym): return e2_ # the first nf are frag sites - e2ss = [0.,0.] - e2os = [0.,0.] - for s in [0,1]: - e2ss[s] += contract_2e(jmax[s], rdm2s_k[2*s], Vs[s], s, sym=4) + e2ss = [0.0, 0.0] + e2os = [0.0, 0.0] + for s in [0, 1]: + e2ss[s] += contract_2e(jmax[s], rdm2s_k[2 * s], Vs[s], s, sym=4) # Calculate the cross-spin two-electron energy contributions V = Vs[2] @@ -438,20 +494,20 @@ def contract_2e(jmaxs, rdm2_, V_, s, sym): e2 = sum(e2ss) + sum(e2os) # Sum the energy contributions - e_ = e1+e2+ec + e_ = e1 + e2 + ec # Initialize temporary energy variables - etmp = 0. - e1_tmp = 0. - e2_tmp = 0. - ec_tmp = 0. + etmp = 0.0 + e1_tmp = 0.0 + e2_tmp = 0.0 + ec_tmp = 0.0 # Calculate the total energy contribution for the specified fragment indices for i in efac[0][1]: - e2_tmp += efac[0][0]*e2[i] - for s in [0,1]: - etmp += efac[s][0]*e_[s][i] - e1_tmp += efac[s][0]*e1[s][i] - ec_tmp += efac[s][0]*ec[s][i] + e2_tmp += efac[0][0] * e2[i] + for s in [0, 1]: + etmp += efac[s][0] * e_[s][i] + e1_tmp += efac[s][0] * e1[s][i] + ec_tmp += efac[s][0] * ec[s][i] - return [e1_tmp,e2_tmp,ec_tmp] + return [e1_tmp, e2_tmp, ec_tmp] diff --git a/molbe/lchain.py b/molbe/lchain.py index df52d0c0..8dc9098e 100644 --- a/molbe/lchain.py +++ b/molbe/lchain.py @@ -1,6 +1,7 @@ # Author(s): Oinam Romesh Meitei -def chain(self,mol, frozen_core=False, closed=False): + +def chain(self, mol, frozen_core=False, closed=False): """ Hard coded linear chain fragment constructor """ @@ -15,199 +16,238 @@ def chain(self,mol, frozen_core=False, closed=False): else: start_ -= coreshift ncore_ = self.core_list[adx] - stop_ -= coreshift+ncore_ - coreshift = ncore_+coreshift + stop_ -= coreshift + ncore_ + coreshift = ncore_ + coreshift sites.append([i for i in range(start_, stop_)]) if closed: - lnext = [i for i in self.kpt if i>1] + lnext = [i for i in self.kpt if i > 1] if not len(lnext) == 0: nk1 = lnext[0] else: - print('Gamma point does not work') + print("Gamma point does not work") sys.exit() Ns = mol.aoslice_by_atom()[-1][3] - - if self.be_type=='be2': - fs=[] + if self.be_type == "be2": + fs = [] if closed: # change back to i,i+1,i+2 - sites_left = [-i-1 for i in sites[0]] - #sites_left = [i for i in sites[-1]] + sites_left = [-i - 1 for i in sites[0]] + # sites_left = [i for i in sites[-1]] ns = len(sites[-1]) - sites_right = [i+ns for i in sites[-1]] + sites_right = [i + ns for i in sites[-1]] if closed: - self.fsites.append(sites_left+sites[0]+sites[1]) + self.fsites.append(sites_left + sites[0] + sites[1]) fs.append([sites_left, sites[0], sites[1]]) - for i in range(mol.natm-2): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]) - fs.append([ sites[i], sites[i+1], sites[i+2]]) + for i in range(mol.natm - 2): + self.fsites.append(sites[i] + sites[i + 1] + sites[i + 2]) + fs.append([sites[i], sites[i + 1], sites[i + 2]]) if closed: - self.fsites.append(sites[-2]+sites[-1]+sites_right) - fs.append([sites[-2],sites[-1],sites_right]) + self.fsites.append(sites[-2] + sites[-1] + sites_right) + fs.append([sites[-2], sites[-1], sites_right]) self.Nfrag = len(self.fsites) if closed: - for i in fs: #[1:-1]: - self.edge.append([i[0],i[-1]]) + for i in fs: # [1:-1]: + self.edge.append([i[0], i[-1]]) else: self.edge.append([fs[0][2]]) # change back 0->1 for i in fs[1:-1]: - self.edge.append([i[0],i[-1]]) + self.edge.append([i[0], i[-1]]) self.edge.append([fs[-1][0]]) if closed: - self.center.append([self.Nfrag-1,1]) - #self.center.append([0,2]) + self.center.append([self.Nfrag - 1, 1]) + # self.center.append([0,2]) else: self.center.append([1]) - for i in range(self.Nfrag-2): - self.center.append([i,i+2]) + for i in range(self.Nfrag - 2): + self.center.append([i, i + 2]) if closed: - #self.center.append([self.Nfrag-3,self.Nfrag-1]) - self.center.append([self.Nfrag-2,0]) + # self.center.append([self.Nfrag-3,self.Nfrag-1]) + self.center.append([self.Nfrag - 2, 0]) else: - self.center.append([self.Nfrag-2]) + self.center.append([self.Nfrag - 2]) if closed: for ix, i in enumerate(self.fsites): tmp_ = [] elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) for i in range(self.Nfrag): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][1]]) - if self.be_type=='be3' and not closed: - fs=[] + if self.be_type == "be3" and not closed: + fs = [] - for i in range(mol.natm-4): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]+ - sites[i+3]+ sites[i+4]) - fs.append([ sites[i], sites[i+1], sites[i+2], sites[i+3], - sites[i+4]]) + for i in range(mol.natm - 4): + self.fsites.append( + sites[i] + sites[i + 1] + sites[i + 2] + sites[i + 3] + sites[i + 4] + ) + fs.append( + [sites[i], sites[i + 1], sites[i + 2], sites[i + 3], sites[i + 4]] + ) self.Nfrag = len(self.fsites) self.edge.append([fs[0][3], fs[0][4]]) # change back 0->1 for i in fs[1:-1]: - self.edge.append([i[0],i[1], i[-2], i[-1]]) + self.edge.append([i[0], i[1], i[-2], i[-1]]) self.edge.append([fs[-1][0], fs[-1][1]]) - self.center.append([1,2]) - self.center.append([0,0,2,3]) - for i in range(self.Nfrag-4): - self.center.append([i, i+1, i+3, i+4]) + self.center.append([1, 2]) + self.center.append([0, 0, 2, 3]) + for i in range(self.Nfrag - 4): + self.center.append([i, i + 1, i + 3, i + 4]) - self.center.append([self.Nfrag-4, self.Nfrag-3, - self.Nfrag-1, self.Nfrag-1]) - self.center.append([self.Nfrag-3, self.Nfrag-2]) + self.center.append( + [self.Nfrag - 4, self.Nfrag - 3, self.Nfrag - 1, self.Nfrag - 1] + ) + self.center.append([self.Nfrag - 3, self.Nfrag - 2]) for i in range(self.Nfrag): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][2]]) - if self.be_type=='be3' and closed: - fs=[] + if self.be_type == "be3" and closed: + fs = [] ns = len(sites[-1]) - sites_left1 = [-i-1 for i in sites[0]] - sites_left0 = [i-ns for i in sites_left1] - sites_right0 = [i+ns for i in sites[-1]] - sites_right1 = [i+ns for i in sites_right0] - - self.fsites.append(sites_left0+sites_left1+sites[0]+sites[1]+sites[2]) - self.fsites.append(sites_left1+sites[0]+sites[1]+sites[2]+sites[3]) - fs.append([sites_left0,sites_left1,sites[0],sites[1],sites[2]]) - fs.append([sites_left1,sites[0],sites[1],sites[2],sites[3]]) - for i in range(len(sites)-4): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]+ - sites[i+3]+ sites[i+4]) - fs.append([ sites[i], sites[i+1], sites[i+2], sites[i+3], - sites[i+4]]) - self.fsites.append(sites[-4]+sites[-3]+sites[-2]+sites[-1]+sites_right0) - self.fsites.append(sites[-3]+sites[-2]+sites[-1]+sites_right0+sites_right1) - fs.append([sites[-4],sites[-3],sites[-2],sites[-1],sites_right0]) - fs.append([sites[-3],sites[-2],sites[-1],sites_right0,sites_right1]) + sites_left1 = [-i - 1 for i in sites[0]] + sites_left0 = [i - ns for i in sites_left1] + sites_right0 = [i + ns for i in sites[-1]] + sites_right1 = [i + ns for i in sites_right0] + + self.fsites.append(sites_left0 + sites_left1 + sites[0] + sites[1] + sites[2]) + self.fsites.append(sites_left1 + sites[0] + sites[1] + sites[2] + sites[3]) + fs.append([sites_left0, sites_left1, sites[0], sites[1], sites[2]]) + fs.append([sites_left1, sites[0], sites[1], sites[2], sites[3]]) + for i in range(len(sites) - 4): + self.fsites.append( + sites[i] + sites[i + 1] + sites[i + 2] + sites[i + 3] + sites[i + 4] + ) + fs.append( + [sites[i], sites[i + 1], sites[i + 2], sites[i + 3], sites[i + 4]] + ) + self.fsites.append(sites[-4] + sites[-3] + sites[-2] + sites[-1] + sites_right0) + self.fsites.append( + sites[-3] + sites[-2] + sites[-1] + sites_right0 + sites_right1 + ) + fs.append([sites[-4], sites[-3], sites[-2], sites[-1], sites_right0]) + fs.append([sites[-3], sites[-2], sites[-1], sites_right0, sites_right1]) self.Nfrag = len(self.fsites) - #self.edge.append([fs[0][3], fs[0][4]]) - #self.edge.append([fs[1][1], fs[1][3], fs[1][4]]) - for i in fs: #[2:-2]: - self.edge.append([i[0],i[1], i[-2], i[-1]]) - #self.edge.append([fs[-2][0],fs[-2][1],fs[-2][3]]) - #self.edge.append([fs[-1][0],fs[-1][1]]) + # self.edge.append([fs[0][3], fs[0][4]]) + # self.edge.append([fs[1][1], fs[1][3], fs[1][4]]) + for i in fs: # [2:-2]: + self.edge.append([i[0], i[1], i[-2], i[-1]]) + # self.edge.append([fs[-2][0],fs[-2][1],fs[-2][3]]) + # self.edge.append([fs[-1][0],fs[-1][1]]) - self.center.append([self.Nfrag-2,self.Nfrag-1, 1,2]) - self.center.append([self.Nfrag-1,0,2,3]) - for i in range(self.Nfrag-4): - self.center.append([i, i+1, i+3, i+4]) - self.center.append([self.Nfrag-4, self.Nfrag-3, self.Nfrag-1, 0]) - self.center.append([self.Nfrag-3, self.Nfrag-2, 0, 1]) + self.center.append([self.Nfrag - 2, self.Nfrag - 1, 1, 2]) + self.center.append([self.Nfrag - 1, 0, 2, 3]) + for i in range(self.Nfrag - 4): + self.center.append([i, i + 1, i + 3, i + 4]) + self.center.append([self.Nfrag - 4, self.Nfrag - 3, self.Nfrag - 1, 0]) + self.center.append([self.Nfrag - 3, self.Nfrag - 2, 0, 1]) for ix, i in enumerate(self.fsites): tmp_ = [] elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) for i in range(self.Nfrag): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][2]]) - if self.be_type=='be4' and not closed: - fs=[] - for i in range(mol.natm-6): - self.fsites.append(sites[i]+ sites[i+1]+ sites[i+2]+ - sites[i+3]+ sites[i+4] + sites[i+5] + sites[i+6]) - fs.append([ sites[i], sites[i+1], sites[i+2], sites[i+3], - sites[i+4],sites[i+5], sites[i+6]]) + if self.be_type == "be4" and not closed: + fs = [] + for i in range(mol.natm - 6): + self.fsites.append( + sites[i] + + sites[i + 1] + + sites[i + 2] + + sites[i + 3] + + sites[i + 4] + + sites[i + 5] + + sites[i + 6] + ) + fs.append( + [ + sites[i], + sites[i + 1], + sites[i + 2], + sites[i + 3], + sites[i + 4], + sites[i + 5], + sites[i + 6], + ] + ) self.Nfrag = len(self.fsites) self.edge.append([fs[0][4], fs[0][5], fs[0][6]]) for i in fs[1:-1]: - self.edge.append([i[0],i[1], i[2], i[-3],i[-2], i[-1]]) + self.edge.append([i[0], i[1], i[2], i[-3], i[-2], i[-1]]) self.edge.append([fs[-1][0], fs[-1][1], fs[-1][2]]) - self.center.append([1,2,3]) - self.center.append([0,0,0,2,3,4]) - self.center.append([0,0,1,3,4,5]) - for i in range(self.Nfrag-6): - self.center.append([i, i+1,i+2, i+4, i+5,i+6]) - - self.center.append([self.Nfrag-6, self.Nfrag-5, self.Nfrag-4, - self.Nfrag-2, self.Nfrag-1, self.Nfrag-1]) - self.center.append([self.Nfrag-5, self.Nfrag-4, self.Nfrag-3, - self.Nfrag-1, self.Nfrag-1, self.Nfrag-1]) - self.center.append([self.Nfrag-4, self.Nfrag-3, self.Nfrag-2]) + self.center.append([1, 2, 3]) + self.center.append([0, 0, 0, 2, 3, 4]) + self.center.append([0, 0, 1, 3, 4, 5]) + for i in range(self.Nfrag - 6): + self.center.append([i, i + 1, i + 2, i + 4, i + 5, i + 6]) + + self.center.append( + [ + self.Nfrag - 6, + self.Nfrag - 5, + self.Nfrag - 4, + self.Nfrag - 2, + self.Nfrag - 1, + self.Nfrag - 1, + ] + ) + self.center.append( + [ + self.Nfrag - 5, + self.Nfrag - 4, + self.Nfrag - 3, + self.Nfrag - 1, + self.Nfrag - 1, + self.Nfrag - 1, + ] + ) + self.center.append([self.Nfrag - 4, self.Nfrag - 3, self.Nfrag - 2]) for i in range(self.Nfrag): self.centerf_idx.append([self.fsites[i].index(j) for j in fs[i][3]]) - if self.be_type=='be4' and closed: - - print('Will add this soon!') + if self.be_type == "be4" and closed: + print("Will add this soon!") sys.exit() if not closed: for ix, i in enumerate(self.fsites): tmp_ = [] - elist_ = [ xx for yy in self.edge[ix] for xx in yy] + elist_ = [xx for yy in self.edge[ix] for xx in yy] for j in i: - if not j in elist_: tmp_.append(i.index(j)) + if not j in elist_: + tmp_.append(i.index(j)) self.ebe_weight.append([1.0, tmp_]) - if not self.be_type=='be1': + if not self.be_type == "be1": for i in range(self.Nfrag): idx = [] for j in self.edge[i]: - idx.append([self.fsites[i].index(k) for k in j]) self.edge_idx.append(idx) @@ -226,9 +266,10 @@ def chain(self,mol, frozen_core=False, closed=False): for i in range(self.Nfrag): idx = [] for j in range(len(self.center[i])): - - idx.append([self.fsites[self.center[i][j]].index(k) - for k in self.edge[i][j]]) + idx.append( + [ + self.fsites[self.center[i][j]].index(k) + for k in self.edge[i][j] + ] + ) self.center_idx.append(idx) - - diff --git a/molbe/lo.py b/molbe/lo.py index 2ab34cff..1f18154a 100644 --- a/molbe/lo.py +++ b/molbe/lo.py @@ -1,77 +1,88 @@ # Author(s): Henry Tran, Oinam Meitei, Shaun Weatherly # from pyscf import lib -import numpy,sys +import numpy, sys from copy import deepcopy from functools import reduce -from molbe.external.lo_helper import get_symm_mat_pow, get_aoind_by_atom, reorder_by_atom_ +from molbe.external.lo_helper import ( + get_symm_mat_pow, + get_aoind_by_atom, + reorder_by_atom_, +) + def dot_gen(A, B, ovlp): return A.T @ B if ovlp is None else A.T @ ovlp @ B -def get_cano_orth_mat(A, thr=1.E-6, ovlp=None): - S = dot_gen(A,A,ovlp) + +def get_cano_orth_mat(A, thr=1.0e-6, ovlp=None): + S = dot_gen(A, A, ovlp) e, u = numpy.linalg.eigh(S) if thr > 0: - idx_keep = e/e[-1] > thr + idx_keep = e / e[-1] > thr else: idx_keep = list(range(e.shape[0])) - U = u[:,idx_keep] * e[idx_keep]**-0.5 + U = u[:, idx_keep] * e[idx_keep] ** -0.5 return U -def cano_orth(A, thr=1.E-6, ovlp=None): - """ Canonically orthogonalize columns of A - """ + +def cano_orth(A, thr=1.0e-6, ovlp=None): + """Canonically orthogonalize columns of A""" U = get_cano_orth_mat(A, thr, ovlp) return A @ U -def get_symm_orth_mat(A, thr=1.E-6, ovlp=None): - S = dot_gen(A,A,ovlp) + +def get_symm_orth_mat(A, thr=1.0e-6, ovlp=None): + S = dot_gen(A, A, ovlp) e, u = numpy.linalg.eigh(S) if int(numpy.sum(e < thr)) > 0: - raise ValueError("Linear dependence is detected in the column space of A: smallest eigenvalue (%.3E) is less than thr (%.3E). Please use 'cano_orth' instead." % (numpy.min(e), thr)) + raise ValueError( + "Linear dependence is detected in the column space of A: smallest eigenvalue (%.3E) is less than thr (%.3E). Please use 'cano_orth' instead." + % (numpy.min(e), thr) + ) U = u @ numpy.diag(e**-0.5) @ u.T return U -def symm_orth(A, thr=1.E-6, ovlp=None): - """ Symmetrically orthogonalize columns of A - """ + +def symm_orth(A, thr=1.0e-6, ovlp=None): + """Symmetrically orthogonalize columns of A""" U = get_symm_orth_mat(A, thr, ovlp) return A @ U def remove_core_mo(Clo, Ccore, S, thr=0.5): - assert(numpy.allclose(Clo.T@S@Clo,numpy.eye(Clo.shape[1]))) - assert(numpy.allclose(Ccore.T@S@Ccore,numpy.eye(Ccore.shape[1]))) + assert numpy.allclose(Clo.T @ S @ Clo, numpy.eye(Clo.shape[1])) + assert numpy.allclose(Ccore.T @ S @ Ccore, numpy.eye(Ccore.shape[1])) - n,nlo = Clo.shape + n, nlo = Clo.shape ncore = Ccore.shape[1] - Pcore = Ccore@Ccore.T @ S + Pcore = Ccore @ Ccore.T @ S Clo1 = (numpy.eye(n) - Pcore) @ Clo pop = numpy.diag(Clo1.T @ S @ Clo1) - idx_keep = numpy.where(pop>thr)[0] - assert(len(idx_keep) == nlo-ncore) - Clo2 = symm_orth(Clo1[:,idx_keep], ovlp=S) + idx_keep = numpy.where(pop > thr)[0] + assert len(idx_keep) == nlo - ncore + Clo2 = symm_orth(Clo1[:, idx_keep], ovlp=S) return Clo2 -def reorder_lo(C, S, idao_by_atom, atom_by_motif, motifname, - ncore_by_motif, thresh=0.5, verbose=3): - """ Reorder localized orbitals +def reorder_lo( + C, S, idao_by_atom, atom_by_motif, motifname, ncore_by_motif, thresh=0.5, verbose=3 +): + """Reorder localized orbitals This function reorders the IAOs and PAOs so that the IAOs and PAOs for each atom are grouped together. """ - pop_by_ao = (Xinv @ C)**2 + pop_by_ao = (Xinv @ C) ** 2 reorder_idx_by_atom = [] for idao in idao_by_atom: - pop = numpy.sum(pop_by_ao[idao], axis = 0) + pop = numpy.sum(pop_by_ao[idao], axis=0) -def get_xovlp(mol, basis='sto-3g'): +def get_xovlp(mol, basis="sto-3g"): """ Gets set of valence orbitals based on smaller (should be minimal) basis inumpy.t: @@ -82,16 +93,18 @@ def get_xovlp(mol, basis='sto-3g'): S22 - Overlap in new basis set """ from pyscf.gto import intor_cross + mol_alt = mol.copy() mol_alt.basis = basis mol_alt.build() - S12 = intor_cross('int1e_ovlp', mol, mol_alt) - S22 = mol_alt.intor('int1e_ovlp') + S12 = intor_cross("int1e_ovlp", mol, mol_alt) + S22 = mol_alt.intor("int1e_ovlp") return S12, S22 -def get_iao(Co, S12, S1, S2 = None): + +def get_iao(Co, S12, S1, S2=None): """ Args: Co: occupied coefficient matrix with core @@ -120,14 +133,15 @@ def get_iao(Co, S12, S1, S2 = None): Potil = Cotil @ numpy.linalg.inv(Stil) @ Cotil.T Ciao = (numpy.eye(n) - (Po + Potil - 2 * Po @ S1 @ Potil) @ S1) @ ptil - Ciao = symm_orth(Ciao, ovlp = S1) + Ciao = symm_orth(Ciao, ovlp=S1) # check span rep_err = numpy.linalg.norm(Ciao @ Ciao.T @ S1 @ Po - Po) - if rep_err > 1.E-10: + if rep_err > 1.0e-10: raise RuntimeError return Ciao + def get_pao(Ciao, S, S12, S2, mol): """ Args: @@ -141,12 +155,12 @@ def get_pao(Ciao, S, S12, S2, mol): """ n = Ciao.shape[0] s12 = numpy.linalg.inv(S) @ S12 - nonval = numpy.eye(n) - s12 @ s12.T # set of orbitals minus valence (orth in working basis) - - - Piao = Ciao @ Ciao.T @ S # projector into IAOs - Cpao = (numpy.eye(n) - Piao) @ nonval # project out IAOs from non-valence basis + nonval = ( + numpy.eye(n) - s12 @ s12.T + ) # set of orbitals minus valence (orth in working basis) + Piao = Ciao @ Ciao.T @ S # projector into IAOs + Cpao = (numpy.eye(n) - Piao) @ nonval # project out IAOs from non-valence basis # begin canonical orthogonalization to get rid of redundant orbitals numpy.o0 = Cpao.shape[1] @@ -155,6 +169,7 @@ def get_pao(Ciao, S, S12, S2, mol): return Cpao + def get_pao_native(Ciao, S, mol, valence_basis): """ Args: @@ -175,7 +190,11 @@ def get_pao_native(Ciao, S, mol, valence_basis): full_ao_labels = mol.ao_labels() valence_ao_labels = mol_alt.ao_labels() - vir_idx = [idx for idx, label in enumerate(full_ao_labels) if (not label in valence_ao_labels)] + vir_idx = [ + idx + for idx, label in enumerate(full_ao_labels) + if (not label in valence_ao_labels) + ] Piao = Ciao @ Ciao.T @ S Cpao = (numpy.eye(n) - Piao)[:, vir_idx] @@ -187,34 +206,44 @@ def get_pao_native(Ciao, S, mol, valence_basis): npao0 = Cpao.shape[1] Cpao = cano_orth(Cpao, ovlp=S) npao1 = Cpao.shape[1] - print("# of PAO: %d --> %d" % (npao0,npao1), flush=True) + print("# of PAO: %d --> %d" % (npao0, npao1), flush=True) print("", flush=True) return Cpao + def get_loc(mol, C, method, pop_method=None, init_guess=None): - if method.upper() == 'ER': + if method.upper() == "ER": from pyscf.lo import ER as Localizer - elif method.upper() == 'PM': + elif method.upper() == "PM": from pyscf.lo import PM as Localizer - elif method.upper() == 'FB' or method.upper() == 'BOYS': + elif method.upper() == "FB" or method.upper() == "BOYS": from pyscf.lo import Boys as Localizer else: - raise NotImplementedError('Localization scheme not understood') + raise NotImplementedError("Localization scheme not understood") mlo = Localizer(mol, C) if pop_method is not None: - mlo.pop_method=str(pop_method) + mlo.pop_method = str(pop_method) mlo.init_guess = init_guess C_ = mlo.kernel() return C_ -def localize(self, lo_method, mol=None, valence_basis='sto-3g', - hstack=False, pop_method=None, init_guess=None, - valence_only=False, nosave=False): - """ Molecular orbital localization + +def localize( + self, + lo_method, + mol=None, + valence_basis="sto-3g", + hstack=False, + pop_method=None, + init_guess=None, + valence_only=False, + nosave=False, +): + """Molecular orbital localization Performs molecular orbital localization computations. For large basis, IAO is recommended augmented with PAO orbitals. @@ -234,119 +263,127 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', from numpy.linalg import eigh from pyscf.lo import iao from pyscf.lo import orth - import scipy.linalg,functools - from .helper import ncore_ - - if lo_method == 'lowdin': + import scipy.linalg, functools + from .helper import ncore_ + if lo_method == "lowdin": es_, vs_ = eigh(self.S) - edx = es_ > 1.e-15 - self.W = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].T) + edx = es_ > 1.0e-15 + self.W = numpy.dot(vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].T) if self.frozen_core: if self.unrestricted: - P_core = [numpy.eye(self.W.shape[0]) - numpy.dot(self.P_core[s], self.S) for s in [0,1]] + P_core = [ + numpy.eye(self.W.shape[0]) - numpy.dot(self.P_core[s], self.S) + for s in [0, 1] + ] C_ = numpy.dot(P_core, self.W) - Cpop = [functools.reduce(numpy.dot, - (C_[s].T, self.S, C_[s])) for s in [0,1]] - Cpop = [numpy.diag(Cpop[s]) for s in [0,1]] - no_core_idx = [numpy.where(Cpop[s] > 0.7)[0] for s in [0,1]] - C_ = [C_[s][:,no_core_idx[s]] for s in [0,1]] - S_ = [functools.reduce(numpy.dot, (C_[s].T, self.S, C_[s])) for s in [0,1]] + Cpop = [ + functools.reduce(numpy.dot, (C_[s].T, self.S, C_[s])) + for s in [0, 1] + ] + Cpop = [numpy.diag(Cpop[s]) for s in [0, 1]] + no_core_idx = [numpy.where(Cpop[s] > 0.7)[0] for s in [0, 1]] + C_ = [C_[s][:, no_core_idx[s]] for s in [0, 1]] + S_ = [ + functools.reduce(numpy.dot, (C_[s].T, self.S, C_[s])) + for s in [0, 1] + ] W_ = [] - for s in [0,1]: + for s in [0, 1]: es_, vs_ = eigh(S_[s]) s_ = numpy.sqrt(es_) - s_ = numpy.diag(1.0/s_) + s_ = numpy.diag(1.0 / s_) W_.append(functools.reduce(numpy.dot, (vs_, s_, vs_.T))) - self.W = [numpy.dot(C_[s], W_[s]) for s in [0,1]] + self.W = [numpy.dot(C_[s], W_[s]) for s in [0, 1]] else: P_core = numpy.eye(self.W.shape[0]) - numpy.dot(self.P_core, self.S) C_ = numpy.dot(P_core, self.W) # NOTE: PYSCF has basis in 1s2s3s2p2p2p3p3p3p format # fix no_core_idx - use population for now - Cpop = functools.reduce(numpy.dot, - (C_.T, self.S, C_)) + Cpop = functools.reduce(numpy.dot, (C_.T, self.S, C_)) Cpop = numpy.diag(Cpop) no_core_idx = numpy.where(Cpop > 0.7)[0] - C_ = C_[:,no_core_idx] + C_ = C_[:, no_core_idx] S_ = functools.reduce(numpy.dot, (C_.T, self.S, C_)) es_, vs_ = eigh(S_) s_ = numpy.sqrt(es_) - s_ = numpy.diag(1.0/s_) - W_ = functools.reduce(numpy.dot, - (vs_, s_, vs_.T)) + s_ = numpy.diag(1.0 / s_) + W_ = functools.reduce(numpy.dot, (vs_, s_, vs_.T)) self.W = numpy.dot(C_, W_) if self.unrestricted: if self.frozen_core: - self.lmo_coeff_a = functools.reduce(numpy.dot, - (self.W[0].T, self.S, self.C_a[:,self.ncore:])) - self.lmo_coeff_b = functools.reduce(numpy.dot, - (self.W[1].T, self.S, self.C_b[:,self.ncore:])) + self.lmo_coeff_a = functools.reduce( + numpy.dot, (self.W[0].T, self.S, self.C_a[:, self.ncore :]) + ) + self.lmo_coeff_b = functools.reduce( + numpy.dot, (self.W[1].T, self.S, self.C_b[:, self.ncore :]) + ) else: - self.lmo_coeff_a = functools.reduce(numpy.dot, - (self.W.T, self.S, self.C_a)) - self.lmo_coeff_b = functools.reduce(numpy.dot, - (self.W.T, self.S, self.C_b)) + self.lmo_coeff_a = functools.reduce( + numpy.dot, (self.W.T, self.S, self.C_a) + ) + self.lmo_coeff_b = functools.reduce( + numpy.dot, (self.W.T, self.S, self.C_b) + ) else: if self.frozen_core: - self.lmo_coeff = functools.reduce(numpy.dot, - (self.W.T, self.S, self.C[:,self.ncore:])) + self.lmo_coeff = functools.reduce( + numpy.dot, (self.W.T, self.S, self.C[:, self.ncore :]) + ) else: - self.lmo_coeff = functools.reduce(numpy.dot, - (self.W.T, self.S, self.C)) - - elif lo_method in ['pipek-mezey','pipek', 'PM']: + self.lmo_coeff = functools.reduce(numpy.dot, (self.W.T, self.S, self.C)) + elif lo_method in ["pipek-mezey", "pipek", "PM"]: es_, vs_ = eigh(self.S) - edx = es_ > 1.e-15 - self.W = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].T) + edx = es_ > 1.0e-15 + self.W = numpy.dot(vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].T) es_, vs_ = eigh(self.S) - edx = es_ > 1.e-15 - W_ = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].T) + edx = es_ > 1.0e-15 + W_ = numpy.dot(vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].T) if self.frozen_core: P_core = numpy.eye(W_.shape[0]) - numpy.dot(self.P_core, self.S) C_ = numpy.dot(P_core, W_) - Cpop = functools.reduce(numpy.dot, - (C_.T, self.S, C_)) + Cpop = functools.reduce(numpy.dot, (C_.T, self.S, C_)) Cpop = numpy.diag(Cpop) no_core_idx = numpy.where(Cpop > 0.55)[0] - C_ = C_[:,no_core_idx] + C_ = C_[:, no_core_idx] S_ = functools.reduce(numpy.dot, (C_.T, self.S, C_)) es_, vs_ = eigh(S_) s_ = numpy.sqrt(es_) - s_ = numpy.diag(1.0/s_) - W_ = functools.reduce(numpy.dot, - (vs_, s_, vs_.T)) + s_ = numpy.diag(1.0 / s_) + W_ = functools.reduce(numpy.dot, (vs_, s_, vs_.T)) W_ = numpy.dot(C_, W_) - self.W = get_loc(self.mol, W_, 'PM', pop_method=pop_method, init_guess=init_guess) + self.W = get_loc( + self.mol, W_, "PM", pop_method=pop_method, init_guess=init_guess + ) if not self.frozen_core: self.lmo_coeff = self.W.T @ self.S @ self.C else: - self.lmo_coeff = self.W.T @ self.S @ self.C[:,self.ncore:] + self.lmo_coeff = self.W.T @ self.S @ self.C[:, self.ncore :] - elif lo_method=='iao': + elif lo_method == "iao": from pyscf import lo import os, h5py - loc_type = 'SO' - val_basis = 'sto-3g' + loc_type = "SO" + val_basis = "sto-3g" # Occupied mo_coeff (with core) - Co = self.C[:,:self.Nocc] + Co = self.C[:, : self.Nocc] # Get necessary overlaps, second arg is IAO basis S12, S2 = get_xovlp(self.mol, basis=val_basis) # Use these to get IAOs - Ciao = get_iao(Co, S12, self.S, S2 = S2) + Ciao = get_iao(Co, S12, self.S, S2=S2) if not valence_only: # Now get PAOs - if loc_type.upper() != 'SO': + if loc_type.upper() != "SO": Cpao = get_pao(Ciao, self.S, S12, S2, self.mol) - elif loc_type.upper() == 'SO': + elif loc_type.upper() == "SO": Cpao = get_pao_native(Ciao, self.S, self.mol, valence_basis=val_basis) # rearrange by atom @@ -358,11 +395,11 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', if self.frozen_core: # Remove core MOs - Cc = self.C[:,:self.ncore] # Assumes core are first + Cc = self.C[:, : self.ncore] # Assumes core are first Ciao = remove_core_mo(Ciao, Cc, self.S) # Localize orbitals beyond symm orth - if loc_type.upper() != 'SO': + if loc_type.upper() != "SO": Ciao = get_loc(self.mol, Ciao, loc_type) if not valence_only: Cpao = get_loc(self.mol, Cpao, loc_type) @@ -370,7 +407,9 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', shift = 0 ncore = 0 if not valence_only: - Wstack = numpy.zeros((Ciao.shape[0], Ciao.shape[1]+Cpao.shape[1])) #-self.ncore)) + Wstack = numpy.zeros( + (Ciao.shape[0], Ciao.shape[1] + Cpao.shape[1]) + ) # -self.ncore)) else: Wstack = numpy.zeros((Ciao.shape[0], Ciao.shape[1])) @@ -379,81 +418,86 @@ def localize(self, lo_method, mol=None, valence_basis='sto-3g', nc = ncore_(self.mol.atom_charge(ix)) ncore += nc niao = len(iaoind_by_atom[ix]) - iaoind_ix = [ i_ - ncore for i_ in iaoind_by_atom[ix][nc:]] - Wstack[:, shift:shift+niao-nc] = Ciao[:, iaoind_ix] - shift += niao-nc + iaoind_ix = [i_ - ncore for i_ in iaoind_by_atom[ix][nc:]] + Wstack[:, shift : shift + niao - nc] = Ciao[:, iaoind_ix] + shift += niao - nc if not valence_only: npao = len(paoind_by_atom[ix]) - Wstack[:,shift:shift+npao] = Cpao[:, paoind_by_atom[ix]] + Wstack[:, shift : shift + npao] = Cpao[:, paoind_by_atom[ix]] shift += npao else: if not hstack: for ix in range(self.mol.natm): niao = len(iaoind_by_atom[ix]) - Wstack[:, shift:shift+niao] = Ciao[:, iaoind_by_atom[ix]] + Wstack[:, shift : shift + niao] = Ciao[:, iaoind_by_atom[ix]] shift += niao if not valence_only: npao = len(paoind_by_atom[ix]) - Wstack[:,shift:shift+npao] = Cpao[:, paoind_by_atom[ix]] + Wstack[:, shift : shift + npao] = Cpao[:, paoind_by_atom[ix]] shift += npao else: Wstack = numpy.hstack((Ciao, Cpao)) if not nosave: self.W = Wstack - assert(numpy.allclose(self.W.T @ self.S @ self.W, numpy.eye(self.W.shape[1]))) + assert numpy.allclose( + self.W.T @ self.S @ self.W, numpy.eye(self.W.shape[1]) + ) else: - assert(numpy.allclose(Wstack.T @ self.S @ Wstack, numpy.eye(Wstack.shape[1]))) + assert numpy.allclose( + Wstack.T @ self.S @ Wstack, numpy.eye(Wstack.shape[1]) + ) return Wstack nmo = self.C.shape[1] - self.ncore nlo = self.W.shape[1] if not valence_only: if nmo > nlo: - Co_nocore = self.C[:,self.ncore:self.Nocc] - Cv = self.C[:,self.Nocc:] + Co_nocore = self.C[:, self.ncore : self.Nocc] + Cv = self.C[:, self.Nocc :] # Ensure that the LOs span the occupied space - assert(numpy.allclose(numpy.sum((self.W.T @ self.S @ Co_nocore)**2.), - self.Nocc - self.ncore)) + assert numpy.allclose( + numpy.sum((self.W.T @ self.S @ Co_nocore) ** 2.0), + self.Nocc - self.ncore, + ) # Find virtual orbitals that lie in the span of LOs u, l, vt = numpy.linalg.svd(self.W.T @ self.S @ Cv, full_matrices=False) nvlo = nlo - self.Nocc - self.ncore - assert(numpy.allclose(numpy.sum(l[:nvlo]), nvlo)) + assert numpy.allclose(numpy.sum(l[:nvlo]), nvlo) C_ = numpy.hstack([Co_nocore, Cv @ vt[:nvlo].T]) self.lmo_coeff = self.W.T @ self.S @ C_ else: - self.lmo_coeff = self.W.T @ self.S @ self.C[:,self.ncore:] + self.lmo_coeff = self.W.T @ self.S @ self.C[:, self.ncore :] else: - self.lmo_coeff = self.W.T @ self.S @ self.C[:,self.ncore:] + self.lmo_coeff = self.W.T @ self.S @ self.C[:, self.ncore :] - elif lo_method == 'boys': + elif lo_method == "boys": from pyscf.lo import Boys + es_, vs_ = eigh(self.S) - edx = es_ > 1.e-15 - W_ = numpy.dot(vs_[:,edx]/numpy.sqrt(es_[edx]), vs_[:,edx].T) + edx = es_ > 1.0e-15 + W_ = numpy.dot(vs_[:, edx] / numpy.sqrt(es_[edx]), vs_[:, edx].T) if self.frozen_core: P_core = numpy.eye(W_.shape[0]) - numpy.dot(self.P_core, self.S) C_ = numpy.dot(P_core, W_) - Cpop = functools.reduce(numpy.dot, - (C_.T, self.S, C_)) + Cpop = functools.reduce(numpy.dot, (C_.T, self.S, C_)) Cpop = numpy.diag(Cpop) no_core_idx = numpy.where(Cpop > 0.55)[0] - C_ = C_[:,no_core_idx] + C_ = C_[:, no_core_idx] S_ = functools.reduce(numpy.dot, (C_.T, self.S, C_)) es_, vs_ = eigh(S_) s_ = numpy.sqrt(es_) - s_ = numpy.diag(1.0/s_) - W_ = functools.reduce(numpy.dot, - (vs_, s_, vs_.T)) + s_ = numpy.diag(1.0 / s_) + W_ = functools.reduce(numpy.dot, (vs_, s_, vs_.T)) W_ = numpy.dot(C_, W_) - self.W = get_loc(self.mol, W_, 'BOYS') + self.W = get_loc(self.mol, W_, "BOYS") if not self.frozen_core: self.lmo_coeff = self.W.T @ self.S @ self.C else: - self.lmo_coeff = self.W.T @ self.S @ self.C[:,self.ncore:] + self.lmo_coeff = self.W.T @ self.S @ self.C[:, self.ncore :] else: - print('lo_method = ',lo_method,' not implemented!',flush=True) - print('exiting',flush=True) + print("lo_method = ", lo_method, " not implemented!", flush=True) + print("exiting", flush=True) sys.exit() diff --git a/molbe/mbe.py b/molbe/mbe.py index d9221a8e..b620db70 100644 --- a/molbe/mbe.py +++ b/molbe/mbe.py @@ -3,15 +3,30 @@ from .pfrag import Frags from .helper import get_core import molbe.be_var as be_var -import numpy,functools,sys, pickle +import numpy, functools, sys, pickle from pyscf import lib -import h5py,os,time +import h5py, os, time + class storeBE: - def __init__(self, Nocc, hf_veff, hcore, - S, C, hf_dm, hf_etot, W, lmo_coeff, - enuc, - E_core, C_core, P_core, core_veff, mo_energy): + def __init__( + self, + Nocc, + hf_veff, + hcore, + S, + C, + hf_dm, + hf_etot, + W, + lmo_coeff, + enuc, + E_core, + C_core, + P_core, + core_veff, + mo_energy, + ): self.Nocc = Nocc self.hf_veff = hf_veff self.hcore = hcore @@ -28,6 +43,7 @@ def __init__(self, Nocc, hf_veff, hcore, self.core_veff = core_veff self.mo_energy = mo_energy + class BE: """ Class for handling bootstrap embedding (BE) calculations. @@ -48,15 +64,29 @@ class BE: Method for orbital localization, default is 'lowdin'. """ - def __init__(self, mf, fobj, eri_file='eri_file.h5', - lo_method='lowdin', pop_method=None, compute_hf=True, - restart=False, save=False, - restart_file='storebe.pk', - mo_energy = None, - save_file='storebe.pk',hci_pt=False, - nproc=1, ompnum=4, scratch_dir=None, - hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, - integral_direct_DF=False, auxbasis = None): + def __init__( + self, + mf, + fobj, + eri_file="eri_file.h5", + lo_method="lowdin", + pop_method=None, + compute_hf=True, + restart=False, + save=False, + restart_file="storebe.pk", + mo_energy=None, + save_file="storebe.pk", + hci_pt=False, + nproc=1, + ompnum=4, + scratch_dir=None, + hci_cutoff=0.001, + ci_coeff_cutoff=None, + select_cutoff=None, + integral_direct_DF=False, + auxbasis=None, + ): """ Constructor for BE object. @@ -94,7 +124,7 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', if restart: # Load previous calculation data from restart file - with open(restart_file, 'rb') as rfile: + with open(restart_file, "rb") as rfile: store_ = pickle.load(rfile) rfile.close() self.Nocc = store_.Nocc @@ -120,7 +150,7 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.auxbasis = auxbasis # Fragment information from fobj - self.frag_type=fobj.frag_type + self.frag_type = fobj.frag_type self.Nfrag = fobj.Nfrag self.fsites = fobj.fsites self.edge = fobj.edge @@ -132,21 +162,21 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.be_type = fobj.be_type self.mol = fobj.mol - self.ebe_hf = 0. - self.ebe_tot = 0. + self.ebe_hf = 0.0 + self.ebe_tot = 0.0 # HCI parameters self.hci_cutoff = hci_cutoff self.ci_coeff_cutoff = ci_coeff_cutoff self.select_cutoff = select_cutoff - self.hci_pt=hci_pt + self.hci_pt = hci_pt self.mf = mf if not restart: self.mo_energy = mf.mo_energy self.mf = mf - self.Nocc = mf.mol.nelectron//2 + self.Nocc = mf.mol.nelectron // 2 self.enuc = mf.energy_nuc() self.hcore = mf.get_hcore() @@ -166,21 +196,21 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.scratch_dir = scratch_dir # Set scratch directory - jobid='' + jobid = "" if be_var.CREATE_SCRATCH_DIR: try: - jobid = str(os.environ['SLURM_JOB_ID']) + jobid = str(os.environ["SLURM_JOB_ID"]) except: - jobid = '' - if not be_var.SCRATCH=='': - self.scratch_dir = be_var.SCRATCH+str(jobid) - os.system('mkdir -p '+self.scratch_dir) + jobid = "" + if not be_var.SCRATCH == "": + self.scratch_dir = be_var.SCRATCH + str(jobid) + os.system("mkdir -p " + self.scratch_dir) else: self.scratch_dir = None - if jobid == '': - self.eri_file = be_var.SCRATCH+eri_file + if jobid == "": + self.eri_file = be_var.SCRATCH + eri_file else: - self.eri_file = self.scratch_dir+'/'+eri_file + self.eri_file = self.scratch_dir + "/" + eri_file self.frozen_core = False if not fobj.frozen_core else True self.ncore = 0 @@ -197,43 +227,70 @@ def __init__(self, mf, fobj, eri_file='eri_file.h5', self.core_list = fobj.core_list if not restart: - self.Nocc -=self.ncore - self.hf_dm = 2.*numpy.dot(self.C[:,self.ncore:self.ncore+self.Nocc], - self.C[:,self.ncore:self.ncore+self.Nocc].T) - self.C_core = self.C[:,:self.ncore] + self.Nocc -= self.ncore + self.hf_dm = 2.0 * numpy.dot( + self.C[:, self.ncore : self.ncore + self.Nocc], + self.C[:, self.ncore : self.ncore + self.Nocc].T, + ) + self.C_core = self.C[:, : self.ncore] self.P_core = numpy.dot(self.C_core, self.C_core.T) - self.core_veff = mf.get_veff(dm = self.P_core*2.) - self.E_core = numpy.einsum('ji,ji->',2.*self.hcore+self.core_veff, self.P_core) + self.core_veff = mf.get_veff(dm=self.P_core * 2.0) + self.E_core = numpy.einsum( + "ji,ji->", 2.0 * self.hcore + self.core_veff, self.P_core + ) self.hf_veff -= self.core_veff self.hcore += self.core_veff if not restart: # Localize orbitals - self.localize(lo_method, pop_method=pop_method, mol=self.mol, valence_basis=fobj.valence_basis, valence_only=fobj.valence_only) - - if fobj.valence_only and lo_method=='iao': - self.Ciao_pao = self.localize(lo_method, pop_method=pop_method, mol=self.mol, valence_basis=fobj.valence_basis, - hstack=True, - valence_only=False, nosave=True) + self.localize( + lo_method, + pop_method=pop_method, + mol=self.mol, + valence_basis=fobj.valence_basis, + valence_only=fobj.valence_only, + ) + + if fobj.valence_only and lo_method == "iao": + self.Ciao_pao = self.localize( + lo_method, + pop_method=pop_method, + mol=self.mol, + valence_basis=fobj.valence_basis, + hstack=True, + valence_only=False, + nosave=True, + ) if save: # Save intermediate results for restart - store_ = storeBE(self.Nocc, self.hf_veff, self.hcore, - self.S, self.C, self.hf_dm, self.hf_etot, - self.W, self.lmo_coeff, self.enuc, - self.E_core, self.C_core, self.P_core, self.core_veff, self.mo_energy) - - with open(save_file, 'wb') as rfile: + store_ = storeBE( + self.Nocc, + self.hf_veff, + self.hcore, + self.S, + self.C, + self.hf_dm, + self.hf_etot, + self.W, + self.lmo_coeff, + self.enuc, + self.E_core, + self.C_core, + self.P_core, + self.core_veff, + self.mo_energy, + ) + + with open(save_file, "wb") as rfile: pickle.dump(store_, rfile, pickle.HIGHEST_PROTOCOL) rfile.close() - - if not restart : + if not restart: # Initialize fragments and perform initial calculations - self.initialize(mf._eri,compute_hf) + self.initialize(mf._eri, compute_hf) else: - self.initialize(None,compute_hf, restart=True) - + self.initialize(None, compute_hf, restart=True) from ._opt import optimize from molbe.external.optqn import get_be_error_jacobian @@ -244,26 +301,23 @@ def print_ini(self): """ Print initialization banner for the MOLBE calculation. """ - print('-----------------------------------------------------------', - flush=True) + print("-----------------------------------------------------------", flush=True) - print(' MMM MMM OOOO LL BBBBBBB EEEEEEE ',flush=True) - print(' M MM MM M OO OO LL BB B EE ',flush=True) - print(' M MM MM M OO OO LL BB B EE ',flush=True) - print(' M MMM M OO OO LL === BBBBBBB EEEEEEE ',flush=True) - print(' M M OO OO LL BB B EE ',flush=True) - print(' M M OO OO LL BB B EE ',flush=True) - print(' M M OOOO LLLLLL BBBBBBB EEEEEEE ',flush=True) + print(" MMM MMM OOOO LL BBBBBBB EEEEEEE ", flush=True) + print(" M MM MM M OO OO LL BB B EE ", flush=True) + print(" M MM MM M OO OO LL BB B EE ", flush=True) + print(" M MMM M OO OO LL === BBBBBBB EEEEEEE ", flush=True) + print(" M M OO OO LL BB B EE ", flush=True) + print(" M M OO OO LL BB B EE ", flush=True) + print(" M M OOOO LLLLLL BBBBBBB EEEEEEE ", flush=True) print(flush=True) - print(' MOLECULAR BOOTSTRAP EMBEDDING',flush=True) - print(' BEn = ',self.be_type,flush=True) - print('-----------------------------------------------------------', - flush=True) + print(" MOLECULAR BOOTSTRAP EMBEDDING", flush=True) + print(" BEn = ", self.be_type, flush=True) + print("-----------------------------------------------------------", flush=True) print(flush=True) - - def initialize(self, eri_,compute_hf, restart=False): + def initialize(self, eri_, compute_hf, restart=False): """ Initialize the Bootstrap Embedding calculation. @@ -281,25 +335,38 @@ def initialize(self, eri_,compute_hf, restart=False): from pyscf import ao2mo from multiprocessing import Pool - if compute_hf: E_hf = 0. + if compute_hf: + E_hf = 0.0 # Create a file to store ERIs if not restart: - file_eri = h5py.File(self.eri_file,'w') + file_eri = h5py.File(self.eri_file, "w") lentmp = len(self.edge_idx) for I in range(self.Nfrag): - if lentmp: - fobjs_ = Frags(self.fsites[I], I, edge=self.edge[I], - eri_file=self.eri_file, - center=self.center[I], edge_idx=self.edge_idx[I], - center_idx=self.center_idx[I],efac=self.ebe_weight[I], - centerf_idx=self.centerf_idx[I]) + fobjs_ = Frags( + self.fsites[I], + I, + edge=self.edge[I], + eri_file=self.eri_file, + center=self.center[I], + edge_idx=self.edge_idx[I], + center_idx=self.center_idx[I], + efac=self.ebe_weight[I], + centerf_idx=self.centerf_idx[I], + ) else: - fobjs_ = Frags(self.fsites[I],I,edge=[],center=[], - eri_file=self.eri_file, - edge_idx=[],center_idx=[],centerf_idx=[], - efac=self.ebe_weight[I]) + fobjs_ = Frags( + self.fsites[I], + I, + edge=[], + center=[], + eri_file=self.eri_file, + edge_idx=[], + center_idx=[], + centerf_idx=[], + efac=self.ebe_weight[I], + ) fobjs_.sd(self.W, self.lmo_coeff, self.Nocc) self.Fobjs.append(fobjs_) @@ -312,27 +379,38 @@ def initialize(self, eri_,compute_hf, restart=False): # No -- Do we have (ij|P) from density fitting? # Yes -- ao2mo, outcore version, using saved (ij|P) # No -- if integral_direct_DF is requested, invoke on-the-fly routine - assert (not eri_ is None) or (hasattr(self.mf, 'with_df')) or (self.integral_direct_DF), "Input mean-field object is missing ERI (mf._eri) or DF (mf.with_df) object AND integral direct DF routine was not requested. Please check your inputs." - if not eri_ is None: # incore ao2mo using saved eri from mean-field calculation + assert ( + (not eri_ is None) + or (hasattr(self.mf, "with_df")) + or (self.integral_direct_DF) + ), "Input mean-field object is missing ERI (mf._eri) or DF (mf.with_df) object AND integral direct DF routine was not requested. Please check your inputs." + if ( + not eri_ is None + ): # incore ao2mo using saved eri from mean-field calculation for I in range(self.Nfrag): eri = ao2mo.incore.full(eri_, self.Fobjs[I].TA, compact=True) file_eri.create_dataset(self.Fobjs[I].dname, data=eri) - elif hasattr(self.mf, 'with_df') and not self.mf.with_df is None: + elif hasattr(self.mf, "with_df") and not self.mf.with_df is None: # pyscf.ao2mo uses DF object in an outcore fashion using (ij|P) in pyscf temp directory for I in range(self.Nfrag): eri = self.mf.with_df.ao2mo(self.Fobjs[I].TA, compact=True) file_eri.create_dataset(self.Fobjs[I].dname, data=eri) else: # If ERIs are not saved on memory, compute fragment ERIs integral-direct - if self.integral_direct_DF: # Use density fitting to generate fragment ERIs on-the-fly + if ( + self.integral_direct_DF + ): # Use density fitting to generate fragment ERIs on-the-fly from .eri_onthefly import integral_direct_DF - integral_direct_DF(self.mf, self.Fobjs, file_eri, auxbasis=self.auxbasis) - else: # Calculate ERIs on-the-fly to generate fragment ERIs + + integral_direct_DF( + self.mf, self.Fobjs, file_eri, auxbasis=self.auxbasis + ) + else: # Calculate ERIs on-the-fly to generate fragment ERIs # TODO: Future feature to be implemented # NOTE: Ideally, we want AO shell pair screening for this. return NotImplementedError else: - eri=None + eri = None for fobjs_ in self.Fobjs: # Process each fragment @@ -349,11 +427,15 @@ def initialize(self, eri_,compute_hf, restart=False): fobjs_.heff = numpy.zeros_like(fobjs_.h1) fobjs_.scf(fs=True, eri=eri) - fobjs_.dm0 = numpy.dot( fobjs_._mo_coeffs[:,:fobjs_.nsocc], - fobjs_._mo_coeffs[:,:fobjs_.nsocc].conj().T) *2. + fobjs_.dm0 = ( + numpy.dot( + fobjs_._mo_coeffs[:, : fobjs_.nsocc], + fobjs_._mo_coeffs[:, : fobjs_.nsocc].conj().T, + ) + * 2.0 + ) if compute_hf: - eh1, ecoul, ef = fobjs_.energy_hf(return_e1=True) E_hf += fobjs_.ebe_hf @@ -361,13 +443,14 @@ def initialize(self, eri_,compute_hf, restart=False): file_eri.close() if compute_hf: - - self.ebe_hf = E_hf+self.enuc+self.E_core + self.ebe_hf = E_hf + self.enuc + self.E_core hf_err = self.hf_etot - self.ebe_hf - print('HF-in-HF error : {:>.4e} Ha'. - format(hf_err), flush=True) - if abs(hf_err)>1.e-5: - print('WARNING!!! Large HF-in-HF energy error') + print( + "HF-in-HF error : {:>.4e} Ha".format(hf_err), + flush=True, + ) + if abs(hf_err) > 1.0e-5: + print("WARNING!!! Large HF-in-HF energy error") print(flush=True) @@ -376,8 +459,16 @@ def initialize(self, eri_,compute_hf, restart=False): fobj.udim = couti couti = fobj.set_udim(couti) - def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False, - scratch_dir=None, **solver_kwargs): + def oneshot( + self, + solver="MP2", + nproc=1, + ompnum=4, + calc_frag_energy=False, + clean_eri=False, + scratch_dir=None, + **solver_kwargs, + ): """ Perform a one-shot bootstrap embedding calculation. @@ -402,31 +493,63 @@ def oneshot(self, solver='MP2', nproc=1, ompnum=4, calc_frag_energy=False, clean print("Calculating Energy by Fragment? ", calc_frag_energy) if nproc == 1: - rets = be_func(None, self.Fobjs, self.Nocc, solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff = self.ci_coeff_cutoff, - select_cutoff = self.select_cutoff, - nproc=ompnum, frag_energy=calc_frag_energy, - ereturn=True, eeval=True, scratch_dir=self.scratch_dir, **self.solver_kwargs) + rets = be_func( + None, + self.Fobjs, + self.Nocc, + solver, + self.enuc, + hf_veff=self.hf_veff, + hci_cutoff=self.hci_cutoff, + ci_coeff_cutoff=self.ci_coeff_cutoff, + select_cutoff=self.select_cutoff, + nproc=ompnum, + frag_energy=calc_frag_energy, + ereturn=True, + eeval=True, + scratch_dir=self.scratch_dir, + **self.solver_kwargs, + ) else: - rets = be_func_parallel(None, self.Fobjs, self.Nocc, solver, self.enuc, hf_veff=self.hf_veff, - hci_cutoff=self.hci_cutoff, - ci_coeff_cutoff = self.ci_coeff_cutoff, - select_cutoff = self.select_cutoff, - ereturn=True, eeval=True, frag_energy=calc_frag_energy, - nproc=nproc, ompnum=ompnum, scratch_dir=self.scratch_dir, **self.solver_kwargs) - - print('-----------------------------------------------------', - flush=True) - print(' One Shot BE ', flush=True) - print(' Solver : ',solver,flush=True) - print('-----------------------------------------------------', - flush=True) + rets = be_func_parallel( + None, + self.Fobjs, + self.Nocc, + solver, + self.enuc, + hf_veff=self.hf_veff, + hci_cutoff=self.hci_cutoff, + ci_coeff_cutoff=self.ci_coeff_cutoff, + select_cutoff=self.select_cutoff, + ereturn=True, + eeval=True, + frag_energy=calc_frag_energy, + nproc=nproc, + ompnum=ompnum, + scratch_dir=self.scratch_dir, + **self.solver_kwargs, + ) + + print("-----------------------------------------------------", flush=True) + print(" One Shot BE ", flush=True) + print(" Solver : ", solver, flush=True) + print("-----------------------------------------------------", flush=True) print(flush=True) if calc_frag_energy: - print("Final Tr(F del g) is : {:>12.8f} Ha".format(rets[1][0]+rets[1][2]), flush=True) - print("Final Tr(V K_approx) is : {:>12.8f} Ha".format(rets[1][1]), flush=True) - print("Final e_corr is : {:>12.8f} Ha".format(rets[0]), flush=True) + print( + "Final Tr(F del g) is : {:>12.8f} Ha".format( + rets[1][0] + rets[1][2] + ), + flush=True, + ) + print( + "Final Tr(V K_approx) is : {:>12.8f} Ha".format(rets[1][1]), + flush=True, + ) + print( + "Final e_corr is : {:>12.8f} Ha".format(rets[0]), + flush=True, + ) self.ebe_tot = rets[0] @@ -456,7 +579,7 @@ def update_fock(self, heff=None): for idx, fobj in self.Fobjs: fobj.fock += heff[idx] - def write_heff(self, heff_file='bepotfile.h5'): + def write_heff(self, heff_file="bepotfile.h5"): """ Write the effective Hamiltonian to a file. @@ -465,13 +588,13 @@ def write_heff(self, heff_file='bepotfile.h5'): heff_file : str, optional Path to the file to store effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, 'w') + filepot = h5py.File(heff_file, "w") for fobj in self.Fobjs: print(fobj.heff.shape, fobj.dname, flush=True) filepot.create_dataset(fobj.dname, data=fobj.heff) filepot.close() - def read_heff(self, heff_file='bepotfile.h5'): + def read_heff(self, heff_file="bepotfile.h5"): """ Read the effective Hamiltonian from a file. @@ -480,13 +603,12 @@ def read_heff(self, heff_file='bepotfile.h5'): heff_file : str, optional Path to the file storing effective Hamiltonian, by default 'bepotfile.h5'. """ - filepot = h5py.File(heff_file, 'r') + filepot = h5py.File(heff_file, "r") for fobj in self.Fobjs: fobj.heff = filepot.get(fobj.dname) filepot.close() - def initialize_pot(Nfrag, edge_idx): """ Initialize the potential array for bootstrap embedding. @@ -509,17 +631,16 @@ def initialize_pot(Nfrag, edge_idx): list of float Initialized potential array with zeros. """ - pot_=[] + pot_ = [] if not len(edge_idx) == 0: for I in range(Nfrag): for i in edge_idx[I]: for j in range(len(i)): for k in range(len(i)): - if j>k: + if j > k: continue - pot_.append(0.) + pot_.append(0.0) - pot_.append(0.) + pot_.append(0.0) return pot_ - diff --git a/molbe/misc.py b/molbe/misc.py index c52b632f..9387f37c 100644 --- a/molbe/misc.py +++ b/molbe/misc.py @@ -6,8 +6,17 @@ from pyscf.lib import chkfile -def libint2pyscf(xyzfile, hcore, basis, hcore_skiprows=1, - use_df=False, unrestricted=False, spin=0, charge=0): + +def libint2pyscf( + xyzfile, + hcore, + basis, + hcore_skiprows=1, + use_df=False, + unrestricted=False, + spin=0, + charge=0, +): """Build a pyscf Mole and RHF/UHF object using the given xyz file and core Hamiltonian (in libint standard format) c.f. @@ -83,7 +92,8 @@ def libint2pyscf(xyzfile, hcore, basis, hcore_skiprows=1, mydf = df.DF(mol).build() mf.with_df = mydf - else: mf = scf.UHF(mol) if unrestricted else scf.RHF(mol) + else: + mf = scf.UHF(mol) if unrestricted else scf.RHF(mol) mf.get_hcore = lambda *args: hcore_pyscf return mol, mf @@ -135,7 +145,12 @@ def be2fcidump(be_obj, fcidump_prefix, basis): raise Exception("Basis should be either embedding or fragment_mo") fcidump.from_integrals( - fcidump_prefix + "f" + str(fidx), h1e, h2e, frag.TA.shape[1], frag.nsocc, ms=0 + fcidump_prefix + "f" + str(fidx), + h1e, + h2e, + frag.TA.shape[1], + frag.nsocc, + ms=0, ) @@ -185,7 +200,12 @@ def ube2fcidump(be_obj, fcidump_prefix, basis): raise Exception("Basis should be either embedding or fragment_mo") fcidump.from_integrals( - fcidump_prefix + "f" + str(fidx) + "a", h1e, h2e, frag.TA.shape[1], frag.nsocc, ms=0 + fcidump_prefix + "f" + str(fidx) + "a", + h1e, + h2e, + frag.TA.shape[1], + frag.nsocc, + ms=0, ) for fidx, frag in enumerate(be_obj.Fobjs_b): @@ -215,7 +235,12 @@ def ube2fcidump(be_obj, fcidump_prefix, basis): raise Exception("Basis should be either embedding or fragment_mo") fcidump.from_integrals( - fcidump_prefix + "f" + str(fidx) + "b", h1e, h2e, frag.TA.shape[1], frag.nsocc, ms=0 + fcidump_prefix + "f" + str(fidx) + "b", + h1e, + h2e, + frag.TA.shape[1], + frag.nsocc, + ms=0, ) @@ -231,15 +256,15 @@ def be2puffin( spin=0, nproc=1, ompnum=1, - be_type='be1', + be_type="be1", df_aux_basis=None, frozen_core=True, - localization_method='lowdin', + localization_method="lowdin", localization_basis=None, unrestricted=False, from_chk=False, checkfile=None, - ecp=None + ecp=None, ): """Front-facing API bridge tailored for SCINE Puffin Returns the CCSD oneshot energies @@ -307,9 +332,9 @@ def be2puffin( mol = gto.M(atom=xyzfile, basis=basis, charge=charge, spin=spin) if not from_chk: - if hcore is None: #from point charges OR with no external potential + if hcore is None: # from point charges OR with no external potential hcore_pyscf = None - else: #from starting Hamiltonian in Libint format + else: # from starting Hamiltonian in Libint format if libint_inp == True: libint2pyscf = [] for labelidx, label in enumerate(mol.ao_labels()): @@ -331,8 +356,12 @@ def be2puffin( hcore_pyscf = hcore if not jk is None: jk_pyscf = ( - jk[0][numpy.ix_(libint2pyscf, libint2pyscf, libint2pyscf, libint2pyscf)], - jk[1][numpy.ix_(libint2pyscf, libint2pyscf, libint2pyscf, libint2pyscf)], + jk[0][ + numpy.ix_(libint2pyscf, libint2pyscf, libint2pyscf, libint2pyscf) + ], + jk[1][ + numpy.ix_(libint2pyscf, libint2pyscf, libint2pyscf, libint2pyscf) + ], ) mol.incore_anyway = True @@ -343,34 +372,52 @@ def be2puffin( if hcore is None: if pts_and_charges: from pyscf import qmmm - print("Using QM/MM Point Charges: Assuming QM structure in Angstrom and MM Coordinates in Bohr !!!") - mf1 = scf.UHF(mol).set(max_cycle = 200) #using SOSCF is more reliable - #mf1 = scf.UHF(mol).set(max_cycle = 200, level_shift = (0.3, 0.2)) + + print( + "Using QM/MM Point Charges: Assuming QM structure in Angstrom and MM Coordinates in Bohr !!!" + ) + mf1 = scf.UHF(mol).set( + max_cycle=200 + ) # using SOSCF is more reliable + # mf1 = scf.UHF(mol).set(max_cycle = 200, level_shift = (0.3, 0.2)) # using level shift helps, but not always. level_shift and # scf.addons.dynamic_level_shift do not seem to work with QM/MM # note: from the SCINE database, the structure is in Angstrom but the MM point charges # are in Bohr !! - mf = qmmm.mm_charge(mf1, pts_and_charges[0], pts_and_charges[1], unit='bohr').newton() #mf object, coordinates, charges + mf = qmmm.mm_charge( + mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr" + ).newton() # mf object, coordinates, charges else: - mf = scf.UHF(mol).set(max_cycle = 200, level_shift = (0.3, 0.2)) + mf = scf.UHF(mol).set(max_cycle=200, level_shift=(0.3, 0.2)) else: - mf = scf.UHF(mol).set(max_cycle = 200).newton() - else: # restricted - if pts_and_charges: # running QM/MM + mf = scf.UHF(mol).set(max_cycle=200).newton() + else: # restricted + if pts_and_charges: # running QM/MM from pyscf import qmmm - print("Using QM/MM Point Charges: Assuming QM structure in Angstrom and MM Coordinates in Bohr !!!") - mf1 = scf.RHF(mol).set(max_cycle = 200) - mf = qmmm.mm_charge(mf1, pts_and_charges[0], pts_and_charges[1], unit='bohr').newton() - print("Setting use_df to false and jk to none: have not tested DF and QM/MM from point charges at the same time") + + print( + "Using QM/MM Point Charges: Assuming QM structure in Angstrom and MM Coordinates in Bohr !!!" + ) + mf1 = scf.RHF(mol).set(max_cycle=200) + mf = qmmm.mm_charge( + mf1, pts_and_charges[0], pts_and_charges[1], unit="bohr" + ).newton() + print( + "Setting use_df to false and jk to none: have not tested DF and QM/MM from point charges at the same time" + ) use_df = False jk = None elif use_df and jk is None: from pyscf import df + mf = scf.RHF(mol).density_fit(auxbasis=df_aux_basis) - else: mf = scf.RHF(mol) + else: + mf = scf.RHF(mol) - if not hcore is None: mf.get_hcore = lambda *args: hcore_pyscf - if not jk is None: mf.get_jk = lambda *args: jk_pyscf + if not hcore is None: + mf.get_hcore = lambda *args: hcore_pyscf + if not jk is None: + mf.get_jk = lambda *args: jk_pyscf if checkfile: print("Saving checkfile to:", checkfile) @@ -384,13 +431,17 @@ def be2puffin( print("Reference HF Unconverged -- stopping the calculation", flush=True) sys.exit() if use_df: - print("Using auxillary basis in density fitting: ", mf.with_df.auxmol.basis, flush=True) + print( + "Using auxillary basis in density fitting: ", + mf.with_df.auxmol.basis, + flush=True, + ) print("DF auxillary nao_nr", mf.with_df.auxmol.nao_nr(), flush=True) print("Time for mf kernel to run: ", time_post_mf - time_pre_mf, flush=True) elif from_chk: print("Running from chkfile", checkfile, flush=True) - scf_result_dic = chkfile.load(checkfile, 'scf') + scf_result_dic = chkfile.load(checkfile, "scf") if unrestricted: mf = scf.UHF(mol) else: @@ -406,43 +457,41 @@ def be2puffin( fobj = fragpart( be_type=be_type, frag_type="autogen", mol=mol, frozen_core=frozen_core - ) + ) time_post_fragpart = time.time() - print("Time for fragmentation to run: ", time_post_fragpart - time_post_mf, flush=True) + print( + "Time for fragmentation to run: ", time_post_fragpart - time_post_mf, flush=True + ) # Run embedding setup if unrestricted: mybe = UBE(mf, fobj, lo_method="lowdin") - solver="UCCSD" + solver = "UCCSD" else: mybe = BE(mf, fobj, lo_method="lowdin") - solver="CCSD" + solver = "CCSD" # Run oneshot embedding and return system energy - mybe.oneshot(solver=solver, nproc=nproc, ompnum=ompnum, calc_frag_energy=True, clean_eri=True) + mybe.oneshot( + solver=solver, nproc=nproc, ompnum=ompnum, calc_frag_energy=True, clean_eri=True + ) return mybe.ebe_tot - def print_energy(ecorr, e_V_Kapprox, e_F_dg, e_hf): - - # Print energy results - print('-----------------------------------------------------', - flush=True) - print(' BE ENERGIES with cumulant-based expression', flush=True) - - print('-----------------------------------------------------', - flush=True) - print(' E_BE = E_HF + Tr(F del g) + Tr(V K_approx)', flush=True) - print(' E_HF : {:>14.8f} Ha'.format(e_hf), flush=True) - print(' Tr(F del g) : {:>14.8f} Ha'.format(e_F_dg), flush=True) - print(' Tr(V K_aprrox) : {:>14.8f} Ha'.format(e_V_Kapprox), flush=True) - print(' E_BE : {:>14.8f} Ha'.format(ecorr + e_hf), flush=True) - print(' Ecorr BE : {:>14.8f} Ha'.format(ecorr), flush=True) - print('-----------------------------------------------------', - flush=True) + print("-----------------------------------------------------", flush=True) + print(" BE ENERGIES with cumulant-based expression", flush=True) + + print("-----------------------------------------------------", flush=True) + print(" E_BE = E_HF + Tr(F del g) + Tr(V K_approx)", flush=True) + print(" E_HF : {:>14.8f} Ha".format(e_hf), flush=True) + print(" Tr(F del g) : {:>14.8f} Ha".format(e_F_dg), flush=True) + print(" Tr(V K_aprrox) : {:>14.8f} Ha".format(e_V_Kapprox), flush=True) + print(" E_BE : {:>14.8f} Ha".format(ecorr + e_hf), flush=True) + print(" Ecorr BE : {:>14.8f} Ha".format(ecorr), flush=True) + print("-----------------------------------------------------", flush=True) print(flush=True) diff --git a/molbe/pfrag.py b/molbe/pfrag.py index e9c45540..85bfcb67 100644 --- a/molbe/pfrag.py +++ b/molbe/pfrag.py @@ -2,8 +2,8 @@ from .solver import schmidt_decomposition from .helper import * -import numpy,h5py -import functools,sys, math +import numpy, h5py +import functools, sys, math from pyscf import ao2mo from functools import reduce @@ -16,11 +16,19 @@ class Frags: fragments for BE calculations. """ - def __init__(self, fsites, ifrag, edge=None, center=None, - edge_idx=None, center_idx=None, efac=None, - eri_file='eri_file.h5', - centerf_idx=None, - unrestricted=False): + def __init__( + self, + fsites, + ifrag, + edge=None, + center=None, + edge_idx=None, + center_idx=None, + efac=None, + eri_file="eri_file.h5", + centerf_idx=None, + unrestricted=False, + ): """Constructor function for `Frags` class. Parameters @@ -54,9 +62,13 @@ def __init__(self, fsites, ifrag, edge=None, center=None, self.h1 = None self.ifrag = ifrag if unrestricted: - self.dname = ['f'+str(ifrag)+'/aa', 'f'+str(ifrag)+'/bb', 'f'+str(ifrag)+'/ab'] + self.dname = [ + "f" + str(ifrag) + "/aa", + "f" + str(ifrag) + "/bb", + "f" + str(ifrag) + "/ab", + ] else: - self.dname = 'f'+str(ifrag) + self.dname = "f" + str(ifrag) self.nao = None self.mo_coeffs = None self._mo_coeffs = None @@ -81,8 +93,8 @@ def __init__(self, fsites, ifrag, edge=None, center=None, self.__rdm2 = None self.rdm1 = None self.genvs = None - self.ebe = 0. - self.ebe_hf = 0. + self.ebe = 0.0 + self.ebe_hf = 0.0 self.efac = efac self.fock = None self.veff = None @@ -90,7 +102,7 @@ def __init__(self, fsites, ifrag, edge=None, center=None, self.dm_init = None self.dm0 = None self.eri_file = eri_file - self.unitcell_nkpt = 1. + self.unitcell_nkpt = 1.0 def sd(self, lao, lmo, nocc, norb=None, return_orb_count=False): """ @@ -114,15 +126,17 @@ def sd(self, lao, lmo, nocc, norb=None, return_orb_count=False): """ if return_orb_count: - TA, n_f, n_b = schmidt_decomposition(lmo, nocc, self.fsites, norb=norb, return_orb_count=return_orb_count) + TA, n_f, n_b = schmidt_decomposition( + lmo, nocc, self.fsites, norb=norb, return_orb_count=return_orb_count + ) else: TA = schmidt_decomposition(lmo, nocc, self.fsites) self.C_lo_eo = TA - TA = numpy.dot(lao,TA) + TA = numpy.dot(lao, TA) self.nao = TA.shape[1] self.TA = TA if return_orb_count: - return [n_f, n_b] + return [n_f, n_b] def cons_h1(self, h1): """ @@ -134,8 +148,7 @@ def cons_h1(self, h1): One-electron Hamiltonian matrix. """ - h1_tmp = functools.reduce(numpy.dot, - (self.TA.T, h1, self.TA)) + h1_tmp = functools.reduce(numpy.dot, (self.TA.T, h1, self.TA)) self.h1 = h1_tmp def cons_fock(self, hf_veff, S, dm, eri_=None): @@ -155,14 +168,15 @@ def cons_fock(self, hf_veff, S, dm, eri_=None): """ if eri_ is None: - eri_ = get_eri(self.dname, self.TA.shape[1], ignore_symm=True, eri_file=self.eri_file) + eri_ = get_eri( + self.dname, self.TA.shape[1], ignore_symm=True, eri_file=self.eri_file + ) veff_ = get_veff(eri_, dm, S, self.TA, hf_veff) self.veff = veff_.real self.fock = self.h1 + veff_.real - - def get_nsocc(self, S, C, nocc,ncore=0): + def get_nsocc(self, S, C, nocc, ncore=0): """ Get the number of occupied orbitals for the fragment. @@ -184,7 +198,7 @@ def get_nsocc(self, S, C, nocc,ncore=0): """ import scipy.linalg - C_ = functools.reduce(numpy.dot,(self.TA.T, S, C[:,ncore:ncore+nocc])) + C_ = functools.reduce(numpy.dot, (self.TA.T, S, C[:, ncore : ncore + nocc])) P_ = numpy.dot(C_, C_.T) nsocc_ = numpy.trace(P_) nsocc = int(numpy.round(nsocc_)) @@ -192,15 +206,15 @@ def get_nsocc(self, S, C, nocc,ncore=0): try: mo_coeffs = scipy.linalg.svd(C_)[0] except: - - mo_coeffs = scipy.linalg.eigh(C_)[1][:,-nsocc:] + mo_coeffs = scipy.linalg.eigh(C_)[1][:, -nsocc:] self._mo_coeffs = mo_coeffs self.nsocc = nsocc return P_ - - def scf(self, heff=None, fs=False, eri=None, dm0 = None, unrestricted=False, spin_ind=None): + def scf( + self, heff=None, fs=False, eri=None, dm0=None, unrestricted=False, spin_ind=None + ): """ Perform self-consistent field (SCF) calculation for the fragment. @@ -221,9 +235,13 @@ def scf(self, heff=None, fs=False, eri=None, dm0 = None, unrestricted=False, spi """ import copy - if self._mf is not None: self._mf = None - if self._mc is not None: self._mc = None - if heff is None: heff = self.heff + + if self._mf is not None: + self._mf = None + if self._mc is not None: + self._mc = None + if heff is None: + heff = self.heff if eri is None: if unrestricted: @@ -233,18 +251,23 @@ def scf(self, heff=None, fs=False, eri=None, dm0 = None, unrestricted=False, spi eri = get_eri(dname, self.nao, eri_file=self.eri_file) if dm0 is None: - dm0 = numpy.dot( self._mo_coeffs[:,:self.nsocc], - self._mo_coeffs[:,:self.nsocc].conj().T) *2. - - mf_ = get_scfObj(self.fock + heff, eri, self.nsocc, dm0 = dm0) + dm0 = ( + numpy.dot( + self._mo_coeffs[:, : self.nsocc], + self._mo_coeffs[:, : self.nsocc].conj().T, + ) + * 2.0 + ) + + mf_ = get_scfObj(self.fock + heff, eri, self.nsocc, dm0=dm0) if not fs: self._mf = mf_ self.mo_coeffs = mf_.mo_coeff.copy() else: self._mo_coeffs = mf_.mo_coeff.copy() - mf_= None + mf_ = None - def update_heff(self,u, cout = None, return_heff=False, only_chem=False): + def update_heff(self, u, cout=None, return_heff=False, only_chem=False): """ Update the effective Hamiltonian for the fragment. """ @@ -256,9 +279,9 @@ def update_heff(self,u, cout = None, return_heff=False, only_chem=False): else: cout = cout - for i,fi in enumerate(self.fsites): + for i, fi in enumerate(self.fsites): if not any(i in sublist for sublist in self.edge_idx): - heff_[i,i] -= u[-1] + heff_[i, i] -= u[-1] if only_chem: self.heff = heff_ @@ -266,13 +289,13 @@ def update_heff(self,u, cout = None, return_heff=False, only_chem=False): if cout is None: return heff_ else: - return(cout, heff_) + return (cout, heff_) return cout for i in self.edge_idx: for j in range(len(i)): for k in range(len(i)): - if j>k :#or j==k: + if j > k: # or j==k: continue heff_[i[j], i[k]] = u[cout] @@ -280,33 +303,38 @@ def update_heff(self,u, cout = None, return_heff=False, only_chem=False): cout += 1 - self.heff = heff_ if return_heff: if cout is None: return heff_ else: - return(cout, heff_) + return (cout, heff_) return cout def set_udim(self, cout): for i in self.edge_idx: for j in range(len(i)): for k in range(len(i)): - if j>k : + if j > k: continue cout += 1 return cout - - def energy(self,rdm2s, eri=None, print_fragE=False): + def energy(self, rdm2s, eri=None, print_fragE=False): ## This function uses old energy expression and will be removed - rdm2s = numpy.einsum("ijkl,pi,qj,rk,sl->pqrs", 0.5*rdm2s, - *([self.mo_coeffs]*4),optimize=True) - - e1 = 2.*numpy.einsum("ij,ij->i",self.h1[:self.nfsites], self._rdm1[:self.nfsites]) - ec = numpy.einsum("ij,ij->i",self.veff[:self.nfsites],self._rdm1[:self.nfsites]) - + rdm2s = numpy.einsum( + "ijkl,pi,qj,rk,sl->pqrs", + 0.5 * rdm2s, + *([self.mo_coeffs] * 4), + optimize=True, + ) + + e1 = 2.0 * numpy.einsum( + "ij,ij->i", self.h1[: self.nfsites], self._rdm1[: self.nfsites] + ) + ec = numpy.einsum( + "ij,ij->i", self.veff[: self.nfsites], self._rdm1[: self.nfsites] + ) if self.TA.ndim == 3: jmax = self.TA[0].shape[1] @@ -314,97 +342,120 @@ def energy(self,rdm2s, eri=None, print_fragE=False): jmax = self.TA.shape[1] if eri is None: - r = h5py.File(self.eri_file,'r') + r = h5py.File(self.eri_file, "r") eri = r[self.dname][()] r.close() - - e2 = numpy.zeros_like(e1) for i in range(self.nfsites): for j in range(jmax): - ij = i*(i+1)//2+j if i > j else j*(j+1)//2+i - Gij = rdm2s[i,j,:jmax,:jmax].copy() + ij = i * (i + 1) // 2 + j if i > j else j * (j + 1) // 2 + i + Gij = rdm2s[i, j, :jmax, :jmax].copy() Gij[numpy.diag_indices(jmax)] *= 0.5 Gij += Gij.T e2[i] += Gij[numpy.tril_indices(jmax)] @ eri[ij] - e_ = e1+e2+ec - etmp = 0. - e1_ = 0. - ec_ = 0. - e2_ = 0. + e_ = e1 + e2 + ec + etmp = 0.0 + e1_ = 0.0 + ec_ = 0.0 + e2_ = 0.0 for i in self.efac[1]: - etmp += self.efac[0]*e_[i] + etmp += self.efac[0] * e_[i] e1_ += self.efac[0] * e1[i] ec_ += self.efac[0] * ec[i] e2_ += self.efac[0] * e2[i] - print('BE Energy Frag-{:>3} {:>12.7f} {:>12.7f} {:>12.7f}; Total : {:>12.7f}'. - format(self.dname, e1_, ec_, e2_, etmp)) + print( + "BE Energy Frag-{:>3} {:>12.7f} {:>12.7f} {:>12.7f}; Total : {:>12.7f}".format( + self.dname, e1_, ec_, e2_, etmp + ) + ) self.ebe = etmp - return (e1+e2+ec) - - def energy_hf(self, rdm_hf=None, mo_coeffs = None, eri=None, return_e1=False, unrestricted = False, spin_ind=None): + return e1 + e2 + ec + + def energy_hf( + self, + rdm_hf=None, + mo_coeffs=None, + eri=None, + return_e1=False, + unrestricted=False, + spin_ind=None, + ): if mo_coeffs is None: mo_coeffs = self._mo_coeffs if rdm_hf is None: - rdm_hf = numpy.dot(mo_coeffs[:,:self.nsocc], - mo_coeffs[:,:self.nsocc].conj().T) + rdm_hf = numpy.dot( + mo_coeffs[:, : self.nsocc], mo_coeffs[:, : self.nsocc].conj().T + ) - unrestricted_fac = 1. if unrestricted else 2. + unrestricted_fac = 1.0 if unrestricted else 2.0 - e1 = unrestricted_fac*numpy.einsum("ij,ij->i", self.h1[:self.nfsites], - rdm_hf[:self.nfsites]) + e1 = unrestricted_fac * numpy.einsum( + "ij,ij->i", self.h1[: self.nfsites], rdm_hf[: self.nfsites] + ) - ec = 0.5 * unrestricted_fac * numpy.einsum("ij,ij->i",self.veff[:self.nfsites], - rdm_hf[:self.nfsites]) + ec = ( + 0.5 + * unrestricted_fac + * numpy.einsum( + "ij,ij->i", self.veff[: self.nfsites], rdm_hf[: self.nfsites] + ) + ) if self.TA.ndim == 3: jmax = self.TA[0].shape[1] else: jmax = self.TA.shape[1] if eri is None: - r = h5py.File(self.eri_file,'r') + r = h5py.File(self.eri_file, "r") if isinstance(self.dname, list): - eri = [r[self.dname[0]][()],r[self.dname[1]][()]] + eri = [r[self.dname[0]][()], r[self.dname[1]][()]] else: eri = r[self.dname][()] r.close() - e2 = numpy.zeros_like(e1) for i in range(self.nfsites): for j in range(jmax): - ij = i*(i+1)//2+j if i > j else j*(j+1)//2+i - Gij = (2.*rdm_hf[i,j]*rdm_hf - - numpy.outer(rdm_hf[i], rdm_hf[j]))[:jmax,:jmax] + ij = i * (i + 1) // 2 + j if i > j else j * (j + 1) // 2 + i + Gij = (2.0 * rdm_hf[i, j] * rdm_hf - numpy.outer(rdm_hf[i], rdm_hf[j]))[ + :jmax, :jmax + ] Gij[numpy.diag_indices(jmax)] *= 0.5 Gij += Gij.T - if unrestricted: #unrestricted ERI file has 3 spin components: a, b, ab - e2[i] += 0.5 * unrestricted_fac * Gij[numpy.tril_indices(jmax)] @ eri[spin_ind][ij] + if ( + unrestricted + ): # unrestricted ERI file has 3 spin components: a, b, ab + e2[i] += ( + 0.5 + * unrestricted_fac + * Gij[numpy.tril_indices(jmax)] + @ eri[spin_ind][ij] + ) else: - e2[i] += 0.5 * unrestricted_fac * Gij[numpy.tril_indices(jmax)] @ eri[ij] + e2[i] += ( + 0.5 * unrestricted_fac * Gij[numpy.tril_indices(jmax)] @ eri[ij] + ) - e_ = e1+e2+ec - etmp = 0. + e_ = e1 + e2 + ec + etmp = 0.0 for i in self.efac[1]: - etmp += self.efac[0]*e_[i] + etmp += self.efac[0] * e_[i] self.ebe_hf = etmp if return_e1: - e_h1 = 0. - e_coul = 0. + e_h1 = 0.0 + e_coul = 0.0 for i in self.efac[1]: + e_h1 += self.efac[0] * e1[i] + e_coul += self.efac[0] * (e2[i] + ec[i]) + return (e_h1, e_coul, e1 + e2 + ec) - e_h1 += self.efac[0]*e1[i] - e_coul += self.efac[0]*(e2[i]+ec[i]) - return(e_h1,e_coul, e1+e2+ec) - - return e1+e2+ec - + return e1 + e2 + ec diff --git a/molbe/rdm.py b/molbe/rdm.py index 72902668..1d5de9dc 100644 --- a/molbe/rdm.py +++ b/molbe/rdm.py @@ -2,9 +2,16 @@ import numpy, sys, scipy -def rdm1_fullbasis(self, return_ao=True, only_rdm1=False, - only_rdm2=False, return_lo=False, - return_RDM2=True, print_energy=False): + +def rdm1_fullbasis( + self, + return_ao=True, + only_rdm1=False, + only_rdm2=False, + return_lo=False, + return_RDM2=True, + print_energy=False, +): """ Compute the one-particle and two-particle reduced density matrices (RDM1 and RDM2). @@ -52,16 +59,26 @@ def rdm1_fullbasis(self, return_ao=True, only_rdm1=False, if return_RDM2: # Adjust the one-particle reduced density matrix (RDM1) drdm1 = fobjs.__rdm1.copy() - drdm1[numpy.diag_indices(fobjs.nsocc)] -= 2. + drdm1[numpy.diag_indices(fobjs.nsocc)] -= 2.0 # Compute the two-particle reduced density matrix (RDM2) and subtract non-connected component - dm_nc = numpy.einsum('ij,kl->ijkl', drdm1, drdm1, dtype=numpy.float64, optimize=True) - \ - 0.5*numpy.einsum('ij,kl->iklj', drdm1, drdm1, dtype=numpy.float64,optimize=True) + dm_nc = numpy.einsum( + "ij,kl->ijkl", drdm1, drdm1, dtype=numpy.float64, optimize=True + ) - 0.5 * numpy.einsum( + "ij,kl->iklj", drdm1, drdm1, dtype=numpy.float64, optimize=True + ) fobjs.__rdm2 -= dm_nc # Generate the projection matrix - cind = [ fobjs.fsites[i] for i in fobjs.efac[1]] - Pc_ = fobjs.TA.T @ self.S @ self.W[:, cind] @ self.W[:, cind].T @ self.S @ fobjs.TA + cind = [fobjs.fsites[i] for i in fobjs.efac[1]] + Pc_ = ( + fobjs.TA.T + @ self.S + @ self.W[:, cind] + @ self.W[:, cind].T + @ self.S + @ fobjs.TA + ) if not only_rdm2: # Compute RDM1 in the localized orbital (LO) basis and transform to AO basis @@ -72,38 +89,68 @@ def rdm1_fullbasis(self, return_ao=True, only_rdm1=False, if not only_rdm1: # Transform RDM2 to AO basis - rdm2s = numpy.einsum("ijkl,pi,qj,rk,sl->pqrs", fobjs.__rdm2, - *([fobjs.mo_coeffs]*4),optimize=True) - rdm2_ao = numpy.einsum('xi,ijkl,px,qj,rk,sl->pqrs', - Pc_, rdm2s, fobjs.TA, fobjs.TA, - fobjs.TA, fobjs.TA, optimize=True) + rdm2s = numpy.einsum( + "ijkl,pi,qj,rk,sl->pqrs", + fobjs.__rdm2, + *([fobjs.mo_coeffs] * 4), + optimize=True, + ) + rdm2_ao = numpy.einsum( + "xi,ijkl,px,qj,rk,sl->pqrs", + Pc_, + rdm2s, + fobjs.TA, + fobjs.TA, + fobjs.TA, + fobjs.TA, + optimize=True, + ) rdm2AO += rdm2_ao if not only_rdm1: # Symmetrize RDM2 and add the non-cumulant part if requested - rdm2AO = (rdm2AO + rdm2AO.T)/2. + rdm2AO = (rdm2AO + rdm2AO.T) / 2.0 if return_RDM2: - nc_AO = numpy.einsum('ij,kl->ijkl', rdm1AO, rdm1AO, dtype=numpy.float64, optimize=True) - \ - numpy.einsum('ij,kl->iklj', rdm1AO, rdm1AO, dtype=numpy.float64, optimize=True)*0.5 + nc_AO = ( + numpy.einsum( + "ij,kl->ijkl", rdm1AO, rdm1AO, dtype=numpy.float64, optimize=True + ) + - numpy.einsum( + "ij,kl->iklj", rdm1AO, rdm1AO, dtype=numpy.float64, optimize=True + ) + * 0.5 + ) rdm2AO = nc_AO + rdm2AO # Transform RDM2 to the molecular orbital (MO) basis if needed if not return_ao: CmoT_S = self.C.T @ self.S - rdm2MO = numpy.einsum("ijkl,pi,qj,rk,sl->pqrs", - rdm2AO, CmoT_S, CmoT_S, - CmoT_S, CmoT_S, optimize=True) + rdm2MO = numpy.einsum( + "ijkl,pi,qj,rk,sl->pqrs", + rdm2AO, + CmoT_S, + CmoT_S, + CmoT_S, + CmoT_S, + optimize=True, + ) # Transform RDM2 to the localized orbital (LO) basis if needed if return_lo: CloT_S = self.W.T @ self.S - rdm2LO = numpy.einsum("ijkl,pi,qj,rk,sl->pqrs", - rdm2AO, CloT_S, CloT_S, - CloT_S, CloT_S, optimize=True) + rdm2LO = numpy.einsum( + "ijkl,pi,qj,rk,sl->pqrs", + rdm2AO, + CloT_S, + CloT_S, + CloT_S, + CloT_S, + optimize=True, + ) if not only_rdm2: # Symmetrize RDM1 - rdm1AO = (rdm1AO + rdm1AO.T)/2. + rdm1AO = (rdm1AO + rdm1AO.T) / 2.0 # Transform RDM1 to the MO basis if needed if not return_ao: @@ -115,24 +162,23 @@ def rdm1_fullbasis(self, return_ao=True, only_rdm1=False, if return_RDM2 and print_energy: # Compute and print energy contributions - Eh1 = numpy.einsum('ij,ij', self.hcore, rdm1AO, optimize=True) - eri = ao2mo.restore(1,self.mf._eri, self.mf.mo_coeff.shape[1]) - E2 = 0.5*numpy.einsum('pqrs,pqrs', eri,rdm2AO, optimize=True) + Eh1 = numpy.einsum("ij,ij", self.hcore, rdm1AO, optimize=True) + eri = ao2mo.restore(1, self.mf._eri, self.mf.mo_coeff.shape[1]) + E2 = 0.5 * numpy.einsum("pqrs,pqrs", eri, rdm2AO, optimize=True) print(flush=True) - print('-----------------------------------------------------', - flush=True) - print(' BE ENERGIES with cumulant-based expression', flush=True) - - print('-----------------------------------------------------', - flush=True) - - print(' 1-elec E : {:>15.8f} Ha'.format(Eh1), flush=True) - print(' 2-elec E : {:>15.8f} Ha'.format(E2), flush=True) - E_tot = Eh1+E2+self.E_core + self.enuc - print(' E_BE : {:>15.8f} Ha'.format(E_tot), flush=True) - print(' Ecorr BE : {:>15.8f} Ha'.format((E_tot)-self.ebe_hf), flush=True) - print('-----------------------------------------------------', - flush=True) + print("-----------------------------------------------------", flush=True) + print(" BE ENERGIES with cumulant-based expression", flush=True) + + print("-----------------------------------------------------", flush=True) + + print(" 1-elec E : {:>15.8f} Ha".format(Eh1), flush=True) + print(" 2-elec E : {:>15.8f} Ha".format(E2), flush=True) + E_tot = Eh1 + E2 + self.E_core + self.enuc + print(" E_BE : {:>15.8f} Ha".format(E_tot), flush=True) + print( + " Ecorr BE : {:>15.8f} Ha".format((E_tot) - self.ebe_hf), flush=True + ) + print("-----------------------------------------------------", flush=True) print(flush=True) if only_rdm1: @@ -151,10 +197,15 @@ def rdm1_fullbasis(self, return_ao=True, only_rdm1=False, if return_lo and not return_ao: return (rdm1MO, rdm2MO, rdm1LO, rdm2LO) - if return_ao: return rdm1AO, rdm2AO - if not return_ao: return rdm1MO, rdm2MO + if return_ao: + return rdm1AO, rdm2AO + if not return_ao: + return rdm1MO, rdm2MO + -def compute_energy_full(self, approx_cumulant=False, use_full_rdm=False, return_rdm=True): +def compute_energy_full( + self, approx_cumulant=False, use_full_rdm=False, return_rdm=True +): """ Compute the total energy using rdms in the full basis. @@ -183,7 +234,9 @@ def compute_energy_full(self, approx_cumulant=False, use_full_rdm=False, return_ from pyscf import scf, ao2mo # Compute the one-particle reduced density matrix (RDM1) and the cumulant (Kumul) in the full basis - rdm1f, Kumul, rdm1_lo, rdm2_lo = self.rdm1_fullbasis(return_lo=True, return_RDM2=False) + rdm1f, Kumul, rdm1_lo, rdm2_lo = self.rdm1_fullbasis( + return_lo=True, return_RDM2=False + ) if not approx_cumulant: # Compute the true two-particle reduced density matrix (RDM2) if not using approximate cumulant @@ -191,8 +244,15 @@ def compute_energy_full(self, approx_cumulant=False, use_full_rdm=False, return_ if return_rdm: # Construct the full RDM2 from RDM1 - RDM2_full = numpy.einsum('ij,kl->ijkl', rdm1f, rdm1f, dtype=numpy.float64, optimize=True) - \ - numpy.einsum('ij,kl->iklj', rdm1f, rdm1f, dtype=numpy.float64, optimize=True)*0.5 + RDM2_full = ( + numpy.einsum( + "ij,kl->ijkl", rdm1f, rdm1f, dtype=numpy.float64, optimize=True + ) + - numpy.einsum( + "ij,kl->iklj", rdm1f, rdm1f, dtype=numpy.float64, optimize=True + ) + * 0.5 + ) # Add the cumulant part to RDM2 if not approx_cumulant: @@ -207,72 +267,75 @@ def compute_energy_full(self, approx_cumulant=False, use_full_rdm=False, return_ veff = scf.hf.get_veff(self.mol, rdm1f, hermi=0) # Compute the one-electron energy - Eh1 = numpy.einsum('ij,ij', self.hcore, rdm1f, optimize=True) + Eh1 = numpy.einsum("ij,ij", self.hcore, rdm1f, optimize=True) # Compute the energy due to the effective potential - EVeff = numpy.einsum('ij,ij',veff, rdm1f, optimize=True) + EVeff = numpy.einsum("ij,ij", veff, rdm1f, optimize=True) # Compute the change in the one-electron energy - Eh1_dg = numpy.einsum('ij,ij',self.hcore, del_gamma, optimize=True) + Eh1_dg = numpy.einsum("ij,ij", self.hcore, del_gamma, optimize=True) # Compute the change in the effective potential energy - Eveff_dg = numpy.einsum('ij,ij',self.hf_veff, del_gamma, optimize=True) + Eveff_dg = numpy.einsum("ij,ij", self.hf_veff, del_gamma, optimize=True) # Restore the electron repulsion integrals (ERI) - eri = ao2mo.restore(1,self.mf._eri, self.mf.mo_coeff.shape[1]) + eri = ao2mo.restore(1, self.mf._eri, self.mf.mo_coeff.shape[1]) # Compute the cumulant part of the two-electron energy - EKumul = numpy.einsum('pqrs,pqrs', eri,Kumul, optimize=True) + EKumul = numpy.einsum("pqrs,pqrs", eri, Kumul, optimize=True) if not approx_cumulant: # Compute the true two-electron energy if not using approximate cumulant - EKumul_T = numpy.einsum('pqrs,pqrs', eri,Kumul_T, optimize=True) + EKumul_T = numpy.einsum("pqrs,pqrs", eri, Kumul_T, optimize=True) if use_full_rdm and return_rdm: # Compute the full two-electron energy using the full RDM2 - E2 = numpy.einsum('pqrs,pqrs', eri,RDM2_full, optimize=True) + E2 = numpy.einsum("pqrs,pqrs", eri, RDM2_full, optimize=True) # Compute the approximate BE total energy - EKapprox = self.ebe_hf + Eh1_dg + Eveff_dg + EKumul/2. + EKapprox = self.ebe_hf + Eh1_dg + Eveff_dg + EKumul / 2.0 self.ebe_tot = EKapprox if not approx_cumulant: # Compute the true BE total energy if not using approximate cumulant - EKtrue = Eh1 + EVeff/2. + EKumul_T/2. + self.enuc + self.E_core + EKtrue = Eh1 + EVeff / 2.0 + EKumul_T / 2.0 + self.enuc + self.E_core self.ebe_tot = EKtrue # Print energy results - print('-----------------------------------------------------', - flush=True) - print(' BE ENERGIES with cumulant-based expression', flush=True) - - print('-----------------------------------------------------', - flush=True) - print(' E_BE = E_HF + Tr(F del g) + Tr(V K_approx)', flush=True) - print(' E_HF : {:>14.8f} Ha'.format(self.ebe_hf), flush=True) - print(' Tr(F del g) : {:>14.8f} Ha'.format(Eh1_dg+Eveff_dg), flush=True) - print(' Tr(V K_aprrox) : {:>14.8f} Ha'.format(EKumul/2.), flush=True) - print(' E_BE : {:>14.8f} Ha'.format(EKapprox), flush=True) - print(' Ecorr BE : {:>14.8f} Ha'.format(EKapprox-self.ebe_hf), flush=True) + print("-----------------------------------------------------", flush=True) + print(" BE ENERGIES with cumulant-based expression", flush=True) + + print("-----------------------------------------------------", flush=True) + print(" E_BE = E_HF + Tr(F del g) + Tr(V K_approx)", flush=True) + print(" E_HF : {:>14.8f} Ha".format(self.ebe_hf), flush=True) + print(" Tr(F del g) : {:>14.8f} Ha".format(Eh1_dg + Eveff_dg), flush=True) + print(" Tr(V K_aprrox) : {:>14.8f} Ha".format(EKumul / 2.0), flush=True) + print(" E_BE : {:>14.8f} Ha".format(EKapprox), flush=True) + print(" Ecorr BE : {:>14.8f} Ha".format(EKapprox - self.ebe_hf), flush=True) if not approx_cumulant: print(flush=True) - print(' E_BE = Tr(F[g] g) + Tr(V K_true)', flush=True) - print(' Tr(h1 g) : {:>14.8f} Ha'.format(Eh1), flush=True) - print(' Tr(Veff[g] g) : {:>14.8f} Ha'.format(EVeff/2.), flush=True) - print(' Tr(V K_true) : {:>14.8f} Ha'.format(EKumul_T/2.), flush=True) - print(' E_BE : {:>14.8f} Ha'.format(EKtrue), flush=True) + print(" E_BE = Tr(F[g] g) + Tr(V K_true)", flush=True) + print(" Tr(h1 g) : {:>14.8f} Ha".format(Eh1), flush=True) + print(" Tr(Veff[g] g) : {:>14.8f} Ha".format(EVeff / 2.0), flush=True) + print(" Tr(V K_true) : {:>14.8f} Ha".format(EKumul_T / 2.0), flush=True) + print(" E_BE : {:>14.8f} Ha".format(EKtrue), flush=True) if use_full_rdm and return_rdm: - print(' E(g+G) : {:>14.8f} Ha'.format(Eh1 + 0.5*E2 + self.E_core + self.enuc), - flush=True) - print(' Ecorr BE : {:>14.8f} Ha'.format(EKtrue-self.ebe_hf), flush=True) + print( + " E(g+G) : {:>14.8f} Ha".format( + Eh1 + 0.5 * E2 + self.E_core + self.enuc + ), + flush=True, + ) + print( + " Ecorr BE : {:>14.8f} Ha".format(EKtrue - self.ebe_hf), flush=True + ) print(flush=True) - print(' True - approx : {:>14.4e} Ha'.format(EKtrue-EKapprox)) - print('-----------------------------------------------------', - flush=True) + print(" True - approx : {:>14.4e} Ha".format(EKtrue - EKapprox)) + print("-----------------------------------------------------", flush=True) print(flush=True) # Return the RDMs if requested - if return_rdm: return(rdm1f, RDM2_full) - + if return_rdm: + return (rdm1f, RDM2_full) diff --git a/molbe/solver.py b/molbe/solver.py index 15714554..80c64582 100644 --- a/molbe/solver.py +++ b/molbe/solver.py @@ -6,14 +6,39 @@ import time import os from molbe import be_var -from molbe.external.ccsd_rdm import make_rdm1_ccsd_t1, make_rdm2_urlx, make_rdm1_uccsd, make_rdm2_uccsd - -def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, - only_chem = False, nproc=4, hci_pt=False, - hci_cutoff=0.001, ci_coeff_cutoff = None, select_cutoff=None, - eeval=False, ereturn=False, frag_energy=False, relax_density=False, - return_vec=False, ecore=0., ebe_hf=0., be_iter=None, use_cumulant=True, - scratch_dir=None, **solver_kwargs): +from molbe.external.ccsd_rdm import ( + make_rdm1_ccsd_t1, + make_rdm2_urlx, + make_rdm1_uccsd, + make_rdm2_uccsd, +) + + +def be_func( + pot, + Fobjs, + Nocc, + solver, + enuc, + hf_veff=None, + only_chem=False, + nproc=4, + hci_pt=False, + hci_cutoff=0.001, + ci_coeff_cutoff=None, + select_cutoff=None, + eeval=False, + ereturn=False, + frag_energy=False, + relax_density=False, + return_vec=False, + ecore=0.0, + ebe_hf=0.0, + be_iter=None, + use_cumulant=True, + scratch_dir=None, + **solver_kwargs, +): """ Perform bootstrap embedding calculations for each fragment. @@ -62,23 +87,22 @@ def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, Depending on the options, it returns the norm of the error vector, the energy, or a combination of these values. """ from pyscf import fci - import h5py,os + import h5py, os from pyscf import ao2mo from .helper import get_frag_energy rdm_return = False if relax_density: rdm_return = True - E = 0. + E = 0.0 if frag_energy or eeval: - total_e = [0.,0.,0.] + total_e = [0.0, 0.0, 0.0] # Loop over each fragment and solve using the specified solver for fobj in Fobjs: # Update the effective Hamiltonian if not pot is None: - heff_ = fobj.update_heff(pot, return_heff=True, - only_chem=only_chem) + heff_ = fobj.update_heff(pot, return_heff=True, only_chem=only_chem) else: heff_ = None @@ -88,23 +112,25 @@ def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, fobj.scf() # Solve using the specified solver - if solver=='MP2': + if solver == "MP2": fobj._mc = solve_mp2(fobj._mf, mo_energy=fobj._mf.mo_energy) - elif solver=='CCSD': + elif solver == "CCSD": if rdm_return: - fobj.t1, fobj.t2, rdm1_tmp, rdm2s = solve_ccsd(fobj._mf, - mo_energy=fobj._mf.mo_energy, - relax=True, use_cumulant=use_cumulant, - rdm2_return=True, - rdm_return=True) + fobj.t1, fobj.t2, rdm1_tmp, rdm2s = solve_ccsd( + fobj._mf, + mo_energy=fobj._mf.mo_energy, + relax=True, + use_cumulant=use_cumulant, + rdm2_return=True, + rdm_return=True, + ) else: - fobj.t1, fobj.t2 = solve_ccsd(fobj._mf, - mo_energy=fobj._mf.mo_energy, - rdm_return=False) + fobj.t1, fobj.t2 = solve_ccsd( + fobj._mf, mo_energy=fobj._mf.mo_energy, rdm_return=False + ) rdm1_tmp = make_rdm1_ccsd_t1(fobj.t1) - - elif solver=='FCI': + elif solver == "FCI": from .helper import get_eri import scipy @@ -112,15 +138,16 @@ def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, efci, civec = mc.kernel() rdm1_tmp = mc.make_rdm1(civec, mc.norb, mc.nelec) - elif solver=='HCI': + elif solver == "HCI": from pyscf import hci from .helper import get_eri # pilot pyscf.hci only in old versions nao, nmo = fobj._mf.mo_coeff.shape - eri = ao2mo.kernel(fobj._mf._eri, fobj._mf.mo_coeff, aosym='s4', - compact=False).reshape(4*((nmo),)) + eri = ao2mo.kernel( + fobj._mf._eri, fobj._mf.mo_coeff, aosym="s4", compact=False + ).reshape(4 * ((nmo),)) ci_ = hci.SCI(fobj._mf.mol) if select_cutoff is None and ci_coeff_cutoff is None: @@ -133,17 +160,20 @@ def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, ci_.ci_coeff_cutoff = ci_coeff_cutoff nelec = (fobj.nsocc, fobj.nsocc) - h1_ = fobj.fock+fobj.heff - h1_ = functools.reduce(numpy.dot, (fobj._mf.mo_coeff.T, h1_, fobj._mf.mo_coeff)) - eci, civec = ci_.kernel(h1_, eri, nmo, nelec) + h1_ = fobj.fock + fobj.heff + h1_ = functools.reduce( + numpy.dot, (fobj._mf.mo_coeff.T, h1_, fobj._mf.mo_coeff) + ) + eci, civec = ci_.kernel(h1_, eri, nmo, nelec) civec = numpy.asarray(civec) - - (rdm1a_, rdm1b_), (rdm2aa, rdm2ab, rdm2bb) = ci_.make_rdm12s(civec, nmo, nelec) + (rdm1a_, rdm1b_), (rdm2aa, rdm2ab, rdm2bb) = ci_.make_rdm12s( + civec, nmo, nelec + ) rdm1_tmp = rdm1a_ + rdm1b_ - rdm2s = rdm2aa + rdm2ab + rdm2ab.transpose(2,3,0,1) + rdm2bb + rdm2s = rdm2aa + rdm2ab + rdm2ab.transpose(2, 3, 0, 1) + rdm2bb - elif solver=='SHCI': + elif solver == "SHCI": from pyscf.shciscf import shci if scratch_dir is None and be_var.CREATE_SCRATCH_DIR: @@ -153,111 +183,140 @@ def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, else: tmp = os.path.join(scratch_dir, str(os.getpid()), str(fobj.dname)) if not os.path.isdir(tmp): - os.system('mkdir -p '+tmp) + os.system("mkdir -p " + tmp) nao, nmo = fobj._mf.mo_coeff.shape nelec = (fobj.nsocc, fobj.nsocc) mch = shci.SHCISCF(fobj._mf, nmo, nelec, orbpath=fobj.dname) - mch.fcisolver.mpiprefix = 'mpirun -np '+str(nproc) + mch.fcisolver.mpiprefix = "mpirun -np " + str(nproc) if hci_pt: mch.fcisolver.stochastic = False mch.fcisolver.epsilon2 = hci_cutoff else: - mch.fcisolver.stochastic = True # this is for PT and doesnt add PT to rdm + mch.fcisolver.stochastic = ( + True # this is for PT and doesnt add PT to rdm + ) mch.fcisolver.nPTiter = 0 mch.fcisolver.sweep_iter = [0] mch.fcisolver.DoRDM = True - mch.fcisolver.sweep_epsilon = [ hci_cutoff ] + mch.fcisolver.sweep_epsilon = [hci_cutoff] mch.fcisolver.scratchDirectory = scratch_dir mch.mc1step() rdm1_tmp, rdm2s = mch.fcisolver.make_rdm12(0, nmo, nelec) - elif solver == 'SCI': + elif solver == "SCI": from pyscf import cornell_shci from pyscf import ao2mo, mcscf nao, nmo = fobj._mf.mo_coeff.shape nelec = (fobj.nsocc, fobj.nsocc) - cas = mcscf.CASCI (fobj._mf, nmo, nelec) + cas = mcscf.CASCI(fobj._mf, nmo, nelec) h1, ecore = cas.get_h1eff(mo_coeff=fobj._mf.mo_coeff) - eri = ao2mo.kernel(fobj._mf._eri, fobj._mf.mo_coeff, aosym='s4', compact=False).reshape(4*((nmo),)) + eri = ao2mo.kernel( + fobj._mf._eri, fobj._mf.mo_coeff, aosym="s4", compact=False + ).reshape(4 * ((nmo),)) ci = cornell_shci.SHCI() - ci.runtimedir=fobj.dname - ci.restart=True - ci.config['var_only'] = True - ci.config['eps_vars'] = [hci_cutoff] - ci.config['get_1rdm_csv'] = True - ci.config['get_2rdm_csv'] = True + ci.runtimedir = fobj.dname + ci.restart = True + ci.config["var_only"] = True + ci.config["eps_vars"] = [hci_cutoff] + ci.config["get_1rdm_csv"] = True + ci.config["get_2rdm_csv"] = True ci.kernel(h1, eri, nmo, nelec) - rdm1_tmp, rdm2s = ci.make_rdm12(0,nmo,nelec) - - elif solver in ['block2', 'DMRG','DMRGCI','DMRGSCF']: + rdm1_tmp, rdm2s = ci.make_rdm12(0, nmo, nelec) + elif solver in ["block2", "DMRG", "DMRGCI", "DMRGSCF"]: solver_kwargs_ = solver_kwargs.copy() if scratch_dir is None and be_var.CREATE_SCRATCH_DIR: tmp = os.path.join(be_var.SCRATCH, str(os.getpid()), str(fobj.dname)) else: tmp = os.path.join(scratch_dir, str(os.getpid()), str(fobj.dname)) if not os.path.isdir(tmp): - os.system('mkdir -p '+tmp) + os.system("mkdir -p " + tmp) try: - rdm1_tmp, rdm2s = solve_block2(fobj._mf, - fobj.nsocc, - frag_scratch = tmp, - **solver_kwargs_) + rdm1_tmp, rdm2s = solve_block2( + fobj._mf, fobj.nsocc, frag_scratch=tmp, **solver_kwargs_ + ) except Exception as inst: - print(f"Fragment DMRG solver failed with Exception: {type(inst)}\n", inst, flush=True) + print( + f"Fragment DMRG solver failed with Exception: {type(inst)}\n", + inst, + flush=True, + ) finally: - if solver_kwargs_.pop('force_cleanup', False): - os.system('rm -r '+ os.path.join(tmp,'F.*')) - os.system('rm -r '+ os.path.join(tmp,'FCIDUMP*')) - os.system('rm -r '+ os.path.join(tmp,'node*')) + if solver_kwargs_.pop("force_cleanup", False): + os.system("rm -r " + os.path.join(tmp, "F.*")) + os.system("rm -r " + os.path.join(tmp, "FCIDUMP*")) + os.system("rm -r " + os.path.join(tmp, "node*")) else: - print('Solver not implemented',flush=True) - print('exiting',flush=True) + print("Solver not implemented", flush=True) + print("exiting", flush=True) sys.exit() - if solver=='MP2': + if solver == "MP2": rdm1_tmp = fobj._mc.make_rdm1() fobj.__rdm1 = rdm1_tmp.copy() - fobj._rdm1 = functools.reduce(numpy.dot, - (fobj.mo_coeffs, - #fobj._mc.make_rdm1(), - rdm1_tmp, - fobj.mo_coeffs.T))*0.5 + fobj._rdm1 = ( + functools.reduce( + numpy.dot, + ( + fobj.mo_coeffs, + # fobj._mc.make_rdm1(), + rdm1_tmp, + fobj.mo_coeffs.T, + ), + ) + * 0.5 + ) if eeval or ereturn: - if solver =='CCSD' and not rdm_return: + if solver == "CCSD" and not rdm_return: with_dm1 = True - if use_cumulant: with_dm1=False + if use_cumulant: + with_dm1 = False rdm2s = make_rdm2_urlx(fobj.t1, fobj.t2, with_dm1=with_dm1) - elif solver == 'MP2': + elif solver == "MP2": rdm2s = fobj._mc.make_rdm2() - elif solver =='FCI': + elif solver == "FCI": rdm2s = mc.make_rdm2(civec, mc.norb, mc.nelec) if use_cumulant: hf_dm = numpy.zeros_like(rdm1_tmp) - hf_dm[numpy.diag_indices(fobj.nsocc)] += 2. + hf_dm[numpy.diag_indices(fobj.nsocc)] += 2.0 del_rdm1 = rdm1_tmp.copy() - del_rdm1[numpy.diag_indices(fobj.nsocc)] -= 2. - nc = numpy.einsum('ij,kl->ijkl',hf_dm, hf_dm) + \ - numpy.einsum('ij,kl->ijkl',hf_dm, del_rdm1) + \ - numpy.einsum('ij,kl->ijkl',del_rdm1, hf_dm) - nc -= (numpy.einsum('ij,kl->iklj',hf_dm, hf_dm) + \ - numpy.einsum('ij,kl->iklj',hf_dm, del_rdm1) + \ - numpy.einsum('ij,kl->iklj',del_rdm1, hf_dm))*0.5 + del_rdm1[numpy.diag_indices(fobj.nsocc)] -= 2.0 + nc = ( + numpy.einsum("ij,kl->ijkl", hf_dm, hf_dm) + + numpy.einsum("ij,kl->ijkl", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->ijkl", del_rdm1, hf_dm) + ) + nc -= ( + numpy.einsum("ij,kl->iklj", hf_dm, hf_dm) + + numpy.einsum("ij,kl->iklj", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->iklj", del_rdm1, hf_dm) + ) * 0.5 rdm2s -= nc fobj.__rdm2 = rdm2s.copy() if frag_energy or eeval: # Find the energy of a given fragment, with the cumulant definition. # Return [e1, e2, ec] as e_f and add to the running total_e. - e_f = get_frag_energy(fobj.mo_coeffs, fobj.nsocc, fobj.nfsites, - fobj.efac, fobj.TA, fobj.h1, hf_veff, rdm1_tmp, - rdm2s, fobj.dname, eri_file=fobj.eri_file, veff0=fobj.veff0) + e_f = get_frag_energy( + fobj.mo_coeffs, + fobj.nsocc, + fobj.nfsites, + fobj.efac, + fobj.TA, + fobj.h1, + hf_veff, + rdm1_tmp, + rdm2s, + fobj.dname, + eri_file=fobj.eri_file, + veff0=fobj.veff0, + ) total_e = [sum(x) for x in zip(total_e, e_f)] fobj.energy_hf() @@ -267,21 +326,33 @@ def be_func(pot, Fobjs, Nocc, solver, enuc, hf_veff=None, if frag_energy: return (Ecorr, total_e) - ernorm, ervec = solve_error(Fobjs,Nocc, only_chem=only_chem) + ernorm, ervec = solve_error(Fobjs, Nocc, only_chem=only_chem) if return_vec: return (ernorm, ervec, [Ecorr, total_e]) if eeval: - print('Error in density matching : {:>2.4e}'.format(ernorm), flush=True) + print("Error in density matching : {:>2.4e}".format(ernorm), flush=True) return ernorm -def be_func_u(pot, Fobjs, solver, enuc, hf_veff=None, - eeval=False, ereturn=False, frag_energy=True, - relax_density=False, ecore=0., ebe_hf=0., - scratch_dir=None, use_cumulant=True, frozen=False): +def be_func_u( + pot, + Fobjs, + solver, + enuc, + hf_veff=None, + eeval=False, + ereturn=False, + frag_energy=True, + relax_density=False, + ecore=0.0, + ebe_hf=0.0, + scratch_dir=None, + use_cumulant=True, + frozen=False, +): """ Perform bootstrap embedding calculations for each fragment with UCCSD. @@ -323,7 +394,7 @@ def be_func_u(pot, Fobjs, solver, enuc, hf_veff=None, Depending on the options, it returns the norm of the error vector, the energy, or a combination of these values. """ from pyscf import scf - import h5py,os + import h5py, os from pyscf import ao2mo from .helper import get_frag_energy_u from molbe.external.unrestricted_utils import make_uhf_obj @@ -331,73 +402,91 @@ def be_func_u(pot, Fobjs, solver, enuc, hf_veff=None, rdm_return = False if relax_density: rdm_return = True - E = 0. + E = 0.0 if frag_energy or eeval: - total_e = [0.,0.,0.] + total_e = [0.0, 0.0, 0.0] # Loop over each fragment and solve using the specified solver - for (fobj_a, fobj_b) in Fobjs: - heff_ = None # No outside chemical potential implemented for unrestricted yet + for fobj_a, fobj_b in Fobjs: + heff_ = None # No outside chemical potential implemented for unrestricted yet fobj_a.scf(unrestricted=True, spin_ind=0) fobj_b.scf(unrestricted=True, spin_ind=1) full_uhf, eris = make_uhf_obj(fobj_a, fobj_b, frozen=frozen) - if solver == 'UCCSD': + if solver == "UCCSD": if rdm_return: - ucc, rdm1_tmp, rdm2s = solve_uccsd(full_uhf, eris, relax=relax_density, - rdm_return=True, rdm2_return=True, - frozen=frozen) + ucc, rdm1_tmp, rdm2s = solve_uccsd( + full_uhf, + eris, + relax=relax_density, + rdm_return=True, + rdm2_return=True, + frozen=frozen, + ) else: - ucc = solve_uccsd(full_uhf, eris, relax = relax_density, rdm_return=False, - frozen=frozen) + ucc = solve_uccsd( + full_uhf, eris, relax=relax_density, rdm_return=False, frozen=frozen + ) rdm1_tmp = make_rdm1_uccsd(ucc, relax=relax_density) else: - print('Solver not implemented',flush=True) - print('exiting',flush=True) + print("Solver not implemented", flush=True) + print("exiting", flush=True) sys.exit() fobj_a.__rdm1 = rdm1_tmp[0].copy() - fobj_b._rdm1 = functools.reduce(numpy.dot, - (fobj_a._mf.mo_coeff, - rdm1_tmp[0], - fobj_a._mf.mo_coeff.T))*0.5 + fobj_b._rdm1 = ( + functools.reduce( + numpy.dot, (fobj_a._mf.mo_coeff, rdm1_tmp[0], fobj_a._mf.mo_coeff.T) + ) + * 0.5 + ) fobj_b.__rdm1 = rdm1_tmp[1].copy() - fobj_b._rdm1 = functools.reduce(numpy.dot, - (fobj_b._mf.mo_coeff, - rdm1_tmp[1], - fobj_b._mf.mo_coeff.T))*0.5 + fobj_b._rdm1 = ( + functools.reduce( + numpy.dot, (fobj_b._mf.mo_coeff, rdm1_tmp[1], fobj_b._mf.mo_coeff.T) + ) + * 0.5 + ) if eeval or ereturn: - if solver =='UCCSD' and not rdm_return: + if solver == "UCCSD" and not rdm_return: with_dm1 = True - if use_cumulant: with_dm1=False + if use_cumulant: + with_dm1 = False rdm2s = make_rdm2_uccsd(ucc, with_dm1=with_dm1) fobj_a.__rdm2 = rdm2s[0].copy() fobj_b.__rdm2 = rdm2s[1].copy() if frag_energy: if frozen: - h1_ab = [full_uhf.h1[0]+full_uhf.full_gcore[0]+full_uhf.core_veffs[0], - full_uhf.h1[1]+full_uhf.full_gcore[1]+full_uhf.core_veffs[1]] + h1_ab = [ + full_uhf.h1[0] + + full_uhf.full_gcore[0] + + full_uhf.core_veffs[0], + full_uhf.h1[1] + + full_uhf.full_gcore[1] + + full_uhf.core_veffs[1], + ] else: h1_ab = [fobj_a.h1, fobj_b.h1] - e_f = get_frag_energy_u((fobj_a._mo_coeffs,fobj_b._mo_coeffs), - (fobj_a.nsocc,fobj_b.nsocc), - (fobj_a.nfsites, fobj_b.nfsites), - (fobj_a.efac, fobj_b.efac), - (fobj_a.TA, fobj_b.TA), - h1_ab, - hf_veff, - rdm1_tmp, - rdm2s, - fobj_a.dname, - eri_file=fobj_a.eri_file, - gcores=full_uhf.full_gcore, - frozen=frozen - ) + e_f = get_frag_energy_u( + (fobj_a._mo_coeffs, fobj_b._mo_coeffs), + (fobj_a.nsocc, fobj_b.nsocc), + (fobj_a.nfsites, fobj_b.nfsites), + (fobj_a.efac, fobj_b.efac), + (fobj_a.TA, fobj_b.TA), + h1_ab, + hf_veff, + rdm1_tmp, + rdm2s, + fobj_a.dname, + eri_file=fobj_a.eri_file, + gcores=full_uhf.full_gcore, + frozen=frozen, + ) total_e = [sum(x) for x in zip(total_e, e_f)] if frag_energy: @@ -430,13 +519,13 @@ def solve_error(Fobjs, Nocc, only_chem=False): import math err_edge = [] - err_chempot = 0. + err_chempot = 0.0 if only_chem: for fobj in Fobjs: # Compute chemical potential error for each fragment for i in fobj.efac[1]: - err_chempot += fobj._rdm1[i,i] + err_chempot += fobj._rdm1[i, i] err_chempot /= Fobjs[0].unitcell_nkpt err = err_chempot - Nocc @@ -444,20 +533,19 @@ def solve_error(Fobjs, Nocc, only_chem=False): # Compute edge and chemical potential errors for fobj in Fobjs: - #match rdm-edge + # match rdm-edge for edge in fobj.edge_idx: - for j_ in range(len(edge)): for k_ in range(len(edge)): - if j_>k_: + if j_ > k_: continue err_edge.append(fobj._rdm1[edge[j_], edge[k_]]) - #chem potential + # chem potential for i in fobj.efac[1]: - err_chempot += fobj._rdm1[i,i] + err_chempot += fobj._rdm1[i, i] err_chempot /= Fobjs[0].unitcell_nkpt - err_edge.append(err_chempot) # far-end edges are included as err_chempot + err_edge.append(err_chempot) # far-end edges are included as err_chempot # Compute center errors err_cen = [] @@ -467,10 +555,9 @@ def solve_error(Fobjs, Nocc, only_chem=False): lenc = len(cens) for j_ in range(lenc): for k_ in range(lenc): - if j_>k_: + if j_ > k_: continue - err_cen.append(Fobjs[fobj.center[cindx]]._rdm1[cens[j_], - cens[k_]]) + err_cen.append(Fobjs[fobj.center[cindx]]._rdm1[cens[j_], cens[k_]]) err_cen.append(Nocc) err_edge = numpy.array(err_edge) @@ -480,10 +567,11 @@ def solve_error(Fobjs, Nocc, only_chem=False): err_vec = err_edge - err_cen # Compute the norm of the error vector - norm_ = numpy.mean(err_vec * err_vec)**0.5 + norm_ = numpy.mean(err_vec * err_vec) ** 0.5 return norm_, err_vec + def solve_mp2(mf, frozen=None, mo_coeff=None, mo_occ=None, mo_energy=None): """ Perform an MP2 (2nd order Moller-Plesset perturbation theory) calculation. @@ -512,14 +600,17 @@ def solve_mp2(mf, frozen=None, mo_coeff=None, mo_occ=None, mo_energy=None): from pyscf import mp # Set default values for optional parameters - if mo_coeff is None: mo_coeff = mf.mo_coeff - if mo_energy is None: mo_energy = mf.mo_energy - if mo_occ is None: mo_occ = mf.mo_occ + if mo_coeff is None: + mo_coeff = mf.mo_coeff + if mo_energy is None: + mo_energy = mf.mo_energy + if mo_occ is None: + mo_occ = mf.mo_occ # Initialize the MP2 object pt__ = mp.MP2(mf, frozen=frozen, mo_coeff=mo_coeff, mo_occ=mo_occ) mf = None - pt__.verbose=0 + pt__.verbose = 0 # Run the MP2 calculation pt__.kernel(mo_energy=mo_energy) @@ -527,9 +618,19 @@ def solve_mp2(mf, frozen=None, mo_coeff=None, mo_occ=None, mo_energy=None): return pt__ - -def solve_ccsd(mf, frozen=None, mo_coeff=None,relax=False, use_cumulant=False, with_dm1=True,rdm2_return = False, - mo_occ=None, mo_energy=None, rdm_return=False, verbose=0): +def solve_ccsd( + mf, + frozen=None, + mo_coeff=None, + relax=False, + use_cumulant=False, + with_dm1=True, + rdm2_return=False, + mo_occ=None, + mo_energy=None, + rdm_return=False, + verbose=0, +): """ Solve the CCSD (Coupled Cluster with Single and Double excitations) equations. @@ -574,33 +675,37 @@ def solve_ccsd(mf, frozen=None, mo_coeff=None,relax=False, use_cumulant=False, w from pyscf.cc.ccsd_rdm import make_rdm2 # Set default values for optional parameters - if mo_coeff is None: mo_coeff = mf.mo_coeff - if mo_energy is None: mo_energy = mf.mo_energy - if mo_occ is None: mo_occ = mf.mo_occ + if mo_coeff is None: + mo_coeff = mf.mo_coeff + if mo_energy is None: + mo_energy = mf.mo_energy + if mo_occ is None: + mo_occ = mf.mo_occ # Initialize the CCSD object - cc__ = cc.CCSD(mf, frozen=frozen, mo_coeff=mo_coeff, - mo_occ = mo_occ) - cc__.verbose=0 + cc__ = cc.CCSD(mf, frozen=frozen, mo_coeff=mo_coeff, mo_occ=mo_occ) + cc__.verbose = 0 mf = None - cc__.incore_complete=True + cc__.incore_complete = True # Prepare the integrals and Fock matrix eris = cc__.ao2mo() - eris.mo_energy=mo_energy + eris.mo_energy = mo_energy eris.fock = numpy.diag(mo_energy) # Solve the CCSD equations try: - cc__.verbose=verbose + cc__.verbose = verbose cc__.kernel(eris=eris) except: print(flush=True) - print('Exception in CCSD -> applying level_shift=0.2, diis_space=25',flush=True) + print( + "Exception in CCSD -> applying level_shift=0.2, diis_space=25", flush=True + ) print(flush=True) - cc__.verbose=4 - cc__.diis_space=25 - cc__.level_shift=0.2 + cc__.verbose = 4 + cc__.diis_space = 25 + cc__.level_shift = 0.2 cc__.kernel(eris=eris) # Extract the CCSD amplitudes @@ -612,21 +717,31 @@ def solve_ccsd(mf, frozen=None, mo_coeff=None,relax=False, use_cumulant=False, w if not relax: l1 = numpy.zeros_like(t1) l2 = numpy.zeros_like(t2) - rdm1a = cc.ccsd_rdm.make_rdm1(cc__, t1, t2, - l1,l2) + rdm1a = cc.ccsd_rdm.make_rdm1(cc__, t1, t2, l1, l2) else: rdm1a = cc__.make_rdm1(with_frozen=False) if rdm2_return: - if use_cumulant: with_dm1 = False - rdm2s = make_rdm2(cc__, cc__.t1, cc__.t2, cc__.l1, cc__.l2, with_frozen=False, ao_repr=False, with_dm1=with_dm1) - return(t1, t2, rdm1a, rdm2s) - return(t1, t2, rdm1a, cc__) + if use_cumulant: + with_dm1 = False + rdm2s = make_rdm2( + cc__, + cc__.t1, + cc__.t2, + cc__.l1, + cc__.l2, + with_frozen=False, + ao_repr=False, + with_dm1=with_dm1, + ) + return (t1, t2, rdm1a, rdm2s) + return (t1, t2, rdm1a, cc__) return (t1, t2) -def solve_block2(mf:object, nocc:int, frag_scratch:str = None, **solver_kwargs): - """ DMRG fragment solver using the pyscf.dmrgscf wrapper. + +def solve_block2(mf: object, nocc: int, frag_scratch: str = None, **solver_kwargs): + """DMRG fragment solver using the pyscf.dmrgscf wrapper. Parameters ---------- @@ -680,7 +795,7 @@ def solve_block2(mf:object, nocc:int, frag_scratch:str = None, **solver_kwargs): use_cumulant = solver_kwargs.pop("use_cumulant", True) norb = solver_kwargs.pop("norb", mf.mo_coeff.shape[1]) - nelec = solver_kwargs.pop("nelec", mf.mo_coeff.shape[1]) + nelec = solver_kwargs.pop("nelec", mf.mo_coeff.shape[1]) lo_method = solver_kwargs.pop("lo_method", None) startM = solver_kwargs.pop("startM", 25) maxM = solver_kwargs.pop("maxM", 500) @@ -688,54 +803,56 @@ def solve_block2(mf:object, nocc:int, frag_scratch:str = None, **solver_kwargs): max_mem = solver_kwargs.pop("max_mem", 100) max_noise = solver_kwargs.pop("max_noise", 1e-3) min_tol = solver_kwargs.pop("min_tol", 1e-8) - twodot_to_onedot = solver_kwargs.pop("twodot_to_onedot", int((5*max_iter)//6)) + twodot_to_onedot = solver_kwargs.pop("twodot_to_onedot", int((5 * max_iter) // 6)) root = solver_kwargs.pop("root", 0) - block_extra_keyword = solver_kwargs.pop("block_extra_keyword", ['fiedler']) - schedule_kwargs = solver_kwargs.pop("schedule_kwargs", {}) + block_extra_keyword = solver_kwargs.pop("block_extra_keyword", ["fiedler"]) + schedule_kwargs = solver_kwargs.pop("schedule_kwargs", {}) if norb <= 2: - block_extra_keyword = ['noreorder'] #Other reordering algorithms explode if the network is too small. + block_extra_keyword = [ + "noreorder" + ] # Other reordering algorithms explode if the network is too small. if lo_method is None: orbs = mf.mo_coeff elif isinstance(lo_method, str): - raise NotImplementedError("Localization within the fragment+bath subspace is currently not supported.") + raise NotImplementedError( + "Localization within the fragment+bath subspace is currently not supported." + ) mc = mcscf.CASCI(mf, norb, nelec) mc.fcisolver = dmrgscf.DMRGCI(mf.mol) ###Sweep scheduling - mc.fcisolver.scheduleSweeps = schedule_kwargs.pop("scheduleSweeps", - [(1*max_iter)//6, - (2*max_iter)//6, - (3*max_iter)//6, - (4*max_iter)//6, - (5*max_iter)//6, - max_iter] - ) - mc.fcisolver.scheduleMaxMs = schedule_kwargs.pop("scheduleMaxMs", - [startM if (startMijkl',hf_dm, hf_dm) + \ - numpy.einsum('ij,kl->ijkl',hf_dm, del_rdm1) + \ - numpy.einsum('ij,kl->ijkl',del_rdm1, hf_dm) - nc -= (numpy.einsum('ij,kl->iklj',hf_dm, hf_dm) + \ - numpy.einsum('ij,kl->iklj',hf_dm, del_rdm1) + \ - numpy.einsum('ij,kl->iklj',del_rdm1, hf_dm))*0.5 + del_rdm1[numpy.diag_indices(nocc)] -= 2.0 + nc = ( + numpy.einsum("ij,kl->ijkl", hf_dm, hf_dm) + + numpy.einsum("ij,kl->ijkl", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->ijkl", del_rdm1, hf_dm) + ) + nc -= ( + numpy.einsum("ij,kl->iklj", hf_dm, hf_dm) + + numpy.einsum("ij,kl->iklj", hf_dm, del_rdm1) + + numpy.einsum("ij,kl->iklj", del_rdm1, hf_dm) + ) * 0.5 rdm2 -= nc return rdm1, rdm2 -def solve_uccsd(mf, eris_inp, frozen=None, mo_coeff=None, relax=False, - use_cumulant=False, with_dm1=True, rdm2_return = False, - mo_occ=None, mo_energy=None, rdm_return=False, verbose=0): + +def solve_uccsd( + mf, + eris_inp, + frozen=None, + mo_coeff=None, + relax=False, + use_cumulant=False, + with_dm1=True, + rdm2_return=False, + mo_occ=None, + mo_energy=None, + rdm_return=False, + verbose=0, +): """ Solve the U-CCSD (Unrestricted Coupled Cluster with Single and Double excitations) equations. @@ -815,7 +948,7 @@ def solve_uccsd(mf, eris_inp, frozen=None, mo_coeff=None, relax=False, from molbe.external.uccsd_eri import make_eris_incore C = mf.mo_coeff - nao = [C[s].shape[0] for s in [0,1]] + nao = [C[s].shape[0] for s in [0, 1]] Vss = eris_inp[:2] Vos = eris_inp[-1] @@ -827,10 +960,10 @@ def ao2mofn(moish): # what spin component we are dealing with by comparing the first # 2-by-2 block of the mo coeff matrix. # Note that this assumes we have at least two basis functions - moish_feature = moish[:2,:2] + moish_feature = moish[:2, :2] s = -1 - for ss in [0,1]: - if numpy.allclose(moish_feature, C[ss][:2,:2]): + for ss in [0, 1]: + if numpy.allclose(moish_feature, C[ss][:2, :2]): s = ss break if s < 0: @@ -838,42 +971,64 @@ def ao2mofn(moish): return ao2mo.incore.full(Vss[s], moish, compact=False) elif isinstance(moish, list) or isinstance(moish, tuple): if len(moish) != 4: - raise RuntimeError("Expect a list/tuple of 4 numpy arrays but get %d of them." % len(moish)) - moish_feature = [mo[:2,:2] for mo in moish] - for s in [0,1]: - Cs_feature = C[s][:2,:2] - if not (numpy.allclose(moish_feature[2*s], Cs_feature) and - numpy.allclose(moish_feature[2*s+1], Cs_feature)): - raise RuntimeError("Expect a list/tuple of 4 numpy arrays in the order (moa,moa,mob,mob).") + raise RuntimeError( + "Expect a list/tuple of 4 numpy arrays but get %d of them." + % len(moish) + ) + moish_feature = [mo[:2, :2] for mo in moish] + for s in [0, 1]: + Cs_feature = C[s][:2, :2] + if not ( + numpy.allclose(moish_feature[2 * s], Cs_feature) + and numpy.allclose(moish_feature[2 * s + 1], Cs_feature) + ): + raise RuntimeError( + "Expect a list/tuple of 4 numpy arrays in the order (moa,moa,mob,mob)." + ) try: return ao2mo.incore.general(Vos, moish, compact=False) except: - return numpy.einsum('ijkl,ip,jq,kr,ls->pqrs', Vos, moish[0], moish[1], moish[2], moish[3], optimize=True) + return numpy.einsum( + "ijkl,ip,jq,kr,ls->pqrs", + Vos, + moish[0], + moish[1], + moish[2], + moish[3], + optimize=True, + ) else: - raise RuntimeError("moish must be either a numpy array or a list/tuple of 4 numpy arrays.") + raise RuntimeError( + "moish must be either a numpy array or a list/tuple of 4 numpy arrays." + ) # Initialize the UCCSD object ucc = cc.uccsd.UCCSD(mf, mo_coeff=mf.mo_coeff, mo_occ=mf.mo_occ) # Prepare the integrals - eris = make_eris_incore(ucc, Vss, Vos, mo_coeff=mf.mo_coeff, ao2mofn=ao2mofn, frozen=frozen) + eris = make_eris_incore( + ucc, Vss, Vos, mo_coeff=mf.mo_coeff, ao2mofn=ao2mofn, frozen=frozen + ) # Solve UCCSD equations: Level shifting options to be tested for unrestricted code - ucc.verbose=verbose + ucc.verbose = verbose ucc.kernel(eris=eris) # Compute and return the density matrices if requested if rdm_return: rdm1 = make_rdm1_uccsd(ucc, relax=relax) if rdm2_return: - if use_cumulant: with_dm1=False + if use_cumulant: + with_dm1 = False rdm2 = make_rdm2_uccsd(ucc, relax=relax, with_dm1=with_dm1) return (ucc, rdm1, rdm2) return (ucc, rdm1, None) return ucc -def schmidt_decomposition(mo_coeff, nocc, Frag_sites, cinv = None, rdm=None, norb=None, return_orb_count=False): +def schmidt_decomposition( + mo_coeff, nocc, Frag_sites, cinv=None, rdm=None, norb=None, return_orb_count=False +): """ Perform Schmidt decomposition on the molecular orbital coefficients. @@ -917,12 +1072,11 @@ def schmidt_decomposition(mo_coeff, nocc, Frag_sites, cinv = None, rdm=None, no # Compute the reduced density matrix (RDM) if not provided if not mo_coeff is None: - C = mo_coeff[:,:nocc] + C = mo_coeff[:, :nocc] if rdm is None: Dhf = numpy.dot(C, C.T) if not cinv is None: - Dhf = functools.reduce(numpy.dot, - (cinv, Dhf, cinv.conj().T)) + Dhf = functools.reduce(numpy.dot, (cinv, Dhf, cinv.conj().T)) else: Dhf = rdm @@ -930,10 +1084,8 @@ def schmidt_decomposition(mo_coeff, nocc, Frag_sites, cinv = None, rdm=None, no Tot_sites = Dhf.shape[0] # Identify environment sites (indices not in Frag_sites) - Env_sites1 = numpy.array([i for i in range(Tot_sites) - if not i in Frag_sites]) - Env_sites = numpy.array([[i] for i in range(Tot_sites) - if not i in Frag_sites]) + Env_sites1 = numpy.array([i for i in range(Tot_sites) if not i in Frag_sites]) + Env_sites = numpy.array([[i] for i in range(Tot_sites) if not i in Frag_sites]) Frag_sites1 = numpy.array([[i] for i in Frag_sites]) # Compute the environment part of the density matrix @@ -962,13 +1114,11 @@ def schmidt_decomposition(mo_coeff, nocc, Frag_sites, cinv = None, rdm=None, no # Initialize the transformation matrix (TA) TA = numpy.zeros([Tot_sites, len(Frag_sites) + len(Bidx)]) - TA[Frag_sites, :len(Frag_sites)] = numpy.eye(len(Frag_sites)) # Fragment part - TA[Env_sites1,len(Frag_sites):] = Evec[:,Bidx] # Environment part + TA[Frag_sites, : len(Frag_sites)] = numpy.eye(len(Frag_sites)) # Fragment part + TA[Env_sites1, len(Frag_sites) :] = Evec[:, Bidx] # Environment part if return_orb_count: # return TA, norbs_frag, norbs_bath return TA, Frag_sites1.shape[0], len(Bidx) else: return TA - - diff --git a/molbe/ube.py b/molbe/ube.py index 9d4547e7..5806a492 100644 --- a/molbe/ube.py +++ b/molbe/ube.py @@ -1,6 +1,6 @@ # Author(s): Minsik Cho, Leah Weisburn -""" 🍠 +"""🍠 Bootstrap Embedding Calculation with an Unrestricted Hartree-Fock Bath @@ -17,6 +17,7 @@ from .pfrag import Frags import molbe.be_var as be_var + class UBE(BE): # 🍠 def __init__( self, @@ -78,9 +79,9 @@ def __init__( self.hcore = mf.get_hcore() self.S = mf.get_ovlp() - self.C = [numpy.array(mf.mo_coeff[0]),numpy.array(mf.mo_coeff[1])] + self.C = [numpy.array(mf.mo_coeff[0]), numpy.array(mf.mo_coeff[1])] self.hf_dm = [mf.make_rdm1()[0], mf.make_rdm1()[1]] - self.hf_veff = [mf.get_veff()[0],mf.get_veff()[1]] + self.hf_veff = [mf.get_veff()[0], mf.get_veff()[1]] self.hf_etot = mf.e_tot self.W = None @@ -115,13 +116,30 @@ def __init__( self.Nocc[0] -= self.ncore self.Nocc[1] -= self.ncore - self.hf_dm = [numpy.dot(self.C[s][:,self.ncore:self.ncore+self.Nocc[s]], - self.C[s][:,self.ncore:self.ncore+self.Nocc[s]].T) for s in [0,1]] - self.C_core = [self.C[s][:, :self.ncore] for s in [0,1]] - self.P_core = [numpy.dot(self.C_core[s], self.C_core[s].T) for s in [0,1]] - self.core_veff = 1.*mf.get_veff(dm = self.P_core) - - self.E_core = sum([numpy.einsum("ji,ji->", 2*self.hcore+self.core_veff[s], self.P_core[s]) for s in [0,1]]) * 0.5 + self.hf_dm = [ + numpy.dot( + self.C[s][:, self.ncore : self.ncore + self.Nocc[s]], + self.C[s][:, self.ncore : self.ncore + self.Nocc[s]].T, + ) + for s in [0, 1] + ] + self.C_core = [self.C[s][:, : self.ncore] for s in [0, 1]] + self.P_core = [numpy.dot(self.C_core[s], self.C_core[s].T) for s in [0, 1]] + self.core_veff = 1.0 * mf.get_veff(dm=self.P_core) + + self.E_core = ( + sum( + [ + numpy.einsum( + "ji,ji->", + 2 * self.hcore + self.core_veff[s], + self.P_core[s], + ) + for s in [0, 1] + ] + ) + * 0.5 + ) # iao ignored for now self.C_a = numpy.array(mf.mo_coeff[0]) @@ -135,21 +153,21 @@ def __init__( valence_only=fobj.valence_only, ) - jobid='' + jobid = "" if be_var.CREATE_SCRATCH_DIR: try: - jobid = str(os.environ['SLURM_JOB_ID']) + jobid = str(os.environ["SLURM_JOB_ID"]) except: - jobid = '' - if not be_var.SCRATCH=='': - self.scratch_dir = be_var.SCRATCH+str(jobid) - os.system('mkdir -p '+self.scratch_dir) + jobid = "" + if not be_var.SCRATCH == "": + self.scratch_dir = be_var.SCRATCH + str(jobid) + os.system("mkdir -p " + self.scratch_dir) else: self.scratch_dir = None - if jobid == '': - self.eri_file = be_var.SCRATCH+eri_file + if jobid == "": + self.eri_file = be_var.SCRATCH + eri_file else: - self.eri_file = self.scratch_dir+'/'+eri_file + self.eri_file = self.scratch_dir + "/" + eri_file self.initialize(mf._eri, compute_hf) @@ -179,7 +197,7 @@ def initialize(self, eri_, compute_hf): center_idx=self.center_idx[I], efac=self.ebe_weight[I], centerf_idx=self.centerf_idx[I], - unrestricted=True + unrestricted=True, ) else: fobjs_a = Frags( @@ -192,7 +210,7 @@ def initialize(self, eri_, compute_hf): center_idx=[], centerf_idx=[], efac=self.ebe_weight[I], - unrestricted=True + unrestricted=True, ) self.Fobjs_a.append(fobjs_a) # beta @@ -208,7 +226,7 @@ def initialize(self, eri_, compute_hf): center_idx=self.center_idx[I], efac=self.ebe_weight[I], centerf_idx=self.centerf_idx[I], - unrestricted=True + unrestricted=True, ) else: fobjs_b = Frags( @@ -221,7 +239,7 @@ def initialize(self, eri_, compute_hf): center_idx=[], centerf_idx=[], efac=self.ebe_weight[I], - unrestricted=True + unrestricted=True, ) self.Fobjs_b.append(fobjs_b) @@ -237,13 +255,29 @@ def initialize(self, eri_, compute_hf): if self.frozen_core: fobj_a.core_veff = self.core_veff[0] fobj_b.core_veff = self.core_veff[1] - orb_count_a.append(fobj_a.sd(self.W[0], self.lmo_coeff_a, self.Nocc[0], return_orb_count=True)) - orb_count_b.append(fobj_b.sd(self.W[1], self.lmo_coeff_b, self.Nocc[1], return_orb_count=True)) + orb_count_a.append( + fobj_a.sd( + self.W[0], self.lmo_coeff_a, self.Nocc[0], return_orb_count=True + ) + ) + orb_count_b.append( + fobj_b.sd( + self.W[1], self.lmo_coeff_b, self.Nocc[1], return_orb_count=True + ) + ) else: fobj_a.core_veff = None fobj_b.core_veff = None - orb_count_a.append(fobj_a.sd(self.W, self.lmo_coeff_a, self.Nocc[0], return_orb_count=True)) - orb_count_b.append(fobj_b.sd(self.W, self.lmo_coeff_b, self.Nocc[1], return_orb_count=True)) + orb_count_a.append( + fobj_a.sd( + self.W, self.lmo_coeff_a, self.Nocc[0], return_orb_count=True + ) + ) + orb_count_b.append( + fobj_b.sd( + self.W, self.lmo_coeff_b, self.Nocc[1], return_orb_count=True + ) + ) all_noccs.append(self.Nocc) @@ -252,21 +286,21 @@ def initialize(self, eri_, compute_hf): if eri_ is None and not self.mf.with_df is None: # NOT IMPLEMENTED: should not be called, as no unrestricted DF tested # for density-fitted integrals; if mf is provided, pyscf.ao2mo uses DF object in an outcore fashion - eri_a = ao2mo.kernel( - self.mf.mol, fobj_a.TA, compact=True) - eri_b = ao2mo.kernel( - self.mf.mol, fobj_b.TA, compact=True) + eri_a = ao2mo.kernel(self.mf.mol, fobj_a.TA, compact=True) + eri_b = ao2mo.kernel(self.mf.mol, fobj_b.TA, compact=True) else: eri_a = ao2mo.incore.full( - eri_, fobj_a.TA, compact=True) # otherwise, do an incore ao2mo - eri_b = ao2mo.incore.full( - eri_, fobj_b.TA, compact=True) + eri_, fobj_a.TA, compact=True + ) # otherwise, do an incore ao2mo + eri_b = ao2mo.incore.full(eri_, fobj_b.TA, compact=True) - Csd_A = fobj_a.TA # may have to add in nibath here + Csd_A = fobj_a.TA # may have to add in nibath here Csd_B = fobj_b.TA # cross-spin ERI term - eri_ab = ao2mo.incore.general(eri_, (Csd_A,Csd_A,Csd_B,Csd_B), compact=True) + eri_ab = ao2mo.incore.general( + eri_, (Csd_A, Csd_A, Csd_B, Csd_B), compact=True + ) file_eri.create_dataset(fobj_a.dname[0], data=eri_a) file_eri.create_dataset(fobj_a.dname[1], data=eri_b) @@ -283,11 +317,14 @@ def initialize(self, eri_, compute_hf): fobj_a.heff = numpy.zeros_like(fobj_a.h1) fobj_a.scf(fs=True, eri=eri_a) fobj_a.dm0 = numpy.dot( - fobj_a._mo_coeffs[:, : fobj_a.nsocc], fobj_a._mo_coeffs[:, : fobj_a.nsocc].conj().T + fobj_a._mo_coeffs[:, : fobj_a.nsocc], + fobj_a._mo_coeffs[:, : fobj_a.nsocc].conj().T, ) if compute_hf: - eh1_a, ecoul_a, ef_a = fobj_a.energy_hf(return_e1=True, unrestricted=True, spin_ind=0) + eh1_a, ecoul_a, ef_a = fobj_a.energy_hf( + return_e1=True, unrestricted=True, spin_ind=0 + ) EH1 += eh1_a ECOUL += ecoul_a E_hf += fobj_a.ebe_hf @@ -302,33 +339,59 @@ def initialize(self, eri_, compute_hf): fobj_b.scf(fs=True, eri=eri_b) fobj_b.dm0 = numpy.dot( - fobj_b._mo_coeffs[:, : fobj_b.nsocc], fobj_b._mo_coeffs[:, : fobj_b.nsocc].conj().T + fobj_b._mo_coeffs[:, : fobj_b.nsocc], + fobj_b._mo_coeffs[:, : fobj_b.nsocc].conj().T, ) if compute_hf: - eh1_b, ecoul_b, ef_b = fobj_b.energy_hf(return_e1=True, unrestricted=True, spin_ind=1) + eh1_b, ecoul_b, ef_b = fobj_b.energy_hf( + return_e1=True, unrestricted=True, spin_ind=1 + ) EH1 += eh1_b ECOUL += ecoul_b E_hf += fobj_b.ebe_hf file_eri.close() print("Number of Orbitals per Fragment:", flush=True) - print("____________________________________________________________________", flush=True) - print("| Fragment | Nocc | Fragment Orbs | Bath Orbs | Schmidt Space |", flush=True) - print("____________________________________________________________________", flush=True) + print( + "____________________________________________________________________", + flush=True, + ) + print( + "| Fragment | Nocc | Fragment Orbs | Bath Orbs | Schmidt Space |", + flush=True, + ) + print( + "____________________________________________________________________", + flush=True, + ) for I in range(self.Nfrag): - print('| {:>2} | ({:>3},{:>3}) | ({:>3},{:>3}) | ({:>3},{:>3}) | ({:>3},{:>3}) |'.format(I, all_noccs[I][0],all_noccs[I][1], - orb_count_a[I][0], orb_count_b[I][0], - orb_count_a[I][1], orb_count_b[I][1], - orb_count_a[I][0]+orb_count_a[I][1], - orb_count_b[I][0]+orb_count_b[I][1]), - flush=True) - print("____________________________________________________________________", flush=True) + print( + "| {:>2} | ({:>3},{:>3}) | ({:>3},{:>3}) | ({:>3},{:>3}) | ({:>3},{:>3}) |".format( + I, + all_noccs[I][0], + all_noccs[I][1], + orb_count_a[I][0], + orb_count_b[I][0], + orb_count_a[I][1], + orb_count_b[I][1], + orb_count_a[I][0] + orb_count_a[I][1], + orb_count_b[I][0] + orb_count_b[I][1], + ), + flush=True, + ) + print( + "____________________________________________________________________", + flush=True, + ) if compute_hf: hf_err = self.hf_etot - (E_hf + self.enuc + self.E_core) self.ebe_hf = E_hf + self.enuc + self.E_core - self.ek - print("HF-in-HF error : {:>.4e} Ha".format(hf_err), flush=True) + print( + "HF-in-HF error : {:>.4e} Ha".format(hf_err), + flush=True, + ) if abs(hf_err) > 1.0e-5: print("WARNING!!! Large HF-in-HF energy error") print("eh1 ", EH1) @@ -346,36 +409,40 @@ def initialize(self, eri_, compute_hf): fobj.udim = couti couti = fobj.set_udim(couti) - def oneshot(self, solver="UCCSD", nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False): - + def oneshot( + self, solver="UCCSD", nproc=1, ompnum=4, calc_frag_energy=False, clean_eri=False + ): from .solver import be_func_u from .be_parallel import be_func_parallel_u if nproc == 1: - E, E_comp = be_func_u(None, - zip(self.Fobjs_a, self.Fobjs_b), - solver, - self.enuc, - hf_veff=self.hf_veff, - eeval=True, - ereturn=True, - relax_density=False, - frag_energy=calc_frag_energy, - frozen=self.frozen_core) + E, E_comp = be_func_u( + None, + zip(self.Fobjs_a, self.Fobjs_b), + solver, + self.enuc, + hf_veff=self.hf_veff, + eeval=True, + ereturn=True, + relax_density=False, + frag_energy=calc_frag_energy, + frozen=self.frozen_core, + ) else: - E, E_comp = be_func_parallel_u(None, - zip(self.Fobjs_a, self.Fobjs_b), - solver, - self.enuc, - hf_veff=self.hf_veff, - eeval=True, - ereturn=True, - relax_density=False, - frag_energy=calc_frag_energy, - frozen=self.frozen_core, - nproc=nproc, - ompnum=ompnum) - + E, E_comp = be_func_parallel_u( + None, + zip(self.Fobjs_a, self.Fobjs_b), + solver, + self.enuc, + hf_veff=self.hf_veff, + eeval=True, + ereturn=True, + relax_density=False, + frag_energy=calc_frag_energy, + frozen=self.frozen_core, + nproc=nproc, + ompnum=ompnum, + ) print("-----------------------------------------------------", flush=True) print(" One Shot BE ", flush=True) @@ -394,6 +461,7 @@ def oneshot(self, solver="UCCSD", nproc=1, ompnum=4, calc_frag_energy=False, cle except: print("Scratch directory not removed") + def initialize_pot(Nfrag, edge_idx): pot_ = [] diff --git a/setup.py b/setup.py index 3ba4725a..e12bb2cd 100644 --- a/setup.py +++ b/setup.py @@ -2,28 +2,29 @@ # Read the contents of your README file from pathlib import Path + this_directory = Path(__file__).parent long_description = (this_directory / "README.md").read_text() setup( - name='quemb', - version='1.0', - description='QuEmb: A framework for efficient simulation of large molecules, surfaces, and solids via Bootstrap Embedding', + name="quemb", + version="1.0", + description="QuEmb: A framework for efficient simulation of large molecules, surfaces, and solids via Bootstrap Embedding", long_description=long_description, - long_description_content_type='text/markdown', - url='https://github.com/oimeitei/quemb', - license='Apache 2.0', + long_description_content_type="text/markdown", + url="https://github.com/oimeitei/quemb", + license="Apache 2.0", packages=find_packages(), - python_requires='>=3.7', + python_requires=">=3.7", install_requires=[ - 'numpy>=1.22.0', - 'scipy>=1.7.0', - 'pyscf>=2.0.0', - 'matplotlib', + "numpy>=1.22.0", + "scipy>=1.7.0", + "pyscf>=2.0.0", + "matplotlib", # TODO this temporarily points to mcocdawc/libdmet_preview. # as soon as this PR: https://github.com/gkclab/libdmet_preview/pull/21 # is merged, let it point to the original libdmet # https://github.com/gkclab/libdmet_preview - 'libdmet @ git+https://github.com/mcocdawc/libdmet_preview.git@add_fixes_for_BE' + "libdmet @ git+https://github.com/mcocdawc/libdmet_preview.git@add_fixes_for_BE", ], ) diff --git a/tests/chem_dm_kBE_test.py b/tests/chem_dm_kBE_test.py index fcbf5576..bcde09d7 100644 --- a/tests/chem_dm_kBE_test.py +++ b/tests/chem_dm_kBE_test.py @@ -9,6 +9,7 @@ from pyscf.pbc import gto, scf, df from kbe import fragpart, BE + class Test_kBE_Full(unittest.TestCase): try: import libdmet @@ -20,50 +21,69 @@ def test_kc2_sto3g_be1_chempot(self): kpt = [1, 1, 1] cell = gto.Cell() - a = 1. - b = 1. - c = 12. + a = 1.0 + b = 1.0 + c = 12.0 lat = numpy.eye(3) - lat[0,0] = a - lat[1,1] = b - lat[2,2] = c + lat[0, 0] = a + lat[1, 1] = b + lat[2, 2] = c cell.a = lat - cell.atom = [['C', (0.,0.,i*3.)] for i in range(2)] + cell.atom = [["C", (0.0, 0.0, i * 3.0)] for i in range(2)] - cell.unit='Angstrom' - cell.basis = 'sto-3g' - cell.verbose=0 + cell.unit = "Angstrom" + cell.basis = "sto-3g" + cell.verbose = 0 cell.build() - self.periodic_test(cell, kpt, 'be1', 'C2 (kBE1)', 'autogen', -74.64695833012868, only_chem = True) + self.periodic_test( + cell, kpt, "be1", "C2 (kBE1)", "autogen", -74.64695833012868, only_chem=True + ) def test_kc4_sto3g_be2_density(self): kpt = [1, 1, 1] cell = gto.Cell() - a = 1. - b = 1. - c = 12. + a = 1.0 + b = 1.0 + c = 12.0 lat = numpy.eye(3) - lat[0,0] = a - lat[1,1] = b - lat[2,2] = c + lat[0, 0] = a + lat[1, 1] = b + lat[2, 2] = c cell.a = lat - cell.atom = [['C', (0.,0.,i*3.)] for i in range(4)] + cell.atom = [["C", (0.0, 0.0, i * 3.0)] for i in range(4)] - cell.unit='Angstrom' - cell.basis = 'sto-3g' - cell.verbose=0 + cell.unit = "Angstrom" + cell.basis = "sto-3g" + cell.verbose = 0 cell.build() - self.periodic_test(cell, kpt, 'be2', 'C4 (kBE2)', 'autogen', -149.4085332249809, only_chem = False) - - def periodic_test(self, cell, kpt, be_type, test_name, frag_type, target, delta = 1e-4, only_chem = True): - + self.periodic_test( + cell, + kpt, + "be2", + "C4 (kBE2)", + "autogen", + -149.4085332249809, + only_chem=False, + ) + + def periodic_test( + self, + cell, + kpt, + be_type, + test_name, + frag_type, + target, + delta=1e-4, + only_chem=True, + ): kpts = cell.make_kpts(kpt, wrap_around=True) mydf = df.GDF(cell, kpts) mydf.build() @@ -74,17 +94,21 @@ def periodic_test(self, cell, kpt, be_type, test_name, frag_type, target, delta kmf.conv_tol = 1e-12 kmf.kernel() - kfrag = fragpart(be_type=be_type, - mol=cell, - frag_type=frag_type, - kpt=kpt, - frozen_core=True) + kfrag = fragpart( + be_type=be_type, mol=cell, frag_type=frag_type, kpt=kpt, frozen_core=True + ) mykbe = BE(kmf, kfrag, kpts=kpts) - mykbe.optimize(solver='CCSD', only_chem=only_chem) + mykbe.optimize(solver="CCSD", only_chem=only_chem) + + self.assertAlmostEqual( + mykbe.ebe_tot, + target, + msg="kBE Correlation Energy for " + + test_name + + " does not match the expected correlation energy!", + delta=delta, + ) - self.assertAlmostEqual(mykbe.ebe_tot, target, - msg = "kBE Correlation Energy for " + test_name - + " does not match the expected correlation energy!", delta = delta) -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/chempot_molBE_test.py b/tests/chempot_molBE_test.py index 447189a0..73757fe5 100644 --- a/tests/chempot_molBE_test.py +++ b/tests/chempot_molBE_test.py @@ -8,40 +8,61 @@ from pyscf import cc, gto, scf from molbe import fragpart, BE + class TestBE_restricted(unittest.TestCase): def test_h8_sto3g_ben(self): # Linear Equidistant (r=1Å) H8 Chain, STO-3G # CCSD Total Energy: -4.306498896 Ha # Target BE Total energies from in-house code mol = gto.M() - mol.atom = [['H', (0.,0.,i)] for i in range(8)] - mol.basis = 'sto-3g' - mol.charge = 0.; mol.spin = 0. + mol.atom = [["H", (0.0, 0.0, i)] for i in range(8)] + mol.basis = "sto-3g" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - self.molecular_restricted_test(mol, 'be2', 'H8 (BE2)', 'hchain_simple', -4.30628355, only_chem = True) - self.molecular_restricted_test(mol, 'be3', 'H8 (BE3)', 'hchain_simple', -4.30649890, only_chem = True) + self.molecular_restricted_test( + mol, "be2", "H8 (BE2)", "hchain_simple", -4.30628355, only_chem=True + ) + self.molecular_restricted_test( + mol, "be3", "H8 (BE3)", "hchain_simple", -4.30649890, only_chem=True + ) - @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + @unittest.skipIf( + os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions" + ) def test_octane_sto3g_ben(self): # Octane, STO-3G # CCSD Total Energy: -310.3344616 Ha mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/octane.xyz') - mol.basis = 'sto-3g' - mol.charge = 0.; mol.spin = 0. + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/octane.xyz") + mol.basis = "sto-3g" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - self.molecular_restricted_test(mol, 'be2', 'Octane (BE2)', 'autogen', -310.33471581, only_chem = True) - self.molecular_restricted_test(mol, 'be3', 'Octane (BE3)', 'autogen', -310.33447096, only_chem = True) + self.molecular_restricted_test( + mol, "be2", "Octane (BE2)", "autogen", -310.33471581, only_chem=True + ) + self.molecular_restricted_test( + mol, "be3", "Octane (BE3)", "autogen", -310.33447096, only_chem=True + ) - def molecular_restricted_test(self, mol, be_type, test_name, frag_type, target, delta = 1e-4, - only_chem = True): - mf = scf.RHF(mol); mf.kernel() - fobj = fragpart(frag_type=frag_type, be_type=be_type, mol = mol) + def molecular_restricted_test( + self, mol, be_type, test_name, frag_type, target, delta=1e-4, only_chem=True + ): + mf = scf.RHF(mol) + mf.kernel() + fobj = fragpart(frag_type=frag_type, be_type=be_type, mol=mol) mybe = BE(mf, fobj) - mybe.optimize(solver='CCSD', method='QN', only_chem=only_chem) - self.assertAlmostEqual(mybe.ebe_tot, target, - msg = "BE Correlation Energy (Chem. Pot. Optimization) for " + test_name - + " does not match the expected correlation energy!", delta = delta) + mybe.optimize(solver="CCSD", method="QN", only_chem=only_chem) + self.assertAlmostEqual( + mybe.ebe_tot, + target, + msg="BE Correlation Energy (Chem. Pot. Optimization) for " + + test_name + + " does not match the expected correlation energy!", + delta=delta, + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/dm_molBE_test.py b/tests/dm_molBE_test.py index 4387a280..d8040ddd 100644 --- a/tests/dm_molBE_test.py +++ b/tests/dm_molBE_test.py @@ -8,27 +8,44 @@ from pyscf import cc, gto, scf from molbe import fragpart, BE + class TestBE_restricted(unittest.TestCase): # TODO: Add test against known values (molecular_restrict_test) def test_h8_sto3g_ben_trustRegion(self): # Test consistency between two QN methods mol = gto.M() - mol.atom = [['H', (0.,0.,i)] for i in range(7)]; mol.atom.append(['H', (0.,0.,4.2)]) - mol.basis = 'sto-3g' - mol.charge = 0.; mol.spin = 0. + mol.atom = [["H", (0.0, 0.0, i)] for i in range(7)] + mol.atom.append(["H", (0.0, 0.0, 4.2)]) + mol.basis = "sto-3g" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - self.molecular_QN_test(mol, 'be2', 'H8 (BE2)', 'hchain_simple', only_chem = False) + self.molecular_QN_test(mol, "be2", "H8 (BE2)", "hchain_simple", only_chem=False) - def molecular_QN_test(self, mol, be_type, test_name, frag_type, delta = 1e-6, only_chem = True): - mf = scf.RHF(mol); mf.max_cycle = 100; mf.kernel() - fobj = fragpart(frag_type=frag_type, be_type=be_type, mol = mol) + def molecular_QN_test( + self, mol, be_type, test_name, frag_type, delta=1e-6, only_chem=True + ): + mf = scf.RHF(mol) + mf.max_cycle = 100 + mf.kernel() + fobj = fragpart(frag_type=frag_type, be_type=be_type, mol=mol) mybe1 = BE(mf, fobj) - mybe1.optimize(solver='CCSD', method='QN', trust_region=False, only_chem=only_chem) + mybe1.optimize( + solver="CCSD", method="QN", trust_region=False, only_chem=only_chem + ) mybe2 = BE(mf, fobj) - mybe2.optimize(solver='CCSD', method='QN', trust_region=True, only_chem=only_chem) - self.assertAlmostEqual(mybe1.ebe_tot, mybe2.ebe_tot, - msg = "BE Correlation Energy (DM) for " + test_name - + " does not return comparable results from two QN optimization!", delta = delta) + mybe2.optimize( + solver="CCSD", method="QN", trust_region=True, only_chem=only_chem + ) + self.assertAlmostEqual( + mybe1.ebe_tot, + mybe2.ebe_tot, + msg="BE Correlation Energy (DM) for " + + test_name + + " does not return comparable results from two QN optimization!", + delta=delta, + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/dmrg_molBE_test.py b/tests/dmrg_molBE_test.py index 2e2bdf12..998377d0 100644 --- a/tests/dmrg_molBE_test.py +++ b/tests/dmrg_molBE_test.py @@ -9,44 +9,66 @@ from pyscf import gto, scf from molbe import fragpart, BE + class TestBE_DMRG(unittest.TestCase): try: from pyscf import dmrgscf except ImportError: - dmrgscf=None + dmrgscf = None @unittest.skipIf(dmrgscf is None, "Module `pyscf.dmrgscf` not imported correctly.") def test_h8_sto3g_pipek(self): mol = gto.M() - mol.atom = [['H', (0.,0.,i*1.2)] for i in range(8)] - mol.basis = 'sto-3g' + mol.atom = [["H", (0.0, 0.0, i * 1.2)] for i in range(8)] + mol.basis = "sto-3g" mol.charge = 0 mol.spin = 0 mol.build() - self.molecular_DMRG_test(mol, 'be1', 100, 'H8 (BE1)', 'hchain_simple', -4.20236532) + self.molecular_DMRG_test( + mol, "be1", 100, "H8 (BE1)", "hchain_simple", -4.20236532 + ) - def molecular_DMRG_test(self, mol, be_type, maxM, test_name, frag_type, target, delta = 1e-4): + def molecular_DMRG_test( + self, mol, be_type, maxM, test_name, frag_type, target, delta=1e-4 + ): with tempfile.TemporaryDirectory() as tmp: - mf = scf.RHF(mol); mf.kernel() + mf = scf.RHF(mol) + mf.kernel() fobj = fragpart(frag_type=frag_type, be_type=be_type, mol=mol) - mybe = BE(mf, fobj, lo_method='pipek', pop_method='lowdin') - mybe.oneshot(solver='block2', - scratch_dir=str(tmp), - maxM=int(maxM), - maxIter=30, - force_cleanup=True) - self.assertAlmostEqual(mybe.ebe_tot, target, - msg = "BE Correlation Energy (Chem. Pot. Optimization) for " + test_name - + " does not match the expected correlation energy!", delta = delta) + mybe = BE(mf, fobj, lo_method="pipek", pop_method="lowdin") + mybe.oneshot( + solver="block2", + scratch_dir=str(tmp), + maxM=int(maxM), + maxIter=30, + force_cleanup=True, + ) + self.assertAlmostEqual( + mybe.ebe_tot, + target, + msg="BE Correlation Energy (Chem. Pot. Optimization) for " + + test_name + + " does not match the expected correlation energy!", + delta=delta, + ) tmpfiles = [] for path, dir_n, file_n in os.walk(tmp, topdown=True): for n in file_n: - if n.startswith('F.') or n.startswith('FCIDUMP') or n.startswith('node'): + if ( + n.startswith("F.") + or n.startswith("FCIDUMP") + or n.startswith("node") + ): tmpfiles.append(n) try: assert len(tmpfiles) == 0 except Exception as inst: - print(f"DMRG tempfiles not cleared correctly in {test_name}:\n", inst, tmpfiles) + print( + f"DMRG tempfiles not cleared correctly in {test_name}:\n", + inst, + tmpfiles, + ) + -if __name__ == '__main__': - unittest.main() \ No newline at end of file +if __name__ == "__main__": + unittest.main() diff --git a/tests/eri_onthefly_test.py b/tests/eri_onthefly_test.py index bd442f6e..35208edb 100644 --- a/tests/eri_onthefly_test.py +++ b/tests/eri_onthefly_test.py @@ -8,21 +8,31 @@ from pyscf import gto, scf from molbe import fragpart, BE + class TestDF_ontheflyERI(unittest.TestCase): - @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + @unittest.skipIf( + os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions" + ) def test_octane_BE2(self): # Octane, cc-pvtz mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/octane.xyz') - mol.basis = 'cc-pvtz' - mol.charge = 0.; mol.spin = 0. + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/octane.xyz") + mol.basis = "cc-pvtz" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - mf = scf.RHF(mol); mf.direct_scf = True; mf.kernel() - fobj = fragpart(frag_type='autogen', be_type='be2', mol = mol) + mf = scf.RHF(mol) + mf.direct_scf = True + mf.kernel() + fobj = fragpart(frag_type="autogen", be_type="be2", mol=mol) mybe = BE(mf, fobj, integral_direct_DF=True) - self.assertAlmostEqual(mybe.ebe_hf, mf.e_tot, - msg = "HF-in-HF energy for Octane (BE2) does not match the HF energy!", - delta = 1e-6) + self.assertAlmostEqual( + mybe.ebe_hf, + mf.e_tot, + msg="HF-in-HF energy for Octane (BE2) does not match the HF energy!", + delta=1e-6, + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/hf-in-hf_BE_test.py b/tests/hf-in-hf_BE_test.py index 0fa5fdeb..343f7332 100644 --- a/tests/hf-in-hf_BE_test.py +++ b/tests/hf-in-hf_BE_test.py @@ -8,47 +8,56 @@ from pyscf import gto, scf from molbe import fragpart, BE + class TestHFinHF_restricted(unittest.TestCase): def test_h8_sto3g_ben(self): # Linear Equidistant (r=1Å) H8 Chain, STO-3G mol = gto.M() - mol.atom = [['H', (0.,0.,i)] for i in range(8)] - mol.basis = 'sto-3g' - mol.charge = 0.; mol.spin = 0. + mol.atom = [["H", (0.0, 0.0, i)] for i in range(8)] + mol.basis = "sto-3g" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - self.molecular_restricted_test(mol, 'be1', 'H8 (BE1)') - self.molecular_restricted_test(mol, 'be2', 'H8 (BE2)') - self.molecular_restricted_test(mol, 'be3', 'H8 (BE3)') + self.molecular_restricted_test(mol, "be1", "H8 (BE1)") + self.molecular_restricted_test(mol, "be2", "H8 (BE2)") + self.molecular_restricted_test(mol, "be3", "H8 (BE3)") def test_h8_ccpvdz_ben(self): # Linear Equidistant (r=1Å) H8 Chain, cc-pVDZ mol = gto.M() - mol.atom = [['H', (0.,0.,i)] for i in range(8)] - mol.basis = 'cc-pvdz' - mol.charge = 0.; mol.spin = 0. + mol.atom = [["H", (0.0, 0.0, i)] for i in range(8)] + mol.basis = "cc-pvdz" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - self.molecular_restricted_test(mol, 'be1', 'H8 (BE1)') - self.molecular_restricted_test(mol, 'be2', 'H8 (BE2)') - self.molecular_restricted_test(mol, 'be3', 'H8 (BE3)') + self.molecular_restricted_test(mol, "be1", "H8 (BE1)") + self.molecular_restricted_test(mol, "be2", "H8 (BE2)") + self.molecular_restricted_test(mol, "be3", "H8 (BE3)") def test_octane_sto3g_ben(self): # Octane, STO-3G mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/octane.xyz') - mol.basis = 'sto-3g' - mol.charge = 0.; mol.spin = 0. + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/octane.xyz") + mol.basis = "sto-3g" + mol.charge = 0.0 + mol.spin = 0.0 mol.build() - self.molecular_restricted_test(mol, 'be1', 'Octane (BE1)') - self.molecular_restricted_test(mol, 'be2', 'Octane (BE2)') - self.molecular_restricted_test(mol, 'be3', 'Octane (BE3)') + self.molecular_restricted_test(mol, "be1", "Octane (BE1)") + self.molecular_restricted_test(mol, "be2", "Octane (BE2)") + self.molecular_restricted_test(mol, "be3", "Octane (BE3)") - def molecular_restricted_test(self, mol, be_type, test_name, delta = 1e-5): - mf = scf.RHF(mol); mf.kernel() - fobj = fragpart(frag_type='autogen', be_type=be_type, mol = mol) + def molecular_restricted_test(self, mol, be_type, test_name, delta=1e-5): + mf = scf.RHF(mol) + mf.kernel() + fobj = fragpart(frag_type="autogen", be_type=be_type, mol=mol) mybe = BE(mf, fobj) - self.assertAlmostEqual(mybe.ebe_hf, mf.e_tot, - msg = "HF-in-HF energy for " + test_name - + " does not match the HF energy!", delta = delta) + self.assertAlmostEqual( + mybe.ebe_hf, + mf.e_tot, + msg="HF-in-HF energy for " + test_name + " does not match the HF energy!", + delta=delta, + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main() diff --git a/tests/ube-oneshot_test.py b/tests/ube-oneshot_test.py index e2f6a639..81a75e3a 100644 --- a/tests/ube-oneshot_test.py +++ b/tests/ube-oneshot_test.py @@ -10,62 +10,109 @@ from pyscf import gto, scf from molbe import fragpart, UBE + class TestOneShot_Unrestricted(unittest.TestCase): - @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + @unittest.skipIf( + os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions" + ) def test_hexene_anion_sto3g_frz_ben(self): # Linear Equidistant (r=1Å) H8 Chain, STO-3G mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/hexene.xyz') - mol.basis = 'sto-3g' - mol.charge = -1; mol.spin = 1 + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/hexene.xyz") + mol.basis = "sto-3g" + mol.charge = -1 + mol.spin = 1 mol.build() - self.molecular_unrestricted_oneshot_test(mol, 'be1', 'Hexene Anion Frz (BE1)', True, -0.35753374) - self.molecular_unrestricted_oneshot_test(mol, 'be2', 'Hexene Anion Frz (BE2)', True, -0.34725961) - self.molecular_unrestricted_oneshot_test(mol, 'be3', 'Hexene Anion Frz (BE3)', True, -0.34300834) + self.molecular_unrestricted_oneshot_test( + mol, "be1", "Hexene Anion Frz (BE1)", True, -0.35753374 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be2", "Hexene Anion Frz (BE2)", True, -0.34725961 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be3", "Hexene Anion Frz (BE3)", True, -0.34300834 + ) - @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + @unittest.skipIf( + os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions" + ) def test_hexene_cation_sto3g_frz_ben(self): # Linear Equidistant (r=1Å) H8 Chain, cc-pVDZ mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/hexene.xyz') - mol.basis = 'sto-3g' - mol.charge = 1; mol.spin = 1 + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/hexene.xyz") + mol.basis = "sto-3g" + mol.charge = 1 + mol.spin = 1 mol.build() - self.molecular_unrestricted_oneshot_test(mol, 'be1', 'Hexene Cation Frz (BE1)', True, -0.40383508) - self.molecular_unrestricted_oneshot_test(mol, 'be2', 'Hexene Cation Frz (BE2)', True, -0.36496690) - self.molecular_unrestricted_oneshot_test(mol, 'be3', 'Hexene Cation Frz (BE3)', True, -0.36996484) + self.molecular_unrestricted_oneshot_test( + mol, "be1", "Hexene Cation Frz (BE1)", True, -0.40383508 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be2", "Hexene Cation Frz (BE2)", True, -0.36496690 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be3", "Hexene Cation Frz (BE3)", True, -0.36996484 + ) - @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + @unittest.skipIf( + os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions" + ) def test_hexene_anion_sto3g_unfrz_ben(self): # Octane, STO-3G mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/hexene.xyz') - mol.basis = 'sto-3g' - mol.charge = -1; mol.spin = 1 + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/hexene.xyz") + mol.basis = "sto-3g" + mol.charge = -1 + mol.spin = 1 mol.build() - self.molecular_unrestricted_oneshot_test(mol, 'be1', 'Hexene Anion Unfrz (BE1)', False, -0.38478279) - self.molecular_unrestricted_oneshot_test(mol, 'be2', 'Hexene Anion Unfrz (BE2)', False, -0.39053689) - self.molecular_unrestricted_oneshot_test(mol, 'be3', 'Hexene Anion Unfrz (BE3)', False, -0.38960174) + self.molecular_unrestricted_oneshot_test( + mol, "be1", "Hexene Anion Unfrz (BE1)", False, -0.38478279 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be2", "Hexene Anion Unfrz (BE2)", False, -0.39053689 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be3", "Hexene Anion Unfrz (BE3)", False, -0.38960174 + ) - @unittest.skipIf(os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions") + @unittest.skipIf( + os.getenv("GITHUB_ACTIONS") == "true", "Skip expensive tests on Github Actions" + ) def test_hexene_cation_sto3g_unfrz_ben(self): mol = gto.M() - mol.atom = os.path.join(os.path.dirname(__file__), 'xyz/hexene.xyz') - mol.basis = 'sto-3g' - mol.charge = 1; mol.spin = 1 + mol.atom = os.path.join(os.path.dirname(__file__), "xyz/hexene.xyz") + mol.basis = "sto-3g" + mol.charge = 1 + mol.spin = 1 mol.build() - self.molecular_unrestricted_oneshot_test(mol, 'be1', 'Hexene Cation Frz (BE1)', False, -0.39471433) - self.molecular_unrestricted_oneshot_test(mol, 'be2', 'Hexene Cation Frz (BE2)', False, -0.39846777) - self.molecular_unrestricted_oneshot_test(mol, 'be3', 'Hexene Cation Frz (BE3)', False, -0.39729184) + self.molecular_unrestricted_oneshot_test( + mol, "be1", "Hexene Cation Frz (BE1)", False, -0.39471433 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be2", "Hexene Cation Frz (BE2)", False, -0.39846777 + ) + self.molecular_unrestricted_oneshot_test( + mol, "be3", "Hexene Cation Frz (BE3)", False, -0.39729184 + ) - def molecular_unrestricted_oneshot_test(self, mol, be_type, test_name, frz, exp_result, delta = 1e-4): - mf = scf.UHF(mol); mf.kernel() - fobj = fragpart(frag_type='autogen', be_type=be_type, mol = mol, frozen_core = frz) + def molecular_unrestricted_oneshot_test( + self, mol, be_type, test_name, frz, exp_result, delta=1e-4 + ): + mf = scf.UHF(mol) + mf.kernel() + fobj = fragpart(frag_type="autogen", be_type=be_type, mol=mol, frozen_core=frz) mybe = UBE(mf, fobj) mybe.oneshot(solver="UCCSD", nproc=1, calc_frag_energy=True, clean_eri=True) - self.assertAlmostEqual(mybe.ebe_tot - mybe.uhf_full_e, exp_result, - msg = "Unrestricted One-Shot Energy for " + test_name - + " is incorrect by" + str(mybe.ebe_tot - mybe.uhf_full_e - exp_result), delta = delta) + self.assertAlmostEqual( + mybe.ebe_tot - mybe.uhf_full_e, + exp_result, + msg="Unrestricted One-Shot Energy for " + + test_name + + " is incorrect by" + + str(mybe.ebe_tot - mybe.uhf_full_e - exp_result), + delta=delta, + ) + -if __name__ == '__main__': +if __name__ == "__main__": unittest.main()