Merge pull request #872 from vanroekel/add_woa23_salinity_restoring

Adds WOA 23 SSS restoring file script

compass/ocean/tests/global_ocean/files_for_e3sm/remap_sea_surface_salinity_restoring.py

@@ -29,11 +29,11 @@ def __init__(self, test_case):
                          ntasks=512, min_tasks=1)
 
         self.add_input_file(
-            filename='PHC2_salx.2004_08_03.filled_double_precision.nc',
-            target='PHC2_salx.2004_08_03.filled_double_precision.nc',
+            filename='woa23_decav_0.25_sss_monthly_extrap.20241101.nc',
+            target='woa23_decav_0.25_sss_monthly_extrap.20241101.nc',
             database='initial_condition_database')
 
-        self.add_output_file(filename='sss.PHC2_monthlyClimatology.nc')
+        self.add_output_file(filename='sss.WOA23_monthlyClimatology.nc')
 
     def run(self):
         """
@@ -44,9 +44,9 @@ def run(self):
         config = self.config
         ntasks = self.ntasks
 
-        in_filename = 'PHC2_salx.2004_08_03.filled_double_precision.nc'
+        in_filename = 'woa23_decav_0.25_sss_monthly_extrap.20241101.nc'
 
-        prefix = 'sss.PHC2_monthlyClimatology'
+        prefix = 'sss.WOA23_monthlyClimatology'
         suffix = f'{self.mesh_short_name}.{self.creation_date}'
 
         remapped_filename = f'{prefix}.nc'

compass/ocean/tests/utility/__init__.py

@@ -1,4 +1,7 @@
 from compass.ocean.tests.utility.combine_topo import CombineTopo
+from compass.ocean.tests.utility.create_salin_restoring import (
+    CreateSalinRestoring,
+)
 from compass.ocean.tests.utility.cull_restarts import CullRestarts
 from compass.ocean.tests.utility.extrap_woa import ExtrapWoa
 from compass.testgroup import TestGroup
@@ -19,3 +22,4 @@ def __init__(self, mpas_core):
         self.add_test_case(CombineTopo(test_group=self))
         self.add_test_case(CullRestarts(test_group=self))
         self.add_test_case(ExtrapWoa(test_group=self))
+        self.add_test_case(CreateSalinRestoring(test_group=self))

compass/ocean/tests/utility/create_salin_restoring/__init__.py

@@ -0,0 +1,29 @@
+from compass.ocean.tests.utility.create_salin_restoring.extrap_salin import (
+    ExtrapSalin,
+)
+from compass.ocean.tests.utility.create_salin_restoring.salinity_restoring import (  # noqa: E501
+    Salinity,
+)
+from compass.testcase import TestCase
+
+
+class CreateSalinRestoring(TestCase):
+    """
+    A test case for first creating monthly sea surface salinity from WOA23
+    then extrapolating the WOA23 data into missing ocean regions such as
+    ice-shelf cavities and coasts
+    """
+
+    def __init__(self, test_group):
+        """
+        Create the test case
+
+        Parameters
+        ----------
+        test_group : compass.ocean.tests.utility.Utility
+            The test group that this test case belongs to
+        """
+        super().__init__(test_group=test_group, name='create_salin_restoring')
+
+        self.add_step(Salinity(test_case=self))
+        self.add_step(ExtrapSalin(test_case=self))

compass/ocean/tests/utility/create_salin_restoring/extrap_salin.py

@@ -0,0 +1,122 @@
+from datetime import datetime
+
+import numpy as np
+import xarray as xr
+from scipy.signal import convolve2d
+
+from compass.step import Step
+
+
+class ExtrapSalin(Step):
+    """
+    Extrapolate WOA 2023 monthly sea surface salinity data into missing ocean
+    regions, including ice cavities and coasts
+
+    Attributes
+    ----------
+    woa_filename : str
+        The name of the output file name after extrapolation
+
+    """
+    def __init__(self, test_case):
+        """
+        Create a new test case
+
+        Parameters
+        ----------
+        test_case : compass.ocean.tests.utility.create_salin_restoring.
+        CreateSalinRestoring
+            The test case this step belongs to
+
+        """
+        super().__init__(test_case=test_case, name='extrap', cpus_per_task=64,
+                         min_cpus_per_task=1, openmp_threads=1)
+
+        self.add_input_file(
+            filename='woa_surface_salinity_monthly.nc',
+            target='../salinity_restoring/woa_surface_salinity_monthly.nc')
+
+        self.woa_filename = None
+
+    def setup(self):
+        """
+        Determine the output filename
+        """
+
+        now = datetime.now()
+
+        datestring = now.strftime("%Y%m%d")
+
+        self.woa_filename = f'woa23_decav_0.25_sss_monthly_extrap.' \
+                            f'{datestring}.nc'
+        self.add_output_file(self.woa_filename)
+
+    def run(self):
+        """
+        Extrapolate WOA 2023 model temperature and salinity into ice-shelf
+        cavities.
+        """
+        # extrapolate horizontally using the ocean mask
+        _extrap(self.woa_filename)
+
+
+def _extrap(out_filename):
+
+    in_filename = 'woa_surface_salinity_monthly.nc'
+    ds = xr.open_dataset(in_filename)
+
+    field = ds.SALT.values.copy()
+
+    # a small averaging kernel
+    x = np.arange(-1, 2)
+    x, y = np.meshgrid(x, x)
+    kernel = np.exp(-0.5 * (x**2 + y**2))
+
+    threshold = 0.01
+    nlon = field.shape[-1]
+
+    lon_with_halo = np.array([nlon - 2, nlon - 1] + list(range(nlon)) + [0, 1])
+    lon_no_halo = list(range(2, nlon + 2))
+
+    for i in range(12):
+        valid = np.isfinite(field[i, :, :])
+        orig_mask = valid
+        prev_fill_count = 0
+        while True:
+            valid_weight_sum = _extrap_with_halo(valid, kernel, valid,
+                                                 lon_with_halo, lon_no_halo)
+
+            new_valid = valid_weight_sum > threshold
+
+            # don't want to overwrite original data but do want ot smooth
+            # extrapolated data
+            fill_mask = np.logical_and(new_valid, np.logical_not(orig_mask))
+
+            fill_count = np.count_nonzero(fill_mask)
+            if fill_count == prev_fill_count:
+                # no change so we're done
+                break
+
+            field_extrap = _extrap_with_halo(field[i, :, :], kernel, valid,
+                                             lon_with_halo, lon_no_halo)
+            field[i, fill_mask] = field_extrap[fill_mask] / \
+                valid_weight_sum[fill_mask]
+
+            valid = new_valid
+            prev_fill_count = fill_count
+
+    attrs = ds.SALT.attrs
+    dims = ds.SALT.dims
+    ds['SALT'] = (dims, field)
+    ds['SALT'].attrs = attrs
+
+    ds.to_netcdf(out_filename)
+
+
+def _extrap_with_halo(field, kernel, valid, lon_with_halo, lon_no_halo):
+    field = field.copy()
+    field[np.logical_not(valid)] = 0.
+    field_with_halo = field[:, lon_with_halo]
+    field_extrap = convolve2d(field_with_halo, kernel, mode='same')
+    field_extrap = field_extrap[:, lon_no_halo]
+    return field_extrap

compass/ocean/tests/utility/create_salin_restoring/salinity_restoring.py

@@ -0,0 +1,113 @@
+import numpy as np
+import xarray as xr
+from mpas_tools.io import write_netcdf
+
+from compass.step import Step
+
+
+class Salinity(Step):
+    """
+    A step for combining January through December sea surface salinity
+    data into a single file for salinity restoring in G-cases.
+    The top level data of the monthly WOA23 is utilized.
+    """
+
+    def __init__(self, test_case):
+        """
+        Create a new step
+
+        Parameters
+        ----------
+        test_case : compass.ocean.tests.utility.create_salin_restoring.
+        CreateSalinRestoring
+        The test case this step belongs to
+        """
+        super().__init__(test_case, name='salinity_restoring', ntasks=1,
+                         min_tasks=1)
+        self.add_output_file(filename='woa_surface_salinity_monthly.nc')
+
+    def setup(self):
+        """
+        Set up the step in the work directory, including downloading any
+        dependencies.
+        """
+        super().setup()
+
+        base_url = \
+            'https://www.ncei.noaa.gov/thredds-ocean/fileServer/woa23/DATA'
+
+        woa_dir = 'salinity/netcdf/decav91C0/0.25'
+
+        woa_files = dict(jan='woa23_decav91C0_s01_04.nc',
+                         feb='woa23_decav91C0_s02_04.nc',
+                         mar='woa23_decav91C0_s03_04.nc',
+                         apr='woa23_decav91C0_s04_04.nc',
+                         may='woa23_decav91C0_s05_04.nc',
+                         jun='woa23_decav91C0_s06_04.nc',
+                         jul='woa23_decav91C0_s07_04.nc',
+                         aug='woa23_decav91C0_s08_04.nc',
+                         sep='woa23_decav91C0_s09_04.nc',
+                         octo='woa23_decav91C0_s10_04.nc',
+                         nov='woa23_decav91C0_s11_04.nc',
+                         dec='woa23_decav91C0_s12_04.nc')
+
+        for month in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul',
+                      'aug', 'sep', 'octo', 'nov', 'dec']:
+            woa_filename = woa_files[month]
+            woa_url = f'{base_url}/{woa_dir}/{woa_filename}'
+
+            self.add_input_file(
+                filename=f'woa_salin_{month}.nc',
+                target=woa_filename,
+                database='initial_condition_database',
+                url=woa_url)
+
+    def run(self):
+        """
+        Run this step of the test case
+        """
+        ds_jan = xr.open_dataset('woa_salin_jan.nc', decode_times=False)
+
+        ds_out = xr.Dataset()
+
+        for var in ['lon', 'lat']:
+            ds_out[var] = ds_jan[var]
+            ds_out[f'{var}_bnds'] = ds_jan[f'{var}_bnds']
+
+        slices = list()
+        for month in ['jan', 'feb', 'mar', 'apr', 'may', 'jun', 'jul',
+                      'aug', 'sep', 'octo', 'nov', 'dec']:
+            ds = xr.open_dataset(
+                f'woa_salin_{month}.nc',
+                decode_times=False).isel(depth=0). \
+                drop_vars('depth')
+            slices.append(ds.s_an)
+
+        ds_out['s_an'] = xr.concat(slices, dim='time')
+        ds_out['s_an'].attrs = ds_jan['s_an'].attrs
+
+        # Change names of variables and alter attributes
+        ds_out = ds_out.rename_dims({'time': 'Time'})
+        ds_out = ds_out.rename_vars({'s_an': 'SALT'})
+        ds_out = ds_out.drop_vars('time')
+
+        # Create a time index
+        time_var = np.arange(1, 12.1, 1)
+        ds_out = ds_out.assign(Time=xr.DataArray(time_var, dims=['Time']))
+
+        # Create a xtime array
+        xtime_list = []
+        for i in range(1, 13):
+            xtime_string = f"0000-{i:02d}-15_00:00:00"
+            xtime_list.append(xtime_string)
+        xtime_out = np.array(xtime_list, dtype='S64')
+
+        ds_out = ds_out.assign(xtime=xr.DataArray(xtime_out, dims=['Time']))
+
+        # Change attributes to be consistent with PHC restoring file
+        ds_out.Time.attrs['long_name'] = "Month Index"
+        ds_out.Time.attrs['units'] = "month"
+        ds_out.Time.attrs['axis'] = "T"
+
+        # Save file and change time dimension to unlimited
+        write_netcdf(ds_out, 'woa_surface_salinity_monthly.nc')

docs/developers_guide/ocean/test_groups/global_ocean.rst

@@ -1406,9 +1406,8 @@ The test case is made up of 10 steps:
     ice-shelf melt flux) compsets.
 
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.remap_sea_surface_salinity_restoring.RemapSeaSurfaceSalinityRestoring`
-    is used to remap sea-surface salinity from the Polar science center
-    Hydrographic Climatology (PHC) to the MPAS mesh.  This dataset is used in
-    E3SM for compsets with data atmospheric/land forcing.
+    is used to remap sea-surface salinity from the World Ocean Atlas (2023) to the MPAS mesh.  
+    This dataset is used in E3SM for compsets with data atmospheric/land forcing.
 
 :py:class:`compass.ocean.tests.global_ocean.files_for_e3sm.remap_iceberg_climatology.RemapIcebergClimatology`
     is used to remap freshwater fluxes from an iceberg climatology from

docs/developers_guide/ocean/test_groups/utility.rst

@@ -114,3 +114,30 @@ avoid this problem, we perform extrapolation in 4 steps.
 
 The resulting file is ready to be placed in compass' initial condition database
 (see :ref:`dev_step_input_download` for details on databases).
+
+create_salin_restoring
+----------------------
+The class :py:class:`compass.ocean.tests.utility.create_salin_restoring.CreateSalinRestoring`
+defines a test case for creating a monthly average sea surface salinity dataset based on the 
+`WOA 2023 <https://www.ncei.noaa.gov/products/world-ocean-atlas>`_ data.  It also extrapolates
+the twelve months of data into ice-shelf cavities and across continents.
+
+salinity_restoring
+------------------
+
+The class :py:class:`compass.ocean.tests.utility.create_salin_restoring.Salinity`
+defines a step to download and combine January through December sea surface salinity into a single
+file that serves as the base dataset for salinity restoring in forced ocean sea-ice cases (FOSI).  
+The surface level of WOA 2023 data is utilized.
+
+The reference date for each month of data is assumed to be the 15th of each month.  In a simulation,
+this implies that for a model start time of January 1, the salinity is restored to the average of 
+the December and January sea surface salinities.
+
+extrap_salin
+------------
+
+The class :py:class:`compass.ocean.tests.utility.create_salin_restoring.ExtrapSalin`
+defines a step to extrapolate the combined January through December sea surface salinities 
+into missing ocean regions such as ice-shelf cavities and across continents.  Since this is only
+extrapolation of surface values, masks are not utilized.

docs/users_guide/ocean/test_groups/utility.rst

@@ -32,3 +32,11 @@ ice-shelf cavities and coastal regions, then into land, grounded ice and below
 bathymetry.  It is provided mainly for developers to update for use on
 future datasets and is not intended for users, so we do not provide full
 documentation here.
+
+create_salin_restoring
+----------------------
+The ``ocean/utility/create_salin_restoring`` test case is used to download monthly 
+average `WOA 2023 <https://www.ncei.noaa.gov/products/world-ocean-atlas>`_ data, 
+extracts the surface layer and extrapolates into ice-shelf cavities and across
+continents.  It is provided mainly for developers to update the salinity restoring
+data and is not intended for users, so we do not provide full documentation here.