roms-tools 1.2.0__py3-none-any.whl → 1.4.0__py3-none-any.whl

roms_tools/setup/datasets.py

@@ -1,15 +1,16 @@
+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,
     interpolate_from_climatology,
-    is_cftime_datetime,
+    get_time_type,
     convert_cftime_to_datetime,
 )
 from roms_tools.setup.download import download_correction_data
@@ -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,10 @@ 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.
+    apply_post_processing: bool
+        Indicates whether to post-process the dataset for futher use. Defaults to True.
 
     Attributes
     ----------
@@ -50,13 +56,9 @@ class Dataset:
     ...     start_time=datetime(2022, 1, 1),
     ...     end_time=datetime(2022, 12, 31),
     ... )
-    >>> dataset.load_data()
-    >>> print(dataset.ds)
-    <xarray.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 +70,8 @@ class Dataset:
         }
     )
     climatology: Optional[bool] = False
+    use_dask: Optional[bool] = True
+    apply_post_processing: Optional[bool] = True
 
     is_global: bool = field(init=False, repr=False)
     ds: xr.Dataset = field(init=False, repr=False)
@@ -78,11 +82,22 @@ class Dataset:
         1. Loads the dataset from the specified filename.
         2. Applies time filtering based on start_time and end_time if provided.
         3. Selects relevant fields as specified by var_names.
-        4. Ensures latitude values are in ascending order.
+        4. Ensures latitude values 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__}."
+            )
+
         ds = self.load_data()
+        ds = self.clean_up(ds)
         self.check_dataset(ds)
 
         # Select relevant times
@@ -90,11 +105,18 @@ class Dataset:
             ds = self.add_time_info(ds)
             ds = self.select_relevant_times(ds)
 
+            if self.dim_names["time"] != "time":
+                ds = ds.rename({self.dim_names["time"]: "time"})
+
         # Select relevant fields
         ds = self.select_relevant_fields(ds)
 
         # Make sure that latitude is ascending
-        ds = self.ensure_latitude_ascending(ds)
+        ds = self.ensure_dimension_is_ascending(ds, dim="latitude")
+
+        if "depth" in self.dim_names:
+            # Make sure that depth is ascending
+            ds = self.ensure_dimension_is_ascending(ds, dim="depth")
 
         # Check whether the data covers the entire globe
         object.__setattr__(self, "is_global", self.check_if_global(ds))
@@ -105,6 +127,9 @@ class Dataset:
 
         object.__setattr__(self, "ds", ds)
 
+        if self.apply_post_processing:
+            self.post_process()
+
     def load_data(self) -> xr.Dataset:
         """
         Load dataset from the specified file.
@@ -118,42 +143,125 @@ class Dataset:
         ------
         FileNotFoundError
             If the specified file does not exist.
+        ValueError
+            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}'."
+        # Precompile the regex for matching wildcard characters
+        wildcard_regex = re.compile(r"[\*\?\[\]]")
+
+        # 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
+        contains_wildcard = False
+        if isinstance(filename_str, str):
+            contains_wildcard = bool(wildcard_regex.search(filename_str))
+            if contains_wildcard:
+                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):
+            contains_wildcard = any(wildcard_regex.search(f) for f in filename_str)
+            if contains_wildcard:
+                matching_files = []
+                for f in filename_str:
+                    files = glob.glob(f)
+                    if not files:
+                        raise FileNotFoundError(
+                            f"No files found matching the pattern '{f}'."
+                        )
+                    matching_files.extend(files)
+            else:
+                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."
+            )
+
+        # Determine the kwargs for combining datasets
+        if contains_wildcard or len(matching_files) == 1:
+            # If there is a wildcard or just one file, use by_coords
+            kwargs = {"combine": "by_coords"}
+        else:
+            # Otherwise, use nested combine based on time
+            kwargs = {"combine": "nested", "concat_dim": self.dim_names["time"]}
+
+        # Base kwargs used for dataset combination
+        combine_kwargs = {
+            "coords": "minimal",
+            "compat": "override",
+            "combine_attrs": "override",
+        }
+
+        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",
-                )
-            else:
-                ds = xr.open_dataset(
-                    self.filename,
-                    chunks=chunks,
+            ds = xr.open_mfdataset(
+                matching_files,
+                chunks=chunks,
+                **combine_kwargs,
+                **kwargs,
+            )
+        else:
+            ds_list = []
+            for file in matching_files:
+                ds = xr.open_dataset(file, chunks=None)
+                ds_list.append(ds)
+
+            if kwargs["combine"] == "by_coords":
+                ds = xr.combine_by_coords(ds_list, **combine_kwargs)
+            elif kwargs["combine"] == "nested":
+                ds = xr.combine_nested(
+                    ds_list, concat_dim=kwargs["concat_dim"], **combine_kwargs
                 )
 
+        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
+
+    def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
+        """
+        Dummy method to be overridden by child classes to clean up the dataset.
+
+        This method is intended as a placeholder and should be implemented in subclasses
+        to provide specific functionality.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            The xarray Dataset to be cleaned up.
+
+        Returns
+        -------
+        xr.Dataset
+            The xarray Dataset cleaned up (as implemented by child classes).
+        """
         return ds
 
     def check_dataset(self, ds: xr.Dataset) -> None:
@@ -206,8 +314,6 @@ class Dataset:
 
         return ds
 
-    import xarray as xr
-
     def add_time_info(self, ds: xr.Dataset) -> xr.Dataset:
         """
         Dummy method to be overridden by child classes to add time information to the dataset.
@@ -229,97 +335,163 @@ class Dataset:
 
     def select_relevant_times(self, ds) -> xr.Dataset:
         """
-        Selects and returns the subset of the dataset corresponding to the specified time range.
+        Select a subset of the dataset based on the specified time range.
+
+        This method filters the dataset to include all records between `start_time` and `end_time`.
+        Additionally, it ensures that one record at or before `start_time` and one record at or
+        after `end_time` are included, even if they fall outside the strict 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`.
+        If no `end_time` is specified, the method will select the time range of
+        [start_time, start_time + 24 hours] and return the closest time entry to `start_time` within that range.
 
         Parameters
         ----------
         ds : xr.Dataset
-            The input dataset to be filtered.
+            The input dataset to be filtered. Must contain a time dimension.
 
         Returns
         -------
         xr.Dataset
-            A dataset containing only the data points within the specified time range.
+            A dataset filtered to the specified time range, including the closest entries
+            at or before `start_time` and at or after `end_time` if applicable.
 
         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 between `start_time` and `start_time + 24 hours`.
 
         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 `start_time` or no records at or after `end_time` are found.
+
+        UserWarning
+            If the dataset does not contain any time dimension or the time dimension is incorrectly named.
+
+        Notes
+        -----
+        - If the `climatology` flag is set and `end_time` is not provided, the method will
+          interpolate initial conditions from climatology data.
+        - If the dataset uses `cftime` datetime objects, these will be converted to standard
+          `np.datetime64` objects before filtering.
         """
 
         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 len(ds[time_dim]) != 12:
+                    raise ValueError(
+                        f"The dataset contains {len(ds[time_dim])} time steps, but the climatology flag is set to True, which requires exactly 12 time steps."
+                    )
                 if not self.end_time:
                     # Interpolate from climatology for initial conditions
                     ds = interpolate_from_climatology(
                         ds, self.dim_names["time"], self.start_time
                     )
             else:
-                if len(ds[time_dim]) == 12:
-                    warnings.warn(
-                        "The dataset contains exactly 12 time steps. This may indicate that it is "
-                        "climatological data. Please verify if climatology is appropriate for your "
-                        "analysis and set the climatology flag to True."
+                time_type = get_time_type(ds[time_dim])
+                if time_type == "int":
+                    raise ValueError(
+                        "The dataset contains integer time values, which are only supported when the climatology flag is set to True. However, your climatology flag is set to False."
                     )
-                if is_cftime_datetime(ds[time_dim]):
+                if time_type == "cftime":
                     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:
+                    # Look in time range [self.start_time, self.start_time + 24h]
+                    end_time = self.start_time + timedelta(days=1)
+                    times = (np.datetime64(self.start_time) <= ds[time_dim]) & (
+                        ds[time_dim] < np.datetime64(end_time)
+                    )
+                    if np.all(~times):
+                        raise ValueError(
+                            f"The dataset does not contain any time entries between the specified start_time: {self.start_time} "
+                            f"and {self.start_time + timedelta(hours=24)}. "
+                            "Please ensure the dataset includes time entries for that range."
+                        )
+
+                    ds = ds.where(times, drop=True)
+                    if ds.sizes[time_dim] > 1:
+                        # Pick the time closest to self.start_time
+                        ds = ds.isel({time_dim: 0})
+                    print(
+                        f"Selected time entry closest to the specified start_time ({self.start_time}) within the range [{self.start_time}, {self.start_time + timedelta(hours=24)}]: {ds[time_dim].values}"
+                    )
         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
 
-    def ensure_latitude_ascending(self, ds: xr.Dataset) -> xr.Dataset:
+    def ensure_dimension_is_ascending(
+        self, ds: xr.Dataset, dim="latitude"
+    ) -> xr.Dataset:
         """
-        Ensure that the latitude dimension is in ascending order.
+        Ensure that the specified dimension in the dataset is in ascending order.
+
+        If the values along the specified dimension are in descending order, this function reverses the order of the dimension to make it ascending.
 
         Parameters
         ----------
         ds : xr.Dataset
-            The xarray Dataset to check.
+            The input `xarray.Dataset` whose dimension is to be checked and, if necessary, reordered.
+        dim : str, optional
+            The name of the dimension to check for ascending order.
+            Defaults to "latitude". The dimension is expected to be one of the keys in `self.dim_names`.
 
         Returns
         -------
-        ds : xr.Dataset
-            The xarray Dataset with latitude in ascending order.
+        xr.Dataset
+            A new `xarray.Dataset` with the specified dimension in ascending order.
+            If the dimension was already in ascending order, the original dataset is returned unchanged.
+            If the dimension was in descending order, the dataset is returned with the dimension reversed.
+
         """
         # Make sure that latitude is ascending
-        lat_diff = np.diff(ds[self.dim_names["latitude"]])
-        if np.all(lat_diff < 0):
-            ds = ds.isel(**{self.dim_names["latitude"]: slice(None, None, -1)})
+        diff = np.diff(ds[self.dim_names[dim]])
+        if np.all(diff < 0):
+            ds = ds.isel(**{self.dim_names[dim]: slice(None, None, -1)})
 
         return ds
 
@@ -384,7 +556,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:
@@ -392,6 +568,17 @@ class Dataset:
 
         return ds_concatenated
 
+    def post_process(self):
+        """
+        Placeholder method to be overridden by subclasses for dataset post-processing.
+
+        Returns
+        -------
+        None
+            This method does not return any value. Subclasses are expected to modify the dataset in-place.
+        """
+        pass
+
     def choose_subdomain(
         self, latitude_range, longitude_range, margin, straddle, return_subdomain=False
     ):
@@ -489,19 +676,6 @@ class Dataset:
         else:
             object.__setattr__(self, "ds", subdomain)
 
-    def convert_to_negative_depth(self):
-        """
-        Converts the depth values in the dataset to negative if they are non-negative.
-
-        This method checks the values in the depth dimension of the dataset (`self.ds[self.dim_names["depth"]]`).
-        If all values are greater than or equal to zero, it negates them and updates the dataset accordingly.
-
-        """
-        depth = self.ds[self.dim_names["depth"]]
-
-        if (depth >= 0).all():
-            self.ds[self.dim_names["depth"]] = -depth
-
 
 @dataclass(frozen=True, kw_only=True)
 class TPXODataset(Dataset):
@@ -557,11 +731,29 @@ class TPXODataset(Dataset):
     ds: xr.Dataset = field(init=False, repr=False)
     reference_date: datetime = datetime(1992, 1, 1)
 
-    def __post_init__(self):
-        # Perform any necessary dataset initialization or modifications here
-        ds = super().load_data()
+    def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
+        """
+        Clean up and standardize the dimensions and coordinates of the dataset for further processing.
+
+        This method performs the following operations:
+        - Assigns new coordinate variables for 'omega', 'longitude', and 'latitude' based on existing dataset variables.
+          - 'omega' is retained as it is.
+          - 'longitude' is derived from 'lon_r', assuming it is constant along the 'ny' dimension.
+          - 'latitude' is derived from 'lat_r', assuming it is constant along the 'nx' dimension.
+        - Renames the dimensions 'nx' and 'ny' to 'longitude' and 'latitude', respectively, for consistency.
+        - Renames the tidal dimension to 'ntides' for standardization.
+        - Updates the `dim_names` attribute of the object to reflect the new dimension names: 'longitude', 'latitude', and 'ntides'.
+
+        Parameters
+        ----------
+        ds : xr.Dataset
+            The input dataset to be cleaned and standardized. It should contain the coordinates 'omega', 'lon_r', 'lat_r', and the tidal dimension.
 
-        # Clean up dataset
+        Returns
+        -------
+        ds : xr.Dataset
+            A cleaned and standardized `xarray.Dataset` with updated coordinates and dimensions.
+        """
         ds = ds.assign_coords(
             {
                 "omega": ds["omega"],
@@ -573,7 +765,9 @@ class TPXODataset(Dataset):
                 ),  # lat_r is constant along nx, i.e., is only a function of ny
             }
         )
