diff --git a/expts_manager/Expts_manager.py b/expts_manager/Expts_manager.py new file mode 100755 index 0000000..bac4393 --- /dev/null +++ b/expts_manager/Expts_manager.py @@ -0,0 +1,1500 @@ +#!/usr/bin/env python3 +""" +ACCESS-OM Experiment Management Tool +This python script manages experiment runs for both ACCESS-OM2 and ACCESS-OM3, +providing functionalities to set up control and perturbation experiments, +modify configuration files, and manage related utilities. + +Latest version: https://github.com/COSIMA/om3-scripts/pull/34 +Author: Minghang Li +Email: minghang.li1@anu.edu.au +License: Apache 2.0 License http://www.apache.org/licenses/LICENSE-2.0.txt +""" + + +# =========================================================================== +import os +import sys +import re +import subprocess +import shutil +import glob +import argparse +import warnings + +try: + import numpy as np + import git + import f90nml + from ruamel.yaml import YAML + + ryaml = YAML() + ryaml.preserve_quotes = True +except ImportError: + print("\nFatal error: modules not available.") + print("On NCI, do the following and try again:") + print(" module use /g/data/vk83/modules && module load payu/1.1.5\n") + raise +# =========================================================================== + + +class LiteralString(str): + pass + + +def represent_literal_str(dumper, data): + return dumper.represent_scalar("tag:yaml.org,2002:str", data, style="|") + + +ryaml.representer.add_representer(LiteralString, represent_literal_str) + + +def update_MOM6_params_override(param_dict_change, commt_dict_change): + """ + Prepends `#override` to parameters for MOM6. + Args: + param_dict_change (dict): dictionary of parameters to override + commt_dict_change (dict): dictionary of comments for parameters + Returns: + tuple: Two dictionaries with `#override` prepended to each key. + """ + override_param_dict_change = { + f"#override {k}": v for k, v in param_dict_change.items() + } + override_commt_dict_change = { + f"#override {k}": v for k, v in commt_dict_change.items() + } + return override_param_dict_change, override_commt_dict_change + + +class Expts_manager(object): + """ + A class to manage ACCESS-OM3 experiment runs, including control and perturbation experiments. + Attributes: + MOM_prefix (str): Prefix for MOM6 parameters. + nml_suffix (str): Suffix for namelist parameters. + runseq_prefix (str): Prefix for the coupling timestep in `nuopc.runseq`. + combo_suffix (str): Suffix for combo perturbation experiments, i.e., multiple-parameter tests. + branch_perturb (str): branch name for the perturbation. + """ + + DIR_MANAGER = os.getcwd() + + def __init__( + self, + force_overwrite_tools: bool = False, + MOM_prefix: str = "MOM_list", + CONFIG_prefix: str = "config_list", + runconfig_suffix1: str = "_attributes", + runconfig_suffix2: str = "_modelio", + nml_suffix: str = "_nml", + runseq_prefix: str = "runseq_list", + combo_suffix: str = "_combo", + branch_perturb: str = "perturb", + ): + + self.dir_manager = self.DIR_MANAGER + self.force_overwrite_tools = force_overwrite_tools + self.MOM_prefix = MOM_prefix + self.CONFIG_prefix = CONFIG_prefix + self.runconfig_suffix1 = runconfig_suffix1 + self.runconfig_suffix2 = runconfig_suffix2 + self.nml_suffix = nml_suffix + self.runseq_prefix = runseq_prefix + self.branch_perturb = branch_perturb + self.combo_suffix = combo_suffix + + def load_variables(self, yamlfile): + """ + Loads variables from the input yaml file + Args: + yamlfile (str): Path to the YAML configuration file, i.e., Expts_manager.yaml. + Attributes: + yamlfile (str): Path to the YAML configuration file. + indata (dict): Data loaded from the YAML file. + utils_url (str): Git url for the om3-utils tool. + utils_branch_name (str): Branch name for the om3-utils tool. + utils_dir_name (str): User-defined directory for the om3-utils tool. + base_url (dict): Git url for the ACCESS-OM3 configuration. + base_commit (str): Specific git commit for the ACCESS-OM3 configuration. + base_dir_name (str): User-defined directory name for the baseline control experiment. + base_branch_name (str): User-defined branch name for the control experiment. + test_path (str): User-defined path for test runs, including control and perturbation experiments. + startfrom (int/str): Restart number of the control experiment used as an initial condition for perturbation tests; use 'rest' to start from the initial state. + startfrom_str (str): String representation of `startfrom`, padded to three digits. + ctrl_nruns (int): Number of control runs. It is associated with total number of output directories that have been generated. + pert_nruns (int): Number of perturbation experiment runs; associated with total number of output directories that have been generated. + """ + self.yamlfile = yamlfile + self.indata = self._read_ryaml(yamlfile) + + self.model = self.indata["model"] + self.force_overwrite_tools = self.indata.get("force_overwrite_tools", False) + self.utils_url = self.indata.get("utils_url", None) + self.utils_dir_name = self.indata.get("utils_dir_name", None) + self.utils_branch_name = self.indata.get("utils_branch_name", None) + + self.base_url = self.indata["base_url"] + self.base_commit = self.indata["base_commit"] + self.base_dir_name = self.indata["base_dir_name"] + self.base_branch_name = self.indata["base_branch_name"] + + self.test_path = self.indata["test_path"] + + self.diag_url = self.indata.get("diag_url", None) + self.diag_branch_name = self.indata.get("diag_branch_name", None) + self.diag_dir_name = self.indata.get("diag_dir_name", None) + self.diag_ctrl = self.indata.get("diag_ctrl", False) + self.diag_pert = self.indata.get("diag_pert", False) + + self.ctrl_nruns = self.indata.get("ctrl_nruns", 0) + self.run_namelists = self.indata.get("run_namelists", False) + self.check_duplicate_jobs = self.indata.get("check_duplicate_jobs", True) + self.check_skipping = self.indata.get("check_skipping", False) + self.force_restart = self.indata.get("force_restart", False) + self.startfrom = self.indata["startfrom"] + self.startfrom_str = str(self.startfrom).strip().lower().zfill(3) + self.nruns = self.indata.get("nruns", 0) + + self._initialise_variables() + + def _initialise_variables(self): + """ + Initialises variables from experiment setups + nml_ctrl (f90nml): f90 namlist for the interested parameters. It is used as a base to modify for perturbation experiments. + tag_model (str): Switch for tuning parameters between f90 namelist and MOM_input. + param_dict_change_list list[dict]: Specific for MOM_input, the list containing tunning parameter dictionaries. + commt_dict_change (dict): Specific for MOM_input, dictionary of comments for parameters. + append_group_list (list): Specific for f90nml, the list containing tunning parameters. + expt_names list(str): Optional user-defined directory names for perturbation experiments. + tmp_count (int): count the number of parameter groups in a single parameter block in process. + group_count (int): total number of parameter groups in a single parameter block. + """ + self.nml_ctrl = None + self.tag_model = None + self.param_dict_change_list = [] + self.commt_dict_change = {} + self.append_group_list = [] + self.previous_key = None + self.expt_names = None + self.diag_path = None + + self.tmp_count = 0 + self.group_count = None + + def load_tools(self): + """ + Loads external tools required for the experiments. + """ + + # currently import from a fork: https://github.com/minghangli-uni/om3-utils + # will update the tool when it is merged to COSIMA/om3-utils + def _clone_repo(branch_name, url, path, tool_name, force_overwrite_tools): + if os.path.exists(path) and os.path.isdir(path): + if force_overwrite_tools: + print( + f"-- Force_overwrite_tools is activated, hence removing existing {tool_name}: {path}" + ) + shutil.rmtree(path) + print(f"Cloning {tool_name} for use!") + command = ( + f"git clone --branch {branch_name} {url} {path} --single-branch" + ) + subprocess.run(command, shell=True, check=True) + else: + print(f"{tool_name} already exists, hence skips cloning!") + else: + print(f"Cloning {tool_name} for use!") + command = ( + f"git clone --branch {branch_name} {url} {path} --single-branch" + ) + subprocess.run(command, shell=True, check=True) + print(f"Finished cloning {tool_name}!") + + # om3-utils is a must for om3 but not required for access-om2. + utils_path = ( + os.path.join(self.dir_manager, self.utils_dir_name) + if self.utils_dir_name + else None + ) + + if self.model == "access-om3" and utils_path is None: + raise ValueError(f"om3-utils tool must be loaded for {self.model}!") + elif self.model == "access-om2" and utils_path: + warnings.warn( + f"om3-utils tool is not used for {self.model}, " + f"hence, is not cloned!", + UserWarning, + ) + + if self.model == "access-om3": + _clone_repo( + self.utils_branch_name, + self.utils_url, + utils_path, + self.utils_dir_name, + self.force_overwrite_tools, + ) + + # make_diag_table is [optional] + self.diag_path = ( + os.path.join(self.dir_manager, self.test_path, self.diag_dir_name) + if self.diag_dir_name + else None + ) + + if self.diag_path is not None: + _clone_repo( + self.diag_branch_name, + self.diag_url, + self.diag_path, + self.diag_dir_name, + self.force_overwrite_tools, + ) + sys.path.extend([utils_path, self.diag_path]) + else: + sys.path.extend([utils_path]) + + if utils_path is not None: + # load modules from om3-utils + from om3utils import MOM6InputParser + from om3utils.nuopc_config import read_nuopc_config, write_nuopc_config + + self.MOM6InputParser = MOM6InputParser + self.read_nuopc_config = read_nuopc_config + self.write_nuopc_config = write_nuopc_config + + def create_test_path(self): + """ + Creates the local test directory for blocks of parameter testing. + """ + if os.path.exists(self.test_path): + print(f"test directory {self.test_path} already exists!") + else: + os.makedirs(self.test_path) + print(f"test directory {self.test_path} is created!") + + def model_selection(self): + """ + Ensures the model to be either "access-om2" or "access-om3" + """ + if self.model not in (("access-om2", "access-om3")): + raise ValueError( + f"{self.model} requires to be either " f"access-om2 or access-om3!" + ) + + def manage_ctrl_expt(self): + """ + Setup and run the control experiment + """ + self.base_path = os.path.join( + self.dir_manager, self.test_path, self.base_dir_name + ) + base_path = self.base_path + ctrl_nruns = self.ctrl_nruns + + # access-om2 specific + ocn_path = os.path.join(base_path, "ocean") + + if os.path.exists(base_path): + print(f"Base path is already created and located at {base_path}") + if not os.path.isfile(os.path.join(base_path, "config.yaml")): + print( + "previous commit fails, please try with an updated commit hash for the control experiment!" + ) + # extract specific configuration via commit hash + self._extract_config_via_commit() + else: + # clone the template repo and setup the control branch + self._clone_template_repo() + + # extract specific configuration via commit hash + self._extract_config_via_commit() + + # [optional] modify diag_table + if self.diag_ctrl and self.diag_path: + if self.model == "access-om3": + self._copy_diag_table(base_path) + elif self.model == "access-om2": + self._copy_diag_table(ocn_path) + + # setup the control experiments + self._setup_ctrl_expt() + + # check existing pbs jobs + pbs_jobs = self._output_existing_pbs_jobs() + + # check duplicated running jobs + if self.check_duplicate_jobs: + duplicated_bool = self._check_duplicated_jobs(pbs_jobs, base_path) + else: + duplicated_bool = False + + if not duplicated_bool: + # Checks the current state of the repo, commits relevant changes. + self._check_and_commit_changes() + + # start control runs, count existing runs and do additional runs if needed + self._start_experiment_runs( + base_path, self.base_dir_name, duplicated_bool, ctrl_nruns + ) + + def _clone_template_repo(self): + """ + Clones the template repo. + """ + print(f"Cloning template from {self.base_url} to {self.base_path}") + command = f"payu clone {self.base_url} {self.base_path}" + subprocess.run(command, shell=True, check=False) + + def _extract_config_via_commit(self): + """ + Extract specific configuration via commit hash. + """ + templaterepo = git.Repo(self.base_path) + print( + f"Check out commit {self.base_commit} and creat new branch {self.base_branch_name}!" + ) + # checkout the new branch from the specific template commit + templaterepo.git.checkout( + "-b", self.base_branch_name, str(self.base_commit) + "^0" + ) + + def _copy_diag_table(self, path): + """ + Copies the diagnostic table (`diag_table`) to the specified path if a path is defined. + """ + if self.diag_path: + command = f"scp {os.path.join(self.diag_path,'diag_table')} {path}" + subprocess.run(command, shell=True, check=False) + print(f"Copy diag_table to {path}") + else: + print( + f"{self.diag_path} is not defined, hence skip copy diag_table to the control experiment" + ) + + def _count_file_nums(self): + """ + Counts the number of file numbers. + """ + return len(os.listdir(self.base_path)) + + def _setup_ctrl_expt(self): + """ + Modifies parameters based on the input YAML configuration for the ctrl experiment. + + Updates configuration files (config.yaml, nuopc.runconfig etc), + namelist and MOM_input for the control experiment if needed. + """ + # for file_name in os.listdir(self.base_path): + for root, dirs, files in os.walk(self.base_path): + dirs[:] = [tmp_d for tmp_d in dirs if ".git" not in tmp_d] + for f in files: + if ".git" in f: + continue + file_name = os.path.relpath(os.path.join(root, f), self.base_path) + yaml_data = self.indata.get(file_name, None) + + if yaml_data: + # Update parameters from namelists + if file_name.endswith("_in") or file_name.endswith(".nml"): + self._update_nml_params(self.base_path, yaml_data, file_name) + + # Update config entries from `nuopc.runconfig` + if file_name == "nuopc.runconfig": + self._update_runconfig_params( + self.base_path, yaml_data, file_name + ) + + # Update config entries from `config_yaml` + if file_name == "config.yaml": + self._update_config_params(self.base_path, yaml_data, file_name) + + # Update and overwrite parameters from and into `MOM_input` + if file_name == "MOM_input": + # parse existing MOM_input + MOM_inputParser = self._parser_mom6_input( + os.path.join(self.base_path, file_name) + ) + param_dict = ( + MOM_inputParser.param_dict + ) # read parameter dictionary + commt_dict = ( + MOM_inputParser.commt_dict + ) # read comment dictionary + param_dict.update(yaml_data) + # overwrite to the same `MOM_input` + MOM_inputParser.writefile_MOM_input( + os.path.join(self.base_path, file_name) + ) + + # Update only coupling timestep from `nuopc.runseq` + if file_name == "nuopc.runseq": + nuopc_runseq_file = os.path.join(self.base_path, file_name) + self._update_cpl_dt_nuopc_seq(nuopc_runseq_file, yaml_data) + + def _check_and_commit_changes(self): + """ + Checks the current state of the repo, stages relevant changes, and commits them. + If no changes are detected, it provides a message indicating that no commit was made. + """ + repo = git.Repo(self.base_path) + print(f"Current base branch is: {repo.active_branch.name}") + deleted_files = self._get_deleted_files(repo) + # remove deleted files or `work` directory + if deleted_files: + repo.index.remove(deleted_files, r=True) + untracked_files = self._get_untracked_files(repo) + changed_files = self._get_changed_files(repo) + staged_files = set(untracked_files + changed_files) + # restore *.swp files in case users open any files during case is are running + self._restore_swp_files(repo, staged_files) + commit_message = f"Control experiment setup: Configure `{self.base_branch_name}` branch by `{self.yamlfile}`\n committed files/directories {staged_files}!" + if staged_files: + repo.index.add(staged_files) + repo.index.commit(commit_message) + else: + print( + f"Nothing changed, hence no further commits to the {self.base_path} repo!" + ) + + def manage_perturb_expt(self): + """ + Sets up perturbation experiments based on the configuration provided in `Expts_manager.yaml`. + + This function processes various parameter blocks defined in the YAML configuration, which may include + 1. namelist files (`_in`, `.nml`), + 2. MOM6 input files (`MOM_input`), + 3. `nuopc.runconfig`, + 4. `nuopc.runseq` (currently only for the coupling timestep). + + Raises: + - Warning: If no namelist configurations are provided, the function issues a warning indicating that no parameter tuning tests will be conducted. + """ + # main section, top level key that groups different namlists + namelists = self.indata["namelists"] + if not namelists: + warnings.warn( + "NO namelists were provided, hence there are no parameter-tunning tests!", + UserWarning, + ) + return + + for k, nmls in namelists.items(): + if not nmls: + continue + + # parameter tunning within one group in a single file + if not k.startswith("cross_block"): + self._process_params_blocks(k, nmls) + else: + # parameter tunning across multiple files + self._process_params_blocks_cross_files(k, namelists) + + def _process_params_blocks(self, k, nmls): + """ + Determines the type of parameter block and processes it accordingly. + + Args: + k (str): The key indicating the type of parameter block. + nmls (dict): The namelist dictionary for the parameter block. + """ + self.tag_model, expt_dir_name = self._determine_block_type(k) + + # parameter groups, in which contains one or more specific parameters + for k_sub in nmls: + self._process_params_group(k, k_sub, nmls, expt_dir_name, self.tag_model) + + def _determine_block_type(self, k): + """ + Determines the type of parameter block based on the key. + + Args: + k (str): The key indicating the type of parameter block. + """ + # parameter blocks, in which contains one or more groups of parameters, + # e.g., input.nml, ice_in etc. + if k.endswith(("_in", ".nml")): + tag_model = "nml" + elif k == "MOM_input": + tag_model = "mom6" + elif k == "nuopc.runseq": + tag_model = "cpl_dt" + elif k == "config.yaml": + tag_model = "config" + elif k == "nuopc.runconfig": + tag_model = "runconfig" + elif k.startswith("cross_block"): + tag_model = "cb" + else: + raise ValueError(f"Unsupported block type: {k}") + # [Optional] The key in the YAML input file specifies a list of + # user-defined directory names related to parameter testing. + expt_dir_name = k + "_dirs" + return tag_model, expt_dir_name + + def _count_second_level_keys(self, tmp_dict, expt_dir_name): + """ + Counts the number of groups + """ + group_count = 0 + for key, value in tmp_dict.items(): + # skip the user-defined expt directory name + if key == expt_dir_name: + continue + if isinstance(value, dict): + for inner_key, inner_value in value.items(): + if isinstance(inner_value, dict): + group_count += 1 + return group_count + + def _process_params_blocks_cross_files(self, k, namelists): + """ + Determines the type of parameter block for cross-blocks and processes them accordingly. + + Args: + k (str): The key indicating the type of parameter block. + namelists (dict): The highest-level namelist dictionary. + """ + self.tag_model, expt_dir_name = self._determine_block_type(k) + self.group_count = self._count_second_level_keys(namelists[k], expt_dir_name) + self.tmp_count = 0 + + # tmp_k => k (equivalent to `k`, when `cross_block` is disabled) + # k: filename + for tmp_k, tmp_nmls in namelists[k].items(): + if tmp_k.startswith(expt_dir_name): + self._set_cross_block_dirs(tmp_nmls) + else: + self._handle_params_cross_files(tmp_k, tmp_nmls) + + # reset user-defined dirs + self._reset_expt_names() + + def _set_cross_block_dirs(self, tmp_nmls): + """ + Sets cross block directories and sets perturbation directory names. + """ + self.expt_names = tmp_nmls # user-defined directories + self.num_expts = len(self.expt_names) # count dirs + + def _handle_params_cross_files(self, tmp_k, tmp_nmls): + """ + Processes all parameters in the namelist. + """ + for k_sub in tmp_nmls: + self.tmp_count += 1 + name_dict = tmp_nmls[k_sub] + if k_sub.endswith(self.combo_suffix): + if tmp_k.startswith("MOM_input"): + MOM_inputParser = self._parser_mom6_input( + os.path.join(self.base_path, "MOM_input") + ) + commt_dict = MOM_inputParser.commt_dict + else: + commt_dict = None + if name_dict is not None: + self._generate_combined_dicts(name_dict, commt_dict, k_sub, tmp_k) + self.setup_expts(tmp_k) + + def _reset_expt_names(self): + """ + Resets user-defined perturbation experiment names. + """ + self.expt_names = None + + def _parser_mom6_input(self, path): + """ + Parses MOM6 input file. + """ + mom6parser = self.MOM6InputParser.MOM6InputParser() + mom6parser.read_input(path) + mom6parser.parse_lines() + return mom6parser + + def _process_params_group(self, k, k_sub, nmls, expt_dir_name, tag_model): + """ + Processes individual parameter groups based on the tag model. + + Args: + k (str): The key indicating the type of parameter block. + k_sub (str): The key for the specific parameter group. + nmls (dict): The namelist dictionary for the parameter block. + expt_dir_name (str, optional): The key in the YAML file specifies a list of user-defined directory names related to parameter testing. + tag_model (str): The tag model indicating the type of parameter block. + """ + if tag_model == "nml": + self._handle_nml_group(k, k_sub, expt_dir_name, nmls) + elif tag_model == "mom6": + self._handle_mom6_group(k, k_sub, expt_dir_name, nmls) + elif tag_model == "cpl_dt": + self._handle_cpl_dt_group(k, k_sub, expt_dir_name, nmls) + elif tag_model == "config": + self._handle_config_group(k, k_sub, expt_dir_name, nmls) + elif tag_model == "runconfig": + self._handle_runconfig_group(k, k_sub, expt_dir_name, nmls) + self.previous_key = k_sub + + def _handle_config_group(self, k, k_sub, expt_dir_name, nmls): + """ + Handles config.yaml and nuopc.runconfig parameter groups specific to `config` tag model. + """ + if k_sub.startswith(self.CONFIG_prefix): + self._process_parameter_group_common(k, k_sub, nmls, expt_dir_name) + elif k_sub.startswith(expt_dir_name): + pass + else: + raise ValueError( + f"groupname must start with `{self.CONFIG_prefix}` " + f"or [optional] user-defined directory key must start with `{expt_dir_name}`!" + ) + + def _handle_runconfig_group(self, k, k_sub, expt_dir_name, nmls): + """ + Handles config.yaml and nuopc.runconfig parameter groups specific to `config` tag model. + """ + if ( + not k_sub.endswith(self.runconfig_suffix1) + or k_sub.endswith(self.runconfig_suffix2) + or k_sub.endswith(self.combo_suffix) + ): + self._process_parameter_group_common(k, k_sub, nmls, expt_dir_name) + elif k_sub.startswith(expt_dir_name): + pass + else: + raise ValueError( + f"groupname must end with either (`{self.runconfig_suffix1}` or `{self.runconfig_suffix2}` for single parameter tunning " + f"or `{self.combo_suffix}` for multiple parameter tunning " + f"or [optional] user-defined directory key must start with `{expt_dir_name}`!" + ) + + def _handle_nml_group(self, k, k_sub, expt_dir_name, nmls): + """ + Handles namelist parameter groups specific to `nml` tag model. + """ + if k_sub.endswith(self.nml_suffix) or k_sub.endswith(self.combo_suffix): + self._process_parameter_group_common(k, k_sub, nmls, expt_dir_name) + elif k_sub.startswith(expt_dir_name): + pass + else: + raise ValueError( + f"groupname must end with either `{self.nml_suffix}` for single parameter tunning " + f"or `{self.combo_suffix}` for multiple parameter tunning " + f"or [optional] user-defined directory key must start with `{expt_dir_name}`!" + ) + + def _handle_mom6_group(self, k, k_sub, expt_dir_name, nmls): + """ + Handles namelist parameter groups specific to `mom6` tag model. + """ + if k_sub.startswith(self.MOM_prefix): + MOM_inputParser = self._parser_mom6_input( + os.path.join(self.base_path, "MOM_input") + ) + commt_dict = MOM_inputParser.commt_dict + self._process_parameter_group_common( + k, + k_sub, + nmls, + expt_dir_name, + commt_dict=commt_dict, + ) + elif k_sub.startswith(expt_dir_name): + pass + else: + raise ValueError( + f"For the MOM6 input, the groupname must start with `{self.MOM_prefix}` " + f"or [optional] user-defined directory key must start with `{expt_dir_name}`!" + ) + + def _handle_cpl_dt_group(self, k, k_sub, expt_dir_name, nmls): + """ + Handles namelist parameter groups specific to `cpl_dt` tag model. + """ + if k_sub.startswith(self.runseq_prefix): + self._process_parameter_group_common(k, k_sub, nmls, expt_dir_name) + elif k_sub.startswith(expt_dir_name): + pass + else: + raise ValueError( + f"groupname must start with `{self.runseq_prefix}` " + f"or [optional] user-defined directory key must start with `{expt_dir_name}`!" + ) + + def _process_parameter_group_common( + self, k, k_sub, nmls, expt_dir_name, commt_dict=None + ): + """ + Processes parameter groups to all tag models. + + Args: + k (str): The key indicating the type of parameter block. + k_sub (str): The key for the specific parameter group. + nmls (dict): The namelist dictionary for the parameter block. + expt_dir_name (str, optional): The key in the YAML file specifies a list of user-defined directory names related to parameter testing. + commt_dict (dict, optional): A dictionary of comments, if applicable. + """ + name_dict = nmls[k_sub] + self._cal_num_expts(name_dict, k_sub) + if self.previous_key and self.previous_key.startswith(expt_dir_name): + self._valid_expt_names(nmls, name_dict) + if k_sub.endswith(self.combo_suffix): + self._generate_combined_dicts(name_dict, commt_dict, k_sub, k) + else: + self._generate_individual_dicts(name_dict, commt_dict, k_sub) + self.setup_expts(k) + + def _cal_num_expts(self, name_dict, k_sub): + """ + Evaluates the number of parameter-tunning experiments. + """ + if k_sub.endswith(self.combo_suffix): + if isinstance(next(iter(name_dict.values())), list): + self.num_expts = len(next(iter(name_dict.values()))) + else: + self.num_expts = 1 + else: + self.num_expts = 0 + for v_s in name_dict.values(): + if isinstance(v_s, list): + self.num_expts += len(v_s) + else: + self.num_expts = 1 + + def _valid_expt_names(self, nmls, name_dict): + """ + Compares the number of parameter-tunning experiments with [optional] user-defined experiment names. + """ + self.expt_names = nmls.get(self.previous_key) + if self.expt_names and len(self.expt_names) != self.num_expts: + raise ValueError( + f"The number of user-defined experiment directories {self.expt_names} " + f"is different from that of tunning parameters {name_dict}!" + f"\nPlease double check the number or leave it/them blank!" + ) + + def _generate_individual_dicts(self, name_dict, commt_dict, k_sub): + """ + Each dictionary has a single key-value pair. + """ + param_dict_change_list = [] + append_group_list = [] + for k, vs in name_dict.items(): + if isinstance(vs, list): + for v in vs: + param_dict_change = {k: v} + param_dict_change_list.append(param_dict_change) + append_group = k_sub + append_group_list.append(append_group) + else: + param_dict_change = {k: vs} + param_dict_change_list.append(param_dict_change) + append_group_list = [k_sub] + self.param_dict_change_list = param_dict_change_list + if self.tag_model == "mom6": + self.commt_dict_change = {k: commt_dict.get(k, "") for k in name_dict} + elif self.tag_model in (("nml", "config", "runconfig")): + self.append_group_list = append_group_list + + def _generate_combined_dicts(self, name_dict, commt_dict, k_sub, parameter_block): + """ + Generates a list of dictionaries where each dictionary contains all keys with values from the same index. + """ + param_dict_change_list = [] + append_group_list = [] + for i in range(self.num_expts): + name_dict = self._preprocess_nested_dicts(name_dict) + param_dict_change = {k: name_dict[k][i] for k in name_dict} + append_group = k_sub + append_group_list.append(append_group) + param_dict_change_list.append(param_dict_change) + self.param_dict_change_list = param_dict_change_list + + if self.tag_model == "mom6" or parameter_block == "MOM_input": + self.commt_dict_change = {k: commt_dict.get(k, "") for k in name_dict} + elif ( + self.tag_model in (("nml", "config", "runconfig", "cpl_dt")) + or parameter_block.endswith(("_in", ".nml")) + or parameter_block in (("config.yaml", "nuopc.runconfig", "nuopc.runseq")) + ): + self.append_group_list = append_group_list + + def _preprocess_nested_dicts(self, input_data): + """ + Pre-processes nested dictionary with lists. + """ + res_dicts = {} + for tmp_key, tmp_values in input_data.items(): + if isinstance(tmp_values, list) and all( + isinstance(v, dict) for v in tmp_values + ): + res_dicts[tmp_key] = [] + num_entries = len(next(iter(tmp_values[0].values()))) + for i in range(num_entries): + entry_list = [] + for submodel in tmp_values: + entry = {} + for k, v in submodel.items(): + entry[k] = v[i] + entry_list.append(entry) + res_dicts[tmp_key].append(entry_list) + else: + res_dicts[tmp_key] = tmp_values + return res_dicts + + def setup_expts(self, parameter_block): + """ + Sets up perturbation experiments based on the YAML input file provided in `Expts_manager.yaml`. + """ + for i, param_dict in enumerate(self.param_dict_change_list): + print(f"-- tunning parameters: {param_dict}") + # generate perturbation experiment directory names + expt_name = self._generate_expt_names(i) + + # perturbation experiment path + expt_path = os.path.join(self.dir_manager, self.test_path, expt_name) + + # generate perturbation experiment directory + if os.path.exists(expt_path): + if self.tmp_count == self.group_count or self.tag_model != "cb": + print(f"-- not creating {expt_path} - already exists!") + else: + if self.tmp_count == 1 or self.tag_model != "cb": + self._generate_expt_directory(expt_path, parameter_block, i) + + if self.tmp_count == self.group_count or self.tag_model != "cb": + # optionally update diag_table for perturbation runs + if self.diag_pert and self.diag_path: + self._copy_diag_table(expt_path) + + # symlink restart directories + restartpath = self._generate_restart_symlink(expt_path) + self._update_metadata_yaml_perturb(expt_path, param_dict, restartpath) + + # update jobname same as perturbation experiment name + self._update_perturb_jobname(expt_path, expt_name) + + # optionally update nuopc.runconfig for perturbation runs + # if there is no parameter tunning under cb or runconfig flags! + if self.tag_model not in (("cb", "runconfig")): + self._update_nuopc_config_perturb(expt_path) + + # update params for each parameter block + if self.tag_model == "mom6" or parameter_block == "MOM_input": + self._update_mom6_params(expt_path, param_dict) + elif ( + self.tag_model == "nml" + or parameter_block.endswith("_in") + or parameter_block.endswith(".nml") + ): + self._update_nml_params(expt_path, param_dict, parameter_block, i) + elif self.tag_model == "cpl_dt" or parameter_block == "nuopc.runseq": + self._update_cpl_dt_params(expt_path, param_dict, parameter_block) + elif self.tag_model == "config" or parameter_block == "config.yaml": + self._update_config_params(expt_path, param_dict, parameter_block) + elif self.tag_model == "runconfig" or parameter_block == "nuopc.runconfig": + self._update_runconfig_params(expt_path, param_dict, parameter_block, i) + + if self.tmp_count == self.group_count or self.tag_model != "cb": + pbs_jobs = self._output_existing_pbs_jobs() + if self.check_duplicate_jobs: + duplicated_bool = self._check_duplicated_jobs(pbs_jobs, expt_path) + else: + duplicated_bool = False + + # start runs, count existing runs and do additional runs if needed + self._start_experiment_runs( + expt_path, expt_name, duplicated_bool, self.nruns + ) + + if self.tag_model != "cb": + # reset to None after the loop to update user-defined perturbation experiment names! + self._reset_expt_names() + + def _generate_expt_names(self, indx): + if self.expt_names is None: + # if `expt_names` does not exist, `expt_names` is set as the tunning parameters appending with associated values. + return "_".join( + [f"{k}_{v}" for k, v in self.param_dict_change_list[indx].items()] + ) + # user-defined directory names for each parameter-tunning experiment. + return self.expt_names[indx] + + def _generate_expt_directory(self, expt_path, parameter_block, indx): + """ + Generates a new experiment directory by cloning the control experiment. + Checks if the tuning parameter matches the control experiment, + this validation currently applies only to `nml` files. + + Args: + expt_path (str): The path to the experiment directory. + param_dict (dict): The dictionary of parameters to update. + """ + if self.check_skipping: + if self.tag_model == "nml": + self._check_skipping( + self.param_dict_change_list[indx], + self.append_group_list[indx], + parameter_block, + expt_path, + ) + elif self.tag_model == "mom6": # TODO + pass + elif self.tag_model == "cpl_dt": # TODO + pass + + print(f"Directory {expt_path} not exists, hence cloning template!") + command = f"payu clone -B {self.base_branch_name} -b {self.branch_perturb} {self.base_path} {expt_path}" # automatically leave a commit with expt uuid + subprocess.run(command, shell=True, check=True) + + def _update_mom6_params(self, expt_path, param_dict): + """ + Updates MOM6 parameters in the 'MOM_override' file. + + Args: + expt_path (str): The path to the experiment directory. + param_dict (dict): The dictionary of parameters to update. + """ + MOM6_or_parser = self._parser_mom6_input( + os.path.join(expt_path, "MOM_override") + ) + MOM6_or_parser.param_dict, MOM6_or_parser.commt_dict = ( + update_MOM6_params_override(param_dict, self.commt_dict_change) + ) + MOM6_or_parser.writefile_MOM_input(os.path.join(expt_path, "MOM_override")) + + def _update_nml_params(self, expt_path, param_dict, parameter_block, indx=None): + """ + Updates namelist parameters and overwrites namelist file. + + Args: + expt_path (str): The path to the experiment directory. + param_dict (dict): The dictionary of parameters to update. + parameter_block (str): The name of the namelist file. + """ + nml_path = os.path.join(expt_path, parameter_block) + if indx is not None: + nml_group = self.append_group_list[indx] + # rename the namlist by removing the suffix if the suffix with `_combo` + if nml_group.endswith(self.combo_suffix): + nml_group = nml_group[: -len(self.combo_suffix)] + patch_dict = {nml_group: {}} + for nml_name, nml_value in param_dict.items(): + if nml_name == "turning_angle": + cosw = np.cos(nml_value * np.pi / 180.0) + sinw = np.sin(nml_value * np.pi / 180.0) + patch_dict[nml_group]["cosw"] = cosw + patch_dict[nml_group]["sinw"] = sinw + else: # for generic parameters + patch_dict[nml_group][nml_name] = nml_value + param_dict = patch_dict + f90nml.patch(nml_path, param_dict, nml_path + "_tmp") + else: + f90nml.patch(nml_path, param_dict, nml_path + "_tmp") + os.rename(nml_path + "_tmp", nml_path) + + self._format_nml_params(nml_path, param_dict) + + def _format_nml_params(self, nml_path, param_dict): + """ + Handles pre-formatted strings or values. + + Args: + nml_path (str): The path to specific f90 namelist file. + param_dict (dict): The dictionary of parameters to update. + e.g., in yaml input file, + ocean/input.nml: + mom_oasis3_interface_nml: + fields_in: "'u_flux', 'v_flux', 'lprec'" + fields_out: "'t_surf', 's_surf', 'u_surf'" + results in, + &mom_oasis3_interface_nml + fields_in = 'u_flux', 'v_flux', 'lprec' + fields_out = 't_surf', 's_surf', 'u_surf' + """ + with open(nml_path, "r") as f: + fileread = f.readlines() + for tmp_group, tmp_subgroups in param_dict.items(): + for tmp_param, tmp_values in tmp_subgroups.items(): + for i in range(len(fileread)): + if tmp_param in fileread[i]: + fileread[i] = f" {tmp_param} = {tmp_values}\n" + break + with open(nml_path, "w") as f: + f.writelines(fileread) + + def _update_config_params(self, expt_path, param_dict, parameter_block): + """ + Updates namelist parameters and overwrites namelist file. + + Args: + expt_path (str): The path to the experiment directory. + param_dict (dict): The dictionary of parameters to update. + parameter_block (str): The name of the namelist file. + """ + nml_path = os.path.join(expt_path, parameter_block) + expt_name = os.path.basename(expt_path) + + file_read = self._read_ryaml(nml_path) + if "jobname" in param_dict: + if param_dict["jobname"] != expt_name: + warnings.warn( + f"\n" + f"-- jobname must be the same as {expt_name}, " + f"hence jobname is forced to be {expt_name}!", + UserWarning, + ) + param_dict["jobname"] = expt_name + self._update_config_entries(file_read, param_dict) + self._write_ryaml(file_read, nml_path) + + def _update_runconfig_params( + self, expt_path, param_dict, parameter_block, indx=None + ): + """ + Updates namelist parameters and overwrites namelist file. + + Args: + expt_path (str): The path to the experiment directory. + param_dict (dict): The dictionary of parameters to update. + parameter_block (str): The name of the namelist file. + """ + nml_path = os.path.join(expt_path, parameter_block) + if indx is not None: + nml_group = self.append_group_list[indx] + # rename the namlist by removing the suffix if the suffix with `_combo` + if nml_group.endswith(self.combo_suffix): + nml_group = nml_group[: -len(self.combo_suffix)] + param_dict = self.nested_dict(nml_group, param_dict) + file_read = self.read_nuopc_config(nml_path) + self._update_config_entries(file_read, param_dict) + self.write_nuopc_config(file_read, nml_path) + + def nested_dict(self, outer_key, inner_dict): + return {outer_key: inner_dict} + + def _update_cpl_dt_params(self, expt_path, param_dict, parameter_block): + """ + Updates coupling timestep parameters. + + Args: + expt_path (str): The path to the experiment directory. + param_dict (dict): The dictionary of parameters to update. + """ + nuopc_runseq_file = os.path.join(expt_path, parameter_block) + self._update_cpl_dt_nuopc_seq( + nuopc_runseq_file, param_dict[next(iter(param_dict.keys()))] + ) + + def _generate_restart_symlink(self, expt_path): + """ + Generates a symlink to the restart directory if needed. + + Args: + expt_path (str): The path to the experiment directory. + """ + if self.startfrom_str != "rest": + link_restart = os.path.join("archive", "restart" + self.startfrom_str) + # restart dir from control experiment + restartpath = os.path.realpath(os.path.join(self.base_path, link_restart)) + # restart dir symlink for each perturbation experiment + dest = os.path.join(expt_path, link_restart) + + # only generate symlink if it doesnt exist or force_restart is enabled + if ( + not os.path.islink(dest) + or self.force_restart + or (os.path.islink(dest) and not os.path.exists(os.readlink(dest))) + ): + if os.path.exists(dest) or os.path.islink(dest): + os.remove(dest) # remove symlink + print(f"-- Remove restart symlink: {dest}") + os.symlink(restartpath, dest) # generate a new symlink + if self.force_restart: + print(f"-- Restart symlink has been forced to be : {dest}") + else: + print(f"-- Restart symlink: {dest}") + # print restart symlink on the screen + else: + print(f"-- Restart symlink: {dest}") + else: + restartpath = "rest" + print(f"-- Restart symlink: {restartpath}") + + return restartpath + + def _update_nuopc_config_perturb(self, path): + """ + Updates nuopc.runconfig for perturbation experiment runs. + """ + nuopc_input = self.indata.get("perturb_run_config", None) + if nuopc_input is not None: + nuopc_file_path = os.path.join(path, "nuopc.runconfig") + nuopc_runconfig = self.read_nuopc_config(nuopc_file_path) + self._update_config_entries(nuopc_runconfig, nuopc_input) + self.write_nuopc_config(nuopc_runconfig, nuopc_file_path) + + def _update_perturb_jobname(self, expt_path, expt_name): + """ + Updates `jobname` only for now. + Args: + expt_path (str): The path to the perturbation experiment directory. + expt_name (str): The name of the perturbation experiment. + """ + config_path = os.path.join(expt_path, "config.yaml") + config_data = self._read_ryaml(config_path) + config_data["jobname"] = expt_name + self._write_ryaml(config_data, config_path) + + def _update_metadata_yaml_perturb(self, expt_path, param_dict, restartpath): + """ + Updates the `metadata.yaml` file with relevant metadata. + + Args: + expt_path (str): The path to the perturbation experiment directory. + param_dict (dict): The dictionary of parameters to include in metadata. + """ + metadata_path = os.path.join(expt_path, "metadata.yaml") + metadata = self._read_ryaml(metadata_path) # load metadata of each perturbation + self._update_metadata_description(metadata, restartpath) # update `description` + + # remove None comments from `description` + self._remove_metadata_comments("description", metadata) + keywords = self._extract_metadata_keywords(param_dict) + + # extract parameters from the change list, and update `keywords` + metadata["keywords"] = ( + f"{self.base_dir_name}, {self.branch_perturb}, {keywords}" + ) + + # remove None comments from `keywords` + self._remove_metadata_comments("keywords", metadata) + + self._write_ryaml(metadata, metadata_path) # write to file + + def _output_existing_pbs_jobs(self): + """ + Checks the existing qstat pbs information. + """ + current_job_status_path = os.path.join(self.dir_manager, "current_job_status") + command = f"qstat -f > {current_job_status_path}" + subprocess.run(command, shell=True, check=False) + + pbs_jobs = {} + current_key = None + current_value = "" + job_id = None + with open(current_job_status_path, "r") as f: + pbs_job_file = f.read() + + pbs_job_file = pbs_job_file.replace("\t", " ") + + for line in pbs_job_file.splitlines(): + line = line.rstrip() + if not line: + continue + if line.startswith("Job Id:"): + job_id = line.split(":", 1)[1].strip() + pbs_jobs[job_id] = {} + current_key = None + current_value = "" + elif line.startswith(" ") and current_key: # 8 indents multi-line + current_value += line.strip() + elif line.startswith(" ") and " = " in line: # 4 indents for new pair + # Save the previous multi-line value + if current_key: + pbs_jobs[job_id][current_key] = current_value.strip() + key, value = line.split(" = ", 1) # save key + current_key = key.strip() + current_value = value.strip() + + # remove the `current_job_status` file + os.remove(current_job_status_path) + + return pbs_jobs + + def _check_duplicated_jobs(self, pbs_jobs, expt_path): + + def extract_current_and_parent_path(tmp_path): + + # extract base_name or expt_name from pbs jobs + folder_path = "/" + "/".join(tmp_path.split("/")[1:-1]) + + # extract test_path from pbs jobs + parent_path = "/" + "/".join(tmp_path.split("/")[1:-2]) + + return folder_path, parent_path + + parent_paths = {} + for job_id, job_info in pbs_jobs.items(): + folder_path, parent_path = extract_current_and_parent_path( + job_info["Error_Path"] + ) + + job_state = job_info["job_state"] + if job_state not in ("F", "S"): + if parent_path not in parent_paths: + parent_paths[parent_path] = [] + parent_paths[parent_path].append(folder_path) + duplicated = False + + for parent_path, folder_paths in parent_paths.items(): + if expt_path in folder_paths: + print( + f"-- You have duplicated runs for folder '{os.path.basename(expt_path)}' in the same folder '{parent_path}', " + f"hence not submitting this job!\n" + ) + duplicated = True + return duplicated + + def _start_experiment_runs(self, expt_path, expt_name, duplicated, num_runs): + """ + Runs perturbation experiments. + + Args: + expt_path (str): The path to the control/perturbation experiment directory. + expt_name (str): The name of the control/perturbation experiment. + """ + + def runs(): + doneruns = len( + glob.glob(os.path.join(expt_path, "archive", "output[0-9][0-9][0-9]*")) + ) + newruns = num_runs - doneruns + if newruns > 0: + print(f"\nRun experiment -n {newruns}\n") + command = f"cd {expt_path} && payu run -n {newruns} -f" + subprocess.run(command, shell=True, check=False) + print("\n") + else: + print( + f"-- `{expt_name}` has already completed {doneruns} runs! Hence, stopping further runs.\n" + ) + + if not duplicated: + + # clean `work` directory for failed jobs + self._clean_workspace(expt_path) + + if num_runs > 0: + runs() + else: + print( + f"-- number of runs is {num_runs}, hence no new experiments will start!\n" + ) + + def _clean_workspace(self, dir_path): + """ + Cleans `work` directory for failed jobs. + """ + work_dir = os.path.join(dir_path, "work") + # in case any failed job + if os.path.islink(work_dir) and os.path.isdir(work_dir): + # Payu sweep && setup to ensure the changes correctly && remove the `work` directory + command = f"payu sweep && payu setup" + subprocess.run(command, shell=True, check=False) + print(f"Clean up a failed job {work_dir} and prepare it for resubmission.") + + def _check_skipping(self, param_dict, nml_group, parameter_block, expt_path): + """ + Checks if the tuning parameter matches the control experiment, + this validation currently applies only to `nml` files. + """ + if self.tag_model == "nml": + # rename the namlist if suffix with `_combo` + if nml_group.endswith(self.combo_suffix): + nml_group = nml_group[: -len(self.combo_suffix)] + nml_name = param_dict.keys() + if len(nml_name) == 1: # one param:value pair + nml_value = param_dict[list(nml_name)[0]] + else: # combination of param:value pairs + nml_value = [param_dict[j] for j in nml_name] + + if "turning_angle" in param_dict: + cosw = np.cos(param_dict["turning_angle"] * np.pi / 180.0) + sinw = np.sin(param_dict["turning_angle"] * np.pi / 180.0) + + # load nml of the control experiment + self.nml_ctrl = f90nml.read(os.path.join(self.base_path, parameter_block)) + + # nml_name (i.e. tunning parameter) may not be found in the control experiment + if all(cn in self.nml_ctrl.get(nml_group, {}) for cn in nml_name): + if "turning_angle" in param_dict: + skip = ( + self.nml_ctrl[nml_group]["cosw"] == cosw + and self.nml_ctrl[nml_group]["sinw"] == sinw + and all( + self.nml_ctrl[nml_group].get(cn) == param_dict[cn] + for cn in nml_name + if cn not in ["cosw", "sinw"] + ) + ) + else: + skip = all( + self.nml_ctrl[nml_group].get(cn) == param_dict[cn] + for cn in nml_name + ) + else: + print( + f"Not all {nml_name} are found in {nml_group}, hence not skipping!" + ) + skip = False + + if skip: + print( + "-- not creating", + expt_path, + "- parameters are identical to the control experiment located at", + self.base_path, + "\n", + ) + return + + # might need MOM_parameter.all, because many parameters are in-default hence not shown up in `MOM_input` + if self.tag_model == "mom6": + # TODO + pass + + def _parser_mom6_input(self, path): + """ + Parses MOM6 input file. + """ + mom6parser = self.MOM6InputParser.MOM6InputParser() + mom6parser.read_input(path) + mom6parser.parse_lines() + return mom6parser + + def _update_config_entries(self, base, change): + """ + Recursively update nuopc_runconfig and config.yaml entries. + """ + for k, v in change.items(): + if isinstance(v, dict) and k in base: + self._update_config_entries(base[k], v) + else: + base[k] = v + + def _update_cpl_dt_nuopc_seq(self, seq_path, update_cpl_dt): + """ + Updates only coupling timestep through nuopc.runseq. + """ + with open(seq_path, "r") as f: + lines = f.readlines() + pattern = re.compile(r"@(\S*)") + update_lines = [] + for l in lines: + matches = pattern.findall(l) + if matches: + update_line = re.sub(r"@(\S+)", f"@{update_cpl_dt}", l) + update_lines.append(update_line) + else: + update_lines.append(l) + with open(seq_path, "w") as f: + f.writelines(update_lines) + + def _get_untracked_files(self, repo): + """ + Gets untracked git files. + """ + return repo.untracked_files + + def _get_changed_files(self, repo): + """ + Gets changed git files. + """ + return [file.a_path for file in repo.index.diff(None)] + + def _get_deleted_files(self, repo): + """ + Gets deleted git files. + """ + return [file.a_path for file in repo.index.diff(None) if file.deleted_file] + + def _restore_swp_files(self, repo, staged_files): + """ + Restores tmp git files. + """ + swp_files = [file for file in staged_files if file.endswith(".swp")] + for file in swp_files: + repo.git.restore(file) + + def _read_ryaml(self, yaml_path): + """ + Reads YAML file and preserve comments. + """ + with open(yaml_path, "r") as f: + return ryaml.load(f) + + def _write_ryaml(self, data, yaml_path): + """ + Writes YAML file and preserve comments. + """ + with open(yaml_path, "w") as f: + ryaml.dump(data, f) + + def _update_metadata_description(self, metadata, restartpath): + """ + Updates metadata description with experiment details. + """ + tmp_string1 = ( + f"\nNOTE: this is a perturbation experiment, but the description above is for the control run." + f"\nThis perturbation experiment is based on the control run {self.base_path} from {self.base_branch_name}" + ) + tmp_string2 = f"\nbut with initial condition {restartpath}." + desc = metadata["description"] + if desc is None: + desc = "" + if tmp_string1.strip() not in desc.strip(): + desc += tmp_string1 + if tmp_string2.strip() not in desc.strip(): + desc += tmp_string2 + metadata["description"] = LiteralString(desc) + + def _remove_metadata_comments(self, key, metadata): + """ + Removes comments after the key in metadata. + """ + if key in metadata: + metadata.ca.items[key] = [None, None, None, None] + + def _extract_metadata_keywords(self, param_change_dict): + """ + Extracts keywords from parameter change dictionary. + """ + keywords = ", ".join(param_change_dict.keys()) + return keywords + + def main(self): + """ + Main function for the program. + """ + parser = argparse.ArgumentParser( + description="Manage either ACCESS-OM2 or ACCESS-OM3 experiments.\ + Latest version and help: https://github.com/COSIMA/om3-scripts/pull/34" + ) + parser.add_argument( + "INPUT_YAML", + type=str, + nargs="?", + default="Expts_manager.yaml", + help="YAML file specifying parameter values for expt runs. Default is Expts_manager.yaml", + ) + args = parser.parse_args() + INPUT_YAML = vars(args)["INPUT_YAML"] + + yamlfile = os.path.join(self.dir_manager, INPUT_YAML) + self.load_variables(yamlfile) + self.create_test_path() + self.model_selection() + self.load_tools() + self.manage_ctrl_expt() + if self.run_namelists: + print("==== Start perturbation experiments ====") + self.manage_perturb_expt() + else: + print("==== No perturbation experiments are prescribed ====") + + +if __name__ == "__main__": + expt_manager = Expts_manager() + expt_manager.main() diff --git a/expts_manager/Expts_manager.yaml b/expts_manager/Expts_manager.yaml new file mode 100644 index 0000000..dbf4a67 --- /dev/null +++ b/expts_manager/Expts_manager.yaml @@ -0,0 +1,286 @@ + +# ===================================================================================== +# YAML Configuration for Expts_manager.py +# ===================================================================================== +# This configuration file defining the parameters and settings required for cloning, +# setting up, and running control and perturbation experiments using `Expts_manager.py`. +# Detailed explanations are provided to ensure the configuration is straightforward. +# ===================================================================================== + + +# ============ Model Selection ======================================================== + +model: access-om3 # Specify the model to be used. Options: 'access-om2', 'access-om3' + +# ============ Utility Tool Configuration (only for access-om3) ======================= +# The following configuration provides the necessary tools to: +# 1. Parse parameters and comments from `MOM_input` in MOM6. +# 2. Read and write the `nuopc.runconfig` file. + +force_overwrite_tools: False +utils_url: git@github.com:minghangli-uni/om3-utils.git # Git URL for the utility tool repository +utils_branch_name: main # The branch for which the utility tool will be checked out +utils_dir_name: om3-utils # Directory name for the utility tool (user-defined) + +# ============ Diagnostic Table (optional) ============================================ +# Configuration for customising the diagnostic table. + +diag_url: git@github.com:minghangli-uni/make_diag_table.git # Git URL for the `make_diag_table` +diag_branch_name: general_scheme1 # Branch for the `make_diag_table` +diag_dir_name: make_diag_table # Directory name for the `make_diag_table` (user-defined) +diag_ctrl: False # Set to 'True' to modify the diagnostic table for the control experiment +diag_pert: False # Set to 'True' to modify the diagnostic table for perturbation experiments + +# ============ Control Experiment Setup =============================================== + +base_url: git@github.com:ACCESS-NRI/access-om3-configs.git # Git URL for the control experiment repository +base_commit: "2bc6107" # Commit hash to use; Please ensure it is a string! +base_dir_name: Ctrl-1deg_jra55do_ryf # Directory name for cloning (user-defined) +base_branch_name: ctrl # Branch name for the experiment (user-defined) +test_path: test # Relative path for all test (control and perturbation) runs (user-defined) + +# ============ Control Experiment Variables =========================================== +# Allows modification of various control experiment settings. +# 1. config.yaml (access-om2 or access-om3) +# 2. all namelists such as with endswith "_in" or ".nml" etc. (access-om2 or access-om3) +# 3. cpl_dt (coupling timestep) (access-om3) +# 4. nuopc.runconfig (access-om3) +# 5. MOM_input (access-om3) +# Below are some examples for the illustration purpose, please modify for your own settings. + +config.yaml: + ncpus: 240 + mem: 960GB + walltime: 24:00:00 + #jobname: # `jobname` will be forced to be the name of the directory, which is `Ctrl-1deg_jra55do_ryf` in this example. + metadata: + enable: True + + runlog: True + restart_freq: 1 + +cpl_dt: 1800.0 # Coupling timestep in the `nuopc_runseq` +nuopc.runconfig: + CLOCK_attributes: + stop_option: ndays + stop_n: 1 + restart_option: ndays + restart_n: 1 + + PELAYOUT_attributes: + atm_ntasks: 24 + cpl_ntasks: 24 + glc_ntasks: 1 + ice_ntasks: 24 + lnd_ntasks: 1 + ocn_ntasks: 216 + ocn_rootpe: 24 + rof_ntasks: 24 + wav_ntasks: 1 + + +ice_in: + domain_nml: + max_blocks: -1 + block_size_x: 15 + block_size_y: 300 + + +input.nml: + diag_manager_nml: + max_axes: 400 + max_files: 200 + max_num_axis_sets: 200 + + +MOM_input: + HFREEZE: 12 + BOUND_SALINITY: False + + +# ============ Namelist Tunning ================================ +# Tune parameters across different model components. + +# Generalised structure +# 1. Single-parameter tunning within single file +# namelists: +# filename: (Required). +# filename_dirs: List of directory names (Optional - user-defined: filename must be appended with "_dirs"; additional strings may follow). +# groupname: (Required: for f90 namelists or nuopc.runconfig, this is simply the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list"). +# parameter1: list of values +# parameter2: list of values +# ... +# ... + +# 2. Multiple-parameter tuning within a single group in a single file +# namelists: +# filename: (Required). +# filename_dirs: List of directory names (Optional - user-defined: filename must be appended with "_dirs"; additional strings may follow). +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, use the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list", then append "_combo"). +# parameter1: list of values +# parameter2: list of values +# ... +# ... + +# 3. Multiple-parameter tuning across different files using user-defined directories +# namelists: +# cross_block: (Required: additional strings may follow "cross_block"). +# cross_block_dirs: List of directory names (Required - user-defined: filename must be appended with "_dirs"; additional strings may follow). +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, use the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list", then append "_combo"). +# parameter1: list of values +# parameter2: list of values +# ... +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, use the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list", then append "_combo"). +# parameter1: list of values +# parameter2: list of values +# ... +# ... +# +# cross_block: (Required: may append other strings after "cross_block"). +# cross_block_dirs: list of directory names (Optional - user-defined: filename must be appended with "_dirs", you may append other strings after "_dirs"). +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, just the group name; for MOM input, "MOM_list", nuopc.runseq, "runseq_list", then append "_combo" at the end). +# parameter1: list of values +# parameter2: list of values +# ... +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, just the group name; for MOM input, "MOM_list", nuopc.runseq, "runseq_list", then append "_combo" at the end). +# parameter1: list of values +# parameter2: list of values +# ... +# ... +# ... + +# The following namelist options are provided as examples. +# Note: Parameters should be based on the underlying physics. +# Some configurations may fail to run or produce invalid results. +# 1.1 Single parameter tuning within a single file using default perturbation directory names. +# 1.2 Single parameter tuning within a single file using user-defined perturbation directory names. +# 2.1 Multi-parameter tuning within a single group in a single file using default perturbation directory names. +# 2.2 Multi-parameter tuning within a single group in a single file using user-defined perturbation directory names. +# 3. Multi-parameter tuning across multiple files. + +namelists: + ice_in: + ponds_nml: + dpscale: [0.002, 0.003] + rfracmax: [1.1, 1.2] + + ice_in_dirs_test: [icet1, icet2, icet3, icet4] + shortwave_nml: + ahmax: [0.2, 0.3] + r_snw: [0.1, 0.2] + + + MOM_input: + MOM_list4: + MINIMUM_DEPTH: [0, 0.5] + MAXIMUM_DEPTH: [6000.0, 7000.0, 8000.0] + + MOM_input_dirs1: [lexpt0, lexpt1] + MOM_list1_combo: + DT_THERM: [3600.0, 108000.0] + DIABATIC_FIRST: [False, False] + THERMO_SPANS_COUPLING: [True, True] + + + cross_block2: + cross_block2_dirs: [cg_1, cg_2] + nuopc.runconfig: + CLOCK_attributes_combo: + restart_n: [1, 1] + restart_option: [ndays, ndays] + stop_n: [1, 1] + stop_option: [ndays, ndays] + ocn_cpl_dt: [1800.0, 7200.0] + + PELAYOUT_attributes_combo: + ocn_ntasks: [168, 120] + + + config.yaml: + config_list1_combo: + ncpus: [192, 144] + mem: [768GB, 576GB] + + + ice_in: + shortwave_nml_combo: + albicei: [0.36, 0.39] + albicev: [0.78, 0.81] + + ponds_nml_combo: + dpscale: [0.002, 0.003] + + + MOM_input: + MOM_list1_combo: + DT_THERM: [3600.0, 7200.0] + DIABATIC_FIRST: [False, False] + THERMO_SPANS_COUPLING: [True, True] + DTBT_RESET_PERIOD: [3600.0, 7200.0] + + + nuopc.runseq: + runseq_list1_combo: + cpl_dt: [1800.0, 7200.0] + + + + + +# ============ Perturbation Experiment Setup (Optional - access-om3) ======================= +# Configure settings for perturbation experiments. Currently, only `nuopc.runconfig` is supported. +# If conducting parameter tuning for `nuopc.runconfig`, any pre-existing settings in this section +# will be purged by the above namelist tunning. + +perturb_run_config: + CLOCK_attributes: + stop_option: nyears + stop_n: 1 + restart_option: nyears + restart_n: 1 + + + +# ============ Control experiment and perturbation Runs =================================== +# This section configures the settings for running control experiments and their corresponding perturbation tests. + +ctrl_nruns: 0 +# Number of control experiment runs. +# Default: 0. +# Adjust this value to set the number of iterations for the control experiment, which serves as the baseline for perturbations. + +run_namelists: True +# Determines whether to run using the specified namelists. +# Default: False. +# Set to 'True' to apply configurations from the namelist section; otherwise, 'False' to skip this step. + +check_duplicate_jobs: True +# Checks if there are duplicate PBS jobs within the same parent directory (`test_path`) based on their names. +# This check is useful when you have existing running jobs and want to add additional tests, which helps avoid conflicts by ensuring new jobs don't duplicate existing ones in the same `test_path`. +# The check will not be triggered if the jobs are located in different `test_path`. It only applies to jobs within the same `test_path` directory. +# Default: True. +# If duplicates are found, a message will be printed indicating the presence of duplicate runs and those runs will be skipped. + +check_skipping: False +# Checks if certain runs should be skipped based on pre-existing conditions or results. Currently only valid for nml type. +# Default: False. +# Set to 'True' if you want the system to skip runs under specific criteria; otherwise, keep it 'False'. Currently only valid for nml type. + +force_restart: False +# Controls whether to force a restart of the experiments regardless of existing initial conditions. +# Default: False. +# Set to 'True' to enforce a restart of the control and perturbation runs. + +startfrom: 'rest' +# Defines the starting point for perturbation tests. +# Options: a specific restart number of the control experiment, or 'rest' to start from the initial state. +# This parameter determines the initial condition for each perturbation test. + +nruns: 0 +# Total number of output directories to generate for each Expts_manager member. +# Default: 0. +# Specifies how many runs of each perturbation experiment will be started; this number correlates with the different parameter combinations defined in the configuration. diff --git a/expts_manager/README.md b/expts_manager/README.md new file mode 100644 index 0000000..38a799c --- /dev/null +++ b/expts_manager/README.md @@ -0,0 +1,65 @@ +# ACCESS-OM Expts Manager +The **ACCESS-OM Experiment Manager** is a Python-based tool designed to streamline the setup and management of either **ACCESS-OM2** or **ACCESS-OM3**. This tool automates the creation of experiment directories, applies parameter changes, and updates relevant configuration files based on user-defined settings in a YAML configuration file. + +## Directory Structure +``` +. +├── Expts_manager.py +├── write_Expts_manager_yaml.py +├── Expts_manager.yaml +└── README.md + +0 directories, 4 files +``` + +### Components: +1. `Expts_manager.py`: + - This file contains the `ExptManager` class, which faciliates the setup, configuration, and execution of experiments. Key functionalities include: + - Configuration management based on user-defined parameters. + - Automation of directory creation, ensuring a smooth workflow for running experiments. + - Support for parameter updates, including: + - MOM6 for **ACCESS-OM3**, + - Fortran namelists for both **ACCESS-OM2** and **ACCESS-OM3**, + - Coupling timesteps (`nuopc.runseq`) for **ACCESS-OM3**, + - Component settings (`nuopc.runconfig`) for **ACCESS-OM3**. + + - Automation of experiment initiation, managing the number of runs to be executed. + - Ability to skip experiment runs if parameters are identical to those of the control experiment. This functionality can be toggled (currently applicable only to `f90nml`). + - Prevention of duplicate PBS jobs within the same `test_path` directory. + - Integration of Git management to track changes during experiments. + - Updating of experiment metadata, including details and descriptions, facilitated by `Payu`. + +2. `Expts_manager.yaml`: + - This YAML configuration file defines the parameters and settings required for managing control and perturbation experiments. It enables users to: + - Clone necessary repositories, + - Setup experiments with customised configurations, + - Manage utility tools, diagnostic tools, and parameter tuning. + +3. `write_Expts_manager_yaml.py`: + - A Python script used to generate `Expts_manager.yaml`. + +## Usage +1. **Edit the YAML File:** Customise experiment parameters and settings by editing `Expts_manager.yaml` or `write_Expts_manager_yaml.py`. Documentation and examples are provided within the file. + - To generate Expts_manager.yaml, simply run: + ``` + ./write_Expts_manager_yaml.py + ``` + +2. **Run the Manager:** Execute the experiment manager script using the following command: + ``` + ./Expts_manager.py + ``` + +4. **View Available Options:** Users can view available options by: + ``` + $ ./Expts_manager.py -h + usage: Expts_manager.py [-h] [INPUT_YAML] + + Manage both ACCESS-OM2 and ACCESS-OM3 experiments. Latest version and help: https://github.com/COSIMA/om3-scripts/pull/34 + + positional arguments: + INPUT_YAML YAML file specifying parameter values for expt runs. Default is Expts_manager.yaml + + options: + -h, --help show this help message and exit + ``` diff --git a/expts_manager/write_Expts_manager_yaml.py b/expts_manager/write_Expts_manager_yaml.py new file mode 100755 index 0000000..43c4d38 --- /dev/null +++ b/expts_manager/write_Expts_manager_yaml.py @@ -0,0 +1,489 @@ +#!/usr/bin/env python3 +try: + from ruamel.yaml import YAML + + ryaml = YAML() + ryaml.preserve_quotes = True + ryaml.indent(mapping=4, sequence=4, offset=2) +except ImportError: + print("\nFatal error: modules not available.") + print("On NCI, do the following and try again:") + print(" module use /g/data/vk83/modules && module load payu/1.1.5\n") + raise + +import os +from io import StringIO + + +def write_config_yaml_file(file_path, description_sections): + """ + Writes the YAML file for Experiment Manager tool. + + Parameters: + - file_path (str): Path to the YAML file to write. + - description_sections (list of tuples): Each tuple contains a section description (str) and its configs (list of dicts). + """ + intro_comment = """ +# ===================================================================================== +# YAML Configuration for Expts_manager.py +# ===================================================================================== +# This configuration file defining the parameters and settings required for cloning, +# setting up, and running control and perturbation experiments using `Expts_manager.py`. +# Detailed explanations are provided to ensure the configuration is straightforward. +# ===================================================================================== +""" + + def dump_block_style(value): + buffer = StringIO() + ryaml.dump(value, buffer) + buffer.seek(0) + return buffer.read() + + def write_value(key, value, file, comment=None, indent=0): + indent_space = " " * (indent * 4) + if isinstance(value, dict): + file.write(f"{indent_space}{key}:\n") + for sub_key, sub_value in value.items(): + write_value(sub_key, sub_value, file, indent=indent + 1) + elif isinstance(value, list): + list_content = ", ".join(map(str, value)) + file.write(f"{indent_space}{key}: [{list_content}]") + else: + file.write(f"{indent_space}{key}: {value}") + if comment: + file.write(f" {comment}") + file.write("\n") + + with open(file_path, "w") as file: + file.write(intro_comment + "\n") + for description, config_list in description_sections: + file.write(f"{description}\n") + for item in config_list: + if isinstance(item, dict): + key = item.get("key", "") + value = item.get("value", "") + comment = item.get("comment", "") + write_value(key, value, file, comment=comment) + + +if __name__ == "__main__": + # Descriptions and configs for Model Selection + descrpt_model_sel = """ +# ============ Model Selection ======================================================== +""" + model_sel = [ + { + "key": "model", + "value": "access-om3", + "comment": "# Specify the model to be used. Options: 'access-om2', 'access-om3'", + }, + ] + + # Descriptions and configs for Utility tool + descrpt_util = """ +# ============ Utility Tool Configuration (only for access-om3) ======================= +# The following configuration provides the necessary tools to: +# 1. Parse parameters and comments from `MOM_input` in MOM6. +# 2. Read and write the `nuopc.runconfig` file. +""" + config_util = [ + { + "key": "force_overwrite_tools", + "value": "False", + }, + { + "key": "utils_url", + "value": "git@github.com:minghangli-uni/om3-utils.git", + "comment": "# Git URL for the utility tool repository", + }, + { + "key": "utils_branch_name", + "value": "main", + "comment": "# The branch for which the utility tool will be checked out", + }, + { + "key": "utils_dir_name", + "value": "om3-utils", + "comment": "# Directory name for the utility tool (user-defined)", + }, + ] + + # Descriptions and configs for Control Experiment Setup + descrpt_diag = """ +# ============ Diagnostic Table (optional) ============================================ +# Configuration for customising the diagnostic table. +""" + config_diag = [ + { + "key": "diag_url", + "value": "git@github.com:minghangli-uni/make_diag_table.git", + "comment": "# Git URL for the `make_diag_table`", + }, + { + "key": "diag_branch_name", + "value": "general_scheme1", + "comment": "# Branch for the `make_diag_table`", + }, + { + "key": "diag_dir_name", + "value": "make_diag_table", + "comment": "# Directory name for the `make_diag_table` (user-defined)", + }, + { + "key": "diag_ctrl", + "value": "False", + "comment": "# Set to 'True' to modify the diagnostic table for the control experiment", + }, + { + "key": "diag_pert", + "value": "False", + "comment": "# Set to 'True' to modify the diagnostic table for perturbation experiments", + }, + ] + + # Descriptions and configs for Control Experiment Setup + descrpt_control = """ +# ============ Control Experiment Setup =============================================== +""" + config_control = [ + { + "key": "base_url", + "value": "git@github.com:ACCESS-NRI/access-om3-configs.git", + "comment": "# Git URL for the control experiment repository", + }, + { + "key": "base_commit", + "value": '"2bc6107"', + "comment": "# Commit hash to use; Please ensure it is a string!", + }, + { + "key": "base_dir_name", + "value": "Ctrl-1deg_jra55do_ryf", + "comment": "# Directory name for cloning (user-defined)", + }, + { + "key": "base_branch_name", + "value": "ctrl", + "comment": "# Branch name for the experiment (user-defined)", + }, + { + "key": "test_path", + "value": "test", + "comment": "# Relative path for all test (control and perturbation) runs (user-defined)", + }, + ] + + # Descriptions and configs for Control Experiment Variables + descrpt_control_expt = """ +# ============ Control Experiment Variables =========================================== +# Allows modification of various control experiment settings. +# 1. config.yaml (access-om2 or access-om3) +# 2. all namelists such as with endswith "_in" or ".nml" etc. (access-om2 or access-om3) +# 3. cpl_dt (coupling timestep) (access-om3) +# 4. nuopc.runconfig (access-om3) +# 5. MOM_input (access-om3) +# Below are some examples for the illustration purpose, please modify for your own settings. +""" + config_control_expt = [ + { + "key": "config.yaml", + "value": { + "ncpus": 240, + "mem": "960GB", + "walltime": "24:00:00", + "#jobname": "# `jobname` will be forced to be the name of the directory, which is `Ctrl-1deg_jra55do_ryf` in this example.", + "metadata": {"enable": True}, + "runlog": True, + "restart_freq": 1, + }, + }, + { + "key": "cpl_dt", + "value": 1800.0, + "comment": "# Coupling timestep in the `nuopc_runseq`", + }, + { + "key": "nuopc.runconfig", + "value": { + "CLOCK_attributes": { + "stop_option": "ndays", + "stop_n": 1, + "restart_option": "ndays", + "restart_n": 1, + }, + "PELAYOUT_attributes": { + "atm_ntasks": 24, + "cpl_ntasks": 24, + "glc_ntasks": 1, + "ice_ntasks": 24, + "lnd_ntasks": 1, + "ocn_ntasks": 216, + "ocn_rootpe": 24, + "rof_ntasks": 24, + "wav_ntasks": 1, + }, + }, + }, + { + "key": "ice_in", + "value": { + "domain_nml": { + "max_blocks": -1, + "block_size_x": 15, + "block_size_y": 300, + }, + }, + }, + { + "key": "input.nml", + "value": { + "diag_manager_nml": { + "max_axes": 400, + "max_files": 200, + "max_num_axis_sets": 200, + }, + }, + }, + { + "key": "MOM_input", + "value": { + "HFREEZE": 12, + "BOUND_SALINITY": False, + }, + }, + ] + + # Descriptions and configs for Namelist Tunning + descrpt_namelists = """ +# ============ Namelist Tunning ================================ +# Tune parameters across different model components. + +# Generalised structure +# 1. Single-parameter tunning within single file +# namelists: +# filename: (Required). +# filename_dirs: List of directory names (Optional - user-defined: filename must be appended with "_dirs"; additional strings may follow). +# groupname: (Required: for f90 namelists or nuopc.runconfig, this is simply the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list"). +# parameter1: list of values +# parameter2: list of values +# ... +# ... + +# 2. Multiple-parameter tuning within a single group in a single file +# namelists: +# filename: (Required). +# filename_dirs: List of directory names (Optional - user-defined: filename must be appended with "_dirs"; additional strings may follow). +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, use the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list", then append "_combo"). +# parameter1: list of values +# parameter2: list of values +# ... +# ... + +# 3. Multiple-parameter tuning across different files using user-defined directories +# namelists: +# cross_block: (Required: additional strings may follow "cross_block"). +# cross_block_dirs: List of directory names (Required - user-defined: filename must be appended with "_dirs"; additional strings may follow). +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, use the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list", then append "_combo"). +# parameter1: list of values +# parameter2: list of values +# ... +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, use the group name; for MOM input, use "MOM_list"; for nuopc.runseq, use "runseq_list", then append "_combo"). +# parameter1: list of values +# parameter2: list of values +# ... +# ... +# +# cross_block: (Required: may append other strings after "cross_block"). +# cross_block_dirs: list of directory names (Optional - user-defined: filename must be appended with "_dirs", you may append other strings after "_dirs"). +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, just the group name; for MOM input, "MOM_list", nuopc.runseq, "runseq_list", then append "_combo" at the end). +# parameter1: list of values +# parameter2: list of values +# ... +# filename: +# groupname_combo: (Required: for f90 namelists or nuopc.runconfig, just the group name; for MOM input, "MOM_list", nuopc.runseq, "runseq_list", then append "_combo" at the end). +# parameter1: list of values +# parameter2: list of values +# ... +# ... +# ... + +# The following namelist options are provided as examples. +# Note: Parameters should be based on the underlying physics. +# Some configurations may fail to run or produce invalid results. +# 1.1 Single parameter tuning within a single file using default perturbation directory names. +# 1.2 Single parameter tuning within a single file using user-defined perturbation directory names. +# 2.1 Multi-parameter tuning within a single group in a single file using default perturbation directory names. +# 2.2 Multi-parameter tuning within a single group in a single file using user-defined perturbation directory names. +# 3. Multi-parameter tuning across multiple files. +""" + config_namelists = [ + { + "key": "namelists", + "value": { + "ice_in": { + "ponds_nml": { + "dpscale": [0.002, 0.003], + "rfracmax": [1.1, 1.2], + }, + "ice_in_dirs_test": ["icet1", "icet2", "icet3", "icet4"], + "shortwave_nml": { + "ahmax": [0.2, 0.3], + "r_snw": [0.1, 0.2], + }, + }, + "MOM_input": { + "MOM_list4": { + "MINIMUM_DEPTH": [0, 0.5], + "MAXIMUM_DEPTH": [6000.0, 7000.0, 8000.0], + }, + "MOM_input_dirs1": ["lexpt0", "lexpt1"], + "MOM_list1_combo": { + "DT_THERM": [3600.0, 108000.0], + "DIABATIC_FIRST": [False, False], + "THERMO_SPANS_COUPLING": [True, True], + }, + }, + "cross_block2": { + "cross_block2_dirs": ["cg_1", "cg_2"], + "nuopc.runconfig": { + "CLOCK_attributes_combo": { + "restart_n": [1, 1], + "restart_option": ["ndays", "ndays"], + "stop_n": [1, 1], + "stop_option": ["ndays", "ndays"], + "ocn_cpl_dt": [1800.0, 7200.0], + }, + "PELAYOUT_attributes_combo": { + "ocn_ntasks": [168, 120], + }, + }, + "config.yaml": { + "config_list1_combo": { + "ncpus": [192, 144], + "mem": ["768GB", "576GB"], + }, + }, + "ice_in": { + "shortwave_nml_combo": { + "albicei": [0.36, 0.39], + "albicev": [0.78, 0.81], + }, + "ponds_nml_combo": { + "dpscale": [0.002, 0.003], + }, + }, + "MOM_input": { + "MOM_list1_combo": { + "DT_THERM": [3600.0, 7200.0], + "DIABATIC_FIRST": [False, False], + "THERMO_SPANS_COUPLING": [True, True], + "DTBT_RESET_PERIOD": [3600.0, 7200.0], + }, + }, + "nuopc.runseq": { + "runseq_list1_combo": { + "cpl_dt": [1800.0, 7200.0], + }, + }, + }, + }, + }, + ] + + # Descriptions and configs for Perturbation Experiment Setup (Optional) + descrpt_perturb_setup = """ +# ============ Perturbation Experiment Setup (Optional - access-om3) ======================= +# Configure settings for perturbation experiments. Currently, only `nuopc.runconfig` is supported. +# If conducting parameter tuning for `nuopc.runconfig`, any pre-existing settings in this section +# will be purged by the above namelist tunning. +""" + config_perturb_setup = [ + { + "key": "perturb_run_config", + "value": { + "CLOCK_attributes": { + "stop_option": "nyears", + "stop_n": 1, + "restart_option": "nyears", + "restart_n": 1, + }, + }, + }, + ] + + descrpt_runs = """ +# ============ Control experiment and perturbation Runs =================================== +# This section configures the settings for running control experiments and their corresponding perturbation tests. + """ + config_runs = [ + { + "key": "ctrl_nruns", + "value": 0, + "comment": "\n# Number of control experiment runs.\ + \n# Default: 0.\ + \n# Adjust this value to set the number of iterations for the control experiment, which serves as the baseline for perturbations.\n", + }, + { + "key": "run_namelists", + "value": True, + "comment": "\n# Determines whether to run using the specified namelists.\ + \n# Default: False.\ + \n# Set to 'True' to apply configurations from the namelist section; otherwise, 'False' to skip this step.\n", + }, + { + "key": "check_duplicate_jobs", + "value": True, + "comment": "\n# Checks if there are duplicate PBS jobs within the same parent directory (`test_path`) based on their names.\ + \n# This check is useful when you have existing running jobs and want to add additional tests, which helps avoid conflicts by ensuring new jobs don't duplicate existing ones in the same `test_path`.\ + \n# The check will not be triggered if the jobs are located in different `test_path`. It only applies to jobs within the same `test_path` directory.\ + \n# Default: True.\ + \n# If duplicates are found, a message will be printed indicating the presence of duplicate runs and those runs will be skipped.\n", + }, + { + "key": "check_skipping", + "value": False, + "comment": "\n# Checks if certain runs should be skipped based on pre-existing conditions or results. Currently only valid for nml type.\ + \n# Default: False.\ + \n# Set to 'True' if you want the system to skip runs under specific criteria; otherwise, keep it 'False'. Currently only valid for nml type.\n", + }, + { + "key": "force_restart", + "value": False, + "comment": "\n# Controls whether to force a restart of the experiments regardless of existing initial conditions.\ + \n# Default: False.\ + \n# Set to 'True' to enforce a restart of the control and perturbation runs.\n", + }, + { + "key": "startfrom", + "value": "'rest'", + "comment": "\n# Defines the starting point for perturbation tests.\ + \n# Options: a specific restart number of the control experiment, or 'rest' to start from the initial state.\ + \n# This parameter determines the initial condition for each perturbation test.\n", + }, + { + "key": "nruns", + "value": 0, + "comment": "\n# Total number of output directories to generate for each Expts_manager member.\ + \n# Default: 0.\ + \n# Specifies how many runs of each perturbation experiment will be started; this number correlates with the different parameter combinations defined in the configuration.", + }, + ] + + yaml_file_path = os.path.join(os.getcwd(), "Expts_manager.yaml") + + write_config_yaml_file( + yaml_file_path, + [ + (descrpt_model_sel, model_sel), + (descrpt_util, config_util), + (descrpt_diag, config_diag), + (descrpt_control, config_control), + (descrpt_control_expt, config_control_expt), + (descrpt_namelists, config_namelists), + (descrpt_perturb_setup, config_perturb_setup), + (descrpt_runs, config_runs), + ], + )