diff --git a/examples/packages/pyscf/test_casscf.py b/examples/packages/pyscf/test_casscf.py index df687ab71..cab8ab928 100644 --- a/examples/packages/pyscf/test_casscf.py +++ b/examples/packages/pyscf/test_casscf.py @@ -29,43 +29,51 @@ # else: # raise RuntimeError("Could not locate src/ directory on path for libra_py import") -from pyscf import gto from libra_py.packages.pyscf.implementations.casscf import CASSCF from libra_py.packages.pyscf.interfaces import ElectronicStructureStrategy, MolecularGeometry import numpy as np -geom1 = MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0,0.0,0.0],[0.0,0.0,0.7746]])) -geom2 = MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0,0.0,0.0],[0.0,0.0,1.0]])) +NTRAJ = 2 +NSTATES = 3 +GRAD_ROOT = 2 +geom_step0 = [ + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.7746]])), + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.9000]])), +] +geom_step1 = [ + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0000]])), + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.1000]])), +] # instantiate the CASSCF strategy class with parameters for the test -casscf = CASSCF(norbcas=2, nelecas=2, nroots=3, basis='sto-3g', charge=1) +casscf = CASSCF(norbcas=2, nelecas=2, nroots=NSTATES, basis='sto-3g', charge=1, ntraj=NTRAJ) -# wrap geometry set and HF in set_geom_and_run_hf method -casscf.set_geom_and_run_hf(geom1) +for traj_id, geom in enumerate(geom_step0): + # First geometry for each trajectory initializes that trajectory's state slot. + casscf.set_geom_and_run_hf(geom, traj_id=traj_id) -# 2) energy for root 0,1,2 -energies1 = [casscf.compute_energy(root) for root in range(3)] -print('Energies at geom1', energies1) + energies = [casscf.compute_energy(root, traj_id=traj_id) for root in range(NSTATES)] + print(f'Trajectory {traj_id} energies at step 0', energies) -# 3) gradient for root 2 -grad1 = casscf.compute_gradient(2) -print('Gradient root 2 geom1', grad1) + grad = casscf.compute_gradient(GRAD_ROOT, traj_id=traj_id) + print(f'Trajectory {traj_id} gradient root {GRAD_ROOT} at step 0', grad) -# 4) set mol2 geometry and run HF -casscf.set_geom_and_run_hf(geom2) +for traj_id, geom in enumerate(geom_step1): + # Second geometry reuses only this trajectory's previous CASSCF/HF state. + casscf.set_geom_and_run_hf(geom, traj_id=traj_id) -# 5) energy for root 0,1,2 -energies2 = [casscf.compute_energy(root) for root in range(3)] -print('Energies at geom2', energies2) + energies = [casscf.compute_energy(root, traj_id=traj_id) for root in range(NSTATES)] + print(f'Trajectory {traj_id} energies at step 1', energies) -# 6) gradient for root 2 -grad2 = casscf.compute_gradient(2) -print('Gradient root 2 geom2', grad2) + grad = casscf.compute_gradient(GRAD_ROOT, traj_id=traj_id) + print(f'Trajectory {traj_id} gradient root {GRAD_ROOT} at step 1', grad) -# 7) time-overlap matrix for 3 roots -overlap = casscf.time_overlap_matrix(3) -print('Time-overlap matrix', overlap) + overlap = casscf.time_overlap_matrix(NSTATES, traj_id=traj_id) + print(f'Trajectory {traj_id} time-overlap matrix', overlap) + + assert overlap.shape == (NSTATES, NSTATES) assert isinstance(casscf, ElectronicStructureStrategy) +assert casscf.ntraj == NTRAJ diff --git a/examples/packages/pyscf/test_casscf_nacv.py b/examples/packages/pyscf/test_casscf_nacv.py index 7097cbafc..19337dd4e 100644 --- a/examples/packages/pyscf/test_casscf_nacv.py +++ b/examples/packages/pyscf/test_casscf_nacv.py @@ -43,9 +43,9 @@ # from libra_py.packages.pyscf.interfaces import ElectronicStructureStrategy, MolecularGeometry NSTATES = 2 -DISTANCE_START_BOHR = 7.0 +DISTANCE_START_BOHR = 7.5 DISTANCE_STOP_BOHR = 15.0 -DISTANCE_STEP_BOHR = 1.0 +DISTANCE_STEP_BOHR = 0.25 basis_dict = {'Li': 'sto-3g', 'F': '6-311+g*'} cas_list = [4, 7, 11, 14, 17]# 0-indexed: 3 (F2pz), 6 (Li2s), 10 (F5pz), 13 (F5s), 16 (F4pz) @@ -60,7 +60,8 @@ basis=basis_dict, unit='Bohr', charge=0, - cas_list=cas_list + cas_list=cas_list, + use_prev_ci=True, ) for step, d in enumerate(distances): diff --git a/examples/packages/pyscf/test_cisd.py b/examples/packages/pyscf/test_cisd.py index d32dba7ff..51ca7de73 100644 --- a/examples/packages/pyscf/test_cisd.py +++ b/examples/packages/pyscf/test_cisd.py @@ -33,30 +33,47 @@ from libra_py.packages.pyscf.interfaces import ElectronicStructureStrategy, MolecularGeometry import numpy as np -geom1 = MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0,0.0,0.0],[0.0,0.0,0.7746]])) -geom2 = MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0,0.0,0.0],[0.0,0.0,1.0]])) +NTRAJ = 2 +NSTATES = 3 +GRAD_ROOT = 2 -cisd = CISD(nroots=3, basis='sto-3g', charge=1) +geom_step0 = [ + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.7746]])), + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 0.9000]])), +] -cisd.set_geom_and_run_hf(geom1) -energies1 = [cisd.compute_energy(root) for root in range(3)] -print('Energies at geom1', energies1) +geom_step1 = [ + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.0000]])), + MolecularGeometry(atom_labels=['He', 'H'], coords_angstrom=np.array([[0.0, 0.0, 0.0], [0.0, 0.0, 1.1000]])), +] -grad1 = cisd.compute_gradient(2) -print('Gradient root 2 geom1', grad1) +cisd = CISD(nroots=NSTATES, basis='sto-3g', charge=1, ntraj=NTRAJ) -cisd.set_geom_and_run_hf(geom2) -energies2 = [cisd.compute_energy(root) for root in range(3)] -print('Energies at geom2', energies2) +for traj_id, geom in enumerate(geom_step0): + cisd.set_geom_and_run_hf(geom, traj_id=traj_id) -grad2 = cisd.compute_gradient(2) -print('Gradient root 2 geom2', grad2) + energies = [cisd.compute_energy(root, traj_id=traj_id) for root in range(NSTATES)] + print(f'Trajectory {traj_id} energies at step 0', energies) -overlap = cisd.time_overlap_matrix(3) -print('Time-overlap matrix', overlap) + grad = cisd.compute_gradient(GRAD_ROOT, traj_id=traj_id) + print(f'Trajectory {traj_id} gradient root {GRAD_ROOT} at step 0', grad) + +for traj_id, geom in enumerate(geom_step1): + cisd.set_geom_and_run_hf(geom, traj_id=traj_id) + + energies = [cisd.compute_energy(root, traj_id=traj_id) for root in range(NSTATES)] + print(f'Trajectory {traj_id} energies at step 1', energies) + + grad = cisd.compute_gradient(GRAD_ROOT, traj_id=traj_id) + print(f'Trajectory {traj_id} gradient root {GRAD_ROOT} at step 1', grad) + + overlap = cisd.time_overlap_matrix(NSTATES, traj_id=traj_id) + print(f'Trajectory {traj_id} time-overlap matrix', overlap) + + assert overlap.shape == (NSTATES, NSTATES) assert isinstance(cisd, ElectronicStructureStrategy) -assert overlap.shape == (3,3) +assert cisd.ntraj == NTRAJ #if __name__ == "__main__": # pass diff --git a/src/libra_py/packages/pyscf/implementations/casscf.py b/src/libra_py/packages/pyscf/implementations/casscf.py index e05446f18..cd8d12103 100644 --- a/src/libra_py/packages/pyscf/implementations/casscf.py +++ b/src/libra_py/packages/pyscf/implementations/casscf.py @@ -25,153 +25,182 @@ @dataclass -class Cache: +class CASSCFTrajectoryState: + mol: Optional[Any] = None + mf: Optional[Any] = None + mc: Optional[Any] = None + ao_overlap: Optional[np.ndarray] = None prev_mol: Optional[Any] = None - prev_mc: Optional[Any] = None prev_mf: Optional[Any] = None + prev_mc: Optional[Any] = None class CASSCF(ElectronicStructureStrategy): """PySCF-based CASSCF backend for the universal ES interface.""" - # 1) when a geom is set the HF is run; the MF is written to the member attribute - # `_mf` - # 2) when the CASSCF energy is requested for a certain root, the CASSCF is run - # (number of roots requested = self._nroots). The resulting MC object is - # stored in the member attribute `mc`, which also contains energies of other - # states. + # Each trajectory stores its own PySCF molecule, HF object, CASSCF object, + # AO overlap, and previous-step state in CASSCFTrajectoryState. def __init__( self, - mol: Optional[Any] = None, norbcas: int = 0, - nelecas:int = 0, + nelecas: int = 0, nroots: int = 1, basis: str = "sto-3g", unit: str = "Angstrom", charge: int = 0, cas_list: Optional[List[int]] = None, + ntraj: int = 1, + use_prev_ci: bool = False, ) -> None: - #setting up the initial state - super().__init__(mol=mol, nroots=nroots, basis=basis, unit=unit, charge=int(charge)) + super().__init__( + nroots=nroots, + basis=basis, + unit=unit, + charge=int(charge), + ntraj=ntraj, + ) + self._norbcas: int = norbcas - self._nelecas: Union[int, Tuple[int, int]] = nelecas - self._nroots: int = nroots - self._basis: str = basis - self._unit: str = unit # default to Angstrom - self._charge: int = int(charge) - self._mf: Optional[Any] = None - self._mc: Optional[Any] = None - self._ao_overlap: Optional[np.ndarray] = None # AO overlap between consecutive geoms for time-overlap computation - self._cache: Cache = Cache() - self._cas_list: Optional[List[int]] = cas_list # list of active space orbital indices (0-based); if None, use the first `n - - def save_cache(self) -> None: - # Cache the current state directly (None is fine for first geometry). - self._cache.prev_mol = self._mol - self._cache.prev_mc = self._mc - self._cache.prev_mf = self._mf - - def run_hf(self) -> None: - # Clear the current state. - self._mc = None - self._ao_overlap = None - - # Set up the new molecule and run HF. - charge: int = self._charge - self._mol = gto.M( + self._nelecas: Union[int, Tuple[int, int]] = nelecas + self._cas_list: Optional[List[int]] = cas_list + self._use_prev_ci: bool = use_prev_ci + + def _get_traj_state(self, traj_id: int) -> CASSCFTrajectoryState: + state = self._traj_states[traj_id] + if state is None: + state = CASSCFTrajectoryState() + self._traj_states[traj_id] = state + return state + + def set_geom(self, geom: MolecularGeometry, traj_id: int = 0) -> None: + state = self._get_traj_state(traj_id) + + state.mol = gto.M( atom=";".join( f"{label} {coord[0]} {coord[1]} {coord[2]}" - for label, coord in zip(self._geom.atom_labels, self._geom.coords_angstrom) + for label, coord in zip(geom.atom_labels, geom.coords_angstrom) ), basis=self._basis, unit=self._unit, - charge=charge, + charge=self._charge, spin=0, ) - #compute the AO overlap between new geom and the previous one - if self._cache.prev_mol is not None: - self._ao_overlap = gto.intor_cross("int1e_ovlp", self._cache.prev_mol, self._mol) - self._mf = scf.RHF(self._mol).run(verbose=0) + state.mc = None + state.ao_overlap = None + + if state.prev_mol is not None: + state.ao_overlap = gto.intor_cross("int1e_ovlp", state.prev_mol, state.mol) - def compute_energy(self, root: int) -> float: - if self._mf is None: + def save_cache(self, traj_id: int = 0) -> None: + # Cache the current state directly (None is fine for first geometry). + state = self._get_traj_state(traj_id) + state.prev_mol = state.mol + state.prev_mc = state.mc + state.prev_mf = state.mf + + def run_hf(self, traj_id: int = 0) -> None: + state = self._get_traj_state(traj_id) + + if state.mol is None: + raise ValueError(f"Geometry for trajectory {traj_id} has not been set.") + + dm0 = None + if state.prev_mf is not None: + try: + dm0 = state.prev_mf.make_rdm1() + except Exception: + dm0 = None + + mf = scf.RHF(state.mol) + mf.verbose = 0 + if dm0 is not None: + mf.kernel(dm0=dm0) + else: + mf.kernel() + + state.mf = mf + + def compute_energy(self, root: int, traj_id: int = 0) -> float: + state = self._get_traj_state(traj_id) + + if state.mf is None: raise ValueError("HF must be run before computing CASSCF energies.") + if root < 0 or root >= self._nroots: + raise IndexError(f"Requested root {root}, but only {self._nroots} roots are available.") + + if state.mc is None: + mc = mcscf.CASSCF(state.mf, self._norbcas, self._nelecas) + mc.fcisolver = fci.direct_spin0.FCI(state.mol) + mc.fcisolver.nroots = self._nroots - if self._mc is None: - self._mc = mcscf.CASSCF(self._mf, self._norbcas, self._nelecas) - self._mc.fcisolver = fci.direct_spin0.FCI(self._mol) - self._mc.fcisolver.nroots = self._nroots if self._nroots > 1: - self._mc = self._mc.state_average_([1.0 / self._nroots] * self._nroots) # equal weights as default; required for gradients in pyscf + mc = mc.state_average_([1.0 / self._nroots] * self._nroots) # hardcoded equal weights for now; - # Correctly use currently converged HF spatial orbitals as initial guess for SA-CASSCF - # Sort orbitals if cas_list is provided - mo_coeff = self._mf.mo_coeff + mo_coeff = state.mf.mo_coeff if self._cas_list is not None: - # Notice: In PySCF mcscf.sort_mo, cas_list must be a list of 1-based indices - # if your cas_list values passed in __init__ are 1-based, directly use it: - mo_coeff = mcscf.sort_mo(self._mc, mo_coeff, self._cas_list) - - # Prepare CI coefficients from previous step (only if use_prev_ci is True) + mo_coeff = mcscf.sort_mo(mc, mo_coeff, self._cas_list) + ci0 = None - if getattr(self, '_use_prev_ci', False) and self._cache.prev_mc is not None: - ci0 = getattr(self._cache.prev_mc, 'ci', None) + if self._use_prev_ci and state.prev_mc is not None: + ci0 = getattr(state.prev_mc, "ci", None) if ci0 is not None: - self._mc.kernel(mo_coeff, ci0=ci0) + mc.kernel(mo_coeff, ci0=ci0) else: - self._mc.kernel(mo_coeff) + mc.kernel(mo_coeff) - e_states: Optional[Sequence[float]] = getattr(self._mc, "e_states", None) + state.mc = mc + + e_states: Optional[Sequence[float]] = getattr(state.mc, "e_states", None) if e_states is not None: if root < 0 or root >= len(e_states): raise IndexError(f"Requested root {root}, but only {len(e_states)} roots are available.") return float(np.asarray(e_states)[root]) + if root != 0: raise IndexError("Only root 0 is available for a single-state CASSCF calculation.") - return float(self._mc.e_tot) + return float(state.mc.e_tot) - def compute_gradient(self, root: int) -> np.ndarray: - if self._mc is None: + def compute_gradient(self, root: int, traj_id: int = 0) -> np.ndarray: + state = self._get_traj_state(traj_id) + if state.mc is None: raise ValueError("CASSCF must be run before computing gradients.") - e_states: Optional[Sequence[float]] = getattr(self._mc, "e_states", None) + e_states: Optional[Sequence[float]] = getattr(state.mc, "e_states", None) if e_states is not None: if root < 0 or root >= len(e_states): raise IndexError(f"Requested root {root}, but only {len(e_states)} roots are available.") - return np.asarray(self._mc.nuc_grad_method(state=root).kernel()) + return np.asarray(state.mc.nuc_grad_method(state=root).kernel()) if root != 0: raise IndexError("Only root 0 is available for a single-state CASSCF calculation.") - return np.asarray(self._mc.nuc_grad_method().kernel()) + return np.asarray(state.mc.nuc_grad_method().kernel()) - def time_overlap_matrix(self, nroots: int) -> np.ndarray: - if self._mc is None: + def time_overlap_matrix(self, nroots: int, traj_id: int = 0) -> np.ndarray: + state = self._get_traj_state(traj_id) + if state.mc is None: raise ValueError("CASSCF must be run before computing time-overlap matrix.") - if self._cache is None or self._cache.prev_mf is None or self._cache.prev_mol is None: - raise ValueError("Previous and current HF/molecule states are required for time-overlap computation.") - if self._ao_overlap is None: + if state.prev_mc is None or state.prev_mol is None: + raise ValueError("Previous and current CASSCF/molecule states are required for time-overlap computation.") + if state.ao_overlap is None: raise ValueError("Cached AO overlap between consecutive geometries is not available.") if nroots <= 0: raise ValueError(f"nroots must be positive, got {nroots}") + if nroots > self._nroots: + raise ValueError(f"Requested {nroots} roots, but CASSCF is configured for {self._nroots} roots") nroots = int(nroots) - if self._mf is None or self._mol is None: - raise ValueError("Current HF/molecule state is required for time-overlap computation.") - - prev_casci = mcscf.CASCI(self._cache.prev_mf, self._norbcas, self._nelecas) - prev_casci.fcisolver = fci.direct_spin0.FCISolver(self._cache.prev_mol) - prev_casci.fcisolver.nroots = nroots - h1prev, _ = prev_casci.get_h1eff(prev_casci.mo_coeff) - h2prev = prev_casci.get_h2cas(prev_casci.mo_coeff) - _, prev_roots = prev_casci.fcisolver.kernel(h1prev, h2prev, prev_casci.ncas, prev_casci.nelecas, nroots=nroots) - - curr_casci = mcscf.CASCI(self._mf, self._norbcas, self._nelecas) - curr_casci.fcisolver = fci.direct_spin0.FCISolver(self._mol) - curr_casci.fcisolver.nroots = nroots - h1curr, _ = curr_casci.get_h1eff(curr_casci.mo_coeff) - h2curr = curr_casci.get_h2cas(curr_casci.mo_coeff) - _, curr_roots = curr_casci.fcisolver.kernel(h1curr, h2curr, curr_casci.ncas, curr_casci.nelecas, nroots=nroots) - + if state.mol is None: + raise ValueError("Current molecule state is required for time-overlap computation.") + + prev_roots = getattr(state.prev_mc, "ci", None) + curr_roots = getattr(state.mc, "ci", None) + if prev_roots is None or curr_roots is None: + raise ValueError("Previous and current CI vectors are required for time-overlap computation.") + + if not isinstance(prev_roots, (list, tuple)): #treat single root as list of one root for uniformity + prev_roots = [prev_roots] + if not isinstance(curr_roots, (list, tuple)): + curr_roots = [curr_roots] if len(prev_roots) < nroots or len(curr_roots) < nroots: raise ValueError( f"Requested {nroots} roots, but only {len(prev_roots)} previous and {len(curr_roots)} current roots are available." @@ -180,38 +209,42 @@ def time_overlap_matrix(self, nroots: int) -> np.ndarray: prev_roots = [np.asarray(v) for v in prev_roots[:nroots]] curr_roots = [np.asarray(v) for v in curr_roots[:nroots]] - mo_prev_act = np.asarray(prev_casci.mo_coeff)[:, prev_casci.ncore : prev_casci.ncore + prev_casci.ncas] - mo_curr_act = np.asarray(curr_casci.mo_coeff)[:, curr_casci.ncore : curr_casci.ncore + curr_casci.ncas] - s12_mo = mo_prev_act.T @ self._ao_overlap @ mo_curr_act + mo_prev_act = np.asarray(state.prev_mc.mo_coeff)[:, state.prev_mc.ncore : state.prev_mc.ncore + state.prev_mc.ncas] + mo_curr_act = np.asarray(state.mc.mo_coeff)[:, state.mc.ncore : state.mc.ncore + state.mc.ncas] + s12_mo = mo_prev_act.T @ state.ao_overlap @ mo_curr_act overlap = np.zeros((nroots, nroots), dtype=float) - nelecas = prev_casci.nelecas + nelecas = state.prev_mc.nelecas for i in range(nroots): for j in range(nroots): overlap[i, j] = fci.addons.overlap( prev_roots[i], curr_roots[j], - prev_casci.ncas, + state.prev_mc.ncas, nelecas, s=s12_mo, ) overlap = np.asarray(np.real_if_close(overlap)) + # Phase correction. Ensure the diagonal elements of the overlap matrix are positive by flipping signs if necessary. for i in range(nroots): if overlap[i, i] < 0: overlap[i, :] = -overlap[i, :] return overlap - def compute_nac_vectors(self, use_etfs: bool = True) -> np.ndarray: - if self._mc is None: + def compute_nac_vectors(self, use_etfs: bool = True, traj_id: int = 0) -> np.ndarray: + state = self._get_traj_state(traj_id) + if state.mc is None: raise ValueError("CASSCF must be run before computing NAC vectors.") + if state.mol is None: + raise ValueError("Current molecule is required before computing NAC vectors.") nstates = self._nroots - natm = int(self._mol.natm) + natm = int(state.mol.natm) - mc_nacs = self._mc.nac_method() + mc_nacs = state.mc.nac_method() nacv = np.zeros((nstates, nstates, natm, 3), dtype=np.float64) for ket in range(nstates): @@ -224,4 +257,3 @@ def compute_nac_vectors(self, use_etfs: bool = True) -> np.ndarray: ) return nacv - diff --git a/src/libra_py/packages/pyscf/implementations/cisd.py b/src/libra_py/packages/pyscf/implementations/cisd.py index 771f0cabb..72c6f85a0 100644 --- a/src/libra_py/packages/pyscf/implementations/cisd.py +++ b/src/libra_py/packages/pyscf/implementations/cisd.py @@ -22,152 +22,205 @@ from typing import Any, Optional, Sequence import numpy as np -from pyscf import ci +from pyscf import ci, gto, scf from libra_py.packages.pyscf.interfaces import ElectronicStructureStrategy, MolecularGeometry -#from libra_py.packages.pyscf.utils import run_rhf_for_geometry + @dataclass -class CISDCache: +class CISDTrajectoryState: + mol: Optional[Any] = None + mf: Optional[Any] = None + ci: Optional[Any] = None + ao_overlap: Optional[np.ndarray] = None prev_mol: Optional[Any] = None prev_mf: Optional[Any] = None prev_ci: Optional[Any] = None + class CISD(ElectronicStructureStrategy): """PySCF-based CISD backend for the universal ES interface.""" def __init__( self, - mol: Optional[Any] = None, nroots: int = 1, basis: str = "sto-3g", unit: str = "Angstrom", charge: int = 0, use_prev_ci: bool = False, + ntraj: int = 1, ) -> None: - super().__init__(mol=mol, nroots=nroots, basis=basis, unit=unit, charge=int(charge)) - self._ci: Optional[Any] = None - self._cache: CISDCache = CISDCache() + super().__init__( + nroots=nroots, + basis=basis, + unit=unit, + charge=int(charge), + ntraj=ntraj, + ) + self._use_prev_ci: bool = use_prev_ci - def save_cache(self) -> None: - self._cache.prev_mol = self._mol - self._cache.prev_mf = self._mf - self._cache.prev_ci = self._ci + def _get_traj_state(self, traj_id: int) -> CISDTrajectoryState: + state = self._traj_states[traj_id] + if state is None: + state = CISDTrajectoryState() + self._traj_states[traj_id] = state + return state + + def set_geom(self, geom: MolecularGeometry, traj_id: int = 0) -> None: + state = self._get_traj_state(traj_id) + state.ci = None + state.ao_overlap = None + + state.mol = gto.M( + atom=";".join( + f"{label} {coord[0]} {coord[1]} {coord[2]}" + for label, coord in zip(geom.atom_labels, geom.coords_angstrom) + ), + basis=self._basis, + unit=self._unit, + charge=self._charge, + spin=0, + ) + + if state.prev_mol is not None: + state.ao_overlap = gto.intor_cross("int1e_ovlp", state.prev_mol, state.mol) + + def save_cache(self, traj_id: int = 0) -> None: + state = self._get_traj_state(traj_id) + state.prev_mol = state.mol + state.prev_mf = state.mf + state.prev_ci = state.ci + + def run_hf(self, traj_id: int = 0) -> None: + state = self._get_traj_state(traj_id) - def run_hf(self) -> None: - self._ci = None - self._ao_overlap = None + if state.mol is None: + raise ValueError(f"Geometry for trajectory {traj_id} has not been set.") dm0 = None - if self._cache.prev_mf is not None: + if state.prev_mf is not None: try: - dm0 = self._cache.prev_mf.make_rdm1() + dm0 = state.prev_mf.make_rdm1() except Exception: dm0 = None - self._mol, self._mf, self._ao_overlap = run_rhf_for_geometry( - self._geom, - basis=self._basis, - unit=self._unit, - charge=self._charge, - prev_mol=self._cache.prev_mol, - overlap_integral="int1e_ovlp", - spin=0, - verbose=0, - dm0=dm0, - ) + mf = scf.RHF(state.mol) + mf.verbose = 0 + if dm0 is not None: + mf.kernel(dm0=dm0) + else: + mf.kernel() + state.mf = mf - def compute_energy(self, root: int) -> float: - if self._mf is None: - raise ValueError("HF must be run before computing CISD energies.") + def compute_energy(self, root: int, traj_id: int = 0) -> float: + state = self._get_traj_state(traj_id) + + if state.mf is None: + raise ValueError("HF must be run before computing CISD energies.") if root < 0 or root >= self._nroots: - raise IndexError(f"Requested root {root}, but only {self._nroots} roots are available.") - - if self._ci is None: - self._ci = ci.cisd.CISD(self._mf) - self._ci.nroots = self._nroots - self._ci.verbose = 0 - # Use CI coefficients from previous step if requested - ci0 = None - if self._use_prev_ci and self._cache.prev_ci is not None: - ci0 = getattr(self._cache.prev_ci, "ci", None) - if ci0 is not None: - self._ci.kernel(ci0=ci0) - else: - self._ci.kernel() - - e_tot: Optional[Sequence[float]] = getattr(self._ci, "e_tot", None) + raise IndexError(f"Requested root {root}, but only {self._nroots} roots are available.") + + if state.ci is None: + cisd = ci.cisd.CISD(state.mf) + cisd.nroots = self._nroots + cisd.verbose = 0 + + ci0 = None + if self._use_prev_ci and state.prev_ci is not None: + ci0 = getattr(state.prev_ci, "ci", None) + + if ci0 is not None: + cisd.kernel(ci0=ci0) + else: + cisd.kernel() + + state.ci = cisd + + e_tot: Optional[Sequence[float]] = getattr(state.ci, "e_tot", None) if isinstance(e_tot, (list, tuple, np.ndarray)): - return float(np.asarray(e_tot)[root]) + return float(np.asarray(e_tot)[root]) if root != 0: - raise IndexError("Only root 0 is available for this CISD calculation.") + raise IndexError("Only root 0 is available for this CISD calculation.") return float(e_tot) - def compute_gradient(self, root: int) -> np.ndarray: - if self._ci is None: - self.compute_energy(root=0) - # normalize CI roots into a list - ci_vectors = self._ci.ci + def compute_gradient(self, root: int, traj_id: int = 0) -> np.ndarray: + state = self._get_traj_state(traj_id) + if state.ci is None: + self.compute_energy(root=0, traj_id=traj_id) + + ci_vectors = state.ci.ci if isinstance(ci_vectors, (list, tuple)): - ci_roots = [np.asarray(vec) for vec in ci_vectors] + ci_roots = [np.asarray(vec) for vec in ci_vectors] else: - ci_roots = [np.asarray(ci_vectors)] + ci_roots = [np.asarray(ci_vectors)] if root < 0 or root >= len(ci_roots): - raise IndexError(f"Requested root {root}, but only {len(ci_roots)} roots are available.") + raise IndexError(f"Requested root {root}, but only {len(ci_roots)} roots are available.") if len(ci_roots) == 1: - if root != 0: - raise IndexError("Only root 0 is available for a single-state CISD calculation.") - return np.asarray(self._ci.nuc_grad_method().kernel()) - return np.asarray(self._ci.nuc_grad_method().kernel(state=root)) + if root != 0: + raise IndexError("Only root 0 is available for a single-state CISD calculation.") + return np.asarray(state.ci.nuc_grad_method().kernel()) + return np.asarray(state.ci.nuc_grad_method().kernel(state=root)) + + def time_overlap_matrix(self, nroots: int, traj_id: int = 0) -> np.ndarray: + state = self._get_traj_state(traj_id) - def time_overlap_matrix(self, nroots: int) -> np.ndarray: if nroots <= 0: - raise ValueError(f"nroots must be positive, got {nroots}") + raise ValueError(f"nroots must be positive, got {nroots}") if nroots > self._nroots: - raise ValueError(f"Requested {nroots} roots, but CISD is configured for {self._nroots} roots") - - if self._ci is None: - raise ValueError("CISD must be run before computing time-overlap matrix.") + raise ValueError(f"Requested {nroots} roots, but CISD is configured for {self._nroots} roots") + if state.ci is None: + raise ValueError("CISD must be run before computing time-overlap matrix.") nroots = int(nroots) - if self._cache.prev_ci is None or self._cache.prev_mf is None or self._cache.prev_mol is None: - raise ValueError("Previous CISD state must exist for time-overlap computation.") - if self._mf is None or self._mol is None: - raise ValueError("Current CISD state must exist for time-overlap computation.") - if self._ao_overlap is None: - raise ValueError("Cached AO overlap between consecutive geometries is not available.") + if state.prev_ci is None or state.prev_mf is None or state.prev_mol is None: + raise ValueError("Previous CISD state must exist for time-overlap computation.") + if state.mf is None or state.mol is None: + raise ValueError("Current CISD state must exist for time-overlap computation.") + if state.ao_overlap is None: + raise ValueError("Cached AO overlap between consecutive geometries is not available.") - s12_mo = reduce(np.dot, (self._cache.prev_mf.mo_coeff.T, self._ao_overlap, self._mf.mo_coeff)) + s12_mo = reduce(np.dot, (state.prev_mf.mo_coeff.T, state.ao_overlap, state.mf.mo_coeff)) - prev_ci_roots = self._cache.prev_ci.ci - curr_ci_roots = self._ci.ci + prev_ci_roots = state.prev_ci.ci + curr_ci_roots = state.ci.ci if isinstance(prev_ci_roots, (list, tuple)): - prev_ci_list = [np.asarray(v) for v in prev_ci_roots] + prev_ci_list = [np.asarray(v) for v in prev_ci_roots] else: - prev_ci_list = [np.asarray(prev_ci_roots)] + prev_ci_list = [np.asarray(prev_ci_roots)] if isinstance(curr_ci_roots, (list, tuple)): - curr_ci_list = [np.asarray(v) for v in curr_ci_roots] + curr_ci_list = [np.asarray(v) for v in curr_ci_roots] else: - curr_ci_list = [np.asarray(curr_ci_roots)] + curr_ci_list = [np.asarray(curr_ci_roots)] + + if len(prev_ci_list) < nroots or len(curr_ci_list) < nroots: + raise ValueError( + f"Requested {nroots} roots, but only {len(prev_ci_list)} previous and {len(curr_ci_list)} current roots are available." + ) - nmo = self._ci.nmo - nelec = self._mol.nelectron // 2 + nmo = state.ci.nmo + nelec = state.mol.nelectron // 2 overlap = np.zeros((nroots, nroots), dtype=float) for i in range(nroots): - for j in range(nroots): - overlap[i, j] = ci.cisd.overlap( - prev_ci_list[i], - curr_ci_list[j], - nmo, - nelec, - s12_mo, - ) - - return np.asarray(np.real_if_close(overlap)) + for j in range(nroots): + overlap[i, j] = ci.cisd.overlap( + prev_ci_list[i], + curr_ci_list[j], + nmo, + nelec, + s12_mo, + ) + + overlap = np.asarray(np.real_if_close(overlap)) + + for i in range(nroots): + if overlap[i, i] < 0: + overlap[i, :] = -overlap[i, :] + + return overlap diff --git a/src/libra_py/packages/pyscf/interfaces.py b/src/libra_py/packages/pyscf/interfaces.py index b36e28285..2ed1d0d2f 100644 --- a/src/libra_py/packages/pyscf/interfaces.py +++ b/src/libra_py/packages/pyscf/interfaces.py @@ -59,49 +59,51 @@ class ElectronicStructureStrategy(ABC): def __init__( self, - mol: Optional[Any] = None, nroots: int = 1, basis: str = "sto-3g", unit: str = "Angstrom", charge: int = 0, + ntraj: int = 1, ) -> None: - self._mol = mol self._nroots = nroots self._basis = basis self._unit = unit self._charge = charge - self._mf = None - self._ao_overlap = None - self._geom = None + self._traj_states = [None] * ntraj # electronic states for each trajectory,either pointers to ram or disk paths to cache files @property def nstates(self) -> int: return self._nroots + + @property + def ntraj(self) -> int: + return len(self._traj_states) # ------------------------------------------------------------------ # Core computation (required) # ------------------------------------------------------------------ - - def save_cache(self) -> None: + @abstractmethod + def save_cache(self, traj_id: int = 0) -> None: + pass + + @abstractmethod + def set_geom(self, geom: MolecularGeometry, traj_id: int = 0) -> None: pass - - def set_geom(self, geom: MolecularGeometry) -> None: - self._geom = geom @abstractmethod - def run_hf(self) -> None: + def run_hf(self, traj_id: int = 0) -> None: pass - def set_geom_and_run_hf(self, geom: MolecularGeometry) -> None: - self.save_cache() - self.set_geom(geom) - self.run_hf() + def set_geom_and_run_hf(self, geom: MolecularGeometry, traj_id: int = 0) -> None: + self.save_cache(traj_id = traj_id) + self.set_geom(geom,traj_id = traj_id) + self.run_hf(traj_id = traj_id) @abstractmethod - def compute_energy(self, root: int) -> float: + def compute_energy(self, root: int, traj_id: int = 0) -> float:#hard coded to single root for now; """Return the total energy (Hartree) for *root*.""" - def compute_gradient(self, root: int) -> np.ndarray: + def compute_gradient(self, root: int, traj_id: int = 0) -> np.ndarray: """Return the nuclear gradient for *root*. Returns @@ -109,9 +111,12 @@ def compute_gradient(self, root: int) -> np.ndarray: np.ndarray Shape ``(natoms, 3)`` in Hartree/Bohr. """ + raise NotImplementedError( + "This backend does not provide nuclear gradients." + ) - def time_overlap_matrix(self, nroots: int) -> np.ndarray: + def time_overlap_matrix(self, nroots: int, traj_id: int = 0) -> np.ndarray: """Return the time-overlap matrix ````. Returns @@ -124,7 +129,7 @@ def time_overlap_matrix(self, nroots: int) -> np.ndarray: "use time-overlap-based NACs instead." ) - def compute_nac_vectors(self, **kwargs: Any) -> np.ndarray: + def compute_nac_vectors(self, *, traj_id: int = 0, **kwargs: Any) -> np.ndarray: """Return all NAC vectors ``d_{ij}`` between states. Returns @@ -142,3 +147,4 @@ def compute_nac_vectors(self, **kwargs: Any) -> np.ndarray: "This backend does not provide explicit NAC vectors; " "use time-overlap-based NACs instead." ) + diff --git a/src/libra_py/packages/pyscf/methods.py b/src/libra_py/packages/pyscf/methods.py index 630e92092..8442a917e 100644 --- a/src/libra_py/packages/pyscf/methods.py +++ b/src/libra_py/packages/pyscf/methods.py @@ -36,7 +36,7 @@ class tmp: pass -def pyscf_compute_adi(q, params, full_id): +def es_compute_adi(q, params, full_id): # ================= Decode trajectory index ================= Id = Cpp2Py(full_id)