def _imports():
global copy, izip, np, pylab, Polygon, Basemap, os
import numpy as np
import pylab
import os
from matplotlib.patches import Polygon
from mpl_toolkits.basemap import Basemap
from itertools import izip
import copy
return
if __name__ != "__main__":
_imports()
def _default_value(value, default):
if value is not None: return value
else: return default
def _slice_union(*slices):
slc_start = min( s.start for s in slices )
slc_stop = max( s.stop for s in slices )
return slice(slc_start, slc_stop)
def _slice_intersection(*slices):
slc_start = max( s.start for s in slices )
slc_stop = min( s.stop for s in slices )
return slice(slc_start, max(slc_start, slc_end))
[docs]class ExperimentGrid(object):
"""
Represents an ARPS Arakawa C grid, with provisions for subsets.
Args:
grid_size (tuple): Tuple of integers (NX, NY) representing the number of grid points in each dimension of the
grid.
grid_spacing (tuple): Tuple of floats (DX, DY) representing the grid spacing in meters in each dimension of the
grid.
bounds (tuple): Tuple of slice objects (SLICEX, SLICEY) representing the bounds in each direction of a subset of
the grid. Default is (slice(None), slice(None))
mpi (tuple): Tuple of integers (NPROCX, NPROCY) representing the MPI configuration for this grid. Default is (1,
1), meaning no MPI.
state_list (list): An optional list of strings containing the two-letter postal codes for each state visible in
the domain. This is used for optimizing the drawPolitical() method.
projection (str): Map projection to use. Only tested value is 'lcc' for Lambert conic conformal.
resolution (str): Resolution of the mapping data to load. Acceptable values are 'c' (crude), 'l' (low), and 'i'
(intermediate). Recommended value is 'i'.
lat_0: Center latitude for the Lambert projection (degrees N)
lat_1: True latitude 1 for the Lambert projection (degrees N)
lat_2: True latitude 2 for the Lambert projection (degrees N)
lon_0: True longitude for the Lambert projection (degrees E)
ctr_lat: Center latitude for the domain (degrees N)
ctr_lon: Center longitude for the domain (degrees E)
"""
def __init__(self, **kwargs):
"""
Constructor for the ExperimentGrid class
"""
self._nx, self._ny = kwargs['grid_size']
self._dx, self._dy = kwargs['grid_spacing']
del kwargs['grid_size']
del kwargs['grid_spacing']
bounds = kwargs.get('bounds', (slice(None), slice(None)))
if bounds is None:
self._bound_x, self._bound_y = slice(0, self._nx), slice(0, self._ny)
else:
bound_x, bound_y = bounds
self._bound_x = slice(_default_value(bound_x.start, 0), _default_value(bound_x.stop, self._nx))
self._bound_y = slice(_default_value(bound_y.start, 0), _default_value(bound_y.stop, self._ny))
self._mpi_x, self._mpi_y = kwargs.get('mpi', (1, 1))
if 'bounds' in kwargs:
del kwargs['bounds']
if 'mpi' in kwargs:
del kwargs['mpi']
kwargs['rsphere'] = 6371000
kwargs['width'] = self._dx * self._nx
kwargs['height'] = self._dy * self._ny
if 'ctr_lat' in kwargs:
self._ctr_lat = kwargs['ctr_lat']
del kwargs['ctr_lat']
else:
self._ctr_lat = kwargs['lat_0']
if 'ctr_lon' in kwargs:
self._ctr_lon = kwargs['ctr_lon']
del kwargs['ctr_lon']
else:
self._ctr_lon = kwargs['lon_0']
if 'state_list' in kwargs:
self._state_list = kwargs['state_list']
del kwargs['state_list']
else:
self._state_list = None
self._proj = copy.copy(kwargs)
map = Basemap(**kwargs)
ctr_x, ctr_y = map(self._ctr_lon, self._ctr_lat)
move_x = ctr_x - kwargs['width'] / 2
move_y = ctr_y - kwargs['height'] / 2
lower_bound_x = _default_value(self._bound_x.start, 0)
lower_bound_y = _default_value(self._bound_y.start, 0)
upper_bound_x = _default_value(self._bound_x.stop, self._nx) - 1.4
upper_bound_y = _default_value(self._bound_y.stop, self._ny) - 1.4
llcrnrlon, llcrnrlat = map(move_x + self._dx * lower_bound_x, move_y + self._dy * lower_bound_y, inverse=True)
urcrnrlon, urcrnrlat = map(move_x + self._dx * upper_bound_x, move_y + self._dy * upper_bound_y, inverse=True)
del kwargs['width']
del kwargs['height']
kwargs['llcrnrlat'] = llcrnrlat
kwargs['llcrnrlon'] = llcrnrlon
kwargs['urcrnrlat'] = urcrnrlat
kwargs['urcrnrlon'] = urcrnrlon
self._map = Basemap(**kwargs)
return
[docs] def get_xy(self, ij=None, pcm=False):
"""
Get x, y coordinates of the center each grid point or a specific grid point.
Args:
ij (tuple): Get the coordinates of the grid point specified by (i, j). Optional.
pcm (bool): Get the coordinates of the corners of the cells, rather than the centers.
Returns:
If ij is not specified, a tuple of 2-dimensional numpy arrays (x, y). Axes are ordered (NX, NY). If ij is
specified, returns a tuple of floats. Both are in meters.
"""
diff = 0 if pcm else 1
off_x = self._dx / 2.
off_y = self._dy / 2.
xs, ys = np.meshgrid(self._dx * (np.arange(self._nx) - diff), self._dy * (np.arange(self._ny) - diff))
xs = xs[self.get_bounds(pcm=pcm)]
xs = xs - xs[0, 0] - off_x
ys = ys[self.get_bounds(pcm=pcm)]
ys = ys - ys[0, 0] - off_y
if ij is None:
return xs, ys
else:
x, y = ij
return xs[x, y], ys[x, y]
[docs] def get_bounds(self, pcm=False):
"""
Get the grid point bounds for this subdomain
Args:
pcm (bool): Whether to assume this will be used in matplotlib's `pcolormesh()` function. Default is False.
Returns:
The bounds as a tuple of slice objects (SLICEY, SLICEX)
"""
bnd_y, bnd_x = self._bound_y, self._bound_x
cores = self.get_cores()
x_slc_core, y_slc_core = [], []
for core in cores:
x_slc, y_slc = self.core_bounds(*core)
x_slc_core.append(x_slc)
y_slc_core.append(y_slc)
cores_bnd_x = _slice_union(*x_slc_core)
cores_bnd_y = _slice_union(*y_slc_core)
bnd_y = slice(bnd_y.start - cores_bnd_y.start, bnd_y.stop - cores_bnd_y.start)
bnd_x = slice(bnd_x.start - cores_bnd_x.start, bnd_x.stop - cores_bnd_x.start)
expand = 1 if pcm else 0
bnd_y = slice(bnd_y.start, bnd_y.stop + expand)
bnd_x = slice(bnd_x.start, bnd_x.stop + expand)
return bnd_y, bnd_x
[docs] def get_center(self):
"""
Get the center lat and lon in the domain.
Returns:
A tuple of floats (LAT, LON) in degrees.
"""
return self._ctr_lat, self._ctr_lon
[docs] def get_boundary_coords(self):
"""
Get the latitudes and longitudes of the boundary of this domain (useful for drawing this domain using another
projection).
Returns:
A tuple of 1-dimensional numpy arrays containing the latitudes and longitudes of the domain boundary in
degrees.
"""
def _create_boundary(nx, ny, lb_x=0, lb_y=0):
boundary_x = lb_x * np.ones((2 * (nx + ny),))
boundary_y = lb_y * np.ones((2 * (nx + ny),))
boundary_x[:nx] = np.arange(lb_x, lb_x + nx)
boundary_x[nx:(nx + ny)] = lb_x + nx
boundary_x[(nx + ny):(2 * nx + ny)] = np.arange(lb_x + nx, lb_x, -1)
boundary_y[nx:(nx + ny)] = np.arange(lb_y, lb_y + ny)
boundary_y[(nx + ny):(2 * nx + ny)] = lb_y + ny
boundary_y[(2 * nx + ny):(2 * (nx + ny))] = np.arange(lb_y + ny, lb_y, -1)
return boundary_x, boundary_y
if self._bound_x != slice(None) and self._bound_y != slice(None):
subset_nx = self._bound_x.stop - self._bound_x.start
subset_ny = self._bound_y.stop - self._bound_y.start
boundary_x, boundary_y = _create_boundary(subset_nx, subset_ny, self._bound_x.start, self._bound_y.start)
else:
boundary_x, boundary_y = _create_boundary(self._nx, self._ny)
boundary_lons, boundary_lats = self(self._dx * boundary_x, self._dy * boundary_y, inverse=True)
return boundary_lats, boundary_lons
[docs] def get_grid_spacing(self):
"""
Get the grid spacing.
Returns:
The grid spacing in meters as a tuple (DX, DY).
"""
return self._dx, self._dy
[docs] def get_grid_size(self):
"""
Get the number of grid points.
Returns:
The number of grid points in the full domain as a tuple (NX, NY).
"""
nx, ny = self._nx, self._ny
return nx, ny
[docs] def get_width_height(self, override=False):
"""
Get the width and height of the domain.
Returns:
The width and height of the domain in meters of the subsetted domain as a tuple (WIDTHX, WIDTHY)
"""
if self._bound_x == slice(None) and self._bound_y == slice(None) or override:
return self._dx * self._nx, self._dy * self._ny
else:
return self._dx * (self._bound_x.stop - self._bound_x.start), self._dy * (self._bound_y.stop - self._bound_y.start)
[docs] def core_bounds(self, icore, jcore):
"""
Get the grid point bounds for a particular core.
Args:
icore (int): i-index of the core on the domain.
jcore (int): j-index of the core on the domain.
Returns:
A tuple of slice objects (BOUNDX, BOUNDY).
"""
nx_core = (self._nx - 3) / self._mpi_x + 3
ny_core = (self._ny - 3) / self._mpi_y + 3
lower_bound_x = (nx_core - 3) * icore
upper_bound_x = (nx_core - 3) * (icore + 1) + 4
lower_bound_y = (ny_core - 3) * jcore
upper_bound_y = (ny_core - 3) * (jcore + 1) + 4
return slice(lower_bound_x, upper_bound_x), slice(lower_bound_y, upper_bound_y)
[docs] def get_cores(self):
"""
Get a list of cores that are in the domain subset.
Returns:
A list of tuples containing the i and j indices of the cores in the domain subset.
"""
lower_bound_x = _default_value(self._bound_x.start, 0)
lower_bound_y = _default_value(self._bound_y.start, 0)
upper_bound_x = _default_value(self._bound_x.stop, self._nx) - 1
upper_bound_y = _default_value(self._bound_y.stop, self._ny) - 1
cores = []
for icore in xrange(self._mpi_x):
for jcore in xrange(self._mpi_y):
core_bounds_x, core_bounds_y = self.core_bounds(icore, jcore)
out_of_bounds_x = (upper_bound_x < core_bounds_x.start)
out_of_bounds_x |= (lower_bound_x > core_bounds_x.stop)
out_of_bounds_y = (upper_bound_y < core_bounds_y.start)
out_of_bounds_y |= (lower_bound_y > core_bounds_y.stop)
if not(out_of_bounds_x or out_of_bounds_y):
cores.append((icore, jcore))
return cores
[docs] def draw_political(self, scale_len=0, overlays=[], color='k', lw=1):
"""
Draw political boundaries on the current matplotlib axes.
Args:
scale_len (float): Length of a scale bar in km to draw in the bottom-left corner of the plot. Default is 0,
meaning no scale.
overlays (list): A list of strings pointing to shapefiles to plot as overlays.
color (str): Color in which to draw the political boundaries. Defaults to 'k' (black).
lw (float): Base line width for the political boundaries. Country boundaries are drawn with width `lw`,
states with width `1.5 * lw`, and counties are drawn with width `0.5 * lw`
"""
self._map.drawcountries(linewidth=lw, color=color)
self._map.drawcoastlines(linewidth=lw, color=color)
shapefile_dir = os.path.join(os.path.dirname(__file__), '..', 'data')
if not hasattr(self._map, 'counties'):
self._map.readshapefile(shapefile_dir + "/counties/countyp020", 'counties', linewidth=(lw * 0.5),
color=color)
else:
for county, data in izip(self._map.counties_info, self._map.counties):
if not self._state_list or county['STATE'] in self._state_list:
pylab.gca().add_patch(Polygon(data, ec=color, fc='none', linewidth=(lw * 0.5)))
if not hasattr(self._map, 'states'):
self._map.readshapefile(shapefile_dir + "/states/tl_2011_us_state", 'states', linewidth=(lw * 1.5),
color=color)
else:
for state, data in izip(self._map.states_info, self._map.states):
if not self._state_list or state['STUSPS'] in self._state_list:
pylab.gca().add_patch(Polygon(data, ec=color, fc='none', linewidth=(lw * 1.5)))
for overlay in overlays:
self._map.readshapefile(overlay, overlay, linewidth=(lw * 0.5), color=color)
if scale_len > 0:
offset_x, offset_y = 0.03, 0.09
ax_trans = pylab.gca().transData + pylab.gca().transAxes.inverted()
coords_x, coords_y = ax_trans.transform(1000 * np.array([[0, scale_len], [0, 0]]))
scale_bar = pylab.plot(coords_x + offset_x, coords_y + offset_y, 'k-', lw=(lw * 2.5),
transform=pylab.gca().transAxes)
scale_txt = pylab.text(coords_x.mean() + offset_x, coords_y.mean() + offset_y - 0.025,
"%d km" % scale_len, color='k', ha='center', va='top', transform=pylab.gca().transAxes)
txt_extent = scale_txt.get_window_extent(renderer=pylab.gcf().canvas.get_renderer()).get_points()
txt_extent_ax_x, txt_extent_ax_y = pylab.gca().transAxes.inverted().transform(txt_extent)
if txt_extent_ax_x[0] < 0.01:
fudge = 0.01 - txt_extent_ax_x[0]
scale_bar[0].set_xdata(scale_bar[0].get_xdata() + fudge)
tp_x, tp_y = scale_txt.get_position()
scale_txt.set_position((tp_x + fudge, tp_y))
return
[docs] def __call__(self, *args, **kwargs):
"""
Convert between latitude, longitude coordinates and x, y coordinates on the domain.
The only meaningful keyword argument is `inverse`, which should be a boolean specifying whether to convert from
latitudes and longitudes to x, y (False) or vice-versa (True). Default is False.
If inverse is False, *args is two numpy arrays of longitudes and latitudes in degrees. If inverse is True, *args
is two numpy arrays of x and y coordinates in meters.
Returns:
A tuple of numpy arrays of either x and y coordinates in meters or longitudes and latitudes in degrees.
"""
return self._map(*args, **kwargs)
@staticmethod
[docs] def from_file(fobj, projection='lcc', resolution='l', mpi=(1, 1), bounds=None):
"""
Construct an ExperimentGrid object from information found in a history file. Should work for WRF and ARPS.
Args:
fobj: File object (such as a NetCDF or HDF file) from which to take the information.
projection (str): Map projection to use. Currently, the only supported value is 'lcc' for Lambert Conic
Conformal.
resolution (str): Resolution of the mapping data to load. Acceptable values are 'c' (crude), 'l' (low), and
'i' (intermediate). Recommended value is 'i'.
mpi (tuple): Tuple of integers (NPROCX, NPROCY) representing the MPI configuration for this grid. Default is
(1, 1), meaning no MPI.
bounds (tuple): Tuple of slice objects (SLICEX, SLICEY) representing the bounds in each direction of a
subset of the grid. Default is (slice(None), slice(None))
Returns:
ExperimentGrid: An ExperimentGrid object with parameters matching those in the input file.
"""
def _get_attrs(fobj, tries):
for var, attr in tries:
if var is None:
if hasattr(fobj, attr):
return getattr(fobj, attr)
else:
if var in fobj.variables and hasattr(fobj.variables[var], attr):
return getattr(fobj.variables[var], attr)
nx = _get_attrs(fobj, [('x_stag', 'shape'), ('XLONG', 'shape'), ('x', 'shape')])
ny = _get_attrs(fobj, [('y_stag', 'shape'), ('XLAT', 'shape'), ('y', 'shape')])
if len(nx) == 1:
nx = nx[0]
else:
nx = nx[-1]
if len(ny) == 1:
ny = ny[0]
else:
ny = ny[-2]
dx = _get_attrs(fobj, [(None, 'DX')])
dy = _get_attrs(fobj, [(None, 'DY')])
ctr_lat = _get_attrs(fobj, [(None, 'CTRLAT'), (None, 'CEN_LAT')])
ctr_lon = _get_attrs(fobj, [(None, 'CTRLON'), (None, 'CEN_LON')])
lat_0 = _get_attrs(fobj, [(None, 'latitude_of_projection_origin'), (None, 'CEN_LAT')])
lon_0 = _get_attrs(fobj, [(None, 'longitude_of_central_meridian'), (None, 'STAND_LON')])
lat_1 = _get_attrs(fobj, [(None, 'TRUELAT1')])
lat_2 = _get_attrs(fobj, [(None, 'TRUELAT2')])
return ExperimentGrid(projection=projection, resolution=resolution,
grid_size=(nx, ny), grid_spacing=(dx, dy),
ctr_lat=ctr_lat, ctr_lon=ctr_lon,
lat_0=lat_0, lon_0=lon_0, lat_1=lat_1, lat_2=lat_2,
mpi=mpi, bounds=bounds)
if __name__ == "__main__":
import matplotlib
import mapproj
matplotlib.use('agg')
_imports()
grid = goshen_1km_grid()
# grid.drawPolitical()
# pylab.savefig("grid_test.png")
ctr_lat, ctr_lon = np.array([41.61975]), np.array([-104.34843])
nx, ny = 255, 255
dx, dy = 1000., 1000.
mapproj.setmapr(2, 1., np.array([30., 60.]), -107.344)
ctr_x, ctr_y = mapproj.lltoxy(ctr_lat, ctr_lon)
sw_x, sw_y = ctr_x - nx * dx / 2, ctr_y - ny * dy / 2
mapproj.setorig(1, sw_x, sw_y)
print grid(ctr_lon, ctr_lat)
print mapproj.lltoxy(ctr_lat, ctr_lon)