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 d93db5a2..daccde3e 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, @@ -3268,7 +3316,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, )