roms-tools 1.1.0__py3-none-any.whl → 1.3.0__py3-none-any.whl

roms_tools/setup/datasets.py

@@ -1,10 +1,11 @@
+import re
 import xarray as xr
 from dataclasses import dataclass, field
 import glob
 from datetime import datetime, timedelta
 import numpy as np
-from typing import Dict, Optional
-import dask
+from typing import Dict, Optional, Union, List
+from pathlib import Path
 import warnings
 from roms_tools.setup.utils import (
     assign_dates_to_climatology,
@@ -22,8 +23,9 @@ class Dataset:
 
     Parameters
     ----------
-    filename : str
-        The path to the data files. Can contain wildcards.
+    filename : Union[str, Path, List[Union[str, Path]]]
+        The path to the data file(s). Can be a single string (with or without wildcards), a single Path object,
+        or a list of strings or Path objects containing multiple files.
     start_time : Optional[datetime], optional
         The start time for selecting relevant data. If not provided, the data is not filtered by start time.
     end_time : Optional[datetime], optional
@@ -35,6 +37,9 @@ class Dataset:
         Dictionary specifying the names of dimensions in the dataset.
     climatology : bool
         Indicates whether the dataset is climatological. Defaults to False.
+    use_dask: bool
+        Indicates whether to use dask for chunking. If True, data is loaded with dask; if False, data is loaded eagerly. Defaults to False.
+
 
     Attributes
     ----------
@@ -56,7 +61,7 @@ class Dataset:
     Dimensions:  ...
     """
 
-    filename: str
+    filename: Union[str, Path, List[Union[str, Path]]]
     start_time: Optional[datetime] = None
     end_time: Optional[datetime] = None
     var_names: Dict[str, str]
@@ -68,6 +73,7 @@ class Dataset:
         }
     )
     climatology: Optional[bool] = False
+    use_dask: Optional[bool] = True
 
     is_global: bool = field(init=False, repr=False)
     ds: xr.Dataset = field(init=False, repr=False)
@@ -118,41 +124,81 @@ class Dataset:
         ------
         FileNotFoundError
             If the specified file does not exist.
+        ValueError
+            If wildcards are found in the filename and use_dask=False.
+            If a list of files is provided but self.dim_names["time"] is not available or use_dask=False.
         """
 
-        # Check if the file exists
-        matching_files = glob.glob(self.filename)
-        if not matching_files:
-            raise FileNotFoundError(
-                f"No files found matching the pattern '{self.filename}'."
+        # Convert Path objects to strings
+        if isinstance(self.filename, (str, Path)):
+            filename_str = str(self.filename)
+        elif isinstance(self.filename, list):
+            filename_str = [str(f) for f in self.filename]
+        else:
+            raise ValueError(
+                "filename must be a string, Path, or a list of strings/Paths."
             )
 
-        # Load the dataset
-        with dask.config.set(**{"array.slicing.split_large_chunks": False}):
-            # Define the chunk sizes
+        # Handle the case when filename is a string
+        if isinstance(filename_str, str):
+            if re.search(r"[\*\?\[\]]", filename_str):
+                if not self.use_dask:
+                    raise ValueError(
+                        "Wildcards detected in the filename but use_dask=False. "
+                        "Wildcards can only be used when use_dask=True."
+                    )
+                matching_files = glob.glob(filename_str)
+                if not matching_files:
+                    raise FileNotFoundError(
+                        f"No files found matching the pattern '{filename_str}'."
+                    )
+            else:
+                matching_files = [filename_str]
+
+        # Handle the case when filename is a list
+        elif isinstance(filename_str, list):
+            if not self.use_dask:
+                raise ValueError(
+                    "A list of files is provided, but use_dask=False. "
+                    "Multiple files can only be loaded when use_dask=True."
+                )
+            matching_files = filename_str
+
+        # Check if time dimension is available when multiple files are provided
+        if isinstance(filename_str, list) and "time" not in self.dim_names:
+            raise ValueError(
+                "A list of files is provided, but time dimension is not available. "
+                "A time dimension must be available to concatenate the files."
+            )
+
+        if self.use_dask:
+
             chunks = {
                 self.dim_names["latitude"]: -1,
                 self.dim_names["longitude"]: -1,
             }
-            if "depth" in self.dim_names.keys():
+            if "depth" in self.dim_names:
                 chunks[self.dim_names["depth"]] = -1
-            if "time" in self.dim_names.keys():
+            if "time" in self.dim_names:
                 chunks[self.dim_names["time"]] = 1
 
-                ds = xr.open_mfdataset(
-                    self.filename,
-                    combine="nested",
-                    concat_dim=self.dim_names["time"],
-                    coords="minimal",
-                    compat="override",
-                    chunks=chunks,
-                    engine="netcdf4",
-                )
+            if re.search(r"[\*\?\[\]]", filename_str) or len(matching_files) == 1:
+                kwargs = {"combine": "by_coords"}
             else:
-                ds = xr.open_dataset(
-                    self.filename,
-                    chunks=chunks,
-                )
+                kwargs = {"combine": "nested", "concat_dim": self.dim_names["time"]}
+
+            ds = xr.open_mfdataset(
+                matching_files,
+                coords="minimal",
+                compat="override",
+                chunks=chunks,
+                **kwargs,
+            )
+        else:
+            ds = xr.open_dataset(matching_files[0], chunks=None)
+
+        if "time" in self.dim_names and self.dim_names["time"] not in ds.dims:
+            ds = ds.expand_dims(self.dim_names["time"])
 
         return ds
 
@@ -231,9 +277,8 @@ class Dataset:
         """
         Selects and returns the subset of the dataset corresponding to the specified time range.
 
-        This function filters the dataset to include only the data points within the specified
-        time range, defined by `self.start_time` and `self.end_time`. If `self.end_time` is not
-        provided, it defaults to one day after `self.start_time`.
+        This function filters the dataset to include all times between `start_time` and `end_time`,
+        plus exactly one record at or before `start_time`, and exactly one record at or after `end_time`.
 
         Parameters
         ----------
@@ -243,22 +288,26 @@ class Dataset:
         Returns
         -------
         xr.Dataset
-            A dataset containing only the data points within the specified time range.
+            A dataset containing the filtered time range.
 
         Raises
         ------
         ValueError
-            If no matching times are found or if the number of matching times does not meet expectations.
+            If no matching times are found.
 
         Warns
         -----
         UserWarning
-            If the dataset contains only 12 time steps but the climatology flag is not set.
+            If the dataset contains exactly 12 time steps but the climatology flag is not set.
             This may indicate that the dataset represents climatology data.
+
+        UserWarning
+            If no records at or before `self.start_time` or no records at or after `self.end_time` are found.
+
         """
 
         time_dim = self.dim_names["time"]
-        if time_dim in ds.coords or time_dim in ds.data_vars:
+        if time_dim in ds.variables:
             if self.climatology:
                 if not self.end_time:
                     # Interpolate from climatology for initial conditions
@@ -276,29 +325,58 @@ class Dataset:
                     ds = ds.assign_coords(
                         {time_dim: convert_cftime_to_datetime(ds[time_dim])}
                     )
-                if not self.end_time:
-                    end_time = self.start_time + timedelta(days=1)
-                else:
+                if self.end_time:
                     end_time = self.end_time
 
-                times = (np.datetime64(self.start_time) <= ds[time_dim]) & (
-                    ds[time_dim] < np.datetime64(end_time)
-                )
-                ds = ds.where(times, drop=True)
+                    # Identify records before or at start_time
+                    before_start = ds[time_dim] <= np.datetime64(self.start_time)
+                    if before_start.any():
+                        closest_before_start = (
+                            ds[time_dim].where(before_start, drop=True).max()
+                        )
+                    else:
+                        warnings.warn("No records found at or before the start_time.")
+                        closest_before_start = ds[time_dim].min()
+
+                    # Identify records after or at end_time
+                    after_end = ds[time_dim] >= np.datetime64(end_time)
+                    if after_end.any():
+                        closest_after_end = (
+                            ds[time_dim].where(after_end, drop=True).min()
+                        )
+                    else:
+                        warnings.warn("No records found at or after the end_time.")
+                        closest_after_end = ds[time_dim].max()
+
+                    # Select records within the time range and add the closest before/after
+                    within_range = (ds[time_dim] > np.datetime64(self.start_time)) & (
+                        ds[time_dim] < np.datetime64(end_time)
+                    )
+                    selected_times = ds[time_dim].where(
+                        within_range
+                        | (ds[time_dim] == closest_before_start)
+                        | (ds[time_dim] == closest_after_end),
+                        drop=True,
+                    )
+                    ds = ds.sel({time_dim: selected_times})
+                else:
+                    end_time = self.start_time + timedelta(days=1)
+                    times = (np.datetime64(self.start_time) <= ds[time_dim]) & (
+                        ds[time_dim] < np.datetime64(end_time)
+                    )
+                    ds = ds.where(times, drop=True)
+
+                    if ds.sizes[time_dim] != 1:
+                        found_times = ds.sizes[time_dim]
+                        raise ValueError(
+                            f"There must be exactly one time matching the start_time. Found {found_times} matching times."
+                        )
+
         else:
             warnings.warn(
                 "Dataset does not contain any time information. Please check if the time dimension "
                 "is correctly named or if the dataset includes time data."
             )
-        if not ds.sizes[time_dim]:
-            raise ValueError("No matching times found in the dataset.")
-
-        if not self.end_time:
-            if ds.sizes[time_dim] != 1:
-                found_times = ds.sizes[time_dim]
-                raise ValueError(
-                    f"There must be exactly one time matching the start_time. Found {found_times} matching times."
-                )
 
         return ds
 
@@ -384,7 +462,11 @@ class Dataset:
                 field = ds[var]
                 field_concatenated = xr.concat(
                     [field, field, field], dim=self.dim_names["longitude"]
-                ).chunk({self.dim_names["longitude"]: -1})
+                )
+                if self.use_dask:
+                    field_concatenated = field_concatenated.chunk(
+                        {self.dim_names["longitude"]: -1}
+                    )
                 field_concatenated[self.dim_names["longitude"]] = lon_concatenated
                 ds_concatenated[var] = field_concatenated
             else:
@@ -584,9 +666,14 @@ class TPXODataset(Dataset):
                 "ntides": self.dim_names["ntides"],
             },
         )
+        self.check_dataset(ds)
+
         # Select relevant fields
         ds = super().select_relevant_fields(ds)
 
+        # Make sure that latitude is ascending
+        ds = super().ensure_latitude_ascending(ds)
+
         # Check whether the data covers the entire globe
         object.__setattr__(self, "is_global", super().check_if_global(ds))
 
@@ -689,7 +776,7 @@ class CESMDataset(Dataset):
     Attributes
     ----------
     ds : xr.Dataset
-        The xarray Dataset containing the GLORYS data on its original grid.
+        The xarray Dataset containing the CESM data on its original grid.
     """
 
     # overwrite load_data method from parent class
@@ -708,38 +795,17 @@ class CESMDataset(Dataset):
             If the specified file does not exist.
         """
 
-        # Check if the file exists
-        matching_files = glob.glob(self.filename)
-        if not matching_files:
-            raise FileNotFoundError(
-                f"No files found matching the pattern '{self.filename}'."
-            )
-
-        # Load the dataset
-        with dask.config.set(**{"array.slicing.split_large_chunks": False}):
-            # Define the chunk sizes
-            chunks = {
-                self.dim_names["latitude"]: -1,
-                self.dim_names["longitude"]: -1,
-            }
+        ds = super().load_data()
 
-            ds = xr.open_mfdataset(
-                self.filename,
-                combine="nested",
-                coords="minimal",
-                compat="override",
-                chunks=chunks,
-                engine="netcdf4",
-            )
-            if "time" not in self.dim_names:
-                if "time" in ds.dims:
-                    self.dim_names["time"] = "time"
+        if "time" not in self.dim_names:
+            if "time" in ds.dims:
+                self.dim_names["time"] = "time"
+            else:
+                if "month" in ds.dims:
+                    self.dim_names["time"] = "month"
                 else:
