Small optimisation to 'nearest' function

Author: evanmason

make_eddy_track_AVISO.py

@@ -161,6 +161,8 @@ def set_index_padding(self, pad=2):
         '''
         print '--- Setting padding indices with pad=%s' %pad
 
+        self.pad = pad
+        
         def get_str(thestr, pad):
             '''
             Get start indices for pad
@@ -238,7 +240,7 @@ def nearest_point(self, lon, lat):
         '''
         Get indices to a point (lon, lat) in the grid
         '''
-        i, j = eddy_tracker.nearest(lon, lat, self._lon, self._lat)
+        i, j = eddy_tracker.nearest(lon, lat, self._lon, self._lat, self._lon.shape)
         return i, j
 
 
@@ -505,6 +507,8 @@ def __init__(self, AVISO_file, lonmin, lonmax, latmin, latmax, with_pad=True, us
         self.make_gridmask(with_pad, use_maskoceans).uvmask()
         self.get_AVISO_f_pm_pn()
         self.set_u_v_eke()
+        pad2 = 2 * self.pad
+        self.shape = (self.f().shape[0] - pad2, self.f().shape[1] - pad2)
 
 
     def get_AVISO_data(self, AVISO_file):
@@ -807,12 +811,12 @@ def pcol_2dxy(self, x, y):
     #savedir = '/marula/emason/aviso_eddy_tracking/pablo_exp/'
     #savedir = '/marula/emason/aviso_eddy_tracking/new_AVISO_test/'
     #savedir = '/marula/emason/aviso_eddy_tracking/new_AVISO_SUBSAMP-3days/'
-    #savedir = '/marula/emason/aviso_eddy_tracking/junk/'
+    savedir = '/marula/emason/aviso_eddy_tracking/junk/'
     #savedir = '/marula/emason/aviso_eddy_tracking/new_AVISO_test/BlackSea/'
     #savedir = '/marula/emason/aviso_eddy_tracking/Corridor_V3_Dec2014/'
     #savedir = '/marula/emason/aviso_eddy_tracking/CLS_test_1_acd66ba338a9/'
     #savedir = '/marula/emason/aviso_eddy_tracking/CLS_test_2/'
-    savedir = '/marula/emason/aviso_eddy_tracking/CLS_test_3/'
+    #savedir = '/marula/emason/aviso_eddy_tracking/CLS_test_3/'
     #savedir = '/path/to/save/your/outputs/'
 
 
@@ -909,10 +913,10 @@ def pcol_2dxy(self, x, y):
         #latmin = -47.
         #latmax = -24.
 
-        #lonmin = -38.     # small test
-        #lonmax = -25.
-        #latmin = 30.
-        #latmax = 38.
+        lonmin = -38.     # small test
+        lonmax = -25.
+        latmin = 30.
+        latmax = 38.
 
 
     elif the_domain in 'MedSea':

make_eddy_tracker_list_obj.py

@@ -47,7 +47,7 @@ def haversine_distance_vector(lon1, lat1, lon2, lat2):
     lat1 = np.asfortranarray(lat1.copy())
     lon2 = np.asfortranarray(lon2.copy())
     lat2 = np.asfortranarray(lat2.copy())
-    dist = np.asfortranarray(np.empty_like(lon1))
+    dist = np.asfortranarray(np.empty(lon1.shape))
     hav.haversine_distvec(lon1, lat1, lon2, lat2, dist)
     return dist
 
@@ -57,28 +57,41 @@ def newPosition(lonin, latin, angle, distance):
     Given the inputs (base lon, base lat, angle, distance) return
      the lon, lat of the new position...
     '''
-    lonin = np.asfortranarray(np.copy(lonin))
-    latin = np.asfortranarray(np.copy(latin))
-    angle = np.asfortranarray(np.copy(angle))
-    distance = np.asfortranarray(np.copy(distance))
-    lon = np.asfortranarray(np.empty_like(lonin))
-    lat = np.asfortranarray(np.empty_like(lonin))
-    hav.waypoint_vec(lonin, latin, angle, distance, lon, lat)
-    return lon, lat
+    lon = np.asfortranarray(lonin.copy())
+    lat = np.asfortranarray(latin.copy())
+    angle = np.asfortranarray(angle.copy())
+    distance = np.asfortranarray(distance.copy())
+    lonout = np.asfortranarray(np.empty(lonin.shape))
+    latout = np.asfortranarray(np.empty(lonin.shape))
+    hav.waypoint_vec(lon, lat, angle, distance, lonout, latout)
+    return lon[0], lat[0]
 
 
-def nearest(lon_pt, lat_pt, lon2d, lat2d):
+#def nearest(lon_pt, lat_pt, lon2d, lat2d):
+    #"""
+    #Return the nearest i, j point to a given lon, lat point
+    #in a lat/lon grid
+    #"""
+    #lon2d, lat2d = lon2d.copy(), lat2d.copy()
+    #lon2d -= lon_pt
+    #lat2d -= lat_pt
+    #d = np.hypot(lon2d, lat2d)
+    #j, i = np.unravel_index(d.argmin(), d.shape)
+    #return i, j
+    
+
+def nearest(lon_pt, lat_pt, lon2d, lat2d, theshape):
     """
     Return the nearest i, j point to a given lon, lat point
     in a lat/lon grid
     """
-    lon2d, lat2d = lon2d.copy(), lat2d.copy()
-    lon2d -= lon_pt
-    lat2d -= lat_pt
-    d = np.hypot(lon2d, lat2d)
-    j, i = np.unravel_index(d.argmin(), d.shape)
+    #print type(lon_pt), lon_pt
+    lon_pt += -lon2d
+    lat_pt += -lat2d
+    d = np.sqrt(lon_pt**2 + lat_pt**2)
+    j, i = np.unravel_index(d.argmin(), theshape)
     return i, j
-    
+
 
 def uniform_resample(x, y, num_fac=2, kind='linear'):
     '''
@@ -1111,17 +1124,18 @@ def get_angle(deg, ang):
             return deg - np.rad2deg(ang)
 
         grdangle = grd.angle()[j,i]
-        
+        #print type(centlon), type(centlat)
         a_lon, a_lat = newPosition(centlon, centlat, get_angle(0, grdangle), radius)
         b_lon, b_lat = newPosition(centlon, centlat, get_angle(90, grdangle), radius)
         c_lon, c_lat = newPosition(centlon, centlat, get_angle(180, grdangle), radius)
         d_lon, d_lat = newPosition(centlon, centlat, get_angle(270, grdangle), radius)
 
         # Get i,j of bounding box around eddy
-        a_i, a_j = nearest(a_lon, a_lat, grd.lon(), grd.lat())
-        b_i, b_j = nearest(b_lon, b_lat, grd.lon(), grd.lat())
-        c_i, c_j = nearest(c_lon, c_lat, grd.lon(), grd.lat())
-        d_i, d_j = nearest(d_lon, d_lat, grd.lon(), grd.lat())
+        #print grd.lon().shape, grd.lat().shape, grd.shape
+        a_i, a_j = nearest(a_lon, a_lat, grd.lon(), grd.lat(), grd.shape)
+        b_i, b_j = nearest(b_lon, b_lat, grd.lon(), grd.lat(), grd.shape)
+        c_i, c_j = nearest(c_lon, c_lat, grd.lon(), grd.lat(), grd.shape)
+        d_i, d_j = nearest(d_lon, d_lat, grd.lon(), grd.lat(), grd.shape)
 
         self.imin = np.maximum(np.min([a_i, b_i, c_i, d_i]) - 5, 0) # Must not go
         self.jmin = np.maximum(np.min([a_j, b_j, c_j, d_j]) - 5, 0) # below zero

py_eddy_tracker_classes.py

@@ -702,12 +702,16 @@ def collection_loop(CS, grd, rtime, A_list_obj, C_list_obj,
                 cx, cy = Eddy.M(contlon_e, contlat_e)
                 centlon_e, centlat_e, eddy_radius_e, aerr = fit_circle(cx, cy)
                 #print centlon_e, centlat_e
+                
                 aerr = np.atleast_1d(aerr)
                 # Filter for shape: >35% (>55%) is not an eddy for Q (SLA)
                 if np.logical_and(aerr >= 0., aerr <= Eddy.shape_err[collind]):
 
-                    # Get centroid in lon lat
+                    # Get centroid in lon lat  
+                    #print type(centlon_e), type(centlat_e), 1111
                     centlon_e, centlat_e = Eddy.M.projtran(centlon_e, centlat_e, inverse=True)
+                    # For some reason centlat_e is transformed to 'float'...
+                    centlon_e, centlat_e = np.float64(centlon_e), np.float64(centlat_e)
 
                     # Get eddy_radius_e (NOTE: if Q, we overwrite eddy_radius_e defined ~8 lines above)
                     if 'Q' in Eddy.diag_type:
@@ -729,10 +733,15 @@ def collection_loop(CS, grd, rtime, A_list_obj, C_list_obj,
 
                     if proceed0:
                         # Get indices of centroid
+                        
+                        #print type(centlon_e), type(centlat_e)
                         centi, centj = eddy_tracker.nearest(centlon_e, centlat_e,
                                                             grd.lon(),
-                                                            grd.lat())
+                                                            grd.lat(), grd.shape)
                         # If condition not True this eddy has already been detected
+                        
+                        #print type(centlon_e), type(centlat_e)
+                        
                         if 'Q' in Eddy.diag_type:
 
                             if xi[centj, centi] != Eddy.fillval:
@@ -983,7 +992,7 @@ def collection_loop(CS, grd, rtime, A_list_obj, C_list_obj,
 
                                     # Get indices of speed-based centroid
                                     centi, centj = eddy_tracker.nearest(centlon_s, centlat_s,
-                                                                        grd.lon(), grd.lat())
+                                                                        grd.lon(), grd.lat(), grd.shape)
 
                                 # Set the bounds
                                 #bounds = np.atleast_2d(np.int16([imin, imax, jmin, jmax]))