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

roms_tools/setup/datasets.py

@@ -10,7 +10,7 @@ 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
@@ -39,7 +39,8 @@ class Dataset:
         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
     ----------
@@ -55,10 +56,6 @@ class Dataset:
     ...     start_time=datetime(2022, 1, 1),
     ...     end_time=datetime(2022, 12, 31),
     ... )
-    >>> dataset.load_data()
-    >>> print(dataset.ds)
-    <xarray.Dataset>
-    Dimensions:  ...
     """
 
     filename: Union[str, Path, List[Union[str, Path]]]
@@ -74,6 +71,7 @@ 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)
@@ -84,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
@@ -96,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))
@@ -111,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.
@@ -125,10 +144,12 @@ 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.
         """
 
+        # 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)
@@ -140,13 +161,10 @@ class Dataset:
             )
 
         # Handle the case when filename is a string
+        contains_wildcard = False
         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."
-                    )
+            contains_wildcard = bool(wildcard_regex.search(filename_str))
+            if contains_wildcard:
                 matching_files = glob.glob(filename_str)
                 if not matching_files:
                     raise FileNotFoundError(
@@ -157,12 +175,18 @@ class Dataset:
 
         # 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
+            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:
@@ -171,6 +195,21 @@ class Dataset:
                 "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 = {
@@ -182,26 +221,49 @@ class Dataset:
             if "time" in self.dim_names:
                 chunks[self.dim_names["time"]] = 1
 
-            if re.search(r"[\*\?\[\]]", filename_str) or len(matching_files) == 1:
-                kwargs = {"combine": "by_coords"}
-            else:
-                kwargs = {"combine": "nested", "concat_dim": self.dim_names["time"]}
-
             ds = xr.open_mfdataset(
                 matching_files,
-                coords="minimal",
-                compat="override",
                 chunks=chunks,
+                **combine_kwargs,
                 **kwargs,
             )
         else:
-            ds = xr.open_dataset(matching_files[0], chunks=None)
+            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:
         """
         Check if the dataset contains the specified variables and dimensions.
@@ -252,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.
@@ -275,25 +335,30 @@ 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 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`.
+        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 the filtered 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.
+            If no matching times are found between `start_time` and `start_time + 24 hours`.
 
         Warns
         -----
@@ -302,26 +367,38 @@ class Dataset:
             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.
+            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.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])}
                     )
@@ -360,18 +437,25 @@ class Dataset:
                     )
                     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)
                     )
-                    ds = ds.where(times, drop=True)
-
-                    if ds.sizes[time_dim] != 1:
-                        found_times = ds.sizes[time_dim]
+                    if np.all(~times):
                         raise ValueError(
-                            f"There must be exactly one time matching the start_time. Found {found_times} matching times."
+                            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 "
@@ -380,24 +464,34 @@ class Dataset:
 
         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
 
@@ -474,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
     ):
@@ -571,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):
@@ -639,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'.
 
-        # Clean up dataset
+        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.
+
+        Returns
+        -------
+        ds : xr.Dataset
+            A cleaned and standardized `xarray.Dataset` with updated coordinates and dimensions.
+        """
         ds = ds.assign_coords(
             {
                 "omega": ds["omega"],
@@ -655,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,
@@ -663,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):
         """
@@ -702,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):
@@ -751,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):
@@ -771,7 +919,7 @@ 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
     ----------
@@ -779,24 +927,20 @@ class CESMDataset(Dataset):
         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.
         """
 
-        ds = super().load_data()
-
         if "time" not in self.dim_names:
             if "time" in ds.dims:
                 self.dim_names["time"] = "time"
@@ -864,7 +1008,7 @@ 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
     ----------
@@ -917,14 +1061,14 @@ 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 (
@@ -955,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):
@@ -1016,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):
@@ -1074,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
@@ -1125,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):