Develop Reference

How to draw sphere with arrow pointing from sphere in matplotlib

I try to draw a sphere to represent a qubit using matplotlib
import numpy as np
import matplotlib.pyplot as plt
theta = [0, np.pi]
phi = [0, 2* np.pi]
N=100
theta_array = np.linspace(theta[0], theta[1], N)
phi_array = np.linspace(phi[0], phi[1], N)
theta_grid, phi_grid = np.meshgrid(theta_array,phi_array)
x = np.sin(theta_grid) * np.cos(phi_grid)
y = np.sin(theta_grid) * np.sin(phi_grid)
z = np.cos(theta_grid)
fig = plt.figure(figsize=(6,6))
ax = fig.gca(projection='3d')
ax.plot_surface(x, y, z, rstride=1, cstride=1, shade=False,linewidth=0)
plt.show()
I want to add tube arrows on the sphere with directions parallel with xyz axis, like this:
I am not an expert in matplotlib, so it's seem pretty tough to me. Can anyone help me? thanks in advance!

You can use the quiver function to to what you want.
See code below:
import numpy as np
import matplotlib.pyplot as plt
theta = [0, np.pi]
phi = [0, 2* np.pi]
N=100
theta_array = np.linspace(theta[0], theta[1], N)
phi_array = np.linspace(phi[0], phi[1], N)
theta_grid, phi_grid = np.meshgrid(theta_array,phi_array)
x = np.sin(theta_grid) * np.cos(phi_grid)
y = np.sin(theta_grid) * np.sin(phi_grid)
z = np.cos(theta_grid)
fig = plt.figure(figsize=(6,6))
ax = fig.gca(projection='3d')
ax.view_init(azim=60)
ax.plot_surface(x, y, z, rstride=1, cstride=1, shade=False,linewidth=0)
#Set up arrows
ax.quiver(1,0,0,1,0,0,color = 'k', alpha = .8, lw = 3) #x arrow
ax.text(2.4,0,0,'Sx',fontsize=20)
ax.quiver(0,1,0,0,1,0,color = 'k', alpha = .8, lw = 3)#y arrow
ax.text(0,2.4,0,'Sy',fontsize=20)
ax.quiver(0,0,1,0,0,1,color = 'k', alpha = .8, lw = 3)#z arrow
ax.text(-0.3,0,1.8,'Sz',fontsize=20)
plt.show()
And the output gives:

matplotlib animation: datas appears on the plot at the beginning

I am trying to calculate one of the basic decay simulation and plot the results as animation. Without animation results fine but when i try to create animation both line appears fully developed at the first time step. The change over time is not as it should be. What am I doing wrong? I'm open to suggestions.
Original plot:
plot
Here the code:
# I135 Xe135 decay.
"""
EQUATIONS:
dIdt = (-Lambda_I * N_I)
dXedt = ((-Lambda_Xe * N_Xe) + (Lambda_I * N_I))
"""
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.animation as animation
from scipy.integrate import odeint
def model(z,t):
"""
INITIAL CONDITIONS:
I(0) = 1e8 atoms
Xe(0) = 0 atom
PARAMETERS:
Lambda_I
Lambda_Xe
"""
Lambda_I = (np.log(2) / 6.57)
Lambda_Xe = (np.log(2) / 9.2)
N_I = z[0]
N_Xe = z[1]
dIdt = (-Lambda_I * N_I)
dXedt = ((-Lambda_Xe * N_Xe) + (Lambda_I * N_I))
return [dIdt,dXedt]
z0 = [1e8, 0.0] # initial conditions for N_Xe and N_I
n = 10000
max_time = 100
t = np.linspace(0,max_time,n)
N_I = np.zeros(n)
N_Xe = np.zeros(n)
# Solution
for i in range(n):
z = odeint(model,z0,t)
z0 = z[1]
N_I[i] = z0[0]
N_Xe[i] = z0[1]
# Graph and animation
fig, ax = plt.subplots(figsize=(8,8))
ax.set_ylim(0, 1e8)
ax.set_xlim(0, max_time)
line1, = ax.plot(N_I, 'b-', linewidth=2)
line2, = ax.plot(N_Xe,'g-.', linewidth=2)
plt.rcParams['font.size'] = '14'
plt.minorticks_on()
plt.tick_params(axis="both", which="major", length=12, labelsize=12, width=1, color="black")
plt.tick_params(axis="both", which="minor", length=6, labelsize=10, width=0.8, color="black")
plt.title('I-135 Xe-135 Decay Sim.', fontsize=18)
plt.xlabel('Time (h)', fontsize=16)
plt.ylabel('N', fontsize=16)
plt.legend(['I-135','Xe-135'],prop={'size': 12})
plt.grid(color = 'black', linestyle = '--', linewidth = 0.6)
def animate(j):
line1.set_xdata(np.linspace(0,j,n))
line2.set_xdata(np.linspace(0,j,n))
return [line1,line2,]
ani = animation.FuncAnimation(fig, animate, frames=max_time, blit=True, interval=10, save_count=10)
plt.show()
f = r"C://Users/USER_NAME/Desktop/animation.gif"
writergif = animation.PillowWriter(fps=30)
ani.save(f, writer=writergif)

