coreglib: clean up at/ct correction code

Author: dshean

demcoreg/coreglib.py

@@ -10,7 +10,6 @@
 import matplotlib.pyplot as plt
 
 from pygeotools.lib import malib, iolib, warplib
-from imview import pltlib
 
 def apply_xy_shift(ds, dx, dy, createcopy=True):
     """
@@ -486,21 +485,23 @@ def find_subpixel_peak_position(corr, subpixel_method='gaussian'):
     return subp_peak_position[0], subp_peak_position[1]
 
 ## functions for along-track cross-track correction and plotting
-def successive_med(a, first_axis=1,first_axis_only=False,sav_filter=False,sg_window=101,sg_poly=2,min_axes_count=350):
+
+#TO DO: add number of iterations as argument
+def successive_med(a, first_axis=1, first_axis_only=False, sav_filter=False, sg_window=101, sg_poly=2, min_axes_count=350):
     """
     Subtract median values from each axis of the input difference map array
     Parameters
     -----------
     a: np.ma.array
         input array
     first_axis: int
-        1 or 2 (which axis) to operate on first
+        0 or 1 (which axis) to operate on first
     first_axis_only: bool
         wether to apply the correction over the first axis only (True) or both the axes (False)
     sav_filter: bool
-        whether to smooth the axis-medians or not using savgol filter
+        whether to smooth the median values along each axis using savgol filter
     sg_window: int
-        odd window sizes to be used in savgol filtering
+        odd window size to be used in savgol filtering
     sg_poly: int
         polynomial order to be used in savgol filtering
     min_axes_count: int
