-
-
Save mcleonard/5351452 to your computer and use it in GitHub Desktop.
import math | |
class Vector(object): | |
def __init__(self, *args): | |
""" Create a vector, example: v = Vector(1,2) """ | |
if len(args)==0: self.values = (0,0) | |
else: self.values = args | |
def norm(self): | |
""" Returns the norm (length, magnitude) of the vector """ | |
return math.sqrt(sum( x*x for x in self )) | |
def argument(self, radians=False): | |
""" Returns the argument of the vector, the angle clockwise from +y. In degress by default, | |
set radians=True to get the result in radians. This only works for 2D vectors. """ | |
arg_in_rad = math.acos(Vector(0, 1)*self/self.norm()) | |
if radians: | |
return arg_in_rad | |
arg_in_deg = math.degrees(arg_in_rad) | |
if self.values[0] < 0: | |
return 360 - arg_in_deg | |
else: | |
return arg_in_deg | |
def normalize(self): | |
""" Returns a normalized unit vector """ | |
norm = self.norm() | |
normed = tuple( x / norm for x in self ) | |
return self.__class__(*normed) | |
def rotate(self, theta): | |
""" Rotate this vector. If passed a number, assumes this is a | |
2D vector and rotates by the passed value in degrees. Otherwise, | |
assumes the passed value is a list acting as a matrix which rotates the vector. | |
""" | |
if isinstance(theta, (int, float)): | |
# So, if rotate is passed an int or a float... | |
if len(self) != 2: | |
raise ValueError("Rotation axis not defined for greater than 2D vector") | |
return self._rotate2D(theta) | |
matrix = theta | |
if not all(len(row) == len(self) for row in matrix) or not len(matrix)==len(self): | |
raise ValueError("Rotation matrix must be square and same dimensions as vector") | |
return self.matrix_mult(matrix) | |
def _rotate2D(self, theta): | |
""" Rotate this vector by theta in degrees. | |
Returns a new vector. | |
""" | |
theta = math.radians(theta) | |
# Just applying the 2D rotation matrix | |
dc, ds = math.cos(theta), math.sin(theta) | |
x, y = self.values | |
x, y = dc*x - ds*y, ds*x + dc*y | |
return self.__class__(x, y) | |
def matrix_mult(self, matrix): | |
""" Multiply this vector by a matrix. Assuming matrix is a list of lists. | |
Example: | |
mat = [[1,2,3],[-1,0,1],[3,4,5]] | |
Vector(1,2,3).matrix_mult(mat) -> (14, 2, 26) | |
""" | |
if not all(len(row) == len(self) for row in matrix): | |
raise ValueError('Matrix must match vector dimensions') | |
# Grab a row from the matrix, make it a Vector, take the dot product, | |
# and store it as the first component | |
product = tuple(Vector(*row)*self for row in matrix) | |
return self.__class__(*product) | |
def inner(self, vector): | |
""" Returns the dot product (inner product) of self and another vector | |
""" | |
if not isinstance(vector, Vector): | |
raise ValueError('The dot product requires another vector') | |
return sum(a * b for a, b in zip(self, vector)) | |
def __mul__(self, other): | |
""" Returns the dot product of self and other if multiplied | |
by another Vector. If multiplied by an int or float, | |
multiplies each component by other. | |
""" | |
if isinstance(other, Vector): | |
return self.inner(other) | |
elif isinstance(other, (int, float)): | |
product = tuple( a * other for a in self ) | |
return self.__class__(*product) | |
else: | |
raise ValueError("Multiplication with type {} not supported".format(type(other))) | |
def __rmul__(self, other): | |
""" Called if 4 * self for instance """ | |
return self.__mul__(other) | |
def __truediv__(self, other): | |
if isinstance(other, Vector): | |
divided = tuple(self[i] / other[i] for i in range(len(self))) | |
elif isinstance(other, (int, float)): | |
divided = tuple( a / other for a in self ) | |
else: | |
raise ValueError("Division with type {} not supported".format(type(other))) | |
return self.__class__(*divided) | |
def __add__(self, other): | |
""" Returns the vector addition of self and other """ | |
if isinstance(other, Vector): | |
added = tuple( a + b for a, b in zip(self, other) ) | |
elif isinstance(other, (int, float)): | |
added = tuple( a + other for a in self ) | |
else: | |
raise ValueError("Addition with type {} not supported".format(type(other))) | |
return self.__class__(*added) | |
def __radd__(self, other): | |
""" Called if 4 + self for instance """ | |
return self.__add__(other) | |
def __sub__(self, other): | |
""" Returns the vector difference of self and other """ | |
if isinstance(other, Vector): | |
subbed = tuple( a - b for a, b in zip(self, other) ) | |
elif isinstance(other, (int, float)): | |
subbed = tuple( a - other for a in self ) | |
else: | |
raise ValueError("Subtraction with type {} not supported".format(type(other))) | |
return self.__class__(*subbed) | |
def __rsub__(self, other): | |
""" Called if 4 - self for instance """ | |
return self.__sub__(other) | |
def __iter__(self): | |
return self.values.__iter__() | |
def __len__(self): | |
return len(self.values) | |
def __getitem__(self, key): | |
return self.values[key] | |
def __repr__(self): | |
return str(self.values) |
I'm somewhat surprised people are using this class. I'll go through the comments and see if there are things I should update.
Are you kidding? This is a VERY helpful piece of code! There are occasions where numpy is not allowed but geometry calculations are needed...
BTW: I have taken your class and changed quite a bit, which you have now included in a somehow other way (but with more or less the same functionality). But some additional functions were very helpful for me, I want to share the code with you. But be aware that they might not run out of the box since you might have used a little other variable names than in my personal fork ;)
def get_x(self):
""" Return the 1st element """
if len(self.values) > 0:
return self.values[0]
else:
return None
x = property(get_x)
def get_y(self):
""" Return the 2nd element """
if len(self.values) > 1:
return self.values[1]
else:
return None
y = property(get_y)
def get_z(self):
""" Return the 3rd element """
if len(self.values) > 2:
return self.values[2]
else:
return None
z = property(get_z)
def distance(self, other, polarcoordinates=False):
"""Returns the distance of self and other.
If polarcoordinates: Returns the 2D norm of the distance
vector in x-y-plane
"""
return (self - other).norm(polarcoordinates)
def is_parallel(self, other, abs_tol=1e-3):
""" Returns true if the angle between self and other is close to
0° or 180° with abs_tol tolerance. """
angle = self.angle(other, True)
return (math.isclose(angle, 0, abs_tol=abs_tol) \
or math.isclose(angle, 180, abs_tol=abs_tol))
def __eq__(self, other, abs_tol=1e-3):
""" Compares self with other including a tolerance """
isequal = []
for i in range(len(self.values)):
isequal.append(math.isclose(self.values[i], other.values[i],
abs_tol=abs_tol))
return all(isequal)
def __ne__(self, other, abs_tol=1e-3):
""" Compares self with other including a tolerance """
return not self.__eq__(other, abs_tol)
def __call__(self, idx=None):
""" Returns the values or only one element if an index is given """
if idx is None:
return self.values
elif idx < len(self.values) and isinstance(idx, int):
return self.values[idx]
else:
return None
Hi @mcleonard ,
where is the license specified?
can you please include license type in the header of the gist?
To anyone who happens to still be using this: There's a bug in rsub: the output needs to multiplied by -1, because right now [1,2] - [3,3]
gives the same result as [3,3] - [1,2]
. This made simulated golfers run away from the ball they were supposed to be hitting!
I updated the class a bit to work with Python 3 better (using
__truediv__
instead of__div__
) and improved some of the code.