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