def plot_wireframe(self, X, Y, Z, *args, **kwargs):
rstride = cbook.popd(kwargs, "rstride", 1)
cstride = cbook.popd(kwargs, "cstride", 1)
had_data = self.has_data()
rows,cols = Z.shape
tX,tY,tZ = nx.transpose(X), nx.transpose(Y), nx.transpose(Z)
rii = [i for i in range(0,rows,rstride)]+[rows-1]
cii = [i for i in range(0,cols,cstride)]+[cols-1]
xlines = [X[i] for i in rii]
ylines = [Y[i] for i in rii]
zlines = [Z[i] for i in rii]
#
txlines = [tX[i] for i in cii]
tylines = [tY[i] for i in cii]
tzlines = [tZ[i] for i in cii]
#
lines = [zip(xl,yl,zl) for xl,yl,zl in zip(xlines,ylines,zlines)]
lines += [zip(xl,yl,zl) for xl,yl,zl in zip(txlines,tylines,tzlines)]
linec = self.add_lines(lines, *args, **kwargs)
self.auto_scale_xyz(X,Y,Z, had_data)
return linec
def center_matrix(M, dim=0):
"""
Return the matrix M with each row having zero mean and unit std
if dim=1, center columns rather than rows
"""
# todo: implement this w/o loop. Allow optional arg to specify
# dimension to remove the mean from
if dim == 1: M = transpose(M)
M = array(M, Float)
if len(M.shape) == 1 or M.shape[0] == 1 or M.shape[1] == 1:
M = M - mean(M)
sigma = std(M)
if sigma > 0:
M = divide(M, sigma)
if dim == 1: M = transpose(M)
return M
for i in range(M.shape[0]):
M[i] -= mean(M[i])
sigma = std(M[i])
if sigma > 0:
M[i] = divide(M[i], sigma)
if dim == 1: M = transpose(M)
return M
def center_matrix(M, dim=0):
"""
Return the matrix M with each row having zero mean and unit std
if dim=1, center columns rather than rows
"""
# todo: implement this w/o loop. Allow optional arg to specify
# dimension to remove the mean from
if dim==1: M = transpose(M)
M = array(M, Float)
if len(M.shape)==1 or M.shape[0]==1 or M.shape[1]==1:
M = M-mean(M)
sigma = std(M)
if sigma>0:
M = divide(M, sigma)
if dim==1: M=transpose(M)
return M
for i in range(M.shape[0]):
M[i] -= mean(M[i])
sigma = std(M[i])
if sigma>0:
M[i] = divide(M[i], sigma)
if dim==1: M=transpose(M)
return M
def line2d_seg_dist(p0, p1, p):
"""distance(s) from line defined by p0 - p1 to point(s) p
p[0] = x(s)
p[1] = y(s)
intersection point p = p0 + u*(p1-p0)
and intersection point lies within segement if u is between 0 and 1
"""
p, p0, p1 = map(nx.asarray, (p[:2], p0[:2], p1[:2]))
x, y = p
x0, y0 = p0
x1, y1 = p1
#
u = ((x - x0) * (x1 - x0) + (y - y0) *
(y1 - y0)) / (dist2d(p0, p1) * dist2d(p0, p1))
#
def dist(p, p0, p1, u):
if u > 0 and u < 1:
pi = p0 + u * (p1 - p0)
return dist2d(p, pi)
else:
return nx.minimum(dist2d(p, p0), dist2d(p, p1))
if not iterable(u):
return dist(p, p0, p1, u)
else:
# for each point calculate dist
pt = nx.transpose(p)
return nx.array([dist(pe, p0, p1, ue) for pe, ue in zip(pt, u)])
def plot_surface(self, X, Y, Z, *args, **kwargs):
had_data = self.has_data()
rows, cols = Z.shape
tX,tY,tZ = nx.transpose(X), nx.transpose(Y), nx.transpose(Z)
rstride = cbook.popd(kwargs, 'rstride', 10)
cstride = cbook.popd(kwargs, 'cstride', 10)
#
polys = []
boxes = []
for rs in nx.arange(0,rows,rstride):
for cs in nx.arange(0,cols,cstride):
ps = []
corners = []
for a,ta in [(X,tX),(Y,tY),(Z,tZ)]:
ztop = a[rs][cs:min(cols-1,cs+cstride)]
zleft = ta[min(cols-1,cs+cstride)][rs:min(rows,rs+rstride+1)]
zbase = a[min(rows-1,rs+rstride)][cs:min(cols,cs+cstride+1):]
zbase = zbase[::-1]
zright = ta[cs][rs:min(rows-1,rs+rstride):]
zright = zright[::-1]
corners.append([ztop[0],ztop[-1],zbase[0],zbase[-1]])
z = nx.concatenate((ztop,zleft,zbase,zright))
ps.append(z)
boxes.append(map(nx.array,zip(*corners)))
polys.append(zip(*ps))
#
lines = []
shade = []
for box in boxes:
n = proj3d.cross(box[0]-box[1],
box[0]-box[2])
n = n/proj3d.mod(n)*5
shade.append(nx.dot(n,[-1,-1,0.5]))
lines.append((box[0],n+box[0]))
#
color = nx.array([0,0,1,1])
norm = normalize(min(shade),max(shade))
colors = [color * (0.5+norm(v)*0.5) for v in shade]
for c in colors: c[3] = 1
polyc = art3d.Poly3DCollection(polys, facecolors=colors, *args, **kwargs)
polyc._zsort = 1
self.add_collection(polyc)
#
self.auto_scale_xyz(X,Y,Z, had_data)
return polyc
def norm(x, y=2):
"""
Norm of a matrix or a vector according to Matlab.
The description is taken from Matlab:
For matrices...
NORM(X) is the largest singular value of X, max(svd(X)).
NORM(X,2) is the same as NORM(X).
NORM(X,1) is the 1-norm of X, the largest column sum,
= max(sum(abs((X)))).
NORM(X,inf) is the infinity norm of X, the largest row sum,
= max(sum(abs((X')))).
NORM(X,'fro') is the Frobenius norm, sqrt(sum(diag(X'*X))).
NORM(X,P) is available for matrix X only if P is 1, 2, inf or 'fro'.
For vectors...
NORM(V,P) = sum(abs(V).^P)^(1/P).
NORM(V) = norm(V,2).
NORM(V,inf) = max(abs(V)).
NORM(V,-inf) = min(abs(V)).
"""
x = asarray(x)
if numerix.mlab.rank(x) == 2:
if y == 2:
return numerix.mlab.max(numerix.mlab.svd(x)[1])
elif y == 1:
return numerix.mlab.max(asum(absolute((x))))
elif y == 'inf':
return numerix.mlab.max(asum(absolute((transpose(x)))))
elif y == 'fro':
return numerix.mlab.sqrt(
asum(numerix.mlab.diag(matrixmultiply(transpose(x), x))))
else:
verbose.report_error('Second argument not permitted for matrices')
return None
else:
if y == 'inf':
return numerix.mlab.max(absolute(x))
elif y == '-inf':
return numerix.mlab.min(absolute(x))
else:
return power(asum(power(absolute(x), y)), 1 / float(y))
def norm(x,y=2):
"""
Norm of a matrix or a vector according to Matlab.
The description is taken from Matlab:
For matrices...
NORM(X) is the largest singular value of X, max(svd(X)).
NORM(X,2) is the same as NORM(X).
NORM(X,1) is the 1-norm of X, the largest column sum,
= max(sum(abs((X)))).
NORM(X,inf) is the infinity norm of X, the largest row sum,
= max(sum(abs((X')))).
NORM(X,'fro') is the Frobenius norm, sqrt(sum(diag(X'*X))).
NORM(X,P) is available for matrix X only if P is 1, 2, inf or 'fro'.
For vectors...
NORM(V,P) = sum(abs(V).^P)^(1/P).
NORM(V) = norm(V,2).
NORM(V,inf) = max(abs(V)).
NORM(V,-inf) = min(abs(V)).
"""
x = asarray(x)
if numerix.mlab.rank(x)==2:
if y==2:
return numerix.mlab.max(numerix.mlab.svd(x)[1])
elif y==1:
return numerix.mlab.max(asum(absolute((x))))
elif y=='inf':
return numerix.mlab.max(asum(absolute((transpose(x)))))
elif y=='fro':
return numerix.mlab.sqrt(asum(numerix.mlab.diag(matrixmultiply(transpose(x),x))))
else:
verbose.report_error('Second argument not permitted for matrices')
return None
else:
if y == 'inf':
return numerix.mlab.max(absolute(x))
elif y == '-inf':
return numerix.mlab.min(absolute(x))
else:
return power(asum(power(absolute(x),y)),1/float(y))
def detrend_linear(x):
"Return x minus best fit line; 'linear' detrending "
# I'm going to regress x on xx=range(len(x)) and return x -
# (b*xx+a). Now that I have polyfit working, I could convert the
# code here, but if it ain't broke, don't fix it!
xx = arange(float(len(x)))
X = transpose(array([xx] + [x]))
C = cov(X)
b = C[0, 1] / C[0, 0]
a = mean(x) - b * mean(xx)
return x - (b * xx + a)
def detrend_linear(x):
"Return x minus best fit line; 'linear' detrending "
# I'm going to regress x on xx=range(len(x)) and return x -
# (b*xx+a). Now that I have polyfit working, I could convert the
# code here, but if it ain't broke, don't fix it!
xx = arange(float(len(x)))
X = transpose(array([xx]+[x]))
C = cov(X)
b = C[0,1]/C[0,0]
a = mean(x) - b*mean(xx)
return x-(b*xx+a)
def corrcoef(*args):
"""
corrcoef(X) where X is a matrix returns a matrix of correlation
coefficients for each numrows observations and numcols variables.
corrcoef(x,y) where x and y are vectors returns the matrix or
correlation coefficients for x and y.
Numeric arrays can be real or complex
The correlation matrix is defined from the covariance matrix C as
r(i,j) = C[i,j] / sqrt(C[i,i]*C[j,j])
"""
if len(args)==2:
X = transpose(array([args[0]]+[args[1]]))
elif len(args)==1:
X = args[0]
else:
raise RuntimeError, 'Only expecting 1 or 2 arguments'
C = cov(X)
if len(args)==2:
d = resize(diagonal(C), (2,1))
denom = numerix.mlab.sqrt(matrixmultiply(d,transpose(d)))
else:
dc = diagonal(C)
N = len(dc)
shape = N,N
vi = resize(dc, shape)
denom = numerix.mlab.sqrt(vi*transpose(vi)) # element wise multiplication
r = divide(C,denom)
try: return r.real
except AttributeError: return r
def corrcoef(*args):
"""
corrcoef(X) where X is a matrix returns a matrix of correlation
coefficients for each numrows observations and numcols variables.
corrcoef(x,y) where x and y are vectors returns the matrix or
correlation coefficients for x and y.
Numeric arrays can be real or complex
The correlation matrix is defined from the covariance matrix C as
r(i,j) = C[i,j] / sqrt(C[i,i]*C[j,j])
"""
if len(args) == 2:
X = transpose(array([args[0]] + [args[1]]))
elif len(args) == 1:
X = args[0]
else:
raise RuntimeError, 'Only expecting 1 or 2 arguments'
C = cov(X)
if len(args) == 2:
d = resize(diagonal(C), (2, 1))
denom = numerix.mlab.sqrt(matrixmultiply(d, transpose(d)))
else:
dc = diagonal(C)
N = len(dc)
shape = N, N
vi = resize(dc, shape)
denom = numerix.mlab.sqrt(vi *
transpose(vi)) # element wise multiplication
r = divide(C, denom)
try:
return r.real
except AttributeError:
return r
def mfuncC(f, x):
"""
mfuncC(f, x) : matrix function with possibly complex eigenvalues.
Note: Numeric defines (v,u) = eig(x) => x*u.T = u.T * Diag(v)
This function is needed by sqrtm and allows further functions.
"""
x = array(x)
(v, u) = numerix.mlab.eig(x)
uT = transpose(u)
V = numerix.mlab.diag(f(v + 0j))
y = matrixmultiply(uT, matrixmultiply(V, linear_algebra.inverse(uT)))
return approx_real(y)
def idwt2(coeffs, wavelet, mode='sym'):
"""
2D Inverse Discrete Wavelet Transform. Reconstruct data from coefficients
arrays.
coeffs - four 2D coefficients arrays arranged as follows:
(approximation,
(horizontal details,
vertical details,
diagonal details)
)
wavelet - wavelet to use (Wavelet object or name string)
mode - signal extension mode, see MODES
"""
if len(coeffs) != 2 or len(coeffs[1]) != 3:
raise ValueError("Invalid coeffs param")
# L -low-pass data, H - high-pass data
LL, (LH, HL, HH) = coeffs
(LL, LH, HL, HH) = (transpose(LL), transpose(LH), transpose(HL), transpose(HH))
for arr in (LL, LH, HL, HH):
if len(arr.shape) != 2:
raise TypeError("All input coefficients arrays must be 2D")
del arr
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
mode = MODES.from_object(mode)
# idwt columns
L = []
append_L = L.append
for rowL, rowH in izip(LL, LH):
append_L(idwt(rowL, rowH, wavelet, mode, 1))
del LL, LH
H = []
append_H = H.append
for rowL, rowH in izip(HL, HH):
append_H(idwt(rowL, rowH, wavelet, mode, 1))
del HL, HH
L = transpose(L)
H = transpose(H)
# idwt rows
data = []
append_data = data.append
for rowL, rowH in izip(L, H):
append_data(idwt(rowL, rowH, wavelet, mode, 1))
return array(data, default_dtype)
def mfuncC(f, x):
"""
mfuncC(f, x) : matrix function with possibly complex eigenvalues.
Note: Numeric defines (v,u) = eig(x) => x*u.T = u.T * Diag(v)
This function is needed by sqrtm and allows further functions.
"""
x = array(x)
(v, u) = numerix.mlab.eig(x)
uT = transpose(u)
V = numerix.mlab.diag(f(v+0j))
y = matrixmultiply(
uT, matrixmultiply(
V, linear_algebra.inverse(uT)))
return approx_real(y)
def polyfit(x,y,N):
"""
Do a best fit polynomial of order N of y to x. Return value is a
vector of polynomial coefficients [pk ... p1 p0]. Eg, for N=2
p2*x0^2 + p1*x0 + p0 = y1
p2*x1^2 + p1*x1 + p0 = y1
p2*x2^2 + p1*x2 + p0 = y2
.....
p2*xk^2 + p1*xk + p0 = yk
Method: if X is a the Vandermonde Matrix computed from x (see
http://mathworld.wolfram.com/VandermondeMatrix.html), then the
polynomial least squares solution is given by the 'p' in
X*p = y
where X is a len(x) x N+1 matrix, p is a N+1 length vector, and y
is a len(x) x 1 vector
This equation can be solved as
p = (XT*X)^-1 * XT * y
where XT is the transpose of X and -1 denotes the inverse.
For more info, see
http://mathworld.wolfram.com/LeastSquaresFittingPolynomial.html,
but note that the k's and n's in the superscripts and subscripts
on that page. The linear algebra is correct, however.
See also polyval
"""
x = asarray(x)+0.
y = asarray(y)+0.
y = reshape(y, (len(y),1))
X = Matrix(vander(x, N+1))
Xt = Matrix(transpose(X))
c = array(linear_algebra.inverse(Xt*X)*Xt*y) # convert back to array
c.shape = (N+1,)
return c
def polyfit(x, y, N):
"""
Do a best fit polynomial of order N of y to x. Return value is a
vector of polynomial coefficients [pk ... p1 p0]. Eg, for N=2
p2*x0^2 + p1*x0 + p0 = y1
p2*x1^2 + p1*x1 + p0 = y1
p2*x2^2 + p1*x2 + p0 = y2
.....
p2*xk^2 + p1*xk + p0 = yk
Method: if X is a the Vandermonde Matrix computed from x (see
http://mathworld.wolfram.com/VandermondeMatrix.html), then the
polynomial least squares solution is given by the 'p' in
X*p = y
where X is a len(x) x N+1 matrix, p is a N+1 length vector, and y
is a len(x) x 1 vector
This equation can be solved as
p = (XT*X)^-1 * XT * y
where XT is the transpose of X and -1 denotes the inverse.
For more info, see
http://mathworld.wolfram.com/LeastSquaresFittingPolynomial.html,
but note that the k's and n's in the superscripts and subscripts
on that page. The linear algebra is correct, however.
See also polyval
"""
x = asarray(x) + 0.
y = asarray(y) + 0.
y = reshape(y, (len(y), 1))
X = Matrix(vander(x, N + 1))
Xt = Matrix(transpose(X))
c = array(linear_algebra.inverse(Xt * X) * Xt * y) # convert back to array
c.shape = (N + 1, )
return c
def prepca(P, frac=0):
"""
Compute the principal components of P. P is a numVars x
numObservations numeric array. frac is the minimum fraction of
variance that a component must contain to be included
Return value are
Pcomponents : a num components x num observations numeric array
Trans : the weights matrix, ie, Pcomponents = Trans*P
fracVar : the fraction of the variance accounted for by each
component returned
"""
U, s, v = svd(P)
varEach = s**2 / P.shape[1]
totVar = asum(varEach)
fracVar = divide(varEach, totVar)
ind = int(asum(fracVar >= frac))
# select the components that are greater
Trans = transpose(U[:, :ind])
# The transformed data
Pcomponents = matrixmultiply(Trans, P)
return Pcomponents, Trans, fracVar[:ind]
def prepca(P, frac=0):
"""
Compute the principal components of P. P is a numVars x
numObservations numeric array. frac is the minimum fraction of
variance that a component must contain to be included
Return value are
Pcomponents : a num components x num observations numeric array
Trans : the weights matrix, ie, Pcomponents = Trans*P
fracVar : the fraction of the variance accounted for by each
component returned
"""
U,s,v = svd(P)
varEach = s**2/P.shape[1]
totVar = asum(varEach)
fracVar = divide(varEach,totVar)
ind = int(asum(fracVar>=frac))
# select the components that are greater
Trans = transpose(U[:,:ind])
# The transformed data
Pcomponents = matrixmultiply(Trans,P)
return Pcomponents, Trans, fracVar[:ind]
def colorbar_classic(self,
mappable,
cax=None,
orientation='vertical',
tickfmt='%1.1f',
cspacing='proportional',
clabels=None,
drawedges=False,
edgewidth=0.5,
edgecolor='k'):
"""
Create a colorbar for mappable image
mappable is the cm.ScalarMappable instance that you want the
colorbar to apply to, e.g. an Image as returned by imshow or a
PatchCollection as returned by scatter or pcolor.
tickfmt is a format string to format the colorbar ticks
cax is a colorbar axes instance in which the colorbar will be
placed. If None, as default axesd will be created resizing the
current aqxes to make room for it. If not None, the supplied axes
will be used and the other axes positions will be unchanged.
orientation is the colorbar orientation: one of 'vertical' | 'horizontal'
cspacing controls how colors are distributed on the colorbar.
if cspacing == 'linear', each color occupies an equal area
on the colorbar, regardless of the contour spacing.
if cspacing == 'proportional' (Default), the area each color
occupies on the the colorbar is proportional to the contour interval.
Only relevant for a Contour image.
clabels can be a sequence containing the
contour levels to be labelled on the colorbar, or None (Default).
If clabels is None, labels for all contour intervals are
displayed. Only relevant for a Contour image.
if drawedges == True, lines are drawn at the edges between
each color on the colorbar. Default False.
edgecolor is the line color delimiting the edges of the colors
on the colorbar (if drawedges == True). Default black ('k')
edgewidth is the width of the lines delimiting the edges of
the colors on the colorbar (if drawedges == True). Default 0.5
return value is the colorbar axes instance
"""
if orientation not in ('horizontal', 'vertical'):
raise ValueError('Orientation must be horizontal or vertical')
if isinstance(mappable, FigureImage) and cax is None:
raise TypeError(
'Colorbars for figure images currently not supported unless you provide a colorbar axes in cax'
)
ax = self.gca()
cmap = mappable.cmap
if cax is None:
l, b, w, h = ax.get_position()
if orientation == 'vertical':
neww = 0.8 * w
ax.set_position((l, b, neww, h), 'both')
cax = self.add_axes([l + 0.9 * w, b, 0.1 * w, h])
else:
newh = 0.8 * h
ax.set_position((l, b + 0.2 * h, w, newh), 'both')
cax = self.add_axes([l, b, w, 0.1 * h])
else:
if not isinstance(cax, Axes):
raise TypeError('Expected an Axes instance for cax')
norm = mappable.norm
if norm.vmin is None or norm.vmax is None:
mappable.autoscale()
cmin = norm.vmin
cmax = norm.vmax
if isinstance(mappable, ContourSet):
# mappable image is from contour or contourf
clevs = mappable.levels
clevs = minimum(clevs, cmax)
clevs = maximum(clevs, cmin)
isContourSet = True
elif isinstance(mappable, ScalarMappable):
# from imshow or pcolor.
isContourSet = False
clevs = linspace(cmin, cmax, cmap.N + 1) # boundaries, hence N+1
else:
raise TypeError("don't know how to handle type %s" %
type(mappable))
N = len(clevs)
C = array([clevs, clevs])
if cspacing == 'linear':
X, Y = meshgrid(clevs, [0, 1])
elif cspacing == 'proportional':
X, Y = meshgrid(linspace(cmin, cmax, N), [0, 1])
else:
raise ValueError("cspacing must be 'linear' or 'proportional'")
if orientation == 'vertical':
args = (transpose(Y), transpose(C), transpose(X), clevs)
else:
args = (C, Y, X, clevs)
#If colors were listed in the original mappable, then
# let contour handle them the same way.
colors = getattr(mappable, 'colors', None)
if colors is not None:
kw = {'colors': colors}
else:
kw = {'cmap': cmap, 'norm': norm}
if isContourSet and not mappable.filled:
CS = cax.contour(*args, **kw)
colls = mappable.collections
for ii in range(len(colls)):
CS.collections[ii].set_linewidth(colls[ii].get_linewidth())
else:
kw['antialiased'] = False
CS = cax.contourf(*args, **kw)
if drawedges:
for col in CS.collections:
col.set_edgecolor(edgecolor)
col.set_linewidth(edgewidth)
mappable.add_observer(CS)
mappable.set_colorbar(CS, cax)
if isContourSet:
if cspacing == 'linear':
ticks = linspace(cmin, cmax, N)
else:
ticks = clevs
if cmin == mappable.levels[0]:
ticklevs = clevs
else: # We are not showing the full ends of the range.
ticks = ticks[1:-1]
ticklevs = clevs[1:-1]
labs = [tickfmt % lev for lev in ticklevs]
if clabels is not None:
for i, lev in enumerate(ticklevs):
if lev not in clabels:
labs[i] = ''
if orientation == 'vertical':
cax.set_xticks([])
cax.yaxis.tick_right()
cax.yaxis.set_label_position('right')
if isContourSet:
cax.set_yticks(ticks)
cax.set_yticklabels(labs)
else:
cax.yaxis.set_major_formatter(FormatStrFormatter(tickfmt))
else:
cax.set_yticks([])
if isContourSet:
cax.set_xticks(ticks)
cax.set_xticklabels(labs)
else:
cax.xaxis.set_major_formatter(FormatStrFormatter(tickfmt))
self.sca(ax)
return cax
def colorbar_classic(self, mappable, cax=None,
orientation='vertical', tickfmt='%1.1f',
cspacing='proportional',
clabels=None, drawedges=False, edgewidth=0.5,
edgecolor='k'):
"""
Create a colorbar for mappable image
mappable is the cm.ScalarMappable instance that you want the
colorbar to apply to, e.g. an Image as returned by imshow or a
PatchCollection as returned by scatter or pcolor.
tickfmt is a format string to format the colorbar ticks
cax is a colorbar axes instance in which the colorbar will be
placed. If None, as default axesd will be created resizing the
current aqxes to make room for it. If not None, the supplied axes
will be used and the other axes positions will be unchanged.
orientation is the colorbar orientation: one of 'vertical' | 'horizontal'
cspacing controls how colors are distributed on the colorbar.
if cspacing == 'linear', each color occupies an equal area
on the colorbar, regardless of the contour spacing.
if cspacing == 'proportional' (Default), the area each color
occupies on the the colorbar is proportional to the contour interval.
Only relevant for a Contour image.
clabels can be a sequence containing the
contour levels to be labelled on the colorbar, or None (Default).
If clabels is None, labels for all contour intervals are
displayed. Only relevant for a Contour image.
if drawedges == True, lines are drawn at the edges between
each color on the colorbar. Default False.
edgecolor is the line color delimiting the edges of the colors
on the colorbar (if drawedges == True). Default black ('k')
edgewidth is the width of the lines delimiting the edges of
the colors on the colorbar (if drawedges == True). Default 0.5
return value is the colorbar axes instance
"""
if orientation not in ('horizontal', 'vertical'):
raise ValueError('Orientation must be horizontal or vertical')
if isinstance(mappable, FigureImage) and cax is None:
raise TypeError('Colorbars for figure images currently not supported unless you provide a colorbar axes in cax')
ax = self.gca()
cmap = mappable.cmap
if cax is None:
l,b,w,h = ax.get_position()
if orientation=='vertical':
neww = 0.8*w
ax.set_position((l,b,neww,h), 'both')
cax = self.add_axes([l + 0.9*w, b, 0.1*w, h])
else:
newh = 0.8*h
ax.set_position((l,b+0.2*h,w,newh), 'both')
cax = self.add_axes([l, b, w, 0.1*h])
else:
if not isinstance(cax, Axes):
raise TypeError('Expected an Axes instance for cax')
norm = mappable.norm
if norm.vmin is None or norm.vmax is None:
mappable.autoscale()
cmin = norm.vmin
cmax = norm.vmax
if isinstance(mappable, ContourSet):
# mappable image is from contour or contourf
clevs = mappable.levels
clevs = minimum(clevs, cmax)
clevs = maximum(clevs, cmin)
isContourSet = True
elif isinstance(mappable, ScalarMappable):
# from imshow or pcolor.
isContourSet = False
clevs = linspace(cmin, cmax, cmap.N+1) # boundaries, hence N+1
else:
raise TypeError("don't know how to handle type %s"%type(mappable))
N = len(clevs)
C = array([clevs, clevs])
if cspacing == 'linear':
X, Y = meshgrid(clevs, [0, 1])
elif cspacing == 'proportional':
X, Y = meshgrid(linspace(cmin, cmax, N), [0, 1])
else:
raise ValueError("cspacing must be 'linear' or 'proportional'")
if orientation=='vertical':
args = (transpose(Y), transpose(C), transpose(X), clevs)
else:
args = (C, Y, X, clevs)
#If colors were listed in the original mappable, then
# let contour handle them the same way.
colors = getattr(mappable, 'colors', None)
if colors is not None:
kw = {'colors': colors}
else:
kw = {'cmap':cmap, 'norm':norm}
if isContourSet and not mappable.filled:
CS = cax.contour(*args, **kw)
colls = mappable.collections
for ii in range(len(colls)):
CS.collections[ii].set_linewidth(colls[ii].get_linewidth())
else:
kw['antialiased'] = False
CS = cax.contourf(*args, **kw)
if drawedges:
for col in CS.collections:
col.set_edgecolor(edgecolor)
col.set_linewidth(edgewidth)
mappable.add_observer(CS)
mappable.set_colorbar(CS, cax)
if isContourSet:
if cspacing == 'linear':
ticks = linspace(cmin, cmax, N)
else:
ticks = clevs
if cmin == mappable.levels[0]:
ticklevs = clevs
else: # We are not showing the full ends of the range.
ticks = ticks[1:-1]
ticklevs = clevs[1:-1]
labs = [tickfmt % lev for lev in ticklevs]
if clabels is not None:
for i, lev in enumerate(ticklevs):
if lev not in clabels:
labs[i] = ''
if orientation=='vertical':
cax.set_xticks([])
cax.yaxis.tick_right()
cax.yaxis.set_label_position('right')
if isContourSet:
cax.set_yticks(ticks)
cax.set_yticklabels(labs)
else:
cax.yaxis.set_major_formatter(FormatStrFormatter(tickfmt))
else:
cax.set_yticks([])
if isContourSet:
cax.set_xticks(ticks)
cax.set_xticklabels(labs)
else:
cax.xaxis.set_major_formatter(FormatStrFormatter(tickfmt))
self.sca(ax)
return cax
def idwt2(coeffs, wavelet, mode='sym'):
"""
2D Inverse Discrete Wavelet Transform. Reconstruct data from coefficients
arrays.
coeffs - four 2D coefficients arrays arranged as follows (in the same way
as dwt2 output -- see dwt2 description for details):
(approximation,
(horizontal details,
vertical details,
diagonal details)
)
wavelet - wavelet to use (Wavelet object or name string)
mode - signal extension mode, see MODES
"""
if len(coeffs) != 2 or len(coeffs[1]) != 3:
raise ValueError("Invalid coeffs param")
# L -low-pass data, H - high-pass data
LL, (LH, HL, HH) = coeffs
if LL is not None:
LL = transpose(LL)
if LH is not None:
LH = transpose(LH)
if HL is not None:
HL = transpose(HL)
if HH is not None:
HH = transpose(HH)
all_none = True
for arr in (LL, LH, HL, HH):
if arr is not None:
all_none = False
if len(arr.shape) != 2:
raise TypeError("All input coefficients arrays must be 2D.")
del arr
if all_none:
raise ValueError(
"At least one input coefficients array must not be None.")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
mode = MODES.from_object(mode)
# idwt columns
L = []
if LL is None and LH is None:
L = None
else:
if LL is None:
# IDWT can handle None input values - equals to zero-array
LL = cycle([None])
if LH is None:
# IDWT can handle None input values - equals to zero-array
LH = cycle([None])
for rowL, rowH in izip(LL, LH):
L.append(idwt(rowL, rowH, wavelet, mode, 1))
del LL, LH
H = []
if HL is None and HH is None:
H = None
else:
if HL is None:
# IDWT can handle None input values - equals to zero-array
HL = cycle([None])
if HH is None:
# IDWT can handle None input values - equals to zero-array
HH = cycle([None])
for rowL, rowH in izip(HL, HH):
H.append(idwt(rowL, rowH, wavelet, mode, 1))
del HL, HH
if L is not None:
L = transpose(L)
if H is not None:
H = transpose(H)
# idwt rows
data = []
if L is None:
# IDWT can handle None input values - equals to zero-array
L = cycle([None])
if H is None:
# IDWT can handle None input values - equals to zero-array
H = cycle([None])
for rowL, rowH in izip(L, H):
data.append(idwt(rowL, rowH, wavelet, mode, 1))
return array(data, default_dtype)
def dwt2(data, wavelet, mode='sym'):
"""
2D Discrete Wavelet Transform.
data - 2D array with input data
wavelet - wavelet to use (Wavelet object or name string)
mode - signal extension mode, see MODES
Returns approximation and three details 2D coefficients arrays.
The result form four 2D coefficients arrays organized in tuples:
(approximation,
(horizontal details,
vertical details,
diagonal details)
)
which sometimes is also interpreted as laid out in one 2D array
of coefficients, where:
-----------------
| | |
| A(LL) | H(LH) |
| | |
(A, (H, V, D)) <---> -----------------
| | |
| V(HL) | D(HH) |
| | |
-----------------
"""
data = as_float_array(data)
if len(data.shape) != 2:
raise ValueError("Expected 2D data array")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
mode = MODES.from_object(mode)
# filter rows
H, L = [], []
for row in data:
cA, cD = dwt(row, wavelet, mode)
L.append(cA)
H.append(cD)
del data
# filter columns
H = transpose(H)
L = transpose(L)
LL, LH = [], []
for row in L:
cA, cD = dwt(array(row, default_dtype), wavelet, mode)
LL.append(cA)
LH.append(cD)
del L
HL, HH = [], []
for row in H:
cA, cD = dwt(array(row, default_dtype), wavelet, mode)
HL.append(cA)
HH.append(cD)
del H
# build result structure
# (approx., (horizontal, vertical, diagonal))
ret = (transpose(LL), (transpose(LH), transpose(HL), transpose(HH)))
return ret
def swt2(data, wavelet, level, start_level=0):
"""
2D Stationary Wavelet Transform.
data - 2D array with input data
wavelet - wavelet to use (Wavelet object or name string)
level - how many decomposition steps to perform
start_level - the level at which the decomposition will start
Returns list of approximation and details coefficients:
[
(cA_n,
(cH_n, cV_n, cD_n)
),
(cA_n+1,
(cH_n+1, cV_n+1, cD_n+1)
),
...,
(cA_n+level,
(cH_n+level, cV_n+level, cD_n+level)
)
]
where cA is approximation, cH is horizontal details, cV is
vertical details, cD is diagonal details and n is start_level.
"""
data = as_float_array(data)
if len(data.shape) != 2:
raise ValueError("Expected 2D data array")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
ret = []
for i in range(start_level, start_level + level):
# filter rows
H, L = [], []
append_L = L.append
append_H = H.append
for row in data:
cA, cD = swt(row, wavelet, level=1, start_level=i)[0]
append_L(cA)
append_H(cD)
del data
# filter columns
H = transpose(H)
L = transpose(L)
LL, LH = [], []
append_LL = LL.append
append_LH = LH.append
for row in L:
cA, cD = swt(array(row, default_dtype),
wavelet,
level=1,
start_level=i)[0]
append_LL(cA)
append_LH(cD)
del L
HL, HH = [], []
append_HL = HL.append
append_HH = HH.append
for row in H:
cA, cD = swt(array(row, default_dtype),
wavelet,
level=1,
start_level=i)[0]
append_HL(cA)
append_HH(cD)
del H
# build result structure
# (approx., (horizontal, vertical, diagonal))
approx = transpose(LL)
ret.append((approx, (transpose(LH), transpose(HL), transpose(HH))))
data = approx # for next iteration
return ret
def colorbar(self, mappable, tickfmt='%1.1f', cax=None, orientation='vertical'):
"""
Create a colorbar for mappable image
tickfmt is a format string to format the colorbar ticks
cax is a colorbar axes instance in which the colorbar will be
placed. If None, as default axesd will be created resizing the
current aqxes to make room for it. If not None, the supplied axes
will be used and the other axes positions will be unchanged.
orientation is the colorbar orientation: one of 'vertical' | 'horizontal'
return value is the colorbar axes instance
"""
if orientation not in ('horizontal', 'vertical'):
raise ValueError('Orientation must be horizontal or vertical')
if isinstance(mappable, FigureImage) and cax is None:
raise TypeError('Colorbars for figure images currently not supported unless you provide a colorbar axes in cax')
ax = self.gca()
cmap = mappable.cmap
norm = mappable.norm
if norm.vmin is None or norm.vmax is None:
mappable.autoscale()
cmin = norm.vmin
cmax = norm.vmax
if cax is None:
l,b,w,h = ax.get_position()
if orientation=='vertical':
neww = 0.8*w
ax.set_position((l,b,neww,h))
cax = self.add_axes([l + 0.9*w, b, 0.1*w, h])
else:
newh = 0.8*h
ax.set_position((l,b+0.2*h,w,newh))
cax = self.add_axes([l, b, w, 0.1*h])
else:
if not isinstance(cax, Axes):
raise TypeError('Expected an Axes instance for cax')
N = cmap.N
c = linspace(cmin, cmax, N)
C = array([c,c])
if orientation=='vertical':
C = transpose(C)
if orientation=='vertical':
extent=(0, 1, cmin, cmax)
else:
extent=(cmin, cmax, 0, 1)
coll = cax.imshow(C,
interpolation='nearest',
#interpolation='bilinear',
origin='lower',
cmap=cmap, norm=norm,
extent=extent)
mappable.add_observer(coll)
mappable.set_colorbar(coll, cax)
if orientation=='vertical':
cax.set_xticks([])
cax.yaxis.tick_right()
cax.yaxis.set_major_formatter(FormatStrFormatter(tickfmt))
else:
cax.set_yticks([])
cax.xaxis.set_major_formatter(FormatStrFormatter(tickfmt))
self.sca(ax)
return cax
def dwt2(data, wavelet, mode='sym'):
"""
2D Discrete Wavelet Transform.
data - 2D array with input data
wavelet - wavelet to use (Wavelet object or name string)
mode - signal extension mode, see MODES
Returns approximaion and three details 2D coefficients arrays.
The result form four 2D coefficients arrays organized in tuples:
(approximation,
(horizontal details,
vertical details,
diagonal details)
)
which sometimes is also interpreted as layed out in one 2D array
of coefficients, where:
-----------------
| | |
| A(LL) | H(LH) |
| | |
(A, (H, V, D)) <---> -----------------
| | |
| V(HL) | D(HH) |
| | |
-----------------
"""
data = as_float_array(data)
if len(data.shape) != 2:
raise ValueError("Expected 2D data array")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
mode = MODES.from_object(mode)
# filter rows
H, L = [], []
append_L = L.append
append_H = H.append
for row in data:
cA, cD = dwt(row, wavelet, mode)
append_L(cA)
append_H(cD)
del data
# filter columns
H = transpose(H)
L = transpose(L)
LL, LH = [], []
append_LL = LL.append
append_LH = LH.append
for row in L:
cA, cD = dwt(array(row, default_dtype), wavelet, mode)
append_LL(cA)
append_LH(cD)
del L
HL, HH = [], []
append_HL = HL.append
append_HH = HH.append
for row in H:
cA, cD = dwt(array(row, default_dtype), wavelet, mode)
append_HL(cA)
append_HH(cD)
del H
# build result structure
# (approx., (horizontal, vertical, diagonal))
ret = (transpose(LL), (transpose(LH), transpose(HL), transpose(HH)))
return ret
def colorbar(self,
mappable,
tickfmt='%1.1f',
cax=None,
orientation='vertical'):
"""
Create a colorbar for mappable image
tickfmt is a format string to format the colorbar ticks
cax is a colorbar axes instance in which the colorbar will be
placed. If None, as default axesd will be created resizing the
current aqxes to make room for it. If not None, the supplied axes
will be used and the other axes positions will be unchanged.
orientation is the colorbar orientation: one of 'vertical' | 'horizontal'
return value is the colorbar axes instance
"""
if orientation not in ('horizontal', 'vertical'):
raise ValueError('Orientation must be horizontal or vertical')
if isinstance(mappable, FigureImage) and cax is None:
raise TypeError(
'Colorbars for figure images currently not supported unless you provide a colorbar axes in cax'
)
ax = self.gca()
cmap = mappable.cmap
norm = mappable.norm
if norm.vmin is None or norm.vmax is None:
mappable.autoscale()
cmin = norm.vmin
cmax = norm.vmax
if cax is None:
l, b, w, h = ax.get_position()
if orientation == 'vertical':
neww = 0.8 * w
ax.set_position((l, b, neww, h))
cax = self.add_axes([l + 0.9 * w, b, 0.1 * w, h])
else:
newh = 0.8 * h
ax.set_position((l, b + 0.2 * h, w, newh))
cax = self.add_axes([l, b, w, 0.1 * h])
else:
if not isinstance(cax, Axes):
raise TypeError('Expected an Axes instance for cax')
N = cmap.N
c = linspace(cmin, cmax, N)
C = array([c, c])
if orientation == 'vertical':
C = transpose(C)
if orientation == 'vertical':
extent = (0, 1, cmin, cmax)
else:
extent = (cmin, cmax, 0, 1)
coll = cax.imshow(C,
interpolation='nearest',
origin='lower',
cmap=cmap,
norm=norm,
extent=extent)
mappable.add_observer(coll)
mappable.set_colorbar(coll, cax)
if orientation == 'vertical':
cax.set_xticks([])
cax.yaxis.tick_right()
cax.yaxis.set_major_formatter(FormatStrFormatter(tickfmt))
else:
cax.set_yticks([])
cax.xaxis.set_major_formatter(FormatStrFormatter(tickfmt))
self.sca(ax)
return cax
def idwt2(coeffs, wavelet, mode='sym'):
"""
2D Inverse Discrete Wavelet Transform. Reconstruct data from coefficients
arrays.
coeffs - four 2D coefficients arrays arranged as follows (in the same way
as dwt2 output -- see dwt2 description for details):
(approximation,
(horizontal details,
vertical details,
diagonal details)
)
wavelet - wavelet to use (Wavelet object or name string)
mode - signal extension mode, see MODES
"""
if len(coeffs) != 2 or len(coeffs[1]) != 3:
raise ValueError("Invalid coeffs param")
# L -low-pass data, H - high-pass data
LL, (LH, HL, HH) = coeffs
if not LL is None: LL = transpose(LL)
if not LH is None: LH = transpose(LH)
if not HL is None: HL = transpose(HL)
if not HH is None: HH = transpose(HH)
all_none = True
for arr in (LL, LH, HL, HH):
if arr is not None:
all_none = False
if len(arr.shape) != 2:
raise TypeError("All input coefficients arrays must be 2D.")
del arr
if all_none:
raise ValueError(
"At least one input coefficients array must not be None.")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
mode = MODES.from_object(mode)
# idwt columns
L = []
append_L = L.append
if LL is None and LH is None:
L = None
else:
if LL is None:
LL = cycle([
None
]) # IDWT can handle None input values - equals to zero-array
if LH is None:
LH = cycle([
None
]) # IDWT can handle None input values - equals to zero-array
for rowL, rowH in izip(LL, LH):
append_L(idwt(rowL, rowH, wavelet, mode, 1))
del LL, LH
H = []
append_H = H.append
if HL is None and HH is None:
H = None
else:
if HL is None:
HL = cycle([
None
]) # IDWT can handle None input values - equals to zero-array
if HH is None:
HH = cycle([
None
]) # IDWT can handle None input values - equals to zero-array
for rowL, rowH in izip(HL, HH):
append_H(idwt(rowL, rowH, wavelet, mode, 1))
del HL, HH
if L is not None:
L = transpose(L)
if H is not None:
H = transpose(H)
# idwt rows
data = []
append_data = data.append
if L is None:
L = cycle(
[None]) # IDWT can handle None input values - equals to zero-array
if H is None:
H = cycle(
[None]) # IDWT can handle None input values - equals to zero-array
for rowL, rowH in izip(L, H):
append_data(idwt(rowL, rowH, wavelet, mode, 1))
return array(data, default_dtype)
def swt2(data, wavelet, level, start_level=0):
"""
2D Stationary Wavelet Transform.
data - 2D array with input data
wavelet - wavelet to use (Wavelet object or name string)
level - how many decomposition steps to perform
start_level - the level at which the decomposition will start
Returns list of approximation and details coefficients:
[
(cA_n,
(cH_n, cV_n, cD_n)
),
(cA_n+1,
(cH_n+1, cV_n+1, cD_n+1)
),
...,
(cA_n+level,
(cH_n+level, cV_n+level, cD_n+level)
)
]
where cA is approximation, cH is horizontal details, cV is
vertical details, cD is diagonal details and n is start_level.
"""
data = as_float_array(data)
if len(data.shape) != 2:
raise ValueError("Expected 2D data array")
if not isinstance(wavelet, Wavelet):
wavelet = Wavelet(wavelet)
ret = []
for i in range(start_level, start_level + level):
# filter rows
H, L = [], []
for row in data:
cA, cD = swt(row, wavelet, level=1, start_level=i)[0]
L.append(cA)
H.append(cD)
del data
# filter columns
H = transpose(H)
L = transpose(L)
LL, LH = [], []
for row in L:
cA, cD = swt(
array(row, default_dtype), wavelet, level=1, start_level=i
)[0]
LL.append(cA)
LH.append(cD)
del L
HL, HH = [], []
for row in H:
cA, cD = swt(
array(row, default_dtype), wavelet, level=1, start_level=i
)[0]
HL.append(cA)
HH.append(cD)
del H
# build result structure
# (approx., (horizontal, vertical, diagonal))
approx = transpose(LL)
ret.append((approx, (transpose(LH), transpose(HL), transpose(HH))))
# for next iteration
data = approx # noqa
return ret