利用PyMOL制作反应势能面示意动画

# coding: utf-8

# ##############################################################################
# 2013-11-14 10:25:11 简单示例
# 2016-09-07 20:23:18 注释, 颜色映射
# ##############################################################################

import math
from pymol import cmd
from pymol.cgo import *

# 势能面函数
def Fxy(x,y):
	if x==0: Fx = 4
	else: Fx = 4*math.sin(x)/x+.1*x
	if y==0: Fy = 4
	else: Fy = 4*math.sin(y)/y
	return Fx+Fy

# 势能面函数法向
def dFxy(x,y):
	if x==0: dFx = 0
	else: dFx = 4*(x*math.cos(x)-math.sin(x))/x**2+0.1
	if y==0: dFy = 0
	else: dFy = 4*(y*math.cos(y)-math.sin(y))/y**2
	Rtmp = 1/math.sqrt(dFx*dFx+dFy*dFy+1)
	return -dFx*Rtmp, -dFy*Rtmp, Rtmp

def RGB(V, Vmin, Vmax):
	dV=Vmax-Vmin; x=(V-Vmin)/dV
	r=1; g=1; b=1
	if   x<0.25: r = 0;     g = 4*x
	elif x<0.50: r = 0;     b = 2-4*x
	elif x<0.75: r = 4*x-2; b = 0
	else:        g = 4-4*x; b = 0

	r=min(max(r,0), 1)
	g=min(max(g,0) ,1)
	b=min(max(b,0) ,1)
	return r, g, b

# 是否使用法向, 显示网格, 颜色映射, 使用Verlet方法计算轨迹
YesNorm=1; YesGrid=1; YesMap=0; YesTrj=1

Xini = 4.5 # 小球初始位置
Xmin = -7; Xmax = 5.5; dX = .5; Nx = int((Xmax-Xmin)/dX)
Ymin = -7; Ymax = 7;   dY = .5; Ny = int((Ymax-Ymin)/dY)

X = [ 0 for i in range(Nx) ]
Y = [ 0 for j in range(Ny) ]
Z = [ [0]*Ny for i in range(Nx) ]
Zx= [ [0]*Ny for i in range(Nx) ]
Zy= [ [0]*Ny for i in range(Nx) ]
Zz= [ [1]*Ny for i in range(Nx) ]

for i in range(Nx): X[i] = Xmin+dX*i
for j in range(Ny): Y[j] = Ymin+dY*j

for i in range(Nx):
	x=X[i]
	for j in range(Ny):
		y=Y[j]
		Z[i][j] = Fxy(x,y)
		if YesNorm: Zx[i][j], Zy[i][j], Zz[i][j] = dFxy(x,y)

# 获取极值, 用于颜色映射
Zmin=min(min(Z)); Zmax=max(max(Z))

PES = []
# 绘制网格
if YesGrid:
	PES.extend( [ COLOR,  0, 0, 0 ] )
	for i in range(Nx):
		PES.extend( [ BEGIN, LINE_STRIP ] )
		for j in range(0,Ny):
			PES.extend( [ NORMAL, Zx[i][j], Zy[i][j], Zz[i][j] ] )
			PES.extend( [ VERTEX, X[i], Y[j], Z[i][j] ] )
		PES.append( END )

	for j in range(Ny):
		PES.extend( [ BEGIN, LINE_STRIP ] )
		for i in range(Nx):
			PES.extend( [ NORMAL, Zx[i][j], Zy[i][j], Zz[i][j] ] )
			PES.extend( [ VERTEX, X[i], Y[j], Z[i][j] ] )
		PES.append( END )

# 绘制表面
PES.extend( [ COLOR, .19, .6, .83 ] )
for j in range(Ny-1):
	PES.extend( [ BEGIN, TRIANGLE_STRIP ] )
	for i in range(Nx):
		if YesMap:
			r, g, b=RGB(Z[i][j+1], Zmin, Zmax)
			PES.extend( [ COLOR,  r, g, b ] )
		PES.extend( [ NORMAL, Zx[i][j+1], Zy[i][j+1], Zz[i][j+1] ] )
		PES.extend( [ VERTEX, X[i], Y[j+1], Z[i][j+1] ] )
		if YesMap:
			r, g, b=RGB(Z[i][j], Zmin, Zmax)
			PES.extend( [ COLOR,  r, g, b ] )
		PES.extend( [ NORMAL, Zx[i][j], Zy[i][j], Zz[i][j] ] )
		PES.extend( [ VERTEX, X[i], Y[j], Z[i][j] ] )
	PES.append( END )

# 绘制路径
PES.extend( [ LINEWIDTH, 5 ] )
PES.extend( [ BEGIN, LINE_STRIP ] )
PES.extend( [ COLOR, 1., 0., .0 ] )
x = -Xini; y = -Xini
while y<=Xini:
	z = Fxy(x,y)
	dFx, dFy, dFz = dFxy(x,y)
	PES.extend( [ NORMAL, dFx, dFy, dFz ] )
	PES.extend( [ VERTEX, x, y, z ] )
	y = y+0.1
PES.append( END )

PES.extend( [ BEGIN, LINE_STRIP] )
PES.extend( [ COLOR, 0, 1., .0 ] )
x = -Xini; y = -Xini
while x<=Xini:
	z = Fxy(x,y)
	dFx, dFy, dFz = dFxy(x,y)
	PES.extend( [ NORMAL, dFx, dFy, dFz ] )
	PES.extend( [ VERTEX, x, y, z ] )
	x = x+0.1
PES.append( END )

# 绘制小球
x = -Xini; y = -Xini; z = Fxy(x,y)
PES.extend ( [ COLOR, 1, 1, 0 ] )
PES.extend ( [ SPHERE, x, y, z, .2 ] )
x = -Xini; y = Xini; z = Fxy(x,y)
PES.extend ( [ COLOR, 1, 0, 0 ] )
PES.extend ( [ SPHERE, x, y, z, .2 ] )
x = Xini; y = -Xini; z = Fxy(x,y)
PES.extend ( [ COLOR, 0, 1, 0 ] )
PES.extend ( [ SPHERE, x, y, z, .2 ] )

cmd.load_cgo(PES, 'PES')

# 模拟小球运行
Rsph = 0.5
if not YesTrj:
	x = -Xini
	for Ifrm in range(30):
		y = -Xini+Ifrm*0.3; z = Fxy(x,y)
		dFx, dFy, dFz = dFxy(x,y)
		SYS = [ COLOR, 1, 1-Ifrm/30., 0 ]
		SYS.extend( [ SPHERE, x+Rsph*dFx, y+Rsph*dFy, z+Rsph*dFz, Rsph ] )
		cmd.load_cgo(SYS, 'SYS', Ifrm)

	y = -Xini
	for Ifrm in range(30):
		x = -Xini+Ifrm*0.3; z = Fxy(x,y)
		dFx, dFy, dFz = dFxy(x,y)
		SYS = [ COLOR, 1-Ifrm/30., 1, 0 ]
		SYS.extend( [ SPHERE, x+Rsph*dFx, y+Rsph*dFy, z+Rsph*dFz, Rsph ] )
		cmd.load_cgo(SYS, 'SYS', Ifrm+30)
else:
	# 初始位置和速度
	x = -Xini; y=-Xini
	vx = 3.28; vy = 0;    Nfrm=130
	vx = 0;    vy = 3.15; Nfrm=120

	dt = 0.05; E0 = 0.5*(vx*vx+vy*vy)+Fxy(x,y)
	dFx0, dFy0, dFz0 = dFxy(x,y)
	for Ifrm in range(Nfrm):
		SYS = [ COLOR, 1, 1, 1 ]
		x = x + (vx + 0.5*dFx0*dt)*dt;
		y = y + (vy + 0.5*dFy0*dt)*dt;
		z = Fxy(x,y)
		dFx, dFy, dFz = dFxy(x,y)
		SYS.extend( [ SPHERE, x+Rsph*dFx, y+Rsph*dFy, z+Rsph*dFz, Rsph ] )
		cmd.load_cgo(SYS, 'SYS', Ifrm)

		vx = vx+0.5*(dFx0+dFx)*dt
		vy = vy+0.5*(dFy0+dFy)*dt
		E = 0.5*(vx*vx+vy*vy)
		Rtmp = math.sqrt(2.*abs(E0-z)/(vx*vx+vy*vy))
		vx = vx*Rtmp; vy = vy*Rtmp
		dFx0, dFy0 = dFx, dFy

cmd.reset()
cmd.zoom('PES', 1.0)
cmd.clip('far', -10.0)
cmd.turn('z', 30)
cmd.turn('x', -60)
cmd.mplay()