diff --git a/.github/workflows/testing.yml b/.github/workflows/testing.yml index feb203e8..efe9811b 100644 --- a/.github/workflows/testing.yml +++ b/.github/workflows/testing.yml @@ -18,7 +18,7 @@ jobs: testing: needs: formatting runs-on: ubuntu-latest - container: ghcr.io/cosima/regional-test-env:updated + container: ghcr.io/crocodile-cesm/crocodile_rm6_test_env:latest defaults: run: shell: bash -el {0} diff --git a/docs/angle_calc.md b/docs/angle_calc.md new file mode 100644 index 00000000..5d8a3e91 --- /dev/null +++ b/docs/angle_calc.md @@ -0,0 +1,45 @@ +# Rotation and angle calculation in RM6 using MOM6 Angle Calculation +This document explains the implementation of MOM6 angle calculation in RM6, which is the process by which RM6 calculates the angle of curved hgrids. + +**Issue:** MOM6 doesn't actually use the user-provided "angle_dx" field in input hgrids, but internally calculates the angle. + +**Solution:** To accomodate this fact, when we rotate our boundary conditions, we implemented MOM6 angle calculation in a file called "rotation.py", and adjusted functions where we regrid the boundary conditions. + + +## MOM6 process of angle calculation (T-point only) +1. Calculate pi/4rads / 180 degrees = Gives a 1/4 conversion of degrees to radians. I.E. multiplying an angle in degrees by this gives the conversion to radians at 1/4 the value. +2. Figure out the longitudunal extent of our domain, or periodic range of longitudes. For global cases it is len_lon = 360, for our regional cases it is given by the hgrid. +3. At each point on our hgrid, we find the q-point to the top left, bottom left, bottom right, top right. We adjust each of these longitudes to be in the range of len_lon around the point itself. (module_around_point) +4. We then find the lon_scale, which is the "trigonometric scaling factor converting changes in longitude to equivalent distances in latitudes". Whatever that actually means is we add the latitude of all four of these points from part 3 and basically average it and convert to radians. We then take the cosine of it. As I understand it, it's a conversion of longitude to equivalent latitude distance. +5. Then we calculate the angle. This is a simple arctan2 so y/x. + 1. The "y" component is the addition of the difference between the diagonals in longitude (adjusted by modulo_around_point in step 3) multiplied by the lon_scale, which is our conversion to latitude. + 2. The "x" component is the same addition of differences in latitude. + 3. Thus, given the same units, we can call arctan to get the angle in degrees + + +## Problem +MOM6 only calculates the angle at t-points. For boundary rotation, we need the angle at the boundary, which is q/u/v points. Because we need the points to the left, right, top, and bottom of the point, this method won't work for the boundary. + + +# Convert this method to boundary angles - 3 Options +1. **GIVEN_ANGLE**: Don't calculate the angle and use the user-provided field in the hgrid called "angle_dx" +2. **EXPAND_GRID**: Calculate another boundary row/column points around the hgrid using simple difference techniques. Use the new points to calculate the angle at the boundaries. This works because we can now access the four points needed to calculate the angle, where previously at boundaries we would be missing at least two. + + +## Code Description + +Most calculation code is implemented in the rotation.py script, and the functional uses are in regrid_velocity_tracers and regrid_tides functions in the segment class of RM6. + + +### Calculation Code (rotation.py) +1. **Rotational Method Definition**: Rotational Methods are defined in the enum class "Rotational Method" in rotation.py. +2. **MOM6 Angle Calculation**: The method is implemented in "mom6_angle_calculation_method" in rotation.py and the direct t-point angle calculation is "initialize_grid_rotation_angle". +3. **Fred's Pseudo Grid Expansion**: The method to add the additional boundary row/columns is referenced in "pseudo_hgrid" functions in rotation.py + +### Implementation Code (regional_mom6.py) +Both regridding functions (regrid_velocity_tracers, regrid_tides) accept a parameter called "rotational_method" which takes the Enum class defining the rotational method. + +We then define each method with a bunch of if statements. Here are the processes: + +1. Given angle is the default method of accepting the hgrid's angle_dx +2. Fred's method is the least code, and we simply swap out the hgrid angle with the generated one we calculate right where we do the rotation. diff --git a/docs/api.rst b/docs/api.rst index 3db049aa..93087763 100644 --- a/docs/api.rst +++ b/docs/api.rst @@ -1,22 +1,45 @@ -=============== - API reference -=============== +regional\_mom6 package +====================== +Submodules +---------- -+++++++++++++++++++ - ``regional_mom6`` -+++++++++++++++++++ +regional\_mom6.regional\_mom6 module +------------------------------------ .. automodule:: regional_mom6.regional_mom6 :members: :undoc-members: - :private-members: + :show-inheritance: +regional\_mom6.regridding module +-------------------------------- -+++++++++++ - ``utils`` -+++++++++++ +.. automodule:: regional_mom6.regridding + :members: + :undoc-members: + :show-inheritance: + +regional\_mom6.rotation module +------------------------------ + +.. automodule:: regional_mom6.rotation + :members: + :undoc-members: + :show-inheritance: + +regional\_mom6.utils module +--------------------------- .. automodule:: regional_mom6.utils :members: :undoc-members: + :show-inheritance: + +Module contents +--------------- + +.. automodule:: regional_mom6 + :members: + :undoc-members: + :show-inheritance: diff --git a/docs/docker_image_dev.md b/docs/docker_image_dev.md new file mode 100644 index 00000000..5c20030a --- /dev/null +++ b/docs/docker_image_dev.md @@ -0,0 +1,35 @@ +# Docker Image & Github Testing (For contributors) + +RM6 uses a docker image in github actions for holding large data. It wasn't directly being used, but for downloading the curvilinear grid for testing, we are using it. This document is a list of helpful commands to work on it. + +The link to the image is here: +https://github.com/COSIMA/regional-mom6/pkgs/container/regional-test-env + +For local development of the image to add data to it for testing, first pull it. +```docker pull ghcr.io/cosima/regional-test-env:updated``` + +Then to do testing of the image, we cd into our cloned version of RM6, and run this command. It mounts our code in the /workspace directory.: +```docker run -it --rm \ -v $(pwd):/workspace \ -w /workspace \ ghcr.io/cosima/regional-test-env:updated \ /bin/bash``` + +The -it flag is for shell access, and the workspace stuff is to get our local code in the container. +You have to download conda, python, pip, and all that business to properly run the tests. + +Getting to adding the data, you should create a folder and add both the data you want to add and a file simple called "Dockerfile". In Dockerfile, we'll get the original image, then copy the data we need to the data folder. + +``` +# Use the base image +FROM ghcr.io/cosima/regional-test-env: + +# Copy your local file into the /data directory in the container +COPY /data/ +``` + +Then, we need to build the image, tag it, and push it + +``` +docker build -t my-custom-image . # IN THE DIRECTORY WITH THE DOCKERFILE +docker tag my-custom-image ghcr.io/cosima/regional-test-env: +docker push ghcr.io/cosima/regional-test-env: +``` + + diff --git a/docs/index.rst b/docs/index.rst index 531ad99c..d8f614ca 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -91,8 +91,10 @@ The bibtex entry for the paper is: installation demos mom6-file-structure-primer + angle_calc api contributing + docker_image_dev Indices and tables diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 4b621e18..1232910a 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -14,7 +14,7 @@ import os import importlib.resources import datetime -from .utils import quadrilateral_areas, ap2ep, ep2ap +from .utils import quadrilateral_areas, ap2ep, ep2ap, is_rectilinear_hgrid import pandas as pd from pathlib import Path import glob @@ -22,6 +22,7 @@ import json import copy from . import regridding as rgd +from . import rotation as rot warnings.filterwarnings("ignore") @@ -694,7 +695,6 @@ def __init__( self.layout = None # This should be a tuple. Leaving in a dummy 'None' makes it easy to remind the user to provide a value later on. self.minimum_depth = minimum_depth # Minimum depth allowed in bathy file self.tidal_constituents = tidal_constituents - if hgrid_type == "from_file": try: self.hgrid = xr.open_dataset(self.mom_input_dir / "hgrid.nc") @@ -1542,6 +1542,7 @@ def setup_ocean_state_boundaries( varnames, arakawa_grid="A", boundary_type="rectangular", + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ This function is a wrapper for `simple_boundary`. Given a list of up to four cardinal directions, @@ -1592,6 +1593,7 @@ def setup_ocean_state_boundaries( orientation ), # A number to identify the boundary; indexes from 1 arakawa_grid=arakawa_grid, + rotational_method=rotational_method, ) def setup_single_boundary( @@ -1602,6 +1604,7 @@ def setup_single_boundary( segment_number, arakawa_grid="A", boundary_type="simple", + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ Here 'simple' refers to boundaries that are parallel to lines of constant longitude or latitude. @@ -1642,7 +1645,9 @@ def setup_single_boundary( repeat_year_forcing=self.repeat_year_forcing, ) - self.segments[orientation].regrid_velocity_tracers() + self.segments[orientation].regrid_velocity_tracers( + rotational_method=rotational_method + ) print("Done.") return @@ -1653,16 +1658,17 @@ def setup_boundary_tides( tpxo_velocity_filepath, tidal_constituents="read_from_expt_init", boundary_type="rectangle", + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ - This function: We subset our tidal data and generate more boundary files! Args: path_to_td (str): Path to boundary tidal file. - tidal_filename: Name of the tpxo product that's used in the tidal_filename. Should be h_{tidal_filename}, u_{tidal_filename} - tidal_constiuents: List of tidal constituents to include in the regridding. Default is [0] which is the M2 constituent. - boundary_type (Optional[str]): Type of boundary. Currently, only ``'rectangle'`` is supported. Here 'rectangle' refers to boundaries that are parallel to lines of constant longitude or latitude. + 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. + Returns: *.nc files: Regridded tidal velocity and elevation files in 'inputdir/forcing' @@ -1670,7 +1676,7 @@ def setup_boundary_tides( This tidal data functions are sourced from the GFDL NWA25 and changed in the following ways: - Converted code for RM6 segment class - Implemented Horizontal Subsetting - - Combined all functions of NWA25 into a four function process (in the style of rm6) (expt.setup_tides_rectangular_boundaries, self.coords, segment.regrid_tides, segment.encode_tidal_files_and_output) + - Combined all functions of NWA25 into a four function process (in the style of rm6) (expt.setup_tides_rectangular_boundaries, coords, segment.regrid_tides, segment.encode_tidal_files_and_output) Original Code was sourced from: @@ -1749,7 +1755,9 @@ def setup_boundary_tides( seg = self.segments[b] # Output and regrid tides - seg.regrid_tides(tpxo_v, tpxo_u, tpxo_h, times) + seg.regrid_tides( + tpxo_v, tpxo_u, tpxo_h, times, rotational_method=rotational_method + ) print("Done") def setup_bathymetry( @@ -2901,129 +2909,41 @@ def __init__( self.segment_name = segment_name self.repeat_year_forcing = repeat_year_forcing - @property - def coords(self): - """ - - - This function: - Allows us to call the self.coords for use in the xesmf.Regridder in the regrid_tides function. self.coords gives us the subset of the hgrid based on the orientation. - - Args: - None - Returns: - xr.Dataset: The correct coordinate space for the orientation - - General Description: - This tidal data functions are sourced from the GFDL NWA25 and changed in the following ways: - - Converted code for RM6 segment class - - Implemented Horizontal Subsetting - - Combined all functions of NWA25 into a four function process (in the style of rm6) (expt.setup_tides_rectangular_boundaries, segment.coords, segment.regrid_tides, segment.encode_tidal_files_and_output) - - - Code adapted from: - Author(s): GFDL, James Simkins, Rob Cermak, etc.. - Year: 2022 - Title: "NWA25: Northwest Atlantic 1/25th Degree MOM6 Simulation" - Version: N/A - Type: Python Functions, Source Code - Web Address: https://github.com/jsimkins2/nwa25 - - """ - # Rename nxp and nyp to locations - if self.orientation == "south": - rcoord = xr.Dataset( - { - "lon": self.hgrid["x"].isel(nyp=0), - "lat": self.hgrid["y"].isel(nyp=0), - "angle": self.hgrid["angle_dx"].isel(nyp=0), - } - ) - rcoord = rcoord.rename_dims({"nxp": f"nx_{self.segment_name}"}) - rcoord.attrs["perpendicular"] = "ny" - rcoord.attrs["parallel"] = "nx" - rcoord.attrs["axis_to_expand"] = ( - 2 ## Need to keep track of which axis the 'main' coordinate corresponds to when re-adding the 'secondary' axis - ) - rcoord.attrs["locations_name"] = ( - f"nx_{self.segment_name}" # Legacy name of nx_... was locations. This provides a clear transform in regrid_tides - ) - elif self.orientation == "north": - rcoord = xr.Dataset( - { - "lon": self.hgrid["x"].isel(nyp=-1), - "lat": self.hgrid["y"].isel(nyp=-1), - "angle": self.hgrid["angle_dx"].isel(nyp=-1), - } - ) - rcoord = rcoord.rename_dims({"nxp": f"nx_{self.segment_name}"}) - rcoord.attrs["perpendicular"] = "ny" - rcoord.attrs["parallel"] = "nx" - rcoord.attrs["axis_to_expand"] = 2 - rcoord.attrs["locations_name"] = f"nx_{self.segment_name}" - elif self.orientation == "west": - rcoord = xr.Dataset( - { - "lon": self.hgrid["x"].isel(nxp=0), - "lat": self.hgrid["y"].isel(nxp=0), - "angle": self.hgrid["angle_dx"].isel(nxp=0), - } - ) - rcoord = rcoord.rename_dims({"nyp": f"ny_{self.segment_name}"}) - rcoord.attrs["perpendicular"] = "nx" - rcoord.attrs["parallel"] = "ny" - rcoord.attrs["axis_to_expand"] = 3 - rcoord.attrs["locations_name"] = f"ny_{self.segment_name}" - elif self.orientation == "east": - rcoord = xr.Dataset( - { - "lon": self.hgrid["x"].isel(nxp=-1), - "lat": self.hgrid["y"].isel(nxp=-1), - "angle": self.hgrid["angle_dx"].isel(nxp=-1), - } - ) - rcoord = rcoord.rename_dims({"nyp": f"ny_{self.segment_name}"}) - rcoord.attrs["perpendicular"] = "nx" - rcoord.attrs["parallel"] = "ny" - rcoord.attrs["axis_to_expand"] = 3 - rcoord.attrs["locations_name"] = f"ny_{self.segment_name}" - - # Make lat and lon coordinates - rcoord = rcoord.assign_coords(lat=rcoord["lat"], lon=rcoord["lon"]) - - return rcoord - - def rotate(self, u, v): + def rotate(self, u, v, radian_angle): # Make docstring """ Rotate the velocities to the grid orientation. - Args: u (xarray.DataArray): The u-component of the velocity. v (xarray.DataArray): The v-component of the velocity. + radian_angle (xarray.DataArray): The angle of the grid in RADIANS Returns: Tuple[xarray.DataArray, xarray.DataArray]: The rotated u and v components of the velocity. """ - angle = self.coords.angle.values * np.pi / 180 - u_rot = u * np.cos(angle) - v * np.sin(angle) - v_rot = u * np.sin(angle) + v * np.cos(angle) + u_rot = u * np.cos(radian_angle) - v * np.sin(radian_angle) + v_rot = u * np.sin(radian_angle) + v * np.cos(radian_angle) return u_rot, v_rot - def regrid_velocity_tracers(self): + def regrid_velocity_tracers(self, rotational_method=rot.RotationMethod.GIVEN_ANGLE): """ Cut out and interpolate the velocities and tracers + Args: + rotational_method (rot.RotationMethod): The method to use for rotation of the velocities. Currently, the default method, GIVEN_ANGLE, works even with non-rotated grids """ - + if rotational_method == rot.RotationMethod.NO_ROTATION: + if not is_rectilinear_hgrid(self.hgrid): + raise ValueError("NO_ROTATION method only works with rectilinear grids") rawseg = xr.open_dataset(self.infile, decode_times=False, engine="netcdf4") + coords = rgd.coords(self.hgrid, self.orientation, self.segment_name) if self.arakawa_grid == "A": rawseg = rawseg.rename({self.x: "lon", self.y: "lat"}) - ## In this case velocities and tracers all on same points + # In this case velocities and tracers all on same points regridder = rgd.create_regridder( rawseg[self.u], coords, @@ -3036,7 +2956,40 @@ def regrid_velocity_tracers(self): [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]) + + ## Angle Calculation & Rotation + if rotational_method == rot.RotationMethod.GIVEN_ANGLE: + rotated_u, rotated_v = self.rotate( + regridded[self.u], + regridded[self.v], + radian_angle=np.radians(coords.angle.values), + ) + + elif rotational_method == rot.RotationMethod.EXPAND_GRID: + + # Recalculate entire hgrid angles + self.hgrid["angle_dx_rm6"] = ( + rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid) + ) + + # Get just the boundary + degree_angle = rgd.coords( + self.hgrid, + self.orientation, + self.segment_name, + angle_variable_name="angle_dx_rm6", + )["angle"] + + # Rotate + rotated_u, rotated_v = self.rotate( + regridded[self.u], + regridded[self.v], + radian_angle=np.radians(degree_angle.values), + ) + elif rotational_method == rot.RotationMethod.NO_ROTATION: + # Just transfer values + rotated_u, rotated_v = regridded[self.u], regridded[self.v] + rotated_ds = xr.Dataset( { self.u: rotated_u, @@ -3064,9 +3017,33 @@ def regrid_velocity_tracers(self): rawseg[[self.u, self.v]].rename({self.xq: "lon", self.yq: "lat"}) ) - velocities_out["u"], velocities_out["v"] = self.rotate( - velocities_out["u"], velocities_out["v"] - ) + # See explanation of the rotational methods in the A grid section + if rotational_method == rot.RotationMethod.GIVEN_ANGLE: + velocities_out["u"], velocities_out["v"] = self.rotate( + velocities_out["u"], + velocities_out["v"], + radian_angle=np.radians(coords.angle.values), + ) + elif rotational_method == rot.RotationMethod.EXPAND_GRID: + self.hgrid["angle_dx_rm6"] = ( + rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid) + ) + degree_angle = rgd.coords( + self.hgrid, + self.orientation, + self.segment_name, + angle_variable_name="angle_dx_rm6", + )["angle"] + velocities_out["u"], velocities_out["v"] = self.rotate( + velocities_out["u"], + velocities_out["v"], + radian_angle=np.radians(degree_angle.values), + ) + elif rotational_method == rot.RotationMethod.NO_ROTATION: + velocities_out["u"], velocities_out["v"] = ( + velocities_out["u"], + velocities_out["v"], + ) segment_out = xr.merge( [ @@ -3105,7 +3082,30 @@ 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]) + # See explanation of the rotational methods in the A grid section + if rotational_method == rot.RotationMethod.GIVEN_ANGLE: + rotated_u, rotated_v = self.rotate( + regridded_u, + regridded_v, + radian_angle=np.radians(coords.angle.values), + ) + elif rotational_method == rot.RotationMethod.EXPAND_GRID: + self.hgrid["angle_dx_rm6"] = ( + rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid) + ) + degree_angle = rgd.coords( + self.hgrid, + self.orientation, + self.segment_name, + angle_variable_name="angle_dx_rm6", + )["angle"] + rotated_u, rotated_v = self.rotate( + regridded_u, + regridded_v, + radian_angle=np.radians(degree_angle.values), + ) + elif rotational_method == rot.RotationMethod.NO_ROTATION: + rotated_u, rotated_v = regridded_u, regridded_v rotated_ds = xr.Dataset( { self.u: rotated_u, @@ -3137,7 +3137,7 @@ def regrid_velocity_tracers(self): # fill in NaNs segment_out = rgd.fill_missing_data(segment_out, self.z) segment_out = rgd.fill_missing_data( - segment_out, f"{self.coords.attrs['parallel']}_{self.segment_name}" + segment_out, f"{coords.attrs['parallel']}_{self.segment_name}" ) times = xr.DataArray( @@ -3194,10 +3194,10 @@ 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"{coords.attrs['parallel']}_{self.segment_name}"] = np.arange( + segment_out[f"{coords.attrs['parallel']}_{self.segment_name}"].size ) - segment_out[f"{self.coords.attrs['perpendicular']}_{self.segment_name}"] = [0] + segment_out[f"{coords.attrs['perpendicular']}_{self.segment_name}"] = [0] encoding_dict = { "time": {"dtype": "double"}, f"nx_{self.segment_name}": { @@ -3222,7 +3222,12 @@ def regrid_velocity_tracers(self): return segment_out, encoding_dict def regrid_tides( - self, tpxo_v, tpxo_u, tpxo_h, times, method="nearest_s2d", periodic=False + self, + tpxo_v, + tpxo_u, + tpxo_h, + times, + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ This function: @@ -3236,6 +3241,7 @@ def regrid_tides( infile_td (str): Raw Tidal File/Dir tpxo_v, tpxo_u, tpxo_h (xarray.Dataset): Specific adjusted for MOM6 tpxo datasets (Adjusted with setup_tides) times (pd.DateRange): The start date of our model period + rotational_method (rot.RotationMethod): The method to use for rotation of the velocities. Currently, the default method, GIVEN_ANGLE, works even with non-rotated grids Returns: *.nc files: Regridded tidal velocity and elevation files in 'inputdir/forcing' @@ -3254,8 +3260,11 @@ def regrid_tides( Type: Python Functions, Source Code Web Address: https://github.com/jsimkins2/nwa25 """ + if rotational_method == rot.RotationMethod.NO_ROTATION: + if not is_rectilinear_hgrid(self.hgrid): + raise ValueError("NO_ROTATION method only works with rectilinear grids") - # Establish Coord + # Establish Coords coords = rgd.coords(self.hgrid, self.orientation, self.segment_name) ########## Tidal Elevation: Horizontally interpolate elevation components ############ @@ -3310,8 +3319,9 @@ def regrid_tides( self.encode_tidal_files_and_output(ds_ap, "tz") ########### Regrid Tidal Velocity ###################### - regrid_u = rgd.create_regridder(tpxo_u[["lon", "lat", "uRe"]], coords, ".temp") - regrid_v = rgd.create_regridder(tpxo_v[["lon", "lat", "vRe"]], coords, ".temp2") + + regrid_u = rgd.create_regridder(tpxo_u[["lon", "lat", "uRe"]], coords) + regrid_v = rgd.create_regridder(tpxo_v[["lon", "lat", "vRe"]], coords) # Interpolate each real and imaginary parts to self. uredest = regrid_u(tpxo_u[["lon", "lat", "uRe"]])["uRe"] @@ -3338,17 +3348,40 @@ def regrid_tides( ucplex = uredest + 1j * uimdest vcplex = vredest + 1j * vimdest - 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, :]) + if rotational_method == rot.RotationMethod.GIVEN_ANGLE: + + # Get user-provided angle + angle = coords["angle"] + + # Rotate + INC -= np.radians(angle.data[np.newaxis, :]) + + elif rotational_method == rot.RotationMethod.EXPAND_GRID: + + # Generate entire hgrid angles using pseudo_hgrid + self.hgrid["angle_dx_rm6"] = ( + rot.initialize_grid_rotation_angles_using_expanded_hgrid(self.hgrid) + ) + + # Get just boundary angles + degree_angle = rgd.coords( + self.hgrid, + self.orientation, + self.segment_name, + angle_variable_name="angle_dx_rm6", + )["angle"] + + # 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( {f"uamp_{self.segment_name}": ua, f"vamp_{self.segment_name}": va} ) @@ -3401,7 +3434,7 @@ def encode_tidal_files_and_output(self, ds, filename): This tidal data functions are sourced from the GFDL NWA25 and changed in the following ways: - Converted code for RM6 segment class - Implemented Horizontal Subsetting - - Combined all functions of NWA25 into a four function process (in the style of rm6) (expt.setup_tides_rectangular_boundaries, self.coords, segment.regrid_tides, segment.encode_tidal_files_and_output) + - Combined all functions of NWA25 into a four function process (in the style of rm6) (expt.setup_tides_rectangular_boundaries, coords, segment.regrid_tides, segment.encode_tidal_files_and_output) Original Code was sourced from: diff --git a/regional_mom6/regridding.py b/regional_mom6/regridding.py index b4026e3f..adebb7a5 100644 --- a/regional_mom6/regridding.py +++ b/regional_mom6/regridding.py @@ -36,15 +36,22 @@ regridding_logger = setup_logger(__name__) -def coords(hgrid: xr.Dataset, orientation: str, segment_name: str) -> xr.Dataset: +def coords( + hgrid: xr.Dataset, + orientation: str, + segment_name: str, + coords_at_t_points=False, + angle_variable_name="angle_dx", +) -> xr.Dataset: """ This function: - Allows us to call the self.coords for use in the xesmf.Regridder in the regrid_tides function. self.coords gives us the subset of the hgrid based on the orientation. + Allows us to call the coords for use in the xesmf.Regridder in the regrid_tides function. self.coords gives us the subset of the hgrid based on the orientation. Args: hgrid (xr.Dataset): The hgrid dataset orientation (str): The orientation of the boundary segment_name (str): The name of the segment + coords_at_t_points (bool, optional): Whether to return the boundary t-points instead of the q/u/v of a general boundary for rotation. Defaults to False. Returns: xr.Dataset: The correct coordinate space for the orientation @@ -57,13 +64,37 @@ def coords(hgrid: xr.Dataset, orientation: str, segment_name: str) -> xr.Dataset Web Address: https://github.com/jsimkins2/nwa25 """ + + dataset_to_get_coords = None + + if coords_at_t_points: + regridding_logger.info("Creating coordinates of the boundary t-points") + + # Calc T Point Info + ds = get_hgrid_arakawa_c_points(hgrid, "t") + + tangle_dx = hgrid[angle_variable_name][(ds.t_points_y, ds.t_points_x)] + # Assign to dataset + dataset_to_get_coords = xr.Dataset( + { + "x": ds.tlon, + "y": ds.tlat, + angle_variable_name: (("nyp", "nxp"), tangle_dx.values), + }, + coords={"nyp": ds.nyp, "nxp": ds.nxp}, + ) + else: + regridding_logger.info("Creating coordinates of the boundary q/u/v points") + # Don't have to do anything because this is the actual boundary. t-points are one-index deep and require managing. + dataset_to_get_coords = hgrid + # Rename nxp and nyp to locations if orientation == "south": rcoord = xr.Dataset( { - "lon": hgrid["x"].isel(nyp=0), - "lat": hgrid["y"].isel(nyp=0), - "angle": hgrid["angle_dx"].isel(nyp=0), + "lon": dataset_to_get_coords["x"].isel(nyp=0), + "lat": dataset_to_get_coords["y"].isel(nyp=0), + "angle": dataset_to_get_coords[angle_variable_name].isel(nyp=0), } ) rcoord = rcoord.rename_dims({"nxp": f"nx_{segment_name}"}) @@ -75,9 +106,9 @@ def coords(hgrid: xr.Dataset, orientation: str, segment_name: str) -> xr.Dataset elif orientation == "north": rcoord = xr.Dataset( { - "lon": hgrid["x"].isel(nyp=-1), - "lat": hgrid["y"].isel(nyp=-1), - "angle": hgrid["angle_dx"].isel(nyp=-1), + "lon": dataset_to_get_coords["x"].isel(nyp=-1), + "lat": dataset_to_get_coords["y"].isel(nyp=-1), + "angle": dataset_to_get_coords[angle_variable_name].isel(nyp=-1), } ) rcoord = rcoord.rename_dims({"nxp": f"nx_{segment_name}"}) @@ -87,9 +118,9 @@ def coords(hgrid: xr.Dataset, orientation: str, segment_name: str) -> xr.Dataset elif orientation == "west": rcoord = xr.Dataset( { - "lon": hgrid["x"].isel(nxp=0), - "lat": hgrid["y"].isel(nxp=0), - "angle": hgrid["angle_dx"].isel(nxp=0), + "lon": dataset_to_get_coords["x"].isel(nxp=0), + "lat": dataset_to_get_coords["y"].isel(nxp=0), + "angle": dataset_to_get_coords[angle_variable_name].isel(nxp=0), } ) rcoord = rcoord.rename_dims({"nyp": f"ny_{segment_name}"}) @@ -99,9 +130,9 @@ def coords(hgrid: xr.Dataset, orientation: str, segment_name: str) -> xr.Dataset elif orientation == "east": rcoord = xr.Dataset( { - "lon": hgrid["x"].isel(nxp=-1), - "lat": hgrid["y"].isel(nxp=-1), - "angle": hgrid["angle_dx"].isel(nxp=-1), + "lon": dataset_to_get_coords["x"].isel(nxp=-1), + "lat": dataset_to_get_coords["y"].isel(nxp=-1), + "angle": dataset_to_get_coords[angle_variable_name].isel(nxp=-1), } ) rcoord = rcoord.rename_dims({"nyp": f"ny_{segment_name}"}) @@ -115,10 +146,64 @@ def coords(hgrid: xr.Dataset, orientation: str, segment_name: str) -> xr.Dataset return rcoord +def get_hgrid_arakawa_c_points(hgrid: xr.Dataset, point_type="t") -> xr.Dataset: + """ + Get the Arakawa C points from the Hgrid, originally written by Fred (Castruccio) and moved to RM6 + Parameters + ---------- + hgrid: xr.Dataset + The hgrid dataset + Returns + ------- + xr.Dataset + The specific points x, y, & point indexes + """ + if point_type not in "uvqth": + raise ValueError("point_type must be one of 'uvqht'") + + regridding_logger.info("Getting {} points..".format(point_type)) + + # Figure out the maths for the offset + k = 2 + kp2 = k // 2 + offset_one_by_two_y = np.arange(kp2, len(hgrid.x.nyp), k) + offset_one_by_two_x = np.arange(kp2, len(hgrid.x.nxp), k) + by_two_x = np.arange(0, len(hgrid.x.nxp), k) + by_two_y = np.arange(0, len(hgrid.x.nyp), k) + + # T point locations + if point_type == "t" or point_type == "h": + points = (offset_one_by_two_y, offset_one_by_two_x) + # U point locations + elif point_type == "u": + points = (offset_one_by_two_y, by_two_x) + # V point locations + elif point_type == "v": + points = (by_two_y, offset_one_by_two_x) + # Corner point locations + elif point_type == "q": + points = (by_two_y, by_two_x) + else: + raise ValueError("Invalid Point Type (u, v, q, or t/h only)") + + point_dataset = xr.Dataset( + { + "{}lon".format(point_type): hgrid.x[points], + "{}lat".format(point_type): hgrid.y[points], + "{}_points_y".format(point_type): points[0], + "{}_points_x".format(point_type): points[1], + } + ) + point_dataset.attrs["description"] = ( + "Arakawa C {}-points of supplied h-grid".format(point_type) + ) + return point_dataset + + def create_regridder( forcing_variables: xr.Dataset, output_grid: xr.Dataset, - outfile: Path = Path(".temp"), + outfile: Path = None, method: str = "bilinear", ) -> xe.Regridder: """ @@ -229,7 +314,7 @@ def generate_dz(ds: xr.Dataset, z_dim_name: str) -> xr.Dataset: def add_secondary_dimension( - ds: xr.Dataset, var: str, coords, segment_name: str + ds: xr.Dataset, var: str, coords, segment_name: str, to_beginning=False ) -> xr.Dataset: """Add the perpendiciular dimension to the dataset, even if it's like one val. It's required. Parameters @@ -242,6 +327,8 @@ def add_secondary_dimension( The coordinates from the function coords... segment_name : str The segment name + to_beginning : bool, optional + Whether to add the perpendicular dimension to the beginning or to the selected position, by default False Returns ------- xr.Dataset @@ -257,12 +344,20 @@ def add_secondary_dimension( "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): - regridding_logger.debug("Bump it by one") - insert_behind_by = 0 + if not to_beginning: + + 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: + # Missing vertical dim or tidal coord means we don't need to offset the perpendicular + insert_behind_by = 1 else: - # Missing vertical dim or tidal coord means we don't need to offset the perpendicular - insert_behind_by = 1 + insert_behind_by = coords.attrs[ + "axis_to_expand" + ] # Just magic to add dim to the beginning regridding_logger.debug(f"Expand dimensions") ds[var] = ds[var].expand_dims( diff --git a/regional_mom6/rotation.py b/regional_mom6/rotation.py new file mode 100644 index 00000000..9c3900dd --- /dev/null +++ b/regional_mom6/rotation.py @@ -0,0 +1,250 @@ +from .utils import setup_logger + +rotation_logger = setup_logger(__name__) +# An Enum is like a dropdown selection for a menu, it essentially limits the type of input parameters. It comes with additional complexity, which of course is always a challenge. +from enum import Enum +import xarray as xr +import numpy as np +from .regridding import get_hgrid_arakawa_c_points + + +class RotationMethod(Enum): + """ + This Enum defines the rotational method to be used in boundary conditions. The main regional mom6 class passes in this enum to regrid_tides and regrid_velocity_tracers to determine the method used. + + EXPAND_GRID: This method is used with the basis that we can find the angles at the q-u-v points by pretending we have another row/column of the hgrid with the same distances as the t-point to u/v points in the actual grid then use the four poitns to calculate the angle the exact same way MOM6 does. + GIVEN_ANGLE: This is the original default RM6 method which expects a pre-given angle called angle_dx + NO_ROTATION: Grids parallel to lat/lon axes, no rotation needed + """ + + EXPAND_GRID = 1 + GIVEN_ANGLE = 2 + NO_ROTATION = 3 + + +def initialize_grid_rotation_angles_using_expanded_hgrid( + hgrid: xr.Dataset, +) -> xr.Dataset: + """ + Calculate the angle_dx in degrees from the true x (east?) direction counterclockwise) and return as dataarray + + Parameters + ---------- + hgrid: xr.Dataset + The hgrid dataset + Returns + ------- + xr.DataArray + The t-point angles + """ + # Get expanded (pseudo) grid + expanded_hgrid = create_expanded_hgrid(hgrid) + + return mom6_angle_calculation_method( + expanded_hgrid.x.max() - expanded_hgrid.x.min(), + expanded_hgrid.isel(nyp=slice(2, None), nxp=slice(0, -2)), + expanded_hgrid.isel(nyp=slice(2, None), nxp=slice(2, None)), + expanded_hgrid.isel(nyp=slice(0, -2), nxp=slice(0, -2)), + expanded_hgrid.isel(nyp=slice(0, -2), nxp=slice(2, None)), + hgrid, + ) + + +def initialize_grid_rotation_angle(hgrid: xr.Dataset) -> xr.DataArray: + """ + Calculate the angle_dx in degrees from the true x (east?) direction counterclockwise) and return as DataArray + Parameters + ---------- + hgrid: xr.Dataset + The hgrid dataset + Returns + ------- + xr.DataArray + The t-point angles + """ + ds_t = get_hgrid_arakawa_c_points(hgrid, "t") + ds_q = get_hgrid_arakawa_c_points(hgrid, "q") + + # Reformat into x, y comps + t_points = xr.Dataset( + { + "x": (("nyp", "nxp"), ds_t.tlon.data), + "y": (("nyp", "nxp"), ds_t.tlat.data), + } + ) + q_points = xr.Dataset( + { + "x": (("nyp", "nxp"), ds_q.qlon.data), + "y": (("nyp", "nxp"), ds_q.qlat.data), + } + ) + + return mom6_angle_calculation_method( + hgrid.x.max() - hgrid.x.min(), + q_points.isel(nyp=slice(1, None), nxp=slice(0, -1)), + q_points.isel(nyp=slice(1, None), nxp=slice(1, None)), + q_points.isel(nyp=slice(0, -1), nxp=slice(0, -1)), + q_points.isel(nyp=slice(0, -1), nxp=slice(1, None)), + t_points, + ) + + +def modulo_around_point(x, xc, Lx): + """ + This function calculates the modulo around a point. Return the modulo value of x in an interval [xc-(Lx/2) xc+(Lx/2)]. If Lx<=0, then it returns x without applying modulo arithmetic. + Parameters + ---------- + x: float + Value to which to apply modulo arithmetic + xc: float + Center of modulo range + Lx: float + Modulo range width + Returns + ------- + float + x shifted by an integer multiple of Lx to be close to xc, + """ + if Lx <= 0: + return x + else: + return ((x - (xc - 0.5 * Lx)) % Lx) - Lx / 2 + xc + + +def mom6_angle_calculation_method( + len_lon, + top_left: xr.DataArray, + top_right: xr.DataArray, + bottom_left: xr.DataArray, + bottom_right: xr.DataArray, + point: xr.DataArray, +) -> xr.DataArray: + """ + Calculate the angle of the point using the MOM6 method in initialize_grid_rotation_angle. Built for vectorized calculations + Parameters + ---------- + len_lon: float + The length of the longitude of the regional domain + top_left, top_right, bottom_left, bottom_right: xr.DataArray + The four points around the point to calculate the angle from the hgrid requires an x and y component + point: xr.DataArray + The point to calculate the angle from the hgrid + Returns + ------- + xr.DataArray + The angle of the point + """ + rotation_logger.info("Calculating grid rotation angle") + # Direct Translation + pi_720deg = ( + np.arctan(1) / 180 + ) # One quarter the conversion factor from degrees to radians + + # Compute lonB for all points + lonB = np.zeros((2, 2, len(point.nyp), len(point.nxp))) + + # Vectorized computation of lonB + # Vectorized computation of lonB + lonB[0][0] = modulo_around_point(bottom_left.x, point.x, len_lon) # Bottom Left + lonB[1][0] = modulo_around_point(top_left.x, point.x, len_lon) # Top Left + lonB[1][1] = modulo_around_point(top_right.x, point.x, len_lon) # Top Right + lonB[0][1] = modulo_around_point(bottom_right.x, point.x, len_lon) # Bottom Right + + # Compute lon_scale + lon_scale = np.cos( + pi_720deg * ((bottom_left.y + bottom_right.y) + (top_right.y + top_left.y)) + ) + + # Compute angle + angle = np.arctan2( + lon_scale * ((lonB[0, 1] - lonB[1, 0]) + (lonB[1, 1] - lonB[0, 0])), + (bottom_left.y - top_right.y) + (top_left.y - bottom_right.y), + ) + # Assign angle to angles_arr + angles_arr = np.rad2deg(angle) - 90 + + # Assign angles_arr to hgrid + t_angles = xr.DataArray( + angles_arr, + dims=["nyp", "nxp"], + coords={ + "nyp": point.nyp.values, + "nxp": point.nxp.values, + }, + ) + return t_angles + + +def create_expanded_hgrid(hgrid: xr.Dataset, expansion_width=1) -> xr.Dataset: + """ + Adds an additional boundary to the hgrid to allow for the calculation of the angle_dx for the boundary points using the method in MOM6 + """ + if expansion_width != 1: + raise NotImplementedError("Only expansion_width = 1 is supported") + + pseudo_hgrid_x = np.full((len(hgrid.nyp) + 2, len(hgrid.nxp) + 2), np.nan) + pseudo_hgrid_y = np.full((len(hgrid.nyp) + 2, len(hgrid.nxp) + 2), np.nan) + + ## Fill Boundaries + pseudo_hgrid_x[1:-1, 1:-1] = hgrid.x.values + pseudo_hgrid_x[0, 1:-1] = hgrid.x.values[0, :] - ( + hgrid.x.values[1, :] - hgrid.x.values[0, :] + ) # Bottom Fill + pseudo_hgrid_x[-1, 1:-1] = hgrid.x.values[-1, :] + ( + hgrid.x.values[-1, :] - hgrid.x.values[-2, :] + ) # Top Fill + pseudo_hgrid_x[1:-1, 0] = hgrid.x.values[:, 0] - ( + hgrid.x.values[:, 1] - hgrid.x.values[:, 0] + ) # Left Fill + pseudo_hgrid_x[1:-1, -1] = hgrid.x.values[:, -1] + ( + hgrid.x.values[:, -1] - hgrid.x.values[:, -2] + ) # Right Fill + + pseudo_hgrid_y[1:-1, 1:-1] = hgrid.y.values + pseudo_hgrid_y[0, 1:-1] = hgrid.y.values[0, :] - ( + hgrid.y.values[1, :] - hgrid.y.values[0, :] + ) # Bottom Fill + pseudo_hgrid_y[-1, 1:-1] = hgrid.y.values[-1, :] + ( + hgrid.y.values[-1, :] - hgrid.y.values[-2, :] + ) # Top Fill + pseudo_hgrid_y[1:-1, 0] = hgrid.y.values[:, 0] - ( + hgrid.y.values[:, 1] - hgrid.y.values[:, 0] + ) # Left Fill + pseudo_hgrid_y[1:-1, -1] = hgrid.y.values[:, -1] + ( + hgrid.y.values[:, -1] - hgrid.y.values[:, -2] + ) # Right Fill + + ## Fill Corners + pseudo_hgrid_x[0, 0] = hgrid.x.values[0, 0] - ( + hgrid.x.values[1, 1] - hgrid.x.values[0, 0] + ) # Bottom Left + pseudo_hgrid_x[-1, 0] = hgrid.x.values[-1, 0] - ( + hgrid.x.values[-2, 1] - hgrid.x.values[-1, 0] + ) # Top Left + pseudo_hgrid_x[0, -1] = hgrid.x.values[0, -1] - ( + hgrid.x.values[1, -2] - hgrid.x.values[0, -1] + ) # Bottom Right + pseudo_hgrid_x[-1, -1] = hgrid.x.values[-1, -1] - ( + hgrid.x.values[-2, -2] - hgrid.x.values[-1, -1] + ) # Top Right + + pseudo_hgrid_y[0, 0] = hgrid.y.values[0, 0] - ( + hgrid.y.values[1, 1] - hgrid.y.values[0, 0] + ) # Bottom Left + pseudo_hgrid_y[-1, 0] = hgrid.y.values[-1, 0] - ( + hgrid.y.values[-2, 1] - hgrid.y.values[-1, 0] + ) # Top Left + pseudo_hgrid_y[0, -1] = hgrid.y.values[0, -1] - ( + hgrid.y.values[1, -2] - hgrid.y.values[0, -1] + ) # Bottom Right + pseudo_hgrid_y[-1, -1] = hgrid.y.values[-1, -1] - ( + hgrid.y.values[-2, -2] - hgrid.y.values[-1, -1] + ) # Top Right + + pseudo_hgrid = xr.Dataset( + { + "x": (["nyp", "nxp"], pseudo_hgrid_x), + "y": (["nyp", "nxp"], pseudo_hgrid_y), + } + ) + return pseudo_hgrid diff --git a/regional_mom6/utils.py b/regional_mom6/utils.py index 91a96c36..d499430e 100644 --- a/regional_mom6/utils.py +++ b/regional_mom6/utils.py @@ -1,6 +1,7 @@ import numpy as np import logging import sys +import xarray as xr def vecdot(v1, v2): @@ -318,3 +319,17 @@ def setup_logger(name: str) -> logging.Logger: # Add the handler to the logger logger.addHandler(handler) return logger + + +def is_rectilinear_hgrid(hgrid: xr.Dataset) -> bool: + """ + Check if the hgrid is a rectilinear grid. + """ + if hgrid.x.shape[0] < 2 or hgrid.x.shape[1] < 2: + raise ValueError("hgrid must have at least 2 points in each direction") + if not np.all(hgrid.y == hgrid.y[:, 0].values[:, np.newaxis]): + return False + if not np.all(hgrid.x == hgrid.x[0, :].values[np.newaxis, :]): + return False + + return True diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..db1bf2cf --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,29 @@ +import pytest +import os +import xarray as xr + +# Define the path where the curvilinear hgrid file is expected in the Docker container +DOCKER_FILE_PATH = "/data/small_curvilinear_hgrid.nc" + + +# Define the local directory where the user might have added the curvilinear hgrid file +LOCAL_FILE_PATH = ( + "/glade/u/home/manishrv/documents/nwa12_0.1/tides_dev/small_curvilinear_hgrid.nc" +) + + +@pytest.fixture +def get_curvilinear_hgrid(): + # Check if the file exists in the Docker-specific location + if os.path.exists(DOCKER_FILE_PATH): + return xr.open_dataset(DOCKER_FILE_PATH) + + # Check if the user has provided the file in a specific local directory + elif os.path.exists(LOCAL_FILE_PATH): + return xr.open_dataset(LOCAL_FILE_PATH) + + # If neither location contains the file, raise an error + else: + pytest.skip( + f"Required file 'hgrid.nc' not found in {DOCKER_FILE_PATH} or {LOCAL_FILE_PATH}" + ) diff --git a/tests/test_config.py b/tests/test_config.py index eefc4297..31d55886 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -102,11 +102,12 @@ def test_load_config(tmp_path): ## Directory where you'll run the experiment from run_dir = Path( + tmp_path, os.path.join( tmp_path, expt_name, "run_files", - ) + ), ) data_path = Path(tmp_path / "data") for path in (run_dir, input_dir, data_path): @@ -126,7 +127,7 @@ def test_load_config(tmp_path): mom_input_dir=input_dir, toolpath_dir="", ) - path = os.path.join(tmp_path, "testing_config.json") + path = tmp_path / "testing_config.json" config_expt = expt.write_config_file(path) new_expt = rmom6.create_experiment_from_config( os.path.join(path), mom_input_folder=tmp_path, mom_run_folder=tmp_path diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index d459aed3..e5554f8e 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -17,8 +17,6 @@ "number_vertical_layers", "layer_thickness_ratio", "depth", - "mom_run_dir", - "mom_input_dir", "toolpath_dir", "hgrid_type", ), @@ -31,8 +29,6 @@ 5, 1, 1000, - "rundir/", - "inputdir/", "toolpath", "even_spacing", ), @@ -46,12 +42,12 @@ def test_setup_bathymetry( number_vertical_layers, layer_thickness_ratio, depth, - mom_run_dir, - mom_input_dir, toolpath_dir, hgrid_type, tmp_path, ): + mom_run_dir = tmp_path / "rundir" + mom_input_dir = tmp_path / "inputdir" expt = experiment( longitude_extent=longitude_extent, latitude_extent=latitude_extent, @@ -166,8 +162,7 @@ def generate_silly_coords( coords = {"silly_lat": silly_lat, "silly_lon": silly_lon, "silly_depth": silly_depth} -mom_run_dir = "rundir/" -mom_input_dir = "inputdir/" + toolpath_dir = "toolpath" hgrid_type = "even_spacing" @@ -198,8 +193,6 @@ def generate_silly_coords( "number_vertical_layers", "layer_thickness_ratio", "depth", - "mom_run_dir", - "mom_input_dir", "toolpath_dir", "hgrid_type", ), @@ -212,8 +205,6 @@ def generate_silly_coords( number_vertical_layers, layer_thickness_ratio, depth, - "rundir/", - "inputdir/", "toolpath", "even_spacing", ), @@ -227,8 +218,6 @@ def test_ocean_forcing( number_vertical_layers, layer_thickness_ratio, depth, - mom_run_dir, - mom_input_dir, toolpath_dir, hgrid_type, temp_dataarray_initial_condition, @@ -246,7 +235,8 @@ def test_ocean_forcing( "silly_lon": silly_lon, "silly_depth": silly_depth, } - + mom_run_dir = tmp_path / "rundir" + mom_input_dir = tmp_path / "inputdir" expt = experiment( longitude_extent=longitude_extent, latitude_extent=latitude_extent, @@ -332,8 +322,6 @@ def test_ocean_forcing( "number_vertical_layers", "layer_thickness_ratio", "depth", - "mom_run_dir", - "mom_input_dir", "toolpath_dir", "hgrid_type", ), @@ -346,8 +334,6 @@ def test_ocean_forcing( 5, 1, 1000, - "rundir/", - "inputdir/", "toolpath", "even_spacing", ), @@ -361,8 +347,6 @@ def test_rectangular_boundaries( number_vertical_layers, layer_thickness_ratio, depth, - mom_run_dir, - mom_input_dir, toolpath_dir, hgrid_type, tmp_path, @@ -443,7 +427,8 @@ def test_rectangular_boundaries( ) eastern_boundary.to_netcdf(tmp_path / "east_unprocessed.nc") eastern_boundary.close() - + mom_run_dir = tmp_path / "rundir" + mom_input_dir = tmp_path / "inputdir" expt = experiment( longitude_extent=longitude_extent, latitude_extent=latitude_extent, diff --git a/tests/test_manish_branch.py b/tests/test_manish_branch.py deleted file mode 100644 index 98755c8d..00000000 --- a/tests/test_manish_branch.py +++ /dev/null @@ -1,338 +0,0 @@ -""" -Test suite for everything involed in pr #12 -""" - -import regional_mom6 as rmom6 -import os -import pytest -import logging -from pathlib import Path -import xarray as xr -import numpy as np -import shutil -import importlib - -IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" - - -@pytest.fixture(scope="module") -def dummy_tidal_data(): - nx = 2160 - ny = 1081 - nc = 15 - nct = 4 - - # Define tidal constituents - con_list = [ - "m2 ", - "s2 ", - "n2 ", - "k2 ", - "k1 ", - "o1 ", - "p1 ", - "q1 ", - "mm ", - "mf ", - "m4 ", - "mn4 ", - "ms4 ", - "2n2 ", - "s1 ", - ] - con_data = np.array([list(con) for con in con_list], dtype="S1") - - # Generate random data for the variables - lon_z_data = np.tile(np.linspace(-180, 180, nx), (ny, 1)).T - lat_z_data = np.tile(np.linspace(-90, 90, ny), (nx, 1)) - ha_data = np.random.rand(nc, nx, ny) - hp_data = np.random.rand(nc, nx, ny) * 360 # Random phases between 0 and 360 - hRe_data = np.random.rand(nc, nx, ny) - hIm_data = np.random.rand(nc, nx, ny) - - # Create the xarray dataset - ds_h = xr.Dataset( - { - "con": (["nc", "nct"], con_data), - "lon_z": (["nx", "ny"], lon_z_data), - "lat_z": (["nx", "ny"], lat_z_data), - "ha": (["nc", "nx", "ny"], ha_data), - "hp": (["nc", "nx", "ny"], hp_data), - "hRe": (["nc", "nx", "ny"], hRe_data), - "hIm": (["nc", "nx", "ny"], hIm_data), - }, - coords={ - "nc": np.arange(nc), - "nct": np.arange(nct), - "nx": np.arange(nx), - "ny": np.arange(ny), - }, - attrs={ - "type": "Fake OTIS tidal elevation file", - "title": "Fake TPXO9.v1 2018 tidal elevation file", - }, - ) - - # Generate random data for the variables for u_tpxo9.v1 - lon_u_data = ( - np.random.rand(nx, ny) * 360 - 180 - ) # Random longitudes between -180 and 180 - lat_u_data = ( - np.random.rand(nx, ny) * 180 - 90 - ) # Random latitudes between -90 and 90 - lon_v_data = ( - np.random.rand(nx, ny) * 360 - 180 - ) # Random longitudes between -180 and 180 - lat_v_data = ( - np.random.rand(nx, ny) * 180 - 90 - ) # Random latitudes between -90 and 90 - Ua_data = np.random.rand(nc, nx, ny) - ua_data = np.random.rand(nc, nx, ny) - up_data = np.random.rand(nc, nx, ny) * 360 # Random phases between 0 and 360 - Va_data = np.random.rand(nc, nx, ny) - va_data = np.random.rand(nc, nx, ny) - vp_data = np.random.rand(nc, nx, ny) * 360 # Random phases between 0 and 360 - URe_data = np.random.rand(nc, nx, ny) - UIm_data = np.random.rand(nc, nx, ny) - VRe_data = np.random.rand(nc, nx, ny) - VIm_data = np.random.rand(nc, nx, ny) - - # Create the xarray dataset for u_tpxo9.v1 - ds_u = xr.Dataset( - { - "con": (["nc", "nct"], con_data), - "lon_u": (["nx", "ny"], lon_u_data), - "lat_u": (["nx", "ny"], lat_u_data), - "lon_v": (["nx", "ny"], lon_v_data), - "lat_v": (["nx", "ny"], lat_v_data), - "Ua": (["nc", "nx", "ny"], Ua_data), - "ua": (["nc", "nx", "ny"], ua_data), - "up": (["nc", "nx", "ny"], up_data), - "Va": (["nc", "nx", "ny"], Va_data), - "va": (["nc", "nx", "ny"], va_data), - "vp": (["nc", "nx", "ny"], vp_data), - "URe": (["nc", "nx", "ny"], URe_data), - "UIm": (["nc", "nx", "ny"], UIm_data), - "VRe": (["nc", "nx", "ny"], VRe_data), - "VIm": (["nc", "nx", "ny"], VIm_data), - }, - coords={ - "nc": np.arange(nc), - "nct": np.arange(nct), - "nx": np.arange(nx), - "ny": np.arange(ny), - }, - attrs={ - "type": "Fake OTIS tidal transport file", - "title": "Fake TPXO9.v1 2018 WE/SN transports/currents file", - }, - ) - - return ds_h, ds_u - - -@pytest.fixture(scope="module") -def dummy_bathymetry_data(): - latitude_extent = [16.0, 27] - longitude_extent = [192, 209] - - bathymetry = np.random.random((100, 100)) * (-100) - bathymetry = xr.DataArray( - bathymetry, - dims=["silly_lat", "silly_lon"], - coords={ - "silly_lat": np.linspace( - latitude_extent[0] - 5, latitude_extent[1] + 5, 100 - ), - "silly_lon": np.linspace( - longitude_extent[0] - 5, longitude_extent[1] + 5, 100 - ), - }, - ) - bathymetry.name = "silly_depth" - return bathymetry - - -class TestAll: - - @pytest.fixture(scope="module") - def full_legit_expt_setup(self, dummy_bathymetry_data, tmp_path): - - expt_name = "testing" - - latitude_extent = [16.0, 27] - longitude_extent = [192, 209] - - date_range = ["2005-01-01 00:00:00", "2005-02-01 00:00:00"] - - ## Place where all your input files go - input_dir = Path( - os.path.join( - tmp_path, - expt_name, - "inputs", - ) - ) - - ## Directory where you'll run the experiment from - run_dir = Path( - os.path.join( - tmp_path, - expt_name, - "run_files", - ) - ) - data_path = Path(tmp_path / "data") - for path in (run_dir, input_dir, data_path): - os.makedirs(str(path), exist_ok=True) - bathy_path = data_path / "bathymetry.nc" - bathymetry = dummy_bathymetry_data - bathymetry.to_netcdf(bathy_path) - self.glorys_path = bathy_path - ## User-1st, test if we can even read the angled nc files. - expt = rmom6.experiment( - longitude_extent=longitude_extent, - latitude_extent=latitude_extent, - date_range=date_range, - resolution=0.05, - number_vertical_layers=75, - layer_thickness_ratio=10, - depth=4500, - minimum_depth=5, - mom_run_dir=run_dir, - mom_input_dir=input_dir, - toolpath_dir="", - ) - return expt - - def test_full_legit_expt_setup(self, tmp_path, dummy_bathymetry_data): - expt_name = "testing" - latitude_extent = [16.0, 27] - longitude_extent = [192, 209] - - date_range = ["2005-01-01 00:00:00", "2005-02-01 00:00:00"] - - ## Place where all your input files go - input_dir = Path( - os.path.join( - tmp_path, - expt_name, - "inputs", - ) - ) - - ## Directory where you'll run the experiment from - run_dir = Path( - os.path.join( - tmp_path, - expt_name, - "run_files", - ) - ) - data_path = Path(tmp_path / "data") - for path in (run_dir, input_dir, data_path): - os.makedirs(str(path), exist_ok=True) - bathy_path = data_path / "bathymetry.nc" - bathymetry = dummy_bathymetry_data - bathymetry.to_netcdf(bathy_path) - ## User-1st, test if we can even read the angled nc files. - expt = rmom6.experiment( - longitude_extent=longitude_extent, - latitude_extent=latitude_extent, - date_range=date_range, - resolution=0.05, - number_vertical_layers=75, - layer_thickness_ratio=10, - depth=4500, - minimum_depth=5, - mom_run_dir=run_dir, - mom_input_dir=input_dir, - toolpath_dir="", - ) - assert str(expt) - - def test_tides(self, dummy_tidal_data, tmp_path): - """ - Test the main setup tides function! - """ - expt_name = "testing" - ## User-1st, test if we can even read the angled nc files. - expt = rmom6.experiment.create_empty( - expt_name=expt_name, - mom_input_dir=tmp_path, - mom_run_dir=tmp_path, - ) - # Generate Fake Tidal Data - ds_h, ds_u = dummy_tidal_data - - # Save to Fake Folder - ds_h.to_netcdf(tmp_path / "h_fake_tidal_data.nc") - ds_u.to_netcdf(tmp_path / "u_fake_tidal_data.nc") - - # Set other required variables needed in setup_tides - - # Lat Long - expt.longitude_extent = (-5, 5) - expt.latitude_extent = (0, 30) - # Grid Type - expt.hgrid_type = "even_spacing" - # Datesƒ - expt.date_range = ("2000-01-01", "2000-01-02") - expt.segments = {} - # Generate Hgrid Data - expt.resolution = 0.1 - expt.hgrid = expt._make_hgrid() - # Create Forcing Folder - os.makedirs(tmp_path / "forcing", exist_ok=True) - - expt.setup_boundary_tides( - tmp_path / "h_fake_tidal_data.nc", - tmp_path / "u_fake_tidal_data.nc", - ) - - def test_change_MOM_parameter(self, tmp_path): - """ - Test the change MOM parameter function, as well as read_MOM_file and write_MOM_file under the hood. - """ - expt_name = "testing" - ## User-1st, test if we can even read the angled nc files. - expt = rmom6.experiment.create_empty( - expt_name=expt_name, - mom_input_dir=tmp_path, - mom_run_dir=tmp_path, - ) - # Copy over the MOM Files to the dump_files_dir - base_run_dir = Path( - os.path.join( - importlib.resources.files("regional_mom6").parent, - "demos", - "premade_run_directories", - ) - ) - shutil.copytree( - base_run_dir / "common_files", expt.mom_run_dir, dirs_exist_ok=True - ) - MOM_override_dict = expt.read_MOM_file_as_dict("MOM_override") - og = expt.change_MOM_parameter("DT", "30", "COOL COMMENT") - MOM_override_dict_new = expt.read_MOM_file_as_dict("MOM_override") - assert MOM_override_dict_new["DT"]["value"] == "30" - assert MOM_override_dict["DT"]["value"] == og - assert MOM_override_dict_new["DT"]["comment"] == "COOL COMMENT\n" - - def test_properties_empty(self, tmp_path): - """ - Test the properties - """ - expt_name = "testing" - ## User-1st, test if we can even read the angled nc files. - expt = rmom6.experiment.create_empty( - expt_name=expt_name, - mom_input_dir=tmp_path, - mom_run_dir=tmp_path, - ) - dss = expt.era5 - dss_2 = expt.tides_boundaries - dss_3 = expt.ocean_state_boundaries - dss_4 = expt.initial_condition - dss_5 = expt.bathymetry_property - print(dss, dss_2, dss_3, dss_4, dss_5) diff --git a/tests/test_regridding.py b/tests/test_regridding.py new file mode 100644 index 00000000..4d163ea5 --- /dev/null +++ b/tests/test_regridding.py @@ -0,0 +1,41 @@ +import regional_mom6 as rm6 +import regional_mom6.rotation as rot +import regional_mom6.regridding as rgd +import pytest +import xarray as xr +import numpy as np + + +# Not testing get_arakawa_c_points, coords, & create_regridder +def test_smoke_untested_funcs(get_curvilinear_hgrid): + hgrid = get_curvilinear_hgrid + assert rgd.get_hgrid_arakawa_c_points(hgrid, "t") + assert rgd.coords(hgrid, "north", "segment_002") + + +def test_fill_missing_data(): + return + + +def test_add_or_update_time_dim(): + return + + +def test_generate_dz(): + return + + +def test_add_secondary_dimension(): + return + + +def test_add_vertical_coordinate_encoding(): + return + + +def test_generate_layer_thickness(): + return + + +def test_generate_encoding(): + return diff --git a/tests/test_rotation.py b/tests/test_rotation.py new file mode 100644 index 00000000..ccf29d2d --- /dev/null +++ b/tests/test_rotation.py @@ -0,0 +1,233 @@ +import regional_mom6 as rm6 +import regional_mom6.rotation as rot +import regional_mom6.regridding as rgd +import pytest +import xarray as xr +import numpy as np +import os + + +def test_get_curvilinear_hgrid_fixture(get_curvilinear_hgrid): + # If the fixture fails to find the file, the test will be skipped. + assert get_curvilinear_hgrid is not None + + +def test_expanded_hgrid_generation(get_curvilinear_hgrid): + hgrid = get_curvilinear_hgrid + expanded_hgrid = rot.create_expanded_hgrid(hgrid) + + # Check Size + assert len(expanded_hgrid.nxp) == (len(hgrid.nxp) + 2) + assert len(expanded_hgrid.nyp) == (len(hgrid.nyp) + 2) + + # Check pseudo_hgrid keeps the same values + assert (expanded_hgrid.x.values[1:-1, 1:-1] == hgrid.x.values).all() + assert (expanded_hgrid.y.values[1:-1, 1:-1] == hgrid.y.values).all() + + # Check extra boundary has realistic values + diff_check = 1 + assert ( + ( + expanded_hgrid.x.values[0, 1:-1] + - (hgrid.x.values[0, :] - (hgrid.x.values[1, :] - hgrid.x.values[0, :])) + ) + < diff_check + ).all() + assert ( + ( + expanded_hgrid.x.values[1:-1, 0] + - (hgrid.x.values[:, 0] - (hgrid.x.values[:, 1] - hgrid.x.values[:, 0])) + ) + < diff_check + ).all() + assert ( + ( + expanded_hgrid.x.values[-1, 1:-1] + - (hgrid.x.values[-1, :] - (hgrid.x.values[-2, :] - hgrid.x.values[-1, :])) + ) + < diff_check + ).all() + assert ( + ( + expanded_hgrid.x.values[1:-1, -1] + - (hgrid.x.values[:, -1] - (hgrid.x.values[:, -2] - hgrid.x.values[:, -1])) + ) + < diff_check + ).all() + + # Check corners for the same... + assert ( + expanded_hgrid.x.values[0, 0] + - (hgrid.x.values[0, 0] - (hgrid.x.values[1, 1] - hgrid.x.values[0, 0])) + ) < diff_check + assert ( + expanded_hgrid.x.values[-1, 0] + - (hgrid.x.values[-1, 0] - (hgrid.x.values[-2, 1] - hgrid.x.values[-1, 0])) + ) < diff_check + assert ( + expanded_hgrid.x.values[0, -1] + - (hgrid.x.values[0, -1] - (hgrid.x.values[1, -2] - hgrid.x.values[0, -1])) + ) < diff_check + assert ( + expanded_hgrid.x.values[-1, -1] + - (hgrid.x.values[-1, -1] - (hgrid.x.values[-2, -2] - hgrid.x.values[-1, -1])) + ) < diff_check + + # Same for y + assert ( + ( + expanded_hgrid.y.values[0, 1:-1] + - (hgrid.y.values[0, :] - (hgrid.y.values[1, :] - hgrid.y.values[0, :])) + ) + < diff_check + ).all() + assert ( + ( + expanded_hgrid.y.values[1:-1, 0] + - (hgrid.y.values[:, 0] - (hgrid.y.values[:, 1] - hgrid.y.values[:, 0])) + ) + < diff_check + ).all() + assert ( + ( + expanded_hgrid.y.values[-1, 1:-1] + - (hgrid.y.values[-1, :] - (hgrid.y.values[-2, :] - hgrid.y.values[-1, :])) + ) + < diff_check + ).all() + assert ( + ( + expanded_hgrid.y.values[1:-1, -1] + - (hgrid.y.values[:, -1] - (hgrid.y.values[:, -2] - hgrid.y.values[:, -1])) + ) + < diff_check + ).all() + + assert ( + expanded_hgrid.y.values[0, 0] + - (hgrid.y.values[0, 0] - (hgrid.y.values[1, 1] - hgrid.y.values[0, 0])) + ) < diff_check + assert ( + expanded_hgrid.y.values[-1, 0] + - (hgrid.y.values[-1, 0] - (hgrid.y.values[-2, 1] - hgrid.y.values[-1, 0])) + ) < diff_check + assert ( + expanded_hgrid.y.values[0, -1] + - (hgrid.y.values[0, -1] - (hgrid.y.values[1, -2] - hgrid.y.values[0, -1])) + ) < diff_check + assert ( + expanded_hgrid.y.values[-1, -1] + - (hgrid.y.values[-1, -1] - (hgrid.y.values[-2, -2] - hgrid.y.values[-1, -1])) + ) < diff_check + + return + + +def test_mom6_angle_calculation_method(get_curvilinear_hgrid): + """ + Check no rotation, up tilt, down tilt. + """ + + # Check no rotation + top_left = xr.Dataset( + { + "x": (("nyp", "nxp"), [[0]]), + "y": (("nyp", "nxp"), [[1]]), + } + ) + top_right = xr.Dataset( + { + "x": (("nyp", "nxp"), [[1]]), + "y": (("nyp", "nxp"), [[1]]), + } + ) + bottom_left = xr.Dataset( + { + "x": (("nyp", "nxp"), [[0]]), + "y": (("nyp", "nxp"), [[0]]), + } + ) + bottom_right = xr.Dataset( + { + "x": (("nyp", "nxp"), [[1]]), + "y": (("nyp", "nxp"), [[0]]), + } + ) + point = xr.Dataset( + { + "x": (("nyp", "nxp"), [[0.5]]), + "y": (("nyp", "nxp"), [[0.5]]), + } + ) + + assert ( + rot.mom6_angle_calculation_method( + 2, top_left, top_right, bottom_left, bottom_right, point + ) + == 0 + ) + + # Angled + hgrid = get_curvilinear_hgrid + ds_t = rgd.get_hgrid_arakawa_c_points(hgrid, "t") + ds_q = rgd.get_hgrid_arakawa_c_points(hgrid, "q") + + # Reformat into x, y comps + t_points = xr.Dataset( + { + "x": (("nyp", "nxp"), ds_t.tlon.data), + "y": (("nyp", "nxp"), ds_t.tlat.data), + } + ) + q_points = xr.Dataset( + { + "x": (("nyp", "nxp"), ds_q.qlon.data), + "y": (("nyp", "nxp"), ds_q.qlat.data), + } + ) + assert ( + ( + rot.mom6_angle_calculation_method( + hgrid.x.max() - hgrid.x.min(), + q_points.isel(nyp=slice(1, None), nxp=slice(0, -1)), + q_points.isel(nyp=slice(1, None), nxp=slice(1, None)), + q_points.isel(nyp=slice(0, -1), nxp=slice(0, -1)), + q_points.isel(nyp=slice(0, -1), nxp=slice(1, None)), + t_points, + ) + - hgrid["angle_dx"].isel(nyp=ds_t.t_points_y, nxp=ds_t.t_points_x).values + ) + < 1 + ).all() + + return + + +def test_initialize_grid_rotation_angle(get_curvilinear_hgrid): + """ + Generate a curvilinear grid and test the grid rotation angle at t_points based on what we pass to generate_curvilinear_grid + """ + hgrid = get_curvilinear_hgrid + angle = rot.initialize_grid_rotation_angle(hgrid) + ds_t = rgd.get_hgrid_arakawa_c_points(hgrid, "t") + assert ( + ( + angle.values + - hgrid["angle_dx"].isel(nyp=ds_t.t_points_y, nxp=ds_t.t_points_x).values + ) + < 1 + ).all() # Angle is correct + assert angle.values.shape == ds_t.tlon.shape # Shape is correct + return + + +def test_initialize_grid_rotation_angle_using_expanded_hgrid(get_curvilinear_hgrid): + """ + Generate a curvilinear grid and test the grid rotation angle at t_points based on what we pass to generate_curvilinear_grid + """ + hgrid = get_curvilinear_hgrid + angle = rot.initialize_grid_rotation_angles_using_expanded_hgrid(hgrid) + + assert (angle.values - hgrid.angle_dx < 1).all() + assert angle.values.shape == hgrid.x.shape + return diff --git a/tests/test_tides_and_parameter.py b/tests/test_tides_and_parameter.py new file mode 100644 index 00000000..8b1e0f99 --- /dev/null +++ b/tests/test_tides_and_parameter.py @@ -0,0 +1,278 @@ +""" +Test suite for everything involed in pr #12 +""" + +import regional_mom6 as rmom6 +import os +import pytest +import logging +from pathlib import Path +import xarray as xr +import numpy as np +import shutil +import importlib + +IN_GITHUB_ACTIONS = os.getenv("GITHUB_ACTIONS") == "true" +# @pytest.mark.skipif( +# IN_GITHUB_ACTIONS, reason="Test doesn't work in Github Actions." +# ) + + +@pytest.fixture(scope="module") +def dummy_tidal_data(): + nx = 100 + ny = 100 + nc = 15 + nct = 4 + + # Define tidal constituents + con_list = [ + "m2 ", + "s2 ", + "n2 ", + "k2 ", + "k1 ", + "o1 ", + "p1 ", + "q1 ", + "mm ", + "mf ", + "m4 ", + "mn4 ", + "ms4 ", + "2n2 ", + "s1 ", + ] + con_data = np.array([list(con) for con in con_list], dtype="S1") + + # Generate random data for the variables + lon_z_data = np.tile(np.linspace(-180, 180, nx), (ny, 1)).T + lat_z_data = np.tile(np.linspace(-90, 90, ny), (nx, 1)) + ha_data = np.random.rand(nc, nx, ny) + hp_data = np.random.rand(nc, nx, ny) * 360 # Random phases between 0 and 360 + hRe_data = np.random.rand(nc, nx, ny) + hIm_data = np.random.rand(nc, nx, ny) + + # Create the xarray dataset + ds_h = xr.Dataset( + { + "con": (["nc", "nct"], con_data), + "lon_z": (["nx", "ny"], lon_z_data), + "lat_z": (["nx", "ny"], lat_z_data), + "ha": (["nc", "nx", "ny"], ha_data), + "hp": (["nc", "nx", "ny"], hp_data), + "hRe": (["nc", "nx", "ny"], hRe_data), + "hIm": (["nc", "nx", "ny"], hIm_data), + }, + coords={ + "nc": np.arange(nc), + "nct": np.arange(nct), + "nx": np.arange(nx), + "ny": np.arange(ny), + }, + attrs={ + "type": "Fake OTIS tidal elevation file", + "title": "Fake TPXO9.v1 2018 tidal elevation file", + }, + ) + + # Generate random data for the variables for u_tpxo9.v1 + lon_u_data = ( + np.random.rand(nx, ny) * 360 - 180 + ) # Random longitudes between -180 and 180 + lat_u_data = ( + np.random.rand(nx, ny) * 180 - 90 + ) # Random latitudes between -90 and 90 + lon_v_data = ( + np.random.rand(nx, ny) * 360 - 180 + ) # Random longitudes between -180 and 180 + lat_v_data = ( + np.random.rand(nx, ny) * 180 - 90 + ) # Random latitudes between -90 and 90 + Ua_data = np.random.rand(nc, nx, ny) + ua_data = np.random.rand(nc, nx, ny) + up_data = np.random.rand(nc, nx, ny) * 360 # Random phases between 0 and 360 + Va_data = np.random.rand(nc, nx, ny) + va_data = np.random.rand(nc, nx, ny) + vp_data = np.random.rand(nc, nx, ny) * 360 # Random phases between 0 and 360 + URe_data = np.random.rand(nc, nx, ny) + UIm_data = np.random.rand(nc, nx, ny) + VRe_data = np.random.rand(nc, nx, ny) + VIm_data = np.random.rand(nc, nx, ny) + + # Create the xarray dataset for u_tpxo9.v1 + ds_u = xr.Dataset( + { + "con": (["nc", "nct"], con_data), + "lon_u": (["nx", "ny"], lon_u_data), + "lat_u": (["nx", "ny"], lat_u_data), + "lon_v": (["nx", "ny"], lon_v_data), + "lat_v": (["nx", "ny"], lat_v_data), + "Ua": (["nc", "nx", "ny"], Ua_data), + "ua": (["nc", "nx", "ny"], ua_data), + "up": (["nc", "nx", "ny"], up_data), + "Va": (["nc", "nx", "ny"], Va_data), + "va": (["nc", "nx", "ny"], va_data), + "vp": (["nc", "nx", "ny"], vp_data), + "URe": (["nc", "nx", "ny"], URe_data), + "UIm": (["nc", "nx", "ny"], UIm_data), + "VRe": (["nc", "nx", "ny"], VRe_data), + "VIm": (["nc", "nx", "ny"], VIm_data), + }, + coords={ + "nc": np.arange(nc), + "nct": np.arange(nct), + "nx": np.arange(nx), + "ny": np.arange(ny), + }, + attrs={ + "type": "Fake OTIS tidal transport file", + "title": "Fake TPXO9.v1 2018 WE/SN transports/currents file", + }, + ) + + return ds_h, ds_u + + +@pytest.fixture(scope="module") +def dummy_bathymetry_data(): + latitude_extent = [16.0, 27] + longitude_extent = [192, 209] + + bathymetry = np.random.random((100, 100)) * (-100) + bathymetry = xr.DataArray( + bathymetry, + dims=["silly_lat", "silly_lon"], + coords={ + "silly_lat": np.linspace( + latitude_extent[0] - 5, latitude_extent[1] + 5, 100 + ), + "silly_lon": np.linspace( + longitude_extent[0] - 5, longitude_extent[1] + 5, 100 + ), + }, + ) + bathymetry.name = "silly_depth" + return bathymetry + + +@pytest.fixture() +def full_expt_setup(dummy_bathymetry_data, tmp_path): + + expt_name = "testing" + + latitude_extent = [16.0, 27] + longitude_extent = [192, 209] + + date_range = ["2005-01-01 00:00:00", "2005-02-01 00:00:00"] + + ## Place where all your input files go + input_dir = Path( + os.path.join( + tmp_path, + expt_name, + "inputs", + ) + ) + + ## Directory where you'll run the experiment from + run_dir = Path( + os.path.join( + tmp_path, + expt_name, + "run_files", + ) + ) + data_path = Path(tmp_path, "data") + for path in (run_dir, input_dir, data_path): + os.makedirs(str(path), exist_ok=True) + bathy_path = data_path / "bathymetry.nc" + bathymetry = dummy_bathymetry_data + bathymetry.to_netcdf(bathy_path) + ## User-1st, test if we can even read the angled nc files. + expt = rmom6.experiment( + longitude_extent=longitude_extent, + latitude_extent=latitude_extent, + date_range=date_range, + resolution=0.05, + number_vertical_layers=75, + layer_thickness_ratio=10, + depth=4500, + minimum_depth=5, + mom_run_dir=run_dir, + mom_input_dir=input_dir, + toolpath_dir="", + ) + return expt + + +def test_full_expt_setup(full_expt_setup): + assert str(full_expt_setup) + + +def test_tides(dummy_tidal_data, tmp_path): + """ + Test the main setup tides function! + """ + expt_name = "testing" + + expt = rmom6.experiment.create_empty( + expt_name=expt_name, + mom_input_dir=tmp_path, + mom_run_dir=tmp_path, + ) + # Generate Fake Tidal Data + ds_h, ds_u = dummy_tidal_data + + # Save to Fake Folder + ds_h.to_netcdf(tmp_path / "h_fake_tidal_data.nc") + ds_u.to_netcdf(tmp_path / "u_fake_tidal_data.nc") + + # Set other required variables needed in setup_tides + + # Lat Long + expt.longitude_extent = (-5, 5) + expt.latitude_extent = (0, 30) + # Grid Type + expt.hgrid_type = "even_spacing" + # Dates + expt.date_range = ("2000-01-01", "2000-01-02") + expt.segments = {} + # Generate Hgrid Data + expt.resolution = 0.1 + expt.hgrid = expt._make_hgrid() + # Create Forcing Folder + os.makedirs(tmp_path / "forcing", exist_ok=True) + + expt.setup_boundary_tides( + tmp_path / "h_fake_tidal_data.nc", + tmp_path / "u_fake_tidal_data.nc", + ) + + +def test_change_MOM_parameter(tmp_path): + """ + Test the change MOM parameter function, as well as read_MOM_file and write_MOM_file under the hood. + """ + expt_name = "testing" + + expt = rmom6.experiment.create_empty( + expt_name=expt_name, + mom_input_dir=tmp_path, + mom_run_dir=tmp_path, + ) + # Copy over the MOM Files to the dump_files_dir + base_run_dir = Path( + os.path.join( + importlib.resources.files("regional_mom6").parent, + "demos", + "premade_run_directories", + ) + ) + shutil.copytree(base_run_dir / "common_files", expt.mom_run_dir, dirs_exist_ok=True) + MOM_override_dict = expt.read_MOM_file_as_dict("MOM_override") + og = expt.change_MOM_parameter("DT", "30", "COOL COMMENT") + MOM_override_dict_new = expt.read_MOM_file_as_dict("MOM_override") + assert MOM_override_dict_new["DT"]["value"] == "30" + assert MOM_override_dict["DT"]["value"] == og + assert MOM_override_dict_new["DT"]["comment"] == "COOL COMMENT\n"