In [ ]:

    

%matplotlib inline
import matplotlib.pyplot as plt
import numpy as np
import pywigner as lsc
import dynamiq_engine as dynq
import dynamiq_engine.potentials as pes


    

In [ ]:

    

tully_V11 = pes.OneDimensionalInteractionModel(
    pes.interactions.TanhInteraction(a=1.6, V0=0.1)
)
tully_V22 = pes.OneDimensionalInteractionModel(
    pes.interactions.TanhInteraction(a=1.6, V0=-0.1)
)
tully_V12 = pes.OneDimensionalInteractionModel(
    pes.interactions.GaussianInteraction(A=0.05, alpha=1.0)
)
tully_matrix = dynq.NonadiabaticMatrix([[tully_V11, tully_V12],
                                        [tully_V12, tully_V22]])
tully = pes.MMSTHamiltonian(tully_matrix)
tully_topology = dynq.Topology(
    masses=np.array([1980.0]),
    potential=tully
)


    

In [ ]:

    

nuclear = lsc.operators.CoherentProjection(
    x0=[-5.0],
    p0=[19.0],
    gamma=[1.0]
)
electronic = lsc.operators.ElectronicCoherentProjection.with_n_dofs(2).excite(dof=1)


    

In [ ]:

    

tully_A_op = nuclear * electronic


    

In [ ]:

    

tully_sampler = tully_A_op.default_sampler()


    

In [ ]:

    

snap = dynq.MMSTSnapshot(
    coordinates=np.array([-5.0]),
    momenta=np.array([19.0]),
    electronic_coordinates=np.array([0.0, 0.0]),
    electronic_momenta=np.array([0.0,0.0]),
    topology=tully_topology
)


    

In [ ]:

    

# do 10k samples
samples = []
old_snap = snap
for i in range(10000):
    new_snap = tully_sampler.generate_initial_snapshot(old_snap)
    samples.append(new_snap)
    old_snap = new_snap


    

In [ ]:

    

hist_R = np.histogram([s.coordinates[0] for s in samples], bins=40, normed=True)
hist_P = np.histogram([s.momenta[0] for s in samples], bins=40, normed=True)
hist_x1 = np.histogram([s.electronic_coordinates[0] for s in samples], bins=40, normed=True)
hist_x2 = np.histogram([s.electronic_coordinates[1] for s in samples], bins=40, normed=True)
hist_p1 = np.histogram([s.electronic_momenta[0] for s in samples], bins=40, normed=True)
hist_p2 = np.histogram([s.electronic_momenta[1] for s in samples], bins=40, normed=True)


    

In [ ]:

    

def midpt(hist):
    return np.array([0.5*(hist[1][i]+hist[1][i+1]) for i in range(len(hist[1])-1)])


    

In [ ]:

    

class NormedGaussian(object):
    def __init__(self, alpha, x0):
        self.alpha = alpha
        self.x0 = x0
    
    def __call__(self, x):
        return np.sqrt(self.alpha / np.pi) * np.exp(-self.alpha * (x-self.x0)**2)


    

In [ ]:

    

gaussR = NormedGaussian(alpha=1.0, x0=-5.0)
R = midpt(hist_R)
plt.plot(R, hist_R[0])
plt.plot(R, gaussR(R))


    

In [ ]:

    

gaussP = NormedGaussian(alpha=1.0, x0=19.0)
P = midpt(hist_P)
plt.plot(P, hist_P[0])
plt.plot(P, gaussP(P))


    

In [ ]:

    

gaussx1 = NormedGaussian(alpha=1.0, x0=0.0)
x1 = midpt(hist_x1)
plt.plot(x1, hist_x1[0])
plt.plot(x1, gaussx1(x1))


    

In [ ]:

    

gaussx2 = NormedGaussian(alpha=1.0, x0=0.0)
x2 = midpt(hist_x2)
plt.plot(x2, hist_x2[0])
plt.plot(x2, gaussx2(x2))


    

In [ ]:

    

gaussp1 = NormedGaussian(alpha=1.0, x0=0.0)
p1 = midpt(hist_p1)
plt.plot(p1, hist_p1[0])
plt.plot(p1, gaussp1(p1))


    

In [ ]:

    

gaussp2 = NormedGaussian(alpha=1.0, x0=0.0)
p2 = midpt(hist_p2)
plt.plot(p2, hist_p2[0])
plt.plot(p2, gaussp2(p2))


    

In [ ]: