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-2009, Enthought, Inc.

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.set(u=u, v=v, w=w)

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

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

@on_trait_change('w')
def update_w(self):
self.flow.mlab_source.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()

```