from __future__ import (absolute_import, division, print_function)

import numpy as np

import wrapt

from .util import iter_left_indexes, to_np

from .py3compat import py3range

from .config import xarray_enabled

from .constants import default_fill

if xarray_enabled():

from xarray import DataArray

def uvmet_left_iter(alg_dtype=np.float64):

"""A decorator to handle iterating over the leftmost dimensions for the

uvmet diagnostic.

For example, if a wrapped function works with three-dimensional arrays, but

the variables include a 4th leftmost dimension for 'Time', this decorator

will iterate over all times, call the 3D Fortran routine, and aggregate the

results in to a 4D output array.

It is also important to note that the final output array is allocated

first, and then views are passed to the wrapped function so that values

do not need to get copied in to the final output array.

Args:

alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used

in the wrapped function.

Returns:

:class:`numpy.ndarray`: The aggregated uvmet output array that includes

all extra leftmost dimensions.

"""

@wrapt.decorator

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

lat_lon_fixed = False

if lat.ndim == 2:

lat_lon_fixed = True

if lon.ndim == 2 and not lat_lon_fixed:

raise ValueError("'lat' and 'lon' shape mismatch")

num_left_dims_u = u.ndim - 2

num_left_dims_lat = lat.ndim - 2

if (num_left_dims_lat > num_left_dims_u):

raise ValueError("number of 'lat' dimensions is greater than 'u'")

if lat_lon_fixed:

mode = 0 # fixed lat/lon

else:

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 = to_np(u)

v_arr = to_np(v)

umissing = default_fill(np.float64)

if isinstance(u_arr, np.ma.MaskedArray):

has_missing = True

umissing = u_arr.fill_value

vmissing = default_fill(np.float64)

if isinstance(v_arr, np.ma.MaskedArray):

has_missing = True

vmissing = v_arr.fill_value

uvmetmissing = umissing

is_stag = 0

if (u.shape[-1] != lat.shape[-1] or u.shape[-2] != lat.shape[-2]):

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 = 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")

# No special left side iteration, return the function result

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

outdims += (2,)

outdims += lat.shape[-2:]

outview_array = np.empty(outdims, alg_dtype)

# Final Output moves the u_v dimension to left side

output_dims = (2,)

output_dims += extra_dims

output_dims += lat.shape[-2:]

output = np.empty(output_dims, orig_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])

u_output_idxs = (0,) + left_idxs + (slice(None),)

v_output_idxs = (1,) + left_idxs + (slice(None),)

u_view_idxs = left_idxs + (0, slice(None))

v_view_idxs = left_idxs + (1, slice(None))

new_u = u[left_and_slice_idxs]

new_v = v[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]

# Skip the possible empty/missing arrays for the join method

skip_missing = False

for arg in (new_u, new_v, new_lat, new_lon):

if isinstance(arg, np.ma.MaskedArray):

if arg.mask.all():

output[u_output_idxs] = uvmetmissing

output[v_output_idxs] = uvmetmissing

skip_missing = True

has_missing = True

if skip_missing:

continue

# Call the numerical routine

result = wrapped(new_u, new_v, new_lat, new_lon, cen_long, cone,

isstag=is_stag, has_missing=has_missing,

umissing=umissing, vmissing=vmissing,

uvmetmissing=uvmetmissing, outview=outview)

# Make sure the result is the same data as what got passed in

# Can delete this once everything works

if (result.__array_interface__["data"][0] !=

outview.__array_interface__["data"][0]):

raise RuntimeError("output array was copied")

output[u_output_idxs] = (

outview_array[u_view_idxs].astype(orig_dtype))

output[v_output_idxs] = (

outview_array[v_view_idxs].astype(orig_dtype))

if has_missing:

output = np.ma.masked_values(output, uvmetmissing)

return output

return func_wrapper

def cape_left_iter(alg_dtype=np.float64):

"""A decorator to handle iterating over the leftmost dimensions for the

cape diagnostic.

For example, if a wrapped function works with three-dimensional arrays, but

the variables include a 4th leftmost dimension for 'Time', this decorator

will iterate over all times, call the 3D Fortran routine, and aggregate the

results in to a 4D output array.

It is also important to note that the final output array is allocated

first, and then views are passed to the wrapped function so that values

do not need to get copied in to the final output array.

Args:

alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used

in the wrapped function.

Returns:

:class:`numpy.ndarray`: The aggregated cape output array that includes

all extra leftmost dimensions.

"""

@wrapt.decorator

def func_wrapper(wrapped, instance, args, kwargs):

# The cape calculations use an ascending vertical pressure coordinate

new_args = list(args)

new_kwargs = dict(kwargs)

p_hpa = args[0]

tk = args[1]

qv = args[2]

ht = args[3]

ter = args[4]

sfp = args[5]

missing = args[6]

i3dflag = args[7]

ter_follow = args[8]

is2d = i3dflag == 0

# Note: This should still work with DataArrays

is1d = np.isscalar(sfp) or np.size(sfp) == 1

# Make sure sfp and terrain are regular floats for 1D case

# This should also work with DataArrays

if is1d:

ter = float(ter)

sfp = float(sfp)

orig_dtype = p_hpa.dtype

if not is1d:

# Need to order in ascending pressure order

flip = False

bot_idxs = (0,) * p_hpa.ndim

top_idxs = list(bot_idxs)

top_idxs[-3] = -1

top_idxs = tuple(top_idxs)

if p_hpa[bot_idxs] > p_hpa[top_idxs]:

flip = True

p_hpa = np.ascontiguousarray(p_hpa[..., ::-1, :, :])

tk = np.ascontiguousarray(tk[..., ::-1, :, :])

qv = np.ascontiguousarray(qv[..., ::-1, :, :])

ht = np.ascontiguousarray(ht[..., ::-1, :, :])

new_args[0] = p_hpa

new_args[1] = tk

new_args[2] = qv

new_args[3] = ht

num_left_dims = p_hpa.ndim - 3

else:

# Need to order in ascending pressure order

flip = False

if p_hpa[0] > p_hpa[-1]:

flip = True

p_hpa = np.ascontiguousarray(p_hpa[::-1])

tk = np.ascontiguousarray(tk[::-1])

qv = np.ascontiguousarray(qv[::-1])

ht = np.ascontiguousarray(ht[::-1])

# Need to make 3D views for the fortran code.

# Going to make these fortran ordered, since the f_contiguous and

# c_contiguous flags are broken in numpy 1.11 (always false). This

# should work across all numpy versions.

new_args[0] = p_hpa.reshape((1, 1, p_hpa.shape[0]), order='F')

new_args[1] = tk.reshape((1, 1, tk.shape[0]), order='F')

new_args[2] = qv.reshape((1, 1, qv.shape[0]), order='F')

new_args[3] = ht.reshape((1, 1, ht.shape[0]), order='F')

new_args[4] = np.full((1, 1), ter, orig_dtype)

new_args[5] = np.full((1, 1), sfp, orig_dtype)

num_left_dims = 0

# No special left side iteration, build the output from the cape,cin

# result

if (num_left_dims == 0):

cape, cin = wrapped(*new_args, **new_kwargs)

output_dims = (2,)

if not is1d:

output_dims += p_hpa.shape[-3:]

else:

output_dims += (p_hpa.shape[0], 1, 1)

output = np.empty(output_dims, orig_dtype)

if flip and not is2d:

output[0, :] = cape[::-1, :, :]

output[1, :] = cin[::-1, :, :]

else:

output[0, :] = cape[:]

output[1, :] = cin[:]

return output

# Initial output is ...,cape_cin,nz,ny,nx to create contiguous views

outdims = p_hpa.shape[0:num_left_dims]

extra_dims = tuple(outdims) # Copy the left-most dims for iteration

outdims += (2,) # cape_cin

outdims += p_hpa.shape[-3:]

outview_array = np.empty(outdims, alg_dtype)

# Create the output array where the leftmost dim is the product type

output_dims = (2,)

output_dims += extra_dims

output_dims += p_hpa.shape[-3:]

output = np.empty(output_dims, orig_dtype)

for left_idxs in iter_left_indexes(extra_dims):

left_and_slice_idxs = left_idxs + (slice(None),)

cape_idxs = left_idxs + (0, slice(None))

cin_idxs = left_idxs + (1, slice(None))

cape_output_idxs = (0,) + left_idxs + (slice(None),)

cin_output_idxs = (1,) + left_idxs + (slice(None),)

view_cape_reverse_idxs = left_idxs + (0, slice(None, None, -1),

slice(None))

view_cin_reverse_idxs = left_idxs + (1, slice(None, None, -1),

slice(None))

new_args[0] = p_hpa[left_and_slice_idxs]

new_args[1] = tk[left_and_slice_idxs]

new_args[2] = qv[left_and_slice_idxs]

new_args[3] = ht[left_and_slice_idxs]

new_args[4] = ter[left_and_slice_idxs]

new_args[5] = sfp[left_and_slice_idxs]

capeview = outview_array[cape_idxs]

cinview = outview_array[cin_idxs]

# Skip the possible empty/missing arrays for the join method

# Note: Masking handled by cape.py or computation.py, so only

# supply the fill values here.

skip_missing = False

for arg in (new_args[0:6]):

if isinstance(arg, np.ma.MaskedArray):

if arg.mask.all():

if flip and not is2d:

output[cape_output_idxs] = missing

output[cin_output_idxs] = missing

else:

output[cape_output_idxs] = missing

output[cin_output_idxs] = missing

skip_missing = True

if skip_missing:

continue

# Call the numerical routine

new_kwargs["capeview"] = capeview

new_kwargs["cinview"] = cinview

cape, cin = wrapped(*new_args, **new_kwargs)

# Make sure the result is the same data as what got passed in

# Can delete this once everything works

if (cape.__array_interface__["data"][0] !=

capeview.__array_interface__["data"][0]):

raise RuntimeError("output array was copied")

if flip and not is2d:

output[cape_output_idxs] = (

outview_array[view_cape_reverse_idxs].astype(orig_dtype))

output[cin_output_idxs] = (

outview_array[view_cin_reverse_idxs].astype(orig_dtype))

else:

output[cape_output_idxs] = (

outview_array[cape_idxs].astype(orig_dtype))

output[cin_output_idxs] = (

outview_array[cin_idxs].astype(orig_dtype))

return output

return func_wrapper

def cloudfrac_left_iter(alg_dtype=np.float64):

"""A decorator to handle iterating over the leftmost dimensions for the

cloud fraction diagnostic.

For example, if a wrapped function works with three-dimensional arrays, but

the variables include a 4th leftmost dimension for 'Time', this decorator

will iterate over all times, call the 3D Fortran routine, and aggregate the

results in to a 4D output array.

It is also important to note that the final output array is allocated

first, and then views are passed to the wrapped function so that values

do not need to get copied in to the final output array.

Args:

alg_dtype (:class:`np.dtype` or :obj:`str`): The numpy data type used

in the wrapped function.

Returns:

:class:`numpy.ndarray`: The aggregated cloud fraction output array

that includes all extra leftmost dimensions.

"""

@wrapt.decorator

def func_wrapper(wrapped, instance, args, kwargs):

new_args = list(args)

new_kwargs = dict(kwargs)

vert = args[0]

rh = args[1]

num_left_dims = vert.ndim - 3

orig_dtype = vert.dtype

# No special left side iteration, build the output from the

# low, mid, high results.

if (num_left_dims == 0):

low, mid, high = wrapped(*new_args, **new_kwargs)

output_dims = (3,)

output_dims += vert.shape[-2:]

output = np.empty(output_dims, orig_dtype)

output[0, :] = low[:]

output[1, :] = mid[:]

output[2, :] = high[:]

return output

# Initial output is ...,low_mid_high,nz,ny,nx to create contiguous

# views

outdims = vert.shape[0:num_left_dims]

extra_dims = tuple(outdims) # Copy the left-most dims for iteration

outdims += (3,) # low_mid_high

outdims += vert.shape[-2:]

outview_array = np.empty(outdims, alg_dtype)

# Create the output array where the leftmost dim is the cloud type

output_dims = (3,)

output_dims += extra_dims

output_dims += vert.shape[-2:]

output = np.empty(output_dims, orig_dtype)

has_missing = False

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))

mid_idxs = left_idxs + (1, slice(None))

high_idxs = left_idxs + (2, slice(None))

low_output_idxs = (0,) + left_idxs + (slice(None),)

mid_output_idxs = (1,) + left_idxs + (slice(None),)

high_output_idxs = (2,) + left_idxs + (slice(None),)

new_args[0] = vert[left_and_slice_idxs]

new_args[1] = rh[left_and_slice_idxs]

# Skip the possible empty/missing arrays for the join method

# Note: Masking handled by cloudfrac.py or computation.py, so only

# supply the fill values here.

skip_missing = False

for arg in (new_args[0:2]):

if isinstance(arg, np.ma.MaskedArray):

if arg.mask.all():

output[low_output_idxs] = missing

output[mid_output_idxs] = missing

output[high_output_idxs] = missing

skip_missing = True

has_missing = True

if skip_missing:

continue

lowview = outview_array[low_idxs]

midview = outview_array[mid_idxs]

highview = outview_array[high_idxs]

new_kwargs["lowview"] = lowview

new_kwargs["midview"] = midview

new_kwargs["highview"] = highview

low, mid, high = wrapped(*new_args, **new_kwargs)

# Make sure the result is the same data as what got passed in

# Can delete this once everything works

if (low.__array_interface__["data"][0] !=

lowview.__array_interface__["data"][0]):

raise RuntimeError("output array was copied")

output[low_output_idxs] = (

outview_array[low_idxs].astype(orig_dtype))

output[mid_output_idxs] = (

outview_array[mid_idxs].astype(orig_dtype))

output[high_output_idxs] = (

outview_array[high_idxs].astype(orig_dtype))

if has_missing:

