| 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()
|
|---|
