Example 1: Relaxation of H2O vibrational state under polynomial PES

run type	wavefunction	backend	Basis	steps
improved relaxation	MPS-SM	Numpy	HO-DVR	20

1. Import modules

• Required in any calculations

from pytdscf import BasInfo, Model, Simulator, __version__
__version__

'1.1.0'

2. Set primitive basis as HO-DVR

• FBR = Finite Basis Representation

  • using analytic integration

• DVR = Discrete Variational Representation

  • using numerical integration

  • diagonal potential operator

from collections import Counter
from math import sqrt

import numpy as np
import sympy
from discvar import HarmonicOscillator as HO
from pympo import AssignManager, OpSite, SumOfProducts
from pympo.utils import export_npz

from pytdscf.dvr_operator_cls import TensorOperator
from pytdscf.hamiltonian_cls import TensorHamiltonian
from pytdscf.potentials.h2o_potential import k_orig

backend = "numpy"

ndim = 3
freqs = [sqrt(k_orig[(k, k)]) for k in range(1, ndim + 1)]  # a.u.
freqs

[0.007556564680635903, 0.017105918773130724, 0.017567584215986715]

nprims = [9] * ndim  # Number of primitive basis

basis = [
    HO(nprim, omega, units="a.u.")
    for nprim, omega in zip(nprims, freqs, strict=True)
]
basinfo = BasInfo([basis])  # Set basis information object

MPS-MCTDH wavefunction

\[\begin{split}|\Psi_{\rm{MPS-MCTDH}}\rangle = \sum_{\mathbf \{j\}}\sum_{\mathbf \{\tau\}} a\substack{j_1 \\ 1\tau_1}a\substack{j_2 \\ \tau_1\tau_2} \cdots a\substack{j_f \\ \tau_{f-1}1} |\varphi_{j_1}^{(1)}(q_1)\rangle|\varphi_{j_2}^{(2)}(q_2)\rangle \cdots|\varphi_{j_f}^{(f)}(q_f)\rangle\end{split}\]

where SPF is

\[\varphi_{j_p}^{(p)}(q_p) = \sum_{i_p=1}^{n_p} c_{i_p}^{j_p}\chi_{i_p}^{(p)}(q_p) \; (j_p = 1,2,\ldots, N_p)\]

Here, select \(\{\chi_{i_p}^{(p)}(q_p)\}\) as Harmonic Oscillator eigenfunction. See detail in documenation. Here one defines \(n_p\) = 9, \(N_p\) = 9. (i.e., \(\phi=\chi\), so-called “standard method (SM)”)

NOTE

• In MPS, \(n = N\) (SM) is usually better than \(n < M\) (MCTDH). Only when using a laptop computer, MCTDH sometimes works better. (because the required RAM in MCTDH is smaller than SM.)

3. Set Hamiltonian (Polynomial Function)

• Here, one uses pre-calculated Polyonimal PES.

• k_orig is the coefficients of Taylor expansion of PES in reference geometry by mass-weighted coordinate \(Q_i\) (The unit of coordinate is NOT AMU but \(\sqrt{m_e}a_0\), i.e. atomic units)

  \[V - V_0 = \frac{k_{11}}{2!} Q_1^2 + \frac{k_{22}}{2!} Q_2^2 + \frac{k_{33}}{2!} Q_3^2 + \frac{k_{111}}{3!} Q_1^3 + \frac{k_{122}}{2!} Q_1Q_2^2 + \cdots\]

k_orig

