Modified pickling procedures

evanmason · evanmason · commit 606f55e65668 · 2015-03-06T17:02:01.000+01:00
diff --git a/make_eddy_track_AVISO.py b/make_eddy_track_AVISO.py
@@ -30,8 +30,8 @@
 Scroll down to line ~640 to get started
 ===============================================================================
 """
-#from matplotlib import use as mpl_use
-#mpl_use('Agg')
+from matplotlib import use as mpl_use
+mpl_use('Agg')
 import sys
 import glob as glob
 from py_eddy_tracker_classes import plt, np, dt, Dataset, time, \
@@ -47,7 +47,7 @@
 from mpl_toolkits.basemap import Basemap
 import yaml
 from datetime import datetime
-import pickle
+import cPickle as pickle
 
 
 def timeit(method):
@@ -538,6 +538,21 @@ def __init__(self, AVISO_FILE, THE_DOMAIN, PRODUCT,
         pad2 = 2 * self.pad
         self.shape = (self.f().shape[0] - pad2, self.f().shape[1] - pad2)
     
+    
+    def __getstate__(self):
+        """
+        Needed for Pickle
+        """
+        print '--- removing unwanted attributes'
+        pops = ('Mx', 'My', '_f', '_angle', '_dx', '_dy', '_gof', '_lon',
+                '_lat', '_pm', '_pn', '_umask', '_vmask', 'eke', 'sla',
+                'mask', 'slacopy', 'sla_coeffs', 'uspd_coeffs')
+        result = self.__dict__.copy()
+        for pop in pops:
+            result.pop(pop)
+        return result
+    
+    
     #@timeit
     def get_AVISO_data(self, AVISO_FILE):
         """
@@ -638,16 +653,6 @@ def set_mask(self, sla):
     
     
     
-    def set_global_mask(self):
-        """
-        """
-        if 'Global' in self.THE_DOMAIN:
-            labels = ndimage.label(np.logical_not(self.sla.mask))[0]
-            
-        else:
-            return
-        return
-    
     
     def fillmask(self, x, mask):
         """
@@ -844,7 +849,6 @@ def set_interp_coeffs(self, sla, uspd):
     DATE_END = config['DATE_END'] = config['DOMAIN']['DATE_END']
     
     AVISO_DT14 = config['AVISO']['AVISO_DT14']
-    AVISO_FILES = config['AVISO']['AVISO_FILES']
     if AVISO_DT14:
         PRODUCT = 'AVISO_DT14'
         AVISO_DT14_SUBSAMP = config['AVISO']['AVISO_DT14_SUBSAMP']
@@ -1040,7 +1044,6 @@ def set_interp_coeffs(self, sla, uspd):
             A_eddy.reset_holding_variables()
             C_eddy.reset_holding_variables()
             
-            #grdmask = grd.mask()[j0:j1,i0:i1]
             
             sla = sla_grd.get_AVISO_data(AVISO_FILE)
             sla_grd.set_mask(sla).uvmask()
@@ -1204,18 +1207,16 @@ def set_interp_coeffs(self, sla, uspd):
                                    A_list_obj=None, C_list_obj=C_eddy,
                                    sign_type=C_eddy.SIGN_TYPE, VERBOSE=C_eddy.VERBOSE)
             
-            
-            
                 
             ymd_str = ''.join((str(yr), str(mo).zfill(2), str(da).zfill(2)))
             
             
             # Test pickling
             with open("".join((SAVE_DIR, 'A_eddy_%s.pkl' % ymd_str)), 'wb') as save_pickle:
-                pickle.dump(A_eddy, save_pickle)
+                pickle.dump(A_eddy, save_pickle, 2)
            
             with open("".join((SAVE_DIR, 'C_eddy_%s.pkl' % ymd_str)), 'wb') as save_pickle:
-                pickle.dump(C_eddy, save_pickle)
+                pickle.dump(C_eddy, save_pickle, 2)
     
             ## Unpickle
             #with open('C_eddy.pkl', 'rb') as load_pickle:
@@ -1276,8 +1277,6 @@ def set_interp_coeffs(self, sla, uspd):
             
             if SAVE_FIGURES: # Make figures for animations
                 
-                tit = ''.join((str(yr), str(mo).zfill(2), str(da).zfill(2)))
-                
                 if 'Q' in DIAGNOSTIC_TYPE:
                     anim_figure(A_eddy, C_eddy, Mx, My, plt.cm.RdBu_r, rtime, DIAGNOSTIC_TYPE, 
                                 SAVE_DIR, 'Q-parameter ' + ymd_str, animax, animax_cbar,
diff --git a/make_eddy_tracker_list_obj.py b/make_eddy_tracker_list_obj.py
@@ -903,15 +903,7 @@ def write2netcdf(self, rtime):
         return self
         
         
-    #def get_distances(self, centlon, centlat):
-        #"""
-        #Return all distances between the current centroid *centlon*,
-        #*centlat* and arrays *self.old_lon*, *self.old_lat*
-        #"""
-        #clon = np.full_like(self.old_lon, centlon)
-        #clat = np.full_like(self.old_lon, centlat)
-        #return haversine_distance_vector(clon, clat,
-                                         #self.old_lon, self.old_lat)
+
 
 
     def insert_at_index(self, xarr, ind, x):
