Option to read geom_file

If geom_file is provided, override the coordinates in grd_file
using those from this file. This is necessary when
land-block elimination is enabled, which leads to NaNs in the
coordinates in grd_file.

mom6_tools/MOM6grid.py

@@ -1,7 +1,7 @@
 import xarray as xr
 
 # -- return MOM6 grid object
-def MOM6grid(grd_file, xrformat=False):
+def MOM6grid(grd_file, geom_file=None, xrformat=False):
     """
     Return an object or xarray Dataset with the MOM6 grid data.
 
@@ -10,6 +10,10 @@ def MOM6grid(grd_file, xrformat=False):
     grd_file : str
     Path to the static file.
 
+    geom_file : str, optional
+    If provided, override the coordinates in grd_file using those from this file. This is necessary when
+    land-block elimination is enabled, which leads to NaNs in grd_file.
+
     xrformat : boolean, optional
     If True, returns an xarray Dataset. Otherwise (default), returns an
     object with numpy arrays.
@@ -27,6 +31,14 @@ def MOM6grid(grd_file, xrformat=False):
     mgeolon_c = mgeolon_u.rename({'mgeolon_c'})
     mgeolon_v = mgeolon.rename({'mgeolon_v'})
 
+    if geom_file:
+      try: geom = xr.open_dataset(geom_file, decode_times=False)
+      except: raise Exception('Could not find file', geom_file)
+      nc['geolat'].values = geom.geolat.values
+      nc['geolon'].values = geom.geolon.values
+      nc['geolat_c'].values = geom.geolatb.values
+      nc['geolon_c'].values = geom.geolonb.values
+
     if xrformat:
       nc['mgeolon']   = mgeolon
       nc['mgeolon_c'] = mgeolon_c