def initialize_rope_from_cloud(xyzs, plotting=False):
xyzs = xyzs.reshape(-1,3)
if len(xyzs) > 500: xyzs = xyzs[::len(xyzs)//500]
pdists = ssd.squareform(ssd.pdist(xyzs,'sqeuclidean'))
G = nx.Graph()
for (i_from, row) in enumerate(pdists):
to_inds = np.flatnonzero(row[:i_from] < MAX_DIST**2)
for i_to in to_inds:
G.add_edge(i_from, i_to, weight = pdists[i_from, i_to])
A = nx.floyd_warshall_numpy(G)
A[np.isinf(A)] = 0
(i_from_long, i_to_long) = np.unravel_index(A.argmax(), A.shape)
nodes = G.nodes();
path = nx.shortest_path(G, source=nodes[i_from_long], target=nodes[i_to_long])
xyz_path = xyzs[path,:]
xyzs_unif = curves.unif_resample(xyz_path,N_CTRL_PTS,tol=.005)
labels = np.ones(len(xyzs_unif)-1,'int')
labels[[0,-1]] = 2
if plotting:
import enthought.mayavi.mlab as mlab
mlab.plot3d(*xyzs_unif.T, tube_radius=.001)
mlab.points3d(*xyzs.T, scale_factor=.01)
mlab.show()
return xyzs_unif, labels
def plot_cuboid(corner_tups):
for tup in corner_tups:
p1 = tup[0]
p2 = tup[1]
mlab.plot3d([p1[0,0],p2[0,0]],[p1[1,0],p2[1,0]],
[p1[2,0],p2[2,0]],color=(1.,1.,0.),
representation='wireframe',tube_radius=None)
def viewImgWithNodes(img, spacing, contours,g, title=''):
mlab.figure(bgcolor=(0, 0, 0), size=(900, 900))
#src = mlab.pipeline.scalar_field(img)
## Our data is not equally spaced in all directions:
#src.spacing = [1, 1, 1]
#src.update_image_data = True
#
#mlab.pipeline.iso_surface(src, contours=contours, opacity=0.2)
nodes = np.array(g.nodes())
dsize = 4*np.ones(nodes.shape[0],dtype='float32')
print(dsize.shape,nodes.shape)
#mlab.points3d(nodes[:,0],nodes[:,1],nodes[:,2],color=(0.0,1.0,0.0))
mlab.points3d(nodes[:,2],nodes[:,1],nodes[:,0],dsize,color=(0.0,0.0,1.0), scale_factor=0.25)
for n1, n2, edge in g.edges(data=True):
path = [n1]+edge['path']+[n2]
pa = np.array(path)
#print pa
mlab.plot3d(pa[:,2],pa[:,1],pa[:,0],color=(0,1,0),tube_radius=0.25)
mlab.view(-125, 54, 'auto','auto')
mlab.roll(-175)
mlab.title(title, height=0.1)
mlab.show()
def drawArrows(self, figure, bbox, number, partner):
"""
draw an 'arrow' from the cell 'number' to 'partner'
"""
#load the center of mass position
n = np.array(number)
p = np.array(partner)
n.shape = (n.size, )
p.shape = (p.size, )
if p.size > 0 and n.size > 0:
com1 = self.file1["features"][str(n[0])]["com"][:]
com2 = self.file2["features"][str(p[0])]["com"][:]
#write the cell label as text
mlab.text3d(com1[2]-bbox[2]+1,com1[1]-bbox[1]+1,com1[0]-bbox[0]+1, str(number), color=(1,1,1), figure=figure)
#plot a point where the current cell is
mlab.points3d([com1[2]-bbox[2]+1],[com1[1]-bbox[1]+1],[com1[0]-bbox[0]+1],color=(0,0,1), figure=figure)
#plot a line to the descendant's center
mlab.plot3d([com1[2]-bbox[2]+1,com2[2]-bbox[2]+1],
[com1[1]-bbox[1]+1,com2[1]-bbox[1]+1],
[com1[0]-bbox[0]+1,com2[0]-bbox[0]+1],
tube_radius=0.2, color=(1,0,0), figure=figure)
#plot a second line, if there is a split
if p.size == 2:
com3 = self.file2["features"][str(p[1])]["com"][:]
mlab.plot3d([com1[2]-bbox[2]+1,com3[2]-bbox[2]+1],
[com1[1]-bbox[1]+1,com3[1]-bbox[1]+1],
[com1[0]-bbox[0]+1,com3[0]-bbox[0]+1],
tube_radius=0.2, color=(1,0,0), figure=figure)
def ChargeDensityFog3D(directory, save_file=None , DrawAtoms=True, DrawCell=True, opacity_range=[0,1]):
# Get data from calculation files
with jasp(directory) as calc:
atoms = calc.get_atoms()
x, y, z, cd = calc.get_charge_density()
mlab.figure(bgcolor=(0,0,0), size=(640,480))
# Draw atoms
if DrawAtoms == True:
for atom in atoms:
mlab.points3d(atom.x,
atom.y,
atom.z,
scale_factor=vdw_radii[atom.number]/10.,
resolution=16,
color=tuple(cpk_colors[atom.number]),
scale_mode='none')
# Draw unit cell
if DrawCell == True:
a1, a2, a3 = atoms.get_cell()
origin = [0,0,0]
cell_matrix = [[origin, a1],
[origin, a2],
[origin, a3],
[a1, a1+a2],
[a1, a1+a3],
[a2, a2+a1],
[a2, a2+a3],
[a3, a1+a3],
[a3, a2+a3],
[a1+a2, a1+a2+a3],
[a2+a3, a1+a2+a3],
[a1+a3, a1+a3+a2]] # contains all points on the box
for p1, p2 in cell_matrix:
mlab.plot3d([p1[0], p2[0]], # x-coords of box
[p1[1], p2[1]], # y-coords
[p1[2], p2[2]]) # z-coords
# Plot the charge density
src = mlab.pipeline.scalar_field(x, y, z, cd) #Source data
vmin = cd.min() #find minimum and maximum value of CD data
vmax = cd.max()
vol = mlab.pipeline.volume(src) # Make a volumetric representation of the data
# Set opacity transfer function
from tvtk.util.ctf import PiecewiseFunction
otf = PiecewiseFunction()
otf.add_point(vmin, opacity_range[0]) #Transparency at zero electron density
otf.add_point(vmax*1, opacity_range[1]) #Transparency at max electron density
vol._otf=otf
vol._volume_property.set_scalar_opacity(otf)
#Show a legend
mlab.colorbar(title="e- density\n(e/Ang^3)", orientation="vertical", nb_labels=5, label_fmt='%.2f')
mlab.view(azimuth=-90, elevation=90, distance='auto') # Set viewing angle
mlab.show()
if save_file != None:
mlab.savefig(save_file)
def _show_button_fired( self ):
p, r = self.data.slice
mlab.figure( 'Continuous fibers' )
mlab.plot3d( p[:, 0], p[:, 1], p[:, 2], r,
tube_radius=20, tube_sides=20, colormap='Spectral' )
def plot(pts,color=(1.,1.,1.), scalar_list=None):
if scalar_list != None:
mlab.plot3d(pts[0,:].A1,pts[1,:].A1,pts[2,:].A1,scalar_list,
representation = 'wireframe', tube_radius = None)
mlab.colorbar()
else:
mlab.plot3d(pts[0,:].A1,pts[1,:].A1,pts[2,:].A1,color=color,
representation = 'wireframe', tube_radius = None)
def plot_cuboid(corner_tups):
for tup in corner_tups:
p1 = tup[0]
p2 = tup[1]
mlab.plot3d([p1[0, 0], p2[0, 0]], [p1[1, 0], p2[1, 0]],
[p1[2, 0], p2[2, 0]],
color=(1., 1., 0.),
representation='wireframe',
tube_radius=None)
def plotPCs3d(v, s, PCAInputMatrix):
#get extents
longest_axis = max(PCAInputMatrix.max(axis=0) - PCAInputMatrix.min(axis=0))
longest_axis /= 2
scalar = s / s[0] * longest_axis
for i, vector in enumerate(v):
scaled_vector = vector * scalar[i]
pc = zip(np.zeros(len(scaled_vector)), scaled_vector)
mlab.plot3d(np.array(pc[0]), np.array(pc[1]), np.array(pc[2]), color=(1, 0, 0))
def plot_cartesian(traj,
xaxis=None,
yaxis=None,
zaxis=None,
color='b',
label='_nolegend_',
linewidth=2,
scatter_size=20):
''' xaxis - x axis for the graph (0,1 or 2)
zaxis - for a 3d plot. not implemented.
'''
import arm_trajectories as at
#if traj.__class__ == at.JointTrajectory:
if isinstance(traj, at.JointTrajectory):
traj = joint_to_cartesian(traj)
pts = np.matrix(traj.p_list).T
label_list = ['X coord (m)', 'Y coord (m)', 'Z coord (m)']
x = pts[xaxis, :].A1.tolist()
y = pts[yaxis, :].A1.tolist()
if zaxis == None:
pl.plot(x, y, c=color, linewidth=linewidth, label=label)
pl.scatter(
x, y, c=color, s=scatter_size, label='_nolegend_', linewidths=0)
pl.xlabel(label_list[xaxis])
pl.ylabel(label_list[yaxis])
pl.legend(loc='best')
pl.axis('equal')
else:
from numpy import array
from enthought.mayavi.api import Engine
engine = Engine()
engine.start()
if len(engine.scenes) == 0:
engine.new_scene()
z = pts[zaxis, :].A1.tolist()
time_list = [t - traj.time_list[0] for t in traj.time_list]
mlab.plot3d(x, y, z, time_list, tube_radius=None, line_width=4)
mlab.axes()
mlab.xlabel(label_list[xaxis])
mlab.ylabel(label_list[yaxis])
mlab.zlabel(label_list[zaxis])
mlab.colorbar(title='Time')
# -------------------------------------------
axes = engine.scenes[0].children[0].children[0].children[1]
axes.axes.position = array([0., 0.])
axes.axes.label_format = '%-#6.2g'
axes.title_text_property.font_size = 4
def plot_cartesian(traj, xaxis=None, yaxis=None, zaxis=None, color='b',label='_nolegend_',
linewidth=2, scatter_size=10, plot_velocity=False):
import arm_trajectories as at
#if traj.__class__ == at.JointTrajectory:
if isinstance(traj,at.JointTrajectory):
traj = joint_to_cartesian(traj)
pts = np.matrix(traj.p_list).T
label_list = ['X coord (m)', 'Y coord (m)', 'Z coord (m)']
x = pts[xaxis,:].A1.tolist()
y = pts[yaxis,:].A1.tolist()
if plot_velocity:
vels = np.matrix(traj.v_list).T
xvel = vels[xaxis,:].A1.tolist()
yvel = vels[yaxis,:].A1.tolist()
if zaxis == None:
mpu.plot_yx(y, x, color, linewidth, '-', scatter_size, label,
axis = 'equal', xlabel = label_list[xaxis],
ylabel = label_list[yaxis],)
if plot_velocity:
mpu.plot_quiver_yxv(y, x, np.matrix([xvel,yvel]),
width = 0.001, scale = 1.)
mpu.legend()
else:
from numpy import array
from enthought.mayavi.api import Engine
engine = Engine()
engine.start()
if len(engine.scenes) == 0:
engine.new_scene()
z = pts[zaxis,:].A1.tolist()
time_list = [t-traj.time_list[0] for t in traj.time_list]
mlab.plot3d(x,y,z,time_list,tube_radius=None,line_width=4)
mlab.axes()
mlab.xlabel(label_list[xaxis])
mlab.ylabel(label_list[yaxis])
mlab.zlabel(label_list[zaxis])
mlab.colorbar(title='Time')
# -------------------------------------------
axes = engine.scenes[0].children[0].children[0].children[1]
axes.axes.position = array([ 0., 0.])
axes.axes.label_format = '%-#6.2g'
axes.title_text_property.font_size=4
def plot_cartesian(traj, xaxis=None, yaxis=None, zaxis=None, color='b',label='_nolegend_',
linewidth=2, scatter_size=10, plot_velocity=False):
import matplotlib_util.util as mpu
import arm_trajectories as at
#if traj.__class__ == at.JointTrajectory:
if isinstance(traj,at.JointTrajectory):
traj = joint_to_cartesian(traj)
pts = np.matrix(traj.p_list).T
label_list = ['X coord (m)', 'Y coord (m)', 'Z coord (m)']
x = pts[xaxis,:].A1.tolist()
y = pts[yaxis,:].A1.tolist()
if plot_velocity:
vels = np.matrix(traj.v_list).T
xvel = vels[xaxis,:].A1.tolist()
yvel = vels[yaxis,:].A1.tolist()
if zaxis == None:
mpu.plot_yx(y, x, color, linewidth, '-', scatter_size, label,
axis = 'equal', xlabel = label_list[xaxis],
ylabel = label_list[yaxis],)
if plot_velocity:
mpu.plot_quiver_yxv(y, x, np.matrix([xvel,yvel]),
width = 0.001, scale = 1.)
mpu.legend()
else:
from numpy import array
from enthought.mayavi.api import Engine
engine = Engine()
engine.start()
if len(engine.scenes) == 0:
engine.new_scene()
z = pts[zaxis,:].A1.tolist()
time_list = [t-traj.time_list[0] for t in traj.time_list]
mlab.plot3d(x,y,z,time_list,tube_radius=None,line_width=4)
mlab.axes()
mlab.xlabel(label_list[xaxis])
mlab.ylabel(label_list[yaxis])
mlab.zlabel(label_list[zaxis])
mlab.colorbar(title='Time')
# -------------------------------------------
axes = engine.scenes[0].children[0].children[0].children[1]
axes.axes.position = array([ 0., 0.])
axes.axes.label_format = '%-#6.2g'
axes.title_text_property.font_size=4
def line():
"""Generates a pretty set of lines."""
x =y = z = numpy.arange(5)
l = plot3d(x, y, z,
#tube_radius=0.025,
color=(0.862745,0.0784314,0.235294))
def _viztest3d(self, extra_bb):
"""
3D vizualization of CPRep for obj
"""
from enthought.mayavi import mlab
import numpy as np
f = mlab.figure(1, fgcolor=(0, 0, 0), bgcolor=(1, 1, 1), size=(640,640))
mlab.clf()
if self._hasParam:
L = self.ParamGrid()
x,y,z = L
#print x,y,z
#s = mlab.mesh(xb, yb, zb, scalars=real((Eplot_bulb*E*uxx).reshape(xb.shape)))
# TODO: mayavi bug, opacity<1
s = mlab.mesh(x, y, z, scalars=z, opacity=1.0)
a,b = self.getBB()
ll = (b+a)/2 - (extra_bb)*(b-a)/2
rr = (b+a)/2 + (extra_bb)*(b-a)/2
for i in range(0,40):
x = np.random.uniform( ll, rr )
cp,dist,bdy,other = self.cp(x)
colt = (0.5,0.5,0.5)
op = 0.3
l = mlab.plot3d([x[0],cp[0]], [x[1],cp[1]], [x[2], cp[2]], color=colt, opacity=op, tube_radius=0.1)
#mlab.title('3d viz test')
mlab.show()
def display(self):
""" Display Electrode in Mayavi """
f = mlab.figure(figure=mlab.gcf(),
bgcolor=(0,0,0),
size=(800,600))
# Build Electrode tip
x, y, z = self.shaft[0,:].T
self.mlab_tip = mlab.points3d(array(x),
array(y),
array(z),
scale_factor=self.diameter,
color=(0.5, 0.5, 0.5),
name='Electrode Tip')
self.mlab_tip.glyph.glyph_source.glyph_source.theta_resolution = 30
self.mlab_tip.glyph.glyph_source.glyph_source.phi_resolution = 16
# Build Electrode Shaft
x, y, z = self.shaft.T
self.mlab_shaft = mlab.plot3d(x.T, y.T, z.T,
color=(0.5, 0.5, 0.5),
tube_radius=self.diameter / 2,
tube_sides=40,
name='Electrode')
self.mlab_shaft.parent.parent.filter.capping=True
def line(pt1, pt2, tmin, tmax, numpts=30):
t = linspace(tmin, tmax, numpts)
x = pt1[0] + (pt2[0] - pt1[0]) * t
y = pt1[1] + (pt2[1] - pt1[1]) * t
z = pt1[2] + (pt2[2] - pt1[2]) * t
obj = plot3d(x, y, z, color=(0.1, 0.1, 1))
return (x, y, z)
def draw_3d_tree(X, Y, Z, index_edges):
for a, b in index_edges:
foo = mlab.plot3d(
[X[a], X[b]],
[Y[a], Y[b]],
[Z[a], Z[b]],
color=(.5, .5, .5))
def renderSnapshot(self, t):
sources = []
for e in self._endEffectors:
px, py, pz = self._generateDataPoints(t, e)
sources.append(mlab.plot3d(px, py, pz, *self._mlab_args,
**self._mlab_kwargs).mlab_source)
return sources
def plot_orbits3d(self,xaxis='x',yaxis='y',clf=True,**kwargs):
"""
Plots a 3D plot of the particles and their orbits with mayavi.
:param bool clf: If True, clear figure before plotting.
kwargs are passed into :func:`enthough.mayavi.mlab.plot3d`.
:returns:
the result of :func:`enthough.mayavi.mlab.plot3d` and the result of
:func:`enthough.mayavi.mlab.quiver3d`
"""
from enthought.mayavi import mlab as M
if clf:
M.clf()
orbarr = self.getOrbits()[1].transpose((1,2,0))
t = self._orbt
pntress = []
for o in orbarr:
pntress.append(M.plot3d(o[0],o[1],o[2],t,**kwargs))
xp,yp,zp = self.particles.position.T
vx,vy,vz = self.particles.velocity.T
quiverres = M.quiver3d(xp,yp,zp,vx,vy,vz)
return pntress,quiverres
def plot(pts, color=(1., 1., 1.), scalar_list=None):
if scalar_list != None:
mlab.plot3d(pts[0, :].A1,
pts[1, :].A1,
pts[2, :].A1,
scalar_list,
representation='wireframe',
tube_radius=None)
mlab.colorbar()
else:
mlab.plot3d(pts[0, :].A1,
pts[1, :].A1,
pts[2, :].A1,
color=color,
representation='wireframe',
tube_radius=None)
def plot_lat_long_lines():
#latitude = linspace(-90,90,19) * pi/180.
#longitude = linspace(0, 360, 37) * pi/180.
latitude = linspace(-20,90,10) * pi/180.
#longitude = linspace(-90,90,19) * pi/180.
longitude = linspace(-270,-90,19) * pi/180.
theta,phi = meshgrid(longitude,latitude)
X = cos(theta)*cos(phi)
Y = sin(theta)*cos(phi)
Z = sin(phi)
for i in range(X.shape[0]):
mlab.plot3d(X[i,:],Y[i,:],Z[i,:],color=(0,0,0),tube_radius=None,line_width=1)
for j in range(X.shape[1]):
mlab.plot3d(X[:,j],Y[:,j],Z[:,j],color=(0,0,0),tube_radius=None,line_width=1)
def plot_canada_map():
"""
Plot coastline as tubes in 3D.
"""
d = defaultdict(list)
gmt = GMT(config={'PAGE_ORIENTATION':'landscape'})
#rng = '-70/-52/40/52.'
rng = '-150/-49/40./89'
scl = 'M10c'
gmt.pscoast(R=rng,J=scl,B='a0.5wsne',D='l',W='thinnest',m=True,A='2/1/1')
a = gmt.output.getvalue().split('\n')
z = 0.
cnt = 0
connections = list()
for _l in a:
if _l.startswith('#'):continue
if _l.startswith('>'):
cnt += 1
continue
try:
d[cnt].append(map(float,_l.split('\t')))
except ValueError:
print _l
for _k in d.keys():
ar = np.array(d[_k])
x = (6371)*np.cos(2*np.pi*ar[:,1]/360.)*np.cos(2*np.pi*ar[:,0]/360.)
y = (6371)*np.cos(2*np.pi*ar[:,1]/360.)*np.sin(2*np.pi*ar[:,0]/360.)
z = (6371)*np.sin(2*np.pi*ar[:,1]/360.)
pts = mlab.plot3d(x,y,z,tube_radius=2.0,color=(0,0,0))
def connection_plot(r1,r2):
segments = 10
dr = (r2 - r1)/segments
'''
for n in xrange(segments):
if not n%3:
x1,y1,z1 = [k[0] for k in zip(r1+n*dr)]
x2,y2,z2 = [k[0] for k in zip(r1+(n+1)*dr)]
mlab.plot3d( [x1,x2], [y1,y2], [z1,z2],color=(0,0,.5),tube_radius=.005 )
'''
X,Y,Z = zip(*[r1 + n*dr for n in xrange(segments)])
C = xrange(len(X))
mlab.plot3d( X,Y,Z, C, tube_radius=.02, tube_sides=fig_res, opacity=.6, colormap="Blues" )
def initialize_rope(label_img, xyz,bgr, plotting=False):
# XXX that sucks
rope_mask = (label_img==1) | ndi.morphology.binary_dilation(label_img==2,np.ones((5,5)))
rope_mask = ndi.binary_opening(rope_mask, np.ones((3,3)))
rope_mask = ndi.binary_dilation(rope_mask, np.ones((15,15)))
skel = mip.skeletonize(rope_mask)
if plotting:
cv2.imshow('bgr',bgr.copy())
cv2.imshow('labels',label_img.astype('uint8')*50)
cv2.imshow("skel", skel.astype('uint8')*50)
cv2.imshow("rope_mask", rope_mask.astype('uint8')*50)
cv2.waitKey(5)
#start_end = (uv_start, uv_end) = get_cc_centers(label_img==2,2)
G = skel2graph(skel)
path2d = longest_shortest_path(G)
path3d = to3d(path2d,xyz)
xyzs_unif = curves.unif_resample(path3d,N_CTRL_PTS,tol=.0025)
#us,vs = xyz2uv(xyzs_unif).T
#labels = label_img[us,vs]
labels = np.ones(xyzs_unif.shape[0]-1,'int')
labels[0] = labels[-1] = 2
if plotting:
xs,ys,zs = xyzs_unif.T
us_skel, vs_skel = np.nonzero(skel)
xs_skel, ys_skel, zs_skel = to3d(np.c_[us_skel, vs_skel], xyz).T
import matplotlib.pyplot as plt, enthought.mayavi.mlab as mlab
mlab.plot3d(xs,ys,zs,np.arange(len(xs)),tube_radius=.0025, colormap='spectral')
mlab.points3d(xs_skel, ys_skel, zs_skel, scale_factor=.02, color=(1,1,1),opacity=.1)
for (i,path) in enumerate(path3d):
mlab.text3d(path[0,0],path[0,1],path[0,2],str(2*i),scale=.01,color=(0,0,0))
mlab.text3d(path[-1,0],path[-1,1],path[-1,2],str(2*i+1),scale=.01,color=(0,0,0))
mlab.show()
return xyzs_unif, labels
def __init__( self, x0=0, y0=0, z0=0, scale=1.0, color=(0,1,0), line_width=3, **kwargs ):
from enthought.mayavi import mlab
fig = mlab.gcf()
render = fig.scene.disable_render
fig.scene.disable_render = True
xx = x0 + scale / 200.0 * np.array( [
[ -49, -40, np.nan ],
[ 51, 60, np.nan ],
[ -60, -51, np.nan ],
[ 40, 49, np.nan ],
[ -30, 50, np.nan ],
[ -40, 40, np.nan ],
[ -50, 30, np.nan ],
] )
yy = y0 + scale / 200.0 * np.array( [
[ 10, 90, np.nan ],
[ 10, 90, np.nan ],
[ -90, -10, np.nan ],
[ -90, -10, np.nan ],
[ 100, 100, np.nan ],
[ 0, 0, np.nan ],
[ -100, -100, np.nan ],
] )
zz = z0 * np.ones_like( xx )
glyphs = [5], [0,2,4,5,6], [0,3], [0,2], [2,4,6], [1,2], [1], [0,2,5,6], [], [2]
hh = []
for g in glyphs:
i = np.array( [ i for i in range(7) if i not in g ] )
h = []
for x in -0.875, 0.125, 0.875:
h += [ mlab.plot3d(
scale * x + xx[i].flatten(), yy[i].flatten(), zz[i].flatten(),
color=color,
tube_radius=None,
line_width=line_width,
**kwargs
) ]
hh += [h]
self.glyphs = hh
x = x0 + scale / 200.0 * np.array( [-81, -79, np.nan, -71, -69] )
y = y0 + scale / 200.0 * np.array( [-60, -40, np.nan, 40, 60] )
z = z0 * np.ones_like( x )
h = mlab.plot3d( x, y, z, color=color, line_width=line_width, tube_radius=None, **kwargs )
self.colon = h
fig.scene.disable_render = render
return
def solenoid(N=10, L=5e0, R=1e0, numpts=1000):
"""Generates a solenoid."""
z = linspace(0, L, numpts)
theta = 2 * pi * N / L * z
x = R * cos(theta)
y = R * sin(theta)
l = plot3d(x, y, z, tube_radius=0.025, colormap='Spectral')
return l
def _viztest3d(self):
"""
3D vizualization of CPRep for obj with boundary
"""
from enthought.mayavi import mlab
import numpy as np
f = mlab.figure(1, fgcolor=(0, 0, 0), bgcolor=(1, 1, 1), size=(500,700))
mlab.clf()
#s = mlab.mesh(xb, yb, zb, scalars=real(evec_b.reshape(xb.shape)))
#l = mlab.plot3d(xs, ys, zs, real(evec_s))
#mlab.title(str(ii) + ' ew=' + str(eval), size=0.2)
#mlab.show()
#mlab.savefig('b_horn' + str(ii) + '_' + str(eval) + '.png')
#s = mlab.mesh(xb, yb, zb, scalars=real((Eplot_bulb*E*uxx).reshape(xb.shape)))
#l = mlab.plot3d(xs, ys, zs, real((Eplot_stem*E*uxx).reshape(xs.shape)))
#(x1,y1),(x2,y2),(x3,y3) = mesh2d(resolution=3)
# TODO: is has param:
if self._hasParam:
L = self.ParamGrid()
x,y,z = L
#print x,y,z
#s = mlab.mesh(xb, yb, zb, scalars=real((Eplot_bulb*E*uxx).reshape(xb.shape)))
# TODO: mayavi bug, opacity<1
s = mlab.mesh(x, y, z, scalars=z, opacity=1.0)
a,b = self.getBB()
ll = (b+a)/2 - (extra_bb)*(b-a)/2
rr = (b+a)/2 + (extra_bb)*(b-a)/2
for i in range(0,200):
x = np.random.uniform( ll, rr )
cp,dist,bdy,other = self.cp(x)
drawplot = False
if bdy==1:
if (np.random.random(1) < 1.0):
drawplot = True
col = 'g'
colt = (0.5,1,.2)
op = 0.9
elif bdy==2:
if (np.random.random(1) < 1.0):
drawplot = True
col = 'b'
colt = (.2,.5,1)
op = 0.9
else:
if ( (np.random.random(1) < 0.5) and (dist <= max((b-a)/10)) ):
drawplot = True
col = 'k'
colt = (0.5,0.5,0.5)
op = 0.3
if drawplot:
l = mlab.plot3d([x[0],cp[0]], [x[1],cp[1]], [x[2], cp[2]], color=colt, opacity=op)#, tube_radius=0.1)
#mlab.title('3d viz test')
mlab.show()
def plot_tree(structure, overlap=0.5):
""" Plots a tree using mayavi. """
from enthought.mayavi import mlab
from numpy import array, sum, any, dot, concatenate, argmin
tree = build_tree(structure, overlap=overlap)
# find original positions
ids, originals = [], []
for i, node in enumerate(tree):
if id(node.center) not in ids:
ids.append(id(node.center))
originals.append(node.pos)
originals = array(originals)
# now adds neighbors, including those outside the unit cell
positions = originals.copy()
others, connections = [], []
for i, node in enumerate(tree):
for neighbor, vector in node:
position = neighbor.pos + dot(structure.cell, vector)
dist = sum((positions-position)**2, axis=1)
if any(dist < 1e-8): connections.append( (i, argmin(dist)) )
else:
positions = concatenate((positions, position[None, :]))
connections.append((i, positions.shape[0]-1))
if any(sum((position[None,:]-originals)**2, axis=1) < 1e-8): continue
if len(others) \
and any(sum((position[None,:]-others)**2, axis=1) < 1e-8): continue
others.append(position)
others = array(others)
a0, a1, a2 = structure.cell.T
box = [ [0, 0, 0], a0, a0+a1, a1, [0,0,0], a2, a2+a0, a2+a1+a0, a2+a1, a2, \
a2 + a0, a0, a0+a1, a0+a1+a2, a1+a2, a1 ]
mlab.plot3d(*array(box).T, tube_radius=0.01, color=(0.8, 0.8, 0.8))
mlab.points3d(*originals.T, color=(0.8, 0.5, 0.5), scale_factor=0.2)
mlab.points3d(*others.T, color=(0.5, 0.8, 0.5), scale_factor=0.1)
# add connections
pts = mlab.points3d(*positions.T, scale_factor=0)
pts.mlab_source.dataset.lines = array(connections)
tube = mlab.pipeline.tube(pts, tube_radius=0.01)
tube.filter.radius_factor = 1.
tube.filter.vary_radius = 'vary_radius_by_scalar'
mlab.pipeline.surface(tube, color=(0.3, 0.3, 0.3))
def draw_z_crossing(x, y):
angles = list(gen_angles(20))
xs = [x + g_crossing_radius*math.cos(a) for a in angles]
ys = [y + g_crossing_radius*math.sin(a) for a in angles]
zs = np.zeros_like(xs)
foo = mlab.plot3d(
xs, ys, zs,
color=(0, 0, 1),
opacity=g_crossing_opacity)
def draw_x_crossing(y, z):
angles = list(gen_angles(20))
ys = [y + g_crossing_radius*math.cos(a) for a in angles]
zs = [z + g_crossing_radius*math.sin(a) for a in angles]
xs = np.zeros_like(ys)
foo = mlab.plot3d(
xs, ys, zs,
color=(1, 0, 0),
opacity=g_crossing_opacity)
def drawArrows(self, figure, bbox, number, partner):
"""
draw an 'arrow' from the cell 'number' to 'partner'
"""
#load the center of mass position
n = np.array(number)
p = np.array(partner)
n.shape = (n.size, )
p.shape = (p.size, )
if p.size > 0 and n.size > 0:
com1 = self.file1["features"][str(n[0])]["com"][:]
com2 = self.file2["features"][str(p[0])]["com"][:]
#write the cell label as text
mlab.text3d(com1[2] - bbox[2] + 1,
com1[1] - bbox[1] + 1,
com1[0] - bbox[0] + 1,
str(number),
color=(1, 1, 1),
figure=figure)
#plot a point where the current cell is
mlab.points3d([com1[2] - bbox[2] + 1], [com1[1] - bbox[1] + 1],
[com1[0] - bbox[0] + 1],
color=(0, 0, 1),
figure=figure)
#plot a line to the descendant's center
mlab.plot3d([com1[2] - bbox[2] + 1, com2[2] - bbox[2] + 1],
[com1[1] - bbox[1] + 1, com2[1] - bbox[1] + 1],
[com1[0] - bbox[0] + 1, com2[0] - bbox[0] + 1],
tube_radius=0.2,
color=(1, 0, 0),
figure=figure)
#plot a second line, if there is a split
if p.size == 2:
com3 = self.file2["features"][str(p[1])]["com"][:]
mlab.plot3d([com1[2] - bbox[2] + 1, com3[2] - bbox[2] + 1],
[com1[1] - bbox[1] + 1, com3[1] - bbox[1] + 1],
[com1[0] - bbox[0] + 1, com3[0] - bbox[0] + 1],
tube_radius=0.2,
color=(1, 0, 0),
figure=figure)
def display_coil(n, r0, R, half=False):
"""
Display a coils in the 3D view.
If half is True, display only one half of the coil.
"""
n, l, m = base_vectors(n)
theta = np.linspace(0, (2-half)*np.pi, 30)
theta = theta[..., np.newaxis]
coil = np.atleast_1d(R)*(np.sin(theta)*l + np.cos(theta)*m)
coil += r0
coil_x = coil[:, 0]
coil_y = coil[:, 1]
coil_z = coil[:, 2]
mlab.plot3d(coil_x, coil_y, coil_z,
tube_radius=0.01,
name='Coil %i' % display_coil.num,
color=(0, 0, 0))
display_coil.num += 1
return coil_x, coil_y, coil_z
def _plot3d(self, a, showCP, npoints):
"""
Internal plot function, use Spline.plot().
"""
if not isplot3d:
raise NotImplementedError('Lacking Mayavi to do the plotting.')
x = self(np.linspace(self.knots[2],self.knots[-3],npoints,endpoint=0))
k = self(np.hstack((self.knots[2:-3],self.knots[-3]-1e-15)))
ml.plot3d(x[:,a[0]],x[:,a[1]],x[:,a[2]],color=(.5,.2,.3))
ml.points3d(k[:,a[0]],k[:,a[1]],k[:,a[2]],
color=(.5,.2,.3),scale_factor=.3)
if showCP:
ml.plot3d(self.cp[:,a[0]], self.cp[:,a[1]],
self.cp[:,a[2]], color=(.2,.4,.4))
ml.points3d(self.cp[:,a[0]], self.cp[:,a[1]],
self.cp[:,a[2]], color=(.2,.4,.4), scale_factor=.3)
ml.show()
def display_coil(n, r0, R, half=False):
"""
Display a coils in the 3D view.
If half is True, display only one half of the coil.
"""
n, l, m = base_vectors(n)
theta = np.linspace(0, (2 - half) * np.pi, 30)
theta = theta[..., np.newaxis]
coil = np.atleast_1d(R) * (np.sin(theta) * l + np.cos(theta) * m)
coil += r0
coil_x = coil[:, 0]
coil_y = coil[:, 1]
coil_z = coil[:, 2]
mlab.plot3d(coil_x,
coil_y,
coil_z,
tube_radius=0.01,
name='Coil %i' % display_coil.num,
color=(0, 0, 0))
display_coil.num += 1
return coil_x, coil_y, coil_z
def renderSolenoidCoil(solenoid, segments=1000, **kwargs):
"""
Visualise L{SolenoidMagneticField} model coil in 3D.
@param solenoid: L{SolenoidMagneticField} to render.
@param segments: Number of coil segments to generate. Larger numbers
result in smoother coils.
@param kwargs: Additional keywords to pass to
L{enthought.mayavi.mlab.plot3d}.
"""
coils = 2*solenoid.b * solenoid.n
z = np.linspace(-solenoid.b,solenoid.b,segments)
theta = np.linspace(0,coils*2*np.pi,segments)
x = solenoid.a * np.cos(theta)
y = solenoid.a * np.sin(theta)
x,y,z = solenoid.transform.apply(np.vstack((x,y,z)))
if not 'tube_radius' in kwargs:kwargs['tube_radius']=0.001
if not 'opacity' in kwargs:kwargs['opacity']=0.5
mlab.plot3d(x,y,z,**kwargs)
def viewGraph(g, sub=3,title=''):
mlab.figure(bgcolor=(0, 0, 0), size=(900, 900))
nodes = g.nodes()
random.shuffle(nodes)
nodes = np.array(nodes[0:100])
#mlab.points3d(nodes[:,2],nodes[:,1],nodes[:,0],color=(0.0,0.0,1.0))
edges = g.edges(data=True)
print(len(edges))
input('continue')
count = 0
for n1, n2, edge in edges:
count += 1
if( count % 100 == 0 ):
print(count)
path = [n1]+edge['path']+[n2]
pa = np.array(path)
#print pa
mlab.plot3d(pa[::sub,2],pa[::sub,1],pa[::sub,0],color=(0,1,0),tube_radius=0.75)
mlab.view(-125, 54, 'auto','auto')
mlab.roll(-175)
mlab.title(title, height=0.1)
mlab.show()
def test_several_scene_models(self):
""" Check that plotting to scene attributes using their
mlab attribute does create objects as children, and does not
unset the current scene
"""
class TestObject(HasTraits):
scene1 = Instance(MlabSceneModel, ())
scene2 = Instance(MlabSceneModel, ())
test_object = TestObject()
x, y, z = np.random.random((3, 10))
plt = mlab.plot3d(x, y, z, figure=test_object.scene1.mayavi_scene)
pts = mlab.points3d(x, y, z, figure=test_object.scene2.mayavi_scene)
# Check that each figure got the module it should have
self.assertEqual(plt.scene, test_object.scene1)
self.assertEqual(pts.scene, test_object.scene2)
def plot_cylinder(start, end, tube_radius=0.1, color=(0, 0, 0)):
mlab.plot3d([start[0], end[0]], [start[1], end[1]], [start[2], end[2]],
tube_radius=tube_radius,
color=color)
for i, j in zip(ii, jj):
if linalg.norm(sens_loc[i] - sens_loc[j]) > min_dist:
con_nodes.append((i, j))
con_val.append(con[i, j])
con_val = np.array(con_val)
# Show the connections as tubes between sensors
vmax = np.max(con_val)
vmin = np.min(con_val)
for val, nodes in zip(con_val, con_nodes):
x1, y1, z1 = sens_loc[nodes[0]]
x2, y2, z2 = sens_loc[nodes[1]]
points = mlab.plot3d([x1, x2], [y1, y2], [z1, z2], [val, val],
vmin=vmin,
vmax=vmax,
tube_radius=0.001,
colormap='RdBu')
points.module_manager.scalar_lut_manager.reverse_lut = True
mlab.scalarbar(title='Phase Lag Index (PLI)', nb_labels=4)
# Add the sensor names for the connections shown
nodes_shown = list(set([n[0] for n in con_nodes] + [n[1] for n in con_nodes]))
for node in nodes_shown:
x, y, z = sens_loc[node]
mlab.text3d(x,
y,
z,
raw.ch_names[picks[node]],
def draw_z_crossing(x, y):
angles = list(gen_angles(20))
xs = [x + g_crossing_radius * math.cos(a) for a in angles]
ys = [y + g_crossing_radius * math.sin(a) for a in angles]
zs = np.zeros_like(xs)
foo = mlab.plot3d(xs, ys, zs, color=(0, 0, 1), opacity=g_crossing_opacity)
def draw_3d_tree(X, Y, Z, index_edges):
for a, b in index_edges:
foo = mlab.plot3d([X[a], X[b]], [Y[a], Y[b]], [Z[a], Z[b]],
color=(.5, .5, .5))
def draw_x_crossing(y, z):
angles = list(gen_angles(20))
ys = [y + g_crossing_radius * math.cos(a) for a in angles]
zs = [z + g_crossing_radius * math.sin(a) for a in angles]
xs = np.zeros_like(ys)
foo = mlab.plot3d(xs, ys, zs, color=(1, 0, 0), opacity=g_crossing_opacity)
mlab.figure(bgcolor=(1, 1, 1))
# plot the atoms as spheres
for atom in atoms:
mlab.points3d(
atom.x,
atom.y,
atom.z,
scale_factor=vdw_radii[atom.number] / 5.,
resolution=20,
# a tuple is required for the color
color=tuple(cpk_colors[atom.number]),
scale_mode='none')
# draw the unit cell - there are 8 corners, and 12 connections
a1, a2, a3 = atoms.get_cell()
origin = [0, 0, 0]
cell_matrix = [[origin, a1], [origin, a2], [origin, a3], [a1, a1 + a2],
[a1, a1 + a3], [a2, a2 + a1], [a2, a2 + a3], [a3, a1 + a3],
[a3, a2 + a3], [a1 + a2, a1 + a2 + a3], [a2 + a3, a1 + a2 + a3],
[a1 + a3, a1 + a3 + a2]]
for p1, p2 in cell_matrix:
mlab.plot3d(
[p1[0], p2[0]], # x-positions
[p1[1], p2[1]], # y-positions
[p1[2], p2[2]], # z-positions
tube_radius=0.02)
# Now plot the charge density
mlab.contour3d(x, y, z, cd, transparent=True)
# this view was empirically found by iteration
mlab.view(azimuth=-90, elevation=90, distance='auto')
mlab.savefig('images/co-centered-cd.png')
mlab.show()
def display(self,show_illum=True,dvol=50,cladding=None,scattering=True,spreading=True,bounds = None):
""" Display optrode in mayavi
"""
if not bounds:
if hasattr(self,"bounds"):
bounds = self.bounds
else:
bounds = [[-500,500],[-500,500],[-500,500]]
self.bounds = bounds
if cladding==None:
if hasattr(self,'mlab_cladding'):
cladding = True
else:
cladding = False
if hasattr(self,'mlab_tube'):
self.mlab_tube.parent.parent.remove()
if hasattr(self,'mlab_illum'):
self.mlab_illum.parent.parent.remove()
if hasattr(self,'mlab_cladding'):
self.mlab_cladding.parent.parent.remove()
try:
from enthought.mayavi import mlab
except:
from mayavi import mlab
cyl = mlab.plot3d(self.x,
self.y,
self.z,
name='optrode',
color=(.9,.9,.9))
self.mlab_tube = cyl.parent.parent
self.mlab_tube.filter.capping = True
self.mlab_tube.filter.number_of_sides = 20
self.mlab_tube.filter.radius = self.radius
if cladding:
from numpy import array
clad= mlab.plot3d(self.x,
self.y,
self.z-array([0.1,0.1]),
name='cladding',
color=(0.5,0.5,0.5),
opacity = 0.5)
self.mlab_cladding = clad.parent.parent
self.mlab_cladding.filter.capping = True
self.mlab_cladding.filter.number_of_sides = 20
self.mlab_cladding.filter.radius = self.radius*2
self.mlab_cladding.children[0].children[0].actor.property.backface_culling = True
if show_illum:
from numpy import diff, mgrid,array,matrix,c_,cos,sin,arctan,ones
from numpy.linalg import norm
x = self.xyz[1,0]
y = self.xyz[1,1]
z = self.xyz[1,2]
X,Y,Z = mgrid[x+bounds[0][0]:x+bounds[0][1]:dvol*1j,
y+bounds[1][0]:y+bounds[1][1]:dvol*1j,
z+bounds[2][0]:z+bounds[2][1]:dvol*1j]
Tx = self.find_illumination(X,Y,Z,spreading,scattering)
self.mlab_illum = mlab.contour3d(X,Y,Z,Tx,
#opacity=0.8,
transparent=True,
vmin=0.001,
vmax=0.1,
contours=[t for t in [0.1,0.01,0.001]])
self.mlab_illum.parent.scalar_lut_manager.use_default_range = False
self.mlab_illum.parent.scalar_lut_manager.data_range = array([ 0.001, 0.1 ])
self.mlab_illum.parent.scalar_lut_manager.lut.scale='log10'
#self.mlab_illum.actor.property.backface_culling = True
self.mlab_illum.actor.property.frontface_culling = True
xshft=bndx/2.0-sum(xs)/len(xs)
yshft=bndy/2.0-sum(ys)/len(ys)
zshft=bndz/2.0-sum(zs)/len(zs)
xs=map(lambda x:x+xshft,xs)
ys=map(lambda y:y+yshft,ys)
zs=map(lambda z:z+zshft,zs)
N=len(xs)
print "Lattice has %d atoms.\n" % N
#Plotting
if ploton==1:
from enthought.mayavi import mlab
fig=mlab.gcf()
mlab.points3d(xs,ys,zs)
mlab.plot3d([0,bndx],[0,0],[0,0])
mlab.plot3d([0,0],[0,bndy],[0,0])
mlab.plot3d([0,0],[0,0],[0,bndz])
mlab.show()
#Ask before overwriting a file.
try:
val=raw_input("Continue? (y/n): ")
if not(val in ('y','Y')):
print "Okey dokey, stopping."
exit(0)
ofil=open(sys.argv[8],"w")
except ValueError:
exit(0)
#Dump the atomic coordinates etc to the file, make it readable for lammps
def draw_camera (R, t, scale = 0.1, opt = 'fancy', figure = None):
""" Draw a camera in world coords according to its pose
If R = id and t = [0 0 0]' show a camera in the world
center pointing towards the +Z axis.
Usage: draw_camera (R, t, scale, opt);
Input:
R - 3x3 Rotation matrix or 3x1 Rodrigues rotation vector (camera to
world rotation)
t - 3x1 Translation vector (camera to world translation)
scale - (optional) scaling parameter, in meters. Default: 0.1.
opt - Visualization option: (default: 'pyr')
'pyr' shows an inverted pyramid in the camera direction
'axis' shows the 3 axis (red for X, green for Y, blue for Z)
"""
# Generate figure if none given
if figure == None:
figure = mlab.figure()
# Enforce 2d column vector of translation
t = np.atleast_2d(t.ravel()).T
# Five points that define the pyramid
# Careful, because np.c_ somehow transposes this matrix, this is 3x6.
pyr_points = scale * np.c_[[ 0, 0, 0], \
[-0.4, -0.4, 1], \
[ 0.4, -0.4, 1], \
[ 0.4, 0.4, 1], \
[-0.4, 0.4, 1], \
[np.nan, np.nan, np.nan]]
pyr_tri_side = np.c_[[0, 1, 2], \
[0, 2, 3], \
[0, 3, 4], \
[0, 1, 4]].T
pyr_tri_front = np.c_[[1, 2, 3], \
[1, 3, 4]].T
# Four points that define the axis in 3-space
axis_points = scale * np.c_[[0, 0, 0], \
[1, 0, 0], \
[0, 1, 0], \
[0, 0, 1]]
# Order in which to draw the points, so that it looks like 3 axis
axis_idx_x = np.r_[1, 2] - 1
axis_idx_y = np.r_[1, 3] - 1
axis_idx_z = np.r_[1, 4] - 1
# Some color constants
RED = (1,0,0)
GREEN = (0,1,0)
BLUE = (0,0,1)
if opt == 'pyr':
# Rotate pyramid and plot it
tx_pyr = np.dot(R, pyr_points) + \
np.tile( t, (1, pyr_points.shape[1]) )
mlab.triangular_mesh(tx_pyr[0, :], \
tx_pyr[1, :], \
tx_pyr[2, :], \
pyr_tri_side, \
color = BLUE, \
figure = figure)
mlab.draw()
mlab.triangular_mesh(tx_pyr[0, :], \
tx_pyr[1, :], \
tx_pyr[2, :], \
pyr_tri_front, \
color = GREEN, \
figure = figure)
elif opt == 'axis':
# Rotate the 3 axis and plot them
tx_axis = np.dot(R, axis_points) + \
np.tile( t, (1, axis_points.shape[1]) )
mlab.plot3d(tx_axis[0, axis_idx_x], \
tx_axis[1, axis_idx_x], \
tx_axis[2, axis_idx_x], \
color = RED, \
tube_radius = .003, \
figure = figure)
mlab.plot3d(tx_axis[0, axis_idx_y], \
tx_axis[1, axis_idx_y], \
tx_axis[2, axis_idx_y], \
color = GREEN, \
tube_radius = .003, \
figure = figure)
mlab.plot3d(tx_axis[0, axis_idx_z], \
tx_axis[1, axis_idx_z], \
tx_axis[2, axis_idx_z], \
color = BLUE, \
tube_radius = .003, \
figure = figure)
elif opt == 'fancy':
# Rotate pyramid and plot it
tx_pyr = np.dot(R, pyr_points) + \
np.tile( t, (1, pyr_points.shape[1]) )
mlab.triangular_mesh(tx_pyr[0, :], \
tx_pyr[1, :], \
tx_pyr[2, :], \
pyr_tri_side, \
color = BLUE, \
figure = figure)
mlab.triangular_mesh(tx_pyr[0, :], \
tx_pyr[1, :], \
tx_pyr[2, :], \
pyr_tri_front, \
color = GREEN, \
figure = figure)
# Rotate the 3 axis and plot them
tx_axis = np.dot(R, axis_points) + \
np.tile( t, (1, axis_points.shape[1]) )
mlab.plot3d(tx_axis[0, axis_idx_x], \
tx_axis[1, axis_idx_x], \
tx_axis[2, axis_idx_x], \
color = RED, \
tube_radius = .003, \
figure = figure)
mlab.plot3d(tx_axis[0, axis_idx_y], \
tx_axis[1, axis_idx_y], \
tx_axis[2, axis_idx_y], \
color = GREEN, \
tube_radius = .003, \
figure = figure)
else:
print('Unrecognized option');
def __init__(self, hs):
"""Plot the half-sarcomere"""
self.hs = hs
# Trigger an update of location data
self.update_locs()
# Do initial plotting of the thick and thin fils
self.thick_tubes = []
for x, yz, s in zip(self.thick_xlocs, self.thick_yzlocs, self.thick_s):
y = np.repeat(yz[0], len(x))
z = np.repeat(yz[1], len(x))
self.thick_tubes.append(
mlab.plot3d(x,
y,
z,
s,
tube_radius=fil_rad,
tube_sides=fil_seg,
colormap=fil_color,
vmin=fil_lims[0],
vmax=fil_lims[1]))
self.thin_tubes = []
for x, yz, s in zip(self.thin_xlocs, self.thin_yzlocs, self.thin_s):
y = np.repeat(yz[0], len(x))
z = np.repeat(yz[1], len(x))
self.thin_tubes.append(
mlab.plot3d(x,
y,
z,
s,
tube_radius=0.6 * fil_rad,
tube_sides=fil_seg,
colormap=fil_color,
vmin=fil_lims[0],
vmax=fil_lims[1]))
# Plot the total force at the end of each filament
self.update_ends()
self.thick_end_cube = []
for yz, s in zip(self.thick_yzlocs, self.thick_end):
x = [0]
y = [yz[0]]
z = [yz[1]]
s = [s]
self.thick_end_cube.append(
mlab.points3d(x,
y,
z,
s,
mode='cube',
scale_mode='none',
scale_factor=1.5 * fil_rad,
colormap=fil_color,
vmin=-50,
vmax=50))
self.thin_end_cube = []
for yz, s in zip(self.thin_yzlocs, self.thin_end):
x = [self.z_line]
y = [yz[0]]
z = [yz[1]]
s = [s]
self.thin_end_cube.append(
mlab.points3d(x,
y,
z,
s,
mode='cube',
scale_mode='none',
scale_factor=1.5 * fil_rad,
colormap=fil_color,
vmin=-50,
vmax=50))
# Plot the cross-bridges
self.update_bound()
self.cross_bridges = []
for fil, x, yz in zip(self.bound, self.thick_xlocs, self.thick_yzlocs):
y = [yz[0]]
z = [yz[1]]
for face in fil:
# TODO: THIS IS WHERE I LEFT OFF ON IMPLEMENTING
# CROSS-BRIDGE PLOTTING. THIS IS A BIT OF A STICKY FELLOW IN
# THAT IT REQUIRES THE CROSS-BRIDGES BE EITHER ALL SHOWN
# (WHICH IS VISUALLY CLUTTERED) OR DELETED AS NEEDED WITH
# EACH REDISPLAY (WHICH IS COMPLICATED). THE WAY TO GO IS
# PROBABLY TO DELETE ALL WITH EACH REDISPLAY AND THEN PLOT
# ALL CROSS-BRIDGES ANEW EACH TIME.
pass
Copyright (c) 2009 Emanuele Olivetti <emanuele_AT_relativita.com>
This library is free software; you can redistribute it and/or modify
it either under the terms of the GNU General Public License version 3
as published by the Free Software Foundation.
"""
from enthought.mayavi import mlab
import streamlines
if __name__ == "__main__":
# filename = "DeterministicTractography/QBALLRecon/hardiO10.trk"
filename = "hardiO10.trk_cross_streamline_id_20.trk"
try:
s
except:
s = streamlines.Streamlines()
print
print "file:", filename
s.loadTrk(filename)
s.printHeaderTrk()
pass
for stream in s.streamline:
mlab.plot3d(stream[:, 0], stream[:, 1], stream[:, 2], tube_radius=0.2)
pass
mlab.show()
con_val = list()
for i, j in zip(ii, jj):
if linalg.norm(sens_loc[i] - sens_loc[j]) > min_dist:
con_nodes.append((i, j))
con_val.append(con[i, j])
con_val = np.array(con_val)
# Show the connections as tubes between sensors
vmax = np.max(con_val)
vmin = np.min(con_val)
for val, nodes in zip(con_val, con_nodes):
x1, y1, z1 = sens_loc[nodes[0]]
x2, y2, z2 = sens_loc[nodes[1]]
mlab.plot3d([x1, x2], [y1, y2], [z1, z2], [val, val],
vmin=vmin,
vmax=vmax,
tube_radius=0.002)
mlab.scalarbar(title='Phase Lag Index (PLI)', nb_labels=4)
# Add the sensor names for the connections shown
nodes_shown = list(set([n[0] for n in con_nodes] + [n[1] for n in con_nodes]))
for node in nodes_shown:
x, y, z = sens_loc[node]
mlab.text3d(x,
y,
z,
raw.ch_names[picks[node]],
scale=0.005,
color=(0, 0, 0))
atoms=atoms)
calc.set_nbands(f=2)
calc.stop_if(calc.potential_energy is None)
x, y, z, lp = calc.get_local_potential()
x, y, z, cd = calc.get_charge_density()
mlab.figure(1, bgcolor=(1, 1, 1)) # make a white figure
# plot the atoms as spheres
for atom in atoms:
mlab.points3d(
atom.x,
atom.y,
atom.z,
scale_factor=vdw_radii[atom.number] / 5.,
resolution=20,
# a tuple is required for the color
color=tuple(cpk_colors[atom.number]),
scale_mode='none')
# plot the bonds. We want a line from C-Br, C-F, etc...
# We create a bond matrix showing which atoms are connected.
bond_matrix = [[0, 1], [0, 2], [0, 3], [0, 4]]
for a1, a2 in bond_matrix:
mlab.plot3d(
atoms.positions[[a1, a2], 0], # x-positions
atoms.positions[[a1, a2], 1], # y-positions
atoms.positions[[a1, a2], 2], # z-positions
[2, 2],
tube_radius=0.02,
colormap='Reds')
mlab.contour3d(x, y, z, lp)
mlab.savefig('images/halogen-ep.png')
mlab.show()
def sensor_connectivity_3d(raw, picks, con, idx, n_con=20, min_dist=0.05,
scale_factor=0.005, tube_radius=0.001):
""" Function to plot sensor connectivity showing strongest
connections(n_con) excluding sensors that are less than min_dist apart.
https://github.com/mne-tools/mne-python/blob/master/examples/connectivity/plot_sensor_connectivity.py
Parameters
----------
raw : Raw object
Instance of mne.io.Raw
picks : list
Picks to be included.
con : ndarray (n_channels, n_channels)
Connectivity matrix.
idx : list
List of indices of sensors of interest.
n_con : int
Number of connections of interest.
min_dist : float
Minimum distance between sensors allowed.
Note: Please modify scale factor and tube radius to appropriate sizes
if the plot looks scrambled.
"""
# Now, visualize the connectivity in 3D
try:
from enthought.mayavi import mlab
except:
from mayavi import mlab
mlab.figure(size=(600, 600), bgcolor=(0.5, 0.5, 0.5))
# Plot the sensor location
sens_loc = [raw.info['chs'][picks[i]]['loc'][:3] for i in idx]
sens_loc = np.array(sens_loc)
pts = mlab.points3d(sens_loc[:, 0], sens_loc[:, 1], sens_loc[:, 2],
color=(1, 1, 1), opacity=1, scale_factor=scale_factor)
# Get the strongest connections
threshold = np.sort(con, axis=None)[-n_con]
ii, jj = np.where(con >= threshold)
# Remove close connections
con_nodes = list()
con_val = list()
for i, j in zip(ii, jj):
if sci.linalg.norm(sens_loc[i] - sens_loc[j]) > min_dist:
con_nodes.append((i, j))
con_val.append(con[i, j])
con_val = np.array(con_val)
# Show the connections as tubes between sensors
vmax = np.max(con_val)
vmin = np.min(con_val)
for val, nodes in zip(con_val, con_nodes):
x1, y1, z1 = sens_loc[nodes[0]]
x2, y2, z2 = sens_loc[nodes[1]]
points = mlab.plot3d([x1, x2], [y1, y2], [z1, z2], [val, val],
vmin=vmin, vmax=vmax, tube_radius=tube_radius,
colormap='RdBu')
points.module_manager.scalar_lut_manager.reverse_lut = True
mlab.scalarbar(title='Phase Lag Index (PLI)', nb_labels=4)
# Add the sensor names for the connections shown
nodes_shown = list(set([n[0] for n in con_nodes] +
[n[1] for n in con_nodes]))
for node in nodes_shown:
x, y, z = sens_loc[node]
mlab.text3d(x, y, z, raw.ch_names[picks[node]], scale=0.005,
color=(0, 0, 0))
view = (-88.7, 40.8, 0.76, np.array([-3.9e-4, -8.5e-3, -1e-2]))
mlab.view(*view)
scale_factor=cube_size * 3,
scale_mode='scalar',
mode='cube',
line_width=lw * 0.75,
name='HI_clumps')
# Finally, trace the tidal tail in HI inside the cube, using a cyan cylinder.
tails_x = np.array([-80, -50, 0, 30, 50, 80, 80, 90, 70, 50])
tails_y = np.array([0, 0, -10, -15, -50, -70, -100, -110, -140, -160])
tails_z = np.array(
[6995, 7010, 7020, 7035, 7070, 7120, 7121, 7170, 7215, 7220])
rad = 4
mlab.plot3d(tails_x,
tails_y,
tails_z,
color=(0, 1, 1),
tube_radius=rad,
name='HCG91a_tail')
# --- !!! ---
# MAYAVI BUG & WORK-AROUND (#2 continued)
# Normally, the diagram would be complete at this point.
# However, for the x3d export, we still need to manually include the axes,
# axes labels, and some tick labels. This is tedious, but here we go.
#
# If you do not care about x3d export remove those lines, and set
#
#ax.visible = True
#
# --- !!! ---
color=(1, 1, 1),
scale_mode='none')
H2 = mlab.points3d(atoms_x[-1:],
atoms_y[-1:],
atoms_z[-1:],
scale_factor=2,
resolution=20,
color=(1, 1, 1),
scale_mode='none')
# The bounds between the atoms, we use the scalar information to give
# color
mlab.plot3d(atoms_x,
atoms_y,
atoms_z, [1, 2, 1],
tube_radius=0.4,
colormap='Reds')
# Display the electron localization function ###################################
# Load the data, we need to remove the first 8 lines and the '\n'
str = ' '.join(file('h2o-elf.cube').readlines()[9:])
data = np.fromstring(str, sep=' ')
data.shape = (40, 40, 40)
source = mlab.pipeline.scalar_field(data)
min = data.min()
max = data.max()
vol = mlab.pipeline.volume(source,
vmin=min + 0.65 * (max - min),
