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 sympy
from discvar import Exponential, HarmonicOscillator, Sine
from pympo import AssignManager, OpSite, SumOfProducts

from pytdscf import Exciton, Model, Simulator, units

\[\begin{split}\begin{align} \hat{H} &= \left[-\frac{1}{2} \frac{\partial^2}{\partial Q_1^2} - \frac{1}{2I} \frac{\partial^2}{\partial \theta_2^2} - \sum_{j=3}^{25} \frac{1}{2} \frac{\partial^2}{\partial Q_{j}^2} \right] \begin{pmatrix} 1 & 0 \\ 0 & 1 \end{pmatrix} \\ &+ \begin{pmatrix} D_e\left(1-\exp\left(-\sqrt{\frac{\Omega^2}{2D_e}}Q_1\right)\right)^2 - \frac{E}{2} & \lambda Q_1 \\ \lambda Q_1 & D_e\left(1-\exp\left(-\sqrt{\frac{\Omega^2}{2D_e}}(Q_1-Q_{\mathrm{ref}})\right)\right)^2 + \frac{E}{2} \end{pmatrix} \\ &+ \begin{pmatrix} \frac{W_0}{2}\left(1-\cos\theta_2\right) & 0\\ 0 & -\frac{W_1}{2}\left(1-\cos\theta_2\right) \end{pmatrix} \\ &+\sum_{j=3}^{25} \begin{pmatrix} \frac{\omega_j^2}{2}Q_j^2 & 0 \\ 0 & \frac{\omega_j^2}{2}Q_j^2+ \kappa_j Q_j \end{pmatrix} \end{align}\end{split}\]

References:

Johnson, Philip JM, et al. “The primary photochemistry of vision occurs at the molecular speed limit.” The Journal of Physical Chemistry B 121.16 (2017): 4040-4047.
Morse potential is our original demonstration.

The second quantized representation is given by

\[\hat{T} = \left( \hat{a}_0 \hat{a}_0^\dagger + \hat{a}_0^\dagger \hat{a}_0 \right) \left( -\frac{1}{2} \frac{\partial^2}{\partial Q_1^2} -\frac{1}{2I} \frac{\partial^2}{\partial \theta_2^2} -\sum_{j=3}^{25} \frac{1}{2} \frac{\partial^2}{\partial Q_j^2} \right)\]

\[\begin{split}\begin{align*} \hat{V} &= \hat{a}_0 \hat{a}_0^\dagger \left( -\frac{E}{2} + D_e\left(1-\exp\left(-\sqrt{\frac{\Omega^2}{2D_e}}Q_1\right)\right)^2 + \frac{W_0}{2}\left(1-\cos\theta_2\right) + \sum_{j=3}^{25} \frac{\omega_j^2}{2} Q_j^2 \right)\\ &+ \hat{a}_0^\dagger \hat{a}_0 \left( \frac{E}{2} + D_e\left(1-\exp\left(-\sqrt{\frac{\Omega^2}{2D_e}}(Q_1-Q_{\mathrm{ref}})\right)\right)^2 - \frac{W_1}{2}\left(1-\cos\theta_2\right) + \sum_{j=3}^{25} \frac{\omega_j^2}{2} Q_j^2 + \kappa_j Q_j \right) \\ &+ \left(\hat{a}_0 + \hat{a}_0^\dagger\right) \lambda Q_1 \end{align*}\end{split}\]

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.")
        )
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]:

$\displaystyle \frac{E a^{\dagger}_{0} a_{0}}{2} - \frac{E a_{0} a^{\dagger}_{0}}{2} + \frac{W_{0} a_{0} a^{\dagger}_{0} (1-\cos \theta_{2})}{2} - \frac{W_{1} a^{\dagger}_{0} a_{0} (1-\cos \theta_{2})}{2} + c_{10} a^{\dagger}_{0} a_{0} q_{10} + c_{11} a^{\dagger}_{0} a_{0} q_{11} + c_{12} a^{\dagger}_{0} a_{0} q_{12} + c_{13} a^{\dagger}_{0} a_{0} q_{13} + c_{14} a^{\dagger}_{0} a_{0} q_{14} + c_{15} a^{\dagger}_{0} a_{0} q_{15} + c_{16} a^{\dagger}_{0} a_{0} q_{16} + c_{17} a^{\dagger}_{0} a_{0} q_{17} + c_{18} a^{\dagger}_{0} a_{0} q_{18} + c_{19} a^{\dagger}_{0} a_{0} q_{19} + c_{20} a^{\dagger}_{0} a_{0} q_{20} + c_{21} a^{\dagger}_{0} a_{0} q_{21} + c_{22} a^{\dagger}_{0} a_{0} q_{22} + c_{23} a^{\dagger}_{0} a_{0} q_{23} + c_{24} a^{\dagger}_{0} a_{0} q_{24} + c_{25} a^{\dagger}_{0} a_{0} q_{25} + c_{3} a^{\dagger}_{0} a_{0} q_{3} + c_{4} a^{\dagger}_{0} a_{0} q_{4} + c_{5} a^{\dagger}_{0} a_{0} q_{5} + c_{6} a^{\dagger}_{0} a_{0} q_{6} + c_{7} a^{\dagger}_{0} a_{0} q_{7} + c_{8} a^{\dagger}_{0} a_{0} q_{8} + c_{9} a^{\dagger}_{0} a_{0} q_{9} + \lambda \left(a^{\dagger}_{0} + a_{0}\right) q_{1} + \frac{\omega_{10}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{10}}{2} + \frac{\omega_{11}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{11}}{2} + \frac{\omega_{12}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{12}}{2} + \frac{\omega_{13}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{13}}{2} + \frac{\omega_{14}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{14}}{2} + \frac{\omega_{15}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{15}}{2} + \frac{\omega_{16}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{16}}{2} + \frac{\omega_{17}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{17}}{2} + \frac{\omega_{18}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{18}}{2} + \frac{\omega_{19}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{19}}{2} + \frac{\omega_{20}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{20}}{2} + \frac{\omega_{21}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{21}}{2} + \frac{\omega_{22}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{22}}{2} + \frac{\omega_{23}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{23}}{2} + \frac{\omega_{24}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{24}}{2} + \frac{\omega_{25}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{25}}{2} + \frac{\omega_{3}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{3}}{2} + \frac{\omega_{4}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{4}}{2} + \frac{\omega_{5}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{5}}{2} + \frac{\omega_{6}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{6}}{2} + \frac{\omega_{7}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{7}}{2} + \frac{\omega_{8}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{8}}{2} + \frac{\omega_{9}^{2} \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) q^{2}_{9}}{2} + a^{\dagger}_{0} a_{0} M^{(1)}_1 + a_{0} a^{\dagger}_{0} M^{(0)}_1$

[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]:

$\displaystyle - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{1} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{10} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{11} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{12} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{13} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{14} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{15} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{16} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{17} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{18} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{19} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{20} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{21} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{22} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{23} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{24} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{25} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{3} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{4} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{5} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{6} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{7} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{8} - 0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{9} - \frac{0.5 \left(a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\right) p^{2}_{2}}{I}$

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)

