Displaced Harmonic Oscillator

Although we have successfully implemented the dynamics with laser pulses in the class QuantumDynamics1D, we still have to write quite a bit of code afterwards to setup a system with multiple harmonic potentials. We can make this process easier by writing a child class of QuantumDynamics1D, which is dedicated to perform dynamics on several harmonic surfaces.

We can also add additional features to this child class, in this case plotting potentials and populations.

We use the class HarmonicOscillator from section 6.2.1 for the child class DisplacedHarmonicOscillator.

class DisplacedHarmonicOscillator(QuantumDynamics1D):
    def __init__(self, m=1.0, omega=np.ones(2), d=np.zeros(2), 
                 e_eg=np.zeros(2), colormap='jet'):
        self.mass = m
        self.omega = omega
        self.nstate = omega.shape[0]
        self.d = d
        self.e_eg = e_eg

        self.ho = []
        for i in range(self.nstate):
            self.ho.append(HarmonicOscillator(self.mass, self.omega[i]))

        cmap = get_cmap(colormap)
        self.colors = cmap(np.linspace(0.0, 1.0, self.nstate))

    def set_potential(self):
        self.vgrid = np.zeros((self.n, self.nstate, self.nstate))
        for state in range(self.nstate):
            self.vgrid[:, state, state] = self.e_eg[state] \
                + self.ho[state].get_potential(self.xgrid - self.d[state])

    def set_coupling(self, v12=1.0, sigma=6.0, ij=((0, 1), (1, 0))):
        self.vcoupling = np.zeros((self.n, self.nstate, self.nstate))
        for kl in ij:
            v = v12 * np.exp(-self.xgrid**2 / (2.0*sigma**2))
            self.vcoupling[:, kl[0], kl[1]]= v
        self.vgrid = self.vgrid + self.vcoupling

    def set_dipole_coupling(self, mu, ij=((0, 1), (1, 0))):
        self.mu = np.zeros((self.n, self.nstate, self.nstate))
        for kl in ij:
            self.mu[:, kl[0], kl[1]] = mu[kl] * np.ones(self.xgrid.shape[0])

    def plot_potentials(self, ax, scalecoupling=100.0):
        for state in range(self.nstate):
            ax.plot(self.xgrid, self.vgrid[:, state, state], 
                    color=self.colors[state], lw=3.0, 
                    label="$V_{{{0}{0}}}$".format(state + 1))
        
        for i in range(self.nstate):
            for j in range(i + 1,self.nstate):
                ax.plot(self.xgrid, scalecoupling * self.vcoupling[:, i, j],
                        color="magenta", lw=3.0, 
                        label="$V_{{{0}{1}}}$".format(i + 1, j + 1))
                ax.fill_between(self.xgrid, scalecoupling * self.vcoupling[:, i, j],
                                 color="magenta", lw=3.0, alpha=0.2)
        
        ax.set_xlabel('$x$')
        ax.set_ylabel('$V(x)$')
        ax.legend()

    def get_state_populations(self):
        ntime = self.wavepackets.shape[0]
        self.populations = self.dx * np.einsum(
            'ijk,ijk->ik', np.conjugate(self.wavepackets), self.wavepackets,
        )

    def plot_populations(self, ax):
        for state in range(self.nstate):
            ax.plot(self.populations[:, state].real, 
                    label="$P_{{{}}}$".format(state + 1))
        
        ax.set_xlabel("frame")
        ax.set_ylabel("population")
        ax.legend()

We then setup the same system using this new class and perform the simulation.

displacements = np.array([3.0, -3.0])
dho = DisplacedHarmonicOscillator(d=displacements)

# setup grid
xmin = -20.0
xmax = -xmin
ngrid = 256
x = np.linspace(xmin, xmax, ngrid)
dho.setup_grid(xmin, xmax, ngrid)

# setup potential
dho.set_potential()

# setup coupling
dho.set_coupling()

# setup transition dipole matrix
mu = np.zeros((2, 2))
mu[0, 1] = 1.0
mu[1, 0] = 1.0
dho.set_dipole_coupling(mu)

# time step
n = 200
nstep = 2000
tstep = (2*np.pi/dho.omega[0]) / n

# setup electric field
def electric_field(t, e0, omega, sigma, td):
    return e0 * np.cos(omega*t) * np.exp(-(t - td)**2/(2.0*sigma**2))

e0 = 5.0
omega_eg = dho.e_eg[1] - dho.e_eg[0] + dho.ho[1].get_potential(dho.d[0] - dho.d[1])
print("omega_eg = {}".format(omega_eg))

sigma_t = 20 * tstep
td = 300 * tstep
dho.set_efield(lambda t: electric_field(t, e0, omega_eg, sigma_t, td))

# setup initial wave function
psi0 = np.zeros((dho.xgrid.shape[0], dho.nstate))
psi00 = dho.ho[1].get_eigenfunction(x - dho.d[1], 0)
psi0[:, 1] = psi00
psi0[:, 0] = np.zeros(x.shape[0])

# visualize and calculate population
dho.propagate(tstep, nstep, psi0, method="TD Split Operator")
dho.get_state_populations()

# plot
fig, axs = plt.subplots(1, 2, figsize=(10, 4))
dho.plot_potentials(axs[0])
dho.plot_populations(axs[1])
fig.tight_layout()

# animate
anim_dho = dho.animate(nstep=nstep, delay=10, scalewf=50.0, plot_potential=True)