diff --git a/.github/workflows/ci-tests.yml b/.github/workflows/ci-tests.yml index 081b673..ca9f210 100644 --- a/.github/workflows/ci-tests.yml +++ b/.github/workflows/ci-tests.yml @@ -29,7 +29,7 @@ jobs: uses: actions/checkout@v4 with: repository: alperaltuntas/CESM - ref: cesm3_0_beta03_gui + ref: cesm3_0_beta06_gui path: CESM #submodules: recursive diff --git a/CONTRIBUTING.md b/CONTRIBUTING.md new file mode 100644 index 0000000..649e1ea --- /dev/null +++ b/CONTRIBUTING.md @@ -0,0 +1,37 @@ +# Contributing to visualCaseGen + +Thank you for your interest in contributing to visualCaseGen! + +We welcome contributions that fix bugs, introduce new features, and enhance the usability, robustness, or performance of the tool. Here's how you can get involved: + +## Getting Started + +Fork the repository and clone your fork. + +Follow the instructions in the visualCaseGen documentation to set up your local visualCaseGen environment: https://esmci.github.io/visualCaseGen/ + +Create a new branch for your feature or bugfix. + +## What You Can Contribute + +- Bug fixes and issue resolutions +- Enhancements to the GUI or core functionality +- Improved documentation or examples +- Code cleanup and test coverage + +## Guidelines + +- Follow the existing code style and structure. +- Include clear commit messages. +- Test your changes before submitting. +- Add tests if you introduce new features. +- Open a Pull Request with a short summary of your changes. +- If addressing an open issue, mention it in your PR (e.g., Closes #42). + +## Reporting Issues + +Please use the GitHub Issues page to report bugs or suggest features. When reporting: + +- Include the version of visualCaseGen you used and details about your system environment. +- Provide clear steps to reproduce the issue. +- Include screenshots or logs if helpful. diff --git a/ProConPy/config_var.py b/ProConPy/config_var.py index 67653bf..01a12a2 100644 --- a/ProConPy/config_var.py +++ b/ProConPy/config_var.py @@ -334,6 +334,11 @@ def update_options_validities(self): return validities_changed + @property + def valid_options(self): + """Returns the list of valid options for this variable.""" + return [opt for opt in self._options if self._options_validities.get(opt, False)] + def _refresh_widget_options(self): """Refresh the widget options list based on information in the current self._options_validities.""" diff --git a/ProConPy/stage.py b/ProConPy/stage.py index 5f7f52b..6216525 100644 --- a/ProConPy/stage.py +++ b/ProConPy/stage.py @@ -352,7 +352,7 @@ def _proceed(self): return # Display the child stage and its siblings by appending them to the current stage's widget - if self.has_children(): + if self.has_children() and next_stage.is_descendant_of(self): self._widget.add_child_stages(first_child=next_stage) # Proceed the csp solver before enabling the next stage @@ -378,20 +378,24 @@ def get_next(self, full_dfs=False): The next stage to visit, if found. Otherwise, None. """ - if self.has_children(): - return self._get_child_to_enable(full_dfs) - elif self._right is not None: + # First try to get a child stage to enable + if (child_to_enable := self._get_child_to_enable(full_dfs)) is not None: + return child_to_enable + + # No child stage to enable. Try to get the right sibling. + if self._right is not None: return self._right - else: # Backtrack - ancestor = self._parent - while ancestor is not None: - if ancestor._right is not None and ( - full_dfs or not ancestor.has_condition() - ): - return ancestor._right - else: - ancestor = ancestor._parent - return None + + # No child or right sibling. Backtrack to find the next stage. + ancestor = self._parent + while ancestor is not None: + if ancestor._right is not None and ( + full_dfs or not ancestor.has_condition() + ): + return ancestor._right + else: + ancestor = ancestor._parent + return None def _get_child_to_enable(self, full_dfs): """Determine the child stage to activate. @@ -401,6 +405,9 @@ def _get_child_to_enable(self, full_dfs): full_dfs : bool If True, visit all the stages in the stage tree. Otherwise, skip stages whose guards are not satisfied.""" + + if self.has_children() is False: + return None child_to_activate = None @@ -412,6 +419,11 @@ def _get_child_to_enable(self, full_dfs): child_to_activate is None ), "Only one child stage can be activated at a time." child_to_activate = child + + if child_to_activate is None: + # No child guard's condition is satisfied. + # Let the caller handle this case (by backtracking). + return None else: # If children are not guards, the first child is activated. # Note the remaining children will be activated in sequence by their siblings. @@ -419,11 +431,7 @@ def _get_child_to_enable(self, full_dfs): # If the child to activate is a Guard, return it's first child if child_to_activate.has_condition(): - return child_to_activate._children[0] - - assert ( - child_to_activate is not None - ), "At least one child stage must be activated." + child_to_activate = child_to_activate._children[0] return child_to_activate diff --git a/environment.yml b/environment.yml index a310f6a..dcf347e 100644 --- a/environment.yml +++ b/environment.yml @@ -6,7 +6,8 @@ channels: dependencies: - python>=3.11.10,<3.12 - #- xesmf + - libxml2>=2.13,<2.14 + - xesmf>=0.8.10,<0.9 - pip - pip: - -e ./external/mom6_bathy/ diff --git a/external/ipyfilechooser b/external/ipyfilechooser index 4c4f94a..67edea7 160000 --- a/external/ipyfilechooser +++ b/external/ipyfilechooser @@ -1 +1 @@ -Subproject commit 4c4f94a8d99c4cd63cb039b75b128f6c550051ad +Subproject commit 67edea7fd41e4c8c31101a27fdcaa924dd2b8a21 diff --git a/external/mom6_bathy b/external/mom6_bathy index c8156d7..aaaca3f 160000 --- a/external/mom6_bathy +++ b/external/mom6_bathy @@ -1 +1 @@ -Subproject commit c8156d7852a970e94769e3549ac97ec168cddd33 +Subproject commit aaaca3f690b41acb6eb4bc0d27d0d4ed289a4031 diff --git a/pyproject.toml b/pyproject.toml index 48ff9d7..4d0c3f9 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,6 +1,6 @@ [project] name = "visualCaseGen" -version = "0.1.1" +version = "0.1.3" authors = [ { name = "Alper Altuntas, NCAR", email = "altuntas@ucar.edu" } ] @@ -9,7 +9,7 @@ readme = "README.md" license = { file = "LICENSE.md" } classifiers = [ "Intended Audience :: Science/Research", - "License :: OSI Approved :: LGPL", + "License :: OSI Approved :: Apache 2.0 License", "Programming Language :: Python", "Framework :: Jupyter" ] @@ -35,4 +35,4 @@ requires = ["setuptools>=42", "wheel"] build-backend = "setuptools.build_meta" [tool.setuptools] -packages = ["ProConPy", "visualCaseGen"] \ No newline at end of file +packages = ["ProConPy", "visualCaseGen"] diff --git a/tests/1_unit/test_cores_case_creator.py b/tests/1_unit/test_cores_case_creator.py index df3425e..8bbb71d 100644 --- a/tests/1_unit/test_cores_case_creator.py +++ b/tests/1_unit/test_cores_case_creator.py @@ -17,11 +17,11 @@ def test_calc_cores_based_on_grid_cases(): assert CaseCreator._calc_cores_based_on_grid(33) == 1 # Test ideal cores amount - assert CaseCreator._calc_cores_based_on_grid(800*128) == 128 + assert CaseCreator._calc_cores_based_on_grid(300*128) == 128 - assert CaseCreator._calc_cores_based_on_grid(800*32) == 128 + assert CaseCreator._calc_cores_based_on_grid(300*32) == 128 - assert CaseCreator._calc_cores_based_on_grid(740 * 780) == 768 + assert CaseCreator._calc_cores_based_on_grid(740 * 780) == 2048 \ No newline at end of file diff --git a/tests/3_system/test_custom_compset_std_grid.py b/tests/3_system/test_custom_compset_std_grid.py new file mode 100755 index 0000000..e591546 --- /dev/null +++ b/tests/3_system/test_custom_compset_std_grid.py @@ -0,0 +1,107 @@ +#!/usr/bin/env python3 + +import pytest +import shutil +import os +from pathlib import Path +import tempfile +import time + +from ProConPy.config_var import ConfigVar, cvars +from ProConPy.stage import Stage +from ProConPy.csp_solver import csp +from visualCaseGen.cime_interface import CIME_interface +from visualCaseGen.initialize_configvars import initialize_configvars +from visualCaseGen.initialize_widgets import initialize_widgets +from visualCaseGen.initialize_stages import initialize_stages +from visualCaseGen.specs.options import set_options +from visualCaseGen.specs.relational_constraints import get_relational_constraints +from visualCaseGen.custom_widget_types.mom6_bathy_launcher import MOM6BathyLauncher +from visualCaseGen.custom_widget_types.case_creator_widget import CaseCreatorWidget +from tests.utils import safe_create_case + + +# do not show logger output +import logging + +logger = logging.getLogger() +logger.setLevel(logging.CRITICAL) + +base_temp_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "temp")) + + +def test_custom_compset_std_grid(): + """Configure a custom compset with a standard grid: 2000_DATM%JRA_SLND_CICE%PRES_MOM6_SROF_SGLC_WW3. Progress through the stages + until the launch stage is reached.""" + + ConfigVar.reboot() + Stage.reboot() + cime = CIME_interface() + initialize_configvars(cime) + initialize_widgets(cime) + initialize_stages(cime) + set_options(cime) + csp.initialize(cvars, get_relational_constraints(cvars), Stage.first()) + + # At initialization, the first stage should be enabled + assert Stage.first().enabled + cvars['COMPSET_MODE'].value = 'Custom' + + # CCOMPSET_MODE is the only variable in the first stage, so assigning a value to it should disable the first stage + assert not Stage.first().enabled + + # The next stge is Custom Component Set, whose first child is Model Time Period + assert Stage.active().title.startswith('Time Period') + cvars['INITTIME'].value = '1850' + + cvars['COMP_OCN'].value = "mom" + cvars['COMP_ICE'].value = "sice" + cvars['COMP_ATM'].value = "cam" + cvars['COMP_ROF'].value = "mosart" + cvars['COMP_LND'].value = "clm" + cvars['COMP_WAV'].value = "swav" + cvars['COMP_GLC'].value = "sglc" + + + assert Stage.active().title.startswith('Component Physics') + + cvars['COMP_ATM_PHYS'].value = "CAM60" + cvars['COMP_LND_PHYS'].value = "CLM60" + + assert Stage.active().title.startswith('Component Options') + + cvars['COMP_ATM_OPTION'].value = "1PCT" + cvars['COMP_LND_OPTION'].value = "BGC" + cvars['COMP_OCN_OPTION'].value = "(none)" + cvars['COMP_ICE_OPTION'].value = "(none)" + cvars['COMP_ROF_OPTION'].value = "(none)" + + assert Stage.active().title.startswith('2. Grid') + + cvars['GRID_MODE'].value = 'Standard' + cvars['GRID'].value = 'f09_t232' + + assert Stage.active().title.startswith('3. Launch') + launch_stage = Stage.active() + + with tempfile.TemporaryDirectory(dir=base_temp_dir) as temp_case_path: + pass # immediately remove the random temporary directory, + # which will become the caseroot directory below + + cvars["CASEROOT"].value = temp_case_path + + case_creator = launch_stage._widget._main_body.children[-1] + assert isinstance(case_creator, CaseCreatorWidget) + + cvars["PROJECT"].value = "12345" + + # *Click* the create_case button + safe_create_case(cime.srcroot, case_creator) + + # sleep for a bit to allow the case to be created + time.sleep(5) + + # remove the caseroot directory + shutil.rmtree(temp_case_path) + + diff --git a/tests/3_system/test_custom_mom6_grid.py b/tests/3_system/test_custom_mom6_grid.py index c8c5457..6e81201 100755 --- a/tests/3_system/test_custom_mom6_grid.py +++ b/tests/3_system/test_custom_mom6_grid.py @@ -141,7 +141,10 @@ def test_custom_mom6_grid(): assert Stage.active().title.startswith("Simple Initial Conditions") cvars["T_REF"].value = 10.0 - # Since land grid gets set automatically, we should be in the Launch stage: + # Since land grid and runoff grid get set automatically, we should be in the runoff to ocn mapping: + assert Stage.active().title.startswith("Runoff to Ocean Mapping") + cvars["ROF_OCN_MAPPING_STATUS"].value = "skip" + assert Stage.active().title.startswith("3. Launch") launch_stage = Stage.active() diff --git a/tests/3_system/test_f2000_custom_grid.py b/tests/3_system/test_f2000_custom_grid.py index 4163006..7ef9124 100755 --- a/tests/3_system/test_f2000_custom_grid.py +++ b/tests/3_system/test_f2000_custom_grid.py @@ -159,6 +159,9 @@ def construct_custom_res_from_std_grids(cime): assert Stage.active().title.startswith("Land Grid") cvars["CUSTOM_LND_GRID"].value = "0.9x1.25" + assert Stage.active().title.startswith("Runoff Grid") + cvars["CUSTOM_ROF_GRID"].value = "r05" + assert Stage.active().title.startswith("3. Launch") launch_stage = Stage.active() @@ -248,6 +251,9 @@ def construct_custom_res_from_modified_clm_grid(cime): # click the "Run Surface Data Modifier" button fsurdat_modifier_launcher._on_launch_clicked(b=None) + assert Stage.active().title.startswith("Runoff Grid") + cvars["CUSTOM_ROF_GRID"].value = "r05" + assert Stage.active().title.startswith("3. Launch") launch_stage = Stage.active() @@ -364,6 +370,9 @@ def construct_custom_res_from_new_mom6_grid_modified_clm_grid(cime): # click the "Run Surface Data Modifier" button fsurdat_modifier_launcher._on_launch_clicked(b=None) + assert Stage.active().title.startswith("Runoff Grid") + cvars["CUSTOM_ROF_GRID"].value = "r05" + assert Stage.active().title.startswith("3. Launch") launch_stage = Stage.active() diff --git a/tests/3_system/test_fhist_custom_grid.py b/tests/3_system/test_fhist_custom_grid.py index d11c4b3..f3a9a99 100755 --- a/tests/3_system/test_fhist_custom_grid.py +++ b/tests/3_system/test_fhist_custom_grid.py @@ -129,6 +129,9 @@ def construct_custom_res_from_std_grids(cime): assert Stage.active().title.startswith("Land Grid") cvars["CUSTOM_LND_GRID"].value = "0.9x1.25" + assert Stage.active().title.startswith("Runoff Grid") + cvars["CUSTOM_ROF_GRID"].value = "r05" + assert Stage.active().title.startswith("3. Launch") launch_stage = Stage.active() @@ -218,6 +221,9 @@ def construct_custom_res_from_modified_clm_grid(cime): # click the "Run Surface Data Modifier" button fsurdat_modifier_launcher._on_launch_clicked(b=None) + assert Stage.active().title.startswith("Runoff Grid") + cvars["CUSTOM_ROF_GRID"].value = "r05" + assert Stage.active().title.startswith("3. Launch") launch_stage = Stage.active() diff --git a/tests/3_system/test_rof_ocn_map.py b/tests/3_system/test_rof_ocn_map.py new file mode 100755 index 0000000..099ce0e --- /dev/null +++ b/tests/3_system/test_rof_ocn_map.py @@ -0,0 +1,159 @@ +#!/usr/bin/env python3 + +import pytest +import os +import tempfile + +from ProConPy.config_var import ConfigVar, cvars +from ProConPy.stage import Stage +from ProConPy.csp_solver import csp +from visualCaseGen.cime_interface import CIME_interface +from visualCaseGen.initialize_configvars import initialize_configvars +from visualCaseGen.initialize_widgets import initialize_widgets +from visualCaseGen.initialize_stages import initialize_stages +from visualCaseGen.specs.options import set_options +from visualCaseGen.specs.relational_constraints import get_relational_constraints + +# do not show logger output +import logging + +logger = logging.getLogger() +logger.setLevel(logging.CRITICAL) + +base_temp_dir = os.path.abspath(os.path.join(os.path.dirname(__file__), "..", "temp")) + +def configure_for_rof_map(compset_alias): + """This function configures a standard compset with custom grids suitable for testing + the runoff to ocean mapping generator.""" + + ConfigVar.reboot() + Stage.reboot() + cime = CIME_interface() + initialize_configvars(cime) + initialize_widgets(cime) + initialize_stages(cime) + set_options(cime) + csp.initialize(cvars, get_relational_constraints(cvars), Stage.first()) + + assert os.path.exists(base_temp_dir), "temp testing directory does not exist" + + # At initialization, the first stage should be enabled + assert Stage.first().enabled + cvars["COMPSET_MODE"].value = "Standard" + + # COMPSET_MODE is the only variable in the first stage, so assigning a value to it should disable the first stage + assert not Stage.first().enabled + + # The next stage is Custom Component Set, whose first child is Model Time Period + assert Stage.active().title.startswith("Support Level") + cvars["SUPPORT_LEVEL"].value = "All" + + # Apply filters + for comp_class in cime.comp_classes: + cvars[f"COMP_{comp_class}_FILTER"].value = "any" + + ## Pick a standard compset + cvars["COMPSET_ALIAS"].value = compset_alias + + # Create a custom grid + assert Stage.active().title.startswith("2. Grid") + cvars["GRID_MODE"].value = "Custom" + + # Set the custom grid path + assert Stage.active().title.startswith("Custom Grid") + +def test_standard_rof_to_ocn_mapping(): + """This test configures a case with a standard runoff to ocean mapping for a custom resolution""" + + configure_for_rof_map("C_JRA") + + assert Stage.active().title.startswith("Custom Grid") + + with tempfile.TemporaryDirectory(dir=base_temp_dir) as temp_grid_path: + + cvars["CUSTOM_GRID_PATH"].value = temp_grid_path + # since this is a JRA run, the atmosphere grid must automatically be set to TL319 + assert cvars["CUSTOM_ATM_GRID"].value == "TL319" + + # Set the custom ocean grid mode + assert Stage.active().title.startswith("Ocean") + cvars["OCN_GRID_MODE"].value = "Standard" + + assert Stage.active().title.startswith("Ocean Grid") + cvars["CUSTOM_OCN_GRID"].value = "tx2_3v2" + + # Land Grid and Runoff grid should be set automatically + assert Stage.active().title.startswith("Runoff to Ocean Mapping") + + # Currently, smoothing parameters should be unset + assert cvars["ROF_OCN_MAPPING_RMAX"].value is None, "ROF_OCN_MAPPING_RMAX should be None" + assert cvars["ROF_OCN_MAPPING_FOLD"].value is None, "ROF_OCN_MAPPING_FOLD should be None" + + # *Click* "Use Standard Map" button + runoffMappingGenerator = Stage.active()._widget._supplementary_widgets[0] + runoffMappingGenerator._btn_use_standard.click() + + map_status = cvars["ROF_OCN_MAPPING_STATUS"].value + assert map_status == "Standard", f"ROF_OCN_MAPPING_STATUS should be 'Standard', but got: {map_status}" + + # revert stage + assert Stage.active().title.startswith("3. Launch") + Stage.active().revert() + assert Stage.active().title.startswith("Runoff to Ocean Mapping") + +# WARNING: for this test to run successfully, MPI must be available. As such, this test +# cannot be run through the derecho login nodes. +@pytest.mark.slow +def test_custom_rof_to_ocn_mapping(): + """This test configures a case with a custom runoff to ocean mapping for a custom resolution""" + + configure_for_rof_map('C_IAF') + + assert Stage.active().title.startswith("Custom Grid") + + with tempfile.TemporaryDirectory(dir=base_temp_dir) as temp_grid_path: + + cvars["CUSTOM_GRID_PATH"].value = temp_grid_path + # since this is a JRA run, the atmosphere grid must automatically be set to TL319 + assert cvars["CUSTOM_ATM_GRID"].value == "T62" + + # Set the custom ocean grid mode + assert Stage.active().title.startswith("Ocean") + cvars["OCN_GRID_MODE"].value = "Standard" + + assert Stage.active().title.startswith("Ocean Grid") + cvars["CUSTOM_OCN_GRID"].value = "tx2_3v2" + + # Land Grid and Runoff grid should be set automatically + assert Stage.active().title.startswith("Runoff to Ocean Mapping") + + # Currently, smoothing parameters should be unset + assert cvars["ROF_OCN_MAPPING_RMAX"].value is None, "ROF_OCN_MAPPING_RMAX should be None" + assert cvars["ROF_OCN_MAPPING_FOLD"].value is None, "ROF_OCN_MAPPING_FOLD should be None" + + # *Click* the Generate New Map button + runoffMappingGenerator = Stage.active()._widget._supplementary_widgets[0] + runoffMappingGenerator._btn_generate_new.click() + + # After clicking the button, the smoothing parameters should be set to suggested values + assert cvars["ROF_OCN_MAPPING_RMAX"].value is not None, "ROF_OCN_MAPPING_RMAX should have been set to a suggested value" + assert cvars["ROF_OCN_MAPPING_FOLD"].value is not None, "ROF_OCN_MAPPING_FOLD should have been set to a suggested value" + + # *Click* the Run mapping generator button + runoffMappingGenerator._btn_run_generate.click() + + map_status = cvars["ROF_OCN_MAPPING_STATUS"].value + assert map_status.startswith("CUSTOM:"), f"ROF_OCN_MAPPING_STATUS should indicate a custom mapping, but got: {map_status}" + map_paths = map_status.split("CUSTOM:")[1] + nn_map_path, nnsm_map_path = map_paths.split(",") + + # check if the mapping file was actually created + assert os.path.isfile(nn_map_path), f"Nearest neighbor map file was not created at {nn_map_path}" + assert os.path.isfile(nnsm_map_path), f"Smoothed map file was not created at {nnsm_map_path}" + + assert Stage.active().title.startswith("3. Launch") + +if __name__ == "__main__": + test_standard_rof_to_ocn_mapping() + test_custom_rof_to_ocn_mapping() + print("All tests passed!") diff --git a/visualCaseGen/cime_interface.py b/visualCaseGen/cime_interface.py index 21848d4..43c326e 100644 --- a/visualCaseGen/cime_interface.py +++ b/visualCaseGen/cime_interface.py @@ -5,7 +5,7 @@ import socket import getpass import subprocess -from collections import namedtuple +from collections import namedtuple, defaultdict from pathlib import Path from ProConPy.dialog import alert_warning @@ -14,7 +14,7 @@ Compset = namedtuple("Compset", ["alias", "lname", "model"]) Resolution = namedtuple("Resolution", ["alias", "compset", "not_compset", "desc"]) -ComponentGrid = namedtuple("ComponentGrid", ["name", "nx", "ny", "mesh", "desc", "compset_constr", "not_compset_constr"]) +ComponentGrid = namedtuple("ComponentGrid", ["name", "nx", "ny", "mesh", "desc", "compset_constr", "not_compset_constr", "is_default"]) class CIME_interface: @@ -73,15 +73,16 @@ def __init__(self, cesmroot=None): for model in self.models[comp_class]: self._retrieve_model_phys_opt(comp_class, model) self._retrieve_domains_and_resolutions() + self._retrieve_maps() self._retrieve_compsets() self._retrieve_machines() self._retrieve_clm_data() - + def _set_cimeroot(self, cesmroot=None): """Sets the cimeroot attribute, This method is called by the __init__ method. The cimeroot attribute is set based on the cesmroot argument, which is either passed to the __init__ method or determined from the CESMROOT environment variable. - + Parameters ---------- cesmroot : str | Path | None @@ -101,7 +102,7 @@ def _set_cimeroot(self, cesmroot=None): cesmroot = Path(Path(filepath).parent.parent) assert cesmroot.is_dir(), "Cannot find CESM root directory!" - # Set cimeroot attribute + # Set cimeroot attribute self.cimeroot = cesmroot / "cime" assert self.cimeroot.is_dir(), f"Cannot find CIME directory at {self.cimeroot}" @@ -110,7 +111,7 @@ def _check_cime_compatibility(self): """Checks the compatibility of the CIME version. This method is called by the __init__ method. The CIME version is determined based on the git tag of the CIME repository. The CIME version is checked for compatibility with visualCaseGen.""" - + # cime git tag: cime_git_tag = subprocess.check_output( ["git", "-C", self.cimeroot, "describe", "--tags"] @@ -352,10 +353,10 @@ def _retrieve_models(self, comp_class): self.models[comp_class].append(model) def _get_domains(self): - """Reads and returns component grids, i.e., domains, from the CIME XML file. The + """Reads and returns component grids, i.e., domains, from the CIME XML file. The compset_constr and not_compset_constr attributes of the ComponentGrid object are not filled in here, but are filled in later in the _retrieve_component_grid_constraints method - + Returns ------- dict @@ -398,12 +399,13 @@ def _get_domains(self): desc=desc+support, compset_constr=set(), # to be filled in later not_compset_constr=set(), # to be filled in later + is_default=False # to be updated later ) return domains def _retrieve_domains_and_resolutions(self): - """Retrieves and stores component grids and model resolutions from the corresponding CIME + """Retrieves and stores component grids and model resolutions from the corresponding CIME XML files. Retrieved resolutions are stored in self.resolutions list and component grids are stored in self.domains dict.""" @@ -413,14 +415,53 @@ def _retrieve_domains_and_resolutions(self): # Get the initial (temporary) dict of domains from the component_grids file domains = self._get_domains() + domain_found = {domain_name: False for domain_name in domains.keys()} grids = self._grids_obj.get_child("grids") - # Component grids dict and resolutions list to be populated in the loop below + # Domains, i.e., component grids, are stored in self.domains dict. The keys of + # self.domains are component names, e.g., "ocnice". The values are dicts where keys are domain names, + # e.g., "tx2_3v2", and values are ComponentGrid named tuples with attributes name, nx, ny, mesh, desc, + # compset_constr, and not_compset_constr. Since these constraints are resolution-specific, and + # not domain-specific, they are initially inserted into sets and then processed appropriately + # in the _process_domain_constraints method, where they are maintained or dropped depending on + # whether they are common across all resolutions that include this domain. self.domains = {comp: {} for comp in self._grids_obj._comp_gridnames} + + # Resolutions, i.e., combinations of component grids, .e.,g "TL319_t232" are stored in self.resolutions list. + # Each resolution is a Resolution named tuple with attributes alias, compset, not_compset, and desc. + # The compset and not_compset attributes are strings that represent the compset constraints for this resolution. self.resolutions = [] - # loop through model grids (resolutions) + # Loop through model grid defaults. i.e., default grids for each model, to populate self.domains + # We read them in case any domain is only listed in the model grid defaults and not in the model grids. + # In that case, we want to make sure that the domain is still included in the self.domains dict. + model_grid_defaults = self._grids_obj.get_child("model_grid_defaults", root=grids) + default_grids = self._grids_obj.get_children("grid", root=model_grid_defaults) + for default_grid in default_grids: + comp_name = self._grids_obj.get(default_grid, "name") # e.g., atm, lnd, ocnice, etc. + compset = self._grids_obj.get(default_grid, "compset") + comp_grid = self._grids_obj.text(default_grid) + if comp_grid == "null": + continue + if comp_grid in domains: + if comp_grid not in self.domains[comp_name]: + self.domains[comp_name][comp_grid] = ComponentGrid( + name=comp_grid, + nx=domains[comp_grid].nx, + ny=domains[comp_grid].ny, + mesh=domains[comp_grid].mesh, + desc=domains[comp_grid].desc, + compset_constr=set(), + not_compset_constr=set(), + is_default=True + ) + domain_found[comp_grid] = True + + # Loop through model grids, i.e., resolutions, to populate self.resolutions. + # Also, for each resolution, loop through the component grids that are part of this resolution + # and add them to the self.domains dict depending on which component name they are associated with. This is + # how we determine which domains (model grids) are part of which component name (e.g., atm, lnd, ocnice, etc.). model_grid_nodes = self._grids_obj.get_children("model_grid", root=grids) for model_grid_node in model_grid_nodes: alias = self._grids_obj.get(model_grid_node, "alias") @@ -428,56 +469,171 @@ def _retrieve_domains_and_resolutions(self): not_compset = self._grids_obj.get(model_grid_node, "not_compset") desc = "" - # loop through component grids that are part of this resolution + # Loop through all component grids that are part of this resolution all_component_grids_found = True grid_nodes = self._grids_obj.get_children("grid", root=model_grid_node) for grid_node in grid_nodes: comp_name = self._grids_obj.get(grid_node, "name") # e.g., atm, lnd, ocnice, etc. comp_grid = self._grids_obj.text(grid_node) - if comp_grid in domains: - if comp_grid not in self.domains[comp_name]: - self.domains[comp_name][comp_grid] = ComponentGrid( - name=comp_grid, - nx=domains[comp_grid].nx, - ny=domains[comp_grid].ny, - mesh=domains[comp_grid].mesh, - desc=domains[comp_grid].desc, - compset_constr=set(), - not_compset_constr=set(), - ) - domain = self.domains[comp_name][comp_grid] - desc += ' | ' + ' ' + comp_name.upper() + ': ' + domain.desc - domain.compset_constr.add(compset) - domain.not_compset_constr.add(not_compset) - else: + + # Skip if the component grid is null. This means that this component is not part of this resolution. + if comp_grid == "null": + continue + + # If the component grid is not found in the domains dict, then this resolution is invalid and we skip it. + if comp_grid not in domains: #logger.warning(f"Domain {comp_grid} not found in component_grids file.") all_component_grids_found = False break - - if not all_component_grids_found: - continue - self.resolutions.append(Resolution(alias, compset, not_compset, desc)) - + # If the component grid is not already in the self.domains dict for this component, add it. + if comp_grid not in self.domains[comp_name]: + self.domains[comp_name][comp_grid] = ComponentGrid( + name=comp_grid, + nx=domains[comp_grid].nx, + ny=domains[comp_grid].ny, + mesh=domains[comp_grid].mesh, + desc=domains[comp_grid].desc, + compset_constr=set(), + not_compset_constr=set(), + is_default=False + ) + domain_found[comp_grid] = True + + # Retrieve the domain object for this component grid and add the compset and not_compset constraints to it. + domain = self.domains[comp_name][comp_grid] + desc += ' | ' + ' ' + comp_name.upper() + ': ' + domain.desc + domain.compset_constr.add(compset) + domain.not_compset_constr.add(not_compset) + + # Add this resolution to the self.resolutions list + if all_component_grids_found: + self.resolutions.append(Resolution(alias, compset, not_compset, desc)) + + # If there are remaining domains that are not found to be belonging to any component, attempt to find out + # which component they belong to by looking at the description of the domain. + descr_tips ={ + 'rof': ('rof', 'runoff'), + 'atm': ('atm', 'atmosphere'), + 'lnd': ('lnd', 'land'), + 'ocnice': ('ocn', 'ocean', 'ice'), + 'glc': ('glc', 'glacier'), + 'wav': ('wav'), + } + for domain in domains: + if not domain_found[domain]: + for comp_name, tips in descr_tips.items(): + if any(tip in domains[domain].desc.lower() for tip in tips): + self.domains[comp_name][domain] = ComponentGrid( + name=domain, + nx=domains[domain].nx, + ny=domains[domain].ny, + mesh=domains[domain].mesh, + desc=domains[domain].desc, + compset_constr=set(), + not_compset_constr=set(), + is_default=False + ) + break + + + # Finally, process the compset and not_compset constraints for each domain (component grid). self._process_domain_constraints() + def get_mesh_path(self, comp_name, domain_name): + """Returns the mesh file path for a given component name and domain name. + + Parameters + ---------- + comp_name : str + component name, e.g., "atm", "lnd", "ocnice", etc. + domain_name : str + domain name, e.g., "tx2_3v2", "gx1v7", etc. + + Returns + ------- + str + mesh file path for the given component name and domain name. + If not found, returns an empty string. + """ + + if comp_name not in self.domains: + logger.error(f"Component {comp_name} not found in domains.") + return '' + if domain_name not in self.domains[comp_name]: + logger.error(f"Domain {domain_name} not found for component {comp_name}.") + return '' + + domain = self.domains[comp_name][domain_name] + mesh_path = domain.mesh + + if 'DIN_LOC_ROOT' in mesh_path: + assert self.din_loc_root is not None, "DIN_LOC_ROOT not set." + mesh_path = mesh_path.replace('$DIN_LOC_ROOT', self.din_loc_root) + mesh_path = mesh_path.replace('${DIN_LOC_ROOT}', self.din_loc_root) + + return mesh_path + + def _retrieve_maps(self): + """Retrieves the grid mapping files from the CIME XML file. The retrieved mapping files are stored + in the self.maps attribute, which is a nested dict where keys are source grids and values + are dicts where keys are destination grids and values are lists of (name, filepath) tuples. + This is currently used only to determine whether a runoff to ocean mapping file needs to be + generated or not. + """ + + assert hasattr(self, '_grids_obj'), "_grids_obj attribute not found. Call _retrieve_domains_and_resolutions() first." + + self.maps = defaultdict(dict) # maps[src_grid][dst_grid] = [(name, filepath), ...] + + gridmaps = self._grids_obj.get_child("gridmaps") + gridmap_nodes = self._grids_obj.get_children("gridmap", root=gridmaps) + for gridmap_node in gridmap_nodes: + comps = list(gridmap_node.attrib.keys()) # not being utilized currently + grids = list(gridmap_node.attrib.values()) + src_grid, dst_grid = grids[0], grids[1] + self.maps[src_grid][dst_grid] = [] + map_nodes = self._grids_obj.get_children("map", root=gridmap_node) + for map_node in map_nodes: + name = self._grids_obj.get(map_node, "name") + path = self._grids_obj.text(map_node) + self.maps[src_grid][dst_grid].append( (name, path) ) + def _process_domain_constraints(self): """Update the compset_constr and not_compset_constr attributes of the ComponentGrid objects in - the self.domains dict. This method is called after the component grids and resolutions have + the self.domains dict. This method is called after the component grids and resolutions have been retrieved. It updates the compset and not_compset constraints for each domain: If a domain is part of one or more resolutions that have no compset/not_compset constraints, then the domain - is deemed unconstrained. Otherwise, the domain is constrained by the disjunction of the - compset/not_compset constraints of all the resolutions it is part of. + is deemed unconstrained. Otherwise, the domain is constrained by the disjunction of the compset + constraints and the conjunction of the not_compset constraints of all the resolutions it is part of. """ for comp_name, domains in self.domains.items(): for domain_name, domain in domains.items(): - final_compset_constr = '' - if None not in domain.compset_constr and len(domain.compset_constr) >= 0: + # compset constraint + if None in domain.compset_constr or len(domain.compset_constr) == 0: + final_compset_constr = '' + else: final_compset_constr = '|'.join(domain.compset_constr) - final_not_compset_constr = '' - if None not in domain.not_compset_constr and len(domain.not_compset_constr) >= 0: - final_not_compset_constr = '|'.join(domain.not_compset_constr) + + # not_compset constraint: collect expressions (i.e., models with or without options) that are + # common across all not_compset_constr sets for this domain. If there are no common expressions + # across all not_compset_constr sets, then the final not_compset_constr is empty. + if None in domain.not_compset_constr or len(domain.not_compset_constr) == 0: + final_not_compset_constr = '' + else: + expr_count = defaultdict(int) + for not_compset_constr in domain.not_compset_constr: + exprs = set(not_compset_constr.split('|')) + for expr in exprs: + expr_count[expr] += 1 + common_exprs = [expr for expr, count in expr_count.items() if count == len(domain.not_compset_constr)] + if common_exprs: + final_not_compset_constr = '|'.join(common_exprs) + else: + final_not_compset_constr = '' + + # Update the domain object with the final compset and not_compset constraints self.domains[comp_name][domain_name] = domain._replace( compset_constr=final_compset_constr, not_compset_constr=final_not_compset_constr @@ -517,7 +673,7 @@ def get_components_from_compset_lname(self, compset_lname): for i, comp_class in enumerate(self.comp_classes): components[comp_class] = compset_lname_split[i+1] - + return components @@ -680,7 +836,7 @@ def _handle_machine_not_ported(self): ) logger.warning( "Please set the CIME machine in the visualCaseGen configuration file." - ) + ) return diff --git a/visualCaseGen/config_vars/grid_vars.py b/visualCaseGen/config_vars/grid_vars.py index 3313279..566ec08 100644 --- a/visualCaseGen/config_vars/grid_vars.py +++ b/visualCaseGen/config_vars/grid_vars.py @@ -151,3 +151,10 @@ def default_fsurdat_area_spec(): # Note: this var isn't a part of any of the stages. ConfigVarStr("FSURDAT_MOD_STATUS", widget_none_val="") # a status variable to prevent the completion of the stage prematurely + # A preexisting ROF grid to be picked for custom grid + ConfigVarStrMS("CUSTOM_ROF_GRID") + + # Runoff to ocean mapping status + ConfigVarStr("ROF_OCN_MAPPING_STATUS", widget_none_val="") + ConfigVarReal("ROF_OCN_MAPPING_RMAX") + ConfigVarReal("ROF_OCN_MAPPING_FOLD") diff --git a/visualCaseGen/custom_widget_types/case_creator.py b/visualCaseGen/custom_widget_types/case_creator.py index 8205857..c4e589b 100644 --- a/visualCaseGen/custom_widget_types/case_creator.py +++ b/visualCaseGen/custom_widget_types/case_creator.py @@ -223,6 +223,8 @@ def _update_modelgrid_aliases(self, custom_grid_path, ocn_grid, do_exec): # Component grid names: atm_grid = cvars["CUSTOM_ATM_GRID"].value lnd_grid = cvars["CUSTOM_LND_GRID"].value + rof_grid = cvars["CUSTOM_ROF_GRID"].value + # modelgrid_aliases xml file that stores resolutions: srcroot = self._cime.srcroot ccs_config_root = Path(srcroot) / "ccs_config" @@ -239,12 +241,19 @@ def _update_modelgrid_aliases(self, custom_grid_path, ocn_grid, do_exec): if not os.access(modelgrid_aliases_xml, os.W_OK): raise RuntimeError(f"Cannot write to {modelgrid_aliases_xml}.") + # Construct the component grids string to be logged: + component_grids_str = f' atm grid: "{atm_grid}" \n' + component_grids_str += f' lnd grid: "{lnd_grid}" \n' + component_grids_str += f' ocn grid: "{ocn_grid}".\n' + if rof_grid is not None and rof_grid != "": + component_grids_str += f' rof grid: "{rof_grid}".\n' + # log the modification of modelgrid_aliases.xml: with self._out: print( f'{BPOINT} Updating ccs_config/modelgrid_aliases_nuopc.xml file to include the new ' f'resolution "{resolution_name}" consisting of the following component grids.\n' - f' atm grid: "{atm_grid}", lnd grid: "{lnd_grid}", ocn grid: "{ocn_grid}".\n' + f'{component_grids_str}' ) # Read in xml file and generate grids object file: @@ -278,6 +287,7 @@ def _update_modelgrid_aliases(self, custom_grid_path, ocn_grid, do_exec): ) new_atm_grid.text = atm_grid + # Add lnd grid to resolution entry: new_lnd_grid = SubElement( new_resolution, "grid", @@ -285,6 +295,7 @@ def _update_modelgrid_aliases(self, custom_grid_path, ocn_grid, do_exec): ) new_lnd_grid.text = lnd_grid + # Add ocn grid to resolution entry: new_ocnice_grid = SubElement( new_resolution, "grid", @@ -292,6 +303,16 @@ def _update_modelgrid_aliases(self, custom_grid_path, ocn_grid, do_exec): ) new_ocnice_grid.text = ocn_grid + # Add rof grid to resolution entry if it exists: + if rof_grid is not None and rof_grid != "": + new_rof_grid = SubElement( + new_resolution, + "grid", + attrib={"name": "rof"}, + ) + new_rof_grid.text = rof_grid + + if not do_exec: return @@ -485,6 +506,28 @@ def _run_create_newcase(self, caseroot, compset, resolution, do_exec): raise RuntimeError("Error creating case.") def _apply_all_xmlchanges(self, do_exec): + """Apply all the necessary xmlchanges to the case. + + Parameters + ---------- + do_exec : bool + If True, execute the commands. If False, only print them. + """ + + # If standard grid is selected, no modifications are needed: + grid_mode = cvars["GRID_MODE"].value + if grid_mode == "Standard": + return # no modifications needed for standard grid + else: + assert grid_mode == "Custom", f"Unknown grid mode: {grid_mode}" + + self._apply_lnd_grid_xmlchanges(do_exec) + self._apply_ocn_grid_xmlchanges(do_exec) + self._apply_runoff_ocn_mapping_xmlchanges(do_exec) + + + 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": @@ -507,29 +550,47 @@ def _apply_all_xmlchanges(self, do_exec): 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}" + + def _apply_ocn_grid_xmlchanges(self, do_exec): + """Apply xmlchanges related to custom ocean grid if needed.""" - # Set NTASKS based on grid size. e.g. NX * NY < max_pts_per_core - if cvars["COMP_OCN"].value == "mom": + # Set NTASKS based on grid size if custom ocn grid. e.g. NX * NY < max_pts_per_core + if cvars["COMP_OCN"].value == "mom" and cvars["OCN_GRID_MODE"].value == "Custom": num_points = int(cvars["OCN_NX"].value) * int(cvars["OCN_NY"].value) cores = CaseCreator._calc_cores_based_on_grid(num_points) with self._out: print(f"{COMMENT}Apply NTASK grid xml changes:{RESET}\n") xmlchange("NTASKS_OCN",cores, do_exec, self._is_non_local(), self._out) + + def _apply_runoff_ocn_mapping_xmlchanges(self, do_exec): + """Apply xmlchanges related to runoff to ocean mapping files if custom mapping is selected.""" + + if (rof_ocn_mapping_status := cvars["ROF_OCN_MAPPING_STATUS"].value) is not None: + if rof_ocn_mapping_status.startswith("CUSTOM:"): + mapping_files = rof_ocn_mapping_status[7:] + nn_map_file, nnsm_map_file = mapping_files.split(",") + with self._out: + print(f"{COMMENT}Apply runoff to ocean mapping xml changes:{RESET}\n") + xmlchange("ROF2OCN_ICE_RMAPNAME", nnsm_map_file, do_exec, self._is_non_local(), self._out) + xmlchange("ROF2OCN_LIQ_RMAPNAME", nnsm_map_file, do_exec, self._is_non_local(), self._out) + @staticmethod - def _calc_cores_based_on_grid( num_points, min_points_per_core = 32, max_points_per_core = 800, ideal_multiple_of_cores_used = 128): + def _calc_cores_based_on_grid( num_points, min_points_per_core = 32, max_points_per_core = 300, ideal_multiple_of_cores_used = 128): """Calculate the number of cores based on the grid size.""" min_cores = math.ceil(num_points/max_points_per_core) - max_cores = math.ceil(num_points/min_points_per_core) + max_cores = math.ceil(num_points/min_points_per_core) + + # If min_cores is less than the first multiple of ideal cores, just return the min_cores + if max_cores < ideal_multiple_of_cores_used: + return min_cores # Request a multiple of the entire core (ideal_multiple_of_cores_used) starting from the min ideal_cores = ((min_cores + ideal_multiple_of_cores_used - 1) // ideal_multiple_of_cores_used) * ideal_multiple_of_cores_used - if ideal_cores <= max_cores: - return ideal_cores - else: - return (max_cores+min_cores)//2 + return ideal_cores + def _apply_user_nl_changes(self, model, var_val_pairs, do_exec, comment=None, log_title=True): """Apply changes to a given user_nl file.""" @@ -586,7 +647,7 @@ def _apply_mom_namelist_changes(self, do_exec): # Determine timesteps based on the grid resolution (assuming coupling frequency of 1800.0 sec): res_x = float(cvars['OCN_LENX'].value) / int(cvars["OCN_NX"].value) res_y = float(cvars['OCN_LENY'].value) / int(cvars["OCN_NY"].value) - dt = 600.0 * min(res_x,res_y) # A 1-deg grid should have ~600 sec tstep (a safe value) + dt = 7200.0 * min(res_x,res_y) # A 1-deg grid should have ~600 sec tstep (a safe value) # Make sure 1800.0 is a multiple of dt and dt is a power of 2 and/or 3: dt = min((1800.0 / n for n in [2**i * 3**j for i in range(10) for j in range(6)] if 1800.0 % n == 0), key=lambda x: abs(dt - x)) # Try setting dt_therm to dt*4, or dt*3, or dt*3, depending on whether 1800.0 becomes a multiple of dt: diff --git a/visualCaseGen/custom_widget_types/mom6_bathy_launcher.py b/visualCaseGen/custom_widget_types/mom6_bathy_launcher.py index 62a9fc4..9d5839a 100644 --- a/visualCaseGen/custom_widget_types/mom6_bathy_launcher.py +++ b/visualCaseGen/custom_widget_types/mom6_bathy_launcher.py @@ -408,6 +408,11 @@ def topo_file_path(): def vgrid_file_path(): custom_ocn_grid_path = MOM6BathyLauncher.get_custom_ocn_grid_path() return custom_ocn_grid_path / f"ocean_vgrid_{MOM6BathyLauncher.nc_file_suffix()}" + + @staticmethod + def scrip_grid_file_path(): + custom_ocn_grid_path = MOM6BathyLauncher.get_custom_ocn_grid_path() + return custom_ocn_grid_path / f"scrip_{MOM6BathyLauncher.nc_file_suffix()}" @staticmethod def esmf_mesh_file_path(): diff --git a/visualCaseGen/custom_widget_types/runoff_mapping_generator.py b/visualCaseGen/custom_widget_types/runoff_mapping_generator.py new file mode 100644 index 0000000..b8c7c30 --- /dev/null +++ b/visualCaseGen/custom_widget_types/runoff_mapping_generator.py @@ -0,0 +1,236 @@ +import os +from ipywidgets import HBox, VBox, Button, Output, Label +from pathlib import Path + +from ProConPy.out_handler import handler as owh +from ProConPy.config_var import cvars +from ProConPy.dialog import alert_warning +from mom6_bathy import mapping +from visualCaseGen.custom_widget_types.mom6_bathy_launcher import MOM6BathyLauncher + +class RunoffMappingGenerator(VBox): + """Widget to generate runoff to ocean mapping for custom grids. + The widget first checks if there exists a standard mapping between the selected + runoff grid and the custom ocean grid. If not, it allows the user to generate + a new mapping using the mom6_bathy mapping module. + """ + + def __init__(self, cime, **kwargs): + super().__init__(**kwargs) + self.cime = cime + + self._btn_use_standard = Button( + description="Use Standard Map", + disabled=True, + tooltip="Use the standard mapping files already available for the selected grids.", + ) + self._btn_use_standard.on_click(self.on_btn_use_standard_clicked) + + self._btn_generate_new = Button( + description="Generate New Map", + disabled=False, + tooltip="Generate a new mapping file using the mom6_bathy mapping module.", + ) + self._btn_generate_new.on_click(self.on_btn_generate_new_clicked) + + self._out = Output() + + self._btn_run_generate = Button( + description="Run mapping generator", + disabled=False, + button_style="success", + tooltip="Run the mapping generator with the specified parameters.", + layout={"width": "260px", "align_self": "center"}, + ) + self._btn_run_generate.on_click(self.on_btn_run_generate_clicked) + + self._generate_new_dialog = VBox([ + cvars["ROF_OCN_MAPPING_RMAX"].widget, + cvars["ROF_OCN_MAPPING_FOLD"].widget, + self._btn_run_generate, + self._out + ], + layout={"display": "none"} + ) + + self.children = [ + HBox([ + Label("Select mapping option:"), + self._btn_use_standard, + self._btn_generate_new + ], + layout={"justify_content": "center", "margin": "10px"} + ), + self._generate_new_dialog + ] + + # Reset the widget when the runoff grid changes + cvars["CUSTOM_ROF_GRID"].observe(self.reset, names='value', type='change') + + @property + def disabled(self): + return super().disabled + + @disabled.setter + def disabled(self, value): + self._btn_use_standard.disabled = value or not self.standard_map_exists() + self._btn_generate_new.disabled = value + for child in self._generate_new_dialog.children: + child.disabled = value + + def reset(self, change): + """Reset all widget children and auxiliary config variables. To be called + when the runoff grid changes.""" + self._out.clear_output() + self._generate_new_dialog.layout.display = "none" + cvars["ROF_OCN_MAPPING_RMAX"].value = None + cvars["ROF_OCN_MAPPING_FOLD"].value = None + + def standard_map_exists(self): + """Check if there exists a standard mapping between the selected + runoff grid and the custom ocean grid. + """ + if cvars["OCN_GRID_MODE"].value != "Standard": + return False + + rof_grid = cvars["CUSTOM_ROF_GRID"].value + ocn_grid = cvars["CUSTOM_OCN_GRID"].value + + if ocn_grid in self.cime.maps[rof_grid]: + return True + + return False + + @owh.out.capture() + def on_btn_use_standard_clicked(self, b): + """Handler for the 'Use Standard Map' button click event. + Sets the ROF_OCN_MAPPING_STATUS variable to indicate that + the standard mapping will be used. + """ + + if not self.standard_map_exists(): + alert_warning( + "No standard mapping exists between the selected runoff grid " + "and the custom ocean grid. Please generate a new mapping." + ) + return + + self._out.clear_output() + self._generate_new_dialog.layout.display = "none" + cvars["ROF_OCN_MAPPING_STATUS"].value = "Standard" + + + def get_rof_grid_and_mesh(self): + """Return the runoff grid name and mesh path.""" + + rof_grid = cvars["CUSTOM_ROF_GRID"].value + rof_mesh_path = self.cime.get_mesh_path("rof", rof_grid) + return rof_grid, rof_mesh_path + + def get_ocn_grid_and_mesh(self): + """Return the ocean grid name and mesh path.""" + + ocn_grid_mode = cvars["OCN_GRID_MODE"].value + + match ocn_grid_mode: + case "Standard": + ocn_grid = cvars["CUSTOM_OCN_GRID"].value + ocn_mesh_path = self.cime.get_mesh_path("ocnice", ocn_grid) + case "Create New": + ocn_grid = cvars["CUSTOM_OCN_GRID_NAME"].value + ocn_mesh_path = MOM6BathyLauncher.esmf_mesh_file_path() + case _: + assert False, f"Unsupported OCN_GRID_MODE: {ocn_grid_mode}" + + return ocn_grid, ocn_mesh_path + + def on_btn_generate_new_clicked(self, b): + """Handler for the 'Generate New Map' button click event. + Sets the ROF_OCN_MAPPING_STATUS variable to indicate that + a new mapping will be generated. + """ + + cvars["ROF_OCN_MAPPING_STATUS"].value = None + self._out.clear_output() + self._generate_new_dialog.layout.display = "" + + rmax = cvars["ROF_OCN_MAPPING_RMAX"].value + fold = cvars["ROF_OCN_MAPPING_FOLD"].value + + # Suggest default values for RMAX and FOLD if not set + if rmax is None and fold is None: + _, ocn_mesh_path = self.get_ocn_grid_and_mesh() + suggested_rmax, suggested_fold = mapping.get_suggested_smoothing_params(ocn_mesh_path) + + cvars["ROF_OCN_MAPPING_RMAX"].value = suggested_rmax + cvars["ROF_OCN_MAPPING_FOLD"].value = suggested_fold + + + @owh.out.capture() + def on_btn_run_generate_clicked(self, b): + """Handler for the 'Run mapping generator' button click event. + Runs the mapping generator with the specified parameters. + """ + + cvars["ROF_OCN_MAPPING_STATUS"].value = None + self._out.clear_output() + + rmax = cvars["ROF_OCN_MAPPING_RMAX"].value + if rmax is None: + alert_warning("Please specify a valid RMAX value.") + return + + fold = cvars["ROF_OCN_MAPPING_FOLD"].value + if fold is None: + alert_warning("Please specify a valid FOLD value.") + return + + rof_grid, rof_mesh_path = self.get_rof_grid_and_mesh() + ocn_grid, ocn_mesh_path = self.get_ocn_grid_and_mesh() + + try: + # disable the widget ahead of running the mapping generator + self.disabled = True + + mapping_file_prefix = f"{rof_grid}_to_{ocn_grid}_map" + output_dir = RunoffMappingGenerator.mapping_dir() + + # Run the mapping generator + with self._out: + mapping.gen_rof_maps( + rof_mesh_path=rof_mesh_path, + ocn_mesh_path=ocn_mesh_path, + output_dir=output_dir, + mapping_file_prefix=mapping_file_prefix, + rmax=rmax, + fold=fold + ) + + nn_map_filepath = mapping.get_nn_map_filepath( + mapping_file_prefix=mapping_file_prefix, + output_dir=output_dir, + ) + + nnsm_map_filepath = mapping.get_smoothed_map_filepath( + mapping_file_prefix=mapping_file_prefix, + output_dir=output_dir, + rmax=rmax, + fold=fold + ) + + # Set the mapping status to CUSTOM after successful generation + cvars["ROF_OCN_MAPPING_STATUS"].value = f"CUSTOM:{nn_map_filepath},{nnsm_map_filepath}" + + except Exception as e: + alert_warning( + f"An error occurred while generating the mapping: {e}" + ) + self.disabled = False + + + @staticmethod + def mapping_dir(): + custom_grid_path = cvars["CUSTOM_GRID_PATH"].value + mapping_dir = Path(custom_grid_path) / "mapping" + os.makedirs(mapping_dir, exist_ok=True) + return mapping_dir \ No newline at end of file diff --git a/visualCaseGen/initialize.py b/visualCaseGen/initialize.py new file mode 100644 index 0000000..e13433c --- /dev/null +++ b/visualCaseGen/initialize.py @@ -0,0 +1,41 @@ +import logging + +from ProConPy.config_var import ConfigVar, cvars +from ProConPy.stage import Stage +from ProConPy.csp_solver import csp +from visualCaseGen.cime_interface import CIME_interface +from visualCaseGen.initialize_configvars import initialize_configvars +from visualCaseGen.initialize_widgets import initialize_widgets +from visualCaseGen.initialize_stages import initialize_stages +from visualCaseGen.specs.options import set_options +from visualCaseGen.specs.relational_constraints import get_relational_constraints + +logger = logging.getLogger('\t'+__name__.split('.')[-1]) + + +def initialize(cesmroot=None): + """Initialize the visualCaseGen system by setting up configuration variables, stages, and widgets. + + Parameters: + ----------- + cesmroot : str, optional + The path to the CESM root directory. If not provided, it will be determined automatically. + + Returns: + -------- + cime : CIME_interface + An instance of the CIME_interface class, initialized with the provided CESM root directory. + """ + + logger.info("Initializing the visualCaseGen system...") + + ConfigVar.reboot() + Stage.reboot() + cime = CIME_interface(cesmroot=cesmroot) + initialize_configvars(cime) + initialize_widgets(cime) + initialize_stages(cime) + set_options(cime) + csp.initialize(cvars, get_relational_constraints(cvars), Stage.first()) + + return cime \ No newline at end of file diff --git a/visualCaseGen/specs/grid_options.py b/visualCaseGen/specs/grid_options.py index 8057dac..2f00983 100644 --- a/visualCaseGen/specs/grid_options.py +++ b/visualCaseGen/specs/grid_options.py @@ -17,6 +17,7 @@ def set_grid_options(cime): set_custom_atm_grid_options(cime) set_custom_ocn_grid_options(cime) set_custom_lnd_grid_options(cime) + set_custom_rof_grid_options(cime) def set_standard_grid_options(cime): @@ -95,7 +96,7 @@ def check_comp_grid(comp_class, proposed_grid, compset_lname): proposed_grid.compset_constr, compset_lname ): return False - if proposed_grid.not_compset_constr and re.search( + if (not proposed_grid.is_default) and proposed_grid.not_compset_constr and re.search( proposed_grid.not_compset_constr, compset_lname ): return False @@ -108,8 +109,12 @@ def set_custom_atm_grid_options(cime): This function is called at initialization.""" # CUSTOM_ATM_GRID options - def custom_atm_grid_options_func(comp_atm, grid_mode): + def custom_atm_grid_options_func(grid_mode): """Return the options and descriptions for the custom ATM grid variable.""" + + if grid_mode != "Custom": + return None, None + compset_lname = cvars["COMPSET_LNAME"].value compatible_atm_grids = [] descriptions = [] @@ -122,7 +127,7 @@ def custom_atm_grid_options_func(comp_atm, grid_mode): cv_custom_atm_grid = cvars["CUSTOM_ATM_GRID"] cv_custom_atm_grid.options_spec = OptionsSpec( - func=custom_atm_grid_options_func, args=(cvars["COMP_ATM"], cvars["GRID_MODE"]) + func=custom_atm_grid_options_func, args=(cvars["GRID_MODE"],) ) @@ -228,3 +233,34 @@ def custom_lnd_grid_options_func(comp_lnd, custom_atm_grid, lnd_grid_mode): cv_lnd_include_nonveg = cvars["LND_INCLUDE_NONVEG"] cv_lnd_include_nonveg.options = ["True", "False"] + + +def set_custom_rof_grid_options(cime): + """Set the options and options specs for the custom ROF grid variable. + This function is called at initialization.""" + + # CUSTOM_ROF_GRID options + def custom_rof_grid_options_func(grid_mode): + """Return the options and descriptions for the custom ROF grid variable.""" + + if grid_mode != "Custom": + return None, None + + if cvars["COMP_ROF"].value == "srof": + return ["null"], ["(When stub ROF is selected, custom ROF grid is set to null.)"] + + # Loop through all ROF grids and check if they are compatible with the compset constraints + compset_lname = cvars["COMPSET_LNAME"].value + compatible_rof_grids = [] + descriptions = [] + for rof_grid in cime.domains["rof"].values(): + if check_comp_grid("ROF", rof_grid, compset_lname) is False: + continue + compatible_rof_grids.append(rof_grid.name) + descriptions.append(rof_grid.desc) + return compatible_rof_grids, descriptions + + cv_custom_rof_grid = cvars["CUSTOM_ROF_GRID"] + cv_custom_rof_grid.options_spec = OptionsSpec( + func=custom_rof_grid_options_func, args=(cvars["GRID_MODE"],) + ) \ No newline at end of file diff --git a/visualCaseGen/specs/relational_constraints.py b/visualCaseGen/specs/relational_constraints.py index 6da1eb6..4bc3c99 100644 --- a/visualCaseGen/specs/relational_constraints.py +++ b/visualCaseGen/specs/relational_constraints.py @@ -24,6 +24,8 @@ def get_relational_constraints(cvars): OCN_NX = cvars['OCN_NX']; OCN_NY = cvars['OCN_NY']; OCN_LENX = cvars['OCN_LENX']; OCN_LENY = cvars['OCN_LENY'] LND_GRID_MODE = cvars['LND_GRID_MODE']; LND_SOIL_COLOR = cvars['LND_SOIL_COLOR']; LND_DOM_PFT = cvars['LND_DOM_PFT'] LND_MAX_SAT_AREA = cvars['LND_MAX_SAT_AREA']; LND_STD_ELEV = cvars['LND_STD_ELEV'] + CUSTOM_ROF_GRID = cvars['CUSTOM_ROF_GRID'] + ROF_OCN_MAPPING_RMAX = cvars['ROF_OCN_MAPPING_RMAX']; ROF_OCN_MAPPING_FOLD = cvars['ROF_OCN_MAPPING_FOLD'] # Return a dictionary of constraints where keys are the z3 boolean expressions corresponding to the constraints # and values are error messages to be displayed when the constraint is violated. @@ -165,7 +167,35 @@ def get_relational_constraints(cvars): "Max fraction of saturated area must be set to a value between 0 and 1.", LND_STD_ELEV >= 0.0: - "Standard deviation of elevation must be a nonnegative number." + "Standard deviation of elevation must be a nonnegative number.", + + # Custom rof grid constraints ------------------ + Implies(In(CUSTOM_ROF_GRID, ["JRA025", "rx1"]), COMP_ROF=="drof"): + "JRA025 and rx1 runoff grids can only be selected if DROF is the runoff component.", + + CUSTOM_ROF_GRID != "r05mz": # mizuroute is no longer available + "r05mz runoff grid can only be selected if MIZUROUTE is the runoff component.", + + In(COMP_ROF_OPTION, ["NYF", "IAF"]) == (CUSTOM_ROF_GRID=="rx1"): + "When Core2 forcing is selected for the ocean component, the runoff grid must be set to rx1.", + + (Contains(COMP_ROF_OPTION, "JRA")) == (CUSTOM_ROF_GRID == "JRA025"): + "When JRA forcing is selected for the ocean component, the runoff grid must be set to JRA025.", + + (COMP_ROF_OPTION == "GLOFAS") == (CUSTOM_ROF_GRID=="GLOFAS"): + "When GLOFAS forcing is selected for the ocean component, the runoff grid must be set to GLOFAS.", + + ROF_OCN_MAPPING_RMAX > 0: + "ROF_OCN_MAPPING_RMAX must be a positive number.", + + ROF_OCN_MAPPING_FOLD > 0: + "ROF_OCN_MAPPING_FOLD must be a positive number.", + + ROF_OCN_MAPPING_RMAX <= 4000: + "ROF_OCN_MAPPING_RMAX must be less than or equal to 4000 km.", + + ROF_OCN_MAPPING_FOLD <= 8000: + "ROF_OCN_MAPPING_FOLD must be less than or equal to 8000 km.", #### Assertions to stress-test the CSP solver diff --git a/visualCaseGen/stages/grid_stages.py b/visualCaseGen/stages/grid_stages.py index a1aad75..4b6db00 100644 --- a/visualCaseGen/stages/grid_stages.py +++ b/visualCaseGen/stages/grid_stages.py @@ -3,6 +3,7 @@ from pathlib import Path import time import os +from z3 import And from ProConPy.config_var import cvars from ProConPy.stage import Stage, Guard @@ -10,6 +11,7 @@ from visualCaseGen.custom_widget_types.stage_widget import StageWidget from visualCaseGen.custom_widget_types.mom6_bathy_launcher import MOM6BathyLauncher from visualCaseGen.custom_widget_types.clm_modifier_launcher import MeshMaskModifierLauncher, FsurdatModifierLauncher +from visualCaseGen.custom_widget_types.runoff_mapping_generator import RunoffMappingGenerator logger = logging.getLogger("\t" + __name__.split(".")[-1]) @@ -296,3 +298,31 @@ def initialize_grid_stages(cime): cvars["FSURDAT_MOD_STATUS"] ], ) + + + stg_custom_rof_grid = Stage( + title="Runoff Grid", + description="From the below list of standard runoff grids, select one to be used as the " + "runoff grid within the new, custom CESM grid.", + widget=StageWidget(VBox), + parent=guard_custom_grid, + varlist=[cvars["CUSTOM_ROF_GRID"]], + auto_set_default_value=False, + ) + + stg_custom_rof_ocn_mapping = Stage( + title="Runoff to Ocean Mapping", + description="If the ocean model is MOM6, and unless there exists a standard mapping between" + "the selected runoff grid and the custom ocean grid, a new mapping must be created using " + "the mom6_bathy mapping module.", + widget=StageWidget( + VBox, + supplementary_widgets=[RunoffMappingGenerator(cime)] + ), + parent=Guard( + title="ROF to OCN Mapping", + parent=stg_custom_rof_grid, + condition=And(cvars["COMP_OCN"] == "mom", cvars["COMP_ROF"] != "srof") + ), + varlist=[cvars["ROF_OCN_MAPPING_STATUS"]], + ) diff --git a/visualCaseGen/widgets/grid_widgets.py b/visualCaseGen/widgets/grid_widgets.py index aeeeb05..8698e1c 100644 --- a/visualCaseGen/widgets/grid_widgets.py +++ b/visualCaseGen/widgets/grid_widgets.py @@ -25,13 +25,15 @@ def initialize_grid_widgets(cime): disabled=False, ) - initialize_standard_grid_widgets(cime) - initialize_custom_atm_grid_widgets(cime) - initialize_custom_ocn_grid_widgets(cime) - initialize_custom_lnd_grid_widgets(cime) - -def initialize_standard_grid_widgets(cime): - # Standard grid options + initialize_standard_grid_widgets() + initialize_custom_grid_path_widget(cime) + initialize_custom_atm_grid_widgets() + initialize_custom_ocn_grid_widgets() + initialize_custom_lnd_grid_widgets() + initialize_custom_rof_grid_widgets() + +def initialize_standard_grid_widgets(): + """Initialize the widgets for the standard grid options.""" cv_grid = cvars["GRID"] cv_grid.widget = MultiCheckbox( description="Grid:", @@ -39,8 +41,8 @@ def initialize_standard_grid_widgets(cime): ) cv_grid.valid_opt_char = chr(int("27A4", base=16)) -def initialize_custom_atm_grid_widgets(cime): - +def initialize_custom_grid_path_widget(cime): + """Initialize the widget for the custom grid path variable.""" default_path = Path.home() if cime.cime_output_root is not None: if (p := Path(cime.cime_output_root)).exists(): @@ -56,14 +58,16 @@ def initialize_custom_atm_grid_widgets(cime): layout={'width': '90%', 'margin': '10px'}, ) +def initialize_custom_atm_grid_widgets(): + """Initialize the widgets for the custom ATM grid options.""" cv_custom_atm_grid = cvars["CUSTOM_ATM_GRID"] cv_custom_atm_grid.widget = MultiCheckbox( description="Custom ATM Grid:", allow_multi_select=False, ) -def initialize_custom_ocn_grid_widgets(cime): - +def initialize_custom_ocn_grid_widgets(): + """Initialize the widgets for the custom OCN grid options.""" cv_custom_ocn_grid_mode = cvars["OCN_GRID_MODE"] cv_custom_ocn_grid_mode.widget = ToggleButtons( description="Ocean Grid Mode:", @@ -177,11 +181,11 @@ def initialize_custom_ocn_grid_widgets(cime): continuous_update=False, ) -def initialize_custom_lnd_grid_widgets(cime): +def initialize_custom_lnd_grid_widgets(): + """Initialize the widgets for the custom LND grid options.""" description_width = "250px" - cv_lnd_grid_mode = cvars["LND_GRID_MODE"] cv_lnd_grid_mode.widget = ToggleButtons( description="LND grid mode:", @@ -321,3 +325,29 @@ def initialize_custom_lnd_grid_widgets(cime): cv_fsurdat_mod_status = cvars["FSURDAT_MOD_STATUS"] cv_fsurdat_mod_status.widget = DisabledText(value='') + + +def initialize_custom_rof_grid_widgets(): + """Initialize the widgets for the custom ROF grid options.""" + cv_custom_rof_grid = cvars["CUSTOM_ROF_GRID"] + cv_custom_rof_grid.widget = MultiCheckbox( + description="Custom ROF Grid:", + allow_multi_select=False, + ) + + cv_rof_ocn_mapping_status = cvars["ROF_OCN_MAPPING_STATUS"] + cv_rof_ocn_mapping_status.widget = DisabledText(value='') + + cv_rof_ocn_mapping_rmax = cvars["ROF_OCN_MAPPING_RMAX"] + cv_rof_ocn_mapping_rmax.widget = Text( + description="Smoothing Rmax (km):", + layout={"width": "370px", "padding": "5px"}, + style={"description_width": "250px"}, + ) + + cv_rof_ocn_mapping_fold = cvars["ROF_OCN_MAPPING_FOLD"] + cv_rof_ocn_mapping_fold.widget = Text( + description="Smoothing Fold (km):", + layout={"width": "370px", "padding": "5px"}, + style={"description_width": "250px"}, + ) \ No newline at end of file