/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
 65import math
 66import sys
 67
 68twopi = math.pi*2.0
 69halfpi = math.pi/2.0
 70
 71def IsComplex(obj):
 72    return hasattr(obj, 're') and hasattr(obj, 'im')
 73
 74def 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
 82def PolarToComplex(r = 0, phi = 0, fullcircle = twopi):
 83    phi = phi * (twopi / fullcircle)
 84    return Complex(math.cos(phi)*r, math.sin(phi)*r)
 85
 86def Re(obj):
 87    if IsComplex(obj):
 88        return obj.re
 89    return obj
 90
 91def Im(obj):
 92    if IsComplex(obj):
 93        return obj.im
 94    return 0
 95
 96class 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
226def exp(z):
227    r = math.exp(z.re)
228    return Complex(math.cos(z.im)*r,math.sin(z.im)*r)
229
230
231def 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
245def 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
319if __name__ == '__main__':
320    test()