Graph theoretic fragmentation via graphgen. (#86)

* Added `graphgen` fragmentation.

* Fix reference to `mf.mol` when `lo_method="pipek"`

* Fix reference to `frag_scratch` when `solver=="DMRG"`

* Ruff fixes.

* Formatting and removed unused code.

* Final formatting

* Update dependencies.

* `mypy` static typing.

* Suppress `mypy` for `networkx` imports.

* Remove debug prints

* Added complete docstring for `graphgen`.

* Add types package for `networkx`

* Unsuppress mypy warnings

* `norm`

* New `FragmentMap` data class.

* Edits to `FragmentMap` typing.

* Remove unused kwargs.

* Fix formatting

* Final formatting

* Additions to `nitpick_exceptions`.

* `FragmentMap` Docstring edits.

* Remove unfinished code.

* Use `np.floating[Any]`

* Organize imports.

* Fixes for `FragmentMap`

* Add unit tests for `autogen` and `graphgen`

* Long line fix.

* Suppress E501 in `fragmentation_test`

* Add ruff exclusion rule for `fragmentation_tests`

* Ruff formatting yet again.

* Add `fragmentation_test` to mypy blacklist

* Remove defaults in `FragmentMap` init.

* Add unit tests for energy comparisons across `autogen` and `graphgen`

* Add checks for IAOs in `graphgen`

* Update `graphgen` docstring

* Formatting.

* Test `intersphinxlink`

* More strict typing for `adjacency_mat`

* Test removing `fragment_map` from nitpick exceptions.

* Rename `valence_basis` to `iao_valence_basis`

* Finish renaming `valence_basis` to `iao_valence_basis`

* Ruff fixes.

* Final formatting.

* Update `molbe_ppp`

* Address Oskar comments.

* Simplified check for `fragpart.mol`

.ruff.toml

@@ -1,6 +1,7 @@
 [lint]
 # see the rules [here](https://docs.astral.sh/ruff/rules/)
 select = ["E", "F", "I", "NPY", "PL", "ARG"]
+exclude = ["tests/fragmentation_test.py"]
 ignore = [
     "S101",
         # https://docs.astral.sh/ruff/rules/assert/

docs/source/conf.py

@@ -58,6 +58,7 @@
     "python": ("https://docs.python.org/3", None),
     "pyscf": ("https://pyscf.org/", None),
     "h5py": ("https://docs.h5py.org/en/stable/", None),
+    "networkx": ("https://networkx.org/documentation/stable/", None),
 }
 
 

example/molbe_ppp.py

@@ -32,7 +32,7 @@
 mf.kernel()
 
 # Define fragments; use IAO scheme with 'sto-3g' as the minimal basis set
-fobj = fragpart(be_type="be2", mol=mol, valence_basis="sto-3g", frozen_core=True)
+fobj = fragpart(be_type="be2", mol=mol, iao_valence_basis="sto-3g", frozen_core=True)
 
 # Initialize BE
 mybe = BE(mf, fobj, lo_method="iao")

mypy.ini

@@ -18,7 +18,7 @@
     disallow_untyped_defs = False
     check_untyped_defs = False
 
-[mypy-tests.chem_dm_kBE_test,tests.chempot_molBE_test,tests.dm_molBE_test,tests.dmrg_molBE_test,tests.eri_onthefly_test,tests.hf-in-hf_BE_test,tests.kbe_polyacetylene_test,tests.molbe_h8_test,tests.molbe_io_fcidump_test,tests.molbe_octane_get_rdms_test,tests.molbe_oneshot_rbe_hcore_test,tests.molbe_oneshot_rbe_qmmm-fromchk_test,tests.ube-oneshot_test]
+[mypy-tests.fragmentation_test,tests.chem_dm_kBE_test,tests.chempot_molBE_test,tests.dm_molBE_test,tests.dmrg_molBE_test,tests.eri_onthefly_test,tests.hf-in-hf_BE_test,tests.kbe_polyacetylene_test,tests.molbe_h8_test,tests.molbe_io_fcidump_test,tests.molbe_octane_get_rdms_test,tests.molbe_oneshot_rbe_hcore_test,tests.molbe_oneshot_rbe_qmmm-fromchk_test,tests.ube-oneshot_test]
     disallow_untyped_defs = False
     check_untyped_defs = False
 

setup.py

@@ -22,6 +22,7 @@
         "numpy>=1.22.0",
         "scipy>=1.7.0",
         "pyscf>=2.0.0",
+        "networkx",
         "matplotlib",
         "libdmet @ git+https://github.com/gkclab/libdmet_preview.git",
         "attrs",

src/quemb/kbe/autofrag.py

@@ -241,7 +241,7 @@ def autogen(
     nx=False,
     ny=False,
     nz=False,
-    valence_basis=None,
+    iao_valence_basis=None,
     interlayer=False,
     print_frags=True,
 ):
@@ -274,7 +274,7 @@ def autogen(
     write_geom : bool, optional
         Whether to write a 'fragment.xyz' file which contains all the fragments in
         Cartesian coordinates. Defaults to False.
-    valence_basis : str, optional
+    iao_valence_basis : str, optional
         Name of minimal basis set for IAO scheme. 'sto-3g' is sufficient for most cases.
         Defaults to None.
     valence_only : bool, optional
@@ -1922,11 +1922,11 @@ def autogen(
                     )
         w.close()
 
-    pao = valence_basis is not None
+    pao = iao_valence_basis is not None
 
     if pao:
         cell2 = cell.copy()
-        cell2.basis = valence_basis
+        cell2.basis = iao_valence_basis
         cell2.build()
 
         bas2list = cell2.aoslice_by_atom()

src/quemb/kbe/fragment.py

@@ -23,7 +23,7 @@ def __init__(
         ny=False,
         nz=False,
         kpt=None,
-        valence_basis=None,
+        iao_valence_basis=None,
         be_type="be2",
         mol=None,
         frozen_core=False,
@@ -56,7 +56,7 @@ def __init__(
         mol : pyscf.pbc.gto.cell.Cell
             pyscf.pbc.gto.cell.Cell object. This is required for the options, 'autogen',
             and 'chain' as frag_type.
-        valence_basis: str
+        iao_valence_basis: str
             Name of minimal basis set for IAO scheme. 'sto-3g' suffice for most cases.
         frozen_core: bool
             Whether to invoke frozen core approximation. This is set to False by default
@@ -90,7 +90,7 @@ def __init__(
         self.frozen_core = frozen_core
         self.self_match = self_match
         self.allcen = allcen
-        self.valence_basis = valence_basis
+        self.iao_valence_basis = iao_valence_basis
         self.kpt = kpt
         self.molecule = False  # remove this
 
@@ -117,7 +117,7 @@ def __init__(
                 kpt,
                 be_type=be_type,
                 frozen_core=frozen_core,
-                valence_basis=valence_basis,
+                iao_valence_basis=iao_valence_basis,
                 unitcell=unitcell,
                 nx=nx,
                 ny=ny,

src/quemb/kbe/lo.py

@@ -45,7 +45,7 @@ class Mixin_k_Localize:
     def localize(
         self,
         lo_method,
-        valence_basis="sto-3g",
+        iao_valence_basis="sto-3g",
         core_basis="sto-3g",
         iao_wannier=True,
         iao_val_core=True,
@@ -60,7 +60,7 @@ def localize(
         lo_method : str
             Localization method in quantum chemistry. 'lowdin', 'boys','iao',
             and 'wannier' are supported.
-        valence_basis : str
+        iao_valence_basis : str
             Name of valence basis set for IAO scheme. 'sto-3g' suffice for most cases.
         core_basis : str
             Name of core basis set for IAO scheme. 'sto-3g' suffice for most cases.
@@ -128,7 +128,7 @@ def localize(
         elif lo_method == "iao":
             if not iao_val_core or not self.frozen_core:
                 Co = self.C[:, :, : self.Nocc].copy()
-                S12, S2 = get_xovlp_k(self.cell, self.kpts, basis=valence_basis)
+                S12, S2 = get_xovlp_k(self.cell, self.kpts, basis=iao_valence_basis)
                 ciao_ = get_iao_k(Co, S12, self.S, S2=S2)
 
                 # tmp - aos are not rearrange and so below is not necessary
@@ -145,7 +145,7 @@ def localize(
                 # Cpao = get_pao_k(Ciao, self.S, S12, S2, self.cell)
                 # get_pao_native_k returns symm orthogonalized orbitals
                 cpao_ = get_pao_native_k(
-                    Ciao_, self.S, self.cell, valence_basis, self.kpts
+                    Ciao_, self.S, self.cell, iao_valence_basis, self.kpts
                 )
 
                 nk, nao, nlo = cpao_.shape
@@ -192,7 +192,7 @@ def localize(
 
                 # Begin valence
                 s12_val_, s2_val = get_xovlp_k(
-                    self.cell, self.kpts, basis=valence_basis
+                    self.cell, self.kpts, basis=iao_valence_basis
                 )
                 C_nocore = self.C[:, :, self.ncore :].copy()
                 C_nocore_occ_ = C_nocore[:, :, : self.Nocc].copy()
@@ -235,7 +235,12 @@ def localize(
                     )
 
                 cpao_ = get_pao_native_k(
-                    c_core_val, self.S, self.cell, valence_basis, self.kpts, ortho=True
+                    c_core_val,
+                    self.S,
+                    self.cell,
+                    iao_valence_basis,
+                    self.kpts,
+                    ortho=True,
                 )
                 nk, nao, nlo = cpao_.shape
                 Cpao_ = zeros((nk, nao, nlo), dtype=complex128)

src/quemb/kbe/lo_k.py

@@ -192,7 +192,7 @@ def get_pao_k(Ciao, S, S12):
     return asarray(Cpao)
 
 
-def get_pao_native_k(Ciao, S, mol, valence_basis, ortho=True):
+def get_pao_native_k(Ciao, S, mol, iao_valence_basis, ortho=True):
     """
 
     Parameters
@@ -203,7 +203,7 @@ def get_pao_native_k(Ciao, S, mol, valence_basis, ortho=True):
         ao ovlp matrix
     mol :
         mol object
-    valence_basis:
+    iao_valence_basis:
         basis used for valence orbitals
 
     Returns
@@ -215,7 +215,7 @@ def get_pao_native_k(Ciao, S, mol, valence_basis, ortho=True):
 
     # Form a mol object with the valence basis for the ao_labels
     mol_alt = mol.copy()
-    mol_alt.basis = valence_basis
+    mol_alt.basis = iao_valence_basis
     mol_alt.build()
 
     full_ao_labels = mol.ao_labels()

src/quemb/kbe/pbe.py

@@ -263,7 +263,7 @@ def __init__(
             # Localize orbitals
             self.localize(
                 lo_method,
-                valence_basis=fobj.valence_basis,
+                iao_valence_basis=fobj.iao_valence_basis,
                 iao_wannier=iao_wannier,
                 iao_val_core=iao_val_core,
             )