swim_mass_ene_spectra

A sample code for SWIM mass-energy matrix

''' A sample code for SWIM mass-energy matrix
'''

import os
import sys
import logging
logging.basicConfig()
import datetime
import math
import dateutil.parser

from optparse import OptionParser

import matplotlib.pyplot as plt
import numpy as np
import scipy as sp
import matplotlib.dates as mdate

from irfpy.swim import swim_mass, swim_start


def main(t0, t1, filter=None, max=None):
    ''' Main code.
    '''
    if filter == None:
        filter = []

    # Mass spectra
    tlist = swim_mass.getobstime(timerange=[t0, t1])
    cntm = np.array([swim_mass.getdata(t, filter=filter).getData() for t in tlist])
    # N x E x D x M

    # Average over the time (N) direction.
    cntm = cntm.mean(0)
    # Now cntm is E16 x D16 x M32

    fig = plt.figure(figsize=[10, 10])
    axs = []
    for iy in range(5):
        for ix in range(4):
            axs.append(fig.add_subplot(5, 4, iy * 4 + ix + 1))

    xax = np.arange(33)  # Mass direction
    yax = np.arange(17)  # Energy direction
    xm, ym = np.meshgrid(xax, yax)

    # Directional plot
    for idir in range(16):
        if max == None:
            vmax = cntm[:, idir, :].max()
        else:
            vmax = max

        axs[idir].pcolor(xm, ym, cntm[:, idir, :], vmin=0, vmax=vmax)  # .T not needed.
        print(idir, cntm[:, idir, :].max())

        axs[idir].set_title('D=%2d' % idir)


#        if idir >= 13:   # Y axis shown
#            plt.setp(axs[idir], yticks=[])
#        elif idir == 12:  # X, Y shown
#            pass
#        elif idir % 4 == 0:   # X axis shown.
#            plt.setp(axs[idir], xticks=[])
#        else:
#            plt.setp(axs[idir], xticks=[], yticks=[])

    # Sum in D
    if max == None:
        vmax = cntm.mean(1).max()
    else:
        vmax = max
    axs[16].pcolor(xm, ym, cntm.mean(1), vmin=0, vmax=max)
    axs[16].set_title('Ave(D)')

    # Show Mask
    onecnt = np.ones([16, 16, 32])
    for f in filter:
        onecnt = f(onecnt)
    onecnt = onecnt.sum(1)
    print(onecnt)
    axs[17].pcolor(xm, ym, onecnt, vmin=0, vmax=16)
    axs[17].set_title('Mask')


    for ax in axs:
        ax.set_xlim(0, 32)
        ax.set_ylim(0, 16)

    for i, ax in enumerate(axs):
        if i % 4 != 0:
            plt.setp(ax, yticks=[])
        else:
            ax.set_ylabel('iE')
        if i < 16:
            plt.setp(ax, xticks=[])
        else:
            ax.set_xlabel('iM')

    return fig


if __name__ == '__main__':
    usage = "%prog [options] t0 t1"
    parser = OptionParser(usage)

    parser.add_option('-o', '--output-prefix', dest='outprefix',
        default='swim_mass_ene_spectra',
        help='Prefix of output file name')
    parser.add_option('-a', '--auto-color-scale', dest='autocolorscale',
        action='store_true', default=False,
        help='Prefix of output file name')
    parser.add_option('-x', '--max', dest='vmax',
        type='float', default=1.0,
        help='Maximum of colorbar')

    options, args = parser.parse_args()

    if len(args) != 2:
        parser.error('t0 and t1 should be specified.')

    if options.autocolorscale:
        vmax = None
    else:
        vmax = options.vmax

    t0 = dateutil.parser.parse(args[0])
    t1 = dateutil.parser.parse(args[1])

    figraw = main(t0, t1, max=vmax)
    figraw.axes[0].set_title('L0: D=0')
    figraw.savefig('%s_l0.png' % options.outprefix)

    figl1 = main(t0, t1, filter=[swim_mass.filter.remove_one_count_event_l1], max=vmax)
    figl1.axes[0].set_title('L1: D=0')
    figl1.savefig('%s_l1.png' % options.outprefix)

    figl2 = main(t0, t1, filter=[swim_mass.filter.remove_one_count_event_l2], max=vmax)
    figl2.axes[0].set_title('L2: D=0')
    figl2.savefig('%s_l2.png' % options.outprefix)

    figl3 = main(t0, t1, filter=[swim_mass.filter.remove_one_count_event_l3], max=vmax)
    figl3.axes[0].set_title('L3: D=0')
    figl3.savefig('%s_l3.png' % options.outprefix)

    figl4 = main(t0, t1, filter=[swim_mass.filter.remove_one_count_event_l4], max=vmax)
    figl4.axes[0].set_title('L4: D=0')
    figl4.savefig('%s_l4.png' % options.outprefix)