defaultdict(float,
            {(1, 1): 5.7101669772633975e-05,
             (2, 2): 0.00029261245707294615,
             (3, 3): 0.0003086200151857856,
             (1, 1, 1): -8.973542624563865e-07,
             (2, 2, 2): -1.8571147341445975e-05,
             (1, 2, 2): 5.028987089424822e-07,
             (1, 1, 2): 1.2870557913839666e-06,
             (1, 3, 3): 2.0063268625796784e-06,
             (2, 3, 3): -1.8853947560756764e-05,
             (1, 1, 1, 1): -2.2778131948543168e-08,
             (2, 2, 2, 2): 1.042951948572713e-06,
             (3, 3, 3, 3): 1.1133748664915738e-06,
             (1, 2, 2, 2): -8.193988329963448e-08,
             (1, 1, 2, 2): -1.852073688081903e-07,
             (1, 1, 1, 2): 5.750959195886642e-08,
             (1, 1, 3, 3): -2.1211138514059556e-07,
             (2, 2, 3, 3): 1.0721581542221527e-06,
             (1, 2, 3, 3): -1.256574051408931e-07})

Setup one particle operator

# Define p^2_i operators
p2_list = [basis_i.get_2nd_derivative_matrix_dvr() for basis_i in basis]
p2_ops = [
    OpSite(r"\hat{p}^2" + f"_{i + 1}", i, value=p2_list[i])
    for i in range(0, ndim)
]

# Define q_i, q^2_i, q^3_i, q^4_i operators
q1_list = [np.array(basis_i.get_grids()) for basis_i in basis]  # DVR grids
q2_list = [q1_int**2 for q1_int in q1_list]
q3_list = [q1_int**3 for q1_int in q1_list]
q4_list = [q1_int**4 for q1_int in q1_list]

q1_ops = [
    OpSite(r"\hat{q}_{" + f"{i + 1}" + "}", i, value=q1_list[i])
    for i in range(0, ndim)
]
q2_ops = [
    OpSite(r"\hat{q}^2_{" + f"{i + 1}" + "}", i, value=q2_list[i])
    for i in range(0, ndim)
]
q3_ops = [
    OpSite(r"\hat{q}^3_{" + f"{i + 1}" + "}", i, value=q3_list[i])
    for i in range(0, ndim)
]
q4_ops = [
    OpSite(r"\hat{q}^4_{" + f"{i + 1}" + "}", i, value=q4_list[i])
    for i in range(0, ndim)
]

qn_list = [q1_ops, q2_ops, q3_ops, q4_ops]

Setup potential and kinetic operator

kin_sop = SumOfProducts()

for p2_op in p2_ops:
    kin_sop += -sympy.Rational(1 / 2) * p2_op

kin_sop.symbol

$\displaystyle - \frac{\hat{p}^2_1}{2} - \frac{\hat{p}^2_2}{2} - \frac{\hat{p}^2_3}{2}$

pot_sop = SumOfProducts()
subs = {}  # substitutions afterwards

for key, coef in k_orig.items():
    if coef == 0:
        continue
    count = Counter(key)
    op = 1
    k = sympy.Symbol(f"k_{key}")
    op = 1
    for isite, order in count.items():
        if order > 0:
            op *= qn_list[order - 1][isite - 1]
        if order > 1:
            op /= sympy.factorial(order)
    pot_sop += k * op
    subs[k] = coef

pot_sop = pot_sop.simplify()
pot_sop.symbol

