Fix data splitting for S2BMS UC

configs/experiment/prediction.yaml

@@ -0,0 +1,47 @@
+# @package _global_
+
+# to execute this experiment run:
+# python train.py experiment=example
+
+defaults:
+  - override /model: s2bms_prediction
+  - override /data: s2bms_prediction
+  - override /metrics: s2bms_predictive
+
+# all parameters below will be merged with parameters from default configurations set above
+# this allows you to overwrite only specified parameters
+
+model:
+  geo_encoder:
+    _target_: src.models.components.geo_encoders.geoclip.GeoClipCoordinateEncoder
+    geo_data_name: coords
+  trainable_modules: [prediction_head]
+  prediction_head:
+    _target_: src.models.components.pred_heads.linear_pred_head.LinearPredictionHead
+
+data:
+  dataset:
+    modalities:
+      coords:
+    use_target_data: true
+    use_aux_data: false
+  caption_builder:
+  pin_memory: false
+  batch_size: 8
+  persistent_workers: false
+  saved_split_file_name: "s2bms_union_val_test.pth"
+
+tags: ["prediction", "geoclip_coords"]
+
+seed: 12345
+
+trainer:
+  min_epochs: 1
+  max_epochs: 100
+
+logger:
+  wandb:
+    tags: ${tags}
+    group: "predictive"
+  aim:
+    experiment: "predictive"

data/s2bms/splits/explanations.txt

@@ -0,0 +1,6 @@
+split_indices_s2bms-unlabelled-20260529_2026-05-29-1433.pth: all train/val/test unlabelled indices (not used, source data)
+split_indices_s2bms+s2bms-unlabelled-20260529_2026-05-29-1438.pth: original S2BMS split and all train/val/test unlabelled indices (not used, source data)
+split_indices_s2bms_2024-08-14-1459.pth: original S2BMS split (not used, source data)
+
+s2bms_union_val_test: S2BMS labelled subset for this study. All val and test split datapoints are available for modalities of: coords, alphaearth and tessera.
+s2bms_unlabelled_union_val_test: s2bms_union_val_test + all train samples from the unlabelled dataset

notebooks/09-GT-aef-tessera-datasplit-filtering.ipynb

@@ -25,13 +25,11 @@
    "outputs": [],
    "source": [
     "# Check how do aef and tessera files overlap\n",
-    "\n",
-    "paths = glob.glob(\"data/s2bms/eo/aef/*.tif\")\n",
+    "paths = glob.glob(\"data/s2bms/eo/aef/*.npy\")\n",
     "aef = []\n",
     "for p in paths:\n",
     "    aef.append(os.path.basename(p).split(\".\")[0].split(\"-\")[-1])\n",
     "\n",
-    "\n",
     "paths = glob.glob(\"data/s2bms/eo/tessera/*.npy\")\n",
     "tsr = []\n",
     "for p in paths:\n",
@@ -42,6 +40,10 @@
     "tsr_left = list(set(tsr) - set(common))  # None\n",
     "aef_left = list(set(aef) - set(common))  # 12\n",
     "\n",
+    "# Filter out unlabelled data\n",
+    "tsr_left = {i for i in tsr_left if \"sample\" not in i}\n",
+    "aef_left = {i for i in aef_left if \"sample\" not in i}\n",
+    "\n",
     "print(\"Tessera:\", tsr_left)\n",
     "print(\"AlphaEarth:\", aef_left)"
    ]
@@ -98,14 +100,12 @@
     "        print(k, len(v))\n",
     "        s += len(v)\n",
     "\n",
-    "if os.path.exists(\n",
-    "    \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth\"\n",
-    "):\n",
+    "if os.path.exists(\"data/s2bms/splits/s2bms_union_val_test.pth\"):\n",
     "    print(\"Already saved\")\n",
     "else:\n",
     "    torch.save(\n",
     "        split_indices,\n",
-    "        \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test.pth\",\n",
+    "        \"data/s2bms/splits/s2bms_union_val_test.pth\",\n",
     "    )"
    ]
   },
@@ -116,59 +116,182 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "# Filter out remaining aef but missing tessera tiles form train splits\n",
+    "# Visualise data split\n",
+    "df = pd.read_csv(\"data/s2bms/model_ready_s2bms.csv\")\n",
     "\n",
