| 1 | |
|---|
| 2 | import math |
|---|
| 3 | |
|---|
| 4 | class Vector: |
|---|
| 5 | |
|---|
| 6 | rep = None |
|---|
| 7 | |
|---|
| 8 | def __init__(self, x, y=None, z=None, w=None): |
|---|
| 9 | if y == None: |
|---|
| 10 | self.rep = tuple(x) |
|---|
| 11 | elif z == None: |
|---|
| 12 | self.rep = (x, y) |
|---|
| 13 | elif w == None: |
|---|
| 14 | self.rep = (x, y, z) |
|---|
| 15 | else: |
|---|
| 16 | self.rep = (x, y, z, w) |
|---|
| 17 | |
|---|
| 18 | # math functions |
|---|
| 19 | |
|---|
| 20 | def __add__(self, other): |
|---|
| 21 | return Vector(map(lambda a, b: a + b, self.rep, other.rep)) |
|---|
| 22 | |
|---|
| 23 | def __sub__(self, other): |
|---|
| 24 | return Vector(map(lambda a, b: a - b, self.rep, other.rep)) |
|---|
| 25 | |
|---|
| 26 | def __mul__(self, value): |
|---|
| 27 | return Vector(map(lambda x: x * value, self.rep)) |
|---|
| 28 | |
|---|
| 29 | def __div__(self, value): |
|---|
| 30 | return Vector(map(lambda x: x / value, self.rep)) |
|---|
| 31 | |
|---|
| 32 | def __abs__(self): |
|---|
| 33 | "absolute gives the length" |
|---|
| 34 | return math.sqrt(self % self) |
|---|
| 35 | |
|---|
| 36 | def __neg__(self): |
|---|
| 37 | return Vector(map(lambda x: -x, self.rep)) |
|---|
| 38 | |
|---|
| 39 | def __mod__(self, other): |
|---|
| 40 | "dot product" |
|---|
| 41 | return self.dot(other) |
|---|
| 42 | |
|---|
| 43 | def __xor__(self, other): |
|---|
| 44 | "3d cross product" |
|---|
| 45 | return self.cross(other) |
|---|
| 46 | |
|---|
| 47 | # coord access functions |
|---|
| 48 | |
|---|
| 49 | def __len__(self): |
|---|
| 50 | return len(self.rep) |
|---|
| 51 | |
|---|
| 52 | def __getitem__(self, index): |
|---|
| 53 | return self.rep[index] |
|---|
| 54 | |
|---|
| 55 | def __getattr__(self, name): |
|---|
| 56 | name_indices = {'x':0, 'y':1, 'z':2, 'w': 3, 'u': 0, 'v': 1} |
|---|
| 57 | if not name_indices.has_key(name): raise AttributeError(name) |
|---|
| 58 | return self.rep[name_indices[name]] |
|---|
| 59 | |
|---|
| 60 | # other useful basic functions |
|---|
| 61 | |
|---|
| 62 | def __eq__(self, other): |
|---|
| 63 | return 0 not in map(lambda a, b: a == b, self.rep, other.rep) |
|---|
| 64 | |
|---|
| 65 | def __ne__(self, other): |
|---|
| 66 | return 0 in map(lambda a, b: a == b, self.rep, other.rep) |
|---|
| 67 | |
|---|
| 68 | def __nonzero__(self): |
|---|
| 69 | return max(map(lambda x: abs(x), self.rep)) > 0 |
|---|
| 70 | |
|---|
| 71 | # string conversion functions |
|---|
| 72 | |
|---|
| 73 | def __str__(self): |
|---|
| 74 | s = ",".join(map(lambda x: "%.3f" % x, self.rep)) |
|---|
| 75 | return '(' + s + ')' |
|---|
| 76 | |
|---|
| 77 | def __repr__(self): |
|---|
| 78 | s = ",".join(map(lambda x: "%.3f" % x, self.rep)) |
|---|
| 79 | return 'Vector(' + s + ')' |
|---|
| 80 | |
|---|
| 81 | # other functions |
|---|
| 82 | |
|---|
| 83 | def dot(self, other): |
|---|
| 84 | temp = map(lambda x, y: x * y, self.rep, other.rep) |
|---|
| 85 | return reduce(lambda a, b: a + b, temp) |
|---|
| 86 | |
|---|
| 87 | def cross(self, other): |
|---|
| 88 | a1, a2, a3 = self |
|---|
| 89 | b1, b2, b3 = other |
|---|
| 90 | return Vector(a2*b3 - a3*b2, a3*b1 - a1*b3, a1*b2 - a2*b1) |
|---|
| 91 | |
|---|
| 92 | def normalize(self): |
|---|
| 93 | len = abs(self) |
|---|
| 94 | if len < 1e-5: raise ValueError('zero vector') |
|---|
| 95 | self.rep = tuple(map(lambda x: x / len, self.rep)) |
|---|
| 96 | |
|---|
| 97 | def unit(self): |
|---|
| 98 | len = abs(self) |
|---|
| 99 | if len < 1e-5: raise ValueError('zero vector') |
|---|
| 100 | return self / len |
|---|
| 101 | |
|---|
| 102 | # test |
|---|
| 103 | if __name__ == "__main__": |
|---|
| 104 | a = Vector(0,3,4) |
|---|
| 105 | b = Vector(1,2,5) |
|---|
| 106 | |
|---|
| 107 | print a, b |
|---|
| 108 | print a + b |
|---|
| 109 | print abs(a) |
|---|
| 110 | print a + Vector(1,0,0) |
|---|
| 111 | print a ^ b |
|---|
| 112 | x, y, z = a |
|---|
| 113 | print x, y, z |
|---|
| 114 | |
|---|