-                    if "month" in ds.dims:
-                        self.dim_names["time"] = "month"
-                    else:
-                        ds = ds.expand_dims({"time": 1})
-                        self.dim_names["time"] = "time"
+                    ds = ds.expand_dims({"time": 1})
+                    self.dim_names["time"] = "time"
 
         return ds
 
@@ -769,6 +835,8 @@ class CESMDataset(Dataset):
             ds = assign_dates_to_climatology(ds, time_dim)
             # rename dimension
             ds = ds.swap_dims({time_dim: "time"})
+            if time_dim in ds.variables:
+                ds = ds.drop_vars(time_dim)
             # Update dimension names
             updated_dim_names = self.dim_names.copy()
             updated_dim_names["time"] = "time"
@@ -801,7 +869,7 @@ class CESMBGCDataset(CESMDataset):
     Attributes
     ----------
     ds : xr.Dataset
-        The xarray Dataset containing the GLORYS data on its original grid.
+        The xarray Dataset containing the CESM data on its original grid.
     """
 
     var_names: Dict[str, str] = field(
@@ -861,20 +929,23 @@ class CESMBGCDataset(CESMDataset):
             # Fill variables that only have data in upper 150m with NaNs below
             if (
                 "z_t_150m" in self.ds.dims
-                and np.equal(self.ds.z_t[:15].values, self.ds.z_t_150m.values).all()
+                and np.equal(
+                    self.ds.z_t[: len(self.ds.z_t_150m)].values, self.ds.z_t_150m.values
+                ).all()
             ):
                 for var in self.var_names:
                     if "z_t_150m" in self.ds[var].dims:
                         self.ds[var] = self.ds[var].rename({"z_t_150m": "z_t"})
-                        self.ds[var] = self.ds[var].chunk({"z_t": -1})
+                        if self.use_dask:
+                            self.ds[var] = self.ds[var].chunk({"z_t": -1})
             # Convert depth from cm to m
             ds = self.ds.assign_coords({"depth": self.ds["z_t"] / 100})
             ds["depth"].attrs["long_name"] = "Depth"
             ds["depth"].attrs["units"] = "m"
             ds = ds.swap_dims({"z_t": "depth"})
-            if "z_t" in ds:
+            if "z_t" in ds.variables:
                 ds = ds.drop_vars("z_t")
-            if "z_t_150m" in ds:
+            if "z_t_150m" in ds.variables:
                 ds = ds.drop_vars("z_t_150m")
             # update dataset
             object.__setattr__(self, "ds", ds)
@@ -909,7 +980,7 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
     Attributes
     ----------
     ds : xr.Dataset
-        The xarray Dataset containing the GLORYS data on its original grid.
+        The xarray Dataset containing the CESM data on its original grid.
     """
 
     var_names: Dict[str, str] = field(
@@ -932,6 +1003,19 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
 
     climatology: Optional[bool] = False
 
+    def post_process(self):
+        """
+        Perform post-processing on the dataset to remove specific variables.
+
+        This method checks if the variable "z_t" exists in the dataset. If it does,
+        the variable is removed from the dataset. The modified dataset is then
+        reassigned to the `ds` attribute of the object.
+        """
+
+        if "z_t" in self.ds.variables:
+            ds = self.ds.drop_vars("z_t")
+            object.__setattr__(self, "ds", ds)
+
 
 @dataclass(frozen=True, kw_only=True)
 class ERA5Dataset(Dataset):
@@ -957,7 +1041,7 @@ class ERA5Dataset(Dataset):
     Attributes
     ----------
     ds : xr.Dataset
-        The xarray Dataset containing the GLORYS data on its original grid.
+        The xarray Dataset containing the ERA5 data on its original grid.
     """
 
     var_names: Dict[str, str] = field(

roms_tools/setup/download.py

@@ -32,14 +32,17 @@ pup_test_data = pooch.create(
     # The registry specifies the files that can be fetched
     registry={
         "GLORYS_test_data.nc": "648f88ec29c433bcf65f257c1fb9497bd3d5d3880640186336b10ed54f7129d2",
+        "GLORYS_coarse_test_data.nc": "ed14ca6aa72810e2472e6ee21c59e5e38f59cd6eb39c14ff6a01ccba05d11d48",
         "ERA5_regional_test_data.nc": "bd12ce3b562fbea2a80a3b79ba74c724294043c28dc98ae092ad816d74eac794",
         "ERA5_global_test_data.nc": "8ed177ab64c02caf509b9fb121cf6713f286cc603b1f302f15f3f4eb0c21dc4f",
         "TPXO_global_test_data.nc": "457bfe87a7b247ec6e04e3c7d3e741ccf223020c41593f8ae33a14f2b5255e60",
         "TPXO_regional_test_data.nc": "11739245e2286d9c9d342dce5221e6435d2072b50028bef2e86a30287b3b4032",
         "CESM_regional_test_data_one_time_slice.nc": "43b578ecc067c85f95d6b97ed7b9dc8da7846f07c95331c6ba7f4a3161036a17",
         "CESM_regional_test_data_climatology.nc": "986a200029d9478fd43e6e4a8bc43e8a8f4407554893c59b5fcc2e86fd203272",
+        "CESM_regional_coarse_test_data_climatology.nc": "5cde5f968fba7900b6ff5bcf135126b5e25185fc3bd842bf66052cc2a6197d81",
         "CESM_surface_global_test_data_climatology.nc": "a072757110c6f7b716a98f867688ef4195a5966741d2f368201ac24617254e35",
         "CESM_surface_global_test_data.nc": "874106ffbc8b1b220db09df1551bbb89d22439d795b4d1e5a24ee775e9a7bf6e",
+        "grid_created_with_matlab.nc": "fd537ef8159fabb18e38495ec8d44e2fa1b7fb615fcb1417dd4c0e1bb5f4e41d",
     },
 )