Create UM SST and sea ice ancillaries from ERA data (#11)

Author: Scott Wales

conda/meta.yaml

@@ -16,7 +16,7 @@ requirements:
         - python
         - cdo>=1.9.4
         - cfunits
-        - dask
+        - dask>=1.2.0
         - mule
         - netcdf4
         - numpy
@@ -25,6 +25,7 @@ requirements:
         - xarray
         - whichcraft # [py2k]
         - hdf5>=1.10.1 # Missing dependency of CDO
+        - iris
 
 test:
     requires:

scripts/um/era_sst.py

@@ -0,0 +1,83 @@
+#!/usr/bin/env python
+#
+# Copyright 2019 Scott Wales
+#
+# Author: Scott Wales <[email protected]>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+Set up ERA-Interim SSTs and sea ice for a UM run
+"""
+
+from coecms.regrid import esmf_generate_weights, regrid
+import argparse
+import xarray
+import iris
+from dask.diagnostics import ProgressBar
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument('--start-date', help='ISO-formatted start date')
+    parser.add_argument('--end-date', help='ISO-formatted end date')
+    parser.add_argument('--output', '-o', help='Output file name', required=True)
+    parser.add_argument('--target-mask', help='Target UM land mask', required=True)
+    parser.add_argument('--frequency', choices=[6, 12, 24],
+                        type=int, help='Update frequency (hours)', default=24)
+    args = parser.parse_args()
+
+    # Read in the source mask
+    tos = xarray.open_mfdataset('/g/data1a/ub4/erai/netcdf/6hr/ocean/'
+                                'oper_an_sfc/v01/tos/'
+                                'tos_6hrs_ERAI_historical_an-sfc_2001*.nc',
+                                coords='all')
+    src_mask = tos.tos.isel(time=0)
+
+    # Read in the target mask
+    mask_iris = iris.load_cube(args.target_mask, iris.AttributeConstraint(STASH='m01s00i030'))
+    mask_iris.coord('latitude').var_name = 'lat'
+    mask_iris.coord('longitude').var_name = 'lon'
+    tgt_mask = xarray.DataArray.from_iris(mask_iris).load()
+    tgt_mask = tgt_mask.where(tgt_mask == 0)
+
+    tgt_mask.lon.attrs['standard_name'] = 'longitude'
+    tgt_mask.lat.attrs['standard_name'] = 'latitude'
+    tgt_mask.lon.attrs['units'] = 'degrees_east'
+    tgt_mask.lat.attrs['units'] = 'degrees_north'
+
+    print(tgt_mask)
+
+    weights = esmf_generate_weights(src_mask, tgt_mask, method='patch')
+
+    with ProgressBar():
+
+        # Read and slice the source data
+        tos = xarray.open_mfdataset('/g/data1a/ub4/erai/netcdf/6hr/ocean/'
+                                    'oper_an_sfc/v01/tos/'
+                                    'tos_6hrs_ERAI_historical_an-sfc_2001*.nc',
+                                    coords='all')
+        sic = xarray.open_mfdataset('/g/data1a/ub4/erai/netcdf/6hr/seaIce/'
+                                    'oper_an_sfc/v01/sic/'
+                                    'sic_6hrs_ERAI_historical_an-sfc_2001*.nc',
+                                    coords='all')
+        ds = xarray.Dataset({'tos': tos.tos, 'sic': sic.sic})
+        ds = ds.sel(time=slice(args.start_date, args.end_date))
+        print(ds)
+
+        newds = regrid(ds, weights=weights)
+
+        newds['time'] = newds['time'].astype('i4')
+        newds.to_netcdf(args.output)
+
+if __name__ == '__main__':
+    main()

setup.py

@@ -24,5 +24,6 @@
             ],
         entry_points = {
             'console_scripts': [
+                'coecms=coecms.cli:cli',
                 ]}
         )

src/coecms/cli/__init__.py

@@ -0,0 +1,20 @@
+#!/usr/bin/env python
+#
+# Copyright 2019 Scott Wales
+#
+# Author: Scott Wales <[email protected]>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .main import cli
+from .um import um

src/coecms/cli/main.py

@@ -0,0 +1,23 @@
+#!/usr/bin/env python
+#
+# Copyright 2019 Scott Wales
+#
+# Author: Scott Wales <[email protected]>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import click
+
+@click.group()
+def cli():
+    pass

src/coecms/cli/um.py

@@ -0,0 +1,106 @@
+#!/usr/bin/env python
+#
+# Copyright 2019 Scott Wales
+#
+# Author: Scott Wales <[email protected]>
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from .main import cli
+from ..grid import UMGrid
+from ..regrid import regrid, esmf_generate_weights
+from ..um.create_ancillary import create_surface_ancillary
+import click
+import pandas
+import mule
+import iris
+import xarray
+from dask.diagnostics import ProgressBar
+import dask.distributed
+import matplotlib.pyplot as plt
+
+@cli.group()
+def um():
+    """
+    Tools for working with the Unified Model
+    """
+    pass
+
+@um.group()
+def ancil():
+    """
+    Tools for working with ancil files
+    """
+    pass
+
+def validate_date(ctx, param, value):
+    """
+    Ensures an argument is a valid date
+    """
+    try:
+        return pandas.to_datetime(value, utc=True, dayfirst=True)
+    except ValueError:
+        raise click.BadParameter(f'unable to parse "{value}" as a date')
+
+def validate_um_ancil(ctx, param, value):
+    """
+    Ensures an argument is a UM file
+    """
+    try:
+        return mule.AncilFile.from_file(value)
+    except:
+        raise click.BadParameter(f'"{value}" does not seem to be a UM ancil file')
+
+@ancil.command()
+@click.option('--start-date', callback=validate_date, required=True)
+@click.option('--end-date', callback=validate_date, required=True)
+@click.option('--target-mask',
+              type=click.Path(exists=True, dir_okay=False))
+@click.option('--output', required=True,
+              type=click.Path(writable=True, dir_okay=False))
+def era_sst(start_date, end_date, target_mask, output):
+    """
+    Create ancil files from ERA reanalysis data
+    """
+
+    um_grid = UMGrid.from_mask(target_mask)
+
+    file_start = start_date - pandas.offsets.MonthBegin()
+    file_end = end_date + pandas.offsets.MonthEnd()
+    file_a = pandas.date_range(file_start,file_end,freq='MS')
+    file_b = file_a + pandas.offsets.MonthEnd()
+
+    dates = [f'{a.strftime("%Y%m%d")}_{b.strftime("%Y%m%d")}'
+             for a,b in zip(file_a, file_b)]
+
+    # Read and slice the source data
+    tos = xarray.open_mfdataset(['/g/data1a/ub4/erai/netcdf/6hr/ocean/'
+                                 'oper_an_sfc/v01/tos/'
+                                 'tos_6hrs_ERAI_historical_an-sfc_'+d+'.nc'
+                                 for d in dates],
+                                 chunks={'time': 1,})
+    sic = xarray.open_mfdataset(['/g/data1a/ub4/erai/netcdf/6hr/seaIce/'
+                                 'oper_an_sfc/v01/sic/'
+                                 'sic_6hrs_ERAI_historical_an-sfc_'+d+'.nc'
+                                 for d in dates],
+                                 chunks={'time': 1,})
+    ds = xarray.Dataset({'tos': tos.tos, 'sic': sic.sic})
+    ds = ds.sel(time=slice(start_date, end_date))
+
+    weights = esmf_generate_weights(tos.tos.isel(time=0), um_grid, method='patch')
+    newds = regrid(ds, weights=weights)
+
+    print(newds)
+
+    ancil = create_surface_ancillary(newds, {'tos': 24, 'sic': 31})
+    ancil.to_file(output)

src/coecms/grid.py

@@ -19,6 +19,7 @@
 import six
 import xarray
 import numpy
+import iris
 
 """
 Different grid types
@@ -174,6 +175,24 @@ def to_scrip(self):
         return scrip
 
 
+class UMGrid(LonLatGrid):
+    @classmethod
+    def from_mask(cls, mask_path):
+        mask = iris.load_cube(mask_path, iris.AttributeConstraint(STASH='m01s00i030'))
+        mask.coord('latitude').var_name = 'lat'
+        mask.coord('longitude').var_name = 'lon'
+
+        mask = xarray.DataArray.from_iris(mask).load()
+        mask = mask.where(mask == 0)
+
+        mask.lon.attrs['standard_name'] = 'longitude'
+        mask.lat.attrs['standard_name'] = 'latitude'
+        mask.lon.attrs['units'] = 'degrees_east'
+        mask.lat.attrs['units'] = 'degrees_north'
+
+        return mask
+
+
 class ScripGrid(Grid):
     def __init__(self, grid):
         self._grid = grid