Orbital refinement for excited states of H4 using state-averaged DMRG

In this tutorial, we will explore the process of orbital refinement for excited states of the H4 molecule by using the NWChem orbitals as guess orbitals and performing state-averaged DMRG calculations. First, the parameters for the multiwavelet representation are defined.

molecule_name = "h4"
n_elec = 4
number_roots = 3
iterations = 3
box_size = 50.0
wavelet_order = 7
madness_thresh = 0.0001
basisset = "6-31g"

1

Number of electrons

2

Number of states (ground state, 1st excited state, 2nd excited state)

3

Number of iterations for orbital refinement

4

Size of the simulation box

5

Order of wavelet basis functions

6

Threshold for numerical precision of function representation

7

Initial basis set for calculation

Run NWChem calculation

To get the initial guess orbitals, we will perform a NWChem calculation. If you are using the FrayedEnds devcontainer or the singularity image, NWChem is already installed. Otherwise, you will need to install NWChem and adjust the path in the code below.

import subprocess as sp

nwchem_input = (
    """
title "molecule"
memory stack 1500 mb heap 100 mb global 1400 mb
charge 0
geometry units angstrom noautosym nocenter
    H 0.0 0.0 -1.5
    H 0.0 0.0 -0.5
    H 0.0 0.0 0.5
    H 0.0 0.0 1.5
end
basis
  * library """
    + basisset
    + """
end
scf
 maxiter 200
end
task scf
"""
)

with open("nwchem", "w") as f:
    f.write(nwchem_input)
programm = sp.call(
    "/opt/anaconda3/envs/frayedends/bin/nwchem nwchem",
    stdout=open("nwchem.out", "w"),
    stderr=open("nwchem_err.log", "w"),
    shell=True,
)

1

Run NWChem calculation

2

Adjust the path to NWChem here if you are not using the FrayedEnds devcontainer or the singularity image.

Convert NWChem AOs and MOs to MRA-Orbitals

Next, we will read the molecular orbitals (MOs) from the NWChem calculation and translate them into multiwavelets by using the NWChem_Converter class in FrayedEnds.

import frayedends as fe

world = fe.MadWorld(ndims=3, L=box_size, k=wavelet_order, thresh=madness_thresh)

converter = fe.NWChem_Converter(world)
converter.read_nwchem_file("nwchem")
orbs = converter.get_mos()
Vnuc = converter.get_Vnuc()
nuclear_repulsion_energy = converter.get_nuclear_repulsion_energy()

n_orbitals = len(orbs)

molecule = fe.MolecularGeometry(units="angstrom")
molecule.add_atom(0.0, 0.0, -1.5, "H")
molecule.add_atom(0.0, 0.0, -0.5, "H")
molecule.add_atom(0.0, 0.0, 0.5, "H")
molecule.add_atom(0.0, 0.0, 1.5, "H")

for i in range(n_orbitals):
    world.cube_plot(f"initial_orb{i}", orbs[i], molecule)

1

Setting up the numerical environment for the MRA calculations by creating a frayedends world object with the specified parameters

2

Create an NWChem_Converter object

3

Read the NWChem output file to extract the orbitals and other relevant information

4

Get the molecular orbitals (MOs) from the converter

5

Get the nuclear potential from the converter

6

Get the nuclear repulsion energy from the converter

7

Define a linear H4 molecule geometry with 1.0 Angstrom spacing between adjacent atoms

MADNESS runtime initialized with 9 threads in the pool and affinity OFF

Calculate initial integrals

After obtaining the orbitals, we can calculate the initial integrals required for the DMRG calculation, including the two-body, kinetic, potential, and overlap integrals.

integrals = fe.Integrals(world)
G = integrals.compute_two_body_integrals(orbs, ordering="chem").elems
T = integrals.compute_kinetic_integrals(orbs)
V = integrals.compute_potential_integrals(orbs, Vnuc)
h1 = T + V
S = integrals.compute_overlap_integrals(orbs)

1

Create an integrals object to compute the initial integrals

2

Compute the two-body, kinetic, potential, and overlap integrals using the orbitals obtained from NWChem

Perform state-averaged DMRG calculation with orbital reordering and extract RDMs

