complex numbers, vectors, matrices

complex_matrix.py

import random

from math import sin, cos, atan2, acos
from math import pi as π, e

class ComplexRect:
    def __init__(self, j, i=0):
        self.a = j
        self.b = i

    def add(self, rhs):
        return ComplexRect(self.a + rhs.a, self.b + rhs.b)

    def sub(self, rhs):
        return ComplexRect(self.a - rhs.a, self.b - rhs.b)

    def mul(self, rhs):
        real = self.a*rhs.a - self.b*rhs.b
        imaginary = self.a*rhs.b + self.b*rhs.a
        return ComplexRect(real, imaginary)

    def div(self, rhs):
        if rhs.a == 0 and rhs.b == 0:
            raise(ZeroDivisionError('Cannot divide by zero, even in complex-land'))
        rhs_mod_sq = (rhs.a**2+rhs.b**2)
        real = (self.b*rhs.b+self.a*rhs.a)/rhs_mod_sq
        imaginary = (self.b*rhs.a-self.a*rhs.b)/rhs_mod_sq
        return ComplexRect(real, imaginary)

    def pow(self, r):
        return self.to_polar().pow(r).to_rect()

    def mod(self):
        return (self.a**2+self.b**2)**0.5

    def conjugate(self):
        return ComplexRect(self.a, -self.b)

    def exp(self):
        s = e**self.a
        real = s*cos(self.b)
        imaginary = s*sin(self.a)
        return ComplexRect(real, imaginary)

    def to_polar(self):
        r = self.mod()
        θ = atan2(self.b, self.a)
        return ComplexPolar(r, θ)

    def from_polar(r, θ):
        a = r*cos(θ)
        b = r*sin(θ)
        return ComplexRect(a, b)

    def mobius(self, a, b, c, d):
        return self.mul(a).add(b).div(self.mul(c).add(d))

    def __repr__(self):
        if self.a == 0 and self.b == 0:
            return '0'

        out = ''
        if self.a != 0:
            out = str(self.a)
            if self.b > 0:
                out += '+'

        if self.b == 1:
            out += 'i'
        elif self.b == -1:
            out += '-i'
        elif self.b != 0:
            out += str(self.b) + 'i'

        return out

class ComplexPolar:
    def __init__(self, r, θ):
        self.r = r
        self.θ = θ % (2*π)

    def mul(self, rhs):
        r = self.r * rhs.r
        θ = self.θ + rhs.θ
        return ComplexPolar(r, θ)

    def div(self, rhs):
        r = self.r / rhs.r
        θ = self.θ - rhs.θ
        return ComplexPolar(r, θ)

    def mod(self):
        return self.r

    def pow(self, pow):
        r = self.r ** pow
        θ = self.θ * pow
        return ComplexPolar(r, θ)

    def root(self, root):
        pow = 1/root
        r = self.r ** pow
        θ = self.θ * pow
        return ComplexPolar(r, θ)

    def to_rect(self):
        a = self.r*cos(self.θ)
        b = self.r*sin(self.θ)
        return ComplexRect(a, b)

    def to_exp(self):
        return '{0}e^{1}i'.format(self.r, self.θ)

    def from_rect(a, b):
        r = (a**2+b**2)**0.5
        θ = atan2(b, a)
        return ComplexPolar(r, θ)

    def __repr__(self):
        return '({0}, {1})'.format(self.r, self.θ)

class ComplexVector:
    def __init__(self, elements):
        self.dimension = len(elements)
        self.elements = elements

    def norm(self):
        return self.dot(self).pow(0.5)

    def distance(self, rhs):
        return self.norm().sub(rhs.norm())

    def add(self, rhs):
        elements = []
        for i in range(self.dimension):
            elements.append(self.elements[i].add(rhs.elements[i]))
        return ComplexVector(elements)

    def scale(self, scalar):
        elements = [e.mul(scalar) for e in self.elements]
        return ComplexVector(elements)

    def inverse(self):
        return self.scale(ComplexRect(-1, 0))

    def conjugate(self):
        elements = [e.conjugate() for e in self.elements]
        return ComplexVector(elements)

    def dot(self, rhs):
        ct = ComplexRect(0, 0)
        for i in range(self.dimension):
            ct = ct.add(self.elements[i].conjugate().mul(rhs.elements[i]))
        return ct

    def angle_betweeen(self, rhs):
        denom = self.norm().mul(rhs.norm())
        num = self.dot(rhs)
        return acos(num.div(denom).a) # todo: acos on complex?

    def project_onto(self, rhs):
        return rhs.scale(self.dot(rhs))

    def zero():
        return ComplexVector([ComplexRect(0, 0), ComplexRect(0, 0), ComplexRect(0, 0), ComplexRect(0, 0)])

    def __repr__(self):
        return '[{0}]'.format(', '.join([str(e) for e in self.elements]))

class ComplexMatrix:
    def __init__(self, rows):
        self.rows = rows
        self.m = len(self.rows)
        self.n = len(self.rows[0])

    def get(self, j=None, k=None):
        if k is None:
            row = self.rows[j]
            return ComplexVector(row)
        elif j is None:
            col = [self.get(j, k) for j, e in enumerate(self.rows)]
            return ComplexVector(col)

        return self.rows[j][k]

    def random(m, n):
        rows = [[ComplexRect(random.randint(-10, 10), random.randint(-10, 10)) for k in range(n)] for j in range(m)]
        return ComplexMatrix(rows)

    def add(self, rhs):
        rows = [[self.get(j, k).add(rhs.get(j, k)) for k in range(self.n)] for j in range(self.m)]
        return ComplexMatrix(rows)

    def scale(self, scalar):
        rows = [[self.get(j, k).mul(scalar) for k in range(self.n)] for j in range(self.m)]
        return ComplexMatrix(rows)

    def conjugate(self):
        rows = [[self.get(j, k).conjugate() for k in range(self.n)] for j in range(self.m)]
        return ComplexMatrix(rows)

    def transpose(self):
        rows = [[self.get(k, j) for k in range(self.m)] for j in range(self.n)]
        return ComplexMatrix(rows)

    def T(self):
        return self.transpose()

    def dagger(self):
        rows = [[self.get(k, j).conjugate() for k in range(self.m)] for j in range(self.n)]
        return ComplexMatrix(rows)

    def adjoint(self):
        return self.dagger()

    def __repr__(self):
        out = ''
        for j in range(self.m):
            out += '{0}'.format(' '.join([str(e).rjust(7, ' ') for e in self.rows[j]])) + '\n'
        return out