From 44dbcfe9a8880eee66a6839ab9b4a64e08e9025f Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Mon, 6 Jul 2026 20:49:05 -0700
Subject: [PATCH 1/8] Unmixing regression first pass

---
 models/spectf/training_spectf_unmixing.py | 380 ++++++++++++++++++++++
 models/spectf/v1/spectf_unmixing.yaml     |  33 ++
 src/cover_class/dataloader/dataloader.py  |  25 +-
 src/cover_class/train.py                  |   8 +-
 4 files changed, 437 insertions(+), 9 deletions(-)
 create mode 100644 models/spectf/training_spectf_unmixing.py
 create mode 100644 models/spectf/v1/spectf_unmixing.yaml

diff --git a/models/spectf/training_spectf_unmixing.py b/models/spectf/training_spectf_unmixing.py
new file mode 100644
index 0000000..6da4a10
--- /dev/null
+++ b/models/spectf/training_spectf_unmixing.py
@@ -0,0 +1,380 @@
+import os
+
+# Set CUBLAS_WORKSPACE_CONFIG for deterministic CUDA behavior
+# pylint: disable=wrong-import-position
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+import getpass
+from datetime import datetime
+import rich_click as click
+import yaml
+import wandb
+import matplotlib.pyplot as plt
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import torch.nn.functional as F
+import schedulefree
+
+from cover_class.train import setup_training_from_config, make_simulation_test_set, banddef_from_config #type: ignore
+from cover_class.utils import seed as sseed, ood_test_set_from_config #type: ignore
+from cover_class.reporting import ModelConfig, Report #type: ignore
+
+from spectf.model import SpecTfEncoder
+from spectf.utils import get_device
+
+ENV_VAR_PREFIX = 'COVER_CLASS_TRAIN_'
+
+class TestDataset(Dataset):
+    def __init__(self, test_X, test_Y):
+        super().__init__()
+
+        self.test_X = test_X
+        self.test_Y = test_Y
+
+    def __len__(self):
+        return len(self.test_Y)
+
+    def __getitem__(self, idx):
+        return self.test_X[idx], self.test_Y[idx]
+
+
+class FocalCategoricalCrossEntropy(nn.Module):
+    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
+        super(FocalCategoricalCrossEntropy, self).__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, inputs, targets):
+        # inputs: (batch_size, n_classes) logits (pre-softmax)
+        # targets: (batch_size, n_classes) fractional ground truth (sum to 1)
+
+        # Compute log softmax
+        log_probs = F.log_softmax(inputs, dim=-1)
+
+        # Standard categorical cross-entropy: -sum(target * log_prob)
+        ce_loss = -(targets * log_probs).sum(dim=-1)
+
+        # Compute p_t for focal weighting
+        probs = torch.exp(log_probs)
+        p_t = (targets * probs).sum(dim=-1)  # Probability assigned to true distribution
+
+        # Focal weight: (1 - p_t)^gamma
+        focal_weight = (1 - p_t) ** self.gamma
+
+        # Apply focal weight
+        focal_loss = focal_weight * ce_loss
+
+        # Apply alpha weighting if specified
+        if self.alpha is not None:
+            # Weight by target distribution's average confidence
+            alpha_weight = self.alpha * targets.max(dim=-1)[0] + (1 - self.alpha) * (1 - targets.max(dim=-1)[0])
+            focal_loss = alpha_weight * focal_loss
+
+        if self.reduction == 'mean':
+            return focal_loss.mean()
+        if self.reduction == 'sum':
+            return focal_loss.sum()
+
+        return focal_loss
+
+
+@click.command()
+@click.option(
+    "--outdir",
+    required=True,
+    type=click.Path(exists=True, dir_okay=True, file_okay=False),
+    help="Output file directory.",
+    envvar=f'{ENV_VAR_PREFIX}OUTDIR'
+)
+@click.option(
+    "--data-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the dataloader.",
+    envvar=f'{ENV_VAR_PREFIX}_DATA_CONFIG'
+)
+@click.option(
+    "--model-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the model architecture.",
+    envvar=f'{ENV_VAR_PREFIX}_MODEL_CONFIG'
+)
+@click.option(
+    "--simulated-test-set-size",
+    required=False,
+    type=int,
+    default=100_000,
+    help="Number of rows to generate for the simulated test set.",
+    envvar=f'{ENV_VAR_PREFIX}_SIMULATED_TEST_SET_SIZE'
+)
+@click.option(
+    "--focal-alpha",
+    required=False,
+    default="0.25",
+    help="Focal loss alpha parameter. Set to 'None' to disable.",
+    envvar=f'{ENV_VAR_PREFIX}_FOCAL_ALPHA'
+)
+@click.option(
+    "--focal-gamma",
+    required=False,
+    type=float,
+    default=2.0,
+    help="Focal loss gamma parameter.",
+    envvar=f'{ENV_VAR_PREFIX}_FOCAL_GAMMA'
+)
+def run_pipeline_unmixing(
+        outdir: str,
+        data_config: str,
+        model_config: str,
+        simulated_test_set_size: int = 100_000,
+        focal_alpha: str = "0.25",
+        focal_gamma: float = 2.0
+    ):
+
+    with open(model_config, 'r', encoding='utf-8') as f:
+        m_config = yaml.safe_load(f)
+
+    # Load data config to get min_frac threshold
+    with open(data_config, 'r', encoding='utf-8') as f:
+        d_config = yaml.safe_load(f)
+
+    min_frac_threshold = d_config.get('simulation', {}).get('min_frac', 0.2)
+
+    dataloader, test_X, test_Y = setup_training_from_config(
+        data_config,
+        m_config['batch_size'],
+        shuffle=True,
+        seed=m_config['random_seed'],
+        subsampled_files_outdir=outdir,
+        return_fractions=True,  # Enable fraction mode
+        misc_dataloader_params={'num_workers': m_config['training']['num_workers']})
+
+    banddef = banddef_from_config(data_config)
+
+    # create simulation eval set
+    sseed(m_config['random_seed'])
+    simulation_x_test, simulation_y_labels, simulation_y_fractions = make_simulation_test_set(
+        dataloader, test_X, test_Y, simulated_test_set_size, one_hot_encode=False
+    )
+
+    # Test set dataloader
+    test_dataset = TestDataset(simulation_x_test, simulation_y_fractions)
+    test_dataloader = DataLoader(test_dataset, batch_size=m_config['batch_size'], shuffle=False)
+
+    # Validation set dataloader for OOD evaluation
+    ood_test_set_x, ood_test_set_y = ood_test_set_from_config(data_config)
+    ood_dataset = TestDataset(ood_test_set_x, ood_test_set_y)
+    ood_dataloader = DataLoader(ood_dataset, batch_size=m_config['batch_size'], shuffle=False)
+
+    # hardcoded GPU 0
+    device = get_device(0)
+
+    # model definition
+    model = SpecTfEncoder(banddef.to(dtype=torch.float32, device=device),
+                         dim_output=m_config['model']['dim_output'],
+                         num_heads=m_config['model']['num_heads'],
+                         dim_proj=m_config['model']['dim_proj'],
+                         dim_ff=m_config['model']['dim_ff'],
+                         dropout=m_config['model']['dropout'],
+                         agg=m_config['model']['agg'],
+                         use_residual=m_config['model']['use_residual'],
+                         num_layers=m_config['model']['num_layers']).to(device)
+
+    # criterion: KLDivLoss or FocalCategoricalCrossEntropy
+    alpha_val = None if focal_alpha == "None" else float(focal_alpha)
+    if alpha_val is None and focal_gamma == 0.0:
+        # Standard categorical cross-entropy (via KLDivLoss)
+        # Note: KLDivLoss expects log-probabilities as input, so we need log_softmax
+        criterion = nn.KLDivLoss(reduction='batchmean')
+        use_kl_div = True
+    else:
+        criterion = FocalCategoricalCrossEntropy(alpha=alpha_val, gamma=focal_gamma)
+        use_kl_div = False
+
+    optimizer = schedulefree.AdamWScheduleFree(
+        (p for p in model.parameters() if p.requires_grad),
+        lr=m_config['training']['learning_rate'],
+        warmup_steps=m_config['training']['warmup_steps']
+    )
+
+    # W&B
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    run = wandb.init(
+        entity=m_config['wandb']['entity'],
+        project=m_config['wandb']['project'],
+        name=f"{timestamp}_unmixing",
+        dir='./',
+        config={
+            "outdir": outdir,
+            "data_config": data_config,
+            "model_config": model_config,
+            "simulated_test_set_size": simulated_test_set_size,
+            "focal_alpha": focal_alpha,
+            "focal_gamma": focal_gamma,
+            "min_frac_threshold": min_frac_threshold,
+            "task": "unmixing"
+        },
+        settings=wandb.Settings(_service_wait=300)
+    )
+
+    # For binary metrics, convert fractions to binary using min_frac threshold
+    simulation_y_test_binary = (simulation_y_fractions >= min_frac_threshold).astype(float)
+
+    report = Report(
+        outdir=outdir,
+        config=data_config,
+        author=getpass.getuser(),
+        model_config=ModelConfig(
+            model=model,
+            model_name=SpecTfEncoder.__name__,
+            hyperparams={
+                "learning_rate": m_config['training']['learning_rate'],
+                "batch_size": m_config['batch_size'],
+                "optimizer": optimizer.__class__.__name__,
+                "focal_alpha": alpha_val,
+                "focal_gamma": focal_gamma,
+                "min_frac_threshold": min_frac_threshold,
+                "params": m_config['model']
+            },
+        ),
+        Y_test=simulation_y_test_binary,
+        Y_ood_test=ood_test_set_y,
+        random_seed=m_config['random_seed'],
+        run_name=timestamp
+    )
+
+    # Epoch-based training loop
+    total_epochs = m_config['training']['total_epochs']
+    steps_per_epoch = m_config['training']['steps_per_epoch']
+    bs = m_config['batch_size']
+
+    ds = report.config['datasets'] # type: ignore
+    class_names = [str(c) for c in ds.keys() if ds[c] is not None and len(ds[c])]
+
+    for epoch in range(total_epochs):
+        model.train()
+        optimizer.train()
+        epoch_loss = 0.0
+        step = 0
+        for batch_X, batch_Y in dataloader:
+            batch_X = batch_X.to(device=device, dtype=torch.float32)
+            batch_X = torch.unsqueeze(batch_X, -1)
+            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
+
+            optimizer.zero_grad()
+            logits = model(batch_X)
+
+            if use_kl_div:
+                # KLDivLoss expects log-probabilities as input
+                log_probs = F.log_softmax(logits, dim=-1)
+                loss = criterion(log_probs, batch_Y)
+            else:
+                loss = criterion(logits, batch_Y)
+
+            loss.backward()
+            optimizer.step()
+
+            nats = loss.cpu().item()
+            epoch_loss += nats
+            run.log({"loss_train": nats, "step": step + epoch * steps_per_epoch})
+
+            step += 1
+            if step >= steps_per_epoch:
+                break
+
+        avg_epoch_loss = epoch_loss / steps_per_epoch
+        run.log({"loss_train_epoch": avg_epoch_loss, "epoch": epoch})
+
+        # Save model at each epoch
+        if (epoch + 1) % 10 == 0:
+            torch.save(model.state_dict(), os.path.join(outdir, f'model_epoch{epoch+1}.pth'))
+
+        # Test loop
+        model.eval()
+        optimizer.eval()
+        y_hat_fractions = np.zeros_like(simulation_y_fractions, dtype=float)
+        test_loss_sum = 0.0
+        test_batches = 0
+        for i, (batch_X, batch_Y) in enumerate(test_dataloader):
+            batch_X = batch_X.to(device=device, dtype=torch.float32)
+            batch_X = torch.unsqueeze(batch_X, -1)
+            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
+            with torch.no_grad():
+                logits = model(batch_X)
+
+                if use_kl_div:
+                    log_probs = F.log_softmax(logits, dim=-1)
+                    batch_loss = criterion(log_probs, batch_Y)
+                else:
+                    batch_loss = criterion(logits, batch_Y)
+
+                test_loss_sum += batch_loss.cpu().item()
+                test_batches += 1
+
+                # Apply softmax to get fractions
+                batch_y_hat = torch.softmax(logits, dim=-1)
+                batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
+                batch_len = len(batch_y_hat)
+                y_hat_fractions[i*bs:i*bs+batch_len] = batch_y_hat
+        avg_test_loss = test_loss_sum / test_batches if test_batches > 0 else float('nan')
+
+        # Convert fractions to binary for metrics
+        y_hat_binary = (y_hat_fractions >= min_frac_threshold).astype(float)
+
+        # OOD loop - still uses binary targets
+        y_hat_ood = np.zeros_like(ood_test_set_y, dtype=float)
+        ood_loss_sum = 0.0
+        ood_batches = 0
+        for i, (batch_X, batch_Y) in enumerate(ood_dataloader):
+            batch_X = batch_X.to(device=device, dtype=torch.float32)
+            batch_X = torch.unsqueeze(batch_X, -1)
+            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
+            with torch.no_grad():
+                logits = model(batch_X)
+
+                # For OOD, we don't have fractions, so skip loss computation
+                # (or treat it as binary classification)
+
+                # Apply softmax to get fractions
+                batch_y_hat = torch.softmax(logits, dim=-1)
+                batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
+                batch_len = len(batch_y_hat)
+                y_hat_ood[i*bs:i*bs+batch_len] = batch_y_hat
+
+        # Apply threshold to OOD predictions for binary metrics
+        y_hat_ood_binary = (y_hat_ood >= min_frac_threshold).astype(float)
+
+        # Calculate test set metrics using the best thresholds for the test set
+        _figs = []
+        test_metrics = report.generate_metrics(simulation_y_test_binary, y_hat_binary, None, _figs, class_names)
+        for f in _figs: plt.close(f)
+        del _figs
+
+        # Extract the thresholds used for the test set
+        test_thresholds = [test_metrics[class_name]['Threshold'] for class_name in class_names]
+
+        # Calculate OOD validation set metrics using the test set thresholds
+        _figs = []
+        ood_metrics = report.generate_metrics(ood_test_set_y, y_hat_ood_binary, test_thresholds, _figs, class_names)
+        for f in _figs: plt.close(f)
+        del _figs
+
+        run.log({
+            "test_metrics": test_metrics,
+            "ood_metrics": ood_metrics,
+            "epoch": epoch,
+            "loss_test_epoch": avg_test_loss,
+        })
+
+    # Generate report at end of training
+    report.make_report(y_hat_binary, y_hat_ood_binary, None, ood_overfit=False)
+
+if __name__ == "__main__":
+    # pylint: disable=no-value-for-parameter
+    run_pipeline_unmixing()
diff --git a/models/spectf/v1/spectf_unmixing.yaml b/models/spectf/v1/spectf_unmixing.yaml
new file mode 100644
index 0000000..63ed390
--- /dev/null
+++ b/models/spectf/v1/spectf_unmixing.yaml
@@ -0,0 +1,33 @@
+# Configuration for training SpecTf model with direct unmixing
+# This config trains the model to predict fractional class abundances directly
+# using softmax + categorical cross-entropy loss
+
+# DataLoader parameters
+batch_size: 64
+
+# Model parameters
+model:
+  dim_output: 5
+  num_heads: 32
+  dim_proj: 128
+  dim_ff: 512
+  dropout: 0.05
+  agg: "mean"
+  use_residual: false
+  num_layers: 2
+
+# Training parameters
+training:
+  learning_rate: 0.0001
+  warmup_steps: 10
+  steps_per_epoch: 1000
+  total_epochs: 1000
+  num_workers: 16
+
+# Random seed for reproducibility
+random_seed: 42
+
+# Weights & Biases logging
+wandb:
+  entity: "jpl-cmml"
+  project: "frac-cover"
diff --git a/src/cover_class/dataloader/dataloader.py b/src/cover_class/dataloader/dataloader.py
index 1847c00..f9fe23e 100644
--- a/src/cover_class/dataloader/dataloader.py
+++ b/src/cover_class/dataloader/dataloader.py
@@ -6,6 +6,7 @@
 
 from cover_class.simulation import args_from_config, SimulationArgs, DataArgs
 import cover_class.simulation as sim
+from cover_class.simulation.simulate import get_fractions_by_class
 from cover_class.utils import read_config
 
 
@@ -20,6 +21,7 @@ class OrchestratorDatasetArgs(Struct):
     static_labels: Optional[torch.Tensor]
 
     num_classes: int = field(default=0)
+    return_fractions: bool = field(default=False)
     _using_static: bool = field(default=False)
     _using_sim: bool   = field(default=False)
 
@@ -96,8 +98,12 @@ def make_one_hot(y:torch.Tensor) -> torch.Tensor:
                 self.static_samples_seen += len(idx)
                 if end >= len(self.args.static_data)-1: # type: ignore
                     self.__reset__()
-                
-                labels = make_one_hot(self.args.static_labels[idx]) # type: ignore
+
+                if self.args.return_fractions:
+                    # For static data, create one-hot-like fractions (1.0 for true class, 0.0 for others)
+                    labels = make_one_hot(self.args.static_labels[idx]).to(dtype=torch.float32) # type: ignore
+                else:
+                    labels = make_one_hot(self.args.static_labels[idx]) # type: ignore
                 self.batch_dirichlet_fraction_store = None
                 yield self.args.static_data[idx], labels # type: ignore
 
@@ -106,7 +112,12 @@ def make_one_hot(y:torch.Tensor) -> torch.Tensor:
                 # mypy doesn't catch self.args._using_sim
                 data, labels, fractions = sim.run_simulation(self.args.sim_config_args, self.args.sim_data_args) # type: ignore
                 self.batch_dirichlet_fraction_store = fractions
-                yield data, make_one_hot(labels)
+
+                if self.args.return_fractions:
+                    # Return fractions directly (already in correct format from run_simulation)
+                    yield data, fractions
+                else:
+                    yield data, make_one_hot(labels)
 
             else: raise StopIteration()
 
@@ -129,11 +140,12 @@ def __use_static_predicate__(self) -> bool:
 
 
 def dataloader_from_config(
-        config: Dict|str, 
+        config: Dict|str,
         spectra:FloatTensor,
         labels:LongTensor,
         batch_size:int,
-        shuffle: bool = True, 
+        shuffle: bool = True,
+        return_fractions: bool = False,
         misc_dataloader_params: dict = {},
     ) -> DataLoader:
 
@@ -146,7 +158,8 @@ def dataloader_from_config(
         sim_config_args,
         sim_data_args,
         spectra,
-        labels.long()
+        labels.long(),
+        return_fractions=return_fractions
     )
     ods = OrchestratorDataset(ods_args, shuffle)
     return DataLoader(ods, batch_size=None, **misc_dataloader_params)
diff --git a/src/cover_class/train.py b/src/cover_class/train.py
index 634e373..6e6d744 100644
--- a/src/cover_class/train.py
+++ b/src/cover_class/train.py
@@ -63,12 +63,13 @@ def train_test_from_config(config: str|Dict, seed: Optional[int] = None):
 
 
 def setup_training_from_config(
-        config: str|Dict, 
+        config: str|Dict,
         batch_size: int,
         shuffle: bool = True,
         seed: Optional[int] = None,
         subsampled_files_outdir: str = '',
         run_name: str = '',
+        return_fractions: bool = False,
         misc_dataloader_params: dict = {},
     ) -> Tuple[DataLoader, FloatTensor, Tensor]:
     """
@@ -86,7 +87,7 @@ def setup_training_from_config(
     for i, d in enumerate(config['datasets']):
         hdf5_list = config['datasets'][d]
         if hdf5_list is None: continue
-        
+
         # subsampling and train test split will happen on a per file basis
         for hdf5 in hdf5_list:
             with h5py.File(hdf5, 'r') as f:
@@ -113,11 +114,12 @@ def setup_training_from_config(
     test_spectra = test_spectra.to(torch.float32)
 
     odl = dataloader_from_config(
-        config, 
+        config,
         FloatTensor(train_spectra),
         LongTensor(train_labels.to(dtype=torch.long)),
         batch_size,
         shuffle,
+        return_fractions,
         misc_dataloader_params,
     )
 

From b547c13945bf1f84134921adbda490b1e08b65d0 Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Mon, 6 Jul 2026 21:27:39 -0700
Subject: [PATCH 2/8] type casting fix

---
 models/spectf/training_spectf_unmixing.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/models/spectf/training_spectf_unmixing.py b/models/spectf/training_spectf_unmixing.py
index 6da4a10..4356300 100644
--- a/models/spectf/training_spectf_unmixing.py
+++ b/models/spectf/training_spectf_unmixing.py
@@ -224,6 +224,7 @@ def run_pipeline_unmixing(
     )
 
     # For binary metrics, convert fractions to binary using min_frac threshold
+    simulation_y_fractions = simulation_y_fractions.cpu().numpy()
     simulation_y_test_binary = (simulation_y_fractions >= min_frac_threshold).astype(float)
 
     report = Report(

From 7ca2cf96b0566897739982cdec4c8887b37af608 Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Mon, 6 Jul 2026 21:51:07 -0700
Subject: [PATCH 3/8] Pass num_classes instead of inferring it from data
 unreliably

---
 src/cover_class/simulation/simulate.py | 9 +++++----
 1 file changed, 5 insertions(+), 4 deletions(-)

diff --git a/src/cover_class/simulation/simulate.py b/src/cover_class/simulation/simulate.py
index 1623d73..dfaa4b8 100644
--- a/src/cover_class/simulation/simulate.py
+++ b/src/cover_class/simulation/simulate.py
@@ -113,7 +113,7 @@ def run_simulation(
             filtered_n_components_per_class = (cumsum_n_components[:, 1:] - cumsum_n_components[:, :-1]).to(dtype=torch.int16, device=device) #type: ignore
 
             # Sum the Dirichlet fractions by class for each row, then look at the target class.
-            fracs_by_class = get_fractions_by_class(filtered_n_components_per_class, classes, dirich_fractions)
+            fracs_by_class = get_fractions_by_class(filtered_n_components_per_class, classes, dirich_fractions, sim_args.n_classes)
 
             keep_rows = ((fracs_by_class[:, force_class] >= lo) & (fracs_by_class[:, force_class] <= hi))
 
@@ -162,6 +162,7 @@ def run_simulation(
                 filtered_n_components_per_class,
                 classes,
                 dirich_fractions,
+                sim_args.n_classes,
             )
             return resulting_real_spectra, classes.long(), fracs_by_class # type: ignore[return-value]
         return resulting_real_spectra, classes.long(), None # type: ignore[return-value]
@@ -404,13 +405,14 @@ def make_positive_definite(A: Tensor, min_eigen=1e-8) -> FloatTensor:
 
 
 def get_fractions_by_class(
-        filtered_n_components_per_class: ShortTensor, 
+        filtered_n_components_per_class: ShortTensor,
         classes: CharTensor,
         dirich_fractions: FloatTensor,
+        num_classes: int,
 
     ) -> FloatTensor:
     # Returns a (n_iter, n_classes) matrix of the sum of dirichlet constributions per class
-    
+
     # first get the dirichlet fractions by number of components
     n_iters, n_classes_per_sim = filtered_n_components_per_class.shape
     n_max_sim_comps = dirich_fractions.size(1)
@@ -424,7 +426,6 @@ def get_fractions_by_class(
     # then assign each of those to a class (one-hot)
     valid = classes >= 0
     cls = classes.clamp_min(0)
-    num_classes = int(cls[valid].max().item()) + 1 if valid.any() else 0
     result = summed_fracs.new_zeros(summed_fracs.size(0), num_classes)
     result.scatter_add_(1, cls.to(dtype=torch.int32), summed_fracs * valid)
 

From 160af4e2a6dff1017918d3ced8151e8b543aa0fd Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Tue, 14 Jul 2026 12:00:17 -0700
Subject: [PATCH 4/8] Initial reg eval script

---
 models/spectf/report_spectf_unmixing.py | 245 ++++++++++++++++++++++++
 1 file changed, 245 insertions(+)
 create mode 100644 models/spectf/report_spectf_unmixing.py

diff --git a/models/spectf/report_spectf_unmixing.py b/models/spectf/report_spectf_unmixing.py
new file mode 100644
index 0000000..56bcec2
--- /dev/null
+++ b/models/spectf/report_spectf_unmixing.py
@@ -0,0 +1,245 @@
+import os
+
+# Set CUBLAS_WORKSPACE_CONFIG for deterministic CUDA behavior
+# pylint: disable=wrong-import-position
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+from datetime import datetime
+import rich_click as click
+import yaml
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, DataLoader
+import matplotlib.pyplot as plt
+from sklearn.metrics import r2_score, mean_absolute_error, mean_squared_error
+
+from cover_class.train import setup_training_from_config, make_simulation_test_set, banddef_from_config
+from cover_class.utils import seed as sseed
+
+from spectf.model import SpecTfEncoder
+from spectf.utils import get_device
+
+ENV_VAR_PREFIX = 'COVER_CLASS_TRAIN_'
+
+
+class TestDataset(Dataset):
+    def __init__(self, test_X, test_Y):
+        super().__init__()
+        self.test_X = test_X
+        self.test_Y = test_Y
+
+    def __len__(self):
+        return len(self.test_Y)
+
+    def __getitem__(self, idx):
+        return self.test_X[idx], self.test_Y[idx]
+
+
+@click.command()
+@click.option(
+    "--outdir",
+    required=True,
+    type=click.Path(exists=True, dir_okay=True, file_okay=False),
+    help="Output file directory.",
+    envvar=f'{ENV_VAR_PREFIX}OUTDIR'
+)
+@click.option(
+    "--data-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the dataloader.",
+    envvar=f'{ENV_VAR_PREFIX}_DATA_CONFIG'
+)
+@click.option(
+    "--model-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the model architecture.",
+    envvar=f'{ENV_VAR_PREFIX}_MODEL_CONFIG'
+)
+@click.option(
+    "--model-weights",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the model weights file (.pth).",
+    envvar=f'{ENV_VAR_PREFIX}_MODEL_WEIGHTS'
+)
+@click.option(
+    "--simulated-test-set-size",
+    required=False,
+    type=int,
+    default=100_000,
+    help="Number of rows to generate for the simulated test set.",
+    envvar=f'{ENV_VAR_PREFIX}_SIMULATED_TEST_SET_SIZE'
+)
+def run_unmixing_evaluation(
+        outdir: str,
+        data_config: str,
+        model_config: str,
+        model_weights: str,
+        simulated_test_set_size: int = 100_000,
+    ):
+
+    # Load model config
+    with open(model_config, 'r', encoding='utf-8') as f:
+        m_config = yaml.safe_load(f)
+
+    # Load data config
+    with open(data_config, 'r', encoding='utf-8') as f:
+        d_config = yaml.safe_load(f)
+
+    # Get class names from data config
+    ds = d_config['datasets']
+    class_names = [str(c) for c in ds.keys() if ds[c] is not None and len(ds[c])]
+    print(f"Classes: {class_names}")
+
+    # Set up dataloader and test sets
+    dataloader, test_X, test_Y = setup_training_from_config(
+        data_config,
+        m_config['batch_size'],
+        shuffle=True,
+        seed=m_config['random_seed'],
+        subsampled_files_outdir=outdir,
+        return_fractions=True,  # Enable fraction mode for unmixing
+        misc_dataloader_params={'num_workers': m_config['training']['num_workers']})
+
+    # Get banddef (wavelength definitions)
+    banddef = banddef_from_config(data_config)
+
+    # Create the simulated test set
+    sseed(m_config['random_seed'])
+    simulation_x_test, simulation_y_labels, simulation_y_fractions = make_simulation_test_set(
+        dataloader, test_X, test_Y, simulated_test_set_size, one_hot_encode=False
+    )
+
+    # Create a dataset/dataloader to feed the test set in batches
+    test_dataset = TestDataset(simulation_x_test, simulation_y_fractions)
+    test_dataloader = DataLoader(test_dataset, batch_size=m_config['batch_size'], shuffle=False)
+
+    # hardcoded GPU 0
+    device = get_device(0)
+
+    # Model definition
+    model = SpecTfEncoder(banddef.to(dtype=torch.float32, device=device),
+                         dim_output=m_config['model']['dim_output'],
+                         num_heads=m_config['model']['num_heads'],
+                         dim_proj=m_config['model']['dim_proj'],
+                         dim_ff=m_config['model']['dim_ff'],
+                         dropout=m_config['model']['dropout'],
+                         agg=m_config['model']['agg'],
+                         use_residual=m_config['model']['use_residual'],
+                         num_layers=m_config['model']['num_layers']).to(device)
+
+    # Load model weights
+    print(f"Loading model weights from {model_weights}")
+    model.load_state_dict(torch.load(model_weights, map_location=device))
+    model.eval()
+
+    # Batch size
+    bs = m_config['batch_size']
+
+    # Test loop - get predictions
+    print("Evaluating on simulated test set...")
+    y_hat_fractions = np.zeros_like(simulation_y_fractions, dtype=float)
+
+    for i, (batch_X, _) in enumerate(test_dataloader):
+        batch_X = batch_X.to(device=device, dtype=torch.float32)
+        batch_X = torch.unsqueeze(batch_X, -1)
+        with torch.no_grad():
+            logits = model(batch_X)
+            # Apply softmax to get fractions
+            batch_y_hat = torch.softmax(logits, dim=-1)
+            batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
+            batch_len = len(batch_y_hat)
+            y_hat_fractions[i*bs:i*bs+batch_len] = batch_y_hat
+
+    # Convert to numpy for metrics computation
+    y_true = simulation_y_fractions.cpu().numpy() if torch.is_tensor(simulation_y_fractions) else simulation_y_fractions
+    y_pred = y_hat_fractions
+
+    # Compute metrics per class
+    print("\n" + "="*80)
+    print("REGRESSION METRICS PER CLASS")
+    print("="*80)
+
+    metrics = {}
+    for i, class_name in enumerate(class_names):
+        r2 = r2_score(y_true[:, i], y_pred[:, i])
+        mae = mean_absolute_error(y_true[:, i], y_pred[:, i])
+        rmse = np.sqrt(mean_squared_error(y_true[:, i], y_pred[:, i]))
+
+        metrics[class_name] = {
+            'R²': r2,
+            'MAE': mae,
+            'RMSE': rmse
+        }
+
+        print(f"\n{class_name.upper()}")
+        print(f"  R² Score:  {r2:.4f}")
+        print(f"  MAE:       {mae:.4f}")
+        print(f"  RMSE:      {rmse:.4f}")
+
+    # Create 5-panel scatter plot
+    fig, axes = plt.subplots(1, 5, figsize=(20, 4))
+
+    for i, (ax, class_name) in enumerate(zip(axes, class_names)):
+        # Scatter plot
+        ax.scatter(y_true[:, i], y_pred[:, i], alpha=0.3, s=1, c='blue', edgecolors='none')
+
+        # 1:1 line
+        ax.plot([0, 1], [0, 1], 'r--', linewidth=2, label='1:1 line')
+
+        # Labels and title
+        ax.set_xlabel('True Fraction', fontsize=12)
+        ax.set_ylabel('Predicted Fraction', fontsize=12)
+        ax.set_title(f'{class_name}', fontsize=14, fontweight='bold')
+
+        # Set axis limits
+        ax.set_xlim([0, 1])
+        ax.set_ylim([0, 1])
+        ax.set_aspect('equal')
+
+        # Add metrics as text
+        metrics_text = (
+            f"R² = {metrics[class_name]['R²']:.3f}\n"
+            f"MAE = {metrics[class_name]['MAE']:.3f}\n"
+            f"RMSE = {metrics[class_name]['RMSE']:.3f}"
+        )
+        ax.text(0.05, 0.95, metrics_text, transform=ax.transAxes,
+                fontsize=10, verticalalignment='top',
+                bbox=dict(boxstyle='round', facecolor='white', alpha=0.8))
+
+        # Grid
+        ax.grid(True, alpha=0.3)
+        ax.legend(loc='lower right', fontsize=9)
+
+    plt.tight_layout()
+
+    # Save figure
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_path = os.path.join(outdir, f"{timestamp}_unmixing_scatter_plots.png")
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    print(f"\n{'='*80}")
+    print(f"Scatter plot saved to: {output_path}")
+    print(f"{'='*80}\n")
+
+    # Save metrics to text file
+    metrics_path = os.path.join(outdir, f"{timestamp}_unmixing_metrics.txt")
+    with open(metrics_path, 'w', encoding='utf-8') as f:
+        f.write("REGRESSION METRICS PER CLASS\n")
+        f.write("="*80 + "\n\n")
+        for class_name in class_names:
+            f.write(f"{class_name.upper()}\n")
+            f.write(f"  R² Score:  {metrics[class_name]['R²']:.6f}\n")
+            f.write(f"  MAE:       {metrics[class_name]['MAE']:.6f}\n")
+            f.write(f"  RMSE:      {metrics[class_name]['RMSE']:.6f}\n\n")
+
+    print(f"Metrics saved to: {metrics_path}")
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    # pylint: disable=no-value-for-parameter
+    run_unmixing_evaluation()

From b2b4ada01d922341723efb1136ed581c61ee473f Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Fri, 24 Jul 2026 22:16:55 -0700
Subject: [PATCH 5/8] OOD training injection implementation

---
 models/spectf/training_spectf_ood.py     | 333 +++++++++++++++++++++++
 models/spectf/v1/data_config.yaml        |   2 +
 src/cover_class/dataloader/dataloader.py | 103 ++++++-
 src/cover_class/train.py                 |  19 +-
 src/cover_class/utils.py                 |  26 +-
 5 files changed, 459 insertions(+), 24 deletions(-)
 create mode 100644 models/spectf/training_spectf_ood.py

diff --git a/models/spectf/training_spectf_ood.py b/models/spectf/training_spectf_ood.py
new file mode 100644
index 0000000..6c29d2a
--- /dev/null
+++ b/models/spectf/training_spectf_ood.py
@@ -0,0 +1,333 @@
+import os
+
+# Set CUBLAS_WORKSPACE_CONFIG for deterministic CUDA behavior
+# pylint: disable=wrong-import-position
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+import getpass
+from datetime import datetime
+import rich_click as click
+import yaml
+import wandb
+import matplotlib.pyplot as plt
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.optim as optim
+from torch.utils.data import Dataset, DataLoader
+import torch.nn.functional as F
+import schedulefree
+
+from cover_class.train import setup_training_from_config, make_simulation_test_set, banddef_from_config #type: ignore
+from cover_class.utils import seed as sseed, ood_test_set_from_config #type: ignore
+from cover_class.reporting import ModelConfig, Report #type: ignore
+
+from spectf.model import SpecTfEncoder
+from spectf.utils import get_device
+
+ENV_VAR_PREFIX = 'COVER_CLASS_TRAIN_'
+
+class TestDataset(Dataset):
+    def __init__(self, test_X, test_Y):
+        super().__init__()
+
+        self.test_X = test_X
+        self.test_Y = test_Y
+
+    def __len__(self):
+        return len(self.test_Y)
+
+    def __getitem__(self, idx):
+        return self.test_X[idx], self.test_Y[idx]
+
+
+class FocalLoss(nn.Module):
+    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
+        super(FocalLoss, self).__init__()
+        self.alpha = alpha
+        self.gamma = gamma
+        self.reduction = reduction
+
+    def forward(self, inputs, targets):
+        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
+        pt = torch.exp(-bce_loss)
+        focal_loss = ((1 - pt) ** self.gamma) * bce_loss
+
+        if self.alpha is not None:
+            alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
+            focal_loss = alpha_t * focal_loss
+
+        if self.reduction == 'mean':
+            return focal_loss.mean()
+        if self.reduction == 'sum':
+            return focal_loss.sum()
+
+        return focal_loss
+
+@click.command()
+@click.option(
+    "--outdir",
+    required=True,
+    type=click.Path(exists=True, dir_okay=True, file_okay=False),
+    help="Output file directory.",
+    envvar=f'{ENV_VAR_PREFIX}OUTDIR'
+)
+@click.option(
+    "--data-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the dataloader.",
+    envvar=f'{ENV_VAR_PREFIX}_DATA_CONFIG'
+)
+@click.option(
+    "--model-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the model architecture.",
+    envvar=f'{ENV_VAR_PREFIX}_MODEL_CONFIG'
+)
+@click.option(
+    "--simulated-test-set-size",
+    required=False,
+    type=int,
+    default=100_000,
+    help="Number of rows to generate for the simulated test set.",
+    envvar=f'{ENV_VAR_PREFIX}_SIMULATED_TEST_SET_SIZE'
+)
+@click.option(
+    "--focal-alpha",
+    required=False,
+    default="0.25",
+    help="Focal loss alpha parameter. Set to 'None' to disable.",
+    envvar=f'{ENV_VAR_PREFIX}_FOCAL_ALPHA'
+)
+@click.option(
+    "--focal-gamma",
+    required=False,
+    type=float,
+    default=2.0,
+    help="Focal loss gamma parameter.",
+    envvar=f'{ENV_VAR_PREFIX}_FOCAL_GAMMA'
+)
+def run_pipeline_classifier(
+        outdir: str,
+        data_config: str,
+        model_config: str,
+        simulated_test_set_size: int = 100_000,
+        focal_alpha: str = "0.25",
+        focal_gamma: float = 2.0
+    ):
+
+    with open(model_config, 'r', encoding='utf-8') as f:
+        m_config = yaml.safe_load(f)
+
+    # inject_ood pulls the 'ood-train-set' spectra into the training dataloader as-is
+    # (unmixed, real labels) to intentionally induce overfitting on OOD data.
+    dataloader, test_X, test_Y = setup_training_from_config(
+        data_config,
+        m_config['batch_size'],
+        shuffle=True,
+        seed=m_config['random_seed'],
+        subsampled_files_outdir=outdir,
+        misc_dataloader_params={'num_workers': m_config['training']['num_workers']},
+        inject_ood=True)
+
+    banddef = banddef_from_config(data_config)
+
+    # create simulation eval set
+    sseed(m_config['random_seed'])
+    simulation_x_test, simulation_y_test, _ = make_simulation_test_set(dataloader, test_X, test_Y, simulated_test_set_size)
+
+    # Test set dataloader
+    test_dataset = TestDataset(simulation_x_test, simulation_y_test)
+    test_dataloader = DataLoader(test_dataset, batch_size=m_config['batch_size'], shuffle=False)
+
+    # Validation set dataloader for OOD evaluation
+    ood_test_set_x, ood_test_set_y = ood_test_set_from_config(data_config)
+    ood_dataset = TestDataset(ood_test_set_x, ood_test_set_y)
+    ood_dataloader = DataLoader(ood_dataset, batch_size=m_config['batch_size'], shuffle=False)
+
+    # hardcoded GPU 0
+    device = get_device(0)
+
+    # model definition
+    model = SpecTfEncoder(banddef.to(dtype=torch.float32, device=device),
+                         dim_output=m_config['model']['dim_output'],
+                         num_heads=m_config['model']['num_heads'],
+                         dim_proj=m_config['model']['dim_proj'],
+                         dim_ff=m_config['model']['dim_ff'],
+                         dropout=m_config['model']['dropout'],
+                         agg=m_config['model']['agg'],
+                         use_residual=m_config['model']['use_residual'],
+                         num_layers=m_config['model']['num_layers']).to(device)
+
+    # criterion = nn.BCEWithLogitsLoss()
+    alpha_val = None if focal_alpha == "None" else float(focal_alpha)
+    if alpha_val is None and focal_gamma == 0.0:
+        criterion = nn.BCEWithLogitsLoss()
+    else:
+        criterion = FocalLoss(alpha=alpha_val, gamma=focal_gamma)
+
+    optimizer = schedulefree.AdamWScheduleFree(
+        (p for p in model.parameters() if p.requires_grad),
+        lr=m_config['training']['learning_rate'],
+        warmup_steps=m_config['training']['warmup_steps']
+    )
+
+    # W&B
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    run = wandb.init(
+        entity=m_config['wandb']['entity'],
+        project=m_config['wandb']['project'],
+        name=timestamp,
+        dir='./',
+        config={
+            "outdir": outdir,
+            "data_config": data_config,
+            "model_config": model_config,
+            "simulated_test_set_size": simulated_test_set_size,
+            "focal_alpha": focal_alpha,
+            "focal_gamma": focal_gamma,
+            "inject_ood": True,
+        },
+        settings=wandb.Settings(_service_wait=300)
+    )
+
+    report = Report(
+        outdir=outdir,
+        config=data_config,
+        author=getpass.getuser(),
+        model_config=ModelConfig(
+            model=model,
+            model_name=SpecTfEncoder.__name__,
+            hyperparams={
+                "learning_rate": m_config['training']['learning_rate'],
+                "batch_size": m_config['batch_size'],
+                "optimizer": optimizer.__class__.__name__,
+                "focal_alpha": alpha_val,
+                "focal_gamma": focal_gamma,
+                "params": m_config['model']
+            },
+        ),
+        Y_test=simulation_y_test,
+        Y_ood_test=ood_test_set_y,
+        random_seed=m_config['random_seed'],
+        run_name=timestamp
+    )
+
+    # Epoch-based training loop
+    total_epochs = m_config['training']['total_epochs']
+    steps_per_epoch = m_config['training']['steps_per_epoch']
+    bs = m_config['batch_size']
+
+    ds = report.config['datasets'] # type: ignore
+    class_names = [str(c) for c in ds.keys() if ds[c] is not None and len(ds[c])]
+
+    for epoch in range(total_epochs):
+        model.train()
+        optimizer.train()
+        epoch_loss = 0.0
+        step = 0
+        for batch_X, batch_Y in dataloader:
+            batch_X = batch_X.to(device=device, dtype=torch.float32)
+            batch_X = torch.unsqueeze(batch_X, -1)
+            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
+
+            optimizer.zero_grad()
+            logits = model(batch_X)
+            # Injected OOD spectra carry a -1 sentinel for unknown/ambiguous entries;
+            # mask those out so they don't contribute to the loss during backprop.
+            mask = batch_Y >= 0
+            loss = criterion(logits[mask], batch_Y[mask])
+            loss.backward()
+            optimizer.step()
+
+            nats = loss.cpu().item()
+            epoch_loss += nats
+            run.log({"loss_train": nats, "step": step + epoch * steps_per_epoch})
+
+            step += 1
+            if step >= steps_per_epoch:
+                break
+
+        avg_epoch_loss = epoch_loss / steps_per_epoch
+        run.log({"loss_train_epoch": avg_epoch_loss, "epoch": epoch})
+
+        # Save model at each epoch
+        if (epoch + 1) % 10 == 0:
+            torch.save(model.state_dict(), os.path.join(outdir, f'model_epoch{epoch+1}.pth'))
+
+        # Test loop
+        model.eval()
+        optimizer.eval()
+        y_hat = np.zeros_like(simulation_y_test, dtype=float)
+        test_loss_sum = 0.0
+        test_batches = 0
+        for i, (batch_X, batch_Y) in enumerate(test_dataloader):
+            batch_X = batch_X.to(device=device, dtype=torch.float32)
+            batch_X = torch.unsqueeze(batch_X, -1)
+            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
+            with torch.no_grad():
+                logits = model(batch_X)
+                batch_loss = criterion(logits, batch_Y)
+                test_loss_sum += batch_loss.cpu().item()
+                test_batches += 1
+                batch_y_hat = torch.sigmoid(logits)
+                batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
+                batch_len = len(batch_y_hat)
+                y_hat[i*bs:i*bs+batch_len] = batch_y_hat
+        avg_test_loss = test_loss_sum / test_batches if test_batches > 0 else float('nan')
+
+        # OOD loop
+        y_hat_ood = np.zeros_like(ood_test_set_y, dtype=float)
+        ood_loss_sum = 0.0
+        ood_batches = 0
+        for i, (batch_X, batch_Y) in enumerate(ood_dataloader):
+            batch_X = batch_X.to(device=device, dtype=torch.float32)
+            batch_X = torch.unsqueeze(batch_X, -1)
+            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
+            with torch.no_grad():
+                logits = model(batch_X)
+                mask = ~torch.isnan(batch_Y)
+                if mask.any():
+                    batch_loss = criterion(logits[mask], batch_Y[mask])
+                    ood_loss_sum += batch_loss.cpu().item()
+                    ood_batches += 1
+                batch_y_hat = torch.sigmoid(logits)
+                batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
+                batch_len = len(batch_y_hat)
+                y_hat_ood[i*bs:i*bs+batch_len] = batch_y_hat
+
+        avg_ood_loss = ood_loss_sum / ood_batches if ood_batches > 0 else float('nan')
+
+
+        # Calculate test set metrics using the best thresholds for the test set
+        _figs = []
+        test_metrics = report.generate_metrics(simulation_y_test, y_hat, None, _figs, class_names)
+        for f in _figs: plt.close(f)
+        del _figs
+
+        # Extract the thresholds used for the test set
+        test_thresholds = [test_metrics[class_name]['Threshold'] for class_name in class_names]
+
+        # Calculate OOD validation set metrics using the test set thresholds
+        _figs = []
+        ood_metrics = report.generate_metrics(ood_test_set_y, y_hat_ood, test_thresholds, _figs, class_names)
+        for f in _figs: plt.close(f)
+        del _figs
+
+        run.log({
+            "test_metrics": test_metrics,
+            "ood_metrics": ood_metrics,
+            "epoch": epoch,
+            "loss_test_epoch": avg_test_loss,
+            "loss_ood_epoch": avg_ood_loss
+        })
+
+    # Generate report at end of training
+    report.make_report(y_hat, y_hat_ood, None, ood_overfit=False)
+
+if __name__ == "__main__":
+    # pylint: disable=no-value-for-parameter
+    run_pipeline_classifier()
diff --git a/models/spectf/v1/data_config.yaml b/models/spectf/v1/data_config.yaml
index 9c2b694..c297a78 100644
--- a/models/spectf/v1/data_config.yaml
+++ b/models/spectf/v1/data_config.yaml
@@ -18,6 +18,7 @@ datasets:
   water:
     - /emit-frac/datasets/20260517-sub/water_kmeans_subsampled_water_water.hdf5
 ood-test-set: /emit-frac/datasets/validation_20260317_wfids_noash_goodwl_subsampled.h5
+ood-train-set: /emit-frac/datasets/ood_train_20260317_wfids_noash_goodwl_subsampled.h5
 simulation:
   sim_mixture_probs_csv: /cover-class/src/cover_class/config/sim-mixture-probs.csv
   n_components:
@@ -93,3 +94,4 @@ subsample:
     n_samples: 10
 dataloader:
   percent-static-data: 0.0
+  percent-ood-data: 0.0
diff --git a/src/cover_class/dataloader/dataloader.py b/src/cover_class/dataloader/dataloader.py
index f9fe23e..b47b434 100644
--- a/src/cover_class/dataloader/dataloader.py
+++ b/src/cover_class/dataloader/dataloader.py
@@ -20,25 +20,56 @@ class OrchestratorDatasetArgs(Struct):
     static_data: Optional[torch.FloatTensor]
     static_labels: Optional[torch.Tensor]
 
+    # Real OOD spectra injected into training as-is (already multi-hot labels, never mixed/simulated).
+    ood_data: Optional[torch.FloatTensor] = field(default=None)
+    ood_labels: Optional[torch.Tensor] = field(default=None)
+    percent_ood: float = field(default=0.0)
+
     num_classes: int = field(default=0)
     return_fractions: bool = field(default=False)
     _using_static: bool = field(default=False)
     _using_sim: bool   = field(default=False)
+    _using_ood: bool   = field(default=False)
 
-    # This tells you the ids out of a 100 batches, which ones will be static/simulated
+    # This tells you the ids out of a 100 batches, which source each will draw from:
+    #   0 -> simulated, 1 -> static, 2 -> ood
     _method_selection_idxs: CharTensor = field(default_factory=lambda: CharTensor(torch.zeros(100, dtype=torch.int8)))
 
     def __post_init__(self):
         self._using_static = (self.static_labels is not None) and (self.static_data is not None)
         self._using_sim = (self.sim_config_args is not None) and (self.sim_data_args is not None)
-        assert self._using_static or self._using_sim, "Need to provide either simulation or static data arguments"
-        if self._using_static and not self._using_sim: self.percent_static = 1.00
-        elif not self._using_static: self.percent_static = 0.
-        assert (self.percent_static is not None) and (0. <= self.percent_static <= 1.00), "'percent_static' needs to be between [0, 1.]"
-
-        self.num_classes = self.sim_config_args.n_classes if self._using_sim else len(torch.unique(self.static_labels))
-
-        self._method_selection_idxs[:int(self.percent_static*100)] = 1
+        self._using_ood = (self.ood_labels is not None) and (self.ood_data is not None)
+        assert self._using_static or self._using_sim or self._using_ood, "Need to provide simulation, static, or ood data arguments"
+
+        if not self._using_static: self.percent_static = 0.
+        if not self._using_ood: self.percent_ood = 0.
+
+        # When there is no simulation source, the static/ood sources must account for all batches.
+        if not self._using_sim:
+            total = self.percent_static + self.percent_ood
+            assert total > 0, "Need a non-zero percent for static and/or ood data when simulation is unavailable"
+            self.percent_static /= total
+            self.percent_ood /= total
+
+        assert 0. <= self.percent_static <= 1.00, "'percent_static' needs to be between [0, 1.]"
+        assert 0. <= self.percent_ood <= 1.00, "'percent_ood' needs to be between [0, 1.]"
+        assert self.percent_static + self.percent_ood <= 1.0 + 1e-6, "'percent_static' + 'percent_ood' must be <= 1."
+
+        if self._using_sim:
+            self.num_classes = self.sim_config_args.n_classes
+        elif self._using_static:
+            self.num_classes = len(torch.unique(self.static_labels))
+        else:
+            self.num_classes = self.ood_labels.shape[1]
+
+        n_static = int(self.percent_static * 100)
+        n_ood = int(self.percent_ood * 100)
+        self._method_selection_idxs[:n_static] = 1
+        self._method_selection_idxs[n_static:n_static + n_ood] = 2
+        if not self._using_sim:
+            # Assign any slots left unallocated by rounding to an available static/ood source.
+            remaining = self._method_selection_idxs == 0
+            self._method_selection_idxs[remaining] = 1 if self._using_static else 2
         self._shuffle_method_selection_idxs()
 
     def _shuffle_method_selection_idxs(self):
@@ -66,17 +97,22 @@ class OrchestratorDataset(IterableDataset):
     static_epoch = 0
     static_epoch_step = 0
     static_samples_seen = 0
+    ood_epoch = 0
+    ood_epoch_step = 0
+    ood_samples_seen = 0
     is_simulated_batch = False
     batch_dirichlet_fraction_store: Optional[FloatTensor] = None
 
     _static_idx_order: Optional[LongTensor] = None
+    _ood_idx_order: Optional[LongTensor] = None
 
-    def __init__(self, 
-                args: OrchestratorDatasetArgs, 
-                shuffle: bool = True, 
+    def __init__(self,
+                args: OrchestratorDatasetArgs,
+                shuffle: bool = True,
             ) -> None:
         self.args = args; self.shuffle = shuffle
         if self.args._using_static: self.__shuffle__()
+        if self.args._using_ood: self.__shuffle_ood__()
 
     def __iter__(self) -> Iterator[Tuple[torch.FloatTensor, torch.Tensor]]:
         ''' This iterator does not stop '''
@@ -107,6 +143,20 @@ def make_one_hot(y:torch.Tensor) -> torch.Tensor:
                 self.batch_dirichlet_fraction_store = None
                 yield self.args.static_data[idx], labels # type: ignore
 
+            elif self.args._using_ood and self.__use_ood_predicate__():
+                # Real OOD spectra provided as-is: labels are already multi-hot and are
+                # NEVER mixed or simulated. Unknown entries carry a -1 sentinel to be masked.
+                self.is_simulated_batch = False
+                start =  (self.ood_epoch_step * self.args.batch_size)
+                end   = ((self.ood_epoch_step+1) * self.args.batch_size)
+                self.ood_epoch_step += 1
+                idx = self._ood_idx_order[start: end] # type: ignore
+                self.ood_samples_seen += len(idx)
+                if end >= len(self.args.ood_data)-1: # type: ignore
+                    self.__reset_ood__()
+                self.batch_dirichlet_fraction_store = None
+                yield self.args.ood_data[idx], self.args.ood_labels[idx].to(dtype=torch.float32) # type: ignore
+
             elif self.args._using_sim:
                 self.is_simulated_batch = True
                 # mypy doesn't catch self.args._using_sim
@@ -126,17 +176,35 @@ def __shuffle__(self) -> None:
             self._static_idx_order = LongTensor(torch.randperm(
                 self.args.static_labels.size(0), # type: ignore # caller's responsibility
                 device=self.args.static_labels.device, # type: ignore
-                dtype=torch.int64, 
+                dtype=torch.int64,
             ))
             self.args._shuffle_method_selection_idxs()
 
+    def __shuffle_ood__(self) -> None:
+        if self.shuffle:
+            self._ood_idx_order = LongTensor(torch.randperm(
+                self.args.ood_labels.size(0), # type: ignore # caller's responsibility
+                device=self.args.ood_labels.device, # type: ignore
+                dtype=torch.int64,
+            ))
+        elif self._ood_idx_order is None:
+            self._ood_idx_order = LongTensor(torch.arange(self.args.ood_labels.size(0), dtype=torch.int64)) # type: ignore
+
     def __reset__(self) -> None:
         self.static_epoch += 1
         self.static_epoch_step = 0
         self.__shuffle__()
 
+    def __reset_ood__(self) -> None:
+        self.ood_epoch += 1
+        self.ood_epoch_step = 0
+        self.__shuffle_ood__()
+
     def __use_static_predicate__(self) -> bool:
-        return bool(self.args._method_selection_idxs[(self.step % 100)])
+        return self.args._method_selection_idxs[(self.step % 100)].item() == 1
+
+    def __use_ood_predicate__(self) -> bool:
+        return self.args._method_selection_idxs[(self.step % 100)].item() == 2
 
 
 def dataloader_from_config(
@@ -147,6 +215,8 @@ def dataloader_from_config(
         shuffle: bool = True,
         return_fractions: bool = False,
         misc_dataloader_params: dict = {},
+        ood_spectra: Optional[FloatTensor] = None,
+        ood_labels: Optional[Tensor] = None,
     ) -> DataLoader:
 
     config = read_config(config)
@@ -159,7 +229,10 @@ def dataloader_from_config(
         sim_data_args,
         spectra,
         labels.long(),
-        return_fractions=return_fractions
+        ood_data=ood_spectra,
+        ood_labels=ood_labels,
+        percent_ood=config["dataloader"].get("percent-ood-data", 0.0),
+        return_fractions=return_fractions,
     )
     ods = OrchestratorDataset(ods_args, shuffle)
     return DataLoader(ods, batch_size=None, **misc_dataloader_params)
diff --git a/src/cover_class/train.py b/src/cover_class/train.py
index 6e6d744..73ec668 100644
--- a/src/cover_class/train.py
+++ b/src/cover_class/train.py
@@ -9,7 +9,7 @@
 from datetime import datetime
 
 from cover_class.dataloader import dataloader_from_config, OrchestratorDataset
-from cover_class.utils import read_config, seed as sseed
+from cover_class.utils import read_config, seed as sseed, ood_test_set_from_config
 from cover_class.subsample import subsample_from_config, train_test_split, drop_bad_bands, drop_bad_banddef
 from cover_class.simulation import run_simulation, SimulationArgs, DataArgs, one_hot_encode_simulated_data
 from cover_class.static.retrieval import make_hdf5
@@ -71,9 +71,12 @@ def setup_training_from_config(
         run_name: str = '',
         return_fractions: bool = False,
         misc_dataloader_params: dict = {},
+        inject_ood: bool = False,
     ) -> Tuple[DataLoader, FloatTensor, Tensor]:
     """
     :param: simulated_test_set_n_rows = 0 means don't return a simulated set
+    :param: inject_ood if True, load the 'ood-train-set' from the config and inject those
+            real OOD spectra into the training dataloader as-is (unmixed) to induce overfitting.
 
     Returns: A tuple of the training dataloader, the test data matrix, and test labels
     """
@@ -113,6 +116,18 @@ def setup_training_from_config(
     train_spectra = train_spectra.to(torch.float32)
     test_spectra = test_spectra.to(torch.float32)
 
+    ood_spectra, ood_labels = None, None
+    if inject_ood:
+        # Unknown (label 2) entries become -1 so they can be masked out of the loss during backprop.
+        ood_spectra, ood_labels = ood_test_set_from_config(
+            config,
+            include_unknown=False,
+            err_on_missed_class=False,
+            key='ood-train-set',
+            unknown_fill=-1.0,
+        )
+        ood_spectra = FloatTensor(ood_spectra.to(torch.float32))
+
     odl = dataloader_from_config(
         config,
         FloatTensor(train_spectra),
@@ -121,6 +136,8 @@ def setup_training_from_config(
         shuffle,
         return_fractions,
         misc_dataloader_params,
+        ood_spectra=ood_spectra,
+        ood_labels=ood_labels,
     )
 
     return odl, FloatTensor(test_spectra), test_labels
diff --git a/src/cover_class/utils.py b/src/cover_class/utils.py
index 55b19d9..e8585e1 100644
--- a/src/cover_class/utils.py
+++ b/src/cover_class/utils.py
@@ -35,15 +35,25 @@ def load_rfl(hdr_fp:str) -> Tuple[np.ndarray, np.ndarray]:
     banddef = np.array(banddef, dtype=float) # type: ignore 
     return rfl, banddef # type: ignore 
 
-def ood_test_set_from_config(c: str|Dict, include_unknown: bool = False, err_on_missed_class: bool = True) -> Tuple[torch.Tensor, torch.Tensor]:
+def ood_test_set_from_config(
+        c: str|Dict,
+        include_unknown: bool = False,
+        err_on_missed_class: bool = True,
+        key: str = 'ood-test-set',
+        unknown_fill: float = float('nan'),
+    ) -> Tuple[torch.Tensor, torch.Tensor]:
     """
-    Load the OOD test set from a configuration file.
+    Load an OOD set (test or train) from a configuration file.
 
     Args:
-        c: Config file path or a dictionary containing 'datasets' and 'ood-test-set'.
-        include_unknown: If True, samples with label 2 (unknown) are included and treated as present.
-                        If False, any sample with a label 2 is discarded from the dataset.
+        c: Config file path or a dictionary containing 'datasets' and the OOD set path.
+        include_unknown: If True, samples with label 2 (unknown) are treated as present (1).
+                        If False, label 2 is replaced with `unknown_fill` so it can be masked out.
         err_on_missed_class: If True, raises a RuntimeError if a class specified in the config is missing from the OOD set.
+        key: The config key holding the OOD set path (e.g. 'ood-test-set' or 'ood-train-set').
+        unknown_fill: The value used to fill unknown (label 2) entries when include_unknown is False.
+                      NaN for the test set (already masked in eval); -1 for the train set so those
+                      entries can be masked out of the loss during backprop.
 
     Returns:
         A tuple containing:
@@ -53,7 +63,7 @@ def ood_test_set_from_config(c: str|Dict, include_unknown: bool = False, err_on_
     config = read_config(c)
     class_order: List[str] = [d for d in config['datasets'].keys() if config['datasets'][d] is not None]
 
-    with h5py.File(config['ood-test-set'], 'r') as f:
+    with h5py.File(config[key], 'r') as f:
         labels = np.asarray(f['labels'][:])
         spectra = np.asarray(f['spectra'][:])
         classes = np.asarray(f.attrs['classes'][:]).astype(str) # type: ignore
@@ -68,7 +78,7 @@ def ood_test_set_from_config(c: str|Dict, include_unknown: bool = False, err_on_
             if include_unknown:
                 Y_np[:, i] = np.where(class_labels == 2, 1, class_labels)
             else:
-                Y_np[:, i] = np.where(class_labels == 2, np.nan, class_labels)
+                Y_np[:, i] = np.where(class_labels == 2, unknown_fill, class_labels)
             if err_on_missed_class and not np.any(Y_np[:, i] == 1):
-                raise RuntimeError(f"No data is found in the OOD Test set for class: '{name}'")
+                raise RuntimeError(f"No data is found in the OOD set for class: '{name}'")
         return X, torch.from_numpy(Y_np).to(torch.float32)

From 50fbe922000735b2913286034459c6b13f2ee007 Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Tue, 18 Aug 2026 18:36:36 -0700
Subject: [PATCH 6/8] reporting script for francisco's dataset

---
 models/spectf/report_spectf_francisco.py | 232 +++++++++++++++++++++++
 1 file changed, 232 insertions(+)
 create mode 100644 models/spectf/report_spectf_francisco.py

diff --git a/models/spectf/report_spectf_francisco.py b/models/spectf/report_spectf_francisco.py
new file mode 100644
index 0000000..27944f8
--- /dev/null
+++ b/models/spectf/report_spectf_francisco.py
@@ -0,0 +1,232 @@
+import os
+
+# Set CUBLAS_WORKSPACE_CONFIG for deterministic CUDA behavior
+# pylint: disable=wrong-import-position
+os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
+
+import csv
+from datetime import datetime
+import rich_click as click
+import yaml
+
+import numpy as np
+import torch
+import matplotlib.pyplot as plt
+
+from cover_class.utils import read_config
+from cover_class.subsample.forward_pipeline import drop_bad_bands, drop_bad_banddef
+
+from spectf.model import SpecTfEncoder
+from spectf.utils import get_device
+
+ENV_VAR_PREFIX = 'COVER_CLASS_FRANCISCO_'
+
+# The francisco fraction CSV only reports these material fractions (soil/pv/npv).
+# True water and snow+ice are always 0 for this dataset; shade is ignored.
+FRACTION_CSV_COLUMNS = {'soil': 'soil', 'pv': 'pv', 'npv': 'npv'}
+
+
+def load_francisco_data(rfl_csv: str, frac_csv: str, class_names: list, drop_wl_ranges):
+    """Load the francisco spectra and align them to their per-plot true fractions.
+
+    Each plot in the fraction CSV maps to multiple spectra in the reflectance CSV
+    (matched by the plot identifier), so those spectra share identical labels.
+
+    Returns:
+        spectra: (N, B) reflectance with bad bands dropped
+        banddef: (B,) wavelengths with bad bands dropped
+        true_fractions: (N, C) true fraction per class (water/snow+ice forced to 0)
+        plot_ids: (N,) plot identifier per spectrum
+    """
+    # --- Reflectance CSV: header row is [plot_num, wl_0, wl_1, ...] ---
+    with open(rfl_csv, 'r', encoding='utf-8') as f:
+        reader = csv.reader(f)
+        header = next(reader)
+        wavelengths = np.array([float(w) for w in header[1:]], dtype=np.float64)
+
+        plot_ids = []
+        spectra = []
+        for row in reader:
+            if not row:
+                continue
+            plot_ids.append(row[0])
+            spectra.append([float(v) for v in row[1:]])
+
+    spectra = np.array(spectra, dtype=np.float32)
+    plot_ids = np.array(plot_ids)
+
+    # Drop bad bands to match the band definition the model was trained on
+    spectra = drop_bad_bands(spectra, wavelengths, drop_wl_ranges)
+    banddef = drop_bad_banddef(wavelengths, drop_wl_ranges)
+
+    # --- Fraction CSV: one row per plot ---
+    plot_to_frac = {}
+    with open(frac_csv, 'r', encoding='utf-8') as f:
+        reader = csv.DictReader(f)
+        for row in reader:
+            frac = np.zeros(len(class_names), dtype=np.float32)
+            for ci, cname in enumerate(class_names):
+                col = FRACTION_CSV_COLUMNS.get(cname)
+                if col is not None and col in row:
+                    frac[ci] = float(row[col])
+                # else: leave at 0 (water, snow+ice)
+            plot_to_frac[row['plot']] = frac
+
+    # Align each spectrum to its plot's fractions
+    missing = sorted({p for p in plot_ids if p not in plot_to_frac})
+    if missing:
+        raise ValueError(f"{len(missing)} plot(s) in reflectance CSV have no fraction row: {missing[:5]}...")
+
+    true_fractions = np.stack([plot_to_frac[p] for p in plot_ids], axis=0)
+
+    return spectra, banddef, true_fractions, plot_ids
+
+
+@click.command()
+@click.option(
+    "--outdir",
+    required=True,
+    type=click.Path(exists=True, dir_okay=True, file_okay=False),
+    help="Output file directory.",
+    envvar=f'{ENV_VAR_PREFIX}OUTDIR'
+)
+@click.option(
+    "--data-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML data config (used for class names and drop-bands).",
+    envvar=f'{ENV_VAR_PREFIX}DATA_CONFIG'
+)
+@click.option(
+    "--model-config",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the YAML config for the model architecture.",
+    envvar=f'{ENV_VAR_PREFIX}MODEL_CONFIG'
+)
+@click.option(
+    "--model-weights",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to the model weights file (.pth).",
+    envvar=f'{ENV_VAR_PREFIX}MODEL_WEIGHTS'
+)
+@click.option(
+    "--rfl-csv",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to emit_pixels_rfl.csv (spectra).",
+    envvar=f'{ENV_VAR_PREFIX}RFL_CSV'
+)
+@click.option(
+    "--frac-csv",
+    required=True,
+    type=click.Path(exists=True, dir_okay=False),
+    help="Path to fraction_output.csv (per-plot fractions).",
+    envvar=f'{ENV_VAR_PREFIX}FRAC_CSV'
+)
+def run_francisco_evaluation(
+        outdir: str,
+        data_config: str,
+        model_config: str,
+        model_weights: str,
+        rfl_csv: str,
+        frac_csv: str,
+    ):
+
+    # Load model config
+    with open(model_config, 'r', encoding='utf-8') as f:
+        m_config = yaml.safe_load(f)
+
+    # Load data config (for class names and drop-bands)
+    d_config = read_config(data_config)
+    ds = d_config['datasets']
+    class_names = [str(c) for c in ds.keys() if ds[c] is not None and len(ds[c])]
+    print(f"Classes: {class_names}")
+    drop_wl_ranges = d_config['drop-bands-wavelengths']
+
+    # Load francisco spectra and align to per-plot fractions
+    print(f"Loading spectra from {rfl_csv}")
+    print(f"Loading fractions from {frac_csv}")
+    spectra, banddef, true_fractions, plot_ids = load_francisco_data(
+        rfl_csv, frac_csv, class_names, drop_wl_ranges
+    )
+    print(f"Loaded {spectra.shape[0]} spectra ({spectra.shape[1]} bands) "
+          f"across {len(np.unique(plot_ids))} plots")
+
+    # Device (cuda / mps / cpu)
+    device = get_device(0)
+
+    # Model definition
+    banddef_t = torch.from_numpy(np.asarray(banddef)).to(dtype=torch.float32, device=device)
+    model = SpecTfEncoder(banddef_t,
+                         dim_output=m_config['model']['dim_output'],
+                         num_heads=m_config['model']['num_heads'],
+                         dim_proj=m_config['model']['dim_proj'],
+                         dim_ff=m_config['model']['dim_ff'],
+                         dropout=m_config['model']['dropout'],
+                         agg=m_config['model']['agg'],
+                         use_residual=m_config['model']['use_residual'],
+                         num_layers=m_config['model']['num_layers']).to(device)
+
+    # Load model weights
+    print(f"Loading model weights from {model_weights}")
+    model.load_state_dict(torch.load(model_weights, map_location=device))
+    model.eval()
+
+    # Inference — this is a BCE (sigmoid) classifier, so apply sigmoid to logits
+    bs = m_config['batch_size']
+    n = spectra.shape[0]
+    y_pred = np.zeros((n, m_config['model']['dim_output']), dtype=float)
+
+    print("Running inference...")
+    spectra_t = torch.from_numpy(spectra).to(dtype=torch.float32)
+    with torch.no_grad():
+        for i in range(0, n, bs):
+            batch = spectra_t[i:i+bs].to(device=device)
+            batch = torch.unsqueeze(batch, -1)
+            logits = model(batch)
+            probs = torch.sigmoid(logits).detach().cpu().numpy().astype(float)
+            y_pred[i:i+probs.shape[0]] = probs
+
+    y_true = true_fractions
+
+    # Scatter plot: one panel per class, true fraction (x) vs predicted posterior (y)
+    n_classes = len(class_names)
+    fig, axes = plt.subplots(1, n_classes, figsize=(4 * n_classes, 4))
+    if n_classes == 1:
+        axes = [axes]
+
+    for i, (ax, class_name) in enumerate(zip(axes, class_names)):
+        ax.scatter(y_true[:, i], y_pred[:, i], alpha=0.3, s=8, c='blue', edgecolors='none')
+
+        # 1:1 line
+        ax.plot([0, 1], [0, 1], 'r--', linewidth=2, label='1:1 line')
+
+        ax.set_xlabel('True Fraction', fontsize=12)
+        ax.set_ylabel('Predicted Posterior', fontsize=12)
+        ax.set_title(f'{class_name}', fontsize=14, fontweight='bold')
+
+        ax.set_xlim([-0.05, 1.05])
+        ax.set_ylim([-0.05, 1.05])
+        ax.set_aspect('equal')
+
+        ax.grid(True, alpha=0.3)
+        ax.legend(loc='lower right', fontsize=9)
+
+    plt.tight_layout()
+
+    # Save figure
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    output_path = os.path.join(outdir, f"{timestamp}_francisco_scatter_plots.png")
+    plt.savefig(output_path, dpi=300, bbox_inches='tight')
+    print(f"\n{'='*80}")
+    print(f"Scatter plot saved to: {output_path}")
+    print(f"{'='*80}\n")
+
+    plt.show()
+
+
+if __name__ == "__main__":
+    # pylint: disable=no-value-for-parameter
+    run_francisco_evaluation()

From fa10d371ed7ada93670277cfb946e51fb812a83d Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Tue, 18 Aug 2026 21:51:49 -0700
Subject: [PATCH 7/8] Added plotting functions

---
 models/spectf/report_spectf_francisco.py | 167 ++++++++++++++++++++++-
 1 file changed, 164 insertions(+), 3 deletions(-)

diff --git a/models/spectf/report_spectf_francisco.py b/models/spectf/report_spectf_francisco.py
index 27944f8..2684500 100644
--- a/models/spectf/report_spectf_francisco.py
+++ b/models/spectf/report_spectf_francisco.py
@@ -12,6 +12,8 @@
 import numpy as np
 import torch
 import matplotlib.pyplot as plt
+from matplotlib.collections import LineCollection
+from matplotlib.lines import Line2D
 
 from cover_class.utils import read_config
 from cover_class.subsample.forward_pipeline import drop_bad_bands, drop_bad_banddef
@@ -82,6 +84,122 @@ def load_francisco_data(rfl_csv: str, frac_csv: str, class_names: list, drop_wl_
     return spectra, banddef, true_fractions, plot_ids
 
 
+def _plot_spectra_group(ax, wavelengths, group_spectra, color, label, alpha=0.5):
+    """Plot a group of spectra on *ax* as semi-transparent *color* lines.
+
+    Uses a LineCollection for efficiency, falling back to individual lines
+    when the group is empty.
+    """
+    if len(group_spectra) == 0:
+        return
+
+    n_spectra = len(group_spectra)
+    # Broadcast wavelengths to match spectra shape: (n_spectra, n_bands)
+    wavelengths_bc = np.broadcast_to(wavelengths, group_spectra.shape)
+    # Build line segments: shape (n_spectra, n_bands, 2) with (wavelength, value)
+    segs = np.stack([wavelengths_bc, group_spectra], axis=-1)  # (n_spectra, n_bands, 2)
+
+    lc = LineCollection(segs, colors=color, alpha=alpha, linewidths=0.5)
+    ax.add_collection(lc)
+    ax.set_xlim(450, 2500)
+    ax.set_ylim(0, 1)
+
+
+def _shade_dropped_bands(ax, drop_wl_ranges, xlim=(450, 2500)):
+    """Shade *drop_wl_ranges* on *ax* as a light gray background, clipped to *xlim*."""
+    if not drop_wl_ranges:
+        return
+    for low, high in drop_wl_ranges:
+        if high < xlim[0] or low > xlim[1]:
+            continue
+        ax.axvspan(max(low, xlim[0]), min(high, xlim[1]), color='gray', alpha=0.15, zorder=0)
+
+
+def plot_spectra_by_class(spectra, banddef, y_pred, class_names, thresholds, outdir,
+                          figure_prefix='spectra', timestamp=None, drop_wl_ranges=None):
+    """Generate one figure per class, plotting spectra colored by prediction vs threshold.
+
+    For classes where *all* spectra contain some fraction (soil, pv, npv):
+      - green if y_pred >= threshold  (true positive)
+      - red   if y_pred <  threshold  (false negative)
+
+    For other classes (snow+ice, water):
+      - red only if y_pred >= threshold (false positive)
+    """
+    # Classes that are always present in this dataset (green/red logic)
+    always_present = {'soil', 'pv', 'npv'}
+    if timestamp is None:
+        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    # banddef is already in nm (from the CSV header wavelengths)
+    wavelength_nm = banddef
+
+    for ci, class_name in enumerate(class_names):
+        threshold = thresholds[ci]
+        y_true_col = y_pred[:, ci]
+
+        print(class_name, threshold)
+
+        if class_name in always_present:
+            # Green: prediction >= threshold (true positive)
+            # Red:   prediction <  threshold (false negative)
+            green_mask = y_true_col >= threshold
+            red_mask = ~green_mask
+        else:
+            # Only plot red (false positives: prediction >= threshold)
+            green_mask = np.zeros(len(y_true_col), dtype=bool)
+            if threshold == 0.0:
+                red_mask = y_true_col > threshold  # strictly greater than 0
+            else:
+                red_mask = y_true_col >= threshold
+        print(y_true_col[:5])
+        print(red_mask.sum())
+
+        fig, ax = plt.subplots(figsize=(12, 6))
+        _shade_dropped_bands(ax, drop_wl_ranges)
+
+        if green_mask.any():
+            _plot_spectra_group(
+                ax, wavelength_nm, spectra[green_mask],
+                color='green', label=f'Pred >= {threshold:.3f} (true positive)'
+            )
+        if red_mask.any():
+            _plot_spectra_group(
+                ax, wavelength_nm, spectra[red_mask],
+                color='red', label=f'Pred {"≤" if class_name in always_present else "≥"} {threshold:.3f}'
+            )
+
+        ax.set_xlabel('Wavelength (nm)', fontsize=12)
+        ax.set_ylabel('Reflectance', fontsize=12)
+        ax.set_title(f'{class_name} Spectra', fontsize=14, fontweight='bold')
+        ax.set_xlim(450, 2500)
+        ax.set_ylim(0, 1)
+        ax.grid(True, alpha=0.3)
+
+        ax.text(
+            0.02, 0.98, f'green: {green_mask.sum()}\nred: {red_mask.sum()}',
+            transform=ax.transAxes, fontsize=10, verticalalignment='top',
+            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8, edgecolor='gray'),
+        )
+
+        if class_name in always_present:
+            legend_elements = [
+                Line2D([0], [0], color='green', lw=2, label=f'Pred >= {threshold:.3f} (true positive)'),
+                Line2D([0], [0], color='red', lw=2, label=f'Pred < {threshold:.3f} (false negative)'),
+            ]
+        else:
+            legend_elements = [
+                Line2D([0], [0], color='red', lw=2, label=f'Pred >= {threshold:.3f} (false positive)'),
+            ]
+
+        ax.legend(handles=legend_elements, loc='upper right', fontsize=9)
+        plt.tight_layout()
+
+        output_path = os.path.join(outdir, f"{timestamp}_{figure_prefix}_{class_name}.png")
+        plt.savefig(output_path, dpi=300, bbox_inches='tight')
+        print(f"Spectra plot saved to: {output_path}")
+        plt.close(fig)
+
+
 @click.command()
 @click.option(
     "--outdir",
@@ -125,6 +243,15 @@ def load_francisco_data(rfl_csv: str, frac_csv: str, class_names: list, drop_wl_
     help="Path to fraction_output.csv (per-plot fractions).",
     envvar=f'{ENV_VAR_PREFIX}FRAC_CSV'
 )
+@click.option(
+    "--thresholds",
+    required=False,
+    type=float,
+    nargs=5,
+    default=[1.0, 0.959, 0.408, 1.000, 0.000],
+    help="List of 5 thresholds, one per class, in order: soil, pv, npv, snow+ice, water.",
+    envvar=f'{ENV_VAR_PREFIX}THRESHOLDS'
+)
 def run_francisco_evaluation(
         outdir: str,
         data_config: str,
@@ -132,7 +259,8 @@ def run_francisco_evaluation(
         model_weights: str,
         rfl_csv: str,
         frac_csv: str,
-    ):
+        thresholds: list,
+):
 
     # Load model config
     with open(model_config, 'r', encoding='utf-8') as f:
@@ -187,10 +315,25 @@ def run_francisco_evaluation(
             batch = torch.unsqueeze(batch, -1)
             logits = model(batch)
             probs = torch.sigmoid(logits).detach().cpu().numpy().astype(float)
+            probs[probs < 1e-6] = 0.0
+            probs[probs > 1-(1e-6)] = 1.0
             y_pred[i:i+probs.shape[0]] = probs
 
     y_true = true_fractions
 
+    # Export y_pred as CSV
+    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    y_pred_csv_path = os.path.join(outdir, f"{timestamp}_y_pred.csv")
+    header = ['plot_num'] + class_names
+    with open(y_pred_csv_path, 'w', encoding='utf-8', newline='') as f:
+        writer = csv.writer(f)
+        writer.writerow(header)
+        for idx, plot_id in enumerate(plot_ids):
+            writer.writerow([plot_id] + [f'{v:.8f}' for v in y_pred[idx]])
+    print(f"\n{'='*80}")
+    print(f"y_pred exported to: {y_pred_csv_path}")
+    print(f"{'='*80}\n")
+
     # Scatter plot: one panel per class, true fraction (x) vs predicted posterior (y)
     n_classes = len(class_names)
     fig, axes = plt.subplots(1, n_classes, figsize=(4 * n_classes, 4))
@@ -216,8 +359,7 @@ def run_francisco_evaluation(
 
     plt.tight_layout()
 
-    # Save figure
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
+    # Save figure (reusing the timestamp from the CSV export for consistency)
     output_path = os.path.join(outdir, f"{timestamp}_francisco_scatter_plots.png")
     plt.savefig(output_path, dpi=300, bbox_inches='tight')
     print(f"\n{'='*80}")
@@ -226,6 +368,25 @@ def run_francisco_evaluation(
 
     plt.show()
 
+    # Plot spectra by class with thresholds
+    thresholds_list = list(thresholds)
+    print(f"\nClass thresholds: {dict(zip(class_names, thresholds_list))}")
+    assert len(thresholds_list) == len(class_names), \
+        f"Expected {len(class_names)} thresholds, got {len(thresholds_list)}"
+
+    print("Generating spectra plots by class...")
+    plot_spectra_by_class(
+        spectra=spectra,
+        banddef=banddef,
+        y_pred=y_pred,
+        class_names=class_names,
+        thresholds=thresholds_list,
+        outdir=outdir,
+        timestamp=timestamp,
+        drop_wl_ranges=drop_wl_ranges,
+    )
+    print("Spectra plots complete.")
+
 
 if __name__ == "__main__":
     # pylint: disable=no-value-for-parameter

From 99e2e3a077a84db91d3e4e1cbab6544617962616 Mon Sep 17 00:00:00 2001
From: Jake Lee <jake.h.lee@jpl.nasa.gov>
Date: Tue, 18 Aug 2026 22:11:03 -0700
Subject: [PATCH 8/8] Cleanup and consolidation

---
 models/spectf/report_spectf_francisco.py | 393 -----------------------
 models/spectf/training_spectf.py         |  17 +-
 models/spectf/training_spectf_ood.py     | 333 -------------------
 3 files changed, 12 insertions(+), 731 deletions(-)
 delete mode 100644 models/spectf/report_spectf_francisco.py
 delete mode 100644 models/spectf/training_spectf_ood.py

diff --git a/models/spectf/report_spectf_francisco.py b/models/spectf/report_spectf_francisco.py
deleted file mode 100644
index 2684500..0000000
--- a/models/spectf/report_spectf_francisco.py
+++ /dev/null
@@ -1,393 +0,0 @@
-import os
-
-# Set CUBLAS_WORKSPACE_CONFIG for deterministic CUDA behavior
-# pylint: disable=wrong-import-position
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-
-import csv
-from datetime import datetime
-import rich_click as click
-import yaml
-
-import numpy as np
-import torch
-import matplotlib.pyplot as plt
-from matplotlib.collections import LineCollection
-from matplotlib.lines import Line2D
-
-from cover_class.utils import read_config
-from cover_class.subsample.forward_pipeline import drop_bad_bands, drop_bad_banddef
-
-from spectf.model import SpecTfEncoder
-from spectf.utils import get_device
-
-ENV_VAR_PREFIX = 'COVER_CLASS_FRANCISCO_'
-
-# The francisco fraction CSV only reports these material fractions (soil/pv/npv).
-# True water and snow+ice are always 0 for this dataset; shade is ignored.
-FRACTION_CSV_COLUMNS = {'soil': 'soil', 'pv': 'pv', 'npv': 'npv'}
-
-
-def load_francisco_data(rfl_csv: str, frac_csv: str, class_names: list, drop_wl_ranges):
-    """Load the francisco spectra and align them to their per-plot true fractions.
-
-    Each plot in the fraction CSV maps to multiple spectra in the reflectance CSV
-    (matched by the plot identifier), so those spectra share identical labels.
-
-    Returns:
-        spectra: (N, B) reflectance with bad bands dropped
-        banddef: (B,) wavelengths with bad bands dropped
-        true_fractions: (N, C) true fraction per class (water/snow+ice forced to 0)
-        plot_ids: (N,) plot identifier per spectrum
-    """
-    # --- Reflectance CSV: header row is [plot_num, wl_0, wl_1, ...] ---
-    with open(rfl_csv, 'r', encoding='utf-8') as f:
-        reader = csv.reader(f)
-        header = next(reader)
-        wavelengths = np.array([float(w) for w in header[1:]], dtype=np.float64)
-
-        plot_ids = []
-        spectra = []
-        for row in reader:
-            if not row:
-                continue
-            plot_ids.append(row[0])
-            spectra.append([float(v) for v in row[1:]])
-
-    spectra = np.array(spectra, dtype=np.float32)
-    plot_ids = np.array(plot_ids)
-
-    # Drop bad bands to match the band definition the model was trained on
-    spectra = drop_bad_bands(spectra, wavelengths, drop_wl_ranges)
-    banddef = drop_bad_banddef(wavelengths, drop_wl_ranges)
-
-    # --- Fraction CSV: one row per plot ---
-    plot_to_frac = {}
-    with open(frac_csv, 'r', encoding='utf-8') as f:
-        reader = csv.DictReader(f)
-        for row in reader:
-            frac = np.zeros(len(class_names), dtype=np.float32)
-            for ci, cname in enumerate(class_names):
-                col = FRACTION_CSV_COLUMNS.get(cname)
-                if col is not None and col in row:
-                    frac[ci] = float(row[col])
-                # else: leave at 0 (water, snow+ice)
-            plot_to_frac[row['plot']] = frac
-
-    # Align each spectrum to its plot's fractions
-    missing = sorted({p for p in plot_ids if p not in plot_to_frac})
-    if missing:
-        raise ValueError(f"{len(missing)} plot(s) in reflectance CSV have no fraction row: {missing[:5]}...")
-
-    true_fractions = np.stack([plot_to_frac[p] for p in plot_ids], axis=0)
-
-    return spectra, banddef, true_fractions, plot_ids
-
-
-def _plot_spectra_group(ax, wavelengths, group_spectra, color, label, alpha=0.5):
-    """Plot a group of spectra on *ax* as semi-transparent *color* lines.
-
-    Uses a LineCollection for efficiency, falling back to individual lines
-    when the group is empty.
-    """
-    if len(group_spectra) == 0:
-        return
-
-    n_spectra = len(group_spectra)
-    # Broadcast wavelengths to match spectra shape: (n_spectra, n_bands)
-    wavelengths_bc = np.broadcast_to(wavelengths, group_spectra.shape)
-    # Build line segments: shape (n_spectra, n_bands, 2) with (wavelength, value)
-    segs = np.stack([wavelengths_bc, group_spectra], axis=-1)  # (n_spectra, n_bands, 2)
-
-    lc = LineCollection(segs, colors=color, alpha=alpha, linewidths=0.5)
-    ax.add_collection(lc)
-    ax.set_xlim(450, 2500)
-    ax.set_ylim(0, 1)
-
-
-def _shade_dropped_bands(ax, drop_wl_ranges, xlim=(450, 2500)):
-    """Shade *drop_wl_ranges* on *ax* as a light gray background, clipped to *xlim*."""
-    if not drop_wl_ranges:
-        return
-    for low, high in drop_wl_ranges:
-        if high < xlim[0] or low > xlim[1]:
-            continue
-        ax.axvspan(max(low, xlim[0]), min(high, xlim[1]), color='gray', alpha=0.15, zorder=0)
-
-
-def plot_spectra_by_class(spectra, banddef, y_pred, class_names, thresholds, outdir,
-                          figure_prefix='spectra', timestamp=None, drop_wl_ranges=None):
-    """Generate one figure per class, plotting spectra colored by prediction vs threshold.
-
-    For classes where *all* spectra contain some fraction (soil, pv, npv):
-      - green if y_pred >= threshold  (true positive)
-      - red   if y_pred <  threshold  (false negative)
-
-    For other classes (snow+ice, water):
-      - red only if y_pred >= threshold (false positive)
-    """
-    # Classes that are always present in this dataset (green/red logic)
-    always_present = {'soil', 'pv', 'npv'}
-    if timestamp is None:
-        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    # banddef is already in nm (from the CSV header wavelengths)
-    wavelength_nm = banddef
-
-    for ci, class_name in enumerate(class_names):
-        threshold = thresholds[ci]
-        y_true_col = y_pred[:, ci]
-
-        print(class_name, threshold)
-
-        if class_name in always_present:
-            # Green: prediction >= threshold (true positive)
-            # Red:   prediction <  threshold (false negative)
-            green_mask = y_true_col >= threshold
-            red_mask = ~green_mask
-        else:
-            # Only plot red (false positives: prediction >= threshold)
-            green_mask = np.zeros(len(y_true_col), dtype=bool)
-            if threshold == 0.0:
-                red_mask = y_true_col > threshold  # strictly greater than 0
-            else:
-                red_mask = y_true_col >= threshold
-        print(y_true_col[:5])
-        print(red_mask.sum())
-
-        fig, ax = plt.subplots(figsize=(12, 6))
-        _shade_dropped_bands(ax, drop_wl_ranges)
-
-        if green_mask.any():
-            _plot_spectra_group(
-                ax, wavelength_nm, spectra[green_mask],
-                color='green', label=f'Pred >= {threshold:.3f} (true positive)'
-            )
-        if red_mask.any():
-            _plot_spectra_group(
-                ax, wavelength_nm, spectra[red_mask],
-                color='red', label=f'Pred {"≤" if class_name in always_present else "≥"} {threshold:.3f}'
-            )
-
-        ax.set_xlabel('Wavelength (nm)', fontsize=12)
-        ax.set_ylabel('Reflectance', fontsize=12)
-        ax.set_title(f'{class_name} Spectra', fontsize=14, fontweight='bold')
-        ax.set_xlim(450, 2500)
-        ax.set_ylim(0, 1)
-        ax.grid(True, alpha=0.3)
-
-        ax.text(
-            0.02, 0.98, f'green: {green_mask.sum()}\nred: {red_mask.sum()}',
-            transform=ax.transAxes, fontsize=10, verticalalignment='top',
-            bbox=dict(boxstyle='round', facecolor='white', alpha=0.8, edgecolor='gray'),
-        )
-
-        if class_name in always_present:
-            legend_elements = [
-                Line2D([0], [0], color='green', lw=2, label=f'Pred >= {threshold:.3f} (true positive)'),
-                Line2D([0], [0], color='red', lw=2, label=f'Pred < {threshold:.3f} (false negative)'),
-            ]
-        else:
-            legend_elements = [
-                Line2D([0], [0], color='red', lw=2, label=f'Pred >= {threshold:.3f} (false positive)'),
-            ]
-
-        ax.legend(handles=legend_elements, loc='upper right', fontsize=9)
-        plt.tight_layout()
-
-        output_path = os.path.join(outdir, f"{timestamp}_{figure_prefix}_{class_name}.png")
-        plt.savefig(output_path, dpi=300, bbox_inches='tight')
-        print(f"Spectra plot saved to: {output_path}")
-        plt.close(fig)
-
-
-@click.command()
-@click.option(
-    "--outdir",
-    required=True,
-    type=click.Path(exists=True, dir_okay=True, file_okay=False),
-    help="Output file directory.",
-    envvar=f'{ENV_VAR_PREFIX}OUTDIR'
-)
-@click.option(
-    "--data-config",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to the YAML data config (used for class names and drop-bands).",
-    envvar=f'{ENV_VAR_PREFIX}DATA_CONFIG'
-)
-@click.option(
-    "--model-config",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to the YAML config for the model architecture.",
-    envvar=f'{ENV_VAR_PREFIX}MODEL_CONFIG'
-)
-@click.option(
-    "--model-weights",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to the model weights file (.pth).",
-    envvar=f'{ENV_VAR_PREFIX}MODEL_WEIGHTS'
-)
-@click.option(
-    "--rfl-csv",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to emit_pixels_rfl.csv (spectra).",
-    envvar=f'{ENV_VAR_PREFIX}RFL_CSV'
-)
-@click.option(
-    "--frac-csv",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to fraction_output.csv (per-plot fractions).",
-    envvar=f'{ENV_VAR_PREFIX}FRAC_CSV'
-)
-@click.option(
-    "--thresholds",
-    required=False,
-    type=float,
-    nargs=5,
-    default=[1.0, 0.959, 0.408, 1.000, 0.000],
-    help="List of 5 thresholds, one per class, in order: soil, pv, npv, snow+ice, water.",
-    envvar=f'{ENV_VAR_PREFIX}THRESHOLDS'
-)
-def run_francisco_evaluation(
-        outdir: str,
-        data_config: str,
-        model_config: str,
-        model_weights: str,
-        rfl_csv: str,
-        frac_csv: str,
-        thresholds: list,
-):
-
-    # Load model config
-    with open(model_config, 'r', encoding='utf-8') as f:
-        m_config = yaml.safe_load(f)
-
-    # Load data config (for class names and drop-bands)
-    d_config = read_config(data_config)
-    ds = d_config['datasets']
-    class_names = [str(c) for c in ds.keys() if ds[c] is not None and len(ds[c])]
-    print(f"Classes: {class_names}")
-    drop_wl_ranges = d_config['drop-bands-wavelengths']
-
-    # Load francisco spectra and align to per-plot fractions
-    print(f"Loading spectra from {rfl_csv}")
-    print(f"Loading fractions from {frac_csv}")
-    spectra, banddef, true_fractions, plot_ids = load_francisco_data(
-        rfl_csv, frac_csv, class_names, drop_wl_ranges
-    )
-    print(f"Loaded {spectra.shape[0]} spectra ({spectra.shape[1]} bands) "
-          f"across {len(np.unique(plot_ids))} plots")
-
-    # Device (cuda / mps / cpu)
-    device = get_device(0)
-
-    # Model definition
-    banddef_t = torch.from_numpy(np.asarray(banddef)).to(dtype=torch.float32, device=device)
-    model = SpecTfEncoder(banddef_t,
-                         dim_output=m_config['model']['dim_output'],
-                         num_heads=m_config['model']['num_heads'],
-                         dim_proj=m_config['model']['dim_proj'],
-                         dim_ff=m_config['model']['dim_ff'],
-                         dropout=m_config['model']['dropout'],
-                         agg=m_config['model']['agg'],
-                         use_residual=m_config['model']['use_residual'],
-                         num_layers=m_config['model']['num_layers']).to(device)
-
-    # Load model weights
-    print(f"Loading model weights from {model_weights}")
-    model.load_state_dict(torch.load(model_weights, map_location=device))
-    model.eval()
-
-    # Inference — this is a BCE (sigmoid) classifier, so apply sigmoid to logits
-    bs = m_config['batch_size']
-    n = spectra.shape[0]
-    y_pred = np.zeros((n, m_config['model']['dim_output']), dtype=float)
-
-    print("Running inference...")
-    spectra_t = torch.from_numpy(spectra).to(dtype=torch.float32)
-    with torch.no_grad():
-        for i in range(0, n, bs):
-            batch = spectra_t[i:i+bs].to(device=device)
-            batch = torch.unsqueeze(batch, -1)
-            logits = model(batch)
-            probs = torch.sigmoid(logits).detach().cpu().numpy().astype(float)
-            probs[probs < 1e-6] = 0.0
-            probs[probs > 1-(1e-6)] = 1.0
-            y_pred[i:i+probs.shape[0]] = probs
-
-    y_true = true_fractions
-
-    # Export y_pred as CSV
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    y_pred_csv_path = os.path.join(outdir, f"{timestamp}_y_pred.csv")
-    header = ['plot_num'] + class_names
-    with open(y_pred_csv_path, 'w', encoding='utf-8', newline='') as f:
-        writer = csv.writer(f)
-        writer.writerow(header)
-        for idx, plot_id in enumerate(plot_ids):
-            writer.writerow([plot_id] + [f'{v:.8f}' for v in y_pred[idx]])
-    print(f"\n{'='*80}")
-    print(f"y_pred exported to: {y_pred_csv_path}")
-    print(f"{'='*80}\n")
-
-    # Scatter plot: one panel per class, true fraction (x) vs predicted posterior (y)
-    n_classes = len(class_names)
-    fig, axes = plt.subplots(1, n_classes, figsize=(4 * n_classes, 4))
-    if n_classes == 1:
-        axes = [axes]
-
-    for i, (ax, class_name) in enumerate(zip(axes, class_names)):
-        ax.scatter(y_true[:, i], y_pred[:, i], alpha=0.3, s=8, c='blue', edgecolors='none')
-
-        # 1:1 line
-        ax.plot([0, 1], [0, 1], 'r--', linewidth=2, label='1:1 line')
-
-        ax.set_xlabel('True Fraction', fontsize=12)
-        ax.set_ylabel('Predicted Posterior', fontsize=12)
-        ax.set_title(f'{class_name}', fontsize=14, fontweight='bold')
-
-        ax.set_xlim([-0.05, 1.05])
-        ax.set_ylim([-0.05, 1.05])
-        ax.set_aspect('equal')
-
-        ax.grid(True, alpha=0.3)
-        ax.legend(loc='lower right', fontsize=9)
-
-    plt.tight_layout()
-
-    # Save figure (reusing the timestamp from the CSV export for consistency)
-    output_path = os.path.join(outdir, f"{timestamp}_francisco_scatter_plots.png")
-    plt.savefig(output_path, dpi=300, bbox_inches='tight')
-    print(f"\n{'='*80}")
-    print(f"Scatter plot saved to: {output_path}")
-    print(f"{'='*80}\n")
-
-    plt.show()
-
-    # Plot spectra by class with thresholds
-    thresholds_list = list(thresholds)
-    print(f"\nClass thresholds: {dict(zip(class_names, thresholds_list))}")
-    assert len(thresholds_list) == len(class_names), \
-        f"Expected {len(class_names)} thresholds, got {len(thresholds_list)}"
-
-    print("Generating spectra plots by class...")
-    plot_spectra_by_class(
-        spectra=spectra,
-        banddef=banddef,
-        y_pred=y_pred,
-        class_names=class_names,
-        thresholds=thresholds_list,
-        outdir=outdir,
-        timestamp=timestamp,
-        drop_wl_ranges=drop_wl_ranges,
-    )
-    print("Spectra plots complete.")
-
-
-if __name__ == "__main__":
-    # pylint: disable=no-value-for-parameter
-    run_francisco_evaluation()
diff --git a/models/spectf/training_spectf.py b/models/spectf/training_spectf.py
index 4940cae..6c29d2a 100644
--- a/models/spectf/training_spectf.py
+++ b/models/spectf/training_spectf.py
@@ -34,7 +34,7 @@ def __init__(self, test_X, test_Y):
 
         self.test_X = test_X
         self.test_Y = test_Y
-    
+
     def __len__(self):
         return len(self.test_Y)
 
@@ -122,13 +122,16 @@ def run_pipeline_classifier(
     with open(model_config, 'r', encoding='utf-8') as f:
         m_config = yaml.safe_load(f)
 
+    # inject_ood pulls the 'ood-train-set' spectra into the training dataloader as-is
+    # (unmixed, real labels) to intentionally induce overfitting on OOD data.
     dataloader, test_X, test_Y = setup_training_from_config(
         data_config,
         m_config['batch_size'],
         shuffle=True,
         seed=m_config['random_seed'],
         subsampled_files_outdir=outdir,
-        misc_dataloader_params={'num_workers': m_config['training']['num_workers']})
+        misc_dataloader_params={'num_workers': m_config['training']['num_workers']},
+        inject_ood=True)
 
     banddef = banddef_from_config(data_config)
 
@@ -186,6 +189,7 @@ def run_pipeline_classifier(
             "simulated_test_set_size": simulated_test_set_size,
             "focal_alpha": focal_alpha,
             "focal_gamma": focal_gamma,
+            "inject_ood": True,
         },
         settings=wandb.Settings(_service_wait=300)
     )
@@ -232,7 +236,10 @@ def run_pipeline_classifier(
 
             optimizer.zero_grad()
             logits = model(batch_X)
-            loss = criterion(logits, batch_Y)
+            # Injected OOD spectra carry a -1 sentinel for unknown/ambiguous entries;
+            # mask those out so they don't contribute to the loss during backprop.
+            mask = batch_Y >= 0
+            loss = criterion(logits[mask], batch_Y[mask])
             loss.backward()
             optimizer.step()
 
@@ -301,7 +308,7 @@ def run_pipeline_classifier(
         for f in _figs: plt.close(f)
         del _figs
 
-        # Extract the thresholds used for the test set 
+        # Extract the thresholds used for the test set
         test_thresholds = [test_metrics[class_name]['Threshold'] for class_name in class_names]
 
         # Calculate OOD validation set metrics using the test set thresholds
@@ -323,4 +330,4 @@ def run_pipeline_classifier(
 
 if __name__ == "__main__":
     # pylint: disable=no-value-for-parameter
-    run_pipeline_classifier()
\ No newline at end of file
+    run_pipeline_classifier()
diff --git a/models/spectf/training_spectf_ood.py b/models/spectf/training_spectf_ood.py
deleted file mode 100644
index 6c29d2a..0000000
--- a/models/spectf/training_spectf_ood.py
+++ /dev/null
@@ -1,333 +0,0 @@
-import os
-
-# Set CUBLAS_WORKSPACE_CONFIG for deterministic CUDA behavior
-# pylint: disable=wrong-import-position
-os.environ["CUBLAS_WORKSPACE_CONFIG"] = ":4096:8"
-
-import getpass
-from datetime import datetime
-import rich_click as click
-import yaml
-import wandb
-import matplotlib.pyplot as plt
-
-import numpy as np
-import torch
-import torch.nn as nn
-import torch.optim as optim
-from torch.utils.data import Dataset, DataLoader
-import torch.nn.functional as F
-import schedulefree
-
-from cover_class.train import setup_training_from_config, make_simulation_test_set, banddef_from_config #type: ignore
-from cover_class.utils import seed as sseed, ood_test_set_from_config #type: ignore
-from cover_class.reporting import ModelConfig, Report #type: ignore
-
-from spectf.model import SpecTfEncoder
-from spectf.utils import get_device
-
-ENV_VAR_PREFIX = 'COVER_CLASS_TRAIN_'
-
-class TestDataset(Dataset):
-    def __init__(self, test_X, test_Y):
-        super().__init__()
-
-        self.test_X = test_X
-        self.test_Y = test_Y
-
-    def __len__(self):
-        return len(self.test_Y)
-
-    def __getitem__(self, idx):
-        return self.test_X[idx], self.test_Y[idx]
-
-
-class FocalLoss(nn.Module):
-    def __init__(self, alpha=0.25, gamma=2.0, reduction='mean'):
-        super(FocalLoss, self).__init__()
-        self.alpha = alpha
-        self.gamma = gamma
-        self.reduction = reduction
-
-    def forward(self, inputs, targets):
-        bce_loss = F.binary_cross_entropy_with_logits(inputs, targets, reduction='none')
-        pt = torch.exp(-bce_loss)
-        focal_loss = ((1 - pt) ** self.gamma) * bce_loss
-
-        if self.alpha is not None:
-            alpha_t = self.alpha * targets + (1 - self.alpha) * (1 - targets)
-            focal_loss = alpha_t * focal_loss
-
-        if self.reduction == 'mean':
-            return focal_loss.mean()
-        if self.reduction == 'sum':
-            return focal_loss.sum()
-
-        return focal_loss
-
-@click.command()
-@click.option(
-    "--outdir",
-    required=True,
-    type=click.Path(exists=True, dir_okay=True, file_okay=False),
-    help="Output file directory.",
-    envvar=f'{ENV_VAR_PREFIX}OUTDIR'
-)
-@click.option(
-    "--data-config",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to the YAML config for the dataloader.",
-    envvar=f'{ENV_VAR_PREFIX}_DATA_CONFIG'
-)
-@click.option(
-    "--model-config",
-    required=True,
-    type=click.Path(exists=True, dir_okay=False),
-    help="Path to the YAML config for the model architecture.",
-    envvar=f'{ENV_VAR_PREFIX}_MODEL_CONFIG'
-)
-@click.option(
-    "--simulated-test-set-size",
-    required=False,
-    type=int,
-    default=100_000,
-    help="Number of rows to generate for the simulated test set.",
-    envvar=f'{ENV_VAR_PREFIX}_SIMULATED_TEST_SET_SIZE'
-)
-@click.option(
-    "--focal-alpha",
-    required=False,
-    default="0.25",
-    help="Focal loss alpha parameter. Set to 'None' to disable.",
-    envvar=f'{ENV_VAR_PREFIX}_FOCAL_ALPHA'
-)
-@click.option(
-    "--focal-gamma",
-    required=False,
-    type=float,
-    default=2.0,
-    help="Focal loss gamma parameter.",
-    envvar=f'{ENV_VAR_PREFIX}_FOCAL_GAMMA'
-)
-def run_pipeline_classifier(
-        outdir: str,
-        data_config: str,
-        model_config: str,
-        simulated_test_set_size: int = 100_000,
-        focal_alpha: str = "0.25",
-        focal_gamma: float = 2.0
-    ):
-
-    with open(model_config, 'r', encoding='utf-8') as f:
-        m_config = yaml.safe_load(f)
-
-    # inject_ood pulls the 'ood-train-set' spectra into the training dataloader as-is
-    # (unmixed, real labels) to intentionally induce overfitting on OOD data.
-    dataloader, test_X, test_Y = setup_training_from_config(
-        data_config,
-        m_config['batch_size'],
-        shuffle=True,
-        seed=m_config['random_seed'],
-        subsampled_files_outdir=outdir,
-        misc_dataloader_params={'num_workers': m_config['training']['num_workers']},
-        inject_ood=True)
-
-    banddef = banddef_from_config(data_config)
-
-    # create simulation eval set
-    sseed(m_config['random_seed'])
-    simulation_x_test, simulation_y_test, _ = make_simulation_test_set(dataloader, test_X, test_Y, simulated_test_set_size)
-
-    # Test set dataloader
-    test_dataset = TestDataset(simulation_x_test, simulation_y_test)
-    test_dataloader = DataLoader(test_dataset, batch_size=m_config['batch_size'], shuffle=False)
-
-    # Validation set dataloader for OOD evaluation
-    ood_test_set_x, ood_test_set_y = ood_test_set_from_config(data_config)
-    ood_dataset = TestDataset(ood_test_set_x, ood_test_set_y)
-    ood_dataloader = DataLoader(ood_dataset, batch_size=m_config['batch_size'], shuffle=False)
-
-    # hardcoded GPU 0
-    device = get_device(0)
-
-    # model definition
-    model = SpecTfEncoder(banddef.to(dtype=torch.float32, device=device),
-                         dim_output=m_config['model']['dim_output'],
-                         num_heads=m_config['model']['num_heads'],
-                         dim_proj=m_config['model']['dim_proj'],
-                         dim_ff=m_config['model']['dim_ff'],
-                         dropout=m_config['model']['dropout'],
-                         agg=m_config['model']['agg'],
-                         use_residual=m_config['model']['use_residual'],
-                         num_layers=m_config['model']['num_layers']).to(device)
-
-    # criterion = nn.BCEWithLogitsLoss()
-    alpha_val = None if focal_alpha == "None" else float(focal_alpha)
-    if alpha_val is None and focal_gamma == 0.0:
-        criterion = nn.BCEWithLogitsLoss()
-    else:
-        criterion = FocalLoss(alpha=alpha_val, gamma=focal_gamma)
-
-    optimizer = schedulefree.AdamWScheduleFree(
-        (p for p in model.parameters() if p.requires_grad),
-        lr=m_config['training']['learning_rate'],
-        warmup_steps=m_config['training']['warmup_steps']
-    )
-
-    # W&B
-    timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
-    run = wandb.init(
-        entity=m_config['wandb']['entity'],
-        project=m_config['wandb']['project'],
-        name=timestamp,
-        dir='./',
-        config={
-            "outdir": outdir,
-            "data_config": data_config,
-            "model_config": model_config,
-            "simulated_test_set_size": simulated_test_set_size,
-            "focal_alpha": focal_alpha,
-            "focal_gamma": focal_gamma,
-            "inject_ood": True,
-        },
-        settings=wandb.Settings(_service_wait=300)
-    )
-
-    report = Report(
-        outdir=outdir,
-        config=data_config,
-        author=getpass.getuser(),
-        model_config=ModelConfig(
-            model=model,
-            model_name=SpecTfEncoder.__name__,
-            hyperparams={
-                "learning_rate": m_config['training']['learning_rate'],
-                "batch_size": m_config['batch_size'],
-                "optimizer": optimizer.__class__.__name__,
-                "focal_alpha": alpha_val,
-                "focal_gamma": focal_gamma,
-                "params": m_config['model']
-            },
-        ),
-        Y_test=simulation_y_test,
-        Y_ood_test=ood_test_set_y,
-        random_seed=m_config['random_seed'],
-        run_name=timestamp
-    )
-
-    # Epoch-based training loop
-    total_epochs = m_config['training']['total_epochs']
-    steps_per_epoch = m_config['training']['steps_per_epoch']
-    bs = m_config['batch_size']
-
-    ds = report.config['datasets'] # type: ignore
-    class_names = [str(c) for c in ds.keys() if ds[c] is not None and len(ds[c])]
-
-    for epoch in range(total_epochs):
-        model.train()
-        optimizer.train()
-        epoch_loss = 0.0
-        step = 0
-        for batch_X, batch_Y in dataloader:
-            batch_X = batch_X.to(device=device, dtype=torch.float32)
-            batch_X = torch.unsqueeze(batch_X, -1)
-            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
-
-            optimizer.zero_grad()
-            logits = model(batch_X)
-            # Injected OOD spectra carry a -1 sentinel for unknown/ambiguous entries;
-            # mask those out so they don't contribute to the loss during backprop.
-            mask = batch_Y >= 0
-            loss = criterion(logits[mask], batch_Y[mask])
-            loss.backward()
-            optimizer.step()
-
-            nats = loss.cpu().item()
-            epoch_loss += nats
-            run.log({"loss_train": nats, "step": step + epoch * steps_per_epoch})
-
-            step += 1
-            if step >= steps_per_epoch:
-                break
-
-        avg_epoch_loss = epoch_loss / steps_per_epoch
-        run.log({"loss_train_epoch": avg_epoch_loss, "epoch": epoch})
-
-        # Save model at each epoch
-        if (epoch + 1) % 10 == 0:
-            torch.save(model.state_dict(), os.path.join(outdir, f'model_epoch{epoch+1}.pth'))
-
-        # Test loop
-        model.eval()
-        optimizer.eval()
-        y_hat = np.zeros_like(simulation_y_test, dtype=float)
-        test_loss_sum = 0.0
-        test_batches = 0
-        for i, (batch_X, batch_Y) in enumerate(test_dataloader):
-            batch_X = batch_X.to(device=device, dtype=torch.float32)
-            batch_X = torch.unsqueeze(batch_X, -1)
-            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
-            with torch.no_grad():
-                logits = model(batch_X)
-                batch_loss = criterion(logits, batch_Y)
-                test_loss_sum += batch_loss.cpu().item()
-                test_batches += 1
-                batch_y_hat = torch.sigmoid(logits)
-                batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
-                batch_len = len(batch_y_hat)
-                y_hat[i*bs:i*bs+batch_len] = batch_y_hat
-        avg_test_loss = test_loss_sum / test_batches if test_batches > 0 else float('nan')
-
-        # OOD loop
-        y_hat_ood = np.zeros_like(ood_test_set_y, dtype=float)
-        ood_loss_sum = 0.0
-        ood_batches = 0
-        for i, (batch_X, batch_Y) in enumerate(ood_dataloader):
-            batch_X = batch_X.to(device=device, dtype=torch.float32)
-            batch_X = torch.unsqueeze(batch_X, -1)
-            batch_Y = batch_Y.to(device=device, dtype=torch.float32)
-            with torch.no_grad():
-                logits = model(batch_X)
-                mask = ~torch.isnan(batch_Y)
-                if mask.any():
-                    batch_loss = criterion(logits[mask], batch_Y[mask])
-                    ood_loss_sum += batch_loss.cpu().item()
-                    ood_batches += 1
-                batch_y_hat = torch.sigmoid(logits)
-                batch_y_hat = batch_y_hat.detach().cpu().numpy().astype(float)
-                batch_len = len(batch_y_hat)
-                y_hat_ood[i*bs:i*bs+batch_len] = batch_y_hat
-
-        avg_ood_loss = ood_loss_sum / ood_batches if ood_batches > 0 else float('nan')
-
-
-        # Calculate test set metrics using the best thresholds for the test set
-        _figs = []
-        test_metrics = report.generate_metrics(simulation_y_test, y_hat, None, _figs, class_names)
-        for f in _figs: plt.close(f)
-        del _figs
-
-        # Extract the thresholds used for the test set
-        test_thresholds = [test_metrics[class_name]['Threshold'] for class_name in class_names]
-
-        # Calculate OOD validation set metrics using the test set thresholds
-        _figs = []
-        ood_metrics = report.generate_metrics(ood_test_set_y, y_hat_ood, test_thresholds, _figs, class_names)
-        for f in _figs: plt.close(f)
-        del _figs
-
-        run.log({
-            "test_metrics": test_metrics,
-            "ood_metrics": ood_metrics,
-            "epoch": epoch,
-            "loss_test_epoch": avg_test_loss,
-            "loss_ood_epoch": avg_ood_loss
-        })
-
-    # Generate report at end of training
-    report.make_report(y_hat, y_hat_ood, None, ood_overfit=False)
-
-if __name__ == "__main__":
-    # pylint: disable=no-value-for-parameter
-    run_pipeline_classifier()