316 lines
9.0 KiB
Python
316 lines
9.0 KiB
Python
import math
|
|
|
|
class vector:
|
|
""" Vector """
|
|
def __init__(self, xi=0.0, yi=0.0, zi=0.0, atom1 = 0, atom2 = 0):
|
|
self.x = xi
|
|
self.y = yi
|
|
self.z = zi
|
|
|
|
if atom1:
|
|
# make vector pointing to atom1
|
|
self.x = atom1.x
|
|
self.y = atom1.y
|
|
self.z = atom1.z
|
|
|
|
if atom2:
|
|
# make inter-atomic vector (atom1 -> atom2)
|
|
self.x = atom2.x - self.x
|
|
self.y = atom2.y - self.y
|
|
self.z = atom2.z - self.z
|
|
|
|
return
|
|
|
|
def __add__(self, other):
|
|
return vector(self.x + other.x,
|
|
self.y + other.y,
|
|
self.z + other.z)
|
|
|
|
def __sub__(self, other):
|
|
return vector(self.x - other.x,
|
|
self.y - other.y,
|
|
self.z - other.z)
|
|
|
|
def __mul__(self, other):
|
|
""" Dot product, scalar and matrix multiplication """
|
|
|
|
if isinstance(other,vector):
|
|
return self.x * other.x + self.y * other.y + self.z * other.z
|
|
elif isinstance(other, matrix4x4):
|
|
return vector(
|
|
xi = other.a11*self.x + other.a12*self.y + other.a13*self.z + other.a14*1.0,
|
|
yi = other.a21*self.x + other.a22*self.y + other.a23*self.z + other.a24*1.0,
|
|
zi = other.a31*self.x + other.a32*self.y + other.a33*self.z + other.a34*1.0
|
|
)
|
|
elif type(other) in [int, float]:
|
|
return vector(self.x * other, self.y * other, self.z * other)
|
|
else:
|
|
print('%s not supported'%type(other))
|
|
raise TypeError
|
|
|
|
def __rmul__(self,other):
|
|
return self.__mul__(other)
|
|
|
|
|
|
def __pow__(self, other):
|
|
""" Cross product """
|
|
return vector(self.y * other.z - self.z * other.y,
|
|
self.z * other.x - self.x * other.z,
|
|
self.x * other.y - self.y * other.x)
|
|
|
|
|
|
def __neg__(self):
|
|
res = vector(xi = -self.x,
|
|
yi = -self.y,
|
|
zi = -self.z)
|
|
return res
|
|
|
|
def sq_length(self):
|
|
return self.x * self.x + self.y * self.y + self.z * self.z
|
|
|
|
def length(self):
|
|
return math.sqrt(self.sq_length())
|
|
|
|
def __str__(self):
|
|
return '%10.4f %10.4f %10.4f'%(self.x, self.y, self.z)
|
|
|
|
def __repr__(self):
|
|
return '<vector>'
|
|
|
|
def orthogonal(self):
|
|
""" Returns a vector orthogonal to self """
|
|
res = vector(self.y, -self.x, 0)
|
|
|
|
if abs(self.y) < abs(self.z):
|
|
res = vector(self.z, 0, -self.x)
|
|
|
|
return res
|
|
|
|
def rescale(self, new_length):
|
|
""" Rescale vector to new length while preserving direction """
|
|
frac = new_length/(self.length())
|
|
res = vector(xi = self.x*frac,
|
|
yi = self.y*frac,
|
|
zi = self.z*frac)
|
|
return res
|
|
|
|
|
|
class matrix4x4:
|
|
def __init__(self,
|
|
a11i=0.0, a12i=0.0, a13i=0.0, a14i=0.0,
|
|
a21i=0.0, a22i=0.0, a23i=0.0, a24i=0.0,
|
|
a31i=0.0, a32i=0.0, a33i=0.0, a34i=0.0,
|
|
a41i=0.0, a42i=0.0, a43i=0.0, a44i=0.0):
|
|
|
|
self.a11 = a11i
|
|
self.a12 = a12i
|
|
self.a13 = a13i
|
|
self.a14 = a14i
|
|
|
|
self.a21 = a21i
|
|
self.a22 = a22i
|
|
self.a23 = a23i
|
|
self.a24 = a24i
|
|
|
|
self.a31 = a31i
|
|
self.a32 = a32i
|
|
self.a33 = a33i
|
|
self.a34 = a34i
|
|
|
|
self.a41 = a41i
|
|
self.a42 = a42i
|
|
self.a43 = a43i
|
|
self.a44 = a44i
|
|
|
|
return
|
|
|
|
|
|
|
|
|
|
|
|
# methods working on vectors
|
|
|
|
|
|
def angle(a, b):
|
|
dot = a * b
|
|
return math.acos(dot / (a.length() * b.length()))
|
|
|
|
|
|
def angle_degrees(a,b):
|
|
return math.degrees(angle(a, b))
|
|
|
|
|
|
def signed_angle_around_axis(a,b, axis):
|
|
na = a**axis
|
|
nb = b**axis
|
|
|
|
v = angle(na,nb)
|
|
|
|
d = b*(a**axis)
|
|
|
|
if d < 0:
|
|
v =-v
|
|
|
|
return v
|
|
|
|
def signed_angle_degrees(a,b):
|
|
return 180/math.pi * signed_angle(a, b)
|
|
|
|
|
|
def rotate_vector_around_an_axis(theta, axis, v):
|
|
#print "# 1. rotate space about the z-axis so that the rotation axis lies in the xz-plane"
|
|
gamma = 0.0
|
|
if axis.y != 0:
|
|
if axis.x != 0:
|
|
gamma = -axis.x/abs(axis.x)*math.asin(axis.y/(math.sqrt(axis.x*axis.x + axis.y*axis.y)))
|
|
else:
|
|
gamma = math.pi/2.0
|
|
|
|
Rz = rotate_atoms_around_z_axis(gamma)
|
|
v = Rz * v
|
|
axis = Rz * axis
|
|
|
|
#print "# 2. rotate space about the y-axis so that the rotation axis lies along the z-axis"
|
|
beta = 0.0
|
|
if axis.x != 0:
|
|
beta = -axis.x/abs(axis.x)*math.acos(axis.z/math.sqrt(axis.x*axis.x + axis.z*axis.z))
|
|
Ry = rotate_atoms_around_y_axis(beta)
|
|
v = Ry * v
|
|
axis = Ry *axis
|
|
|
|
#print "# 3. perform the desired rotation by theta about the z-axis"
|
|
Rz = rotate_atoms_around_z_axis(theta)
|
|
v = Rz * v
|
|
|
|
#print "# 4. apply the inverse of step 2."
|
|
Ry = rotate_atoms_around_y_axis(-beta)
|
|
v = Ry * v
|
|
|
|
#print "# 5. apply the inverse of step 1."
|
|
Rz = rotate_atoms_around_z_axis(-gamma)
|
|
v = Rz * v
|
|
|
|
return v
|
|
|
|
def rotate_atoms_around_z_axis(angle):
|
|
Rz = matrix4x4(
|
|
a11i = math.cos(angle), a12i = -math.sin(angle), a13i = 0.0, a14i = 0.0,
|
|
a21i = math.sin(angle), a22i = math.cos(angle), a23i = 0.0, a24i = 0.0,
|
|
a31i = 0.0 , a32i = 0.0 , a33i = 1.0, a34i = 0.0,
|
|
a41i = 0.0 , a42i = 0.0 , a43i = 0.0, a44i = 1.0
|
|
)
|
|
|
|
return Rz
|
|
|
|
|
|
def rotate_atoms_around_y_axis(angle):
|
|
Ry = matrix4x4(
|
|
a11i = math.cos(angle), a12i = 0.0, a13i = math.sin(angle), a14i = 0.0,
|
|
a21i = 0.0 , a22i = 1.0, a23i = 0.0 , a24i = 0.0,
|
|
a31i = -math.sin(angle), a32i = 0.0, a33i = math.cos(angle), a34i = 0.0,
|
|
a41i = 0.0 , a42i = 0.0, a43i = 0.0 , a44i = 1.0
|
|
)
|
|
|
|
return Ry
|
|
|
|
|
|
|
|
class multi_vector:
|
|
def __init__(self, atom1=0, atom2=0):
|
|
self.vectors = []
|
|
self.keys = []
|
|
|
|
# store vectors for all configurations of atoms
|
|
if atom1!=0:
|
|
self.keys = lib.get_sorted_configurations(atom1.configurations.keys())
|
|
if atom2!=0:
|
|
keys2 = lib.get_sorted_configurations(atom2.configurations.keys())
|
|
if self.keys != keys2:
|
|
raise 'Cannot make multi vector: Atomic configurations mismatch for\n %s\n %s\n'%(atom1,atom2)
|
|
for key in self.keys:
|
|
atom1.setConfiguration(key)
|
|
if atom2!=0:
|
|
atom2.setConfiguration(key)
|
|
v = vector(atom1=atom1, atom2=atom2)
|
|
self.vectors.append(v)
|
|
#print(key,v)
|
|
return
|
|
|
|
def __getattribute__(self,name):
|
|
try:
|
|
return object.__getattribute__(self, name)
|
|
except AttributeError:
|
|
return self.do_job(name)
|
|
|
|
def __str__(self):
|
|
res = ''
|
|
for i in range(len(self.keys)):
|
|
res += '%s %s\n'%(self.keys[i], self.vectors[i])
|
|
return res
|
|
|
|
|
|
def do_job(self, job):
|
|
#print(job)
|
|
self.res = multi_vector()
|
|
for i in range(len(self.vectors)):
|
|
self.res.vectors.append(eval('self.vectors[%d].%s()'%(i,job)))
|
|
self.res.keys.append(self.keys[i])
|
|
return self.get_result
|
|
|
|
def get_result(self):
|
|
return self.res
|
|
|
|
def generic_operation(self, operation, other):
|
|
if self.keys != other.keys:
|
|
raise 'Incompatable keys'
|
|
|
|
self.res = multi_vector()
|
|
for i in range(len(self.vectors)):
|
|
self.res.vectors.append(eval('self.vectors[%d] %s other.vectors[%d]'%(i,operation,i)))
|
|
self.res.keys.append(self.keys[i])
|
|
return
|
|
|
|
def __add__(self, other):
|
|
self.generic_operation('+',other)
|
|
return self.res
|
|
|
|
def __sub__(self, other):
|
|
self.generic_operation('-',other)
|
|
return self.res
|
|
|
|
def __mul__(self, other):
|
|
self.generic_operation('*',other)
|
|
return self.res
|
|
|
|
def __pow__(self, other):
|
|
self.generic_operation('**',other)
|
|
return self.res
|
|
|
|
def generic_self_operation(self, operation):
|
|
return
|
|
|
|
def __neg__(self):
|
|
self.generic_operation('*',-1.0)
|
|
return self.res
|
|
|
|
def rescale(self, new_length):
|
|
self.res = multi_vector()
|
|
for i in range(len(self.vectors)):
|
|
self.res.vectors.append(self.vectors[i].rescale(new_length))
|
|
self.res.keys.append(self.keys[i])
|
|
return self.res
|
|
|
|
|
|
def rotate_multi_vector_around_an_axis(theta, axis, v):
|
|
""" both axis ans v must be multi_vectors """
|
|
|
|
if axis.keys != v.keys:
|
|
raise 'Incompatible keys in rotate multi_vector'
|
|
|
|
res = multi_vector()
|
|
for i in range(len(v.keys)):
|
|
res.vectors.append(rotate_vector_around_an_axis(theta, axis.vectors[i], v.vectors[i]))
|
|
res.keys.append(v.keys[i])
|
|
|
|
return res
|