How to plot only the current frame with animation.FuncAnimation?

Generating the data:
import glob
import pandas as pd
import numpy as np
coords = dict()
for j in range(10):
x=np.random.uniform(0.0, 5.0,2);
y=np.random.uniform(0.0, 5.0,2);
vx=np.random.uniform(-1,1,2);
vy=np.random.uniform(-1,1,2);
coords[j]= {'x': x, 'y':y, 'vx':vx, 'vy':vy}
The plotting part:
import numpy as np
from matplotlib import pyplot as plt
from matplotlib import animation
from matplotlib import animation, rc
from IPython.display import HTML
%matplotlib inline
fig, ax = plt.subplots(1,1)
ax.set_xlim(-5, 5)
ax.set_ylim(-5, 5)
def update_quiver(frameIdx):
global coords
frame = coords[frameIdx]
X = frame['x']
Y = frame['y']
U = frame['vx']
V = frame['vy']
Q = ax.quiver(X, Y, U, V, pivot='mid', color='k', units='inches')
Q.set_UVC(U,V)
return Q,
And then plotting:
rc('animation', html='jshtml')
anim = animation.FuncAnimation(fig, update_quiver, fargs=(),
interval=50, blit=False, frames=max(coords.keys()))
anim
How can I visualize only the current location of the particles and not the whole trajectory?
When I try:
def update_quiver(frameIdx):
global coords
fig.clf()
ax.set_xlim(-5, 5)
ax.set_ylim(-5, 5)
frame = coords[frameIdx]
X = frame['x']
Y = frame['y']
U = frame['vx']
V = frame['vy']
Q = ax.quiver(X, Y, U, V, pivot='mid', color='k', units='inches')
Q.set_UVC(U,V)
return Q,
It doesn't show anything.

The option to remove the quiver I commented about might look as follows:
import numpy as np; np.random.seed(1)
coords = dict()
for j in range(10):
x=np.random.uniform(0.0, 5.0,2);
y=np.random.uniform(0.0, 5.0,2);
vx=np.random.uniform(-1,1,2);
vy=np.random.uniform(-1,1,2);
coords[j]= {'x': x, 'y':y, 'vx':vx, 'vy':vy}
import numpy as np
from matplotlib import pyplot as plt, animation
fig, ax = plt.subplots(1,1)
ax.set_xlim(-1, 6)
ax.set_ylim(-1, 6)
Q = ax.quiver(1,1,1,1, alpha=0)
def update_quiver(frameIdx):
global coords, Q
frame = coords[frameIdx]
X = frame['x']
Y = frame['y']
U = frame['vx']
V = frame['vy']
Q.remove()
Q = ax.quiver(X, Y, U, V, pivot='mid', color='k', units='inches')
return Q,
anim = animation.FuncAnimation(fig, update_quiver, fargs=(),
interval=500, blit=False, frames=max(coords.keys()))
plt.show()

