from math import sqrt from operator import (add, sub, mul, div, floordiv, mod, pow as opow, truediv, lt, le, eq, ne, gt, ge) class VectorError(Exception): pass # base exception class InvalidDimension(VectorError): pass InvalidDimensionError = InvalidDimension ELEMENTS = dict(zip("xyzw", xrange(4))) class Vector(list): """ Vector class meant to mimic classical mathematical notation The difference is that column vectors are defined instead as just lists or tuples """ def __init__(self, *args): if len(args) == 1: args = args[0] list.__init__(self, args) if len(self) < 2: # no < 2D vectors self.extend([0]*(2-len(self))) self._assert_contents() # check for a valid vector def __repr__(self): return repr(tuple(map(float, self))) def _get_mag(self): return sqrt(sum(map(lambda x: x**2, self))) def _set_mag(self, value): self.normalize() self *= value return value length = magnitude = property(_get_mag, _set_mag, doc="The length of the vector") def _assert_contents(self): for i in xrange(len(self)): x = self[i] if not isinstance(x, (int, float, long, complex)): raise ValueError('Invalid Vector Value: (%r)'%x) elif isinstance(x, complex): pass else: self[i] = float(x) ### ==== LIST OPERATORS ==== ### def __setitem__(self, key, value): if isinstance(key, slice): stop = key.stop else: stop = key while len(self) < (stop+1): self.append(0) list.__setitem__(self, key, float(value)) #self._assert_contents() ### ==== SWIZZLING ==== ### def __getattr__(self, name): if not len([x for x in name if x in ELEMENTS]) == len(name): raise AttributeError("Vector object has no attribute '%s'"%name) v = [] for x in name: if len(self) <= ELEMENTS[x]: raise AttributeError("Vector object has no element '%s'"%x) v.append(self[ELEMENTS[x]]) return Vector(v) if len(v) > 1 else v[0] def __setattr__(self, name, value): made_up_of_elements = len([x for x in name if x in ELEMENTS])==len(name) if not made_up_of_elements and name not in self.__dict__: raise AttributeError("Vector object has no attribute '%s'"%name) elif len(name) > 1 and (not hasattr(value, '__iter__') or len(name) != len(value)): raise TypeError("Vector size does not match swizzle") elif name in self.__dict__: self.__dict__[name] = value if not hasattr(value, "__iter__"): value = [value] for i in range(len(name)): vname, val = name[i], value[i] idx = ELEMENTS[vname] self[idx] = float(val) ### ==== OPERATORS ==== ### def _test_length(self, other): if not len(self) == len(other): raise InvalidDimensionError def _apply_op(self, other, op): vec = Vector() self._test_length(other) for i in xrange(len(self)): vec[i] = op(self[i], other[i]) return vec def __add__(self, other): return self._apply_op(other, add) def __sub__(self, other): return self._apply_op(other, sub) def __mul__(self, other): if not isinstance(other, (int, float)): return self.dot(other) return Vector(map(lambda x: mul(x, other), self)) def __div__(self, other): return Vector(map(lambda x: div(x, other), self)) def __truediv__(self, other): return Vector(map(lambda x: truediv(x, other), self)) def __floordiv__(self, other): return Vector(map(lambda x: floordiv(x, other), self)) def __mod__(self, other): return Vector(map(lambda x: mod(x, other), self)) def __pow__(self, other, modulo=None): v = Vector(map(lambda x: opow(x, other), self)) if modulo is not None: v % modulo return v ## == Right Operators == ## def __rmul__(self, other): if not isinstance(other, (int, float)): return Vector(other).dot(self) return Vector(map(lambda x: mul(other, x), self)) def __rdiv__(self, other): return Vector(map(lambda x: div(other, x), self)) def __rtruediv__(self, other): return Vector(map(lambda x: truediv(other, x), self)) def __rfloordiv__(self, other): return Vector(map(lambda x: floordiv(other, x), self)) ## == Other Operators == ## def __iadd__(self, other): self = add(self, other) return self def __isub__(self, other): self = sub(self, other) return self def __imul__(self, other): self = mul(self, other) return self def __idiv__(self, other): self = div(self, other) return self def __itruediv__(self, other): self = truediv(self, other) return self def __ifloordiv__(self, other): self = floordiv(self, other) return self def __imod__(self, other): self = mod(self, other) return self def __ipow__(self, other, modulo=None): self = pow(self, other, modulo) return self def __neg__(self): return self * -1 def __pos__(self): return self * +1 def __abs__(self): return Vector(map(lambda x: abs(x), self)) def __invert__(self): return 1.0 / self ## == Comparison Operators == ## def __lt__(self, other): return lt(self.magnitude, other.magnitude) def __le__(self, other): return le(self.magnitude, other.magnitude) def __eq__(self, other): return all(self._apply_op(other, eq)) def __ne__(self, other): return all(self._apply_op(other, ne)) def __gt__(self, other): return gt(self.magnitude, other.magnitude) def __ge__(self, other): return ge(self.magnitude, other.magnitude) ## == Vector Operations == ## def apply(self, func): return Vector(map(func, self)) def copy(self): return Vector(iter(self)) def zero(self): for i in xrange(len(self)): self[i] = 0.0 return self def normalize(self): mag = self.magnitude for i in xrange(len(self)): self[i] /= mag return self def negate(self): self *= -1 return self def reflect(self, mirror): self = mirror - (2 * self * self.dot(mirror)) return self def _make_length(self, n): if len(self) < n: self.extend([0.0]*(n-len(self))) return True elif len(self) > n: del self[n:] return False def resize2D(self): self._make_length(2) return self def resize3D(self): self._make_length(3) return self def resize4D(self): if self._make_length(4): self[3] = 1.0 return self def cross(self, other): pass def dot(self, other): self._test_length(other) return sum([self[i] * other[i] for i in xrange(len(self))], 0.0) if __name__ == "__main__": v = Vector(1, 2, 3) u = Vector(5, 10, 15)