To improve the efficiency of the DMRG calculation, an orbital reordering is performed using the one-body integral \(T + V\) and the two-body integral \(G\) obtained from the previous step. We then perform a state-averaged DMRG calculation using the reordered integrals and extract the one-body and two-body reduced density matrices (rdms). Finally, the resulting RDMs are transformed back into the original orbital ordering.

import numpy as np
from pyblock2.driver.core import DMRGDriver, SymmetryTypes

driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SU2, n_threads=8)
driver.initialize_system(n_sites=n_orbitals, n_elec=n_elec, spin=0)
mpo = driver.get_qc_mpo(h1e=h1, g2e=G, ecore=nuclear_repulsion_energy, iprint=0)
ket = driver.get_random_mps(tag="KET", bond_dim=100, nroots=number_roots)
energies = driver.dmrg(
    mpo, ket, n_sweeps=10, bond_dims=[100], noises=[1e-5] * 4 + [0], thrds=[1e-10] * 8, iprint=1
)

idx = driver.orbital_reordering(h1, G)
h1_new = h1[idx][:, idx]
g2_new = G[idx][:, idx][:, :, idx][:, :, :, idx]

driver.initialize_system(n_sites=n_orbitals, n_elec=n_elec, spin=0)
mpo = driver.get_qc_mpo(h1e=h1_new, g2e=g2_new, ecore=nuclear_repulsion_energy, iprint=0)
ket = driver.get_random_mps(tag="KET", bond_dim=100, nroots=number_roots)
energies = driver.dmrg(
    mpo, ket, n_sweeps=10, bond_dims=[100], noises=[1e-5] * 4 + [0], thrds=[1e-10] * 8, iprint=1
)
print("State-averaged MPS energies = [%s]" % " ".join("%20.15f" % x for x in energies))

kets = [driver.split_mps(ket, ir, tag="KET-%d" % ir) for ir in range(ket.nroots)]
sa_1pdm = np.mean([driver.get_1pdm(k) for k in kets], axis=0)
sa_2pdm = np.mean([driver.get_2pdm(k) for k in kets], axis=0).transpose(0, 3, 1, 2)
print(
    "Energy from SA-pdms = %20.15f"
    % (
        np.einsum("ij,ij->", sa_1pdm, h1_new)
        + 0.5 * np.einsum("ijkl,ijkl->", sa_2pdm, g2_new)
        + nuclear_repulsion_energy
    )
)

idx_back = np.zeros(len(idx), dtype=int)
for i in range(len(idx)):
    idx_back[idx[i]] = i

sa_1pdm = sa_1pdm[idx_back][:, idx_back]
sa_2pdm = sa_2pdm[idx_back][:, idx_back][:, :, idx_back][:, :, :, idx_back]
sa_2pdm_phys = sa_2pdm.swapaxes(1, 2)

1

Perform State Average (SA) DMRG calculation

2

Perform an orbital reordering using the one-body and two-body integrals

3

Perform a second State Average (SA) DMRG calculation with the reordered integrals

4

Extract reduced density matrices (rdms)

5

Compute the state-average one-body reduced density matrix

6

Compute the state average two-body reduced density matrix

7

Transform the rdms back into the original orbital ordering

Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.163 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DW = 5.67495e-20

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.586 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = 1.78e-15 | DW = 5.89849e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.995 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = 1.78e-15 | DW = 2.32726e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.524 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = 8.88e-16 | DW = 5.23044e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      2.935 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = 8.88e-16 | DW = 2.53627e-20


Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.367 | E[  3] =      -2.2248639286     -1.8539249223     -1.8307814236 | DW = 4.66699e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.709 | E[  3] =      -2.2248639286     -1.8539249223     -1.8307814236 | DE = 2.66e-15 | DW = 4.26176e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.959 | E[  3] =      -2.2248639286     -1.8539249223     -1.8307814236 | DE = -4.44e-15 | DW = 7.54028e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.218 | E[  3] =      -2.2248639286     -1.8539249223     -1.8307814236 | DE = -4.44e-15 | DW = 4.47832e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      1.286 | E[  3] =      -2.2248639286     -1.8539249223     -1.8307814236 | DE = 0.00e+00 | DW = 2.09600e-20

State-averaged MPS energies = [  -2.224863928635028   -1.853924922268516   -1.830781423564431]
Energy from SA-pdms =   -1.969856758155991

Orbital refinement

Finally, we will perform the orbital refinement by using the state-averaged 1-body and 2-body reduced density matrices obtained from the DMRG calculation. The refined orbitals can then be used for further state-averaged DMRG calculations to improve the accuracy of the excited state energies. The orbital refinement is repeated for a specified number of iterations (here: 3), as defined at the beginning of the tutorial.

import time

for iter in range(iterations):
    iter_start = time.perf_counter()

    opti = fe.OrbitalRefinement(world, Vnuc, nuclear_repulsion_energy)
    orbs = opti.get_orbitals(orbitals=orbs, rdm1=sa_1pdm, rdm2=sa_2pdm_phys, opt_thresh=0.001, occ_thresh=0.001)

    for i in range(n_orbitals):
        world.cube_plot(f"orb{i}", orbs[i], molecule)

    G = integrals.compute_two_body_integrals(orbs, ordering="chem").elems
    T = integrals.compute_kinetic_integrals(orbs)
    V = integrals.compute_potential_integrals(orbs, Vnuc)
    h1 = T + V
    S = integrals.compute_overlap_integrals(orbs)

    driver = DMRGDriver(scratch="./tmp", symm_type=SymmetryTypes.SU2, n_threads=8)
    driver.initialize_system(n_sites=n_orbitals, n_elec=n_elec, spin=0)
    mpo = driver.get_qc_mpo(h1e=h1, g2e=G, ecore=nuclear_repulsion_energy, iprint=0)
    ket = driver.get_random_mps(tag="KET", bond_dim=100, nroots=number_roots)
    energies = driver.dmrg(
        mpo, ket, n_sweeps=10, bond_dims=[100], noises=[1e-5] * 4 + [0], thrds=[1e-10] * 8, iprint=1
    )

    idx = driver.orbital_reordering(h1, G)
    h1_new = h1[idx][:, idx]
    g2_new = G[idx][:, idx][:, :, idx][:, :, :, idx]

    driver.initialize_system(n_sites=n_orbitals, n_elec=n_elec, spin=0)
    mpo = driver.get_qc_mpo(h1e=h1_new, g2e=g2_new, ecore=nuclear_repulsion_energy, iprint=0)
    ket = driver.get_random_mps(tag="KET", bond_dim=100, nroots=number_roots)
    energies = driver.dmrg(
        mpo, ket, n_sweeps=10, bond_dims=[100], noises=[1e-5] * 4 + [0], thrds=[1e-10] * 8, iprint=1
    )
    print("State-averaged MPS energies after refinement = [%s]" % " ".join("%20.15f" % x for x in energies))

    kets = [driver.split_mps(ket, ir, tag="KET-%d" % ir) for ir in range(ket.nroots)]
    sa_1pdm = np.mean([driver.get_1pdm(k) for k in kets], axis=0)
    sa_2pdm = np.mean([driver.get_2pdm(k) for k in kets], axis=0).transpose(0, 3, 1, 2)
    print(
        "Energy from SA-pdms = %20.15f"
        % (
            np.einsum("ij,ij->", sa_1pdm, h1_new)
            + 0.5 * np.einsum("ijkl,ijkl->", sa_2pdm, g2_new)
            + nuclear_repulsion_energy
        )
    )
    idx_back = np.zeros(len(idx), dtype=int)
    for i in range(len(idx)):
        idx_back[idx[i]] = i

    sa_1pdm = sa_1pdm[idx_back][:, idx_back]
    sa_2pdm = sa_2pdm[idx_back][:, idx_back][:, :, idx_back][:, :, :, idx_back]
    sa_2pdm_phys = sa_2pdm.swapaxes(1, 2)

fe.cleanup(globals())

1

Create an optimization object for orbital refinement and get the refined orbitals using the state-averaged one-body and two-body rdms

2

Plot the refined orbitals

3

Calculate the integrals with the refined orbitals

4

