wrf-python/src/wrf/specialdec.py

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

import numpy as np

import wrapt 

#from .destag import destagger
from .util import iter_left_indexes, npvalues
from .py3compat import py3range
from .config import xarray_enabled
from .constants import Constants

if xarray_enabled():
    from xarray import DataArray


# 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 _move_dim_to_left(arr, dimidx):
        
    if isinstance(arr, np.ma.MaskedArray):
        has_missing = True
        missing = result.fill_value
    
    shape = list(arr.shape)
    move_dim_size = shape.pop(dimidx)
    
    output_dims = [move_dim_size] + shape
    
    output = np.empty(output_dims, arr.dtype)
    
    if dimidx < 0:
        right_ndim = -dimidx
    else:
        right_ndim = arr.ndim - dimidx
        
    rightdims = [slice(None)] * right_ndim
    
    for i in py3range(move_dim_size):
        rightdims[0] = i
        outidxs = (Ellipsis,) + tuple(rightdims)
        output[i,:] = outview_array[outidxs]

    if has_missing:
        output = np.ma.masked_values(output, missing)
    
    return output


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
        
        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 = u
        if isinstance(u, DataArray):
            u_arr = npvalues(u)
            
        v_arr = v
        if isinstance(v, DataArray):
            v_arr = npvalues(v)
          
        umissing = Constants.DEFAULT_FILL  
        if isinstance(u_arr, np.ma.MaskedArray):
            has_missing = True
            umissing = u_arr.fill_value
        
        vmissing = Constants.DEFAULT_FILL 
        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]
            
            # 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):
    """Decorator to handle iterating over leftmost dimensions when using 
    multiple files and/or multiple times with the cape product.
    
    
    """
    @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 = False if i3dflag != 0 else True
            
        # 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
        orig_dtype = p_hpa.dtype
        
        # 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,)
            output_dims += p_hpa.shape[-3:]
            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]
            
            # 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):
    """Decorator to handle iterating over leftmost dimensions when using 
    multiple files and/or multiple times with the cloudfrac product.
    
    
    """
    @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 = args[0]
        rh = args[1]
            
        num_left_dims = p.ndim - 3
        orig_dtype = p.dtype
        
        # No special left side iteration, build the output from the cape,cin
        # result
        if (num_left_dims == 0):
            low, med, high = wrapped(*new_args, **new_kwargs)
            
            output_dims = (3,)
            output_dims += p.shape[-2:]
            output = np.empty(output_dims, orig_dtype)
            
            output[0,:] = low[:]
            output[1,:] = med[:]
            output[2,:] = high[:]
            
            return output
                

        # Initial output is ...,cape_cin,nz,ny,nx to create contiguous views
        outdims = p.shape[0:num_left_dims]
        extra_dims = tuple(outdims) # Copy the left-most dims for iteration
        
        outdims += (3,) # low_mid_high
        
        outdims += p.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 += p.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),)
            low_idxs = left_idxs + (0, slice(None))
            med_idxs = left_idxs + (1, slice(None))
            high_idxs = left_idxs + (2, slice(None))
            
            low_output_idxs = (0,) + left_idxs + (slice(None),)
            med_output_idxs = (1,) + left_idxs + (slice(None),)
            high_output_idxs = (2,) + left_idxs + (slice(None),)
            
            new_args[0] = p[left_and_slice_idxs]
            new_args[1] = rh[left_and_slice_idxs]
            
            lowview = outview_array[low_idxs]
            medview = outview_array[med_idxs]
            highview = outview_array[high_idxs]
            
            new_kwargs["lowview"] = lowview
            new_kwargs["medview"] = medview
            new_kwargs["highview"] = highview
            
            low, med, 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[med_output_idxs] = (
                            outview_array[med_idxs].astype(orig_dtype))
            output[high_output_idxs] = (
                            outview_array[high_idxs].astype(orig_dtype))
        
        return output
    
    return func_wrapper