def drawArrows(file1, file2, figure, bbox, index, descendant):
"""
Draw an 'arrow' from the cell 'index' to 'descendant'
'descendant' is assumed to be a 1D np.array with size 0, 1 or 2.
"""
#load the center of mass position
if descendant.size > 0 and descendant[0] != 0:
com1 = file1["features"][str(index[0])]["com"][:]
com2 = file2["features"][str(descendant[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(index), 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 descendant.size == 2:
com3 = file2["features"][str(descendant[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 mlab_show_test_dataset():
from mayavi import mlab
mlab.points3d(bs_i, bs_j, bs_k,mode="cube",color=(1,1,0))
# Left pre-central
mlab.points3d(left_pre_cen60_1_i, left_pre_cen60_1_j,
left_pre_cen60_1_k,mode="cube",color=(0,1,1),
opacity=0.5)
mlab.points3d(left_pre_cen60_2_i, left_pre_cen60_2_j,
left_pre_cen60_2_k,mode="cube",color=(1,0,1),
opacity=0.5)
# Right pre-central
mlab.points3d(right_pre_cen60_1_i, right_pre_cen60_1_j,
right_pre_cen60_1_k,mode="cube",color=(0,1,0),
opacity=0.5)
mlab.points3d(right_pre_cen60_2_i, right_pre_cen60_2_j,
right_pre_cen60_2_k,mode="cube",color=(0,0,1),
opacity=0.5)
tds1.render_tracks = True
tds1.dynamic_color_clusters = False
tds1.static_color = "red"
tds1.draw_tracks()
tds2.render_tracks = True
tds2.dynamic_color_clusters = False
tds2.static_color = "blue"
tds2.draw_tracks()
def main():
parser = argparse.ArgumentParser()
parser.add_argument("point_x", help="x coordinate of point.", type=int)
parser.add_argument("point_y", help="y coordinate of point.", type=int)
parser.add_argument("file", help="The bathymetry file.")
parser.add_argument("--halo", help="Size of halo.", type=int, default=50)
args = parser.parse_args()
f = nc.Dataset(args.file)
topo = f.variables['depth'][:]
# Calculate the extents of the topo array to show. We don't just add halo to (point_x, point_y)
# because it may run off the edge of the map.
north_ext = min(args.point_y + args.halo, topo.shape[0])
south_ext = max(args.point_y - args.halo, 0)
east_ext = min(args.point_x + args.halo, topo.shape[1])
west_ext = max(args.point_x - args.halo, 0)
width = east_ext - west_ext
height = north_ext - south_ext
# The origin of the figure in global coordinates.
origin_x = west_ext + width / 2
origin_y = south_ext + height / 2
point_y = args.point_y - origin_y
point_x = args.point_x - origin_x
mlab.surf(topo[south_ext:north_ext, west_ext:east_ext], warp_scale=0.005)
mlab.points3d([point_y], [point_x], [0], color=(1, 0, 0), scale_factor=1.0)
mlab.show()
def draw_artifact(artifact):
draw_artifact_ports(artifact)
center = artifact.pos
mlab.points3d(center[0],center[1],center[2],color=(1,1,1))
wire_box(center=artifact.pos)
func_points = np.array([func.pos for func in artifact.subfunctions])
mlab.points3d(func_points[:,0],
func_points[:,1],
func_points[:,2],
color=(0,0,1),scale_factor=0.4)
for port in artifact.in_ports:
fcn = port['parent']
src = fcn.pos
for output in fcn.outputs:
snk = output['flow'].sink().pos
print src,snk
line(src,snk,color=(1,0,0))
for port in artifact.out_ports:
fcn = port['parent']
snk = fcn.pos
for input in fcn.inputs:
src = input['flow'].source().pos
line(src,snk,color=(0,1,0))
def plot_predicted_labels(points, labels):
print '[plot_points] Plotting points!'
xs = np.array([int(point._x) for point in points])
ys = np.array([int(point._y) for point in points])
zs = np.array([int(point._z) for point in points])
mlab.points3d(xs, ys, zs, labels, scale_factor = .4, mode='cube')
mlab.show()
def standardView(figure, atoms, resolution=32, scale=1):
from mayavi import mlab
input_atoms = atoms.copy()
del input_atoms.constraints
cell_center = cellCenter(input_atoms)
numbers = input_atoms.get_atomic_numbers()
u_numbers = np.unique(numbers)
vmin, vmax = input_atoms.positions[:,2].min(), input_atoms.positions[:,2].max()
vmin += (vmax - vmin)/2
for number in u_numbers:
take = np.array(numbers == number, bool)
atoms = input_atoms[take]
points = atoms.positions
points = points.transpose()
color = tuple(jmol_colors[number])
radius = my_radii[number]
if np.isnan(radius):
radius = 3.00
radius *= 2*scale
mlab.points3d(points[0], points[1], points[2], color=color,
scale_factor=radius,
resolution=resolution,
)
def plot_pos_ori(pos, ori, color=(0., 0., 0.)):
mlab.points3d(pos[:, 0], pos[:, 1], pos[:, 2], scale_factor=0.005,
color=color)
mlab.quiver3d(pos[:, 0], pos[:, 1], pos[:, 2],
ori[:, 0], ori[:, 1], ori[:, 2],
scale_factor=0.03,
color=color)
def colorAndSizeEveryAtom(
figure,
positions,
sizes,
colors,
colormap='hot',
resolution=32,
vmin=None,
vmax=None,
):
from mayavi import mlab
for position, size, color in zip(positions, sizes, colors):
x, y, z = position
color = tuple(color)
mlab.points3d(np.array([x]),
np.array([y]),
np.array([z]),
color=color,
scale_mode='none',
scale_factor=np.array([size]),
resolution=resolution,
reset_zoom=False,
figure=figure,
)
return figure
def mayaView(
atoms,
figure=None,
colormap='hot',
resolution=32,
vmax=None,
depth=3,
):
from mayavi import mlab
numbers = atoms.get_atomic_numbers()
u_numbers = np.unique(numbers)
if vmax is None:
vmax = atoms.positions[:,2].max()
vmin = vmax - depth
my_colors = {1:(1,0,0), 6:(0.5,0.5,1)}
for number in u_numbers:
take = numbers == number
element_atoms = atoms[take]
points = element_atoms.positions
points = points.transpose()
radius = my_radii[number]*2
mlab.points3d(points[0], points[1], points[2], points[2], scale_mode='none', scale_factor=radius, colormap=colormap, vmin=vmin, vmax=vmax,
resolution=resolution,
reset_zoom=False,
figure=figure,
)
return figure
def make_figures(coor, fun, interp_results, error):
'''Produce MayaVi figures for interpolation results'''
mlab.figure()
vmax, vmin = np.max(fun.fine), np.min(fun.fine)
mlab.mesh(coor.fine.x, coor.fine.y, coor.fine.z,
scalars=fun.fine, vmax=vmax, vmin=vmin)
mlab.points3d(coor.coarse.x, coor.coarse.y, coor.coarse.z, fun.coarse,
scale_factor=0.1, scale_mode='none', vmax=vmax, vmin=vmin)
mlab.colorbar(title='Spherical Harmonic', orientation='vertical')
mlab.savefig('Figures/poorfunctionsphere.png')
# Figure showing results of rbf interpolation
mlab.figure()
mlab.mesh(coor.fine.x, coor.fine.y, coor.fine.z,
scalars=interp_results, vmax=vmax, vmin=vmin)
mlab.points3d(coor.coarse.x, coor.coarse.y, coor.coarse.z, fun.coarse,
scale_factor=0.1, scale_mode='none', vmax=vmax, vmin=vmin)
mlab.colorbar(title='Interpolation', orientation='vertical')
mlab.savefig('Figures/poorinterpsphere.png')
mlab.figure()
mlab.mesh(coor.fine.x, coor.fine.y, coor.fine.z,
scalars=error.errors, vmax=error.max, vmin=-error.max)
mlab.colorbar(title='Error', orientation='vertical')
# mlab.points3d(coor.coarse.x, coor.coarse.y, coor.coarse.z, scalars, scale_factor=0.1, scale_mode='none',vmax=vmax, vmin=vmin)
mlab.savefig('Figures/poorerrorsphere.png')
mlab.show()
def xmas_balls(connectivity, node_data=None, edge_data=False):
"""
Plots coloured balls at the region centres of connectivity, colour and
size is determined by a vector of length number of regions (node_data).
Optional: adds the connections between pair of nodes.
"""
centres = connectivity.centres
edges = numpy.array(numpy.nonzero(connectivity.weights))
edges = numpy.array([(start, stop) for (start, stop) in edges.T if start != stop])
if node_data is not None:
data_scale = 13.0 / node_data.max()
pts = mlab.points3d(centres[:, 0], centres[:, 1], centres[:, 2],
node_data, transparent=True,
scale_factor=data_scale,
colormap='Blues')
mlab.colorbar(orientation="vertical")
else:
#NOTE: the magic numbers are used to align region centers and surface representation.
#Do not ask ...
pts = mlab.points3d(centres[:, 0] * 1.13, centres[:, 1] * 1.13 + 15, centres[:, 2] - 25)
if edge_data:
pts.mlab_source.dataset.lines = edges
tube = mlab.pipeline.tube(pts, tube_radius=0.5)
mlab.pipeline.surface(tube, colormap='binary', opacity=0.142)
def render(self, marker_size=1, marker_face_colour=(1, 1, 1), **kwargs):
from mayavi import mlab
mlab.points3d(
self.points[:, 0], self.points[:, 1], self.points[:, 2],
figure=self.figure, scale_factor=marker_size,
color=marker_face_colour)
return self
def render(self, scale_factor=1.0, text_scale=1.0, **kwargs):
import mayavi.mlab as mlab
# disabling the rendering greatly speeds up this for loop
self.figure.scene.disable_render = True
positions = []
for label in self.lmark_group:
p = self.lmark_group[label]
for i, p in enumerate(p.points):
positions.append(p)
l = '%s_%d' % (label, i)
# TODO: This is due to a bug in mayavi that won't allow
# rendering text to an empty figure
mlab.points3d(p[0], p[1], p[2], scale_factor=scale_factor)
mlab.text3d(p[0], p[1], p[2], l, figure=self.figure,
scale=text_scale)
positions = np.array(positions)
os = np.zeros_like(positions)
os[:, 2] = 1
mlab.quiver3d(positions[:, 0], positions[:, 1], positions[:, 2],
os[:, 0], os[:, 1], os[:, 2], figure=self.figure)
self.figure.scene.disable_render = False
# Ensure everything fits inside the camera viewport
mlab.get_engine().current_scene.scene.reset_zoom()
return self
def make_figures(coor, fun, interp_results, error):
# Figure of harmoinc function on sphere in fine cordinates
# Points3d showing interpolation training points coloured to their value
mlab.figure()
vmax, vmin = np.max(fun.fine), np.min(fun.fine)
mlab.mesh(coor.fine.x, coor.fine.y, coor.fine.z,
scalars=fun.fine, vmax=vmax, vmin=vmin)
mlab.points3d(coor.coarse.x, coor.coarse.y, coor.coarse.z, fun.coarse,
scale_factor=0.1, scale_mode='none', vmax=vmax, vmin=vmin)
mlab.colorbar(title='Spherical Harmonic', orientation='vertical')
# mlab.savefig('interppointssphere.png')
# Figure showing results of rbf interpolation
mlab.figure()
mlab.mesh(coor.fine.x, coor.fine.y, coor.fine.z,
scalars=interp_results, vmax=vmax, vmin=vmin)
mlab.points3d(coor.coarse.x, coor.coarse.y, coor.coarse.z, fun.coarse,
scale_factor=0.1, scale_mode='none', vmax=vmax, vmin=vmin)
mlab.colorbar(title='Interpolation', orientation='vertical')
# mlab.savefig('interpsphere.png')
mlab.figure()
mlab.mesh(coor.fine.x, coor.fine.y, coor.fine.z,
scalars=error.errors, vmax=error.max, vmin=-error.max)
mlab.colorbar(title='Error', orientation='vertical')
# mlab.points3d(coor.coarse.x, coor.coarse.y, coor.coarse.z, scalars, scale_factor=0.1, scale_mode='none',vmax=vmax, vmin=vmin)
# mlab.savefig('interpsphere.png')
mlab.show()
def _sphere_viz_default(self):
# different between wx and qt
try:
color_tuple = self.sphere_color.toTuple()
except:
color_tuple = self.sphere_color
try:
pts = mlab.points3d(
self.sphere_coords[:,0],
self.sphere_coords[:,1],
self.sphere_coords[:,2],
mode='cube',
scale_factor=1,
figure = self.scene3d.mayavi_scene,
color = (color_tuple[0]/255.,
color_tuple[1]/255.,
color_tuple[2]/255.)
)
except:
pts = mlab.points3d(
self.sphere_coords[:,0],
self.sphere_coords[:,1],
self.sphere_coords[:,2],
mode='cube',
scale_factor=1,
figure = self.scene3d.mayavi_scene,
color = (1.,0.,0.)
)
return pts
def plot(self):
"""Plots the geometry using the package Mayavi."""
print('\nPlot the three-dimensional geometry ...')
from mayavi import mlab
x_init = self.gather_coordinate('x', position='initial')
y_init = self.gather_coordinate('y', position='initial')
z_init = self.gather_coordinate('z', position='initial')
x = self.gather_coordinate('x')
y = self.gather_coordinate('y')
z = self.gather_coordinate('z')
figure = mlab.figure('body', size=(600, 600))
figure.scene.disable_render = False
same = (numpy.allclose(x, x_init, rtol=1.0E-06) and
numpy.allclose(y, y_init, rtol=1.0E-06) and
numpy.allclose(z, z_init, rtol=1.0E-06))
if not same:
mlab.points3d(x_init, y_init, z_init, name='initial',
scale_factor=0.01, color=(0, 0, 1))
mlab.points3d(x, y, z, name='current',
scale_factor=0.01, color=(1, 0, 0))
mlab.axes()
mlab.orientation_axes()
mlab.outline()
figure.scene.disable_render = True
mlab.show()
def plot(ply):
'''
Plot vertices and triangles from a PlyData instance. Assumptions:
`ply' has a 'vertex' element with 'x', 'y', and 'z'
properties;
`ply' has a 'face' element with an integral list property
'vertex_indices', all of whose elements have length 3.
'''
vertex = ply['vertex']
(x, y, z) = (vertex[t] for t in ('x', 'y', 'z'))
mlab.points3d(x, y, z, color=(1, 1, 1), mode='point')
if 'face' in ply:
tri_idx = ply['face']['vertex_indices']
idx_dtype = tri_idx[0].dtype
triangles = numpy.fromiter(tri_idx, [('data', idx_dtype, (3,))],
count=len(tri_idx))['data']
mlab.triangular_mesh(x, y, z, triangles,
color=(1, 0, 0.4), opacity=0.5)
def _showSimple():
max_interpol_pts = 10
def interpolate_section(section):
sec_start = section.get_distal_npa4()
sec_end = section.get_proximal_npa4()
length = section.get_length()
rad = min(section.d_r, section.p_r)
n = min(max(int(lToRRatio * length / rad), 1), max_interpol_pts)
j_vec_steps = (sec_end - sec_start) / n
int_pts = [sec_start + k * j_vec_steps for k in range(0, n)]
return int_pts
lbs = []
for morph in self.morphs:
lb = SeqUtils.flatten(ListBuilderSectionVisitor(functor=interpolate_section, morph=morph) ())
lbs.extend(lb)
pts = numpy.array(lbs)
x = pts[:, 0]
y = pts[:, 1]
z = pts[:, 2]
s = pts[:, 3]
mlab.points3d(x, y, z, s, colormap=self.colormap, scale_factor=self.scale_factor)
mlab.outline()
def to_line_plot(self, size=(400, 350)):
"""
Return a mayavi figure instance with histone and linkers shown
"""
if maya_imported is True:
fig = mlab.figure(bgcolor=(1., 1., 1.), size=size)
if hasattr(self, "histones"):
# ax = fig.add_subplot(111, projection='3d')
histones = []
for histone in self.histones:
pos = np.array([bp.position for bp in histone.basepairs])
mlab.plot3d(pos[:, 0], pos[:, 1], pos[:, 2],
color=(1., .8, 0),
tube_radius = 11.5)
histones.append(histone.position)
histones = np.array(histones)
mlab.points3d(histones[:, 0], histones[:, 1], histones[:, 2],
color=(0, 0, 1.), opacity=0.4, scale_factor=70)
for linker in self.linkers:
pos = np.array([bp.position for bp in linker.basepairs])
mlab.plot3d(pos[:, 0], pos[:, 1], pos[:, 2],
color=(0, .8, 0),
tube_radius = 11.5)
else:
pos = np.array([bp.position for bp in self.basepairs])
mlab.plot3d(pos[:, 0], pos[:, 1], pos[:, 2], color=(1., 0, 0),
tube_radius = 11.5)
return fig
else:
print("MayaVi not imporrted, cannot produce this plot")
return None
def plot_element_values(n, nodes, values, resample_n=None,
node_tuples=None, show_nodes=False):
dims = len(nodes)
orig_nodes = nodes
orig_values = values
if resample_n is not None:
import modepy as mp
basis = mp.simplex_onb(dims, n)
fine_nodes = mp.equidistant_nodes(dims, resample_n)
values = np.dot(mp.resampling_matrix(basis, fine_nodes, nodes), values)
nodes = fine_nodes
n = resample_n
from pytools import generate_nonnegative_integer_tuples_summing_to_at_most \
as gnitstam
if dims == 1:
import matplotlib.pyplot as pt
pt.plot(nodes[0], values)
if show_nodes:
pt.plot(orig_nodes[0], orig_values, "x")
pt.show()
elif dims == 2:
import mayavi.mlab as mlab
mlab.triangular_mesh(
nodes[0], nodes[1], values, submesh(list(gnitstam(n, 2))))
if show_nodes:
mlab.points3d(orig_nodes[0], orig_nodes[1], orig_values,
scale_factor=0.05)
mlab.show()
else:
raise RuntimeError("unsupported dimensionality %d" % dims)
def show_sensors(pth, red=None, green=None):
"""
show sensors stored in a h5.
:param red: a list of labels to be shown in red
"""
if red is None:
red = []
if green is None:
green = []
mlab.figure()
sensors = h5py.File(pth)
d = numpy.array(sensors['locations'])
labels = list(sensors['labels'])
highlight = numpy.ones(d.shape[0])
for i, l in enumerate(labels):
if l in red:
highlight[i] = 5.0
elif l in green:
highlight[i] = 7.0
else:
highlight[i] = 2.0
mlab.points3d(d[:,0], d[:,1], d[:,2], highlight, scale_mode='none')
mlab.axes()
def plot_sr(self, mlab, lc=0, sr=0):
lc = self.lc_arr[lc]
gd = self.geo_data_dict
mlab.points3d(gd['X'], gd['Y'], gd['Z'], lc[:, sr],
# colormap = "YlOrBr",
mode="cube",
scale_factor=0.1)
def plot_static(color):
global max_p, min_p, ary_q, ary_p
# --- Set labels and ticks. ---
# ax.set_xlabel(r'$k$', fontsize=16)
# ax.set_ylabel(r'$\theta$', fontsize=16)
# ax.set_zlabel(r'$p_{\theta}$', fontsize=16)
# plt.title('Standard map')
# unit = 0.5
# q_tick = np.arange(0, 2.0+unit, unit)
# q_label = [r"$0$",r"$\frac{1}{2}\pi$",r"$\pi$",r"$\frac{3}{2}\pi$",r"$2\pi$"]
# ax.set_ylim(0, 2*pi)
# ax.set_yticks(q_tick*pi)
# ax.set_yticklabels(q_label, fontsize=18)
# ax.set_xlim(0, kmax)
# if max(ary_p[ntransient:]) > max_p:
# max_p = max(ary_p[ntransient:])
# if min(ary_p[ntransient:]) < min_p:
# min_p = min(ary_p[ntransient:])
# margin = (max_p-min_p)*0.05
# ax.set_zlim(min_p-margin, max_p+margin)
# --- Plot the standard map. ---
mlab.points3d([0.1, ] * (nmax - ntransient + 1),
ary_q[ntransient:], ary_p[ntransient:])
plt.show()
def drawgranules(gids, gcolor, *arg):
from params import granule_diam as diam
x = []
y = []
z = []
if self.colors:
s = []
if self.colors and not not_in_weights:
gids = gids.intersection(self.colors.keys())
for gid in gids:
if self.projected:
if scale_mode == 'scalar':
u, p = gpo(gid)[-2:]
if self.colors:
try:
s.append(self.colors[gid] / 20 * 20 * 0.9 + 10)
except KeyError:
continue
dep = -(100 - s[-1]) / 20 * params.grid_dim # depth from colors
p = [ dep * u[0] + p[0], dep * u[1] + p[1], dep * u[2] + p[2]]
diam = 0.42
elif scale_mode == 'none':
p = gpo(gid)[2]
else:
p = gpo(gid)[0]
if self.colors and not (self.projected and scale_mode == 'scalar'):
try:
color = self.colors[gid]
except KeyError:
color = 0
if gthreshold:
if color >= gthreshold:
s.append(1)
else:
s.append(0)
else:
s.append(color)
r = params.ranstream(gid, 0)
r.uniform(-params.grid_dim * .5, params.grid_dim * .5)
x.append(p[0] + r.repick())
y.append(p[1] + r.repick())
z.append(p[2] + r.repick())
#print 'drawn mitral:',len(x)
if self.colors:
if self.vmin != None and self.vmax != None:
return points3d(x, y, z, s, scale_factor=diam, vmin=self.vmin, vmax=self.vmax, scale_mode=scale_mode, colormap=gpalette)
return points3d(x, y, z, s, scale_factor=diam, scale_mode=scale_mode,colormap=gpalette)
return points3d(x, y, z, scale_factor=diam, color=gcolor,colormap=gpalette)
def plot_intersections(self,):
'''Plot the intersection points on the plasma surface
SRH: 12July2013
'''
from mayavi import mlab
mlab.points3d(self.intersection1[self.valid_channels,0], self.intersection1[self.valid_channels,1], self.intersection1[self.valid_channels,2],scale_factor=0.02, color=(0,0,1))
mlab.points3d(self.intersection2[self.valid_channels,0], self.intersection2[self.valid_channels,1], self.intersection2[self.valid_channels,2],scale_factor=0.02,color=(1,0,0))
def bigtest():
from jds_image_proc.clouds import voxel_downsample
from jds_image_proc.pcd_io import load_xyz
import mayavi.mlab as mlab
pts = load_xyz("/home/joschu/Data/scp/three_objs_ds.pcd")
#pts = voxel_downsample(xyz, .03, False)
mlab.clf()
mlab.points3d(pts[:,0], pts[:,1], pts[:,2], color = (1,1,1), scale_factor=.01)
clus = []
labels = decompose(pts, .025)
for i in xrange(labels.max()+1):
clu = np.flatnonzero(labels == i)
clus.append(clu)
for clu in sorted(clus, key=len, reverse=True):
if len(clu) < 10: break
dirs = ss.get_sphere_points(1)
sup_pd = np.dot(pts[clu,:], dirs.T)
best_d = sup_pd.max(axis=0)
print "max deficit",(sup_pd - best_d[None,:]).max(axis=1).min()
mlab.points3d(pts[clu,0], pts[clu,1], pts[clu,2], color = (rand(),rand(),rand()), scale_factor=.01)
raw_input()
def plot_matrix(connectmat_file, centers_file, threshold_pct=5, weight_edges=False,
node_scale_factor=2, edge_radius=.5, resolution=8, name_scale_factor=1,
names_file=None, node_indiv_colors=[], highlight_nodes=[], fliplr=False):
"""
Given a connectivity matrix and a (x,y,z) centers file for each region, plot the 3D network
"""
matrix = core.file_reader(connectmat_file)
nodes = core.file_reader(centers_file)
if names_file:
names = core.file_reader(names_file,1)
num_nodes = len(nodes)
edge_thresh_pct = threshold_pct / 100.0
matrix_flat = np.array(matrix).flatten()
edge_thresh = np.sort(matrix_flat)[len(matrix_flat)-int(len(matrix_flat)*edge_thresh_pct)]
matrix = core.file_reader(connectmat_file)
ma = np.array(matrix)
thresh = scipy.stats.scoreatpercentile(ma.ravel(),100-threshold_pct)
ma_thresh = ma*(ma > thresh)
if highlight_nodes:
nr = ma.shape[0]
subset_mat = np.zeros((nr, nr))
for i in highlight_nodes:
subset_mat[i,:] = 1
subset_mat[:,i] = 1
ma_thresh = ma_thresh * subset_mat
if fliplr:
new_nodes = []
for node in nodes:
new_nodes.append([45-node[0],node[1],node[2]]) # HACK
nodes = new_nodes
mlab.figure(bgcolor=(1, 1, 1), size=(400, 400))
for count,(x,y,z) in enumerate(nodes):
if node_indiv_colors:
mlab.points3d(x,y,z, color=colors[node_indiv_colors[count]], scale_factor=node_scale_factor, resolution=resolution)
else:
mlab.points3d(x,y,z, color=(0,1,0), scale_factor=node_scale_factor, resolution=resolution)
if names_file:
width = .025*name_scale_factor*len(names[count])
print width
print names[count]
mlab.text(x, y,names[count], z=z,width=.025*len(names[count]),color=(0,0,0))
for i in range(num_nodes-1):
x0,y0,z0 = nodes[i]
for j in range(i+1, num_nodes):
#if matrix[i][j] > edge_thresh:
if ma_thresh[i][j] > edge_thresh:
x1,y1,z1 = nodes[j]
if weight_edges:
mlab.plot3d([x0,x1], [y0,y1], [z0,z1],
tube_radius=matrix[i][j]/matrix_flat.max(),
color=(1,1,1))
else:
mlab.plot3d([x0,x1], [y0,y1], [z0,z1],
tube_radius=edge_radius,
color=(1,1,1))
def generar(self, nivel, cursor, distancia, dim):
if nivel == 0:
mlab.points3d([cursor.x, cursor.x+3, cursor.x, cursor.x+3],[cursor.y, cursor.y, cursor.y -3,cursor.y -3],
[cursor.z, cursor.z, cursor.z -3,cursor.z -3],colormap="copper", scale_factor=.25)
cursor.x = cursor.x + distancia*np.cos(cursor.ang)
cursor.y = cursor.y + distancia*np.sin(cursor.ang)
cursor.z = cursor.z + distancia*np.sin(cursor.ang)
#a = [(cursor.x, cursor.x, cursor.x)] #para trabajar con 1D
#b = (cursor.x + distancia*np.cos(cursor.ang),cursor.y + distancia*np.sin(cursor.ang))
#mlab.points3d(a, a, a, colormap="copper", scale_factor=.25)
#cursor.x = b[0]
#cursor.y = b[1]
else:
print "nivel = ", nivel
for i in self.gramatica[self.inicial]:
if i == 'A':
self.generar(nivel-1, cursor, distancia/3.0, dim)
elif i == 'B':
cursor.x += 3*distancia
print "la x de cursor tiene: ",cursor.x
def testVisually(self):
'''blocks selected visually.'''
# if self.comm.rank == 0:
g2z,zvec = PETSc.Scatter().toZero(self.bnd.gindBlockWBand)
g2z.scatter(self.bnd.gindBlockWBand,zvec, PETSc.InsertMode.INSERT)
x = self.bnd.BlockSub2CenterCarWithoutBand(\
self.bnd.BlockInd2SubWithoutBand(zvec.getArray()) )
lx = self.bnd.BlockSub2CenterCarWithoutBand(\
self.bnd.BlockInd2SubWithoutBand(self.bnd.gindBlockWBand.getArray()))
try:
try:
from mayavi import mlab
except ImportError:
from enthought.mayavi import mlab
if self.comm.rank == 0:
mlab.figure()
mlab.points3d(x[:,0],x[:,1],x[:,2])
mlab.figure()
mlab.points3d(lx[:,0],lx[:,1],lx[:,2])
mlab.show()
#fig.add(pts1)
#fig.add(pts2)
except ImportError:
import pylab as pl
from mpl_toolkits.mplot3d import Axes3D #@UnusedImport
fig = pl.figure()
ax = fig.add_subplot(111, projection='3d')
ax.scatter3D(x[:,0],x[:,1],x[:,2],c='blue',marker='o')
ax.scatter3D(lx[:,0],lx[:,1],lx[:,2],c='red',marker='D')
pl.savefig('testVis{0}.png'.format(self.comm.rank))
pl.show()
def show3d(self, style={'A':1}):
from mayavi import mlab
scalar = []
pos = []
def rep(lst):
for (x,y,z) in lst:
yield x,y,z
if x == 0:
yield 1,y,z
if y == 0:
yield 1,1,z
if y == 0:
yield x,1,z
if z == 0:
yield x,1,1
if z == 0:
yield x,y,1
if x == 0:
yield 1,y,1
if x == y == z == 0:
yield 1,1,1
for atom, site in self.sites:
site = list(rep(site))
scalar.extend([style[atom]]*len(site))
pos.extend(self.v(site))
scalar = np.asarray(scalar)
pos = np.asarray(pos)
mlab.points3d(pos[:,0],pos[:,1],pos[:,2],scalar,scale_factor=1, resolution=10)
mlab.show()
###############################################################################
# However, only sources that lie in the plotted MRI slices are shown.
# Let's write a few lines of mayavi to see all sources.
import numpy as np # noqa
from mayavi import mlab # noqa
from surfer import Brain # noqa
brain = Brain('sample', 'lh', 'inflated', subjects_dir=subjects_dir)
surf = brain._geo
vertidx = np.where(src[0]['inuse'])[0]
mlab.points3d(surf.x[vertidx],
surf.y[vertidx],
surf.z[vertidx],
color=(1, 1, 0),
scale_factor=1.5)
###############################################################################
# Compute forward solution
# ------------------------
#
# We can now compute the forward solution.
# To reduce computation we'll just compute a single layer BEM (just inner
# skull) that can then be used for MEG (not EEG).
#
# We specify if we want a one-layer or a three-layer BEM using the
# conductivity parameter.
#
# The BEM solution requires a BEM model which describes the geometry
def plot_undirected_cnx(mat, labels, parc, fig=None, lup_title=None,
ldown_title=None, rup_title=None, rdown_title=None,
figsize=(3840,2160), lineres=1000,
subjects_dir="/home/jev/hdd/freesurfer/subjects",
alpha_max=None, alpha_min=None, uniform_weight=False,
surface="inflated", alpha=1, top_cnx=50, bot_cnx=None,
color=(1,0,0)):
if top_cnx is not None:
matflat = mat.flatten()
thresh = np.sort(matflat[matflat>0])[-top_cnx]
mat[mat<thresh] = 0
if bot_cnx is not None:
matflat = np.abs(mat.flatten())
thresh = np.sort(matflat[matflat>0])[-bot_cnx]
mat[np.abs(mat)>thresh] = 0
lingrad = np.linspace(0,1,lineres)
if fig is None:
fig = mlab.figure(size=figsize)
brain = Brain('fsaverage', 'both', surface, alpha=alpha,
subjects_dir=subjects_dir, figure=fig)
if lup_title:
brain.add_text(0, 0.8, lup_title, "lup", font_size=40)
if ldown_title:
brain.add_text(0, 0, ldown_title, "ldown", font_size=40)
if rup_title:
brain.add_text(0.7, 0.8, rup_title, "rup", font_size=40)
if rdown_title:
brain.add_text(0.7, 0., rdown_title, "rdown", font_size=40)
brain.add_annotation(parc,color="black")
rrs = np.array([brain.geo[l.hemi].coords[l.center_of_mass()] for l in labels])
if alpha_max is None:
alpha_max = np.abs(mat).max()
if alpha_min is None:
alpha_min = np.abs(mat[mat!=0]).min()
inds = np.where(mat>0)
origins = rrs[inds[0],]
dests = rrs[inds[1],]
areas = np.zeros(len(labels))
np.add.at(areas,inds[0],1)
np.add.at(areas,inds[1],1)
area_weights = areas/areas.max()
lengths = np.linalg.norm(origins-dests, axis=1)
lengths = np.broadcast_to(lengths,(3,len(lengths))).T
midpoints = (origins+dests)/2
midpoint_units = midpoints/np.linalg.norm(midpoints,axis=1,keepdims=True)
spline_mids = midpoints + midpoint_units*lengths*2
if uniform_weight:
alphas = np.ones(len(inds[0]))*0.8
area_weights[area_weights>0] = 0.8
else:
alphas = ((np.abs(mat[inds[0],inds[1]])-alpha_min)/(alpha_max-alpha_min))
alphas[alphas<0],alphas[alphas>1] = 0, 1
mlab.points3d(origins[:,0],origins[:,1],origins[:,2],
alphas,scale_factor=10,color=color,transparent=True)
mlab.points3d(dests[:,0],dests[:,1],dests[:,2],
alphas,scale_factor=10,color=color,transparent=True)
for l_idx, l in enumerate(labels):
if area_weights[l_idx] == 0:
continue
brain.add_label(l,color=color, alpha=area_weights[l_idx])
spl_pts = np.empty((len(origins),3,lineres))
for idx in range(len(origins)):
curve = bezier.Curve(np.array([[origins[idx,0],spline_mids[idx,0],dests[idx,0]],
[origins[idx,1],spline_mids[idx,1],dests[idx,1]],
[origins[idx,2],spline_mids[idx,2],dests[idx,2]]]),
degree=2)
spl_pts[idx,] = curve.evaluate_multi(lingrad)
mlab.plot3d(spl_pts[idx,0,],spl_pts[idx,1,],spl_pts[idx,2,],
lingrad*255,tube_radius=alphas[idx]*2,color=color,
opacity=alphas[idx])
return brain
def display_simple_using_mayavi_2(vf_list,
pointcloud_list,
minmax=(-1, 1),
mayavi_wireframe=False,
opacity=1.0,
separate=True,
gradients_at=None,
gradients_from_iobj=None,
pointsizes=None,
pointcloud_opacity=1.):
"""Two separate panels"""
print("Mayavi.")
sys.stdout.flush()
from mayavi import mlab
if pointsizes is None:
pointsizes = [0.2] * 10
if type(opacity) is list:
opacities = opacity # 1.0
else:
opacities = [opacity] + [0.2] * (len(vf_list) - 1) # 1.0, 0.2 #0.1
for fi in range(len(vf_list)):
if separate:
mlab.figure()
vf = vf_list[fi]
verts, faces = vf
if verts is None:
continue
if verts.size == 0:
print("Warning: empty vertices")
continue
if faces.size == 0:
print("Warning: no faces")
continue
assert verts.ndim == 2
assert faces.ndim == 2
assert verts.shape == (verts.shape[0], 3), str(verts.shape)
assert faces.shape == (faces.shape[0], 3), str(faces.shape)
if type(mayavi_wireframe) is list:
wire_frame1 = mayavi_wireframe[fi]
assert len(mayavi_wireframe) == len(vf_list)
else:
wire_frame1 = mayavi_wireframe
mlab.triangular_mesh(
[vert[0] for vert in verts], [vert[1] for vert in verts],
[vert[2] for vert in verts],
faces,
representation="surface" if not wire_frame1 else "wireframe",
opacity=opacities[fi],
scale_factor=100.0)
color_list = [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 0, 1), (0, 1, 0)]
i = 0
for c in pointcloud_list:
mlab.points3d(c[:, 0],
c[:, 1],
c[:, 2],
color=color_list[i],
scale_factor=pointsizes[i],
opacity=pointcloud_opacity)
i += 1
del i
if minmax is not None:
(RANGE_MIN, RANGE_MAX) = minmax
x = np.linspace(RANGE_MIN, RANGE_MAX, 2).reshape(2, 1)
y = np.zeros((2, 1))
z = np.zeros((2, 1))
mlab.plot3d(x, y, z, line_width=3, name="x-axis")
mlab.plot3d(y, x, z, line_width=3, name="y-axis")
mlab.plot3d(z, y, x, line_width=3, name="z-axis")
mlab.text3d(RANGE_MAX, 0, 0, "x", scale=0.3)
mlab.text3d(0, RANGE_MAX, 0, "y", scale=0.3)
mlab.text3d(0, 0, RANGE_MAX, "z", scale=0.3)
mlab.text3d(RANGE_MIN, 0, 0, "-x", scale=0.3)
mlab.text3d(0, RANGE_MIN, 0, "-y", scale=0.3)
mlab.text3d(0, 0, RANGE_MIN, "-z", scale=0.3)
def add_random_interior_points(ax, iobj, avg_edge_len):
""" Adding random points """
n = 10000
import basic_functions
ampl = avg_edge_len
x = basic_functions.make_random_vector_vectorized(n,
ampl,
1,
type="rand",
normalize=False)
v = iobj.implicitFunction(x)
x_sel = x[v >= 0, :]
if x_sel.size == 0:
print("No points")
return
ax.points3d(x_sel[:, 0],
x_sel[:, 1],
x_sel[:, 2],
color=(0, 0, 0),
scale_factor=0.2)
if gradients_at is not None:
verts1, faces1 = vf_list[0]
avg_edge_len = compute_average_edge_length(verts, faces)
visualise_gradients(mlab, gradients_at, gradients_from_iobj,
avg_edge_len / 20.)
if gradients_from_iobj is not None:
add_random_interior_points(mlab, gradients_from_iobj, avg_edge_len)
mlab.show()
return
def main(arg):
"""Main function of the script"""
client = carla.Client(arg.host, arg.port)
client.set_timeout(10.0)
client.load_world('Town05')
client.reload_world()
world = client.get_world()
try:
original_settings = world.get_settings()
settings = world.get_settings()
traffic_manager = client.get_trafficmanager(8000)
traffic_manager.set_synchronous_mode(True)
delta = 0.05
settings.fixed_delta_seconds = delta
settings.synchronous_mode = True
settings.no_rendering_mode = arg.no_rendering
world.apply_settings(settings)
blueprint_library = world.get_blueprint_library()
# vehicle_bp = blueprint_library.filter(arg.filter)[0]
# vehicle_transform = random.choice(world.get_map().get_spawn_points())
# vehicle = world.spawn_actor(vehicle_bp, vehicle_transform)
# vehicle.set_autopilot(arg.no_autopilot)
lidar_bp = generate_lidar_bp(arg, world, blueprint_library, delta)
lidar_transform = carla.Transform(carla.Location(x=-65.0, y=3.0,
z=6.0))
lidar = world.spawn_actor(lidar_bp, lidar_transform)
fig = mlab.figure(size=(960, 540), bgcolor=(0.05, 0.05, 0.05))
vis = mlab.points3d(0, 0, 0, 0, mode='point', figure=fig)
mlab.view(distance=25)
buf = {'pts': np.zeros((1, 3)), 'intensity': np.zeros(1)}
# @mlab.animate(delay=100)
def anim():
i = 0
while True:
vis.mlab_source.reset(x=buf['pts'][:, 0],
y=buf['pts'][:, 1],
z=buf['pts'][:, 2],
scalars=buf['intensity'])
mlab.savefig(f'{i%10}.png', figure=fig)
time.sleep(0.1)
i += 1
if arg.semantic:
lidar.listen(lambda data: semantic_lidar_callback(data, buf))
else:
lidar.listen(lambda data: lidar_callback(data, buf))
loopThread = threading.Thread(target=carlaEventLoop,
args=[world],
daemon=True)
loopThread.start()
anim()
# mlab.show()
finally:
world.apply_settings(original_settings)
traffic_manager.set_synchronous_mode(False)
vehicle.destroy()
lidar.destroy()
mlab.clf()
###############################################################################
# Display points at city positions
import numpy as np
coords = np.array(coords)
# First we have to convert latitude/longitude information to 3D
# positioning.
lat, long = coords.T * np.pi / 180
x = np.cos(long) * np.cos(lat)
y = np.cos(long) * np.sin(lat)
z = np.sin(long)
points = mlab.points3d(x,
y,
z,
scale_mode='none',
scale_factor=0.03,
color=(0, 0, 1))
###############################################################################
# Display connections between cities
connections = np.array(connections)
# We add lines between the points that we have previously created by
# directly modifying the VTK dataset.
points.mlab_source.dataset.lines = connections
points.mlab_source.reset()
# To represent the lines, we use the surface module. Using a wireframe
# representation allows to control the line-width.
mlab.pipeline.surface(points,
color=(1, 1, 1),
representation='wireframe',
def viewer_pointcloud(pointcloud):
mlab.figure(bgcolor=(1, 1, 1))
mlab.points3d(pointcloud[:, 0], pointcloud[:, 1], pointcloud[:, 2], color=(0, 1, 0), mode='sphere', scale_factor = 0.025)
mlab.show()
return
#prev_velo = None
#reg = pcl.IterativeClosestPointNonLinear()
fig_scale = 300
fig_ratio = [4, 3]
plt_fig = plt.figure(1, figsize=(fig_ratio[0], fig_ratio[1]), dpi=fig_scale)
ax_fig = plt_fig.add_axes([0, 0, 1, 1])
ax_fig.axis('off')
plt.show(block=False)
fig = mlab.figure(bgcolor=(0, 0, 0),
size=(fig_ratio[0] * fig_scale, fig_ratio[1] * fig_scale))
for i, velo in enumerate(dataset.velo):
mlab.clf()
cen = np.vstack((cen, [0, 0, 0, 1]))
mlab.points3d(cen[:, 0],
cen[:, 1],
cen[:, 2],
color=(1, 0, 0),
scale_factor=0.5)
velo = filter_ground(velo, cen, th=1000)
# draw annotated objects
filled_idx = np.zeros((velo.shape[0], ), dtype=bool)
for j, box in enumerate(tracklet_rects[i]):
draw_class.draw_box(box, tracklet_ids[i][j])
idx = in_hull(velo[:, :3], box[:3, :].T)
draw_class.draw_cluster(velo[idx, :], tracklet_ids[i][j])
filled_idx |= idx
# print other points
draw_class.draw_cluster(velo[~filled_idx, :])
def anim():
f = mlab.gcf()
while True:
for i in range(N):
plot_points = [(i + j) % N for j in range(10)]
# print("x: ",x[plot_points],", y: ",y[plot_points], ", z: ",z[plot_points], ", s: ", scalars[plot_points])
pts.mlab_source.reset(x=x[plot_points],
y=y[plot_points],
z=z[plot_points],
scalars=scalars[plot_points])
pts.module_manager.scalar_lut_manager.lut.table = colors[
plot_points]
yield
pts = mlab.points3d(x, y, z, scalars, mode='sphere',
scale_factor=10.0) # Create points
pts.glyph.color_mode = 'color_by_scalar' # Color by scalar
pts.glyph.scale_mode = 'scale_by_vector'
# Set look-up table and redraw
pts.module_manager.scalar_lut_manager.lut.table = colors
base_pts = mlab.points3d(x,
y,
z,
scalars,
mode='sphere',
opacity=0.2,
scale_factor=10.0) # Create points
base_pts.glyph.color_mode = 'color_by_scalar' # Color by scalar
base_pts.glyph.scale_mode = 'scale_by_vector'
# Set look-up table and redraw
def update_plot(self):
## PLot to Show
outline = mayavi.mlab.outline(line_width=3)
outline.outline_mode = 'cornered'
def Rx(theta_x):
return np.array([[1, 0, 0], [0,
np.cos(theta_x), -np.sin(theta_x)],
[0, np.sin(theta_x),
np.cos(theta_x)]])
def Ry(theta_y):
return np.array([[np.cos(theta_y), 0,
np.sin(theta_y)], [0, 1, 0],
[-np.sin(theta_y), 0,
np.cos(theta_y)]])
def Rz(theta_z):
return np.array([[np.cos(theta_z), -np.sin(theta_z), 0],
[np.sin(theta_z),
np.cos(theta_z), 0], [0, 0, 1]])
HKL_aligned = np.array([0, 1, 2])
mu = 0.34
structures[0].lattice.set_E_keV(energy_keV)
angles = structures[0].lattice.angles(HKL_aligned, mu)
theta = angles[0][0]
gamma = angles[0][1]
delta = angles[0][2]
print('theta =', theta)
print('gamma =', gamma)
print('delta =', delta)
theta = angles[1][0]
gamma = angles[1][1]
delta = angles[1][2]
print('theta =', theta)
print('gamma =', gamma)
print('delta =', delta)
#ki = Ry(np.deg2rad(mu)).dot(Rz(np.deg2rad(theta)).dot([k0, 0, 0]))
ki = np.dot(Rz(np.deg2rad(theta)),
np.dot(Ry(np.deg2rad(mu)), [k0, 0, 0]))
#kf = np.dot(Rz(np.deg2rad(theta)), np.dot(Ry(np.deg2rad(mu)), np.dot(Rz(np.deg2rad(delta)), np.dot(Ry(np.deg2rad(gamma)),[k0, 0, 0]))))
kf = np.dot(
Rz(np.deg2rad(delta)),
np.dot(
Ry(np.deg2rad(gamma)),
np.dot(Rz(np.deg2rad(theta)),
np.dot(Ry(np.deg2rad(mu)), [k0, 0, 0]))))
#kf = Rz(np.deg2rad(theta-delta)).dot(Ry(np.deg2rad(mu-gamma)).dot([k0, 0, 0]))
#ki = rsp.RotationMatrix(0, np.deg2rad(mu), np.deg2rad(theta)).dot([k0,0,0])
#kf = rsp.RotationMatrix(0, np.deg2rad(-gamma), np.deg2rad(delta)).dot([k0,0,0])
mlab.plot3d([-ki[0], 0], [-ki[1], 0], [-ki[2], 0], color=(0, 0, 1))
mlab.plot3d([-ki[0], -ki[0] + kf[0]], [-ki[1], -ki[1] + kf[1]],
[-ki[2], -ki[2] + kf[2]],
color=(0, 0, 1))
sphere = mlab.points3d(-ki[0],
-ki[1],
-ki[2],
scale_mode='none',
scale_factor=2 * k0,
color=(0.67, 0.77, 0.93),
resolution=50,
opacity=0.7)
sphere.actor.property.specular = 0.45
sphere.actor.property.specular_power = 5
sphere.actor.property.backface_culling = True
peaks = []
space_plots = []
for i in range(len(structures)):
struc = structures[i]
space_plots.append(rsplt.space_plot(struc.lattice))
if (struc.plot_peaks):
peaks.append(space_plots[i].plot_peaks(
qx_lims=qx_lims,
qy_lims=qy_lims,
qz_lims=qz_lims,
q_inplane_lim=q_inplane_lim,
mag_q_lims=mag_q_lims,
color=struc.color))
if (struc.plot_rods):
space_plots[i].plot_rods(qx_lims=qx_lims,
qy_lims=qy_lims,
qz_lims=qz_lims,
q_inplane_lim=q_inplane_lim,
color=struc.color)
if (struc.plot_grid):
space_plots[i].plot_grid(qx_lims=qx_lims,
qy_lims=qy_lims,
qz_lims=qz_lims,
q_inplane_lim=q_inplane_lim,
color=struc.color)
if (struc.plot_unitcell):
space_plots[i].plot_unit_cell()
if (plot_axes):
q1 = structures[0].lattice.q([1, 0, 0])
q2 = structures[0].lattice.q([0, 1, 0])
q3 = structures[0].lattice.q([0, 0, 1])
rsplt.Arrow_From_A_to_B(0,
0,
0,
q1[0],
q1[1],
q1[2],
color=(0, 0, 0))
rsplt.Arrow_From_A_to_B(0,
0,
0,
q2[0],
q2[1],
q2[2],
color=(0, 0, 0))
rsplt.Arrow_From_A_to_B(0,
0,
0,
q3[0],
q3[1],
q3[2],
color=(0, 0, 0))
number_of_peaks = []
for peak in peaks:
number_of_peaks.append(peak[0].glyph.glyph_source.glyph_source.
output.points.to_array().shape[0])
mayavi.mlab.figure(figure=mayavi.mlab.gcf(engine=None),
bgcolor=(0, 0, 0))
def display_simple_using_mayavi_2(vf_list, pointcloud_list, minmax=(-1, 1), mayavi_wireframe=False, opacity=1.0, separate=True,
gradients_at=None, gradients_from_iobj=None, pointsizes=None, pointcloud_opacity=1.):
"""Two separate panels"""
sys.stderr.write("Mayavi.")
sys.stdout.flush()
if pointsizes is None:
pointsizes = [0.2]*10
if type(opacity) is list:
opacities = opacity # 1.0
else:
opacities = [opacity] + [0.2]*(len(vf_list)-1) # 1.0, 0.2 #0.1
for fi in range(len(vf_list)):
if separate:
mlab.figure()
vf = vf_list[fi]
vertex, faces = vf
if vertex is None:
continue
if vertex.size == 0:
sys.stderr.write("Warning: empty vertex")
continue
if faces.size == 0:
sys.stderr.write("Warning: no faces")
continue
assert vertex.ndim == 2
assert faces.ndim == 2
assert vertex.shape == (vertex.shape[0], 3), str(vertex.shape)
assert faces.shape == (faces.shape[0], 3), str(faces.shape)
if type(mayavi_wireframe) is list:
wire_frame1 = mayavi_wireframe[fi]
assert len(mayavi_wireframe) == len(vf_list)
else:
wire_frame1 = mayavi_wireframe
mlab.triangular_mesh([vert[0] for vert in vertex],
[vert[1] for vert in vertex],
[vert[2] for vert in vertex],
faces,
representation="surface" if not wire_frame1 else "wireframe",
opacity=opacities[fi], scale_factor=100.0)
color_list = [(1, 0, 0), (0, 0, 0), (1, 1, 0), (0, 0, 1), (0, 1, 0)]
i = 0
for c in pointcloud_list:
mlab.points3d(c[:, 0], c[:, 1], c[:, 2], color=color_list[i], scale_factor=pointsizes[i], opacity=pointcloud_opacity)
i += 1
del i
if minmax is not None:
(RANGE_MIN, RANGE_MAX) = minmax
x = np.linspace(RANGE_MIN, RANGE_MAX, 2).reshape(2, 1)
y = np.zeros((2, 1))
z = np.zeros((2, 1))
mlab.plot3d(x, y, z, line_width=3, name="x-axis")
mlab.plot3d(y, x, z, line_width=3, name="y-axis")
mlab.plot3d(z, y, x, line_width=3, name="z-axis")
mlab.text3d(RANGE_MAX, 0, 0, "x", scale=0.3)
mlab.text3d(0, RANGE_MAX, 0, "y", scale=0.3)
mlab.text3d(0, 0, RANGE_MAX, "z", scale=0.3)
mlab.text3d(RANGE_MIN, 0, 0, "-x", scale=0.3)
mlab.text3d(0, RANGE_MIN, 0, "-y", scale=0.3)
mlab.text3d(0, 0, RANGE_MIN, "-z", scale=0.3)
mlab.show()
return
scale_factor = 0.25
outputDir = 'panels'
arrowOffset = 0.25
ao = arrowOffset
X, Y = np.meshgrid(np.arange(nX), np.arange(nY))
X = X.ravel()
Y = Y.ravel()
ZE = np.zeros(nX * nY)
ZI = np.zeros(nX * nY) + ZOffset
mlab.figure(bgcolor=(1, 1, 1), size=(610*2, 582*2))
# Layers
mlab.points3d(X, Y, ZE, scale_factor=scale_factor, color=(1, 0, 0),
resolution=resolution)
mlab.points3d(X, Y, ZI, scale_factor=scale_factor, color=(0, 0, 1),
resolution=resolution)
# gE, gI arrows
x = [ao * nX, (0.9-ao)*nX]
y = [ao * nY, (0.9-ao)*nY]
zStart = [0.2 * ZOffset, 0.8 * ZOffset]
zComp = np.array([0.6 * ZOffset, -0.6 * ZOffset])
mlab.quiver3d(x, y, zStart, [0, 0], [0, 0], zComp, color=(0, 0, 0), line_width=20,
scale_factor=1, mode='arrow', resolution=resolution)
textScale = 2.0
mlab.text3d(x[0] - 0.075*nX, y[0] - 0.075*nY, ZOffset/2, "gE", orientation=[0, 0, 0],
orient_to_camera=True, color=(0, 0, 0), line_width=3, scale=textScale)
qmPos, tmpCols = getActiveAtoms(pos, rcols)
qmx, qmy, qmz = qm[0].T
posx, posy, posz = qmPos[0].T
# Separate positions into carbon and hydrogen
carbonPos, _ = getActiveAtoms(qmPos, tmpCols, "C")
hydrogPos, _ = getActiveAtoms(qmPos, tmpCols, "H")
neonPos, _ = getActiveAtoms(pos, rcols, "Ne")
# Init forces
activeFrc, _ = getActiveAtoms(frc, fcols)
frcx, frcy, frcz = activeFrc[0].T
#Init ad Momentum
stateMom, _ = load_tintf.find_in_histF(am, amcols, {"state": 0, "step_num": 0})
amx, amy, amz = stateMom.T
# Plot atoms
cPts = mlab.points3d(*carbonPos[0].T, scale_factor=sizes['C'], color=(0, 0, 0))
hPts = mlab.points3d(*hydrogPos[0].T, scale_factor=sizes['H'], color=(1, 1, 0))
nePts = mlab.points3d(*hydrogPos[0].T,
scale_factor=sizes['Ne'],
color=(1, 1, 1))
# Plot vectors
#qmPts = mlab.quiver3d(posx, posy, posz, qmx, qmy, qmz)
frcPts = mlab.quiver3d(posx, posy, posz, frcx, frcy, frcz)
#amPts = mlab.quiver3d(posx, posy, posz, amx, amy, amz)
@mlab.animate(delay=10)
def anim():
for i in range(nsteps):
cPts.mlab_source.points = carbonPos[i]
hPts.mlab_source.points = hydrogPos[i]
m0 = 4; m1 = 3; m2 = 2; m3 = 3; m4 = 1; m5 = 2; m6 = 2; m7 = 4;
s = sin(m0*phi)**m1 + cos(m2*phi)**m3 + sin(m4*theta)**m5 + cos(m6*theta)**m7
x = sin(phi)*cos(theta)
y = cos(phi)
z = sin(phi)*sin(theta)
################################################################################
# Plot the data
from mayavi import mlab
# A first plot in 3D
fig = mlab.figure(1)
mlab.clf()
mesh = mlab.mesh(x, y, z, scalars=s)
cursor3d = mlab.points3d(0., 0., 0., mode='axes',
color=(0, 0, 0),
scale_factor=0.5)
mlab.title('Click on the ball')
# A second plot, flat
fig2d = mlab.figure(2)
mlab.clf()
im = mlab.imshow(s)
cursor = mlab.points3d(0, 0, 0, mode='2dthick_cross',
color=(0, 0, 0),
scale_factor=10)
mlab.view(90, 0)
################################################################################
# Some logic to select 'mesh' and the data index when picking.
x = data[:, 0]
y = data[:, 1]
z = data[:, 2]
thresh = -1.4
# bgcolor=(0.898,0.803,1.0)
f = mlab.figure(figure="LidarData",
size=(1920, 1080),
bgcolor=(0, 0, 0),
fgcolor=(0., 1.0, 0))
# f.view(azimuth=None, elevation=None, distance=5, focalpoint=None,roll=None, reset_roll=True, figure=None)
f.scene.movie_maker.record = True
plt = mlab.points3d(x, y, z, mode="point", colormap="copper", scale_factor=100)
#below '@' stuff is known as a 'decorator' in python
@mlab.animate(delay=100)
def anim():
for i in range(143):
print("file number = " + str(i))
filename = os.path.join(foldername, '%06d.bin' % i)
data = np.fromfile(filename, np.float32)
data = np.array(data.reshape(data.shape[0] // 4, 4))
x_ = data[:, 0]
y_ = data[:, 1]
z_ = data[:, 2]
# Create data with x and y random in the [-2, 2] segment, and z a
# Gaussian function of x and y.
np.random.seed(12345)
x = 4 * (np.random.random(500) - 0.5)
y = 4 * (np.random.random(500) - 0.5)
def f(x, y):
return np.exp(-(x**2 + y**2))
z = f(x, y)
mlab.figure(1, fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))
# Visualize the points
pts = mlab.points3d(x, y, z, z, scale_mode='none', scale_factor=0.2)
# Create and visualize the mesh
mesh = mlab.pipeline.delaunay2d(pts)
surf = mlab.pipeline.surface(mesh)
mlab.view(47, 57, 8.2, (0.1, 0.15, 0.14))
mlab.show()
### EXAMPLE 4
### surface with spheres
x, y = np.meshgrid(np.linspace(-2.5, 2), np.linspace(-2, 2))
f = lambda x, y: .4 * np.sin(2 * np.pi * x) * np.sin(2 * np.pi * y)
z = f(x, y)
mlab.surf(x.T, y.T, z.T, colormap="copper")
px, py = np.meshgrid(np.arange(-2, 2) + .25, np.arange(-2, 2) + .75)
px, py = px.flatten(), py.flatten()
def show_points(fn,
fne=None,
cm="Blues",
mode="point",
color=None,
swap=True,
scaled=True,
dscale=1,
what=3,
opacity=1.0,
dumb=False):
L = []
S = []
D = []
if scaled:
from dataset import WurzelInfo
info = WurzelInfo(fn)
print "Scale: ", info.scale
wireframe = False
if what == "wireframe":
wireframe = True
what = 3
print "Show Point3D `%s'" % fn
with open(info.datapath + "/" + fn) as f:
for line in f.readlines():
line = line.split()
#if len(line)<7:
# continue
L.append([float(x) for x in line[:3]]) # 0..2: coordinate axes
S.append(float(line[what])) # 3: mass, 4: diameter
D.append([float(x) for x in line[4:]]) # D: not used
S = np.array(S)
L = np.vstack(L)
D = np.vstack(D)
if wireframe:
S[:] = 1
print "NumPoints: ", L.shape[0]
if scaled:
L /= info.scale # L is in mm, now it matches raw data again
if what == 4:
S /= info.scale # S is in mm, now it matches raw data again
S *= dscale
#S[S<0.5]=0.5
#mlab.quiver3d(L[:,0],L[:,1],L[:,2], D[:,0],D[:,1],D[:,2], scale_factor=3.)
if (L < 0).sum() > 0 and False:
""" this is in cm """
L[:, 0] = (-L[:, 0]) / 100. * 256 + 120
L[:, 1] = (L[:, 1]) / 100. * 256 + 110
L[:, 2] = (L[:, 2]) / 131. * 256 + -0
#L[:,0] = (-L[:,0])/100.*256 + 110
#L[:,1] = (L[:,1])/100.*256 + 120
#L[:,2] = (L[:,2])/131.*256 + -5
gt = True
else:
L[:] += 0.5
gt = False
#L = xyz2pol(L)
#mlab.points3d(L[:,0].squeeze(),L[:,1].squeeze(),L[:,2].squeeze(),S,scale_factor=1.5,colormap="bone",scale_mode="none",mode="2dtriangle")
#pts = mlab.points3d(L[:,0],L[:,1],L[:,2],S,scale_factor=0.8,colormap=cm,scale_mode="none",resolution="20",mode=mode)
#pts = mlab.points3d(L[:,0],L[:,1],L[:,2],S,scale_factor=0.1,colormap=cm,scale_mode="none",resolution="20",mode=mode)
assert np.isnan(S).sum() == 0
if dumb:
pts = mlab.points3d(L[:, 0],
L[:, 1],
L[:, 2],
S,
scale_factor=1,
resolution="20",
mode=mode,
color=color)
else:
pts = mlab.points3d(L[:, 0],
L[:, 1],
L[:, 2],
S,
scale_factor=1,
colormap=cm,
resolution="20",
mode=mode,
scale_mode='scalar')
pts.actor.visible = False
print "Done."
if None == fne: return
print "Show Edges3D"
E = []
thresh = 100
with open(info.datapath + "/" + fne) as f:
for line in f.readlines():
vec = line.split()
line = [int(x) for x in vec[:2]]
try:
if np.linalg.norm(L[line[0], :] - L[line[1], :]) > thresh:
continue
except:
pass
E.append(line)
E = np.vstack(E)
pts.mlab_source.dataset.lines = E
tube = mlab.pipeline.tube(pts,
tube_sides=7,
tube_radius=1,
name="root tubes")
#tube.filter.radius_factor = 1.
tube.filter.vary_radius = 'vary_radius_by_absolute_scalar'
#tube.filter.vary_radius = 'vary_radius_by_scalar'
tube.filter.capping = True
#tube.filter.use_default_normal = True
if color == None:
color = (1, 0, 0) if gt else (0, 1, 0)
#color = (0,1,0) if gt else (1,0,0)
res = mlab.pipeline.surface(tube, color=color)
res.actor.property.opacity = opacity
print "Done."
def selectPoints(self, visu_factor, iso):
plot = Visualization.DataVisulization(
ndimage.gaussian_filter(self.image1, visu_factor), iso)
plot.contour3d()
plot2 = Visualization.DataVisulization(
ndimage.gaussian_filter(self.image2, visu_factor), iso)
plot2.contour3d()
nodes, elements = plot.surf_points, plot.surf_elements
nodes2, elements2 = plot2.surf_points, plot2.surf_elements
whole_nodes = np.concatenate((nodes, nodes2), axis=0)
sx, sy = nodes.shape
elements3 = elements2 + sx
whole_elements = np.concatenate((elements, elements3), axis=0)
#### help function #####
def reshapePoint(point1, point2): # reshape the points array
xpos = np.ndarray((2))
ypos = np.ndarray((2))
zpos = np.ndarray((2))
xpos[0] = point1[0]
xpos[1] = point2[0]
ypos[0] = point1[1]
ypos[1] = point2[1]
zpos[0] = point1[2]
zpos[1] = point2[2]
color = 0.5 * np.ones(xpos.shape)
return xpos, ypos, zpos, color
# A first plot in 3D
fig = mlab.figure(2)
mlab.clf()
mesh = mlab.triangular_mesh(whole_nodes[:, 0], whole_nodes[:, 1],
whole_nodes[:, 2], whole_elements)
cursor3d = mlab.points3d(0.,
0.,
0.,
mode='axes',
color=(0, 0, 0),
scale_factor=0.5)
mlab.title('Click on the LC Region')
## A second plot, flat projection image, this is under development
#fig2d = mlab.figure(2)
#mlab.clf()
#im = mlab.imshow(s)
#cursor = mlab.points3d(0, 0, 0, mode='2dthick_cross',
# color=(0, 0, 0),
# scale_factor=10)
#mlab.view(90, 0)
################################################################################
def picker_callback(picker_obj):
print "selecting a new point"
picked = picker_obj.actors
if mesh.actor.actor._vtk_obj in [o._vtk_obj for o in picked]:
# m.mlab_source.points is the points array underlying the vtk
# dataset. GetPointId return the index in this array.
print picker_obj.point_id
x_ = np.lib.index_tricks.unravel_index(picker_obj.point_id,
whole_nodes[:, 0].shape)
print whole_nodes[:, 0][x_], whole_nodes[:, 1][
x_], whole_nodes[:, 2][x_]
cursor3d.mlab_source.reset(x=whole_nodes[:, 0][x_],
y=whole_nodes[:, 1][x_],
z=whole_nodes[:, 2][x_])
fig.on_mouse_pick(picker_callback)
mlab.show()
map_actor = tvtk.Actor2D()
map_actor.mapper = mapper
map_actor.position = (0, 0)
# Opacity doesn't seem to be supported, see
# http://vtk.1045678.n5.nabble.com/Opacity-of-vtkActor2D-td1238026.html
map_actor.property.opacity = 0.5
renderer = scene.renderer
renderer.add_actor2d(map_actor)
window = scene.render_window
def fit_image_to_window():
resize.output_dimensions = fit_to(imsize[0], imsize[1], window.size[0],
window.size[1])
fit_image_to_window()
def on_window_modified(obj, evt):
#print 'window size : ', window.size
fit_image_to_window()
window.add_observer('ModifiedEvent', on_window_modified)
mlab.points3d(1, 1, 1)
##
def extract_frustum_data_rgb_detection(det_filename,
split,
output_filename,
viz=False,
type_whitelist=['Car'],
img_height_threshold=25,
lidar_point_threshold=5):
''' Extract point clouds in frustums extruded from 2D detection boxes.
Update: Lidar points and 3d boxes are in *rect camera* coord system
(as that in 3d box label files)
Input:
det_filename: string, each line is
img_path typeid confidence xmin ymin xmax ymax
split: string, either trianing or testing
output_filename: string, the name for output .pickle file
type_whitelist: a list of strings, object types we are interested in.
img_height_threshold: int, neglect image with height lower than that.
lidar_point_threshold: int, neglect frustum with too few points.
Output:
None (will write a .pickle file to the disk)
'''
dataset = kitti_object(os.path.join(ROOT_DIR, 'dataset/KITTI/object'),
split)
det_id_list, det_type_list, det_box2d_list, det_prob_list = \
read_det_file(det_filename)
cache_id = -1
cache = None
id_list = []
type_list = []
box2d_list = []
prob_list = []
input_list = [] # channel number = 4, xyz,intensity in rect camera coord
frustum_angle_list = [] # angle of 2d box center from pos x-axis
for det_idx in range(len(det_id_list)):
data_idx = det_id_list[det_idx]
# print('det idx: %d/%d, data idx: %d' % \
# (det_idx, len(det_id_list), data_idx))
if cache_id != data_idx:
calib = dataset.get_calibration(data_idx) # 3 by 4 matrix
pc_velo = dataset.get_lidar(data_idx)
pc_rect = np.zeros_like(pc_velo)
pc_rect[:, 0:3] = calib.project_velo_to_rect(pc_velo[:, 0:3])
pc_rect[:, 3] = pc_velo[:, 3]
img = dataset.get_image(data_idx)
img_height, img_width, img_channel = img.shape
_, pc_image_coord, img_fov_inds = get_lidar_in_image_fov(\
pc_velo[:,0:3], calib, 0, 0, img_width, img_height, True, clip_distance=0.0)
cache = [calib, pc_rect, pc_image_coord, img_fov_inds]
cache_id = data_idx
else:
calib, pc_rect, pc_image_coord, img_fov_inds = cache
if det_type_list[det_idx] not in type_whitelist: continue
# 2D BOX: Get pts rect backprojected
xmin, ymin, xmax, ymax = det_box2d_list[det_idx]
box_fov_inds = (pc_image_coord[:,0]<xmax) & \
(pc_image_coord[:,0]>=xmin) & \
(pc_image_coord[:,1]<ymax) & \
(pc_image_coord[:,1]>=ymin)
box_fov_inds = box_fov_inds & img_fov_inds
pc_in_box_fov = pc_rect[box_fov_inds, :]
# Get frustum angle (according to center pixel in 2D BOX)
box2d_center = np.array([(xmin + xmax) / 2.0, (ymin + ymax) / 2.0])
uvdepth = np.zeros((1, 3))
uvdepth[0, 0:2] = box2d_center
uvdepth[0, 2] = 20 # some random depth
box2d_center_rect = calib.project_image_to_rect(uvdepth)
frustum_angle = -1 * np.arctan2(box2d_center_rect[0, 2],
box2d_center_rect[0, 0])
# Pass objects that are too small
if ymax-ymin<img_height_threshold or \
len(pc_in_box_fov)<lidar_point_threshold:
continue
id_list.append(data_idx)
type_list.append(det_type_list[det_idx])
box2d_list.append(det_box2d_list[det_idx])
prob_list.append(det_prob_list[det_idx])
input_list.append(pc_in_box_fov)
frustum_angle_list.append(frustum_angle)
with open(output_filename, 'wb') as fp:
pickle.dump(id_list, fp)
pickle.dump(box2d_list, fp)
pickle.dump(input_list, fp)
pickle.dump(type_list, fp)
pickle.dump(frustum_angle_list, fp)
pickle.dump(prob_list, fp)
if viz:
import mayavi.mlab as mlab
for i in range(10):
p1 = input_list[i]
fig = mlab.figure(figure=None,
bgcolor=(0.4, 0.4, 0.4),
fgcolor=None,
engine=None,
size=(500, 500))
mlab.points3d(p1[:, 0],
p1[:, 1],
p1[:, 2],
p1[:, 1],
mode='point',
colormap='gnuplot',
scale_factor=1,
figure=fig)
fig = mlab.figure(figure=None,
bgcolor=(0.4, 0.4, 0.4),
fgcolor=None,
engine=None,
size=(500, 500))
mlab.points3d(p1[:, 2],
-p1[:, 0],
-p1[:, 1],
seg,
mode='point',
colormap='gnuplot',
scale_factor=1,
figure=fig)
raw_input()
#H=nx.krackhardt_kite_graph()
#H=nx.Graph();H.add_edge('a','b');H.add_edge('a','c');H.add_edge('a','d')
#H=nx.grid_2d_graph(4,5)
H=nx.cycle_graph(20)
# reorder nodes from 0,len(G)-1
G=nx.convert_node_labels_to_integers(H)
# 3d spring layout
pos=nx.spring_layout(G,dim=3)
# numpy array of x,y,z positions in sorted node order
xyz=np.array([pos[v] for v in sorted(G)])
# scalar colors
scalars=np.array(G.nodes())+5
mlab.figure(1, bgcolor=(0, 0, 0))
mlab.clf()
pts = mlab.points3d(xyz[:,0], xyz[:,1], xyz[:,2],
scalars,
scale_factor=0.1,
scale_mode='none',
colormap='Blues',
resolution=20)
pts.mlab_source.dataset.lines = np.array(G.edges())
tube = mlab.pipeline.tube(pts, tube_radius=0.01)
mlab.pipeline.surface(tube, color=(0.8, 0.8, 0.8))
mlab.savefig('mayavi2_spring.png')
mlab.show() # interactive window
import numpy as np from mayavi import mlab #建立数据 t = np.linspace(0, 4 * np.pi, 20) x = np.sin(2 * t) y = np.cos(t) z = np.cos(2 * t) s = 2 + np.sin(t) #对数据进行可视化 points = mlab.points3d(x, y, z, s, colormap="Reds", scale_factor=.25) mlab.show()
print(len(names))
co_orb = co_orbiting_sats(names)
counter_orb = counter_orbiting_sats(names)
pot_lmc = lmc_mw_ratio_potetial()
savetxt('pot_lmc_mw3d.txt', pot_lmc)
mlab.figure(bgcolor=(0.0, 0.0, 0.0))
tt = mlab.contour3d(pot_lmc,
contours=20,
opacity=.3,
extent=[-300, 300, -300, 300, -300, 300],
transparent=True,
colormap='viridis')
mlab.points3d(x_sats, y_sats, z_sats, color=(1,0,0),\
scale_factor=10, opacity=1)
mlab.quiver3d(x_sats, y_sats, z_sats, vx_sats, vy_sats,\
vz_sats, color=(1,0,0), scale_factor=30, opacity=1)
mlab.points3d(x_sats[co_orb], y_sats[co_orb], z_sats[co_orb],\
color=(1,1,0), scale_factor=10, opacity=1)
mlab.quiver3d(x_sats[co_orb], y_sats[co_orb], z_sats[co_orb],\
vx_sats[co_orb], vy_sats[co_orb], vz_sats[co_orb],\
color=(1,1,0), scale_factor=30, opacity=1)
mlab.points3d(x_sats[counter_orb], y_sats[counter_orb], z_sats[counter_orb],\
color=(1,0,1), scale_factor=10, opacity=1)
mlab.quiver3d(x_sats[counter_orb], y_sats[counter_orb], z_sats[counter_orb],\
vx_sats[counter_orb], vy_sats[counter_orb], vz_sats[counter_orb],\
color=(1,0,1), scale_factor=30, opacity=1)
#scan_dir = f'data/kitti/object/training/velodyne_reduced/{file_id}.bin'
scan_dir = f'velodyne_reduced_reduced/{file_id}.bin'
if not os.path.exists(scan_dir):
print("skip %s for pc" % file_id)
continue
scan = np.fromfile(scan_dir, dtype=np.float32).reshape(-1, 4)
# load labels
#label_dir = f'data/kitti/object/training/label_2/{file_id}.txt'
label_dir = f'label_2_reduced/{file_id}.txt'
with open(label_dir, 'r') as f:
labels = f.readlines()
# draw point cloud
plot = mlab.points3d(scan[:, 0],
scan[:, 1],
scan[:, 2],
mode="point",
figure=fig,
color=(1, 0, 0))
for line in labels:
draw_box(line, colors)
mlab.view(azimuth=230, distance=100)
mlab.savefig(filename='reduced_output/%s_GT.png' % file_id)
print("Saved %s_GT.png" % file_id)
mlab.clf(fig)
mlab.close(all=True)
def extract_frustum_data(idx_filename,
split,
output_filename,
viz=False,
perturb_box2d=False,
augmentX=1,
type_whitelist=['Car']):
''' Extract point clouds and corresponding annotations in frustums
defined generated from 2D bounding boxes
Lidar points and 3d boxes are in *rect camera* coord system
(as that in 3d box label files)
Input:
idx_filename: string, each line of the file is a sample ID
split: string, either trianing or testing
output_filename: string, the name for output .pickle file
viz: bool, whether to visualize extracted data
perturb_box2d: bool, whether to perturb the box2d
(used for data augmentation in train set)
augmentX: scalar, how many augmentations to have for each 2D box.
type_whitelist: a list of strings, object types we are interested in.
Output:
None (will write a .pickle file to the disk)
'''
dataset = kitti_object(os.path.join(ROOT_DIR, 'dataset/KITTI/object'),
split)
data_idx_list = [int(line.rstrip()) for line in open(idx_filename)]
id_list = [] # int number
box2d_list = [] # [xmin,ymin,xmax,ymax]
box3d_list = [] # (8,3) array in rect camera coord
input_list = [] # channel number = 4, xyz,intensity in rect camera coord
label_list = [] # 1 for roi object, 0 for clutter
type_list = [] # string e.g. Car
heading_list = [] # ry (along y-axis in rect camera coord) radius of
# (cont.) clockwise angle from positive x axis in velo coord.
box3d_size_list = [] # array of l,w,h
frustum_angle_list = [] # angle of 2d box center from pos x-axis
pos_cnt = 0
all_cnt = 0
for data_idx in data_idx_list:
print('------------- ', data_idx)
calib = dataset.get_calibration(data_idx) # 3 by 4 matrix
objects = dataset.get_label_objects(data_idx)
pc_velo = dataset.get_lidar(data_idx)
pc_rect = np.zeros_like(pc_velo)
pc_rect[:, 0:3] = calib.project_velo_to_rect(pc_velo[:, 0:3])
pc_rect[:, 3] = pc_velo[:, 3]
img = dataset.get_image(data_idx)
img_height, img_width, img_channel = img.shape
_, pc_image_coord, img_fov_inds = get_lidar_in_image_fov(
pc_velo[:, 0:3],
calib,
0,
0,
img_width,
img_height,
True,
clip_distance=0.0)
for obj_idx in range(len(objects)):
if objects[obj_idx].type not in type_whitelist: continue
# BEV cuboids
if objects[obj_idx].type == 'Car':
obj = objects[obj_idx]
box3d_pts_2d, box3d_pts_3d = utils.compute_box_3d(obj, calib.P)
corners_3d_velo = calib.project_rect_to_velo(box3d_pts_3d)
xmax = np.max(corners_3d_velo[:, 0])
xmin = np.min(corners_3d_velo[:, 0])
ymax = np.max(corners_3d_velo[:, 1])
ymin = np.min(corners_3d_velo[:, 1])
for _ in range(augmentX):
box2d = np.array([xmin, ymin, xmax, ymax])
if perturb_box2d:
xmin, ymin, ymax, ymax = random_shift_box2d(box2d)
box_fov_inds = (pc_velo[:, 0] < xmax) & (
pc_velo[:, 0] >= xmin) & (pc_velo[:, 1] < ymax) & (
pc_velo[:, 1] >= ymin)
# box_fov_inds = box_fov_inds & img_fov_inds
pc_in_box_fov = pc_rect[box_fov_inds, :]
xc = (xmin + xmax) / 2.0
yc = (ymin + ymax) / 2.0
frustum_angle = -1 * np.arctan2(xc, -yc)
_, inds = extract_pc_in_box3d(pc_in_box_fov, box3d_pts_3d)
label = np.zeros((pc_in_box_fov.shape[0]))
label[inds] = 1
heading_angle = obj.ry
# Get 3D BOX size
box3d_size = np.array([obj.l, obj.w, obj.h])
if np.sum(label) == 0:
continue
id_list.append(data_idx)
box2d_list.append(np.array([xmin, ymin, xmax, ymax]))
box3d_list.append(box3d_pts_3d)
input_list.append(pc_in_box_fov)
label_list.append(label)
type_list.append(objects[obj_idx].type)
heading_list.append(heading_angle)
box3d_size_list.append(box3d_size)
frustum_angle_list.append(frustum_angle)
# collect statistics
pos_cnt += np.sum(label)
all_cnt += pc_in_box_fov.shape[0]
# 2D BOX: Get pts rect backprojected
box2d = objects[obj_idx].box2d
for _ in range(augmentX):
# Augment data by box2d perturbation
if perturb_box2d:
xmin, ymin, xmax, ymax = random_shift_box2d(box2d)
else:
xmin, ymin, xmax, ymax = box2d
box_fov_inds = (pc_image_coord[:,0]<xmax) & \
(pc_image_coord[:,0]>=xmin) & \
(pc_image_coord[:,1]<ymax) & \
(pc_image_coord[:,1]>=ymin)
box_fov_inds = box_fov_inds & img_fov_inds
pc_in_box_fov = pc_rect[box_fov_inds, :]
# Get frustum angle (according to center pixel in 2D BOX)
box2d_center = np.array([(xmin + xmax) / 2.0,
(ymin + ymax) / 2.0])
uvdepth = np.zeros((1, 3))
uvdepth[0, 0:2] = box2d_center
uvdepth[0, 2] = 20 # some random depth
box2d_center_rect = calib.project_image_to_rect(uvdepth)
frustum_angle = -1 * np.arctan2(box2d_center_rect[0, 2],
box2d_center_rect[0, 0])
# 3D BOX: Get pts velo in 3d box
obj = objects[obj_idx]
box3d_pts_2d, box3d_pts_3d = utils.compute_box_3d(obj, calib.P)
_, inds = extract_pc_in_box3d(pc_in_box_fov, box3d_pts_3d)
label = np.zeros((pc_in_box_fov.shape[0]))
label[inds] = 1
# Get 3D BOX heading
heading_angle = obj.ry
# Get 3D BOX size
box3d_size = np.array([obj.l, obj.w, obj.h])
# Reject too far away object or object without points
if ymax - ymin < 25 or np.sum(label) == 0:
continue
id_list.append(data_idx)
box2d_list.append(np.array([xmin, ymin, xmax, ymax]))
box3d_list.append(box3d_pts_3d)
input_list.append(pc_in_box_fov)
label_list.append(label)
type_list.append(objects[obj_idx].type)
heading_list.append(heading_angle)
box3d_size_list.append(box3d_size)
frustum_angle_list.append(frustum_angle)
# collect statistics
pos_cnt += np.sum(label)
all_cnt += pc_in_box_fov.shape[0]
print('Average pos ratio: %f' % (pos_cnt / float(all_cnt)))
print('Average npoints: %f' % (float(all_cnt) / len(id_list)))
with open(output_filename, 'wb') as fp:
pickle.dump(id_list, fp)
pickle.dump(box2d_list, fp)
pickle.dump(box3d_list, fp)
pickle.dump(input_list, fp)
pickle.dump(label_list, fp)
pickle.dump(type_list, fp)
pickle.dump(heading_list, fp)
pickle.dump(box3d_size_list, fp)
pickle.dump(frustum_angle_list, fp)
if viz:
import mayavi.mlab as mlab
for i in range(10):
p1 = input_list[i]
seg = label_list[i]
fig = mlab.figure(figure=None,
bgcolor=(0.4, 0.4, 0.4),
fgcolor=None,
engine=None,
size=(500, 500))
mlab.points3d(p1[:, 0],
p1[:, 1],
p1[:, 2],
seg,
mode='point',
colormap='gnuplot',
scale_factor=1,
figure=fig)
fig = mlab.figure(figure=None,
bgcolor=(0.4, 0.4, 0.4),
fgcolor=None,
engine=None,
size=(500, 500))
mlab.points3d(p1[:, 2],
-p1[:, 0],
-p1[:, 1],
seg,
mode='point',
colormap='gnuplot',
scale_factor=1,
figure=fig)
raw_input()
def generatePlotScatterData(self):
mlab.clf(figure=self.scatter_scene.mayavi_scene)
self.points = mlab.points3d(self.x, self.y, self.__scale_z_data(), self.s, mode="point", figure=self.scatter_scene.mayavi_scene, colormap="seismic", vmin=np.min(self.s), vmax=np.max(self.s))
def plot_graph( G,
node_size = 1, node_color = (1, 1, 1), scale_node_keyword = None,
node_color_keyword = None, tube_radius = 1.4, edge_color = (1, 1, 1),
edge_color_keyword = None, edge_radius_keyword = None,
edge_colormap = 'BrBG', **kwargs):
""" 3D plot of a 3D network, this function uses a list of coordinates to visualize a network which might represent a 3D skeleton
For both, edges and nodes the size and the color can be used to visualize a parameter of the attribute dictionary,
for edges this needs to be either a a number per edge, or a sequence with the length equal to the length of coordinates
Parameters
----------
G : networkx.Graph
nodes and edges must have coordinates stored in attribute dictionary named 'x','y','z'
node_size : float
size of spheres
node_color : tuple
color of sphears
edge_color : tuple
color of tubes
scale_node_keyword : string or None
if None, constant sizes are used, otherwise
the nodes are scaled with float given in G.node[i][scale_node_keyword] could also be 'degree'
node_color_keyword: string or None
if None is given node spheres have the same color, otherwise the float value of G.node[i][node_color_keyword] is used in cobination with a defauld colormap
edge_color_keyword : string or None
if None use edgecolor, otherwise G[i][j][edge_color_keyword] is used to colorecode this value or list of values
edge_radius_keyword : string or None
if None use edge_radius, otherwise Ge[i][j][edge_radius_keyword] is used to vary the radius according to this value or list of values
Returns
-------
tube_surf : tvtk actor
actor of the edges
pts : tvtk actor
actor of the nodes
"""
print('plot graph')
from tvtk.api import tvtk
cell_array = tvtk.CellArray()
# initialize coordinates of edge segments
xstart = []
ystart = []
zstart = []
xstop = []
ystop = []
zstop = []
edge_r = [] #for radius
edge_c = [] #for color
edgecount = 0
lines = []
# initialize coordinates of nodes
xn = []
yn = []
zn = []
edge_node_n1 = []
edge_node_n2 = []
edge_angle = []
edge_length = []
edge_projection = []
for node in G.nodes():
xn.append(G.nodes[node]['x'])
yn.append(G.nodes[node]['y'])
zn.append(G.nodes[node]['z'])
# add edge segments
i = 0 #number of added segments
for n1, n2, edgedict in G.edges(data=True):
edgecount += 1
xstart.extend(edgedict['x'])
xstop.extend(edgedict['x'])
ystart.extend(edgedict['y'])
ystop.extend(edgedict['y'])
zstart.extend(edgedict['z'])
zstop.extend(edgedict['z'])
#how many segments in line?
l = len(edgedict['x'])
line = list(range(i, i + l))
line.insert(0, l)
lines.extend(line)
i += l
#print("edgedict[edge_color_keyword] is {0} \n:".format(edgedict[edge_color_keyword]))
#ori
#edgedict_temp = edgedict[edge_color_keyword]
#color by location
edgedict_temp = np.ones(len(edgedict[edge_color_keyword]))*np.mean(edgedict[edge_color_keyword])
#gradient color
#edgedict_temp = np.arange(len(edgedict[edge_color_keyword]))
#same color
#edgedict_temp = np.ones(len(edgedict[edge_color_keyword]))
#color
if edge_color_keyword is None:
edge_c += [1] * (l)
elif edge_color_keyword not in edgedict.keys():
edge_c += [1] * (l)
else:
if len(edgedict[edge_color_keyword]) == 1:
#edge_c.extend([edgedict[edge_color_keyword]] * (l))
edge_c.extend([edgedict_temp] * (l))
else:
assert len(edgedict[edge_color_keyword]) == len(edgedict['x'])
#edge_c.extend(edgedict[edge_color_keyword].tolist())
edge_c.extend(edgedict_temp.tolist())
#print("edge_c after is {0} \n:".format(edge_c))
#radius
if edge_radius_keyword in edgedict.keys():
if np.size(edgedict[edge_radius_keyword]) == 1:
edge_r.extend([edgedict[edge_radius_keyword]] * (l))
else:
assert len(edgedict[edge_radius_keyword]) == len(edgedict['x'])
edge_r.extend(edgedict[edge_radius_keyword])
else:
edge_r += [tube_radius] * (l)
#coord_n1 = (xn[n1], yn[n1], zn[n1])
#coord_n2 = (xn[n2], yn[n2], zn[n2])
#compute angle for each edge
vector_n1_n2 = (xn[n2]-xn[n1] , yn[n2]-yn[n1], zn[n2]-zn[n1])
(r_data, theta_data, phi_data) = trace_angle(xn[n2]-xn[n1] , yn[n2]-yn[n1], zn[n2]-zn[n1])
reference_vector = (1, 0, 0)
reference_vector_angle = angle_between(vector_n1_n2, reference_vector)
#print("Angle is {0}:".format(reference_vector_angle))
#print("r_data = {0}, theta_data = {1}, phi_data = {2}:\n".format(r_data, theta_data, phi_data))
#record every node and edge properities
edge_node_n1.append(n1)
edge_node_n2.append(n2)
edge_angle.append((reference_vector_angle))
edge_length.append(abs(edgedict['weight']))
edge_projection.append(r_data)
lower_limit, upper_limit = find_anomalies(np.asarray(edge_length))
print("edge_count is {0} \n:".format(edgecount))
print("edge_length is {0} \n:".format(edge_length))
print("lower_limit upper_limit is {0} {1}:\n".format(lower_limit,upper_limit))
#print("lines is {0} :\n".format(len(lines)))
#pdb.set_trace()
#pot network
fig = mlab.gcf()
fig = mlab.figure(1, bgcolor = (1,1,1))
disable_render = fig.scene.disable_render
fig.scene.disable_render = True
# ------------------------------------------------------------------------
# render edges
#
edge_c = np.array(edge_c)
if np.isinf(edge_c).any():
edge_c[np.isinf(edge_c)] = np.nan
xv = np.array([xstart])
yv = np.array([ystart])
zv = np.array([zstart])
# pdb.set_trace()
edges_src = mlab.pipeline.scalar_scatter(xv.ravel(), yv.ravel(), zv.ravel(), edge_c.ravel())
edges_src.mlab_source.dataset.point_data.add_array(np.abs(np.array(edge_r).ravel()))
edges_src.mlab_source.dataset.point_data.get_array(1).name = 'radius'
cell_array.set_cells(edgecount, lines)
edges_src.mlab_source.dataset.lines = cell_array
edges_src.mlab_source.update()
if tube_radius is not None:
#edges_src.mlab_source.dataset.point_data.vectors = np.abs(
# np.array([edge_r,edge_r,edge_r]).T)
edges_src.mlab_source.dataset.point_data.add_array(np.abs(np.array(edge_r).ravel()))
edges_src.mlab_source.dataset.point_data.get_array(1).name = 'radius'
tubes = mlab.pipeline.tube(mlab. pipeline.set_active_attribute(edges_src, point_scalars='radius'), tube_radius = tube_radius,)
#tubes.filter.radius=edge_r
tubes.filter.vary_radius = 'vary_radius_by_scalar'
else:
tubes = edges_src
tube_surf = mlab.pipeline.surface(mlab.pipeline.set_active_attribute(tubes, point_scalars='scalars'), colormap = edge_colormap, color = edge_color, **kwargs)
tubes.children[0].point_scalars_name='scalars'
if edge_color_keyword is None:
tube_surf.actor.mapper.scalar_visibility = False
else:
tube_surf.actor.mapper.scalar_visibility = True
#pdb.set_trace()
if not node_size:
return tube_surf, None
# ------------------------------------------------------------------------
# Plot the nodes
if node_size is not None:
nodes = G.nodes()
L = len(G.nodes())
if node_color_keyword is None and scale_node_keyword is None:
nodes = G.nodes()
L = len(G.nodes())
s = np.ones(len(xn))
# pdb.set_trace()
pts = mlab.points3d(xn, yn, zn, s, scale_factor=node_size,
# colormap=node_colormap,
resolution=16,
color=node_color)
else:
if scale_node_keyword is not None:
if scale_node_keyword == 'degree':
dic = G.degree(nodes)
s = []
for n in nodes:
s.append(dic[n])
else:
s= list(nx.get_node_attributes(G,scale_node_keyword).values() )
else:
s = np.ones(len(xn)) * node_size
if node_color_keyword is not None:
node_color_scalar = list(nx.get_node_attributes(
G, node_color_keyword ).values())
else:
node_color_scalar = len(xn) * [1]
pts = mlab.quiver3d(np.array(xn).ravel(), np.array(yn).ravel(), np.array(zn).ravel(),s,s,s, scalars=np.array(node_color_scalar).ravel(), mode = 'sphere', scale_factor = node_size,resolution=16)
pts.glyph.glyph_source.glyph_position = 'center'
pts.glyph.color_mode = 'color_by_scalar'
#pts.glyph.scale_mode = 'scale_by_vector'
# ------------------------------------------------------------------------
# Plot the edge index at the first node location
'''
print("edge_node_n1: {0}\n".format(edge_node_n1))
print("edge_node_n2: {0}\n".format(edge_node_n2))
print("length: {0} {1}\n".format(len(edge_angle),len(edge_node_n1)))
'''
# ------------------------------------------------------------------------
# Select edges from shared nodes
duplicate_node = duplicates(edge_node_n1)
idx = indices(edge_node_n1, duplicate_node)
#print("duplicate_node: {0}\n".format(duplicate_node))
#print("duplicate index: {0} \n".format(idx))
#connect node index in edge_node_n1 list
connect_code_idx = []
angle_thresh = 100
dis_thresh = (lower_limit + upper_limit)*0.8
for i in range(len(idx)):
#print("Index = {0}:".format(i))
#print("Value = {0}:".format(idx[duplicate_node[i]]))
angle_rec = []
for element in idx[duplicate_node[i]]:
#print("element = {0}:".format(element))
angle_rec.append(edge_angle[element])
#print(edge_node_n1[element])
#print(edge_node_n2[element])
angle_diff = np.diff(angle_rec)
#print("angle_rec = {0}".format(angle_rec))
#print("angle_diff = {0}".format(angle_diff))
for j in range(len(angle_diff)):
if angle_diff[j] < angle_thresh:
'''
print(angle_rec[j],angle_rec[j+1])
print("Index of duplicated: {0} \n".format(idx[duplicate_node[i]][j]))
print("Index of duplicated: {0} \n".format(idx[duplicate_node[i]][j+1]))
print("Index of duplicated: {0} {1} \n".format(j, j+1))
'''
connect_code_idx.append(idx[duplicate_node[i]][j])
#connect_code.append(idx[duplicate_node[i]][j+1])
#connect_node_unique = list(set(connect_code))
#connect_node_unique = connect_code
print("connect_node: {0}\n".format(connect_code_idx))
edge_node_n1_select = [j for i, j in enumerate(edge_node_n1) if i not in connect_code_idx]
edge_node_n2_select = [j for i, j in enumerate(edge_node_n2) if i not in connect_code_idx]
edge_angle_select = [j for i, j in enumerate(edge_angle) if i not in connect_code_idx]
edge_length_select = [j for i, j in enumerate(edge_length) if i not in connect_code_idx]
'''
print("edge_node_n1_select: {0}\n".format(edge_node_n1_select))
print("edge_node_n2_select: {0}\n".format(edge_node_n2_select))
print("edge_angle_select: {0}\n".format(edge_angle_select))
print("edge_length_select: {0}\n".format(edge_length_select))
print("edge_node_n1_select length: {0}\n".format(len(edge_node_n2_select)))
'''
# ------------------------------------------------------------------------
# Select edges with threshold from distance between edges
#
close_edge_nd = []
for i in range(len(edge_node_n1_select)):
idx_n1 = edge_node_n1_select[i]
idx_n2 = edge_node_n2_select[i]
#n1_coord = np.array([xn[idx_n1], yn[idx_n1], zn[idx_n1]])
#n2_coord = np.array([xn[idx_n2], yn[idx_n2], zn[idx_n2]])
n12_mid = np.mean([np.array([xn[idx_n1], yn[idx_n1], zn[idx_n1]]), np.array([xn[idx_n2], yn[idx_n2], zn[idx_n2]])], axis = 0)
for j in range(len(edge_node_n1_select)):
if j == i :
dis_P2line = 10000
#print("self_dis = {0}\n".format(dis_P2line))
else:
n1_coord_line = np.array([xn[edge_node_n1_select[j]], yn[edge_node_n1_select[j]], zn[edge_node_n1_select[j]]])
n2_coord_line = np.array([xn[edge_node_n2_select[j]], yn[edge_node_n2_select[j]], zn[edge_node_n2_select[j]]])
dis_P2line = lineseg_dist(n12_mid, n1_coord_line, n2_coord_line)
#print("n1_coord: {0}\n".format(n1_coord_line))
#print("n2_coord: {0}\n".format(n2_coord_line))
#print("n12_mid: {0}\n".format(n12_mid))
#print("distance {0} {1} = {2}\n".format(i, j, dis_P2line))
try:
dif_angle = abs(edge_angle_select[i] - edge_angle_select[j])
except (IndexError, ValueError):
dif_angle = 100
try:
if (dis_P2line < dis_thresh) and (dif_angle) < angle_thresh:
#if (dis_P2line < dis_thresh) or (edge_length_select[i] < dis_thresh):
close_edge_nd.append(i)
except (IndexError, ValueError):
dif_angle = 100
close_edge_nd = np.unique(close_edge_nd)
print("close_edge_nd index = {0}\n".format(close_edge_nd))
#combined_nd = connect_code_idx + close_edge_nd
edge_node_n1_select_final = [j for i, j in enumerate(edge_node_n1_select) if i not in close_edge_nd]
edge_node_n2_select_final = [j for i, j in enumerate(edge_node_n2_select) if i not in close_edge_nd]
edge_angle_select_final = [j for i, j in enumerate(edge_angle_select) if i not in close_edge_nd]
edge_length_select_final = [j for i, j in enumerate(edge_length_select) if i not in close_edge_nd]
'''
print("edge_node_n1_select_final: {0}\n".format(edge_node_n1_select_final))
print("edge_node_n2_select final: {0}\n".format(edge_node_n2_select_final))
'''
angle_select = []
length_select = []
projection_select = []
#render seletced nodes with text index
for i in range(len(edge_node_n1_select_final)):
idx = edge_node_n1_select_final[i]
idx_2 = edge_node_n2_select_final[i]
if i in duplicate_node:
#pts = mlab.text3d(xn[idx], yn[idx], zn[idx], str(i), scale=(2, 2, 2), color=(1, 0.0, 0.0))
print("mlab.text3d")
#pts = mlab.plot3d([xn[idx], xn[idx_2]], [yn[idx], yn[idx_2]], [zn[idx], zn[idx_2]], color = (1, 0.0, 0.0), tube_radius = 0.4)
else:
#pts = mlab.text3d(xn[idx], yn[idx], zn[idx], str(i), scale=(2, 2, 2))
print("mlab.text3d")
#pts = mlab.plot3d([xn[idx], xn[idx_2]], [yn[idx], yn[idx_2]], [zn[idx], zn[idx_2]], color = (0, 0.0, 1.0), tube_radius = 0.4)
'''
try:
pts = mlab.text3d(xn[idx], yn[idx], zn[idx], str(edge_length_select_final[i]), scale=(2, 2, 2), color=(1, 0.0, 0.0))
except (IndexError, ValueError):
pass
'''
try:
angle_select.append(edge_angle[i])
length_select.append(edge_length[i])
projection_select.append(edge_projection[i])
except IndexError:
pass
angle_select = [0 if math.isnan(x) else x for x in angle_select]
#print("angle_select = {0}\n ".format(angle_select))
#print("length_select = {0}\n ".format(length_select))
#print("projection_select = {0}\n ".format(projection_select))
print('graph rendering finished')
fig.scene.disable_render = disable_render
#return tube_surf, pts, edge_node_n1_select_final, edge_node_n2_select_final, edge_angle_select_final, edge_length_select_final, projection_select
return edge_node_n1_select_final, edge_node_n2_select_final, edge_angle_select_final, edge_length_select_final, projection_select
def show_mpl_heatmap(self, **kwargs): # pylint: disable=invalid-name,too-many-arguments,too-many-locals,too-many-statements,too-many-branches
"""
Show a heatmap of the trajectory with matplotlib.
"""
import numpy as np
from scipy import stats
try:
from mayavi import mlab
except ImportError:
raise ImportError(
'Unable to import the mayavi package, that is required to'
'use the plotting feature you requested. '
'Please install it first and then call this command again '
'(note that the installation of mayavi is quite complicated '
'and requires that you already installed the python numpy '
'package, as well as the vtk package')
from ase.data.colors import jmol_colors
from ase.data import atomic_numbers
# pylint: disable=invalid-name
def collapse_into_unit_cell(point, cell):
"""
Applies linear transformation to coordinate system based on crystal
lattice, vectors. The inverse of that inverse transformation matrix with the
point given results in the point being given as a multiples of lattice vectors
Than take the integer of the rows to find how many times you have to shift
the point back"""
invcell = np.matrix(cell).T.I
# point in crystal coordinates
points_in_crystal = np.dot(invcell, point).tolist()[0]
#point collapsed into unit cell
points_in_unit_cell = [i % 1 for i in points_in_crystal]
return np.dot(cell.T, points_in_unit_cell).tolist()
elements = kwargs.pop('elements', None)
mintime = kwargs.pop('mintime', None)
maxtime = kwargs.pop('maxtime', None)
stepsize = kwargs.pop('stepsize', None) or 1
contours = np.array(kwargs.pop('contours', None) or (0.1, 0.5))
sampling_stepsize = int(kwargs.pop('sampling_stepsize', None) or 0)
times = self.get_times()
if mintime is None:
minindex = 0
else:
minindex = np.argmax(times > mintime)
if maxtime is None:
maxindex = len(times)
else:
maxindex = np.argmin(times < maxtime)
positions = self.get_positions()[minindex:maxindex:stepsize]
try:
if self.get_attribute('units|positions') in ('bohr', 'atomic'):
bohr_to_ang = 0.52917720859
positions *= bohr_to_ang
except KeyError:
pass
symbols = self.symbols
if elements is None:
elements = set(symbols)
cells = self.get_cells()
if cells is None:
raise ValueError(
'No cell parameters have been supplied for TrajectoryData')
else:
cell = np.array(cells[0])
storage_dict = {s: {} for s in elements}
for ele in elements:
storage_dict[ele] = [np.array([]), np.array([]), np.array([])]
for iat, ele in enumerate(symbols):
if ele in elements:
for idim in range(3):
storage_dict[ele][idim] = np.concatenate(
(storage_dict[ele][idim], positions[:, iat,
idim].flatten()))
for ele in elements:
storage_dict[ele] = np.array(storage_dict[ele]).T
storage_dict[ele] = np.array([
collapse_into_unit_cell(pos, cell) for pos in storage_dict[ele]
]).T
white = (1, 1, 1)
mlab.figure(bgcolor=white, size=(1080, 720))
for i1, a in enumerate(cell):
i2 = (i1 + 1) % 3
i3 = (i1 + 2) % 3
for b in [np.zeros(3), cell[i2]]:
for c in [np.zeros(3), cell[i3]]:
p1 = b + c
p2 = p1 + a
mlab.plot3d([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]],
tube_radius=0.1)
for ele, data in storage_dict.items():
kde = stats.gaussian_kde(data, bw_method=0.15)
_x = data[0, :]
_y = data[1, :]
_z = data[2, :]
xmin, ymin, zmin = _x.min(), _y.min(), _z.min()
xmax, ymax, zmax = _x.max(), _y.max(), _z.max()
_xi, _yi, _zi = np.mgrid[xmin:xmax:60j, ymin:ymax:30j,
zmin:zmax:30j] # pylint: disable=invalid-slice-index
coords = np.vstack([item.ravel() for item in [_xi, _yi, _zi]])
density = kde(coords).reshape(_xi.shape)
# Plot scatter with mayavi
#~ figure = mlab.figure('DensityPlot')
grid = mlab.pipeline.scalar_field(_xi, _yi, _zi, density)
#~ min = density.min()
maxdens = density.max()
#~ mlab.pipeline.volume(grid, vmin=min, vmax=min + .5*(max-min))
surf = mlab.pipeline.iso_surface(grid,
opacity=0.5,
colormap='cool',
contours=(maxdens *
contours).tolist())
lut = surf.module_manager.scalar_lut_manager.lut.table.to_array()
# The lut is a 255x4 array, with the columns representing RGBA
# (red, green, blue, alpha) coded with integers going from 0 to 255.
# We modify the alpha channel to add a transparency gradient
lut[:, -1] = np.linspace(100, 255, 256)
lut[:, 0:3] = 255 * jmol_colors[atomic_numbers[ele]]
# and finally we put this LUT back in the surface object. We could have
# added any 255*4 array rather than modifying an existing LUT.
surf.module_manager.scalar_lut_manager.lut.table = lut
if sampling_stepsize > 0:
mlab.points3d(_x[::sampling_stepsize],
_y[::sampling_stepsize],
_z[::sampling_stepsize],
color=tuple(
jmol_colors[atomic_numbers[ele]].tolist()),
scale_mode='none',
scale_factor=0.3,
opacity=0.3)
mlab.view(azimuth=155, elevation=70, distance='auto')
mlab.show()
dTheta1 = dist/R1 dTheta2 = dist/R2 for h in np.linspace(0,H,ceil(H/dist)+1): allPoints.extend([[R1*cos(theta1), R1*sin(theta1), h] for theta1 in np.arange(0,2*pi,dTheta1)]) allPoints.extend([[R2*cos(theta2), R2*sin(theta2), h] for theta2 in np.arange(0,2*pi,dTheta2)]) # FLAT SURFACE allPoints.append((0,0,0)) for r in np.linspace(dist,R2,ceil(R2/dist)+1): dTheta = dist/r if r < R1: allPoints.extend([[r*cos(theta), r*sin(theta), 0] for theta in np.arange(0,2*pi,dTheta)]) if r > R1: allPoints.extend([[r*cos(theta), r*sin(theta), H] for theta in np.arange(0,2*pi,dTheta)]) allPoints = np.asarray(allPoints) mlab.points3d(allPoints[:,0],allPoints[:,1],allPoints[:,2]) return allPoints if __name__ == '__main__' : allPoints = genFiducialPC() print len(allPoints) from mayavi import mlab X = allPoints[:,0] Y = allPoints[:,1] Z = allPoints[:,2] mlab.points3d(X, Y, Z, mode='point', scale_factor=.25) mlab.show()
def mayavi_plot_surfaces(surface_plots: List[SurfacePlot],
animationDelay: int = 10,
show=True):
"""
Plot a set of surfaces (surface with style, trajectory, animation)
:param animationDelay:
:param show:
:param surface_plots:
:return:
"""
mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0, 0, 0), size=(800, 800))
mlab.clf()
# list of animated points with trajectory
animated_solids = []
time_iterations = None
for splot in surface_plots:
# if surface is
if splot.smesh is not None and splot.showSurface:
# load style
try:
style_set = mayavi_plot_surface_styles[splot.surfaceStyle]
except IndexError:
raise Exception("plot_style undefined")
# plot surface with styles
for style in style_set:
mlab.mesh(*splot.smesh, **style)
# add trajectory
if splot.trajectory is not None:
# todo fix error when animation does not work when showTrajectory is False
if splot.showTrajectory:
mlab.plot3d(splot.trajectory[:, 0],
splot.trajectory[:, 1],
splot.trajectory[:, 2],
color=splot.trajectoryColor,
tube_radius=0.01)
if splot.showSolid:
solid = mlab.points3d(*splot.trajectory[0],
color=splot.solidColor,
scale_factor=0.1)
if splot.animate:
# get time iterations
time_iterations = splot.trajectory.shape[0]
# animate trajectory
animated_solids.append((solid, splot.trajectory))
mlab.orientation_axes()
# add animation
if animated_solids:
@mlab.animate(delay=animationDelay)
def anim():
while True:
for i in range(0, time_iterations, int(splot.timeFactor)):
for solid, traj in animated_solids:
x, y, z = traj[i]
solid.mlab_source.set(x=x, y=y, z=z)
yield
if show:
anim()
elif show:
mlab.show()
data = mat2['sceneVox']
x, y, z = np.where(data == 1)
X = final[0]
Y = final[1]
Z = final[2]
coords = []
for idx, x_val in enumerate(x):
x_val = x[idx]
y_val = y[idx]
z_val = z[idx]
x_coord = X[x_val][y_val][z_val]
y_coord = Y[x_val][y_val][z_val]
z_coord = Z[x_val][y_val][z_val]
coord = np.array([x_coord, y_coord, z_coord])
coords.append(coord)
coords = np.array(coords)
mlab.points3d(*tuple(coords.T), color=(0, 1, 1))
for idx, bbox in enumerate(mat2['modelBboxes'][0]):
new_bbox = []
bbox = bbox.T
arr = np.array([x for x in itertools.product(*tuple(bbox.T))])
mlab.points3d(*tuple(arr.T),
color=(np.random.random_sample(), np.random.random_sample(),
np.random.random_sample()))
mlab.show()