BUG: [mlab] Fix corner cases for the shape of arguments to surf, and add

	    tests.
BUG: [mlab] Add the forgotten scalars traits to mlab.quiver3d.
BUG: [mlab] make warp_scale play well with extents.

DOC: Add the wigner example to demonstrate extent, warp_scale and
	    ranges.

docs/source/mayavi/auto/mlab_helper_functions.rst

@@ -126,6 +126,8 @@ quiver3d
         :scale_mode: the scaling mode for the glyphs
                      ('vector', 'scalar', or 'none').
 
+        :scalars: optional scalar data.
+
         :mode: the mode of the glyphs. Must be '2darrow' or '2dcircle' or
                '2dcross' or '2ddash' or '2ddiamond' or '2dhooked_arrow' or
                '2dsquare' or '2dthick_arrow' or '2dthick_cross' or
@@ -310,8 +312,13 @@ surf
         :warp_scale: scale of the z axis (warped from
                      the value of the scalar). By default this scale
                      is calculated to give a pleasant aspect ratio to
-                     the plot. You can overright this behavoir by
+                     the plot. You can overright this behavior by
                      specifying a float value.
+                     If you specify a value for warp_scale in
+                     addition to an extent, the warp scale will be
+                     determined by the warp_scale, and the plot be
+                     positioned along the z axis with the zero of the
+                     data centered on the center of the extent.
 
         :name: the name of the vtk object created.
 

docs/source/mayavi/auto/mlab_other_functions.rst

@@ -35,13 +35,19 @@ axes
                  Ranges of the labels displayed on the axes.
                  Default is the object's extents.
 
+        :x_axis_visibility: Whether or not the x axis is visible (boolean)
+
+        :y_axis_visibility: Whether or not the y axis is visible (boolean)
+
         :xlabel: the label of the x axis
 
         :extent: [xmin, xmax, ymin, ymax, zmin, zmax]
                  Default is the object's extents.
 
         :ylabel: the label of the y axis
 
+        :z_axis_visibility: Whether or not the z axis is visible (boolean)
+
 
 
 

enthought/mayavi/tools/helper_functions.py

@@ -114,8 +114,8 @@ def get_all_traits(self):
 
 ############################################################################# 
 