@@ -934,103 +926,7 @@ def insert_at_index(self, xarr, ind, x):
         setattr(self, xarr, val)
         
         
-        '''
-        if strcompare('new_lon', xarr):
-            try:
-                self.new_lon[ind] = x
-            except:
-                self.new_lon.extend([0] *
-                    (ind - len(self.new_lon) + 1))
-                self.new_lon[ind] = x
-        elif strcompare('new_lat', xarr):
-            try:
-                self.new_lat[ind] = x
-            except:
-                self.new_lat.extend([0] *
-                    (ind - len(self.new_lat) + 1))
-                self.new_lat[ind] = x
-        elif strcompare('new_radii_s', xarr):
-            try:
-                self.new_radii_s[ind] = x
-            except:
-                self.new_radii_s.extend([0] *
-                    (ind - len(self.new_radii_s) + 1))
-                self.new_radii_s[ind] = x
-        elif strcompare('new_radii_e', xarr):
-            try:
-                self.new_radii_e[ind] = x
-            except:
-                self.new_radii_e.extend([0] *
-                    (ind - len(self.new_radii_e) + 1))
-                self.new_radii_e[ind] = x
-        elif strcompare('new_amp', xarr):
-            try:
-                self.new_amp[ind] = x
-            except:
-                self.new_amp.extend([0] *
-                    (ind - len(self.new_amp) + 1))
-                self.new_amp[ind] = x
-        elif strcompare('new_uavg', xarr):
-            try:
-                self.new_uavg[ind] = x
-            except:
-                self.new_uavg.extend([0] *
-                    (ind - len(self.new_uavg) + 1))
-                self.new_uavg[ind] = x
-        elif strcompare('new_teke', xarr):
-            try:
-                self.new_teke[ind] = x
-            except:
-                self.new_teke.extend([0] *
-                    (ind - len(self.new_teke) + 1))
-                self.new_teke[ind] = x
-        elif strcompare('new_temp', xarr):
-            try:
-                self.new_temp[ind] = x
-            except:
-                self.new_temp.extend([0] *
-                    (ind - len(self.new_temp) + 1))
-                self.new_temp[ind] = x
-        elif strcompare('new_salt', xarr):
-            try:
-                self.new_salt[ind] = x
-            except:
-                self.new_salt.extend([0] *
-                    (ind - len(self.new_salt) + 1))
-                self.new_salt[ind] = x
-        elif strcompare('new_shape_error', xarr):
-            try:
-                self.new_shape_error[ind] = x
-            except:
-                self.new_shape_error.extend([0] *
-                    (ind - len(self.new_shape_error) + 1))
-                self.new_shape_error[ind] = x
-
-        elif strcompare('new_contour_e', xarr):
-            try:
-                self.new_contour_e[ind] = x
-            except:
-                self.new_contour_e.append([0] *
-                    (ind - len(self.new_contour_e) + 1))
-                self.new_contour_e[ind] = x
-        elif strcompare('new_contour_s', xarr):
-            try:
-                self.new_contour_s[ind] = x
-            except:
-                self.new_contour_s.append([0] *
-                    (ind - len(self.new_contour_s) + 1))
-                self.new_contour_s[ind] = x
-        elif strcompare('new_uavg_profile', xarr):
-            try:
-                self.new_uavg_profile[ind] = x
-            except:
-                self.new_uavg_profile.append([0] *
-                    (ind - len(self.new_uavg_profile) + 1))
-                self.new_uavg_profile[ind] = x
-        else:
-            raise Exception
-        
-        return self'''
+
     
 
     
diff --git a/py_eddy_tracker_classes.py b/py_eddy_tracker_classes.py
@@ -880,30 +880,19 @@ def track_eddies(Eddy, first_record):
     
     far_away = 1e9
     
