Create subset of LIGHT particle initialization file to only seed specific region of ocean

subset_particle_init.py

"""
Subset Particle Init
--------------------
Author: Riley X. Brady
Date: July 9th, 2018

Given a LIGHT particle initialization file and latitude/longitude bounds,
this script will subset the global seeding and only seed a specific
subregion of the ocean.
"""
import argparse
import numpy as np
import xarray as xr
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import matplotlib.pyplot as plt


def map_subset(lon, lat):
    """
    Plots subset particles and displays.
    """
    _, ax = plt.subplots(subplot_kw=dict(projection=ccrs.PlateCarree()))
    gl = ax.gridlines(draw_labels=True, color="#a9a9a9")
    gl.xlabels_top = False
    gl.ylabels_right = False
    ax.add_feature(cfeature.LAND, facecolor='k')
    lon = lon[0::10]
    lat = lat[0::10]
    ax.scatter(lon, lat, s=0.1, transform=ccrs.PlateCarree())
    ax.set_title('Subset particle initialization')
    plt.draw()


def xyz_to_latlon(x, y, z):
    """
    Convert from xyz coordinates to lat/lon coordinates.
    """
    sq = np.sqrt(x**2 + y**2 + z**2)
    lat = np.arcsin(z / sq)
    lon = np.arctan2(y, x)
    return lon * 180./np.pi, \
        lat * 180./np.pi


def main():
    ds = xr.open_dataset(args['input'])
    ds['lon'], ds['lat'] = xyz_to_latlon(ds.xParticle, ds.yParticle,
                                         ds.zParticle)
    if args['lon'] is not None:
        ds = ds.where((ds.lon > x0) & (ds.lon < x1), drop=True)
    if args['lat'] is not None:
        ds = ds.where((ds.lat > y0) & (ds.lat < y1), drop=True)
    if args['display']:
        map_subset(ds.lon, ds.lat)
    print('saving netCDF...')
    del ds['lat']
    del ds['lon']
    ds.to_netcdf(args['output'])
    plt.show()


if __name__ == '__main__':
    parser = argparse.ArgumentParser(
        description="Subsets a LIGHT initialization file to only seed a \
                     specific section of the ocean")
    parser.add_argument('--lat', nargs=2, default=None, type=float,
                        help="Input latitude bounds (-90 to 90) separated \
                        by a space with the minimum value first")
    parser.add_argument('--lon', nargs=2, default=None, type=float,
                        help="Input longitude bounds (-180 to 180) separated \
                        by a space with the minimum value first")
    parser.add_argument('-i', '--input', type=str, required=True,
                        help="Path to input particle file.")
    parser.add_argument('-o', '--output', type=str, required=True,
                        help="Path to output edited file (including filename)")
    parser.add_argument('-d', '--display', default=False,
                        action='store_true',
                        help="If enabled, displays map of subset particles.")
    args = vars(parser.parse_args())
    # Check entries for various constraints, parse input.
    if args['lat'] is None and args['lon'] is None:
        raise ValueError("""
                          Please input either a lon or lat subset. Otherwise
                          you do not need to be subsetting the init file.
                         """)
    if not args['input'].endswith('.nc'):
        raise ValueError('Please input a netCDF file (*.nc)')
    if not args['output'].endswith('.nc'):
        raise ValueError('Please output a netCDF file (*.nc)')
    if args['lat'] is not None:
        y0, y1 = args['lat']
    if args['lon'] is not None:
        x0, x1 = args['lon']
    main()