plot_synthpro_output.py

#!/usr/bin/env python2.7

"""
Plotting script for the evaluation of profiles extracted by SynthPro. 

This script is intended to work with synthetic profiles in the EN4 format 
that have been interpolated to observed depths. It may behave unexpectedly
if other data types are plotted.

"""

import argparse

import matplotlib
matplotlib.use('AGG')
import matplotlib.pyplot as plt
from mpl_toolkits.basemap import Basemap
    
from synthpro.synthpro import *


def get_args():
    """
    Get arguments from command line.
    
    """
    parser = argparse.ArgumentParser(
        description='Plot comparisons of observed and synthetic profiles generated by SynthPro')
    parser.add_argument(
        'month', type=int, help='Month used in file names.')
    parser.add_argument(
        'year', type=int, help='Year used in file names.')
    parser.add_argument(
        'namelist', type=str, help='Path to namelist.ini')
    parser.add_argument(
        '-d', '--day', type=int, help='Day used in file names.', default=None)
    parser.add_argument(
        '-o', '--outdir', type=str, help='Directory to save plots.', default='./')
        
    args = parser.parse_args()

    return args


def latlon(coord,lon=False):
    """ Returns lat or lon as a string with the suffix S, N, E or W."""
  
    if lon:
        if coord < 0:
            string =  '%iW' % (-coord)
        else:
            string =  '%iE' % (coord)
    else:
        if coord < 0:
            string =  '%iS' % (-coord) 
        else:
            string =  '%iN' % (coord)
            
    return string


def lbound(val, minval):
    """ Apply a lower bound """
    if val < minval:
        return minval
    else:
        return val


def ubound(val, maxval):
    """ Apply an upper bound """
    if val > maxval:
        return maxval
    else:
        return val


def plot_profile_locations(args, config, obsDat, synthDat, modelDat,
                           figsize=(10,6), latlonpad=0.02, proj='cyl',
                           plot_global=True):
    """ Plot location of observed and extracted locations of profiles. """    

    # Get i, j indices
    synthDat.load_i()
    synthDat.load_j()
    imin = config.getint('model_temp', 'imin')
    jmin = config.getint('model_temp', 'jmin')
    nanind = np.isnan(synthDat.i) == False
    ii = np.array(synthDat.i[nanind],dtype=np.int) - imin
    jj = np.array(synthDat.j[nanind],dtype=np.int) - jmin

    # Get model lat, lon values
    mlats = modelDat.lats[jj, ii]
    mlons = modelDat.lons[jj, ii]

    # Get profile lat, lon values
    slons, slats = synthDat.lons[nanind], synthDat.lats[nanind]

    # Set basemap
    fig = plt.figure(figsize=figsize)

    if plot_global:
        llcrnrlon = -180
        llcrnrlat = -90
        urcrnrlon = 180
        urcrnrlat = 90
    else:
        llcrnrlon = lbound(mlons.min() - latlonpad * (mlons.max() - mlons.min()), -180)
        llcrnrlat = lbound(mlats.min() - latlonpad * (mlats.max() - mlats.min()), -90)
        urcrnrlon = ubound(mlons.max() + latlonpad * (mlons.max() - mlons.min()), 360)
        urcrnrlat = ubound(mlats.max() + latlonpad * (mlats.max() - mlats.min()), 90)

    m = Basemap(llcrnrlon=llcrnrlon, llcrnrlat=llcrnrlat,
                urcrnrlon=urcrnrlon, urcrnrlat=urcrnrlat,
                projection=proj)
    m.fillcontinents(color='darkseagreen')
    
    # Project locations
    mx, my = m(mlons, mlats)
    sx, sy = m(slons, slats)
    
    # Add to basemap
    plt.plot(mx, my, '.r', markeredgecolor='none', markersize=6, label='Synthetic profile locations')
    plt.plot(sx, sy, '.k', markeredgecolor='none', markersize=3, label='Observed profile locations')

    # Add notation
    parallels = np.arange(-80.,81,20.)
    m.drawparallels(parallels,labels=[False,True,True,False])
    meridians = np.arange(0.,351.,30.)
    m.drawmeridians(meridians,labels=[True,False,False,True])
    plt.legend(loc=2, numpoints=1, fontsize=8)
    plt.title('Comparison of observed and synthetic profile locations - %4i/%02i' % (args.year, args.month))
    
    # Save figure
    savef = args.outdir + config.get('synth_profiles', 'file_name').split('/')[-1].replace('.nc', '.location_map.png')
    printmsg.message(config,'SAVING: ' + savef)
    fig.savefig(savef, resolution=100)
    
    
