95 watervapor lsm util func

climanet/utils.py

@@ -789,7 +789,9 @@ def data_preparation(
 
     # rechunk data
     input_da = input_da.chunk({"M": 1, "T": -1, "lat": 100, "lon": 100})
-    input_da_nan_mask = input_da_nan_mask.chunk({"M": 1, "T": -1, "lat": 100, "lon": 100})
+    input_da_nan_mask = input_da_nan_mask.chunk(
+        {"M": 1, "T": -1, "lat": 100, "lon": 100}
+    )
     monthly_da = monthly_da.chunk({"M": 1, "lat": 100, "lon": 100})
     padded_days_mask = padded_days_mask.chunk({"M": 1})
     time_features = time_features.chunk({"M": 1})
@@ -891,3 +893,29 @@ def read_st_data(data_path=".", var_name="tos"):
 
     # if one of the datasets is None, we need to compute them
     return input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features
+
+
+def coarsen_land_mask(
+    input_lsm: xr.DataArray | xr.Dataset, coarse_factor: int = 2, threshold: float = 0.5
+):
+    """Coarsen spatial resolution of land mask data by coarse_factor.
+
+    It also applies a threshold to values outside of [0,1].
+    see https://confluence.ecmwf.int/spaces/FUG/pages/673550380/Section+2A.1.3.1+Land-Sea+mask
+
+    Args:
+        input_lsm (xarray.DataArray): Land-sea mask from ERA5-Land data.
+        coarse_factor (int, optional): Factor by which to coarsen the resolution. Defaults to 2.
+        threshold (float, optional): Threshold for determining mask value. Defaults to 0.5.
+
+    Returns:
+        xarray.DataArray | xarray.Dataset : Coarse land-sea mask.
+    """
+    coarse_lsm = input_lsm.coarsen(
+        lat=coarse_factor, lon=coarse_factor, boundary="trim"
+    ).mean()
+
+    # Apply threshold
+    coarse_lsm = coarse_lsm >= threshold
+
+    return coarse_lsm

notebooks/watervapor/training_hourly_watervapor.ipynb

@@ -2,7 +2,7 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": null,
    "id": "7e6644bb-2d85-47ec-af53-f0960f9314e6",
    "metadata": {},
    "outputs": [],
@@ -12,7 +12,7 @@
     "import torch\n",
     "import torch.nn.functional\n",
     "from climanet.st_encoder_decoder import SpatioTemporalModel\n",
-    "from climanet.utils import set_seed, configure_compute_resources, plot_results, plot_histograms, plot_loss, data_preparation, read_st_data, add_month_day_dims\n",
+    "from climanet.utils import set_seed, configure_compute_resources, plot_results, plot_histograms, plot_loss, data_preparation, read_st_data, add_month_day_dims, coarsen_land_mask\n",
     "from climanet.train import train_monthly_model, TrainConfig\n",
     "from climanet.predict import predict_monthly_var, PredictionConfig\n",
     "from climanet.dataset import STDataset, DataLoaderConfig\n",
@@ -32,28 +32,6 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "5bdd29f8",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "# # Make land-sea mask at 0.5 degree resolution from 0.25 degree resolution used in SST\n",
-    "\n",
-    "# # down sample lsm_mask_025 by a factor of 2 by averaging\n",
-    "# lsm_mask_025 = xr.open_dataset(\"../eso4clima/sst/era5_lsm_bool.nc\")\n",
-    "# lsm_mask_05 = lsm_mask_025.coarsen(lat=2, lon=2, boundary='trim').mean()\n",
-    "\n",
-    "# # make >=0.5 to 1, and <0.5 to 0\n",
-    "# lsm_mask_05 = lsm_mask_05.where(lsm_mask_05 >= 0.5, 0)\n",
-    "# lsm_mask_05 = lsm_mask_05.where(lsm_mask_05 < 0.5, 1)\n",
-    "# lsm_mask_05['lsm'] = lsm_mask_05['lsm'].astype(bool)\n",
-    "\n",
-    "# lsm_mask_05.to_netcdf(\"./era5_lsm_bool_05.nc\")\n",
-    "# lsm_mask_05.isel(time=0)['lsm'].plot()"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 2,
    "id": "f383c4c8-0199-4092-be12-07be525fdec1",
    "metadata": {
     "scrolled": true
@@ -74,9 +52,11 @@
     "monthly_data_validation = xr.open_mfdataset(data_folder / f\"202104_mon_ERA5_full_{var_name}.nc\")\n",
     "monthly_data_test = xr.open_mfdataset(data_folder / f\"202204_mon_ERA5_full_{var_name}.nc\")\n",
     "\n",
-    "# water vapor uses a dedicated coarser-resolution land-sea mask\n",
-    "file_name = \"./era5_lsm_bool_05.nc\"\n",
-    "lsm_mask = xr.open_dataset(file_name)"
+    "# Make land-sea mask for watervapor from sst mask\n",
+    "file_name = data_folder / \"era5_lsm_bool.nc\"  # downloded from era5 and regridded using the function `regrid_to_boundary_centered_grid`\n",
+    "lsm_mask_sst = xr.open_dataset(file_name)\n",
+    "lsm_mask = coarsen_land_mask(lsm_mask_sst)\n",
+    "lsm_mask"
    ]
   },
   {
@@ -725,9 +705,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "CLIMANET Kernel",
+   "display_name": "ClimaNet (3.14.0.final.0)",
    "language": "python",
-   "name": "my-kernel"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -739,7 +719,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.14.3"
+   "version": "3.14.0"
   }
  },
  "nbformat": 4,

tests/test_utils.py

@@ -2,7 +2,7 @@
 import xarray as xr
 from tbparse import SummaryReader
 
-from climanet.utils import data_preparation, setup_logging
+from climanet.utils import coarsen_land_mask, data_preparation, setup_logging
 
 
 def test_setup_logging(tmp_path):
@@ -22,7 +22,6 @@ def test_setup_logging(tmp_path):
     assert reader.scalars["value"].iloc[0] == 42  # check scalar
 
 
-
 def _make_datasets():
     # 4x4 dataset with, 6 days in one month
     lat = 4
@@ -67,12 +66,14 @@ def _make_datasets():
 def test_data_preparation():
     daily_da, monthly_da = _make_datasets()
 
-    input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features = data_preparation(
-        daily_da, monthly_da, calculate_residuals=False, save_to_zarr=False
+    input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features = (
+        data_preparation(
+            daily_da, monthly_da, calculate_residuals=False, save_to_zarr=False
+        )
     )
     assert input_da.shape == (1, 31, 4, 4)  # (M, T=31, H=4, W=4)
     assert isinstance(input_da, xr.DataArray)
-    assert input_da_nan_mask[0, 1, 1, 1] == True  # check that the NaN mask is correctly set
+    assert input_da_nan_mask[0, 1, 1, 1]  # check that the NaN mask is correctly set
     assert isinstance(input_da_nan_mask, xr.DataArray)
     assert monthly_da.shape == (1, 4, 4)  # (M, H=4, W=4)
     assert isinstance(monthly_da, xr.DataArray)
@@ -86,7 +87,11 @@ def test_data_preparation_to_zarr(tmp_path):
     daily_da, monthly_da = _make_datasets()
 
     _ = data_preparation(
-        daily_da, monthly_da, run_dir=tmp_path, calculate_residuals=False, save_to_zarr=True
+        daily_da,
+        monthly_da,
+        run_dir=tmp_path,
+        calculate_residuals=False,
+        save_to_zarr=True,
     )
     assert (tmp_path / "input_da.zarr").exists()
     assert (tmp_path / "input_da_nan_mask.zarr").exists()
@@ -99,7 +104,11 @@ def test_data_preparation_from_zarr(tmp_path):
     daily_da, monthly_da = _make_datasets()
 
     _ = data_preparation(
-        daily_da, monthly_da, run_dir=tmp_path, calculate_residuals=False, save_to_zarr=True
+        daily_da,
+        monthly_da,
+        run_dir=tmp_path,
+        calculate_residuals=False,
+        save_to_zarr=True,
     )
 
     # Now load from zarr
@@ -111,11 +120,44 @@ def test_data_preparation_from_zarr(tmp_path):
 
     assert input_da["tos"].shape == (1, 31, 4, 4)  # (M, T=31, H=4, W=4)
     assert isinstance(input_da, xr.Dataset)
-    assert input_da_nan_mask["tos"][0, 1, 1, 1] == True  # check that the NaN mask is correctly set
+    assert input_da_nan_mask["tos"][
+        0, 1, 1, 1
+    ]  # check that the NaN mask is correctly set
     assert isinstance(input_da_nan_mask, xr.Dataset)
     assert monthly_da["tos"].shape == (1, 4, 4)  # (M, H=4, W=4)
     assert isinstance(monthly_da, xr.Dataset)
     assert padded_days_mask["tos"].shape == (1, 31)  # (M, T=31)
     assert isinstance(padded_days_mask, xr.Dataset)
-    assert time_features["tos"].shape == (1, 31, 3)  # (M, T=31, 2) for month and day features
+    assert time_features["tos"].shape == (
+        1,
+        31,
+        3,
+    )  # (M, T=31, 2) for month and day features
     assert isinstance(time_features, xr.Dataset)
+
+
+def test_coarsen_land_mask():
+    """Downsample a mock lsm with a factor of 2."""
+    mask_values = np.zeros((1, 8, 4))  # Mock values for the land-sea mask
+    mask_values[0, 0:2, 0:2] = 0.7  # Set 4 values to be above 0.5
+    mask_values[0, 2:4, 2:4] = 0.2  # Set another 4 values to be below 0.5
+    lsm_mask_025_mock = xr.DataArray(
+        mask_values,
+        dims=["time", "lat", "lon"],
+        coords={
+            "time": [0],
+            "lat": np.linspace(-2.0, 2.0, 8),
+            "lon": np.linspace(-1.0, 1.0, 4),
+        },
+        name="lsm",
+    )
+    lsm_mask_05 = coarsen_land_mask(lsm_mask_025_mock)
+
+    assert np.all(
+        lsm_mask_05.values[0, 0:1, 0:1]  # after downsampling 0:2 -> 0:1
+    )  # all values above 0.5 should be True
+    assert not np.any(
+        lsm_mask_05.values[0, 1:2, 1:2]  # after downsampling 2:4 -> 1:2
+    )  # all values below 0.5 should be False
+    assert lsm_mask_05.shape == (1, 4, 2)  # By default downsample with a factor of 2
+    assert lsm_mask_05.dtype == bool  # Mask should be of boolean type