roms-tools 2.6.1__py3-none-any.whl → 2.7.0__py3-none-any.whl

roms_tools/setup/grid.py

@@ -8,6 +8,7 @@ import matplotlib.pyplot as plt
 import yaml
 import importlib.metadata
 from typing import Dict, Union, List
+from roms_tools.constants import R_EARTH
 from roms_tools.utils import save_datasets
 from roms_tools.setup.topography import _add_topography
 from roms_tools.setup.mask import _add_mask, _add_velocity_masks
@@ -918,8 +919,6 @@ class Grid:
             - lonq, latq: 2D arrays of longitudes and latitudes at cell corners.
         """
 
-        r_earth = 6371315.0
-
         # initially define the domain to be longer in x-direction (dimension "length")
         # than in y-direction (dimension "width") to keep grid distortion minimal
         if self.size_y > self.size_x:
@@ -929,8 +928,8 @@ class Grid:
             domain_length, domain_width = self.size_x * 1e3, self.size_y * 1e3  # in m
             nl, nw = self.nx, self.ny
 
-        domain_length_in_degrees = domain_length / r_earth
-        domain_width_in_degrees = domain_width / r_earth
+        domain_length_in_degrees = domain_length / R_EARTH
+        domain_width_in_degrees = domain_width / R_EARTH
 
         # Generate 1D longitude arrays at cell centers and corners
         lon_array_1d_in_degrees = domain_length_in_degrees * (
@@ -1098,6 +1097,45 @@ class Grid:
 
         return ds
 
+    def _infer_nominal_horizontal_resolution(self, lat=None):
+        """Estimate the nominal horizontal resolution of the grid in degrees at a
+        specified latitude.
+
+        This method calculates the nominal horizontal resolution of the grid by first
+        determining the average grid spacing in meters. The spacing is then converted
+        to degrees, accounting for the Earth's curvature, and the latitude where the
+        resolution is being computed.
+
+        Parameters
+        ----------
+        lat : float, optional
+            Latitude (in degrees) at which to estimate the horizontal resolution.
+            If not provided, the resolution is calculated at the average latitude of
+            the grid (`lat_rho`).
+
+        Returns
+        -------
+        float
+            The estimated horizontal resolution in degrees, adjusted for the Earth's curvature.
+        """
+        if lat is None:
+            # Center latitude in degrees
+            lat = (self.ds.lat_rho.max() + self.ds.lat_rho.min()) / 2
+
+        # Convert latitude to radians
+        lat_rad = np.deg2rad(lat)
+
+        # Mean resolution in meters
+        resolution_in_m = ((1 / self.ds.pm).mean() + (1 / self.ds.pn).mean()) / 2
+
+        # Meters per degree at the equator
+        meters_per_degree = 2 * np.pi * R_EARTH / 360
+
+        # Correct for latitude by multiplying by cos(latitude) for longitude
+        resolution_in_degrees = resolution_in_m / (meters_per_degree * np.cos(lat_rad))
+
+        return resolution_in_degrees
+
 
 def _rotate(coords, rot):
     """Rotate grid counterclockwise relative to surface of Earth by rot degrees."""

roms_tools/setup/initial_conditions.py

@@ -21,7 +21,7 @@ from roms_tools.vertical_coordinate import (
     compute_depth_coordinates,
     compute_depth,
 )
-from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
+from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset, UnifiedBGCDataset
 from roms_tools.setup.utils import (
     nan_check,
     substitute_nans_by_fillvalue,
@@ -29,6 +29,7 @@ from roms_tools.setup.utils import (
     get_target_coords,
     rotate_velocities,
     compute_barotropic_velocity,
+    compute_missing_bgc_variables,
     _to_yaml,
     _from_yaml,
 )
@@ -122,6 +123,12 @@ class InitialConditions:
 
         if self.bgc_source is not None:
             processed_fields = self._process_data(processed_fields, type="bgc")
+            processed_fields = compute_missing_bgc_variables(processed_fields)
+
+        for var_name in processed_fields:
+            processed_fields[var_name] = transpose_dimensions(
+                processed_fields[var_name]
+            )
 
         d_meta = get_variable_metadata()
         ds = self._write_into_dataset(processed_fields, d_meta)
@@ -158,19 +165,42 @@ class InitialConditions:
         self._set_variable_info(data, type=type)
         attr_name = f"variable_info_{type}"
         variable_info = getattr(self, attr_name)
-        var_names = variable_info.keys()
+
+        # Create the var_names dictionary, associating each variable with its location
+        # Avoid looping over processed_fields.keys() directly, as they may already contain
+        # finalized physics variables. This is especially important when transitioning
+        # to processing biogeochemical (BGC) variables, ensuring that only relevant
+        # variables are processed.
+        var_names = {
+            var: {
+                "name": data.var_names[var],
+                "location": variable_info[var]["location"],
+            }
+            for var in data.var_names.keys()
+            if data.var_names[var] in data.ds.data_vars
+        }
+        # Update the dictionary with optional variables and their locations
+        var_names.update(
+            {
+                var: {
+                    "name": data.opt_var_names[var],
+                    "location": variable_info[var]["location"],
+                }
+                for var in data.opt_var_names.keys()
+                if data.opt_var_names[var] in data.ds.data_vars
+            }
+        )
 
         # lateral regridding
         lateral_regrid = LateralRegridToROMS(target_coords, data.dim_names)
 
         for var_name in var_names:
-            if var_name in data.var_names.keys():
-                processed_fields[var_name] = lateral_regrid.apply(
-                    data.ds[data.var_names[var_name]]
-                )
+            processed_fields[var_name] = lateral_regrid.apply(
+                data.ds[var_names[var_name]["name"]]
+            )
 
         # rotation of velocities and interpolation to u/v points
-        if "u" in variable_info and "v" in variable_info:
+        if "u" in var_names and "v" in var_names:
             processed_fields["u"], processed_fields["v"] = rotate_velocities(
                 processed_fields["u"],
                 processed_fields["v"],
@@ -178,18 +208,9 @@ class InitialConditions:
                 interpolate=True,
             )
 
-        var_names_dict = {
-            location: [
-                name
-                for name, info in variable_info.items()
-                if info["location"] == location and info["is_3d"]
-            ]
-            for location in ["rho", "u", "v"]
-        }
-
         if type == "bgc":
             # Ensure time coordinate matches that of physical variables
-            for var_name in variable_info.keys():
+            for var_name in var_names:
                 processed_fields[var_name] = processed_fields[var_name].assign_coords(
                     {"time": processed_fields["temp"]["time"]}
                 )
@@ -200,17 +221,23 @@ class InitialConditions:
         )
 
         for location in ["rho", "u", "v"]:
-            if len(var_names_dict[location]) > 0:
+            # Filter var_names by location and check for 3D variables
+            filtered_vars = [
+                var_name
+                for var_name, info in var_names.items()
+                if info["location"] == location and variable_info[var_name]["is_3d"]
+            ]
+
+            if filtered_vars:
+                # Handle depth coordinates
                 self._get_depth_coordinates(zeta, location, "layer")
 
-        # Vertical regridding
-        for location in ["rho", "u", "v"]:
-            if len(var_names_dict[location]) > 0:
+                # Vertical regridding
                 vertical_regrid = VerticalRegridToROMS(
                     self.ds_depth_coords[f"layer_depth_{location}"],
                     data.ds[data.dim_names["depth"]],
                 )
-                for var_name in var_names_dict[location]:
+                for var_name in filtered_vars:
                     if var_name in processed_fields:
                         field = processed_fields[var_name]
                         if self.use_dask:
@@ -220,7 +247,7 @@ class InitialConditions:
                         processed_fields[var_name] = vertical_regrid.apply(field)
 
         # Compute barotropic velocities
-        if "u" in variable_info and "v" in variable_info:
+        if "u" in var_names and "v" in var_names:
             for location in ["u", "v"]:
                 self._get_depth_coordinates(zeta, location, "interface")
                 processed_fields[f"{location}bar"] = compute_barotropic_velocity(
@@ -228,11 +255,6 @@ class InitialConditions:
                     self.ds_depth_coords[f"interface_depth_{location}"],
                 )
 
-        for var_name in processed_fields.keys():
-            processed_fields[var_name] = transpose_dimensions(
-                processed_fields[var_name]
-            )
-
         return processed_fields
 
     def _input_checks(self):
@@ -295,9 +317,17 @@ class InitialConditions:
                 climatology=self.bgc_source["climatology"],
                 use_dask=self.use_dask,
             )
+        elif self.bgc_source["name"] == "UNIFIED":
+            data = UnifiedBGCDataset(
+                filename=self.bgc_source["path"],
+                start_time=self.ini_time,
+                climatology=self.bgc_source["climatology"],
+                use_dask=self.use_dask,
+            )
+
         else:
             raise ValueError(
-                'Only "CESM_REGRIDDED" is a valid option for bgc_source["name"].'
+                'Only "CESM_REGRIDDED" and "UNIFIED" are valid options for bgc_source["name"].'
             )
 
         return data
@@ -384,7 +414,10 @@ class InitialConditions:
 
         if type == "bgc":
             variable_info = {}
-            for var_name in data.var_names.keys():
+
+            for var_name in list(data.var_names.keys()) + list(
+                data.opt_var_names.keys()
+            ):
                 if var_name == "ALK":
                     variable_info[var_name] = {**default_info, "validate": True}
                 else:
@@ -453,10 +486,13 @@ class InitialConditions:
         # save in new dataset
         ds = xr.Dataset()
 
-        for var_name in processed_fields.keys():
-            ds[var_name] = processed_fields[var_name].astype(np.float32)
-            ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
-            ds[var_name].attrs["units"] = d_meta[var_name]["units"]
+        for var_name in processed_fields:
+            if var_name in d_meta:
+
+                # drop auxiliary variables
+                ds[var_name] = processed_fields[var_name].astype(np.float32)
+                ds[var_name].attrs["long_name"] = d_meta[var_name]["long_name"]
+                ds[var_name].attrs["units"] = d_meta[var_name]["units"]
 
         # initialize vertical velocity to zero
         ds["w"] = xr.zeros_like(

roms_tools/setup/river_forcing.py

@@ -9,7 +9,8 @@ from pathlib import Path
 import matplotlib.pyplot as plt
 import matplotlib.cm as cm
 from roms_tools import Grid
-from roms_tools.plot import _get_projection, _add_field_to_ax
+from roms_tools.constants import NUM_TRACERS
+from roms_tools.plot import _plot, _get_projection
 from roms_tools.utils import save_datasets
 from roms_tools.setup.datasets import DaiRiverDataset
 from roms_tools.setup.utils import (
@@ -19,6 +20,7 @@ from roms_tools.setup.utils import (
     convert_to_roms_time,
     _to_yaml,
     _from_yaml,
+    add_tracer_metadata_to_ds,
     get_variable_metadata,
 )
 
@@ -44,6 +46,7 @@ class RiverForcing:
             - A single string (with or without wildcards).
             - A single Path object.
             - A list of strings or Path objects containing multiple files.
+
           - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
 
         The default is the Dai and Trenberth global river dataset (updated in May 2019), which does not require a path.
@@ -296,7 +299,7 @@ class RiverForcing:
         ds = ds.assign_coords({"nriver": nriver})
 
         if self.include_bgc:
-            ntracers = 2 + 32
+            ntracers = NUM_TRACERS
         else:
             ntracers = 2
         tracer_data = np.zeros(
@@ -310,52 +313,8 @@ class RiverForcing:
             tracer_data, dims=("river_time", "ntracers", "nriver")
         )
         river_tracer.attrs["long_name"] = "River tracer data"
-
-        if self.include_bgc:
-            tracer_names = xr.DataArray(
-                [
-                    "temp",
-                    "salt",
-                    "PO4",
-                    "NO3",
-                    "SiO3",
-                    "NH4",
-                    "Fe",
-                    "Lig",
-                    "O2",
-                    "DIC",
-                    "DIC_ALT_CO2",
-                    "ALK",
-                    "ALK_ALT_CO2",
-                    "DOC",
-                    "DON",
-                    "DOP",
-                    "DOPr",
-                    "DONr",
-                    "DOCr",
-                    "zooC",
-                    "spChl",
-                    "spC",
-                    "spP",
-                    "spFe",
-                    "spCaCO3",
-                    "diatChl",
-                    "diatC",
-                    "diatP",
-                    "diatFe",
-                    "diatSi",
-                    "diazChl",
-                    "diazC",
-                    "diazP",
-                    "diazFe",
-                ],
-                dims="ntracers",
-            )
-        else:
-            tracer_names = xr.DataArray(["temp", "salt"], dims="ntracers")
-        tracer_names.attrs["long_name"] = "Tracer name"
-        river_tracer.coords["tracer_name"] = tracer_names
         ds["river_tracer"] = river_tracer
+        ds = add_tracer_metadata_to_ds(ds, self.include_bgc)
 
         ds, time = convert_to_roms_time(
             ds, self.model_reference_date, self.climatology, time_name="river_time"
@@ -441,24 +400,25 @@ class RiverForcing:
         return river_indices
 
     def _write_indices_into_dataset(self, ds):
-        """Adds river location indices to the dataset as the "river_flux" variable.
+        """Adds river location indices to the dataset as the "river_index" and
+        "river_fraction" variables.
 
-        This method creates a new "river_flux" variable
-        using river station indices from `self.indices` and assigns it to the dataset.
-        The indices specify the river station locations in terms of eta_rho and xi_rho grid cell indices.
+        This method creates new "river_index" and "river_fraction" variables using river station indices
+        from `self.indices` and assigns them to the dataset.
 
         Parameters
         ----------
         ds : xarray.Dataset
-            The dataset to which the "river_flux" variable will be added.
+            The dataset to which the "river_index" and "river_fraction" variables will be added.
 
         Returns
         -------
         xarray.Dataset
-            The modified dataset with the "river_flux" variable added.
+            The modified dataset with the "river_index" and "river_fraction" variables added.
         """
 
-        river_locations = xr.zeros_like(self.grid.ds.h)
+        river_index = xr.zeros_like(self.grid.ds.h)
+        river_fraction = xr.zeros_like(self.grid.ds.h)
 
         for nriver in ds.nriver:
             river_name = str(ds.river_name.sel(nriver=nriver).values)
@@ -467,14 +427,18 @@ class RiverForcing:
 
             for eta_index, xi_index in indices:
 
-                river_locations[eta_index, xi_index] = (
-                    nriver  # assign unique nriver ID (Fortran-based indexing)
-                    + fraction  # Fractional contribution for multiple grid points
-                )
+                # Assign unique nriver ID (Fortran-based indexing)
+                river_index[eta_index, xi_index] = nriver
+                # Fractional contribution for multiple grid points
+                river_fraction[eta_index, xi_index] = fraction
 
-        river_locations.attrs["long_name"] = "River ID plus local volume fraction"
-        river_locations.attrs["units"] = "none"
-        ds["river_flux"] = river_locations
+        river_index.attrs["long_name"] = "River ID"
+        river_index.attrs["units"] = "none"
+        ds["river_index"] = river_index
+
+        river_fraction.attrs["long_name"] = "River volume fraction"
+        river_fraction.attrs["units"] = "none"
+        ds["river_fraction"] = river_fraction
 
         ds = ds.drop_vars(["lat_rho", "lon_rho"])
 
