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 show(cuds):
""" Show the cuds objects using the default visualisation.
Parameters
----------
cuds :
A top level cuds object (e.g. a mesh). The method will detect
the type of object and create the appropriate visualisation.
"""
if isinstance(cuds, (ABCMesh, ABCParticles, ABCLattice)):
source = CUDSSource(cuds=cuds)
else:
msg = 'Provided object {} is not of any known cuds type'
raise TypeError(msg.format(type(cuds)))
modules = default_module(source)
# ensure that a new scene is made
mayavi_engine = mlab.get_engine()
mayavi_engine.new_scene()
# add source
mayavi_engine.add_source(source)
# add default modules
for module in modules:
mayavi_engine.add_module(module)
mlab.show()
def setUp(self):
# set up source
sgrid = datasets.generateStructuredGrid()
source = VTKDataSource(data=sgrid)
self.engine = mlab.get_engine()
# set up scene, first scene is empty
# second scene has the settings we want to restore
for _ in range(2):
fig = mlab.figure()
fig.scene.off_screen_rendering = True
# add source
self.engine.add_source(source)
# add more modules
self.engine.add_module(IsoSurface())
self.engine.add_module(Text3D())
self.modules = source.children[0].children
# set camera
self.view = (25., 14., 20., [0., 0., 2.5])
mlab.view(*self.view)
# save the visualisation
self.temp_dir = tempfile.mkdtemp()
self.filename = os.path.join(self.temp_dir, "test_vis.mv2")
self.engine.save_visualization(self.filename)
# save the scene as an image for comparison later
self.ref_saved_filename = os.path.join(self.temp_dir, "ref_saved.png")
mlab.savefig(self.ref_saved_filename)
def __init__(self, engine_name="", engine=None, mayavi_engine=None):
'''
Parameters
----------
engine_name : str
Name of the Simphony Modeling Engine wrapper
engine : ABCModelingEngine
Simphony Modeling Engine wrapper
mayavi_engine : mayavi.api.engine
Default to be mayavi.mlab.get_engine()
'''
# Traits initialisation
HasTraits.__init__(self)
if mayavi_engine is None:
# Standalone Mayavi Engine
mayavi_engine = mlab.get_engine()
else:
mayavi_engine = mayavi_engine
# Add panels
self.panels = TabbedPanelCollection.create(
add_engine=AddEnginePanel(engine_manager=self),
add_source=AddSourcePanel(engine_name=self.engine_name,
engine=self.engine,
mayavi_engine=mayavi_engine),
run_and_animate=RunAndAnimatePanel(engine=self.engine,
mayavi_engine=mayavi_engine))
if engine and engine_name:
self.add_engine(engine_name, engine)
def mayavi_scraper(block, block_vars, gallery_conf):
"""Scrape Mayavi images.
Parameters
----------
block : tuple
A tuple containing the (label, content, line_number) of the block.
block_vars : dict
Dict of block variables.
gallery_conf : dict
Contains the configuration of Sphinx-Gallery
Returns
-------
rst : str
The ReSTructuredText that will be rendered to HTML containing
the images. This is often produced by
:func:`sphinx_gallery.gen_rst.figure_rst`.
"""
from mayavi import mlab
image_path_iterator = block_vars['image_path_iterator']
image_paths = list()
e = mlab.get_engine()
for scene, image_path in zip(e.scenes, image_path_iterator):
mlab.savefig(image_path, figure=scene)
# make sure the image is not too large
scale_image(image_path, image_path, 850, 999)
image_paths.append(image_path)
mlab.close(all=True)
return figure_rst(image_paths, gallery_conf['src_dir'])
def picker_callback(picker):
""" Picker callback: this get called when on pick events.
"""
engine = mlab.get_engine()
scene = engine.scenes[0]
vtk_scene = scene.scene
interactor = vtk_scene.interactor
original_mouse_position = interactor.event_position
render_window = vtk_scene.render_window
#print(original_mouse_position)
while(True):
current_mouse_posiiton = interactor.event_position
print(current_mouse_posiiton)
wx.Yield()
num_shaded_pixels = (abs(current_mouse_posiiton[0] - original_mouse_position[0]) + 1)*(abs(current_mouse_posiiton[1] - original_mouse_position[1]) + 1)
render_window.set_pixel_data(original_mouse_position[0], original_mouse_position[1], current_mouse_posiiton[0], current_mouse_posiiton[1], [1,1,1]*num_shaded_pixels, 1)
if picker.actor in red_glyphs.actor.actors:
# Find which data point corresponds to the point picked:
# we have to account for the fact that each data point is
# represented by a glyph with several points
point_id = picker.point_id/glyph_points.shape[0]
# If the no points have been selected, we have '-1'
if point_id != -1:
# Retrieve the coordinnates coorresponding to that data
# point
x, y, z = x1[point_id], y1[point_id], z1[point_id]
def test_figure(self):
""" Various tests for mlab.figure().
"""
# Test when specifying figure instances
f1 = mlab.figure()
e = mlab.get_engine()
self.assertTrue(e.current_scene is f1)
f2 = mlab.figure()
self.assertTrue(e.current_scene is f2)
mlab.figure(f1)
self.assertTrue(e.current_scene is f1)
# Test when specifying figure numbers
f1 = mlab.figure(3)
self.assertTrue(e.current_scene is f1)
f2 = mlab.figure(4)
self.assertTrue(e.current_scene is f2)
mlab.figure(3)
self.assertTrue(e.current_scene is f1)
# Test when specifying figure names
f1 = mlab.figure('Test 1')
self.assertTrue(e.current_scene is f1)
f2 = mlab.figure('Test 2')
self.assertTrue(e.current_scene is f2)
mlab.figure('Test 1')
self.assertTrue(e.current_scene is f1)
def tearDown(self):
# Check that the NullEngine is still the mlab engine
if not mlab.get_engine() is self.e:
raise AssertionError, \
"The NullEngine has been overridden"
engine_manager.current_engine = None
# Unregistering the engine, to avoid side-effects between tests
self.e.stop()
registry.unregister_engine(self.e)
def new_func(test_case):
try:
func(test_case)
finally:
num_scenes = len(mlab.get_engine().scenes)
test_case.assertNotEqual(num_scenes, 0,
"No scene is opened")
# close everything
mlab.close(all=True)
def close():
"""Close the scene."""
f = mlab.gcf()
e = mlab.get_engine()
e.window.workbench.prompt_on_exit = False
e.window.close()
mlab.options.backend = 'auto'
# Hack: on Linux the splash screen does not go away so we force it.
GUI.invoke_after(500, e.window.workbench.application.gui.stop_event_loop)
def serve_udp(engine=None, port=9007, logto=sys.stdout):
"""Serve the `M2UDP` protocol using the given `engine` on the
specified `port` logging messages to given `logto` which is a
file-like object. This function will block till the service is
closed. There is no need to call `mlab.show()` after or before
this. The Mayavi UI will be fully responsive.
**Parameters**
:engine: Mayavi engine to use. If this is `None`,
`mlab.get_engine()` is used to find an appropriate engine.
:port: int: port to serve on.
:logto: file : File like object to log messages to. If this is
`None` it disables logging.
**Examples**
Here is a very simple example::
from mayavi import mlab
from mayavi.tools import server
mlab.test_plot3d()
server.serve_udp()
Test it like so::
import socket
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.bind(('', 9008))
s.sendto('camera.azimuth(10)', ('', 9007))
**Warning**
Data sent is exec'd so this is a security hole.
"""
from mayavi import mlab
e = engine or mlab.get_engine()
# Setup the protocol with the right attributes.
proto = M2UDP()
proto.engine = e
proto.scene = e.current_scene.scene
proto.mlab = mlab
if logto is not None:
log.startLogging(logto)
log.msg('Serving Mayavi2 UDP server on port', port)
log.msg('Using Engine', e)
# Register the running wxApp.
reactor.registerWxApp(wx.GetApp())
# Listen on port 9007 using above protocol.
reactor.listenUDP(port, proto)
# Run the server + app. This will block.
reactor.run()
def tearDown(self):
# Check that the NullEngine was not set as the default mlab engine.
if not mlab.get_engine() is self._non_null_engine:
raise AssertionError("The NullEngine has overridden the default one")
engine_manager.current_engine = None
# Unregistering all unused engines.
registry.unregister_engine(self._non_null_engine)
for engine in list(registry.engines):
registry.unregister_engine(engine)
def test_save_load_visualization_with_mlab(self):
# test mlab.get_engine
engine = mlab.get_engine()
try:
self.check_save_load_visualization(engine)
finally:
mlab.clf()
mlab.close(all=True)
def __init__(self, hex_list, np_file='/tmp/magnetic_ground_truth.np', robot_height=40, width=800, height=600,
start_point=(0, 0, 0), message='experiment default message...'):
self.debug = False
self.animator = None
self.movement_mode = 0
self.start_point = start_point
self.robot_height = robot_height
self.message = message
self.start_time = int(time.time() * 1000)
self.width = width
self.height = height
self.f = mlab.figure(size=(self.width, self.height))
visual.set_viewer(self.f)
v = mlab.view(270, 180)
#print v
engine = mlab.get_engine()
self.s = engine.current_scene
self.s.scene.interactor.add_observer('KeyPressEvent', self.keypress_callback)
self.robots = []
colors = list(PathBatterySimulator.color_codes)
for key, local_hex_list in sorted(hex_list['internal_routes'].items()):
color = colors.pop(0)
ball = visual.sphere(color=color, radius=PathBatterySimulator.ball_radius)
ball.x = self.start_point[0]
ball.y = self.start_point[1]
ball.z = self.start_point[2]
r, g, b = color
rt = r + (0.25 * (1 - r))
gt = g + (0.25 * (1 - g))
bt = b + (0.25 * (1 - b))
curve_color = (rt, gt, bt)
curve = visual.curve(color=curve_color, radius=PathBatterySimulator.curve_radius)
r_ball = RobotBall(key, local_hex_list, hex_list['external_routes'][key], ball, curve)
self.robots.append(r_ball)
x = np.linspace(0, self.width, 1)
y = np.linspace(0, self.height, 1)
z = np.loadtxt(np_file)
z *= 255.0/z.max()
mlab.surf(x, y, z)
self.master_cmd = MasterCommand(self.robots)
def tearDown(self):
# Check that the NullEngine is still the mlab engine
current_engine = mlab.get_engine()
engine_overridden = current_engine is not self.e
engine_manager.current_engine = None
self.e.stop()
registry.unregister_engine(self.e)
if engine_overridden:
current_engine.stop()
registry.unregister_engine(current_engine)
raise AssertionError("The NullEngine has been overridden")
def view(src):
""" Open up a mayavi scene and display the dataset in it.
"""
from mayavi import mlab
mayavi = mlab.get_engine()
fig = mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0, 0, 0),)
mayavi.add_source(src)
mlab.pipeline.surface(src, opacity=0.1)
mlab.pipeline.surface(mlab.pipeline.extract_edges(src),
color=(0, 0, 0), )
def test_add_dataset_works_with_vtk_datasets(self):
# Given
pd = vtk.vtkPolyData()
# When
mlab.pipeline.add_dataset(pd)
# Then
e = mlab.get_engine()
src = e.scenes[0].children[0]
from mayavi.sources.vtk_data_source import VTKDataSource
self.assertTrue(isinstance(src, VTKDataSource))
self.assertEqual(tvtk.to_vtk(src.data), pd)
def view(dataset):
"""
Open up a mayavi scene and display the cubeset in it.
"""
engine = mlab.get_engine()
#fig = mlab.figure(bgcolor=(0, 0, 0), fgcolor=(1, 1, 1),
# figure=dataset.class_name[3:])
src = VTKDataSource(data=dataset)
engine.add_source(src)
# TODO : make some cubes more redish to show some "activity"
mlab.pipeline.surface(src, colormap='gray')
def test_init_default_mayavi_engine(self):
# given
engine = DummyEngine()
manager = EngineManagerStandaloneUI("test", engine)
# then
with self.assertRaises(AttributeError):
self.mayavi_engine
for panel in manager.panels:
if hasattr(panel, "mayavi_engine"):
self.assertEqual(panel.mayavi_engine, mlab.get_engine())
def save_figures(image_path, fig_count, gallery_conf):
"""Save all open matplotlib figures of the example code-block
Parameters
----------
image_path : str
Path where plots are saved (format string which accepts figure number)
fig_count : int
Previous figure number count. Figure number add from this number
gallery_conf : dict
Contains the configuration of Sphinx-Gallery
Returns
-------
images_rst : str
rst code to embed the images in the document
fig_num : int
number of figures saved
"""
figure_list = []
for fig_num in plt.get_fignums():
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
fig = plt.figure(fig_num)
kwargs = {}
to_rgba = matplotlib.colors.colorConverter.to_rgba
for attr in ['facecolor', 'edgecolor']:
fig_attr = getattr(fig, 'get_' + attr)()
default_attr = matplotlib.rcParams['figure.' + attr]
if to_rgba(fig_attr) != to_rgba(default_attr):
kwargs[attr] = fig_attr
current_fig = image_path.format(fig_count + fig_num)
fig.savefig(current_fig, **kwargs)
figure_list.append(current_fig)
if gallery_conf.get('find_mayavi_figures', False):
from mayavi import mlab
e = mlab.get_engine()
last_matplotlib_fig_num = fig_count + len(figure_list)
total_fig_num = last_matplotlib_fig_num + len(e.scenes)
mayavi_fig_nums = range(last_matplotlib_fig_num + 1, total_fig_num + 1)
for scene, mayavi_fig_num in zip(e.scenes, mayavi_fig_nums):
current_fig = image_path.format(mayavi_fig_num)
mlab.savefig(current_fig, figure=scene)
# make sure the image is not too large
scale_image(current_fig, current_fig, 850, 999)
figure_list.append(current_fig)
mlab.close(all=True)
return figure_rst(figure_list, gallery_conf['src_dir'])
def save_figures(image_path, fig_count, gallery_conf):
"""Save all open matplotlib figures of the example code-block
Parameters
----------
image_path : str
Path where plots are saved (format string which accepts figure number)
fig_count : int
Previous figure number count. Figure number add from this number
gallery_conf : dict
Contains the configuration of Sphinx-Gallery
Returns
-------
images_rst : str
rst code to embed the images in the document
fig_num : int
number of figures saved
"""
figure_list = []
for fig_num in plt.get_fignums():
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
fig = plt.figure(fig_num)
kwargs = {}
to_rgba = matplotlib.colors.colorConverter.to_rgba
for attr in ['facecolor', 'edgecolor']:
fig_attr = getattr(fig, 'get_' + attr)()
default_attr = matplotlib.rcParams['figure.' + attr]
if to_rgba(fig_attr) != to_rgba(default_attr):
kwargs[attr] = fig_attr
current_fig = image_path.format(fig_count + fig_num)
fig.savefig(current_fig, **kwargs)
figure_list.append(current_fig)
if gallery_conf.get('find_mayavi_figures', False):
from mayavi import mlab
e = mlab.get_engine()
last_matplotlib_fig_num = fig_count + len(figure_list)
total_fig_num = last_matplotlib_fig_num + len(e.scenes)
mayavi_fig_nums = range(last_matplotlib_fig_num + 1, total_fig_num + 1)
for scene, mayavi_fig_num in zip(e.scenes, mayavi_fig_nums):
current_fig = image_path.format(mayavi_fig_num)
mlab.savefig(current_fig, figure=scene)
# make sure the image is not too large
scale_image(current_fig, current_fig, 850, 999)
figure_list.append(current_fig)
mlab.close(all=True)
return figure_rst(figure_list, gallery_conf['src_dir'])
def _startup(self):
"""
Open a VTK file and return the engine that is visualizing it.
Whenever the figure is shown, startup needs to be called again,
similar to how showing a matplotlib figure will destroy the instance.
"""
assert self.fid, "File ID must be specified before generating figure"
# Instantiate mlab
self.fig = mlab.figure(size=self.figsize)
self.engine = mlab.get_engine()
self.engine.scenes[0].scene.background = colors["w"]
self.vtkfr = self.engine.open(self.fid)
def setCam(cam=None):
if cam:
e = mlab.get_engine()
c = e.current_scene
c.scene.camera.position = cam['position']
c.scene.camera.focal_point = cam['focal_point']
c.scene.camera.view_angle = cam['view_angle']
c.scene.camera.view_up = cam['view_up']
c.scene.camera.clipping_range = cam['clipping_range']
c.scene.camera.compute_view_plane_normal()
c.scene.render()
def imayavi_clear_data(scenes=None):
"""Workaround for Mayavi / VTK memory leak
This is needed when Mayavi/VTK keeps a reference to source data
when you would expect it to be freed like on a call to `mlab.clf()`
or when removing sources from the pipeline.
Note:
This must be called when the pipeline still has the source, so
before a call to `mlab.clf()`, etc.
1. Set release_data_flag on all sources' data
2. Remove reference to the data
3. Remove the data source
Args:
scene (None, mayavi.core.scene.Scene, or 'all'): if None, gets
current scene; if Scene object, just that one; if 'all',
act on all scenes in the current engine. Can also be a list
of Scene objects
"""
if scenes is None:
scenes = [mlab.get_engine().current_scene]
elif scenes == "all":
scenes = mlab.get_engine().scenes
if not isinstance(scenes, (list, tuple)):
scenes = [scenes]
if all(s is None for s in scenes):
return
for s in scenes:
s.stop()
for child in list(s.children):
imayavi_remove_source(child)
s.start()
return
def PlotMeshwOverlay(v, f, y, a):
# plot a surface (vert & faces) as a 3D patch (trisurf) with overlay
fig = mlab.figure(1, bgcolor=(0, 0, 0))
pts = mlab.triangular_mesh(v[:, 0],
v[:, 1],
v[:, 2],
f,
scalars=y[:, 0],
opacity=a)
mlab.get_engine().scenes[0].scene.x_plus_view()
mlab.view(0., 0.)
mlab.colorbar(title="overlay")
# limits = [v.min(), v.max()]
# #cmap = 'coolwarm'
# fig = plt.figure()
# ax = fig.add_subplot(111, projection='3d', xlim=limits, ylim=limits)
# ax.view_init(elev=0, azim=0)
# ax.set_axis_off()
#
# #cmp = cm.jet(y)
# #cmp=np.squeeze(cmp)
#
# collec = ax.plot_trisurf(v[:, 0], v[:, 1], v[:, 2],
# triangles=f, linewidth=0.,
# antialiased=False,
# cmap=cm.jet,alpha=a)
# #shade=False,
# #yf = y[f]
# #colors = np.amax(yf,axis=1)
#
# colors = np.mean(y[f], axis=1) # map vertex cols to face cols!
# newy=colors[:,0]
# collec.set_array(newy)
# ax.add_collection(collec)
# fig.colorbar(collec, ax=ax)
return pts, fig #ax, collec, fig
def test_test_backend_clf(self):
"""Test if setting the backend to 'test' works."""
mlab.options.backend = 'test'
mlab.test_contour3d()
e = mlab.get_engine()
self.assertEqual(len(e.scenes), 1)
self.assertEqual(len(e.scenes[0].children), 1)
mlab.clf()
self.assertEqual(len(e.scenes), 1)
self.assertEqual(len(e.scenes[0].children), 0)
mlab.pipeline.open(get_example_data('cube.vti'))
mlab.clf()
self.assertEqual(len(e.scenes), 1)
self.assertEqual(len(e.scenes[0].children), 0)
def save_figures(image_path, fig_count, gallery_conf):
"""Save all open matplotlib figures of the example code-block
Parameters
----------
image_path : str
Path where plots are saved (format string which accepts figure number)
fig_count : int
Previous figure number count. Figure number add from this number
Returns
-------
list of strings containing the full path to each figure
"""
figure_list = []
fig_managers = matplotlib._pylab_helpers.Gcf.get_all_fig_managers()
for fig_mngr in fig_managers:
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
fig = plt.figure(fig_mngr.num)
kwargs = {}
to_rgba = matplotlib.colors.colorConverter.to_rgba
for attr in ["facecolor", "edgecolor"]:
fig_attr = getattr(fig, "get_" + attr)()
default_attr = matplotlib.rcParams["figure." + attr]
if to_rgba(fig_attr) != to_rgba(default_attr):
kwargs[attr] = fig_attr
current_fig = image_path.format(fig_count + fig_mngr.num)
fig.savefig(current_fig, **kwargs)
figure_list.append(current_fig)
if gallery_conf.get("find_mayavi_figures", False):
from mayavi import mlab
e = mlab.get_engine()
last_matplotlib_fig_num = len(figure_list)
total_fig_num = last_matplotlib_fig_num + len(e.scenes)
mayavi_fig_nums = range(last_matplotlib_fig_num, total_fig_num)
for scene, mayavi_fig_num in zip(e.scenes, mayavi_fig_nums):
current_fig = image_path.format(mayavi_fig_num)
mlab.savefig(current_fig, figure=scene)
# make sure the image is not too large
scale_image(current_fig, current_fig, 850, 999)
figure_list.append(current_fig)
mlab.close(all=True)
return figure_list
def save_figures(image_path, fig_count, gallery_conf):
"""Save all open matplotlib figures of the example code-block
Parameters
----------
image_path : str
Path where plots are saved (format string which accepts figure number)
fig_count : int
Previous figure number count. Figure number add from this number
Returns
-------
list of strings containing the full path to each figure
"""
figure_list = []
fig_managers = matplotlib._pylab_helpers.Gcf.get_all_fig_managers()
for fig_mngr in fig_managers:
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
fig = plt.figure(fig_mngr.num)
kwargs = {}
to_rgba = matplotlib.colors.colorConverter.to_rgba
for attr in ['facecolor', 'edgecolor']:
fig_attr = getattr(fig, 'get_' + attr)()
default_attr = matplotlib.rcParams['figure.' + attr]
if to_rgba(fig_attr) != to_rgba(default_attr):
kwargs[attr] = fig_attr
current_fig = image_path.format(fig_count + fig_mngr.num)
fig.savefig(current_fig, **kwargs)
figure_list.append(current_fig)
if gallery_conf.get('find_mayavi_figures', False):
from mayavi import mlab
e = mlab.get_engine()
last_matplotlib_fig_num = len(figure_list)
total_fig_num = last_matplotlib_fig_num + len(e.scenes)
mayavi_fig_nums = range(last_matplotlib_fig_num, total_fig_num)
for scene, mayavi_fig_num in zip(e.scenes, mayavi_fig_nums):
current_fig = image_path.format(mayavi_fig_num)
mlab.savefig(current_fig, figure=scene)
# make sure the image is not too large
scale_image(current_fig, current_fig, 850, 999)
figure_list.append(current_fig)
mlab.close(all=True)
return figure_list
def run_mlab_file(filename, image_file):
## XXX: Monkey-patch mlab.show, so that we keep control of the
## the mainloop
old_show = mlab.show
def my_show(func=None):
pass
mlab.show = my_show
mlab.clf()
e = mlab.get_engine()
e.close_scene(mlab.gcf())
execfile(filename, {'__name__': '__main__'})
mlab.savefig(image_file)
size = mlab.gcf().scene.get_size()
for scene in e.scenes:
e.close_scene(scene)
mlab.show = old_show
def view(src):
""" Open up a mayavi scene and display the dataset in it.
"""
from mayavi import mlab
mayavi = mlab.get_engine()
fig = mlab.figure(
bgcolor=(1, 1, 1),
fgcolor=(0, 0, 0),
)
mayavi.add_source(src)
mlab.pipeline.surface(src, opacity=0.1)
mlab.pipeline.surface(
mlab.pipeline.extract_edges(src),
color=(0, 0, 0),
)
def make_plot(phi_min = 0, phi_max = 2.*np.pi):
vmec_filename = '/home/srh112/code/python/h1_eq_generation/results7/kh0.100-kv1.000fixed/wout_kh0.100-kv1.000fixed.nc'
f = mlab.figure(1, fgcolor=(0, 0, 0), bgcolor=(1, 1, 1))
from mayavi.sources.poly_data_reader import PolyDataReader
a = PolyDataReader()
a.initialize("/home/srh112/code/python/python-h1/h1/h1model/helical.stl")
e = mlab.get_engine()
e.add_source(a)
from mayavi.modules.surface import Surface
s = Surface()
e.add_module(s)
x, y, z, B = plot_vmec(vmec_filename=vmec_filename, phi_min = phi_min, phi_max = phi_max)
#x, y, z, B = extract_VMEC_surface_data(vmec_filename, s_ind=-1, phi_min = phi_min, phi_max = phi_max)
#pts = mlab.mesh(x[:,:], y[:,:], z[:,:], opacity = 1.0, scalars = B, colormap = 'hot', representation='surface')
#plot the TFC's
include_coils = range(5,27,1)
tfc_thickness=0.075;tfc_width=0.15; tfc_radius = 0.383
#tfc_mesh_props = {'opacity':0.3,'color':(1,0.,0.)}
tfc_mesh_props = {'opacity':1.0,'color':(0.5,0.5,0.5)}
plot_tfc(include_coils, tfc_thickness=0.075, tfc_width=0.15, tfc_radius = 0.383, tfc_mesh_props = None)
pfc_mesh_props = {'opacity':1.,'color':(0.5,0.5,0.5)}
plot_pfc(pfc_thickness = 0.11, pfc_width = 0.11, pfc_radius = 1.0, pfc_mesh_props=pfc_mesh_props)
ovc_mesh_props = {'opacity':1,'color':(255/255., 255/255., 51/255.)}
#plot_ovc(ovc_mesh_props=ovc_mesh_props)
ivc_mesh_props = {'opacity':1,'color':(0.5,0.5,0.5)}
#plot_ivc(ivc_mesh_props=ivc_mesh_props)
#Plot the HMA and poloidal Mirnov arrays as cubes joined by a line
show_hma = 1; show_pol_array1 = 1; show_pol_array2 = 1
HMA_x, HMA_y, HMA_z, pol_array1_x, pol_array1_y, pol_array1_z, pol_array2_x, pol_array2_y, pol_array2_z = mirnov_locations()
if show_hma:
mlab.plot3d(HMA_x, HMA_y, HMA_z,line_width=1,tube_radius=0.01)
#mlab.points3d(HMA_x, HMA_y, HMA_z, scale_mode='none', scale_factor = 0.04, color=(0.5,0.5,0.5),mode='cube')
mlab.points3d(HMA_x, HMA_y, HMA_z, scale_mode='none', scale_factor = 0.04, color=(0,0.,1),mode='cube')
if show_pol_array1:
mlab.plot3d(pol_array1_x, pol_array1_y, pol_array1_z,line_width=1,tube_radius=0.02)
mlab.points3d(pol_array1_x, pol_array1_y, pol_array1_z, scale_mode='none', scale_factor = 0.04, color=(0,1,0),mode='cube')
if show_pol_array2:
mlab.plot3d(pol_array2_x, pol_array2_y, pol_array2_z,line_width=1,tube_radius=0.02)
mlab.points3d(pol_array2_x, pol_array2_y, pol_array2_z, scale_mode='none', scale_factor = 0.04, color=(0.,1.,0.),mode='cube')
return x, y, z
def run_mlab_file(filename, image_file):
## XXX: Monkey-patch mlab.show, so that we keep control of the
## the mainloop
old_show = mlab.show
def my_show(func=None):
pass
mlab.show = my_show
mlab.clf()
e = mlab.get_engine()
e.close_scene(mlab.gcf())
exec(compile(open(filename).read(), filename, 'exec'),
{'__name__': '__main__'})
mlab.savefig(image_file)
size = mlab.gcf().scene.get_size()
for scene in e.scenes:
e.close_scene(scene)
mlab.show = old_show
def drawScalarCutPlane(planeNorm=(0,0,0), planeOri=(0,0,0), filterNorm=None,
tubing=False, viewControls=True, engine=None, scene=None):
"""Draws scalar cut plane in a Mayavi figure
Parameters
----------
planeNorm : 3-tuple
Normal vector to the scalar cut plane
planeOri : 3-tuple
The coordinates in world coordinates where the plane should be placed
filterNorm : 3-tuple
?? (default is `None`) if `None` then this is assigned to the `planeNorm`
tubing : boolean
Whether or not to put a tube surrounding the plane. Default is `False` i.e. no tubing
viewControls: boolean
Enable (if True, which is default behavior) or disable GUI based control of the cut-plane
engine : Mayavi Engine
Default = None
scene : Mayavi scene object
Default = None
Returns
-------
scp : scalar cut plane object
See also `drawScalarCutPlaneUsingPipeline`
"""
if not engine:
engine = mlab.get_engine()
if not scene:
scene = engine.scenes[0]
scp = ScalarCutPlane()
engine.add_module(scp)
scp.implicit_plane.widget.origin = planeOri
scp.implicit_plane.normal = planeNorm
if filterNorm:
scp.warp_scalar.filter.normal = filterNorm
else:
scp.warp_scalar.filter.normal = planeNorm # I'm currently not sure what it does ... it generally follows the normal
scp.implicit_plane.widget.enabled = viewControls
if viewControls:
scp.implicit_plane.widget.tubing = tubing
return scp
def showvolume(Vin, currentfigurenum):
mlab.figure(currentfigurenum, bgcolor=(1, 1, 1), fgcolor=(1, 1, 1))
mlab.clf()
p = mlab.contour3d(Vin.vol, color=(1, 0, 0))
mlab.text(0.05,
0.95,
'Please close the window to continue calculations.',
color=(0, 0, 0),
width=0.9)
mlab.text(0.3, 0.05, 'Rotate using click&drag', color=(0, 0, 0), width=0.4)
c_scene = mlab.get_engine().current_scene
# c_scene.scene.light_manager.light_mode = 'vtk';
c_scene.scene.camera.position = [0, 0, -128]
c_scene.scene.camera.view_up = [-1, 0, 0]
c_scene.scene.render()
mlab.show()
return p
def vtk_SaveAllFigsAsFiles(filename, **kwargs):
tag = kwargs.get('tag', False)
figext = kwargs.get('format', 'png')
savedir = kwargs.get('dir', '.')
figs = kwargs.get('figs', None) # list of figures number
verbose = kwargs.get('verbose', False)
scale = kwargs.get('scale', 1.0)
#if not isMayavi():
#print('Needs Mayavi ...')
#return
if tag:
Softname = 'Python'
V = sys.version.split(' ')[0]
Release = V.replace('.', '')
Tag = Softname + Release
# FullTag=Tag+'_Mayavi'+mayavi.__version__.replace('.','')
if not os.path.exists(savedir):
os.makedirs(savedir)
set_colorbars_option('label_text_property', color=(0, 0, 0))
set_colorbars_option('title_text_property', color=(0, 0, 0))
scenes = mlab.get_engine().scenes
for i in range(len(scenes)):
sc = scenes[i]
fig = mlab.figure(sc) # set as current
nfig = sc.name[13::] # sc.name is 'Mayavi Scene <NUM>'
#nfig=figs[i]
if tag:
File = savedir + os.path.sep + filename + '_fig' + nfig + '_' + Tag + '.' + figext
else:
File = savedir + os.path.sep + filename + '_fig' + nfig + '.' + figext
old_bg = fig.scene.background
fig.scene.background = (1, 1, 1)
#old_fg=fig.scene.foreground
#fig.scene.foreground=(0, 0, 0)
if verbose:
print(' Save Mayavi Scene %s in %s' % (nfig, File))
mlab.savefig(File, magnification=scale)
fig.scene.background = old_bg
def make_base_sphere(seed=None):
mlab.figure(
figure=None, fgcolor=None, engine=None, size=(300, 350))
cap = threed.make_cap(
r=radius, cap_angle=np.pi, direction=np.array([0, 0, 1]),
usteps=50, vsteps=50)
mlab.mesh(*cap, representation='surface', color=white)
engine = mlab.get_engine()
from mayavi.modules.outline import Outline
o = Outline()
engine.add_module(o)
from mayavi.modules.axes import Axes
ax = Axes()
engine.add_module(ax)
if seed is not None:
np.random.seed(seed)
def get_colorbars(**kwargs):
enable = kwargs.pop('enable', None)
figure = kwargs.pop('figure', None)
if figure is None:
F = find_Colors_and_legends(mlab.get_engine().scenes, [])
else:
assert (isinstance(figure,
mayavi.core.scene.Scene)) # use mlab.figure(1)
F = find_Colors_and_legends(figure, [])
CBs = []
for f in F:
slm = f.trait_get()['scalar_lut_manager']
if enable is None:
CBs.append(slm)
elif slm.show_scalar_bar == enable:
CBs.append(slm)
slm = f.trait_get()['vector_lut_manager']
if enable is None:
CBs.append(slm)
elif slm.show_scalar_bar == enable:
CBs.append(slm)
return CBs
def clear_data(figures=None):
"""Workaround for Mayavi / VTK memory leak
This is needed when Mayavi/VTK keeps a reference to source data
when you would expect it to be freed like on a call to `mlab.clf()`
or when removing sources from the pipeline.
Note:
This must be called when the pipeline still has the source, so
before a call to `mlab.clf()`, etc.
1. Set release_data_flag on all sources' data
2. Remove reference to the data
3. Remove the data source
Args:
figures (None, mayavi.core.scene.Scene, or 'all'): if None,
gets current scene; if Scene object, just that one; if
'all', act on all scenes in the current engine. Can also be
a list of Scene objects
"""
if figures is None:
figures = [mlab.gcf()]
elif figures == "all":
figures = mlab.get_engine().scenes
if not isinstance(figures, (list, tuple)):
figures = [figures]
if all(fig is None for fig in figures):
return
for fig in figures:
# # fig stop / start kills mayavi now, not sure why
# fig.stop()
for child in list(fig.children):
remove_source(child)
# fig.start()
return
def mlab_view(dataset):
fig = mlab.figure(bgcolor=(1, 1, 1),
fgcolor=(0, 0, 0),
figure=dataset.class_name[3:])
engine = mlab.get_engine()
scene = engine.scenes[0]
scene.scene.z_plus_view()
src = mlab.pipeline.add_dataset(dataset)
warp_vector = mlab.pipeline.warp_vector(src)
surf = mlab.pipeline.surface(warp_vector)
etc = ExtractTensorComponents()
engine.add_filter(etc, warp_vector)
surface2 = Surface()
engine.add_filter(surface2, etc)
etc.filter.scalar_mode = 'component'
lut = etc.children[0]
lut.scalar_lut_manager.show_scalar_bar = True
lut.scalar_lut_manager.show_legend = True
lut.scalar_lut_manager.scalar_bar.height = 0.8
lut.scalar_lut_manager.scalar_bar.width = 0.17
lut.scalar_lut_manager.scalar_bar.position = np.array([0.82, 0.1])
def visStack(v, opacity=.5, color=(1, 0, 0), mode=''):
if mode != 'same':
mlab.figure(bgcolor=(1, 1, 1))
s = mlab.get_engine() # Returns the running mayavi engine.
scene = s.current_scene # Returns the current scene.
scene.scene.disable_render = True # for speed
origion = [0, 0, 0]
label = 'Segmentation'
A = ArraySource(scalar_data=v)
A.origin = np.array(origion)
D = s.add_source(A) # add the data to the Mayavi engine
#Apply gaussain 3d filter to smooth visualization
# F=mlab.pipeline.user_defined(D, filter='ImageGaussianSmooth')
# F.filter.set_standard_deviation(0,0,0)
contour = Contour()
s.add_filter(contour)
# smooth = mlab.pipeline.user_defined(contour, filter='SmoothPolyDataFilter')
# smooth.filter.number_of_iterations = 1
# smooth.filter.relaxation_factor = 0
surface = Surface()
s.add_module(surface)
surface.module_manager.scalar_lut_manager.lut_mode = u'coolwarm'
surface.module_manager.scalar_lut_manager.reverse_lut = True
surface.actor.property.opacity = opacity
surface.actor.mapper.scalar_visibility = False
surface.actor.property.color = color #color
return surface
def plot_events_sliding(xs,
ys,
ts,
ps,
args,
dt=None,
sdt=None,
frames=None,
frame_ts=None,
padding=True):
skip = max(len(xs) // args.num_show, 1)
xs, ys, ts, ps = xs[::skip], ys[::skip], ts[::skip], ps[::skip]
t0 = ts[0]
sx, sy, st, sp = [], [], [], []
if padding:
for i in np.arange(ts[0] - dt, ts[0], sdt):
sx.append(0)
sy.append(0)
st.append(i)
sp.append(0)
print(len(sx))
print(st)
print(ts)
xs = np.concatenate((np.array(sx), xs))
ys = np.concatenate((np.array(sy), ys))
ts = np.concatenate((np.array(st), ts))
ps = np.concatenate((np.array(sp), ps))
print(ts)
ts += -st[0]
frame_ts += -st[0]
t0 += -st[0]
print(ts)
f = mlab.figure(bgcolor=(1, 1, 1), size=(1080, 720))
engine = mlab.get_engine()
scene = engine.scenes[0]
scene.scene.camera.position = [
373.1207907160101, 5353.96218497846, 7350.065665045519
]
scene.scene.camera.focal_point = [
228.0033999234376, 37.75424682790012, 3421.439332472788
]
scene.scene.camera.view_angle = 30.0
scene.scene.camera.view_up = [
0.9997493712140433, -0.02027499237784438, -0.009493125997461629
]
scene.scene.camera.clipping_range = [2400.251302762254, 11907.415293888362]
scene.scene.camera.compute_view_plane_normal()
print("ts from {} to {}, imgs from {} to {}".format(
ts[0], ts[-1], frame_ts[0], frame_ts[-1]))
frame_ts = np.array([t0] + list(frame_ts[0:-1]))
if dt is None:
dt = (ts[-1] - ts[0]) / 10
sdt = dt / 10
print("Using dt={}, sdt={}".format(dt, sdt))
if frames is not None:
sensor_size = frames[0].shape
else:
sensor_size = [max(ys), max(xs)]
if len(frame_ts.shape) == 2:
frame_ts = frame_ts[:, 1]
for i, t0 in enumerate(tqdm(np.arange(ts[0], ts[-1] - dt, sdt))):
te = t0 + dt
eidx0 = np.searchsorted(ts, t0)
eidx1 = np.searchsorted(ts, te)
fidx0 = np.searchsorted(frame_ts, t0)
fidx1 = np.searchsorted(frame_ts, te)
#print("{}:{} = {}".format(frame_ts[fidx0], ts[eidx0], fidx0))
wxs, wys, wts, wps = xs[eidx0:eidx1], ys[eidx0:eidx1], ts[
eidx0:eidx1], ps[eidx0:eidx1],
if fidx0 == fidx1:
wframes = []
wframe_ts = []
else:
wframes = frames[fidx0:fidx1]
wframe_ts = frame_ts[fidx0:fidx1]
save_path = os.path.join(args.output_path,
"frame_{:010d}.jpg".format(i))
plot_events(wxs,
wys,
wts,
wps,
save_path=save_path,
num_show=-1,
event_size=args.event_size,
imgs=wframes,
img_ts=wframe_ts,
show_events=not args.hide_events,
azim=args.azim,
elev=args.elev,
show_frames=not args.hide_frames,
crop=args.crop,
compress_front=args.compress_front,
invert=args.invert,
num_compress=args.num_compress,
show_plot=args.show_plot,
img_size=sensor_size,
show_axes=args.show_axes,
ts_scale=args.ts_scale)
if save_path is not None:
ensure_dir(save_path)
#mlab.savefig(save_path, figure=f, magnification=10)
#GUI().process_events()
#img = mlab.screenshot(figure=f, mode='rgba', antialiased=True)
#print(img.shape)
mlab.savefig(save_path, figure=f, magnification=8)
mlab.clf()
bc=[left_x, right_x, left_y],
record={
'strain': Vis3DTensorField(var='eps_ab'),
# 'damage': Vis3DStateField(var='omega_a'),
# 'kinematic hardening': Vis3DStateField(var='z_a')
})
s = m.sim
s.tloop.k_max = 1000
s.tline.step = 0.1
s.tloop.verbose = True
s.run()
print('area', thickness)
F_ti = s.hist.F_t
print('left')
print(np.sum(F_ti[-1, right_x.dofs]))
print('right')
print(np.sum(F_ti[-1, left_x.dofs]))
mlab.options.backend = 'envisage'
f_strain = mlab.figure()
scene = mlab.get_engine().scenes[-1]
scene.name = 'strain'
strain_viz = Viz3DTensorField(vis3d=s.hist['strain'])
strain_viz.setup()
decorate_figure(f_strain, strain_viz, 200, [70, 20, 0])
mlab.show()
def close():
"""Close the scene."""
f = mlab.gcf()
e = mlab.get_engine()
v = e.get_viewer(f)
v.close()
def generate_file_rst(fname, target_dir, src_dir, plot_gallery):
""" Generate the rst file for a given example.
"""
base_image_name = os.path.splitext(fname)[0]
image_fname = '%s_%%s.png' % base_image_name
this_template = rst_template
last_dir = os.path.split(src_dir)[-1]
# to avoid leading . in file names, and wrong names in links
if last_dir == '.' or last_dir == 'examples':
last_dir = ''
else:
last_dir += '_'
short_fname = last_dir + fname
src_file = os.path.join(src_dir, fname)
example_file = os.path.join(target_dir, fname)
shutil.copyfile(src_file, example_file)
# The following is a list containing all the figure names
figure_list = []
image_dir = os.path.join(target_dir, 'images')
thumb_dir = os.path.join(image_dir, 'thumb')
if not os.path.exists(image_dir):
os.makedirs(image_dir)
if not os.path.exists(thumb_dir):
os.makedirs(thumb_dir)
image_path = os.path.join(image_dir, image_fname)
stdout_path = os.path.join(image_dir,
'stdout_%s.txt' % base_image_name)
time_path = os.path.join(image_dir,
'time_%s.txt' % base_image_name)
thumb_file = os.path.join(thumb_dir, fname[:-3] + '.png')
time_elapsed = 0
if plot_gallery:
# generate the plot as png image if file name
# starts with plot and if it is more recent than an
# existing image.
first_image_file = image_path % 1
if os.path.exists(stdout_path):
stdout = open(stdout_path).read()
else:
stdout = ''
if os.path.exists(time_path):
time_elapsed = float(open(time_path).read())
if (not os.path.exists(first_image_file) or
os.stat(first_image_file).st_mtime
<= os.stat(src_file).st_mtime):
# We need to execute the code
print 'plotting %s' % fname
t0 = time()
import matplotlib.pyplot as plt
plt.close('all')
try:
from mayavi import mlab
except Exception, e:
from enthought.mayavi import mlab
mlab.close(all=True)
cwd = os.getcwd()
try:
# First CD in the original example dir, so that any file
# created by the example get created in this directory
orig_stdout = sys.stdout
os.chdir(os.path.dirname(src_file))
my_buffer = StringIO()
my_stdout = Tee(sys.stdout, my_buffer)
sys.stdout = my_stdout
my_globals = {'pl': plt}
execfile(os.path.basename(src_file), my_globals)
time_elapsed = time() - t0
sys.stdout = orig_stdout
my_stdout = my_buffer.getvalue()
# get variables so we can later add links to the documentation
example_code_obj = {}
for var_name, var in my_globals.iteritems():
if not hasattr(var, '__module__'):
continue
if not isinstance(var.__module__, basestring):
continue
if var.__module__.split('.')[0] not in DOCMODULES:
continue
# get the type as a string with other things stripped
tstr = str(type(var))
tstr = (tstr[tstr.find('\'')
+ 1:tstr.rfind('\'')].split('.')[-1])
# get shortened module name
module_short = get_short_module_name(var.__module__,
tstr)
cobj = {'name': tstr, 'module': var.__module__,
'module_short': module_short,
'obj_type': 'object'}
example_code_obj[var_name] = cobj
# find functions so we can later add links to the documentation
funregex = re.compile('[\w.]+\(')
with open(src_file, 'rt') as fid:
for line in fid.readlines():
if line.startswith('#'):
continue
for match in funregex.findall(line):
fun_name = match[:-1]
try:
exec('this_fun = %s' % fun_name, my_globals)
except Exception as err:
print 'extracting function failed'
print err
continue
this_fun = my_globals['this_fun']
if not callable(this_fun):
continue
if not hasattr(this_fun, '__module__'):
continue
if not isinstance(this_fun.__module__, basestring):
continue
if (this_fun.__module__.split('.')[0]
not in DOCMODULES):
continue
# get shortened module name
fun_name_short = fun_name.split('.')[-1]
module_short = get_short_module_name(
this_fun.__module__, fun_name_short)
cobj = {'name': fun_name_short,
'module': this_fun.__module__,
'module_short': module_short,
'obj_type': 'function'}
example_code_obj[fun_name] = cobj
fid.close()
if len(example_code_obj) > 0:
# save the dictionary, so we can later add hyperlinks
codeobj_fname = example_file[:-3] + '_codeobj.pickle'
with open(codeobj_fname, 'wb') as fid:
cPickle.dump(example_code_obj, fid,
cPickle.HIGHEST_PROTOCOL)
fid.close()
if '__doc__' in my_globals:
# The __doc__ is often printed in the example, we
# don't with to echo it
my_stdout = my_stdout.replace(my_globals['__doc__'],
'')
my_stdout = my_stdout.strip()
if my_stdout:
output_lines = my_stdout.split('\n')
if len(output_lines) > MAX_NB_LINES_STDOUT:
output_lines = output_lines[:MAX_NB_LINES_STDOUT]
output_lines.append('...')
stdout = ('**Script output**::\n\n %s\n\n'
% ('\n '.join(output_lines)))
open(stdout_path, 'w').write(stdout)
open(time_path, 'w').write('%f' % time_elapsed)
os.chdir(cwd)
# In order to save every figure we have two solutions :
# * iterate from 1 to infinity and call plt.fignum_exists(n)
# (this requires the figures to be numbered
# incrementally: 1, 2, 3 and not 1, 2, 5)
# * iterate over [fig_mngr.num for fig_mngr in
# matplotlib._pylab_helpers.Gcf.get_all_fig_managers()]
last_fig_num = 0
for fig_num in (fig_mngr.num for fig_mngr in
matplotlib._pylab_helpers.Gcf.get_all_fig_managers()):
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
plt.figure(fig_num)
# hack to keep black bg
facecolor = plt.gcf().get_facecolor()
if facecolor == (0.0, 0.0, 0.0, 1.0):
plt.savefig(image_path % fig_num, facecolor='black')
else:
plt.savefig(image_path % fig_num)
figure_list.append(image_fname % fig_num)
last_fig_num = fig_num
e = mlab.get_engine()
for scene in e.scenes:
last_fig_num += 1
mlab.savefig(image_path % last_fig_num)
figure_list.append(image_fname % last_fig_num)
mlab.close(scene)
except:
print 80 * '_'
print '%s is not compiling:' % fname
traceback.print_exc()
print 80 * '_'
finally:
os.chdir(cwd)
sys.stdout = orig_stdout
print " - time elapsed : %.2g sec" % time_elapsed
else:
figure_list = [f[len(image_dir):]
for f in glob.glob(image_path % '[1-9]')]
# generate thumb file
this_template = plot_rst_template
if os.path.exists(first_image_file):
make_thumbnail(first_image_file, thumb_file, 180, 120)
def ToyModel3d(sample):
"""
This script configure the 3D render motor (Mayavi) to show an interactive
reconstruction of the asphalt mixture sample
"""
src = mlab.pipeline.scalar_field(sample)
inverse_lut = False
colors = 5
iso = mlab.pipeline.iso_surface(src, contours=[1], opacity=0.4, colormap = 'blue-red')
iso.module_manager.scalar_lut_manager.reverse_lut = inverse_lut
iso.module_manager.scalar_lut_manager.number_of_colors = colors
ipw = mlab.pipeline.image_plane_widget(src, plane_orientation='y_axes', slice_index=10, colormap = 'blue-red')
ipw.module_manager.scalar_lut_manager.reverse_lut = inverse_lut
ipw.module_manager.scalar_lut_manager.number_of_colors = colors
scp = mlab.pipeline.scalar_cut_plane(src, colormap = 'blue-red')
scp.module_manager.scalar_lut_manager.reverse_lut = inverse_lut
scp.module_manager.scalar_lut_manager.number_of_colors = colors
#Set the Mayavi Colorbar Ranges
scp.module_manager.scalar_lut_manager.use_default_range = False
scp.module_manager.scalar_lut_manager.scalar_bar.position2 = array([ 0.1, 0.8])
scp.module_manager.scalar_lut_manager.scalar_bar.position = array([ 0.01, 0.15])
scp.module_manager.scalar_lut_manager.data_range = array([ 0., 2.])
scp.module_manager.scalar_lut_manager.scalar_bar.position2 = array([ 0.1, 0.8])
scp.module_manager.scalar_lut_manager.scalar_bar.position = array([ 0.01, 0.15])
scp.module_manager.scalar_lut_manager.data_range = array([ 0., 2.])
engine = mlab.get_engine()
textAggregate = Text()
textMastic = Text()
textVoids = Text()
engine.add_filter(textAggregate, scp.module_manager)
engine.add_filter(textMastic, ipw.module_manager)
engine.add_filter(textVoids, iso.module_manager)
textAggregate.text = 'Aggregate'
textMastic.text = 'Mastic'
textVoids.text = 'Air Voids'
textAggregate.actor.text_scale_mode = 'viewport'
textMastic.actor.text_scale_mode = 'viewport'
textVoids.actor.text_scale_mode = 'viewport'
textAggregate.actor.minimum_size = array([ 1, 10])
textMastic.actor.minimum_size = array([ 1, 10])
textVoids.actor.minimum_size = array([ 1, 10])
textAggregate.actor.position = array([ 0.115, 0.7 ])
textMastic.actor.position = array([ 0.115, 0.45])
textVoids.actor.position = array([ 0.115, 0.23])
mlab.orientation_axes()
mlab.title("Asphalt Mixture Reconstruction", size=0.25)
mlab.colorbar(title='Material', orientation='vertical', nb_labels=0, nb_colors=3)
mlab.show()
def setUp(self):
self._stop_unregister_all_engines()
mlab.options.backend = 'test'
e = mlab.get_engine()
self.e = e
def ctmr_gauss_plot(tri,
vert,
color=(0.8, 0.8, 0.8),
elecs=None,
weights=None,
opacity=1.0,
representation='surface',
line_width=1.0,
gsp=10,
cmap=mpl.cm.get_cmap('RdBu_r'),
show_colorbar=True,
new_fig=True,
vmin=None,
vmax=None,
ambient=0.4225,
specular=0.333,
specular_power=66,
diffuse=0.6995,
interpolation='phong'):
''' This function plots the 3D brain surface mesh
Parameters
----------
color : tuple
(n,n,n) tuple of floats between 0.0 and 1.0, background color of brain
elecs : array-like
[nchans x 3] matrix of electrode coordinate values in 3D
weights : array-like
[nchans x 1] - if [elecs] is also given, this will color the brain vertices
according to these weights
msize : float
size of the electrode. default = 2
opacity : float (0.0 - 1.0)
opacity of the brain surface (value from 0.0 - 1.0)
cmap : str or mpl.colors.LinearSegmentedColormap
colormap to use when plotting gaussian weights with [elecs]
and [weights]
representation : {'surface', 'wireframe'}
surface representation
line_width : float
width of lines for triangular mesh
gsp : float
gaussian smoothing parameter, larger makes electrode activity
more spread out across the surface if specified
Returns
-------
mesh : mayavi mesh (actor)
mlab : mayavi mlab scene
'''
# if color is another iterable, make it a tuple.
color = tuple(color)
brain_color = []
#c = np.zeros(vert.shape[0],)
if elecs is not None:
brain_color = np.zeros(vert.shape[0], )
for i in np.arange(elecs.shape[0]):
b_z = np.abs(vert[:, 2] - elecs[i, 2])
b_y = np.abs(vert[:, 1] - elecs[i, 1])
b_x = np.abs(vert[:, 0] - elecs[i, 0])
gauss_wt = np.nan_to_num(weights[i] * np.exp(
(-(b_x**2 + b_z**2 + b_y**2)) / gsp)) #gaussian
brain_color = brain_color + gauss_wt
#scale the colors so that it matches the weights that were passed in
brain_color = brain_color * (np.abs(weights).max() /
np.abs(brain_color).max())
if vmin == None and vmax == None:
vmin, vmax = -np.abs(brain_color).max(), np.abs(brain_color).max()
# plot cortex and begin display
if new_fig:
mlab.figure(fgcolor=(0, 0, 0), bgcolor=(1, 1, 1), size=(1200, 900))
if elecs is not None:
kwargs = {}
if type(cmap) == str:
kwargs.update(colormap=cmap)
mesh = mlab.triangular_mesh(vert[:, 0],
vert[:, 1],
vert[:, 2],
tri,
representation=representation,
opacity=opacity,
line_width=line_width,
scalars=brain_color,
vmin=vmin,
vmax=vmax,
**kwargs)
if type(cmap) == mpl.colors.LinearSegmentedColormap:
mesh.module_manager.scalar_lut_manager.lut.table = (
cmap(np.linspace(0, 1, 255)) * 255).astype('int')
else:
mesh = mlab.triangular_mesh(vert[:, 0],
vert[:, 1],
vert[:, 2],
tri,
color=color,
representation=representation,
opacity=opacity,
line_width=line_width)
# cell_data = mesh.mlab_source.dataset.cell_data
# cell_data.scalars = brain_color
# cell_data.scalars.name = 'Cell data'
# cell_data.update()
#mesh2 = mlab.pipeline.set_active_attribute(mesh, cell_scalars = 'Cell data')
#mlab.pipeline.surface(mesh)
if weights is not None and show_colorbar:
mlab.colorbar()
# change OpenGL mesh properties for phong point light shading
mesh.actor.property.ambient = ambient
mesh.actor.property.specular = specular
mesh.actor.property.specular_power = specular_power
mesh.actor.property.diffuse = diffuse
mesh.actor.property.interpolation = interpolation
mesh.scene.light_manager.light_mode = 'vtk'
if opacity < 1.0:
mesh.scene.renderer.set(
use_depth_peeling=True
) #, maximum_number_of_peels=100, occlusion_ratio=0.0
# Make the mesh look smoother
for child in mlab.get_engine().scenes[0].children:
poly_data_normals = child.children[0]
try:
poly_data_normals.filter.feature_angle = 80.0 # Feature angle says which angles are considered hard corners
except:
pass
return mesh, mlab
def __init__(self,
hex_list,
np_file='/tmp/magnetic_ground_truth.np',
robot_height=40,
width=800,
height=600,
start_point=(0, 0, 0),
message='experiment default message...',
battery=99999):
self.debug = False
self.animator = None
self.movement_mode = 0
self.start_point = start_point
self.robot_height = robot_height
self.message = message
self.start_time = int(time.time() * 1000)
self.timestep = 0
self.width = width
self.height = height
self.f = mlab.figure(size=(self.width, self.height))
visual.set_viewer(self.f)
v = mlab.view(270, 180)
#print v
engine = mlab.get_engine()
self.s = engine.current_scene
self.s.scene.interactor.add_observer('KeyPressEvent',
self.keypress_callback)
self.robots = []
colors = list(PathBatterySimulator.color_codes)
for key, local_hex_list in sorted(hex_list['internal_routes'].items()):
color = colors.pop(0)
ball = visual.sphere(color=color,
radius=PathBatterySimulator.ball_radius)
ball.x = self.start_point[0]
ball.y = self.start_point[1]
ball.z = self.start_point[2]
r, g, b = color
rt = r + (0.25 * (1 - r))
gt = g + (0.25 * (1 - g))
bt = b + (0.25 * (1 - b))
curve_color = (rt, gt, bt)
curve = visual.curve(color=curve_color,
radius=PathBatterySimulator.curve_radius)
r_ball = RobotBall(key,
local_hex_list,
hex_list['external_routes'][key],
ball,
curve,
battery=battery)
self.robots.append(r_ball)
x = np.linspace(0, self.width, 1)
y = np.linspace(0, self.height, 1)
z = np.loadtxt(np_file)
z *= 255.0 / z.max()
# HARDCODED
# Todo: REMOVE THIS CODE ON THE FINAL RELEASE
for xx in xrange(0, 200):
for yy in xrange(400, 600):
z[yy][xx] = 0
mlab.surf(x, y, z)
self.master_cmd = MasterCommand(self.robots)
self.robot_pos = open('/tmp/robot_log.txt', 'a')
def close():
"""Close the scene."""
f = mlab.gcf()
e = mlab.get_engine()
v = e.get_viewer(f)
v.close()
def save_figures(image_path, fig_count, gallery_conf):
"""Save all open matplotlib figures of the example code-block
Parameters
----------
image_path : str
Path where plots are saved (format string which accepts figure number)
fig_count : int
Previous figure number count. Figure number add from this number
gallery_conf : dict
Contains the configuration of Sphinx-Gallery
Returns
-------
figure_list : list of str
strings containing the full path to each figure
images_rst : str
rst code to embed the images in the document
"""
figure_list = []
fig_numbers = plt.get_fignums()
for fig_num in fig_numbers:
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
fig = plt.figure(fig_num)
kwargs = {}
to_rgba = matplotlib.colors.colorConverter.to_rgba
for attr in ['facecolor', 'edgecolor']:
fig_attr = getattr(fig, 'get_' + attr)()
default_attr = matplotlib.rcParams['figure.' + attr]
if to_rgba(fig_attr) != to_rgba(default_attr):
kwargs[attr] = fig_attr
current_fig = image_path.format(fig_count + fig_num)
fig.savefig(current_fig, **kwargs)
figure_list.append(current_fig)
if gallery_conf.get('find_mayavi_figures', False):
from mayavi import mlab
e = mlab.get_engine()
last_matplotlib_fig_num = fig_count + len(figure_list)
total_fig_num = last_matplotlib_fig_num + len(e.scenes)
mayavi_fig_nums = range(last_matplotlib_fig_num + 1, total_fig_num + 1)
for scene, mayavi_fig_num in zip(e.scenes, mayavi_fig_nums):
current_fig = image_path.format(mayavi_fig_num)
mlab.savefig(current_fig, figure=scene)
# make sure the image is not too large
scale_image(current_fig, current_fig, 850, 999)
figure_list.append(current_fig)
mlab.close(all=True)
# Depending on whether we have one or more figures, we're using a
# horizontal list or a single rst call to 'image'.
images_rst = ""
if len(figure_list) == 1:
figure_name = os.path.relpath(figure_list[0], gallery_conf['src_dir'])
images_rst = SINGLE_IMAGE % figure_name.lstrip('/')
elif len(figure_list) > 1:
images_rst = HLIST_HEADER
for figure_name in figure_list:
figure_name = os.path.relpath(figure_name, gallery_conf['src_dir'])
images_rst += HLIST_IMAGE_TEMPLATE % figure_name.lstrip('/')
return figure_list, images_rst
def _set_keyframe(self):
t = self.current_frame / float(self._last_frame)
self._camera_interpolator.add_camera(t, mlab.get_engine().current_scene.scene.camera)
self._t_keyframes[self.current_frame] = t
def setUp(self):
self._stop_unregister_all_engines()
mlab.options.backend = 'test'
e = mlab.get_engine()
self.e = e
def plot_map(map,
affine,
cut_coords=None,
anat=None,
anat_affine=None,
slicer='ortho',
figure=None,
axes=None,
title=None,
threshold=None,
annotate=True,
draw_cross=True,
do3d=False,
threshold_3d=None,
view_3d=(38.5, 70.5, 300, (-2.7, -12, 9.1)),
black_bg=False,
**kwargs):
""" Plot three cuts of a given activation map (Frontal, Axial, and Lateral)
Parameters
----------
map : 3D ndarray
The activation map, as a 3D image.
affine : 4x4 ndarray
The affine matrix going from image voxel space to MNI space.
cut_coords: None, or a tuple of floats
The MNI coordinates of the point where the cut is performed, in
MNI coordinates and order.
If slicer is 'ortho', this should be a 3-tuple: (x, y, z)
For slicer == 'x', 'y', or 'z', then these are the
coordinates of each cut in the corresponding direction.
If None is given, the cuts is calculated automaticaly.
anat : 3D ndarray or False, optional
The anatomical image to be used as a background. If None, the
MNI152 T1 1mm template is used. If False, no anat is displayed.
anat_affine : 4x4 ndarray, optional
The affine matrix going from the anatomical image voxel space to
MNI space. This parameter is not used when the default
anatomical is used, but it is compulsory when using an
explicite anatomical image.
slicer: {'ortho', 'x', 'y', 'z'}
Choose the direction of the cuts. With 'ortho' three cuts are
performed in orthogonal directions
figure : integer or matplotlib figure, optional
Matplotlib figure used or its number. If None is given, a
new figure is created.
axes : matplotlib axes or 4 tuple of float: (xmin, xmax, ymin, ymin), optional
The axes, or the coordinates, in matplotlib figure space,
of the axes used to display the plot. If None, the complete
figure is used.
title : string, optional
The title dispayed on the figure.
threshold : a number, None, or 'auto'
If None is given, the maps are not thresholded.
If a number is given, it is used to threshold the maps:
values below the threshold are plotted as transparent. If
auto is given, the threshold is determined magically by
analysis of the map.
annotate: boolean, optional
If annotate is True, positions and left/right annotation
are added to the plot.
draw_cross: boolean, optional
If draw_cross is True, a cross is drawn on the plot to
indicate the cut plosition.
do3d: {True, False or 'interactive'}, optional
If True, Mayavi is used to plot a 3D view of the
map in addition to the slicing. If 'interactive', the
3D visualization is displayed in an additional interactive
window.
threshold_3d:
The threshold to use for the 3D view (if any). Defaults to
the same threshold as that used for the 2D view.
view_3d: tuple,
The view used to take the screenshot: azimuth, elevation,
distance and focalpoint, see the docstring of mlab.view.
black_bg: boolean, optional
If True, the background of the image is set to be black. If
you whish to save figures with a black background, you
will need to pass "facecolor='k', edgecolor='k'" to pylab's
savefig.
kwargs: extra keyword arguments, optional
Extra keyword arguments passed to pylab.imshow
Notes
-----
Arrays should be passed in numpy convention: (x, y, z)
ordered.
Use masked arrays to create transparency:
import numpy as np
map = np.ma.masked_less(map, 0.5)
plot_map(map, affine)
"""
map, affine = _xyz_order(map, affine)
nan_mask = np.isnan(np.asarray(map))
if np.any(nan_mask):
map = map.copy()
map[nan_mask] = 0
del nan_mask
# Deal with automatic settings of plot parameters
if threshold == 'auto':
# Threshold epsilon above a percentile value, to be sure that some
# voxels are indeed threshold
threshold = _fast_abs_percentile(map) + 1e-5
if do3d:
try:
try:
from mayavi import version
except ImportError:
from enthought.mayavi import version
if not int(version.version[0]) > 2:
raise ImportError
except ImportError:
warnings.warn('Mayavi > 3.x not installed, plotting only 2D')
do3d = False
slicer = SLICERS[slicer].init_with_figure(data=map,
affine=affine,
threshold=threshold,
cut_coords=cut_coords,
figure=figure,
axes=axes,
black_bg=black_bg,
leave_space=do3d)
# Use Mayavi for the 3D plotting
if do3d:
from .maps_3d import plot_map_3d, m2screenshot
try:
from tvtk.api import tvtk
except ImportError:
from enthought.tvtk.api import tvtk
version = tvtk.Version()
offscreen = True
if (version.vtk_major_version, version.vtk_minor_version) < (5, 2):
offscreen = False
if do3d == 'interactive':
offscreen = False
cmap = kwargs.get('cmap', pl.cm.cmap_d[pl.rcParams['image.cmap']])
# Computing vmin and vmax is costly in time, and is needed
# later, so we compute them now, and store them for future
# use
vmin = kwargs.get('vmin', map.min())
kwargs['vmin'] = vmin
vmax = kwargs.get('vmax', map.max())
kwargs['vmax'] = vmax
try:
from mayavi import mlab
except ImportError:
from enthought.mayavi import mlab
if threshold_3d is None:
threshold_3d = threshold
plot_map_3d(np.asarray(map),
affine,
cut_coords=cut_coords,
anat=anat,
anat_affine=anat_affine,
offscreen=offscreen,
cmap=cmap,
threshold=threshold_3d,
view=view_3d,
vmin=vmin,
vmax=vmax)
ax = slicer.axes.values()[0].ax.figure.add_axes((0.001, 0, 0.29, 1))
ax.axis('off')
m2screenshot(mpl_axes=ax)
if offscreen:
# Clean up, so that the offscreen engine doesn't become the
# default
mlab.clf()
engine = mlab.get_engine()
try:
from mayavi.core.registry import registry
except:
from enthought.mayavi.core.registry import registry
for key, value in registry.engines.iteritems():
if value is engine:
registry.engines.pop(key)
break
if threshold:
map = np.ma.masked_inside(map, -threshold, threshold, copy=False)
_plot_anat(slicer,
anat,
anat_affine,
title=title,
annotate=annotate,
draw_cross=draw_cross)
slicer.plot_map(map, affine, **kwargs)
return slicer
def grid(self, **kwargs):
x = kwargs.get('x', self._obj.SCHISM_hgrid_node_x[:].values)
y = kwargs.get('y', self._obj.SCHISM_hgrid_node_y[:].values)
try:
t = kwargs.get('t', self._obj.time.values)
except:
pass
tri3 = kwargs.get(
'tri3',
self._obj.SCHISM_hgrid_face_nodes.values[:, :3].astype(int))
dim = kwargs.get('dim', '2D')
R = kwargs.get('R', 1.)
if dim == '3D':
px = np.cos(y / 180 * np.pi) * np.cos(x / 180 * np.pi) * R
py = np.cos(y / 180 * np.pi) * np.sin(x / 180 * np.pi) * R
pz = np.sin(y / 180 * np.pi) * R
else:
px = x
py = y
pz = np.zeros(x.shape[0])
mlab.figure(1,
size=(3840 / 2, 2160 / 2),
bgcolor=(0, 0, 0),
fgcolor=(1., 1., 1.))
mlab.clf()
bcolor = kwargs.get('bcolor', (0., 0., 0.))
if dim == '3D': self.globe(R - .002, bcolor=bcolor)
# 3D triangular mesh surface (like trisurf)
grd = mlab.triangular_mesh(px,
py,
pz,
tri3,
representation='wireframe',
opacity=1.0)
coast = kwargs.get('coastlines', None)
if coast is not None:
try:
del kwargs['coastlines', 'R']
except:
pass
src, lines = self.c3d(coast, R=R, **kwargs)
mlab.pipeline.surface(src,
color=(1, 0, 0),
line_width=10,
opacity=0.8)
engine = mlab.get_engine()
scene = engine.scenes[0]
scene.scene.z_plus_view()
mlab.show()
return
def ToyModel3d(sample):
"""
This script configure the 3D render motor (Mayavi) to show an interactive
reconstruction of the asphalt mixture sample
"""
src = mlab.pipeline.scalar_field(sample)
inverse_lut = False
colors = 5
iso = mlab.pipeline.iso_surface(src,
contours=[1],
opacity=0.4,
colormap='blue-red')
iso.module_manager.scalar_lut_manager.reverse_lut = inverse_lut
iso.module_manager.scalar_lut_manager.number_of_colors = colors
ipw = mlab.pipeline.image_plane_widget(src,
plane_orientation='y_axes',
slice_index=10,
colormap='blue-red')
ipw.module_manager.scalar_lut_manager.reverse_lut = inverse_lut
ipw.module_manager.scalar_lut_manager.number_of_colors = colors
scp = mlab.pipeline.scalar_cut_plane(src, colormap='blue-red')
scp.module_manager.scalar_lut_manager.reverse_lut = inverse_lut
scp.module_manager.scalar_lut_manager.number_of_colors = colors
#Set the Mayavi Colorbar Ranges
scp.module_manager.scalar_lut_manager.use_default_range = False
scp.module_manager.scalar_lut_manager.scalar_bar.position2 = array(
[0.1, 0.8])
scp.module_manager.scalar_lut_manager.scalar_bar.position = array(
[0.01, 0.15])
scp.module_manager.scalar_lut_manager.data_range = array([0., 2.])
scp.module_manager.scalar_lut_manager.scalar_bar.position2 = array(
[0.1, 0.8])
scp.module_manager.scalar_lut_manager.scalar_bar.position = array(
[0.01, 0.15])
scp.module_manager.scalar_lut_manager.data_range = array([0., 2.])
engine = mlab.get_engine()
textAggregate = Text()
textMastic = Text()
textVoids = Text()
engine.add_filter(textAggregate, scp.module_manager)
engine.add_filter(textMastic, ipw.module_manager)
engine.add_filter(textVoids, iso.module_manager)
textAggregate.text = 'Aggregate'
textMastic.text = 'Mastic'
textVoids.text = 'Air Voids'
textAggregate.actor.text_scale_mode = 'viewport'
textMastic.actor.text_scale_mode = 'viewport'
textVoids.actor.text_scale_mode = 'viewport'
textAggregate.actor.minimum_size = array([1, 10])
textMastic.actor.minimum_size = array([1, 10])
textVoids.actor.minimum_size = array([1, 10])
textAggregate.actor.position = array([0.115, 0.7])
textMastic.actor.position = array([0.115, 0.45])
textVoids.actor.position = array([0.115, 0.23])
mlab.orientation_axes()
mlab.title("Asphalt Mixture Reconstruction", size=0.25)
mlab.colorbar(title='Material',
orientation='vertical',
nb_labels=0,
nb_colors=3)
mlab.show()
def generate_file_rst(fname, target_dir, src_dir, plot_gallery):
""" Generate the rst file for a given example.
"""
base_image_name = os.path.splitext(fname)[0]
image_fname = '%s_%%s.png' % base_image_name
this_template = rst_template
last_dir = os.path.split(src_dir)[-1]
# to avoid leading . in file names, and wrong names in links
if last_dir == '.' or last_dir == 'examples':
last_dir = ''
else:
last_dir += '_'
short_fname = last_dir + fname
src_file = os.path.join(src_dir, fname)
example_file = os.path.join(target_dir, fname)
shutil.copyfile(src_file, example_file)
# The following is a list containing all the figure names
figure_list = []
image_dir = os.path.join(target_dir, 'images')
thumb_dir = os.path.join(image_dir, 'thumb')
if not os.path.exists(image_dir):
os.makedirs(image_dir)
if not os.path.exists(thumb_dir):
os.makedirs(thumb_dir)
image_path = os.path.join(image_dir, image_fname)
stdout_path = os.path.join(image_dir, 'stdout_%s.txt' % base_image_name)
time_path = os.path.join(image_dir, 'time_%s.txt' % base_image_name)
thumb_file = os.path.join(thumb_dir, fname[:-3] + '.png')
time_elapsed = 0
if plot_gallery:
# generate the plot as png image if file name
# starts with plot and if it is more recent than an
# existing image.
first_image_file = image_path % 1
if os.path.exists(stdout_path):
stdout = open(stdout_path).read()
else:
stdout = ''
if os.path.exists(time_path):
time_elapsed = float(open(time_path).read())
if (not os.path.exists(first_image_file)
or os.stat(first_image_file).st_mtime <=
os.stat(src_file).st_mtime):
# We need to execute the code
print 'plotting %s' % fname
t0 = time()
import matplotlib.pyplot as plt
plt.close('all')
try:
from mayavi import mlab
except Exception, e:
from enthought.mayavi import mlab
mlab.close(all=True)
cwd = os.getcwd()
try:
# First CD in the original example dir, so that any file
# created by the example get created in this directory
orig_stdout = sys.stdout
os.chdir(os.path.dirname(src_file))
my_buffer = StringIO()
my_stdout = Tee(sys.stdout, my_buffer)
sys.stdout = my_stdout
my_globals = {'pl': plt}
execfile(os.path.basename(src_file), my_globals)
time_elapsed = time() - t0
sys.stdout = orig_stdout
my_stdout = my_buffer.getvalue()
# get variables so we can later add links to the documentation
example_code_obj = {}
for var_name, var in my_globals.iteritems():
if not hasattr(var, '__module__'):
continue
if not isinstance(var.__module__, basestring):
continue
if var.__module__.split('.')[0] not in DOCMODULES:
continue
# get the type as a string with other things stripped
tstr = str(type(var))
tstr = (tstr[tstr.find('\'') +
1:tstr.rfind('\'')].split('.')[-1])
# get shortened module name
module_short = get_short_module_name(var.__module__, tstr)
cobj = {
'name': tstr,
'module': var.__module__,
'module_short': module_short,
'obj_type': 'object'
}
example_code_obj[var_name] = cobj
# find functions so we can later add links to the documentation
funregex = re.compile('[\w.]+\(')
with open(src_file, 'rt') as fid:
for line in fid.readlines():
if line.startswith('#'):
continue
for match in funregex.findall(line):
fun_name = match[:-1]
try:
exec('this_fun = %s' % fun_name, my_globals)
except Exception as err:
print 'extracting function failed'
print err
continue
this_fun = my_globals['this_fun']
if not callable(this_fun):
continue
if not hasattr(this_fun, '__module__'):
continue
if not isinstance(this_fun.__module__, basestring):
continue
if (this_fun.__module__.split('.')[0]
not in DOCMODULES):
continue
# get shortened module name
fun_name_short = fun_name.split('.')[-1]
module_short = get_short_module_name(
this_fun.__module__, fun_name_short)
cobj = {
'name': fun_name_short,
'module': this_fun.__module__,
'module_short': module_short,
'obj_type': 'function'
}
example_code_obj[fun_name] = cobj
fid.close()
if len(example_code_obj) > 0:
# save the dictionary, so we can later add hyperlinks
codeobj_fname = example_file[:-3] + '_codeobj.pickle'
with open(codeobj_fname, 'wb') as fid:
cPickle.dump(example_code_obj, fid,
cPickle.HIGHEST_PROTOCOL)
fid.close()
if '__doc__' in my_globals:
# The __doc__ is often printed in the example, we
# don't with to echo it
my_stdout = my_stdout.replace(my_globals['__doc__'], '')
my_stdout = my_stdout.strip()
if my_stdout:
output_lines = my_stdout.split('\n')
if len(output_lines) > MAX_NB_LINES_STDOUT:
output_lines = output_lines[:MAX_NB_LINES_STDOUT]
output_lines.append('...')
stdout = ('**Script output**::\n\n %s\n\n' %
('\n '.join(output_lines)))
open(stdout_path, 'w').write(stdout)
open(time_path, 'w').write('%f' % time_elapsed)
os.chdir(cwd)
# In order to save every figure we have two solutions :
# * iterate from 1 to infinity and call plt.fignum_exists(n)
# (this requires the figures to be numbered
# incrementally: 1, 2, 3 and not 1, 2, 5)
# * iterate over [fig_mngr.num for fig_mngr in
# matplotlib._pylab_helpers.Gcf.get_all_fig_managers()]
last_fig_num = 0
for fig_num in (fig_mngr.num for fig_mngr in matplotlib.
_pylab_helpers.Gcf.get_all_fig_managers()):
# Set the fig_num figure as the current figure as we can't
# save a figure that's not the current figure.
plt.figure(fig_num)
# hack to keep black bg
facecolor = plt.gcf().get_facecolor()
if facecolor == (0.0, 0.0, 0.0, 1.0):
plt.savefig(image_path % fig_num, facecolor='black')
else:
plt.savefig(image_path % fig_num)
figure_list.append(image_fname % fig_num)
last_fig_num = fig_num
e = mlab.get_engine()
for scene in e.scenes:
last_fig_num += 1
mlab.savefig(image_path % last_fig_num)
figure_list.append(image_fname % last_fig_num)
mlab.close(scene)
except:
print 80 * '_'
print '%s is not compiling:' % fname
traceback.print_exc()
print 80 * '_'
finally:
os.chdir(cwd)
sys.stdout = orig_stdout
print " - time elapsed : %.2g sec" % time_elapsed
else:
figure_list = [
f[len(image_dir):] for f in glob.glob(image_path % '[1-9]')
]
# generate thumb file
this_template = plot_rst_template
if os.path.exists(first_image_file):
make_thumbnail(first_image_file, thumb_file, 180, 120)
colormap='RdBu',
transparent=False,
opacity=1)
# Upper layers
for i in range(0, max_slice):
mlab.pipeline.image_plane_widget(scalar,
plane_orientation='z_axes',
slice_index=50,
colormap='RdBu',
transparent=True,
opacity=0.5)
engine = mlab.get_engine()
streamlines = engine.scenes[0].children[1].children[0].children[0].children[0]
streamlines.seed.widget = streamlines.seed.widget_list[2]
streamlines.seed.widget.enabled = False
streamlines.seed.widget.interactor = None
# streamlines.seed.widget = <tvtk.tvtk_classes.plane_widget.PlaneWidget object at 0x136851d70>
streamlines.seed.widget.enabled = False
streamlines.seed.widget.point1 = array([64.5, 96.25, 32.75])
streamlines.seed.widget.origin = array([32.75, 32.75, 64.5])
streamlines.seed.widget.center = array([64.5, 64.5, 64.5])
streamlines.seed.widget.normal = array([0., 0., 1.])
streamlines.seed.widget.point2 = array([32.75, 96.25, 64.5])
def draw(surf_vertices,
surf_triangles,
surf_scalars,
surf_cmap,
surf_opacity,
node_coords=None,
node_sizes=None,
node_colors=None,
conn_list=None,
conn_pct=None,
conn_cmap=None):
"""
Draws a brain graph superimposed on the brain surface
Parameters
----------
surf_vertices: numpy.ndarray, shape:(n_vert, 3)
Three-dimensional coordinates of each vertex of the surface
surf_triangles: numpy.ndarray, shape:(n_tria, 3)
Sets of vertex indices forming a triangle face of the surface
surf_scalars: numpy.ndarray, shape:(n_vert,)
Scalars representing magnitude or intensity at each vertex
surf_cmap: str
Colormap used to map surf_scalars to color
surf_opacity: float, 0.0<=x<=1.0
Opacity of the colormap
node_coords: numpy.ndarray, shape:(n_node, 3)
XYZ MNI coordinates of the network nodes
Perhaps centroid location of the ROIs
node_sizes: numpy.ndarray, shape:(n_node,)
Radii of spheres representing nodes
node_colors: numpy.ndarray, shape:(n_node, 4)
RGBA values to fill spheres representing nodes
conn_list: numpy.ndarray, shape:(n_node*(n_node-1) / 2,)
List of connections between nodes in the network
Assumes list represents the upper triangle of the adjacency matrix
conn_pct: tuple, shape:(2,)
Percentile range of connections to plot
conn_cmap: str
Colormap used to map connection strengths to color
"""
### Setup the engine and scene
my_engine = mlab.get_engine()
fig = mlab.figure(size=(1000, 1000),
bgcolor=(1.0, 1.0, 1.0),
engine=my_engine)
if node_coords is not None:
### Plot the nodes on the brain
n_node = node_coords.shape[0]
for n_i in xrange(n_node):
node_source = mlab.pipeline.scalar_scatter(node_coords[n_i, 0],
node_coords[n_i, 1],
node_coords[n_i, 2],
figure=fig)
node_surface = mlab.pipeline.glyph(node_source,
scale_mode='none',
scale_factor=node_sizes[n_i],
mode='sphere',
color=tuple(
node_colors[n_i][:3]),
opacity=node_colors[n_i][3],
figure=fig)
if conn_list is not None:
### Plot the connections on the brain
n_conn = conn_list.shape[0]
# Generate connection vectors
e_start = np.zeros((n_conn, 3))
e_vec = np.zeros((n_conn, 3))
triu_ix, triu_iy = np.triu_indices(n_node, k=1)
for ii, (ix, iy) in enumerate(zip(triu_ix, triu_iy)):
e_start[ii, :] = node_coords[ix, :]
e_vec[ii, :] = (node_coords[iy, :] - node_coords[ix, :])
# Threshold the connections
lo_thr = np.percentile(conn_list, conn_pct[0])
hi_thr = np.percentile(conn_list, conn_pct[1])
conn_thr_idx = np.flatnonzero((conn_list >= lo_thr)
& (conn_list <= hi_thr))
# Import vectors into pipeline
conn_source = mlab.pipeline.vector_scatter(e_start[:, 0],
e_start[:, 1],
e_start[:, 2],
e_vec[:, 0],
e_vec[:, 1],
e_vec[:, 2],
figure=fig)
conn_source.mlab_source.dataset.point_data.scalars = conn_list
#conn_thresh = mlab.pipeline.threshold(conn_source, low=lo_thr, up=hi_thr, figure=fig)
# Change connection attributes
conn_surface = mlab.pipeline.vectors(conn_source,
colormap=conn_cmap,
scale_factor=1,
line_width=4,
transparent=True,
scale_mode='vector',
figure=fig)
conn_surface.glyph.glyph.clamping = False
#conn_surface.actor.property.opacity = 0.1
conn_surface.module_manager.vector_lut_manager.reverse_lut = False
conn_surface.glyph.glyph_source.glyph_source = (\
conn_surface.glyph.glyph_source.glyph_dict['glyph_source2d'])
conn_surface.glyph.glyph_source.glyph_source.glyph_type = 'dash'
### Plot the colored brain regions
surf_source = mlab.pipeline.triangular_mesh_source(surf_vertices[:, 0],
surf_vertices[:, 1],
surf_vertices[:, 2],
surf_triangles,
scalars=surf_scalars,
opacity=1.0,
figure=fig)
surf_norms = mlab.pipeline.poly_data_normals(surf_source, figure=fig)
surf_norms.filter.splitting = True
surf_surf = mlab.pipeline.surface(surf_norms, figure=fig)
surf_surf.parent.scalar_lut_manager.set(lut_mode=surf_cmap,
data_range=[0, 1],
use_default_range=False)
lut = surf_surf.module_manager.scalar_lut_manager.lut.table.to_array()
lut[:, 3] = surf_opacity
surf_surf.module_manager.scalar_lut_manager.lut.table = lut
surf_surf.actor.property.backface_culling = True
if conn_list is not None:
return my_engine, conn_source
else:
return my_engine
def serve_tcp(engine=None, port=8007, logto=sys.stdout, max_connect=1):
"""Serve the `M2TCP` protocol using the given `engine` on the
specified `port` logging messages to given `logto` which is a
file-like object. This function will block till the service is
closed. There is no need to call `mlab.show()` after or before
this. The Mayavi UI will be fully responsive.
**Parameters**
:engine: Mayavi engine to use. If this is `None`,
`mlab.get_engine()` is used to find an appropriate engine.
:port: int: port to serve on.
:logto: file: File like object to log messages to. If this is
`None` it disables logging.
:max_connect: int: Maximum number of simultaneous connections to
support.
**Examples**
Here is a very simple example::
from mayavi import mlab
from mayavi.tools import server
mlab.test_plot3d()
server.serve_tcp()
The TCP server will listen on port 8007 by default in the above.
Any data sent to the server is simply exec'd, meaning you can do
pretty much anything you want. The `engine`, `scene`, `camera` and
`mlab` are all available and can be used. For example after running
the above you can do this::
$ telnet localhost 8007
Trying 127.0.0.1...
Connected to localhost.
Escape character is '^]'.
scene.camera.azimuth(45)
mlab.clf()
mlab.test_contour3d()
scene.camera.zoom(1.5)
**Warning**
Data sent is exec'd so this is a security hole.
"""
from mayavi import mlab
e = engine or mlab.get_engine()
# Setup the factory with the right attributes.
factory = Factory()
factory.protocol = M2TCP
factory.maxConnect = max_connect
factory.numConnect = 0
factory.engine = e
factory.scene = e.current_scene.scene
factory.mlab = mlab
if logto is not None:
log.startLogging(logto)
log.msg('Serving Mayavi2 TCP server on port', port)
log.msg('Using Engine', e)
# Register the running wxApp.
reactor.registerWxApp(wx.GetApp())
# Listen on port 9007 using above protocol.
reactor.listenTCP(port, factory)
# Run the server + app. This will block.
reactor.run()
def _engine_default(self):
return m2_mlab.get_engine()
