Example 11 Time-dependent reduced density
This example takes a while. I recommend execution on HPC and select JAX as wrapper
This example is the same model as Example 10.
The differences are
Example 10 employs electronic states as a direct product state.
Therefore, two MPS are simulated.
Example 11 employs electronic states as a MPS site.
Therefore, single but longer MPS is simulated.
Example 10 construct MPO by SVD while Example 11 construct MPO analytically (use symbolic scheme).
Therefore, MPO constructed in Example 11 is free from numerical error.
Example 10 is slightly accurate but Example 11 is much computationally feasible when we simulate a couple of or dozens of electronic states.
Because direct product electronic state is accurate but suffers from exponential complexity.
Furthermore, single MPS embedded in a single Array is a suitable architecture for computation.
[1]:
!nvidia-smi
zsh:1: command not found: nvidia-smi
[2]:
import platform
import sys
import pytdscf
print(sys.version)
print(f"pytdscf version = {pytdscf.__version__}")
print(platform.platform())
import jax.extend
print(jax.extend.backend.get_backend().platform)
3.13.2 (main, Feb 4 2025, 14:51:09) [Clang 16.0.0 (clang-1600.0.26.6)]
pytdscf version = 1.1.0
macOS-15.4.1-arm64-arm-64bit-Mach-O
cpu
[3]:
import matplotlib.pyplot as plt
import numpy as np
import random
import sympy
from discvar import Exponential, HarmonicOscillator, Sine
from pympo import AssignManager, OpSite, SumOfProducts
from pytdscf import BasInfo, Exciton, Model, Simulator, units
from pytdscf.dvr_operator_cls import TensorOperator
from pytdscf.hamiltonian_cls import TensorHamiltonian
References:
Morse potential is our original demonstration.
The second quantized representation is given by
Setup numerical parameters
[4]:
backend = "jax"
use_bath = True
Ω = 1532 / units.au_in_cm1
K = 0.19 / units.au_in_eV
De = 2.285278e-1 # Dissociation energy 2.285278E-1 Eh = 600 kJ /mol
Q_ref = K / (Ω**1.5)
I = sympy.Symbol("I")
W0 = sympy.Symbol("W_0")
W1 = sympy.Symbol("W_1")
λ = sympy.Symbol("lambda")
E = sympy.Symbol("E")
subs = {}
subs[I] = 1 / (
4.84e-04 / units.au_in_eV
) # inertial mass of C=C torsion (isomerization)
subs[W0] = 3.56 / units.au_in_eV # isomerization term
subs[W1] = 1.19 / units.au_in_eV # isomerization term
subs[λ] = 0.19 / units.au_in_eV * np.sqrt(Ω) # Linear vibronic coupling term
subs[E] = 2.58 / units.au_in_eV - K**2 / 2 / Ω
if use_bath:
bath_ω = (
np.array(
[
792.8,
842.8,
866.2,
882.4,
970.3,
976.0,
997.0,
1017.1,
1089.6,
1189.0,
1214.7,
1238.1,
1267.9,
1317.0,
1359.0,
1389.0,
1428.4,
1434.9,
1451.8,
1572.8,
1612.1,
1629.2,
1659.1,
]
)
/ units.au_in_cm1
)
dimless_displace = np.array(
[
0.175,
0.2,
0.175,
0.225,
0.55,
0.3,
0.33,
0.45,
0.125,
0.175,
0.44,
0.5,
0.475,
0.238,
0.25,
0.25,
0.25,
0.225,
0.225,
0.25,
0.225,
0.125,
0.225,
]
) # c=κ/ω
ω_symbols = [None, None, None]
c_symbols = [None, None, None]
for i, (ω, c) in enumerate(zip(bath_ω, dimless_displace, strict=True)):
ω_symbols.append(sympy.Symbol(f"omega_{i + 3}"))
c_symbols.append(sympy.Symbol(f"c_{i + 3}"))
subs[ω_symbols[-1]] = ω
subs[c_symbols[-1]] = c * (ω**1.5)
Setup basis for wavefunction
[5]:
basis = [
Exciton(nstate=2),
Sine(ngrid=2**6, length=175.0, x0=-55.0, units="a.u."),
Exponential(ngrid=2**8 - 1, length=2 * np.pi, x0=-np.pi),
]
if use_bath:
for ω, c in zip(bath_ω, dimless_displace, strict=False):
ho_grids = HarmonicOscillator(10, ω, units="a.u.").get_grids()
left_grid = ho_grids[0] - c / np.sqrt(ω)
length = ho_grids[-1] - ho_grids[0] + c / np.sqrt(ω) ## c=κ/ω
basis.append(
Sine(ngrid=2**5, length=length, x0=left_grid, units="a.u.")
)
basinfo = BasInfo([basis])
ndim = len(basis)
Setup one particle operator
[6]:
q1_list = [None] + [np.array(basis_i.get_grids()) for basis_i in basis[1:]]
q2_list = [None] + [q1_int**2 for q1_int in q1_list[1:]]
p2_list = [None] + [
basis_i.get_2nd_derivative_matrix_dvr() for basis_i in basis[1:]
]
morse_0 = De * (1.0 - np.exp(-1.0 * np.sqrt(Ω**2 / 2.0 / De) * q1_list[1])) ** 2
morse_1 = (
De
* (1.0 - np.exp(-1.0 * np.sqrt(Ω**2 / 2.0 / De) * (q1_list[1] - Q_ref)))
** 2
)
one_cos = 1.0 - np.cos(q1_list[2])
a = basis[0].get_annihilation_matrix()
adag = basis[0].get_creation_matrix()
p2_ops = [None] + [
OpSite(f"p^2_{i}", i, value=p2_list[i]) for i in range(1, ndim)
]
q2_ops = [None] + [
OpSite(f"q^2_{i}", i, value=q2_list[i]) for i in range(1, ndim)
]
q1_ops = [None] + [
OpSite(f"q_{i}", i, value=q1_list[i]) for i in range(1, ndim)
]
morse_op_0 = OpSite("M^{(0)}_1", 1, value=morse_0)
morse_op_1 = OpSite("M^{(1)}_1", 1, value=morse_1)
cos_op = OpSite(r"(1-\cos \theta_{2})", 2, value=one_cos)
a_op = OpSite("a_0", 0, value=a)
adag_op = OpSite(r"a^\dagger_0", 0, value=adag)
Setup potential and kinetic operator
[7]:
pot_sop = SumOfProducts()
pot_sop += (-E / 2) * a_op * adag_op
pot_sop += a_op * adag_op * morse_op_0
pot_sop += (W0 / 2)* a_op * adag_op * cos_op
for i in range(3, ndim):
pot_sop += (ω_symbols[i] ** 2 / 2) * a_op * adag_op * q2_ops[i]
pot_sop += (E / 2) * adag_op * a_op
pot_sop += adag_op * a_op * morse_op_1
pot_sop += (-W1 / 2) * adag_op * a_op * cos_op
for i in range(3, ndim):
pot_sop += (ω_symbols[i] ** 2 / 2) * adag_op * a_op * q2_ops[i]
pot_sop += (c_symbols[i]) * adag_op * a_op * q1_ops[i]
pot_sop += λ * (adag_op + a_op) * q1_ops[1]
pot_sop = pot_sop.simplify()
pot_sop.symbol
[7]:
[8]:
kin_sop = SumOfProducts()
kin_sop += (-1 / 2) * (a_op * adag_op + adag_op * a_op) * p2_ops[1]
kin_sop += (-1 / 2 / I) * (a_op * adag_op + adag_op * a_op) * p2_ops[2]
for i in range(3, ndim):
kin_sop += (-1 / 2) * (a_op * adag_op + adag_op * a_op) * p2_ops[i]
kin_sop = kin_sop.simplify()
kin_sop.symbol
[8]:
Setup MPO
[9]:
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)
W0 W1 W2 W3 W4 W5 W6 W7 W8 W9 W10 W11 W12 W13 W14 W15 W16 W17 W18 W19 W20 W21 W22 W23 W24 W25 =
[10]:
am_kin = AssignManager(kin_sop)
am_kin.assign()
display(*am_kin.Wsym)
W_prod = sympy.Mul(*am_kin.Wsym)
print(*[f"W{i}" for i in range(am_kin.ndim)], "=")
display(W_prod[0].expand())
kin_mpo = am_kin.numerical_mpo(subs=subs)
W0 W1 W2 W3 W4 W5 W6 W7 W8 W9 W10 W11 W12 W13 W14 W15 W16 W17 W18 W19 W20 W21 W22 W23 W24 W25 =
Setup Hamiltonian
[11]:
# MPO has legs on (0,1,2, ... ,f-1) sites. This legs are given by tuple key
potential = [
[
{
(((0, 0),) + tuple(i for i in range(1, ndim))): TensorOperator(
mpo=pot_mpo
)
}
]
] # key is ((0,0), 1, 2, ..., ndim-1)
kinetic = [
[{tuple((i, i) for i in range(ndim)): TensorOperator(mpo=kin_mpo)}]
] # key is ((0, 0), (1, 1), ..., (ndim-1, ndim-1))
H = TensorHamiltonian(
ndof=len(basis), potential=potential, kinetic=kinetic, backend=backend
)
operators = {"hamiltonian": H}
Setup Model (basis, operators, initial states)
[12]:
model = Model(basinfo=basinfo, operators=operators)
model.m_aux_max = 20
[13]:
def get_initial_state(ω, x):
weight = np.exp(-ω / 2 * (np.array(x)) ** 2)
weight /= np.linalg.norm(weight)
return weight.tolist()
ω_list = [Ω, 2 / (0.15228**2)]
if use_bath:
ω_list += bath_ω.tolist()
# I wonder where 0.15228 comes from???
# but https://pubs.acs.org/doi/epdf/10.1021/acs.jctc.2c00209 adopts this value
# as an initial Gaussian width of the torsional mode
init_vib = [
get_initial_state(ω, bas.get_grids())
for ω, bas in zip(ω_list, basis[1:], strict=True)
]
model.init_HartreeProduct = [[[0.0, 1.0], *init_vib]]
Execution
[14]:
jobname = "isomerization_exciton"
if use_bath:
jobname += "_bath"
simulator = Simulator(jobname=jobname, model=model, backend=backend, verbose=2)
simulator.propagate(
maxstep=4000,
stepsize=0.05,
reduced_density=(
[(0, 1, 2), (0,)],
40,
), # we want to know diagonal_element of (|0><0| |1><1| |2><2|) and (|0><0|)
energy=False,
autocorr=False,
)
# calculate (0,1) reduced density every 20 steps
00:35:04 | INFO | Log file is ./isomerization_exciton_bath_prop/main.log
00:35:04 | INFO | Wave function is saved in wf_isomerization_exciton_bath.pkl
00:35:04 | INFO | Start initial step 0.000 [fs]
00:35:09 | INFO | End 0 step; propagated 0.050 [fs]; AVG Krylov iteration: 6.46
00:36:49 | INFO | End 100 step; propagated 5.050 [fs]; AVG Krylov iteration: 6.00
00:38:28 | INFO | End 200 step; propagated 10.050 [fs]; AVG Krylov iteration: 5.88
00:40:05 | INFO | End 300 step; propagated 15.050 [fs]; AVG Krylov iteration: 5.85
00:41:40 | INFO | End 400 step; propagated 20.050 [fs]; AVG Krylov iteration: 5.85
00:43:16 | INFO | End 500 step; propagated 25.050 [fs]; AVG Krylov iteration: 5.85
00:44:51 | INFO | End 600 step; propagated 30.050 [fs]; AVG Krylov iteration: 5.85
00:46:28 | INFO | End 700 step; propagated 35.050 [fs]; AVG Krylov iteration: 5.85
00:48:06 | INFO | End 800 step; propagated 40.050 [fs]; AVG Krylov iteration: 5.85
00:49:48 | INFO | End 900 step; propagated 45.050 [fs]; AVG Krylov iteration: 5.85
00:51:26 | INFO | Saved wavefunction 49.950 [fs]
00:51:29 | INFO | End 1000 step; propagated 50.050 [fs]; AVG Krylov iteration: 5.85
00:53:10 | INFO | End 1100 step; propagated 55.050 [fs]; AVG Krylov iteration: 5.85
00:54:51 | INFO | End 1200 step; propagated 60.050 [fs]; AVG Krylov iteration: 5.85
00:56:34 | INFO | End 1300 step; propagated 65.050 [fs]; AVG Krylov iteration: 5.85
00:58:17 | INFO | End 1400 step; propagated 70.050 [fs]; AVG Krylov iteration: 5.85
00:59:58 | INFO | End 1500 step; propagated 75.050 [fs]; AVG Krylov iteration: 5.85
01:01:35 | INFO | End 1600 step; propagated 80.050 [fs]; AVG Krylov iteration: 5.85
01:03:10 | INFO | End 1700 step; propagated 85.050 [fs]; AVG Krylov iteration: 5.85
01:04:46 | INFO | End 1800 step; propagated 90.050 [fs]; AVG Krylov iteration: 5.85
01:06:21 | INFO | End 1900 step; propagated 95.050 [fs]; AVG Krylov iteration: 5.85
01:07:55 | INFO | Saved wavefunction 99.950 [fs]
01:07:57 | INFO | End 2000 step; propagated 100.050 [fs]; AVG Krylov iteration: 5.85
01:09:33 | INFO | End 2100 step; propagated 105.050 [fs]; AVG Krylov iteration: 5.85
01:11:10 | INFO | End 2200 step; propagated 110.050 [fs]; AVG Krylov iteration: 5.77
01:12:49 | INFO | End 2300 step; propagated 115.050 [fs]; AVG Krylov iteration: 5.69
01:14:28 | INFO | End 2400 step; propagated 120.050 [fs]; AVG Krylov iteration: 5.27
01:16:05 | INFO | End 2500 step; propagated 125.050 [fs]; AVG Krylov iteration: 5.19
01:17:41 | INFO | End 2600 step; propagated 130.050 [fs]; AVG Krylov iteration: 5.23
01:19:18 | INFO | End 2700 step; propagated 135.050 [fs]; AVG Krylov iteration: 5.19
01:20:55 | INFO | End 2800 step; propagated 140.050 [fs]; AVG Krylov iteration: 5.12
01:22:31 | INFO | End 2900 step; propagated 145.050 [fs]; AVG Krylov iteration: 5.12
01:24:06 | INFO | Saved wavefunction 149.950 [fs]
01:24:08 | INFO | End 3000 step; propagated 150.050 [fs]; AVG Krylov iteration: 5.08
01:25:45 | INFO | End 3100 step; propagated 155.050 [fs]; AVG Krylov iteration: 5.08
01:27:21 | INFO | End 3200 step; propagated 160.050 [fs]; AVG Krylov iteration: 5.08
01:28:58 | INFO | End 3300 step; propagated 165.050 [fs]; AVG Krylov iteration: 5.08
01:30:35 | INFO | End 3400 step; propagated 170.050 [fs]; AVG Krylov iteration: 5.08
01:32:10 | INFO | End 3500 step; propagated 175.050 [fs]; AVG Krylov iteration: 5.08
01:33:48 | INFO | End 3600 step; propagated 180.050 [fs]; AVG Krylov iteration: 5.08
01:35:25 | INFO | End 3700 step; propagated 185.050 [fs]; AVG Krylov iteration: 5.08
01:37:00 | INFO | End 3800 step; propagated 190.050 [fs]; AVG Krylov iteration: 5.08
01:38:34 | INFO | End 3900 step; propagated 195.050 [fs]; AVG Krylov iteration: 5.08
01:40:09 | INFO | Saved wavefunction 199.950 [fs]
01:40:10 | INFO | End 3999 step; propagated 199.950 [fs]; AVG Krylov iteration: 5.08
01:40:10 | INFO | End simulation and save wavefunction
01:40:10 | INFO | Wave function is saved in wf_isomerization_exciton_bath.pkl
[14]:
(None, <pytdscf.wavefunction.WFunc at 0x12fab1160>)
Check results (reduced densities)
[15]:
from pytdscf.util import read_nc
data = read_nc(f"{jobname}_prop/reduced_density.nc", [(0,), (0,1,2)])
time_data = data["time"]
density_data_0 = data[(0,)].real
density_data_012 = data[(0,1,2)].real
[16]:
plt.plot(time_data, density_data_0[:, 0], label="S0")
plt.plot(time_data, density_data_0[:, 1], label="S1")
plt.legend()
plt.show()

