WRF Post-Process to CF Example

wrf_run_post_cf_loop.py

#!/usr/bin/env python

from itertools import product
from multiprocessing import Pool, Process

import numpy as np
import xarray as xr
import metpy.calc as mpcalc

from datetime import datetime
import cartopy.crs as ccrs
import cartopy.feature as cfeature
import pyproj
from metpy.plots.mapping import CFProjection

def generate_cf_projection_from_wrf(wrf_dataset):
    """
    Create a MetPy CFProjection object from the attributes of a WRF dataset (as an
    xarray Dataset).
    
    Right now only works with LCC
    """
    cf_attrs = {
        'grid_mapping_name': 'lambert_conformal_conic',
        'earth_radius': 6370000,
        'standard_parallel': [wrf_dataset.TRUELAT1, wrf_dataset.TRUELAT2],
        'longitude_of_central_meridian': wrf_dataset.STAND_LON,
        'latitude_of_projection_origin': wrf_dataset.MOAD_CEN_LAT
    }
    return CFProjection(cf_attrs)

def generate_dimension_coordinates_from_wrf(wrf_dataset):
    """
    Create coordinates corresponding to the bottom_top, south_north, and west_east
    dimensions in a WRF dataset (as an xarray Dataset).
    
    Must already have a MetPy crs coordinate added. Cannot be a moving domain.
    """
    # Obtain CRSs
    data_crs = wrf_dataset['metpy_crs'].metpy.cartopy_crs
    latlon_crs = ccrs.Geodetic(globe=ccrs.Globe(ellipse='WGS84'))
    
    # Obtain grid parameters
    dx, dy = wrf_dataset.DX, wrf_dataset.DY
    nx, ny = (wrf_dataset.dims[i] for i in ['west_east', 'south_north'])
    
    # Compute easting and northing of domain center
    e, n = data_crs.transform_point(wrf_dataset.CEN_LON, wrf_dataset.CEN_LAT, latlon_crs)
    
    # Find easting and northing of lower-left corner of domain
    x0 = -(nx - 1) / 2. * dx + e
    y0 = -(ny - 1) / 2. * dy + n
    
    # Build horizontal grid values
    x, y = np.arange(nx) * dx + x0, np.arange(ny) * dy + y0
    
    # Build DataArrays
    eta_attrs = wrf_dataset['ZNU'].attrs
    eta_attrs['axis'] = 'Z'
    eta_attrs['standard_name'] = 'atmosphere_hybrid_sigma_pressure_coordinate'
    bottom_top = xr.DataArray(wrf_dataset['ZNU'][0], dims='bottom_top', attrs=eta_attrs)
    south_north = xr.DataArray(y, dims='south_north', attrs={'axis': 'Y', 'units': 'm',
                                                             'standard_name': 'projection_y_coordinate'})
    west_east = xr.DataArray(x, dims='west_east', attrs={'axis': 'X', 'units': 'm',
                                                         'standard_name': 'projection_x_coordinate'})
    
    return bottom_top, south_north, west_east

