From d8b31979ce5af38440cec9ed66bbbf43e7877ab6 Mon Sep 17 00:00:00 2001 From: "Navid C. Constantinou" Date: Thu, 9 May 2024 14:03:29 +0300 Subject: [PATCH 1/8] attempt to simplify IC --- demos/reanalysis-forced.ipynb | 12 +----------- regional_mom6/regional_mom6.py | 21 +++++++++++++++++++-- 2 files changed, 20 insertions(+), 13 deletions(-) diff --git a/demos/reanalysis-forced.ipynb b/demos/reanalysis-forced.ipynb index 82607591..a84ba3fa 100644 --- a/demos/reanalysis-forced.ipynb +++ b/demos/reanalysis-forced.ipynb @@ -288,17 +288,7 @@ " glorys_path / \"ic_unprocessed.nc\", # directory where the unprocessed initial condition is stored, as defined earlier\n", " ocean_varnames,\n", " arakawa_grid=\"A\"\n", - " )\n", - "\n", - "# Now iterate through our four boundaries \n", - "for i, orientation in enumerate([\"south\", \"north\", \"west\", \"east\"]):\n", - " expt.rectangular_boundary(\n", - " glorys_path / (orientation + \"_unprocessed.nc\"),\n", - " ocean_varnames,\n", - " orientation, # Needs to know the cardinal direction of the boundary\n", - " i + 1, # Just a number to identify the boundary. Indexes from 1 \n", - " arakawa_grid=\"A\"\n", - " )" + " ) " ] }, { diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index a28e8d4b..0104c7ee 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -567,7 +567,12 @@ def _make_vgrid(self): return vcoord def initial_condition( - self, ic_path, varnames, arakawa_grid="A", vcoord_type="height" + self, + ic_path, + varnames, + arakawa_grid="A", + vcoord_type="height", + boundaries=["south", "north", "west", "east"], ): """ Reads the initial condition from files in ``ic_path``, interpolates to the @@ -874,11 +879,23 @@ def initial_condition( "eta_t": {"_FillValue": None}, }, ) - print("done setting up initial condition.") self.ic_eta = eta_out self.ic_tracers = tracers_out self.ic_vels = vel_out + + # Now iterate through our four boundaries + for i, orientation in enumerate(boundaries, start=1): + self.rectangular_boundary( + glorys_path / (orientation + "_unprocessed.nc"), + ocean_varnames, + orientation, # The cardinal direction of the boundary + i, # A number to identify the boundary; indexes from 1 + arakawa_grid=arakawa_grid, + ) + + print("done setting up initial condition.") + return def rectangular_boundary( From 5ebb005445ba79b3618861d121c54a814e2a6eb9 Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 16:49:38 +1000 Subject: [PATCH 2/8] add a wrapper function that iterates over four simple boundaries --- demos/reanalysis-forced.ipynb | 10 +++++- regional_mom6/regional_mom6.py | 61 +++++++++++++++++++++++++--------- 2 files changed, 54 insertions(+), 17 deletions(-) diff --git a/demos/reanalysis-forced.ipynb b/demos/reanalysis-forced.ipynb index a84ba3fa..72f56527 100644 --- a/demos/reanalysis-forced.ipynb +++ b/demos/reanalysis-forced.ipynb @@ -288,7 +288,15 @@ " glorys_path / \"ic_unprocessed.nc\", # directory where the unprocessed initial condition is stored, as defined earlier\n", " ocean_varnames,\n", " arakawa_grid=\"A\"\n", - " ) " + " ) \n", + "\n", + "# Set up the four boundary conditions. Remember that in the glorys_path, we have four boundary files names north_unprocessed.nc etc. \n", + "expt.rectangular_boundaries(\n", + " glorys_path,\n", + " ocean_varnames,\n", + " boundaries = [\"south\", \"north\", \"west\", \"east\"],\n", + " arakawa_grid = \"A\"\n", + " )" ] }, { diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 0104c7ee..ba208f58 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -568,18 +568,17 @@ def _make_vgrid(self): def initial_condition( self, - ic_path, + raw_ic_path, varnames, arakawa_grid="A", vcoord_type="height", - boundaries=["south", "north", "west", "east"], ): """ Reads the initial condition from files in ``ic_path``, interpolates to the model grid, fixes up metadata, and saves back to the input directory. Args: - ic_path (Union[str, Path]): Path to initial condition file. + raw_ic_path (Union[str, Path]): Path to raw initial condition file to read in. varnames (Dict[str, str]): Mapping from MOM6 variable/coordinate names to the names in the input dataset. For example, ``{'xq': 'lonq', 'yh': 'lath', 'salt': 'so', ...}``. arakawa_grid (Optional[str]): Arakawa grid staggering type of the initial condition. @@ -591,7 +590,7 @@ def initial_condition( # Remove time dimension if present in the IC. # Assume that the first time dim is the intended on if more than one is present - ic_raw = xr.open_dataset(ic_path) + ic_raw = xr.open_dataset(raw_ic_path) if varnames["time"] in ic_raw.dims: ic_raw = ic_raw.isel({varnames["time"]: 0}) if varnames["time"] in ic_raw.coords: @@ -766,7 +765,7 @@ def initial_condition( .bfill("lat") ) - ## Make our three horizontal regrideers + ## Make our three horizontal regridders regridder_u = xe.Regridder( ic_raw_u, ugrid, @@ -884,26 +883,53 @@ def initial_condition( self.ic_tracers = tracers_out self.ic_vels = vel_out + print("done setting up initial condition.") + + return + + def rectangular_boundaries( + self, + raw_boundaries_path, + varnames, + boundaries=["south", "north", "west", "east"], + arakawa_grid="A", + ): + """ + This function is a wrapper for `simple_boundary`. Given a list of up to four cardinal directions, + it creates a boundary forcing file for each one. Ensure that the raw boundaries are all saved in the same directory, + and that they are named using the format `east_unprocessed.nc` + + Args: + raw_boundaries_path (str): Path to the directory containing the raw boundary forcing files. + varnames (Dict[str, str]): Mapping from MOM6 variable/coordinate names to the name in the + input dataset. + boundaries (List[str]): List of cardinal directions for which to create boundary forcing files. + Default is `["south", "north", "west", "east"]`. + arakawa_grid (Optional[str]): Arakawa grid staggering type of the boundary forcing. + Either ``'A'`` (default), ``'B'``, or ``'C'``. + """ + for i in boundaries: + if i not in ["south", "north", "west", "east"]: + raise ValueError( + f"Invalid boundary direction: {i}. Must be one of ['south', 'north', 'west', 'east']" + ) + # Now iterate through our four boundaries for i, orientation in enumerate(boundaries, start=1): - self.rectangular_boundary( - glorys_path / (orientation + "_unprocessed.nc"), - ocean_varnames, + self.simple_boundary( + Path(raw_boundaries_path) / (orientation + "_unprocessed.nc"), + varnames, orientation, # The cardinal direction of the boundary i, # A number to identify the boundary; indexes from 1 arakawa_grid=arakawa_grid, ) - print("done setting up initial condition.") - - return - - def rectangular_boundary( + def simple_boundary( self, path_to_bc, varnames, orientation, segment_number, arakawa_grid="A" ): """ - Set up a boundary forcing file for a given orientation. Here the term 'rectangular' - means boundaries along lines of constant latitude or longitude. + Here 'simple' refers to boundaries that are parallel to lines of constant longitude or latitude. + Set up a boundary forcing file for a given orientation. Args: path_to_bc (str): Path to boundary forcing file. Ideally this should be a pre cut-out @@ -921,7 +947,10 @@ def rectangular_boundary( """ print("Processing {} boundary...".format(orientation), end="") - + if not path_to_bc.exists(): + raise FileNotFoundError( + f"Boundary file not found at {path_to_bc}. Please ensure that the files are named in the format `east_unprocessed.nc`." + ) seg = segment( hgrid=self.hgrid, infile=path_to_bc, # location of raw boundary From 9754e3c6c0b7c71a6d655f6acc9181e6ffffabc2 Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 16:53:11 +1000 Subject: [PATCH 3/8] black --- regional_mom6/regional_mom6.py | 1 - 1 file changed, 1 deletion(-) diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index ef008ee2..ceda46f3 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -890,7 +890,6 @@ def initial_condition( }, ) - self.ic_eta = eta_out self.ic_tracers = tracers_out self.ic_vels = vel_out From 7d850f8faddf80a468c2476ffc262050df4543e1 Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 16:59:42 +1000 Subject: [PATCH 4/8] additional error handling for new method --- regional_mom6/regional_mom6.py | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index ceda46f3..d70e60ef 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -925,6 +925,15 @@ def rectangular_boundaries( f"Invalid boundary direction: {i}. Must be one of ['south', 'north', 'west', 'east']" ) + if len(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" + ) + + if len(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." + ) # Now iterate through our four boundaries for i, orientation in enumerate(boundaries, start=1): self.simple_boundary( From 172c3aee04a9f7c66324d8418997988fd7bf2241 Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 17:01:43 +1000 Subject: [PATCH 5/8] update test function to match --- tests/test_expt_class.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index c484a355..e36b6443 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -355,7 +355,7 @@ def test_ocean_forcing( ), ], ) -def test_rectangular_boundary( +def test_rectangular_boundaries( longitude_extent, latitude_extent, date_range, @@ -471,4 +471,4 @@ def test_rectangular_boundary( "tracers": {"temp": "temp", "salt": "salt"}, } - expt.rectangular_boundary(tmp_path / "east_unprocessed", varnames, "east", 1) + expt.rectangular_boundaries(tmp_path, varnames) From e22a0a754d868d9830b6f708d3bae32b31bf3360 Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 17:07:01 +1000 Subject: [PATCH 6/8] only test east --- tests/test_expt_class.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index e36b6443..b2bead57 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -471,4 +471,4 @@ def test_rectangular_boundaries( "tracers": {"temp": "temp", "salt": "salt"}, } - expt.rectangular_boundaries(tmp_path, varnames) + expt.rectangular_boundaries(tmp_path, varnames,["east"]) From 1edf09ba58def2f010e05fe8cfddb4ecbd6c5e7b Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 17:08:32 +1000 Subject: [PATCH 7/8] black --- tests/test_expt_class.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index b2bead57..855437f3 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -471,4 +471,4 @@ def test_rectangular_boundaries( "tracers": {"temp": "temp", "salt": "salt"}, } - expt.rectangular_boundaries(tmp_path, varnames,["east"]) + expt.rectangular_boundaries(tmp_path, varnames, ["east"]) From 46791e314f30e5a7fa25d87512261267a706230f Mon Sep 17 00:00:00 2001 From: ashjbarnes Date: Tue, 11 Jun 2024 17:16:34 +1000 Subject: [PATCH 8/8] include .nc extension to test --- tests/test_expt_class.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index 855437f3..d0713c2f 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -443,7 +443,7 @@ def test_rectangular_boundaries( ), } ) - eastern_boundary.to_netcdf(tmp_path / "east_unprocessed") + eastern_boundary.to_netcdf(tmp_path / "east_unprocessed.nc") eastern_boundary.close() expt = experiment(