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

roms_tools/setup/grid.py

@@ -24,7 +24,7 @@ from roms_tools.setup.utils import extract_single_value
 from pathlib import Path
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class Grid:
     """A single ROMS grid, used for creating, plotting, and then saving a new ROMS
     domain grid.
@@ -131,7 +131,7 @@ class Grid:
 
     def _input_checks(self):
         if self.topography_source is None:
-            object.__setattr__(self, "topography_source", {"name": "ETOPO5"})
+            self.topography_source = {"name": "ETOPO5"}
 
         if "name" not in self.topography_source:
             raise ValueError(
@@ -157,7 +157,7 @@ class Grid:
                 "========================================================================================================"
             )
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def update_topography(
         self, topography_source=None, hmin=None, verbose=False
@@ -223,9 +223,9 @@ class Grid:
             )
 
         # Update the grid's dataset and related attributes
-        object.__setattr__(self, "ds", ds)
-        object.__setattr__(self, "topography_source", topography_source)
-        object.__setattr__(self, "hmin", hmin)
+        self.ds = ds
+        self.topography_source = topography_source
+        self.hmin = hmin
 
     def update_vertical_coordinate(
         self, N=None, theta_s=None, theta_b=None, hc=None, verbose=False
@@ -321,11 +321,11 @@ class Grid:
                 "========================================================================================================"
             )
 
-        object.__setattr__(self, "ds", ds)
-        object.__setattr__(self, "theta_s", theta_s)
-        object.__setattr__(self, "theta_b", theta_b)
-        object.__setattr__(self, "hc", hc)
-        object.__setattr__(self, "N", N)
+        self.ds = ds
+        self.theta_s = theta_s
+        self.theta_b = theta_b
+        self.hc = hc
+        self.N = N
 
     def _straddle(self) -> None:
         """Check if the Greenwich meridian goes through the domain.
@@ -342,9 +342,9 @@ class Grid:
             np.abs(self.ds.lon_rho.diff("xi_rho")).max() > 300
             or np.abs(self.ds.lon_rho.diff("eta_rho")).max() > 300
         ):
-            object.__setattr__(self, "straddle", True)
+            self.straddle = True
         else:
