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 # (I suppose it would be easy to implement a cmath module.)
  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
  66 import types, math
  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 type(obj) == types.TupleType:
  78                 return apply(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         else:
  90                 return obj
  91
  92 def Im(obj):
  93         if IsComplex(obj):
  94                 return obj.im
  95         else:
  96                 return obj
  97
  98 class Complex:
  99
 100         def __init__(self, re=0, im=0):
 101                 if IsComplex(re):
 102                         im = i + Complex(0, re.im)
 103                         re = re.re
 104                 if IsComplex(im):
 105                         re = re - im.im
 106                         im = im.re
 107                 self.__dict__['re'] = re
 108                 self.__dict__['im'] = im
 109
 110         def __setattr__(self, name, value):
 111                         raise TypeError, 'Complex numbers are immutable'
 112
 113         def __hash__(self):
 114                 if not self.im: return hash(self.re)
 115                 mod = sys.maxint + 1L
 116                 return int((hash(self.re) + 2L*hash(self.im) + mod) % (2L*mod) - mod)
 117
 118         def __repr__(self):
 119                 if not self.im:
 120                         return 'Complex(%s)' % `self.re`
 121                 else:
 122                         return 'Complex(%s, %s)' % (`self.re`, `self.im`)
 123
 124         def __str__(self):
 125                 if not self.im:
 126                         return `self.re`
 127                 else:
 128                         return 'Complex(%s, %s)' % (`self.re`, `self.im`)
 129
 130         def __neg__(self):
 131                 return Complex(-self.re, -self.im)
 132
 133         def __pos__(self):
 134                 return self
 135
 136         def __abs__(self):
 137                 # XXX could be done differently to avoid overflow!
 138                 return math.sqrt(self.re*self.re + self.im*self.im)
 139
 140         def __int__(self):
 141                 if self.im:
 142                         raise ValueError, "can't convert Complex with nonzero im to int"
 143                 return int(self.re)
 144
 145         def __long__(self):
 146                 if self.im:
 147                         raise ValueError, "can't convert Complex with nonzero im to long"
 148                 return long(self.re)
 149
 150         def __float__(self):
 151                 if self.im:
 152                         raise ValueError, "can't convert Complex with nonzero im to float"
 153                 return float(self.re)
 154
 155         def __cmp__(self, other):
 156                 other = ToComplex(other)
 157                 return cmp((self.re, self.im), (other.re, other.im))
 158
 159         def __rcmp__(self, other):
 160                 other = ToComplex(other)
 161                 return cmp(other, self)
 162
 163         def __nonzero__(self):
 164                 return not (self.re == self.im == 0)
 165
 166         abs = radius = __abs__
 167
 168         def angle(self, fullcircle = twopi):
 169                 return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi)
 170
 171         phi = angle
 172
 173         def __add__(self, other):
 174                 other = ToComplex(other)
 175                 return Complex(self.re + other.re, self.im + other.im)
 176
 177         __radd__ = __add__
 178
 179         def __sub__(self, other):
 180                 other = ToComplex(other)
 181                 return Complex(self.re - other.re, self.im - other.im)
 182
 183         def __rsub__(self, other):
 184                 other = ToComplex(other)
 185                 return other - self
 186
 187         def __mul__(self, other):
 188                 other = ToComplex(other)
 189                 return Complex(self.re*other.re - self.im*other.im,
 190                                self.re*other.im + self.im*other.re)
 191
 192         __rmul__ = __mul__
 193
 194         def __div__(self, other):
 195                 other = ToComplex(other)
 196                 d = float(other.re*other.re + other.im*other.im)
 197                 if not d: raise ZeroDivisionError, 'Complex division'
 198                 return Complex((self.re*other.re + self.im*other.im) / d,
 199                                (self.im*other.re - self.re*other.im) / d)
 200
 201         def __rdiv__(self, other):
 202                 other = ToComplex(other)
 203                 return other / self
 204
 205         def __pow__(self, n, z=None):
 206                 if z is not None:
 207                         raise TypeError, 'Complex does not support ternary pow()'
 208                 if IsComplex(n):
 209                         if n.im:
 210                           if self.im: raise TypeError, 'Complex to the Complex power'
 211                           else: return exp(math.log(self.re)*n)
 212                         n = n.re
 213                 r = pow(self.abs(), n)
 214                 phi = n*self.angle()
 215                 return Complex(math.cos(phi)*r, math.sin(phi)*r)
 216
 217         def __rpow__(self, base):
 218                 base = ToComplex(base)
 219                 return pow(base, self)
 220
 221 def exp(z):
 222         r = math.exp(z.re)
 223         return Complex(math.cos(z.im)*r,math.sin(z.im)*r)
 224
 225
 226 def checkop(expr, a, b, value, fuzz = 1e-6):
 227         import sys
 228         print '       ', a, 'and', b,
 229         try:
 230                 result = eval(expr)
 231         except:
 232                 result = sys.exc_type
 233         print '->', result
 234         if (type(result) == type('') or type(value) == type('')):
 235                 ok = result == value
 236         else:
 237                 ok = abs(result - value) <= fuzz
 238         if not ok:
 239                 print '!!\t!!\t!! should be', value, 'diff', abs(result - value)
 240
 241
 242 def test():
 243         testsuite = {
 244                 'a+b': [
 245                         (1, 10, 11),
 246                         (1, Complex(0,10), Complex(1,10)),
 247                         (Complex(0,10), 1, Complex(1,10)),
 248                         (Complex(0,10), Complex(1), Complex(1,10)),
 249                         (Complex(1), Complex(0,10), Complex(1,10)),
 250                 ],
 251                 'a-b': [
 252                         (1, 10, -9),
 253                         (1, Complex(0,10), Complex(1,-10)),
 254                         (Complex(0,10), 1, Complex(-1,10)),
 255                         (Complex(0,10), Complex(1), Complex(-1,10)),
 256                         (Complex(1), Complex(0,10), Complex(1,-10)),
 257                 ],
 258                 'a*b': [
 259                         (1, 10, 10),
 260                         (1, Complex(0,10), Complex(0, 10)),
 261                         (Complex(0,10), 1, Complex(0,10)),
 262                         (Complex(0,10), Complex(1), Complex(0,10)),
 263                         (Complex(1), Complex(0,10), Complex(0,10)),
 264                 ],
 265                 'a/b': [
 266                         (1., 10, 0.1),
 267                         (1, Complex(0,10), Complex(0, -0.1)),
 268                         (Complex(0, 10), 1, Complex(0, 10)),
 269                         (Complex(0, 10), Complex(1), Complex(0, 10)),
 270                         (Complex(1), Complex(0,10), Complex(0, -0.1)),
 271                 ],
 272                 'pow(a,b)': [
 273                         (1, 10, 1),
 274                         (1, Complex(0,10), 1),
 275                         (Complex(0,10), 1, Complex(0,10)),
 276                         (Complex(0,10), Complex(1), Complex(0,10)),
 277                         (Complex(1), Complex(0,10), 1),
 278                         (2, Complex(4,0), 16),
 279                 ],
 280                 'cmp(a,b)': [
 281                         (1, 10, -1),
 282                         (1, Complex(0,10), 1),
 283                         (Complex(0,10), 1, -1),
 284                         (Complex(0,10), Complex(1), -1),
 285                         (Complex(1), Complex(0,10), 1),
 286                 ],
 287         }
 288         exprs = testsuite.keys()
 289         exprs.sort()
 290         for expr in exprs:
 291                 print expr + ':'
 292                 t = (expr,)
 293                 for item in testsuite[expr]:
 294                         apply(checkop, t+item)
 295
 296
 297 if __name__ == '__main__':
 298         test()