output = np.ma.masked_values(output, missing)

return output

return func_wrapper

def interplevel_left_iter(is2dlev, alg_dtype=np.float64):

@wrapt.decorator

def func_wrapper(wrapped, instance, args, kwargs):

new_args = list(args)

new_kwargs = dict(kwargs)

field3d = args[0]

z = args[1]

levels = args[2]

num_left_dims = z.ndim - 3

orig_dtype = field3d.dtype

left_dims = z.shape[0:num_left_dims]

multiproduct = True if field3d.ndim - z.ndim == 1 else False

# No special left side iteration, build the output from the

# low, mid, high results.

if (num_left_dims == 0):

if multiproduct:

if not is2dlev:

outshape = (field3d.shape[0:-3] + levels.shape +

field3d.shape[-2:])

else:

outshape = (field3d.shape[0:-3] + field3d.shape[-2:])

output = np.empty(outshape, dtype=alg_dtype)

for i in py3range(field3d.shape[0]):

new_args[0] = field3d[i, :]

new_kwargs["outview"] = output[i, :]

_ = wrapped(*new_args, **new_kwargs)

else:

output = wrapped(*args, **kwargs)

return output

if multiproduct:

outdims = field3d.shape[0:1] + left_dims

else:

outdims = left_dims

extra_dims = tuple(outdims)

if not is2dlev:

outdims += levels.shape

outdims += z.shape[-2:]

outview_array = np.empty(outdims, alg_dtype)

for left_idxs in iter_left_indexes(extra_dims):

field_out_slice_idxs = left_idxs + (slice(None),)

if multiproduct:

z_slice_idxs = left_idxs[1:] + (slice(None),)

else:

z_slice_idxs = left_idxs + (slice(None),)

new_args[0] = field3d[field_out_slice_idxs]

new_args[1] = z[z_slice_idxs]

if is2dlev:

if levels.ndim > 2:

new_args[2] = levels[z_slice_idxs]

new_kwargs["outview"] = outview_array[field_out_slice_idxs]

_ = wrapped(*new_args, **new_kwargs)

output = outview_array.astype(orig_dtype)

return output

return func_wrapper

def check_cape_args():

"""A decorator to check that the cape_3d arguments are valid.

An exception is raised when an invalid argument is found.

Returns:

None

Raises:

:class:`ValueError`: Raised when an invalid argument is detected.

"""

@wrapt.decorator

def func_wrapper(wrapped, instance, args, kwargs):

p_hpa = args[0]

tk = args[1]

qv = args[2]

ht = args[3]

ter = args[4]

sfp = args[5]

missing = args[6]

i3dflag = args[7]

ter_follow = args[8]

is2d = False if i3dflag != 0 else True

is1d = ((np.isscalar(sfp) or np.size(sfp) == 1) or

(np.isscalar(ter) or np.size(ter) == 1))

if not (p_hpa.shape == tk.shape == qv.shape == ht.shape):

raise ValueError("arguments 0, 1, 2, 3 must be the same shape")

# 2D CAPE does not allow for scalars

if not is1d:

if ter.ndim != p_hpa.ndim-1 or sfp.ndim != p_hpa.ndim-1:

raise ValueError("arguments 4 and 5 must have "

"{} dimensions".format(p_hpa.ndim-1))

else:

if np.size(ter) != np.size(sfp):

raise ValueError("arguments 4 and 5 must both be scalars or "

"both be arrays")

# Only need to test p_hpa since we assured args 0-3 have same ndim

if p_hpa.ndim != 1:

raise ValueError("arguments 0-3 "

"must be 1-dimensional when "

"arguments 4 and 5 are scalars")

return wrapped(*args, **kwargs)

return func_wrapper

def check_interplevel_args(is2dlev):

"""A decorator to check that the interplevel arguments are valid.

An exception is raised when an invalid argument is found.

Returns:

None

Raises:

:class:`ValueError`: Raised when an invalid argument is detected.

"""

@wrapt.decorator

def func_wrapper(wrapped, instance, args, kwargs):

field3d = args[0]

z = args[1]

levels = args[2]

multiproduct = True if (field3d.ndim - z.ndim) == 1 else False

if not multiproduct:

if field3d.shape != z.shape:

raise ValueError("arguments 0 and 1 must have the same shape")

else:

if field3d.shape[1:] != z.shape:

raise ValueError("argument 0 and 1 must have same rightmost "

"dimensions")

if is2dlev:

if levels.ndim != 2:

if (levels.shape[0:-2] != z.shape[0:-3] or

levels.shape[-2:] != z.shape[-2:]):

raise ValueError("argument 1 and 2 must have "

"the same leftmost and rightmost "

"dimensions")

return wrapped(*args, **kwargs)

return func_wrapper