|
1 | 1 | # Complex numbers |
| 2 | +# --------------- |
2 | 3 |
|
| 4 | +# This module represents complex numbers as instances of the class Complex. |
| 5 | +# A Complex instance z has two data attribues, z.re (the real part) and z.im |
| 6 | +# (the imaginary part). In fact, z.re and z.im can have any value -- all |
| 7 | +# arithmetic operators work regardless of the type of z.re and z.im (as long |
| 8 | +# as they support numerical operations). |
| 9 | +# |
| 10 | +# The following functions exist (Complex is actually a class): |
| 11 | +# Complex([re [,im]) -> creates a complex number from a real and an imaginary part |
| 12 | +# IsComplex(z) -> true iff z is a complex number (== has .re and .im attributes) |
| 13 | +# ToComplex(z) -> a complex number equal to z; z itself if IsComplex(z) is true |
| 14 | +# if z is a tuple(re, im) it will also be converted |
| 15 | +# PolarToComplex([r [,phi [,fullcircle]]]) -> |
| 16 | +# the complex number z for which r == z.radius() and phi == z.angle(fullcircle) |
| 17 | +# (r and phi default to 0) |
| 18 | +# |
| 19 | +# Complex numbers have the following methods: |
| 20 | +# z.abs() -> absolute value of z |
| 21 | +# z.radius() == z.abs() |
| 22 | +# z.angle([fullcircle]) -> angle from positive X axis; fullcircle gives units |
| 23 | +# z.phi([fullcircle]) == z.angle(fullcircle) |
| 24 | +# |
| 25 | +# These standard functions and unary operators accept complex arguments: |
| 26 | +# abs(z) |
| 27 | +# -z |
| 28 | +# +z |
| 29 | +# not z |
| 30 | +# repr(z) == `z` |
| 31 | +# str(z) |
| 32 | +# hash(z) -> a combination of hash(z.re) and hash(z.im) such that if z.im is zero |
| 33 | +# the result equals hash(z.re) |
| 34 | +# Note that hex(z) and oct(z) are not defined. |
| 35 | +# |
| 36 | +# These conversions accept complex arguments only if their imaginary part is zero: |
| 37 | +# int(z) |
| 38 | +# long(z) |
| 39 | +# float(z) |
| 40 | +# |
| 41 | +# The following operators accept two complex numbers, or one complex number |
| 42 | +# and one real number (int, long or float): |
| 43 | +# z1 + z2 |
| 44 | +# z1 - z2 |
| 45 | +# z1 * z2 |
| 46 | +# z1 / z2 |
| 47 | +# pow(z1, z2) |
| 48 | +# cmp(z1, z2) |
| 49 | +# Note that z1 % z2 and divmod(z1, z2) are not defined, |
| 50 | +# nor are shift and mask operations. |
| 51 | +# |
| 52 | +# The standard module math does not support complex numbers. |
| 53 | +# (I suppose it would be easy to implement a cmath module.) |
| 54 | +# |
| 55 | +# Idea: |
| 56 | +# add a class Polar(r, phi) and mixed-mode arithmetic which |
| 57 | +# chooses the most appropriate type for the result: |
| 58 | +# Complex for +,-,cmp |
| 59 | +# Polar for *,/,pow |
3 | 60 |
|
4 | | -from math import sqrt |
5 | 61 |
|
| 62 | +import types, math |
6 | 63 |
|
7 | | -class complex: |
| 64 | +twopi = math.pi*2.0 |
| 65 | +halfpi = math.pi/2.0 |
8 | 66 |
|
9 | | - def __init__(self, re, im): |
10 | | - self.re = float(re) |
11 | | - self.im = float(im) |
| 67 | +def IsComplex(obj): |
| 68 | + return hasattr(obj, 're') and hasattr(obj, 'im') |
12 | 69 |
|
13 | | - def __coerce__(self, other): |
14 | | - if type(other) == type(self): |
15 | | - if other.__class__ == self.__class__: |
16 | | - return self, other |
17 | | - else: |
18 | | - raise TypeError, 'cannot coerce to complex' |
19 | | - else: |
20 | | - # The cast to float() may raise an exception! |
21 | | - return self, complex(float(other), 0.0) |
| 70 | +def ToComplex(obj): |
| 71 | + if IsComplex(obj): |
| 72 | + return obj |
| 73 | + elif type(obj) == types.TupleType: |
| 74 | + return apply(Complex, obj) |
| 75 | + else: |
| 76 | + return Complex(obj) |
| 77 | + |
| 78 | +def PolarToComplex(r = 0, phi = 0, fullcircle = twopi): |
| 79 | + phi = phi * (twopi / fullcircle) |
| 80 | + return Complex(math.cos(phi)*r, math.sin(phi)*r) |
| 81 | + |
| 82 | +def Re(obj): |
| 83 | + if IsComplex(obj): |
| 84 | + return obj.re |
| 85 | + else: |
| 86 | + return obj |
| 87 | + |
| 88 | +def Im(obj): |
| 89 | + if IsComplex(obj): |
| 90 | + return obj.im |
| 91 | + else: |
| 92 | + return obj |
| 93 | + |
| 94 | +class Complex: |
| 95 | + |
| 96 | + def __init__(self, re=0, im=0): |
| 97 | + if IsComplex(re): |
| 98 | + im = i + Complex(0, re.im) |
| 99 | + re = re.re |
| 100 | + if IsComplex(im): |
| 101 | + re = re - im.im |
| 102 | + im = im.re |
| 103 | + self.__dict__['re'] = re |
| 104 | + self.__dict__['im'] = im |
| 105 | + |
| 106 | + def __setattr__(self, name, value): |
| 107 | + raise TypeError, 'Complex numbers are immutable' |
| 108 | + |
| 109 | + def __hash__(self): |
| 110 | + if not self.im: return hash(self.re) |
| 111 | + mod = sys.maxint + 1L |
| 112 | + return int((hash(self.re) + 2L*hash(self.im) + mod) % (2L*mod) - mod) |
22 | 113 |
|
23 | 114 | def __repr__(self): |
24 | | - return 'complex' + `self.re, self.im` |
| 115 | + if not self.im: |
| 116 | + return 'Complex(%s)' % `self.re` |
| 117 | + else: |
| 118 | + return 'Complex(%s, %s)' % (`self.re`, `self.im`) |
25 | 119 |
|
26 | | - def __cmp__(a, b): |
27 | | - a = a.__abs__() |
28 | | - b = b.__abs__() |
29 | | - return (a > b) - (a < b) |
| 120 | + def __str__(self): |
| 121 | + if not self.im: |
| 122 | + return `self.re` |
| 123 | + else: |
| 124 | + return 'Complex(%s, %s)' % (`self.re`, `self.im`) |
| 125 | + |
| 126 | + def __neg__(self): |
| 127 | + return Complex(-self.re, -self.im) |
| 128 | + |
| 129 | + def __pos__(self): |
| 130 | + return self |
| 131 | + |
| 132 | + def __abs__(self): |
| 133 | + # XXX could be done differently to avoid overflow! |
| 134 | + return math.sqrt(self.re*self.re + self.im*self.im) |
| 135 | + |
| 136 | + def __int__(self): |
| 137 | + if self.im: |
| 138 | + raise ValueError, "can't convert Complex with nonzero im to int" |
| 139 | + return int(self.re) |
| 140 | + |
| 141 | + def __long__(self): |
| 142 | + if self.im: |
| 143 | + raise ValueError, "can't convert Complex with nonzero im to long" |
| 144 | + return long(self.re) |
30 | 145 |
|
31 | 146 | def __float__(self): |
32 | 147 | if self.im: |
33 | | - raise ValueError, 'cannot convert complex to float' |
| 148 | + raise ValueError, "can't convert Complex with nonzero im to float" |
34 | 149 | return float(self.re) |
35 | 150 |
|
36 | | - def __long__(self): |
37 | | - return long(float(self)) |
| 151 | + def __cmp__(self, other): |
| 152 | + other = ToComplex(other) |
| 153 | + return cmp((self.re, self.im), (other.re, other.im)) |
38 | 154 |
|
39 | | - def __int__(self): |
40 | | - return int(float(self)) |
| 155 | + def __rcmp__(self, other): |
| 156 | + other = ToComplex(other) |
| 157 | + return cmp(other, self) |
| 158 | + |
| 159 | + def __nonzero__(self): |
| 160 | + return not (self.re == self.im == 0) |
41 | 161 |
|
42 | | - def __abs__(self): |
43 | | - # XXX overflow? |
44 | | - return sqrt(self.re*self.re + self.im*self.im) |
| 162 | + abs = radius = __abs__ |
45 | 163 |
|
46 | | - def __add__(a, b): |
47 | | - return complex(a.re + b.re, a.im + b.im) |
| 164 | + def angle(self, fullcircle = twopi): |
| 165 | + return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi) |
48 | 166 |
|
49 | | - def __sub__(a, b): |
50 | | - return complex(a.re - b.re, a.im - b.im) |
| 167 | + phi = angle |
51 | 168 |
|
52 | | - def __mul__(a, b): |
53 | | - return complex(a.re*b.re - a.im*b.im, a.re*b.im + a.im*b.re) |
| 169 | + def __add__(self, other): |
| 170 | + other = ToComplex(other) |
| 171 | + return Complex(self.re + other.re, self.im + other.im) |
54 | 172 |
|
55 | | - def __div__(a, b): |
56 | | - q = (b.re*b.re + b.im*b.im) |
57 | | - re = (a.re*b.re + a.im*b.im) / q |
58 | | - im = (a.im*b.re - b.im*a.re) / q |
59 | | - return complex(re, im) |
| 173 | + __radd__ = __add__ |
60 | 174 |
|
61 | | - def __neg__(self): |
62 | | - return complex(-self.re, -self.im) |
| 175 | + def __sub__(self, other): |
| 176 | + other = ToComplex(other) |
| 177 | + return Complex(self.re - other.re, self.im - other.im) |
| 178 | + |
| 179 | + def __rsub__(self, other): |
| 180 | + other = ToComplex(other) |
| 181 | + return other - self |
| 182 | + |
| 183 | + def __mul__(self, other): |
| 184 | + other = ToComplex(other) |
| 185 | + return Complex(self.re*other.re - self.im*other.im, |
| 186 | + self.re*other.im + self.im*other.re) |
| 187 | + |
| 188 | + __rmul__ = __mul__ |
| 189 | + |
| 190 | + def __div__(self, other): |
| 191 | + other = ToComplex(other) |
| 192 | + d = float(other.re*other.re + other.im*other.im) |
| 193 | + if not d: raise ZeroDivisionError, 'Complex division' |
| 194 | + return Complex((self.re*other.re + self.im*other.im) / d, |
| 195 | + (self.im*other.re - self.re*other.im) / d) |
| 196 | + |
| 197 | + def __rdiv__(self, other): |
| 198 | + other = ToComplex(other) |
| 199 | + return other / self |
| 200 | + |
| 201 | + def __pow__(self, n, z=None): |
| 202 | + if z is not None: |
| 203 | + raise TypeError, 'Complex does not support ternary pow()' |
| 204 | + if IsComplex(n): |
| 205 | + if n.im: raise TypeError, 'Complex to the Complex power' |
| 206 | + n = n.re |
| 207 | + r = pow(self.abs(), n) |
| 208 | + phi = n*self.angle() |
| 209 | + return Complex(math.cos(phi)*r, math.sin(phi)*r) |
| 210 | + |
| 211 | + def __rpow__(self, base): |
| 212 | + base = ToComplex(base) |
| 213 | + return pow(base, self) |
| 214 | + |
| 215 | + |
| 216 | +def checkop(expr, a, b, value, fuzz = 1e-6): |
| 217 | + import sys |
| 218 | + print ' ', a, 'and', b, |
| 219 | + try: |
| 220 | + result = eval(expr) |
| 221 | + except: |
| 222 | + result = sys.exc_type |
| 223 | + print '->', result |
| 224 | + if (type(result) == type('') or type(value) == type('')): |
| 225 | + ok = result == value |
| 226 | + else: |
| 227 | + ok = abs(result - value) <= fuzz |
| 228 | + if not ok: |
| 229 | + print '!!\t!!\t!! should be', value, 'diff', abs(result - value) |
63 | 230 |
|
64 | 231 |
|
65 | 232 | def test(): |
66 | | - a = complex(2, 0) |
67 | | - b = complex(3, 4) |
68 | | - print a |
69 | | - print b |
70 | | - print a+b |
71 | | - print a-b |
72 | | - print a*b |
73 | | - print a/b |
74 | | - print b+a |
75 | | - print b-a |
76 | | - print b*a |
77 | | - print b/a |
78 | | - i = complex(0, 1) |
79 | | - print i, i*i, i*i*i, i*i*i*i |
80 | | - j = complex(1, 1) |
81 | | - print j, j*j, j*j*j, j*j*j*j |
82 | | - print abs(j), abs(j*j), abs(j*j*j), abs(j*j*j*j) |
83 | | - print i/-i |
84 | | - |
85 | | -test() |
| 233 | + testsuite = { |
| 234 | + 'a+b': [ |
| 235 | + (1, 10, 11), |
| 236 | + (1, Complex(0,10), Complex(1,10)), |
| 237 | + (Complex(0,10), 1, Complex(1,10)), |
| 238 | + (Complex(0,10), Complex(1), Complex(1,10)), |
| 239 | + (Complex(1), Complex(0,10), Complex(1,10)), |
| 240 | + ], |
| 241 | + 'a-b': [ |
| 242 | + (1, 10, -9), |
| 243 | + (1, Complex(0,10), Complex(1,-10)), |
| 244 | + (Complex(0,10), 1, Complex(-1,10)), |
| 245 | + (Complex(0,10), Complex(1), Complex(-1,10)), |
| 246 | + (Complex(1), Complex(0,10), Complex(1,-10)), |
| 247 | + ], |
| 248 | + 'a*b': [ |
| 249 | + (1, 10, 10), |
| 250 | + (1, Complex(0,10), Complex(0, 10)), |
| 251 | + (Complex(0,10), 1, Complex(0,10)), |
| 252 | + (Complex(0,10), Complex(1), Complex(0,10)), |
| 253 | + (Complex(1), Complex(0,10), Complex(0,10)), |
| 254 | + ], |
| 255 | + 'a/b': [ |
| 256 | + (1., 10, 0.1), |
| 257 | + (1, Complex(0,10), Complex(0, -0.1)), |
| 258 | + (Complex(0, 10), 1, Complex(0, 10)), |
| 259 | + (Complex(0, 10), Complex(1), Complex(0, 10)), |
| 260 | + (Complex(1), Complex(0,10), Complex(0, -0.1)), |
| 261 | + ], |
| 262 | + 'pow(a,b)': [ |
| 263 | + (1, 10, 1), |
| 264 | + (1, Complex(0,10), 'TypeError'), |
| 265 | + (Complex(0,10), 1, Complex(0,10)), |
| 266 | + (Complex(0,10), Complex(1), Complex(0,10)), |
| 267 | + (Complex(1), Complex(0,10), 'TypeError'), |
| 268 | + (2, Complex(4,0), 16), |
| 269 | + ], |
| 270 | + 'cmp(a,b)': [ |
| 271 | + (1, 10, -1), |
| 272 | + (1, Complex(0,10), 1), |
| 273 | + (Complex(0,10), 1, -1), |
| 274 | + (Complex(0,10), Complex(1), -1), |
| 275 | + (Complex(1), Complex(0,10), 1), |
| 276 | + ], |
| 277 | + } |
| 278 | + exprs = testsuite.keys() |
| 279 | + exprs.sort() |
| 280 | + for expr in exprs: |
| 281 | + print expr + ':' |
| 282 | + t = (expr,) |
| 283 | + for item in testsuite[expr]: |
| 284 | + apply(checkop, t+item) |
| 285 | + |
| 286 | + |
| 287 | +if __name__ == '__main__': |
| 288 | + test() |
0 commit comments