Sounding Plotter

import posixpath

import numpy as np
import matplotlib.pyplot as plt
import metpy.calc as mpcalc
from metpy.plots import add_metpy_logo, add_timestamp, SkewT
from metpy.units import units

from siphon.simplewebservice.wyoming import WyomingUpperAir

class DriveUploader:
    def __init__(self, credsfile='mycreds.txt'):
        from pydrive.drive import GoogleDrive
        self.gdrive = GoogleDrive(self._get_auth(credsfile))

    def _get_auth(self, credsfile):
        from pydrive.auth import GoogleAuth
        gauth = GoogleAuth()

        # Try to load saved client credentials
        gauth.LoadCredentialsFile(credsfile)
        if gauth.credentials is None:
            # Authenticate if they're not there
            gauth.LocalWebserverAuth()
        elif gauth.access_token_expired:
            # Refresh them if expired
            gauth.Refresh()
        else:
            # Initialize the saved creds
            gauth.Authorize()
        # Save the current credentials to a file
        gauth.SaveCredentialsFile(credsfile)
        return gauth

    def _get_first_file_id(self, title, parent, **kwargs):
        query = "title='{}' and '{}' in parents".format(title, parent)
        for k, v in kwargs.items():
            query += " and {}='{}'".format(k, v)
        res = next(self.gdrive.ListFile({'q': query}))
        if res:
            return res[0]['id']
        return None

    def get_folder(self, path):
        parent = 'root'
        for part in path.split('/'):
            if not part:
                continue
            parent = self._get_first_file_id(part, parent,
                                             mimeType='application/vnd.google-apps.folder')
        return parent

    def create_or_get_file(self, path):
        pathname, filename = posixpath.split(path)
        folder = self.get_folder(pathname)
        create_file_args = {'parents': [{'kind': 'drive#fileLink', 'id': folder}]}

        file_id = self._get_first_file_id(filename, folder)
        if file_id is not None:
            create_file_args['id'] = file_id
        return self.gdrive.CreateFile(create_file_args)

    def upload_to(self, local_path, remote_path):
        f = self.create_or_get_file(remote_path)
        f.SetContentFile(local_path)
        f['title'] = posixpath.basename(remote_path)
        f.Upload()

def plot_skewt(df):
    # We will pull the data out of the example dataset into individual variables
    # and assign units.
    pres = df['pressure'].values * units.hPa
    temp = df['temperature'].values * units.degC
    dewpoint = df['dewpoint'].values * units.degC
    wind_speed = df['speed'].values * units.knots
    wind_dir = df['direction'].values * units.degrees
    u, v = mpcalc.wind_components(wind_speed, wind_dir)

    # Create a new figure. The dimensions here give a good aspect ratio.
    fig = plt.figure(figsize=(9, 9))
    skew = SkewT(fig, rotation=45)

    # Plot temperature, dewpoint and wind barbs
    skew.plot(pres, temp, 'red')
    skew.plot(pres, dewpoint, 'green')
    
    # Plot wind barbs
    my_interval = np.arange(100, 1000, 50) * units('hPa') #set spacing interval
    ix = mpcalc.resample_nn_1d(pres, my_interval) #find nearest indices for chosen interval
    skew.plot_barbs(pres[ix], u[ix], v[ix], xloc=1) #plot values closest to chosen interval
    
    # Improve labels and set axis limits
    skew.ax.set_xlabel('Temperature (\N{DEGREE CELSIUS})')
    skew.ax.set_ylabel('Pressure (hPa)')
    skew.ax.set_ylim(1000, 100)
    skew.ax.set_xlim(-40, 59)

    # Calculate LCL height and plot as black dot.
    lcl_pressure, lcl_temperature = mpcalc.lcl(pres[0], temp[0], dewpoint[0]) #index 0 is chosen to lift parcel from the surface
    skew.plot(lcl_pressure, lcl_temperature, 'ko', markerfacecolor='black')

    # Calculate full parcel profile and add to plot as black line
    prof = mpcalc.parcel_profile(pres, temp[0], dewpoint[0]).to('degC')
    skew.plot(pres, prof, 'black', linewidth=2)

    # Shade areas of CAPE and CIN
    skew.shade_cin(pres, temp, prof, dewpoint)
    skew.shade_cape(pres, temp, prof)

    # Add emphasis to 0 degree isotherm with color change
    skew.ax.axvline(0, color='c', linestyle='--', linewidth=2)

    # Add the relevant special lines throughout the figure
    skew.plot_dry_adiabats(t0=np.arange(233, 533, 15) * units.K, alpha=0.25, color='orangered')
    skew.plot_moist_adiabats(t0=np.arange(233, 400, 10) * units.K, alpha=0.25, color='tab:green')
    skew.plot_mixing_lines(pressure=np.arange(1000, 99, -25) * units.hPa, linestyle='dotted', color='tab:blue')

    # Add the MetPy logo!
    fig = plt.gcf()
    add_metpy_logo(fig, 115, 100, size='small');
    
    return skew


def make_name(site, time):
    return '{site}_{dt:%Y%m%d_%H%M}.png'.format(site=site, dt=time)

if __name__ == '__main__':
    import argparse
    from datetime import datetime, timedelta, UTC
    import tempfile

    # Set up argument parsing for the script. Provides one argument for the site, and another
    # that controls whether the plot should be shown or saved as an image.
    parser = argparse.ArgumentParser(description='Download sounding data and plot.')
    parser.add_argument('-s', '--site', help='Site to obtain data for', type=str,
                        default='DDC')
    parser.add_argument('--show', help='Whether to show the plot rather than save to disk',
                        action='store_true')
    parser.add_argument('-d', '--date', help='Date and time to request data for in YYYYMMDDHH.'
                        ' Defaults to most recent 00/12 hour.', type=str)
    parser.add_argument('-g', '--gdrive', help='Google Drive upload path', type=str)
    parser.add_argument('-f', '--filename', help='Image filename', type=str)
    args = parser.parse_args()

    if args.date:
        request_time = datetime.strptime(args.date, '%Y%m%d%H')
    else:
        # Figure out the most recent sounding, 00 or 12. Subtracting two hours
        # helps ensure that we choose a time with data available.
        now = datetime.now(UTC) - timedelta(hours=2)
        request_time = now.replace(hour=(now.hour // 12) * 12, minute=0, second=0)

    # Request the data and plot
    df = WyomingUpperAir.request_data(request_time, args.site)
    skewt = plot_skewt(df)

    # Add the timestamp for the data to the plot
    add_timestamp(skewt.ax, request_time, y=1.02, x=0, ha='left', fontsize='large')
    skewt.ax.set_title(args.site)

    if args.show:
        plt.show()
    else:
        fname = args.filename if args.filename else make_name(args.site, request_time)
        if args.gdrive:
            uploader = DriveUploader()
            with tempfile.NamedTemporaryFile(suffix='.png') as f:
                skewt.ax.figure.savefig(f.name)
                uploader.upload_to(f.name, posixpath.join(args.gdrive, fname))
        else:
            skewt.ax.figure.savefig(make_name(args.site, request_time))