Skip to content

Commit

Permalink
Merge branch 'main' into crr_changes
Browse files Browse the repository at this point in the history
  • Loading branch information
manishvenu committed Oct 14, 2024
2 parents 3298094 + 5ff128c commit 08caa9f
Show file tree
Hide file tree
Showing 3 changed files with 103 additions and 26 deletions.
36 changes: 35 additions & 1 deletion demos/reanalysis-forced.ipynb
Original file line number Diff line number Diff line change
Expand Up @@ -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!"
]
},
{
Expand Down Expand Up @@ -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": {},
Expand Down
86 changes: 64 additions & 22 deletions regional_mom6/regional_mom6.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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",
Expand Down Expand Up @@ -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):
Expand Down Expand Up @@ -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,
Expand Down Expand Up @@ -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)
Expand All @@ -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,
Expand Down Expand Up @@ -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(
Expand All @@ -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

Expand Down Expand Up @@ -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,
)
Expand Down
7 changes: 4 additions & 3 deletions tests/test_manish_branch.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
"""
Expand Down

0 comments on commit 08caa9f

Please sign in to comment.