la.py

   1 #!/usr/local/bin/python
   2
   3 """
   4 A very simple 3D vector, matrix, and linear algebra library.
   5 """
   6
   7 import math
   8 import types
   9
  10
  11 class Vector:
  12
  13     """A 3 element vector."""
  14
  15     ZERO = None
  16     I = None
  17     J = None
  18     K = None
  19
  20     def __init__(self, x, y, z):
  21         self.x, self.y, self.z = float(x), float(y), float(z)
  22
  23     def normSquared(self):
  24         """Return the squared vector norm.  If it is not necessary to
  25         actually calculate the norm, this is more efficient since it
  26         does not involve a square root."""
  27         return self.x**2 + self.y**2 + self.z**2
  28
  29     def norm(self):
  30         """Return the vector norm."""
  31         return math.sqrt(self.normSquared())
  32
  33     def __len__(self): return 3
  34
  35     def __getitem__(self, index):
  36         return (self.x, self.y, self.z)[index]
  37
  38     def dot(self, other):
  39         """Return the dot product, u dot v."""
  40         return self.x*other.x + self.y*other.y + self.z*other.z
  41
  42     def cross(self, other):
  43         """Return the cross product, u cross v."""
  44         return Vector(self.y*other.z - self.z*other.y, \
  45                       self.z*other.x - self.x*other.z, \
  46                       self.x*other.y - self.y*other.x)
  47
  48     def unit(self):
  49         """Return a vector in the same direction but whose length is
  50         unity."""
  51         norm = self.norm()
  52         return Vector(self.x/norm, self.y/norm, self.z/norm)
  53
  54     def bound(self, norm):
  55         """Return a vector in the same direction, but whose length is no
  56         greater than norm."""
  57         if self.normSquared() > norm**2:
  58             return norm*self.unit()
  59         else:
  60             return self
  61
  62     def __neg__(self):
  63         """Return -v."""
  64         return Vector(-self.x, -self.y, -self.z)
  65
  66     def __add__(self, other):
  67         """Return u + v."""
  68         return Vector(self.x + other.x, self.y + other.y, self.z + other.z)
  69
  70     def __sub__(self, other):
  71         """Return u - v."""
  72         return Vector(self.x - other.x, self.y - other.y, self.z - other.z)
  73
  74     def __mul__(self, scalar):
  75         """Return v r."""
  76         return Vector(self.x*scalar, self.y*scalar, self.z*scalar)
  77
  78     def __rmul__(self, scalar):
  79         """Return r v."""
  80         return Vector.__mul__(self, scalar)
  81
  82     def __div__(self, scalar):
  83         """Return v/r."""
  84         return Vector(self.x/scalar, self.y/scalar, self.z/scalar)
  85
  86     def __eq__(self, other):
  87         """Return u == v."""
  88         return self.x == other.x and self.y == other.y and self.z == other.z
  89
  90     def __ne__(self, other):
  91         """Return u != v."""
  92         return self.x != other.x or self.y != other.y or self.z != other.z
  93
  94     def __lt__(self, other): raise TypeError, "no ordering on vectors"
  95     def __gt__(self, other): raise TypeError, "no ordering on vectors"
  96     def __le__(self, other): raise TypeError, "no ordering on vectors"
  97     def __ge__(self, other): raise TypeError, "no ordering on vectors"
  98
  99     def __str__(self):
 100         return '(%s %s %s)' % (self.x, self.y, self.z)
 101
 102     def __repr__(self):
 103         return '<%s @ 0x%x %s>' % (self.__class__, id(self), str(self))
 104
 105 Vector.ZERO = Vector(0.0, 0.0, 0.0)
 106 Vector.I = Vector(1.0, 0.0, 0.0)
 107 Vector.J = Vector(0.0, 1.0, 0.0)
 108 Vector.K = Vector(0.0, 0.0, 1.0)
 109
 110 def PolarVector(rho, theta, phi):
 111     """Return a polar vector (r, theta, phi)."""
 112     r = rho*math.cos(phi)
 113     z = rho*math.sin(phi)
 114     return Vector(r*math.cos(theta), r*math.sin(theta), z)
 115
 116
 117 class Matrix:
 118
 119     """A 3x3 square matrix."""
 120
 121     ZERO = None
 122     IDENTITY = None
 123
 124     def __init__(self, a, b, c, d, e, f, g, h, i):
 125         self.a, self.b, self.c, \
 126                  self.d, self.e, self.f, \
 127                  self.g, self.h, self.i = \
 128                  float(a), float(b), float(c), \
 129                  float(d), float(e), float(f), \
 130                  float(g), float(h), float(i)
 131
 132     def trace(self):
 133         """Return the trace of the matrix, tr M."""
 134         return self.a*self.e*self.i
 135
 136     def determinant(self):
 137         """Return the determinant of the matrix, det M."""
 138         return self.a*self.e*self.i - self.a*self.f*self.h + \
 139                self.b*self.f*self.g - self.b*self.d*self.i + \
 140                self.c*self.d*self.h - self.c*self.e*self.g
 141
 142     def isSingular(self):
 143         """Is the matrix singular?"""
 144         return self.determinant() == 0
 145
 146     def rows(self):
 147         """Return the rows of the matrix as a sequence of vectors."""
 148         return [Vector(self.a, self.b, self.c), \
 149                 Vector(self.d, self.e, self.f), \
 150                 Vector(self.g, self.h, self.i)]
 151
 152     def row(self, index):
 153         """Get the nth row as a vector."""
 154         return self.rows()[index]
 155
 156     def columns(self):
 157         """Return the columns of the matrix as a sequence of vectors."""
 158         return [Vector(self.a, self.b, self.g), \
 159                 Vector(self.b, self.e, self.h), \
 160                 Vector(self.c, self.f, self.i)]
 161
 162     def column(self, index):
 163         """Get the nth column as a vector."""
 164         return self.columns()[index]
 165
 166     def __len__(self): return 3
 167     def __getitem__(self, index): return self.row(index)
 168
 169     def transpose(self):
 170         """Return the transpose of the matrix, M^T."""
 171         return Matrix(self.a, self.d, self.g, \
 172                       self.b, self.e, self.h, \
 173                       self.c, self.f, self.i)
 174
 175     def __add__(self, other):
 176         """Return M + N."""
 177         return Matrix(self.a + other.a, self.b + other.b, self.c + other.c, \
 178                       self.d + other.d, self.e + other.e, self.f + other.f, \
 179                       self.g + other.g, self.h + other.h, self.i + other.i)
 180
 181     def __sub__(self, other):
 182         """Return M - N."""
 183         return Matrix(self.a - other.a, self.b - other.b, self.c - other.c, \
 184                       self.d - other.d, self.e - other.e, self.f - other.f, \
 185                       self.g - other.g, self.h - other.h, self.i - other.i)
 186
 187     def __mul__(self, scalar):
 188         """Return M r."""
 189         return Matrix(self.a*scalar, self.b*scalar, self.c*scalar, \
 190                       self.d*scalar, self.e*scalar, self.f*scalar, \
 191                       self.g*scalar, self.h*scalar, self.i*scalar)
 192
 193     def __rmul__(self, scalar):
 194         """Return r M."""
 195         return Matrix.__mul__(self, scalar)
 196
 197     def __div__(self, scalar):
 198         """Return M/r."""
 199         return Matrix(self.a/scalar, self.b/scalar, self.c/scalar, \
 200                       self.d/scalar, self.e/scalar, self.f/scalar, \
 201                       self.g/scalar, self.h/scalar, self.i/scalar)
 202
 203     def __pow__(self, exponent):
 204         """Return M^n."""
 205         if type(exponent) != types.IntType and \
 206            type(exponent) != types.LongType:
 207             raise TypeError, "exponent must be integral"
 208         if exponent <= 0:
 209             raise ValueError, "exponent must be positive"
 210         result = self
 211         for i in range(exponent - 1):
 212             result = result.concatenate(self)
 213         return result
 214
 215     def concatenate(self, other):
 216         """Return M1 M2."""
 217         return Matrix(self.a*other.a + self.b*other.d + self.c*other.g, \
 218                       self.a*other.b + self.b*other.e + self.c*other.h, \
 219                       self.a*other.c + self.b*other.f + self.c*other.i, \
 220                       self.d*other.a + self.e*other.d + self.f*other.g, \
 221                       self.d*other.b + self.e*other.e + self.f*other.h, \
 222                       self.d*other.c + self.e*other.f + self.f*other.i, \
 223                       self.g*other.a + self.h*other.d + self.i*other.g, \
 224                       self.g*other.b + self.h*other.e + self.i*other.h, \
 225                       self.g*other.c + self.h*other.f + self.i*other.i)
 226
 227     def transform(self, vector):
 228         """Return v M."""
 229         return Vector(self.a*vector.x + self.b*vector.y + self.c*vector.z, \
 230                       self.d*vector.x + self.e*vector.y + self.f*vector.z, \
 231                       self.g*vector.x + self.h*vector.y + self.i*vector.z)
 232
 233     def __cofactorSign(self, i, j):
 234         """Compute (-1)^(i + j) to help in determining cofactors."""
 235         if (i + j + 2) % 2 == 0:
 236             return 1
 237         else:
 238             return -1
 239
 240     def __minorDeterminant(self, minor):
 241         """Find the determinant of a 2x2 matrix represented as [a b c d]."""
 242         a, b, c, d = minor
 243         return a*d - b*c
 244
 245     def __cofactor(self, i, j):
 246         """Return the cofactor for the (i, j)-th element."""
 247         minor = []
 248         for y in range(3):
 249             for x in range(3):
 250                 if x == i or y == j:
 251                     continue
 252                 minor.append(self[y][x])
 253         return self.__cofactorSign(i, j)*self.__minorDeterminant(minor)
 254
 255     def adjoint(self):
 256         """Return adj M."""
 257         coefficients = []
 258         for j in range(3):
 259             for i in range(3):
 260                 coefficients.append(self.__cofactor(j, i)) # built-n transpose
 261         return Matrix(*coefficients)
 262
 263     def inverse(self):
 264         """Return M^-1 where M M^-1 = M^-1 M = I."""
 265         return self.adjoint()/self.determinant()
 266
 267     def __eq__(self, other):
 268         """Return M == N."""
 269         return \
 270             self.a == other.a and self.b == other.b and self.c == other.c and \
 271             self.d == other.d and self.e == other.e and self.f == other.f and \
 272             self.g == other.g and self.h == other.h and self.i == other.i
 273
 274     def __ne__(self, other):
 275         """Return M != N."""
 276         return \
 277             self.a != other.a or self.b != other.b or self.c != other.c or \
 278             self.d != other.d or self.e != other.e or self.f != other.f or \
 279             self.g != other.g or self.h != other.h or self.i != other.i
 280
 281     def __str__(self):
 282         return '[%s %s %s; %s %s %s; %s %s %s]' % \
 283                (self.a, self.b, self.c, \
 284                 self.d, self.e, self.f, \
 285                 self.g, self.h, self.i)
 286
 287     def __repr__(self):
 288         return '<%s @ 0x%x %s>' % (self.__class__, id(self), str(self))
 289
 290 Matrix.ZERO = Matrix(0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0)
 291 Matrix.IDENTITY = Matrix(1.0, 0.0, 0.0, 0.0, 1.0, 0.0, 0.0, 0.0, 1.0)
 292
 293 def ScaleMatrix(sx, sy=None, sz=None):
 294     """Return a scaling matrix."""
 295     if sy is None:
 296         sy = sx
 297     if sz is None:
 298         sz = sy
 299     return Matrix(sx, 0, 0, 0, sy, 0, 0, 0, sz)
 300
 301 ### more matrix creators
 302 ### __hash__ methods for vectors and matrices