-    "split_indices[\"train_indices\"] = split_indices[\"train_indices\"][\n",
-    "    ~split_indices[\"train_indices\"].isin(aef_left_str)\n",
-    "]\n",
+    "train_idx = split_indices[\"train_indices\"]\n",
+    "train_df = df[df[\"name_loc\"].isin(map(str, train_idx))]\n",
+    "\n",
+    "val_idx = split_indices[\"val_indices\"]\n",
+    "val_df = df[df[\"name_loc\"].isin(map(str, val_idx))]\n",
+    "\n",
+    "test_idx = split_indices[\"test_indices\"]\n",
+    "test_df = df[df[\"name_loc\"].isin(map(str, test_idx))]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "6",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import folium\n",
+    "\n",
+    "m = folium.Map(location=[df.lat.mean(), df.lon.mean()], zoom_start=3)\n",
+    "\n",
+    "# Train = blue\n",
+    "for _, row in train_df.iterrows():\n",
+    "    folium.CircleMarker(location=[row.lat, row.lon], radius=2, color=\"blue\", fill=True).add_to(m)\n",
+    "\n",
+    "# Val = orange\n",
+    "for _, row in val_df.iterrows():\n",
+    "    folium.CircleMarker(location=[row.lat, row.lon], radius=2, color=\"orange\", fill=True).add_to(m)\n",
+    "\n",
+    "# Test = red\n",
+    "for _, row in test_df.iterrows():\n",
+    "    folium.CircleMarker(location=[row.lat, row.lon], radius=2, color=\"red\", fill=True).add_to(m)\n",
+    "\n",
+    "m"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "7",
+   "metadata": {},
+   "source": [
+    "# Re-do Unlabelled split"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "8",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "df_unlabelled = pd.read_csv(\"data/s2bms/model_ready_s2bms-unlabelled-20260529.csv\")\n",
+    "df_unlabelled = df_unlabelled[df_unlabelled.split == \"train\"]\n",
+    "\n",
+    "train_unlabelled = list(df_unlabelled.name_loc)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "9",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "pth = \"data/s2bms/splits/split_indices_s2bms+s2bms-unlabelled-20260529_2026-05-29-1438.pth\"\n",
+    "split_indices = torch.load(pth, weights_only=False)\n",
     "\n",
-    "print(\"After test/val/train cleaning:\")\n",
+    "print(\"Before test/val/train:\")\n",
+    "s = 0\n",
     "for k, v in split_indices.items():\n",
     "    if not k == \"clusters\":\n",
     "        print(k, len(v))\n",
-    "        s += len(v)\n",
-    "\n",
-    "if os.path.exists(\n",
-    "    \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth\"\n",
-    "):\n",
-    "    print(\"Already saved\")\n",
-    "else:\n",
+    "        s += len(v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "10",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "split_indices[\"train_indices\"] = pd.Series(\n",
+    "    [i for i in list(split_indices[\"train_indices\"]) if i in train_unlabelled or \"UKBMS\" in i]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "11",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "val_test_indices = torch.load(\n",
+    "    \"data/s2bms/splits/s2bms_union_val_test.pth\",\n",
+    "    weights_only=False,\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "12",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "split_indices[\"val_indices\"] = pd.Series(\n",
+    "    [i for i in list(val_test_indices[\"val_indices\"]) if \"UKBMS\" in i]\n",
+    ")\n",
+    "split_indices[\"test_indices\"] = pd.Series(\n",
+    "    [i for i in list(val_test_indices[\"test_indices\"]) if \"UKBMS\" in i]\n",
+    ")"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "13",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "print(\"After test/val/train:\")\n",
+    "s = 0\n",
+    "for k, v in split_indices.items():\n",
+    "    if not k == \"clusters\":\n",
+    "        print(k, len(v))\n",
+    "        s += len(v)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "id": "14",
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "if not os.path.exists(\"data/s2bms/splits/s2bms_unlabelled_union_val_test.pth\"):\n",
     "    torch.save(\n",
     "        split_indices,\n",
-    "        \"data/s2bms/splits/split_indices_s2bms_2024-08-14-1459_union_aef_tsr_val_test_train.pth\",\n",
-    "    )"
+    "        \"data/s2bms/splits/s2bms_unlabelled_union_val_test.pth\",\n",
+    "    )\n",
+    "else:\n",
+    "    print(\"Already saved\")"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "6",
+   "id": "15",
    "metadata": {},
    "outputs": [],
    "source": [
     "# Visualise data split\n",
-    "df = pd.read_csv(\"data/s2bms/model_ready_s2bms.csv\")\n",
+    "df_merged = pd.concat([df, df_unlabelled])\n",
     "\n",
     "train_idx = split_indices[\"train_indices\"]\n",
-    "train_df = df[df[\"name_loc\"].isin(map(str, train_idx))]\n",
-    "\n",
+    "train_df = df_merged[df_merged[\"name_loc\"].isin(map(str, train_idx))]\n",
     "\n",
     "val_idx = split_indices[\"val_indices\"]\n",
-    "val_df = df[df[\"name_loc\"].isin(map(str, val_idx))]\n",
+    "val_df = df_merged[df_merged[\"name_loc\"].isin(map(str, val_idx))]\n",
     "\n",
     "test_idx = split_indices[\"test_indices\"]\n",
-    "test_df = df[df[\"name_loc\"].isin(map(str, test_idx))]"
+    "test_df = df_merged[df_merged[\"name_loc\"].isin(map(str, test_idx))]"
    ]
   },
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "7",
+   "id": "16",
    "metadata": {},
    "outputs": [],
    "source": [
-    "import folium\n",
-    "\n",
     "m = folium.Map(location=[df.lat.mean(), df.lon.mean()], zoom_start=3)\n",
     "\n",
     "# Train = blue\n",
@@ -189,13 +312,19 @@
   {
    "cell_type": "code",
    "execution_count": null,
-   "id": "8",
+   "id": "17",
    "metadata": {},
    "outputs": [],
    "source": []
   }
  ],
- "metadata": {},
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  }
+ },
  "nbformat": 4,
  "nbformat_minor": 5
 }

src/data/base_datamodule.py

@@ -2,7 +2,7 @@
 import os
 import time
 from functools import partial
-from typing import Any, Dict, List, Tuple
+from typing import Any, Tuple
 
 import numpy as np
 import pandas as pd
@@ -67,8 +67,6 @@ def __init__(
             self.caption_builder.sync_with_dataset(self.dataset)
             self.concept_configs = caption_builder.concepts
 
-        self.split_data()
-
     @property
     def tabular_dim(self):
         return self.dataset.tabular_dim
@@ -86,6 +84,7 @@ def setup(self, stage: str = "fit") -> None:
 
         # Set up the dataset (download requested modalities)
         self.dataset.setup()
+        self.split_data()
 
     @property
     def batch_size_per_device(self) -> None:
@@ -251,10 +250,19 @@ def split_data(self) -> None:
             if test_indices is not None and not isinstance(test_indices, pd.Series):
                 raise NotImplementedError("Expected a pd series of name_locs for data splits.")
 
-            train_indices = np.where(self.dataset.df["name_loc"].isin(train_indices))[0]
-            val_indices = np.where(self.dataset.df["name_loc"].isin(val_indices))[0]
+            ds_records_names = [i["name_loc"] for i in self.dataset.records]
+            records_name_to_idx = {name: idx for idx, name in enumerate(ds_records_names)}
+
+            train_indices = np.array(
+                [records_name_to_idx[n] for n in train_indices if n in records_name_to_idx]
+            )
+            val_indices = np.array(
+                [records_name_to_idx[n] for n in val_indices if n in records_name_to_idx]
+            )
             if test_indices is not None:
-                test_indices = np.where(self.dataset.df["name_loc"].isin(test_indices))[0]
+                test_indices = np.array(
+                    [records_name_to_idx[n] for n in test_indices if n in records_name_to_idx]
+                )
 
             print(f"Dataset was split using indices from file: {self.saved_split_file_path}")
         else:
@@ -264,13 +272,10 @@ def split_data(self) -> None:
 
         if split_data_from_inds:
             self.data_train = torch.utils.data.Subset(self.dataset, train_indices)
-            self.data_train.dataset.mode = "train"
             self.data_val = torch.utils.data.Subset(self.dataset, val_indices)
-            self.data_val.dataset.mode = "val"
 
             if test_indices is not None:
                 self.data_test = torch.utils.data.Subset(self.dataset, test_indices)
-                self.data_test.dataset.mode = "test"
             else:
                 self.data_test = None