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 topo(j_data, file_name=None, size=(600,600)):
u"""Creates the topological view of the molecule.
** Parameters **
j_data : dict
Data on the molecule, as deserialized from the scanlog format.
file_name : str, optional
Base name of the file in which to save the image.
size : tuple(int, int), optional
The size of the image to save.
** Returns **
figure : mayavi.core.scene.Scene
The MayaVi scene containing the visualization.
"""
figure, normal = _init_scene(j_data)
geom = np.array(j_data["results"]["geometry"]["elements_3D_coords_converged"]).reshape((-1,3))/A_to_a0
## Show labels and numbers ( = indices + 1 )
for i, atom in enumerate(j_data["molecule"]["atoms_Z"]):
P, label = geom[i], tab[atom][1]
mlab.text3d(P[0] - normal[0], P[1] - normal[1], P[2] - normal[2], label + str(i + 1), color=(0,0,0), scale=0.5, figure=figure)
if file_name is not None:
mlab.savefig("{}-TOPO.png".format(file_name), figure=figure, size=size)
return figure
def plot_3d(lats,lons,depth,vs,runlat,runlon,vtkname='rayleigh_c_u.vtk',
annotate_depth=False,coastline=False,annotate_lat=True,annotate_lon=True):
"""
plot 3d results using mayavi
"""
if coastline:
plot_canada_map()
sgrid = get_tvtk_grid(lats,lons,depth,vs)
d = mlab.pipeline.add_dataset(sgrid)
#sf = mlab.pipeline.surface(d)
#gx = mlab.pipeline.grid_plane(d)
#gx.grid_plane.axis = 'x'
gy = mlab.pipeline.grid_plane(d)
gy.grid_plane.axis = 'x'
gy.module_manager.scalar_lut_manager.show_scalar_bar = True
gy.module_manager.scalar_lut_manager.lut_mode = 'jet'
gy.module_manager.scalar_lut_manager.data_range = np.array([ 2. , 4.8])
gy.module_manager.scalar_lut_manager.scalar_bar_representation.maximum_size = np.array([100000, 100000])
gy.module_manager.scalar_lut_manager.scalar_bar_representation.minimum_size = np.array([1, 1])
gy.module_manager.scalar_lut_manager.scalar_bar_representation.position2 = np.array([ 0.08796009, 0.56264591])
gy.module_manager.scalar_lut_manager.scalar_bar_representation.position = np.array([ 0.03396896, 0.39182879])
gy.actor.mapper.progress = 1.0
gy.actor.mapper.scalar_range = np.array([ 0., 1.])
gy.actor.mapper.scalar_visibility = True
gy.actor.property.representation = 'surface'
gy.grid_plane.position = 6
#gz = mlab.pipeline.grid_plane(d)
#gz.grid_plane.axis = 'z'
if annotate_lat:
for lat in runlat:
x,y,z = convert_pt(lat,-58.,10.)
txt = mlab.text3d(x,y,z,'%d'%(lat),color=(0,0,0),line_width=10.0)
txt.scale = [20,20,20]
if annotate_lon:
for lon in runlon[1::]:
x,y,z = convert_pt(49.,lon,10.)
txt = mlab.text3d(x,y,z,'%d'%(lon),color=(0,0,0),line_width=10.0)
txt.scale = [20,20,20]
if annotate_depth:
for dp in [-10,-40,-80,-120]:
x,y,z = convert_pt(49,-68.,dp)
txt = mlab.text3d(x,y,z,'%d km'%(dp),color=(0,0,0),line_width=10.0)
txt.scale = [20,20,20]
### Include 3D screenshot in matplotlib
#arr = mlab.screenshot()
#import pylab as pl
#pl.imshow(arr)
#pl.show()
mlab.text(0.76,0.86,'49N',width=0.1)
mlab.show()
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 test_text3d(self):
""" Test the text3d module.
"""
data = np.random.random((3, 3, 3))
src = mlab.pipeline.scalar_field(data)
t = mlab.text3d(0, 0, 0, 'foo', opacity=0.5, scale=2,
orient_to_camera=False, color=(0, 0, 0),
orientation=(90, 0, 0))
def text3d( x, y, z, s, bcolor=None, bwidth=0.5, bn=16, **kwargs ):
"""
Mayavi text3d command augmented with poor man's bold.
"""
from enthought.mayavi import mlab
h = []
if bcolor is not None:
args = kwargs.copy()
args['color'] = bcolor
for i in range( bn ):
phi = 2.0 * np.pi * i / bn
x_ = x + bwidth * np.cos( phi )
y_ = y + bwidth * np.sin( phi )
h += [ mlab.text3d( x_, y_, z, s, **args ) ]
h[-1].actor.property.lighting = False
h += [ mlab.text3d( x_, y_, z, s, **kwargs ) ]
h[-1].actor.property.lighting = False
return h
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 test_text3d(self):
""" Test the text3d module.
"""
data = np.random.random((3, 3, 3))
src = mlab.pipeline.scalar_field(data)
t = mlab.text3d(0,
0,
0,
'foo',
opacity=0.5,
scale=2,
orient_to_camera=False,
color=(0, 0, 0),
orientation=(90, 0, 0))
def test_den():
# P = np.random.random((10,10))
# P=np.load(r'c:\maxdenP.np.npy')
# myfilestr=r'c:\structfactors_density.dat'
# x,y,z=np.loadtxt(myfilestr).T
# P=z.reshape((101,101))
# print P.shape
fig = mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))
x, y, z = np.array([[1, 0, 1], [0, -1, 1], [-1, 0, -1], [0, 1, -1], [1, 1, 0], [-1, -1, 0]], "float64").T
# view along z-axis
pts_as = mlab.points3d(x, y, z, color=(1, 0, 0), colormap="gist_rainbow", figure=fig, scale_factor=0.3)
# x,y,z=gen_fe()
print "x", x, x.shape
print "y", y, y.shape
print "z", z, z.shape
x, y, z = np.array([[0, 1, 1], [0, -1, -1]], "float64").T
pts_FE = mlab.points3d(x, y, z, color=(0, 1, 0), colormap="gist_rainbow", figure=fig, scale_factor=0.3)
mlab.text3d(0, 1, 1, "011", scale=0.2, color=(1, 1, 0))
mlab.text3d(1, 1, 0, "110", scale=0.2)
mlab.text3d(-1, -1, 0, "-1-10", scale=0.2)
mlab.text3d(1, 0, 1, "101", scale=0.2)
mlab.text3d(0, -1, 1, "0-11", scale=0.2)
mlab.text3d(-1, 0, -1, "-10-1", scale=0.2)
mlab.text3d(0, 1, -1, "01-1", scale=0.2)
# pts_fe=mlab.points3d(x,y,z-.125,color=(0,1,0),colormap='gist_rainbow',figure=fig,scale_factor=.02)
# x,y,z=gen_sr()
# print 'x',x,x.shape
# pts_sr=mlab.points3d(x,y,z-.125,color=(0,0,1),colormap='gist_rainbow',figure=fig)
# outline=mlab.outline(figure=fig,extent=[0,1,0,1,-1,0])
outline = mlab.outline(figure=fig)
mlab.orientation_axes(figure=fig, xlabel="a", ylabel="b", zlabel="c")
# print 'shape',P.shape
# P = P[:,:, np.newaxis]
# print 'after',P.shape
# src = mlab.pipeline.scalar_field(P)
# src = mlab.pipeline.array2d_source(P)
# surf = mlab.pipeline.surface(src,figure=fig,extent=[0,1,0,1,-1,0],name='surf2',opacity=0.4)
# surf.transform.GetTransform().RotateX(90)
print "done"
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)
def test_den():
#P = np.random.random((10,10))
#P=np.load(r'c:\maxdenP.np.npy')
#myfilestr=r'c:\structfactors_density.dat'
#x,y,z=np.loadtxt(myfilestr).T
#P=z.reshape((101,101))
#print P.shape
fig = mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))
x, y, z = np.array([[1, 0, 1], [0, -1, 1], [-1, 0, -1], [0, 1, -1],
[1, 1, 0], [-1, -1, 0]], 'float64').T
#view along z-axis
pts_as = mlab.points3d(x,
y,
z,
color=(1, 0, 0),
colormap='gist_rainbow',
figure=fig,
scale_factor=.3)
#x,y,z=gen_fe()
print 'x', x, x.shape
print 'y', y, y.shape
print 'z', z, z.shape
x, y, z = np.array([
[0, 1, 1],
[0, -1, -1],
], 'float64').T
pts_FE = mlab.points3d(x,
y,
z,
color=(0, 1, 0),
colormap='gist_rainbow',
figure=fig,
scale_factor=.3)
mlab.text3d(0, 1, 1, '011', scale=.2, color=(1, 1, 0))
mlab.text3d(1, 1, 0, '110', scale=.2)
mlab.text3d(-1, -1, 0, '-1-10', scale=.2)
mlab.text3d(1, 0, 1, '101', scale=.2)
mlab.text3d(0, -1, 1, '0-11', scale=.2)
mlab.text3d(-1, 0, -1, '-10-1', scale=.2)
mlab.text3d(0, 1, -1, '01-1', scale=.2)
#pts_fe=mlab.points3d(x,y,z-.125,color=(0,1,0),colormap='gist_rainbow',figure=fig,scale_factor=.02)
#x,y,z=gen_sr()
#print 'x',x,x.shape
#pts_sr=mlab.points3d(x,y,z-.125,color=(0,0,1),colormap='gist_rainbow',figure=fig)
#outline=mlab.outline(figure=fig,extent=[0,1,0,1,-1,0])
outline = mlab.outline(figure=fig)
mlab.orientation_axes(figure=fig, xlabel='a', ylabel='b', zlabel='c')
#print 'shape',P.shape
#P = P[:,:, np.newaxis]
#print 'after',P.shape
#src = mlab.pipeline.scalar_field(P)
#src = mlab.pipeline.array2d_source(P)
#surf = mlab.pipeline.surface(src,figure=fig,extent=[0,1,0,1,-1,0],name='surf2',opacity=0.4)
#surf.transform.GetTransform().RotateX(90)
print 'done'
# 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]],
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)
def plot_3d_spectrum_mayavi(fs, fignum=None, vmin=None, vmax=None,
pop_ids=None):
"""
Logarithmic heatmap of single 3d FS.
This method relies on MayaVi2's mlab interface. See http://code.enthought.com/projects/mayavi/docs/development/html/mayavi/mlab.html . To edit plot
properties, click leftmost icon in the toolbar.
If you get an ImportError upon calling this function, it is likely that you
don't have mayavi installed.
fs: FS to plot
vmin: Values in fs below vmin are masked in plot.
vmax: Values in fs above vmax saturate the color spectrum.
fignum: Figure number to plot into. If None, a new figure will be created.
Note that these are MayaVi figures, which are separate from
matplotlib figures.
pop_ids: If not None, override pop_ids stored in Spectrum.
"""
from enthought.mayavi import mlab
fig = mlab.figure(fignum, bgcolor=(1,1,1))
mlab.clf(fig)
if vmin is None:
vmin = fs.min()
if vmax is None:
vmax = fs.max()
# Which entries should I plot?
toplot = numpy.logical_not(fs.mask)
toplot = numpy.logical_and(toplot, fs.data >= vmin)
# For the color mapping
normalized = (numpy.log(fs)-numpy.log(vmin))\
/(numpy.log(vmax)-numpy.log(vmin))
normalized = numpy.minimum(normalized, 1)
xs,ys,zs = numpy.indices(fs.shape)
flat_xs = xs.flatten()
flat_ys = ys.flatten()
flat_zs = zs.flatten()
flat_toplot = toplot.flatten()
mlab.barchart(flat_xs[flat_toplot], flat_ys[flat_toplot],
flat_zs[flat_toplot], normalized.flatten()[flat_toplot],
colormap='hsv', scale_mode='none', lateral_scale=1,
figure=fig)
if pop_ids is None:
if fs.pop_ids is not None:
pop_ids = fs.pop_ids
else:
pop_ids = ['pop0','pop1','pop2']
a = mlab.axes(xlabel=pop_ids[0],ylabel=pop_ids[1],zlabel=pop_ids[2],
figure=fig, color=(0,0,0))
a.axes.label_format = ""
a.title_text_property.color = (0,0,0)
mlab.text3d(fs.sample_sizes[0],fs.sample_sizes[1],fs.sample_sizes[2]+1,
'(%i,%i,%i)'%tuple(fs.sample_sizes), scale=0.75, figure=fig,
color=(0,0,0))
mlab.view(azimuth=-40, elevation=65, distance='auto', focalpoint='auto')
mlab.show()
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]], 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)
def initialize_rope(label_img, xyz,bgr, plotting=False):
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)))
if plotting:
cv2.imshow('bgr',bgr.copy())
cv2.imshow('labels',label_img.astype('uint8')*50)
cv2.waitKey(5)
start_end = (uv_start, uv_end) = get_cc_centers(label_img==2,2)
skel = mip.skeletonize(rope_mask)
skel = skeletons.removeSpurs1(skel,10)
if plotting: cv2.imshow('skel',skel*100)
paths2d = skeletons.get_paths_2d(skel)
endpts2d = []
for path in paths2d:
endpts2d.append(path[0])
endpts2d.append(path[-1])
i_start = nearest_neighbor(uv_start, endpts2d)
i_end = nearest_neighbor(uv_end, endpts2d)
print "start,end",i_start,i_end
paths3d = [to3d(path,xyz) for path in paths2d]
if plotting:
starts_ends = []
for path in paths2d:
starts_ends.append(path[0])
starts_ends.append(path[-1])
plot_colorskel(paths2d,bgr,start_end,starts_ends)
C = skeletons.make_cost_matrix(paths3d)
PG = skeletons.make_path_graph(C, [len(path) for path in paths3d])
(_,nodes) = skeletons.longest_path_between(PG,i_start, set([]), i_end)
if nodes is None: raise Exception("longest path failed")
print nodes
total_path = []
for node in nodes[::2]:
if node%2 == 0:
total_path.extend(paths3d[node//2])
else:
total_path.extend(paths3d[node//2][::-1])
total_path = np.array(total_path)
xyzs_unif = curves.unif_resample(total_path,N_SEGS+1,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(paths3d):
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 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)
## For surface rendering
field=mlab.contour3d(sregion,colormap='gist_ncar') # Generate a scalar field
#field.contour.maximum_contour = 13000
field.contour.minimum_contour = 2000
field.actor.property.opacity = 0.3
mlab.outline()
mlab.xlabel('RA(J2000)')
mlab.ylabel('DEC(J2000)')
mlab.zlabel('Velocity')
#mlab.axes()
#mlab.colorbar()
## Add 3-d text in the cube
mlab.text3d(30,80,23,'Outflows in L1551 IRS5',scale=2.)
#mlab.text(45,45,"outflows!",width=0.2,z=20)
## Draw arrows to show outflows etc.
vector1=mlab.quiver3d(55,55,30,1,1,0.2,mode='arrow',scale_factor=10.0,color=(0.0,0.0,1.0))
vector2=mlab.quiver3d(55,40,30,1,-1.2,0.2,mode='arrow',scale_factor=10.0,color=(0.0,0.0,1.0))
vector3=mlab.quiver3d(45,50,30,-0.8,1,0.2,mode='2dthick_arrow',scale_factor=10.0,color=(0.0,0.0,1.0))
vector4=mlab.quiver3d(49,48,30,0.2,0,1,mode='2darrow',scale_factor=8.0,color=(0.0,0.0,1.0))
vector5=mlab.quiver3d(47,58,20,-1,1.8,-0.2,mode='arrow',scale_factor=10.0,color=(1.0,0.0,0.0))
vector6=mlab.quiver3d(45,45,20,-1.2,-1,-0.2,mode='arrow',scale_factor=10.0,color=(1.0,0.0,0.0))
vector7=mlab.quiver3d(55,45,20,1,-1.4,-0.2,mode='2dthick_arrow',scale_factor=10.0,color=(1.0,0.0,0.0))
vector8=mlab.quiver3d(48,48,20,-0.2,0,-1,mode='2darrow',scale_factor=8.0,color=(1.0,0.0,0.0))
# Insert the continuum image as a background
img=pyfits.open('L1551.cont.fits') # Read the fits data
def atom_chooser(file_name = "C:/LanthanideSMMs/GdCu SMMs/7.cif", atoms = [], fract_coords = []):
"""display a list of atoms using a list of fractional coordinate arrays
unit cell parameters will be extracted from the CIF-file
inputs:
file_name: a string with front slashes for file tree delimiters that directs to the location of the CIF file
atoms: string list of atoms index in same manner as fract_coords
fract_coords: a list of numpy arrays containing the fractional coordinates
output:
chosen_atom_indices
"""
import CifFile #existing library for reading in CIF format
import os #operating system
import re #regular expressions
import numpy as np #numeric python
from enthought.mayavi import mlab
#convert file_name to acceptable format
file_name = "file://" + file_name
file_name = os.path.normpath(file_name)
#read contents of CIF file
cf = CifFile.ReadCif(file_name)
sample_name = cf.dictionary.keys()[0]
#unit cell in meters (converted from Angstroms)
a = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_length_a']))) * 1e-10
b = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_length_b']))) * 1e-10
c = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_length_c']))) * 1e-10
#angles in degrees
alpha = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_angle_alpha'])))
beta = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_angle_beta'])))
gamma = float(re.sub('\)','',re.sub('\(','',cf[sample_name]['_cell_angle_gamma'])))
#angles in radians
alpha_ = alpha*np.pi/180
beta_ = beta*np.pi/180
gamma_ = gamma*np.pi/180
#need to generate real space vectors...
a1 = np.array([a, 0, 0])
a2 = np.array([b*np.cos(gamma_), b*np.sin(gamma_)*np.sqrt(1-(np.cos(gamma_)*np.cos(beta_)-np.cos(alpha_))**2/(np.sin(gamma_)**2*np.sin(beta_)**2)), b*np.sin(gamma_)*(np.cos(gamma_)*np.cos(beta_)-np.cos(alpha_))/(np.sin(gamma_)*np.sin(beta_))])
a3 = np.array([c*np.cos(beta_), 0, c*np.sin(beta_)])
a1_A = a1*1e10
a2_A = a2*1e10
a3_A = a3*1e10
#start the mayavi pipeline
# mlab.figure(fgcolor = (0,0,0), bgcolor = (1,1,1))
#display the bounding box of the unit cell
mlab.plot3d([0,a1_A[0]], [0,a1_A[1]], [0,a1_A[2]], line_width = 4)
mlab.plot3d([0,a2_A[0]], [0,a2_A[1]], [0,a2_A[2]], line_width = 4)
mlab.plot3d([a1_A[0], a1_A[0]+a2_A[0]], [a1_A[1], a1_A[1]+a2_A[1]], [a1_A[2], a1_A[2]+a2_A[2]], line_width = 4)
mlab.plot3d([a2_A[0], a1_A[0]+a2_A[0]], [a2_A[1], a1_A[1]+a2_A[1]], [a2_A[2], a1_A[2]+a2_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+0,a1_A[0]+a3_A[0]], [a3_A[1]+0,a1_A[1]+a3_A[1]], [a3_A[2]+0,a1_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+0,a2_A[0]+a3_A[0]], [a3_A[1]+0,a2_A[1]+a3_A[1]], [a3_A[2]+0,a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+a1_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a1_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a1_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+a2_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a2_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a2_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([0,a3_A[0]], [0,a3_A[1]], [0,a3_A[2]], line_width = 4)
mlab.plot3d([a1_A[0],a1_A[0]+a3_A[0]], [a1_A[1],a1_A[1]+a3_A[1]], [a1_A[2],a1_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a2_A[0],a2_A[0]+a3_A[0]], [a2_A[1],a2_A[1]+a3_A[1]], [a2_A[2],a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a2_A[0]+a1_A[0],a2_A[0]+a1_A[0]+a3_A[0]], [a2_A[1]+a1_A[1],a2_A[1]+a1_A[1]+a3_A[1]], [a2_A[2]+a1_A[2],a2_A[2]+a1_A[2]+a3_A[2]], line_width = 4)
r_coords = np.zeros(np.shape(fract_coords))
#take the fractional coordinates and generate real coordinates using the basis vectors
for i in range(np.size(atoms)):
r_coords[i] = a1_A * fract_coords[i][0] + a2_A * fract_coords[i][1] + a3_A * fract_coords[i][2]
for i in range(np.size(atoms)):
mlab.points3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.15, scale_factor = 1, color = (0,0,1), mode = 'cube')
mlab.text3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.35, scale = 0.5, text = atoms[i], color = (0,0,1))
# placeholder, delete when ready
chosen_atom_indices = []
mlab.show()
return chosen_atom_indices
def _nodepicked_callback(self, picker_obj, evt):
""" The callback function to determine what to do with a picked node """
from enthought.tvtk.api import tvtk
picker_obj = tvtk.to_tvtk(picker_obj)
picked = picker_obj.actors
# if picked actor are the glyphs
if self.nodes.actor.actor._vtk_obj in [o._vtk_obj for o in picked]:
pos = picker_obj.picked_positions[0]
dist_small = sys.maxint # the biggest integer
for i, points in enumerate(self.nodes.mlab_source.points):
# if picked positions ON a node
dist = np.linalg.norm((np.array(pos)- points))
# save the smallest distance node
if dist <= dist_small:
j = i
dist_small = dist
# get the id of the picked node
picknodeid = 'n' + str(j+1)
if self.nodepicker.show_info:
# generate a nice view for the information
nodedata = self.network.graph.node[picknodeid]
nodedata['dn_internal_id'] = picknodeid
deg = self.network.graph.degree(picknodeid)
from types import IntType, StringType
if type(deg) == IntType or type(deg) == StringType:
nodedata['degree'] = deg
# does this node has a self-loop?
if self.network.graph.has_edge(picknodeid, picknodeid):
nodedata['has_selfloop'] = True
else:
nodedata['has_selfloop'] = True
mynode = Node(nodedata=nodedata)
mynode.configure_traits()
elif self.nodepicker.toggle_selection:
# toggles the selection of a node
if self.datasourcemanager._srcobj.selected_nodes[j] == 1:
# unpick
self._select_nodes(selection_node_array = np.array([j]), \
first_is_selected = True, change_vector = False)
else:
# pick
self._select_nodes(selection_node_array = np.array([j]), activate = True, \
first_is_selected = True, change_vector = False)
elif self.nodepicker.select_node:
# toggle the picking state
if self.datasourcemanager._srcobj.selected_nodes[j] == 1:
# unpick
self._select_nodes(selection_node_array = np.array([j]), \
first_is_selected = True)
else:
# pick
self._select_nodes(selection_node_array = np.array([j]), activate = True, \
first_is_selected = True)
elif self.nodepicker.select_node2:
# get the neighbors of the picked node in an array
nbrs = self.datasourcemanager._srcobj.relabled_graph.neighbors(j)
# put the two lists together
cur_nodes = np.append(j, np.array(nbrs))
if self.datasourcemanager._srcobj.selected_nodes[j] == 1:
# unpick
self._select_nodes(selection_node_array = cur_nodes, \
first_is_selected = True)
else:
# pick
self._select_nodes(selection_node_array = cur_nodes, activate = True ,\
first_is_selected = True)
elif self.nodepicker.toggle_label:
if self.node_labels.has_key(j):
# we want to remove the node label from the pipeline first
# and then delete the entry in the dictionary
a = self.node_labels[j]
a.remove()
del self.node_labels[j]
else:
srcobj = self.datasourcemanager._srcobj
# create the text3d instance and add it to the dictionary
from enthought.mayavi import mlab
a = mlab.text3d(srcobj.positions[j,0], srcobj.positions[j,1], \
srcobj.positions[j,2], \
' ' + srcobj.labels[j], # white spaces are hackish
name = 'Node ' + srcobj.labels[j])
self.node_labels[j] = a
else:
logger.info("No valid pick operation.")
else: C.append('g')
fig = mlab.figure(1, bgcolor=(1, 1, 1), size=(500, 500))
fig_res = 100
S = [1 for n in xrange(len(X))]
mlab.points3d(X, Y, Z, resolution=fig_res)
for x,y,z,l in zip(X,Y,Z,labels):
if int(l)<8:
d = .17
mlab.text3d(x-d,y+d*1.25,z,l,line_width=3.5, scale=.15,color=(0,0,0), opacity=.9)
elif int(l) in [9,8]:
mlab.text3d(x-d/2,y+d*1.3,z,l,line_width=3.5, scale=.15,color=(0,0,0), opacity=.9)
else:
mlab.text3d(x-d/2,y+d*1.5,z,l,line_width=3.5, scale=.15,color=(0,0,0), opacity=.9)
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)]
pts,s = latlondep2geo(depth,lat,lon,vs)
sgrid = tvtk.StructuredGrid(dimensions=(9,13,81),points=pts)
sgrid.points = pts
sgrid.point_data.scalars = ravel(s.copy())
sgrid.point_data.scalars.name = 'scalars'
d = mlab.pipeline.add_dataset(sgrid)
sf = mlab.pipeline.surface(d)
gx = mlab.pipeline.grid_plane(d)
gx.grid_plane.axis = 'x'
gy = mlab.pipeline.grid_plane(d)
gy.grid_plane.axis = 'y'
gz = mlab.pipeline.grid_plane(d)
gz.grid_plane.axis = 'z'
for lat in runlat:
x,y,z = convert_pt(lat,-68.,10.)
txt = mlab.text3d(x,y,z,'%d'%(lat),color=(0,0,0),line_width=10.0)
txt.scale = [30,30,30]
for lon in runlon[1::]:
x,y,z = convert_pt(43.,lon,10.)
txt = mlab.text3d(x,y,z,'%d'%(lon),color=(0,0,0),line_width=10.0)
txt.scale = [30,30,30]
for dp in [-10,-40]:
x,y,z = convert_pt(43.,-68.5,dp)
txt = mlab.text3d(x,y,z,'%d km'%dp,color=(0,0,0),line_width=10.0)
txt.scale = [30,30,30]
#plot_canada_map()
mlab.show()
#arr = mlab.screenshot()
#import pylab as pl
#pl.imshow(arr)
def mshplot3D(mesh, scale=[1, 1, 1], elm_ids=[], labels=[0, 0, 0]):
''' Plot 3D meshes (and optionally label vertices, elements, and edges)
@param mesh A 3D mesh object
@param elm_ids Numpy array of elements that will be plotted. If left empty,
all elements are plotted.
@param labels (\c bool) List of three elements indicateting what should be
labeled. By default, nothing is labeled
labels[0]=True labels vertices
labels[1]=True labels elements
labels[2]=True labels edges
@param scale (\c int) List indicating the scaling of the vertices -- this
is used to scale the z-direction for a thin mesh. By
default, there is no scaling, i.e. scale = [1, 1, 1]
@see msh.mesh.Mesh3D
@author Matt Ueckermann
'''
if type(elm_ids) is not int:
if len(elm_ids) == 0:
elm_ids = np.arange(len(mesh.elm))
#Figure out how many vertices are in each element type
n_vert_in_type = [len(int_el_pqr(element=element, dim=3)[0]) \
for element in mesh.u_elm_type]
#Now create the TVTK data-structures for the 'cells' and 'offsets'
#cells give [#verts, vert1id, vert2id...,vertnid, #verts ...]
#offsets gives the id's in the cells list where #verts are listed. So above
# it is [0, n+1]
cells = np.array([], dtype=int)
offset = np.array([], dtype=int)
for elm, elm_type in zip(mesh.elm[elm_ids, :], mesh.elm_type[elm_ids]):
n_vt = n_vert_in_type[elm_type]
offset = np.append(offset, len(cells))
cells = np.append(cells, [n_vt] + elm[:n_vt].tolist())
#Also have to create a list of element-types -- to do that we have to
#convert from my numbering system to the TVTK numbering system
if mesh.dim == 3:
type_convert = np.array([tvtk.Tetra().cell_type,\
tvtk.Hexahedron().cell_type, tvtk.Wedge().cell_type])
elif mesh.dim == 2:
type_convert = np.array([tvtk.Triangle().cell_type,\
tvtk.Quad().cell_type])
cell_types = mesh.u_elm_type[mesh.elm_type[elm_ids]]
cell_types = type_convert[cell_types]
#To help visualize the mesh, we color it be the distance from the origin
if mesh.dim == 3:
x, y, z = mesh.vert[:].T
elif mesh.dim == 2:
x, y = mesh.vert[:, :2].T
z = np.zeros_like(x)
dist = np.sqrt(x**2 + y**2 + z**2)
#and we scale the x, y, z, coordinates according the the assigned 'scale'
points = np.column_stack((x * scale[0], y * scale[1], z * scale[2]))
#Now we create the data-structures
cell_array = tvtk.CellArray()
cell_array.set_cells(len(offset), cells)
ug = tvtk.UnstructuredGrid(points=points)
ug.set_cells(cell_types, offset, cell_array)
ug.point_data.scalars = dist.ravel()
ug.point_data.scalars.name = 'Distance from Origin'
#Next we set the new data-structures as a mayavi source
src = sources.vtk_data_source.VTKDataSource(data=ug)
#Extract the edges from the grid
edges = mlab.pipeline.extract_edges(src)
#Use a shorter name for the elements
elm = mesh.elm[elm_ids, :]
if any(labels) and len(elm) > 20:
string = "WARNING:\nAre you sure you want to label more than 20" + \
"elements in 3D? -- Labels are very memory" + \
"(apparently -- who would have guessed?) \nY(es) to proceed:"
answer = raw_input(string)
if answer.lower() not in ['yes', 'y', 'yes']:
labels = [0, 0, 0]
#Next add labels if desired:
maxdist = np.max(dist) / 2.
if labels[0]: #Vertex labels
for i in xrange(len(offset)):
for vnum in elm[i, :]:
if vnum >= 0:
mlab.text3d(points[vnum, 0], points[vnum, 1], \
points[vnum, 2], '%d' % vnum, \
color=(0, 0, 0), scale=0.02*maxdist)
if labels[1]: #element labels
for i in xrange(len(offset)):
if elm_ids[i] < 0:
elm_label = mesh.numelm + elm_ids[i]
else:
elm_label = elm_ids[i]
x = np.mean(points[elm[i, :cells[offset[i]]], 0])
y = np.mean(points[elm[i, :cells[offset[i]]], 1])
z = np.mean(points[elm[i, :cells[offset[i]]], 2])
mlab.text3d(x, y, z, '%d' % elm_label, color=(0.3, 0.3, 0.9), \
scale=0.03*maxdist)
if labels[2]: #edge labels
if len(elm_ids) < len(mesh.elm):
elm2ed = connect_elm2ed(mesh.elm2elm[:], mesh.ed2ed)
ed_ids = np.unique(elm2ed[elm_ids])
ed_ids = ed_ids[ed_ids >= 0]
ed2ed = mesh.ed2ed[ed_ids, :]
else:
ed_ids = np.arange(len(mesh.ed2ed))
ed2ed = mesh.ed2ed[:]
for i in xrange(len(ed2ed)):
n = 8
if ed2ed[i, -1] < 0:
n = 7
x = np.mean(points[ed2ed[i, 4:n], 0])
y = np.mean(points[ed2ed[i, 4:n], 1])
z = np.mean(points[ed2ed[i, 4:n], 2])
mlab.text3d(x, y, z, '%d' % ed_ids[i], color=(0.3, 0.9, 0.3), \
scale=0.03*maxdist)
#And plot them!
mlab.pipeline.surface(edges, opacity=0.4, line_width=2)
mlab.axes()
mlab.title('3D mesh', size=0.5, height=0.95)
#mlab.show()
return cells, offset, cell_types
#display the bounding box of the unit cell
mlab.plot3d([0,a1_A[0]], [0,a1_A[1]], [0,a1_A[2]], line_width = 4)
mlab.plot3d([0,a2_A[0]], [0,a2_A[1]], [0,a2_A[2]], line_width = 4)
mlab.plot3d([a1_A[0], a1_A[0]+a2_A[0]], [a1_A[1], a1_A[1]+a2_A[1]], [a1_A[2], a1_A[2]+a2_A[2]], line_width = 4)
mlab.plot3d([a2_A[0], a1_A[0]+a2_A[0]], [a2_A[1], a1_A[1]+a2_A[1]], [a2_A[2], a1_A[2]+a2_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+0,a1_A[0]+a3_A[0]], [a3_A[1]+0,a1_A[1]+a3_A[1]], [a3_A[2]+0,a1_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+0,a2_A[0]+a3_A[0]], [a3_A[1]+0,a2_A[1]+a3_A[1]], [a3_A[2]+0,a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+a1_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a1_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a1_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a3_A[0]+a2_A[0], a1_A[0]+a2_A[0]+a3_A[0]], [a3_A[1]+a2_A[1], a1_A[1]+a2_A[1]+a3_A[1]], [a3_A[2]+a2_A[2], a1_A[2]+a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([0,a3_A[0]], [0,a3_A[1]], [0,a3_A[2]], line_width = 4)
mlab.plot3d([a1_A[0],a1_A[0]+a3_A[0]], [a1_A[1],a1_A[1]+a3_A[1]], [a1_A[2],a1_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a2_A[0],a2_A[0]+a3_A[0]], [a2_A[1],a2_A[1]+a3_A[1]], [a2_A[2],a2_A[2]+a3_A[2]], line_width = 4)
mlab.plot3d([a2_A[0]+a1_A[0],a2_A[0]+a1_A[0]+a3_A[0]], [a2_A[1]+a1_A[1],a2_A[1]+a1_A[1]+a3_A[1]], [a2_A[2]+a1_A[2],a2_A[2]+a1_A[2]+a3_A[2]], line_width = 4)
r_coords = np.zeros(np.shape(fract_coords))
#take the fractional coordinates and generate real coordinates using the basis vectors
for i in range(np.size(atoms)):
r_coords[i] = a1_A * fract_coords[i][0] + a2_A * fract_coords[i][1] + a3_A * fract_coords[i][2]
for i in range(np.size(atoms)):
mlab.points3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.15, scale_factor = 1, color = color_dict[atoms[i]], mode = 'cube')
mlab.text3d(r_coords[i][0], r_coords[i][1], r_coords[i][2], opacity = 0.35, scale = 0.5, text = atoms[i]+', '+str(i), color = color_dict[atoms[i]])
# placeholder, delete when ready
chosen_atom_indices = []
mlab.show()
