From 9cf9a11460e09ea5b39051c7a5735fa992831d95 Mon Sep 17 00:00:00 2001 From: Aaron LI Date: Tue, 25 Oct 2016 21:21:05 +0800 Subject: grid.py: Override "make_coordinate_grid()" using JIT-optimized version Remove the original "make_coordinate_grid()" function, and rename the "make_coordinate_grid_fast()" to override the former one. --- fg21sim/utils/grid.py | 48 ++---------------------------------------------- 1 file changed, 2 insertions(+), 46 deletions(-) (limited to 'fg21sim') diff --git a/fg21sim/utils/grid.py b/fg21sim/utils/grid.py index bf4adca..7a81ea2 100644 --- a/fg21sim/utils/grid.py +++ b/fg21sim/utils/grid.py @@ -30,56 +30,12 @@ def _wrap_latitudes(lat): return lat -def make_coordinate_grid(center, size, resolution): - """Make a rectangle, Cartesian coordinate grid. - - Parameters - ---------- - center : 2-float tuple - Center coordinate (longitude, latitude) of the grid, - with longitude [0, 360) degree, latitude [-90, 90] degree. - size : float, or 2-float tuple - The sizes (size_lon, size_lat) of the grid along the longitude - and latitude directions. If only one float specified, then the - grid is square. - resolution : float - The grid resolution, unit [ degree ]. - - Returns - ------- - lon : 2D `~numpy.ndarray` - The array with elements representing the longitudes of each grid - pixel. The array is odd-sized, with the input center locating at - the exact grid central pixel. - Also, the longitudes are fixed to be in the valid range [0, 360). - lat : 2D `~numpy.ndarray` - The array with elements representing the latitudes of each grid - pixel. - Also, the latitudes are fixed to be in the valid range [-90, 90]. - """ - lon0, lat0 = center - try: - size_lon, size_lat = size - except (TypeError, ValueError): - size_lon = size_lat = size - # Half number of pixels (excluding the center) - hn_lon = np.ceil(0.5*size_lon / resolution).astype(np.int) - hn_lat = np.ceil(0.5*size_lat / resolution).astype(np.int) - idx_lon = lon0 + np.arange(-hn_lon, hn_lon+1) * resolution - idx_lat = lat0 + np.arange(-hn_lat, hn_lat+1) * resolution - # Fix the longitudes and latitudes to be in the valid ranges - idx_lon = _wrap_longitudes(idx_lon) - idx_lat = _wrap_latitudes(idx_lat) - lon, lat = np.meshgrid(idx_lon, idx_lat) - return (lon, lat) - - @nb.jit(nb.types.UniTuple(nb.float64[:, :], 2)( nb.types.UniTuple(nb.float64, 2), nb.types.UniTuple(nb.float64, 2), nb.float64), nopython=True) -def make_coordinate_grid_fast(center, size, resolution): +def make_coordinate_grid(center, size, resolution): """Make a rectangle, Cartesian coordinate grid. This is the ``numba.jit`` optimized version of ``make_coordinate_grid``. @@ -170,7 +126,7 @@ def make_grid_ellipse(center, size, resolution, rotation=None): """ size_major = max(size) size = (size_major, size_major) - lon, lat = make_coordinate_grid_fast(center, size, resolution) + lon, lat = make_coordinate_grid(center, size, resolution) shape = lon.shape # Fill the ellipse into the grid r0, c0 = np.floor(np.array(shape) / 2.0).astype(np.int) -- cgit v1.2.2