73 water vapor notebook

climanet/predict.py

@@ -124,6 +124,10 @@ def predict_monthly_var(
     model.to(device)
     model.eval()
 
+    # Length of the dataset should >=1
+    if len(dataset) < 1:
+        raise ValueError("The dataset is empty. Please provide a non-empty dataset.")
+
     use_cuda = device == "cuda"
     dataloader = DataLoader(
         dataset,

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/example_daily.ipynb → notebooks/sst/training_daily_sst.ipynb

@@ -182,14 +182,14 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "id": "ca9b6c44-1f4a-4ff2-b15e-7b330bbc501e",
    "metadata": {
     "jp-MarkdownHeadingCollapsed": true
    },
    "outputs": [],
    "source": [
-    "# read data\n",
+    "# read data \n",
     "data_dir = f\"{run_dir}/data_train\"\n",
     "input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features = read_st_data(data_path=data_dir, var_name=var_name)"
    ]
@@ -259,12 +259,12 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 11,
    "id": "d02dace1-c194-49bd-a283-94f7a04e666b",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# read data\n",
+    "# read data \n",
     "data_dir = f\"{run_dir}/data_validation\"\n",
     "input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features = read_st_data(data_path=data_dir, var_name=var_name)"
    ]
@@ -475,19 +475,19 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 18,
    "id": "4c30511a-49fc-4a4d-8bfc-61ec3857fde0",
    "metadata": {},
    "outputs": [],
    "source": [
-    "# read data\n",
+    "# read data \n",
     "data_dir = f\"{run_dir}/data_test\"\n",
     "input_da, input_da_nan_mask, monthly_da, padded_days_mask, time_features = read_st_data(data_path=data_dir, var_name=var_name)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 20,
    "id": "ac99fdf5-dfd8-44a0-957b-dea177ac92a1",
    "metadata": {},
    "outputs": [
@@ -503,7 +503,7 @@
     }
    ],
    "source": [
-    "spatial_crop_size = monthly_da.shape[1:]  # the whole dataset\n",
+    "spatial_patch_size = monthly_da.shape[1:]  # the whole dataset \n",
     "\n",
     "dataset_test = STDataset(\n",
     "    input_da=input_da,\n",
@@ -513,7 +513,7 @@
     "    time_features=time_features,\n",
     "    land_mask=lsm_subset[\"lsm\"],\n",
     "    model_patch_size=patch_size,\n",
-    "    crop_size=(1, *spatial_crop_size),\n",
+    "    crop_size=(1, *spatial_crop_size), \n",
     "    stride=None,  # no stride in inference\n",
     "    sh_embed_dim=96,\n",
     "    sh_order_L = 10,\n",
@@ -574,7 +574,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 23,
    "id": "7c2deb40-bee8-4973-80f0-9d9485eabf0c",
    "metadata": {},
    "outputs": [
@@ -591,7 +591,7 @@
     "# inference on test data, verbose is True\n",
     "predictions = predict_monthly_var(\n",
     "    model=f\"{run_dir}/best_model.pth\",\n",
-    "    dataset=dataset_test,\n",
+    "    dataset=dataset_test,  \n",
     "    dataloader_config=dataloader_config,\n",
     "    prediction_config=prediction_config,\n",
     "    run_dir=run_dir,\n",
@@ -647,7 +647,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 27,
    "id": "7ba5d204-aefc-4103-8475-c802c603ad3f",
    "metadata": {},
    "outputs": [
@@ -668,7 +668,7 @@
     "\n",
     "ocean = ~lsm_subset[\"lsm\"].values\n",
     "masked_err_baseline = err_baseline.where(ocean)\n",
-    "masked_err_predictions = err_predictions.where(ocean)\n",
+    "masked_err_predictions = err_predictions.where(ocean) \n",
     "\n",
     "plot_results(masked_err_baseline, masked_err_predictions, label=\"error K\", title=(\"err_Baseline\", \"err_Prediction\"), error=True)"
    ]

notebooks/example_hourly.ipynb → notebooks/sst/training_hourly_sst.ipynb

@@ -670,9 +670,9 @@
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "climanet",
+   "display_name": "Python 3 (ipykernel)",
    "language": "python",
-   "name": "climanet"
+   "name": "python3"
   },
   "language_info": {
    "codemirror_mode": {
@@ -684,7 +684,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.14.6"
+   "version": "3.11.14"
   }
  },
  "nbformat": 4,