-    # We will need these in m for ellipse method below
+    # We will need these in M for ellipse method below
     old_x, old_y = Eddy.M(np.array(Eddy.old_lon), np.array(Eddy.old_lat))
     new_x, new_y = Eddy.M(np.array(Eddy.new_lon_tmp),
                           np.array(Eddy.new_lat_tmp))
     
-    X_old = np.array([Eddy.old_lon, Eddy.old_lat]).T
-    X_new = np.array([Eddy.new_lon_tmp, Eddy.new_lat_tmp]).T
+    X_old = np.asfortranarray([Eddy.old_lon, Eddy.old_lat]).T
+    X_new = np.asfortranarray([Eddy.new_lon_tmp, Eddy.new_lat_tmp]).T
     
     # Use haversine distance for distance matrix between every old and new eddy
-    #dist_mat = np.empty((X_old.shape[0], X_new.shape[0]))
-    dist_mat = np.ones((X_old.shape[0], X_new.shape[0]), dtype=np.float64)
+    dist_mat = np.empty((X_old.shape[0], X_new.shape[0]), order='F')
     
-    #print 'BEFORE dist_mat', dist_mat.min(), dist_mat.max()
-    #print dist_mat.flags
-    #print '--------------'
-    dist_mat = np.asfortranarray(dist_mat)
-    haversine.distance_matrix(np.asfortranarray(X_old),
-                              np.asfortranarray(X_new),
-                              dist_mat)
+    haversine.distance_matrix(X_old, X_new, dist_mat)
     dist_mat = np.ascontiguousarray(dist_mat)
-    #print '+++++++++++++++'
-    #print 'AFTER dist_mat', dist_mat.min(), dist_mat.max()
-    #print dist_mat.flags
-    #print 'JUNK', dist_mat_junk.min(), dist_mat_junk.max()
     
     dist_mat_copy = dist_mat.copy()
     
@@ -980,28 +969,26 @@ def track_eddies(Eddy, first_record):
         # Loop over separation distances between old and new
         for new_ind, new_dist in enumerate(dist_mat[old_ind]):
                 
-            sep_proceed = False
+            within_range = False
             
             if new_dist < far_away:
                 
                 if 'ellipse' in Eddy.SEPARATION_METHOD:
                     if Eddy.search_ellipse.ellipse_path.contains_point(
                                             (new_x[new_ind], new_y[new_ind])):
-                        sep_proceed = True
-                    #else:
-                        #sep_proceed = False
+                        within_range = True
                 
                 elif 'sum_radii' in Eddy.SEPARATION_METHOD:
                     sep_dist = Eddy.new_radii_tmp_e[new_ind]
                     sep_dist += Eddy.old_radii_e[old_ind]
                     sep_dist *= Eddy.sep_dist_fac
-                    sep_proceed = new_dist <= sep_dist
+                    within_range = new_dist <= sep_dist
                 
                 else:
                     Exception
             
             # Pass only the eddies within ellipse or sep_dist
-            if sep_proceed:
+            if within_range:
                 
                 old_amplitude[:] = Eddy.old_amp[old_ind]
                 new_amplitude[:] = Eddy.new_amp_tmp[new_ind]
@@ -1102,8 +1089,6 @@ def track_eddies(Eddy, first_record):
                     delta_salt = np.r_[delta_salt,
                         np.abs(np.diff([Eddy.old_salt[old_ind], new_st[i]]))]
             
-            #Eddy.search_ellipse.view_search_ellipse(Eddy)
-            #print dist_arr, DIST0
             # This from Penven etal (2005)
             deltaX = np.sqrt((delta_area / AREA0)**2 +
                              (delta_amp / AMP0)**2 +
@@ -1159,18 +1144,8 @@ def track_eddies(Eddy, first_record):
             
             # Use backup_ind to reinsert distances into dist_mat for the unused eddy/eddies
             for bind in backup_ind[dx_unused]:
-                #print dist_mat.shape
-                #print dist_mat[:, bind]
-                #print dist_mat_copy[:, bind]
                 dist_mat[:, bind] = dist_mat_copy[:, bind]
-                #print dist_mat[:, bind]
-                #print '\n'
-                #print 'B new_eddy_inds', new_eddy_inds
                 new_eddy_inds[bind] = True
-                #print 'A new_eddy_inds', new_eddy_inds
-                #print '\n'
-                
-                
                 
             if debug_dist:
                 print 'backup_ind', backup_ind
