|
| 1 | +"""CLMODIS to clmodis converter |
| 2 | +
|
| 3 | +TODO: Convert this module into the new handlers API (handler.py, handlers.yaml, |
| 4 | +formulas.py). |
| 5 | +""" |
| 6 | + |
| 7 | +from __future__ import absolute_import, annotations, division, unicode_literals |
| 8 | + |
| 9 | +import logging |
| 10 | +from typing import Dict, List, Union |
| 11 | + |
| 12 | +import cmor |
| 13 | +import numpy as np |
| 14 | +import xarray as xr |
| 15 | + |
| 16 | +from e3sm_to_cmip._logger import _setup_child_logger |
| 17 | +from e3sm_to_cmip.util import print_message, setup_cmor |
| 18 | + |
| 19 | +logger = _setup_child_logger(__name__) |
| 20 | + |
| 21 | +# list of raw variable names needed |
| 22 | +RAW_VARIABLES = [str("CLMODIS")] |
| 23 | +VAR_NAME = str("clmodis") |
| 24 | +VAR_UNITS = str("%") |
| 25 | +TABLE = str("CMIP6_CFmon.json") |
| 26 | + |
| 27 | + |
| 28 | +def handle( # noqa: C901 |
| 29 | + vars_to_filepaths: Dict[str, List[str]], |
| 30 | + tables: str, |
| 31 | + metadata_path: str, |
| 32 | + cmor_log_dir: str, |
| 33 | + table: str | None, |
| 34 | +) -> str | None: |
| 35 | + """Transform E3SM.CLMODIS into CMIP.clmodis |
| 36 | +
|
| 37 | + clmodis = CLMODIS |
| 38 | +
|
| 39 | + NOTE: `table` is not used and is a placeholder because VarHandler.cmorize |
| 40 | + also includes the `freq` arg. |
| 41 | +
|
| 42 | + Parameters |
| 43 | + ---------- |
| 44 | + vars_to_filepaths : Dict[str, List[str]] |
| 45 | + A dictionary mapping E3SM raw variables to a list of filepath(s). |
| 46 | + tables_path : str |
| 47 | + The path to directory containing CMOR Tables directory. |
| 48 | + metadata_path : str |
| 49 | + The path to user json file for CMIP6 metadata |
| 50 | + cmor_log_dir : str |
| 51 | + The directory that stores the CMOR logs. |
| 52 | + table : str | None |
| 53 | + The CMOR table filename, derived from a custom `freq`, by default |
| 54 | + None. |
| 55 | +
|
| 56 | + Returns |
| 57 | + ------- |
| 58 | + str | None |
| 59 | + If CMORizing was successful, return the output CMIP variable name |
| 60 | + to indicate success. If failed, return None |
| 61 | + .""" |
| 62 | + logging.info(f"Starting {VAR_NAME}") |
| 63 | + |
| 64 | + nonzero = False |
| 65 | + for variable in RAW_VARIABLES: |
| 66 | + if len(vars_to_filepaths[variable]) == 0: |
| 67 | + msg = f"{variable}: Unable to find input files for {RAW_VARIABLES}" |
| 68 | + print_message(msg) |
| 69 | + logging.error(msg) |
| 70 | + nonzero = True |
| 71 | + if nonzero: |
| 72 | + return None |
| 73 | + |
| 74 | + msg = f"{VAR_NAME}: running with input files: {vars_to_filepaths}" |
| 75 | + logger.debug(msg) |
| 76 | + |
| 77 | + setup_cmor( |
| 78 | + var_name=VAR_NAME, |
| 79 | + table_path=tables, |
| 80 | + table_name=TABLE, |
| 81 | + user_input_path=metadata_path, |
| 82 | + cmor_log_dir=cmor_log_dir, |
| 83 | + ) |
| 84 | + |
| 85 | + msg = f"{VAR_NAME}: CMOR setup complete" |
| 86 | + logger.info(msg) |
| 87 | + |
| 88 | + data: Dict[str, Union[np.ndarray, xr.DataArray]] = {} |
| 89 | + |
| 90 | + # assuming all year ranges are the same for every variable |
| 91 | + num_files_per_variable = len(vars_to_filepaths["CLMODIS"]) |
| 92 | + |
| 93 | + # sort the input files for each variable |
| 94 | + vars_to_filepaths["CLMODIS"].sort() |
| 95 | + |
| 96 | + for index in range(num_files_per_variable): |
| 97 | + ds = xr.open_dataset(vars_to_filepaths["CLMODIS"][index], decode_times=False) |
| 98 | + |
| 99 | + tau = ds["cosp_tau_modis"].values |
| 100 | + tau[-1] = 100.0 |
| 101 | + tau_bnds = ds["cosp_tau_modis_bnds"].values |
| 102 | + tau_bnds[-1] = [60.0, 100000.0] |
| 103 | + |
| 104 | + # Units of cosp_pr changed from hPa to Pa |
| 105 | + unit_conv_fact = 1 |
| 106 | + if ds["cosp_prs"].units == "hPa": |
| 107 | + unit_conv_fact = 100 |
| 108 | + |
| 109 | + # load |
| 110 | + data = { |
| 111 | + "CLMODIS": ds["CLMODIS"].values, |
| 112 | + "lat": ds["lat"], |
| 113 | + "lon": ds["lon"], |
| 114 | + "lat_bnds": ds["lat_bnds"], |
| 115 | + "lon_bnds": ds["lon_bnds"], |
| 116 | + "time": ds["time"].values, |
| 117 | + "time_bnds": ds["time_bnds"].values, |
| 118 | + "plev7c": ds["cosp_prs"].values * unit_conv_fact, |
| 119 | + "plev7c_bnds": ds["cosp_prs_bnds"].values * unit_conv_fact, |
| 120 | + "tau": tau, |
| 121 | + "tau_bnds": tau_bnds, |
| 122 | + } |
| 123 | + |
| 124 | + # create the cmor variable and axis |
| 125 | + axes = [ |
| 126 | + {str("table_entry"): str("time"), str("units"): ds["time"].units}, |
| 127 | + { |
| 128 | + str("table_entry"): str("plev7c"), |
| 129 | + str("units"): str("Pa"), |
| 130 | + str("coord_vals"): data["plev7c"], |
| 131 | + str("cell_bounds"): data["plev7c_bnds"], |
| 132 | + }, |
| 133 | + { |
| 134 | + str("table_entry"): str("tau"), |
| 135 | + str("units"): str("1"), |
| 136 | + str("coord_vals"): data["tau"], |
| 137 | + str("cell_bounds"): data["tau_bnds"], |
| 138 | + }, |
| 139 | + { |
| 140 | + str("table_entry"): str("latitude"), |
| 141 | + str("units"): ds["lat"].units, |
| 142 | + str("coord_vals"): data["lat"].values, # type: ignore |
| 143 | + str("cell_bounds"): data["lat_bnds"].values, # type: ignore |
| 144 | + }, |
| 145 | + { |
| 146 | + str("table_entry"): str("longitude"), |
| 147 | + str("units"): ds["lon"].units, |
| 148 | + str("coord_vals"): data["lon"].values, # type: ignore |
| 149 | + str("cell_bounds"): data["lon_bnds"].values, # type: ignore |
| 150 | + }, |
| 151 | + ] |
| 152 | + |
| 153 | + axis_ids = list() |
| 154 | + for axis in axes: |
| 155 | + axis_id = cmor.axis(**axis) |
| 156 | + axis_ids.append(axis_id) |
| 157 | + |
| 158 | + varid = cmor.variable(VAR_NAME, VAR_UNITS, axis_ids) |
| 159 | + |
| 160 | + # write out the data |
| 161 | + msg = f"{VAR_NAME}: time {data['time_bnds'][0][0]:1.1f} - {data['time_bnds'][-1][-1]:1.1f}" |
| 162 | + logger.info(msg) |
| 163 | + |
| 164 | + cmor.write( |
| 165 | + varid, |
| 166 | + data["CLMODIS"], |
| 167 | + time_vals=data["time"], |
| 168 | + time_bnds=data["time_bnds"], |
| 169 | + ) |
| 170 | + |
| 171 | + msg = f"{VAR_NAME}: write complete, closing" |
| 172 | + logger.info(msg) |
| 173 | + |
| 174 | + cmor.close() |
| 175 | + msg = f"{VAR_NAME}: file close complete" |
| 176 | + logger.info(msg) |
| 177 | + |
| 178 | + return VAR_NAME |
0 commit comments