-        ds = ds.rename({"nx": "longitude", "ny": "latitude"})
+        ds = ds.rename(
+            {"nx": "longitude", "ny": "latitude", self.dim_names["ntides"]: "ntides"}
+        )
 
         object.__setattr__(
             self,
@@ -581,25 +775,11 @@ class TPXODataset(Dataset):
             {
                 "latitude": "latitude",
                 "longitude": "longitude",
-                "ntides": self.dim_names["ntides"],
+                "ntides": "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))
-
-        # If dataset is global concatenate three copies of field along longitude dimension
-        if self.is_global:
-            ds = super().concatenate_longitudes(ds)
-
-        object.__setattr__(self, "ds", ds)
+        return ds
 
     def check_number_constituents(self, ntides: int):
         """
@@ -620,6 +800,29 @@ class TPXODataset(Dataset):
                 f"The dataset contains fewer than {ntides} tidal constituents."
             )
 
+    def post_process(self):
+        """
+        Apply a depth-based mask to the dataset, ensuring only positive depths are retained.
+
+        This method checks if the 'depth' variable is present in the dataset. If found, a mask is created where
+        depths greater than 0 are considered valid (mask value of 1). This mask is applied to all data variables
+        in the dataset, replacing values at invalid depths (depth ≤ 0) with NaN. The mask itself is also stored
+        in the dataset under the variable 'mask'.
+
+        Returns
+        -------
+        None
+            The dataset is modified in-place by applying the mask to each variable.
+        """
+
+        if "depth" in self.var_names.keys():
+            mask = xr.where(self.ds["depth"] > 0, 1, 0)
+
+            for var in self.ds.data_vars:
+                self.ds[var] = xr.where(mask == 1, self.ds[var], np.nan)
+
+            self.ds["mask"] = mask
+
 
 @dataclass(frozen=True, kw_only=True)
 class GLORYSDataset(Dataset):
@@ -669,6 +872,33 @@ class GLORYSDataset(Dataset):
 
     climatology: Optional[bool] = False
 
+    def post_process(self):
+        """
+        Apply a mask to the dataset based on the 'zeta' variable, with 0 where 'zeta' is NaN.
+
+        This method creates a mask based on the
+        first time step (time=0) of 'zeta'. The mask has 1 for valid data and 0 where 'zeta' is NaN. This mask is applied
+        to all data variables, replacing values with NaN where 'zeta' is NaN at time=0.
+        The mask itself is stored in the dataset under the variable 'mask'.
+
+        Returns
+        -------
+        None
+            The dataset is modified in-place by applying the mask to each variable.
+
+        """
+
+        mask = xr.where(
+            self.ds[self.var_names["zeta"]].isel({self.dim_names["time"]: 0}).isnull(),
+            0,
+            1,
+        )
+
+        for var in self.ds.data_vars:
+            self.ds[var] = xr.where(mask == 1, self.ds[var], np.nan)
+
+        self.ds["mask"] = mask
+
 
 @dataclass(frozen=True, kw_only=True)
 class CESMDataset(Dataset):
@@ -689,62 +919,37 @@ class CESMDataset(Dataset):
     dim_names: Dict[str, str], optional
         Dictionary specifying the names of dimensions in the dataset.
     climatology : bool
-        Indicates whether the dataset is climatological. Defaults to True.
+        Indicates whether the dataset is climatological. Defaults to False.
 
     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
-    def load_data(self) -> xr.Dataset:
+    # overwrite clean_up method from parent class
+    def clean_up(self, ds: xr.Dataset) -> xr.Dataset:
         """
-        Load dataset from the specified file.
+        Ensure the dataset's time dimension is correctly defined and standardized.
+
+        This method verifies that the time dimension exists in the dataset and assigns it appropriately. If the "time" dimension is missing, the method attempts to assign an existing "time" or "month" dimension. If neither exists, it expands the dataset to include a "time" dimension with a size of one.
 
         Returns
         -------
         ds : xr.Dataset
-            The loaded xarray Dataset containing the forcing data.
+            The xarray Dataset with the correct time dimension assigned or added.
 
-        Raises
-        ------
-        FileNotFoundError
-            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 = 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
 
@@ -803,12 +1008,12 @@ class CESMBGCDataset(CESMDataset):
     dim_names: Dict[str, str], optional
         Dictionary specifying the names of dimensions in the dataset.
     climatology : bool
-        Indicates whether the dataset is climatological. Defaults to True.
+        Indicates whether the dataset is climatological. Defaults to False.
 
     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(
@@ -856,24 +1061,27 @@ class CESMBGCDataset(CESMDataset):
         }
     )
 
-    climatology: Optional[bool] = True
+    climatology: Optional[bool] = False
 
     def post_process(self):
         """
         Processes and converts CESM data values as follows:
         - Convert depth values from cm to m.
+        - Apply a mask to the dataset based on the 'P04' variable at the surface.
         """
-
         if self.dim_names["depth"] == "z_t":
             # 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"
@@ -891,6 +1099,19 @@ class CESMBGCDataset(CESMDataset):
             updated_dim_names["depth"] = "depth"
             object.__setattr__(self, "dim_names", updated_dim_names)
 
+        mask = xr.where(
+            self.ds[self.var_names["PO4"]]
+            .isel({self.dim_names["time"]: 0, self.dim_names["depth"]: 0})
+            .isnull(),
+            0,
+            1,
+        )
+
+        for var in self.ds.data_vars:
+            self.ds[var] = xr.where(mask == 1, self.ds[var], np.nan)
+
+        self.ds["mask"] = mask
+
 
 @dataclass(frozen=True, kw_only=True)
 class CESMBGCSurfaceForcingDataset(CESMDataset):
@@ -916,7 +1137,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(
@@ -952,6 +1173,19 @@ class CESMBGCSurfaceForcingDataset(CESMDataset):
             ds = self.ds.drop_vars("z_t")
             object.__setattr__(self, "ds", ds)
 
+        mask = xr.where(
+            self.ds[self.var_names["pco2_air"]]
+            .isel({self.dim_names["time"]: 0})
+            .isnull(),
+            0,
+            1,
+        )
+
+        for var in self.ds.data_vars:
+            self.ds[var] = xr.where(mask == 1, self.ds[var], np.nan)
+
+        self.ds["mask"] = mask
+
 
 @dataclass(frozen=True, kw_only=True)
 class ERA5Dataset(Dataset):
@@ -977,7 +1211,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(
@@ -1010,6 +1244,7 @@ class ERA5Dataset(Dataset):
         - Convert rainfall from meters to cm/day.
         - Convert temperature from Kelvin to Celsius.
         - Compute relative humidity if not present, convert to absolute humidity.
+        - Use SST to create mask.
         """
         # Translate radiation to fluxes. ERA5 stores values integrated over 1 hour.
         # Convert radiation from J/m^2 to W/m^2
@@ -1061,6 +1296,8 @@ class ERA5Dataset(Dataset):
             for var in self.ds.data_vars:
                 self.ds[var] = xr.where(mask == 1, self.ds[var], np.nan)
 
+            self.ds["mask"] = mask
+
 
 @dataclass(frozen=True, kw_only=True)
 class ERA5Correction(Dataset):