Skip to content

Commit 938ce79

Browse files
MIWdlBMIWdlB
authored
Simpler localized system object (#184)
In preparation for handling Restricted open, this PR includes some simplifications to the LocalizedSystem dataclass. The idea is that in future we will want to handle spin localizations of different numbers of electrons in each spin, without having to explicitly check for this downstream. This should make it easier to handle the change in array dimension between MO coefficient matrix and occupancies for, ROKS, RKS, UKS respectively. property\method RKS ROKS UKS Occupancy (n) (2,n) (2,n) MO Coeff (M,M) (M,M) (2,M,M) This gets extra complicated when we allow any combination of R,RO,U as each of input and output, with variable number of electrons. To reduce the variability in the objects that are being stored and tracked, occupancy arrays are now boolean arrays rather than arrays of indices. This preserves the shape (n) for each spin. Rather than pass c_active and c_enviro, we instead pass density matrices, which preserve the shape (M,M) for each dm, rather than ( M , n a c t ) .
1 parent 921116e commit 938ce79

File tree

14 files changed

+1932
-1667
lines changed

14 files changed

+1932
-1667
lines changed

‎docs/jupyter_execute/notebooks/6. Long form Notes.ipynb‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -515,8 +515,8 @@
515515
"metadata": {},
516516
"outputs": [],
517517
"source": [
518-
"print(f\"active inds: {localized_system.active_mo_inds}\")\n",
519-
"print(f\"enviro inds: {localized_system.enviro_mo_inds}\")"
518+
"print(f\"active inds: {localized_system.active_occ_inds}\")\n",
519+
"print(f\"enviro inds: {localized_system.enviro_occ_inds}\")"
520520
]
521521
},
522522
{
@@ -763,7 +763,7 @@
763763
" charge=charge,\n",
764764
" spin=spin,\n",
765765
")\n",
766-
"rks_emb_mol.nelectron = 2 * len(localized_system.active_mo_inds) # <--change no e-\n",
766+
"rks_emb_mol.nelectron = 2 * len(localized_system.active_occ_inds) # <--change no e-\n",
767767
"rks_emb_mol.build()\n",
768768
"\n",
769769
"rks_emb = scf.RKS(rks_emb_mol)\n",
@@ -855,7 +855,7 @@
855855
" spin=spin,\n",
856856
")\n",
857857
"full_system_mol.nelectron = 2 * len(\n",
858-
" localized_system.active_mo_inds\n",
858+
" localized_system.active_occ_inds\n",
859859
") # <- change number of e-\n",
860860
"full_system_mol.build()\n",
861861
"\n",
@@ -1038,7 +1038,7 @@
10381038
" spin=spin,\n",
10391039
")\n",
10401040
"full_system_mol_HUZ.nelectron = 2 * len(\n",
1041-
" localized_system.active_mo_inds\n",
1041+
" localized_system.active_occ_inds\n",
10421042
") # <- change number of e-\n",
10431043
"full_system_mol_HUZ.build()\n",
10441044
"\n",
@@ -1139,7 +1139,7 @@
11391139
"\n",
11401140
"overlap_by_size = overlap.argsort()[::-1]\n",
11411141
"\n",
1142-
"n_env_mo = len(localized_system.enviro_mo_inds)\n",
1142+
"n_env_mo = len(localized_system.enviro_occ_inds)\n",
11431143
"frozen_enviro_orb_inds = overlap_by_size[:n_env_mo]\n",
11441144
"active_MOs_occ_and_virt_embedded = [\n",
11451145
" mo_i\n",
@@ -1367,8 +1367,8 @@
13671367
"# note we only take MO environment indices!\n",
13681368
"ortho_proj = np.einsum(\n",
13691369
" \"ik,jk->ij\",\n",
1370-
" c_loc_ortho[:, localized_system.enviro_mo_inds],\n",
1371-
" c_loc_ortho[:, localized_system.enviro_mo_inds],\n",
1370+
" c_loc_ortho[:, localized_system.enviro_occ_inds],\n",
1371+
" c_loc_ortho[:, localized_system.enviro_occ_inds],\n",
13721372
")"
13731373
]
13741374
},