@@ -507,56 +471,25 @@ class RiverForcing:
         """Plots the original and updated river locations on a map projection."""
 
         field = self.grid.ds.mask_rho
-        vmax = 3
-        vmin = 0
-        cmap = plt.colormaps.get_cmap("Blues")
-        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
-
         lon_deg = self.grid.ds.lon_rho
         lat_deg = self.grid.ds.lat_rho
-
-        # check if North or South pole are in domain
-        if lat_deg.max().values > 89 or lat_deg.min().values < -89:
-            raise NotImplementedError(
-                "Plotting is not implemented for the case that the domain contains the North or South pole."
-            )
-
         if self.grid.straddle:
             lon_deg = xr.where(lon_deg > 180, lon_deg - 360, lon_deg)
         field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
-        trans = _get_projection(lon_deg, lat_deg)
+        vmax = 6
+        vmin = 0
+        cmap = plt.colormaps.get_cmap("Blues")
+        kwargs = {"vmax": vmax, "vmin": vmin, "cmap": cmap}
 
-        lon_deg = lon_deg.values
-        lat_deg = lat_deg.values
+        trans = _get_projection(lon_deg, lat_deg)
 
         fig, axs = plt.subplots(
             1, 2, figsize=(13, 13), subplot_kw={"projection": trans}
         )
 
         for ax in axs:
-            _add_field_to_ax(
-                ax,
-                lon_deg,
-                lat_deg,
-                field,
-                add_colorbar=False,
-                kwargs=kwargs,
-            )
-            # Add gridlines with labels for latitude and longitude
-            gridlines = ax.gridlines(
-                draw_labels=True, linewidth=0.5, color="gray", alpha=0.7, linestyle="--"
-            )
-            gridlines.top_labels = False  # Hide top labels
-            gridlines.right_labels = False  # Hide right labels
-            gridlines.xlabel_style = {
-                "size": 10,
-                "color": "black",
-            }  # Customize longitude label style
-            gridlines.ylabel_style = {
-                "size": 10,
-                "color": "black",
-            }  # Customize latitude label style
+            _plot(field, kwargs=kwargs, ax=ax, c=None, add_colorbar=False)
 
         proj = ccrs.PlateCarree()
 

roms_tools/setup/surface_forcing.py

@@ -8,13 +8,14 @@ from pathlib import Path
 import logging
 from typing import Dict, Union, List, Optional
 from roms_tools import Grid
-from roms_tools.utils import save_datasets
+from roms_tools.utils import save_datasets, transpose_dimensions
 from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.plot import _plot
 from roms_tools.setup.datasets import (
     ERA5Dataset,
     ERA5Correction,
     CESMBGCSurfaceForcingDataset,
+    UnifiedBGCSurfaceDataset,
 )
 from roms_tools.setup.utils import (
     get_target_coords,
@@ -24,6 +25,7 @@ from roms_tools.setup.utils import (
     get_variable_metadata,
     group_dataset,
     rotate_velocities,
+    compute_missing_surface_bgc_variables,
     convert_to_roms_time,
     _to_yaml,
     _from_yaml,
@@ -160,6 +162,15 @@ class SurfaceForcing:
         if self.type == "physics" and self.correct_radiation:
             processed_fields = self._apply_correction(processed_fields, data)
 
+        if self.type == "bgc":
+            processed_fields = compute_missing_surface_bgc_variables(processed_fields)
+
+        # Reorder dimensions
+        for var_name in processed_fields:
+            processed_fields[var_name] = transpose_dimensions(
+                processed_fields[var_name]
+            )
+
         d_meta = get_variable_metadata()
 
         ds = self._write_into_dataset(processed_fields, data, d_meta)
@@ -265,9 +276,12 @@ class SurfaceForcing:
             if self.source["name"] == "CESM_REGRIDDED":
 
                 data = CESMBGCSurfaceForcingDataset(**data_dict)
+            elif self.source["name"] == "UNIFIED":
+
+                data = UnifiedBGCSurfaceDataset(**data_dict)
             else:
                 raise ValueError(
-                    'Only "CESM_REGRIDDED" is a valid option for source["name"] when type is "bgc".'
+                    'Only "CESM_REGRIDDED" and "UNIFIED" are valid options for source["name"] when type is "bgc".'
                 )
 
         return data
@@ -331,7 +345,9 @@ class SurfaceForcing:
             }
         elif self.type == "bgc":
             variable_info = {}
-            for var_name in data.var_names.keys():
+            for var_name in list(data.var_names.keys()) + list(
+                data.opt_var_names.keys()
+            ):
                 variable_info[var_name] = default_info
                 if var_name == "pco2_air":
                     variable_info[var_name] = {**default_info, "validate": True}