Example 2: Operation of dipole moment operator to H2O vibrational ground state

run type	restart	wavefunction	backend	Basis	max iteration
operation	True (file suffix _gs)	MPS-SM (restart)	Numpy	HO-DVR	10

1. Import modules

• Required in any calculations

from pytdscf import BasInfo, Model, Simulator, __version__
__version__

'1.2.0'

2. Set DVR primitive basis

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 define \(n_p\) = 9, \(N_p\) = 9.

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.)

• Basis information must be the same as restart one.

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, get_eye_site

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.

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

basis = [
    HO(nprim, omega, units="a.u.")
    for nprim, omega in zip(nprims, freqs, strict=True)
]
ndim = len(basis)  # Number of degree of freedom, H2O has 3 DOFs
basinfo = BasInfo([basis])  # Set basis information object

3. Set Dipole Operator (Polynomial Function)

Here, one use pre-calculated Polyonimal PES and DMS.

from pytdscf.potentials.h2o_dipole import mu

mu  # Dipole has (x,y,z) elements

{(): [-1.69908e-15, 1.24913e-14, -1.93795],
 (3,): [-2.20831e-17, 0.00853527, -8.32759e-16],
 (2,): [1.50857e-17, 2.08217e-15, -0.00326347],
 (1,): [6.37588e-18, 8.65662e-16, 0.0142383],
 (3, 3): [3.5274e-18, -1.35302e-15, -2.31565e-05],
 (2, 3): [3.46044e-18, -0.000294259, -7.3169e-16],
 (2, 2): [-1.5306e-18, -1.42229e-15, 0.00020955],
 (1, 3): [1.45261e-17, 0.000327409, -2.99469e-17],
 (1, 2): [3.90656e-18, 1.26166e-16, -0.000112968],
 (1, 1): [-6.45481e-18, 6.79098e-16, 0.000192831],
 (3, 3, 3): [-1.34711e-21, 7.33335e-06, 9.41511e-22],
 (2, 3, 3): [2.2067e-22, -3.92968e-22, 3.0608e-06],
 (1, 3, 3): [-2.55725e-22, 4.55392e-22, -3.54702e-06],
 (2, 2, 3): [6.16547e-22, -3.35633e-06, -4.3091e-22],
 (2, 2, 2): [1.69378e-22, -3.01627e-22, 2.34936e-06],
 (1, 2, 2): [3.17065e-22, -5.64628e-22, 4.39783e-06],
 (1, 1, 3): [-1.08836e-21, 5.92476e-06, 7.60666e-22],
 (1, 1, 2): [-2.92033e-23, 5.20049e-23, -4.05062e-07],
 (1, 1, 1): [5.60185e-22, -9.97572e-22, 7.77e-06]}

Setup one particle operator

q1_list = [np.array(basis_i.get_grids()) for basis_i in basis]
q2_list = [q1_int**2 for q1_int in q1_list]
q3_list = [q1_int**3 for q1_int in q1_list]

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

qn_list = [q1_ops, q2_ops, q3_ops]

subs = {}
Evec = np.array([1.0, 1.0, 1.0]) * 1e-02

Setup Potential Operator

muE_sop = SumOfProducts()

for key, dmom in mu.items():
    if (coef := np.dot(dmom, Evec)) == 0.0:
        continue
    count = Counter(key)
    muE = sympy.Symbol(r"\mu_{" + f"{key}" + r"}\cdot E")
    subs[muE] = coef
    if key == ():
        # Scalar term
        muE_sop += muE * get_eye_site(0, q1_list[0].shape[0])
        continue
    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)
    muE_sop += muE * op

muE_sop = muE_sop.simplify()
muE_sop.symbol

