diff --git a/demos/reanalysis-forced.ipynb b/demos/reanalysis-forced.ipynb index 35f7eabe..2fd34621 100644 --- a/demos/reanalysis-forced.ipynb +++ b/demos/reanalysis-forced.ipynb @@ -225,7 +225,9 @@ "cell_type": "markdown", "metadata": {}, "source": [ - "### Check out your domain:" + "### Check out your domain!\n", + "\n", + "Calling `expt.bathymetry` returns an xarray dataset, which can be plotted as usual. If you haven't yet run setup_bathymetry, calling `expt.bathymetry` will return `None` and prompt you to do so!" ] }, { @@ -309,6 +311,38 @@ " )" ] }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### Check out your initial condition data" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expt.init_tracers.salt.isel(zl = 0).plot()" + ] + }, + { + "cell_type": "markdown", + "metadata": {}, + "source": [ + "### You can plot your segment data too" + ] + }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "expt.segment_001.u_segment_001.isel(time = 5).plot()" + ] + }, { "cell_type": "markdown", "metadata": {}, diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 4e3e647f..c5f15dc7 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -14,13 +14,8 @@ import os import importlib.resources import datetime -from .utils import ( - quadrilateral_areas, - ap2ep, - ep2ap, -) +from .utils import quadrilateral_areas, ap2ep, ep2ap import pandas as pd -import re from pathlib import Path import glob from collections import defaultdict @@ -699,7 +694,7 @@ def __init__( depth, mom_run_dir, mom_input_dir, - toolpath_dir, + toolpath_dir=None, longitude_extent=None, latitude_extent=None, hgrid_type="even_spacing", @@ -800,12 +795,66 @@ def __str__(self) -> str: return json.dumps(self.write_config_file(export=False, quiet=True), indent=4) def __getattr__(self, name): + + ## First, check whether the attribute is an input file + + if name == "bathymetry": + if (self.mom_input_dir / "bathymetry.nc").exists(): + return xr.open_dataset( + self.mom_input_dir / "bathymetry.nc", + decode_cf=False, + decode_times=False, + ) + else: + print( + f"bathymetry.nc file not found! Make sure you've successfully run the setup_bathmetry method, or copied your own bathymetry.nc file into {self.mom_input_dir}." + ) + return None + elif name == "init_velocities": + if (self.mom_input_dir / "init_vel.nc").exists(): + return xr.open_dataset( + self.mom_input_dir / "init_vel.nc", + decode_cf=False, + decode_times=False, + ) + else: + print( + f"init_vel.nc file not found! Make sure you've successfully run the setup_initial_condition method, or copied your own init_vel.nc file into {self.mom_input_dir}." + ) + return + + elif name == "init_tracers": + if (self.mom_input_dir / "init_tracers.nc").exists(): + return xr.open_dataset( + self.mom_input_dir / "init_tracers.nc", + decode_cf=False, + decode_times=False, + ) + else: + print( + f"init_tracers.nc file not found! Make sure you've successfully run the setup_initial_condition method, or copied your own init_tracers.nc file into {self.mom_input_dir}." + ) + return + + elif "segment" in name: + try: + xr.open_mfdataset( + str(self.mom_input_dir / f"*{name}*.nc"), + decode_times=False, + decode_cf=False, + ) + except: + print( + f"{name} files not found! Make sure you've successfully run the setup_ocean_state_boundaries method, or copied your own segment files file into {self.mom_input_dir}." + ) + return None + + ## If we get here, attribute wasn't found + available_methods = [ method for method in dir(self) if not method.startswith("__") ] - error_message = ( - f"{name} method not found. Available methods are: {available_methods}" - ) + error_message = f"{name} not found. Available methods and attributes are: {available_methods}" raise AttributeError(error_message) def _make_hgrid(self): @@ -1067,7 +1116,6 @@ def write_config_file(self, path=None, export=True, quiet=False): "minimum_depth": self.minimum_depth, "vgrid": str(vgrid_path), "hgrid": str(hgrid_path), - "bathymetry": self.bathymetry_property, "ocean_state": self.ocean_state_boundaries, "tides": self.tides_boundaries, "initial_conditions": self.initial_condition, @@ -1849,7 +1897,7 @@ def setup_bathymetry( tgrid = xr.Dataset( data_vars={ "depth": ( - ["nx", "ny"], + ["ny", "nx"], np.zeros( self.hgrid.x.isel( nxp=slice(1, None, 2), nyp=slice(1, None, 2) @@ -1859,13 +1907,13 @@ def setup_bathymetry( }, coords={ "lon": ( - ["nx", "ny"], + ["ny", "nx"], self.hgrid.x.isel( nxp=slice(1, None, 2), nyp=slice(1, None, 2) ).values, ), "lat": ( - ["nx", "ny"], + ["ny", "nx"], self.hgrid.y.isel( nxp=slice(1, None, 2), nyp=slice(1, None, 2) ).values, @@ -2520,15 +2568,10 @@ def change_MOM_parameter( "If not deleting a parameter, you must specify a new value for it." ) - MOM_input_dict = self.read_MOM_file_as_dict("MOM_input") MOM_override_dict = self.read_MOM_file_as_dict("MOM_override") original_val = "No original val" if not delete: - # We don't want to keep any parameters in MOM_input that we change. We want to clearly list them in MOM_override. - if param_name in MOM_input_dict.keys(): - original_val = MOM_override_dict[param_name]["value"] - print("Removing original value {} from MOM_input".format(original_val)) - del MOM_input_dict[param_name] + if param_name in MOM_override_dict.keys(): original_val = MOM_override_dict[param_name]["value"] print( @@ -2550,7 +2593,6 @@ def change_MOM_parameter( param_name ) ) - self.write_MOM_file(MOM_input_dict) self.write_MOM_file(MOM_override_dict) return original_val @@ -3272,7 +3314,7 @@ def regrid_velocity_tracers(self): ) segment_out[f"lat_{self.segment_name}"] = ( [f"ny_{self.segment_name}", f"nx_{self.segment_name}"], - self.coords.lon.expand_dims( + self.coords.lat.expand_dims( dim="blank", axis=self.coords.attrs["axis_to_expand"] - 2 ).data, ) diff --git a/tests/test_manish_branch.py b/tests/test_manish_branch.py index aa7ecb8f..0c293e61 100644 --- a/tests/test_manish_branch.py +++ b/tests/test_manish_branch.py @@ -282,10 +282,11 @@ def test_change_MOM_parameter(self): shutil.copytree( base_run_dir / "common_files", self.expt.mom_run_dir, dirs_exist_ok=True ) - self.expt.change_MOM_parameter("OBC_SEGMENT_001", "adasd", "COOL COMMENT") + og = self.expt.change_MOM_parameter("MINIMUM_DEPTH", "adasd", "COOL COMMENT") MOM_override_dict = self.expt.read_MOM_file_as_dict("MOM_override") - assert MOM_override_dict["OBC_SEGMENT_001"]["value"] == "adasd" - assert MOM_override_dict["OBC_SEGMENT_001"]["comment"] == "COOL COMMENT\n" + assert MOM_override_dict["MINIMUM_DEPTH"]["value"] == "adasd" + assert MOM_override_dict["original"]["OBC_SEGMENT_001"]["value"] == og + assert MOM_override_dict["MINIMUM_DEPTH"]["comment"] == "COOL COMMENT\n" def test_properties_empty(self): """