roms-tools 1.5.0__py3-none-any.whl → 1.6.1__py3-none-any.whl

roms_tools/setup/initial_conditions.py

@@ -15,11 +15,9 @@ from roms_tools.setup.utils import (
     get_target_coords,
     rotate_velocities,
     compute_barotropic_velocity,
-    _extrapolate_deepest_to_bottom,
     transpose_dimensions,
 )
-from roms_tools.setup.fill import _lateral_fill
-from roms_tools.setup.regrid import _lateral_regrid, _vertical_regrid
+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
@@ -92,31 +90,31 @@ class InitialConditions:
 
         self._input_checks()
 
-        data_vars = {}
-        data_vars = self._process_data(data_vars, type="physics")
+        processed_fields = {}
+        processed_fields = self._process_data(processed_fields, type="physics")
 
         if self.bgc_source is not None:
-            data_vars = self._process_data(data_vars, type="bgc")
+            processed_fields = self._process_data(processed_fields, type="bgc")
 
-        for var in data_vars.keys():
-            data_vars[var] = transpose_dimensions(data_vars[var])
+        for var_name in processed_fields.keys():
+            processed_fields[var_name] = transpose_dimensions(
+                processed_fields[var_name]
+            )
 
         d_meta = get_variable_metadata()
-        ds = self._write_into_dataset(data_vars, d_meta)
+        ds = self._write_into_dataset(processed_fields, d_meta)
 
         ds = self._add_global_metadata(ds)
 
-        ds["zeta"].load()
-        # NaN values at wet points indicate that the raw data did not cover the domain, and the following will raise a ValueError
-        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+        self._validate(ds)
 
         # substitute NaNs over land by a fill value to avoid blow-up of ROMS
-        for var in ds.data_vars:
-            ds[var] = substitute_nans_by_fillvalue(ds[var])
+        for var_name in ds.data_vars:
+            ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
         object.__setattr__(self, "ds", ds)
 
-    def _process_data(self, data_vars, type="physics"):
+    def _process_data(self, processed_fields, type="physics"):
 
         target_coords = get_target_coords(self.grid)
 
@@ -126,35 +124,29 @@ class InitialConditions:
             data = self._get_bgc_data()
 
         data.choose_subdomain(
-            latitude_range=[
-                target_coords["lat"].min().values,
-                target_coords["lat"].max().values,
-            ],
-            longitude_range=[
-                target_coords["lon"].min().values,
-                target_coords["lon"].max().values,
-            ],
-            margin=2,
-            straddle=target_coords["straddle"],
+            target_coords,
+            buffer_points=20,  # lateral fill needs good buffer from data margin
         )
 
-        variable_info = self._set_variable_info(data, type=type)
-
-        data_vars = _extrapolate_deepest_to_bottom(data_vars, data)
+        data.extrapolate_deepest_to_bottom()
+        data.apply_lateral_fill()
 
-        data_vars = _lateral_fill(data_vars, data)
+        variable_info = self._set_variable_info(data, type=type)
+        var_names = variable_info.keys()
 
         # lateral regridding
-        var_names = variable_info.keys()
-        data_vars = _lateral_regrid(
-            data, target_coords["lon"], target_coords["lat"], data_vars, var_names
-        )
+        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(
+                    data.ds[data.var_names[var_name]]
+                )
 
         # rotation of velocities and interpolation to u/v points
         if "u" in variable_info and "v" in variable_info:
-            (data_vars["u"], data_vars["v"],) = rotate_velocities(
-                data_vars["u"],
-                data_vars["v"],
+            (processed_fields["u"], processed_fields["v"],) = rotate_velocities(
+                processed_fields["u"],
+                processed_fields["v"],
                 target_coords["angle"],
                 interpolate=True,
             )
@@ -167,28 +159,32 @@ class InitialConditions:
                 if info["location"] == location and info["is_3d"]
             ]
             if len(var_names) > 0:
