Although NumPy offers a faster option, it is still instructive to code a class for vectors in pure Python. The following code defines the Vector2D
class and tests it for various operations.
import math
class Vector2D:
"""A two-dimensional vector with Cartesian coordinates."""
def __init__(self, x, y):
self.x, self.y = x, y
def __str__(self):
"""Human-readable string representation of the vector."""
return '{:g}i + {:g}j'.format(self.x, self.y)
def __repr__(self):
"""Unambiguous string representation of the vector."""
return repr((self.x, self.y))
def dot(self, other):
"""The scalar (dot) product of self and other. Both must be vectors."""
if not isinstance(other, Vector2D):
raise TypeError('Can only take dot product of two Vector2D objects')
return self.x * other.x + self.y * other.y
# Alias the __matmul__ method to dot so we can use a @ b as well as a.dot(b).
__matmul__ = dot
def __sub__(self, other):
"""Vector subtraction."""
return Vector2D(self.x - other.x, self.y - other.y)
def __add__(self, other):
"""Vector addition."""
return Vector2D(self.x + other.x, self.y + other.y)
def __mul__(self, scalar):
"""Multiplication of a vector by a scalar."""
if isinstance(scalar, int) or isinstance(scalar, float):
return Vector2D(self.x*scalar, self.y*scalar)
raise NotImplementedError('Can only multiply Vector2D by a scalar')
def __rmul__(self, scalar):
"""Reflected multiplication so vector * scalar also works."""
return self.__mul__(scalar)
def __neg__(self):
"""Negation of the vector (invert through origin.)"""
return Vector2D(-self.x, -self.y)
def __truediv__(self, scalar):
"""True division of the vector by a scalar."""
return Vector2D(self.x / scalar, self.y / scalar)
def __mod__(self, scalar):
"""One way to implement modulus operation: for each component."""
return Vector2D(self.x % scalar, self.y % scalar)
def __abs__(self):
"""Absolute value (magnitude) of the vector."""
return math.sqrt(self.x**2 + self.y**2)
def distance_to(self, other):
"""The distance between vectors self and other."""
return abs(self - other)
def to_polar(self):
"""Return the vector's components in polar coordinates."""
return self.__abs__(), math.atan2(self.y, self.x)
if __name__ == '__main__':
v1 = Vector2D(2, 5/3)
v2 = Vector2D(3, -1.5)
print('v1 = ', v1)
print('repr(v2) = ', repr(v2))
print('v1 + v2 = ', v1 + v2)
print('v1 - v2 = ', v1 - v2)
print('abs(v2 - v1) = ', abs(v2 - v1))
print('-v2 = ', -v2)
print('v1 * 3 = ', v1 * 3)
print('7 * v2 = ', 7 * v1)
print('v2 / 2.5 = ', v2 / 2.5)
print('v1 % 1 = ', v1 % 1)
print('v1.dot(v2) = v1 @ v2 = ', v1 @ v2)
print('v1.distance_to(v2) = ',v1.distance_to(v2))
print('v1 as polar vector, (r, theta) =', v1.to_polar())
The output should be:
v1 = 2i + 1.66667j
repr(v2) = (3, -1.5)
v1 + v2 = 5i + 0.166667j
v1 - v2 = -1i + 3.16667j
abs(v2 - v1) = 3.3208098075285464
-v2 = -3i + 1.5j
v1 * 3 = 6i + 5j
7 * v2 = 14i + 11.6667j
v2 / 2.5 = 1.2i + -0.6j
v1 % 1 = 0i + 0.666667j
v1.dot(v2) = v1 @ v2 = 3.5
v1.distance_to(v2) = 3.3208098075285464
v1 as polar vector, (r, theta) = (2.6034165586355518, 0.6947382761967033)