Repository: troyvvgroup/quemb
Pull request #153: Improved init of pfrag
Author: login: mcocdawc
Created: 2025-04-23T13:26:44Z | Merged: 2025-04-25T20:50:58Z

- Consistent output of `relAO_per_edge_per_frag` even if there are no edges. It will now return a list of empty lists and not an empty list. This means one can just loop over the empty list of edges and not special case for the case of no edges being present.

- Enables a nicer initialisation of the pfrag objects.


- [x] Does it break UBE? I cannot test that myself @lweisburn . :-( 


Diff:
diff --git a/src/quemb/kbe/autofrag.py b/src/quemb/kbe/autofrag.py
index 0e6c66eb2..cca667ba9 100644
--- a/src/quemb/kbe/autofrag.py
+++ b/src/quemb/kbe/autofrag.py
@@ -2351,6 +2351,12 @@ def autogen(
 
         relAO_in_ref_per_edge.append(idx)
 
+    if not AO_per_edge:
+        AO_per_edge = [[] for _ in range(n_frag)]
+        ref_frag_idx_per_edge = [[] for _ in range(n_frag)]
+        relAO_per_edge = [[] for _ in range(n_frag)]
+        relAO_in_ref_per_edge = [[] for _ in range(n_frag)]
+
     return (
         AO_per_frag,
         AO_per_edge,
diff --git a/src/quemb/kbe/fragment.py b/src/quemb/kbe/fragment.py
index ba2d41c4a..8fbc64d5f 100644
--- a/src/quemb/kbe/fragment.py
+++ b/src/quemb/kbe/fragment.py
@@ -6,6 +6,7 @@
 from pyscf.pbc.gto.cell import Cell
 
 from quemb.kbe.autofrag import autogen
+from quemb.kbe.pfrag import Frags
 from quemb.molbe.chemfrag import ChemGenArgs, chemgen
 from quemb.molbe.helper import get_core
 from quemb.shared.typing import (
@@ -13,6 +14,7 @@
     GlobalAOIdx,
     ListOverEdge,
     ListOverFrag,
+    PathLike,
     RelAOIdx,
     RelAOIdxInRef,
 )
@@ -117,6 +119,21 @@ def all_centers_are_origins(self) -> bool:
             )
         )
 
+    def to_Frags(self, I: int, eri_file: PathLike, unitcell_nkpt: int) -> Frags:
+        return Frags(
+            AO_in_frag=self.AO_per_frag[I],
+            ifrag=I,
+            AO_per_edge=self.AO_per_edge[I],
+            eri_file=eri_file,
+            ref_frag_idx_per_edge=self.ref_frag_idx_per_edge[I],
+            relAO_per_edge=self.relAO_per_edge[I],
+            relAO_in_ref_per_edge=self.relAO_in_ref_per_edge[I],
+            centerweight_and_relAO_per_center=self.centerweight_and_relAO_per_center[I],
+            relAO_per_origin=self.relAO_per_origin[I],
+            unitcell=self.unitcell,
+            unitcell_nkpt=unitcell_nkpt,
+        )
+
 
 def fragmentate(
     mol: Cell,
diff --git a/src/quemb/kbe/pbe.py b/src/quemb/kbe/pbe.py
index 95589d9ff..76acaa458 100644
--- a/src/quemb/kbe/pbe.py
+++ b/src/quemb/kbe/pbe.py
@@ -480,42 +480,9 @@ def initialize(self, compute_hf: bool, restart: bool = False) -> None:
         # Create a file to store ERIs
         if not restart:
             file_eri = h5py.File(self.eri_file, "w")
-        lentmp = len(self.fobj.relAO_per_edge)
         transform_parallel = False  # hard set for now
         for fidx in range(self.fobj.n_frag):
-            if lentmp:
-                fobjs_ = Frags(
-                    self.fobj.AO_per_frag[fidx],
-                    fidx,
-                    edge=self.fobj.AO_per_edge[fidx],
-                    eri_file=self.eri_file,
-                    ref_frag_idx_per_edge=self.fobj.ref_frag_idx_per_edge[fidx],
-                    relAO_per_edge=self.fobj.relAO_per_edge[fidx],
-                    relAO_in_ref_per_edge=self.fobj.relAO_in_ref_per_edge[fidx],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        fidx
-                    ],
-                    relAO_per_origin=self.fobj.relAO_per_origin[fidx],
-                    unitcell=self.fobj.unitcell,
-                    unitcell_nkpt=self.unitcell_nkpt,
-                )
-            else:
-                fobjs_ = Frags(
-                    self.fobj.AO_per_frag[fidx],
-                    fidx,
-                    edge=[],
-                    ref_frag_idx_per_edge=[],
-                    eri_file=self.eri_file,
-                    relAO_per_edge=[],
-                    relAO_in_ref_per_edge=[],
-                    relAO_per_origin=[],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        fidx
-                    ],
-                    unitcell=self.fobj.unitcell,
-                    unitcell_nkpt=self.unitcell_nkpt,
-                )
-
+            fobjs_ = self.fobj.to_Frags(fidx, self.eri_file, self.unitcell_nkpt)
             fobjs_.sd(
                 self.W,
                 self.lmo_coeff,
diff --git a/src/quemb/kbe/pfrag.py b/src/quemb/kbe/pfrag.py
index c898d31b3..e537de27f 100644
--- a/src/quemb/kbe/pfrag.py
+++ b/src/quemb/kbe/pfrag.py
@@ -1,5 +1,6 @@
 # Author(s): Oinam Romesh Meitei
 
+from collections.abc import Sequence
 
 import h5py
 import numpy as np
@@ -8,6 +9,7 @@
     complex128,
     diag_indices,
     einsum,
+    float64,
     outer,
     result_type,
     trace,
@@ -22,6 +24,16 @@
 from quemb.kbe.solver import schmidt_decomp_svd
 from quemb.molbe.helper import get_eri, get_scfObj
 from quemb.shared.helper import unused
+from quemb.shared.typing import (
+    FragmentIdx,
+    GlobalAOIdx,
+    Matrix,
+    PathLike,
+    RelAOIdx,
+    RelAOIdxInRef,
+    SeqOverEdge,
+    Tensor3D,
+)
 
 
 class Frags:
@@ -34,19 +46,19 @@ class Frags:
 
     def __init__(
         self,
-        AO_per_frag,
-        ifrag,
-        edge=None,
-        ref_frag_idx_per_edge=None,
-        relAO_per_edge=None,
-        relAO_in_ref_per_edge=None,
-        centerweight_and_relAO_per_center=None,
-        eri_file="eri_file.h5",
-        unitcell_nkpt=1,
-        ewald_ek=None,
-        relAO_per_origin=None,
-        unitcell=1,
-    ):
+        *,
+        AO_in_frag: Sequence[GlobalAOIdx],
+        ifrag: int,
+        AO_per_edge: SeqOverEdge[Sequence[GlobalAOIdx]],
+        ref_frag_idx_per_edge: SeqOverEdge[FragmentIdx],
+        relAO_per_edge: SeqOverEdge[Sequence[RelAOIdx]],
+        relAO_in_ref_per_edge: SeqOverEdge[Sequence[RelAOIdxInRef]],
+        centerweight_and_relAO_per_center: tuple[float, Sequence[RelAOIdx]],
+        relAO_per_origin: Sequence[RelAOIdx],
+        eri_file: PathLike,
+        unitcell_nkpt: int,
+        unitcell: int,
+    ) -> None:
         """Constructor function for :python:`Frags` class.
 
         Parameters
@@ -57,7 +69,7 @@ def __init__(
             Read more detailed description in :class:`quemb.molbe.autofrag.FragPart`.
         ifrag : int
             fragment index (∈ [0, pbe.n_frag - 1])
-        edge : list, optional
+        AO_per_edge : list, optional
             list of lists of edge site AOs for each atom in the fragment,
             by default None
             Read more detailed description in :class:`quemb.molbe.autofrag.FragPart`.
@@ -82,19 +94,26 @@ def __init__(
             indices of the origin in the fragment, by default None
         """
 
-        self.AO_per_frag = AO_per_frag
+        self.AO_in_frag = AO_in_frag
         self.unitcell = unitcell
         self.unitcell_nkpt = unitcell_nkpt
-        self.n_frag = len(AO_per_frag)
-        self.TA = None
-        self.TA_lo_eo = None
-        self.h1 = None
-        self.ifrag = ifrag
+        self.n_frag = len(AO_in_frag)
         self.dname = "f" + str(ifrag)
-        self.nao = None
-        self.mo_coeffs = None
-        self._mo_coeffs = None
-        self.nsocc = None
+        self.AO_per_edge = AO_per_edge
+        self.ref_frag_idx_per_edge = ref_frag_idx_per_edge
+        self.relAO_per_edge = relAO_per_edge
+        self.relAO_in_ref_per_edge = relAO_in_ref_per_edge
+        self.relAO_per_origin = relAO_per_origin
+        self.centerweight_and_relAO_per_center = centerweight_and_relAO_per_center
+        self.ifrag = ifrag
+
+        self.TA: Matrix[float64]
+        self.TA_lo_eo: Matrix[float64]
+        self.h1 = None
+        self.nao: int
+        self.mo_coeffs: Matrix[float64]
+        self._mo_coeffs: Matrix[float64]
+        self.nsocc: int
         self._mf = None
         self._mc = None
 
