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