diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..3b17260 --- /dev/null +++ b/.gitignore @@ -0,0 +1,57 @@ +# Byte-compiled / optimized / DLL files +__pycache__/ +evcont/__pycache__/ +*.py[cod] + +# Run examples +*.png +*.npy +*.out +*.sh + +# C extensions +*.so + +# Distribution / packaging +bin/ +build/ +develop-eggs/ +dist/ +eggs/ +lib/ +lib64/ +parts/ +sdist/ +var/ +*.egg-info/ +.installed.cfg +*.egg + +# Installer logs +pip-log.txt +pip-delete-this-directory.txt + +# Unit test / coverage reports +.tox/ +.coverage +.cache +nosetests.xml +coverage.xml + +# Translations +*.mo + +# Mr Developer +.mr.developer.cfg +.project +.pydevproject + +# Rope +.ropeproject + +# Django stuff: +*.log +*.pot + +# Sphinx documentation +docs/_build/ diff --git a/README.md b/README.md index 064c4d9..5c4472e 100644 --- a/README.md +++ b/README.md @@ -1,7 +1,10 @@ # EVCont -This package bundles different scripts and tools for the application of the eigenvector continuation from few variational *ab initio* states, as presented in: Rath, Y., & Booth, G. H. (2025). Interpolating numerically exact many-body wave functions for accelerated molecular dynamics. *Nature Communications*, 16(1), 2005. [1] - +This package bundles different scripts and tools for the application of the eigenvector continuation from few variational *ab initio* states, as presented in: +Rath, Y., & Booth, G. H. (2025) [1] +& Atalar, K., Rath, Y., Crespo-Otero, R. & Booth, G. H. (2024) [2] +The codebase also includes low-rank compression protocols for two-body reduced density matrices and functionality to infer from compressed states, as described in: +Atalar, K., Burton, H. G. A., Grüneis, A. & Booth, G. H. (2026) [3] ## Installation The project comes with a pyproject.toml. @@ -18,14 +21,27 @@ Additional optional dependencies are not installed automatically and need to be Optional dependencies include: - [block2](https://github.com/block-hczhai/block2-preview): required for the continuation from MPS - [pygnme](https://github.com/BoothGroup/pygnme/blob/master/README.md?plain=1): required for the continuation from CAS states -- [dscribe](https://github.com/SINGROUP/dscribe): Required for GAP predictions +- [quantel](https://github.com/hgaburton/quantel): for alternative state-specific CASSCF solutions +- [dscribe](https://github.com/SINGROUP/dscribe): Required for GAP predictions - as used in comparison scripts in [scripts](./scripts) + +The codebase is interfaced to: +- [Newton-X](https://newtonx.org): for nonadiabatic molecular dynamics simulations + +which require separate installations. + +The [nx-interface](./nx-interface/) folder contains the driver script used to call this codebase within Newton-X workflows. + ## Code -This repository is merely a collection of utility functions and scripts for the eigenvector continuation (in particular, interfacing different codes). -The [evcont](./evcont) folder contains common helper functions and bundles the main utilities (which can be installed). +This repository is a collection of utility functions and scripts for the eigenvector continuation (in particular, interfacing different codes). + +The [evcont](./evcont) folder contains various classes for training with different wavefunction solvers, common helper functions and bundles the main utilities (which can be installed). + +The [examples](./examples) folder contains example scripts for demonstrating the general functionality of the codebase. + The [scripts](./scripts) folder contains scripts generate the data from our manuscript [1]. -These should also serve as a good first point of entry into the general functionality of the codebase. +They also serve as example workflows for the codebase. The following scripts are included: - [scripts/PES_H_chain/H6_PES/H6_continuation.py](./scripts/PES_H_chain/H6_PES/H6_continuation.py): Prediction of the PES for a 6-atom H chain from different training points as depicted in Fig. (1) of [1] @@ -37,7 +53,11 @@ The following scripts are included: ## Contact -Questions? Feel free to contact us via [yannic.rath@npl.co.uk](mailto:yannic.rath@npl.co.uk). +Questions? Feel free to contact us via [kemal.atalar@kcl.ac.uk](mailto:kemal.atalar@kcl.ac.uk), [yannic.rath@npl.co.uk](mailto:yannic.rath@npl.co.uk) or [george.booth@kcl.ac.uk](mailto:george.booth@kcl.ac.uk). -## Manuscript +## Manuscripts [1] Rath, Y., & Booth, G. H. (2025). Interpolating numerically exact many-body wave functions for accelerated molecular dynamics. *Nature Communications*, 16(1), 2005. + +[2] Atalar, K., Rath, Y., Crespo-Otero, R. & Booth, G. H. (2024). Fast and accurate nonadiabatic molecular dynamics enabled through variational interpolation of correlated electron wavefunctions. *Faraday Discussions*, 254, 542-569. + +[3] Atalar, K., Burton, H. G. A., Grüneis, A. & Booth, G. H. (2026). Low-rank compression of two-electron reduced density matrices. [arXiv:2605.11253](https://arxiv.org/abs/2605.11253) diff --git a/evcont/CASCI_EVCont.py b/evcont/CASCI_EVCont.py index 4ff911a..23df3ec 100644 --- a/evcont/CASCI_EVCont.py +++ b/evcont/CASCI_EVCont.py @@ -1,14 +1,32 @@ import numpy as np +import pickle -from evcont.electron_integral_utils import get_basis +from evcont.electron_integral_utils import get_basis, get_integrals + +from evcont.low_rank_utils import reduce_2rdm, vectorize_lowrank from pygnme import wick, utils -from pyscf.mcscf.casci import CASCI +#from pyscf.mcscf.casci import CASCI +from pyscf import scf, mcscf, gto from mpi4py import MPI from tqdm import tqdm +import sys +import os, re + +########################################################################### +# Load Quantel if available +try: + from quantel.ints.pyscf_integrals import PySCFMolecule, PySCFIntegrals + from quantel.wfn.ss_casscf import SS_CASSCF + from quantel.opt.mode_controlling import ModeControl + + QUANTEL_FOUND = True +except: + QUANTEL_FOUND = False +########################################################################### rank = MPI.COMM_WORLD.Get_rank() @@ -96,266 +114,1843 @@ class CAS_EVCont_obj: CAS_EVCont_obj holds the data structure for the continuation from CAS states. """ - def __init__(self, ncas, neleca, casci_solver=CASCI): + def __init__(self, ncas, neleca, + nroots=1, solver='SS-CASSCF', + software='pyscf', quantel_path=None, solutions_to_reconverge=None, + lowrank=False, + **kwargs): """ Initialize the CAS_EVCont_obj. Args: ncas (int): Number of CAS orbitals. neleca (int): Number of active space electrons. - casci_solver (object): CASCI solver object from PySCF (can also be CASSCF). + nroots (int): Number of states to be continued. + solver (object): CAS solver type. Options: CASCI, SA-CASSCF, SS-CASSCF. + software (str): Software backend to use ('pyscf' or 'quantel'). + quantel_path (str): Path to Quantel solution to continue from. + solutions_to_reconverge (list): List of solutions to reconverge in the Quantel path. + lowrank (bool): Whether to use low-rank approximation for 2-body t-RDMs. + **kwargs: Additional keyword arguments for low-rank settings. Attributes: ncas (int): Number of CAS orbitals. neleca (int): Number of alpha electrons. - cascis (list): List to store CASCI objects. + overlap (ndarray): Overlap matrix. one_rdm (ndarray): One-electron t-RDM. - two_rdm (ndarray): Two-electron t-RDM. - casci_solver (object): CASCI solver object. + two_rdm (ndarray): Two-electron t-RDM. (if not using low-rank) + vecs_lowrank (dict): Low-rank decomposition of 2-body t-RDMs. (if using low-rank) + + mols (list): List of molecule objects for each state. + mo_coeffs (list): List of MO coefficient matrices for each state. + cis (list): List of CI vectors for each state. + trafos (list): List of transformation matrices for each state. + + """ self.ncas = ncas self.neleca = neleca - self.cascis = [] self.overlap = None self.one_rdm = None self.two_rdm = None - self.casci_solver = casci_solver + # OBSOLETE: Keeping for the old routines, new routines use mo_coeffs and cis + self.cascis = [] - def append_to_rdms(self, mol): + self.mols = [] + self.mo_coeffs = [] + self.cis = [] + self.trafos = [] + + #self.casci_solver = casci_solver + self.nroots = nroots + + # Checks and sets solver/software related attributes + self._input_checks(solver, software, nroots, quantel_path, solutions_to_reconverge) + + # Use each determinant as a separate state + # EXPERIMENTAL: will turn into an input in the future + self.uncontracted = False + + # Internal: Set flags for using add_state vs append_to_rdms + # (to prevent double addition into self.cascis or missing states in tRDMs) + self.use_rdm = None + + ### Initialize low-rank attributes + ### Initialize low-rank attributes + self.lowrank = lowrank + if lowrank: + #self.truncation_style = kwargs['truncation_style'] + self.kwargs = kwargs + + # Diagonals of 2-cumulants ([nbra, nket, 3, norb, norb]) + self.diagonal_lr = None + # Low rank eigendecomposition of the rest of 2-rdm + # Old version: dictionary[(nbra, nket)] = (vals_trunc, vecs_trunc) + # New version: dictionary['vals': np.array([nbra, nket, nvec]), + # 'vecs': np.array([nbra, nket, nvec, nao, nao])] + self.vecs_lowrank = {} + + # Precomputation for OTF Hamiltonian + self.precompute = False + self.inv_OAO_all = [] + self.mb_all = None + self.occ_strings_all = [] + + def _input_checks(self, solver, software, nroots, quantel_path, solutions_to_reconverge): + if solver in ['CASCI','SS-CASSCF','SA-CASSCF', 'casci','ss-casscf','sa-casscf']: + self.solver = solver + elif solver in ['CASSCF', 'casscf']: + if nroots == 1: + self.solver = 'SS-CASSCF' + else: + print('Warning: Solver should specificy state-averaged vs state-specific. Defaulting to state-averaged solver.') + self.solver = 'SA-CASSCF' + else: + print(f'Unknown solver "{solver}" in CAS_EVCont_obj') + sys.exit() + + # Check for the software + if software in ['pyscf']: + self.software = software + elif software in ['quantel']: + if QUANTEL_FOUND: + if solver in ['SS-CASSCF']: + self.software = software + self.quantel_path = quantel_path + self.solutions_to_reconverge = solutions_to_reconverge + else: + print('Unsupported solver for Quantel backend.') + sys.exit() + else: + print('Quantel package not found. Install Quantel or use pyscf as software backend.') + sys.exit() + + def vectorize_lowrank(self,hermitian=True): + vectorize_lowrank(self,hermitian=hermitian) + + def append_to_rdms(self, mol, state=None, quantel_tag='ref', debug=False): """ Append a new training geometry. See pygnme examples for more information about the evaluation of the t-RDMs. Args: mol (object): Molecular object of the training geometry. + state (list, optional): List of precomputed states to be added. If None, new states will be computed. + quantel_tag (str, optional): Tag for quantel states if using quantel software. Default is 'ref' folder in self.quantel_path + debug (bool, optional): If True, print debug information. Defaults to False. Raises: AssertionError: If the mean-field calculation is not converged. """ - overlap = self.overlap - one_rdm = self.one_rdm - two_rdm = self.two_rdm + # Some checks + if self.use_rdm is None: + use_rdm = True + elif not self.use_rdm: + print('Error in append_to_rdms: already using add_state') + sys.exit() + + lowrank = self.lowrank + + ## Preliminaries before state iterations + # AO-SAO transformation + ovlp_bra = mol.intor_symmetric("int1e_ovlp") + basis_OAO_bra = get_basis(mol) + + if self.software == 'pyscf' and state is None: + # Run mean field calculations for the orbitals + #mf = mol.copy().RHF() + mf = scf.RHF(mol.copy()) + #mf.level_shift = 0.5 + #mf.damp = 0.2 + #mf.diis_space = 12 + mf.kernel() + + assert mf.converged + + #MPI.COMM_WORLD.Bcast(mf.mo_coeff) + + if self.solver == 'SA-CASSCF': + cas_sa = mcscf.CASSCF(mf, self.ncas, self.neleca) + if self.nroots > 1: + cas_sa = cas_sa.state_average_([1/self.nroots]*self.nroots) + cas_sa.kernel() + #mo_sacasscf = cas_sa.mo_coeff + assert cas_sa.converged + + elif self.solver == 'CASCI': + mc_casci = mcscf.CASCI(mf, self.ncas, self.neleca) + mc_casci.fcisolver.nroots = self.nroots + mc_casci.kernel() + + assert mc_casci.converged + + elif self.software == 'quantel' and state is None: + # Quantel molecule object + mol_q = PySCFMolecule(mol.atom, mol.basis, mol.unit) + ints = PySCFIntegrals(mol_q) + #metric = ints.overlap_matrix() + #hcore = ints.oei_matrix() + + def convert_to_mcscf(mol,wfn, ncas, neleca): + mc = mcscf.CASCI(mol, ncas, neleca) + mc.fcisolver.max_cycle = 1 + mc.casci(wfn.mo_coeff,ci0=wfn.mat_ci[:,0]) + return mc + + # Save the new tag path + # Create new geometry directory + new_tag = create_next_geom_dir(self.quantel_path) + + # Save the geometry and integrals + mol_q.tofile(os.path.join(self.quantel_path, new_tag, 'molecule.xyz')) + + h1_q = ints.oei_ao_to_mo(basis_OAO_bra, basis_OAO_bra) + h2_q = np.einsum('pi,qj,pqrs,rk,sl->ijkl', basis_OAO_bra, basis_OAO_bra, ints.tei_array(), basis_OAO_bra, basis_OAO_bra,optimize=True) + + np.savetxt(os.path.join(self.quantel_path, new_tag, 'oei.dat'), h1_q) + np.save(os.path.join(self.quantel_path, new_tag, 'tei.npy'), h2_q) + + # Iterate over different states + if state is None: + if self.software == 'quantel': + nroots = len(self.solutions_to_reconverge) + else: + nroots = self.nroots + else: + nroots = len(state) + + for istate in range(nroots): + + # Read the DM representation from existing training states + overlap = self.overlap + one_rdm = self.one_rdm + if not lowrank: + two_rdm = self.two_rdm + else: + diagonal_lr = self.diagonal_lr + vecs_lowrank = self.vecs_lowrank + + if state is None and self.software == 'pyscf': + if self.solver == 'CASCI': + #casci_bra = mcscf.CASCI(mf, self.ncas, self.neleca).state_specific_(istate) + mo_coeff_bra = mc_casci.mo_coeff + mol_bra = mc_casci.mol + + if self.nroots > 1: + ci_bra = mc_casci.ci[istate] + e = mc_casci.e_tot[istate] + else: + ci_bra = mc_casci.ci + e = mc_casci.e_tot + + ncas = mc_casci.ncas + ncore = mc_casci.ncore + + elif self.solver == 'SA-CASSCF': + # casci_bra = mcscf.CASCI(mf, self.ncas, self.neleca).state_specific_(istate) + # casci_bra.casci(mo_sacasscf) + mo_coeff_bra = cas_sa.mo_coeff + mol_bra = cas_sa.mol + if self.nroots > 1: + ci_bra = cas_sa.ci[istate] + e = cas_sa.e_states[istate] + else: + ci_bra = cas_sa.ci + e = cas_sa.e_tot + + ncas = cas_sa.ncas + ncore = cas_sa.ncore + + elif self.solver == 'SS-CASSCF': + cas_ss = mcscf.CASSCF(mf, self.ncas, self.neleca).state_specific_(istate) + cas_ss.kernel() + #casci_bra = mcscf.CASCI(mf, self.ncas, self.neleca).state_specific_(istate) + #casci_bra.casci(cas_ss.mo_coeff) + + mo_coeff_bra = cas_ss.mo_coeff + mol_bra = cas_ss.mol + ci_bra = cas_ss.ci + + e = cas_ss.e_tot + + assert cas_ss.converged + + ncas = cas_ss.ncas + ncore = cas_ss.ncore + + #else: + + + elif self.software == 'quantel' and state is None: + # Quantel molecule object + state_ind = self.solutions_to_reconverge[istate]+1 + wfn = SS_CASSCF(ints, (self.ncas,self.neleca)) + wfn.initialise(np.genfromtxt(os.path.join(self.quantel_path, quantel_tag, f'{state_ind:04d}.mo_coeff')), + np.genfromtxt(os.path.join(self.quantel_path, quantel_tag, f'{state_ind:04d}.mat_ci'))) + + # Reconverge solution at the new geometry + ModeControl().run(wfn) + + # Save the reconverged wavefunction + state_path = os.path.join(self.quantel_path, new_tag) + np.savetxt(os.path.join(state_path,f'{state_ind:04d}.mo_coeff'), wfn.mo_coeff) + np.savetxt(os.path.join(state_path, f'{state_ind:04d}.mat_ci'), wfn.mat_ci) + + casci_bra = convert_to_mcscf(mol,wfn, self.ncas, self.neleca) + mo_coeff_bra = casci_bra.mo_coeff + mol_bra = casci_bra.mol + ci_bra = casci_bra.ci + + #e = wfn.energy + out = casci_bra.kernel() + e = out[0] + + ncas = casci_bra.ncas + ncore = casci_bra.ncore + nelec = mol_bra.nelec + + else: + casci_bra = state[istate] + + mo_coeff_bra = casci_bra.mo_coeff + mol_bra = casci_bra.mol + ci_bra = casci_bra.ci + + out = casci_bra.kernel() + e = out[0] + + assert np.all(casci_bra.fcisolver.converged) + + if hasattr(casci_bra, "converged"): + assert casci_bra.converged + + ncas = casci_bra.ncas + ncore = casci_bra.ncore + nelec = mol_bra.nelec + + nelec = mol_bra.nelec + + # Check if this new state is already stored (mo_coeff_bra and ci_bra are within a threshold of any element of self.mo_coeffs and self.cis) + state_exists = False + for i in range(len(self.cis)): + mo_coeff_existing = self.mo_coeffs[i] + ci_existing = self.cis[i] + + mo_diff = np.linalg.norm(mo_coeff_existing - mo_coeff_bra) + ci_diff = np.linalg.norm(ci_existing - ci_bra) + + if mo_diff < 1e-6 and ci_diff < 1e-6: + print('Warning: The appended state is already stored in the training set. Skipping addition.') + state_exists = True + break + + if state_exists: + continue + + # New version: store MO coeffs and CI vectors separately + self.mo_coeffs.append(mo_coeff_bra) + self.cis.append(ci_bra) + self.mols.append(mol_bra) + + trafo_bra = basis_OAO_bra.T.dot(ovlp_bra).dot(mo_coeff_bra) + + self.trafos.append(trafo_bra) + + mo_coeffs = self.mo_coeffs + cis = self.cis + mols = self.mols + trafos = self.trafos + n_cascis = len(cis) + + MPI.COMM_WORLD.Bcast(ci_bra) + MPI.COMM_WORLD.Bcast(mo_coeff_bra) + + bra_ref_state = wick.reference_state[float]( + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mol_bra.nelec[0], + ncas, + ncore, + owndata(mo_coeff_bra), + ) - mf = mol.copy().RHF() - mf.kernel() + if rank == 0: + overlap_new = np.zeros((n_cascis, n_cascis)) + if overlap is not None: + overlap_new[:-1, :-1] = overlap + one_rdm_new = np.zeros( + (n_cascis, n_cascis, mo_coeff_bra.shape[0], mo_coeff_bra.shape[0]) + ) + if one_rdm is not None: + one_rdm_new[:-1, :-1, :, :] = one_rdm + + # Only define two_rdm if not lowrank + if not lowrank: + two_rdm_new = np.zeros( + ( + n_cascis, + n_cascis, + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + ) + ) + if two_rdm is not None: + two_rdm_new[:-1, :-1, :, :, :, :] = two_rdm + + else: + diagonal_lr_new = np.ones( + (n_cascis, + n_cascis, 3, mo_coeff_bra.shape[0], mo_coeff_bra.shape[0]) + ) + if diagonal_lr is not None: + diagonal_lr_new[:-1, :-1, :, :, :] = diagonal_lr + + else: + overlap_new = one_rdm_new = two_rdm_new = None + + bra_occ_strings = utils.fci_bitset_list( + mol_bra.nelec[0] - ncore, ncas + ) - assert mf.converged + for i in range(n_cascis): + mo_coeff_ket = mo_coeffs[i] + ci_ket = cis[i] - MPI.COMM_WORLD.Bcast(mf.mo_coeff) - casci_bra = self.casci_solver(mf, self.ncas, self.neleca) + trafo_ket = trafos[i] - self.cascis.append(casci_bra) + trafo_ket_bra = basis_OAO_bra.dot(trafo_ket) - cascis = self.cascis - n_cascis = len(cascis) + ket_ref_state = wick.reference_state[float]( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + nelec[0], + ncas, + ncore, + owndata(trafo_ket_bra), + ) - casci_bra.kernel() + orbitals = wick.wick_orbitals[float, float]( + bra_ref_state, ket_ref_state, owndata(ovlp_bra) + ) - assert casci_bra.fcisolver.converged + wick_mb = wick.wick_rscf[float, float, float](orbitals, 0.0) - if hasattr(casci_bra, "converged"): - assert casci_bra.converged + ket_occ_strings = utils.fci_bitset_list( + nelec[0] - ncore, ncas + ) - MPI.COMM_WORLD.Bcast(casci_bra.ci) - MPI.COMM_WORLD.Bcast(casci_bra.mo_coeff) + rdm1_tmp = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) + rdm1 = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) + rdm2_tmp = np.zeros( + ( + mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + ) + ) + rdm2 = np.zeros( + ( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + ) + ) + overlap_accumulate = 0.0 + + all_ids = np.array( + [ + [iabra, ibbra, iaket, ibket] + for iabra in range(len(bra_occ_strings)) + for ibbra in range(len(bra_occ_strings)) + for iaket in range(len(ket_occ_strings)) + for ibket in range(len(ket_occ_strings)) + ] + ) - mo_coeff_bra = casci_bra.mo_coeff - mol_bra = casci_bra.mol + n_ranks = MPI.COMM_WORLD.Get_size() - ovlp_bra = mol_bra.intor_symmetric("int1e_ovlp") - basis_OAO_bra = get_basis(mol_bra) - trafo_bra = basis_OAO_bra.T.dot(ovlp_bra).dot(mo_coeff_bra) + all_ids_local = np.array_split(all_ids, n_ranks)[rank] - bra_ref_state = wick.reference_state[float]( - mo_coeff_bra.shape[0], - mo_coeff_bra.shape[0], - mol_bra.nelec[0], - casci_bra.ncas, - casci_bra.ncore, - owndata(mo_coeff_bra), - ) + if rank == 0: + pbar = tqdm(total=len(all_ids_local)) + + for ids in all_ids_local: + iabra, ibbra, iaket, ibket = ids + stringabra = bra_occ_strings[iabra] + stringbbra = bra_occ_strings[ibbra] + stringaket = ket_occ_strings[iaket] + stringbket = ket_occ_strings[ibket] + + rdm1_tmp.fill(0.0) + rdm2_tmp.fill(0.0) + o = wick_mb.evaluate_rdm12( + stringabra, + stringbbra, + stringaket, + stringbket, + 1.0, + rdm1_tmp, + rdm2_tmp, + ) + overlap_accumulate += ( + o * ci_bra[iabra, ibbra] * ci_ket[iaket, ibket] + ) + + rdm1 += ( + rdm1_tmp * ci_bra[iabra, ibbra] * ci_ket[iaket, ibket] + ) + rdm2 += ( + rdm2_tmp.reshape(rdm2.shape) + * ci_bra[iabra, ibbra] + * ci_ket[iaket, ibket] + ) + + if rank == 0: + pbar.update(1) - if rank == 0: - overlap_new = np.zeros((n_cascis, n_cascis)) - if overlap is not None: - overlap_new[:-1, :-1] = overlap - one_rdm_new = np.zeros( - (n_cascis, n_cascis, mo_coeff_bra.shape[0], mo_coeff_bra.shape[0]) - ) - if one_rdm is not None: - one_rdm_new[:-1, :-1, :, :] = one_rdm - two_rdm_new = np.zeros( - ( - n_cascis, - n_cascis, - mo_coeff_bra.shape[0], - mo_coeff_bra.shape[0], + if rank == 0: + pbar.close() + + overlap_accumulate = MPI.COMM_WORLD.allreduce( + overlap_accumulate, op=MPI.SUM + ) + + MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm1, op=MPI.SUM) + MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm2, op=MPI.SUM) + + if rank == 0: + + overlap_new[-1, i] = overlap_accumulate + overlap_new[i, -1] = overlap_accumulate.conj() + rdm1 = np.einsum( + "...ij,ai,bj->...ab", rdm1, trafo_ket, trafo_bra, optimize="optimal" + ) + rdm2 = np.einsum( + "...ijkl,ai,bj,ck,dl->...abcd", + rdm2, + trafo_bra, + trafo_ket, + trafo_bra, + trafo_ket, + optimize="optimal", + ) + + if debug: + np.save('rdm2_%i_%i.npy'%(n_cascis-1, i),rdm2) + + one_rdm_new[-1, i, :, :] = rdm1 + one_rdm_new[i, -1, :, :] = rdm1.conj().T + + if not lowrank: + two_rdm_new[-1, i, :, :, :, :] = rdm2 + two_rdm_new[i, -1, :, :, :, :] = np.einsum('ijkl->lkji',rdm2.conj()) + + # Low rank + else: + # Get low rank representation + print('States: %i %i, overlap: %f' % (n_cascis-1, i, overlap_accumulate)) + + lowrank_vecs, diagonals, use_joint = \ + reduce_2rdm(rdm1, rdm2, overlap_accumulate, + mol=mol, train_en=e, + **self.kwargs) + + diagonal_lr_new[-1, i, :, :, :] = diagonals + try: + # This gives an error if diagonals are not saved and set to None by reduce_2rdm + diagonal_lr_new[i, -1, :, :, :] = diagonals.conj() + except: + diagonal_lr_new[i, -1, :, :, :] = diagonals + + #diagonal_lr_new[i, -1, :, :, :] = diagonals_conj + + # Data structure [(bra,ket)]; eval, leftvec, rightvec, use_joint + vecs_lowrank[(n_cascis-1, i)] = lowrank_vecs[0], lowrank_vecs[1], lowrank_vecs[2], use_joint + vecs_lowrank[(i,n_cascis-1)] = lowrank_vecs[0].conj(), lowrank_vecs[1].conj(), lowrank_vecs[2].conj(), use_joint + + + self.overlap = overlap_new + self.one_rdm = one_rdm_new + if not lowrank: + self.two_rdm = two_rdm_new + else: + self.diagonal_lr = diagonal_lr_new + self.vecs_lowrank = vecs_lowrank + + # Experimental: Uncontracted CAS continuation. Can be combined with append_to_rdms once tested + def append_to_rdms_separate_determinants(self, mol, state=None, debug=False): + """ + Append a new training geometry with each determinant as a separate state. + Modified version that creates separate states for each determinant instead of summing them. + + Args: + mol (object): Molecular object of the training geometry. + + Raises: + AssertionError: If the mean-field calculation is not converged. + """ + # Some checks + if self.use_rdm is None: + use_rdm = True + elif not self.use_rdm: + print('Error in append_to_rdms_separate_determinants: already using add_state') + sys.exit() + + lowrank = self.lowrank + + # Run mean field calculations for the orbitals + mf = scf.RHF(mol.copy()) + mf.kernel() + + assert mf.converged + + MPI.COMM_WORLD.Bcast(mf.mo_coeff) + + if state is None: + if self.solver == 'SA-CASSCF': + cas_sa = mcscf.CASSCF(mf, self.ncas, self.neleca).state_average_([1/self.nroots]*self.nroots) + cas_sa.kernel() + assert cas_sa.converged + + elif self.solver == 'CASCI': + mc_casci = mcscf.CASCI(mf, self.ncas, self.neleca) + mc_casci.fcisolver.nroots = self.nroots + mc_casci.kernel() + assert mc_casci.converged + + # Iterate over different states + if state is None: + nroots = self.nroots + else: + nroots = len(state) + + for istate in range(nroots): + # Read the DM representation from existing training states + overlap = self.overlap + one_rdm = self.one_rdm + if not lowrank: + two_rdm = self.two_rdm + else: + diagonal_lr = self.diagonal_lr + vecs_lowrank = self.vecs_lowrank + + if state is None: + if self.solver == 'CASCI': + mo_coeff_bra = mc_casci.mo_coeff + mol_bra = mc_casci.mol + ci_bra = mc_casci.ci[istate] + e = mc_casci.e_tot[istate] + ncas = mc_casci.ncas + ncore = mc_casci.ncore + + elif self.solver == 'SA-CASSCF': + mo_coeff_bra = cas_sa.mo_coeff + mol_bra = cas_sa.mol + ci_bra = cas_sa.ci[istate] + e = cas_sa.e_states[istate] + ncas = cas_sa.ncas + ncore = cas_sa.ncore + + elif self.solver == 'SS-CASSCF': + cas_ss = mcscf.CASSCF(mf, self.ncas, self.neleca).state_specific_(istate) + cas_ss.kernel() + mo_coeff_bra = cas_ss.mo_coeff + mol_bra = cas_ss.mol + ci_bra = cas_ss.ci + e = cas_ss.e_tot + assert cas_ss.converged + ncas = cas_ss.ncas + ncore = cas_ss.ncore + else: + casci_bra = state[istate] + mo_coeff_bra = casci_bra.mo_coeff + mol_bra = casci_bra.mol + ci_bra = casci_bra.ci + out = casci_bra.kernel() + e = out[0] + assert np.all(casci_bra.fcisolver.converged) + if hasattr(casci_bra, "converged"): + assert casci_bra.converged + ncas = casci_bra.ncas + ncore = casci_bra.ncore + + nelec = mol_bra.nelec + + # Get determinant strings for this state + bra_occ_strings = utils.fci_bitset_list(mol_bra.nelec[0] - ncore, ncas) + + + # Efficiently handle the single-determinant case: only process the nonzero element + nz = np.argwhere(np.abs(ci_bra) > 1e-10) + print(f"State {istate}: Found {len(nz)} significant determinants") + for det_idx, (iabra, ibbra) in enumerate(nz): + ovlp_bra = mol_bra.intor_symmetric("int1e_ovlp") + basis_OAO_bra = get_basis(mol_bra) + trafo_bra = basis_OAO_bra.T.dot(ovlp_bra).dot(mo_coeff_bra) + + self.mo_coeffs.append(mo_coeff_bra) + ci_single_det = np.zeros_like(ci_bra) + ci_single_det[iabra, ibbra] = 1.0 + self.cis.append(ci_single_det) + self.mols.append(mol_bra) + self.trafos.append(trafo_bra) + + mo_coeffs = self.mo_coeffs + cis = self.cis + mols = self.mols + trafos = self.trafos + n_cascis = len(cis) + + MPI.COMM_WORLD.Bcast(ci_single_det) + MPI.COMM_WORLD.Bcast(mo_coeff_bra) + + bra_ref_state = wick.reference_state[float]( mo_coeff_bra.shape[0], mo_coeff_bra.shape[0], + mol_bra.nelec[0], + ncas, + ncore, + owndata(mo_coeff_bra), ) + + if rank == 0: + overlap_new = np.zeros((n_cascis, n_cascis)) + if overlap is not None: + overlap_new[:-1, :-1] = overlap + one_rdm_new = np.zeros( + (n_cascis, n_cascis, mo_coeff_bra.shape[0], mo_coeff_bra.shape[0]) + ) + if one_rdm is not None: + one_rdm_new[:-1, :-1, :, :] = one_rdm + + if not lowrank: + two_rdm_new = np.zeros( + ( + n_cascis, + n_cascis, + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + ) + ) + if two_rdm is not None: + two_rdm_new[:-1, :-1, :, :, :, :] = two_rdm + else: + diagonal_lr_new = np.ones( + (n_cascis, n_cascis, 3, mo_coeff_bra.shape[0], mo_coeff_bra.shape[0]) + ) + if diagonal_lr is not None: + diagonal_lr_new[:-1, :-1, :, :, :] = diagonal_lr + else: + overlap_new = one_rdm_new = two_rdm_new = None + + # Only need to process the single nonzero determinant for bra + iabra_bra, ibbra_bra = iabra, ibbra + for i in range(n_cascis): + mo_coeff_ket = mo_coeffs[i] + ci_ket = cis[i] + trafo_ket = trafos[i] + + trafo_ket_bra = basis_OAO_bra.dot(trafo_ket) + + ket_ref_state = wick.reference_state[float]( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + nelec[0], + ncas, + ncore, + owndata(trafo_ket_bra), + ) + + orbitals = wick.wick_orbitals[float, float]( + bra_ref_state, ket_ref_state, owndata(ovlp_bra) + ) + + wick_mb = wick.wick_rscf[float, float, float](orbitals, 0.0) + + ket_occ_strings = utils.fci_bitset_list(nelec[0] - ncore, ncas) + + rdm1_tmp = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) + rdm1 = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) + rdm2_tmp = np.zeros( + ( + mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + ) + ) + rdm2 = np.zeros( + ( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + ) + ) + overlap_accumulate = 0.0 + + # Only process the nonzero element in ci_ket + nz_ket = np.argwhere(np.abs(ci_ket) > 1e-10) + for iabra_ket, ibbra_ket in nz_ket: + stringabra = bra_occ_strings[iabra_bra] + stringbbra = bra_occ_strings[ibbra_bra] + stringaket = ket_occ_strings[iabra_ket] + stringbket = ket_occ_strings[ibbra_ket] + + rdm1_tmp.fill(0.0) + rdm2_tmp.fill(0.0) + o = wick_mb.evaluate_rdm12( + stringabra, + stringbbra, + stringaket, + stringbket, + 1.0, + rdm1_tmp, + rdm2_tmp, + ) + overlap_accumulate += o * ci_ket[iabra_ket, ibbra_ket] + rdm1 += rdm1_tmp * ci_ket[iabra_ket, ibbra_ket] + rdm2 += rdm2_tmp.reshape(rdm2.shape) * ci_ket[iabra_ket, ibbra_ket] + + overlap_accumulate = MPI.COMM_WORLD.allreduce(overlap_accumulate, op=MPI.SUM) + MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm1, op=MPI.SUM) + MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm2, op=MPI.SUM) + + if rank == 0: + overlap_new[-1, i] = overlap_accumulate + overlap_new[i, -1] = overlap_accumulate.conj() + rdm1 = np.einsum( + "...ij,ai,bj->...ab", rdm1, trafo_ket, trafo_bra, optimize="optimal" + ) + rdm2 = np.einsum( + "...ijkl,ai,bj,ck,dl->...abcd", + rdm2, + trafo_bra, + trafo_ket, + trafo_bra, + trafo_ket, + optimize="optimal", + ) + + if debug: + np.save('rdm2_det_%i_%i_%i.npy'%(istate, det_idx, i), rdm2) + + one_rdm_new[-1, i, :, :] = rdm1 + one_rdm_new[i, -1, :, :] = rdm1.conj().T + + if not lowrank: + two_rdm_new[-1, i, :, :, :, :] = rdm2 + two_rdm_new[i, -1, :, :, :, :] = np.einsum('ijkl->klij', rdm2.conj()) + else: + print(f"Determinant {det_idx} of state {istate}, overlap with state {i}: {overlap_accumulate}") + lowrank_vecs, diagonals, use_joint = reduce_2rdm( + rdm1, rdm2, overlap_accumulate, + mol=mol, train_en=e, + **self.kwargs + ) + + diagonal_lr_new[-1, i, :, :, :] = diagonals + try: + diagonal_lr_new[i, -1, :, :, :] = diagonals.conj() + except: + diagonal_lr_new[i, -1, :, :, :] = diagonals + + vecs_lowrank[(n_cascis-1, i)] = lowrank_vecs[0], lowrank_vecs[1], lowrank_vecs[2], use_joint + vecs_lowrank[(i, n_cascis-1)] = lowrank_vecs[0].conj(), lowrank_vecs[1].conj(), lowrank_vecs[2].conj(), use_joint + + self.overlap = overlap_new + self.one_rdm = one_rdm_new + if not lowrank: + self.two_rdm = two_rdm_new + else: + if getattr(self, 'kwargs', {}).get('save_diag', True): + self.diagonal_lr = diagonal_lr_new + else: + self.diagonal_lr = None + self.vecs_lowrank = vecs_lowrank + + overlap = overlap_new + one_rdm = one_rdm_new + if not lowrank: + two_rdm = two_rdm_new + else: + diagonal_lr = diagonal_lr_new + + def otf_hamiltonian(self, h1, h2): + """ + Generate subspace Hamiltonian on the fly from precomputed training states (self.cascis) + Note: Still need to test if the MPI version works + + Args: + h1 (np.array): 1-electron integrals at the test geometry in SAO basis. + h2 (np.array): 2-electron integrals at the test geometry in SAO basis. + """ + states = self.cascis + + nwf = len(states) + H = np.zeros([nwf,nwf]) + S = np.zeros([nwf,nwf]) + + #time_pre_bra = 0. + #time_pre_ket_worb = 0. + #n_bra_pre = 0 + #n_ket_pre = 0 + + #st = time() + + # Iterate over bra states + for a, casci_bra in enumerate(states): + + MPI.COMM_WORLD.Bcast(casci_bra.ci) + MPI.COMM_WORLD.Bcast(casci_bra.mo_coeff) + + #st_bra = time() + + mo_coeff_bra = casci_bra.mo_coeff + mol_bra = casci_bra.mol + + ovlp_bra = mol_bra.intor_symmetric("int1e_ovlp") + basis_OAO_bra = get_basis(mol_bra) + trafo_bra = basis_OAO_bra.T.dot(ovlp_bra).dot(mo_coeff_bra) + #print('bra',ovlp_bra.shape,basis_OAO_bra.shape,trafo_bra.shape) + + bra_ref_state = wick.reference_state[float]( + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mol_bra.nelec[0], + casci_bra.ncas, + casci_bra.ncore, + owndata(mo_coeff_bra), ) - if two_rdm is not None: - two_rdm_new[:-1, :-1, :, :, :, :] = two_rdm - else: - overlap_new = one_rdm_new = two_rdm_new = None - bra_occ_strings = utils.fci_bitset_list( - mol_bra.nelec[0] - casci_bra.ncore, casci_bra.ncas - ) + bra_occ_strings = utils.fci_bitset_list( + mol_bra.nelec[0] - casci_bra.ncore, casci_bra.ncas + ) - for i in range(n_cascis): - casci_ket = cascis[i] - mo_coeff_ket = casci_ket.mo_coeff - mol_ket = casci_ket.mol + # Generate and transform 1- and 2-electron integrals + # AO(test) to AO(bra) transformation + #basis_test_bra = np.dot(get_basis(mol),np.linalg.inv(basis_OAO_bra)) + #h1e, h2e = get_integrals(mol, basis_test_bra) - ovlp_ket = mol_ket.intor_symmetric("int1e_ovlp") - basis_OAO_ket = get_basis(mol_ket) - trafo_ket = basis_OAO_ket.T.dot(ovlp_ket).dot(mo_coeff_ket) + # Transform 1- and 2-electron integrals into AO basis of bra + inv_basis_OAO_bra = np.linalg.inv(basis_OAO_bra) - trafo_ket_bra = basis_OAO_bra.dot(trafo_ket) + #time_pre_bra += time()-st_bra + #n_bra_pre += 1 - ket_ref_state = wick.reference_state[float]( - mo_coeff_ket.shape[0], - mo_coeff_ket.shape[0], - mol_ket.nelec[0], - casci_ket.ncas, - casci_ket.ncore, - owndata(trafo_ket_bra), + h1e = np.einsum( + "ia,jb,ij->ab", inv_basis_OAO_bra, inv_basis_OAO_bra, h1, optimize="optimal" ) - orbitals = wick.wick_orbitals[float, float]( - bra_ref_state, - ket_ref_state, - owndata(mol_bra.intor_symmetric("int1e_ovlp")), + h2e = np.einsum( + "ia,jb,kc,ld,ijkl->abcd", + inv_basis_OAO_bra, + inv_basis_OAO_bra, + inv_basis_OAO_bra, + inv_basis_OAO_bra, + h2, + optimize="optimal", ) - wick_mb = wick.wick_rscf[float, float, float](orbitals, 0.0) + MPI.COMM_WORLD.Bcast(h1e) + MPI.COMM_WORLD.Bcast(h2e) + + # Iterate over ket states + #for b, casci_ket in enumerate(states): + for b in range(a, nwf): + casci_ket = states[b] + #st_ket = time() + + # Prepare ket state + mo_coeff_ket = casci_ket.mo_coeff + mol_ket = casci_ket.mol + + ovlp_ket = mol_ket.intor_symmetric("int1e_ovlp") + basis_OAO_ket = get_basis(mol_ket) + trafo_ket = basis_OAO_ket.T.dot(ovlp_ket).dot(mo_coeff_ket) + + trafo_ket_bra = basis_OAO_bra.dot(trafo_ket) + #print('bra',ovlp_ket.shape,basis_OAO_ket.shape,trafo_ket.shape, trafo_ket_bra.shape) + + ket_ref_state = wick.reference_state[float]( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mol_ket.nelec[0], + casci_ket.ncas, + casci_ket.ncore, + owndata(trafo_ket_bra), + ) + + ket_occ_strings = utils.fci_bitset_list( + mol_ket.nelec[0] - casci_ket.ncore, casci_ket.ncas + ) + orbitals = wick.wick_orbitals[float, float]( + bra_ref_state, ket_ref_state, owndata(ovlp_bra) + ) + + mb = wick.wick_rscf[float, float, float](orbitals, 0.0) + + #time_pre_ket_worb += time()-st_ket + #n_ket_pre += 1 + + # Add one- and two-body contributions + h1e = owndata(h1e) + h2e = owndata(h2e.reshape(h1e.shape[0]**2, h1e.shape[0]**2)) + mb.add_one_body(h1e) + mb.add_two_body(h2e) + + overlap_accumulate = 0.0 + hamiltonian_accumulate = 0.0 + + all_ids = np.array( + [ + [iabra, ibbra, iaket, ibket] + for iabra in range(len(bra_occ_strings)) + for ibbra in range(len(bra_occ_strings)) + for iaket in range(len(ket_occ_strings)) + for ibket in range(len(ket_occ_strings)) + ] + ) + + n_ranks = MPI.COMM_WORLD.Get_size() + + all_ids_local = np.array_split(all_ids, n_ranks)[rank] + + if rank == 0: + pbar = tqdm(total=len(all_ids_local)) + + for ids in all_ids_local: + iabra, ibbra, iaket, ibket = ids + + # Compute S and H contribution for this pair of determinants + stmp, htmp = mb.evaluate(bra_occ_strings[iabra], + bra_occ_strings[ibbra], + ket_occ_strings[iaket], + ket_occ_strings[ibket], + 1.0) + + hamiltonian_accumulate += htmp * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + overlap_accumulate += stmp * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + + if rank == 0: + pbar.update(1) + + if rank == 0: + pbar.close() + + overlap_accumulate = MPI.COMM_WORLD.allreduce( + overlap_accumulate, op=MPI.SUM + ) + hamiltonian_accumulate = MPI.COMM_WORLD.allreduce( + hamiltonian_accumulate, op=MPI.SUM + ) + + if rank == 0: + #print(hamiltonian_accumulate, overlap_accumulate) + + H[a,b] = hamiltonian_accumulate + S[a,b] = overlap_accumulate + + H[b,a] = np.conj(hamiltonian_accumulate) + S[b,a] = np.conj(overlap_accumulate) + + #print('----------------------------------------------') + #print('Time per Hamiltonian: %.5f'%(time()-st)) + #print('Time available for precomputation: %.5f'%(time_pre_bra+time_pre_ket_worb)) + #print('Bra preparation time: %.5f, (%.5f per each bra)'%(time_pre_bra, time_pre_bra/n_bra_pre)) + #print('Ket preparation time with orbital object definition: %.5f, (%.5f per each ket)'%(time_pre_ket_worb, time_pre_ket_worb/n_ket_pre)) + #print('----------------------------------------------') + return H, S + + def otf_hamiltonia_precomputed(self, h1, h2): + """ + Generate subspace Hamiltonian on the fly from precomputed training states (self.cascis) + Using precomputed quantities for speedup + Note: Still need to test if the MPI version works + + Args: + h1 (np.array): 1-electron integrals at the test geometry in SAO basis. + h2 (np.array): 2-electron integrals at the test geometry in SAO basis. + """ + + if self.precompute == False: + print('Precomputations were not available. Precomputing now.') + self.precompute_for_otf() + + states = self.cascis + inv_OAO_all = self.inv_OAO_all + occ_strings_all = self.occ_strings_all + #mb_all = self.mb_all + + #print(mb_all[0][0]) + nwf = len(states) + H = np.zeros([nwf,nwf]) + S = np.zeros([nwf,nwf]) + + #mb_all = [[i,j] for i in range(nwf) for j in range(nwf)] + + #time_pre_bra = 0. + #time_pre_ket_worb = 0. + #n_bra_pre = 0 + #n_ket_pre = 0 + + #st = time() + + # Iterate over bra states + for a, casci_bra in enumerate(states): + + MPI.COMM_WORLD.Bcast(casci_bra.ci) + #MPI.COMM_WORLD.Bcast(inv_OAO_all[a]) + #MPI.COMM_WORLD.Bcast(occ_strings_all[a]) + #MPI.COMM_WORLD.Bcast(mb_all[a,:]) - ket_occ_strings = utils.fci_bitset_list( - mol_ket.nelec[0] - casci_ket.ncore, casci_ket.ncas + bra_occ_strings = occ_strings_all[a] + + # Generate and transform 1- and 2-electron integrals + # AO(test) to AO(bra) transformation + #basis_test_bra = np.dot(get_basis(mol),np.linalg.inv(basis_OAO_bra)) + #h1e, h2e = get_integrals(mol, basis_test_bra) + + # Transform 1- and 2-electron integrals into AO basis of bra + inv_basis_OAO_bra = inv_OAO_all[a] + + #time_pre_bra += time()-st_bra + #n_bra_pre += 1 + + h1e = np.einsum( + "ia,jb,ij->ab", inv_basis_OAO_bra, inv_basis_OAO_bra, h1, optimize="optimal" ) - rdm1_tmp = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) - rdm1 = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) - rdm2_tmp = np.zeros( - ( - mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], - mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + h2e = np.einsum( + "ia,jb,kc,ld,ijkl->abcd", + inv_basis_OAO_bra, + inv_basis_OAO_bra, + inv_basis_OAO_bra, + inv_basis_OAO_bra, + h2, + optimize="optimal", + ) + + MPI.COMM_WORLD.Bcast(h1e) + MPI.COMM_WORLD.Bcast(h2e) + + # Iterate over ket states + #for b, casci_ket in enumerate(states): + for b in range(a, nwf): + casci_ket = states[b] + #st_ket = time() + + ket_occ_strings = occ_strings_all[b] + + + #mb = wick.wick_rscf[float, float, float](orbitals, 0.0) + + mb = self.mb_all[a][b] + + #time_pre_ket_worb += time()-st_ket + #n_ket_pre += 1 + + # Add one- and two-body contributions + h1e = owndata(h1e) + h2e = owndata(h2e.reshape(h1e.shape[0]**2, h1e.shape[0]**2)) + mb.add_one_body(h1e) + mb.add_two_body(h2e) + + overlap_accumulate = 0.0 + hamiltonian_accumulate = 0.0 + + + all_ids = np.array( + [ + [iabra, ibbra, iaket, ibket] + for iabra in range(len(bra_occ_strings)) + for ibbra in range(len(bra_occ_strings)) + for iaket in range(len(ket_occ_strings)) + for ibket in range(len(ket_occ_strings)) + ] + ) + + n_ranks = MPI.COMM_WORLD.Get_size() + + all_ids_local = np.array_split(all_ids, n_ranks)[rank] + + if rank == 0: + pbar = tqdm(total=len(all_ids_local)) + + for ids in all_ids_local: + iabra, ibbra, iaket, ibket = ids + + # Compute S and H contribution for this pair of determinants + stmp, htmp = mb.evaluate(bra_occ_strings[iabra], + bra_occ_strings[ibbra], + ket_occ_strings[iaket], + ket_occ_strings[ibket], + 1.0) + + hamiltonian_accumulate += htmp * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + overlap_accumulate += stmp * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + + if rank == 0: + pbar.update(1) + + if rank == 0: + pbar.close() + + overlap_accumulate = MPI.COMM_WORLD.allreduce( + overlap_accumulate, op=MPI.SUM + ) + hamiltonian_accumulate = MPI.COMM_WORLD.allreduce( + hamiltonian_accumulate, op=MPI.SUM ) + + if rank == 0: + #print(hamiltonian_accumulate, overlap_accumulate) + + H[a,b] = hamiltonian_accumulate + S[a,b] = overlap_accumulate + + H[b,a] = np.conj(hamiltonian_accumulate) + S[b,a] = np.conj(overlap_accumulate) + + + #print('----------------------------------------------') + #print('Time per Hamiltonian: %.5f'%(time()-st)) + #print('Time available for precomputation: %.5f'%(time_pre_bra+time_pre_ket_worb)) + #print('Bra preparation time: %.5f, (%.5f per each bra)'%(time_pre_bra, time_pre_bra/n_bra_pre)) + #print('Ket preparation time with orbital object definition: %.5f, (%.5f per each ket)'%(time_pre_ket_worb, time_pre_ket_worb/n_ket_pre)) + #print('----------------------------------------------') + return H, S + + def precompute_for_otf(self): + """ + Precompute and setup the on-the-fly computation of the subspace Hamiltonian + beforehand to save time during test evaluations + """ + + assert len(self.cascis) != 0 + self.precompute = True + + states = self.cascis + + nwf = len(states) + + #self.mb_all = np.ones((nwf,nwf))*np.nan + self.mb_all = [] + self.orbitals_all = [] + + # Iterate over bra states + for a, casci_bra in enumerate(states): + + mb_a = [] + orb_a = [] + MPI.COMM_WORLD.Bcast(casci_bra.ci) + MPI.COMM_WORLD.Bcast(casci_bra.mo_coeff) + + #st_bra = time() + + mo_coeff_bra = casci_bra.mo_coeff + mol_bra = casci_bra.mol + + ovlp_bra = mol_bra.intor_symmetric("int1e_ovlp") + basis_OAO_bra = get_basis(mol_bra) + trafo_bra = basis_OAO_bra.T.dot(ovlp_bra).dot(mo_coeff_bra) + #print('bra',ovlp_bra.shape,basis_OAO_bra.shape,trafo_bra.shape) + + bra_ref_state = wick.reference_state[float]( + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mol_bra.nelec[0], + casci_bra.ncas, + casci_bra.ncore, + owndata(mo_coeff_bra), ) - rdm2 = np.zeros( - ( - mo_coeff_ket.shape[0], - mo_coeff_ket.shape[0], + + bra_occ_strings = utils.fci_bitset_list( + mol_bra.nelec[0] - casci_bra.ncore, casci_bra.ncas + ) + self.occ_strings_all.append(bra_occ_strings) + + # Transform 1- and 2-electron integrals into AO basis of bra + inv_basis_OAO_bra = np.linalg.inv(basis_OAO_bra) + self.inv_OAO_all.append(inv_basis_OAO_bra) + + # Iterate over ket states + #for b, casci_ket in enumerate(states): + #for b in range(a, nwf): + for b in range(nwf): + casci_ket = states[b] + #st_ket = time() + + # Prepare ket state + mo_coeff_ket = casci_ket.mo_coeff + mol_ket = casci_ket.mol + + ovlp_ket = mol_ket.intor_symmetric("int1e_ovlp") + basis_OAO_ket = get_basis(mol_ket) + trafo_ket = basis_OAO_ket.T.dot(ovlp_ket).dot(mo_coeff_ket) + + trafo_ket_bra = basis_OAO_bra.dot(trafo_ket) + #print('bra',ovlp_ket.shape,basis_OAO_ket.shape,trafo_ket.shape, trafo_ket_bra.shape) + + ket_ref_state = wick.reference_state[float]( mo_coeff_ket.shape[0], mo_coeff_ket.shape[0], + mol_ket.nelec[0], + casci_ket.ncas, + casci_ket.ncore, + owndata(trafo_ket_bra), + ) + + ket_occ_strings = utils.fci_bitset_list( + mol_ket.nelec[0] - casci_ket.ncore, casci_ket.ncas + ) + orbitals = wick.wick_orbitals[float, float]( + bra_ref_state, ket_ref_state, owndata(ovlp_bra) ) + orb_a.append(orbitals) + mb = wick.wick_rscf[float, float, float](orbitals, 0.0) + mb_a.append(mb) + + #self.mb_all[a,b] = mb + #self.mb_all[b,a] = mb + + self.mb_all.append(mb_a) + self.orbitals_all.append(orb_a) + #self.mb_all = np.array(self.mb_all) + + return 1 + + def add_state(self, mol): + """ + Compute the wavefunctions and store them in this object for on-the-fly continuation + later on. + ALTERNATIVE to append_to_rdms + + Args: + mol (object): Molecular object of the training geometry. + + Raises: + AssertionError: If the mean-field calculation is not converged. + """ + # Some checks + if self.use_rdm is None: + use_rdm = False + elif self.use_rdm: + print('Error in add_state: already using append_to_rdms') + sys.exit() + + # Run mean field calculations for the orbitals + mf = scf.RHF(mol.copy()) + mf.kernel() + + assert mf.converged + + MPI.COMM_WORLD.Bcast(mf.mo_coeff) + + # Specificy the CAS solver for the current state + if self.solver == 'SA-CASSCF': + mc = mcscf.CASSCF(mf, self.ncas, self.neleca).state_average_([1/self.nroots]*self.nroots) + mc.kernel() + mo_sacasscf = mc.mo_coeff + + # Iterate over different states + for istate in range(self.nroots): + + if self.solver == 'CASCI': + casci_bra = mcscf.CASCI(mf, self.ncas, self.neleca).state_specific_(istate) + elif self.solver == 'SS-CASSCF': + cas_ss = mcscf.CASSCF(mf, self.ncas, self.neleca).state_specific_(istate) + cas_ss.kernel() + casci_bra = mcscf.CASCI(mf, self.ncas, self.neleca).state_specific_(istate) + casci_bra.casci(cas_ss.mo_coeff) + else: + casci_bra = mcscf.CASCI(mf, self.ncas, self.neleca).state_specific_(istate) + casci_bra.casci(mo_sacasscf) + + self.cascis.append(casci_bra) + + def states_to_rdms(self): + """ + Construct transition RDMs between given training states (self.cascis) + + Raises: + AssertionError: If the mean-field calculation is not converged. + """ + # Some checks + assert len(self.cascis) > 0 + + assert self.two_rdm is None and self.one_rdm is None + + states = self.cascis + + n_cascis = len(states) + + if rank == 0: + overlap_new = np.zeros((n_cascis, n_cascis)) + one_rdm_new = np.zeros((n_cascis, n_cascis, states[0].mo_coeff.shape[0], states[0].mo_coeff.shape[0])) + if not self.lowrank: + two_rdm_new = np.zeros((n_cascis, n_cascis, + states[0].mo_coeff.shape[0], + states[0].mo_coeff.shape[0], + states[0].mo_coeff.shape[0], + states[0].mo_coeff.shape[0])) + else: + diagonal_lr_new = np.ones((n_cascis, n_cascis, 3, + states[0].mo_coeff.shape[0], + states[0].mo_coeff.shape[0])) + vecs_lowrank = {} + else: + overlap_new = one_rdm_new = two_rdm_new = None + if self.lowrank: + diagonal_lr_new = None + vecs_lowrank = None + + + # Iterate over bra states + for a, casci_bra in enumerate(states): + + MPI.COMM_WORLD.Bcast(casci_bra.ci) + MPI.COMM_WORLD.Bcast(casci_bra.mo_coeff) + + #st_bra = time() + + mo_coeff_bra = casci_bra.mo_coeff + mol_bra = casci_bra.mol + + ovlp_bra = mol_bra.intor_symmetric("int1e_ovlp") + basis_OAO_bra = get_basis(mol_bra) + trafo_bra = basis_OAO_bra.T.dot(ovlp_bra).dot(mo_coeff_bra) + #print('bra',ovlp_bra.shape,basis_OAO_bra.shape,trafo_bra.shape) + + bra_ref_state = wick.reference_state[float]( + mo_coeff_bra.shape[0], + mo_coeff_bra.shape[0], + mol_bra.nelec[0], + casci_bra.ncas, + casci_bra.ncore, + owndata(mo_coeff_bra), ) - overlap_accumulate = 0.0 - - all_ids = np.array( - [ - [iabra, ibbra, iaket, ibket] - for iabra in range(len(bra_occ_strings)) - for ibbra in range(len(bra_occ_strings)) - for iaket in range(len(ket_occ_strings)) - for ibket in range(len(ket_occ_strings)) - ] + + bra_occ_strings = utils.fci_bitset_list( + mol_bra.nelec[0] - casci_bra.ncore, casci_bra.ncas ) - n_ranks = MPI.COMM_WORLD.Get_size() + # Iterate over ket states + #for b, casci_ket in enumerate(states): + for b in range(n_cascis): + casci_ket = states[b] + #st_ket = time() - all_ids_local = np.array_split(all_ids, n_ranks)[rank] + # Prepare ket state + mo_coeff_ket = casci_ket.mo_coeff + mol_ket = casci_ket.mol - if rank == 0: - pbar = tqdm(total=len(all_ids_local)) - - for ids in all_ids_local: - iabra, ibbra, iaket, ibket = ids - stringabra = bra_occ_strings[iabra] - stringbbra = bra_occ_strings[ibbra] - stringaket = ket_occ_strings[iaket] - stringbket = ket_occ_strings[ibket] - - rdm1_tmp.fill(0.0) - rdm2_tmp.fill(0.0) - o = wick_mb.evaluate_rdm12( - stringabra, - stringbbra, - stringaket, - stringbket, - 1.0, - rdm1_tmp, - rdm2_tmp, + ovlp_ket = mol_ket.intor_symmetric("int1e_ovlp") + basis_OAO_ket = get_basis(mol_ket) + trafo_ket = basis_OAO_ket.T.dot(ovlp_ket).dot(mo_coeff_ket) + + trafo_ket_bra = basis_OAO_bra.dot(trafo_ket) + #print('bra',ovlp_ket.shape,basis_OAO_ket.shape,trafo_ket.shape, trafo_ket_bra.shape) + + ket_ref_state = wick.reference_state[float]( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mol_ket.nelec[0], + casci_ket.ncas, + casci_ket.ncore, + owndata(trafo_ket_bra), + ) + + ket_occ_strings = utils.fci_bitset_list( + mol_ket.nelec[0] - casci_ket.ncore, casci_ket.ncas ) - overlap_accumulate += ( - o * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + orbitals = wick.wick_orbitals[float, float]( + bra_ref_state, ket_ref_state, owndata(ovlp_bra) ) - rdm1 += ( - rdm1_tmp * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + wick_mb = wick.wick_rscf[float, float, float](orbitals, 0.0) + + rdm1_tmp = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) + rdm1 = np.zeros((mo_coeff_ket.shape[0], mo_coeff_ket.shape[0])) + rdm2_tmp = np.zeros( + ( + mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0] * mo_coeff_ket.shape[0], + ) + ) + rdm2 = np.zeros( + ( + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + mo_coeff_ket.shape[0], + ) ) - rdm2 += ( - rdm2_tmp.reshape(rdm2.shape) - * casci_bra.ci[iabra, ibbra] - * casci_ket.ci[iaket, ibket] + overlap_accumulate = 0.0 + + all_ids = np.array( + [ + [iabra, ibbra, iaket, ibket] + for iabra in range(len(bra_occ_strings)) + for ibbra in range(len(bra_occ_strings)) + for iaket in range(len(ket_occ_strings)) + for ibket in range(len(ket_occ_strings)) + ] ) - if rank == 0: - pbar.update(1) + n_ranks = MPI.COMM_WORLD.Get_size() - if rank == 0: - pbar.close() + all_ids_local = np.array_split(all_ids, n_ranks)[rank] - overlap_accumulate = MPI.COMM_WORLD.allreduce( - overlap_accumulate, op=MPI.SUM - ) + if rank == 0: + pbar = tqdm(total=len(all_ids_local)) + + for ids in all_ids_local: + iabra, ibbra, iaket, ibket = ids + stringabra = bra_occ_strings[iabra] + stringbbra = bra_occ_strings[ibbra] + stringaket = ket_occ_strings[iaket] + stringbket = ket_occ_strings[ibket] + + rdm1_tmp.fill(0.0) + rdm2_tmp.fill(0.0) + o = wick_mb.evaluate_rdm12( + stringabra, + stringbbra, + stringaket, + stringbket, + 1.0, + rdm1_tmp, + rdm2_tmp, + ) + overlap_accumulate += ( + o * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + ) + + rdm1 += ( + rdm1_tmp * casci_bra.ci[iabra, ibbra] * casci_ket.ci[iaket, ibket] + ) + rdm2 += ( + rdm2_tmp.reshape(rdm2.shape) + * casci_bra.ci[iabra, ibbra] + * casci_ket.ci[iaket, ibket] + ) + + if rank == 0: + pbar.update(1) - MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm1, op=MPI.SUM) - MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm2, op=MPI.SUM) + if rank == 0: + pbar.close() - if rank == 0: - overlap_new[-1, i] = overlap_accumulate - overlap_new[i, -1] = overlap_accumulate.conj() - rdm1 = np.einsum( - "...ij,ai,bj->...ab", rdm1, trafo_ket, trafo_bra, optimize="optimal" - ) - rdm2 = np.einsum( - "...ijkl,ai,bj,ck,dl->...abcd", - rdm2, - trafo_bra, - trafo_ket, - trafo_bra, - trafo_ket, - optimize="optimal", + overlap_accumulate = MPI.COMM_WORLD.allreduce( + overlap_accumulate, op=MPI.SUM ) - one_rdm_new[-1, i, :, :] = rdm1 - one_rdm_new[i, -1, :, :] = rdm1.conj() - two_rdm_new[-1, i, :, :, :, :] = rdm2 - two_rdm_new[i, -1, :, :, :, :] = rdm2.conj() + MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm1, op=MPI.SUM) + MPI.COMM_WORLD.Allreduce(MPI.IN_PLACE, rdm2, op=MPI.SUM) + + if rank == 0: + overlap_new[a, b] = overlap_accumulate + overlap_new[b, a] = overlap_accumulate.conj() + rdm1 = np.einsum( + "...ij,ai,bj->...ab", rdm1, trafo_ket, trafo_bra, optimize="optimal" + ) + rdm2 = np.einsum( + "...ijkl,ai,bj,ck,dl->...abcd", + rdm2, + trafo_bra, + trafo_ket, + trafo_bra, + trafo_ket, + optimize="optimal", + ) + + one_rdm_new[a, b, :, :] = rdm1 + one_rdm_new[b, a, :, :] = rdm1.conj() + + if not self.lowrank: + two_rdm_new[a, b, :, :, :, :] = rdm2 + two_rdm_new[b, a, :, :, :, :] = np.einsum('ijkl->klij', rdm2.conj()) + else: + lowrank_vecs, diagonals, use_joint = reduce_2rdm( + rdm1, + rdm2, + overlap_accumulate, + mol=mol_bra, + train_en=casci_bra.e_tot, + **self.kwargs, + ) + + diagonal_lr_new[a, b, :, :, :] = diagonals + try: + diagonal_lr_new[b, a, :, :, :] = diagonals.conj() + except: + diagonal_lr_new[b, a, :, :, :] = diagonals + + vecs_lowrank[(a, b)] = ( + lowrank_vecs[0], + lowrank_vecs[1], + lowrank_vecs[2], + use_joint, + ) + vecs_lowrank[(b, a)] = ( + lowrank_vecs[0].conj(), + lowrank_vecs[1].conj(), + lowrank_vecs[2].conj(), + use_joint, + ) + self.overlap = overlap_new self.one_rdm = one_rdm_new - self.two_rdm = two_rdm_new + if not self.lowrank: + self.two_rdm = two_rdm_new + else: + self.diagonal_lr = diagonal_lr_new + self.vecs_lowrank = vecs_lowrank def prune_datapoints(self, keep_ids): """ - Prunes training points from the continuation object based on the given keep_ids. + Prune training points (states/geometries) from the continuation object. - Args: - keep_ids (list): List of indices to keep. + This adapts to the newer storage model where individual state data are + held in parallel lists (mo_coeffs, cis, trafos, mols) and low-rank data + are stored in dictionaries keyed by (i,j) pairs. - Returns: - None + Parameters + ---------- + keep_ids : sequence[int] or sequence[bool] + Indices to keep (integer list) or a boolean mask of length n_states. + Order is preserved as given. """ + # Normalize keep_ids: allow boolean mask + import numpy as _np + if isinstance(keep_ids, (list, tuple)) and len(keep_ids) > 0 and isinstance(keep_ids[0], (bool, _np.bool_)): + keep_ids = [i for i, flag in enumerate(keep_ids) if flag] + else: + keep_ids = list(keep_ids) + + if len(keep_ids) == 0: + raise ValueError("prune_datapoints: keep_ids is empty; refusing to drop all datapoints.") + + # Ensure indices are within range + n_states = len(self.mo_coeffs) + if any((i < 0 or i >= n_states) for i in keep_ids): + raise IndexError("prune_datapoints: keep_ids contains out-of-range indices.") + + # Deduplicate while preserving order + seen = set(); ordered_keep = [] + for i in keep_ids: + if i not in seen: + ordered_keep.append(i); seen.add(i) + keep_ids = ordered_keep + + # Core square matrices/tensors if self.overlap is not None: - self.overlap = self.overlap[np.ix_(keep_ids, keep_ids)] + self.overlap = self.overlap[_np.ix_(keep_ids, keep_ids)] if self.one_rdm is not None: - self.one_rdm = self.one_rdm[np.ix_(keep_ids, keep_ids)] + # shape (n,n,nao,nao) + self.one_rdm = self.one_rdm[_np.ix_(keep_ids, keep_ids)] if self.two_rdm is not None: - self.two_rdm = self.two_rdm[np.ix_(keep_ids, keep_ids)] - self.cascis = [self.cascis[i] for i in keep_ids] + self.two_rdm = self.two_rdm[_np.ix_(keep_ids, keep_ids)] + if self.lowrank and self.diagonal_lr is not None: + self.diagonal_lr = self.diagonal_lr[_np.ix_(keep_ids, keep_ids)] + + # Parallel lists of per-state data + self.mo_coeffs = [self.mo_coeffs[i] for i in keep_ids] + self.cis = [self.cis[i] for i in keep_ids] + self.trafos = [self.trafos[i] for i in keep_ids] + if hasattr(self, 'mols') and self.mols is not None and len(self.mols) == n_states: + self.mols = [self.mols[i] for i in keep_ids] + + # Precompute-related arrays (if present) + # inv_OAO_all, occ_strings_all, mb_all, orbitals_all created in precompute_for_otf + if getattr(self, 'precompute', False): + if hasattr(self, 'inv_OAO_all') and len(self.inv_OAO_all) == n_states: + self.inv_OAO_all = [self.inv_OAO_all[i] for i in keep_ids] + if hasattr(self, 'occ_strings_all') and len(self.occ_strings_all) == n_states: + self.occ_strings_all = [self.occ_strings_all[i] for i in keep_ids] + if hasattr(self, 'mb_all') and len(self.mb_all) == n_states: + self.mb_all = [self.mb_all[i] for i in keep_ids] + if hasattr(self, 'orbitals_all') and len(self.orbitals_all) == n_states: + self.orbitals_all = [self.orbitals_all[i] for i in keep_ids] + + # Low-rank dictionary remapping + if self.lowrank and self.vecs_lowrank is not None: + vecs_lowrank_new = {} + for new_i, old_i in enumerate(keep_ids): + for new_j, old_j in enumerate(keep_ids): + key_old = (old_i, old_j) + if key_old in self.vecs_lowrank: + vecs_lowrank_new[(new_i, new_j)] = self.vecs_lowrank[key_old] + self.vecs_lowrank = vecs_lowrank_new + + # Sanity: update counts if stored elsewhere + # (No explicit n_states attribute; len(self.mo_coeffs) is authoritative.) + return + + def save(self, filename): + """ + Save the attributes of this CAS_EVCont_obj to a file. + + Args: + filename (str): Path to the output file (will be saved as pickle) + + Returns: + None + """ + # Collect all the essential attributes + cas_data = { + # Basic parameters + 'ncas': self.ncas, + 'neleca': self.neleca, + 'nroots': self.nroots, + 'solver': self.solver, + 'lowrank': self.lowrank, + 'software': self.software, + 'quantel_path': getattr(self, 'quantel_path', None), + 'solutions_to_reconverge': getattr(self, 'solutions_to_reconverge', None), + + # RDM and overlap data + 'overlap': self.overlap, + 'one_rdm': self.one_rdm, + 'two_rdm': self.two_rdm, + + # Low-rank specific data (if applicable) + 'diagonal_lr': self.diagonal_lr if self.lowrank else None, + 'vecs_lowrank': self.vecs_lowrank if self.lowrank else None, + 'kwargs': self.kwargs if self.lowrank else None, + + # State information + 'mo_coeffs': self.mo_coeffs, + 'cis': self.cis, + 'trafos': self.trafos, + + # Additional flags + 'uncontracted': self.uncontracted, + 'use_rdm': self.use_rdm, + 'precompute': self.precompute, + + # Save pyscf Molecule object geometries (not the full mols since it's not pickleable) + 'molecule_geometries': [mol.atom for mol in self.mols], + 'molecule_basis': [mol.basis for mol in self.mols], + 'molecule_unit': [mol.unit for mol in self.mols], + } + + + # Save to pickle file + with open(filename, 'wb') as f: + pickle.dump(cas_data, f, protocol=pickle.HIGHEST_PROTOCOL) + + if rank == 0: + print(f"CAS object saved to {filename}") + + @classmethod + def load(cls, filename): + """ + Load CAS_EVCont_obj attributes from a file and reinitialize the object. + + Args: + filename (str): Path to the input file (pickle format) + + Returns: + CAS_EVCont_obj: Reinitialized CAS_EVCont_obj instance + """ + # Load the saved data + with open(filename, 'rb') as f: + cas_data = pickle.load(f) + + # Reinitialize the CAS object with basic parameters + # Extract optional parameters with defaults for backward compatibility + software = cas_data.get('software', 'pyscf') + quantel_path = cas_data.get('quantel_path', None) + solutions_to_reconverge = cas_data.get('solutions_to_reconverge', None) + + if cas_data['lowrank']: + cas_obj = cls( + cas_data['ncas'], + cas_data['neleca'], + nroots=cas_data['nroots'], + solver=cas_data['solver'], + software=software, + quantel_path=quantel_path, + solutions_to_reconverge=solutions_to_reconverge, + lowrank=True, + **cas_data['kwargs'] + ) + else: + cas_obj = cls( + cas_data['ncas'], + cas_data['neleca'], + nroots=cas_data['nroots'], + solver=cas_data['solver'], + software=software, + quantel_path=quantel_path, + solutions_to_reconverge=solutions_to_reconverge, + lowrank=False + ) + + # Restore RDM and overlap data + cas_obj.overlap = cas_data['overlap'] + cas_obj.one_rdm = cas_data['one_rdm'] + cas_obj.two_rdm = cas_data['two_rdm'] + + # Restore low-rank data if applicable + if cas_data['lowrank']: + cas_obj.diagonal_lr = cas_data['diagonal_lr'] + cas_obj.vecs_lowrank = cas_data['vecs_lowrank'] + + # Restore state information + cas_obj.mo_coeffs = cas_data['mo_coeffs'] + cas_obj.cis = cas_data['cis'] + cas_obj.trafos = cas_data['trafos'] + + # Restore additional flags + cas_obj.uncontracted = cas_data['uncontracted'] + cas_obj.use_rdm = cas_data['use_rdm'] + cas_obj.precompute = cas_data['precompute'] + + # Restore molecule objects if geometries were saved + if 'molecule_geometries' in cas_data: + cas_obj.mols = [] + for geom, basis, unit in zip(cas_data['molecule_geometries'], cas_data['molecule_basis'], cas_data['molecule_unit']): + mol = gto.Mole() + mol.build(atom=geom, basis=basis, unit=unit, verbose=0) + cas_obj.mols.append(mol) + + if rank == 0: + print(f"CAS object loaded from {filename}") + print(f" ncas={cas_obj.ncas}, neleca={cas_obj.neleca}, nroots={cas_obj.nroots}") + print(f" solver={cas_obj.solver}, lowrank={cas_obj.lowrank}, software={cas_obj.software}") + print(f" Number of states: {len(cas_obj.cis)}") + if cas_obj.software == 'quantel': + print(f" quantel_path={cas_obj.quantel_path}") + print(f" solutions_to_reconverge={cas_obj.solutions_to_reconverge}") + + return cas_obj + + +# Quantel specific parser functions +def create_next_geom_dir(quantel_path): + """ + Create the next geometry directory in the Quantel path. + Args: + quantel_path (str): Path to the Quantel directory. + Returns: + str: Name of the newly created geometry directory. + """ + nums = [] + for name in os.listdir(quantel_path): + if os.path.isdir(os.path.join(quantel_path, name)): + m = re.fullmatch(r'geom(\d+)', name) + if m: + nums.append(int(m.group(1))) + next_n = (max(nums) + 1) if nums else 1 + new_name = f"geom{next_n}" + os.makedirs(os.path.join(quantel_path, new_name)) + return new_name \ No newline at end of file diff --git a/evcont/DMRG_EVCont.py b/evcont/DMRG_EVCont.py index 4967756..2805dfa 100644 --- a/evcont/DMRG_EVCont.py +++ b/evcont/DMRG_EVCont.py @@ -21,6 +21,8 @@ def append_to_rdms_OAO_basis( one_rdm=None, two_rdm=None, converge_dmrg_fun=converge_dmrg, + nroots=1, + roots_train=[1], mem=5, ): """ @@ -30,6 +32,10 @@ def append_to_rdms_OAO_basis( """ mol_bra = mols[-1] + mol_ind = len(mols)-1 + mol_tag = mol_ind + + new_tags = tags h1, h2 = get_integrals(mol_bra, get_basis(mol_bra, basis_type="OAO")) @@ -53,38 +59,79 @@ def append_to_rdms_OAO_basis( ) mps_solver.initialize_system(norb, n_elec=nelec, spin=mol_bra.spin) - converge_dmrg_fun(h1, h2, mol_bra.nelec, "MPS_{}".format(tags[-1])) + converge_dmrg_fun(h1, h2, mol_bra.nelec, "MPS_{}".format(mol_tag), nroots=nroots) - bra = mps_solver.load_mps("MPS_{}".format(tags[-1])) + bra = mps_solver.load_mps("MPS_{}".format(mol_tag),nroots=nroots) - overlap_new = np.ones((len(mols), len(mols))) - if overlap is not None: - overlap_new[:-1, :-1] = overlap - one_rdm_new = np.ones((len(mols), len(mols), norb, norb)) - if one_rdm is not None: - one_rdm_new[:-1, :-1, :, :] = one_rdm - two_rdm_new = np.ones((len(mols), len(mols), norb, norb, norb, norb)) - if two_rdm is not None: - two_rdm_new[:-1, :-1, :, :, :, :] = two_rdm + # If ground state + if nroots == 1: + bras = [bra] + else: + bras = [mps_solver.split_mps(bra, ir, tag="MPS_%d_%d"%(mol_tag,ir)) for ir in range(nroots)] - for i, mol_ket in enumerate(mols): - ket = mps_solver.load_mps("MPS_{}".format(tags[i])) + # Iterate over ground and excited states include them in the training + # if their index is in self.roots_train + for ind in range(nroots): + if ind in roots_train: - ovlp = ( - np.array(mps_solver.expectation(bra, mps_solver.get_identity_mpo(), ket)) - / n_ranks - ) - o_RDM = np.array(mps_solver.get_1pdm(ket, bra=bra)) - t_RDM = np.array(np.transpose(mps_solver.get_2pdm(ket, bra=bra), (0, 3, 1, 2))) + # Add tags at each time + if nroots == 1: + new_tags.append(mol_tag) + else: + tag_ir = "%d_%d"%(mol_ind,ind) + new_tags.append(tag_ir) + + # Initialize intermediate representation + nvec = len(new_tags) + overlap_new = np.ones((nvec, nvec)) + if overlap is not None: + overlap_new[:-1, :-1] = overlap + one_rdm_new = np.ones((nvec, nvec, norb, norb)) + if one_rdm is not None: + one_rdm_new[:-1, :-1, :, :] = one_rdm + two_rdm_new = np.ones((nvec, nvec, norb, norb, norb, norb)) + if two_rdm is not None: + two_rdm_new[:-1, :-1, :, :, :, :] = two_rdm + + # Iterate over kets and add the new state to the representation + print(new_tags) + print(overlap_new.shape) + for i, tag_i in enumerate(new_tags): + + ket = mps_solver.load_mps("MPS_{}".format(tag_i))#,nroots=nroots) + + ovlp = ( + np.array(mps_solver.expectation(bras[ind], mps_solver.get_identity_mpo(), ket)) + / n_ranks + ) - overlap_new[-1, i] = ovlp - overlap_new[i, -1] = ovlp.conj() - one_rdm_new[-1, i, :, :] = o_RDM - one_rdm_new[i, -1, :, :] = o_RDM.conj() - two_rdm_new[-1, i, :, :, :, :] = t_RDM - two_rdm_new[i, -1, :, :, :, :] = t_RDM.conj() + order_1 = (1, 0) + o_RDM = np.array(mps_solver.get_trans_1pdm(bras[ind],ket)).transpose(order_1) + o_RDM_conj = np.array(mps_solver.get_trans_1pdm(ket,bras[ind])).transpose(order_1) - return overlap_new, one_rdm_new, two_rdm_new + order_2 = (0, 3, 1, 2) + #order = (3, 0, 2, 1) + t_RDM = np.array(np.transpose(mps_solver.get_trans_2pdm(bras[ind],ket), order_2)) + t_RDM_conj = np.array(np.transpose(mps_solver.get_trans_2pdm(ket,bras[ind]), order_2)) + + overlap_new[-1, i] = ovlp + overlap_new[i, -1] = ovlp.conj() + + one_rdm_new[-1, i, :, :] = o_RDM + one_rdm_new[i, -1, :, :] = o_RDM_conj + #one_rdm_new[i, -1, :, :] = np.transpose(o_RDM.conj(),(1,0)) + #one_rdm_new[i, -1, :, :] = o_RDM.conj() + + two_rdm_new[-1, i, :, :, :, :] = t_RDM + two_rdm_new[i, -1, :, :, :, :] = t_RDM_conj + #two_rdm_new[i, -1, :, :, :, :] = np.transpose(t_RDM.conj(), (1,0,3,2)) + + # Update for next iteration + overlap = overlap_new.copy() + two_rdm = two_rdm_new.copy() + one_rdm = one_rdm_new.copy() + + return overlap_new, one_rdm_new, two_rdm_new, new_tags def append_to_rdms_rerun( @@ -438,6 +485,8 @@ def __init__( self, dmrg_converge_fun=converge_dmrg, append_method=append_to_rdms_OAO_basis, + nroots=1, + roots_train=None, mem=5, ): """ @@ -445,15 +494,24 @@ def __init__( Args: dmrg_converge_fun: The function to converge DMRG at each training point. - append_method: The method to append to rdms (see implementations above). - mem: The size of usable memory for block2 (in GB). + append_method : The method to append to rdms (see implementations above). + nroots (int) : Number of states to be solved. Default is 1, the ground state. + roots_train (list): Indices of states to include in the continuation + mem : The size of usable memory for block2 (in GB). """ self.solver = dmrg_converge_fun self.append_method = append_method + self.nroots = nroots + if roots_train == None: + self.roots_train = list(range(nroots)) + else: + self.roots_train = roots_train + assert isinstance(roots_train,list) + self.mols = [] self.tags = [] - self.max_tag = 0 + #self.max_tag = 0 self.overlap = None self.one_rdm = None self.two_rdm = None @@ -467,15 +525,17 @@ def append_to_rdms(self, mol): mol: The molecule to append. """ self.mols.append(mol) - self.tags.append(self.max_tag) - self.max_tag += 1 - self.overlap, self.one_rdm, self.two_rdm = self.append_method( + #self.tags.append(self.max_tag) + #self.max_tag += 1 + self.overlap, self.one_rdm, self.two_rdm, self.tags = self.append_method( self.mols, self.tags, overlap=self.overlap, one_rdm=self.one_rdm, two_rdm=self.two_rdm, converge_dmrg_fun=self.solver, + nroots=self.nroots, + roots_train=self.roots_train, mem=self.mem, ) diff --git a/evcont/FCI_EVCont.py b/evcont/FCI_EVCont.py index 26e963c..6bbd9d6 100644 --- a/evcont/FCI_EVCont.py +++ b/evcont/FCI_EVCont.py @@ -1,11 +1,16 @@ import numpy as np +import sys +import itertools from evcont.electron_integral_utils import get_basis, get_integrals -from pyscf import fci +from pyscf import scf, ao2mo, fci, symm from pyscf.fci.addons import transform_ci +from evcont.ab_initio_gradients_loewdin import get_loewdin_trafo + +from evcont.low_rank_utils import reduce_2rdm, vectorize_lowrank, unpack_vectorized_lowrank class FCI_EVCont_obj: """ @@ -18,6 +23,9 @@ def __init__( cibasis='canonical', nroots=1, roots_train=None, + irrep_name=None, + lowrank=False, + **kwargs ): """ Initializes the FCI_EVCont_obj class. @@ -28,6 +36,8 @@ def __init__( Note that after computation, the basis is converted to OAO nroots: Number of states to be solved. Default is 1, the ground state. roots_train (list): Indices of states to include in the continuation + irrep_name (string): If not None, only include states corresponding + to the given symmetry irreducible representation Attributes: fcivecs (list): The FCI training states. @@ -47,14 +57,47 @@ def __init__( self.roots_train = roots_train assert isinstance(roots_train,list) + # Symmetry + if irrep_name == None: + self.use_symmetry = False + self.irrep_name = None + else: + self.use_symmetry = True + self.irrep_name = irrep_name + + # Need canonical basis + assert cibasis == 'canonical' + # Initialize attributes self.fcivecs = [] self.ens = [] + self.ens_nuc = [] self.mol_index = [] self.overlap = None self.one_rdm = None self.two_rdm = None + + ### Initialize low-rank attributes + self.lowrank = lowrank + if lowrank: + #self.truncation_style = kwargs['truncation_style'] + self.kwargs = kwargs + + # Diagonals of 2-cumulants ([nbra, nket, 3, norb, norb]) + self.diagonal_lr = None + # Low rank eigendecomposition of the rest of 2-cumulant + # Old version: dictionary[(nbra, nket)] = (vals_trunc, vecs_trunc) + # New version: dictionary['vals': np.array([nbra, nket, nvec]), + # 'vecs': np.array([nbra, nket, nvec, nao, nao])] + + self.vecs_lowrank = {} + + def vectorize_lowrank(self,hermitian=True): + vectorize_lowrank(self,hermitian=hermitian) + def unpack_vectorized_lowrank(self): + unpack_vectorized_lowrank(self) + def append_to_rdms(self, mol): """ Append a new training geometry by growing the t-RDMs. @@ -63,18 +106,31 @@ def append_to_rdms(self, mol): mol (object): Molecular object of the training geometry. """ + # Relevant matrices for SAO basis + #S = mol.intor("int1e_ovlp") + #ao_mo_trafo = get_loewdin_trafo(S) + basis = get_basis(mol,basis_type=self.cibasis) h1, h2 = get_integrals(mol, basis) nroots_train = max(self.roots_train)+1 - e_all, fcivec_all = self.cisolver.kernel(h1, h2, mol.nao, mol.nelec, - nroots=nroots_train) + + # Get FCI energies and wavefunctions in SAO basis + if not self.use_symmetry: + e_all, fcivec_all = self.cisolver.kernel(h1, h2, mol.nao, mol.nelec, + nroots=nroots_train) + + else: + orbsym = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, basis) + + e_all, fcivec_all = self.cisolver.kernel(h1, h2, mol.nao, mol.nelec, + nroots=nroots_train, orbsym=orbsym,) # If ground state if nroots_train == 1: e_all = [e_all] fcivec_all = [fcivec_all] - + # Transform to OAO basis if self.cibasis != 'OAO': S = mol.intor("int1e_ovlp") @@ -84,6 +140,12 @@ def append_to_rdms(self, mol): fcivec_all = [transform_ci(fcivec_i,mol.nelec,u) for fcivec_i in fcivec_all] + # Fix gauge; probably not necessary + for fcivec_i in fcivec_all: + # Set maximum element to be positive + idx = np.unravel_index(np.argmax(np.abs(fcivec_i.real)),fcivec_i.shape) + fcivec_i *= np.sign(fcivec_i[idx]) + # Setting molecular index if len(self.mol_index) == 0: mindex = 0 @@ -100,20 +162,36 @@ def append_to_rdms(self, mol): self.fcivecs.append(fcivec) - self.ens.append(e + mol.energy_nuc()) + self.ens.append(e) + self.ens_nuc.append(mol.energy_nuc()) self.mol_index.append(mindex) - - overlap_new = np.ones((len(self.fcivecs), len(self.fcivecs))) + + new_ntrain = len(self.fcivecs) + + overlap_new = np.ones((new_ntrain, new_ntrain)) if self.overlap is not None: overlap_new[:-1, :-1] = self.overlap one_rdm_new = np.ones((len(self.fcivecs), len(self.fcivecs), mol.nao, mol.nao)) if self.one_rdm is not None: one_rdm_new[:-1, :-1, :, :] = self.one_rdm - two_rdm_new = np.ones( - (len(self.fcivecs), len(self.fcivecs), mol.nao, mol.nao, mol.nao, mol.nao) - ) - if self.two_rdm is not None: - two_rdm_new[:-1, :-1, :, :, :, :] = self.two_rdm + + # Only define two_rdm if not lowrank + if not self.lowrank: + two_rdm_new = np.ones( + (len(self.fcivecs), len(self.fcivecs), mol.nao, mol.nao, mol.nao, mol.nao) + ) + if self.two_rdm is not None: + two_rdm_new[:-1, :-1, :, :, :, :] = self.two_rdm + + else: + diagonal_lr_new = np.ones( + (len(self.fcivecs), len(self.fcivecs), 3, mol.nao, mol.nao) + ) + if self.diagonal_lr is not None: + diagonal_lr_new[:-1, :-1, :, :, :] = self.diagonal_lr + + + # Iterate over training states to add RDMs to the existing states for i in range(len(self.fcivecs)): ovlp = self.fcivecs[-1].flatten().conj().dot(self.fcivecs[i].flatten()) overlap_new[-1, i] = ovlp @@ -121,14 +199,55 @@ def append_to_rdms(self, mol): rdm1, rdm2 = self.cisolver.trans_rdm12( self.fcivecs[-1], self.fcivecs[i], mol.nao, mol.nelec ) + #rdm1_conj, rdm2_conj = self.cisolver.trans_rdm12( + # self.fcivecs[i], self.fcivecs[-1], mol.nao, mol.nelec + #) one_rdm_new[-1, i, :, :] = rdm1 - one_rdm_new[i, -1, :, :] = rdm1.conj() - two_rdm_new[-1, i, :, :, :, :] = rdm2 - two_rdm_new[i, -1, :, :, :, :] = rdm2.conj() + one_rdm_new[i, -1, :, :] = rdm1.conj().T + #one_rdm_new[i, -1, :, :] = rdm1_conj + + if not self.lowrank: + two_rdm_new[-1, i, :, :, :, :] = rdm2 + two_rdm_new[i, -1, :, :, :, :] = np.einsum('ijkl->lkji',rdm2.conj()) + #two_rdm_new[i, -1, :, :, :, :] = rdm2_conj + + # Low rank + else: + print('States: %i %i, overlap: %f' % (new_ntrain-1, i, ovlp)) + + # Get low rank representation + lowrank_vecs, diagonals, use_joint = \ + reduce_2rdm(rdm1, rdm2, ovlp, + mol=mol, train_en=e, + **self.kwargs) + + #lowrank_vecs_conj, diagonals_conj = \ + # reduce_2rdm(rdm1_conj, rdm2_conj, ovlp, + # mol=mol, train_en=e, + # **self.kwargs) + + diagonal_lr_new[-1, i, :, :, :] = diagonals + #diagonal_lr_new[i, -1, :, :, :] = diagonals_conj + try: + # This gives an error if diagonals are not saved and set to None by reduce_2rdm + diagonal_lr_new[i, -1, :, :, :] = diagonals.conj() + except: + diagonal_lr_new[i, -1, :, :, :] = diagonals + + #self.vecs_lowrank[(new_ntrain-1, i)] = lowrank_vecs + #self.vecs_lowrank[(i, new_ntrain-1)] = lowrank_vecs_conj + + self.vecs_lowrank[(new_ntrain-1, i)] = lowrank_vecs[0], lowrank_vecs[1], lowrank_vecs[2], use_joint + #vecs_lowrank[(i,n_cascis-1)] = lowrank_vecs_conj + self.vecs_lowrank[(i,new_ntrain-1)] = lowrank_vecs[0].conj(), lowrank_vecs[1].conj(), lowrank_vecs[2].conj(), use_joint + self.overlap = overlap_new self.one_rdm = one_rdm_new - self.two_rdm = two_rdm_new + if not self.lowrank: + self.two_rdm = two_rdm_new + else: + self.diagonal_lr = diagonal_lr_new def prune_datapoints(self, keep_ids): """ @@ -141,6 +260,10 @@ def prune_datapoints(self, keep_ids): None """ + if self.nroots > 1 or self.lowrank: + print('Error in prune_datapoints: Pruning has not been implemented for excited states or low rank implementations') + sys.exit() + if self.overlap is not None: self.overlap = self.overlap[np.ix_(keep_ids, keep_ids)] if self.one_rdm is not None: diff --git a/evcont/FCI_NAC.py b/evcont/FCI_NAC.py new file mode 100644 index 0000000..5416915 --- /dev/null +++ b/evcont/FCI_NAC.py @@ -0,0 +1,461 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Mon Nov 13 13:59:17 2023 + +Full CI nonadiabatic coupling vectors in SAO basis + +@author: katalar +""" + +import numpy as np + +from pyscf import scf, ao2mo, grad, fci, symm + +from pyscf.fci.addons import transform_ci + +from evcont.electron_integral_utils import get_basis, get_integrals + +from evcont.ab_initio_gradients_loewdin import ( + get_one_and_two_el_grad, + get_loewdin_trafo, + get_orbital_derivative_coupling, + get_grad_elec_from_gradH, + get_grad_elec_OAO +) + +def get_FCI_energy_with_grad_and_NAC(mol, fcisolver, cibasis='canonical', nroots=1, savemem=True, hermitian=True): + """ + Calculates the potential energiesm its gradient w.r.t. nuclear positions of a + molecule and nonadiabatic couplings from full CI in SAO basis + + Args: + mol : pyscf.gto.Mole + The molecule object. + fcisolver : pyscf.fci solver + The solver object for full CI + nroots (optional): int + Number of states in the solver. + hermitian (optional): bool + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple (en, grad_all, nac_all, nac_all_hfonly) + A tuple containing the total potential energies, its gradients and NACs: + + en: ndarray(nroot,) + Total potential energies for both ground and excited states + + grad_all: list of ndarray(nat,3) + Gradients of multistate energies + + nac_all: dictionary of np.darray(nat,3) + Nonadiabatic coupling vectors between all states, + e.g. nac_all['02'] is NAC along ground state and 2nd excited state + + nac_all_hfonly: dictionary of ndarray(nat,3) + Hellman-Feynmann contribution to NACs + """ + + # Construct h1 and h2 + basis = get_basis(mol,basis_type=cibasis) + h1, h2 = get_integrals(mol, basis) + + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + #h1 = np.linalg.multi_dot((ao_mo_trafo.T, scf.hf.get_hcore(mol), ao_mo_trafo)) + #h2 = ao2mo.restore(1, ao2mo.kernel(mol, ao_mo_trafo), mol.nao) + + # Get FCI energies and wavefunctions in SAO basis + en, fcivec = fcisolver.kernel(h1, h2, mol.nao, mol.nelec) + + # Transform to SAO basis + if cibasis != 'OAO': + S = mol.intor("int1e_ovlp") + basis_oao = ao_mo_trafo #get_basis(mol) + + u = np.einsum('ji,jk,kl->il',basis,S,basis_oao) + + fcivec = [transform_ci(fcivec_i,mol.nelec,u) for fcivec_i in fcivec] + + # Get the gradient of one and two-electron integrals before contracting onto + # rdms and trmds of different states + if not savemem: + h1_jac, h2_jac = get_one_and_two_el_grad(mol,ao_mo_trafo=ao_mo_trafo) + + # Get the orbital derivative coupling for NACs + orb_deriv = get_orbital_derivative_coupling(mol,ao_mo_trafo=ao_mo_trafo) + + # Nuclear part of the gradient + grad_nuc = grad.RHF(scf.RHF(mol)).grad_nuc() + + grad_elec_all = [] + nac_all = {} + nac_all_hfonly = {} + # Iterate over pairs of eigenstates + for i_state in range(nroots): + ci_i = fcivec[i_state] + + for j_state in range(nroots): + ci_2 = fcivec[j_state] + + # FCI 1- and 2-particle tRDMs + one_rdm, two_rdm = fcisolver.trans_rdm12(ci_i, ci_2, mol.nao, mol.nelec) + + if savemem: + grad_elec = get_grad_elec_OAO( + mol, one_rdm, two_rdm + ) + else: + # d\dR of subspace Hamiltonian + grad_elec = get_grad_elec_from_gradH( + one_rdm, two_rdm, h1_jac, h2_jac + ) + + # Energy gradients + if i_state == j_state: + grad_elec_all.append(grad_elec) + + # Nonadiabatic couplings + else: + nac_orb = np.einsum("ij,ijkl->kl",one_rdm, orb_deriv, optimize="optimal") + + nac_hf = grad_elec/(en[j_state]-en[i_state]) + nac_ij = nac_hf + nac_orb + nac_all[str(i_state)+str(j_state)] = nac_ij + nac_all_hfonly[str(i_state)+str(j_state)] = nac_hf + + # Add the nuclear contribution to gradient + grad_all = np.array(grad_elec_all) + grad_nuc + + return ( + en.real + mol.energy_nuc(), + grad_all, + nac_all, + nac_all_hfonly + ) + +# For testing, when it works - merge with previous one +def get_FCI_energy_with_grad_and_NAC_withsym(mol, fcisolver, nroots=1, hermitian=True, irrep_name=None): + """ + Calculates the potential energiesm its gradient w.r.t. nuclear positions of a + molecule and nonadiabatic couplings from full CI in SAO basis + + Args: + mol : pyscf.gto.Mole + The molecule object. + fcisolver : pyscf.fci solver + The solver object for full CI + nroots (optional): int + Number of states in the solver. + hermitian (optional): bool + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple (en, grad_all, nac_all, nac_all_hfonly) + A tuple containing the total potential energies, its gradients and NACs: + + en: ndarray(nroot,) + Total potential energies for both ground and excited states + + grad_all: list of ndarray(nat,3) + Gradients of multistate energies + + nac_all: dictionary of np.darray(nat,3) + Nonadiabatic coupling vectors between all states, + e.g. nac_all['02'] is NAC along ground state and 2nd excited state + + nac_all_hfonly: dictionary of ndarray(nat,3) + Hellman-Feynmann contribution to NACs + """ + # Loewdin orbitals + S = mol.intor("int1e_ovlp") + ao_mo_trafo = get_loewdin_trafo(S) + + if irrep_name == None: + # OAO basis + basis = ao_mo_trafo + + else: + # Use canonical for symmetry adapted + myhf = scf.RHF(mol) + _ = myhf.scf() + basis = myhf.mo_coeff + + # Construct h1 and h2 + h1 = np.linalg.multi_dot((basis.T, scf.hf.get_hcore(mol), basis)) + h2 = ao2mo.restore(1, ao2mo.kernel(mol, basis), basis.shape[1]) + + # Get FCI energies and wavefunctions in SAO basis + if irrep_name == None: + en, fcivec = fcisolver.kernel(h1, h2, mol.nao, mol.nelec) + + else: + orbsym = symm.label_orb_symm(mol, mol.irrep_id, mol.symm_orb, basis) + + en, fcivec = fcisolver.kernel(h1, h2, mol.nao, mol.nelec, orbsym=orbsym) + + u = np.einsum('ji,jk,kl->il',basis,S,ao_mo_trafo) + + fcivec = [transform_ci(fcivec_i,mol.nelec,u) for fcivec_i in fcivec] + + # Get the gradient of one and two-electron integrals before contracting onto + # rdms and trmds of different states + h1_jac, h2_jac = get_one_and_two_el_grad(mol,ao_mo_trafo=ao_mo_trafo) + + # Get the orbital derivative coupling for NACs + orb_deriv = get_orbital_derivative_coupling(mol,ao_mo_trafo=ao_mo_trafo) + + # Nuclear part of the gradient + grad_nuc = grad.RHF(scf.RHF(mol)).grad_nuc() + + grad_elec_all = [] + nac_all = {} + nac_all_hfonly = {} + # Iterate over pairs of eigenstates + for i_state in range(nroots): + ci_i = fcivec[i_state] + + for j_state in range(nroots): + ci_2 = fcivec[j_state] + + # FCI 1- and 2-particle tRDMs + one_rdm, two_rdm = fcisolver.trans_rdm12(ci_i, ci_2, mol.nao, mol.nelec) + + # d\dR of subspace Hamiltonian + grad_elec = get_grad_elec_from_gradH( + one_rdm, two_rdm, h1_jac, h2_jac + ) + + # Energy gradients + if i_state == j_state: + grad_elec_all.append(grad_elec) + + # Nonadiabatic couplings + else: + nac_orb = np.einsum("ij,ijkl->kl",one_rdm, orb_deriv, optimize="optimal") + + nac_hf = grad_elec/(en[j_state]-en[i_state]) + nac_ij = nac_hf + nac_orb + nac_all[str(i_state)+str(j_state)] = nac_ij + nac_all_hfonly[str(i_state)+str(j_state)] = nac_hf + + # Add the nuclear contribution to gradient + grad_all = np.array(grad_elec_all) + grad_nuc + + return ( + en.real + mol.energy_nuc(), + grad_all, + nac_all, + nac_all_hfonly + ) + +if __name__ == '__main__': + + from pyscf import gto + import pickle + + nstate = 2 # Max no of state to compute NACs up to + + natom = 8 + + check_spin = True + withMolcas = False + fix_singlet = True + + fix_sym = 'A1g' + #fix_sym = None + + if fix_sym == None: + mol_sym = False + else: + mol_sym = True + + test_range = np.linspace(0.8, 3.0,40) + + def get_mol(positions): + mol = gto.Mole() + + mol.build( + atom=[("H", pos) for pos in positions], + basis="sto-3g", + #basis="6-31g", + symmetry=mol_sym, + unit="Bohr", + verbose=0 + ) + + return mol + + mol_dummy = get_mol([(x, 0.0, 0.0) for x in test_range[0] * np.arange(natom)]) + # Set fci solver to be used + if fix_sym == None: + fcisolver = fci.direct_spin0.FCI() + else: + fcisolver = fci.direct_spin0_symm.FCI(mol_dummy) + fcisolver.wfnsym = fix_sym + + fcisolver.nroots = nstate+4 + + if fix_singlet: + fci.addons.fix_spin_(fcisolver,ss=0) # Fix spin + + # Prediction on test dataset and comparison against FCI results + fci_en = np.zeros([len(test_range),nstate+4]) + fci_nac = [] + fci_cionly_nac = [] + + for i, test_dist in enumerate(test_range): + print(i) + positions = [(x, 0.0, 0.0) for x in test_dist * np.arange(natom)] + + mol = get_mol(positions) + #h1, h2 = get_integrals(mol, get_basis(mol)) + + # Continuation + en_f, grad_f, nac_f, nac_f_hfonly = get_FCI_energy_with_grad_and_NAC_withsym( + mol, + fcisolver, + nroots=nstate+1, + irrep_name=fix_sym + ) + + fci_en[i,:] = en_f + fci_nac += [nac_f] + fci_cionly_nac += [nac_f_hfonly] + + + ####################################################################### + if withMolcas: + # Read NACs computed from openMOLCAS + fname = 'test_NACs.pkl' + with open(fname,'rb') as f: + test_NACs = pickle.load(f) + + # MOLCAS energies for comparison + fname = 'test_en.npy' + with open(fname,'rb') as f: + molcas_en = np.load(f) + + # Separate for geometry for plotting + molcas_nac = [] + for key in test_NACs.keys(): + all_NACs = test_NACs[key] + nac_i = {} + + for keyj in all_NACs.keys(): + # 0 - CI contribution (not divided by energy difference) + # 1 - CSF contribution + # 2 - Full NACs + ci_NAC = all_NACs[keyj][2] + nac_i[keyj] = ci_NAC + molcas_nac.append(nac_i) + + + # CI only part of the molcas NACs + molcas_cionly_nac = [] + for key in test_NACs.keys(): + all_NACs = test_NACs[key] + nac_i = {} + + for keyj in all_NACs.keys(): + ci_NAC = all_NACs[keyj][2] - all_NACs[keyj][1] + nac_i[keyj] = ci_NAC + molcas_cionly_nac.append(nac_i) + + ####################################################################### + # Plot NAC comparison + import matplotlib.pylab as plt + + fci_absh = {} + fci_cionly_absh = {} + molcas_absh = {} + molcas_hf_absh = {} + for istate in range(nstate+1): + for jstate in range(nstate+1): + if istate != jstate: + + st_label = str(istate)+str(jstate) + fci_absh[st_label] = np.array([np.abs(fci_nac[i][st_label]).sum() for i in range(len(test_range))]) + fci_cionly_absh[st_label] = np.array([np.abs(fci_cionly_nac[i][st_label]).sum() for i in range(len(test_range))]) + + if withMolcas: + molcas_absh[st_label] = np.array([np.abs(molcas_nac[i][st_label]).sum() for i in range(len(test_range))]) + molcas_hf_absh[st_label] = np.array([np.abs(molcas_cionly_nac[i][st_label]).sum() for i in range(len(test_range))]) + + # Colors + clr_st = {'01':'b', '10':'b', + '02':'r','20':'r', + '03':'pink','30':'pink', + '13':'y','31':'y', + '23':'violet','32':'violet', + '12':'g','21':'g'} + + labelsize = 15 + interfont=12 + + if withMolcas: + # Plot + fig, axes = plt.subplots(nrows=2,ncols=3,sharex=True, + figsize=[15,10],gridspec_kw={'hspace':0.,'wspace':0.1}, + height_ratios=[1,1]) + + axes[0][0].plot(test_range,molcas_en,'k',alpha=0.8,label=['CASSCF-molcas']+[None]*(molcas_en.shape[1]-1)) + axes[0][0].plot(test_range,fci_en,'r--',alpha=0.8,label=['FCI-pyscf']+[None]*(fci_en.shape[1]-1)) + + axes[0][1].plot(test_range,fci_en,alpha=0.8,label=['FCI']+[None]*(fci_en.shape[1]-1)) + #axes[0][1].plot(test_range,cont_en,'k',alpha=0.8) + #axes[0][1].plot(trainig_dists, np.array(train_en),'xr') + axes[0][2].plot(test_range, np.abs(fci_en[:,:nstate+1]-molcas_en[:,:nstate+1])) + + + for key, el in fci_absh.items(): + axes[1][0].plot(test_range,molcas_absh[key],label=key,c=clr_st[key]) + axes[1][1].plot(test_range,fci_absh[key],label=key,c=clr_st[key]) + axes[1][2].plot(test_range,np.abs(molcas_absh[key]-fci_absh[key]),label=key,c=clr_st[key]) + + axes[1][0].legend(loc='upper right',fontsize=interfont) + axes[0][0].legend(loc='upper right',fontsize=interfont) + + axes[0][0].set_title('CASSCF - MOLCAS',fontsize=labelsize) + axes[0][1].set_title('FCI - new implementation',fontsize=labelsize) + axes[0][2].set_title('Diff',fontsize=labelsize) + axes[1][0].set_ylabel(r'$||\mathbf{d}_{ij}||$ (a$_0$$^{-1}$)',fontsize=labelsize) + axes[0][0].set_ylabel(r'Energy (Hartree)',fontsize=labelsize) + + + axes[1][0].set_ylim(ymin=-0.2, ymax=min(10,axes[1][0].get_ylim()[1])) + + for axcol in axes: + for ax in axcol: + ax.yaxis.grid(color='gray', linestyle='dashed') + ax.xaxis.grid(color='gray', linestyle='dashed') + + plt.show() + + else: + # Plot + fig, axes = plt.subplots(nrows=1,ncols=2,sharex=True,#sharey='row', + figsize=[10,10],gridspec_kw={'hspace':0.,'wspace':0}, + height_ratios=[1]) + + axes[0].plot(test_range,fci_en,'r--',alpha=0.8,label=['FCI-pyscf']+[None]*(fci_en.shape[1]-1)) + #axes[0][1].plot(test_range,cont_en,'k',alpha=0.8) + #axes[0][1].plot(trainig_dists, np.array(train_en),'xr') + + for key, el in fci_absh.items(): + #axes[1][0].plot(test_range,molcas_absh[key],label=key,c=clr_st[key]) + axes[1].plot(test_range,fci_absh[key],label=key,c=clr_st[key]) + + axes[1].legend(loc='upper right',fontsize=interfont) + + #axes[1][0].set_ylabel(r'$||\mathbf{d}_{ij}||$ (a$_0$$^{-1}$)',fontsize=labelsize) + #axes[0][0].set_ylabel(r'Energy (Hartree)',fontsize=labelsize) + + + #axes[1][0].set_ylim(ymin=-0.2, ymax=min(10,axes[1][0].get_ylim()[1])) + + plt.show() + + ####################################################################### + diff --git a/evcont/MD_utils.py b/evcont/MD_utils.py index c613833..cdc8bda 100644 --- a/evcont/MD_utils.py +++ b/evcont/MD_utils.py @@ -67,7 +67,7 @@ def get_trajectory( init_veloc=None, hermitian=True, trajectory_output=None, - energy_output=None, + data_output=None, ): """ Helper function to compute an MD trajectory from eigenvector continuation with @@ -114,7 +114,7 @@ def get_trajectory( incore_anyway=True, frames=frames, trajectory_output=trajectory_output, - energy_output=energy_output, + data_output=data_output, verbose=0, ) myintegrator.run() @@ -195,7 +195,7 @@ def converge_EVCont_MD( EVCont_obj.two_rdm, steps=steps, trajectory_output=trajectory_out, - energy_output=en_out, + data_output=en_out, dt=dt, ) @@ -245,7 +245,7 @@ def converge_EVCont_MD( EVCont_obj.two_rdm, steps=steps, trajectory_output=trajectory_out, - energy_output=en_out, + data_output=en_out, dt=dt, ) else: @@ -436,7 +436,7 @@ def converge_EVCont_MD( EVCont_obj.two_rdm, steps=steps, trajectory_output=trajectory_out, - energy_output=en_out, + data_output=en_out, dt=dt, ) diff --git a/evcont/NAMD_utils.py b/evcont/NAMD_utils.py new file mode 100644 index 0000000..e6f0d9a --- /dev/null +++ b/evcont/NAMD_utils.py @@ -0,0 +1,1104 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Mon Dec 11 09:56:23 2023 + +Functions to read Newton-X trajectories and their properties as well as +active learning the training geometries for continuation + +@author: katalar +""" + +import numpy as np +import os +import sys +import subprocess +import glob + +from scipy.signal import find_peaks + +from evcont.electron_integral_utils import get_integrals, get_basis + +############################################################################## +# NX I/O FUNCTIONS +############################################################################## + +def read_population(fname,nstat): + population = [[] for i in range(nstat)] + with open(fname) as f: + for line in f: + if 'Population' in line: + splt = line.split() + population[int(splt[1])-1].append(float(splt[2])) + + return np.array(population).T + +def read_dyn(dynout_f, natm): + + pos_all = [] + read_pos = False + with open(dynout_f) as f: + pos_i = []; pos_line = 0 + for line in f: + + # Read positions + if read_pos: + pos_line += 1 + + pos_i.append([float(i) for i in line.split()[2:5]]) + + # Stop reading and reset if all atoms are read + if pos_line == natm: + # Save + pos_all.append(np.array(pos_i)) + # Reset + read_pos = False + pos_line = 0 + pos_i = [] # + + # Initiate reading + if 'geometry' in line: + read_pos = True + + return pos_all + +def read_traj(): + + geom_all = np.load('TEMP/traj_geom.npy') + + # Clean the duplicates after hopping + v, c = np.unique(geom_all, return_counts=True, axis=0) + + dupl = v[c > 1] + + for dupl_i in dupl: + indices = np.argwhere((dupl_i == geom_all).all(axis=2).all(axis=1)) + # If duplicates are consecutive, remove the duplicate + if abs(indices[1]-indices[0]) == 1: + geom_all = np.delete(geom_all, indices[1], axis=0) + + return geom_all + +def read_NX(path,pos='TEMP'): + + # Record current directory and change directory into path + cwd = os.getcwd() + os.chdir(path) + + # Read energies + en = np.loadtxt('RESULTS/en.dat') + + tprob_all = np.loadtxt('RESULTS/tprob',skiprows=1) + tprob = tprob_all[:,3:] + randi, substep, step = tprob_all[:,0], tprob_all[:,1], tprob_all[:,2] + + dynall = np.loadtxt('RESULTS/typeofdyn.log',usecols=[2,7]) + tim, pes = dynall[:,0], dynall[:,1] + + # Number of states + nstat = en[:,1:-2].shape[1] + + # Read populations + populations = read_population('RESULTS/sh.out',nstat) + + geom_f = 'geom' + + atom_f = [] + with open(geom_f,'r') as f: + for line in f.readlines(): + splt = line.split() + #sym, atomic no, xc, yc, zc, mass + atom_f.append((splt[0], np.array(splt[2:5],dtype=np.float64))) + #atom_f.append((splt[0], [float(i) for i in splt[2:5]])) + + natm = len(atom_f) + + if pos != 'TEMP': + pos_all = read_dyn('RESULTS/dyn.out',natm) + else: + pos_all = read_traj() + + # Return to original directory + os.chdir(cwd) + + return natm, nstat, en, [randi, substep, step, tprob], [tim, pes], populations, pos_all + +def write_model(overlap, one_rdm, two_rdm=None, vecs_lowrank=None, diagonal_lr=None, model_path=None): + """ + Write the intermediate data that will be used for predictions, namely: + - Overlap of training wavefunctions, S + - 1-el reduced transition density matrices of training wavefunctions + - 2-el reduced transition density matrices of training wavefunctions + OR low-rank vectors if using low-rank approximation + + Args: + overlap (ndarray) + one_rdm (ndarray) + two_rdm (ndarray, optional): Full 2-RDM if not using low-rank + vecs_lowrank (dict, optional): Low-rank vectors if using low-rank + diagonal_lr (ndarray, optional): Diagonal components if using low-rank + """ + import pickle + + # Allow caller to specify a custom model path; fall back to original location if not provided + if model_path is None: + model_path = 'sample/JOB_NAD' + + os.makedirs(model_path, exist_ok=True) + + np.save(os.path.join(model_path,'overlap_final.npy'),overlap) + np.save(os.path.join(model_path,'one_rdm_final.npy'),one_rdm) + + if two_rdm is not None: + # Full 2-RDM case + np.save(os.path.join(model_path,'two_rdm_final.npy'),two_rdm) + + if vecs_lowrank is not None: + # Low-rank case - save as pickle + with open(os.path.join(model_path,'lowrank_vecs.pkl'), 'wb') as f: + pickle.dump(vecs_lowrank, f, protocol=pickle.HIGHEST_PROTOCOL) + + if diagonal_lr is not None: + np.save(os.path.join(model_path,'diagonal_lr.npy'), diagonal_lr) + +def read_model(path): + """ + Read the intermediate data that will be used for predictions, namely: + - Overlap of training wavefunctions, S + - 1-el reduced transition density matrices of training wavefunctions + - 2-el reduced transition density matrices of training wavefunctions, + either as a full tensor (two_rdm_final.npy) or low-rank vectors (lowrank_vecs.pkl) + + Args: + path (str): path to the model files + + Returns: + overlap (ndarray) + one_rdm (ndarray) + two_rdm (ndarray or dict): Full 2-RDM or low-rank vectors + diagonals (ndarray or None): Diagonals if they exist, None otherwise + """ + import pickle + + overlap = np.load(os.path.join(path,'overlap_final.npy')) + one_rdm = np.load(os.path.join(path,'one_rdm_final.npy')) + + two_rdm_npy = os.path.join(path, 'two_rdm_final.npy') + two_rdm_pkl = os.path.join(path, 'lowrank_vecs.pkl') + + if os.path.exists(two_rdm_npy): + two_rdm = np.load(two_rdm_npy) + elif os.path.exists(two_rdm_pkl): + with open(two_rdm_pkl, 'rb') as f: + two_rdm = pickle.load(f) + else: + raise FileNotFoundError("Neither 'two_rdm_final.npy' nor 'lowrank_vecs.pkl' was found in the specified path.") + + # Try to load diagonals (optional, mainly for low-rank models) + diag_file = os.path.join(path, 'diagonal_lr.npy') + diagonals = None + if os.path.exists(diag_file): + diagonals = np.load(diag_file, allow_pickle=True) + + return overlap, one_rdm, two_rdm, diagonals + +def remove_model(path): + """ + Remove intermediate representation files from the path given IF they exist + """ + ov_path = os.path.join(path,'overlap_final.npy') + + if os.path.isfile(ov_path): + os.remove(ov_path) + os.remove(os.path.join(path,'one_rdm_final.npy')) + + # Remove two_rdm if exists + two_rdm_path = os.path.join(path,'two_rdm_final.npy') + if os.path.isfile(two_rdm_path): + os.remove(two_rdm_path) + + # Remove low-rank files if they exist + lowrank_path = os.path.join(path,'lowrank_vecs.pkl') + if os.path.isfile(lowrank_path): + os.remove(lowrank_path) + + diagonal_path = os.path.join(path,'diagonal_lr.npy') + if os.path.isfile(diagonal_path): + os.remove(diagonal_path) + +def clean_traj(): + """ + Clean the memory intensive files after each trajectory finishes + """ + existing_ind = [int(i.split('_')[-1].split('.')[0]) for i in glob.glob('ham_dist*')] + + if len(existing_ind) > 0: + max_it = max(existing_ind) + + for nit in range(max_it-1): + path_i = 'TRAJ_%i'%nit + + # Remove from TRAJ/DEBUG/TEMP + remove_model(os.path.join(path_i,'DEBUG','TEMP')) + + # Remove from TRAJ/JOB_NAD + remove_model(os.path.join(path_i,'JOB_NAD')) + + # Remove from TRAJ/TEMP/JOB* + remove_model(os.path.join(path_i,'TEMP','JOB_NAD')) + remove_model(os.path.join(path_i,'TEMP','JOB_AD')) + remove_model(os.path.join(path_i,'TEMP')) + + # Remove from TRAJ/INFO_RESTART/JOB_NAD + remove_model(os.path.join(path_i,'INFO_RESTART','JOB_NAD')) + +############################################################################## +# ACTIVE LEARNING +############################################################################## + +# Pseudocode + +#1 Set initial tRDMs and overlaps for the first run (e.g. starting geometry) +#2 Run NX dynamics trajectory for n steps with timestep dt +### Save the trajectory info; geometries and energies and anything else? +#3 Find geometry with farthrest ham distance and add it to the representation +#4 Repeat 2-3 until convergence + +# Convergence is achieved when the energies between consecutive iterations do +# not change for more than a threshold (mHa) for at 2(?) iterations + +# Details of step 2 +#i + +# Function input +#a + +def converge_NAMD_traj( + EVCont_obj, + init_mol, + steps=100, + dt=0.1, + nstat=3, + nstatdyn=2, + iseed=8, + convergence_thresh=1.0e-3, + nconv=2, + max_iter=100, + data_addition='weighted_highest_peak_ham', + nx_path=None, + reconverge_from_closest_hdist=False, + append_as_HPC_job=False, + run_command='sbatch $NX/moldyn.pl', + run_append_command='qsub append_states.sh', + compute_hamdist_during_traj=False, + compute_hamdist_as_HPC_job=False, + run_hamdist_command='qsub compute_hamdist.sh', + solver='CAS' + ): + """ + Converging eigenvector continuation training set for Newton-X nonadiabatic + dynamics trajectories. On-the-fly learning of which geometries to add by + converging the trajectory. + + Call and run from a separate directory that contains a directory + called 'sample' with sample NX input files. Otherwise, the calculation will be stuck + + Args: + EVCont_obj: + The data structure for the eigenvector continuation. + init_mol: + The initial molecule object. + steps (int): + Number of MD simulation steps. Default is 100. + dt (float): + Time step for the simulation. Default is 0.1 ns. + nstat (int): + Total number of states to be used in the NAMD simulation. + nstatdyn (int): + The state at which dynamics start from. + iseed (int): + Random seed of the NX trajectory. This achieves direct comparison. + Choose a value > 1 as NX has different meaning for iseed = 0,1. + convergence_thresh (float): + Energy convergence threshold to terminate the training. Default is 1.0e-3. + nconv (int): + Number of consecutive iterations required for convergence + max_iter (int): + Calculation is stopped after max_iter iterations if convergence + is not achieved. + data_addition (str): + Criterion for adding new data points. Can be "farthest_point_ham" (default), + in which case data is added based on electron integral difference, + "farthest_point", in which case data is added based on the farthest point + according to Euclidean distance, or "energy", in which case data is added + based on the energy difference. + nx_path (str): + Path for the 'bin' folder of the installed Newton-X code. If not defined, + the $NX environment variable will be expected to be predefined in the terminal. + Otherwise, the calculation will crash. + reconverge_from_closest_hdist (bool): + Whether to reconverge from the closest geometry based on Hamiltonian distance. Default is False. + run_command (str): + Terminal command to use for running NX. Based on the HPC, different + commands may be required. + run_append_command (str): + Terminal command to use for appending new states to the continuation object. + Based on the HPC, different commands may be required. + compute_hamdist_during_traj (bool): + Whether to compute Hamiltonian distances locally while the trajectory is running. + compute_hamdist_as_HPC_job (bool): + Whether to offload Hamiltonian distance computation to a separate HPC job. + run_hamdist_command (str): + Command used to submit/run the Hamiltonian distance job. + solver (str): + The type of solver used for the continuation object. Used to determine continuation object type + when appending new states as a HPC job. Default is 'CAS'. + + Returns: + 1 + + """ + + # Set Newton-X path + if nx_path is not None: + os.environ['NX'] = nx_path + + # Check if it is a restart calculation or a new calculation + existing_ind = [int(i.split('_')[-1].split('.')[0]) for i in glob.glob('ham_dist*')] + + # Current iteration of the convergence + if len(existing_ind) > 0: + nit = max(existing_ind) + 1 + else: + nit = 0 + + # Setup models directory - if it doesn't exist + OBJECT_CAN_BE_SAVED = False + if hasattr(EVCont_obj, 'load'): + OBJECT_CAN_BE_SAVED = True + + if not os.path.exists('iterative-models'): + os.mkdir('iterative-models') + + CONT_OBJ_FILENAME = f'iterative-models/continuation_object-{nit}.pkl' + + print('NAMD convergence - Starting iteration {}'.format(nit)) + + # Update the model from the last iteration or initialize one if it doesn't exist + if EVCont_obj.overlap is None: + + if append_as_HPC_job and OBJECT_CAN_BE_SAVED: + # Create an empty object and append the first geometry as an HPC job + EMPTY_OBJ_FILENAME = f'iterative-models/continuation_object-empty.pkl' + EVCont_obj.save(EMPTY_OBJ_FILENAME) + + EVCont_obj = run_append_states(init_mol.copy(), EMPTY_OBJ_FILENAME, CONT_OBJ_FILENAME, solver, run_append_command, quantel_tag='ref') + trn_geometries = [init_mol.atom_coords()] + np.save('trn_geometries.npy', trn_geometries) + + else: + EVCont_obj.append_to_rdms(init_mol.copy()) + trn_geometries = [init_mol.atom_coords()] + np.save('trn_geometries.npy', trn_geometries) + + # Optionally save the continuation object if possible + if OBJECT_CAN_BE_SAVED: + try: + print(f"Saving continuation object to {CONT_OBJ_FILENAME}") + EVCont_obj.save(CONT_OBJ_FILENAME) + except Exception as e: + print(f"Warning: Could not save continuation object: {e}") + + else: + # Read initial training geometries + trn_geometries = np.load('trn_geometries.npy') + + # Write model to iteration-specific directory and point evcont.in to it + model_dir = os.path.join('iterative-models', f'model_{nit}') # user-requested path pattern + abs_model_dir = os.path.abspath(model_dir) + os.makedirs(model_dir, exist_ok=True) + + # Update sample/JOB_NAD/evcont.in to reference this model directory via trdm_path + evcont_in_path = os.path.join('sample', 'JOB_NAD', 'evcont.in') + try: + if os.path.isfile(evcont_in_path): + with open(evcont_in_path, 'r') as f: + lines = f.readlines() + found = False + for i, line in enumerate(lines): + stripped = line.strip() + if stripped.startswith('trdm_path') or stripped.startswith('#trdm_path'): + lines[i] = f'trdm_path = {abs_model_dir}\n' + found = True + if not found: + # Prepend if no existing trdm_path directive + lines.insert(0, f'trdm_path = {abs_model_dir}\n') + with open(evcont_in_path, 'w') as f: + f.writelines(lines) + else: + print(f"Warning: evcont.in not found at {evcont_in_path}; cannot set trdm_path.") + except Exception as e: + print(f"Warning: could not modify evcont.in to set trdm_path: {e}") + + if EVCont_obj.lowrank: + EVCont_obj.vectorize_lowrank() + write_model(EVCont_obj.overlap, + EVCont_obj.one_rdm, + two_rdm=None, + vecs_lowrank=EVCont_obj.lowrank_vectorized, + diagonal_lr=EVCont_obj.diagonal_vectorized, + model_path=model_dir) + else: + write_model(EVCont_obj.overlap, + EVCont_obj.one_rdm, + two_rdm=EVCont_obj.two_rdm, + model_path=model_dir) + + ########################################################################### + # Setup and run NAMD trajectory + inp_par = [steps, dt, nstat, nstatdyn, iseed] + run_trajectory(nit, inp_par, run_command, + init_mol=init_mol, + trn_geometries=trn_geometries, + compute_hamdist_during_traj=compute_hamdist_during_traj) + + # Read output of current and previous trajectory + out_n = read_NX('TRAJ_%i'%nit,pos='dyn') + trajectory = out_n[-1] + ########################################################################### + # Check convergence + converged = False + + # Only check convergence after 0th iteration + if nit > 0: + out_prev = read_NX('TRAJ_%i'%(nit-1)) + + en = out_n[2][:,1:nstat+1] # Energies of all states + en_prev = out_prev[2][:,1:nstat+1] + + # If lens are different, e.g. when different commensurate timesteps are used + no_en = en.shape[0] + no_en_prev = en_prev.shape[0] + if no_en > no_en_prev: + en = en[::round(no_en/no_en_prev),:] + + elif no_en < no_en_prev: + en_prev = en_prev[::round(no_en_prev/no_en),:] + + # Mean energy difference across all states + en_diff = np.abs(en-en_prev).mean(axis=1) + np.savetxt('en_diff_{}.txt'.format(nit), en_diff) + + print('Current max(en_diff) is {:.4f} Ha'.format(max(en_diff))) + + if max(en_diff) < convergence_thresh: + converged = True + + # Check for previous iterations + for i in range(max(1,nit-nconv),nit): + # Check for the previous iteration as well + en_diff_prev = np.loadtxt('en_diff_{}.txt'.format(i)) + + if max(en_diff_prev) < convergence_thresh and converged: + converged = True + else: + converged = False + + if converged: + print('NAMD trajectory is converged within specified threshold of {:.4f} Ha for {} consecutive iterations'.format(convergence_thresh,nconv)) + + if nit >= max_iter: + print('Convergence was NOT achieved within the specified limit of {} iterations'.format(max_iter)) + # Stop the calculation + converged = True + + else: + # Set a very high energy difference, implying very far away from + # convergence, for later heuristics + en_diff = np.array([100.]) + + ########################################################################### + if converged: + return 1 + + else: + ###################################################################### + ##### SELECTION OF NEW TRAINING GEOMETRY + ###################################################################### + hamdist_file = f'ham_dist_{nit}.txt' + + # Check if distances were already computed during trajectory + if os.path.isfile(hamdist_file): + print(f"Loading Hamiltonian distances computed during trajectory from {hamdist_file}") + hamiltonian_distance_all = np.loadtxt(hamdist_file) + + # Compute via HPC job submission + elif compute_hamdist_as_HPC_job and append_as_HPC_job: + # Use in-memory trajectory instead of relying on TEMP/traj_geom.npy + hamiltonian_distance_all = run_hamdist_job( + nit, + init_mol, + run_hamdist_command, + trajectory + ) + + # Compute locally after trajectory completes + else: + if compute_hamdist_as_HPC_job and not append_as_HPC_job: + print("Warning: When compute_hamdist_as_HPC_job is True, append_as_HPC_job must also be True. Continuing to local hamdist computation.") + # Compute locally (also produce argmin indices for reuse) + hamiltonian_distance_all, argmins_all = hamiltonian_similarity_argmin(init_mol, trajectory, trn_geometries) + np.savetxt(hamdist_file, hamiltonian_distance_all) + try: + np.savetxt(f"ham_argmin_{nit}.txt", argmins_all, fmt='%d') + except Exception as e: + print(f"Warning: could not write argmin indices ham_argmin_{nit}.txt: {e}") + + ###################################################################### + ##### SELECTION OF NEW TRAINING GEOMETRY + ###################################################################### + addgeom_ind = select_active_learning_geometry(hamiltonian_distance_all, data_addition, en_diff, convergence_thresh) + + ###################################################################### + ##### AFTER SELECTION + ###################################################################### + # Add the new geometry to the training set + new_geom = trajectory[addgeom_ind] + new_trn_geometries = np.concatenate((trn_geometries,[new_geom])) + mol_new = init_mol.copy().set_geom_(new_geom) + + # Write + np.save('trn_geometries.npy',new_trn_geometries) + + ### Add to continuation + if reconverge_from_closest_hdist: + if hasattr(EVCont_obj, 'software') and EVCont_obj.software == 'quantel': + # Prefer using precomputed closest training indices to avoid recomputation + argmin_file = f"ham_argmin_{nit}.txt" + if os.path.isfile(argmin_file): + argmins_all = np.loadtxt(argmin_file, dtype=int) + closest_ind = int(argmins_all[addgeom_ind])+1 + closest_geom_tag = f'geom{closest_ind}' + else: + # No argmin file; fallback to original integral-based computation + oei_new, tei_new = get_integrals(mol_new, get_basis(mol_new)) + files = glob.glob(os.path.join(EVCont_obj.quantel_path, 'geom*')) + oei_trn = [] + tei_trn = [] + for f in files: + ind = int(f.split('geom')[-1]) + oei_trn.append(np.loadtxt(os.path.join(f,'oei.dat'))) + tei_trn.append(np.load(os.path.join(f,'tei.npy'))) + hamiltonian_distance_to_new = hamiltonian_distance( + oei_new, + tei_new, + np.array(oei_trn), + np.array(tei_trn) + ) + closest_ind = np.argmin(hamiltonian_distance_to_new) + closest_geom_tag = f'geom{files[closest_ind].split("geom")[-1]}' + + print('Re-converging from geometry closest in ham distance to the new geometry ({})'.format(closest_geom_tag)) + if append_as_HPC_job and OBJECT_CAN_BE_SAVED: + NEW_OBJ_FILENAME = f'iterative-models/continuation_object-{nit+1}.pkl' + EVCont_obj = run_append_states(mol_new, CONT_OBJ_FILENAME, NEW_OBJ_FILENAME, solver, run_append_command, quantel_tag=closest_geom_tag) + else: + EVCont_obj.append_to_rdms(mol_new, quantel_tag=closest_geom_tag) + + else: + print('Re-converging from closest ham distance is only implemented for Quantel software.') + sys.exit() + else: + if append_as_HPC_job and OBJECT_CAN_BE_SAVED: + NEW_OBJ_FILENAME = f'iterative-models/continuation_object-{nit+1}.pkl' + EVCont_obj = run_append_states(mol_new, CONT_OBJ_FILENAME, NEW_OBJ_FILENAME, solver, run_append_command, quantel_tag='ref') + else: + EVCont_obj.append_to_rdms(mol_new) + + # Save the continuation object if possible + if OBJECT_CAN_BE_SAVED and not append_as_HPC_job: + CONT_OBJ_FILENAME = f'iterative-models/continuation_object-{nit+1}.pkl' + try: + print(f"Saving continuation object to {CONT_OBJ_FILENAME}") + EVCont_obj.save(CONT_OBJ_FILENAME) + except Exception as e: + print(f"Warning: Could not save continuation object: {e}") + + # Go to next iteration + converge_NAMD_traj( + EVCont_obj, + init_mol, + steps=steps, + dt=dt, + nstat=nstat, + nstatdyn=nstatdyn, + iseed=iseed, + convergence_thresh=convergence_thresh, + nconv=nconv, + max_iter=max_iter, + data_addition=data_addition, + nx_path=nx_path, + reconverge_from_closest_hdist=reconverge_from_closest_hdist, + append_as_HPC_job=append_as_HPC_job, + run_command=run_command, + run_append_command=run_append_command, + compute_hamdist_during_traj=compute_hamdist_during_traj, + compute_hamdist_as_HPC_job=compute_hamdist_as_HPC_job, + run_hamdist_command=run_hamdist_command, + solver=solver + ) + +def run_append_states(mol, cont_obj_path, newcont_obj_path, solver, run_append_command, quantel_tag='None'): + """ + Run a job to append new geometries to the continuation object + """ + + # Save geometry + # Check if file exists and change name if necessary + geomfname = 'iterative-models/geom_0.xyz' + count = 0 + while os.path.isfile(geomfname): + count += 1 + geomfname = f'iterative-models/geom_{count}.xyz' + mol.tofile(geomfname) + + # Submit job with the correct arguments + os.system(f'{run_append_command} {cont_obj_path} {newcont_obj_path} {solver} {geomfname} {mol.basis} {quantel_tag}') + + # Check until the job finishes + poll_seconds = 20 # minimal polling interval + print(f"Waiting for appended continuation object file: {newcont_obj_path}") + while not os.path.isfile(newcont_obj_path): + os.system(f'sleep {poll_seconds}') + print(f"Detected file {newcont_obj_path}. Loading updated continuation object.") + + # Read in the appended continuation object + if solver == 'CAS': + try: + from evcont.CASCI_EVCont import CAS_EVCont_obj + cont_obj = CAS_EVCont_obj.load(newcont_obj_path) + except Exception as e: + print(f"Error loading appended continuation object: {e}") + sys.exit(1) + + return cont_obj + +def run_hamdist_job(nit, init_mol, run_hamdist_command, trajectory): + """Submit a job to compute Hamiltonian distances for a trajectory. + + Instead of relying on an existing TEMP/traj_geom.npy file, this function + receives the in-memory `trajectory` array, writes it to + TRAJ_/traj_geom.npy, and passes that path to the job script. + + Arguments passed to the external script (compute_hamdist.sh): + geom_file basis trajectory_npy trn_geometries_npy output_file [cache_prefix] + + Parameters: + nit (int): iteration / trajectory index + init_mol: Mole object (used for basis retrieval) + run_hamdist_command (str): submission command (e.g. 'qsub compute_hamdist.sh') + trajectory (ndarray): geometries from the just-completed NX trajectory + Returns: + distances (ndarray): Hamiltonian distances (loaded from ham_dist_.txt) + """ + + # Any geometry file to initialize the molecule - should already exist + geom_file = 'iterative-models/geom_0.xyz' + + # Write provided trajectory to a new file at top-level of TRAJ_ + traj_dir = f'TRAJ_{nit}' + if not os.path.isdir(traj_dir): + print(f"Error: trajectory directory {traj_dir} not found") + sys.exit(1) + traj_fname = os.path.join(traj_dir, 'traj_geom.npy') + try: + np.save(traj_fname, trajectory) + except Exception as e: + print(f"Error writing trajectory to {traj_fname}: {e}") + sys.exit(1) + + # Training geometries and hamdist output file + trn_fname = 'trn_geometries.npy' + hamdist_outfile = f'ham_dist_{nit}.txt' + + basis = init_mol.basis + cache_prefix = 'training_integrals' + + # Build command with positional arguments (compatible with HPC script) + cmd = f"{run_hamdist_command} {geom_file} {basis} {traj_fname} {trn_fname} {hamdist_outfile} {cache_prefix}" + + print(f"Submitting Hamiltonian distance job: {cmd}") + os.system(cmd) + + # Poll for output files + poll_seconds = 20 + print(f"Waiting for Hamiltonian distance output file: {hamdist_outfile}") + while not os.path.isfile(hamdist_outfile): + os.system(f'sleep {poll_seconds}') + + # Also wait (with a timeout) for the argmin file written by the job script + argmin_outfile = f'ham_argmin_{nit}.txt' + print(f"Waiting for closest training indices file: {argmin_outfile}") + max_polls = 30 # ~10 minutes + polls = 0 + while not os.path.isfile(argmin_outfile) and polls < max_polls: + os.system(f'sleep {poll_seconds}') + polls += 1 + if not os.path.isfile(argmin_outfile): + print(f"Warning: {argmin_outfile} not detected after waiting. Will proceed without it and fall back later if needed.") + + try: + distances = np.loadtxt(hamdist_outfile) + print(f"Loaded Hamiltonian distances from {hamdist_outfile}") + except Exception as e: + print(f"Error reading Hamiltonian distance output {hamdist_outfile}: {e}") + sys.exit(1) + return distances + +def run_trajectory(traj_ind, inp_par, run_command, init_mol=None, trn_geometries=None, compute_hamdist_during_traj=False): + """ + Run a Newton-X calculation with sample input files from the 'sample' directory. + + Optionally compute Hamiltonian distances on-the-fly during the trajectory. + + Args: + traj_ind (int): Trajectory index + inp_par (list): Input parameters [steps, dt, nstat, nstatdyn, iseed] + run_command (str): Command to run NX + init_mol (Mole, optional): Initial molecule object for hamdist computation + trn_geometries (ndarray, optional): Training geometries for hamdist computation + compute_hamdist_during_traj (bool): Whether to compute hamdist during trajectory + """ + # Copy sample files into a separate directory - named TRAJ_ind + new_path = 'TRAJ_%i'%traj_ind + os.system('cp -r sample %s'%new_path) + + # Clean scratch data from other trajectories + os.system('sleep 10') + clean_traj() + + print(run_command) + + # Record current working directory and change it to TRAJ_ind + cwd = os.getcwd() + os.chdir(new_path) + + # Check if the calculation has already finished + status = check_status() + if status == 'Success': + print(f"Trajectory {traj_ind} has already finished. Skipping run.") + os.chdir(cwd) + return + + # Setup input files + # TODO - for now, they remain the same as sample + + # Run + os.system(run_command) + os.system('sleep 100') + + # Wait until calculation finishes or crashes + # Optionally compute Hamiltonian distances during the trajectory + hamdist_outfile = None + prev_traj_length = 0 + hamiltonian_distances = [] + + if compute_hamdist_during_traj and init_mol is not None and trn_geometries is not None: + hamdist_outfile = os.path.join(cwd, f'ham_dist_{traj_ind}.txt') + print(f"Will compute Hamiltonian distances on-the-fly and save to {hamdist_outfile}") + + while True: + os.system('sleep 100') + + status = check_status() + + if status == 'Error': + print('NX has crushed - check the end of output at DEBUG/runnx.error') + sys.exit() + elif status == 'Success': + break + + # Compute Hamiltonian distances for new trajectory points + if compute_hamdist_during_traj and status == 'Running': + traj_file = 'TEMP/traj_geom.npy' + if os.path.isfile(traj_file): + try: + current_traj = np.load(traj_file) + current_length = len(current_traj) + + # Only compute if new geometries have been added + if current_length > prev_traj_length: + print(f"Computing Hamiltonian distances for new geometries (total: {current_length})...") + + # Compute distances for new geometries only + new_geoms = current_traj[prev_traj_length:current_length] + new_distances = hamiltonian_similarity(init_mol, new_geoms, trn_geometries) + hamiltonian_distances.extend(new_distances.tolist()) + + # Save updated distances to file + np.savetxt(hamdist_outfile, np.array(hamiltonian_distances)) + print(f" Computed distances for geometries {prev_traj_length} to {current_length-1}") + + prev_traj_length = current_length + + except Exception as e: + print(f"Warning: Could not compute Hamiltonian distances during trajectory: {e}") + + # TODO: Add an early exit condition for an early peak detection and killing the trajectory + # e.g. if a peak in hamdist is detected, and that peak is higher than previous hamdist peaks + # Need to add a function to kill the NX job; copy bits of addgeom_ind from convergence loop; etc. + # Need to iteratively check for convergence as well since if max(en_diff) < convergence_thresh early on, + # the trajectory needs to continue. + + # Final computation if any geometries were missed + if compute_hamdist_during_traj and hamdist_outfile is not None: + traj_file = 'TEMP/traj_geom.npy' + if os.path.isfile(traj_file): + try: + final_traj = np.load(traj_file) + final_length = len(final_traj) + + if final_length > prev_traj_length: + print(f"Computing final Hamiltonian distances (total: {final_length})...") + new_geoms = final_traj[prev_traj_length:final_length] + new_distances = hamiltonian_similarity(init_mol, new_geoms, trn_geometries) + hamiltonian_distances.extend(new_distances.tolist()) + np.savetxt(hamdist_outfile, np.array(hamiltonian_distances)) + print(f" Final distances computed for geometries {prev_traj_length} to {final_length-1}") + except Exception as e: + print(f"Warning: Could not compute final Hamiltonian distances: {e}") + + # Return to original directory for next iteration + os.chdir(cwd) + +def check_status(): + """ + Check the status Newton-X calculation + + Returns: + + """ + # Read last line + try: + line = str(subprocess.check_output(['tail', '-1', 'RESULTS/nx.log'])) + line2 = str(subprocess.check_output(['tail', '-2', 'RESULTS/nx.log'])) + except: + # If file hasn't been created yet + return 'Starting' + + # Return the status + if 'DEBUG/runnx.error' in line2: + return 'Error' + elif 'NEWTON-X ends here' in line: + return 'Success' + else: + return 'Running' + +def select_active_learning_geometry(hamiltonian_distance_all, data_addition, en_diff=None, convergence_thresh=None, exponent=0.5): + """ + Select which geometry to add to the training set based on the hamiltonian + distance metric and specified data addition method. + + Args: + hamiltonian_distance_all (ndarray): + Array of Hamiltonian distances for all geometries in the trajectory. + data_addition (str): + Criterion for adding new data points. Can be + "farthest_point_ham": in which case geometry that is furthest away from the training set is added, + "first_peak_ham": in which case data is added based on the first peak in Hamiltonian distance, + "weighted_highest_peak_ham": in which case data is added based on a weighted peak selection in + Hamiltonian distance, + "variable_weight_peak_ham": in which case the exponent of the weighting function is varied + based on current convergence. + en_diff (ndarray): + Array of energy differences for all geometries between the two previous trajectory iterations. + convergence_thresh (float): + Energy convergence threshold to terminate the training. + + Returns: + addgeom_ind (int): + Index of the geometry to be added to the training set. + """ + # Disregards peaks with hamdist less than this threshold + # (make this a parameter later on, also depends on norb) + threshold = 0.01 + + # Find the index of the geometry with maximum H_dist + addgeom_ind = np.argmax(hamiltonian_distance_all) + + # Find which new geometry to add to the training set + # Select the geometry that's the furtherst in hamiltonian distance + if data_addition == "farthest_point_ham": + + # Find the index of the geometry with maximum H_dist + addgeom_ind = np.argmax(hamiltonian_distance_all) + + # Select the temporally first peak in the hamiltonian distance + elif data_addition == "first_peak_ham": + + # If no peak is found, choose the index with largest ham distance + addgeom_ind = np.argmax(hamiltonian_distance_all) + + # Find all peaks + peaks = find_peaks(hamiltonian_distance_all)[0] + + # If peaks above a threshold exist, choose that over farthest + if len(peaks) > 0: + ind_above_thr = np.where(hamiltonian_distance_all[peaks] > threshold) + if len(ind_above_thr[0]) > 0: + addgeom_ind = peaks[ind_above_thr][0] + + elif data_addition in ["weighted_highest_peak_ham", "variable_weight_peak_ham"]: + # Exponent of the time penalty function + # (0 - chooses max, -->inf chooses 1st peak) + if data_addition == "weighted_highest_peak_ham": + exponent = exponent + else: + # Variable exponent based on current convergence (still experimental) + scaling = 0.01 + scaled_exp = scaling * en_diff.max()/convergence_thresh + exponent = min(max(scaled_exp, 0.), 5.) # Limit between 0 and 5 + + # If no peak is found, choose the index with largest ham distance + addgeom_ind = np.argmax(hamiltonian_distance_all) + + # Find all peaks + peaks = find_peaks(hamiltonian_distance_all)[0] + + # Add the max point to the peaks as a possible selection geometry + if addgeom_ind not in peaks: + peaks = np.append(peaks, addgeom_ind) + + # If peaks above a threshold exist, choose that over farthest + ind_above_thr = np.where((hamiltonian_distance_all[peaks] > threshold) & (peaks > 0)) + + if len(ind_above_thr[0]) > 0: + # Peaks above threshold + peaks_above_thr = peaks[ind_above_thr] + + # penalty function ranging from 0 (favourable) to 1 (unfavourable) + penalty = (peaks_above_thr/len(hamiltonian_distance_all))**exponent + + step_weighted_hamdist = hamiltonian_distance_all[peaks_above_thr]/penalty + + addgeom_ind = peaks_above_thr[np.argmax(step_weighted_hamdist)] + + else: + print('The data_addition method {} is not implemented.'.format(data_addition)) + sys.exit() + + return addgeom_ind + +def hamiltonian_distance(oei1, tei1, oei2, tei2): + """ + Calculate a distance metric between two sets of one-electron and two-electron integrals. + + Parameters: + oei1, tei1 : numpy.ndarray + One- and two-electron integrals for a geometry. + oei2, tei2 : numpy.ndarray + One- and two-electron integrals for a different geometry + (can be an array for different geometries in which case an array of distances is returned). + + Returns: + float + A scalar distance metric quantifying the difference between the two Hamiltonians. + """ + nbasis = oei1.shape[-1] + + N1 = nbasis**2 # number of 1e integral elements + N2 = nbasis**4 # number of 2e integral elements + + distance = ( + np.sum(abs(oei1 - oei2)**2, axis=(-1, -2)) / N1 + + 0.5 * np.sum(abs(tei1 - tei2)**2, axis=(-1, -2, -3, -4)) / N2 + ) + # Rescale for the upcoming heuristics (e.g. peak detection, etc.) + return distance*1000 + +def hamiltonian_similarity(init_mol, trajectory, trn_geometries): + """ + Compute the minimum Hamiltonian distance of a trajectory to a set + of training geometries + """ + # Initialize hamiltonians of the training set + h1_trn = np.zeros((len(trn_geometries), init_mol.nao, init_mol.nao)) + h2_trn = np.zeros( + ( + len(trn_geometries), + init_mol.nao, + init_mol.nao, + init_mol.nao, + init_mol.nao, + ) + ) + + # Compute 1- and 2-electron integrals for all training geometries + for j, trn_geom in enumerate(trn_geometries): + mol = init_mol.copy().set_geom_(trn_geom) + h1, h2 = get_integrals(mol, get_basis(mol)) + h1_trn[j] = h1 + h2_trn[j] = h2 + + # Compute min Hamiltonian distance to the training geometries for the new traj + min_dist_l = [] + for j, geometry in enumerate(trajectory): + mol = init_mol.copy().set_geom_(geometry) + h1, h2 = get_integrals(mol, get_basis(mol)) + + distance = hamiltonian_distance(h1, h2, h1_trn, h2_trn) + min_dist = np.min(distance) + min_dist_l += [min_dist] + + return np.array(min_dist_l) + +def hamiltonian_similarity_argmin(init_mol, trajectory, trn_geometries): + """ + Compute the minimum Hamiltonian distance and the argmin training index + of a trajectory to a set of training geometries. + + Returns: + min_distances (ndarray) + argmins (ndarray[int]) + """ + # Initialize hamiltonians of the training set + h1_trn = np.zeros((len(trn_geometries), init_mol.nao, init_mol.nao)) + h2_trn = np.zeros( + ( + len(trn_geometries), + init_mol.nao, + init_mol.nao, + init_mol.nao, + init_mol.nao, + ) + ) + + # Compute 1- and 2-electron integrals for all training geometries + for j, trn_geom in enumerate(trn_geometries): + mol = init_mol.copy().set_geom_(trn_geom) + h1, h2 = get_integrals(mol, get_basis(mol)) + h1_trn[j] = h1 + h2_trn[j] = h2 + + # Compute min Hamiltonian distance and argmin training index + min_dist_l = [] + argmin_l = [] + for j, geometry in enumerate(trajectory): + mol = init_mol.copy().set_geom_(geometry) + h1, h2 = get_integrals(mol, get_basis(mol)) + + distance = hamiltonian_distance(h1, h2, h1_trn, h2_trn) + min_dist = np.min(distance) + min_dist_l.append(min_dist) + argmin_l.append(int(np.argmin(distance))) + + return np.array(min_dist_l), np.array(argmin_l, dtype=int) + + +if __name__ == '__main__': + print('yes') + + + + + diff --git a/evcont/_append_to_rdms.py b/evcont/_append_to_rdms.py new file mode 100644 index 0000000..59ed2d1 --- /dev/null +++ b/evcont/_append_to_rdms.py @@ -0,0 +1,100 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +_append_to_rdms.py + +Helper script to append new geometries to an existing continuation object. +Can be submitted as an HPC job within an active learning workflow (as used in NAMD_utils.py). + +Inputs (via CLI args): + CONT_OBJ_FILENAME Path to existing continuation object pickle file + NEW_OBJ_FILENAME Path to save updated continuation object + CONT_OBJ_TYPE Type of continuation object ('CAS', 'FCI', etc.) + GEOM_FILENAME XYZ geometry file for the new geometry + BASIS Basis set string (e.g., '6-31g') + TAG Optional tag for reading in previous solutions (e.g., 'ref', 'geom0') + +Operation: + 1. Load existing continuation object from pickle file. + 2. Build molecule from geometry file + basis. + 3. Append new geometry to continuation object (with quantum chemistry calculation). + 4. Save updated continuation object to new file. + +Exit codes: + 0 success + 1 failure (exception) + +Author: Kemal Atalar +Date: November 2025 +""" + +import os +import sys +import numpy as np +from pyscf import gto + + +def main(): + # Parse positional arguments (compatible with HPC job script like append_states.sh) + if len(sys.argv) < 7: + print("Usage: python _append_to_rdms.py ") + return 1 + + cont_obj_filename = sys.argv[1] + new_obj_filename = sys.argv[2] + cont_obj_type = sys.argv[3] + geom_filename = sys.argv[4] + basis = sys.argv[5] + tag = sys.argv[6] + + try: + # Import continuation object class based on type + if cont_obj_type == 'CAS': + from evcont.CASCI_EVCont import CAS_EVCont_obj as CONT_OBJ_CLASS + else: + print(f"Error: Unsupported continuation object type: {cont_obj_type}") + return 1 + + # Load continuation object + print(f"Loading continuation object from {cont_obj_filename}") + try: + cont_obj = CONT_OBJ_CLASS.load(cont_obj_filename) + except Exception as e: + print(f"Error: Failed to load continuation object from {cont_obj_filename}") + print(f" {e}") + return 1 + + # Build molecule from geometry file + print(f"Building molecule from {geom_filename} with basis {basis}") + mol = gto.Mole() + mol.build( + atom=geom_filename, + basis=basis, + unit='Angstrom', + verbose=0 + ) + + # Append the new geometry to the continuation object + print(f"Appending new geometry to continuation object (tag: {tag})...") + if hasattr(cont_obj, 'software') and cont_obj.software == 'quantel': + cont_obj.append_to_rdms(mol, quantel_tag=tag) + else: + cont_obj.append_to_rdms(mol) + + # Save updated continuation object + print(f"Saving updated continuation object to {new_obj_filename}") + cont_obj.save(new_obj_filename) + + print(f"Successfully appended geometry from {geom_filename}") + print("DONE") + return 0 + + except Exception as e: + print(f"Error appending geometry: {e}") + import traceback + traceback.print_exc() + return 1 + +if __name__ == '__main__': + sys.exit(main()) + diff --git a/evcont/_compute_hamdist.py b/evcont/_compute_hamdist.py new file mode 100644 index 0000000..e30d0f5 --- /dev/null +++ b/evcont/_compute_hamdist.py @@ -0,0 +1,213 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +_compute_hamdist.py + +HPC/offline helper to compute Hamiltonian distances between a trajectory and +existing training geometries for EVCont active learning. + +Inputs (via CLI args): + GEOM_FILENAME XYZ geometry file (initial geometry to build molecule) + BASIS Basis set string (e.g., '6-31g') + TRAJECTORY_NPY Numpy file containing trajectory geometries (nsteps, natm, 3) + TRN_GEOMETRIES_NPY Numpy file containing training geometries (ntrn, natm, 3) + OUTPUT_FILE Output text filename for ham distances + [CACHE_PREFIX] Optional prefix for cached training integrals (default: training_integrals) + [--force-recompute] Optional flag to ignore existing cached integrals + +Operation: + 1. Build molecule from geometry file + basis (like _append_to_rdms.py). + 2. Load training geometries and trajectory geometries from numpy files. + 3. (Re)compute one- and two-electron integrals for training geometries unless + cached arrays exist (_h1.npy / _h2.npy). + 4. Loop over trajectory geometries, compute integrals, evaluate minimum + Hamiltonian distance to training set (using evcont.NAMD_utils.hamiltonian_distance). + 5. Write distances to plain text file (np.savetxt). + +The script deliberately avoids modifying any continuation object; it is purely +computational. Caching prevents repeated recomputation of training integrals +across iterations. + +Exit codes: + 0 success + 1 failure (exception) + +Author: Kemal Atalar +Date: November 2025 +""" + +import os +import sys +import numpy as np + +from pyscf import gto +from evcont.electron_integral_utils import get_integrals, get_basis +from evcont.NAMD_utils import hamiltonian_distance + + +def build_molecule(geom_filename, basis): + """Build PySCF molecule from geometry file and basis string.""" + mol = gto.Mole() + mol.build( + atom=geom_filename, + basis=basis, + unit='Angstrom', + verbose=0 + ) + return mol + + +def compute_training_integrals(init_mol, trn_geometries, cache_prefix, force): + """Return (h1_trn, h2_trn) arrays, appending to cache when training set grows. + + Behavior: + - If cache exists and --force-recompute is not given, reuse cached entries and + only compute integrals for newly added training geometries, then extend the cache. + - If basis/nao changes (shape mismatch), ignore cache and recompute fully. + """ + h1_cache = f"{cache_prefix}_h1.npy" + h2_cache = f"{cache_prefix}_h2.npy" + + ntrn = len(trn_geometries) + nao = init_mol.nao + + cached_len = 0 + h1_cached = None + h2_cached = None + + if os.path.isfile(h1_cache) and os.path.isfile(h2_cache) and not force: + try: + h1_cached = np.load(h1_cache) + h2_cached = np.load(h2_cache) + # Validate cached shapes (other than length dim) + if h1_cached.shape[1:] != (nao, nao) or h2_cached.shape[1:] != (nao, nao, nao, nao): + print("Warning: cached integrals have incompatible shape (basis/nao changed); ignoring cache and recomputing.") + h1_cached = None + h2_cached = None + else: + cached_len = h1_cached.shape[0] + if cached_len == ntrn: + print(f"Loaded cached training integrals from {h1_cache}, {h2_cache}") + return h1_cached, h2_cached + elif cached_len > ntrn: + # Truncate to match current trn_geometries length + print(f"Warning: cache has more entries ({cached_len}) than trn_geometries ({ntrn}); truncating cache in-memory.") + h1_cached = h1_cached[:ntrn] + h2_cached = h2_cached[:ntrn] + cached_len = ntrn + else: + print(f"Extending cache from {cached_len} -> {ntrn} training geometries...") + except Exception as e: + print(f"Warning: failed loading cached integrals; recomputing. Reason: {e}") + h1_cached = None + h2_cached = None + + # Allocate full arrays + h1_trn = np.zeros((ntrn, nao, nao)) + h2_trn = np.zeros((ntrn, nao, nao, nao, nao)) + + # If we have valid cached data, copy it into the front + if h1_cached is not None and h2_cached is not None and cached_len > 0: + h1_trn[:cached_len] = h1_cached + h2_trn[:cached_len] = h2_cached + + # Compute missing integrals (either full or only the tail) + start_i = cached_len if (h1_cached is not None and h2_cached is not None) else 0 + todo = ntrn - start_i + if todo > 0: + print(f"Computing integrals for {todo} training geometries (from index {start_i})...") + for i in range(start_i, ntrn): + geom = trn_geometries[i] + mol = init_mol.copy().set_geom_(geom) + h1, h2 = get_integrals(mol, get_basis(mol)) + h1_trn[i] = h1 + h2_trn[i] = h2 + else: + print("No new training geometries to compute; using cached integrals.") + + # Persist/extend cache for future iterations + try: + np.save(h1_cache, h1_trn) + np.save(h2_cache, h2_trn) + if todo > 0: + print(f"Updated cached training integrals to {h1_cache}, {h2_cache}") + else: + print(f"Cache verified: {h1_cache}, {h2_cache}") + except Exception as e: + print(f"Warning: failed to save training integrals cache: {e}") + + return h1_trn, h2_trn + + +def compute_distances(init_mol, trajectory, h1_trn, h2_trn): + """Compute minimum Hamiltonian distance and argmin training index for each trajectory geometry. + + Returns: + distances (ndarray): min distance per traj geometry + argmins (ndarray[int]): index of closest training geometry per traj geometry + """ + distances = [] + argmins = [] + print(f"Computing Hamiltonian distances for {len(trajectory)} trajectory geometries...") + for geom in trajectory: + mol = init_mol.copy().set_geom_(geom) + h1, h2 = get_integrals(mol, get_basis(mol)) + d_all = hamiltonian_distance(h1, h2, h1_trn, h2_trn) + distances.append(np.min(d_all)) + argmins.append(int(np.argmin(d_all))) + return np.array(distances), np.array(argmins, dtype=int) + + +def main(): + # Parse positional arguments (compatible with HPC job script like append_states.sh) + if len(sys.argv) < 6: + print("Usage: python _compute_hamdist.py [cache_prefix] [--force-recompute]") + return 1 + + geom_filename = sys.argv[1] + basis = sys.argv[2] + trajectory_npy = sys.argv[3] + trn_geometries_npy = sys.argv[4] + output_file = sys.argv[5] + + # Optional cache prefix (default: training_integrals) + cache_prefix = sys.argv[6] if len(sys.argv) > 6 and not sys.argv[6].startswith('--') else 'training_integrals' + + # Optional force recompute flag + force_recompute = '--force-recompute' in sys.argv + + try: + print(f"Building molecule from {geom_filename} with basis {basis}") + init_mol = build_molecule(geom_filename, basis) + + print(f"Loading trajectory from {trajectory_npy}") + trajectory = np.load(trajectory_npy) + + print(f"Loading training geometries from {trn_geometries_npy}") + trn_geometries = np.load(trn_geometries_npy) + + h1_trn, h2_trn = compute_training_integrals(init_mol, trn_geometries, cache_prefix, force_recompute) + distances, argmins = compute_distances(init_mol, trajectory, h1_trn, h2_trn) + + # Write distances + np.savetxt(output_file, distances) + print(f"Hamiltonian distances written to {output_file}") + + # Also write argmin indices to a sibling file for downstream use + argmin_file = output_file.replace('ham_dist', 'ham_argmin') + try: + np.savetxt(argmin_file, argmins, fmt='%d') + print(f"Closest training indices written to {argmin_file}") + except Exception as e: + print(f"Warning: could not write argmin indices to {argmin_file}: {e}") + print("DONE") + return 0 + except Exception as e: + print(f"Error computing Hamiltonian distances: {e}") + import traceback + traceback.print_exc() + return 1 + + +if __name__ == '__main__': + sys.exit(main()) diff --git a/evcont/ab_initio_eigenvector_continuation.py b/evcont/ab_initio_eigenvector_continuation.py index ed4e008..f7f2594 100644 --- a/evcont/ab_initio_eigenvector_continuation.py +++ b/evcont/ab_initio_eigenvector_continuation.py @@ -1,5 +1,7 @@ import numpy as np +import sys + from scipy.linalg import eigh, eig from evcont.electron_integral_utils import ( @@ -8,6 +10,9 @@ compress_electron_exchange_symmetry, ) +from evcont.low_rank_utils import lowrank_hamiltonian + +from pyscf.lib import safe_eigh def approximate_ground_state(h1, h2, one_RDM, two_RDM, S, hermitian=True): """ @@ -72,7 +77,8 @@ def approximate_ground_state(h1, h2, one_RDM, two_RDM, S, hermitian=True): if hermitian is True: # Solve the generalized eigenvalue problem for Hermitian Hamiltonian - vals, vecs = eigh(H, S) + #vals, vecs = eigh(H, S) + vals, vecs, _ = safe_eigh(H, S) else: # Solve the generalized eigenvalue problem for non-Hermitian Hamiltonian vals, vecs = eig(H, S) @@ -89,8 +95,78 @@ def approximate_ground_state(h1, h2, one_RDM, two_RDM, S, hermitian=True): return en_approx, gs_approx +def solve_subspace(H, S, nroots=1, hermitian=True, lindep=1e-12): + """ + Diagonalize the subspace Hamiltonian + """ + + if hermitian is True: + # Solve the generalized eigenvalue problem for Hermitian Hamiltonian + #vals, vecs = eigh(H, S) + vals, vecs, _ = safe_eigh(H, S, lindep=lindep) + + else: + # Solve the generalized eigenvalue problem for non-Hermitian Hamiltonian + vals, vecs = eig(H, S) + + # Filter out imaginary eigenvalues + valid_vals = abs(vals.imag) < 1.0e-5 + + # Make sure nroots isn't higher than available eigenstates + assert vals[valid_vals].shape[0] >= nroots + + # Find the index of the minimum GS eigenvalue + argroots = np.argsort(vals[valid_vals].real)[:nroots] -def approximate_multistate(h1, h2, one_RDM, two_RDM, S, nroots=1, hermitian=True): + # Get the energy approximation and ground state approximation + en_approx = vals[valid_vals][argroots].real + evec_approx = vecs[:, valid_vals][:, argroots].real.T + + return en_approx, evec_approx + +def approximate_multistate_lowrank(mol, one_RDM, lowrank_vecs, cum_diagonal, S, + nroots=1, Jdiag_only=True, sao_diag=True, + hermitian=True, density_fit=True, df_basis=None, + lindep=1e-4): + """ + Returns multiple approximate electronic states from solving the generalised + eigenvalue problem defined via the one- and two-body transition RDMs. + + Args: + h1 (np.ndarray): One-electron integrals. + h2 (np.ndarray): Two-electron integrals. + one_RDM (np.ndarray): One-body t-RDM. + two_RDM (np.ndarray): Two-body t-RDM. + nroots: Number of states to be solved. Default is 1, the ground state. + S (np.ndarray): Overlap matrix. + hermitian (bool, optional): + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + Tuple[float, np.ndarray]: Energy approximation and ground state approximation. + """ + + # Calculate the Hamiltonian matrix + H = lowrank_hamiltonian(mol, one_RDM, S, lowrank_vecs, cum_diagonal, + density_fit=density_fit, df_basis=df_basis, + Jdiag_only=Jdiag_only, sao_diag=sao_diag) + + #print(' Hamiltonian') + #print(H.tolist()) + #plt.figure() + #sb.heatmap(H, annot=True) + #plt.show() + + #plt.savefig('hamiltonian_%i.png'%(rdm_computation)) + #print(' Overlap') + #print(S.tolist()) + + #print(H, S) + en_approx, evec_approx = solve_subspace(H, S, nroots=nroots, hermitian=hermitian, lindep=lindep) + + return en_approx, evec_approx + +def approximate_multistate(h1, h2, one_RDM, two_RDM, S, nroots=1, hermitian=True, lindep=1e-12): """ Returns multiple approximate electronic states from solving the generalised eigenvalue problem defined via the one- and two-body transition RDMs. @@ -151,26 +227,63 @@ def approximate_multistate(h1, h2, one_RDM, two_RDM, S, nroots=1, hermitian=True else: assert False + + en_approx, evec_approx = solve_subspace(H, S, nroots=nroots, hermitian=hermitian, lindep=lindep) - if hermitian is True: - # Solve the generalized eigenvalue problem for Hermitian Hamiltonian - vals, vecs = eigh(H, S) - else: - # Solve the generalized eigenvalue problem for non-Hermitian Hamiltonian - vals, vecs = eig(H, S) + return en_approx, evec_approx - # Filter out imaginary eigenvalues - valid_vals = abs(vals.imag) < 1.0e-5 - # Make sure nroots isn't higher than available eigenstates - assert vals[valid_vals].shape[0] >= nroots +#import matplotlib.pylab as plt +#import seaborn as sb - # Find the index of the minimum GS eigenvalue - argroots = np.argsort(vals[valid_vals].real)[:nroots] +def approximate_multistate_otf(h1, h2, one_RDM=None, two_RDM=None, S=None, otf_hamiltonian=None, nroots=1, hermitian=True, mol=None): + """ + Returns multiple approximate electronic states from solving the generalised + eigenvalue problem defined via the one- and two-body transition RDMs. - # Get the energy approximation and ground state approximation - en_approx = vals[valid_vals][argroots].real - evec_approx = vecs[:, valid_vals][:, argroots].real.T + Args: + h1 (np.ndarray): One-electron integrals. + h2 (np.ndarray): Two-electron integrals. + one_RDM (np.ndarray): One-body t-RDM. + two_RDM (np.ndarray): Two-body t-RDM. + nroots: Number of states to be solved. Default is 1, the ground state. + S (np.ndarray): Overlap matrix. + hermitian (bool, optional): + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + Tuple[float, np.ndarray]: Energy approximation and ground state approximation. + """ + + # Check that either RDMs or the otf generator is given + if (one_RDM is None or two_RDM is None or S is None) and otf_hamiltonian is None: + print('Error in approximate_multistate_otf: Neither RDMs or OTF generator is given') + sys.exit() + + rdm_computation = False + if otf_hamiltonian is None: + rdm_computation = True + + if rdm_computation: + # Calculate the Hamiltonian matrix + H = np.einsum("ijkl,kl->ij", one_RDM, h1, optimize="optimal") + 0.5 * np.einsum( + "ijklmn,klmn->ij", two_RDM, h2, optimize="optimal" + ) + else: + if mol is None: + H, S = otf_hamiltonian(h1, h2) + else: + H, S = otf_hamiltonian(mol) + + #print(' Hamiltonian') + #print(H) + #plt.figure() + #sb.heatmap(H, annot=True) + #plt.savefig('hamiltonian_%i.png'%(rdm_computation)) + #print(' Overlap') + #print(S) + + en_approx, evec_approx = solve_subspace(H, S, nroots=nroots, hermitian=hermitian) return en_approx, evec_approx @@ -210,8 +323,7 @@ def approximate_ground_state_OAO(mol, one_RDM, two_RDM, S, hermitian=True): return total_energy, vec - -def approximate_multistate_OAO(mol, one_RDM, two_RDM, S, nroots=1, hermitian=True): +def approximate_multistate_OAO(mol, one_RDM, two_RDM, S, nroots=1, hermitian=True, lindep=1e-12): """ This function approximates multiple state energies and wavefunctions of a given molecule from an eigenvector continuation with t-RDMS and the overlap matrix S. @@ -241,10 +353,90 @@ def approximate_multistate_OAO(mol, one_RDM, two_RDM, S, nroots=1, hermitian=Tru # Approximate the ground state energy and wavefunction in projected subspace en, vec = approximate_multistate( - h1, h2, one_RDM, two_RDM, S, nroots=nroots, hermitian=hermitian + h1, h2, one_RDM, two_RDM, S, nroots=nroots, hermitian=hermitian, lindep=lindep ) # Calculate the total energy by adding the nuclear repulsion energy total_energy = en.real + mol.energy_nuc() return total_energy, vec + + +def approximate_multistate_lowrank_OAO(mol, one_RDM, lowrank_vecs, cum_diagonal, S, + nroots=1, Jdiag_only=True, sao_diag=True, + hermitian=True, density_fit=True, df_basis=None, + lindep=1e-4): + """ + This function approximates multiple state energies and wavefunctions of a given + molecule from an eigenvector continuation with t-RDMS and the overlap matrix S. + + Args: + mol (Molecule): The molecule object representing the system. + one_RDM (ndarray): The one-electron t-RDM. + two_RDM (ndarray): The two-electron t-RDM. + S (ndarray): The overlap matrix. + nroots: Number of states to be solved. Default is 1, the ground state. + hermitian (bool, optional): + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple: A tuple containing the approximate ground state energy and the + ground state wavefunction in the learning subspace as a vector of expansion + coefficients. + + """ + # Construct h1 and h2 + #h1, h2 = get_integrals(mol, get_basis(mol)) + + # Approximate the ground state energy and wavefunction in projected subspace + #en, vec = approximate_multistate(h1, h2, one_RDM, two_RDM, S, nroots=nroots, hermitian=hermitian) + en, vec = approximate_multistate_lowrank(mol, one_RDM, lowrank_vecs, cum_diagonal, S, + nroots=nroots, hermitian=hermitian, + density_fit=density_fit, df_basis=df_basis, + Jdiag_only=Jdiag_only, sao_diag=sao_diag, lindep=lindep) + # Calculate the total energy by adding the nuclear repulsion energy + total_energy = en.real + mol.energy_nuc() + + return total_energy, vec + + +def approximate_multistate_otf_OAO(mol, one_RDM=None, two_RDM=None, S=None, otf_hamiltonian=None, nroots=1, hermitian=True, passmol=False): + """ + This function approximates multiple state energies and wavefunctions of a given + molecule from an eigenvector continuation with t-RDMS and the overlap matrix S. + + Args: + mol (Molecule): The molecule object representing the system. + one_RDM (ndarray): The one-electron t-RDM. + two_RDM (ndarray): The two-electron t-RDM. + S (ndarray): The overlap matrix. + nroots: Number of states to be solved. Default is 1, the ground state. + hermitian (bool, optional): + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple: A tuple containing the approximate ground state energy and the + ground state wavefunction in the learning subspace as a vector of expansion + coefficients. + + """ + + # Check that either RDMs or the otf generator is given + if (one_RDM is None or two_RDM is None or S is None) and otf_hamiltonian is None: + print('Error in approximate_multistate_otf_OAO: Neither RDMs or OTF generator is given') + sys.exit() + + # Construct h1 and h2 + if not passmol: + h1, h2 = get_integrals(mol, get_basis(mol)) + + # Approximate the ground state energy and wavefunction in projected subspace + en, vec = approximate_multistate_otf(h1, h2, one_RDM, two_RDM, S, otf_hamiltonian, nroots=nroots, hermitian=hermitian) + else: + # Approximate the ground state energy and wavefunction in projected subspace + en, vec = approximate_multistate_otf(None, None, one_RDM, two_RDM, S, otf_hamiltonian, nroots=nroots, hermitian=hermitian, mol=mol) + + # Calculate the total energy by adding the nuclear repulsion energy + total_energy = en.real + mol.energy_nuc() + + return total_energy, vec diff --git a/evcont/ab_initio_gradients_loewdin.py b/evcont/ab_initio_gradients_loewdin.py index 7bc1544..35c2055 100644 --- a/evcont/ab_initio_gradients_loewdin.py +++ b/evcont/ab_initio_gradients_loewdin.py @@ -1,14 +1,28 @@ import numpy as np -from pyscf import scf, ao2mo, grad +from pyscf import scf, ao2mo, grad, lib -from evcont.ab_initio_eigenvector_continuation import approximate_ground_state +from evcont.ab_initio_eigenvector_continuation import ( + approximate_ground_state, + approximate_multistate, + solve_subspace +) from evcont.electron_integral_utils import ( get_loewdin_trafo, restore_electron_exchange_symmetry, + get_df_integrals +) + +#from evcont.low_rank_utils_mpi import ( +from evcont.low_rank_utils import ( + get_jk_builds, unstack_tril ) +from evcont.logging_utils import logger, log_time, timeit + +import sys + def get_overlap_grad(mol): """ @@ -37,14 +51,16 @@ def get_overlap_grad(mol): # Transpose the return value to match the desired ordering of indices return np.transpose(deriv, (2, 3, 1, 0)) - -def loewdin_trafo_grad(overlap_mat): +@timeit +def loewdin_trafo_grad(overlap_mat, degeneracy_precision=7): """ Calculate the gradient of the Loewdin transformation. This also takes care of degeneracies by resorting to degenerate perturbation theory. Parameters: overlap_mat (np.ndarray): Matrix representing the overlap between atomic orbitals. + degeneracy_precision (int): repurposed to set an absolute tolerance tol = 10**(-degeneracy_precision) + for treating eigenvalues as degenerate. Lower values make the tolerance smaller. Returns: np.ndarray: Gradient of the Loewdin transformation. @@ -52,20 +68,28 @@ def loewdin_trafo_grad(overlap_mat): vals, vecs = np.linalg.eigh(overlap_mat) - rounded_vals = np.round(vals, decimals=5) - degenerate_vals = np.unique(rounded_vals) + # Determine degeneracy by small differences in eigenvalues rather than rounding + tol = 10.0 ** (-degeneracy_precision) + n = vals.shape[0] + assigned = np.zeros(n, dtype=bool) + degenerate_groups = [] + for i in range(n): + if assigned[i]: + continue + same = np.where(np.abs(vals - vals[i]) <= tol)[0] + assigned[same] = True + degenerate_groups.append(same) U_full = np.zeros((*overlap_mat.shape, *overlap_mat.shape)) degenerate_subspace = np.zeros(overlap_mat.shape, dtype=bool) # Take care of degeneracies - for val in degenerate_vals: - degenerate_ids = (np.argwhere(rounded_vals == val)).flatten() + for degenerate_ids in degenerate_groups: subspace = vecs[:, degenerate_ids] - V_projected = 0.5 * np.einsum( + V_projected = 0.5 * lib.einsum( "ai,bj->abij", subspace, subspace - ) + 0.5 * np.einsum("bi,aj->abij", subspace, subspace) + ) + 0.5 * lib.einsum("bi,aj->abij", subspace, subspace, optimize='optimal') # Get rotation to diagonalise V in degenerate subspace _, U = np.linalg.eigh(V_projected) @@ -79,11 +103,11 @@ def loewdin_trafo_grad(overlap_mat): ] = U degenerate_subspace[np.ix_(degenerate_ids, degenerate_ids)] = True - vecs_rotated = np.einsum("ij,abjk->abik", vecs, U_full) + vecs_rotated = lib.einsum("ij,abjk->abik", vecs, U_full, optimize='optimal') - Vji = 0.5 * np.einsum( - "abai,abbj->abij", vecs_rotated, vecs_rotated - ) + 0.5 * np.einsum("abbi,abaj->abij", vecs_rotated, vecs_rotated) + Vji = 0.5 * lib.einsum( + "abai,abbj->abij", vecs_rotated, vecs_rotated, optimize='optimal' + ) + 0.5 * lib.einsum("abbi,abaj->abij", vecs_rotated, vecs_rotated, optimize='optimal') Zji = np.zeros((*overlap_mat.shape, *overlap_mat.shape)) Zji[:, :, ~degenerate_subspace] = ( @@ -91,7 +115,7 @@ def loewdin_trafo_grad(overlap_mat): / ((vals - np.expand_dims(vals, -1))[~degenerate_subspace]) ) - dvecs = np.einsum("abij,abjk->abik", vecs_rotated, Zji) + dvecs = lib.einsum("abij,abjk->abik", vecs_rotated, Zji, optimize='optimal') dvals = Vji[:, :, np.arange(Vji.shape[2]), np.arange(Vji.shape[3])] transformed_vals = np.where(vals > 1.0e-15, 1 / np.sqrt(vals), 0.0) @@ -99,19 +123,20 @@ def loewdin_trafo_grad(overlap_mat): np.where(vals > 1.0e-15, -(0.5 / np.sqrt(vals) ** 3), 0.0) * dvals ) dS = ( - np.einsum("abij, abkj->abik", dvecs * transformed_vals, vecs_rotated) - + np.einsum( + lib.einsum("abij, abkj->abik", dvecs * transformed_vals, vecs_rotated, optimize='optimal') + + lib.einsum( "abij, abkj->abik", vecs_rotated * np.expand_dims(d_transformed_vals, axis=-2), vecs_rotated, + optimize='optimal' ) - + np.einsum("abij, abkj->abik", vecs_rotated * transformed_vals, dvecs) + + lib.einsum("abij, abkj->abik", vecs_rotated * transformed_vals, dvecs, optimize='optimal') ) # Transpose the return value to match the desired ordering of indices return np.transpose(dS, (2, 3, 0, 1)) - +@timeit def get_derivative_ao_mo_trafo(mol): """ Calculates the derivatives of the atomic orbital to molecular orbital @@ -125,7 +150,7 @@ def get_derivative_ao_mo_trafo(mol): """ overlap_grad = get_overlap_grad(mol) - trafo_grad = np.einsum( + trafo_grad = lib.einsum( "ijkl, ijmn->klmn", loewdin_trafo_grad(mol.intor("int1e_ovlp")), overlap_grad, @@ -133,7 +158,14 @@ def get_derivative_ao_mo_trafo(mol): return trafo_grad - +def fix_gauge(vec): + """ + Make so that the first element is always positive + """ + for vec_i in vec: + idx = np.unravel_index(np.argmax(np.abs(vec_i.real)),vec_i.shape) + vec_i *= -np.sign(vec_i[idx]) + def get_one_el_grad_ao(mol): """ Calculate the one-electron integral derivatives in the AO basis. @@ -142,7 +174,8 @@ def get_one_el_grad_ao(mol): mol (pyscf.gto.Mole): The molecular system. Returns: - np.ndarray: The one-electron integral derivatives in the AO basis. + np.ndarray(nel,nel,nat,3): + The one-electron integral derivatives in the AO basis. """ hcore_gen = grad.RHF(scf.RHF(mol)).hcore_generator() @@ -152,7 +185,7 @@ def get_one_el_grad_ao(mol): return np.transpose(return_val, (2, 3, 0, 1)) -def get_one_el_grad(mol, ao_mo_trafo=None, ao_mo_trafo_grad=None): +def get_one_el_grad(mol, h1_ao=None, ao_mo_trafo=None, ao_mo_trafo_grad=None): """ Calculate the gradient of the one-electron integrals with respect to nuclear coordinates. @@ -167,26 +200,26 @@ def get_one_el_grad(mol, ao_mo_trafo=None, ao_mo_trafo_grad=None): orbitals. Returns: - numpy.ndarray + numpy.ndarray(nel,nel,nat,3): The gradient of the one-electron integrals. """ if ao_mo_trafo is None: ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) - h1_ao = scf.hf.get_hcore(mol) + if h1_ao is None: + h1_ao = scf.hf.get_hcore(mol) if ao_mo_trafo_grad is None: ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) h1_grad_ao = get_one_el_grad_ao(mol) - h1_grad = np.einsum("ijkl,im,mn->jnkl", ao_mo_trafo_grad, h1_ao, ao_mo_trafo) + h1_grad = lib.einsum("ijkl,im,mn->jnkl", ao_mo_trafo_grad, h1_ao, ao_mo_trafo) h1_grad += np.swapaxes(h1_grad, 0, 1) - h1_grad += np.einsum("ij,iklm,kn->jnlm", ao_mo_trafo, h1_grad_ao, ao_mo_trafo) + h1_grad += lib.einsum("ij,iklm,kn->jnlm", ao_mo_trafo, h1_grad_ao, ao_mo_trafo) return h1_grad - def two_el_grad(h2_ao, two_rdm, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_slices): """ Calculate the two-electron integral gradient. @@ -207,7 +240,7 @@ def two_el_grad(h2_ao, two_rdm, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_ """ - two_el_contraction = np.einsum( + two_el_contraction = lib.einsum( "ijkl,abcd,aimn,bj,ck,dl->mn", two_rdm + np.transpose(two_rdm, (1, 0, 2, 3)) @@ -221,7 +254,7 @@ def two_el_grad(h2_ao, two_rdm, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_ optimize="optimal", ) - two_rdm_ao = np.einsum( + two_rdm_ao = lib.einsum( "ijkl,ai,bj,ck,dl->abcd", two_rdm, ao_mo_trafo, @@ -231,8 +264,8 @@ def two_el_grad(h2_ao, two_rdm, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_ optimize="optimal", ) - two_el_contraction_from_grad = np.einsum( - "nmbcd,abcd->nma", + two_el_contraction_from_grad_traced = lib.einsum( + "nmbcd,mbcd->nm", h2_ao_deriv, two_rdm_ao + np.transpose(two_rdm_ao, (1, 0, 3, 2)) @@ -241,15 +274,15 @@ def two_el_grad(h2_ao, two_rdm, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_ optimize="optimal", ) - h2_grad_ao_b = np.zeros((3, len(atm_slices), two_rdm.shape[0], two_rdm.shape[1])) + h2_grad_ao_sum = np.zeros((len(atm_slices),3)) for i, slice in enumerate(atm_slices): # Subtract the gradient contribution from the contraction part - h2_grad_ao_b[:, i, slice[0] : slice[1], :] -= two_el_contraction_from_grad[ - :, slice[0] : slice[1], : - ] + h2_grad_ao_sum[i,:] -= two_el_contraction_from_grad_traced[ + :, slice[0] : slice[1] + ].sum(axis=1) # Return the two-electron integral gradient - return two_el_contraction + np.einsum("nmbb->mn", h2_grad_ao_b) + return h2_grad_ao_sum + two_el_contraction def get_grad_elec_OAO(mol, one_rdm, two_rdm, ao_mo_trafo=None, ao_mo_trafo_grad=None): @@ -298,11 +331,65 @@ def get_grad_elec_OAO(mol, one_rdm, two_rdm, ao_mo_trafo=None, ao_mo_trafo_grad= ) grad_elec = ( - np.einsum("ij,ijkl->kl", one_rdm, h1_jac, optimize="optimal") + lib.einsum("ij,ijkl->kl", one_rdm, h1_jac, optimize="optimal") + + 0.5 * two_el_gradient + ) + + return grad_elec + + +def get_grad_elec_OAO_customERI(mol, h2_ao, h2_ao_deriv, one_rdm, two_rdm, ao_mo_trafo=None, ao_mo_trafo_grad=None): + """ + Calculates the gradient of the electronic energy based on one- and two-rdms + in the OAO. + + Args: + mol (object): Molecule object. + one_rdm (ndarray): One-electron reduced density matrix. + two_rdm (ndarray): Two-electron reduced density matrix. + ao_mo_trafo (ndarray, optional): + AO to MO transformation matrix. Is computed if not provided. + ao_mo_trafo_grad (ndarray, optional): + Gradient of AO to MO transformation matrix. Is computed if not provided. + + Returns: + ndarray: Electronic gradient. + """ + + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + if ao_mo_trafo_grad is None: + ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) + + h1_jac = get_one_el_grad( + mol, ao_mo_trafo=ao_mo_trafo, ao_mo_trafo_grad=ao_mo_trafo_grad + ) + + #h2_ao = mol.intor("int2e") + #h2_ao_deriv = mol.intor("int2e_ip1", comp=3) + + two_el_gradient = two_el_grad( + h2_ao, + two_rdm, + ao_mo_trafo, + ao_mo_trafo_grad, + h2_ao_deriv, + tuple( + [ + (mol.aoslice_by_atom()[i][2], mol.aoslice_by_atom()[i][3]) + for i in range(mol.natm) + ] + ), + ) + + grad_elec = ( + lib.einsum("ij,ijkl->kl", one_rdm, h1_jac, optimize="optimal") + 0.5 * two_el_gradient ) return grad_elec + #return 0.5 * two_el_gradient def get_energy_with_grad( @@ -377,3 +464,1341 @@ def get_energy_with_grad( en.real + mol.energy_nuc(), grad_elec + grad.RHF(scf.RHF(mol)).grad_nuc(), ) + +def get_energy_with_grad_cpuefficient(mol, one_RDM, two_RDM, S, hermitian=True, return_density_matrices=False): + """ + Calculates the potential energy and its gradient w.r.t. nuclear positions of a + molecule from the eigenvector continuation. + + Args: + mol : pyscf.gto.Mole + The molecule object. + one_RDM : numpy.ndarray + The one-electron t-RDM. + two_RDM : numpy.ndarray + The two-electron t-RDM. + S : numpy.ndarray + The overlap matrix. + hermitian (bool, optional): + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple + A tuple containing the total potential energy and its gradient. + """ + # Construct h1 and h2 + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + h1 = np.linalg.multi_dot((ao_mo_trafo.T, scf.hf.get_hcore(mol), ao_mo_trafo)) + h2 = ao2mo.restore(1, ao2mo.kernel(mol, ao_mo_trafo), mol.nao) + + en, vec = approximate_ground_state(h1, h2, one_RDM, two_RDM, S, hermitian=hermitian) + + # Get the gradient of one and two-electron integrals before contracting onto + # rdms of different states + h1_jac, h2_jac = get_one_and_two_el_grad(mol,ao_mo_trafo=ao_mo_trafo) + + one_rdm_predicted = np.tensordot(np.outer(vec, vec), one_RDM, axes=2) + + if len(two_RDM.shape) == 2 or len(two_RDM.shape) == 5: + # symmetry in data points + eigenvec_mat = 2 * np.outer(vec, vec) + + np.fill_diagonal(eigenvec_mat, 0.5 * np.diag(eigenvec_mat)) + + two_rdm_predicted = np.tensordot( + eigenvec_mat[np.tril_indices(len(vec))], two_RDM, axes=1 + ) + + else: + two_rdm_predicted = np.tensordot(np.outer(vec, vec), two_RDM, axes=2) + + if len(two_rdm_predicted.shape) != 4: + two_rdm_predicted = restore_electron_exchange_symmetry( + two_rdm_predicted, mol.nao + ) + + grad_elec = get_grad_elec_from_gradH( + one_rdm_predicted, two_rdm_predicted, h1_jac, h2_jac + ) + + return ( + en.real + mol.energy_nuc(), + grad_elec + grad.RHF(scf.RHF(mol)).grad_nuc(), + ) + +############################################################ +# NEW MULTISTATE - SEPERATED FROM REST FOR TESTING +############################################################ + +def get_two_el_grad(h2_ao, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_slices): + """ + Calculate the two-electron integral gradient. + + Args: + h2_ao (np.ndarray): Two-electron integrals in atomic orbital basis. + two_rdm (np.ndarray): Two-electron reduced density matrix. + ao_mo_trafo (np.ndarray): + Transformation matrix from atomic orbital to molecular orbital basis. + ao_mo_trafo_grad (np.ndarray): Gradient of the transformation matrix. + h2_ao_deriv (np.ndarray): + Derivative of the two-electron integrals with respect to nuclear + coordinates. + atm_slices (list): List of atom index slices. + + Returns: + np.ndarray: The two-electron gradient. + + """ + + two_el_contraction_ao = lib.einsum( + "abcd,aimn,bj,ck,dl->ijklmn", + h2_ao + h2_ao.transpose(1,0,3,2) + + h2_ao.transpose(2,3,0,1) + h2_ao.transpose(3,2,0,1), + ao_mo_trafo_grad, + ao_mo_trafo, + ao_mo_trafo, + ao_mo_trafo, + optimize="optimal", + ) + + # h2_ao + h2_ao.transpose(1,0,2,3) + # + h2_ao.transpose(3,2,1,0) + h2_ao.transpose(2,3,0,1), + + h2_grad_ao_sum = np.zeros((h2_ao.shape[0], h2_ao.shape[1], h2_ao.shape[2], h2_ao.shape[3], len(atm_slices),3)) + for i, slice in enumerate(atm_slices): + + two_el_ao = lib.einsum( + "nmbcd,mi,bj,ck,dl->ijkln", + h2_ao_deriv[:,slice[0] : slice[1],:,:,:], + ao_mo_trafo[slice[0] : slice[1],:], + ao_mo_trafo, + ao_mo_trafo, + ao_mo_trafo, + optimize="optimal", + ) + + # Subtract the gradient contribution from the contraction part + h2_grad_ao_sum[:,:,:,:,i,:] -= two_el_ao + two_el_ao.transpose(1, 0, 2, 3, 4) \ + + two_el_ao.transpose(3, 2, 1, 0, 4) + two_el_ao.transpose(2, 3, 0, 1, 4) + + h2_grad = two_el_contraction_ao + h2_grad_ao_sum + + # Return the two-electron integral gradient + return h2_grad + +def get_two_el_grad_new(h2_ao, ao_mo_trafo, ao_mo_trafo_grad, h2_ao_deriv, atm_slices): + """ + Calculate the two-electron integral gradient. + + Args: + h2_ao (np.ndarray): Two-electron integrals in atomic orbital basis. + two_rdm (np.ndarray): Two-electron reduced density matrix. + ao_mo_trafo (np.ndarray): + Transformation matrix from atomic orbital to molecular orbital basis. + ao_mo_trafo_grad (np.ndarray): Gradient of the transformation matrix. + h2_ao_deriv (np.ndarray): + Derivative of the two-electron integrals with respect to nuclear + coordinates. + atm_slices (list): List of atom index slices. + + Returns: + np.ndarray: The two-electron gradient. + + """ + + two_el_contraction_ao = lib.einsum( + "abcd,aimn,bj,ck,dl->ijklmn", + h2_ao, + ao_mo_trafo_grad, + ao_mo_trafo, + ao_mo_trafo, + ao_mo_trafo, + optimize="optimal", + ) + + two_el_contraction_ao += \ + np.transpose(two_el_contraction_ao,(1, 0, 2, 3,4,5)) +\ + np.transpose(two_el_contraction_ao,(3, 2, 1, 0,4,5)) +\ + np.transpose(two_el_contraction_ao,(2, 3, 0, 1,4,5)) + + h2_grad_ao_sum = np.zeros((h2_ao.shape[0], h2_ao.shape[1], h2_ao.shape[2], h2_ao.shape[3], len(atm_slices),3)) + for i, slice in enumerate(atm_slices): + + two_el_ao = lib.einsum( + "nmbcd,mi,bj,ck,dl->ijkln", + h2_ao_deriv[:,slice[0] : slice[1],:,:,:], + ao_mo_trafo[slice[0] : slice[1],:], + ao_mo_trafo, + ao_mo_trafo, + ao_mo_trafo, + optimize="optimal", + ) + + # Subtract the gradient contribution from the contraction part + h2_grad_ao_sum[:,:,:,:,i,:] -= two_el_ao + two_el_ao.transpose(1, 0, 2, 3, 4) \ + + two_el_ao.transpose(3, 2, 1, 0, 4) + two_el_ao.transpose(2, 3, 0, 1, 4) + + h2_grad = two_el_contraction_ao + h2_grad_ao_sum + + # Return the two-electron integral gradient + return h2_grad + + +def get_one_and_two_el_grad(mol,ao_mo_trafo=None, ao_mo_trafo_grad=None): + """ + Calculates the gradient of the one- and two-electron integrals + in the OAO. + + Args: + mol (object): Molecule object. + ao_mo_trafo (ndarray, optional): + AO to MO transformation matrix. Is computed if not provided. + ao_mo_trafo_grad (ndarray, optional): + Gradient of AO to MO transformation matrix. Is computed if not provided. + + Returns: + tuple of np.ndarray: + One- and two-electron gradients. + """ + + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + if ao_mo_trafo_grad is None: + ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) + + h1_jac = get_one_el_grad( + mol, ao_mo_trafo=ao_mo_trafo, ao_mo_trafo_grad=ao_mo_trafo_grad + ) + + h2_ao = mol.intor("int2e") + h2_ao_deriv = mol.intor("int2e_ip1", comp=3) + + h2_jac = get_two_el_grad( + h2_ao, + ao_mo_trafo, + ao_mo_trafo_grad, + h2_ao_deriv, + tuple( + [ + (mol.aoslice_by_atom()[i][2], mol.aoslice_by_atom()[i][3]) + for i in range(mol.natm) + ] + ), + ) + + + return (h1_jac, h2_jac) + + +def get_grad_elec_from_gradH(one_rdm, two_rdm, h1_jac, h2_jac): + """ + Calculates the continuation gradient from the one- and two-electron + integrals derivatives. + + Args: + mol (object): Molecule object. + one_RDM : numpy.ndarray + The one-electron t-RDM. + two_RDM : numpy.ndarray + The two-electron t-RDM. + h1_jac : numpy.ndarray + The gradient of the one-electron integrals. + h2_jac : numpy.ndarray + The gradient of the two-electron integrals. + + Returns: + ndarray: Electronic gradient. + """ + + + two_el_gradient = lib.einsum( + "ijkl,ijklmn->mn", + two_rdm, + h2_jac, + optimize="optimal", + ) + + grad_elec = ( + lib.einsum("ij,ijkl->kl", one_rdm, h1_jac, optimize="optimal") + + 0.5 * two_el_gradient + ) + + return grad_elec + + +def get_orbital_derivative_coupling(mol,ao_mo_trafo=None, ao_mo_trafo_grad=None, ovlp=None): + """ + For orbital contribution to nonadiabatic coupling vectors; + < Phi_a | d/dR Phi_b> where Phi are MOs + + Args: + mol (object): Molecule object. + ao_mo_trafo (ndarray, optional): + AO to MO transformation matrix. Is computed if not provided. + ao_mo_trafo_grad (ndarray, optional): + Gradient of AO to MO transformation matrix. Is computed if not provided. + + Returns: + tuple of np.ndarray (nbasis, nbasis, nat,3): + Orbital derivative coupling (to be contracted with 1-trdm). + """ + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + if ao_mo_trafo_grad is None: + ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) + + if ovlp is None: + ovlp = mol.intor("int1e_ovlp") + + # Contraction of the SAO transformation derivative + # \sum_{ik} dC_{ij}/dR * C_{kl} * s_{ik} + trafo_deriv_contraction = lib.einsum("ijAx,ik,kl->jlAx",ao_mo_trafo_grad, ovlp, ao_mo_trafo,optimize="optimal") + + # Orbital derivative contraction + # \sum_{ik} C_{ij} * C_{kl} * < bas_i | d bas_j/dR > + atm_slices = tuple( + [ + (mol.aoslice_by_atom()[i][2], mol.aoslice_by_atom()[i][3]) + for i in range(mol.natm) + ] + ) + + #""" + orb_deriv_contraction = np.zeros((mol.nao, mol.nao, mol.natm, 3)) + ipovlp = mol.intor("int1e_ipovlp", comp=3) # shape (3, nao, nao) + + for i, (start, stop) in enumerate(atm_slices): + tmp = lib.einsum("ij,xik,kl->jlx", + ao_mo_trafo[start:stop], + ipovlp[:, start:stop, :], + ao_mo_trafo, + optimize=True, + ) # shape (nb, nb, 3) + orb_deriv_contraction[:, :, i, :] -= tmp + + """ + deriv_ov = np.zeros((len(atm_slices),3,mol.nao,mol.nao)) + for i, slice in enumerate(atm_slices): + deriv_ov[i,:, slice[0] : slice[1], :] -= mol.intor("int1e_ipovlp")[:, slice[0] : slice[1], :] + + orb_deriv_contraction = lib.einsum("ij,Axik,kl->jlAx",ao_mo_trafo, deriv_ov, ao_mo_trafo,optimize="optimal") + """ + return trafo_deriv_contraction + orb_deriv_contraction + +def get_multistate_energy_with_grad(mol, one_RDM, two_RDM, S, nroots=1, hermitian=True, return_density_matrices=False): + """ + Calculates the potential energy and its gradient w.r.t. nuclear positions of a + molecule from the eigenvector continuation. + + Args: + mol : pyscf.gto.Mole + The molecule object. + one_RDM : numpy.ndarray + The one-electron t-RDM. + two_RDM (np.ndarray): Two-body t-RDM. Can have different shape depending on whether + symmetry-compressed representations are used or not: + No symmetries: shape(two_RDM) = (Ntrn, Ntrn, Norb, Norb, Norb, Norb) + Data symmetry only: shape(two_RDM) = (Ntrn * (Ntrn + 1)/2, Norb, Norb, Norb, Norb) + RDM electron exchange symmetry only: shape(two_RDM) = (Ntrn, Ntrn, (Norb**2 * (Norb**2 +1)/2) + RDM electron exchange symmetry + data symmetry; shape(two_RDM) = (Ntrn * (Ntrn + 1)/2, (Norb**2 * (Norb**2 +1)/2)) + S : numpy.ndarray + The overlap matrix. + nroots (optional): int + Number of states in the solver. + hermitian (bool, optional): + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple + A tuple containing the total potential energies and its gradients. + """ + # Construct h1 and h2 + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + h1 = np.linalg.multi_dot((ao_mo_trafo.T, scf.hf.get_hcore(mol), ao_mo_trafo)) + h2 = ao2mo.restore(1, ao2mo.kernel(mol, ao_mo_trafo), mol.nao) + + en, vec = approximate_multistate(h1, h2, one_RDM, two_RDM, S, nroots=nroots, hermitian=hermitian) + + # Get the gradient of one and two-electron integrals before contracting onto + # rdms of different states + h1_jac, h2_jac = get_one_and_two_el_grad(mol,ao_mo_trafo=ao_mo_trafo) + + grad_elec_all = [] + one_rdm_predicted_all = [] + two_rdm_predicted_all = [] + for i_state in range(nroots): + vec_i = vec[i_state,:] + + #one_rdm_predicted = lib.einsum("i,ijkl,j->kl", vec_i, one_RDM, vec_i, optimize="optimal") + #two_rdm_predicted = lib.einsum( + # "i,ijklmn,j->klmn", vec_i, two_RDM, vec_i, optimize="optimal" + #) + + one_rdm_predicted = np.tensordot(np.outer(vec_i, vec_i), one_RDM, axes=2) + if len(two_RDM.shape) == 2 or len(two_RDM.shape) == 5: + # symmetry in data points + eigenvec_mat = 2 * np.outer(vec_i, vec_i) + + np.fill_diagonal(eigenvec_mat, 0.5 * np.diag(eigenvec_mat)) + + two_rdm_predicted = np.tensordot( + eigenvec_mat[np.tril_indices(len(vec_i))], two_RDM, axes=1 + ) + + else: + two_rdm_predicted = np.tensordot(np.outer(vec_i, vec_i), two_RDM, axes=2) + + if len(two_rdm_predicted.shape) != 4: + two_rdm_predicted = restore_electron_exchange_symmetry( + two_rdm_predicted, mol.nao + ) + + + grad_elec = get_grad_elec_from_gradH( + one_rdm_predicted, two_rdm_predicted, h1_jac, h2_jac + ) + + one_rdm_predicted_all.append(one_rdm_predicted) + two_rdm_predicted_all.append(two_rdm_predicted) + grad_elec_all.append(grad_elec) + + grad_elec_all = np.array(grad_elec_all) + + if return_density_matrices: + return ( + en.real + mol.energy_nuc(), + grad_elec + grad.RHF(scf.RHF(mol)).grad_nuc(), + one_rdm_predicted_all, + two_rdm_predicted_all, + ) + + else: + return ( + en.real + mol.energy_nuc(), + grad_elec + grad.RHF(scf.RHF(mol)).grad_nuc(), + ) + +@timeit +def get_multistate_energy_with_grad_and_NAC(mol, one_RDM, two_RDM, S, nroots=1, + savemem=True, hermitian=True): + """ + Calculates the potential energiesm its gradient w.r.t. nuclear positions of a + molecule and nonadiabatic couplings from eigenvector continuation for both + ground and excited states. + + Args: + mol : pyscf.gto.Mole + The molecule object. + one_RDM : numpy.ndarray + The one-electron t-RDM. + two_RDM : numpy.ndarray + The two-electron t-RDM. + S : numpy.ndarray + The overlap matrix. + nroots (optional): int + Number of states in the solver. + hermitian (optional): bool + Whether problem is solved with eigh or with eig. Defaults to True. + + Returns: + tuple (vec, en, grad_all, nac_all, nac_all_hfonly) + A tuple containing the continuation eigenvector, total potential energies, its gradients and NACs: + + vec: ndarray(ntrain,) + Coefficients of the linear expansion + + en: ndarray(nroot,) + Total potential energies for both ground and excited states + + grad_all: list of ndarray(nat,3) + Gradients of multistate energies + + nac_all: dictionary of np.darray(nat,3) + Nonadiabatic coupling vectors between all states, + e.g. nac_all['02'] is NAC along ground state and 2nd excited state + + nac_all_hfonly: dictionary of ndarray(nat,3) + Hellman-Feynmann contribution to NACs + """ + + # Construct h1 and h2 + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + h1 = np.linalg.multi_dot((ao_mo_trafo.T, scf.hf.get_hcore(mol), ao_mo_trafo)) + h2 = ao2mo.restore(1, ao2mo.kernel(mol, ao_mo_trafo), mol.nao) + + # Diagonalization of the subspace Hamiltonian for the continuation of + # energies and eigenstates + en, vec = approximate_multistate(h1, h2, one_RDM, two_RDM, S, nroots=nroots, hermitian=hermitian) + fix_gauge(vec) + + if not savemem: + # Get the gradient of one and two-electron integrals before contracting onto + # rdms and trmds of different states + h1_jac, h2_jac = get_one_and_two_el_grad(mol,ao_mo_trafo=ao_mo_trafo) + + # Get the orbital derivative coupling for NACs + orb_deriv = get_orbital_derivative_coupling(mol,ao_mo_trafo=ao_mo_trafo) + + # Nuclear part of the gradient + grad_nuc = grad.RHF(scf.RHF(mol)).grad_nuc() + + grad_elec_all = [] + nac_all = {} + nac_all_hfonly = {} + # Iterate over pairs of eigenstates + for i_state in range(nroots): + vec_i = vec[i_state,:] + + for j_state in range(nroots): + vec_j = vec[j_state,:] + + # Contracting to subspace eigenstate in hand + #one_rdm_predicted = lib.einsum("i,ijkl,j->kl", vec_i, one_RDM, vec_j, optimize="optimal") + #two_rdm_predicted = lib.einsum( + # "i,ijklmn,j->klmn", vec_i, two_RDM, vec_j, optimize="optimal" + #) + one_rdm_predicted = np.tensordot(np.outer(vec_i, vec_j), one_RDM, axes=2) + if len(two_RDM.shape) == 2 or len(two_RDM.shape) == 5: + # symmetry in data points + eigenvec_mat = 2 * np.outer(vec_i, vec_j) + + np.fill_diagonal(eigenvec_mat, 0.5 * np.diag(eigenvec_mat)) + + two_rdm_predicted = np.tensordot( + eigenvec_mat[np.tril_indices(len(vec_i))], two_RDM, axes=1 + ) + + else: + two_rdm_predicted = np.tensordot(np.outer(vec_i, vec_j), two_RDM, axes=2) + + if len(two_rdm_predicted.shape) != 4: + two_rdm_predicted = restore_electron_exchange_symmetry( + two_rdm_predicted, mol.nao + ) + + # d\dR of subspace Hamiltonian + if savemem: + grad_elec = get_grad_elec_OAO( + mol, one_rdm_predicted, two_rdm_predicted + ) + else: + grad_elec = get_grad_elec_from_gradH( + one_rdm_predicted, two_rdm_predicted, h1_jac, h2_jac + ) + + + # Energy gradients + if i_state == j_state: + grad_elec_all.append(grad_elec) + + # Nonadiabatic couplings + else: + # Hellman-Feynman contribution to NAC + nac_hf = grad_elec/(en[j_state]-en[i_state]) + + #if i_state == 0 and j_state == 1: + # print(i_state, j_state, grad_elec) + + # Orbital contribution to NAC + nac_orb = lib.einsum("ij,ijkl->kl",one_rdm_predicted, orb_deriv, optimize="optimal") + + # Total NAC + nac_ij = nac_hf + nac_orb + + # Save to dictionaries + nac_all[str(i_state)+str(j_state)] = nac_ij + nac_all_hfonly[str(i_state)+str(j_state)] = nac_hf + + # Add the nuclear contribution to gradient + grad_all = np.array(grad_elec_all) + grad_nuc + + return ( + vec, + en.real + mol.energy_nuc(), + grad_all, + nac_all, + nac_all_hfonly + ) + + +############################################################################## +def two_el_grad_lowrank(mol, lowrank_vecs, ED_builds, SVD_builds, vec_i, vec_j, + ao_mo_trafo=None, ao_mo_trafo_grad=None): + """ + Computing the gradient of the electronic Hamiltonian wrt atomic coordinates + using the low-rank representation of the 2-tRDM + + ### OUTDATED; use state_resolved_version, or adapt it here + + """ + # AO to SAO basis transformation + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + if ao_mo_trafo_grad is None: + ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) + + # Preliminaries + ntrain = vec_i.shape[0] + norb = ao_mo_trafo.shape[-1] + + # Unpack the preliminaries + if lowrank_vecs['has_ed']: + lr_vecs, lr_vecs_ao, vhf, vhf_grad, vhf_grad_t = ED_builds + + if lowrank_vecs['has_svd']: + svd_lvecs, svd_rvecs, svd_lvecs_ao, svd_rvecs_ao, \ + vj_left, vj_right, \ + vj_l_grad, vj_r_grad = SVD_builds + + #print('Error in two_el_grad_lowrank: SVD not implemented yet') + #sys.exit() + + + # Basis functions indices for each atom + atm_slices = tuple( + [ + (mol.aoslice_by_atom()[i][2], mol.aoslice_by_atom()[i][3]) + for i in range(mol.natm) + ] + ) + + # Pulay terms + """ + pulay_term = 4*lib.einsum('xj,ABawx,ABaij,ABa,A,B->wi', ao_mo_trafo, + vhf, lr_vecs, lowrank_vecs['vals'][:,:,:vhf.shape[2]], + vec_i, vec_j, optimize='optimal') + + grad_i = lib.einsum('wimn,wi->mn',ao_mo_trafo_grad,pulay_term,optimize='optimal') + """ + pulay_term = np.zeros([ntrain,ntrain,norb,norb]) + if lowrank_vecs['has_ed']: + pulay_term += 2*lib.einsum('xj,ABawx,ABaij,ABa->ABwi', ao_mo_trafo, + vhf, lr_vecs, lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + pulay_term += 2*lib.einsum('wi,ABawx,ABaij,ABa->ABxj', ao_mo_trafo, + vhf, lr_vecs, lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + if lowrank_vecs['has_svd']: + + pulay_term += lib.einsum('xj,ABawx,ABaij,ABa->ABwi', ao_mo_trafo, + vj_left, svd_rvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + pulay_term += lib.einsum('wi,ABawx,ABaij,ABa->ABxj', ao_mo_trafo, + vj_left, svd_rvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + pulay_term += lib.einsum('xj,ABawx,ABaij,ABa->ABwi', ao_mo_trafo, + vj_right, svd_lvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + pulay_term += lib.einsum('wi,ABawx,ABaij,ABa->ABxj', ao_mo_trafo, + vj_right, svd_lvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + if lowrank_vecs['hermitian']: + # Set the upper triangle + pulay_term[np.triu_indices(ntrain)] = pulay_term.transpose(1,0,2,3)[np.triu_indices(ntrain)].conj() + + grad_i = lib.einsum('wimn,ABwi,A,B->mn',ao_mo_trafo_grad,pulay_term,vec_i, vec_j, optimize='optimal') + + #grad_i = np.zeros_like(grad_i) # Setting the previous part to zero for testing + + # JK Grad terms + grad_h2 = np.zeros([ntrain, ntrain, 3, norb]) + if lowrank_vecs['has_ed']: + + grad_h2 += 2*lib.einsum('ABanij,ABaij,ABa->ABni', + vhf_grad_t, lr_vecs_ao,lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + grad_h2 += 2*lib.einsum('ABanij,ABaji,ABa->ABni', + vhf_grad, lr_vecs_ao,lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + if lowrank_vecs['has_svd']: + + #print(vj_l_grad.shape,svd_rvecs_ao.shape ) + grad_h2 += lib.einsum('ABanij,ABaij,ABa->ABni', + vj_l_grad, svd_rvecs_ao+svd_rvecs_ao.transpose(0,1,2,4,3),lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + grad_h2 += lib.einsum('ABanij,ABaij,ABa->ABni', + vj_r_grad, svd_lvecs_ao+svd_lvecs_ao.transpose(0,1,2,4,3),lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + if lowrank_vecs['hermitian']: + # Set the upper triangle + grad_h2[np.triu_indices(ntrain)] = grad_h2.transpose(1,0,2,3)[np.triu_indices(ntrain)].conj() + + grad_el_traced = lib.einsum('ABni,A,B->ni',grad_h2, vec_i, vec_j, optimize='optimal') + + # Sum contributions from each orbital on atom site, i + for i, slice in enumerate(atm_slices): + grad_i[i,:] += grad_el_traced[:,slice[0] : slice[1]].sum(axis=1) + + return grad_i + +@timeit +def state_resolved_two_el_grad_lowrank(mol, lowrank_vecs, ED_builds, SVD_builds, + diag_builds=None, diagonals=None, + ao_mo_trafo=None, ao_mo_trafo_grad=None, + ints_SAO=None, + df_response=False): + """ + Computing the gradient of the electronic Hamiltonian wrt atomic coordinates + using the low-rank representation of the 2-tRDM + + Note: This function does the same thing as two_el_grad_lowrank() but returns + the 2-el gradient with the training state indices such that this can only be + computed once, and reused for the force and NAC calculation of each electronic + excited state. + + """ + # AO to SAO basis transformation + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + if ao_mo_trafo_grad is None: + ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) + + # Preliminaries + # Get ntrain from pairloc (max index + 1) - might just make ntrain an input arg to avoid this + if lowrank_vecs['has_ed']: + ntrain_ed = max([i for i, j in lowrank_vecs['pairloc'].keys()]) + 1 + else: + ntrain_ed = 0 + if lowrank_vecs['has_svd']: + ntrain_svd = max([i for i, j in lowrank_vecs['pairloc_svd'].keys()]) + 1 + else: + ntrain_svd = 0 + ntrain = max(ntrain_ed, ntrain_svd) + + norb = ao_mo_trafo.shape[-1] + + # Unpack the preliminaries + if lowrank_vecs['has_ed']: + lr_vecs, lr_vecs_ao, vhf, vhf_grad, vhf_grad_t = ED_builds + + if lowrank_vecs['has_svd']: + svd_lvecs, svd_rvecs, svd_lvecs_ao, svd_rvecs_ao, \ + vj_left, vj_right, \ + vj_l_grad, vj_r_grad = SVD_builds + + # Unpack diagonal builds + use_diag, Jdiag_only = False, True + if diag_builds is not None: + use_diag = True + (vj, vj_grad, vk, vk_t, vk_grad) = diag_builds + if vk is not None: + Jdiag_only = False + + # If SAO implementation for diagonal corrections + if ints_SAO is not None: + use_diag = True + sao_diag = True + Lpq, Lpq_sao, Lpq_grad_sao = ints_SAO + + if len(Lpq_grad_sao.shape) == 5: + Jdiag_only = False + Lpq_grad_sao_diag = np.einsum('nPiaa->nPia',Lpq_grad_sao) + + elif len(Lpq_grad_sao.shape) == 4: + Jdiag_only = True + Lpq_grad_sao_diag = Lpq_grad_sao + + else: + print('Error in state_resolved_two_el_grad_lowrank: Lpq_grad_sao has incorrect number of indices.') + sys.exit() + else: + sao_diag = False + + # Basis functions indices for each atom + atm_slices = tuple( + [ + (mol.aoslice_by_atom()[i][2], mol.aoslice_by_atom()[i][3]) + for i in range(mol.natm) + ] + ) + + # Pulay terms + pulay_term = np.zeros([ntrain,ntrain,norb,norb]) + if lowrank_vecs['has_ed']: + pulay_term += 2*lib.einsum('xj,ABawx,ABaij,ABa->ABwi', ao_mo_trafo, + vhf, lr_vecs, lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + pulay_term += 2*lib.einsum('wi,ABawx,ABaij,ABa->ABxj', ao_mo_trafo, + vhf, lr_vecs, lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + if lowrank_vecs['has_svd']: + + pulay_term += lib.einsum('xj,ABawx,ABaij,ABa->ABwi', ao_mo_trafo, + vj_left, svd_rvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + pulay_term += lib.einsum('wi,ABawx,ABaij,ABa->ABxj', ao_mo_trafo, + vj_left, svd_rvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + pulay_term += lib.einsum('xj,ABawx,ABaij,ABa->ABwi', ao_mo_trafo, + vj_right, svd_lvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + pulay_term += lib.einsum('wi,ABawx,ABaij,ABa->ABxj', ao_mo_trafo, + vj_right, svd_lvecs, lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + if use_diag: + + if not sao_diag: + pulay_term += 2*lib.einsum('xi,ABiwx->ABwi', + ao_mo_trafo, vj + vj.transpose(0,1,2,4,3), + optimize='optimal') + + # Not needed! Same as above, so just times the above by 2 + #pulay_term += lib.einsum('zj,ABjyz->AByj', + # ao_mo_trafo, vj + vj.transpose(0,1,2,4,3), + # optimize='optimal') + if not Jdiag_only: + + pulay_term += 2*lib.einsum('xi,ABiwx->ABwi', + ao_mo_trafo, vk + vk_t, + optimize='optimal') + + else: + diagJ_unpack = unstack_tril(diagonals[0],hermitian=False) + + pulay_term += 2*lib.einsum('xi,ABij,Pwx,Pjj->ABwi', + ao_mo_trafo, + diagJ_unpack + diagJ_unpack.transpose(0,1,3,2), + Lpq, + Lpq_sao, + optimize='optimal') + + if not Jdiag_only: + # TODO: diag_K contributions + diagK_unpack = unstack_tril(diagonals[1],hermitian=False) + + pulay_term += 2*lib.einsum('xj,ABij,Pwx,Pij->ABwi', + ao_mo_trafo, + diagK_unpack + diagK_unpack.transpose(0,1,3,2), + Lpq, + Lpq_sao, + optimize='optimal') + + if lowrank_vecs['hermitian']: + # Set the upper triangle + pulay_term[np.triu_indices(ntrain)] = pulay_term.transpose(1,0,2,3)[np.triu_indices(ntrain)].conj() + + intermediate_pulay = lib.einsum('wimn,ABwi->ABmn',ao_mo_trafo_grad,pulay_term, optimize='optimal') + + # JK Grad terms + grad_h2 = np.zeros([ntrain, ntrain, 3, norb]) + + if lowrank_vecs['has_ed']: + + grad_h2 += 2*lib.einsum('ABanij,ABaij,ABa->ABni', + vhf_grad_t, lr_vecs_ao,lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + grad_h2 += 2*lib.einsum('ABanij,ABaji,ABa->ABni', + vhf_grad, lr_vecs_ao,lowrank_vecs['vals'][:,:,:vhf.shape[2]], + optimize='optimal') + + if lowrank_vecs['has_svd']: + + #print(vj_l_grad.shape,svd_rvecs_ao.shape ) + grad_h2 += lib.einsum('ABanij,ABaij,ABa->ABni', + vj_l_grad, svd_rvecs_ao+svd_rvecs_ao.transpose(0,1,2,4,3),lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + grad_h2 += lib.einsum('ABanij,ABaij,ABa->ABni', + vj_r_grad, svd_lvecs_ao+svd_lvecs_ao.transpose(0,1,2,4,3),lowrank_vecs['vals'][:,:,:vj_left.shape[2]], + optimize='optimal') + + if use_diag: + + ## AO Builds + if not sao_diag: + grad_h2 += 4*lib.einsum('ABinyz,yi,zi->ABny', + vj_grad, ao_mo_trafo, ao_mo_trafo, + optimize='optimal') + + # Same result as above + #grad_h2 += 2*lib.einsum('ABinyz,yi,zi->ABny', + # vj_grad_t, ao_mo_trafo, ao_mo_trafo, + # optimize='optimal') + + if not Jdiag_only: + + grad_h2 += 2*lib.einsum('ABinwy,wi,yi->ABnw', + vk_grad, ao_mo_trafo, ao_mo_trafo, + optimize='optimal') + + # Included this term in the K build + #grad_h2 += 2*lib.einsum('ABinwy,wi,yi->ABnw', + # vk_grad_t, ao_mo_trafo, ao_mo_trafo, + # optimize='optimal') + + + else: + # df_grad_4c_ints = np.einsum('xijP,Pkl->xijkl', deriv_cderi, cd_array) + diagJ_unpack = unstack_tril(diagonals[0],hermitian=False) + + # NOTE: Unlike AO, the minus sign is not included in the integrals + # so we need to subtract these contributions in SAO + grad_h2 -= 2 * lib.einsum('ABab,nPia,Pbb->ABni',diagJ_unpack + diagJ_unpack.transpose(0,1,3,2), Lpq_grad_sao_diag, Lpq_sao) + + if not Jdiag_only: + + diagK_unpack = unstack_tril(diagonals[1],hermitian=False) + + grad_h2 -= 2 * lib.einsum('ABab,nPiab,Pab->ABni',diagK_unpack + diagK_unpack.transpose(0,1,3,2), Lpq_grad_sao, Lpq_sao) + + if lowrank_vecs['hermitian']: + # Set the upper triangle + grad_h2[np.triu_indices(ntrain)] = grad_h2.transpose(1,0,2,3)[np.triu_indices(ntrain)].conj() + + # Sum contributions from each orbital on atom site, i + intermediate_h2 = np.zeros([ntrain, ntrain, mol.natm, 3]) + for i, slice in enumerate(atm_slices): + intermediate_h2[:,:,i,:] += grad_h2[:,:,:,slice[0] : slice[1]].sum(axis=3) + + # Response of the auxillary basis (Currently NOT WORKING) + if df_response: + intermediate_df = np.zeros([ntrain, ntrain, mol.natm, 3]) + if lowrank_vecs['has_ed']: + #vhf_aux = lib.einsum('aamn,a->mn',vj_list.aux - vk_list.aux*.5,lr_vals[ii],optimize='optimal') + + intermediate_df += vhf_grad_t.aux + intermediate_df += vhf_grad.aux + + #intermediate_df += lib.einsum('ABABamn,ABa->ABmn', + # vhf_grad.aux,lowrank_vecs['vals'][:,:,:vhf.shape[2]], + # optimize='optimal') + + if lowrank_vecs['has_svd']: + + intermediate_df += vj_l_grad.aux + intermediate_df += vj_r_grad.aux + + return intermediate_h2 + intermediate_pulay + intermediate_df + else: + return intermediate_h2 + intermediate_pulay + +@timeit +def get_lowrank_en_with_grad_and_NAC(mol, one_RDM, S, lowrank_vecs, + diagonals=None, Jdiag_only=True, sao_diag=False, + nroots=1, + density_fit=True, df_basis=None, + ao_mo_trafo=None, ao_mo_trafo_grad=None, + df_response=False, + hermitian=True, + lindep=1e-6): + """ + Construct subspace Hamiltonian using the low-rank decomposition of + 2-transition-cumulant + + Input: + mol (Mole object): pySCF mole object at the test geometry + + """ + ### Preliminaries + ntrain = S.shape[0] + + norb = one_RDM.shape[-1] + norb_sq = norb * norb + + with log_time('Grad preliminaries'): + # Initiate the mean field object to use DF integrals (no need to use kernel) + #mol.symmetry = False + mf = scf.RHF(mol).density_fit(auxbasis=df_basis) + + # AO to SAO basis transformation + ovlp_ao = mol.intor("int1e_ovlp") + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(ovlp_ao) + + if ao_mo_trafo_grad is None: + ao_mo_trafo_grad = get_derivative_ao_mo_trafo(mol) + + # 1-electron integrals with DF + h1_ao = mf.get_hcore() + h1e_sao = lib.einsum('ai,ab,bj->ij', ao_mo_trafo, h1_ao, ao_mo_trafo) + + if df_response: + print('Error in get_lowrank_en_with_grad_and_NAC(): \n \ + Response of the auxillary basis is not fully implemented. Set df_response=False.') + sys.exit() + + ###################################################### + # Get preliminaries + ED_builds, SVD_builds, diag_builds = get_jk_builds( + mol, lowrank_vecs, + diagonals=diagonals, + Jdiag_only=Jdiag_only, + sao_diag=sao_diag, + ao_mo_trafo=ao_mo_trafo, + density_fit=density_fit, + df_basis=df_basis, + df_response=df_response + ) + + if lowrank_vecs['has_ed']: + lr_vecs, lr_vecs_ao, vhf, vhf_grad, vhf_grad_t = ED_builds + + if lowrank_vecs['has_svd']: + svd_lvecs, svd_rvecs, svd_lvecs_ao, svd_rvecs_ao, \ + vj_left, vj_right, \ + vj_l_grad, vj_r_grad = SVD_builds + + use_diag = False + if diag_builds is not None: + use_diag = True + (vj, vj_grad, vk, vk_t, vk_grad) = diag_builds + + # if sao_diag + elif diagonals is not None: + use_diag = True + + # SAO grads for gradients + if sao_diag: + + with log_time('SAO transf. integrals'): + + # Get the integrals in DF + Lpq, Lpq_grad = get_df_integrals(mol, auxbasis=df_basis, grad=True) + norb = Lpq.shape[-1] + + # Transform to SAO basis + Lpq = lib.pack_tril(Lpq) + Lpq_sao = ao2mo._ao2mo.nr_e2(Lpq, ao_mo_trafo, + (0, ao_mo_trafo.shape[1], 0, ao_mo_trafo.shape[1]),aosym="s2",mosym="s2") + Lpq_sao = lib.unpack_tril(Lpq_sao) + + # Transform the derivatives to SAO basis + """ + Lpq_grad = np.einsum('xijP->xPij',Lpq_grad) + orig_shap = Lpq_grad.shape + #Lpq_grad = Lpq_grad.reshape([-1, norb, norb]) + Lpq_grad = Lpq_grad.reshape([orig_shap[0]*orig_shap[1], norb, norb]) + + Lpq_grad_sao = ao2mo._ao2mo.nr_e2(Lpq_grad, ao_mo_trafo, + (0, norb, 0, norb), aosym='s1', mosym='s1') + Lpq_grad_sao = Lpq_grad_sao.reshape(orig_shap) + """ + # Explicit transformation (for testing, comment out later) + #Lpq_grad_sao = np.einsum('xabP,ia,jb->xPij',Lpq_grad,ao_mo_trafo,ao_mo_trafo) + + if Jdiag_only: + Lpq_grad_sao = np.einsum('ia,ja,nijP->nPia', ao_mo_trafo, ao_mo_trafo, Lpq_grad, optimize='optimal') + else: + Lpq_grad_sao = np.einsum('ia,jb,nijP->nPiab', ao_mo_trafo, ao_mo_trafo, Lpq_grad, optimize='optimal') + + ints_sao = (lib.unpack_tril(Lpq), Lpq_sao, Lpq_grad_sao) + else: + ints_sao = None + + ###################################################### + ######### CONSTRUCT SUBSPACE HAMILTONIAN AND GRADS + ###################################################### + with log_time('Low-rank Hamiltonian'): + ### 1-body contributions + subspace_h = lib.einsum('...kl,kl->...', one_RDM, h1e_sao) + + # Contruction for subspace Hamiltonian + if lowrank_vecs['has_ed']: + subspace_h += 0.5*lib.einsum('xyaij,xyaij,xya->xy', vhf, lr_vecs_ao, lowrank_vecs['vals'][:,:,:vhf.shape[2]],optimize='optimal') + + if lowrank_vecs['has_svd']: + subspace_h += 0.5*lib.einsum('xyaij,xyaij,xya->xy', vj_right, svd_lvecs_ao, lowrank_vecs['vals'][:,:,:vj_right.shape[2]],optimize='optimal') + + if use_diag: + + if not sao_diag: + subspace_h += 0.5 * np.einsum('yj,zj,XYjyz->XY', ao_mo_trafo, ao_mo_trafo, vj) + if not Jdiag_only: + subspace_h += 0.5 * np.einsum('xj,zj,XYjxz->XY', ao_mo_trafo, ao_mo_trafo, vk) + + else: + diagJ_unpack = unstack_tril(diagonals[0],hermitian=False) + subspace_h += 0.5 * np.einsum('XYij,Pii,Pjj->XY',diagJ_unpack, Lpq_sao, Lpq_sao) + if not Jdiag_only: + diagK_unpack = unstack_tril(diagonals[1],hermitian=False) + subspace_h += 0.5 * np.einsum('XYij,Pij,Pij->XY',diagK_unpack, Lpq_sao, Lpq_sao) + + if lowrank_vecs['hermitian']: + # Set the upper triangle + subspace_h[np.triu_indices(ntrain)] = subspace_h.T[np.triu_indices(ntrain)].conj() + + # Diagonalize + en, vec = solve_subspace(subspace_h, S, hermitian=hermitian, nroots=nroots, lindep=lindep) + fix_gauge(vec) + + ###################################################### + ######### GET GRADIENTS + ###################################################### + #with log_time('Grad one-body'): + # Get the orbital derivative coupling for NACs + orb_deriv = get_orbital_derivative_coupling(mol, + ao_mo_trafo=ao_mo_trafo, + ao_mo_trafo_grad=ao_mo_trafo_grad, + ovlp=ovlp_ao) + + # 1-el grad + h1_jac = get_one_el_grad( + mol, + h1_ao=h1_ao, + ao_mo_trafo=ao_mo_trafo, + ao_mo_trafo_grad=ao_mo_trafo_grad + ) + + # Nuclear part of the gradient + grad_nuc = grad.RHF(scf.RHF(mol)).grad_nuc() + + # 2-el grad + grad_elec_h2_state = state_resolved_two_el_grad_lowrank( + mol, lowrank_vecs, + ED_builds, SVD_builds, diag_builds, + diagonals=diagonals, + ao_mo_trafo=ao_mo_trafo, + ao_mo_trafo_grad=ao_mo_trafo_grad, + #sao_diag=sao_diag, + ints_SAO=ints_sao, + df_response=df_response + ) + + grad_elec_all = [] + nac_all = {} + nac_all_hfonly = {} + # Iterate over pairs of eigenstates + for i_state in range(nroots): + vec_i = vec[i_state,:] + + for j_state in range(nroots): + vec_j = vec[j_state,:] + + # Contracting to subspace eigenstate in hand + #one_rdm_predicted = lib.einsum("i,ijkl,j->kl", vec_i, one_RDM, vec_j, optimize="optimal") + #two_rdm_predicted = lib.einsum( + # "i,ijklmn,j->klmn", vec_i, two_RDM, vec_j, optimize="optimal" + #) + one_rdm_predicted = np.tensordot(np.outer(vec_i, vec_j), one_RDM, axes=2) + + # 1-electron contributions to gradient + grad_elec_h1 = lib.einsum("ij,ijkl->kl", one_rdm_predicted, h1_jac, optimize="optimal") + + #grad_elec_h2 = two_el_grad_lowrank(mol, lowrank_vecs, ED_builds, SVD_builds, vec_i, vec_j, + # ao_mo_trafo=ao_mo_trafo, ao_mo_trafo_grad=ao_mo_trafo_grad) + grad_elec_h2 = lib.einsum('ABmn,A,B->mn',grad_elec_h2_state,vec_i, vec_j,optimize="optimal") + + grad_elec = grad_elec_h1 + 0.5*grad_elec_h2 + #grad_elec = 0.5*grad_elec_h2 # For testing 2-el part only + + # Energy gradients + if i_state == j_state: + grad_elec_all.append(grad_elec) + + # Nonadiabatic couplings + else: + # Hellman-Feynman contribution to NAC + nac_hf = grad_elec/(en[j_state]-en[i_state]) + + #if i_state == 0 and j_state == 1: + # print(i_state, j_state, grad_elec) + + # Orbital contribution to NAC + nac_orb = lib.einsum("ij,ijkl->kl",one_rdm_predicted, orb_deriv, optimize="optimal") + + # Total NAC + nac_ij = nac_hf + nac_orb + + # Save to dictionaries + nac_all[str(i_state)+str(j_state)] = nac_ij + nac_all_hfonly[str(i_state)+str(j_state)] = nac_hf + + # Add the nuclear contribution to gradient + grad_all = np.array(grad_elec_all) + grad_nuc + + return ( + vec, + en.real + mol.energy_nuc(), + grad_all, + nac_all, + nac_all_hfonly + ) + +if __name__ == '__main__': + + # Some initial checks for the code + from pyscf import gto, fci + + from pyscf.fci.addons import fix_spin_ + from evcont.FCI_EVCont import FCI_EVCont_obj + + from pyscf.mcscf import CASCI + + from time import time + + CASE = 'NAC' + #CASE = 'Exc-Grad' + #CASE = 'Grad' + + nstate = 2 #1st excited state + nroots_evcont = 3 + cibasis = 'OAO' + + natom = 8 + + test_range = np.linspace(0.8, 3.0,20) + + def get_mol(positions): + mol = gto.Mole() + + mol.build( + atom=[("H", pos) for pos in positions], + basis="sto-6g", + #basis="6-31g", + #basis='ccpvdz', + symmetry=False, + unit="Bohr", + verbose=0 + ) + + return mol + + + # training geometries + equilibrium_dist = 1.78596 + + equilibrium_pos = np.array([(x * equilibrium_dist, 0.0, 0.0) for x in range(10)]) + + training_stretches = np.array([0.0, 0.5, -0.5, 1.0, -1.0]) + + trainig_dists = equilibrium_dist + training_stretches + + #trainig_dists = [1.0, 1.8, 2.6] + + continuation_object = FCI_EVCont_obj(nroots=nroots_evcont, + cibasis=cibasis) + + + # Generate training data + prepare training models + for i, dist in enumerate(trainig_dists): + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + mol = get_mol(positions) + continuation_object.append_to_rdms(mol) + + print('Finished training') + + if CASE == 'Grad': + + st = time() + # Test the ground and excite state forces and gradients at training points + for i, dist in enumerate(trainig_dists): + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + mol = get_mol(positions) + + # Predictions from mutistate continuation + en_continuation_ms, grad_continuation_ms = get_energy_with_grad( + mol, + continuation_object.one_rdm, + continuation_object.two_rdm, + continuation_object.overlap + ) + + # Predictions from Hartree-Fock + hf_energy, hf_grad = mol.RHF().nuc_grad_method().as_scanner()(mol) + + # Fci reference values + #en_exact, grad_exact = CASCI(mol.RHF(), natom, natom).nuc_grad_method().as_scanner()(mol) + + # Fci excited state reference values + mc = CASCI(mol.RHF(), mol.nao, mol.nelectron) + #mc = CASCI(mol.RHF(), 10,6) Li2 + mc.fcisolver = fci.direct_spin0.FCI() + + ci_scan_0 = mc.nuc_grad_method().as_scanner() + + en_exact, grad_exact = ci_scan_0(mol) + + # Checks + print(i) + + assert np.allclose(en_exact,en_continuation_ms) + assert np.allclose(grad_exact,grad_continuation_ms) + + print(f'Time taken: {time()-st:.1f} sec') + + + if CASE == 'Exc-Grad': + + st = time() + # Test the ground and excite state forces and gradients at training points + for i, dist in enumerate(trainig_dists): + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + mol = get_mol(positions) + + # Predictions from mutistate continuation + _, en_continuation_ms, grad_continuation_ms, nac_continuation, _ = get_multistate_energy_with_grad_and_NAC( + mol, + continuation_object.one_rdm, + continuation_object.two_rdm, + continuation_object.overlap, + nroots=nstate+1 + ) + + # Predictions from Hartree-Fock + hf_energy, hf_grad = mol.RHF().nuc_grad_method().as_scanner()(mol) + + # Fci reference values + #en_exact, grad_exact = CASCI(mol.RHF(), natom, natom).nuc_grad_method().as_scanner()(mol) + + # Fci excited state reference values + mc = CASCI(mol.RHF(), mol.nao, mol.nelectron) + #mc = CASCI(mol.RHF(), 10,6) Li2 + mc.fcisolver = fci.direct_spin0.FCI() + #mc.fcisolver = fci.direct_spin1.FCI() + #fix_spin_(mc.fcisolver,shift=0.9,ss=0) + mc.fcisolver.nroots = nstate+3 + #mc.fcisolver.nroots = 6 + #mc.fcisolver.conv_tol = 1.e-14 + #mc.fcisolver.max_space=30 + #mc.fcisolver.max_cycle=250 + + ci_scan_exc = mc.nuc_grad_method().as_scanner(state=nstate) + ci_scan_0 = mc.nuc_grad_method().as_scanner(state=0) + + en_exc_exact, grad_exc_exact = ci_scan_exc(mol) + en_exact, grad_exact = ci_scan_0(mol) + + # Get the reference numerical FCI NACs + #nac_all = nac_fd_FCI(mc,nroots=nstate+1) + #print(nac_all) + + # Checks + print(i) + + assert np.allclose(en_exact,en_continuation_ms[0]) + assert np.allclose(grad_exact,grad_continuation_ms[0]) + + assert np.allclose(en_exc_exact,en_continuation_ms[nstate]) + assert np.allclose(grad_exc_exact,grad_continuation_ms[nstate],atol=1e-5) + + print(f'Time taken: {time()-st:.1f} sec') + + # Test NACs + if CASE == 'NAC': + + # TODO: Write new tests using FCI implementation + pass diff --git a/evcont/converge_dmrg.py b/evcont/converge_dmrg.py index 757f123..1ed1bb3 100644 --- a/evcont/converge_dmrg.py +++ b/evcont/converge_dmrg.py @@ -17,6 +17,7 @@ def converge_dmrg( noises=np.append(np.logspace(-2, -7, num=4), 0), tolerance=1.0e-4, restart_tag=None, + nroots=1, mem=5, reorder=None, ): @@ -66,9 +67,9 @@ def converge_dmrg( "nodex/{}-mps_info.bin".format(restart_tag) ): # Load a previous DMRG calculation if restart tag is provided - ket = mps_solver.load_mps(restart_tag) + ket = mps_solver.load_mps(restart_tag, nroots=nroots) else: - ket = mps_solver.get_random_mps(tag, bond_dim=bond_dim_schedule[0], nroots=1) + ket = mps_solver.get_random_mps(tag, bond_dim=bond_dim_schedule[0], nroots=nroots) final_energies = [] @@ -88,7 +89,7 @@ def converge_dmrg( tol=tolerance, ) bnd_dms, dws, ens = mps_solver.get_dmrg_results() - final_energies.append(ens[-1][0]) + final_energies.append(ens[-1][:]) if rank == 0: print(bnd_dms, final_energies, dws) with open("DMRG_result_{}.txt".format(tag), "a") as fl: @@ -101,7 +102,7 @@ def converge_dmrg( "{} {} {} {}\n".format(bnd_dms[j], ens[j][0], dws[j], noise) ) if len(final_energies) > 1: - if abs(final_energies[-1] - final_energies[-2]) < tolerance: + if np.max(np.abs(final_energies[-1] - final_energies[-2])) < tolerance: break return ket, final_energies[-1] diff --git a/evcont/electron_integral_utils.py b/evcont/electron_integral_utils.py index 3b58a7d..05b5c9f 100644 --- a/evcont/electron_integral_utils.py +++ b/evcont/electron_integral_utils.py @@ -1,6 +1,6 @@ import numpy as np -from pyscf import scf, lo, ao2mo +from pyscf import scf, lo, ao2mo, df def get_loewdin_trafo(overlap_mat): @@ -125,7 +125,7 @@ def get_integrals(mol, basis): Parameters: mol (pyscf.gto.Mole): The molecule object. - basis (numpy.ndarray): The basis set. + basis (numpy.ndarray): The basis set (AO->MO transformation coefficients). Returns: h1 (numpy.ndarray): The one-electron integrals. @@ -136,3 +136,67 @@ def get_integrals(mol, basis): h2 = ao2mo.restore(1, ao2mo.kernel(mol, basis), basis.shape[1]) return h1, h2 + + +def get_df_integrals(mol, basis=None, auxbasis=None, grad=False): + """ + Compute the density-fitted ERIs and ERI gradients in a specified basis. + + Parameters: + mol (pyscf.gto.Mole): The molecule object. + basis (numpy.ndarray): The basis set (AO->MO transformation coefficients). + If None, uses AO basis. + auxbasis (str): The auxiliary basis for density fitting. + grad (bool): If True, also compute gradient integrals. + + Returns: + If grad=False: + cd_array (numpy.ndarray): Cholesky decomposed ERIs in the specified basis. + If grad=True: + cd_array, deriv_cderi (tuple): Cholesky decomposed ERIs and their gradients. + """ + + # Set auxillary basis + auxmol = df.addons.make_auxmol(mol, auxbasis=auxbasis) + naux = auxmol.nao + + # ints_3c is the 3-center integral tensor (ij|P), where i and j are the + # indices of AO basis and P is the auxiliary basis + ints_3c2e = df.incore.aux_e2(mol, auxmol, intor='int3c2e') + # ints_2c2e is the (P|Q) integrals + ints_2c2e = auxmol.intor('int2c2e') + vals, vecs = np.linalg.eigh(ints_2c2e) + assert(len(vals[vals < 1.e-15]) == 0) # PSD + + metric = np.array(np.dot(vecs * (1 / np.sqrt(vals)) , vecs.conj().T)) + cd_array = np.einsum('PQ,ijP->Qij', metric, ints_3c2e) + + # Transform to the specified basis if provided + if basis is not None: + cd_array = np.einsum('Pij,ai,bj->Pab', cd_array, basis, basis, optimize='optimal') + + # Full 4c integrals can be reconstructed as: + #explicit_df_eri = np.einsum('Pij,Pkl->ijkl', cd_array, cd_array) + + if grad: + # Now consider gradient integrals. The 4c integrals we want to approximate are: + # (d/dx i j | k l) + # grad_4c_ints = mol.intor("int2e_ip1", comp=3) + + # We can get the integrals ( d/dx i, j | P) + ints_3c2e_ip1 = df.incore.aux_e2(mol, auxmol, intor='int3c2e_ip1', comp=3) + # Use the same metric as before + deriv_cderi = np.einsum('PQ,xijP -> xijQ', metric, ints_3c2e_ip1) + + # Transform to the specified basis if provided + if basis is not None: + deriv_cderi = np.einsum('Pxij,ai,bj->Pxab', deriv_cderi, basis, basis, optimize='optimal') + + # Full 4c derivative integrals can be reconstructed as: + # df_grad_4c_ints = np.einsum('xijP,Pkl->xijkl', deriv_cderi, cd_array) + + return cd_array, deriv_cderi + + else: + return cd_array + \ No newline at end of file diff --git a/evcont/excited_utils.py b/evcont/excited_utils.py new file mode 100644 index 0000000..deb1bc9 --- /dev/null +++ b/evcont/excited_utils.py @@ -0,0 +1,314 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Tue Oct 24 15:01:50 2023 + +@author: katalar +""" + +import numpy as np + +import numpy.linalg as LA + +from pyscf import gto + +from pyscf.scf import hf + +from evcont.FCI_EVCont import FCI_EVCont_obj + +from evcont.ab_initio_eigenvector_continuation import approximate_multistate_OAO + +from evcont.electron_integral_utils import get_basis, get_integrals + +import sys + + +def make_trdm1(mol, one_rdm, vec_i, vec_j): + """ + Predicted 1-body transition reduced density matrices from eigenvector continuation + in the SAO basis (orbital basis) + + Input: + mol: pyscf.gto.Mole() + Molecule object + one_rdm: Training one-rdm matrix ndarray(ntrain, ntrain, nao, nao) + One or a list of 1-body reduced density matrices (in SAO basis) + + Returns: + predicted_one_trdm_sao: ndarray(nao, nao) + Predicted 1-body transition reduced density matrix in SAO basis + """ + + predicted_one_trdm = np.einsum("i,ijkl,j->kl", vec_i, one_rdm, vec_j, optimize="optimal") + + return predicted_one_trdm + +def make_rdm1(mol, one_rdm, vec): + """ + Predicted 1-body reduced density matrices from eigenvector continuation + in the SAO basis (orbital basis) + """ + + predicted_one_rdm = make_trdm1(mol, one_rdm, vec, vec) + + return predicted_one_rdm + + +def trans_dip_moment(mol, one_trdm_inp, ref=None): + """ + Compute transition dipole moment from 1-body transition reduced density matrices + in SAO basis (electronic contribution only, no nuclear dipole) + + Input: + mol: pyscf.gto.Mole() + Molecule object + one_trdm_inp: list of ndarray(mol.nao, mol.nao) or ndarray(mol.nao, mol.nao) + One or a list of 1-body transition reduced density matrices in SAO basis + ref: ndarray(3,), optional + Reference point for dipole gauge. If None, uses nuclear charge center. + + Returns: + mol_tdip or mol_tdip_l: ndarray(3,) or list of ndarray(3,) + Transition dipole moments (units A.U.) + """ + + # Set gauge for dipole integrals + if ref is None: + charges = mol.atom_charges() + coords = mol.atom_coords() + nuc_charge_center = np.einsum('z,zx->x', charges, coords) / charges.sum() + ref = nuc_charge_center + + # Get dipole integrals in AO basis and transform to SAO basis + basis = get_basis(mol) + + with mol.with_common_orig(ref): + dip_ints_ao = mol.intor_symmetric('int1e_r', comp=3) + + # Transform dipole integrals from AO to SAO: dip_ints_sao[x,i,j] = basis[a,i] * dip_ints_ao[x,a,b] * basis[b,j] + dip_ints = np.einsum('ai,xab,bj->xij', basis, dip_ints_ao, basis, optimize='optimal') + + single = True + if not isinstance(one_trdm_inp, list): + one_trdm_l = [one_trdm_inp] + else: + one_trdm_l = one_trdm_inp + single = False + + # Iterate over transition RDMs + mol_tdip_l = [] + for one_trdm in one_trdm_l: + el_tdip = np.einsum('xij,ji->x', dip_ints, one_trdm).real + mol_tdip_l += [el_tdip] + + if not single: + return mol_tdip_l + else: + return mol_tdip_l[0] + + +def dip_moment(mol, one_rdm_inp): + """ + Compute dipole moment from 1-body reduced density matrices in SAO basis + (Based on pyscf implementation) + + Input: + mol: pyscf.gto.Mole() + Molecule object + one_rdm_inp: list of ndarray(mol.nao, mol.nao) or ndarray(mol.nao, mol.nao) + One or a list of 1-body reduced density matrices in SAO basis + + Returns: + mol_dip or mol_dip_l: ndarray(3,) or list of ndarray(3,) + Dipole moments of the molecule (units A.U.) + """ + + charges = mol.atom_charges() + coords = mol.atom_coords() + + # Set gauge to nuclear charge center and transform integrals to SAO basis + nuc_charge_center = np.einsum('z,zx->x', charges, coords) / charges.sum() + basis = get_basis(mol) + + #nucl_dip = np.einsum('i,ix->x', charges, coords) + nucl_dip = np.einsum('i,ix->x', charges, coords - nuc_charge_center) + + with mol.with_common_orig(nuc_charge_center): + ao_dip = mol.intor_symmetric('int1e_r', comp=3) + + # Transform dipole integrals from AO to SAO basis + dip_ints = np.einsum('ai,xab,bj->xij', basis, ao_dip, basis, optimize='optimal') + + single = True + if not isinstance(one_rdm_inp,list): + one_rdm_l = [one_rdm_inp] + else: + one_rdm_l = one_rdm_inp + single = False + + # Iterate over one_rdm s + mol_dip_l = [] + for one_rdm in one_rdm_l: + el_dip = np.einsum('xij,ji->x', dip_ints, one_rdm).real + mol_dip = nucl_dip - el_dip + mol_dip_l += [mol_dip] + + if not single: + return mol_dip_l + else: + return mol_dip_l[0] + + #dipole_mom = hf.dip_moment(mol, one_rone_rdm, unit='au') + #return dipole_mom + +def oscillator_strength(en_i, en_j, tran_dipole): + """ + Compute the oscillator strength from the energies and transition dipole moment + f_ij = 2/3 (E_i - E_j) * | t_ij |^2 + + """ + + return 2/3. * (en_i-en_j)* LA.norm(tran_dipole)**2. + + +def print_excited(en, dip_mom, trans_dip, osc_str=None): + """ + Print excited state information from eigenvector continuations + """ + + nstate = len(en) + # Add dummy first element for ground state index + trans_dip_n = [np.array([0.,0.,0.])] + trans_dip + if osc_str is not None: + osc_str_n = [0.] + osc_str + + if osc_str is None: + pr_format = '{:5} {:8.5f} {:6.3f} {:6.3f} {:6.3f} |{:8.5f} {:6.3f} {:6.3f} {:6.3f} |{:8.5f}' + print('{:6} {:15} {:25} {:30}'.format('state','DeltaE (H)','dipole moment (au)','transition dipole moment (au)')) + for i in range(nstate): + print(pr_format.format(i, en[i]-en[0],*dip_mom[i],LA.norm(dip_mom[i]),*trans_dip_n[i],LA.norm(trans_dip_n[i]))) + print() + + else: + pr_format = '{:5} {:8.5f} {:6.3f} {:6.3f} {:6.3f} |{:8.5f} {:6.3f} {:6.3f} {:6.3f} |{:8.5f} {:8.5f} ' + print('{:6} {:15} {:25} {:30}'.format('state','DeltaE (H)','dipole moment (au)','transition dipole moment (au)')) + for i in range(nstate): + print(pr_format.format(i, en[i]-en[0],*dip_mom[i],LA.norm(dip_mom[i]),*trans_dip_n[i],LA.norm(trans_dip_n[i]),osc_str_n[i])) + print() + + +if __name__ == '__main__': + + from pyscf import fci + + cibasis='canonical' + nroots_evcont = 6 + cisolver = fci.direct_spin0.FCI() + + natom = 2 # Fixed, don't change + + def get_mol(geometry): + mol = gto.Mole() + + mol.build( + atom=[ + ("Li", geometry[0]), + ("H", geometry[1]), + ], + #basis="aug-cc-pVDZ", + basis="cc-pVDZ", + #basis="631-G*", + #basis="sto-6g", + symmetry=True, + unit="Bohr", + verbose=0 + ) + + return mol + + ang2bohr = 1.88973 + equilibrium_dist = 1.5957*ang2bohr + + #equilibrium_dist = 1.78596 + + equilibrium_pos = np.array([(x * equilibrium_dist, 0.0, 0.0) for x in range(10)]) + + training_stretches = np.array([0.0, 0.5, -0.5, 1.0, -1.0]) + + trainig_dists = equilibrium_dist + training_stretches + + #trainig_dists = [1.0, 1.8, 2.6] + + continuation_object = FCI_EVCont_obj(nroots=nroots_evcont, + cibasis=cibasis, + cisolver=cisolver) + + + # Generate training data + prepare training models + for i, dist in enumerate(trainig_dists): + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + mol = get_mol(positions) + continuation_object.append_to_rdms(mol) + + for di, dist in enumerate(trainig_dists): + print() + print('Training dist: %.3f'%dist) + print() + + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + mol = get_mol(positions) + + basis = get_basis(mol,basis_type=cibasis) + h1, h2 = get_integrals(mol, basis) + + # Continuation + en, vec = approximate_multistate_OAO(mol, + continuation_object.one_rdm, + continuation_object.two_rdm, + continuation_object.overlap, + nroots=nroots_evcont) + + # Continuation dipole moment from one-rdms of individual states + one_rdm_predicted_l = [make_rdm1(mol, continuation_object.one_rdm, vec[ii]) for ii in range(len(vec))] + dipole_moment_predicted_l = dip_moment(mol,one_rdm_predicted_l) + + # Transition rdms and dipole moments from ground to excited states + one_trdm_predicted_l = [make_trdm1(mol, continuation_object.one_rdm, vec[0], vec[ii]) for ii in range(1,len(vec))] + transition_dipmoment_predicted_l = trans_dip_moment(mol, one_trdm_predicted_l) + + # Oscillator strength + osc_strength_predicted_l = [oscillator_strength(en[ii], en[0], transition_dipmoment_predicted_l[ii-1]) for ii in range(1,len(vec))] + + print_excited(en,dipole_moment_predicted_l, transition_dipmoment_predicted_l,osc_strength_predicted_l) + #print() + + # FCI + e_all, fcivec_all = cisolver.kernel(h1, h2, mol.nao, mol.nelec, + nroots=nroots_evcont,) + #tol = 1.e-14,max_space=30, + #max_cycle=250) + e_all += mol.energy_nuc() + assert np.allclose(e_all, en) + + one_rdm_l = [cisolver.make_rdm1(fcivec_all[ii],mol.nao,mol.nelec) for ii in range(len(fcivec_all))] + one_rdm_l = [basis.dot(one_rdm_l[ii]).dot(basis.T) for ii in range(len(one_rdm_l))] + + #dipole_moment_l = [dip_moment(mol,one_rdm_l[i]) for i in range(len(one_rdm_l))] + dipole_moment_l = dip_moment(mol,one_rdm_l) + + # Transition dipole moment + one_trdm_l = [cisolver.trans_rdm1(fcivec_all[0],fcivec_all[ii],mol.nao,mol.nelec) for ii in range(1,len(fcivec_all))] + one_trdm_l = [basis.dot(one_trdm_l[ii]).dot(basis.T) for ii in range(len(one_trdm_l))] + #one_trdm_l = [np.einsum("ji,jk,kl->il",basis, one_trdm_l[ii], basis, optimize="optimal") for ii in range(len(one_trdm_l))] + + transition_dipmoment_l = trans_dip_moment(mol, one_trdm_l) + + osc_strength_l = [oscillator_strength(e_all[ii], e_all[0], transition_dipmoment_l[ii-1]) for ii in range(1,len(fcivec_all))] + + print_excited(e_all,dipole_moment_l, transition_dipmoment_l,osc_strength_l) + + assert np.allclose(dipole_moment_l,dipole_moment_predicted_l) + + #assert np.allclose(transition_dipmoment_l,transition_dipmoment_predicted_l) + + #1/0 diff --git a/evcont/logging_utils.py b/evcont/logging_utils.py new file mode 100644 index 0000000..7e7ee3c --- /dev/null +++ b/evcont/logging_utils.py @@ -0,0 +1,43 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Wed Jul 16 11:13:37 2025 + +Logging utilities + +@author: Kemal Atalar +""" + +import logging +import time +from contextlib import contextmanager + +# ---- Logger Setup ---- +logger = logging.getLogger("timing_logger") +logger.setLevel(logging.INFO) + +if not logger.hasHandlers(): # prevent duplicate handlers if imported multiple times + ch = logging.StreamHandler() + formatter = logging.Formatter('[%(asctime)s] %(message)s', datefmt='%H:%M:%S') + ch.setFormatter(formatter) + logger.addHandler(ch) + +# ---- Timing Context Manager ---- +@contextmanager +def log_time(section_name): + start = time.time() + logger.info(f"Start: {section_name}") + yield + end = time.time() + logger.info(f"End: {section_name} | Elapsed: {end - start:.3f} s") + +# ---- Function Timing Decorator ---- +def timeit(func): + def wrapper(*args, **kwargs): + logger.info(f"Start: {func.__name__}") + t0 = time.time() + result = func(*args, **kwargs) + logger.info(f"End: {func.__name__} | Elapsed: {time.time() - t0:.3f} s") + return result + return wrapper + diff --git a/evcont/low_rank_utils.py b/evcont/low_rank_utils.py new file mode 100644 index 0000000..28c674e --- /dev/null +++ b/evcont/low_rank_utils.py @@ -0,0 +1,1732 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Fri Jun 6 17:13:27 2025 + +Low-rank decomposition of 2-body (transition) reduced density matrices + +Mixed decomposition: + Joint ED: Joint eigenvalue decomposition where each low-rank vector + contribute to both Coulomb and exchange channels. + Coulomb SVD: SVD in the Coulomb grouping of the state indices + +Function here include: + - Static and dynamic truncation of the decomposition based on + eval magnitude/Hamiltonian error + - Subspace Hamiltonian construction from low-rank vectors + - Vectorizing of the low-rank vectors for fast inference + - Computation of JK and JK grad builds from low-rank vectors + - Gradient inference from the low-rank representation + +Note: Not compatible with complex RDMs as it stands; e.g. vectors from SVD are +assumed to be real-valued. Can be changed in the future if necessary. + +@author: Kemal Atalar +""" + +import multiprocessing as mp +mp.set_start_method("fork", force=True) + +from multiprocessing import Process, Pipe + +import numpy as np +import sys +import itertools +import time +import math +import os + +import scipy +from scipy.linalg import eigh, svd +from scipy.sparse.linalg import eigsh, svds + +import pyscf +from pyscf import scf, gto, ao2mo, fci, lib, df + +from pyscf.df.grad.rhf import get_jk as get_jk_grad_df +from pyscf.df.grad.rhf import Gradients as mf_grad_df + +from pyscf.grad.rhf import get_jk as get_jk_grad_nodf +from pyscf.grad.rhf import Gradients as mf_grad_nodf + +from evcont.electron_integral_utils import get_loewdin_trafo, get_integrals, get_df_integrals, get_basis +from evcont.logging_utils import logger, log_time, timeit + +######################################################################## +# Iterative solver with timeout fallback to full diagonalization +def _eigsh_worker(mat, k, conn, which='LM',use_svd=False): + try: + if not use_svd: + conn.send(eigsh(mat, k=k, which=which)) + else: + conn.send(svds(mat, k=k, which=which)) + except Exception as e: + conn.send(("error", str(e))) + finally: + conn.close() + +def try_iterative_diag(mat, k, which='LM', use_svd=False, max_time=100000): + parent_conn, child_conn = Pipe() + p = Process(target=_eigsh_worker, args=(mat, k, child_conn, which, use_svd)) + p.start() + p.join(timeout=max_time) + + if p.is_alive(): + p.terminate() + p.join() + print("[timeout] Falling back to full diagonalization...") + + if parent_conn.poll(): + result = parent_conn.recv() + + if isinstance(result, tuple) and isinstance(result[0], str) and result[0] == "error": + print(f"[error] {result[1]}") + else: + print("[iterative] Success") + return result + + # Full fallback + print("[full] Running full eigh/svds...") + if not use_svd: + return eigh(mat) + else: + return svd(mat) + +######################################################################## +@timeit +def reduce_2rdm(rdm1, rdm2, ovlp, + truncation_style='eigval',nvecs=10, eval_thr=0.1, ham_thr=0.001, + save_diag=False, Jdiag_only=True, + relax_amp=True, opt_no_diag=True, relax_after=True, + use_svd=False, svd_weight = 1.0, + iterative=False, nit=None, max_iter_time=10000, + mol=None,train_en=None, + min_eval=None): + """ + Function to compress the 2-transition-RDM between a pair of training + states using joint decomposition and diagonal corrections, and + a variety of truncation criteria. + + Input: + rdm1 (np.array([n,n])): 1-body reduced density matrix between two training states + rdm2 (np.array([n,n,n,n])): 2-body reduced density matrix between two training states + ovlp (float): Overlap between the training state pair + truncation_style (str): + Criteria of low rank truncation. + Available options: {'eigval' (default) : Choose vectors whose eigenvalue**2 is more than eval_thr, + 'nvec' : Choose 'nvec' highest eval**2 number of vectors, + 'ham' : Choose the minimum number of vectors such that error in the subspace + Hamiltonian matrix elements is less than 'ham_thr' + 'ham_en' : Choose the minimum number of vectors such that error in the subspace + Hamiltonian*overlap matrix elements is less than 'ham_thr'} + nvecs (int): Number of low rank vectors to include + eval_thr (float): Threshold to choose vectors based on their eval**2 + ham_thr (float): Threshold to choose vectors based on their H matrix elements (Hartree units) + + save_diag (bool): Whether to save the diagonal corrections to make low-rank 2RDM diagonal elements exact. + Jdiag_only (bool): Whether only Coulomb diagonals are used in the inference when determining the + truncation + + # Parameters for amplitude relaxation of joint decomposition + relax_amp (bool): Whether to perform amplitude relaxation after selecting the low-rank vectors + in the joint decomposition. + opt_no_diag (bool): Whether to remove diagonal elements from the relaxation. + Only applicable if relax_amp is True and save_diag is True. + relax_after (bool): Whether to perform amplitude relaxation after selecting the low-rank + vectors based on Hamiltonian error vs relaxing during the selection process. + Only applicable if relax_amp is True and truncation_style is 'ham' or 'ham_en'. + + # Parameters relating to SVD decomposition + use_svd (bool): Whether to perform SVD in addition to the joint decomposition and choose the more compact representation. + svd_weight (float): Weight for choosing the SVD decomposition over joint ED. + + # Parameters for iterative diagonalization of the 2RDM for larger systems + iterative (bool): Whether to use iterative diagonalization for the low-rank decomposition. + nit (int): Number of eigenvalues/vectors to compute in the iterative diagonalization. Only applicable if iterative is True. + max_iter_time (float): Maximum time in seconds to allow for the iterative diagonalization + before falling back to full diagonalization. + + # Parameters relevant for Hamiltonian error truncation + mol (pyscf Mole object): Molecule object that is used for computing the Hamiltonian error + train_en (float): Energy of the training geometry used for the truncation + min_eval (float or None): When using Hamiltonian-based truncation, include all tied vectors + whose eigenvalue magnitude equals the boundary eigenvalue magnitude. If provided, + ties are detected against this value; otherwise the boundary value is used. + + Output: + lowrank_vecs (vals_trunc, vecs_trunc): Low rank eigenvalues and eigenvectors of the 2tRDM decomposition. + diagonals (np.array([3,n,n])): Diagonal correction matrices to be added to the low-rank reconstructed 2RDM. + joint (bool): Whether the low-rank vectors are from the joint decomposition (True) or SVD (False) + + """ + + # Matrix to decompose + mat_decomp = rdm2 + + norb = rdm1.shape[0] + norb_sq = norb * norb + + # Ensure the matrix is hermitian + if not np.allclose(rdm2.reshape((norb_sq, norb_sq)), rdm2.reshape((norb_sq,norb_sq)).T): + print('Warning: 2RDM was not Hermitian.') + # Hermitise + rdm2 = 0.5 * (rdm2 + np.einsum('...abcd->...cdab',rdm2.conj())) + mat_decomp = rdm2 + + # Matrix to decompose in the joint decomposition + # Refactor the 2(t)RDM such that its eigenvectors solely correponds to Coulomb grouping + mat_decomp = 4/3*mat_decomp + 2/3*np.einsum('ijkl->ilkj',mat_decomp) + + # Check the nit is given is iterative is True + if iterative and nit is None: + print('Error in reduce_rdm: nit is not given for iterative diagonalization') + sys.exit() + + # Diagonalize + if not iterative: + evals, evecs = scipy.linalg.eigh(mat_decomp.reshape((norb_sq, norb_sq))) + else: + evals, evecs = try_iterative_diag(mat_decomp.reshape((norb_sq, norb_sq)), + k=nit, + which='LM', + max_time=max_iter_time) + + rightvecs = None + joint = True # Joint decomp + + ######################################################################## + #### Select low rank vectors + + # Choose at least one vector (lower bound for dynamic truncation) + min_nvecs = 1 + + # Unless the 2RDM is close to zero + if np.linalg.norm(rdm2) < 1e-8: + print("Warning in reduce_2rdm: RDM2 is close to zero; setting base_count to 0 to avoid selecting vectors based on noise.") + min_nvecs = 0 + + # Make sure the mol and training energy is given for this truncation + if truncation_style in ['ham','ham_en']: + if mol is None or train_en is None: + print('Error in reduce_2rdm: Insufficient input for decomposition based on Hamiltonian error.') + sys.exit() + + elif truncation_style not in ['eigval','nvec']: + print('Unknown truncation_style in reduce_2rdm: %s'%truncation_style) + sys.exit() + + # Fixed truncation based on 'nvecs' parameter + # Or dynamic truncation based on the eigenvalue magnitude / Hamiltonian error + if truncation_style in ['eigval','nvec']: + + lowrank_vecs_joint = select_lowrank(evals, evecs, norb, rightvecs=rightvecs, + truncation_style=truncation_style, + nvecs=nvecs, eval_thr=eval_thr, min_nvec=min_nvecs, + relax_amp=relax_amp, rdm2=rdm2, + remove_diagopt=(opt_no_diag and save_diag), + jdiag_only=Jdiag_only) + + elif truncation_style in ['ham','ham_en']: + + ham_selector = select_lowrank_ham_relaxed if (joint and relax_amp and not relax_after) else select_lowrank_ham + out_ham = ham_selector(evals, evecs, joint, norb, + rdm2, rdm1, ovlp,save_diag, + mol, train_en, Jdiag_only, + rightvecs=rightvecs, + truncation_style=truncation_style, + ham_thr=ham_thr, min_nvec=min_nvecs, + min_eval=min_eval, + relax_amp=relax_amp, + remove_diagopt=(opt_no_diag and save_diag)) + lowrank_vecs_joint = out_ham[:-1] + #ham_err_joint = out_ham[-1] + + # Choose between joint decomposition and SVD based on the compactness of the representation (weighted by svd_weight) and/or Hamiltonian error + if not use_svd: + lowrank_vecs = lowrank_vecs_joint + + else: + + if not iterative: + evecs2, evals2, rightvecs2 = scipy.linalg.svd(rdm2.reshape((norb_sq, norb_sq))) + else: + evecs2, evals2, rightvecs2 = svds(rdm2.reshape((norb_sq, norb_sq)), k=nit, which='LM') + + if truncation_style in ['eigval','nvec']: + + lowrank_vecs_svd = select_lowrank(evals2, evecs2, norb, rightvecs=rightvecs2, + truncation_style=truncation_style, + nvecs=nvecs, eval_thr=eval_thr, min_nvec=min_nvecs, + relax_amp=False) + + elif truncation_style in ['ham','ham_en']: + + out_ham = select_lowrank_ham(evals2, evecs2, False, norb, + rdm2, rdm1, ovlp,save_diag, + mol, train_en, Jdiag_only, + rightvecs=rightvecs2, + truncation_style=truncation_style, + ham_thr=ham_thr, min_nvec=min_nvecs, + min_eval=min_eval, relax_amp=False) + lowrank_vecs_svd = out_ham[:-1] + #ham_err_svd = out_ham[-1] + + # Check which one is more compact + if truncation_style in ['ham','ham_en']: + if len(lowrank_vecs_joint[0])*svd_weight < len(lowrank_vecs_svd[0]): + lowrank_vecs = lowrank_vecs_joint + #ham_err = ham_err_joint + else: + print('**Using SVD') + lowrank_vecs = lowrank_vecs_svd + joint = False + #ham_err = ham_err_svd + + # TODO: Add considerations for norm error; not just compactness + # SVD as well + elif truncation_style in ['eigval']: + if len(lowrank_vecs_joint[0])*svd_weight < len(lowrank_vecs_svd[0]): + lowrank_vecs = lowrank_vecs_joint + else: + print('**Using SVD') + lowrank_vecs = lowrank_vecs_svd + joint = False + #ham_err = ham_err_svd + else: + if np.abs(lowrank_vecs_joint[0]).max() < np.abs(lowrank_vecs_svd[0]).max(): + lowrank_vecs = lowrank_vecs_joint + else: + print('**Using SVD') + lowrank_vecs = lowrank_vecs_svd + joint = False + + ######################################################################## + if not save_diag: + diagonals = None + + else: + + remainder = rdm2 - reconstruct_rdm2_joint(lowrank_vecs,joint=joint) + + # Save diagonals of the remainder + diagonals = np.zeros([3,norb,norb]) + for (i,j) in itertools.product(range(norb), range(norb)): + diagonals[0, i, j] = remainder[ i, i, j, j] + if (not Jdiag_only) and i != j: + diagonals[1, i, j] = remainder[ i, j, i, j] + diagonals[2, i, j] = remainder[ i, j, j, i] + + reconstructed_rdm2 = reconstruct_rdm2_joint(lowrank_vecs, diagonals, joint=joint) + if mol is not None: + # Compute the energy error as well + h1, h2 = get_integrals(mol, get_basis(mol)) + ham_inferred = 0.5*np.einsum('pqrs,pqrs->', reconstructed_rdm2, h2,optimize='optimal') + \ + np.einsum('pq,pq->', rdm1, h1) + # Recompute train hamiltonian + train_ham = 0.5*np.einsum('pqrs,pqrs->', rdm2, h2,optimize='optimal') + \ + np.einsum('pq,pq->', rdm1, h1) + + #ham_error = ham_inferred - train_en*ovlp + ham_error = ham_inferred - train_ham + #print(f'-- Verifying en_train*ovlp - ham_train: {train_en*ovlp - train_ham:.1e} or en_train*|ovlp| - ham_train: {train_en*np.abs(ovlp) - train_ham:.1e}') + print(f'-- Norm error: {np.linalg.norm(reconstructed_rdm2 - rdm2):.1e}, Hamiltonian error: {ham_error:.1e}, nvecs: {len(lowrank_vecs[0])}') + + else: + print(f'-- Norm error: {np.linalg.norm(reconstructed_rdm2 - rdm2):.1e}, nvecs: {len(lowrank_vecs[0])}') + + return lowrank_vecs, diagonals, joint + +def reconstruct_rdm2_joint(lowrank_vecs, diagonals=None, joint=True): + """ + Reconstructing the 2RDM + """ + lr_vals, lr_vecs, lr_rightvecs = lowrank_vecs + rdm2_i = np.einsum('ija,a,akl->ijkl',lr_vecs, lr_vals, lr_rightvecs.conj(),optimize='optimal') + # Add exchange part as well + if joint: + rdm2_i -= 0.5*np.einsum('kja,a,ail->ijkl',lr_vecs, lr_vals, lr_rightvecs.conj(),optimize='optimal') + + if diagonals is not None: + norb = diagonals.shape[-1] + for (i,j) in itertools.product(range(norb), range(norb)): + rdm2_i[ i, i, j, j] += diagonals[0, i, j] + if i != j: + rdm2_i[ i, j, i, j] += diagonals[1, i, j] + rdm2_i[ i, j, j, i] += diagonals[2, i, j] + + return rdm2_i + +""" +def reconstruct_subspace(lowrank_vecs, h2, ntrain=3): + subspace_h = np.zeros([ntrain, ntrain]) + for vecs in lowrank_vecs.values(): + rdm2_i = reconstruct_rdm2_joint(vecs) + subspace_h += +""" + +######################################################################## +@timeit +def lowrank_hamiltonian(mol, one_RDM, S, lowrank_vecs, diagonals=None, + sao_basis=None, density_fit=True, df_basis=None, + Jdiag_only=True, sao_diag=False, + hermitian=True, + debug=False): + """ + Construct subspace Hamiltonian using the low-rank decomposition of + 2-transition-cumulant + + Input: + mol (Mole object): pySCF mole object at the test geometry + + """ + + ### Preliminaries + ntrain = S.shape[0] + + norb = one_RDM.shape[-1] + norb_sq = norb * norb + + if diagonals is not None: + use_diag = True + else: + use_diag = False + + # Initiate the mean field object to use DF integrals (no need to use kernel) + #mol.symmetry = False + if density_fit: + mf = scf.RHF(mol).density_fit(auxbasis=df_basis) + mf_grad = mf_grad_df + get_jk = mf.with_df.get_jk + + else: + mf = scf.RHF(mol) + mf_grad = mf_grad_nodf + get_jk = mf.get_jk + + # AO to SAO basis transformation + if sao_basis is None: + sao_basis = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + # 1-electron integrals with DF + h1_ao = mf.get_hcore() + h1e_sao = np.einsum('ai,ab,bj->ij', sao_basis, h1_ao, sao_basis) + + # Get ERIs in SAO basis (using density fitting) for debugging + if debug or sao_diag: + # Run calculation to fill the MF object + #mf.scf() + #Lpq = mf.with_df._cderi + #print(Lpq.shape) + + # Alternative - without using MF object + Lpq = get_df_integrals(mol,auxbasis=df_basis) + Lpq = lib.pack_tril(Lpq) + #print(Lpq.shape) + + Lpq_sao = ao2mo._ao2mo.nr_e2(Lpq, sao_basis, + (0, sao_basis.shape[1], 0, sao_basis.shape[1]),aosym="s2",mosym="s2") + Lpq_sao = lib.unpack_tril(Lpq_sao) + df_eri_sao = lib.einsum('Pij,Pkl->ijkl', Lpq_sao, Lpq_sao,optimize='optimal') + + ### 1-body contributions + subspace_h = np.einsum('...kl,kl->...', one_RDM, h1e_sao,optimize='optimal') + + # Check if low-rank vectors have been vectorized + if ('vals' in lowrank_vecs): + vectorized = True + # Whether to use training point symmetry; keep original 'hermitian' arg as default + hermitian = lowrank_vecs.get('hermitian', hermitian) + + else: + vectorized = False + + # Construct the subspace Hamiltonian + if not vectorized: + + # Vectorized version is more efficient but keeping this here for testing purposes + for bra in range(ntrain): + if hermitian: + ket_max = bra+1 + else: + ket_max = ntrain + for ket in range(ntrain): + + nvec = lowrank_vecs[(bra,ket)][0].shape[0] + use_joint = lowrank_vecs[(bra, ket)][3] + + # Transform the low-rank vecs into + lr_vecs_group = np.ascontiguousarray(lowrank_vecs[(bra, ket)][1].transpose((2,0,1))) + sao_basis_T = np.ascontiguousarray(sao_basis.T) + lr_vecs_ao = ao2mo._ao2mo.nr_e2(lr_vecs_group, sao_basis_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + lr_vecs_ao = lr_vecs_ao.reshape((nvec,norb,norb)) + + # JK build + if use_joint: + vj_list, vk_list = get_jk(dm=lr_vecs_ao.transpose(0,2,1), hermi=0) # Specify hermiticity per case + subspace_h[bra,ket] += 0.5*np.einsum('aij,aij,a->', vj_list - 0.5 * vk_list, lr_vecs_ao.conj(), lowrank_vecs[(bra, ket)][0]) + + # J build from SVD + else: + lr_rightvecs_group = np.ascontiguousarray(lowrank_vecs[(bra, ket)][2]) + sao_basis_arr = np.ascontiguousarray(sao_basis) + lr_rightvecs_ao = ao2mo._ao2mo.nr_e2(lr_rightvecs_group, sao_basis_arr, + (0, norb, 0, norb), aosym='s1', mosym='s1') + lr_rightvecs_ao = lr_rightvecs_ao.reshape((nvec,norb,norb)) + + # For reference; direct contraction: + #rdm2_i = np.einsum('ija,a,akl->ijkl',lr_vecs, lr_vals, lr_rightvecs.conj(),optimize='optimal') + + # For test purposes, explicitly reconstruct RDM and contract with ERIs + if debug: + lowrank_vecs_i = lowrank_vecs[(bra, ket)][0],lowrank_vecs[(bra, ket)][1],lowrank_vecs[(bra, ket)][2] + rdm2_i = reconstruct_rdm2_joint(lowrank_vecs_i, None, joint=use_joint) + subspace_h[bra,ket] += 0.5*np.einsum('ijkl,ijkl->', rdm2_i, df_eri_sao) + + else: + vj_list, vk_list = get_jk(dm=lr_rightvecs_ao, hermi=0, with_k=False) # Specify hermiticity per case + subspace_h[bra,ket] += 0.5*np.einsum('aij,aij,a->', vj_list, lr_vecs_ao, lowrank_vecs[(bra, ket)][0]) + + if use_diag: + + # Avoid transforming DF array, and do everything via Coulomb and exchange builds + diag_ao_1 = np.einsum('ij,wi,xi->jwx',diagonals[bra, ket, 0, :, :], sao_basis, sao_basis) + vj = get_jk(dm = diag_ao_1, hermi=0, with_k=False)[0] + subspace_h[bra, ket] += 0.5 * np.einsum('yj,zj,jyz->', sao_basis, sao_basis, vj) + + if not Jdiag_only: + diag_ao_23 = np.einsum('ij,wi,yi->jwy',diagonals[bra, ket, 1, :, :] + diagonals[bra, ket, 2, :, :], sao_basis, sao_basis) + vk = get_jk(dm = diag_ao_23, hermi=0, with_j=False)[1] + subspace_h[bra, ket] += 0.5 * np.einsum('xj,zj,jxz->', sao_basis, sao_basis, vk) + + else: + nvec = lowrank_vecs['vals'].shape[2] + + ### Joint ED inference + if lowrank_vecs['has_ed']: + # Grouped JK builds + lr_vecs_grouped = lowrank_vecs['vecs_stacked'] + + # Transform the low-rank vecs into + lr_vecs_grouped_c = np.ascontiguousarray(lr_vecs_grouped) + sao_basis_T = np.ascontiguousarray(sao_basis.T) + lr_vecs_ao = ao2mo._ao2mo.nr_e2(lr_vecs_grouped_c, sao_basis_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + lr_vecs_ao = lr_vecs_ao.reshape((lr_vecs_grouped.shape[0],norb,norb)) + + # JK build + vj_list, vk_list = get_jk(dm=lr_vecs_ao.transpose(0,2,1), hermi=0) # Specify hermiticity per case + vhf = vj_list - 0.5*vk_list + + # Reindex to separate bra, ket, nvec indices + vhf = unpack_vec(vhf, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + lr_vecs_ao = unpack_vec(lr_vecs_ao, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + + # Contruction for subspace Hamiltonian + subspace_h += 0.5*np.einsum('xyaij,xyaij,xya->xy', vhf, lr_vecs_ao, lowrank_vecs['vals'][:,:,:vhf.shape[2]],optimize='optimal') + + ### Coulomb SVD inference + if lowrank_vecs['has_svd']: + + # Grouped J Builds + svd_vecs_grouped = lowrank_vecs['vecs_svd_stacked'] + svd_rightvecs_grouped = lowrank_vecs['rightvecs_stacked'] + + # Transform the low-rank vecs into AO basis + svd_rightvecs_grouped_c = np.ascontiguousarray(svd_rightvecs_grouped) + sao_basis_T = np.ascontiguousarray(sao_basis.T) + svd_rightvecs_ao = ao2mo._ao2mo.nr_e2(svd_rightvecs_grouped_c, sao_basis_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + svd_rightvecs_ao = svd_rightvecs_ao.reshape((svd_rightvecs_grouped.shape[0],norb,norb)) + + svd_vecs_grouped_c = np.ascontiguousarray(svd_vecs_grouped) + sao_basis_T = np.ascontiguousarray(sao_basis.T) + svd_vecs_ao = ao2mo._ao2mo.nr_e2(svd_vecs_grouped_c, sao_basis_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + svd_vecs_ao = svd_vecs_ao.reshape((svd_vecs_grouped.shape[0],norb,norb)) + + # J build + vj_list, _ = get_jk(dm=svd_rightvecs_ao, hermi=0, with_k=False) # Specify hermiticity per case + + # Reindex to separate bra, ket, nvec indices + vj = unpack_vec(vj_list, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + svd_vecs_ao = unpack_vec(svd_vecs_ao, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + + subspace_h += 0.5*np.einsum('xyaij,xyaij,xya->xy', vj, svd_vecs_ao, lowrank_vecs['vals'][:,:,:vj.shape[2]],optimize='optimal') + + if use_diag: + + if not sao_diag: + # Transform the low-rank vecs into AO basis for J builds + diagJ_ao = np.einsum('Nij,wi,xi->Njwx', diagonals[0], sao_basis, sao_basis) + + # Flatten + orig_shape = diagJ_ao.shape[:2] + flat_diagJ_ao = diagJ_ao.reshape(orig_shape[0] * orig_shape[1], *diagJ_ao.shape[2:]) + + # JK Builds + vj_list = get_jk(dm=flat_diagJ_ao, hermi=0, with_k=False)[0] + + # Unflatten + vj_unflat = vj_list.reshape(*orig_shape, *flat_diagJ_ao.shape[1:]) + vj = unstack_tril(vj_unflat, hermitian=hermitian) + + subspace_h += 0.5 * np.einsum('yj,zj,XYjyz->XY', sao_basis, sao_basis, vj) + + else: + diagJ_unpack = unstack_tril(diagonals[0], hermitian=hermitian) + subspace_h += 0.5 * np.einsum('XYij,Pii,Pjj->XY', diagJ_unpack, Lpq_sao, Lpq_sao) + + if not Jdiag_only: + + if not sao_diag: + # Transform the low-rank vecs into AO basis for K builds + diagK_ao = np.einsum('Nij,wi,yi->Njwy', diagonals[1], sao_basis, sao_basis) + + # Flatten + orig_shape = diagK_ao.shape[:2] + flat_diagK_ao = diagK_ao.reshape(orig_shape[0] * orig_shape[1], *diagK_ao.shape[2:]) + + # JK Builds + vk_list = get_jk(dm=flat_diagK_ao, hermi=0, with_j=False)[1] + + # Unflatten + vk_unflat = vk_list.reshape(*orig_shape, *flat_diagK_ao.shape[1:]) + vk = unstack_tril(vk_unflat, hermitian=hermitian) + + subspace_h += 0.5 * np.einsum('xj,zj,XYjxz->XY', sao_basis, sao_basis, vk) + + else: + diagK_unpack = unstack_tril(diagonals[1], hermitian=hermitian) + subspace_h += 0.5 * np.einsum('XYij,Pij,Pij->XY', diagK_unpack, Lpq_sao, Lpq_sao) + + if hermitian: + # Set the upper triangle + subspace_h[np.triu_indices(ntrain)] = subspace_h.T[np.triu_indices(ntrain)].conj() + + # Check that hermitian + #assert np.allclose(subspace_h, subspace_h.T.conj()) + + return subspace_h + +############################################################################### +@timeit +def get_jk_builds(mol, lowrank_vecs, + diagonals=None, Jdiag_only=True, sao_diag=False, + ao_mo_trafo=None, + density_fit=False, df_basis=None, + df_response=False, verbose=False): + """ + Precompute the J(K) builds for the low-rank vectors for fast inference + """ + # AO to SAO basis transformation + if ao_mo_trafo is None: + ao_mo_trafo = get_loewdin_trafo(mol.intor("int1e_ovlp")) + + # Initiate the mean field object to use DF integrals (no need to use kernel) + #mol.symmetry = False + if density_fit: + mf = scf.RHF(mol).density_fit(auxbasis=df_basis) + mf_grad = mf_grad_df + get_jk = mf.with_df.get_jk + #get_jk_grad = get_jk_grad_df + + else: + mf = scf.RHF(mol) + mf_grad = mf_grad_nodf + get_jk = mf.get_jk + #get_jk_grad = get_jk_grad_nodf + + # Check if diagonals are given + if diagonals is None: + use_diag = False + else: + use_diag = True + + ###################################################### + # Check if low-rank vectors have been vectorized + if ('vals' in lowrank_vecs): + vectorized = True + # Whether to use training point symmetry + hermitian = lowrank_vecs['hermitian'] + + else: + print('Error in get_jk_builds: Lowrank vectors are not in the vectorized format. Run "continuation_object.vectorize_lowrank()".') + sys.exit() + + ###################################################### + if lowrank_vecs['has_ed']: + norb = lowrank_vecs['vecs_stacked'].shape[-1] + else: + norb = lowrank_vecs['vecs_svd_stacked'].shape[-1] + nvec = lowrank_vecs['vals'].shape[2] + ntrain = lowrank_vecs['vals'].shape[0] + + ###################################################### + ######### COMPUTE PRELIMINARIES & FOCK BUILDS + ###################################################### + # Initiate grad object + grad_obj = mf_grad(mf) + # Set verbose level for detailed timing output (6 or higher shows timer_debug1) + if verbose: + grad_obj.verbose = 6 + # TODO: Add auxbasis_response in the future, for now ignore it + grad_obj.auxbasis_response = df_response + + ### Joint ED inference + if lowrank_vecs['has_ed']: + # Grouped JK builds + lr_vecs_grouped = lowrank_vecs['vecs_stacked'] + + with log_time("AO transform (1)"): + # Transform the low-rank vecs into + lr_vecs_grouped_c = np.ascontiguousarray(lr_vecs_grouped) + ao_mo_trafo_T = np.ascontiguousarray(ao_mo_trafo.T) + lr_vecs_ao = ao2mo._ao2mo.nr_e2(lr_vecs_grouped_c, ao_mo_trafo_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + lr_vecs_ao = lr_vecs_ao.reshape((lr_vecs_grouped.shape[0],norb,norb)) + + # JK build + with log_time("JK Builds (1)"): + # Stick to PySCF's own threading; joblib overhead slowed plain get_jk + vj_list, vk_list = get_jk(dm=lr_vecs_ao.transpose(0,2,1), hermi=0) + + vhf = vj_list - 0.5*vk_list + + # Grad JK builds + # TODO: Add auxbasis_response in the future, for now ignore it + with log_time("JK Grad Builds (2)"): + lr_vecs_ao = np.ascontiguousarray(lr_vecs_ao) + lr_vecs_ao_T = np.ascontiguousarray(lr_vecs_ao.transpose(0,2,1)) + + vj_grad_list, vk_grad_list = grad_obj.get_jk(dm=lr_vecs_ao, hermi=0) + vj_grad_list_t, vk_grad_list_t = grad_obj.get_jk(dm=lr_vecs_ao_T, hermi=0) + + + vhf_grad = vj_grad_list - 0.5*vk_grad_list + vhf_grad_t = vj_grad_list_t - 0.5*vk_grad_list_t + + #vhf_aux = np.einsum('aamn,a->mn',vj_list.aux - vk_list.aux*.5,lr_vals[ii],optimize='optimal') + #vhf_aux = (vj_list.aux - vk_list.aux*.5).sum((0,1))#,lr_vals[ii]) + #grad_i += vhf_aux + # Reindex to separate bra, ket, nvec indices + vhf = unpack_vec(vhf, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + lr_vecs_ao = unpack_vec(lr_vecs_ao, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + lr_vecs = unpack_vec(lr_vecs_grouped, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + vhf_grad = unpack_grad_vec(vhf_grad, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + vhf_grad_t = unpack_grad_vec(vhf_grad_t, lowrank_vecs['pairloc'],hermitian=hermitian, nbra=ntrain) + + if df_response: + vhf_aux = unpack_grad_aux(vj_grad_list.aux - 0.5*vk_grad_list.aux, + lowrank_vecs['pairloc'], + lowrank_vecs['vals'],hermitian=hermitian) + + vhf_aux_t = unpack_grad_aux(vj_grad_list_t.aux - 0.5*vk_grad_list_t.aux, + lowrank_vecs['pairloc'], + lowrank_vecs['vals'],hermitian=hermitian) + + vhf_grad = lib.tag_array(vhf_grad, aux=np.array(vhf_aux)) + vhf_grad_t = lib.tag_array(vhf_grad_t, aux=np.array(vhf_aux_t)) + + + ### Coulomb SVD inference + if lowrank_vecs['has_svd']: + + # Grouped J Builds + svd_vecs_grouped = lowrank_vecs['vecs_svd_stacked'] + svd_rightvecs_grouped = lowrank_vecs['rightvecs_stacked'] + + with log_time("AO transform (2)"): + # Transform the low-rank vecs into AO basis + svd_rightvecs_grouped_c = np.ascontiguousarray(svd_rightvecs_grouped) + ao_mo_trafo_T = np.ascontiguousarray(ao_mo_trafo.T) + svd_rightvecs_ao = ao2mo._ao2mo.nr_e2(svd_rightvecs_grouped_c, ao_mo_trafo_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + svd_rightvecs_ao = svd_rightvecs_ao.reshape((svd_rightvecs_grouped.shape[0],norb,norb)) + + svd_vecs_grouped_c = np.ascontiguousarray(svd_vecs_grouped) + ao_mo_trafo_T = np.ascontiguousarray(ao_mo_trafo.T) + svd_vecs_ao = ao2mo._ao2mo.nr_e2(svd_vecs_grouped_c, ao_mo_trafo_T, + (0, norb, 0, norb), aosym='s1', mosym='s1') + svd_vecs_ao = svd_vecs_ao.reshape((svd_vecs_grouped.shape[0],norb,norb)) + + svd_vecs_ao = np.ascontiguousarray(svd_vecs_ao) + svd_rightvecs_ao = np.ascontiguousarray(svd_rightvecs_ao) + + # J builds + with log_time("J Builds (2)"): + vj_r_list, _ = get_jk(dm=svd_rightvecs_ao, hermi=0, with_k=False) + vj_l_list, _ = get_jk(dm=svd_vecs_ao, hermi=0, with_k=False) + + # Grad JK builds + # TODO: Add auxbasis_response in the future, for now ignore it + with log_time("J Grad Builds (2)"): + vj_lgrad_list = grad_obj.get_j(dm=svd_vecs_ao, hermi=0) + vj_rgrad_list = grad_obj.get_j(dm=svd_rightvecs_ao, hermi=0) + + # Reindex to separate bra, ket, nvec indices + vj_right = unpack_vec(vj_r_list, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + vj_left = unpack_vec(vj_l_list, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + svd_lvecs_ao = unpack_vec(svd_vecs_ao, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + svd_rvecs_ao = unpack_vec(svd_rightvecs_ao, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + svd_lvecs = unpack_vec(svd_vecs_grouped, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + svd_rvecs = unpack_vec(svd_rightvecs_grouped, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + + vj_l_grad = unpack_grad_vec(vj_lgrad_list, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + vj_r_grad = unpack_grad_vec(vj_rgrad_list, lowrank_vecs['pairloc_svd'],hermitian=hermitian, nbra=ntrain) + + if df_response: + vj_l_aux = unpack_grad_aux(vj_l_grad.aux, + lowrank_vecs['pairloc_svd'], + lowrank_vecs['vals'],hermitian=hermitian) + + vj_r_aux = unpack_grad_aux(vj_r_grad.aux, + lowrank_vecs['pairloc_svd'], + lowrank_vecs['vals'],hermitian=hermitian) + + vj_l_grad = lib.tag_array(vj_l_grad, aux=np.array(vj_l_aux)) + vj_r_grad = lib.tag_array(vhf_grad_t, aux=np.array(vj_r_aux)) + + # Diagonal JK builds + if use_diag and not sao_diag: + + ### First the Coulomb build + + # Expand out the 'i' indices for J builds + diagJ_ao = np.einsum('Nij,wj,xj->Niwx',diagonals[0], ao_mo_trafo, ao_mo_trafo,optimize='optimal') + + # Flatten + orig_shape = diagJ_ao.shape[:2] + flat_diagJ_ao = diagJ_ao.reshape(orig_shape[0] * orig_shape[1], *diagJ_ao.shape[2:]) + + # JK Builds + with log_time("Diag J Builds (1)"): + vj_list = get_jk(dm=flat_diagJ_ao, hermi=0, with_k=False)[0] + + # JK grad builds + with log_time("Diag Grad J Builds (1)"): + vj_grad_list = grad_obj.get_j(dm=flat_diagJ_ao.transpose(0,2,1), hermi=0) + + # Unflatten + vj_unflat = vj_list.reshape(*orig_shape, *vj_list.shape[1:]) + vj = unstack_tril(vj_unflat,hermitian=hermitian) + + vj_grad_unflat = vj_grad_list.reshape(*orig_shape, *vj_grad_list.shape[1:]) + vj_grad = unstack_tril(vj_grad_unflat,hermitian=hermitian) + + if not Jdiag_only: + # Transform the low-rank vecs into + diagK_ao = np.einsum('Nij,wi,yi->Njwy',diagonals[1], ao_mo_trafo, ao_mo_trafo,optimize='optimal') + diagK_ao_T = np.einsum('Nji,wi,yi->Njwy',diagonals[1], ao_mo_trafo, ao_mo_trafo,optimize='optimal') + + # Flatten + orig_shape = diagK_ao.shape[:2] + flat_diagK_ao = diagK_ao.reshape(orig_shape[0] * orig_shape[1], *diagK_ao.shape[2:]) + flat_diagK_ao_T = diagK_ao_T.reshape(orig_shape[0] * orig_shape[1], *diagK_ao.shape[2:]) + + # JK Builds + with log_time("Diag K Builds (1)"): + vk_list = get_jk(dm=flat_diagK_ao, hermi=0, with_j=False)[1] + vk_t_list = get_jk(dm=flat_diagK_ao_T, hermi=0, with_j=False)[1] + + # JK grad builds + with log_time("Diag Grad K Builds (1)"): + vk_grad_list = grad_obj.get_jk(dm=flat_diagK_ao.transpose(0,2,1), hermi=0, with_j=False)[1] + + # Unflatten + vk_unflat = vk_list.reshape(*orig_shape, *flat_diagK_ao.shape[1:]) # (3, 4, 5, 6) + vk = unstack_tril(vk_unflat,hermitian=hermitian) + + vk_t_unflat = vk_t_list.reshape(*orig_shape, *flat_diagK_ao.shape[1:]) # (3, 4, 5, 6) + vk_t = unstack_tril(vk_t_unflat,hermitian=hermitian) + + vk_grad_unflat = vk_grad_list.reshape(*orig_shape, *vk_grad_list.shape[1:]) + vk_grad = unstack_tril(vk_grad_unflat,hermitian=hermitian) + + else: + vk, vk_t, vk_grad = None, None, None + + # Function OUTPUT + ed_return, svd_return, diag_return = None, None, None + + if lowrank_vecs['has_ed']: + ed_return = (lr_vecs, lr_vecs_ao, vhf, vhf_grad, vhf_grad_t) + + if lowrank_vecs['has_svd']: + svd_return = (svd_lvecs, svd_rvecs, svd_lvecs_ao, svd_rvecs_ao, vj_left, vj_right, vj_l_grad, vj_r_grad) + + if use_diag and not sao_diag: + diag_return = (vj, vj_grad, vk, vk_t, vk_grad) + + return ed_return, svd_return, diag_return + + +############################################################################### +# Amplitude relaxation for joint decomposition +def relax_coef(rdm2, vec, diag=False, jdiag_only=True): + # Basis functions + B = np.einsum('ija,kla->aijkl', vec, vec, optimize='optimal') \ + - 0.5 * np.einsum('ila,kja->aijkl', vec, vec, optimize='optimal') + + # Build masked tensors first (target and basis), then form normal equations. + if diag: + norb_loc = rdm2.shape[0] + mask = np.ones((norb_loc, norb_loc, norb_loc, norb_loc), dtype=rdm2.dtype) + + for (i, j) in itertools.product(range(norb_loc), range(norb_loc)): + mask[i, i, j, j] = 0.0 + if not jdiag_only and i != j: + mask[i, j, i, j] = 0.0 + mask[i, j, j, i] = 0.0 + + rdm2_opt = rdm2 * mask + B_opt = B * mask[None, :, :, :, :] + else: + rdm2_opt = rdm2 + B_opt = B + + # Gram matrix and projection in the masked space + G = np.einsum('aijkl,bijkl->ab', B_opt, B_opt, optimize='optimal') + b = np.einsum('ijkl,aijkl->a', rdm2_opt, B_opt, optimize='optimal') + + # Solve the linear system + coef, res, rank, sval = np.linalg.lstsq(G, b, rcond=None) + + return coef + +def select_lowrank(evals, evecs, norb, + rightvecs=None, + truncation_style='eigval',nvecs=10, eval_thr=0.1, min_nvec=0, + relax_amp=True, rdm2=None, remove_diagopt=True, jdiag_only=True): + """ + Function to select low-rank vectors from the eigendecomposition + """ + # Check if right eigenvectors are given + if rightvecs is None: + rightvecs = evecs.T + + # Sort the eigenstates by the square of their eigenvalue + idx = (-np.power(evals, 2)).argsort() + evals_sort = evals[idx] + evecs_sort = evecs[:,idx] + rightvecs_sort = rightvecs[idx,:] + + # Truncate through either eigvals or a given number of vectors + if truncation_style == 'eigval': + nvecs = len(evals_sort[np.power(evals_sort,2) > eval_thr]) + nvecs = max(min_nvec, nvecs) + + #norb = np.sqrt(evecs_sort.shape[0],dtype=int) + vals_trunc = evals_sort[:nvecs] + vecs_trunc = evecs_sort[:,:nvecs].reshape((norb, norb, nvecs)) + rightvecs_trunc = rightvecs_sort[:nvecs,:].reshape((nvecs,norb, norb)) + + # Relax amplitudes + if relax_amp: + vals_new = relax_coef(rdm2, vecs_trunc, diag=remove_diagopt, jdiag_only=jdiag_only) + print('Relaxing amplitudes; norm of change in vals:', np.linalg.norm(vals_new - vals_trunc)) + vals_trunc = vals_new + + return vals_trunc, vecs_trunc, rightvecs_trunc + +def select_lowrank_ham(evals, evecs, joint, norb, + rdm2, rdm1, ovlp, save_diag, + mol, training_energy, Jdiag_only, + rightvecs=None, + density_fit=True, + truncation_style='ham', + ham_thr=0.001, + min_nvec=1, min_eval=None, + relax_amp=True, + remove_diagopt=False, + ): + """ + Select a low rank decomposition of the RDM based on the error on + subspace hamiltonian. + + Strategy: + 1. Sort eigenstates by the square of their eigenvalue + 2. Choose a base set based on min_nvec and min_eval (using magnitude to handle negative evals) + 3. Compute initial rdm2_cum from this base set using reconstruct_rdm2_joint + 4. Continue adding vectors and checking convergence on hamiltonian error until converged + """ + # Check if right eigenvectors are given + if rightvecs is None: + rightvecs = evecs.T + + # Sort the eigenstates by the square of their eigenvalue + idx = (-np.power(evals, 2)).argsort() + evals_sort = evals[idx] + evecs_sort = evecs[:,idx] + rightvecs_sort = rightvecs[idx,:] + + # For direct contraction + h1, h2 = get_integrals(mol, get_basis(mol)) + e1 = lib.einsum('ij,ij->', h1, rdm1) + + # Exact element of subspace Hamiltonian - compute explicitly instead of using train_en*ovlp + ham_training = 0.5*lib.einsum('pqrs,pqrs->', rdm2, h2, optimize='optimal') + \ + lib.einsum('pq,pq->', rdm1, h1) + #ham_training = ovlp * training_energy + + max_nvec = min(norb * norb, evals_sort.shape[0]) + + # Step 2: Determine base set from min_nvec and min_eval + base_count = min_nvec + if min_eval is not None: + # Include all vectors with eigenvalue magnitude^2 >= min_eval^2 + # evals_sort is already sorted by magnitude^2 descending + min_eval_mag2 = np.power(min_eval, 2) + eval_mag2 = np.power(evals_sort, 2) + + # Find how many vectors have magnitude^2 >= threshold + count_above_threshold = np.sum(eval_mag2 >= min_eval_mag2) + base_count = max(min_nvec, count_above_threshold) + + # Step 3: Compute initial rdm2_cum from base set + base_vals = evals_sort[:base_count] + base_vecs = evecs_sort[:, :base_count].reshape((norb, norb, base_count)) + base_rightvecs = rightvecs_sort[:base_count, :].reshape((base_count, norb, norb)) + + rdm2_cum = reconstruct_rdm2_joint((base_vals, base_vecs, base_rightvecs), + diagonals=None, joint=joint) + + # Apply diagonal corrections to base set if requested + if save_diag: + for (i, j) in itertools.product(range(norb), range(norb)): + rdm2_cum[i, i, j, j] = rdm2[i, i, j, j] + if not Jdiag_only and i != j: + rdm2_cum[i, j, i, j] = rdm2[i, j, i, j] + rdm2_cum[i, j, j, i] = rdm2[i, j, j, i] + + # Compute Hamiltonian error for the base set + e2_base = 0.5 * lib.einsum('ijkl,ijkl->', h2, rdm2_cum,optimize='optimal') + ham_base = e1 + e2_base + + if truncation_style == 'ham': + ham_err_base = ham_training - ham_base + elif truncation_style == 'ham_en': + ham_err_base = training_energy - ham_base / ovlp + + # Step 4: Continue from base_count, checking convergence + ham_err_list = [ham_err_base] + nvec_select = base_count + + for k in range(base_count, max_nvec): + # Add contribution from k-th vector + v_left = evecs_sort[:, k].reshape(norb, norb) + v_right = rightvecs_sort[k, :].reshape(norb, norb) + contrib = evals_sort[k] * np.einsum('ij,kl->ijkl', v_left, v_right.conj()) + if joint: + contrib -= 0.5 * evals_sort[k] * np.einsum('kj,il->ijkl', v_left, v_right.conj()) + rdm2_cum += contrib + + # Apply diagonal corrections if requested + if save_diag: + for (i, j) in itertools.product(range(norb), range(norb)): + rdm2_cum[i, i, j, j] = rdm2[i, i, j, j] + if (not Jdiag_only) and i != j: + rdm2_cum[i, j, i, j] = rdm2[i, j, i, j] + rdm2_cum[i, j, j, i] = rdm2[i, j, j, i] + + # Compute Hamiltonian and error + e2 = 0.5 * lib.einsum('ijkl,ijkl->', h2, rdm2_cum,optimize='optimal') + ham_new = e1 + e2 + + if truncation_style == 'ham': + ham_err = ham_training - ham_new + elif truncation_style == 'ham_en': + ham_err = training_energy - ham_new / ovlp + + ham_err_list.append(ham_err) + + # Check for convergence (two-step robustness check) + # Need at least 2 errors to compare, and both current and previous must be below threshold + if len(ham_err_list) >= 2: + if abs(ham_err_list[-1]) <= ham_thr and abs(ham_err_list[-2]) <= ham_thr: + # Both current and previous errors are below threshold + # This means convergence was achieved at the previous iteration k-1 + nvec_select = k + break + + # Check for last iteration to set to max if never converged + if k == max_nvec - 1: + nvec_select = max_nvec + + # Return truncated decomposition + vals_trunc = evals_sort[:nvec_select] + vecs_trunc = evecs_sort[:, :nvec_select].reshape((norb, norb, nvec_select)) + rightvecs_trunc = rightvecs_sort[:nvec_select, :].reshape((nvec_select, norb, norb)) + + # Relax amplitudes + if relax_amp: + vals_new = relax_coef(rdm2, vecs_trunc, diag=remove_diagopt, jdiag_only=Jdiag_only) + print('Relaxing amplitudes; norm of change in vals:', np.linalg.norm(vals_new - vals_trunc)) + vals_trunc = vals_new + + return vals_trunc, vecs_trunc, rightvecs_trunc, ham_err_list[-1] if ham_err_list else 0.0 + + +def select_lowrank_ham_relaxed(evals, evecs, joint, norb, + rdm2, rdm1, ovlp, save_diag, + mol, training_energy, Jdiag_only, + rightvecs=None, + density_fit=True, + truncation_style='ham', + ham_thr=0.001, + min_nvec=1, min_eval=None, + relax_amp=True, + remove_diagopt=False, + ): + """ + Select a joint low-rank decomposition based on Hamiltonian error while + re-optimizing the selected amplitudes after each truncation step. + + This path is intended only for the joint-decomposition case with + amplitude relaxation enabled. + """ + # This selector is intentionally narrow to keep behavior predictable. + if not joint: + raise ValueError('select_lowrank_ham_relaxed only supports joint=True') + if not relax_amp: + raise ValueError('select_lowrank_ham_relaxed requires relax_amp=True') + + # For Hermitian ED, right vectors are the transpose if not supplied. + if rightvecs is None: + rightvecs = evecs.T + + # Process candidates from largest |eigenvalue|^2 to smallest. + idx = (-np.power(evals, 2)).argsort() + evals_sort = evals[idx] + evecs_sort = evecs[:, idx] + rightvecs_sort = rightvecs[idx, :] + + # Precompute one-electron and exact training Hamiltonian pieces once. + h1, h2 = get_integrals(mol, get_basis(mol)) + e1 = lib.einsum('ij,ij->', h1, rdm1) + ham_training = 0.5 * lib.einsum('pqrs,pqrs->', rdm2, h2, optimize='optimal') + \ + lib.einsum('pq,pq->', rdm1, h1) + + max_nvec = min(norb * norb, evals_sort.shape[0]) + + base_count = min_nvec + if min_eval is not None: + # Optionally force all vectors above a minimum |eigenvalue| into the base set. + min_eval_mag2 = np.power(min_eval, 2) + eval_mag2 = np.power(evals_sort, 2) + count_above_threshold = np.sum(eval_mag2 >= min_eval_mag2) + base_count = max(min_nvec, count_above_threshold) + + def build_relaxed_prefix(prefix_count): + # Build the current prefix basis and re-fit amplitudes on the full target rdm2. + vecs_prefix = evecs_sort[:, :prefix_count].reshape((norb, norb, prefix_count)) + rightvecs_prefix = rightvecs_sort[:prefix_count, :].reshape((prefix_count, norb, norb)) + vals_prefix = relax_coef( + rdm2, + vecs_prefix, + diag=remove_diagopt, + jdiag_only=Jdiag_only, + ) + + # Reconstruct the relaxed 2RDM from fitted amplitudes. + rdm2_prefix = reconstruct_rdm2_joint( + (vals_prefix, vecs_prefix, rightvecs_prefix), + diagonals=None, + joint=True, + ) + + # If diagonal terms are saved separately in production, mirror that here + # so selection is based on the same Hamiltonian expression used downstream. + if save_diag: + for (i, j) in itertools.product(range(norb), range(norb)): + rdm2_prefix[i, i, j, j] = rdm2[i, i, j, j] + if not Jdiag_only and i != j: + rdm2_prefix[i, j, i, j] = rdm2[i, j, i, j] + rdm2_prefix[i, j, j, i] = rdm2[i, j, j, i] + + # Evaluate the Hamiltonian error of this relaxed prefix. + ham_prefix = e1 + 0.5 * lib.einsum('ijkl,ijkl->', h2, rdm2_prefix, optimize='optimal') + if truncation_style == 'ham': + ham_err = ham_training - ham_prefix + elif truncation_style == 'ham_en': + ham_err = training_energy - ham_prefix / ovlp + else: + raise ValueError(f'Unknown truncation_style in select_lowrank_ham_relaxed: {truncation_style}') + + return vals_prefix, vecs_prefix, rightvecs_prefix, ham_err + + ham_err_list = [] + vals_trunc = np.zeros((0,)) + vecs_trunc = np.zeros((norb, norb, 0)) + rightvecs_trunc = np.zeros((0, norb, norb)) + nvec_select = base_count + + # Grow the prefix one vector at a time, relax, then test Hamiltonian error. + for prefix_count in range(base_count, max_nvec + 1): + vals_prefix, vecs_prefix, rightvecs_prefix, ham_err = build_relaxed_prefix(prefix_count) + ham_err_list.append(ham_err) + vals_trunc = vals_prefix + vecs_trunc = vecs_prefix + rightvecs_trunc = rightvecs_prefix + nvec_select = prefix_count + + # Two-step robustness check: require two consecutive below-threshold errors. + if len(ham_err_list) >= 2: + if abs(ham_err_list[-1]) <= ham_thr and abs(ham_err_list[-2]) <= ham_thr: + nvec_select = prefix_count + break + + # Defensive slicing; arrays are already prefix-sized, but this keeps output explicit. + vals_trunc = vals_trunc[:nvec_select] + vecs_trunc = vecs_trunc[:, :, :nvec_select] + rightvecs_trunc = rightvecs_trunc[:nvec_select] + + return vals_trunc, vecs_trunc, rightvecs_trunc, ham_err_list[-1] if ham_err_list else 0.0 + +############################################################################### +def stack_lowrank(vecs_lowrank, hermitian=True): + """ + Function to group dynamically chosen low-rank eigenstates for different + bra,ket pairs into a compound index for efficient inference + """ + # Prelim + nbra = list(vecs_lowrank.keys())[-1][0]+1 + norb = vecs_lowrank[(0,0)][1].shape[1] + + # Store the locations of bra,ket pairs in the composite index + pair_loc = {} + + # Start stacking + vecs_lr = [] + vals_lr = [] + nvec_tot = 0 + + # Have a separate on for SVD vectors that only needs J builds + pair_svd_loc = {} + vecs_svd_lr = [] + rightvecs_svd_lr = [] + vals_svd_lr = [] + nsvd_tot = 0 + + for i in range(nbra): + + # Only iterarte through lower triangular indices + if hermitian: + jmax = i+1 + else: + jmax = nbra + + for j in range(jmax): + lr_i = vecs_lowrank[(i,j)] + + nvec_i = lr_i[0].shape[-1] + + # Joint ED + if lr_i[-1]: + vecs_lr.append(lr_i[1].transpose(2,0,1)) + vals_lr.append(lr_i[0]) + + pair_loc[(i,j)] = [nvec_tot, nvec_tot + nvec_i] + + nvec_tot += nvec_i + + # Coulomb SVD + else: + vecs_svd_lr.append(lr_i[1].transpose(2,0,1)) + rightvecs_svd_lr.append(lr_i[2]) + vals_svd_lr.append(lr_i[0]) + + pair_svd_loc[(i,j)] = [nsvd_tot, nsvd_tot + nvec_i] + + nsvd_tot += nvec_i + + # Check if any (t)RDM used ED + has_ed = True + if len(vecs_lr) == 0: + has_ed = False + elif len(np.concatenate(vecs_lr)) == 0: + has_ed = False + + # Check if any (t)RDM used SVD + has_svd = True + if len(vecs_svd_lr) == 0: + has_svd = False + elif len(np.concatenate(vecs_svd_lr)) == 0: + has_svd = False + + # Set up the final dictionary + stacked_lowrank = {} + stacked_lowrank['hermitian'] = hermitian + + if has_ed: + # Joint ED vectors + vecs_stacked = np.concatenate(vecs_lr,axis=0) + vals_stacked = np.concatenate(vals_lr) + + stacked_lowrank['vals'] = vals_stacked + stacked_lowrank['vecs'] = vecs_stacked + stacked_lowrank['pairloc'] = pair_loc + + if has_svd: + # Coulomb SVD vectors + vecs_svd_stacked = np.concatenate(vecs_svd_lr,axis=0) + rightvecs_svd_stacked = np.concatenate(rightvecs_svd_lr,axis=0) + vals_svd_stacked = np.concatenate(vals_svd_lr) + + stacked_lowrank['vals_svd'] = vals_svd_stacked + stacked_lowrank['vecs_svd'] = vecs_svd_stacked + stacked_lowrank['rightvecs_svd'] = rightvecs_svd_stacked + stacked_lowrank['pairloc_svd'] = pair_svd_loc + + return stacked_lowrank, has_svd, has_ed + + +def stack_tril(arr, hermitian=True): + """ + Stack selected blocks from a (n, n, x, x) array into a compact (m, x, x) array. + + If hermitian=True, stacks only the lower triangle (i >= j), + assuming the array is Hermitian in its (n, n) block structure. + + If hermitian=False, stacks the full (i, j) grid in row-major order. + + Parameters: + arr : np.ndarray + Input array of shape (n, n, x, x) + hermitian : bool + Whether to restrict to lower-triangular blocks only + + Returns: + stacked : np.ndarray + Stacked array of shape (m, x, x) where m = n*(n+1)//2 if Hermitian, + or m = n*n if not. + """ + n, _, x, _ = arr.shape + packed = [] + for i in range(n): + jmax = i + 1 if hermitian else n + for j in range(jmax): + packed.append(arr[i, j]) + return np.array(packed) + + +def unstack_tril(packed, hermitian=True): + """ + Unpacks a stacked array of shape (m, ...) into shape (n, n, ...), where the first + axis was previously packed using only the lower triangle (if hermitian=True) or the full (n,n) grid. + + Parameters: + packed : np.ndarray + Input array with shape (m, ...) where m = n*(n+1)//2 (hermitian) or n*n (full) + hermitian : bool + Whether the packed data was from the lower triangle only + + Returns: + arr : np.ndarray + Output array of shape (n, n, ...) + """ + # If the input already appears to be in full (n, n, ...) grid form, just + # return it (but ensure Hermitian symmetry is enforced when requested). + if packed.ndim >= 2 and packed.shape[0] == packed.shape[1]: + arr = packed.copy() + if hermitian: + n = arr.shape[0] + for i in range(n): + for j in range(i + 1, n): + try: + arr[j, i] = arr[i, j].conj() + except Exception: + arr[j, i] = arr[i, j] + return arr + + m = packed.shape[0] + rest_shape = packed.shape[1:] + + # The packed length m can sometimes be both a triangular number and a perfect square + # (e.g., m=36 -> triangular for n=8, square for n=6). Use the `hermitian` flag to + # disambiguate: when hermitian=True prefer triangular (lower-triangle) packing, + # otherwise prefer full-grid (n*n) packing. + + # Check triangular possibility: solve n(n+1)/2 == m + tri_n = int((math.isqrt(1 + 8 * m) - 1) // 2) + is_tri = (tri_n * (tri_n + 1) // 2 == m) + + # Check square possibility: m == n*n + sq_n = int(math.isqrt(m)) + is_sq = (sq_n * sq_n == m) + + # Prefer triangular when requested or when square interpretation is impossible + if is_tri and (hermitian or not is_sq): + n = tri_n + arr = np.zeros((n, n) + rest_shape, dtype=packed.dtype) + idx = 0 + for i in range(n): + for j in range(i + 1): # j <= i + val = packed[idx] + arr[i, j] = val + # Mirror to the upper triangle to restore full matrix symmetry + try: + arr[j, i] = val.conj() + except Exception: + arr[j, i] = val + idx += 1 + return arr + + # If square packing fits (n*n) interpret as full-grid + if is_sq: + n = sq_n + arr = packed.reshape((n, n) + rest_shape).copy() + if hermitian: + for i in range(n): + for j in range(i + 1, n): + try: + arr[j, i] = arr[i, j].conj() + except Exception: + arr[j, i] = arr[i, j] + return arr + + # If none of the interpretations matched, raise an informative error + raise ValueError(f"Invalid packed shape for unstacking: not triangular nor square (m={m})") + + +def stack_diagonal(diagonals, hermitian=True): + """ + Stack 2(t)RDM diagonals for a vectorized inference + """ + + diag_J = diagonals[:,:,0] + diag_K = diagonals[:,:,1] + diagonals[:,:,2] + + stacked_diagJ = stack_tril(diag_J, hermitian=hermitian) + stacked_diagK = stack_tril(diag_K, hermitian=hermitian) + + return (stacked_diagJ, stacked_diagK) + +def unpack_vec(vecs,pair_loc,hermitian=True,nbra=None): + """ + Function to unpack vectors stacked using "stack_lowrank" function + + """ + if nbra is None: + nbra = list(pair_loc.keys())[-1][0]+1 + norb = vecs.shape[1] + nvec_max = np.max([j-i for i,j in pair_loc.values()]) + + vecs_unpacked = np.zeros([nbra, nbra, nvec_max,norb,norb]) + + # Precompute index arrays for batch assignment + for (i, j), (start, end) in pair_loc.items(): + nv = end - start + vecs_unpacked[i, j, :nv] = vecs[start:end] + + return vecs_unpacked + + +def unpack_grad_vec(vecs,pair_loc,hermitian=True,nbra=None): + """ + Function to unpack vectors stacked using "stack_lowrank" function + + """ + if nbra is None: + nbra = list(pair_loc.keys())[-1][0]+1 + norb = vecs.shape[2] + nvec_max = np.max([j-i for i,j in pair_loc.values()]) + + vecs_unpacked = np.zeros([nbra, nbra, nvec_max, 3, norb, norb]) + + # Precompute index arrays for batch assignment + for (i, j), (start, end) in pair_loc.items(): + nv = end - start + vecs_unpacked[i, j, :nv] = vecs[start:end] + + """ + # Check if auxbasis response is computed + #try: + print('auxbasis unpacking') + auxvec = vecs.aux + nat = auxvec.shape[-2] + + aux_unpacked = np.zeros([nbra, nbra, nbra, nbra, nvec_max, nat, 3]) + + # Precompute index arrays for batch assignment + for (i, j), (start, end) in pair_loc.items(): + nv = end - start + aux_unpacked[i,i,j, j, :nv] = auxvec[start:end,start:end] + + vecs_unpacked = lib.tag_array(vecs_unpacked, aux=np.array(aux_unpacked)) + + #except: + # print('No auxbasis') + # None + """ + return vecs_unpacked + + +def unpack_grad_aux(vecs,pair_loc,vals,hermitian=True): + """ + Function to unpack vectors stacked using "stack_lowrank" function + + """ + nbra = list(pair_loc.keys())[-1][0]+1 + + auxvec = vecs + nat = auxvec.shape[-2] + + aux_unpacked = np.zeros([nbra, nbra, nat, 3]) + + # Precompute index arrays for batch assignment + for (i, j), (start, end) in pair_loc.items(): + nv = end - start + aux_unpacked[i,j] = np.einsum('aamn,a->mn',auxvec[start:end,start:end],vals[i,j,:nv],optimize='optimal') + + return aux_unpacked + +def unpack_lowrank(stacked_lowrank,hermitian=True): + """ + For testing; function to unpack both eigenvectors and eigenvectors + from the stacked_lowrank dictionary + """ + + vals_stacked = stacked_lowrank['vals'] + vecs_stacked = stacked_lowrank['vecs'] + pair_loc = stacked_lowrank['pairloc'] + + # Prelim + nbra = list(pair_loc.keys())[-1][0]+1 + + unpacked_vecs = {} + if hermitian: + for i in range(nbra): + for j in range(i+1): + st, en = pair_loc[(i,j)] + vals_i = vals_stacked[st:en] + vecs_i = vecs_stacked[st:en].transpose(1,2,0) + + unpacked_vecs[(i,j)] = vals_i, vecs_i + + else: + for i, j in itertools.product(range(nbra), range(nbra)): + st, en = pair_loc[(i,j)] + vals_i = vals_stacked[st:en] + vecs_i = vecs_stacked[st:en].transpose(1,2,0) + + unpacked_vecs[(i,j)] = vals_i, vecs_i + + return unpacked_vecs + +# Attribute function to vectorize low-rank vectors for EVCont solver classes +def vectorize_lowrank(self, hermitian=True): + + # Make sure a low-rank decomposition has been performed + assert len(self.vecs_lowrank.items()) != 0 + + # Find the largest number of vectors for each bra,ket pair + nbra = self.overlap.shape[0] + norb = self.one_rdm.shape[-1] + # Determine maximum vectors per pair quickly + nvec_max = max((lr[0].shape[-1] for lr in self.vecs_lowrank.values()), default=0) + + # Use stack_lowrank to get packed/staged arrays and flags + stacked, has_svd, has_ed = stack_lowrank(self.vecs_lowrank, hermitian=hermitian) + + # Allocate padded arrays for per-pair fast indexing + vals_lr = np.zeros((nbra, nbra, nvec_max), dtype=float) + vecs_lr = np.zeros((nbra, nbra, nvec_max, norb, norb), dtype=float) + # store number of vectors per pair explicitly to make unpack lossless + nvecs_per_pair = np.zeros((nbra, nbra), dtype=int) + # Only allocate rightvecs if any pair uses ED (joint decomposition) + rightvecs_lr = None + if has_ed: + rightvecs_lr = np.zeros((nbra, nbra, nvec_max, norb, norb), dtype=float) + + # Fill padded arrays from the original per-pair dict (single pass) + for (i, j), lr_i in self.vecs_lowrank.items(): + vals = lr_i[0] + vecs = lr_i[1] + rvecs = lr_i[2] + nvec_i = vals.shape[-1] + + if nvec_i == 0: + continue + + vals_lr[i, j, :nvec_i] = vals + nvecs_per_pair[i, j] = nvec_i + # stored as (norb, norb, nvec) in original, need to transpose to (nvec, norb, norb) + vecs_lr[i, j, :nvec_i] = vecs.transpose(2, 0, 1) + if rightvecs_lr is not None: + rightvecs_lr[i, j, :nvec_i] = rvecs + + # Vectorize diagonal corrections: store both the packed J and the packed K-sum + # for compatibility, but also keep packed components of diag[1] and diag[2] + if ( getattr(self, 'diagonal_lr', None) is not None ) and ( not np.isnan(self.diagonal_lr).all() ): + diagJ = self.diagonal_lr[:, :, 0] + diagK1 = self.diagonal_lr[:, :, 1] + diagK2 = self.diagonal_lr[:, :, 2] + + stacked_diagJ = stack_tril(diagJ, hermitian=hermitian) + stacked_diagKsum = stack_tril(diagK1 + diagK2, hermitian=hermitian) + # Also store separated components for lossless roundtrip + stacked_diagK1 = stack_tril(diagK1, hermitian=hermitian) + stacked_diagK2 = stack_tril(diagK2, hermitian=hermitian) + + self.diagonal_vectorized = np.stack((stacked_diagJ, stacked_diagKsum)) + # components stored separately to enable exact reconstruction + self.diagonal_vectorized_components = np.stack((stacked_diagK1, stacked_diagK2)) + else: + self.diagonal_vectorized = None + self.diagonal_vectorized_components = None + #self.diagonal_K = diagK + + # Set this low-rank description + # Build final dict + self.lowrank_vectorized = { + 'vals': vals_lr, + #'vecs': vecs_lr, + 'hermitian': hermitian, + 'has_ed': has_ed, + 'has_svd': has_svd, + 'nvecs': nvecs_per_pair, + 'norb': norb, + 'ntrain': nbra + } + + if has_ed: + #self.lowrank_vectorized['rightvecs'] = rightvecs_lr + self.lowrank_vectorized['vecs_stacked'] = stacked['vecs'] + self.lowrank_vectorized['pairloc'] = stacked['pairloc'] + + if has_svd: + self.lowrank_vectorized['rightvecs_stacked'] = stacked['rightvecs_svd'] + self.lowrank_vectorized['vecs_svd_stacked'] = stacked['vecs_svd'] + self.lowrank_vectorized['pairloc_svd'] = stacked['pairloc_svd'] + + +############################################################################### + +def unpack_vectorized_lowrank(self): + """ + Reconstruct per-(bra,ket) lowrank dict (`self.vecs_lowrank`) and + `self.diagonal_lr` from the vectorized representations + (`self.lowrank_vectorized` and `self.diagonal_vectorized`). This + enables appending new training points when the object is already + vectorized. Mirrors the helper previously attached to the FCI_EVCont_obj. + """ + if not getattr(self, 'lowrank_vectorized', None): + return + + lv = self.lowrank_vectorized + vals = lv.get('vals') + vecs = lv.get('vecs') + rightvecs = lv.get('rightvecs', None) + pairloc = lv.get('pairloc', {}) + pairloc_svd = lv.get('pairloc_svd', {}) + hermitian = lv.get('hermitian', True) + nvecs_per_pair = lv.get('nvecs', None) + + # Basic shapes + nbra = vals.shape[0] + norb = vecs.shape[-1] + + # Reconstruct per-pair vecs_lowrank + vecs_lowrank = {} + for i, j in itertools.product(range(nbra), range(nbra)): + vals_ij = vals[i, j] + vecs_ij = vecs[i, j] + + # Determine nvec for this pair: prefer explicit stored count if present + if nvecs_per_pair is not None: + nvec = int(nvecs_per_pair[i, j]) + else: + mask = np.any(np.abs(vecs_ij) > 1e-12, axis=(1, 2)) + nvec = int(mask.sum()) + + if nvec == 0: + # keep zero-length arrays for consistency + vals_i = np.zeros((0,)) + vecs_i = np.zeros((norb, norb, 0)) + right_i = np.zeros((0, norb, norb)) if rightvecs is not None else np.zeros((0, norb, norb)) + else: + vals_i = vals_ij[:nvec].copy() + # stored in vectorize_lowrank as (nvec, norb, norb) + vecs_i = vecs_ij[:nvec].transpose(1, 2, 0).copy() + if rightvecs is not None: + right_i = rightvecs[i, j, :nvec].copy() + else: + # try to recover from stacked SVD/rightvecs_stacked if present + if 'rightvecs_stacked' in lv and (i, j) in pairloc_svd: + st, en = pairloc_svd[(i, j)] + right_i = lv['rightvecs_stacked'][st:en].copy() + elif 'rightvecs_stacked' in lv and (j, i) in pairloc_svd: + st, en = pairloc_svd[(j, i)] + right_i = lv['rightvecs_stacked'][st:en].copy() + else: + right_i = np.zeros((nvec, norb, norb)) + + # Determine whether this pair used joint ED (present in pairloc in either order) + use_joint = (i, j) in pairloc or (j, i) in pairloc + + vecs_lowrank[(i, j)] = (vals_i, vecs_i, right_i, use_joint) + + # Reconstruct diagonals from diagonal_vectorized if available + if getattr(self, 'diagonal_vectorized', None) is not None: + diag_stack = self.diagonal_vectorized + # diag_stack shape (2, m, norb, norb) + diagJ_packed = diag_stack[0] + diagKsum_packed = diag_stack[1] + + # If explicit components were stored at vectorization time, use them to reconstruct exact diag[1] and diag[2] + comp = getattr(self, 'diagonal_vectorized_components', None) + if comp is not None: + diagK1_packed = comp[0] + diagK2_packed = comp[1] + + diagJ_unpacked = unstack_tril(diagJ_packed, hermitian=hermitian) + diagK1_unpacked = unstack_tril(diagK1_packed, hermitian=hermitian) + diagK2_unpacked = unstack_tril(diagK2_packed, hermitian=hermitian) + + diagonal_lr = np.zeros((nbra, nbra, 3, norb, norb), dtype=diagJ_unpacked.dtype) + diagonal_lr[:, :, 0, :, :] = diagJ_unpacked + diagonal_lr[:, :, 1, :, :] = diagK1_unpacked + diagonal_lr[:, :, 2, :, :] = diagK2_unpacked + + else: + # Fallback: only K-sum available + diagJ_unpacked = unstack_tril(diagJ_packed, hermitian=hermitian) + diagK_unpacked = unstack_tril(diagKsum_packed, hermitian=hermitian) + + # Recreate original 3-component diagonal array: [J, K_sum, 0] + diagonal_lr = np.zeros((nbra, nbra, 3, norb, norb), dtype=diagJ_unpacked.dtype) + diagonal_lr[:, :, 0, :, :] = diagJ_unpacked + # diagK_unpacked contains sum of original diag[1] + diag[2] + diagonal_lr[:, :, 1, :, :] = diagK_unpacked + diagonal_lr[:, :, 2, :, :] = 0.0 + + else: + diagonal_lr = None + + # Assign back to self so append_to_rdms can operate on the per-pair structures + self.vecs_lowrank = vecs_lowrank + self.diagonal_lr = diagonal_lr + + + +def rdm2_from_rdm1(rdm1, ovlp): + """ + 1-body contribution to the 2-(transition) reduced density matrices + """ + rdm1_contribution = ( np.einsum('ij,kl->jilk', rdm1, rdm1) - 0.5 * np.einsum('kj,il->jilk', rdm1, rdm1) ) * 1/ovlp + return rdm1_contribution + +def build_diag_mask(norb): + """ + Function that returns a mask array for diagonal matrices of + 4D tensor with dimensions norb^4 + """ + # Build training overlaps and (t)RDMs (note that hermiticity should be used for performant code, as well as no norb^4 objects stored). + diag_mask = np.zeros((norb, norb, norb, norb)) + for (i,j) in itertools.product(range(norb), range(norb)): + diag_mask[i,i,j,j] = diag_mask[i,j,i,j] = diag_mask[i,j,j,i] = 1.0 + + return diag_mask + + diff --git a/examples/NAC/00-fci_nac.py b/examples/NAC/00-fci_nac.py new file mode 100644 index 0000000..2c0fd17 --- /dev/null +++ b/examples/NAC/00-fci_nac.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 +""" +Minimal example: compute FCI nonadiabatic couplings (NACs) with PySCF + evcont. + +This script: +1) builds a small H-chain geometry, +2) computes direct FCI energies, gradients, and NACs, +3) computes continuation energies, gradients, and NACs, +4) compares continuation results against direct FCI. +""" + +import numpy as np +from pyscf import fci, gto + +from evcont.FCI_EVCont import FCI_EVCont_obj +from evcont.FCI_NAC import get_FCI_energy_with_grad_and_NAC +from evcont.ab_initio_gradients_loewdin import get_multistate_energy_with_grad_and_NAC + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def pair_labels(nroots): + """All unique state-pair labels ipq", mo_cas, rdm1_cas, mo_cas, optimize="optimal") +rdm2_ao = np.einsum( + "pi,qj,rk,sl,ijkl->pqrs", + mo_cas, + mo_cas, + mo_cas, + mo_cas, + rdm2_cas, + optimize="optimal", +) + +basis_oao = get_basis(mol) +basis_inv = np.linalg.inv(basis_oao) + +rdm1 = np.einsum("ap,bq,pq->ab", basis_inv, basis_inv, rdm1_ao, optimize="optimal") +rdm2 = np.einsum( + "ap,bq,cr,ds,pqrs->abcd", + basis_inv, + basis_inv, + basis_inv, + basis_inv, + rdm2_ao, + optimize="optimal", +) + +# ----------------------------------------------------------------------- +# 3) Reference energy from OAO integrals and OAO RDMs +# ----------------------------------------------------------------------- + +h1, h2 = get_integrals(mol, basis_oao) +e_ref = np.einsum("pq,pq->", h1, rdm1) + 0.5 * np.einsum("pqrs,pqrs->", h2, rdm2) + +print(f"CAS({nelecas},{ncas})CI electronic energy: {e_casci_elec:+.10f} Ha") +print() + +ovlp = 1.0 + +# ----------------------------------------------------------------------- +# 4) Compare eigval vs ham, both without diagonal correction +# ----------------------------------------------------------------------- + +print("Truncation style: eigval (eval_thr=1e-12), no diagonal correction") +lowrank_vecs, diagonals, joint = reduce_2rdm( + rdm1, + rdm2, + ovlp, + truncation_style="eigval", + eval_thr=1e-12, + save_diag=False, + Jdiag_only=True, + mol=mol, + train_en=e_ref, +) +rdm2_rec = reconstruct_rdm2_joint(lowrank_vecs, diagonals=diagonals, joint=joint) +e_rec = np.einsum("pq,pq->", h1, rdm1) + 0.5 * np.einsum("pqrs,pqrs->", h2, rdm2_rec) +print( + f" rank = {len(lowrank_vecs[0])} / {rdm2.shape[0]**2}," + f" ||dRDM2|| = {np.linalg.norm(rdm2_rec - rdm2):.4e}," + f" |dE| = {abs(e_rec - e_ref):.4e} Ha" +) +print() + +print("Truncation style: ham (ham_thr=1e-5), no diagonal correction") +lowrank_vecs, diagonals, joint = reduce_2rdm( + rdm1, + rdm2, + ovlp, + truncation_style="ham", + ham_thr=1e-5, + save_diag=False, + Jdiag_only=True, + mol=mol, + train_en=e_ref, +) +rdm2_rec = reconstruct_rdm2_joint(lowrank_vecs, diagonals=diagonals, joint=joint) +e_rec = np.einsum("pq,pq->", h1, rdm1) + 0.5 * np.einsum("pqrs,pqrs->", h2, rdm2_rec) +print( + f" rank = {len(lowrank_vecs[0])} / {rdm2.shape[0]**2}," + f" ||dRDM2|| = {np.linalg.norm(rdm2_rec - rdm2):.4e}," + f" |dE| = {abs(e_rec - e_ref):.4e} Ha" +) +print() diff --git a/examples/low_rank/01-single_2rdm_truncations_ccsd.py b/examples/low_rank/01-single_2rdm_truncations_ccsd.py new file mode 100644 index 0000000..7feb9f4 --- /dev/null +++ b/examples/low_rank/01-single_2rdm_truncations_ccsd.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +""" +Example: single CCSD 2RDM compression with Hamiltonian-error threshold truncation. + +This script compares three diagonal-correction choices: +1) no diagonal correction +2) diagonal J correction +3) diagonal J+K correction + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import cc, gto, scf + +from evcont.electron_integral_utils import get_basis, get_integrals +from evcont.logging_utils import logger as evcont_logger +from evcont.low_rank_utils import reduce_2rdm, reconstruct_rdm2_joint +# Keep example output focused on the comparison lines below. +evcont_logger.disabled = True + +# ----------------------------------------------------------------------- +# 1) Build molecule and run RHF + CCSD +# ----------------------------------------------------------------------- + +mol = gto.M( + atom=""" + C 0.000000 0.000000 0.000000 + H 0.000000 0.000000 1.089000 + H 1.026719 0.000000 -0.363000 + H -0.513360 -0.889165 -0.363000 + H -0.513360 0.889165 -0.363000 + """, + basis="cc-pvdz", + unit="Angstrom", + symmetry=False, + verbose=0, +) + +mf = scf.RHF(mol) +mf.kernel() + +mycc = cc.CCSD(mf) +mycc.kernel() + +e_ccsd_elec = mycc.e_tot - mol.energy_nuc() + +# ----------------------------------------------------------------------- +# 2) Build AO RDMs, then transform to OAO basis +# +# evcont works in the OAO basis (get_basis / get_integrals use it), +# so we transform the CCSD AO density matrices accordingly. +# ----------------------------------------------------------------------- + +rdm1_ao = mycc.make_rdm1(ao_repr=True) +rdm2_ao = mycc.make_rdm2(ao_repr=True) + +# OAO transformation matrix S^{-1/2}; its inverse maps AO -> OAO indices +basis_oao = get_basis(mol) +basis_inv = np.linalg.inv(basis_oao) + +rdm1 = np.einsum( + "ap,bq,pq->ab", basis_inv, basis_inv, rdm1_ao, optimize="optimal" +) +rdm2 = np.einsum( + "ap,bq,cr,ds,pqrs->abcd", + basis_inv, basis_inv, basis_inv, basis_inv, + rdm2_ao, + optimize="optimal", +) + +# ----------------------------------------------------------------------- +# 3) Reference energy from OAO integrals and OAO RDMs +# ----------------------------------------------------------------------- + +h1, h2 = get_integrals(mol, basis_oao) + +# E = sum_pq h1_pq rdm1_pq + 1/2 sum_pqrs h2_pqrs rdm2_pqrs +e_ref = ( + np.einsum("pq,pq->", h1, rdm1) + + 0.5 * np.einsum("pqrs,pqrs->", h2, rdm2) +) + +print(f"CCSD electronic energy: {e_ref:+.10f} Ha") +print() + +# For a single 2RDM the self-overlap is 1. +ovlp = 1.0 + +# ----------------------------------------------------------------------- +# 4) HAM threshold truncation with three diagonal-correction modes +# ----------------------------------------------------------------------- + +print("Truncation style: ham (ham_thr=1e-3), no diagonal correction") +lowrank_vecs, diagonals, joint = reduce_2rdm( + rdm1, rdm2, ovlp, + truncation_style='ham', + ham_thr=1e-3, + save_diag=False, Jdiag_only=True, + mol=mol, train_en=e_ref, +) +rdm2_rec = reconstruct_rdm2_joint( + lowrank_vecs, diagonals=diagonals, joint=joint +) +e_rec = ( + np.einsum('pq,pq->', h1, rdm1) + + 0.5 * np.einsum('pqrs,pqrs->', h2, rdm2_rec) +) +print( + f" rank = {len(lowrank_vecs[0])} / {rdm2.shape[0]**2}," + f" ||dRDM2|| = {np.linalg.norm(rdm2_rec - rdm2):.4e}," + f" |dE| = {abs(e_rec - e_ref):.4e} Ha" +) +print() + +print("Truncation style: ham (ham_thr=1e-3), diagonal J correction") +lowrank_vecs, diagonals, joint = reduce_2rdm( + rdm1, rdm2, ovlp, + truncation_style='ham', + ham_thr=1e-3, + save_diag=True, Jdiag_only=True, + mol=mol, train_en=e_ref, +) +rdm2_rec = reconstruct_rdm2_joint( + lowrank_vecs, diagonals=diagonals, joint=joint +) +e_rec = ( + np.einsum('pq,pq->', h1, rdm1) + + 0.5 * np.einsum('pqrs,pqrs->', h2, rdm2_rec) +) +print( + f" rank = {len(lowrank_vecs[0])} / {rdm2.shape[0]**2}," + f" ||dRDM2|| = {np.linalg.norm(rdm2_rec - rdm2):.4e}," + f" |dE| = {abs(e_rec - e_ref):.4e} Ha" +) +print() + +print("Truncation style: ham (ham_thr=1e-3), diagonal J+K correction") +lowrank_vecs, diagonals, joint = reduce_2rdm( + rdm1, rdm2, ovlp, + truncation_style='ham', + ham_thr=1e-3, + save_diag=True, Jdiag_only=False, + mol=mol, train_en=e_ref, +) +rdm2_rec = reconstruct_rdm2_joint( + lowrank_vecs, diagonals=diagonals, joint=joint +) +e_rec = ( + np.einsum('pq,pq->', h1, rdm1) + + 0.5 * np.einsum('pqrs,pqrs->', h2, rdm2_rec) +) +print( + f" rank = {len(lowrank_vecs[0])} / {rdm2.shape[0]**2}," + f" ||dRDM2|| = {np.linalg.norm(rdm2_rec - rdm2):.4e}," + f" |dE| = {abs(e_rec - e_ref):.4e} Ha" +) +print() + + diff --git a/examples/low_rank/02-cas_cont_lowrank.py b/examples/low_rank/02-cas_cont_lowrank.py new file mode 100644 index 0000000..980367d --- /dev/null +++ b/examples/low_rank/02-cas_cont_lowrank.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +""" +Minimal working example: low-rank continuation of CAS states. +- eigenvalue truncation of joint decomposition + +This script demonstrates the smallest end-to-end workflow: +1) Build CAS training data at a few H-chain geometries. +2) Build both full and low-rank continuation models. +3) Predict state energies at a test geometry. +4) Compare against a direct CASCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, scf, mcscf + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import ( + approximate_multistate_OAO, + approximate_multistate_lowrank_OAO, +) + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def casci_reference_energies(mol, ncas, neleca, nroots): + """Direct CASCI reference energies (total energies, including E_nuc).""" + mf = scf.RHF(mol) + mf.kernel() + if not mf.converged: + raise RuntimeError("RHF did not converge for reference calculation.") + + mc = mcscf.CASCI(mf, ncas, neleca) + mc.fcisolver.nroots = nroots + mc.kernel() + + return np.array(mc.e_tot, dtype=float) + + +# Problem setup kept intentionally small so this runs quickly. +natom = 4 +nroots = 2 +ncas = 4 +neleca = 2 +basis = "6-31g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +# Simple low-rank setting for demonstration. +lowrank_kwargs = { + "truncation_style": "eigval", + "eval_thr": 1e-12, + "save_diag": False, +} + +# Low-rank and full models for side-by-side comparison. +cont_lr = CAS_EVCont_obj( + ncas, + neleca, + nroots=nroots, + solver="CASCI", + lowrank=True, + **lowrank_kwargs, +) +cont_full = CAS_EVCont_obj( + ncas, + neleca, + nroots=nroots, + solver="CASCI", + lowrank=False, +) + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont_lr.append_to_rdms(mol) + cont_full.append_to_rdms(mol) + +# Vectorize low-rank representation for fast inference. +cont_lr.vectorize_lowrank(hermitian=True) + +# Low-rank representation details +nvecs = cont_lr.lowrank_vectorized['nvecs'] + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_lr, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr.one_rdm, + cont_lr.lowrank_vectorized, + cont_lr.diagonal_vectorized, + cont_lr.overlap, + nroots=nroots, + density_fit=False, + Jdiag_only=True, + sao_diag=False, +) + +e_full, _ = approximate_multistate_OAO( + test_mol, + cont_full.one_rdm, + cont_full.two_rdm, + cont_full.overlap, + nroots=nroots, +) + +e_ref = casci_reference_energies(test_mol, ncas=ncas, neleca=neleca, nroots=nroots) + +print("=" * 72) +print("Minimal CAS Low-Rank Continuation Example") +print("=" * 72) +print(f"System: H{natom}, basis={basis}, CAS({ncas}, {neleca}), nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 72) +print(f"Low-rank representation:") +print(f" truncation style: {lowrank_kwargs['truncation_style']}") +print(f" threshold: {lowrank_kwargs['eval_thr']:.1e}") +print(f" {nvecs.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("-" * 72) +print("state CASCI ref (Ha) Full EVCont (Ha) Low-rank EVCont (Ha)") +print("-" * 72) + +for i in range(nroots): + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_full[i]:>16.8f} {e_lr[i]:>19.8f}") + +print("-" * 72) +print("Absolute errors vs CASCI (mHa):") +for i in range(nroots): + err_full_mha = 1000.0 * abs(e_full[i] - e_ref[i]) + err_lr_mha = 1000.0 * abs(e_lr[i] - e_ref[i]) + print(f"state {i}: full={err_full_mha:8.3f} mHa, low-rank={err_lr_mha:8.3f} mHa") +print("=" * 72) + diff --git a/examples/low_rank/03-fci_cont_lowrank_hamthr.py b/examples/low_rank/03-fci_cont_lowrank_hamthr.py new file mode 100644 index 0000000..66f152c --- /dev/null +++ b/examples/low_rank/03-fci_cont_lowrank_hamthr.py @@ -0,0 +1,124 @@ +#!/usr/bin/env python3 +""" +Minimal working example: low-rank continuation of FCI states. +- Hamiltonian-error truncation of the low-rank decomposition + +Workflow: +1) Build FCI training data at a few H-chain geometries. +2) Build both full and low-rank continuation models. +3) Predict state energies at a test geometry. +4) Compare against a direct FCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, fci + +from evcont.electron_integral_utils import get_basis, get_integrals +from evcont.FCI_EVCont import FCI_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import ( + approximate_multistate_OAO, + approximate_multistate_lowrank_OAO, +) + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def fci_reference_energies(mol, nroots): + """Direct FCI reference energies (total energies, including E_nuc).""" + h1, h2 = get_integrals(mol, get_basis(mol)) + cisolver = fci.direct_spin0.FCI() + e_ref, _ = cisolver.kernel(h1, h2, mol.nao, mol.nelec, nroots=nroots) + + if nroots == 1: + e_ref = np.array([e_ref], dtype=float) + else: + e_ref = np.array(e_ref, dtype=float) + + return e_ref + mol.energy_nuc() + + +# Problem setup kept intentionally small so this runs quickly. +natom = 4 +nroots = 2 +basis = "sto-3g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +# Hamiltonian-error threshold used during low-rank truncation. +lowrank_kwargs = { + "truncation_style": "ham", + "ham_thr": 1e-3, + "save_diag": True, + "Jdiag_only": True, +} + +# Low-rank and full models for side-by-side comparison. +cont_lr = FCI_EVCont_obj(nroots=nroots, lowrank=True, **lowrank_kwargs) +cont_full = FCI_EVCont_obj(nroots=nroots, lowrank=False) + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont_lr.append_to_rdms(mol) + cont_full.append_to_rdms(mol) + +# Vectorize low-rank representation for fast inference. +cont_lr.vectorize_lowrank(hermitian=True) +nvecs = cont_lr.lowrank_vectorized["nvecs"] + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_lr, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr.one_rdm, + cont_lr.lowrank_vectorized, + cont_lr.diagonal_vectorized, + cont_lr.overlap, + nroots=nroots, + Jdiag_only=True, + sao_diag=False, +) + +e_full, _ = approximate_multistate_OAO( + test_mol, + cont_full.one_rdm, + cont_full.two_rdm, + cont_full.overlap, + nroots=nroots, +) + +e_ref = fci_reference_energies(test_mol, nroots=nroots) + +print("=" * 72) +print("Minimal FCI Low-Rank Continuation Example (Hamiltonian Threshold)") +print("=" * 72) +print(f"System: H{natom}, basis={basis}, nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 72) +print("Low-rank representation:") +print(f" truncation style: {lowrank_kwargs['truncation_style']}") +print(f" Hamiltonian threshold: {lowrank_kwargs['ham_thr']:.1e} Ha") +print(f" {nvecs.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("-" * 72) +print("state FCI ref (Ha) Full EVCont (Ha) Low-rank EVCont (Ha)") +print("-" * 72) + +for i in range(nroots): + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_full[i]:>16.8f} {e_lr[i]:>19.8f}") + +print("-" * 72) +print("Absolute errors vs FCI (mHa):") +for i in range(nroots): + err_full_mha = 1000.0 * abs(e_full[i] - e_ref[i]) + err_lr_mha = 1000.0 * abs(e_lr[i] - e_ref[i]) + print(f"state {i}: full={err_full_mha:8.3f} mHa, low-rank={err_lr_mha:8.3f} mHa") +print("=" * 72) diff --git a/examples/low_rank/04-lowrank_NACs.py b/examples/low_rank/04-lowrank_NACs.py new file mode 100644 index 0000000..bada944 --- /dev/null +++ b/examples/low_rank/04-lowrank_NACs.py @@ -0,0 +1,160 @@ +#!/usr/bin/env python3 +""" +Minimal working example: low-rank CAS continuation inference of +nonadiabatic coupling vectors and energy gradients. + +This script demonstrates an end-to-end workflow: +1) Build CAS training data at a few H-chain geometries. +2) Build both full and low-rank continuation models. +3) Predict energies/gradients/NACs at a test geometry. +4) Compare low-rank vs full continuation results. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_gradients_loewdin import ( + get_lowrank_en_with_grad_and_NAC, + get_multistate_energy_with_grad_and_NAC, +) +# Keep example output focused on the comparison lines below. +from evcont.logging_utils import logger as evcont_logger +evcont_logger.disabled = True + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def pair_labels(nroots): + """All unique state-pair labels i3d} {e_full[i]:>14.8f} {e_lr[i]:>14.8f} {de_mha:>11.4f}") +print("-" * 80) + +print("Gradient comparison (Ha/a0)") +for i in range(nroots): + dg = np.linalg.norm(g_lr[i] - g_full[i]) + print(f"state {i}: |Delta grad| = {dg:.6e}") +print("-" * 80) + +print("NAC comparison (a0^-1)") +for label in labels: + nac_full_norm = np.linalg.norm(nac_full[label]) + nac_lr_norm = np.linalg.norm(nac_lr[label]) + + # NAC vectors can differ by a global sign due to phase/gauge choices. + nac_err = min( + np.linalg.norm(nac_lr[label] - nac_full[label]), + np.linalg.norm(nac_lr[label] + nac_full[label]), + ) + + print( + f"pair {label}: |Full|={nac_full_norm:.6e}, " + f"|Low-rank|={nac_lr_norm:.6e}, " + f"|Delta|={nac_err:.6e}" + ) + +print("=" * 80) +print("Note: Difference purely due to low-rank inference using density fitting") diff --git a/examples/low_rank/05-amplitude_relaxation.py b/examples/low_rank/05-amplitude_relaxation.py new file mode 100644 index 0000000..c5cd4e3 --- /dev/null +++ b/examples/low_rank/05-amplitude_relaxation.py @@ -0,0 +1,157 @@ +#!/usr/bin/env python3 +""" +Minimal working example: +comparing low-rank continuation with and without amplitude +relaxation after eigenvalue truncation. + +Workflow: +1) Build FCI training data at a few H-chain geometries. +2) Build both full, low-rank and low-rank with amplitude relaxation continuation models. +3) Predict state energies at a test geometry. +4) Compare against a direct FCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, fci + +from evcont.electron_integral_utils import get_basis, get_integrals +from evcont.FCI_EVCont import FCI_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import ( + approximate_multistate_OAO, + approximate_multistate_lowrank_OAO, +) + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def fci_reference_energies(mol, nroots): + """Direct FCI reference energies (total energies, including E_nuc).""" + h1, h2 = get_integrals(mol, get_basis(mol)) + cisolver = fci.direct_spin0.FCI() + e_ref, _ = cisolver.kernel(h1, h2, mol.nao, mol.nelec, nroots=nroots) + + if nroots == 1: + e_ref = np.array([e_ref], dtype=float) + else: + e_ref = np.array(e_ref, dtype=float) + + return e_ref + mol.energy_nuc() + + +# Problem setup kept intentionally small so this runs quickly. +natom = 8 +nroots = 2 +basis = "sto-3g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +# Hamiltonian-error threshold used during low-rank truncation. +lowrank_kwargs = { + "truncation_style": "eigval", + "eval_thr": 1e-1, + "save_diag": True, + "Jdiag_only": True, + "relax_amp": False, + "opt_no_diag": False, +} + +# With amplitude relaxation after eigenvalue truncation. +lowrank_relax_kwargs = { + "truncation_style": "eigval", + "eval_thr": 1e-1, + "save_diag": True, + "Jdiag_only": True, + "relax_amp": True, + "opt_no_diag": True, +} + +# Low-rank and full models for side-by-side comparison. +cont_lr = FCI_EVCont_obj(nroots=nroots, lowrank=True, **lowrank_kwargs) +cont_lr_relax = FCI_EVCont_obj(nroots=nroots, lowrank=True, **lowrank_relax_kwargs) +cont_full = FCI_EVCont_obj(nroots=nroots, lowrank=False) + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont_lr.append_to_rdms(mol) + cont_lr_relax.append_to_rdms(mol) + cont_full.append_to_rdms(mol) + +# Vectorize low-rank representation for fast inference. +cont_lr.vectorize_lowrank(hermitian=True) +cont_lr_relax.vectorize_lowrank(hermitian=True) +nvecs = cont_lr.lowrank_vectorized["nvecs"] +nvecs_relax = cont_lr_relax.lowrank_vectorized["nvecs"] + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_lr, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr.one_rdm, + cont_lr.lowrank_vectorized, + cont_lr.diagonal_vectorized, + cont_lr.overlap, + nroots=nroots, + Jdiag_only=True, + sao_diag=False, +) + +e_lr_relax, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr_relax.one_rdm, + cont_lr_relax.lowrank_vectorized, + cont_lr_relax.diagonal_vectorized, + cont_lr_relax.overlap, + nroots=nroots, + Jdiag_only=True, + sao_diag=False, +) + +e_full, _ = approximate_multistate_OAO( + test_mol, + cont_full.one_rdm, + cont_full.two_rdm, + cont_full.overlap, + nroots=nroots, +) + +e_ref = fci_reference_energies(test_mol, nroots=nroots) + +print("=" * 72) +print("Minimal FCI Low-Rank Continuation Example (Eigenvalue Truncation)") +print("=" * 72) +print(f"System: H{natom}, basis={basis}, nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 72) +print("Low-rank representation:") +print(f" truncation style: {lowrank_kwargs['truncation_style']}") +print(f" Eval threshold: {lowrank_kwargs['eval_thr']:.1e}") +print(f" {nvecs.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("Low-rank representation with amplitude relaxation:") +print(f" truncation style: {lowrank_relax_kwargs['truncation_style']}") +print(f" Eval threshold: {lowrank_relax_kwargs['eval_thr']:.1e}") +print(f" {nvecs_relax.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("-" * 72) +print("state FCI ref (Ha) Full EVCont (Ha) Low-rank EVCont (Ha) Low-rank EVCont with Relax (Ha)") +print("-" * 72) + +for i in range(nroots): + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_full[i]:>16.8f} {e_lr[i]:>19.8f} {e_lr_relax[i]:>19.8f}") + +print("-" * 72) +print("Absolute errors vs FCI (mHa):") +for i in range(nroots): + err_full_mha = 1000.0 * abs(e_full[i] - e_ref[i]) + err_lr_mha = 1000.0 * abs(e_lr[i] - e_ref[i]) + err_lr_relax_mha = 1000.0 * abs(e_lr_relax[i] - e_ref[i]) + print(f"state {i}: full={err_full_mha:8.3f} mHa, low-rank={err_lr_mha:8.3f} mHa, low-rank with relax={err_lr_relax_mha:8.3f} mHa") +print("=" * 72) diff --git a/examples/low_rank/06-relax_Herror.py b/examples/low_rank/06-relax_Herror.py new file mode 100644 index 0000000..3ac3506 --- /dev/null +++ b/examples/low_rank/06-relax_Herror.py @@ -0,0 +1,156 @@ +#!/usr/bin/env python3 +""" +Minimal working example: +comparing low-rank continuation with and without amplitude relaxation +after Hamiltonian error truncation. + +Workflow: +1) Build FCI training data at a few H-chain geometries. +2) Build both full, low-rank and low-rank with amplitude relaxation continuation models. +3) Predict state energies at a test geometry. +4) Compare against a direct FCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, fci + +from evcont.electron_integral_utils import get_basis, get_integrals +from evcont.FCI_EVCont import FCI_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import ( + approximate_multistate_OAO, + approximate_multistate_lowrank_OAO, +) + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def fci_reference_energies(mol, nroots): + """Direct FCI reference energies (total energies, including E_nuc).""" + h1, h2 = get_integrals(mol, get_basis(mol)) + cisolver = fci.direct_spin0.FCI() + e_ref, _ = cisolver.kernel(h1, h2, mol.nao, mol.nelec, nroots=nroots) + + if nroots == 1: + e_ref = np.array([e_ref], dtype=float) + else: + e_ref = np.array(e_ref, dtype=float) + + return e_ref + mol.energy_nuc() + + +# Problem setup kept intentionally small so this runs quickly. +natom = 8 +nroots = 2 +basis = "sto-3g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +# Hamiltonian-error threshold used during low-rank truncation. +lowrank_kwargs = { + "truncation_style": "ham", + "ham_thr": 1e-3, + "save_diag": True, + "Jdiag_only": True, + "relax_amp": False, +} + +# With relaxation +lowrank_relax_kwargs = { + "truncation_style": "ham", + "ham_thr": 1e-3, + "save_diag": True, + "Jdiag_only": True, + "relax_amp": True, + "opt_no_diag": True, # Remove diagonals from optimization since save_diag=True (default is True if save_diag=True) +} + +# Low-rank and full models for side-by-side comparison. +cont_lr = FCI_EVCont_obj(nroots=nroots, lowrank=True, **lowrank_kwargs) +cont_lr_relax = FCI_EVCont_obj(nroots=nroots, lowrank=True, **lowrank_relax_kwargs) +cont_full = FCI_EVCont_obj(nroots=nroots, lowrank=False) + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont_lr.append_to_rdms(mol) + cont_lr_relax.append_to_rdms(mol) + cont_full.append_to_rdms(mol) + +# Vectorize low-rank representation for fast inference. +cont_lr.vectorize_lowrank(hermitian=True) +cont_lr_relax.vectorize_lowrank(hermitian=True) +nvecs = cont_lr.lowrank_vectorized["nvecs"] +nvecs_relax = cont_lr_relax.lowrank_vectorized["nvecs"] + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_lr, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr.one_rdm, + cont_lr.lowrank_vectorized, + cont_lr.diagonal_vectorized, + cont_lr.overlap, + nroots=nroots, + Jdiag_only=True, + sao_diag=False, +) + +e_lr_relax, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr_relax.one_rdm, + cont_lr_relax.lowrank_vectorized, + cont_lr_relax.diagonal_vectorized, + cont_lr_relax.overlap, + nroots=nroots, + Jdiag_only=True, + sao_diag=False, +) + +e_full, _ = approximate_multistate_OAO( + test_mol, + cont_full.one_rdm, + cont_full.two_rdm, + cont_full.overlap, + nroots=nroots, +) + +e_ref = fci_reference_energies(test_mol, nroots=nroots) + +print("=" * 72) +print("Minimal FCI Low-Rank Continuation Example (Hamiltonian Threshold)") +print("=" * 72) +print(f"System: H{natom}, basis={basis}, nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 72) +print("Low-rank representation:") +print(f" truncation style: {lowrank_kwargs['truncation_style']}") +print(f" Hamiltonian threshold: {lowrank_kwargs['ham_thr']:.1e} Ha") +print(f" {nvecs.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("Low-rank representation with amplitude relaxation:") +print(f" truncation style: {lowrank_relax_kwargs['truncation_style']}") +print(f" Hamiltonian threshold: {lowrank_relax_kwargs['ham_thr']:.1e} Ha") +print(f" {nvecs_relax.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("-" * 72) +print("state FCI ref (Ha) Full EVCont (Ha) Low-rank EVCont (Ha) Low-rank EVCont with Relax (Ha)") +print("-" * 72) + +for i in range(nroots): + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_full[i]:>16.8f} {e_lr[i]:>19.8f} {e_lr_relax[i]:>19.8f}") + +print("-" * 72) +print("Absolute errors vs FCI (mHa):") +for i in range(nroots): + err_full_mha = 1000.0 * abs(e_full[i] - e_ref[i]) + err_lr_mha = 1000.0 * abs(e_lr[i] - e_ref[i]) + err_lr_relax_mha = 1000.0 * abs(e_lr_relax[i] - e_ref[i]) + print(f"state {i}: full={err_full_mha:8.3f} mHa, low-rank={err_lr_mha:8.3f} mHa, low-rank with relax={err_lr_relax_mha:8.3f} mHa") +print("=" * 72) diff --git a/examples/low_rank/07-relax_Herror_cas.py b/examples/low_rank/07-relax_Herror_cas.py new file mode 100644 index 0000000..127deea --- /dev/null +++ b/examples/low_rank/07-relax_Herror_cas.py @@ -0,0 +1,158 @@ +#!/usr/bin/env python3 +""" +Minimal working example: +comparing low-rank continuation with and without amplitude +relaxation after Hamiltonian error truncation for CAS states. + +Truncation is based on relaxed amplitude Hamiltonian error (relax_after=False). + +Workflow: +1) Build CASCI training data at a few H-chain geometries. +2) Build both full and low-rank continuation models. +3) Predict state energies at a test geometry. +4) Compare against a direct CASCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, scf, mcscf + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import ( + approximate_multistate_OAO, + approximate_multistate_lowrank_OAO, +) + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def casci_reference_energies(mol, ncas, neleca, nroots): + """Direct CASCI reference energies (total energies, including E_nuc).""" + mf = scf.RHF(mol) + mf.kernel() + if not mf.converged: + raise RuntimeError("RHF did not converge for reference calculation.") + + mc = mcscf.CASCI(mf, ncas, neleca) + mc.fcisolver.nroots = nroots + mc.kernel() + + return np.array(mc.e_tot, dtype=float) + + +# Problem setup kept intentionally small so this runs quickly. +natom = 8 +nroots = 2 +basis = "6-31g" +ncas = 4 +neleca = 2 + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +# Hamiltonian-error threshold used during low-rank truncation. +lowrank_kwargs = { + "truncation_style": "ham", + "ham_thr": 1e-3, + "save_diag": False, + "relax_amp": False, +} + +lowrank_relax_kwargs = { + "truncation_style": "ham", + "ham_thr": 1e-3, + "save_diag": False, + "relax_amp": True, + "relax_after": False, # Relax after selecting low-rank vectors based on Hamiltonian error (vs relaxing during selection) +} + +# Low-rank and full models for side-by-side comparison. +cont_lr = CAS_EVCont_obj(ncas, neleca, nroots=nroots, solver="CASCI", lowrank=True, **lowrank_kwargs) +cont_lr_relax = CAS_EVCont_obj(ncas, neleca, nroots=nroots, solver="CASCI", lowrank=True, **lowrank_relax_kwargs) +cont_full = CAS_EVCont_obj(ncas, neleca, nroots=nroots, solver="CASCI", lowrank=False) + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont_lr.append_to_rdms(mol) + cont_lr_relax.append_to_rdms(mol) + cont_full.append_to_rdms(mol) + +# Vectorize low-rank representation for fast inference. +cont_lr.vectorize_lowrank(hermitian=True) +cont_lr_relax.vectorize_lowrank(hermitian=True) +nvecs = cont_lr.lowrank_vectorized["nvecs"] +nvecs_relax = cont_lr_relax.lowrank_vectorized["nvecs"] + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_lr, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr.one_rdm, + cont_lr.lowrank_vectorized, + cont_lr.diagonal_vectorized, + cont_lr.overlap, + nroots=nroots, + density_fit=False, + Jdiag_only=True, + sao_diag=False, +) + +e_lr_relax, _ = approximate_multistate_lowrank_OAO( + test_mol, + cont_lr_relax.one_rdm, + cont_lr_relax.lowrank_vectorized, + cont_lr_relax.diagonal_vectorized, + cont_lr_relax.overlap, + nroots=nroots, + density_fit=False, + Jdiag_only=True, + sao_diag=False, +) + +e_full, _ = approximate_multistate_OAO( + test_mol, + cont_full.one_rdm, + cont_full.two_rdm, + cont_full.overlap, + nroots=nroots, +) + +e_ref = casci_reference_energies(test_mol, ncas=ncas, neleca=neleca, nroots=nroots) + +print("=" * 72) +print("Minimal CASCI Low-Rank Continuation Example (Hamiltonian Threshold)") +print("=" * 72) +print(f"System: H{natom}, basis={basis}, CAS({ncas}, {neleca}), nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 72) +print("Low-rank representation:") +print(f" truncation style: {lowrank_kwargs['truncation_style']}") +print(f" Hamiltonian threshold: {lowrank_kwargs['ham_thr']:.1e} Ha") +print(f" {nvecs.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("Low-rank representation with amplitude relaxation:") +print(f" truncation style: {lowrank_relax_kwargs['truncation_style']}") +print(f" Hamiltonian threshold: {lowrank_relax_kwargs['ham_thr']:.1e} Ha") +print(f" {nvecs_relax.mean():.1f} vectors per 2RDM (out of max {test_mol.nao**2})") +print("-" * 72) +print("state CASCI ref (Ha) Full EVCont (Ha) Low-rank EVCont (Ha) Low-rank EVCont with Relax (Ha)") +print("-" * 72) + +for i in range(nroots): + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_full[i]:>16.8f} {e_lr[i]:>19.8f} {e_lr_relax[i]:>19.8f}") + +print("-" * 72) +print("Absolute errors vs CASCI (mHa):") +for i in range(nroots): + err_full_mha = 1000.0 * abs(e_full[i] - e_ref[i]) + err_lr_mha = 1000.0 * abs(e_lr[i] - e_ref[i]) + err_lr_relax_mha = 1000.0 * abs(e_lr_relax[i] - e_ref[i]) + print(f"state {i}: full={err_full_mha:8.3f} mHa, low-rank={err_lr_mha:8.3f} mHa, low-rank with relax={err_lr_relax_mha:8.3f} mHa") +print("=" * 72) diff --git a/examples/low_rank/08-SVD_lowrank.py b/examples/low_rank/08-SVD_lowrank.py new file mode 100644 index 0000000..4cc5fe2 --- /dev/null +++ b/examples/low_rank/08-SVD_lowrank.py @@ -0,0 +1,165 @@ +#!/usr/bin/env python3 +""" +Minimal working example: +low-rank continuation with SVD compression in Coulomb grouping +in contrast to joint decomposition shown in earlier examples. +(same script as 04-lowrank_NACs.py but with SVD-based low-rank representation) + +This script demonstrates an end-to-end workflow: +1) Build CAS training data at a few H-chain geometries. +2) Build both full and low-rank continuation models. +3) Predict energies/gradients/NACs at a test geometry. +4) Compare low-rank vs full continuation results. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_gradients_loewdin import ( + get_lowrank_en_with_grad_and_NAC, + get_multistate_energy_with_grad_and_NAC, +) +# Keep example output focused on the comparison lines below. +from evcont.logging_utils import logger as evcont_logger +evcont_logger.disabled = True + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def pair_labels(nroots): + """All unique state-pair labels i3d} {e_full[i]:>14.8f} {e_lr[i]:>14.8f} {de_mha:>11.4f}") +print("-" * 80) + +print("Gradient comparison (Ha/a0)") +for i in range(nroots): + dg = np.linalg.norm(g_lr[i] - g_full[i]) + print(f"state {i}: |Delta grad| = {dg:.6e}") +print("-" * 80) + +print("NAC comparison (a0^-1)") +for label in labels: + nac_full_norm = np.linalg.norm(nac_full[label]) + nac_lr_norm = np.linalg.norm(nac_lr[label]) + + # NAC vectors can differ by a global sign due to phase/gauge choices. + nac_err = min( + np.linalg.norm(nac_lr[label] - nac_full[label]), + np.linalg.norm(nac_lr[label] + nac_full[label]), + ) + + print( + f"pair {label}: |Full|={nac_full_norm:.6e}, " + f"|Low-rank|={nac_lr_norm:.6e}, " + f"|Delta|={nac_err:.6e}" + ) + +print("=" * 80) diff --git a/examples/low_rank/09-NAC_with_diag.py b/examples/low_rank/09-NAC_with_diag.py new file mode 100644 index 0000000..ad9777d --- /dev/null +++ b/examples/low_rank/09-NAC_with_diag.py @@ -0,0 +1,164 @@ +#!/usr/bin/env python3 +""" +Minimal working example: low-rank continuation inference for NACs +with diagonal corrections. + +This script demonstrates an end-to-end workflow: +1) Build CAS training data at a few H-chain geometries. +2) Build both full and low-rank continuation models. +3) Predict energies/gradients/NACs at a test geometry. +4) Compare low-rank vs full continuation results. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_gradients_loewdin import ( + get_lowrank_en_with_grad_and_NAC, + get_multistate_energy_with_grad_and_NAC, +) +# Keep example output focused on the comparison lines below. +from evcont.logging_utils import logger as evcont_logger +evcont_logger.disabled = True + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def pair_labels(nroots): + """All unique state-pair labels i3d} {e_full[i]:>14.8f} {e_lr[i]:>14.8f} {de_mha:>11.4f}") +print("-" * 80) + +print("Gradient comparison (Ha/a0)") +for i in range(nroots): + dg = np.linalg.norm(g_lr[i] - g_full[i]) + print(f"state {i}: |Delta grad| = {dg:.6e}") +print("-" * 80) + +print("NAC comparison (a0^-1)") +for label in labels: + nac_full_norm = np.linalg.norm(nac_full[label]) + nac_lr_norm = np.linalg.norm(nac_lr[label]) + + # NAC vectors can differ by a global sign due to phase/gauge choices. + nac_err = min( + np.linalg.norm(nac_lr[label] - nac_full[label]), + np.linalg.norm(nac_lr[label] + nac_full[label]), + ) + + print( + f"pair {label}: |Full|={nac_full_norm:.6e}, " + f"|Low-rank|={nac_lr_norm:.6e}, " + f"|Delta|={nac_err:.6e}" + ) + +print("=" * 80) diff --git a/examples/solvers/00-fci_continuation.py b/examples/solvers/00-fci_continuation.py new file mode 100644 index 0000000..2e0465d --- /dev/null +++ b/examples/solvers/00-fci_continuation.py @@ -0,0 +1,83 @@ +#!/usr/bin/env python3 +""" +Minimal working example: FCI eigenvector continuation. + +Workflow: +1) Build FCI training data at a few H-chain geometries. +2) Predict state energies at a test geometry via EVCont. +3) Compare against a direct FCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, fci + +from evcont.electron_integral_utils import get_basis, get_integrals +from evcont.FCI_EVCont import FCI_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import approximate_multistate_OAO + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def fci_reference_energies(mol, nroots): + """Direct FCI reference energies (total energies, including E_nuc).""" + h1, h2 = get_integrals(mol, get_basis(mol)) + cisolver = fci.direct_spin0.FCI() + e_ref, _ = cisolver.kernel(h1, h2, mol.nao, mol.nelec, nroots=nroots) + + if nroots == 1: + e_ref = np.array([e_ref], dtype=float) + else: + e_ref = np.array(e_ref, dtype=float) + + return e_ref + mol.energy_nuc() + + +# Problem setup kept intentionally small so this runs quickly. +natom = 4 +nroots = 2 +basis = "sto-3g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +cont = FCI_EVCont_obj(nroots=nroots) + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont.append_to_rdms(mol) + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_cont, _ = approximate_multistate_OAO( + test_mol, + cont.one_rdm, + cont.two_rdm, + cont.overlap, + nroots=nroots, +) + +e_ref = fci_reference_energies(test_mol, nroots=nroots) + +print("=" * 60) +print("Minimal FCI Eigenvector Continuation Example") +print("=" * 60) +print(f"System: H{natom}, basis={basis}, nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 60) +print("state FCI ref (Ha) EVCont (Ha) Error (mHa)") +print("-" * 60) + +for i in range(nroots): + err_mha = 1000.0 * abs(e_cont[i] - e_ref[i]) + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_cont[i]:>14.8f} {err_mha:>10.3f}") + +print("=" * 60) diff --git a/examples/solvers/01-casci_continuation.py b/examples/solvers/01-casci_continuation.py new file mode 100644 index 0000000..b53919a --- /dev/null +++ b/examples/solvers/01-casci_continuation.py @@ -0,0 +1,87 @@ +#!/usr/bin/env python3 +""" +Minimal working example: CAS eigenvector continuation. + +Workflow: +1) Build CASCI training data at a few H-chain geometries. +2) Predict state energies at a test geometry via EVCont. +3) Compare against a direct CASCI reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, scf, mcscf + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import approximate_multistate_OAO + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def casci_reference_energies(mol, ncas, neleca, nroots): + """Direct CASCI reference energies (total energies, including E_nuc).""" + mf = scf.RHF(mol) + mf.kernel() + if not mf.converged: + raise RuntimeError("RHF did not converge for reference calculation.") + + mc = mcscf.CASCI(mf, ncas, neleca) + mc.fcisolver.nroots = nroots + mc.kernel() + + if nroots == 1: + return np.array([mc.e_tot], dtype=float) + return np.array(mc.e_tot, dtype=float) + + +# Problem setup kept intentionally small so this runs quickly. +natom = 4 +nroots = 1 +ncas = 4 +neleca = 2 +basis = "6-31g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +cont = CAS_EVCont_obj(ncas, neleca, nroots=nroots, solver="CASCI") + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont.append_to_rdms(mol) + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_cont, _ = approximate_multistate_OAO( + test_mol, + cont.one_rdm, + cont.two_rdm, + cont.overlap, + nroots=nroots, +) + +e_ref = casci_reference_energies(test_mol, ncas=ncas, neleca=neleca, nroots=nroots) + +print("=" * 60) +print("Minimal CASCI Eigenvector Continuation Example") +print("=" * 60) +print(f"System: H{natom}, basis={basis}, CAS({ncas}, {neleca}), nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 60) +print("state CASCI ref (Ha) EVCont (Ha) Error (mHa)") +print("-" * 60) + +for i in range(nroots): + err_mha = 1000.0 * abs(e_cont[i] - e_ref[i]) + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_cont[i]:>14.8f} {err_mha:>10.3f}") + +print("=" * 60) + diff --git a/examples/solvers/02-casscf_continuation.py b/examples/solvers/02-casscf_continuation.py new file mode 100644 index 0000000..b3d0133 --- /dev/null +++ b/examples/solvers/02-casscf_continuation.py @@ -0,0 +1,91 @@ +#!/usr/bin/env python3 +""" +Minimal working example: CAS eigenvector continuation. + +Workflow: +1) Build CASSCF training data at a few H-chain geometries. +2) Predict state energies at a test geometry via EVCont. +3) Compare against a direct CASSCF reference. + +Author: Kemal Atalar +""" + +import numpy as np +from pyscf import gto, scf, mcscf + +from evcont.CASCI_EVCont import CAS_EVCont_obj +from evcont.ab_initio_eigenvector_continuation import approximate_multistate_OAO + + +def build_h_chain(natom, spacing_bohr, basis="sto-3g"): + """Create a linear hydrogen chain with fixed spacing in Bohr.""" + atom = [("H", (i * spacing_bohr, 0.0, 0.0)) for i in range(natom)] + return gto.M(atom=atom, basis=basis, unit="Bohr", symmetry=False, verbose=0) + + +def casscf_reference_energies(mol, ncas, neleca, nroots): + """Direct CASSCF reference energies (total energies, including E_nuc).""" + mf = scf.RHF(mol) + mf.kernel() + if not mf.converged: + raise RuntimeError("RHF did not converge for reference calculation.") + + mc = mcscf.CASSCF(mf, ncas, neleca) + if nroots > 1: + # Optimize a common orbital set for all roots. + weights = [1.0 / nroots] * nroots + mc = mc.state_average_(weights) + mc.kernel() + + if nroots == 1: + return np.array([mc.e_tot], dtype=float) + + return np.array(mc.e_states, dtype=float) + + +# Problem setup kept intentionally small so this runs quickly. +natom = 4 +nroots = 2 +ncas = 4 +neleca = 2 +basis = "6-31g" + +train_spacings = [1.2, 1.8] +test_spacing = 1.5 + +cont = CAS_EVCont_obj(ncas, neleca, nroots=nroots, solver="sa-casscf") + +# Build training set. +for spacing in train_spacings: + mol = build_h_chain(natom=natom, spacing_bohr=spacing, basis=basis) + cont.append_to_rdms(mol) + +# Predict at test geometry. +test_mol = build_h_chain(natom=natom, spacing_bohr=test_spacing, basis=basis) + +e_cont, _ = approximate_multistate_OAO( + test_mol, + cont.one_rdm, + cont.two_rdm, + cont.overlap, + nroots=nroots, +) + +e_ref = casscf_reference_energies(test_mol, ncas=ncas, neleca=neleca, nroots=nroots) + +print("=" * 60) +print("Minimal CASSCF Eigenvector Continuation Example") +print("=" * 60) +print(f"System: H{natom}, basis={basis}, CAS({ncas}, {neleca}), nroots={nroots}") +print(f"Training spacings (Bohr): {train_spacings}") +print(f"Test spacing (Bohr): {test_spacing}") +print("-" * 60) +print("state CASSCF ref (Ha) EVCont (Ha) Error (mHa)") +print("-" * 60) + +for i in range(nroots): + err_mha = 1000.0 * abs(e_cont[i] - e_ref[i]) + print(f"{i:>3d} {e_ref[i]:>16.8f} {e_cont[i]:>14.8f} {err_mha:>10.3f}") + +print("=" * 60) + diff --git a/nx-interface/run-evcont-driver.py b/nx-interface/run-evcont-driver.py new file mode 100644 index 0000000..baaf3ac --- /dev/null +++ b/nx-interface/run-evcont-driver.py @@ -0,0 +1,583 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Fri Oct 27 16:04:25 2023 + +Driver to interface Newton-X CS 2.4 for non-adiabatic molecular dynamics +simulations with eigenvector continuation + +For the interface, we need: + - Reading geometry + - Computing and writing necessary quantities: + + Dynamics: Multistate energies, gradients and nonadiabatic coupling vectors (NAC) + + InitCond: Multistate energies and oscillator strengths + +@author: Kemal Atalar +""" + +import numpy as np + +import os +import sys +import pickle +from pathlib import Path + +############################ +# Checks for evcont and pyscf +try: + from evcont.ab_initio_gradients_loewdin import get_multistate_energy_with_grad_and_NAC, get_lowrank_en_with_grad_and_NAC + from evcont.FCI_NAC import get_FCI_energy_with_grad_and_NAC, get_FCI_energy_with_grad_and_NAC_withsym + from evcont.NAMD_utils import read_model +except: + print('Error in run-evcont-driver: evcont is not installed!') + sys.exit() + +try: + from pyscf import gto, fci, scf, mcscf, lib, grad, df + +except: + print('Error in run-evcont-driver: pyscf is not installed!') + sys.exit() + + +# Get parameters from nx-interface +NSTAT = int(sys.argv[1]) +NSTATDYN = int(sys.argv[2]) + +############################ + +def read_mol(basis, mol_sym): + """ + Read the current geometry from the trajectory and build the molecule object + """ + # Assumes geom file is in the current directory + geom_f = 'geom' + + atom_f = [] + with open(geom_f,'r') as f: + for line in f.readlines(): + splt = line.split() + #sym, atomic no, xc, yc, zc, mass + atom_f.append((splt[0], np.array(splt[2:5],dtype=np.float64))) + #atom_f.append((splt[0], [float(i) for i in splt[2:5]])) + + #print(atom_f) + + # Create the molecule + mol = gto.Mole() + + mol.build( + atom=atom_f, + basis=basis, + symmetry=mol_sym, + unit="Bohr", + verbose=0 + ) + + return mol + +############################ +""" +def read_input_file(filename='evcont.in'): + + # Default input parameters + defaults = { + 'basis': 'sto-6g', + 'trdm_path': None, + 'fix_singlet' : False, + 'fix_sym' : None, + 'lowrank' : False, + 'density_fit' : False, + 'df_basis' : None, + 'use_pyscf': False, + 'use_quantel' : False, + 'pyscf_solver': None, + } + + variables = defaults.copy() + input_path = os.path.join(os.getcwd(), filename) + + if not os.path.exists(input_path): + print(f"Warning: '{filename}' not found in the current directory. Using all default values.") + return variables + + with open(input_path, 'r') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#'): + continue + if '=' in line: + key, value = line.split('=', 1) + key = key.strip() + value = value.strip() + + # Attempt type conversion based on defaults + if key in defaults: + expected_type = type(defaults[key]) + try: + if expected_type == bool: + value = value.lower() == 'true' + else: + value = expected_type(value) + except ValueError: + print(f"Warning: Could not convert '{key}' to {expected_type.__name__}, using default.") + continue + variables[key] = value + + return variables +""" + +def read_input_file(filename='evcont.in'): + """ + Reads key=value pairs from an input file and fills in defaults. + + Args: + filename (str): Path to input file. Defaults to 'evcont.in'. + defaults (dict): Dictionary of default values. + required_keys (list): Keys that must be present in input or defaults. + + Returns: + dict: Dictionary of input parameters. + + Raises: + FileNotFoundError: If input file is not found. + ValueError: If required keys are missing. + """ + # Default input parameters + defaults = { + 'basis': 'sto-6g', + 'use_pyscf': False, + 'trdm_path': None, + 'fix_singlet' : False, + 'fix_sym' : None, + 'lowrank' : False, + 'density_fit' : False, + 'df_basis' : None, + 'pyscf_solver' : None, + 'use_quantel' : False, + 'ncas' : None, + 'nelec' : None, + 'jdiag_only' : True + } + + required_keys=[] + + if not os.path.isfile(filename): + raise FileNotFoundError(f"Input file '{filename}' not found.") + + user_inputs = {} + with open(filename, 'r') as f: + for line in f: + line = line.strip() + if not line or line.startswith('#'): + continue + if '=' not in line: + continue + key, value = map(str.strip, line.split('=', 1)) + value_lower = value.lower() + + # If value is "none", interpret as Python None + if value_lower == 'none': + converted_value = None + elif key in defaults: + default_value = defaults[key] + if default_value is None: + converted_value = value + else: + expected_type = type(default_value) + try: + if expected_type == bool: + converted_value = value_lower == 'true' + else: + converted_value = expected_type(value) + except ValueError: + print(f"Warning: Could not convert '{key}' to {expected_type.__name__}, using raw value.") + converted_value = value + else: + # No default given — store as string or None + converted_value = None if value_lower == 'none' else value + + user_inputs[key] = converted_value + + # Fill in defaults for missing keys + for key, val in defaults.items(): + if key not in user_inputs: + user_inputs[key] = val + + # Enforce required keys + missing = [k for k in required_keys if k not in user_inputs or user_inputs[k] is None] + # Specialized enforcement + if user_inputs['use_pyscf'] and user_inputs['pyscf_solver'] is None: + missing.append('pyscf_solver') + + if user_inputs['pyscf_solver'] is not None and 'cas' in user_inputs['pyscf_solver']: + if user_inputs['ncas'] is None: + missing.append('ncas') + if user_inputs['nelec'] is None: + missing.append('nelec') + + if missing: + raise ValueError(f"Missing required input(s): {', '.join(missing)}") + + return user_inputs + +def run_training(path): + """ + TODO: Automate training, it needs pretraining for now + """ + pass + +def load_pickle(filename): + with open(filename, 'rb') as f: + return pickle.load(f) + +def get_phase(old,new): + + #norm_old = np.linalg.norm(old) + #norm_new = np.linalg.norm(new) + + cosq = np.einsum("ij,ij",old, new) + + if cosq >= 0: + return 1. + else: + return -1. + +def adjust_phase(natm): + + # Read old and current NACs + currentnac = np.loadtxt('nad_vectors') + try: + oldnac = np.loadtxt('oldh') + except: + oldnac = currentnac + + # Compute the overlap and adjust the phase + n_nac = int(oldnac.shape[0]/natm) + + adjusted_nacs = [] + for i in range(n_nac): + oldi= oldnac[i*natm : (i+1)*natm, :] + curri = currentnac[i*natm : (i+1)*natm, :] + + phase = get_phase(oldi,curri) + adjusted_nacs.append(phase * curri) + + # Write the adjusted NACs + np.savetxt('nad_vectors',np.vstack(adjusted_nacs)) + +def write_traj(mol): + ''' + Write positions along the trajectory to a separate file, 'traj_geom.npy' + (to retain more precision than 'dyn.out') + + ''' + fnam = 'traj_geom.npy' + + if not os.path.isfile(fnam): + # Create the first instance + np.save(fnam, [mol.atom_coords()]) + + else: + # Load + coord = np.load(fnam) + + # Add the new geometry + new_traj = np.concatenate((coord,[mol.atom_coords()])) + + # Write to file + np.save(fnam, new_traj) + +def write_cont(vec): + ''' + Write positions along the trajectory to a separate file, 'traj_geom.npy' + (to retain more precision than 'dyn.out') + + ''' + fnam = 'traj_vec.npy' + + if not os.path.isfile(fnam): + # Create the first instance + np.save(fnam, [vec]) + + else: + # Write to file + np.save(fnam, np.concatenate((np.load(fnam),[vec]))) + +def sacasscf_en_with_grad_and_nac(mol, cas, nroots=1, + fix_singlet=True, + anneal=True, + compute_grad=True, compute_nac=True, + density_fit=False): + """ + Wrapper for running pyscf CASSCF calculation for each molecular geometry + along a NAMD trajectory + + Args: + mol: pyscf Mole object + cas: tuple of (ncas, nelec) + nroots: number of roots + fix_singlet: fix spin to singlet + anneal: use orbital annealing + compute_grad: compute gradients + compute_nac: compute nonadiabatic couplings + density_fit: use density fitting for integrals + """ + # CAS + ncas, nelec = cas + + # Validate CAS parameters + if ncas > mol.nao: + raise ValueError(f"ncas={ncas} exceeds number of orbitals ({mol.nao})") + if nelec > mol.nelectron: + raise ValueError(f"nelec={nelec} exceeds number of electrons ({mol.nelectron})") + + # Setup calculation + mf = scf.RHF (mol) + if density_fit: + mf = df.density_fit(mf) + mf.run() + mc = mcscf.CASSCF (mf, ncas, nelec) + if fix_singlet: + mc.fix_spin_(ss=0, shift=1) + + mc = mc.state_average ([1/nroots for i in range(nroots)]) + mc.conv_tol = 1e-10 + + # Setup orbitals and run CASSCF calculation + orb_path = "cas_orbitals.npy" + if anneal and Path(orb_path).exists(): + mo_prev = np.load(orb_path) + mo_proj = mcscf.project_init_guess(mc, mo_prev) + mc.kernel(mo_proj) + + elif anneal: + print('File with orbitals from previous CAS iteration (cas_orbitals.npy) not found. Using mean-field starting point.') + mc.run() + + else: + mc.run() + + # Check convergence + if not mc.converged: + print(f"WARNING: SA-CASSCF calculation did not converge! (energy={mc.e_tot})") + + # Save orbitals for next iteration + if anneal: + np.save(orb_path,mc.mo_coeff) + + # Verify we have enough states + if len(mc.e_states) < nroots: + raise RuntimeError(f"CASSCF returned {len(mc.e_states)} states, but {nroots} were requested") + + # Compute energy, grad and NAC + en = mc.e_states[:nroots] + #print(en) + grad_all = [] + nac_all = {} + + # Set grad and NAC objects only if needed + if compute_grad: + mc_grads = mc.Gradients() + if compute_nac: + mc_nacs = mc.nac_method() + + for state in range(nroots): + # Gradients + if compute_grad: + grad_all.append( mc_grads.kernel(state=state)) + + # NACs + for jstate in range(state): + if compute_nac: + nac_all[str(state)+str(jstate)] = mc_nacs.kernel (state=(state,jstate)) + + return en, np.array(grad_all), nac_all + +def evcont_feed_nx(mode, adjustphase=True): + ''' + Call evcont at the geometry to extract energies, gradients and nonadiabatic + coupling vectors (can be extended to other properties) + + Modified from run-mlatom-driver.py in Newton-X MLAtom interface + + Args: + mode (int): + 0 - initcond + Only modifies oscillator strengths (Not implemented yet) + 1 - dynamics + Updates energies, gradients and NACs + ''' + + # Read the input parameters + inputs = read_input_file() + + trdm_path = inputs['trdm_path'] + use_pyscf = inputs['use_pyscf'] + use_quantel = inputs['use_quantel'] + pyscf_solver = inputs['pyscf_solver'] + fix_sym = inputs['fix_sym'] + fix_singlet = inputs['fix_singlet'] + + # Symmetry + if fix_sym == None or not use_pyscf: + mol_sym = False + else: + mol_sym = True + + # Get the mol object for continuation + mol = read_mol(inputs['basis'], mol_sym) + + # Set FCI solver if use_pyscf + if use_pyscf: + # Set fci solver to be used + + if fix_sym == None: + FCISOLVER = fci.direct_spin0.FCI() + else: + FCISOLVER = fci.direct_spin0_symm.FCI(mol) + FCISOLVER.wfnsym = fix_sym + + FCISOLVER.nroots = NSTAT+1 + + if fix_singlet: + fci.addons.fix_spin_(FCISOLVER,ss=0) # Fix spin + + # Add the current geometry to list of geometries along the trajectory + write_traj(mol) + + # Get energies, gradients, NAC + if not use_pyscf: + print('Implementation: evcont') + + # Read the intermediate state from continuation training + cwd = os.getcwd() + if trdm_path is None: + cont_ovlp, cont_1rdm, cont_2rdm, cont_diag = read_model(cwd) + else: + cont_ovlp, cont_1rdm, cont_2rdm, cont_diag = read_model(trdm_path) + + + # From eigenvector continuation + if inputs['lowrank']: + if inputs['density_fit']: + print('Low-rank inference - with density fitting (%s basis)'%inputs['df_basis']) + else: + print('Low-rank inference - w/out density fitting') + + vec_cont, en_cont, grad_cont, nac_cont, _ = get_lowrank_en_with_grad_and_NAC( + mol, + cont_1rdm, + cont_ovlp, + cont_2rdm, + cont_diag, + nroots=NSTAT, + density_fit=inputs['density_fit'], + df_basis=inputs['df_basis'], + Jdiag_only=inputs['jdiag_only'] + ) + else: + vec_cont, en_cont, grad_cont, nac_cont, _ = get_multistate_energy_with_grad_and_NAC( + mol, + cont_1rdm, cont_2rdm, cont_ovlp, + nroots=NSTAT + ) + + write_cont(vec_cont) + + else: + + if pyscf_solver in ['fci','FCI']: + print('Implementation: pyscf FCI - sym_%s'%fix_sym) + # FCI results in SAO basis + en_cont, grad_cont, nac_cont, _ = get_FCI_energy_with_grad_and_NAC_withsym( + mol, + FCISOLVER, + nroots=NSTAT+1, + irrep_name=fix_sym + ) + + elif pyscf_solver in ['sacasscf', 'SACASSCF','sa-casscf','SA-CASSCF']: + print('Implementation: pyscf CASSCF - sym_%s'%fix_sym) + + cas = (int(inputs['ncas']), int(inputs['nelec'])) + en_cont, grad_cont, nac_cont = sacasscf_en_with_grad_and_nac( + mol, cas, nroots=NSTAT, fix_singlet=fix_singlet, + density_fit=inputs['density_fit'] + ) + + # Checks - write to output (going into EVCont.out) + print('geom',mol.atom_coords()) + print() + if not (use_pyscf or use_quantel): + print('vec', vec_cont, vec_cont.shape) + print() + print('en',en_cont) + print() + print('grad', grad_cont) + print() + print('nac',nac_cont) + print() + + # Write energies and gradients + with open('epot', 'w') as fepot, open('grad.all', 'w') as fgradall, open('grad', 'w') as fgrad: + for istate in range(1,NSTAT+1): + fepot.writelines(' %.13f\n' % en_cont[istate-1]) + + for iatom in range(mol.natm): + current = grad_cont[istate-1,iatom,:] + + fgradall.writelines(' %.13f %.13f %.13f\n' % (current[0],current[1],current[2])) + if (istate == NSTATDYN): + fgrad.writelines(' %.13f %.13f %.13f\n' % (current[0],current[1],current[2])) + + # Write nonadiabatic coupling vectors + with open('nad_vectors', 'w') as fnad: + for ii in range(NSTAT): + for jj in range(ii): + nac_str = str(ii)+str(jj) + + for iatom in range(mol.natm): + current = nac_cont[nac_str][iatom,:] + + fnad.writelines(' %.13f %.13f %.13f\n' % (current[0],current[1],current[2])) + + if adjustphase: + adjust_phase(mol.natm) + + # TODO: Transition moments, Oscillator strengths, energy gaps, etc. + + return 1 + +if __name__ == '__main__': + + check = False + #mol = read_mol(basis=BASIS) + + # Dynamics only for now + evcont_feed_nx(1) + + # Try for specific cases + if check: + import os + + drc = '/Users/katalar/Code/newtonx/Analysis/H8/S2-dt01/evcont-ntrain11/TEMP' + + cwd = os.getcwd() + + os.chdir(drc) + #evcont_feed_nx(1) + + mol = read_mol('sto-3g',False) + + tmpd = os.getcwd() + cont_ovlp, cont_1rdm, cont_2rdm, _ = read_model(tmpd) + + # From eigenvector continuation + en_cont, grad_cont, nac_cont, _ = get_multistate_energy_with_grad_and_NAC( + mol, + cont_1rdm, cont_2rdm, cont_ovlp, + nroots=NSTAT+1 + ) + + os.chdir(cwd) diff --git a/pyproject.toml b/pyproject.toml index ece94bb..aeed382 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -5,9 +5,12 @@ build-backend = "hatchling.build" [project] name = "EVCont" version = "0.0.1" -authors = [{name = "Yannic Rath", email="yannic.rath@npl.co.uk"}] +authors = [{name = "Yannic Rath", email="yannic.rath@kcl.ac.uk"},{name = "Kemal Atalar", email="kemal.atalar@kcl.ac.uk"}] readme = "README.md" dependencies = ["numpy", "scipy", "pyscf"] +[tool.hatch.metadata] +allow-direct-references = true + [tool.hatch.build.targets.wheel] -packages = ["evcont"] \ No newline at end of file +packages = ["evcont"] diff --git a/scripts/NAC/H4/compare.py b/scripts/NAC/H4/compare.py new file mode 100644 index 0000000..b126231 --- /dev/null +++ b/scripts/NAC/H4/compare.py @@ -0,0 +1,346 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Tue Oct 31 15:12:36 2023 + +Benchmark FCI NACs against OpenMOLCAS CASSCF NACs with full active space + +@author: katalar +""" + +import numpy as np + +from pyscf import gto, fci + +from evcont.FCI_EVCont import FCI_EVCont_obj + +from evcont.electron_integral_utils import get_basis, get_integrals + +from evcont.ab_initio_gradients_loewdin import get_multistate_energy_with_grad_and_NAC + +from pyscf.mcscf import CASCI + +import pickle + +import matplotlib.pylab as plt + + +nstate = 2 #1st excited state +nroots_evcont = 3 +cibasis = 'OAO' + +natom = 4 + +check_spin = False +plot_extensive = True + +test_range = np.linspace(0.8, 3.0,40) +#test_range = np.linspace(0.4, 1.5,20) + +# Set fcisolver - fix spin +myci = fci.direct_spin0.FCI() +#myci.nroots = nstate+4 +fci.addons.fix_spin_(myci,ss=0) + +def get_mol(positions): + mol = gto.Mole() + + mol.build( + atom=[("H", pos) for pos in positions], + basis="sto-3g", + #basis="6-31g", + symmetry=False, + unit="Bohr", + verbose=0 + ) + + return mol + + +equilibrium_dist = 1.78596 +#equilibrium_dist = + +equilibrium_pos = np.array([(x * equilibrium_dist, 0.0, 0.0) for x in range(10)]) + +#training_stretches = np.array([0.0, 0.5, -0.5, 1.0, -1.0]) +#trainig_dists = equilibrium_dist + training_stretches + +#trainig_dists = [1.0, 1.8, 2.6] +#trainig_dists = [1.0, 1.8] + +trainig_dists = np.linspace(0.8,3.0,6) +# New angstrom ones +#trainig_dists = [0.5,0.9,1.3] + +continuation_object = FCI_EVCont_obj(nroots=nroots_evcont, + cibasis=cibasis,cisolver=myci) + +# Generate training data + prepare training models +for i, dist in enumerate(trainig_dists): + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + mol = get_mol(positions) + continuation_object.append_to_rdms(mol) + +print('Finished training') + +train_nac = [] +train_grad = [] +train_en = [] +for i, test_dist in enumerate(trainig_dists): + print(i) + positions = [(x, 0.0, 0.0) for x in test_dist * np.arange(natom)] + + mol = get_mol(positions) + h1, h2 = get_integrals(mol, get_basis(mol)) + + # Continuation + en_continuation_ms, grad_continuation, nac_continuation, _ = get_multistate_energy_with_grad_and_NAC( + mol, + continuation_object.one_rdm, + continuation_object.two_rdm, + continuation_object.overlap, + nroots=nstate+1 + ) + + #cont_en[i,:] = en_continuation_ms + train_nac += [nac_continuation] + train_grad += [grad_continuation] + train_en += [en_continuation_ms] + +# Prediction on test dataset and comparison against FCI results +fci_en = np.zeros([len(test_range),nstate+4]) +fci_nac = [] +cont_en = np.zeros([len(test_range),nstate+1]) +cont_nac = [] +cont_nac_hf = [] +fci_spin = np.zeros([len(test_range),nstate+4]) + +for i, test_dist in enumerate(test_range): + print(i) + positions = [(x, 0.0, 0.0) for x in test_dist * np.arange(natom)] + + mol = get_mol(positions) + h1, h2 = get_integrals(mol, get_basis(mol)) + + # Continuation + en_continuation_ms, _, nac_continuation, nac_cont_hfonly = get_multistate_energy_with_grad_and_NAC( + mol, + continuation_object.one_rdm, + continuation_object.two_rdm, + continuation_object.overlap, + nroots=nstate+1 + ) + + cont_en[i,:] = en_continuation_ms + cont_nac += [nac_continuation] + cont_nac_hf += [nac_cont_hfonly] + + # FCI ones + mc = CASCI(mol.RHF(), natom, natom) + mc.fcisolver = fci.direct_spin0.FCI() + mc.fcisolver.nroots = nstate+4 + fci.addons.fix_spin_(mc.fcisolver,ss=0) # Fix spin + #ci_scan_exc = mc.nuc_grad_method().as_scanner(state=nstate) + #ci_scan_0 = mc.nuc_grad_method().as_scanner(state=0) + + #en_exc_exact, grad_exc_exact = ci_scan_exc(mol) + #en_exact, grad_exact = ci_scan_0(mol) + en_exact, fcivec_pos = mc.fcisolver.kernel(h1, h2, mol.nao, mol.nelec) + en_exact += mol.energy_nuc() + + # [spin, 2S+1] + spin_exact = [mc.fcisolver.spin_square(fcivec_i, mol.nao, mol.nelec)[0] for fcivec_i in fcivec_pos] + + # Get the reference numerical FCI NACs (update this part) + #nac_all = nac_continuation + + fci_en[i,:] = en_exact + fci_spin[i,:] = spin_exact + #fci_nac += [nac_all] + +####################################################################### +if check_spin: + + fig, axes = plt.subplots(nrows=2,ncols=1,sharex=True,sharey='row', + figsize=[6,10],gridspec_kw={'hspace':0.,'wspace':0}, + height_ratios=[3,1]) + + axes[0].plot(test_range,fci_en,alpha=0.8,label=['FCI-pyscf']+[None]*(fci_en.shape[1]-1)) + + axes[1].plot(test_range,fci_spin) + + axes[0].set_ylabel('E (Ha)') + axes[1].set_ylabel(r'$S^2$') + + plt.show() + +####################################################################### +# Read NACs computed from openMOLCAS +fname = 'test_NACs.pkl' +with open(fname,'rb') as f: + test_NACs = pickle.load(f) + +# MOLCAS energies for comparison +fname = 'test_en.npy' +with open(fname,'rb') as f: + molcas_en = np.load(f) + +# Separate for geometry for plotting +fci_nac = [] +for key in test_NACs.keys(): + all_NACs = test_NACs[key] + nac_i = {} + + for keyj in all_NACs.keys(): + # 0 - CI contribution (not divided by energy difference) + # 1 - CSF contribution + # 2 - Full NACs + ci_NAC = all_NACs[keyj][2] + nac_i[keyj] = ci_NAC + fci_nac.append(nac_i) + + +# CI only part of the molcas NACs +fci_cionly_nac = [] +for key in test_NACs.keys(): + all_NACs = test_NACs[key] + nac_i = {} + + for keyj in all_NACs.keys(): + ci_NAC = all_NACs[keyj][2] - all_NACs[keyj][1] + nac_i[keyj] = ci_NAC + fci_cionly_nac.append(nac_i) + +####################################################################### +# Plot NAC comparison +import matplotlib.pylab as plt + +fci_absh = {} +fci_cionly_absh = {} +cont_absh = {} +cont_hf_absh = {} +for istate in range(nstate+1): + for jstate in range(nstate+1): + if istate != jstate: + + st_label = str(istate)+str(jstate) + fci_absh[st_label] = np.array([np.abs(fci_nac[i][st_label]).sum() for i in range(len(test_range))]) + fci_cionly_absh[st_label] = np.array([np.abs(fci_cionly_nac[i][st_label]).sum() for i in range(len(test_range))]) + cont_absh[st_label] = np.array([np.abs(cont_nac[i][st_label]).sum() for i in range(len(test_range))]) + cont_hf_absh[st_label] = np.array([np.abs(cont_nac_hf[i][st_label]).sum() for i in range(len(test_range))]) + +# Colors +clr_st = {'01':'b', '10':'b', + '02':'r','20':'r', + '03':'pink','30':'pink', + '13':'y','31':'y', + '23':'violet','32':'violet', + '12':'g','21':'g'} + +labelsize = 15 +interfont=12 + +if not plot_extensive: + # Plot + fig, axes = plt.subplots(nrows=2,ncols=2,sharex=True,sharey='row', + figsize=[10,10],gridspec_kw={'hspace':0.,'wspace':0}, + height_ratios=[1,1]) + + axes[0][0].plot(test_range,molcas_en,'k',alpha=0.8,label=['CASSCF-molcas']+[None]*(molcas_en.shape[1]-1)) + axes[0][0].plot(test_range,fci_en,'r--',alpha=0.8,label=['FCI-pyscf']+[None]*(fci_en.shape[1]-1)) + axes[0][1].plot(test_range,cont_en,'k',alpha=0.8) + axes[0][1].plot(trainig_dists, np.array(train_en),'xr') + + for key, el in fci_absh.items(): + axes[1][0].plot(test_range,fci_absh[key],label=key,c=clr_st[key]) + axes[1][1].plot(test_range,cont_absh[key],label=key,c=clr_st[key]) + + axes[1][0].legend(loc='upper right',fontsize=interfont) + axes[0][0].legend(loc='upper right',fontsize=interfont) + + axes[0][0].set_title('CASSCF',fontsize=labelsize) + axes[0][1].set_title('EVcont',fontsize=labelsize) + axes[1][0].set_ylabel(r'$||\mathbf{d}_{ij}||$ (a$_0$$^{-1}$)',fontsize=labelsize) + axes[0][0].set_ylabel(r'Energy (Hartree)',fontsize=labelsize) + + + axes[1][0].set_ylim(ymin=-0.2, ymax=min(10,axes[1][0].get_ylim()[1])) + + for axcol in axes[1:]: + ylims = axcol[1].get_ylim() + axcol[1].vlines(x=trainig_dists,ymin=ylims[0],ymax=ylims[1],ls='--',color='gray') + + plt.show() + +####################################################################### +else: + # Plot + fig, axes = plt.subplots(nrows=4,ncols=3,sharex=True, + figsize=[15,20],gridspec_kw={'hspace':0.,'wspace':0.2}, + height_ratios=[1,1,1,1]) + + #axes[0][1].plot(test_range,molcas_en,'k',alpha=0.8,label=['CASSCF-molcas']+[None]*(molcas_en.shape[1]-1)) + axes[0][0].plot(test_range,molcas_en,'k',alpha=0.8,label=['CASSCF-molcas']+[None]*(molcas_en.shape[1]-1)) + axes[0][0].plot(test_range,fci_en,'r--',alpha=0.8,label=['FCI-pyscf']+[None]*(fci_en.shape[1]-1)) + + axes[0][1].plot(test_range,cont_en,'k',alpha=0.8) + axes[0][1].plot(trainig_dists, np.array(train_en),'xr') + + axes[0][2].plot(test_range, np.abs(cont_en-molcas_en[:,:nstate+1]),'k') + + for key, el in fci_absh.items(): + # FCI + axes[1][0].plot(test_range,fci_absh[key],label=key,c=clr_st[key]) + axes[2][0].plot(test_range,fci_cionly_absh[key],label=key,c=clr_st[key]) + fci_corr = np.array(fci_absh[key])-np.array(fci_cionly_absh[key]) + axes[3][0].plot(test_range,fci_corr,label=key,c=clr_st[key]) + # EVcont + axes[1][1].plot(test_range,cont_absh[key],label=key,c=clr_st[key]) + axes[2][1].plot(test_range,cont_hf_absh[key],label=key,c=clr_st[key]) + cont_corr = np.array(cont_absh[key])-np.array(cont_hf_absh[key]) + axes[3][1].plot(test_range,cont_corr,label=key,c=clr_st[key]) + # Diff + axes[1][2].plot(test_range,np.abs(cont_absh[key]-fci_absh[key]),label=key,c=clr_st[key]) + axes[2][2].plot(test_range,np.abs(cont_hf_absh[key]-fci_cionly_absh[key]),label=key,c=clr_st[key]) + #cont_corr = np.array(cont_absh[key])-np.array(cont_hf_absh[key]) + axes[3][2].plot(test_range,np.abs(cont_corr-fci_corr),label=key,c=clr_st[key]) + + axes[1][0].legend(loc='upper right',fontsize=interfont) + axes[0][0].legend(loc='upper right',fontsize=interfont) + + axes[0][0].set_title('CASSCF',fontsize=labelsize) + #axes[0][1].set_title('CASSCF - CI only',fontsize=labelsize) + #axes[0][2].set_title('CASSCF - CSF correction',fontsize=labelsize) + axes[0][1].set_title('EVcont',fontsize=labelsize) + #axes[0][4].set_title('EVcont - HF only',fontsize=labelsize) + #axes[0][5].set_title('EVcont - basis correction',fontsize=labelsize) + axes[0][2].set_title('|EVcont - FCI|',fontsize=labelsize) + + axes[0][0].set_ylabel(r'Energy (Hartree)',fontsize=labelsize) + axes[1][0].set_ylabel(r'$||\mathbf{d}_{ij}||$ (a$_0$$^{-1}$)',fontsize=labelsize) + axes[2][0].set_ylabel(r'$||\mathbf{d}_{ij}||$ (a$_0$$^{-1}$) - CI only',fontsize=labelsize) + axes[3][0].set_ylabel(r'$||\mathbf{d}_{ij}||$ (a$_0$$^{-1}$) - correction',fontsize=labelsize) + + axes[1][0].set_ylim(ymin=-0.2, ymax=min(10,axes[1][0].get_ylim()[1])) + # Equate ylims + for i in range(0,4): + axes[i][0].set_ylim(axes[i][1].get_ylim()) + axes[i][2].set_yscale('log') + + for axcol in axes[1:]: + for ax in axcol: + ylims = ax.get_ylim() + ax.set_ylim(ylims) + ax.vlines(x=trainig_dists,ymin=ylims[0],ymax=ylims[1],ls='--',color='k') + + for axcol in axes: + for ax in axcol: + ax.yaxis.grid(color='gray', linestyle='dashed') + ax.xaxis.grid(color='gray', linestyle='dashed') + + #plt.savefig('compare-nac.png',dpi=1000,bbox_inches='tight') + plt.show() + + + + diff --git a/scripts/NAC/H4/nroot5/test_NACs.pkl b/scripts/NAC/H4/nroot5/test_NACs.pkl new file mode 100644 index 0000000..8b79364 Binary files /dev/null and b/scripts/NAC/H4/nroot5/test_NACs.pkl differ diff --git a/scripts/NAC/H4/run_NAC_molcas.py b/scripts/NAC/H4/run_NAC_molcas.py new file mode 100644 index 0000000..e22a6e7 --- /dev/null +++ b/scripts/NAC/H4/run_NAC_molcas.py @@ -0,0 +1,141 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Tue Oct 31 16:26:12 2023 + +@author: katalar +""" + +import numpy as np + +import os + +import pickle + +from itertools import product + +# Bohr +test_range = np.linspace(0.8, 3.0,40) + +nroots = 3 + +inp_f_text = \ + '&gateway\nCoord\n \ +4\nHydrogen chain coordinates in Bohr\nH 0.000000 0.000000 0.000000\n\ +H xxx 0.000000 0.000000\nH yyy 0.000000 0.000000\n\ +H zzz 0.000000 0.000000\nBasis=sto-3g\nGroup=Nosymm\n\ +&SEWARD\n&RASSCF\nnactel = 4 0 0 \ninactive = 0 \nras2 = 4\nciroot =%i %i 1\n\ +&ALASKA\n NAC = 1 2\n\n'%(nroots+1, nroots+1) + +fname = 'H4.input' + +########################################################################### +def read_nac_molcas(outf): + ''' + + ''' + # Will contain 3 lists - CI NAC, CSF NAC, and their total sum + all_NAC = [] + + lines = 0 + read_ci_derivative = False + + tmp_list = [] + with open(outf) as f: + for line in f.readlines(): + if 'derivative coupling' in line: + read_ci_derivative = True + + if read_ci_derivative: + + if lines == 2 and len(line.split()) > 2: + # Read the values + tmp_list.append([float(i) for i in line.split()[1:]]) + + if '------' in line: + lines += 1 + + if lines == 3: + lines = 0 + read_ci_derivative = False + all_NAC.append(tmp_list) + tmp_list = [] + + #print(all_NAC) + return [np.array(i) for i in all_NAC] + +def read_en_molcas(outf): + + en = [] + with open(outf) as f: + for line in f.readlines(): + + if '::' in line: + en += [float(line.split()[-1])] + + return en +########################################################################### +test_NACs = {} +test_energies = [] +for i, dist_i in enumerate(test_range): + print(i) + save_en = True + # Save the current directory + cwd = os.getcwd() + + # Create a subdirectory for the geometry + d_drc = 'd-%.4f'%dist_i + os.mkdir(d_drc) + os.chdir(d_drc) + cwd_mid = os.getcwd() + + # New input file for openMOLCAS + new_inp = inp_f_text.replace('xxx','%.6f'%dist_i).replace('yyy','%.6f'%(2*dist_i)).replace('zzz','%.6f'%(3*dist_i)) + + nac_i = {} + for i,j in product(range(nroots),range(nroots)): + if i != j: + # Create directories for different NACs + nac_drc = 'nac-%i%i'%(i,j) + os.mkdir(nac_drc) + os.chdir(nac_drc) + + # Make input file + nac_inp = new_inp.replace('NAC = 1 2','NAC = %i %i'%(i+1,j+1)) + + with open(fname, 'w+') as f: + f.write(nac_inp) + + # Run calculation + os.system('$MOLCAS/pymolcas %s > out'%fname) + + #os.system('sleep 2') + + # Read NACs + allNAC = read_nac_molcas('out') + nac_i['%i%i'%(i,j)] = allNAC + + # Read energies + if save_en: + en_i = read_en_molcas('out') + test_energies.append(en_i) + save_en = False # Same for different NACs + + os.chdir(cwd_mid) + + # Add to the dataset + test_NACs[dist_i] = nac_i + + # Return to original test directory + os.chdir(cwd) + +# Write all to file +with open('test_NACs.pkl','wb') as f: + pickle.dump(test_NACs,f) + +# Write energies +with open('test_en.npy', 'wb') as f: + np.save(f, np.array(test_energies)) + + + \ No newline at end of file diff --git a/scripts/NAC/H4/test_NACs.pkl b/scripts/NAC/H4/test_NACs.pkl new file mode 100644 index 0000000..56cc767 Binary files /dev/null and b/scripts/NAC/H4/test_NACs.pkl differ diff --git a/scripts/low_rank/hydrogen_chain.py b/scripts/low_rank/hydrogen_chain.py new file mode 100644 index 0000000..0789244 --- /dev/null +++ b/scripts/low_rank/hydrogen_chain.py @@ -0,0 +1,774 @@ +#!/usr/bin/env python3 +# -*- coding: utf-8 -*- +""" +Created on Tue May 7 15:02:20 2024 + +Script to test low rank construction of eigenvector continuation + +@author: Kemal Atalar +""" + +import numpy as np +import time +import pickle + +from pyscf import gto, fci, scf, lib, ao2mo, mcscf, df + +from evcont.FCI_EVCont import FCI_EVCont_obj +from evcont.CASCI_EVCont import CAS_EVCont_obj + +from evcont.electron_integral_utils import get_basis, get_integrals, get_loewdin_trafo, get_df_integrals + +#from evcont.ab_initio_gradients_loewdin import get_multistate_energy_with_grad_and_NAC + +from evcont.ab_initio_eigenvector_continuation import ( + approximate_multistate_lowrank_OAO, + approximate_multistate_OAO, + approximate_multistate +) + +from evcont.ab_initio_gradients_loewdin import ( + get_lowrank_en_with_grad_and_NAC , + get_multistate_energy_with_grad_and_NAC, + get_grad_elec_OAO_customERI +) + +#from evcont.FCI_NAC import get_FCI_energy_with_grad_and_NAC_withsym +#from pyscf.mcscf import CASCI +#import pickle + +import matplotlib.pylab as plt +#import matplotlib as mpl +plt.style.use('default') + +nroots_evcont = 2 +cibasis = 'canonical' +#cibasis = 'OAO' + +density_fit = True + +#df_basis = 'weigend' +#df_basis = 'cc-pvdz-jkfit' +#df_basis = 'cc-pvdz-ri' +df_basis = None + +natom = 4 + +cont_solver = 'CAS' +#cont_solver = 'FCI' + +#cassolver='SA-CASSCF' +cassolver='CASCI' +ncas, neleca = 4,2 +figsave = True +figshow = True + +#plot_extensive = False +fix_singlet = False +#withMolcas = False + +# Whether to use FCI as a comparison if a CAS solver is used +fci_done = False + +fix_sym = 'A1g' +fix_sym = None + +if fix_sym == None: + mol_sym = False +else: + mol_sym = True + +use_diag = False +Jdiag_only = True # Only use diagonal corrections that contribute as J builds +sao_diag = False # Diagonal inference in SAO basis + +lowrank_kwargs = { + 'truncation_style':'nvec', + 'nvecs':10, + 'iterative':True, + 'nit':12, + 'max_iter_time':10, + } + +#lowrank_kwargs = {'truncation_style':'eigval', 'eval_thr':1e-12} +lowrank_kwargs = {'truncation_style':'eigval', 'eval_thr':1e-3} +#lowrank_kwargs = {'truncation_style':'ham', 'ham_thr':0.001, 'save_diag':use_diag} +#lowrank_kwargs = {'truncation_style':'ham_en', 'ham_thr':0.0002} + +#lowrank_kwargs = {'truncation_style':'eigval', 'eval_thr':1e-1, 'save_diag':use_diag} +#lowrank_kwargs = {'truncation_style':'nvec', 'nvecs':3, 'save_diag':use_diag} + +vectorize = True + +# For testing, use reconstructed 2tRDM instead of full evcont +# - to see if fast inference is working as intended +compare_to_reconstruct = False + +compare_gradients = False +# If True, also compare against the numerical gradient through FD +compare_to_numerical = False + +# If true, remove other contributions (nuclear terms) +remove_nuclear_grad = False +remove_nondiag = False # only from reconstructed RDM inference +remove_Jdiag = False # Set J diagonals of the low-rank representation to zero for testing + + +#test_range = np.linspace(0.8, 3.0,40) +test_range = np.linspace(0.8, 3.0, 15) +#test_range = np.linspace(0.4, 1.5,20) + +def get_mol(positions): + mol = gto.Mole() + + mol.build( + atom=[("H", pos) for pos in positions], + #basis="sto-3g", + #basis="sto-6g", + basis="6-31g", + #basis='cc-pvdz', + symmetry=mol_sym, + unit="Bohr", + verbose=0 + ) + + return mol + +def load_pickle(filename): + with open(filename, 'rb') as f: + return pickle.load(f) + +def save_pickle(filename, data_dict): + with open(filename, 'wb') as f: + pickle.dump(data_dict, f, protocol=pickle.HIGHEST_PROTOCOL) + +def numerical_gradients(mol, energy_func, h=1e-6): + """ + Compute numerical nuclear gradients of a molecule for one or multiple states. + + Args: + mol: PySCF Mole object + energy_func: callable(mol) -> float or array-like + Function that takes a mol and returns energy (float) or list/array of energies + h: finite difference step size (Bohr) + + Returns: + grads: (nstates, natm, 3) array of gradients + """ + coords = mol.atom_coords() + ref_energies = np.atleast_1d(energy_func(mol)) + nstates = len(ref_energies) + natm = mol.natm + grads = np.zeros((nstates, natm, 3)) + + eye_atoms = np.eye(natm) + + for i in range(natm): + for a in range(3): # x, y, z + disp = np.zeros(3) + disp[a] = h + + # +h displacement + mol.set_geom_(coords + eye_atoms[i][:, None] * disp, unit='Bohr') + e_plus = np.atleast_1d(energy_func(mol)) + + # -h displacement + mol.set_geom_(coords - eye_atoms[i][:, None] * disp, unit='Bohr') + e_minus = np.atleast_1d(energy_func(mol)) + + grads[:, i, a] = (e_plus - e_minus) / (2 * h) + + # reset geometry + mol.set_geom_(coords, unit='Bohr') + + return grads if nstates > 1 else grads[0] + +mol_dummy = get_mol([(x, 0.0, 0.0) for x in test_range[0] * np.arange(natom)]) + +# Set fci solver to be used +if fix_sym is None: + myci = fci.direct_spin0.FCI() +else: + myci = fci.direct_spin0_symm.FCI(mol_dummy) + myci.wfnsym = fix_sym + +#myci = fci.direct_spin0.FCI() +#myci = fci.direct_spin1.FCISolver() + + +if fix_singlet: + fci.addons.fix_spin_(myci,ss=0) # Fix spin + +equilibrium_dist = 1.78596 + +equilibrium_pos = np.array([(x * equilibrium_dist, 0.0, 0.0) for x in range(10)]) + +trainig_dists = [0.97, 1.76]#, 2.60] +#trainig_dists = np.linspace(0.97,2.60,5) + +if cont_solver == 'FCI': + continuation_object = FCI_EVCont_obj(nroots=nroots_evcont, + cibasis=cibasis,cisolver=myci, + irrep_name=fix_sym, + lowrank=True, + **lowrank_kwargs) + + continuation_object_full = FCI_EVCont_obj(nroots=nroots_evcont, + cibasis=cibasis,cisolver=myci, + irrep_name=fix_sym) + +else: + continuation_object = CAS_EVCont_obj(ncas, neleca,nroots=nroots_evcont, + solver=cassolver, + lowrank=True, + **lowrank_kwargs) + + continuation_object_full = CAS_EVCont_obj(ncas, neleca,nroots=nroots_evcont, + solver=cassolver) + +trn_geometries = [] +# Generate training data + prepare training models +for i, dist in enumerate(trainig_dists): + positions = [(x, 0.0, 0.0) for x in dist * np.arange(natom)] + # Add the geometry to set of training points + trn_geometries.append(positions) + + # Build molecule + mol = get_mol(positions) + if cont_solver == 'CAS': + continuation_object.append_to_rdms(mol,debug=False) + else: + continuation_object.append_to_rdms(mol) + + #continuation_object.append_to_rdms_new(mol) + continuation_object_full.append_to_rdms(mol) + +# If vectorize +if vectorize: + continuation_object.vectorize_lowrank(hermitian=False) + vecs_lr = continuation_object.lowrank_vectorized + diag_lr = continuation_object.diagonal_vectorized +else: + vecs_lr = continuation_object.vecs_lowrank + diag_lr = continuation_object.diagonal_lr + +if remove_Jdiag: + diag_lr[0] = np.zeros_like(diag_lr[0]) + +# Reconstruct the two_rdm for testing +two_rdm_rec = np.zeros_like(continuation_object_full.two_rdm) + +from evcont.low_rank_utils import unstack_tril +# Only use the J coul +diagonals_rec = np.zeros_like(continuation_object.diagonal_lr) +if use_diag: + if not remove_Jdiag: + diagonals_rec[:,:,0] = unstack_tril(diag_lr[0], False) #continuation_object.diagonal_lr[:,:,0] + + if not Jdiag_only: + diagonals_rec[:,:,1:] = continuation_object.diagonal_lr[:,:,1:] + +from evcont.low_rank_utils import reconstruct_rdm2_joint + +norb = two_rdm_rec.shape[-1] +for (i,j), vi in continuation_object.vecs_lowrank.items(): + + if not remove_nondiag: + rdm2_i = reconstruct_rdm2_joint(vi[:3], diagonals=diagonals_rec[i,j],joint=vi[-1]) + else: + # Diagonal only - for testing + rdm2_i = reconstruct_rdm2_joint([np.zeros([2]),np.zeros([norb,norb,2]),np.zeros([2,norb,norb])], diagonals=diagonals_rec[i,j],joint=vi[-1]) + + two_rdm_rec[i,j] = rdm2_i + two_rdm_rec[j,i] = np.einsum('ijkl->jilk',rdm2_i.conj()) + +if compare_to_reconstruct: + two_rdm_to_comp = two_rdm_rec +else: + two_rdm_to_comp = continuation_object_full.two_rdm + +# Save +i = 'final' +np.save("overlap_{}.npy".format(i), continuation_object.overlap) +np.save("one_rdm_{}.npy".format(i), continuation_object.one_rdm) + +np.save("diagonal_lr_{}.npy".format(i), diag_lr) +np.save("lowrank_vecs_{}.npy".format(i), continuation_object.vecs_lowrank) +save_pickle('vecs_lr.pkl', vecs_lr) + +np.save('trn_geometries_{}.npy'.format(i), trn_geometries) + +# Plot the low-rank decomposition distribution for training state pairs +from matplotlib.patches import Patch + +# Expansion limit +key_clrs = ['tab:blue', 'tab:orange'] +no_vec_dic = {} +clr_dic = {} +for key, item in continuation_object.vecs_lowrank.items(): + #print(key, item[0].shape) + if item[-1]: + kclr = key_clrs[0] + else: + kclr = key_clrs[1] + if key[1] >= key[0]: + no_vec_dic[','.join([str(i) for i in key])] = item[0].shape[0] + clr_dic[','.join([str(i) for i in key])] = kclr + +fig, ax = plt.subplots(figsize=[4,7]) +ax.grid(alpha=0.5) + +# Extract keys, values, and corresponding colors +labels, values = zip(*no_vec_dic.items()) +colors = [clr_dic[label] for label in labels] + +# Plot with specified colors +ax.barh(labels, values, color=colors) +#D = {u'Label1':26, u'Label2': 17, u'Label3':30} +#ax.barh(*zip(*no_vec_dic.items())) +ax.set_xlabel('Number of vectors (max %i)'%(continuation_object.one_rdm.shape[-1]**2)) +ax.set_ylabel('(bra, ket) index') + +# Add legend +legend_elements = [ + Patch(facecolor=key_clrs[0], label='Joint ED'), + Patch(facecolor=key_clrs[1], label='Coulomb SVD') +] +ax.legend(handles=legend_elements, loc='best') + +if figsave: + plt.savefig('nvecs_H%i_%s_roots%i_%s'%(natom,cont_solver,nroots_evcont,lowrank_kwargs['truncation_style'])+'.png',bbox_inches='tight',dpi=500) + +if figshow: + plt.show() + +if use_diag: + diags = diag_lr +else: + diags = None + +lr_tot = 0.; lr_n_eval = 0 +train_lowrank_en = [] +train_en = [] +for i, test_dist in enumerate(trainig_dists): + print(i) + positions = [(x, 0.0, 0.0) for x in test_dist * np.arange(natom)] + + mol = get_mol(positions) + h1, h2 = get_integrals(mol, get_basis(mol)) + + # Continuation + start = time.time() + + en_continuation_ms, vec = approximate_multistate_lowrank_OAO( + mol, + continuation_object.one_rdm, + vecs_lr, + diags, + continuation_object.overlap, + Jdiag_only=Jdiag_only, + sao_diag=sao_diag, + nroots=nroots_evcont + ) + lr_tot += (time.time()-start); lr_n_eval += 1 + + train_lowrank_en += [en_continuation_ms] + + en_continuation_ms, vec = approximate_multistate_OAO( + mol, + continuation_object_full.one_rdm, + continuation_object_full.two_rdm, + continuation_object_full.overlap, + nroots=nroots_evcont + ) + + train_en += [en_continuation_ms] + + +############## +# Testing + +def lowrank_en(mol): + return approximate_multistate_lowrank_OAO( + mol, + continuation_object.one_rdm, + vecs_lr, + diags, + continuation_object.overlap, + Jdiag_only=Jdiag_only, + sao_diag=sao_diag, + nroots=nroots_evcont+1 + )[0] + + +# Prediction on test dataset and comparison against FCI results +nroots_to_compute = nroots_evcont + 1 + +fci_en = np.zeros([len(test_range),nroots_evcont]) +ref_en = np.zeros([len(test_range),nroots_evcont]) +hf_en = np.zeros([len(test_range)]) +cont_en = np.zeros([len(test_range),nroots_to_compute]) +cont_lowrank_en = np.zeros([len(test_range),nroots_to_compute]) + +ref_grad = np.zeros([len(test_range),nroots_evcont, mol.natm, 3]) +cont_grad = np.zeros([len(test_range),nroots_to_compute, mol.natm, 3]) +cont_lr_grad = np.zeros([len(test_range),nroots_to_compute, mol.natm, 3]) +if compare_to_numerical: + num_lr_grad = np.zeros([len(test_range),nroots_to_compute, mol.natm, 3]) + +for i, test_dist in enumerate(test_range): + print(i) + positions = [(x, 0.0, 0.0) for x in test_dist * np.arange(natom)] + + mol = get_mol(positions) + h1, h2 = get_integrals(mol, get_basis(mol,'canonical')) + + print(' low rank - start') + # Continuation + start = time.time() + if not compare_gradients: + en_continuation_ms, vec = approximate_multistate_lowrank_OAO( + mol, + continuation_object.one_rdm, + vecs_lr, + diags, + continuation_object.overlap, + nroots=nroots_to_compute, + Jdiag_only=Jdiag_only, + sao_diag=sao_diag, + df_basis=df_basis + ) + cont_lowrank_en[i,:] += en_continuation_ms + + else: + out = get_lowrank_en_with_grad_and_NAC(mol, continuation_object.one_rdm, + continuation_object.overlap, + vecs_lr, diags, + sao_diag=sao_diag, + nroots=nroots_to_compute, + density_fit=density_fit, + Jdiag_only=Jdiag_only, + df_basis=df_basis) + + cont_lowrank_en[i,:] = out[1] + + lr_tot += (time.time()-start); lr_n_eval += 1 + + print(' low rank - finish - %.1f sec'%(time.time()-start)) + + if compare_to_numerical: + grad_num = numerical_gradients(mol, lowrank_en) + num_lr_grad[i,:] = grad_num + + ## HF and FCI + mf = scf.RHF(mol).density_fit(auxbasis=df_basis) + # Note that in performant code, we don't actually need to run HF at the training points, + # just have access to the get_jk function. + ehf = mf.scf() + hf_en[i] = ehf + assert(mf.converged) + + # Do FCI for the exact energy + h1_ao = mf.get_hcore() + Lpq_ao = lib.unpack_tril(mf.with_df._cderi) + Lpq_mo = lib.einsum('pi,qj,Lpq->Lij', mf.mo_coeff, mf.mo_coeff, Lpq_ao) + df_eri = lib.einsum('Pij,Pkl->ijkl', Lpq_mo, Lpq_mo) + h1e_mo = np.einsum('ai,ab,bj->ij', mf.mo_coeff, h1_ao, mf.mo_coeff) + #print(h1e_mo, df_eri, mol.nao, mol.nelec) + + Lpq_ao, deriv_cderi = get_df_integrals(mol, auxbasis=df_basis, grad=True) + """ + # DF-ERI gradients + auxmol = df.addons.make_auxmol(mol, df_basis) + + # ints_3c is the 3-center integral tensor (ij|P), where i and j are the + # indices of AO basis and P is the auxiliary basis + ints_3c2e = df.incore.aux_e2(mol, auxmol, intor='int3c2e') + #ints_2c2e is the (P|Q) integrals + ints_2c2e = auxmol.intor('int2c2e') + vals, vecs = np.linalg.eigh(ints_2c2e) + assert(len(vals[vals < 1.e-15]) == 0) # PSD + metric = np.array(np.dot(vecs * (1 / np.sqrt(vals)) , vecs.conj().T)) + cd_array = np.einsum('PQ,ijP->Qij', metric, ints_3c2e) + + # We can get the integrals ( d/dx i, j | P) + ints_3c2e_ip1 = df.incore.aux_e2(mol, auxmol, intor='int3c2e_ip1', comp=3) + # Use the same metric as before + deriv_cderi = np.einsum('PQ,xijP -> xijQ', metric, ints_3c2e_ip1) + """ + if compare_gradients: + # To reconstruct the full 4c derivative integrals, we need to contract with the previous cderi integrals + df_grad_4c_ints = np.einsum('xijP,Pkl->xijkl', deriv_cderi, Lpq_ao) + + h2_ao_deriv = df_grad_4c_ints + h2_ao = lib.einsum('Pij,Pkl->ijkl', Lpq_ao, Lpq_ao) + + h2_ao_nondf = mol.intor("int2e") + h2_ao_deriv_nondf = mol.intor("int2e_ip1", comp=3) + + print("Max error in 4c ERI derivative:", np.max(np.abs(h2_ao_deriv - h2_ao_deriv_nondf))) + + grad_nuc = df.grad.RHF(mf).grad_nuc() + #grad_nuc = grad.RHF(scf.RHF(mol)).grad_nuc() + + #cont_lr_grad[i] = out[2] + if remove_nuclear_grad: + cont_lr_grad[i] = out[2] - grad_nuc + else: + cont_lr_grad[i] = out[2] + + # Only do FCI if number of orbitals is less than 16 + if mol.nao < 16 and (cont_solver == 'FCI' or fci_done): + e_fci, c_fci = myci.kernel(h1e_mo, df_eri, mol.nao, mol.nelec, nroots=nroots_evcont) + e_fci += mol.energy_nuc() + fci_en[i,:] = e_fci + + if compare_gradients: + # Gradients + mc = mcscf.CASCI(mf, ncas=mf.mo_coeff.shape[0], nelecas=natom) + out_mc = mc.kernel() + e_fci = out_mc[0] + + assert mc.converged + + grad_method = mc.Gradients() + + grad_ref_l = [] + for refi in range(nroots_evcont): + grad_ref_l.append(grad_method.kernel(state=refi))# - grad_method.grad_nuc()) + + ref_grad[i] = np.array(grad_ref_l) + + if remove_nuclear_grad: + grad_ref = grad_method.kernel(state=0) - grad_method.grad_nuc() + else: + grad_ref = grad_method.kernel(state=0) #- grad_method.grad_nuc() + + else: + fci_done = False + + if cont_solver != 'FCI': + # CAS reference + #mf2 = scf.RHF(mol.copy()).run() + mf2 = mf + if cassolver == 'CASCI': + mc = mcscf.CASCI(mf2, ncas, neleca) + mc.fcisolver.nroots = nroots_evcont + mc.kernel() + ref_en[i,:] = mc.e_tot + elif cassolver == 'SA-CASSCF': + mc_sa = mcscf.CASSCF(mf2, ncas, neleca).state_average_([1/nroots_evcont]*nroots_evcont) + mc_sa.kernel() + #mc = mcscf.CASCI(mf2, ncas, neleca) + #mc.casci(mc_sa.mo_coeff) + #mc.fcisolver.nroots = nroots_evcont + #e_cas = mc.kernel()[1] + ref_en[i,:] = mc_sa.e_states + elif cassolver == 'SS-CASSCF': + e_cas = [] + for istate in range(nroots_evcont): + mc_ss = mcscf.CASSCF(mf2, ncas, neleca).state_specific_(istate) + mc_ss.kernel() + mc = mcscf.CASCI(mf2, ncas, neleca).state_specific_(istate) + mc.casci(mc_ss.mo_coeff) + #mc.fcisolver.nroots = nroots_evcont + e_cas.append(mc.kernel()[0]) + ref_en[i,:] = np.array(e_cas) #+ mol.energy_nuc() + + else: + ref_en[i,:] = e_fci + + # Full continuation + print(' full') + # Find h1 and eris in SAO basis + sao_basis = get_loewdin_trafo(mol.intor("int1e_ovlp")) + h1e_sao = np.einsum('ai,ab,bj->ij', sao_basis, h1_ao, sao_basis) + #Lpq_sao = lib.einsum('pi,qj,Lpq->Lij', sao_basis, sao_basis, Lpq_ao) + Lpq_sao = ao2mo._ao2mo.nr_e2(mf.with_df._cderi, sao_basis, + (0, sao_basis.shape[1], 0, sao_basis.shape[1]),aosym="s2",mosym="s2") + Lpq_sao = lib.unpack_tril(Lpq_sao) + df_eri_sao = lib.einsum('Pij,Pkl->ijkl', Lpq_sao, Lpq_sao) + + if not compare_gradients: + en_continuation_ms, vec = approximate_multistate( + h1e_sao, + df_eri_sao, + continuation_object_full.one_rdm, + two_rdm_to_comp,#continuation_object_full.two_rdm, + continuation_object_full.overlap, + nroots=nroots_to_compute + ) + cont_en[i,:] = en_continuation_ms + mol.energy_nuc() + + else: + # Get grad + grad_cont = [] + for i_state in range(nroots_to_compute): + vec_i = vec[i_state,:] + vec_j = vec_i + + one_rdm_predicted = np.tensordot(np.outer(vec_i, vec_j), continuation_object_full.one_rdm, axes=2) + two_rdm_predicted = np.tensordot(np.outer(vec_i, vec_j), + two_rdm_to_comp,#continuation_object_full.two_rdm, + axes=2) + + grad_i = get_grad_elec_OAO_customERI(mol, h2_ao, h2_ao_deriv, + one_rdm_predicted, + two_rdm_predicted) + grad_cont.append(grad_i + grad_nuc) + + if remove_nuclear_grad: + cont_grad[i] = np.array(grad_cont) - grad_nuc + else: + cont_grad[i] = np.array(grad_cont) #- grad_nuc + + out_full = get_multistate_energy_with_grad_and_NAC(mol, + continuation_object_full.one_rdm, + two_rdm_to_comp, #continuation_object_full.two_rdm, + continuation_object_full.overlap, + nroots=nroots_to_compute, savemem=True) + + cont_en[i,:] = out_full[1] + + if cont_solver == 'CAS': + if fci_done: + print(ehf, e_fci, ref_en[i], cont_en[i], cont_lowrank_en[i]) + else: + print(ehf, ref_en[i], cont_en[i], cont_lowrank_en[i]) + + else: + print(ehf, ref_en[i,:], cont_en[i], cont_lowrank_en[i]) + if compare_gradients: + print('grad', np.linalg.norm(grad_ref-out[2][0]), + np.linalg.norm(grad_cont[0]-out[2][0]), + np.linalg.norm(out_full[2][0]-out[2][0]), + np.linalg.norm(grad_num[0]-out[2][0]), + ) + #print('nac','\n', out[4],'\n', out_full[4]) + #print(' \n', grad_ref, '\n', out[2][0],'\n', grad_cont[0] ) + #1/0 + +print('Time per low-rank (s): %.2f'%(lr_tot/lr_n_eval)) + +# PLOT +fig, [ax1,ax2,ax3] = plt.subplots(nrows=3,sharex=True,figsize=[4,7],height_ratios=[3,1.5,1.5], + gridspec_kw={'hspace':0.,'wspace':0.}) + +ax1.plot(test_range, hf_en,'orange',label='HF') +if nroots_evcont > 1: + if (cont_solver == 'FCI' or fci_done): + ax1.plot(test_range,fci_en,'k',label=['FCI']+[None]*(nroots_evcont-1)) + if cont_solver != 'FCI': + ax1.plot(test_range,ref_en,'green',label=[cont_solver]+[None]*(nroots_evcont-1)) + ax1.plot(test_range,cont_en,'b',label=['full evcont']+[None]*(cont_en.shape[-1]-1)) + ax1.plot(test_range,cont_lowrank_en,'--r',label=['low rank evcont']+[None]*(cont_lowrank_en.shape[-1]-1)) +else: + if (cont_solver == 'FCI' or fci_done): + ax1.plot(test_range,fci_en,'k',label='FCI') + if cont_solver != 'FCI': + ax1.plot(test_range,ref_en,'green',label=cont_solver) + ax1.plot(test_range,cont_en,'b',label='full evcont') + ax1.plot(test_range,cont_lowrank_en,'--r',label='low rank evcont') + + +ax1.plot(trainig_dists,train_en,'xb') +ax1.plot(trainig_dists,train_lowrank_en,'xr') +ax1.legend() + +if (cont_solver == 'FCI' or fci_done): + ax2.plot(test_range,cont_en[:,:nroots_evcont] - fci_en,'b') + ax2.plot(test_range,cont_lowrank_en[:,:nroots_evcont] - fci_en,'--r') + +if cont_solver != 'FCI': + ax3.plot(test_range,cont_en[:,:nroots_evcont] - ref_en,'b') + ax3.plot(test_range,cont_lowrank_en[:,:nroots_evcont] - ref_en,'--r') + ax3.set_ylabel(r'$E_{cont}$ - $E_{%s}$ (Ha)'%cont_solver) +else: + + ax3.plot(test_range,cont_lowrank_en - cont_en,'--r') + ax3.set_ylabel(r'$E_{cont}$ - $E_{lowrank}$ (Ha)') + +ax1.set_ylabel('Energy (Ha)') +ax2.set_ylabel(r'$E_{cont}$ - $E_{FCI}$ (Ha)') +ax3.set_xlabel('Atomic separation ($a_0$)') + +if figsave: + plt.savefig('H%i_%s_roots%i_%s'%(natom,cont_solver,nroots_evcont,lowrank_kwargs['truncation_style'])+'.png',bbox_inches='tight',dpi=500) +else: + plt.show() + + +# PLOT - with grads +fig, axes = plt.subplots(nrows=3, ncols=2, sharex=True,figsize=[8,7],height_ratios=[3,1.5,1.5], + gridspec_kw={'hspace':0.,'wspace':0.2}) + +[ax1,ax2,ax3] = axes[:,0] +[ax4,ax5,ax6] = axes[:,1] + +ax1.plot(test_range, hf_en,'orange',label='HF') +if nroots_evcont > 1: + if (cont_solver == 'FCI' or fci_done): + ax1.plot(test_range,fci_en,'k',label=['FCI']+[None]*(nroots_evcont-1)) + ax4.plot(test_range,np.linalg.norm(ref_grad,axis=(2,3)),'k',label=['FCI']+[None]*(nroots_evcont-1)) + if cont_solver != 'FCI': + ax1.plot(test_range,ref_en,'green',label=[cont_solver]+[None]*(nroots_evcont-1)) + ax1.plot(test_range,cont_en,'b',label=['full evcont']+[None]*(cont_en.shape[-1]-1)) + ax1.plot(test_range,cont_lowrank_en,'--r',label=['low rank evcont']+[None]*(cont_lowrank_en.shape[-1]-1)) + ax4.plot(test_range,np.linalg.norm(cont_grad,axis=(2,3)),'b',label=['full evcont']+[None]*(cont_en.shape[-1]-1)) + ax4.plot(test_range,np.linalg.norm(cont_lr_grad,axis=(2,3)),'--r',label=['low rank evcont']+[None]*(cont_lowrank_en.shape[-1]-1)) + if compare_to_numerical: + ax4.plot(test_range,np.linalg.norm(num_lr_grad,axis=(2,3)),'tab:orange',ls='--',lw=2,alpha=0.7,label=['numerical lowrank']+[None]*(cont_en.shape[-1]-1)) + +else: + if (cont_solver == 'FCI' or fci_done): + ax1.plot(test_range,fci_en,'k',label='FCI') + if cont_solver != 'FCI': + ax1.plot(test_range,ref_en,'green',label=cont_solver) + ax1.plot(test_range,cont_en,'b',label='full evcont') + ax1.plot(test_range,cont_lowrank_en,'--r',label='low rank evcont') + + +ax1.plot(trainig_dists,train_en,'xb') +ax1.plot(trainig_dists,train_lowrank_en,'xr') +ax1.legend() +ax4.legend() + +if (cont_solver == 'FCI' or fci_done): + ax2.plot(test_range,cont_en[:,:nroots_evcont] - fci_en,'b') + ax2.plot(test_range,cont_lowrank_en[:,:nroots_evcont] - fci_en,'--r') + + ax5.plot(test_range,np.linalg.norm(cont_grad[:,:nroots_evcont] - ref_grad,axis=(2,3)),'b') + ax5.plot(test_range,np.linalg.norm(cont_lr_grad[:,:nroots_evcont] - ref_grad,axis=(2,3)),'--r') + if compare_to_numerical: + ax5.plot(test_range,np.linalg.norm(num_lr_grad[:,:nroots_evcont] - ref_grad,axis=(2,3)),'tab:orange',ls='--',lw=2,alpha=0.7) + +if cont_solver != 'FCI': + ax3.plot(test_range,cont_en[:,:nroots_evcont] - ref_en,'b') + ax3.plot(test_range,cont_lowrank_en[:,:nroots_evcont] - ref_en,'--r') + ax3.set_ylabel(r'$E_{cont}$ - $E_{%s}$ (Ha)'%cont_solver) + +else: + + ax3.plot(test_range,cont_lowrank_en - cont_en,'--r') + ax3.set_ylabel(r'$E_{cont}$ - $E_{lowrank}$ (Ha)') + + ax6.plot(test_range,np.linalg.norm(cont_grad[:,:] - cont_lr_grad[:,:],axis=(2,3)),'--r') + ax6.plot(test_range,np.linalg.norm(cont_grad[:,:] - cont_lr_grad[:,:],axis=(2,3)),'--r') + if compare_to_numerical: + ax6.plot(test_range,np.linalg.norm(num_lr_grad[:,:] - cont_lr_grad[:,:],axis=(2,3)),'tab:orange',ls='--',lw=2,alpha=0.7) + + +ax1.set_title('Energy') +ax4.set_title('Gradient') + +ax1.set_ylabel('Energy (Ha)') +ax2.set_ylabel(r'$E_{cont}$ - $E_{FCI}$ (Ha)') +ax3.set_xlabel('Atomic separation ($a_0$)') + +if figsave: + plt.savefig('H%i_%s_roots%i_%s'%(natom,cont_solver,nroots_evcont,lowrank_kwargs['truncation_style'])+'.png',bbox_inches='tight',dpi=500) +else: + plt.show() + + +