$\displaystyle \frac{k_{(1, 1)} \hat{q}^2_{1}}{2} + \frac{k_{(1, 1, 1)} \hat{q}^3_{1}}{6} + \frac{k_{(1, 1, 1, 1)} \hat{q}^4_{1}}{24} + \frac{k_{(1, 1, 1, 2)} \hat{q}^3_{1} \hat{q}_{2}}{6} + \frac{k_{(1, 1, 2)} \hat{q}^2_{1} \hat{q}_{2}}{2} + \frac{k_{(1, 1, 2, 2)} \hat{q}^2_{1} \hat{q}^2_{2}}{4} + \frac{k_{(1, 1, 3, 3)} \hat{q}^2_{1} \hat{q}^2_{3}}{4} + \frac{k_{(1, 2, 2)} \hat{q}_{1} \hat{q}^2_{2}}{2} + \frac{k_{(1, 2, 2, 2)} \hat{q}_{1} \hat{q}^3_{2}}{6} + \frac{k_{(1, 2, 3, 3)} \hat{q}_{1} \hat{q}_{2} \hat{q}^2_{3}}{2} + \frac{k_{(1, 3, 3)} \hat{q}_{1} \hat{q}^2_{3}}{2} + \frac{k_{(2, 2)} \hat{q}^2_{2}}{2} + \frac{k_{(2, 2, 2)} \hat{q}^3_{2}}{6} + \frac{k_{(2, 2, 2, 2)} \hat{q}^4_{2}}{24} + \frac{k_{(2, 2, 3, 3)} \hat{q}^2_{2} \hat{q}^2_{3}}{4} + \frac{k_{(2, 3, 3)} \hat{q}_{2} \hat{q}^2_{3}}{2} + \frac{k_{(3, 3)} \hat{q}^2_{3}}{2} + \frac{k_{(3, 3, 3, 3)} \hat{q}^4_{3}}{24}$

Setup MPO

am_kin = AssignManager(kin_sop)
am_kin.assign()
display(*am_kin.Wsym)
kin_mpo = am_kin.numerical_mpo(subs=subs)

$\displaystyle \left[\begin{matrix}\hat{p}^2_1 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}- \frac{1}{2} & 0\\- \frac{\hat{p}^2_2}{2} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1\\- \frac{\hat{p}^2_3}{2}\end{matrix}\right]$

am_pot = AssignManager(pot_sop)
am_pot.assign()
display(*am_pot.Wsym)
W_prod = sympy.Mul(*am_pot.Wsym)
print(*[f"W{i}" for i in range(am_pot.ndim)], "=")
display(W_prod[0].expand())
pot_mpo = am_pot.numerical_mpo(subs=subs)

$\displaystyle \left[\begin{matrix}\hat{q}^3_{1} & \hat{q}_{1} & \hat{q}^4_{1} & \hat{q}^2_{1} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & \frac{k_{(1, 1, 1)}}{6} + \frac{k_{(1, 1, 1, 2)} \hat{q}_{2}}{6} & 0\\\frac{k_{(1, 2, 3, 3)} \hat{q}_{2}}{2} + \frac{k_{(1, 3, 3)}}{2} & \frac{k_{(1, 2, 2)} \hat{q}^2_{2}}{2} + \frac{k_{(1, 2, 2, 2)} \hat{q}^3_{2}}{6} & 0\\0 & \frac{k_{(1, 1, 1, 1)}}{24} & 0\\\frac{k_{(1, 1, 3, 3)}}{4} & \frac{k_{(1, 1)}}{2} + \frac{k_{(1, 1, 2)} \hat{q}_{2}}{2} + \frac{k_{(1, 1, 2, 2)} \hat{q}^2_{2}}{4} & 0\\\frac{k_{(2, 2, 3, 3)} \hat{q}^2_{2}}{4} + \frac{k_{(2, 3, 3)} \hat{q}_{2}}{2} + \frac{k_{(3, 3)}}{2} & \frac{k_{(2, 2)} \hat{q}^2_{2}}{2} + \frac{k_{(2, 2, 2)} \hat{q}^3_{2}}{6} + \frac{k_{(2, 2, 2, 2)} \hat{q}^4_{2}}{24} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}\hat{q}^2_{3}\\1\\\frac{k_{(3, 3, 3, 3)} \hat{q}^4_{3}}{24}\end{matrix}\right]$

W0 W1 W2 =

