Skip to content

Instantly share code, notes, and snippets.

@mcleonard
Last active February 22, 2024 12:30
Show Gist options
  • Save mcleonard/5351452 to your computer and use it in GitHub Desktop.
Save mcleonard/5351452 to your computer and use it in GitHub Desktop.
A vector class in pure python.
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)
@hillexed
Copy link

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!

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment