Generalise variable removal

adjust.py

@@ -234,6 +234,7 @@ def main(args):
         use_cftime=False,
         valid_min=args.valid_min,
         valid_max=args.valid_max,
+        drop_vars=args.drop_vars,
     )
     var = args.rename_var if args.rename_var else args.var
 
@@ -388,6 +389,13 @@ def main(args):
         default=False,
         help="Use wildcards to shorten the file lists in output_file history attribute",
     )
+    parser.add_argument(
+        "--drop_vars",
+        type=str,
+        nargs='*',
+        default=[],
+        help="File variables to drop when reading infiles"
+    )
     args = parser.parse_args()
     log_level = logging.INFO if args.verbose else logging.WARNING
     logging.basicConfig(level=log_level)

train.py

@@ -129,6 +129,7 @@ def main(args):
         output_units=args.output_units,
         valid_min=args.valid_min,
         valid_max=args.valid_max,
+        drop_vars=args.hist_drop_vars,
     )
     calendar_hist = type(ds_hist['time'].values[0])
     ds_ref = utils.read_data(
@@ -142,6 +143,7 @@ def main(args):
         output_calendar=calendar_hist,
         valid_min=args.valid_min,
         valid_max=args.valid_max,
+        drop_vars=args.ref_drop_vars,
     )
     ds_out = train(
         ds_hist,
@@ -307,6 +309,20 @@ def main(args):
         default=False,
         help="Use wildcards to shorten the file lists in output_file history attribute",
     )
+    parser.add_argument(
+        "--hist_drop_vars",
+        type=str,
+        nargs='*',
+        default=[],
+        help="File variables to drop when reading hist files"
+    )
+    parser.add_argument(
+        "--ref_drop_vars",
+        type=str,
+        nargs='*',
+        default=[],
+        help="File variables to drop when reading ref files"
+    )
     args = parser.parse_args()
     log_level = logging.INFO if args.verbose else logging.WARNING
     logging.basicConfig(level=log_level)

utils.py

@@ -192,6 +192,7 @@ def read_data(
     output_calendar=None,
     valid_min=None,
     valid_max=None,
+    drop_vars=[],
 ):
     """Read and process an input dataset.
 
@@ -225,6 +226,8 @@ def read_data(
         Clip data to valid minimum value
     valid_max : float, optional
         Clip data to valid maximum value
+    drop_vars : list, optional
+        List of variables to drop when reading file
 
     Returns
     -------
@@ -243,10 +246,9 @@ def read_data(
         except ValueError:
             ds = xr.open_mfdataset(infiles)
 
-    drop_coords = ['height', 'sigma', 'level_height', 'model_level_number']
-    for coord in drop_coords:
+    for drop_var in drop_vars:
         try:
-            ds = ds.drop(coord)
+            ds = ds.drop(drop_var)
         except ValueError:
             pass