From d92d6d72ab14b9261e533757230ee78839db0667 Mon Sep 17 00:00:00 2001 From: Rob Knapen Date: Thu, 7 May 2026 17:17:53 +0200 Subject: [PATCH 1/2] Adds support for dual-year tessera fusion for the crop yield use case --- src/data/base_datamodule.py | 107 ++++--- src/data/yield_africa_dataset.py | 85 +++++- .../yield_africa_pipeline_tui.py | 169 ++++++++--- .../yield_africa_tessera_preprocess.py | 285 +++++++++++++----- .../yield_africa_tessera_rename.py | 115 +++++++ src/models/base_model.py | 8 + .../geo_encoders/average_encoder.py | 8 +- .../geo_encoders/tabular_encoder.py | 3 + tests/test_dynamic_gate_fusion_encoder.py | 61 ++++ .../test_yield_africa_dual_tessera_dataset.py | 255 ++++++++++++++++ 10 files changed, 931 insertions(+), 165 deletions(-) create mode 100644 src/data_preprocessing/yield_africa_tessera_rename.py create mode 100644 tests/test_yield_africa_dual_tessera_dataset.py diff --git a/src/data/base_datamodule.py b/src/data/base_datamodule.py index 6d1147a..9b4ef50 100644 --- a/src/data/base_datamodule.py +++ b/src/data/base_datamodule.py @@ -1,6 +1,4 @@ -import copy import os -import time from functools import partial from typing import Any, Dict, List, Tuple @@ -8,7 +6,6 @@ import pandas as pd import torch from lightning import LightningDataModule -from sklearn.cluster import DBSCAN from sklearn.model_selection import GroupShuffleSplit from torch.utils.data import DataLoader, random_split @@ -48,8 +45,9 @@ def __init__( :param save_split: if to save split file :param saved_split_file_name: file name to save split file :param caption_builder: instance of BaseCaptionBuilder for generating textual captions - :param spatial_split_distance_m: minimum distance in metres between clusters when - split_mode is 'spatial_clusters'. Default 1000 m. + :param spatial_split_distance_m: grid cell size in metres when split_mode is + 'spatial_clusters'. Samples within the same cell are kept together and assigned to the + same split. Default 1000 m. """ super().__init__() self.save_hyperparameters(logger=False) @@ -122,33 +120,43 @@ def split_data(self) -> None: elif self.hparams.split_mode == "spatial_clusters": min_dist = self.hparams.spatial_split_distance_m - coords = np.array([self.dataset.df.lat, self.dataset.df.lon]).T + # Use records (not df): records is already filtered (e.g. missing tiles + # dropped), so len(records) <= len(df). Indices must be into records + # because __getitem__ and __len__ both operate on self.records. + # lat/lon come from df (always present) keyed by name_loc so the + # coordinate array stays aligned with records regardless of modalities. + records = self.dataset.records + _nl_to_coords = dict( + zip( + self.dataset.df["name_loc"], + zip(self.dataset.df["lat"], self.dataset.df["lon"]), + ) + ) + coords = np.array( + [ + [_nl_to_coords[r["name_loc"]][0] for r in records], + [_nl_to_coords[r["name_loc"]][1] for r in records], + ] + ).T n = len(coords) + # Grid-based spatial grouping: assign each sample to a geographic + # cell of size spatial_split_distance_m × spatial_split_distance_m. + # GroupShuffleSplit then distributes whole cells across splits, so + # geographically close samples stay together while split proportions + # remain balanced (unlike DBSCAN, which chain-links dense data into + # a few giant clusters and produces wildly uneven splits). + _METERS_PER_DEG_LAT = 111_000.0 + lat_step = min_dist / _METERS_PER_DEG_LAT + lon_step = min_dist / (_METERS_PER_DEG_LAT * np.cos(np.radians(np.mean(coords[:, 0])))) + grid_ids = np.floor(coords[:, 0] / lat_step).astype(np.int64) * 1_000_000 + np.floor( + coords[:, 1] / lon_step + ).astype(np.int64) + _, clusters = np.unique(grid_ids, return_inverse=True) + n_clusters = int(clusters.max()) + 1 print( - f"Splitting {n} samples into spatial clusters " - f"(eps={min_dist / 1000:.1f} km, haversine, n_jobs=-1)..." + f"Splitting {n} samples into {n_clusters} spatial grid cells " + f"(cell size ≈ {min_dist / 1000:.0f} km). Creating splits..." ) - # Convert (lat, lon) degrees to radians for sklearn's haversine metric. - # haversine returns arc length on the unit sphere, so eps must be in radians. - _EARTH_RADIUS_M = 6_371_000 - coords_rad = np.radians(coords) - eps_rad = min_dist / _EARTH_RADIUS_M - t0 = time.time() - clustering = DBSCAN( - eps=eps_rad, - metric="haversine", - algorithm="ball_tree", - min_samples=2, - n_jobs=-1, - ).fit(coords_rad) - print(f"DBSCAN done in {time.time() - t0:.1f}s. Creating splits...") - # Non-clustered points are labeled -1. Change to new cluster label. - clusters = copy.deepcopy(clustering.labels_) - new_cl = np.max(clusters) + 1 - for i, cl in enumerate(clusters): - if cl == -1: - clusters[i] = new_cl - new_cl += 1 gss = GroupShuffleSplit( n_splits=1, @@ -197,13 +205,15 @@ def split_data(self) -> None: ) print( - f"Created {len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} test indices using DBSCAN spatial clustering with {min_dist} m minimum distance between clusters." + f"Created {len(train_indices)} train, {len(val_indices)} val, " + f"{len(test_indices)} test indices across {n_clusters} spatial grid cells " + f"(cell size ≈ {min_dist / 1000:.0f} km)." ) if self.hparams.save_split: split_indices = { - "train_indices": self.dataset.df.name_loc[train_indices], - "val_indices": self.dataset.df.name_loc[val_indices], - "test_indices": self.dataset.df.name_loc[test_indices], + "train_indices": pd.Series([records[i]["name_loc"] for i in train_indices]), + "val_indices": pd.Series([records[i]["name_loc"] for i in val_indices]), + "test_indices": pd.Series([records[i]["name_loc"] for i in test_indices]), "clusters": clusters, } @@ -231,10 +241,21 @@ 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] + # Map name_locs → records-level indices (not df row indices). + # self.records may be shorter than self.df when records are + # dropped (e.g. missing tessera_prev tiles in config B), so + # df row indices would be out of range in __getitem__. + _name_loc_to_rec_idx = {r["name_loc"]: i for i, r in enumerate(self.dataset.records)} + train_indices = np.array( + [_name_loc_to_rec_idx[nl] for nl in train_indices if nl in _name_loc_to_rec_idx] + ) + val_indices = np.array( + [_name_loc_to_rec_idx[nl] for nl in val_indices if nl in _name_loc_to_rec_idx] + ) if test_indices is not None: - test_indices = np.where(self.dataset.df["name_loc"].isin(test_indices))[0] + test_indices = np.array( + [_name_loc_to_rec_idx[nl] for nl in test_indices if nl in _name_loc_to_rec_idx] + ) print(f"Dataset was split using indices from file: {self.saved_split_file_path}") else: @@ -277,11 +298,14 @@ def _compute_tabular_normalisation_stats(self) -> None: return train_indices = self.data_train.indices - train_df = self.dataset.df.iloc[train_indices][feat_names] + train_df = pd.DataFrame( + [[self.dataset.records[i][k] for k in feat_names] for i in train_indices], + columns=feat_names, + ) mean = train_df.mean(axis=0).values std = train_df.std(axis=0).values - std = np.where(std == 0, 1.0, std) # avoid division by zero for constant features + std = np.where((std == 0) | np.isnan(std), 1.0, std) self.tabular_normalisation_stats = ( torch.tensor(mean, dtype=torch.float32), @@ -304,11 +328,14 @@ def _compute_target_normalisation_stats(self) -> None: return train_indices = self.data_train.indices - train_df = self.dataset.df.iloc[train_indices][target_names] + train_df = pd.DataFrame( + [[self.dataset.records[i][k] for k in target_names] for i in train_indices], + columns=target_names, + ) mean = train_df.mean(axis=0).values std = train_df.std(axis=0).values - std = np.where(std == 0, 1.0, std) # avoid division by zero for constant targets + std = np.where((std == 0) | np.isnan(std), 1.0, std) self.target_normalisation_stats = ( torch.tensor(mean, dtype=torch.float32), diff --git a/src/data/yield_africa_dataset.py b/src/data/yield_africa_dataset.py index 839179d..ad09555 100644 --- a/src/data/yield_africa_dataset.py +++ b/src/data/yield_africa_dataset.py @@ -52,9 +52,12 @@ class YieldAfricaDataset(BaseDataset): Modality design note -------------------- - `implemented_mod = {"coords"}` because tabular features live directly in - the model-ready CSV and are picked up via the `feat_` column prefix. - They do NOT need to be listed in `modalities`. + Tabular features live directly in the model-ready CSV and are picked up + via the `feat_` column prefix. They do NOT need to be listed in + `modalities`. Implemented spatial modalities: ``coords``, ``tessera`` + (year-Y embedding), ``tessera_prev`` (year-Y−1 embedding for dual-year + fusion). Adding ``tessera_prev`` to modalities activates dual-year + loading; single-year runs are unaffected. In addition to the CSV feat_* columns, the following features are injected: - ``feat_year`` : normalised year (zero-mean, unit-std) @@ -88,6 +91,7 @@ def __init__( years: List[int] | None = None, exclude_countries: List[str] | None = None, exclude_years: List[int] | None = None, + require_prev_year_tessera: bool = True, ) -> None: super().__init__( data_dir=data_dir, @@ -97,11 +101,12 @@ def __init__( dataset_name="yield_africa", seed=seed, cache_dir=cache_dir, - implemented_mod={"coords", "tessera"}, + implemented_mod={"coords", "tessera", "tessera_prev"}, mock=mock, use_features=use_features, csv_name=csv_name, ) + self.require_prev_year_tessera = require_prev_year_tessera # Inject year and country one-hot columns as feat_* so that # get_records() picks them up automatically. Build all new columns in @@ -144,6 +149,23 @@ def __init__( self.df = pd.concat([self.df, pd.DataFrame(new_cols, index=self.df.index)], axis=1) + # Build a cross-year tessera path index from the full unfiltered df. + # Must happen before the country/year filter below so that year-Y records + # can resolve year-Y−1 paths even when those rows are excluded by a + # years= filter. Keys: (lat_rounded, lon_rounded, year) → path. + if "tessera_prev" in self.modalities: + _tessera_dir_full = os.path.join(self.data_dir, "eo", "tessera") + _year_path_index: dict[tuple[float, float, int], str] = {} + _name_loc_coords: dict[str, tuple[float, float, int]] = {} + for _, _r in self.df.iterrows(): + _lat_r = round(float(_r["lat"]), 6) + _lon_r = round(float(_r["lon"]), 6) + _year_r = int(_r["year"]) + _year_path_index[(_lat_r, _lon_r, _year_r)] = os.path.join( + _tessera_dir_full, f"tessera_{_r['name_loc']}_{_year_r}.npy" + ) + _name_loc_coords[_r["name_loc"]] = (_lat_r, _lon_r, _year_r) + # Apply country/year filters to self.df and rebuild records. # BaseDataset.__init__ has already loaded the CSV; filtering here avoids # touching BaseDataset and keeps the logic use-case specific. @@ -180,6 +202,21 @@ def __init__( # self.feat_names and self.tabular_dim. self.records = self.get_records() + # Rewrite tessera paths to the year-suffixed convention + # (tessera_{name_loc}_{year}.npy). BaseDataset.add_modality_paths_to_df() + # generates paths without a year; this override is local to + # YieldAfricaDataset and leaves BaseDataset unchanged. + if "tessera" in self.modalities: + _tessera_dir = os.path.join(self.data_dir, "eo", "tessera") + _name_loc_to_year: dict[str, int] = dict( + zip(self.df["name_loc"], self.df["year"].astype(int)) + ) + for rec in self.records: + year = _name_loc_to_year[rec["name_loc"]] + rec["tessera_path"] = os.path.join( + _tessera_dir, f"tessera_{rec['name_loc']}_{year}.npy" + ) + # Drop records whose TESSERA tile is absent so the model is never # trained or evaluated on zero-padded stand-ins. if "tessera" in self.modalities: @@ -193,6 +230,38 @@ def __init__( before, ) + # Resolve tessera_prev_path for each record using the cross-year index. + # Records whose year-1 tile is absent are dropped when + # require_prev_year_tessera=True (default), or retained with + # tessera_prev_path=None when False. + if "tessera_prev" in self.modalities: + resolved = [] + for rec in self.records: + lat_r, lon_r, year_r = _name_loc_coords[rec["name_loc"]] + key = (lat_r, lon_r, year_r - 1) + prev_path = _year_path_index.get(key) + if prev_path is not None and os.path.exists(prev_path): + resolved.append({**rec, "tessera_prev_path": prev_path}) + else: + # Fall back to synthetic tile produced by --include-prev-year: + # tessera_{name_loc}_prev_{year-1}.npy + synth_path = os.path.join( + _tessera_dir_full, + f"tessera_{rec['name_loc']}_prev_{year_r - 1}.npy", + ) + if os.path.exists(synth_path): + resolved.append({**rec, "tessera_prev_path": synth_path}) + elif not self.require_prev_year_tessera: + resolved.append({**rec, "tessera_prev_path": None}) + dropped = len(self.records) - len(resolved) + if dropped: + log.warning( + "Dropped %d/%d records: no year-1 TESSERA tile found.", + dropped, + len(self.records), + ) + self.records = resolved + def setup(self) -> None: """Check for requested modality data; warn if TESSERA tiles are absent. @@ -204,13 +273,14 @@ def setup(self) -> None: single fixed year for bulk download, which is incompatible with the multi-year nature of this dataset. """ - if "tessera" in self.modalities: + if "tessera" in self.modalities or "tessera_prev" in self.modalities: tessera_dir = os.path.join(self.data_dir, "eo", "tessera") if not os.path.exists(tessera_dir) or len(os.listdir(tessera_dir)) == 0: log.warning( "TESSERA tiles not found at %s. " "Run src/data_preprocessing/yield_africa_tessera_preprocess.py " - "to pre-fetch tiles. Records with missing tiles are excluded from the dataset.", + "to pre-fetch tiles. For tessera_prev, also pass --include-prev-year. " + "Records with missing tiles are excluded from the dataset.", tessera_dir, ) @@ -226,6 +296,9 @@ def __getitem__(self, idx: int) -> Dict[str, Any]: ) elif modality == "tessera": sample["eo"]["tessera"] = self.load_tessera(row["tessera_path"]) + elif modality == "tessera_prev": + if row.get("tessera_prev_path") is not None: + sample["eo"]["tessera_prev"] = self.load_tessera(row["tessera_prev_path"]) if self.use_features and self.feat_names: sample["eo"]["tabular"] = torch.tensor( diff --git a/src/data_preprocessing/yield_africa_pipeline_tui.py b/src/data_preprocessing/yield_africa_pipeline_tui.py index e4e6440..2ce4542 100644 --- a/src/data_preprocessing/yield_africa_pipeline_tui.py +++ b/src/data_preprocessing/yield_africa_pipeline_tui.py @@ -280,28 +280,63 @@ class PipelineStep: "you can safely interrupt and re-run. Use more --workers for speed." ), ), + PipelineStep( + id="tessera_rename", + title="2c-migrate. Rename TESSERA tiles (one-time)", + description=( + "One-time migration: renames existing tessera_{name_loc}.npy files\n" + "to tessera_{name_loc}_{year}.npy by looking up each record's year\n" + "from the model-ready CSV. Must be run once on an existing tile\n" + "directory before using the year-suffixed naming convention introduced\n" + "by the dual-year tessera feature.\n\n" + "Dry-run is the default — no files are touched until --no-dry-run is\n" + "passed. Safe to re-run: already-renamed files are skipped." + ), + script="src/data_preprocessing/yield_africa_tessera_rename.py", + required=False, + depends_on=["make_ready"], + args_hint=( + "--data_dir data/\n" + "[--no-dry-run] # omit to preview only; review output before applying" + ), + output_hint=( + "data/yield_africa/eo/tessera/tessera_{name_loc}_{year}.npy\n" + "(files renamed in place — no copies written, instant on most filesystems)" + ), + advice=( + "Run this ONCE before the tessera download step (2c) if you already\n" + "have tiles on disk in the old tessera_{name_loc}.npy format.\n" + "New installations that have no tiles yet can skip this step.\n\n" + "Always run without --no-dry-run first to review what would change.\n" + "Then re-run with --no-dry-run to apply the renames.\n" + "Files not found in the model-ready CSV are left untouched with a warning." + ), + ), PipelineStep( id="tessera", title="2c. Download TESSERA embeddings", description=( "Downloads per-record, year-specific satellite embeddings from the\n" "GeoTessera service. Each record gets a small EO tile (default 9 px)\n" - "that captures local land-cover and phenology at the plot location." + "that captures local land-cover and phenology at the plot location.\n" + "With --include-prev-year, also fetches year-1 tiles needed for\n" + "dual-year tessera fusion (configs/data/yield_africa_tessera_dual.yaml)." ), script="src/data_preprocessing/yield_africa_tessera_preprocess.py", required=False, depends_on=["make_ready"], args_hint=( "--data_dir data/\n" - "[--countries KEN RWA] # subset of countries\n" - "[--years 2019 2020] # subset of years\n" - "[--tile_size 9] # pixels around plot centre\n" - "[--workers 1] # parallel download threads\n" - "[--retry-stuck] # retry records that stalled previously" + "[--countries KEN RWA] # subset of countries\n" + "[--years 2019 2020] # subset of years\n" + "[--tile_size 9] # pixels around plot centre\n" + "[--workers 1] # parallel download threads\n" + "[--retry-stuck] # retry records that stalled previously\n" + "[--include-prev-year] # also fetch year-1 tiles for dual-year fusion" ), output_hint=( - "data/yield_africa/eo/tessera/tessera_{name_loc}.npy\n" - "(one NumPy array per plot)\n" + "data/yield_africa/eo/tessera/tessera_{name_loc}_{year}.npy\n" + "(one NumPy array per plot, year-suffixed)\n" "data/yield_africa/eo/tessera/stuck.txt (records that stalled)" ), extra_deps=["geotessera (install with: uv sync --extra geotessera)"], @@ -312,7 +347,14 @@ class PipelineStep: "enough disk space (can be on an external drive).\n" "This step is resumable — existing .npy files are skipped on rerun.\n" "Records whose download stalls are written to stuck.txt and skipped\n" - "on subsequent runs. Use --retry-stuck to attempt them again." + "on subsequent runs. Use --retry-stuck to attempt them again.\n\n" + "If upgrading from an earlier version that used tessera_{name_loc}.npy\n" + "naming, run step 2c-migrate (tessera_rename) first to rename existing\n" + "tiles before downloading new ones.\n\n" + "Use --include-prev-year to also fetch year-1 tiles. These are needed\n" + "for the dual-year tessera experiment configs (yield_africa_tessera_dual).\n" + "Locations whose year-1 tile was already fetched (multi-year survey sites)\n" + "are skipped automatically." ), ), PipelineStep( @@ -458,17 +500,38 @@ class PipelineStep: STEP_INDEX = {s.id: s for s in STEPS} EXPERIMENTS = [ - ("yield_africa_coords_reg", "Coordinates only (baseline)"), - ("yield_africa_tabular_reg", "Tabular soil/climate features"), - ("yield_africa_fusion_reg", "Tabular + coordinate fusion"), - ("yield_africa_tabular_spatial", "Tabular + spatial split"), - ("yield_africa_fusion_spatial", "Fusion + spatial split"), - ("yield_africa_tabular_loco", "Tabular + LOCO split"), - ("yield_africa_fusion_loco", "Fusion + LOCO split"), - ("yield_africa_tessera_reg", "TESSERA embeddings"), - ("yield_africa_tessera_fusion_reg", "TESSERA + tabular fusion"), - ("yield_africa_tessera_fusion_spatial", "TESSERA fusion + spatial split"), - ("yield_africa_tessera_fusion_loco", "TESSERA fusion + LOCO split"), + # Tabular-only baselines (no EO data) + ("yield_africa_tabular_reg", "Tabular features only (all countries)"), + ("yield_africa_tabular_spatial", "Tabular only + spatial-cluster split"), + ("yield_africa_tabular_loco", "Tabular only + leave-one-country-out split"), + # TESSERA spatial embeddings — current year + ("yield_africa_tessera_reg", "TESSERA embeddings only (no tabular)"), + ("yield_africa_tessera_fusion_reg", "TESSERA + tabular fusion (all countries)"), + ("yield_africa_tessera_fusion_spatial", "TESSERA + tabular + spatial-cluster split"), + ("yield_africa_tessera_fusion_loco", "TESSERA + tabular + leave-one-country-out split"), + # Dual-year TESSERA (year Y + year Y−1) — requires prev-year tiles + ( + "yield_africa_tessera_dual_fusion_reg", + "Dual-year TESSERA + tabular, dynamic gate (all countries)", + ), + # Ablation study — fusion strategies for dual-year TESSERA (Kenya only) + ("yield_africa_ablation_dual_A_tessera_KEN", "Ablation A: tessera year-Y + tabular [KEN]"), + ("yield_africa_ablation_dual_B_prev_KEN", "Ablation B: tessera year-Y−1 + tabular [KEN]"), + ( + "yield_africa_ablation_dual_C_concat_KEN", + "Ablation C: dual-year TESSERA + tabular, concat [KEN]", + ), + ( + "yield_africa_ablation_dual_D_gated_KEN", + "Ablation D: dual-year TESSERA + tabular, gated [KEN]", + ), + ( + "yield_africa_ablation_dual_E_dynamic_KEN", + "Ablation E: dual-year TESSERA + tabular, dynamic gate [KEN]", + ), + # Hyperparameter validation — anchor configs for HP search (Kenya only) + ("yield_africa_val_TO_1_baseline_KEN", "Val TO-1: tessera-only baseline [KEN]"), + ("yield_africa_val_FM_1_baseline_KEN", "Val FM-1: full model baseline [KEN]"), ] # --------------------------------------------------------------------------- @@ -820,7 +883,7 @@ def show_experiments() -> None: console.print() console.print( "[dim]Tip: override the split file on the command line, e.g.:\n" - " python src/train.py experiment=yield_africa_fusion_spatial \\\n" + " python src/train.py experiment=yield_africa_tessera_fusion_spatial \\\n" " data.saved_split_file_name=split_spatial_25km.pth[/]" ) @@ -920,12 +983,13 @@ def run_step_prompt(step: PipelineStep) -> None: " [bold cyan]4[/] Step 1 — Build model-ready CSV [bold yellow](required)[/]\n" " [bold cyan]5[/] Step 2a — Augment with NDVI\n" " [bold cyan]6[/] Step 2b — Augment with AgERA5 climate\n" - " [bold cyan]7[/] Step 2c — Download TESSERA embeddings\n" - " [bold cyan]8[/] Step 2d — Merge augmented CSVs [dim](optional)[/]\n" - " [bold cyan]9[/] Step 2e — Compare all augmentations [dim](optional)[/]\n" - " [bold cyan]a[/] Step 3a — Generate spatial splits\n" - " [bold cyan]b[/] Step 3b — Generate LOCO splits\n" - " [bold cyan]c[/] Training experiments reference\n" + " [bold cyan]7[/] Step 2c-migrate — Rename TESSERA tiles [dim](one-time migration)[/]\n" + " [bold cyan]8[/] Step 2c — Download TESSERA embeddings\n" + " [bold cyan]9[/] Step 2d — Merge augmented CSVs [dim](optional)[/]\n" + " [bold cyan]a[/] Step 2e — Compare all augmentations [dim](optional)[/]\n" + " [bold cyan]b[/] Step 3a — Generate spatial splits\n" + " [bold cyan]c[/] Step 3b — Generate LOCO splits\n" + " [bold cyan]d[/] Training experiments reference\n" " [bold cyan]q[/] Quit", title="[bold cyan]Menu[/]", border_style="cyan", @@ -936,7 +1000,7 @@ def run_step_prompt(step: PipelineStep) -> None: choice = Prompt.ask( "[bold]Choice[/]", - choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "a", "b", "c", "q"], + choices=["1", "2", "3", "4", "5", "6", "7", "8", "9", "a", "b", "c", "d", "q"], show_choices=False, ) @@ -959,16 +1023,18 @@ def run_step_prompt(step: PipelineStep) -> None: elif choice == "6": run_step_prompt(STEP_INDEX["augment_agera5"]) elif choice == "7": - run_step_prompt(STEP_INDEX["tessera"]) + run_step_prompt(STEP_INDEX["tessera_rename"]) elif choice == "8": - run_step_prompt(STEP_INDEX["merge_augmentations"]) + run_step_prompt(STEP_INDEX["tessera"]) elif choice == "9": - run_step_prompt(STEP_INDEX["check_augmentations"]) + run_step_prompt(STEP_INDEX["merge_augmentations"]) elif choice == "a": - run_step_prompt(STEP_INDEX["spatial_splits"]) + run_step_prompt(STEP_INDEX["check_augmentations"]) elif choice == "b": - run_step_prompt(STEP_INDEX["loco_splits"]) + run_step_prompt(STEP_INDEX["spatial_splits"]) elif choice == "c": + run_step_prompt(STEP_INDEX["loco_splits"]) + elif choice == "d": show_experiments() console.print() @@ -1072,10 +1138,18 @@ def _menu_ansi() -> str: + RESET, BOLD + FG_CYAN + " 5" + RESET + " Step 2a — Augment with NDVI", BOLD + FG_CYAN + " 6" + RESET + " Step 2b — Augment with AgERA5 climate", - BOLD + FG_CYAN + " 7" + RESET + " Step 2c — Download TESSERA embeddings", BOLD + FG_CYAN - + " 8" + + " 7" + + RESET + + " Step 2c-migrate — Rename TESSERA tiles " + + DIM + + "(one-time migration)" + + RESET, + BOLD + FG_CYAN + " 8" + RESET + " Step 2c — Download TESSERA embeddings", + BOLD + + FG_CYAN + + " 9" + RESET + " Step 2d — Merge augmented CSVs " + DIM @@ -1083,15 +1157,15 @@ def _menu_ansi() -> str: + RESET, BOLD + FG_CYAN - + " 9" + + " a" + RESET + " Step 2e — Compare all augmentations " + DIM + "(optional)" + RESET, - BOLD + FG_CYAN + " a" + RESET + " Step 3a — Generate spatial splits", - BOLD + FG_CYAN + " b" + RESET + " Step 3b — Generate LOCO splits", - BOLD + FG_CYAN + " c" + RESET + " Training experiments reference", + BOLD + FG_CYAN + " b" + RESET + " Step 3a — Generate spatial splits", + BOLD + FG_CYAN + " c" + RESET + " Step 3b — Generate LOCO splits", + BOLD + FG_CYAN + " d" + RESET + " Training experiments reference", BOLD + FG_CYAN + " q" + RESET + " Quit", ] return _box("Menu", lines, width=w, border=FG_CYAN) @@ -1278,7 +1352,9 @@ def _experiments_ansi() -> None: print() print(DIM + " Tip: override the split file on the command line, e.g.:" + RESET) - print(DIM + " python src/train.py experiment=yield_africa_fusion_spatial \\" + RESET) + print( + DIM + " python src/train.py experiment=yield_africa_tessera_fusion_spatial \\" + RESET + ) print(DIM + " data.saved_split_file_name=split_spatial_25km.pth" + RESET) @@ -1301,12 +1377,13 @@ def run_plain() -> None: "4": lambda: _run_step_ansi(STEP_INDEX["make_ready"]), "5": lambda: _run_step_ansi(STEP_INDEX["augment_ndvi"]), "6": lambda: _run_step_ansi(STEP_INDEX["augment_agera5"]), - "7": lambda: _run_step_ansi(STEP_INDEX["tessera"]), - "8": lambda: _run_step_ansi(STEP_INDEX["merge_augmentations"]), - "9": lambda: _run_step_ansi(STEP_INDEX["check_augmentations"]), - "a": lambda: _run_step_ansi(STEP_INDEX["spatial_splits"]), - "b": lambda: _run_step_ansi(STEP_INDEX["loco_splits"]), - "c": _experiments_ansi, + "7": lambda: _run_step_ansi(STEP_INDEX["tessera_rename"]), + "8": lambda: _run_step_ansi(STEP_INDEX["tessera"]), + "9": lambda: _run_step_ansi(STEP_INDEX["merge_augmentations"]), + "a": lambda: _run_step_ansi(STEP_INDEX["check_augmentations"]), + "b": lambda: _run_step_ansi(STEP_INDEX["spatial_splits"]), + "c": lambda: _run_step_ansi(STEP_INDEX["loco_splits"]), + "d": _experiments_ansi, } while True: diff --git a/src/data_preprocessing/yield_africa_tessera_preprocess.py b/src/data_preprocessing/yield_africa_tessera_preprocess.py index 3a1e770..6f0390d 100644 --- a/src/data_preprocessing/yield_africa_tessera_preprocess.py +++ b/src/data_preprocessing/yield_africa_tessera_preprocess.py @@ -3,10 +3,12 @@ Location: src/data_preprocessing/yield_africa_tessera_preprocess.py Tiles are saved as NumPy arrays to: - {data_dir}/yield_africa/eo/tessera/tessera_{name_loc}.npy + {data_dir}/yield_africa/eo/tessera/tessera_{name_loc}_{year}.npy -This matches the path built by BaseDataset.add_modality_paths_to_df() and -loaded by BaseDataset.setup_tessera() at training time. +The year suffix makes each tile's harvest year unambiguous and enables +dual-year fusion (year Y and year Y-1) without a separate directory. +YieldAfricaDataset overrides the default path built by +BaseDataset.add_modality_paths_to_df() to match this convention. Unlike tessera_from_df() (which takes a single fixed year), this script uses each record's own `year` column so that per-record inter-annual @@ -29,6 +31,10 @@ # Smaller tile size (faster, less context) python src/data_preprocessing/yield_africa_tessera_preprocess.py \\ --data_dir data/ --tile_size 5 + + # Also fetch year-1 tiles for dual-year tessera fusion + python src/data_preprocessing/yield_africa_tessera_preprocess.py \\ + --data_dir data/ --include-prev-year """ import argparse @@ -36,8 +42,8 @@ import multiprocessing import os import socket -import time import sys +import time from pathlib import Path # Ensure the project root is on sys.path when the script is run directly. @@ -45,7 +51,11 @@ import pandas as pd -from src.data_preprocessing.tessera_embeds import NoTileError, PartialTileError, get_tessera_embeds +from src.data_preprocessing.tessera_embeds import ( + NoTileError, + PartialTileError, + get_tessera_embeds, +) log = logging.getLogger(__name__) @@ -65,6 +75,7 @@ def _init_worker(cache_dir: str, use_local_registry: bool, registry_dir: str) -> None: """Pool initializer: runs once per worker process to set up GeoTessera.""" from geotessera import GeoTessera + global _process_gt socket.setdefaulttimeout(60) embeddings_dir = str(Path(cache_dir) / "raw") @@ -78,17 +89,97 @@ def _init_worker(cache_dir: str, use_local_registry: bool, registry_dir: str) -> def _worker_fetch(args: tuple) -> str: """Multiprocessing worker — reuses the per-process GeoTessera instance. - Must be a top-level function so it is picklable across processes. - Returning normally means success; raising means error (logged by caller). + Must be a top-level function so it is picklable across processes. Returning normally means + success; raising means error (logged by caller). """ lon, lat, name_loc, year, save_dir, tile_size = args get_tessera_embeds( - lon=lon, lat=lat, name_loc=name_loc, year=year, - save_dir=save_dir, tile_size=tile_size, tessera_con=_process_gt, + lon=lon, + lat=lat, + name_loc=name_loc, + year=year, + save_dir=save_dir, + tile_size=tile_size, + tessera_con=_process_gt, ) return name_loc +def _run_fetch_loop( + items: list[tuple[str, tuple]], + pool_initargs: tuple, + stuck_file: Path, + workers: int, +) -> None: + """Execute one fetch pass through the multiprocessing pool. + + :param items: list of ``(display_name, worker_args)`` pairs. ``display_name`` + is used in progress messages and appended to ``stuck_file`` on timeout. + ``worker_args`` is the tuple accepted by ``_worker_fetch``: + ``(lon, lat, name_loc_stem, year, save_dir_str, tile_size)``. + :param pool_initargs: forwarded verbatim to ``_init_worker`` on pool start. + :param stuck_file: path to the stuck log; timed-out display names are appended. + :param workers: number of pool worker processes. + :raises KeyboardInterrupt: re-raised after pool cleanup so the caller can exit. + """ + HEARTBEAT = 15 # seconds between "still fetching" log lines + TILE_TIMEOUT = 60 # seconds per record before the worker process is killed + + n = len(items) + pool = multiprocessing.Pool( + processes=workers, initializer=_init_worker, initargs=pool_initargs + ) + try: + for done, (display_name, args) in enumerate(items, 1): + result = pool.apply_async(_worker_fetch, (args,)) + start = time.monotonic() + timed_out = False + while True: + try: + result.get(timeout=HEARTBEAT) + break # completed successfully + except multiprocessing.TimeoutError: + elapsed = int(time.monotonic() - start) + if elapsed >= TILE_TIMEOUT: + timed_out = True + break + print(f" ... fetching {display_name} ({elapsed}s)") + except NoTileError: + print(f" Skipped {display_name}: no TESSERA data for this location/year") + break + except PartialTileError: + print( + f" Skipped {display_name}: tile too close to mosaic edge, not enough context" + ) + break + except Exception as exc: + print(f" ERROR fetching {display_name}: {exc}") + break + + if timed_out: + pool.terminate() + pool.join() + pool = multiprocessing.Pool( + processes=workers, initializer=_init_worker, initargs=pool_initargs + ) + with open(stuck_file, "a") as fh: + fh.write(display_name + "\n") + lon, lat, _, year = args[0], args[1], args[2], args[3] + print(f" Stuck: {display_name} lon={lon:.4f} lat={lat:.4f} year={year}") + + if done % 100 == 0 or done == n: + print(f" {done}/{n}") + + except KeyboardInterrupt: + print("\nInterrupted.") + pool.terminate() + pool.join() + raise + + pool.close() + pool.join() + + def fetch_tessera_tiles( data_dir: str, tile_size: int = DEFAULT_TILE_SIZE, @@ -97,6 +188,7 @@ def fetch_tessera_tiles( cache_dir: str | None = None, workers: int = 1, retry_stuck: bool = False, + include_prev_year: bool = False, ) -> None: """Fetch TESSERA tiles for every record in the yield_africa CSV. @@ -114,6 +206,10 @@ def fetch_tessera_tiles( own GeoTessera instance and can be killed on timeout. Default: 1. :param retry_stuck: if True, clear stuck.txt and retry previously-stuck records instead of skipping them. + :param include_prev_year: if True, after the main year-Y fetch, also fetch + year-Y−1 tiles needed for dual-year tessera fusion. Tiles that already + exist on disk are skipped. For locations without a CSV row at year-1, a + synthetic stem ``{name_loc}_prev_{year-1}`` is used and a warning is logged. """ dataset_dir = Path(data_dir) / DATASET_NAME csv_path = dataset_dir / MODEL_READY_CSV @@ -131,7 +227,9 @@ def fetch_tessera_tiles( embeddings_dir = str(Path(cache_dir) / "raw") + # Keep the full unfiltered df for the prev-year lat/lon index. df = pd.read_csv(csv_path) + df_full = df # Optional filters (consistent with YieldAfricaDataset filter params) if countries is not None: @@ -143,7 +241,9 @@ def fetch_tessera_tiles( n_total = len(df) n_existing = sum( - 1 for _, row in df.iterrows() if (save_dir / f"tessera_{row.name_loc}.npy").exists() + 1 + for _, row in df.iterrows() + if (save_dir / f"tessera_{row.name_loc}_{int(row.year)}.npy").exists() ) n_to_fetch = n_total - n_existing @@ -157,17 +257,6 @@ def fetch_tessera_tiles( _default_registry_dir = Path.home() / ".cache" / "geotessera" _use_local_registry = (_default_registry_dir / "registry.parquet").exists() - # Per-task timeout: when a worker process exceeds this, it is killed and - # the record added to stuck.txt. Multiprocessing (unlike threading) allows - # true process termination, so stuck downloads cannot block forward progress. - HEARTBEAT = 15 # seconds between "still fetching" log lines - TILE_TIMEOUT = 60 # seconds per record before the worker process is killed - # Note: stuck workers spin at 100% CPU and leak ~55 MB/s of rasterio - # MemoryFile objects inside GeoTessera's fetch loop. 60s caps peak memory - # at ~3 GB per stuck record and recovers 3x faster than the old 180s limit. - - # Records that caused a stall in a previous run are skipped unless - # --retry-stuck is passed, which clears stuck.txt before the run. stuck_file = save_dir / "stuck.txt" stuck_records: set[str] = set() if stuck_file.exists(): @@ -177,64 +266,107 @@ def fetch_tessera_tiles( else: stuck_records = set(stuck_file.read_text().splitlines()) if stuck_records: - print(f" Skipping {len(stuck_records)} previously-stuck record(s): {sorted(stuck_records)}") + print( + f" Skipping {len(stuck_records)} previously-stuck record(s): {sorted(stuck_records)}" + ) - rows = [row for _, row in df.iterrows() if row.name_loc not in stuck_records] - done = 0 + # Build main fetch items; display_name == row.name_loc for backward compat + # with existing stuck.txt entries (which store bare name_loc values). + main_items: list[tuple[str, tuple]] = [ + ( + row.name_loc, + ( + row.lon, + row.lat, + f"{row.name_loc}_{int(row.year)}", + int(row.year), + str(save_dir), + tile_size, + ), + ) + for _, row in df.iterrows() + if row.name_loc not in stuck_records + ] _pool_initargs = (cache_dir, _use_local_registry, str(_default_registry_dir)) - pool = multiprocessing.Pool(processes=workers, initializer=_init_worker, initargs=_pool_initargs) - try: - for row in rows: - args = (row.lon, row.lat, row.name_loc, int(row.year), str(save_dir), tile_size) - result = pool.apply_async(_worker_fetch, (args,)) - start = time.monotonic() - timed_out = False - while True: - try: - result.get(timeout=HEARTBEAT) - break # completed successfully - except multiprocessing.TimeoutError: - elapsed = int(time.monotonic() - start) - if elapsed >= TILE_TIMEOUT: - timed_out = True - break - print(f" ... fetching {row.name_loc} ({elapsed}s)") - except NoTileError: - print(f" Skipped {row.name_loc}: no TESSERA data for this location/year") - break - except PartialTileError: - print(f" Skipped {row.name_loc}: tile too close to mosaic edge, not enough context") - break - except Exception as exc: - print(f" ERROR fetching {row.name_loc}: {exc}") - break - if timed_out: - pool.terminate() - pool.join() - pool = multiprocessing.Pool(processes=workers, initializer=_init_worker, initargs=_pool_initargs) - with open(stuck_file, "a") as fh: - fh.write(row.name_loc + "\n") - print( - f" Stuck: {row.name_loc} " - f"lon={row.lon:.4f} lat={row.lat:.4f} year={int(row.year)}" - ) + try: + _run_fetch_loop(main_items, _pool_initargs, stuck_file, workers) + except KeyboardInterrupt: + return - done += 1 - if done % 100 == 0 or done == len(rows): - print(f" {done}/{n_total}") + print(f"Done. Tiles saved to: {save_dir}") - except KeyboardInterrupt: - print("\nInterrupted.") - pool.terminate() - pool.join() + if not include_prev_year: return - pool.close() - pool.join() + # ------------------------------------------------------------------ # + # Prev-year fetch: tiles for year Y-1 needed for dual-year fusion. # + # ------------------------------------------------------------------ # + print("\nBuilding prev-year tile list...") - print(f"Done. Tiles saved to: {save_dir}") + # Refresh stuck records — main loop may have appended new entries. + stuck_records_prev: set[str] = set() + if stuck_file.exists(): + stuck_records_prev = set(stuck_file.read_text().splitlines()) + + # Build (lat, lon, year) → name_loc index from the full unfiltered CSV so + # that rows excluded by a countries/years filter can still supply their + # tile paths as prev-year sources. + loc_year_index: dict[tuple[float, float, int], str] = { + (round(float(r.lat), 6), round(float(r.lon), 6), int(r.year)): r.name_loc + for _, r in df_full.iterrows() + } + + prev_items: list[tuple[str, tuple]] = [] + n_prev_already_present = 0 + n_synthetic = 0 + + for _, row in df.iterrows(): + year_prev = int(row.year) - 1 + key = (round(float(row.lat), 6), round(float(row.lon), 6), year_prev) + + prev_name_loc = loc_year_index.get(key) + if prev_name_loc is not None: + # A CSV row exists for this location at year-1; reuse its name_loc. + stem = f"{prev_name_loc}_{year_prev}" + else: + # No CSV row at year-1 (first-year location or gap year). + # Use a synthetic stem so the tile can be audited. + stem = f"{row.name_loc}_prev_{year_prev}" + n_synthetic += 1 + + tile_path = save_dir / f"tessera_{stem}.npy" + if tile_path.exists(): + n_prev_already_present += 1 + continue + + if stem in stuck_records_prev: + continue + + args = (float(row.lon), float(row.lat), stem, year_prev, str(save_dir), tile_size) + prev_items.append((stem, args)) + + if n_synthetic: + log.warning( + "%d prev-year record(s) have no matching CSV row at year-1 " + "(synthetic stems used; audit tiles named *_prev_* in %s).", + n_synthetic, + save_dir, + ) + + print( + f"Prev-year tiles: {n_prev_already_present} already present, " + f"{len(prev_items)} to fetch, {n_synthetic} with synthetic stems." + ) + + if prev_items: + try: + _run_fetch_loop(prev_items, _pool_initargs, stuck_file, workers) + except KeyboardInterrupt: + return + + print("Prev-year fetch done.") def main() -> None: @@ -299,12 +431,22 @@ def main() -> None: default=False, help="Clear stuck.txt and retry previously-stuck records instead of skipping them.", ) + parser.add_argument( + "--include-prev-year", + action="store_true", + default=False, + help=( + "After the main year-Y fetch, also fetch year-Y-1 tiles for dual-year " + "tessera fusion. Tiles already on disk are skipped. " + "Use with configs/data/yield_africa_tessera_dual.yaml." + ), + ) args = parser.parse_args() print( f"Fetching TESSERA tiles data_dir={args.data_dir} " f"tile_size={args.tile_size} countries={args.countries or 'all'} " - f"years={args.years or 'all'}" + f"years={args.years or 'all'} include_prev_year={args.include_prev_year}" ) fetch_tessera_tiles( data_dir=args.data_dir, @@ -314,6 +456,7 @@ def main() -> None: cache_dir=args.cache_dir, workers=args.workers, retry_stuck=args.retry_stuck, + include_prev_year=args.include_prev_year, ) diff --git a/src/data_preprocessing/yield_africa_tessera_rename.py b/src/data_preprocessing/yield_africa_tessera_rename.py new file mode 100644 index 0000000..c6bb4ee --- /dev/null +++ b/src/data_preprocessing/yield_africa_tessera_rename.py @@ -0,0 +1,115 @@ +"""Rename existing tessera tiles to include a year suffix. + +Existing naming: tessera_{name_loc}.npy +New naming: tessera_{name_loc}_{year}.npy + +Run this script once before deploying Tasks 2-7 of the dual-year tessera plan. +Tiles already carrying a four-digit year suffix are skipped automatically, so +the script is safe to re-run. + +Usage +----- + # Preview what would be renamed (default — no files touched) + python src/data_preprocessing/yield_africa_tessera_rename.py \\ + --data_dir /path/to/external/data + + # Actually rename + python src/data_preprocessing/yield_africa_tessera_rename.py \\ + --data_dir /path/to/external/data --no-dry-run +""" + +import argparse +import logging +import re +import sys +from pathlib import Path + +sys.path.insert(0, str(Path(__file__).resolve().parents[2])) + +import pandas as pd + +log = logging.getLogger(__name__) + +DATASET_NAME = "yield_africa" +MODEL_READY_CSV = f"model_ready_{DATASET_NAME}.csv" + +# Matches a four-digit year at the end of the stem, e.g. _2020 in tessera_RWA_24803_2020 +_YEAR_SUFFIX_RE = re.compile(r"_\d{4}$") + + +def rename_tessera_tiles(data_dir: str, dry_run: bool = True) -> None: + dataset_dir = Path(data_dir) / DATASET_NAME + csv_path = dataset_dir / MODEL_READY_CSV + tessera_dir = dataset_dir / "eo" / "tessera" + + if not csv_path.exists(): + raise FileNotFoundError(f"Model-ready CSV not found: {csv_path}") + if not tessera_dir.exists(): + raise FileNotFoundError(f"Tessera directory not found: {tessera_dir}") + + df = pd.read_csv(csv_path) + name_loc_to_year: dict[str, int] = dict(zip(df["name_loc"], df["year"].astype(int))) + + n_renamed = 0 + n_already_done = 0 + n_unmatched = 0 + + for src in sorted(tessera_dir.glob("tessera_*.npy")): + name_loc = src.stem.removeprefix("tessera_") + + if _YEAR_SUFFIX_RE.search(name_loc): + n_already_done += 1 + continue + + year = name_loc_to_year.get(name_loc) + if year is None: + log.warning("No year found in CSV for %s — leaving untouched", src.name) + n_unmatched += 1 + continue + + dst = src.with_name(f"tessera_{name_loc}_{year}.npy") + print(f" {'[dry-run] ' if dry_run else ''}rename {src.name} -> {dst.name}") + if not dry_run: + src.rename(dst) + n_renamed += 1 + + print( + f"\nSummary: {n_renamed} to rename, " + f"{n_already_done} already up-to-date, " + f"{n_unmatched} unmatched." + ) + if dry_run and n_renamed > 0: + print("Re-run with --no-dry-run to apply changes.") + + +def main() -> None: + logging.basicConfig(level=logging.WARNING, format="%(levelname)s: %(message)s") + + parser = argparse.ArgumentParser( + description=( + "Rename tessera_.npy tiles to tessera__.npy. " + "Dry-run is the default — pass --no-dry-run to apply changes." + ) + ) + parser.add_argument( + "--data_dir", + type=str, + required=True, + help="Root data directory (same as paths.data_dir in configs).", + ) + parser.add_argument( + "--dry-run", + action=argparse.BooleanOptionalAction, + default=True, + help="Preview renames without touching files (default: True). Pass --no-dry-run to apply.", + ) + args = parser.parse_args() + + if args.dry_run: + print("DRY RUN — no files will be renamed. Pass --no-dry-run to apply.\n") + + rename_tessera_tiles(data_dir=args.data_dir, dry_run=args.dry_run) + + +if __name__ == "__main__": + main() diff --git a/src/models/base_model.py b/src/models/base_model.py index d0f92e8..c327e44 100644 --- a/src/models/base_model.py +++ b/src/models/base_model.py @@ -249,6 +249,14 @@ def on_save_checkpoint(self, checkpoint): if any(k.startswith(part) for part in self.trainable_modules) } + # Also save all non-None buffers (normalisation stats such as target_mean, + # target_std, feat_mean, feat_std are not trainable parameters so they + # never match the trainable_modules filter above, but they must survive + # checkpointing so resumed runs and standalone inference stay correct). + for name, buf in self.named_buffers(): + if buf is not None: + checkpoint["state_dict"][name] = buf + # Update model configurations checkpoint["hyper_parameters"].update( { diff --git a/src/models/components/geo_encoders/average_encoder.py b/src/models/components/geo_encoders/average_encoder.py index 7e31b44..23a2363 100644 --- a/src/models/components/geo_encoders/average_encoder.py +++ b/src/models/components/geo_encoders/average_encoder.py @@ -17,7 +17,7 @@ def __init__( """ super().__init__() - self.dict_n_bands_default = {"s2": 4, "aef": 64, "tessera": 128} + self.dict_n_bands_default = {"s2": 4, "aef": 64, "tessera": 128, "tessera_prev": 128} self.allowed_geo_data_names: list[str] = list(self.dict_n_bands_default.keys()) assert ( geo_data_name in self.allowed_geo_data_names @@ -36,7 +36,11 @@ def _setup(self) -> List[str]: def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: """Data forward pass through the encoder.""" tile = batch.get("eo", {}).get(self.geo_data_name) - feats = self.geo_encoder(tile.mean(dim=(-2, -1))) + # nanmean so that masked/nodata pixels (NaN from rasterio merge fill) are + # excluded from the spatial average rather than poisoning the whole channel. + # nan_to_num handles the edge case where an entire channel is masked (all + # pixels NaN), which nanmean would otherwise return as NaN. + feats = self.geo_encoder(torch.nan_to_num(tile.nanmean(dim=(-2, -1)), nan=0.0)) if self.extra_projector: feats = self.extra_projector(feats) return feats diff --git a/src/models/components/geo_encoders/tabular_encoder.py b/src/models/components/geo_encoders/tabular_encoder.py index e0572eb..70cb988 100644 --- a/src/models/components/geo_encoders/tabular_encoder.py +++ b/src/models/components/geo_encoders/tabular_encoder.py @@ -79,6 +79,9 @@ def forward(self, batch: Dict[str, torch.Tensor]) -> torch.Tensor: if self.feat_mean is not None: tab_data = (tab_data - self.feat_mean) / self.feat_std + # Replace NaN (e.g. missing feat_ndvi_month_* values) with 0, which is + # the training mean in z-score space — equivalent to mean imputation. + tab_data = torch.nan_to_num(tab_data, nan=0.0) feats = self.geo_encoder(tab_data) diff --git a/tests/test_dynamic_gate_fusion_encoder.py b/tests/test_dynamic_gate_fusion_encoder.py index 6dff786..29da7db 100644 --- a/tests/test_dynamic_gate_fusion_encoder.py +++ b/tests/test_dynamic_gate_fusion_encoder.py @@ -195,6 +195,67 @@ def test_with_tabular_encoder(): assert out.shape == (4, 32) +def test_three_branch_dual_tessera(): + """Three-branch EncoderWrapper matching the dual-tessera model config. + + Branches: tessera (year Y), tessera_prev (year Y-1), tabular — all projected to dim=256. + Verifies: output shape [batch, 256], gate weights sum to 1.0 per sample, all weights > 0. + """ + from src.models.components.geo_encoders.mlp_projector import MLPProjector + + tabular_dim = 20 + stub_dim = 128 # TESSERA embedding width + + wrapper = EncoderWrapper( + encoder_branches=[ + { + "encoder": _StubEncoder(stub_dim, key="tessera"), + "projector": MLPProjector(nn_layers=1, output_dim=256), + }, + { + "encoder": _StubEncoder(stub_dim, key="tessera_prev"), + "projector": MLPProjector(nn_layers=1, output_dim=256), + }, + {"encoder": TabularEncoder(output_dim=256, dropout_prob=0.2)}, + ], + fusion_strategy="dynamic_gate", + ) + wrapper.set_tabular_input_dim(tabular_dim) + wrapper.setup() + + batch_size = 4 + batch = { + "eo": { + "tessera": torch.randn(batch_size, stub_dim), + "tessera_prev": torch.randn(batch_size, stub_dim), + "tabular": torch.randn(batch_size, tabular_dim), + } + } + + out = wrapper(batch) + assert out.shape == (batch_size, 256) + + # Replicate the gating arithmetic to inspect branch weights. + with torch.no_grad(): + branch_feats = [] + for i, branch in enumerate(wrapper.encoder_branches): + feats = branch["encoder"](batch) + if "projector" in branch: + feats = branch["projector"](feats) + feats = wrapper.branch_norms[i](feats) + branch_feats.append(feats) + + stacked = torch.stack(branch_feats, dim=1) + gate_input = stacked.flatten(start_dim=1) + weights = torch.softmax(wrapper.dynamic_gate_mlp(gate_input), dim=1) + + # Weights must sum to 1.0 per sample across all 3 branches. + assert weights.shape == (batch_size, 3) + assert weights.sum(dim=1).allclose(torch.ones(batch_size), atol=1e-6) + # Softmax is strictly positive — all branches have non-zero weight. + assert (weights > 0).all() + + def test_existing_gated_unaffected(): """Ensure static gated strategy still works after the dynamic_gate additions.""" wrapper = EncoderWrapper( diff --git a/tests/test_yield_africa_dual_tessera_dataset.py b/tests/test_yield_africa_dual_tessera_dataset.py new file mode 100644 index 0000000..3bc9d6a --- /dev/null +++ b/tests/test_yield_africa_dual_tessera_dataset.py @@ -0,0 +1,255 @@ +"""Tests for dual-year TESSERA loading in YieldAfricaDataset. + +All data (CSV + .npy tiles) is created synthetically via tmp_path — no +GeoTessera service access required. Run with: pytest --use-mock + +Mock layout +----------- +Three records are created: + + LOC_A (lat=1.0, lon=1.0, year=2019) — serves as the prev-year source for LOC_B + LOC_B (lat=1.0, lon=1.0, year=2020) — same physical location; prev-year = LOC_A + LOC_C (lat=2.0, lon=2.0, year=2020) — unique location; no prev-year tile on disk + +Tiles written to {mock_dir}/eo/tessera/: + tessera_LOC_A_2019.npy shape [TILE_SIZE, TILE_SIZE, 128] + tessera_LOC_B_2020.npy + tessera_LOC_C_2020.npy + (no tile for lat=2.0, lon=2.0, year=2019) +""" + +import os +from pathlib import Path + +import numpy as np +import pandas as pd +import pytest +import torch + +from src.data.yield_africa_dataset import YieldAfricaDataset + +# --------------------------------------------------------------------------- +# Constants +# --------------------------------------------------------------------------- + +TILE_SIZE = 3 +TILE_CHANNELS = 128 # fixed by the geotessera model + +DUAL_MODALITIES = { + "tessera": {"size": TILE_SIZE, "format": "npy"}, + "tessera_prev": {"size": TILE_SIZE, "format": "npy"}, +} +SINGLE_MODALITIES = { + "tessera": {"size": TILE_SIZE, "format": "npy"}, +} + +# Minimal feature columns (must start with feat_) +_FEAT_COLS = { + "feat_map": [820, 750, 700], + "feat_mat": [22.1, 21.5, 21.0], +} + +# --------------------------------------------------------------------------- +# Helpers +# --------------------------------------------------------------------------- + + +def _make_tile(rng: np.random.Generator | None = None) -> np.ndarray: + """Return a synthetic tessera tile with shape [H, W, C] = [TILE_SIZE, TILE_SIZE, 128].""" + gen = rng or np.random.default_rng(0) + return gen.random((TILE_SIZE, TILE_SIZE, TILE_CHANNELS)).astype(np.float32) + + +def _write_mock_dir(tmp_path: Path) -> tuple[str, Path]: + """Create the mock CSV and tessera tiles; return (data_dir, tessera_dir).""" + mock_dir = tmp_path / "mock" + tessera_dir = mock_dir / "eo" / "tessera" + tessera_dir.mkdir(parents=True, exist_ok=True) + + df = pd.DataFrame( + { + "name_loc": ["LOC_A", "LOC_B", "LOC_C"], + "lat": [1.0, 1.0, 2.0], + "lon": [1.0, 1.0, 2.0], + "year": [2019, 2020, 2020], + "country": ["ETH", "ETH", "ETH"], + "target_yld_ton_ha": [2.0, 2.5, 1.8], + **_FEAT_COLS, + } + ) + df.to_csv(mock_dir / "model_ready_mock.csv", index=False) + + rng = np.random.default_rng(42) + np.save(tessera_dir / "tessera_LOC_A_2019.npy", _make_tile(rng)) + np.save(tessera_dir / "tessera_LOC_B_2020.npy", _make_tile(rng)) + np.save(tessera_dir / "tessera_LOC_C_2020.npy", _make_tile(rng)) + # No tile for (lat=2.0, lon=2.0, year=2019) — LOC_C has no prev-year tile. + + return str(tmp_path), tessera_dir + + +def _make_dataset( + data_dir: str, + tmp_path: Path, + modalities: dict, + **kwargs, +) -> YieldAfricaDataset: + return YieldAfricaDataset( + data_dir=data_dir, + cache_dir=str(tmp_path / "cache"), + modalities=modalities, + use_target_data=True, + use_aux_data="none", + seed=42, + mock=True, + use_features=True, + **kwargs, + ) + + +# --------------------------------------------------------------------------- +# Fixtures +# --------------------------------------------------------------------------- + + +@pytest.fixture +def yield_africa_dual_csv(request, tmp_path) -> tuple[str, Path]: + """Mock CSV + tessera tiles; returns (data_dir, tessera_dir).""" + use_mock = request.config.getoption("--use-mock") + if not use_mock: + assert False, "Real data not available in test environment." + return _write_mock_dir(tmp_path) + + +# --------------------------------------------------------------------------- +# Tests +# --------------------------------------------------------------------------- + + +def test_tessera_filename_includes_year(yield_africa_dual_csv, tmp_path): + """tessera_path in records must end with _{year}.npy after path rewriting.""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset(data_dir, tmp_path, SINGLE_MODALITIES) + for rec in ds.records: + assert rec["tessera_path"].endswith(".npy") + # Filename stem is tessera_{name_loc}_{year} + fname = os.path.basename(rec["tessera_path"]) + assert fname.startswith("tessera_") + # Year appears as a four-digit suffix immediately before .npy + stem = fname.removesuffix(".npy") + year_part = stem.split("_")[-1] + assert year_part.isdigit() and len(year_part) == 4, f"Expected year suffix in '{fname}'" + + +def test_single_year_unaffected(yield_africa_dual_csv, tmp_path): + """Dataset with only `tessera` modality loads without tessera_prev_path in records.""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset(data_dir, tmp_path, SINGLE_MODALITIES) + + assert len(ds) == 3 # all records have year-Y tiles on disk + for rec in ds.records: + assert "tessera_path" in rec + assert "tessera_prev_path" not in rec + + sample = ds[0] + assert "tessera" in sample["eo"] + assert "tessera_prev" not in sample["eo"] + + +def test_prev_path_points_to_existing_tile(yield_africa_dual_csv, tmp_path): + """tessera_prev_path for LOC_B must point to LOC_A's tessera file — no copy.""" + data_dir, tessera_dir = yield_africa_dual_csv + # require=True → only LOC_B survives (A has no 2018 tile; C has no 2019 tile) + ds = _make_dataset(data_dir, tmp_path, DUAL_MODALITIES, require_prev_year_tessera=True) + + assert len(ds) == 1 + rec = ds.records[0] + assert rec["tessera_path"].endswith("tessera_LOC_B_2020.npy") + + prev_path = rec["tessera_prev_path"] + assert prev_path.endswith("tessera_LOC_A_2019.npy") + # Must be the original file — no copy written + assert os.path.exists(prev_path) + assert os.path.dirname(prev_path) == str(tessera_dir) + + +def test_dual_year_both_keys_in_sample(yield_africa_dual_csv, tmp_path): + """Sample['eo'] must contain both 'tessera' and 'tessera_prev' tensors.""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset(data_dir, tmp_path, DUAL_MODALITIES, require_prev_year_tessera=True) + + assert len(ds) == 1 + sample = ds[0] + assert "tessera" in sample["eo"] + assert "tessera_prev" in sample["eo"] + assert isinstance(sample["eo"]["tessera"], torch.Tensor) + assert isinstance(sample["eo"]["tessera_prev"], torch.Tensor) + + +def test_tensor_shape_matches_size_config(yield_africa_dual_csv, tmp_path): + """Both tessera tensors must have shape [channels, size, size].""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset(data_dir, tmp_path, DUAL_MODALITIES, require_prev_year_tessera=True) + + sample = ds[0] + expected = (TILE_CHANNELS, TILE_SIZE, TILE_SIZE) + assert sample["eo"]["tessera"].shape == expected + assert sample["eo"]["tessera_prev"].shape == expected + + +def test_missing_prev_tile_dropped_by_default(yield_africa_dual_csv, tmp_path): + """Records with no resolvable year-1 tile are dropped when require_prev_year_tessera=True.""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset(data_dir, tmp_path, DUAL_MODALITIES, require_prev_year_tessera=True) + + # LOC_A has no 2018 tile; LOC_C has no 2019 tile at its location. + # Only LOC_B (whose prev-year tile is LOC_A's 2019 tile) survives. + assert len(ds) == 1 + assert ds.records[0]["tessera_path"].endswith("tessera_LOC_B_2020.npy") + + +def test_missing_prev_tile_kept_when_not_required(yield_africa_dual_csv, tmp_path): + """Records with no prev tile are retained with tessera_prev_path=None when require=False.""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset(data_dir, tmp_path, DUAL_MODALITIES, require_prev_year_tessera=False) + + # All 3 records are retained (prev tile missing → None, not dropped). + assert len(ds) == 3 + + prev_paths = [rec["tessera_prev_path"] for rec in ds.records] + none_count = sum(1 for p in prev_paths if p is None) + resolved_count = sum(1 for p in prev_paths if p is not None) + assert none_count == 2 # LOC_A (no 2018 tile) + LOC_C (no 2019 tile at its location) + assert resolved_count == 1 # LOC_B resolves to LOC_A's 2019 tile + + +def test_year1_row_excluded_by_filter_still_resolves(yield_africa_dual_csv, tmp_path): + """LOC_B resolves its year-1 path even when years=[2020] excludes LOC_A from training.""" + data_dir, _ = yield_africa_dual_csv + ds = _make_dataset( + data_dir, + tmp_path, + DUAL_MODALITIES, + require_prev_year_tessera=True, + years=[2020], + ) + + # years=[2020] removes LOC_A (2019) from self.df after filtering. + # But the cross-year index was built before the filter → LOC_B still resolves. + # LOC_C (2020, unique location) still has no prev tile → dropped. + assert len(ds) == 1 + rec = ds.records[0] + assert rec["tessera_path"].endswith("tessera_LOC_B_2020.npy") + assert rec["tessera_prev_path"].endswith("tessera_LOC_A_2019.npy") + assert os.path.exists(rec["tessera_prev_path"]) + + +def test_no_duplicate_file_written(yield_africa_dual_csv, tmp_path): + """Dataset initialisation must not write any new files to the tessera directory.""" + data_dir, tessera_dir = yield_africa_dual_csv + files_before = set(tessera_dir.iterdir()) + + _make_dataset(data_dir, tmp_path, DUAL_MODALITIES, require_prev_year_tessera=False) + + files_after = set(tessera_dir.iterdir()) + assert files_after == files_before, f"Unexpected new files: {files_after - files_before}" From 28469edfffc032457338272230ba0dbac30dbbc8 Mon Sep 17 00:00:00 2001 From: Rob Knapen Date: Tue, 7 Jul 2026 14:20:51 +0200 Subject: [PATCH 2/2] Restore DBSCAN spatial split as default alongside grid-based spatial_grid mode --- src/data/base_datamodule.py | 119 ++++++++++++++++++++++++++++++++++-- 1 file changed, 115 insertions(+), 4 deletions(-) diff --git a/src/data/base_datamodule.py b/src/data/base_datamodule.py index 9b4ef50..9b8d45a 100644 --- a/src/data/base_datamodule.py +++ b/src/data/base_datamodule.py @@ -1,4 +1,6 @@ +import copy import os +import time from functools import partial from typing import Any, Dict, List, Tuple @@ -6,6 +8,7 @@ import pandas as pd import torch from lightning import LightningDataModule +from sklearn.cluster import DBSCAN from sklearn.model_selection import GroupShuffleSplit from torch.utils.data import DataLoader, random_split @@ -41,13 +44,15 @@ def __init__( :param num_workers: number of workers for dataloader :param pin_memory: pin memory for dataloader :param dataset_name: dataset name - :param split_mode: data split mode: random/from_file + :param split_mode: data split mode: random/spatial_clusters/spatial_grid/from_file :param save_split: if to save split file :param saved_split_file_name: file name to save split file :param caption_builder: instance of BaseCaptionBuilder for generating textual captions - :param spatial_split_distance_m: grid cell size in metres when split_mode is - 'spatial_clusters'. Samples within the same cell are kept together and assigned to the - same split. Default 1000 m. + :param spatial_split_distance_m: distance in metres used for spatial splitting. When + split_mode is 'spatial_clusters', this is the DBSCAN eps (max distance between samples + for them to be considered part of the same cluster). When split_mode is 'spatial_grid', + this is the grid cell size; samples within the same cell are kept together and assigned + to the same split. Default 1000 m. """ super().__init__() self.save_hyperparameters(logger=False) @@ -119,6 +124,112 @@ def split_data(self) -> None: } elif self.hparams.split_mode == "spatial_clusters": + min_dist = self.hparams.spatial_split_distance_m + # Use records (not df): records is already filtered (e.g. missing tiles + # dropped), so len(records) <= len(df). Indices must be into records + # because __getitem__ and __len__ both operate on self.records. + # lat/lon come from df (always present) keyed by name_loc so the + # coordinate array stays aligned with records regardless of modalities. + records = self.dataset.records + _nl_to_coords = dict( + zip( + self.dataset.df["name_loc"], + zip(self.dataset.df["lat"], self.dataset.df["lon"]), + ) + ) + coords = np.array( + [ + [_nl_to_coords[r["name_loc"]][0] for r in records], + [_nl_to_coords[r["name_loc"]][1] for r in records], + ] + ).T + n = len(coords) + print( + f"Splitting {n} samples into spatial clusters " + f"(eps={min_dist / 1000:.1f} km, haversine, n_jobs=-1)..." + ) + # Convert (lat, lon) degrees to radians for sklearn's haversine metric. + # haversine returns arc length on the unit sphere, so eps must be in radians. + _EARTH_RADIUS_M = 6_371_000 + coords_rad = np.radians(coords) + eps_rad = min_dist / _EARTH_RADIUS_M + t0 = time.time() + clustering = DBSCAN( + eps=eps_rad, + metric="haversine", + algorithm="ball_tree", + min_samples=2, + n_jobs=-1, + ).fit(coords_rad) + print(f"DBSCAN done in {time.time() - t0:.1f}s. Creating splits...") + # Non-clustered points are labeled -1. Change to new cluster label. + clusters = copy.deepcopy(clustering.labels_) + new_cl = np.max(clusters) + 1 + for i, cl in enumerate(clusters): + if cl == -1: + clusters[i] = new_cl + new_cl += 1 + + gss = GroupShuffleSplit( + n_splits=1, + test_size=self.hparams.train_val_test_split[2], + random_state=self.hparams.seed, + ) + train_val_indices, test_indices = next( + gss.split(np.arange(len(coords)), groups=clusters) + ) + gss_2 = GroupShuffleSplit( + n_splits=1, + test_size=( + self.hparams.train_val_test_split[1] + / (self.hparams.train_val_test_split[0] + self.hparams.train_val_test_split[1]) + ), + random_state=self.hparams.seed, + ) + tmp_train_indices, tmp_val_indices = next( + gss_2.split(train_val_indices, groups=clusters[train_val_indices]) + ) + train_indices = train_val_indices[tmp_train_indices] + val_indices = train_val_indices[tmp_val_indices] + clusters_train = clusters[train_indices] + clusters_val = clusters[val_indices] + clusters_test = clusters[test_indices] + # assert no overlap in indices: + assert len(np.intersect1d(train_indices, val_indices)) == 0, np.intersect1d( + train_indices, val_indices + ) + assert len(np.intersect1d(train_indices, test_indices)) == 0, np.intersect1d( + train_indices, test_indices + ) + assert len(np.intersect1d(val_indices, test_indices)) == 0, np.intersect1d( + val_indices, test_indices + ) + + # assert no overlap in clusters: + assert len(np.intersect1d(clusters_train, clusters_val)) == 0, np.intersect1d( + clusters_train, clusters_val + ) + assert len(np.intersect1d(clusters_train, clusters_test)) == 0, np.intersect1d( + clusters_train, clusters_test + ) + assert len(np.intersect1d(clusters_val, clusters_test)) == 0, np.intersect1d( + clusters_val, clusters_test + ) + + print( + f"Created {len(train_indices)} train, {len(val_indices)} val, {len(test_indices)} " + f"test indices using DBSCAN spatial clustering with {min_dist} m minimum " + f"distance between clusters." + ) + if self.hparams.save_split: + split_indices = { + "train_indices": pd.Series([records[i]["name_loc"] for i in train_indices]), + "val_indices": pd.Series([records[i]["name_loc"] for i in val_indices]), + "test_indices": pd.Series([records[i]["name_loc"] for i in test_indices]), + "clusters": clusters, + } + + elif self.hparams.split_mode == "spatial_grid": min_dist = self.hparams.spatial_split_distance_m # Use records (not df): records is already filtered (e.g. missing tiles # dropped), so len(records) <= len(df). Indices must be into records