argo_data_explore.py

import glob
from xarray import open_dataset
import matplotlib.pyplot as plt
import numpy as np
# from tqdm import trange
from pandas import to_datetime, isnull
from mpl_toolkits.basemap import Basemap


def argo_subset_by_szn(ncfile, months, option=1):
    # Get subsetted latitude and longitude data by time
    # Want to get data from a certain season only
    # options: 1: return lat and lon subsetted
    #          2: return depth subsetted

    ncdata = open_dataset(ncfile)

    time_pd = to_datetime(ncdata.JULD.data)
    subsetter = np.where((time_pd.month >= months[0]) & (time_pd.month <= months[-1]))
    
    if option == 1:
        lat_subset = ncdata.LATITUDE.data[subsetter]
        lon_subset = ncdata.LONGITUDE.data[subsetter]
        return lat_subset, lon_subset
    
    elif option == 2:
        # Get deepest pressure in addition to subset by month
        time_subset = time_pd[subsetter]
        pres_subset = np.nanmax(ncdata.PRES_ADJUSTED.data, axis=1)[subsetter]
        return time_subset, pres_subset


def argo_map_dist(nclist, output_folder, szn, left_lon, bot_lat,
                  right_lon, top_lat, var='oxygen', verbose=False):
    # TAKES DATA FILE PATH AS INPUT instead of arrays of data
    # Plot spatial distribution of data on a map using the Basemap package
    # See if some geographic regions are underrepresented
    # Use the Basemap package for creating maps
    
    # nclist: list of netCDF file paths to be read in with xarray.open_dataset
    # output_folder: full path of output folder for files generated by this function
    # instrument: 'BOT' for bottle, 'PCTD' for profiling CTD, 'MCTD' for
    #             moored CTD
    # left_lon, bot_lat, right_lon, top_lat: corner coordinates for the
    #                                        Basemap map
    # szn: 'Winter', 'Spring', 'Summer', 'Fall', or 'All'
    # verbose: Prints out messages during code running if True. Default False
    
    # Check the multifile option
    if len(nclist) > 1:
        multifile = True
    else:
        multifile = False
    
    if verbose:
        print('multifile', multifile)
    
    # Assign months to plot
    if szn == 'Winter':
        months = np.arange(1, 4)
    elif szn == 'Spring':
        months = np.arange(4, 7)
    elif szn == 'Summer':
        months = np.arange(7, 10)
    elif szn == 'Fall':
        months = np.arange(10, 13)
    elif szn == 'All':
        months = np.arange(1, 13)
    else:
        print('Invalid value for szn:', szn)
    
    lat_subset, lon_subset = argo_subset_by_szn(nclist[0], months, option=1)
    
    # Print accounting statistics:
    print('Printing counting statistics:')
    print('Min/max latitude bounds:', min(lat_subset), max(lat_subset))
    print('Min/max longitude bounds:', min(lon_subset), max(lon_subset))
    
    # Set up Lambert conformal map
    m = Basemap(llcrnrlon=left_lon, llcrnrlat=bot_lat,
                urcrnrlon=right_lon, urcrnrlat=top_lat, projection='lcc',
                resolution='h', lat_0=0.5 * (bot_lat + top_lat),
                lon_0=0.5 * (left_lon + right_lon))
    
    # Use NASA's "Blue Marble" image
    # m.bluemarble()
    
    # Plot lat and lon data as red markers
    
    # Initialize figure
    fig = plt.figure(num=None, figsize=(8, 6), dpi=100)
    ax = plt.subplot(1, 1, 1)  # need in order to add points to plot iteratively?
    m.drawcoastlines(linewidth=0.2)
    m.drawmapboundary(fill_color='white')
    m.fillcontinents(color='0.8')
    
    # Plot the locations of the samples
    # REPEATED IF MULTIFILE
    x, y = m(lon_subset, lat_subset)
    # Plot on the subplot ax
    m.scatter(x, y, marker='o', color='r', s=0.5)
    
    # Add scatter points to plot from other files
    if multifile:
        for i in range(1, len(nclist)):
            lat_subset, lon_subset = argo_subset_by_szn(nclist[i], months)
            # Plot the locations of the samples
            x, y = m(lon_subset, lat_subset)
            # Plot on the subplot ax
            m.scatter(x, y, marker='o', color='r', s=0.5)
    
    if verbose:
        print('Map and data plotted')
    
    plt.title('Argo {} 1991-2020 {}'.format(var, szn))
    
    png_name = output_folder + 'Argo_{}_spatial_dist_{}.png'.format(var, szn)
    plt.savefig(png_name, dpi=400)
    
    # Close the figure so it can't be written to any more
    plt.close(fig)
    
    return png_name


def argo_depth_scatter(nclist, output_folder, szn, var='oxygen', verbose=False,
                       multifile=False):
    # USED THIS FUNCTION
    # Scatter plot of maximum profile depth vs time
    
    # Assign months to plot
    if szn == 'Winter':
        months = np.arange(1, 4)
    elif szn == 'Spring':
        months = np.arange(4, 7)
    elif szn == 'Summer':
        months = np.arange(7, 10)
    elif szn == 'Fall':
        months = np.arange(10, 13)
    elif szn == 'All':
        months = np.arange(1, 13)
    else:
        print('Invalid value for szn:', szn)
    
    # Subset pressure by season
    time_subset, pres_subset = argo_subset_by_szn(nclist[0], months, option=2)
    
    if verbose:
        print('Depth subsetted by input season')
    
    # Make scatter plot
    fig = plt.figure(num=None, figsize=(8, 6), dpi=100)
    ax = plt.subplot(1, 1, 1)
    plt.scatter(time_subset, pres_subset, s=1, color='b')
    # Adjust axes to show all data in nclist
    xmin = to_datetime('2004-01-01')
    xmax = to_datetime('2020-12-31')
    ax.set_xlim(xmin, xmax)

    # Add scatter points to plot from other files
    if multifile:
        for i in range(1, len(nclist)):
            time_subset, pres_subset = argo_subset_by_szn(nclist[i], months, option=2)
            plt.scatter(time_subset, pres_subset, s=1, color='b')
    
    plt.ylabel('Pressure (dbar)')
    plt.title('Argo {} maximum profile depth: {}'.format(var, szn))
    
    # Invert the y axis (depth)
    plt.gca().invert_yaxis()
    
    png_name = output_folder + 'Argo_{}_max_depths_{}.png'.format(var, szn)
    plt.savefig(png_name, dpi=400)
    plt.close(fig)
    
    return png_name


# Examine the Argo netCDF files that Sam Huntington downloaded
fpath = '/home/hourstonh/Documents/climatology/data/oxy_clim/Coriolis_Argo/'
infiles = glob.glob(fpath + '*.nc', recursive=False)
infiles.sort()

# Output folder
dest_dir = '/home/hourstonh/Documents/climatology/data_explore/Argo/'

szns = ['Winter', 'Spring', 'Summer', 'Fall']

# Map the files
for s in szns:
    argo_map_dist(infiles, dest_dir, s, left_lon=-160,
                  right_lon=-102, bot_lat=25, top_lat=62, var='oxygen')
    
# Check if all Argo files actually contain oxygen
for f in infiles:
    sdata = open_dataset(f)
    print(sdata.JULD.data[0], sdata.JULD.data[-1])
    sdata.close()
    
# Maximum depth plots
for s in szns:
    argo_depth_scatter(infiles, dest_dir, s, multifile=True)
    
    
# Look at NODC PFL data, see if it's the same as the
# Coriolis Argo data that Sam already downloaded