diff --git a/pywafo/src/wafo/demos.py b/pywafo/src/wafo/demos.py index 46801ea..97b70f7 100644 --- a/pywafo/src/wafo/demos.py +++ b/pywafo/src/wafo/demos.py @@ -10,20 +10,70 @@ __all__ = ['peaks', 'humps', 'magic'] def magic(n): ''' - Return magic square for n of odd orders. + Return magic square for n of any orders > 2. A magic square has the property that the sum of every row and column, as well as both diagonals, is the same number. - + Examples + -------- + >>> magic(3) + array([[8, 1, 6], + [3, 5, 7], + [4, 9, 2]]) + + >>> magic(4) + array([[16, 2, 3, 13], + [ 5, 11, 10, 8], + [ 9, 7, 6, 12], + [ 4, 14, 15, 1]]) + + >>> magic(6) + array([[35, 1, 6, 26, 19, 24], + [ 3, 32, 7, 21, 23, 25], + [31, 9, 2, 22, 27, 20], + [ 8, 28, 33, 17, 10, 15], + [30, 5, 34, 12, 14, 16], + [ 4, 36, 29, 13, 18, 11]]) ''' - if np.mod(n,1)==1: # odd order + if np.mod(n,2)==1: # odd order ix = np.arange(n) + 1 J, I = np.meshgrid(ix, ix) A = np.mod(I + J - (n + 3) / 2, n) B = np.mod(I + 2 * J - 2, n) M = n * A + B + 1 + elif np.mod(n,4)==0: # doubly even order + M = np.arange(1,n*n+1).reshape(n,n) + ix = np.mod(np.arange(n) + 1,4)//2 + J, I = np.meshgrid(ix, ix) + iz = np.flatnonzero(I==J) + M.put(iz, n*n+1-M.flat[iz]) + else: # singly even order + p = n//2 + M0 = magic(p) + M = np.hstack((np.vstack((M0, M0+3*p*p)),np.vstack((M0+2*p*p, M0+p*p)))) + + if n>2: + k = (n-2)//4 + Jvec = np.hstack((np.arange(k), np.arange(n-k+1, n))) + for i in range(p): + for j in Jvec: + temp = M[i][j] + M[i][j]=M[i+p][j] + M[i+p][j] = temp + + i=k + j=0 + temp = M[i][j]; + M[i][j] = M[i+p][j] + M[i+p][j] = temp; + + j=i; + temp=M[i+p][j]; + M[i+p][j]=M[i][j]; + M[i][j]=temp; + return M def peaks(x=None, y=None, n=51): @@ -35,7 +85,8 @@ def peaks(x=None, y=None, n=51): ------- >>> import matplotlib.pyplot as plt >>> x,y,z = peaks() - >>> h = plt.contourf(x,y,z) + + h = plt.contourf(x,y,z) ''' if x is None: @@ -60,7 +111,8 @@ def humps(x=None): >>> import matplotlib.pyplot as plt >>> x = np.linspace(0,1) >>> y = humps(x) - >>> h = plt.plot(x,y) + + h = plt.plot(x,y) ''' if x is None: y = np.linspace(0, 1) @@ -69,5 +121,9 @@ def humps(x=None): return 1.0 / ((y - 0.3) ** 2 + 0.01) + 1.0 / ((y - 0.9) ** 2 + 0.04) + 2 * y - 5.2 +def test_docstrings(): + import doctest + doctest.testmod() + if __name__ == '__main__': - pass + test_docstrings()