# Complex numbers | |
# --------------- | |
# [Now that Python has a complex data type built-in, this is not very | |
# useful, but it's still a nice example class] | |
# This module represents complex numbers as instances of the class Complex. | |
# A Complex instance z has two data attribues, z.re (the real part) and z.im | |
# (the imaginary part). In fact, z.re and z.im can have any value -- all | |
# arithmetic operators work regardless of the type of z.re and z.im (as long | |
# as they support numerical operations). | |
# | |
# The following functions exist (Complex is actually a class): | |
# Complex([re [,im]) -> creates a complex number from a real and an imaginary part | |
# IsComplex(z) -> true iff z is a complex number (== has .re and .im attributes) | |
# ToComplex(z) -> a complex number equal to z; z itself if IsComplex(z) is true | |
# if z is a tuple(re, im) it will also be converted | |
# PolarToComplex([r [,phi [,fullcircle]]]) -> | |
# the complex number z for which r == z.radius() and phi == z.angle(fullcircle) | |
# (r and phi default to 0) | |
# exp(z) -> returns the complex exponential of z. Equivalent to pow(math.e,z). | |
# | |
# Complex numbers have the following methods: | |
# z.abs() -> absolute value of z | |
# z.radius() == z.abs() | |
# z.angle([fullcircle]) -> angle from positive X axis; fullcircle gives units | |
# z.phi([fullcircle]) == z.angle(fullcircle) | |
# | |
# These standard functions and unary operators accept complex arguments: | |
# abs(z) | |
# -z | |
# +z | |
# not z | |
# repr(z) == `z` | |
# str(z) | |
# hash(z) -> a combination of hash(z.re) and hash(z.im) such that if z.im is zero | |
# the result equals hash(z.re) | |
# Note that hex(z) and oct(z) are not defined. | |
# | |
# These conversions accept complex arguments only if their imaginary part is zero: | |
# int(z) | |
# long(z) | |
# float(z) | |
# | |
# The following operators accept two complex numbers, or one complex number | |
# and one real number (int, long or float): | |
# z1 + z2 | |
# z1 - z2 | |
# z1 * z2 | |
# z1 / z2 | |
# pow(z1, z2) | |
# cmp(z1, z2) | |
# Note that z1 % z2 and divmod(z1, z2) are not defined, | |
# nor are shift and mask operations. | |
# | |
# The standard module math does not support complex numbers. | |
# The cmath modules should be used instead. | |
# | |
# Idea: | |
# add a class Polar(r, phi) and mixed-mode arithmetic which | |
# chooses the most appropriate type for the result: | |
# Complex for +,-,cmp | |
# Polar for *,/,pow | |
import math | |
import sys | |
twopi = math.pi*2.0 | |
halfpi = math.pi/2.0 | |
def IsComplex(obj): | |
return hasattr(obj, 're') and hasattr(obj, 'im') | |
def ToComplex(obj): | |
if IsComplex(obj): | |
return obj | |
elif isinstance(obj, tuple): | |
return Complex(*obj) | |
else: | |
return Complex(obj) | |
def PolarToComplex(r = 0, phi = 0, fullcircle = twopi): | |
phi = phi * (twopi / fullcircle) | |
return Complex(math.cos(phi)*r, math.sin(phi)*r) | |
def Re(obj): | |
if IsComplex(obj): | |
return obj.re | |
return obj | |
def Im(obj): | |
if IsComplex(obj): | |
return obj.im | |
return 0 | |
class Complex: | |
def __init__(self, re=0, im=0): | |
_re = 0 | |
_im = 0 | |
if IsComplex(re): | |
_re = re.re | |
_im = re.im | |
else: | |
_re = re | |
if IsComplex(im): | |
_re = _re - im.im | |
_im = _im + im.re | |
else: | |
_im = _im + im | |
# this class is immutable, so setting self.re directly is | |
# not possible. | |
self.__dict__['re'] = _re | |
self.__dict__['im'] = _im | |
def __setattr__(self, name, value): | |
raise TypeError, 'Complex numbers are immutable' | |
def __hash__(self): | |
if not self.im: | |
return hash(self.re) | |
return hash((self.re, self.im)) | |
def __repr__(self): | |
if not self.im: | |
return 'Complex(%r)' % (self.re,) | |
else: | |
return 'Complex(%r, %r)' % (self.re, self.im) | |
def __str__(self): | |
if not self.im: | |
return repr(self.re) | |
else: | |
return 'Complex(%r, %r)' % (self.re, self.im) | |
def __neg__(self): | |
return Complex(-self.re, -self.im) | |
def __pos__(self): | |
return self | |
def __abs__(self): | |
return math.hypot(self.re, self.im) | |
def __int__(self): | |
if self.im: | |
raise ValueError, "can't convert Complex with nonzero im to int" | |
return int(self.re) | |
def __long__(self): | |
if self.im: | |
raise ValueError, "can't convert Complex with nonzero im to long" | |
return long(self.re) | |
def __float__(self): | |
if self.im: | |
raise ValueError, "can't convert Complex with nonzero im to float" | |
return float(self.re) | |
def __cmp__(self, other): | |
other = ToComplex(other) | |
return cmp((self.re, self.im), (other.re, other.im)) | |
def __rcmp__(self, other): | |
other = ToComplex(other) | |
return cmp(other, self) | |
def __nonzero__(self): | |
return not (self.re == self.im == 0) | |
abs = radius = __abs__ | |
def angle(self, fullcircle = twopi): | |
return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi) | |
phi = angle | |
def __add__(self, other): | |
other = ToComplex(other) | |
return Complex(self.re + other.re, self.im + other.im) | |
__radd__ = __add__ | |
def __sub__(self, other): | |
other = ToComplex(other) | |
return Complex(self.re - other.re, self.im - other.im) | |
def __rsub__(self, other): | |
other = ToComplex(other) | |
return other - self | |
def __mul__(self, other): | |
other = ToComplex(other) | |
return Complex(self.re*other.re - self.im*other.im, | |
self.re*other.im + self.im*other.re) | |
__rmul__ = __mul__ | |
def __div__(self, other): | |
other = ToComplex(other) | |
d = float(other.re*other.re + other.im*other.im) | |
if not d: raise ZeroDivisionError, 'Complex division' | |
return Complex((self.re*other.re + self.im*other.im) / d, | |
(self.im*other.re - self.re*other.im) / d) | |
def __rdiv__(self, other): | |
other = ToComplex(other) | |
return other / self | |
def __pow__(self, n, z=None): | |
if z is not None: | |
raise TypeError, 'Complex does not support ternary pow()' | |
if IsComplex(n): | |
if n.im: | |
if self.im: raise TypeError, 'Complex to the Complex power' | |
else: return exp(math.log(self.re)*n) | |
n = n.re | |
r = pow(self.abs(), n) | |
phi = n*self.angle() | |
return Complex(math.cos(phi)*r, math.sin(phi)*r) | |
def __rpow__(self, base): | |
base = ToComplex(base) | |
return pow(base, self) | |
def exp(z): | |
r = math.exp(z.re) | |
return Complex(math.cos(z.im)*r,math.sin(z.im)*r) | |
def checkop(expr, a, b, value, fuzz = 1e-6): | |
print ' ', a, 'and', b, | |
try: | |
result = eval(expr) | |
except: | |
result = sys.exc_type | |
print '->', result | |
if isinstance(result, str) or isinstance(value, str): | |
ok = (result == value) | |
else: | |
ok = abs(result - value) <= fuzz | |
if not ok: | |
print '!!\t!!\t!! should be', value, 'diff', abs(result - value) | |
def test(): | |
print 'test constructors' | |
constructor_test = ( | |
# "expect" is an array [re,im] "got" the Complex. | |
( (0,0), Complex() ), | |
( (0,0), Complex() ), | |
( (1,0), Complex(1) ), | |
( (0,1), Complex(0,1) ), | |
( (1,2), Complex(Complex(1,2)) ), | |
( (1,3), Complex(Complex(1,2),1) ), | |
( (0,0), Complex(0,Complex(0,0)) ), | |
( (3,4), Complex(3,Complex(4)) ), | |
( (-1,3), Complex(1,Complex(3,2)) ), | |
( (-7,6), Complex(Complex(1,2),Complex(4,8)) ) ) | |
cnt = [0,0] | |
for t in constructor_test: | |
cnt[0] += 1 | |
if ((t[0][0]!=t[1].re)or(t[0][1]!=t[1].im)): | |
print " expected", t[0], "got", t[1] | |
cnt[1] += 1 | |
print " ", cnt[1], "of", cnt[0], "tests failed" | |
# test operators | |
testsuite = { | |
'a+b': [ | |
(1, 10, 11), | |
(1, Complex(0,10), Complex(1,10)), | |
(Complex(0,10), 1, Complex(1,10)), | |
(Complex(0,10), Complex(1), Complex(1,10)), | |
(Complex(1), Complex(0,10), Complex(1,10)), | |
], | |
'a-b': [ | |
(1, 10, -9), | |
(1, Complex(0,10), Complex(1,-10)), | |
(Complex(0,10), 1, Complex(-1,10)), | |
(Complex(0,10), Complex(1), Complex(-1,10)), | |
(Complex(1), Complex(0,10), Complex(1,-10)), | |
], | |
'a*b': [ | |
(1, 10, 10), | |
(1, Complex(0,10), Complex(0, 10)), | |
(Complex(0,10), 1, Complex(0,10)), | |
(Complex(0,10), Complex(1), Complex(0,10)), | |
(Complex(1), Complex(0,10), Complex(0,10)), | |
], | |
'a/b': [ | |
(1., 10, 0.1), | |
(1, Complex(0,10), Complex(0, -0.1)), | |
(Complex(0, 10), 1, Complex(0, 10)), | |
(Complex(0, 10), Complex(1), Complex(0, 10)), | |
(Complex(1), Complex(0,10), Complex(0, -0.1)), | |
], | |
'pow(a,b)': [ | |
(1, 10, 1), | |
(1, Complex(0,10), 1), | |
(Complex(0,10), 1, Complex(0,10)), | |
(Complex(0,10), Complex(1), Complex(0,10)), | |
(Complex(1), Complex(0,10), 1), | |
(2, Complex(4,0), 16), | |
], | |
'cmp(a,b)': [ | |
(1, 10, -1), | |
(1, Complex(0,10), 1), | |
(Complex(0,10), 1, -1), | |
(Complex(0,10), Complex(1), -1), | |
(Complex(1), Complex(0,10), 1), | |
], | |
} | |
for expr in sorted(testsuite): | |
print expr + ':' | |
t = (expr,) | |
for item in testsuite[expr]: | |
checkop(*(t+item)) | |
if __name__ == '__main__': | |
test() |