diff --git a/.gitignore b/.gitignore index f79f8e1..c070e32 100644 --- a/.gitignore +++ b/.gitignore @@ -43,5 +43,6 @@ Thumbs.db *.project *.pydevproject doc/build +doc/source/api/generated test/ipynb/.ipynb_checkpoints diff --git a/doc/source/_templates/product_table.txt b/doc/source/_templates/product_table.txt new file mode 100644 index 0000000..9f0748c --- /dev/null +++ b/doc/source/_templates/product_table.txt @@ -0,0 +1,238 @@ ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| Variable Name | Description | Available Units | Additional Keyword Arguments | ++====================+===============================================================+=============================+===============================================================================================+ +| avo | Absolute Vorticity | 10-5 s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| eth/theta_e | Equivalent Potential Temperature | K | **units** (str) : Set to desired units. Default is *'K'*. | +| | | | | +| | | degC | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| cape_2d | 2D cape (mcape/mcin/lcl/lfc) | J kg-1 ; J kg-1 ; m ; m | **missing** (float): Fill value for output only | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| cape_3d | 3D cape and cin | J kg-1 | **missing** (float): Fill value for output only | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| ctt | Cloud Top Temperature | degC | **units** (str) : Set to desired units. Default is *'degC'*. | +| | | | | +| | | K | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| cloudfrac | Cloud Fraction | % | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| dbz | Reflectivity | dBZ | **do_variant** (boolean): Set to True to enable variant calculation. Default is *False*. | +| | | | | +| | | | **do_liqskin** (boolean): Set to True to enable liquid skin calculation. Default is *False*. | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| mdbz | Maximum Reflectivity | dBZ | **do_variant** (boolean): Set to True to enable variant calculation. Default is *False*. | +| | | | | +| | | | **do_liqskin** (boolean): Set to True to enable liquid skin calculation. Default is *False*. | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| geopt/geopotential | Full Model Geopotential | m2 s-2 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| helicity | Storm Relative Helicity | m2 s-2 | **top** (float): The top level for the calculation in meters. Default is *3000.0*. | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| lat | Latitude | decimal degrees | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| lon | Longitude | decimal degrees | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| omg/omega | Omega | Pa s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| p/pres | Full Model Pressure | Pa | **units** (str) : Set to desired units. Default is *'Pa'*. | +| | | | | +| | (in specified units) | hPa | | +| | | | | +| | | mb | | +| | | | | +| | | torr | | +| | | | | +| | | mmhg | | +| | | | | +| | | atm | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| pressure | Full Model Pressure (hPa) | hPa | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| pvo | Potential Vorticity | PVU | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| pw | Precipitable Water | kg m-2 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| rh2 | 2m Relative Humidity | % | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| slp | Sea Level Pressure | hPa | **units** (str) : Set to desired units. Default is *'hPa'*. | +| | | | | +| | | hPa | | +| | | | | +| | | mb | | +| | | | | +| | | torr | | +| | | | | +| | | mmhg | | +| | | | | +| | | atm | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| ter | Model Terrain Height | m | **units** (str) : Set to desired units. Default is *'m'*. | +| | | | | +| | | km | | +| | | | | +| | | dm | | +| | | | | +| | | ft | | +| | | | | +| | | mi | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| td2 | 2m Dew Point Temperature | degC | **units** (str) : Set to desired units. Default is *'degC'*. | +| | | | | +| | | K | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| td | Dew Point Temperature | degC | **units** (str) : Set to desired units. Default is *'degC'*. | +| | | | | +| | | K | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| tc | Temperature in Celsius | degC | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| th/theta | Potential Temperature | K | **units** (str) : Set to desired units. Default is *'K'*. | +| | | | | +| | | degC | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| temp | Temperature (in specified units) | K | **units** (str) : Set to desired units. Default is *'K'*. | +| | | | | +| | | degC | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| tk | Temperature in Kelvin | K | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| times | Times in the File or Sequence | | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| xtimes | XTIME Coordinate | minutes since | | +| | | | | +| | (if applicable) | start of | | +| | | | | +| | | model run | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| tv | Virtual Temperature | K | **units** (str) : Set to desired units. Default is *'K'*. | +| | | | | +| | | degC | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| twb | Wet Bulb Temperature | K | **units** (str) : Set to desired units. Default is *'K'*. | +| | | | | +| | | degC | | +| | | | | +| | | degF | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| updraft_helicity | Updraft Helicity | m2 s-2 | **bottom** (float): The bottom level for the calculation in meters. Default is *2000.0*. | +| | | | | +| | | | **top** (float): The top level for the calculation in meters. Default is *5000.0*. | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| ua | U-component of Wind on Mass Points | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| va | V-component of Wind on Mass Points | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| wa | W-component of Wind on Mass Points | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| uvmet10 | 10 m U and V Components of Wind | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | Rotated to Earth Coordinates | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| uvmet | U and V Components of Wind | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | Rotated to Earth Coordinates | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| wspd_wdir | Wind Speed and Direction | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | in Grid Coordinates | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| wspd_wdir10 | 10m Wind Speed and Direction | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | in Grid Coordinates | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| uvmet_wspd_wdir | Wind Speed and Direction | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | Rotated to Earth Coordinates | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| uvmet10_wspd_wdir | 10m Wind Speed and Direction | m s-1 | **units** (str) : Set to desired units. Default is *'m s-1'*. | +| | | | | +| | Rotated to Earth Coordinates | km h-1 | | +| | | | | +| | | mi h-1 | | +| | | | | +| | | kt | | +| | | | | +| | | ft s-1 | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ +| z/height | Full Model Height | m | **msl** (boolean): Set to False to return AGL values. True is for MSL. Default is *True*. | +| | | | | +| | | km | **units** (str) : Set to desired units. Default is *'m'*. | +| | | | | +| | | dm | | +| | | | | +| | | ft | | +| | | | | +| | | mi | | ++--------------------+---------------------------------------------------------------+-----------------------------+-----------------------------------------------------------------------------------------------+ + + + diff --git a/doc/source/_templates/wrf_class.rst b/doc/source/_templates/wrf_class.rst new file mode 100644 index 0000000..d999287 --- /dev/null +++ b/doc/source/_templates/wrf_class.rst @@ -0,0 +1,32 @@ +{{ fullname }} +{{ underline }} + +.. currentmodule:: {{ module }} + +.. autoclass:: {{ objname }} + + {% block methods %} + + {% if methods %} + .. rubric:: Methods + + .. autosummary:: + {% for item in methods %} + {% if item != '__init__' %} + ~{{ name }}.{{ item }} + {% endif %} + {%- endfor %} + {% endif %} + {% endblock %} + + {% block attributes %} + {% if attributes %} + .. rubric:: Attributes + + .. autosummary:: + {% for item in attributes %} + ~{{ name }}.{{ item }} + {%- endfor %} + {% endif %} + {% endblock %} + \ No newline at end of file diff --git a/doc/source/api/generated/README b/doc/source/api/generated/README new file mode 100644 index 0000000..4d0652a --- /dev/null +++ b/doc/source/api/generated/README @@ -0,0 +1 @@ +This location contains the RST files created by Sphinx AutoSummary for the API. diff --git a/doc/source/api/index.rst b/doc/source/api/index.rst new file mode 100644 index 0000000..25a5967 --- /dev/null +++ b/doc/source/api/index.rst @@ -0,0 +1,55 @@ +API Reference +============= + +Diagnostic Routines +------------------- + +.. autosummary:: + :nosignatures: + :toctree: generated/ + + wrf.getvar + + +Interpolation Routines +---------------------- + +.. autosummary:: + :nosignatures: + :toctree: generated/ + + wrf.interplevel + wrf.vertcross + wrf.interpline + wrf.vinterp + + +Extraction Routines +------------------- + +.. autosummary:: + :nosignatures: + :toctree: generated/ + + wrf.npvalues + wrf.util.either + wrf.util.is_moving_domain + + +Raw Computation Routines +------------------------ + +.. autosummary:: + :nosignatures: + :toctree: generated/ + + wrf.xy + wrf.interp1d + wrf.interp2dxy + wrf.slp + wrf.uvmet + wrf.cape_2d + wrf.cape_3d + + + \ No newline at end of file diff --git a/doc/source/conf.py b/doc/source/conf.py index c1ec1b8..4cccfc0 100644 --- a/doc/source/conf.py +++ b/doc/source/conf.py @@ -15,6 +15,36 @@ import sys import os +try: + from unittest.mock import MagicMock +except ImportError: + from mock import Mock as MagicMock + +class Mock(MagicMock): + @classmethod + def __getattr__(cls, name): + return Mock() + +MOCK_MODULES = ["numpy", "numpy.ma", "xarray", + "pandas", "matplotlib", "netCDF4", "mpl_toolkits.basemap", + "wrf._wrffortran"] +sys.modules.update((mod_name, Mock()) for mod_name in MOCK_MODULES) + +consts = {"DEFAULT_FILL" : 9.9692099683868690E36} + +class MockWrfConstants(object): + def __init__(self): + self.__dict__ = consts + +def mock_asscalar(val): + return float(val) + +sys.modules["wrf._wrffortran"].wrf_constants = MockWrfConstants() +sys.modules["numpy"].asscalar = mock_asscalar + + + + # If extensions (or modules to document with autodoc) are in another directory, # add these directories to sys.path here. If the directory is relative to the # documentation root, use os.path.abspath to make it absolute, like shown here. @@ -29,9 +59,26 @@ import os # extensions coming with Sphinx (named 'sphinx.ext.*') or your custom # ones. extensions = [ - 'sphinx.ext.autodoc', 'sphinx.ext.napoleon' + 'sphinx.ext.autodoc', + 'sphinx.ext.napoleon', + 'sphinx.ext.autosummary', + 'sphinx.ext.intersphinx', ] +intersphinx_mapping = { + 'matplotlib': ('http://matplotlib.org/', None), + 'python': ('http://docs.python.org/3/', None), + 'numpy': ('http://docs.scipy.org/doc/numpy/', None), + 'xarray': ('http://xarray.pydata.org/en/stable/', None), + 'wrapt': ('http://wrapt.readthedocs.io/en/latest/', + None) + } + +napoleon_use_admonition_for_examples = True +napoleon_include_special_with_doc = True + +autosummary_generate = True + # Add any paths that contain templates here, relative to this directory. templates_path = ['_templates'] @@ -48,7 +95,7 @@ master_doc = 'index' # General information about the project. project = u'wrf-python' -copyright = u'2016, UCAR' +copyright = u'2016, University Corporation for Atmospheric Research' author = u'Bill Ladwig' # The version info for the project you're documenting, acts as replacement for @@ -56,9 +103,10 @@ author = u'Bill Ladwig' # built documents. # # The short X.Y version. -version = u'0.0.1' +import wrf +version = wrf.__version__ # The full version, including alpha/beta/rc tags. -release = u'0.0.1' +release = wrf.__version__ # The language for content autogenerated by Sphinx. Refer to documentation # for a list of supported languages. @@ -110,7 +158,12 @@ todo_include_todos = False # The theme to use for HTML and HTML Help pages. See the documentation for # a list of builtin themes. -html_theme = 'alabaster' +import sphinx_rtd_theme + +html_theme = "sphinx_rtd_theme" + +html_theme_path = [sphinx_rtd_theme.get_html_theme_path()] +#html_theme = 'alabaster' # Theme options are theme-specific and customize the look and feel of a theme # further. For a list of options available for each theme, see the diff --git a/doc/source/index.rst b/doc/source/index.rst index 1365f5d..ed8c07a 100644 --- a/doc/source/index.rst +++ b/doc/source/index.rst @@ -3,28 +3,16 @@ You can adapt this file completely to your liking, but it should at least contain the root `toctree` directive. -Welcome to wrf-python's documentation! -====================================== +User Guide for wrf-python +========================== Contents: .. toctree:: - :maxdepth: 2 + :maxdepth: 1 + + api/index -Extraction Routine ------------------- -.. autofunction:: wrf.getvar - -Interpolation Routines ----------------------- - -.. autofunction:: wrf.interplevel - -.. autofunction:: wrf.vertcross - -.. autofunction:: wrf.interpline - -.. autofunction:: wrf.vinterp Indices and tables diff --git a/ncl_reference/getvar_table.txt b/ncl_reference/getvar_table.txt deleted file mode 100644 index 37ab148..0000000 --- a/ncl_reference/getvar_table.txt +++ /dev/null @@ -1,74 +0,0 @@ - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | Variable Name | Description | Units | Additional Keyword Arguments | - +====================+===============================================================+=====================+===============================================================================================+ - | avo | Absolute Vorticity | 10-5 s-1 | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | eth/theta_e | Equivalent Potential Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | cape_2d | 2D cape (mcape/mcin/lcl/lfc) | J/kg / J/kg / m / m | missing: Fill value for output only (float) | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | cape_3d | 3D cape and cin | J/kg | missing: Fill value for output only (float) | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | ctt | Cloud Top Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | dbz | Reflectivity | dBz | do_variant: Set to True to enable variant calculation. Default is False. | - | | | | do_liqskin : Set to True to enable liquid skin calculation. Default is False. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | mdbz | Maximum Reflectivity | dBz | do_variant: Set to True to enable variant calculation. Default is False. | - | | | | do_liqskin: Set to True to enable liquid skin calculation. Default is False. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | geopt/geopotential | Full Model Geopotential | m2 s-2 | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | helicity | Storm Relative Helicity | m-2/s-2 | top: The top level for the calculation in meters (float). Default is 3000.0. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | lat | Latitude | decimal degrees | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | lon | Longitude | decimal degrees | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | omg/omega | Omega | Pa/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | p/pres | Full Model Pressure | Pa | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | pressure | Full Model Pressure | hPa | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | pvo | Potential Vorticity | PVU | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | pw | Precipitable Water | kg m-2 | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | rh2 | 2m Relative Humidity | % | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | slp | Sea Level Pressure | hPa | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | ter | Model Terrain Height | m | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | td2 | 2m Dew Point Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | td | Dew Point Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | tc | Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | th/theta | Potential Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | tk | Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | times | Times in the File or Sequence | | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | tv | Virtual Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | twb | Wet Bulb Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | updraft_helicity | Updraft Helicity | m-2/s-2 | bottom: The bottom level for the calculation in meters (float). Default is 2000.0. | - | | | | top: The top level for the calculation in meters (float). Default is 5000.0. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | ua | U-component of Wind on Mass Points | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | va | V-component of Wind on Mass Points | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | wa | W-component of Wind on Mass Points | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | uvmet10 | 10 m U and V Components of Wind Rotated to Earth Coordinates | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | uvmet | U and V Components of Wind Rotated to Earth Coordinates | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | z/height | Full Model Height | m | msl: Set to False to return AGL values. Otherwise, MSL. Default is True. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ diff --git a/ncl_reference/product_table.txt b/ncl_reference/product_table.txt new file mode 100644 index 0000000..b47b491 --- /dev/null +++ b/ncl_reference/product_table.txt @@ -0,0 +1,77 @@ ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| Variable Name | Description | Units | Additional Keyword Arguments | ++====================+===============================================================+=====================+===============================================================================================+ +| avo | Absolute Vorticity | 10-5 s-1 | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| eth/theta_e | Equivalent Potential Temperature | K | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| cape_2d | 2D cape (mcape/mcin/lcl/lfc) | J/kg / J/kg / m / m | missing: Fill value for output only (float) | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| cape_3d | 3D cape and cin | J/kg | missing: Fill value for output only (float) | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| ctt | Cloud Top Temperature | C | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| cloudfrac | Cloud Fraction | % | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| dbz | Reflectivity | dBz | do_variant: Set to True to enable variant calculation. Default is False. | +| | | | do_liqskin : Set to True to enable liquid skin calculation. Default is False. | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| mdbz | Maximum Reflectivity | dBz | do_variant: Set to True to enable variant calculation. Default is False. | +| | | | do_liqskin: Set to True to enable liquid skin calculation. Default is False. | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| geopt/geopotential | Full Model Geopotential | m2 s-2 | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| helicity | Storm Relative Helicity | m2 s-2 | top: The top level for the calculation in meters (float). Default is 3000.0. | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| lat | Latitude | decimal degrees | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| lon | Longitude | decimal degrees | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| omg/omega | Omega | Pa/s | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| p/pres | Full Model Pressure | Pa | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| pressure | Full Model Pressure | hPa | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| pvo | Potential Vorticity | PVU | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| pw | Precipitable Water | kg m-2 | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| rh2 | 2m Relative Humidity | % | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| slp | Sea Level Pressure | hPa | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| ter | Model Terrain Height | m | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| td2 | 2m Dew Point Temperature | C | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| td | Dew Point Temperature | C | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| tc | Temperature | C | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| th/theta | Potential Temperature | K | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| tk | Temperature | K | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| times | Times in the File or Sequence | | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| tv | Virtual Temperature | K | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| twb | Wet Bulb Temperature | K | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| updraft_helicity | Updraft Helicity | m2 s-2 | bottom: The bottom level for the calculation in meters (float). Default is 2000.0. | +| | | | top: The top level for the calculation in meters (float). Default is 5000.0. | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| ua | U-component of Wind on Mass Points | m/s | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| va | V-component of Wind on Mass Points | m/s | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| wa | W-component of Wind on Mass Points | m/s | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| uvmet10 | 10 m U and V Components of Wind Rotated to Earth Coordinates | m/s | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| uvmet | U and V Components of Wind Rotated to Earth Coordinates | m/s | | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ +| z/height | Full Model Height | m | msl: Set to False to return AGL values. Otherwise, MSL. Default is True. | ++--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ + diff --git a/setup.py b/setup.py index afa319d..2e1c387 100755 --- a/setup.py +++ b/setup.py @@ -1,3 +1,5 @@ +import os +import sys import setuptools import numpy.distutils.core @@ -25,10 +27,24 @@ ext1 = numpy.distutils.core.Extension( with open("src/wrf/version.py") as f: exec(f.read()) -requirements = [ - "numpy>=1.9.0", - "wrapt>=1.10" - ] +on_rtd = os.environ.get("READTHEDOCS", None) == "True" +#on_rtd=True +if on_rtd: + if sys.version_info < (3,3): + requirements = ["mock"] # for python2 and python < 3.3 + else: + requirements = [] # for >= python3.3 + ext_modules = [] + +else: + # Place install_requires into the text file "requirements.txt" + with open("requirements.txt") as f2: + requirements = f2.read().strip().splitlines() + + #if sys.version_info < (3,3): + # requirements.append("mock") + ext_modules = [ext1] + numpy.distutils.core.setup( author = "Bill Ladwig", @@ -61,7 +77,7 @@ numpy.distutils.core.setup( name = "wrf", version = __version__, packages = setuptools.find_packages("src"), - ext_modules = [ext1], + ext_modules = ext_modules, package_dir = {"" : "src"}, #namespace_packages=["wrf"], # Note: If this doesn't work, you need to add the file to MANIFEST diff --git a/src/wrf/__init__.py b/src/wrf/__init__.py index 828d0c1..ae2d6f1 100755 --- a/src/wrf/__init__.py +++ b/src/wrf/__init__.py @@ -6,12 +6,3 @@ from .api import * __all__ = [] __all__.extend(api.__all__) - - - - - - - - - \ No newline at end of file diff --git a/src/wrf/cache.py b/src/wrf/cache.py index 98672a8..e005e21 100644 --- a/src/wrf/cache.py +++ b/src/wrf/cache.py @@ -9,6 +9,38 @@ from .config import get_cache_size _local_storage = local() def cache_item(key, product, value): + """Store an item in the cache. + + The cache should be viewed as two nested dictionaries. The outer key is + usually the id for the sequence where the cached item was generated. The + inner key is the product type. + + Storing a cached item behaves like: + + cache[key][product] = value + + The cache is thread local, so stored items are only available in + the thread that cached them. + + Args: + + key (:obj:`int`): The outer dictionary cache key, which is typically + the id of the sequence where the cached item was generated. + + product (:obj:`str`): The inner dictionary cache key, which is a + string for the product type. + + value (:obj:`object`): The object to store in the cache. + + Returns: + + None. + + See Also: + + :meth:`get_cached_item` + + """ global _local_storage if key is None: @@ -32,6 +64,36 @@ def cache_item(key, product, value): def get_cached_item(key, product): + """Return an item from the cache. + + The cache should be viewed as two nested dictionaries. The outer key is + usually the id for the sequence where the cached item was generated. The + inner key is the product type. + + Retrieving a cached item behaves like: + + value = cache[key][product] + + The cache is thread local, so stored items are only available in + the thread that cached them. + + Args: + + key (:obj:`int`): The outer dictionary cache key, which is typically + the id of the sequence where the cached item was generated. + + product (:obj:`str`): The inner dictionary cache key, which is a + string for the product type. + + Returns: + + :obj:`object`: The cached object. + + See Also: + + :meth:`cache_item` + + """ if key is None: return None @@ -50,6 +112,15 @@ def get_cached_item(key, product): return result def _get_cache(): + """Return the cache. + + This is primarily used for testing. + + Returns: + + :class:`threading.local` + + """ return getattr(_local_storage, "cache", None) diff --git a/src/wrf/cape.py b/src/wrf/cape.py index f7daf65..7eba2c5 100755 --- a/src/wrf/cape.py +++ b/src/wrf/cape.py @@ -11,12 +11,75 @@ from .util import extract_vars from .metadecorators import set_cape_metadata @set_cape_metadata(is2d=True) -def get_2dcape(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, missing=Constants.DEFAULT_FILL): - """Return the 2d fields of cape, cin, lcl, and lfc""" + """Return the 2d fields of CAPE, CIN, LCL, and LFC. + + The leftmost dimension of the returned array represents four different + quantities: + + - return_val[0,...] will contain CAPE [J kg-1] + - return_val[1,...] will contain CIN [J kg-1] + - return_val[2,...] will contain LCL [m] + - return_val[3,...] will contain LFC [m] + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + missing (:obj:`float`): The fill value to use for the output. + Default is :data:`wrf.Constants.DEFAULT_FILL`. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + cape, cin, lcl, and lfc values as an array whose + leftmost dimension is 4 (0=CAPE, 1=CIN, 2=LCL, 3=LFC). + 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. + + """ varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] @@ -65,12 +128,72 @@ def get_2dcape(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, @set_cape_metadata(is2d=False) -def get_3dcape(wrfnc, timeidx=0, method="cat", +def get_3dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, missing=Constants.DEFAULT_FILL): - """Return the 3d fields of cape and cin""" + """Return the three-dimensional CAPE and CIN. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain CAPE [J kg-1] + - return_val[1,...] will contain CIN [J kg-1] + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + missing (:obj:`float`): The fill value to use for the output. + Default is :data:`wrf.Constants.DEFAULT_FILL`. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + CAPE and CIN as an array whose leftmost dimension is 2 (0=CAPE, 1=CIN). + 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. + + """ varnames = ("T", "P", "PB", "QVAPOR", "PH", "PHB", "HGT", "PSFC") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] diff --git a/src/wrf/cloudfrac.py b/src/wrf/cloudfrac.py index 62848c6..16eee8d 100644 --- a/src/wrf/cloudfrac.py +++ b/src/wrf/cloudfrac.py @@ -7,10 +7,70 @@ from .metadecorators import set_cloudfrac_metadata from .util import extract_vars @set_cloudfrac_metadata() -def get_cloudfrac(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): + """Return the cloud fraction. - ncvars = extract_vars(wrfnc, timeidx, ("P", "PB", "QVAPOR", "T"), + The leftmost dimension of the returned array represents three different + quantities: + + - return_val[0,...] will contain LOW level cloud fraction + - return_val[1,...] will contain MID level cloud fraction + - return_val[2,...] will contain HIGH level cloud fraction + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + cloud fraction array whose leftmost dimension is 3 (LOW=0, MID=1, + HIGH=2). + 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. + + """ + + ncvars = extract_vars(wrfin, timeidx, ("P", "PB", "QVAPOR", "T"), method, squeeze, cache, meta=False, _key=_key) diff --git a/src/wrf/computation.py b/src/wrf/computation.py index 3fe0b5c..27c8984 100644 --- a/src/wrf/computation.py +++ b/src/wrf/computation.py @@ -10,6 +10,7 @@ from .extension import (_interpz3d, _interp2dxy, _interp1d, _slp, _tk, _td, _srhel, _udhel, _avo, _pvo, _eth, _wetbulb, _tv, _omega, _pw) from .util import from_var +from .decorators import convert_units from .metadecorators import (set_alg_metadata, set_uvmet_alg_metadata, set_interp_metadata, set_cape_alg_metadata, set_cloudfrac_alg_metadata) @@ -18,52 +19,714 @@ from .interputils import get_xy @set_interp_metadata("xy") def xy(field, pivot_point=None, angle=None, start_point=None, end_point=None, meta=True): + """Return the x,y points for a line within a two-dimensional grid. + + This function is primarily used to obtain the x,y points when making a + cross section. + + Args: + + field3d (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + field with at least two dimensions. + + pivot_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, + in the form of [x, y] (or [west_east, south_north]), which + indicates the x,y location through which the plane will pass. + Must also specify `angle`. + + angle (:obj:`float`, optional): Only valid for cross sections where + a plane will be plotted through + a given point on the model domain. 0.0 represents a S-N cross + section. 90.0 is a W-E cross section. + + start_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the start + x,y location through which the plane will pass. + + end_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the end x,y + location through which the plane will pass. + + meta (:obj:`bool`, optional): Set to False to disable metadata and + return :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: An array of + x,y points, which has shape num_points x 2. + 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. + + Examples: + + Example 1: Using Pivot Point and Angle + + .. code-block:: python + + from wrf import getvar, xy + from netCDF4 import Dataset + + wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") + field = wrf.getvar(wrfnc, "slp") + + # Use the center of the grid + pivot = (field.shape[-1]/2.0, field.shape[-2]/2.0) + + # West-East + angle = 90.0 + + xy_points = xy(field, pivot_point=pivot, angle=angle) + + Example 2: Using Start Point and End Point + + .. code-block:: python + + from wrf import getvar, xy + from netCDF4 import Dataset + + wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") + field = wrf.getvar(wrfnc, "slp") + + # Make a diagonal of lower left to upper right + start = (0, 0) + end = (-1, -1) + + xy_points = xy(field, start_point=start, end_point=end) + + + """ return get_xy(field, pivot_point, angle, start_point, end_point) @set_interp_metadata("1d") def interp1d(field, z_in, z_out, missingval=Constants.DEFAULT_FILL, meta=True): + """Return the linear interpolation of a one-dimensional variable. + + This function is typically used to interpolate a variable in a vertical + column, but the coordinate system need not be a vertical coordinate + system. Multiple interpolation points may be specified in the *z_out* + parameter. + + Args: + + field (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + one-dimensional field. Metadata for *field* is only copied + to the output if *field* is a :class:`xarray.DataArray` object. + + z_in (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The + one-dimensional coordinates associated with *field* (usually the + vertical coordinates, either height or pressure). + + z_out (:class:`xarray.DataArray`, :class:`numpy.ndarray`): A + one-dimensional array of *z_in* coordinate points to interpolate + to. Must be the same type as *z_in*. + + missingval (:obj:`float`, optional): The fill value to use for the + output. Default is :data:`wrf.Constants.DEFAULT_FILL`. + + meta (:obj:`bool`, optional): Set to False to disable metadata and + return :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: An array with the + same dimensionality as *z_out* containing the interpolated values. + 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. + + + Examples: + + Example 1: Calculate the 850 hPa and 500 hPa values at location \ + x,y = (100,200) + + .. code-block:: python + + import numpy as np + from wrf import getvar, interp1d + from netCDF4 import Dataset + + wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") + + # Get a 1D vertical column for pressure at location x,y = 100,200 + p_1d = wrf.getvar(wrfnc, "pres", units="hPa")[:,200,100] + + # Get a 1D vertical column for height at location 100,200 + ht_1d = wrf.getvar(wrfnc, "z", units="dm")[:,200,100] + + # Want the heights (in decameters) at 850, 500 hPa + levels = np.asarray([850., 500.]) + + # Get the 850 hPa and 500 hPa values at location 100,200. + interp_vals = interp1d(p_1d, ht_1d, levels) + + """ return _interp1d(field, z_in, z_out, missingval) @set_interp_metadata("2dxy") def interp2dxy(field3d, xy, meta=True): + """Return a cross section for a three-dimensional field. + + The returned array will hold the vertical cross section data along the + line described by *xy*. + + This method differs from :meth:`wrf.vertcross` in that it will return + all vertical levels found in *field3d*. :meth:`wrf.vertcross` includes + an additional interpolation to set the output to a fixed number of + vertical levels. Also, a :class:`numpy.ma.MaskedArray` is not created + and this routine should be considered as low-level access to the underlying + Fortran routine. + + See Also: + + :meth:`wrf.xy`, :meth:`wrf.vertcross` + + Args: + + field3d (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The + array to interpolate with at least three dimensions, whose + rightmost dimensions are nz x ny x nx. + + xy (:class:`xarray.DataArray` or :class:`numpy.ndarray`): An array + of one less dimension than *field3d*, whose rightmost dimensions + are nxy x 2. This array holds the x,y pairs of a line across the + model domain. The requested vertical cross section will be + extracted from *field3d* along this line. + + meta (:obj:`bool`, optional): Set to False to disable metadata and + return :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: An array + containing the vertical cross section along the line *xy*. The + returned dimensions will be the same as *xy*, but with the rightmost + dimensions being nz x nxy. 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. + + Examples: + + Example 1: Calculate the vertical cross section for RH for a diagonal + line from the lower left to the upper right of the domain. + + .. code-block:: python + + from wrf import getvar, xy, interp2dxy + from netCDF4 import Dataset + + wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") + + rh = getvar(wrfnc, "rh") + start = (0, 0) + end = (-1, -1) + xy_line = xy(rh, start_point=start, end_point=end) + + vert_cross = interp2dxy(rh, xy_line) + + """ return _interp2dxy(field3d, xy) @set_interp_metadata("horiz") -def interpz3d(field3d, z, desiredloc, missingval=Constants.DEFAULT_FILL, +def interpz3d(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, meta=True): - return _interpz3d(field3d, z, desiredloc, missingval) + """Return the field interpolated to a specified pressure or height level. + + This function is roughly equivalent to :meth:`interplevel`, but does not + handle multi-product diagnostics (uvmet, cape_3d, etc) that contain an + additional leftmost dimension for the product type. Also, a + :class:`numpy.ma.MaskedArray` is not created and this routine should + be considered as low-level access to the underlying Fortran routine. + + See Also: + + :meth:`wrf.interplevel` + + Args: + + field3d (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional field to interpolate, with the rightmost + dimensions of nz x ny x nx. + + vert (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional array for the vertical coordinate, typically + pressure or height. This array must have the same dimensionality + as *field3d*. + + desiredlev (:obj:`float`): The desired vertical level. + Must be in the same units as the *vert* parameter. + + missing (:obj:`float`): The fill value to use for the output. + Default is :data:`wrf.Constants.DEFAULT_FILL`. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + interpolated variable. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + Example: + + Example 1: Interpolate Geopotential Height to 500 hPa + + .. code-block:: python + + from netCDF4 import Dataset + from wrf import getvar, interpz3d + + wrfin = Dataset("wrfout_d02_2010-06-13_21:00:00") + + p = getvar(wrfin, "pressure") + ht = getvar(wrfin, "z", units="dm") + + ht_500 = interpz3d(ht, p, 500.0) + + """ + return _interpz3d(field3d, vert, desiredlev, missing) -@set_alg_metadata(2, "pres", refvarndims=3, units="hPa", - description="sea level pressure") -def slp(height, tkel, pres, qv, meta=True): +@set_alg_metadata(2, "pres", refvarndims=3, description="sea level pressure") +@convert_units("pressure", "hpa") +def slp(height, tkel, pres, qv, meta=True, units="hPa"): + """Return the sea level pressure. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + height (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] with the rightmost dimensions being + bottom_top x south_north x west_east. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *height*. + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with + the same dimensionality as *height*. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *height*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'slp'. Default + is 'hPa'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + sea level pressure. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.temp`, :meth:`wrf.tk` + + """ return _slp(height, tkel, pres, qv) -@set_alg_metadata(3, "pres", units="K", - description="temperature") -def tk(pres, theta, meta=True): +@set_alg_metadata(3, "pres", description="temperature") +@convert_units("temp", "k") +def temp(pres, theta, meta=True, units="K"): + """Return the temperature in the specified units. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with at least + three dimensions. The rightmost dimensions are bottom_top x + south_north x west_east. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + theta (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Potential + temperature (perturbation plus reference temperature) in [K] with + the same dimensionality as *pres*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'temp'. Default + is 'K'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + temperature in the specified units. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.tk` + + """ + return _tk(pres, theta) -@set_alg_metadata(3, "pres", units="degC", - description="dew point temperature") -def td(pres, qv, meta=True): +def tk(pres, theta, meta=True): + """Return the temperature in Kelvin. + + This is an alias for :meth:`wrf.temp(pres, theta, meta, units='K'). This + function is here for users that are familiar with NCL's wrf_tk. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with at least + three dimensions. The rightmost dimensions are bottom_top x + south_north x west_east. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + theta (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Potential + temperature (perturbation plus reference temperature) in [K] with + the same dimensionality as *pres*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + temperature in Kelvin. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.temp` + + """ + return temp(pres, theta, meta, units="K") + + +@set_alg_metadata(3, "pres", description="dew point temperature") +@convert_units("temp", "c") +def td(pres, qv, meta=True, units="degC"): + """Return the dewpoint temperature. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with at least + three dimensions. The rightmost dimensions are bottom_top x + south_north x west_east. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *pres*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'dp'. Default + is 'degC'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + dewpoint temperature. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.rh` + + """ return _td(pres, qv) -@set_alg_metadata(3, "pres", - description="relative humidity", units=None) +@set_alg_metadata(3, "pres", description="relative humidity", units=None) def rh(qv, pres, tkel, meta=True): + """Return the relative humidity. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with at least three dimensions. The + rightmost dimensions are bottom_top x south_north x west_east. + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with the + same dimensionality as *qv*. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *qv*. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + relative humidity. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.td` + + """ return _rh(qv, pres, tkel) @set_uvmet_alg_metadata(latarg="lat", windarg="u") -def uvmet(u, v, lat, lon, cen_long, cone, meta=True): +@convert_units("wind", "m s-1") +def uvmet(u, v, lat, lon, cen_long, cone, meta=True, units="m s-1"): + """Return the u,v components of the wind rotated to earth coordinates. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain U + - return_val[1,...] will contain V + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + u (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The u + component of the wind [m s-1]. This variable can be staggered or + unstaggered, but must be at least two dimensions. If staggered, + the rightmost dimensions are south_north x west east. + + If staggered, the rightmost dimensions are south_north x + west_east_stag. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + v (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The v + component of the wind [m s-1]. This variable can be staggered or + unstaggered, but must be at least two dimensions. If staggered, + the rightmost dimensions are south_north x west east. + + IF staggered, the rightmost dimensions are south_north_stag x + west_east. + + lat (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The + latitude array. + + This array can either be: + + - two-dimensional of size south_north x west_east. + - multi-dimensional with the same number of dimensions as *u* + and *v*, but with rightmost dimensions south_north x + west_east and the same leftmost dimensions as *u* and *v* + - multi-dimensional with one fewer dimensions as *u* and *v*, + with rightmost dimensions south_north x west_east and the same + leftmost dimensions as *u* and *v*, minus the + third-from-the-right dimension of *u* and *v*. + + Note: + + This variable must also be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + lon (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The + longitude array. + + This array can either be: + + - two-dimensional of size south_north x west_east. + - multi-dimensional with the same number of dimensions as *u* + and *v*, but with rightmost dimensions south_north x + west_east and the same leftmost dimensions as *u* and *v* + - multi-dimensional with one fewer dimensions as *u* and *v*, + with rightmost dimensions south_north x west_east and the same + leftmost dimensions as *u* and *v*, minus the + third-from-the-right dimension of *u* and *v*. + + + cen_long (:obj:`float`): The standard longitude for the map projection. + + cone (:obj:`float`): The cone factor used for the map project. IF the + projection is not a conic projection, the *cone* is simply 1.0. + For conic projections, the cone factor is given by: + + .. code-block:: python + + if((fabs(true_lat1 - true_lat2) > 0.1) and + (fabs(true_lat2 - 90.) > 0.1)): + cone = (log(cos(true_lat1*radians_per_degree)) + - log(cos(true_lat2*radians_per_degree))) + + cone = (cone / + (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) + - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) + else: + cone = sin(fabs(true_lat1)*radians_per_degree) + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet'. Default + is 'm s-1'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + u,v components of the wind rotated to earth coordinates. The leftmost + dimension size is 2, for u and v. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar` + + """ return _uvmet(u, v, lat, lon, cen_long, cone) @@ -71,19 +734,149 @@ def uvmet(u, v, lat, lon, cen_long, cone, meta=True): description=from_var("field", "description"), units=from_var("field", "units")) def smooth2d(field, passes, meta=True): + """Return the field smoothed. + + This routine does not modify the original data. + + Args: + + field (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The field + to smooth, which must be at least two dimensions. Missing/fill + values will be ignored as long as the type is either a + :class:`numpy.ma.MaskedArray or a :class:`xarray.DataArray` with + a *_FillValue* attribute. + + passes (:obj:`int`): The number of smoothing passes. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Returns: + + :class:`xarray.DataArray`, :class:`numpy.ma.MaskedArray` or \ + :class:`numpy.ndarray`): The smoothed field. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be either a :class:`numpy.ndarray` or a :class:`numpy.ma.MaskedArray` + depending on the type for *field*. + + + """ return _smooth2d(field, passes) @set_cape_alg_metadata(is2d=True, copyarg="pres_hpa") -def cape_2d(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, ter_follow, - missing=Constants.DEFAULT_FILL, meta=True): +def cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, + missingval=Constants.DEFAULT_FILL, meta=True): + """Return the two-dimensional CAPE, CIN, LCL, and LFC. + + This function calculates the maximum convective available potential + energy (CAPE), maximum convective inhibition (CIN), + lifted condensation level (LCL), and level of free convection (LFC). This + function uses the RIP [Read/Interpolate/plot] code to calculate + potential energy (CAPE) and convective inhibition + (CIN) [J kg-1] only for the parcel with max theta-e + in the column (i.e. something akin to Colman's MCAPE). CAPE is defined as + the accumulated buoyant energy from the level of free convection (LFC) to + the equilibrium level (EL). CIN is defined as the accumulated negative + buoyant energy from the parcel starting point to the LFC. The word 'parcel' + here refers to a 500 meter deep parcel, with actual temperature and + moisture averaged over that depth. + + The leftmost dimension of the returned array represents four different + quantities: + + - return_val[0,...] will contain CAPE [J kg-1] + - return_val[1,...] will contain CIN [J kg-1] + - return_val[2,...] will contain LCL [m] + - return_val[3,...] will contain LFC [m] + + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [hPa] with at + least three dimensions. The rightmost dimensions can be + top_bottom x south_north x west_east or bottom_top x south_north x + west_east. + + Note: + + The units for *pres_hpa* are [hPa]. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *pres_hpa*. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *pres_hpa*. + + height (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] with the same dimensionality as + *pres_hpa*. + + terrain (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Terrain height in [m]. This is at least a two-dimensional array + with the same dimensionality as *pres_hpa*, excluding the vertical + (bottom_top/top_bottom) dimension. + + psfc_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + The surface pressure in [hPa]. This is at least a two-dimensional + array with the same dimensionality as *pres_hpa*, excluding the + vertical (bottom_top/top_bottom) dimension. + + Note: + + The units for *psfc_hpa* are [hPa]. + + ter_follow (:obj:`bool`): A boolean that should be set to True if the + data uses terrain following coordinates (WRF data). Set to + False for pressure level data. + + missingval (:obj:`float`, optional): The fill value to use for the + output. Default is :data:`wrf.Constants.DEFAULT_FILL`. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + cape, cin, lcl, and lfc values as an array whose + leftmost dimension is 4 (0=CAPE, 1=CIN, 2=LCL, 3=LFC) . If xarray is + enabled and the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.cape_3d` + + """ if isinstance(ter_follow, bool): ter_follow = 1 if ter_follow else 0 i3dflag = 0 - cape_cin = _cape(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, - missing, i3dflag, ter_follow) + cape_cin = _cape(pres_hpa, tkel, qv, height, terrain, psfc_hpa, + missingval, i3dflag, ter_follow) left_dims = cape_cin.shape[1:-3] right_dims = cape_cin.shape[-2:] @@ -100,31 +893,250 @@ def cape_2d(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, ter_follow, result[2,...,:,:] = cape_cin[1,...,-2,:,:] result[3,...,:,:] = cape_cin[1,...,-3,:,:] - return ma.masked_values(result, missing) + return ma.masked_values(result, missingval) @set_cape_alg_metadata(is2d=False, copyarg="pres_hpa") -def cape_3d(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, ter_follow, - missing=Constants.DEFAULT_FILL, meta=True): +def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, + missingval=Constants.DEFAULT_FILL, meta=True): + """Return the three-dimensional CAPE and CIN. + + This function calculates the maximum convective available potential + energy (CAPE) and maximum convective inhibition (CIN). This + function uses the RIP [Read/Interpolate/plot] code to calculate + potential energy (CAPE) and convective inhibition + (CIN) [J kg-1] for every grid point in the entire 3D domain + (treating each grid point as a parcel). + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain CAPE [J kg-1] + - return_val[1,...] will contain CIN [J kg-1] + + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [hPa] with at + least three dimensions. The rightmost dimensions can be + top_bottom x south_north x west_east or bottom_top x south_north x + west_east. + + Note: + + The units for *pres_hpa* are [hPa]. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *pres_hpa*. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *pres_hpa*. + + height (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] with the same dimensionality as + *pres_hpa*. + + terrain (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Terrain height in [m]. This is at least a two-dimensional array + with the same dimensionality as *pres_hpa*, excluding the vertical + (bottom_top/top_bottom) dimension. + + psfc_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + The surface pressure in [hPa]. This is at least a two-dimensional + array with the same dimensionality as *pres_hpa*, excluding the + vertical (bottom_top/top_bottom) dimension. + + Note: + + The units for *psfc_hpa* are [hPa]. + + ter_follow (:obj:`bool`): A boolean that should be set to True if the + data uses terrain following coordinates (WRF data). Set to + False for pressure level data. + + missingval (:obj:`float`, optional): The fill value to use for the + output. Default is :data:`wrf.Constants.DEFAULT_FILL`. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + CAPE and CIN as an array whose + leftmost dimension is 2 (0=CAPE, 1=CIN). If xarray is + enabled and the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.cape_2d` + + """ if isinstance(ter_follow, bool): ter_follow = 1 if ter_follow else 0 i3dflag = 1 - cape_cin = _cape(pres_hpa, tkel, qvapor, height, terrain, psfc_hpa, - missing, i3dflag, ter_follow) + cape_cin = _cape(pres_hpa, tkel, qv, height, terrain, psfc_hpa, + missingval, i3dflag, ter_follow) - return ma.masked_values(cape_cin, missing) + return ma.masked_values(cape_cin, missingval) @set_cloudfrac_alg_metadata(copyarg="pres") def cloudfrac(pres, relh, meta=True): + """Return the cloud fraction. + + The leftmost dimension of the returned array represents three different + quantities: + + - return_val[0,...] will contain LOW level cloud fraction + - return_val[1,...] will contain MID level cloud fraction + - return_val[2,...] will contain HIGH level cloud fraction + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa], with the + rightmost dimensions as bottom_top x south_north x west_east + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + relh (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Relative + humidity with the same dimensionality as *pres* + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + cloud fraction array whose leftmost dimension is 3 (LOW=0, MID=1, + HIGH=2). If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.rh` + + """ return _cloudfrac(pres, relh) -@set_alg_metadata(2, "pres_hpa", refvarndims=3, units="degC", +@set_alg_metadata(2, "pres_hpa", refvarndims=3, description="cloud top temperature") -def ctt(pres_hpa, tkel, qv, qcld, height, terrain, qice=None, meta=True): +@convert_units("temp", "c") +def ctt(pres_hpa, tkel, qv, qcld, height, terrain, qice=None, meta=True, + units="degC"): + """Return the cloud top temperature. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres_hpa (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [hPa], with the + rightmost dimensions as bottom_top x south_north x west_east + + Note: + + The units for *psfc_hpa* are [hPa]. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *pres_hpa*. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *pres_hpa*. + + qcld (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Cloud water + vapor mixing ratio in [kg/kg] with the same dimensionality as + *pres_hpa*. + + height (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] with the same dimensionality as + *pres_hpa*. + + terrain (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Terrain height in [m]. This is at least a two-dimensional array + with the same dimensionality as *pres_hpa*, excluding the vertical + (bottom_top/top_bottom) dimension. + + qice (:class:`xarray.DataArray` or :class:`numpy.ndarray`, optional): + Ice mixing ratio in [kg/kg] with the same dimensionality as + *pres_hpa*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'ctt'. Default + is 'degC'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + cloud top temperature. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.cfrac` + + """ # Qice and QCLD need to be in g/kg if qice is None: @@ -141,6 +1153,82 @@ def ctt(pres_hpa, tkel, qv, qcld, height, terrain, qice=None, meta=True): description="radar reflectivity") def dbz(pres, tkel, qv, qr, qs=None, qg=None, use_varint=False, use_liqskin=False, meta=True): + """Return the simulated radar reflectivity. + + This function computes equivalent reflectivity factor [dBZ] at each + model grid point assuming spherical particles of constant density, + with exponential size distributions. This function is based on + "dbzcalc.f" in RIP. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa], with the + rightmost dimensions as bottom_top x south_north x west_east + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *pres*. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *pres_hpa*. + + qr (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Rain water + vapor mixing ratio in [kg/kg] with the same dimensionality as + *pres*. + + qs (:class:`xarray.DataArray` or :class:`numpy.ndarray`, optional): + Snow mixing ratio in [kg/kg] with the same dimensionality as + *pres*. + + qg (:class:`xarray.DataArray` or :class:`numpy.ndarray`, optional): + Graupel mixing ratio in [kg/kg] with the same dimensionality as + *pres*. + + use_varint (:obj:`bool`, optional): When set to False, + the intercept parameters are assumed constant + (as in MM5's Reisner-2 bulk microphysical scheme). + When set to True, the variable intercept + parameters are used as in the more recent version of Reisner-2 + (based on Thompson, Rasmussen, and Manning, 2004, Monthly weather + Review, Vol. 132, No. 2, pp. 519-542.). + + use_liqskin (:obj:`bool`, optional): When set to True, frozen particles + that are at a temperature above freezing are assumed to scatter + as a liquid particle. Set to False to disable. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + simulated radar reflectivity. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.cfrac` + + """ if qs is None: qs = np.zeros(qv.shape, qv.dtype) @@ -157,19 +1245,167 @@ def dbz(pres, tkel, qv, qr, qs=None, qg=None, use_varint=False, @set_alg_metadata(2, "terrain", units="m2 s-2", description="storm relative helicity") -def srhel(u, v, z, terrain, top=3000.0, meta=True): +def srhel(u, v, height, terrain, top=3000.0, meta=True): + """Return the storm relative helicity. + + This function calculates storm relative helicity from WRF ARW output. + SRH (Storm Relative Helicity) is a measure of the potential for cyclonic + updraft rotation in right-moving supercells, and is calculated for the + lowest 1-km and 3-km layers above ground level. There is no clear threshold + value for SRH when forecasting supercells, since the formation of + supercells appears to be related more strongly to the deeper layer + vertical shear. Larger values of 0-3 km SRH (greater than 250 m2 s-2) + and 0-1 km SRH (greater than 100 m2 s-2), however, do suggest an + increased threat of tornadoes with supercells. For SRH, larger values are + generally better, but there are no clear "boundaries" between non-tornadic + and significant tornadic supercells. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + u (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The u + component of the wind that must have at least three dimensions. + The rightmost dimensions are bottom_top x south_north x west_east. + + v (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The v + component of the wind with the same dimensionality as *u*. + + height (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] with the same dimensionality as + *u*. + + terrain (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Terrain height in [m]. This is at least a two-dimensional array + with the same dimensionality as *u*, excluding the bottom_top + dimension. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + top (:obj:`float`): The height of the layer below which helicity is + calculated (meters above ground level). + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + storm relative helicity. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.udhel` + + """ + # u, v get swapped in vertical _u = np.ascontiguousarray(u[...,::-1,:,:]) _v = np.ascontiguousarray(v[...,::-1,:,:]) - _z = np.ascontiguousarray(z[...,::-1,:,:]) + _height = np.ascontiguousarray(height[...,::-1,:,:]) - return _srhel(_u, _v, _z, terrain, top) + return _srhel(_u, _v, _height, terrain, top) @set_alg_metadata(2, "u", refvarndims=3, units="m2 s-2", description="updraft helicity") def udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom=2000.0, top=5000.0, meta=True): + """Return the updraft helicity. + + This function calculates updraft helicity to detect + rotating updrafts. The formula follows Kain et al., 2008, Wea. and + Forecasting, 931-952, but this version has controls for the limits of + integration, *bottom* to *top*, in m AGL. Kain et al used 2000 to 5000 m. + The expected range is 25 to 250 m-2/s-2. Keith Brewster, CAPS/Univ. of + Oklahoma ; March, 2010 + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + zstag (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] that is at least three dimensions with + a staggered vertical dimension. The rightmost dimensions are + bottom_top_stag x south_north x west_east. + + mapfct (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the mass grid. An array of at least + two dimensions, whose rightmost two dimensions must be + south_north x west_east. If this array is more than two dimensions, + they must be the same as *zstag*'s leftmost dimensions. + + u (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The u + component of the wind [m s-1] whose rightmost three dimensions + must be bottom_top x south_north x west_east. The leftmost + dimensions must be the same as zp's leftmost dimensions. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + v (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The v + component of the wind [m s-1] whose rightmost three dimensions + must be bottom_top x south_north x west_east. The leftmost + dimensions must be the same as *zstag*'s leftmost dimensions. + + wstag (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The z + component of the wind [m s-1] with the same dimensionality as + *zstag*. + + dx (:obj:`float`): The distance between x grid points. + + dy (:obj:`float`): The distance between y grid points. + + bottom (:obj:`float`, optional): The bottom limit of integration. + Default is 2000.0. + + top (:obj:`float`, optional): The upper limit of integration. + Default is 5000.0. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + updraft helicity. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.srhel` + + """ return _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top) @@ -178,42 +1414,481 @@ def udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom=2000.0, top=5000.0, stagdim=-1, stagsubvar="vstag", description="absolute vorticity") def avo(ustag, vstag, msfu, msfv, msfm, cor, dx, dy, meta=True): + """Return the absolute vorticity. + + This function returns absolute vorticity [10-5 s-1], which is the sum of + the relative vorticity at each grid point and the Coriolis parameter + at the latitude. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + ustag (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + The u component of the wind in [m s-1] that is at least three + dimensions with a staggered west_east dimension. The rightmost + dimensions are bottom_top x south_north x west_east_stag. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + vstag (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + The v component of the wind in [m s-1] that is at least three + dimensions with a staggered south_north dimension. The rightmost + dimensions are bottom_top x south_north_stag x west_east. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + msfu (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the u-grid that is at least + two dimensions, whose rightmost two dimensions must be + the same as *ustag*. If this array contains more than two + dimensions, they must be the same as *ustag* and *vstag*'s leftmost + dimensions. + + msfv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the v-grid that is at least + two dimensions, whose rightmost two dimensions must be + the same as *vstag*. If this array contains more than two + dimensions, they must be the same as *ustag* and *vstag*'s leftmost + dimensions. + + msfm (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the mass grid that is at least + two dimensions, whose rightmost two dimensions must be + south_north x west_east. If this array contains more than two + dimensions, they must be the same as *ustag* and *vstag*'s leftmost + dimensions. + + cor (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The Coriolis + sine latitude array that is at least + two dimensions, whose dimensions must be the same as *msfm*. + + dx (:obj:`float`): The distance between x grid points. + + dy (:obj:`float`): The distance between y grid points. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + absolute vorticity. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.pvo` + + """ return _avo(ustag, vstag, msfu, msfv, msfm, cor, dx, dy) -@set_alg_metadata(3, "tkel", units="PVU", +@set_alg_metadata(3, "theta", units="PVU", description="potential vorticity") -def pvo(ustag, vstag, tkel, pres, msfu, msfv, msfm, cor, dx, dy, meta=True): - return _pvo(ustag, vstag, tkel, pres, msfu, msfv, msfm, cor, dx, dy) +def pvo(ustag, vstag, theta, pres, msfu, msfv, msfm, cor, dx, dy, meta=True): + """Return the potential vorticity. + + This function calculates the potential vorticity [PVU] at each grid point. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + ustag (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + The u component of the wind in [m s-1] that is at least three + dimensions with a staggered west_east dimension. The rightmost + dimensions are bottom_top x south_north x west_east_stag. + + + vstag (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + The v component of the wind in [m s-1] that is at least three + dimensions with a staggered south_north dimension. The rightmost + dimensions are bottom_top x south_north_stag x west_east. + + theta (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The + potential temperature field [K] whose rightmost dimensions are + bottom_top x south_north x west_east and whose leftmost dimensions + are the same as *ustag*. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa], with the + same dimensions as *theta*. + + msfu (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the u-grid that is at least + two dimensions, whose rightmost two dimensions must be + the same as *ustag*. If this array contains more than two + dimensions, they must be the same as *ustag* and *vstag*'s leftmost + dimensions. + + msfv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the v-grid that is at least + two dimensions, whose rightmost two dimensions must be + the same as *vstag*. If this array contains more than two + dimensions, they must be the same as *ustag* and *vstag*'s leftmost + dimensions. + + msfm (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The map + scale factor on the mass grid that is at least + two dimensions, whose rightmost two dimensions must be + south_north x west_east. If this array contains more than two + dimensions, they must be the same as *ustag* and *vstag*'s leftmost + dimensions. + + cor (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The Coriolis + sine latitude array that is at least + two dimensions, whose dimensions must be the same as *msfm*. + + dx (:obj:`float`): The distance between x grid points. + + dy (:obj:`float`): The distance between y grid points. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + potential vorticity. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.avo` + + """ + return _pvo(ustag, vstag, theta, pres, msfu, msfv, msfm, cor, dx, dy) -@set_alg_metadata(3, "qv", units="K", +@set_alg_metadata(3, "qv", description="equivalent potential temperature") -def eth(qv, tkel, pres, meta=True): +@convert_units("temp", "k") +def eth(qv, tkel, pres, meta=True, units="K"): + """Return the equivalent potential temperature. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] that is at least three dimensions, with + the rightmost dimensions of bottom_top x south_north x west_east. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *qv*. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with the + same dimensionality as *qv*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'eth'. Default + is 'K'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + equivalent potential temperature. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.temp`, :meth:`wrf.wetbulb`, + :meth:`tvirtual` + + """ + return _eth(qv, tkel, pres) -@set_alg_metadata(3, "pres", units="K", +@set_alg_metadata(3, "pres", description="wetbulb temperature") -def wetbulb(pres, tkel, qv, meta=True): +@convert_units("temp", "k") +def wetbulb(pres, tkel, qv, meta=True, units="K"): + """Return the wetbulb temperature. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa], with the + rightmost dimensions as bottom_top x south_north x west_east + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with same dimensionality as *pres*. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as + *pres* + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'twb'. Default + is 'K'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + wetbulb temperature. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.temp`, :meth:`wrf.eth`, + :meth:`tvirtual` + + """ return _wetbulb(pres, tkel, qv) -@set_alg_metadata(3, "tkel", units="K", +@set_alg_metadata(3, "tkel", description="virtual temperature") -def tvirtual(tkel, qv, meta=True): +@convert_units("temp", "k") +def tvirtual(tkel, qv, meta=True, units="K"): + """Return the virtual temperature. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`) Temperature + in [K] with the rightmost dimensions as bottom_top x south_north + x west_east. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as + *tkel* + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'tv'. Default + is 'K'. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + virtual temperature. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`wrf.temp`, :meth:`wrf.eth`, + :meth:`wetbulb` + + """ return _tv(tkel, qv) @set_alg_metadata(3, "qv", units="Pa s-1", description="omega") def omega(qv, tkel, w, pres, meta=True): + """Return omega. + + This function calculates omega (dp/dt) [Pa s-1]. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the rightmost dimensions as + bottom_top x south_north x west_east. + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Temperature + in [K] with the same dimensionality as *qv*. + + w (:class:`xarray.DataArray` or :class:`numpy.ndarray`): The vertical + velocity [m s-1] with the same dimensionality as *qv*. + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa] with the + same dimensionality as *qv*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: Omega. + If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar`, :meth:`uvmet` + + """ return _omega(qv, tkel, w, pres) @set_alg_metadata(2, "pres", refvarndims=3, units="kg m-2", description="precipitable water") def pw(pres, tkel, qv, height, meta=True): + """Return the precipitable water. + + This is the raw computational algorithm and does not extract any variables + from WRF output files. Use :meth:`wrf.getvar` to both extract and compute + diagnostic variables. + + Args: + + pres (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Full + pressure (perturbation + base state pressure) in [Pa], with the + rightmost dimensions as bottom_top x south_north x west_east + + Note: + + This variable must be + supplied as a :class:`xarray.DataArray` in order to copy the + dimension names to the output. Otherwise, default names will + be used. + + tkel (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Temperature + in [K] with the same dimensionality as *pres*. + + qv (:class:`xarray.DataArray` or :class:`numpy.ndarray`): Water vapor + mixing ratio in [kg/kg] with the same dimensionality as *pres* + + height (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + Geopotential height in [m] with the same dimensionality as + *pres*. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Warning: + + The input arrays must not contain any missing/fill values or + :data:`numpy.nan` values. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The precipitable + water [kg m-2]. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + See Also: + + :meth:`wrf.getvar` + + """ tv = _tv(tkel, qv) return _pw(pres, tv, qv, height) diff --git a/src/wrf/config.py b/src/wrf/config.py index 56d636c..0537949 100644 --- a/src/wrf/config.py +++ b/src/wrf/config.py @@ -32,70 +32,126 @@ except ImportError: def xarray_enabled(): + """Return True if xarray is installed and enabled. + + Returns: + + :obj:`bool`: True if xarray is installed and enabled. + + """ global _local_config return _local_config.xarray_enabled def disable_xarray(): + """Disable xarray.""" global _local_config _local_config.xarray_enabled = False def enable_xarray(): + """Enable xarray.""" global _local_config _local_config.xarray_enabled = True def cartopy_enabled(): + """Return True if cartopy is installed and enabled. + + Returns: + + :obj:`bool`: True if cartopy is installed and enabled. + + """ global _local_config return _local_config.cartopy_enabled def enable_cartopy(): + """Enable cartopy.""" global _local_config _local_config.cartopy_enabled = True def disable_cartopy(): + """Disable cartopy.""" global _local_config _local_config.cartopy_enabled = True def basemap_enabled(): + """Return True if basemap is installed and enabled. + + Returns: + + :obj:`bool`: True if basemap is installed and enabled. + + """ global _local_config return _local_config.basemap_enabled def enable_basemap(): + """Enable basemap.""" global _local_config _local_config.basemap_enabled = True def disable_basemap(): + """Disable basemap.""" global _local_config _local_config.basemap_enabled = True def pyngl_enabled(): + """Return True if pyngl is installed and enabled. + + Returns: + + :obj:`bool`: True if pyngl is installed and enabled. + + """ global _local_config return _local_config.pyngl_enabled def enable_pyngl(): + """Enable pyngl.""" global _local_config _local_config.pyngl_enabled = True def disable_pyngl(): + """Disable pyngl.""" global _local_config _local_config.pyngl_enabled = True def set_cache_size(size): + """Set the maximum number of items that the threadsafe cache can retain. + + This cache is primarily used for coordinate variables. + + Args: + + size (:obj:`int`): The number of items to retain in the cache. + + Returns: + + None + + """ global _local_config _local_config.cache_size = size def get_cache_size(): + """Return the maximum number of items that the threadsafe cache can retain. + + Returns: + + :obj:`int`: The maximum number of items the cache can retain. + + """ global _local_config return int(_local_config.cache_size) diff --git a/src/wrf/coordpair.py b/src/wrf/coordpair.py index 787b358..ad42332 100644 --- a/src/wrf/coordpair.py +++ b/src/wrf/coordpair.py @@ -5,7 +5,36 @@ from .py3compat import py2round def _binary_operator(operator): + """Function wrapper for binary operators. + + Args: + + operator (method): The operator to wrap. + + Returns: + + method: An implementation for the *operator* type. + + """ def func(self, other): + """Operator implementation. + + Operator action is performed across the same class attributes when + the *other* object is a :class:`CoordPair`. If the *other* object is + a scalar value, the operator action is performed across all + attributes with the scalar value. + + Args: + + other (:class:`CoordPair` or scalar): A separate :class:`CoordPair` + object or scalar value. + + Returns: + + :class:`CoordPair`: A new :class:`CoordPair` object that is the + result of the operator action. + + """ if isinstance(other, CoordPair): args = [ None if getattr(self, attr) is None or getattr(other, attr) is None @@ -23,7 +52,28 @@ def _binary_operator(operator): def _unary_operator(operator): + """Function wrapper for unary operators. + + Args: + + operator (method): The operator to wrap. + + Returns: + + method: An implementation for the *operator* type. + + """ def func(self): + """Operator implementation. + + Operator action is performed across all class attributes. + + Returns: + + :class:`CoordPair`: A new :class:`CoordPair` object that is the + result of the operator action. + + """ args = [None if getattr(self, attr) is None else getattr(getattr(self, attr), operator)() for attr in ("x", "y", "lat", "lon")] @@ -34,7 +84,30 @@ def _unary_operator(operator): def _cmp_operator(operator): + """Function wrapper for comparison operators. + + Args: + + operator (method): The operator to wrap. + + Returns: + + method: An implementation for the *operator* type. + + """ + def func(self, other): + """Operator implementation. + + Performs a comparison operation across all of the same class + attributes, and returns True if all these operations are True. + + Returns: + + :obj:`boot`: Returns True if all comparisons across class + attributes returns True, otherwise False. + + """ vals = [getattr(getattr(self, attr), operator)(getattr(other, attr)) for attr in ("x", "y", "lat", "lon") if getattr(self, attr) is not None] @@ -45,7 +118,35 @@ def _cmp_operator(operator): class CoordPair(object): + """A class that stores (x, y) and/or (latitude, longitude) + coordinate pairs. + + Most math operators are supplied. When the other operand is a + :class:`CoordPair`, the operation is performed with the same attribute. + When a math operation uses a scalar as the other operand, the + operation is applied across all attributes. + + Attributes: + + x (:obj:`float`): The x-coordinate. + y (:obj:`float`): The y-coordinate. + lat (:obj:`float`): The latitude coordinate. + lon (:obj:`float`): The longitude coordinate. + + + """ def __init__(self, x=None, y=None, lat=None, lon=None): + """Initialize a :class:`CoordPair` object. + + Args: + + x (:obj:`float`, optional): The x-coordinate. + y (:obj:`float`, optional): The y-coordinate. + lat (:obj:`float`, optional): The latitude coordinate. + lon (:obj:`float`, optional): The longitude coordinate. + + + """ self.x = x self.y = y self.lat = lat @@ -72,6 +173,17 @@ class CoordPair(object): def xy_str(self, fmt="{:.4f}, {:.4f}"): + """Return a :obj:`str` for the (x,y) coordinate pair. + + Args: + + fmt (:obj:`str`): The format string. Default is '{:.4f}, {:.4f}' + + Returns: + + :obj:`str`: A string for the (x,y) coordinate pair + + """ if self.x is None or self.y is None: return None @@ -79,15 +191,38 @@ class CoordPair(object): def latlon_str(self, fmt="{:.4f}, {:.4f}"): + """Return a :obj:`str` for the (latitude, longitude) coordinate pair. + + Args: + + fmt (:obj:`str`): The format string. Default is '{:.4f}, {:.4f}' + + Returns: + + :obj:`str`: A string for the (latitude, longitude) coordinate pair + + """ if self.lat is None or self.lon is None: return None return fmt.format(self.lat, self.lon) - def __round__(self, d=None): + def __round__(self, ndigits=None): + """Return a new :class:`CoordPair` object with all coordinate values + rounded to the nearest integer. + + Args: + + ndigits (:obj:`int`): The number of digits. + + Returns: + + :class:`CoordPair`: A CoordPair object. + + """ args = [None if getattr(self, attr) is None - else py2round(getattr(self, attr), d) + else py2round(getattr(self, attr), ndigits) for attr in ("x", "y", "lat", "lon")] return CoordPair(*args) diff --git a/src/wrf/ctt.py b/src/wrf/ctt.py index f0d671b..51c2b49 100644 --- a/src/wrf/ctt.py +++ b/src/wrf/ctt.py @@ -17,14 +17,67 @@ from .util import extract_vars description="cloud top temperature", MemoryOrder="XY") @convert_units("temp", "c") -def get_ctt(wrfnc, timeidx=0, method="cat", +def get_ctt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="degC"): """Return the cloud top temperature. + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'ctt'. Default + is 'degC'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + cloud top temperature. + 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. + """ varnames = ("T", "P", "PB", "PH", "PHB", "HGT", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] @@ -36,7 +89,7 @@ def get_ctt(wrfnc, timeidx=0, method="cat", haveqci = 1 try: - icevars = extract_vars(wrfnc, timeidx, "QICE", + icevars = extract_vars(wrfin, timeidx, "QICE", method, squeeze, cache, meta=False, _key=_key) except KeyError: @@ -46,7 +99,7 @@ def get_ctt(wrfnc, timeidx=0, method="cat", qice = icevars["QICE"] * 1000.0 #g/kg try: - cldvars = extract_vars(wrfnc, timeidx, "QCLOUD", + cldvars = extract_vars(wrfin, timeidx, "QCLOUD", method, squeeze, cache, meta=False, _key=_key) except KeyError: diff --git a/src/wrf/dbz.py b/src/wrf/dbz.py index 8315557..d7372dc 100755 --- a/src/wrf/dbz.py +++ b/src/wrf/dbz.py @@ -13,21 +13,75 @@ from .metadecorators import copy_and_set_metadata @copy_and_set_metadata(copy_varname="T", name="dbz", description="radar reflectivity", units="dBZ") -def get_dbz(wrfnc, timeidx=0, method="cat", +def get_dbz(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, use_varint=False, use_liqskin=False): - """ Return the dbz + """Return the simulated radar reflectivity. - use_varint - do variable intercept (if False, constants are used. Otherwise, - intercepts are calculated using a technique from Thompson, Rasmussen, - and Manning (2004, Monthly Weather Review, Vol. 132, No. 2, pp. 519-542.) + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. - use_liqskin - do liquid skin for snow (frozen particles above freezing scatter - as liquid) + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + use_varint (:obj:`bool`, optional): When set to False, + the intercept parameters are assumed constant + (as in MM5's Reisner-2 bulk microphysical scheme). + When set to True, the variable intercept + parameters are used as in the more recent version of Reisner-2 + (based on Thompson, Rasmussen, and Manning, 2004, Monthly weather + Review, Vol. 132, No. 2, pp. 519-542.). + + use_liqskin (:obj:`bool`, optional): When set to True, frozen particles + that are at a temperature above freezing are assumed to scatter + as a liquid particle. Set to False to disable. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The simulated + radar reflectivity. + 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. """ varnames = ("T", "P", "PB", "QVAPOR", "QRAIN") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] @@ -36,7 +90,7 @@ def get_dbz(wrfnc, timeidx=0, method="cat", qr = ncvars["QRAIN"] try: - snowvars = extract_vars(wrfnc, timeidx, "QSNOW", + snowvars = extract_vars(wrfin, timeidx, "QSNOW", method, squeeze, cache, meta=False, _key=_key) except KeyError: @@ -45,7 +99,7 @@ def get_dbz(wrfnc, timeidx=0, method="cat", qs = snowvars["QSNOW"] try: - graupvars = extract_vars(wrfnc, timeidx, "QGRAUP", + graupvars = extract_vars(wrfin, timeidx, "QGRAUP", method, squeeze, cache, meta=False, _key=_key) except KeyError: @@ -70,10 +124,77 @@ def get_dbz(wrfnc, timeidx=0, method="cat", description="maximum radar reflectivity", units="dBZ", MemoryOrder="XY") -def get_max_dbz(wrfnc, timeidx=0, method="cat", +def get_max_dbz(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, - do_varint=False, do_liqskin=False): - return np.amax(get_dbz(wrfnc, timeidx, method, squeeze, cache, meta, - _key, do_varint, do_liqskin), + use_varint=False, use_liqskin=False): + """Return the maximum simulated radar reflectivity. + + This function returns the maximum reflectivity found in the column for + each grid point. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + use_varint (:obj:`bool`, optional): When set to False, + the intercept parameters are assumed constant + (as in MM5's Reisner-2 bulk microphysical scheme). + When set to True, the variable intercept + parameters are used as in the more recent version of Reisner-2 + (based on Thompson, Rasmussen, and Manning, 2004, Monthly weather + Review, Vol. 132, No. 2, pp. 519-542.). + + use_liqskin (:obj:`bool`, optional): When set to True, frozen particles + that are at a temperature above freezing are assumed to scatter + as a liquid particle. Set to False to disable. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The maximum + simulated radar reflectivity. + 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. + + """ + return np.amax(get_dbz(wrfin, timeidx, method, squeeze, cache, meta, + _key, use_varint, use_liqskin), axis=-3) diff --git a/src/wrf/decorators.py b/src/wrf/decorators.py index c13af3b..25301e2 100644 --- a/src/wrf/decorators.py +++ b/src/wrf/decorators.py @@ -17,13 +17,23 @@ if xarray_enabled(): from xarray import DataArray def convert_units(unit_type, alg_unit): - """A decorator that applies unit conversion to a function's output array. + """A decorator that converts the units from the wrapped function's output. - Arguments: + The desired units are determined from the wrapped function's arguments. - - unit_type - the unit type category (wind, pressure, etc) - - alg_unit - the units that the function returns by default + Args: + unit_type (:obj:`str`): The unit type. Choices are: 'wind', + 'pressure', 'temp', or 'height'. + + alg_unit (:obj:`str`): The units returned by the wrapped function, + which is usually the units returned by the Fortran routine. + + Returns: + + :class:`numpy.ndarray`: The wrapped function's output in the desired + units. + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): @@ -39,10 +49,15 @@ def convert_units(unit_type, alg_unit): return func_wrapper -def _calc_out_dims(outvar, left_dims): - left_dims = [x for x in left_dims] - right_dims = [x for x in outvar.shape] - return left_dims + right_dims +#def _calc_out_dims(outvar, left_dims): +# """ +# +# """ +# #left_dims = [x for x in left_dims] +# #right_dims = [x for x in outvar.shape] +# #return left_dims + right_dims +# +# return left_dims + outvar.shape def left_iteration(ref_var_expected_dims, @@ -55,33 +70,63 @@ def left_iteration(ref_var_expected_dims, outviews="outview", alg_dtype=np.float64, cast_output=True): - """Decorator to handle iterating over leftmost dimensions when using - multiple files and/or multiple times. - - Arguments: - - ref_var_expected_dims - the number of dimensions that the Fortran - algorithm is expecting for the reference variable - - ref_var_right_ndims - the number of dims from the right to keep for - the reference variable when making the output. Can also be a - combine_dims instance if the sizes are determined from multiple - variables. - - insert_dims - a sequence of dimensions to insert between the left - dimenions (e.g. time) and the kept right dimensions - - ref_var_idx - the index in args used as the reference variable for - calculating leftmost dimensions - - ref_var_name - the keyword argument name for kwargs used as the - reference varible for calculating leftmost dimensions - - alg_out_fixed_dims - additional fixed dimension sizes for the - numerical algorithm (e.g. uvmet has a fixed left dimsize of 2) - - ignore_args - indexes of any arguments which should be passed - directly without any slicing - - ignore_kargs - keys of any keyword arguments which should be passed - directly without slicing - - outviews - a single key or sequence of keys indicating the keyword - argument used as the output variable(s) - - algtype - the data type used in the numerical routine - - cast_output - cast the final output to the ref_var data type - + """A decorator to handle iterating over the leftmost dimensions. + + For example, if a wrapped function works with three-dimensional arrays, but + the variables include a 4th leftmost dimension for 'Time', this decorator + will iterate over all times, call the 3D Fortran routine, and aggregate the + results in to a 4D output array. + + It is also important to note that the final output array is allocated + first, and then views are passed to the wrapped function so that values + do not need to get copied in to the final output array. + + Args: + + ref_var_expected_dims (:obj:`int`): The number of dimensions that the + Fortran routine is expecting for the reference variable. + + ref_var_right_ndims (:obj:`int`): The number of dimensions from the + right to keep for the reference variable when making the output. + Can also be a :class:`combine_dims` object if the sizes are + determined from multiple variables. + + insert_dims (sequence of :obj:`int`, optional): A sequence of + dimensions to insert between the left dimensions (e.g. time) and + the kept right dimensions. Default is None. + + ref_var_idx (:obj:`int`, optional): The index in the wrapped function's + positional arguments to be used as the reference variable for + determining the leftmost dimensions. Must be specified if + *ref_var_name* is None. Default is None. + + ref_var_name (:obj:`str`, optional): The keyword argument name for the + wrapped function's keyword arguments to be used as the reference + variable for calculating the leftmost dimensions. Must be + specified if *ref_var_idx* is None. Default is None. + + ignore_args (sequence of :obj:`int`): Indexes of any arguments that + should be ignored when creating the sliced views that are + passed to the Fortran routine. + + ignore_kargs (sequence of :obj:`str`): Keys of any keyword arguments + that should be ignored when creating the sliced views that are + passed to the Fortran routine. + + outviews (:obj:`str` or a sequence): A single key or sequence of keys + that indicate the wrapped function's keyword argument to use + as the output variable(s) in the wrapped function. + + alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used + in the wrapped function. + + cast_output (:obj:`bool`): Set to True to cast the wrapped function's + output to the same type as the reference variable. + + Returns: + + :class:`numpy.ndarray`: The aggregated output array that includes + all extra leftmost dimensions found in the reference variable. """ @wrapt.decorator @@ -204,10 +249,39 @@ def left_iteration(ref_var_expected_dims, def cast_type(ref_idx=0, arg_idxs=None, karg_names=None, - alg_dtype=np.float64, - outviews="outview"): - """Decorator to handle casting to/from required dtype used in - numerical routine. + alg_dtype=np.float64, outviews="outview"): + """A decorator to handle type casting. + + This decorator is used to cast variables to and from the required + :class:`numpy.dtype` used in the wrapped function. + + Args: + + ref_idx (:obj:`int`, optional): The index in the wrapped function's + positional arguments to be used as the reference variable for + determining the :class:`numpy.dtype` to return. Default is 0. + + arg_idxs (sequence of :obj:`int`, optional): A sequence of indexes in the + wrapped function's positional arguments that indicate which + arguments to cast. Must be specified if *karg_names* is None. + Default is None. + + karg_names (sequence of :obj:`str`): A sequence of keyword arguments + in the wrapped function's keyword arguments that indicate the + arguments to cast. Must be specified if *arg_idxs* is None. + Default is None. + + alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used + in the wrapped function. + + outviews (:obj:`str` or a sequence): A single key or sequence of keys + that indicate the wrapped function's keyword argument to use + as the output variable(s) in the wrapped function. + + Returns: + + :class:`numpy.ndarray`: The wrapped function's output cast to the + same :class:`numpy.dtype` as the reference variable. """ @wrapt.decorator @@ -255,6 +329,27 @@ def cast_type(ref_idx=0, arg_idxs=None, karg_names=None, def _extract_and_transpose(arg, do_transpose): + """Return a transposed view of the :class:`numpy.ndarray` inside of a + :class:`xarray.DataArray` object. + + If the *arg* parameter is not a :class:`xarray.DataArray` object, then + *arg* is returned. + + Args: + + arg (:class:`xarray.DataArray` or :obj:`object`): Can be any object + type. + + do_transpose: Set to False to only extract the variable. When True, + the extracted array will also be transposed to a Fortran view if + it is not already Fortran contiguous. + + Returns: + + :class:`numpy.ndarray`: A numpy array. If *do_transpose* is True, + the numpy array will also be a Fortran contiguous view. + + """ if xarray_enabled(): if isinstance(arg, DataArray): @@ -269,8 +364,25 @@ def _extract_and_transpose(arg, do_transpose): def extract_and_transpose(do_transpose=True, outviews="outview"): - """Decorator to extract the data array from a DataArray and also - transposes the view of the data if the data is not fortran contiguous. + """A decorator to extract the data array from a :class:`xarray.DataArray` + + This decorator also transposes the view of the data to Fortran + contiguous if *do_transpose* is True. + + Args: + + do_transpose: Set to False to only extract the variable. When True, + the extracted array will also be transposed to a Fortran view if + it is not already Fortran contiguous. + + outviews (:obj:`str` or a sequence): A single key or sequence of keys + that indicate the wrapped function's keyword argument to use + as the output variable(s) in the wrapped function. + + Returns: + + :class:`numpy.ndarray`: A numpy array. If *do_transpose* is True, + the numpy array will also be a Fortran contiguous view. """ @wrapt.decorator @@ -311,16 +423,37 @@ def extract_and_transpose(do_transpose=True, outviews="outview"): def check_args(refvaridx, refvarndim, rightdims, stagger=None, refstagdim=None): - """ + """A decorator to check that the wrapped function's arguments are valid. + + An exception is raised when an invalid argument is found. + + Args: + + refvaridx (:obj:`int`): The wrapped function's positional argument + index to use as the reference variable. + + refvarndim (:obj:`int`): The number of dimensions for the reference + variable that is expected by the wrapped function. + + rightdims (sequence of :obj:`int`): The expected number of right + dimensions for each argument. - refvaridx - reference variable to check for presence of left dimensions - refvarndim - the number of dimensions for the references variable - expected by the fortran routine - rightdims - the expected number of right dimensions for each argument - stagger - the dimension which is staggered for each argument. Use - None to indicate no staggering. - refstagdim - If the reference variable is staggered, indicate the dim that - is staggered + stagger (sequence of :obj:`int` or :obj:`None`, optional): The + dimension that is staggered for each argument in the wrapped + function. Use :obj:`None` in the sequence to indicate no + staggering for that argument. Default is None. + + refstagdim (:obj:`int`, optional): The staggered dimension for the + reference variable, if applicable. Default is None. + + Returns: + + None + + Raises: + + :class:`ValueError`: Raised when an invalid argument is detected. + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): diff --git a/src/wrf/destag.py b/src/wrf/destag.py index e8e4c45..169481c 100755 --- a/src/wrf/destag.py +++ b/src/wrf/destag.py @@ -8,14 +8,31 @@ from .metadecorators import set_destag_metadata @set_destag_metadata() @extract_and_transpose(do_transpose=False) def destagger(var, stagger_dim, meta=False): - """ De-stagger the variable. - - Arguments: - - var is a numpy array for the variable - - stagger_dim is the dimension of the numpy array to de-stagger - (e.g. 0, 1, 2). Note: negative values are acceptable to choose - a dimensions from the right hand side (e.g. -1, -2, -3) - - meta - set to True to include 'var' metadata + """Return the variable on the unstaggered grid. + + This function destaggers the variable by taking the average of the + values located on either side of the grid box. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A variable + on a staggered grid. + + stagger_dim (:obj:`int`): The dimension index to destagger. + Negative values can be used to choose dimensions referenced + from the right hand side (-1 is the rightmost dimension). + + meta (:obj:`bool`, optional): Set to False to disable metadata and + return :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: + The destaggered variable. 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. """ var_shape = var.shape diff --git a/src/wrf/dewpoint.py b/src/wrf/dewpoint.py index 95bc4b9..87914ed 100755 --- a/src/wrf/dewpoint.py +++ b/src/wrf/dewpoint.py @@ -11,11 +11,67 @@ from .util import extract_vars @copy_and_set_metadata(copy_varname="QVAPOR", name="td", description="dew point temperature") @convert_units("temp", "c") -def get_dp(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_dp(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="degC"): + """Return the dewpoint temperature. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'td'. Default + is 'degC'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + dewpoint temperature. + 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. + + """ varnames=("P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) p = ncvars["P"] @@ -32,10 +88,66 @@ def get_dp(wrfnc, timeidx=0, method="cat", squeeze=True, @copy_and_set_metadata(copy_varname="Q2", name="td2", description="2m dew point temperature") @convert_units("temp", "c") -def get_dp_2m(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_dp_2m(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="degC"): + """Return the 2m dewpoint temperature. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'td2'. Default + is 'degC'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + 2m dewpoint temperature. + 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. + + """ varnames=("PSFC", "Q2") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) # Algorithm requires hPa diff --git a/src/wrf/extension.py b/src/wrf/extension.py index d4ea1cb..2781ddc 100755 --- a/src/wrf/extension.py +++ b/src/wrf/extension.py @@ -21,25 +21,52 @@ from .specialdec import (uvmet_left_iter_nocopy, cape_left_iter, cloudfrac_left_iter) class FortranError(Exception): + """Raised when an error occurs in a Fortran routine.""" def __init__(self, message=None): + """Initialize a :class:`wrf.FortranError` objection. + + Args: + + message (:obj:`str`): The error message. + + """ self._msg = message def __str__(self): return self._msg def __call__(self, message): + """Callable method to make the exception object raise itself. + + This allows the exception to be thrown from inside Fortran routines + by using f2py's callback mechanism. This is no longer used within + wrf-python, but may be useful to other users. + + See Also: + + `f2py doc `_ + + """ raise self.__class__(message) +# The routines below are thin wrappers around the Fortran functions. These +# are not meant to be called by end users. Use the public API instead for +# that purpose. # IMPORTANT! Unless otherwise noted, all variables used in the routines -# below assume that fortran-ordered views are being used. This allows -# f2py to pass the array pointers directly to the fortran routine. +# below assume that Fortran-ordered views are being used. This allows +# f2py to pass the array pointers directly to the Fortran routine. @check_args(0, 3, (3, 3, None, None)) @left_iteration(3, 2, ref_var_idx=0, ignore_args=(2,3)) @cast_type(arg_idxs=(0,1)) @extract_and_transpose() def _interpz3d(field3d, z, desiredloc, missingval, outview=None): + """Wrapper for dinterp3dz. + + Located in wrf_user.f90. + + """ if outview is None: outview = np.empty(field3d.shape[0:2], np.float64, order="F") @@ -57,6 +84,11 @@ def _interpz3d(field3d, z, desiredloc, missingval, outview=None): @cast_type(arg_idxs=(0,1)) @extract_and_transpose() def _interp2dxy(field3d, xy, outview=None): + """Wrapper for dinterp2dxy. + + Located in wrf_user.f90. + + """ if outview is None: outview = np.empty((xy.shape[-1], field3d.shape[-1]), np.float64, order="F") @@ -72,7 +104,11 @@ def _interp2dxy(field3d, xy, outview=None): @cast_type(arg_idxs=(0,1,2)) @extract_and_transpose() def _interp1d(v_in, z_in, z_out, missingval, outview=None): + """Wrapper for dinterp1d. + + Located in wrf_user.f90. + """ if outview is None: outview = np.empty_like(z_out) @@ -90,6 +126,14 @@ def _interp1d(v_in, z_in, z_out, missingval, outview=None): @cast_type(arg_idxs=(0,)) @extract_and_transpose(do_transpose=False) def _vertcross(field3d, xy, var2dz, z_var2d, missingval, outview=None): + """Return the vertical cross section. + + This routine was originally written in scripted NCL code and doesn't + directly wrap a Fortran routine. + + Located in WRFUserARW.ncl. + + """ # Note: This is using C-ordering if outview is None: outview = np.empty((z_var2d.shape[0], xy.shape[0]), dtype=var2dz.dtype) @@ -107,6 +151,14 @@ def _vertcross(field3d, xy, var2dz, z_var2d, missingval, outview=None): @cast_type(arg_idxs=(0,)) @extract_and_transpose(do_transpose=False) def _interpline(field2d, xy, outview=None): + """Return the two-dimensional field interpolated to a line. + + This routine was originally written in scripted NCL code and doesn't + directly wrap a Fortran routine. + + Located in WRFUserARW.ncl. + + """ # Note: This is using C-ordering if outview is None: outview = np.empty(xy.shape[0], dtype=field2d.dtype) @@ -127,7 +179,11 @@ def _interpline(field2d, xy, outview=None): @cast_type(arg_idxs=(0,1,2,3)) @extract_and_transpose() def _slp(z, t, p, q, outview=None): + """Wrapper for dcomputeseaprs. + Located in wrf_user.f90. + + """ t_surf = np.zeros(z.shape[0:2], np.float64, order="F") t_sea_level = np.zeros(z.shape[0:2], np.float64, order="F") level = np.zeros(z.shape[0:2], np.int32, order="F") @@ -160,6 +216,11 @@ def _slp(z, t, p, q, outview=None): @cast_type(arg_idxs=(0,1)) @extract_and_transpose() def _tk(pressure, theta, outview=None): + """Wrapper for dcomputetk. + + Located in wrf_user.f90. + + """ # No need to transpose here since operations on 1D array shape = pressure.shape if outview is None: @@ -177,6 +238,11 @@ def _tk(pressure, theta, outview=None): @cast_type(arg_idxs=(0,1)) @extract_and_transpose() def _td(pressure, qv_in, outview=None): + """Wrapper for dcomputetd. + + Located in wrf_user.f90. + + """ shape = pressure.shape if outview is None: outview = np.empty_like(pressure) @@ -194,6 +260,11 @@ def _td(pressure, qv_in, outview=None): @cast_type(arg_idxs=(0,1,2)) @extract_and_transpose() def _rh(qv, q, t, outview=None): + """Wrapper for dcomputerh. + + Located in wrf_user.f90. + + """ shape = qv.shape if outview is None: outview = np.empty_like(qv) @@ -216,6 +287,11 @@ def _rh(qv, q, t, outview=None): @cast_type(arg_idxs=(0,1,2,3,4,5)) @extract_and_transpose() def _avo(u, v, msfu, msfv, msfm, cor, dx, dy, outview=None): + """Wrapper for dcomputeabsvort. + + Located in wrf_pvo.f90. + + """ if outview is None: outshape = (v.shape[0],) + u.shape[1:] outview = np.empty(outshape, np.float64, order="F") @@ -240,6 +316,11 @@ def _avo(u, v, msfu, msfv, msfm, cor, dx, dy, outview=None): @cast_type(arg_idxs=(0,1,2,3,4,5,6,7)) @extract_and_transpose() def _pvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy, outview=None): + """Wrapper for dcomputepv. + + Located in wrf_pvo.f90. + + """ if outview is None: outview = np.empty_like(prs) @@ -263,6 +344,11 @@ def _pvo(u, v, theta, prs, msfu, msfv, msfm, cor, dx, dy, outview=None): @cast_type(arg_idxs=(0,1,2)) @extract_and_transpose() def _eth(qv, tk, p, outview=None): + """Wrapper for deqthecalc. + + Located in eqthecalc.f90. + + """ if outview is None: outview = np.empty_like(qv) @@ -280,6 +366,11 @@ def _eth(qv, tk, p, outview=None): def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, umissing=Constants.DEFAULT_FILL, vmissing=Constants.DEFAULT_FILL, uvmetmissing=Constants.DEFAULT_FILL, outview=None): + """Wrapper for dcomputeuvmet. + + Located in wrf_user.f90. + + """ longca = np.zeros(lat.shape[0:2], np.float64, order="F") longcb = np.zeros(lon.shape[0:2], np.float64, order="F") rpd = Constants.PI/180. @@ -312,6 +403,11 @@ def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, @cast_type(arg_idxs=(0,1,2,3)) @extract_and_transpose() def _omega(qv, tk, w, p, outview=None): + """Wrapper for omgcalc. + + Located in wrf_rip_phys_routines.f90. + + """ if outview is None: outview = np.empty_like(qv) @@ -329,6 +425,11 @@ def _omega(qv, tk, w, p, outview=None): @cast_type(arg_idxs=(0,1)) @extract_and_transpose() def _tv(tk, qv, outview=None): + """Wrapper for virtual_temp. + + Located in wrf_rip_phys_routines.f90. + + """ if outview is None: outview = np.empty_like(tk) @@ -344,6 +445,11 @@ def _tv(tk, qv, outview=None): @cast_type(arg_idxs=(0,1,2)) @extract_and_transpose() def _wetbulb(p, tk, qv, psafile=psafilepath(), outview=None): + """Wrapper for wetbulbcalc. + + Located in wrf_rip_phys_routines.f90. + + """ if outview is None: outview = np.empty_like(p) @@ -369,6 +475,11 @@ def _wetbulb(p, tk, qv, psafile=psafilepath(), outview=None): @cast_type(arg_idxs=(0,1,2,3)) @extract_and_transpose() def _srhel(u, v, z, ter, top, outview=None): + """Wrapper for dcalrelhl. + + Located in wrf_relhl.f90. + + """ if outview is None: outview = np.empty_like(ter) @@ -387,6 +498,11 @@ def _srhel(u, v, z, ter, top, outview=None): @cast_type(arg_idxs=(0,1,2,3,4)) @extract_and_transpose() def _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top, outview=None): + """Wrapper for dcalcuh. + + Located in calc_uh.f90. + + """ if outview is None: outview = np.empty_like(mapfct) @@ -416,7 +532,11 @@ def _udhel(zstag, mapfct, u, v, wstag, dx, dy, bottom, top, outview=None): @cast_type(arg_idxs=(0,1,2,3)) @extract_and_transpose() def _pw(p, tv, qv, ht, outview=None): + """Wrapper for dcomputepw. + + Located in wrf_pw.f90. + """ if outview is None: outview = np.empty(p.shape[0:2], p.dtype, order="F") @@ -434,6 +554,11 @@ def _pw(p, tv, qv, ht, outview=None): @cast_type(arg_idxs=(0,1,2,3,4,5)) @extract_and_transpose() def _dbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin, outview=None): + """Wrapper for calcdbz. + + Located in wrf_user_dbz.f90. + + """ if outview is None: outview = np.empty_like(p) @@ -457,7 +582,11 @@ def _dbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin, outview=None): @extract_and_transpose(outviews=("capeview", "cinview")) def _cape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow, psafile=psafilepath(), capeview=None, cinview=None): + """Wrapper for dcapecalc3d. + Located in rip_cape.f90. + + """ if capeview is None: capeview = np.zeros(p_hpa.shape[0:3], p_hpa.dtype, order="F") @@ -493,7 +622,11 @@ def _cape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow, @cast_type(arg_idxs=(0, 1), outviews=("lowview", "medview", "highview")) @extract_and_transpose(outviews=("lowview", "medview", "hightview")) def _cloudfrac(p, rh, lowview=None, medview=None, highview=None): + """Wrapper for dcloudfrace. + + Located in wrf_cloud_fracf.f90. + """ if lowview is None: lowview = np.zeros(p.shape[0:2], p.dtype, order="F") @@ -516,7 +649,11 @@ def _lltoxy(map_proj, truelat1, truelat2, stdlon, lat1, lon1, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, loninc, lat, lon, outview=None): + """Wrapper for dlltoij. + Located in wrf_user_latlon_routines.f90. + + """ if outview is None: outview = np.zeros((2), dtype=np.float64, order="F") @@ -552,7 +689,11 @@ def _lltoxy(map_proj, truelat1, truelat2, stdlon, def _xytoll(map_proj, truelat1, truelat2, stdlon, lat1, lon1, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, loninc, x, y, outview=None): + """Wrapper for dijtoll. + + Located in wrf_user_latlon_routines.f90. + """ if outview is None: outview = np.zeros((2), dtype=np.float64, order="F") @@ -590,6 +731,11 @@ def _xytoll(map_proj, truelat1, truelat2, stdlon, lat1, lon1, @cast_type(arg_idxs=(0,1,2,3,4,5,6)) @extract_and_transpose() def _ctt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci, outview=None): + """Wrapper for wrfcttcalc. + + Located in wrf_fctt.f90. + + """ if outview is None: outview = np.empty_like(ter) @@ -611,6 +757,11 @@ def _ctt(p_hpa, tk, qice, qcld, qv, ght, ter, haveqci, outview=None): @cast_type(arg_idxs=(0,)) @extract_and_transpose() def _smooth2d(field, passes, outview=None): + """Wrapper for dfilter2d. + + Located in wrf_user.f90. + + """ # Unlike NCL, this routine will not modify the values in place, but # copies the original data before modifying it. @@ -639,6 +790,11 @@ def _smooth2d(field, passes, outview=None): @cast_type(arg_idxs=(0,1,2)) @extract_and_transpose() def _monotonic(var, lvprs, coriolis, idir, delta, icorsw, outview=None): + """Wrapper for wrf_monotonic. + + Located in wrf_vinterp.f90. + + """ # If icorsw is not 0, then the input variable might get modified by the # fortran routine. We don't want this, so make a copy and pass that on. var = var.copy(order="A") if icorsw != 0 else var @@ -667,7 +823,11 @@ def _monotonic(var, lvprs, coriolis, idir, delta, icorsw, outview=None): def _vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp, vcarray, interp_levels, icase, extrap, vcor, logp, missing, outview=None): + """Wrapper for wrf_vintrp. + Located in wrf_vinterp.f90. + + """ if outview is None: outdims = field.shape[0:2] + interp_levels.shape outview = np.empty(outdims, field.dtype, order="F") @@ -700,10 +860,3 @@ def _vintrp(field, pres, tk, qvp, ght, terrain, sfp, smsfp, return result - - - - - - - diff --git a/src/wrf/geoht.py b/src/wrf/geoht.py index 0f78ec7..d99dd88 100755 --- a/src/wrf/geoht.py +++ b/src/wrf/geoht.py @@ -7,19 +7,79 @@ from .decorators import convert_units from .metadecorators import set_height_metadata from .util import extract_vars, either -def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, +def _get_geoht(wrfin, timeidx, method="cat", squeeze=True, cache=None, meta=True, _key=None, height=True, msl=True): - """Return the geopotential in units of m2 s-2 if height is False, - otherwise return the geopotential height in meters. If height is True, - then if msl is True the result will be in MSL, otherwise AGL (the terrain - height is subtracted). + """Return the geopotential or geopotential height. + + If *height* is False, then geopotential is returned in units of + [m2 s-2]. If *height* is True, then geopotential height is + returned in units of [m]. If *msl* is True, then geopotential height + is return as Mean Sea Level (MSL). If *msl* is False, then geopotential + height is returned as Above Ground Level (AGL). + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + height (:obj:`bool`, optional): Set to True to return geopotential + height instead of geopotential. Default is True. + + msl (:obj:`bool`, optional): Set to True to return geopotential height + as Mean Sea Level (MSL). Set to False to return the + geopotential height as Above Ground Level (AGL) by subtracting + the terrain height. Default is True. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + geopotential or geopotential height. + 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. """ - varname = either("PH", "GHT")(wrfnc) + varname = either("PH", "GHT")(wrfin) if varname == "PH": - ph_vars = extract_vars(wrfnc, timeidx, ("PH", "PHB", "HGT"), + ph_vars = extract_vars(wrfin, timeidx, ("PH", "PHB", "HGT"), method, squeeze, cache, meta=False, _key=_key) ph = ph_vars["PH"] @@ -28,7 +88,7 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, geopt = ph + phb geopt_unstag = destagger(geopt, -3) else: - ght_vars = extract_vars(wrfnc, timeidx, ("GHT", "HGT_M"), + ght_vars = extract_vars(wrfin, timeidx, ("GHT", "HGT_M"), method, squeeze, cache, meta=False, _key=_key) geopt_unstag = ght_vars["GHT"] * Constants.G @@ -51,18 +111,137 @@ def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, @set_height_metadata(geopt=True) -def get_geopt(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_geopt(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - return _get_geoht(wrfnc, timeidx, method, squeeze, cache, meta, _key, + """Return the geopotential. + + The geopotential is returned in units of [m2 s-2]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + geopotential. + 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. + + """ + return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key, False, True) @set_height_metadata(geopt=False) @convert_units("height", "m") -def get_height(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_height(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, msl=True, units="m"): + """Return the geopotential height. + + If *msl* is True, then geopotential height is returned as Mean Sea Level + (MSL). If *msl* is False, then geopotential height is returned as + Above Ground Level (AGL) by subtracting the terrain height. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + msl (:obj:`bool`, optional): Set to True to return geopotential height + as Mean Sea Level (MSL). Set to False to return the + geopotential height as Above Ground Level (AGL) by subtracting + the terrain height. Default is True. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'z'. Default + is 'm'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + geopotential height. + 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. + + """ - return _get_geoht(wrfnc, timeidx, method, squeeze, cache, meta, _key, + return _get_geoht(wrfin, timeidx, method, squeeze, cache, meta, _key, True, msl) diff --git a/src/wrf/helicity.py b/src/wrf/helicity.py index 67335ff..2d2b6a8 100755 --- a/src/wrf/helicity.py +++ b/src/wrf/helicity.py @@ -12,11 +12,68 @@ from .metadecorators import copy_and_set_metadata @copy_and_set_metadata(copy_varname="HGT", name="srh", description="storm relative helicity", units="m2 s-2") -def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_srh(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, top=3000.0): + """Return the storm relative helicity. + + The *top* argument specifies the top of the integration in [m]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + top (:obj:`float`, optional): The top of the integration in [m]. + Default is 3000.0. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + storm relative helicity. + 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. + + """ # Top can either be 3000 or 1000 (for 0-1 srh or 0-3 srh) - ncvars = extract_vars(wrfnc, timeidx, ("HGT", "PH", "PHB"), + ncvars = extract_vars(wrfin, timeidx, ("HGT", "PH", "PHB"), method, squeeze, cache, meta=False, _key=_key) @@ -25,13 +82,13 @@ def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True, phb = ncvars["PHB"] # As coded in NCL, but not sure this is possible - varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("U", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = destagger(u_vars[varname], -1) - varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("V", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = destagger(v_vars[varname], -2) @@ -52,11 +109,73 @@ def get_srh(wrfnc, timeidx=0, method="cat", squeeze=True, @copy_and_set_metadata(copy_varname="MAPFAC_M", name="updraft_helicity", description="updraft helicity", units="m2 s-2") -def get_uh(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_uh(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, bottom=2000.0, top=5000.0): - ncvars = extract_vars(wrfnc, timeidx, ("W", "PH", "PHB", "MAPFAC_M"), + """Return the updraft helicity. + + The *bottom* and *top* arguments specify the bottom and top limits + for the integration in [m]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + bottom (:obj:`float`, optional): The bottom limit for the integration + in [m]. Default is 2000.0. + + top (:obj:`float`, optional): The top limit for the integration in [m]. + Default is 5000.0. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + updraft helicity. + 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. + + """ + + ncvars = extract_vars(wrfin, timeidx, ("W", "PH", "PHB", "MAPFAC_M"), method, squeeze, cache, meta=False, _key=None) wstag = ncvars["W"] @@ -64,18 +183,18 @@ def get_uh(wrfnc, timeidx=0, method="cat", squeeze=True, phb = ncvars["PHB"] mapfct = ncvars["MAPFAC_M"] - attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) + attrs = extract_global_attrs(wrfin, attrs=("DX", "DY")) dx = attrs["DX"] dy = attrs["DY"] # As coded in NCL, but not sure this is possible - varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("U", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = destagger(u_vars[varname], -1) - varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("V", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = destagger(v_vars[varname], -2) diff --git a/src/wrf/interp.py b/src/wrf/interp.py index b7d939e..dd76ea5 100755 --- a/src/wrf/interp.py +++ b/src/wrf/interp.py @@ -20,135 +20,169 @@ from .pressure import get_pressure # Note: Extension decorator is good enough to handle left dims @set_interp_metadata("horiz") -def interplevel(field3d, z, desiredlev, missing=Constants.DEFAULT_FILL, +def interplevel(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, meta=True): - """Interpolates a three-dimensional field specified pressure or - height level. + """Return the three-dimensional field horizontally interpolated to a + specified vertical level. - Parameters - ---------- - field3d : `xarray.DataArray` or `numpy.ndarray` - A three-dimensional field. - - z : `xarray.DataArray` or `numpy.ndarray` - A three-dimensional array for the vertical coordinate, typically - pressure or height. - - desiredlev : float - The desired vertical level. Must be in the same units as the `z` - parameter. - - missing : float - The fill value to use for the output. - `Default is wrf.Constants.DEFAULT_FILL`. - - meta : {True, False} - Set to False to disable metadata and return `numpy.ndarray` instead of - `xarray.DataArray`. Default is True. - - Returns - ------- - `xarray.DataArray` or `numpy.ndarray` - Returns the interpolated variable. If xarray is enabled and - the meta parameter is True, then the result will be an - `xarray.DataArray` object. Otherwise, the result will be a - `numpy.ndarray` object with no metadata. + Args: + + field3d (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional field to interpolate, with the rightmost + dimensions of nz x ny x nx. + + vert (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional array for the vertical coordinate, typically + pressure or height. This array must have the same dimensionality + as *field3d*. + + desiredlev (:obj:`float`): The desired vertical level. + Must be in the same units as the *vert* parameter. + + missing (:obj:`float`): The fill value to use for the output. + Default is :data:`wrf.Constants.DEFAULT_FILL`. + + meta (:obj:`bool`): Set to False to disable metadata and return + :class:`numpy.ndarray` instead of + :class:`xarray.DataArray`. Default is True. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + interpolated variable. If xarray is enabled and + the *meta* parameter is True, then the result will be an + :class:`xarray.DataArray` object. Otherwise, the result will + be a :class:`numpy.ndarray` object with no metadata. + + Example: + + Interpolate Geopotential Height to 500 hPa + + .. code-block:: python + + from netCDF4 import Dataset + from wrf import getvar, interplevel + + wrfin = Dataset("wrfout_d02_2010-06-13_21:00:00") + + p = getvar(wrfin, "pressure") + ht = getvar(wrfin, "z", units="dm") + + ht_500 = interplevel(ht, p, 500.0) + """ # Some fields like uvmet have an extra left dimension for the product # type, we'll handle that iteration here. - multi = True if field3d.ndim - z.ndim == 1 else False + multi = True if field3d.ndim - vert.ndim == 1 else False if not multi: - result = _interpz3d(field3d, z, desiredlev, missing) + result = _interpz3d(field3d, vert, desiredlev, missing) else: outshape = field3d.shape[0:-3] + field3d.shape[-2:] result = np.empty(outshape, dtype=field3d.dtype) for i in py3range(field3d.shape[0]): result[i,:] = ( - _interpz3d(field3d[i,:], z, desiredlev, missing)[:]) + _interpz3d(field3d[i,:], vert, desiredlev, missing)[:]) return ma.masked_values (result, missing) @set_interp_metadata("cross") -def vertcross(field3d, z, missing=Constants.DEFAULT_FILL, +def vertcross(field3d, vert, missing=Constants.DEFAULT_FILL, pivot_point=None, angle=None, start_point=None, end_point=None, latlon=False, cache=None, meta=True): - """Return the vertical cross section for a 3D field, interpolated - to a vertical plane defined by a horizontal line. - - The horizontal line is defined by either including the - `pivot_point` and `angle` parameters, or the `start_point` and - `end_point` parameters. - - Parameters - ---------- - field3d : `xarray.DataArray` or `numpy.ndarray` - A three-dimensional field. - - z : `xarray.DataArray` or `numpy.ndarray` - A three-dimensional array for the vertical coordinate, typically - pressure or height - - pivot_point : tuple or list, optional - A tuple or list with two entries, in the form of [x, y] - (or [west_east, south_north]), which indicates the x,y location - through which the plane will pass. Must also specify `angle`. - - angle : float, optional - Only valid for cross sections where a plane will be plotted through - a given point on the model domain. 0.0 represents a S-N cross section - and 90.0 a W-E cross section. - - start_point : tuple or list, optional - A tuple or list with two entries, in the form of [x, y] - (or [west_east, south_north]), which indicates the start x,y location - through which the plane will pass. - - end_point : tuple or list, optional - A tuple or list with two entries, in the form of [x, y] - (or [west_east, south_north]), which indicates the end x,y location - through which the plane will pass. - - latlon : {True, False}, optional - Set to True to also interpolate the two-dimensional latitude and - longitude coordinates along the same horizontal line and include - this information in the metadata (if enabled). This can be - helpful for plotting. Default is False. - - cache : dict, optional - A dictionary of (varname, numpy.ndarray) pairs which can be used to - supply pre-extracted NetCDF variables to the computational routines. - This can be used to prevent the repeated variable extraction from large - sequences of data files. Default is None. - - meta : {True, False}, optional - Set to False to disable metadata and return `numpy.ndarray` instead of - `xarray.DataArray`. Default is True. - - Returns - ------- - `xarray.DataArray` or `numpy.ndarray` - Returns the interpolated variable. If xarray is enabled and - the meta parameter is True, then the result will be an - `xarray.DataArray` object. Otherwise, the result will be a - `numpy.ndarray` object with no metadata. + """Return the vertical cross section for a three-dimensional field. + + The cross section is defined by a horizontal line through the domain. + This horizontal line is defined by either including the + *pivot_point* and *angle* parameters, or the *start_point* and + *end_point* parameters. + + The vertical levels for the cross section are fixed, and are determined by + dividing the vertical coordinate in to grid boxes of roughly 1% of the + maximum vertical distance from top to bottom. If all vertical levels are + desired, use the lower level :meth:`wrf.interp2dxy` function. + + See Also: + + :meth:`wrf.interp2dxy` + + Args: + + field3d (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional field to interpolate, whose + rightmost dimensions are nz x ny x nx. + + vert (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional variable for the vertical coordinate, typically + pressure or height. This array must have the same dimensionality + as *field3d* + + missing (:obj:`float`): The fill value to use for the output. + Default is :data:`wrf.Constants.DEFAULT_FILL`. + + pivot_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, + in the form of [x, y] (or [west_east, south_north]), which + indicates the x,y location through which the plane will pass. + Must also specify `angle`. + + angle (:obj:`float`, optional): Only valid for cross sections where + a plane will be plotted through + a given point on the model domain. 0.0 represents a S-N cross + section. 90.0 is a W-E cross section. + + start_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the start + x,y location through which the plane will pass. + + end_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the end x,y + location through which the plane will pass. + + latlon (:obj:`bool`, optional): Set to True to also interpolate the + two-dimensional latitude and longitude coordinates along the same + horizontal line and include this information in the metadata + (if enabled). This can be helpful for plotting. Default is False. + + cache (: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. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: + The interpolated variable. 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. """ # Some fields like uvmet have an extra left dimension for the product # type, we'll handle that iteration here. - multi = True if field3d.ndim - z.ndim == 1 else False + multi = True if field3d.ndim - vert.ndim == 1 else False try: xy = cache["xy"] var2dz = cache["var2dz"] z_var2d = cache["z_var2d"] except (KeyError, TypeError): - xy, var2dz, z_var2d = get_xy_z_params(npvalues(z), pivot_point, angle, - start_point, end_point) + xy, var2dz, z_var2d = get_xy_z_params(npvalues(vert), pivot_point, + angle, start_point, end_point) if not multi: result = _vertcross(field3d, xy, var2dz, z_var2d, missing) @@ -170,55 +204,59 @@ def interpline(field2d, pivot_point=None, cache=None, meta=True): """Return the two-dimensional field interpolated along a line. - Parameters - ---------- - field2d : `xarray.DataArray` or `numpy.ndarray` - A two-dimensional field. - - pivot_point : tuple or list, optional - A tuple or list with two entries, in the form of [x, y] - (or [west_east, south_north]), which indicates the x,y location - through which the plane will pass. Must also specify `angle`. - - angle : float, optional - Only valid for cross sections where a plane will be plotted through - a given point on the model domain. 0.0 represents a S-N cross section - and 90.0 a W-E cross section. - - start_point : tuple or list, optional - A tuple or list with two entries, in the form of [x, y] - (or [west_east, south_north]), which indicates the start x,y location - through which the plane will pass. - - end_point : tuple or list, optional - A tuple or list with two entries, in the form of [x, y] - (or [west_east, south_north]), which indicates the end x,y location - through which the plane will pass. - - latlon : {True, False}, optional - Set to True to also interpolate the two-dimensional latitude and - longitude coordinates along the same horizontal line and include - this information in the metadata (if enabled). This can be - helpful for plotting. Default is False. - - cache : dict, optional - A dictionary of (varname, numpy.ndarray) pairs which can be used to - supply pre-extracted NetCDF variables to the computational routines. - This can be used to prevent the repeated variable extraction from large - sequences of data files. Default is None. - - meta : {True, False}, optional - Set to False to disable metadata and return `numpy.ndarray` instead of - `xarray.DataArray`. Default is True. - - - Returns - ------- - `xarray.DataArray` or `numpy.ndarray` - Returns the interpolated variable. If xarray is enabled and - the meta parameter is True, then the result will be an - `xarray.DataArray` object. Otherwise, the result will be a - `numpy.ndarray` object with no metadata. + Args: + + field2d (:class:`xarray.DataArray` or :class:`numpy.ndarray`): + A two-dimensional field. + + pivot_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, + in the form of [x, y] (or [west_east, south_north]), which + indicates the x,y location through which the plane will pass. + Must also specify `angle`. + + angle (:obj:`float`, optional): Only valid for cross sections where + a plane will be plotted through + a given point on the model domain. 0.0 represents a S-N cross + section. 90.0 is a W-E cross section. + + start_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the start + x,y location through which the plane will pass. + + end_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the end x,y + location through which the plane will pass. + + latlon (:obj:`bool`, optional): Set to True to also interpolate the + two-dimensional latitude and longitude coordinates along the same + horizontal line and include this information in the metadata + (if enabled). This can be helpful for plotting. Default is False. + + cache (: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. + + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: + The interpolated variable. 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. + """ @@ -231,80 +269,92 @@ def interpline(field2d, pivot_point=None, @set_interp_metadata("vinterp") -def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, +def vinterp(wrfin, field, vert_coord, interp_levels, extrapolate=False, field_type=None, log_p=False, timeidx=0, method="cat", squeeze=True, cache=None, meta=True): """Return the field vertically interpolated to the given the type of surface and a set of new levels. - Parameters - ---------- - wrfnc : `netCD4F.Dataset`, `Nio.NioFile`, or a sequence - Input WRF ARW NetCDF data as a `netCDF4.Dataset`, `Nio.NioFile` or an - iterable sequence of the aforementioned types. - - field : `xarray.DataArray` or `numpy.ndarray` - A three-dimensional field. - - vert_coord : {'pressure', 'pres', 'p', 'ght_msl', 'ght_agl', 'theta', 'th', 'theta-e', 'thetae', 'eth'} - A string indicating the vertical coordinate type to interpolate to. - - Valid strings are: - * 'pressure', 'pres', 'p': pressure [hPa] - * 'ght_msl': grid point height msl [km] - * 'ght_agl': grid point height agl [km] - * 'theta', 'th': potential temperature [K] - * 'theta-e', 'thetae', 'eth': equivalent potential temperature [K] + Args: + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`,\ + or an iterable): + Input WRF ARW NetCDF data as a :class:`netCDF4.Dataset`, + :class:`Nio.NioFile` or an iterable sequence of the + aforementioned types. - interp_levels : sequence - A 1D sequence of vertical levels to interpolate to. - - extrapolate : {True, False}, optional - Set to True to extrapolate values below ground. Default is False. - - field_type : {'none', 'pressure', 'pres', 'p', 'z', 'tc', 'tk', 'theta', 'th', 'theta-e', 'thetae', 'eth', 'ght'}, optional - The type of field. Default is None. - - log_p : {True, False} - Use the log of the pressure for interpolation instead of just pressure. - Default is False. - - timeidx : int, optional - The time index to use when extracting auxiallary variables used in - the interpolation. This value must be set to match the same value - used when the `field` variable was extracted. Default is 0. - - method : {'cat', 'join'}, optional - The aggregation method to use for sequences, either 'cat' or 'join'. - 'cat' combines the data along the Time index. 'join' is creates a new - index for the file. This must be set to the same method used when - extracting the `field` variable. The default is 'cat'. - - squeeze : {True, False}, optional - Set to False to prevent dimensions with a size of 1 from being removed - from the shape of the output. Default is True. - - cache : dict, optional - A dictionary of (varname, ndarray) which can be used to supply - pre-extracted NetCDF variables to the computational routines. This can - be used to prevent the repeated variable extraction from large - sequences of data files. Default is None. - - meta : {True, False}, optional - Set to False to disable metadata and return `numpy.ndarray` instead of - `xarray.DataArray`. Default is True. - - - Returns - ------- - `xarray.DataArray` or `numpy.ndarray` - Returns the interpolated variable. If xarray is enabled and - the meta parameter is True, then the result will be an - `xarray.DataArray` object. Otherwise, the result will be a - `numpy.ndarray` object with no metadata. + field (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + three-dimensional field. + + vert_coord (:obj:`str`): A string indicating the vertical coordinate + type to interpolate to. + + Valid strings are: + * 'pressure', 'pres', 'p': pressure [hPa] + * 'ght_msl': grid point height msl [km] + * 'ght_agl': grid point height agl [km] + * 'theta', 'th': potential temperature [K] + * 'theta-e', 'thetae', 'eth': equivalent potential temperature \ + [K] + + interp_levels (sequence): A 1D sequence of vertical levels to + interpolate to. + + extrapolate (:obj:`bool`, optional): Set to True to extrapolate + values below ground. Default is False. + + field_type (:obj:`str`, optional): + The type of field. Default is None. + + Valid strings are: + * 'none': None + * 'pressure', 'pres', 'p': pressure + * 'z', 'ght': geopotential height + * 'tc': temperature [degC] + * 'tk': temperature [K] + * 'theta', 'th': potential temperature [K] + * 'theta-e', 'thetae', 'eth': equivalent potential temperature + + log_p (:obj:`bool`, optional): Use the log of the pressure for + interpolation instead of pressure. Default is False. + + timeidx (:obj:`int`, optional): + The time index to use when extracting auxiallary variables used in + the interpolation. This value must be set to match the same value + used when the `field` variable was extracted. Default is 0. + + method (: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. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: + The interpolated variable. 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. """ - _key = get_id(wrfnc) + _key = get_id(wrfin) # Remove case sensitivity field_type = field_type.lower() if field_type is not None else "none" @@ -341,7 +391,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, interp_levels = np.asarray(interp_levels, np.float64) # TODO: Check if field is staggered - if is_staggered(field, wrfnc): + if is_staggered(wrfin, field): raise RuntimeError("Please unstagger field in the vertical") # Check for valid coord @@ -364,29 +414,29 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, # Extract vriables #timeidx = -1 # Should this be an argument? - ncvars = extract_vars(wrfnc, timeidx, ("PSFC", "QVAPOR", "F"), + ncvars = extract_vars(wrfin, timeidx, ("PSFC", "QVAPOR", "F"), method, squeeze, cache, meta=False, _key=_key) sfp = ncvars["PSFC"] * ConversionFactors.PA_TO_HPA qv = ncvars["QVAPOR"] coriolis = ncvars["F"] - terht = get_terrain(wrfnc, timeidx, units="m", + terht = get_terrain(wrfin, timeidx, units="m", method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) - t = get_theta(wrfnc, timeidx, units="k", + t = get_theta(wrfin, timeidx, units="k", method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) - tk = get_temp(wrfnc, timeidx, units="k", + tk = get_temp(wrfin, timeidx, units="k", method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) - p = get_pressure(wrfnc, timeidx, units="pa", + p = get_pressure(wrfin, timeidx, units="pa", method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) - ght = get_height(wrfnc, timeidx, msl=True, units="m", + ght = get_height(wrfin, timeidx, msl=True, units="m", method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) - ht_agl = get_height(wrfnc, timeidx, msl=False, units="m", + ht_agl = get_height(wrfin, timeidx, msl=False, units="m", method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) @@ -427,7 +477,7 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, idir = 1 delta = 0.01 - eth = get_eth(wrfnc, timeidx, method=method, squeeze=squeeze, + eth = get_eth(wrfin, timeidx, method=method, squeeze=squeeze, cache=cache, meta=False, _key=_key) p_hpa = p * ConversionFactors.PA_TO_HPA diff --git a/src/wrf/interputils.py b/src/wrf/interputils.py index 645e4ef..a41d2b5 100644 --- a/src/wrf/interputils.py +++ b/src/wrf/interputils.py @@ -10,26 +10,65 @@ from .py3compat import py3range def to_positive_idxs(shape, coord): + """Return the positive index values. + + This function converts negative index values to positive index values. + + Args: + + shape (indexable sequence): The array shape. + + coord (indexable sequence): The coordinate pair for x and y. + + Returns: + + :obj:`list`: The coordinate values with all positive indexes. + + """ if (coord[-2] >= 0 and coord[-1] >= 0): return coord - return [x if (x >= 0) else shape[-i-1]+x for (i,x) in enumerate(coord) ] + return [x if (x >= 0) else shape[-i-1]+x for (i,x) in enumerate(coord)] -def calc_xy(xdim, ydim, pivot_point=None, angle=None, +def _calc_xy(xdim, ydim, pivot_point=None, angle=None, start_point=None, end_point=None): - """Returns the x,y points for the horizontal cross section line. + """Return the x,y points for the horizontal cross section line. - xdim - maximum x-dimension - ydim - maximum y-dimension - pivot_point - a pivot point of (south_north, west_east) - (must be used with angle) - angle - the angle through the pivot point in degrees - start_point - a start_point sequence of (x, y) - end_point - an end point sequence of (x, y) + Args: + + xdim (:obj:`int`): The x-dimension size. + + ydim (:obj:`int`): The y-dimension size. + + pivot_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, + in the form of [x, y] (or [west_east, south_north]), which + indicates the x,y location through which the plane will pass. + Must also specify `angle`. + + angle (:obj:`float`, optional): Only valid for cross sections where + a plane will be plotted through + a given point on the model domain. 0.0 represents a S-N cross + section. 90.0 is a W-E cross section. + + start_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the start + x,y location through which the plane will pass. + + end_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the end x,y + location through which the plane will pass. + + Returns: - """ + :class:`np.ndarray`: A two-dimensional array with the left index + representing each point along the line, and the rightmost dimension + having two values for the x and y coordinates [0=X, 1=Y]. + """ # Have a pivot point with an angle to find cross section if pivot_point is not None and angle is not None: xp = pivot_point[-2] @@ -139,8 +178,52 @@ def calc_xy(xdim, ydim, pivot_point=None, angle=None, return xy + def get_xy_z_params(z, pivot_point=None, angle=None, start_point=None, end_point=None): + """Return the cross section parameters. + + This function returns the xy horizontal cross section line coordinates, + the xy x z vertical values interpolated along the xy cross section + line, and the fixed vertical levels to be used by the cross section + algorithm (at ~1% increments for the minimum to maximum vertical + span). + + Args: + + z (:class:`numpy.ndarray`): The vertical coordinate, whose rightmost + dimensions are bottom_top x south_north x west_east. + + pivot_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, + in the form of [x, y] (or [west_east, south_north]), which + indicates the x,y location through which the plane will pass. + Must also specify `angle`. + + angle (:obj:`float`, optional): Only valid for cross sections where + a plane will be plotted through + a given point on the model domain. 0.0 represents a S-N cross + section. 90.0 is a W-E cross section. + + start_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the start + x,y location through which the plane will pass. + + end_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the end x,y + location through which the plane will pass. + + Returns: + + :obj:`tuple`: A tuple containing the xy horizontal cross section + coordinates, the vertical values interpolated along the xy cross + section line, and the fixed vertical levels used by the + cross section algorithm at ~1% increments of minimum to maximum + vertical span. + + """ xy = get_xy(z, pivot_point, angle, start_point, end_point) @@ -173,8 +256,44 @@ def get_xy_z_params(z, pivot_point=None, angle=None, return xy, var2dz, z_var2d + def get_xy(var, pivot_point=None, angle=None, start_point=None, end_point=None): + """Return the x,y points for the horizontal cross section line. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A variable + that contains a :attr:`shape` attribute. + + pivot_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, + in the form of [x, y] (or [west_east, south_north]), which + indicates the x,y location through which the plane will pass. + Must also specify `angle`. + + angle (:obj:`float`, optional): Only valid for cross sections where + a plane will be plotted through + a given point on the model domain. 0.0 represents a S-N cross + section. 90.0 is a W-E cross section. + + start_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the start + x,y location through which the plane will pass. + + end_point (:obj:`tuple` or :obj:`list`, optional): A + :obj:`tuple` or :obj:`list` with two entries, in the form of + [x, y] (or [west_east, south_north]), which indicates the end x,y + location through which the plane will pass. + + Returns: + + :class:`np.ndarray`: A two-dimensional array with the left index + representing each point along the line, and the rightmost dimension + having two values for the x and y coordinates [0=X, 1=Y]. + + """ if pivot_point is not None: pos_pivot = to_positive_idxs(var.shape[-2:], pivot_point) else: @@ -193,6 +312,6 @@ def get_xy(var, pivot_point=None, angle=None, xdim = var.shape[-1] ydim = var.shape[-2] - xy = calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) + xy = _calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end) return xy diff --git a/src/wrf/latlon.py b/src/wrf/latlon.py index 4522c45..c333ac9 100755 --- a/src/wrf/latlon.py +++ b/src/wrf/latlon.py @@ -6,37 +6,218 @@ from .latlonutils import (_lat_varname, _lon_varname, _ll_to_xy, _xy_to_ll) from .metadecorators import set_latlon_metadata -def get_lat(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_lat(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, stagger=None): + """Return the two dimensional latitude coordinate variable. - varname = _lat_varname(wrfnc, stagger) - lat_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + stagger (:obj:`str`): By default, the latitude is returned on the mass + grid, but a staggered grid can be chosen with the following + options: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + two dimensional latitude coordinate variable. + 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. + + """ + + varname = _lat_varname(wrfin, stagger) + lat_var = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta, _key) return lat_var[varname] -def get_lon(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_lon(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, stagger=None): + """Return the two dimensional longitude coordinate variable. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: - varname = _lon_varname(wrfnc, stagger) - lon_var = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + stagger (:obj:`str`): By default, the longitude is returned on the mass + grid, but a staggered grid can be chosen with the following + options: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + two dimensional longitude coordinate variable. + 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. + + """ + + varname = _lon_varname(wrfin, stagger) + lon_var = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta, _key) return lon_var[varname] # TODO: Do we need the user to know about method, squeeze, cache for this? -# Can either use wrfnc as a single file or sequence, or provide -# projection parameters (which don't allow for moving domains) @set_latlon_metadata(xy=True) -def ll_to_xy(wrfnc, latitude, longitude, timeidx=0, stagger=None, method="cat", - squeeze=True, cache=None, meta=True, as_int=True): - _key = get_id(wrfnc) - return _ll_to_xy(latitude, longitude, wrfnc, timeidx, stagger, method, - squeeze, cache, _key, as_int, **{}) +def ll_to_xy(wrfin, latitude, longitude, timeidx=0, + squeeze=True, meta=True, stagger=None, as_int=True): + """Return the x,y coordinates for a specified latitude and longitude. + + The *latitude* and *longitude* arguments can be a single value or a + sequence of values. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain the X (west_east) values. + - return_val[1,...] will contain the Y (south_north) values. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input WRF ARW NetCDF + data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` + or an iterable sequence of the aforementioned types. + + latitude (:obj:`float` or sequence): A single latitude or a sequence + of latitude values to be converted. + + longitude (:obj:`float` or sequence): A single longitude or a sequence + of latitude values to be converted. + + 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. + + 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. + + stagger (:obj:`str`): By default, the latitude is returned on the mass + grid, but a staggered grid can be chosen with the following + options: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + as_int (:obj:`bool`): Set to True to return the x,y values as + :obj:`int`, otherwise they will be returned as :obj:`float`. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + x,y coordinate value(s) whose leftmost dimension is 2 (0=X, 1=Y). + 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. + + """ + _key = get_id(wrfin) + return _ll_to_xy(latitude, longitude, wrfin, timeidx, stagger, "cat", + squeeze, None, _key, as_int, **{}) @set_latlon_metadata(xy=True) @@ -45,7 +226,84 @@ def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, as_int=True, ref_lat=None, ref_lon=None, pole_lat=None, pole_lon=None, known_x=None, known_y=None, dx=None, dy=None, latinc=None, loninc=None): + """Return the x, y coordinates for a specified latitude and longitude. + + The *latitude* and *longitude* arguments can be a single value or a + sequence of values. This version of the ll_to_xy routine allows users + to manually specify projection parameters. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain the X (west_east) values. + - return_val[1,...] will contain the Y (south_north) values. + + Args: + + latitude (:obj:`float` or sequence): A single latitude or a sequence + of latitude values to be converted. + + longitude (:obj:`float` or sequence): A single longitude or a sequence + of latitude values to be converted. + + 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. + + as_int (:obj:`bool`): Set to True to return the x,y values as + :obj:`int`, otherwise they will be returned as :obj:`float`. + + map_proj (:obj:`int`): Model projection [1=Lambert Conformal, + 2=Polar Stereographic, 3=Mercator, 6=Lat-Lon]. Required. + + truelat1 (:obj:`float`): True latitude 1. Required for + map_proj = 1, 2, 3 (defaults to 0 otherwise). + + truelat2 (:obj:`float`): True latitude 2. Optional for + map_proj = 1 (defaults to 0 otherwise). + + stand_lon (:obj:`float`): Standard longitude. Required. + + ref_lat (:obj:`float`): A reference latitude. Required. + + ref_lon (:obj:`float`): A reference longitude. Required. + + known_x (:obj:`float`): The known x-coordinate associated with + *ref_lon*. Required. + + known_y (:obj:`float`): The known y-coordinate associated with + *ref_lat*. Required. + + pole_lat (:obj:`float`): Pole latitude. Optional for + *map_proj* = 6 (defaults to 90 otherwise). + + pole_lon (:obj:`float`): Pole longitude. Optional for + *map_proj* = 6 (defaults to 0 otherwise). + + dx (:obj:`float`): The x spacing in meters at the true latitude. + Required for *map_proj* = 1, 2, 3 (defaults to 0 otherwise). + + dy (:obj:`float`) - The y spacing in meters at the true latitude. + Required for *map_proj* = 1, 2, 3 (defaults to 0 otherwise). + + latinc (:obj:`float`): The lower left corner latitude at (0,0). + Required for *map_proj* = 6. + + loninc (:obj:`float`) - The lower left corner longitude at (0,0). + Required for *map_proj* = 6. + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + x,y coordinate value(s) whose leftmost dimension is 2 (0=X, 1=Y). + 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. + + """ loc = locals() projparams = {name : loc[name] for name in ("map_proj", "truelat1", "truelat2", "stand_lon", "ref_lat", @@ -53,15 +311,73 @@ def ll_to_xy_proj(latitude, longitude, meta=True, squeeze=True, as_int=True, "known_x", "known_y", "dx", "dy", "latinc", "loninc")} - return _ll_to_xy(latitude, longitude, None, 0, squeeze, "cat", True, None, + return _ll_to_xy(latitude, longitude, None, 0, True, "cat", squeeze, None, None, as_int, **projparams) @set_latlon_metadata(xy=False) -def xy_to_ll(wrfnc, x, y, timeidx=0, stagger=None, method="cat", squeeze=True, - cache=None, meta=True): - _key = get_id(wrfnc) - return _xy_to_ll(x, y, wrfnc, timeidx, stagger, method, squeeze, cache, +def xy_to_ll(wrfin, x, y, timeidx=0, stagger=None, squeeze=True, meta=True): + """Return the latitude and longitude for specified x,y coordinates. + + The *x* and *y* arguments can be a single value or a sequence of values. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain the latitude values. + - return_val[1,...] will contain the longitude values. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input WRF ARW NetCDF + data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` + or an iterable sequence of the aforementioned types. + + x (:obj:`float` or sequence): A single x-coordinate or a sequence + of x-coordinate values to be converted. + + y (:obj:`float` or sequence): A single y-coordinate or a sequence + of y-coordinate values to be converted. + + 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. + + 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. + + stagger (:obj:`str`): By default, the latitude is returned on the mass + grid, but a staggered grid can be chosen with the following + options: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + as_int (:obj:`bool`): Set to True to return the x,y values as + :obj:`int`, otherwise they will be returned as :obj:`float`. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + latitude and longitude values whose leftmost dimension is 2 + (0=latitude, 1=longitude). + 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. + + """ + _key = get_id(wrfin) + return _xy_to_ll(x, y, wrfin, timeidx, stagger, "cat", True, None, _key, **{}) @@ -70,6 +386,84 @@ def xy_to_ll_proj(x, y, meta=True, squeeze=True, map_proj=None, truelat1=None, truelat2=None, stand_lon=None, ref_lat=None, ref_lon=None, pole_lat=None, pole_lon=None, known_x=None, known_y=None, dx=None, dy=None, latinc=None, loninc=None): + """Return the latitude and longitude for the specified x,y coordinates. + + The *x* and *y* arguments can be a single value or a + sequence of values. This version of the xy_to_ll routine allows users + to manually specify map projection parameters. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain the latitude values. + - return_val[1,...] will contain the longitude values. + + Args: + + x (:obj:`float` or sequence): A single x-coordinate or a sequence + of x-coordinate values to be converted. + + y (:obj:`float` or sequence): A single y-coordinate or a sequence + of y-coordinate values to be converted. + + 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. + + as_int (:obj:`bool`): Set to True to return the x,y values as + :obj:`int`, otherwise they will be returned as :obj:`float`. + + map_proj (:obj:`int`): Model projection [1=Lambert Conformal, + 2=Polar Stereographic, 3=Mercator, 6=Lat-Lon]. Required. + + truelat1 (:obj:`float`): True latitude 1. Required for + map_proj = 1, 2, 3 (defaults to 0 otherwise). + + truelat2 (:obj:`float`): True latitude 2. Optional for + map_proj = 1 (defaults to 0 otherwise). + + stand_lon (:obj:`float`): Standard longitude. Required. + + ref_lat (:obj:`float`): A reference latitude. Required. + + ref_lon (:obj:`float`): A reference longitude. Required. + + known_x (:obj:`float`): The known x-coordinate associated with + *ref_lon*. Required. + + known_y (:obj:`float`): The known y-coordinate associated with + *ref_lat*. Required. + + pole_lat (:obj:`float`): Pole latitude. Optional for + *map_proj* = 6 (defaults to 90 otherwise). + + pole_lon (:obj:`float`): Pole longitude. Optional for + *map_proj* = 6 (defaults to 0 otherwise). + + dx (:obj:`float`): The x spacing in meters at the true latitude. + Required for *map_proj* = 1, 2, 3 (defaults to 0 otherwise). + + dy (:obj:`float`) - The y spacing in meters at the true latitude. + Required for *map_proj* = 1, 2, 3 (defaults to 0 otherwise). + + latinc (:obj:`float`): The lower left corner latitude at (0,0). + Required for *map_proj* = 6. + + loninc (:obj:`float`) - The lower left corner longitude at (0,0). + Required for *map_proj* = 6. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + latitude and longitude values whose leftmost dimension is 2 + (0=latitude, 1=longitude). + 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. + """ loc = locals() projparams = {name : loc[name] for name in ("map_proj", "truelat1", "truelat2", "stand_lon", "ref_lat", diff --git a/src/wrf/latlonutils.py b/src/wrf/latlonutils.py index 6a0aabe..527f445 100644 --- a/src/wrf/latlonutils.py +++ b/src/wrf/latlonutils.py @@ -12,9 +12,32 @@ from .util import (extract_vars, extract_global_attrs, iter_left_indexes, is_mapping, is_multi_file) from .py3compat import viewkeys, viewitems -def _lat_varname(wrfnc, stagger): +def _lat_varname(wrfin, stagger): + """Return the latitude variable name for the specified stagger type. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input WRF ARW NetCDF + data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` + or an iterable sequence of the aforementioned types. + + stagger (:obj:`str`): The staggered grid type which is one of the + following: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + Returns: + + :obj:`str`: The latitude variable name. + + """ if stagger is None or stagger.lower() == "m": - varname = either("XLAT", "XLAT_M")(wrfnc) + varname = either("XLAT", "XLAT_M")(wrfin) elif stagger.lower() == "u" or stagger.lower() == "v": varname = "XLAT_{}".format(stagger.upper()) else: @@ -22,9 +45,32 @@ def _lat_varname(wrfnc, stagger): return varname -def _lon_varname(wrfnc, stagger): +def _lon_varname(wrfin, stagger): + """Return the longitude variable name for the specified stagger type. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input WRF ARW NetCDF + data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` + or an iterable sequence of the aforementioned types. + + stagger (:obj:`str`): The staggered grid type, which is one of the + following: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + Returns: + + :obj:`str`: The latitude variable name. + + """ if stagger is None or stagger.lower() == "m": - varname = either("XLONG", "XLONG_M")(wrfnc) + varname = either("XLONG", "XLONG_M")(wrfin) elif stagger.lower() == "u" or stagger.lower() == "v": varname = "XLONG_{}".format(stagger.upper()) else: @@ -32,11 +78,61 @@ def _lon_varname(wrfnc, stagger): return varname -def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache, _key): +def _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key): + """Return the map projection parameters. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. The default is 0. + + stagger (:obj:`str`): The staggered grid type, which is one of the + following: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + + + """ if timeidx < 0: raise ValueError("'timeidx' must be greater than 0") - attrs = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", "TRUELAT1", + attrs = extract_global_attrs(wrfin, attrs=("MAP_PROJ", "TRUELAT1", "TRUELAT2", "STAND_LON", "DX", "DY")) map_proj = attrs["MAP_PROJ"] @@ -47,7 +143,7 @@ def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache, _key): dy = attrs["DY"] if map_proj == ProjectionTypes.LAT_LON: - pole_attrs = extract_global_attrs(wrfnc, attrs=("POLE_LAT", + pole_attrs = extract_global_attrs(wrfin, attrs=("POLE_LAT", "POLE_LON")) pole_lat = pole_attrs["POLE_LAT"] pole_lon = pole_attrs["POLE_LON"] @@ -59,31 +155,32 @@ def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache, _key): latinc = 0.0 loninc = 0.0 - latvar = _lat_varname(wrfnc, stagger) - lonvar = _lon_varname(wrfnc, stagger) + latvar = _lat_varname(wrfin, stagger) + lonvar = _lon_varname(wrfin, stagger) lat_timeidx = timeidx - is_moving = is_moving_domain(wrfnc, latvar=latvar, lonvar=lonvar) + is_moving = is_moving_domain(wrfin, latvar=latvar, lonvar=lonvar, + _key=_key) # Only need one file and one time if the domain is not moving if not is_moving: if is_multi_time_req(timeidx): lat_timeidx = 0 - if is_multi_file(wrfnc): - if not is_mapping(wrfnc): - wrfnc = next(iter(wrfnc)) # only need one file + if is_multi_file(wrfin): + if not is_mapping(wrfin): + wrfin = next(iter(wrfin)) # only need one file else: - first_entry = next(iter(viewkeys(wrfnc))) - wrfnc = wrfnc[first_entry] + first_entry = next(iter(viewkeys(wrfin))) + wrfin = wrfin[first_entry] key = _key[first_entry] - return _get_proj_params(wrfnc, timeidx, stagger, + return _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, key) - xlat = extract_vars(wrfnc, lat_timeidx, (latvar,), method, squeeze, cache, + xlat = extract_vars(wrfin, lat_timeidx, (latvar,), method, squeeze, cache, meta=False, _key=_key)[latvar] - xlon = extract_vars(wrfnc, lat_timeidx, (lonvar,), method, squeeze, cache, + xlon = extract_vars(wrfin, lat_timeidx, (lonvar,), method, squeeze, cache, meta=False, _key=_key)[lonvar] ref_lat = np.ravel(xlat[..., 0, 0]) @@ -98,10 +195,35 @@ def _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, cache, _key): def _dict_keys_to_upper(d): + """Return a dictionary with the keys changed to uppercase. + + Args: + + d (:obj:`dict`): A dictionary. + + Returns: + + :obj:`dict`: A dictionary with uppercase keys. + + """ return {key.upper(): val for key, val in viewitems(d)} # known_x and known_y are 0-based def _kwarg_proj_params(**projparams): + """Return the map projection parameters. + + This function aggregates the projection parameter keyword + arguments and also performs sanity checking on them. + + Args: + + **projparams: Projection parameter keyword arguments. + + Returns: + + :obj:`tuple`: The map projection parameters. + + """ projparams = _dict_keys_to_upper(projparams) map_proj = projparams.get("MAP_PROJ") @@ -169,14 +291,89 @@ def _kwarg_proj_params(**projparams): # Will return 0-based indexes -def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0, +def _ll_to_xy(latitude, longitude, wrfin=None, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None, _key=None, as_int=True, **projparams): + """Return the x,y coordinates for a specified latitude and longitude. + + The *latitude* and *longitude* arguments can be a single value or a + sequence of values. - if wrfnc is not None: + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain the X (west_east) values. + - return_val[1,...] will contain the Y (south_north) values. + + Args: + + latitude (:obj:`float` or sequence): A single latitude or a sequence + of latitude values to be converted. + + longitude (:obj:`float` or sequence): A single longitude or a sequence + of latitude values to be converted. + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. The default is 0. + + stagger (:obj:`str`): By default, the latitude and longitude are + returned on the mass grid, but a staggered grid can be chosen + with the following options: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + as_int (:obj:`bool`): Set to True to return the x,y values as + :obj:`int`, otherwise they will be returned as :obj:`float`. + + **projparams: Map projection keyword arguments to set manually. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + x,y coordinate value(s) whose leftmost dimension is 2 (0=X, 1=Y). + 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 wrfin is not None: (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, - loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, + loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key) else: (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, @@ -255,14 +452,92 @@ def _ll_to_xy(latitude, longitude, wrfnc=None, timeidx=0, return result # X and Y should be 0-based -def _xy_to_ll(x, y, wrfnc=None, timeidx=0, stagger=None, +def _xy_to_ll(x, y, wrfin=None, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None, _key=None, **projparams): + + """Return the latitude and longitude for specified x,y coordinates. + + The *x* and *y* arguments can be a single value or a sequence of values. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain the latitude values. + - return_val[1,...] will contain the longitude values. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input WRF ARW NetCDF + data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` + or an iterable sequence of the aforementioned types. + + x (:obj:`float` or sequence): A single x-coordinate or a sequence + of x-coordinate values to be converted. + + y (:obj:`float` or sequence): A single y-coordinate or a sequence + of y-coordinate values to be converted. + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. The default is 0. + + stagger (:obj:`str`): By default, the latitude and longitude are + returned on the mass grid, but a staggered grid can be chosen + with the following options: + + - 'm': Use the mass grid (default). + - 'u': Use the same staggered grid as the u wind component, + which has a staggered west_east (x) dimension. + - 'v': Use the same staggered grid as the v wind component, + which has a staggered south_north (y) dimension. + + 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + **projparams: Map projection keyword arguments to set manually. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + latitude and longitude values whose leftmost dimension is 2 + (0=latitude, 1=longitude). + 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 wrfnc is not None: + if wrfin is not None: (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, pole_lat, pole_lon, known_x, known_y, dx, dy, latinc, - loninc) = _get_proj_params(wrfnc, timeidx, stagger, method, squeeze, + loninc) = _get_proj_params(wrfin, timeidx, stagger, method, squeeze, cache, _key) else: (map_proj, truelat1, truelat2, stdlon, ref_lat, ref_lon, diff --git a/src/wrf/metadecorators.py b/src/wrf/metadecorators.py index f635f0b..8314349 100644 --- a/src/wrf/metadecorators.py +++ b/src/wrf/metadecorators.py @@ -7,7 +7,7 @@ import numpy as np import numpy.ma as ma from .extension import _interpline -from .util import (extract_vars, combine_with, either, from_args, arg_location, +from .util import (extract_vars, either, from_args, arg_location, is_coordvar, latlon_coordvars, npvalues, from_var, iter_left_indexes) from .coordpair import CoordPair @@ -22,11 +22,58 @@ if xarray_enabled(): def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, remove_dims=None, dimnames=None, coords=None, **fixed_attrs): - """Decorator to set the metadata for a WRF method. + """A decorator that sets the metadata for a wrapped function's output. - A cache is inserted/updated to include the extracted variable that will - have its metadata copied. This prevents the variable being extracted more - than once. This extraction can be slow with sequences of large files. + Generally, the metadata is copied from the variable specified by + *copy_varname*, with other fixed fields set by the other decorator + arguments. + + The *cache* argument used by most diagnostic routines is supplied by + this decorator when the *copy_varname* variable is extracted, in order + to prevent the variable from being extracted again by the wrapped function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + copy_varname (:obj:`str`, optional): The NetCDF variable name to copy. + Default is None. + + delete_attrs (sequence of :obj:`str`, optional): A sequence of key + names to remove from the :attr:`xarray.DataArray.attrs` attribute + in the wrapped function output (after being copied from + *copy_varname*). Default is None. + + name (:obj:`str`): The name to use for the + :attr:`xarray.DataArray.name` attribute. + + remove_dims (sequence of :obj:`int`, optional): A sequence of dimension + indexes to be removed from the wrapped function output (after being + copied from *copy_varname*). This is useful when the copy + variable is three dimensional but the wrapped function output + is two dimensional, and you still want to keep the names of + the rightmost two dimensions. Default is None. + + dimnames (sequence of :obj:`str`, optional): A sequence of dimension + names in order to manually set the + :attr:`xarray.DataArray.dims` attribute. Default is None. + + coords (:obj:`dict`): A mapping of coordinate name to coordinate + value to manually specify the :attr:`xarray.DataArray.coords` + attribute. Default is None. + + **fixed_attrs: These keyword arguments are added to the + :attr:`xarray.DataArray.attrs` attribute. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + function output with or without metadata. 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. + """ @wrapt.decorator @@ -39,12 +86,12 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, if not xarray_enabled() or not do_meta: return wrapped(*args, **kwargs) - argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", + argvars = from_args(wrapped, ("wrfin", "timeidx", "method", "squeeze", "cache", "units", "meta", "_key"), *args, **kwargs) - wrfnc = argvars["wrfnc"] + wrfin = argvars["wrfin"] timeidx = argvars["timeidx"] units = argvars["units"] method = argvars["method"] @@ -56,7 +103,7 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, # Note: can't modify nonlocal var if (callable(copy_varname)): - _copy_varname = copy_varname(wrfnc) + _copy_varname = copy_varname(wrfin) else: _copy_varname = copy_varname @@ -66,7 +113,7 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, if var_to_copy is None: - var_to_copy = extract_vars(wrfnc, timeidx, (_copy_varname,), + var_to_copy = extract_vars(wrfin, timeidx, (_copy_varname,), method, squeeze, cache, meta=True, _key=_key)[_copy_varname] @@ -90,11 +137,8 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, outname = ucode(var_to_copy.name) if dimnames is not None: - if isinstance(dimnames, combine_with): - outdimnames, outcoords = dimnames(var_to_copy) - else: - outdimnames = dimnames - outcoords = coords + outdimnames = dimnames + outcoords = coords else: outdimnames += var_to_copy.dims outcoords.update(var_to_copy.coords) @@ -140,6 +184,49 @@ def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, def set_wind_metadata(copy_varname, name, description, wind_ncvar=False, two_d=False, wspd_wdir=False): + """A decorator that sets the metadata for a wrapped wind function's output. + + This is a special metadata decorator for working with wind functions, which + include wind extraction routines, uvmet routines, and + wind speed / wind direction routines. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + copy_varname (:obj:`str`, optional): The NetCDF variable name to copy. + Default is None. + + name (:obj:`str`): The name to use for the + :attr:`xarray.DataArray.name` attribute. + + description (:obj:`str`): The description to use for the 'description' + key in the :attr:`xarray.DataArray.attrs` attribute. + + wind_ncvar (:obj:`bool`, optional): Set to True when the wrapped + function is simply extracting a wind variable (U, V, W) from the + NetCDF file. Set to False for other types of wind algorithms + (uvmet, wspd_wdir, etc). Default is False. + + two_d (:obj:`bool`, optional): Set to True if the wind field is + two-dimensional. Set to False for a three-dimensional wind field. + Default is False. + + wspd_wdir (:obj:`bool`): Set to True if the wrapped function is a + wind speed / wind direction algorithm. Otherwise, set to False. + Default is False. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + wind function output with or without metadata. 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. + + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -150,11 +237,11 @@ def set_wind_metadata(copy_varname, name, description, if not xarray_enabled() or not do_meta: return wrapped(*args, **kwargs) - argvars = from_args(wrapped, ("wrfnc", "timeidx", "units", + argvars = from_args(wrapped, ("wrfin", "timeidx", "units", "method", "squeeze", "ten_m", "cache", "_key"), *args, **kwargs) - wrfnc = argvars["wrfnc"] + wrfin = argvars["wrfin"] timeidx = argvars["timeidx"] units = argvars["units"] method = argvars["method"] @@ -166,11 +253,11 @@ def set_wind_metadata(copy_varname, name, description, cache = {} if isinstance(copy_varname, either): - _copy_varname = copy_varname(wrfnc) + _copy_varname = copy_varname(wrfin) else: _copy_varname = copy_varname - copy_var = extract_vars(wrfnc, timeidx, _copy_varname, + copy_var = extract_vars(wrfin, timeidx, _copy_varname, method, squeeze, cache, meta=True, _key=_key)[_copy_varname] @@ -236,6 +323,28 @@ def set_wind_metadata(copy_varname, name, description, return func_wrapper def set_cape_metadata(is2d): + """A decorator that sets the metadata for a wrapped CAPE function's output. + + This is a special metadata decorator for working with CAPE functions. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + is2d (:obj:`bool`): Set to True if the wrapped function is for a + two-dimensional CAPE routine. Set to False for a + three-dimensional CAPE routine. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + CAPE function output with or without metadata. 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -246,10 +355,10 @@ def set_cape_metadata(is2d): if not xarray_enabled() or not do_meta: return wrapped(*args, **kwargs) - argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", "squeeze", + argvars = from_args(wrapped, ("wrfin", "timeidx", "method", "squeeze", "cache", "_key", "missing"), *args, **kwargs) - wrfnc = argvars["wrfnc"] + wrfin = argvars["wrfin"] timeidx = argvars["timeidx"] method = argvars["method"] squeeze = argvars["squeeze"] @@ -260,7 +369,7 @@ def set_cape_metadata(is2d): cache = {} _copy_varname = "P" - copy_var = extract_vars(wrfnc, timeidx, _copy_varname, method, squeeze, + copy_var = extract_vars(wrfin, timeidx, _copy_varname, method, squeeze, cache, meta=True, _key=_key)[_copy_varname] # Make a copy so we don't modify a user supplied cache @@ -330,6 +439,24 @@ def set_cape_metadata(is2d): def set_cloudfrac_metadata(): + """A decorator that sets the metadata for a wrapped cloud fraction + function's output. + + This is a special metadata decorator for working with cloud fraction + functions. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + cloud fraction function output with or without metadata. 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -340,10 +467,10 @@ def set_cloudfrac_metadata(): if not xarray_enabled() or not do_meta: return wrapped(*args, **kwargs) - argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", "squeeze", + argvars = from_args(wrapped, ("wrfin", "timeidx", "method", "squeeze", "cache", "_key"), *args, **kwargs) - wrfnc = argvars["wrfnc"] + wrfin = argvars["wrfin"] timeidx = argvars["timeidx"] method = argvars["method"] squeeze = argvars["squeeze"] @@ -353,7 +480,7 @@ def set_cloudfrac_metadata(): cache = {} _copy_varname = "P" - copy_var = extract_vars(wrfnc, timeidx, _copy_varname, method, squeeze, + copy_var = extract_vars(wrfin, timeidx, _copy_varname, method, squeeze, cache, meta=True, _key=_key)[_copy_varname] # Make a copy so we don't modify a user supplied cache @@ -403,6 +530,29 @@ def set_cloudfrac_metadata(): return func_wrapper def set_latlon_metadata(xy=False): + """A decorator that sets the metadata for a wrapped latlon function's + output. + + This is a special metadata decorator for working with latlon functions. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + xy (:obj:`bool`, optional): Set to True if the wrapped function returns + xy values (ll_to_xy). Set to False if the wrapped function returns + latlon values (xy_to_ll). Default is False. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + latlon function output with or without metadata. 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): @@ -465,6 +615,29 @@ def set_latlon_metadata(xy=False): return func_wrapper def set_height_metadata(geopt=False): + """A decorator that sets the metadata for a wrapped height function's + output. + + This is a special metadata decorator for working with height functions. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + geopt (:obj:`bool`, optional): Set to True if the wrapped function + returns geopotential. Set to True if the wrapped function + returns geopotential height. Default is False. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + height function output with or without metadata. 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -475,11 +648,11 @@ def set_height_metadata(geopt=False): if not xarray_enabled() or not do_meta: return wrapped(*args, **kwargs) - argvars = from_args(wrapped, ("wrfnc", "timeidx", "method", + argvars = from_args(wrapped, ("wrfin", "timeidx", "method", "squeeze", "units", "msl", "cache", "_key"), *args, **kwargs) - wrfnc = argvars["wrfnc"] + wrfin = argvars["wrfin"] timeidx = argvars["timeidx"] units = argvars["units"] method = argvars["method"] @@ -493,8 +666,8 @@ def set_height_metadata(geopt=False): # For height, either copy the met_em GHT variable or copy and modify # pressure (which has the same dims as destaggered height) - ht_metadata_varname = either("P", "GHT")(wrfnc) - ht_var = extract_vars(wrfnc, timeidx, ht_metadata_varname, + ht_metadata_varname = either("P", "GHT")(wrfin) + ht_var = extract_vars(wrfin, timeidx, ht_metadata_varname, method, squeeze, cache, meta=True, _key=_key) ht_metadata_var = ht_var[ht_metadata_varname] @@ -531,13 +704,46 @@ def set_height_metadata(geopt=False): dims=outdimnames, coords=outcoords, attrs=outattrs) return func_wrapper -def _set_horiz_meta(wrapped, instance, args, kwargs): - argvars = from_args(wrapped, ("field3d", "z", "desiredlev", +def _set_horiz_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped + horizontal interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + horiztontal interpolation function output with or without metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ + argvars = from_args(wrapped, ("field3d", "vert", "desiredlev", "missing"), *args, **kwargs) field3d = argvars["field3d"] - z = argvars["z"] + z = argvars["vert"] desiredloc = argvars["desiredlev"] missingval = argvars["missing"] @@ -590,15 +796,48 @@ def _set_horiz_meta(wrapped, instance, args, kwargs): return DataArray(result, name=outname, dims=outdimnames, coords=outcoords, attrs=outattrs) -def _set_cross_meta(wrapped, instance, args, kwargs): - argvars = from_args(wrapped, ("field3d", "z", "latlon", "missing", +def _set_cross_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped cross \ + section interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + cross section interpolation function output with or without metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ + argvars = from_args(wrapped, ("field3d", "vert", "latlon", "missing", "pivot_point", "angle", "start_point", "end_point", "cache"), *args, **kwargs) field3d = argvars["field3d"] - z = argvars["z"] + z = argvars["vert"] inc_latlon = argvars["latlon"] missingval = argvars["missing"] pivot_point = argvars["pivot_point"] @@ -736,7 +975,40 @@ def _set_cross_meta(wrapped, instance, args, kwargs): -def _set_line_meta(wrapped, instance, args, kwargs): +def _set_line_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped line + interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + line interpolation function output with or without metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ argvars = from_args(wrapped, ("field2d", "pivot_point", "angle", "start_point", "end_point", "latlon", "cache"), @@ -870,8 +1142,42 @@ def _set_line_meta(wrapped, instance, args, kwargs): coords=outcoords, attrs=outattrs) -def _set_vinterp_meta(wrapped, instance, args, kwargs): - argvars = from_args(wrapped, ("wrfnc", "field", "vert_coord", +def _set_vinterp_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped + vertical coordinate interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + vertical coordinate interpolation function output with or without + metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ + argvars = from_args(wrapped, ("wrfin", "field", "vert_coord", "interp_levels", "extrapolate", "field_type", "log_p", "timeidx", "method", "squeeze", @@ -917,6 +1223,40 @@ def _set_vinterp_meta(wrapped, instance, args, kwargs): def _set_2dxy_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped line + cross section interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + line cross section interpolation function output with or without + metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ argvars = from_args(wrapped, ("field3d", "xy"), *args, **kwargs) field3d = argvars["field3d"] @@ -989,6 +1329,39 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): def _set_1d_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped 1D + interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + 1D interpolation function output with or without metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ argvars = from_args(wrapped, ("field", "z_in", "z_out", "missingval"), *args, **kwargs) @@ -1042,6 +1415,39 @@ def _set_1d_meta(wrapped, instance, args, kwargs): def _set_xy_meta(wrapped, instance, args, kwargs): + """A decorator implementation that sets the metadata for a wrapped xy line + interpolation function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + wrapped: The wrapped function which in turns needs to be called by your + wrapper function. + + instance: The object to which the wrapped function was bound when it + was called. + + args: The list of positional arguments supplied when the decorated + function was called. + + kwargs: The dictionary of keyword arguments supplied when the decorated + function was called. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + xy line interpolation function output with or without metadata. + 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. + + See Also: + + :mod:`wrapt` + + """ argvars = from_args(wrapped, ("field", "pivot_point", "angle", "start_point", "end_point"), *args, **kwargs) @@ -1078,6 +1484,27 @@ def _set_xy_meta(wrapped, instance, args, kwargs): def set_interp_metadata(interp_type): + """A decorator that sets the metadata for a wrapped interpolation + function. + + If the wrapped function's *meta* argument is False, then this decorator + returns the wrapped function output without applying the metadata. + + Args: + + interp_type (:obj:`str`): The type of interpolation routine. Choices + are: 'horiz', 'cross', 'line', 'vinterp', '2dxy', '1d', + 'xy'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + interpolation function output with or without metadata. + 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -1110,19 +1537,51 @@ def set_alg_metadata(alg_ndims, refvarname, refvarndims=None, missingarg=None, stagdim=None, stagsubvar=None, units=None, description=None): - """ - alg_ndims: number of dimensions returned by the algorithm - refvarndims: number of right dimensions for the refernce var, used - when the result has less dimensions than reference - refvarname: argument name for the reference variable - missingarg: argument name for the missing value - stagdim: staggered dimension in reference - stagsubvar: the variable name to use to supply the unstaggered - dimension size + """A decorator that sets the metadata for a wrapped raw diagnostic + function. + + Args: + + alg_ndims (:obj:`int`): The number of dimensions returned by the + wrapped function. + + refvarname (:obj:`str`): The wrapped function argument name for the + reference variable. + + refvarndims (:obj:`int`, optional): The number of right dimensions for + the reference variable. This paramter is required when the + wrapped function result has less dimensions than reference. + Default is None. + missingarg (:obj:`str`, optional): The wrapped function argument name + for the missing value variable. Default is None. + + stagdim (:obj`int`, optional): The staggered dimension for the + reference. This is only needed if the reference variable is + on a staggered grid. Default is None. + + stagsubvar (:obj:`str`, optional): The wrapped function argument name + to use to supply the unstaggered dimension name for the staggered + dimension in the reference. This is needed if *stagdim* is not + None. It is primarily used for absolute vorticity. Default is None. + + units (:obj:`str`, optional): The units to use if if there is no + 'units' argument for the wrapped function. Default is None. + + description (:obj:`str`, optional): A description for the wrapped + algorithm, which is stored in the :attr:`xarray.DataArray.attrs` + attribute under the 'description' key. Default is None. + - """ + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + numerical function output with or without metadata. + 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -1160,6 +1619,13 @@ def set_alg_metadata(alg_ndims, refvarname, outattrs["units"] = _units else: outattrs["units"] = units + + else: + # Check for a units argument, if not, just ignore + _units = from_args(wrapped, ("units",), *args, **kwargs)["units"] + if _units is not None: + outattrs["units"] = _units + if description is not None: if isinstance(description, from_var): @@ -1224,9 +1690,35 @@ def set_alg_metadata(alg_ndims, refvarname, return func_wrapper -def set_uvmet_alg_metadata(units="m s-1", description="earth rotated u,v", +def set_uvmet_alg_metadata(units=None, description="earth rotated u,v", latarg="lat", windarg="u"): + """A decorator that sets the metadata for the wrapped raw UVMET diagnostic + function. + Args: + + units (:obj:`str`, optional): The units to use if if there is no + 'units' argument for the wrapped function. Default is None. + + description (:obj:`str`, optional): A description for the wrapped + algorithm, which is stored in the :attr:`xarray.DataArray.attrs` + attribute under the 'description' key. Default is None. + + latarg (:obj:'str`, optional): The wrapped function argument name for + latitude. Default is 'lat'. + + windarg (:obj:`str`, optional): The wrapped function argument name for + the u wind component. Default is 'u'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + UVMET function output with or without metadata. + 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -1247,6 +1739,10 @@ def set_uvmet_alg_metadata(units="m s-1", description="earth rotated u,v", if units is not None: outattrs["units"] = units + else: + _units = from_args(wrapped, ("units",), *args, **kwargs)["units"] + if _units is not None: + outattrs["units"] = _units if description is not None: outattrs["description"] = description @@ -1254,7 +1750,7 @@ def set_uvmet_alg_metadata(units="m s-1", description="earth rotated u,v", latvar = from_args(wrapped, latarg, *args, **kwargs)[latarg] uvar = from_args(wrapped, windarg, *args, **kwargs)[windarg] - if isinstance(uvar, DataArray): + if isinstance(uvar, DataArray) and isinstance(latvar, DataArray): # Right dims come from latvar outdims[-2:] = latvar.dims[-2:] @@ -1275,7 +1771,26 @@ def set_uvmet_alg_metadata(units="m s-1", description="earth rotated u,v", def set_cape_alg_metadata(is2d, copyarg="pres_hpa"): + """A decorator that sets the metadata for the wrapped raw CAPE diagnostic + function. + Args: + + is2d (:obj:`bool`): Set to True for the two-dimensional CAPE + calculation, False for three-dimensional CAPE. + + copyarg (:obj:`str`): The wrapped function argument to use for + copying dimension names. Default is 'pres_hpa'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + CAPE function output with or without metadata. + 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] @@ -1342,6 +1857,23 @@ def set_cape_alg_metadata(is2d, copyarg="pres_hpa"): def set_cloudfrac_alg_metadata(copyarg="pres"): + """A decorator that sets the metadata for the wrapped raw cloud fraction + diagnostic function. + + Args: + + copyarg (:obj:`str`): The wrapped function argument to use for + copying dimension names. Default is 'pres'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + cloud fraction function output with or without metadata. + 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): @@ -1389,6 +1921,18 @@ def set_cloudfrac_alg_metadata(copyarg="pres"): def set_destag_metadata(): + """A decorator that sets the metadata for the wrapped raw destaggering + function. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wrapped + destaggering function output with or without metadata. + 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. + + """ @wrapt.decorator def func_wrapper(wrapped, instance, args, kwargs): diff --git a/src/wrf/omega.py b/src/wrf/omega.py index f602c69..2c5761d 100755 --- a/src/wrf/omega.py +++ b/src/wrf/omega.py @@ -11,10 +11,62 @@ from .metadecorators import copy_and_set_metadata @copy_and_set_metadata(copy_varname="T", name="omega", description="omega", units="Pa s-1") -def get_omega(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_omega(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): + """Return Omega. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: Omega. + 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. + + """ varnames=("T", "P", "W", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] diff --git a/src/wrf/precip.py b/src/wrf/precip.py index a7f4eea..6bb3973 100755 --- a/src/wrf/precip.py +++ b/src/wrf/precip.py @@ -5,8 +5,8 @@ from .util import extract_vars __all__ = ["get_accum_precip", "get_precip_diff"] -def get_accum_precip(wrfnc, timeidx=0): - ncvars = extract_vars(wrfnc, timeidx, varnames=("RAINC", "RAINNC")) +def get_accum_precip(wrfin, timeidx=0): + ncvars = extract_vars(wrfin, timeidx, varnames=("RAINC", "RAINNC")) rainc = ncvars["RAINC"] rainnc = ncvars["RAINNC"] @@ -14,9 +14,9 @@ def get_accum_precip(wrfnc, timeidx=0): return rainsum -def get_precip_diff(wrfnc1, wrfnc2, timeidx=0): - vars1 = extract_vars(wrfnc1, timeidx, varnames=("RAINC", "RAINNC")) - vars2 = extract_vars(wrfnc2, timeidx, varnames=("RAINC", "RAINNC")) +def get_precip_diff(wrfin1, wrfin2, timeidx=0): + vars1 = extract_vars(wrfin1, timeidx, varnames=("RAINC", "RAINNC")) + vars2 = extract_vars(wrfin2, timeidx, varnames=("RAINC", "RAINNC")) rainc1 = vars1["RAINC"] rainnc1 = vars1["RAINNC"] diff --git a/src/wrf/pressure.py b/src/wrf/pressure.py index 4d960b1..25c219a 100755 --- a/src/wrf/pressure.py +++ b/src/wrf/pressure.py @@ -9,30 +9,140 @@ from .util import extract_vars, either @copy_and_set_metadata(copy_varname=either("P", "PRES"), name="pressure", description="pressure") @convert_units("pressure", "pa") -def get_pressure(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_pressure(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="Pa"): - varname = either("P", "PRES")(wrfnc) + """Return the pressure in the specified units. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'pres'. Default + is 'Pa'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The pressure in + the specified units. + 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. + + """ + varname = either("P", "PRES")(wrfin) if varname == "P": - p_vars = extract_vars(wrfnc, timeidx, ("P", "PB"), + p_vars = extract_vars(wrfin, timeidx, ("P", "PB"), method, squeeze, cache, meta=False, _key=_key) p = p_vars["P"] pb = p_vars["PB"] pres = p + pb else: - pres = extract_vars(wrfnc, timeidx, "PRES", + pres = extract_vars(wrfin, timeidx, "PRES", method, squeeze, cache, meta=False, _key=_key)["PRES"] return pres -def get_pressure_hpa(wrfnc, timeidx=0, method="cat", squeeze=True, - cache=None, meta=True, _key=None, - units="hPa"): - return get_pressure(wrfnc, timeidx, method, squeeze, cache, meta, _key, - units) +def get_pressure_hpa(wrfin, timeidx=0, method="cat", squeeze=True, + cache=None, meta=True, _key=None): + """Return the pressure in [hPa]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The pressure in + [hPa]. + 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. + + """ + return get_pressure(wrfin, timeidx, method, squeeze, cache, meta, _key, + units="hPa") diff --git a/src/wrf/projection.py b/src/wrf/projection.py index b2a72c2..95e0502 100644 --- a/src/wrf/projection.py +++ b/src/wrf/projection.py @@ -18,11 +18,39 @@ if pyngl_enabled(): if cartopy_enabled(): class MercatorWithLatTS(crs.Mercator): + """A :class:`cartopy.crs.Mercator` subclass that adds support for + a latitude of true scale parameter. + + See Also: + + :class:`cartopy.crs.Mercator` + + """ def __init__(self, central_longitude=0.0, latitude_true_scale=0.0, min_latitude=-80.0, max_latitude=84.0, globe=None): + """Initialize a :class:`wrf.MercatorWithLatTS` object. + + Args: + + central_longitude (:obj:`float`, optional): The central + longitude. Default is 0.0. + + latitude_true_scale (:obj:`float`, optional): The latitude + of true scale. Default is 0.0. + + min_latitude (:obj:`float`, optional): The maximum southerly + extent of the projection. Default is -80.0. + + max_latitude (:obj:`float`, optional): The maximum northerly + extent of the projection. Default is 84.0.. + + globe (:class:`cartopy.crs.Globe`, optional): A globe object. + If omitted, a default globe is created. + + """ proj4_params = [("proj", "merc"), ("lon_0", central_longitude), ("lat_ts", latitude_true_scale), @@ -51,11 +79,83 @@ if cartopy_enabled(): self._threshold = np.diff(self.x_limits)[0] / 720 def _ismissing(val): + """Return True if a value is None, greater than 90.0, or less than -90. + + This function is used to check for invalid latitude values. + + Args: + + val (numeric): A numeric value. + + Returns: + + :obj:`bool`: True if the value is None, greater than 90.0, or less + than -90.0. Otherwise, False is returned. + + """ return val is None or val > 90. or val < -90. class WrfProj(object): + """A base class for storing map projection information from WRF data. + + Subclasses of this type will be stored in the 'projection' attribute + entry within a :attr:`xarray.DataArray.attrs` dictionary. This base class + contains the methods required to extract the appropriate projection class + for PyNGL, matplotlib basemap, and cartopy. + + Attributes: + + ll_lat (:obj:`float): Lower left corner latitude. + + ll_lat (:obj:`float): Lower left corner longitude. + + ur_lat (:obj:`float): Upper right corner latitude. + + ur_lon (:obj:`float): Upper right corner longitude. + + bottom_left (indexable sequence): A pair of (ll_lat, ll_lon). + + top_right (indexable sequence): A pair of (ur_lat, ur_lon). + + """ def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Initialize a :class:`wrf.WrfProj` object. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner + as a (latitude, longitude) pair. Must also specify *top_right* + if used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'MAP_PROJ': The map projection type as an integer. + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + """ + if bottom_left is not None and top_right is not None: self.ll_lat = bottom_left[0] self.ll_lon = bottom_left[1] @@ -119,11 +219,31 @@ class WrfProj(object): semiminor_axis=Constants.WRF_EARTH_RADIUS)) def cartopy_xlim(self): - """Return the x extents in projected coordinates (for cartopy)""" + """Return the x extents in projected coordinates for cartopy. + + Returns: + + :obj:`list`: A pair of [xmin, xmax]. + + See Also: + + :mod:`cartopy`, :mod:`matplotlib` + + """ return self._cart_extents()[0] def cartopy_ylim(self): - """Return the y extents in projected coordinates (for cartopy)""" + """Return the y extents in projected coordinates for cartopy. + + Returns: + + :obj:`list`: A pair of [ymin, ymax]. + + See Also: + + :mod:`cartopy`, :mod:`matplotlib` + + """ return self._cart_extents()[1] def __repr__(self): @@ -138,41 +258,102 @@ class WrfProj(object): return "{}({})".format(self.__class__.__name__, args) def basemap(self, resolution='l'): - """Return a mpl_toolkits.basemap.Basemap instance for the - projection""" + """Return a :class:`matplotlib.mpl_toolkits.basemap.Basemap` object + for the map projection. + + Arguments: + + resolution (:obj:`str`): The map resolution type. + + Returns: + + :class:`matplotlib.mpl_toolkits.basemap.Basemap`: A Basemap + object for the projection. + + See Also: + + :class:`matplotlib.mpl_toolkits.basemap.Basemap` + + """ if not basemap_enabled(): raise RuntimeError("'mpl_toolkits.basemap' is not " "installed or is disabled") return self._basemap(resolution) def cartopy(self): - """Return a cartopy.crs.Projection subclass for the - projection""" + """Return a :class:`cartopy.crs.Projection` subclass for the + map projection. + + Returns: + + :class:`cartopy.crs.Projection`: A Projection subclass for the + map projection. + + See Also: + + :class:`cartopy.crs.Projection` + + """ if not cartopy_enabled(): raise RuntimeError("'cartopy' is not " "installed or is disabled") return self._cartopy() def pyngl(self): - """Return the PyNGL resources for the projection""" + """Return a :class:`Ngl.Resources` object for the map projection. + + Returns: + + :class:`Ngl.Resources`: A dict-like object that contains the + PyNGL resources for the map projection. + + See Also: + + `PyNGL `_ + + """ if not pyngl_enabled(): raise RuntimeError("'pyngl' is not " "installed or is disabled") return self._pyngl() def proj4(self): - """Return the proj4 string for the map projection""" + """Return the PROJ.4 string for the map projection. + + Returns: + + :obj:`str`: A string suitable for use with the PROJ.4 library. + + See Also: + + PROJ.4 `_ + + """ return self._proj4() def cf(self): """Return a dictionary with the NetCDF CF parameters for the - projection""" + projection. + + Returns: + + :obj:`dict`: A dictionary with the NetCDF CF parameter names and + projection parameter values. + + """ return self._cf_params() # Used for 'missing' projection values during the 'join' method class NullProjection(WrfProj): + """A :class:`wrf.WrfProj` subclass for empty projections. + + The :class:`NullProjection` is primarily used for creating missing + projections when using the 'join' method. + + """ def __init__(self): + """Initialize a :class:`wrf.NullProjection` object.""" pass def __repr__(self): @@ -180,8 +361,51 @@ class NullProjection(WrfProj): class LambertConformal(WrfProj): + """A :class:`wrf.WrfProj` subclass for Lambert Conformal Conic projections. + + See Also: + + :class:`wrf.WrfProj`, :class:`wrf.LatLon`, + :class:`wrf.PolarStereographic`, + :class:`Mercator`, :class:`RotatedLatLon` + + """ def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Initialize a :class:`wrf.LambertConformal` object. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner + as a (latitude, longitude) pair. Must also specify *top_right* + if used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + """ super(LambertConformal, self).__init__(bottom_left, top_right, lats, lons, **proj_params) @@ -283,8 +507,51 @@ class LambertConformal(WrfProj): return _proj4 class Mercator(WrfProj): + """A :class:`wrf.WrfProj` subclass for Mercator projections. + + See Also: + + :class:`wrf.WrfProj`, :class:`wrf.LatLon`, + :class:`wrf.PolarStereographic`, + :class:`RotatedLatLon`, :class:`LambertConformal` + + """ def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Initialize a :class:`wrf.Mercator` object. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner + as a (latitude, longitude) pair. Must also specify *top_right* + if used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + """ super(Mercator, self).__init__(bottom_left, top_right, lats, lons, **proj_params) @@ -383,8 +650,52 @@ class Mercator(WrfProj): return _proj4 class PolarStereographic(WrfProj): + """A :class:`wrf.WrfProj` subclass for Polar Stereographic projections. + + See Also: + + :class:`wrf.WrfProj`, :class:`wrf.LatLon`, + :class:`wrf.RotatedLatLon`, + :class:`Mercator`, :class:`LambertConformal` + + """ + def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Initialize a :class:`wrf.PolarStereographic` object. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner + as a (latitude, longitude) pair. Must also specify *top_right* + if used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + """ super(PolarStereographic, self).__init__(bottom_left, top_right, lats, lons, **proj_params) self._hemi = -90. if self.truelat1 < 0 else 90. @@ -482,8 +793,51 @@ class PolarStereographic(WrfProj): class LatLon(WrfProj): + """A :class:`wrf.WrfProj` subclass for Lat Lon projections. + + See Also: + + :class:`wrf.WrfProj`, :class:`wrf.RotatedLatLon`, + :class:`wrf.PolarStereographic`, + :class:`Mercator`, :class:`LambertConformal` + + """ def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Initialize a :class:`wrf.LatLon` object. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner + as a (latitude, longitude) pair. Must also specify *top_right* + if used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + """ super(LatLon, self).__init__(bottom_left, top_right, lats, lons, **proj_params) @@ -563,8 +917,7 @@ class LatLon(WrfProj): # 4) In basemap, lon_0 should be set to the central (standard) longitude. # 5) In either cartopy, basemap or pyngl, I'm not sure that projections with # a pole_lon not equal to 0 or 180 can be plotted. Hopefully people -# follow the WPS instructions, otherwise I need to see a sample file and -# a lot of rum. +# follow the WPS instructions, otherwise I need to see a sample file. # 6) For items in 3 - 4, the "longitude" (lon_0 or pole_longitude) is # determined by WRF's # STAND_LON values, with the following calculations based on hemisphere: @@ -583,8 +936,51 @@ class LatLon(WrfProj): # to keep things in the -180 to 180, -90 to 90 range. # 12) This projection makes me sad. class RotatedLatLon(WrfProj): + """A :class:`wrf.WrfProj` subclass for Rotated Lat Lon projections. + + See Also: + + :class:`wrf.WrfProj`, :class:`wrf.LatLon`, + :class:`wrf.PolarStereographic`, + :class:`Mercator`, :class:`LambertConformal` + + """ def __init__(self, bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Initialize a :class:`wrf.RotatedLatLon` object. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner + as a (latitude, longitude) pair. Must also specify *top_right* + if used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + """ super(RotatedLatLon, self).__init__(bottom_left, top_right, lats, lons, **proj_params) @@ -717,6 +1113,49 @@ class RotatedLatLon(WrfProj): def getproj(bottom_left=None, top_right=None, lats=None, lons=None, **proj_params): + """Return a :class:`wrf.WrfProj` subclass. + + This functions serves as a factory function for returning a + :class:`wrf.WrfProj` subclass from the specified map projection parameters. + + Args: + + bottom_left (indexable sequence, optional): The lower left corner as + a (latitude, longitude) pair. Must also specify *top_right* if + used. Default is None. + + top_right (indexable sequence): The upper right corner as a + (latitude, longitude) pair. Must also specify *bottom_left* + if used. Default is None. + + lats (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the latitude values. Must + also specify *lons* if used. Default is None. + + lons (:class:`numpy.ndarray`, optional): An array of at least + two dimensions containing all of the longitude values. Must + also specify *lats* if used. Default is None. + + **proj_params: Map projection optional keyword arguments, that + have the same names as found in WRF output NetCDF global + attributes: + + - 'MAP_PROJ': The map projection type as an integer. + - 'CEN_LAT': Center latitude. + - 'CEN_LON': Center longitude. + - 'TRUELAT1': True latitude 1. + - 'TRUELAT2': True latitude 2. + - 'MOAD_CEN_LAT': Mother of all domains center latitude. + - 'STAND_LON': Standard longitude. + - 'POLE_LAT': Pole latitude. + - 'POLE_LON': Pole longitude. + + Returns: + + :class:`wrf.WrfProj`: A :class:`wrf.WrfProj` subclass for the + specified map projection parameters. + + """ proj_type = proj_params.get("MAP_PROJ", 0) if proj_type == ProjectionTypes.LAMBERT_CONFORMAL: diff --git a/src/wrf/pw.py b/src/wrf/pw.py index d536a15..031fd23 100755 --- a/src/wrf/pw.py +++ b/src/wrf/pw.py @@ -13,10 +13,62 @@ from .metadecorators import copy_and_set_metadata description="precipitable water", MemoryOrder="XY", units="kg m-2") -def get_pw(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_pw(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): + """Return the preciptiable water. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The preciptable + water. 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. + + """ varnames=("T", "P", "PB", "PH", "PHB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] diff --git a/src/wrf/py3compat.py b/src/wrf/py3compat.py index d8004fc..84b8ef8 100644 --- a/src/wrf/py3compat.py +++ b/src/wrf/py3compat.py @@ -6,6 +6,17 @@ from math import floor, copysign # Dictionary python 2-3 compatibility stuff def viewitems(d): + """Return either the items or viewitems method for a dictionary. + + Args: + + d (:obj:`dict`): A dictionary. + + Returns: + + view method: Either the items or viewitems method. + + """ func = getattr(d, "viewitems", None) if func is None: func = d.items @@ -13,6 +24,17 @@ def viewitems(d): def viewkeys(d): + """Return either the keys or viewkeys method for a dictionary. + + Args: + + d (:obj:`dict`): A dictionary. + + Returns: + + view method: Either the keys or viewkeys method. + + """ func = getattr(d, "viewkeys", None) if func is None: func = d.keys @@ -20,12 +42,34 @@ def viewkeys(d): def viewvalues(d): + """Return either the values or viewvalues method for a dictionary. + + Args: + + d (:obj:`dict`): A dictionary. + + Returns: + + view method: Either the values or viewvalues method. + + """ func = getattr(d, "viewvalues", None) if func is None: func = d.values return func() def isstr(s): + """Return True if the object is a string type. + + Args: + + s (string): A string (str, unicode, bytes). + + Returns: + + :obj:`bool`: True if the object is a type of string. Otherwise, False. + + """ try: return isinstance(s, basestring) except NameError: @@ -34,12 +78,48 @@ def isstr(s): # Python 2 rounding behavior def _round2(x, d=None): + """Return the result of Python 2.x rounding, which is to round the number + to the nearest integer. + + Python 3.x uses banker's rounding, which is not applicable for nearest + neighbor approaches with grid boxes. + + Args: + + x (:obj:`float`): A number, usually a float. + + d (:obj:`int`, optional): The number of digits. Default is None, + which indicates the nearest integer. + + Returns: + + :obj:`float`: The rounded number. + + """ d = 0 if d is None else d p = 10 ** d return float(floor((x * p) + copysign(0.5, x)))/p def py2round(x, d=None): + """Return the result of Python 2.x rounding, which is to round the number + to the nearest integer. + + Python 3.x uses banker's rounding, which is not applicable for nearest + neighbor approaches with grid boxes. + + Args: + + x (:obj:`float`): A number, usually a float. + + d (:obj:`int`, optional): The number of digits. Default is None, + which indicates the nearest integer. + + Returns: + + :obj:`float`: The rounded number. + + """ if version_info >= (3,): return _round2(x, d) @@ -47,6 +127,19 @@ def py2round(x, d=None): def py3range(*args): + """Return the equivalent of the range function in Python 3.x. + + For Python 2.x, this is the same as the xrange function. + + Args: + + *args: The function arguments for range or xrange. + + Returns: + + iterable: An iterable sequence. + + """ if version_info >= (3,): return range(*args) @@ -54,6 +147,21 @@ def py3range(*args): def ucode(*args, **kwargs): + """Return a Python 3.x unicode string. + + For Python 2.x, this is accomplished by using the unicode function. + + Args: + + *args: The function positional arguments for str or unicode. + + **kwargs: The function keyword arguments for str or unicode. + + Returns: + + string: A unicode string. + + """ if version_info >= (3, ): return str(*args, **kwargs) diff --git a/src/wrf/rh.py b/src/wrf/rh.py index fc2292b..f38fa58 100755 --- a/src/wrf/rh.py +++ b/src/wrf/rh.py @@ -11,10 +11,62 @@ from .metadecorators import copy_and_set_metadata @copy_and_set_metadata(copy_varname="T", name="rh", description="relative humidity", units="%") -def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_rh(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): + """Return the relative humidity. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The relative + humidity. 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. + + """ varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] @@ -33,10 +85,62 @@ def get_rh(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, @copy_and_set_metadata(copy_varname="T2", name="rh2", description="2m relative humidity", units="%") -def get_rh_2m(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_rh_2m(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): + """Return the 2m relative humidity. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The 2m relative + humidity. 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. + + """ varnames=("T2", "PSFC", "Q2") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t2 = ncvars["T2"] psfc = ncvars["PSFC"] diff --git a/src/wrf/routines.py b/src/wrf/routines.py index 21df24f..547ae53 100644 --- a/src/wrf/routines.py +++ b/src/wrf/routines.py @@ -65,8 +65,8 @@ _FUNC_MAP = {"cape2d" : get_2dcape, "pres" : get_pressure, "wspd_wdir" : get_destag_wspd_wdir, "wspd_wdir10" : get_destag_wspd_wdir10, - "wspd_wdir_uvmet" : get_uvmet_wspd_wdir, - "wspd_wdir_uvmet10" : get_uvmet10_wspd_wdir, + "uvmet_wspd_wdir" : get_uvmet_wspd_wdir, + "uvmet10_wspd_wdir" : get_uvmet10_wspd_wdir, "ctt" : get_ctt, "cloudfrac" : get_cloudfrac } @@ -108,8 +108,8 @@ _VALID_KARGS = {"cape2d" : ["missing"], "pressure" : ["units"], "wspd_wdir" : ["units"], "wspd_wdir10" : ["units"], - "wspd_wdir_uvmet" : ["units"], - "wspd_wdir_uvmet10" : ["units"], + "uvmet_wspd_wdir" : ["units"], + "uvmet10_wspd_wdir" : ["units"], "ctt" : [], "cloudfrac" : [], "default" : [] @@ -131,7 +131,9 @@ _ALIASES = {"cape_2d" : "cape2d", "updraft_helicity" : "uhel", "td" : "dp", "td2" : "dp2m", - "cfrac" : "cloudfrac" + "cfrac" : "cloudfrac", + "wspd_wdir_uvmet" : "uvmet_wspd_wdir", + "wspd_wdir_uvmet10" : "uvmet10_wspd_wdir" } class ArgumentError(Exception): @@ -151,155 +153,134 @@ def _undo_alias(alias): def _check_kargs(var, kargs): for arg in viewkeys(kargs): if arg not in _VALID_KARGS[var]: - raise ArgumentError("'%s' is an invalid keyword " + raise ValueError("'%s' is an invalid keyword " "argument for '%s'" % (arg, var)) -def getvar(wrfnc, varname, timeidx=0, +def getvar(wrfin, varname, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, - **kargs): + **kwargs): """Returns basic diagnostics from the WRF ARW model output. - Below is a table of available diagnostics. - - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | Variable Name | Description | Units | Additional Keyword Arguments | - +====================+===============================================================+=====================+===============================================================================================+ - | avo | Absolute Vorticity | 10-5 s-1 | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | eth/theta_e | Equivalent Potential Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | cape_2d | 2D cape (mcape/mcin/lcl/lfc) | J/kg / J/kg / m / m | missing: Fill value for output only (float) | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | cape_3d | 3D cape and cin | J/kg | missing: Fill value for output only (float) | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | ctt | Cloud Top Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | cloudfrac | Cloud Fraction | % | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | dbz | Reflectivity | dBz | do_variant: Set to True to enable variant calculation. Default is False. | - | | | | do_liqskin : Set to True to enable liquid skin calculation. Default is False. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | mdbz | Maximum Reflectivity | dBz | do_variant: Set to True to enable variant calculation. Default is False. | - | | | | do_liqskin: Set to True to enable liquid skin calculation. Default is False. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | geopt/geopotential | Full Model Geopotential | m2 s-2 | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | helicity | Storm Relative Helicity | m-2/s-2 | top: The top level for the calculation in meters (float). Default is 3000.0. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | lat | Latitude | decimal degrees | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | lon | Longitude | decimal degrees | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | omg/omega | Omega | Pa/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | p/pres | Full Model Pressure | Pa | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | pressure | Full Model Pressure | hPa | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | pvo | Potential Vorticity | PVU | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | pw | Precipitable Water | kg m-2 | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | rh2 | 2m Relative Humidity | % | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | slp | Sea Level Pressure | hPa | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | ter | Model Terrain Height | m | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | td2 | 2m Dew Point Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | td | Dew Point Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | tc | Temperature | C | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | th/theta | Potential Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | tk | Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | times | Times in the File or Sequence | | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | tv | Virtual Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | twb | Wet Bulb Temperature | K | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | updraft_helicity | Updraft Helicity | m-2/s-2 | bottom: The bottom level for the calculation in meters (float). Default is 2000.0. | - | | | | top: The top level for the calculation in meters (float). Default is 5000.0. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | ua | U-component of Wind on Mass Points | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | va | V-component of Wind on Mass Points | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | wa | W-component of Wind on Mass Points | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | uvmet10 | 10 m U and V Components of Wind Rotated to Earth Coordinates | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | uvmet | U and V Components of Wind Rotated to Earth Coordinates | m/s | | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - | z/height | Full Model Height | m | msl: Set to False to return AGL values. Otherwise, MSL. Default is True. | - +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ - + A table of all available diagnostics is below. + + .. include:: ../../_templates/product_table.txt - Parameters - ---------- - wrfnc : `netCD4F.Dataset`, `Nio.NioFile`, or a sequence - Input WRF ARW NetCDF data as a `netCDF4.Dataset`, `Nio.NioFile` or an - iterable sequence of the aforementioned types. + Args: + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input WRF ARW NetCDF + data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` + or an iterable sequence of the aforementioned types. - varname : str - The variable name. + varname (:obj:`str`) : The variable name. - timeidx : int or `wrf.ALL_TIMES`, optional - The desired time index. This value can be a positive integer, - negative integer, or `wrf.ALL_TIMES` (an alias for None) to return - all times in the file or sequence. The default is 0. + 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 : {'cat', 'join'}, optional - The aggregation method to use for sequences, either 'cat' or 'join'. - 'cat' combines the data along the Time index. 'join' is creates a new - index for the file. The default is 'cat'. + 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 : {True, False}, optional - Set to False to prevent dimensions with a size of 1 from being removed - from the shape of the output. Default is True. + 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 : dict, optional - A dictionary of (varname, ndarray) which can be used to supply - pre-extracted NetCDF variables to the computational routines. This can - be used to prevent the repeated variable extraction from large - sequences of data files. Default is None. + 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 : {True, False}, optional - Set to False to disable metadata and return `numpy.ndarray` instead of - `xarray.DataArray`. 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. + + **kwargs: Optional keyword arguments for certain diagnostics. + See table above. - Returns - ------- - `xarray.DataArray` or `numpy.ndarray` - Returns the specified diagnostics output. If xarray is enabled and - the meta parameter is True, then the result will be an - `xarray.DataArray` object. Otherwise, the result will be a - `numpy.ndarray` object with no metadata. + 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. + + Raises: + :class:`ValueError`: Raised when an invalid diagnostic type or + keyword argument is passed to the routine. + :class:`FortranError`: Raised when a problem occurs during a Fortran + calculation. + + See Also: + + :class:`numpy.ndarray`, :class:`xarray.DataArray` + + + Examples: + Using netCDF4 + + .. code-block:: python + + from netCDF4 import Dataset + from wrf import getvar + + wrfnc = Dataset("wrfout_d02_2010-06-13_21:00:00") + slp = getvar(wrfnc, "slp") + + Using PyNIO + + .. code-block:: python + + from Nio import open_file + from wrf import getvar + + wrfnc = open_file("wrfout_d02_2010-06-13_21:00:00"+".nc", "r") + slp = getvar(wrfnc, "slp") + + Using Iterables: + + .. code-block:: python + + import os + from netCDF4 import Dataset + from wrf import getvar + + filedir = "/path/to/wrf/files" + wrfin = [Dataset(f) for f in os.listdir(filedir) + if f.startswith("wrfout_d02_")] + + uvmet = getvar(wrfin, "uvmet", timeidx=3, units="kt") + + """ - _key = get_id(wrfnc) + _key = get_id(wrfin) - wrfnc = get_iterable(wrfnc) + wrfin = get_iterable(wrfin) - if is_standard_wrf_var(wrfnc, varname) and varname != "Times": - return extract_vars(wrfnc, timeidx, varname, + if is_standard_wrf_var(wrfin, varname) and varname != "Times": + return extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta, _key)[varname] elif varname == "Times": varname = "times" # Diverting to the get_times routine actual_var = _undo_alias(varname) if actual_var not in _VALID_KARGS: - raise ArgumentError("'%s' is not a valid variable name" % (varname)) + raise ValueError("'%s' is not a valid variable name" % (varname)) - _check_kargs(actual_var, kargs) - return _FUNC_MAP[actual_var](wrfnc, timeidx, method, squeeze, cache, - meta, _key, **kargs) + _check_kargs(actual_var, kwargs) + + return _FUNC_MAP[actual_var](wrfin, timeidx, method, squeeze, cache, + meta, _key, **kwargs) diff --git a/src/wrf/slp.py b/src/wrf/slp.py index d0263f2..78ae930 100755 --- a/src/wrf/slp.py +++ b/src/wrf/slp.py @@ -15,11 +15,69 @@ from .util import extract_vars description="sea level pressure", MemoryOrder="XY") @convert_units("pressure", "hpa") -def get_slp(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_slp(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="hPa"): + + """Return the sea level pressure in the specified units. + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'slp'. Default + is 'Pa'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The sea level + pressure in the specified units. 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. + + """ varnames=("T", "P", "PB", "QVAPOR", "PH", "PHB") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] diff --git a/src/wrf/specialdec.py b/src/wrf/specialdec.py index ef911a0..cb553ad 100644 --- a/src/wrf/specialdec.py +++ b/src/wrf/specialdec.py @@ -14,8 +14,27 @@ if xarray_enabled(): def uvmet_left_iter_nocopy(alg_dtype=np.float64): - """Decorator to handle iterating over leftmost dimensions when using - multiple files and/or multiple times with the uvmet product. + """A decorator to handle iterating over the leftmost dimensions for the + uvmet diagnostic. + + For example, if a wrapped function works with three-dimensional arrays, but + the variables include a 4th leftmost dimension for 'Time', this decorator + will iterate over all times, call the 3D Fortran routine, and aggregate the + results in to a 4D output array. + + It is also important to note that the final output array is allocated + first, and then views are passed to the wrapped function so that values + do not need to get copied in to the final output array. + + Args: + + alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used + in the wrapped function. + + Returns: + + :class:`numpy.ndarray`: The aggregated uvmet output array that includes + all extra leftmost dimensions. """ @wrapt.decorator @@ -159,9 +178,27 @@ def uvmet_left_iter_nocopy(alg_dtype=np.float64): def cape_left_iter(alg_dtype=np.float64): - """Decorator to handle iterating over leftmost dimensions when using - multiple files and/or multiple times with the cape product. + """A decorator to handle iterating over the leftmost dimensions for the + cape diagnostic. + + For example, if a wrapped function works with three-dimensional arrays, but + the variables include a 4th leftmost dimension for 'Time', this decorator + will iterate over all times, call the 3D Fortran routine, and aggregate the + results in to a 4D output array. + + It is also important to note that the final output array is allocated + first, and then views are passed to the wrapped function so that values + do not need to get copied in to the final output array. + Args: + + alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used + in the wrapped function. + + Returns: + + :class:`numpy.ndarray`: The aggregated cape output array that includes + all extra leftmost dimensions. """ @wrapt.decorator @@ -289,9 +326,27 @@ def cape_left_iter(alg_dtype=np.float64): return func_wrapper def cloudfrac_left_iter(alg_dtype=np.float64): - """Decorator to handle iterating over leftmost dimensions when using - multiple files and/or multiple times with the cloudfrac product. + """A decorator to handle iterating over the leftmost dimensions for the + cloud fraction diagnostic. + + For example, if a wrapped function works with three-dimensional arrays, but + the variables include a 4th leftmost dimension for 'Time', this decorator + will iterate over all times, call the 3D Fortran routine, and aggregate the + results in to a 4D output array. + + It is also important to note that the final output array is allocated + first, and then views are passed to the wrapped function so that values + do not need to get copied in to the final output array. + + Args: + + alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used + in the wrapped function. + + Returns: + :class:`numpy.ndarray`: The aggregated cloud fraction output array + that includes all extra leftmost dimensions. """ @wrapt.decorator diff --git a/src/wrf/temp.py b/src/wrf/temp.py index 029bf1a..de2751a 100755 --- a/src/wrf/temp.py +++ b/src/wrf/temp.py @@ -11,12 +11,68 @@ from .util import extract_vars @copy_and_set_metadata(copy_varname="T", name="theta", description="potential temperature") @convert_units("temp", "k") -def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_theta(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="K"): + """Return the potential temperature. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'theta'. Default + is 'K'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + potential temperature. + 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. + + """ varnames = ("T",) - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] full_t = t + Constants.T_BASE @@ -27,13 +83,68 @@ def get_theta(wrfnc, timeidx=0, method="cat", squeeze=True, @copy_and_set_metadata(copy_varname="T", name="temp", description="temperature") @convert_units("temp", "k") -def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_temp(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="K"): - """Return the temperature in Kelvin or Celsius""" + """Return the temperature in the specified units. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'temp'. Default + is 'K'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + temperature in the specified units. + 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. + + """ varnames=("T", "P", "PB") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] @@ -49,13 +160,68 @@ def get_temp(wrfnc, timeidx=0, method="cat", squeeze=True, @copy_and_set_metadata(copy_varname="T", name="theta_e", description="equivalent potential temperature") @convert_units("temp", "K") -def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_eth(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="K"): - "Return equivalent potential temperature (Theta-e) in Kelvin" + """Return the equivalent potential temperature. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'eth'. Default + is 'K'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + equivalent potential temperature. + 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. + + """ varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] @@ -74,13 +240,68 @@ def get_eth(wrfnc, timeidx=0, method="cat", squeeze=True, @copy_and_set_metadata(copy_varname="T", name="tv", description="virtual temperature") @convert_units("temp", "k") -def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_tv(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="K"): - "Return the virtual temperature (tv) in Kelvin or Celsius" + """Return the virtual temperature. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'tv'. Default + is 'K'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + virtual temperature. + 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. + + """ varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] @@ -100,13 +321,68 @@ def get_tv(wrfnc, timeidx=0, method="cat", squeeze=True, @copy_and_set_metadata(copy_varname="T", name="twb", description="wetbulb temperature") @convert_units("temp", "k") -def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_tw(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="K"): - "Return the wetbulb temperature (tw)" + """Return the wetbulb temperature. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'twb'. Default + is 'K'. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + wetbulb temperature. + 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. + + """ varnames=("T", "P", "PB", "QVAPOR") - ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, + ncvars = extract_vars(wrfin, timeidx, varnames, method, squeeze, cache, meta=False, _key=_key) t = ncvars["T"] p = ncvars["P"] @@ -122,15 +398,119 @@ def get_tw(wrfnc, timeidx=0, method="cat", squeeze=True, return tw -def get_tk(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_tk(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key, + """Return the temperature in [K]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + temperature in [K]. + 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. + + """ + return get_temp(wrfin, timeidx, method, squeeze, cache, meta, _key, units="K") -def get_tc(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_tc(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - return get_temp(wrfnc, timeidx, method, squeeze, cache, meta, _key, + """Return the temperature in [degC]. + + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The + temperature in [degC]. + 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. + + """ + return get_temp(wrfin, timeidx, method, squeeze, cache, meta, _key, units="degC") diff --git a/src/wrf/terrain.py b/src/wrf/terrain.py index 54d7c59..9c4588b 100755 --- a/src/wrf/terrain.py +++ b/src/wrf/terrain.py @@ -10,11 +10,68 @@ from .util import extract_vars, either @copy_and_set_metadata(copy_varname=either("HGT", "HGT_M"), name="terrain", description="terrain height") @convert_units("height", "m") -def get_terrain(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_terrain(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=False, _key=None, units="m"): - varname = either("HGT", "HGT_M")(wrfnc) + """Return the terrain height in the specified units. - return extract_vars(wrfnc, timeidx, varname, + This functions extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'ter'. Default + is 'm'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The terrain + height in the specified units. 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. + + """ + varname = either("HGT", "HGT_M")(wrfin) + + return extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key)[varname] diff --git a/src/wrf/times.py b/src/wrf/times.py index 943dec7..9200c01 100755 --- a/src/wrf/times.py +++ b/src/wrf/times.py @@ -4,14 +4,119 @@ from __future__ import (absolute_import, division, print_function, from .util import extract_times -def get_times(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_times(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - return extract_times(wrfnc, timeidx, method, squeeze, cache, + """Return a sequence of time objects. + + This function reads the 'Times' variable and creates a sequence of + :class:`datetime.datetime` objects. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: A sequence of + :class:`datetime.datetime` objects. If *meta* is True, the sequence + will be of type :class:`xarray.DataArray`, otherwise the sequence is + :class:`numpy.ndarray`. + + """ + return extract_times(wrfin, timeidx, method, squeeze, cache, meta=meta, do_xtime=False) -def get_xtimes(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_xtimes(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - return extract_times(wrfnc, timeidx, method, squeeze, cache, + """Return a sequence of time objects. + + This function reads the 'XTIME' variable and creates a sequence of + :obj:`float` objects. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. + + _key (:obj:`int`, optional): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: A sequence of + :obj:`float` objects. If *meta* is True, the sequence will be of type + :class:`xarray.DataArray`, otherwise the sequence is + :class:`numpy.ndarray`. + + Raises: + + :class:`KeyError`: Raised if the 'XTIME' variable is not present + in the NetCDF file. + + """ + return extract_times(wrfin, timeidx, method, squeeze, cache, meta=meta, do_xtime=True) diff --git a/src/wrf/units.py b/src/wrf/units.py index b0a0734..3d7eca7 100755 --- a/src/wrf/units.py +++ b/src/wrf/units.py @@ -4,8 +4,28 @@ from __future__ import (absolute_import, division, print_function, from .constants import Constants, ConversionFactors -# Handles unit conversions that only differ by multiplication factors def _apply_conv_fact(var, vartype, var_unit, dest_unit): + """Return the variable converted to different units using a conversion + factor. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + variable. + + vartype (:obj:`str`): The type of variable. Choices are: 'wind', + 'pressure', 'temp', or 'height'. + + var_unit (:obj:`str`): The variable's current units. + + dest_unit (:obj:`str`): The desired units. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in + the desired units. + + """ if var_unit == dest_unit: return var @@ -22,6 +42,21 @@ def _apply_conv_fact(var, vartype, var_unit, dest_unit): def _to_kelvin(var, var_unit): + """Return the variable in Kelvin. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + variable. + + var_unit (:obj:`str`): The variable's current units. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in + Kelvin. + + """ if var_unit == "c": return var + Constants.CELKEL elif var_unit == "f": @@ -29,15 +64,58 @@ def _to_kelvin(var, var_unit): def _k_to_c(var): + """Return the variable in Celsius. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + variable in units of Kelvin. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in + Celsius. + + """ return var - Constants.CELKEL def _k_to_f(var): + """Return the variable in Fahrenheit. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + variable in units of Kelvin. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in + Fahrenheit. + + """ return 1.8 * _k_to_c(var) + 32.0 -# Temperature is a more complicated operation so requires functions def _apply_temp_conv(var, var_unit, dest_unit): + """Return the variable converted to different units using a temperature + conversion algorithm. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + variable. + + var_unit (:obj:`str`): The variable's current units. + + dest_unit (:obj:`str`): The desired units. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in + the desired units. + + """ if dest_unit == var_unit: return var @@ -51,6 +129,7 @@ def _apply_temp_conv(var, var_unit, dest_unit): return (_TEMP_CONV_METHODS[dest_unit])(var) +# A mapping of unit names to their dictionary key names _UNIT_ALIASES = {"mps" : "m s-1", "m/s" : "m s-1", "ms-1" : "m s-1", @@ -148,24 +227,28 @@ _UNIT_ALIASES = {"mps" : "m s-1", } +# A mapping of unit types to the avaible units _VALID_UNITS = {"wind" : ["m s-1", "kt", "mi h-1", "km h-1", "ft s-1"], "pressure" : ["pa", "hpa", "mb", "torr", "mmhg", "atm"], "temp" : ["k", "f", "c"], "height" : ["m", "km", "dm", "ft", "mi"] } +# Conversion factor map for wind from base units _WIND_BASE_FACTORS = {"kt" : ConversionFactors.MPS_TO_KTS, "km h-1" : ConversionFactors.MPS_TO_KMPH, "mi h-1" : ConversionFactors.MPS_TO_MPH, "ft s-1" : ConversionFactors.MPS_TO_FPS } +# Conversion factor map to base units _WIND_TOBASE_FACTORS = {"kt" : 1.0/ConversionFactors.MPS_TO_KTS, "km h-1" : 1.0/ConversionFactors.MPS_TO_KMPH, "mi h-1" : 1.0/ConversionFactors.MPS_TO_MPH, "ft s-1" : 1.0/ConversionFactors.MPS_TO_FPS } +# Conversion factor map for pressure from base units _PRES_BASE_FACTORS = {"hpa" : ConversionFactors.PA_TO_HPA, "mb" : ConversionFactors.PA_TO_HPA, "torr" : ConversionFactors.PA_TO_TORR, @@ -173,6 +256,7 @@ _PRES_BASE_FACTORS = {"hpa" : ConversionFactors.PA_TO_HPA, "atm" : ConversionFactors.PA_TO_ATM } +# Conversion factor map for pressure to base units _PRES_TOBASE_FACTORS = {"hpa" : 1.0/ConversionFactors.PA_TO_HPA, "mb" : 1.0/ConversionFactors.PA_TO_HPA, "torr" : 1.0/ConversionFactors.PA_TO_TORR, @@ -180,24 +264,28 @@ _PRES_TOBASE_FACTORS = {"hpa" : 1.0/ConversionFactors.PA_TO_HPA, "atm" : 1.0/ConversionFactors.PA_TO_ATM } +# Conversion factor map for height from base units _HEIGHT_BASE_FACTORS = {"km" : ConversionFactors.M_TO_KM, "dm" : ConversionFactors.M_TO_DM, "ft" : ConversionFactors.M_TO_FT, "mi" : ConversionFactors.M_TO_MILES } +# Conversion factor map for height to base units _HEIGHT_TOBASE_FACTORS = {"km" : 1.0/ConversionFactors.M_TO_KM, "dm" : 1.0/ConversionFactors.M_TO_DM, "ft" : 1.0/ConversionFactors.M_TO_FT, "mi" : 1.0/ConversionFactors.M_TO_MILES } +# Mapping of unit type to base unit type _BASE_UNITS = {"wind" : "m s-1", "pressure" : "pa", "temp" : "k", "height" : "m" } +# A mapping of unit type to a mapping of to/from base conversion factors _CONV_FACTORS = {"wind" : {"to_dest" : _WIND_BASE_FACTORS, "to_base" : _WIND_TOBASE_FACTORS}, "pressure" : {"to_dest" : _PRES_BASE_FACTORS, @@ -206,11 +294,23 @@ _CONV_FACTORS = {"wind" : {"to_dest" : _WIND_BASE_FACTORS, "to_base" : _HEIGHT_TOBASE_FACTORS} } +# A mapping of temperature type to the conversion function _TEMP_CONV_METHODS = {"c" : _k_to_c, "f" : _k_to_f } def dealias_and_clean_unit(unit): + """Return the properly cleaned and dealiased unit name. + + Args: + + unit (:obj:`str`): The unit name. + + Returns: + + :obj:`str`: A unit name suitable for dictionary key lookups. + + """ cleaned_unit = " ".join(unit.lower().split()) dealiased = _UNIT_ALIASES.get(cleaned_unit, None) @@ -218,12 +318,49 @@ def dealias_and_clean_unit(unit): def check_units(unit, unit_type): + """Raise an exception if the unit name is invalid. + + Args: + + unit (:obj:`str`): The unit name. + + unit_type (:obj:`str`): The type of unit. + + Returns: + + None + + Raises: + + :class:`ValueError`: Raised when the unit name is invalid. + + """ u_cleaned = dealias_and_clean_unit(unit) if u_cleaned not in _VALID_UNITS[unit_type]: raise ValueError("invalid unit type '%s'" % unit) def do_conversion(var, vartype, var_unit, dest_unit): + """Return the variable converted to different units. + + Args: + + var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A + variable. + + vartype (:obj:`str`): The type of variable. Choices are: 'wind', + 'pressure', 'temp', or 'height'. + + var_unit (:obj:`str`): The variable's current units. + + dest_unit (:obj:`str`): The desired units. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The variable in + the desired units. + + """ u_cleaned = dealias_and_clean_unit(dest_unit) if vartype != "temp": return _apply_conv_fact(var, vartype, var_unit.lower(), u_cleaned) diff --git a/src/wrf/util.py b/src/wrf/util.py index 55cfd5d..79df6b3 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -57,30 +57,138 @@ _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(varname): - return _COORD_PAIR_MAP[varname] +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:`strq): 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(wrfnc): - return (isinstance(wrfnc, Iterable) and not isstr(wrfnc)) +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): Input 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 "XTIME" in wrfnc.variables +def has_time_coord(wrfin): + """Return True if the input file or sequence contains the time + coordinate variable. + + The time coordinate is named 'XTIME'. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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 contains the time coordinate + variable, False otherwise. + + """ + return "XTIME" in wrfin.variables -def is_mapping(wrfnc): - return isinstance(wrfnc, Mapping) +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): Input 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__ try: argvals = getargvalues(gen.gi_frame) @@ -125,17 +233,34 @@ class TestGen(object): return self.next() -def latlon_coordvars(d): +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(d): + if name in viewkeys(ncvars): lat_coord = name break for name in _LON_COORDS: - if name in viewkeys(d): + if name in viewkeys(ncvars): lon_coord = name break @@ -143,16 +268,53 @@ def latlon_coordvars(d): 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 which returns - a new iterator from the start of the sequence when __iter__ is called""" + """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. + + Args: + + wrapped (an iterable object): Any iterable object that contains the + *__iter__* method. + + Note: + + Do not increment the wrapped iterable outside of this wrapper. + + """ def __init__(self, wrapped): 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): return _generator_copy(self._wrapped) else: @@ -161,7 +323,22 @@ class IterWrapper(Iterable): def get_iterable(wrfseq): - """Returns a resetable iterable object.""" + """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: @@ -176,85 +353,165 @@ def get_iterable(wrfseq): if isinstance(wrfseq, dict): return wrfseq else: - return dict(wrfseq) # generator/custom iterable class + return dict(wrfseq) # generator/custom iterable dict class # Helper to extract masked arrays from DataArrays that convert to NaN -def npvalues(array_type): - if not isinstance(array_type, DataArray): - result = array_type +def npvalues(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: - try: - fill_value = array_type.attrs["_FillValue"] - except KeyError: - result = array_type.values - else: - result = ma.masked_invalid(array_type.values, copy=False) - result.set_fill_value(fill_value) + 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): Input 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, wrfnc): - if is_multi_file(wrfnc): - if not is_mapping(wrfnc): - wrfnc = next(iter(wrfnc)) + def __call__(self, wrfin): + if is_multi_file(wrfin): + if not is_mapping(wrfin): + wrfin = next(iter(wrfin)) else: - entry = wrfnc[next(iter(viewkeys(wrfnc)))] + entry = wrfin[next(iter(viewkeys(wrfin)))] return self(entry) for varname in self.varnames: - if varname in wrfnc.variables: + if varname in wrfin.variables: return varname raise ValueError("{} are not valid variable names".format( self.varnames)) - - -class combine_with(object): - # Remove remove_idx first, then insert_idx is applied to removed set - def __init__(self, varname, remove_dims=None, insert_before=None, - new_dimnames=None, new_coords=None): - self.varname = varname - self.remove_dims = remove_dims - self.insert_before = insert_before - self.new_dimnames = new_dimnames if new_dimnames is not None else [] - self.new_coords = (new_coords if new_coords is not None - else OrderedDict()) - - def __call__(self, var): - new_dims = list(var.dims) - new_coords = OrderedDict(var.coords) - - if self.remove_dims is not None: - for dim in self.remove_dims: - new_dims.remove(dim) - del new_coords[dim] - - if self.insert_before is not None: - insert_idx = new_dims.index(self.insert_before) - new_dims = (new_dims[0:insert_idx] + self.new_dimnames + - new_dims[insert_idx:]) - elif self.new_dimnames is not None: - new_dims = self.new_dimnames - - if self.new_coords is not None: - new_coords.update(self.new_coords) - - return new_dims, new_coords # 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): @@ -272,7 +529,48 @@ class combine_dims(object): 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 @@ -282,14 +580,41 @@ class from_var(object): var = None if vard is not None: var = vard[self.varname] - - if not isinstance(var, DataArray): - return None - return var.attrs.get(self.attribute, None) + 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][:] @@ -312,28 +637,66 @@ def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): return False -def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), +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): Input 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(wrfseq) + latvar = latvar(wrfin) if isinstance(lonvar, either): - lonvar = lonvar(wrfseq) + lonvar = lonvar(wrfin) # In case it's just a single file - if not is_multi_file(wrfseq): - wrfseq = [wrfseq] + 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 wrfseq is not a list/tuple. - if not is_mapping(wrfseq): - wrf_iter = iter(wrfseq) + # 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(wrfseq))) - entry = wrfseq[dict_key] + 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) @@ -405,6 +768,20 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), 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 @@ -414,41 +791,123 @@ def _get_global_attr(wrfnc, attr): return val -def extract_global_attrs(wrfnc, attrs): +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): Input 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(wrfnc) + multifile = is_multi_file(wrfin) if multifile: - if not is_mapping(wrfnc): - wrfnc = next(iter(wrfnc)) + if not is_mapping(wrfin): + wrfin = next(iter(wrfin)) else: - entry = wrfnc[next(iter(viewkeys(wrfnc)))] + entry = wrfin[next(iter(viewkeys(wrfin)))] return extract_global_attrs(entry, attrs) - return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist} + 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. -def extract_dim(wrfnc, dim): - if is_multi_file(wrfnc): - if not is_mapping(wrfnc): - wrfnc = next(iter(wrfnc)) + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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 = wrfnc[next(iter(viewkeys(wrfnc)))] + entry = wrfin[next(iter(viewkeys(wrfin)))] return extract_dim(entry, dim) - d = wrfnc.dimensions[dim] + d = wrfin.dimensions[dim] if not isinstance(d, int): return len(d) #netCDF4 return d # PyNIO def _combine_dict(wrfdict, varname, timeidx, method, meta, _key): - """Dictionary combination creates a new left index for each key, then - does a cat or join for the list of files for that 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) @@ -541,6 +1000,19 @@ def _combine_dict(wrfdict, varname, timeidx, method, meta, _key): 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: @@ -560,6 +1032,18 @@ def _find_coord_names(coords): 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 @@ -577,6 +1061,36 @@ def _find_max_time_size(wrfseq): 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 @@ -742,11 +1256,40 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile, 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 @@ -776,6 +1319,39 @@ def _find_forward(wrfseq, varname, timeidx, is_moving, meta, _key): 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: @@ -809,18 +1385,94 @@ def _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key): else: comboidx += numtimes - raise IndexError("timeidx {} is out of bounds".format(timeidx)) - - -def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta, _key): + 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) - -# TODO: implement in C + 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) @@ -1017,6 +1669,21 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): 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: @@ -1024,8 +1691,56 @@ def _get_numfiles(wrfseq): return sum(1 for _ in wrf_iter) -# TODO: implement in C 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) @@ -1256,12 +1971,72 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): return outarr -def combine_files(wrfseq, varname, timeidx, is_moving=None, +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(wrfseq) + wrfseq = get_iterable(wrfin) # Dictionary is unique if is_mapping(wrfseq): @@ -1277,11 +2052,62 @@ def combine_files(wrfseq, varname, timeidx, is_moving=None, return outarr.squeeze() if squeeze else outarr -# Cache is a dictionary of already extracted variables -def _extract_var(wrfnc, varname, timeidx, is_moving, +def _extract_var(wrfin, varname, timeidx, is_moving, method, squeeze, cache, meta, _key): - # Mainly used internally so variables don't get extracted multiple times, - # particularly to copy metadata. This can be slow. + """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] @@ -1294,80 +2120,258 @@ def _extract_var(wrfnc, varname, timeidx, is_moving, return cache_var multitime = is_multi_time_req(timeidx) - multifile = is_multi_file(wrfnc) + multifile = is_multi_file(wrfin) if is_time_coord_var(varname): - return extract_times(wrfnc, timeidx, method, squeeze, cache, + 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(wrfnc, varname, _key=_key) - result = _build_data_array(wrfnc, varname, timeidx, is_moving, + 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 = wrfnc.variables[varname][timeidx,:] + result = wrfin.variables[varname][timeidx,:] result = result[np.newaxis, :] # So that no squeeze works else: - result = wrfnc.variables[varname][:] + result = wrfin.variables[varname][:] else: # Squeeze handled in this routine, so just return it - return combine_files(wrfnc, varname, timeidx, is_moving, + return combine_files(wrfin, varname, timeidx, is_moving, method, squeeze, meta, _key) return result.squeeze() if squeeze else result -def extract_vars(wrfnc, timeidx, varnames, method="cat", squeeze=True, +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. + + 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: + + :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(wrfnc, var, timeidx, None, + return {var:_extract_var(wrfin, var, timeidx, None, method, squeeze, cache, meta, _key) for var in varlist} -# Python 3 compatability + 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. + + """ return dt.datetime.strptime("".join(npbytes_to_str(timearr)), "%Y-%m-%d_%H:%M:%S") -def _file_times(wrfnc, do_xtime): +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): Input 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 = wrfnc.variables["Times"][:,:] + times = wrfin.variables["Times"][:,:] for i in py3range(times.shape[0]): yield _make_time(times[i,:]) else: - xtimes = wrfnc.variables["XTIME"][:] + xtimes = wrfin.variables["XTIME"][:] for i in py3range(xtimes.shape[0]): yield xtimes[i] -def _extract_time_map(wrfnc, timeidx, do_xtime, meta=False): +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): Input 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(wrfnc)} + for key, wrfseq in viewitems(wrfin)} -def extract_times(wrfnc, timeidx, method="cat", squeeze=True, cache=None, +def extract_times(wrfin, timeidx, method="cat", squeeze=True, cache=None, meta=False, do_xtime=False): - if is_mapping(wrfnc): - return _extract_time_map(wrfnc, timeidx, do_xtime) + + """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): Input 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(wrfnc) + multi_file = is_multi_file(wrfin) if not multi_file: - wrf_list = [wrfnc] + wrf_list = [wrfin] else: - wrf_list = wrfnc + wrf_list = wrfin try: if method.lower() == "cat": @@ -1415,22 +2419,61 @@ def extract_times(wrfnc, timeidx, method="cat", squeeze=True, cache=None, return outarr -def is_standard_wrf_var(wrfnc, var): - multifile = is_multi_file(wrfnc) +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): Input 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(wrfnc): - wrfnc = next(iter(wrfnc)) + if not is_mapping(wrfin): + wrfin = next(iter(wrfin)) else: - entry = wrfnc[next(iter(viewkeys(wrfnc)))] - return is_standard_wrf_var(entry, var) + entry = wrfin[next(iter(viewkeys(wrfin)))] + return is_standard_wrf_var(entry, varname) - return var in wrfnc.variables + return varname in wrfin.variables -def is_staggered(var, wrfnc): - we = extract_dim(wrfnc, "west_east") - sn = extract_dim(wrfnc, "south_north") - bt = extract_dim(wrfnc, "bottom_top") +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): Input 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 @@ -1438,33 +2481,51 @@ def is_staggered(var, wrfnc): return False -def get_left_indexes(ref_var, expected_dims): - """Returns the extra left side dimensions for a variable with an expected - shape. +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: - For example, if a 2D variable contains an additional left side dimension - for time, this will return the time dimension size. + :obj:`tuple`: The shape for the extra leftmost dimensions. """ - extra_dim_num = ref_var.ndim - expected_dims + extra_dim_num = var.ndim - expected_dims if (extra_dim_num == 0): return [] - return tuple([ref_var.shape[x] for x in py3range(extra_dim_num)]) + return tuple([var.shape[x] for x in py3range(extra_dim_num)]) def iter_left_indexes(dims): - """A generator which yields the iteration tuples for a sequence of - dimensions sizes. + """Yield the iteration tuples for a sequence of dimensions sizes. - For example, if an array shape is (3,3), then this will yield: + For example, if *dims* is (3,3), then this will yield: (0,0), (0,1), (1,0), (1,1) - Arguments: + This is primarily used to iterate over the leftmost index values. + + Args: + + dims (indexable sequence): A sequence of dimension sizes. + + Yields: - - dims - a sequence of dimensions sizes (e.g. ndarry.shape) + :obj:`tuple`: The leftmost indexing iteration sizes. + """ arg = [py3range(dim) for dim in dims] @@ -1473,6 +2534,27 @@ def iter_left_indexes(dims): 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 @@ -1481,8 +2563,34 @@ def get_right_slices(var, right_ndims, fixed_val=0): [slice(None)]*right_ndims) -def get_proj_params(wrfnc, timeidx=0, varname=None): - proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", +def get_proj_params(wrfin, timeidx=0, varname=None): + """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): Input 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", @@ -1495,25 +2603,49 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): if varname is not None: if not is_coordvar(varname): - coord_names = getattr(wrfnc.variables[varname], + coord_names = getattr(wrfin.variables[varname], "coordinates").split() lon_coord = coord_names[0] lat_coord = coord_names[1] else: lat_coord, lon_coord = get_coord_pairs(varname) else: - lat_coord, lon_coord = latlon_coordvars(wrfnc.variables) + lat_coord, lon_coord = latlon_coordvars(wrfin.variables) - return (wrfnc.variables[lat_coord][time_idx_or_slice,:], - wrfnc.variables[lon_coord][time_idx_or_slice,:], + return (wrfin.variables[lat_coord][time_idx_or_slice,:], + wrfin.variables[lon_coord][time_idx_or_slice,:], proj_params) def from_args(func, argnames, *args, **kwargs): - """Parses the function args and kargs looking for the desired argument - value. Otherwise, the value is taken from the default keyword argument - using the arg spec. + """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 + *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] @@ -1533,6 +2665,30 @@ def from_args(func, argnames, *args, **kwargs): 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 + *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 @@ -1554,6 +2710,30 @@ def _args_to_list2(func, args, kwargs): 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 + *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) bound.apply_defaults() @@ -1563,7 +2743,28 @@ def _args_to_list3(func, args, kwargs): # Note: Doesn't allow for **kwargs or *args def args_to_list(func, args, kwargs): - """Converts the mixed args/kwargs to a single list of args""" + """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 + *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: @@ -1573,6 +2774,33 @@ def 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) @@ -1587,6 +2815,33 @@ def _arg_location2(func, argname, args, kwargs): 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()) @@ -1601,10 +2856,28 @@ def _arg_location3(func, argname, args, kwargs): def arg_location(func, argname, args, kwargs): - """Parses the function args, kargs and signature looking for the desired - argument location (either in args, kargs, or argspec.defaults), - and returns a list containing representing all arguments in the - correct order with defaults filled in. + """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,): @@ -1616,16 +2889,44 @@ def 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_id(seq): - if not is_mapping(seq): - return id(seq) +def get_id(obj): + """Return the object id. + + The object 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. Otherwise, only the + object id is returned. + + Args: + + obj (:obj:`object`): Any object type. + + 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_mapping(obj): + return id(obj) - # For each key in the mapping, recurisvely call get_id until + # For each key in the mapping, recursively call get_id until # until a non-mapping is found - return {key : get_id(val) for key,val in viewitems(seq)} + return {key : get_id(val) for key,val in viewitems(obj)} diff --git a/src/wrf/uvmet.py b/src/wrf/uvmet.py index 0007e82..3783bc4 100755 --- a/src/wrf/uvmet.py +++ b/src/wrf/uvmet.py @@ -16,36 +16,101 @@ from .util import extract_vars, extract_global_attrs, either @convert_units("wind", "m s-1") -def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, +def _get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, ten_m=False, units ="m s-1"): - """ Return a tuple of u,v with the winds rotated in to earth space""" + """Return the u,v wind components rotated to earth coordinates. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain U_EARTH + - return_val[1,...] will contain V_EARTH + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + ten_m (:obj:`bool`, optional): Set to True to use the 10m surface + winds, rather than the three-dimensional wind field. Default is + False. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet'. Default + is 'm s-1'. + + Returns: + + :class:`numpy.ndarray`: The u,v components of the wind rotated to + earth coordinates, whose leftmost dimensions is 2 (0=U, 1=V). The + *meta* parameter is ignored for this function, so only a + :class:`numpy.ndarray` is returned. + + """ if not ten_m: - varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("U", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = destagger(u_vars[varname], -1) - varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("V", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = destagger(v_vars[varname], -2) else: - varname = either("U10", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("U10", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = (u_vars[varname] if varname == "U10" else destagger(u_vars[varname][...,0,:,:], -1)) - varname = either("V10", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("V10", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = (v_vars[varname] if varname == "V10" else destagger(v_vars[varname][...,0,:,:], -2)) - map_proj_attrs = extract_global_attrs(wrfnc, attrs="MAP_PROJ") + map_proj_attrs = extract_global_attrs(wrfin, attrs="MAP_PROJ") map_proj = map_proj_attrs["MAP_PROJ"] # 1 - Lambert @@ -75,7 +140,7 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, return result elif map_proj in (1,2): - lat_attrs = extract_global_attrs(wrfnc, attrs=("TRUELAT1", + lat_attrs = extract_global_attrs(wrfin, attrs=("TRUELAT1", "TRUELAT2")) radians_per_degree = Constants.PI/180.0 # Rotation needed for Lambert and Polar Stereographic @@ -83,10 +148,10 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, true_lat2 = lat_attrs["TRUELAT2"] try: - lon_attrs = extract_global_attrs(wrfnc, attrs="STAND_LON") + lon_attrs = extract_global_attrs(wrfin, attrs="STAND_LON") except AttributeError: try: - cen_lon_attrs = extract_global_attrs(wrfnc, attrs="CEN_LON") + cen_lon_attrs = extract_global_attrs(wrfin, attrs="CEN_LON") except AttributeError: raise RuntimeError("longitude attributes not found in NetCDF") else: @@ -95,14 +160,14 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, cen_lon = lon_attrs["STAND_LON"] - varname = either("XLAT_M", "XLAT")(wrfnc) - xlat_var = extract_vars(wrfnc, timeidx, varname, + varname = either("XLAT_M", "XLAT")(wrfin) + xlat_var = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) lat = xlat_var[varname] - varname = either("XLONG_M", "XLONG")(wrfnc) - xlon_var = extract_vars(wrfnc, timeidx, varname, + varname = either("XLONG_M", "XLONG")(wrfin) + xlon_var = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) lon = xlon_var[varname] @@ -133,11 +198,75 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, description="earth rotated u,v", two_d=False, wspd_wdir=False) -def get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_uvmet(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): + """Return the u,v wind components rotated to earth coordinates. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain U_EARTH + - return_val[1,...] will contain V_EARTH + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: - return _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, _key, + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet'. Default + is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The u,v components + of the wind rotated to earth coordinates, whose leftmost dimensions is + 2 (0=U_EARTH, 1=V_EARTH). + 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. + + """ + + return _get_uvmet(wrfin, timeidx, method, squeeze, cache, meta, _key, False, units) @@ -146,11 +275,77 @@ def get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, description="10m earth rotated u,v", two_d=True, wspd_wdir=False) -def get_uvmet10(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_uvmet10(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): + """Return the u,v components for the 10m winds rotated to earth + coordinates. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain U10_EARTH + - return_val[1,...] will contain V10_EARTH + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet10'. + Default is 'm s-1'. + + Returns: - return _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, _key, + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The u,v components + of the 10m wind rotated to earth coordinates, whose leftmost dimensions + is 2 (0=U10_EARTH, 1=V10_EARTH). + 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. + + """ + + return _get_uvmet(wrfin, timeidx, method, squeeze, cache, meta, _key, True, units) @@ -159,11 +354,78 @@ def get_uvmet10(wrfnc, timeidx=0, method="cat", squeeze=True, description="earth rotated wspd,wdir", two_d=False, wspd_wdir=True) -def get_uvmet_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_uvmet_wspd_wdir(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): + """Return the wind speed and wind direction for the wind rotated to + earth coordinates. + + This function should be used when comparing with observations. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain WSPD_EARTH + - return_val[1,...] will contain WDIR_EARTH + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: - uvmet = _get_uvmet(wrfnc, timeidx, method, squeeze, + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet_wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wind speed and + wind direction for the wind rotated to earth coordinates, whose + leftmost dimensions is 2 (0=WSPD_EARTH, 1=WDIR_EARTH). 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. + + """ + + uvmet = _get_uvmet(wrfin, timeidx, method, squeeze, cache, meta, _key, False, units) return _calc_wspd_wdir(uvmet[0,...,:,:,:], uvmet[1,...,:,:,:], @@ -175,11 +437,78 @@ def get_uvmet_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True, description="10m earth rotated wspd,wdir", two_d=True, wspd_wdir=True) -def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_uvmet10_wspd_wdir(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): + """Return the wind speed and wind direction for the 10m wind rotated to + earth coordinates. + + This function should be used when comparing with observations. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain WSPD10_EARTH + - return_val[1,...] will contain WDIR10_EARTH + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for + 'uvmet10_wspd_wdir'. Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wind speed and + wind direction for the 10m wind rotated to earth coordinates, whose + leftmost dimensions is 2 (0=WSPD10_EARTH, 1=WDIR10_EARTH). 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. + + """ - uvmet10 = _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, _key, + uvmet10 = _get_uvmet(wrfin, timeidx, method, squeeze, cache, meta, _key, True, units) return _calc_wspd_wdir(uvmet10[0,...,:,:], uvmet10[1,...,:,:], True, units) diff --git a/src/wrf/version.py b/src/wrf/version.py index bc1e6b3..d830b24 100644 --- a/src/wrf/version.py +++ b/src/wrf/version.py @@ -1,2 +1,2 @@ -__version__ = "1.0a2" +__version__ = "1.0a3" diff --git a/src/wrf/vorticity.py b/src/wrf/vorticity.py index 28b149f..3fad2a1 100755 --- a/src/wrf/vorticity.py +++ b/src/wrf/vorticity.py @@ -9,14 +9,67 @@ from .metadecorators import copy_and_set_metadata @copy_and_set_metadata(copy_varname="T", name="avo", description="absolute vorticity", units="10-5 s-1") -def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_avo(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "MAPFAC_U", + """Return the absolute vorticity. + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The absolute + vorticity. 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. + + """ + ncvars = extract_vars(wrfin, timeidx, ("U", "V", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", "F"), method, squeeze, cache, meta=False, _key=_key) - attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) + attrs = extract_global_attrs(wrfin, attrs=("DX", "DY")) u = ncvars["U"] v = ncvars["V"] msfu = ncvars["MAPFAC_U"] @@ -33,14 +86,67 @@ def get_avo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, @copy_and_set_metadata(copy_varname="T", name="pvo", description="potential vorticity", units="PVU") -def get_pvo(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, +def get_pvo(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None): - ncvars = extract_vars(wrfnc, timeidx, ("U", "V", "T", "P", + """Return the potential vorticity. + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The potential + vorticity. 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. + + """ + ncvars = extract_vars(wrfin, timeidx, ("U", "V", "T", "P", "PB", "MAPFAC_U", "MAPFAC_V", "MAPFAC_M", "F"), method, squeeze, cache, meta=False, _key=_key) - attrs = extract_global_attrs(wrfnc, attrs=("DX", "DY")) + attrs = extract_global_attrs(wrfin, attrs=("DX", "DY")) u = ncvars["U"] v = ncvars["V"] diff --git a/src/wrf/wind.py b/src/wrf/wind.py index c5dd57f..fdd91de 100755 --- a/src/wrf/wind.py +++ b/src/wrf/wind.py @@ -12,15 +12,73 @@ from .metadecorators import set_wind_metadata @convert_units("wind", "m s-1") def _calc_wspd(u, v, units="m s-1"): + """Return the wind speed. + + Args: + + u (:class:`numpy.ndarray`): The u component of the wind. + + v (:class:`numpy.ndarray`): The v component of the wind. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet_wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`numpy.ndarray`: The wind speed. + + """ return np.sqrt(u**2 + v**2) def _calc_wdir(u, v): + """Return the wind direction. + + Args: + + u (:class:`numpy.ndarray`): The u component of the wind. + + v (:class:`numpy.ndarray`): The v component of the wind. + + Returns: + + :class:`numpy.ndarray`: The wind direction. + + """ wdir = 270.0 - np.arctan2(v,u) * (180.0/Constants.PI) return np.remainder(wdir, 360.0) def _calc_wspd_wdir(u, v, two_d, units): + """Return the wind speed and wind direction. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain WSPD + - return_val[1,...] will contain WDIR + + Args: + + u (:class:`numpy.ndarray`): The u component of the wind. + + v (:class:`numpy.ndarray`): The v component of the wind. + + two_d (:obj:`bool`): Set to True if the u,v wind components are + for a two-dimensional array (no height dimension). Otherwise, + set to False. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'uvmet_wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`numpy.ndarray`: The wind speed and wind direction, whose + leftmost dimension is 2 (0=WSPD, 1=WDIR). + + """ wspd = _calc_wspd(u, v, units) wdir = _calc_wdir(u, v) @@ -51,11 +109,65 @@ def _calc_wspd_wdir(u, v, two_d, units): two_d=False, wspd_wdir=False) @convert_units("wind", "m s-1") -def get_u_destag(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_u_destag(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): - varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + """Return the u-component of the wind on mass points. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The u-component + of the wind. + 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. + + """ + varname = either("U", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = destagger(u_vars[varname], -1) @@ -69,11 +181,65 @@ def get_u_destag(wrfnc, timeidx=0, method="cat", squeeze=True, wind_ncvar=True, wspd_wdir=False) @convert_units("wind", "m s-1") -def get_v_destag(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_v_destag(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): - varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + """Return the v-component of the wind on mass points. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The v-component + of the wind. + 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. + + """ + varname = either("V", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = destagger(v_vars[varname], -2) @@ -87,10 +253,64 @@ def get_v_destag(wrfnc, timeidx=0, method="cat", squeeze=True, wind_ncvar=True, wspd_wdir=False) @convert_units("wind", "m s-1") -def get_w_destag(wrfnc, timeidx=0, method="cat", squeeze=True, +def get_w_destag(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): - w_vars = extract_vars(wrfnc, timeidx, "W", method, squeeze, cache, + """Return the w-component of the wind on mass points. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The w-component + of the wind. + 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. + + """ + w_vars = extract_vars(wrfin, timeidx, "W", method, squeeze, cache, meta=False, _key=_key) w = destagger(w_vars["W"], -3) return w @@ -101,16 +321,86 @@ def get_w_destag(wrfnc, timeidx=0, method="cat", squeeze=True, description="wspd,wdir in projection space", two_d=False, wspd_wdir=True) -def get_destag_wspd_wdir(wrfnc, timeidx=0, method="cat", +def get_destag_wspd_wdir(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): - varname = either("U", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + """Return the wind speed and wind direction for the wind in the projected + coordinate space. + + The wind speed and direction from this function will be relative to the + grid. This function should not be used to compare with observations, + instead use :meth:`wrf.uvmet10_wspd_wdir` to get the earth relative wind + speed and direction. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain WSPD + - return_val[1,...] will contain WDIR + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for 'wspd_wdir'. + Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wind speed and + wind direction for the wind in projected space, whose + leftmost dimensions is 2 (0=WSPD, 1=WDIR). 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. + + """ + varname = either("U", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = destagger(u_vars[varname], -1) - varname = either("V", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("V", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = destagger(v_vars[varname], -2) @@ -122,18 +412,88 @@ def get_destag_wspd_wdir(wrfnc, timeidx=0, method="cat", description="10m wspd,wdir in projection space", two_d=False, wspd_wdir=True) -def get_destag_wspd_wdir10(wrfnc, timeidx=0, method="cat", +def get_destag_wspd_wdir10(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, units="m s-1"): + """Return the wind speed and wind direction for the 10m wind in + projected coordinate space. + + The wind speed and direction from this function will be relative to the + grid. This function should not be used to compare with observations, + instead use :meth:`wrf.uvmet10_wspd_wdir` to get the earth relative wind + speed and direction. + + The leftmost dimension of the returned array represents two different + quantities: + + - return_val[0,...] will contain WSPD10 + - return_val[1,...] will contain WDIR10 + + This function extracts the necessary variables from the NetCDF file + object in order to perform the calculation. + + Args: + + wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ + iterable): Input 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. 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): A caching key. This is used for internal + purposes only. Default is None. + + units (:obj:`str`): The desired units. Refer to the :meth:`getvar` + product table for a list of available units for + 'wspd_wdir10'. Default is 'm s-1'. + + Returns: + + :class:`xarray.DataArray` or :class:`numpy.ndarray`: The wind speed and + wind direction for the 10m wind in projected space, whose + leftmost dimensions is 2 (0=WSPD10, 1=WDIR10). 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. + + """ - varname = either("U10", "UU")(wrfnc) - u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("U10", "UU")(wrfin) + u_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) u = (u_vars[varname] if varname == "U10" else destagger(u_vars[varname][...,0,:,:], -1)) - varname = either("V10", "VV")(wrfnc) - v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache, + varname = either("V10", "VV")(wrfin) + v_vars = extract_vars(wrfin, timeidx, varname, method, squeeze, cache, meta=False, _key=_key) v = (v_vars[varname] if varname == "V10" else destagger(v_vars[varname][...,0,:,:], -2)) diff --git a/test/reduce_file.py b/test/reduce_file.py new file mode 100644 index 0000000..35e61d9 --- /dev/null +++ b/test/reduce_file.py @@ -0,0 +1,34 @@ +from netCDF4 import Dataset as NetCDF + +f = "/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-02-25_18_00_00" +outfilename = "/Users/ladwig/Documents/wrf_files/rotated_pole_test.nc" + +in_nc = NetCDF(f, mode='r', format="NETCDF3_CLASSIC") +out_nc = NetCDF(outfilename, mode='w', format="NETCDF3_CLASSIC") + +# Copy Global Attributes +for att_name in in_nc.ncattrs(): + setattr(out_nc, att_name, getattr(in_nc, att_name)) + +# Copy Dimensions, but modify the vertical dimensions +for dimname, dim in in_nc.dimensions.iteritems(): + out_nc.createDimension(dimname, len(dim)) + +# Copy Variables and their Attributes, using the reduced vertical dimension +for varname, var in in_nc.variables.iteritems(): + if varname in ("T2", "XLAT", "XLONG", "XTIME"): + datatype = var.datatype + dimensions = var.dimensions + shape = var.shape + + new_shape = shape + + new_var = out_nc.createVariable(varname, datatype, dimensions) + + new_var[:] = var[:] + + for att_name in var.ncattrs(): + setattr(new_var, att_name, getattr(var, att_name)) + +in_nc.close() +out_nc.close() \ No newline at end of file