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

code
Author

Thuy Truong

Published

March 10, 2026

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(
    "/Users/timo/miniforge3/envs/fe_test/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.MadWorld3D(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)

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

  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

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.Integrals3D(world)
G = integrals.compute_two_body_integrals(orbs, ordering="chem").elems
T = integrals.compute_kinetic_integrals(orbs)
V = integrals.compute_potential_integrals(orbs, Vnuc)
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 and extract rdms

Now we will perform a state-averaged DMRG calculation using the integrals obtained from the previous step and extract the one-body and two-body reduced density matrices (rdms).

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=T + V, 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
)
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, T + V) + 0.5 * np.einsum("ijkl,ijkl->", sa_2pdm, G) + nuclear_repulsion_energy)
)
sa_2pdm_phys = sa_2pdm.swapaxes(1, 2)
1
Performe State Average (SA) DMRG calculation and extract rdms
2
Compute the state-average 1-body reduced density matrix
3
Compute the state average 2-body reduced density matrix 
Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.228 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DW = 1.98183e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.370 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = 2.22e-14 | DW = 4.53005e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.619 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = 0.00e+00 | DW = 2.83627e-19

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

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      0.989 | E[  3] =      -2.2248639286     -1.8539249222     -1.8307814235 | DE = -1.78e-15 | DW = 2.52047e-20

State-averaged MPS energies = [  -2.224863928644570   -1.853924922247054   -1.830781423545666]
Energy from SA-pdms =   -1.969856758145756

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.Optimization3D(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)
    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=T + V, 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
    )
    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, T + V) + 0.5 * np.einsum("ijkl,ijkl->", sa_2pdm, G) + nuclear_repulsion_energy)
    )
    sa_2pdm_phys = sa_2pdm.swapaxes(1, 2)

fe.cleanup(globals())
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.333 | E[  3] =      -2.2370817475     -1.9070522038     -1.8537965136 | DW = 1.06290e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.527 | E[  3] =      -2.2370817475     -1.9070522038     -1.8537965136 | DE = -8.88e-16 | DW = 6.30342e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.877 | E[  3] =      -2.2370817475     -1.9070522038     -1.8537965136 | DE = -3.55e-15 | DW = 6.35324e-20

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.331 | E[  3] =      -2.2370817475     -1.9070522038     -1.8537965136 | DE = 8.88e-16 | DW = 6.51451e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      2.563 | E[  3] =      -2.2370817475     -1.9070522038     -1.8537965136 | DE = -4.44e-15 | DW = 2.47761e-20

State-averaged MPS energies after refinement = [  -2.237081747471295   -1.907052203844466   -1.853796513604087]
Energy from SA-pdms =   -1.999310154973277

Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.295 | E[  3] =      -2.2358492990     -1.9227693171     -1.8543688642 | DW = 1.20569e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.721 | E[  3] =      -2.2358492990     -1.9227693171     -1.8543688642 | DE = 0.00e+00 | DW = 4.49642e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.234 | E[  3] =      -2.2358492990     -1.9227693171     -1.8543688642 | DE = 1.78e-15 | DW = 8.81862e-20

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      2.747 | E[  3] =      -2.2358492990     -1.9227693171     -1.8543688642 | DE = 1.78e-15 | DW = 4.32698e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      2.934 | E[  3] =      -2.2358492990     -1.9227693171     -1.8543688642 | DE = -6.22e-15 | DW = 2.77518e-20

State-averaged MPS energies after refinement = [  -2.235849299027298   -1.922769317104890   -1.854368864176128]
Energy from SA-pdms =   -2.004329160102771

Sweep =    0 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      0.990 | E[  3] =      -2.2348149251     -1.9295229542     -1.8545234556 | DW = 1.90691e-19

Sweep =    1 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.463 | E[  3] =      -2.2348149251     -1.9295229542     -1.8545234556 | DE = -3.55e-15 | DW = 5.31156e-20

Sweep =    2 | Direction =  forward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.667 | E[  3] =      -2.2348149251     -1.9295229542     -1.8545234556 | DE = 0.00e+00 | DW = 1.51803e-19

Sweep =    3 | Direction = backward | Bond dimension =  100 | Noise =  1.00e-05 | Dav threshold =  1.00e-10
Time elapsed =      1.813 | E[  3] =      -2.2348149251     -1.9295229542     -1.8545234556 | DE = 1.78e-15 | DW = 6.25487e-20

Sweep =    4 | Direction =  forward | Bond dimension =  100 | Noise =  0.00e+00 | Dav threshold =  1.00e-10
Time elapsed =      1.959 | E[  3] =      -2.2348149251     -1.9295229542     -1.8545234556 | DE = -3.55e-15 | DW = 3.34387e-20

State-averaged MPS energies after refinement = [  -2.234814925146770   -1.929522954166136   -1.854523455619775]
Energy from SA-pdms =   -2.006287111644225
  1. Create an optimization object for orbital refinement and get the refined orbitals using the state-averaged 1-body and 2-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 and print the energies after refinement
  5. Compute the state-average 1-body rdm with the refined orbitals
  6. Compute the state-average 2-body rdm with the refined orbitals

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.