roms-tools 2.2.1__py3-none-any.whl → 2.4.0__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -3,28 +3,35 @@ import numpy as np
 import importlib.metadata
 from dataclasses import dataclass, field
 from typing import Dict, Union, List, Optional
-from roms_tools.setup.grid import Grid
+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,
+    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.vertical_coordinate import compute_depth
 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,
-    transpose_dimensions,
-    interpolate_from_rho_to_u,
-    interpolate_from_rho_to_v,
     _to_yaml,
     _from_yaml,
 )
-from roms_tools.setup.regrid import LateralRegrid, VerticalRegrid
-from roms_tools.setup.plot import _plot, _section_plot, _profile_plot, _line_plot
-import matplotlib.pyplot as plt
-from pathlib import Path
 
 
 @dataclass(frozen=True, kw_only=True)
@@ -62,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
@@ -90,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)
@@ -98,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")
@@ -132,7 +150,6 @@ class InitialConditions:
             target_coords,
             buffer_points=20,  # lateral fill needs good buffer from data margin
         )
-
         data.extrapolate_deepest_to_bottom()
         data.apply_lateral_fill()
 
@@ -143,6 +160,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(
@@ -151,60 +169,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"]
+        }
 
-        # 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
+        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
+        )
+
+        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():
@@ -247,6 +266,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):
 
@@ -368,82 +393,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".
 
-        layer_vars = []
-        for location in ["rho"] + additional_locations:
-            layer_vars.append(f"{type}_depth_{location}")
-
-        if all(layer_var in self.grid.ds for layer_var in layer_vars):
-            # Vertical coordinate data already exists
-            pass
-
-        elif f"{type}_depth_rho" in self.grid.ds:
-            depth = self.grid.ds[f"{type}_depth_rho"]
-
-            if "u" in additional_locations or "v" in additional_locations:
-                # interpolation
-                if "u" in additional_locations:
-                    depth_u = interpolate_from_rho_to_u(depth)
-                    depth_u.attrs["long_name"] = f"{type} depth at u-points"
-                    depth_u.attrs["units"] = "m"
-                    self.grid.ds[f"{type}_depth_u"] = depth_u
-                if "v" in additional_locations:
-                    depth_v = interpolate_from_rho_to_v(depth)
-                    depth_v.attrs["long_name"] = f"{type} depth at v-points"
-                    depth_v.attrs["units"] = "m"
-                    self.grid.ds[f"{type}_depth_v"] = depth_v
-        else:
-            h = self.grid.ds["h"]
-            if type == "layer":
-                depth = compute_depth(
-                    0, h, self.grid.hc, self.grid.ds.Cs_r, self.grid.ds.sigma_r
-                )
+        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:
-                depth = compute_depth(
-                    0, h, self.grid.hc, self.grid.ds.Cs_w, self.grid.ds.sigma_w
+                raise ValueError(
+                    f"Invalid depth_type: {depth_type}. Choose 'layer' or 'interface'."
                 )
 
-            depth.attrs["long_name"] = f"{type} depth at rho-points"
-            depth.attrs["units"] = "m"
-            self.grid.ds[f"{type}_depth_rho"] = depth
-
-            if "u" in additional_locations or "v" in additional_locations:
-                # interpolation
-                depth_u = interpolate_from_rho_to_u(depth)
-                depth_u.attrs["long_name"] = f"{type} depth at u-points"
-                depth_u.attrs["units"] = "m"
-                depth_v = interpolate_from_rho_to_v(depth)
-                depth_v.attrs["long_name"] = f"{type} depth at v-points"
-                depth_v.attrs["units"] = "m"
-                self.grid.ds[f"{type}_depth_u"] = depth_u
-                self.grid.ds[f"{type}_depth_v"] = depth_v
+            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):
 
@@ -457,7 +461,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)
@@ -554,6 +558,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"]
@@ -656,18 +663,24 @@ class InitialConditions:
             If the field specified by `var_name` is 2D and both `eta` and `xi` are specified.
         """
 
-        if len(self.ds[var_name].squeeze().dims) == 3 and not any(
-            [s is not None, eta is not None, xi is not None]
-        ):
+        field = self.ds[var_name].squeeze()
+
+        if len(field.dims) == 3:
+            if not any([s is not None, eta is not None, xi is not None]):
+                raise ValueError(
+                    "Invalid input: For 3D fields, you must specify at least one of the dimensions 's', 'eta', or 'xi'."
+                )
+            if all([s is not None, eta is not None, xi is not None]):
+                raise ValueError(
+                    "Ambiguous input: For 3D fields, specify at most two of 's', 'eta', or 'xi'. Specifying all three is not allowed."
+                )
+
+        if len(field.dims) == 2 and all([eta is not None, xi is not None]):
             raise ValueError(
-                "For 3D fields, at least one of s, eta, or xi must be specified."
+                "Conflicting input: For 2D fields, specify only one dimension, either 'eta' or 'xi', not both."
             )
 
-        if len(self.ds[var_name].squeeze().dims) == 2 and all(
-            [eta is not None, xi is not None]
-        ):
-            raise ValueError("For 2D fields, specify either eta or xi, not both.")
-
+        # Load the data
         if self.use_dask:
             from dask.diagnostics import ProgressBar
 
@@ -675,54 +688,55 @@ class InitialConditions:
                 self.ds[var_name].load()
 
         field = self.ds[var_name].squeeze()
-        if s is not None:
-            layer_contours = False
 
+        # Get correct mask and horizontal coordinates
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
-            if layer_contours:
-                if "interface_depth_rho" in self.grid.ds:
-                    interface_depth = self.grid.ds.interface_depth_rho
-                else:
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=[]
-                    )
-            layer_depth = self.grid.ds.layer_depth_rho
-            mask = self.grid.ds.mask_rho
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_rho, "lat": self.grid.ds.lat_rho}
-            )
-
+            loc = "rho"
         elif all(dim in field.dims for dim in ["eta_rho", "xi_u"]):
-            if layer_contours:
-                if "interface_depth_u" in self.grid.ds:
-                    interface_depth = self.grid.ds.interface_depth_u
-                else:
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=["u", "v"]
-                    )
-            layer_depth = self.grid.ds.layer_depth_u
-            mask = self.grid.ds.mask_u
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_u, "lat": self.grid.ds.lat_u}
-            )
-
+            loc = "u"
         elif all(dim in field.dims for dim in ["eta_v", "xi_rho"]):
-            if layer_contours:
-                if "interface_depth_v" in self.grid.ds:
-                    interface_depth = self.grid.ds.interface_depth_v
-                else:
-                    self.get_vertical_coordinates(
-                        type="interface", additional_locations=["u", "v"]
-                    )
-            layer_depth = self.grid.ds.layer_depth_v
-            mask = self.grid.ds.mask_v
-            field = field.assign_coords(
-                {"lon": self.grid.ds.lon_v, "lat": self.grid.ds.lat_v}
-            )
+            loc = "v"
         else:
             ValueError("provided field does not have two horizontal dimension")
 
-        # slice the field as desired
+        mask = self.grid.ds[f"mask_{loc}"]
+        lat_deg = self.grid.ds[f"lat_{loc}"]
+        lon_deg = self.grid.ds[f"lon_{loc}"]
+
+        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})
+
+        # Retrieve depth coordinates
+        if s is not None:
+            layer_contours = False
+        # Note that `layer_depth_{loc}` has already been computed during `__post_init__`.
+        layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"].squeeze()
+
+        # Slice the field as desired
+        def _slice_and_assign(
+            field,
+            mask,
+            layer_depth,
+            title,
+            dim_name,
+            dim_values,
+            idx,
+        ):
+            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 "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, title
+
         title = field.long_name
         if s is not None:
             title = title + f", s_rho = {field.s_rho[s].item()}"
@@ -733,49 +747,28 @@ class InitialConditions:
             depth_contours = False
 
         if eta is not None:
-            if "eta_rho" in field.dims:
-                title = title + f", eta_rho = {field.eta_rho[eta].item()}"
-                field = field.isel(eta_rho=eta)
-                layer_depth = layer_depth.isel(eta_rho=eta)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(eta_rho=eta)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            elif "eta_v" in field.dims:
-                title = title + f", eta_v = {field.eta_v[eta].item()}"
-                field = field.isel(eta_v=eta)
-                layer_depth = layer_depth.isel(eta_v=eta)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(eta_v=eta)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{var_name}]."
-                )
+            field, mask, layer_depth, title = _slice_and_assign(
+                field,
+                mask,
+                layer_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,
+            )
+
         if xi is not None:
-            if "xi_rho" in field.dims:
-                title = title + f", xi_rho = {field.xi_rho[xi].item()}"
-                field = field.isel(xi_rho=xi)
-                layer_depth = layer_depth.isel(xi_rho=xi)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(xi_rho=xi)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            elif "xi_u" in field.dims:
-                title = title + f", xi_u = {field.xi_u[xi].item()}"
-                field = field.isel(xi_u=xi)
-                layer_depth = layer_depth.isel(xi_u=xi)
-                if layer_contours:
-                    interface_depth = interface_depth.isel(xi_u=xi)
-                if "s_rho" in field.dims:
-                    field = field.assign_coords({"layer_depth": layer_depth})
-            else:
-                raise ValueError(
-                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{var_name}]."
-                )
+            field, mask, layer_depth, title = _slice_and_assign(
+                field,
+                mask,
+                layer_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,
+            )
 
-        # chose colorbar
+        # Choose colorbar
         if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
@@ -792,26 +785,59 @@ class InitialConditions:
 
         if eta is None and xi is None:
             _plot(
-                self.grid.ds,
                 field=field.where(mask),
-                straddle=self.grid.straddle,
                 depth_contours=depth_contours,
                 title=title,
                 kwargs=kwargs,
                 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,
                     interface_depth=interface_depth,
@@ -821,40 +847,22 @@ class InitialConditions:
                 )
             else:
                 if "s_rho" in field.dims:
-                    _profile_plot(field, title=title, ax=ax)
+                    _profile_plot(field.where(mask), title=title, ax=ax)
                 else:
-                    _line_plot(field, 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:
+                    _line_plot(field.where(mask), title=title, ax=ax)
 
-          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
@@ -864,17 +872,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