wrf-python/src/wrf/interputils.py

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

from math import floor, ceil

import numpy as np

from .extension import _interp2dxy
from .util import py3range

__all__ = ["to_positive_idxs", "calc_xy", "get_xy_z_params", "get_xy"]

def to_positive_idxs(shape, coord):
    if (coord[-2] >= 0 and coord[-1] >= 0):
        return coord
    
    return [x if (x >= 0) else shape[i]+x for (i,x) in enumerate(coord) ]

def calc_xy(xdim, ydim, pivot_point=None, angle=None, 
           start_point=None, end_point=None):
    """Returns the x,y points for the horizontal cross section line.
    
    xdim - maximum x-dimension
    ydim - maximum y-dimension
    pivot_point - a pivot point of (south_north, west_east) 
                  (must be used with angle)
    angle - the angle through the pivot point in degrees
    start_point - a start_point sequence of [south_north1, west_east1]
    end_point - an end point sequence of [south_north2, west_east2]
    
    """ 
    
    # Have a pivot point with an angle to find cross section
    if pivot_point is not None and angle is not None:
        xp = pivot_point[-1]
        yp = pivot_point[-2]
        
        if (angle > 315.0 or angle < 45.0 
            or ((angle > 135.0) and (angle < 225.0))):
            
            #x = y*slope + intercept
            slope = -(360.-angle)/45.
            if( angle < 45. ):
                slope = angle/45.
            if( angle > 135.):
                slope = (angle-180.)/45.
            
            intercept = xp - yp*slope
            
            # find intersections with domain boundaries
            y0 = 0.
            x0 = y0*slope + intercept
            
            if( x0 < 0.):  # intersect outside of left boundary
                x0 = 0.
                y0 =  (x0 - intercept)/slope
            if( x0 > xdim-1):  #intersect outside of right boundary
                x0 = xdim-1
                y0 =  (x0 - intercept)/slope
            y1 = ydim-1.  #need to make sure this will be a float?
            x1 = y1*slope + intercept
            
            if( x1 < 0.):  # intersect outside of left boundary
                x1 = 0.
                y1 =  (x1 - intercept)/slope
            
            if( x1 > xdim-1):  # intersect outside of right boundary
                x1 = xdim-1
                y1 =  (x1 - intercept)/slope
        else:
            #  y = x*slope + intercept
            slope = (90.-angle)/45.
            if( angle > 225. ):
                slope = (270.-angle)/45.
            intercept = yp - xp*slope

            #find intersections with domain boundaries
            x0 = 0.
            y0 = x0*slope + intercept
            
            if( y0 < 0.):  # intersect outside of bottom boundary
                y0 = 0.
                x0 =  (y0 - intercept)/slope
            
            if( y0 > ydim-1):  # intersect outside of top boundary
                y0 = ydim-1
                x0 =  (y0 - intercept)/slope
            
            x1 = xdim-1.  #  need to make sure this will be a float?
            y1 = x1*slope + intercept
            
            if( y1 < 0.):  # intersect outside of bottom boundary
                y1 = 0.
                x1 =  (y1 - intercept)/slope
            
            if( y1 > ydim-1):# intersect outside of top boundary
                y1 = ydim-1
                x1 =  (y1 - intercept)/slope
    elif start_point is not None and end_point is not None:
        x0 = start_point[-1]
        y0 = start_point[-2]
        x1 = end_point[-1]
        y1 = end_point[-2]
        if ( x1 > xdim-1 ): 
            x1 = xdim
        if ( y1 > ydim-1): 
            y1 = ydim
    else:
        raise ValueError("invalid start/end or pivot/angle arguments")
    
    dx = x1 - x0
    dy = y1 - y0
    distance = (dx*dx + dy*dy)**0.5
    npts = int(distance)
    dxy = distance/npts
    
    xy = np.zeros((npts,2), "float")

    dx = dx/npts
    dy = dy/npts
    
    for i in py3range(npts):
        xy[i,0] = x0 + i*dx
        xy[i,1] = y0 + i*dy
        
    return xy

def get_xy_z_params(z, pivot_point=None, angle=None,
                    start_point=None, end_point=None):
    
    xy = get_xy(z, pivot_point, angle, start_point, end_point)
    
    # Interp z
    var2dz = _interp2dxy(z, xy)
    
    extra_dim_num = z.ndim - 3
    idx1 = tuple([0]*extra_dim_num + [0,0])
    idx2 = tuple([0]*extra_dim_num + [-1,0])
    
    #  interp to constant z grid
    if(var2dz[idx1] > var2dz[idx2]):  # monotonically decreasing coordinate
        z_max = floor(np.amax(z) / 10) * 10     # bottom value
        z_min = ceil(np.amin(z) / 10) * 10      # top value
        dz = 10
        nlevels = int((z_max-z_min) / dz)
        z_var2d = np.zeros((nlevels), dtype=z.dtype)
        z_var2d[0] = z_max
        dz = -dz
    else:
        z_max = np.amax(z)
        z_min = 0.
        dz = 0.01 * z_max
        nlevels = int(z_max / dz)
        z_var2d = np.zeros((nlevels), dtype=z.dtype)
        z_var2d[0] = z_min
    
    for i in py3range(1,nlevels):
        z_var2d[i] = z_var2d[0] + i*dz
        
    return xy, var2dz, z_var2d

def get_xy(var, pivot_point=None, angle=None, 
           start_point=None, end_point=None):
    if pivot_point is not None:
        pos_pivot = to_positive_idxs(var.shape[-2:], pivot_point)
    else:
        pos_pivot = pivot_point
        
    if start_point is not None:
        pos_start = to_positive_idxs(var.shape[-2:], start_point)
    else:
        pos_start = start_point
    
    if end_point is not None:
        pos_end = to_positive_idxs(var.shape[-2:], end_point)
    else:
        pos_end = start_point   
        
    xdim = var.shape[-1]
    ydim = var.shape[-2]
    
    xy = calc_xy(xdim, ydim, pos_pivot, angle, pos_start, pos_end)
    
    return xy