/Demo/classes/Complex.py

http://unladen-swallow.googlecode.com/ · Python · 320 lines · 212 code · 43 blank · 65 comment · 34 complexity · d3bc787e6f30205c58ed3212c2d4b44f MD5 · raw file

  1. # Complex numbers
  2. # ---------------
  3. # [Now that Python has a complex data type built-in, this is not very
  4. # useful, but it's still a nice example class]
  5. # This module represents complex numbers as instances of the class Complex.
  6. # A Complex instance z has two data attribues, z.re (the real part) and z.im
  7. # (the imaginary part). In fact, z.re and z.im can have any value -- all
  8. # arithmetic operators work regardless of the type of z.re and z.im (as long
  9. # as they support numerical operations).
  10. #
  11. # The following functions exist (Complex is actually a class):
  12. # Complex([re [,im]) -> creates a complex number from a real and an imaginary part
  13. # IsComplex(z) -> true iff z is a complex number (== has .re and .im attributes)
  14. # ToComplex(z) -> a complex number equal to z; z itself if IsComplex(z) is true
  15. # if z is a tuple(re, im) it will also be converted
  16. # PolarToComplex([r [,phi [,fullcircle]]]) ->
  17. # the complex number z for which r == z.radius() and phi == z.angle(fullcircle)
  18. # (r and phi default to 0)
  19. # exp(z) -> returns the complex exponential of z. Equivalent to pow(math.e,z).
  20. #
  21. # Complex numbers have the following methods:
  22. # z.abs() -> absolute value of z
  23. # z.radius() == z.abs()
  24. # z.angle([fullcircle]) -> angle from positive X axis; fullcircle gives units
  25. # z.phi([fullcircle]) == z.angle(fullcircle)
  26. #
  27. # These standard functions and unary operators accept complex arguments:
  28. # abs(z)
  29. # -z
  30. # +z
  31. # not z
  32. # repr(z) == `z`
  33. # str(z)
  34. # hash(z) -> a combination of hash(z.re) and hash(z.im) such that if z.im is zero
  35. # the result equals hash(z.re)
  36. # Note that hex(z) and oct(z) are not defined.
  37. #
  38. # These conversions accept complex arguments only if their imaginary part is zero:
  39. # int(z)
  40. # long(z)
  41. # float(z)
  42. #
  43. # The following operators accept two complex numbers, or one complex number
  44. # and one real number (int, long or float):
  45. # z1 + z2
  46. # z1 - z2
  47. # z1 * z2
  48. # z1 / z2
  49. # pow(z1, z2)
  50. # cmp(z1, z2)
  51. # Note that z1 % z2 and divmod(z1, z2) are not defined,
  52. # nor are shift and mask operations.
  53. #
  54. # The standard module math does not support complex numbers.
  55. # The cmath modules should be used instead.
  56. #
  57. # Idea:
  58. # add a class Polar(r, phi) and mixed-mode arithmetic which
  59. # chooses the most appropriate type for the result:
  60. # Complex for +,-,cmp
  61. # Polar for *,/,pow
  62. import math
  63. import sys
  64. twopi = math.pi*2.0
  65. halfpi = math.pi/2.0
  66. def IsComplex(obj):
  67. return hasattr(obj, 're') and hasattr(obj, 'im')
  68. def ToComplex(obj):
  69. if IsComplex(obj):
  70. return obj
  71. elif isinstance(obj, tuple):
  72. return Complex(*obj)
  73. else:
  74. return Complex(obj)
  75. def PolarToComplex(r = 0, phi = 0, fullcircle = twopi):
  76. phi = phi * (twopi / fullcircle)
  77. return Complex(math.cos(phi)*r, math.sin(phi)*r)
  78. def Re(obj):
  79. if IsComplex(obj):
  80. return obj.re
  81. return obj
  82. def Im(obj):
  83. if IsComplex(obj):
  84. return obj.im
  85. return 0
  86. class Complex:
  87. def __init__(self, re=0, im=0):
  88. _re = 0
  89. _im = 0
  90. if IsComplex(re):
  91. _re = re.re
  92. _im = re.im
  93. else:
  94. _re = re
  95. if IsComplex(im):
  96. _re = _re - im.im
  97. _im = _im + im.re
  98. else:
  99. _im = _im + im
  100. # this class is immutable, so setting self.re directly is
  101. # not possible.
  102. self.__dict__['re'] = _re
  103. self.__dict__['im'] = _im
  104. def __setattr__(self, name, value):
  105. raise TypeError, 'Complex numbers are immutable'
  106. def __hash__(self):
  107. if not self.im:
  108. return hash(self.re)
  109. return hash((self.re, self.im))
  110. def __repr__(self):
  111. if not self.im:
  112. return 'Complex(%r)' % (self.re,)
  113. else:
  114. return 'Complex(%r, %r)' % (self.re, self.im)
  115. def __str__(self):
  116. if not self.im:
  117. return repr(self.re)
  118. else:
  119. return 'Complex(%r, %r)' % (self.re, self.im)
  120. def __neg__(self):
  121. return Complex(-self.re, -self.im)
  122. def __pos__(self):
  123. return self
  124. def __abs__(self):
  125. return math.hypot(self.re, self.im)
  126. def __int__(self):
  127. if self.im:
  128. raise ValueError, "can't convert Complex with nonzero im to int"
  129. return int(self.re)
  130. def __long__(self):
  131. if self.im:
  132. raise ValueError, "can't convert Complex with nonzero im to long"
  133. return long(self.re)
  134. def __float__(self):
  135. if self.im:
  136. raise ValueError, "can't convert Complex with nonzero im to float"
  137. return float(self.re)
  138. def __cmp__(self, other):
  139. other = ToComplex(other)
  140. return cmp((self.re, self.im), (other.re, other.im))
  141. def __rcmp__(self, other):
  142. other = ToComplex(other)
  143. return cmp(other, self)
  144. def __nonzero__(self):
  145. return not (self.re == self.im == 0)
  146. abs = radius = __abs__
  147. def angle(self, fullcircle = twopi):
  148. return (fullcircle/twopi) * ((halfpi - math.atan2(self.re, self.im)) % twopi)
  149. phi = angle
  150. def __add__(self, other):
  151. other = ToComplex(other)
  152. return Complex(self.re + other.re, self.im + other.im)
  153. __radd__ = __add__
  154. def __sub__(self, other):
  155. other = ToComplex(other)
  156. return Complex(self.re - other.re, self.im - other.im)
  157. def __rsub__(self, other):
  158. other = ToComplex(other)
  159. return other - self
  160. def __mul__(self, other):
  161. other = ToComplex(other)
  162. return Complex(self.re*other.re - self.im*other.im,
  163. self.re*other.im + self.im*other.re)
  164. __rmul__ = __mul__
  165. def __div__(self, other):
  166. other = ToComplex(other)
  167. d = float(other.re*other.re + other.im*other.im)
  168. if not d: raise ZeroDivisionError, 'Complex division'
  169. return Complex((self.re*other.re + self.im*other.im) / d,
  170. (self.im*other.re - self.re*other.im) / d)
  171. def __rdiv__(self, other):
  172. other = ToComplex(other)
  173. return other / self
  174. def __pow__(self, n, z=None):
  175. if z is not None:
  176. raise TypeError, 'Complex does not support ternary pow()'
  177. if IsComplex(n):
  178. if n.im:
  179. if self.im: raise TypeError, 'Complex to the Complex power'
  180. else: return exp(math.log(self.re)*n)
  181. n = n.re
  182. r = pow(self.abs(), n)
  183. phi = n*self.angle()
  184. return Complex(math.cos(phi)*r, math.sin(phi)*r)
  185. def __rpow__(self, base):
  186. base = ToComplex(base)
  187. return pow(base, self)
  188. def exp(z):
  189. r = math.exp(z.re)
  190. return Complex(math.cos(z.im)*r,math.sin(z.im)*r)
  191. def checkop(expr, a, b, value, fuzz = 1e-6):
  192. print ' ', a, 'and', b,
  193. try:
  194. result = eval(expr)
  195. except:
  196. result = sys.exc_type
  197. print '->', result
  198. if isinstance(result, str) or isinstance(value, str):
  199. ok = (result == value)
  200. else:
  201. ok = abs(result - value) <= fuzz
  202. if not ok:
  203. print '!!\t!!\t!! should be', value, 'diff', abs(result - value)
  204. def test():
  205. print 'test constructors'
  206. constructor_test = (
  207. # "expect" is an array [re,im] "got" the Complex.
  208. ( (0,0), Complex() ),
  209. ( (0,0), Complex() ),
  210. ( (1,0), Complex(1) ),
  211. ( (0,1), Complex(0,1) ),
  212. ( (1,2), Complex(Complex(1,2)) ),
  213. ( (1,3), Complex(Complex(1,2),1) ),
  214. ( (0,0), Complex(0,Complex(0,0)) ),
  215. ( (3,4), Complex(3,Complex(4)) ),
  216. ( (-1,3), Complex(1,Complex(3,2)) ),
  217. ( (-7,6), Complex(Complex(1,2),Complex(4,8)) ) )
  218. cnt = [0,0]
  219. for t in constructor_test:
  220. cnt[0] += 1
  221. if ((t[0][0]!=t[1].re)or(t[0][1]!=t[1].im)):
  222. print " expected", t[0], "got", t[1]
  223. cnt[1] += 1
  224. print " ", cnt[1], "of", cnt[0], "tests failed"
  225. # test operators
  226. testsuite = {
  227. 'a+b': [
  228. (1, 10, 11),
  229. (1, Complex(0,10), Complex(1,10)),
  230. (Complex(0,10), 1, Complex(1,10)),
  231. (Complex(0,10), Complex(1), Complex(1,10)),
  232. (Complex(1), Complex(0,10), Complex(1,10)),
  233. ],
  234. 'a-b': [
  235. (1, 10, -9),
  236. (1, Complex(0,10), Complex(1,-10)),
  237. (Complex(0,10), 1, Complex(-1,10)),
  238. (Complex(0,10), Complex(1), Complex(-1,10)),
  239. (Complex(1), Complex(0,10), Complex(1,-10)),
  240. ],
  241. 'a*b': [
  242. (1, 10, 10),
  243. (1, Complex(0,10), Complex(0, 10)),
  244. (Complex(0,10), 1, Complex(0,10)),
  245. (Complex(0,10), Complex(1), Complex(0,10)),
  246. (Complex(1), Complex(0,10), Complex(0,10)),
  247. ],
  248. 'a/b': [
  249. (1., 10, 0.1),
  250. (1, Complex(0,10), Complex(0, -0.1)),
  251. (Complex(0, 10), 1, Complex(0, 10)),
  252. (Complex(0, 10), Complex(1), Complex(0, 10)),
  253. (Complex(1), Complex(0,10), Complex(0, -0.1)),
  254. ],
  255. 'pow(a,b)': [
  256. (1, 10, 1),
  257. (1, Complex(0,10), 1),
  258. (Complex(0,10), 1, Complex(0,10)),
  259. (Complex(0,10), Complex(1), Complex(0,10)),
  260. (Complex(1), Complex(0,10), 1),
  261. (2, Complex(4,0), 16),
  262. ],
  263. 'cmp(a,b)': [
  264. (1, 10, -1),
  265. (1, Complex(0,10), 1),
  266. (Complex(0,10), 1, -1),
  267. (Complex(0,10), Complex(1), -1),
  268. (Complex(1), Complex(0,10), 1),
  269. ],
  270. }
  271. for expr in sorted(testsuite):
  272. print expr + ':'
  273. t = (expr,)
  274. for item in testsuite[expr]:
  275. checkop(*(t+item))
  276. if __name__ == '__main__':
  277. test()