def compute_ensemble(qs, t0, t1, t2, F, mu, mol_I, max_steps = 1000000):
muz = zeros(qs.shape[0])
for i in xrange(qs.shape[0]):
qt = ar.asym_rotor(t0, t1, t2, max_steps, qs[i,:], F, mu, mol_I)
muz[i] = ar.muz_bar(qt,F,mu,mol_I)
print "computing (", i, "/", qs.shape[0], ") <muz> = ", muz[i]
return muz
def animate(coords, mass, dipole, temp, F, cycles, w0=None, max_tstep=100000, no_anim=True):
"""
UNITS:
coords angstrom
mass amu
t ps
dipole e * angstrom
temp kelvin
F amu * angstrom / ps**2.0 / e
"""
cm_coords = coords - tile(center_of_mass(coords, mass), (coords.shape[0], 1))
mol_I, mol_Ix = eig(inertia_tensor(cm_coords, mass))
mol_I.sort()
print "principal moments of inertia"
print mol_I
q0, w0 = initial_cond(coords, mass, dipole, temp, F)
print "initial conditions"
print " q0 = ", q0, " norm(q0) = ", norm(q0)
print " w0 = ", w0, " norm(w0) = ", norm(w0)
# q0 = array([0., 0., 1., 0.])
# w0 = array([0.,1.,0.], dtype=float)
print "time in picoseconds"
# t0 = 2.0 * pi / norm(w0) * cycles[0]
# t1 = t0 + 2.0 * pi / norm(w0) * cycles[1]
# t2 = t1 + 2.0 * pi / norm(w0) * cycles[2]
t0 = cycles[0]
t1 = t0 + cycles[1]
t2 = t1 + cycles[2]
print "t0 = ", t0
print "t1 = ", t1
print "t2 = ", t2
print "total time = ", t0 + t1 + t2
print "integrating motion of rotor"
qt = ar.asym_rotor(t0, t1, t2, max_tstep, r_[q0, w0], F, dipole, mol_I)
print "qt.shape = ", qt.shape
# now extract the orientation quaternions
tsr = qt[:, 0]
qtr = qt[:, 1:5]
wtr = qt[:, 5:8]
# diagnostic plots
# field in rk time steps
ft = zeros(len(tsr))
ft[tsr < t0] = 0.0
ft[(tsr > t0) * (tsr < t1)] = F * tsr / (t1 - t0)
ft[tsr > t1] = F
fp = pylab.figure(figsize=(3, 4))
# total energy
pylab.subplot(4, 1, 1)
pylab.plot(tsr, array([total_energy(qt[i, :], ft[i], dipole, mol_I) for i in xrange(len(tsr))]))
pylab.title("Total Energy")
# angular momentum
pylab.subplot(4, 1, 2)
L = r_[[angular_momentum(qt[i, :], ft[i], dipole, mol_I) for i in xrange(len(tsr))]]
print "L.shape = ", L.shape
# pylab.plot(tsr, array([norm(L[i,:]) for i in xrange(len(tsr))]), 'r-')
pylab.plot(tsr, array([L[i, 0] for i in xrange(len(tsr))]), "b-")
pylab.plot(tsr, array([L[i, 1] for i in xrange(len(tsr))]), "g-")
pylab.plot(tsr, array([L[i, 2] for i in xrange(len(tsr))]), "r-")
pylab.title("Angular Momentum")
# dipole projection
pylab.subplot(4, 1, 3)
dp = array([dipole_projection(qt[i, :], ft[i], dipole, mol_I) for i in xrange(len(tsr))])
dpavg = cumsum(dp) / (arange(len(dp)) + 1.0)
pylab.plot(tsr, dpavg)
pylab.title("Projected Dipole (polarization)")
# field strength
pylab.subplot(4, 1, 4)
pylab.plot(tsr, ft)
pylab.title("Field Strength")
pylab.subplots_adjust(hspace=0.95)
pylab.show()
if no_anim:
return
# frame rate is 10 ps to 1 s
frames = floor(t2 / 10 * 30)
print "total time is ", t2, " rendering ", frames, " frames"
ts, qt = interpolate_quat(tsr, qtr, frames)
# field
ft = zeros(len(ts))
ft[ts < t0] = 0.0
ft[(ts > t0) * (ts < t1)] = F * ts / (t1 - t0)
ft[ts > t1] = F
# render the molecule
f = mlab.figure(size=(500, 500), bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
# # draw the molecule
molecule = mlab.points3d(
cm_coords[:, 0], cm_coords[:, 1], cm_coords[:, 2], scale_factor=2, color=(1, 1, 0), resolution=20
)
# and the dipole moment
dip = mlab.quiver3d(
[0], [0], [0], [dipole[0]], [dipole[1]], [dipole[2]], scale_factor=1, line_width=2.0, color=(0, 1, 0)
)
mlab.view(distance=45.0)
# timelab = mlab.text(0.1, 0.9, 'Time: 0 ps', width=0.4)
# fieldlab = mlab.text(0.5, 0.9, 'Field: 0 kV/cm', width=0.4)
ms = molecule.mlab_source
dms = dip.mlab_source
for t, fi, q in zip(ts, ft, qt):
# timelab.text = "Time: %d ps" % t
# fieldlab.text = "Field: %f" % fi
M = q_to_rotmatrix(q)
cp = dot(cm_coords, M)
dp = dot(dipole, M) * 12.0
ms.set(x=cp[:, 0], y=cp[:, 1], z=cp[:, 2])
dms.set(u=[dp[0]], v=[dp[1]], w=[dp[2]])
