143 lines
4.4 KiB
Python
143 lines
4.4 KiB
Python
import math
|
|
|
|
class Vector:
|
|
''' a class that looks like a list, but has some extras for
|
|
vector operations
|
|
'''
|
|
def __init__(self, x=0, y=0, z=0):
|
|
if type(x).__name__ == 'list':
|
|
if len(x) == 3:
|
|
self.x = x[0]
|
|
self.y = x[1]
|
|
self.z = x[2]
|
|
elif len(x) == 2:
|
|
self.x = x[0]
|
|
self.y = x[1]
|
|
self.z = 0
|
|
elif len(x) == 1:
|
|
self.x = x[0]
|
|
self.y = 0
|
|
self.z = 0
|
|
elif type(x).__name__ in ['int', 'float', 'float64']:
|
|
self.x = x
|
|
self.y = y
|
|
self.z = z
|
|
else:
|
|
#TODO: make proper type check here!
|
|
raise TypeError('Vector() doesn\'t accept {} as type.'.format(type(x).__name__))
|
|
|
|
self.v = [self.x, self.y, self.z]
|
|
|
|
#
|
|
# standard maths and python operations
|
|
|
|
|
|
def __add__(self, other):
|
|
return Vector(self.x + other.x, self.y + other.y, self.z + other.z)
|
|
|
|
def __radd__(self, other):
|
|
return self.__add__(other)
|
|
|
|
def __sub__(self, other):
|
|
return self.__add__(other * -1)
|
|
|
|
def __mul__(self, other):
|
|
if type(other).__name__ == 'int' or type(other).__name__ == 'float':
|
|
return Vector([i*other for i in self.v])
|
|
else:
|
|
raise TypeError('__mul__ only works with int and float as 2nd operand!')
|
|
|
|
def __rmul__(self, other):
|
|
return self.__mul__(other)
|
|
|
|
def __div__(self, other):
|
|
return self.__truediv__(other)
|
|
|
|
def __truediv__(self, other):
|
|
#TODO: make it with __mul__
|
|
if type(other).__name__ == 'int' or type(other).__name__ == 'float':
|
|
return Vector(self.x / other, self.y / other, self.z / other, )
|
|
else:
|
|
raise TypeError('__truediv__ only works with int and float as 2nd operand!')
|
|
|
|
def __len__(self):
|
|
return 3
|
|
|
|
def __repr__(self):
|
|
return str(self.v)
|
|
|
|
def __getitem__(self, key):
|
|
return [self.x, self.y, self.z][key]
|
|
|
|
def getlist(self):
|
|
return [self.x, self.y, self.z]
|
|
|
|
|
|
|
|
#
|
|
# special vetor operations
|
|
|
|
def magn(self):
|
|
''' returns the "length" of this vector '''
|
|
return math.sqrt(math.pow(self.x,2) + math.pow(self.y,2) + math.pow(self.z,2))
|
|
|
|
def magn2d(self):
|
|
''' returns the "length" of this vector, omitting the z comonent'''
|
|
return math.sqrt(math.pow(self.x,2) + math.pow(self.y,2))
|
|
|
|
def norm(self):
|
|
''' returns a vector that has lenght of 1 '''
|
|
return Vector([i/self.magn() for i in self.v])
|
|
|
|
def dotp(self, other):
|
|
''' return the dot product(skalarprodukt) '''
|
|
return self.x * other.x + self.y * other.y + self.z * other.z
|
|
|
|
def proj(self, other):
|
|
''' return the projection of this vector on the 'other' vector '''
|
|
dp = self.dotp(other)
|
|
nrm = other / math.pow(other.magn(),2)
|
|
return dp * nrm
|
|
|
|
def rot(self, a):
|
|
''' return the vector that has been rotated around its z-axis '''
|
|
x = self.x * math.cos(a) - self.y * math.sin(a)
|
|
y = self.x * math.sin(a) + self.y * math.cos(a)
|
|
return Vector(x, y, self.z)
|
|
|
|
def angle(self, a):
|
|
''' return angle between two vectors
|
|
(but it lacks orientation!) '''
|
|
return math.acos(self.dotp(a) / ( self.magn() * a.magn()))
|
|
|
|
|
|
def tospherical(self):
|
|
''' return cartesian self to spherical coordinates
|
|
returns: radial, azimuth, polar '''
|
|
|
|
radial = math.sqrt(self.x **2 + self.y **2 + self.z **2 )
|
|
azimuth = math.acos(self.z / radial)
|
|
polar = math.atan2(self.y, self.x)
|
|
|
|
return [radial, azimuth, polar]
|
|
|
|
|
|
def tocartesian(self):
|
|
''' if for some reason this vector is defined in sperical coordinates,
|
|
this function will return the corresponding cartesian coordinates.
|
|
'''
|
|
x = self.x * math.sin(self.y) * math.cos(self.z)
|
|
y = self.x * math.sin(self.y) * math.sin(self.z)
|
|
z = self.x * math.cos(self.y)
|
|
return Vector(x, y, z)
|
|
|
|
|
|
#
|
|
# testing area
|
|
if __name__ == "__main__":
|
|
|
|
v0 = Vector(1,1,1)
|
|
|
|
#print Vector(v0.tospherical()).tocartesian()
|
|
|
|
|