Lorenz ui example — mayavi 4.7.2 documentation

# Lorenz ui example

This example displays the trajectories for the Lorenz system of equations using mlab along with the z-nullcline. It provides a simple UI where a user can change the parameters and the system of equations on the fly. This primarily demonstrates how one can build powerful tools with a UI using Traits and Mayavi.

For explanations and more examples of interactive application building with Mayavi, please refer to section Building applications using Mayavi.

Python source code: `lorenz_ui.py`

    # 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.
        ######################################################################

        # 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, scene.activated')
        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

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