Demo/classes/Complex.py

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