plot_flexpart.py

#!/usr/bin/env/python3
# -*- coding: utf-8 -*-

"""
Plotting script for FlexPart output for the FAZE-In project

Project: FAZE-In

Usage : plotting_flexpart_output.py startdate datapath

Args:
    startdate: Start date of data
    datapath: Description

"""

__author__ = 'C. Symonds'
__copyright__ = 'Copyright 2020, University of Leeds'
__credits__ = ['C. Symonds', 'L. Kiely', 'R.Rigby']
__license__ = 'MIT'
__version__ = '0.3.0'
__maintainer__ = 'C. Symonds'
__email__ = 'C.C.Symonds@leeds.ac.uk'
__status__ = 'Development'

# Built-in/Generic Imports
from os import path, makedirs
import dateutil.parser as dateparse
from sys import stderr

# Libs
import matplotlib as mpl
mpl.use('Agg')
from cartopy.mpl.ticker import LongitudeFormatter, LatitudeFormatter
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import iris
import matplotlib.pyplot as plt
import numpy as np
import argparse
from math import log10,floor

#iris.FUTURE.netcdf_promote = True


class ArgumentsError(Exception):
    '''
    Exception raised when there is an error detected in the argument list.
    '''
    def __init__(self, msg):
        stderr.write('[FATAL ERROR] : %s' % msg )
        exit(9)


class FatalError(Exception):
    '''
    Exception raised when there is an unrecoverable error encountered.
    '''
    def __init__(self, msg):
        stderr.write('[FATAL ERROR] : %s' % msg )
        exit(9)


class FileError(Exception):
    '''
    Exception raised when contents of files are not as expected
    '''
    def __init__(self,msg):
        stderr.write('[FILE ERROR] : %s' % msg )
        exit(9)


def getargs():
    '''
    Plotting data generated by FLEXPART run between two dates -
    Retrieves the two dates from command line arguments.

    In : None
    Out: Date(s) of data in "YYYY-MM-DD" format
    '''

    parser = argparse.ArgumentParser(description=(
        'Retrieve data from ECMWF and convert to RELEASES text format.\n'+
        'Requires date for GFAS data in YYYY-MM-DD format as input argument.'
        ))

    parser.add_argument('startdate',
                        type=str,
                        help='[REQUIRED]\nFirst date for which data is available. Should be in format "YYYY-MM-DD"\n'+
                             'If data from only one date required, give that date here')

    parser.add_argument('datapath',
                        type=str,
                        help='Location of FLEXPART nc file. PNGs will be put here also.')

    parser.add_argument('--var_name', '-v',
                        type=str,
                        help='Name of the variable to be plotted',
                        default='BC')

    args = parser.parse_args()

    # Validate the date

    startdate=dateparse.parse(args.startdate).date()

    if ( startdate.year != int(args.startdate[:4]) and startdate.month <= 12 and startdate.day <= 12 ):
        raise ArgumentsError("Start date was not in ISO 8601 format (YYYY-MM-DD), and date "+
        "parser cannot infer true date as both month and day are below 12 and "+
        "could be month-first or day-first format.\n"+
        "Please retry using the recommended ISO 8601 format (YYYY-MM-DD)")
    else:
        strtstr = startdate.strftime("%Y%m%d")

    # Validate the output path, or generate it from default value

    datapath = args.datapath

    if not path.exists(datapath):
        print('Directory to write png files to'
              + ' does not exist\nAttempting to create:')
        try:
            makedirs(datapath)
        except:
            raise FatalError('Could not create directory '+ datapath +'\n')
        else:
            print ("Success!\n")

    if datapath and not path.isdir(datapath):
        raise FatalError(datapath + ' exists but is not a directory\n')

    var_name = args.var_name

    return (strtstr,datapath,var_name)


def main():

    (date,datapath,var_name) = getargs()

    constrain_at = iris.Constraint(name=var_name, height=lambda cell: cell.point <= 20.0)

    nc_file = path.join(datapath,'grid_conc_'+date+'000000.nc')

    try:
        cube = iris.load_cube(nc_file, constrain_at)
    except:
        raise FileError("Could not find file {}".format(path.join(datapath,'grid_conc_'+date+'000000.nc')))

    cube_lat = cube.coord('grid_latitude')
    cube_lon = cube.coord('grid_longitude')
    cube_time = cube.coord('time')
    cube_height = cube.coord('height')

    lats = cube_lat.points
    lons = cube_lon.points
    heights = cube_height.points

    lat_min = int(np.ceil(lats.min()))
    lat_max = int(np.floor(lats.max()))
    lon_min = int(np.ceil(lons.min()))
    lon_max = int(np.floor(lons.max()))

    xx, yy = np.meshgrid(lons, lats)

    cube_datetime = cube_time.units.num2date(cube_time.points)

    ##average up to 20m###
    data = cube[0,0,:,:,:,:].data
    for index, height in enumerate(heights):

        if index == 0:
            delta_h = height
        else:
            delta_h = height - heights[index-1]

        data[:,index,:,:] *= (delta_h/(1000*heights[-1]))


    data = np.sum(data,axis = 1)

    i_height = int(round(heights[-1]))

    data_min = np.nanmin(data)
    data_max = np.nanmax(data)

    col_bound_min = 0
    col_bound_max = 100
    vals = np.array([25,50,100,150,200,250,300,350,400,450,500,600,700,800,900,1000,1500,2000,2500,3000,3500,4000,4500,5000,6000,7000,8000,10000,15000,20000])
    col_bounds = vals/200
    cmap = plt.get_cmap('RdYlBu_r', 100)
    newcolors = cmap(np.linspace(0, 1, 100))
    newcmp = mpl.colors.ListedColormap(newcolors)
    norm = mpl.colors.BoundaryNorm(col_bounds,newcmp.N)

    font = {'weight' : 'bold',
            'size'   : 4}

    mpl.rc('font', **font)


    for t in range(data.shape[0]):

        t_data = data[t,:,:]

        if np.where(np.isnan(t_data) == False)[0].shape[0] < 5:
            continue

        filename_time = cube_datetime[t].strftime('%Y%m%d%H%M')
        title_time = cube_datetime[t].strftime('%Y-%m-%d %H:%M')

        plt.figure(figsize = (3.375,3))


        plot_ax = plt.axes(projection=ccrs.PlateCarree())

        plot_ax.set_ylim(bottom=lat_min, top=lat_max)
        plot_ax.set_xlim(left=lon_min, right=lon_max)
        plot_ax.set_xticks(range(lon_min + 4, lon_max, 10), crs=ccrs.PlateCarree())
        plot_ax.set_yticks(range(lat_min + 4, lat_max, 10), crs=ccrs.PlateCarree())
        lon_formatter = LongitudeFormatter(number_format='.1f', degree_symbol='')
        lat_formatter = LatitudeFormatter(number_format='.1f', degree_symbol='')
        #cbarformatter = "{:4.{prec}f}".format(x, prec=min(2,2-floor(log10(abs(x)))))
        plot_ax.xaxis.set_major_formatter(lon_formatter)
        plot_ax.yaxis.set_major_formatter(lat_formatter)

        plot_ax.add_feature(cfeature.LAND)
        plot_ax.add_feature(cfeature.OCEAN)
        plot_ax.add_feature(cfeature.COASTLINE)

        plt.contourf(xx,
                     yy,
                     t_data,
                     col_bounds,
                     cmap=newcmp,
                     vmin=data_min,
                     vmax=data_max,
                     zorder=10,
                     alpha = 0.4,
                     extend='both',
                     norm = norm)

        plt.clim(col_bound_min, col_bound_max)
        plt.colorbar(orientation='horizontal')

        title = '{:,}m {}'.format(i_height, title_time)
        plot_ax.set_title(title)

        filename = path.join(datapath,'{}m-{}.png'.format(i_height, filename_time))
        plt.savefig(filename, dpi=200, format='png', bbox_inches='tight')

        plt.close('all')


if __name__ == "__main__":
    main()