-            object.__setattr__(self, "straddle", False)
+            self.straddle = False
 
     def _coarsen(self):
         """Update the grid by adding grid variables that are coarsened versions of the
@@ -391,7 +391,7 @@ class Grid:
             ds[coarse_var].attrs["long_name"] = f"{long_name} on coarsened grid"
             ds[coarse_var].attrs["units"] = ds[fine_var].attrs["units"]
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def plot(
         self, bathymetry: bool = True, title: str = None, with_dim_names: bool = False
@@ -461,7 +461,7 @@ class Grid:
         xi : int, optional
             The xi-index to plot. Default is None.
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
+            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots.
 
         Returns
         -------
@@ -581,14 +581,14 @@ class Grid:
         grid = cls.__new__(cls)
 
         # Set the dataset for the grid instance
-        object.__setattr__(grid, "ds", ds)
+        grid.ds = ds
 
         # Check if the Greenwich meridian goes through the domain.
         grid._straddle()
 
         if not all(coord in grid.ds for coord in ["lat_u", "lon_u", "lat_v", "lon_v"]):
             ds = _add_lat_lon_at_velocity_points(grid.ds, grid.straddle)
-            object.__setattr__(grid, "ds", ds)
+            grid.ds = ds
 
         # Coarsen the grid if necessary
         if not all(
@@ -606,7 +606,7 @@ class Grid:
         for var in ["lat_rho", "lon_rho", "lat_coarse", "lon_coarse"]:
             if var not in ds.coords:
                 ds = grid.ds.set_coords(var)
-                object.__setattr__(grid, "ds", ds)
+                grid.ds = ds
 
         # Update vertical coordinate if necessary
         if not all(var in grid.ds for var in ["Cs_r", "Cs_w"]):
@@ -620,14 +620,14 @@ class Grid:
                 N=N, theta_s=theta_s, theta_b=theta_b, hc=hc, verbose=True
             )
         else:
-            object.__setattr__(grid, "theta_s", ds.attrs["theta_s"].item())
-            object.__setattr__(grid, "theta_b", ds.attrs["theta_b"].item())
-            object.__setattr__(grid, "hc", ds.attrs["hc"].item())
-            object.__setattr__(grid, "N", len(ds.s_rho))
+            grid.theta_s = ds.attrs["theta_s"].item()
+            grid.theta_b = ds.attrs["theta_b"].item()
+            grid.hc = ds.attrs["hc"].item()
+            grid.N = len(ds.s_rho)
 
         # Manually set the remaining attributes by extracting parameters from dataset
-        object.__setattr__(grid, "nx", ds.sizes["xi_rho"] - 2)
-        object.__setattr__(grid, "ny", ds.sizes["eta_rho"] - 2)
+        grid.nx = ds.sizes["xi_rho"] - 2
+        grid.ny = ds.sizes["eta_rho"] - 2
         if "center_lon" in ds.attrs:
             center_lon = ds.attrs["center_lon"]
         elif "tra_lon" in ds:
@@ -637,7 +637,7 @@ class Grid:
                 "Missing grid information: 'center_lon' attribute or 'tra_lon' variable "
                 "must be present in the dataset."
             )
-        object.__setattr__(grid, "center_lon", center_lon)
+        grid.center_lon = center_lon
         if "center_lat" in ds.attrs:
             center_lat = ds.attrs["center_lat"]
         elif "tra_lat" in ds:
@@ -647,7 +647,7 @@ class Grid:
                 "Missing grid information: 'center_lat' attribute or 'tra_lat' variable "
                 "must be present in the dataset."
             )
-        object.__setattr__(grid, "center_lat", center_lat)
+        grid.center_lat = center_lat
         if "rot" in ds.attrs:
             rot = ds.attrs["rot"]
         elif "rotate" in ds:
@@ -657,7 +657,7 @@ class Grid:
                 "Missing grid information: 'rot' attribute or 'rotate' variable "
                 "must be present in the dataset."
             )
-        object.__setattr__(grid, "rot", rot)
+        grid.rot = rot
 
         for attr in [
             "size_x",
@@ -848,7 +848,7 @@ class Grid:
                 "========================================================================================================"
             )
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _add_global_metadata(self, ds):
         """Add global metadata and attributes to the dataset.

roms_tools/setup/initial_conditions.py

@@ -8,7 +8,7 @@ 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.regrid import LateralRegridToROMS, VerticalRegridToROMS
 from roms_tools.plot import _plot, _section_plot, _profile_plot, _line_plot
 from roms_tools.utils import (
     transpose_dimensions,
@@ -34,7 +34,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class InitialConditions:
     """Represents initial conditions for ROMS, including physical and biogeochemical
     data.
@@ -115,7 +115,7 @@ class InitialConditions:
 
         self._input_checks()
         # Dataset for depth coordinates
-        object.__setattr__(self, "ds_depth_coords", xr.Dataset())
+        self.ds_depth_coords = xr.Dataset()
 
         processed_fields = {}
         processed_fields = self._process_data(processed_fields, type="physics")
@@ -135,7 +135,7 @@ class InitialConditions:
         for var_name in ds.data_vars:
             ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _process_data(self, processed_fields, type="physics"):
 
@@ -161,7 +161,7 @@ class InitialConditions:
         var_names = variable_info.keys()
 
         # lateral regridding
-        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        lateral_regrid = LateralRegridToROMS(target_coords, data.dim_names)
 
         for var_name in var_names:
             if var_name in data.var_names.keys():
@@ -206,7 +206,7 @@ class InitialConditions:
         # Vertical regridding
         for location in ["rho", "u", "v"]:
             if len(var_names_dict[location]) > 0:
-                vertical_regrid = VerticalRegrid(
+                vertical_regrid = VerticalRegridToROMS(
                     self.ds_depth_coords[f"layer_depth_{location}"],
                     data.ds[data.dim_names["depth"]],
                 )
@@ -247,14 +247,10 @@ class InitialConditions:
         if "path" not in self.source.keys():
             raise ValueError("`source` must include a 'path'.")
         # set self.source["climatology"] to False if not provided
-        object.__setattr__(
-            self,
-            "source",
-            {
-                **self.source,
-                "climatology": self.source.get("climatology", False),
-            },
-        )
+        self.source = {
+            **self.source,
+            "climatology": self.source.get("climatology", False),
+        }
         if self.bgc_source is not None:
             if "name" not in self.bgc_source.keys():
                 raise ValueError(
@@ -265,14 +261,10 @@ class InitialConditions:
                     "`bgc_source` must include a 'path' if it is provided."
                 )
             # set self.bgc_source["climatology"] to False if not provided
-            object.__setattr__(
-                self,
-                "bgc_source",
-                {
-                    **self.bgc_source,
-                    "climatology": self.bgc_source.get("climatology", False),
-                },
-            )
+            self.bgc_source = {
+                **self.bgc_source,
+                "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:
@@ -655,7 +647,7 @@ class InitialConditions:
             visualize the layering of the water column. For clarity, the number of layer
             contours displayed is limited to a maximum of 10. Default is False.
         ax : matplotlib.axes.Axes, optional
-            The axes to plot on. If None, a new figure is created. Note that this argument does not work for horizontal plots that display the eta- and xi-dimensions at the same time.
+            The axes to plot on. If None, a new figure is created. Note that this argument does not work for 2D horizontal plots. Default is None.
 
         Returns
         -------
@@ -687,13 +679,6 @@ class InitialConditions:
                 "Conflicting input: For 2D fields, specify only one dimension, either 'eta' or 'xi', not both."
             )
 
-        # Load the data
-        if self.use_dask:
-            from dask.diagnostics import ProgressBar
-
-            with ProgressBar():
-                self.ds[var_name].load()
-
         field = self.ds[var_name].squeeze()
 
         # Get correct mask and horizontal coordinates
@@ -715,11 +700,18 @@ class InitialConditions:
 
         field = field.assign_coords({"lon": lon_deg, "lat": lat_deg})
 
+        # Load the data
+        if self.use_dask:
+            from dask.diagnostics import ProgressBar
+
+            with ProgressBar():
+                self.ds[var_name].load()
+
         # 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()
+        layer_depth = self.ds_depth_coords[f"layer_depth_{loc}"].squeeze().load()
 
         # Slice the field as desired
         def _slice_and_assign(

roms_tools/setup/nesting.py

@@ -20,7 +20,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class ChildGrid(Grid):
     """Represents a ROMS child grid that is compatible with the provided parent grid.
 
@@ -87,7 +87,7 @@ class ChildGrid(Grid):
             self.metadata["period"],
         )
 
-        object.__setattr__(self, "ds_nesting", ds_nesting)
+        self.ds_nesting = ds_nesting
 
     def _modify_child_topography_and_mask(self):
         """Adjust the child grid's topography and mask to align with the parent grid.
@@ -107,7 +107,7 @@ class ChildGrid(Grid):
             parent_grid_ds, child_grid_ds
         )
 
-        object.__setattr__(self, "ds", child_grid_ds)
+        self.ds = child_grid_ds
 
     def update_topography(
         self, topography_source=None, hmin=None, verbose=False

roms_tools/setup/river_forcing.py

@@ -23,7 +23,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class RiverForcing:
     """Represents river forcing input data for ROMS.
 
@@ -118,13 +118,13 @@ class RiverForcing:
                 raise ValueError(
                     "No relevant rivers found. Consider increasing domain size or using a different river dataset."
                 )
-            object.__setattr__(self, "original_indices", original_indices)
+            self.original_indices = original_indices
             updated_indices = self._move_rivers_to_closest_coast(target_coords, data)
-            object.__setattr__(self, "indices", updated_indices)
+            self.indices = updated_indices
 
         else:
             logging.info("Use provided river indices.")
-            object.__setattr__(self, "original_indices", self.indices)
+            self.original_indices = self.indices
             check_river_locations_are_along_coast(self.grid.ds.mask_rho, self.indices)
             data.extract_named_rivers(self.indices)
 
@@ -135,11 +135,11 @@ class RiverForcing:
         for var_name in ds.data_vars:
             ds[var_name] = substitute_nans_by_fillvalue(ds[var_name], fill_value=0.0)
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _input_checks(self):
         if self.source is None:
-            object.__setattr__(self, "source", {"name": "DAI"})
+            self.source = {"name": "DAI"}
 
         if "name" not in self.source:
             raise ValueError("`source` must include a 'name'.")
@@ -148,11 +148,10 @@ class RiverForcing:
                 raise ValueError("`source` must include a 'path'.")
 
         # Set 'climatology' to False if not provided in 'source'
-        object.__setattr__(
-            self,
-            "source",
-            {**self.source, "climatology": self.source.get("climatology", False)},
-        )
+        self.source = {
+            **self.source,
+            "climatology": self.source.get("climatology", False),
+        }
 
         # Check if 'indices' is provided and has the correct format
         if self.indices is not None:
@@ -255,23 +254,23 @@ class RiverForcing:
             - `river_tracer`: A `DataArray` representing tracer data for temperature, salinity and BGC tracers (if specified) for each river over time.
         """
         if self.source["climatology"]:
-            object.__setattr__(self, "climatology", True)
+            self.climatology = True
         else:
             if self.convert_to_climatology in ["never", "if_any_missing"]:
                 data_ds = data.select_relevant_times(data.ds)
                 if self.convert_to_climatology == "if_any_missing":
                     if data_ds[data.var_names["flux"]].isnull().any():
                         data.compute_climatology()
-                        object.__setattr__(self, "climatology", True)
+                        self.climatology = True
                     else:
-                        object.__setattr__(data, "ds", data_ds)
-                        object.__setattr__(self, "climatology", False)
+                        data.ds = data_ds
+                        self.climatology = False
                 else:
-                    object.__setattr__(data, "ds", data_ds)
-                    object.__setattr__(self, "climatology", False)
+                    data.ds = data_ds
+                    self.climatology = False
             elif self.convert_to_climatology == "always":
                 data.compute_climatology()
-                object.__setattr__(self, "climatology", True)
+                self.climatology = True
 
         ds = xr.Dataset()
 
@@ -442,21 +441,21 @@ class RiverForcing:
         return river_indices
 
     def _write_indices_into_dataset(self, ds):
-        """Adds river location indices to the dataset as the "river_location" variable.
+        """Adds river location indices to the dataset as the "river_flux" variable.
 
-        This method creates a new "river_location" variable
+        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.
 
         Parameters
         ----------
         ds : xarray.Dataset
-            The dataset to which the "river_location" variable will be added.
+            The dataset to which the "river_flux" variable will be added.
 
         Returns
         -------
         xarray.Dataset
-            The modified dataset with the "river_location" variable added.
+            The modified dataset with the "river_flux" variable added.
         """
 
         river_locations = xr.zeros_like(self.grid.ds.h)
@@ -475,7 +474,7 @@ class RiverForcing:
 
         river_locations.attrs["long_name"] = "River ID plus local volume fraction"
         river_locations.attrs["units"] = "none"
-        ds["river_location"] = river_locations
+        ds["river_flux"] = river_locations
 
         ds = ds.drop_vars(["lat_rho", "lon_rho"])
 

roms_tools/setup/surface_forcing.py

@@ -9,7 +9,7 @@ import logging
 from typing import Dict, Union, List, Optional
 from roms_tools import Grid
 from roms_tools.utils import save_datasets
