diff --git a/fre/__init__.py b/fre/__init__.py index d762baf0..d0903ee0 100644 --- a/fre/__init__.py +++ b/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 * \ No newline at end of file diff --git a/fre/app/__init__.py b/fre/app/__init__.py new file mode 100644 index 00000000..8715de49 --- /dev/null +++ b/fre/app/__init__.py @@ -0,0 +1,3 @@ +from .maskAtmosPlevel import maskAtmosPlevel_subtool + +__all__ = ["maskAtmosPlevel_subtool"] diff --git a/fre/app/maskAtmosPlevel.py b/fre/app/maskAtmosPlevel.py new file mode 100755 index 00000000..fd6d5896 --- /dev/null +++ b/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 diff --git a/fre/fre.py b/fre/fre.py index efc54545..0869feaa 100644 --- a/fre/fre.py +++ b/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) + ############################################# """ diff --git a/meta.yaml b/meta.yaml index 2782765d..f831c92b 100644 --- a/meta.yaml +++ b/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