def plot_dist2ob_hist(args, config, synthDat, figsize=(10,6)):
    """
    Plot histogram of distance between observed and synthetic profiles
    
    """
    # Load distance data
    dist2ob = synthDat.read_var('distance_to_ob')

    # Remove masked data
    try:
        dist2ob = dist2ob[dist2ob.mask == False]
    except AttributeError:
        pass
        
    # Set figure
    fig = plt.figure(figsize=figsize)
    n, bins, patches = plt.hist(dist2ob/1000., normed=True)
    plt.xlabel('Distance (km)')
    plt.ylabel('Normalized frequency')
    plt.title('Histogram of distances between observed and synthetic profiles - %4i/%02i' % (args.year, args.month))

    # Save figure
    savef = args.outdir + config.get('synth_profiles', 'file_name').split('/')[-1].replace('.nc', '.distances_histogram.png')
    printmsg.message(config,'SAVING: ' + savef)
    fig.savefig(savef, resolution=100)
    
    
def get_qc_ind(proDat, var, reject='4'):
    """ Return index for profile data passing quality control """
    try:
        dat_qc = proDat.read_var(var + '_QC')
        pos_qc = proDat.read_var('POSITION_QC')
        pos_qc = pos_qc != reject
        dat_qc = return_1d_qc(dat_qc, reject=reject)
        qc_ind = np.where(pos_qc & dat_qc)[0]
    except:
        printmsg.message(config, 'Could not find data for %s. Plotting all data.' % var)
        qc_ind = np.arange(len(proDat.lats))
     
    return qc_ind


def return_1d_qc(dat_qc, reject='4'):
    """ For each profile, return True (accept) or False (reject). """
    nqs = np.arange(dat_qc.shape[0])
    qc_out = []
    
    for nq in nqs:
        if any(dat_qc[nq] == reject):
            qc_out.append(False)
        else:
            qc_out.append(True)
    
    return np.array(qc_out)


def plot_obs_vs_synth(args, config, obsDat, synthDat, figsize=(10,6)):
    """ Plot observed vs synthetic profile data  """ 
    
    # Load var names
    tvar = config.get('obs_profiles', 'temp_var')
    svar = config.get('obs_profiles', 'sal_var')

    for var in [tvar, svar]:
        
        #Load data
        odat = obsDat.read_var(var)[get_qc_ind(obsDat, var)]
        sdat = synthDat.read_var(var)[get_qc_ind(synthDat, var)]
        
        # Plot data
        fig = plt.figure(figsize=figsize)
        plt.plot(odat, sdat, '.r')        
        
        # Add annotations
        plt.xlabel('Observations')
        plt.ylabel('Synthetic data')
        plt.title('Observed vs synthetic %s profiles - %4i/%02i' % (var, args.year, args.month))
        
        # Save figure
        savef = args.outdir + config.get('synth_profiles', 'file_name').split('/')[-1].replace('.nc', '.%s_obs_vs_synth.png' % var)
        printmsg.message(config,'SAVING: ' + savef)
        fig.savefig(savef, resolution=100)
            