[17]:
from matplotlib.animation import FuncAnimation
fig, (ax0, ax1) = plt.subplots(1, 2, figsize=(16, 6))
Q, θ = np.meshgrid(basis[1].get_grids(), basis[2].get_grids())
def update(i):
ax0.cla()
ax1.cla()
im0 = ax0.contour(
Q,
θ,
density_data_012[i, 0, :, :].T,
levels=[1.0e-05, 1.0e-04, 1.0e-03],
colors=["lightblue", "deepskyblue", "navy"],
)
# cmap='jet', vmin=0, vmax=0.05, )
im1 = ax1.contour(
Q,
θ,
density_data_012[i, 1, :, :].T,
levels=[1.0e-05, 1.0e-04, 1.0e-03],
colors=["lightblue", "deepskyblue", "navy"],
)
# fix coordinate of title
fig.suptitle(f"time = {time_data[i]:4.2f} [fs]")
ax0.set_title(
f"S0 \n population = {np.sum(density_data_012[i, 0, :, :]):0.2f}"
)
ax1.set_title(
f"S1 \n population = {np.sum(density_data_012[i, 1, :, :]):0.2f}"
)
ax0.set_xlabel("mass-weighted $Q_0$ [a.u.]")
ax1.set_xlabel("mass-weighted $Q_0$ [a.u.]")
ax0.set_ylabel(r"$\theta_1$ [rad]")
ax1.set_ylabel(r"$\theta_1$ [rad]")
ax0.set_yticks(
[-np.pi, -np.pi / 2.0, 0.0, np.pi / 2.0, np.pi],
[
"$-\\pi$ \n trans",
r"$-\frac{\pi}{2}$",
"0 \n cis",
r"$\frac{\pi}{2}$",
"$\\pi$ \n trans",
],
)
ax1.set_yticks(
[-np.pi, -np.pi / 2.0, 0.0, np.pi / 2.0, np.pi],
[
"$-\\pi$ \n trans",
r"$-\frac{\pi}{2}$",
"0 \n cis",
r"$\frac{\pi}{2}$",
"$\\pi$ \n trans",
],
)
return im0, im1
anim = FuncAnimation(fig, update, frames=range(0, len(time_data)), interval=100)
if use_bath:
gif_name = "reduced_density_exciton_bath.gif"
else:
gif_name = "reduced_density_exciton_zerotemp.gif"
anim.save(gif_name, writer="pillow")

[ ]: