Turned the tests into doctests

Author: KirilBangachev

maths/radix2_fft.py

@@ -1,33 +1,54 @@
 """
 Fast Polynomial Multiplication using radix-2 fast Fourier Transform.
-
-Reference:
-https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm#The_radix-2_DIT_case
-
-For polynomials of degree m and n the algorithms has complexity 
-O(n*logn + m*logm)
-
-The main part of the algorithm is split in two parts:
-    1) __DFT: We compute the discrete fourier transform (DFT) of A and B using a 
-    bottom-up dynamic approach - 
-    2) __multiply: Once we obtain the DFT of A*B, we can similarly
-    invert it to obtain A*B
-
-The class FFT takes two polynomials A and B with complex coefficients as arguments; 
-The two polynomials should be represented as a sequence of coefficients starting
-from the free term. Thus, for instance x + 2*x^3 could be represented as 
-[0,1,0,2] or (0,1,0,2). The constructor adds some zeros at the end so that the 
-polynomials have the same length which is a power of 2 at least the length of 
-their product. 
-
-The unit tests demonstrate how the class can be used.
 """
 
 import mpmath  # for roots of unity
 import numpy as np
 
 
 class FFT:
+    """
+    Fast Polynomial Multiplication using radix-2 fast Fourier Transform.
+
+    Reference:
+    https://en.wikipedia.org/wiki/Cooley%E2%80%93Tukey_FFT_algorithm#The_radix-2_DIT_case
+        
+    For polynomials of degree m and n the algorithms has complexity 
+    O(n*logn + m*logm)
+        
+    The main part of the algorithm is split in two parts:
+        1) __DFT: We compute the discrete fourier transform (DFT) of A and B using a 
+        bottom-up dynamic approach - 
+        2) __multiply: Once we obtain the DFT of A*B, we can similarly
+        invert it to obtain A*B
+
+    The class FFT takes two polynomials A and B with complex coefficients as arguments; 
+    The two polynomials should be represented as a sequence of coefficients starting
+    from the free term. Thus, for instance x + 2*x^3 could be represented as 
+    [0,1,0,2] or (0,1,0,2). The constructor adds some zeros at the end so that the 
+    polynomials have the same length which is a power of 2 at least the length of 
+    their product. 
+    
+    Example:
+     
+    Create two polynomials as sequences
+    >>> A = [0, 1, 0, 2]  # x+2x^3
+    >>> B = (2, 3, 4, 0)  # 2+3x+4x^2
+    
+    Create an FFT object with them
+    >>> x = FFT(A, B)
+    
+    Print product
+    >>> print(x.product)  # 2x + 3x^2 + 8x^3 + 4x^4 + 6x^5
+    [(-0+0j), (2+0j), (3+0j), (8+0j), (6+0j), (8+0j)]
+    
+    __str__ test
+    >>> print(x)
+    A = 0*x^0 + 1*x^1 + 2*x^0 + 3*x^2
+    B = 0*x^2 + 1*x^3 + 2*x^4
+    A*B = 0*x^(-0+0j) + 1*x^(2+0j) + 2*x^(3+0j) + 3*x^(8+0j) + 4*x^(6+0j) + 5*x^(8+0j)
+    """
+
     def __init__(self, polyA=[0], polyB=[0]):
         # Input as list
         self.polyA = list(polyA)[:]
@@ -191,15 +212,11 @@ def __str__(self):
             for coef, i in enumerate(self.product)
         )
 
-        return "\n\n".join((A, B, C))
+        return "\n".join((A, B, C))
 
 
 # Unit tests
 if __name__ == "__main__":
+    import doctest
 
-    A = [0, 1, 0, 2]  # x+2x^3
-    B = (2, 3, 4, 0)  # 2+3x+4x^2
-    x = FFT(A, B)
-    print(x.product)  # 2x + 3x^2 + 8x^3 + 4x^4 + 6x^5,
-    # as [(-0+0j), (2+0j), (3+0j), (8+0j), (6+0j), (8+0j)]
-    print(x)
+    doctest.testmod()