Source code for irfpy.cena.plot.axis_cena

''' Module to return an axis that has energy-time diagram for CENA.
'''
import datetime

import numpy as np
import matplotlib.pyplot as plt
from matplotlib.dates import date2num

import logging

import irfpy.cena.cena_mass2 as cena_mass

import irfpy.cena.energy
import irfpy.cena.cena_time

from irfpy.cena.plot import axis_timeseries

def getSimpleEnergyTime(t0, t1, rect=None, vmin=0, vmax=100, validation=None, **kwds):
    ''' Return the energy time diagram for count rate summed over.
    '''
    if validation is None:
        validation = []

    tlist = cena_mass.getobstime(timerange=[t0, t1])

    if len(tlist) == 0:
        raise ValueError('No time interval for CENA mass mode.')

    ### Energy boundaries with 17 values
    yarr = irfpy.cena.energy.getBoundaryE16()
    
    ### Data
    cnt = []

    for t in tlist:
        fulldata = cena_mass.getdataE16(t, validation=validation).getData()  # 16x7x128
        # Sum over mass and direction.
        espec = fulldata.sum(axis=2).sum(axis=1)  # 16 array
        cnt.append(espec)
    cnt = np.ma.masked_array(cnt)

    xarr = [date2num(t) for t in tlist]
    cenatres = 4.5 / 86400. # in second, for redunduncy, 4.5 is used.
    xarr, cnt = axis_timeseries.insert_datagap(xarr, cnt, gap=cenatres)

    ### Extend tlist into number_of_data + 1.
    t00 = xarr[0] - cenatres
    tlist.insert(0, t00)

    ax = axis_timeseries.axis_time_spectra(xarr, yarr, cnt, vmin=vmin, vmax=vmax, rect=rect, **kwds)
    ax.set_yscale('log')
    ax.set_ylim(5, 4000)
    ax.set_xlim(xarr[0], xarr[-1])

    return ax
    

def getSimpleDirectionTime(t0, t1, rect=None, vmin=0, vmax=100, validation=None, **kwds):
    if validation is None:
        validation = []

    tlist = cena_mass.getobstime(timerange=[t0, t1])

    if len(tlist) == 0:
        raise ValueError('No time interval for CENA mass mode.')

    yarr = list(range(7, -1, -1))  # 7 direction.  (+1 for plotting)

    ### Count rate dat
    cnt = []

    for t in tlist:
        fulldata = cena_mass.getdataE16(t, validation=validation).getData()  # 16 x7 x 128
        # Sum over the mass and energy
        dspec = fulldata.sum(axis=2).sum(axis=0)  # 7-length array
        cnt.append(dspec)
    cnt = np.ma.masked_array(cnt)

    # handle the time gap
    xarr = [date2num(t) for t in tlist]
    cenatres = 4.5/86400.
    xarr, cnt = axis_timeseries.insert_datagap(xarr, cnt, gap=cenatres)

    ### Extend tlist into number_of_data + 1.
    t00 = xarr[0] - cenatres
    tlist.insert(0, t00)

    ax = axis_timeseries.axis_time_spectra(xarr, yarr, cnt, vmin=vmin, vmax=vmax, rect=rect, **kwds)
    ax.set_ylim(-0.1, 7.1)
    ax.set_xlim(xarr[0], xarr[-1])

    return ax


def getSimpleMassTime(t0, t1, rect=None, vmin=0, vmax=100, validation=None, **kwds):
    if validation is None:
        validation = []

    tlist = cena_mass.getobstime(timerange=[t0, t1])

    if len(tlist) == 0:
        raise ValueError('No time interval for CENA mass mode.')

    yarr = list(range(0, 129))

    cnt = []

    for t in tlist:
        fulldata = cena_mass.getdataE16(t, validation=validation).getData()
        mspec = fulldata.sum(axis=1).sum(axis=0)
        cnt.append(mspec)

    cnt = np.ma.masked_array(cnt)

    # handle the time gap
    xarr = [date2num(t) for t in tlist]
    cenatres = 4.5/86400.
    xarr, cnt = axis_timeseries.insert_datagap(xarr, cnt, gap=cenatres)

    ### Extend tlist into number_of_data + 1.
    t00 = xarr[0] - cenatres
    tlist.insert(0, t00)

    ax = axis_timeseries.axis_time_spectra(xarr, yarr, cnt, vmin=vmin, vmax=vmax, rect=rect, **kwds)
    ax.set_ylim(-0.1, 128.1)
    ax.set_xlim(xarr[0], xarr[-1])

    return ax