split supergrid file

Author: rhaegar325

src/access_mopper/ocean_supergrid.py

@@ -1,3 +1,5 @@
+import os
+
 import numpy as np
 import xarray as xr
 
@@ -7,8 +9,14 @@ class Supergrid(object):
     # https://gist.github.com/rbeucher/b67c2b461557bc215a70017ea8dd337b
 
     def __init__(self, supergrid_file):
-        self.supergrid_file = supergrid_file
-        self.supergrid = xr.open_dataset(supergrid_file)
+        if isinstance(supergrid_file, list):
+            self.supergrid_file_list = supergrid_file
+            self.supergrid = xr.open_mfdataset(self.supergrid_file_list)
+        elif os.path.isfile(supergrid_file):
+            self.supergrid_file = supergrid_file
+            self.supergrid = xr.open_dataset(supergrid_file)
+        else:
+            raise ValueError("Couldn't find supergrid file")
 
         # T point locations
         self.xt = self.supergrid["x"][1::2, 1::2]

tests/data/om3/supergrid/x.nc

@@ -1,3 +1,4 @@
+import glob
 import importlib.resources as resources
 from pathlib import Path
 
@@ -139,7 +140,7 @@ def test_cmorise_OM3_2d(model_om3, cmor_name):
         DATA_DIR / f"om3/2d/access-om3.mom6.2d.{cmor_name}.1mon.mean.1902_01.nc"
     )
     try:
-        model_om3.supergrid = DATA_DIR / "supergrid/ocean_hgrid.nc"
+        model_om3.supergrid = glob.glob(str(DATA_DIR / "om3/supergrid/*.nc"))
         model_om3.cmorise(
             file_paths=file_pattern,
             compound_name="Omon." + cmor_name,
@@ -158,7 +159,7 @@ def test_cmorise_OM3_3d(model_om3, cmor_name):
         DATA_DIR / f"om3/3d/access-om3.mom6.3d.{cmor_name}.1mon.mean.1924_01.nc"
     )
     try:
-        model_om3.supergrid = DATA_DIR / "supergrid/ocean_hgrid.nc"
+        model_om3.supergrid = glob.glob(str(DATA_DIR / "om3/supergrid/*.nc"))
         model_om3.cmorise(
             file_paths=file_pattern,
             compound_name="Omon." + cmor_name,