Work on ocean support

Author: rbeucher

CHANGELOG.rst

@@ -8,5 +8,6 @@ v2.0.0
 ------
 
 - Fully supported by ACCESS-NRI
+- Remove dependency on CMOR
 - Remove MOPPeR database
 - Remove integration with NCI queue system (PBS)

GettingStartedOcean.ipynb

@@ -0,0 +1,162 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": 1,
+   "id": "66ea26fc-9f2f-4667-92ff-a46ea27a16d5",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "Loaded Configuration:\n",
+      "Creator Name: Romain Beucher\n",
+      "Organisation: ACCESS-NRI\n",
+      "Creator Email: romain.beucher@anu.edu.au\n",
+      "Creator URL: 0000-0003-3891-5444\n"
+     ]
+    }
+   ],
+   "source": [
+    "from access_mopper.vocabulary_processors import CMIP6Vocabulary"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 2,
+   "id": "e1601178-2f7c-4538-ab98-8e856d4185c2",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "input_paths=\"/home/romain/PROJECTS/ACCESS-MOPPeR/ACCESS-MOPPeR/tests/data/esm1-6/ocean/ocean-2d-surface_temp-1monthly-mean-ym_0966_01.nc\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 3,
+   "id": "fd293ca3-d6e2-4fdc-935c-bd43d06d3110",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "parent_experiment_config = {\n",
+    "        \"parent_experiment_id\": \"piControl\",\n",
+    "        \"parent_activity_id\": \"CMIP\",\n",
+    "        \"parent_source_id\": \"ACCESS-ESM1-5\",\n",
+    "        \"parent_variant_label\": \"r1i1p1f1\",\n",
+    "        \"parent_time_units\": \"days since 0001-01-01 00:00:00\",\n",
+    "        \"parent_mip_era\": \"CMIP6\",\n",
+    "        \"branch_time_in_child\": 0.0,\n",
+    "        \"branch_time_in_parent\": 54786.0,\n",
+    "        \"branch_method\": \"standard\"\n",
+    "}"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 4,
+   "id": "fb371092-d40d-487d-8b4d-28cede52e6bb",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "voc = CMIP6Vocabulary(\n",
+    "    compound_name=\"Omon.tos\",\n",
+    "    experiment_id=\"historical\",\n",
+    "    source_id=\"ACCESS-ESM1-5\",\n",
+    "    variant_label=\"r1i1p1f1\",\n",
+    "    grid_label=\"gn\",\n",
+    "    activity_id=\"CMIP\",\n",
+    "    parent_info=parent_experiment_config)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 5,
+   "id": "d7ad164d-46a4-4d73-ab2d-c09fc6c1e257",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from access_mopper.cmip6_cmoriser import CMIP6_Ocean_CMORiser\n",
+    "from access_mopper.utilities import load_cmip6_mappings\n",
+    "from access_mopper.ocean_supergrid import Supergrid"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 6,
+   "id": "f5e56ce6-ee0e-4c38-9526-f08746fbc49e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "variable_mapping = load_cmip6_mappings(\"Omon.tos\")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
+   "id": "397687e8-46de-4d3f-87fd-0682c00afede",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "supergrid = \"/home/romain/PROJECTS/ACCESS-MOPPeR/grids/access-om2/input_20201102/mom_1deg/ocean_hgrid.nc\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 8,
+   "id": "51cbaca6-bf14-4849-b074-dd709f4c822e",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "cmoriser = CMIP6_Ocean_CMORiser(\n",
+    "    input_paths=input_paths,\n",
+    "    output_path=\".\",\n",
+    "    cmor_name=voc.cmor_name,\n",
+    "    cmip6_vocab=voc,\n",
+    "    variable_mapping=variable_mapping,\n",
+    "    supergrid_path=supergrid\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "54efceb4-1d4f-4f51-ba5d-8a2644aa7652",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CMORised output written to tos_Omon_ACCESS-ESM1-5_historical_r1i1p1f1_gn_096601-096612.nc\n"
+     ]
+    }
+   ],
+   "source": [
+    "cmoriser.run()"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.11.11"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 5
+}

src/access_mopper/cmip6_cmoriser.py

@@ -7,6 +7,7 @@
 import xarray as xr
 from cftime import num2date
 
+from access_mopper.ocean_supergrid import Supergrid
 from access_mopper.utilities import load_cmip6_mappings, type_mapping
 from access_mopper.vocabulary_processors import CMIP6Vocabulary
 
@@ -303,7 +304,10 @@ def write(self):
             for k, v in attrs.items():
                 dst.setncattr(k, v)
             for dim, size in self.ds.sizes.items():
-                dst.createDimension(dim, size)
+                if dim == "time":
+                    dst.createDimension(dim, None)  # Unlimited dimension
+                else:
+                    dst.createDimension(dim, size)
             for var in self.ds.variables:
                 vdat = self.ds[var]
                 fill = None if var.endswith("_bnds") else vdat.attrs.get("_FillValue")
@@ -329,6 +333,152 @@ def run(self):
         self.write()
 
 
+class CMIP6_Ocean_CMORiser(CMIP6_CMORiser):
+    """
+    CMORiser subclass for ocean variables using curvilinear supergrid coordinates.
+    """
+
+    def __init__(
+        self,
+        input_paths: Union[str, List[str]],
+        output_path: str,
+        cmor_name: str,
+        cmip6_vocab: Any,
+        variable_mapping: Dict[str, Any],
+        supergrid_path: Union[str, Path],
+        drs_root: Optional[Path] = None,
+    ):
+        super().__init__(
+            input_paths=input_paths,
+            output_path=output_path,
+            cmor_name=cmor_name,
+            cmip6_vocab=cmip6_vocab,
+            variable_mapping=variable_mapping,
+            drs_root=drs_root,
+        )
+        self.supergrid = Supergrid(supergrid_path)
+        self.grid_info = None
+        self.grid_type = None
+
+    def infer_grid_type(self):
+        coord_sets = {
+            "T": {"xt_ocean", "yt_ocean"},
+            "U": {"xu_ocean", "yu_ocean"},
+            "V": {"xv_ocean", "yv_ocean"},
+            "Q": {"xq_ocean", "yq_ocean"},
+        }
+        present_coords = set(self.ds.coords)
+        for grid, required in coord_sets.items():
+            if required.issubset(present_coords):
+                return grid
+        raise ValueError("Could not infer grid type from dataset coordinates.")
+
+    def select_and_process_variables(self):
+        input_vars = self.mapping[self.cmor_name]["model_variables"]
+        calc = self.mapping[self.cmor_name]["calculation"]
+
+        self.ds = self.ds[input_vars]
+
+        if calc["type"] == "direct":
+            self.ds[self.cmor_name] = self.ds[input_vars[0]]
+        elif calc["type"] == "formula":
+            local = {var: self.ds[var] for var in input_vars}
+            self.ds[self.cmor_name] = eval(calc["formula"], {}, local)
+        else:
+            raise ValueError(f"Unsupported calculation type: {calc['type']}")
+
+        dim_rename = {
+            "xt_ocean": "i",
+            "yt_ocean": "j",
+            "xu_ocean": "i",
+            "yu_ocean": "j",
+            "xq_ocean": "i",
+            "yq_ocean": "j",
+            "xv_ocean": "i",
+            "yv_ocean": "j",
+        }
+        dims_to_rename = {
+            k: v for k, v in dim_rename.items() if k in self.ds[self.cmor_name].dims
+        }
+        self.ds[self.cmor_name] = self.ds[self.cmor_name].rename(dims_to_rename)
+        self.ds[self.cmor_name] = self.ds[self.cmor_name].transpose("time", "j", "i")
+
+        self.grid_type = self.infer_grid_type()
+        # Drop all other data variables except the CMOR variable
+        self.ds = self.ds[[self.cmor_name]]
+
+        # Drop unused coordinates
+        used_coords = set()
+        for dim in self.ds[self.cmor_name].dims:
+            if dim in self.ds.coords:
+                used_coords.add(dim)
+            else:
+                # Might be implicit dimension (e.g. from formula), check all coords
+                for coord in self.ds.coords:
+                    if dim in self.ds[coord].dims:
+                        used_coords.add(coord)
+        self.ds = self.ds.drop_vars([c for c in self.ds.coords if c not in used_coords])
+
+    def update_attributes(self):
+        grid_type = self.grid_type
+        self.grid_info = self.supergrid.extract_grid(grid_type)
+
+        self.ds = self.ds.assign_coords(
+            {
+                "i": self.grid_info["i"],
+                "j": self.grid_info["j"],
+                "vertices": self.grid_info["vertices"],
+            }
+        )
+
+        self.ds["latitude"] = self.grid_info["latitude"]
+        self.ds["longitude"] = self.grid_info["longitude"]
+        self.ds["vertices_latitude"] = self.grid_info["vertices_latitude"]
+        self.ds["vertices_longitude"] = self.grid_info["vertices_longitude"]
+
+        self.ds["latitude"].attrs.update(
+            {
+                "standard_name": "latitude",
+                "units": "degrees_north",
+                "bounds": "vertices_latitude",
+            }
+        )
+        self.ds["longitude"].attrs.update(
+            {
+                "standard_name": "longitude",
+                "units": "degrees_east",
+                "bounds": "vertices_longitude",
+            }
+        )
+        self.ds["vertices_latitude"].attrs.update(
+            {"standard_name": "latitude", "units": "degrees_north"}
+        )
+        self.ds["vertices_longitude"].attrs.update(
+            {"standard_name": "longitude", "units": "degrees_east"}
+        )
+
+        self.ds.attrs = {
+            k: v
+            for k, v in self.vocab.get_required_global_attributes().items()
+            if v not in (None, "")
+        }
+
+        if "nv" in self.ds.dims:
+            self.ds = self.ds.rename_dims({"nv": "bnds"}).rename_vars({"nv": "bnds"})
+            self.ds["bnds"].attrs.update(
+                {"long_name": "vertex number of the bounds", "units": "1"}
+            )
+
+        cmor_attrs = self.vocab.variable
+        self.ds[self.cmor_name].attrs.update(
+            {k: v for k, v in cmor_attrs.items() if v not in (None, "")}
+        )
+        var_type = cmor_attrs.get("type", "double")
+        self.ds[self.cmor_name] = self.ds[self.cmor_name].astype(
+            self.type_mapping.get(var_type, np.float64)
+        )
+
+
 class CMIP6Workflow:
     """
     Coordinates the CMORisation process using CMIP6Vocabulary and CMORiser.

src/access_mopper/mappings/Mappings_CMIP6_Omon.json

@@ -663,19 +663,19 @@
         }
     },
     "tos": {
-        "dimensions": [
-            "time",
-            "yt_ocean",
-            "xt_ocean"
-        ],
+        "dimensions": {
+            "time": "time",
+            "lon": "yt_ocean",
+            "lat": "xt_ocean"
+        },
         "units": "K",
         "positive": null,
         "model_variables": [
-            "sst"
+            "surface_temp"
         ],
         "calculation": {
-            "type": "formula",
-            "formula": "K_degC(sst)"
+            "type": "direct",
+            "formula": "surface_temp"
         }
     },
     "umo": {