diff --git a/modules/util/lon_lat_grid.py b/modules/util/lon_lat_grid.py index 73cffdf3d19b837a75767483fbded46e9df923d6..9b910a0a4d628726e3fdeb3f787cd365ff02bc5b 100644 --- a/modules/util/lon_lat_grid.py +++ b/modules/util/lon_lat_grid.py @@ -21,13 +21,15 @@ class LonLatGrid: # Incoming longitude must be in range: 0 - 360 degrees # Can have NaN for off Earth grid points (these are handled internally). # closeness_threshold: if < distance of located_point to target, return off grid - def __init__(self, grd_lons, grd_lats, closeness_threshold=2000): + def __init__(self, grd_lons, grd_lats, closeness_threshold=5000, reduce=1, leaf_size=40): if grd_lons.shape != grd_lats.shape: raise MyGenericException('incoming lons,lats must have same shape') self.grd_lons = grd_lons self.grd_lats = grd_lats - self.shape = grd_lons.shape + self.grd_lons = grd_lons[::reduce, ::reduce] + self.grd_lats = grd_lats[::reduce, ::reduce] + self.shape = self.grd_lons.shape self.lenx = self.shape[1] self.leny = self.shape[0] self.fnc_dct = {} @@ -49,8 +51,9 @@ class LonLatGrid: points = np.stack([flons, flats], axis=1) - self.kd = BallTree(np.deg2rad(points), leaf_size=500, metric='haversine') + self.kd = BallTree(np.deg2rad(points), leaf_size=leaf_size, metric='haversine') self.closeness_threshold = closeness_threshold + self.reduce = reduce # locate nearest neighbor for incoming target in earth coordinates (Should not have NaNs) # lons, lats can be flat or 2D @@ -68,7 +71,10 @@ class LonLatGrid: xy = np.stack([lons, lats], axis=1) - dist, indices = self.kd.query(np.deg2rad(xy)) + # Assumes nn for a query pont returned in order of increasing distance + dist, indices = self.kd.query(np.deg2rad(xy), k=1) + dist = dist[:, 0] + indices = indices[:, 0] dist *= 6370000 # convert unit radius to meters valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold) @@ -90,9 +96,10 @@ class LonLatGrid: xy = np.stack([lons, lats], axis=1) - dist, indices = self.kd.query(np.deg2rad(xy)) - dist = dist[0] - indices = indices[0] + # Assumes nn for a query pont returned in order of increasing distance + dist, indices = self.kd.query(np.deg2rad(xy), k=1) + dist = dist[:, 0] + indices = indices[:, 0] dist *= 6370000 # convert unit radius to meters valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold) @@ -100,6 +107,9 @@ class LonLatGrid: ll = np.where(indices >= 0, (indices / self.lenx).astype(np.int32), indices) cc = np.where(indices >= 0, indices % self.lenx, indices) + cc *= self.reduce + ll *= self.reduce + return cc, ll # bi-linear interpolation of a function on the grid, the range of which is grd_zvals for target lons,lats.