Add arguments, improve output processing, keep time attributes

src/landusedata/__main__.py

@@ -1,5 +1,5 @@
 from landusedata._main import main
 
-# Gaurd against import time side effects
+# Guard against import time side effects
 if __name__ == '__main__':
     raise SystemExit(main())

src/landusedata/_main.py

@@ -1,8 +1,25 @@
 import argparse
+import os
 
 from landusedata.luh2 import main as luh2main
 from landusedata.landusepft import main as lupftmain
 
+def _shared_arguments(parser):
+    parser.add_argument(
+        'regrid_target_file',
+        help='target surface data file with desired grid resolution',
+    )
+    parser.add_argument(
+        'luh2_static_file',
+        help='luh2 static data file',
+    )
+    parser.add_argument(
+        '--overwrite',
+        action='store_true',
+        help='overwrite existing output file, if it exists',
+    )
+    return parser
+
 def main(argv=None):
 
     # Define top level parser
@@ -23,11 +40,11 @@ def main(argv=None):
     luh2_parser.set_defaults(func=luh2main)
     lupft_parser.set_defaults(func=lupftmain)
 
+    # Shared arguments
+    luh2_parser = _shared_arguments(luh2_parser)
+    lupft_parser = _shared_arguments(lupft_parser)
+
     # LUH2 subparser arguments
-    luh2_parser.add_argument('regrid_target_file',
-                             help='target surface data file with desired grid resolution')
-    luh2_parser.add_argument("luh2_static_file",
-                             help = "luh2 static data file")
     luh2_parser.add_argument('luh2_states_file',
                              help = "full path of luh2 raw states file")
     luh2_parser.add_argument('luh2_transitions_file',
@@ -50,10 +67,6 @@ def main(argv=None):
                              help = "output filename")
 
     # Landuse x pft subparser arguments
