roms-tools 2.3.0__py3-none-any.whl → 2.5.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -5,17 +5,27 @@ from dataclasses import dataclass, field
 from typing import Dict, Union, List, Optional
 import matplotlib.pyplot as plt
 from pathlib import Path
+import logging
 from datetime import datetime
 from roms_tools import Grid
 from roms_tools.regrid import LateralRegrid, VerticalRegrid
 from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
-from roms_tools.utils import transpose_dimensions
+from roms_tools.utils import (
+    transpose_dimensions,
+    save_datasets,
+    get_dask_chunks,
+    interpolate_from_rho_to_u,
+    interpolate_from_rho_to_v,
+)
+from roms_tools.vertical_coordinate import (
+    compute_depth_coordinates,
+    compute_depth,
+)
 from roms_tools.setup.datasets import GLORYSDataset, CESMBGCDataset
 from roms_tools.setup.utils import (
     nan_check,
     substitute_nans_by_fillvalue,
     get_variable_metadata,
-    save_datasets,
     get_target_coords,
     rotate_velocities,
     compute_barotropic_velocity,
@@ -59,10 +69,17 @@ class InitialConditions:
             - A list of strings or Path objects containing multiple files.
           - "climatology" (bool): Indicates if the data is climatology data. Defaults to False.
 
+    adjust_depth_for_sea_surface_height : bool, optional
+        Whether to account for sea surface height variations when computing depth coordinates.
+        Defaults to `False`.
     model_reference_date : datetime, optional
         The reference date for the model. Defaults to January 1, 2000.
     use_dask: bool, optional
         Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
+    horizontal_chunk_size : int, optional
+        The chunk size used for horizontal partitioning for the vertical regridding when `use_dask = True`. Defaults to 50.
+        A larger number results in a bigger memory footprint but faster computations.
+        A smaller number results in a smaller memory footprint but slower computations.
     bypass_validation: bool, optional
         Indicates whether to skip validation checks in the processed data. When set to True,
         the validation process that ensures no NaN values exist at wet points
@@ -87,7 +104,9 @@ class InitialConditions:
     source: Dict[str, Union[str, Path, List[Union[str, Path]]]]
     bgc_source: Optional[Dict[str, Union[str, Path, List[Union[str, Path]]]]] = None
     model_reference_date: datetime = datetime(2000, 1, 1)
+    adjust_depth_for_sea_surface_height: bool = False
     use_dask: bool = False
+    horizontal_chunk_size: int = 50
     bypass_validation: bool = False
 
     ds: xr.Dataset = field(init=False, repr=False)
@@ -95,6 +114,8 @@ class InitialConditions:
     def __post_init__(self):
 
         self._input_checks()
+        # Dataset for depth coordinates
+        object.__setattr__(self, "ds_depth_coords", xr.Dataset())
 
         processed_fields = {}
         processed_fields = self._process_data(processed_fields, type="physics")
@@ -129,7 +150,8 @@ class InitialConditions:
             target_coords,
             buffer_points=20,  # lateral fill needs good buffer from data margin
         )
-
+        # Enforce double precision to ensure reproducibility
+        data.convert_to_float64()
         data.extrapolate_deepest_to_bottom()
         data.apply_lateral_fill()
 
@@ -140,6 +162,7 @@ class InitialConditions:
 
         # lateral regridding
         lateral_regrid = LateralRegrid(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(
@@ -148,60 +171,61 @@ class InitialConditions:
 
         # rotation of velocities and interpolation to u/v points
         if "u" in variable_info and "v" in variable_info:
-            (processed_fields["u"], processed_fields["v"],) = rotate_velocities(
+            processed_fields["u"], processed_fields["v"] = rotate_velocities(
                 processed_fields["u"],
                 processed_fields["v"],
                 target_coords["angle"],
                 interpolate=True,
             )
 
-        var_names_dict = {}
-        for location in ["rho", "u", "v"]:
-            var_names_dict[location] = [
+        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():
+                processed_fields[var_name] = processed_fields[var_name].assign_coords(
+                    {"time": processed_fields["temp"]["time"]}
+                )
+
+        # Get depth coordinates
+        zeta = (
+            processed_fields["zeta"] if self.adjust_depth_for_sea_surface_height else 0
+        )
 
-        # compute layer depth coordinates
-        if len(var_names_dict["u"]) > 0 or len(var_names_dict["v"]) > 0:
-            self._get_vertical_coordinates(
-                type="layer",
-                additional_locations=["u", "v"],
-            )
-        else:
-            if len(var_names_dict["rho"]) > 0:
-                self._get_vertical_coordinates(type="layer", additional_locations=[])
-        # vertical regridding
+        for location in ["rho", "u", "v"]:
+            if len(var_names_dict[location]) > 0:
+                self._get_depth_coordinates(zeta, location, "layer")
+
+        # Vertical regridding
         for location in ["rho", "u", "v"]:
             if len(var_names_dict[location]) > 0:
                 vertical_regrid = VerticalRegrid(
-                    self.grid.ds[f"layer_depth_{location}"],
+                    self.ds_depth_coords[f"layer_depth_{location}"],
                     data.ds[data.dim_names["depth"]],
                 )
                 for var_name in var_names_dict[location]:
                     if var_name in processed_fields:
-                        processed_fields[var_name] = vertical_regrid.apply(
-                            processed_fields[var_name]
-                        )
-
-        # compute barotropic velocities
+                        field = processed_fields[var_name]
+                        if self.use_dask:
+                            field = field.chunk(
+                                get_dask_chunks(location, self.horizontal_chunk_size)
+                            )
+                        processed_fields[var_name] = vertical_regrid.apply(field)
+
+        # Compute barotropic velocities
         if "u" in variable_info and "v" in variable_info:
-            self._get_vertical_coordinates(
-                type="interface",
-                additional_locations=["u", "v"],
-            )
             for location in ["u", "v"]:
+                self._get_depth_coordinates(zeta, location, "interface")
                 processed_fields[f"{location}bar"] = compute_barotropic_velocity(
                     processed_fields[location],
-                    self.grid.ds[f"interface_depth_{location}"],
-                )
-
-        if type == "bgc":
-            # Ensure time coordinate matches that of physical variables
-            for var_name in variable_info.keys():
-                processed_fields[var_name] = processed_fields[var_name].assign_coords(
-                    {"time": processed_fields["temp"]["time"]}
+                    self.ds_depth_coords[f"interface_depth_{location}"],
                 )
 
         for var_name in processed_fields.keys():
@@ -212,6 +236,11 @@ class InitialConditions:
         return processed_fields
 
     def _input_checks(self):
+        # Check that ini_time is not None
+        if self.ini_time is None:
+            raise ValueError(
+                "`ini_time` must be a valid datetime object and cannot be None."
+            )
 
         if "name" not in self.source.keys():
             raise ValueError("`source` must include a 'name'.")
@@ -244,6 +273,12 @@ class InitialConditions:
                     "climatology": self.bgc_source.get("climatology", False),
                 },
             )
+        if self.adjust_depth_for_sea_surface_height:
+            logging.info("Sea surface height will be used to adjust depth coordinates.")
+        else:
+            logging.info(
+                "Sea surface height will NOT be used to adjust depth coordinates."
+            )
 
     def _get_data(self):
 
@@ -365,35 +400,61 @@ class InitialConditions:
 
         object.__setattr__(self, f"variable_info_{type}", variable_info)
 
-    def _get_vertical_coordinates(self, type, additional_locations=["u", "v"]):
-        """Retrieve layer and interface depth coordinates.
-
-        This method computes and updates the layer and interface depth coordinates. It handles depth calculations for rho points and
-        additional specified locations (u and v).
+    def _get_depth_coordinates(
+        self, zeta: xr.DataArray | float, location: str, depth_type: str = "layer"
+    ) -> None:
+        """Ensure depth coordinates are computed and stored for a given location and
+        depth type.
 
         Parameters
         ----------
-        type : str
-            The type of depth coordinate to retrieve. Valid options are:
-            - "layer": Retrieves layer depth coordinates.
-            - "interface": Retrieves interface depth coordinates.
-
-        additional_locations : list of str, optional
-            Specifies additional locations to compute depth coordinates for. Default is ["u", "v"].
-            Valid options include:
-            - "u": Computes depth coordinates for u points.
-            - "v": Computes depth coordinates for v points.
-
-        Updates
-        -------
-        self.grid.ds : xarray.Dataset
-            The dataset is updated with the following vertical depth coordinates:
-            - f"{type}_depth_rho": Depth coordinates at rho points.
-            - f"{type}_depth_u": Depth coordinates at u points (if applicable).
-            - f"{type}_depth_v": Depth coordinates at v points (if applicable).
+        zeta : xr.DataArray or float
+            Free-surface elevation (can be a scalar or a DataArray).
+        location : str
+            Grid location for depth computation ("rho", "u", or "v").
+        depth_type : str, optional
+            Type of depth coordinates to compute, by default "layer".
+
+        Notes
+        ------
+        Rather than calling compute_depth_coordinates from the vertical_coordinate.py module,
+        this method computes the depth coordinates from scratch because of optional chunking.
         """
+        key = f"{depth_type}_depth_{location}"
+
+        if key not in self.ds_depth_coords:
+            # Select the appropriate depth computation parameters
+            if depth_type == "layer":
+                Cs = self.grid.ds["Cs_r"]
+                sigma = self.grid.ds["sigma_r"]
+            elif depth_type == "interface":
+                Cs = self.grid.ds["Cs_w"]
+                sigma = self.grid.ds["sigma_w"]
+            else:
+                raise ValueError(
+                    f"Invalid depth_type: {depth_type}. Choose 'layer' or 'interface'."
+                )
 
-        self.grid.compute_depth_coordinates(type, additional_locations)
+            h = self.grid.ds["h"]
+
+            # Interpolate h and zeta to the specified location
+            if location == "u":
+                h = interpolate_from_rho_to_u(h)
+                if isinstance(zeta, xr.DataArray):
+                    zeta = interpolate_from_rho_to_u(zeta)
+            elif location == "v":
+                h = interpolate_from_rho_to_v(h)
+                if isinstance(zeta, xr.DataArray):
+                    zeta = interpolate_from_rho_to_v(zeta)
+
+            if self.use_dask:
+                h = h.chunk(get_dask_chunks(location, self.horizontal_chunk_size))
+                if self.adjust_depth_for_sea_surface_height:
+                    zeta = zeta.chunk(
+                        get_dask_chunks(location, self.horizontal_chunk_size)
+                    )
+            depth = compute_depth(zeta, h, self.grid.ds.attrs["hc"], Cs, sigma)
+            self.ds_depth_coords[key] = depth
 
     def _write_into_dataset(self, processed_fields, d_meta):
 
@@ -407,7 +468,7 @@ class InitialConditions:
 
         # initialize vertical velocity to zero
         ds["w"] = xr.zeros_like(
-            self.grid.ds["interface_depth_rho"].expand_dims(
+            (self.grid.ds["Cs_w"] * self.grid.ds["h"]).expand_dims(
                 time=processed_fields["u"].time
             )
         ).astype(np.float32)
@@ -504,6 +565,9 @@ class InitialConditions:
         ds.attrs["roms_tools_version"] = roms_tools_version
         ds.attrs["ini_time"] = str(self.ini_time)
         ds.attrs["model_reference_date"] = str(self.model_reference_date)
+        ds.attrs["adjust_depth_for_sea_surface_height"] = str(
+            self.adjust_depth_for_sea_surface_height
+        )
         ds.attrs["source"] = self.source["name"]
         if self.bgc_source is not None:
             ds.attrs["bgc_source"] = self.bgc_source["name"]
@@ -637,7 +701,6 @@ class InitialConditions:
             loc = "rho"
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
             loc = "u"
-
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
             loc = "v"
         else:
@@ -656,47 +719,30 @@ class InitialConditions:
         if s is not None:
             layer_contours = False
         # Note that `layer_depth_{loc}` has already been computed during `__post_init__`.
-        layer_depth = self.grid.ds[f"layer_depth_{loc}"]
-        if layer_contours:
-            if f"interface_depth_{loc}" not in self.grid.ds:
-                if loc == "rho":
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=[]
-                    )
-                else:
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=["u", "v"]
-                    )
-            interface_depth = self.grid.ds[f"interface_depth_{loc}"]
-        else:
-            interface_depth = None
+        layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"].squeeze()
 
         # Slice the field as desired
         def _slice_and_assign(
             field,
             mask,
             layer_depth,
-            interface_depth,
             title,
             dim_name,
             dim_values,
             idx,
-            layer_contours=False,
         ):
             if dim_name in field.dims:
                 title = title + f", {dim_name} = {dim_values[idx].item()}"
                 field = field.isel(**{dim_name: idx})
                 mask = mask.isel(**{dim_name: idx})
                 layer_depth = layer_depth.isel(**{dim_name: idx})
-                if layer_contours:
-                    interface_depth = interface_depth.isel(**{dim_name: idx})
                 if "s_rho" in field.dims:
                     field = field.assign_coords({"layer_depth": layer_depth})
             else:
                 raise ValueError(
                     f"None of the expected dimensions ({dim_name}) found in field."
                 )
-            return field, mask, layer_depth, interface_depth, title
+            return field, mask, layer_depth, title
 
         title = field.long_name
         if s is not None:
@@ -708,29 +754,25 @@ class InitialConditions:
             depth_contours = False
 
         if eta is not None:
-            field, mask, layer_depth, interface_depth, title = _slice_and_assign(
+            field, mask, layer_depth, title = _slice_and_assign(
                 field,
                 mask,
                 layer_depth,
-                interface_depth,
                 title,
                 "eta_rho" if "eta_rho" in field.dims else "eta_v",
                 field.eta_rho if "eta_rho" in field.dims else field.eta_v,
                 eta,
-                layer_contours,
             )
 
         if xi is not None:
-            field, mask, layer_depth, interface_depth, title = _slice_and_assign(
+            field, mask, layer_depth, title = _slice_and_assign(
                 field,
                 mask,
                 layer_depth,
-                interface_depth,
                 title,
                 "xi_rho" if "xi_rho" in field.dims else "xi_u",
                 field.xi_rho if "xi_rho" in field.dims else field.xi_u,
                 xi,
-                layer_contours,
             )
 
         # Choose colorbar
@@ -757,19 +799,54 @@ class InitialConditions:
                 c="g",
             )
         else:
-            if not layer_contours:
-                interface_depth = None
-            else:
-                # restrict number of layer_contours to 10 for the sake of plot clearity
-                nr_layers = len(interface_depth["s_w"])
-                selected_layers = np.linspace(
-                    0, nr_layers - 1, min(nr_layers, 10), dtype=int
-                )
-                interface_depth = interface_depth.isel(s_w=selected_layers)
-
             if len(field.dims) == 2:
+                if layer_contours:
+                    if loc == "rho":
+                        # interface_depth_rho has not been computed yet
+                        interface_depth = compute_depth_coordinates(
+                            self.grid.ds,
+                            self.ds.zeta,
+                            depth_type="interface",
+                            location=loc,
+                            eta=eta,
+                            xi=xi,
+                        )
+                    elif loc == "u":
+                        index_kwargs = {}
+                        if eta is not None:
+                            index_kwargs["eta_rho"] = eta
+                        if xi is not None:
+                            index_kwargs["xi_u"] = xi
+
+                        interface_depth = (
+                            self.ds_depth_coords[f"interface_depth_{loc}"]
+                            .isel(**index_kwargs)
+                            .squeeze()
+                        )
+                    elif loc == "v":
+                        index_kwargs = {}
+                        if eta is not None:
+                            index_kwargs["eta_v"] = eta
+                        if xi is not None:
+                            index_kwargs["xi_rho"] = xi
+
+                        interface_depth = (
+                            self.ds_depth_coords[f"interface_depth_{loc}"]
+                            .isel(**index_kwargs)
+                            .squeeze()
+                        )
+
+                    # restrict number of layer_contours to 10 for the sake of plot clearity
+                    nr_layers = len(interface_depth["s_w"])
+                    selected_layers = np.linspace(
+                        0, nr_layers - 1, min(nr_layers, 10), dtype=int
+                    )
+                    interface_depth = interface_depth.isel(s_w=selected_layers)
+                else:
+                    interface_depth = None
+
                 _section_plot(
-                    field.where(mask),
+                    field,
                     interface_depth=interface_depth,
                     title=title,
                     kwargs=kwargs,
@@ -781,36 +858,18 @@ class InitialConditions:
                 else:
                     _line_plot(field.where(mask), title=title, ax=ax)
 
-    def save(
-        self, filepath: Union[str, Path], np_eta: int = None, np_xi: int = None
-    ) -> None:
-        """Save the initial conditions information to a netCDF4 file.
-
-        This method supports saving the dataset in two modes:
-
-          1. **Single File Mode (default)**:
-
-            If both `np_eta` and `np_xi` are `None`, the entire dataset is saved as a single netCDF4 file
-            with the base filename specified by `filepath.nc`.
-
-          2. **Partitioned Mode**:
-
-            - If either `np_eta` or `np_xi` is specified, the dataset is divided into spatial tiles along the eta-axis and xi-axis.
-            - Each spatial tile is saved as a separate netCDF4 file.
+    def save(self, filepath: Union[str, Path]) -> None:
+        """Save the initial conditions information to one netCDF4 file.
 
         Parameters
         ----------
         filepath : Union[str, Path]
             The base path or filename where the dataset should be saved.
-        np_eta : int, optional
-            The number of partitions along the `eta` direction. If `None`, no spatial partitioning is performed.
-        np_xi : int, optional
-            The number of partitions along the `xi` direction. If `None`, no spatial partitioning is performed.
 
         Returns
         -------
-        List[Path]
-            A list of Path objects for the filenames that were saved.
+        Path
+            A `Path` object representing the location of the saved file.
         """
 
         # Ensure filepath is a Path object
@@ -820,17 +879,11 @@ class InitialConditions:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                self.ds.load()
-
         dataset_list = [self.ds]
         output_filenames = [str(filepath)]
 
         saved_filenames = save_datasets(
-            dataset_list, output_filenames, np_eta=np_eta, np_xi=np_xi
+            dataset_list, output_filenames, use_dask=self.use_dask
         )
 
         return saved_filenames
@@ -852,7 +905,6 @@ class InitialConditions:
         cls,
         filepath: Union[str, Path],
         use_dask: bool = False,
-        bypass_validation: bool = False,
     ) -> "InitialConditions":
         """Create an instance of the InitialConditions class from a YAML file.
 
@@ -862,10 +914,6 @@ class InitialConditions:
             The path to the YAML file from which the parameters will be read.
         use_dask: bool, optional
             Indicates whether to use dask for processing. If True, data is processed with dask; if False, data is processed eagerly. Defaults to False.
-        bypass_validation: bool, optional
-            Indicates whether to skip validation checks in the processed data. When set to True,
-            the validation process that ensures no NaN values exist at wet points
-            in the processed dataset is bypassed. Defaults to False.
 
         Returns
         -------
@@ -880,5 +928,4 @@ class InitialConditions:
             grid=grid,
             **initial_conditions_params,
             use_dask=use_dask,
-            bypass_validation=bypass_validation,
         )