$\displaystyle \frac{k_{(1, 1)} \hat{q}^2_{1}}{2} + \frac{k_{(1, 1, 1)} \hat{q}^3_{1}}{6} + \frac{k_{(1, 1, 1, 1)} \hat{q}^4_{1}}{24} + \frac{k_{(1, 1, 1, 2)} \hat{q}^3_{1} \hat{q}_{2}}{6} + \frac{k_{(1, 1, 2)} \hat{q}^2_{1} \hat{q}_{2}}{2} + \frac{k_{(1, 1, 2, 2)} \hat{q}^2_{1} \hat{q}^2_{2}}{4} + \frac{k_{(1, 1, 3, 3)} \hat{q}^2_{1} \hat{q}^2_{3}}{4} + \frac{k_{(1, 2, 2)} \hat{q}_{1} \hat{q}^2_{2}}{2} + \frac{k_{(1, 2, 2, 2)} \hat{q}_{1} \hat{q}^3_{2}}{6} + \frac{k_{(1, 2, 3, 3)} \hat{q}_{1} \hat{q}_{2} \hat{q}^2_{3}}{2} + \frac{k_{(1, 3, 3)} \hat{q}_{1} \hat{q}^2_{3}}{2} + \frac{k_{(2, 2)} \hat{q}^2_{2}}{2} + \frac{k_{(2, 2, 2)} \hat{q}^3_{2}}{6} + \frac{k_{(2, 2, 2, 2)} \hat{q}^4_{2}}{24} + \frac{k_{(2, 2, 3, 3)} \hat{q}^2_{2} \hat{q}^2_{3}}{4} + \frac{k_{(2, 3, 3)} \hat{q}_{2} \hat{q}^2_{3}}{2} + \frac{k_{(3, 3)} \hat{q}^2_{3}}{2} + \frac{k_{(3, 3, 3, 3)} \hat{q}^4_{3}}{24}$

Setup Hamiltonian

# Save MPO: list[np.ndarray] for the succeeding propagation
export_npz(pot_mpo, "h2o_pot_mpo.npz")
export_npz(kin_mpo, "h2o_kin_mpo.npz")

# for potential MPO each core W has a single site "leg"
# | | |
# W-W-W
# while for kinetic MPO each core W has two site "legs"
# | | |
# W-W-W
# | | |

potential = [
    [
        {
            (tuple((i,) for i in range(0, ndim))): TensorOperator(mpo=pot_mpo),
            (tuple((i, i) for i in range(0, ndim))): TensorOperator(
                mpo=kin_mpo
            ),
        }
    ]
]

H = TensorHamiltonian(
    ndof=len(basis), potential=potential, kinetic=None, backend=backend
)

operators = {"hamiltonian": H}

4. Set wavefunction (MPS) and All Model

• m_aux_max is a bond dimension of MPS (maximum rank of auxiliary index \(\tau_p\))

model = Model(basinfo, operators)
model.m_aux_max = 9

5. Execute Calculation

• time step width is defined by stepsize=0.1 fs

F.Y.I., See also documentation about Simulator

NOTE

• Runtype cannot be rebound. If you change the runtype, you should restart the kernel.

• Improved relaxation, i.e., diagonalization-based variational optimization, is much faster than pure imaginary time evolution.

• When simulating larger systems, (f>6, m>10), contraction of tensors can be overhead, in such a situation JAX, especially supporting GPU, is recommended to use as a backend.

jobname = "h2o_polynomial"
simulator = Simulator(jobname, model, ci_type="MPS", backend="Numpy", verbose=4)
simulator.relax(savefile_ext="_gs", maxstep=20, stepsize=0.1)

17:04:31 | INFO | 
     ____     __________   .____ ____   _____
    / _  |   /__  __/ _ \ / ___ / _  \ / ___/
   / /_) /_  __/ / / / ||/ /__ / / )_// /__
  /  ___/ / / / / / / / |.__  / |  __/ ___/
 /  /  / /_/ / / / /_/ /___/ /| \_/ / /
/__/   \__, /_/ /_____/_____/ \____/_/
      /____/

17:04:31 | INFO | Log file is ./h2o_polynomial_relax/main.log
17:04:31 | INFO | Wave function is saved in wf_h2o_polynomial_gs.pkl
17:04:31 | INFO | Start initial step    0.000 [fs]

