|
|
@ -8,7 +8,6 @@ from collections import Iterable, Mapping, OrderedDict |
|
|
|
from itertools import product |
|
|
|
from itertools import product |
|
|
|
from types import GeneratorType |
|
|
|
from types import GeneratorType |
|
|
|
import datetime as dt |
|
|
|
import datetime as dt |
|
|
|
from math import floor, copysign |
|
|
|
|
|
|
|
from inspect import getmodule |
|
|
|
from inspect import getmodule |
|
|
|
|
|
|
|
|
|
|
|
try: |
|
|
|
try: |
|
|
@ -27,8 +26,7 @@ import numpy.ma as ma |
|
|
|
from .config import xarray_enabled |
|
|
|
from .config import xarray_enabled |
|
|
|
from .projection import getproj, NullProjection |
|
|
|
from .projection import getproj, NullProjection |
|
|
|
from .constants import Constants, ALL_TIMES |
|
|
|
from .constants import Constants, ALL_TIMES |
|
|
|
from .py3compat import (viewitems, viewkeys, viewvalues, isstr, py2round, |
|
|
|
from .py3compat import (viewitems, viewkeys, isstr, py3range, ucode) |
|
|
|
py3range, ucode) |
|
|
|
|
|
|
|
from .cache import cache_item, get_cached_item |
|
|
|
from .cache import cache_item, get_cached_item |
|
|
|
|
|
|
|
|
|
|
|
if xarray_enabled(): |
|
|
|
if xarray_enabled(): |
|
|
@ -73,12 +71,15 @@ def is_multi_time_req(timeidx): |
|
|
|
def is_multi_file(wrfnc): |
|
|
|
def is_multi_file(wrfnc): |
|
|
|
return (isinstance(wrfnc, Iterable) and not isstr(wrfnc)) |
|
|
|
return (isinstance(wrfnc, Iterable) and not isstr(wrfnc)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def has_time_coord(wrfnc): |
|
|
|
def has_time_coord(wrfnc): |
|
|
|
return "XTIME" in wrfnc.variables |
|
|
|
return "XTIME" in wrfnc.variables |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def is_mapping(wrfnc): |
|
|
|
def is_mapping(wrfnc): |
|
|
|
return isinstance(wrfnc, Mapping) |
|
|
|
return isinstance(wrfnc, Mapping) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _generator_copy(gen): |
|
|
|
def _generator_copy(gen): |
|
|
|
funcname = gen.__name__ |
|
|
|
funcname = gen.__name__ |
|
|
|
try: |
|
|
|
try: |
|
|
@ -96,10 +97,12 @@ def _generator_copy(gen): |
|
|
|
|
|
|
|
|
|
|
|
return res |
|
|
|
return res |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def test(): |
|
|
|
def test(): |
|
|
|
q = [1,2,3] |
|
|
|
q = [1,2,3] |
|
|
|
for i in q: |
|
|
|
for i in q: |
|
|
|
yield i |
|
|
|
yield i |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class TestGen(object): |
|
|
|
class TestGen(object): |
|
|
|
def __init__(self, count=3): |
|
|
|
def __init__(self, count=3): |
|
|
@ -121,6 +124,7 @@ class TestGen(object): |
|
|
|
def __next__(self): |
|
|
|
def __next__(self): |
|
|
|
return self.next() |
|
|
|
return self.next() |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def latlon_coordvars(d): |
|
|
|
def latlon_coordvars(d): |
|
|
|
lat_coord = None |
|
|
|
lat_coord = None |
|
|
|
lon_coord = None |
|
|
|
lon_coord = None |
|
|
@ -137,9 +141,11 @@ def latlon_coordvars(d): |
|
|
|
|
|
|
|
|
|
|
|
return lat_coord, lon_coord |
|
|
|
return lat_coord, lon_coord |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def is_coordvar(varname): |
|
|
|
def is_coordvar(varname): |
|
|
|
return varname in _COORD_VARS |
|
|
|
return varname in _COORD_VARS |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class IterWrapper(Iterable): |
|
|
|
class IterWrapper(Iterable): |
|
|
|
"""A wrapper class for generators and custom iterable classes which returns |
|
|
|
"""A wrapper class for generators and custom iterable classes which returns |
|
|
|
a new iterator from the start of the sequence when __iter__ is called""" |
|
|
|
a new iterator from the start of the sequence when __iter__ is called""" |
|
|
@ -172,6 +178,7 @@ def get_iterable(wrfseq): |
|
|
|
else: |
|
|
|
else: |
|
|
|
return dict(wrfseq) # generator/custom iterable class |
|
|
|
return dict(wrfseq) # generator/custom iterable class |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# Helper to extract masked arrays from DataArrays that convert to NaN |
|
|
|
# Helper to extract masked arrays from DataArrays that convert to NaN |
|
|
|
def npvalues(array_type): |
|
|
|
def npvalues(array_type): |
|
|
|
if not isinstance(array_type, DataArray): |
|
|
|
if not isinstance(array_type, DataArray): |
|
|
@ -187,6 +194,7 @@ def npvalues(array_type): |
|
|
|
|
|
|
|
|
|
|
|
return result |
|
|
|
return result |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# Helper utilities for metadata |
|
|
|
# Helper utilities for metadata |
|
|
|
class either(object): |
|
|
|
class either(object): |
|
|
|
def __init__(self, *varnames): |
|
|
|
def __init__(self, *varnames): |
|
|
@ -207,6 +215,7 @@ class either(object): |
|
|
|
raise ValueError("{} are not valid variable names".format( |
|
|
|
raise ValueError("{} are not valid variable names".format( |
|
|
|
self.varnames)) |
|
|
|
self.varnames)) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
class combine_with(object): |
|
|
|
class combine_with(object): |
|
|
|
# Remove remove_idx first, then insert_idx is applied to removed set |
|
|
|
# Remove remove_idx first, then insert_idx is applied to removed set |
|
|
|
def __init__(self, varname, remove_dims=None, insert_before=None, |
|
|
|
def __init__(self, varname, remove_dims=None, insert_before=None, |
|
|
@ -239,6 +248,7 @@ class combine_with(object): |
|
|
|
|
|
|
|
|
|
|
|
return new_dims, new_coords |
|
|
|
return new_dims, new_coords |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# This should look like: |
|
|
|
# This should look like: |
|
|
|
# [(0, (-3,-2)), # variable 1 |
|
|
|
# [(0, (-3,-2)), # variable 1 |
|
|
|
# (1, -1)] # variable 2 |
|
|
|
# (1, -1)] # variable 2 |
|
|
@ -387,6 +397,7 @@ def is_moving_domain(wrfseq, varname=None, latvar=either("XLAT", "XLAT_M"), |
|
|
|
|
|
|
|
|
|
|
|
return False |
|
|
|
return False |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _get_global_attr(wrfnc, attr): |
|
|
|
def _get_global_attr(wrfnc, attr): |
|
|
|
val = getattr(wrfnc, attr, None) |
|
|
|
val = getattr(wrfnc, attr, None) |
|
|
|
|
|
|
|
|
|
|
@ -395,6 +406,7 @@ def _get_global_attr(wrfnc, attr): |
|
|
|
if len(val) == 1: |
|
|
|
if len(val) == 1: |
|
|
|
return val[0] |
|
|
|
return val[0] |
|
|
|
return val |
|
|
|
return val |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def extract_global_attrs(wrfnc, attrs): |
|
|
|
def extract_global_attrs(wrfnc, attrs): |
|
|
|
if isstr(attrs): |
|
|
|
if isstr(attrs): |
|
|
@ -413,6 +425,7 @@ def extract_global_attrs(wrfnc, attrs): |
|
|
|
|
|
|
|
|
|
|
|
return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist} |
|
|
|
return {attr:_get_global_attr(wrfnc, attr) for attr in attrlist} |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def extract_dim(wrfnc, dim): |
|
|
|
def extract_dim(wrfnc, dim): |
|
|
|
if is_multi_file(wrfnc): |
|
|
|
if is_multi_file(wrfnc): |
|
|
|
if not is_mapping(wrfnc): |
|
|
|
if not is_mapping(wrfnc): |
|
|
@ -426,6 +439,7 @@ def extract_dim(wrfnc, dim): |
|
|
|
return len(d) #netCDF4 |
|
|
|
return len(d) #netCDF4 |
|
|
|
return d # PyNIO |
|
|
|
return d # PyNIO |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _combine_dict(wrfdict, varname, timeidx, method, meta, _key): |
|
|
|
def _combine_dict(wrfdict, varname, timeidx, method, meta, _key): |
|
|
|
"""Dictionary combination creates a new left index for each key, then |
|
|
|
"""Dictionary combination creates a new left index for each key, then |
|
|
|
does a cat or join for the list of files for that key""" |
|
|
|
does a cat or join for the list of files for that key""" |
|
|
@ -656,7 +670,7 @@ def _build_data_array(wrfnc, varname, timeidx, is_moving_domain, is_multifile, |
|
|
|
time_coord_vals = None |
|
|
|
time_coord_vals = None |
|
|
|
if time_coord is not None: |
|
|
|
if time_coord is not None: |
|
|
|
# If not from a multifile sequence, then cache the time |
|
|
|
# If not from a multifile sequence, then cache the time |
|
|
|
# coordinate. Otherwise, handled in cat/join/ |
|
|
|
# coordinate. Otherwise, handled in cat/join/ |
|
|
|
if not is_multifile: |
|
|
|
if not is_multifile: |
|
|
|
time_coord_vals = get_cached_item(_key, time_coord) |
|
|
|
time_coord_vals = get_cached_item(_key, time_coord) |
|
|
|
|
|
|
|
|
|
|
@ -798,6 +812,7 @@ def _find_arr_for_time(wrfseq, varname, timeidx, is_moving, meta, _key): |
|
|
|
else: |
|
|
|
else: |
|
|
|
return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key) |
|
|
|
return _find_reverse(wrfseq, varname, timeidx, is_moving, meta, _key) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# TODO: implement in C |
|
|
|
# TODO: implement in C |
|
|
|
def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): |
|
|
|
def _cat_files(wrfseq, varname, timeidx, is_moving, squeeze, meta, _key): |
|
|
|
if is_moving is None: |
|
|
|
if is_moving is None: |
|
|
@ -1234,6 +1249,7 @@ def _join_files(wrfseq, varname, timeidx, is_moving, meta, _key): |
|
|
|
|
|
|
|
|
|
|
|
return outarr |
|
|
|
return outarr |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def combine_files(wrfseq, varname, timeidx, is_moving=None, |
|
|
|
def combine_files(wrfseq, varname, timeidx, is_moving=None, |
|
|
|
method="cat", squeeze=True, meta=True, |
|
|
|
method="cat", squeeze=True, meta=True, |
|
|
|
_key=None): |
|
|
|
_key=None): |
|
|
@ -1416,7 +1432,6 @@ def is_staggered(var, wrfnc): |
|
|
|
return False |
|
|
|
return False |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_left_indexes(ref_var, expected_dims): |
|
|
|
def get_left_indexes(ref_var, expected_dims): |
|
|
|
"""Returns the extra left side dimensions for a variable with an expected |
|
|
|
"""Returns the extra left side dimensions for a variable with an expected |
|
|
|
shape. |
|
|
|
shape. |
|
|
@ -1432,6 +1447,7 @@ def get_left_indexes(ref_var, expected_dims): |
|
|
|
|
|
|
|
|
|
|
|
return tuple([ref_var.shape[x] for x in py3range(extra_dim_num)]) |
|
|
|
return tuple([ref_var.shape[x] for x in py3range(extra_dim_num)]) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def iter_left_indexes(dims): |
|
|
|
def iter_left_indexes(dims): |
|
|
|
"""A generator which yields the iteration tuples for a sequence of |
|
|
|
"""A generator which yields the iteration tuples for a sequence of |
|
|
|
dimensions sizes. |
|
|
|
dimensions sizes. |
|
|
@ -1448,6 +1464,7 @@ def iter_left_indexes(dims): |
|
|
|
arg = [py3range(dim) for dim in dims] |
|
|
|
arg = [py3range(dim) for dim in dims] |
|
|
|
for idxs in product(*arg): |
|
|
|
for idxs in product(*arg): |
|
|
|
yield idxs |
|
|
|
yield idxs |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_right_slices(var, right_ndims, fixed_val=0): |
|
|
|
def get_right_slices(var, right_ndims, fixed_val=0): |
|
|
|
extra_dim_num = var.ndim - right_ndims |
|
|
|
extra_dim_num = var.ndim - right_ndims |
|
|
@ -1457,6 +1474,7 @@ def get_right_slices(var, right_ndims, fixed_val=0): |
|
|
|
return tuple([fixed_val]*extra_dim_num + |
|
|
|
return tuple([fixed_val]*extra_dim_num + |
|
|
|
[slice(None)]*right_ndims) |
|
|
|
[slice(None)]*right_ndims) |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def get_proj_params(wrfnc, timeidx=0, varname=None): |
|
|
|
def get_proj_params(wrfnc, timeidx=0, varname=None): |
|
|
|
proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", |
|
|
|
proj_params = extract_global_attrs(wrfnc, attrs=("MAP_PROJ", |
|
|
|
"CEN_LAT", "CEN_LON", |
|
|
|
"CEN_LAT", "CEN_LON", |
|
|
@ -1556,8 +1574,9 @@ def from_args(func, argnames, *args, **kwargs): |
|
|
|
|
|
|
|
|
|
|
|
return result |
|
|
|
return result |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _args_to_list2(func, args, kwargs): |
|
|
|
def _args_to_list2(func, args, kwargs): |
|
|
|
argspec = getargspec(func) |
|
|
|
argspec = getargspec(func) |
|
|
|
|
|
|
|
|
|
|
|
# Build the full tuple with defaults filled in |
|
|
|
# Build the full tuple with defaults filled in |
|
|
|
outargs = [None]*len(argspec.args) |
|
|
|
outargs = [None]*len(argspec.args) |
|
|
@ -1576,14 +1595,16 @@ def _args_to_list2(func, args, kwargs): |
|
|
|
|
|
|
|
|
|
|
|
return outargs |
|
|
|
return outargs |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _args_to_list3(func, args, kwargs): |
|
|
|
def _args_to_list3(func, args, kwargs): |
|
|
|
sig = signature(func) |
|
|
|
sig = signature(func) |
|
|
|
bound = sig.bind(*args, **kwargs) |
|
|
|
bound = sig.bind(*args, **kwargs) |
|
|
|
bound.apply_defaults() |
|
|
|
bound.apply_defaults() |
|
|
|
|
|
|
|
|
|
|
|
return [x for x in bound.arguments.values()] |
|
|
|
return [x for x in bound.arguments.values()] |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
# Note: Doesn't allow for **kwargs or *args |
|
|
|
def args_to_list(func, args, kwargs): |
|
|
|
def args_to_list(func, args, kwargs): |
|
|
|
"""Converts the mixed args/kwargs to a single list of args""" |
|
|
|
"""Converts the mixed args/kwargs to a single list of args""" |
|
|
|
if version_info > (3,): |
|
|
|
if version_info > (3,): |
|
|
@ -1607,6 +1628,7 @@ def _arg_location2(func, argname, args, kwargs): |
|
|
|
|
|
|
|
|
|
|
|
return list_args, result_idx |
|
|
|
return list_args, result_idx |
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
def _arg_location3(func, argname, args, kwargs): |
|
|
|
def _arg_location3(func, argname, args, kwargs): |
|
|
|
sig = signature(func) |
|
|
|
sig = signature(func) |
|
|
|
params = list(sig.parameters.keys()) |
|
|
|
params = list(sig.parameters.keys()) |
|
|
|