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/demos/premade_run_directories/README.md b/demos/premade_run_directories/README.md index a3488a41..4a1da513 100644 --- a/demos/premade_run_directories/README.md +++ b/demos/premade_run_directories/README.md @@ -1,4 +1,4 @@ -## Premade Run Directories +# Premade Run Directories These directories are used for the demo notebooks, and can be used as templates for setting up a new experiment. The [documentation](https://regional-mom6.readthedocs.io/en/latest/mom6-file-structure-primer.html) explains what all the files are for. diff --git a/demos/reanalysis-forced.ipynb b/demos/reanalysis-forced.ipynb index a4d161f2..245b0964 100644 --- a/demos/reanalysis-forced.ipynb +++ b/demos/reanalysis-forced.ipynb @@ -44,10 +44,12 @@ }, { "cell_type": "code", - "execution_count": 3, + "execution_count": 1, "metadata": {}, "outputs": [], "source": [ + "import warnings\n", + "warnings.filterwarnings(\"ignore\")\n", "import regional_mom6 as rmom6\n", "\n", "import os\n", @@ -55,6 +57,17 @@ "from dask.distributed import Client" ] }, + { + "cell_type": "code", + "execution_count": null, + "metadata": {}, + "outputs": [], + "source": [ + "# Currently, only the regional_mom6 module reports logging information to the info level, for more detailed output\n", + "# import logging\n", + "# logging.basicConfig(level=logging.INFO) \n" + ] + }, { "cell_type": "markdown", "metadata": {}, @@ -191,7 +204,7 @@ "metadata": {}, "outputs": [], "source": [ - "expt.get_glorys_rectangular(\n", + "expt.get_glorys(\n", " raw_boundaries_path=glorys_path\n", ")" ] @@ -428,7 +441,7 @@ ], "metadata": { "kernelspec": { - "display_name": "crr_dev", + "display_name": "CrocoDash", "language": "python", "name": "python3" }, @@ -442,7 +455,7 @@ "name": "python", "nbconvert_exporter": "python", "pygments_lexer": "ipython3", - "version": "3.12.7" + "version": "3.12.8" } }, "nbformat": 4, diff --git a/docs/angle_calc.md b/docs/angle_calc.md new file mode 100644 index 00000000..089cda3f --- /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. The EXPAND_GRID 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..ce2095ea 100644 --- a/docs/index.rst +++ b/docs/index.rst @@ -91,8 +91,11 @@ The bibtex entry for the paper is: installation demos mom6-file-structure-primer + rm6_workflow + angle_calc api contributing + docker_image_dev Indices and tables diff --git a/docs/mom6-file-structure-primer.md b/docs/mom6-file-structure-primer.md index 708fdebd..102c019f 100644 --- a/docs/mom6-file-structure-primer.md +++ b/docs/mom6-file-structure-primer.md @@ -106,3 +106,9 @@ These files can be big, so it is usually helpful to store them somewhere without confusing, and getting them wrong can likewise cause some cryptic error messages! These boundaries do not have to follow lines of constant longitude and latitude, but it is much easier to set things up if they do. For an example of a curved boundary, see this [Northwest Atlantic experiment](https://github.com/jsimkins2/nwa25/tree/main). + +* `forcing/{tz/tu}_segment**` + The boundary tidal segments, numbered the same way as in `MOM_input`. The dimensions and coordinates are fairly + confusing, and getting them wrong can likewise cause some cryptic error messages! These boundaries do not have to + follow lines of constant longitude and latitude, but it is much easier to set things up if they do. For an example + of a curved boundary, see this [Northwest Atlantic experiment](https://github.com/jsimkins2/nwa25/tree/main). diff --git a/docs/rm6_workflow.md b/docs/rm6_workflow.md new file mode 100644 index 00000000..647a12d2 --- /dev/null +++ b/docs/rm6_workflow.md @@ -0,0 +1,32 @@ +# Regional MOM6 Workflow + +Regional MOM6(RM6) sets up all the data and files for running a basic regional case of MOM6. + +It includes: + +1. Run Files like MOM_override, MOM_input, diag_table +2. BC File like velocity, tracers, tides +3. Basic input files like hgrid & bathymetry. +4. Initial Condition files + + +To set up a case with all the files, RM6 has its own way of grabbing and organizing files for the user. + +RM6 organizes all files into two directories, a "input" directory, and a "run" directory. The input directory includes all of the data we need for our regional case. The run directory includes all of the parameters and outputs (diags) we want for our model. Please see the structure primer document for more information. + +The other folders are the data directories. RM6 needs the user to collect the initial condition & boundary condition data and put them into a folder(s). + +Therefore, to start for the user to use RM6, they should have two (empty or not) directories for the input and run files, as well as directories for their input data. + +Depending on the computer used (NCAR-Derecho/Casper doesn't need this), the user may also need to provide a path to FRE_tools. + +To create all these files, RM6 using a class called "Experiment" to hold all of the parameters and functions. Users can follow a few quick steps to setup their cases: +1. Initalize the experiment object with all the directories and parameters wanted. The initalization can also create the hgrid and vertical coordinate or accept two files called "hgrid.nc" and "vcoord.nc" in the input directory. +2. Call different "setup..." functions to setup all the data needed for the case (bathymetry, initial condition, velocity, tracers, tides). +3. Finally, call "setup_run_directory" to setup the run files like MOM_override for their cases. +4. Based on how MOM6 is setup on the computer, there are follow-up steps unique to each situation. RM6 provides all of what the user needs to run MOM6. + + +There are a few convience functions to help support the process. +1. Very light read and write config file functions to easily save experiments +2. A change_MOM_parameter function to easily adjust MOM parameter values from python. \ No newline at end of file diff --git a/regional_mom6/regional_mom6.py b/regional_mom6/regional_mom6.py index 80986c43..4bda003c 100644 --- a/regional_mom6/regional_mom6.py +++ b/regional_mom6/regional_mom6.py @@ -14,13 +14,22 @@ import os import importlib.resources import datetime -from .utils import quadrilateral_areas, ap2ep, ep2ap import pandas as pd from pathlib import Path import glob from collections import defaultdict import json import copy +from regional_mom6 import regridding as rgd +from regional_mom6 import rotation as rot +from regional_mom6.utils import ( + quadrilateral_areas, + ap2ep, + ep2ap, + is_rectilinear_hgrid, + rotate, +) + warnings.filterwarnings("ignore") @@ -589,7 +598,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"], ): @@ -655,7 +664,7 @@ def __init__( vgrid_path=None, 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"], @@ -695,7 +704,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": if hgrid_path is None: hgrid_path = self.mom_input_dir / "hgrid.nc" @@ -932,6 +940,11 @@ def _make_vgrid(self, thicknesses=None): of largest to smallest layer thickness (``layer_thickness_ratio``) and the total ``depth`` parameters. (All these parameters are specified at the class level.) + + Args: + thicknesses (Optional[np.ndarray]): An array of layer thicknesses. If not provided, + the layer thicknesses are generated using the :func:`~hyperbolictan_thickness_profile` + function. """ if thicknesses is None: @@ -939,6 +952,9 @@ def _make_vgrid(self, thicknesses=None): self.number_vertical_layers, self.layer_thickness_ratio, self.depth ) + if not isinstance(thicknesses, np.ndarray): + raise ValueError("thicknesses must be a numpy array") + zi = np.cumsum(thicknesses) zi = np.insert(zi, 0, 0.0) # add zi = 0.0 as first interface @@ -1134,6 +1150,7 @@ def setup_initial_condition( varnames, arakawa_grid="A", vcoord_type="height", + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ Reads the initial condition from files in ``ic_path``, interpolates to the @@ -1147,6 +1164,7 @@ def setup_initial_condition( Either ``'A'`` (default), ``'B'``, or ``'C'``. vcoord_type (Optional[str]): The type of vertical coordinate used in the forcing files. Either ``'height'`` or ``'thickness'``. + rotational_method (Optional[RotationMethod]): The method used to rotate the velocities. """ # Remove time dimension if present in the IC. @@ -1269,28 +1287,6 @@ def setup_initial_condition( + "Terminating!" ) - ## Construct the xq, yh and xh, yq grids - ugrid = ( - self.hgrid[["x", "y"]] - .isel(nxp=slice(None, None, 2), nyp=slice(1, None, 2)) - .rename({"x": "lon", "y": "lat"}) - .set_coords(["lat", "lon"]) - ) - vgrid = ( - self.hgrid[["x", "y"]] - .isel(nxp=slice(1, None, 2), nyp=slice(None, None, 2)) - .rename({"x": "lon", "y": "lat"}) - .set_coords(["lat", "lon"]) - ) - - ## Construct the cell centre grid for tracers (xh, yh). - tgrid = ( - self.hgrid[["x", "y"]] - .isel(nxp=slice(1, None, 2), nyp=slice(1, None, 2)) - .rename({"x": "lon", "y": "lat", "nxp": "nx", "nyp": "ny"}) - .set_coords(["lat", "lon"]) - ) - # NaNs might be here from the land mask of the model that the IC has come from. # If they're not removed then the coastlines from this other grid will be retained! # The land mask comes from the bathymetry file, so we don't need NaNs @@ -1330,38 +1326,64 @@ def setup_initial_condition( .bfill("lat") ) + self.hgrid["lon"] = self.hgrid["x"] + self.hgrid["lat"] = self.hgrid["y"] + tgrid = ( + rgd.get_hgrid_arakawa_c_points(self.hgrid, "t") + .rename({"tlon": "lon", "tlat": "lat", "nxp": "nx", "nyp": "ny"}) + .set_coords(["lat", "lon"]) + ) + ## Make our three horizontal regridders - regridder_u = xe.Regridder( - ic_raw_u, - ugrid, - "bilinear", + + regridder_u = rgd.create_regridder( + ic_raw_u, self.hgrid, locstream_out=False, method="bilinear" ) - regridder_v = xe.Regridder( - ic_raw_v, - vgrid, - "bilinear", + regridder_v = rgd.create_regridder( + ic_raw_v, self.hgrid, locstream_out=False, method="bilinear" ) + regridder_t = rgd.create_regridder( + ic_raw_tracers, tgrid, locstream_out=False, method="bilinear" + ) # Doesn't need to be rotated, so we can regrid to just tracers - regridder_t = xe.Regridder( - ic_raw_tracers, - tgrid, - "bilinear", - ) + # ugrid= rgd.get_hgrid_arakawa_c_points(self.hgrid, "u").rename({"ulon": "lon", "ulat": "lat"}).set_coords(["lat", "lon"]) + # vgrid = rgd.get_hgrid_arakawa_c_points(self.hgrid, "v").rename({"vlon": "lon", "vlat": "lat"}).set_coords(["lat", "lon"]) - print("INITIAL CONDITIONS") + ## Construct the cell centre grid for tracers (xh, yh). + print("Setting up Initial Conditions") ## Regrid all fields horizontally. print("Regridding Velocities... ", end="") + regridded_u = regridder_u(ic_raw_u) + regridded_v = regridder_v(ic_raw_v) + rotated_u, rotated_v = rotate( + regridded_u, + regridded_v, + radian_angle=np.radians( + rot.get_rotation_angle(rotational_method, self.hgrid).values + ), + ) + # Slice the velocites to the u and v grid. + u_points = rgd.get_hgrid_arakawa_c_points(self.hgrid, "u") + v_points = rgd.get_hgrid_arakawa_c_points(self.hgrid, "v") + rotated_v = rotated_v[:, v_points.v_points_y.values, v_points.v_points_x.values] + rotated_u = rotated_u[:, u_points.u_points_y.values, u_points.u_points_x.values] + rotated_u["lon"] = u_points.ulon + rotated_u["lat"] = u_points.ulat + rotated_v["lon"] = v_points.vlon + rotated_v["lat"] = v_points.vlat + + # Merge Vels vel_out = xr.merge( [ - regridder_u(ic_raw_u) - .rename({"lon": "xq", "lat": "yh", "nyp": "ny", varnames["zl"]: "zl"}) - .rename("u"), - regridder_v(ic_raw_v) - .rename({"lon": "xh", "lat": "yq", "nxp": "nx", varnames["zl"]: "zl"}) - .rename("v"), + rotated_u.rename( + {"lon": "xq", "lat": "yh", "nyp": "ny", varnames["zl"]: "zl"} + ).rename("u"), + rotated_v.rename( + {"lon": "xh", "lat": "yq", "nxp": "nx", varnames["zl"]: "zl"} + ).rename("v"), ] ) @@ -1420,14 +1442,11 @@ def setup_initial_condition( eta_out.attrs = ic_raw_eta.attrs ## Regrid the fields vertically - if ( vcoord_type == "thickness" ): ## In this case construct the vertical profile by summing thickness tracers_out["zl"] = tracers_out["zl"].diff("zl") - dz = tracers_out[self.z].diff(self.z) - dz.name = "dz" - dz = xr.concat([dz, dz[-1]], dim=self.z) + dz = rgd.generate_dz(tracers_out, self.z) tracers_out = tracers_out.interp({"zl": self.vgrid.zl.values}) vel_out = vel_out.interp({"zl": self.vgrid.zl.values}) @@ -1472,7 +1491,7 @@ def setup_initial_condition( return - def get_glorys_rectangular(self, raw_boundaries_path): + def get_glorys(self, raw_boundaries_path): """ This function is a wrapper for `get_glorys_data`, calling this function once for each of the rectangular boundary segments and the initial condition. For more complex boundary shapes, call `get_glorys_data` directly for each of your boundaries that aren't parallel to lines of constant latitude or longitude. For example, for an angled Northern boundary that spans multiple latitudes, you'll need to download a wider rectangle containing the entire boundary. @@ -1570,6 +1589,8 @@ def setup_ocean_state_boundaries( varnames, arakawa_grid="A", boundary_type="rectangular", + bathymetry_path=None, + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ This function is a wrapper for `simple_boundary`. Given a list of up to four cardinal directions, @@ -1584,11 +1605,13 @@ def setup_ocean_state_boundaries( Default is `["south", "north", "west", "east"]`. 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. + boundary_type (Optional[str]): Type of box around region. Currently, only ``'rectangular'`` or ``'curvilinear'`` is supported. + bathymetry_path (Optional[str]): Path to the bathymetry file. Default is None, in which case the BC is not masked. + rotational_method (Optional[str]): Method to use for rotating the boundary velocities. Default is 'GIVEN_ANGLE'. """ - if boundary_type != "rectangular": + if not (boundary_type == "rectangular" or boundary_type == "curvilinear"): raise ValueError( - "Only rectangular boundaries are supported by this method. To set up more complex boundary shapes you can manually call the 'simple_boundary' method for each boundary." + "Only rectangular or curvilinear boundaries are supported by this method. To set up more complex boundary shapes you can manually call the 'simple_boundary' method for each boundary." ) for i in self.boundaries: if i not in ["south", "north", "west", "east"]: @@ -1620,6 +1643,8 @@ 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, ) def setup_single_boundary( @@ -1629,7 +1654,8 @@ def setup_single_boundary( orientation, segment_number, arakawa_grid="A", - boundary_type="simple", + bathymetry_path=None, + rotational_method=rot.RotationMethod.GIVEN_ANGLE, ): """ Here 'simple' refers to boundaries that are parallel to lines of constant longitude or latitude. @@ -1648,18 +1674,20 @@ def setup_single_boundary( the ``MOM_input``. 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 BC is not masked + rotational_method (Optional[str]): Method to use for rotating the boundary velocities. Default is 'GIVEN_ANGLE'. """ - print("Processing {} boundary...".format(orientation), end="") + print( + "Processing {} boundary velocity & tracers...".format(orientation), end="" + ) if not path_to_bc.exists(): raise FileNotFoundError( f"Boundary file not found at {path_to_bc}. Please ensure that the files are named in the format `east_unprocessed.nc`." ) - if boundary_type != "simple": - 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, @@ -1670,35 +1698,41 @@ 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 def setup_boundary_tides( self, - path_to_td, - tidal_filename, + tpxo_elevation_filepath, + tpxo_velocity_filepath, tidal_constituents="read_from_expt_init", - boundary_type="rectangle", + boundary_type="rectangular", + bathymetry_path=None, + 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. + tpxo_elevation_filepath: Filepath to the TPXO elevation product. Generally of the form h_tidalversion.nc + tpxo_velocity_filepath: Filepath to the TPXO velocity product. Generally of the form u_tidalversion.nc + 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. Curvilinear is also suported. + bathymetry_path (str): Path to the bathymetry file. Default is None, in which case the BC is not masked + 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' + .nc files: Regridded tidal velocity and elevation files in 'inputdir/forcing' 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) + - 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, coords, segment.regrid_tides, segment.encode_tidal_files_and_output) Original Code was sourced from: @@ -1709,14 +1743,14 @@ def setup_boundary_tides( Type: Python Functions, Source Code Web Address: https://github.com/jsimkins2/nwa25 """ - if boundary_type != "rectangle": + if not (boundary_type == "rectangular" or boundary_type == "curvilinear"): raise ValueError( - "Only rectangular boundaries are supported by this method." + "Only rectangular or curvilinear boundaries are supported by this method." ) if tidal_constituents != "read_from_expt_init": self.tidal_constituents = tidal_constituents tpxo_h = ( - xr.open_dataset(Path(path_to_td / f"h_{tidal_filename}")) + xr.open_dataset(Path(tpxo_elevation_filepath)) .rename({"lon_z": "lon", "lat_z": "lat", "nc": "constituent"}) .isel( constituent=convert_to_tpxo_tidal_constituents(self.tidal_constituents) @@ -1727,7 +1761,7 @@ def setup_boundary_tides( tpxo_h["hRe"] = np.real(h) tpxo_h["hIm"] = np.imag(h) tpxo_u = ( - xr.open_dataset(Path(path_to_td / f"u_{tidal_filename}")) + xr.open_dataset(Path(tpxo_velocity_filepath)) .rename({"lon_u": "lon", "lat_u": "lat", "nc": "constituent"}) .isel( constituent=convert_to_tpxo_tidal_constituents(self.tidal_constituents) @@ -1738,7 +1772,7 @@ def setup_boundary_tides( tpxo_u["uRe"] = np.real(u) tpxo_u["uIm"] = np.imag(u) tpxo_v = ( - xr.open_dataset(Path(path_to_td / f"u_{tidal_filename}")) + xr.open_dataset(Path(tpxo_velocity_filepath)) .rename({"lon_v": "lon", "lat_v": "lat", "nc": "constituent"}) .isel( constituent=convert_to_tpxo_tidal_constituents(self.tidal_constituents) @@ -1759,25 +1793,24 @@ def setup_boundary_tides( print("Processing {} boundary...".format(b), end="") # If the GLORYS ocean_state has already created segments, we don't create them again. - if b not in self.segments.keys(): # I.E. not set yet - seg = segment( - hgrid=self.hgrid, - infile=None, # location of raw boundary - outfolder=self.mom_input_dir, - varnames=None, - segment_name="segment_{:03d}".format( - self.find_MOM6_rectangular_orientation(b) - ), - orientation=b, # orienataion - startdate=self.date_range[0], - repeat_year_forcing=self.repeat_year_forcing, - ) - self.segments[b] = seg - else: - seg = self.segments[b] + seg = segment( + hgrid=self.hgrid, + bathymetry_path=bathymetry_path, + infile=None, # location of raw boundary + outfolder=self.mom_input_dir, + varnames=None, + segment_name="segment_{:03d}".format( + self.find_MOM6_rectangular_orientation(b) + ), + orientation=b, # orienataion + startdate=self.date_range[0], + repeat_year_forcing=self.repeat_year_forcing, + ) # 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( @@ -2264,7 +2297,8 @@ def setup_run_directory( if not premade_rundir_path.exists(): print("Could not find premade run directories at ", premade_rundir_path) print( - "Perhaps the package was imported directly rather than installed with conda. Checking if this is the case... " + "Perhaps the package was imported directly rather than installed with conda. Checking if this is the case... ", + end="", ) premade_rundir_path = Path( @@ -2301,13 +2335,13 @@ def setup_run_directory( # Check if we can implement tides if with_tides: - tidal_files_exist = any( - "tidal" in filename - for filename in ( - os.listdir(Path(self.mom_input_dir / "forcing")) - + os.listdir(Path(self.mom_input_dir)) - ) - ) + if not (self.mom_input_dir / "forcing").exists(): + all_files = os.listdir(Path(self.mom_input_dir)) + else: + all_files = os.listdir( + Path(self.mom_input_dir / "forcing") + ) + os.listdir(Path(self.mom_input_dir)) + tidal_files_exist = any("tidal" in filename for filename in all_files) if not tidal_files_exist: raise ( "No files with 'tidal' in their names found in the forcing or input directory. If you meant to use tides, please run the setup_tides_rectangle_boundaries method first. That does output some tidal files." @@ -2554,6 +2588,19 @@ def setup_run_directory( 0, ] nml.write(self.mom_run_dir / "input.nml", force=True) + + # Edit Diag Table Date + # Read the file + with open(self.mom_run_dir / "diag_table", "r") as file: + lines = file.readlines() + + # The date is the second line + lines[1] = self.date_range[0].strftime("%Y %-m %-d %-H %-M %-S\n") + + # Write the file + with open(self.mom_run_dir / "diag_table", "w") as file: + file.writelines(lines) + return def change_MOM_parameter( @@ -2692,13 +2739,13 @@ def write_MOM_file(self, MOM_file_dict): ) print( - "Changed", - var, - "from", - original_MOM_file_dict[var], - "to", - MOM_file_dict[var], - "in {}!".format(MOM_file_dict["filename"]), + "Changed " + + str(var) + + " from " + + str(original_MOM_file_dict[var]["value"]) + + " to " + + str(MOM_file_dict[var]["value"]) + + "in {}!".format(str(MOM_file_dict["filename"])) ) # Add new fields @@ -2898,6 +2945,7 @@ def __init__( segment_name, orientation, startdate, + bathymetry_path=None, arakawa_grid="A", time_units="days", repeat_year_forcing=False, @@ -2947,170 +2995,73 @@ 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 - @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): - # 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. - - 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) - 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 """ 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 - regridder = xe.Regridder( + # In this case velocities and tracers all on same points + regridder = rgd.create_regridder( rawseg[self.u], - self.coords, - "bilinear", - locstream_out=True, - reuse_weights=False, - filename=self.outfolder + coords, + self.outfolder / f"weights/bilinear_velocity_weights_{self.orientation}.nc", + method="bilinear", ) - segment_out = xr.merge( - [ - regridder( - rawseg[ - [self.u, self.v, self.eta] - + [self.tracers[i] for i in self.tracers] - ] - ) + regridded = regridder( + rawseg[ + [self.u, self.v, self.eta] + [self.tracers[i] for i in self.tracers] ] ) + ## Angle Calculation & Rotation + rotated_u, rotated_v = rotate( + regridded[self.u], + regridded[self.v], + radian_angle=np.radians( + rot.get_rotation_angle( + rotational_method, self.hgrid, orientation=self.orientation + ).values + ), + ) + + 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])]) + if self.arakawa_grid == "B": ## All tracers on one grid, all velocities on another - regridder_velocity = xe.Regridder( + regridder_velocity = rgd.create_regridder( rawseg[self.u].rename({self.xq: "lon", self.yq: "lat"}), - self.coords, - "bilinear", - locstream_out=True, - reuse_weights=False, - filename=self.outfolder + coords, + self.outfolder / f"weights/bilinear_velocity_weights_{self.orientation}.nc", ) - - regridder_tracer = xe.Regridder( + regridder_tracer = rgd.create_regridder( rawseg[self.tracers["salt"]].rename({self.xh: "lon", self.yh: "lat"}), - self.coords, - "bilinear", - locstream_out=True, - reuse_weights=False, - filename=self.outfolder + coords, + self.outfolder / f"weights/bilinear_tracer_weights_{self.orientation}.nc", ) @@ -3118,8 +3069,15 @@ 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 + velocities_out["u"], velocities_out["v"] = rotate( + velocities_out["u"], + velocities_out["v"], + radian_angle=np.radians( + rot.get_rotation_angle( + rotational_method, self.hgrid, orientation=self.orientation + ).values + ), ) segment_out = xr.merge( @@ -3135,40 +3093,50 @@ def regrid_velocity_tracers(self): if self.arakawa_grid == "C": ## All tracers on one grid, all velocities on another - regridder_uvelocity = xe.Regridder( + regridder_uvelocity = rgd.create_regridder( rawseg[self.u].rename({self.xq: "lon", self.yh: "lat"}), - self.coords, - "bilinear", - locstream_out=True, - reuse_weights=False, - filename=self.outfolder + coords, + self.outfolder / f"weights/bilinear_uvelocity_weights_{self.orientation}.nc", ) - regridder_vvelocity = xe.Regridder( + regridder_vvelocity = rgd.create_regridder( rawseg[self.v].rename({self.xh: "lon", self.yq: "lat"}), - self.coords, - "bilinear", - locstream_out=True, - reuse_weights=False, - filename=self.outfolder + coords, + self.outfolder / f"weights/bilinear_vvelocity_weights_{self.orientation}.nc", ) - regridder_tracer = xe.Regridder( + regridder_tracer = rgd.create_regridder( rawseg[self.tracers["salt"]].rename({self.xh: "lon", self.yh: "lat"}), - self.coords, - "bilinear", - locstream_out=True, - reuse_weights=False, - filename=self.outfolder + coords, + self.outfolder / f"weights/bilinear_tracer_weights_{self.orientation}.nc", ) + regridded_u = regridder_uvelocity(rawseg[[self.u]]) + regridded_v = regridder_vvelocity(rawseg[[self.v]]) + + # See explanation of the rotational methods in the A grid section + rotated_u, rotated_v = rotate( + regridded_u, + regridded_v, + radian_angle=np.radians( + rot.get_rotation_angle( + rotational_method, self.hgrid, orientation=self.orientation + ).values + ), + ) + + rotated_ds = xr.Dataset( + { + self.u: rotated_u, + self.v: rotated_v, + } + ) segment_out = xr.merge( [ - regridder_vvelocity(rawseg[[self.v]]), - regridder_uvelocity(rawseg[[self.u]]), + rotated_ds, regridder_tracer( rawseg[[self.eta] + [self.tracers[i] for i in self.tracers]] ), @@ -3177,9 +3145,10 @@ def regrid_velocity_tracers(self): ## segment out now contains our interpolated boundary. ## Now, we need to fix up all the metadata and save + segment_out = segment_out.rename( + {"lon": f"lon_{self.segment_name}", "lat": f"lat_{self.segment_name}"} + ) - del segment_out["lon"] - 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})) @@ -3189,44 +3158,29 @@ def regrid_velocity_tracers(self): segment_out[self.tracers["temp"]].attrs["units"] = "degrees Celsius" # fill in NaNs - segment_out = ( - segment_out.ffill(self.z) - .interpolate_na(f"{self.coords.attrs['parallel']}_{self.segment_name}") - .ffill(f"{self.coords.attrs['parallel']}_{self.segment_name}") - .bfill(f"{self.coords.attrs['parallel']}_{self.segment_name}") + # segment_out = rgd.fill_missing_data(segment_out, self.z) + segment_out = rgd.fill_missing_data( + segment_out, + xdim=f"{coords.attrs['parallel']}_{self.segment_name}", + zdim=self.z, ) - time = 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, + 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"], ) - - segment_out = segment_out.assign_coords({"time": time}) + # This to change the time coordinate. + 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}", } - # Dictionary we built for encoding the netcdf at end - encoding_dict = { - "time": { - "dtype": "double", - }, - f"nx_{self.segment_name}": { - "dtype": "int32", - }, - f"ny_{self.segment_name}": { - "dtype": "int32", - }, - } - - ### Generate the dz variable; needs to be in layer thicknesses - dz = segment_out[self.z].diff(self.z) - dz.name = "dz" - dz = xr.concat([dz, dz[-1]], dim=self.z) - # 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 @@ -3245,135 +3199,70 @@ def regrid_velocity_tracers(self): ## Rename each variable in dataset segment_out = segment_out.rename({allfields[var]: v}) - ## Rename vertical coordinate for this variable - segment_out[f"{var}_{self.segment_name}"] = segment_out[ - f"{var}_{self.segment_name}" - ].rename({self.z: f"nz_{self.segment_name}_{var}"}) - - ## Replace the old depth coordinates with incremental integers - segment_out[f"nz_{self.segment_name}_{var}"] = np.arange( - segment_out[f"nz_{self.segment_name}_{var}"].size + segment_out = rgd.vertical_coordinate_encoding( + segment_out, v, self.segment_name, self.z ) - ## Re-add the secondary dimension (even though it represents one value..) - segment_out[v] = segment_out[v].expand_dims( - f"{self.coords.attrs['perpendicular']}_{self.segment_name}", - axis=self.coords.attrs["axis_to_expand"], + segment_out = rgd.add_secondary_dimension( + segment_out, v, coords, self.segment_name ) - ## Add the layer thicknesses - segment_out[f"dz_{v}"] = ( - [ - "time", - f"nz_{v}", - f"ny_{self.segment_name}", - f"nx_{self.segment_name}", - ], - da.broadcast_to( - dz.data[None, :, None, None], - segment_out[v].shape, - chunks=( - 1, - None, - None, - None, - ), ## Chunk in each time, and every 5 vertical layers - ), + segment_out = rgd.generate_layer_thickness( + segment_out, v, self.segment_name, self.z ) - encoding_dict[v] = { - "_FillValue": netCDF4.default_fillvals["f8"], - "zlib": True, - # "chunksizes": tuple(s), - } - encoding_dict[f"dz_{v}"] = { - "_FillValue": netCDF4.default_fillvals["f8"], - "zlib": True, - # "chunksizes": tuple(s), - } - - ## appears to be another variable just with integers?? - encoding_dict[f"nz_{self.segment_name}_{var}"] = {"dtype": "int32"} - ## 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}"}) - encoding_dict[f"eta_{self.segment_name}"] = { - "_FillValue": netCDF4.default_fillvals["f8"], - } - segment_out[f"eta_{self.segment_name}"] = segment_out[ - f"eta_{self.segment_name}" - ].expand_dims( - f"{self.coords.attrs['perpendicular']}_{self.segment_name}", - axis=self.coords.attrs["axis_to_expand"] - 1, + + segment_out = rgd.add_secondary_dimension( + segment_out, f"eta_{self.segment_name}", coords, self.segment_name ) # 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['perpendicular']}_{self.segment_name}"] = [0] - if self.orientation == "north": - self.hgrid_seg = self.hgrid.isel(nyp=[-1]) - self.perpendicular = "ny" - self.parallel = "nx" - - if self.orientation == "south": - self.hgrid_seg = self.hgrid.isel(nyp=[0]) - self.perpendicular = "ny" - self.parallel = "nx" - - if self.orientation == "east": - self.hgrid_seg = self.hgrid.isel(nxp=[-1]) - self.perpendicular = "nx" - self.parallel = "ny" - - if self.orientation == "west": - self.hgrid_seg = self.hgrid.isel(nxp=[0]) - self.perpendicular = "nx" - self.parallel = "ny" - - # Store actual lat/lon values here as variables rather than coordinates - segment_out[f"lon_{self.segment_name}"] = ( - [f"ny_{self.segment_name}", f"nx_{self.segment_name}"], - self.coords.lon.expand_dims( - dim="blank", axis=self.coords.attrs["axis_to_expand"] - 2 - ).data, - ) - segment_out[f"lat_{self.segment_name}"] = ( - [f"ny_{self.segment_name}", f"nx_{self.segment_name}"], - self.coords.lat.expand_dims( - dim="blank", axis=self.coords.attrs["axis_to_expand"] - 2 - ).data, - ) - - # Add units to the lat / lon to keep the `categorize_axis_from_units` checker from throwing warnings - segment_out[f"lat_{self.segment_name}"].attrs = { - "units": "degrees_north", - } - segment_out[f"lon_{self.segment_name}"].attrs = { - "units": "degrees_east", - } - segment_out[f"ny_{self.segment_name}"].attrs = { - "units": "degrees_north", - } - segment_out[f"nx_{self.segment_name}"].attrs = { - "units": "degrees_east", + 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"}, + f"nx_{self.segment_name}": { + "dtype": "int32", + }, + f"ny_{self.segment_name}": { + "dtype": "int32", + }, } - # If repeat-year forcing, add modulo coordinate - if self.repeat_year_forcing: - segment_out["time"] = segment_out["time"].assign_attrs({"modulo": " "}) + segment_out = rgd.mask_dataset( + segment_out, + self.hgrid, + self.bathymetry, + self.orientation, + self.segment_name, + ) + encoding_dict = rgd.generate_encoding( + segment_out, + encoding_dict, + default_fill_value=1.0e20, + ) - with ProgressBar(): - segment_out.load().to_netcdf( - self.outfolder / f"forcing_obc_{self.segment_name}.nc", - encoding=encoding_dict, - unlimited_dims="time", - ) + segment_out.load().to_netcdf( + self.outfolder / f"forcing_obc_{self.segment_name}.nc", + encoding=encoding_dict, + unlimited_dims="time", + ) 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: @@ -3387,14 +3276,15 @@ 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' + .nc files: Regridded tidal velocity and elevation files in 'inputdir/forcing' 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) + - 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) Original Code was sourced from: @@ -3406,70 +3296,66 @@ def regrid_tides( Web Address: https://github.com/jsimkins2/nwa25 """ + # Establish Coords + coords = rgd.coords(self.hgrid, self.orientation, self.segment_name) + ########## Tidal Elevation: Horizontally interpolate elevation components ############ - regrid = xe.Regridder( + regrid = rgd.create_regridder( tpxo_h[["lon", "lat", "hRe"]], - self.coords, - method="nearest_s2d", - locstream_out=True, - periodic=False, - filename=Path( + coords, + Path( self.outfolder / "forcing" / f"regrid_{self.segment_name}_tidal_elev.nc" ), - reuse_weights=False, ) + redest = regrid(tpxo_h[["lon", "lat", "hRe"]]) imdest = regrid(tpxo_h[["lon", "lat", "hIm"]]) # Fill missing data. # Need to do this first because complex would get converted to real - redest = redest.ffill(dim=self.coords.attrs["locations_name"], limit=None)[ - "hRe" - ] - imdest = imdest.ffill(dim=self.coords.attrs["locations_name"], limit=None)[ - "hIm" - ] + redest = rgd.fill_missing_data( + redest, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) + redest = redest["hRe"] + imdest = rgd.fill_missing_data( + imdest, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) + imdest = imdest["hIm"] # Convert complex cplex = redest + 1j * imdest # Convert to real amplitude and phase. ds_ap = xr.Dataset({f"zamp_{self.segment_name}": np.abs(cplex)}) + # np.angle doesn't return dataarray ds_ap[f"zphase_{self.segment_name}"] = ( - ("constituent", self.coords.attrs["locations_name"]), + ("constituent", f"{coords.attrs['parallel']}_{self.segment_name}"), -1 * np.angle(cplex), ) # radians # Add time coordinate and transpose so that time is first, # so that it can be the unlimited dimension - ds_ap, _ = xr.broadcast(ds_ap, times) + times = xr.DataArray( + pd.date_range( + self.startdate, periods=1 + ), # 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 = ds_ap.transpose( - "time", "constituent", self.coords.attrs["locations_name"] + "time", "constituent", f"{coords.attrs['parallel']}_{self.segment_name}" ) self.encode_tidal_files_and_output(ds_ap, "tz") ########### Regrid Tidal Velocity ###################### - regrid_u = xe.Regridder( - tpxo_u[["lon", "lat", "uRe"]], - self.coords, - method=method, - locstream_out=True, - periodic=periodic, - reuse_weights=False, - ) - regrid_v = xe.Regridder( - tpxo_v[["lon", "lat", "vRe"]], - self.coords, - method=method, - locstream_out=True, - periodic=periodic, - reuse_weights=False, - ) + 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 segment. + # Interpolate each real and imaginary parts to self. uredest = regrid_u(tpxo_u[["lon", "lat", "uRe"]])["uRe"] uimdest = regrid_u(tpxo_u[["lon", "lat", "uIm"]])["uIm"] vredest = regrid_v(tpxo_v[["lon", "lat", "vRe"]])["vRe"] @@ -3477,10 +3363,18 @@ def regrid_tides( # Fill missing data. # Need to do this first because complex would get converted to real - uredest = uredest.ffill(dim=self.coords.attrs["locations_name"], limit=None) - uimdest = uimdest.ffill(dim=self.coords.attrs["locations_name"], limit=None) - vredest = vredest.ffill(dim=self.coords.attrs["locations_name"], limit=None) - vimdest = vimdest.ffill(dim=self.coords.attrs["locations_name"], limit=None) + uredest = rgd.fill_missing_data( + uredest, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) + uimdest = rgd.fill_missing_data( + uimdest, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) + vredest = rgd.fill_missing_data( + vredest, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) + vimdest = rgd.fill_missing_data( + vimdest, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) # Convert to complex, remaining separate for u and v. ucplex = uredest + 1j * uimdest @@ -3491,34 +3385,44 @@ def regrid_tides( # and convert ellipse back to amplitude and phase. SEMA, ECC, INC, PHA = ap2ep(ucplex, vcplex) - # Rotate to the model grid by adjusting the inclination. - # Requries that angle is in radians. + # Rotate + INC -= np.radians( + rot.get_rotation_angle( + rotational_method, self.hgrid, orientation=self.orientation + ).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} ) - # up, vp aren't dataarrays + # up, vp aren't dataarraysf ds_ap[f"uphase_{self.segment_name}"] = ( - ("constituent", self.coords.attrs["locations_name"]), + ("constituent", f"{coords.attrs['parallel']}_{self.segment_name}"), up, ) # radians ds_ap[f"vphase_{self.segment_name}"] = ( - ("constituent", self.coords.attrs["locations_name"]), + ("constituent", f"{coords.attrs['parallel']}_{self.segment_name}"), vp, ) # radians - ds_ap, _ = xr.broadcast(ds_ap, times) - - # Need to transpose so that time is first, - # so that it can be the unlimited dimension + times = xr.DataArray( + pd.date_range( + self.startdate, periods=1 + ), # 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 = ds_ap.transpose( - "time", "constituent", self.coords.attrs["locations_name"] + "time", "constituent", f"{coords.attrs['parallel']}_{self.segment_name}" ) # Some things may have become missing during the transformation - ds_ap = ds_ap.ffill(dim=self.coords.attrs["locations_name"], limit=None) + ds_ap = rgd.fill_missing_data( + ds_ap, xdim=f"{coords.attrs['parallel']}_{self.segment_name}", zdim=None + ) self.encode_tidal_files_and_output(ds_ap, "tu") @@ -3527,23 +3431,23 @@ def regrid_tides( def encode_tidal_files_and_output(self, ds, filename): """ This function: - - Expands the dimensions (with the segment name) - - Renames some dimensions to be more specific to the segment - - Provides an output file encoding - - Exports the files. + - Expands the dimensions (with the segment name) + - Renames some dimensions to be more specific to the segment + - Provides an output file encoding + - Exports the files. Args: self.outfolder (str/path): The output folder to save the tidal files into dataset (xarray.Dataset): The processed tidal dataset filename (str): The output file name Returns: - *.nc files: Regridded [FILENAME] files in 'self.outfolder/[filename]_[segmentname].nc' + .nc files: Regridded [FILENAME] files in 'self.outfolder/[filename]_[segmentname].nc' 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) + - 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, coords, segment.regrid_tides, segment.encode_tidal_files_and_output) Original Code was sourced from: @@ -3556,48 +3460,34 @@ def encode_tidal_files_and_output(self, ds, filename): """ + coords = rgd.coords(self.hgrid, self.orientation, self.segment_name) ## Expand Tidal Dimensions ## - if "z" in ds.coords or "constituent" in ds.dims: - offset = 0 - else: - offset = 1 - if self.orientation in ["south", "north"]: - ds = ds.expand_dims(f"ny_{self.segment_name}", 2 - offset) - elif self.orientation in ["west", "east"]: - ds = ds.expand_dims(f"nx_{self.segment_name}", 3 - offset) + + for var in ds: + + ds = rgd.add_secondary_dimension(ds, str(var), coords, self.segment_name) ## Rename Tidal Dimensions ## ds = ds.rename( {"lon": f"lon_{self.segment_name}", "lat": f"lat_{self.segment_name}"} ) - if "z" in ds.coords: - ds = ds.rename({"z": f"nz_{self.segment_name}"}) - if self.orientation in ["south", "north"]: - ds = ds.rename( - {self.coords.attrs["locations_name"]: f"nx_{self.segment_name}"} - ) - elif self.orientation in ["west", "east"]: - ds = ds.rename( - {self.coords.attrs["locations_name"]: f"ny_{self.segment_name}"} - ) + ds = rgd.mask_dataset( + ds, self.hgrid, self.bathymetry, self.orientation, self.segment_name + ) ## Perform Encoding ## - for v in ds: - ds[v].encoding["_FillValue"] = 1.0e20 + fname = f"{filename}_{self.segment_name}.nc" # Set format and attributes for coordinates, including time if it does not already have calendar attribute # (may change this to detect whether time is a time type or a float). # Need to include the fillvalue or it will be back to nan encoding = { - "time": dict(_FillValue=1.0e20), + "time": dict(dtype="float64", calendar="gregorian", _FillValue=1.0e20), f"lon_{self.segment_name}": dict(dtype="float64", _FillValue=1.0e20), f"lat_{self.segment_name}": dict(dtype="float64", _FillValue=1.0e20), } - if "calendar" not in ds["time"].attrs and "modulo" not in ds["time"].attrs: - encoding.update( - {"time": dict(dtype="float64", calendar="gregorian", _FillValue=1.0e20)} - ) + encoding = rgd.generate_encoding(ds, encoding, default_fill_value=1.0e20) ## Export Files ## ds.to_netcdf( @@ -3606,4 +3496,4 @@ def encode_tidal_files_and_output(self, ds, filename): encoding=encoding, unlimited_dims="time", ) - return + return ds diff --git a/regional_mom6/regridding.py b/regional_mom6/regridding.py new file mode 100644 index 00000000..963aa20b --- /dev/null +++ b/regional_mom6/regridding.py @@ -0,0 +1,715 @@ +""" +Helper Functions to take the user though the regridding of boundary conditions and encoding for MOM6. Built for RM6 + +Steps: +1. Initial Regridding -> Find the boundary of the hgrid, and regrid the forcing variables to that boundary. Call (initial_regridding) and then use the xesmf Regridder with whatever datasets you need. +2. Work on some data issues + + 1. For temperature - Make sure it's in Celsius + 2. FILL IN NANS -> this is important for MOM6 (fill_missing_data) -> This diverges between + +3. For tides, we split the tides into an amplitude and a phase +4. In some cases, here is a great place to rotate the velocities to match a curved grid (tidal_velocity), velocity is also a good place to do this. +5. We then add the time coordinate +6. For vars that are not just surface variables, we need to add several depth related variables + 1. Add a dz variable in layer thickness + 2. Some metadata issues later on +7. Now we do up the metadata +8. Rename variables to var_segment_num +9. (IF VERTICAL EXISTS) Rename the vertical coordinate of the variable to nz_segment_num_var +10. (IF VERTICAL EXISTS) Declare this new vertical coordiante as a increasing series of integers +11. Re-add the "perpendicular" dimension +12. ....Add layer thickness of dz to the vertical forcings +13. Add to encoding_dict a fill value(_FillValue) and zlib, dtype, for time, lat long, ....and each variable (no type needed though) + + + +""" + +import xesmf as xe +import xarray as xr +from pathlib import Path +import dask.array as da +import numpy as np +import netCDF4 +from regional_mom6.utils import setup_logger + + +regridding_logger = setup_logger(__name__, set_handler=False) + + +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 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 + + 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 + + """ + + 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": 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}"}) + 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 + ) + elif orientation == "north": + rcoord = xr.Dataset( + { + "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}"}) + rcoord.attrs["perpendicular"] = "ny" + rcoord.attrs["parallel"] = "nx" + rcoord.attrs["axis_to_expand"] = 2 + elif orientation == "west": + rcoord = xr.Dataset( + { + "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}"}) + rcoord.attrs["perpendicular"] = "nx" + rcoord.attrs["parallel"] = "ny" + rcoord.attrs["axis_to_expand"] = 3 + elif orientation == "east": + rcoord = xr.Dataset( + { + "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}"}) + rcoord.attrs["perpendicular"] = "nx" + rcoord.attrs["parallel"] = "ny" + rcoord.attrs["axis_to_expand"] = 3 + + # Make lat and lon coordinates + rcoord = rcoord.assign_coords(lat=rcoord["lat"], lon=rcoord["lon"]) + + 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 = None, + method: str = "bilinear", + locstream_out: bool = True, + periodic: bool = False, +) -> xe.Regridder: + """ + Basic Regridder for any forcing variables, this just wraps the xesmf regridder for a few defaults + Parameters + ---------- + forcing_variables : xr.Dataset + The dataset of the forcing variables + output_grid : xr.Dataset + The dataset of the output grid -> this is the boundary of the hgrid + outfile : Path, optional + The path to the output file for weights I believe, by default Path(".temp") + method : str, optional + The regridding method, by default "bilinear" + locstream_out : bool, optional + Whether to output the locstream, by default True + periodic : bool, optional + Whether the grid is periodic, by default False + Returns + ------- + xe.Regridder + The regridding object + """ + regridding_logger.info("Creating Regridder") + regridder = xe.Regridder( + forcing_variables, + output_grid, + method=method, + locstream_out=locstream_out, + periodic=periodic, + filename=outfile, + reuse_weights=False, + ) + return regridder + + +def fill_missing_data( + ds: xr.Dataset, xdim: str = "locations", zdim: str = "z", fill: str = "b" +) -> xr.Dataset: + """ + Fill in missing values, taken from GFDL NWA 25 Repo + Parameters + ---------- + ds : xr.Dataset + The dataset to fill in + z_dim_name : str + The name of the z dimension + Returns + ------- + xr.Dataset + The filled in dataset + """ + regridding_logger.info("Filling in missing data horizontally, then vertically") + if fill == "f": + filled = ds.ffill(dim=xdim, limit=None) + elif fill == "b": + filled = ds.bfill(dim=xdim, limit=None) + if zdim is not None: + filled = filled.ffill(dim=zdim, limit=None).fillna(0) + return filled + + +def add_or_update_time_dim(ds: xr.Dataset, times) -> xr.Dataset: + """ + Add the time dimension to the dataset, in tides case can be one time step. + Parameters + ---------- + ds : xr.Dataset + The dataset to add the time dimension to + times : list, np.Array, xr.DataArray + The list of times + Returns + ------- + xr.Dataset + The dataset with the time dimension added + """ + regridding_logger.info("Adding time dimension") + + regridding_logger.debug(f"Times: {times}") + regridding_logger.debug(f"Make sure times is a DataArray") + # 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 + else: + regridding_logger.debug("Time not in dataset, xr.Broadcasting time dimension") + ds, _ = xr.broadcast(ds, times) + + # Make sure time is first.... + regridding_logger.debug("Transposing time to first dimension") + new_dims = ["time"] + [dim for dim in ds.dims if dim != "time"] + ds = ds.transpose(*new_dims) + + return ds + + +def generate_dz(ds: xr.Dataset, z_dim_name: str) -> xr.Dataset: + """ + For vertical coordinates, you need to have the layer thickness or something. Generate the dz variable for the dataset + Parameters + ---------- + ds : xr.Dataset + The dataset to get the z variable from + z_dim_name : str + The name of the z dimension + Returns + ------- + xr.Dataset + the dz variable + """ + dz = ds[z_dim_name].diff(z_dim_name) + dz.name = "dz" + dz = xr.concat([dz, dz[-1]], dim=z_dim_name) + return dz + + +def add_secondary_dimension( + 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 + ----------- + ds : xr.Dataset + The dataset to add the perpendicular dimension to + var : str + The variable to add the perpendicular dimension to + coords : xr.Dataset + 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 + 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" + ) + 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: + 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( + f"{coords.attrs['perpendicular']}_{segment_name}", + axis=coords.attrs["axis_to_expand"] - insert_behind_by, + ) + return ds + + +def vertical_coordinate_encoding( + ds: xr.Dataset, var: str, segment_name: str, old_vert_coord_name: str +) -> xr.Dataset: + """ + Rename vertical coordinate to nz[additional-text] then change it to regular increments + + Parameters + ---------- + ds : xr.Dataset + The dataset to rename the vertical coordinate in + var : str + The variable to rename the vertical coordinate in + segment_name : str + The segment name + old_vert_coord_name : str + The old vertical coordinate name + """ + + regridding_logger.info("Renaming vertical coordinate to nz_... in {}".format(var)) + section = "_seg" + base_var = var[: var.find(section)] if section in var else var + ds[var] = ds[var].rename({old_vert_coord_name: f"nz_{segment_name}_{base_var}"}) + + ## Replace the old depth coordinates with incremental integers + regridding_logger.info("Replacing old depth coordinates with incremental integers") + ds[f"nz_{segment_name}_{base_var}"] = np.arange( + ds[f"nz_{segment_name}_{base_var}"].size + ) + + return ds + + +def generate_layer_thickness( + ds: xr.Dataset, var: str, segment_name: str, old_vert_coord_name: str +) -> xr.Dataset: + """ + Generate Layer Thickness Variable, needed for vars with vertical dimensions + Parameters + ---------- + ds : xr.Dataset + The dataset to generate the layer thickness for + var : str + The variable to generate the layer thickness for + segment_name : str + The segment name + old_vert_coord_name : str + The old vertical coordinate name + Returns + ------- + xr.Dataset + The dataset with the layer thickness variable added + """ + regridding_logger.debug("Generating layer thickness variable for {}".format(var)) + dz = generate_dz(ds, old_vert_coord_name) + ds[f"dz_{var}"] = ( + [ + "time", + f"nz_{var}", + f"ny_{segment_name}", + f"nx_{segment_name}", + ], + da.broadcast_to( + dz.data[None, :, None, None], + ds[var].shape, + chunks=( + 1, + None, + None, + None, + ), ## Chunk in each time, and every 5 vertical layers + ), + ) + + 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", + add_land_exceptions=True, +) -> 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 + add_land_exceptions : bool + Add the corners and 3 coast point exceptions + Returns + ------- + np.ndarray + The boundary mask + """ + + # Hide the bathy as an hgrid so we can take advantage of the coords function to get the boundary points. + + # First rename bathy dims to nyp and nxp + 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_as_hgrid = hgrid.copy(deep=True) + + # Create new depth field + bathy_as_hgrid["depth"] = bathy_as_hgrid["angle_dx"] + bathy_as_hgrid["depth"][:, :] = np.nan + + # Fill at t_points (what bathy is determined at) + ds_t = get_hgrid_arakawa_c_points(hgrid, "t") + + # Identify any extra dimension like ntiles + extra_dim = None + for dim in bathy.dims: + if dim not in ["nyp", "nxp"]: + extra_dim = dim + break + # Select the first index along the extra dimension if it exists + if extra_dim: + bathy = bathy.isel({extra_dim: 0}) + + bathy_as_hgrid["depth"][ + ds_t.t_points_y.values, ds_t.t_points_x.values + ] = bathy.depth + + bathy_as_hgrid_coords = coords( + bathy_as_hgrid, + side, + segment_name, + angle_variable_name="depth", + coords_at_t_points=True, + ) + + # Get the Boundary Depth -> we're done with the hgrid now + bathy_as_hgrid_coords["boundary_depth"] = bathy_as_hgrid_coords["angle"] + + # Mask Fill Values + land = 0.5 + ocean = 1.0 + zero_out = 0.0 + + # Create Mask + boundary_mask = np.full(np.shape(coords(hgrid, side, segment_name).angle), ocean) + + # Fill with MOM6 version of mask + for i in range(len(bathy_as_hgrid_coords["boundary_depth"])): + if bathy_as_hgrid_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 + + if add_land_exceptions: + # Land points that can't be NaNs: Corners & 3 points at the coast + + # 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 + coasts_lower_index = [] + coasts_higher_index = [] + for index in range(1, len(boundary_mask) - 1): + if boundary_mask[index - 1] == land and boundary_mask[index] == ocean: + coasts_lower_index.append(index) + elif boundary_mask[index + 1] == land and boundary_mask[index] == ocean: + coasts_higher_index.append(index) + + # Remove 3 land points from the coast, and make them zeroed out real values + for i in range(3): + for coast in coasts_lower_index: + if coast - 1 - i >= 0: + boundary_mask[coast - 1 - i] = zero_out + for coast in coasts_higher_index: + if coast + 1 + i < len(boundary_mask): + boundary_mask[coast + 1 + i] = zero_out + + # Corner Q-points defined as land should be zeroed out + if boundary_mask[0] == land: + boundary_mask[0] = zero_out + if boundary_mask[-1] == land: + boundary_mask[-1] = zero_out + + # Convert land points to nans + boundary_mask[np.where(boundary_mask == land)] = np.nan + + 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", + add_land_exceptions=True, +) -> 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_land_exceptions : bool + To add the corner and 3 point coast exception + """ + ## 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, + add_land_exceptions=add_land_exceptions, + ) + 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) + elif np.isnan(dataset_reduce_dim[0]): # The corner is nan in the data + ds[var] = ds[var].copy() + ds[var][..., 0, 0] = 0 + elif np.isnan(dataset_reduce_dim[-1]): # The corner is nan in the data + ds[var] = ds[var].copy() + if orientation in ["east", "west"]: + ds[var][..., -1, 0] = 0 + else: + ds[var][..., 0, -1] = 0 + + ## Remove Nans if needed ## + ds[var] = ds[var].fillna(0) + + ## Apply the mask ## # Multiplication allows us to use 1, 0, and nan in 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"] +) -> dict: + """ + Generate the encoding dictionary for the dataset + Parameters + ---------- + ds : xr.Dataset + The dataset to generate the encoding for + encoding_dict : dict + The starting encoding dict with some specifications needed for time and other vars, this will be updated with encodings in this function + default_fill_value : float, optional + The default fill value, by default 1.0e20 + Returns + ------- + dict + The encoding dictionary + """ + regridding_logger.info("Generating encoding dictionary") + for var in ds: + if "_segment_" in var and not "nz" in var: + encoding_dict[var] = { + "_FillValue": default_fill_value, + } + for var in ds.coords: + if "nz_" in var: + encoding_dict[var] = { + "dtype": "int32", + } + + return encoding_dict diff --git a/regional_mom6/rotation.py b/regional_mom6/rotation.py new file mode 100644 index 00000000..7193779d --- /dev/null +++ b/regional_mom6/rotation.py @@ -0,0 +1,326 @@ +from regional_mom6 import utils +from regional_mom6.regridding import get_hgrid_arakawa_c_points, coords + +rotation_logger = utils.setup_logger(__name__, set_handler=False) +# 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 + + +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 + + +def get_rotation_angle( + rotational_method: RotationMethod, hgrid: xr.Dataset, orientation=None +): + """ + This function returns the rotation angle - THIS IS ALWAYS BASED ON THE ASSUMPTION OF DEGREES - based on the rotational method and provided hgrid, if orientation & coords are provided, it will assume the boundary is requested + Parameters + ---------- + rotational_method: RotationMethod + The rotational method to use + hgrid: xr.Dataset + The hgrid dataset + orientation: xr.Dataset + The orientation, which also lets us now that we are on a boundary + Returns + ------- + xr.DataArray + angle in degrees + """ + rotation_logger.info("Getting rotation angle") + boundary = False + if orientation != None: + rotation_logger.debug( + "The rotational angle is requested for the boundary: {}".format(orientation) + ) + boundary = True + + if rotational_method == RotationMethod.NO_ROTATION: + rotation_logger.debug("Using NO_ROTATION method") + if not utils.is_rectilinear_hgrid(hgrid): + raise ValueError("NO_ROTATION method only works with rectilinear grids") + angles = xr.zeros_like(hgrid.x) + + if boundary: + # Subset to just boundary + # Add zeroes to hgrid + hgrid["zero_angle"] = angles + + # Cut to boundary + zero_angle = coords( + hgrid, + orientation, + "doesnt_matter", + angle_variable_name="zero_angle", + )["angle"] + + return zero_angle + else: + return angles + elif rotational_method == RotationMethod.GIVEN_ANGLE: + rotation_logger.debug("Using GIVEN_ANGLE method") + if boundary: + return coords( + hgrid, orientation, "doesnt_matter", angle_variable_name="angle_dx" + )["angle"] + else: + return hgrid["angle_dx"] + elif rotational_method == RotationMethod.EXPAND_GRID: + rotation_logger.debug("Using EXPAND_GRID method") + hgrid["angle_dx_rm6"] = initialize_grid_rotation_angles_using_expanded_hgrid( + hgrid + ) + + if boundary: + degree_angle = coords( + hgrid, + orientation, + "doesnt_matter", + angle_variable_name="angle_dx_rm6", + )["angle"] + return degree_angle + else: + return hgrid["angle_dx_rm6"] + else: + raise ValueError("Invalid rotational method") diff --git a/regional_mom6/utils.py b/regional_mom6/utils.py index 447a2e4f..5f393c5e 100644 --- a/regional_mom6/utils.py +++ b/regional_mom6/utils.py @@ -1,4 +1,8 @@ import numpy as np +import logging +import sys +import xarray as xr +from regional_mom6 import regridding as rgd def vecdot(v1, v2): @@ -294,3 +298,87 @@ def ep2ap(SEMA, ECC, INC, PHA): vp = -np.angle(cv) return ua, va, up, vp + + +def setup_logger( + name: str, set_handler=False, log_level=logging.INFO +) -> logging.Logger: + """ + Setup general config for a logger. + """ + logger = logging.getLogger(name) + logger.setLevel(log_level) + if set_handler and not logger.hasHandlers(): + # Create a handler to print to stdout (Jupyter captures stdout) + handler = logging.StreamHandler(sys.stdout) + handler.setLevel(log_level) + + # Create a formatter (optional) + formatter = logging.Formatter("%(name)s.%(funcName)s:%(levelname)s:%(message)s") + handler.setFormatter(formatter) + + # Add the handler to the logger + logger.addHandler(handler) + return logger + + +def rotate_complex(u, v, radian_angle): + """ + Rotate velocities to grid orientation using complex number math (Same as rotate) + 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. + """ + + # express velocity in the complex plan + vel = u + v * 1j + # rotate velocity using grid angle theta + vel = vel * np.exp(1j * radian_angle) + + # From here you can easily get the rotated u, v, or the magnitude/direction of the currents: + u = np.real(vel) + v = np.imag(vel) + + return u, v + + +def rotate(u, v, radian_angle): + """ + 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. + """ + + 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 is_rectilinear_hgrid(hgrid: xr.Dataset, rtol: float = 1e-3) -> bool: + """ + Check if the hgrid is a rectilinear grid. From mom6_bathy.grid.is_rectangular by Alper (Altuntas + ) + Check if the grid is a rectangular lat-lon grid by comparing the first and last rows and columns of the tlon and tlat arrays. + + Args: + hgrid (xarray.Dataset): The horizontal grid dataset. + rtol (float): Relative tolerance. Default is 1e-3. + """ + ds_t = rgd.get_hgrid_arakawa_c_points(hgrid) + if ( + np.allclose(ds_t.tlon[:, 0], ds_t.tlon[0, 0], rtol=rtol) + and np.allclose(ds_t.tlon[:, -1], ds_t.tlon[0, -1], rtol=rtol) + and np.allclose(ds_t.tlat[0, :], ds_t.tlat[0, 0], rtol=rtol) + and np.allclose(ds_t.tlat[-1, :], ds_t.tlat[-1, 0], rtol=rtol) + ): + return True + return False diff --git a/tests/conftest.py b/tests/conftest.py new file mode 100644 index 00000000..ab580660 --- /dev/null +++ b/tests/conftest.py @@ -0,0 +1,283 @@ +import pytest +import os +import xarray as xr +import numpy as np +import regional_mom6 as rm6 + +# 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, skip test + else: + pytest.skip( + f"Required file 'hgrid.nc' not found in {DOCKER_FILE_PATH} or {LOCAL_FILE_PATH}" + ) + + +@pytest.fixture +def get_rectilinear_hgrid(): + lat = np.linspace(0, 10, 7) + lon = np.linspace(0, 10, 13) + rect_hgrid = rm6.generate_rectangular_hgrid(lat, lon) + return rect_hgrid + + +@pytest.fixture() +def generate_silly_vt_dataset(): + latitude_extent = [30, 40] + longitude_extent = [-80, -70] + eastern_boundary = xr.Dataset( + { + "temp": xr.DataArray( + np.random.random((100, 5, 10, 10)), + dims=["silly_lat", "silly_lon", "silly_depth", "time"], + coords={ + "silly_lat": np.linspace( + latitude_extent[0] - 5, latitude_extent[1] + 5, 100 + ), + "silly_lon": np.linspace( + longitude_extent[1] - 0.5, longitude_extent[1] + 0.5, 5 + ), + "silly_depth": np.linspace(0, 1000, 10), + "time": np.linspace(0, 1000, 10), + }, + ), + "eta": xr.DataArray( + np.random.random((100, 5, 10)), + dims=["silly_lat", "silly_lon", "time"], + coords={ + "silly_lat": np.linspace( + latitude_extent[0] - 5, latitude_extent[1] + 5, 100 + ), + "silly_lon": np.linspace( + longitude_extent[1] - 0.5, longitude_extent[1] + 0.5, 5 + ), + "time": np.linspace(0, 1000, 10), + }, + ), + "salt": xr.DataArray( + np.random.random((100, 5, 10, 10)), + dims=["silly_lat", "silly_lon", "silly_depth", "time"], + coords={ + "silly_lat": np.linspace( + latitude_extent[0] - 5, latitude_extent[1] + 5, 100 + ), + "silly_lon": np.linspace( + longitude_extent[1] - 0.5, longitude_extent[1] + 0.5, 5 + ), + "silly_depth": np.linspace(0, 1000, 10), + "time": np.linspace(0, 1000, 10), + }, + ), + "u": xr.DataArray( + np.random.random((100, 5, 10, 10)), + dims=["silly_lat", "silly_lon", "silly_depth", "time"], + coords={ + "silly_lat": np.linspace( + latitude_extent[0] - 5, latitude_extent[1] + 5, 100 + ), + "silly_lon": np.linspace( + longitude_extent[1] - 0.5, longitude_extent[1] + 0.5, 5 + ), + "silly_depth": np.linspace(0, 1000, 10), + "time": np.linspace(0, 1000, 10), + }, + ), + "v": xr.DataArray( + np.random.random((100, 5, 10, 10)), + dims=["silly_lat", "silly_lon", "silly_depth", "time"], + coords={ + "silly_lat": np.linspace( + latitude_extent[0] - 5, latitude_extent[1] + 5, 100 + ), + "silly_lon": np.linspace( + longitude_extent[1] - 0.5, longitude_extent[1] + 0.5, 5 + ), + "silly_depth": np.linspace(0, 1000, 10), + "time": np.linspace(0, 1000, 10), + }, + ), + } + ) + return eastern_boundary + + +@pytest.fixture() +def generate_silly_ic_dataset(): + def _generate_silly_ic_dataset( + longitude_extent, + latitude_extent, + resolution, + number_vertical_layers, + depth, + temp_dataarray_initial_condition, + ): + nx, ny = number_of_gridpoints(longitude_extent, latitude_extent, resolution) + silly_lat, silly_lon, silly_depth = generate_silly_coords( + longitude_extent, latitude_extent, resolution, depth, number_vertical_layers + ) + + dims = ["silly_lat", "silly_lon", "silly_depth"] + + coords = { + "silly_lat": silly_lat, + "silly_lon": silly_lon, + "silly_depth": silly_depth, + } + # initial condition includes, temp, salt, eta, u, v + initial_cond = xr.Dataset( + { + "eta": xr.DataArray( + np.random.random((ny, nx)), + dims=["silly_lat", "silly_lon"], + coords={ + "silly_lat": silly_lat, + "silly_lon": silly_lon, + }, + ), + "temp": temp_dataarray_initial_condition, + "salt": xr.DataArray( + np.random.random((ny, nx, number_vertical_layers)), + dims=dims, + coords=coords, + ), + "u": xr.DataArray( + np.random.random((ny, nx, number_vertical_layers)), + dims=dims, + coords=coords, + ), + "v": xr.DataArray( + np.random.random((ny, nx, number_vertical_layers)), + dims=dims, + coords=coords, + ), + } + ) + return initial_cond + + return _generate_silly_ic_dataset + + +@pytest.fixture() +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 + + +def get_temperature_dataarrays( + longitude_extent, latitude_extent, resolution, number_vertical_layers, depth +): + + silly_lat, silly_lon, silly_depth = generate_silly_coords( + longitude_extent, latitude_extent, resolution, depth, number_vertical_layers + ) + + dims = ["silly_lat", "silly_lon", "silly_depth"] + + coords = { + "silly_lat": silly_lat, + "silly_lon": silly_lon, + "silly_depth": silly_depth, + } + + toolpath_dir = "toolpath" + hgrid_type = "even_spacing" + + nx, ny = number_of_gridpoints(longitude_extent, latitude_extent, resolution) + + temp_in_C, temp_in_C_masked, temp_in_K, temp_in_K_masked = ( + generate_temperature_arrays(nx, ny, number_vertical_layers) + ) + + temp_C = xr.DataArray(temp_in_C, dims=dims, coords=coords) + temp_K = xr.DataArray(temp_in_K, dims=dims, coords=coords) + temp_C_masked = xr.DataArray(temp_in_C_masked, dims=dims, coords=coords) + temp_K_masked = xr.DataArray(temp_in_K_masked, dims=dims, coords=coords) + + maximum_temperature_in_C = np.max(temp_in_C) + return [temp_C, temp_C_masked, temp_K, temp_K_masked] + + +def number_of_gridpoints(longitude_extent, latitude_extent, resolution): + nx = int((longitude_extent[-1] - longitude_extent[0]) / resolution) + ny = int((latitude_extent[-1] - latitude_extent[0]) / resolution) + + return nx, ny + + +def generate_silly_coords( + longitude_extent, latitude_extent, resolution, depth, number_vertical_layers +): + nx, ny = number_of_gridpoints(longitude_extent, latitude_extent, resolution) + + horizontal_buffer = 5 + + silly_lat = np.linspace( + latitude_extent[0] - horizontal_buffer, + latitude_extent[1] + horizontal_buffer, + ny, + ) + silly_lon = np.linspace( + longitude_extent[0] - horizontal_buffer, + longitude_extent[1] + horizontal_buffer, + nx, + ) + silly_depth = np.linspace(0, depth, number_vertical_layers) + + return silly_lat, silly_lon, silly_depth + + +def generate_temperature_arrays(nx, ny, number_vertical_layers): + + # temperatures close to 0 ᵒC + temp_in_C = np.random.randn(ny, nx, number_vertical_layers) + + temp_in_C_masked = np.copy(temp_in_C) + if int(ny / 4 + 4) < ny - 1 and int(nx / 3 + 4) < nx + 1: + temp_in_C_masked[ + int(ny / 3) : int(ny / 3 + 5), int(nx) : int(nx / 4 + 4), : + ] = float("nan") + else: + raise ValueError("use bigger domain") + + temp_in_K = np.copy(temp_in_C) + 273.15 + temp_in_K_masked = np.copy(temp_in_C_masked) + 273.15 + + # ensure we didn't mask the minimum temperature + if np.nanmin(temp_in_C_masked) == np.min(temp_in_C): + return temp_in_C, temp_in_C_masked, temp_in_K, temp_in_K_masked + else: + return generate_temperature_arrays(nx, ny, number_vertical_layers) diff --git a/tests/test_config.py b/tests/test_config.py index 005b685c..31d55886 100644 --- a/tests/test_config.py +++ b/tests/test_config.py @@ -6,7 +6,7 @@ import shutil -def test_write_config(): +def test_write_config(tmp_path): expt_name = "testing" latitude_extent = [16.0, 27] @@ -17,6 +17,7 @@ def test_write_config(): ## Place where all your input files go input_dir = Path( os.path.join( + tmp_path, expt_name, "inputs", ) @@ -25,11 +26,12 @@ def test_write_config(): ## Directory where you'll run the experiment from run_dir = Path( os.path.join( + tmp_path, expt_name, "run_files", ) ) - data_path = Path("data") + data_path = Path(tmp_path / "data") for path in (run_dir, input_dir, data_path): os.makedirs(str(path), exist_ok=True) @@ -49,7 +51,7 @@ def test_write_config(): expt_name="test", boundaries=["south", "north"], ) - config_dict = expt.write_config_file() + config_dict = expt.write_config_file(tmp_path / "testing_config.json") assert config_dict["longitude_extent"] == tuple(longitude_extent) assert config_dict["latitude_extent"] == tuple(latitude_extent) assert config_dict["date_range"] == date_range @@ -61,7 +63,18 @@ def test_write_config(): assert config_dict["expt_name"] == "test" assert config_dict["hgrid_type"] == "even_spacing" assert config_dict["repeat_year_forcing"] == False - assert config_dict["tidal_constituents"] == ["M2"] + assert config_dict["tidal_constituents"] == [ + "M2", + "S2", + "N2", + "K2", + "K1", + "O1", + "P1", + "Q1", + "MM", + "MF", + ] assert config_dict["expt_name"] == "test" assert config_dict["boundaries"] == ["south", "north"] shutil.rmtree(run_dir) @@ -69,7 +82,7 @@ def test_write_config(): shutil.rmtree(data_path) -def test_load_config(): +def test_load_config(tmp_path): expt_name = "testing" @@ -81,6 +94,7 @@ def test_load_config(): ## Place where all your input files go input_dir = Path( os.path.join( + tmp_path, expt_name, "inputs", ) @@ -88,12 +102,14 @@ def test_load_config(): ## 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("data") + data_path = Path(tmp_path / "data") for path in (run_dir, input_dir, data_path): os.makedirs(str(path), exist_ok=True) @@ -111,9 +127,11 @@ def test_load_config(): mom_input_dir=input_dir, toolpath_dir="", ) - 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)) + new_expt = rmom6.create_experiment_from_config( + os.path.join(path), mom_input_folder=tmp_path, mom_run_folder=tmp_path + ) assert str(new_expt) == str(expt) print(new_expt.vgrid) print(expt.vgrid) @@ -125,8 +143,3 @@ def test_load_config(): assert os.path.exists(new_expt.mom_input_dir / "hgrid.nc") & os.path.exists( new_expt.mom_input_dir / "vcoord.nc" ) - shutil.rmtree(run_dir) - shutil.rmtree(input_dir) - shutil.rmtree(data_path) - shutil.rmtree(new_expt.mom_run_dir) - shutil.rmtree(new_expt.mom_input_dir) diff --git a/tests/test_expt_class.py b/tests/test_expt_class.py index aff3a4b8..eb288ef5 100644 --- a/tests/test_expt_class.py +++ b/tests/test_expt_class.py @@ -2,6 +2,14 @@ import pytest from regional_mom6 import experiment import xarray as xr +import xesmf as xe +import dask +from .conftest import ( + generate_temperature_arrays, + generate_silly_coords, + number_of_gridpoints, + get_temperature_dataarrays, +) ## Note: ## When creating test dataarrays we use 'silly' names for coordinates to @@ -17,8 +25,6 @@ "number_vertical_layers", "layer_thickness_ratio", "depth", - "mom_run_dir", - "mom_input_dir", "toolpath_dir", "hgrid_type", ), @@ -31,8 +37,6 @@ 5, 1, 1000, - "rundir/", - "inputdir/", "toolpath", "even_spacing", ), @@ -46,12 +50,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, @@ -60,8 +64,8 @@ def test_setup_bathymetry( number_vertical_layers=number_vertical_layers, layer_thickness_ratio=layer_thickness_ratio, depth=depth, - mom_run_dir=mom_run_dir, - mom_input_dir=mom_input_dir, + mom_run_dir=tmp_path / mom_run_dir, + mom_input_dir=tmp_path / mom_input_dir, toolpath_dir=toolpath_dir, hgrid_type=hgrid_type, ) @@ -98,58 +102,6 @@ def test_setup_bathymetry( bathymetry_file.unlink() -def number_of_gridpoints(longitude_extent, latitude_extent, resolution): - nx = int((longitude_extent[-1] - longitude_extent[0]) / resolution) - ny = int((latitude_extent[-1] - latitude_extent[0]) / resolution) - - return nx, ny - - -def generate_temperature_arrays(nx, ny, number_vertical_layers): - - # temperatures close to 0 ᵒC - temp_in_C = np.random.randn(ny, nx, number_vertical_layers) - - temp_in_C_masked = np.copy(temp_in_C) - if int(ny / 4 + 4) < ny - 1 and int(nx / 3 + 4) < nx + 1: - temp_in_C_masked[ - int(ny / 3) : int(ny / 3 + 5), int(nx) : int(nx / 4 + 4), : - ] = float("nan") - else: - raise ValueError("use bigger domain") - - temp_in_K = np.copy(temp_in_C) + 273.15 - temp_in_K_masked = np.copy(temp_in_C_masked) + 273.15 - - # ensure we didn't mask the minimum temperature - if np.nanmin(temp_in_C_masked) == np.min(temp_in_C): - return temp_in_C, temp_in_C_masked, temp_in_K, temp_in_K_masked - else: - return generate_temperature_arrays(nx, ny, number_vertical_layers) - - -def generate_silly_coords( - longitude_extent, latitude_extent, resolution, depth, number_vertical_layers -): - nx, ny = number_of_gridpoints(longitude_extent, latitude_extent, resolution) - - horizontal_buffer = 5 - - silly_lat = np.linspace( - latitude_extent[0] - horizontal_buffer, - latitude_extent[1] + horizontal_buffer, - ny, - ) - silly_lon = np.linspace( - longitude_extent[0] - horizontal_buffer, - longitude_extent[1] + horizontal_buffer, - nx, - ) - silly_depth = np.linspace(0, depth, number_vertical_layers) - - return silly_lat, silly_lon, silly_depth - - longitude_extent = [-5, 3] latitude_extent = (0, 10) date_range = ["2003-01-01 00:00:00", "2003-01-01 00:00:00"] @@ -158,36 +110,12 @@ def generate_silly_coords( layer_thickness_ratio = 1 depth = 1000 -silly_lat, silly_lon, silly_depth = generate_silly_coords( - longitude_extent, latitude_extent, resolution, depth, number_vertical_layers -) - -dims = ["silly_lat", "silly_lon", "silly_depth"] - -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" - -nx, ny = number_of_gridpoints(longitude_extent, latitude_extent, resolution) - -temp_in_C, temp_in_C_masked, temp_in_K, temp_in_K_masked = generate_temperature_arrays( - nx, ny, number_vertical_layers -) - -temp_C = xr.DataArray(temp_in_C, dims=dims, coords=coords) -temp_K = xr.DataArray(temp_in_K, dims=dims, coords=coords) -temp_C_masked = xr.DataArray(temp_in_C_masked, dims=dims, coords=coords) -temp_K_masked = xr.DataArray(temp_in_K_masked, dims=dims, coords=coords) - -maximum_temperature_in_C = np.max(temp_in_C) - @pytest.mark.parametrize( "temp_dataarray_initial_condition", - [temp_C, temp_C_masked, temp_K, temp_K_masked], + get_temperature_dataarrays( + longitude_extent, latitude_extent, resolution, number_vertical_layers, depth + ), ) @pytest.mark.parametrize( ( @@ -198,8 +126,6 @@ def generate_silly_coords( "number_vertical_layers", "layer_thickness_ratio", "depth", - "mom_run_dir", - "mom_input_dir", "toolpath_dir", "hgrid_type", ), @@ -212,8 +138,6 @@ def generate_silly_coords( number_vertical_layers, layer_thickness_ratio, depth, - "rundir/", - "inputdir/", "toolpath", "even_spacing", ), @@ -227,26 +151,15 @@ 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, tmp_path, + generate_silly_ic_dataset, ): - - silly_lat, silly_lon, silly_depth = generate_silly_coords( - longitude_extent, latitude_extent, resolution, depth, number_vertical_layers - ) - - dims = ["silly_lat", "silly_lon", "silly_depth"] - - coords = { - "silly_lat": silly_lat, - "silly_lon": silly_lon, - "silly_depth": silly_depth, - } - + dask.config.set(scheduler="single-threaded") + mom_run_dir = tmp_path / "rundir" + mom_input_dir = tmp_path / "inputdir" expt = experiment( longitude_extent=longitude_extent, latitude_extent=latitude_extent, @@ -255,48 +168,22 @@ def test_ocean_forcing( number_vertical_layers=number_vertical_layers, layer_thickness_ratio=layer_thickness_ratio, depth=depth, - mom_run_dir=mom_run_dir, - mom_input_dir=mom_input_dir, + mom_run_dir=tmp_path / mom_run_dir, + mom_input_dir=tmp_path / mom_input_dir, toolpath_dir=toolpath_dir, hgrid_type=hgrid_type, ) - ## Generate some initial condition to test on - - nx, ny = number_of_gridpoints(longitude_extent, latitude_extent, resolution) - # initial condition includes, temp, salt, eta, u, v - initial_cond = xr.Dataset( - { - "eta": xr.DataArray( - np.random.random((ny, nx)), - dims=["silly_lat", "silly_lon"], - coords={ - "silly_lat": silly_lat, - "silly_lon": silly_lon, - }, - ), - "temp": temp_dataarray_initial_condition, - "salt": xr.DataArray( - np.random.random((ny, nx, number_vertical_layers)), - dims=dims, - coords=coords, - ), - "u": xr.DataArray( - np.random.random((ny, nx, number_vertical_layers)), - dims=dims, - coords=coords, - ), - "v": xr.DataArray( - np.random.random((ny, nx, number_vertical_layers)), - dims=dims, - coords=coords, - ), - } + initial_cond = generate_silly_ic_dataset( + longitude_extent, + latitude_extent, + resolution, + number_vertical_layers, + depth, + temp_dataarray_initial_condition, ) - # Generate boundary forcing - initial_cond.to_netcdf(tmp_path / "ic_unprocessed") initial_cond.close() varnames = { @@ -318,9 +205,10 @@ def test_ocean_forcing( # ensure that temperature is in degrees C assert np.nanmin(expt.ic_tracers["temp"]) < 100.0 - + maximum_temperature_in_C = np.max(temp_dataarray_initial_condition) # max(temp) can be less maximum_temperature_in_C due to re-gridding assert np.nanmax(expt.ic_tracers["temp"]) <= maximum_temperature_in_C + dask.config.set(scheduler=None) @pytest.mark.parametrize( @@ -332,8 +220,6 @@ def test_ocean_forcing( "number_vertical_layers", "layer_thickness_ratio", "depth", - "mom_run_dir", - "mom_input_dir", "toolpath_dir", "hgrid_type", ), @@ -346,8 +232,6 @@ def test_ocean_forcing( 5, 1, 1000, - "rundir/", - "inputdir/", "toolpath", "even_spacing", ), @@ -361,8 +245,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 +325,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, @@ -452,8 +335,8 @@ def test_rectangular_boundaries( number_vertical_layers=number_vertical_layers, layer_thickness_ratio=layer_thickness_ratio, depth=depth, - mom_run_dir=mom_run_dir, - mom_input_dir=mom_input_dir, + mom_run_dir=tmp_path / mom_run_dir, + mom_input_dir=tmp_path / mom_input_dir, toolpath_dir=toolpath_dir, hgrid_type=hgrid_type, boundaries=["east"], diff --git a/tests/test_manish_branch.py b/tests/test_manish_branch.py deleted file mode 100644 index b2865512..00000000 --- a/tests/test_manish_branch.py +++ /dev/null @@ -1,289 +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: - @classmethod - def setup_class(self): # tmp_path is a pytest fixture - expt_name = "testing" - ## User-1st, test if we can even read the angled nc files. - self.dump_files_dir = Path("testing_outputs") - os.makedirs(self.dump_files_dir, exist_ok=True) - self.expt = rmom6.experiment.create_empty( - expt_name=expt_name, - mom_input_dir=self.dump_files_dir, - mom_run_dir=self.dump_files_dir, - ) - - @classmethod - def teardown_class(cls): - shutil.rmtree(cls.dump_files_dir) - - @pytest.fixture(scope="module") - def full_legit_expt_setup(self, 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( - expt_name, - "inputs", - ) - ) - - ## Directory where you'll run the experiment from - run_dir = Path( - os.path.join( - expt_name, - "run_files", - ) - ) - data_path = 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, full_legit_expt_setup): - assert str(full_legit_expt_setup) - - # @pytest.mark.skipif( - # IN_GITHUB_ACTIONS, reason="Test doesn't work in Github Actions." - # ) - def test_tides(self, dummy_tidal_data): - """ - Test the main setup tides function! - """ - - # Generate Fake Tidal Data - ds_h, ds_u = dummy_tidal_data - - # Save to Fake Folder - ds_h.to_netcdf(self.dump_files_dir / "h_fake_tidal_data.nc") - ds_u.to_netcdf(self.dump_files_dir / "u_fake_tidal_data.nc") - - # Set other required variables needed in setup_tides - - # Lat Long - self.expt.longitude_extent = (-5, 5) - self.expt.latitude_extent = (0, 30) - # Grid Type - self.expt.hgrid_type = "even_spacing" - # Dates - self.expt.date_range = ("2000-01-01", "2000-01-02") - self.expt.segments = {} - # Generate Hgrid Data - self.expt.resolution = 0.1 - self.expt.hgrid = self.expt._make_hgrid() - # Create Forcing Folder - os.makedirs(self.dump_files_dir / "forcing", exist_ok=True) - - self.expt.setup_boundary_tides(self.dump_files_dir, "fake_tidal_data.nc") - - def test_change_MOM_parameter(self): - """ - Test the change MOM parameter function, as well as read_MOM_file and write_MOM_file under the hood. - """ - - # 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", self.expt.mom_run_dir, dirs_exist_ok=True - ) - MOM_override_dict = self.expt.read_MOM_file_as_dict("MOM_override") - og = self.expt.change_MOM_parameter("DT", "30", "COOL COMMENT") - MOM_override_dict_new = self.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): - """ - Test the properties - """ - dss = self.expt.era5 - dss_2 = self.expt.tides_boundaries - dss_3 = self.expt.ocean_state_boundaries - dss_4 = self.expt.initial_condition - dss_5 = self.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..6ba00dec --- /dev/null +++ b/tests/test_regridding.py @@ -0,0 +1,279 @@ +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, generate_silly_vt_dataset): + hgrid = get_curvilinear_hgrid + ds = generate_silly_vt_dataset + ds["lat"] = ds.silly_lat + ds["lon"] = ds.silly_lat + assert rgd.get_hgrid_arakawa_c_points(hgrid, "t") + assert rgd.coords(hgrid, "north", "segment_002") + assert rgd.create_regridder(ds, ds) + + +def test_fill_missing_data(generate_silly_vt_dataset): + """ + Only testing forward fill for now + """ + ds = generate_silly_vt_dataset + ds["temp"][0, 0, 6:10, 0] = np.nan + + ds = rgd.fill_missing_data(ds, "silly_depth", fill="f") + + assert ( + ds["temp"][0, 0, 6:10, 0] == (ds["temp"][0, 0, 5, 0]) + ).all() # Assert if we are forward filling in time + + ds_2 = generate_silly_vt_dataset + ds_2["temp"][0, 0, 6:10, 0] = ds["temp"][0, 0, 5, 0] + assert (ds["temp"] == (ds_2["temp"])).all() # Assert everything else is the same + + +def test_add_or_update_time_dim(generate_silly_vt_dataset): + ds = generate_silly_vt_dataset + ds = rgd.add_or_update_time_dim(ds, xr.DataArray([0])) + + assert ds["time"].values == [0] # Assert time is added + assert ds["temp"].dims[0] == "time" # Check time is first dim + + +def test_generate_dz(generate_silly_vt_dataset): + ds = generate_silly_vt_dataset + dz = rgd.generate_dz(ds, "silly_depth") + z = np.linspace(0, 1000, 10) + dz_check = np.full(z.shape, z[1] - z[0]) + assert ( + (dz.values - dz_check) < 0.00001 + ).all() # Assert dz is generated correctly (some rounding leniency) + + +def test_add_secondary_dimension(get_curvilinear_hgrid, generate_silly_vt_dataset): + ds = generate_silly_vt_dataset + hgrid = get_curvilinear_hgrid + + # N/S Boundary + coords = rgd.coords(hgrid, "north", "segment_002") + ds = rgd.add_secondary_dimension(ds, "temp", coords, "segment_002") + assert ds["temp"].dims == ( + "silly_lat", + "ny_segment_002", + "silly_lon", + "silly_depth", + "time", + ) + + # E/W Boundary + coords = rgd.coords(hgrid, "east", "segment_003") + ds = generate_silly_vt_dataset + ds = rgd.add_secondary_dimension(ds, "v", coords, "segment_003") + assert ds["v"].dims == ( + "silly_lat", + "silly_lon", + "nx_segment_003", + "silly_depth", + "time", + ) + + # Beginning + ds = generate_silly_vt_dataset + ds = rgd.add_secondary_dimension( + ds, "temp", coords, "segment_003", to_beginning=True + ) + assert ds["temp"].dims[0] == "nx_segment_003" + + # NZ dim E/W Boundary + ds = generate_silly_vt_dataset + ds = ds.rename({"silly_depth": "nz"}) + ds = rgd.add_secondary_dimension(ds, "u", coords, "segment_003") + assert ds["u"].dims == ( + "silly_lat", + "silly_lon", + "nz", + "nx_segment_003", + "time", + ) + + +def test_vertical_coordinate_encoding(generate_silly_vt_dataset): + ds = generate_silly_vt_dataset + ds = rgd.vertical_coordinate_encoding(ds, "temp", "segment_002", "silly_depth") + assert "nz_segment_002_temp" in ds["temp"].dims + assert "nz_segment_002_temp" in ds + assert ( + ds["nz_segment_002_temp"] == np.arange(ds[f"nz_segment_002_temp"].size) + ).all() + + +def test_generate_layer_thickness(generate_silly_vt_dataset): + ds = generate_silly_vt_dataset + ds["temp"] = ds["temp"].transpose("time", "silly_depth", "silly_lat", "silly_lon") + ds = rgd.generate_layer_thickness(ds, "temp", "segment_002", "silly_depth") + assert "dz_temp" in ds + assert ds["dz_temp"].dims == ("time", "nz_temp", "ny_segment_002", "nx_segment_002") + assert ( + ds["temp"]["silly_depth"].shape == ds["dz_temp"]["nz_temp"].shape + ) # Make sure the depth dimension was broadcasted correctly + + +def test_generate_encoding(generate_silly_vt_dataset): + ds = generate_silly_vt_dataset + encoding_dict = {} + ds["temp_segment_002"] = ds["temp"] + ds.coords["temp_segment_003_nz_"] = ds.silly_depth + encoding_dict = rgd.generate_encoding(ds, encoding_dict, default_fill_value="-3") + assert ( + encoding_dict["temp_segment_002"]["_FillValue"] == "-3" + and "dtype" not in encoding_dict["temp_segment_002"] + ) + assert encoding_dict["temp_segment_003_nz_"]["dtype"] == "int32" + + +def test_get_boundary_mask(get_curvilinear_hgrid): + hgrid = get_curvilinear_hgrid + t_points = rgd.get_hgrid_arakawa_c_points(hgrid, "t") + bathy = hgrid.isel(nyp=t_points.t_points_y, nxp=t_points.t_points_x) + bathy["depth"] = (("t_points_y", "t_points_x"), (np.full(bathy.x.shape, 0))) + north_mask = rgd.get_boundary_mask( + hgrid, + bathy, + "north", + "segment_002", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + south_mask = rgd.get_boundary_mask( + hgrid, + bathy, + "south", + "segment_001", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + east_mask = rgd.get_boundary_mask( + hgrid, + bathy, + "east", + "segment_003", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + west_mask = rgd.get_boundary_mask( + hgrid, + bathy, + "west", + "segment_004", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + + # Check corner property of mask, and ensure each direction is following what we expect + for mask in [north_mask, south_mask, east_mask, west_mask]: + assert ( + mask[0] == 0 and mask[-1] == 0 + ) # Ensure Corners are oceans and set to zero for zeroing out values + assert np.isnan(mask[1:-1]).all() # Ensure all other points are land + assert north_mask.shape == (hgrid.x[-1].shape) # Ensure mask is the right shape + assert south_mask.shape == (hgrid.x[0].shape) # Ensure mask is the right shape + assert east_mask.shape == (hgrid.x[:, -1].shape) # Ensure mask is the right shape + assert west_mask.shape == (hgrid.x[:, 0].shape) # Ensure mask is the right shape + + ## Now we check if the coast masking is correct (remember we make 3 cells into the coast be ocean) + start_ind = 6 + end_ind = 9 + for i in range(start_ind, end_ind + 1): + bathy["depth"][-1][i] = 15 + north_mask = rgd.get_boundary_mask( + hgrid, + bathy, + "north", + "segment_002", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + assert ( + north_mask[0] == 0 and north_mask[-1] == 0 + ) # Ensure Corners are oceans and zeroed out if land + assert ( + north_mask[(((start_ind * 2) + 1)) : (((end_ind * 2) + 1) + 1)] == 1 + ).all() # Ensure coasts are ocean with a 3 cell buffer (remeber mask is on the hgrid boundary) so (6 *2 +2) - 3 -> (9 *2 +2) + 3 + assert ( + north_mask[(((start_ind * 2) + 1) - 3) : (((start_ind * 2) + 1))] == 0 + ).all() # Left Side + assert ( + north_mask[(((end_ind * 2) + 1) + 1) : (((end_ind * 2) + 1) + 3 + 1)] == 0 + ).all() # Right Side + + ## On E/W + start_ind = 6 + end_ind = 9 + for i in range(start_ind, end_ind + 1): + bathy["depth"][:, 0][i] = 15 + west_mask = rgd.get_boundary_mask( + hgrid, + bathy, + "west", + "segment_004", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + assert west_mask[0] == 0 and west_mask[-1] == 0 # Ensure Corners are oceans + assert ( + west_mask[(((start_ind * 2) + 1)) : (((end_ind * 2) + 1) + 1)] == 1 + ).all() # Ensure coasts are ocean with a 3 cell buffer (remeber mask is on the hgrid boundary) so (6 *2 +2) - 3 -> (9 *2 +2) + 3 + assert ( + west_mask[(((start_ind * 2) + 1) - 3) : (((start_ind * 2) + 1))] == 0 + ).all() # Ensure left side is zeroed out + assert ( + west_mask[(((end_ind * 2) + 1) + 1) : (((end_ind * 2) + 1) + 3 + 1)] == 0 + ).all() # Right Side is zeroed out + + +def test_mask_dataset(get_curvilinear_hgrid): + hgrid = get_curvilinear_hgrid + t_points = rgd.get_hgrid_arakawa_c_points(hgrid, "t") + bathy = hgrid.isel(nyp=t_points.t_points_y, nxp=t_points.t_points_x) + bathy["depth"] = (("t_points_y", "t_points_x"), (np.full(bathy.x.shape, 0))) + ds = hgrid.copy(deep=True) + ds = ds.drop_vars(("tile", "area", "y", "x", "angle_dx", "dy", "dx")) + ds["temp"] = (("t_points_y", "t_points_x"), (np.full(hgrid.x.shape, 100))) + ds["temp"] = ds["temp"].isel(t_points_y=-1) + start_ind = 6 + end_ind = 9 + for i in range(start_ind, end_ind + 1): + bathy["depth"][-1][i] = 15 + + ds["temp"][ + start_ind * 2 + 2 + ] = ( + np.nan + ) # Add a missing value not in the land mask to make sure it is filled with a dummy value + ds["temp"] = ds["temp"].expand_dims("nz_temp", axis=0) + ds = rgd.mask_dataset( + ds, + hgrid, + bathy, + "north", + "segment_002", + y_dim_name="t_points_y", + x_dim_name="t_points_x", + ) + + assert ( + np.isnan(ds["temp"][0][start_ind * 2 + 2]) == False + ) # Ensure missing value was filled + assert ( + np.isnan( + ds["temp"][0][(((start_ind * 2) + 1) - 3) : (((end_ind * 2) + 1) + 3 + 1)] + ) + ).all() == False # Ensure data is kept in ocean area + assert ( + np.isnan(ds["temp"][0][1 : (((start_ind * 2) + 1) - 3)]) + ).all() == True and ( + np.isnan(ds["temp"][0][(((end_ind * 2) + 1) + 3 + 1) : -1]) + ).all() == True # Ensure data is not in land area diff --git a/tests/test_rotation.py b/tests/test_rotation.py new file mode 100644 index 00000000..5604e5c3 --- /dev/null +++ b/tests/test_rotation.py @@ -0,0 +1,286 @@ +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 + + +def test_get_rotation_angle(get_curvilinear_hgrid, get_rectilinear_hgrid): + """ + Generate a curvilinear grid and test the grid rotation angle at t_points based on what we pass to generate to generate_curvilinear_grid + """ + curved_hgrid = get_curvilinear_hgrid + rect_hgrid = get_rectilinear_hgrid + + o = None + rotational_method = rot.RotationMethod.NO_ROTATION + angle = rot.get_rotation_angle(rotational_method, rect_hgrid, orientation=o) + assert angle.shape == rect_hgrid.x.shape + assert (angle.values == 0).all() + + rotational_method == rot.RotationMethod.NO_ROTATION + with pytest.raises( + ValueError, match="NO_ROTATION method only works with rectilinear grids" + ): + angle = rot.get_rotation_angle(rotational_method, curved_hgrid, orientation=o) + + rotational_method = rot.RotationMethod.GIVEN_ANGLE + angle = rot.get_rotation_angle(rotational_method, curved_hgrid, orientation=o) + assert angle.shape == curved_hgrid.x.shape + assert (angle.values == curved_hgrid.angle_dx).all() + angle = rot.get_rotation_angle(rotational_method, rect_hgrid, orientation=o) + assert angle.shape == rect_hgrid.x.shape + assert (angle.values == 0).all() + + rotational_method = rot.RotationMethod.EXPAND_GRID + angle = rot.get_rotation_angle(rotational_method, curved_hgrid, orientation=o) + assert angle.shape == curved_hgrid.x.shape + assert ( + abs(angle.values - curved_hgrid.angle_dx) < 1 + ).all() # There shouldn't be large differences + angle = rot.get_rotation_angle(rotational_method, rect_hgrid, orientation=o) + assert angle.shape == rect_hgrid.x.shape + assert (angle.values == 0).all() + + # Check if o is boundary that the shape is of a boundary + o = "north" + rotational_method = rot.RotationMethod.NO_ROTATION + angle = rot.get_rotation_angle(rotational_method, rect_hgrid, orientation=o) + assert angle.shape == rect_hgrid.x[-1].shape + assert (angle.values == 0).all() + rotational_method = rot.RotationMethod.EXPAND_GRID + angle = rot.get_rotation_angle(rotational_method, rect_hgrid, orientation=o) + assert angle.shape == rect_hgrid.x[-1].shape + assert (angle.values == 0).all() + rotational_method = rot.RotationMethod.GIVEN_ANGLE + angle = rot.get_rotation_angle(rotational_method, rect_hgrid, orientation=o) + assert angle.shape == rect_hgrid.x[-1].shape + assert (angle.values == 0).all() diff --git a/tests/test_tides_and_parameter.py b/tests/test_tides_and_parameter.py new file mode 100644 index 00000000..1eca029c --- /dev/null +++ b/tests/test_tides_and_parameter.py @@ -0,0 +1,202 @@ +""" +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 + + +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"