17:04:31 | INFO | End     0 step; propagated    0.100 [fs]; AVG Krylov iteration: 3.00
17:04:32 | INFO | Saved wavefunction    0.900 [fs]
17:04:32 | INFO | Saved wavefunction    1.900 [fs]
17:04:32 | INFO | End    19 step; propagated    1.900 [fs]; AVG Krylov iteration: 0.00
17:04:32 | INFO | End simulation and save wavefunction
17:04:32 | INFO | Wave function is saved in wf_h2o_polynomial_gs.pkl

(0.020855716347418774, <pytdscf.wavefunction.WFunc at 0x7f1557d90e60>)

6. Check Log file

See h2o_polynomial_relax/main.log, which is defined in jobname

!tail h2o_polynomial_relax/main.log

2025-06-18 17:04:32 | DEBUG | pytdscf.properties:_export_properties:395 - | pop 1.0000 | ene[eV]:  0.5675130 | time[fs]:    1.400 | elapsed[sec]:     0.17 | ci:  0.2  (ci_exp:  0.1|ci_rnm:  0.0|ci_etc:  0.0 d) |    0 MFLOPS (  0.0 s)
2025-06-18 17:04:32 | DEBUG | pytdscf.properties:_export_properties:395 - | pop 1.0000 | ene[eV]:  0.5675130 | time[fs]:    1.500 | elapsed[sec]:     0.18 | ci:  0.2  (ci_exp:  0.1|ci_rnm:  0.0|ci_etc:  0.0 d) |    0 MFLOPS (  0.0 s)
2025-06-18 17:04:32 | DEBUG | pytdscf.properties:_export_properties:395 - | pop 1.0000 | ene[eV]:  0.5675130 | time[fs]:    1.600 | elapsed[sec]:     0.18 | ci:  0.2  (ci_exp:  0.1|ci_rnm:  0.0|ci_etc:  0.0 d) |    0 MFLOPS (  0.0 s)
2025-06-18 17:04:32 | DEBUG | pytdscf.properties:_export_properties:395 - | pop 1.0000 | ene[eV]:  0.5675130 | time[fs]:    1.700 | elapsed[sec]:     0.19 | ci:  0.2  (ci_exp:  0.2|ci_rnm:  0.0|ci_etc:  0.0 d) |    0 MFLOPS (  0.0 s)
2025-06-18 17:04:32 | DEBUG | pytdscf.properties:_export_properties:395 - | pop 1.0000 | ene[eV]:  0.5675130 | time[fs]:    1.800 | elapsed[sec]:     0.19 | ci:  0.2  (ci_exp:  0.2|ci_rnm:  0.1|ci_etc:  0.0 d) |    0 MFLOPS (  0.0 s)
2025-06-18 17:04:32 | INFO | pytdscf.simulator_cls:_execute:390 - Saved wavefunction    1.900 [fs]
2025-06-18 17:04:32 | DEBUG | pytdscf.properties:_export_properties:395 - | pop 1.0000 | ene[eV]:  0.5675130 | time[fs]:    1.900 | elapsed[sec]:     0.20 | ci:  0.2  (ci_exp:  0.2|ci_rnm:  0.1|ci_etc:  0.0 d) |    0 MFLOPS (  0.0 s)
2025-06-18 17:04:32 | INFO | pytdscf.simulator_cls:_execute:433 - End    19 step; propagated    1.900 [fs]; AVG Krylov iteration: 0.00
2025-06-18 17:04:32 | INFO | pytdscf.simulator_cls:_execute:434 - End simulation and save wavefunction
2025-06-18 17:04:32 | INFO | pytdscf.simulator_cls:save_wavefunction:558 - Wave function is saved in wf_h2o_polynomial_gs.pkl

Vibrational ground state energy is found to be ``0.5675130`` eV!