appl121105_iotorus.main_image_etdiagram_plot

''' Plotting the E-t diagram
'''
import os
import sys

import re
import datetime
import dateutil.parser

import numpy as np
import matplotlib.pyplot as plt
import matplotlib.dates as mpldates
import matplotlib.colors as mplcolors

import pyana.util.datafile

rj = 71492.


def main(dirname):

#    fnconvention = re.compile('(2031-05-08T0[0-9]:[0-9][0-9]:[0-9][0-9])-m([0-9]+)-([0-9]+).txt')
    fnconvention = re.compile('([0-9][0-9][0-9][0-9]-[0-9][0-9]-[0-9][0-9]T[0-9][0-9]:[0-9][0-9]:[0-9][0-9])-m([0-9]+)-([0-9]+).txt')

    fn = os.listdir(dirname)
#    print fn
    fn.sort()

    timelist = []
    masslist = []
    energylist = []
    filenamelist = []

    for f in fn:
        filename = os.path.join(dirname, f)
#        print f, fnconvention.match(f), fnconvention.search(f).groups()
        fnsearch = fnconvention.search(f)
        if fnsearch == None:
#            print 'Not a valid file name, %s' % f
            continue
        else:
            print('Valid file', f)
        t0, mass, eidx = fnsearch.groups()

        t0 = dateutil.parser.parse(t0)
        mass = float(mass)
        eidx = int(eidx)

        timelist.append(t0)
        masslist.append(mass)
        energylist.append(eidx)
        filenamelist.append(filename)

    maxestep = max(energylist) + 1
    print('Max energy step =', maxestep)

    tlist = np.unique(timelist)
    tlen = len(tlist)
    print('nTime =', tlen)

    mlist = np.unique(sorted(masslist))
    nmass = len(mlist)
    print('nMass = ', nmass)

    diff_flux = np.zeros([tlen, maxestep, nmass]) - 1e20
    iopos = np.zeros([3, tlen]) - 1e20
    scpos = np.zeros([3, tlen]) - 1e20

    for t0, ie, mass, fn in zip(timelist, energylist, masslist, filenamelist):
        tidx = tlist.tolist().index(t0)
        midx = mlist.tolist().index(mass)

        fp = open(fn)
        dfile = pyana.util.datafile.Datafile()
        dfile.readheader(fp)  # Only header analysis is done by API. Data part is in the following
        drfile = pyana.util.datafile.DatafileReader(dfile, missing_return="Unknown")
        vals = np.loadtxt(fp)   # (N, 9) array

        diff_flux[tidx, ie, midx] = vals[:, 8].mean()

        # IO/SC position
        iop = np.array([float(i) for i in dfile.header['SUN_SYS3, IO'].split()[1:]])
        scp = np.array([float(i) for i in dfile.header['SCPOS'].split()])
        
        iopos[:, tidx] = iop
        scpos[:, tidx] = scp
        
    return diff_flux, tlist.tolist(), iopos, scpos


if __name__ == "__main__":
    diff_flux, tlist, iopos, scpos = main(sys.argv[1])
    fig = plt.figure()
    diff_flux = np.ma.masked_where(diff_flux < 0, diff_flux)

    print('DIFFFLUX.shape=', diff_flux.shape)
    tused = diff_flux.shape[0]
    timelist = tlist[:tused]
    timelist.append(timelist[-1] + datetime.timedelta(minutes=10))
    timelist = np.array(mpldates.date2num(timelist))
    t0 = timelist[0]
    timelist = (timelist - t0) * 24
    print(timelist)    

    effective_crosssection = [1.18, 1.88, 0.88] 
                    # Reference Xsection 1e-15. For H, O, S.
    plasma_fraction = [0.1, 0.41, 0.25]   # Density fraction of H+, On+, Sn+

    estep = np.array([10 * 1.5 ** i for i in range(17)])

    for i in range(len(diff_flux[0, 0, :])):
        ax = fig.add_subplot(diff_flux.shape[2] + 1, 1, i+1)

        ### Differential flux is #/cm2 sr ev s
        dflx = effective_crosssection[i] * plasma_fraction[i] * diff_flux[:, :, i]
        ### Convert to Energy flux
        eflx = dflx * estep[:-1]
        ### Convert to count rate
        gfact = 5e-5   # in cm2 sr eV/eV
        cntr = gfact * eflx

        ### Now, the plotting value
#        tbp = dflx
#        vmax = 4
#        vmin = -1
#        clbl = 'Differential Flux [#/cm2 sr eV s]'

#        tbp = eflx
#        vmax = 6
#        vmin = 0
#        clbl = 'Energy Flux [eV/cm2 sr eV s]'

        tbp = cntr
        vmax = 2
        vmin = -4
        clbl = 'Count rate [counts/s]'

#        pc = ax.pcolor(timelist, estep, np.log10(tbp).T, vmax=vmax, vmin=vmin)
        pc = ax.pcolor(timelist, estep, tbp.T, norm=mplcolors.LogNorm(vmax=10 ** vmax, vmin=10 ** vmin))
        ax.set_yscale('log')
        ax.set_ylabel('E [eV]')
        ax.set_ylim(10, 10 * 1.5 ** 16)
        ax.set_xlim(0, timelist.max())

    ax = fig.add_subplot(diff_flux.shape[2] + 1, 1, diff_flux.shape[2] + 1)

    ### Distances from Jupiter, and SC-Io distance
    sclen = np.sqrt((scpos ** 2).sum(0))
    scpos_io = scpos - iopos
    scpos_io_len = np.sqrt((scpos_io ** 2).sum(0))
    ax.plot(timelist[:-1], sclen / rj, 'b')
    ax.plot(timelist[:-1], scpos_io_len / rj, 'g')
    iolen = np.sqrt((iopos ** 2).sum(0))
#    ax.plot(timelist[:-1], iolen / rj, 'r')
    ax.set_xlim(0, timelist.max())
    ax.set_ylim(0, 50)
    ax.set_ylabel('Distance [Rj]')

    ### Jupiter-Io-SC Angle.
    # Conver to io frame
    jupiter_io = -iopos
    jupiter_io_len = iolen

    ax2 = ax.twinx()

    # Angle
    innpr = (scpos_io * jupiter_io).sum(0) / (jupiter_io_len * scpos_io_len)
    print(innpr.max(), innpr.min())
    ang = np.rad2deg(np.arccos(innpr))
    # Sign
    outer = np.cross(jupiter_io.T, scpos_io.T)[:, 2]   # z-component
    sign = -np.sign(outer)  # If negative, the angle is positive.
    angp = np.ma.masked_where(sign >= 0, ang)
    angn = np.ma.masked_where(sign < 0, ang)
#    ang = (ang + 360) % 360  # -180 to 180 => 0 to 360
    print(angp)
    print(angn)
    ax2.plot(timelist[:-1], angp, 'r')
    ax2.plot(timelist[:-1], angn, 'k')
    ax2.set_ylabel('J-I-S angle')
    ax2.set_ylim(0, 180)
    ax2.yaxis.set_ticks(list(range(0, 181, 45)))
    
    ax.set_xlabel('Time [h] after %s' % tlist[0].strftime('%FT%T'))

    plt.subplots_adjust(bottom=0.1, right=0.75, top=0.9)
    cax = plt.axes([0.85, 0.1, 0.02, 0.8])
    clb = fig.colorbar(pc, cax=cax)
    clb.set_label(clbl)
    

    fig.savefig('etdiagram.png')
    fig.savefig('etdiagram.eps')