generate_grid.py

import numpy as np
import pyproj
import xarray as xr

# Generate LCC Grid
def generate_lcc_grid(nx = 532, ny = 532, dx = 3000., dy = 3000.,
                      central_latitude = 39.05, central_longitude = -95.68,
                      standard_parallels = 39.05,
                      globe_kwargs = {'a': 6370000, 'b': 6370000}):
    # Inputs default to Squiteri and Gallus (2016) WRF grid

    # Make sure we have standard_parallels in proper format
    try:
        assert len(standard_parallels) == 2
    except TypeError:
        standard_parallels = [standard_parallels, standard_parallels]
    except AssertionError:
        standard_parallels = [standard_parallels[0], standard_parallels[-1]]

    # Set up our projections
    lcc_proj = pyproj.Proj(proj='lcc',
                           lat_1=standard_parallels[0], lat_2=standard_parallels[1],
                           lat_0=central_latitude, lon_0=central_longitude,
                           **globe_kwargs)
    wgs_proj = pyproj.Proj(proj='latlong', datum='WGS84')

    # Get center easting and northing
    e, n = pyproj.transform(wgs_proj, lcc_proj, central_longitude, central_latitude)

    # Set lower-left corner of the domain
    x0 = -(nx - 1) / 2. * dx + e
    y0 = -(ny - 1) / 2. * dy + n

    # Create the 2D LCC Grid
    x = np.arange(nx) * dx + x0
    y = np.arange(ny) * dy + y0
    xx, yy = np.meshgrid(x, y)

    # Transform to LatLon
    lons, lats = pyproj.transform(lcc_proj, wgs_proj, xx, yy)
    lons += 360.

    # Create dataset
    lcc_ds = xr.Dataset(coords={'lon': (['y', 'x'], lons),
                                'lat': (['y', 'x'], lats),
                                'y': y,
                                'x': x})
    lcc_ds['lon'].attrs = {'standard_name': 'longitude',
                         'units': 'degrees_east'}
    lcc_ds['lat'].attrs = {'standard_name': 'latitude',
                         'units': 'degrees_north'}
    lcc_ds['y'].attrs = {'standard_name': 'projection_x_coordinate',
                         'units': 'meter'}
    lcc_ds['x'].attrs = {'standard_name': 'projection_y_coordinate',
                         'units': 'meter'}
    
    return lcc_ds