SHARE
TWEET

vector.py

piguy123 Mar 26th, 2013 47 Never
  1. from math import sqrt
  2. from operator import (add, sub, mul, div, floordiv, mod, pow as opow, truediv,
  3.                       lt, le, eq, ne, gt, ge)
  4.  
  5. class VectorError(Exception): pass # base exception
  6. class InvalidDimension(VectorError): pass
  7. InvalidDimensionError = InvalidDimension
  8.  
  9. ELEMENTS = dict(zip("xyzw", xrange(4)))
  10.  
  11. class Vector(list):
  12.     """
  13.    Vector class meant to mimic classical mathematical notation
  14.  
  15.    The difference is that column vectors are defined instead as just lists or tuples
  16.    """
  17.  
  18.     def __init__(self, *args):
  19.         if len(args) == 1:
  20.             args = args[0]
  21.         list.__init__(self, args)
  22.  
  23.         if len(self) < 2: # no < 2D vectors
  24.             self.extend([0]*(2-len(self)))
  25.         self._assert_contents() # check for a valid vector
  26.        
  27.     def __repr__(self):
  28.         return repr(tuple(map(float, self)))
  29.  
  30.     def _get_mag(self):
  31.         return sqrt(sum(map(lambda x: x**2, self)))
  32.  
  33.     def _set_mag(self, value):
  34.         self.normalize()
  35.         self *= value
  36.         return value
  37.  
  38.     length = magnitude = property(_get_mag, _set_mag, doc="The length of the vector")
  39.  
  40.     def _assert_contents(self):
  41.         for i in xrange(len(self)):
  42.             x = self[i]
  43.             if not isinstance(x, (int, float, long, complex)):
  44.                 raise ValueError('Invalid Vector Value: (%r)'%x)
  45.             elif isinstance(x, complex):
  46.                 pass
  47.             else:
  48.                 self[i] = float(x)
  49.        
  50.  
  51.     ### ==== LIST OPERATORS ==== ###
  52.  
  53.     def __setitem__(self, key, value):
  54.         if isinstance(key, slice):
  55.             stop = key.stop
  56.         else:
  57.             stop = key
  58.        
  59.         while len(self) < (stop+1):
  60.             self.append(0)
  61.            
  62.         list.__setitem__(self, key, float(value))
  63.         #self._assert_contents()
  64.  
  65.     ### ==== SWIZZLING ==== ###
  66.            
  67.     def __getattr__(self, name):
  68.         if not len([x for x in name if x in ELEMENTS]) == len(name):
  69.             raise AttributeError("Vector object has no attribute '%s'"%name)
  70.         v = []
  71.         for x in name:
  72.             if len(self) <= ELEMENTS[x]:
  73.                 raise AttributeError("Vector object has no element '%s'"%x)
  74.             v.append(self[ELEMENTS[x]])
  75.         return Vector(v) if len(v) > 1 else v[0]
  76.  
  77.     def __setattr__(self, name, value):
  78.         made_up_of_elements = len([x for x in name if x in ELEMENTS])==len(name)
  79.         if not made_up_of_elements and name not in self.__dict__:
  80.             raise AttributeError("Vector object has no attribute '%s'"%name)
  81.         elif len(name) > 1 and (not hasattr(value, '__iter__') or len(name) != len(value)):
  82.             raise TypeError("Vector size does not match swizzle")
  83.         elif name in self.__dict__:
  84.             self.__dict__[name] = value
  85.         if not hasattr(value, "__iter__"):
  86.             value = [value]
  87.         for i in range(len(name)):
  88.             vname, val = name[i], value[i]
  89.             idx = ELEMENTS[vname]
  90.             self[idx] = float(val)
  91.  
  92.     ### ==== OPERATORS ==== ###
  93.  
  94.     def _test_length(self, other):
  95.         if not len(self) == len(other):
  96.             raise InvalidDimensionError
  97.  
  98.     def _apply_op(self, other, op):
  99.         vec = Vector()
  100.         self._test_length(other)
  101.         for i in xrange(len(self)):
  102.             vec[i] = op(self[i], other[i])
  103.         return vec
  104.    
  105.     def __add__(self, other):
  106.         return self._apply_op(other, add)
  107.    
  108.     def __sub__(self, other):
  109.         return self._apply_op(other, sub)
  110.  
  111.     def __mul__(self, other):
  112.         if not isinstance(other, (int, float)):
  113.             return self.dot(other)
  114.         return Vector(map(lambda x: mul(x, other), self))
  115.  
  116.     def __div__(self, other):
  117.         return Vector(map(lambda x: div(x, other), self))
  118.  
  119.     def __truediv__(self, other):
  120.         return Vector(map(lambda x: truediv(x, other), self))
  121.  
  122.     def __floordiv__(self, other):
  123.         return Vector(map(lambda x: floordiv(x, other), self))
  124.  
  125.     def __mod__(self, other):
  126.         return Vector(map(lambda x: mod(x, other), self))
  127.  
  128.     def __pow__(self, other, modulo=None):
  129.         v = Vector(map(lambda x: opow(x, other), self))
  130.         if modulo is not None:
  131.             v % modulo
  132.         return v
  133.  
  134.     ## == Right Operators == ##
  135.  
  136.     def __rmul__(self, other):
  137.         if not isinstance(other, (int, float)):
  138.             return Vector(other).dot(self)
  139.         return Vector(map(lambda x: mul(other, x), self))
  140.    
  141.     def __rdiv__(self, other):
  142.         return Vector(map(lambda x: div(other, x), self))
  143.  
  144.     def __rtruediv__(self, other):
  145.         return Vector(map(lambda x: truediv(other, x), self))
  146.  
  147.     def __rfloordiv__(self, other):
  148.         return Vector(map(lambda x: floordiv(other, x), self))
  149.    
  150.     ## == Other Operators == ##
  151.  
  152.     def __iadd__(self, other):
  153.         self = add(self, other)
  154.         return self
  155.  
  156.     def __isub__(self, other):
  157.         self = sub(self, other)
  158.         return self
  159.  
  160.     def __imul__(self, other):
  161.         self = mul(self, other)
  162.         return self
  163.  
  164.     def __idiv__(self, other):
  165.         self = div(self, other)
  166.         return self
  167.  
  168.     def __itruediv__(self, other):
  169.         self = truediv(self, other)
  170.         return self
  171.  
  172.     def __ifloordiv__(self, other):
  173.         self = floordiv(self, other)
  174.         return self
  175.  
  176.     def __imod__(self, other):
  177.         self = mod(self, other)
  178.         return self
  179.  
  180.     def __ipow__(self, other, modulo=None):
  181.         self = pow(self, other, modulo)
  182.         return self
  183.  
  184.     def __neg__(self):
  185.         return self * -1
  186.  
  187.     def __pos__(self):
  188.         return self * +1
  189.  
  190.     def __abs__(self):
  191.         return Vector(map(lambda x: abs(x), self))
  192.  
  193.     def __invert__(self):
  194.         return 1.0 / self
  195.  
  196.     ## == Comparison Operators == ##
  197.  
  198.     def __lt__(self, other):
  199.         return lt(self.magnitude, other.magnitude)
  200.        
  201.     def __le__(self, other):
  202.         return le(self.magnitude, other.magnitude)
  203.        
  204.     def __eq__(self, other):
  205.         return all(self._apply_op(other, eq))
  206.        
  207.     def __ne__(self, other):
  208.         return all(self._apply_op(other, ne))
  209.        
  210.     def __gt__(self, other):
  211.         return gt(self.magnitude, other.magnitude)
  212.        
  213.     def __ge__(self, other):
  214.         return ge(self.magnitude, other.magnitude)
  215.  
  216.     ## == Vector Operations == ##
  217.  
  218.     def apply(self, func):
  219.         return Vector(map(func, self))
  220.  
  221.     def copy(self):
  222.         return Vector(iter(self))
  223.  
  224.     def zero(self):
  225.         for i in xrange(len(self)):
  226.             self[i] = 0.0
  227.         return self
  228.  
  229.     def normalize(self):
  230.         mag = self.magnitude
  231.         for i in xrange(len(self)):
  232.             self[i] /= mag
  233.         return self
  234.  
  235.     def negate(self):
  236.         self *= -1
  237.         return self
  238.  
  239.     def reflect(self, mirror):
  240.         self = mirror - (2 * self * self.dot(mirror))
  241.         return self
  242.  
  243.     def _make_length(self, n):
  244.         if len(self) < n:
  245.             self.extend([0.0]*(n-len(self)))
  246.             return True
  247.         elif len(self) > n:
  248.             del self[n:]
  249.             return False
  250.    
  251.     def resize2D(self):
  252.         self._make_length(2)
  253.         return self
  254.  
  255.     def resize3D(self):
  256.         self._make_length(3)
  257.         return self
  258.  
  259.     def resize4D(self):
  260.         if self._make_length(4):
  261.             self[3] = 1.0
  262.         return self
  263.    
  264.     def cross(self, other):
  265.         pass
  266.  
  267.     def dot(self, other):
  268.         self._test_length(other)
  269.         return sum([self[i] * other[i] for i in xrange(len(self))], 0.0)
  270.    
  271. if __name__ == "__main__":
  272.     v = Vector(1, 2, 3)
  273.     u = Vector(5, 10, 15)
RAW Paste Data
Top