Metrics configurable through yaml

README.md

@@ -196,16 +196,13 @@ We follow the directory structure from the [Hydra-Lightning template](https://gi
 └── README.md
 ```
 
-
 ## Architecture Overview
 
 The diagram below shows how configuration files, datasets, and model components
 relate to each other in the AETHER framework.
 
 ![Framework Architecture](docs/figures/diagram.png)
 
-
-
 ## Attribution
 
 This repo is based on the [Hydra-Lightning template](https://github.com/ashleve/lightning-hydra-template).

configs/data/heat_guatemala.yaml

@@ -19,6 +19,5 @@ split_mode: "from_file"
 save_split: false
 saved_split_file_name: "split_indices_heat_guatemala_2026-02-20-1148.pth"
 
-
 train_val_test_split: [0.7, 0.15, 0.15]
-seed: ${seed}
+seed: ${seed}

configs/experiment/alignment.yaml

@@ -6,6 +6,7 @@
 defaults:
   - override /model: geoclip_alignment
   - override /data: butterfly_coords_text
+  - override /metrics: contrastive_similarities
 
 # all parameters below will be merged with parameters from default configurations set above
 # this allows you to overwrite only specified parameters

configs/experiment/alignment_llm2clip.yaml

@@ -6,6 +6,7 @@
 defaults:
   - override /model: geoclip_llm2clip_alignment
   - override /data: butterfly_coords_text
+  - override /metrics: contrastive_similarities
 
 # all parameters below will be merged with parameters from default configurations set above
 # this allows you to overwrite only specified parameters

configs/experiment/heat_guatemala_coords_reg.yaml

@@ -2,10 +2,10 @@
 # configs/experiment/heat_guatemala_coords_reg.yaml
 # Variant A: GeoClip coordinates only
 
-
 defaults:
   - override /model: heat_geoclip_reg
   - override /data: heat_guatemala
+  - override /metrics: guatemala_regression
 
 tags: ["heat_island", "guatemala", "coords_only", "regression"]
 seed: 12345

configs/experiment/heat_guatemala_fusion_reg.yaml

@@ -2,10 +2,10 @@
 # configs/experiment/heat_guatemala_fusion_reg.yaml
 # Variant C: GeoClip + tabular fusion
 
-
 defaults:
   - override /model: heat_fusion_reg
   - override /data: heat_guatemala
+  - override /metrics: guatemala_regression
 
 tags: ["heat_island", "guatemala", "fusion", "regression"]
 seed: 12345

configs/experiment/heat_guatemala_tabular_reg.yaml

@@ -2,10 +2,10 @@
 # configs/experiment/heat_guatemala_tabular_reg.yaml
 # Variant B: tabular features only
 
-
 defaults:
   - override /model: heat_tabular_reg
   - override /data: heat_guatemala
+  - override /metrics: guatemala_regression
 
 tags: ["heat_island", "guatemala", "tabular_only", "regression"]
 seed: 12345

configs/experiment/prediction.yaml

@@ -6,6 +6,7 @@
 defaults:
   - override /model: predictive_geoclip
   - override /data: butterfly_coords
+  - override /metrics: butterfly_predictive
 
 # all parameters below will be merged with parameters from default configurations set above
 # this allows you to overwrite only specified parameters

configs/metrics/butterfly_predictive.yaml

@@ -0,0 +1,6 @@
+_target_: src.models.components.metrics.metrics_wrapper.MetricsWrapper
+
+metrics:
+  - _target_: src.models.components.loss_fns.mse_loss.MSELoss
+  - _target_: src.models.components.metrics.top_k_accuracy.TopKAccuracy
+    k_list: [1, 5, 10]

configs/metrics/contrastive_similarities.yaml

@@ -0,0 +1,4 @@
+_target_: src.models.components.metrics.metrics_wrapper.MetricsWrapper
+
+metrics:
+  - _target_: src.models.components.metrics.contrastive_similarities.CosineSimilarities

configs/metrics/guatemala_regression.yaml

@@ -0,0 +1,7 @@
+_target_: src.models.components.metrics.metrics_wrapper.MetricsWrapper
+
+metrics:
+  - _target_: src.models.components.loss_fns.mse_loss.MSELoss
+  - _target_: src.models.components.loss_fns.rmse_loss.RMSELoss
+  - _target_: src.models.components.loss_fns.mae_loss.MAELoss
+  - _target_: src.models.components.metrics.r2.RSquared

configs/model/geoclip_alignment.yaml

@@ -9,6 +9,8 @@ text_encoder:
 
 trainable_modules: [text_encoder.projector, loss_fn.log_temp]
 
+metrics: ${metrics}
+
 optimizer:
   _target_: torch.optim.Adam
   _partial_: true

configs/model/geoclip_llm2clip_alignment.yaml

@@ -9,6 +9,8 @@ text_encoder:
 
 trainable_modules: [text_encoder.projector, loss_fn.log_temp]
 
+metrics: ${metrics}
+
 optimizer:
   _target_: torch.optim.Adam
   _partial_: true

configs/model/heat_fusion_reg.yaml

@@ -6,13 +6,13 @@
 # Variant: GeoClip (coords) + tabular fusion
 #   Encoder output = GeoClip (512) + tabular projection (64) = 576-dim
 
-_target_: src.models.predictive_model_regression.PredictiveRegressionModel
+_target_: src.models.predictive_model.PredictiveModel
 
 eo_encoder:
   _target_: src.models.components.eo_encoders.multimodal_encoder.MultiModalEncoder
   use_coords: true
   use_tabular: true
-  tab_embed_dim: 64
+#  tab_embed_dim: 64
 
 prediction_head:
   _target_: src.models.components.pred_heads.mlp_regression_head.MLPRegressionPredictionHead
@@ -22,6 +22,8 @@ prediction_head:
 # GeoClip frozen; tabular projection + head are trained.
 trainable_modules: [eo_encoder, prediction_head]
 
+metrics: ${metrics}
+
 optimizer:
   _target_: torch.optim.Adam
   _partial_: true

configs/model/heat_geoclip_reg.yaml

@@ -6,7 +6,7 @@
 #
 # Variant: coords only (GeoClip encodes lat/lon → 512-dim embedding)
 
-_target_: src.models.predictive_model_regression.PredictiveRegressionModel
+_target_: src.models.predictive_model.PredictiveModel
 
 eo_encoder:
   _target_: src.models.components.eo_encoders.multimodal_encoder.MultiModalEncoder
@@ -21,6 +21,8 @@ prediction_head:
 # Only the prediction head is trained; GeoClip encoder is frozen.
 trainable_modules: [prediction_head]
 
+metrics: ${metrics}
+
 optimizer:
   _target_: torch.optim.Adam
   _partial_: true

configs/model/heat_tabular_reg.yaml

@@ -12,7 +12,9 @@
 #   3. PredictiveRegressionModel.setup() calls
 #      self.eo_encoder.build_tabular_branch(self.trainer.datamodule.tabular_dim)
 
-_target_: src.models.predictive_model_regression.PredictiveRegressionModel
+_target_: src.models.predictive_model.PredictiveModel
+
+metrics: ${metrics}
 
 eo_encoder:
   _target_: src.models.components.eo_encoders.multimodal_encoder.MultiModalEncoder

configs/model/predictive_cnn_s2.yaml

@@ -11,6 +11,8 @@ prediction_head:
 
 trainable_modules: [eo_encoder, prediction_head]
 
+metrics: ${metrics}
+
 optimizer:
   _target_: torch.optim.Adam
   _partial_: true

configs/model/predictive_geoclip.yaml

@@ -8,6 +8,8 @@ prediction_head:
 
 trainable_modules: [prediction_head]
 
+metrics: ${metrics}
+
 optimizer:
   _target_: torch.optim.Adam
   _partial_: true

configs/train.yaml

@@ -4,14 +4,15 @@
 # order of defaults determines the order in which configs override each other
 defaults:
   - _self_
-  - data: butterfly_coords_text
+  - data: butterfly_coords
   - model: predictive_geoclip
   - callbacks: default
   - logger: ${oc.env:LOGGER,wandb}
   - trainer: ${oc.env:TRAINER_PROFILE,default}
   - paths: ${oc.env:STORAGE_MODE,local}
   - extras: default
   - hydra: default
+  - metrics: butterfly_predictive
 
   # experiment configs allow for version control of specific hyperparameters
   # e.g. best hyperparameters for given model and datamodule

src/data/base_datamodule.py

@@ -68,11 +68,7 @@ def __init__(
 
     @property
     def tabular_dim(self):
-        dataset = self.data_train
-        # Unwrap Subset wrappers (e.g. from random_split)
-        while hasattr(dataset, "dataset"):
-            dataset = dataset.dataset
-        return dataset.tabular_dim
+        return self.dataset.tabular_dim
 
     @property
     def num_classes(self) -> int:

src/data/base_dataset.py

@@ -74,15 +74,17 @@ def __init__(
         self.df: pd.DataFrame = pd.read_csv(path_csv)
 
         # Other attributes or placeholders
-        self.seed = seed
+        self.pooch_cli = None
         self.num_classes = None
-        self.mode: str = mode  # 'train', 'val', 'test'
+        self.tabular_dim = None
+        self.seed = seed
         self.use_target_data: bool = use_target_data
         self.use_aux_data: bool = use_aux_data
-        self.records: dict[str, Any] = self.get_records()
-        self.pooch_cli = None
         self.use_features = use_features
 
+        self.mode: str = mode  # 'train', 'val', 'test'
+        self.records: dict[str, Any] = self.get_records()
+
     @final
     def get_records(self) -> dict[str, Any]:
         """Gets record dictionary from the dataframe based on what is needed for the model (aux,
@@ -115,10 +117,12 @@ def get_records(self) -> dict[str, Any]:
         if self.use_aux_data:
             self.aux_names = [c for c in self.df.columns if "aux_" in c]
             columns.extend(self.aux_names)
-        
+
         # Include tabular features
-        self.feat_names = [c for c in self.df.columns if c.startswith("feat_")]
-        columns.extend(self.feat_names)
+        if self.use_features:
+            self.feat_names = [c for c in self.df.columns if c.startswith("feat_")]
+            columns.extend(self.feat_names)
+            self.tabular_dim = len(self.feat_names)
 
         return self.df.loc[:, columns].to_dict("records")
 
@@ -127,13 +131,6 @@ def __len__(self) -> int:
         """Returns the length of the dataset."""
         return len(self.records)
 
-    @final
-    @property
-    def tabular_dim(self) -> int:
-        if not self.use_features:
-            return 0
-        return len(getattr(self, "feat_names", []))
-
     @abstractmethod
     def __getitem__(self, idx: int) -> Dict[str, Any]:
         """Returns a single item from the dataset."""

src/data/heat_guatemala_dataset.py

@@ -1,5 +1,4 @@
-"""
-Heat Guatemala LST dataset.
+"""Heat Guatemala LST dataset.
 
 Location: src/data/heat_guatemala_dataset.py
 
@@ -21,8 +20,7 @@
 
 
 class HeatGuatemalaDataset(BaseDataset):
-    """
-    Dataset for the urban heat island use case (Guatemala City, LST regression).
+    """Dataset for the urban heat island use case (Guatemala City, LST regression).
 
     CSV layout expected (produced by scripts/make_model_ready_heat_guatemala.py):
       - name_loc          : unique location identifier
@@ -63,15 +61,6 @@ def __init__(
             use_features=use_features,
         )
 
-    # ------------------------------------------------------------------
-    # Properties
-    # ------------------------------------------------------------------
-
-    @property
-    def tabular_dim(self) -> int:
-        """Number of tabular features (feat_* columns). 0 if none."""
-        return len(self.feat_names)
-
     # ------------------------------------------------------------------
     # Required overrides
     # ------------------------------------------------------------------

src/data_preprocessing/make_model_ready_heat_guatemala.py

@@ -1,6 +1,5 @@
-"""
-Build model-ready CSV for the Heat Guatemala use case (data/heat_guatemala/model_ready_heat_guatemala.csv).
-"""
+"""Build model-ready CSV for the Heat Guatemala use case
+(data/heat_guatemala/model_ready_heat_guatemala.csv)."""
 
 import argparse
 import re
@@ -87,19 +86,23 @@ def main(source_csv: str, out_csv: str, drop_zero_lst: bool = True) -> None:
     # ------------------------------------------------------------------ #
     # 2. Build output skeleton                                            #
     # ------------------------------------------------------------------ #
-    out = pd.DataFrame({
-        "name_loc":   [f"heat_{i:06d}" for i in range(len(df))],
-        "lat":        df["LAT"].astype(float),
-        "lon":        df["LONG"].astype(float),
-        "target_lst": df[target_col].astype(float),
-    })
+    out = pd.DataFrame(
+        {
+            "name_loc": [f"heat_{i:06d}" for i in range(len(df))],
+            "lat": df["LAT"].astype(float),
+            "lon": df["LONG"].astype(float),
+            "target_lst": df[target_col].astype(float),
+        }
+    )
 
     # Verify target is clean
     assert out["target_lst"].isna().sum() == 0, "BUG: NaN in target after cleaning"
-    print(f"Target stats:  mean={out['target_lst'].mean():.2f}  "
-          f"std={out['target_lst'].std():.2f}  "
-          f"min={out['target_lst'].min():.2f}  "
-          f"max={out['target_lst'].max():.2f}")
+    print(
+        f"Target stats:  mean={out['target_lst'].mean():.2f}  "
+        f"std={out['target_lst'].std():.2f}  "
+        f"min={out['target_lst'].min():.2f}  "
+        f"max={out['target_lst'].max():.2f}"
+    )
 
     # ------------------------------------------------------------------ #
     # 3. Numeric features — impute or drop                                #
@@ -168,7 +171,7 @@ def main(source_csv: str, out_csv: str, drop_zero_lst: bool = True) -> None:
 if __name__ == "__main__":
     ap = argparse.ArgumentParser()
     ap.add_argument("--source_csv", required=True)
-    ap.add_argument("--out_csv",    required=True)
+    ap.add_argument("--out_csv", required=True)
     ap.add_argument("--drop_zero_lst", type=lambda x: x.lower() != "false", default=True)
     args = ap.parse_args()
-    main(args.source_csv, args.out_csv, args.drop_zero_lst)
+    main(args.source_csv, args.out_csv, args.drop_zero_lst)