nilyin
/
wrf-python
forked from 3rdparty/wrf-python
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity
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.
404 Commits
5 Branches
47 Tags
40 MiB
 
 
 
 
 
 
Tree: b40d1d6e2d
Branches Tags
${ item.name }
Create tag ${ searchTerm }
Create branch ${ searchTerm }
from 'b40d1d6e2d'
${ noResults }
wrf-python/src/wrf/util.py

3927 lines
132 KiB
Raw Blame History

from __future__ import (absolute_import, division, print_function)
import os
from sys import version_info
from copy import copy
from collections import Iterable, Mapping, OrderedDict
from itertools import product, tee
from types import GeneratorType
import datetime as dt
from inspect import getmodule
try:
from inspect import signature
except ImportError:
from inspect import getargspec
try:
from inspect import getargvalues
except ImportError:
from inspect import getgeneratorlocals
# Needed for Python 3.4 to use apply_defaults
try:
from inspect import _empty, _VAR_POSITIONAL, _VAR_KEYWORD
except ImportError:
pass
import numpy as np
import numpy.ma as ma
from .config import xarray_enabled
from .constants import default_fill, ALL_TIMES
from .py3compat import (viewitems, viewkeys, isstr, py3range)
from .cache import cache_item, get_cached_item
from .geobnds import GeoBounds, NullGeoBounds
from .coordpair import CoordPair
from .projection import getproj
if xarray_enabled():
from xarray import DataArray
_COORD_PAIR_MAP = {"XLAT": ("XLAT", "XLONG"),
"XLONG": ("XLAT", "XLONG"),
"XLAT_M": ("XLAT_M", "XLONG_M"),
"XLONG_M": ("XLAT_M", "XLONG_M"),
"XLAT_U": ("XLAT_U", "XLONG_U"),
"XLONG_U": ("XLAT_U", "XLONG_U"),
"XLAT_V": ("XLAT_V", "XLONG_V"),
"XLONG_V": ("XLAT_V", "XLONG_V"),
"CLAT": ("CLAT", "CLONG"),
"CLONG": ("CLAT", "CLONG")}
_COORD_VARS = ("XLAT", "XLONG", "XLAT_M", "XLONG_M", "XLAT_U", "XLONG_U",
"XLAT_V", "XLONG_V", "CLAT", "CLONG")
_LAT_COORDS = ("XLAT", "XLAT_M", "XLAT_U", "XLAT_V", "CLAT")
_LON_COORDS = ("XLONG", "XLONG_M", "XLONG_U", "XLONG_V", "CLONG")
_TIME_COORD_VARS = ("XTIME",)
def is_time_coord_var(varname):
"""Return True if the input variable name is a time coordinate.
Args:
varname (:obj:`str`): The input variable name.
Returns:
:obj:`bool`: True if the input variable is a time coordinate,
otherwise False.
"""
return varname in _TIME_COORD_VARS
def get_coord_pairs(coord_varname):
"""Return a :obj:`tuple` for the variable names of the coordinate pair used
for the 2D curvilinear coordinate variable.
For example, the 'XLAT' variable will have coordinate variables of
('XLAT', 'XLONG') since the 'XLAT' variable itself is two-dimensional.
Args:
coord_varname (:obj:`str`): The coordinate variable name.
Returns:
:obj:`bool`: True if the time index is :data:`wrf.ALL_TIMES` or
:obj:`None`, otherwise False.
"""
return _COORD_PAIR_MAP[coord_varname]
def is_multi_time_req(timeidx):
"""Return True if the requested time index is for :data:`wrf.ALL_TIMES` or
:obj:`None`.
Args:
timeidx (:obj:`int`, :data:`wrf.ALL_TIMES`, or :obj:`None`): The
requested time index.
Returns:
:obj:`bool`: True if the time index is :data:`wrf.ALL_TIMES` or
:obj:`None`, otherwise False.
"""
return timeidx is None
def is_multi_file(wrfin):
"""Return True if the input argument is an iterable.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
Returns:
:obj:`bool`: True if the input is an iterable. False if the input
is a single NetCDF file object.
"""
return (isinstance(wrfin, Iterable) and not isstr(wrfin))
def has_time_coord(wrfnc):
"""Return True if the input file or sequence contains the time
coordinate variable.
The time coordinate is named 'XTIME'.
Args:
wrfnc (:class:`netCDF4.Dataset` or :class:`Nio.NioFile`): A single
NetCDF file object.
Returns:
:obj:`bool`: True if the netcdf file contains the time coordinate
variable, False otherwise.
"""
return "XTIME" in wrfnc.variables
def is_mapping(wrfin):
"""Return True if the input file or sequence is a mapping type.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
Returns:
:obj:`bool`: True if the input is a mapping type, False otherwise.
"""
return isinstance(wrfin, Mapping)
def _generator_copy(gen):
"""Return a copy of a generator.
This function instantiates a new generator object using the arguments
passed to the original.
Args:
gen (:class:`types.GeneratorType`): A generator object.
Note:
In order for this to work correctly, the generator cannot modify
the original construction arguments.
Returns:
:class:`types.GeneratorType`: A copy of the generator object.
"""
funcname = gen.__name__
# This is for generator expressions. Only solution is to tee the
# original generator.
if funcname == "<genexpr>":
return tee(gen)
try:
argvals = getargvalues(gen.gi_frame)
except NameError:
argvals = getgeneratorlocals(gen)
module = getmodule(gen.gi_frame)
if module is not None:
try:
try:
argd = {key: argvals.locals[key] for key in argvals.args}
res = module.get(funcname)(**argd)
except AttributeError:
res = getattr(module, funcname)(**argd)
except Exception:
# This is the old way it used to work, but it looks like this was
# fixed by Python.
try:
res = module.get(funcname)(**argvals.locals)
except AttributeError:
res = getattr(module, funcname)(**argvals.locals)
else:
# Created in jupyter or the python interpreter
import __main__
try:
argd = {key: argvals.locals[key] for key in argvals.args}
res = getattr(__main__, funcname)(**argd)
except Exception:
# This was the old way it used to work, but appears to have
# been fixed by Python.
res = getattr(__main__, funcname)(**argvals.locals)
return res
def test():
q = [1, 2, 3]
for i in q:
yield i
class TestGen(object):
def __init__(self, count=3):
self._total = count
self._i = 0
def __iter__(self):
return self
def next(self):
if self._i >= self._total:
raise StopIteration
else:
val = self._i
self._i += 1
return val
# Python 3
def __next__(self):
return self.next()
def latlon_coordvars(ncvars):
"""Return the first found latitude and longitude coordinate names from a
NetCDF variable dictionary.
This function searches the dictionary structure for NetCDF variables
and returns the first found latitude and longitude coordinate
names (typically 'XLAT' and 'XLONG').
Args:
ncvars (:obj:`dict`): A NetCDF variable dictionary object.
Returns:
:obj:`tuple`: The latitude and longitude coordinate name pairs as
(lat_coord_name, lon_coord_name).
"""
lat_coord = None
lon_coord = None
for name in _LAT_COORDS:
if name in viewkeys(ncvars):
lat_coord = name
break
for name in _LON_COORDS:
if name in viewkeys(ncvars):
lon_coord = name
break
return lat_coord, lon_coord
def is_coordvar(varname):
"""Returns True if the variable is a coordinate variable.
Args:
varname (:obj:`str`): The variable name.
Returns:
:obj:`bool`: True if the variable is a coordinate variable,
otherwise False.
"""
return varname in _COORD_VARS
class IterWrapper(Iterable):
"""A wrapper class for generators and custom iterable classes that returns
a new iterator to the start of the sequence when
:meth:`IterWrapper.__iter__` is called.
If the wrapped object is a generator, a copy of the generator is
constructed and returned when :meth:`IterWrapper.__iter__` is called.
If the wrapped object is a custom type, then the :meth:`copy.copy` is
called and a new instance is returned. In both cases, the original
iterable object is unchanged.
Note:
Do not increment the wrapped iterable outside of this wrapper.
"""
def __init__(self, wrapped):
"""Initialize an :class:`wrf.IterWrapper` object.
Args:
wrapped (an iterable object): Any iterable object that contains the
*__iter__* method.
"""
self._wrapped = wrapped
def __iter__(self):
"""Return an iterator to the start of the sequence.
Returns:
An iterator to the start of the sequence.
"""
if isinstance(self._wrapped, GeneratorType):
gen_copy = _generator_copy(self._wrapped)
# If a tuple comes back, then this is a generator expression,
# so store the first tee'd item, then return the other
if isinstance(gen_copy, tuple):
self._wrapped = gen_copy[0]
return gen_copy[1]
return gen_copy
else:
obj_copy = copy(self._wrapped)
return obj_copy.__iter__()
def get_iterable(wrfseq):
"""Returns a resettable iterable object.
In this context, resettable means that when :meth:`object.__iter__`
is called, the iterable returned always points to the first element
in the sequence, similar to how the list and tuple behave.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
Returns:
iterable: A resettable iterator object.
"""
if not is_multi_file(wrfseq):
return wrfseq
else:
if not is_mapping(wrfseq):
if isinstance(wrfseq, (list, tuple, IterWrapper)):
return wrfseq
else:
return IterWrapper(wrfseq) # generator/custom iterable class
else:
if isinstance(wrfseq, dict):
return wrfseq
else:
return dict(wrfseq) # generator/custom iterable dict class
# Helper to extract masked arrays from DataArrays that convert to NaN
def to_np(array):
"""Return the :class:`numpy.ndarray` contained in an
:class:`xarray.DataArray` instance.
If the :class:`xarray.DataArray` instance does not contain a *_FillValue*
or *missing_value* attribute, then this routine simply returns the
:attr:`xarray.DataArray.values` attribute. If the
:class:`xarray.DataArray` object contains a *_FillValue* or *missing_value*
attribute, then this routine returns a :class:`numpy.ma.MaskedArray`
instance, where the NaN values (used by xarray to represent missing data)
are replaced with the fill value.
If the object passed in to this routine is not an
:class:`xarray.DataArray` instance, then this routine simply returns the
passed in object. This is useful in situations where you do not know
if you have an :class:`xarray.DataArray` or a :class:`numpy.ndarray` and
simply want a :class:`numpy.ndarray` returned.
Args:
array (:class:`xarray.DataArray`, :class:`numpy.ndarray`, or any \
object): Can be any object type, but is generally
used with :class:`xarray.DataArray` or :class:`numpy.ndarray`.
Returns:
:class:`numpy.ndarray` or :class:`numpy.ma.MaskedArray`: The
extracted array or the *array* object if *array* is not a
:class:`xarray.DataArray` object..
"""
try:
fill_value = array.attrs["_FillValue"]
except AttributeError:
result = array # Not a DataArray
except KeyError:
result = array.values # Does not have missing values
else:
result = ma.masked_invalid(array.values, copy=False)
result.set_fill_value(fill_value)
return result
# Helper utilities for metadata
class either(object):
"""A callable class that determines which variable is present in the
file.
This is used in situations where the same variable type has different names
depending on the type of file used. For example, in a WRF output file,
'P' is used for pressure, whereas in a met_em file, pressure is named
'PRES'.
Methods:
__call__(wrfin): Return the variable that is present in the file.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
Returns:
:obj:`str`: The variable name that is present in the file.
Attributes:
varnames (sequence): A sequence of possible variable names.
"""
def __init__(self, *varnames):
"""Initialize an :class:`either` object.
Args:
*varnames (sequence): A sequence of possible variable names.
"""
self.varnames = varnames
def __call__(self, wrfin):
if is_multi_file(wrfin):
if not is_mapping(wrfin):
wrfin = next(iter(wrfin))
else:
entry = wrfin[next(iter(viewkeys(wrfin)))]
return self(entry)
for varname in self.varnames:
if varname in wrfin.variables:
return varname
raise ValueError("{} are not valid variable names".format(
self.varnames))
# This should look like:
# [(0, (-3,-2)), # variable 1
# (1, -1)] # variable 2
class combine_dims(object):
"""A callable class that mixes dimension sizes from different function
arguments.
This callable object is used for determining the output size for the
extension module functions. The class is initialized with a sequence of
pairs, where the first value is the function argument index. The second
value is the dimension index(es) to use. The second value can be a
single integer or a sequence if multiple dimensions are used.
Methods:
__call__(*args): Return a tuple with the combined dimension sizes.
Args:
*args: The function arguments from which to extract the
dimensions sizes.
Returns:
:obj:`tuple`: The shape for the combined dimensions.
Attributes:
pairs (sequence): A sequence representing how to combine the
dimensions.
Example:
.. code-block:: python
# Take the -3, and -2 dimension sizes from argument 0
# Then take the -1 dimension size from argument 1
pairs = [(0, (-3, -2), (1, -1)]
combo = combine_dims(pairs)
"""
def __init__(self, pairs):
"""Initialize a :class:`combine_dims` object.
Args:
pairs (sequence): A sequence where each element is a pair
(:obj:`tuple`), with the first element being the function
argument index and the second value being either an integer
or sequence for the dimension size indexes to use.
"""
self.pairs = pairs
def __call__(self, *args):
result = []
for pair in self.pairs:
var = args[pair[0]]
dimidxs = pair[1]
if isinstance(dimidxs, Iterable):
for dimidx in dimidxs:
result.append(var.shape[dimidx])
else:
result.append(var.shape[dimidxs])
return tuple(result)
class from_var(object):
"""A callable class that retrieves attributes from the function argument.
If the function argument is not of type :class:`xarray.DataArray`, then
None will be returned.
It is assumed that the function has been wrapped using the :mod:`wrapt`
module.
Methods:
__call__(wrapped, *args, **kwargs): Return the attribute found in \
the function arguments.
Args:
wrapped: The wrapped function, as used by the :mod:`wrapt`
module.
*args: The function arguments.
**kwargs: The function keyword arguments.
Returns:
:obj:`object`: The requested attribute.
Attributes:
varname (:obj:`str`): The variable name.
attribute (:obj:`str`): The attribute name.
"""
def __init__(self, varname, attribute):
"""Initialize a :class:`from_var` object.
Args:
varname (:obj:`str`): The variable name.
attribute (:obj:`str`): The attribute name.
"""
self.varname = varname
self.attribute = attribute
def __call__(self, wrapped, *args, **kwargs):
vard = from_args(wrapped, (self.varname,), *args, **kwargs)
var = None
if vard is not None:
var = vard[self.varname]
try:
return var.attrs.get(self.attribute, None)
except AttributeError:
return None
def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar):
"""Return True if the corner points have moved.
This function is used to test for a moving domain, since the WRF output
does not set any flags in the file for this. The test will be performed
for all time steps in the NetCDF file.
Args:
wrfnc (:class:`netCDF4.Dataset` or :class:`Nio.NioFile`): A single
NetCDF file object.
first_ll_corner (:obj:`tuple`): A (latitude, longitude) pair for the
lower left corner found in the initial file.
first_ur_corner (:obj:`tuple`): A (latitude, longitude) pair for the
upper right corner found in the initial file.
latvar (:obj:`str`): The latitude variable name to use.
lonvar (:obj:`str`: The longitude variable name to use.
Returns:
:obj:`bool`: True if the corner points have moved, False otherwise.
"""
lats = wrfnc.variables[latvar][:]
lons = wrfnc.variables[lonvar][:]
# Need to check all times
for i in py3range(lats.shape[-3]):
start_idxs = [0] * len(lats.shape) # PyNIO does not support ndim
start_idxs[-3] = i
start_idxs = tuple(start_idxs)
end_idxs = [-1] * len(lats.shape)
end_idxs[-3] = i
end_idxs = tuple(end_idxs)
if (first_ll_corner[0] != lats[start_idxs]
or first_ll_corner[1] != lons[start_idxs]
or first_ur_corner[0] != lats[end_idxs]
or first_ur_corner[1] != lons[end_idxs]):
return True
return False
def is_moving_domain(wrfin, varname=None, latvar=either("XLAT", "XLAT_M"),
lonvar=either("XLONG", "XLONG_M"), _key=None):
"""Return True if the domain is a moving nest.
This function is used to test for a moving domain, since the WRF output
does not set any flags in the file for this. The test will be performed
for all files in any sequences and across all times in each file.
This result is cached internally, so this potentially lengthy check is
only done one time for any given *wrfin* parameter.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
varname (:obj:`str`, optional): A specific NetCDF variable to test,
which must contain a 'coordinates' attribute. If unspecified,
The *latvar* and *lonvar* parameters are used. Default is None.
first_ur_corner (:obj:`tuple`): A (latitude, longitude) pair for the
upper right corner found in the initial file.
latvar (:obj:`str` or :class:`either`, optional): The latitude variable
name. Default is :class:`either('XLAT', 'XLAT_M')`.
lonvar (:obj:`str` or :class:`either`, optional): The latitude variable
name. Default is :class:`either('XLAT', 'XLAT_M')`.
_key (:obj:`int`, optional): A caching key. This is used for internal
purposes only. Default is None.
Returns:
:obj:`bool`: True if the domain is a moving nest, False otherwise.
"""
if isinstance(latvar, either):
latvar = latvar(wrfin)
if isinstance(lonvar, either):
lonvar = lonvar(wrfin)
# In case it's just a single file
if not is_multi_file(wrfin):
wrfin = [wrfin]
# Slow, but safe. Compare the corner points to the first item and see
# any move. User iterator protocol in case wrfin is not a list/tuple.
if not is_mapping(wrfin):
wrf_iter = iter(wrfin)
first_wrfnc = next(wrf_iter)
else:
# Currently only checking the first dict entry.
dict_key = next(iter(viewkeys(wrfin)))
entry = wrfin[dict_key]
key = _key[dict_key] if _key is not None else None
return is_moving_domain(entry, varname, latvar, lonvar, key)
# The long way of checking all lat/lon corner points. Doesn't appear
# to be a shortcut in the netcdf files.
if varname is not None:
try:
coord_str = getattr(first_wrfnc.variables[varname], "coordinates")
# scipy.io.netcdf stores attributes as bytes rather than str
if isinstance(coord_str, str):
coord_names = coord_str.split()
else:
coord_names = coord_str.decode().split()
except AttributeError:
# Variable doesn't have a coordinates attribute, use the
# arguments
lon_coord = lonvar
lat_coord = latvar
else:
for name in coord_names:
if name in _LAT_COORDS:
lat_coord = name
continue
if name in _LON_COORDS:
lon_coord = name
continue
else:
lon_coord = lonvar
lat_coord = latvar
# See if there is a cached value
product = "is_moving_{}_{}".format(lat_coord, lon_coord)
moving = get_cached_item(_key, product)
if moving is not None:
return moving
# Need to search all the files
lats = first_wrfnc.variables[lat_coord][:]
lons = first_wrfnc.variables[lon_coord][:]
zero_idxs = [0]*len(lats.shape) # PyNIO doesn't have ndim
last_idxs = list(zero_idxs)
last_idxs[-2:] = [-1]*2
zero_idxs = tuple(zero_idxs)
last_idxs = tuple(last_idxs)
lat0 = lats[zero_idxs]
lat1 = lats[last_idxs]
lon0 = lons[zero_idxs]
lon1 = lons[last_idxs]
ll_corner = (lat0, lon0)
ur_corner = (lat1, lon1)
# Need to check if the first file is moving, might be a single
# file with multiple times
if _corners_moved(first_wrfnc, ll_corner, ur_corner, lat_coord, lon_coord):
cache_item(_key, product, True)
return True
# Now check any additional files
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
if _corners_moved(wrfnc, ll_corner, ur_corner,
lat_coord, lon_coord):
cache_item(_key, product, True)
return True
cache_item(_key, product, False)
return False
def _get_global_attr(wrfnc, attr):
"""Return the global attribute.
Args:
wrfnc (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`): A single
WRF NetCDF file object.
attr (:obj:`str`): The attribute name.
Returns:
:obj:`object`: The global attribute.
"""
val = getattr(wrfnc, attr, None)
# PyNIO puts single values in to an array
if isinstance(val, np.ndarray):
if len(val) == 1:
return val[0]
return val
def extract_global_attrs(wrfin, attrs):
"""Return the global attribute(s).
If the *wrfin* parameter is a sequence, then only the first element is
used, so the entire sequence must have the same global attributes.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
attrs (:obj:`str`, sequence): The attribute name
or a sequence of attribute names.
Returns:
:obj:`dict`: A mapping of attribute_name to value.
"""
if isstr(attrs):
attrlist = [attrs]
else:
attrlist = attrs
multifile = is_multi_file(wrfin)
if multifile:
if not is_mapping(wrfin):
wrfin = next(iter(wrfin))
else:
entry = wrfin[next(iter(viewkeys(wrfin)))]
return extract_global_attrs(entry, attrs)
return {attr: _get_global_attr(wrfin, attr) for attr in attrlist}
def extract_dim(wrfin, dim):
"""Return the dimension size for the specified dimension name.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
dim (:obj:`str`): The dimension name.
Returns:
:obj:`int`: The dimension size.
"""
if is_multi_file(wrfin):
if not is_mapping(wrfin):
wrfin = next(iter(wrfin))
else:
entry = wrfin[next(iter(viewkeys(wrfin)))]
return extract_dim(entry, dim)
d = wrfin.dimensions[dim]
if not isinstance(d, int):
try:
return len(d) # netCDF4
except TypeError: # scipy.io.netcdf
# Scipy can't handled unlimited dimensions, so now we have to
# figure it out
try:
s = wrfin.variables["P"].shape
except Exception:
raise ValueError("unsupported NetCDF reader")
else:
return s[-4]
return d # PyNIO
def _combine_dict(wrfdict, varname, timeidx, method, meta, _key):
"""Return an array object from a mapping input.
The resulting array is the combination of all sequences associated with
each key in the dictionary. The leftmost dimension will be the keys. The
rightmost dimensions are the dimensions for the aggregated sequences of
arrays, using either the 'cat' or 'join' *method* parameter value.
If dictionaries are nested, then the outermost dictionary keys will be the
leftmost dimension, followed by each subsequent dictionary's keys.
If the order of the keys (and leftmost dimension ordering) matters, it is
recommended that an :class:`OrderedDict` be used instead of a normal
dictionary. Otherwise, the leftmost dimension will be ordered by the
iteration order.
Args:
wrfdict (mapping): A mapping of key to an array or
key to a sequence of arrays.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
keynames = []
numkeys = len(wrfdict)
key_iter = iter(viewkeys(wrfdict))
first_key = next(key_iter)
keynames.append(first_key)
is_moving = is_moving_domain(wrfdict, varname, _key=_key)
_cache_key = _key[first_key] if _key is not None else None
first_array = _extract_var(wrfdict[first_key], varname,
timeidx, is_moving=is_moving, method=method,
squeeze=False, cache=None, meta=meta,
_key=_cache_key)
# Create the output data numpy array based on the first array
outdims = [numkeys]
outdims += first_array.shape
outdata = np.empty(outdims, first_array.dtype)
outdata[0, :] = first_array[:]
idx = 1
while True:
try:
key = next(key_iter)
except StopIteration:
break
else:
keynames.append(key)
_cache_key = _key[key] if _key is not None else None
vardata = _extract_var(wrfdict[key], varname, timeidx,
is_moving=is_moving, method=method,
squeeze=False, cache=None, meta=meta,
_key=_cache_key)
if outdata.shape[1:] != vardata.shape:
raise ValueError("data sequences must have the "
"same size for all dictionary keys")
outdata[idx, :] = to_np(vardata)[:]
idx += 1
if xarray_enabled() and meta:
outname = str(first_array.name)
# Note: assumes that all entries in dict have same coords
outcoords = OrderedDict(first_array.coords)
# First find and store all the existing key coord names/values
# This is applicable only if there are nested dictionaries.
key_coordnames = []
coord_vals = []
existing_cnt = 0
while True:
key_coord_name = "key_{}".format(existing_cnt)
if key_coord_name not in first_array.dims:
break
key_coordnames.append(key_coord_name)
coord_vals.append(to_np(first_array.coords[key_coord_name]))
existing_cnt += 1
# Now add the key coord name and values for THIS dictionary.
# Put the new key_n name at the bottom, but the new values will
# be at the top to be associated with key_0 (left most). This
# effectively shifts the existing 'key_n' coordinate values to the
# right one dimension so *this* dicionary's key coordinate values
# are at 'key_0'.
key_coordnames.append(key_coord_name)
coord_vals.insert(0, keynames)
# make it so that key_0 is leftmost
outdims = key_coordnames + list(first_array.dims[existing_cnt:])
# Create the new 'key_n', value pairs
for coordname, coordval in zip(key_coordnames, coord_vals):
outcoords[coordname] = coordval
outattrs = OrderedDict(first_array.attrs)
outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdims, attrs=outattrs)
else:
outarr = outdata
return outarr
def _find_coord_names(coords):
"""Return the coordinate variables names found in a
:attr:`xarray.DataArray.coords` mapping.
Args:
coords (mapping): A :attr:`xarray.DataArray.coords` mapping object.
Returns:
:obj:`tuple`: The latitude, longitude, and xtime variable names used
in the coordinate mapping.
"""
try:
lat_coord = [name for name in _COORD_VARS[0::2] if name in coords][0]
except IndexError:
lat_coord = None
try:
lon_coord = [name for name in _COORD_VARS[1::2] if name in coords][0]
except IndexError:
lon_coord = None
try:
xtime_coord = [name for name in _TIME_COORD_VARS if name in coords][0]
except IndexError:
xtime_coord = None
return lat_coord, lon_coord, xtime_coord
def _find_max_time_size(wrfseq):
"""Return the maximum number of times found in a sequence of
WRF files.
Args:
wrfseq (sequence): A sequence of WRF NetCDF file objects.
Returns:
:obj:`int`: The maximum number of times found in a file.
"""
wrf_iter = iter(wrfseq)
max_times = 0
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
t = extract_dim(wrfnc, "Time")
max_times = t if t >= max_times else max_times
return max_times
def _get_coord_names(wrfin, varname):
# Need only the first real file
if is_multi_file(wrfin):
if not is_mapping(wrfin):
wrfnc = next(iter(wrfin))
else:
entry = wrfin[next(iter(viewkeys(wrfin)))]
return _get_coord_names(entry, varname)
else:
wrfnc = wrfin
lat_coord = None
lon_coord = None
time_coord = None
var = wrfnc.variables[varname]
# WRF variables will have a coordinates attribute. MET_EM files have
# a stagger attribute which indicates the coordinate variable.
try:
# WRF files
coord_attr = getattr(var, "coordinates")
except AttributeError:
if is_coordvar(varname):
# Coordinate variable (most likely XLAT or XLONG)
lat_coord, lon_coord = get_coord_pairs(varname)
time_coord = None
if has_time_coord(wrfnc):
time_coord = "XTIME"
elif is_time_coord_var(varname):
lon_coord = None
lat_coord = None
time_coord = None
else:
try:
# met_em files or old WRF files
stag_attr = getattr(var, "stagger")
except AttributeError:
lon_coord = None
lat_coord = None
# Let's just check for xlat and xlong in this case
if "XLAT" in wrfnc.variables:
lat_coord = "XLAT"
lon_coord = "XLONG"
else:
# For met_em files, use the stagger name to get the lat/lon var
lat_coord = "XLAT_{}".format(stag_attr)
lon_coord = "XLONG_{}".format(stag_attr)
# If this coord name is missing, it might be an old WRF file
if lat_coord not in wrfnc.variables:
lat_coord = None
lon_coord = None
if "XLAT" in wrfnc.variables:
lat_coord = "XLAT"
lon_coord = "XLONG"
else:
if isinstance(coord_attr, str):
coord_names = coord_attr.split()
else:
coord_names = coord_attr.decode().split()
for name in coord_names:
if name in _LAT_COORDS:
lat_coord = name
continue
if name in _LON_COORDS:
lon_coord = name
continue
if name in _TIME_COORD_VARS:
time_coord = name
continue
if time_coord is not None:
# Make sure the time variable wasn't removed
try:
_ = wrfnc.variables[time_coord]
except KeyError:
time_coord = None
return lat_coord, lon_coord, time_coord
def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile,
_key):
"""Return a :class:`xarray.DataArray` object for the desired variable in
a single NetCDF file object.
Args:
wrfnc (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`): A single
WRF NetCDF file object.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
is_moving_domain (:obj:`bool`): A boolean type that indicates if the
NetCDF file object came from a moving nest.
is_multifile (:obj:`bool`): A boolean type that indicates if the NetCDF
file object came from a sequence.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray`: An array object that contains metadata.
"""
# Note: wrfnc is always a single netcdf file object
# is_moving_domain and is_multifile are arguments indicating if the
# single file came from a sequence, and if that sequence is has a moving
# domain. Both arguments are used mainly for coordinate extraction and
# caching.
multitime = is_multi_time_req(timeidx)
time_idx_or_slice = timeidx if not multitime else slice(None)
var = wrfnc.variables[varname]
if len(var.shape) > 1:
data = var[time_idx_or_slice, :]
else:
data = var[time_idx_or_slice]
# Want to preserve the time dimension
if not multitime:
if len(var.shape) > 1:
data = data[np.newaxis, :]
else:
data = data[np.newaxis]
attrs = OrderedDict()
for dkey, val in viewitems(var.__dict__):
# scipy.io adds these but don't want them
if dkey in ("data", "_shape", "_size", "_typecode", "_attributes",
"maskandscale", "dimensions"):
continue
_dkey = dkey if isinstance(dkey, str) else dkey.decode()
if isstr(val):
_val = val
else:
if isinstance(val, bytes):
_val = val.decode() # scipy.io.netcdf
else:
_val = val
attrs[_dkey] = _val
dimnames = var.dimensions[-data.ndim:]
lat_coord = lon_coord = time_coord = None
try:
if dimnames[-2] == "south_north" and dimnames[-1] == "west_east":
lat_coord, lon_coord, time_coord = _get_coord_names(wrfnc, varname)
except IndexError:
pass
coords = OrderedDict()
# Handle lat/lon coordinates and projection information if available
if lon_coord is not None and lat_coord is not None:
# Using a cache for coordinate variables so the extraction only happens
# once.
lon_coord_dimkey = lon_coord + "_dim"
lon_coord_valkey = lon_coord + "_val"
lat_coord_dimkey = lat_coord + "_dim"
lat_coord_valkey = lat_coord + "_val"
lon_coord_dims = get_cached_item(_key, lon_coord_dimkey)
lon_coord_vals = get_cached_item(_key, lon_coord_valkey)
if lon_coord_dims is None or lon_coord_vals is None:
lon_var = wrfnc.variables[lon_coord]
lon_coord_dims = lon_var.dimensions
lon_coord_vals = lon_var[:]
# Only cache here if the domain is not moving, otherwise
# caching is handled by cat/join
if not is_moving_domain:
cache_item(_key, lon_coord_dimkey, lon_coord_dims)
cache_item(_key, lon_coord_valkey, lon_coord_vals)
lat_coord_dims = get_cached_item(_key, lat_coord_dimkey)
lat_coord_vals = get_cached_item(_key, lat_coord_valkey)
if lat_coord_dims is None or lat_coord_vals is None:
lat_var = wrfnc.variables[lat_coord]
lat_coord_dims = lat_var.dimensions
lat_coord_vals = lat_var[:]
# Only cache here if the domain is not moving, otherwise
# caching is done in cat/join
if not is_moving_domain:
cache_item(_key, lat_coord_dimkey, lat_coord_dims)
cache_item(_key, lat_coord_valkey, lat_coord_vals)
time_coord_vals = None
if time_coord is not None:
# If not from a multifile sequence, then cache the time
# coordinate. Otherwise, handled in cat/join/
if not is_multifile:
time_coord_vals = get_cached_item(_key, time_coord)
if time_coord_vals is None:
time_coord_vals = wrfnc.variables[time_coord][:]
if not is_multifile:
cache_item(_key, time_coord, time_coord_vals)
else:
time_coord_vals = wrfnc.variables[time_coord][:]
if multitime:
if is_moving_domain:
# Special case with a moving domain in a multi-time file,
# otherwise the projection parameters don't change
coords[lon_coord] = lon_coord_dims, lon_coord_vals
coords[lat_coord] = lat_coord_dims, lat_coord_vals
else:
coords[lon_coord] = (lon_coord_dims[1:],
lon_coord_vals[0, :])
coords[lat_coord] = (lat_coord_dims[1:],
lat_coord_vals[0, :])
if time_coord is not None:
coords[time_coord] = (lon_coord_dims[0], time_coord_vals)
else:
coords[lon_coord] = (lon_coord_dims[1:],
lon_coord_vals[timeidx, :])
coords[lat_coord] = (lat_coord_dims[1:],
lat_coord_vals[timeidx, :])
if time_coord is not None:
coords[time_coord] = (lon_coord_dims[0],
[time_coord_vals[timeidx]])
proj_params = get_proj_params(wrfnc)
proj = getproj(**proj_params)
attrs["projection"] = proj
if dimnames[0] == "Time":
t = extract_times(wrfnc, timeidx, meta=False, do_xtime=False)
if not multitime:
t = [t]
coords[dimnames[0]] = t
data_array = DataArray(data, name=varname, dims=dimnames, coords=coords,
attrs=attrs)
return data_array
def _find_forward(wrfseq, varname, timeidx, is_moving, meta, _key):
"""Find and return the array object within a sequence for a specific time
index.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int`): The desired time index. Must be positive.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
wrf_iter = iter(wrfseq)
comboidx = 0
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
numtimes = extract_dim(wrfnc, "Time")
if timeidx < comboidx + numtimes:
filetimeidx = timeidx - comboidx
if meta:
return _build_data_array(wrfnc, varname, filetimeidx,
is_moving, True, _key)
else:
var = wrfnc.variables[varname]
if len(var.shape) > 1:
result = var[filetimeidx, :]
return result[np.newaxis, :] # So that nosqueeze works
else:
result = var[filetimeidx]
return result[np.newaxis] # So that nosqueeze works
else:
comboidx += numtimes
raise IndexError("timeidx {} is out of bounds".format(timeidx))
def _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key):
"""Find and return the array object within a sequence for a specific time
index.
The sequence is searched in reverse.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int`): The desired time index. Must be negative.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
try:
revwrfseq = reversed(wrfseq)
except TypeError:
revwrfseq = reversed(list(wrfseq))
wrf_iter = iter(revwrfseq)
revtimeidx = -timeidx - 1
comboidx = 0
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
numtimes = extract_dim(wrfnc, "Time")
if revtimeidx < comboidx + numtimes:
# Finds the "forward" sequence index, then counts that
# number back from the back of the ncfile times,
# since the ncfile needs to be iterated backwards as well.
filetimeidx = numtimes - (revtimeidx - comboidx) - 1
if meta:
return _build_data_array(wrfnc, varname, filetimeidx,
is_moving, True, _key)
else:
result = wrfnc.variables[varname][filetimeidx, :]
return result[np.newaxis, :] # So that nosqueeze works
else:
comboidx += numtimes
raise IndexError("timeidx {} is out of bounds".format(timeidx))
def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta, _key):
"""Find and return the array object within a sequence for a specific time
index.
The sequence is searched in forward or reverse based on the time index
chosen.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int`): The desired time index.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
if timeidx >= 0:
return _find_forward(wrfseq, varname, timeidx, is_moving, meta, _key)
else:
return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key)
def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key):
"""Return an array object from a sequence of files using the concatenate
method.
The concatenate method aggregates all files in the sequence along the
'Time' dimension, which will be the leftmost dimension. No sorting is
performed, so all files in the sequence must be sorted prior to calling
this method.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
if is_moving is None:
is_moving = is_moving_domain(wrfseq, varname, _key=_key)
file_times = extract_times(wrfseq, ALL_TIMES, meta=False, do_xtime=False)
multitime = is_multi_time_req(timeidx)
# For single times, just need to find the ncfile and appropriate
# time index, and return that array
if not multitime:
return _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta,
_key)
# If all times are requested, need to build a new array and cat together
# all of the arrays in the sequence
wrf_iter = iter(wrfseq)
if xarray_enabled() and meta:
first_var = _build_data_array(next(wrf_iter), varname,
ALL_TIMES, is_moving, True, _key)
else:
first_var = (next(wrf_iter)).variables[varname][:]
outdims = [len(file_times)]
# Making a new time dim, so ignore this one
outdims += first_var.shape[1:]
outdata = np.empty(outdims, first_var.dtype)
numtimes = first_var.shape[0]
startidx = 0
endidx = numtimes
if first_var.ndim > 1:
outdata[startidx:endidx, :] = first_var[:]
else:
outdata[startidx:endidx] = first_var[:]
if xarray_enabled() and meta:
latname, lonname, timename = _find_coord_names(first_var.coords)
timecached = False
latcached = False
loncached = False
outxtimes = None
outlats = None
outlons = None
timekey = timename+"_cat" if timename is not None else None
latkey = latname + "_cat" if latname is not None else None
lonkey = lonname + "_cat" if lonname is not None else None
if timename is not None:
outxtimes = get_cached_item(_key, timekey)
if outxtimes is None:
outxtimes = np.empty(outdims[0])
outxtimes[startidx:endidx] = to_np(
first_var.coords[timename][:])
else:
timecached = True
if is_moving:
outcoorddims = outdims[0:1] + outdims[-2:]
if latname is not None:
# Try to pull from the coord cache
outlats = get_cached_item(_key, latkey)
if outlats is None:
outlats = np.empty(outcoorddims, first_var.dtype)
outlats[startidx:endidx, :] = to_np(
first_var.coords[latname][:])
else:
latcached = True
if lonname is not None:
outlons = get_cached_item(_key, lonkey)
if outlons is None:
outlons = np.empty(outcoorddims, first_var.dtype)
outlons[startidx:endidx, :] = to_np(
first_var.coords[lonname][:])
else:
loncached = True
startidx = endidx
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
vardata = wrfnc.variables[varname][:]
numtimes = vardata.shape[0]
endidx = startidx + numtimes
if vardata.ndim > 1:
outdata[startidx:endidx, :] = vardata[:]
else:
outdata[startidx:endidx] = vardata[:]
if xarray_enabled() and meta:
if timename is not None and not timecached:
xtimedata = wrfnc.variables[timename][:]
outxtimes[startidx:endidx] = xtimedata[:]
if is_moving:
if latname is not None and not latcached:
latdata = wrfnc.variables[latname][:]
outlats[startidx:endidx, :] = latdata[:]
if lonname is not None and not loncached:
londata = wrfnc.variables[lonname][:]
outlons[startidx:endidx, :] = londata[:]
startidx = endidx
if xarray_enabled() and meta:
# Cache the coords if applicable
if not latcached and outlats is not None:
cache_item(_key, latkey, outlats)
if not loncached and outlons is not None:
cache_item(_key, lonkey, outlons)
if not timecached and outxtimes is not None:
cache_item(_key, timekey, outxtimes)
outname = first_var.name
outattrs = OrderedDict(first_var.attrs)
outcoords = OrderedDict(first_var.coords)
outdimnames = list(first_var.dims)
if "Time" not in outdimnames:
outdimnames.insert(0, "Time")
if not multitime:
file_times = [file_times[timeidx]]
outcoords[outdimnames[0]] = file_times
outcoords["datetime"] = outdimnames[0], file_times
if timename is not None:
outxtimes = outxtimes[:]
outcoords[timename] = outdimnames[0], outxtimes
# If the domain is moving, need to create the lat/lon coords
# since they can't be copied
if is_moving:
outlatdims = [outdimnames[0]] + outdimnames[-2:]
if latname is not None:
outlats = outlats[:]
outcoords[latname] = outlatdims, outlats
if lonname is not None:
outlons = outlons[:]
outcoords[lonname] = outlatdims, outlons
outdata = outdata[:]
outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdimnames, attrs=outattrs)
else:
outdata = outdata[:]
outarr = outdata
return outarr
def _get_numfiles(wrfseq):
"""Return the number of files in the sequence.
This function will first try to call the builtin :meth:`len` function, but
if that fails, the entire squence will be iterated over and counted.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
Returns:
:obj:`int`: The number of files in the sequence.
"""
try:
return len(wrfseq)
except TypeError:
wrf_iter = iter(wrfseq)
return sum(1 for _ in wrf_iter)
def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key):
"""Return an array object from a sequence of files using the join
method.
The join method creates a new leftmost dimension for the file/sequence
index. In situations where there are multiple files with multiple times,
and the last file contains less times than the previous files, the
remaining arrays will be arrays filled with missing values. There are
checks in place within the wrf-python algorithms to look for these missing
arrays, but be careful when calling compiled routines outside of
wrf-python.
In general, join is rarely used, so the concatenate method should be used
for most cases.
Args:
wrfseq (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
if is_moving is None:
is_moving = is_moving_domain(wrfseq, varname, _key=_key)
multitime = is_multi_time_req(timeidx)
numfiles = _get_numfiles(wrfseq)
maxtimes = _find_max_time_size(wrfseq)
time_idx_or_slice = timeidx if not multitime else slice(None)
file_times_less_than_max = False
file_idx = 0
# wrfseq might be a generator
wrf_iter = iter(wrfseq)
wrfnc = next(wrf_iter)
numtimes = extract_dim(wrfnc, "Time")
if xarray_enabled() and meta:
first_var = _build_data_array(wrfnc, varname, ALL_TIMES, is_moving,
True, _key)
time_coord = np.full((numfiles, maxtimes), np.datetime64("NaT"),
"datetime64[ns]")
time_coord[file_idx, 0:numtimes] = first_var.coords["Time"][:]
else:
first_var = wrfnc.variables[varname][:]
if numtimes < maxtimes:
file_times_less_than_max = True
# Out dimensions will be the number of files, maxtimes, then the
# non-time shapes from the first variable
outdims = [numfiles]
outdims += [maxtimes]
outdims += first_var.shape[1:]
# For join, always need to start with full masked values
outdata = np.full(outdims, default_fill(first_var.dtype), first_var.dtype)
if first_var.ndim > 1:
outdata[file_idx, 0:numtimes, :] = first_var[:]
else:
outdata[file_idx, 0:numtimes] = first_var[:]
# Create the secondary coordinate arrays
if xarray_enabled() and meta:
latname, lonname, timename = _find_coord_names(first_var.coords)
outcoorddims = outdims[0:2] + outdims[-2:]
timecached = False
latcached = False
loncached = False
outxtimes = None
outlats = None
outlons = None
timekey = timename+"_join" if timename is not None else None
latkey = latname + "_join" if latname is not None else None
lonkey = lonname + "_join" if lonname is not None else None
if timename is not None:
outxtimes = get_cached_item(_key, timekey)
if outxtimes is None:
outxtimes = np.full(outdims[0:2],
default_fill(first_var.dtype),
first_var.dtype)
outxtimes[file_idx, 0:numtimes] = first_var.coords[timename][:]
else:
timecached = True
if is_moving:
if latname is not None:
outlats = get_cached_item(_key, latkey)
if outlats is None:
outlats = np.full(outcoorddims,
default_fill(first_var.dtype),
first_var.dtype)
outlats[file_idx, 0:numtimes, :] = (
first_var.coords[latname][:])
else:
latcached = True
if lonname is not None:
outlons = get_cached_item(_key, lonkey)
if outlons is None:
outlons = np.full(outcoorddims,
default_fill(first_var.dtype),
first_var.dtype)
outlons[file_idx, 0:numtimes, :] = (
first_var.coords[lonname][:])
else:
loncached = True
file_idx = 1
while True:
try:
wrfnc = next(wrf_iter)
except StopIteration:
break
else:
numtimes = extract_dim(wrfnc, "Time")
if numtimes < maxtimes:
file_times_less_than_max = True
outvar = wrfnc.variables[varname][:]
if not multitime:
outvar = outvar[np.newaxis, :]
if outvar.ndim > 1:
outdata[file_idx, 0:numtimes, :] = outvar[:]
else:
outdata[file_idx, 0:numtimes] = outvar[:]
if xarray_enabled() and meta:
# For join, the times are a function of fileidx
file_times = extract_times(wrfnc, ALL_TIMES, meta=False,
do_xtime=False)
time_coord[file_idx, 0:numtimes] = np.asarray(
file_times, "datetime64[ns]")[:]
if timename is not None and not timecached:
xtimedata = wrfnc.variables[timename][:]
outxtimes[file_idx, 0:numtimes] = xtimedata[:]
if is_moving:
if latname is not None and not latcached:
latdata = wrfnc.variables[latname][:]
outlats[file_idx, 0:numtimes, :] = latdata[:]
if lonname is not None and not loncached:
londata = wrfnc.variables[lonname][:]
outlons[file_idx, 0:numtimes, :] = londata[:]
# Need to update coords here
file_idx += 1
# If any of the output files contain less than the max number of times,
# then a mask array is needed to flag all the missing arrays with
# missing values
if file_times_less_than_max:
outdata = np.ma.masked_values(outdata, default_fill(outdata.dtype))
if xarray_enabled() and meta:
# Cache the coords if applicable
if not latcached and outlats is not None:
cache_item(_key, latkey, outlats)
if not loncached and outlons is not None:
cache_item(_key, lonkey, outlons)
if not timecached and outxtimes is not None:
cache_item(_key, timekey, outxtimes)
outname = first_var.name
outcoords = OrderedDict(first_var.coords)
outattrs = OrderedDict(first_var.attrs)
# New dimensions
outdimnames = ["file"] + list(first_var.dims)
outcoords["file"] = [i for i in py3range(numfiles)]
# Time needs to be multi dimensional, so use the default dimension
del outcoords["Time"]
time_coord = time_coord[:, time_idx_or_slice]
if not multitime:
time_coord = time_coord[:, np.newaxis]
outcoords["datetime"] = outdimnames[0:2], time_coord
if isinstance(outdata, np.ma.MaskedArray):
outattrs["_FillValue"] = default_fill(outdata.dtype)
outattrs["missing_value"] = default_fill(outdata.dtype)
if timename is not None:
outxtimes = outxtimes[:, time_idx_or_slice]
if not multitime:
outxtimes = outxtimes[:, np.newaxis]
outcoords[timename] = outdimnames[0:2], outxtimes[:]
# If the domain is moving, need to create the lat/lon coords
# since they can't be copied
if is_moving:
outlatdims = outdimnames[0:2] + outdimnames[-2:]
if latname is not None:
outlats = outlats[:, time_idx_or_slice, :]
if not multitime:
outlats = outlats[:, np.newaxis, :]
outcoords[latname] = outlatdims, outlats
if lonname is not None:
outlons = outlons[:, time_idx_or_slice, :]
if not multitime:
outlons = outlons[:, np.newaxis, :]
outcoords[lonname] = outlatdims, outlons
if not multitime:
outdata = outdata[:, timeidx, :]
outdata = outdata[:, np.newaxis, :]
outarr = DataArray(outdata, name=outname, coords=outcoords,
dims=outdimnames, attrs=outattrs)
else:
if not multitime:
outdata = outdata[:, timeidx, :]
outdata = outdata[:, np.newaxis, :]
outarr = outdata
return outarr
def combine_files(wrfin, varname, timeidx, is_moving=None,
method="cat", squeeze=True, meta=True,
_key=None):
"""Combine and return an array object for the sequence of WRF output
files.
Two aggregation methodologies are available to combine the sequence:
- 'cat': Concatenate the files along the 'Time' dimension. The Time
dimension will be the leftmost dimension. No sorting is performed,
so files must be properly ordered in the sequence prior to calling
this function.
- 'join': Join the files by creating a new leftmost dimension for the
file index. In situations where there are multiple files with
multiple times, and the last file contains less times than the
previous files, the remaining arrays will be arrays filled with
missing values. There are checks in place within the wrf-python
algorithms to look for these missing arrays, but be careful when
calling compiled routines outside of wrf-python.
Args:
wrfin (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
# Handles generators, single files, lists, tuples, custom classes
wrfseq = get_iterable(wrfin)
# Dictionary is unique
if is_mapping(wrfseq):
outarr = _combine_dict(wrfseq, varname, timeidx, method, meta, _key)
elif method.lower() == "cat":
outarr = _cat_files(wrfseq, varname, timeidx, is_moving,
squeeze, meta, _key)
elif method.lower() == "join":
outarr = _join_files(wrfseq, varname, timeidx, is_moving, meta, _key)
else:
raise ValueError("method must be 'cat' or 'join'")
return outarr.squeeze() if squeeze else outarr
def _extract_var(wrfin, varname, timeidx, is_moving,
method, squeeze, cache, meta, _key):
"""Extract a variable from a NetCDF file object or a sequence of NetCDF
file objects.
Args:
wrfin (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varname (:obj:`str`) : The variable name.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
is_moving (:obj:`bool`): A boolean type that indicates if the
sequence is a moving nest.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: If xarray is
enabled and the *meta* parameter is True, then the result will be a
:class:`xarray.DataArray` object. Otherwise, the result will be a
:class:`numpy.ndarray` object with no metadata.
"""
if cache is not None:
try:
cache_var = cache[varname]
except KeyError:
pass
else:
if not meta:
return to_np(cache_var)
return cache_var
multitime = is_multi_time_req(timeidx)
multifile = is_multi_file(wrfin)
if is_time_coord_var(varname):
return extract_times(wrfin, timeidx, method, squeeze, cache,
meta, do_xtime=True)
if not multifile:
if xarray_enabled() and meta:
if is_moving is None:
is_moving = is_moving_domain(wrfin, varname, _key=_key)
result = _build_data_array(wrfin, varname, timeidx, is_moving,
multifile, _key)
else:
if not multitime:
result = wrfin.variables[varname][timeidx, :]
result = result[np.newaxis, :] # So that no squeeze works
else:
result = wrfin.variables[varname][:]
else:
# Squeeze handled in this routine, so just return it
return combine_files(wrfin, varname, timeidx, is_moving,
method, squeeze, meta, _key)
return result.squeeze() if squeeze else result
def extract_vars(wrfin, timeidx, varnames, method="cat", squeeze=True,
cache=None, meta=True, _key=None):
"""Extract variables from a NetCDF file object or a sequence of NetCDF
file objects.
Args:
wrfin (iterable): An iterable type, which includes lists, tuples,
dictionaries, generators, and user-defined classes.
varnames (sequence of :obj:`str`) : A sequence of variable names.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. The default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
meta (:obj:`bool`, optional): Set to False to disable metadata and
return :class:`numpy.ndarray` instead of
:class:`xarray.DataArray`. Default is True.
_key (:obj:`int`, optional): Cache key for the coordinate variables.
This is used for internal purposes only. Default is None.
Returns:
:obj:`dict`: A mapping of variable name to an array object. If xarray
is enabled and the *meta* parameter is True, then the array object will
be a :class:`xarray.DataArray` object. Otherwise, the array object
will be a :class:`numpy.ndarray` object with no metadata.
"""
if isstr(varnames):
varlist = [varnames]
else:
varlist = varnames
return {var: _extract_var(wrfin, var, timeidx, None,
method, squeeze, cache, meta, _key)
for var in varlist}
def npbytes_to_str(var):
"""Return a :obj:`bytes` object for the raw character array.
Args:
var (:class:`numpy.ndarray`): An array of characters.
Returns:
:obj:`bytes`: A string of bytes.
"""
return (bytes(c).decode("utf-8") for c in var[:])
def _make_time(timearr):
"""Return a :class:`datetime.datetime` object for the array of characters.
Args:
timearr (:class:`numpy.ndarray`): An array of characters.
Returns:
:class:`datetime.datetime`: A datetime object.
"""
try:
return dt.datetime.strptime("".join(npbytes_to_str(timearr)),
"%Y-%m-%d_%H:%M:%S")
except ValueError:
return np.datetime64("NaT")
def _file_times(wrfin, do_xtime):
"""Yield a time object for the times found in a sequence of files.
If *do_xtime* to True, a :class:`datetime.datetime` object is yielded.
Otherwise, a :obj:`float` object is yielded.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
do_xtime (:obj:`bool`): Set to True to parse the 'XTIME' variable
instead of the 'Times' variable.
Yields:
:class:`datetime.datetime` or :obj:`float`: A
:class:`datetime.datetime` object if *do_xtime* is False,
otherwise a :obj:`float`.
"""
if not do_xtime:
times = wrfin.variables["Times"][:, :]
for i in py3range(times.shape[0]):
yield _make_time(times[i, :])
else:
xtimes = wrfin.variables["XTIME"][:]
for i in py3range(xtimes.shape[0]):
yield xtimes[i]
def _extract_time_map(wrfin, timeidx, do_xtime, meta=False):
"""Return a mapping of key to a sequence of time objects.
This function is used when *wrfin* is a mapping.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence.
do_xtime (:obj:`bool`): Set to True to parse the 'XTIME' variable
instead of the 'Times' variable.
meta (:obj:`bool`, optional): Set to False to disable metadata.
Returns:
:obj:`dict`: A mapping of key to a sequence of time objects. If
*meta* is True, the sequence will be of type :class:`xarray.DataArray`,
otherwise the sequence is :class:`numpy.ndarray`.
"""
return {key: extract_times(wrfseq, timeidx, do_xtime, meta)
for key, wrfseq in viewitems(wrfin)}
def extract_times(wrfin, timeidx, method="cat", squeeze=True, cache=None,
meta=False, do_xtime=False):
"""Return a sequence of time objects.
If *do_xtime* is False, the 'XTIME' variable is used and each time object
is a :obj:`float`. Otherwise, the 'Times' variable is used, and each
time object is a :class:`datetime.datetime` object.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
meta (:obj:`bool`, optional): Set to False to disable metadata.
do_xtime (:obj:`bool`): Set to True to parse the 'XTIME' variable
instead of the 'Times' variable. Default is False.
Returns:
:class:`xarray.DataArray` or :class:`numpy.ndarray`: A sequence of time
objects. If *meta* is True, the sequence will be of type
:class:`xarray.DataArray`, otherwise the sequence is
:class:`numpy.ndarray`.
"""
if is_mapping(wrfin):
return _extract_time_map(wrfin, timeidx, do_xtime)
multitime = is_multi_time_req(timeidx)
multi_file = is_multi_file(wrfin)
if not multi_file:
wrf_list = [wrfin]
else:
wrf_list = wrfin
dt = "datetime64[ns]" if not do_xtime else np.float64
fill_value = (np.datetime64('NaT') if not do_xtime else
default_fill(np.float64))
try:
if method.lower() == "cat":
time_list = [file_time
for wrf_file in wrf_list
for file_time in _file_times(wrf_file, do_xtime)]
time_arr = np.asarray(time_list, dtype=dt)
elif method.lower() == "join":
time_list = [[file_time
for file_time in _file_times(wrf_file, do_xtime)]
for wrf_file in wrf_list]
num_rows = len(time_list)
num_cols = len(time_list[0])
time_arr = np.full((num_rows, num_cols), fill_value, dtype=dt)
for i, row in enumerate(time_list):
if len(row) == num_cols:
time_arr[i, :] = row[:]
else:
for j, val in enumerate(row):
time_arr[i, j] = val
time_arr = ma.masked_values(time_arr, fill_value)
else:
raise ValueError("invalid method argument '{}'".format(method))
except KeyError:
return None # Thrown for pre-3.7 XTIME not existing
if xarray_enabled() and meta:
outattrs = OrderedDict()
outcoords = None
if method.lower() == "cat":
outdimnames = ["Time"]
else:
outdimnames = ["fileidx", "Time"]
outattrs["missing_value"] = fill_value
outattrs["_FillValue"] = fill_value
if not do_xtime:
outname = "times"
outattrs["description"] = "model times [np.datetime64]"
else:
ncfile = next(iter(wrf_list))
var = ncfile.variables["XTIME"]
outattrs.update(var.__dict__)
outname = "XTIME"
outarr = DataArray(time_arr, name=outname, coords=outcoords,
dims=outdimnames, attrs=outattrs)
else:
outarr = time_arr
if not multitime:
return outarr[timeidx]
return outarr
def is_standard_wrf_var(wrfin, varname):
"""Return True if the variable is a standard WRF variable and not a
diagnostic.
If *wrfin* is a sequence, only the first file is used.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
varname (:obj:`str`): The variable name.
Returns:
:obj:`bool`: True if the variable is a standard WRF variable,
otherwise False.
"""
multifile = is_multi_file(wrfin)
if multifile:
if not is_mapping(wrfin):
wrfin = next(iter(wrfin))
else:
entry = wrfin[next(iter(viewkeys(wrfin)))]
return is_standard_wrf_var(entry, varname)
return varname in wrfin.variables
def is_staggered(wrfin, var):
"""Return True if the variable is on a staggered grid.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
var (array): An array object which contains a :attr:`shape`
attribute.
Returns:
:obj:`bool`: True if the variable is on a staggered grid, otherwise
False.
"""
we = extract_dim(wrfin, "west_east")
sn = extract_dim(wrfin, "south_north")
bt = extract_dim(wrfin, "bottom_top")
if (var.shape[-1] != we or var.shape[-2] != sn or var.shape[-3] != bt):
return True
return False
def get_left_indexes(var, expected_dims):
"""Returns a tuple for the extra leftmost dimension sizes.
For example, if an algorithm expects a 3 dimensional variable, but the
variable includes an additional left dimension for Time, and
this Time dimension has 3 values, then this function will return (3,).
Args:
var (array): An array object that contains the :attr:`ndim`
and :attr:`shape` attributes.
expected_dims (:obj:`int`): The expected number of dimensions (usually
for a computational algorithm).
Returns:
:obj:`tuple`: The shape for the extra leftmost dimensions.
"""
extra_dim_num = var.ndim - expected_dims
if (extra_dim_num == 0):
return []
return var.shape[0:extra_dim_num]
def iter_left_indexes(dims):
"""Yield the iteration tuples for a sequence of dimensions sizes.
For example, if *dims* is (2,2), then this will yield:
(0,0), (0,1), (1,0), (1,1)
This is primarily used to iterate over the leftmost index values.
Args:
dims (indexable sequence): A sequence of dimension sizes.
Yields:
:obj:`tuple`: The leftmost indexing iteration sizes.
"""
arg = [py3range(dim) for dim in dims]
for idxs in product(*arg):
yield idxs
def get_right_slices(var, right_ndims, fixed_val=0):
"""Return an indexing tuple where the left dimensions are held to a
fixed value and the right dimensions are set to slice objects.
For example, if *var* is a 5D variable, and the desired indexing sequence
for a numpy array is (0,0,0,:,:), then *right_ndims* should be set to 2
and *fixed_val* set to 0.
Args:
var (:class:`numpy.ndarray`): A numpy array.
right_ndims (:obj:`int`): The number of right dimensions to be sliced.
fixed_val (:obj:`int`): The value to hold the left dimensions to.
Returns:
:obj:`tuple`: An indexing tuple that can be used to index a
:class:`numpy.ndarray`.
"""
extra_dim_num = var.ndim - right_ndims
if extra_dim_num == 0:
return [slice(None)] * right_ndims
return tuple([fixed_val]*extra_dim_num +
[slice(None)]*right_ndims)
def get_proj_params(wrfin):
"""Return a tuple of latitude, longitude, and projection parameters from
a WRF output file object or a sequence of WRF output file objects.
Args:
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
varname (:obj:`str`, optional): The variable name to extract the
coordinate variable names from. Default is None, which will
use the default coordinate variable names ('XLAT', 'XLONG').
Returns:
:obj:`tuple`: A tuple of the latitude coordinate variable,
longitude coordinate, and global projection attributes.
"""
proj_params = extract_global_attrs(wrfin,
attrs=("MAP_PROJ",
"CEN_LAT", "CEN_LON",
"TRUELAT1", "TRUELAT2",
"MOAD_CEN_LAT", "STAND_LON",
"POLE_LAT", "POLE_LON",
"DX", "DY"))
return proj_params
def from_args(func, argnames, *args, **kwargs):
"""Return a mapping of argument name to value for the called function.
This function parses the function args and kwargs to obtain the
desired argument value. If the argument has not been passed in, the value
is taken from the default keyword argument value.
This func is usually called from within a decorator.
Note:
This function currently does not work with functions that contain
variable length args or kwargs arguments.
Args:
func (function): The function to examine (usually the function that is
wrapped).
argnames (iterable): An iterable sequence of argument names.
*args: The positional arguments.
**kwargs: The keyword arguments.
Returns:
:obj:`dict`: A mapping of argument name to argument value.
"""
if isstr(argnames):
arglist = [argnames]
else:
arglist = argnames
result = OrderedDict()
for argname in arglist:
arg_loc = arg_location(func, argname, args, kwargs)
if arg_loc is not None:
result[argname] = arg_loc[0][arg_loc[1]]
else:
result[argname] = None
return result
def _args_to_list2(func, args, kwargs):
"""Return all of the function arguments, including defaults, as a list.
The result can then be passed to the function via *result. This version
uses :meth:`inspect.argspec`, so is only applicable for Python 2.7.
Note:
This function currently does not work with functions that contain
variable length args or kwargs arguments.
Args:
func (function): The function to examine (usually the function
that is wrapped).
args (:obj:`tuple`): The positional arguments.
kwargs (:obj:`dict`): The keyword arguments.
Returns:
:obj:`list`: A list of all argument values, including defaults.
"""
argspec = getargspec(func)
# Build the full tuple with defaults filled in
outargs = [None]*len(argspec.args)
if argspec.defaults is not None:
for i, default in enumerate(argspec.defaults[::-1], 1):
outargs[-i] = default
# Add the supplied args
for i, arg in enumerate(args):
outargs[i] = arg
# Fill in the supplied kargs
for argname, val in viewitems(kwargs):
argidx = argspec.args.index(argname)
outargs[argidx] = val
return outargs
# For Python 3.4, this will run the BoundArguments.apply_defaults method that
# was introduced in Python 3.5.
def _apply_defaults(bound):
"""Set default values for missing arguments.
For variable-positional arguments (*args) the default is an
empty tuple.
For variable-keyword arguments (**kwargs) the default is an
empty dict.
Args:
bound (:class:`inspect.BoundArguments`): A BoundArguments object.
Returns:
None
"""
arguments = bound.arguments
new_arguments = []
for name, param in bound._signature.parameters.items():
try:
new_arguments.append((name, arguments[name]))
except KeyError:
if param.default is not _empty:
val = param.default
elif param.kind is _VAR_POSITIONAL:
val = ()
elif param.kind is _VAR_KEYWORD:
val = {}
else:
# This BoundArguments was likely produced by
# Signature.bind_partial().
continue
new_arguments.append((name, val))
bound.arguments = OrderedDict(new_arguments)
def _args_to_list3(func, args, kwargs):
"""Return all of the function arguments, including defaults, as a list.
The result can then be passed to the function via *result. This version
uses :meth:`inspect.signature`, so is only applicable for Python 3.4+.
Note:
This function currently does not work with functions that contain
variable length args or kwargs arguments.
Args:
func (function): The function to examine (usually the function
that is wrapped).
args (:obj:`tuple`): The positional arguments.
kwargs (:obj:`dict`): The keyword arguments.
Returns:
:obj:`list`: A list of all argument values, including defaults.
"""
sig = signature(func)
bound = sig.bind(*args, **kwargs)
try:
bound.apply_defaults()
except AttributeError:
_apply_defaults(bound)
return [x for x in bound.arguments.values()]
# Note: Doesn't allow for **kwargs or *args
def args_to_list(func, args, kwargs):
"""Return all of the function arguments, including defaults, as a list.
The result can then be passed to the function via *result*.
Note:
This function currently does not work with functions that contain
variable length args or kwargs arguments.
Args:
func (function): The function to examine (usually the function
that is wrapped).
args (:obj:`tuple`): The positional arguments.
kwargs (:obj:`dict`): The keyword arguments.
Returns:
:obj:`list`: A list of all argument values, including defaults.
"""
if version_info > (3,):
_args_to_list = _args_to_list3
else:
_args_to_list = _args_to_list2
return _args_to_list(func, args, kwargs)
def _arg_location2(func, argname, args, kwargs):
"""Return the function arguments as a single list along with the
index within that list for a specified argument name.
This function parses the args, kargs and signature looking for the
location of *argname*, and returns a list containing all arguments, along
with the argument location in that list.
This function requires :meth:`inspect.getargspec`, so it is only
applicable for Python 2.7.
Args:
func (function): The function to examine (usually the function
that is wrapped).
argname (:obj:`str`): The argument name to locate.
args (:obj:`tuple`): The positional arguments.
kwargs (:obj:`dict`): The keyword arguments.
Returns:
:obj:`tuple`: A tuple containing the list of all argument values along
with the index for location of *argname*.
"""
argspec = getargspec(func)
list_args = _args_to_list2(func, args, kwargs)
# Return the new sequence and location
if argname not in argspec.args and argname not in kwargs:
return None
result_idx = argspec.args.index(argname)
return list_args, result_idx
def _arg_location3(func, argname, args, kwargs):
"""Return the function arguments as a single list along with the
index within that list for a specified argument name.
This function parses the args, kargs and signature looking for the
location of *argname*, and returns a list containing all arguments, along
with the argument location in that list.
This function requires :meth:`inspect.signature`, so it is only
applicable for Python 3.4 and higher.
Args:
func (function): The function to examine (usually the function
that is wrapped).
argname (:obj:`str`): The argument name to locate.
args (:obj:`tuple`): The positional arguments.
kwargs (:obj:`dict`): The keyword arguments.
Returns:
:obj:`tuple`: A tuple containing the list of all argument values along
with the index for location of *argname*.
"""
sig = signature(func)
params = list(sig.parameters.keys())
list_args = _args_to_list3(func, args, kwargs)
try:
result_idx = params.index(argname)
except ValueError:
return None
return list_args, result_idx
def arg_location(func, argname, args, kwargs):
"""Return the function arguments as a single list along with the
index within that list for a specified argument name.
This function parses the args, kargs and signature looking for the
location of *argname*, and returns a list containing all arguments, along
with the argument location in that list.
Args:
func (function): The function to examine (usually the function
that is wrapped).
argname (:obj:`str`): The argument name to locate.
args (:obj:`tuple`): The positional arguments.
kwargs (:obj:`dict`): The keyword arguments.
Returns:
:obj:`tuple`: A tuple containing the list of all argument values along
with the index for location of *argname*.
"""
if version_info > (3,):
_arg_location = _arg_location3
else:
_arg_location = _arg_location2
return _arg_location(func, argname, args, kwargs)
def psafilepath():
"""Return the full path to the 'psadilookup.dat' file.
The 'psadilookup.dat' file contains the lookup table for the cape
routines.
Returns:
:obj:`str`: The full path to the 'psadilookup.dat' file.
"""
return os.path.join(os.path.dirname(__file__), "data", "psadilookup.dat")
def get_filepath(obj):
"""Return the file path for the specified object.
This is used to return the file path for a netcdf object. If the
particular object does not have the appropriate file path information,
then one is created based on the timestep in the file.
Args:
obj: An object.
Returns:
:obj:`str`: A string for a file path.
"""
try:
path = obj.filepath()
except AttributeError:
try:
path = obj.file.path
except AttributeError:
# Let's make up a filename from the first file time
found = False
times = extract_times(obj, None, meta=False, do_xtime=False)
for t in times:
path = "wrfout_{}".format(str(t))
found = True
break
if not found:
raise ValueError("file contains no path information")
return path
def get_id(obj, prefix=''):
"""Return the cache id.
The cache id is used as a caching key for various routines. If the
object type is a mapping, then the result will also be a
mapping of each key to the object id for the value.
Args:
obj (:obj:`object`): Any object type.
prefix (:obj:`str`): A string to help with recursive calls.
Returns:
:obj:`int` or :obj:`dict`: If the *obj* parameter is not a mapping,
then the object id is returned. Otherwise, a mapping of each
key to the object id for the value is returned.
"""
if not is_multi_file(obj):
return hash(prefix + get_filepath(obj))
# For sequences, the hashing string will be the list ID and the
# path for the first file in the sequence
if not is_mapping(obj):
_obj = get_iterable(obj)
_next = next(iter(_obj))
return get_id(_next, prefix + str(id(obj)))
# For each key in the mapping, recursively call get_id until
# until a non-mapping is found
return {key: get_id(val, prefix) for key, val in viewitems(obj)}
def geo_bounds(var=None, wrfin=None, varname=None, timeidx=0, method="cat",
squeeze=True, cache=None):
"""Return the geographic boundaries for the variable or file(s).
When using a :class:`xarray.DataArray` as the *var* parameter, the variable
must contain latitude and longitude coordinates. If these coordinate
dimensions are greater than two dimensions, then an array of
:class:`wrf.GeoBounds` objects will be returned with the same shape as the
leftmost dimensions of the coordinate arrays.
When using a WRF file, or sequence of WRF files, by supplying the
*wrfin* parameter, an array of :class:`wrf.GeoBounds` objects will be
returned if the domain is moving and :data:`wrf.ALL_TIMES` is selected as
the *timeidx* parameter when using *wrfin*. Otherwise, a single
:class:`wrf.GeoBounds` object is returned.
Args:
var (:class:`xarray.DataArray`, optional): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`, optional): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index when *wrfin* is not None. This value can be a
positive integer, negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0. This value is
ignored when *var* is used.
method (:obj:`str`, optional): The aggregation method to use for
sequences when *wrfin* is not None. Must be either 'cat' or
'join'. 'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Only used when *wrfin* is used. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files. Only used when *wrfin* is used.
Default is None.
Returns:
:class:`wrf.GeoBounds`: The domain geographic bounds.
"""
if var is None and wrfin is None:
raise ValueError("'var' or 'wrfin' parameter is required")
# Getting lat/lon from xarray coordinates
if var is not None:
if not xarray_enabled():
raise ValueError("xarray is not installed or is disabled")
is_moving = None
try:
var_coords = var.coords
except AttributeError:
raise ValueError("'var' object does not contain coordinate "
"attributes")
latname, lonname, _ = _find_coord_names(var_coords)
try:
lats = to_np(var_coords[latname])
except KeyError:
raise ValueError("'var' object does not contain a latitude "
"coordinate")
try:
lons = to_np(var_coords[lonname])
except KeyError:
raise ValueError("'var' object does not contain a longitude "
"coordinate")
# Getting lat/lon from the file
elif wrfin is not None:
_key = get_id(wrfin)
is_moving = is_moving_domain(wrfin, varname=varname,
latvar=either("XLAT", "XLAT_M"),
lonvar=either("XLONG", "XLONG_M"),
_key=_key)
if varname is not None:
if xarray_enabled():
var = extract_vars(wrfin, timeidx, varname, method, squeeze,
cache, meta=True, _key=_key)[varname]
return geo_bounds(var)
else:
lat_coord, lon_coord, _ = _get_coord_names(wrfin, varname)
else:
lat_coord = either("XLAT", "XLAT_M")(wrfin)
lon_coord = either("XLONG", "XLONG_M")(wrfin)
# If requesting all times but the domain isn't moving, just
# extract one time
_timeidx = timeidx
if timeidx is None and not is_moving:
_timeidx = 0
coord_data = extract_vars(wrfin, _timeidx, (lat_coord, lon_coord),
method, squeeze, cache, meta=False,
_key=_key)
lats = coord_data[lat_coord]
lons = coord_data[lon_coord]
# Moving domains
if lats.ndim > 2:
# Requesting all times, works for 'cat' and 'join' data
# and always for xarray data
extra_dims = lats.shape[0:-2]
out_geo = np.full(extra_dims, NullGeoBounds(), np.object)
for left_idxs in iter_left_indexes(extra_dims):
latlon_idx = left_idxs + (slice(None),)
out_geo[left_idxs] = GeoBounds(lats=lats[latlon_idx],
lons=lons[latlon_idx])
return out_geo
# Non-moving domains
return GeoBounds(lats=lats, lons=lons)
def _get_wrf_proj_geobnds(var, wrfin, varname, timeidx, method, squeeze,
cache):
"""Return the :class:`wrf.WrfProj` subclass and :class:`wrf.GeoBounds`.
Args:
var (:class:`xarray.DataArray`): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
Returns:
:obj:`tuple`: A tuple of :class:`wrf.WrfProj`
and :class:`wrf.GeoBounds`
"""
# Using a variable
if var is not None:
if not xarray_enabled():
raise ValueError("xarray is not installed or is disabled")
geobnds = geo_bounds(var)
try:
wrf_proj = var.attrs["projection"]
except AttributeError:
raise ValueError("variable does not contain projection "
"information")
else:
geobnds = geo_bounds(wrfin=wrfin, varname=varname, timeidx=timeidx,
method=method, cache=cache)
proj_params = get_proj_params(wrfin)
wrf_proj = getproj(**proj_params)
return wrf_proj, geobnds
def _get_proj_obj(ob_type, var, wrfin, varname, timeidx, method, squeeze,
cache, **kwargs):
"""Return the desired mapping object for the plotting type.
Args:
ob_type (:obj:`str`): Must be 'cartopy', 'basemap', or 'pyngl'.
var (:class:`xarray.DataArray`): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
**kwargs: Additional keyword arguments for the plotting type.
Returns:
mapping object: Will be either :class:`cartopy.crs.Projection`,
:class:`matplotlib.mpl_toolkits.basemap.Basemap` or
:class:`Ngl.Resources`.
"""
wrf_proj, geobnds = _get_wrf_proj_geobnds(var, wrfin, varname, timeidx,
method, squeeze, cache)
if ob_type == "cartopy":
proj_obj = wrf_proj.cartopy()
elif ob_type == "basemap":
try:
_ = len(geobnds)
except TypeError: # Only a single object
proj_obj = wrf_proj.basemap(geobnds, **kwargs)
else:
proj_obj = np.empty(geobnds.shape, np.object)
for idxs, geobnd_val in np.ndenumerate(geobnds):
proj_obj[idxs] = wrf_proj.basemap(geobnd_val, **kwargs)
elif ob_type == "pyngl":
try:
_ = len(geobnds)
except TypeError: # Only a single object
proj_obj = wrf_proj.pyngl(geobnds, **kwargs)
else:
proj_obj = np.empty(geobnds.shape, np.object)
for idxs, geobnd_val in np.ndenumerate(geobnds):
proj_obj[idxs] = wrf_proj.pyngl(geobnd_val, **kwargs)
return proj_obj
def latlon_coords(var, as_np=False):
"""Return the latitude and longitude coordinates from a
:class:`xarray.DataArray` object.
Args:
var (:class:`xarray.DataArray`): A variable.
as_np (:obj:`bool`): Set to True to return the coordinates as
:class:`numpy.ndarray` objects instead of
:class:`xarray.DataArray` objects.
Returns:
:obj:`tuple`: The latitude and longitude coordinate variables.
"""
if not xarray_enabled():
raise ValueError("xarray is not installed or is disabled")
try:
var_coords = var.coords
except AttributeError:
raise ValueError("'var' object does not contain coordinate "
"attributes")
latname, lonname, _ = _find_coord_names(var_coords)
try:
lats = var_coords[latname]
except KeyError:
raise ValueError("'var' object does not contain a latitude "
"coordinate")
try:
lons = var_coords[lonname]
except KeyError:
raise ValueError("'var' object does not contain a longitude "
"coordinate")
if as_np:
return to_np(lats), to_np(lons)
return lats, lons
def get_cartopy(var=None, wrfin=None, varname=None, timeidx=0, method="cat",
squeeze=True, cache=None):
"""Return a :class:`cartopy.crs.Projection` subclass for the
map projection.
Args:
var (:class:`xarray.DataArray`, optional): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
geobounds (:class:`wrf.GeoBounds`, optional): The geobounds to
get the extents. If set to None and using the *var* parameter,
the geobounds will be taken from the variable. If using a
file, then the geobounds will be taken from the native grid.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`, optional): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
Returns:
:class:`cartopy.crs.Projection`: A Projection subclass for the
map projection.
See Also:
:class:`cartopy.crs.Projection`
"""
return _get_proj_obj("cartopy", var, wrfin, varname, timeidx, method,
squeeze, cache)
def get_basemap(var=None, wrfin=None, varname=None, timeidx=0, method="cat",
squeeze=True, cache=None, **kwargs):
"""Return a :class:`matplotlib.mpl_toolkits.basemap.Basemap` object
for the map projection.
Args:
var (:class:`xarray.DataArray`, optional): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
geobounds (:class:`wrf.GeoBounds`, optional): The geobounds to
get the extents. If set to None and using the *var* parameter,
the geobounds will be taken from the variable. If using a
file, then the geobounds will be taken from the native grid.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`, optional): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
**kwargs: Keyword arguments for creating a
:class:`matplotlib.mpl_toolkits.basemap.Basemap`. By default,
the domain bounds will be set to the native projection, the
resolution will be set to 'l', and the other projection
parameters will be set by the information in the file.
Returns:
:class:`cartopy.crs.Projection`: A Projection subclass for the
map projection.
Returns:
:class:`matplotlib.mpl_toolkits.basemap.Basemap`: A Basemap
object for the projection.
See Also:
:class:`matplotlib.mpl_toolkits.basemap.Basemap`
"""
return _get_proj_obj("basemap", var, wrfin, varname, timeidx, method,
squeeze, cache, **kwargs)
def get_pyngl(var=None, wrfin=None, varname=None, timeidx=0, method="cat",
squeeze=True, cache=None, **kwargs):
"""Return a :class:`Ngl.Resources` object for the map projection.
Args:
var (:class:`xarray.DataArray`, optional): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
geobounds (:class:`wrf.GeoBounds`, optional): The geobounds to
get the extents. If set to None and using the *var* parameter,
the geobounds will be taken from the variable. If using a
file, then the geobounds will be taken from the native grid.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`, optional): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
**kwargs: Additional PyNGL resources to set while creating the
:class:`Ngl.Resources` object.
Returns:
:class:`Ngl.Resources`: A dict-like object that contains the
PyNGL resources for the map projection.
See Also:
`PyNGL <https://www.pyngl.ucar.edu/>`_
"""
return _get_proj_obj("pyngl", var, wrfin, varname, timeidx, method,
squeeze, cache)
def cartopy_xlim(var=None, geobounds=None, wrfin=None, varname=None, timeidx=0,
method="cat", squeeze=True, cache=None):
"""Return the x-axis limits in the projected coordinates.
For some map projections, like :class`wrf.RotatedLatLon`, the
:meth:`cartopy.GeoAxes.set_extent` method does not work correctly. This
method is equivalent to:
.. code-block:: python
pc = crs.PlateCarree()
xs, ys, _ = self._cartopy().transform_points(pc,
np.array([geobounds.bottom_left.lon,
geobounds.top_right.lon]),
np.array([geobounds.bottom_left.lat,
geobounds.top_right.lat])).T
_xlimits = xs.tolist()
_ylimits = ys.tolist()
return (_xlimits, _ylimits)[0]
Args:
var (:class:`xarray.DataArray`, optional): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
geobounds (:class:`wrf.GeoBounds`, optional): The geobounds to
get the extents. If set to None and using the *var* parameter,
the geobounds will be taken from the variable. If using a
file, then the geobounds will be taken from the native grid.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`, optional): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
Returns:
:obj:`list`: A list of [start_x, end_x] in the projected coordinate
system.
"""
wrf_proj, native_geobnds = _get_wrf_proj_geobnds(var, wrfin, varname,
timeidx, method, squeeze,
cache)
if geobounds is not None:
return wrf_proj.cartopy_xlim(geobounds)
return wrf_proj.cartopy_xlim(native_geobnds)
def cartopy_ylim(var=None, geobounds=None, wrfin=None, varname=None, timeidx=0,
method="cat", squeeze=True, cache=None):
"""Return the y-axis limits in the projected coordinates.
For some map projections, like :class`wrf.RotatedLatLon`, the
:meth:`cartopy.GeoAxes.set_extent` method does not work correctly. This
method is equivalent to:
.. code-block:: python
pc = crs.PlateCarree()
xs, ys, _ = self._cartopy().transform_points(pc,
np.array([geobounds.bottom_left.lon,
geobounds.top_right.lon]),
np.array([geobounds.bottom_left.lat,
geobounds.top_right.lat])).T
_xlimits = xs.tolist()
_ylimits = ys.tolist()
return (_xlimits, _ylimits)[1]
Args:
var (:class:`xarray.DataArray`, optional): A :class:`xarray.DataArray`
variable that includes latitude,longitude coordinate information.
If not used, then *wrfin* must be provided.
geobounds (:class:`wrf.GeoBounds`, optional): The geobounds to
get the extents. If set to None and using the *var* parameter,
the geobounds will be taken from the variable. If using a
file, then the geobounds will be taken from the native grid.
wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \
iterable, optional): WRF-ARW NetCDF
data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile`
or an iterable sequence of the aforementioned types. If not used,
then *var* must be provided.
varname (:obj:`str`, optional): If using *wrfin*, then this will be the
variable name to use to determine the geobounds. The variable
can be a coordinate variable, or a regular variable that contains
coordinate attributes. If None,
then the 'XLAT', 'XLAT_M', 'XLONG', 'XLONG_M' variables
will be used.
timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The
desired time index. This value can be a positive integer,
negative integer, or
:data:`wrf.ALL_TIMES` (an alias for None) to return
all times in the file or sequence. Default is 0.
method (:obj:`str`, optional): The aggregation method to use for
sequences. Must be either 'cat' or 'join'.
'cat' combines the data along the Time dimension.
'join' creates a new dimension for the file index.
The default is 'cat'.
squeeze (:obj:`bool`, optional): Set to False to prevent dimensions
with a size of 1 from being automatically removed from the shape
of the output. Default is True.
cache (:obj:`dict`, optional): A dictionary of (varname, ndarray)
that can be used to supply pre-extracted NetCDF variables to the
computational routines. It is primarily used for internal
purposes, but can also be used to improve performance by
eliminating the need to repeatedly extract the same variables
used in multiple diagnostics calculations, particularly when using
large sequences of files.
Default is None.
Returns:
:obj:`list`: A list of [start_y, end_y] in the projected coordinate
system.
"""
wrf_proj, native_geobnds = _get_wrf_proj_geobnds(var, wrfin, varname,
timeidx, method, squeeze,
cache)
if geobounds is not None:
return wrf_proj.cartopy_ylim(geobounds)
return wrf_proj.cartopy_ylim(native_geobnds)
def ll_points(lat, lon):
"""Return the lower left latitude and longitude point(s).
This functions extracts the lower left corner points and returns the result
as either a single :class:`CoordPair` object, or a list of
:class:`CoordPair` objects.
This is primarily used for testing or constructing the corner point objects
from the XLAT and XLONG variables.
Args:
lat (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The latitude
array. Must be at least two dimensions.
lon (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The
longitude array. Must be at least two dimensions.
Returns:
:class:`wrf.CoordPair` or :obj:`list`: A single :class:`wrf.CoordPair`
object or a list of :class:`wrf.CoordPair` objects.
"""
latvals = np.ravel(to_np(lat)[..., 0, 0])
lonvals = np.ravel(to_np(lon)[..., 0, 0])
if latvals.shape[0] == 1:
return CoordPair(lat=float(latvals), lon=float(lonvals))
else:
return [CoordPair(lat=latvals[i], lon=lonvals[i])
for i in py3range(latvals.shape[0])]
def pairs_to_latlon(pairs):
"""Return latitude and longitude arrays from a sequence of \
:class:`wrf.CoordPair` objects.
This function converts a sequence of :class:`wrf.CoordPair` objects into
lists of latitude and longitude points. If the *pairs* argument is a
single :class:`wrf.CoordPair` object, then a single latitude and
longitude value is returned.
Args:
pairs (:class:`wrf.CoordPair` or sequence): A single
:class:`wrf.CoordPair` or sequence of :class:`wrf.CoordPair`.
Returns:
:obj:`tuple`: A tuple of (lat, lon), where lat and lon are single
values or lists of values.
"""
if isinstance(pairs, CoordPair):
return (pairs.lat, pairs.lon)
else:
lats = [pair.lat for pair in pairs]
lons = [pair.lon for pair in pairs]
return lats, lons
def is_latlon_pair(pair):
"""Return True if the :class:`wrf.CoordPair` is a lat/lon pair
Args:
pair (:class:`wrf.CoordPair`): A single :class:`wrf.CoordPair` object.
Returns:
:obj:`bool`: True if the pair is a lat/lon pair.
"""
if pair is not None:
return (pair.lat is not None and pair.lon is not None)
else:
return False