diff options
| author | Aaron LI <aaronly.me@outlook.com> | 2016-10-26 15:50:19 +0800 |
|---|---|---|
| committer | Aaron LI <aaronly.me@outlook.com> | 2016-10-26 15:50:19 +0800 |
| commit | 60259c97845f7c96645f130dae770c05ef51c193 (patch) | |
| tree | 03eb287d02af5859dbfc6977c811a3ccff40545b /fg21sim | |
| parent | 7cf0e5977e4a885595ca0931f26a5fef7dc8f82e (diff) | |
| download | fg21sim-60259c97845f7c96645f130dae770c05ef51c193.tar.bz2 | |
utils/grid.py: JIT optimize "make_grid_ellipse()"
NOTE that the original optional paramter "rotation" is now required.
Diffstat (limited to 'fg21sim')
| -rw-r--r-- | fg21sim/utils/grid.py | 25 |
1 files changed, 15 insertions, 10 deletions
diff --git a/fg21sim/utils/grid.py b/fg21sim/utils/grid.py index ca9510c..a1d7d2c 100644 --- a/fg21sim/utils/grid.py +++ b/fg21sim/utils/grid.py @@ -84,7 +84,11 @@ def make_coordinate_grid(center, size, resolution): return (lon, lat) -def make_grid_ellipse(center, size, resolution, rotation=None): +@nb.jit(nb.types.UniTuple(nb.float64[:, :], 3)( + nb.types.UniTuple(nb.float64, 2), nb.types.UniTuple(nb.float64, 2), + nb.float64, nb.float64), + nopython=True) +def make_grid_ellipse(center, size, resolution, rotation=0.0): """Make a square coordinate grid just containing the specified (rotated) ellipse. @@ -97,7 +101,7 @@ def make_grid_ellipse(center, size, resolution, rotation=None): The (major, minor) axes of the filling ellipse, unit [ degree ]. resolution : float The grid resolution, unit [ degree ]. - rotation : float, optional + rotation : float The rotation angle (unit [ degree ]) of the filling ellipse. Returns @@ -129,15 +133,16 @@ def make_grid_ellipse(center, size, resolution, rotation=None): 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) - r_radius, c_radius = np.ceil(0.5*np.array(size)/resolution).astype(np.int) - rr, cc = ellipse(r0, c0, r_radius, c_radius, shape=shape) + r0, c0 = np.floor(np.array(shape) / 2.0).astype(np.int64) + radii = np.ceil(0.5*np.array(size)/resolution).astype(np.int64) + rr, cc = ellipse(r0, c0, radii[0], radii[1], shape=shape) gridmap = np.zeros(shape) - gridmap[rr, cc] = 1.0 - if rotation is not None: - # Rotate the ellipse - gridmap = rotate_center(gridmap, angle=rotation, interp=True, - reshape=False, fill_value=0.0) + # XXX: ``numba`` only support one advanced index + for ri, ci in zip(rr, cc): + gridmap[ri, ci] = 1.0 + # Rotate the ellipse about the grid center + gridmap = rotate_center(gridmap, angle=rotation, interp=True, + reshape=False, fill_value=0.0) return (lon, lat, gridmap) |