From 90e90a6aef31197d1d7728eedffbdc51254650d4 Mon Sep 17 00:00:00 2001 From: Manish Venumuddula <80477243+manishvenu@users.noreply.github.com> Date: Thu, 24 Oct 2024 16:02:49 -0600 Subject: [PATCH] Change Boundary Arguments to __init__ and adjust respective functions + tests (#30) * Package type changes for CRR * Black Formatting * Change MOM Param * Change MOM_param to Work * black * Revert back to reg * Add tide boundary flexibility * Update Config to only vars * Update Write and Read Config for computer independence * Update Read/Write Config for Computer Independenc * Fix Change MOM Param Issues * Fix bug in __getattr__ and comment out matplotlib dependencies in demo * Package type changes for CRR * Black Formatting * Add Boundaries to Init part 1 * black formatting * Changing Boundary Def to __init__ pt 2 * Minor Bugs * Make Direction Dir a one-time calculation because keys() doesn't gaurentee order * Rearrange handling boundaries in __init__ to a boundary array and segments dict in the boundary functions ocean_state and tides * Clean * Improve commenting on setuop_ocean_state_boundaries * Fix bug in declaring number of OBC segments * Minor Big in override * Add boundaries to config --- __init__.py | 0 demos/reanalysis-forced.ipynb | 11 ++- regional_mom6/regional_mom6.py | 140 +++++++++++++++++---------------- tests/test_config.py | 2 + tests/test_expt_class.py | 4 +- tests/test_manish_branch.py | 3 +- 6 files changed, 82 insertions(+), 78 deletions(-) create mode 100644 __init__.py diff --git a/__init__.py b/__init__.py new file mode 100644 index 00000000..e69de29b diff --git a/demos/reanalysis-forced.ipynb b/demos/reanalysis-forced.ipynb index a8195511..a4d161f2 100644 --- a/demos/reanalysis-forced.ipynb +++ b/demos/reanalysis-forced.ipynb @@ -138,7 +138,8 @@ " minimum_depth = 5,\n", " mom_run_dir = run_dir,\n", " mom_input_dir = input_dir,\n", - " toolpath_dir = toolpath_dir\n", + " toolpath_dir = toolpath_dir,\n", + " boundaries=[\"north\", \"south\", \"east\", \"west\"]\n", ")" ] }, @@ -191,8 +192,7 @@ "outputs": [], "source": [ "expt.get_glorys_rectangular(\n", - " raw_boundaries_path=glorys_path,\n", - " boundaries=[\"north\", \"south\", \"east\", \"west\"],\n", + " raw_boundaries_path=glorys_path\n", ")" ] }, @@ -285,7 +285,6 @@ "expt.setup_ocean_state_boundaries(\n", " glorys_path,\n", " ocean_varnames,\n", - " boundaries = [\"south\", \"north\", \"west\", \"east\"],\n", " arakawa_grid = \"A\"\n", " )" ] @@ -429,7 +428,7 @@ ], "metadata": { "kernelspec": { - "display_name": "vroom_clean_env", + "display_name": "crr_dev", "language": "python", "name": "python3" }, @@ -443,7 +442,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.6" + "version": "3.12.7" } }, "nbformat": 4, diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 9b8735e6..cc843959 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -73,34 +73,6 @@ def convert_to_tpxo_tidal_constituents(tidal_constituents): return list_of_ints -def find_MOM6_rectangular_orientation(input): - """ - Convert between MOM6 boundary and the specific segment number needed, or the inverse - """ - direction_dir = { - "south": 1, - "north": 2, - "west": 3, - "east": 4, - } - direction_dir_inv = {v: k for k, v in direction_dir.items()} - - if type(input) == str: - try: - return direction_dir[input] - except: - raise ValueError( - "Invalid Input. Did you spell the direction wrong, it should be lowercase?" - ) - elif type(input) == int: - try: - return direction_dir_inv[input] - except: - raise ValueError("Invalid Input. Did you pick a number 1 through 4?") - else: - raise ValueError("Invalid type of Input, can only be string or int.") - - ## Load Experiment Function @@ -168,6 +140,7 @@ def create_experiment_from_config( expt.layout = None expt.minimum_depth = config_dict["minimum_depth"] expt.tidal_constituents = config_dict["tidal_constituents"] + expt.boundaries = config_dict["boundaries"] if create_hgrid_and_vgrid: print("Creating hgrid and vgrid....") @@ -618,6 +591,7 @@ def create_empty( minimum_depth=4, tidal_constituents=["M2"], expt_name=None, + boundaries=["south", "north", "west", "east"], ): """ Substitute init method to creates an empty expirement object, with the opportunity to override whatever values wanted. @@ -659,6 +633,7 @@ def create_empty( expt.ocean_mask = None expt.layout = None self.segments = {} + self.boundaries = boundaries return expt def __init__( @@ -681,6 +656,7 @@ def __init__( tidal_constituents=["M2"], create_empty=False, expt_name=None, + boundaries=["south", "north", "west", "east"], ): # Creates empty experiment object for testing and experienced user manipulation. @@ -753,9 +729,8 @@ def __init__( else: self.vgrid = self._make_vgrid() - self.segments = ( - {} - ) # Holds segements for use in setting up the ocean state boundary conditions (GLORYS) and the tidal boundary conditions (TPXO) + self.segments = {} + self.boundaries = boundaries # create additional directories and links (self.mom_input_dir / "weights").mkdir(exist_ok=True) @@ -834,6 +809,33 @@ def __getattr__(self, name): error_message = f"{name} not found. Available methods and attributes are: {available_methods}" raise AttributeError(error_message) + def find_MOM6_rectangular_orientation(self, input): + """ + Convert between MOM6 boundary and the specific segment number needed, or the inverse + """ + + direction_dir = {} + counter = 1 + for b in self.boundaries: + direction_dir[b] = counter + counter += 1 + direction_dir_inv = {v: k for k, v in direction_dir.items()} + + if type(input) == str: + try: + return direction_dir[input] + except: + raise ValueError( + "Invalid Input. Did you spell the direction wrong, it should be lowercase?" + ) + elif type(input) == int: + try: + return direction_dir_inv[input] + except: + raise ValueError("Invalid Input. Did you pick a number 1 through 4?") + else: + raise ValueError("Invalid type of Input, can only be string or int.") + def _make_hgrid(self): """ Set up a horizontal grid based on user's specification of the domain. @@ -1088,6 +1090,7 @@ def write_config_file(self, path=None, export=True, quiet=False): "layout": self.layout, "minimum_depth": self.minimum_depth, "tidal_constituents": self.tidal_constituents, + "boundaries": self.boundaries, } if export: if path is not None: @@ -1448,9 +1451,7 @@ def setup_initial_condition( return - def get_glorys_rectangular( - self, raw_boundaries_path, boundaries=["south", "north", "west", "east"] - ): + def get_glorys_rectangular(self, raw_boundaries_path): """ This function is a wrapper for `get_glorys_data`, calling this function once for each of the rectangular boundary segments and the initial condition. For more complex boundary shapes, call `get_glorys_data` directly for each of your boundaries that aren't parallel to lines of constant latitude or longitude. For example, for an angled Northern boundary that spans multiple latitudes, you'll need to download a wider rectangle containing the entire boundary. @@ -1472,7 +1473,7 @@ def get_glorys_rectangular( download_path=raw_boundaries_path, modify_existing=False, # This is the first line, so start bash script anew ) - if "east" in boundaries: + if "east" in self.boundaries: get_glorys_data( longitude_extent=[ float(self.hgrid.x.isel(nxp=-1).min()), @@ -1486,7 +1487,7 @@ def get_glorys_rectangular( segment_name="east_unprocessed", download_path=raw_boundaries_path, ) - if "west" in boundaries: + if "west" in self.boundaries: get_glorys_data( longitude_extent=[ float(self.hgrid.x.isel(nxp=0).min()), @@ -1500,7 +1501,7 @@ def get_glorys_rectangular( segment_name="west_unprocessed", download_path=raw_boundaries_path, ) - if "south" in boundaries: + if "south" in self.boundaries: get_glorys_data( longitude_extent=[ float(self.hgrid.x.isel(nyp=0).min()), @@ -1514,7 +1515,7 @@ def get_glorys_rectangular( segment_name="south_unprocessed", download_path=raw_boundaries_path, ) - if "north" in boundaries: + if "north" in self.boundaries: get_glorys_data( longitude_extent=[ float(self.hgrid.x.isel(nyp=-1).min()), @@ -1538,7 +1539,6 @@ def setup_ocean_state_boundaries( self, raw_boundaries_path, varnames, - boundaries=["south", "north", "west", "east"], arakawa_grid="A", boundary_type="rectangular", ): @@ -1557,27 +1557,28 @@ def setup_ocean_state_boundaries( Either ``'A'`` (default), ``'B'``, or ``'C'``. boundary_type (Optional[str]): Type of box around region. Currently, only ``'rectangular'`` is supported. """ - for i in boundaries: + if boundary_type != "rectangular": + raise ValueError( + "Only rectangular boundaries are supported by this method. To set up more complex boundary shapes you can manually call the 'simple_boundary' method for each boundary." + ) + for i in self.boundaries: if i not in ["south", "north", "west", "east"]: raise ValueError( f"Invalid boundary direction: {i}. Must be one of ['south', 'north', 'west', 'east']" ) - if len(boundaries) < 4: + if len(self.boundaries) < 4: print( - "NOTE: the 'setup_run_directories' method assumes that you have four boundaries. You'll need to modify the MOM_input file manually to reflect the number of boundaries you have, and their orientations. You should be able to find the relevant section in the MOM_input file by searching for 'segment_'. Ensure that the segment names match those in your inputdir/forcing folder" + "NOTE: the 'setup_run_directories' method does understand the less than four boundaries but be careful. Please check the MOM_input/override file carefully to reflect the number of boundaries you have, and their orientations. You should be able to find the relevant section in the MOM_input/override file by searching for 'segment_'. Ensure that the segment names match those in your inputdir/forcing folder" ) - if len(boundaries) > 4: + if len(self.boundaries) > 4: raise ValueError( "This method only supports up to four boundaries. To set up more complex boundary shapes you can manually call the 'simple_boundary' method for each boundary." ) - if boundary_type != "rectangular": - raise ValueError( - "Only rectangular boundaries are supported by this method. To set up more complex boundary shapes you can manually call the 'simple_boundary' method for each boundary." - ) + # Now iterate through our four boundaries - for orientation in boundaries: + for orientation in self.boundaries: self.setup_single_boundary( Path( os.path.join( @@ -1586,7 +1587,7 @@ def setup_ocean_state_boundaries( ), varnames, orientation, # The cardinal direction of the boundary - find_MOM6_rectangular_orientation( + self.find_MOM6_rectangular_orientation( orientation ), # A number to identify the boundary; indexes from 1 arakawa_grid=arakawa_grid, @@ -1628,7 +1629,7 @@ def setup_single_boundary( ) if boundary_type != "simple": raise ValueError("Only simple boundaries are supported by this method.") - seg = segment( + self.segments[orientation] = segment( hgrid=self.hgrid, infile=path_to_bc, # location of raw boundary outfolder=self.mom_input_dir, @@ -1640,10 +1641,8 @@ def setup_single_boundary( repeat_year_forcing=self.repeat_year_forcing, ) - seg.regrid_velocity_tracers() + self.segments[orientation].regrid_velocity_tracers() - # Save Segment to Experiment - self.segments[orientation] = seg print("Done.") return @@ -1726,26 +1725,25 @@ def setup_boundary_tides( ), # Import pandas for this shouldn't be a big deal b/c it's already required in rm6 dependencies dims=["time"], ) - boundaries = ["south", "north", "west", "east"] - # Initialize or find boundary segment - for b in boundaries: + for b in self.boundaries: print("Processing {} boundary...".format(b), end="") # If the GLORYS ocean_state has already created segments, we don't create them again. - if b not in self.segments: + if b not in self.segments.keys(): # I.E. not set yet seg = segment( hgrid=self.hgrid, infile=None, # location of raw boundary outfolder=self.mom_input_dir, varnames=None, segment_name="segment_{:03d}".format( - find_MOM6_rectangular_orientation(b) + self.find_MOM6_rectangular_orientation(b) ), orientation=b, # orienataion startdate=self.date_range[0], repeat_year_forcing=self.repeat_year_forcing, ) + self.segments[b] = seg else: seg = self.segments[b] @@ -2190,7 +2188,6 @@ def setup_run_directory( using_payu=False, overwrite=False, with_tides=False, - boundaries=["south", "north", "west", "east"], ): """ Set up the run directory for MOM6. Either copy a pre-made set of files, or modify @@ -2213,6 +2210,7 @@ def setup_run_directory( / "demos" / "premade_run_directories" ) + if not premade_rundir_path.exists(): print("Could not find premade run directories at ", premade_rundir_path) print( @@ -2382,7 +2380,7 @@ def setup_run_directory( # Define number of OBC segments MOM_override_dict["OBC_NUMBER_OF_SEGMENTS"]["value"] = len( - boundaries + self.boundaries ) # This means that each SEGMENT_00{num} has to be configured to point to the right file, which based on our other functions needs to be specified. # More OBC Consts @@ -2396,18 +2394,18 @@ def setup_run_directory( MOM_override_dict["BRUSHCUTTER_MODE"]["value"] = "True" # Define Specific Segments - for ind, seg in enumerate(boundaries): - ind_seg = ind + 1 + for seg in self.boundaries: + ind_seg = self.find_MOM6_rectangular_orientation(seg) key_start = "OBC_SEGMENT_00" + str(ind_seg) ## Position and Config key_POSITION = key_start - if find_MOM6_rectangular_orientation(seg) == 1: + if seg == "south": index_str = '"J=0,I=0:N' - elif find_MOM6_rectangular_orientation(seg) == 2: + elif seg == "north": index_str = '"J=N,I=N:0' - elif find_MOM6_rectangular_orientation(seg) == 3: + elif seg == "west": index_str = '"I=0,J=N:0' - elif find_MOM6_rectangular_orientation(seg) == 4: + elif seg == "east": index_str = '"I=N,J=0:N' MOM_override_dict[key_POSITION]["value"] = ( index_str + ',FLATHER,ORLANSKI,NUDGED,ORLANSKI_TAN,NUDGED_TAN"' @@ -2420,7 +2418,7 @@ def setup_run_directory( # Data Key key_DATA = key_start + "_DATA" file_num_obc = str( - find_MOM6_rectangular_orientation(seg) + self.find_MOM6_rectangular_orientation(seg) ) # 1,2,3,4 for rectangular boundaries, BUT if we have less than 4 segments we use the index to specific the number, but keep filenames as if we had four boundaries MOM_override_dict[key_DATA][ "value" @@ -2509,7 +2507,7 @@ def setup_run_directory( return def change_MOM_parameter( - self, param_name, param_value=None, comment=None, delete=False + self, param_name, param_value=None, comment=None, override=True, delete=False ): """ *Requires already copied MOM parameter files in the run directory* @@ -2546,6 +2544,7 @@ def change_MOM_parameter( MOM_override_dict[param_name]["value"] = param_value MOM_override_dict[param_name]["comment"] = comment + MOM_override_dict[param_name]["override"] = override else: if param_name in MOM_override_dict.keys(): original_val = MOM_override_dict[param_name]["value"] @@ -2617,6 +2616,11 @@ def write_MOM_file(self, MOM_file_dict): if "#override" in var: var = var.replace("#override", "") var = var.strip() + else: + # As in there wasn't an override before but we want one + if MOM_file_dict[var]["override"]: + lines[jj] = "#override " + lines[jj] + print("Added override to variable " + var + "!") if var in MOM_file_dict.keys() and ( str(MOM_file_dict[var]["value"]) ) != str(original_MOM_file_dict[var]["value"]): diff --git a/tests/test_config.py b/tests/test_config.py index 6c2f35c9..005b685c 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -47,6 +47,7 @@ def test_write_config(): mom_input_dir=input_dir, toolpath_dir="", expt_name="test", + boundaries=["south", "north"], ) config_dict = expt.write_config_file() assert config_dict["longitude_extent"] == tuple(longitude_extent) @@ -62,6 +63,7 @@ def test_write_config(): assert config_dict["repeat_year_forcing"] == False assert config_dict["tidal_constituents"] == ["M2"] assert config_dict["expt_name"] == "test" + assert config_dict["boundaries"] == ["south", "north"] shutil.rmtree(run_dir) shutil.rmtree(input_dir) shutil.rmtree(data_path) diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index 2aa2465a..aff3a4b8 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -456,6 +456,7 @@ def test_rectangular_boundaries( mom_input_dir=mom_input_dir, toolpath_dir=toolpath_dir, hgrid_type=hgrid_type, + boundaries=["east"], ) varnames = { @@ -468,5 +469,4 @@ def test_rectangular_boundaries( "v": "v", "tracers": {"temp": "temp", "salt": "salt"}, } - - expt.setup_ocean_state_boundaries(tmp_path, varnames, ["east"]) + expt.setup_ocean_state_boundaries(tmp_path, varnames) diff --git a/tests/test_manish_branch.py b/tests/test_manish_branch.py index 63cc6de9..b2865512 100644 --- a/tests/test_manish_branch.py +++ b/tests/test_manish_branch.py @@ -9,7 +9,6 @@ from pathlib import Path import xarray as xr import numpy as np -from tests.test_expt_class import generate_silly_coords, number_of_gridpoints import shutil import importlib @@ -246,7 +245,7 @@ def test_tides(self, dummy_tidal_data): self.expt.hgrid_type = "even_spacing" # Dates self.expt.date_range = ("2000-01-01", "2000-01-02") - self.expt.segments = [] + self.expt.segments = {} # Generate Hgrid Data self.expt.resolution = 0.1 self.expt.hgrid = self.expt._make_hgrid()