$\displaystyle \left[\begin{matrix}a^{\dagger}_{0} + a_{0} & a^{\dagger}_{0} a_{0} & a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0} & a_{0} a^{\dagger}_{0}\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 0 & 0 & \lambda q_{1}\\1 & 0 & 0 & M^{(1)}_1\\0 & 0 & 1 & 0\\0 & 1 & 0 & M^{(0)}_1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & - \frac{W_{1} (1-\cos \theta_{2})}{2}\\0 & 0 & - \frac{E}{2} + \frac{W_{0} (1-\cos \theta_{2})}{2}\\0 & 1 & 0\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & c_{3} q_{3}\\0 & 1 & \frac{\omega_{3}^{2} q^{2}_{3}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & c_{4} q_{4}\\1 & 0 & \frac{\omega_{4}^{2} q^{2}_{4}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & \frac{\omega_{5}^{2} q^{2}_{5}}{2}\\1 & 0 & c_{5} q_{5}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & c_{6} q_{6}\\0 & 1 & \frac{\omega_{6}^{2} q^{2}_{6}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & c_{7} q_{7}\\1 & 0 & \frac{\omega_{7}^{2} q^{2}_{7}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & \frac{\omega_{8}^{2} q^{2}_{8}}{2}\\1 & 0 & c_{8} q_{8}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & c_{9} q_{9}\\1 & 0 & \frac{\omega_{9}^{2} q^{2}_{9}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & \frac{\omega_{10}^{2} q^{2}_{10}}{2}\\1 & 0 & c_{10} q_{10}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & c_{11} q_{11}\\0 & 1 & \frac{\omega_{11}^{2} q^{2}_{11}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & c_{12} q_{12}\\1 & 0 & \frac{\omega_{12}^{2} q^{2}_{12}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{13}^{2} q^{2}_{13}}{2}\\0 & 1 & c_{13} q_{13}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{14}^{2} q^{2}_{14}}{2}\\0 & 1 & c_{14} q_{14}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{15}^{2} q^{2}_{15}}{2}\\0 & 1 & c_{15} q_{15}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & \frac{\omega_{16}^{2} q^{2}_{16}}{2}\\1 & 0 & c_{16} q_{16}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & c_{17} q_{17}\\0 & 1 & \frac{\omega_{17}^{2} q^{2}_{17}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & c_{18} q_{18}\\1 & 0 & \frac{\omega_{18}^{2} q^{2}_{18}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{19}^{2} q^{2}_{19}}{2}\\0 & 1 & c_{19} q_{19}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{20}^{2} q^{2}_{20}}{2}\\0 & 1 & c_{20} q_{20}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{21}^{2} q^{2}_{21}}{2}\\0 & 1 & c_{21} q_{21}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0 & \frac{\omega_{22}^{2} q^{2}_{22}}{2}\\0 & 1 & c_{22} q_{22}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & \frac{\omega_{23}^{2} q^{2}_{23}}{2}\\1 & 0 & c_{23} q_{23}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}0 & 1 & c_{24} q_{24}\\1 & 0 & \frac{\omega_{24}^{2} q^{2}_{24}}{2}\\0 & 0 & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}\frac{\omega_{25}^{2} q^{2}_{25}}{2}\\\frac{E}{2} + c_{25} q_{25}\\1\end{matrix}\right]$

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 =

$\displaystyle \frac{E a^{\dagger}_{0} a_{0}}{2} - \frac{E a_{0} a^{\dagger}_{0}}{2} + \frac{W_{0} a_{0} a^{\dagger}_{0} (1-\cos \theta_{2})}{2} - \frac{W_{1} a^{\dagger}_{0} a_{0} (1-\cos \theta_{2})}{2} + c_{10} a^{\dagger}_{0} a_{0} q_{10} + c_{11} a^{\dagger}_{0} a_{0} q_{11} + c_{12} a^{\dagger}_{0} a_{0} q_{12} + c_{13} a^{\dagger}_{0} a_{0} q_{13} + c_{14} a^{\dagger}_{0} a_{0} q_{14} + c_{15} a^{\dagger}_{0} a_{0} q_{15} + c_{16} a^{\dagger}_{0} a_{0} q_{16} + c_{17} a^{\dagger}_{0} a_{0} q_{17} + c_{18} a^{\dagger}_{0} a_{0} q_{18} + c_{19} a^{\dagger}_{0} a_{0} q_{19} + c_{20} a^{\dagger}_{0} a_{0} q_{20} + c_{21} a^{\dagger}_{0} a_{0} q_{21} + c_{22} a^{\dagger}_{0} a_{0} q_{22} + c_{23} a^{\dagger}_{0} a_{0} q_{23} + c_{24} a^{\dagger}_{0} a_{0} q_{24} + c_{25} a^{\dagger}_{0} a_{0} q_{25} + c_{3} a^{\dagger}_{0} a_{0} q_{3} + c_{4} a^{\dagger}_{0} a_{0} q_{4} + c_{5} a^{\dagger}_{0} a_{0} q_{5} + c_{6} a^{\dagger}_{0} a_{0} q_{6} + c_{7} a^{\dagger}_{0} a_{0} q_{7} + c_{8} a^{\dagger}_{0} a_{0} q_{8} + c_{9} a^{\dagger}_{0} a_{0} q_{9} + \lambda a^{\dagger}_{0} q_{1} + \lambda a_{0} q_{1} + \frac{\omega_{10}^{2} a^{\dagger}_{0} a_{0} q^{2}_{10}}{2} + \frac{\omega_{10}^{2} a_{0} a^{\dagger}_{0} q^{2}_{10}}{2} + \frac{\omega_{11}^{2} a^{\dagger}_{0} a_{0} q^{2}_{11}}{2} + \frac{\omega_{11}^{2} a_{0} a^{\dagger}_{0} q^{2}_{11}}{2} + \frac{\omega_{12}^{2} a^{\dagger}_{0} a_{0} q^{2}_{12}}{2} + \frac{\omega_{12}^{2} a_{0} a^{\dagger}_{0} q^{2}_{12}}{2} + \frac{\omega_{13}^{2} a^{\dagger}_{0} a_{0} q^{2}_{13}}{2} + \frac{\omega_{13}^{2} a_{0} a^{\dagger}_{0} q^{2}_{13}}{2} + \frac{\omega_{14}^{2} a^{\dagger}_{0} a_{0} q^{2}_{14}}{2} + \frac{\omega_{14}^{2} a_{0} a^{\dagger}_{0} q^{2}_{14}}{2} + \frac{\omega_{15}^{2} a^{\dagger}_{0} a_{0} q^{2}_{15}}{2} + \frac{\omega_{15}^{2} a_{0} a^{\dagger}_{0} q^{2}_{15}}{2} + \frac{\omega_{16}^{2} a^{\dagger}_{0} a_{0} q^{2}_{16}}{2} + \frac{\omega_{16}^{2} a_{0} a^{\dagger}_{0} q^{2}_{16}}{2} + \frac{\omega_{17}^{2} a^{\dagger}_{0} a_{0} q^{2}_{17}}{2} + \frac{\omega_{17}^{2} a_{0} a^{\dagger}_{0} q^{2}_{17}}{2} + \frac{\omega_{18}^{2} a^{\dagger}_{0} a_{0} q^{2}_{18}}{2} + \frac{\omega_{18}^{2} a_{0} a^{\dagger}_{0} q^{2}_{18}}{2} + \frac{\omega_{19}^{2} a^{\dagger}_{0} a_{0} q^{2}_{19}}{2} + \frac{\omega_{19}^{2} a_{0} a^{\dagger}_{0} q^{2}_{19}}{2} + \frac{\omega_{20}^{2} a^{\dagger}_{0} a_{0} q^{2}_{20}}{2} + \frac{\omega_{20}^{2} a_{0} a^{\dagger}_{0} q^{2}_{20}}{2} + \frac{\omega_{21}^{2} a^{\dagger}_{0} a_{0} q^{2}_{21}}{2} + \frac{\omega_{21}^{2} a_{0} a^{\dagger}_{0} q^{2}_{21}}{2} + \frac{\omega_{22}^{2} a^{\dagger}_{0} a_{0} q^{2}_{22}}{2} + \frac{\omega_{22}^{2} a_{0} a^{\dagger}_{0} q^{2}_{22}}{2} + \frac{\omega_{23}^{2} a^{\dagger}_{0} a_{0} q^{2}_{23}}{2} + \frac{\omega_{23}^{2} a_{0} a^{\dagger}_{0} q^{2}_{23}}{2} + \frac{\omega_{24}^{2} a^{\dagger}_{0} a_{0} q^{2}_{24}}{2} + \frac{\omega_{24}^{2} a_{0} a^{\dagger}_{0} q^{2}_{24}}{2} + \frac{\omega_{25}^{2} a^{\dagger}_{0} a_{0} q^{2}_{25}}{2} + \frac{\omega_{25}^{2} a_{0} a^{\dagger}_{0} q^{2}_{25}}{2} + \frac{\omega_{3}^{2} a^{\dagger}_{0} a_{0} q^{2}_{3}}{2} + \frac{\omega_{3}^{2} a_{0} a^{\dagger}_{0} q^{2}_{3}}{2} + \frac{\omega_{4}^{2} a^{\dagger}_{0} a_{0} q^{2}_{4}}{2} + \frac{\omega_{4}^{2} a_{0} a^{\dagger}_{0} q^{2}_{4}}{2} + \frac{\omega_{5}^{2} a^{\dagger}_{0} a_{0} q^{2}_{5}}{2} + \frac{\omega_{5}^{2} a_{0} a^{\dagger}_{0} q^{2}_{5}}{2} + \frac{\omega_{6}^{2} a^{\dagger}_{0} a_{0} q^{2}_{6}}{2} + \frac{\omega_{6}^{2} a_{0} a^{\dagger}_{0} q^{2}_{6}}{2} + \frac{\omega_{7}^{2} a^{\dagger}_{0} a_{0} q^{2}_{7}}{2} + \frac{\omega_{7}^{2} a_{0} a^{\dagger}_{0} q^{2}_{7}}{2} + \frac{\omega_{8}^{2} a^{\dagger}_{0} a_{0} q^{2}_{8}}{2} + \frac{\omega_{8}^{2} a_{0} a^{\dagger}_{0} q^{2}_{8}}{2} + \frac{\omega_{9}^{2} a^{\dagger}_{0} a_{0} q^{2}_{9}}{2} + \frac{\omega_{9}^{2} a_{0} a^{\dagger}_{0} q^{2}_{9}}{2} + a^{\dagger}_{0} a_{0} M^{(1)}_1 + a_{0} a^{\dagger}_{0} M^{(0)}_1$

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

