Downscale CMIP6 (dc6) sandbox scratch script to get tasmin, tasmax parameter files for each GCM and combine to create DTR parameter file.

make_dtr.py

# Get tasmin, tasmax parameter files for each GCM and combine to create DTR parameter file.
# This is all rough and likely to break. Sorry.
import logging
import pathlib
import dearprudence
import dearprudence.validation


PARAMS_DIR = pathlib.Path("/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters")


def validate_params_file(urlpath, check_inputs_exist=False):
    runs = dearprudence.read_params(urlpath)
    return dearprudence.validation.validate_params(runs, check_inputs_exist=check_inputs_exist)


def _make_target_experiment_id_dict(runs):
    """Dict of runs keyed on experiment_id"""
    d = {getattr(r, r.target).experiment_id: r for r in runs}
    if len(d.keys()) != len(runs):
        raise ValueError(f"Duplicate target 'experiment_id's found in {runs=}")
    return d


def make_dtrruns(tasmin_runs, tasmax_runs):
    """
    Make DtrRuns from a list of tasmin SimpleRuns and a list of tasmax SimpleRuns
    """
    # Key up runs on their target's experiment_id field so we can find common experiment_ids.
    tmins = _make_target_experiment_id_dict(tasmin_runs)
    tmaxs = _make_target_experiment_id_dict(tasmax_runs)

    # Get common target experiments for input variables. Explicitly start with tasmax, as we've
    # run these and the params are more likely to be correct.
    common_experiment_ids = set(tmaxs.keys()).intersection(tmins.keys())

    dtrs = []
    for experiment in common_experiment_ids:

        target=None
        if experiment == "historical":
            target = "historical"
        elif experiment[:3] == "ssp":
            target = "ssp"
        else:
            raise ValueError(f"{experiment=}, must be 'historical' or 'ssp*'")

        dtrs.append(
            dearprudence.DtrRun(
                target=target,
                variable_id="dtr",
                tasmin=tmins[experiment],
                tasmax=tmaxs[experiment]
            )
        )
    return dtrs


def make_dtrruns_file(out_pathurl, source_id, params_dir, flname_template="{source_id}-{variable_id}.yaml"):
    if isinstance(params_dir, str):
        params_dir = pathlib.Path(params_dir)

    tasmin_path = params_dir.joinpath(flname_template.format(source_id=source_id, variable_id="tasmin"))
    tasmax_path = params_dir.joinpath(flname_template.format(source_id=source_id, variable_id="tasmax"))

    dtrs = make_dtrruns(
        tasmin_runs=dearprudence.read_params(tasmin_path),
        tasmax_runs=dearprudence.read_params(tasmax_path)
    )
    dearprudence.write_params(out_pathurl, dtrs)


## This is a good test of the code:
# outpath="/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/ACCESS-ESM1-5-dtr-TEST.yaml"
# make_dtrruns_file(
#     out_pathurl=outpath, 
#     source_id="ACCESS-ESM1-5", 
#     params_dir="/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/"
# )
# validate_params_file(outpath, check_inputs_exist=True)  # This raises an exception because missing tasmin inputs.

## Checking tasmax:
# skip_files = [
#     "/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/GFDL-CM4-tasmax.yaml",
#     "/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/UKESM1-0-LL-tasmax.yaml"
# ]
# for fl in params_dir.glob("*tasmax.yaml"):
#     print(f"Working on {fl}")

#     if str(fl) in skip_files:
#         print("Skipping")
#         continue

#     validate_params_file(fl, check_inputs_exist=True)
#     print("Done")

# outpath="/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/ACCESS-ESM1-5-dtr-TEST.yaml"
# make_dtrruns_file(
#     out_pathurl=outpath, 
#     source_id="ACCESS-ESM1-5", 
#     params_dir="/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/"
# )
# validate_params_file(outpath, check_inputs_exist=True)  # This raises an exception because missing tasmin inputs.

## Another good test of the code:
# outpath="/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/EC-Earth3-dtr.yaml"
# make_dtrruns_file(
#     out_pathurl=outpath, 
#     source_id="EC-Earth3", 
#     params_dir="/Users/bmalevich/Projects/downscaleCMIP6/workflows/parameters/"
# )
# validate_params_file(outpath, check_inputs_exist=True)  # This raises an exception because missing tasmin inputs.


def make_dtr_for_all():
    """Make DTR from all parameter files"""
    logging.info("Starting")

    all_source_ids = [
        "ACCESS-CM2",
        "ACCESS-ESM1-5",
        "AWI-CM-1-1-MR",
        "BCC-CSM2-MR",
        "CAMS-CSM1-0",
        "CMCC-CM2-SR5",
        "CMCC-ESM2",
        "CanESM5",
        "EC-Earth3",
        "EC-Earth3-Veg",
        "EC-Earth3-Veg-LR",
        "FGOALS-g3",
        "GFDL-CM4",
        "GFDL-ESM4",
        "INM-CM4-8",
        "INM-CM5-0",
        "KACE-1-0-G",
        "KIOST-ESM",
        "MIROC-ES2L",
        "MIROC6",
        "MPI-ESM-1-2-HAM",
        "MPI-ESM1-2-HR",
        "MPI-ESM1-2-LR",
        "MRI-ESM2-0",
        "NESM3",
        "NorESM2-LM",
        "NorESM2-MM",
        # "UKESM1-0-LL"
    ]

    for s in all_source_ids:
        logging.info(f"Processing {s}")

        outpath=PARAMS_DIR.joinpath(f"{s}-dtr.yaml")

        logging.debug(f"Writing output to {outpath}")

        try:
            make_dtrruns_file(
                out_pathurl=outpath, 
                source_id=s, 
                params_dir=PARAMS_DIR
            )
        except (FileNotFoundError) as e:
            logging.error(e)
            logging.info("Error reading file - skipping")
            continue


        logging.info(f"Output written to {outpath}")

        try:
            validate_params_file(outpath, check_inputs_exist=True)
        except (dearprudence.ParameterFileValidationError, dearprudence.Cmip6CatalogError) as e:
            logging.error(e)
            continue

        logging.info(f"Processed {s}")

    logging.info("Done")



def validate_all_existing_params():
    """Validate all parameter files"""
    logging.info("Starting")

    skip_files = []
    for fl in PARAMS_DIR.glob("*.yaml"):
        logging.info(f"Working on {fl}")

        if str(fl) in skip_files:
            logging.info("Skipping")
            continue

        try:
            validate_params_file(fl, check_inputs_exist=True)
        except (dearprudence.ParameterFileValidationError, dearprudence.Cmip6CatalogError) as e:
            logging.error(e)
            continue

        logging.info(f"Validated {fl}")
    
    logging.info("Done")


if __name__ == "__main__":
    logging.basicConfig(level="INFO")
    # validate_all_existing_params()
    make_dtr_for_all()