Last active
October 22, 2024 12:31
-
-
Save mcleonard/5351452 to your computer and use it in GitHub Desktop.
A vector class in pure python.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
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) |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
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!