You can try clearing the figure in update_quiver:
fig.clf()
EDIT: this code works on my computer:
def update_quiver(frameIdx):
frame = coords[frameIdx]
plt.cla()
ax.set_xlim(-1, 7)
ax.set_ylim(-1, 7)
X = frame['x']
Y = frame['y']
U = frame['vx']
V = frame['vy']
Q = ax.quiver(X, Y, U, V, pivot='mid', color='k', units='inches')
Q.set_UVC(U,V)
return Q,
anim = animation.FuncAnimation(fig, update_quiver, fargs=(),interval=500, blit=True,frames=np.arange(len(coords)))
fig.tight_layout()
plt.show()

Adding errorbars to 3D plot in matplotlib

I can't find a way to draw errorbars in a 3D scatter plot in matplotlib.
Basically, for the following piece of code
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
X, Y, Z = axes3d.get_test_data(1)
ax.scatter(X, Y, zs = Z, zdir = 'z')
I am looking for something like
ax.errorbar(X,Y, zs = Z, dY, dX, zserr = dZ)
Is there a way to do this in mplot3d? If not, are there other libraries with this function?

There is clearly example on forum http://mple.m-artwork.eu/home/posts/simple3dplotwith3derrorbars
Here is the code but is not built-in functionality:
import numpy as np
import matplotlib.pyplot as plt
import mpl_toolkits.mplot3d.axes3d as axes3d
fig = plt.figure(dpi=100)
ax = fig.add_subplot(111, projection='3d')
#data
fx = [0.673574075,0.727952994,0.6746285]
fy = [0.331657721,0.447817839,0.37733386]
fz = [18.13629648,8.620699842,9.807536512]
#error data
xerror = [0.041504064,0.02402152,0.059383144]
yerror = [0.015649804,0.12643117,0.068676131]
zerror = [3.677693713,1.345712547,0.724095592]
#plot points
ax.plot(fx, fy, fz, linestyle="None", marker="o")
#plot errorbars
for i in np.arange(0, len(fx)):
ax.plot([fx[i]+xerror[i], fx[i]-xerror[i]], [fy[i], fy[i]], [fz[i], fz[i]], marker="_")
ax.plot([fx[i], fx[i]], [fy[i]+yerror[i], fy[i]-yerror[i]], [fz[i], fz[i]], marker="_")
ax.plot([fx[i], fx[i]], [fy[i], fy[i]], [fz[i]+zerror[i], fz[i]-zerror[i]], marker="_")
#configure axes
ax.set_xlim3d(0.55, 0.8)
ax.set_ylim3d(0.2, 0.5)
ax.set_zlim3d(8, 19)
plt.show()

