Lib/Complex.py

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