‎docs/jupyter_execute/notebooks/localization.ipynb‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -317,8 +317,8 @@
317317
"n_env_mos = n_occupied_orbitals - n_act_mos\n",
318318
"\n",
319319
"# get active and enviro indices\n",
320-
"active_mo_inds = np.arange(n_act_mos)\n",
321-
"enviro_mo_inds = np.arange(n_act_mos, n_act_mos + n_env_mos)\n",
320+
"active_occ_inds = np.arange(n_act_mos)\n",
321+
"enviro_occ_inds = np.arange(n_act_mos, n_act_mos + n_env_mos)\n",
322322
"\n",
323323
"# Defining active and environment orbitals and density\n",
324324
"c_active = occupied_orbitals @ right_vectors.T[:, :n_act_mos]\n",
@@ -343,7 +343,7 @@
343343
}
344344
],
345345
"source": [
346-
"enviro_mo_inds"
346+
"enviro_occ_inds"
347347
]
348348
},
349349
{
@@ -357,8 +357,8 @@
357357
"text": [
358358
"n_act_mos=np.int64(9)\n",
359359
"n_env_mos=np.int64(3)\n",
360-
"active_mo_inds=array([0, 1, 2, 3, 4, 5, 6, 7, 8])\n",
361-
"enviro_mo_inds=array([ 9, 10, 11])\n",
360+
"active_occ_inds=array([0, 1, 2, 3, 4, 5, 6, 7, 8])\n",
361+
"enviro_occ_inds=array([ 9, 10, 11])\n",
362362
"c_active.shape=(18, 9)\n",
363363
"c_enviro.shape=(18, 3)\n"
364364
]
@@ -368,8 +368,8 @@
368368
"print(f\"{n_act_mos=}\")\n",
369369
"print(f\"{n_env_mos=}\")\n",
370370
"\n",
371-
"print(f\"{active_mo_inds=}\")\n",
372-
"print(f\"{enviro_mo_inds=}\")\n",
371+
"print(f\"{active_occ_inds=}\")\n",
372+
"print(f\"{enviro_occ_inds=}\")\n",
373373
"\n",
374374
"print(f\"{c_active.shape=}\")\n",
375375
"print(f\"{c_enviro.shape=}\")"

‎docs/jupyter_execute/notebooks/publications/A Scalable Approach to Quantum Simulation via Projection-based Embedding.ipynb‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -878,7 +878,7 @@
878878
"\n",
879879
" exp_data = {\n",
880880
" \"bond_length\": R,\n",
881-
" \"N active MOs\": len(driver.localized_system.active_mo_inds),\n",
881+
" \"N active MOs\": len(driver.localized_system.active_occ_inds),\n",
882882
" \"global_DFT\": driver._global_ks.e_tot,\n",
883883
" \"global_FCI\": E_diag,\n",
884884
" \"global_qham\": [global_qubit_Ham],\n",
@@ -1220,7 +1220,7 @@
12201220
" result[active][\"mu\"][\"classical_energy\"] = driver.mu[\"classical_energy\"]\n",
12211221
" result[active][\"mu\"][\"ground\"] = None\n",
12221222
" result[active][\"mu\"][\"e_ccsd\"] = driver.mu[\"e_ccsd\"]\n",
1223-
" result[active][\"mu\"][\"nmos\"] = len(driver.localized_system.active_mo_inds)\n",
1223+
" result[active][\"mu\"][\"nmos\"] = len(driver.localized_system.active_occ_inds)\n",
12241224
" print(\"Mu finished.\")\n",
12251225
"\n",
12261226
" # Huzinaga result\n",
@@ -1237,7 +1237,7 @@
12371237
" result[active][\"huzinaga\"][\"classical_energy\"] = driver.huzinaga[\"classical_energy\"]\n",
12381238
" result[active][\"huzinaga\"][\"ground\"] = None\n",
12391239
" result[active][\"huzinaga\"][\"e_ccsd\"] = driver.huzinaga[\"e_ccsd\"]\n",
1240-
" result[active][\"huzinaga\"][\"nmos\"] = len(driver.localized_system.active_mo_inds)\n",
1240+
" result[active][\"huzinaga\"][\"nmos\"] = len(driver.localized_system.active_occ_inds)\n",
12411241
" print(\"Huzinaga finished.\")\n",
12421242
"\n",
12431243
" # untapered_mu = mu_builder.build(taper=False)\n",

‎docs/notebooks/6. Long form Notes.ipynb‎

Lines changed: 8 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -502,8 +502,8 @@
502502
"metadata": {},
503503
"outputs": [],
504504
"source": [
505-
"print(f\"active inds: {localized_system.active_mo_inds}\")\n",
506-
"print(f\"enviro inds: {localized_system.enviro_mo_inds}\")"
505+
"print(f\"active inds: {localized_system.active_occ_inds}\")\n",
506+
"print(f\"enviro inds: {localized_system.enviro_occ_inds}\")"
507507
]
508508
},
509509
{
@@ -750,7 +750,7 @@
750750
" charge=charge,\n",
751751
" spin=spin,\n",
752752
")\n",
753-
"rks_emb_mol.nelectron = 2 * len(localized_system.active_mo_inds) # <--change no e-\n",
753+
"rks_emb_mol.nelectron = 2 * len(localized_system.active_occ_inds) # <--change no e-\n",
754754
"rks_emb_mol.build()\n",
755755
"\n",
756756
"rks_emb = scf.RKS(rks_emb_mol)\n",
@@ -842,7 +842,7 @@
842842
" spin=spin,\n",
843843
")\n",
844844
"full_system_mol.nelectron = 2 * len(\n",
845-
" localized_system.active_mo_inds\n",
845+
" localized_system.active_occ_inds\n",
846846
") # <- change number of e-\n",
847847
"full_system_mol.build()\n",
848848
"\n",
@@ -1025,7 +1025,7 @@
10251025
" spin=spin,\n",
10261026
")\n",
10271027
"full_system_mol_HUZ.nelectron = 2 * len(\n",
1028-
" localized_system.active_mo_inds\n",
1028+
" localized_system.active_occ_inds\n",
10291029
") # <- change number of e-\n",
10301030
"full_system_mol_HUZ.build()\n",
10311031
"\n",
@@ -1126,7 +1126,7 @@
11261126
"\n",
11271127
"overlap_by_size = overlap.argsort()[::-1]\n",
11281128
"\n",
1129-
"n_env_mo = len(localized_system.enviro_mo_inds)\n",
1129+
"n_env_mo = len(localized_system.enviro_occ_inds)\n",
11301130
"frozen_enviro_orb_inds = overlap_by_size[:n_env_mo]\n",
11311131
"active_MOs_occ_and_virt_embedded = [\n",
11321132
" mo_i\n",
@@ -1354,8 +1354,8 @@
13541354
"# note we only take MO environment indices!\n",
13551355
"ortho_proj = np.einsum(\n",
13561356
" \"ik,jk->ij\",\n",
1357-
" c_loc_ortho[:, localized_system.enviro_mo_inds],\n",
1358-
" c_loc_ortho[:, localized_system.enviro_mo_inds],\n",
1357+
" c_loc_ortho[:, localized_system.enviro_occ_inds],\n",
1358+
" c_loc_ortho[:, localized_system.enviro_occ_inds],\n",
13591359
")"
13601360
]
13611361
},