$\displaystyle \left[\begin{matrix}a^{\dagger}_{0} a_{0} + a_{0} a^{\dagger}_{0}\end{matrix}\right]$

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

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

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{3} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{4} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{5} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{6} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{7} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{8} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{9} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{10} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{11} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{12} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{13} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{14} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{15} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{16} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{17} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{18} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{19} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{20} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{21} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{22} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{23} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1 & 0\\- 0.5 p^{2}_{24} & 1\end{matrix}\right]$

$\displaystyle \left[\begin{matrix}1\\- 0.5 p^{2}_{25}\end{matrix}\right]$

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 =

$\displaystyle - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{1} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{10} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{11} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{12} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{13} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{14} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{15} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{16} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{17} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{18} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{19} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{20} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{21} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{22} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{23} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{24} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{25} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{3} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{4} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{5} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{6} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{7} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{8} - 0.5 a^{\dagger}_{0} a_{0} p^{2}_{9} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{1} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{10} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{11} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{12} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{13} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{14} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{15} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{16} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{17} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{18} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{19} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{20} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{21} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{22} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{23} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{24} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{25} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{3} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{4} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{5} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{6} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{7} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{8} - 0.5 a_{0} a^{\dagger}_{0} p^{2}_{9} - \frac{0.5 a^{\dagger}_{0} a_{0} p^{2}_{2}}{I} - \frac{0.5 a_{0} a^{\dagger}_{0} p^{2}_{2}}{I}$

Setup Hamiltonian

[11]:

operators = {"potential": pot_mpo, "kinetic": kin_mpo}

Setup Model (basis, operators, initial states)

[12]:

model = Model(basis, operators=operators, bond_dim=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()

../_images/notebook_TD_reduced_density_exciton_27_0.png

[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")

../_images/notebook_TD_reduced_density_exciton_28_0.png

[ ]: