[WIP] DMET Demo

demonstrations/tutorial_dmet_embedding.py

@@ -0,0 +1,277 @@
+r"""Density Matrix Embedding Theory (DMET)
+=========================================
+Materials simulation presents a crucial challenge in quantum chemistry, as understanding and predicting the properties of 
+complex materials is essential for advancements in technology and science. While Density Functional Theory (DFT) is 
+the current workhorse in this field due to its balance between accuracy and computational efficiency, it often falls short in 
+accurately capturing the intricate electron correlation effects found in strongly correlated materials. As a result, 
+researchers often turn to more sophisticated methods, such as full configuration interaction or coupled cluster theory, 
+which provide better accuracy but come at a significantly higher computational cost. 
+
+Embedding theories provide a balanced 
+midpoint solution that enhances our ability to simulate materials accurately and efficiently. The core idea behind embedding 
+is that the system is divided into two parts: impurity, strongly correlated subsystem that requires an exact description, and 
+its environment, which can be treated with an approximate but computationally efficient method.
+Density matrix embedding theory (DMET) is an efficient wave-function-based embedding approach to treat strongly 
+correlated systems. Here, we present a demonstration of how to run DMET calculations through an existing library called 
+libDMET, along with the instructions on how we can use the generated DMET Hamiltonian with PennyLane to use it with quantum 
+computing algorithms. We begin by providing a high-level introduction to DMET, followed by a tutorial on how to set up 
+a DMET calculation.
+
+.. figure:: ../_static/demo_thumbnails/opengraph_demo_thumbnails/OGthumbnail_how_to_build_spin_hamiltonians.png
+    :align: center
+    :width: 70%
+    :target: javascript:void(0)
+"""
+
+######################################################################
+# Theory
+# ------
+# DMET is a wavefunction based embedding approach, which uses density matrices for combining the low-level description 
+# of the environment with a high-level description of the impurity. DMET relies on Schmidt decomposition,
+# which allows us to analyze the degree of entanglement between the two subsystems. The state, :math:`\ket{\Psi}` of 
+# the partitioned system can be represented as the tensor product of the Hilbert space of the two subsystems.
+# Singular value decomposition (SVD) of the coefficient tensor, :math:`\psi_{ij}`, of this tensor product 
+# thus allows us to identify the states
+# in the environment which have overlap with the impurity. This helps truncate the size of the Hilbert space of the 
+# environment to be equal to the size of the impurity, and thus define a set of states referred to as bath. We are
+# then able to project the full Hamiltonian to the space of impurity and bath states, known as embedding space.
+# .. math::
+#
+#      \hat{H}^{imp} = \hat{P} \hat{H}^{sys}\hat{P}
+#
+# where P is the projection operator.
+# We must note here that the Schmidt decomposition requires  apriori knowledge of the wavefunction. DMET, therefore, 
+# operates through a systematic iterative approach, starting with a meanfield description of the wavefunction and 
+# refining it through feedback from solution of impurity Hamiltonian.
+#
+# The DMET procedure starts by getting an approximate description of the system, which is used to partition the system
+# into impurity and bath. We are then able to project the original Hamiltonian to this embedded space and  
+# solve it using a highly accurate method. This high-level description of impurity is then used to 
+# embed the updated correlation back into the full system, thus improving the initial approximation 
+# self-consistently. Let's take a look at the implementation of these steps.
+#
+######################################################################
+# Implementation
+# --------------
+# We now use what we have learned to set up a DMET calculation for $H_6$ system.
+#
+# Constructing the system
+# ^^^^^^^^^^^^^^^^^^^^^^^
+# We begin by defining a periodic system using PySCF [#pyscf]_ to create a cell object
+# representing a hydrogen chain with 6 atoms. Each unit cell contains two Hydrogen atoms at a bond 
+# distance of 0.75 Å.  Finally, we construct a Lattice object from the libDMET library, associating it with 
+# the defined cell and k-mesh, which allows for the use of DMET in studying the properties of
+#  the hydrogen chain system.
+import numpy as np
+from pyscf.pbc import gto, df, scf, tools
+from libdmet.system import lattice
+
+cell = gto.Cell()
+cell.a = ''' 10.0    0.0     0.0                                                                                                                                                                                      
+             0.0     10.0    0.0                                                                                                                                                                                      
+             0.0     0.0     1.5 ''' # lattice vectors for unit cell
+cell.atom = ''' H 0.0      0.0      0.0                                                                                                                                                                               
+                H 0.0      0.0      0.75 ''' # coordinates of atoms in unit cell
+cell.basis = '321g'
+cell.build(unit='Angstrom')
+
+kmesh = [1, 1, 3] # number of k-points in xyz direction
+lat = lattice.Lattice(cell, kmesh)
+filling = cell.nelectron / (lat.nscsites*2.0)
+kpts = lat.kpts
+
+######################################################################
+# We perform a mean-field calculation on the whole system through Hartree-Fock with density
+# fitted integrals using PySCF.
+gdf = df.GDF(cell, kpts)
+gdf._cderi_to_save = 'gdf_ints.h5' #output file for density fitted integral tensor
+gdf.build() #compute the density fitted integrals
+
+kmf = scf.KRHF(cell, kpts, exxdiv=None).density_fit()
+kmf.with_df = gdf #use density-fitted integrals
+kmf.with_df._cderi = 'gdf_ints.h5' #input file for density fitted integrals
+kmf.kernel() #run Hartree-Fock
+
+# Paritioning of Orbital Space
+# ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+# Now we have a description of our system and can start obtaining the impurity and bath orbitals.
+# This requires the localization of the basis of orbitals, we could use any localized basis here, 
+# for example, molecular orbitals(MO), intrinsic atomic orbitals(IAO), etc [#SWouters]_. The use of
+# localized basis here provides a mathematically convenient way to understand the contribution of
+# each atom to properties of the full system. Here, we 
+# rotate the one-electron and two-electron integrals into IAO basis.
+
+from libdmet.basis_transform import make_basis
+
+c_ao_iao, _, _, lo_labels = make_basis.get_C_ao_lo_iao(lat, kmf, minao="MINAO", full_return=True, return_labels=True)
+c_ao_lo = lat.symmetrize_lo(c_ao_iao)
+lat.set_Ham(kmf, gdf, c_ao_lo, eri_symmetry=4) #rotate integral tensors to IAO basis
+
+######################################################################
+# In ab initio systems, we can choose the bath and impurity by looking at the
+# labels of orbitals. We can get the orbitals for each atom in the unit cell by using
+# aoslice_by_atom function. This information helps us identify the orbitals to be included 
+# in the impurity, bath and unentangled environment.
+# In this example, we choose to keep the :math:`1s` orbitals in the unit cell in the
+# impurity, while the bath contains the :math:`2s` orbitals, and the orbitals belonging to the
+# rest of the supercell become part of the unentangled environment. These can be separated by 
+# getting the valence and virtual labels from get_labels function.
+from libdmet.lo.iao import get_labels
+
+aoind = cell.aoslice_by_atom()
+labels, val_labels, virt_labels = get_labels(cell, minao="MINAO")
+ncore = 0
+lat.set_val_virt_core(len(val_labels), len(virt_labels), ncore)
+print("Valence orbitals: ", val_labels)
+print("Virtual orbitals: ", virt_labels)
+
+######################################################################
+# Self-Consistent DMET
+# ^^^^^^^^^^^^^^^^^^^^
+# Now that we have a description of our impurity and bath orbitals, we can implement DMET. 
+# We implement each step of the process in a function and 
+# then call these functions to perform the calculations. This can be done once for one iteration,
+# referred to as single-shot DMET or we can call them iteratively to perform self-consistent DMET.
+# Let's start by constructing the impurity Hamiltonian,
+def construct_impurity_hamiltonian(lat, v_cor, filling, mu, last_dmu, int_bath=True):
+
+    rho, mu, scf_result = dmet.HartreeFock(lat, v_cor, filling, mu,
+                                    ires=True, labels=lo_labels)
+    imp_ham, _, basis = dmet.ConstructImpHam(lat, rho, v_cor, int_bath=int_bath)
+    imp_ham = dmet.apply_dmu(lat, imp_ham, basis, last_dmu)
+
+    return rho, mu, scf_result, imp_ham, basis
+
+# Next, we solve this impurity Hamiltonian with a high-level method, the following function defines
+# the electronic structure solver for the impurity, provides an initial point for the calculation and 
+# passes the Lattice information to the solver. The solver then calculates the energy and density matrix
+# for the impurity.
+def solve_impurity_hamiltonian(lat, cell, basis, imp_ham, last_dmu, scf_result):
+
+    solver = dmet.impurity_solver.FCI(restricted=True, tol=1e-13)
+    basis_k = lat.R2k_basis(basis) #basis in k-space
+
+    solver_args = {"nelec": min((lat.ncore+lat.nval)*2, lat.nkpts*cell.nelectron), \
+               "dm0": dmet.foldRho_k(scf_result["rho_k"], basis_k)}
+
+    rho_emb, energy_emb, imp_ham, dmu = dmet.SolveImpHam_with_fitting(lat, filling, 
+        imp_ham, basis, solver, solver_args=solver_args)
+
+    last_dmu += dmu
+    return rho_emb, energy_emb, imp_ham, last_dmu, [solver, solver_args]
+
+# Final step in single-shot DMET is to include the effect of environment in the final expectation value,
+#  so we define a function for the same which returns the density matrix and energy for the whole system.
+def solve_full_system(lat, rho_emb, energy_emb, basis, imp_ham, last_dmu, solver_info, lo_labels):
+    rho_full, energy_full, nelec_full = \
+            dmet.transformResults(rho_emb, energy_emb, basis, imp_ham, \
+            lattice=lat, last_dmu=last_dmu, int_bath=True, \
+                                  solver=solver_info[0], solver_args=solver_info[1], labels=lo_labels)
+    energy_full *= lat.nscsites
+    return rho_full, energy_full
+
+# We must note here that the effect of environment included in the previous step is
+# at the meanfield level, and will give the results for single-shot DMET. 
+# We can look at a more advanced version of DMET and improve this interaction 
+# with the use of self-consistency, referred to
+# as self-consistent DMET, where a correlation potential is introduced to account for the interactions 
+# between the impurity and its environment. We start with an initial guess of zero for this correlation
+#  potential and optimize it by minimizing the difference between density matrices obtained from the
+#  mean-field Hamiltonian and the impurity Hamiltonian. Let's initialize the correlation potential
+# and define a function to optimize it.
+def initialize_vcor(lat):
+    v_cor = dmet.VcorLocal(restricted=True, bogoliubov=False, nscsites=lat.nscsites)
+    v_cor.assign(np.zeros((2, lat.nscsites, lat.nscsites)))
+    return v_cor
+
+def fit_correlation_potential(rho_emb, lat, basis, v_cor):
+    vcor_new, err = dmet.FitVcor(rho_emb, lat, basis, \
+                v_cor, beta=np.inf, filling=filling, MaxIter1=300, MaxIter2=0)
+
+    dVcor_per_ele = np.max(np.abs(vcor_new.param - v_cor.param))
+    v_cor.update(vcor_new.param)
+    return v_cor, dVcor_per_ele
+
+# Now, we have defined all the ingredients of DMET, we can set up the self-consistency loop to get
+# the full execution. We set up this loop by defining the maximum number of iterations and a convergence
+# criteria. Here, we are using both energy and correlation potential as our convergence parameters, so we 
+# define the initial values and convergence tolerance for both.
+import libdmet.dmet.Hubbard as dmet
+
+max_iter = 10 # maximum number of iterations
+e_old = 0.0 # initial value of energy
+v_cor = initialize_vcor(lat) # initial value of correlation potential
+dVcor_per_ele = None # initial value of correlation potential per electron
+vcor_tol = 1.0e-5 # tolerance for correlation potential convergence
+energy_tol = 1.0e-5 # tolerance for energy convergence
+mu = 0 # initial chemical potential
+last_dmu = 0.0 # change in chemical potential
+for i in range(max_iter):
+    rho, mu, scf_result, imp_ham, basis = construct_impurity_hamiltonian(lat,
+                                     v_cor, filling, mu, last_dmu) # construct impurity Hamiltonian
+    rho_emb, energy_emb, imp_ham, last_dmu, solver_info = solve_impurity_hamiltonian(lat, cell,
+                                     basis, imp_ham, last_dmu, scf_result) # solve impurity Hamiltonian
+    rho_full, energy_full = solve_full_system(lat, rho_emb, energy_emb, basis, imp_ham,
+                                     last_dmu, solver_info, lo_labels) # include the environment interactions
+    v_cor, dVcor_per_ele = fit_correlation_potential(rho_emb,
+                                     lat, basis, v_cor) # fit correlation potential
+
+    dE = energy_full - e_old
+    e_old = energy_full
+    if dVcor_per_ele < vcor_tol and abs(dE) < energy_tol:
+        print("DMET Converged")
+        print("DMET Energy per cell: ", energy_full)
+        break
+
+# This concludes the DMET procedure. At this point, we should note that we are still limited by the number
+#  of orbitals we can have in the impurity because the cost of using a high-level solver such as FCI increases
+#  exponentially with increase in system size. One way to solve this problem could be through the use of 
+# quantum computing algorithm as solver. Next, we see how we can convert this impurity Hamiltonian to a 
+# qubit Hamiltonian through PennyLane to pave the path for using it with quantum algorithms.
+#  The ImpHam object generated above provides us with one-body and two-body integrals along with the 
+# nuclear repulsion energy which can be accessed as follows:
+from pyscf import ao2mo
+norb = imp_ham.norb
+h1 = imp_ham.H1["cd"]
+h2 = imp_ham.H2["ccdd"][0]
+h2 = ao2mo.restore(1, h2, norb) # Get the correct shape based on permutation symmetry
+
+# These one-body and two-body integrals can then be used to generate the qubit Hamiltonian for PennyLane.
+import pennylane as qml
+from pennylane.qchem import one_particle, two_particle, observable
+
+t = one_particle(h1[0])
+v = two_particle(np.swapaxes(h2, 1, 3)) # Swap to physicist's notation
+qubit_op = observable([t,v], mapping="jordan_wigner")
+eigval_qubit = qml.eigvals(qml.SparseHamiltonian(qubit_op.sparse_matrix(), wires = qubit_op.wires))
+print("eigenvalue from PennyLane: ", eigval_qubit)
+print("embedding energy: ", energy_emb)
+
+# We obtained the qubit Hamiltonian for embedded system here and diagonalized it to get the eigenvalues,
+#  and show that this eigenvalue matches the energy we obtained for the embedded system above.
+#  We can also get ground state energy for the system from this value
+#  by solving for the full system as done above in the self-consistency loop using solve_full_system function.
+
+######################################################################
+# Conclusion
+# ^^^^^^^^^^^^^^
+# Density matrix embedding theory is a robust method, designed to tackle simulation of complex systems, 
+# by dividing them into subsystems. It is specifically suited for studying the ground state properties of
+# a system. It provides for a computationally efficient alternative to dynamic quantum embedding schemes 
+# such as dynamic meanfield theory(DMFT), as it uses density matrix for embedding instead of the Green's function
+# and has limited number of bath orbitals. 
+# It has been successfully used for studying various strongly correlated molecular and periodic systems. 
+#
+# References
+# ----------
+#
+# .. [#SWouters]
+#    Sebastian Wouters, Carlos A. Jiménez-Hoyos, *et al.*, 
+#    "A practical guide to density matrix embedding theory in quantum chemistry."
+#    `ArXiv <https://arxiv.org/pdf/1603.08443>`__.
+#     
+#     
+# .. [#pyscf]
+#     Qiming Sun, Xing Zhang, *et al.*, "Recent developments in the PySCF program package."
+#     `ArXiv <https://arxiv.org/pdf/2002.12531>`__.
+#