def plot_example_profiles(args, config, obsDat, synthDat, modelTemp, 
                          modelSal, figsize=(9,12)):
    """ Plot observed, synthetic and model profiles for 12 random profiles """
    
    # Load config data
    nx = 4
    ny = 3
    N = nx * ny 
    
    tvar = config.get('obs_profiles', 'temp_var')
    svar = config.get('obs_profiles', 'sal_var')
    imin = config.getint('model_temp', 'imin')
    jmin = config.getint('model_temp', 'jmin')

    # Load profile depths    
    obsz = obsDat.depths/1000.
    synz = synthDat.depths/1000.
    mdlz = modelTemp.depths/1000.
        
    # Load model i,j indices
    synthDat.load_i()
    synthDat.load_j()   
    syninds = np.where(np.isnan(synthDat.i) == False)[0]

    # Load and plot T/S data
    for var, modelDat in zip([tvar, svar], [modelTemp, modelSal]):
        
        # Load data
        obsdat = obsDat.read_var(var)
        syndat = synthDat.read_var(var)
        
        # Find indices of good profiles
        obsinds = get_qc_ind(obsDat,var) 
        inds = np.intersect1d(obsinds, syninds)
        
        # Set figure
        fig = plt.figure(figsize=figsize)
        matplotlib.rcParams['font.size'] = 8
        fig.suptitle('Comparison of %i randomly selected %s profiles - %4i/%02i'
                     % (N, var, args.year, args.month))
               
        # Select and plot profiles
        if len(inds) > 0:
            for nplt in range(N):
                # Select index
                ind = np.random.choice(inds)
                i = np.int(synthDat.i[ind]) - imin
                j = np.int(synthDat.j[ind]) - jmin
                lon, lat = latlon(obsDat.lons[ind], lon=True), latlon(obsDat.lats[ind])
        
                # Plot data
                fig.add_subplot(ny, nx, nplt+1)
                plt.plot(modelDat.data[:,j,i], -mdlz, 'r', label='Model')
                plt.plot(obsdat[ind,:], -obsz[ind,:], 'kx-', label='Obs')
                plt.plot(syndat[ind,:], -synz[ind,:], 'o', mfc='none', mec='r', label='Synth')
                plt.ylim([-obsz[ind,:].max(), 0])
                plt.title('%s, %s (i=%i, j=%i)' % (lon, lat, i, j))
           
                if nplt == 0:  
                    plt.legend(loc=4, fontsize=10, numpoints=1)      

            # Add annotation
            fig.tight_layout()
            fig.subplots_adjust(bottom=0.05, top=0.95, left=0.1)         
            plt.text(0.5, 0.025, var, ha='center', va='center', transform=fig.transFigure)
            plt.text(0.025, 0.5, 'Depth (km)', ha='center', va='center', transform=fig.transFigure, rotation='vertical')
        
            # Save figure
            savef = args.outdir + config.get('synth_profiles', 'file_name').split('/')[-1].replace('.nc', '.%s_profiles.png' % var)
            printmsg.message(config,'SAVING: ' + savef)
            fig.savefig(savef, resolution=200)


if __name__ == '__main__':
    
    # Build paths to input data files
    args = get_args()   
    config = namelist.get_namelist(args)    
    printmsg.message(config, 'Loading input data')
    config = tools.build_file_name(args, config, 'obs_profiles')
    config = tools.build_file_name(args, config, 'synth_profiles')
    config = tools.build_file_name(args, config, 'model_temp')
    config = tools.build_file_name(args, config, 'model_sal') 
    printmsg.inputs(config) 
    
    # Load data objects     
    printmsg.loading(config)
    obsDat = profiles.assoc_profiles(config, 'obs_profiles')
    synthDat = profiles.assoc_profiles(config, 'synth_profiles', read_only=True)
    modelTemp = model.assoc_model(config, 'model_temp')
    modelSal = model.assoc_model(config, 'model_sal')
    
    # Create plots
    printmsg.message(config, 'Generating comparison plots...')
    plot_profile_locations(args, config, obsDat, synthDat, modelTemp)
    plot_dist2ob_hist(args, config, synthDat)
    plot_obs_vs_synth(args, config, obsDat, synthDat)
    plot_example_profiles(args, config, obsDat, synthDat, modelTemp, modelSal)