diff --git a/visualCaseGen/custom_widget_types/case_creator.py b/visualCaseGen/custom_widget_types/case_creator.py index c4e589b..eefc112 100644 --- a/visualCaseGen/custom_widget_types/case_creator.py +++ b/visualCaseGen/custom_widget_types/case_creator.py @@ -11,7 +11,7 @@ from ProConPy.config_var import cvars from visualCaseGen.custom_widget_types.mom6_bathy_launcher import MOM6BathyLauncher from visualCaseGen.custom_widget_types.dummy_output import DummyOutput -from visualCaseGen.custom_widget_types.case_tools import xmlchange, run_case_setup, append_user_nl +from visualCaseGen.custom_widget_types.case_tools import xmlchange, run_case_setup, append_user_nl, is_ccs_config_writeable COMMENT = "\033[01;96m" # bold, cyan SUCCESS = "\033[1;32m" # bold, green @@ -23,7 +23,7 @@ class CaseCreator: """The base class for CaseCreatorWidget. Here, backend functionalities are implemented.""" - def __init__(self, cime, output=None, allow_xml_override=False): + def __init__(self, cime, output=None, allow_xml_override=False, add_grids_to_ccs_config = True): """Initialize CaseCreator object. Parameters @@ -40,23 +40,25 @@ def __init__(self, cime, output=None, allow_xml_override=False): self._cime = cime self._out = DummyOutput() if output is None else output self._allow_xml_override = allow_xml_override + self._add_grids_to_ccs_config = add_grids_to_ccs_config # By default, visualCaseGen assigns grids through ccs_config, if not possible (which can happen if the user does not own the sandbox), it is possible to apply grid changes through xmlchanges instead. + assert is_ccs_config_writeable or not add_grids_to_ccs_config, "Cannot write to ccs_config xml files. Please set add_grids_to_ccs_config to False to apply grid changes through xmlchanges." def revert_launch(self, do_exec=True): """This function is called when the case creation fails. It reverts the changes made to the ccs_config xml files.""" - - mg = "ccs_config/modelgrid_aliases_nuopc.xml" - if (Path(self._cime.srcroot) / f"{mg}.orig").exists(): - shutil.move( - Path(self._cime.srcroot) / f"{mg}.orig", - Path(self._cime.srcroot) / f"{mg}" - ) - cg = "ccs_config/component_grids_nuopc.xml" - if (Path(self._cime.srcroot) / f"{cg}.orig").exists(): - shutil.move( - Path(self._cime.srcroot) / f"{cg}.orig", - Path(self._cime.srcroot) / f"{cg}" - ) + if self._add_grids_to_ccs_config: + mg = "ccs_config/modelgrid_aliases_nuopc.xml" + if (Path(self._cime.srcroot) / f"{mg}.orig").exists(): + shutil.move( + Path(self._cime.srcroot) / f"{mg}.orig", + Path(self._cime.srcroot) / f"{mg}" + ) + cg = "ccs_config/component_grids_nuopc.xml" + if (Path(self._cime.srcroot) / f"{cg}.orig").exists(): + shutil.move( + Path(self._cime.srcroot) / f"{cg}.orig", + Path(self._cime.srcroot) / f"{cg}" + ) def _remove_orig_xml_files(self): """This function is called when the case creation and modification process is successful. @@ -130,7 +132,10 @@ def create_case(self, do_exec): if cvars["GRID_MODE"].value == "Standard": resolution = cvars["GRID"].value elif cvars["GRID_MODE"].value == "Custom": - resolution = Path(cvars["CUSTOM_GRID_PATH"].value).name + if self._add_grids_to_ccs_config: + resolution = Path(cvars["CUSTOM_GRID_PATH"].value).name + else: + resolution = "USER_RES" # Set to a default visualCaseGen Resolution since grids are changed through xml changes else: raise RuntimeError(f"Unknown grid mode: {cvars['GRID_MODE'].value}") @@ -139,11 +144,16 @@ def create_case(self, do_exec): print(f"{COMMENT}Creating case...{RESET}\n") # First, update ccs_config xml files to add custom grid information if needed: - self._update_ccs_config(do_exec) + if self._add_grids_to_ccs_config: + self._update_ccs_config(do_exec) # Run create_newcase self._run_create_newcase(caseroot, compset, resolution, do_exec) + # If we don't pick the grids through ccs_config, use xml changes + if not self._add_grids_to_ccs_config: + self._update_grids_via_xmlchange(do_exec) + # Navigate to the case directory: with self._out: print(f"{COMMENT}Navigating to the case directory:{RESET}\n") @@ -160,7 +170,8 @@ def create_case(self, do_exec): # Clean up: if do_exec: - self._remove_orig_xml_files() + if self._add_grids_to_ccs_config: + self._remove_orig_xml_files() cvars["CASE_CREATOR_STATUS"].value = "OK" with self._out: caseroot = cvars["CASEROOT"].value @@ -238,7 +249,7 @@ def _update_modelgrid_aliases(self, custom_grid_path, ocn_grid, do_exec): modelgrid_aliases_xml = modelgrid_aliases_xml.as_posix() # confirm that modelgrid_aliases xml file is writeable: - if not os.access(modelgrid_aliases_xml, os.W_OK): + if not is_ccs_config_writeable(self._cime): raise RuntimeError(f"Cannot write to {modelgrid_aliases_xml}.") # Construct the component grids string to be logged: @@ -530,7 +541,7 @@ def _apply_lnd_grid_xmlchanges(self, do_exec): """Apply xmlchanges related to custom land grid if needed.""" lnd_grid_mode = cvars["LND_GRID_MODE"].value - if lnd_grid_mode == "Modified": + if self._add_grids_to_ccs_config and lnd_grid_mode == "Modified": if cvars["COMP_OCN"].value != "mom": with self._out: print(f"{COMMENT}Apply custom land grid xml changes:{RESET}\n") @@ -793,3 +804,117 @@ def _apply_clm_namelist_changes(self, do_exec): ]) self._apply_user_nl_changes("clm", user_nl_clm_changes, do_exec) + + def _update_grids_via_xmlchange(self, do_exec): + """Update the case with custom grid information if needed. + This function is called after running create_newcase.""" + + if cvars["GRID_MODE"].value == "Standard": + return + else: + assert ( + cvars["GRID_MODE"].value == "Custom" + ), f"Unknown grid mode: {cvars['GRID_MODE'].value}" + + # check if custom grid path exists: + ocn_grid_mode = cvars["OCN_GRID_MODE"].value + lnd_grid_mode = cvars["LND_GRID_MODE"].value + custom_grid_path = Path(cvars["CUSTOM_GRID_PATH"].value) + if not custom_grid_path.exists(): + if ocn_grid_mode != "Standard" or lnd_grid_mode != "Standard": + raise RuntimeError(f"Custom grid path {custom_grid_path} does not exist.") + + ocn_grid = None + if ocn_grid_mode == "Standard": + ocn_grid = cvars["CUSTOM_OCN_GRID"].value + elif ocn_grid_mode in ["Modify Existing", "Create New"]: + ocn_grid = cvars["CUSTOM_OCN_GRID_NAME"].value + else: + raise RuntimeError(f"Unknown ocean grid mode: {ocn_grid_mode}") + if ocn_grid is None: + raise RuntimeError("No ocean grid specified.") + + self._update_component_grids_xml(custom_grid_path, ocn_grid, ocn_grid_mode, do_exec) + + + def _update_component_grids_xml( + self, custom_grid_path, ocn_grid, ocn_grid_mode, do_exec + ): + """Update the component_grids xml file with custom ocnice grid information. + This function is called before running create_newcase. + + Parameters + ---------- + custom_grid_path : Path + The path to the custom grid directory. + ocn_grid : str + The name of the custom ocean grid. + ocn_grid_mode : str + The ocean grid mode. It can be "Standard", "Modify Existing", or "Create New". + do_exec : bool + If True, execute the commands. If False, only print them. + """ + + if ocn_grid_mode == "Create New": + ocn_dir = custom_grid_path / "ocnice" + assert ocn_dir.exists(), f"Ocean grid directory {ocn_dir} does not exist." + + ocn_mesh = ( + ocn_dir / f"ESMF_mesh_{ocn_grid}_{cvars['MB_ATTEMPT_ID'].value}.nc" + ) + assert ocn_mesh.exists(), f"Ocean mesh file {ocn_mesh} does not exist." + + + # log the modification of component_grids.xml: + with self._out: + print( + f'{BPOINT} Updating case xml variables to include ' + f'newly generated ocean grid "{ocn_grid}" with the following properties:\n' + f' nx: {cvars["OCN_NX"].value}, ny: {cvars["OCN_NY"].value}.' + f' ocean mesh: {ocn_mesh}.{RESET}\n' + ) + + xmlchange("OCN_NX", cvars["OCN_NX"].value, do_exec, self._is_non_local(), self._out) + + xmlchange("OCN_NY", cvars["OCN_NY"].value, do_exec, self._is_non_local(), self._out) + + xmlchange("OCN_DOMAIN_MESH", ocn_mesh.as_posix(), do_exec, self._is_non_local(), self._out) + + xmlchange("ICE_DOMAIN_MESH", ocn_mesh.as_posix(), do_exec, self._is_non_local(), self._out) + + xmlchange("MASK_MESH", ocn_mesh.as_posix(), do_exec, self._is_non_local(), self._out) + + xmlchange("ATM_GRID", cvars["CUSTOM_ATM_GRID"].value, do_exec, self._is_non_local(), self._out) + + xmlchange("LND_GRID", cvars["CUSTOM_LND_GRID"].value, do_exec, self._is_non_local(), self._out) + + xmlchange("ATM_DOMAIN_MESH", self._cime.get_mesh_path("atm",cvars["CUSTOM_ATM_GRID"].value), do_exec, self._is_non_local(), self._out) + + xmlchange("LND_DOMAIN_MESH", self._cime.get_mesh_path("lnd",cvars["CUSTOM_LND_GRID"].value), do_exec, self._is_non_local(), self._out) + + if cvars["CUSTOM_ROF_GRID"].value is not None and cvars["CUSTOM_ROF_GRID"].value != "" and cvars["CUSTOM_ROF_GRID"].value != "null": + xmlchange("ROF_GRID", cvars["CUSTOM_ROF_GRID"].value, do_exec, self._is_non_local(), self._out) + xmlchange("ROF_DOMAIN_MESH", self._cime.get_mesh_path("rof",cvars["CUSTOM_ROF_GRID"].value), do_exec, self._is_non_local(), self._out) + + + lnd_grid_mode = cvars["LND_GRID_MODE"].value + if lnd_grid_mode == "Modified": + if cvars["COMP_OCN"].value != "mom": + with self._out: + print(f"{COMMENT}Apply custom land grid xml changes:{RESET}\n") + + # TODO: NO LONGER RELEVANT - OCEAN GRIDS ARE DONE THROUGH XML CHANGES AS WELL: instead of xmlchanges, these changes should be made via adding the new lnd domain mesh to + # component_grids_nuopc.xml and modelgrid_aliases_nuopc.xml (just like how we handle new ocean grids) + + # lnd domain mesh + xmlchange("LND_DOMAIN_MESH", cvars["INPUT_MASK_MESH"].value, do_exec, self._is_non_local(), self._out) + + # mask mesh (if modified) + base_lnd_grid = cvars["CUSTOM_LND_GRID"].value + custom_grid_path = Path(cvars["CUSTOM_GRID_PATH"].value) + lnd_dir = custom_grid_path / "lnd" + modified_mask_mesh = lnd_dir / f"{base_lnd_grid}_mesh_mask_modifier.nc" # TODO: the way we get this filename is fragile + assert modified_mask_mesh.exists(), f"Modified mask mesh file {modified_mask_mesh} does not exist." + xmlchange("MASK_MESH", modified_mask_mesh, do_exec, self._is_non_local(), self._out) + else: + assert lnd_grid_mode in [None, "", "Standard"], f"Unknown land grid mode: {lnd_grid_mode}" diff --git a/visualCaseGen/custom_widget_types/case_tools.py b/visualCaseGen/custom_widget_types/case_tools.py index 433b88f..cb5bfdb 100644 --- a/visualCaseGen/custom_widget_types/case_tools.py +++ b/visualCaseGen/custom_widget_types/case_tools.py @@ -1,5 +1,6 @@ from pathlib import Path import subprocess +import os from ProConPy.config_var import cvars from visualCaseGen.custom_widget_types.dummy_output import DummyOutput @@ -8,6 +9,15 @@ RESET = "\033[0m" +def is_ccs_config_writeable(cime): + srcroot = cime.srcroot + ccs_config_root = Path(srcroot) / "ccs_config" + assert ( + ccs_config_root.exists() + ), f"ccs_config_root {ccs_config_root} does not exist." + modelgrid_aliases_xml = ccs_config_root / "modelgrid_aliases_nuopc.xml" + return os.access(modelgrid_aliases_xml, os.W_OK) + def run_case_setup(do_exec, is_non_local=False, out=None): """Run the case.setup script to set up the case instance.