snapshot...

modules/util/lon_lat_grid.py

@@ -20,7 +20,8 @@ class LonLatGrid:
     # grd_lons, grd_lats: the longitude, latitude of each grid point (must be 2D grids), must have same shape
     # Incoming longitude must be in range: 0 - 360 degrees
     # Can have NaN for off Earth grid points (these are handled internally).
-    def __init__(self, grd_lons, grd_lats):
+    # closeness_threshold: if < distance of located_point to target, return off grid
+    def __init__(self, grd_lons, grd_lats, closeness_threshold=2000):
         if grd_lons.shape != grd_lats.shape:
             raise MyGenericException('incoming lons,lats must have same shape')
 
@@ -47,13 +48,13 @@ class LonLatGrid:
         points = np.stack([flons, flats], axis=1)
 
         self.kd = BallTree(np.deg2rad(points), leaf_size=500, metric='haversine')
+        self.closeness_threshold = closeness_threshold
 
     # locate nearest neighbor for incoming target in earth coordinates (Should not have NaNs)
     # lons, lats can be flat or 2D
-    # closeness_threshold: if < distance of located_point to target, return off grid
     # incoming longitude must be in range: 0 - 360 degrees
     # returns NN indexes relative to the grid determined in the ctr.
-    def value_to_index(self, lons, lats, closeness_threshold=2000):
+    def value_to_index(self, lons, lats):
         if lons.shape != lats.shape:
             raise MyGenericException('incoming lons,lats must have same shape')
 
@@ -68,11 +69,33 @@ class LonLatGrid:
         dist, indices = self.kd.query(np.deg2rad(xy))
         dist *= 6370000  # convert unit radius to meters
 
-        valid = (indices < self.map_indexes.size) & (dist < closeness_threshold)
+        valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold)
         indices = indices[valid]
 
         return self.map_indexes[indices], indices
 
+    def earth_to_lc_s(self, lons, lats):
+        if lons.shape != lats.shape:
+            raise MyGenericException('incoming lons,lats must have same shape')
+
+        if len(lons) == 0 or len(lats) == 0:
+            return None, None
+
+        lons = lons.flatten()
+        lats = lats.flatten()
+
+        xy = np.stack([lons, lats], axis=1)
+
+        dist, indices = self.kd.query(np.deg2rad(xy))
+        dist *= 6370000  # convert unit radius to meters
+
+        valid = (indices < self.map_indexes.size) & (dist < self.closeness_threshold)
+        indices = np.where(valid, indices, -1)
+        ll = np.where(indices >= 0, (indices / self.lenx).astype(np.int32), indices)
+        cc = np.where(indices >= 0, indices % self.lenx, indices)
+
+        return cc, ll
+
     # bi-linear interpolation of a function on the grid, the range of which is grd_zvals for target lons,lats.
     # returns the interpolated values at each target point.
     # target lons, lats should not have NaNs, and lons must be in range 0 - 360 degrees