enthought / mayavi

3D visualization of scientific data in Python
http://docs.enthought.com/mayavi/mayavi/
Other
1.3k stars 284 forks source link

How to use mayavi to show the three-dimensional (x,y,z) cords? #1140

Open TianSong1991 opened 2 years ago

TianSong1991 commented 2 years ago

Hello, I want to use mayavi to show the three-dimensional (x,y,z) cords? But I don't find the api of the mayavi have the function of showing three-dimensional (x,y,z) cords. I want to know whether the mayavi has the function of showing three-dimensional (x,y,z) cords?Thanks for everyone answering the question.

drocheam commented 2 years ago

I asked myself the same, but did not find an inbuilt solution for this. But what you can do, is add a text overlay to your scene that gets updated with a position everytime you press a specified mouse button in the scene. This is done by registering an on_mouse_pick() event.

As an example I modified the Lorenz Ui Example from https://docs.enthought.com/mayavi/mayavi/auto/example_lorenz_ui.html#example-lorenz-ui

I added a onSpacePick() function in lines 134-136 that displays the position in the text overlay. The registration of this function is done in scene_active() in lines 67-71 that gets called when the scene is activated. Note that I also needed to change the trait in line 78 for this to work.

Here's the full code:

# Author: Prabhu Ramachandran <prabhu@aero.iitb.ac.in>
# Copyright (c) 2008-2020, Enthought, Inc.
# License: BSD Style.

import numpy as np
import scipy

from traits.api import HasTraits, Range, Instance, \
        on_trait_change, Array, Tuple, Str
from traitsui.api import View, Item, HSplit, Group

from mayavi import mlab
from mayavi.core.ui.api import MayaviScene, MlabSceneModel, \
    SceneEditor

################################################################################
# `Lorenz` class.
################################################################################
class Lorenz(HasTraits):

    # The parameters for the Lorenz system, defaults to the standard ones.
    s = Range(0.0, 20.0, 10.0, desc='the parameter s', enter_set=True,
              auto_set=False)
    r = Range(0.0, 50.0, 28.0, desc='the parameter r', enter_set=True,
              auto_set=False)
    b = Range(0.0, 10.0, 8./3., desc='the parameter b', enter_set=True,
              auto_set=False)

    # These expressions are evaluated to compute the right hand sides of
    # the ODE.  Defaults to the Lorenz system.
    u = Str('s*(y-x)', desc='the x component of the velocity',
            auto_set=False, enter_set=True)
    v = Str('r*x - y - x*z', desc='the y component of the velocity',
            auto_set=False, enter_set=True)
    w = Str('x*y - b*z', desc='the z component of the velocity',
            auto_set=False, enter_set=True)

    # Tuple of x, y, z arrays where the field is sampled.
    points = Tuple(Array, Array, Array)

    # The mayavi(mlab) scene.
    scene = Instance(MlabSceneModel, args=())

    # The "flow" which is a Mayavi streamline module.
    flow = Instance(HasTraits)

    ########################################
    # The UI view to show the user.
    view = View(HSplit(
                    Group(
                        Item('scene', editor=SceneEditor(scene_class=MayaviScene),
                             height=500, width=500, show_label=False)),
                    Group(
                        Item('s'),
                        Item('r'),
                        Item('b'),
                        Item('u'), Item('v'), Item('w')),
                    ),
                resizable=True
                )

    ######################################################################
    # Trait handlers.
    ######################################################################

    @on_trait_change('scene.activated')
    def scene_active(self):
        self.update_flow()
        self.scene.mlab.gcf().on_mouse_pick(self.onSpacePick, button='Right')
        self.SpaceText = self.scene.mlab.text(0.02, 0.02, "")

    # Note that in the `on_trait_change` call below we listen for the
    # `scene.activated` trait.  This conveniently ensures that the flow
    # is generated as soon as the mlab `scene` is activated (which
    # happens when the configure/edit_traits method is called).  This
    # eliminates the need to manually call the `update_flow` method etc.
    @on_trait_change('s, r, b')
    def update_flow(self):
        x, y, z = self.points
        u, v, w = self.get_uvw()
        self.flow.mlab_source.trait_set(u=u, v=v, w=w)

    @on_trait_change('u')
    def update_u(self):
        self.flow.mlab_source.trait_set(u=self.get_vel('u'))

    @on_trait_change('v')
    def update_v(self):
        self.flow.mlab_source.trait_set(v=self.get_vel('v'))

    @on_trait_change('w')
    def update_w(self):
        self.flow.mlab_source.trait_set(w=self.get_vel('w'))

    def get_uvw(self):
        return self.get_vel('u'), self.get_vel('v'), self.get_vel('w')

    def get_vel(self, comp):
        """This function basically evaluates the user specified system
        of equations using scipy.
        """
        func_str = getattr(self, comp)
        try:
            g = scipy.__dict__
            x, y, z = self.points
            s, r, b = self.s, self.r, self.b
            val = eval(func_str, g,
                        {'x': x, 'y': y, 'z': z,
                         's':s, 'r':r, 'b': b})
        except:
            # Mistake, so return the original value.
            val = getattr(self.flow.mlab_source, comp)
        return val

    ######################################################################
    # Private interface.
    ######################################################################
    def _points_default(self):
        x, y, z = np.mgrid[-50:50:100j,-50:50:100j,-10:60:70j]
        return x, y, z

    def _flow_default(self):
        x, y, z = self.points
        u, v, w = self.get_uvw()
        f = self.scene.mlab.flow(x, y, z, u, v, w)
        f.stream_tracer.integration_direction = 'both'
        f.stream_tracer.maximum_propagation = 200
        src = f.mlab_source.m_data
        o = mlab.outline()
        mlab.view(120, 60, 150)
        return f

    def onSpacePick(self, picker_obj):
        pos = picker_obj.pick_position
        self.SpaceText.text = f"Pick Position: ({pos[0]:>9.6g}, {pos[1]:>9.6g}, {pos[2]:>9.6g})"

if __name__ == '__main__':
    # Instantiate the class and configure its traits.
    lor = Lorenz()
    lor.configure_traits()

When your right click anywhere in the scene, the bottom text on the left gets updated to the picked position, see the image below:

screenshot

TianSong1991 commented 1 year ago

@drocheam Wow, it's very useful for me! Thanks for you answering!