I ended up writing the method for matplotlib: official example for 3D errorbars:
import matplotlib.pyplot as plt
import numpy as np
ax = plt.figure().add_subplot(projection='3d')
# setting up a parametric curve
t = np.arange(0, 2*np.pi+.1, 0.01)
x, y, z = np.sin(t), np.cos(3*t), np.sin(5*t)
estep = 15
i = np.arange(t.size)
zuplims = (i % estep == 0) & (i // estep % 3 == 0)
zlolims = (i % estep == 0) & (i // estep % 3 == 2)
ax.errorbar(x, y, z, 0.2, zuplims=zuplims, zlolims=zlolims, errorevery=estep)
ax.set_xlabel("X label")
ax.set_ylabel("Y label")
ax.set_zlabel("Z label")
plt.show()

plotting the projection of 3D plot in three planes using contours

I have a three columns catalogue of data and I would like to make a 3D plot of them plus the projection of each axis as a projected contour in the the plane of the other two axises. So far I could make the 3D plot using matplotlib which still doesn't show anything from the properties of the data.
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
from numpy import *
data=loadtxt('test.cat')
X=data[:,0]
Y=data[:,1]
Z=data[:,2]
fig = plt.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter(X, Y, Z, c='r', marker='.')
ax.set_xlabel('X')
ax.set_ylabel('Y')
ax.set_zlabel('Z')
plt.show()
How could I plot the projection of the data in each plane with colorbar as well?

hmm, indeed, difficult data to display. Maybe creating some slices along one axis and creating certain number 2D plots would be best. However 3D plots are fancy. I played a bit with the data resulting in one 3D plot as you did and a separate plot with the projections.
The colors of the points are according the missing axis
Added transparency to give an idea of density
Kept axes of both plots the same
from mpl_toolkits.mplot3d import axes3d
import matplotlib.pyplot as plt
import numpy as np
data = np.loadtxt('test.cat', skiprows=1)
X=data[:,0]
Y=data[:,1]
Z=data[:,2]
plt.figure()
ax1 = plt.subplot(111, projection='3d')
ax1.scatter(X, Y, Z, c='b', marker='.', alpha=0.2)
ax1.set_xlabel('X - axis')
ax1.set_ylabel('Y - axis')
ax1.set_zlabel('Z - axis')
plt.figure()
ax2 = plt.subplot(111, projection='3d')
plt.hot()
cx = np.ones_like(X) * ax1.get_xlim3d()[0]
cy = np.ones_like(X) * ax1.get_ylim3d()[1]
cz = np.ones_like(Z) * ax1.get_zlim3d()[0]
ax2.scatter(X, Y, cz, c=Z, marker='.', lw=0, alpha=0.2)
ax2.scatter(X, cy, Z, c=-Y, marker='.', lw=0, alpha=0.2)
ax2.scatter(cx, Y, Z, c=X, marker='.', lw=0, alpha=0.2)
ax2.set_xlim3d(ax1.get_xlim3d())
ax2.set_ylim3d(ax1.get_ylim3d())
ax2.set_zlim3d(ax1.get_zlim3d())
ax2.set_xlabel('X - axis')
ax2.set_ylabel('Y - axis')
ax2.set_zlabel('Z - axis')

According to what you want to do you need to use the zdir parameter for the contour and contourf functions. Here an example:
import matplotlib.pyplot as plt
import numpy as np
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
ax = Axes3D(fig)
X = np.arange(-4, 4, 0.25)
Y = np.arange(-4, 4, 0.25)
X, Y = np.meshgrid(X, Y)
R = np.sqrt(X ** 2 + Y ** 2)
Z = np.sin(R)
ax.contourf(X, Y, Z, zdir='x', offset=-4, cmap=plt.cm.hot)
ax.contour(X, Y, Z, zdir='x', offset=-4, colors='k')
ax.contourf(X, Y, Z, zdir='y', offset=4, cmap=plt.cm.hot)
ax.contour(X, Y, Z, zdir='y', offset=4, colors='k')
ax.contourf(X, Y, Z, zdir='z', offset=-1, cmap=plt.cm.hot)
ax.contour(X, Y, Z, zdir='z', offset=-1, colors='k')
plt.show()
With result:

Here's hoping I'm not shooting a mosquito with a [very ugly] cannon, I recently made a Binning_Array class for a school project and I think it might be able to help you:
import numpy as np
import matplotlib.pyplot as plt
import binning_array as ba
from mpl_toolkits.mplot3d import axes3d
data = np.loadtxt('test.cat')
X = data[:,0]
Y = data[:,1]
Z = data[:,2]
n_points = data.shape[0]
X_min = np.round(np.min(data[:,0])-0.5)
X_max = np.round(np.max(data[:,0])+0.5)
Y_min = np.round(np.min(data[:,1])-0.5)
Y_max = np.round(np.max(data[:,1])+0.5)
Z_min = np.round(np.min(data[:,2])-0.5)
Z_max = np.round(np.max(data[:,2])+0.5)
n_min_bins = 25
step = min([(X_max-X_min)/n_min_bins, (Y_max-Y_min)/n_min_bins, (Z_max-Z_min)/n_min_bins])
# Using three Binners
BinnerXY = ba.Binning_Array([[X_min, X_max, step],
[Y_min, Y_max, step]])
BinnerYZ = ba.Binning_Array([[Y_min, Y_max, step],
[Z_min, Z_max, step]])
BinnerXZ = ba.Binning_Array([[X_min, X_max, step],
[Z_min, Z_max, step]])
for point in data:
BinnerXY.add_value([point[0], point[1]])
BinnerXZ.add_value([point[0], point[2]])
BinnerYZ.add_value([point[1], point[2]])
fig = plt.figure()
ax = [fig.add_subplot(221, projection='3d'),
fig.add_subplot(222),
fig.add_subplot(223),
fig.add_subplot(224)]
# Plot 2D projections on the 3D graph
vmin = np.min([BinnerXZ.bin_min(), BinnerYZ.bin_min(), BinnerXY.bin_min()])
vmax = np.max([BinnerXZ.bin_max(), BinnerYZ.bin_max(), BinnerXY.bin_max()])
levels = np.linspace(vmin,vmax,20)
xs_c = np.arange(*BinnerXZ.limits[0])
zs_c = np.arange(*BinnerXZ.limits[1])
ZS_C, XS_C = np.meshgrid(zs_c,xs_c)
ax[0].contourf(X=XS_C, Y=BinnerXZ.bins, Z=ZS_C,
zdir='y', offset=Y_max,
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm, levels=levels,
alpha=0.5)
xs_c = np.arange(*BinnerXY.limits[0])
ys_c = np.arange(*BinnerXY.limits[1])
YS_C, XS_C = np.meshgrid(ys_c,xs_c)
ax[0].contourf(X=XS_C, Y=YS_C, Z=BinnerXY.bins,
zdir='z', offset=Z_min,
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm, levels=levels,
alpha=0.5)
ys_c = np.arange(*BinnerYZ.limits[0])
zs_c = np.arange(*BinnerYZ.limits[1])
ZS_C, YS_C = np.meshgrid(zs_c, ys_c)
ax[0].contourf(X=BinnerYZ.bins, Y=YS_C, Z=ZS_C,
zdir='x', offset=X_min,
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm, levels=levels,
alpha=0.5)
# Plot scatter of all data
ax[0].scatter(X, Y, Z, c='g', marker='.', alpha=0.2)
ax[0].set_xlabel(r"$x$")
ax[0].set_ylabel(r"$y$")
ax[0].set_zlabel(r"$z$")
max_range = max([X_max-X_min, Y_max-Y_min, Z_max-Z_min]) / 2.
pos = [(X_max+X_min)/2., (Y_max+Y_min)/2., (Z_max+Z_min)/2.]
ax[0].set_xlim(pos[0] - max_range, pos[0] + max_range)
ax[0].set_ylim(pos[1] - max_range, pos[1] + max_range)
ax[0].set_zlim(pos[2] - max_range, pos[2] + max_range)
# Plot 2D histograms
BinnerXZ.plot_2d_slice(fig=fig, ax=ax[1], xlabel=r"$x$", ylabel=r'$z$')
BinnerXY.plot_2d_slice(fig=fig, ax=ax[2], xlabel=r"$x$", ylabel=r'$y$')
BinnerYZ.plot_2d_slice(fig=fig, ax=ax[3], xlabel=r"$y$", ylabel=r'$z$')
plt.show()
You can also use only one Binner, but notice that you will get artifacts where the planes intersect:
# ...
# Using three Binners
# ...
# Using only one Binner (adds a small error! see comments!)
Binner = ba.Binning_Array([[X_min, X_max, step],
[Y_min, Y_max, step],
[Z_min, Z_max, step]])
for point in data:
Binner.add_value([point[0], point[1], Z_min])
Binner.add_value([point[0], Y_max-step, point[2]])
Binner.add_value([X_min, point[1], point[2]])
fig = plt.figure()
ax = [fig.add_subplot(221, projection='3d'),
fig.add_subplot(222),
fig.add_subplot(223),
fig.add_subplot(224)]
ax[0].scatter(X, Y, Z, c='g', marker='.', alpha=0.2)
Binner.plot_slices(others={0:X_min, 1:Y_max, 2:Z_min}, fig=fig, ax=ax)
plt.show()
The binning_array.py was made for a school project and is not entirely polished, but it's enough for what you want.
import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import axes3d
class Binning_Array:
def __init__(self, limits=[[-1.,1.,1.],[-1.,1.,1.]]):
"""Create a new binning array.
The amount of given limits determines the dimension of the array,
although only 2 and 3D have been tested.
Each limit must be a list of start, stop and step for the
axis it represents (x, y or z)."""
self.limits = np.array(limits)
self._shape = []
for i in xrange(len(self.limits)):
self._shape.append((self.limits[i][1]-self.limits[i][0]) / \
float(self.limits[i][2]))
self._shape = tuple(self._shape)
self.dimensions = len(self._shape)
self.bins = np.zeros(self._shape)
self.outside = 0.
self._normalized = 1.
def __repr__(self):
"""Representation method. <<REVIEW>>"""
return "Binning Array! Hurray!"
def __getitem__(self, index):
"""Direct acess to read self.bins (use something like self[index])
Direct acess to write would be given by __setitem__
"""
return self.bins.__getitem__(index)
def position2index(self,position,axis=0):
"""Convert a given position to an index in axis.
If it is outside, it returns -1.
"""
if self.limits[axis][0] <= position < self.limits[axis][1]:
return int((position - self.limits[axis][0]) / self.limits[axis][2])
else: return -1
def index2position(self, index, axis=0):
"""Convert a given index to a position in axis.
If it is outisde, it returns -1.
"""
if 0 <= index < self._shape[axis]:
return self.limits[axis][0] + self.limits[axis][2] * index
else:
return -1
def add_value(self, position, value=1., verbose=False):
"""Add a given value to a specified position.
If verbose it returns a list with the axies at which the position
is outside the scope of this Binning_Array.
Not very efficient because of that (verbose was for debugging).
"""
indexs = []
outside = False
if verbose:
outs = []
for i in xrange(self.dimensions):
# using self.dimensions serves as a filter
# if position has non valid shape
index = self.position2index(position[i],i)
if index == -1:
if verbose:
outside = True
outs.append(i)
else:
self.outside += value / self._normalized
return None # nothing, as it is not verbose
else:
indexs.append(index)
if outside: # the only way to get here is if verbose is True...
self.outside += value / self._normalized
return outs # so I can just return outs...
else:
self.bins[tuple(indexs)] += value / self._normalized
if verbose:
return outs
def get_value(self, position, verbose=False):
"""Return the value at the specified position.
If verbose it alse returns a list with the axies at which the position
is outside the scope of this Binning_Array.
"""
indexs = []
outside = False
if verbose:
outs = []
for i in xrange(self.dimensions):
index = self.position2index(position[i],i)
if index == -1:
if verbose:
outside = True
outs.append[i]
else:
return self.outside
else:
indexs.append(index)
if outside: # the only way to get here is if verbose is True
return self.outside, outs # so I can just return outs...
else:
if verbose:
return self.bins[tuple(indexs)], outs
else:
return self.bins[tuple(indexs)]
def normalize(self, total=None):
"""Divide the entire array by the sum of its values (and outside).
Any value added after this will be normalized by the same factor.
"""
if total is None:
total = self.n_counts()
self.bins /= total
self.outside /= total
self.normalize *= total
def n_counts(self):
"""Return the number of counts."""
return np.sum(self.bins) + self.outside
def bin_max(self):
"""Return the value of the largest bin."""
return np.max(self.bins)
def bin_min(self):
"""Return the value of the largest bin."""
return np.min(self.bins)
def plot_2d_slice(self, cuts=[0,1], others={},
fig=None, ax=None, show=True, **kwargs):
"""Plot a 2D slice."""
x = min(cuts)
y = max(cuts)
xs = np.arange(self.limits[x][0],
self.limits[x][1] + self.limits[x][2],
self.limits[x][2])
ys = np.arange(self.limits[y][0],
self.limits[y][1] + self.limits[y][2],
self.limits[y][2])
index = []
title = ''
for i in xrange(self.dimensions):
if i in cuts:
appendix = slice(self._shape[i]+1)
else:
appendix = others.get(i,(self.limits[i][0]+
self.limits[i][1]) / 2.)
title += '%d:%.4e\t' % (i,appendix)
appendix = self.position2index(appendix,i)
index.append(appendix)
index = tuple(index)
if fig is None:
fig, ax = plt.subplots(1,1)
YS,XS = np.meshgrid(ys, xs)
graph = ax.pcolormesh (XS, YS, self.bins[index], cmap=plt.cm.coolwarm)
fig.colorbar(graph, ax=ax)
ax.axis('equal')
ax.set_xlim(self.limits[x][0], self.limits[x][1])
ax.set_ylim(self.limits[y][0], self.limits[y][1])
if 'xticks' in kwargs:
ax.set_xticks(kwargs['xticks'])
if 'yticks' in kwargs.keys():
ax.set_yticks(kwargs['yticks'])
if 'xlabel' in kwargs:
ax.set_xlabel(kwargs['xlabel'])
if 'ylabel' in kwargs:
ax.set_ylabel(kwargs['ylabel'])
if 'xlim' in kwargs:
ax.set_xlim(*kwargs['xlim'])
if 'ylim' in kwargs:
ax.set_ylim(*kwargs['ylim'])
if show:
fig.tight_layout()
fig.show()
def plot_slices(self, others={}, fig=None, ax=None,
show=True, projections=True):
index = []
pos = []
title = ''
for i in xrange(self.dimensions):
temp = others.get(i,(self.limits[i][0]+self.limits[i][1])/2.)
title += '%d:%.4e\t' % (i,temp)
pos.append(temp)
index.append(self.position2index(temp,i))
if self.dimensions == 3:
if fig is None:
fig = plt.figure()
if projections:
ax = [fig.add_subplot(221, projection='3d'),
fig.add_subplot(222),
fig.add_subplot(223),
fig.add_subplot(224)]
else:
ax = fig.add_subplot(111, projection='3d')
if projections:
xs = np.arange(self.limits[0][0],
self.limits[0][1] + self.limits[0][2],
self.limits[0][2])
ys = np.arange(self.limits[1][0],
self.limits[1][1] + self.limits[1][2],
self.limits[1][2])
zs = np.arange(self.limits[2][0],
self.limits[2][1] + self.limits[2][2],
self.limits[2][2])
xs_c = np.arange(*self.limits[0])
ys_c = np.arange(*self.limits[1])
zs_c = np.arange(*self.limits[2])
vmin = np.min(self.bins)
vmax = np.max(self.bins)
levels = np.linspace(vmin,vmax,20)
#graph 0 (3D)
ax[0].set_xlabel(r"$x$")
ax[0].set_ylabel(r"$y$")
ax[0].set_zlabel(r"$z$")
#ax[0].axis('equal') #not supported in 3D:
#http://stackoverflow.com/questions/13685386/\
#matplotlib-equal-unit-length-with-equal-aspect-ratio-z-axis-is-not-equal-to
max_range = max([xs[-1]-xs[0],ys[-1]-ys[0],zs[-1]-zs[0]]) / 2.
# x_mean = (xs[-1] + xs[0])/2.
# y_mean = (ys[-1] + ys[0])/2.
# z_mean = (zs[-1] +zs[0])/2.
ax[0].set_xlim(pos[0] - max_range, pos[0] + max_range)
ax[0].set_ylim(pos[1] - max_range, pos[1] + max_range)
ax[0].set_zlim(pos[2] - max_range, pos[2] + max_range)
# to understand holes in contour plot:
#http://stackoverflow.com/questions/18897950/\
#matplotlib-pyplot-contourf-function-introduces-holes-or-gaps-when-plotting-regul
# graph 1 (2D)
ZS, XS = np.meshgrid(zs,xs)
ZS_C, XS_C = np.meshgrid(zs_c,xs_c)
ax[1].pcolormesh(XS, ZS, self.bins[:,index[1],:],
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm)
ax[0].contourf(X=XS_C, Y=self.bins[:,index[1],:], Z=ZS_C,
zdir='y', offset=pos[1],
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm, levels=levels,
alpha=0.5)
ax[1].set_xlabel(r"$x$")
ax[1].set_ylabel(r"$z$")
ax[1].set_xlim(xs[0],xs[-1])
ax[1].set_ylim(zs[0],zs[-1])
ax[1].axis('equal')
# graph 2 (2D)
YS, XS = np.meshgrid(ys,xs)
YS_C, XS_C = np.meshgrid(ys_c,xs_c)
ax[2].pcolormesh(XS, YS, self.bins[:,:,index[2]],
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm)
ax[0].contourf(X=XS_C, Y=YS_C, Z=self.bins[:,:,index[2]],
zdir='z', offset=pos[2],
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm, levels=levels,
alpha=0.5)
ax[2].set_xlabel(r"$x$")
ax[2].set_ylabel(r"$y$")
ax[2].set_xlim(xs[0],xs[-1])
ax[2].set_ylim(ys[0],ys[-1])
ax[2].axis('equal')
# graph 3 (2D)
ZS, YS = np.meshgrid(zs, ys)
ZS_C, YS_C = np.meshgrid(zs_c, ys_c)
ax[3].pcolormesh(YS, ZS, self.bins[index[0],:,:],
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm)
ax[0].contourf(X=self.bins[index[0],:,:], Y=YS_C, Z=ZS_C,
zdir='x', offset=pos[0],
vmin=vmin, vmax=vmax,
cmap=plt.cm.coolwarm, levels=levels,
alpha=0.5)
ax[3].set_xlabel(r"$y$")
ax[3].set_ylabel(r"$z$")
ax[3].set_xlim(ys[0],ys[-1])
ax[3].set_ylim(zs[0],zs[-1])
ax[3].axis('equal')
else:
# update to draw a given slice, use it to plot eaxh axes above!
ax.plot(self.XS,self.YS,self.ZS)
ax.set_zlabel(r"$z$")
ax.set_xlabel(r"$x$")
ax.set_ylabel(r"$y$")
ax.axis('equal')
else:
if fig is None:
fig, ax = plt.subplots(1)
xs = np.arange(self.limits[0][0],
self.limits[0][1] + self.limits[0][2],
self.limits[0][2])
ys = np.arange(self.limits[1][0],
self.limits[1][1] + self.limits[1][2],
self.limits[1][2],)
YS, XS = np.meshgrid(ys, xs)
graph = ax.pcolormesh(XS, YS, self.bins, cmap=plt.cm.coolwarm)
fig.colorbar(graph)
ax.set_xlim(self.limits[0][0], self.limits[0][1])
ax.set_ylim(self.limits[1][0], self.limits[1][1])
ax.set_title('Energy Distribution')
ax.set_xlabel(r"$x$")
ax.set_ylabel(r"$y$")
ax.axis('equal')
if show:
fig.tight_layout()
fig.show()
return fig, ax
If anything on the code is wrong or bugged please say so and I will edit the above (the school project was already graded, so you won't be doing my homework).
Also, if anything is less than clear please say so I add comments or explanations as needed.