def process_run(run_date, config, domain):
    """Function to process an individual run."""
    print('\n------------------\n')
    print(run_date, config, domain)
    print('\n------------------\n')
    ds = xr.open_dataset(run_date.strftime('/glade/scratch/jthielen/'
                                         '{config}/%Y%m%d/wrfout_{domain}_%Y-%m-%d_12:00:00'.format(config=config, domain=domain)))
    ds = ds.assign_coords(metpy_crs=generate_cf_projection_from_wrf(ds))

    ds['bottom_top'], ds['south_north'], ds['west_east'] = generate_dimension_coordinates_from_wrf(ds)
    ds = ds.swap_dims({'Time': 'XTIME'}).rename({'XTIME': 'time'})
    ds = ds.assign_coords(latitude=ds['XLAT'][0].drop(['time', 'XLAT', 'XLONG']),
                          longitude=ds['XLONG'][0].drop(['time', 'XLAT', 'XLONG'])).drop(
        ['XLAT', 'XLONG', 'XLAT_U', 'XLAT_V', 'XLONG_U', 'XLONG_V'])

    eta_attrs = ds['ZNW'].attrs
    eta_attrs['axis'] = 'Z'
    eta_attrs['standard_name'] = 'atmosphere_hybrid_sigma_pressure_coordinate'
    ds['bottom_top_stag'] = xr.DataArray(ds['ZNW'][0], dims='bottom_top_stag', attrs=eta_attrs)

    print(ds.coords)

    ds_output = xr.Dataset(coords=ds.coords)

    print('land_use_category')
    ds_output['land_use_category'] = ds['LU_INDEX'][0].drop('time')

    print('u')
    ds_output['u_at_10m'] = ds['U10']
    ds_output['u_at_10m'].attrs = {'units': 'm/s', 'standard_name': 'x_wind'}
    ds_output['u'] = ((ds['U'][..., :-1] + ds['U'][..., 1:]) / 2).rename({'west_east_stag': 'west_east'}) # Destagger
    ds_output['u'].attrs = {'units': 'm/s', 'standard_name': 'x_wind'}

    print('v')
    ds_output['v_at_10m'] = ds['V10']
    ds_output['v_at_10m'].attrs = {'units': 'm/s', 'standard_name': 'y_wind'}
    ds_output['v'] = ((ds['V'][..., :-1, :] + ds['V'][..., 1:, :]) / 2).rename({'south_north_stag': 'south_north'}) # Destagger
    ds_output['v'].attrs = {'units': 'm/s', 'standard_name': 'y_wind'}

    print('w')
    ds_output['w'] = ds['W']
    ds_output['w'].attrs = {'units': 'm/s', 'standard_name': 'upward_air_velocity'}

    print('geopotential')
    ds_output['geopotential'] = ds['PHB'] + ds['PH']
    ds_output['geopotential'].attrs = {'units': 'meter^2 / second^2', 'standard_name': 'geopotential'}

    print('potential_temperature')
    ds_output['potential_temperature_at_2m'] = ds['TH2']
    ds_output['potential_temperature_at_2m'].attrs = {'units': 'K', 'standard_name': 'air_potential_temperature'}
    ds_output['potential_temperature'] = ds['T'] + 300
    ds_output['potential_temperature'].attrs = {'units': 'K', 'standard_name': 'air_potential_temperature'}

    print('temperature')
    ds_output['temperature_at_2m'] = ds['T2']
    ds_output['temperature_at_2m'].attrs = {'units': 'K', 'standard_name': 'air_temperature'}

    print('pressure')
    ds_output['pressure_at_surface'] = ds['PSFC']
    ds_output['pressure_at_surface'].attrs = {'units': 'Pa', 'standard_name': 'air_pressure'}
    ds_output['pressure'] = ds['P'] + ds['PB']
    ds_output['pressure'].attrs = {'units': 'Pa', 'standard_name': 'air_pressure'}

    print('turbulence_kinetic_energy')
    ds_output['turbulence_kinetic_energy'] = ds['QKE'] / 2
    ds_output['turbulence_kinetic_energy'].attrs = {'units': 'm**2 / s**2', 'long_name': 'turbulent_kinetic_energy_from_mynn'}

    print('length_scale')
    ds_output['length_scale'] = ds['EL_PBL']
    ds_output['length_scale'].attrs = {'units': 'm', 'long_name': 'length_scale_from_pbl'}

    print('terrain_height')
    ds_output['terrain_height'] = ds['HGT'][0].drop('time')
    ds_output['terrain_height'].attrs = {'units': 'm', 'standard_name': 'surface_altitude'}

    print('friction_velocity')
    ds_output['friction_velocity'] = ds['UST']
    ds_output['friction_velocity'].attrs = {'units': 'm/s', 'long_name': 'u*_in_similarity_theory'}

    print('pbl_height')
    ds_output['pbl_height'] = ds['PBLH']
    ds_output['pbl_height'].attrs = {'units': 'm', 'standard_name': 'atmosphere_boundary_layer_thickness'}

    print('mixing_ratio')
    ds_output['mixing_ratio_at_2m'] = ds['Q2']
    ds_output['mixing_ratio_at_2m'].attrs = {'units': '', 'standard_name': 'humidity_mixing_ratio'}
    ds_output['mixing_ratio'] = ds['QVAPOR']
    ds_output['mixing_ratio'].attrs = {'units': '', 'standard_name': 'humidity_mixing_ratio'}

    print('vegetation_category')
    ds_output['vegetation_category'] = ds['IVGTYP'][0].drop('time')
    ds_output['vegetation_category'].attrs = {'long_name': 'dominant_vegetation_category'}

    print('soil_category')
    ds_output['soil_category'] = ds['ISLTYP'][0].drop('time')
    ds_output['soil_category'].attrs = {'long_name': 'dominant_soil_category'}

    print('vegetation_fraction')
    ds_output['vegetation_fraction'] = ds['VEGFRA']
    ds_output['vegetation_fraction'].attrs = {'standard_name': 'vegetation_area_fraction'}

    print('leaf_area')
    ds_output['leaf_area'] = ds['LAI']
    ds_output['leaf_area'].attrs = {'standard_name': 'leaf_area_index'}

    print('precipitation_accumulation')
    ds_output['precipitation_accumulation'] = ds['RAINNC']
    ds_output['precipitation_accumulation'].attrs = {'units': 'mm', 'standard_name': 'integral_of_lwe_precipitation_rate_wrt_time'}

    print('cloud_fraction')
    ds_output['cloud_fraction'] = ds['CLDFRA']
    ds_output['cloud_fraction'].attrs = {'units': '', 'standard_name': 'cloud_area_fraction_in_atmosphere_layer'}

    print('surface_soil_temperature')
    ds_output['surface_soil_temperature'] = ds['TMN']
    ds_output['surface_soil_temperature'].attrs = {'units': 'K', 'long_name': 'surface_soil_temperature'}

    print('surface_heat_flux')
    ds_output['surface_heat_flux'] = ds['HFX']
    ds_output['surface_heat_flux'].attrs = {'units': 'W m**-2', 'standard_name': 'surface_upward_heat_flux_in_air'}

    print('surface_moisture_flux')
    ds_output['surface_moisture_flux'] = ds['QFX']
    ds_output['surface_moisture_flux'].attrs = {'units': 'W m**-2', 'standard_name': 'surface_upward_heat_flux_in_air'}

    print('surface_latent_heat_flux')
    ds_output['surface_latent_heat_flux'] = ds['LH']
    ds_output['surface_latent_heat_flux'].attrs = {'units': 'W m**-2', 'standard_name': 'surface_upward_latent_heat_flux_in_air'}

    print('accumulated_surface_heat_flux')
    ds_output['accumulated_surface_heat_flux'] = ds['ACHFX']
    ds_output['accumulated_surface_heat_flux'].attrs = {'units': 'J m**-2', 'standard_name': 'integral_of_surface_upward_heat_flux_in_air_wrt_time'}

    print('accumulated_surface_latent_heat_flux')
    ds_output['accumulated_surface_latent_heat_flux'] = ds['ACLHF']
    ds_output['accumulated_surface_latent_heat_flux'].attrs = {'units': 'J m**-2', 'standard_name': 'integral_of_surface_upward_latent_heat_flux_in_air_wrf_time'}

    print('potential_temperature_variance')
    ds_output['potential_temperature_variance'] = ds['TSQ']
    ds_output['potential_temperature_variance'].attrs = {'units': 'K**2', 'long_name': 'variance_of_air_potential_temperature'}

    print('bulk_richardson_number')
    ds_output['bulk_richardson_number'] = ds['BR']
    ds_output['bulk_richardson_number'].attrs = {'units': '1', 'long_name': 'bulk_richardson_number'}

    print('z_over_l')
    ds_output['z_over_l'] = ds['ZOL']
    ds_output['z_over_l'].attrs = {'units': '1', 'long_name': 'dimensionless_most_z_over_l'}

    print('skin_temperature')
    ds_output['skin_temperature'] = ds['TSK']
    ds_output['skin_temperature'].attrs = {'units': 'K', 'long_name': 'surface_skin_temperature'}

    print('surface_regime')
    ds_output['surface_regime'] = ds['REGIME']
    ds_output['surface_regime'].attrs = {'long_name': 'surface_stability_regime_class'}

    print('roughness_length')
    ds_output['roughness_length'] = ds['ZNT']
    ds_output['roughness_length'].attrs = {'units': 'm', 'long_name': 'time_varying_roughness_length'}

    print('timestep')
    ds_output['timestep'] = ds['ITIMESTEP']
    ds_output['timestep'].attrs = {'long_name': 'model_time_index'}

    for varname in ds_output:
        ds_output[varname].attrs['grid_mapping'] = 'lambert_conformal'
    ds_output['lambert_conformal'] = xr.DataArray(0, name='lambert_conformal',
                                                  attrs=ds_output['u'].metpy.crs.to_dict())

    ds_output.attrs['domain'] = domain
    ds_output.attrs['IBC'] = config

    print(ds_output)

    ds_output = ds_output.drop('metpy_crs')
    ds_output.to_netcdf(run_date.strftime('/glade/scratch/jthielen/cf_post/wrf_{config}_{domain}_%Y%m%d_1200_UTC.nc'.format(config=config, domain=domain)))
    
    ds.close()
    ds_output.close()

    del ds
    del ds_output

# Configuration for loop
"""
run_dates = [datetime(2017, 11, 19),
             datetime(2017, 11, 23),
             datetime(2017, 11, 26),
             datetime(2017, 12, 7)]
configs = ['namanl', 'namfcst', 'gfsanl', 'gfsfcst']
domains = ['d01', 'd02']
"""
run_dates = [
    datetime(2017, 12, 2),
]
configs = ['prototype_3', 'prototype_4']
#configs = ['baseline_4.2', 'baseline_4.2_mixlength_2', 'proto2_jahn_zl_and_sun']
#configs = ['proto1_jahn_zl_and_sun']
#domains = ['d01', 'd02']
domains = ['d02']

if __name__ == '__main__':
    workers = 4
    p = Pool(workers)
    p.starmap(process_run, product(run_dates, configs, domains), chunksize=workers)    
    
print('All done.')