Feature/add preprocessing smoothing

pycascadia/remove_restore.py

@@ -5,7 +5,7 @@
 It's just a way for us to learn how to use pyGMT and how to read/write the sample data.
 """
 
-from pygmt import blockmedian, surface, grdtrack, grdcut, grdfilter
+from pygmt import blockmedian, surface, grdtrack, grdfilter
 import os
 import pandas as pd
 import matplotlib.pyplot as plt
@@ -25,13 +25,14 @@ def calc_diff_grid(base_grid, update_grid, diff_threshold=0.0):
     print("Find z in base grid")
     base_pts = grdtrack(bmd, base_grid.grid, 'base_z', interpolation='l')
 
-    print ("Create difference grid")
+    print("Create difference grid")
     diff = pd.DataFrame()
     diff['x'] = base_pts['x']
     diff['y'] = base_pts['y']
     diff['z'] = base_pts['z'] - base_pts['base_z']
 
     diff[diff.z.abs() < diff_threshold]['z'] = 0.0 # Filter out small differences
+    diff[diff.z.isnull()] = 0.0 # Filter out NaNs.
 
     diff_xyz_fname = "diff.xyz"
     diff_grid_fname = "diff.nc"
@@ -51,19 +52,39 @@ def load_base_grid(fname, region=None, spacing=None):
     base_grid = Grid(fname, convert_to_xyz=False)
     if region:
         base_grid.crop(region)
-    if spacing:
-        base_grid.resample(spacing)
-
     return base_grid
 
+def preprocess_base_grid(base_grid, update_grid, final_spacing):
+    """
+    Combines and smooths data from the base and update grid,
+    as described in workflow steps B+C of the GEBCO cookbook.
+    This is a preprocessing step to remove-restore.
+    """
+    xyz1 = base_grid.as_xyz()
+    xyz2 = update_grid.as_xyz()
+    xyz1.append(xyz2, ignore_index=True)
+
+    max_spacing = max(update_grid.spacing, final_spacing)
+    minimal_region = min_regions(update_grid.region, base_grid.region)
+    bmd = blockmedian(xyz1, spacing=4*max_spacing, region=minimal_region, Q=True)
+    combined = surface(bmd.x, bmd.y, bmd.z, spacing=4*max_spacing, region=minimal_region)
+
+    filter_combined = grdfilter(combined, D=0, F=f"c{12*max_spacing}")
+    fname = "filtered.nc"
+    filter_combined.to_netcdf(fname)
+    combined_grid = Grid(fname)
+    os.remove(fname)
+    combined_grid.resample(max_spacing)
+
+    return combined_grid
 
 def main():
     # Handle arguments
     parser = argparse.ArgumentParser(description='Combine multiple bathymmetry sources into a single grid')
     parser.add_argument('filenames', nargs='+', help='sources to combine with the base grid')
     parser.add_argument('--base', required=True, help='base grid')
     parser.add_argument('--spacing', type=float, help='output grid spacing')
-    parser.add_argument('--diff_threshold', default=0.0, help='value above which differences will be added to the base grid')
+    parser.add_argument('--diff_threshold', type=float, default=0.0, help='value above which differences will be added to the base grid')
     parser.add_argument('--plot', action='store_true', help='plot final output before saving')
     parser.add_argument('--output', required=True, help='filename of final output')
     parser.add_argument('--region_of_interest', required=False, nargs=4, type=float,
@@ -84,23 +105,27 @@ def main():
         print("Loading update grid")
         update_grid = Grid(fname, convert_to_xyz=True)
 
-        diff_grid = calc_diff_grid(base_grid, update_grid, diff_threshold=diff_threshold)
+        print("Combining grids")
+        combined_base_grid = preprocess_base_grid(base_grid, update_grid, args.spacing)
+
+        diff_grid = calc_diff_grid(combined_base_grid, update_grid, diff_threshold=diff_threshold)
 
         print("Update base grid")
-        base_grid.grid.values += diff_grid.values
+        combined_base_grid.grid.values += diff_grid.values
 
-    base_grid.save_grid(args.output)
+    # TODO How should this work with several input?!
+    combined_base_grid.save_grid(args.output)
 
     if args.plot:
         fig, axes = plt.subplots(2,2)
         initial_base_grid = load_base_grid(base_fname, region=args.region_of_interest)
         initial_base_grid.plot(ax=axes[0,0])
         axes[0,0].set_title("Initial Grid")
-        base_grid.plot(ax=axes[0,1])
+        combined_base_grid.plot(ax=axes[0,1])
         axes[0,1].set_title("Final Grid")
-        base_grid.grid.differentiate('x').plot(ax=axes[1,0])
+        combined_base_grid.grid.differentiate('x').plot(ax=axes[1,0])
         axes[1,0].set_title("x Derivative of Final Grid")
-        base_grid.grid.differentiate('y').plot(ax=axes[1,1])
+        combined_base_grid.grid.differentiate('y').plot(ax=axes[1,1])
         axes[1,1].set_title("y Derivative of Final Grid")
         plt.show()