‎docs/notebooks/publications/A Scalable Approach to Quantum Simulation via Projection-based Embedding.ipynb‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -849,7 +849,7 @@
849849
"\n",
850850
" exp_data = {\n",
851851
" \"bond_length\": R,\n",
852-
" \"N active MOs\": len(driver.localized_system.active_mo_inds),\n",
852+
" \"N active MOs\": len(driver.localized_system.active_occ_inds),\n",
853853
" \"global_DFT\": driver._global_ks.e_tot,\n",
854854
" \"global_FCI\": E_diag,\n",
855855
" \"global_qham\": [global_qubit_Ham],\n",
@@ -1191,7 +1191,7 @@
11911191
" result[active][\"mu\"][\"classical_energy\"] = driver.mu[\"classical_energy\"]\n",
11921192
" result[active][\"mu\"][\"ground\"] = None\n",
11931193
" result[active][\"mu\"][\"e_ccsd\"] = driver.mu[\"e_ccsd\"]\n",
1194-
" result[active][\"mu\"][\"nmos\"] = len(driver.localized_system.active_mo_inds)\n",
1194+
" result[active][\"mu\"][\"nmos\"] = len(driver.localized_system.active_occ_inds)\n",
11951195
" print(\"Mu finished.\")\n",
11961196
"\n",
11971197
" # Huzinaga result\n",
@@ -1208,7 +1208,7 @@
12081208
" result[active][\"huzinaga\"][\"classical_energy\"] = driver.huzinaga[\"classical_energy\"]\n",
12091209
" result[active][\"huzinaga\"][\"ground\"] = None\n",
12101210
" result[active][\"huzinaga\"][\"e_ccsd\"] = driver.huzinaga[\"e_ccsd\"]\n",
1211-
" result[active][\"huzinaga\"][\"nmos\"] = len(driver.localized_system.active_mo_inds)\n",
1211+
" result[active][\"huzinaga\"][\"nmos\"] = len(driver.localized_system.active_occ_inds)\n",
12121212
" print(\"Huzinaga finished.\")\n",
12131213
"\n",
12141214
" # untapered_mu = mu_builder.build(taper=False)\n",

