Add parameter tunning for config.yaml and nuopc.runconfig

expts_manager/Expts_manager.py

@@ -469,6 +469,12 @@ def _determine_block_type(self, k):
         elif k == "nuopc.runseq":
             tag_model = "cpl_dt"
             expt_dir_name = k[-6:]
+        elif k == "config.yaml":
+            tag_model = "config"
+            expt_dir_name = k[:6] + "_input"
+        elif k == "nuopc.runconfig":
+            tag_model = "runconfig"
+            expt_dir_name = k[-9:] + "_input"
         else:
             raise ValueError(f"Unsupported block type: {k}")
         return tag_model, expt_dir_name
@@ -490,8 +496,26 @@ def _process_params_group(self, k, k_sub, nmls, expt_dir_name, tag_model):
             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("config_list"):
+            self._process_parameter_group_common(k, k_sub, nmls, 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.startswith("runconfig_input"):
+            self._process_parameter_group_common(k, k_sub, nmls, expt_dir_name)
+
     def _handle_nml_group(self, k, k_sub, expt_dir_name, nmls):
         """
         Handles namelist parameter groups specific to `nml` tag model.
@@ -509,7 +533,11 @@ def _handle_mom6_group(self, k, k_sub, expt_dir_name, nmls):
             )
             commt_dict = MOM_inputParser.commt_dict
             self._process_parameter_group_common(
-                k, k_sub, nmls, expt_dir_name, commt_dict=commt_dict
+                k,
+                k_sub,
+                nmls,
+                expt_dir_name,
+                commt_dict=commt_dict,
             )
 
     def _handle_cpl_dt_group(self, k, k_sub, expt_dir_name, nmls):
@@ -591,7 +619,7 @@ def _generate_individual_dicts(self, name_dict, commt_dict, 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 == "nml":
+        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):
@@ -608,9 +636,51 @@ def _generate_combined_dicts(self, name_dict, commt_dict, 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 == "nml":
+        elif self.tag_model in (("nml", "config", "runconfig")):
             self.append_group_list = append_group_list
 
+    def _generate_combined_dicts2(self, name_dict, commt_dict, k_sub):
+        """
+        Generates a list of dictionaries where each dictionary contains all keys with values from the same index.
+        Handles nested dictionaries with list values to create distinct parameter dictionaries per index.
+        """
+        param_dict_change_list = []
+        append_group_list = []
+
+        # Extract keys from nested dictionaries to determine the number of experiments
+        first_key = next(iter(name_dict))
+        first_subkey = next(iter(name_dict[first_key]))
+
+        # Determine the number of experiments from the length of the first subkey list
+        num_expts = len(name_dict[first_key][first_subkey])
+
+        for i in range(num_expts):
+            # Construct a dictionary for the current experiment
+            param_dict_change = {}
+            for k, v in name_dict.items():
+                # Check if the value is a nested dictionary
+                if isinstance(v, dict):
+                    # Construct a new dictionary for this nested dictionary
+                    nested_dict = {
+                        sub_key: (
+                            sub_value[i]
+                            if isinstance(sub_value, list) and i < len(sub_value)
+                            else sub_value
+                        )
+                        for sub_key, sub_value in v.items()
+                    }
+                    param_dict_change[k] = nested_dict
+                else:
+                    # For non-dictionary values, assign directly
+                    param_dict_change[k] = v
+
+            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
+        self.append_group_list = append_group_list
+
     def setup_expts(self, parameter_block):
         """
         Sets up perturbation experiments based on the YAML input file provided in `Expts_manager.yaml`.
@@ -633,10 +703,14 @@ def setup_expts(self, parameter_block):
             # update params for each parameter block
             if self.tag_model == "mom6":
                 self._update_mom6_params(expt_path, param_dict)
-            elif self.tag_model == "nml":
+            elif self.tag_model in "nml":
                 self._update_nml_params(expt_path, param_dict, parameter_block, i)
             elif self.tag_model == "cpl_dt":
                 self._update_cpl_dt_params(expt_path, param_dict)
+            elif self.tag_model == "config":
+                self._update_config_params(expt_path, param_dict, parameter_block, i)
+            elif self.tag_model == "runconfig":
+                self._update_runconfig_params(expt_path, param_dict, parameter_block, i)
 
             # optionally update diag_table for perturbation runs
             if self.diag_pert and self.diag_path:
@@ -750,6 +824,50 @@ def _update_nml_params(self, expt_path, param_dict, parameter_block, indx):
         f90nml.patch(nml_path, patch_dict, nml_path + "_tmp")
         os.rename(nml_path + "_tmp", nml_path)
 
+    def _update_config_params(self, expt_path, param_dict, parameter_block, indx):
+        """
+        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)
+        nml_group = self.append_group_list[indx]
+
+        if nml_group.endswith(self.combo_suffix):
+            nml_group = nml_group[
+                : -len(self.combo_suffix)
+            ]  # rename the namlist by removing the suffix if the suffix with `_combo`
+        file_read = self._read_ryaml(nml_path)
+        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):
+        """
+        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)
+        nml_group = self.append_group_list[indx]
+
+        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):
         """
         Updates coupling timestep parameters.
@@ -770,7 +888,7 @@ def _generate_restart_symlink(self, expt_path):
         Args:
             expt_path (str): The path to the experiment directory.
         """
-        if self.startfrom_str != 'rest':
+        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))
@@ -787,7 +905,7 @@ def _generate_restart_symlink(self, expt_path):
                     os.remove(dest)  # remove symlink
                 os.symlink(restartpath, dest)  # generate symlink
         else:
-            restartpath = 'rest'
+            restartpath = "rest"
 
         return restartpath
 

expts_manager/Expts_manager.yaml

@@ -80,11 +80,11 @@ MOM_input:
 # ============ Perturbation Experiment Setup (Optional) =======================
 # Configure perturbation experiments, currently supports `nuopc.runconfig` only.
 perturb_run_config:
-    CLOCK_attributes:
-        stop_option: nyears
-        stop_n: 1
-        restart_option: nyears
-        restart_n: 1
+   CLOCK_attributes:
+       stop_option: nyears
+       stop_n: 1
+       restart_option: nyears
+       restart_n: 1
 
 # ============ Namelist Tunning ================================
 # Allows fine-tuning of various parameters across different model components 
@@ -198,6 +198,23 @@ perturb_run_config:
 # So, some configurations may not run successfully or produce valid results.
 
 namelists:
+    config.yaml:
+        config_list1:
+            ncpus: [480, 960]
+            mem: [1800GB, 3600GB]
+        config_list2:
+            ncpus: [1080, 2160]
+            mem: [1800GB, 3600GB]
+        config_input_dir3: [config_input1, config_input2]
+        config_list3_combo:
+            ncpus: [1080, 2160]
+            mem: [1800GB, 3600GB]
+    nuopc.runconfig:
+        runconfig_input_dir1: [nuopcrunconfig_input1, nuopcrunconfig_input2]
+        PELAYOUT_attributes_combo:
+                atm_ntasks: [24,48]
+                cpl_ntasks: [24,48]
+                glc_ntasks: [1,1]
     ice_in:
         shortwave_nml_combo:
             albicei: [ 0.36, 0.39, 0.47 ]
@@ -237,10 +254,10 @@ namelists:
         MOM_list3:
             HFREEZE: [12,14]
 
-        MOM_input_dir10: ['min_depth0', 'min_depth0_5', 'max_depth_6000m', 'max_depth_7000m']
+        MOM_input_dir10: ['min_depth0', 'min_depth0_5', 'max_depth_6000m', 'max_depth_7000m', 'max_depth_8000m']
         MOM_list4:
             MINIMUM_DEPTH: [0,0.5]
-            MAXIMUM_DEPTH: [6000.0, 7000.0]
+            MAXIMUM_DEPTH: [6000.0, 7000.0, 8000.0]
 
         MOM_input_dir25: ['timestep_therm_test']
         MOM_list5_combo: