wrf-python/src/wrf/specialdec.py

from __future__ import (absolute_import, division, print_function, 
                        unicode_literals)

import numpy as np

import wrapt 

from .util import iter_left_indexes, to_np
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
        is1d = np.isscalar(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 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
        
        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 is2d:
            if np.isscalar(ter) or np.isscalar(sfp):
                raise ValueError("arguments 4 and 5 cannot be scalars with "
                                 "cape_2d routine")
                
            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))
                
        # 3D cape is allowed to be just a vertical column with scalars
        # for terrain and psfc_hpa.
        else:
            if np.isscalar(ter) and np.isscalar(sfp):
                if p_hpa.ndim != 1:
                    raise ValueError("arguments 0-3 "
                                     "must be 1-dimensional when "
                                     "arguments 4 and 5 are scalars")
                if is2d:
                    raise ValueError("argument 7 must be 0 when using 1D data")
            else:
                if ((np.isscalar(ter) and not np.isscalar(sfp)) or 
                    (not np.isscalar(ter) and np.isscalar(sfp))):
                    raise ValueError("arguments 4 and 5 must both be scalars")
                else:
                    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))
      
        return wrapped(*args, **kwargs)
    
    return func_wrapper