From 49092f6ac11af50b74c0671bd68f945cfb9974b1 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Mon, 17 Oct 2016 15:27:42 -0600 Subject: [PATCH] Added docstrings to all functions. Modified some argument names for better consistency. Added unit support to raw computational algorithms. Modified setup.py to work with readthedocs. Modifyied conf.py to use mocks for readthedocs. Modified gitignore to ignore the generated documentation folder. --- .gitignore | 1 + doc/source/_templates/product_table.txt | 238 +++ doc/source/_templates/wrf_class.rst | 32 + doc/source/api/generated/README | 1 + doc/source/api/index.rst | 55 + doc/source/conf.py | 63 +- doc/source/index.rst | 22 +- ncl_reference/getvar_table.txt | 74 - ncl_reference/product_table.txt | 77 + setup.py | 26 +- src/wrf/__init__.py | 9 - src/wrf/cache.py | 71 + src/wrf/cape.py | 135 +- src/wrf/cloudfrac.py | 64 +- src/wrf/computation.py | 1751 ++++++++++++++++++++++- src/wrf/config.py | 56 + src/wrf/coordpair.py | 139 +- src/wrf/ctt.py | 61 +- src/wrf/dbz.py | 149 +- src/wrf/decorators.py | 233 ++- src/wrf/destag.py | 33 +- src/wrf/dewpoint.py | 120 +- src/wrf/extension.py | 171 ++- src/wrf/geoht.py | 203 ++- src/wrf/helicity.py | 145 +- src/wrf/interp.py | 492 ++++--- src/wrf/interputils.py | 143 +- src/wrf/latlon.py | 430 +++++- src/wrf/latlonutils.py | 323 ++++- src/wrf/metadecorators.py | 642 ++++++++- src/wrf/omega.py | 56 +- src/wrf/precip.py | 10 +- src/wrf/pressure.py | 128 +- src/wrf/projection.py | 461 +++++- src/wrf/pw.py | 56 +- src/wrf/py3compat.py | 108 ++ src/wrf/rh.py | 112 +- src/wrf/routines.py | 237 ++- src/wrf/slp.py | 62 +- src/wrf/specialdec.py | 67 +- src/wrf/temp.py | 416 +++++- src/wrf/terrain.py | 63 +- src/wrf/times.py | 113 +- src/wrf/units.py | 141 +- src/wrf/util.py | 1639 ++++++++++++++++++--- src/wrf/uvmet.py | 381 ++++- src/wrf/version.py | 2 +- src/wrf/vorticity.py | 118 +- src/wrf/wind.py | 396 ++++- test/reduce_file.py | 34 + 50 files changed, 9547 insertions(+), 1012 deletions(-) create mode 100644 doc/source/_templates/product_table.txt create mode 100644 doc/source/_templates/wrf_class.rst create mode 100644 doc/source/api/generated/README create mode 100644 doc/source/api/index.rst delete mode 100644 ncl_reference/getvar_table.txt create mode 100644 ncl_reference/product_table.txt create mode 100644 test/reduce_file.py 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