added qt3t with mlab interface

import matplotlib as mpl

import matplotlib.pyplot as plt

from collections import namedtuple

from numpy import flipud, fliplr

from pysen import version, get_theta, get_q

from pysen.config import L2, XDET, YDET

MAX_LEVELS = 20

Side = namedtuple('Side', ('index', 'zdir', 'offset', 'label', 'colorbar', 'data'))

def define_grid(nt=4, nx=3, ny=3, lmax=8.0, dlam=3.0):

    "define z(wavelength), x, y, grid"

    dx   = XDET/nx # default is 30.0/32=0.009375 # m

    levels = findlevels(np.min(Q), np.max(Q))

    H = np.histogram(Q.flatten(), bins=levels)

    D = [

        (11, 'tx[y=1] ', 1, Q[: ,0 ,: ]),           #

        ( 3, 'tx[y=ny]', 0, np.flipud(Q[: ,0 ,: ])),#

        #

        (14, 'yt[x=1] ', 0, np.fliplr(Q[:,:,0].T)), #

        (16, 'yt[x=nx]', 0, Q[: ,: ,nx-1].T),       #

        #

        (15, 'yx[t=1] ', 0, Q[0         ]),         #

        ( 7, 'yx[t=nt]', 0, Q[nt-1      ]),         #

    S = [

        #  XC-LC (top)

        Side(index=11, zdir='', offset=0, label='tx[y=1] ', colorbar=1, data=       Q[: , 0,: ]   ), #

        #  XC-LC (bot)

        Side(index= 3, zdir='', offset=0, label='tx[y=ny]', colorbar=0, data=flipud(Q[: , 0,: ])  ), #

        # LC-YC (high-Q)

        Side(index=14, zdir='', offset=0, label='yt[x=1] ', colorbar=0, data=fliplr(Q[: ,: ,0 ].T)), #

        # LC-YC (low-Q)

        Side(index=16, zdir='', offset=0, label='yt[x=nx]', colorbar=0, data=       Q[: ,: ,-1].T ), #

        # XC-YC front

        Side(index=15, zdir='', offset=0, label='yx[t=1] ', colorbar=0, data=       Q[ 0,: ,: ]   ), #

        # XC-YC back

        Side(index= 7, zdir='', offset=0, label='yx[t=nt]', colorbar=0, data=       Q[-1,: ,: ]   ), #

        ]

    return {'sides': D, 'histogram':H, 'q': Q,

    return {'sides': S, 'histogram':H, 'q': Q, 'grid': (LC,YC,XC),

            'qmin' : qmin, 'lmax': lmax, 'theta': theta}

def plot_side(res, side=0):

    "plot one side of a cube"

    if side==0:

        plot_cube_2d(res)

        return

    _, levels    = res.get('histogram')

    _, lbl, _, d = res.get('sides')[side-1]

    fig, ax = plt.subplots(figsize=(8, 8), dpi=100)

    fig.suptitle(r'$\lambda_{max}$=%g $q_{min}$=%g %s]' % (res.get('lmax'), res.get('qmin'), lbl))

    cnt = ax.contourf(d, levels)

    ax.contour (d, levels, colors='k')

    ax.set_ylabel(lbl[0])

    ax.set_xlabel(lbl[1])

    #ax.set_aspect(1.0)

    fig.colorbar(cnt, ax=ax, orientation='vertical')

def plot_cube_2d(res): # D, H, qmin, lmax):

    "plot 2d cube, to be later assembled with scissors and glue"

def plot_cube_unfolded(res): # D, H, qmin, lmax):

    "plot 2d cube unfolded, to be later assembled with scissors and glue"

    hist, levels = res.get('histogram')

    width  =  levels[1]   - levels[0]

    ax.bar(center, hist, align = 'center', width = width)

    #ax.set_yscale('log')

    ax.grid(True)

    for k, _label, cbar, d in res.get('sides'):

        ax = plt.subplot(4, 4, k)

        cnt = ax.contourf(d, levels)

        ax.contour (d, levels, colors='k', alpha=0.33)

        if cbar:

    for side in res.get('sides'):

        ax = plt.subplot(4, 4, side.index)

        cnt = ax.contourf(side.data, levels)

        ax.contour (side.data, levels, colors='k', alpha=0.33)

        if side.colorbar:

            fig.colorbar(cnt, orientation='horizontal', ax=ax,

                         pad=-0.15, shrink=0.8)

        if k==15:

        if side.index==15:

            ax.text(0.5, 0.5, r'Q=%.3f, $\lambda$=%.1f $\theta$=%.1f' %

                    (res.get('qmin'), res.get('lmax'), np.degrees(res.get('theta'))))

        ax.text(0.01, 0.01, r'IDX=%d' % side.index, size=24)

        ax.get_xaxis().set_ticks([])

        ax.get_yaxis().set_ticks([])

    fig.subplots_adjust(hspace=0.0, wspace=0.0)

def plot_cube_2d(res, side=0):

    "plot one side of a cube"

    if side==0:

        plot_cube_unfolded(res)

        return

    _, levels    = res.get('histogram')

    #_, lbl, _, d = res.get('sides')[side-1]

    xside = res.get('sides')[side-1]

    fig, ax = plt.subplots(figsize=(8, 8), dpi=100)

    fig.suptitle(r'$\lambda_{max}$=%g $q_{min}$=%g %s' % (res.get('lmax'), res.get('qmin'), xside.label))

    cnt = ax.contourf(xside.data, levels)

    ax.contour (xside.data, levels, colors='k')

    ax.set_ylabel(xside.label[0])

    ax.set_xlabel(xside.label[1])

    #ax.set_aspect(1.0)

    fig.colorbar(cnt, ax=ax, orientation='vertical')

def plot_cube_3d(res): # D, H, qmin, lmax):

    "plot 2d cube, to be later assembled with scissors and glue"

    hist, levels = res.get('histogram')

    fig = plt.figure(figsize=(8, 8), dpi=100)

    ax = fig.add_subplot(111, projection='3d')

    # Plot contour surfaces

    Q = res.get('q')

    LC,YC,XC = res.get('grid')

    # Set limits of the plot from coord limits

    xmin, xmax = XC.min(), XC.max()

    ymin, ymax = YC.min(), YC.max()

    zmin, zmax = LC.min(), LC.max()

    #c = ax.contourf(XC[-1,:,:], YC[-1,:,:],  Q[-1,:,:], zdir='z', offset=zmax, levels=levels)

    #c = ax.contourf(XC[:,-1,:],  Q[:,-1,:], LC[:,-1,:], zdir='y', offset=ymin, levels=levels)

    #c = ax.contourf( Q[:,:, 0], YC[:,:, 0], LC[:,:, 0], zdir='x', offset=xmax, levels=levels)

    c = ax.contourf(XC[ 0,:,:], YC[ 0,:,:], Q[ 0,:,:], zdir='z', offset=zmin, levels=levels)

    c = ax.contourf(XC[:, 0,:], Q[:, 0,:], LC[:, 0,:], zdir='y', offset=ymax, levels=levels)

    c = ax.contourf(Q[:,:,-1], YC[:,:,-1], LC[:,:,-1], zdir='x', offset=xmin, levels=levels)

    # Plot edges

    edges_kw = dict(color='k', linewidth=0.5, zorder=1e3)

    ax.plot([xmin, xmax], [ymax, ymax], [zmin, zmin], **edges_kw)

    ax.plot([xmin, xmax], [ymin, ymin], [zmin, zmin], **edges_kw)

    ax.plot([xmin, xmin], [ymin, ymax], [zmin, zmin], **edges_kw)

    ax.plot([xmax, xmax], [ymin, ymax], [zmin, zmin], **edges_kw)

    ax.plot([xmin, xmin], [ymax, ymax], [zmin, zmax], **edges_kw)

    ax.plot([xmax, xmax], [ymax, ymax], [zmin, zmax], **edges_kw)

    ax.set(xlim=[xmin, xmax], ylim=[ymin, ymax], zlim=[zmin, zmax])

    ax.set(xlabel='X', ylabel='Y', zlabel='T',

           xticks=[],  yticks=[],  zticks=[])

    ax.view_init(150, 30, -180, vertical_axis='y')

    #ax.set_box_aspect(None, zoom=0.9)

    ax.set_axis_off()

    fig.colorbar(c, ax=ax, fraction=0.02, pad=0.1, label=r'Q [$\AA^{-1}$]')

def main():

    "main script"

    nchan = (42, 32, 32) # FIXME (hard coded dimensions)

    parser = argparse.ArgumentParser(description='plot xyt cube')

    parser.set_defaults(mode='flat', qmin=0.05, lmax=8.0, dlam=3.1, theta=None, side=0, cmap='hot')

    parser.set_defaults(mode='flat', qmin=0.05, lmax=8.0, dlam=3.1, theta=None, side=0, cmap=None)

    parser.add_argument('--qmin', '-q',  dest='qmin', type=float,

                        help='set qmin          (default=%(default)s)')

    parser.add_argument('--lmax', '-l',  dest='lmax', type=float,

                        help='set the scattering ange (overrides qmin, default=%(default)s)')

    parser.add_argument('--dlam', '-d',  dest='dlam', type=float,

                        help='set delta lambda  (default=%(default)s)')

    #parser.add_argument('--3d'  , '-3',  dest='mode', action='store_const', const='3d'  ,

    #                        help='plot 3-dimensional cube (requires mayavi)')

    #parser.add_argument('--2d'  , '-2',  dest='mode', action='store_const', const='flat',

    #                        help='plot "flat" cube (cut-and-build)')

    parser.add_argument('--3d'  , '-3',  dest='mode', action='store_const', const='3d'  ,

                            help='plot 3-dimensional cube')

    parser.add_argument('--2d'  , '-2',  dest='mode', action='store_const', const='flat',

                            help='plot "flat" cube (cut-and-build)')

    parser.add_argument('--side', '-S',  dest='side', type=int, choices=range(0,7),

                        help='plot only one side (flat mode)')

    parser.add_argument('--cmap', '-C',  dest='cmap',

    res = calc_cube(args.qmin, lmax=args.lmax, dlam=args.dlam, nchan=nchan)

    if args.mode == 'flat':

        if args.side:

            plot_side(res, args.side)

        else:

            plot_cube_2d(res)

        plot_cube_2d(res, args.side)

    elif args.mode == '3d':

        plot_cube_3d(res)

    else:

        raise RuntimeError('Unknown mode %s' % args.mode)