Pressure masking netcdf modifying tool

fre/__init__.py

@@ -6,4 +6,5 @@
 from fre.frerun import *
 from fre.frecatalog import *
 from fre.freyamltools import *
-from .cmor import *
+from .app import *
+from .cmor import *

fre/app/__init__.py

@@ -0,0 +1,3 @@
+from .maskAtmosPlevel import maskAtmosPlevel_subtool
+
+__all__ = ["maskAtmosPlevel_subtool"]

fre/app/maskAtmosPlevel.py

@@ -0,0 +1,154 @@
+#!/usr/bin/env python
+
+# This script contains the refineDiags that produce data at the same
+# frequency as the input data (no reduction) such as surface albedo,
+# masking fields,...
+# It can accept any file and will only compute refineDiags in fields
+# are present.
+
+import os
+import netCDF4 as nc
+import xarray as xr
+import click
+
+@click.command()
+def maskAtmosPlevel_subtool(infile, outfile, psfile):
+    # Error if outfile exists
+    if os.path.exists(outfile):
+        raise Exception(f"ERROR: Output file '{outfile}' already exists")
+
+    # Open ps dataset
+    if not os.path.exists(psfile):
+        raise Exception(f"ERROR: Input surface pressure file '{psfile}' does not exist")
+    ds_ps = xr.open_dataset(psfile)
+
+    # Exit with message if "ps" not available
+    if "ps" not in list(ds_ps.variables):
+        print(f"WARNING: File '{infile}' does not contain surface pressure, so exiting")
+        return None
+
+    # Open input dataset
+    if not os.path.exists(infile):
+        raise Exception(f"ERROR: Input file '{infile}' does not exist")
+    ds_in = xr.open_dataset(infile)
+
+    # The trigger for atmos masking is a variable attribute "needs_atmos_masking = True".
+    # In the future this will be set within the model, but for now and testing,
+    # we'll add the attribute for variables that end with "_unmsk".
+    ds_in = preprocess(ds_in)
+
+    ds_out = xr.Dataset()
+
+    # Process all variables with attribute "needs_atmos_masking = True"
+    for var in list(ds_in.variables):
+        if 'needs_atmos_masking' in ds_in[var].attrs:
+            del ds_in[var].attrs['needs_atmos_masking']
+            ds_out[var] = mask_field_above_surface_pressure(ds_in, var, ds_ps)
+        else:
+            continue
+
+    # Write the output file if anything was done
+    if ds_out.variables:
+        print(f"Modifying variables '{list(ds_out.variables)}', appending into new file '{outfile}'")
+        write_dataset(ds_out, ds_in, outfile)
+    else:
+        print(f"No variables modified, so not writing output file '{outfile}'")
+    return None
+
+
+def preprocess(ds):
+    """add needs_atmos_masking attribute if var ends with _unmsk"""
+
+    for var in list(ds.variables):
+        if var.endswith('_unmsk'):
+            ds[var].attrs['needs_atmos_masking'] = True
+
+    return ds
+
+
+def mask_field_above_surface_pressure(ds, var, ds_ps):
+    """mask data with pressure larger than surface pressure"""
+
+    plev = pressure_coordinate(ds, var)
+
+    # broadcast pressure coordinate and surface pressure to
+    # the dimensions of the variable to mask
+    plev_extended, _ = xr.broadcast(plev, ds[var])
+    ps_extended, _ = xr.broadcast(ds_ps["ps"], ds[var])
+    # masking do not need looping
+    masked = xr.where(plev_extended > ps_extended, 1.0e20, ds[var])
+    # copy attributes, but it doesn't include the missing values
+    attrs = ds[var].attrs.copy()
+    # add the missing values back
+    attrs['missing_value'] = 1.0e20
+    attrs['_FillValue'] = 1.0e20
+    masked.attrs = attrs
+    # transpose dims like the original array
+    masked = masked.transpose(*ds[var].dims)
+
+    print(f"Processed {var}")
+
+    return masked
+
+
+def pressure_coordinate(ds, varname, verbose=False):
+    """check if dataArray has pressure coordinate fitting requirements
+    and return it"""
+
+    pressure_coord = None
+
+    for dim in list(ds[varname].dims):
+        if dim in list(ds.variables):  # dim needs to have values in file
+            if ds[dim].attrs["long_name"] == "pressure":
+                pressure_coord = ds[dim]
+            elif ("coordinates" in ds.attrs) and (ds[dim].attrs["units"] == "Pa"):
+                pressure_coord = ds[dim]
+
+    return pressure_coord
+
+
+def write_dataset(ds, template, outfile):
+    """prepare the dataset and dump into netcdf file"""
+
+    # copy global attributes
+    ds.attrs = template.attrs.copy()
+
+    # copy all variables and their attributes
+    # except those already processed
+    for var in list(template.variables):
+        if var in list(ds.variables):
+            continue
+        else:
+            ds[var] = template[var]
+            ds[var].attrs = template[var].attrs.copy()
+
+    ds.to_netcdf(outfile, unlimited_dims="time")
+
+    return None
+
+
+def set_netcdf_encoding(ds, pressure_vars):
+    """set preferred options for netcdf encoding"""
+
+    all_vars = list(ds.variables)
+    encoding = {}
+
+    for var in do_not_encode_vars + pressure_vars:
+        if var in all_vars:
+            encoding.update({var: dict(_FillValue=None)})
+
+    return encoding
+
+
+def post_write(filename, ds, var_with_bounds, bounds_variables):
+    """fix a posteriori attributes that xarray.to_netcdf
+    did not do properly using low level netcdf lib"""
+
+    f = nc.Dataset(filename, "a")
+
+    for var, bndvar in zip(var_with_bounds, bounds_variables):
+        f.variables[var].setncattr("bounds", bndvar)
+
+    f.close()
+
+    return None

fre/fre.py

@@ -31,6 +31,7 @@
 from fre import freyamltools
 from fre.freyamltools.freyamltools import *
 
+from .app import maskAtmosPlevel_subtool 
 from .cmor import run_subtool
 
 #############################################
@@ -91,6 +92,33 @@ def frecmor():
     """ - access fre cmor subcommands"""
     pass
 
+@fre.group('app')
+def freapp():
+    """access fre app subcommands"""
+    pass
+
+#############################################
+
+"""
+fre app subcommands to be processed
+"""
+@freapp.command()
+@click.option("-i", "--infile",
+              type=str,
+              help="Input NetCDF file containing pressure-level output to be masked",
+              required=True)
+@click.option("-o", "--outfile",
+              type=str,
+              help="Output file",
+              required=True)
+@click.option("-p", "--psfile",
+              help="Input NetCDF file containing surface pressure (ps)",
+              required=True)
+@click.pass_context
+def mask_atmos_plevel(context, infile, outfile, psfile):
+    """Mask out pressure level diagnostic output below land surface"""
+    context.forward(maskAtmosPlevel_subtool)
+
 #############################################
 
 """

meta.yaml

@@ -42,6 +42,7 @@ test:
     - fre.frepp.status
     - fre.frepp.run
     - fre.frepp.validate
+    - fre.app
     - fre.cmor
   commands:
     - fre --help
@@ -50,6 +51,8 @@ test:
     - fre pp status --help
     - fre pp run --help
     - fre pp validate --help
+    - fre app --help
+    - fre app mask-atmos-plevel --help
     - fre cmor --help
     - fre cmor run --help