-from roms_tools.regrid import LateralRegrid
+from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.plot import _plot
 from roms_tools.setup.datasets import (
     ERA5Dataset,
@@ -30,7 +30,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class SurfaceForcing:
     """Represents surface forcing input data for ROMS.
 
@@ -121,10 +121,10 @@ class SurfaceForcing:
             logging.info("Data will be interpolated onto grid coarsened by factor 2.")
         else:
             logging.info("Data will be interpolated onto fine grid.")
-        object.__setattr__(self, "use_coarse_grid", use_coarse_grid)
+        self.use_coarse_grid = use_coarse_grid
 
         target_coords = get_target_coords(self.grid, self.use_coarse_grid)
-        object.__setattr__(self, "target_coords", target_coords)
+        self.target_coords = target_coords
 
         data.choose_subdomain(
             target_coords,
@@ -140,7 +140,7 @@ class SurfaceForcing:
 
         processed_fields = {}
         # lateral regridding
-        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        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(
@@ -171,7 +171,7 @@ class SurfaceForcing:
         for var_name in ds.data_vars:
             ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _input_checks(self):
         # Check that start_time and end_time are both None or none of them is
@@ -197,11 +197,10 @@ class SurfaceForcing:
             raise ValueError("`source` must include a 'path'.")
 
         # Set 'climatology' to False if not provided in 'source'
-        object.__setattr__(
-            self,
-            "source",
-            {**self.source, "climatology": self.source.get("climatology", False)},
-        )
+        self.source = {
+            **self.source,
+            "climatology": self.source.get("climatology", False),
+        }
 
         # Validate 'coarse_grid_mode'
         valid_modes = ["auto", "always", "never"]
@@ -339,7 +338,7 @@ class SurfaceForcing:
                 else:
                     variable_info[var_name] = {**default_info, "validate": False}
 
-        object.__setattr__(self, "variable_info", variable_info)
+        self.variable_info = variable_info
 
     def _apply_correction(self, processed_fields, data):
 
@@ -382,7 +381,9 @@ class SurfaceForcing:
             )
 
         # Spatial regridding
-        lateral_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
+        lateral_regrid = LateralRegridToROMS(
+            self.target_coords, correction_data.dim_names
+        )
         corr_factor = lateral_regrid.apply(corr_factor)
 
         processed_fields["swrad"] = processed_fields["swrad"] * corr_factor

roms_tools/setup/tides.py

@@ -8,7 +8,7 @@ from pathlib import Path
 from dataclasses import dataclass, field
 from roms_tools import Grid
 from roms_tools.plot import _plot
-from roms_tools.regrid import LateralRegrid
+from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.utils import save_datasets
 from roms_tools.setup.datasets import TPXODataset
 from roms_tools.setup.utils import (
@@ -25,7 +25,7 @@ from roms_tools.setup.utils import (
 )
 
 
-@dataclass(frozen=True, kw_only=True)
+@dataclass(kw_only=True)
 class TidalForcing:
     """Represents tidal forcing for ROMS.
 
@@ -88,7 +88,7 @@ class TidalForcing:
         data.convert_to_float64()
 
         # select desired number of constituents
-        object.__setattr__(data, "ds", data.ds.isel(ntides=slice(None, self.ntides)))
+        data.ds = data.ds.isel(ntides=slice(None, self.ntides))
         self._correct_tides(data)
 
         data.apply_lateral_fill()
@@ -98,7 +98,7 @@ class TidalForcing:
 
         processed_fields = {}
         # lateral regridding
-        lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+        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(
@@ -144,7 +144,7 @@ class TidalForcing:
         for var_name in ds.data_vars:
             ds[var_name] = substitute_nans_by_fillvalue(ds[var_name])
 
-        object.__setattr__(self, "ds", ds)
+        self.ds = ds
 
     def _input_checks(self):
 
@@ -218,7 +218,7 @@ class TidalForcing:
             },
         }
 
-        object.__setattr__(self, "variable_info", variable_info)
+        self.variable_info = variable_info
 
     def _write_into_dataset(self, processed_fields, d_meta):
 
@@ -504,7 +504,7 @@ class TidalForcing:
         var_names.pop("sal_Re", None)  # Remove "sal_Re" if it exists
         var_names.pop("sal_Im", None)  # Remove "sal_Im" if it exists
 
-        object.__setattr__(data, "var_names", var_names)
+        data.var_names = var_names
 
 
 def modified_julian_days(year, month, day, hour=0):

roms_tools/setup/topography.py

@@ -6,7 +6,7 @@ import gcm_filters
 from roms_tools.setup.utils import handle_boundaries
 import warnings
 from itertools import count
-from roms_tools.regrid import LateralRegrid
+from roms_tools.regrid import LateralRegridToROMS
 from roms_tools.setup.datasets import ETOPO5Dataset, SRTM15Dataset
 
 
@@ -150,7 +150,7 @@ def _make_raw_topography(
 
     if verbose:
         start_time = time.time()
-    lateral_regrid = LateralRegrid(target_coords, data.dim_names)
+    lateral_regrid = LateralRegridToROMS(target_coords, data.dim_names)
     hraw = lateral_regrid.apply(data.ds[data.var_names["topo"]], method=method)
     if verbose:
         logging.info(