$\displaystyle \hat{1}_0 \mu_{()}\cdot E + \frac{\mu_{(1, 1)}\cdot E q^{2}_{1}}{2} + \frac{\mu_{(1, 1, 1)}\cdot E q^{3}_{1}}{6} + \frac{\mu_{(1, 1, 2)}\cdot E q^{2}_{1} q^{1}_{2}}{2} + \frac{\mu_{(1, 1, 3)}\cdot E q^{2}_{1} q^{1}_{3}}{2} + \mu_{(1, 2)}\cdot E q^{1}_{1} q^{1}_{2} + \frac{\mu_{(1, 2, 2)}\cdot E q^{1}_{1} q^{2}_{2}}{2} + \mu_{(1, 3)}\cdot E q^{1}_{1} q^{1}_{3} + \frac{\mu_{(1, 3, 3)}\cdot E q^{1}_{1} q^{2}_{3}}{2} + \mu_{(1,)}\cdot E q^{1}_{1} + \frac{\mu_{(2, 2)}\cdot E q^{2}_{2}}{2} + \frac{\mu_{(2, 2, 2)}\cdot E q^{3}_{2}}{6} + \frac{\mu_{(2, 2, 3)}\cdot E q^{2}_{2} q^{1}_{3}}{2} + \mu_{(2, 3)}\cdot E q^{1}_{2} q^{1}_{3} + \frac{\mu_{(2, 3, 3)}\cdot E q^{1}_{2} q^{2}_{3}}{2} + \mu_{(2,)}\cdot E q^{1}_{2} + \frac{\mu_{(3, 3)}\cdot E q^{2}_{3}}{2} + \frac{\mu_{(3, 3, 3)}\cdot E q^{3}_{3}}{6} + \mu_{(3,)}\cdot E q^{1}_{3}$

Setup MPO

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

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

$\displaystyle \left[\begin{matrix}\mu_{()}\cdot E + \frac{\mu_{(2, 2)}\cdot E q^{2}_{2}}{2} + \frac{\mu_{(2, 2, 2)}\cdot E q^{3}_{2}}{6} + \mu_{(2,)}\cdot E q^{1}_{2} & \frac{\mu_{(2, 3, 3)}\cdot E q^{1}_{2}}{2} + \frac{\mu_{(3, 3)}\cdot E}{2} & \frac{\mu_{(2, 2, 3)}\cdot E q^{2}_{2}}{2} + \mu_{(2, 3)}\cdot E q^{1}_{2} + \mu_{(3,)}\cdot E & 1\\\frac{\mu_{(1, 1, 1)}\cdot E}{6} & 0 & 0 & 0\\\frac{\mu_{(1, 1)}\cdot E}{2} + \frac{\mu_{(1, 1, 2)}\cdot E q^{1}_{2}}{2} & 0 & \frac{\mu_{(1, 1, 3)}\cdot E}{2} & 0\\\mu_{(1, 2)}\cdot E q^{1}_{2} + \frac{\mu_{(1, 2, 2)}\cdot E q^{2}_{2}}{2} + \mu_{(1,)}\cdot E & \frac{\mu_{(1, 3, 3)}\cdot E}{2} & \mu_{(1, 3)}\cdot E & 0\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1\\q^{2}_{3}\\q^{1}_{3}\\\frac{\mu_{(3, 3, 3)}\cdot E q^{3}_{3}}{6}\end{matrix}\right]$

W0 W1 W2 =