-                data_vars = _vertical_regrid(
-                    data,
+                vertical_regrid = VerticalRegrid(
                     self.grid.ds[f"layer_depth_{location}"],
-                    data_vars,
-                    var_names,
+                    data.ds[data.dim_names["depth"]],
                 )
+                for var_name in var_names:
+                    if var_name in processed_fields:
+                        processed_fields[var_name] = vertical_regrid.apply(
+                            processed_fields[var_name]
+                        )
 
         # compute barotropic velocities
         if "u" in variable_info and "v" in variable_info:
-            for var in ["u", "v"]:
-                data_vars[f"{var}bar"] = compute_barotropic_velocity(
-                    data_vars[var], self.grid.ds[f"interface_depth_{var}"]
+            for var_name in ["u", "v"]:
+                processed_fields[f"{var_name}bar"] = compute_barotropic_velocity(
+                    processed_fields[var_name],
+                    self.grid.ds[f"interface_depth_{var_name}"],
                 )
 
         if type == "bgc":
             # Ensure time coordinate matches that of physical variables
-            for var in variable_info.keys():
-                data_vars[var] = data_vars[var].assign_coords(
-                    {"time": data_vars["temp"]["time"]}
+            for var_name in variable_info.keys():
+                processed_fields[var_name] = processed_fields[var_name].assign_coords(
+                    {"time": processed_fields["temp"]["time"]}
                 )
 
-        return data_vars
+        return processed_fields
 
     def _input_checks(self):
 
@@ -314,24 +310,26 @@ class InitialConditions:
             }
         elif type == "bgc":
             variable_info = {}
-            for var in data.var_names.keys():
-                variable_info[var] = default_info
+            for var_name in data.var_names.keys():
+                variable_info[var_name] = default_info
 
         return variable_info
 
-    def _write_into_dataset(self, data_vars, d_meta):
+    def _write_into_dataset(self, processed_fields, d_meta):
 
         # save in new dataset
         ds = xr.Dataset()
 
-        for var in data_vars.keys():
-            ds[var] = data_vars[var].astype(np.float32)
-            ds[var].attrs["long_name"] = d_meta[var]["long_name"]
-            ds[var].attrs["units"] = d_meta[var]["units"]
+        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"]
 
         # initialize vertical velocity to zero
         ds["w"] = xr.zeros_like(
-            self.grid.ds["interface_depth_rho"].expand_dims(time=data_vars["u"].time)
+            self.grid.ds["interface_depth_rho"].expand_dims(
+                time=processed_fields["u"].time
+            )
         ).astype(np.float32)
         ds["w"].attrs["long_name"] = d_meta["w"]["long_name"]
         ds["w"].attrs["units"] = d_meta["w"]["units"]
@@ -351,7 +349,7 @@ class InitialConditions:
             "layer_depth_v",
             "interface_depth_v",
         ]
-        existing_vars = [var for var in variables_to_drop if var in ds]
+        existing_vars = [var_name for var_name in variables_to_drop if var_name in ds]
         ds = ds.drop_vars(existing_vars)
 
         ds["Cs_r"] = self.grid.ds["Cs_r"]
@@ -375,6 +373,29 @@ class InitialConditions:
 
         return ds
 
+    def _validate(self, ds):
+        """Validates the dataset by checking for NaN values in SSH at wet points, which
+        would indicate missing raw data coverage over the target domain.
+
+        Parameters
+        ----------
+        ds : xarray.Dataset
+            The dataset to validate.
+
+        Raises
+        ------
+        ValueError
+            If NaN values are found in any of the specified variables at wet points,
+            indicating incomplete data coverage.
+
+        Notes
+        -----
+        This check is only applied to the 2D variable SSH to improve performance.
+        """
+
+        ds["zeta"].load()
+        nan_check(ds["zeta"].squeeze(), self.grid.ds.mask_rho)
+
     def _add_global_metadata(self, ds):
 
         ds.attrs["title"] = "ROMS initial conditions file created by ROMS-Tools"
@@ -398,7 +419,7 @@ class InitialConditions:
 
     def plot(
         self,
-        varname,
+        var_name,
         s=None,
         eta=None,
         xi=None,
@@ -409,7 +430,7 @@ class InitialConditions:
 
         Parameters
         ----------
-        varname : str
+        var_name : str
             The name of the initial conditions field to plot. Options include:
 
             - "temp": Potential temperature.
@@ -480,25 +501,30 @@ class InitialConditions:
         Raises
         ------
         ValueError
-            If the specified `varname` is not one of the valid options.
-            If the field specified by `varname` is 3D and none of `s`, `eta`, or `xi` are specified.
-            If the field specified by `varname` is 2D and both `eta` and `xi` are specified.
+            If the specified `var_name` is not one of the valid options.
+            If the field specified by `var_name` is 3D and none of `s`, `eta`, or `xi` are specified.
+            If the field specified by `var_name` is 2D and both `eta` and `xi` are specified.
         """
 
-        if len(self.ds[varname].squeeze().dims) == 3 and not any(
+        if len(self.ds[var_name].squeeze().dims) == 3 and not any(
             [s is not None, eta is not None, xi is not None]
         ):
             raise ValueError(
                 "For 3D fields, at least one of s, eta, or xi must be specified."
             )
 
-        if len(self.ds[varname].squeeze().dims) == 2 and all(
+        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.")
 
-        self.ds[varname].load()
-        field = self.ds[varname].squeeze()
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds[var_name].load()
+
+        field = self.ds[var_name].squeeze()
 
         if all(dim in field.dims for dim in ["eta_rho", "xi_rho"]):
             interface_depth = self.grid.ds.interface_depth_rho
@@ -553,7 +579,7 @@ class InitialConditions:
                     field = field.assign_coords({"layer_depth": layer_depth})
             else:
                 raise ValueError(
-                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{varname}]."
+                    f"None of the expected dimensions (eta_rho, eta_v) found in ds[{var_name}]."
                 )
         if xi is not None:
             if "xi_rho" in field.dims:
@@ -572,18 +598,18 @@ class InitialConditions:
                     field = field.assign_coords({"layer_depth": layer_depth})
             else:
                 raise ValueError(
-                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{varname}]."
+                    f"None of the expected dimensions (xi_rho, xi_u) found in ds[{var_name}]."
                 )
 
         # chose colorbar
-        if varname in ["u", "v", "w", "ubar", "vbar", "zeta"]:
+        if var_name in ["u", "v", "w", "ubar", "vbar", "zeta"]:
             vmax = max(field.max().values, -field.min().values)
             vmin = -vmax
             cmap = plt.colormaps.get_cmap("RdBu_r")
         else:
             vmax = field.max().values
             vmin = field.min().values
-            if varname in ["temp", "salt"]:
+            if var_name in ["temp", "salt"]:
                 cmap = plt.colormaps.get_cmap("YlOrRd")
             else:
                 cmap = plt.colormaps.get_cmap("YlGn")
@@ -660,7 +686,13 @@ class InitialConditions:
         if filepath.suffix == ".nc":
             filepath = filepath.with_suffix("")
 
-        dataset_list = [self.ds.load()]
+        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(
@@ -698,15 +730,18 @@ class InitialConditions:
 
         initial_conditions_data = {
             "InitialConditions": {
-                "source": self.source,
                 "ini_time": self.ini_time.isoformat(),
-                "model_reference_date": self.model_reference_date.isoformat(),
+                "source": self.source,
             }
         }
         # Include bgc_source if it's not None
         if self.bgc_source is not None:
             initial_conditions_data["InitialConditions"]["bgc_source"] = self.bgc_source
 
+        initial_conditions_data["InitialConditions"][
+            "model_reference_date"
+        ] = self.model_reference_date.isoformat()
+
         yaml_data = {
             **grid_yaml_data,
             **initial_conditions_data,
@@ -716,7 +751,7 @@ class InitialConditions:
             # Write header
             file.write(header)
             # Write YAML data
-            yaml.dump(yaml_data, file, default_flow_style=False)
+            yaml.dump(yaml_data, file, default_flow_style=False, sort_keys=False)
 
     @classmethod
     def from_yaml(

roms_tools/setup/regrid.py

@@ -2,45 +2,38 @@ import xarray as xr
 
 
 class LateralRegrid:
-    """Applies lateral fill and regridding to data.
+    """Handles lateral regridding of data onto a new spatial grid."""
 
-    This class fills missing values in ocean data and interpolates it onto a new grid
-    defined by the provided longitude and latitude.
-
-    Parameters
-    ----------
-    data : DataContainer
-        Container with variables to be interpolated, including a `mask` and dimension names.
-    lon : xarray.DataArray
-        Target longitude coordinates.
-    lat : xarray.DataArray
-        Target latitude coordinates.
-    """
-
-    def __init__(self, data, lon, lat):
-        """Initializes the lateral fill and target grid coordinates.
+    def __init__(self, target_coords, source_dim_names):
+        """Initialize target grid coordinates and names for grid dimensions.
 
         Parameters
         ----------
-        data : DataContainer
-            Data with dimensions and mask for filling.
-        lon : xarray.DataArray
-            Longitude for new grid.
-        lat : xarray.DataArray
-            Latitude for new grid.
+        target_coords : dict
+            Dictionary containing 'lon' and 'lat' as xarray.DataArrays representing
+            the longitude and latitude values of the target grid.
+        source_dim_names : dict
+            Dictionary specifying names for the latitude and longitude dimensions,
+            typically using keys like "latitude" and "longitude" to align with the dataset conventions.
+
+        Attributes
+        ----------
+        coords : dict
+            Maps the dimension names to the corresponding latitude and longitude
+            DataArrays, providing easy access to target grid coordinates.
         """
 
         self.coords = {
-            data.dim_names["latitude"]: lat,
-            data.dim_names["longitude"]: lon,
+            source_dim_names["latitude"]: target_coords["lat"],
+            source_dim_names["longitude"]: target_coords["lon"],
         }
 
-    def apply(self, var):
+    def apply(self, da):
         """Fills missing values and regrids the variable.
 
         Parameters
         ----------
-        var : xarray.DataArray
+        da : xarray.DataArray
             Input data to fill and regrid.
 
         Returns
@@ -48,38 +41,48 @@ class LateralRegrid:
         xarray.DataArray
             Regridded data with filled values.
         """
-        regridded = var.interp(self.coords, method="linear").drop_vars(
+        regridded = da.interp(self.coords, method="linear").drop_vars(
             list(self.coords.keys())
         )
         return regridded
 
 
 class VerticalRegrid:
-    """Performs vertical interpolation of data onto new depth coordinates.
+    """Interpolates data onto new vertical (depth) coordinates.
 
     Parameters
     ----------
-    data : DataContainer
-        Container holding the data to be regridded, with relevant dimension names.
-    target_depth : xarray.DataArray
-        Target depth coordinates for interpolation.
+    target_depth_coords : xarray.DataArray
+        Depth coordinates for the target grid.
+    source_depth_coords : xarray.DataArray
+        Depth coordinates for the source grid.
     """
 
-    def __init__(self, data, target_depth):
-        """Initializes vertical regridding with specified depth coordinates.
+    def __init__(self, target_depth_coords, source_depth_coords):
+        """Initialize regridding factors for interpolation.
 
         Parameters
         ----------
-        data : DataContainer
-            Container holding the data to be regridded, with relevant dimension names.
-        target_depth : xarray.DataArray
-            Target depth coordinates for interpolation.
+        target_depth_coords : xarray.DataArray
+            Depth coordinates for the target grid.
+        source_depth_coords : xarray.DataArray
+            Depth coordinates for the source grid.
+
+        Attributes
+        ----------
+        coeff : xarray.Dataset
+            Dataset containing:
+            - `is_below` : Boolean mask for depths just below target.
+            - `is_above` : Boolean mask for depths just above target.
+            - `upper_mask`, `lower_mask` : Masks for valid interpolation bounds.
+            - `factor` : Weight for blending values between levels.
         """
 
-        self.depth_dim = data.dim_names["depth"]
+        self.depth_dim = source_depth_coords.dims[0]
+        source_depth = source_depth_coords
         dims = {"dim": self.depth_dim}
 
-        dlev = data.ds[data.dim_names["depth"]] - target_depth
+        dlev = source_depth - target_depth_coords
         is_below = dlev == dlev.where(dlev >= 0).min(**dims)
         is_above = dlev == dlev.where(dlev <= 0).max(**dims)
         p_below = dlev.where(is_below).sum(**dims)
@@ -138,61 +141,3 @@ class VerticalRegrid:
             )
 
         return result
-
-
-def _lateral_regrid(data, lon, lat, data_vars, var_names):
-    """Laterally regrid specified variables onto new latitude and longitude coordinates.
-
-    Parameters
-    ----------
-    data : Dataset
-        Input data containing the information about source dimensions.
-    lon : xarray.DataArray
-        Target longitude coordinates.
-    lat : xarray.DataArray
-        Target latitude coordinates.
-    data_vars : dict of str : xarray.DataArray
-        Dictionary of variables to regrid.
-    var_names : list of str
-        Names of variables to regrid.
-
-    Returns
-    -------
-    dict of str : xarray.DataArray
-        Updated data_vars with regridded variables.
-    """
-    lateral_regrid = LateralRegrid(data, lon, lat)
-
-    for var_name in var_names:
-        if var_name in data_vars:
-            data_vars[var_name] = lateral_regrid.apply(data_vars[var_name])
-
-    return data_vars
-
-
-def _vertical_regrid(data, target_depth, data_vars, var_names):
-    """Vertically regrid specified variables onto new depth coordinates.
-
-    Parameters
-    ----------
-    data : Dataset
-        Input dataset containing the variables and source depth information.
-    target_depth : xarray.DataArray
-        Target depth coordinates for regridding.
-    data_vars : dict of str : xarray.DataArray
-        Dictionary of variables to be regridded.
-    var_names : list of str
-        Names of variables to regrid.
-
-    Returns
-    -------
-    dict of str : xarray.DataArray
-        Updated data_vars with variables regridded onto the target depth coordinates.
-    """
-    vertical_regrid = VerticalRegrid(data, target_depth)
-
-    for var_name in var_names:
-        if var_name in data_vars:
-            data_vars[var_name] = vertical_regrid.apply(data_vars[var_name])
-
-    return data_vars