From e0560ba4f0d92fa6dd12e99992f67fe33a73da0c Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Wed, 22 Jun 2016 16:48:38 -0600 Subject: [PATCH] Added a conda recipe. Metadata cleanup for wind and interpolation. Wind algorithms now put the u_v dimension as left-most, like in NCL. --- conda_recipe/build.sh | 9 + conda_recipe/meta.yaml | 32 + fortran/wrf_user.f90 | 58 +- fortran/wrffortran.pyf | 9 +- ncl_reference/getvar_table.txt | 74 ++ src/wrf/__init__.py | 20 +- src/wrf/api.py | 33 + src/wrf/cape.py | 2 +- src/wrf/computation.py | 67 ++ src/wrf/constants.py | 2 - src/wrf/ctt.py | 1 - src/wrf/dbz.py | 1 - src/wrf/decorators.py | 116 ++- src/wrf/destag.py | 6 +- src/wrf/dewpoint.py | 1 - src/wrf/extension.py | 39 +- src/wrf/geoht.py | 2 - src/wrf/helicity.py | 2 - src/wrf/interp.py | 265 +++++-- src/wrf/interputils.py | 21 +- src/wrf/latlon.py | 4 +- src/wrf/metadecorators.py | 477 ++++++++++-- src/wrf/omega.py | 1 - src/wrf/pressure.py | 1 - src/wrf/projection.py | 137 ++-- src/wrf/psadlookup.py | 1 - src/wrf/pw.py | 1 - src/wrf/rh.py | 1 - src/wrf/routines.py | 145 +++- src/wrf/slp.py | 1 - src/wrf/temp.py | 2 - src/wrf/terrain.py | 1 - src/wrf/times.py | 1 - src/wrf/units.py | 1 - src/wrf/util.py | 76 +- src/wrf/uvdecorator.py | 239 +++--- src/wrf/uvmet.py | 24 +- src/wrf/vorticity.py | 1 - src/wrf/wind.py | 11 +- test/ipynb/WRF_Workshop_Demo.ipynb | 340 +++++++++ test/ipynb/WRF_python_demo.ipynb | 163 +++- test/ipynb/nocopy_test.ipynb | 1147 +++++++++++++++++++++++++++- test/projtest.py | 13 +- test/utests.py | 34 +- 44 files changed, 3122 insertions(+), 460 deletions(-) create mode 100644 conda_recipe/build.sh create mode 100644 conda_recipe/meta.yaml create mode 100644 ncl_reference/getvar_table.txt create mode 100644 src/wrf/api.py create mode 100644 src/wrf/computation.py create mode 100644 test/ipynb/WRF_Workshop_Demo.ipynb diff --git a/conda_recipe/build.sh b/conda_recipe/build.sh new file mode 100644 index 0000000..caea697 --- /dev/null +++ b/conda_recipe/build.sh @@ -0,0 +1,9 @@ +#!/bin/sh + +if [ `uname` == Darwin ]; then + LDFLAGS="$LDFLAGS -undefined dynamic_lookup -bundle" +fi + +pip install . + + diff --git a/conda_recipe/meta.yaml b/conda_recipe/meta.yaml new file mode 100644 index 0000000..e12fe34 --- /dev/null +++ b/conda_recipe/meta.yaml @@ -0,0 +1,32 @@ +package: + name: wrf-python + version: "0.0.1" + +source: + git_url: https://github.com/NCAR/wrf-python.git + git_rev: tutorial_backup_5_24_2016 + +build: + number: 1 + detect_binary_files_with_prefix: true + +requirements: + build: + - numpy x.x + - gcc + - python + + run: + - numpy x.x + - python + - wrapt + +test: + requires: + - nose + imports: + - wrf + #commands: + #- cd $SRC_DIR/src/examples && for file in xy2.py; do echo $file ; nosetests $file ; done | tee pyngl-test.log + + diff --git a/fortran/wrf_user.f90 b/fortran/wrf_user.f90 index 7de0a98..0b8c475 100644 --- a/fortran/wrf_user.f90 +++ b/fortran/wrf_user.f90 @@ -691,7 +691,7 @@ END SUBROUTINE DGETIJLATLONG ! NCLFORTSTART SUBROUTINE DCOMPUTEUVMET(u, v, uvmet, longca,longcb,flong,flat, & - cen_long, cone, rpd, nx, ny, nz, nxp1, nyp1, & + cen_long, cone, rpd, nx, ny, nxp1, nyp1, & istag, is_msg_val, umsg, vmsg, uvmetmsg) IMPLICIT NONE @@ -701,22 +701,22 @@ SUBROUTINE DCOMPUTEUVMET(u, v, uvmet, longca,longcb,flong,flat, & !f2py threadsafe !f2py intent(in,out) :: uvmet - INTEGER,INTENT(IN) :: nx,ny,nz,nxp1,nyp1,istag + INTEGER,INTENT(IN) :: nx,ny,nxp1,nyp1,istag LOGICAL,INTENT(IN) :: is_msg_val - REAL(KIND=8), DIMENSION(nxp1,ny,nz), INTENT(IN):: u - REAL(KIND=8), DIMENSION(nx,nyp1,nz), INTENT(IN) :: v + REAL(KIND=8), DIMENSION(nxp1,ny), INTENT(IN):: u + REAL(KIND=8), DIMENSION(nx,nyp1), INTENT(IN) :: v REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: flong REAL(KIND=8), DIMENSION(nx,ny), INTENT(IN) :: flat REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: longca REAL(KIND=8), DIMENSION(nx,ny), INTENT(INOUT) :: longcb REAL(KIND=8), INTENT(IN) :: cen_long,cone,rpd REAL(KIND=8), INTENT(IN) :: umsg,vmsg,uvmetmsg - REAL(KIND=8), DIMENSION(nx,ny,nz,2), INTENT(OUT) :: uvmet + REAL(KIND=8), DIMENSION(nx,ny,2), INTENT(OUT) :: uvmet ! NCLEND - INTEGER :: i,j,k + INTEGER :: i,j REAL(KIND=8) :: uk,vk ! msg stands for missing value in this code @@ -745,32 +745,30 @@ SUBROUTINE DCOMPUTEUVMET(u, v, uvmet, longca,longcb,flong,flat, & END DO ! WRITE (6,FMT=*) ' computing velocities ' - DO k = 1,nz - DO j = 1,ny - DO i = 1,nx - IF (istag.EQ.1) THEN - IF (is_msg_val .AND. (u(i,j,k) .EQ. umsg .OR. v(i,j,k) .EQ. vmsg & - .OR. u(i+1,j,k) .EQ. umsg .OR. v(i,j+1,k) .EQ. vmsg)) THEN - uvmet(i,j,k,1) = uvmetmsg - uvmet(i,j,k,2) = uvmetmsg - ELSE - uk = 0.5D0* (u(i,j,k)+u(i+1,j,k)) - vk = 0.5D0* (v(i,j,k)+v(i,j+1,k)) - uvmet(i,j,k,1) = vk*longcb(i,j) + uk*longca(i,j) - uvmet(i,j,k,2) = vk*longca(i,j) - uk*longcb(i,j) - END IF + DO j = 1,ny + DO i = 1,nx + IF (istag.EQ.1) THEN + IF (is_msg_val .AND. (u(i,j) .EQ. umsg .OR. v(i,j) .EQ. vmsg & + .OR. u(i+1,j) .EQ. umsg .OR. v(i,j+1) .EQ. vmsg)) THEN + uvmet(i,j,1) = uvmetmsg + uvmet(i,j,2) = uvmetmsg ELSE - IF (is_msg_val .AND. (u(i,j,k) .EQ. umsg .OR. v(i,j,k) .EQ. vmsg)) THEN - uvmet(i,j,k,1) = uvmetmsg - uvmet(i,j,k,2) = uvmetmsg - ELSE - uk = u(i,j,k) - vk = v(i,j,k) - uvmet(i,j,k,1) = vk*longcb(i,j) + uk*longca(i,j) - uvmet(i,j,k,2) = vk*longca(i,j) - uk*longcb(i,j) - END IF + uk = 0.5D0* (u(i,j)+u(i+1,j)) + vk = 0.5D0* (v(i,j)+v(i,j+1)) + uvmet(i,j,1) = vk*longcb(i,j) + uk*longca(i,j) + uvmet(i,j,2) = vk*longca(i,j) - uk*longcb(i,j) END IF - END DO + ELSE + IF (is_msg_val .AND. (u(i,j) .EQ. umsg .OR. v(i,j) .EQ. vmsg)) THEN + uvmet(i,j,1) = uvmetmsg + uvmet(i,j,2) = uvmetmsg + ELSE + uk = u(i,j) + vk = v(i,j) + uvmet(i,j,1) = vk*longcb(i,j) + uk*longca(i,j) + uvmet(i,j,2) = vk*longca(i,j) - uk*longcb(i,j) + END IF + END IF END DO END DO diff --git a/fortran/wrffortran.pyf b/fortran/wrffortran.pyf index 9e38207..d2a2698 100644 --- a/fortran/wrffortran.pyf +++ b/fortran/wrffortran.pyf @@ -127,11 +127,11 @@ python module _wrffortran ! in integer, optional,intent(in),check(shape(lat_array,1)==ny),depend(lat_array) :: ny=shape(lat_array,1) integer intent(in) :: imsg end subroutine dgetijlatlong - subroutine dcomputeuvmet(u,v,uvmet,longca,longcb,flong,flat,cen_long,cone,rpd,nx,ny,nz,nxp1,nyp1,istag,is_msg_val,umsg,vmsg,uvmetmsg) ! in :_wrffortran:wrf_user.f90 + subroutine dcomputeuvmet(u,v,uvmet,longca,longcb,flong,flat,cen_long,cone,rpd,nx,ny,nxp1,nyp1,istag,is_msg_val,umsg,vmsg,uvmetmsg) ! in :_wrffortran:wrf_user.f90 threadsafe - real(kind=8) dimension(nxp1,ny,nz),intent(in) :: u - real(kind=8) dimension(nx,nyp1,nz),intent(in),depend(nz) :: v - real(kind=8) dimension(nx,ny,nz,2),intent(out,in),depend(nx,ny,nz) :: uvmet + real(kind=8) dimension(nxp1,ny),intent(in) :: u + real(kind=8) dimension(nx,nyp1),intent(in) :: v + real(kind=8) dimension(nx,ny,2),intent(out,in),depend(nx,ny) :: uvmet real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: longca real(kind=8) dimension(nx,ny),intent(inout),depend(nx,ny) :: longcb real(kind=8) dimension(nx,ny),intent(in),depend(nx,ny) :: flong @@ -141,7 +141,6 @@ python module _wrffortran ! in real(kind=8) intent(in) :: rpd integer, optional,intent(in),check(shape(v,0)==nx),depend(v) :: nx=shape(v,0) integer, optional,intent(in),check(shape(u,1)==ny),depend(u) :: ny=shape(u,1) - integer, optional,intent(in),check(shape(u,2)==nz),depend(u) :: nz=shape(u,2) integer, optional,intent(in),check(shape(u,0)==nxp1),depend(u) :: nxp1=shape(u,0) integer, optional,intent(in),check(shape(v,1)==nyp1),depend(v) :: nyp1=shape(v,1) integer intent(in) :: istag diff --git a/ncl_reference/getvar_table.txt b/ncl_reference/getvar_table.txt new file mode 100644 index 0000000..17caac9 --- /dev/null +++ b/ncl_reference/getvar_table.txt @@ -0,0 +1,74 @@ + +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ + | 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/src/wrf/__init__.py b/src/wrf/__init__.py index 2b5e45f..828d0c1 100755 --- a/src/wrf/__init__.py +++ b/src/wrf/__init__.py @@ -1,25 +1,11 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -import warnings - -from . import config -from .config import * -from . import routines -from .routines import * -from . import util -from .util import * -from . import interp -from .interp import * -from . import projection -from .projection import * +from . import api +from .api import * __all__ = [] -__all__.extend(routines.__all__) -__all__.extend(interp.__all__) -__all__.extend(config.__all__) -__all__.extend(util.__all__) -__all__.extend(projection.__all__) +__all__.extend(api.__all__) diff --git a/src/wrf/api.py b/src/wrf/api.py new file mode 100644 index 0000000..21779e2 --- /dev/null +++ b/src/wrf/api.py @@ -0,0 +1,33 @@ +from .config import (xarray_enabled, disable_xarray, enable_xarray, + cartopy_enabled, enable_cartopy, enable_cartopy, + basemap_enabled, disable_basemap, enable_basemap, + pyngl_enabled, enable_pyngl, disable_pyngl) +from .constants import ALL_TIMES, Constants, ConversionFactors +from .destag import destagger +from .routines import getvar +from .computation import (xy, interp1d, interp2dxy, interpz3d, slp, tk, td, rh, + uvmet, smooth2d) +from .interp import (interplevel, vertcross, interpline, vinterp) +from .util import (viewitems, viewkeys, viewvalues, isstr, py2round, + py3range, ucode, npvalues, extract_global_attrs, + extract_dim, extract_vars, extract_times, combine_files, + is_staggered, get_left_indexes, iter_left_indexes, + get_right_slices, get_proj_params) + +__all__ = [] +__all__ += ["xarray_enabled", "disable_xarray", "enable_xarray", + "cartopy_enabled", "enable_cartopy", "enable_cartopy", + "basemap_enabled", "disable_basemap", "enable_basemap", + "pyngl_enabled", "enable_pyngl", "disable_pyngl"] +__all__ += ["ALL_TIMES", "Constants", "ConversionFactors"] +__all__ += ["destagger"] +__all__ += ["getvar"] +__all__ += ["xy", "interp1d", "interp2dxy", "interpz3d", "slp", "tk", "td", "rh", + "uvmet", "smooth2d"] +__all__ += ["interplevel", "vertcross", "interpline", "vinterp"] +__all__ += ["viewitems", "viewkeys", "viewvalues", "isstr", "py2round", + "py3range", "ucode", "npvalues", "extract_global_attrs", + "extract_dim", "extract_vars", "extract_times", "combine_files", + "is_staggered", "get_left_indexes", "iter_left_indexes", + "get_right_slices", "get_proj_params"] + diff --git a/src/wrf/cape.py b/src/wrf/cape.py index 1809b4d..e0a9382 100755 --- a/src/wrf/cape.py +++ b/src/wrf/cape.py @@ -11,7 +11,6 @@ from .constants import Constants, ConversionFactors from .util import extract_vars, combine_with from .metadecorators import copy_and_set_metadata -__all__ = ["get_2dcape", "get_3dcape"] @copy_and_set_metadata(copy_varname="T", name="cape_2d", @@ -25,6 +24,7 @@ def get_2dcape(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, missing=Constants.DEFAULT_FILL): """Return the 2d fields of cape, cin, lcl, and lfc""" + varnames = ("T", "P", "PB", "QVAPOR", "PH","PHB", "HGT", "PSFC") ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache, meta=False) diff --git a/src/wrf/computation.py b/src/wrf/computation.py new file mode 100644 index 0000000..159f058 --- /dev/null +++ b/src/wrf/computation.py @@ -0,0 +1,67 @@ + +from .constants import Constants +from .extension import (_interpz3d, _interp2dxy, _interp1d, _slp, _tk, _td, + _rh, _uvmet, _smooth2d) +from .util import from_var +from .metadecorators import (set_alg_metadata, set_uvmet_alg_metadata, + set_interp_metadata) +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 get_xy(field, pivot_point, angle, start_point, end_point) + + +@set_interp_metadata("1d") +def interp1d(v_in, z_in, z_out, missingval=Constants.DEFAULT_FILL, + meta=True): + return _interp1d(v_in, z_in, z_out, missingval) + + +@set_interp_metadata("2dxy") +def interp2dxy(field3d, xy, meta=True): + return _interp2dxy(field3d, xy) + + +@set_interp_metadata("horiz") +def interpz3d(field3d, z, desiredloc, missingval=Constants.DEFAULT_FILL, + meta=True): + return _interpz3d(field3d, z, desiredloc, missingval) + + +@set_alg_metadata(2, refvarname="z", refvarndims=3, units="hpa", + description="sea level pressure") +def slp(z, t, p, q, meta=True): + return _slp(z, t, p, q) + + +@set_alg_metadata(3, refvarname="pressure", units="K", + description="temperature") +def tk(pressure, theta, meta=True): + return _tk(pressure, theta) + + +@set_alg_metadata(3, refvarname="pressure", units="degC", + description="dew point temperature") +def td(pressure, qv_in, meta=True): + return _td(pressure, qv_in) + + +@set_alg_metadata(3, refvarname="pressure", + description="relative humidity", units=None) +def rh(qv, q, t, meta=True): + return _rh(qv, q, t, meta) + + +@set_uvmet_alg_metadata() +def uvmet(u, v, lat, lon, cen_long, cone, meta=True): + return _uvmet(u, v, lat, lon, cen_long, cone) + + +@set_alg_metadata(2, + refvarname="field", + description=from_var("field", "description"), + units=from_var("field", "units")) +def smooth2d(field, passes, meta=True): + return _smooth2d(field, passes) diff --git a/src/wrf/constants.py b/src/wrf/constants.py index e096b42..4620a04 100755 --- a/src/wrf/constants.py +++ b/src/wrf/constants.py @@ -1,8 +1,6 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -__all__ = ["ALL_TIMES", "Constants", "ConversionFactors"] - ALL_TIMES = None class Constants(object): diff --git a/src/wrf/ctt.py b/src/wrf/ctt.py index 58db5f9..8185703 100644 --- a/src/wrf/ctt.py +++ b/src/wrf/ctt.py @@ -11,7 +11,6 @@ from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars -__all__ = ["get_ctt"] @copy_and_set_metadata(copy_varname="T", name="ctt", remove_dims=("bottom_top",), diff --git a/src/wrf/dbz.py b/src/wrf/dbz.py index a08603f..bb9f64d 100755 --- a/src/wrf/dbz.py +++ b/src/wrf/dbz.py @@ -9,7 +9,6 @@ from .constants import Constants from .util import extract_vars from .metadecorators import copy_and_set_metadata -__all__ = ["get_dbz", "get_max_dbz"] @copy_and_set_metadata(copy_varname="T", name="dbz", description="radar reflectivity", diff --git a/src/wrf/decorators.py b/src/wrf/decorators.py index c2eaf7c..56d3723 100644 --- a/src/wrf/decorators.py +++ b/src/wrf/decorators.py @@ -11,12 +11,10 @@ from .units import do_conversion, check_units from .util import (iter_left_indexes, viewitems, viewvalues, from_args, npvalues, py3range, combine_dims, isstr) from .config import xarray_enabled +from .constants import Constants if xarray_enabled(): from xarray import DataArray - -__all__ = ["convert_units", "handle_left_iter", "uvmet_left_iter", - "handle_casting", "handle_extract_transpose"] def convert_units(unit_type, alg_unit): """A decorator that applies unit conversion to a function's output array. @@ -89,6 +87,7 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, # Start by getting the left-most 'extra' dims extra_dims = ref_var_shape[0:extra_dim_num] + mask_output = False out_inited = False for left_idxs in iter_left_indexes(extra_dims): # Make the left indexes plus a single slice object @@ -106,6 +105,32 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, if key not in _ignore_kargs else val) for key,val in viewitems(kwargs)} + # Skip the possible empty/missing arrays for the join method + skip_missing = False + if out_inited: + for i,arg in enumerate(new_args): + if isinstance(arg, DataArray): + arr = npvalues(arg) + elif isinstance(arg, np.ndarray): + arr = arg + else: + continue + + if isinstance(arr, np.ma.MaskedArray): + if arr.mask.all(): + if isinstance(output, np.ndarray): + output[left_and_slice_idxs] = ( + Constants.DEFAULT_FILL) + else: + for arr in output: + arr[left_and_slice_idxs] = ( + Constants.DEFAULT_FILL) + skip_missing = True + mask_output = True + + if skip_missing: + continue + # Call the numerical routine result = wrapped(*new_args, **new_kargs) @@ -156,6 +181,13 @@ def handle_left_iter(ref_var_expected_dims, ref_var_idx=-1, output[i].mask[left_and_slice_idxs] = outarr.mask[:] + if mask_output: + if isinstance(output, np.ndarray): + output = ma.masked_values(output, Constants.DEFAULT_FILL) + else: + output = tuple(ma.masked_values(arr, Constants.DEFAULT_FILL) + for arr in output) + return output return func_wrapper @@ -178,7 +210,7 @@ def left_iter_nocopy(ref_var_expected_dims, 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 ndims from the right to keep for + - 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. @@ -236,8 +268,9 @@ def left_iter_nocopy(ref_var_expected_dims, if "outview" not in kwargs: outd = OrderedDict((outkey, np.empty(outdims, alg_dtype)) - for outkey in _outkeys) + for outkey in _outkeys) + mask_output = False for left_idxs in iter_left_indexes(extra_dims): # Make the left indexes plus a single slice object # The single slice will handle all the dimensions to @@ -255,6 +288,28 @@ def left_iter_nocopy(ref_var_expected_dims, if key not in _ignore_kargs else val) for key,val in viewitems(kwargs)} + # Skip the possible empty/missing arrays for the join method + skip_missing = False + for arg in new_args: + if isinstance(arg, DataArray): + arr = npvalues(arg) + elif isinstance(arg, np.ndarray): + arr = arg + else: + continue + + if isinstance(arr, np.ma.MaskedArray): + if arr.mask.all(): + + for output in viewvalues(outd): + output[left_and_slice_idxs] = ( + Constants.DEFAULT_FILL) + skip_missing = True + mask_output = True + + if skip_missing: + continue + # Insert the output views if one hasn't been provided if "outview" not in new_kargs: @@ -269,18 +324,29 @@ def left_iter_nocopy(ref_var_expected_dims, if (result.__array_interface__["data"][0] != outview.__array_interface__["data"][0]): raise RuntimeError("output array was copied") - + + if len(outd) == 1: output = next(iter(viewvalues(outd))) else: output = tuple(arr for arr in viewvalues(outd)) - + + + if cast_output: if isinstance(output, np.ndarray): output = output.astype(ref_var_dtype) else: output = tuple(arr.astype(ref_var_dtype) for arr in output) + # Mostly when used with join + if mask_output: + if isinstance(output, np.ndarray): + output = np.ma.masked_values(output, Constants.DEFAULT_FILL) + else: + output = tuple(masked_values(arr, Constants.DEFAULT_FILL) + for arr in output) + return output return func_wrapper @@ -422,6 +488,42 @@ def handle_extract_transpose(do_transpose=True, outviews="outview"): return result return func_wrapper + + +def check_args(refvaridx, refvarndim, rightdims): + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + + refvar = args[refvaridx] + extra_dims = refvar.ndim - refvarndim + ref_right_sizes = refvar.shape[extra_dims:] + + for i,ndim in enumerate(rightdims): + if ndim is None: + continue + + var = args[i] + + right_var_ndims = rightdims[i] + + # Check that the number of dims is correct + if (var.ndim - extra_dims != right_var_ndims): + raise ValueError("invalid number of dimensions for argument " + "{} (got {}, expected {}).".format(i, + var.ndim, + right_var_ndims + extra_dims)) + + # Check that right dimensions are lined up + if (var.shape[-right_var_ndims:] != + ref_right_sizes[-right_var_ndims:]): + raise ValueError("invalid shape for argument " + "{} (got {}, expected {})".format(i, + var.shape[-right_var_ndims:], + ref_right_sizes[-right_var_ndims:])) + + return wrapped(*args, **kwargs) + + return func_wrapper diff --git a/src/wrf/destag.py b/src/wrf/destag.py index fd8e706..a53811f 100755 --- a/src/wrf/destag.py +++ b/src/wrf/destag.py @@ -2,11 +2,12 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) from .decorators import handle_extract_transpose +from .metadecorators import set_destag_metadata -__all__ = ["destagger"] +@set_destag_metadata() @handle_extract_transpose(do_transpose=False) -def destagger(var, stagger_dim): +def destagger(var, stagger_dim, meta=False): """ De-stagger the variable. Arguments: @@ -14,6 +15,7 @@ def destagger(var, stagger_dim): - 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 """ var_shape = var.shape diff --git a/src/wrf/dewpoint.py b/src/wrf/dewpoint.py index e2659b8..4cb318c 100755 --- a/src/wrf/dewpoint.py +++ b/src/wrf/dewpoint.py @@ -7,7 +7,6 @@ from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars -__all__ = ["get_dp", "get_dp_2m"] @copy_and_set_metadata(copy_varname="QVAPOR", name="td", description="dew point temperature") diff --git a/src/wrf/extension.py b/src/wrf/extension.py index 576af76..cbdd0c5 100755 --- a/src/wrf/extension.py +++ b/src/wrf/extension.py @@ -23,9 +23,9 @@ from ._wrffortran import (dcomputetk, dinterp3dz, dinterp2dxy, dinterp1d, from .decorators import (handle_left_iter, handle_casting, handle_extract_transpose) -from .decorators import (left_iter_nocopy) +from .decorators import (left_iter_nocopy, check_args) from .util import py3range, combine_dims, _npbytes_to_str -from .uvdecorator import uvmet_left_iter, uvmet_left_iter_nocopy +from .uvdecorator import uvmet_left_iter_nocopy __all__ = [] # __all__ += ["FortranException", "computeslp", "computetk", "computetd", @@ -57,6 +57,7 @@ class FortranException(Exception): # missingval) # return result +@check_args(0, 3, (3, 3, None, None)) @left_iter_nocopy(3, 2, ref_var_idx=0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() @@ -79,6 +80,7 @@ def _interpz3d(field3d, z, desiredloc, missingval, outview=None): # xy) # return result +@check_args(0, 3, (3,)) @left_iter_nocopy(3, combine_dims([(0,-3),(1,-2)]), ref_var_idx=0, ignore_args=(1,)) @handle_casting(arg_idxs=(0,1)) @@ -104,14 +106,15 @@ def _interp2dxy(field3d, xy, outview=None): # # return result -@left_iter_nocopy(1, 1, ref_var_idx=0, ignore_args=(2,3)) +@check_args(0, 1, (1,1,None,None)) +@left_iter_nocopy(1, combine_dims([(2,0)]), ref_var_idx=0, ignore_args=(2,3)) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() def _interp1d(v_in, z_in, z_out, missingval, outview=None): if outview is None: outview = np.empty_like(z_out) - + result = dinterp1d(v_in, outview, z_in, @@ -201,7 +204,8 @@ def _interpline(field2d, xy, outview=None): # # return result - + +@check_args(0, 3, (3,3,3,3)) @left_iter_nocopy(3, 2, ref_var_idx=0) @handle_casting(arg_idxs=(0,1,2,3)) @handle_extract_transpose() @@ -246,7 +250,9 @@ def _slp(z, t, p, q, outview=None): # return result -# Note: No left iteration decorator needed with 1D arrays + +@check_args(0, 3, (3,3)) +@left_iter_nocopy(3, 3, ref_var_idx=0) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() def _tk(pressure, theta, outview=None): @@ -259,7 +265,7 @@ def _tk(pressure, theta, outview=None): theta.ravel(order="A")) result = np.reshape(result, shape, order="F") - return result + return result # Note: No left iteration decorator needed with 1D arrays # @handle_casting(arg_idxs=(0,1)) @@ -271,13 +277,15 @@ def _tk(pressure, theta, outview=None): # result = np.reshape(result, shape, order="F") # return result -# Note: No left iteration decorator needed with 1D arrays +@check_args(0, 2, (2,2)) +@left_iter_nocopy(2, 2, ref_var_idx=0) @handle_casting(arg_idxs=(0,1)) @handle_extract_transpose() def _td(pressure, qv_in, outview=None): shape = pressure.shape if outview is None: outview = np.empty_like(pressure) + result = dcomputetd(outview.ravel(order="A"), pressure.ravel(order="A"), qv_in.ravel(order="A")) @@ -297,6 +305,8 @@ def _td(pressure, qv_in, outview=None): # return result # Note: No left iteration decorator needed with 1D arrays +@check_args(0, 2, (2,2,2)) +@left_iter_nocopy(2, 2, ref_var_idx=0) @handle_casting(arg_idxs=(0,1,2)) @handle_extract_transpose() def _rh(qv, q, t, outview=None): @@ -399,11 +409,6 @@ def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, 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. - - # Make the 2D array a 3D array with 1 dimension - if u.ndim < 3: - u = u[:, :, np.newaxis] - v = v[:, :, np.newaxis] if outview is None: outdims = u.shape + (2,) @@ -425,7 +430,7 @@ def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, vmissing, uvmetmissing) - return np.squeeze(result) + return result @handle_left_iter(3,0) @handle_casting(arg_idxs=(0,1,2,3)) @@ -543,10 +548,8 @@ def computedbz(p, tk, qv, qr, qs, qg, sn0, ivarint, iliqskin): @handle_casting(arg_idxs=(0,1,2,3,4,5)) @handle_extract_transpose() def computecape(p_hpa, tk, qv, ht, ter, sfp, missing, i3dflag, ter_follow): - flip_cape = np.zeros((p_hpa.shape[0], p_hpa.shape[1], p_hpa.shape[2]), - np.float64, order="F") - flip_cin = np.zeros((p_hpa.shape[0], p_hpa.shape[1], p_hpa.shape[2]), - np.float64, order="F") + flip_cape = np.zeros(p_hpa.shape[0:3], np.float64, order="F") + flip_cin = np.zeros(p_hpa.shape[0:3], np.float64, order="F") PSADITHTE, PSADIPRS, PSADITMK = get_lookup_tables() PSADITMK = PSADITMK.T diff --git a/src/wrf/geoht.py b/src/wrf/geoht.py index d56303c..4ad0c71 100755 --- a/src/wrf/geoht.py +++ b/src/wrf/geoht.py @@ -7,8 +7,6 @@ from .decorators import convert_units from .metadecorators import set_height_metadata from .util import extract_vars, either -__all__ = ["get_geopt", "get_height"] - def _get_geoht(wrfnc, timeidx, method="cat", squeeze=True, cache=None, meta=True, height=True, msl=True): diff --git a/src/wrf/helicity.py b/src/wrf/helicity.py index d623d33..ce461c3 100755 --- a/src/wrf/helicity.py +++ b/src/wrf/helicity.py @@ -8,8 +8,6 @@ from .destag import destagger from .util import extract_vars, extract_global_attrs, either from .metadecorators import copy_and_set_metadata -__all__ = ["get_srh", "get_uh"] - @copy_and_set_metadata(copy_varname="HGT", name="srh", description="storm relative helicity", units="m-2/s-2") diff --git a/src/wrf/interp.py b/src/wrf/interp.py index e19e07e..b410447 100755 --- a/src/wrf/interp.py +++ b/src/wrf/interp.py @@ -20,42 +20,114 @@ from .geoht import get_height from .temp import get_theta, get_temp, get_eth from .pressure import get_pressure -__all__ = ["interplevel", "vertcross", "interpline", "vinterp"] # Note: Extension decorator is good enough to handle left dims @set_interp_metadata("horiz") -def interplevel(field3d, z, desiredloc, missingval=Constants.DEFAULT_FILL, +def interplevel(field3d, z, desiredlev, missing=Constants.DEFAULT_FILL, meta=True): - """Return the horizontally interpolated data at the provided level + """Interpolates a three-dimensional field specified pressure or + height level. + + Parameters + ---------- + field3d : `xarray.DataArray` or `numpy.ndarray` + A three-dimensional field. - field3d - the 3D field to interpolate - z - the vertical values (height or pressure) - desiredloc - the vertical level to interpolate at (must be same units as - zdata) - missingval - the missing data value (which will be masked on return) + 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 + ------- + out : '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. """ - r1 = _interpz3d(field3d, z, desiredloc, missingval) - masked_r1 = ma.masked_values (r1, missingval) + r1 = _interpz3d(field3d, z, desiredlev, missing) + masked_r1 = ma.masked_values (r1, missing) return masked_r1 @set_interp_metadata("cross") -def vertcross(field3d, z, missingval=Constants.DEFAULT_FILL, +def vertcross(field3d, z, missing=Constants.DEFAULT_FILL, pivot_point=None, angle=None, start_point=None, end_point=None, + include_latlon=False, cache=None, meta=True): """Return the vertical cross section for a 3D field, interpolated to a verical plane defined by a horizontal line. - Arguments: - field3d - a 3D data field - z - 3D height field - 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 tuple of (south_north1,west_east1) - end_point - an end point tuple of (south_north2,west_east2) + 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. + + include_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 + ------- + out : '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. """ @@ -67,23 +139,67 @@ def vertcross(field3d, z, missingval=Constants.DEFAULT_FILL, xy, var2dz, z_var2d = get_xy_z_params(z, pivot_point, angle, start_point, end_point) - res = _vertcross(field3d, xy, var2dz, z_var2d, missingval) + res = _vertcross(field3d, xy, var2dz, z_var2d, missing) - return ma.masked_values(res, missingval) + return ma.masked_values(res, missing) @set_interp_metadata("line") def interpline(field2d, pivot_point=None, - angle=None, start_point=None, - end_point=None, cache=None, meta=True): - """Return the 2D field interpolated along a line. + angle=None, start_point=None, + end_point=None, include_latlon=False, + cache=None, meta=True): + """Return the two-dimensional field interpolated along a line. + + Parameters + ---------- + field2d : `xarray.DataArray` or `numpy.ndarray` + A two-dimensional field. + + pivot_point : tuple or list, optional + A tuple or list with two entries, in the form of [x, y] + (or [west_east, south_north]), which indicates the x,y location + through which the plane will pass. Must also specify `angle`. - Arguments: - field2d - a 2D data field - pivot_point - a pivot point of (south_north,west_east) - angle - the angle through the pivot point in degrees - start_point - a start_point tuple of (south_north1,west_east1) - end_point - an end point tuple of (south_north2,west_east2) + 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. + + include_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 + ------- + out : '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. """ @@ -99,6 +215,79 @@ def interpline(field2d, pivot_point=None, def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, field_type=None, log_p=False, timeidx=0, method="cat", squeeze=True, cache=None, meta=True): + """Return the field vertically interpolated to the given the type of + surface and a set of new levels. + + Parameters + ---------- + wrfnc : `netCD4F.Dataset`, `Nio.NioFile`, or a sequence + Input WRF ARW NetCDF data as a `netCDF4.Dataset`, `Nio.NioFile` or an + iterable sequence of the aforementioned types. + + field2d : `xarray.DataArray` or `numpy.ndarray` + A two-dimensional field. + + vert_coord : {"pressure", "pres", "p", "ght_msl", + "ght_agl", "theta", "th", "theta-e", + "thetae", "eth"} + A string indicating the vertical coordinate type to interpolate to. + Valid strings are: + * "pressure", "pres", "p": pressure [hPa] + * "ght_msl": grid point height msl [km] + * "ght_agl": grid point height agl [km] + * "theta", "th": potential temperature [K] + * "theta-e", "thetae", "eth": equivalent potential temperature [K] + + interp_levels : sequence + A 1D sequence of vertical levels to interpolate to. + + extrapolate : {True, False}, optional + Set to True to extrapolate values below ground. Default is False. + + field_type : {"none", "pressure", "pres", "p", "z", "tc", "tk", "theta", + "th", "theta-e", "thetae", "eth", "ght"}, optional + The type of field. Default is None. + + log_p : {True, False} + Use the log of the pressure for interpolation instead of just pressure. + Default is False. + + timeidx : int, optional + The time index to use when extracting auxiallary variables used in + the interpolation. This value must be set to match the same value + used when the `field` variable was extracted. Default is 0. + + method : {'cat', 'join'}, optional + The aggregation method to use for sequences, either 'cat' or 'join'. + 'cat' combines the data along the Time index. 'join' is creates a new + index for the file. This must be set to the same method used when + extracting the `field` variable. The default is 'cat'. + + squeeze : {True, False}, optional + Set to False to prevent dimensions with a size of 1 from being removed + from the shape of the output. Default is True. + + cache : dict, optional + A dictionary of (varname, ndarray) which can be used to supply + pre-extracted NetCDF variables to the computational routines. This can + be used to prevent the repeated variable extraction from large + sequences of data files. Default is None. + + meta : {True, False}, optional + Set to False to disable metadata and return `numpy.ndarray` instead of + `xarray.DataArray`. Default is True. + + + Returns + ------- + out : '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. + + """ + # Remove case sensitivity field_type = field_type.lower() if field_type is not None else "none" vert_coord = vert_coord.lower() if vert_coord is not None else "none" @@ -241,22 +430,6 @@ def vinterp(wrfnc, field, vert_coord, interp_levels, extrapolate=False, return ma.masked_values(res, missing) -# Thin wrappers around fortran calls which allow for metadata -# Move to the new routines module -# TODO: Rename after the extensions are renamed -@set_interp_metadata("horiz") -def wrap_interpz3d(field3d, z, desiredloc, missingval, meta=True): - return _interpz3d(field3d, z, desiredloc, missingval) - - -@set_interp_metadata("2dxy") -def wrap_interp2dxy(field3d, xy, meta=True): - return _interp2dxy(field3d, xy) - - -@set_interp_metadata("1d") -def wrap_interp1d(v_in, z_in, z_out, missingval, meta=True): - return _interp1d(v_in, z_in, z_out, missingval) diff --git a/src/wrf/interputils.py b/src/wrf/interputils.py index 5bc91e2..abd157f 100644 --- a/src/wrf/interputils.py +++ b/src/wrf/interputils.py @@ -8,13 +8,12 @@ import numpy as np from .extension import _interp2dxy from .util import py3range -__all__ = ["to_positive_idxs", "calc_xy", "get_xy_z_params", "get_xy"] def to_positive_idxs(shape, coord): if (coord[-2] >= 0 and coord[-1] >= 0): return coord - return [x if (x >= 0) else shape[i]+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, start_point=None, end_point=None): @@ -25,15 +24,15 @@ def calc_xy(xdim, ydim, pivot_point=None, angle=None, 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 [south_north1, west_east1] - end_point - an end point sequence of [south_north2, west_east2] + start_point - a start_point sequence of (x, y) + end_point - an end point sequence of (x, 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[-1] - yp = pivot_point[-2] + xp = pivot_point[-2] + yp = pivot_point[-1] if (angle > 315.0 or angle < 45.0 or ((angle > 135.0) and (angle < 225.0))): @@ -97,10 +96,10 @@ def calc_xy(xdim, ydim, pivot_point=None, angle=None, y1 = ydim-1 x1 = (y1 - intercept)/slope elif start_point is not None and end_point is not None: - x0 = start_point[-1] - y0 = start_point[-2] - x1 = end_point[-1] - y1 = end_point[-2] + x0 = start_point[-2] + y0 = start_point[-1] + x1 = end_point[-2] + y1 = end_point[-1] if ( x1 > xdim-1 ): x1 = xdim if ( y1 > ydim-1): @@ -174,7 +173,7 @@ def get_xy(var, pivot_point=None, angle=None, if end_point is not None: pos_end = to_positive_idxs(var.shape[-2:], end_point) else: - pos_end = start_point + pos_end = start_point xdim = var.shape[-1] ydim = var.shape[-2] diff --git a/src/wrf/latlon.py b/src/wrf/latlon.py index cca4473..b33d95c 100755 --- a/src/wrf/latlon.py +++ b/src/wrf/latlon.py @@ -6,9 +6,6 @@ from .latlonutils import (_lat_varname, _lon_varname, ll_to_ij, ij_to_ll) from .metadecorators import set_latlon_metadata -__all__ = ["get_lat", "get_lon", "get_ij", "get_ll"] - - def get_lat(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, stagger=None): @@ -31,6 +28,7 @@ def get_lon(wrfnc, timeidx=0, method="cat", squeeze=True, return lon_var[varname] +# TODO: Rename these! @set_latlon_metadata(ij=True) def get_ij(wrfnc, latitude, longitude, timeidx=0, stagger=None, method="cat", squeeze=True, cache=None, meta=True): diff --git a/src/wrf/metadecorators.py b/src/wrf/metadecorators.py index 785f13f..0f1eed6 100644 --- a/src/wrf/metadecorators.py +++ b/src/wrf/metadecorators.py @@ -6,10 +6,11 @@ from collections import OrderedDict import numpy as np import numpy.ma as ma - +from .extension import _interpline from .util import (viewkeys, viewitems, extract_vars, combine_with, either, from_args, arg_location, - _is_coord_var, CoordPair, npvalues, py3range, ucode) + is_coordvar, latlon_coordvars, CoordPair, npvalues, + py3range, ucode, from_var, iter_left_indexes) from .interputils import get_xy_z_params, get_xy from .latlonutils import ij_to_ll, ll_to_ij from .config import xarray_enabled @@ -17,9 +18,6 @@ from .config import xarray_enabled if xarray_enabled(): from xarray import DataArray -__all__ = ["copy_and_set_metadata", "set_wind_metadata", - "set_latlon_metadata", "set_height_metadata", - "set_interp_metadata"] def copy_and_set_metadata(copy_varname=None, delete_attrs=None, name=None, remove_dims=None, dimnames=None, @@ -187,24 +185,24 @@ def set_wind_metadata(copy_varname, name, description, outattrs = OrderedDict() outdimnames = list(copy_var.dims) - outcoords.update(copy_var.coords) + #outcoords.update(copy_var.coords) outattrs.update(copy_var.attrs) - if wind_ncvar: - pass + if wind_ncvar: + outcoords.update(copy_var.coords) elif not wspd_wdir: if not two_d: - outdimnames.insert(-3, "u_v") + outdimnames.insert(0, "u_v") else: - outdimnames.insert(-2, "u_v") + outdimnames.insert(0, "u_v") outattrs["MemoryOrder"] = "XY" outcoords["u_v"] = ["u", "v"] else: if not two_d: - outdimnames.insert(-3, "wspd_wdir") + outdimnames.insert(0, "wspd_wdir") else: - outdimnames.insert(-2, "wspd_wdir") + outdimnames.insert(0, "wspd_wdir") outattrs["MemoryOrder"] = "XY" outcoords["wspd_wdir"] = ["wspd", "wdir"] @@ -212,6 +210,19 @@ def set_wind_metadata(copy_varname, name, description, if units is not None: outattrs["units"] = units + # xarray doesn't line up coordinate dimensions based on + # names, it just remembers the index it originally mapped to. + # So, need to rebuild the XLAT, XLONG, coordinates again since the + # leftmost index changed. + if not wind_ncvar: + for key,dataarray in viewitems(copy_var.coords): + if is_coordvar(key): + outcoords[key] = dataarray.dims, npvalues(dataarray) + elif key == "XTIME": + outcoords[key] = dataarray.dims, npvalues(dataarray) + elif key == "Time": + outcoords[key] = npvalues(dataarray) + outname = name outattrs["description"] = description @@ -341,15 +352,15 @@ 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", "desiredloc", - "missingval"), +def _set_horiz_meta(wrapped, instance, args, kwargs): + argvars = from_args(wrapped, ("field3d", "z", "desiredlev", + "missing"), *args, **kwargs) field3d = argvars["field3d"] z = argvars["z"] - desiredloc = argvars["desiredloc"] - missingval = argvars["missingval"] + desiredloc = argvars["desiredlev"] + missingval = argvars["missing"] result = wrapped(*args, **kwargs) @@ -384,7 +395,7 @@ def _set_horiz_meta(wrapped, instance, args, kwargs): outname = "{0}_{1}".format(field3d.name, name_levelstr) else: - outname = "field3d_{0}".format(levelstr) + outname = "field3d_{0}".format(name_levelstr) outattrs = OrderedDict() outattrs["PlotLevelID"] = levelstr @@ -400,8 +411,8 @@ 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", "missingval", +def _set_cross_meta(wrapped, instance, args, kwargs): + argvars = from_args(wrapped, ("field3d", "z", "include_latlon", "missing", "pivot_point", "angle", "start_point", "end_point", "cache"), @@ -409,7 +420,8 @@ def _set_cross_meta(wrapped, instance, args, kwargs): field3d = argvars["field3d"] z = argvars["z"] - missingval = argvars["missingval"] + inc_latlon = argvars["include_latlon"] + missingval = argvars["missing"] pivot_point = argvars["pivot_point"] angle = argvars["angle"] start_point = argvars["start_point"] @@ -447,7 +459,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs): ed_x = xy[-1,0] ed_y = xy[-1,1] - cross_str = "cross-section: ({0}, {1}) to ({2}, {3})".format(st_x, st_y, + cross_str = "Cross-Section: ({0}, {1}) to ({2}, {3})".format(st_x, st_y, ed_x, ed_y) if angle is not None: cross_str += " ; center={0} ; angle={1}".format(pivot_point, @@ -463,8 +475,8 @@ def _set_cross_meta(wrapped, instance, args, kwargs): del outcoords[field3d.dims[i]] - # Delete any lat,lon coords - delkeys = [key for key in viewkeys(outcoords) if _is_coord_var(key)] + # Delete any lat,lon coords + delkeys = [key for key in viewkeys(outcoords) if is_coordvar(key)] for key in delkeys: del outcoords[key] @@ -479,9 +491,56 @@ def _set_cross_meta(wrapped, instance, args, kwargs): del outattrs[key] except KeyError: pass + + # Interpolate to get the lat/lon coords, if desired + if inc_latlon: + latcoordname, loncoordname = latlon_coordvars(field3d.coords) - outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) - for i in py3range(xy.shape[-2])]) + if latcoordname is not None and loncoordname is not None: + latcoord = field3d.coords[latcoordname] + loncoord = field3d.coords[loncoordname] + + if latcoord.ndim == 2: + lats = _interpline(latcoord, xy) + lons = _interpline(loncoord, xy) + + outcoords["xy_loc"] = ("xy", + np.asarray(tuple( + CoordPair(x=xy[i,0], y=xy[i,1], + lat=lats[i], lon=lons[i]) + for i in py3range(xy.shape[-2]))) + ) + + else: + extra_dims = latcoord.shape[0:-2] + outdims = extra_dims + xy.shape[-2:-1] + + latlon_loc = np.empty(outdims, np.object_) + for left_dims in iter_left_indexes(extra_dims): + idxs = left_dims + (slice(None),) + lats = _interpline(latcoord[idxs], xy) + lons = _interpline(loncoord[idxs], xy) + + latlon_loc[idxs] = np.asarray(tuple( + CoordPair(x=xy[i,0], y=xy[i,1], + lat=lats[i], lon=lons[i]) + for i in py3range(xy.shape[-2])) + )[:] + + + extra_dimnames = latcoord.dims[0:-2] + loc_dimnames = extra_dimnames + ("xy",) + outcoords["xy_loc"] = (loc_dimnames, latlon_loc) + + else: + outcoords["xy_loc"] = ("xy", np.asarray(tuple( + CoordPair(xy[i,0], xy[i,1]) + for i in py3range(xy.shape[-2])))) + + else: + outcoords["xy_loc"] = ("xy", np.asarray(tuple( + CoordPair(xy[i,0], xy[i,1]) + for i in py3range(xy.shape[-2])))) outcoords["vertical"] = z_var2d[:] @@ -489,7 +548,7 @@ def _set_cross_meta(wrapped, instance, args, kwargs): outname = "field3d_cross" outattrs = OrderedDict() - outattrs["orientation"] = cross_str + outattrs["Orientation"] = cross_str outattrs["missing_value"] = missingval outattrs["_FillValue"] = missingval @@ -498,9 +557,10 @@ def _set_cross_meta(wrapped, instance, args, kwargs): -def _set_line_meta(wrapped, instance, args, kwargs): +def _set_line_meta(wrapped, instance, args, kwargs): argvars = from_args(wrapped, ("field2d", "pivot_point", "angle", - "start_point", "end_point", "cache"), + "start_point", "end_point", "include_latlon", + "cache"), *args, **kwargs) field2d = argvars["field2d"] @@ -508,6 +568,7 @@ def _set_line_meta(wrapped, instance, args, kwargs): angle = argvars["angle"] start_point = argvars["start_point"] end_point = argvars["end_point"] + inc_latlon = argvars["include_latlon"] cache = argvars["cache"] if cache is None: @@ -554,7 +615,7 @@ def _set_line_meta(wrapped, instance, args, kwargs): del outcoords[field2d.dims[i]] # Delete any lat,lon coords - delkeys = [key for key in viewkeys(outcoords) if _is_coord_var(key)] + delkeys = [key for key in viewkeys(outcoords) if is_coordvar(key)] for key in delkeys: del outcoords[key] @@ -569,20 +630,67 @@ def _set_line_meta(wrapped, instance, args, kwargs): except KeyError: pass - outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) - for i in py3range(xy.shape[-2])]) + # Interpolate to get the lat/lon coords, if desired + if inc_latlon: + latcoordname, loncoordname = latlon_coordvars(field2d.coords) + + if latcoordname is not None and loncoordname is not None: + latcoord = field2d.coords[latcoordname] + loncoord = field2d.coords[loncoordname] + + if latcoord.ndim == 2: + lats = _interpline(latcoord, xy) + lons = _interpline(loncoord, xy) + + outcoords["xy_loc"] = ("xy", + np.asarray(tuple( + CoordPair(x=xy[i,0], y=xy[i,1], + lat=lats[i], lon=lons[i]) + for i in py3range(xy.shape[-2]))) + ) + + else: + extra_dims = latcoord.shape[0:-2] + outdims = extra_dims + xy.shape[-2:-1] + + latlon_loc = np.empty(outdims, np.object_) + for left_dims in iter_left_indexes(extra_dims): + idxs = left_dims + (slice(None),) + lats = _interpline(latcoord[idxs], xy) + lons = _interpline(loncoord[idxs], xy) + + latlon_loc[idxs] = np.asarray(tuple( + CoordPair(x=xy[i,0], y=xy[i,1], + lat=lats[i], lon=lons[i]) + for i in py3range(xy.shape[-2])) + )[:] + + + extra_dimnames = latcoord.dims[0:-2] + loc_dimnames = extra_dimnames + ("xy",) + outcoords["xy_loc"] = (loc_dimnames, latlon_loc) + + else: + outcoords["xy_loc"] = ("xy", np.asarray(tuple( + CoordPair(xy[i,0], xy[i,1]) + for i in py3range(xy.shape[-2])))) + + else: + outcoords["xy_loc"] = ("xy", np.asarray(tuple( + CoordPair(xy[i,0], xy[i,1]) + for i in py3range(xy.shape[-2])))) else: outname = "field2d_line" outattrs = OrderedDict() - outattrs["orientation"] = cross_str + outattrs["Orientation"] = cross_str return DataArray(result, name=outname, dims=outdimnames, coords=outcoords, attrs=outattrs) -def _set_vinterp_meta(wrapped, instance, args, kwargs): +def _set_vinterp_meta(wrapped, instance, args, kwargs): argvars = from_args(wrapped, ("wrfnc", "field", "vert_coord", "interp_levels", "extrapolate", "field_type", "log_p", @@ -634,6 +742,7 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): field3d = argvars["field3d"] xy = argvars["xy"] + xy = npvalues(xy) result = wrapped(*args, **kwargs) @@ -652,14 +761,30 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): outattrs = OrderedDict() outdimnames = list(field3d.dims) outcoords.update(field3d.coords) + for i in py3range(-2,0,1): - outdimnames.remove(field3d.dims[i]) del outcoords[field3d.dims[i]] + outdimnames.remove(field3d.dims[i]) + + # Need to remove XLAT, XLONG... + delkeys = (key for key,arr in viewitems(field3d.coords) + if arr.ndim > 1) - outdimnames[-2] = "xy" - outattrs.update(field3d.attrs) + for key in delkeys: + del outcoords[key] + + outdimnames.append("xy") + #outattrs.update(field3d.attrs) - outname = "{0}_xy".format(field3d.name) + desc = field3d.attrs.get("description", None) + if desc is not None: + outattrs["description"] = desc + + units = field3d.attrs.get("units", None) + if units is not None: + outattrs["units"] = units + + outname = "{0}_2dxy".format(field3d.name) outcoords["xy_loc"] = ("xy", [CoordPair(xy[i,0], xy[i,1]) for i in py3range(xy.shape[-2])]) @@ -671,7 +796,7 @@ def _set_2dxy_meta(wrapped, instance, args, kwargs): pass else: - outname = "field3d_xy" + outname = "field3d_2dxy" outattrs["Orientation"] = cross_str @@ -695,19 +820,32 @@ def _set_1d_meta(wrapped, instance, args, kwargs): outcoords = OrderedDict() outattrs = OrderedDict() outdimnames = list(v_in.dims) - outcoords.update(v_in.coords) + #outcoords.update(v_in.coords) + outdimnames.pop(-1) + + for name in outdimnames: + try: + outcoords[name] = v_in.coords[name] + except KeyError: + continue - outdimnames.remove(v_in.dims[-1]) - del outcoords[v_in.dims[-1]] outdimnames.append("z") outname = "{0}_z".format(v_in.name) outcoords["z"] = z_out - outattrs.update(v_in.attrs) + #outattrs.update(v_in.attrs) outattrs["_FillValue"] = missingval outattrs["missing_value"] = missingval + desc = v_in.attrs.get("description", None) + if desc is not None: + outattrs["description"] = desc + + units = v_in.attrs.get("units", None) + if units is not None: + outattrs["units"] = units + else: outname = "v_in_z" @@ -715,6 +853,42 @@ def _set_1d_meta(wrapped, instance, args, kwargs): return DataArray(result, name=outname, dims=outdimnames, coords=outcoords, attrs=outattrs) + +def _set_xy_meta(wrapped, instance, args, kwargs): + argvars = from_args(wrapped, ("field", "pivot_point", "angle", + "start_point", "end_point"), + *args, **kwargs) + + field = argvars["field"] + pivot_point = argvars["pivot_point"] + angle = argvars["angle"] + start_point = argvars["start_point"] + end_point = argvars["end_point"] + + result = wrapped(*args, **kwargs) + + if isinstance(field, DataArray): + outname = "{0}_xy".format(field.name) + else: + outname = "xy" + + outdimnames = ["idx", "x_y"] + outcoords = OrderedDict() + outattrs = OrderedDict() + + outcoords["x_y"] = ["x", "y"] + + if pivot_point is not None and angle is not None: + outattrs["pivot_point"] = pivot_point + outattrs["angle"] = angle + + if start_point is not None and end_point is not None: + outattrs["start_point"] = start_point + outattrs["end_point"] = end_point + + return DataArray(result, name=outname, dims=outdimnames, + coords=outcoords, attrs=outattrs) + def set_interp_metadata(interp_type): @wrapt.decorator @@ -739,5 +913,222 @@ def set_interp_metadata(interp_type): return _set_2dxy_meta(wrapped, instance, args, kwargs) elif interp_type == "1d": return _set_1d_meta(wrapped, instance, args, kwargs) + elif interp_type == "xy": + return _set_xy_meta(wrapped, instance, args, kwargs) + + return func_wrapper + + +def set_alg_metadata(alg_ndims, right_dimnames=None, + refvarndims=None, + refvarname=None, missingarg=None, + insert_dimnames=None, + units=None, description=None, squeeze=False): + + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] + + if do_meta is None: + do_meta = True + + if not xarray_enabled() or not do_meta: + return wrapped(*args, **kwargs) + + + result = wrapped(*args, **kwargs) + + # Default dimension names + outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] + + if missingarg is not None: + missingval = from_args(wrapped, (missingarg,), + *args, **kwargs)[missingarg] + else: + missingval = None + + if missingval is not None: + outattrs["_FillValue"] = missingval + outattrs["missing_value"] = missingval + result = np.ma.masked_values(result, missingval) + + outname = wrapped.__name__ + outattrs = OrderedDict() + + if units is not None: + if isinstance(description, from_var): + _units = units(wrapped, *args, **kwargs) + if uts is not None: + outattrs["units"] = _units + else: + outattrs["units"] = units + + if description is not None: + if isinstance(description, from_var): + desc = description(wrapped, *args, **kwargs) + if desc is not None: + outattrs["description"] = desc + else: + outattrs["description"] = description + + + # Copy the dimnames from the reference variable, otherwise, use + # the supplied dimnames + if refvarname is not None: + refvar = from_args(wrapped, (refvarname,), + *args, **kwargs)[refvarname] + else: + refvar = None + + if isinstance(refvar, DataArray): + + # If right dims are provided, use them first + if right_dimnames is not None: + outdims[-alg_ndims:] = right_dimnames[-alg_ndims:] + else: + # Copy the right dims + outdims[-alg_ndims:] = refvar.dims[-alg_ndims:] + + # Left dims + if refvarndims is None: + # Used when result and reference are aligned on right + if result.ndim > alg_ndims: + result_extra = result.ndim - alg_ndims + + for i in py3range(1, result_extra + 1): + idx = -alg_ndims - i + if -idx <= refvar.ndim: + outdims[idx] = refvar.dims[idx] + else: + continue + # When reference and result aren't exactly aligned (slp,uvmet) + else: + ref_extra = refvar.ndim - refvarndims + ref_left_dimnames = refvar.dims[0:ref_extra] + + for i,dimname in enumerate(ref_left_dimnames[::-1], 1): + idx = -i + if -idx <= result.shape: + outdims[idx] = dimname + else: + continute + + if insert_dimnames is not None: + for pair in insert_dimnames: + outdims.insert(pair[0], pair[1]) + + out = DataArray(result, name=outname, dims=outdims, + attrs=outattrs) + + if squeeze: + return out.squeeze() + + return out + + return func_wrapper + +def set_uvmet_alg_metadata(units="mps", description="earth rotated u,v"): + + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] + + if do_meta is None: + do_meta = True + + if not xarray_enabled() or not do_meta: + return wrapped(*args, **kwargs) + + result = wrapped(*args, **kwargs) + + # Default dimension names + outdims = ["dim_{}".format(i) for i in py3range(result.ndim)] + outname = "uvmet" + outattrs = OrderedDict() + + if units is not None: + outattrs["units"] = units + + if description is not None: + outattrs["description"] = description + + latvar = from_args(wrapped, ("lat",), *args, **kwargs)["lat"] + uvar = from_args(wrapped, ("u",), *args, **kwargs)["u"] + + if isinstance(uvar, DataArray): + # Right dims come from latvar + outdims[-2:] = latvar.dims[-2:] + + # Left dims come from u-var + outdims[1:-2] = uvar.dims[0:-2] + + # Left-most is always u_v + outdims[0] = "u_v" + + out = DataArray(result, name=outname, dims=outdims, + attrs=outattrs) + + return out + return func_wrapper + +def set_destag_metadata(): + + @wrapt.decorator + def func_wrapper(wrapped, instance, args, kwargs): + do_meta = from_args(wrapped, ("meta",), *args, **kwargs)["meta"] + + if do_meta is None: + do_meta = True + + if not xarray_enabled() or not do_meta: + return wrapped(*args, **kwargs) + + result = wrapped(*args, **kwargs) + + # Default dimension names + outdims = ["dim{}".format(i) for i in py3range(result.ndim)] + + destag_args = from_args(wrapped, ("var", "stagger_dim"), + *args, **kwargs) + var = destag_args["var"] + destag_dim = destag_args["stagger_dim"] + + if isinstance(var, DataArray): + + if var.name is not None: + outname = "destag_{}".format(var.name) + else: + outnames = "destag_var" + + outattrs = OrderedDict() + outattrs.update(var.attrs) + + outattrs["destag_dim"] = destag_dim + + outdims = [] + outdims += var.dims + + destag_dim_name = outdims[destag_dim] + if destag_dim_name.find("_stag") >= 0: + new_dim_name = destag_dim_name.replace("_stag", "") + else: + if destag_dim >= 0: + new_dim_name = "dim_{}".format(destag_dim) + else: + dim_num = result.ndim + destag_dim + new_dim_name = "dim_{}".format(dim_num) + + outdims[destag_dim] = new_dim_name + + out = DataArray(result, name=outname, dims=outdims, attrs=outattrs) + + return out + + return func_wrapper + + + + + diff --git a/src/wrf/omega.py b/src/wrf/omega.py index ec4de71..6da289f 100755 --- a/src/wrf/omega.py +++ b/src/wrf/omega.py @@ -8,7 +8,6 @@ from .extension import computeomega, _tk from .util import extract_vars from .metadecorators import copy_and_set_metadata -__all__ = ["get_omega"] @copy_and_set_metadata(copy_varname="T", name="omega", description="omega", diff --git a/src/wrf/pressure.py b/src/wrf/pressure.py index fdf6b39..3ba8c4d 100755 --- a/src/wrf/pressure.py +++ b/src/wrf/pressure.py @@ -5,7 +5,6 @@ from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars, either -__all__ = ["get_pressure", "get_pressure_hpa"] @copy_and_set_metadata(copy_varname=either("P", "PRES"), name="pressure", description="pressure") diff --git a/src/wrf/projection.py b/src/wrf/projection.py index b598015..853fa9e 100644 --- a/src/wrf/projection.py +++ b/src/wrf/projection.py @@ -15,10 +15,6 @@ if basemap_enabled(): if pyngl_enabled(): from Ngl import Resources -__all__ = ["WrfProj", "NullProjection", "LambertConformal", "Mercator", - "PolarStereographic", "LatLon", "RotatedLatLon", - "getproj"] - if cartopy_enabled(): class MercatorWithLatTS(crs.Mercator): @@ -98,31 +94,24 @@ class WrfProj(object): if self.stand_lon is None: self.stand_lon = self._cen_lon - @property - def _basemap(self): + def _basemap(self, resolution='l'): return None - @property def _cf_params(self): return None - @property def _cartopy(self): return None - @property def _cart_extents(self): return ([self.ll_lon, self.ur_lon], [self.ll_lat, self.ur_lat]) - @property def _pyngl(self): return None - @property def _proj4(self): return None - @property def _globe(self): return (None if not cartopy_enabled() else crs.Globe(ellipse=None, @@ -131,11 +120,11 @@ class WrfProj(object): def cartopy_xlim(self): """Return the x extents in projected coordinates (for cartopy)""" - return self._cart_extents[0] + return self._cart_extents()[0] def cartopy_ylim(self): """Return the y extents in projected coordinates (for cartopy)""" - return self._cart_extents[1] + return self._cart_extents()[1] def __repr__(self): args = ("bottom_left={}, top_right={}, " @@ -148,13 +137,13 @@ class WrfProj(object): self.pole_lon)) return "{}({})".format(self.__class__.__name__, args) - def basemap(self): + def basemap(self, resolution='l'): """Return a mpl_toolkits.basemap.Basemap instance for the projection""" if not basemap_enabled(): raise RuntimeError("'mpl_toolkits.basemap' is not " "installed or is disabled") - return self._basemap + return self._basemap(resolution) def cartopy(self): """Return a cartopy.crs.Projection subclass for the @@ -162,23 +151,23 @@ class WrfProj(object): if not cartopy_enabled(): raise RuntimeError("'cartopy' is not " "installed or is disabled") - return self._cartopy + return self._cartopy() def pyngl(self): """Return the PyNGL resources for the projection""" if not pyngl_enabled(): raise RuntimeError("'pyngl' is not " "installed or is disabled") - return self._pyngl + return self._pyngl() def proj4(self): """Return the proj4 string for the map projection""" - return self._proj4 + return self._proj4() def cf(self): """Return a dictionary with the NetCDF CF parameters for the projection""" - return self._cf_params + return self._cf_params() # Used for 'missing' projection values during the 'join' method @@ -200,7 +189,7 @@ class LambertConformal(WrfProj): if self.truelat2 is not None: self._std_parallels.append(self.truelat2) - @property + def _cf_params(self): _cf_params = {} _cf_params["grid_mapping_name"] = "lambert_conformal_conic"; @@ -211,7 +200,7 @@ class LambertConformal(WrfProj): return _cf_params - @property + def _pyngl(self): if not pyngl_enabled(): return None @@ -234,8 +223,8 @@ class LambertConformal(WrfProj): return _pyngl - @property - def _basemap(self): + + def _basemap(self, resolution='l'): if not basemap_enabled(): return None @@ -249,11 +238,10 @@ class LambertConformal(WrfProj): llcrnrlon = self.ll_lon, urcrnrlon = self.ur_lon, rsphere = Constants.WRF_EARTH_RADIUS, - resolution = 'l') + resolution = resolution) return _basemap - @property def _cartopy(self): if not cartopy_enabled(): return None @@ -262,11 +250,10 @@ class LambertConformal(WrfProj): central_longitude = self.stand_lon, central_latitude = self.moad_cen_lat, standard_parallels = self._std_parallels, - globe = self._globe) + globe = self._globe()) return _cartopy - @property def _cart_extents(self): # Need to modify the extents for the new projection pc = crs.PlateCarree() @@ -280,7 +267,6 @@ class LambertConformal(WrfProj): return (_xlimits, _ylimits) - @property def _proj4(self): truelat2 = (self.truelat1 if _ismissing(self.truelat2) @@ -306,7 +292,7 @@ class Mercator(WrfProj): if self.truelat1 == 0. or _ismissing(self.truelat1) else self.truelat1) - @property + def _cf_params(self): _cf_params = {} @@ -316,7 +302,7 @@ class Mercator(WrfProj): return _cf_params - @property + def _pyngl(self): if not pyngl_enabled(): return None @@ -334,8 +320,8 @@ class Mercator(WrfProj): return _pyngl - @property - def _basemap(self): + + def _basemap(self, resolution='l'): if not basemap_enabled(): return None @@ -348,11 +334,11 @@ class Mercator(WrfProj): llcrnrlon = self.ll_lon, urcrnrlon = self.ur_lon, rsphere = Constants.WRF_EARTH_RADIUS, - resolution = 'l') + resolution = resolution) return _basemap - @property + def _cartopy(self): if not cartopy_enabled(): return None @@ -360,17 +346,17 @@ class Mercator(WrfProj): if self._lat_ts == 0.0: _cartopy = crs.Mercator( central_longitude = self.stand_lon, - globe = self._globe) + globe = self._globe()) else: _cartopy = MercatorWithLatTS( central_longitude = self.stand_lon, latitude_true_scale = self._lat_ts, - globe = self._globe) + globe = self._globe()) return _cartopy - @property + def _cart_extents(self): # Need to modify the extents for the new projection @@ -384,7 +370,7 @@ class Mercator(WrfProj): return (_xlimits, _ylimits) - @property + def _proj4(self): _proj4 = ("+proj=merc +units=meters +a={} +b={} " @@ -406,7 +392,7 @@ class PolarStereographic(WrfProj): if _ismissing(self.truelat1) else self.truelat1) - @property + def _cf_params(self): _cf_params = {} _cf_params["grid_mapping_name"] = "polar_stereographic" @@ -417,7 +403,7 @@ class PolarStereographic(WrfProj): return _cf_params - @property + def _pyngl(self): if not pyngl_enabled(): return None @@ -439,8 +425,8 @@ class PolarStereographic(WrfProj): return _pyngl - @property - def _basemap(self): + + def _basemap(self, resolution='l'): if not basemap_enabled(): return None @@ -453,11 +439,11 @@ class PolarStereographic(WrfProj): llcrnrlon = self.ll_lon, urcrnrlon = self.ur_lon, rsphere = Constants.WRF_EARTH_RADIUS, - resolution = 'l') + resolution = resolution) return _basemap - @property + def _cartopy(self): if not cartopy_enabled(): return None @@ -465,10 +451,10 @@ class PolarStereographic(WrfProj): _cartopy = crs.Stereographic(central_latitude=self._hemi, central_longitude=self.stand_lon, true_scale_latitude=self._lat_ts, - globe=self._globe) + globe=self._globe()) return _cartopy - @property + def _cart_extents(self): # Need to modify the extents for the new projection pc = crs.PlateCarree() @@ -481,7 +467,7 @@ class PolarStereographic(WrfProj): return (_xlimits, _ylimits) - @property + def _proj4(self): _proj4 = ("+proj=stere +units=meters +a={} +b={} " "+lat0={} +lon_0={} +lat_ts={}".format( @@ -501,13 +487,13 @@ class LatLon(WrfProj): super(LatLon, self).__init__(bottom_left, top_right, lats, lons, **proj_params) - @property + def _cf_params(self): _cf_params = {} _cf_params["grid_mapping_name"] = "latitude_longitude" return _cf_params - @property + def _pyngl(self): if not pyngl_enabled(): return None @@ -525,8 +511,8 @@ class LatLon(WrfProj): return _pyngl - @property - def _basemap(self): + + def _basemap(self, resolution='l'): if not basemap_enabled(): return None @@ -538,25 +524,25 @@ class LatLon(WrfProj): llcrnrlon = self.ll_lon, urcrnrlon = self.ur_lon, rsphere = Constants.WRF_EARTH_RADIUS, - resolution = 'l') + resolution = resolution) return _basemap - @property + def _cartopy(self): if not cartopy_enabled(): return None _cartopy = crs.PlateCarree(central_longitude=self.stand_lon, - globe=self._globe) + globe=self._globe()) return _cartopy - @property + def _cart_extents(self): return ([self.ll_lon, self.ur_lon], [self.ll_lat, self.ur_lat]) - @property + def _proj4(self): _proj4 = ("+proj=eqc +units=meters +a={} +b={} " "+lon_0={}".format(Constants.WRF_EARTH_RADIUS, @@ -568,8 +554,7 @@ class LatLon(WrfProj): # Each projection system handles this differently. # 1) In WRF, if following the WPS instructions, POLE_LON is mainly used to # determine north or south hemisphere. In other words, it determines if -# the globe is tipped toward or away from you. If a non-0 or non-180 -# value is chosen, PyNGL cannot plot it. +# the globe is tipped toward or away from you. # 2) In WRF, POLE_LAT is always positive, but should be negative in the # proj4 based systems when using the southern hemisphere projections. # 3) In cartopy, pole_longitude is used to describe the dateline, which @@ -645,7 +630,7 @@ class RotatedLatLon(WrfProj): self._cart_pole_lon = (-self.stand_lon - 180.0 if self._north else -self.stand_lon) - @property + def _cf_params(self): _cf_params = {} # Assuming this follows the same guidelines as cartopy @@ -656,7 +641,7 @@ class RotatedLatLon(WrfProj): return _cf_params - @property + def _pyngl(self): if not pyngl_enabled(): return None @@ -674,8 +659,8 @@ class RotatedLatLon(WrfProj): return _pyngl - @property - def _basemap(self): + + def _basemap(self, resolution='l'): if not basemap_enabled(): return None @@ -688,10 +673,10 @@ class RotatedLatLon(WrfProj): urcrnrlon = self.ur_lon, lon_0 = self._bm_lon_0, rsphere = Constants.WRF_EARTH_RADIUS, - resolution = 'l') + resolution = resolution) return _basemap - @property + def _cartopy(self): if not cartopy_enabled(): return None @@ -701,10 +686,10 @@ class RotatedLatLon(WrfProj): pole_latitude=self._bm_cart_pole_lat, central_rotated_longitude=( 180.0 - self.pole_lon), # Probably - globe = self._globe) + globe = self._globe()) return _cartopy - @property + def _cart_extents(self): # Need to modify the extents for the new projection pc = crs.PlateCarree() @@ -717,21 +702,21 @@ class RotatedLatLon(WrfProj): return (_xlimits, _ylimits) - @property + def _proj4(self): _proj4 = ("+proj=ob_tran +o_proj=latlon " - "+a={} +b={} +to_meter={} +o_lon_p={} +o_lat_p={} " - "+lon_0={}".format(Constants.WRF_EARTH_RADIUS, - Constants.WRF_EARTH_RADIUS, - math.radians(1), - 180.0 - self.pole_lon, - self._bm_cart_pole_lat, - 180.0 + self._cart_pole_lon)) + "+a={} +b={} +to_meter={} +o_lon_p={} +o_lat_p={} " + "+lon_0={}".format(Constants.WRF_EARTH_RADIUS, + Constants.WRF_EARTH_RADIUS, + math.radians(1), + 180.0 - self.pole_lon, + self._bm_cart_pole_lat, + 180.0 + self._cart_pole_lon)) return _proj4 def getproj(bottom_left=None, top_right=None, - lats=None, lons=None, **proj_params): + lats=None, lons=None, **proj_params): proj_type = proj_params.get("MAP_PROJ", 0) if proj_type == 1: diff --git a/src/wrf/psadlookup.py b/src/wrf/psadlookup.py index 2565f9a..55d14c3 100755 --- a/src/wrf/psadlookup.py +++ b/src/wrf/psadlookup.py @@ -3,7 +3,6 @@ from __future__ import (absolute_import, division, print_function, import numpy as n -__all__ = ["get_lookup_tables"] _THETA_DIM = 150 _PRS_DIM = 150 diff --git a/src/wrf/pw.py b/src/wrf/pw.py index e7f2fe1..d7c5c6b 100755 --- a/src/wrf/pw.py +++ b/src/wrf/pw.py @@ -7,7 +7,6 @@ from .constants import Constants from .util import extract_vars from .metadecorators import copy_and_set_metadata -__all__ = ["get_pw"] @copy_and_set_metadata(copy_varname="T", name="pw", remove_dims=("bottom_top",), diff --git a/src/wrf/rh.py b/src/wrf/rh.py index 9b39a14..251ded4 100755 --- a/src/wrf/rh.py +++ b/src/wrf/rh.py @@ -7,7 +7,6 @@ from .extension import _rh, _tk from .util import extract_vars from .metadecorators import copy_and_set_metadata -__all__ = ["get_rh", "get_rh_2m"] @copy_and_set_metadata(copy_varname="T", name="rh", description="relative humidity", diff --git a/src/wrf/routines.py b/src/wrf/routines.py index 096cbd4..d3992ab 100644 --- a/src/wrf/routines.py +++ b/src/wrf/routines.py @@ -22,7 +22,6 @@ from .wind import (get_destag_wspd_wdir, get_destag_wspd_wdir10, get_u_destag, get_v_destag, get_w_destag) from .times import get_times -__all__ = ["getvar"] # func is the function to call. kargs are required arguments that should # not be altered by the user @@ -103,9 +102,10 @@ _VALID_KARGS = {"cape2d" : ["missing"], "lon" : [], "pres" : ["units"], "pressure" : ["units"], - "wspddir" : ["units"], - "wspddir_uvmet" : ["units"], - "wspddir_uvmet10" : ["units"], + "wspd_wdir" : ["units"], + "wspd_wdir10" : ["units"], + "wspd_wdir_uvmet" : ["units"], + "wspd_wdir_uvmet10" : ["units"], "ctt" : [], "default" : [] } @@ -149,20 +149,145 @@ def _check_kargs(var, kargs): "argument for '%s'" % (arg, var)) -def getvar(wrfnc, var, timeidx=0, +def getvar(wrfnc, varname, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, **kargs): + """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 | | + +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ + | 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. | + +--------------------+---------------------------------------------------------------+---------------------+-----------------------------------------------------------------------------------------------+ + + + 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. + + varname : 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. + + 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'. + + 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 + ------- + out : 'xarray.DataArray` or `numpy.ndarray` + Returns the specififed 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. + + """ + wrfnc = _get_iterable(wrfnc) - if is_standard_wrf_var(wrfnc, var): - return extract_vars(wrfnc, timeidx, var, - method, squeeze, cache, meta)[var] + if is_standard_wrf_var(wrfnc, varname): + return extract_vars(wrfnc, timeidx, varname, + method, squeeze, cache, meta)[varname] - actual_var = _undo_alias(var) + actual_var = _undo_alias(varname) if actual_var not in _VALID_KARGS: - raise ArgumentError("'%s' is not a valid variable name" % (var)) + raise ArgumentError("'%s' is not a valid variable name" % (varname)) _check_kargs(actual_var, kargs) return _FUNC_MAP[actual_var](wrfnc, timeidx, method, squeeze, cache, meta, **kargs) + diff --git a/src/wrf/slp.py b/src/wrf/slp.py index 3b9865f..d7fb3d2 100755 --- a/src/wrf/slp.py +++ b/src/wrf/slp.py @@ -9,7 +9,6 @@ from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars -__all__ = ["get_slp"] @copy_and_set_metadata(copy_varname="T", name="slp", remove_dims=("bottom_top",), diff --git a/src/wrf/temp.py b/src/wrf/temp.py index f02eda3..28da822 100755 --- a/src/wrf/temp.py +++ b/src/wrf/temp.py @@ -8,8 +8,6 @@ from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars -__all__ = ["get_theta", "get_temp", "get_eth", "get_tv", "get_tw", - "get_tk", "get_tc"] @copy_and_set_metadata(copy_varname="T", name="theta", description="potential temperature") diff --git a/src/wrf/terrain.py b/src/wrf/terrain.py index 972b6cb..ad8685e 100755 --- a/src/wrf/terrain.py +++ b/src/wrf/terrain.py @@ -5,7 +5,6 @@ from .decorators import convert_units from .metadecorators import copy_and_set_metadata from .util import extract_vars, either -__all__ = ["get_terrain"] # Need to handle either @copy_and_set_metadata(copy_varname=either("HGT", "HGT_M"), name="terrain", diff --git a/src/wrf/times.py b/src/wrf/times.py index 3045213..dcfc5a7 100755 --- a/src/wrf/times.py +++ b/src/wrf/times.py @@ -3,7 +3,6 @@ from __future__ import (absolute_import, division, print_function, from .util import extract_times -__all__ = ["get_times"] def get_times(wrfnc, timeidx=0): return extract_times(wrfnc, timeidx) diff --git a/src/wrf/units.py b/src/wrf/units.py index be2795f..1738b10 100755 --- a/src/wrf/units.py +++ b/src/wrf/units.py @@ -3,7 +3,6 @@ from __future__ import (absolute_import, division, print_function, from .constants import Constants, ConversionFactors -__all__ = ["check_units", "do_conversion"] # Handles unit conversions that only differ by multiplication factors def _apply_conv_fact(var, vartype, var_unit, dest_unit): diff --git a/src/wrf/util.py b/src/wrf/util.py index 2d17fa3..383213b 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -8,8 +8,8 @@ from itertools import product from types import GeneratorType import datetime as dt from math import floor, copysign - from inspect import getmodule + try: from inspect import signature except ImportError: @@ -56,10 +56,11 @@ _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), _COORD_VARS = ("XLAT", "XLONG", "XLAT_M", "XLONG_M", "XLAT_U", "XLONG_U", "XLAT_V", "XLONG_V", "CLAT", "CLONG") -_TIME_COORD_VARS = ("XTIME",) +_LAT_COORDS = ("XLAT", "XLAT_M", "XLAT_U", "XLAT_V", "CLAT") -def _is_coord_var(varname): - return varname in _COORD_VARS +_LON_COORDS = ("XLONG", "XLONG_M", "XLONG_U","XLONG_V", "CLONG") + +_TIME_COORD_VARS = ("XTIME",) def _is_time_coord_var(varname): @@ -125,6 +126,25 @@ class TestGen(object): def __next__(self): return self.next() +def latlon_coordvars(d): + lat_coord = None + lon_coord = None + + for name in _LAT_COORDS: + if name in viewkeys(d): + lat_coord = name + break + + for name in _LON_COORDS: + if name in viewkeys(d): + lon_coord = name + break + + return lat_coord, lon_coord + +def is_coordvar(varname): + 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""" @@ -300,6 +320,24 @@ class combine_dims(object): return tuple(result) + +class from_var(object): + def __init__(self, varname, attribute): + self.varname = varname + self.attribute = attribute + + def __call__(self, wrapped, *args, **kwargs): + vard = from_args(wrapped, (self.varname,), *args, **kwargs) + + var = None + if vard is not None: + var = vard[self.varname] + + if not isinstance(var, DataArray): + return None + + return var.attrs.get(self.attribute, None) + def _corners_moved(wrfnc, first_ll_corner, first_ur_corner, latvar, lonvar): lats = wrfnc.variables[latvar] @@ -571,7 +609,7 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain): # WRF files coord_attr = getattr(var, "coordinates") except AttributeError: - if _is_coord_var(varname): + if is_coordvar(varname): # Coordinate variable (most likely XLAT or XLONG) lat_coord, lon_coord = _get_coord_pairs(varname) time_coord = None @@ -1056,7 +1094,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta): outdata = outdata[:, np.newaxis, :] outarr = outdata - + return outarr def combine_files(wrfseq, varname, timeidx, is_moving=None, @@ -1091,8 +1129,7 @@ def _extract_var(wrfnc, varname, timeidx, is_moving, pass else: if not meta: - if isinstance(cache_var, DataArray): - return cache_var.values + return npvalues(cache_var) return cache_var @@ -1238,7 +1275,7 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): time_idx_or_slice = slice(None) if varname is not None: - if not _is_coord_var(varname): + if not is_coordvar(varname): coord_names = getattr(wrfnc.variables[varname], "coordinates").split() lon_coord = coord_names[0] @@ -1246,8 +1283,7 @@ def get_proj_params(wrfnc, timeidx=0, varname=None): else: lat_coord, lon_coord = _get_coord_pairs(varname) else: - lat_coord = "XLAT" - lon_coord = "XLONG" + lat_coord, lon_coord = latlon_coordvars(wrfnc.variables) return (wrfnc.variables[lat_coord][time_idx_or_slice,:], wrfnc.variables[lon_coord][time_idx_or_slice,:], @@ -1284,6 +1320,24 @@ class CoordPair(object): def __str__(self): return self.__repr__() + def xy_str(self, fmt="{:.4f}, {:.4f}"): + if self.x is None or self.y is None: + return None + + return fmt.format(self.x, self.y) + + def latlon_str(self, fmt="{:.4f}, {:.4f}"): + if self.lat is None or self.lon is None: + return None + + return fmt.format(self.lat, self.lon) + + def ij_str(self, fmt="{:.4f}, {:.4f}"): + if self.i is None or self.j is None: + return None + + return fmt.format(self.i, self.j) + def from_args(func, argnames, *args, **kwargs): """Parses the function args and kargs looking for the desired argument diff --git a/src/wrf/uvdecorator.py b/src/wrf/uvdecorator.py index eff190a..619a21e 100644 --- a/src/wrf/uvdecorator.py +++ b/src/wrf/uvdecorator.py @@ -5,7 +5,7 @@ import numpy as np import wrapt -from .destag import destagger +#from .destag import destagger from .util import iter_left_indexes, py3range, npvalues from .config import xarray_enabled from .constants import Constants @@ -13,10 +13,84 @@ from .constants import Constants if xarray_enabled(): from xarray import DataArray -__all__ = ["uvmet_left_iter"] +# def uvmet_left_iter(): +# """Decorator to handle iterating over leftmost dimensions when using +# multiple files and/or multiple times with the uvmet product. +# +# """ +# @wrapt.decorator +# def func_wrapper(wrapped, instance, args, kwargs): +# u = args[0] +# v = args[1] +# lat = args[2] +# lon = args[3] +# cen_long = args[4] +# cone = args[5] +# +# if u.ndim == lat.ndim: +# num_right_dims = 2 +# is_3d = False +# else: +# num_right_dims = 3 +# is_3d = True +# +# is_stag = False +# if ((u.shape[-1] != lat.shape[-1]) or +# (u.shape[-2] != lat.shape[-2])): +# is_stag = True +# +# if is_3d: +# extra_dim_num = u.ndim - 3 +# else: +# extra_dim_num = u.ndim - 2 +# +# if is_stag: +# u = destagger(u,-1) +# v = destagger(v,-2) +# +# # No special left side iteration, return the function result +# if (extra_dim_num == 0): +# return wrapped(u, v, lat, lon, cen_long, cone) +# +# # Start by getting the left-most 'extra' dims +# outdims = u.shape[0:extra_dim_num] +# extra_dims = list(outdims) # Copy the left-most dims for iteration +# +# # Append the right-most dimensions +# outdims += [2] # For u/v components +# +# #outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] +# outdims += list(u.shape[-num_right_dims:]) +# +# output = np.empty(outdims, u.dtype) +# +# for left_idxs in iter_left_indexes(extra_dims): +# # Make the left indexes plus a single slice object +# # The single slice will handle all the dimensions to +# # the right (e.g. [1,1,:]) +# left_and_slice_idxs = tuple([x for x in left_idxs] + [slice(None)]) +# +# new_u = u[left_and_slice_idxs] +# new_v = v[left_and_slice_idxs] +# new_lat = lat[left_and_slice_idxs] +# new_lon = lon[left_and_slice_idxs] +# +# # Call the numerical routine +# result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) +# +# # Note: The 2D version will return a 3D array with a 1 length +# # dimension. Numpy is unable to broadcast this without +# # sqeezing first. +# result = np.squeeze(result) +# +# output[left_and_slice_idxs] = result[:] +# +# return output +# +# return func_wrapper -def uvmet_left_iter(): +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. @@ -30,90 +104,28 @@ def uvmet_left_iter(): cen_long = args[4] cone = args[5] - if u.ndim == lat.ndim: - num_right_dims = 2 - is_3d = False - else: - num_right_dims = 3 - is_3d = True - - is_stag = False - if ((u.shape[-1] != lat.shape[-1]) or - (u.shape[-2] != lat.shape[-2])): - is_stag = True + orig_dtype = u.dtype - if is_3d: - extra_dim_num = u.ndim - 3 - else: - extra_dim_num = u.ndim - 2 + lat_lon_fixed = False + if lat.ndim == 2: + lat_lon_fixed = True - if is_stag: - u = destagger(u,-1) - v = destagger(v,-2) - - # No special left side iteration, return the function result - if (extra_dim_num == 0): - return wrapped(u, v, lat, lon, cen_long, cone) + if lon.ndim == 2 and not lat_lon_fixed: + raise ValueError("'lat' and 'lon' shape mismatch") - # Start by getting the left-most 'extra' dims - outdims = u.shape[0:extra_dim_num] - extra_dims = list(outdims) # Copy the left-most dims for iteration + num_left_dims_u = u.ndim - 2 + num_left_dims_lat = lat.ndim - 2 - # Append the right-most dimensions - outdims += [2] # For u/v components + if (num_left_dims_lat > num_left_dims_u): + raise ValueError("number of 'lat' dimensions is greater than 'u'") - #outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] - outdims += list(u.shape[-num_right_dims:]) - - output = np.empty(outdims, u.dtype) - - for left_idxs in iter_left_indexes(extra_dims): - # Make the left indexes plus a single slice object - # The single slice will handle all the dimensions to - # the right (e.g. [1,1,:]) - left_and_slice_idxs = tuple([x for x in left_idxs] + [slice(None)]) - - new_u = u[left_and_slice_idxs] - new_v = v[left_and_slice_idxs] - new_lat = lat[left_and_slice_idxs] - new_lon = lon[left_and_slice_idxs] - - # Call the numerical routine - result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone) - - # Note: The 2D version will return a 3D array with a 1 length - # dimension. Numpy is unable to broadcast this without - # sqeezing first. - result = np.squeeze(result) - - output[left_and_slice_idxs] = result[:] - - return output - - return func_wrapper - -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. - - """ - @wrapt.decorator - def func_wrapper(wrapped, instance, args, kwargs): - u = args[0] - v = args[1] - lat = args[2] - lon = args[3] - cen_long = args[4] - cone = args[5] - - orig_dtype = u.dtype - - if u.ndim == lat.ndim: - num_right_dims = 2 - is_3d = False + if lat_lon_fixed: + mode = 0 # fixed lat/lon else: - num_right_dims = 3 - is_3d = True + if num_left_dims_u == num_left_dims_lat: + mode = 1 # lat/lon same as u + else: + mode = 2 # probably 3D with 2D lat/lon plus Time has_missing = False u_arr = u @@ -135,56 +147,56 @@ def uvmet_left_iter_nocopy(alg_dtype=np.float64): vmissing = v_arr.fill_value uvmetmissing = umissing - - is_stag = False + + is_stag = 0 if (u.shape[-1] != lat.shape[-1] or u.shape[-2] != lat.shape[-2]): - is_stag = True + is_stag = 1 # Sanity check if (v.shape[-1] == lat.shape[-1] or v.shape[-2] == lat.shape[-2]): raise ValueError("u is staggered but v is not") if (v.shape[-1] != lat.shape[-1] or v.shape[-2] != lat.shape[-2]): - is_stag = True + is_stag = 1 # Sanity check if (u.shape[-1] == lat.shape[-1] or u.shape[-2] == lat.shape[-2]): raise ValueError("v is staggered but u is not") - if is_3d: - extra_dim_num = u.ndim - 3 - else: - extra_dim_num = u.ndim - 2 - + # No special left side iteration, return the function result - if (extra_dim_num == 0): + if (num_left_dims_u == 0): return wrapped(u, v, lat, lon, cen_long, cone, isstag=is_stag, has_missing=has_missing, umissing=umissing, vmissing=vmissing, uvmetmissing=uvmetmissing) + + # Initial output is time,nz,2,ny,nx to create contiguous views + outdims = u.shape[0:num_left_dims_u] + extra_dims = tuple(outdims) # Copy the left-most dims for iteration - # Start by getting the left-most 'extra' dims - outdims = u.shape[0:extra_dim_num] - extra_dims = list(outdims) # Copy the left-most dims for iteration - - # Append the right-most dimensions - if not is_3d: - outdims += (2,1) # Fortran routine needs 3 dimensions - else: - outdims += (2,) + outdims += (2,) - #outdims += [u.shape[x] for x in py3range(-num_right_dims,0,1)] - outdims += u.shape[-num_right_dims:] + outdims += lat.shape[-2:] - output = np.empty(outdims, alg_dtype) + outview_array = np.empty(outdims, alg_dtype) for left_idxs in iter_left_indexes(extra_dims): - left_and_slice_idxs = left_idxs + (slice(None),) - + + if mode == 0: + lat_left_and_slice = (slice(None),) + elif mode == 1: + lat_left_and_slice = left_and_slice_idxs + elif mode == 2: + # Only need the left-most + lat_left_and_slice = tuple(left_idx + for left_idx in left_idxs[0:num_left_dims_lat]) + + new_u = u[left_and_slice_idxs] new_v = v[left_and_slice_idxs] - new_lat = lat[left_and_slice_idxs] - new_lon = lon[left_and_slice_idxs] - outview = output[left_and_slice_idxs] + new_lat = lat[lat_left_and_slice] + new_lon = lon[lat_left_and_slice] + outview = outview_array[left_and_slice_idxs] # Call the numerical routine result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone, @@ -198,13 +210,20 @@ def uvmet_left_iter_nocopy(alg_dtype=np.float64): outview.__array_interface__["data"][0]): raise RuntimeError("output array was copied") + # Need to reshape this so that u_v is left dim, then time (or others), + # then nz, ny, nz + output_dims = (2,) + output_dims += extra_dims + output_dims += lat.shape[-2:] + output = np.empty(output_dims, orig_dtype) - output = output.astype(orig_dtype) + output[0,:] = outview_array[...,0,:,:].astype(orig_dtype) + output[1,:] = outview_array[...,1,:,:].astype(orig_dtype) if has_missing: output = np.ma.masked_values(output, uvmetmissing) - return output.squeeze() + return output return func_wrapper diff --git a/src/wrf/uvmet.py b/src/wrf/uvmet.py index 3502d96..301853c 100755 --- a/src/wrf/uvmet.py +++ b/src/wrf/uvmet.py @@ -14,9 +14,6 @@ from .decorators import convert_units from .metadecorators import set_wind_metadata from .util import extract_vars, extract_global_attrs, either -__all__=["get_uvmet", "get_uvmet10", "get_uvmet_wspd_wdir", - "get_uvmet10_wspd_wdir"] - @convert_units("wind", "mps") def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, @@ -63,15 +60,15 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, # u,v aggregated in to one array end_idx = -3 if not ten_m else -2 - resdim = list(u.shape[0:end_idx]) + [2] + list(u.shape[end_idx:]) + resdim = (2,) + u.shape[0:end_idx] + u.shape[end_idx:] # Make a new output array for the result res = np.empty(resdim, u.dtype) - # For 2D array, this makes (...,0,:,:) and (...,1,:,:) - # For 3D array, this makes (...,0,:,:,:) and (...,1,:,:,:) - idx0 = tuple([Ellipsis] + [0] + [slice(None)]*(-end_idx)) - idx1 = tuple([Ellipsis] + [1] + [slice(None)]*(-end_idx)) + # For 2D array, this makes (0,...,:,:) and (1,...,:,:) + # For 3D array, this makes (0,...,:,:,:) and (1,...,:,:,:) + idx0 = (0,) + (Ellipsis,) + (slice(None),)*(-end_idx) + idx1 = (1,) + (Ellipsis,) + (slice(None),)*(-end_idx) res[idx0] = u[:] res[idx1] = v[:] @@ -123,7 +120,7 @@ def _get_uvmet(wrfnc, timeidx=0, method="cat", squeeze=True, result = _uvmet(u, v, lat, lon, cen_lon, cone) - return result + return result.squeeze() @set_wind_metadata(copy_varname=either("P", "PRES"), @@ -161,10 +158,11 @@ def get_uvmet_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, units="mps"): - uvmet = _get_uvmet(wrfnc, timeidx, False, units, method, squeeze, cache, - meta) + uvmet = _get_uvmet(wrfnc, timeidx, method, squeeze, + cache, meta, False, units) - return _calc_wspd_wdir(uvmet[...,0,:,:,:], uvmet[...,1,:,:,:], False, units) + return _calc_wspd_wdir(uvmet[0,...,:,:,:], uvmet[1,...,:,:,:], + False, units) @set_wind_metadata(copy_varname=either("PSFC", "F"), @@ -179,7 +177,7 @@ def get_uvmet10_wspd_wdir(wrfnc, timeidx=0, method="cat", squeeze=True, uvmet10 = _get_uvmet(wrfnc, timeidx, method, squeeze, cache, meta, True, units) - return _calc_wspd_wdir(uvmet10[...,0,:,:], uvmet10[...,1,:,:], True, units) + return _calc_wspd_wdir(uvmet10[0,...,:,:], uvmet10[1,...,:,:], True, units) diff --git a/src/wrf/vorticity.py b/src/wrf/vorticity.py index e1324e8..74fd62d 100755 --- a/src/wrf/vorticity.py +++ b/src/wrf/vorticity.py @@ -5,7 +5,6 @@ from .extension import computeavo, computepvo from .util import extract_vars, extract_global_attrs from .metadecorators import copy_and_set_metadata -__all__ = ["get_avo", "get_pvo"] @copy_and_set_metadata(copy_varname="T", name="avo", description="absolute vorticity", diff --git a/src/wrf/wind.py b/src/wrf/wind.py index c5d9b63..1cc36b9 100755 --- a/src/wrf/wind.py +++ b/src/wrf/wind.py @@ -9,9 +9,6 @@ from .util import extract_vars, either from .decorators import convert_units from .metadecorators import set_wind_metadata -__all__ = ["get_u_destag", "get_v_destag", "get_w_destag", - "get_destag_wspd_wdir", "get_destag_wspd_wdir10"] - @convert_units("wind", "mps") def _calc_wspd(u, v, units="mps"): @@ -33,13 +30,13 @@ def _calc_wspd_wdir(u, v, two_d, units): res = np.zeros(outdims, wspd.dtype) - idxs0 = ((Ellipsis, 0, slice(None), slice(None), slice(None)) + idxs0 = ((0,Ellipsis, slice(None), slice(None), slice(None)) if not two_d else - (Ellipsis, 0, slice(None), slice(None))) + (1,Ellipsis, slice(None), slice(None))) - idxs1 = ((Ellipsis, 1, slice(None), slice(None), slice(None)) + idxs1 = ((1, Ellipsis, slice(None), slice(None), slice(None)) if not two_d else - (Ellipsis, 1, slice(None), slice(None))) + (0, Ellipsis, slice(None), slice(None))) res[idxs0] = wspd[:] res[idxs1] = wdir[:] diff --git a/test/ipynb/WRF_Workshop_Demo.ipynb b/test/ipynb/WRF_Workshop_Demo.ipynb new file mode 100644 index 0000000..c3961d4 --- /dev/null +++ b/test/ipynb/WRF_Workshop_Demo.ipynb @@ -0,0 +1,340 @@ +{ + "cells": [ + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "# Jupyter Notebook specific command to allow matplotlib images to be shown inline\n", + "%matplotlib inline" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.cm import get_cmap\n", + "from Nio import open_file\n", + "from wrf import getvar, npvalues\n", + "\n", + "#filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-02-25_18_00_00\"\n", + "filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\"\n", + "pynio_filename = filename + \".nc\"\n", + "ncfile = open_file(pynio_filename)\n", + "\n", + "# Extract the terrain height, which will be returned as an xarray.DataArray object (if available). DataArray\n", + "# objects include meta data, similar to NCL's variables.\n", + "# Note: can also use the 'ter' variable\n", + "terrainx = getvar(ncfile, \"HGT\", timeidx=0)\n", + "\n", + "print (terrainx)\n", + "\n", + "# To extract the numpy array, use the npvalues function\n", + "terrain_numpy = npvalues(terrainx)\n", + "print (\"\\nExtracted numpy array:\\n\")\n", + "print (terrain_numpy)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.cm import get_cmap\n", + "from Nio import open_file\n", + "from wrf import getvar, npvalues\n", + "\n", + "\n", + "#filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-02-25_18_00_00\"\n", + "filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\"\n", + "pynio_filename = filename + \".nc\"\n", + "ncfile = open_file(pynio_filename)\n", + "\n", + "# Extract the terrain height, which will be returned as an xarray.DataArray object (if available). DataArray\n", + "# objects include meta data, similar to NCL's variables.\n", + "# Note: can also use the 'ter' variable\n", + "terrainx = getvar(ncfile, \"HGT\", timeidx=0)\n", + "\n", + "# Use npvalues to extract the numpy array, since matplotlib does not handle xarray.DataArray natively\n", + "terrain_data = npvalues(terrainx)\n", + "\n", + "# Get the lat/lon 2D coordinate arrays. Use npvalues to extract the numpy array since basemap does not\n", + "# handle xarray.DataArray natively.\n", + "lons = npvalues(terrainx.coords[\"XLONG\"])\n", + "lats = npvalues(terrainx.coords[\"XLAT\"])\n", + "\n", + "# Extract the basemap object from the projection information\n", + "wrf_proj = terrainx.attrs[\"projection\"]\n", + "bm = wrf_proj.basemap()\n", + "\n", + "# Convert the lat/lon coordinates to projected x,y\n", + "x,y = bm(lons, lats)\n", + "\n", + "# Create the figure\n", + "fig = plt.figure(figsize=(16,16))\n", + "ax = fig.add_axes([0.1,0.1,0.8,0.8])\n", + "\n", + "# Draw filled contours from 200 to 3000 m, every 200 meters.\n", + "levels = np.arange(250, 3000, 200)\n", + "bm.contourf(x, y, terrain_data, levels=levels, extend=\"max\")\n", + "\n", + "# Draw the coastlines and country borders.\n", + "bm.drawcoastlines()\n", + "bm.drawcountries()\n", + "\n", + "# Draw the color bar\n", + "plt.colorbar(ax=ax, shrink=.7)\n", + "\n", + "# Add a title\n", + "plt.title(\"Terrain Height (m)\", {\"fontsize\" : 20})\n", + "\n", + "plt.show()\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "scrolled": false + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.cm import get_cmap\n", + "from Nio import open_file\n", + "from wrf import getvar, npvalues\n", + "\n", + "\n", + "filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-02-25_18_00_00\"\n", + "#filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\"\n", + "pynio_filename = filename + \".nc\"\n", + "ncfile = open_file(pynio_filename)\n", + "\n", + "# Extract the dewpoint, which will be returned as an xarray.DataArray object (if available). DataArray\n", + "# objects include meta data, similar to NCL's variables.\n", + "dewpointx = getvar(ncfile, \"td\", timeidx=0)\n", + "\n", + "\n", + "# Dewpoint is a 3D variable, so let's just use the lowest level\n", + "dewpointx_sfc = dewpointx[0,:,:]\n", + "\n", + "# Use npvalues to extract the numpy array, since matplotlib does not handle xarray.DataArray natively\n", + "dewpoint_ndarray = npvalues(dewpointx_sfc)\n", + "\n", + "# Get the lat/lon 2D coordinate arrays. Use npvalues to extract the numpy array since basemap does not\n", + "# handle xarray.DataArray natively.\n", + "lons = npvalues(dewpointx_sfc.coords[\"XLONG\"])\n", + "lats = npvalues(dewpointx_sfc.coords[\"XLAT\"])\n", + "\n", + "# Extract the basemap object from the projection information\n", + "wrf_proj = dewpointx_sfc.attrs[\"projection\"]\n", + "bm = wrf_proj.basemap()\n", + "\n", + "# Convert the lat/lon coordinates to projected x,y\n", + "x,y = bm(lons, lats)\n", + "\n", + "# Create the figure\n", + "fig = plt.figure(figsize=(16,16))\n", + "ax = fig.add_axes([0.1,0.1,0.8,0.8])\n", + "\n", + "# Draw filled contours from 200 to 3000 m, every 200 meters.\n", + "levels = np.arange(-40, 40, 5)\n", + "bm.contourf(x, y, dewpoint_ndarray, levels=levels, extend=\"both\")\n", + "\n", + "# Draw the coastlines and country borders.\n", + "bm.drawcoastlines()\n", + "bm.drawcountries()\n", + "\n", + "# Draw the color bar\n", + "plt.colorbar(ax=ax, shrink=.7)\n", + "\n", + "# Add a title\n", + "plt.title(\"Dewpoint Temperature (C)\", {\"fontsize\" : 20})\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false, + "scrolled": false + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.cm import get_cmap\n", + "from Nio import open_file\n", + "from wrf import getvar, vertcross, npvalues\n", + "\n", + "# Open the output netcdf file\n", + "#filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-02-25_18_00_00\"\n", + "filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\"\n", + "pynio_filename = filename + \".nc\"\n", + "ncfile = open_file(pynio_filename)\n", + "\n", + "# Extract pressure and model height\n", + "p = getvar(ncfile, \"pressure\")\n", + "z = getvar(ncfile, \"z\")\n", + "\n", + "# Define a horizontal cross section going left to right using a pivot point in the center of the grid\n", + "# with an angle of 90.0 degrees (0 degrees is vertical).\n", + "# Pivot point is a tuple of (x, y)\n", + "pivot_point = (z.shape[-1] / 2, z.shape[-2] / 2) \n", + "angle = 90.0\n", + "\n", + "# Compute the vertical cross-section interpolation. Include the lat/lon points along the cross-section.\n", + "p_vertx = vertcross(p, z, pivot_point=pivot_point, angle=angle, include_latlon=True)\n", + "\n", + "# Extract the numpy array\n", + "p_vert_array = npvalues(p_vertx)\n", + "\n", + "# Create the figure\n", + "fig = plt.figure(figsize=(20,8))\n", + "ax = plt.axes([0.1,0.1,0.8,0.8])\n", + "\n", + "# Define the levels [0, 50, 100, 150, ....]\n", + "levels = [0 + 50*n for n in range(20)]\n", + "\n", + "# Make the contour plot\n", + "plt.contour(p_vert_array, levels=levels)\n", + "\n", + "# Add the color bar\n", + "plt.colorbar(ax=ax)\n", + "\n", + "# Set the x-ticks to use latitude and longitude labels.\n", + "coord_pairs = npvalues(p_vertx.coords[\"xy_loc\"])\n", + "x_ticks = np.arange(coord_pairs.shape[0])\n", + "x_labels = [pair.latlon_str() for pair in npvalues(coord_pairs)]\n", + "plt.xticks(xy_vals[::100], x_labels[::100]) # Only use every 100th tick.\n", + "\n", + "# Set the y-ticks to be height.\n", + "vert_vals = npvalues(p_vertx.coords[\"vertical\"])\n", + "v_ticks = np.arange(vert_vals.shape[0])\n", + "plt.yticks(v_ticks[::10], vert_vals[::10]) # Only use every 10th tick.\n", + "\n", + "# Set the x-axis and y-axis labels\n", + "plt.xlabel(\"Latitude, Longitude\", fontsize=14)\n", + "plt.ylabel(\"Height (m)\", fontsize=14)\n", + "\n", + "# Add a title\n", + "plt.title(\"Vertical Cross-Section of Pressure (hPa)\", {\"fontsize\" : 20})\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "import numpy as np\n", + "import matplotlib.pyplot as plt\n", + "from matplotlib.cm import get_cmap\n", + "from Nio import open_file\n", + "from wrf import getvar, interplevel, npvalues\n", + "\n", + "# Open the output netcdf file\n", + "#filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2016-02-25_18_00_00\"\n", + "filename = \"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\"\n", + "pynio_filename = filename + \".nc\"\n", + "ncfile = open_file(pynio_filename)\n", + "\n", + "# Extract pressure, model height, destaggered u and v winds\n", + "p = getvar(ncfile, \"pressure\")\n", + "z = getvar(ncfile, \"z\", units=\"dm\")\n", + "ua = getvar(ncfile, \"ua\", units=\"kts\")\n", + "va = getvar(ncfile, \"va\", units=\"kts\")\n", + "\n", + "# Interpolate height, u, and v to to 500 hPa\n", + "ht_500 = interplevel(z, p, 500)\n", + "u_500 = interplevel(ua, p, 500)\n", + "v_500 = interplevel(va, p, 500)\n", + "\n", + "# Get the projection\n", + "wrf_proj = p.attrs[\"projection\"]\n", + "bm = wrf_proj.basemap()\n", + "\n", + "# Basemap needs numpy arrays, extract with npvalues\n", + "lons = npvalues(ht_500.coords[\"XLONG\"])\n", + "lats = npvalues(ht_500.coords[\"XLAT\"])\n", + "\n", + "# Convert the lat/lon coordinates to projected x,y\n", + "x,y = bm(lons, lats)\n", + "\n", + "fig = plt.figure(figsize=(20,20))\n", + "ax = plt.axes([0.1,0.1,0.8,0.8])\n", + "\n", + "# Draw the coastlines and country borders.\n", + "bm.drawcoastlines()\n", + "bm.drawcountries()\n", + "bm.drawstates()\n", + "\n", + "# Make the height contours\n", + "bm.contour(x, y, npvalues(ht_500), 10)\n", + "\n", + "# Make the wind barbs. Only use every 50th in each direction.\n", + "bm.barbs(x[::50,::50], y[::50,::50], npvalues(u_500[::50, ::50]), npvalues(v_500[::50, ::50]))\n", + "\n", + "# Make the color bar\n", + "plt.colorbar(ax=ax, shrink=.7)\n", + "\n", + "# Make the title\n", + "plt.title(\"500 MB Heights (dm) and Wind Barbs (kts)\", {\"fontsize\" : 20})\n", + "\n", + "plt.show()" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [] + } + ], + "metadata": { + "kernelspec": { + "display_name": "Python 2", + "language": "python", + "name": "python2" + }, + "language_info": { + "codemirror_mode": { + "name": "ipython", + "version": 2 + }, + "file_extension": ".py", + "mimetype": "text/x-python", + "name": "python", + "nbconvert_exporter": "python", + "pygments_lexer": "ipython2", + "version": "2.7.11" + } + }, + "nbformat": 4, + "nbformat_minor": 0 +} diff --git a/test/ipynb/WRF_python_demo.ipynb b/test/ipynb/WRF_python_demo.ipynb index 4c9370c..5a649c9 100644 --- a/test/ipynb/WRF_python_demo.ipynb +++ b/test/ipynb/WRF_python_demo.ipynb @@ -311,7 +311,8 @@ "\n", "vard = {varname: getvar(ncfile, varname, method=\"join\", squeeze=False) for varname in wrf_vars}\n", "for varname in wrf_vars:\n", - " print(vard[varname])\n" + " print(vard[varname])\n", + " print (\"\\n\")\n" ] }, { @@ -401,16 +402,17 @@ "\n", "z = getvar(ncfile, \"z\")\n", "p = getvar(ncfile, \"pressure\")\n", - "pivot_point = (z.shape[-2] / 2, z.shape[-1] / 2) \n", + "pivot_point = (z.shape[-1] / 2, z.shape[-2] / 2) \n", "angle = 90.0\n", "\n", "p_vert = vertcross(p, z, pivot_point=pivot_point, angle=angle)\n", "print(p_vert)\n", + "print (\"\\n\")\n", "del p_vert\n", "\n", "# Pressure using start_point and end_point\n", - "start_point = (z.shape[-2]/2, 0)\n", - "end_point = (z.shape[-2]/2, -1)\n", + "start_point = (0, z.shape[-2]/2)\n", + "end_point = (-1, z.shape[-2]/2)\n", "\n", "p_vert = vertcross(p, z, start_point=start_point, end_point=end_point)\n", "print(p_vert)\n", @@ -440,7 +442,7 @@ "angle = 90.0\n", "\n", "t2_line = interpline(t2, pivot_point=pivot_point, angle=angle)\n", - "print(t2_line)\n", + "print(t2_line, \"\\n\")\n", "\n", "del t2_line\n", "\n", @@ -449,7 +451,7 @@ "end_point = (t2.shape[-2]/2, -1)\n", "\n", "t2_line = interpline(t2, start_point=start_point, end_point=end_point)\n", - "print(t2_line)\n", + "print(t2_line, \"\\n\")\n", "\n", "del t2_line, t2" ] @@ -618,8 +620,8 @@ "from wrf import npvalues, getvar\n", "\n", "slp = getvar(ncfile, \"slp\")\n", - "lons = slp.coords[\"XLONG\"]\n", - "lats = slp.coords[\"XLAT\"]\n", + "lons = npvalues(slp.coords[\"XLONG\"])\n", + "lats = npvalues(slp.coords[\"XLAT\"])\n", "\n", "wrf_proj = slp.attrs[\"projection\"]\n", "cart_proj = wrf_proj.cartopy()\n", @@ -703,11 +705,23 @@ }, { "cell_type": "code", - "execution_count": null, + "execution_count": 1, "metadata": { "collapsed": false }, - "outputs": [], + "outputs": [ + { + "ename": "NameError", + "evalue": "name 'ncfile' is not defined", + "output_type": "error", + "traceback": [ + "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m", + "\u001b[0;31mNameError\u001b[0m Traceback (most recent call last)", + "\u001b[0;32m\u001b[0m in \u001b[0;36m\u001b[0;34m()\u001b[0m\n\u001b[1;32m 6\u001b[0m \u001b[0;32mfrom\u001b[0m \u001b[0mwrf\u001b[0m \u001b[0;32mimport\u001b[0m \u001b[0mgetvar\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mvertcross\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0mnpvalues\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 7\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0;32m----> 8\u001b[0;31m \u001b[0mp\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetvar\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mncfile\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"pressure\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[0m\u001b[1;32m 9\u001b[0m \u001b[0mz\u001b[0m \u001b[0;34m=\u001b[0m \u001b[0mgetvar\u001b[0m\u001b[0;34m(\u001b[0m\u001b[0mncfile\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0;34m\"z\"\u001b[0m\u001b[0;34m,\u001b[0m \u001b[0munits\u001b[0m\u001b[0;34m=\u001b[0m\u001b[0;34m\"dm\"\u001b[0m\u001b[0;34m)\u001b[0m\u001b[0;34m\u001b[0m\u001b[0m\n\u001b[1;32m 10\u001b[0m \u001b[0;34m\u001b[0m\u001b[0m\n", + "\u001b[0;31mNameError\u001b[0m: name 'ncfile' is not defined" + ] + } + ], "source": [ "# Cross-section of pressure using xarray's builtin plotting\n", "import numpy as np\n", @@ -719,7 +733,7 @@ "p = getvar(ncfile, \"pressure\")\n", "z = getvar(ncfile, \"z\", units=\"dm\")\n", "\n", - "pivot_point = (z.shape[-2] / 2, z.shape[-1] / 2) \n", + "pivot_point = (z.shape[-1] / 2, z.shape[-2] / 2) \n", "angle = 90.0\n", "\n", "p_vert = vertcross(p, z, pivot_point=pivot_point, angle=angle)\n", @@ -756,6 +770,7 @@ "dir = \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi\"\n", "ncfilenames = [os.path.join(dir, x) for x in os.listdir(dir) if x.find(\"_d02_\") > 0]\n", "ncfiles = [nc(x) for x in ncfilenames]\n", + "\n", "#print (ncfiles[0].variables[\"XLONG\"][0,0,-1], ncfiles[0].variables[\"XLONG\"][-1,0,-1])\n", "#print (ncfiles[1].variables[\"XLONG\"][0,0,-1], ncfiles[1].variables[\"XLONG\"][-1,0,-1])\n", "#print (ncfiles[-1].variables[\"XLONG\"][0,0,-1], ncfiles[-1].variables[\"XLONG\"][-1,0,-1])\n" @@ -867,10 +882,132 @@ " \"pvo\", \"pw\", \"rh2\", \"rh\", \"slp\", \"ter\", \"td2\", \"td\", \"tc\",\n", " \"theta\", \"tk\", \"tv\", \"twb\", \"updraft_helicity\", \"ua\", \"va\", \n", " \"wa\", \"uvmet10\", \"uvmet\", \"z\", \"ctt\"]\n", + "#wrf_vars = [\"cape_2d\"]\n", "\n", - "vard = {varname: getvar(ncfiles, varname, method=\"join\", squeeze=False) for varname in wrf_vars}\n", + "vard = {}\n", + "for varname in wrf_vars:\n", + " print (varname)\n", + " vard[varname] = getvar(ncfiles, varname, timeidx=None, method=\"cat\", squeeze=False)\n", + " \n", + "#vard = {varname: getvar(ncfiles, varname, method=\"join\", squeeze=False) for varname in wrf_vars}\n", "for varname in wrf_vars:\n", - " print(vard[varname])" + " print(vard[varname])\n", + " " + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "import os\n", + "from wrf import getvar\n", + "from netCDF4 import Dataset as nc\n", + "\n", + "dir = \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi\"\n", + "ncfilenames = [os.path.join(dir, x) for x in os.listdir(dir) if x.find(\"_d02_\") > 0]\n", + "ncfiles = [nc(x) for x in ncfilenames]\n", + "\n", + "# Pressure using pivot and angle\n", + "from wrf import getvar, vertcross\n", + "\n", + "timeidx = 0\n", + "z = getvar(ncfiles, \"z\", timeidx, method=\"join\")\n", + "p = getvar(ncfiles, \"pressure\", timeidx, method=\"join\")\n", + "pivot_point = (z.shape[-1] / 2, z.shape[-2] / 2) \n", + "angle = 40.0\n", + "\n", + "p_vert = vertcross(p, z, pivot_point=pivot_point, angle=angle)\n", + "print(p_vert)\n", + "print (\"\\n\")\n", + "del p_vert\n", + "\n", + "# Pressure using start_point and end_point\n", + "start_point = (0, z.shape[-2]/2)\n", + "end_point = (-1, z.shape[-2]/2)\n", + "\n", + "p_vert = vertcross(p, z, start_point=start_point, end_point=end_point)\n", + "print(p_vert)\n", + "del p_vert, p, z" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "import os\n", + "from wrf import getvar\n", + "from netCDF4 import Dataset as nc\n", + "\n", + "dir = \"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi\"\n", + "ncfilenames = [os.path.join(dir, x) for x in os.listdir(dir) if x.find(\"_d02_\") > 0]\n", + "ncfiles = [nc(x) for x in ncfilenames]\n", + "\n", + "timeidx = None\n", + "\n", + "# T2 using pivot and angle\n", + "from wrf import interpline, getvar, npvalues\n", + "\n", + "t2 = getvar(ncfiles, \"T2\", timeidx)\n", + "pivot_point = (t2.shape[-2] / 2, t2.shape[-1] / 2) \n", + "angle = 90.0\n", + "\n", + "t2_line = interpline(t2, pivot_point=pivot_point, angle=angle, inc_latlon=True)\n", + "print(t2_line)\n", + "print(\"\\n\")\n", + "\n", + "\n", + "del t2_line\n", + "\n", + "# T2 using start_point and end_point\n", + "start_point = (t2.shape[-2]/2, 0)\n", + "end_point = (t2.shape[-2]/2, -1)\n", + "\n", + "t2_line = interpline(t2, start_point=start_point, end_point=end_point, inc_latlon=True)\n", + "print(t2_line)\n", + "print(\"\\n\")\n", + "\n", + "del t2_line, t2" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "from wrf import getvar\n" + ] + }, + { + "cell_type": "code", + "execution_count": 2, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "getvar?" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": true + }, + "outputs": [], + "source": [ + "\n" ] }, { diff --git a/test/ipynb/nocopy_test.ipynb b/test/ipynb/nocopy_test.ipynb index 4243163..38e1031 100644 --- a/test/ipynb/nocopy_test.ipynb +++ b/test/ipynb/nocopy_test.ipynb @@ -197,7 +197,1152 @@ "from netCDF4 import Dataset as nc\n", "lambertnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n", "uvmet = getvar([lambertnc,lambertnc], \"uvmet\", timeidx=None)\n", - "print (uvmet)" + "uvmet10 = getvar([lambertnc,lambertnc], \"uvmet10\", timeidx=None)\n", + "uv_wspd_wdir = getvar([lambertnc,lambertnc], \"wspd_wdir_uvmet\", timeidx=None)\n", + "uv_wspd_wdir10 = getvar([lambertnc,lambertnc], \"wspd_wdir_uvmet10\", timeidx=None)\n", + "wspd_wdir = getvar([lambertnc,lambertnc], \"wspd_wdir\", timeidx=None)\n", + "wspd_wdir10 = getvar([lambertnc,lambertnc], \"wspd_wdir10\", timeidx=None)\n", + "print (uvmet)\n", + "print (uvmet10)\n", + "print (uv_wspd_wdir)\n", + "print (uv_wspd_wdir10)\n", + "print (wspd_wdir)\n", + "print (wspd_wdir10)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from wrf import (ALL_TIMES, npvalues, Constants, getvar, extract_vars, destagger, \n", + " interp1d, interp2dxy, interpz3d, \n", + " slp, tk, td, rh, uvmet, smooth2d)\n", + "from netCDF4 import Dataset as nc\n", + "\n", + "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrf_vortex_multi/wrfout_d01_2005-08-28_00:00:00\")\n", + "timeidx = ALL_TIMES\n", + "method = \"cat\"\n", + "squeeze = True\n", + "cache = None\n", + "\n" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "\n", + "varnames=(\"T\", \"P\", \"PB\", \"QVAPOR\", \"PH\", \"PHB\")\n", + "ncvars = extract_vars(wrfnc, timeidx, varnames, method, squeeze, cache,\n", + " meta=True)\n", + "\n", + "t = ncvars[\"T\"]\n", + "p = ncvars[\"P\"]\n", + "pb = ncvars[\"PB\"]\n", + "ph = ncvars[\"PH\"]\n", + "phb = ncvars[\"PHB\"]\n", + "\n", + "\n", + "ncvars = extract_vars(wrfnc, timeidx, (\"QVAPOR\",), method, squeeze, cache,\n", + " meta=False)\n", + "qvapor = ncvars[\"QVAPOR\"]\n", + "\n", + " \n", + "full_t = t + Constants.T_BASE\n", + "full_p = p + pb\n", + "qvapor[qvapor < 0] = 0.\n", + " \n", + "full_ph = (ph + phb) / Constants.G\n", + " \n", + "destag_ph = destagger(full_ph, -3)\n", + " \n", + "_tk = tk(full_p, full_t)\n", + "_slp = slp(destag_ph, _tk, full_p, qvapor)\n", + "_td = td(full_p, qvapor)\n", + "_smooth2d = smooth2d(npvalues(_slp), 3)\n", + "\n", + "_uvmet = getvar(wrfnc, \"uvmet\", timeidx=1)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "print (_tk)\n", + "print (\"\\n\")\n", + "print (_slp)\n", + "print (\"\\n\")\n", + "print (_td)\n", + "print (\"\\n\")\n", + "print (_smooth2d)\n", + "print (_uvmet)" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": { + "collapsed": false + }, + "outputs": [], + "source": [ + "from wrf import (ALL_TIMES, npvalues, Constants, getvar, extract_vars, destagger, \n", + " interp1d, interp2dxy, interpz3d, \n", + " slp, tk, td, rh, uvmet, smooth2d, extract_global_attrs)\n", + "from math import fabs, log, tan, sin, cos\n", + "from wrf.util import either\n", + "from netCDF4 import Dataset as nc\n", + "\n", + "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n", + "wrfnc = [wrfnc, wrfnc, wrfnc]\n", + "ten_m = True\n", + "method = \"cat\"\n", + "squeeze=True\n", + "cache=None\n", + "timeidx = ALL_TIMES\n", + "\n", + "if not ten_m:\n", + " varname = either(\"U\", \"UU\")(wrfnc)\n", + " u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n", + " meta=True)\n", + " \n", + " #renamed = u_vars[varname].rename({\"west_east_stag\" : \"test\"})\n", + " u = destagger(u_vars[varname], -1, meta=True)\n", + " #print (renamed)\n", + " #u = destagger(renamed, -1, meta=True)\n", + " #print (u)\n", + "\n", + " varname = either(\"V\", \"VV\")(wrfnc)\n", + " v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n", + " meta=True)\n", + " v = destagger(v_vars[varname], -2)\n", + "else:\n", + " varname = either(\"U10\", \"UU\")(wrfnc)\n", + " u_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n", + " meta=True)\n", + " u = (u_vars[varname] if varname == \"U10\" else \n", + " destagger(u_vars[varname][...,0,:,:], -1)) \n", + "\n", + " varname = either(\"V10\", \"VV\")(wrfnc)\n", + " v_vars = extract_vars(wrfnc, timeidx, varname, method, squeeze, cache,\n", + " meta=True)\n", + " v = (v_vars[varname] if varname == \"V10\" else \n", + " destagger(v_vars[varname][...,0,:,:], -2))\n", + "\n", + "map_proj_attrs = extract_global_attrs(wrfnc, attrs=\"MAP_PROJ\")\n", + "map_proj = map_proj_attrs[\"MAP_PROJ\"]\n", + "\n", + "# 1 - Lambert\n", + "# 2 - Polar Stereographic\n", + "# 3 - Mercator\n", + "# 6 - Lat/Lon\n", + "# Note: NCL has no code to handle other projections (0,4,5) as they \n", + "# don't appear to be supported any longer\n", + "\n", + "if map_proj in (0,3,6):\n", + " # No rotation needed for Mercator and Lat/Lon, but still need\n", + " # u,v aggregated in to one array\n", + "\n", + " end_idx = -3 if not ten_m else -2\n", + " resdim = (2,) + u.shape[0:end_idx] + u.shape[end_idx:]\n", + "\n", + " # Make a new output array for the result\n", + " res = np.empty(resdim, u.dtype)\n", + "\n", + " # For 2D array, this makes (0,...,:,:) and (1,...,:,:)\n", + " # For 3D array, this makes (0,...,:,:,:) and (1,...,:,:,:)\n", + " idx0 = (0,) + (Ellipsis,) + (slice(None),)*(-end_idx)\n", + " idx1 = (1,) + (Ellipsis,) + (slice(None),)*(-end_idx)\n", + "\n", + " res[idx0] = u[:]\n", + " res[idx1] = v[:]\n", + "\n", + " result = res\n", + "elif map_proj in (1,2):\n", + " lat_attrs = extract_global_attrs(wrfnc, attrs=(\"TRUELAT1\",\n", + " \"TRUELAT2\"))\n", + " radians_per_degree = Constants.PI/180.0\n", + " # Rotation needed for Lambert and Polar Stereographic\n", + " true_lat1 = lat_attrs[\"TRUELAT1\"]\n", + " true_lat2 = lat_attrs[\"TRUELAT2\"]\n", + "\n", + " try:\n", + " lon_attrs = extract_global_attrs(wrfnc, attrs=\"STAND_LON\")\n", + " except AttributeError:\n", + " try:\n", + " cen_lon_attrs = extract_global_attrs(wrfnc, attrs=\"CEN_LON\")\n", + " except AttributeError:\n", + " raise RuntimeError(\"longitude attributes not found in NetCDF\")\n", + " else:\n", + " cen_lon = cen_lon_attrs[\"CEN_LON\"]\n", + " else:\n", + " cen_lon = lon_attrs[\"STAND_LON\"]\n", + "\n", + "\n", + " varname = either(\"XLAT_M\", \"XLAT\")(wrfnc)\n", + " xlat_var = extract_vars(wrfnc, timeidx, varname, \n", + " method, squeeze, cache, meta=True)\n", + " lat = xlat_var[varname]\n", + "\n", + " varname = either(\"XLONG_M\", \"XLONG\")(wrfnc)\n", + " xlon_var = extract_vars(wrfnc, timeidx, varname, \n", + " method, squeeze, cache, meta=False)\n", + " lon = xlon_var[varname]\n", + "\n", + " if map_proj == 1:\n", + " if((fabs(true_lat1 - true_lat2) > 0.1) and\n", + " (fabs(true_lat2 - 90.) > 0.1)): \n", + " cone = (log(cos(true_lat1*radians_per_degree)) \n", + " - log(cos(true_lat2*radians_per_degree)))\n", + " cone = (cone / \n", + " (log(tan((45.-fabs(true_lat1/2.))*radians_per_degree)) \n", + " - log(tan((45.-fabs(true_lat2/2.))*radians_per_degree)))) \n", + " else:\n", + " cone = sin(fabs(true_lat1)*radians_per_degree)\n", + " else:\n", + " cone = 1\n", + " \n", + " result = uvmet(u, v, lat, lon, cen_lon, cone, meta=True)\n", + " \n", + "print (result)" + ] + }, + { + "cell_type": "code", + "execution_count": 1, + "metadata": { + "collapsed": false + }, + "outputs": [ + { + "name": "stdout", + "output_type": "stream", + "text": [ + "\n", + "array([[[ 1284.37878418, 1286.88208008, 1288.14526367, ...,\n", + " 1042.98413086, 1043.49230957, 1042.76550293],\n", + " [ 1290.37072754, 1305.41650391, 1315.56066895, ...,\n", + " 1058.42687988, 1054.85864258, 1048.94299316],\n", + " [ 1297.45019531, 1319.68493652, 1332.59887695, ...,\n", + " 1067.23571777, 1063.85803223, 1054.76965332],\n", + " ..., \n", + " [ 2282.65576172, 2248.74243164, 2226.04858398, ...,\n", + " 680.19281006, 715.02349854, 752.45947266],\n", + " [ 2261.47558594, 2233.76464844, 2212.33129883, ...,\n", + " 673.80822754, 712.69897461, 748.10424805],\n", + " [ 2241.63647461, 2225.63354492, 2209.28881836, ...,\n", + " 665.93286133, 703.76361084, 743.64093018]],\n", + "\n", + " [[ 1284.37878418, 1286.88208008, 1288.14526367, ...,\n", + " 1042.98413086, 1043.49230957, 1042.76550293],\n", + " [ 1290.37072754, 1305.41650391, 1315.56066895, ...,\n", + " 1058.42687988, 1054.85864258, 1048.94299316],\n", + " [ 1297.45019531, 1319.68493652, 1332.59887695, ...,\n", + " 1067.23571777, 1063.85803223, 1054.76965332],\n", + " ..., \n", + " [ 2282.65576172, 2248.74243164, 2226.04858398, ...,\n", + " 680.19281006, 715.02349854, 752.45947266],\n", + " [ 2261.47558594, 2233.76464844, 2212.33129883, ...,\n", + " 673.80822754, 712.69897461, 748.10424805],\n", + " [ 2241.63647461, 2225.63354492, 2209.28881836, ...,\n", + " 665.93286133, 703.76361084, 743.64093018]]], dtype=float32)\n", + "Coordinates:\n", + " XLONG (south_north, west_east) float32 -127.749 -127.627 -127.504 ...\n", + " XLAT (south_north, west_east) float32 13.69 13.726 13.7617 ...\n", + " * Time (Time) datetime64[ns] 2010-06-13T21:00:00 2010-06-13T21:00:00\n", + " * south_north (south_north) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 ...\n", + " * west_east (west_east) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ...\n", + " datetime (Time) datetime64[ns] 2010-06-13T21:00:00 2010-06-13T21:00:00\n", + "Attributes:\n", + " FieldType: 104\n", + " units: Pa\n", + " stagger: \n", + " coordinates: XLONG XLAT\n", + " projection: LambertConformal(bottom_left=(13.69003, -127.74881), top_right=(53.277489, -53.45282), stand_lon=-101.0, moad_cen_lat=39.0000038147, pole_lat=90.0, pole_lon=0.0)\n", + " PlotLevelID: 500 m\n", + " missing_value: 9.96920996839e+36\n", + " _FillValue: 9.96920996839e+36\n", + "\n", + "\n", + "\n", + "array([[ 0., 162.],\n", + " [ 1., 162.],\n", + " [ 2., 162.],\n", + " [ 3., 162.],\n", + " [ 4., 162.],\n", + " [ 5., 162.],\n", + " [ 6., 162.],\n", + " [ 7., 162.],\n", + " [ 8., 162.],\n", + " [ 9., 162.],\n", + " [ 10., 162.],\n", + " [ 11., 162.],\n", + " [ 12., 162.],\n", + " [ 13., 162.],\n", + " [ 14., 162.],\n", + " [ 15., 162.],\n", + " [ 16., 162.],\n", + " [ 17., 162.],\n", + " [ 18., 162.],\n", + " [ 19., 162.],\n", + " [ 20., 162.],\n", + " [ 21., 162.],\n", + " [ 22., 162.],\n", + " [ 23., 162.],\n", + " [ 24., 162.],\n", + " [ 25., 162.],\n", + " [ 26., 162.],\n", + " [ 27., 162.],\n", + " [ 28., 162.],\n", + " [ 29., 162.],\n", + " [ 30., 162.],\n", + " [ 31., 162.],\n", + " [ 32., 162.],\n", + " [ 33., 162.],\n", + " [ 34., 162.],\n", + " [ 35., 162.],\n", + " [ 36., 162.],\n", + " [ 37., 162.],\n", + " [ 38., 162.],\n", + " [ 39., 162.],\n", + " [ 40., 162.],\n", + " [ 41., 162.],\n", + " [ 42., 162.],\n", + " [ 43., 162.],\n", + " [ 44., 162.],\n", + " [ 45., 162.],\n", + " [ 46., 162.],\n", + " [ 47., 162.],\n", + " [ 48., 162.],\n", + " [ 49., 162.],\n", + " [ 50., 162.],\n", + " [ 51., 162.],\n", + " [ 52., 162.],\n", + " [ 53., 162.],\n", + " [ 54., 162.],\n", + " [ 55., 162.],\n", + " [ 56., 162.],\n", + " [ 57., 162.],\n", + " [ 58., 162.],\n", + " [ 59., 162.],\n", + " [ 60., 162.],\n", + " [ 61., 162.],\n", + " [ 62., 162.],\n", + " [ 63., 162.],\n", + " [ 64., 162.],\n", + " [ 65., 162.],\n", + " [ 66., 162.],\n", + " [ 67., 162.],\n", + " [ 68., 162.],\n", + " [ 69., 162.],\n", + " [ 70., 162.],\n", + " [ 71., 162.],\n", + " [ 72., 162.],\n", + " [ 73., 162.],\n", + " [ 74., 162.],\n", + " [ 75., 162.],\n", + " [ 76., 162.],\n", + " [ 77., 162.],\n", + " [ 78., 162.],\n", + " [ 79., 162.],\n", + " [ 80., 162.],\n", + " [ 81., 162.],\n", + " [ 82., 162.],\n", + " [ 83., 162.],\n", + " [ 84., 162.],\n", + " [ 85., 162.],\n", + " [ 86., 162.],\n", + " [ 87., 162.],\n", + " [ 88., 162.],\n", + " [ 89., 162.],\n", + " [ 90., 162.],\n", + " [ 91., 162.],\n", + " [ 92., 162.],\n", + " [ 93., 162.],\n", + " [ 94., 162.],\n", + " [ 95., 162.],\n", + " [ 96., 162.],\n", + " [ 97., 162.],\n", + " [ 98., 162.],\n", + " [ 99., 162.],\n", + " [ 100., 162.],\n", + " [ 101., 162.],\n", + " [ 102., 162.],\n", + " [ 103., 162.],\n", + " [ 104., 162.],\n", + " [ 105., 162.],\n", + " [ 106., 162.],\n", + " [ 107., 162.],\n", + " [ 108., 162.],\n", + " [ 109., 162.],\n", + " [ 110., 162.],\n", + " [ 111., 162.],\n", + " [ 112., 162.],\n", + " [ 113., 162.],\n", + " [ 114., 162.],\n", + " [ 115., 162.],\n", + " [ 116., 162.],\n", + " [ 117., 162.],\n", + " [ 118., 162.],\n", + " [ 119., 162.],\n", + " [ 120., 162.],\n", + " [ 121., 162.],\n", + " [ 122., 162.],\n", + " [ 123., 162.],\n", + " [ 124., 162.],\n", + " [ 125., 162.],\n", + " [ 126., 162.],\n", + " [ 127., 162.],\n", + " [ 128., 162.],\n", + " [ 129., 162.],\n", + " [ 130., 162.],\n", + " [ 131., 162.],\n", + " [ 132., 162.],\n", + " [ 133., 162.],\n", + " [ 134., 162.],\n", + " [ 135., 162.],\n", + " [ 136., 162.],\n", + " [ 137., 162.],\n", + " [ 138., 162.],\n", + " [ 139., 162.],\n", + " [ 140., 162.],\n", + " [ 141., 162.],\n", + " [ 142., 162.],\n", + " [ 143., 162.],\n", + " [ 144., 162.],\n", + " [ 145., 162.],\n", + " [ 146., 162.],\n", + " [ 147., 162.],\n", + " [ 148., 162.],\n", + " [ 149., 162.],\n", + " [ 150., 162.],\n", + " [ 151., 162.],\n", + " [ 152., 162.],\n", + " [ 153., 162.],\n", + " [ 154., 162.],\n", + " [ 155., 162.],\n", + " [ 156., 162.],\n", + " [ 157., 162.],\n", + " [ 158., 162.],\n", + " [ 159., 162.],\n", + " [ 160., 162.],\n", + " [ 161., 162.],\n", + " [ 162., 162.],\n", + " [ 163., 162.],\n", + " [ 164., 162.],\n", + " [ 165., 162.],\n", + " [ 166., 162.],\n", + " [ 167., 162.],\n", + " [ 168., 162.],\n", + " [ 169., 162.],\n", + " [ 170., 162.],\n", + " [ 171., 162.],\n", + " [ 172., 162.],\n", + " [ 173., 162.],\n", + " [ 174., 162.],\n", + " [ 175., 162.],\n", + " [ 176., 162.],\n", + " [ 177., 162.],\n", + " [ 178., 162.],\n", + " [ 179., 162.],\n", + " [ 180., 162.],\n", + " [ 181., 162.],\n", + " [ 182., 162.],\n", + " [ 183., 162.],\n", + " [ 184., 162.],\n", + " [ 185., 162.],\n", + " [ 186., 162.],\n", + " [ 187., 162.],\n", + " [ 188., 162.],\n", + " [ 189., 162.],\n", + " [ 190., 162.],\n", + " [ 191., 162.],\n", + " [ 192., 162.],\n", + " [ 193., 162.],\n", + " [ 194., 162.],\n", + " [ 195., 162.],\n", + " [ 196., 162.],\n", + " [ 197., 162.],\n", + " [ 198., 162.],\n", + " [ 199., 162.],\n", + " [ 200., 162.],\n", + " [ 201., 162.],\n", + " [ 202., 162.],\n", + " [ 203., 162.],\n", + " [ 204., 162.],\n", + " [ 205., 162.],\n", + " [ 206., 162.],\n", + " [ 207., 162.],\n", + " [ 208., 162.],\n", + " [ 209., 162.],\n", + " [ 210., 162.],\n", + " [ 211., 162.],\n", + " [ 212., 162.],\n", + " [ 213., 162.],\n", + " [ 214., 162.],\n", + " [ 215., 162.],\n", + " [ 216., 162.],\n", + " [ 217., 162.],\n", + " [ 218., 162.],\n", + " [ 219., 162.],\n", + " [ 220., 162.],\n", + " [ 221., 162.],\n", + " [ 222., 162.],\n", + " [ 223., 162.],\n", + " [ 224., 162.],\n", + " [ 225., 162.],\n", + " [ 226., 162.],\n", + " [ 227., 162.],\n", + " [ 228., 162.],\n", + " [ 229., 162.],\n", + " [ 230., 162.],\n", + " [ 231., 162.],\n", + " [ 232., 162.],\n", + " [ 233., 162.],\n", + " [ 234., 162.],\n", + " [ 235., 162.],\n", + " [ 236., 162.],\n", + " [ 237., 162.],\n", + " [ 238., 162.],\n", + " [ 239., 162.],\n", + " [ 240., 162.],\n", + " [ 241., 162.],\n", + " [ 242., 162.],\n", + " [ 243., 162.],\n", + " [ 244., 162.],\n", + " [ 245., 162.],\n", + " [ 246., 162.],\n", + " [ 247., 162.],\n", + " [ 248., 162.],\n", + " [ 249., 162.],\n", + " [ 250., 162.],\n", + " [ 251., 162.],\n", + " [ 252., 162.],\n", + " [ 253., 162.],\n", + " [ 254., 162.],\n", + " [ 255., 162.],\n", + " [ 256., 162.],\n", + " [ 257., 162.],\n", + " [ 258., 162.],\n", + " [ 259., 162.],\n", + " [ 260., 162.],\n", + " [ 261., 162.],\n", + " [ 262., 162.],\n", + " [ 263., 162.],\n", + " [ 264., 162.],\n", + " [ 265., 162.],\n", + " [ 266., 162.],\n", + " [ 267., 162.],\n", + " [ 268., 162.],\n", + " [ 269., 162.],\n", + " [ 270., 162.],\n", + " [ 271., 162.],\n", + " [ 272., 162.],\n", + " [ 273., 162.],\n", + " [ 274., 162.],\n", + " [ 275., 162.],\n", + " [ 276., 162.],\n", + " [ 277., 162.],\n", + " [ 278., 162.],\n", + " [ 279., 162.],\n", + " [ 280., 162.],\n", + " [ 281., 162.],\n", + " [ 282., 162.],\n", + " [ 283., 162.],\n", + " [ 284., 162.],\n", + " [ 285., 162.],\n", + " [ 286., 162.],\n", + " [ 287., 162.],\n", + " [ 288., 162.],\n", + " [ 289., 162.],\n", + " [ 290., 162.],\n", + " [ 291., 162.],\n", + " [ 292., 162.],\n", + " [ 293., 162.],\n", + " [ 294., 162.],\n", + " [ 295., 162.],\n", + " [ 296., 162.],\n", + " [ 297., 162.],\n", + " [ 298., 162.],\n", + " [ 299., 162.],\n", + " [ 300., 162.],\n", + " [ 301., 162.],\n", + " [ 302., 162.],\n", + " [ 303., 162.],\n", + " [ 304., 162.],\n", + " [ 305., 162.],\n", + " [ 306., 162.],\n", + " [ 307., 162.],\n", + " [ 308., 162.],\n", + " [ 309., 162.],\n", + " [ 310., 162.],\n", + " [ 311., 162.],\n", + " [ 312., 162.],\n", + " [ 313., 162.],\n", + " [ 314., 162.],\n", + " [ 315., 162.],\n", + " [ 316., 162.],\n", + " [ 317., 162.],\n", + " [ 318., 162.],\n", + " [ 319., 162.],\n", + " [ 320., 162.],\n", + " [ 321., 162.],\n", + " [ 322., 162.],\n", + " [ 323., 162.],\n", + " [ 324., 162.],\n", + " [ 325., 162.],\n", + " [ 326., 162.],\n", + " [ 327., 162.],\n", + " [ 328., 162.],\n", + " [ 329., 162.],\n", + " [ 330., 162.],\n", + " [ 331., 162.],\n", + " [ 332., 162.],\n", + " [ 333., 162.],\n", + " [ 334., 162.],\n", + " [ 335., 162.],\n", + " [ 336., 162.],\n", + " [ 337., 162.],\n", + " [ 338., 162.],\n", + " [ 339., 162.],\n", + " [ 340., 162.],\n", + " [ 341., 162.],\n", + " [ 342., 162.],\n", + " [ 343., 162.],\n", + " [ 344., 162.],\n", + " [ 345., 162.],\n", + " [ 346., 162.],\n", + " [ 347., 162.],\n", + " [ 348., 162.],\n", + " [ 349., 162.],\n", + " [ 350., 162.],\n", + " [ 351., 162.],\n", + " [ 352., 162.],\n", + " [ 353., 162.],\n", + " [ 354., 162.],\n", + " [ 355., 162.],\n", + " [ 356., 162.],\n", + " [ 357., 162.],\n", + " [ 358., 162.],\n", + " [ 359., 162.],\n", + " [ 360., 162.],\n", + " [ 361., 162.],\n", + " [ 362., 162.],\n", + " [ 363., 162.],\n", + " [ 364., 162.],\n", + " [ 365., 162.],\n", + " [ 366., 162.],\n", + " [ 367., 162.],\n", + " [ 368., 162.],\n", + " [ 369., 162.],\n", + " [ 370., 162.],\n", + " [ 371., 162.],\n", + " [ 372., 162.],\n", + " [ 373., 162.],\n", + " [ 374., 162.],\n", + " [ 375., 162.],\n", + " [ 376., 162.],\n", + " [ 377., 162.],\n", + " [ 378., 162.],\n", + " [ 379., 162.],\n", + " [ 380., 162.],\n", + " [ 381., 162.],\n", + " [ 382., 162.],\n", + " [ 383., 162.],\n", + " [ 384., 162.],\n", + " [ 385., 162.],\n", + " [ 386., 162.],\n", + " [ 387., 162.],\n", + " [ 388., 162.],\n", + " [ 389., 162.],\n", + " [ 390., 162.],\n", + " [ 391., 162.],\n", + " [ 392., 162.],\n", + " [ 393., 162.],\n", + " [ 394., 162.],\n", + " [ 395., 162.],\n", + " [ 396., 162.],\n", + " [ 397., 162.],\n", + " [ 398., 162.],\n", + " [ 399., 162.],\n", + " [ 400., 162.],\n", + " [ 401., 162.],\n", + " [ 402., 162.],\n", + " [ 403., 162.],\n", + " [ 404., 162.],\n", + " [ 405., 162.],\n", + " [ 406., 162.],\n", + " [ 407., 162.],\n", + " [ 408., 162.],\n", + " [ 409., 162.],\n", + " [ 410., 162.],\n", + " [ 411., 162.],\n", + " [ 412., 162.]])\n", + "Coordinates:\n", + " * x_y (x_y) \n", + "array([[ 0., 162.],\n", + " [ 1., 162.],\n", + " [ 2., 162.],\n", + " [ 3., 162.],\n", + " [ 4., 162.],\n", + " [ 5., 162.],\n", + " [ 6., 162.],\n", + " [ 7., 162.],\n", + " [ 8., 162.],\n", + " [ 9., 162.],\n", + " [ 10., 162.],\n", + " [ 11., 162.],\n", + " [ 12., 162.],\n", + " [ 13., 162.],\n", + " [ 14., 162.],\n", + " [ 15., 162.],\n", + " [ 16., 162.],\n", + " [ 17., 162.],\n", + " [ 18., 162.],\n", + " [ 19., 162.],\n", + " [ 20., 162.],\n", + " [ 21., 162.],\n", + " [ 22., 162.],\n", + " [ 23., 162.],\n", + " [ 24., 162.],\n", + " [ 25., 162.],\n", + " [ 26., 162.],\n", + " [ 27., 162.],\n", + " [ 28., 162.],\n", + " [ 29., 162.],\n", + " [ 30., 162.],\n", + " [ 31., 162.],\n", + " [ 32., 162.],\n", + " [ 33., 162.],\n", + " [ 34., 162.],\n", + " [ 35., 162.],\n", + " [ 36., 162.],\n", + " [ 37., 162.],\n", + " [ 38., 162.],\n", + " [ 39., 162.],\n", + " [ 40., 162.],\n", + " [ 41., 162.],\n", + " [ 42., 162.],\n", + " [ 43., 162.],\n", + " [ 44., 162.],\n", + " [ 45., 162.],\n", + " [ 46., 162.],\n", + " [ 47., 162.],\n", + " [ 48., 162.],\n", + " [ 49., 162.],\n", + " [ 50., 162.],\n", + " [ 51., 162.],\n", + " [ 52., 162.],\n", + " [ 53., 162.],\n", + " [ 54., 162.],\n", + " [ 55., 162.],\n", + " [ 56., 162.],\n", + " [ 57., 162.],\n", + " [ 58., 162.],\n", + " [ 59., 162.],\n", + " [ 60., 162.],\n", + " [ 61., 162.],\n", + " [ 62., 162.],\n", + " [ 63., 162.],\n", + " [ 64., 162.],\n", + " [ 65., 162.],\n", + " [ 66., 162.],\n", + " [ 67., 162.],\n", + " [ 68., 162.],\n", + " [ 69., 162.],\n", + " [ 70., 162.],\n", + " [ 71., 162.],\n", + " [ 72., 162.],\n", + " [ 73., 162.],\n", + " [ 74., 162.],\n", + " [ 75., 162.],\n", + " [ 76., 162.],\n", + " [ 77., 162.],\n", + " [ 78., 162.],\n", + " [ 79., 162.],\n", + " [ 80., 162.],\n", + " [ 81., 162.],\n", + " [ 82., 162.],\n", + " [ 83., 162.],\n", + " [ 84., 162.],\n", + " [ 85., 162.],\n", + " [ 86., 162.],\n", + " [ 87., 162.],\n", + " [ 88., 162.],\n", + " [ 89., 162.],\n", + " [ 90., 162.],\n", + " [ 91., 162.],\n", + " [ 92., 162.],\n", + " [ 93., 162.],\n", + " [ 94., 162.],\n", + " [ 95., 162.],\n", + " [ 96., 162.],\n", + " [ 97., 162.],\n", + " [ 98., 162.],\n", + " [ 99., 162.],\n", + " [ 100., 162.],\n", + " [ 101., 162.],\n", + " [ 102., 162.],\n", + " [ 103., 162.],\n", + " [ 104., 162.],\n", + " [ 105., 162.],\n", + " [ 106., 162.],\n", + " [ 107., 162.],\n", + " [ 108., 162.],\n", + " [ 109., 162.],\n", + " [ 110., 162.],\n", + " [ 111., 162.],\n", + " [ 112., 162.],\n", + " [ 113., 162.],\n", + " [ 114., 162.],\n", + " [ 115., 162.],\n", + " [ 116., 162.],\n", + " [ 117., 162.],\n", + " [ 118., 162.],\n", + " [ 119., 162.],\n", + " [ 120., 162.],\n", + " [ 121., 162.],\n", + " [ 122., 162.],\n", + " [ 123., 162.],\n", + " [ 124., 162.],\n", + " [ 125., 162.],\n", + " [ 126., 162.],\n", + " [ 127., 162.],\n", + " [ 128., 162.],\n", + " [ 129., 162.],\n", + " [ 130., 162.],\n", + " [ 131., 162.],\n", + " [ 132., 162.],\n", + " [ 133., 162.],\n", + " [ 134., 162.],\n", + " [ 135., 162.],\n", + " [ 136., 162.],\n", + " [ 137., 162.],\n", + " [ 138., 162.],\n", + " [ 139., 162.],\n", + " [ 140., 162.],\n", + " [ 141., 162.],\n", + " [ 142., 162.],\n", + " [ 143., 162.],\n", + " [ 144., 162.],\n", + " [ 145., 162.],\n", + " [ 146., 162.],\n", + " [ 147., 162.],\n", + " [ 148., 162.],\n", + " [ 149., 162.],\n", + " [ 150., 162.],\n", + " [ 151., 162.],\n", + " [ 152., 162.],\n", + " [ 153., 162.],\n", + " [ 154., 162.],\n", + " [ 155., 162.],\n", + " [ 156., 162.],\n", + " [ 157., 162.],\n", + " [ 158., 162.],\n", + " [ 159., 162.],\n", + " [ 160., 162.],\n", + " [ 161., 162.],\n", + " [ 162., 162.],\n", + " [ 163., 162.],\n", + " [ 164., 162.],\n", + " [ 165., 162.],\n", + " [ 166., 162.],\n", + " [ 167., 162.],\n", + " [ 168., 162.],\n", + " [ 169., 162.],\n", + " [ 170., 162.],\n", + " [ 171., 162.],\n", + " [ 172., 162.],\n", + " [ 173., 162.],\n", + " [ 174., 162.],\n", + " [ 175., 162.],\n", + " [ 176., 162.],\n", + " [ 177., 162.],\n", + " [ 178., 162.],\n", + " [ 179., 162.],\n", + " [ 180., 162.],\n", + " [ 181., 162.],\n", + " [ 182., 162.],\n", + " [ 183., 162.],\n", + " [ 184., 162.],\n", + " [ 185., 162.],\n", + " [ 186., 162.],\n", + " [ 187., 162.],\n", + " [ 188., 162.],\n", + " [ 189., 162.],\n", + " [ 190., 162.],\n", + " [ 191., 162.],\n", + " [ 192., 162.],\n", + " [ 193., 162.],\n", + " [ 194., 162.],\n", + " [ 195., 162.],\n", + " [ 196., 162.],\n", + " [ 197., 162.],\n", + " [ 198., 162.],\n", + " [ 199., 162.],\n", + " [ 200., 162.],\n", + " [ 201., 162.],\n", + " [ 202., 162.],\n", + " [ 203., 162.],\n", + " [ 204., 162.],\n", + " [ 205., 162.],\n", + " [ 206., 162.],\n", + " [ 207., 162.],\n", + " [ 208., 162.],\n", + " [ 209., 162.],\n", + " [ 210., 162.],\n", + " [ 211., 162.],\n", + " [ 212., 162.],\n", + " [ 213., 162.],\n", + " [ 214., 162.],\n", + " [ 215., 162.],\n", + " [ 216., 162.],\n", + " [ 217., 162.],\n", + " [ 218., 162.],\n", + " [ 219., 162.],\n", + " [ 220., 162.],\n", + " [ 221., 162.],\n", + " [ 222., 162.],\n", + " [ 223., 162.],\n", + " [ 224., 162.],\n", + " [ 225., 162.],\n", + " [ 226., 162.],\n", + " [ 227., 162.],\n", + " [ 228., 162.],\n", + " [ 229., 162.],\n", + " [ 230., 162.],\n", + " [ 231., 162.],\n", + " [ 232., 162.],\n", + " [ 233., 162.],\n", + " [ 234., 162.],\n", + " [ 235., 162.],\n", + " [ 236., 162.],\n", + " [ 237., 162.],\n", + " [ 238., 162.],\n", + " [ 239., 162.],\n", + " [ 240., 162.],\n", + " [ 241., 162.],\n", + " [ 242., 162.],\n", + " [ 243., 162.],\n", + " [ 244., 162.],\n", + " [ 245., 162.],\n", + " [ 246., 162.],\n", + " [ 247., 162.],\n", + " [ 248., 162.],\n", + " [ 249., 162.],\n", + " [ 250., 162.],\n", + " [ 251., 162.],\n", + " [ 252., 162.],\n", + " [ 253., 162.],\n", + " [ 254., 162.],\n", + " [ 255., 162.],\n", + " [ 256., 162.],\n", + " [ 257., 162.],\n", + " [ 258., 162.],\n", + " [ 259., 162.],\n", + " [ 260., 162.],\n", + " [ 261., 162.],\n", + " [ 262., 162.],\n", + " [ 263., 162.],\n", + " [ 264., 162.],\n", + " [ 265., 162.],\n", + " [ 266., 162.],\n", + " [ 267., 162.],\n", + " [ 268., 162.],\n", + " [ 269., 162.],\n", + " [ 270., 162.],\n", + " [ 271., 162.],\n", + " [ 272., 162.],\n", + " [ 273., 162.],\n", + " [ 274., 162.],\n", + " [ 275., 162.],\n", + " [ 276., 162.],\n", + " [ 277., 162.],\n", + " [ 278., 162.],\n", + " [ 279., 162.],\n", + " [ 280., 162.],\n", + " [ 281., 162.],\n", + " [ 282., 162.],\n", + " [ 283., 162.],\n", + " [ 284., 162.],\n", + " [ 285., 162.],\n", + " [ 286., 162.],\n", + " [ 287., 162.],\n", + " [ 288., 162.],\n", + " [ 289., 162.],\n", + " [ 290., 162.],\n", + " [ 291., 162.],\n", + " [ 292., 162.],\n", + " [ 293., 162.],\n", + " [ 294., 162.],\n", + " [ 295., 162.],\n", + " [ 296., 162.],\n", + " [ 297., 162.],\n", + " [ 298., 162.],\n", + " [ 299., 162.],\n", + " [ 300., 162.],\n", + " [ 301., 162.],\n", + " [ 302., 162.],\n", + " [ 303., 162.],\n", + " [ 304., 162.],\n", + " [ 305., 162.],\n", + " [ 306., 162.],\n", + " [ 307., 162.],\n", + " [ 308., 162.],\n", + " [ 309., 162.],\n", + " [ 310., 162.],\n", + " [ 311., 162.],\n", + " [ 312., 162.],\n", + " [ 313., 162.],\n", + " [ 314., 162.],\n", + " [ 315., 162.],\n", + " [ 316., 162.],\n", + " [ 317., 162.],\n", + " [ 318., 162.],\n", + " [ 319., 162.],\n", + " [ 320., 162.],\n", + " [ 321., 162.],\n", + " [ 322., 162.]])\n", + "Coordinates:\n", + " * x_y (x_y) \n", + "array([[[ 3167.0234375 , 3151.203125 , 3129.890625 , ...,\n", + " 1585.2734375 , 1635.03125 , 1629.5234375 ],\n", + " [ 3165.59375 , 3140.640625 , 3120.1875 , ...,\n", + " 1586.25 , 1633.4140625 , 1623.125 ],\n", + " [ 3146.46875 , 3124.3125 , 3105.953125 , ...,\n", + " 1571.5234375 , 1615.1171875 , 1610.8828125 ],\n", + " ..., \n", + " [ 47.16455078, 44.47119141, 44.27294922, ...,\n", + " 18.95068359, 19.76757812, 19.65283203],\n", + " [ 26.86279297, 24.75195312, 24.64208984, ...,\n", + " 10.53710938, 10.97021484, 10.92431641],\n", + " [ 9.92041016, 7.88720703, 7.86523438, ...,\n", + " 3.37207031, 3.5078125 , 3.48535156]],\n", + "\n", + " [[ 3167.0234375 , 3151.203125 , 3129.890625 , ...,\n", + " 1585.2734375 , 1635.03125 , 1629.5234375 ],\n", + " [ 3165.59375 , 3140.640625 , 3120.1875 , ...,\n", + " 1586.25 , 1633.4140625 , 1623.125 ],\n", + " [ 3146.46875 , 3124.3125 , 3105.953125 , ...,\n", + " 1571.5234375 , 1615.1171875 , 1610.8828125 ],\n", + " ..., \n", + " [ 47.16455078, 44.47119141, 44.27294922, ...,\n", + " 18.95068359, 19.76757812, 19.65283203],\n", + " [ 26.86279297, 24.75195312, 24.64208984, ...,\n", + " 10.53710938, 10.97021484, 10.92431641],\n", + " [ 9.92041016, 7.88720703, 7.86523438, ...,\n", + " 3.37207031, 3.5078125 , 3.48535156]]], dtype=float32)\n", + "Coordinates:\n", + " * Time (Time) datetime64[ns] 2010-06-13T21:00:00 2010-06-13T21:00:00\n", + " * bottom_top (bottom_top) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 ...\n", + " datetime (Time) datetime64[ns] 2010-06-13T21:00:00 2010-06-13T21:00:00\n", + " xy_loc (xy) object CoordPair(x=0.0, y=162.0) ...\n", + " * xy (xy) int64 0 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 ...\n", + "Attributes:\n", + " description: perturbation pressure\n", + " units: Pa\n", + " Orientation: (0.0,162.0) to (322.0,162.0)\n", + "\n", + "\n", + "\n", + "array([[ 1398.91162109, 1371.33605957, 1342.515625 , 1284.37878418,\n", + " 1155.76074219, 611.89489746],\n", + " [ 1398.91162109, 1371.33605957, 1342.515625 , 1284.37878418,\n", + " 1155.76074219, 611.89489746]], dtype=float32)\n", + "Coordinates:\n", + " * Time (Time) datetime64[ns] 2010-06-13T21:00:00 2010-06-13T21:00:00\n", + " * z (z) float32 100.0 200.0 300.0 500.0 1000.0 5000.0\n", + "Attributes:\n", + " _FillValue: 9.96920996839e+36\n", + " missing_value: 9.96920996839e+36\n", + " description: perturbation pressure\n", + " units: Pa\n", + "\n", + "\n" + ] + } + ], + "source": [ + "import numpy as np\n", + "from wrf import (ALL_TIMES, npvalues, Constants, getvar, extract_vars, destagger, \n", + " interp1d, interp2dxy, interpz3d, \n", + " slp, tk, td, rh, uvmet, smooth2d, extract_global_attrs, xy)\n", + "from math import fabs, log, tan, sin, cos\n", + "from wrf.util import either\n", + "from netCDF4 import Dataset as nc\n", + "\n", + "timeidx = None\n", + "wrfnc = nc(\"/Users/ladwig/Documents/wrf_files/wrfout_d01_2010-06-13_21:00:00\")\n", + "\n", + "wrfnc = [wrfnc, wrfnc]\n", + "field3d = getvar(wrfnc, \"P\", timeidx, meta=True)\n", + "\n", + "z = getvar(wrfnc, \"z\", timeidx)\n", + "\n", + "interp = interpz3d(field3d, z, 500, missingval=Constants.DEFAULT_FILL, meta=True)\n", + "print (interp)\n", + "print (\"\\n\")\n", + "\n", + "# 2dxy test\n", + "pivot_point = (field3d.shape[-1] / 2, field3d.shape[-2] / 2 ) \n", + "angle = 90.0\n", + "\n", + "start_point = (0, z.shape[-2]/2)\n", + "end_point = (-1, z.shape[-2]/2)\n", + "\n", + "#start_point = (0, 0)\n", + "#end_point = (-1, -1)\n", + "\n", + "_xy = xy(field3d, pivot_point=pivot_point, angle=angle)\n", + "print(_xy)\n", + "print (\"\\n\")\n", + "\n", + "_xy = xy(field3d, start_point=start_point, end_point=end_point)\n", + "print(_xy)\n", + "print (\"\\n\")\n", + "\n", + "interpxy = interp2dxy(field3d, _xy)\n", + "print (interpxy)\n", + "print (\"\\n\")\n", + "\n", + "# 1D test\n", + "only_height = [0]*field3d.ndim\n", + "only_height[-3] = slice(None)\n", + "only_height = (Ellipsis,) + tuple(only_height[-3:])\n", + "\n", + "v_in = field3d[only_height]\n", + "z_in = z[only_height]\n", + "z_out = np.asarray([100.,200.,300.,500.,1000.,5000.], field3d.dtype)\n", + "\n", + "int1d = interp1d(v_in, z_in, z_out, missingval=Constants.DEFAULT_FILL, meta=True)\n", + "print(int1d)\n", + "print (\"\\n\")\n" ] }, { diff --git a/test/projtest.py b/test/projtest.py index 2c4c984..2e0b38c 100644 --- a/test/projtest.py +++ b/test/projtest.py @@ -27,15 +27,18 @@ except ImportError: CARTOPY = False -from wrf.var import (get_proj_params, getproj, RotLatLonProj, - PolarStereographicProj) +from wrf import get_proj_params +from wrf.projection import getproj, RotatedLatLon, PolarStereographic FILE_DIR = "/Users/ladwig/Documents/wrf_files/" WRF_FILES = [ + join(FILE_DIR, "rotated_pole", "EAS_geo_em.d01.nc"), + join(FILE_DIR, "rotated_pole", "EUR_geo_em.d01.nc"), join(FILE_DIR,"wrfout_d01_2016-02-25_18_00_00"), join(FILE_DIR, "wrfout_d01_2008-09-29_23-30-00"), join(FILE_DIR, "wrfout_d01_2010-06-13_21:00:00")] + def nz_proj(): lats = np.array([[-47.824014, -47.824014], [-32.669853, -32.669853]]) @@ -142,7 +145,7 @@ def make_test(wrf_file=None, fixed_case=None): elif fixed_case == "dateline_rot": lats,lons,proj = dateline_rot_proj() - print "wrf proj4: {}".format(proj.proj4()) + print ("wrf proj4: {}".format(proj.proj4())) if PYNGL: # PyNGL plotting wks_type = "png" @@ -172,7 +175,7 @@ def make_test(wrf_file=None, fixed_case=None): bm.drawcoastlines(linewidth=.5) #bm.drawparallels(parallels,labels=[1,1,1,1],fontsize=10) #bm.drawmeridians(meridians,labels=[1,1,1,1],fontsize=10) - print "basemap proj4: {}".format(bm.proj4string) + print ("basemap proj4: {}".format(bm.proj4string)) plt.savefig("basemap_{}.png".format(name_suffix)) plt.close(fig) @@ -180,7 +183,7 @@ def make_test(wrf_file=None, fixed_case=None): # Cartopy plotting fig = plt.figure(figsize=(10,10)) ax = plt.axes([0.1,0.1,0.8,0.8], projection=proj.cartopy()) - print "cartopy proj4: {}".format(proj.cartopy().proj4_params) + print ("cartopy proj4: {}".format(proj.cartopy().proj4_params)) ax.coastlines('50m', linewidth=0.8) #print proj.x_extents() diff --git a/test/utests.py b/test/utests.py index 0ff8124..afe9698 100644 --- a/test/utests.py +++ b/test/utests.py @@ -47,6 +47,7 @@ def setUpModule(): # http://eli.thegreenplace.net/2014/04/02/dynamically-generating-python-test-cases def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): def test(self): + try: from netCDF4 import Dataset as NetCDF except: @@ -88,7 +89,10 @@ def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): ref_vals = refnc.variables[varname][:] else: data = refnc.variables[varname][:] - new_dims = [repeat] + [x for x in data.shape] + if varname != "uvmet" and varname != "uvmet10": + new_dims = [repeat] + [x for x in data.shape] + else: + new_dims = [2] + [repeat] + [x for x in data.shape[1:]] masked=False if (isinstance(data, ma.core.MaskedArray)): masked=True @@ -99,10 +103,20 @@ def make_test(varname, wrf_in, referent, multi=False, repeat=3, pynio=False): ref_vals = ma.asarray(n.zeros(new_dims, data.dtype)) for i in xrange(repeat): - ref_vals[i,:] = data[:] + if varname != "uvmet" and varname != "uvmet10": + ref_vals[i,:] = data[:] - if masked: - ref_vals.mask[i,:] = data.mask[:] + if masked: + ref_vals.mask[i,:] = data.mask[:] + else: + ref_vals[0, i, :] = data[0,:] + ref_vals[1, i, :] = data[1,:] + + if masked: + ref_vals.mask[0,i,:] = data.mask[:] + ref_vals.mask[1,i,:] = data.mask[:] + + if (varname == "tc"): my_vals = getvar(in_wrfnc, "temp", timeidx=timeidx, units="c") @@ -217,7 +231,7 @@ def make_interp_test(varname, wrf_in, referent, multi=False, hts = getvar(in_wrfnc, "z", timeidx=timeidx) p = getvar(in_wrfnc, "pressure", timeidx=timeidx) - pivot_point = (hts.shape[-2] / 2, hts.shape[-1] / 2) + pivot_point = (hts.shape[-1] / 2, hts.shape[-2] / 2) ht_cross = vertcross(hts, p, pivot_point=pivot_point, angle=90.) nt.assert_allclose(npvalues(ht_cross), ref_ht_cross, rtol=.01) @@ -229,8 +243,8 @@ def make_interp_test(varname, wrf_in, referent, multi=False, ref_p_cross, rtol=.01) # Test point to point - start_point = (hts.shape[-2]/2, 0) - end_point = (hts.shape[-2]/2, -1) + start_point = (0,hts.shape[-2]/2) + end_point = (-1,hts.shape[-2]/2) p_cross2 = vertcross(p,hts,start_point=start_point, end_point=end_point) @@ -243,15 +257,15 @@ def make_interp_test(varname, wrf_in, referent, multi=False, ref_t2_line = _get_refvals(referent, "t2_line", repeat, multi) t2 = getvar(in_wrfnc, "T2", timeidx=timeidx) - pivot_point = (t2.shape[-2] / 2, t2.shape[-1] / 2) + pivot_point = (t2.shape[-1] / 2, t2.shape[-2] / 2) t2_line1 = interpline(t2, pivot_point=pivot_point, angle=90.0) nt.assert_allclose(npvalues(t2_line1), ref_t2_line) # Test point to point - start_point = (t2.shape[-2]/2, 0) - end_point = (t2.shape[-2]/2, -1) + start_point = (0, t2.shape[-2]/2) + end_point = (-1, t2.shape[-2]/2) t2_line2 = interpline(t2, start_point=start_point, end_point=end_point)