build_domain.py

#!/usr/bin/env python
"""
Select out the necessary locations from a larger domain file
and output the pared down file for use in MetSim.  This script
is intended to be used when using ascii or binary input.  The 
usage can be found by running `./build_domain.py -h` or simply
by inspecting the USAGE command.

This script requires the input domain file to have the 'elev' 
field filled out, and the forcing files to be a subset of
the input domain file.
"""

import os
import sys

import numpy as np
import xarray as xr
from netCDF4 import Dataset

USAGE = """
Insufficient data provided! Execution format:
    ./build_domain.py DOMAIN_FILE FORCING_DIR OUT_FILE [reverse]

If the reverse argument is given, then file formats are 
expected as FILE_LON_LAT, otherwise they should be FILE_LAT_LON
"""

def main(in_domain: str, out_file: str, in_forcings: list, reverse: bool=False):
    # Record some information about the input locations
    coords = [(l.split('_')[1], l.split('_')[-1]) for l in in_forcings]
    lats = np.unique([float(l[0]) for l in coords]) 
    lons = np.unique([float(l[1]) for l in coords])
    dimnames = ('lat', 'lon')
   
    # Get the input dataset information
    old_domain = xr.open_dataset(in_domain) 
    old_lats = np.array(old_domain.lat.values)
    old_lons = np.array(old_domain.lon.values)
    if reverse:
        lats, lons = lons, lats
        dimnames = dimnames[::-1]
 
    # Build our base dataset
    new_domain = Dataset(out_file, 'w')
    new_domain.createDimension('lat', len(lats))
    new_domain.createDimension('lon', len(lons))
    new_domain.createVariable('elev', 'd', dimnames)
    new_domain.createVariable('mask', 'd', dimnames)
    new_domain.createVariable('latitude', 'd', ('lat',))
    new_domain.createVariable('longitude', 'd', ('lon',))
    
    # Add in some data about the latitudes and longitudes
    new_domain.variables['latitude'][:] = lats
    new_domain.variables['longitude'][:] = lons

    def find_elev(lat, lon):
        """Find elevation using nearest-neighbor"""
        i = int(np.abs(old_lats - lat).argmin())
        j = int(np.abs(old_lons - lon).argmin())
        return old_domain.elev.isel(lat=i, lon=j)

    # Fill in the information
    for i, lat in enumerate(lats):
        for j, lon in enumerate(lons):
            elev = find_elev(lat, lon).values
            mask = elev > 0
            new_domain.variables['elev'][j, i] = elev
            new_domain.variables['mask'][j, i] = mask

    # Close out the dataset
    new_domain.close()


if __name__ == "__main__":
    if (len(sys.argv) < 4 or len(sys.argv) > 5
            or not os.path.isfile(sys.argv[1]) 
            or not os.path.isdir(sys.argv[2])
            or os.path.isfile(sys.argv[3])):
        exit(USAGE)
    reverse = len(sys.argv) == 5
    in_domain_path = sys.argv[1]
    in_forcing_path = sys.argv[2]
    out_file_path = sys.argv[3]
    in_forcings = os.listdir(in_forcing_path)
    main(in_domain_path, out_file_path, in_forcings, reverse)