@@ -102,13 +121,8 @@ def __init__(
         self.t1 = None
         self.t2 = None
 
-        self.heff = None
-        self.edge = edge
-        self.ref_frag_idx_per_edge = ref_frag_idx_per_edge
-        self.relAO_per_edge = relAO_per_edge
-        self.relAO_in_ref_per_edge = relAO_in_ref_per_edge
-        self.relAO_per_origin = relAO_per_origin
-        self.udim = None
+        self.heff: Matrix[float64]
+        self.udim: int
 
         self._rdm1 = None
         self.rdm1__ = None
@@ -118,17 +132,15 @@ def __init__(
         self.genvs = None
         self.ebe = 0.0
         self.ebe_hf = 0.0
-        self.centerweight_and_relAO_per_center = centerweight_and_relAO_per_center
-        self.ewald_ek = ewald_ek
-        self.fock = None
+        self.fock: Matrix[float64]
         self.veff = None
         self.veff0 = None
         self.dm_init = None
-        self.dm0 = None
+        self.dm0: Matrix[float64]
         self.eri_file = eri_file
         self.pot = None
         self.ebe_hf0 = 0.0
-        self.rdm1_lo_k = None
+        self.rdm1_lo_k: Tensor3D[float64]
 
     def sd(
         self,
@@ -140,7 +152,7 @@ def sd(
         kpts=None,
         kmesh=None,
         h1=None,
-    ):
+    ) -> None:
         """
         Perform Schmidt decomposition for the fragment.
 
@@ -173,7 +185,7 @@ def sd(
         else:
             raise ValueError(f"Imaginary density in Full SD {max_val}")
 
-        Sites = [i + (nlo * 0) for i in self.AO_per_frag]
+        Sites = [i + (nlo * 0) for i in self.AO_in_frag]
 
         TA_R = schmidt_decomp_svd(supcell_rdm, Sites, thr_bath=thr_bath)
         teo = TA_R.shape[-1]
@@ -369,7 +381,7 @@ def update_heff(
             cout = self.udim
 
         if do_chempot:
-            for i, fi in enumerate(self.AO_per_frag):
+            for i, fi in enumerate(self.AO_in_frag):
                 if not any(i in sublist for sublist in self.relAO_per_edge):
                     heff_[i, i] -= u[-1]
 
diff --git a/src/quemb/molbe/autofrag.py b/src/quemb/molbe/autofrag.py
index 69957d2e0..7f8ea5b13 100644
--- a/src/quemb/molbe/autofrag.py
+++ b/src/quemb/molbe/autofrag.py
@@ -15,6 +15,7 @@
 from pyscf.pbc.gto import Cell
 
 from quemb.molbe.helper import are_equal, get_core
+from quemb.molbe.pfrag import Frags
 from quemb.shared.helper import unused
 from quemb.shared.typing import (
     AtomIdx,
@@ -26,6 +27,7 @@
     ListOverMotif,
     MotifIdx,
     OriginIdx,
+    PathLike,
     RelAOIdx,
     RelAOIdxInRef,
     Vector,
@@ -158,6 +160,20 @@ def all_centers_are_origins(self) -> bool:
             )
         )
 
+    def to_Frags(self, I: int, eri_file: PathLike, unrestricted: bool = False) -> Frags:
+        return Frags(
+            self.AO_per_frag[I],
+            I,
+            AO_per_edge=self.AO_per_edge[I],
+            eri_file=eri_file,
+            ref_frag_idx_per_edge=self.ref_frag_idx_per_edge[I],
+            relAO_per_edge=self.relAO_per_edge[I],
+            relAO_in_ref_per_edge=self.relAO_in_ref_per_edge[I],
+            centerweight_and_relAO_per_center=self.centerweight_and_relAO_per_center[I],
+            relAO_per_origin=self.relAO_per_origin[I],
+            unrestricted=unrestricted,
+        )
+
 
 @define(frozen=True, kw_only=True)
 class GraphGenArgs:
@@ -291,13 +307,14 @@ def remove_nonnunique_frags(self, natm: int) -> None:
 
     def to_FragPart(self, mol: Mole, n_BE: int, frozen_core: bool) -> FragPart:
         MISSING = []  # type: ignore[var-annotated]
+        MISSING_PER_FRAG = [[] for _ in range(len(self.AO_per_frag))]  # type: ignore[var-annotated]
         return FragPart(
             mol=mol,
             frag_type="graphgen",
             n_BE=n_BE,
             AO_per_edge=self.AO_per_edge,  # type: ignore[arg-type]
-            relAO_per_edge=MISSING,
-            relAO_in_ref_per_edge=MISSING,
+            relAO_per_edge=MISSING_PER_FRAG,
+            relAO_in_ref_per_edge=MISSING_PER_FRAG,
             relAO_per_origin=self.relAO_per_origin,  # type: ignore[arg-type]
             AO_per_frag=self.AO_per_frag,  # type: ignore[arg-type]
             ref_frag_idx_per_edge=self.ref_frag_idx_per_edge,  # type: ignore[arg-type]
@@ -1017,6 +1034,12 @@ def autogen(
 
             relAO_in_ref_per_edge.append(idx)
 
+    if not AO_per_edge:
+        AO_per_edge = [[] for _ in range(n_frag)]
+        ref_frag_idx_per_edge = [[] for _ in range(n_frag)]
+        relAO_per_edge = [[] for _ in range(n_frag)]
+        relAO_in_ref_per_edge = [[] for _ in range(n_frag)]
+
     return FragPart(
         mol=mol,
         frag_type="autogen",
diff --git a/src/quemb/molbe/be_parallel.py b/src/quemb/molbe/be_parallel.py
index b3a955f25..9c3fd74ed 100644
--- a/src/quemb/molbe/be_parallel.py
+++ b/src/quemb/molbe/be_parallel.py
@@ -465,7 +465,7 @@ def be_func_parallel(
             fobj.update_heff(pot, only_chem=only_chem)
 
     with Pool(nprocs) as pool_:
-        results = []
+        results = []  # type: ignore[var-annotated]
         # Run solver in parallel for each fragment
         for fobj in Fobjs:
             assert (
diff --git a/src/quemb/molbe/helper.py b/src/quemb/molbe/helper.py
index 1e7949db4..8aa097bca 100644
--- a/src/quemb/molbe/helper.py
+++ b/src/quemb/molbe/helper.py
@@ -22,6 +22,7 @@
 from pyscf.pbc.gto.cell import Cell
 
 from quemb.shared.helper import ncore_
+from quemb.shared.typing import Matrix
 
 
 def get_veff(eri_, dm, S, TA, hf_veff):
@@ -70,11 +71,11 @@ def get_veff(eri_, dm, S, TA, hf_veff):
 
 # create pyscf pbc scf object
 def get_scfObj(
-    h1,
+    h1: Matrix[float64],
     Eri,
-    nocc,
+    nocc: int,
     dm0=None,
-):
+) -> scf.hf.RHF:
     """Initialize and run a restricted Hartree-Fock (RHF) calculation.
 
     This function sets up an SCF (Self-Consistent Field) object using the provided
diff --git a/src/quemb/molbe/mbe.py b/src/quemb/molbe/mbe.py
index a3e13164a..0d3d3189d 100644
--- a/src/quemb/molbe/mbe.py
+++ b/src/quemb/molbe/mbe.py
@@ -369,7 +369,7 @@ def rdm1_fullbasis(
 
             # Generate the projection matrix
             cind = [
-                fobjs.AO_per_frag[i] for i in fobjs.centerweight_and_relAO_per_center[1]
+                fobjs.AO_in_frag[i] for i in fobjs.centerweight_and_relAO_per_center[1]
             ]
             Pc_ = (
                 fobjs.TA.T
@@ -829,36 +829,8 @@ def initialize(
         # Create a file to store ERIs
         if not restart:
             file_eri = h5py.File(self.eri_file, "w")
-        lentmp = len(self.fobj.relAO_per_edge)
         for I in range(self.fobj.n_frag):
-            if lentmp:
-                fobjs_ = Frags(
-                    self.fobj.AO_per_frag[I],
-                    I,
-                    AO_per_edge=self.fobj.AO_per_edge[I],
-                    eri_file=self.eri_file,
-                    ref_frag_idx_per_edge=self.fobj.ref_frag_idx_per_edge[I],
-                    relAO_per_edge=self.fobj.relAO_per_edge[I],
-                    relAO_in_ref_per_edge=self.fobj.relAO_in_ref_per_edge[I],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        I
-                    ],
-                    relAO_per_origin=self.fobj.relAO_per_origin[I],
-                )
-            else:
-                fobjs_ = Frags(
-                    self.fobj.AO_per_frag[I],
-                    I,
-                    AO_per_edge=[],
-                    ref_frag_idx_per_edge=[],
-                    eri_file=self.eri_file,
-                    relAO_per_edge=[],
-                    relAO_in_ref_per_edge=[],
-                    relAO_per_origin=[],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        I
-                    ],
-                )
+            fobjs_ = self.fobj.to_Frags(I, eri_file=self.eri_file)
             fobjs_.sd(self.W, self.lmo_coeff, self.Nocc, thr_bath=self.thr_bath)
 
             self.Fobjs.append(fobjs_)
diff --git a/src/quemb/molbe/pfrag.py b/src/quemb/molbe/pfrag.py
index 9cfeb2e70..555a5246a 100644
--- a/src/quemb/molbe/pfrag.py
+++ b/src/quemb/molbe/pfrag.py
@@ -1,5 +1,7 @@
 # Author(s): Oinam Romesh Meitei
 
+from collections.abc import Sequence
+
 import h5py
 import numpy as np
 import scipy.linalg
@@ -19,7 +21,15 @@
 from numpy.linalg import eigh, multi_dot
 
 from quemb.molbe.helper import get_eri, get_scfObj, get_veff
-from quemb.shared.typing import Matrix
+from quemb.shared.typing import (
+    FragmentIdx,
+    GlobalAOIdx,
+    Matrix,
+    PathLike,
+    RelAOIdx,
+    RelAOIdxInRef,
+    SeqOverEdge,
+)
 
 
 class Frags:
@@ -32,17 +42,17 @@ class Frags:
 
     def __init__(
         self,
-        AO_per_frag,
+        AO_in_frag: Sequence[GlobalAOIdx],
         ifrag: int,
-        AO_per_edge=None,
-        ref_frag_idx_per_edge=None,
-        relAO_per_edge=None,
-        relAO_in_ref_per_edge=None,
-        centerweight_and_relAO_per_center=None,
-        eri_file="eri_file.h5",
-        relAO_per_origin=None,
-        unrestricted=False,
-    ):
+        AO_per_edge: SeqOverEdge[Sequence[GlobalAOIdx]],
+        ref_frag_idx_per_edge: SeqOverEdge[FragmentIdx],
+        relAO_per_edge: SeqOverEdge[Sequence[RelAOIdx]],
+        relAO_in_ref_per_edge: SeqOverEdge[Sequence[RelAOIdxInRef]],
+        centerweight_and_relAO_per_center: tuple[float, Sequence[RelAOIdx]],
+        relAO_per_origin: Sequence[RelAOIdx],
+        eri_file: PathLike = "eri_file.h5",
+        unrestricted: bool = False,
+    ) -> None:
         """Constructor function for :python:`Frags` class.
 
         Parameters
@@ -74,11 +84,16 @@ def __init__(
             unrestricted calculation, by default False
         """
 
-        self.AO_per_frag = AO_per_frag
-        self.n_frag = len(AO_per_frag)
-        self.TA: Matrix[float64] | None = None
-        self.TA_lo_eo: Matrix[float64] | None = None
-        self.h1 = None
+        self.AO_in_frag = AO_in_frag
+        self.n_frag = len(AO_in_frag)
+        self.AO_per_edge = AO_per_edge
+        self.ref_frag_idx_per_edge = ref_frag_idx_per_edge
+        self.relAO_per_edge = relAO_per_edge
+        self.relAO_in_ref_per_edge = relAO_in_ref_per_edge
+        self.relAO_per_origin = relAO_per_origin
+        self.centerweight_and_relAO_per_center = centerweight_and_relAO_per_center
+        self.eri_file = eri_file
+
         self.ifrag = ifrag
         if unrestricted:
             self.dname: str | list[str] = [
@@ -88,10 +103,15 @@ def __init__(
             ]
         else:
             self.dname = "f" + str(ifrag)
-        self.nao = None
-        self.mo_coeffs = None
-        self._mo_coeffs = None
-        self.nsocc = None
+
+        self.TA: Matrix[float64]
+        self.TA_lo_eo: Matrix[float64]
+
+        self.h1: Matrix[float64]
+        self.nao: int
+        self.mo_coeffs: Matrix[float64]
+        self._mo_coeffs: Matrix[float64]
+        self.nsocc: int
         self._mf = None
         self._mc = None
 
@@ -99,12 +119,7 @@ def __init__(
         self.t1 = None
         self.t2 = None
 
-        self.heff: Matrix[float64] | None = None
-        self.AO_per_edge = AO_per_edge
-        self.ref_frag_idx_per_edge = ref_frag_idx_per_edge
-        self.relAO_per_edge = relAO_per_edge
-        self.relAO_in_ref_per_edge = relAO_in_ref_per_edge
-        self.relAO_per_origin = relAO_per_origin
+        self.heff: Matrix[float64]
         self.udim: int | None = None
 
         self._rdm1 = None
@@ -114,16 +129,21 @@ def __init__(
         self.genvs = None
         self.ebe = 0.0
         self.ebe_hf = 0.0
-        self.centerweight_and_relAO_per_center = centerweight_and_relAO_per_center
         self.fock = None
         self.veff = None
         self.veff0 = None
         self.dm_init = None
-        self.dm0 = None
-        self.eri_file = eri_file
+        self.dm0: Matrix[float64]
         self.unitcell_nkpt = 1.0
 
-    def sd(self, lao, lmo, nocc, thr_bath, norb=None, return_orb_count=False):
+    def sd(
+        self,
+        lao: Matrix[float64],
+        lmo: Matrix[float64],
+        nocc: int,
+        thr_bath: float,
+        norb: int | None = None,
+    ) -> None:
         """
         Perform Schmidt decomposition for the fragment.
 
@@ -141,28 +161,16 @@ def sd(self, lao, lmo, nocc, thr_bath, norb=None, return_orb_count=False):
             Specify number of bath orbitals.
             Used for UBE, where different number of alpha and beta orbitals
             Default is None, allowing orbitals to be chosen by threshold
-        return_orb_count : bool, optional
-            Retrun the number of orbitals in each space, for UBE use/
-            Default is False
         """
-
-        if return_orb_count:
-            TA, n_f, n_b = schmidt_decomposition(
-                lmo,
-                nocc,
-                self.AO_per_frag,
-                thr_bath=thr_bath,
-                norb=norb,
-                return_orb_count=return_orb_count,
-            )
-        else:
-            TA = schmidt_decomposition(lmo, nocc, self.AO_per_frag, thr_bath=thr_bath)
-        self.C_lo_eo = TA
-        TA = lao @ TA
-        self.nao = TA.shape[1]
-        self.TA = TA
-        if return_orb_count:
-            return [n_f, n_b]
+        self.TA_lo_eo, self.n_f, self.n_b = schmidt_decomposition(
+            lmo,
+            nocc,
+            self.AO_in_frag,
+            thr_bath=thr_bath,
+            norb=norb,
+        )
+        self.TA = lao @ self.TA_lo_eo
+        self.nao = self.TA.shape[1]
 
     def cons_fock(self, hf_veff, S, dm, eri_=None):
         """
@@ -279,7 +287,7 @@ def update_heff(self, u, cout=None, only_chem=False):
         if cout is None:
             cout = self.udim
 
-        for i, fi in enumerate(self.AO_per_frag):
+        for i, fi in enumerate(self.AO_in_frag):
             if not any(i in sublist for sublist in self.relAO_per_edge):
                 heff_[i, i] -= u[-1]
 
@@ -386,15 +394,14 @@ def update_ebe_hf(
 
 
 def schmidt_decomposition(
-    mo_coeff,
-    nocc,
-    Frag_sites,
-    thr_bath=1.0e-10,
-    cinv=None,
-    rdm=None,
-    norb=None,
-    return_orb_count=False,
-):
+    mo_coeff: Matrix[float64],
+    nocc: int,
+    AO_in_frag: Sequence[GlobalAOIdx],
+    thr_bath: float = 1.0e-10,
+    cinv: Matrix[float64] | None = None,
+    rdm: Matrix[float64] | None = None,
+    norb: int | None = None,
+) -> tuple[Matrix[float64], int, int]:
     """
     Perform Schmidt decomposition on the molecular orbital coefficients.
 
@@ -404,34 +411,29 @@ def schmidt_decomposition(
 
     Parameters
     ----------
-    mo_coeff : numpy.ndarray
+    mo_coeff :
         Molecular orbital coefficients.
-    nocc : int
+    nocc :
         Number of occupied orbitals.
     Frag_sites : list of int
         List of fragment sites (indices).
-    thr_bath : float,
+    thr_bath :
         Threshold for bath orbitals in Schmidt decomposition
-    cinv : numpy.ndarray, optional
+    cinv :
         Inverse of the transformation matrix. Defaults to None.
-    rdm : numpy.ndarray, optional
+    rdm :
         Reduced density matrix. If not provided, it will be computed from the molecular
         orbitals. Defaults to None.
-    norb : int, optional
+    norb :
         Specifies number of bath orbitals. Used for UBE to make alpha and beta
         spaces the same size. Defaults to None
-    return_orb_count : bool, optional
-        Return more information about the number of orbitals. Used in UBE.
-        Defaults to False
 
     Returns
     -------
-    numpy.ndarray
+    tuple:
+        TA, norbs_frag, norbs_bath
+
         Transformation matrix (TA) including both fragment and entangled bath orbitals.
-    if return_orb_count:
-        numpy.ndarray, int, int
-        returns TA (above), number of orbitals in the fragment space, and number of
-        orbitals in bath space
     """
 
     # Compute the reduced density matrix (RDM) if not provided
@@ -448,9 +450,9 @@ def schmidt_decomposition(
     Tot_sites = Dhf.shape[0]
 
     # Identify environment sites (indices not in Frag_sites)
-    Env_sites1 = array([i for i in range(Tot_sites) if i not in Frag_sites])
-    Env_sites = array([[i] for i in range(Tot_sites) if i not in Frag_sites])
-    Frag_sites1 = array([[i] for i in Frag_sites])
+    Env_sites1 = array([i for i in range(Tot_sites) if i not in AO_in_frag])
+    Env_sites = array([[i] for i in range(Tot_sites) if i not in AO_in_frag])
+    Frag_sites1 = array([[i] for i in AO_in_frag])
 
     # Compute the environment part of the density matrix
     Denv = Dhf[Env_sites, Env_sites.T]
@@ -477,12 +479,9 @@ def schmidt_decomposition(
                 Bidx.append(i)
 
     # Initialize the transformation matrix (TA)
-    TA = zeros([Tot_sites, len(Frag_sites) + len(Bidx)])
-    TA[Frag_sites, : len(Frag_sites)] = eye(len(Frag_sites))  # Fragment part
-    TA[Env_sites1, len(Frag_sites) :] = Evec[:, Bidx]  # Environment part
+    TA = zeros([Tot_sites, len(AO_in_frag) + len(Bidx)])
+    TA[AO_in_frag, : len(AO_in_frag)] = eye(len(AO_in_frag))  # Fragment part
+    TA[Env_sites1, len(AO_in_frag) :] = Evec[:, Bidx]  # Environment part
 
-    if return_orb_count:
-        # return TA, norbs_frag, norbs_bath
-        return TA, Frag_sites1.shape[0], len(Bidx)
-    else:
-        return TA
+    # return TA, norbs_frag, norbs_bath
+    return TA, Frag_sites1.shape[0], len(Bidx)
diff --git a/src/quemb/molbe/ube.py b/src/quemb/molbe/ube.py
index b969955cd..149caf082 100644
--- a/src/quemb/molbe/ube.py
+++ b/src/quemb/molbe/ube.py
@@ -173,79 +173,18 @@ def initialize(self, eri_, compute_hf):
         ECOUL = 0.0
 
         file_eri = h5py.File(self.eri_file, "w")
-        lentmp = len(self.fobj.relAO_per_edge)
-
         # alpha orbitals
-        for I in range(self.fobj.n_frag):
-            if lentmp:
-                fobjs_a = Frags(
-                    self.fobj.AO_per_frag[I],
-                    I,
-                    AO_per_edge=self.fobj.AO_per_edge[I],
-                    eri_file=self.eri_file,
-                    ref_frag_idx_per_edge=self.fobj.ref_frag_idx_per_edge[I],
-                    relAO_per_edge=self.fobj.relAO_per_edge[I],
-                    relAO_in_ref_per_edge=self.fobj.relAO_in_ref_per_edge[I],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        I
-                    ],
-                    relAO_per_origin=self.fobj.relAO_per_origin[I],
-                    unrestricted=True,
-                )
-            else:
-                fobjs_a = Frags(
-                    self.fobj.AO_per_frag[I],
-                    I,
-                    AO_per_edge=[],
-                    ref_frag_idx_per_edge=[],
-                    eri_file=self.eri_file,
-                    relAO_per_edge=[],
-                    relAO_in_ref_per_edge=[],
-                    relAO_per_origin=[],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        I
-                    ],
-                    unrestricted=True,
-                )
-            self.Fobjs_a.append(fobjs_a)
+        self.Fobjs_a = [
+            self.fobj.to_Frags(I, eri_file=self.eri_file, unrestricted=True)
+            for I in range(self.fobj.n_frag)
+        ]
         # beta
-        for I in range(self.fobj.n_frag):
-            if lentmp:
-                fobjs_b = Frags(
-                    self.fobj.AO_per_frag[I],
-                    I,
-                    AO_per_edge=self.fobj.AO_per_edge[I],
-                    eri_file=self.eri_file,
-                    ref_frag_idx_per_edge=self.fobj.ref_frag_idx_per_edge[I],
-                    relAO_per_edge=self.fobj.relAO_per_edge[I],
-                    relAO_in_ref_per_edge=self.fobj.relAO_in_ref_per_edge[I],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        I
-                    ],
-                    relAO_per_origin=self.fobj.relAO_per_origin[I],
-                    unrestricted=True,
-                )
-            else:
-                fobjs_b = Frags(
-                    self.fobj.AO_per_frag[I],
-                    I,
-                    AO_per_edge=[],
-                    ref_frag_idx_per_edge=[],
-                    eri_file=self.eri_file,
-                    relAO_per_edge=[],
-                    relAO_in_ref_per_edge=[],
-                    relAO_per_origin=[],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        I
-                    ],
-                    unrestricted=True,
-                )
-            self.Fobjs_b.append(fobjs_b)
+        self.Fobjs_b = [
+            self.fobj.to_Frags(I, eri_file=self.eri_file, unrestricted=True)
+            for I in range(self.fobj.n_frag)
+        ]
 
-        orb_count_a = []
-        orb_count_b = []
-
-        all_noccs = []
+        all_noccs = [self.Nocc for _ in range(self.fobj.n_frag)]
 
         for I in range(self.fobj.n_frag):
             fobj_a = self.Fobjs_a[I]
@@ -254,48 +193,34 @@ def initialize(self, eri_, compute_hf):
             if self.frozen_core:
                 fobj_a.core_veff = self.core_veff[0]
                 fobj_b.core_veff = self.core_veff[1]
-                orb_count_a.append(
-                    fobj_a.sd(
-                        self.W[0],
-                        self.lmo_coeff_a,
-                        self.Nocc[0],
-                        return_orb_count=True,
-                        thr_bath=self.thr_bath,
-                    )
+                fobj_a.sd(
+                    self.W[0],
+                    self.lmo_coeff_a,
+                    self.Nocc[0],
+                    thr_bath=self.thr_bath,
                 )
-                orb_count_b.append(
-                    fobj_b.sd(
-                        self.W[1],
-                        self.lmo_coeff_b,
-                        self.Nocc[1],
-                        return_orb_count=True,
-                        thr_bath=self.thr_bath,
-                    )
+                fobj_b.sd(
+                    self.W[1],
+                    self.lmo_coeff_b,
+                    self.Nocc[1],
+                    thr_bath=self.thr_bath,
                 )
             else:
                 fobj_a.core_veff = None
                 fobj_b.core_veff = None
-                orb_count_a.append(
-                    fobj_a.sd(
-                        self.W,
-                        self.lmo_coeff_a,
-                        self.Nocc[0],
-                        return_orb_count=True,
-                        thr_bath=self.thr_bath,
-                    )
+                fobj_a.sd(
+                    self.W,
+                    self.lmo_coeff_a,
+                    self.Nocc[0],
+                    thr_bath=self.thr_bath,
                 )
-                orb_count_b.append(
-                    fobj_b.sd(
-                        self.W,
-                        self.lmo_coeff_b,
-                        self.Nocc[1],
-                        return_orb_count=True,
-                        thr_bath=self.thr_bath,
-                    )
+                fobj_b.sd(
+                    self.W,
+                    self.lmo_coeff_b,
+                    self.Nocc[1],
+                    thr_bath=self.thr_bath,
                 )
 
-            all_noccs.append(self.Nocc)
-
             if eri_ is None and self.mf.with_df is not None:
                 # NOT IMPLEMENTED: should not be called, as no unrestricted DF tested
                 # for density-fitted integrals; if mf is provided, pyscf.ao2mo uses DF
@@ -369,6 +294,10 @@ def initialize(self, eri_, compute_hf):
                 EH1 += eh1_b
                 ECOUL += ecoul_b
                 E_hf += fobj_b.ebe_hf
+
+        orb_count_a = [(frag.n_f, frag.n_b) for frag in self.Fobjs_a]
+        orb_count_b = [(frag.n_f, frag.n_b) for frag in self.Fobjs_b]
+
         file_eri.close()
 
         print("Number of Orbitals per Fragment:", flush=True)
diff --git a/src/quemb/shared/external/optqn.py b/src/quemb/shared/external/optqn.py
index bab8b9a08..5fbf0441f 100644
--- a/src/quemb/shared/external/optqn.py
+++ b/src/quemb/shared/external/optqn.py
@@ -5,6 +5,7 @@
 #   (except for the trust region routine)
 #         The code has been slightly modified.
 #
+from collections.abc import Sequence
 
 from numpy import array, empty, float64, outer, zeros
 from numpy.linalg import inv, norm, pinv
@@ -16,7 +17,7 @@
 from quemb.shared.external.cphf_utils import cphf_kernel_batch, get_rhf_dP_from_u
 from quemb.shared.external.cpmp2_utils import get_dPmp2_batch_r
 from quemb.shared.external.jac_utils import get_dPccsdurlx_batch_u
-from quemb.shared.typing import Matrix
+from quemb.shared.typing import GlobalAOIdx, Matrix, RelAOIdx, SeqOverEdge
 
 
 def line_search_LF(func, xold, fold, dx, iter_):
@@ -326,7 +327,7 @@ def get_atbe_Jblock_frag(
         and fobj.heff is not None
         and fobj.nao is not None
     )
-    vpots = get_vpots_frag(fobj.nao, fobj.relAO_per_edge, fobj.AO_per_frag)
+    vpots = get_vpots_frag(fobj.nao, fobj.relAO_per_edge, fobj.AO_in_frag)
     eri_ = get_eri(fobj.dname, fobj.nao, eri_file=fobj.eri_file)
     dm0 = 2.0 * (fobj._mo_coeffs[:, : fobj.nsocc] @ fobj._mo_coeffs[:, : fobj.nsocc].T)
     mf_ = get_scfObj(fobj.fock + fobj.heff, eri_, fobj.nsocc, dm0=dm0)
@@ -359,7 +360,7 @@ def get_atbe_Jblock_frag(
 
                             tmpje_.append(dPs[cout][edge_[j__], edge_[k__]])
                 y_ = 0.0
-                for fidx, fval in enumerate(fobj.AO_per_frag):
+                for fidx, fval in enumerate(fobj.AO_in_frag):
                     if not any(fidx in sublist for sublist in fobj.relAO_per_edge):
                         y_ += dPs[cout][fidx, fidx]
 
@@ -368,11 +369,11 @@ def get_atbe_Jblock_frag(
                 tmpjc_ = []
                 # center on the same fragment
                 # for cen in fobj.efac[1]:
-                for j__ in fobj.relAO_per_origin:
-                    for k__ in fobj.relAO_per_origin:
-                        if j__ > k__:
+                for j_relAO in fobj.relAO_per_origin:
+                    for k_relAO in fobj.relAO_per_origin:
+                        if j_relAO > k_relAO:
                             continue
-                        tmpjc_.append(-dPs[cout][j__, k__])
+                        tmpjc_.append(-dPs[cout][j_relAO, k_relAO])
 
                 Je.append(tmpje_)
 
@@ -384,15 +385,15 @@ def get_atbe_Jblock_frag(
                 cout += 1
 
     alpha = 0.0
-    for fidx, _ in enumerate(fobj.AO_per_frag):
+    for fidx, _ in enumerate(fobj.AO_in_frag):
         if not any(fidx in sublist for sublist in fobj.relAO_per_edge):
             alpha += dP_mu[fidx, fidx]
 
-    for j__ in fobj.relAO_per_origin:
-        for k__ in fobj.relAO_per_origin:
-            if j__ > k__:
+    for j_relAO in fobj.relAO_per_origin:
+        for k_relAO in fobj.relAO_per_origin:
+            if j_relAO > k_relAO:
                 continue
-            xc.append(-dP_mu[j__, k__])
+            xc.append(-dP_mu[j_relAO, k_relAO])
 
     return array(Je).T, array(Jc).T, xe, xc, y, alpha, cout
 
@@ -467,11 +468,13 @@ def ccsdres_func(mf, vpots, eri, nsocc):
 
 
 def get_vpots_frag(
-    nao: int, rel_AO_per_edge_per_frag: list[list[int]], AO_per_frag: list[list[int]]
+    nao: int,
+    rel_AO_per_edge: SeqOverEdge[Sequence[RelAOIdx]],
+    AO_in_frag: Sequence[GlobalAOIdx],
 ) -> list[Matrix[float64]]:
     vpots: list[Matrix[float64]] = []
 
-    for edge_ in rel_AO_per_edge_per_frag:
+    for edge_ in rel_AO_per_edge:
         lene = len(edge_)
         for j__ in range(lene):
             for k__ in range(lene):
@@ -485,8 +488,8 @@ def get_vpots_frag(
     # only the centers
     # outer edges not included
     tmppot = zeros((nao, nao))
-    for fidx, fval in enumerate(AO_per_frag):
-        if not any(fidx in sublist for sublist in rel_AO_per_edge_per_frag):
+    for fidx, fval in enumerate(AO_in_frag):
+        if not any(fidx in sublist for sublist in rel_AO_per_edge):
             tmppot[fidx, fidx] = -1
 
     vpots.append(tmppot)
diff --git a/src/quemb/shared/typing.py b/src/quemb/shared/typing.py
index 2dc5e23af..7cdb16cb2 100644
--- a/src/quemb/shared/typing.py
+++ b/src/quemb/shared/typing.py
@@ -10,6 +10,7 @@
 """
 
 import os
+from collections.abc import Sequence
 from typing import Any, NewType, TypeAlias, TypeVar
 
 import numpy as np
@@ -141,3 +142,9 @@
 ListOverEdge: TypeAlias = list
 ListOverCenter: TypeAlias = list
 ListOverMotif: TypeAlias = list
+
+
+SeqOverFrag: TypeAlias = Sequence
+SeqOverEdge: TypeAlias = Sequence
+SeqOverCenter: TypeAlias = Sequence
+SeqOverMotif: TypeAlias = Sequence
diff --git a/tests/fragmentation_test.py b/tests/fragmentation_test.py
index 9f1bb95cb..70e370554 100644
--- a/tests/fragmentation_test.py
+++ b/tests/fragmentation_test.py
@@ -26,8 +26,8 @@ def test_autogen_h_linear_be1(self):
 
         target = {
             "AO_per_frag": [[0], [1], [2], [3], [4], [5], [6], [7]],
-            "AO_per_edge": [],
-            "ref_frag_idx_per_edge": [],
+            "AO_per_edge": [[], [], [], [], [], [], [], []],
+            "ref_frag_idx_per_edge": [[], [], [], [], [], [], [], []],
             "relAO_per_origin": [[0], [0], [0], [0], [0], [0], [0], [0]],
             "centerweight_and_relAO_per_center": [
                 (1.0, [0]),
@@ -158,8 +158,8 @@ def test_autogen_octane_be1(self):
                 [42, 43, 44, 45, 46, 52, 54, 57],
                 [47, 48, 49, 50, 51, 53, 55, 56],
             ],
-            "AO_per_edge": [],
-            "ref_frag_idx_per_edge": [],
+            "AO_per_edge": [[], [], [], [], [], [], [], []],
+            "ref_frag_idx_per_edge": [[], [], [], [], [], [], [], []],
             "relAO_per_origin": [
                 [0, 1, 2, 3, 4, 5, 6],
                 [0, 1, 2, 3, 4, 5, 6],


Reviews:
user: login: mcocdawc
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-23T16:30:54Z
user: login: mcocdawc
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T13:20:23Z
user: login: mcocdawc
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T13:25:24Z
user: login: mcocdawc
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-24T14:07:38Z
user: login: lweisburn
state: APPROVED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T19:08:31Z
user: login: mcocdawc
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T20:23:19Z
user: login: mcocdawc
state: COMMENTED
author_association: CONTRIBUTOR
submitted_at: 2025-04-25T20:24:42Z

Review comments:
diff_hunk: @@ -174,8 +186,9 @@ def sd(
         else:
             raise ValueError(f"Imaginary density in Full SD {max_val}")
 
-        Sites = [i + (nlo * 0) for i in self.AO_per_frag]
-        if not frag_type == "autogen":
+        Sites = [i + (nlo * 0) for i in self.AO_in_frag]
+        # TODO: Discuss with Minsik
path: src/quemb/kbe/pfrag.py
user: login: mcocdawc
body: the following sort might be relevant for #144 @mscho527 
created_at: 2025-04-23T16:30:54Z
author_association: CONTRIBUTOR
diff_hunk: @@ -141,28 +161,16 @@ def sd(self, lao, lmo, nocc, thr_bath, norb=None, return_orb_count=False):
             Specify number of bath orbitals.
             Used for UBE, where different number of alpha and beta orbitals
             Default is None, allowing orbitals to be chosen by threshold
-        return_orb_count : bool, optional
-            Retrun the number of orbitals in each space, for UBE use/
-            Default is False
         """
-
-        if return_orb_count:
-            TA, n_f, n_b = schmidt_decomposition(
-                lmo,
-                nocc,
-                self.AO_per_frag,
-                thr_bath=thr_bath,
-                norb=norb,
-                return_orb_count=return_orb_count,
-            )
-        else:
-            TA = schmidt_decomposition(lmo, nocc, self.AO_per_frag, thr_bath=thr_bath)
-        self.C_lo_eo = TA
path: src/quemb/molbe/pfrag.py
user: login: mcocdawc
body: I assume this is a typo and they meant `TA_lo_eo` instead of `C_lo_eo`
created_at: 2025-04-24T13:20:23Z
author_association: CONTRIBUTOR
diff_hunk: @@ -174,8 +186,9 @@ def sd(
         else:
             raise ValueError(f"Imaginary density in Full SD {max_val}")
 
-        Sites = [i + (nlo * 0) for i in self.AO_per_frag]
-        if not frag_type == "autogen":
+        Sites = [i + (nlo * 0) for i in self.AO_in_frag]
+        # TODO: Discuss with Minsik
path: src/quemb/kbe/pfrag.py
user: login: mcocdawc
body: It was already fixed 
https://github.com/troyvvgroup/quemb/pull/144/files#diff-059f900a6633877c395553bbd7e267247aef938a4d8f0b50f4494690a85178aeL172
:-D 
created_at: 2025-04-24T13:25:24Z
author_association: CONTRIBUTOR
diff_hunk: @@ -290,13 +306,14 @@ def remove_nonnunique_frags(self, natm: int) -> None:
 
     def to_FragPart(self, mol: Mole, n_BE: int, frozen_core: bool) -> FragPart:
         MISSING = []  # type: ignore[var-annotated]
+        MISSING_PER_FRAG = [[] for _ in range(len(self.AO_per_frag))]  # type: ignore[var-annotated]
path: src/quemb/molbe/autofrag.py
user: login: mcocdawc
body: @ShaunWeatherly 
Previously, when there were no edges, `relAO_per_edge` was just a flat empty list rather than a list of empty lists per fragment (also in autogen). This led to a lot of special casing in downstream code.

I changed the output from an empty list to an empty list of empty lists also for `graphgen`.
created_at: 2025-04-24T14:07:38Z
author_association: CONTRIBUTOR
diff_hunk: @@ -480,49 +480,15 @@ def initialize(self, compute_hf: bool, restart: bool = False) -> None:
         # Create a file to store ERIs
         if not restart:
             file_eri = h5py.File(self.eri_file, "w")
-        lentmp = len(self.fobj.relAO_per_edge)
         transform_parallel = False  # hard set for now
         for fidx in range(self.fobj.n_frag):
-            if lentmp:
-                fobjs_ = Frags(
-                    self.fobj.AO_per_frag[fidx],
-                    fidx,
-                    edge=self.fobj.AO_per_edge[fidx],
-                    eri_file=self.eri_file,
-                    ref_frag_idx_per_edge=self.fobj.ref_frag_idx_per_edge[fidx],
-                    relAO_per_edge=self.fobj.relAO_per_edge[fidx],
-                    relAO_in_ref_per_edge=self.fobj.relAO_in_ref_per_edge[fidx],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        fidx
-                    ],
-                    relAO_per_origin=self.fobj.relAO_per_origin[fidx],
-                    unitcell=self.fobj.unitcell,
-                    unitcell_nkpt=self.unitcell_nkpt,
-                )
-            else:
-                fobjs_ = Frags(
-                    self.fobj.AO_per_frag[fidx],
-                    fidx,
-                    edge=[],
-                    ref_frag_idx_per_edge=[],
-                    eri_file=self.eri_file,
-                    relAO_per_edge=[],
-                    relAO_in_ref_per_edge=[],
-                    relAO_per_origin=[],
-                    centerweight_and_relAO_per_center=self.fobj.centerweight_and_relAO_per_center[
-                        fidx
-                    ],
-                    unitcell=self.fobj.unitcell,
-                    unitcell_nkpt=self.unitcell_nkpt,
-                )
-
+            fobjs_ = self.fobj.to_Frags(fidx, self.eri_file, self.unitcell_nkpt)
path: src/quemb/kbe/pbe.py
user: login: lweisburn
body: happy to see this all moved to a more concise function!
created_at: 2025-04-25T18:44:14Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,19 +46,19 @@ class Frags:
 
     def __init__(
         self,
-        AO_per_frag,
-        ifrag,
-        edge=None,
-        ref_frag_idx_per_edge=None,
-        relAO_per_edge=None,
-        relAO_in_ref_per_edge=None,
-        centerweight_and_relAO_per_center=None,
-        eri_file="eri_file.h5",
-        unitcell_nkpt=1,
-        ewald_ek=None,
-        relAO_per_origin=None,
-        unitcell=1,
-    ):
+        *,
+        AO_in_frag: Sequence[GlobalAOIdx],
path: src/quemb/kbe/pfrag.py
user: login: lweisburn
body: In the documentation, this is still called `AO_per_frag`!
created_at: 2025-04-25T18:59:15Z
author_association: CONTRIBUTOR
diff_hunk: @@ -118,6 +120,21 @@ def all_centers_are_origins(self) -> bool:
             )
         )
 
+    def to_Frags(self, I: int, eri_file: PathLike, unitcell_nkpt: int) -> Frags:
+        return Frags(
+            AO_in_frag=self.AO_per_frag[I],
path: src/quemb/kbe/fragment.py
user: login: lweisburn
body: Just to clarify, there are both `AO_in_frag` and `AO_per_frag` in different codes. In the molecular and most of the periodic code, the previous uses of `AO_per_frag` (and also the `Frag_sites` in the Schmidt decompositions, though not autogen) are becoming `AO_in_frag`. 

Should this be moved to `AO_in_frag` as well? I haven't looked yet in detail at the differences between `_in_frag` and `_per_frag`, but I am a little confused about this distinction. Are they the same thing, just passed from the fragment codes into the later embedding stuff?
created_at: 2025-04-25T19:07:55Z
author_association: CONTRIBUTOR
diff_hunk: @@ -34,19 +46,19 @@ class Frags:
 
     def __init__(
         self,
-        AO_per_frag,
-        ifrag,
-        edge=None,
-        ref_frag_idx_per_edge=None,
-        relAO_per_edge=None,
-        relAO_in_ref_per_edge=None,
-        centerweight_and_relAO_per_center=None,
-        eri_file="eri_file.h5",
-        unitcell_nkpt=1,
-        ewald_ek=None,
-        relAO_per_origin=None,
-        unitcell=1,
-    ):
+        *,
+        AO_in_frag: Sequence[GlobalAOIdx],
path: src/quemb/kbe/pfrag.py
user: login: mcocdawc
body: There were even a couple more discrepancies. Thank you for spotting and fixed.
created_at: 2025-04-25T20:23:19Z
author_association: CONTRIBUTOR
diff_hunk: @@ -118,6 +120,21 @@ def all_centers_are_origins(self) -> bool:
             )
         )
 
+    def to_Frags(self, I: int, eri_file: PathLike, unitcell_nkpt: int) -> Frags:
+        return Frags(
+            AO_in_frag=self.AO_per_frag[I],
path: src/quemb/kbe/fragment.py
user: login: mcocdawc
body: Already clarified in person.

> AO_in_frag are the AO indices in a particular fragment. (pfrags.Frags) It is of type Sequence[GlobalAOIdx]

> AO_per_frag are the AO indices per fragments, basically a list of AO_in_frag. (FragPart) It is of type Sequence[Sequence[GlobalAOIdx]]
created_at: 2025-04-25T20:24:42Z
author_association: CONTRIBUTOR

PR comments:
user: login: lweisburn
created_at: 2025-04-25T19:12:04Z
body: >  Does it break UBE? I cannot test that myself @lweisburn . :-( 

The UBE tests show that this PR should be okay, but also that the newer versions of PySCF/numpy have issues when running the unrestricted calculations. These will have to be addressed in a later PR (raising issue #157 )


author_association: CONTRIBUTOR