-    lupft_parser.add_argument('regrid_target_file',
-                             help='target surface data file with desired grid resolution')
-    lupft_parser.add_argument('luh2_static_file',
-                             help = "luh2 static data file")
     lupft_parser.add_argument('clm_luhforest_file',
                               help = "CLM5_current_luhforest_deg025.nc")
     lupft_parser.add_argument('clm_luhpasture_file',
@@ -69,12 +82,24 @@ def main(argv=None):
     # Parse the arguments
     args = parser.parse_args(argv)
 
+    # Only overwrite existing output if --overwrite specified
+    args.output = os.path.realpath(args.output)
+    if os.path.exists(args.output) and not args.overwrite:
+        raise FileExistsError(f"Output file exists; specify --overwrite to overwrite: {args.output}")
+
+    # Create output directory, if needed. Otherwise, check write access.
+    output_directory = os.path.dirname(args.output)
+    if not os.path.exists(output_directory):
+        os.makedirs(output_directory)
+    elif not os.access(output_directory, os.W_OK):
+        raise PermissionError("No write permissions in " + output_directory)
+
     # Call the default function for the given subcommand
     args.func(args)
 
     # Return successful completion
     return 0
 
-# Gaurd against import time side effects
+# Guard against import time side effects
 if __name__ == '__main__':
     raise SystemExit(main())

src/landusedata/landusepft.py

@@ -85,7 +85,7 @@ def main(args):
     # this will contain different data from a future CLM landuse x pft update
     ds_output['frac_secnd'] = ds_output.frac_primr.copy(deep=True)
 
-    # ds_regrid = ds_regrid.rename_dims(dims_dict={'lat':'lsmlat','lon':'lsmlon'})
+    # ds_regrid = ds_regrid.rename({'lat':'lsmlat','lon':'lsmlon'})
 
     # Output dataset to netcdf file
     print('Writing fates landuse x pft dataset to file')

src/landusedata/luh2.py

@@ -74,6 +74,7 @@ def main(args):
             regrid_luh2["LATIXY"] = ds_regrid_target["LATIXY"] # TO DO: double check if this is strictly necessary
 
         # Rename the dimensions for the output.  This needs to happen after the "LONGXY/LATIXY" assignment
+        # and should not touch the respective coordinate variables (hence rename_dims instead of rename).
         if (not 'lsmlat' in list(regrid_luh2.dims)):
             regrid_luh2 = regrid_luh2.rename_dims({'lat':'lsmlat','lon':'lsmlon'})
 

src/landusedata/regrid.py

@@ -35,30 +35,31 @@ def RegridLoop(ds_to_regrid, regridder):
     # Loop through the variables one at a time to conserve memory
     ds_varnames = list(ds_to_regrid.variables.keys())
     varlen = len(ds_to_regrid.variables)
-    first_var = False
-    for i in range(varlen-1):
-
-        # Skip time variable
-        if (not "time" in ds_varnames[i]):
-
-            # Only regrid variables that match the lat/lon shape.
-            if (ds_to_regrid[ds_varnames[i]][0].shape == (ds_to_regrid.lat.shape[0], ds_to_regrid.lon.shape[0])):
-                print("regridding variable {}/{}: {}".format(i+1, varlen, ds_varnames[i]))
-
-                # For the first non-coordinate variable, copy and regrid the dataset as a whole.
-                # This makes sure to correctly include the lat/lon in the regridding.
-                if (not(first_var)):
-                    ds_regrid = ds_to_regrid[ds_varnames[i]].to_dataset() # convert data array to dataset
-                    ds_regrid = regridder(ds_regrid)
-                    first_var = True
-
-                # Once the first variable has been included, then we can regrid by variable
-                else:
-                    ds_regrid[ds_varnames[i]] = regridder(ds_to_regrid[ds_varnames[i]])
-            else:
-                print("skipping variable {}/{}: {}".format(i+1, varlen, ds_varnames[i]))
+    first_var = True
+    for i, var in enumerate(ds_to_regrid):
+        msg_text = " variable {}/{}: {}\n".format(i+1, varlen, var)
+
+        # Skip time variable and only regrid variables that match the lat/lon shape.
+        do_regrid = not "time" in var
+        do_regrid = do_regrid and \
+            ds_to_regrid[var][0].shape == (ds_to_regrid.lat.shape[0], ds_to_regrid.lon.shape[0])
+        if not do_regrid :
+            print("skipping" + msg_text)
+            continue
+
+        print("regridding" + msg_text)
+
+        # For the first non-coordinate variable, copy and regrid the dataset as a whole.
+        # This makes sure to correctly include the lat/lon in the regridding.
+        if first_var:
+            ds_regrid = ds_to_regrid[var].to_dataset() # convert data array to dataset
+            ds_regrid = regridder(ds_regrid)
+            first_var = False
+
+        # Once the first variable has been included, then we can regrid by variable
         else:
-            print("skipping variable {}/{}: {}".format(i+1, varlen, ds_varnames[i]))
+            # Somehow regridder() can affect its input in-place! We don't want this,
+            # so we pass it in as a copy.
+            ds_regrid[var] = regridder(ds_to_regrid[var].copy())
 
-    print("\n")
     return(ds_regrid)

src/landusedata/utils.py

@@ -21,7 +21,17 @@ def _RegridTargetPrep(regrid_target):
     by xesmf.  As such, this function renames the dimensions.  It also adds
     lat/lon coordinates based on the LONGXY and LATIXY variables.
     """
-    regrid_target = regrid_target.rename_dims(dims_dict={'lsmlat':'lat','lsmlon':'lon'})
+
+    # Skip if lat and lon are already the dimension names
+    if "lat" in regrid_target.dims and "lon" in regrid_target.dims:
+        return regrid_target
+
+    # Drop dimension variables if they already exist
+    regrid_target = regrid_target.drop_vars("lat", errors="ignore")
+    regrid_target = regrid_target.drop_vars("lon", errors="ignore")
+
+    # Rename dimensions and add coordinates
+    regrid_target = regrid_target.rename({'lsmlat':'lat','lsmlon':'lon'})
     regrid_target['lon'] = regrid_target.LONGXY.isel(lat=0)
     regrid_target['lat'] = regrid_target.LATIXY.isel(lon=0)