Last active
December 17, 2015 12:59
-
-
Save gabraganca/5613851 to your computer and use it in GitHub Desktop.
Do a contour plot on a 3D array
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
import scipy.stats as st | |
import numpy as np | |
import matplotlib.pytlot as plt | |
def contour_plot(array3d, **kwargs): | |
""" | |
Do the contour plot of a X, Y, Z vector. | |
Parameters | |
---------- | |
array3d : A 3D array. | |
scale : If equals to 'log', do a logarithmic scale | |
xlabel : Set x-label of plot | |
ylabel : Set y-label of plot | |
title : Set plot title | |
save : Save as a different name. If not set, save as 'contour.pdf' | |
""" | |
print 'Plotting the countour plot' | |
x = st.itemfreq(array3d[:, 0])[:, 0] | |
y = st.itemfreq(array3d[:, 1])[:, 0] | |
X, Y = np.meshgrid(x, y) | |
Z = np.zeros([len(y), len(x)], float) | |
for i in range(len(X)): | |
for j in range(len(X[0])): | |
for k in array3d: | |
if X[i, j] == k[0] and \ | |
np.round(Y[i, j], 2) == np.round(k[1], 2): | |
Z[i, j] = k[2] | |
#Apply scale, if defined | |
if kwargs.has_key('scale') and kwargs['scale'] == 'log': | |
Z = np.log(Z) | |
#Do contour plot | |
plt.figure() | |
CS = plt.contour(X, Y, Z) | |
plt.clabel(CS, inline=1, fontsize=10) | |
plt.axis([15000, 29000, 3.5, 4.5]) | |
if kwargs.has_key('xlabel'): | |
plt.xlabel(kwargs['xlabel']) | |
if kwargs.has_key('ylabel'): | |
plt.ylabel(kwargs['ylabel']) | |
if kwargs.has_key('title'): | |
plt.title(kwargs['title']) | |
if kwargs.has_key('save'): | |
plt.savefig(kwargs['save']) | |
else: | |
plt.savefig('contour.pdf') | |
plt.clf() | |
#============================================================================= |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment