You cannot select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
138 lines
3.2 KiB
Python
138 lines
3.2 KiB
Python
'''
|
|
Created on 20. jan. 2011
|
|
|
|
@author: pab
|
|
'''
|
|
import numpy as np
|
|
from numpy import exp, meshgrid
|
|
__all__ = ['peaks', 'humps', 'magic']
|
|
|
|
|
|
def magic(n):
|
|
'''
|
|
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
|
|
--------
|
|
>>> np.allclose(magic(3),
|
|
... [[8, 1, 6],
|
|
... [3, 5, 7],
|
|
... [4, 9, 2]])
|
|
True
|
|
|
|
>>> np.allclose(magic(4),
|
|
... [[16, 2, 3, 13],
|
|
... [ 5, 11, 10, 8],
|
|
... [ 9, 7, 6, 12],
|
|
... [ 4, 14, 15, 1]])
|
|
True
|
|
|
|
>>> np.allclose(magic(6),
|
|
... [[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]])
|
|
True
|
|
'''
|
|
if (n < 3):
|
|
raise ValueError('n must be greater than 2.')
|
|
|
|
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):
|
|
'''
|
|
Return the "well" known MatLab (R) peaks function
|
|
evaluated in the [-3,3] x,y range
|
|
|
|
Example
|
|
-------
|
|
>>> import matplotlib.pyplot as plt
|
|
>>> x,y,z = peaks()
|
|
|
|
h = plt.contourf(x,y,z)
|
|
|
|
'''
|
|
if x is None:
|
|
x = np.linspace(-3, 3, n)
|
|
if y is None:
|
|
y = np.linspace(-3, 3, n)
|
|
|
|
[x1, y1] = meshgrid(x, y)
|
|
|
|
z = (3 * (1 - x1) ** 2 * exp(-(x1 ** 2) - (y1 + 1) ** 2) -
|
|
10 * (x1 / 5 - x1 ** 3 - y1 ** 5) * exp(-x1 ** 2 - y1 ** 2) -
|
|
1. / 3 * exp(-(x1 + 1) ** 2 - y1 ** 2))
|
|
|
|
return x1, y1, z
|
|
|
|
|
|
def humps(x=None):
|
|
'''
|
|
Computes a function that has three roots, and some humps.
|
|
|
|
Example
|
|
-------
|
|
>>> import matplotlib.pyplot as plt
|
|
>>> x = np.linspace(0,1)
|
|
>>> y = humps(x)
|
|
|
|
h = plt.plot(x,y)
|
|
'''
|
|
if x is None:
|
|
y = np.linspace(0, 1)
|
|
else:
|
|
y = np.asarray(x)
|
|
|
|
return 1.0 / ((y - 0.3) ** 2 + 0.01) + 1.0 / ((y - 0.9) ** 2 + 0.04) + \
|
|
2 * y - 5.2
|
|
|
|
|
|
if __name__ == '__main__':
|
|
from wafo.testing import test_docstrings
|
|
test_docstrings(__file__)
|