## vectors and associated methods

from math import cos, sin, acos, sqrt, pi

from random import random

# List of names imported from this module with import *

__all__ = ['adjust_axis', 'adjust_up', 'comp', 'cross', 'diff_angle', 'dot',

'hat', 'mag', 'mag2', 'norm', 'object_rotate', 'proj', 'rotate',

'vector']

class vector(object):

'vector class'

@staticmethod

def random():

return vector(-1.0 + 2.0*random(), -1.0 + 2.0*random(), -1.0 + 2.0*random())

def __init__(self, *args):

if len(args) == 3:

self._x = float(args[0]) # make sure it's a float; could be numpy.float64

self._y = float(args[1])

self._z = float(args[2])

elif len(args) == 1 and isinstance(args[0], vector): # make a copy of a vector

other = args[0]

self._x = other._x

self._y = other._y

self._z = other._z

else:

raise TypeError('A vector needs 3 components.')

self.on_change = self.ignore

def ignore(self):

pass

@property

def value(self):

return [self._x, self._y, self._z]

@value.setter

def value(self,other): ## ensures a copy; other is a vector

self._x = other._x

self._y = other._y

self._z = other._z

def __neg__(self):

return vector(-self._x, -self._y, -self._z)

def __pos__(self):

return self

def __str__(self):

return '<{:.6g}, {:.6g}, {:.6g}>'.format(self._x, self._y, self._z)

def __repr__(self):

return 'vector({:.6g}, {:.6g}, {:.6g})'.format(self._x, self._y, self._z)

def __add__(self, other):

if type(other) is vector:

return vector(self._x + other._x, self._y + other._y, self._z + other._z)

return NotImplemented

def __sub__(self, other):

if type(other) is vector:

return vector(self._x - other._x, self._y - other._y, self._z - other._z)

return NotImplemented

def __truediv__(self, other): # used by Python 3, and by Python 2 in the presence of __future__ division

if isinstance(other, (float, int)):

return vector(self._x / other, self._y / other, self._z / other)

return NotImplemented

def __mul__(self, other):

if isinstance(other, (float, int)):

return vector(self._x * other, self._y * other, self._z * other)

return NotImplemented

def __rmul__(self, other):

if isinstance(other, (float, int)):

return vector(self._x * other, self._y * other, self._z * other)

return NotImplemented

def __eq__(self,other):

if type(self) is vector and type(other) is vector:

return self.equals(other)

return False

def __ne__(self,other):

if type(self) is vector and type(other) is vector:

return not self.equals(other)

return True

@property

def x(self):

return self._x

@x.setter

def x(self,value):

self._x = value

self.on_change()

@property

def y(self):

return self._y

@y.setter

def y(self,value):

self._y = value

self.on_change()

@property

def z(self):

return self._z

@z.setter

def z(self,value):

self._z = value

self.on_change()

@property

def mag(self):

return sqrt(self._x**2+self._y**2+self._z**2)

@mag.setter

def mag(self,value):

normA = self.hat

self._x = value * normA._x

self._y = value * normA._y

self._z = value * normA._z

self.on_change()

@property

def mag2(self):

return self._x**2+self._y**2+self._z**2

@mag2.setter

def mag2(self,value):

normA = self.hat

v = sqrt(value)

self._x = v * normA._x

self._y = v * normA._y

self._z = v * normA._z

self.on_change()

@property

def hat(self):

smag = self.mag

if ( smag > 0. ):

return self / smag

else:

return vector(0., 0., 0.)

@hat.setter

def hat(self, value):

smg = self.mag

normA = value.hat

self._x = smg * normA._x

self._y = smg * normA._y

self._z = smg * normA._z

self.on_change()

def norm(self):

return self.hat

def dot(self,other):

return ( self._x*other._x + self._y*other._y + self._z*other._z )

def cross(self,other):

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 proj(self,other):

normB = other.hat

return self.dot(normB) * normB

def equals(self,other):

return self._x == other._x and self._y == other._y and self._z == other._z

def comp(self,other): ## result is a scalar

normB = other.hat

return self.dot(normB)

def diff_angle(self, other):

a = self.hat.dot(other.hat)

if a > 1: # avoid roundoff problems

return 0

if a < -1:

return pi

return acos(a)

def rotate(self, angle=0., axis=None):

if axis is None:

u = vector(0,0,1)

else:

u = axis.hat

c = cos(angle)

s = sin(angle)

t = 1.0 - c

x = u._x

y = u._y

z = u._z

m11 = t*x*x+c

m12 = t*x*y-z*s

m13 = t*x*z+y*s

m21 = t*x*y+z*s

m22 = t*y*y+c

m23 = t*y*z-x*s

m31 = t*x*z-y*s

m32 = t*y*z+x*s

m33 = t*z*z+c

sx = self._x

sy = self._y

sz = self._z

return vector( (m11*sx + m12*sy + m13*sz),

(m21*sx + m22*sy + m23*sz),

(m31*sx + m32*sy + m33*sz) )

def rotate_in_place(self, angle=0., axis=None):

if axis is None:

u = vector(0,0,1)

else:

u = axis.hat

c = cos(angle)

s = sin(angle)

t = 1.0 - c

x = u._x

y = u._y

z = u._z

m11 = t*x*x+c

m12 = t*x*y-z*s

m13 = t*x*z+y*s

m21 = t*x*y+z*s

m22 = t*y*y+c

m23 = t*y*z-x*s

m31 = t*x*z-y*s

m32 = t*y*z+x*s

m33 = t*z*z+c

sx = self._x

sy = self._y

sz = self._z

self._x = m11*sx + m12*sy + m13*sz

self._y = m21*sx + m22*sy + m23*sz

self._z = m31*sx + m32*sy + m33*sz

def object_rotate(objaxis, objup, angle, axis):

u = axis.hat

c = cos(angle)

s = sin(angle)

t = 1.0 - c

x = u._x

y = u._y

z = u._z

m11 = t*x*x+c

m12 = t*x*y-z*s

m13 = t*x*z+y*s

m21 = t*x*y+z*s

m22 = t*y*y+c

m23 = t*y*z-x*s

m31 = t*x*z-y*s

m32 = t*y*z+x*s

m33 = t*z*z+c

sx = objaxis._x

sy = objaxis._y

sz = objaxis._z

objaxis._x = m11*sx + m12*sy + m13*sz # avoid creating a new vector object

objaxis._y = m21*sx + m22*sy + m23*sz

objaxis._z = m31*sx + m32*sy + m33*sz

sx = objup._x

sy = objup._y

sz = objup._z

objup._x = m11*sx + m12*sy + m13*sz

objup._y = m21*sx + m22*sy + m23*sz

objup._z = m31*sx + m32*sy + m33*sz

def mag(A):

return A.mag

def mag2(A):

return A.mag2

def norm(A):

return A.hat

def hat(A):

return A.hat

def dot(A,B):

return A.dot(B)

def cross(A,B):

return A.cross(B)

def proj(A,B):

return A.proj(B)

def comp(A,B):

return A.comp(B)

def diff_angle(A,B):

return A.diff_angle(B)

def rotate(A, angle=0., axis = None):

return A.rotate(angle,axis)

def adjust_up(oldaxis, newaxis, up, save_oldaxis): # adjust up when axis is changed

if abs(newaxis._x) + abs(newaxis._y) + abs(newaxis._z) == 0:

# If axis has changed to <0,0,0>, must save the old axis to restore later

if save_oldaxis is None: save_oldaxis = oldaxis

return save_oldaxis

if save_oldaxis is not None:

# Restore saved oldaxis now that newaxis is nonzero

oldaxis._x = save_oldaxis._x # avoid creating a new vector

oldaxis._y = save_oldaxis._y

oldaxis._z = save_oldaxis._z

save_oldaxis = None

if newaxis.dot(up) != 0: # axis and up not orthogonal

angle = oldaxis.diff_angle(newaxis)

if angle > 1e-6: # smaller angles lead to catastrophes

# If axis is flipped 180 degrees, cross(oldaxis,newaxis) is <0,0,0>:

if abs(angle-pi) < 1e-6:

up._x = -up._x

up._y = -up._y

up._z = -up._z

else:

rotaxis = oldaxis.cross(newaxis)

up.rotate_in_place(angle=angle, axis=rotaxis) # avoid creating a new vector

oldaxis._x = newaxis._x # avoid creating a new vector

oldaxis._y = newaxis._y

oldaxis._z = newaxis._z

def adjust_axis(oldup, newup, axis, save_oldup): # adjust axis when up is changed

if abs(newup._x) + abs(newup._y) + abs(newup._z) == 0:

# If up will be set to <0,0,0>, must save the old up to restore later

if save_oldup is None: save_oldup = oldup

return save_oldup

if save_oldup is not None:

# Restore saved oldup now that newup is nonzero

oldup = save_oldup

save_oldup = None

if newup.dot(axis) != 0: # axis and up not orthogonal

angle = oldup.diff_angle(newup)

if angle > 1e-6: # smaller angles lead to catastrophes

# If up is flipped 180 degrees, cross(oldup,newup) is <0,0,0>:

if abs(angle-pi) < 1e-6:

axis._x = -axis._x

axis._y = -axis._y

axis._z = -axis._z

else:

rotaxis = oldup.cross(newup)

axis.rotate_in_place(angle=angle, axis=rotaxis) # avoid creating a new vector

oldup._x = newup._x # avoid creating a new vector

oldup._y = newup._y

oldup._z = newup._z

return save_oldup