Repository: CWorthy-ocean/roms-tools
Issue category: Performance Issue

Issue:
Radiation correction is not done in parallel (even when using Dask)
## Issue

If the user chooses to apply a radiation correction for the surface forcing while using Dask via
```
SurfaceForcing(
    grid=grid,
    ...,
    correct_radiation=True,
    use_dask=True
)
```
this radiation correction is not done in parallel. The reason for this is the internal use of the function 
https://github.com/CWorthy-ocean/roms-tools/blob/67df5fc400b1daeb903ba9b2efe11899d8b55dfa/roms_tools/setup/utils.py#L128-L150

which interpolates a pre-computed climatology for ERA5 shortwave downward radiation (i.e., 12 time stamps) to the hourly ERA5 data of interest. Specifically, this line is the issue:

https://github.com/CWorthy-ocean/roms-tools/blob/67df5fc400b1daeb903ba9b2efe11899d8b55dfa/roms_tools/setup/utils.py#L184-L187

Both `data_array_concat` and `day_of_year` are chunked in time with chunk size 1, where `data_array_concat.time` has a length of 12 + 2, and `day_of_year.time` a length of however many ERA5 times there are: for a month of ERA5 data, that is 744 hourly time slices, i.e., 744 time chunks. However, the `xarray.interp` function cannot deal with chunked data along the core dimension and will return interpolated data with a single time chunk (containing 744 time slices). This means that this operation is not done in parallel.

Note that substituting this by something like
```
        data_array_interpolated = xr.apply_ufunc(
            np.interp, 
            day_of_year, 
            time_concat, 
            data_array_concat, 
            input_core_dims=[[time_dim_name], [time_dim_name], [time_dim_name]],
            output_core_dims=[[time_dim_name]],
            vectorize=True,
            dask="allowed",
            exclude_dims=set((time_dim_name,)),
        )
```
would not help because we are trying to chunk along a core dimension, which `xr.apply_ufunc` is not equipped for (I think).

## Implications for validation step and memory footprint

In the validation step during the surface forcing initialization, a single time slice of the corrected shortwave radiation is checked for NaNs:

https://github.com/CWorthy-ocean/roms-tools/blob/67df5fc400b1daeb903ba9b2efe11899d8b55dfa/roms_tools/setup/surface_forcing.py#L366

However, due to the issue above, this will imply a huge memory footprint because the shortwave radiation **correction** is computed for all time slices (e.g., 744 time stamps), not just one.

This can be seen in the following example, where we use Grid A from https://github.com/CWorthy-ocean/roms-tools/issues/201 and one month of hourly ERA5 data:

```
SurfaceForcing(
    grid=grid,
    ...,
    correct_radiation=True,
    use_dask=True,
    bypass_validation=True
)
```
has the following performance:
```
peak memory: 562.75 MiB, increment: 25.68 MiB
CPU times: user 54.8 s, sys: 1.26 s, total: 56.1 s
Wall time: 30.1 s
```
The memory footprint is minimal because the validation step was skipped.

However,

```
SurfaceForcing(
    grid=grid,
    ...,
    correct_radiation=True,
    use_dask=True,
    bypass_validation=False
)
```
has the following performance and big memory footprint:
```
peak memory: 32898.70 MiB, increment: 32335.71 MiB
CPU times: user 2min 24s, sys: 5.01 s, total: 2min 29s
Wall time: 43.8 s
```





Hints:

Patch:
diff --git a/docs/releases.md b/docs/releases.md
index 587315aa..1d5430af 100644
--- a/docs/releases.md
+++ b/docs/releases.md
@@ -10,10 +10,15 @@
 
 ### Internal Changes
 