@@ -510,39 +511,43 @@ def successive_med(a, first_axis=1,first_axis_only=False,sav_filter=False,sg_win
     b: np.ma.array
         Corrected array 
     med_first: np.array
-        first_axis  corrections
+        first_axis 1D corrections
     med_second: np.array
-        second_axis  corrections
+        second_axis 1D corrections
     first_correction_surface: np.array
-        correction surface along input first axis
-    second_correction_surface:np.array
-        correction surface along input second axis
+        2D correction surface using 1D corrections along first axis
+    second_correction_surface: np.array
+        2D correction surface using 1D corrections along second axis
     med_first_smooth: np.array
-        smoothend first_axis corrections (optional, if sav_filter=True)
+        smoothed first_axis corrections (optional, if sav_filter=True)
     med_second_smooth: np.array
-        smoothened second_axis  corrections (optional, if sav_filter=True)
+        smoothed second_axis corrections (optional, if sav_filter=True)
     
     """
     import scipy.signal
     ### Parts of the function was first written by David for the Arctic DEM snow drift correction 
+    # Determine specified axis order
     second_axis = 0
     if first_axis == 0:
         second_axis = 1
-    # compute axis-wide median and count metrics
+
+    # compute median and count metrics along first axis
     med_first = np.ma.median(a, axis=first_axis)
-    count_first = np.ma.count(a,axis=first_axis)
-    
-    # if an axis has lower than the min_axes_count the number of pixels, the median offset for that axis is set to 0, assuming it to be unreliable
-    idx_first = count_first<min_axes_count
+    count_first = np.ma.count(a, axis=first_axis)
+   
+    # Each row/col must have a minimum number of samples for reliable statistics 
+    # DES note TODO: this should be set to masked, not 0
+    idx_first = count_first < min_axes_count
     med_first[idx_first] = 0
 
-    # apply savgol filter to smoothen the signal
+    # apply savgol filter to smooth the 1D correction 
     if sav_filter:
-        med_first_smooth = scipy.signal.savgol_filter(med_first,window_length=sg_window, polyorder=sg_poly, mode='nearest')
-        first_correction_surface = np.expand_dims(med_first_smooth,axis=first_axis)
+        med_first_smooth = scipy.signal.savgol_filter(med_first, window_length=sg_window, polyorder=sg_poly, mode='nearest')
+        first_correction_surface = np.expand_dims(med_first_smooth, axis=first_axis)
     else:
-        first_correction_surface = np.expand_dims(med_first,axis=first_axis)
-    # correct the first axis with the smoothed curve
+        first_correction_surface = np.expand_dims(med_first, axis=first_axis)
+
+    # correct the array using the corrections along the first axis 
     b = a - first_correction_surface
 
     # if correction is only to be performed over first axis, return a zero magnitude signal
@@ -553,28 +558,32 @@ def successive_med(a, first_axis=1,first_axis_only=False,sav_filter=False,sg_win
             shp_axis = 1
         med_second = np.zeros(a.shape[shp_axis])
 
-    # if correction is to be performed both axes
+    # Compute a correction along the second axis
     else:
-        # compute axis-wide median and count metrics
+        #DES note why b and a here? Is this because b is filled with 0 above?
         med_second = np.ma.median(b, axis=second_axis)
-        count_second = np.ma.count(a,axis=second_axis)
-        # if an axis has lower than the min_axes_count the number of pixels, the median offset for that axis is set to 0, assuming it to be unreliable
-        idx_second = count_second<min_axes_count
+        count_second = np.ma.count(a, axis=second_axis)
+
+        # Each row/col must have a minimum number of samples for reliable statistics 
+        idx_second = count_second < min_axes_count
         med_second[idx_second] = 0
 
-    # apply savgol filter to smoothen the signal 
+    # apply savgol filter to smooth the signal 
     if sav_filter:
-        med_second_smooth = scipy.signal.savgol_filter(med_second,window_length=sg_window, polyorder=sg_poly, mode='nearest')
-        second_correction_surface = np.expand_dims(med_second_smooth,axis=second_axis)
+        med_second_smooth = scipy.signal.savgol_filter(med_second, window_length=sg_window, polyorder=sg_poly, mode='nearest')
+        second_correction_surface = np.expand_dims(med_second_smooth, axis=second_axis)
     else:
-        second_correction_surface = np.expand_dims(med_second_smooth,axis=second_axis)
+        second_correction_surface = np.expand_dims(med_second_smooth, axis=second_axis)
+
+    # correct the array along the second axis 
     b = b - second_correction_surface
-    out = [b,med_first,med_second,first_correction_surface,second_correction_surface]
+
+    out = [b, med_first, med_second, first_correction_surface, second_correction_surface]
     if sav_filter:
-        out.extend([med_first_smooth,med_second_smooth])
+        out.extend([med_first_smooth, med_second_smooth])
     return out
 
-def plot_ct_at_dh_map(ax,dh_init,clim_dh,ct_correction_surface,at_correction_surface,dh_final):
+def plot_ct_at_dh_map(ax, dh_init, clim_dh, ct_correction_surface, at_correction_surface, dh_final):
     """
     Plot initial elevation difference map, Across-track (Row-wise) and Along-track (Column-wise) correction surface and corrected elevation difference map
     Parameters
@@ -592,16 +601,17 @@ def plot_ct_at_dh_map(ax,dh_init,clim_dh,ct_correction_surface,at_correction_sur
     dh_final: np.ma.array
         final dh map (after correction applied)
     """
-    pltlib.iv(dh_init,cmap='RdBu',clim=clim_dh,label='Elevation difference (m)',title='dh before',ax=ax[0])
-    pltlib.add_scalebar(ax=ax[0],res=1)
+    from imview import pltlib
+    pltlib.iv(dh_init, cmap='RdBu', clim=clim_dh, label='Elevation difference (m)', title='dh before', ax=ax[0])
+    pltlib.add_scalebar(ax=ax[0], res=1)
     #across_track_clim = malib.calcperc_sym(ct_correction_surface,(2,98))
-    pltlib.iv(np.zeros(dh_init.shape)+ct_correction_surface,cmap='RdBu',clim=clim_dh,label='Elevation difference (m)',title='Row-wise correction surface',ax=ax[1])
+    pltlib.iv(np.zeros(dh_init.shape)+ct_correction_surface, cmap='RdBu', clim=clim_dh, label='Elevation difference (m)', title='Row-wise correction surface', ax=ax[1])
     #along_track_clim = malib.calcperc_sym(at_correction_surface,(2,98))
-    pltlib.iv(np.zeros(dh_init.shape)+at_correction_surface,cmap='RdBu',clim=clim_dh,label='Elevation difference (m)',title='Column-wise correction surface',ax=ax[2])
-    pltlib.iv(dh_final,cmap='RdBu',clim=clim_dh,label='Elevation difference (m)',title='dh after',ax=ax[3])
+    pltlib.iv(np.zeros(dh_init.shape)+at_correction_surface, cmap='RdBu', clim=clim_dh, label='Elevation difference (m)', title='Column-wise correction surface', ax=ax[2])
+    pltlib.iv(dh_final, cmap='RdBu', clim=clim_dh, label='Elevation difference (m)', title='dh after', ax=ax[3])
     plt.tight_layout()
 
-def plot_ct_at_dh_fits(f,ct_med,ct_smooth,at_med,at_smooth,clim_dh=None):
+def plot_ct_at_dh_fits(f, ct_med, ct_smooth, at_med, at_smooth, clim_dh=None):
     """
     Plot Across-track (Row-wise) and Along-track (Column-wise) correction fits
     Parameters
@@ -613,26 +623,26 @@ def plot_ct_at_dh_fits(f,ct_med,ct_smooth,at_med,at_smooth,clim_dh=None):
     ct_smooth: np.array
         Smooth fit to median error per-row computed using Sav-Golay fit
     at_med: np.array
-        median error per column (1,dh_map.shape[0])
+        median error per column (1, dh_map.shape[0])
     at_smooth:np.array
         Smooth fit to median error per-column computed using Sav-Golay fit
     clim_dh: tuple
         symmetrical min/max values to limit correction fits
         
     """
     ax1 = plt.subplot(1,2,1)
-    ax1.plot(ct_med,np.arange(len(ct_med)),c='k',label='median correction')
-    ax1.plot(ct_smooth,np.arange(len(ct_med)),c='r',label='smooth Sav-Golay fit')
+    ax1.plot(ct_med, np.arange(len(ct_med)), c='k', label='median correction')
+    ax1.plot(ct_smooth, np.arange(len(ct_med)), c='r', label='smooth Sav-Golay fit')
     ax1.set_title('Row-wise correction')
     ax1.set_xlabel('Elevation difference (m)')
-    ax1.axvline(x=0,ls='--',alpha=0.6,c='teal')
+    ax1.axvline(x=0, ls='--', alpha=0.6, c='teal')
     ax1.set_ylabel('Row number')
     ax1.legend()
 
     ax2 = plt.subplot(1,2,2)
-    ax2.plot(np.arange(len(at_med)),at_med,c='k',label='median correction')
-    ax2.plot(np.arange(len(at_med)),at_smooth,c='r',label='smooth Sav-Golay fit')
-    ax2.axhline(y=0,ls='--',alpha=0.6,c='teal')
+    ax2.plot(np.arange(len(at_med)), at_med, c='k', label='median correction')
+    ax2.plot(np.arange(len(at_med)), at_smooth, c='r', label='smooth Sav-Golay fit')
+    ax2.axhline(y=0, ls='--', alpha=0.6, c='teal')
     ax2.set_title('Column-wise correction')
     ax2.set_ylabel('Elevation difference (m)')
     ax2.set_xlabel('Col number')
@@ -647,7 +657,8 @@ def plot_ct_at_dh_fits(f,ct_med,ct_smooth,at_med,at_smooth,clim_dh=None):
         ax2.set_ylim(clim_dh)
     plt.tight_layout()
 
-def ct_at_correction_wrapper(src_dem_fn,dh_fn,dh_filt_fn,ct_only=False,sg_window=101,sg_poly=2,min_axes_count=350,outdir=None):
+#DES TO DO: add number of iterations as argument
+def ct_at_correction_wrapper(src_dem_fn, dh_fn, dh_filt_fn, ct_only=False, sg_window=101, sg_poly=2, min_axes_count=350, outdir=None):
     """
     Wrapper function to apply across-track (row-wise) and along-track correction to difference maps and src DEM
     Parameters
@@ -672,54 +683,49 @@ def ct_at_correction_wrapper(src_dem_fn,dh_fn,dh_filt_fn,ct_only=False,sg_window
         Write corrected source DEM, difference maps and plots to disc (not used currently)
     """
     # warp the difference maps to extent and resolution of source DEMs
-    print("Step1: Warping difference maps to extent of source DEM")
-    ds_list = warplib.memwarp_multi_fn([src_dem_fn,dh_fn,dh_filt_fn],res='first',extent='first',r='cubic')
+    print("Warping difference maps to extent of source DEM")
+    ds_list = warplib.memwarp_multi_fn([src_dem_fn, dh_fn, dh_filt_fn], res='first', extent='first', r='cubic')
 
     # read into memory
-    src_dem,dh,dh_filt = [iolib.ds_getma(ds) for ds in ds_list]
+    src_dem, dh, dh_filt = [iolib.ds_getma(ds) for ds in ds_list]
 
     # calculate clim
-    clim_dh = malib.calcperc_sym(dh_filt,(5,95))
+    clim_dh = malib.calcperc_sym(dh_filt, (5,95))
     # perform the correction
     print("Computing Across-track (Row-wise) and Along-track (Column-wise) correction")
-    dh_filt_corr,ct_med,at_med,ct_correction_surface,at_correction_surface, ct_med_smooth, at_med_smooth = successive_med(dh_filt,
-                                                                                                                          first_axis_only=ct_only,sav_filter=True,
-                                                                                                                          sg_window=sg_window,sg_poly=sg_poly,
-                                                                                                                          min_axes_count=min_axes_count)
+    dh_filt_corr, ct_med, at_med, ct_correction_surface, at_correction_surface, ct_med_smooth, at_med_smooth = \
+            successive_med(dh_filt, first_axis_only=ct_only, sav_filter=True, sg_window=sg_window, sg_poly=sg_poly, min_axes_count=min_axes_count)
+
     # prepare the plots
     out_dh_fig = os.path.splitext(dh_fn)[0] + '_ct_at_dh_map.png'
     print(f"Creating Across-track (Row-wise) and Along-track (Column-wise) difference map figure at {out_dh_fig}")
-    f,ax = plt.subplots(1,4,figsize=(12,5))
+    f, ax = plt.subplots(1, 4, figsize=(12,5))
 
-    plot_ct_at_dh_map(ax,dh_filt,clim_dh,ct_correction_surface,at_correction_surface,dh_filt_corr)
-    f.savefig(out_dh_fig, dpi=300,bbox_inches='tight', pad_inches=0.1)
+    plot_ct_at_dh_map(ax, dh_filt, clim_dh, ct_correction_surface, at_correction_surface, dh_filt_corr)
+    f.savefig(out_dh_fig, dpi=300, bbox_inches='tight', pad_inches=0.1)
 
     out_lineplot_fig = os.path.splitext(dh_fn)[0] + '_ct_at_correction_fit.png'
     print(f"Creating Across-track (Row-wise) and Along-track (Column-wise) correction fits figure at {out_lineplot_fig}")
     fig = plt.figure(figsize=(8,4))
-    plot_ct_at_dh_fits(fig,ct_med,ct_med_smooth,at_med,at_med_smooth,clim_dh=clim_dh)
-    fig.savefig(out_lineplot_fig,dpi=300,bbox_inches='tight', pad_inches=0.1)
+    plot_ct_at_dh_fits(fig, ct_med, ct_med_smooth, at_med, at_med_smooth, clim_dh=clim_dh)
+    fig.savefig(out_lineplot_fig, dpi=300, bbox_inches='tight', pad_inches=0.1)
 
     # Correct source DEM
-    src_dem_corrected = src_dem - ct_correction_surface
-    src_dem_corrected = src_dem_corrected - at_correction_surface
+    src_dem_corrected = src_dem - ct_correction_surface - at_correction_surface
     src_dem_corrected_fn = os.path.splitext(src_dem_fn)[0]+'_ct_at_corrected.tif'
     print(f"Writing out corrected source DEM at {src_dem_corrected_fn}")
-    iolib.writeGTiff(src_dem_corrected,src_dem_corrected_fn,src_ds=ds_list[0])
+    iolib.writeGTiff(src_dem_corrected, src_dem_corrected_fn, src_ds=ds_list[0])
 
     # correct difference map
     # this is the entire difference map from which the filtered map is derived (containing glacier etc)
-    dh_corrected = dh - ct_correction_surface
-    dh_corrected = dh_corrected - at_correction_surface
+    dh_corrected = dh - ct_correction_surface - at_correction_surface
     dh_corrected_fn = os.path.splitext(dh_fn)[0]+'_ct_at_corrected.tif'
     print(f"Writing out corrected elevation difference map at {dh_corrected_fn}")
-    iolib.writeGTiff(dh_corrected,dh_corrected_fn,src_ds=ds_list[1])
+    iolib.writeGTiff(dh_corrected, dh_corrected_fn, src_ds=ds_list[1])
 
     # write out filtered difference map
     dh_filt_corrected_fn = os.path.splitext(dh_filt_fn)[0]+'_ct_at_corrected.tif'
     print(f"Writing out corrected elevation difference map at {dh_filt_fn}")
-    iolib.writeGTiff(dh_filt_corr,dh_filt_corrected_fn,src_ds=ds_list[2])
+    iolib.writeGTiff(dh_filt_corr, dh_filt_corrected_fn, src_ds=ds_list[2])
 
     print("Across-track (row-wise), Along-track (col-wise) correction complete")
-
-