Python module for vector and matrix handling

linalg.py

# linalg.py - vector and matrix handling
#
# Copyright 2011 Jonathan D. Page. All rights reserved.
#
# Redistribution and use in source and binary forms, with or without modification, are
# permitted provided that the following conditions are met:
#
#    1. Redistributions of source code must retain the above copyright notice, this list of
#       conditions and the following disclaimer.
#
#    2. Redistributions in binary form must reproduce the above copyright notice, this list
#       of conditions and the following disclaimer in the documentation and/or other materials
#       provided with the distribution.
#
# THIS SOFTWARE IS PROVIDED BY <COPYRIGHT HOLDER> ''AS IS'' AND ANY EXPRESS OR IMPLIED
# WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND
# FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL <COPYRIGHT HOLDER> OR
# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
# CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON
# ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
# NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF
# ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.


from __future__ import division
from math import *

class vector(object):
    def __init__(self, data): self.data = tuple(data)
    def __repr__(self): return "<" + ", ".join(str(x) for x in self.data) + ">"
    def __add__(self, v): return vector(a + b for a, b in zip(self.data, v.data))
    def __sub__(self, v): return vector(a - b for a, b in zip(self.data, v.data))
    def __mul__(self, n): return vector(a * n for a in self.data)
    __rmul__ = __mul__
    def __div__(self, n): return vector(a / n for a in self.data)
    __truediv__ = __div__
    def magnitude(self): return sqrt(sum(x**2 for x in self.data))
    mag = magnitude
    __abs__ = magnitude
    def unit(self): return self / self.magnitude()
    def pad(self, n): return vector(self.data[x] if x < len(self.data) else 0 for x in range(n))
    def dot(self, v): return sum(a * b for a, b in zip(self.data, v.data))
    def __len__(self): return len(self.data)
    def __getitem__(self, k): return self.data[k]
    def cross(a, b): return vector((a[1]*b[2]-a[2]*b[1], a[2]*b[0]-a[0]*b[2], a[0]*b[1]-a[1]*b[0]))
    def dcross(a, b): return matrix([[basis(0, 3), basis(1, 3), basis(2, 3)], a, b]).det()

def basis(n, order):
    return vector(1 if x == n else 0 for x in range(order))

I = basis(0, 3)
J = basis(1, 3)
K = basis(2, 3)

class matrix(object):
    def __init__(self, data): self.data = tuple(tuple(n) for n in data)
    def __repr__(self): return "\n".join("[" + ", ".join(str(x) for x in n) + "]" for n in self.data)
    def cols(self): return len(self.data[0]) if self.rows() > 0 else 0
    def rows(self): return len(self.data)
    def getcol(self, n): return tuple(r[n] for r in self.data)
    def getcols(self): return [self.getcol(n) for n in range(self.cols())]
    def getrow(self, n): return self.data[n]
    def getrows(self): return list(self.data)
    __getitem__ = getrow
    def __mul__(self, n):
        if type(n) is matrix:
            return self.mmul(n)
        return matrix((n * a for a in b) for b in self.data)
    __rmul__ = __mul__
    def mmul(a, b):
        if a.cols() != b.rows():
            return None
        return matrix((vector(i).dot(vector(j)) for j in b.getcols()) for i in a.getrows())
        
    def __add__(self, m): return matrix((a + b for a, b in zip(c, d)) for c, d in zip(self.data, m.data))
    def __sub__(self, m): return matrix((a - b for a, b in zip(c, d)) for c, d in zip(self.data, m.data))
    def determinant(self):
            if self.cols() != self.rows(): return None
            if self.cols() == 1: return self.data[0][0]
            return reduce(lambda a, b: a + b, (self.data[0][n]*self._c(0, n) for n in range(len(self.data))))
    def _c(self, i, j): return (-1)**(i+j) * self._m(i, j).det()
    def _m(self, i, j):
        data = list(list(n) for n in self.data)
        del data[i]
        for l in data:
            del l[j]
        return matrix(data)
    det = determinant

def deg(rad): return rad * 180 / pi
def rad(deg): return deg * pi / 180

def vector2a(magnitude, rad): return vector((magnitude*cos(rad), magnitude*sin(rad)))
def vector2(x, y): return vector((x, y))
def vector3(x, y, z): return vector((x, y, z))