‎docs/source/notebooks/localization.ipynb‎

Lines changed: 7 additions & 7 deletions
Original file line numberDiff line numberDiff line change
@@ -317,8 +317,8 @@
317317
"n_env_mos = n_occupied_orbitals - n_act_mos\n",
318318
"\n",
319319
"# get active and enviro indices\n",
320-
"active_mo_inds = np.arange(n_act_mos)\n",
321-
"enviro_mo_inds = np.arange(n_act_mos, n_act_mos + n_env_mos)\n",
320+
"active_occ_inds = np.arange(n_act_mos)\n",
321+
"enviro_occ_inds = np.arange(n_act_mos, n_act_mos + n_env_mos)\n",
322322
"\n",
323323
"# Defining active and environment orbitals and density\n",
324324
"c_active = occupied_orbitals @ right_vectors.T[:, :n_act_mos]\n",
@@ -343,7 +343,7 @@
343343
}
344344
],
345345
"source": [
346-
"enviro_mo_inds"
346+
"enviro_occ_inds"
347347
]
348348
},
349349
{
@@ -357,8 +357,8 @@
357357
"text": [
358358
"n_act_mos=np.int64(9)\n",
359359
"n_env_mos=np.int64(3)\n",
360-
"active_mo_inds=array([0, 1, 2, 3, 4, 5, 6, 7, 8])\n",
361-
"enviro_mo_inds=array([ 9, 10, 11])\n",
360+
"active_occ_inds=array([0, 1, 2, 3, 4, 5, 6, 7, 8])\n",
361+
"enviro_occ_inds=array([ 9, 10, 11])\n",
362362
"c_active.shape=(18, 9)\n",
363363
"c_enviro.shape=(18, 3)\n"
364364
]
@@ -368,8 +368,8 @@
368368
"print(f\"{n_act_mos=}\")\n",
369369
"print(f\"{n_env_mos=}\")\n",
370370
"\n",
371-
"print(f\"{active_mo_inds=}\")\n",
372-
"print(f\"{enviro_mo_inds=}\")\n",
371+
"print(f\"{active_occ_inds=}\")\n",
372+
"print(f\"{enviro_occ_inds=}\")\n",
373373
"\n",
374374
"print(f\"{c_active.shape=}\")\n",
375375
"print(f\"{c_enviro.shape=}\")"

‎nbed/driver.py‎

Lines changed: 17 additions & 17 deletions
Original file line numberDiff line numberDiff line change
@@ -266,18 +266,22 @@ def _init_embedded_mol(self) -> gto.Mole:
266266
gto.Mole: An embedded molecule object.
267267
"""
268268
embedded_mol: gto.Mole = self._build_mol()
269-
match self.localized_system.active_mo_inds.ndim:
269+
match self.localized_system.active_occ_inds.ndim:
270270
case 1:
271-
n_elec = len(self.localized_system.active_mo_inds)
271+
n_elec = np.count_nonzero(self.localized_system.active_occ_inds)
272272
logger.debug(f"embedded nelec {n_elec}")
273273

274274
embedded_mol.nelectron = 2 * n_elec
275275
embedded_mol.nelec = (n_elec, n_elec)
276276
embedded_mol.spin = 0
277277
self._electron = embedded_mol.nelectron
278278
case 2:
279-
n_elec_alpha = len(self.localized_system.active_mo_inds[0, :])
280-
n_elec_beta = len(self.localized_system.active_mo_inds[1, :])
279+
n_elec_alpha = np.count_nonzero(
280+
self.localized_system.active_occ_inds[0, :]
281+
)
282+
n_elec_beta = np.count_nonzero(
283+
self.localized_system.active_occ_inds[1, :]
284+
)
281285
logger.debug(f"embedded nelec {n_elec_alpha, n_elec_beta}")
282286

283287
embedded_mol.nelectron = n_elec_alpha + n_elec_beta
@@ -596,7 +600,7 @@ def _huzinaga_embed(
596600
v_emb = huzinaga_op_std + embedding_potential
597601
active_scf.get_hcore = lambda *args: hcore_std + v_emb
598602

599-
if localized_system.c_active.ndim == 3:
603+
if localized_system.dm_active.ndim == 3:
600604
active_scf.energy_elec = lambda *args: energy_elec(active_scf, *args)
601605

602606
active_scf.mo_occ = active_scf.get_occ(mo_embedded_energy, c_active_embedded)
@@ -654,9 +658,9 @@ def _delete_environment(
654658
"""
655659
logger.debug("Deleting environment from SCF object.")
656660

657-
match localized_system.c_enviro.ndim:
661+
match localized_system.dm_enviro.ndim:
658662
case 2:
659-
n_env_mos = localized_system.c_enviro.shape[-1]
663+
n_env_mos = np.sum(localized_system.c_enviro.shape)
660664
logger.debug(f"{n_env_mos=}")
661665
scf.mo_coeff, scf.mo_energy, scf.mo_occ = self._delete_spin_environment(
662666
projector,
@@ -668,19 +672,18 @@ def _delete_environment(
668672
)
669673
case 3:
670674
#
671-
n_env_mos = len(
672-
set(localized_system.enviro_mo_inds[0]).union(
673-
localized_system.enviro_mo_inds[1]
674-
)
675-
)
675+
n_env_mos = [
676+
np.sum(localized_system.enviro_occ_inds[0]),
677+
np.sum(localized_system.enviro_occ_inds[1]),
678+
]
676679
logger.debug(f"{n_env_mos=}")
677680
(
678681
mo_coeff_alpha,
679682
mo_energy_alpha,
680683
mo_occ_alpha,
681684
) = self._delete_spin_environment(
682685
projector,
683-
n_env_mos,
686+
n_env_mos[0],
684687
scf.mo_coeff[0],
685688
scf.mo_energy[0],
686689
scf.mo_occ[0],
@@ -689,7 +692,7 @@ def _delete_environment(
689692
(mo_coeff_beta, mo_energy_beta, mo_occ_beta) = (
690693
self._delete_spin_environment(
691694
projector,
692-
n_env_mos,
695+
n_env_mos[1],
693696
scf.mo_coeff[1],
694697
scf.mo_energy[1],
695698
scf.mo_occ[1],
@@ -832,9 +835,6 @@ def embed(
832835
self.n_mo_overwrite = self.config.n_mo_overwrite
833836

834837
self.localized_system = self._localize()
835-
logger.info("Indices of embedded electrons:")
836-
logger.info(self.localized_system.active_mo_inds)
837-
logger.info(self.localized_system.enviro_mo_inds)
838838

839839
# Run subsystem DFT (calls localized rks)
840840
self.e_act, self.e_env, self.two_e_cross = self._subsystem_dft(

0 commit comments

Comments
 (0)