from rook.utils.fixes_utils import convert_calendar_to_gregorian
from rook.utils.data_utils.attr_utils import (
edit_global_attrs,
edit_var_attrs,
)
from rook.utils.data_utils.coord_utils import add_coord, add_scalar_coord
from rook.utils.data_utils.var_utils import add_data_var
model_specific_global_attrs = {
"CMCC-CM2-SR5": {
"forcing_description": "f1, CMIP6 historical forcings",
"physics_description": "physics from the standard model configuration, with no additional tuning or different parametrization", # noqa
"initialization_description": "hindcast initialized based on observations and using historical forcing", # noqa
},
"EC-Earth3": {
"forcing_description": "f1, CMIP6 historical forcings",
"physics_description": "physics from the standard model configuration, with no additional tuning or different parametrization", # noqa
"initialization_description": "Atmosphere initialization based on full-fields from ERA-Interim (s1979-s2018) or ERA-40 (s1960-s1978); ocean/sea-ice initialization based on full-fields from NEMO/LIM assimilation run nudged towards ORA-S4 (s1960-s2018)", # noqa
},
"HadGEM3-GC31-MM": {
"forcing_description": "f2, CMIP6 v6.2.0 forcings; no ozone remapping",
"physics_description": "physics from the standard model configuration, with no additional tuning or different parametrization", # noqa
"initialization_description": "hindcast initialized based on observations and using historical forcing", # noqa
},
"MPI-ESM1-2-HR": {
"forcing_description": "f1, CMIP6 historical forcings",
"physics_description": "physics from the standard model configuration, with no additional tuning or different parametrization", # noqa
"initialization_description": "hindcast initialized based on observations and using historical forcing", # noqa
},
"MPI-ESM1-2-LR": {
"forcing_description": "f1, CMIP6 historical forcings",
"physics_description": "physics from the standard model configuration, with no additional tuning or different parametrization", # noqa
"initialization_description": "hindcast initialized based on observations and using historical forcing", # noqa
},
}
[docs]
def get_decadal_model_attr_from_dict(ds_id, ds, attr):
# TODO: method originally ported from legacy decadal utility code.
# Add the model-specific global attr
model = ds_id.split(".")[3]
value = model_specific_global_attrs[model][attr]
return value
[docs]
def apply_decadal_fixes(ds_id, ds, output_dir=None):
ds_mod = decadal_fix_calendar(ds_id, ds, output_dir=output_dir)
ds_mod = decadal_fix_1(ds_id, ds_mod)
ds_mod = decadal_fix_2(ds_id, ds_mod)
ds_mod = decadal_fix_3(ds_id, ds_mod)
ds_mod = decadal_fix_4(ds_id, ds_mod)
ds_mod = decadal_fix_5(ds_id, ds_mod)
return ds_mod
[docs]
def decadal_fix_1(ds_id, ds):
operands = {"var_id": "time", "attrs": {"long_name": "valid_time"}}
ds_mod = edit_var_attrs(ds_id, ds, **operands)
return ds_mod
[docs]
def decadal_fix_2(ds_id, ds):
operands = {
"attrs": {
"forcing_description": get_decadal_model_attr_from_dict(
ds_id, ds, "forcing_description"
), # noqa
"physics_description": get_decadal_model_attr_from_dict(
ds_id, ds, "physics_description"
), # noqa
"initialization_description": get_decadal_model_attr_from_dict(
ds_id, ds, "initialization_description"
), # noqa
"startdate": "derive: rook.utils.decadal_utils.get_sub_experiment_id",
"sub_experiment_id": "derive: rook.utils.decadal_utils.get_sub_experiment_id",
}
}
ds_mod = edit_global_attrs(ds_id, ds, **operands)
return ds_mod
[docs]
def decadal_fix_3(ds_id, ds):
operands = {
"var_id": "reftime",
"value": "derive: rook.utils.decadal_utils.get_reftime",
"dtype": "datetime64[ns]",
"attrs": {
"long_name": "Start date of the forecast",
"standard_name": "forecast_reference_time",
},
"encoding": {
"dtype": "int32",
"units": "days since 1850-01-01",
"calendar": "derive: rook.utils.decadal_utils.get_time_calendar",
},
}
ds_mod = add_scalar_coord(ds_id, ds, **operands)
return ds_mod
[docs]
def decadal_fix_4(ds_id, ds):
operands = {
"var_id": "leadtime",
"value": "derive: rook.utils.decadal_utils.get_lead_times",
"dim": ["time"],
"dtype": "float64",
"attrs": {
"long_name": "Time elapsed since the start of the forecast",
"standard_name": "forecast_period",
"units": "days",
},
"encoding": {"dtype": "double"},
}
ds_mod = add_coord(ds_id, ds, **operands)
return ds_mod
[docs]
def decadal_fix_5(ds_id, ds):
operands = {
"var_id": "realization",
"value": ds.attrs.get("realization_index"),
"dtype": "int32",
"attrs": {
"long_name": "realization",
"comment": "For more information on the ripf, refer to the variant_label, initialization_description, physics_description and forcing_description global attributes", # noqa
},
}
ds_mod = add_data_var(ds_id, ds, **operands)
return ds_mod
[docs]
def decadal_fix_calendar(ds_id, ds, output_dir=None):
# set proleptic_gregorian calendar to gregorian (standard).
# the proleptic gregorian calendar extends the gregorin backward in time before 1582.
calendar = ds.time.encoding.get("calendar", "standard")
if calendar == "proleptic_gregorian":
# ds.time.encoding["calendar"] = "standard"
ds = convert_calendar_to_gregorian(ds)
return ds
