Manually coarsen horizontal seeding pattern in LIGHT

manual_coarsen.py

# Manual Coarsen
# --------------
# Author : Riley X. Brady
# Date   : 11/14/2018
# --------------
# The pyAMG method of coarsening our particles traverses in 3D, rather than
# just removing full 2D columns. This method simply strides through the
# lat and lon list, removing full columns.
#
# NOTE: Run make_particle_file.py first without any coarsening turned on.
# This will take in that particles.nc file and coarsen it.
import numpy as np
import argparse
import netCDF4


def _xyz_to_lat_lon(x, y, z):
    x, y, z = np.asarray(x), np.asarray(y), np.asarray(z)
    plat = np.arcsin(z / np.sqrt(x**2 + y**2 + z**2))
    plong = np.arctan2(y, x)
    # longitude adjustment
    plong[plong < 0.0] = 2*np.pi + plong[plong < 0.0]
    return plong * 180./np.pi, \
           plat * 180./np.pi


def _get_dtype(f_in, var):
    # Get datatype
    if f_in[var].datatype is np.dtype('float64'):
        dtype = 'd'
    else:
        dtype = 'i'
    # Return tuple dimensions
    return dtype


def _stride_variable(f_in, var, particleList):
    """
    Coarsen all variables that come through.
    """
    single_dim = ['indexToParticleID', 'resetTime', 'currentBlockReset',
              'currentCellReset', 'xParticleReset', 'yParticleReset',
              'zParticleReset', 'zLevelParticleReset']
    if var in single_dim:
        temp_var = f_in[var][particleList]
    else:
        temp_var = f_in[var][:, particleList]
    return temp_var


def main(init, stride, output):
    f_in = netCDF4.Dataset(init, 'r')

    # get cell positions
    x = f_in.variables['xParticle'][:]
    y = f_in.variables['yParticle'][:]
    z = f_in.variables['zParticle'][:]

    # get their lat/lon locations
    plong, plat = _xyz_to_lat_lon(x,y,z)

    # find unique values (i.e. locations of columns of particles)
    un_lon, un_lat = np.unique(plong), np.unique(plat)

    # stride through values
    lonlist, latlist = un_lon[0::stride], un_lat[0::stride]
    particleList = []
    for i, (lon, lat) in enumerate(zip(plong.squeeze(), plat.squeeze())):
        if (lon in lonlist) and (lat in latlist):
            particleList.append(i)

    # length of new nParticles dimension
    N = len(particleList)

    # CREATE NEW FILE
    nf = netCDF4.Dataset(output, 'w', format='NETCDF3_64BIT_OFFSET')

    # Set up dimensions
    nf.createDimension('Time', None)
    nf.createDimension('nParticles', N)

    for var in f_in.variables.keys():
        # Create variable in new netCDF file
        print(var + '...')
        temp_var = nf.createVariable(var, _get_dtype(f_in, var),
                                     f_in[var].dimensions)
        # Fill with appropriate values
        temp_var[:] = _stride_variable(f_in, var, particleList)
    # Finish out
    nf.close()
    f_in.close()


if __name__ == '__main__':
    ap = argparse.ArgumentParser(
            description="Coarsens particle init file horizontally by a determined stride.")
    ap.add_argument('-i', '--init', required=True, type=str,
            help="Particle init file (netCDF)")
    ap.add_argument('-s', '--stride', required=True, type=int,
            help="Number of particles to skip in lat and lon")
    ap.add_argument('-o', '--output', required=True, type=str,
            help="Coarsened particle output file (.nc at end)")
    args = vars(ap.parse_args())
    init = args['init']
    stride = args['stride']
    output = args['output']
    main(init, stride, output)