A collection of diagnostic and interpolation routines for use with output from the Weather Research and Forecasting (WRF-ARW) Model.
You can not select more than 25 topics Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
 
 
 
 
 
 

462 lines
17 KiB

from __future__ import (absolute_import, division, print_function,
unicode_literals)
import numpy as np
import numpy.ma as ma
from .extension import (_interpz3d, _vertcross, _interpline, _smooth2d,
_monotonic, _vintrp)
from .metadecorators import set_interp_metadata
from .util import extract_vars, is_staggered, get_id, npvalues
from .py3compat import py3range
from .interputils import get_xy, get_xy_z_params
from .constants import Constants, ConversionFactors
from .terrain import get_terrain
from .geoht import get_height
from .temp import get_theta, get_temp, get_eth
from .pressure import get_pressure
# Note: Extension decorator is good enough to handle left dims
@set_interp_metadata("horiz")
def interplevel(field3d, z, desiredlev, missing=Constants.DEFAULT_FILL,
meta=True):
"""Interpolates a three-dimensional field specified pressure or
height level.
Parameters
----------
field3d : `xarray.DataArray` or `numpy.ndarray`
A three-dimensional field.
z : `xarray.DataArray` or `numpy.ndarray`
A three-dimensional array for the vertical coordinate, typically
pressure or height.
desiredlev : float
The desired vertical level. Must be in the same units as the `z`
parameter.
missing : float
The fill value to use for the output.
`Default is wrf.Constants.DEFAULT_FILL`.
meta : {True, False}
Set to False to disable metadata and return `numpy.ndarray` instead of
`xarray.DataArray`. Default is True.
Returns
-------
`xarray.DataArray` or `numpy.ndarray`
Returns the interpolated variable. If xarray is enabled and
the meta parameter is True, then the result will be an
`xarray.DataArray` object. Otherwise, the result will be a
`numpy.ndarray` object with no metadata.
"""
# Some fields like uvmet have an extra left dimension for the product
# type, we'll handle that iteration here.
multi = True if field3d.ndim - z.ndim == 1 else False
if not multi:
result = _interpz3d(field3d, z, desiredlev, missing)
else:
outshape = field3d.shape[0:-3] + field3d.shape[-2:]
result = np.empty(outshape, dtype=field3d.dtype)
for i in py3range(field3d.shape[0]):
result[i,:] = (
_interpz3d(field3d[i,:], z, desiredlev, missing)[:])
return ma.masked_values (result, missing)
@set_interp_metadata("cross")
def vertcross(field3d, z, missing=Constants.DEFAULT_FILL,
pivot_point=None, angle=None,
start_point=None, end_point=None,
latlon=False, cache=None, meta=True):
"""Return the vertical cross section for a 3D field, interpolated
to a vertical plane defined by a horizontal line.
The horizontal line is defined by either including the
`pivot_point` and `angle` parameters, or the `start_point` and
`end_point` parameters.
Parameters
----------
field3d : `xarray.DataArray` or `numpy.ndarray`
A three-dimensional field.
z : `xarray.DataArray` or `numpy.ndarray`
A three-dimensional array for the vertical coordinate, typically
pressure or height
pivot_point : tuple or list, optional
A tuple or list with two entries, in the form of [x, y]
(or [west_east, south_north]), which indicates the x,y location
through which the plane will pass. Must also specify `angle`.
angle : float, optional
Only valid for cross sections where a plane will be plotted through
a given point on the model domain. 0.0 represents a S-N cross section
and 90.0 a W-E cross section.
start_point : tuple or list, optional
A tuple or list with two entries, in the form of [x, y]
(or [west_east, south_north]), which indicates the start x,y location
through which the plane will pass.
end_point : tuple or list, optional
A tuple or list with two entries, in the form of [x, y]
(or [west_east, south_north]), which indicates the end x,y location
through which the plane will pass.
latlon : {True, False}, optional
Set to True to also interpolate the two-dimensional latitude and
longitude coordinates along the same horizontal line and include
this information in the metadata (if enabled). This can be
helpful for plotting. Default is False.
cache : dict, optional
A dictionary of (varname, numpy.ndarray) pairs which can be used to
supply pre-extracted NetCDF variables to the computational routines.
This can be used to prevent the repeated variable extraction from large
sequences of data files. Default is None.
meta : {True, False}, optional
Set to False to disable metadata and return `numpy.ndarray` instead of
`xarray.DataArray`. Default is True.
Returns
-------
`xarray.DataArray` or `numpy.ndarray`
Returns the interpolated variable. If xarray is enabled and
the meta parameter is True, then the result will be an
`xarray.DataArray` object. Otherwise, the result will be a
`numpy.ndarray` object with no metadata.
"""
# Some fields like uvmet have an extra left dimension for the product
# type, we'll handle that iteration here.
multi = True if field3d.ndim - z.ndim == 1 else False
try:
xy = cache["xy"]
var2dz = cache["var2dz"]
z_var2d = cache["z_var2d"]
except (KeyError, TypeError):
xy, var2dz, z_var2d = get_xy_z_params(npvalues(z), pivot_point, angle,
start_point, end_point)
if not multi:
result = _vertcross(field3d, xy, var2dz, z_var2d, missing)
else:
outshape = field3d.shape[0:-3] + (z_var2d.shape[0], xy.shape[0])
result = np.empty(outshape, dtype=field3d.dtype)
for i in py3range(field3d.shape[0]):
result[i,:] = _vertcross(field3d[i,:], xy, var2dz, z_var2d,
missing)[:]
return ma.masked_values(result, missing)
@set_interp_metadata("line")
def interpline(field2d, pivot_point=None,
angle=None, start_point=None,
end_point=None, latlon=False,
cache=None, meta=True):
"""Return the two-dimensional field interpolated along a line.
Parameters
----------
field2d : `xarray.DataArray` or `numpy.ndarray`
A two-dimensional field.
pivot_point : tuple or list, optional
A tuple or list with two entries, in the form of [x, y]
(or [west_east, south_north]), which indicates the x,y location
through which the plane will pass. Must also specify `angle`.
angle : float, optional
Only valid for cross sections where a plane will be plotted through
a given point on the model domain. 0.0 represents a S-N cross section
and 90.0 a W-E cross section.
start_point : tuple or list, optional
A tuple or list with two entries, in the form of [x, y]
(or [west_east, south_north]), which indicates the start x,y location
through which the plane will pass.
end_point : tuple or list, optional
A tuple or list with two entries, in the form of [x, y]
(or [west_east, south_north]), which indicates the end x,y location
through which the plane will pass.
latlon : {True, False}, optional
Set to True to also interpolate the two-dimensional latitude and
longitude coordinates along the same horizontal line and include
this information in the metadata (if enabled). This can be
helpful for plotting. Default is False.
cache : dict, optional
A dictionary of (varname, numpy.ndarray) pairs which can be used to
supply pre-extracted NetCDF variables to the computational routines.
This can be used to prevent the repeated variable extraction from large
sequences of data files. Default is None.
meta : {True, False}, optional
Set to False to disable metadata and return `numpy.ndarray` instead of
`xarray.DataArray`. Default is True.
Returns
-------
`xarray.DataArray` or `numpy.ndarray`
Returns the interpolated variable. If xarray is enabled and
the meta parameter is True, then the result will be an
`xarray.DataArray` object. Otherwise, the result will be a
`numpy.ndarray` object with no metadata.
"""
try:
xy = cache["xy"]
except (KeyError, TypeError):
xy = get_xy(field2d, pivot_point, angle, start_point, end_point)
return _interpline(field2d, xy)
@set_interp_metadata("vinterp")
def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False,
field_type=None, log_p=False, timeidx=0, method="cat",
squeeze=True, cache=None, meta=True):
"""Return the field vertically interpolated to the given the type of
surface and a set of new levels.
Parameters
----------
wrfnc : `netCD4F.Dataset`, `Nio.NioFile`, or a sequence
Input WRF ARW NetCDF data as a `netCDF4.Dataset`, `Nio.NioFile` or an
iterable sequence of the aforementioned types.
field : `xarray.DataArray` or `numpy.ndarray`
A three-dimensional field.
vert_coord : {'pressure', 'pres', 'p', 'ght_msl', 'ght_agl', 'theta', 'th', 'theta-e', 'thetae', 'eth'}
A string indicating the vertical coordinate type to interpolate to.
Valid strings are:
* 'pressure', 'pres', 'p': pressure [hPa]
* 'ght_msl': grid point height msl [km]
* 'ght_agl': grid point height agl [km]
* 'theta', 'th': potential temperature [K]
* 'theta-e', 'thetae', 'eth': equivalent potential temperature [K]
interp_levels : sequence
A 1D sequence of vertical levels to interpolate to.
extrapolate : {True, False}, optional
Set to True to extrapolate values below ground. Default is False.
field_type : {'none', 'pressure', 'pres', 'p', 'z', 'tc', 'tk', 'theta', 'th', 'theta-e', 'thetae', 'eth', 'ght'}, optional
The type of field. Default is None.
log_p : {True, False}
Use the log of the pressure for interpolation instead of just pressure.
Default is False.
timeidx : int, optional
The time index to use when extracting auxiallary variables used in
the interpolation. This value must be set to match the same value
used when the `field` variable was extracted. Default is 0.
method : {'cat', 'join'}, optional
The aggregation method to use for sequences, either 'cat' or 'join'.
'cat' combines the data along the Time index. 'join' is creates a new
index for the file. This must be set to the same method used when
extracting the `field` variable. The default is 'cat'.
squeeze : {True, False}, optional
Set to False to prevent dimensions with a size of 1 from being removed
from the shape of the output. Default is True.
cache : dict, optional
A dictionary of (varname, ndarray) which can be used to supply
pre-extracted NetCDF variables to the computational routines. This can
be used to prevent the repeated variable extraction from large
sequences of data files. Default is None.
meta : {True, False}, optional
Set to False to disable metadata and return `numpy.ndarray` instead of
`xarray.DataArray`. Default is True.
Returns
-------
`xarray.DataArray` or `numpy.ndarray`
Returns the interpolated variable. If xarray is enabled and
the meta parameter is True, then the result will be an
`xarray.DataArray` object. Otherwise, the result will be a
`numpy.ndarray` object with no metadata.
"""
_key = get_id(wrfnc)
# Remove case sensitivity
field_type = field_type.lower() if field_type is not None else "none"
vert_coord = vert_coord.lower() if vert_coord is not None else "none"
valid_coords = ("pressure", "pres", "p", "ght_msl",
"ght_agl", "theta", "th", "theta-e", "thetae", "eth")
valid_field_types = ("none", "pressure", "pres", "p", "z",
"tc", "tk", "theta", "th", "theta-e", "thetae",
"eth", "ght")
icase_lookup = {"none" : 0,
"p" : 1,
"pres" : 1,
"pressure" : 1,
"z" : 2,
"ght" : 2,
"tc" : 3,
"tk" : 4,
"theta" : 5,
"th" : 5,
"theta-e" : 6,
"thetae" : 6,
"eth" : 6}
# These constants match what's in the fortran code.
rgas = Constants.RD
ussalr = Constants.USSALR
sclht = Constants.SCLHT
# interp_levels might be a list or tuple, make a numpy array
if not isinstance(interp_levels, np.ndarray):
interp_levels = np.asarray(interp_levels, np.float64)
# TODO: Check if field is staggered
if is_staggered(field, wrfnc):
raise RuntimeError("Please unstagger field in the vertical")
# Check for valid coord
if vert_coord not in valid_coords:
raise RuntimeError("'%s' is not a valid vertical "
"coordinate type" % vert_coord)
# Check for valid field type
if field_type not in valid_field_types:
raise RuntimeError("'%s' is not a valid field type" % field_type)
log_p_int = 1 if log_p else 0
icase = 0
extrap = 0
if extrapolate:
extrap = 1
icase = icase_lookup[field_type]
# Extract vriables
#timeidx = -1 # Should this be an argument?
ncvars = extract_vars(wrfnc, timeidx, ("PSFC", "QVAPOR", "F"),
method, squeeze, cache, meta=False, _key=_key)
sfp = ncvars["PSFC"] * ConversionFactors.PA_TO_HPA
qv = ncvars["QVAPOR"]
coriolis = ncvars["F"]
terht = get_terrain(wrfnc, timeidx, units="m",
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
t = get_theta(wrfnc, timeidx, units="k",
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
tk = get_temp(wrfnc, timeidx, units="k",
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
p = get_pressure(wrfnc, timeidx, units="pa",
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
ght = get_height(wrfnc, timeidx, msl=True, units="m",
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
ht_agl = get_height(wrfnc, timeidx, msl=False, units="m",
method=method, squeeze=squeeze, cache=cache,
meta=False, _key=_key)
smsfp = _smooth2d(sfp, 3)
vcor = 0
if vert_coord in ("pressure", "pres", "p"):
vcor = 1
vcord_array = p * ConversionFactors.PA_TO_HPA
elif vert_coord == "ght_msl":
vcor = 2
vcord_array = np.exp(-ght/sclht)
elif vert_coord == "ght_agl":
vcor = 3
vcord_array = np.exp(-ht_agl/sclht)
elif vert_coord in ("theta", "th"):
vcor = 4
idir = 1
icorsw = 0
delta = 0.01
p_hpa = p * ConversionFactors.PA_TO_HPA
vcord_array = _monotonic(t, p_hpa, coriolis, idir, delta, icorsw)
# We only extrapolate temperature fields below ground
# if we are interpolating to pressure or height vertical surfaces.
icase = 0
elif vert_coord in ("theta-e", "thetae", "eth"):
vcor = 5
icorsw = 0
idir = 1
delta = 0.01
eth = get_eth(wrfnc, timeidx, method=method, squeeze=squeeze,
cache=cache, meta=False, _key=_key)
p_hpa = p * ConversionFactors.PA_TO_HPA
vcord_array = _monotonic(eth, p_hpa, coriolis, idir, delta, icorsw)
# We only extrapolate temperature fields below ground if we are
# interpolating to pressure or height vertical surfaces
icase = 0
# Set the missing value
if isinstance(field, ma.MaskedArray):
missing = field.fill_value
else:
missing = Constants.DEFAULT_FILL
if (field.shape != p.shape):
raise ValueError("'field' shape does not match other variable shapes. "
"Verify that the 'timeidx' parameter matches the "
"same value used when extracting the 'field' "
"variable.")
res = _vintrp(field, p, tk, qv, ght, terht, sfp, smsfp,
vcord_array, interp_levels,
icase, extrap, vcor, log_p_int, missing)
return ma.masked_values(res, missing)