From 6239fc16fd9312e6c37f4ad3b5363801f12c8d63 Mon Sep 17 00:00:00 2001 From: Bill Ladwig Date: Fri, 8 Dec 2017 16:29:03 -0700 Subject: [PATCH] Fixed numerous issues with reading scalars and non-grid variables from the NetCDF file using getvar. Adds unit tests for reading file variables. Fixes #37 --- fortran/wrf_constants.f90 | 10 +++- src/wrf/api.py | 6 +- src/wrf/cape.py | 10 ++-- src/wrf/cloudfrac.py | 8 ++- src/wrf/computation.py | 20 +++---- src/wrf/constants.py | 39 +++++++++++++ src/wrf/decorators.py | 8 +-- src/wrf/extension.py | 8 +-- src/wrf/interp.py | 12 ++-- src/wrf/specialdec.py | 8 +-- src/wrf/util.py | 116 ++++++++++++++++++++++++++++---------- test/test_filevars.py | 51 +++++++++++++++++ 12 files changed, 227 insertions(+), 69 deletions(-) create mode 100644 test/test_filevars.py diff --git a/fortran/wrf_constants.f90 b/fortran/wrf_constants.f90 index 617e63f..336c0e3 100644 --- a/fortran/wrf_constants.f90 +++ b/fortran/wrf_constants.f90 @@ -9,7 +9,15 @@ MODULE wrf_constants REAL(KIND=8), PARAMETER :: PI = 3.1415926535897932384626433D0 REAL(KIND=8), PARAMETER :: RAD_PER_DEG = PI/180.D0 REAL(KIND=8), PARAMETER :: DEG_PER_RAD = 180.D0/PI - REAL(KIND=8), PARAMETER :: DEFAULT_FILL = 9.9692099683868690D36 + REAL(KIND=8), PARAMETER :: DEFAULT_FILL = 9.9692099683868690E36 + INTEGER(KIND=1), PARAMETER :: DEFAULT_FILL_INT8 = -127 + INTEGER(KIND=2), PARAMETER :: DEFAULT_FILL_INT16 = -32767 + INTEGER(KIND=4), PARAMETER :: DEFAULT_FILL_INT32 = -2147483647 + INTEGER(KIND=8), PARAMETER :: DEFAULT_FILL_INT64 = INT(-9223372036854775806D0, KIND=8) + REAL(KIND=4), PARAMETER :: DEFAULT_FILL_FLOAT = 9.9692099683868690E36 + REAL(KIND=8), PARAMETER :: DEFAULT_FILL_DOUBLE = 9.9692099683868690E36 + CHARACTER(LEN=1), PARAMETER :: DEFAULT_FILL_CHAR = ACHAR(0) + REAL(KIND=8), PARAMETER :: P1000MB = 100000.D0 ! j/k/kg diff --git a/src/wrf/api.py b/src/wrf/api.py index 877096d..9333b2d 100644 --- a/src/wrf/api.py +++ b/src/wrf/api.py @@ -3,7 +3,8 @@ from .config import (xarray_enabled, disable_xarray, enable_xarray, basemap_enabled, disable_basemap, enable_basemap, pyngl_enabled, enable_pyngl, disable_pyngl, set_cache_size, get_cache_size) -from .constants import ALL_TIMES, Constants, ConversionFactors, ProjectionTypes +from .constants import (ALL_TIMES, Constants, ConversionFactors, + ProjectionTypes, default_fill) from .destag import destagger from .routines import getvar from .computation import (xy, interp1d, interp2dxy, interpz3d, slp, tk, td, rh, @@ -58,7 +59,8 @@ __all__ += ["xarray_enabled", "disable_xarray", "enable_xarray", "basemap_enabled", "disable_basemap", "enable_basemap", "pyngl_enabled", "enable_pyngl", "disable_pyngl", "set_cache_size", "get_cache_size"] -__all__ += ["ALL_TIMES", "Constants", "ConversionFactors", "ProjectionTypes"] +__all__ += ["ALL_TIMES", "Constants", "ConversionFactors", "ProjectionTypes", + "default_fill"] __all__ += ["destagger"] __all__ += ["getvar"] __all__ += ["xy", "interp1d", "interp2dxy", "interpz3d", "slp", "tk", "td", diff --git a/src/wrf/cape.py b/src/wrf/cape.py index 7cb7efd..5f5ef10 100755 --- a/src/wrf/cape.py +++ b/src/wrf/cape.py @@ -6,13 +6,13 @@ import numpy.ma as ma from .extension import _tk, _cape from .destag import destagger -from .constants import Constants, ConversionFactors +from .constants import default_fill, Constants, ConversionFactors from .util import extract_vars from .metadecorators import set_cape_metadata @set_cape_metadata(is2d=True) def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, - meta=True, _key=None, missing=Constants.DEFAULT_FILL): + meta=True, _key=None, missing=default_fill(np.float64)): """Return the 2d fields of CAPE, CIN, LCL, and LFC. The leftmost dimension of the returned array represents four different @@ -66,7 +66,7 @@ def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, purposes only. Default is None. missing (:obj:`float`): The fill value to use for the output. - Default is :data:`wrf.Constants.DEFAULT_FILL`. + Default is :data:`wrf.default_fill(np.float64)`. Returns: :class:`xarray.DataArray` or :class:`numpy.ndarray`: The @@ -130,7 +130,7 @@ def get_2dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, @set_cape_metadata(is2d=False) def get_3dcape(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, - _key=None, missing=Constants.DEFAULT_FILL): + _key=None, missing=default_fill(np.float64)): """Return the three-dimensional CAPE and CIN. The leftmost dimension of the returned array represents two different @@ -182,7 +182,7 @@ def get_3dcape(wrfin, timeidx=0, method="cat", purposes only. Default is None. missing (:obj:`float`): The fill value to use for the output. - Default is :data:`wrf.Constants.DEFAULT_FILL`. + Default is :data:`wrf.default_fill(np.float64)`. Returns: :class:`xarray.DataArray` or :class:`numpy.ndarray`: The diff --git a/src/wrf/cloudfrac.py b/src/wrf/cloudfrac.py index 38805c6..f093146 100644 --- a/src/wrf/cloudfrac.py +++ b/src/wrf/cloudfrac.py @@ -1,19 +1,21 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) -from .constants import Constants +import numpy as np +import numpy.ma as ma + +from .constants import Constants, default_fill from .extension import _tk, _rh, _cloudfrac from .metadecorators import set_cloudfrac_metadata from .util import extract_vars from .geoht import _get_geoht -import numpy.ma as ma @set_cloudfrac_metadata() def get_cloudfrac(wrfin, timeidx=0, method="cat", squeeze=True, cache=None, meta=True, _key=None, vert_type="pres", low_thresh=None, mid_thresh=None, - high_thresh=None, missing=Constants.DEFAULT_FILL): + high_thresh=None, missing=default_fill(np.float64)): """Return the cloud fraction for low, mid, and high level clouds. The leftmost dimension of the returned array represents three different diff --git a/src/wrf/computation.py b/src/wrf/computation.py index 1a0e1c3..bce2c75 100644 --- a/src/wrf/computation.py +++ b/src/wrf/computation.py @@ -4,7 +4,7 @@ from __future__ import (absolute_import, division, print_function, import numpy as np import numpy.ma as ma -from .constants import Constants +from .constants import default_fill from .extension import (_interpz3d, _interp2dxy, _interp1d, _slp, _tk, _td, _rh, _uvmet, _smooth2d, _cape, _cloudfrac, _ctt, _dbz, _srhel, _udhel, _avo, _pvo, _eth, _wetbulb, _tv, @@ -104,7 +104,7 @@ def xy(field, pivot_point=None, angle=None, start_point=None, end_point=None, @set_interp_metadata("1d") -def interp1d(field, z_in, z_out, missing=Constants.DEFAULT_FILL, +def interp1d(field, z_in, z_out, missing=default_fill(np.float64), meta=True): """Return the linear interpolation of a one-dimensional variable. @@ -128,7 +128,7 @@ def interp1d(field, z_in, z_out, missing=Constants.DEFAULT_FILL, to. Must be the same type as *z_in*. missing (:obj:`float`, optional): The fill value to use for the - output. Default is :data:`wrf.Constants.DEFAULT_FILL`. + output. Default is :data:`wrf.default_fill(np.float64)`. meta (:obj:`bool`, optional): Set to False to disable metadata and return :class:`numpy.ndarray` instead of @@ -251,7 +251,7 @@ def interp2dxy(field3d, xy, meta=True): @set_interp_metadata("horiz") -def interpz3d(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, +def interpz3d(field3d, vert, desiredlev, missing=default_fill(np.float64), meta=True): """Return the field interpolated to a specified pressure or height level. @@ -280,7 +280,7 @@ def interpz3d(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, Must be in the same units as the *vert* parameter. missing (:obj:`float`): The fill value to use for the output. - Default is :data:`wrf.Constants.DEFAULT_FILL`. + Default is :data:`wrf.default_fill(numpy.float64)`. meta (:obj:`bool`): Set to False to disable metadata and return :class:`numpy.ndarray` instead of @@ -713,7 +713,7 @@ def smooth2d(field, passes, meta=True): @set_cape_alg_metadata(is2d=True, copyarg="pres_hpa") def cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, - missing=Constants.DEFAULT_FILL, meta=True): + missing=default_fill(np.float64), meta=True): """Return the two-dimensional CAPE, CIN, LCL, and LFC. This function calculates the maximum convective available potential @@ -790,7 +790,7 @@ def cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, False for pressure level data. missing (:obj:`float`, optional): The fill value to use for the - output. Default is :data:`wrf.Constants.DEFAULT_FILL`. + output. Default is :data:`wrf.default_fill(numpy.float64)`. meta (:obj:`bool`): Set to False to disable metadata and return :class:`numpy.ndarray` instead of @@ -843,7 +843,7 @@ def cape_2d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, @set_cape_alg_metadata(is2d=False, copyarg="pres_hpa") def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, - missing=Constants.DEFAULT_FILL, meta=True): + missing=default_fill(np.float64), meta=True): """Return the three-dimensional CAPE and CIN. This function calculates the maximum convective available potential @@ -926,7 +926,7 @@ def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, False for pressure level data. missing (:obj:`float`, optional): The fill value to use for the - output. Default is :data:`wrf.Constants.DEFAULT_FILL`. + output. Default is :data:`wrf.default_fill(numpy.float64)`. meta (:obj:`bool`): Set to False to disable metadata and return :class:`numpy.ndarray` instead of @@ -964,7 +964,7 @@ def cape_3d(pres_hpa, tkel, qv, height, terrain, psfc_hpa, ter_follow, @set_cloudfrac_alg_metadata(copyarg="vert") def cloudfrac(vert, relh, vert_inc_w_height, low_thresh, mid_thresh, - high_thresh, missing=Constants.DEFAULT_FILL, meta=True): + high_thresh, missing=default_fill(np.float64), meta=True): """Return the cloud fraction. The leftmost dimension of the returned array represents three different diff --git a/src/wrf/constants.py b/src/wrf/constants.py index 8c62d37..96dcc02 100755 --- a/src/wrf/constants.py +++ b/src/wrf/constants.py @@ -1,6 +1,8 @@ from __future__ import (absolute_import, division, print_function, unicode_literals) +from sys import version_info +import struct import numpy as np from .py3compat import viewitems @@ -42,5 +44,42 @@ class ProjectionTypes(object): POLAR_STEREOGRAPHIC = 2 MERCATOR = 3 LAT_LON = 6 + +# Create the default fill mapping based on type. +_DEFAULT_FILL_MAP = {None: Constants.DEFAULT_FILL, + np.dtype(np.bool_) : False, + np.dtype(np.intc) : Constants.DEFAULT_FILL_INT32, # Usually true + np.dtype(np.int8) : Constants.DEFAULT_FILL_INT8, + np.dtype(np.uint8) : 255, + np.dtype(np.int16) : Constants.DEFAULT_FILL_INT16, + np.dtype(np.uint16) : 65535, + np.dtype(np.int32) : Constants.DEFAULT_FILL_INT32, + np.dtype(np.uint32) : 4294967295, + np.dtype(np.int64) : Constants.DEFAULT_FILL_INT64, + np.dtype(np.uint64) : 18446744073709551614, + np.dtype(np.float_) : Constants.DEFAULT_FILL_DOUBLE, + np.dtype(np.float32) : Constants.DEFAULT_FILL_FLOAT, + np.dtype(np.float64) : Constants.DEFAULT_FILL_DOUBLE + } + +if version_info >= (3, ): + _DEFAULT_FILL_MAP[np.int_] = Constants.DEFAULT_FILL_INT64 +else: + _DEFAULT_FILL_MAP[np.int_] = Constants.DEFAULT_FILL_INT32 + +if (struct.calcsize("P") == 8): + _DEFAULT_FILL_MAP[np.intp] = Constants.DEFAULT_FILL_INT64 +else: + _DEFAULT_FILL_MAP[np.intp] = Constants.DEFAULT_FILL_INT32 + + +# Add the integers based on python 2.x or 3.x +def default_fill(dtype=None): + dt = np.dtype(dtype) if dtype is not None else None + return _DEFAULT_FILL_MAP.get(dt, 0) + + + + \ No newline at end of file diff --git a/src/wrf/decorators.py b/src/wrf/decorators.py index 406dcec..7b1797f 100644 --- a/src/wrf/decorators.py +++ b/src/wrf/decorators.py @@ -11,7 +11,7 @@ from .units import do_conversion, check_units, dealias_and_clean_unit from .util import iter_left_indexes, from_args, to_np, combine_dims from .py3compat import viewitems, viewvalues, isstr from .config import xarray_enabled -from .constants import Constants +from .constants import default_fill if xarray_enabled(): from xarray import DataArray @@ -201,7 +201,7 @@ def left_iteration(ref_var_expected_dims, if all_masked: for output in viewvalues(outd): output[left_and_slice_idxs] = ( - Constants.DEFAULT_FILL) + default_fill(np.float64)) skip_missing = True mask_output = True break @@ -240,9 +240,9 @@ def left_iteration(ref_var_expected_dims, # Mostly when used with join if mask_output: if isinstance(output, np.ndarray): - output = ma.masked_values(output, Constants.DEFAULT_FILL) + output = ma.masked_values(output, default_fill(np.float64)) else: - output = tuple(ma.masked_values(arr, Constants.DEFAULT_FILL) + output = tuple(ma.masked_values(arr, default_fill(np.float64)) for arr in output) return output diff --git a/src/wrf/extension.py b/src/wrf/extension.py index cd4acd5..a9c31a2 100755 --- a/src/wrf/extension.py +++ b/src/wrf/extension.py @@ -3,7 +3,7 @@ from __future__ import (absolute_import, division, print_function, import numpy as np -from .constants import Constants +from .constants import Constants, default_fill from ._wrffortran import (dcomputetk, dinterp3dz, dinterp2dxy, dinterp1d, dcomputeseaprs, dfilter2d, dcomputerh, dcomputeuvmet, @@ -381,8 +381,8 @@ def _eth(qv, tk, p, outview=None): @cast_type(arg_idxs=(0,1,2,3)) @extract_and_transpose() def _uvmet(u, v, lat, lon, cen_long, cone, isstag=0, has_missing=False, - umissing=Constants.DEFAULT_FILL, vmissing=Constants.DEFAULT_FILL, - uvmetmissing=Constants.DEFAULT_FILL, outview=None): + umissing=default_fill(np.float64), vmissing=default_fill(np.float64), + uvmetmissing=default_fill(np.float64), outview=None): """Wrapper for dcomputeuvmet. Located in wrf_user.f90. @@ -795,7 +795,7 @@ def _smooth2d(field, passes, outview=None): if isinstance(field, np.ma.MaskedArray): missing = field.fill_value else: - missing = Constants.DEFAULT_FILL + missing = default_fill(np.float64) if outview is None: outview = field.copy(order="A") diff --git a/src/wrf/interp.py b/src/wrf/interp.py index 8741596..bf9e34b 100755 --- a/src/wrf/interp.py +++ b/src/wrf/interp.py @@ -11,7 +11,7 @@ from .metadecorators import set_interp_metadata from .util import extract_vars, is_staggered, get_id, to_np from .py3compat import py3range from .interputils import get_xy, get_xy_z_params, to_xy_coords -from .constants import Constants, ConversionFactors +from .constants import Constants, default_fill, ConversionFactors from .terrain import get_terrain from .geoht import get_height from .temp import get_theta, get_temp, get_eth @@ -20,7 +20,7 @@ from .pressure import get_pressure # Note: Extension decorator is good enough to handle left dims @set_interp_metadata("horiz") -def interplevel(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, +def interplevel(field3d, vert, desiredlev, missing=default_fill(np.float64), meta=True): """Return the three-dimensional field interpolated to a horizontal plane at the specified vertical level. @@ -40,7 +40,7 @@ def interplevel(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, Must be in the same units as the *vert* parameter. missing (:obj:`float`): The fill value to use for the output. - Default is :data:`wrf.Constants.DEFAULT_FILL`. + Default is :data:`wrf.default_fill(numpy.float64)`. meta (:obj:`bool`): Set to False to disable metadata and return :class:`numpy.ndarray` instead of @@ -90,7 +90,7 @@ def interplevel(field3d, vert, desiredlev, missing=Constants.DEFAULT_FILL, @set_interp_metadata("cross") -def vertcross(field3d, vert, levels=None, missing=Constants.DEFAULT_FILL, +def vertcross(field3d, vert, levels=None, missing=default_fill(np.float64), wrfin=None, timeidx=0, stagger=None, projection=None, pivot_point=None, angle=None, start_point=None, end_point=None, @@ -133,7 +133,7 @@ def vertcross(field3d, vert, levels=None, missing=Constants.DEFAULT_FILL, Default is None. missing (:obj:`float`): The fill value to use for the output. - Default is :data:`wrf.Constants.DEFAULT_FILL`. + Default is :data:`wrf.default_fill(numpy.float64)`. wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ iterable, optional): WRF-ARW NetCDF @@ -653,7 +653,7 @@ def vinterp(wrfin, field, vert_coord, interp_levels, extrapolate=False, if isinstance(field, ma.MaskedArray): missing = field.fill_value else: - missing = Constants.DEFAULT_FILL + missing = default_fill(np.float64) if (field.shape != p.shape): raise ValueError("'field' shape does not match other variable shapes. " diff --git a/src/wrf/specialdec.py b/src/wrf/specialdec.py index 8d5efc7..f72899b 100644 --- a/src/wrf/specialdec.py +++ b/src/wrf/specialdec.py @@ -7,7 +7,7 @@ import wrapt from .util import iter_left_indexes, to_np from .config import xarray_enabled -from .constants import Constants +from .constants import default_fill if xarray_enabled(): from xarray import DataArray @@ -74,12 +74,12 @@ def uvmet_left_iter(alg_dtype=np.float64): v_arr = to_np(v) - umissing = Constants.DEFAULT_FILL + umissing = default_fill(np.float64) if isinstance(u_arr, np.ma.MaskedArray): has_missing = True umissing = u_arr.fill_value - vmissing = Constants.DEFAULT_FILL + vmissing = default_fill(np.float64) if isinstance(v_arr, np.ma.MaskedArray): has_missing = True vmissing = v_arr.fill_value @@ -452,7 +452,7 @@ def cloudfrac_left_iter(alg_dtype=np.float64): output = np.empty(output_dims, orig_dtype) has_missing = False - missing = Constants.DEFAULT_FILL + missing = default_fill(np.float64) for left_idxs in iter_left_indexes(extra_dims): left_and_slice_idxs = left_idxs + (slice(None),) low_idxs = left_idxs + (0, slice(None)) diff --git a/src/wrf/util.py b/src/wrf/util.py index fff88ec..cba14aa 100644 --- a/src/wrf/util.py +++ b/src/wrf/util.py @@ -30,7 +30,7 @@ import numpy as np import numpy.ma as ma from .config import xarray_enabled -from .constants import Constants, ALL_TIMES +from .constants import default_fill, ALL_TIMES from .py3compat import (viewitems, viewkeys, isstr, py3range, ucode) from .cache import cache_item, get_cached_item from .geobnds import GeoBounds, NullGeoBounds @@ -39,7 +39,6 @@ from .projection import getproj if xarray_enabled(): from xarray import DataArray - from pandas import NaT _COORD_PAIR_MAP = {"XLAT" : ("XLAT", "XLONG"), @@ -1097,6 +1096,7 @@ def _find_max_time_size(wrfseq): return max_times + def _get_coord_names(wrfin, varname): # Need only the first real file @@ -1219,16 +1219,29 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile, multitime = is_multi_time_req(timeidx) time_idx_or_slice = timeidx if not multitime else slice(None) var = wrfnc.variables[varname] - data = var[time_idx_or_slice, :] + if len(var.shape) > 1: + data = var[time_idx_or_slice, :] + else: + data = var[time_idx_or_slice] # Want to preserve the time dimension if not multitime: - data = data[np.newaxis, :] + if len(var.shape) > 1: + data = data[np.newaxis, :] + else: + data = data[np.newaxis] attrs = OrderedDict(var.__dict__) dimnames = var.dimensions[-data.ndim:] - lat_coord, lon_coord, time_coord = _get_coord_names(wrfnc, varname) + lat_coord = lon_coord = time_coord = None + + try: + if dimnames[-2] == "south_north" and dimnames[-1] == "west_east": + lat_coord, lon_coord, time_coord = _get_coord_names(wrfnc, varname) + except IndexError: + pass + coords = OrderedDict() @@ -1307,10 +1320,9 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile, coords[time_coord] = (lon_coord_dims[0], [time_coord_vals[timeidx]]) - proj_params = get_proj_params(wrfnc) - proj = getproj(**proj_params) - attrs["projection"] = proj - + proj_params = get_proj_params(wrfnc) + proj = getproj(**proj_params) + attrs["projection"] = proj if dimnames[0] == "Time": t = extract_times(wrfnc, timeidx, meta=False, do_xtime=False) @@ -1374,8 +1386,13 @@ def _find_forward(wrfseq, varname, timeidx, is_moving, meta, _key): return _build_data_array(wrfnc, varname, filetimeidx, is_moving, True, _key) else: - result = wrfnc.variables[varname][filetimeidx, :] - return result[np.newaxis, :] # So that nosqueeze works + var = wrfnc.variables[varname] + if len(var.shape) > 1: + result = var[filetimeidx, :] + return result[np.newaxis, :] # So that nosqueeze works + else: + result = var[filetimeidx] + return result[np.newaxis] # So that nosqueeze works else: comboidx += numtimes @@ -1574,7 +1591,10 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): startidx = 0 endidx = numtimes - outdata[startidx:endidx, :] = first_var[:] + if first_var.ndim > 1: + outdata[startidx:endidx, :] = first_var[:] + else: + outdata[startidx:endidx] = first_var[:] if xarray_enabled() and meta: latname, lonname, timename = _find_coord_names(first_var.coords) @@ -1633,7 +1653,10 @@ def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): endidx = startidx + numtimes - outdata[startidx:endidx, :] = vardata[:] + if vardata.ndim > 1: + outdata[startidx:endidx, :] = vardata[:] + else: + outdata[startidx:endidx] = vardata[:] if xarray_enabled() and meta: if timename is not None and not timecached: @@ -1795,7 +1818,8 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): if xarray_enabled() and meta: first_var = _build_data_array(wrfnc, varname, ALL_TIMES, is_moving, True, _key) - time_coord = np.full((numfiles, maxtimes), int(NaT), "datetime64[ns]") + time_coord = np.full((numfiles, maxtimes), np.datetime64("NaT"), + "datetime64[ns]") time_coord[file_idx, 0:numtimes] = first_var.coords["Time"][:] else: first_var = wrfnc.variables[varname][:] @@ -1810,8 +1834,11 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): outdims += first_var.shape[1:] # For join, always need to start with full masked values - outdata = np.full(outdims, Constants.DEFAULT_FILL, first_var.dtype) - outdata[file_idx, 0:numtimes, :] = first_var[:] + outdata = np.full(outdims, default_fill(first_var.dtype), first_var.dtype) + if first_var.ndim > 1: + outdata[file_idx, 0:numtimes, :] = first_var[:] + else: + outdata[file_idx, 0:numtimes] = first_var[:] # Create the secondary coordinate arrays if xarray_enabled() and meta: @@ -1833,8 +1860,9 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): if timename is not None: outxtimes = get_cached_item(_key, timekey) if outxtimes is None: - outxtimes = np.full(outdims[0:2], Constants.DEFAULT_FILL, - first_var.dtype) + outxtimes = np.full(outdims[0:2], + default_fill(first_var.dtype), + first_var.dtype) outxtimes[file_idx, 0:numtimes] = first_var.coords[timename][:] else: timecached = True @@ -1843,8 +1871,9 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): if latname is not None: outlats = get_cached_item(_key, latkey) if outlats is None: - outlats = np.full(outcoorddims, Constants.DEFAULT_FILL, - first_var.dtype) + outlats = np.full(outcoorddims, + default_fill(first_var.dtype), + first_var.dtype) outlats[file_idx, 0:numtimes, :] = ( first_var.coords[latname][:]) else: @@ -1853,8 +1882,9 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): if lonname is not None: outlons = get_cached_item(_key, lonkey) if outlons is None: - outlons = np.full(outcoorddims, Constants.DEFAULT_FILL, - first_var.dtype) + outlons = np.full(outcoorddims, + default_fill(first_var.dtype), + first_var.dtype) outlons[file_idx, 0:numtimes, :] = ( first_var.coords[lonname][:]) else: @@ -1874,8 +1904,11 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): if not multitime: outvar = outvar[np.newaxis, :] - - outdata[file_idx, 0:numtimes, :] = outvar[:] + + if outvar.ndim > 1: + outdata[file_idx, 0:numtimes, :] = outvar[:] + else: + outdata[file_idx, 0:numtimes] = outvar[:] if xarray_enabled() and meta: # For join, the times are a function of fileidx @@ -1904,7 +1937,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): # then a mask array is needed to flag all the missing arrays with # missing values if file_times_less_than_max: - outdata = np.ma.masked_values(outdata, Constants.DEFAULT_FILL) + outdata = np.ma.masked_values(outdata, default_fill(outdata.dtype)) if xarray_enabled() and meta: # Cache the coords if applicable @@ -1931,8 +1964,8 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): outcoords["datetime"] = outdimnames[0:2], time_coord if isinstance(outdata, np.ma.MaskedArray): - outattrs["_FillValue"] = Constants.DEFAULT_FILL - outattrs["missing_value"] = Constants.DEFAULT_FILL + outattrs["_FillValue"] = default_fill(outdata.dtype) + outattrs["missing_value"] = default_fill(outdata.dtype) if timename is not None: outxtimes = outxtimes[:, time_idx_or_slice] @@ -2378,15 +2411,35 @@ def extract_times(wrfin, timeidx, method="cat", squeeze=True, cache=None, else: wrf_list = wrfin + dt = "datetime64[ns]" if not do_xtime else np.float64 + fill_value = (np.datetime64('NaT') if not do_xtime else + default_fill(np.float64)) + try: if method.lower() == "cat": time_list = [file_time for wrf_file in wrf_list for file_time in _file_times(wrf_file, do_xtime)] + time_arr = np.asarray(time_list, dtype=dt) + elif method.lower() == "join": time_list = [[file_time for file_time in _file_times(wrf_file, do_xtime)] - for wrf_file in wrf_list] + for wrf_file in wrf_list] + + num_rows = len(time_list) + num_cols = len(time_list[0]) + + time_arr = np.full((num_rows, num_cols), fill_value, dtype=dt) + for i,row in enumerate(time_list): + if len(row) == num_cols: + time_arr[i,:] = row[:] + else: + for j,val in enumerate(row): + time_arr[i,j] = val + + time_arr = ma.masked_values(time_arr, fill_value) + else: raise ValueError("invalid method argument '{}'".format(method)) except KeyError: @@ -2400,6 +2453,8 @@ def extract_times(wrfin, timeidx, method="cat", squeeze=True, cache=None, outdimnames = ["Time"] else: outdimnames = ["fileidx", "Time"] + outattrs["missing_value"] = fill_value + outattrs["_FillValue"] = fill_value if not do_xtime: outname = "times" @@ -2412,11 +2467,12 @@ def extract_times(wrfin, timeidx, method="cat", squeeze=True, cache=None, outname = "XTIME" - outarr = DataArray(time_list, name=outname, coords=outcoords, + + outarr = DataArray(time_arr, name=outname, coords=outcoords, dims=outdimnames, attrs=outattrs) else: - outarr = np.asarray(time_list, dtype="datetime64[ns]") + outarr = time_arr if not multitime: return outarr[timeidx] diff --git a/test/test_filevars.py b/test/test_filevars.py new file mode 100644 index 0000000..29381f9 --- /dev/null +++ b/test/test_filevars.py @@ -0,0 +1,51 @@ +import unittest as ut +import numpy.testing as nt +import numpy as np +import numpy.ma as ma +import os, sys +import subprocess + +from wrf import getvar, ALL_TIMES + +TEST_DIR = "/Users/ladwig/Documents/wrf_files/wrf_vortex_multi" +TEST_FILENAMES = ["wrfout_d02_2005-08-28_00:00:00", + "wrfout_d02_2005-08-28_12:00:00", + "wrfout_d02_2005-08-29_00:00:00"] +TEST_FILES = [os.path.join(TEST_DIR, x) for x in TEST_FILENAMES] + +# Python 3 +if sys.version_info > (3,): + xrange = range + + +class WRFFileVarsTest(ut.TestCase): + longMessage = True + +def make_test(ncfiles, varname): + def test(self): + t1 = getvar(ncfiles, varname, 0) + t2 = getvar(ncfiles, varname, 0, meta=False) + t3 = getvar(ncfiles, varname, ALL_TIMES) + t4 = getvar(ncfiles, varname, ALL_TIMES, meta=False) + t5 = getvar(ncfiles, varname, ALL_TIMES, method="join") + + return test + + +if __name__ == "__main__": + from netCDF4 import Dataset + + ncfiles = [Dataset(x) for x in TEST_FILES] + + file_vars = ncfiles[0].variables.keys() + + ignore_vars = [] + + for var in file_vars: + if var in ignore_vars: + continue + + test_func1 = make_test(ncfiles, var) + setattr(WRFFileVarsTest, 'test_{0}'.format(var), test_func1) + + ut.main() \ No newline at end of file