Merge Tests

regional_mom6/regional_mom6.py

@@ -23,6 +23,7 @@
 import copy
 from . import regridding as rgd
 from . import rotation as rot
+from . import rotation as rot
 
 warnings.filterwarnings("ignore")
 
@@ -1542,6 +1543,7 @@ def setup_ocean_state_boundaries(
         varnames,
         arakawa_grid="A",
         boundary_type="rectangular",
+        bathymetry_path=None,
         rotational_method=rot.RotationMethod.GIVEN_ANGLE,
     ):
         """
@@ -1558,6 +1560,8 @@ def setup_ocean_state_boundaries(
             arakawa_grid (Optional[str]): Arakawa grid staggering type of the boundary forcing.
                 Either ``'A'`` (default), ``'B'``, or ``'C'``.
             boundary_type (Optional[str]): Type of box around region. Currently, only ``'rectangular'`` is supported.
+            bathymetry_path (Optional[str]): Path to the bathymetry file. Default is None, in which case the bathymetry file is assumed to be in the input directory.
+            rotational_method (Optional[str]): Method to use for rotating the boundary velocities. Default is 'GIVEN_ANGLE'.
         """
         if boundary_type != "rectangular":
             raise ValueError(
@@ -1578,6 +1582,8 @@ def setup_ocean_state_boundaries(
             raise ValueError(
                 "This method only supports up to four boundaries. To set up more complex boundary shapes you can manually call the 'simple_boundary' method for each boundary."
             )
+        if bathymetry_path is None:
+            bathymetry_path = self.mom_input_dir / "bathymetry.nc"
 
         # Now iterate through our four boundaries
         for orientation in self.boundaries:
@@ -1593,6 +1599,7 @@ def setup_ocean_state_boundaries(
                     orientation
                 ),  # A number to identify the boundary; indexes from 1
                 arakawa_grid=arakawa_grid,
+                bathymetry_path=bathymetry_path,
                 rotational_method=rotational_method,
             )
 
@@ -1604,6 +1611,7 @@ def setup_single_boundary(
         segment_number,
         arakawa_grid="A",
         boundary_type="simple",
+        bathymetry_path=None,
         rotational_method=rot.RotationMethod.GIVEN_ANGLE,
     ):
         """
@@ -1624,6 +1632,8 @@ def setup_single_boundary(
             arakawa_grid (Optional[str]): Arakawa grid staggering type of the boundary forcing.
                 Either ``'A'`` (default), ``'B'``, or ``'C'``.
             boundary_type (Optional[str]): Type of boundary. Currently, only ``'simple'`` is supported. Here 'simple' refers to boundaries that are parallel to lines of constant longitude or latitude.
+            bathymetry_path (Optional[str]): Path to the bathymetry file. Default is None, in which case the bathymetry file is assumed to be in the input directory.
+            rotational_method (Optional[str]): Method to use for rotating the boundary velocities. Default is 'GIVEN_ANGLE'.
         """
 
         print("Processing {} boundary...".format(orientation), end="")
@@ -1635,6 +1645,7 @@ def setup_single_boundary(
             raise ValueError("Only simple boundaries are supported by this method.")
         self.segments[orientation] = segment(
             hgrid=self.hgrid,
+            bathymetry_path=bathymetry_path,
             infile=path_to_bc,  # location of raw boundary
             outfolder=self.mom_input_dir,
             varnames=varnames,
@@ -1658,6 +1669,7 @@ def setup_boundary_tides(
         tpxo_velocity_filepath,
         tidal_constituents="read_from_expt_init",
         boundary_type="rectangle",
+        bathymetry_path=None,
         rotational_method=rot.RotationMethod.GIVEN_ANGLE,
     ):
         """
@@ -1668,7 +1680,8 @@ def setup_boundary_tides(
             tidal_filename: Name of the tpxo product that's used in the tidal_filename. Should be h_tidal_filename, u_tidal_filename
             tidal_constituents: List of tidal constituents to include in the regridding. Default is [0] which is the M2 constituent.
             boundary_type (str): Type of boundary. Currently, only rectangle is supported. Here rectangle refers to boundaries that are parallel to lines of constant longitude or latitude.
-
+            bathymetry_path (str): Path to the bathymetry file. Default is None, in which case the bathymetry file is assumed to be in the input directory.
+            rotational_method (str): Method to use for rotating the tidal velocities. Default is 'GIVEN_ANGLE'.
         Returns:
             *.nc files: Regridded tidal velocity and elevation files in 'inputdir/forcing'
 
@@ -1693,6 +1706,8 @@ def setup_boundary_tides(
             )
         if tidal_constituents != "read_from_expt_init":
             self.tidal_constituents = tidal_constituents
+        if bathymetry_path is None:
+            bathymetry_path = self.mom_input_dir / "bathymetry.nc"
         tpxo_h = (
             xr.open_dataset(Path(tpxo_elevation_filepath))
             .rename({"lon_z": "lon", "lat_z": "lat", "nc": "constituent"})
@@ -1740,6 +1755,7 @@ def setup_boundary_tides(
             if b not in self.segments.keys():  # I.E. not set yet
                 seg = segment(
                     hgrid=self.hgrid,
+                    bathymetry_path=bathymetry_path,
                     infile=None,  # location of raw boundary
                     outfolder=self.mom_input_dir,
                     varnames=None,
@@ -2851,6 +2867,7 @@ def __init__(
         self,
         *,
         hgrid,
+        bathymetry_path,
         infile,
         outfolder,
         varnames,
@@ -2906,6 +2923,10 @@ def __init__(
         self.infile = infile
         self.outfolder = outfolder
         self.hgrid = hgrid
+        try:
+            self.bathymetry = xr.open_dataset(bathymetry_path)
+        except:
+            self.bathymetry = None
         self.segment_name = segment_name
         self.repeat_year_forcing = repeat_year_forcing
 
@@ -3197,6 +3218,9 @@ def regrid_velocity_tracers(self, rotational_method=rot.RotationMethod.GIVEN_ANG
         segment_out[f"{coords.attrs['parallel']}_{self.segment_name}"] = np.arange(
             segment_out[f"{coords.attrs['parallel']}_{self.segment_name}"].size
         )
+        segment_out[f"{coords.attrs['parallel']}_{self.segment_name}"] = np.arange(
+            segment_out[f"{coords.attrs['parallel']}_{self.segment_name}"].size
+        )
         segment_out[f"{coords.attrs['perpendicular']}_{self.segment_name}"] = [0]
         encoding_dict = {
             "time": {"dtype": "double"},
@@ -3379,7 +3403,6 @@ def regrid_tides(
             # Rotate
             INC -= np.radians(degree_angle.data[np.newaxis, :])
         ua, va, up, vp = ep2ap(SEMA, ECC, INC, PHA)
-
         # Convert to real amplitude and phase.
 
         ds_ap = xr.Dataset(
@@ -3460,6 +3483,9 @@ def encode_tidal_files_and_output(self, ds, filename):
             {"lon": f"lon_{self.segment_name}", "lat": f"lat_{self.segment_name}"}
         )
 
+        ds = rgd.mask_dataset(
+            ds, self.hgrid, self.bathymetry, self.orientation, self.segment_name
+        )
         ## Perform Encoding ##
 
         fname = f"{filename}_{self.segment_name}.nc"
@@ -3475,8 +3501,6 @@ def encode_tidal_files_and_output(self, ds, filename):
             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)
         ## Export Files ##
         ds.to_netcdf(
             Path(self.outfolder / "forcing" / fname),

regional_mom6/regridding.py

@@ -443,9 +443,188 @@ def generate_layer_thickness(
     return ds
 
 
+def get_boundary_mask(
+    hgrid: xr.Dataset,
+    bathy: xr.Dataset,
+    side: str,
+    segment_name: str,
+    minimum_depth=0,
+    x_dim_name="lonh",
+    y_dim_name="lath",
+) -> np.ndarray:
+    """
+    Mask out the boundary conditions based on the bathymetry. We don't want to have boundary conditions on land.
+    Parameters
+    ----------
+    hgrid : xr.Dataset
+        The hgrid dataset
+    bathy : xr.Dataset
+        The bathymetry dataset
+    side : str
+        The side of the boundary, "north", "south", "east", or "west"
+    segment_name : str
+        The segment name
+    minimum_depth : float, optional
+        The minimum depth to consider land, by default 0
+    Returns
+    -------
+    np.Array
+        The boundary mask
+    """
+
+    # Hide the bathy as an hgrid so we can take advantage of the coords function to get the boundary points.
+    try:
+        bathy = bathy.rename({y_dim_name: "nyp", x_dim_name: "nxp"})
+    except:
+        try:
+            bathy = bathy.rename({"ny": "nyp", "nx": "nxp"})
+        except:
+            regridding_logger.error("Could not rename bathy to nyp and nxp")
+            raise ValueError("Please provide the bathymetry x and y dimension names")
+
+    # Copy Hgrid
+    bathy_2 = hgrid.copy(deep=True)
+
+    # Create new depth field
+    bathy_2["depth"] = bathy_2["angle_dx"]
+    bathy_2["depth"][:, :] = np.nan
+
+    # Fill at t_points (what bathy is determined at)
+    ds_t = get_hgrid_arakawa_c_points(hgrid, "t")
+    bathy_2["depth"][ds_t.t_points_y.values, ds_t.t_points_x.values] = bathy.depth
+
+    bathy_2_coords = coords(
+        bathy_2,
+        side,
+        segment_name,
+        angle_variable_name="depth",
+        coords_at_t_points=True,
+    )
+
+    # Get the Boundary Depth
+    bathy_2_coords["boundary_depth"] = bathy_2_coords["angle"]
+    land = 0
+    ocean = 1.0
+    boundary_mask = np.full(np.shape(coords(hgrid, side, segment_name).angle), ocean)
+
+    ## Mask2DCu is the mask for the u/v points on the hgrid and is set to OBCmaskCy as well...
+    for i in range(len(bathy_2_coords["boundary_depth"])):
+        if bathy_2_coords["boundary_depth"][i] <= minimum_depth:
+            # The points to the left and right of this t-point are land points
+            boundary_mask[(i * 2) + 2] = land
+            boundary_mask[(i * 2) + 1] = (
+                land  # u/v point on the second level just like mask2DCu
+            )
+            boundary_mask[(i * 2)] = land
+
+    # Looks like in the boundary between land and ocean - in NWA for example - we basically need to remove 3 points closest to ocean as a buffer.
+    # Search for intersections
+    beaches_before = []
+    beaches_after = []
+    for index in range(1, len(boundary_mask) - 1):
+        if boundary_mask[index - 1] == land and boundary_mask[index] == ocean:
+            beaches_before.append(index)
+        elif boundary_mask[index + 1] == land and boundary_mask[index] == ocean:
+            beaches_after.append(index)
+    for beach in beaches_before:
+        for i in range(3):
+            if beach - 1 - i >= 0:
+                boundary_mask[beach - 1 - i] = ocean
+    for beach in beaches_after:
+        for i in range(3):
+            if beach + 1 + i < len(boundary_mask):
+                boundary_mask[beach + 1 + i] = ocean
+    boundary_mask[np.where(boundary_mask == land)] = np.nan
+
+    # Corner Q-points defined as wet
+    boundary_mask[0] = ocean
+    boundary_mask[-1] = ocean
+
+    return boundary_mask
+
+
+def mask_dataset(
+    ds: xr.Dataset,
+    hgrid: xr.Dataset,
+    bathymetry: xr.Dataset,
+    orientation,
+    segment_name: str,
+    y_dim_name="lath",
+    x_dim_name="lonh",
+) -> xr.Dataset:
+    """
+    This function masks the dataset to the provided bathymetry. If bathymetry is not provided, it fills all NaNs with 0.
+    Parameters
+    ----------
+    ds : xr.Dataset
+        The dataset to mask
+    hgrid : xr.Dataset
+        The hgrid dataset
+    bathymetry : xr.Dataset
+        The bathymetry dataset
+    orientation : str
+        The orientation of the boundary
+    segment_name : str
+        The segment name
+    """
+    ## Add Boundary Mask ##
+    if bathymetry is not None:
+        regridding_logger.info(
+            "Masking to bathymetry. If you don't want this, set bathymetry_path to None in the segment class."
+        )
+        mask = get_boundary_mask(
+            hgrid,
+            bathymetry,
+            orientation,
+            segment_name,
+            minimum_depth=0,
+            x_dim_name=x_dim_name,
+            y_dim_name=y_dim_name,
+        )
+        if orientation in ["east", "west"]:
+            mask = mask[:, np.newaxis]
+        else:
+            mask = mask[np.newaxis, :]
+
+        for var in ds.data_vars.keys():
+
+            ## Compare the dataset to the mask by reducing dims##
+            dataset_reduce_dim = ds[var]
+            for index in range(ds[var].ndim - 2):
+                dataset_reduce_dim = dataset_reduce_dim[0]
+            if orientation in ["east", "west"]:
+                dataset_reduce_dim = dataset_reduce_dim[:, 0]
+                mask_reduce = mask[:, 0]
+            else:
+                dataset_reduce_dim = dataset_reduce_dim[0, :]
+                mask_reduce = mask[0, :]
+            loc_nans_data = np.where(np.isnan(dataset_reduce_dim))
+            loc_nans_mask = np.where(np.isnan(mask_reduce))
+
+            ## Check if all nans in the data are in the mask without corners ##
+            if not np.isin(loc_nans_data[1:-1], loc_nans_mask[1:-1]).all():
+                regridding_logger.warning(
+                    f"NaNs in {var} not in mask. This values are filled with zeroes b/c they could cause issues with boundary conditions."
+                )
+
+                ## Remove Nans if needed ##
+                ds[var] = ds[var].fillna(0)
+
+            ## Apply the mask ##
+            ds[var] = ds[var] * mask
+    else:
+        regridding_logger.warning(
+            "All NaNs filled b/c bathymetry wasn't provided to the function. Add bathymetry_path to the segment class to avoid this"
+        )
+        ds = ds.fillna(
+            0
+        )  # Without bathymetry, we can't assume the nans will be allowed in Boundary Conditions
+    return ds
+
+
 def generate_encoding(
     ds: xr.Dataset, encoding_dict, default_fill_value=netCDF4.default_fillvals["f8"]
-) -> xr.Dataset:
+) -> dict:
     """
     Generate the encoding dictionary for the dataset
     Parameters