Skip to content

Commit

Permalink
Change default tides to use all constituents in TPXO allowed in MOM6
Browse files Browse the repository at this point in the history
  • Loading branch information
@manishvenu
manishvenu committed Nov 7, 2024
1 parent df3b1d3 commit 7dc3cae
Show file tree
Hide file tree
Showing 3 changed files with 53 additions and 75 deletions.
110 changes: 44 additions & 66 deletions regional_mom6/regional_mom6.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,7 @@
import json
import copy
from . import regridding as rgd

warnings.filterwarnings("ignore")

__all__ = [
Expand Down Expand Up @@ -589,7 +590,7 @@ def create_empty(
hgrid_type="even_spacing",
repeat_year_forcing=False,
minimum_depth=4,
tidal_constituents=["M2"],
tidal_constituents=["M2", "S2", "N2", "K2", "K1", "O1", "P1", "Q1", "MM", "MF"],
expt_name=None,
boundaries=["south", "north", "west", "east"],
):
Expand Down Expand Up @@ -653,7 +654,7 @@ def __init__(
vgrid_type="hyperbolic_tangent",
repeat_year_forcing=False,
minimum_depth=4,
tidal_constituents=["M2"],
tidal_constituents=["M2", "S2", "N2", "K2", "K1", "O1", "P1", "Q1", "MM", "MF"],
create_empty=False,
expt_name=None,
boundaries=["south", "north", "west", "east"],
Expand Down Expand Up @@ -3013,8 +3014,8 @@ def rotate(self, u, v):

def regrid_velocity_tracers(self):
"""
Cut out and interpolate the velocities and tracers
"""
Cut out and interpolate the velocities and tracers
"""

rawseg = xr.open_dataset(self.infile, decode_times=False, engine="netcdf4")
coords = rgd.coords(self.hgrid, self.orientation, self.segment_name)
Expand All @@ -3032,22 +3033,17 @@ def regrid_velocity_tracers(self):

regridded = regridder(
rawseg[
[self.u, self.v, self.eta]
+ [self.tracers[i] for i in self.tracers]
[self.u, self.v, self.eta] + [self.tracers[i] for i in self.tracers]
]
)
rotated_u, rotated_v = self.rotate(
regridded[self.u], regridded[self.v]
)
rotated_u, rotated_v = self.rotate(regridded[self.u], regridded[self.v])
rotated_ds = xr.Dataset(
{
self.u: rotated_u,
self.v: rotated_v,
}
)
segment_out = xr.merge(
[rotated_ds, regridded.drop_vars([self.u, self.v])]
)
segment_out = xr.merge([rotated_ds, regridded.drop_vars([self.u, self.v])])

if self.arakawa_grid == "B":
## All tracers on one grid, all velocities on another
Expand All @@ -3058,19 +3054,14 @@ def regrid_velocity_tracers(self):
/ f"weights/bilinear_velocity_weights_{self.orientation}.nc",
)
regridder_tracer = rgd.create_regridder(
rawseg[self.tracers["salt"]].rename(
{self.xh: "lon", self.yh: "lat"}
),
rawseg[self.tracers["salt"]].rename({self.xh: "lon", self.yh: "lat"}),
coords,
self.outfolder
/ f"weights/bilinear_tracer_weights_{self.orientation}.nc",
)


velocities_out = regridder_velocity(
rawseg[[self.u, self.v]].rename(
{self.xq: "lon", self.yq: "lat"}
)
rawseg[[self.u, self.v]].rename({self.xq: "lon", self.yq: "lat"})
)

velocities_out["u"], velocities_out["v"] = self.rotate(
Expand All @@ -3097,7 +3088,6 @@ def regrid_velocity_tracers(self):
/ f"weights/bilinear_uvelocity_weights_{self.orientation}.nc",
)


regridder_vvelocity = rgd.create_regridder(
rawseg[self.v].rename({self.xh: "lon", self.yq: "lat"}),
coords,
Expand All @@ -3106,9 +3096,7 @@ def regrid_velocity_tracers(self):
)

regridder_tracer = rgd.create_regridder(
rawseg[self.tracers["salt"]].rename(
{self.xh: "lon", self.yh: "lat"}
),
rawseg[self.tracers["salt"]].rename({self.xh: "lon", self.yh: "lat"}),
coords,
self.outfolder
/ f"weights/bilinear_tracer_weights_{self.orientation}.nc",
Expand All @@ -3117,9 +3105,7 @@ def regrid_velocity_tracers(self):
regridded_u = regridder_uvelocity(rawseg[[self.u]])
regridded_v = regridder_vvelocity(rawseg[[self.v]])

rotated_u, rotated_v = self.rotate(
regridded_u[self.u], regridded_v[self.v]
)
rotated_u, rotated_v = self.rotate(regridded_u[self.u], regridded_v[self.v])
rotated_ds = xr.Dataset(
{
self.u: rotated_u,
Expand All @@ -3130,9 +3116,7 @@ def regrid_velocity_tracers(self):
[
rotated_ds,
regridder_tracer(
rawseg[
[self.eta] + [self.tracers[i] for i in self.tracers]
]
rawseg[[self.eta] + [self.tracers[i] for i in self.tracers]]
),
]
)
Expand All @@ -3144,9 +3128,7 @@ def regrid_velocity_tracers(self):
del segment_out["lat"]
## Convert temperatures to celsius # use pint
if (
np.nanmin(
segment_out[self.tracers["temp"]].isel({self.time: 0, self.z: 0})
)
np.nanmin(segment_out[self.tracers["temp"]].isel({self.time: 0, self.z: 0}))
> 100
):
segment_out[self.tracers["temp"]] -= 273.15
Expand All @@ -3157,21 +3139,23 @@ def regrid_velocity_tracers(self):
segment_out = rgd.fill_missing_data(
segment_out, f"{self.coords.attrs['parallel']}_{self.segment_name}"
)

times = xr.DataArray(np.arange(
0, #! Indexing everything from start of experiment = simple but maybe counterintutive?
segment_out[self.time].shape[
0
], ## Time is indexed from start date of window
dtype=float,
), # Import pandas for this shouldn't be a big deal b/c it's already kinda required somewhere deep in some tree.
dims=["time"])

times = xr.DataArray(
np.arange(
0, #! Indexing everything from start of experiment = simple but maybe counterintutive?
segment_out[self.time].shape[
0
], ## Time is indexed from start date of window
dtype=float,
), # Import pandas for this shouldn't be a big deal b/c it's already kinda required somewhere deep in some tree.
dims=["time"],
)
# This to change the time coordinate.
segment_out = rgd.add_or_update_time_dim(segment_out,times)
segment_out = rgd.add_or_update_time_dim(segment_out, times)
segment_out.time.attrs = {
"calendar": "julian",
"units": f"{self.time_units} since {self.startdate}",
}
"calendar": "julian",
"units": f"{self.time_units} since {self.startdate}",
}
# Here, keep in mind that 'var' keeps track of the mom6 variable names we want, and self.tracers[var]
# will return the name of the variable from the original data

Expand Down Expand Up @@ -3202,7 +3186,6 @@ def regrid_velocity_tracers(self):
segment_out, v, self.segment_name, self.z
)


## Treat eta separately since it has no vertical coordinate. Do the same things as for the surface variables above
segment_out = segment_out.rename({self.eta: f"eta_{self.segment_name}"})

Expand All @@ -3211,12 +3194,8 @@ def regrid_velocity_tracers(self):
)

# Overwrite the actual lat/lon values in the dimensions, replace with incrementing integers
segment_out[f"{self.coords.attrs['parallel']}_{self.segment_name}"] = (
np.arange(
segment_out[
f"{self.coords.attrs['parallel']}_{self.segment_name}"
].size
)
segment_out[f"{self.coords.attrs['parallel']}_{self.segment_name}"] = np.arange(
segment_out[f"{self.coords.attrs['parallel']}_{self.segment_name}"].size
)
segment_out[f"{self.coords.attrs['perpendicular']}_{self.segment_name}"] = [0]
encoding_dict = {
Expand All @@ -3229,7 +3208,9 @@ def regrid_velocity_tracers(self):
},
}
encoding_dict = rgd.generate_encoding(
segment_out, encoding_dict, default_fill_value=netCDF4.default_fillvals["f8"]
segment_out,
encoding_dict,
default_fill_value=netCDF4.default_fillvals["f8"],
)

segment_out.load().to_netcdf(
Expand Down Expand Up @@ -3282,9 +3263,7 @@ def regrid_tides(
tpxo_h[["lon", "lat", "hRe"]],
coords,
Path(
self.outfolder
/ "forcing"
/ f"regrid_{self.segment_name}_tidal_elev.nc"
self.outfolder / "forcing" / f"regrid_{self.segment_name}_tidal_elev.nc"
),
)

Expand Down Expand Up @@ -3328,7 +3307,7 @@ def regrid_tides(
"time", "constituent", f"{coords.attrs['parallel']}_{self.segment_name}"
)

self.encode_tidal_files_and_output( ds_ap, "tz")
self.encode_tidal_files_and_output(ds_ap, "tz")

########### Regrid Tidal Velocity ######################
regrid_u = rgd.create_regridder(tpxo_u[["lon", "lat", "uRe"]], coords, ".temp")
Expand Down Expand Up @@ -3359,16 +3338,13 @@ def regrid_tides(
ucplex = uredest + 1j * uimdest
vcplex = vredest + 1j * vimdest



angle = coords["angle"] # Fred's grid is in degrees

angle = coords["angle"] # Fred's grid is in degrees

# Convert complex u and v to ellipse,
# rotate ellipse from earth-relative to model-relative,
# and convert ellipse back to amplitude and phase.
SEMA, ECC, INC, PHA = ap2ep(ucplex, vcplex)

INC -= np.radians(angle.data[np.newaxis, :])
ua, va, up, vp = ep2ap(SEMA, ECC, INC, PHA)

Expand All @@ -3392,7 +3368,7 @@ def regrid_tides(
), # Import pandas for this shouldn't be a big deal b/c it's already kinda required somewhere deep in some tree.
dims=["time"],
)
ds_ap = rgd.add_or_update_time_dim(ds_ap,times)
ds_ap = rgd.add_or_update_time_dim(ds_ap, times)
ds_ap = ds_ap.transpose(
"time", "constituent", f"{coords.attrs['parallel']}_{self.segment_name}"
)
Expand All @@ -3402,7 +3378,7 @@ def regrid_tides(
ds_ap, f"{coords.attrs['parallel']}_{self.segment_name}"
)

self.encode_tidal_files_and_output( ds_ap, "tu")
self.encode_tidal_files_and_output(ds_ap, "tu")

return

Expand Down Expand Up @@ -3443,7 +3419,7 @@ def encode_tidal_files_and_output(self, ds, filename):
## Expand Tidal Dimensions ##

for var in ds:

ds = rgd.add_secondary_dimension(ds, str(var), coords, self.segment_name)

## Rename Tidal Dimensions ##
Expand All @@ -3462,7 +3438,9 @@ def encode_tidal_files_and_output(self, ds, filename):
f"lon_{self.segment_name}": dict(dtype="float64", _FillValue=1.0e20),
f"lat_{self.segment_name}": dict(dtype="float64", _FillValue=1.0e20),
}
encoding = rgd.generate_encoding(ds, encoding, default_fill_value=netCDF4.default_fillvals["f8"])
encoding = rgd.generate_encoding(
ds, encoding, default_fill_value=netCDF4.default_fillvals["f8"]
)

# Can't have nas in the land segments and such cuz it crashes
ds = ds.fillna(0)
Expand Down
16 changes: 7 additions & 9 deletions regional_mom6/regridding.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -193,7 +193,6 @@ def add_or_update_time_dim(ds: xr.Dataset, times) -> xr.Dataset:
# Make sure times is an xr.DataArray
times = xr.DataArray(times)


if "time" in ds.dims:
regridding_logger.debug("Time already in dataset, overwriting with new values")
ds["time"] = times
Expand Down Expand Up @@ -247,19 +246,18 @@ def add_secondary_dimension(
-------
xr.Dataset
The dataset with the perpendicular dimension added
"""

# Check if we need to insert the dim earlier or later
regridding_logger.info("Adding perpendicular dimension to {}".format(var))


regridding_logger.debug("Checking if nz or constituent is in dimensions, then we have to bump the perpendicular dimension up by one")
regridding_logger.debug(
"Checking if nz or constituent is in dimensions, then we have to bump the perpendicular dimension up by one"
)
insert_behind_by = 0
if any(
coord.startswith("nz") or coord == "constituent" for coord in ds[var].dims
):
if any(coord.startswith("nz") or coord == "constituent" for coord in ds[var].dims):
regridding_logger.debug("Bump it by one")
insert_behind_by = 0
else:
Expand Down
2 changes: 2 additions & 0 deletions regional_mom6/utils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@
import logging
import sys


def vecdot(v1, v2):
"""Return the dot product of vectors ``v1`` and ``v2``.
``v1`` and ``v2`` can be either numpy vectors or numpy.ndarrays
Expand Down Expand Up @@ -296,6 +297,7 @@ def ep2ap(SEMA, ECC, INC, PHA):

return ua, va, up, vp


def setup_logger(name: str) -> logging.Logger:
"""
Setup general config for a logger.
Expand Down

0 comments on commit 7dc3cae

Please sign in to comment.