From 8e006bd9f88c573ce1363ad4bd31e100ae54eb60 Mon Sep 17 00:00:00 2001 From: Bo Dong postdoc Date: Mon, 25 Nov 2024 14:43:41 -0800 Subject: [PATCH] staged --- pcmdi_metrics/monsoon_wang/__init__.py | 2 + .../monsoon_wang/basic_monsoon_wang_param.py | 57 ++++ .../monsoon_wang/monsoon_precip_index_fncs.py | 153 +++++++++ .../monsoon_wang/monsoon_wang_driver.py | 297 ++++++++++++++++++ pcmdi_metrics/monsoon_wang/run.py | 5 + 5 files changed, 514 insertions(+) create mode 100644 pcmdi_metrics/monsoon_wang/__init__.py create mode 100644 pcmdi_metrics/monsoon_wang/basic_monsoon_wang_param.py create mode 100644 pcmdi_metrics/monsoon_wang/monsoon_precip_index_fncs.py create mode 100644 pcmdi_metrics/monsoon_wang/monsoon_wang_driver.py create mode 100644 pcmdi_metrics/monsoon_wang/run.py diff --git a/pcmdi_metrics/monsoon_wang/__init__.py b/pcmdi_metrics/monsoon_wang/__init__.py new file mode 100644 index 000000000..d5e92d462 --- /dev/null +++ b/pcmdi_metrics/monsoon_wang/__init__.py @@ -0,0 +1,2 @@ +from .monsoon_precip_index_fncs import mpd, mpi_skill_scores # noqa +from .monsoon_wang_driver import create_monsoon_wang_parser, monsoon_wang_runner # noqa diff --git a/pcmdi_metrics/monsoon_wang/basic_monsoon_wang_param.py b/pcmdi_metrics/monsoon_wang/basic_monsoon_wang_param.py new file mode 100644 index 000000000..2755464b1 --- /dev/null +++ b/pcmdi_metrics/monsoon_wang/basic_monsoon_wang_param.py @@ -0,0 +1,57 @@ +# +# OPTIONS ARE SET BY USER IN THIS FILE AS INDICATED BELOW BY: +# +# + +# LIST OF MODEL VERSIONS TO BE TESTED + +# trouble model ICON-ESM-LR' + +modnames = [ + "ACCESS-CM2", + "ACCESS-ESM1-5", + "AWI-ESM-1-1-LR", + "BCC-CSM2-MR", + "BCC-ESM1", + "CAMS-CSM1-0", + "CanESM5", + "CAS-ESM2-0", + "CESM2", + "CESM2-WACCM", + "CIESM", + "CMCC-ESM2", + "E3SM-1-0", + "E3SM-1-1", + "E3SM-2-0", + "EC-Earth3", + "FGOALS-f3-L", + "GFDL-CM4", + "GFDL-ESM4", + "GISS-E2-1-G", + "GISS-E2-1-H", + "IPSL-CM5A2-INCA", + "IPSL-CM6A-LR", + "KIOST-ESM", + "MIROC6", + "MPI-ESM1-2-LR", + "MRI-ESM2-0", + "NESM3", + "NorESM2-LM", + "TaiESM1", +] + + +# ROOT PATH FOR MODELS CLIMATOLOGIES +test_data_path = "/p/user_pub/pmp/pmp_results/pmp_v1.1.2/diagnostic_results/CMIP_CLIMS/cmip6/historical/v20230823/pr/cmip6.historical.%(model).r1i1p1f1.mon.pr.198101-200512.AC.v20230823.nc" + +# ROOT PATH FOR OBSERVATIONS +reference_data_path = "/p/user_pub/PCMDIobs/obs4MIPs_legacy/PCMDIobs2_clims/atmos/pr/TRMM-3B43v-7/pr_mon_TRMM-3B43v-7_BE_gn_199801-201712.v20200421.AC.nc" + +# DIRECTORY WHERE TO PUT RESULTS +results_dir = "$OUTPUT_DIR$/monsoon_wang" + +# Threshold +threshold = 2.5 / 86400 + +# monsoon domain mask based on observations +obs_mask = True diff --git a/pcmdi_metrics/monsoon_wang/monsoon_precip_index_fncs.py b/pcmdi_metrics/monsoon_wang/monsoon_precip_index_fncs.py new file mode 100644 index 000000000..b36ebdb1f --- /dev/null +++ b/pcmdi_metrics/monsoon_wang/monsoon_precip_index_fncs.py @@ -0,0 +1,153 @@ +from typing import Union + +import numpy as np +import xarray as xr + +# SEASONAL RANGE - USING ANNUAL CYCLE CLIMATOLGIES 0=Jan, 11=Dec + + +def da_to_ds(d: Union[xr.Dataset, xr.DataArray], var: str = "variable") -> xr.Dataset: + """Convert xarray DataArray to Dataset + + Parameters + ---------- + d : Union[xr.Dataset, xr.DataArray] + Input dataArray. If dataset is given, no process will be done + var : str, optional + Name of dataArray, by default "variable" + + Returns + ------- + xr.Dataset + xarray Dataset + + Raises + ------ + TypeError + Raised when given input is not xarray based variables + """ + if isinstance(d, xr.Dataset): + return d.copy() + elif isinstance(d, xr.DataArray): + return d.to_dataset(name=var).bounds.add_missing_bounds().copy() + else: + raise TypeError( + "Input must be an instance of either xarrary.DataArray or xarrary.Dataset" + ) + + +def regrid(da_in, da_grid, data_var="pr"): + ds_in = da_to_ds(da_in, data_var) + ds_grid = da_to_ds(da_grid, data_var) + + ds_out = ds_in.regridder.horizontal(data_var, ds_grid, tool="regrid2") + da_out = ds_out[data_var] + + return da_out + + +def compute_season(data, season_indices, weights): + out = np.ma.zeros(data.shape[1:], dtype=data.dtype) + N = 0 + for i in season_indices: + out += data[i] * weights[i] + N += weights[i] + return out / N + + +def mpd(data): + """Monsoon precipitation intensity and annual range calculation + .. describe:: Input + * data + * Assumes climatology array with 12 times step first one January + """ + months_length = [ + 31.0, + 28.0, + 31.0, + 30.0, + 31.0, + 30.0, + 31.0, + 31.0, + 30.0, + 31.0, + 30.0, + 31.0, + ] + mjjas = compute_season(data, [4, 5, 6, 7, 8], months_length) + ndjfm = compute_season(data, [10, 11, 0, 1, 2], months_length) + ann = compute_season(data, list(range(12)), months_length) + + data_map = data.isel(time=0) + + annrange = mjjas - ndjfm + + da_annrange = xr.DataArray(annrange, coords=data_map.coords, dims=data_map.dims) + da_annrange = da_annrange.where(da_annrange.lat >= 0, da_annrange * -1) + + mpi = np.divide(da_annrange.values, ann, where=ann.astype(bool)) + + da_mpi = xr.DataArray(mpi, coords=data_map.coords, dims=data_map.dims) + + return da_annrange, da_mpi + + +def mpi_skill_scores(annrange_mod_dom, annrange_obs_dom, threshold=2.5 / 86400.0): + """Monsoon precipitation index skill score calculation + see Wang et al., doi:10.1007/s00382-010-0877-0 + + .. describe:: Input + + * annrange_mod_dom + + * Model Values Range (summer - winter) + + * annrange_obs_dom + + * Observations Values Range (summer - winter) + + * threshold [default is 2.5/86400.] + + * threshold in same units as inputs + """ + mt = np.ma.greater(annrange_mod_dom, threshold) + ot = np.ma.greater(annrange_obs_dom, threshold) + + hitmap = mt * ot # only where both mt and ot are True + hit = float(hitmap.sum()) + + xor = np.ma.logical_xor(mt, ot) + missmap = xor * ot + missed = float(missmap.sum()) + + falarmmap = xor * mt + falarm = float(falarmmap.sum()) + + if (hit + missed + falarm) > 0.0: + score = hit / (hit + missed + falarm) + else: + score = 1.0e20 + + xr_hitmap = xr.DataArray( + hitmap, + name="hitmap", + coords={"lat": annrange_mod_dom.lat, "lon": annrange_mod_dom.lon}, + dims=["lat", "lon"], + ) + + xr_missmap = xr.DataArray( + missmap, + name="missmap", + coords={"lat": annrange_mod_dom.lat, "lon": annrange_mod_dom.lon}, + dims=["lat", "lon"], + ) + + xr_falarmmap = xr.DataArray( + falarmmap, + name="falarmmap", + coords={"lat": annrange_mod_dom.lat, "lon": annrange_mod_dom.lon}, + dims=["lat", "lon"], + ) + + return hit, missed, falarm, score, xr_hitmap, xr_missmap, xr_falarmmap diff --git a/pcmdi_metrics/monsoon_wang/monsoon_wang_driver.py b/pcmdi_metrics/monsoon_wang/monsoon_wang_driver.py new file mode 100644 index 000000000..623764a1a --- /dev/null +++ b/pcmdi_metrics/monsoon_wang/monsoon_wang_driver.py @@ -0,0 +1,297 @@ +#!/usr/bin/env python + +import collections +import os +import sys + +import cdms2 +import numpy +import numpy as np +import xarray as xr +from monsoon_precip_index_fncs import mpd, mpi_skill_scores, regrid + +import pcmdi_metrics +from pcmdi_metrics import resources +from pcmdi_metrics.io import load_regions_specs, region_subset +from pcmdi_metrics.mean_climate.lib.pmp_parser import PMPParser +from pcmdi_metrics.utils import StringConstructor + + +def create_monsoon_wang_parser(): + P = PMPParser() + + P.use("--modnames") + P.use("--results_dir") + P.use("--reference_data_path") + P.use("--test_data_path") + P.use("--obs_mask") + + P.add_argument( + "--outnj", + "--outnamejson", + type=str, + dest="outnamejson", + default="monsoon_wang.json", + help="Output path for jsons", + ) + P.add_argument( + "-e", + "--experiment", + type=str, + dest="experiment", + default="historical", + help="AMIP, historical or picontrol", + ) + P.add_argument( + "-c", "--MIP", type=str, dest="mip", default="CMIP5", help="put options here" + ) + P.add_argument( + "--ovar", dest="obsvar", default="pr", help="Name of variable in obs file" + ) + P.add_argument( + "-v", + "--var", + dest="modvar", + default="pr", + help="Name of variable in model files", + ) + P.add_argument( + "-t", + "--threshold", + default=2.5 / 86400.0, + type=float, + help="Threshold for a hit when computing skill score", + ) + P.add_argument( + "--cmec", + dest="cmec", + default=False, + action="store_true", + help="Use to save CMEC format metrics JSON", + ) + P.add_argument( + "--no_cmec", + dest="cmec", + default=False, + action="store_false", + help="Do not save CMEC format metrics JSON", + ) + P.set_defaults(cmec=False) + return P + + +def monsoon_wang_runner(args): + # args = P.parse_args(sys.argv[1:]) + modpath = StringConstructor(args.test_data_path) + modpath.variable = args.modvar + outpathdata = args.results_dir + if isinstance(args.modnames, str): + mods = eval(args.modnames) + else: + mods = args.modnames + + json_filename = args.outnamejson + + if json_filename == "CMIP_MME": + json_filename = "/MPI_" + args.mip + "_" + args.experiment + + # VAR IS FIXED TO BE PRECIP FOR CALCULATING MONSOON PRECIPITATION INDICES + var = args.modvar + thr = args.threshold + sig_digits = ".3f" + + # Get flag for CMEC output + cmec = args.cmec + + ######################################### + # PMP monthly default PR obs + cdms2.axis.longitude_aliases.append("longitude_prclim_mpd") + cdms2.axis.latitude_aliases.append("latitude_prclim_mpd") + fobs = xr.open_dataset(args.reference_data_path, decode_times=False) + dobs_orig = fobs[args.obsvar] + fobs.close() + + ######################################## + + # FCN TO COMPUTE GLOBAL ANNUAL RANGE AND MONSOON PRECIP INDEX + + annrange_obs, mpi_obs = mpd(dobs_orig) + + # create monsoon domain mask based on observations: annual range > 2.5 mm/day + if args.obs_mask: + domain_mask_obs = xr.where(annrange_obs > thr, 1, 0) + domain_mask_obs.name = "mask" + mpi_obs = mpi_obs.where(domain_mask_obs) + nout_mpi_obs = os.path.join(outpathdata, "mpi_obs_masked.nc") + + ######################################### + # SETUP WHERE TO OUTPUT RESULTING DATA (netcdf) + nout = os.path.join( + outpathdata, "_".join([args.experiment, args.mip, "wang-monsoon"]) + ) + try: + os.makedirs(nout) + except BaseException: + pass + + # SETUP WHERE TO OUTPUT RESULTS (json) + jout = outpathdata + try: + os.makedirs(nout) + except BaseException: + pass + + gmods = [] # "Got" these MODS + for i, mod in enumerate(mods): + modpath.model = mod + for k in modpath.keys(): + try: + val = getattr(args, k) + except Exception: + continue + if not isinstance(val, (list, tuple)): + setattr(modpath, k, val) + else: + setattr(modpath, k, val[i]) + l1 = modpath() + if os.path.isfile(l1) is True: + gmods.append(mod) + + if len(gmods) == 0: + raise RuntimeError("No model file found!") + + ######################################### + + egg_pth = resources.resource_path() + globals = {} + locals = {} + exec( + compile( + open(os.path.join(egg_pth, "default_regions.py")).read(), + os.path.join(egg_pth, "default_regions.py"), + "exec", + ), + globals, + locals, + ) + + regions_specs = locals["regions_specs"] + doms = ["AllMW", "AllM", "NAMM", "SAMM", "NAFM", "SAFM", "ASM", "AUSM"] + + mpi_stats_dic = {} + for i, mod in enumerate(gmods): + modpath.model = mod + for k in modpath.keys(): + try: + val = getattr(args, k) + except Exception: + continue + if not isinstance(val, (list, tuple)): + setattr(modpath, k, val) + else: + setattr(modpath, k, val[i]) + modelFile = modpath() + + mpi_stats_dic[mod] = {} + + print("modelFile = ", modelFile) + f = xr.open_dataset(modelFile) + d_orig = f[var] + + annrange_mod, mpi_mod = mpd(d_orig) + + domain_mask_mod = xr.where(annrange_mod > thr, 1, 0) + mpi_mod = mpi_mod.where(domain_mask_mod) + + lats = annrange_obs.lat[0] + latn = annrange_obs.lat[-1] + lone = annrange_obs.lon[-1] + lonw = annrange_obs.lon[0] + + annrange_obs = regrid(annrange_obs, annrange_mod) + + mpi_obs = regrid(mpi_obs, mpi_mod) + + regions_specs = load_regions_specs() + + for dom in doms: + mpi_stats_dic[mod][dom] = {} + + print("dom = ", dom) + + mpi_obs_reg = region_subset(mpi_obs, dom) + mpi_obs_reg_sd = mpi_obs_reg.std(dim=["lat", "lon"]) + mpi_mod_reg = region_subset(mpi_mod, dom) + + da1_flat = mpi_mod_reg.values.ravel() + da2_flat = mpi_obs_reg.values.ravel() + + cor = np.ma.corrcoef( + np.ma.masked_invalid(da1_flat), np.ma.masked_invalid(da2_flat) + )[0, 1] + + squared_diff = (mpi_mod_reg - mpi_obs_reg) ** 2 + mean_squared_error = squared_diff.mean(skipna=True) + rms = np.sqrt(mean_squared_error) + + rmsn = rms / mpi_obs_reg_sd + + # DOMAIN SELECTED FROM GLOBAL ANNUAL RANGE FOR MODS AND OBS + + annrange_mod_dom = region_subset(annrange_mod, dom) + annrange_obs_dom = region_subset(annrange_obs, dom) + + # SKILL SCORES + # HIT/(HIT + MISSED + FALSE ALARMS) + hit, missed, falarm, score, hitmap, missmap, falarmmap = mpi_skill_scores( + annrange_mod_dom, annrange_obs_dom, thr + ) + + # POPULATE DICTIONARY FOR JSON FILES + mpi_stats_dic[mod][dom] = {} + mpi_stats_dic[mod][dom]["cor"] = format(cor, sig_digits) + mpi_stats_dic[mod][dom]["rmsn"] = format(rmsn, sig_digits) + mpi_stats_dic[mod][dom]["threat_score"] = format(score, sig_digits) + + # SAVE ANNRANGE AND HIT MISS AND FALSE ALARM FOR EACH MOD DOM + fm = os.path.join(nout, "_".join([mod, dom, "wang-monsoon_xcdat.nc"])) + ds_out = xr.Dataset( + { + "obsmap": annrange_obs_dom, + "modmap": annrange_mod_dom, + "hitmap": hitmap, + "missmap": missmap, + "falarmmap": falarmmap, + } + ) + ds_out.to_netcdf(fm) + f.close() + + if np.isnan(cor): + print("invalid correlation values") + sys.exit() + + # OUTPUT METRICS TO JSON FILE + OUT = pcmdi_metrics.io.base.Base(os.path.abspath(jout), json_filename) + + disclaimer = open(os.path.join(egg_pth, "disclaimer.txt")).read() + + metrics_dictionary = collections.OrderedDict() + metrics_dictionary["DISCLAIMER"] = disclaimer + metrics_dictionary["REFERENCE"] = ( + "The statistics in this file are based on" + + " Wang, B., Kim, HJ., Kikuchi, K. et al. " + + "Clim Dyn (2011) 37: 941. doi:10.1007/s00382-010-0877-0" + ) + metrics_dictionary["RESULTS"] = mpi_stats_dic # collections.OrderedDict() + + OUT.var = var + OUT.write( + metrics_dictionary, + json_structure=["model", "domain", "statistic"], + indent=4, + separators=(",", ": "), + ) + if cmec: + print("Writing cmec file") + OUT.write_cmec(indent=4, separators=(",", ": ")) diff --git a/pcmdi_metrics/monsoon_wang/run.py b/pcmdi_metrics/monsoon_wang/run.py new file mode 100644 index 000000000..004963366 --- /dev/null +++ b/pcmdi_metrics/monsoon_wang/run.py @@ -0,0 +1,5 @@ +from monsoon_wang_driver import create_monsoon_wang_parser, monsoon_wang_runner + +P = create_monsoon_wang_parser() +args = P.get_parameter(argparse_vals_only=False) +monsoon_wang_runner(args)