add pao to localization notebook but results bad??

Author: MIWdlB

docs/source/notebooks/localization.ipynb

@@ -44,6 +44,7 @@ class VirtualLocalizerTypes(Enum):
 
     CONCENTRIC = "cl"
     PROJECTED_AO = "pao"
+    DISABLE = "disable"
 
 
 XYZGeometry = Annotated[

nbed/config.py

@@ -93,7 +93,7 @@ def _build_mol(self) -> gto.mole:
         logger.info("Molecule input geometry: %s", self.config.geometry)
         # geometry is raw xyz string
         full_mol = gto.Mole(
-            atom=self.config.geometry[3:],
+            atom=self.config.geometry[2:],
             basis=self.config.basis,
             charge=self.config.charge,
             unit=self.config.unit,
@@ -966,6 +966,9 @@ def collect_results(
                     norm_cutoff=self.config.norm_cutoff,
                     overlap_cutoff=self.config.overlap_cutoff,
                 )
+                result["scf"] = result["pao"].localize_virtual()
+            case VirtualLocalizerTypes.DISABLE:
+                logger.debug("Not performing virtual localization.")
             case _:
                 logger.debug(
                     f"Driver does not have a method implemented for {self.config.virtual_localization}"

nbed/driver.py

@@ -45,7 +45,7 @@ def __init__(
         n_active_atoms: int,
     ):
         """Initialise class."""
-        logger.debug("Initialising LocalizerEnum.")
+        logger.debug("Initialising OccupiedLocalizerTypes.")
         if global_scf.mo_coeff is None:
             logger.debug("SCF method not initialised, running now...")
             global_scf.run()

nbed/localizers/occupied/base.py

@@ -139,14 +139,16 @@ def _localize_spin(
             np.ndarray: Localized C matrix of environment orbitals.
             np.ndarray: Localized C matrix of all occpied orbitals.
         """
-        logger.debug("Localising with SPADE.")
+        logger.debug("Localising spin with SPADE.")
+        logger.debug(f"{c_matrix.shape=}")
+        logger.debug(f"{occupancy=}")
+        logger.debug(f"{n_mo_overwrite=}")
 
         # We want the same partition for each spin.
         # It wouldn't make sense to have different spin states be localized differently.
 
         n_occupied_orbitals = np.count_nonzero(occupancy)
         occupied_orbitals = c_matrix[:, :n_occupied_orbitals]
-        logger.debug(f"{occupancy=}")
         logger.debug(f"{n_occupied_orbitals} occupied AOs.")
 
         n_act_aos = self._global_scf.mol.aoslice_by_atom()[self._n_active_atoms - 1][-1]
@@ -174,7 +176,14 @@ def _localize_spin(
             n_act_mos: int = n_mo_overwrite
         else:
             value_diffs = sigma[:-1] - sigma[1:]
-            n_act_mos: int = np.argmax(value_diffs) + 1
+            logger.debug("Singular value differences %s", value_diffs)
+            # It is possible to choose an active subsystem for which all
+            # singular values are 1 (i.e. the whole system)
+            # we want to avoid numerical error forcing random orbital assignment
+            if np.allclose(value_diffs, [0] * len(value_diffs)):
+                n_act_mos = len(sigma)
+            else:
+                n_act_mos: int = np.argmax(value_diffs) + 1
 
         n_env_mos = n_occupied_orbitals - n_act_mos
         logger.debug(f"{n_act_mos} active MOs.")

nbed/localizers/occupied/spade.py

@@ -135,6 +135,9 @@ def _localize_virtual_spin(
             np.ndarray: The update MO coefficient matrix
         """
         logger.debug("Running concentric localiztion for single spin.")
+        logger.debug(f"{occ=}")
+        logger.debug(f"{mo_coeff.shape=}")
+        logger.debug(f"{fock_operator.shape=}")
 
         effective_virt = mo_coeff[:, occ == 0]
         logger.debug(f"N effective virtuals: {effective_virt.shape}")

nbed/localizers/virtual/concentric.py

@@ -7,7 +7,7 @@
 from pyscf.lib import StreamObject
 from scipy.linalg import fractional_matrix_power
 
-from .base import VirtualLocalizer
+from nbed.localizers.virtual.base import VirtualLocalizer
 
 logger = logging.getLogger(__name__)
 
@@ -19,7 +19,7 @@ def __init__(
         self,
         embedded_scf: StreamObject,
         n_active_atoms: int,
-        c_loc_occ: tuple(NDArray, NDArray | None),
+        c_loc_occ: tuple[NDArray, NDArray | None],
         norm_cutoff: float = 0.05,
         overlap_cutoff=1e-5,
     ):
@@ -38,56 +38,70 @@ def localize_virtual(self) -> StreamObject:
 
         if self.c_loc_occ[1] is None:
             logger.debug("Runing PAO for spinless system.")
-            virtuals = _localize__spin_pao(
+            virtuals = _localize_virtual_spin_pao(
                 self.c_loc_occ[0],
                 ao_overlap,
                 n_act_aos,
                 self.norm_cutoff,
                 self.overlap_cutoff,
             )
-
+            logger.debug(f"{virtuals.shape=}")
             occ_mo_coeff = self.embedded_scf.mo_coeff[:, self.embedded_scf.mo_occ > 0]
-            self.embedded_scf.mo_coeff = np.hstack((occ_mo_coeff, virtuals))
-            self.embedded_scf.mo_occ = self.embedded_scf.mo_occ[
-                : self.embedded_scf.mo_coeff.shape[-1]
-            ]
+            mo_coeff = np.hstack((occ_mo_coeff, virtuals))
+
+            n_mos = mo_coeff.shape[-1]
+            n_occ = np.count_nonzero(self.embedded_scf.mo_occ[0])
+            mo_occ = np.array(n_occ * [1] + (n_mos - n_occ) * [0])
 
         else:  # Restricted open shell
             logger.debug("Running PAO for each spin separately.")
-            alpha_virtuals = _localize__spin_pao(
+            alpha_virtuals = _localize_virtual_spin_pao(
                 self.c_loc_occ[0],
                 ao_overlap,
                 n_act_aos,
                 self.norm_cutoff,
                 self.overlap_cutoff,
             )
-            beta_virtuals = _localize__spin_pao(
+            beta_virtuals = _localize_virtual_spin_pao(
                 self.c_loc_occ[1],
                 ao_overlap,
                 n_act_aos,
                 self.norm_cutoff,
                 self.overlap_cutoff,
             )
-
-            occ_mo_coeff = self.embedded_scf.mo_coeff[
-                :, :, self.embedded_scf.mo_occ > 0
+            logger.debug(f"{alpha_virtuals.shape=}")
+            logger.debug(f"{beta_virtuals.shape=}")
+            alpha_occ_mo_coeff = self.embedded_scf.mo_coeff[0][
+                :, self.embedded_scf.mo_occ[0] > 0
+            ]
+            beta_occ_mo_coeff = self.embedded_scf.mo_coeff[1][
+                :, self.embedded_scf.mo_occ[1] > 0
             ]
-            self.embedded_scf.mo_coeff = np.vstack(
+            mo_coeff = np.array(
+                [
+                    np.hstack((alpha_occ_mo_coeff, alpha_virtuals)),
+                    np.hstack((beta_occ_mo_coeff, beta_virtuals)),
+                ]
+            )
+            n_mos = mo_coeff.shape[-1]
+            alpha_n_occ = np.count_nonzero(self.embedded_scf.mo_occ[0])
+            beta_n_occ = np.count_nonzero(self.embedded_scf.mo_occ[1])
+            mo_occ = np.vstack(
                 (
-                    np.hstack((occ_mo_coeff[0], alpha_virtuals)),
-                    np.hstack((occ_mo_coeff[1], beta_virtuals)),
+                    np.array(alpha_n_occ * [1] + (n_mos - alpha_n_occ) * [0]),
+                    np.array(beta_n_occ * [1] + (n_mos - beta_n_occ) * [0]),
                 )
             )
-            self.embedded_scf.mo_occ = self.embedded_scf.mo_occ[
-                :, : self.embedded_scf.mo_coeff.shape[-1]
-            ]
+
+        self.embedded_scf.mo_coeff = mo_coeff
+        self.embedded_scf.mo_occ = mo_occ
 
         return self.embedded_scf
 
     # where should the cutoff values come from?
 
 
-def _localize__spin_pao(
+def _localize_virtual_spin_pao(
     c_loc_occ: NDArray,
     ao_overlap: NDArray,
     n_act_aos: int,
@@ -124,6 +138,7 @@ def _localize__spin_pao(
 
     renormalized_paos = s_half @ truncated_paos
     logger.debug(f"{renormalized_paos=}")
+    logger.debug(f"{renormalized_paos.shape=}")
     logger.debug(f"{np.einsum("ij,ij->j", renormalized_paos, renormalized_paos)=}")
     renormalized_paos = renormalized_paos / np.sqrt(
         np.einsum("ij,ij->j", renormalized_paos, renormalized_paos)
@@ -134,14 +149,22 @@ def _localize__spin_pao(
 
     diagonalized_overlap = renormalized_paos.T @ ao_overlap @ renormalized_paos
 
+    logger.debug(f"{diagonalized_overlap.shape=}")
     logger.debug(f"{diagonalized_overlap=}")
 
     eigvals, eigvecs = np.linalg.eigh(diagonalized_overlap)
 
-    logger.debug(eigvals)
+    logger.debug(f"Overlap eigenvalues {eigvals}")
+    logger.debug(f"{overlap_cutoff=}")
 
+    logger.debug(f"{eigvecs.shape=}")
     # How to transform the truncated paos?
-    final_paos = eigvecs[:, eigvals > overlap_cutoff]
-    logger.debug(final_paos)
-
+    final_paos = renormalized_paos[:, eigvals > overlap_cutoff]
+    logger.debug(f"{final_paos=}")
+
+    if final_paos.shape[-1] == 0:
+        logger.warning("No projected atomic orbitals!")
+        logger.warning(
+            "This suggests your active region has no virtual Atomic Orbitals."
+        )
     return final_paos