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