roms-tools 3.3.0__py3-none-any.whl → 3.5.0__py3-none-any.whl

roms_tools/{setup → datasets}/lat_lon_datasets.py

@@ -16,22 +16,32 @@ import numpy as np
 import xarray as xr
 
 from roms_tools.constants import R_EARTH
-from roms_tools.download import (
+from roms_tools.datasets.download import (
     download_correction_data,
     download_sal_data,
     download_topo,
 )
-from roms_tools.setup.fill import LateralFill
+from roms_tools.datasets.utils import (
+    check_dataset,
+    convert_to_float64,
+    extrapolate_deepest_to_bottom,
+    select_relevant_fields,
+    select_relevant_times,
+    validate_start_end_time,
+)
+from roms_tools.fill import LateralFill
 from roms_tools.setup.utils import (
     Timed,
     assign_dates_to_climatology,
-    check_dataset,
     get_target_coords,
+)
+from roms_tools.utils import (
+    get_dask_chunks,
+    get_pkg_error_msg,
+    has_gcsfs,
     interpolate_cyclic_time,
-    one_dim_fill,
-    select_relevant_times,
+    load_data,
 )
-from roms_tools.utils import get_dask_chunks, get_pkg_error_msg, has_gcsfs, load_data
 
 TConcatEndTypes = Literal["lower", "upper", "both"]
 REPO_ROOT = Path(__file__).resolve().parents[2]
@@ -121,11 +131,11 @@ class LatLonDataset:
     var_names: dict[str, str]
     opt_var_names: dict[str, str] = field(default_factory=dict)
     climatology: bool = False
-    needs_lateral_fill: bool | None = True
+    needs_lateral_fill: bool = True
     use_dask: bool = False
     read_zarr: bool = False
     allow_flex_time: bool = False
-    apply_post_processing: bool | None = True
+    apply_post_processing: bool = True
 
     ds_loader_fn: Callable[[], xr.Dataset] | None = None
     is_global: bool = field(init=False, repr=False)
@@ -140,21 +150,11 @@ class LatLonDataset:
         4. Ensures latitude, longitude, and depth values are in ascending order.
         5. Checks if the dataset covers the entire globe and adjusts if necessary.
         """
-        # Validate start_time and end_time
-        if self.start_time is not None and not isinstance(self.start_time, datetime):
-            raise TypeError(
-                f"start_time must be a datetime object, but got {type(self.start_time).__name__}."
-            )
-        if self.end_time is not None and not isinstance(self.end_time, datetime):
-            raise TypeError(
-                f"end_time must be a datetime object, but got {type(self.end_time).__name__}."
-            )
-
+        validate_start_end_time(self.start_time, self.end_time)
         ds = self.load_data()
         ds = self.clean_up(ds)
-        self.check_dataset(ds)
+        check_dataset(ds, self.dim_names, self.var_names, self.opt_var_names)
 
-        # Select relevant fields
         ds = self.select_relevant_fields(ds)
 
         # Select relevant times
@@ -226,44 +226,30 @@ class LatLonDataset:
         """
         return ds  # Default behavior (no-op, subclasses should override)
 
-    def check_dataset(self, ds: xr.Dataset) -> None:
-        """Check if the dataset contains the specified variables and dimensions.
-
-        Parameters
-        ----------
-        ds : xr.Dataset
-            The xarray Dataset to check.
-
-        Raises
-        ------
-        ValueError
-            If the dataset does not contain the specified variables or dimensions.
+    def select_relevant_fields(self, ds: xr.Dataset) -> xr.Dataset:
         """
-        check_dataset(ds, self.dim_names, self.var_names)
+        Return a subset of the dataset containing only the required and optional
+        variables defined for this object.
 
-    def select_relevant_fields(self, ds: xr.Dataset) -> xr.Dataset:
-        """Selects and returns a subset of the dataset containing only the variables
-        specified in `self.var_names`.
+        Variables retained are those listed in ``self.var_names`` and
+        ``self.opt_var_names``. Any other data variables are removed, except for
+        the special variable ``"mask"``, which is always preserved if present.
 
         Parameters
         ----------
         ds : xr.Dataset
-            The input dataset from which variables will be selected.
+            The input dataset from which relevant variables will be selected.
 
         Returns
         -------
         xr.Dataset
-            A dataset containing only the variables specified in `self.var_names`.
+            A new dataset containing only the required variables specified in
+            ``self.var_names`` and the optional variables specified in
+            ``self.opt_var_names``, along with ``"mask"`` if present.
         """
-        for var in ds.data_vars:
-            if (
-                var not in self.var_names.values()
-                and var not in self.opt_var_names.values()
-                and var != "mask"
-            ):
-                ds = ds.drop_vars(var)
-
-        return ds
+        return select_relevant_fields(
+            ds, [*self.var_names.values(), *self.opt_var_names.values()]
+        )
 
     def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
         """Dummy method to be overridden by child classes to add time information to the
@@ -336,12 +322,13 @@ class LatLonDataset:
             raise ValueError("select_relevant_times called but start_time is None.")
 
         ds = select_relevant_times(
-            ds,
-            time_dim,
-            self.start_time,
-            self.end_time,
-            self.climatology,
-            self.allow_flex_time,
+            ds=ds,
+            time_dim=time_dim,
+            time_coord=time_dim,
+            start_time=self.start_time,
+            end_time=self.end_time,
+            climatology=self.climatology,
+            allow_flex_time=self.allow_flex_time,
         )
 
         return ds
@@ -543,7 +530,7 @@ class LatLonDataset:
         None
             This method modifies the dataset in place and does not return anything.
         """
-        ds = self.ds.astype({var: "float64" for var in self.ds.data_vars})
+        ds = convert_to_float64(self.ds)
         self.ds = ds
 
         return None
@@ -635,7 +622,7 @@ class LatLonDataset:
         if self.needs_lateral_fill:
             lateral_fill = LateralFill(
                 self.ds["mask"],
-                [self.dim_names["latitude"], self.dim_names["longitude"]],
+                (self.dim_names["latitude"], self.dim_names["longitude"]),
             )
 
             separate_fill_for_velocities = False
@@ -644,7 +631,7 @@ class LatLonDataset:
             if "mask_vel" in self.ds.data_vars:
                 lateral_fill_vel = LateralFill(
                     self.ds["mask_vel"],
-                    [self.dim_names["latitude"], self.dim_names["longitude"]],
+                    (self.dim_names["latitude"], self.dim_names["longitude"]),
                 )
                 separate_fill_for_velocities = True
 
@@ -672,11 +659,19 @@ class LatLonDataset:
         propagating the deepest available data downward.
         """
         if "depth" in self.dim_names:
-            for var_name in self.ds.data_vars:
-                if self.dim_names["depth"] in self.ds[var_name].dims:
-                    self.ds[var_name] = one_dim_fill(
-                        self.ds[var_name], self.dim_names["depth"], direction="forward"
-                    )
+            self.ds = extrapolate_deepest_to_bottom(self.ds, self.dim_names["depth"])
+
+    def rotate_velocities_to_east_and_north(self) -> None:
+        """
+        Rotate velocity components to east/north directions.
+
+        For lat-lon datasets, velocity components are already defined in
+        earth-relative east/north coordinates. Therefore, no rotation is
+        required and this method is a no-op.
+        This method is provided for API compatibility with ROMSDataset,
+        where an explicit rotation using the grid angle is necessary.
+        """
+        return None
 
     @classmethod
     def from_ds(cls, original_dataset: LatLonDataset, ds: xr.Dataset) -> LatLonDataset:
@@ -1106,7 +1101,7 @@ class UnifiedDataset(LatLonDataset):
     attribute is set to `False`.
     """
 
-    needs_lateral_fill: bool | None = False
+    needs_lateral_fill: bool = False
 
     # overwrite clean_up method from parent class
     def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
@@ -1147,7 +1142,8 @@ class UnifiedDataset(LatLonDataset):
                         if var_name in ds.data_vars and "season" in ds[var_name].dims:
                             ds[var_name] = interpolate_cyclic_time(
                                 ds[var_name],
-                                time_dim_name="season",
+                                time_dim="season",
+                                time_coord="season",
                                 day_of_year=ds["month"],
                             )
 
@@ -1671,6 +1667,7 @@ class ETOPO5Dataset(LatLonDataset):
     dim_names: dict[str, str] = field(
         default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
     )
+    needs_lateral_fill: bool = False
 
     def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
         """Assign lat and lon as coordinates.
@@ -1706,6 +1703,32 @@ class SRTM15Dataset(LatLonDataset):
     dim_names: dict[str, str] = field(
         default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
     )
+    needs_lateral_fill: bool = False
+
+
+@dataclass(kw_only=True)
+class EMODDataset(LatLonDataset):
+    """Represents topography data on the original grid from the EMOD dataset."""
+
+    var_names: dict[str, str] = field(
+        default_factory=lambda: {
+            "topo": "elevation",
+        }
+    )
+    dim_names: dict[str, str] = field(
+        default_factory=lambda: {"longitude": "lon", "latitude": "lat"}
+    )
+    needs_lateral_fill: bool = True
+
+    def post_process(self) -> None:
+        """Assign land mask."""
+        mask = xr.where(
+            self.ds[self.var_names["topo"]].isnull(),
+            0,
+            1,
+        )
+
+        self.ds["mask"] = mask
 
 
 @dataclass
@@ -2331,7 +2354,7 @@ def choose_subdomain(
     resolution: float,
     is_global: bool,
     target_coords: Mapping[str, Any],
-    buffer_points: int = 20,
+    buffer_points: int = DEFAULT_NR_BUFFER_POINTS,
     use_dask: bool = False,
 ) -> xr.Dataset:
     """
@@ -2433,6 +2456,7 @@ def choose_subdomain(
                 if lon_min - margin < 0:
                     lon_min += 360
                     lon_max += 360
+
     # Select the subdomain in longitude direction
     subdomain = subdomain.sel(
         **{

roms_tools/{setup → datasets}/river_datasets.py

@@ -7,12 +7,11 @@ from pathlib import Path
 import numpy as np
 import xarray as xr
 
-from roms_tools.download import download_river_data
+from roms_tools.datasets.download import download_river_data
+from roms_tools.datasets.utils import check_dataset, select_relevant_times
 from roms_tools.setup.utils import (
     assign_dates_to_climatology,
-    check_dataset,
     gc_dist,
-    select_relevant_times,
 )
 from roms_tools.utils import load_data
 
@@ -198,7 +197,13 @@ class RiverDataset:
         """
         time_dim = self.dim_names["time"]
 
-        ds = select_relevant_times(ds, time_dim, self.start_time, self.end_time, False)
+        ds = select_relevant_times(
+            ds=ds,
+            time_dim=time_dim,
+            time_coord=time_dim,
+            start_time=self.start_time,
+            end_time=self.end_time,
+        )
 
         return ds