Skip to content

Commit

Permalink
Merge pull request #39 from CROCODILE-CESM/visualCaseGen_integration
Browse files Browse the repository at this point in the history 
Minor changes ahead of CrocoDash-visualCaseGen integration
  • Loading branch information
@alperaltuntas
alperaltuntas authored Jan 7, 2025
2 parents 90e90a6 + 9665892 commit 53ad69e
Show file tree
Hide file tree
Showing 2 changed files with 85 additions and 33 deletions.
2 changes: 2 additions & 0 deletions pyproject.toml
View file
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ dependencies = [
"xarray <= 2024.7.0",
"xesmf >= 0.8.4",
"f90nml >= 1.4.1",
"aiohttp >= 3.9.5,<3.10.0",
"copernicusmarine >= 1.2.4,<1.3.0"
]

[build-system]
Expand Down
116 changes: 83 additions & 33 deletions regional_mom6/regional_mom6.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -575,7 +575,7 @@ class experiment:

@classmethod
def create_empty(
self,
cls,
longitude_extent=None,
latitude_extent=None,
date_range=None,
Expand All @@ -596,7 +596,7 @@ def create_empty(
"""
Substitute init method to creates an empty expirement object, with the opportunity to override whatever values wanted.
"""
expt = self(
expt = cls(
longitude_extent=None,
latitude_extent=None,
date_range=None,
Expand Down Expand Up @@ -632,8 +632,8 @@ def create_empty(
expt.longitude_extent = longitude_extent
expt.ocean_mask = None
expt.layout = None
self.segments = {}
self.boundaries = boundaries
cls.segments = {}
cls.boundaries = boundaries
return expt

def __init__(
Expand All @@ -650,7 +650,9 @@ def __init__(
longitude_extent=None,
latitude_extent=None,
hgrid_type="even_spacing",
hgrid_path=None,
vgrid_type="hyperbolic_tangent",
vgrid_path=None,
repeat_year_forcing=False,
minimum_depth=4,
tidal_constituents=["M2"],
Expand All @@ -673,7 +675,7 @@ def __init__(

self.mom_run_dir = Path(mom_run_dir)
self.mom_input_dir = Path(mom_input_dir)
self.toolpath_dir = Path(toolpath_dir)
self.toolpath_dir = Path(toolpath_dir) if toolpath_dir is not None else None

self.mom_run_dir.mkdir(exist_ok=True)
self.mom_input_dir.mkdir(exist_ok=True)
Expand All @@ -695,8 +697,12 @@ def __init__(
self.tidal_constituents = tidal_constituents

if hgrid_type == "from_file":
if hgrid_path is None:
hgrid_path = self.mom_input_dir / "hgrid.nc"
else:
hgrid_path = Path(hgrid_path)
try:
self.hgrid = xr.open_dataset(self.mom_input_dir / "hgrid.nc")
self.hgrid = xr.open_dataset(hgrid_path)
self.longitude_extent = (
float(self.hgrid.x.min()),
float(self.hgrid.x.max()),
Expand All @@ -712,21 +718,35 @@ def __init__(
)
raise ValueError
else:
if hgrid_path:
raise ValueError(
"hgrid_path can only be set if hgrid_type is 'from_file'."
)
self.longitude_extent = tuple(longitude_extent)
self.latitude_extent = tuple(latitude_extent)
self.hgrid = self._make_hgrid()

if vgrid_type == "from_file":
if vgrid_path is None:
vgrid_path = self.mom_input_dir / "vgrid.nc"
else:
vgrid_path = Path(vgrid_path)

try:
self.vgrid = xr.open_dataset(self.mom_input_dir / "vcoord.nc")
vgrid_from_file = xr.open_dataset(vgrid_path)

except:
print(
"Error while reading in existing vertical coordinates!\n\n"
+ f"Make sure `vcoord.nc`exists in {self.mom_input_dir} directory."
)
raise ValueError
self.vgrid = self._make_vgrid(vgrid_from_file.dz.data)
else:
if vgrid_path:
raise ValueError(
"vgrid_path can only be set if vgrid_type is 'from_file'."
)
self.vgrid = self._make_vgrid()

self.segments = {}
Expand Down Expand Up @@ -906,17 +926,18 @@ def _make_hgrid(self):

return hgrid

def _make_vgrid(self):
def _make_vgrid(self, thicknesses=None):
"""
Generates a vertical grid based on the ``number_vertical_layers``, the ratio
of largest to smallest layer thickness (``layer_thickness_ratio``) and the
total ``depth`` parameters.
(All these parameters are specified at the class level.)
"""

thicknesses = hyperbolictan_thickness_profile(
self.number_vertical_layers, self.layer_thickness_ratio, self.depth
)
if thicknesses is None:
thicknesses = hyperbolictan_thickness_profile(
self.number_vertical_layers, self.layer_thickness_ratio, self.depth
)

zi = np.cumsum(thicknesses)
zi = np.insert(zi, 0, 0.0) # add zi = 0.0 as first interface
Expand Down Expand Up @@ -1531,7 +1552,15 @@ def get_glorys_rectangular(self, raw_boundaries_path):
)

print(
f"script `get_glorys_data.sh` has been created at {raw_boundaries_path}.\n Run this script via bash to download the data from a terminal with internet access. \nYou will need to enter your Copernicus Marine username and password.\nIf you don't have an account, make one here:\nhttps://data.marine.copernicus.eu/register"
f"The script `get_glorys_data.sh` has been generated at:\n {raw_boundaries_path}.\n"
f"To download the data, run this script using `bash` in a terminal with internet access.\n\n"
f"Important instructions:\n"
f"1. You will need your Copernicus Marine username and password.\n"
f" If you do not have an account, you can create one here: \n"
f" https://data.marine.copernicus.eu/register\n"
f"2. You will be prompted to enter your Copernicus Marine credentials multiple times: once for each dataset.\n"
f"3. Depending on the dataset size, the download process may take significant time and resources.\n"
f"4. Thus, on certain systems, you may need to run this script as a batch job.\n"
)
return

Expand Down Expand Up @@ -1760,6 +1789,7 @@ def setup_bathymetry(
vertical_coordinate_name="elevation", # This is to match GEBCO
fill_channels=False,
positive_down=False,
write_to_file=True,
):
"""
Cut out and interpolate the chosen bathymetry and then fill inland lakes.
Expand All @@ -1783,6 +1813,7 @@ def setup_bathymetry(
but can also connect extra islands to land. Default: ``False``.
positive_down (Optional[bool]): If ``True``, it assumes that
bathymetry vertical coordinate is positive down. Default: ``False``.
write_to_file (Optional[bool]): Whether to write the bathymetry to a file. Default: ``True``.
"""

## Convert the provided coordinate names into a dictionary mapping to the
Expand Down Expand Up @@ -1856,9 +1887,12 @@ def setup_bathymetry(
)
bathymetry_output.depth.attrs["long_name"] = "Elevation relative to sea level"
bathymetry_output.depth.attrs["coordinates"] = "lon lat"
bathymetry_output.to_netcdf(
self.mom_input_dir / "bathymetry_original.nc", mode="w", engine="netcdf4"
)
if write_to_file:
bathymetry_output.to_netcdf(
self.mom_input_dir / "bathymetry_original.nc",
mode="w",
engine="netcdf4",
)

tgrid = xr.Dataset(
data_vars={
Expand Down Expand Up @@ -1894,10 +1928,13 @@ def setup_bathymetry(
tgrid.lat.attrs["_FillValue"] = 1e20
tgrid.depth.attrs["units"] = "meters"
tgrid.depth.attrs["coordinates"] = "lon lat"
tgrid.to_netcdf(
self.mom_input_dir / "bathymetry_unfinished.nc", mode="w", engine="netcdf4"
)
tgrid.close()
if write_to_file:
tgrid.to_netcdf(
self.mom_input_dir / "bathymetry_unfinished.nc",
mode="w",
engine="netcdf4",
)
tgrid.close()

bathymetry_output = bathymetry_output.load()

Expand All @@ -1914,19 +1951,27 @@ def setup_bathymetry(
)
regridder = xe.Regridder(bathymetry_output, tgrid, "bilinear", parallel=False)
bathymetry = regridder(bathymetry_output)
bathymetry.to_netcdf(
self.mom_input_dir / "bathymetry_unfinished.nc", mode="w", engine="netcdf4"
)
if write_to_file:
bathymetry.to_netcdf(
self.mom_input_dir / "bathymetry_unfinished.nc",
mode="w",
engine="netcdf4",
)
print(
"Regridding successful! Now calling `tidy_bathymetry` method for some finishing touches..."
)

self.tidy_bathymetry(fill_channels, positive_down)
self.tidy_bathymetry(fill_channels, positive_down, bathymetry=bathymetry)
print("setup bathymetry has finished successfully.")
return

return bathymetry

def tidy_bathymetry(
self, fill_channels=False, positive_down=False, vertical_coordinate_name="depth"
self,
fill_channels=False,
positive_down=False,
vertical_coordinate_name="depth",
bathymetry=None,
):
"""
An auxiliary function for bathymetry used to fix up the metadata and remove inland
Expand All @@ -1944,16 +1989,20 @@ def tidy_bathymetry(
but can also connect extra islands to land. Default: ``False``.
positive_down (Optional[bool]): If ``False`` (default), assume that
bathymetry vertical coordinate is positive down, as is the case in GEBCO for example.
bathymetry (Optional[xr.Dataset]): The bathymetry dataset to tidy up. If not provided,
it will read the bathymetry from the file ``bathymetry_unfinished.nc`` in the input directory
that was created by :func:`~setup_bathymetry`.
"""

## reopen bathymetry to modify
print(
"Tidy bathymetry: Reading in regridded bathymetry to fix up metadata...",
end="",
)
bathymetry = xr.open_dataset(
self.mom_input_dir / "bathymetry_unfinished.nc", engine="netcdf4"
)
if read_bathy_from_file := bathymetry is None:
bathymetry = xr.open_dataset(
self.mom_input_dir / "bathymetry_unfinished.nc", engine="netcdf4"
)

## Ensure correct encoding
bathymetry = xr.Dataset(
Expand Down Expand Up @@ -2115,11 +2164,12 @@ def tidy_bathymetry(
~(bathymetry.depth <= self.minimum_depth), self.minimum_depth + 0.1
)

bathymetry.expand_dims({"ntiles": 1}).to_netcdf(
self.mom_input_dir / "bathymetry.nc",
mode="w",
encoding={"depth": {"_FillValue": None}},
)
if read_bathy_from_file:
bathymetry.expand_dims({"ntiles": 1}).to_netcdf(
self.mom_input_dir / "bathymetry.nc",
mode="w",
encoding={"depth": {"_FillValue": None}},
)

print("done.")
return
Expand Down

0 comments on commit 53ad69e

Please sign in to comment.