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() 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_unmixing.py b/models/spectf/training_spectf_unmixing.py new file mode 100644 index 0000000..4356300 --- /dev/null +++ b/models/spectf/training_spectf_unmixing.py @@ -0,0 +1,381 @@ +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_fractions = simulation_y_fractions.cpu().numpy() + 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/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/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..b47b434 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 @@ -19,24 +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): @@ -64,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 ''' @@ -96,17 +134,40 @@ 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 + 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 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() @@ -115,26 +176,47 @@ 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( - 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 = {}, + ood_spectra: Optional[FloatTensor] = None, + ood_labels: Optional[Tensor] = None, ) -> DataLoader: config = read_config(config) @@ -146,7 +228,11 @@ def dataloader_from_config( sim_config_args, sim_data_args, spectra, - labels.long() + labels.long(), + 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/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) diff --git a/src/cover_class/train.py b/src/cover_class/train.py index 634e373..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 @@ -63,16 +63,20 @@ 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 = {}, + 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 """ @@ -86,7 +90,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: @@ -112,13 +116,28 @@ 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, + config, FloatTensor(train_spectra), LongTensor(train_labels.to(dtype=torch.long)), batch_size, 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)