$\displaystyle \mu_{()}\cdot E + \frac{\mu_{(1, 1)}\cdot E q^{2}_{1}}{2} + \frac{\mu_{(1, 1, 1)}\cdot E q^{3}_{1}}{6} + \frac{\mu_{(1, 1, 2)}\cdot E q^{2}_{1} q^{1}_{2}}{2} + \frac{\mu_{(1, 1, 3)}\cdot E q^{2}_{1} q^{1}_{3}}{2} + \mu_{(1, 2)}\cdot E q^{1}_{1} q^{1}_{2} + \frac{\mu_{(1, 2, 2)}\cdot E q^{1}_{1} q^{2}_{2}}{2} + \mu_{(1, 3)}\cdot E q^{1}_{1} q^{1}_{3} + \frac{\mu_{(1, 3, 3)}\cdot E q^{1}_{1} q^{2}_{3}}{2} + \mu_{(1,)}\cdot E q^{1}_{1} + \frac{\mu_{(2, 2)}\cdot E q^{2}_{2}}{2} + \frac{\mu_{(2, 2, 2)}\cdot E q^{3}_{2}}{6} + \frac{\mu_{(2, 2, 3)}\cdot E q^{2}_{2} q^{1}_{3}}{2} + \mu_{(2, 3)}\cdot E q^{1}_{2} q^{1}_{3} + \frac{\mu_{(2, 3, 3)}\cdot E q^{1}_{2} q^{2}_{3}}{2} + \mu_{(2,)}\cdot E q^{1}_{2} + \frac{\mu_{(3, 3)}\cdot E q^{2}_{3}}{2} + \frac{\mu_{(3, 3, 3)}\cdot E q^{3}_{3}}{6} + \mu_{(3,)}\cdot E q^{1}_{3}$

Setup Operator

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

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

operators = {"hamiltonian": H}

4. Set wave function (MPS) and All Model

• m_aux_max is MPS bond dimenison (maximum of auxiliary index \(\tau_p\))

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

5. Execute Calculation

! ls wf_h2o_polynomial_gs.pkl

wf_h2o_polynomial_gs.pkl

F.Y.I., See also about Simulator

This run type prepare \(|\Psi_{\mu}\rangle\) by variationally optimizing

\[\langle \delta\Psi_{\mu}|\Psi_{\mu} - \hat{\mu}\Psi_{\mathrm{gs}}\rangle=0\]

where \(\Psi_{\mathrm{gs}}\) is a vibrational ground state wavefunction.

NOTE

• Runtype cannnot rebind. If you want to change runtype, you should restart kernel.

jobname = "h2o_polynomial"
simulator = Simulator(jobname, model, backend="numpy", verbose=4)
simulator.operate(
    loadfile_ext="_gs", savefile_ext="_dipole", restart=True, maxstep=10
)

13:53:45 | INFO | Log file is ./h2o_polynomial_operate/main.log
13:53:45 | INFO | Wave function is loaded from wf_h2o_polynomial_gs.pkl
13:53:45 | INFO | Start: apply operator to wave function
13:53:45 | INFO | Wave function is saved in wf_h2o_polynomial_dipole.pkl
13:53:45 | INFO | End  : apply operator to wave function

(0.019386775177213933, <pytdscf.wavefunction.WFunc at 0x7f71af476ea0>)

6. Check Log file

See h2o_polynomial_operate/main.log, which is defined by jobname.

!tail h2o_polynomial_operate/main.log

2025-09-23 13:53:45 | INFO | pytdscf.simulator_cls:get_initial_wavefunction:507 - Wave function is loaded from wf_h2o_polynomial_gs.pkl
2025-09-23 13:53:45 | INFO | pytdscf.simulator_cls:_execute:357 - Start: apply operator to wave function
2025-09-23 13:53:45 | DEBUG | pytdscf.wavefunction:apply_dipole:338 - ----------------------------------------
iterations: 0 norm: 0.019386775177213912
2025-09-23 13:53:45 | DEBUG | pytdscf.wavefunction:_is_converged:297 - convergence : 0.9977392532567779
2025-09-23 13:53:45 | DEBUG | pytdscf.wavefunction:apply_dipole:338 - ----------------------------------------
iterations: 1 norm: 0.019386775177213933
2025-09-23 13:53:45 | DEBUG | pytdscf.wavefunction:_is_converged:297 - convergence : 0.9999999999999999
2025-09-23 13:53:45 | INFO | pytdscf.simulator_cls:save_wavefunction:589 - Wave function is saved in wf_h2o_polynomial_dipole.pkl
2025-09-23 13:53:45 | INFO | pytdscf.simulator_cls:_execute:360 - End  : apply operator to wave function