plot_input_fortracc.py

import gzip as gz
import matplotlib.pyplot as plt
from string import Template
from matplotlib.colors import Normalize
from metpy.plots import colortables as ct
import os
import numpy as np
import pandas as pd
from matplotlib.colorbar import make_axes
import datetime as dt


class Params():
    def __init__(self, year, month, day, hour, minute):
        self.year = year
        self.month = month
        self.day = day
        self.hour = hour
        self.minute = minute

    def __str__(self):
        return str(self.to_int())

    def __add__(self, delta):
        cdatetime = Template(
            "$year-$month-$day $hour:$minute"
        ).substitute(**self.to_str())
        cdatetime = dt.datetime.strptime(cdatetime, '%Y-%m-%d %H:%M')
        deltatime = dt.timedelta(minutes=delta*12)
        ndatetime = cdatetime+deltatime
        temp = Params(
            year=ndatetime.year, month=ndatetime.month, day=ndatetime.day,
            hour=ndatetime.hour, minute=ndatetime.minute)
        return temp

    def __sub__(self, delta):
        cdatetime = Template(
            "$year-$month-$day $hour:$minute"
        ).substitute(**self.to_str())
        cdatetime = dt.datetime.strptime(cdatetime, '%Y-%m-%d %H:%M')
        deltatime = dt.timedelta(minutes=delta*12)
        ndatetime = cdatetime-deltatime
        temp = Params(
            year=ndatetime.year, month=ndatetime.month, day=ndatetime.day,
            hour=ndatetime.hour, minute=ndatetime.minute)
        return temp

    def to_int(self):
        return dict(
            year=int(self.year),
            month=int(self.month),
            day=int(self.day),
            hour=int(self.hour),
            minute=int(self.minute),
        )

    def to_str(self):
        return dict(
            year=str(self.year).zfill(4),
            month=str(self.month).zfill(2),
            day=str(self.day).zfill(2),
            hour=str(self.hour).zfill(2),
            minute=str(self.minute).zfill(2),
        )

    def explode_int(self):
        # return [self.year, self.month, self.day, self.hour, self.minute]
        return list(self.to_int().values())

    def explode_str(self):
        return list(self.to_str().values())


def get_all_filepaths(base_path_in, base_path_ou):
    """
    base_path: string
        directory of output fortracc or output IDL files.
    """

    all_filepaths_input = list(
        os.path.join(base_path_in, file) for file in os.listdir(base_path_in)
    )

    all_filepaths_output = list(
        os.path.join(base_path_ou, file) for file in os.listdir(base_path_ou)
    )

    return all_filepaths_input, all_filepaths_output


def load_raw_input_fortracc(param, files_in, ext='.raw'):
    substr = Template(
        "_$year$month$day$hour$minute$ext"
    ).substitute(ext=ext, **param.to_str())

    files = sorted(list(
        filter(lambda file: file.endswith(substr), files_in)
    ))

    if ext == '.raw.gz':
        with gz.open(files[0], 'r') as file:
            mat_reflect = np.frombuffer(
                file.read(), dtype=np.float32
            ).reshape((241, 241)).copy()
    else:
        with open(files[0], 'rb') as file:
            mat_reflect = np.frombuffer(
                file.read(), dtype=np.float32
            ).reshape((241, 241)).copy()

    mat_reflect[mat_reflect == -99] = np.nan
    mat_reflect = np.flip(mat_reflect, axis=0)

    return mat_reflect


def load_raw_output_fortracc(param, files_ou, ext='.raw'):
    substr = Template(
        "_$year$month$day$hour$minute$ext"
    ).substitute(ext=ext, **param.to_str())

    files = sorted(list(
        filter(
            lambda file: file.endswith(substr), files_ou
        )
    ))

    if ext == '.raw.gz':
        with gz.open(files[0], 'r') as file:
            mat_clusters = np.frombuffer(
                file.read(), dtype=np.int16
            ).reshape((241, 241)).copy()
    else:
        with open(files[0], 'rb') as file:
            mat_clusters = np.frombuffer(
                file.read(), dtype=np.int16
            ).reshape((241, 241)).copy().astype(np.float)

    mat_clusters[mat_clusters == 0] = np.nan

    return mat_clusters


def load_grid(filepath_lon, filepath_lat):
    with open(filepath_lon, 'rb') as file:
        lons = np.frombuffer(file.read(), dtype=np.float32).reshape((241, 241))

    with open(filepath_lat, 'rb') as file:
        lats = np.frombuffer(file.read(), dtype=np.float32).reshape((241, 241))

    return lons, lats


def run_fortracc_input():
    BASE_PATH_RAW_INPUT_FORTRACC = "/home/adriano/thunderstorm/data-cappi/" \
                                   "test/mat_raw_reflectivity"
    BASE_PATH_RAW_OUTPUT_FORTRACC = "/home/adriano/thunderstorm/data-cappi/" \
                                    "test/mat_raw_cluster"
    PATH_LON = "/home/adriano/thunderstorm/data-cappi/test/nav-te.lon"
    PATH_LAT = "/home/adriano/thunderstorm/data-cappi/test/nav-te.lat"

    files_in, files_ou = get_all_filepaths(
        BASE_PATH_RAW_INPUT_FORTRACC, BASE_PATH_RAW_OUTPUT_FORTRACC
    )

    param = Params(year=2014, month=1, day=18, hour=13, minute=48)

    matr_input = load_raw_input_fortracc(param, files_in)
    matc_output = load_raw_output_fortracc(param, files_ou)
    lons_grid, lats_grid = load_grid(PATH_LON, PATH_LAT)

    fig, (ax, ax2) = plt.subplots(1, 2, figsize=(10, 5))

    ax.set_title(Template(
        "Input ForTraCC Reflectivity\n$year-$month-$day $hour:$minute"
    ).substitute(**param.to_str()))

    cmap = ct.get_colortable('NWSReflectivity')
    norm = Normalize(0, 65)

    plot = ax.pcolor(lons_grid, lats_grid, matr_input, cmap=cmap, norm=norm)
    cax, kw = make_axes(ax, location='bottom', pad=0.075, shrink=.9)
    out = fig.colorbar(plot, cax=cax, extend='both', **kw)
    out.set_label('Reflectivity in dBZ', size=10)

    ax2.set_title(Template(
        "Output ForTaCC Clusters\n$year-$month-$day $hour:$minute"
    ).substitute(**param.to_str()))

    plot = ax2.pcolor(lons_grid, lats_grid, matc_output, cmap='cubehelix')
    cax, kw = make_axes(ax2, location='bottom', pad=0.075, shrink=.9)
    out = fig.colorbar(plot, cax=cax, extend='both', **kw)
    out.set_label('Clusters ID', size=10)

    plt.savefig('plot_input_output_fortracc.png', dpi=72, bbox_inches='tight',
                transparent=False, pad_inches=0.1)

    plt.show()


if __name__ == "__main__":
    run_fortracc_input()