from __future__ import (absolute_import, division, print_function) from math import floor, ceil import numpy as np from .extension import _interp2dxy from .py3compat import py3range from .coordpair import CoordPair from .constants import Constants, ProjectionTypes from .latlonutils import _ll_to_xy from .util import pairs_to_latlon def to_positive_idxs(shape, coord): """Return the positive index values. This function converts negative index values to positive index values. Args: shape (indexable sequence): The array shape. coord (indexable sequence): The coordinate pair for x and y. Returns: :obj:`list`: The coordinate values with all positive indexes. """ if (coord[-2] >= 0 and coord[-1] >= 0): return coord return [x if (x >= 0) else shape[-i-1]+x for (i,x) in enumerate(coord)] def _calc_xy(xdim, ydim, pivot_point=None, angle=None, start_point=None, end_point=None): """Return the x,y points for the horizontal cross section line. Args: xdim (:obj:`int`): The x-dimension size. ydim (:obj:`int`): The y-dimension size. pivot_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the x,y location through which the plane will pass. Must also specify `angle`. angle (:obj:`float`, optional): Only valid for cross sections where a plane will be plotted through a given point on the model domain. 0.0 represents a S-N cross section. 90.0 is a W-E cross section. start_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the start x,y location through which the plane will pass. end_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the end x,y location through which the plane will pass. Returns: :class:`np.ndarray`: A two-dimensional array with the left index representing each point along the line, and the rightmost dimension having two values for the x and y coordinates [0=X, 1=Y]. """ # 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[-2] yp = pivot_point[-1] if xp >= xdim or yp >= ydim: raise ValueError("pivot point {} is outside of domain " "with shape {}".format(pivot_point, (xdim, ydim))) 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[-2] y0 = start_point[-1] x1 = end_point[-2] y1 = end_point[-1] if x0 >= xdim or y0 >= ydim: raise ValueError("start_point {} is outside of domain " "with shape {}".format(start_point, (xdim, ydim))) if x1 >= xdim or y1 >= ydim: raise ValueError("end_point {} is outside of domain " "with shape {}".format(end_point, (xdim, 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) + 1 xy = np.zeros((npts,2), "float") dx = dx/(npts-1) dy = dy/(npts-1) 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, levels=None, autolevels=100): """Return the cross section parameters. This function returns the xy horizontal cross section line coordinates, the xy x z vertical values interpolated along the xy cross section line, and the fixed vertical levels to be used by the cross section algorithm (at ~1% increments for the minimum to maximum vertical span). Args: z (:class:`numpy.ndarray`): The vertical coordinate, whose rightmost dimensions are bottom_top x south_north x west_east. pivot_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the x,y location through which the plane will pass. Must also specify `angle`. angle (:obj:`float`, optional): Only valid for cross sections where a plane will be plotted through a given point on the model domain. 0.0 represents a S-N cross section. 90.0 is a W-E cross section. start_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the start x,y location through which the plane will pass. end_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the end x,y location through which the plane will pass. levels (sequence): A sequence of :obj:`float` for the desired vertical levels in the output array. If None, a fixed set of vertical levels is provided. Default is None. autolevels(:obj:`int`, optional): The number of evenly spaced automatically chosen vertical levels to use when *levels* is None. Default is 100. Returns: :obj:`tuple`: A tuple containing the xy horizontal cross section coordinates, the vertical values interpolated along the xy cross section line, and the fixed vertical levels used by the cross section algorithm at ~1% increments of minimum to maximum vertical span. """ 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]) if levels is None: # 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 = (1.0/autolevels) * (z_max - z_min) z_var2d = np.zeros((autolevels), dtype=z.dtype) z_var2d[0] = z_max dz = -dz else: z_max = np.amax(z) z_min = 0. dz = (1.0/autolevels)*z_max z_var2d = np.zeros((autolevels), dtype=z.dtype) z_var2d[0] = z_min for i in py3range(1,autolevels): z_var2d[i] = z_var2d[0] + i*dz else: z_var2d = np.asarray(levels, z.dtype) return xy, var2dz, z_var2d def get_xy(var, pivot_point=None, angle=None, start_point=None, end_point=None): """Return the x,y points for the horizontal cross section line. Args: var (:class:`xarray.DataArray` or :class:`numpy.ndarray`): A variable that contains a :attr:`shape` attribute. pivot_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the x,y location through which the plane will pass. Must also specify `angle`. angle (:obj:`float`, optional): Only valid for cross sections where a plane will be plotted through a given point on the model domain. 0.0 represents a S-N cross section. 90.0 is a W-E cross section. start_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the start x,y location through which the plane will pass. end_point (:obj:`tuple` or :obj:`list`, optional): A :obj:`tuple` or :obj:`list` with two entries, in the form of [x, y] (or [west_east, south_north]), which indicates the end x,y location through which the plane will pass. Returns: :class:`np.ndarray`: A two-dimensional array with the left index representing each point along the line, and the rightmost dimension having two values for the x and y coordinates [0=X, 1=Y]. """ 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 def to_xy_coords(pairs, wrfin=None, timeidx=0, stagger=None, projection=None, ll_point=None): """Return the coordinate pairs in grid space. This function converts latitude,longitude coordinate pairs to x,y coordinate pairs. Args: pairs (:class:`CoordPair` or sequence): A single coordinate pair or a sequence of coordinate pairs to be converted. wrfin (:class:`netCDF4.Dataset`, :class:`Nio.NioFile`, or an \ iterable, optional): WRF-ARW NetCDF data as a :class:`netCDF4.Dataset`, :class:`Nio.NioFile` or an iterable sequence of the aforementioned types. This is used to obtain the map projection when using latitude,longitude coordinates. Should not be used when working with x,y coordinates. Default is None. timeidx (:obj:`int` or :data:`wrf.ALL_TIMES`, optional): The desired time index when obtaining map boundary information from moving nests. This value can be a positive integer, negative integer, or :data:`wrf.ALL_TIMES` (an alias for None) to return all times in the file or sequence. Only required when *wrfin* is specified and the nest is moving. Default is 0. stagger (:obj:`str`): If using latitude, longitude coordinate pairs for *start_point*, *end_point*, or *pivot_point*, set the appropriate grid staggering type for *field2d*. By default, the mass grid is used. The options are: - 'm': Use the mass grid (default). - 'u': Use the same staggered grid as the u wind component, which has a staggered west_east (x) dimension. - 'v': Use the same staggered grid as the v wind component, which has a staggered south_north (y) dimension. projection (:class:`wrf.WrfProj`, optional): The map projection object to use when working with latitude, longitude coordinates, and must be specified if *wrfin* is None. Default is None. ll_point (:class:`wrf.CoordPair`, sequence of :class:`wrf.CoordPair`, \ optional): The lower left latitude, longitude point for your domain, and must be specified if *wrfin* is None. If the domain is a moving nest, this should be a sequence of :class:`wrf.CoordPair`. Default is None. Returns: :class:`wrf.CoordPair` or sequence: The coordinate pair(s) in x,y grid coordinates. """ if (wrfin is None and (projection is None or ll_point is None)): raise ValueError ("'wrfin' parameter or " "'projection' and 'll_point' parameters " "are required") lat, lon = pairs_to_latlon(pairs) if wrfin is not None: xy_vals = _ll_to_xy(lat, lon, wrfin=wrfin, timeidx=timeidx, squeeze=True, meta=False, stagger=stagger, as_int=True) else: map_proj = projection.map_proj if map_proj == ProjectionTypes.LAT_LON: pole_lat = projection.pole_lat pole_lon = projection.pole_lon latinc = ((projection.dy*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS) loninc = ((projection.dx*360.0)/2.0 / Constants.PI/Constants.WRF_EARTH_RADIUS) else: pole_lat = 90.0 pole_lon = 0.0 latinc = 0.0 loninc = 0.0 ll_lat, ll_lon = pairs_to_latlon(ll_point) xy_vals = _ll_to_xy(lat, lon, meta=False, squeeze=True, as_int=True, map_proj=projection.map_proj, truelat1=projection.truelat1, truelat2=projection.truelat2, stand_lon=projection.stand_lon, ref_lat=ll_lat, ref_lon=ll_lon, pole_lat=pole_lat, pole_lon=pole_lon, known_x=0, known_y=0, dx=projection.dx, dy=projection.dy, latinc=latinc, loninc=loninc) xy_vals = xy_vals.squeeze() if xy_vals.ndim == 1: return CoordPair(x=xy_vals[0], y=xy_vals[1]) else: return [CoordPair(x=xy_vals[0,i], y=xy_vals[1,i]) for i in py3range(xy_vals.shape[1])]