Perform a state-averaged DMRG calculation with the refined orbitals

5

Perform an orbital reordering using the one-body and two-body integrals

6

Perform a second state-averaged DMRG calculation with the reordered integrals

7

Extract reduced density matrices (rdms)

8

Compute the state-average one-body rdm with the refined orbitals

9

Compute the state-average two-body rdm with the refined orbitals

10

Transform the rdms back into the original orbital ordering

Output from the orbital refinement loop

Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.486 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DW = 1.19185e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.914 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = -8.88e-16 | DW = 5.21074e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.206 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = -9.77e-15 | DW = 1.08084e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.630 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = 8.88e-16 | DW = 6.40138e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      2.814 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = 3.55e-15 | DW = 4.57325e-20


Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.612 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DW = 1.84445e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.891 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = -8.88e-16 | DW = 2.78420e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.100 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = 8.88e-16 | DW = 2.04572e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.466 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = -8.88e-15 | DW = 8.63092e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      2.601 | E[  3] =      -2.2370817416     -1.9070522959     -1.8537964808 | DE = 8.88e-16 | DW = 2.21081e-20

State-averaged MPS energies after refinement = [  -2.237081741646887   -1.907052295896205   -1.853796480842818]
Energy from SA-pdms =   -1.999310172795306

Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.557 | E[  3] =      -2.2358492873     -1.9227693728     -1.8543687643 | DW = 1.16552e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.924 | E[  3] =      -2.2358492873     -1.9227693728     -1.8543687643 | DE = -8.88e-16 | DW = 4.99879e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.280 | E[  3] =      -2.2358492873     -1.9227693728     -1.8543687643 | DE = 4.44e-15 | DW = 1.10977e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.721 | E[  3] =      -2.2358492873     -1.9227693728     -1.8543687643 | DE = -8.88e-16 | DW = 6.15774e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      3.146 | E[  3] =      -2.2358492873     -1.9227693728     -1.8543687643 | DE = -7.99e-15 | DW = 3.35582e-20


Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.093 | E[  3] =      -2.2358492873     -1.9227693727     -1.8543687643 | DW = 2.38636e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.386 | E[  3] =      -2.2358492873     -1.9227693727     -1.8543687643 | DE = 5.33e-15 | DW = 1.36163e-19

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.830 | E[  3] =      -2.2358492873     -1.9227693727     -1.8543687643 | DE = 0.00e+00 | DW = 1.21338e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      3.113 | E[  3] =      -2.2358492873     -1.9227693727     -1.8543687643 | DE = 0.00e+00 | DW = 3.78870e-19

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      3.411 | E[  3] =      -2.2358492873     -1.9227693727     -1.8543687643 | DE = -1.78e-15 | DW = 4.08407e-20

State-averaged MPS energies after refinement = [  -2.235849287317252   -1.922769372728945   -1.854368764314865]
Energy from SA-pdms =   -2.004329141453686

Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.174 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DW = 1.80167e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.375 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = 2.66e-15 | DW = 4.85680e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.930 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = 5.33e-15 | DW = 1.11411e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      3.440 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = 0.00e+00 | DW = 7.07280e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      3.943 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = -8.88e-16 | DW = 3.45844e-20


Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.383 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DW = 1.14703e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.670 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = 1.24e-14 | DW = 1.14906e-18

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.166 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = 8.88e-16 | DW = 2.78132e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.459 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = -8.88e-16 | DW = 3.43592e-19

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      2.948 | E[  3] =      -2.2348149044     -1.9295230913     -1.8545233285 | DE = 8.88e-16 | DW = 2.10052e-20

State-averaged MPS energies after refinement = [  -2.234814904414891   -1.929523091285816   -1.854523328461209]
Energy from SA-pdms =   -2.006287108053977

The orbital refinement process is repeated for the specified number of iterations. At the end of the orbital refinement, the energies of the ground state and excited states should be improved compared to the initial DMRG calculation with the NWChem orbitals as guess orbitals.

The initial orbitals obtained from NWChem and the refined orbitals after the orbital optimization can be visualized using the cube files and plotted interactively using py3Dmol.

3Dmol.js failed to load for some reason. Please check your browser console for error messages.