diff --git a/helper_scripts/generate_cupid_config_for_cesm_case.py b/helper_scripts/generate_cupid_config_for_cesm_case.py new file mode 100755 index 0000000..ce2c87d --- /dev/null +++ b/helper_scripts/generate_cupid_config_for_cesm_case.py @@ -0,0 +1,113 @@ +#!/usr/bin/env python3 +from __future__ import annotations + +import argparse +import os +import sys + +import yaml + + +def _parse_args(): + """Parse command line arguments""" + + parser = argparse.ArgumentParser( + description="Generate config.yml based on an existing CUPID example", + formatter_class=argparse.ArgumentDefaultsHelpFormatter, + ) + + # Command line argument for location of CESM source code (required) + parser.add_argument( + "--cesm-root", + action="store", + dest="cesm_root", + required=True, + help="Location of CESM source code", + ) + + # Command line argument for CUPiD example from which to get config.yml + parser.add_argument( + "--cupid-example", + action="store", + dest="cupid_example", + default="key_metrics", + help="CUPiD example to use as template for config.yml", + ) + + # Command line argument location of CESM case directory + parser.add_argument( + "--case-root", + action="store", + dest="case_root", + default=os.getcwd(), + help="CESM case directory", + ) + + return parser.parse_args() + + +def generate_cupid_config(case_root, cesm_root, cupid_example): + sys.path.append(os.path.join(cesm_root, "cime")) + from CIME.case import Case + + # Is cupid_example a valid value? + cupid_root = os.path.join(cesm_root, "tools", "CUPiD") + cupid_examples = os.path.join(cupid_root, "examples") + valid_examples = [ + example + for example in next(os.walk(cupid_examples))[1] + if example not in ["ilamb", "nblibrary"] + ] + if cupid_example not in valid_examples: + error_msg = f"argument --cupid-example: invalid choice '{cupid_example}'" + raise KeyError( + f"{error_msg} (choose from subdirectories of {cupid_examples}: {valid_examples})", + ) + + with Case(case_root, read_only=False, record=True) as cesm_case: + case = cesm_case.get_value("CASE") + dout_s_root = cesm_case.get_value("DOUT_S_ROOT") + + # Additional options we need to get from env_cupid.xml + nyears = 1 + start_date = "0001-01-01" + end_date = f"{nyears+1:04d}-01-01" + climo_nyears = 1 + base_case_output_dir = "/glade/campaign/cesm/development/cross-wg/diagnostic_framework/CESM_output_for_testing" + base_nyears = 100 + base_end_date = f"{base_nyears+1:04d}-01-01" + base_climo_nyears = 40 + + with open(os.path.join(cupid_root, "examples", cupid_example, "config.yml")) as f: + my_dict = yaml.safe_load(f) + + my_dict["global_params"]["case_name"] = case + my_dict["global_params"]["start_date"] = start_date + my_dict["global_params"]["end_date"] = end_date + my_dict["global_params"]["climo_nyears"] = climo_nyears + my_dict["global_params"]["base_case_output_dir"] = base_case_output_dir + my_dict["global_params"]["base_end_date"] = base_end_date + my_dict["global_params"]["base_climo_nyears"] = base_climo_nyears + my_dict["timeseries"]["case_name"] = case + my_dict["timeseries"]["atm"]["end_years"] = [nyears, base_nyears] + + # replace with environment variable + my_dict["global_params"]["CESM_output_dir"] = dout_s_root + + # create new file, make it writeable + with open("config.yml", "w") as f: + # Header of file is a comment logging provenance + f.write(f"# This file has been auto-generated for use with {case}\n") + f.write(f"# It is based off of examples/{cupid_example}/config.yml\n") + f.write("# Arguments used:\n") + f.write(f"# cesm_root = {cesm_root}\n") + f.write(f"# case_root = {case_root}\n") + f.write(f"# cupid_example= {cupid_example}\n") + + # enter in each element of the dictionary into the new file + yaml.dump(my_dict, f, sort_keys=False) + + +if __name__ == "__main__": + args = vars(_parse_args()) + generate_cupid_config(**args)