Merge pull request #2 from MrBago/matplotlib_optional

RF - moved hot into try-except so that we can avoid if block later

dipy/viz/projections.py

@@ -13,109 +13,109 @@
     import matplotlib.pyplot as plt
     import matplotlib.tri as tri
     has_mpl = True
+    hot = matplotlib.cm.hot
 except ImportError:
     e_s = "You do not have Matplotlib installed. Some visualization functions"
     e_s += " might not work for you."
     print(e_s)
     has_mpl=False
+    hot = None
 
 import dipy.core.geometry as geo
 
-# Only define these functions if matplotlib is present, otherwise, this module
-# returns 'empty':
-if has_mpl:
-    def sph_project(vertices, val, ax=None, vmin=None, vmax=None,
-                    cmap=matplotlib.cm.hot, cbar=True, triang=False):
 
-        """Draw a signal on a 2D projection of the sphere.
+def sph_project(vertices, val, ax=None, vmin=None, vmax=None,
+                cmap=hot, cbar=True, triang=False):
 
-        Parameters
-        ----------
+    """Draw a signal on a 2D projection of the sphere.
 
-        vertices : (N,3) ndarray
-                    unit vector points of the sphere
+    Parameters
+    ----------
 
-        val: (N) ndarray
-            Function values.
+    vertices : (N,3) ndarray
+                unit vector points of the sphere
 
-        ax : mpl axis, optional
-            If specified, draw onto this existing axis instead.
+    val: (N) ndarray
+        Function values.
 
-        vmin, vmax: floats
-           Values to cut the z
+    ax : mpl axis, optional
+        If specified, draw onto this existing axis instead.
 
-        cmap: mpl colormap
+    vmin, vmax: floats
+       Values to cut the z
 
-        cbar: Whether to add the color-bar to the figure
+    cmap: mpl colormap
 
-        triang: Whether to display the plot triangulated as a pseudo-color plot.
+    cbar: Whether to add the color-bar to the figure
 
-        Returns
-        -------
-        fig : figure
-            Matplotlib figure
+    triang: Whether to display the plot triangulated as a pseudo-color plot.
 
-        Examples
-        --------
-        >>> from dipy.data import get_sphere
-        >>> verts,faces=get_sphere('symmetric724')
-        >>> sph_project(verts,np.random.rand(len(verts)))
+    Returns
+    -------
+    fig : figure
+        Matplotlib figure
 
-        """
-        if ax is None:
-            _, ax = plt.subplots(1)
-        fig = ax.get_figure()
+    Examples
+    --------
+    >>> from dipy.data import get_sphere
+    >>> verts,faces=get_sphere('symmetric724')
+    >>> sph_project(verts,np.random.rand(len(verts)))
 
-        x = vertices[:, 0]
-        y = vertices[:, 1]
+    """
+    if ax is None:
+        _, ax = plt.subplots(1)
+    fig = ax.get_figure()
 
-        my_min = np.nanmin(val)
-        if vmin is not None:
-            my_min = vmin
+    x = vertices[:, 0]
+    y = vertices[:, 1]
 
-        my_max = np.nanmax(val)
-        if vmax is not None:
-            my_max = vmax
+    my_min = np.nanmin(val)
+    if vmin is not None:
+        my_min = vmin
 
-        r = (val - my_min)/float(my_max-my_min)
+    my_max = np.nanmax(val)
+    if vmax is not None:
+        my_max = vmax
 
-        # Enforce the maximum and minumum boundaries, if there are values
-        # outside those boundaries:
-        r[r<0]=0
-        r[r>1]=1
+    r = (val - my_min)/float(my_max-my_min)
 
-        if triang:
-            triang = tri.Triangulation(x, y)
-            ax.tripcolor(triang, r, cmap=cmap)
-        else:
-            cmap_data = cmap._segmentdata
-            red_interp, blue_interp, green_interp = (
-                    interp.interp1d(np.array(cmap_data[gun])[:,0],
-                                    np.array(cmap_data[gun])[:,1]) for gun in
-                                    ['red', 'blue','green'])
+    # Enforce the maximum and minumum boundaries, if there are values
+    # outside those boundaries:
+    r[r<0]=0
+    r[r>1]=1
 
+    if triang:
+        triang = tri.Triangulation(x, y)
+        ax.tripcolor(triang, r, cmap=cmap)
+    else:
+        cmap_data = cmap._segmentdata
+        red_interp, blue_interp, green_interp = (
+                interp.interp1d(np.array(cmap_data[gun])[:,0],
+                                np.array(cmap_data[gun])[:,1]) for gun in
+                                ['red', 'blue','green'])
 
-            for this_x, this_y, this_r in zip(x,y,r):
-                red = red_interp(this_r)
-                blue = blue_interp(this_r)
-                green = green_interp(this_r)
-                ax.plot(this_x, this_y, 'o',
-                        c=[red.item(), green.item(), blue.item()])
 
+        for this_x, this_y, this_r in zip(x,y,r):
+            red = red_interp(this_r)
+            blue = blue_interp(this_r)
+            green = green_interp(this_r)
+            ax.plot(this_x, this_y, 'o',
+                    c=[red.item(), green.item(), blue.item()])
 
-        ax.set_aspect('equal')
-        ax.set_axis_off()
-        if cbar:
-            mappable = matplotlib.cm.ScalarMappable(cmap=cmap)
-            mappable.set_array([my_min, my_max])
-            # setup colorbar axes instance.
-            pos = ax.get_position()
-            l, b, w, h = pos.bounds
-            # setup colorbar axes
-            cax = fig.add_axes([l+w+0.075, b, 0.05, h], frameon=False)
-            fig.colorbar(mappable, cax=cax) # draw colorbar
 
-        ax.set_xlim([-1.1, 1.1])
-        ax.set_ylim([-1.1, 1.1])
+    ax.set_aspect('equal')
+    ax.set_axis_off()
+    if cbar:
+        mappable = matplotlib.cm.ScalarMappable(cmap=cmap)
+        mappable.set_array([my_min, my_max])
+        # setup colorbar axes instance.
+        pos = ax.get_position()
+        l, b, w, h = pos.bounds
+        # setup colorbar axes
+        cax = fig.add_axes([l+w+0.075, b, 0.05, h], frameon=False)
+        fig.colorbar(mappable, cax=cax) # draw colorbar
 
-        return ax
+    ax.set_xlim([-1.1, 1.1])
+    ax.set_ylim([-1.1, 1.1])
+
+    return ax