Update support for ocean in OM3 and ESM1-6

Author: rbeucher

notebooks/Getting_started.ipynb

@@ -251,7 +251,7 @@
    "id": "59b4dedb-aca7-4969-a285-9a07561c3cea",
    "metadata": {},
    "source": [
-    "## CMORising files for an Atmosphere variable (zos)\n"
+    "## CMORising files for an Atmosphere variable (hur)\n"
    ]
   },
   {
@@ -299,6 +299,56 @@
     "                         mip_table=\"CMIP6_Amon.json\")"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "158f4c2f-5e1c-4c25-b27d-129bade71078",
+   "metadata": {},
+   "source": [
+    "## CMORising files for an Ocean variable (zos) (ACCESS-OM2 In development)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 9,
+   "id": "7c7708e7-d72c-4000-b4f4-b7afb45c4a6b",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "file_pattern = \"/home/romain/PROJECTS/ACCESS-MOPPeR/Test_data/cj877/history/ocn/ocean-2d-sea_level-1-monthly-mean-ym_0326_01.nc\""
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 11,
+   "id": "41ea0fec-0315-4be9-8d9d-0dc24b637ed3",
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "Stored in: MOPPeR_outputs/CMIP6/CMIP/CSIRO/ACCESS-ESM1-5/piControl-spinup/r1i1p1f1/Omon/zos/gn/v20250423/zos_Omon_ACCESS-ESM1-5_piControl-spinup_r1i1p1f1_gn_032601-032606.nc\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "\n",
+      "! ------\n",
+      "! All files were closed successfully. \n",
+      "! ------\n",
+      "! \n"
+     ]
+    }
+   ],
+   "source": [
+    "val = experiment.cmorise_ocean(file_paths=glob.glob(file_pattern),\n",
+    "                               compound_name= \"Omon.zos\", \n",
+    "                               cmor_dataset_json= \"experiment.json\",\n",
+    "                               mip_table=\"CMIP6_Omon.json\")"
+   ]
+  },
   {
    "attachments": {},
    "cell_type": "markdown",
@@ -310,7 +360,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 12,
    "id": "c7b6beda-3d28-41ce-8bc7-02ce3402dfbc",
    "metadata": {},
    "outputs": [],
@@ -335,7 +385,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 13,
    "id": "aae20d6a-8ea6-44d6-9e25-ead8f2c93108",
    "metadata": {},
    "outputs": [
@@ -353,7 +403,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 13,
+   "execution_count": 14,
    "id": "350920e3-7ad0-44e5-98e8-653c4830f4e5",
    "metadata": {},
    "outputs": [
@@ -380,7 +430,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 14,
+   "execution_count": 15,
    "id": "e228b219-6c5b-4ade-9d7d-fbd5d9e13ddb",
    "metadata": {},
    "outputs": [],
@@ -390,7 +440,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 15,
+   "execution_count": 16,
    "id": "3d75e6c4-c902-44f0-ba75-7c835483bf85",
    "metadata": {},
    "outputs": [
@@ -451,6 +501,14 @@
     "            cmor_dataset_json= \"experiment.json\",\n",
     "            mip_table=\"CMIP6_Omon.json\")"
    ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "d4c44a43-ef66-4345-9f51-3c45ab53ae03",
+   "metadata": {},
+   "outputs": [],
+   "source": []
   }
  ],
  "metadata": {

src/access_mopper/configurations.py

@@ -9,7 +9,7 @@
 from .calc_atmos import level_to_height
 from .calc_land import average_tile, calc_landcover, calc_topsoil, extract_tilefrac
 from .dataclasses import CMIP6_Experiment
-from .ocean_supergrid import om3_grid
+from .ocean_supergrid import om2_grid, om3_grid
 
 # Supported operators
 OPERATORS = {
@@ -24,6 +24,7 @@
 @dataclass
 class ACCESS_ESM16_CMIP6(CMIP6_Experiment):
     mapping_file_prefix: str = "Mappings_CMIP6_"
+    supergrid = om2_grid
 
     def __post_init__(self):
         self.initialise("ACCESS-ESM1-5")
@@ -136,6 +137,129 @@ def cmorise(self, file_paths, compound_name, cmor_dataset_json, mip_table):
 
         cmor.close()
 
+    def cmorise_ocean(self, file_paths, compound_name, cmor_dataset_json, mip_table):
+        _, cmor_name = compound_name.split(".")
+
+        # Open the matching files with xarray
+        ds = xr.open_mfdataset(file_paths, combine="by_coords", decode_times=False)
+
+        # Extract required variables and coordinates
+        mapping = self.get_mapping(compound_name=compound_name)
+
+        if mapping["calculation"]["type"] == "direct":
+            access_var = mapping["calculation"]["formula"]
+            variable_units = mapping["units"]
+            positive = mapping["positive"]
+            var = ds[access_var]
+        else:
+            raise ValueError("Not implemented")
+
+        dim_mapping = mapping["dimensions"]
+        _ = (
+            self.supergrid.h_cells()
+            if "xt_ocean" in var.dims
+            else self.supergrid.q_cells()
+        )
+        axes = {dim_mapping.get(axis, axis): axis for axis in var.dims}
+
+        i_axis = axes.pop("longitude")
+        i_axis = ds[i_axis].values
+        j_axis = axes.pop("latitude")
+        j_axis = ds[j_axis].values
+        x = np.arange(i_axis.size, dtype="float")
+        x_bnds = np.array([[x_ - 0.5, x_ + 0.5] for x_ in x])
+        y = np.arange(j_axis.size, dtype="float")
+        y_bnds = np.array([[y_ - 0.5, y_ + 0.5] for y_ in y])
+
+        data = var.values
+        lat = self.supergrid.lat
+        lat_bnds = self.supergrid.lat_bnds
+
+        lon = self.supergrid.lon
+        lon_bnds = self.supergrid.lon_bnds
+
+        # Convert time to numeric values
+        time_axis = axes.pop("time")
+        time_numeric = ds[time_axis].values
+        time_units = ds[time_axis].attrs["units"]
+        time_bnds = ds[ds[time_axis].attrs["bounds"]].values
+        # TODO: Check that the calendar is the same than the one defined in the model.json
+        # Convert if not.
+        # calendar = ds[time_axis].attrs["calendar"]
+
+        # CMOR setup
+        ipth = "Test"
+        cmor.setup(
+            inpath=ipth,
+            set_verbosity=cmor.CMOR_NORMAL,
+            netcdf_file_action=cmor.CMOR_REPLACE,
+        )
+
+        cmor.dataset_json(cmor_dataset_json)
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        grid_table_id = cmor.load_table(
+            os.path.join(current_dir, "cmor_tables", "CMIP6_grids.json")
+        )
+        cmor.set_table(grid_table_id)
+
+        cmor_axes = []
+        # Define CMOR axes
+        yaxis_id = cmor.axis(
+            table_entry="j_index", units="1", coord_vals=y, cell_bounds=y_bnds
+        )
+        xaxis_id = cmor.axis(
+            table_entry="i_index", units="1", coord_vals=x, cell_bounds=x_bnds
+        )
+
+        grid_id = cmor.grid(
+            axis_ids=np.array([yaxis_id, xaxis_id]),
+            latitude=lat,
+            longitude=lon,
+            latitude_vertices=lat_bnds,
+            longitude_vertices=lon_bnds,
+        )
+        cmor_axes.append(grid_id)
+
+        # Now, load the Omon table to set up the time axis and variable
+        current_dir = os.path.dirname(os.path.abspath(__file__))
+        mip_table = os.path.join(current_dir, "cmor_tables", mip_table)
+        omon_table_id = cmor.load_table(mip_table)
+        cmor.set_table(omon_table_id)
+
+        cmorTime = cmor.axis(
+            "time", coord_vals=time_numeric, cell_bounds=time_bnds, units=time_units
+        )
+        cmor_axes.append(cmorTime)
+
+        if axes:
+            for axis, dim in axes.items():
+                coord_vals = var[dim].values
+                try:
+                    cell_bounds = var[var[dim].attrs["bounds"]].values
+                except KeyError:
+                    cell_bounds = None
+                axis_units = var[dim].attrs["units"]
+                cmor_axis = cmor.axis(
+                    axis,
+                    coord_vals=coord_vals,
+                    cell_bounds=cell_bounds,
+                    units=axis_units,
+                )
+                cmor_axes.append(cmor_axis)
+
+        # Define CMOR variable
+        cmorVar = cmor.variable(cmor_name, variable_units, cmor_axes, positive=positive)
+
+        # Write data to CMOR
+        data = np.moveaxis(data, 0, -1)
+        cmor.write(cmorVar, data, ntimes_passed=len(time_numeric))
+
+        # Finalize and save the file
+        filename = cmor.close(cmorVar, file_name=True)
+        print("Stored in:", filename)
+
+        cmor.close()
+
 
 @dataclass
 class ACCESS_OM3_CMIP6(CMIP6_Experiment):
@@ -176,11 +300,11 @@ def cmorise(self, file_paths, compound_name, cmor_dataset_json, mip_table):
         y_bnds = np.array([[y_ - 0.5, y_ + 0.5] for y_ in y])
 
         data = var.values
-        lat = om3_grid.lat
-        lat_bnds = om3_grid.lat_bnds
+        lat = self.supergrid.lat
+        lat_bnds = self.supergrid.lat_bnds
 
-        lon = om3_grid.lon
-        lon_bnds = om3_grid.lon_bnds
+        lon = self.supergrid.lon
+        lon_bnds = self.supergrid.lon_bnds
 
         # Convert time to numeric values
         time_axis = axes.pop("time")