+* Parallelize computation of radiation correction, leading to a hugely improved memory footprint for surface forcing generation ([#227](https://github.com/CWorthy-ocean/roms-tools/pull/227))
+* For computation of radiation correction, swap order of temporal and spatial interpolation to further improve memory footprint ([#227](https://github.com/CWorthy-ocean/roms-tools/pull/227))
+
 ### Documentation
 
 ### Bugfixes
 
+* Fix bug in validation step for surface forcing ([#227](https://github.com/CWorthy-ocean/roms-tools/pull/227))
+
 ## v2.3.0
 
 ### New Features
diff --git a/roms_tools/setup/datasets.py b/roms_tools/setup/datasets.py
index 634e2734..7b0ce441 100644
--- a/roms_tools/setup/datasets.py
+++ b/roms_tools/setup/datasets.py
@@ -2058,6 +2058,7 @@ def _select_relevant_times(
                 )
             if not end_time:
                 # Interpolate from climatology for initial conditions
+                ds["time"] = ds["time"].dt.days
                 ds = interpolate_from_climatology(ds, time_dim, start_time)
         else:
             time_type = get_time_type(ds[time_dim])
diff --git a/roms_tools/setup/surface_forcing.py b/roms_tools/setup/surface_forcing.py
index cff894af..3ce9f9dd 100644
--- a/roms_tools/setup/surface_forcing.py
+++ b/roms_tools/setup/surface_forcing.py
@@ -238,14 +238,14 @@ def _set_variable_info(self, data):
                 "qair": {**default_info, "validate": True},
                 "rain": {**default_info, "validate": False},
                 "uwnd": {
-                    "location": "u",
+                    "location": "rho",
                     "is_vector": True,
                     "vector_pair": "vwnd",
                     "is_3d": False,
                     "validate": True,
                 },
                 "vwnd": {
-                    "location": "v",
+                    "location": "rho",
                     "is_vector": True,
                     "vector_pair": "uwnd",
                     "is_3d": False,
@@ -266,7 +266,7 @@ def _set_variable_info(self, data):
     def _apply_correction(self, processed_fields, data):
 
         correction_data = self._get_correction_data()
-        # choose same subdomain as forcing data so that we can use same mask
+        # Match subdomain to forcing data to reuse the mask
         coords_correction = {
             "lat": data.ds[data.dim_names["latitude"]],
             "lon": data.ds[data.dim_names["longitude"]],
@@ -275,23 +275,39 @@ def _apply_correction(self, processed_fields, data):
             coords_correction, straddle=self.target_coords["straddle"]
         )
         correction_data.ds["mask"] = data.ds["mask"]  # use mask from ERA5 data
+        correction_data.ds["time"] = correction_data.ds["time"].dt.days
+
         correction_data.apply_lateral_fill()
-        # regrid
-        lateral_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
-        corr_factor = lateral_regrid.apply(
-            correction_data.ds[correction_data.var_names["swr_corr"]]
-        )
 
-        # temporal interpolation
-        corr_factor = interpolate_from_climatology(
-            corr_factor,
-            correction_data.dim_names["time"],
-            time=processed_fields["swrad"].time,
-        )
+        # Temporal interpolation: Perform before spatial regridding for better performance
+        if self.use_dask:
+            # Perform temporal interpolation for each time slice to enforce chunking in time.
+            # This reduces memory usage by processing one time step at a time.
+            # The interpolated slices are then concatenated along the "time" dimension.
+            corr_factor = xr.concat(
+                [
+                    interpolate_from_climatology(
+                        correction_data.ds[correction_data.var_names["swr_corr"]],
+                        correction_data.dim_names["time"],
+                        time=time,
+                    )
+                    for time in processed_fields["swrad"].time
+                ],
+                dim="time",
+            )
+        else:
+            # Interpolate across all time steps at once
+            corr_factor = interpolate_from_climatology(
+                correction_data.ds[correction_data.var_names["swr_corr"]],
+                correction_data.dim_names["time"],
+                time=processed_fields["swrad"].time,
+            )
 
-        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
+        # Spatial regridding
+        lateral_regrid = LateralRegrid(self.target_coords, correction_data.dim_names)
+        corr_factor = lateral_regrid.apply(corr_factor)
 
-        del corr_factor
+        processed_fields["swrad"] = processed_fields["swrad"] * corr_factor
 
         return processed_fields
 
@@ -357,12 +373,8 @@ def _validate(self, ds):
 
         for var_name in ds.data_vars:
             if self.variable_info[var_name]["validate"]:
-                if self.variable_info[var_name]["location"] == "rho":
-                    mask = self.target_coords["mask"]
-                elif self.variable_info[var_name]["location"] == "u":
-                    mask = self.target_coords["mask_u"]
-                elif self.variable_info[var_name]["location"] == "v":
-                    mask = self.target_coords["mask_v"]
+                # all variables are at rho-points
+                mask = self.target_coords["mask"]
                 nan_check(ds[var_name].isel(time=0), mask)
 
     def _add_global_metadata(self, ds=None):
diff --git a/roms_tools/setup/utils.py b/roms_tools/setup/utils.py
index 6823f543..10e37afa 100644
--- a/roms_tools/setup/utils.py
+++ b/roms_tools/setup/utils.py
@@ -130,23 +130,36 @@ def interpolate_from_climatology(
     time_dim_name: str,
     time: Union[xr.DataArray, pd.DatetimeIndex],
 ) -> Union[xr.DataArray, xr.Dataset]:
-    """Interpolates the given field temporally based on the specified time points.
+    """Temporally interpolates a field based on specified time points.
 
-    If `field` is an xarray.Dataset, this function applies the interpolation to all data variables in the dataset.
+    This function performs temporal interpolation on the input `field` to match the provided `time` values.
+    If the input `field` is an `xarray.Dataset`, the interpolation is applied to all its data variables individually.
 
     Parameters
     ----------
     field : xarray.DataArray or xarray.Dataset
-        The field data to be interpolated. Can be a single DataArray or a Dataset.
+        The input field to be interpolated.
+        - If `field` is an `xarray.DataArray`, it should have a time dimension identified by `time_dim_name`.
+        - If `field` is an `xarray.Dataset`, all variables within the dataset are interpolated along the specified time dimension.
+        The time dimension is assumed to represent `day_of_year` for climatological purposes.
     time_dim_name : str
-        The name of the dimension in `field` that represents time.
+        The name of the time dimension in the `field`. This dimension is used for interpolation.
     time : xarray.DataArray or pandas.DatetimeIndex
-        The target time points for interpolation.
+        The target time points for interpolation. The time values should be compatible with the time format used in the `field`.
 
     Returns
     -------
     xarray.DataArray or xarray.Dataset
-        The field values interpolated to the specified time points. The type matches the input type.
+        The interpolated field, with the same type as the input (`xarray.DataArray` or `xarray.Dataset`),
+        but aligned to the specified `time` values.
+
+    Notes
+    -----
+    - The interpolation assumes the time dimension in `field` corresponds to `day_of_year`.
+      If the input time values are in a datetime format, ensure they are converted to `day_of_year` before calling this function.
+      For example, you can preprocess the time as follows:
+
+      >>> field["time"] = field["time"].dt.dayofyear
     """
 
     def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
@@ -156,11 +169,11 @@ def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
             day_of_year = time.dt.dayofyear
         else:
             if np.size(time) == 1:
-                day_of_year = time.timetuple().tm_yday
+                # Convert single datetime64 object to pandas.Timestamp
+                day_of_year = pd.Timestamp(time).dayofyear
             else:
-                day_of_year = np.array([t.timetuple().tm_yday for t in time])
-
-        data_array[time_dim_name] = data_array[time_dim_name].dt.days
+                # Convert each datetime64 object in the array to pandas.Timestamp
+                day_of_year = np.array([pd.Timestamp(t).dayofyear for t in time])
 
         # Concatenate across the beginning and end of the year
         time_concat = xr.concat(
@@ -200,6 +213,7 @@ def interpolate_single_field(data_array: xr.DataArray) -> xr.DataArray:
             for var, data_array in field.data_vars.items()
         }
         return xr.Dataset(interpolated_data_vars, attrs=field.attrs)
+
     else:
         raise TypeError("Input 'field' must be an xarray.DataArray or xarray.Dataset.")
 
@@ -633,8 +647,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         mask = grid.ds.get("mask_coarse")
         if mask is not None:
             mask = mask.rename({"eta_coarse": "eta_rho", "xi_coarse": "xi_rho"})
-            mask_u = interpolate_from_rho_to_u(mask, method="multiplicative")
-            mask_v = interpolate_from_rho_to_v(mask, method="multiplicative")
 
         lat_psi = grid.ds.get("lat_psi_coarse")
         lon_psi = grid.ds.get("lon_psi_coarse")
@@ -644,8 +656,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         lon = grid.ds.lon_rho
         angle = grid.ds.angle
         mask = grid.ds.get("mask_rho")
-        mask_u = grid.ds.get("mask_u")
-        mask_v = grid.ds.get("mask_v")
         lat_psi = grid.ds.get("lat_psi")
         lon_psi = grid.ds.get("lon_psi")
 
@@ -668,8 +678,6 @@ def get_target_coords(grid, use_coarse_grid=False):
         "lon_psi": lon_psi,
         "angle": angle,
         "mask": mask,
-        "mask_u": mask_u,
-        "mask_v": mask_v,
         "straddle": straddle,
     }
 


Test patch:
diff --git a/roms_tools/tests/test_setup/test_utils.py b/roms_tools/tests/test_setup/test_utils.py
index de4faad6..1012b839 100644
--- a/roms_tools/tests/test_setup/test_utils.py
+++ b/roms_tools/tests/test_setup/test_utils.py
@@ -15,6 +15,7 @@ def test_interpolate_from_climatology(use_dask):
 
     climatology = ERA5Correction(use_dask=use_dask)
     field = climatology.ds["ssr_corr"]
+    field["time"] = field["time"].dt.days
 
     interpolated_field = interpolate_from_climatology(field, "time", era5_times)
     assert len(interpolated_field.time) == len(era5_times)