apps130311_water.mapcs

''' Simple plotting routine for south pole region.
'''
import sys
import datetime
import pickle as pickle
import gzip

import numpy as np
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap

import irfpy.util.gridsphere
import irfpy.cena.cena_mass_time

import mapinmedata


def create_data(md):
    flux = np.zeros_like(md.lonc) - 1e20
    nobs = np.zeros_like(md.lonc) - 1e20

    t0 = datetime.datetime.max
    t1 = datetime.datetime.min

    countrate_map = irfpy.util.gridsphere.SimpleGridSphere(nlon=md.nlon, nlat=md.nlat)

    print("Getting time range")
    for idx in range(len(md)):
        dat = md[idx]
        if len(dat) != 0:
            t0local = min(dat)
            t1local = max(dat)
            if t0local < t0: t0 = t0local
            if t1local > t1: t1 = t1local
    print("... Getting time range, done.", t0, t1)

    try:
        print("Getting data")
        tlist, timeseries_cnt = irfpy.cena.cena_mass_time.getfulldata(t0, t1)
        print("... Getting data, done.", timeseries_cnt.shape)
    except RuntimeError as e:
        return (np.ma.masked_less(flux, 0).reshape([md.nlat, md.nlon]),
                np.ma.masked_less(nobs, 0.1).reshape([md.nlat, md.nlon]),
                None)

    print("Calculating count rate for center 3 channels in SV2")
    for idx in range(len(md)):
        dat = md[idx]
        aveflx = 0
        numobs = 0
        if len(dat) != 0:
            for t in dat:
                try:
                    i = tlist.index(t)
                except:
                    continue
                f = timeseries_cnt[4:12, 2:5, :64, i].sum(2).sum(1)   # (8,)
                countrate_map.append_value(idx, f)
                if not f.mask.any():
                    aveflx += f.sum()
                    numobs += 1
            if numobs != 0:
                flux[idx] = aveflx / numobs
                nobs[idx] = numobs

    print("... Calculation done", nobs.shape, flux.shape)
                    
    flux = np.ma.masked_less(flux, 0)
    flux.shape = [md.nlat, md.nlon]
    print(flux.shape)

    nobs.shape = [md.nlat, md.nlon]

    return flux, nobs, countrate_map
#    return flux, countrate_map.get_counts()

def main():
    '''Main script'''

    md = mapinmedata.orbitdata_reader(sys.argv[1])
    lonarr, latarr = md.get_cgrid()
    nlist = md.get_counts()

    fig = plt.figure(figsize=(12, 6))
    ax1 = fig.add_subplot(1, 2, 1)
    ax2 = fig.add_subplot(1, 2, 2)

    cnts, nobs, countratemap = create_data(md)
    
    m = Basemap(projection='spaeqd', boundinglat=-10, lon_0=90, ax=ax1)
    m.drawcoastlines(color='0.7')
#    m.fillcontinents(color='coral', lake_color='aqua')
    m.drawparallels(np.arange(-80, 81, 10))
    m.drawmeridians(np.arange(-180, 181, 30), labels=[1, 0, 1, 1])
#    m.drawmapboundary(fill_color='aqua')
    x, y = m(lonarr, latarr)
    nobs = np.ma.masked_less(nobs, 0.1)
    if not nobs.mask.all():
        m.pcolor(x, y, nobs)
    
    m = Basemap(projection='spaeqd', boundinglat=-10, lon_0=90, ax=ax2)
#    m = Basemap(projection='spaeqd', boundinglat=0, lon_0=90, ax=ax2)
    m.drawcoastlines(color='0.7')
#    m.fillcontinents(color='coral', lake_color='aqua')
    m.drawparallels(np.arange(-80, 81, 10))
    m.drawmeridians(np.arange(-180, 181, 30), labels=[0, 1, 1, 1])
#    m.drawmapboundary(fill_color='aqua')
    x, y = m(lonarr, latarr)
    if not cnts.mask.all():
        m.pcolor(x, y, np.ma.log(np.ma.masked_invalid(cnts)))
    
    fig.savefig(sys.argv[1] + '_mc.png')

    if countratemap != None:
        fff = gzip.open(sys.argv[1] + "_sv2cnt.txt.gz", 'w')
        pickle.dump(countratemap, fff)
        fff.close()
    return md, cnts

if __name__ == "__main__":
    md, dat = main()