-quiver3d = Pipeline( 
-            doc="""
+class Quiver3D(Pipeline):
+    """
     Plots glyphs (like arrows) indicating the direction of the vectors
     for a 3D volume of data supplied as arguments.
 
@@ -129,11 +129,16 @@ def get_all_traits(self):
     made from the indices of vectors.
     
     If 4 positional arguments are passed the last one must be a callable, f, 
-    that returns vectors. """,
-            _source_function=vector_scatter,
-            _pipeline=[VectorsFactory, ] 
-        )
-quiver3d = document_pipeline(quiver3d)
+    that returns vectors. """
+
+    scalars = Array(help="""optional scalar data.""")
+
+    _source_function = Callable(vector_scatter)
+
+    _pipeline = [VectorsFactory, ]
+
+
+quiver3d = document_pipeline(Quiver3D())
 
 
 def test_quiver3d():
@@ -548,8 +553,15 @@ class Surf(Pipeline):
     warp_scale = Any('auto', help="""scale of the z axis (warped from
                         the value of the scalar). By default this scale
                         is calculated to give a pleasant aspect ratio to
-                        the plot. You can overright this behavoir by
-                        specifying a float value.""")
+                        the plot. You can overright this behavior by
+                        specifying a float value. 
+                        
+                        If you specify a value for warp_scale in
+                        addition to an extent, the warp scale will be
+                        determined by the warp_scale, and the plot be
+                        positioned along the z axis with the zero of the
+                        data centered on the center of the extent.
+                        """)
 
     mask = Array(help="boolean mask array to suppress some data points.")
 
@@ -558,14 +570,39 @@ def __call__(self, *args, **kwargs):
         """
         self.source = self._source_function(*args, **kwargs)
         kwargs.pop('name', None)
-        if self.warp_scale == 'auto' and not 'extent' in kwargs:
+        # Deal with both explicit warp scale and extent, this is 
+        # slightly hairy. The wigner example is a good test case for
+        # this.
+        if 'warp_scale' in kwargs and not kwargs['warp_scale']=='auto' \
+                and 'extent' in kwargs:
+            # XXX: I should use the logging module.
+            print 'Warning: both warp_scale and extent keyword argument' \
+            'specified, the z bounds of the extents will be overridden'
+            xi, xf, yi, yf, zi, zf = kwargs['extent']
+            zo = 0.5*(zi + zf)
             try:
-                xi, xf, yi, yf, zi, zf = self.source.data.bounds
+                si, sf = self.source.data.scalar_range
             except AttributeError:
-                xi, xf, yi, yf, zi, zf = self.source.image_data.bounds
-            zf = 0.3*((xf - xi) + (yf - yi))
+                si, sf = self.source.image_data.scalar_range
+            z_span = kwargs['warp_scale'] * abs(sf - si)
+            zi = zo + si * kwargs['warp_scale'] 
+            zf = zi + z_span
             kwargs['extent'] = (xi, xf, yi, yf, zi, zf)
             kwargs['warp_scale'] = 1.
+        elif self.warp_scale == 'auto':
+                if 'extent' in kwargs and 'warp_scale' in kwargs:
+                    print "Warning: extent specified, warp_scale='auto' " \
+                    "ignored."
+                else:
+                    try:
+                        xi, xf, yi, yf, _, _ = self.source.data.bounds
+                        zi, zf = self.source.data.scalar_range
+                    except AttributeError:
+                        xi, xf, yi, yf, _, _ = self.source.image_data.bounds
+                        zi, zf = self.source.image_data.scalar_range
+                zf = 0.3*((xf - xi) + (yf - yi))
+                kwargs['extent'] = (xi, xf, yi, yf, zi, zf)
+                kwargs['warp_scale'] = 1.
         self.store_kwargs(kwargs)
 
         # Copy the pipeline so as not to modify it for the next call

enthought/mayavi/tools/sources.py

@@ -455,8 +455,13 @@ def reset(self, **traits):
             z = numpy.array([0])
             self.set(x=x, y=y, z=z, trait_change_notify=False)
 
-        dx = x[1, 0] - x[0, 0]
-        dy = y[0, 1] - y[0, 0]
+        # Do some magic to extract the first row/column, independently of
+        # the shape of x and y
+        x = numpy.atleast_2d(x.squeeze().T)[0, :].squeeze()
+        y = numpy.atleast_2d(y.squeeze())[0, :].squeeze()
+
+        dx = x[1] - x[0]
+        dy = y[1] - y[0]
 
         if self.m_data is None:
             ds = ArraySource(transpose_input_array=True)
@@ -712,6 +717,8 @@ def vector_scatter(*args, **kwargs):
     x, y, z, u, v, w = process_regular_vectors(*args)
 
     scalars = kwargs.pop('scalars', None)
+    if scalars is not None:
+        scalars = scalars.ravel()
     name = kwargs.pop('name', 'VectorScatter')
 
     data_source = MGlyphSource()

examples/mayavi/wigner.py

@@ -0,0 +1,73 @@
+"""
+In this example we want to display 3 functions of x and y with a
+surf plot (also called carpet plot).
+
+The important aspect of this example is that the 3 functions should not
+be displayed on top of each other, but side by side. For this we use the
+extent keyword argument. 
+
+In addition, the relative scale between the different plots is important.
+This is why we also use the warp_scale keyword argument, to have the same
+scale on all plots.
+
+Finally, we have to adjust the data bounds: as we want the "horizon" of
+the wigner function in the middle of our extents, we put this to zero.
+
+We add a set of axes and outlines to the plot. We have to play we extents
+and ranges in order to make them fit with the data.
+"""
+
+import numpy
+from enthought.mayavi import mlab
+
+def cat(x, y, alpha, eta=1, purity=1):
+    """ Multiphoton shrodinger cat. eta is the fidelity, alpha the number
+        of photons"""
+    cos = numpy.cos
+    exp = numpy.exp
+    return (1 + eta*(exp(-x**2 -(y-alpha)**2)  + exp(-x**2 -
+    (y+alpha)**2) + 2 * purity * exp(-x**2 - y**2) * cos(2* alpha * x))/(2
+    * (1 + exp(- alpha**2))))/2
+
+x, y = numpy.mgrid[-4:4.15:0.1, -4:4.15:0.1]
+
+mlab.figure(1, size=(500, 250), fgcolor=(1, 1, 1), 
+                                    bgcolor=(0.5, 0.5, 0.5))
+mlab.clf()
+
+cat1 = cat(x, y, 1)
+cat2 = cat(x, y, 2)
+cat3 = cat(x, y, 3)
+
+# The cats lie in a [0, 1] interval, with .5 being the assymptotique
+# value. We want to reposition this value to 0, so as to put it in the
+# center of our extents.
+cat1 -= 0.5
+cat2 -= 0.5
+cat3 -= 0.5
+
+
+cat1_extent = (-14,-6, -4,4, 0,5)
+surf_cat1 = mlab.surf(x-10, y, cat1, colormap='Spectral', warp_scale=5, 
+            extent=cat1_extent, vmin=-0.5, vmax=0.5)
+mlab.outline(surf_cat1, color=(.7, .7, .7), extent=cat1_extent)
+mlab.axes(surf_cat1, color=(.7, .7, .7), extent=cat1_extent, 
+            ranges=(0,1, 0,1, 0,1), xlabel='', ylabel='',
+            zlabel='Probability',
+            x_axis_visibility=False, z_axis_visibility=False)
+
+cat2_extent = (-4,4, -4,4, 0,5)
+surf_cat2 = mlab.surf(x, y, cat2, colormap='Spectral', warp_scale=5, 
+            extent=cat2_extent, vmin=-0.5, vmax=0.5)
+mlab.outline(surf_cat2, color=(0.7, .7, .7), extent=cat2_extent)
+
+cat3_extent = (6,14, -4,4, 0,5)
+surf_cat3 = mlab.surf(x+10, y, cat3, colormap='Spectral', warp_scale=5, 
+            extent=cat3_extent, vmin=-0.5, vmax=0.5)
+mlab.outline(surf_cat3, color=(.7, .7, .7), extent=cat3_extent)
+
+mlab.title('Multi-photons cats Wigner function')
+
+mlab.view(142, -72, 32)
+
+mlab.show()

integrationtests/mayavi/test_mlab.py

@@ -33,6 +33,22 @@ def do(self):
             mlab.clf()
             func()
 
+        ############################################################
+        # Test some specific corner-cases
+        import numpy
+        x, y, z = numpy.mgrid[1:10, 1:10, 1:10]
+        u, v, w = numpy.mgrid[1:10, 1:10, 1:10]
+        s = numpy.sqrt(u**2 + v**2)
+
+        mlab.clf()
+        # Test the extra argument "scalars"
+        mlab.quiver3d(x,y,z,u,v,w,scalars=s)
+
+        # Test surf with strange-shaped inputs
+        X, Y = numpy.ogrid[-10:10, -10:10]
+        Z = X**2 + Y**2
+        mlab.surf(X, Y, Z)
+        mlab.surf(X.ravel(), Y.ravel(), Z)
 
 if __name__ == "__main__":
     t = TestMlab()