Adds slice_max method and update example

example.py

 import numpy as np
-import projection_matrix.projection_matrix as pmp
+import projection_matrix as pmp
 
 # Generate example data
 x = np.linspace(0, 1, 50)
@@ -11,10 +11,14 @@ sigmax = 0.1
 sigmay = 1
 sigmaz = 0.2
 D = (np.exp(-(X-x0)**2/sigmax**2)
+     + 0.5*np.exp(-(X-0.1)**2/sigmax**2)
      + np.exp(-(Y-y0)**2/sigmay**2)
      + np.exp(-(Z-z0)**2/sigmaz**2))
 
-fig, axes = pmp(
-    D, xyz=[x, y, z], labels=['x', 'y', 'z', 'D'], projection=np.max)
+fig, axes = pmp.projection_matrix(
+    D, xyz=[x, y, z], labels=['x', 'y', 'z', 'D'],
+    projection=pmp.slice_max,
+    #projection=np.max
+    )
 fig.savefig('example')
 

projection_matrix/__init__.py

-from .projection_matrix import projection_matrix
+from .projection_matrix import projection_matrix, slice_max

projection_matrix/projection_matrix.py

@@ -3,7 +3,17 @@ import matplotlib.pyplot as plt
 from matplotlib.ticker import MaxNLocator
 
 
-def projection_matrix(D, xyz, labels=None, projection=np.max, max_n_ticks=4,
+def slice_max(D, axis):
+    """ Return the slice along the given axis """
+    idxs = [range(D.shape[j]) for j in range(D.ndim)]
+    max_idx = list(np.unravel_index(D.argmax(), D.shape))
+    for k in np.atleast_1d(axis):
+        idxs[k] = [max_idx[k]]
+    res = np.squeeze(D[np.ix_(*tuple(idxs))])
+    return res
+
+
+def projection_matrix(D, xyz, labels=None, projection=slice_max, max_n_ticks=4,
                       factor=3):
     """ Generate a projection matrix plot
 
@@ -20,7 +30,8 @@ def projection_matrix(D, xyz, labels=None, projection=np.max, max_n_ticks=4,
         labels, the final label is for the dependent variable.
     projection: func
         Function to use for projection, must take an `axis` argument. Default
-        is `np.max()`, to project out a slice along the maximum.
+        is `projection_matrix.slice_max()`, to project out a slice along the
+        maximum.
     max_n_ticks: int
         Number of ticks for x and y axis of the `pcolormesh` plots
     factor: float
@@ -100,3 +111,6 @@ def projection_1D(ax, x, D, xidx, projection):
     ax.yaxis.tick_right()
     ax.yaxis.set_label_position("right")
     return ax
+
+
+