diff --git a/RuBR/README.md b/RuBR/README.md new file mode 100644 index 00000000..f0e498d9 --- /dev/null +++ b/RuBR/README.md @@ -0,0 +1,267 @@ +# RuBR: Model + Experiments Repository + +RuBR contains reusable model components in `model/` and experiment entry scripts in `experiments/`. + +## Directory structure + +- `model/` + - `dann_model.py` (DANN architecture) + - `data.py` (data loading and normalization helpers) + - `layers.py` (shared rotation-invariant and GRL layers) + - `callbacks.py` (training callbacks such as F1-based early stopping) + - `metadata.py` (metadata helpers used by evaluation/analysis) +- `experiments/` + - `rotinv/` + - `train_rotinv_with_features.py` + - `evaluate_rotinv_with_features.py` + - `domain_adaption/` + - `train_dann.py` + - `test_dann.py` + - `train_control_baseline.py` + - `test_control_baseline.py` + - `analysis/` + - `compare_magnitude_histograms.py` + +## Setup + +Requires Python 3.11. + +Install dependencies: + +```bash +pip install -r requirements.txt +``` + +## Data download + +[RAPID Pipeline Products — Public Access](https://caltech-ipac-rapid.readthedocs.io/en/latest/prod/products.html#public-access) + +This page contains instructions to download RAPID pipeline data. + +## Dataset format requirements + +Required keys in every dataset file: + +- `X` +- `feats` +- `y` +- `metadata` + +Required labels: + +- `y.shape == (N,)` +- `y` contains `0` and `1` + +Required image format: + +- `X.shape == (N, 64, 64, 3)` or `X.shape == (N, 3, 64, 64)` +- channel order: `ref, sci, diff` + +Required metadata: + +- `len(metadata) == N` +- for `experiments/rotinv/train_rotinv_with_features.py` and `experiments/rotinv/evaluate_rotinv_with_features.py`: + - each `metadata[i]` must contain `filter` + - allowed `filter` values: `F184`, `H158`, `J129`, `K213`, `R062`, `Y106`, `Z087` + +Required file format: + +- repo-native rotinv scripts accept `.npz` and `.npy` +- repo-native domain-adaptation scripts accept `.npz` +- `X` and `feats` must not contain NaNs + +### Per-script feature requirements + +| Script / model | Required `feats` shape | Required feature order | +|---|---:|---| +| `experiments/rotinv/train_rotinv_with_features.py` | `(N, F)` | exact column order from input dataset; train/val/test must match | +| `experiments/rotinv/evaluate_rotinv_with_features.py` | `(N, F)` | exact column order used during training | +| `experiments/domain_adaption/train_control_baseline.py` | `(N, F)` | exact column order from input dataset; train/val/test must match | +| `experiments/domain_adaption/test_control_baseline.py` | `(N, F)` | exact column order used during training | +| `experiments/domain_adaption/train_dann.py` | `(N, F)` | exact column order from source/target training datasets; all splits must match | +| `experiments/domain_adaption/test_dann.py` | `(N, F)` | exact column order used during training | +| `scripts/eval_comb_author.py` | `(N, 5)` | `flux, mag, npix, roundness, sharpness` | +| `scripts/run_author_test_compat.py control` | `(N, 6)` | `mag, roundness, sharpness, peak, flux, npix` | +| `scripts/run_author_test_compat.py dann` | `(N, 6)` | `mag, roundness, sharpness, peak, flux, npix` | + +### Object-dtype `feats` + +If `feats` is stored as an object array of dicts: + +- repo-native loaders use `list(f.values())` +- `scripts/eval_comb_author.py` uses sorted keys and expects: + - `flux` + - `mag` + - `npix` + - `roundness` + - `sharpness` + +## Experiment details and usage + +### 1) RotInv training + +Script: `experiments/rotinv/train_rotinv_with_features.py` + +Purpose: +- Trains rotationally invariant hybrid classifier (image branch + tabular features). + +Main input: +- `--data_path`: training dataset (`.npz` or supported format in script) + +Main outputs: +- Saved model checkpoints/final model +- Training history plot and evaluation artifacts under `--output_dir` + +Example: + +```bash +python -m experiments.rotinv.train_rotinv_with_features \ + --data_path /path/to/train_data.npz \ + --output_dir ./outputs/rotinv_train +``` + +### 2) RotInv evaluation + +Script: `experiments/rotinv/evaluate_rotinv_with_features.py` + +Purpose: +- Evaluates a trained RotInv model over batched test data. + +Main inputs: +- `--data_dir`: directory containing test batches +- `--model_path`: trained model file + +Main outputs: +- Precision/recall threshold plot +- Magnitude histogram and summary metrics under `--output_dir` + +Example: + +```bash +python -m experiments.rotinv.evaluate_rotinv_with_features \ + --data_dir /path/to/test_batches \ + --model_path /path/to/rotinv_model.h5 \ + --output_dir ./outputs/rotinv_eval +``` + +### 3) DANN training (domain adaptation) + +Script: `experiments/domain_adaption/train_dann.py` + +Purpose: +- Trains domain-adversarial model using source + target domains. + +Main inputs: +- `--source_data`: source-domain training set +- `--target_data`: target-domain training set + +Main outputs: +- Best/final DANN models +- Training curves and summary text under `--output_dir` + +Example: + +```bash +python -m experiments.domain_adaption.train_dann \ + --source_data /path/to/source_train.npz \ + --target_data /path/to/target_train.npz \ + --output_dir ./outputs/dann_train +``` + +### 4) DANN evaluation + +Script: `experiments/domain_adaption/test_dann.py` + +Purpose: +- Evaluates trained DANN model on source and target test sets. + +Main inputs: +- `--model_path` +- `--source_test` +- `--target_test` + +Main outputs: +- Per-domain metrics reports, plots, and summary files under `--output_dir` + +Example: + +```bash +python -m experiments.domain_adaption.test_dann \ + --model_path /path/to/dann_model.h5 \ + --source_test /path/to/source_test.npz \ + --target_test /path/to/target_test.npz \ + --output_dir ./outputs/dann_eval +``` + +### 5) Control baseline training + +Script: `experiments/domain_adaption/train_control_baseline.py` + +Purpose: +- Trains non-adversarial control baseline for direct comparison with DANN. + +Main input: +- `--train_data` + +Main outputs: +- Control model checkpoints/final model and training summaries under `--output_dir` + +Example: + +```bash +python -m experiments.domain_adaption.train_control_baseline \ + --train_data /path/to/control_train.npz \ + --output_dir ./outputs/control_train +``` + +### 6) Control baseline evaluation + +Script: `experiments/domain_adaption/test_control_baseline.py` + +Purpose: +- Evaluates control baseline on test data and reports ROC/PR + confusion metrics. + +Main inputs: +- `--model_path` +- `--test_data` + +Main outputs: +- Test metrics report and evaluation plots under `--output_dir` + +Example: + +```bash +python -m experiments.domain_adaption.test_control_baseline \ + --model_path /path/to/control_model.h5 \ + --test_data /path/to/control_test.npz \ + --output_dir ./outputs/control_eval +``` + +### 7) DANN vs Control analysis + +Script: `experiments/analysis/compare_magnitude_histograms.py` + +Purpose: +- Produces magnitude-based comparison plots between DANN and control models. + +Main inputs: +- `--dann_model` +- `--control_model` +- `--test_data` + +Main outputs: +- Comparison plots and text summaries under `--output_dir` + +Example: + +```bash +python -m experiments.analysis.compare_magnitude_histograms \ + --dann_model /path/to/dann_model.h5 \ + --control_model /path/to/control_model.h5 \ + --test_data /path/to/compare_test.npz \ + --output_dir ./outputs/model_comparison +``` + +## Notes + +- Use module execution (`python -m ...`) from the repository root. diff --git a/RuBR/experiments/analysis/compare_magnitude_histograms.py b/RuBR/experiments/analysis/compare_magnitude_histograms.py new file mode 100644 index 00000000..40caa182 --- /dev/null +++ b/RuBR/experiments/analysis/compare_magnitude_histograms.py @@ -0,0 +1,914 @@ +""" +Compare Magnitude Distributions: DANN vs Control Model + +This script creates magnitude histograms comparing detection performance +between DANN and control models on the target domain. +""" + +import argparse +import os +import numpy as np +import tensorflow as tf +from matplotlib import pyplot as plt +import scienceplots +from model.data import load_dataset, normalize_arrays +from model.layers import ( + rot90_k1, + rot90_k2, + rot90_k3, + gradient_reversal, + GradientReversalLayer, +) +from model.metadata import get_transient_magnitude + +def load_data(data_path): + return load_dataset(data_path, mmap=False, allow_npy_dict=False) + + +def get_predictions(model, X, feats): + """Get model predictions.""" + # Normalize data + X_norm, feats_norm = normalize_arrays(X, feats) + + # Get predictions + predictions = model.predict([X_norm, feats_norm], verbose=0) + + # Handle different output formats (DANN has 2 outputs) + if isinstance(predictions, list): + y_pred_prob = predictions[0].flatten() + else: + y_pred_prob = predictions.flatten() + + return y_pred_prob + + +def main(): + parser = argparse.ArgumentParser( + description="Compare magnitude histograms between DANN and control models" + ) + + parser.add_argument( + "--dann_model", + type=str, + default="./outputs/dann_train/best_model.h5", + help="Path to DANN model (.h5 file)", + ) + parser.add_argument( + "--control_model", + type=str, + default="./outputs/control_train/best_model.h5", + help="Path to control model (.h5 file)", + ) + parser.add_argument( + "--test_data", + type=str, + default="./data/target_test.npz", + help="Path to target domain test data (.npz file)", + ) + parser.add_argument( + "--output_dir", + type=str, + default="./model_comparison_plots", + help="Directory to save plots (default: ./model_comparison_plots)", + ) + parser.add_argument( + "--threshold", + type=float, + default=0.5, + help="Classification threshold for both models (default: 0.5). Overridden by dann_threshold/control_threshold if specified.", + ) + parser.add_argument( + "--dann_threshold", + type=float, + default=None, + help="Classification threshold for DANN model (default: uses --threshold value)", + ) + parser.add_argument( + "--control_threshold", + type=float, + default=None, + help="Classification threshold for control model (default: uses --threshold value)", + ) + parser.add_argument( + "--gpu", + type=int, + default=None, + help="GPU device ID to use (default: auto-select)", + ) + + args = parser.parse_args() + + # Set individual thresholds (use --threshold as default if not specified) + dann_threshold = args.dann_threshold if args.dann_threshold is not None else args.threshold + control_threshold = args.control_threshold if args.control_threshold is not None else args.threshold + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Set GPU device if specified + if args.gpu is not None: + physical_devices = tf.config.experimental.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_visible_devices( + physical_devices[args.gpu], "GPU" + ) + tf.config.experimental.set_memory_growth(physical_devices[args.gpu], True) + + print("=" * 80) + print("MODEL COMPARISON - MAGNITUDE HISTOGRAMS") + print("=" * 80) + print(f"DANN model: {args.dann_model}") + print(f"Control model: {args.control_model}") + print(f"Test data: {args.test_data}") + print(f"DANN threshold: {dann_threshold}") + print(f"Control threshold: {control_threshold}") + print(f"Output directory: {args.output_dir}") + print("=" * 80) + + # Load models + print("\nLoading DANN model...") + dann_model = tf.keras.models.load_model( + args.dann_model, + custom_objects={ + "GradientReversalLayer": GradientReversalLayer, + "rot90_k1": rot90_k1, + "rot90_k2": rot90_k2, + "rot90_k3": rot90_k3 + } + ) + print("DANN model loaded successfully!") + + print("\nLoading control model...") + control_model = tf.keras.models.load_model( + args.control_model, + custom_objects={ + "rot90_k1": rot90_k1, + "rot90_k2": rot90_k2, + "rot90_k3": rot90_k3 + } + ) + print("Control model loaded successfully!") + + # Load test data + print("\n--- Loading Target Domain Test Data ---") + X, feats, y, metadata = load_data(args.test_data) + + # Remove NaN values + mask = np.isnan(X).any(axis=(1, 2, 3)) | np.isnan(feats).any(axis=1) + if mask.any(): + print(f"Removing {mask.sum()} samples with NaN values") + X = X[~mask] + feats = feats[~mask] + y = y[~mask] + metadata = metadata[~mask] + + # Get predictions from both models + print("\nGenerating predictions from DANN model...") + dann_pred_prob = get_predictions(dann_model, X, feats) + + print("Generating predictions from control model...") + control_pred_prob = get_predictions(control_model, X, feats) + + # Extract magnitude values for ground truth positives (y == 1) + print("\nExtracting magnitude values for ground truth positives...") + mag_values = [] + dann_probs = [] + control_probs = [] + filter_values = [] + + for idx, m in enumerate(metadata): + # CRITICAL: Only process ground truth positives (y == 1) + if y[idx] != 1: + continue + + try: + if isinstance(m, dict) and 'match_id' in m and 'jid_folder' in m: + mag = get_transient_magnitude(m['match_id'], m['jid_folder']) + if not np.isnan(mag): + mag_values.append(mag) + dann_probs.append(dann_pred_prob[idx]) + control_probs.append(control_pred_prob[idx]) + # Extract filter information + filter_name = m.get('filter', m.get('band', 'unknown')) + filter_values.append(filter_name) + except Exception as e: + print(f"Error getting magnitude for sample {idx}: {e}") + + if len(mag_values) == 0: + print("No valid magnitude values found. Exiting.") + return + + mag_values = np.array(mag_values) + dann_probs = np.array(dann_probs) + control_probs = np.array(control_probs) + filter_values = np.array(filter_values) + + print(f"Total ground truth positives with valid magnitudes: {len(mag_values)}") + print(f"Unique filters found: {np.unique(filter_values)}") + + # Apply thresholds to get true positives for each model + dann_tp_mask = dann_probs > dann_threshold + control_tp_mask = control_probs > control_threshold + + mag_all_positives = mag_values + mag_dann_tp = mag_values[dann_tp_mask] + mag_control_tp = mag_values[control_tp_mask] + + print(f"DANN true positives: {len(mag_dann_tp)} ({len(mag_dann_tp)/len(mag_values)*100:.1f}%)") + print(f"Control true positives: {len(mag_control_tp)} ({len(mag_control_tp)/len(mag_values)*100:.1f}%)") + + # Create comparison histogram + print("\nCreating comparison histogram...") + + # Apply scienceplots style + plt.style.use(['science', 'ieee', 'no-latex']) + plt.rcParams.update({ + 'font.size': 10, + 'font.family': 'serif', + 'axes.labelsize': 11, + 'axes.titlesize': 11, + 'xtick.labelsize': 9, + 'ytick.labelsize': 9, + 'legend.fontsize': 9, + 'lines.linewidth': 1.5, + 'grid.linewidth': 0.5, + 'axes.linewidth': 0.8, + }) + + fig, ax = plt.subplots(figsize=(3.5, 2.8), dpi=600) + + # Colors for the three distributions + colors = { + 'all': '#000000', # Black - all ground truth + 'dann': '#0000FF', # Blue - DANN detections + 'control': '#FF0000', # Red - Control detections + } + + # Use common bins based on all positives range + bins = np.linspace(mag_all_positives.min(), mag_all_positives.max(), 30) + + # Plot all ground truth positives + ax.hist(mag_all_positives, + bins=bins, + label=f"Ground Truth (N={len(mag_all_positives)})", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=2) + + # Plot DANN true positives + if len(mag_dann_tp) > 0: + ax.hist(mag_dann_tp, + bins=bins, + label=f"Domain Adversarial Training (N={len(mag_dann_tp)})", + color=colors['dann'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=3) + + # Plot control true positives + if len(mag_control_tp) > 0: + ax.hist(mag_control_tp, + bins=bins, + label=f"No Domain Adaptation (N={len(mag_control_tp)})", + color=colors['control'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=4) + + # Set log scale + # ax.set_yscale("log") + + # Labels + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Count", fontsize=10) + ax.set_title("Target Domain - Model Comparison", fontsize=10, pad=8) + + # Legend + ax.legend(loc='best', frameon=True, + fontsize=8, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.5, labelspacing=0.3) + + # Grid + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Styling + ax.set_facecolor('white') + fig.patch.set_facecolor('white') + + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + plt.tight_layout(pad=0.3) + + # Save + plot_path = os.path.join(args.output_dir, "model_comparison_magnitude_histogram.png") + plt.savefig(plot_path, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(plot_path.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f"\nComparison histogram saved to: {plot_path}") + print(f"PDF version: {plot_path.replace('.png', '.pdf')}") + + # Create cumulative histogram + print("\nCreating cumulative histogram...") + + fig, ax = plt.subplots(figsize=(3.5, 2.8), dpi=600) + + # Plot cumulative distributions + ax.hist(mag_all_positives, + bins=bins, + label=f"Ground Truth (N={len(mag_all_positives)})", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=2) + + if len(mag_dann_tp) > 0: + ax.hist(mag_dann_tp, + bins=bins, + label=f"Domain Adversarial Training (N={len(mag_dann_tp)})", + color=colors['dann'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=3) + + if len(mag_control_tp) > 0: + ax.hist(mag_control_tp, + bins=bins, + label=f"No Domain Adaptation (N={len(mag_control_tp)})", + color=colors['control'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=4) + + # Labels + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Cumulative Count", fontsize=10) + ax.set_title("Target Domain - Model Comparison (Cumulative)", fontsize=10, pad=8) + + # Legend + ax.legend(loc='best', frameon=True, + fontsize=8, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.5, labelspacing=0.3) + + # Grid + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Styling + ax.set_facecolor('white') + fig.patch.set_facecolor('white') + + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + plt.tight_layout(pad=0.3) + + # Save + cumulative_plot_path = os.path.join(args.output_dir, "model_comparison_magnitude_histogram_cumulative.png") + plt.savefig(cumulative_plot_path, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(cumulative_plot_path.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f"\nCumulative histogram saved to: {cumulative_plot_path}") + print(f"PDF version: {cumulative_plot_path.replace('.png', '.pdf')}") + + # Print summary statistics + print("\n" + "=" * 80) + print("SUMMARY STATISTICS") + print("=" * 80) + print(f"Ground Truth Positives: {len(mag_all_positives)}") + print(f"Magnitude range: [{mag_all_positives.min():.2f}, {mag_all_positives.max():.2f}]") + print() + print(f"DANN Model:") + print(f" True Positives: {len(mag_dann_tp)}") + print(f" Detection Rate: {len(mag_dann_tp)/len(mag_all_positives)*100:.1f}%") + if len(mag_dann_tp) > 0: + print(f" Magnitude range: [{mag_dann_tp.min():.2f}, {mag_dann_tp.max():.2f}]") + print() + print(f"Control Model:") + print(f" True Positives: {len(mag_control_tp)}") + print(f" Detection Rate: {len(mag_control_tp)/len(mag_all_positives)*100:.1f}%") + if len(mag_control_tp) > 0: + print(f" Magnitude range: [{mag_control_tp.min():.2f}, {mag_control_tp.max():.2f}]") + print("=" * 80) + + # Save statistics to file + stats_path = os.path.join(args.output_dir, "comparison_statistics.txt") + with open(stats_path, 'w') as f: + f.write("MODEL COMPARISON - MAGNITUDE STATISTICS\n") + f.write("=" * 80 + "\n\n") + f.write(f"DANN Model: {args.dann_model}\n") + f.write(f"Control Model: {args.control_model}\n") + f.write(f"Test Data: {args.test_data}\n") + f.write(f"DANN Threshold: {dann_threshold}\n") + f.write(f"Control Threshold: {control_threshold}\n\n") + f.write(f"Ground Truth Positives: {len(mag_all_positives)}\n") + f.write(f"Magnitude range: [{mag_all_positives.min():.2f}, {mag_all_positives.max():.2f}]\n\n") + f.write(f"DANN Model:\n") + f.write(f" True Positives: {len(mag_dann_tp)}\n") + f.write(f" Detection Rate: {len(mag_dann_tp)/len(mag_all_positives)*100:.1f}%\n") + if len(mag_dann_tp) > 0: + f.write(f" Magnitude range: [{mag_dann_tp.min():.2f}, {mag_dann_tp.max():.2f}]\n") + f.write(f"\nControl Model:\n") + f.write(f" True Positives: {len(mag_control_tp)}\n") + f.write(f" Detection Rate: {len(mag_control_tp)/len(mag_all_positives)*100:.1f}%\n") + if len(mag_control_tp) > 0: + f.write(f" Magnitude range: [{mag_control_tp.min():.2f}, {mag_control_tp.max():.2f}]\n") + + print(f"\nStatistics saved to: {stats_path}") + + # Create filter-wise histograms + print("\n" + "=" * 80) + print("CREATING FILTER-WISE HISTOGRAMS") + print("=" * 80) + + unique_filters = np.unique(filter_values) + print(f"Filters to process: {unique_filters}") + + for filter_name in unique_filters: + print(f"\nProcessing filter: {filter_name}") + + # Get mask for this filter + filter_mask = filter_values == filter_name + + mag_filter = mag_values[filter_mask] + dann_probs_filter = dann_probs[filter_mask] + control_probs_filter = control_probs[filter_mask] + + if len(mag_filter) == 0: + print(f" No samples for filter {filter_name}, skipping.") + continue + + print(f" Total samples: {len(mag_filter)}") + + # Apply thresholds + dann_tp_mask_filter = dann_probs_filter > dann_threshold + control_tp_mask_filter = control_probs_filter > control_threshold + + mag_dann_tp_filter = mag_filter[dann_tp_mask_filter] + mag_control_tp_filter = mag_filter[control_tp_mask_filter] + + print(f" DANN detections: {len(mag_dann_tp_filter)} ({len(mag_dann_tp_filter)/len(mag_filter)*100:.1f}%)") + print(f" Control detections: {len(mag_control_tp_filter)} ({len(mag_control_tp_filter)/len(mag_filter)*100:.1f}%)") + + # Create histogram for this filter + fig, ax = plt.subplots(figsize=(3.5, 2.8), dpi=600) + + # Use common bins based on filter data range + bins = np.linspace(mag_filter.min(), mag_filter.max(), 30) + + # Plot all ground truth positives for this filter + ax.hist(mag_filter, + bins=bins, + label=f"Ground Truth (N={len(mag_filter)})", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=2) + + # Plot DANN true positives + if len(mag_dann_tp_filter) > 0: + ax.hist(mag_dann_tp_filter, + bins=bins, + label=f"Domain Adversarial Training (N={len(mag_dann_tp_filter)})", + color=colors['dann'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=3) + + # Plot control true positives + if len(mag_control_tp_filter) > 0: + ax.hist(mag_control_tp_filter, + bins=bins, + label=f"No Domain Adaptation (N={len(mag_control_tp_filter)})", + color=colors['control'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=4) + + # Labels + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Count", fontsize=10) + ax.set_title(f"Target Domain - Model Comparison ({filter_name} filter)", fontsize=10, pad=8) + + # Legend + ax.legend(loc='best', frameon=True, + fontsize=8, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.5, labelspacing=0.3) + + # Grid + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Styling + ax.set_facecolor('white') + fig.patch.set_facecolor('white') + + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + plt.tight_layout(pad=0.3) + + # Save + filter_safe = str(filter_name).replace('/', '_').replace(' ', '_') + plot_path_filter = os.path.join(args.output_dir, f"model_comparison_magnitude_histogram_{filter_safe}.png") + plt.savefig(plot_path_filter, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(plot_path_filter.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f" Saved: {plot_path_filter}") + + # Create cumulative histogram for this filter + fig, ax = plt.subplots(figsize=(3.5, 2.8), dpi=600) + + # Plot cumulative distributions + ax.hist(mag_filter, + bins=bins, + label=f"Ground Truth (N={len(mag_filter)})", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=2) + + if len(mag_dann_tp_filter) > 0: + ax.hist(mag_dann_tp_filter, + bins=bins, + label=f"Domain Adversarial Training (N={len(mag_dann_tp_filter)})", + color=colors['dann'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=3) + + if len(mag_control_tp_filter) > 0: + ax.hist(mag_control_tp_filter, + bins=bins, + label=f"No Domain Adaptation (N={len(mag_control_tp_filter)})", + color=colors['control'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=4) + + # Labels + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Cumulative Count", fontsize=10) + ax.set_title(f"Target Domain - Model Comparison ({filter_name} filter, Cumulative)", fontsize=10, pad=8) + + # Legend + ax.legend(loc='best', frameon=True, + fontsize=8, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.5, labelspacing=0.3) + + # Grid + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Styling + ax.set_facecolor('white') + fig.patch.set_facecolor('white') + + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + plt.tight_layout(pad=0.3) + + # Save + plot_path_filter_cumulative = os.path.join(args.output_dir, f"model_comparison_magnitude_histogram_{filter_safe}_cumulative.png") + plt.savefig(plot_path_filter_cumulative, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(plot_path_filter_cumulative.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f" Saved cumulative: {plot_path_filter_cumulative}") + + # Save filter-wise statistics to file + filter_stats_path = os.path.join(args.output_dir, "comparison_statistics_by_filter.txt") + with open(filter_stats_path, 'w') as f: + f.write("MODEL COMPARISON - MAGNITUDE STATISTICS BY FILTER\n") + f.write("=" * 80 + "\n\n") + f.write(f"DANN Model: {args.dann_model}\n") + f.write(f"Control Model: {args.control_model}\n") + f.write(f"Test Data: {args.test_data}\n") + f.write(f"DANN Threshold: {dann_threshold}\n") + f.write(f"Control Threshold: {control_threshold}\n\n") + + for filter_name in unique_filters: + filter_mask = filter_values == filter_name + mag_filter = mag_values[filter_mask] + dann_probs_filter = dann_probs[filter_mask] + control_probs_filter = control_probs[filter_mask] + + if len(mag_filter) == 0: + continue + + dann_tp_mask_filter = dann_probs_filter > dann_threshold + control_tp_mask_filter = control_probs_filter > control_threshold + + mag_dann_tp_filter = mag_filter[dann_tp_mask_filter] + mag_control_tp_filter = mag_filter[control_tp_mask_filter] + + f.write(f"\nFILTER: {filter_name}\n") + f.write("-" * 80 + "\n") + f.write(f"Ground Truth Positives: {len(mag_filter)}\n") + f.write(f"Magnitude range: [{mag_filter.min():.2f}, {mag_filter.max():.2f}]\n\n") + f.write(f"DANN Model:\n") + f.write(f" True Positives: {len(mag_dann_tp_filter)}\n") + f.write(f" Detection Rate: {len(mag_dann_tp_filter)/len(mag_filter)*100:.1f}%\n") + if len(mag_dann_tp_filter) > 0: + f.write(f" Magnitude range: [{mag_dann_tp_filter.min():.2f}, {mag_dann_tp_filter.max():.2f}]\n") + f.write(f"\nControl Model:\n") + f.write(f" True Positives: {len(mag_control_tp_filter)}\n") + f.write(f" Detection Rate: {len(mag_control_tp_filter)/len(mag_filter)*100:.1f}%\n") + if len(mag_control_tp_filter) > 0: + f.write(f" Magnitude range: [{mag_control_tp_filter.min():.2f}, {mag_control_tp_filter.max():.2f}]\n") + + print(f"\nFilter-wise statistics saved to: {filter_stats_path}") + + # Create combined multi-panel figure with all filters + print("\n" + "=" * 80) + print("CREATING COMBINED MULTI-PANEL FIGURE") + print("=" * 80) + + n_filters = len(unique_filters) + if n_filters > 0: + # Determine layout (rows x cols) + if n_filters == 1: + n_rows, n_cols = 1, 1 + elif n_filters == 2: + n_rows, n_cols = 1, 2 + elif n_filters <= 4: + n_rows, n_cols = 2, 2 + elif n_filters <= 6: + n_rows, n_cols = 2, 3 + elif n_filters <= 9: + n_rows, n_cols = 3, 3 + else: + n_rows = int(np.ceil(n_filters / 3)) + n_cols = 3 + + fig, axes = plt.subplots(n_rows, n_cols, figsize=(3.5*n_cols, 2.8*n_rows), dpi=600) + + # Handle single subplot case + if n_filters == 1: + axes = np.array([axes]) + axes = axes.flatten() + + for idx, filter_name in enumerate(unique_filters): + ax = axes[idx] + + # Get mask for this filter + filter_mask = filter_values == filter_name + + mag_filter = mag_values[filter_mask] + dann_probs_filter = dann_probs[filter_mask] + control_probs_filter = control_probs[filter_mask] + + if len(mag_filter) == 0: + ax.text(0.5, 0.5, f"No data\n({filter_name})", + ha='center', va='center', transform=ax.transAxes) + ax.set_xticks([]) + ax.set_yticks([]) + continue + + # Apply thresholds + dann_tp_mask_filter = dann_probs_filter > dann_threshold + control_tp_mask_filter = control_probs_filter > control_threshold + + mag_dann_tp_filter = mag_filter[dann_tp_mask_filter] + mag_control_tp_filter = mag_filter[control_tp_mask_filter] + + # Use common bins based on filter data range + bins = np.linspace(mag_filter.min(), mag_filter.max(), 30) + + # Plot all ground truth positives for this filter + ax.hist(mag_filter, + bins=bins, + label=f"Ground Truth (N={len(mag_filter)})", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=2) + + # Plot DANN true positives + if len(mag_dann_tp_filter) > 0: + ax.hist(mag_dann_tp_filter, + bins=bins, + label=f"DANN (N={len(mag_dann_tp_filter)})", + color=colors['dann'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=3) + + # Plot control true positives + if len(mag_control_tp_filter) > 0: + ax.hist(mag_control_tp_filter, + bins=bins, + label=f"Control (N={len(mag_control_tp_filter)})", + color=colors['control'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=4) + + # Labels + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Count", fontsize=10) + ax.set_title(f"{filter_name} filter", fontsize=10, pad=8) + + # Legend with smaller font for multi-panel + ax.legend(loc='best', frameon=True, + fontsize=7, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.4, labelspacing=0.2) + + # Grid + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Styling + ax.set_facecolor('white') + + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + # Hide unused subplots + for idx in range(n_filters, len(axes)): + axes[idx].axis('off') + + fig.patch.set_facecolor('white') + plt.suptitle("Model Comparison by Filter - Target Domain", + fontsize=12, fontweight='bold', y=0.995) + plt.tight_layout(pad=0.5, rect=[0, 0, 1, 0.99]) + + # Save combined figure + combined_plot_path = os.path.join(args.output_dir, "model_comparison_magnitude_histogram_all_filters.png") + plt.savefig(combined_plot_path, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(combined_plot_path.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f"\nCombined multi-panel figure saved to: {combined_plot_path}") + print(f"PDF version: {combined_plot_path.replace('.png', '.pdf')}") + + # Create combined multi-panel CUMULATIVE figure + print("\nCreating combined multi-panel cumulative figure...") + + fig, axes = plt.subplots(n_rows, n_cols, figsize=(3.5*n_cols, 2.8*n_rows), dpi=600) + + # Handle single subplot case + if n_filters == 1: + axes = np.array([axes]) + axes = axes.flatten() + + for idx, filter_name in enumerate(unique_filters): + ax = axes[idx] + + # Get mask for this filter + filter_mask = filter_values == filter_name + + mag_filter = mag_values[filter_mask] + dann_probs_filter = dann_probs[filter_mask] + control_probs_filter = control_probs[filter_mask] + + if len(mag_filter) == 0: + ax.text(0.5, 0.5, f"No data\n({filter_name})", + ha='center', va='center', transform=ax.transAxes) + ax.set_xticks([]) + ax.set_yticks([]) + continue + + # Apply thresholds + dann_tp_mask_filter = dann_probs_filter > dann_threshold + control_tp_mask_filter = control_probs_filter > control_threshold + + mag_dann_tp_filter = mag_filter[dann_tp_mask_filter] + mag_control_tp_filter = mag_filter[control_tp_mask_filter] + + # Use common bins based on filter data range + bins = np.linspace(mag_filter.min(), mag_filter.max(), 30) + + # Plot cumulative distributions + ax.hist(mag_filter, + bins=bins, + label=f"Ground Truth (N={len(mag_filter)})", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=2) + + if len(mag_dann_tp_filter) > 0: + ax.hist(mag_dann_tp_filter, + bins=bins, + label=f"DANN (N={len(mag_dann_tp_filter)})", + color=colors['dann'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=3) + + if len(mag_control_tp_filter) > 0: + ax.hist(mag_control_tp_filter, + bins=bins, + label=f"Control (N={len(mag_control_tp_filter)})", + color=colors['control'], + histtype="step", + linewidth=1.5, + linestyle='-', + cumulative=True, + zorder=4) + + # Labels + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Cumulative Count", fontsize=10) + ax.set_title(f"{filter_name} filter (Cumulative)", fontsize=10, pad=8) + + # Legend with smaller font for multi-panel + ax.legend(loc='best', frameon=True, + fontsize=7, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.4, labelspacing=0.2) + + # Grid + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Styling + ax.set_facecolor('white') + + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + # Hide unused subplots + for idx in range(n_filters, len(axes)): + axes[idx].axis('off') + + fig.patch.set_facecolor('white') + plt.suptitle("Model Comparison by Filter - Target Domain (Cumulative)", + fontsize=12, fontweight='bold', y=0.995) + plt.tight_layout(pad=0.5, rect=[0, 0, 1, 0.99]) + + # Save combined cumulative figure + combined_cumulative_plot_path = os.path.join(args.output_dir, "model_comparison_magnitude_histogram_all_filters_cumulative.png") + plt.savefig(combined_cumulative_plot_path, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(combined_cumulative_plot_path.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f"\nCombined multi-panel cumulative figure saved to: {combined_cumulative_plot_path}") + print(f"PDF version: {combined_cumulative_plot_path.replace('.png', '.pdf')}") + + print("\nComparison completed successfully!") + + +if __name__ == "__main__": + main() diff --git a/RuBR/experiments/domain_adaption/test_control_baseline.py b/RuBR/experiments/domain_adaption/test_control_baseline.py new file mode 100644 index 00000000..99acc695 --- /dev/null +++ b/RuBR/experiments/domain_adaption/test_control_baseline.py @@ -0,0 +1,314 @@ +""" +Test Control Model on Test Data + +This script evaluates a trained standard (non-domain-adversarial) model +on test data to establish a baseline performance for comparison. +""" + +import argparse +import os +import tensorflow as tf +import numpy as np +from sklearn.metrics import ( + confusion_matrix, + classification_report, + roc_auc_score, + precision_recall_curve, + auc, + roc_curve +) +from matplotlib import pyplot as plt +from model.data import load_dataset +from model.layers import rot90_k1, rot90_k2, rot90_k3 + + +def load_data(data_path): + """ + Load the dataset from the specified path. + + Args: + data_path (str): Path to the dataset file (.npz file). + + Returns: + tuple: X, feats, y, metadata as numpy arrays. + """ + return load_dataset(data_path, mmap=False, allow_npy_dict=False) + + +def evaluate_model(model, X, feats, y, dataset_name, output_dir): + """ + Evaluate model on test data. + + Args: + model: Trained model + X: Image data + feats: Tabular features + y: Labels + dataset_name: Name of the dataset (for logging) + output_dir: Directory to save results + + Returns: + dict: Dictionary of evaluation metrics + """ + print(f"\n{'='*60}") + print(f"Evaluating on {dataset_name}") + print(f"{'='*60}") + + # Remove NaN values + mask = np.isnan(X).any(axis=(1, 2, 3)) | np.isnan(feats).any(axis=1) + if mask.any(): + print(f"Removing {mask.sum()} samples with NaN values") + X = X[~mask] + feats = feats[~mask] + y = y[~mask] + + # Normalize data (using same approach as training) + X = (X - X.mean(axis=(0, 1, 2), keepdims=True)) / (X.std(axis=(0, 1, 2), keepdims=True) + 1e-6) + feats = (feats - feats.mean(axis=0)) / (feats.std(axis=0) + 1e-6) + + # Get predictions + y_pred_prob = model.predict([X, feats], verbose=0).flatten() + y_pred = (y_pred_prob > 0.5).astype(int) + + # Calculate metrics + accuracy = (y_pred == y).mean() + + # Confusion matrix + cm = confusion_matrix(y, y_pred) + tn, fp, fn, tp = cm.ravel() + + precision = tp / (tp + fp) if (tp + fp) > 0 else 0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0 + f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 + + # Specificity (True Negative Rate) + specificity = tn / (tn + fp) if (tn + fp) > 0 else 0 + + # ROC AUC + try: + roc_auc = roc_auc_score(y, y_pred_prob) + fpr, tpr, thresholds = roc_curve(y, y_pred_prob) + except: + roc_auc = 0.0 + fpr, tpr, thresholds = None, None, None + + # Precision-Recall AUC + try: + precision_curve, recall_curve, _ = precision_recall_curve(y, y_pred_prob) + pr_auc = auc(recall_curve, precision_curve) + except: + pr_auc = 0.0 + precision_curve, recall_curve = None, None + + # Print results + print(f"\nMetrics for {dataset_name}:") + print(f"Accuracy: {accuracy:.4f}") + print(f"Precision: {precision:.4f}") + print(f"Recall (Sensitivity): {recall:.4f}") + print(f"Specificity: {specificity:.4f}") + print(f"F1 Score: {f1:.4f}") + print(f"ROC AUC: {roc_auc:.4f}") + print(f"PR AUC: {pr_auc:.4f}") + + print(f"\nConfusion Matrix:") + print(f" Predicted") + print(f" Neg Pos") + print(f"Actual Neg {tn:5d} {fp:5d}") + print(f" Pos {fn:5d} {tp:5d}") + + print(f"\nClassification Report:") + print(classification_report(y, y_pred)) + + # Plot ROC curve + if fpr is not None and tpr is not None: + plt.figure(figsize=(10, 5)) + + # ROC Curve + plt.subplot(1, 2, 1) + plt.plot(fpr, tpr, label=f'ROC (AUC = {roc_auc:.3f})', linewidth=2) + plt.plot([0, 1], [0, 1], 'k--', label='Random', linewidth=1) + plt.xlabel('False Positive Rate') + plt.ylabel('True Positive Rate') + plt.title(f'ROC Curve - {dataset_name}') + plt.legend() + plt.grid(True, alpha=0.3) + + # Precision-Recall Curve + plt.subplot(1, 2, 2) + if precision_curve is not None and recall_curve is not None: + plt.plot(recall_curve, precision_curve, label=f'PR (AUC = {pr_auc:.3f})', linewidth=2) + plt.xlabel('Recall') + plt.ylabel('Precision') + plt.title(f'Precision-Recall Curve - {dataset_name}') + plt.legend() + plt.grid(True, alpha=0.3) + + plt.tight_layout() + plt.savefig( + os.path.join(output_dir, f"{dataset_name.replace(' ', '_')}_curves.png"), + dpi=300, bbox_inches="tight" + ) + plt.close() + + # Save results to file + results_path = os.path.join(output_dir, f"{dataset_name.replace(' ', '_')}_test_results.txt") + with open(results_path, "w") as f: + f.write(f"Test Results for {dataset_name}\n") + f.write("=" * 60 + "\n\n") + f.write(f"Sample size: {len(y)}\n") + f.write(f"Class distribution: {np.bincount(y.astype(int))}\n\n") + f.write(f"Accuracy: {accuracy:.4f}\n") + f.write(f"Precision: {precision:.4f}\n") + f.write(f"Recall (Sensitivity): {recall:.4f}\n") + f.write(f"Specificity: {specificity:.4f}\n") + f.write(f"F1 Score: {f1:.4f}\n") + f.write(f"ROC AUC: {roc_auc:.4f}\n") + f.write(f"PR AUC: {pr_auc:.4f}\n\n") + f.write(f"Confusion Matrix:\n") + f.write(f" Predicted\n") + f.write(f" Neg Pos\n") + f.write(f"Actual Neg {tn:5d} {fp:5d}\n") + f.write(f" Pos {fn:5d} {tp:5d}\n\n") + f.write(f"Classification Report:\n") + f.write(classification_report(y, y_pred)) + + print(f"\nResults saved to: {results_path}") + + # Return metrics + return { + "accuracy": accuracy, + "precision": precision, + "recall": recall, + "specificity": specificity, + "f1": f1, + "roc_auc": roc_auc, + "pr_auc": pr_auc, + "confusion_matrix": cm, + "n_samples": len(y), + "n_positive": int(y.sum()), + "n_negative": int((1 - y).sum()) + } + + +def main(): + parser = argparse.ArgumentParser( + description="Test control model on test data" + ) + + parser.add_argument( + "--model_path", + type=str, + required=True, + help="Path to trained model (.h5 file)", + ) + parser.add_argument( + "--test_data", + type=str, + default="./data/target_test.npz", + help="Path to test data (.npz file)", + ) + parser.add_argument( + "--dataset_name", + type=str, + default="Test Data", + help="Name of the test dataset (for labeling, default: 'Test Data')", + ) + parser.add_argument( + "--output_dir", + type=str, + default="./control_test_results_snr", + help="Directory to save test results (default: ./control_test_results)", + ) + parser.add_argument( + "--gpu", + type=int, + default=None, + help="GPU device ID to use (default: auto-select)", + ) + + args = parser.parse_args() + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Set GPU device if specified + if args.gpu is not None: + physical_devices = tf.config.experimental.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_visible_devices( + physical_devices[args.gpu], "GPU" + ) + tf.config.experimental.set_memory_growth(physical_devices[args.gpu], True) + + print("=" * 80) + print("CONTROL MODEL TESTING") + print("=" * 80) + print(f"Model: {args.model_path}") + print(f"Test data: {args.test_data}") + print(f"Dataset name: {args.dataset_name}") + print(f"Output directory: {args.output_dir}") + print("=" * 80) + + # Load model + print("\nLoading model...") + model = tf.keras.models.load_model( + args.model_path, + custom_objects={ + "rot90_k1": rot90_k1, + "rot90_k2": rot90_k2, + "rot90_k3": rot90_k3 + } + ) + print("Model loaded successfully!") + model.summary() + + # Load test data + print(f"\n--- Loading Test Data ({args.dataset_name}) ---") + X, feats, y, metadata = load_data(args.test_data) + + # Evaluate model + metrics = evaluate_model( + model, X, feats, y, + args.dataset_name, args.output_dir + ) + + # Summary + print("\n" + "=" * 80) + print("SUMMARY") + print("=" * 80) + print(f"Dataset: {args.dataset_name}") + print(f"Samples: {metrics['n_samples']} (Positive: {metrics['n_positive']}, Negative: {metrics['n_negative']})") + print(f"Accuracy: {metrics['accuracy']:.4f}") + print(f"Precision: {metrics['precision']:.4f}") + print(f"Recall: {metrics['recall']:.4f}") + print(f"Specificity: {metrics['specificity']:.4f}") + print(f"F1 Score: {metrics['f1']:.4f}") + print(f"ROC AUC: {metrics['roc_auc']:.4f}") + print(f"PR AUC: {metrics['pr_auc']:.4f}") + print("=" * 80) + + # Save summary + summary_path = os.path.join(args.output_dir, "summary.txt") + with open(summary_path, "w") as f: + f.write("CONTROL MODEL - TEST SUMMARY\n") + f.write("=" * 80 + "\n\n") + f.write(f"Model: {args.model_path}\n") + f.write(f"Test data: {args.test_data}\n\n") + f.write(f"Dataset: {args.dataset_name}\n") + f.write(f"Samples: {metrics['n_samples']} ") + f.write(f"(Positive: {metrics['n_positive']}, Negative: {metrics['n_negative']})\n\n") + f.write("Performance Metrics:\n") + f.write(f" Accuracy: {metrics['accuracy']:.4f}\n") + f.write(f" Precision: {metrics['precision']:.4f}\n") + f.write(f" Recall: {metrics['recall']:.4f}\n") + f.write(f" Specificity: {metrics['specificity']:.4f}\n") + f.write(f" F1 Score: {metrics['f1']:.4f}\n") + f.write(f" ROC AUC: {metrics['roc_auc']:.4f}\n") + f.write(f" PR AUC: {metrics['pr_auc']:.4f}\n") + + print(f"\nSummary saved to: {summary_path}") + print(f"\nTesting completed successfully!") + + +if __name__ == "__main__": + main() diff --git a/RuBR/experiments/domain_adaption/test_dann.py b/RuBR/experiments/domain_adaption/test_dann.py new file mode 100644 index 00000000..bec8277b --- /dev/null +++ b/RuBR/experiments/domain_adaption/test_dann.py @@ -0,0 +1,901 @@ +""" +Test Domain Adversarial model on both source and target domains. + +This script evaluates a trained DANN model on test data from both domains +and provides comprehensive metrics. +""" + +import argparse +import os +import numpy as np +import pandas as pd +import tensorflow as tf +from sklearn.metrics import ( + confusion_matrix, + classification_report, + roc_auc_score, + precision_recall_curve, + roc_curve, + auc +) +from matplotlib import pyplot as plt +import scienceplots +from model.data import load_dataset, normalize_arrays +from model.metadata import get_transient_magnitude + +# Import model components from flat package module +from model.dann_model import ( + rot90_k1, + rot90_k2, + rot90_k3, + GradientReversalLayer +) +def load_data(data_path): + return load_dataset(data_path, mmap=False, allow_npy_dict=False) + +def evaluate_domain(model, X, feats, y, domain_name, output_dir, metadata=None, threshold=0.5): + """ + Evaluate model on a specific domain. + + Args: + model: Trained DANN model + X: Image data + feats: Tabular features + y: Labels + domain_name: Name of the domain (for logging) + output_dir: Directory to save results + metadata: Optional metadata containing magnitude information + threshold: Classification threshold for histogram (default: 0.5) + + Returns: + dict: Dictionary of evaluation metrics + """ + print(f"\n{'='*60}") + print(f"Evaluating on {domain_name} Domain") + print(f"{'='*60}") + + # Remove NaN values + mask = np.isnan(X).any(axis=(1, 2, 3)) | np.isnan(feats).any(axis=1) + if mask.any(): + print(f"Removing {mask.sum()} samples with NaN values") + X = X[~mask] + feats = feats[~mask] + y = y[~mask] + if metadata is not None: + metadata = metadata[~mask] + + # Normalize data (using same approach as training) + X, feats = normalize_arrays(X, feats) + + # Get predictions (only label output matters for testing) + predictions = model.predict([X, feats], verbose=0) + + # Handle different output formats + if isinstance(predictions, list): + y_pred_prob = predictions[0] # Label predictions + y_pred_domain = predictions[1] # Domain predictions + else: + y_pred_prob = predictions + y_pred_domain = None + + y_pred_prob = y_pred_prob.flatten() + y_pred = (y_pred_prob > 0.5).astype(int) + + # Calculate metrics + accuracy = (y_pred == y).mean() + + # Confusion matrix + cm = confusion_matrix(y, y_pred) + tn, fp, fn, tp = cm.ravel() + + precision = tp / (tp + fp) if (tp + fp) > 0 else 0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0 + f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 + + # ROC AUC + try: + roc_auc = roc_auc_score(y, y_pred_prob) + fpr, tpr, _ = roc_curve(y, y_pred_prob) + except: + roc_auc = 0.0 + fpr, tpr = None, None + + # Precision-Recall AUC + try: + precision_curve, recall_curve, _ = precision_recall_curve(y, y_pred_prob) + pr_auc = auc(recall_curve, precision_curve) + except: + pr_auc = 0.0 + precision_curve, recall_curve = None, None + + # Find optimal F1 threshold + print(f"\nFinding optimal F1 threshold...") + thresholds_f1 = np.linspace(0, 1, 1000) + f1_scores = [] + precision_scores = [] + + for thresh in thresholds_f1: + y_pred_thresh = (y_pred_prob > thresh).astype(int) + cm_thresh = confusion_matrix(y, y_pred_thresh, labels=[0, 1]) + + if cm_thresh.size == 4: + tn_t, fp_t, fn_t, tp_t = cm_thresh.ravel() + elif cm_thresh.size == 1: + if y_pred_thresh.sum() == 0: + tn_t, fp_t, fn_t, tp_t = cm_thresh[0, 0], 0, y.sum(), 0 + else: + tn_t, fp_t, fn_t, tp_t = 0, (y == 0).sum(), 0, y.sum() + else: + tn_t, fp_t, fn_t, tp_t = 0, 0, 0, 0 + + prec_t = tp_t / (tp_t + fp_t) if (tp_t + fp_t) > 0 else 0 + rec_t = tp_t / (tp_t + fn_t) if (tp_t + fn_t) > 0 else 0 + f1_t = 2 * prec_t * rec_t / (prec_t + rec_t) if (prec_t + rec_t) > 0 else 0 + f1_scores.append(f1_t) + precision_scores.append(prec_t) + + f1_scores = np.array(f1_scores) + precision_scores = np.array(precision_scores) + + # Find optimal F1 threshold + best_f1_idx = np.argmax(f1_scores) + best_f1_threshold = thresholds_f1[best_f1_idx] + best_f1_score = f1_scores[best_f1_idx] + + # Calculate metrics at optimal F1 threshold + y_pred_best_f1 = (y_pred_prob > best_f1_threshold).astype(int) + cm_best_f1 = confusion_matrix(y, y_pred_best_f1) + tn_best, fp_best, fn_best, tp_best = cm_best_f1.ravel() + + accuracy_best_f1 = (y_pred_best_f1 == y).mean() + precision_best_f1 = tp_best / (tp_best + fp_best) if (tp_best + fp_best) > 0 else 0 + recall_best_f1 = tp_best / (tp_best + fn_best) if (tp_best + fn_best) > 0 else 0 + + # Find threshold where precision is closest to 60% + print(f"\nFinding threshold for 60% precision...") + target_precision = 0.60 + prec_60_idx = np.argmin(np.abs(precision_scores - target_precision)) + prec_60_threshold = thresholds_f1[prec_60_idx] + + # Calculate metrics at 60% precision threshold + y_pred_prec_60 = (y_pred_prob > prec_60_threshold).astype(int) + cm_prec_60 = confusion_matrix(y, y_pred_prec_60) + tn_p60, fp_p60, fn_p60, tp_p60 = cm_prec_60.ravel() + + accuracy_prec_60 = (y_pred_prec_60 == y).mean() + precision_prec_60 = tp_p60 / (tp_p60 + fp_p60) if (tp_p60 + fp_p60) > 0 else 0 + recall_prec_60 = tp_p60 / (tp_p60 + fn_p60) if (tp_p60 + fn_p60) > 0 else 0 + f1_prec_60 = 2 * precision_prec_60 * recall_prec_60 / (precision_prec_60 + recall_prec_60) if (precision_prec_60 + recall_prec_60) > 0 else 0 + + # Print results at default threshold (0.5) + print(f"\nMetrics for {domain_name} (at threshold 0.5):") + print(f"Accuracy: {accuracy:.4f}") + print(f"Precision: {precision:.4f}") + print(f"Recall: {recall:.4f}") + print(f"F1 Score: {f1:.4f}") + print(f"ROC AUC: {roc_auc:.4f}") + print(f"PR AUC: {pr_auc:.4f}") + + print(f"\nConfusion Matrix (at threshold 0.5):") + print(cm) + + # Print results at optimal F1 threshold + print(f"\n{'='*60}") + print(f"Metrics at Optimal F1 Threshold ({best_f1_threshold:.4f}):") + print(f"{'='*60}") + print(f"F1 Score: {best_f1_score:.4f}") + print(f"Accuracy: {accuracy_best_f1:.4f}") + print(f"Precision: {precision_best_f1:.4f}") + print(f"Recall: {recall_best_f1:.4f}") + print(f"\nConfusion Matrix (at threshold {best_f1_threshold:.4f}):") + print(cm_best_f1) + + # Print results at 60% precision threshold + print(f"\n{'='*60}") + print(f"Metrics at 60% Precision Threshold ({prec_60_threshold:.4f}):") + print(f"{'='*60}") + print(f"Precision: {precision_prec_60:.4f} (target: 0.60)") + print(f"Recall: {recall_prec_60:.4f}") + print(f"F1 Score: {f1_prec_60:.4f}") + print(f"Accuracy: {accuracy_prec_60:.4f}") + print(f"\nConfusion Matrix (at threshold {prec_60_threshold:.4f}):") + print(cm_prec_60) + + print(f"\nClassification Report (at threshold 0.5):") + print(classification_report(y, y_pred)) + + # Plot ROC and Precision-Recall curves + if (fpr is not None and tpr is not None) or (precision_curve is not None and recall_curve is not None): + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + + # ROC Curve + if fpr is not None and tpr is not None: + axes[0].plot(fpr, tpr, linewidth=2, label=f'ROC (AUC = {roc_auc:.3f})') + axes[0].plot([0, 1], [0, 1], 'k--', linewidth=1, label='Random') + axes[0].set_xlabel('False Positive Rate', fontsize=12) + axes[0].set_ylabel('True Positive Rate', fontsize=12) + axes[0].set_title(f'ROC Curve - {domain_name} Domain', fontsize=13) + axes[0].legend(fontsize=10) + axes[0].grid(True, alpha=0.3) + + # Precision-Recall Curve + if precision_curve is not None and recall_curve is not None: + axes[1].plot(recall_curve, precision_curve, linewidth=2, label=f'PR (AUC = {pr_auc:.3f})') + axes[1].axhline(y=y.mean(), color='k', linestyle='--', linewidth=1, label=f'Baseline ({y.mean():.3f})') + axes[1].set_xlabel('Recall', fontsize=12) + axes[1].set_ylabel('Precision', fontsize=12) + axes[1].set_title(f'Precision-Recall Curve - {domain_name} Domain', fontsize=13) + axes[1].legend(fontsize=10) + axes[1].grid(True, alpha=0.3) + + plt.tight_layout() + curves_path = os.path.join(output_dir, f"{domain_name}_curves.png") + plt.savefig(curves_path, dpi=300, bbox_inches="tight") + plt.close() + print(f"ROC and PR curves saved to: {curves_path}") + + # Plot Precision vs Threshold and Recall vs Threshold + print(f"\nGenerating Precision and Recall vs Threshold plots...") + thresholds = np.linspace(0, 1, 100) + precisions = [] + recalls = [] + + for thresh in thresholds: + y_pred_thresh = (y_pred_prob > thresh).astype(int) + cm_thresh = confusion_matrix(y, y_pred_thresh, labels=[0, 1]) + + if cm_thresh.size == 4: # Full confusion matrix + tn_t, fp_t, fn_t, tp_t = cm_thresh.ravel() + elif cm_thresh.size == 1: # Edge case: only one class predicted + if y_pred_thresh.sum() == 0: # All predicted as 0 + tn_t, fp_t, fn_t, tp_t = cm_thresh[0, 0], 0, y.sum(), 0 + else: # All predicted as 1 + tn_t, fp_t, fn_t, tp_t = 0, (y == 0).sum(), 0, y.sum() + else: + # Handle other edge cases + tn_t, fp_t, fn_t, tp_t = 0, 0, 0, 0 + + prec_t = tp_t / (tp_t + fp_t) if (tp_t + fp_t) > 0 else 0 + rec_t = tp_t / (tp_t + fn_t) if (tp_t + fn_t) > 0 else 0 + + precisions.append(prec_t) + recalls.append(rec_t) + + precisions = np.array(precisions) + recalls = np.array(recalls) + + # Create plot + fig, axes = plt.subplots(1, 2, figsize=(14, 5)) + + # Precision vs Threshold + axes[0].plot(thresholds, precisions, linewidth=2, color='blue', label='Precision') + axes[0].axvline(x=0.5, color='red', linestyle='--', linewidth=1, label='Default Threshold (0.5)') + axes[0].set_xlabel('Threshold', fontsize=12) + axes[0].set_ylabel('Precision', fontsize=12) + axes[0].set_title(f'Precision vs Threshold - {domain_name} Domain', fontsize=13) + axes[0].legend(fontsize=10) + axes[0].grid(True, alpha=0.3) + axes[0].set_xlim([0, 1]) + axes[0].set_ylim([0, 1.05]) + + # Recall vs Threshold + axes[1].plot(thresholds, recalls, linewidth=2, color='green', label='Recall') + axes[1].axvline(x=0.5, color='red', linestyle='--', linewidth=1, label='Default Threshold (0.5)') + axes[1].set_xlabel('Threshold', fontsize=12) + axes[1].set_ylabel('Recall', fontsize=12) + axes[1].set_title(f'Recall vs Threshold - {domain_name} Domain', fontsize=13) + axes[1].legend(fontsize=10) + axes[1].grid(True, alpha=0.3) + axes[1].set_xlim([0, 1]) + axes[1].set_ylim([0, 1.05]) + + plt.tight_layout() + threshold_curves_path = os.path.join(output_dir, f"{domain_name}_precision_recall_vs_threshold.png") + plt.savefig(threshold_curves_path, dpi=300, bbox_inches="tight") + plt.close() + print(f"Precision and Recall vs Threshold plots saved to: {threshold_curves_path}") + + # Save results to file + results_path = os.path.join(output_dir, f"{domain_name}_test_results.txt") + with open(results_path, "w") as f: + f.write(f"Test Results for {domain_name} Domain\n") + f.write("=" * 60 + "\n\n") + f.write("Metrics at Default Threshold (0.5):\n") + f.write("-" * 60 + "\n") + f.write(f"Accuracy: {accuracy:.4f}\n") + f.write(f"Precision: {precision:.4f}\n") + f.write(f"Recall: {recall:.4f}\n") + f.write(f"F1 Score: {f1:.4f}\n") + f.write(f"ROC AUC: {roc_auc:.4f}\n") + f.write(f"PR AUC: {pr_auc:.4f}\n\n") + f.write(f"Confusion Matrix:\n{cm}\n\n") + f.write(f"Classification Report:\n") + f.write(classification_report(y, y_pred)) + f.write("\n" + "=" * 60 + "\n\n") + f.write(f"Metrics at Optimal F1 Threshold ({best_f1_threshold:.4f}):\n") + f.write("-" * 60 + "\n") + f.write(f"F1 Score: {best_f1_score:.4f}\n") + f.write(f"Accuracy: {accuracy_best_f1:.4f}\n") + f.write(f"Precision: {precision_best_f1:.4f}\n") + f.write(f"Recall: {recall_best_f1:.4f}\n\n") + f.write(f"Confusion Matrix:\n{cm_best_f1}\n\n") + f.write(f"Classification Report:\n") + f.write(classification_report(y, y_pred_best_f1)) + f.write("\n" + "=" * 60 + "\n\n") + f.write(f"Metrics at 60% Precision Threshold ({prec_60_threshold:.4f}):\n") + f.write("-" * 60 + "\n") + f.write(f"Precision: {precision_prec_60:.4f} (target: 0.60)\n") + f.write(f"Recall: {recall_prec_60:.4f}\n") + f.write(f"F1 Score: {f1_prec_60:.4f}\n") + f.write(f"Accuracy: {accuracy_prec_60:.4f}\n\n") + f.write(f"Confusion Matrix:\n{cm_prec_60}\n\n") + f.write(f"Classification Report:\n") + f.write(classification_report(y, y_pred_prec_60)) + + print(f"\nResults saved to: {results_path}") + + # Save misclassification examples + false_positives = (y == 0) & (y_pred == 1) + false_negatives = (y == 1) & (y_pred == 0) + + n_fp = false_positives.sum() + n_fn = false_negatives.sum() + + print(f"\nMisclassifications:") + print(f" False Positives: {n_fp}") + print(f" False Negatives: {n_fn}") + + # Save examples of misclassifications (up to 10 of each type) + n_examples = 10 + + if n_fp > 0: + fp_indices = np.where(false_positives)[0] + n_show_fp = min(n_examples, len(fp_indices)) + selected_fp = np.random.choice(fp_indices, n_show_fp, replace=False) + + # Create figure for false positives - show 3 channels separately + fig, axes = plt.subplots(n_show_fp, 3, figsize=(12, 3 * n_show_fp)) + fig.suptitle(f'{domain_name} Domain - False Positives (Predicted Real, Actually Bogus)', fontsize=14) + + if n_show_fp == 1: + axes = axes.reshape(1, -1) + + channel_names = ['Science', 'Reference', 'Difference'] + + for idx, sample_idx in enumerate(selected_fp): + img = X[sample_idx] + + # Plot each channel separately + for ch in range(3): + axes[idx, ch].imshow(img[:, :, ch], cmap='gray', interpolation='nearest') + if idx == 0: + axes[idx, ch].set_title(f'{channel_names[ch]}', fontsize=11) + axes[idx, ch].axis('off') + + # Add prediction info on the left + axes[idx, 0].text(-0.15, 0.5, f'Pred: {y_pred_prob[sample_idx]:.3f}\nTrue: {y[sample_idx]}', + transform=axes[idx, 0].transAxes, fontsize=10, + verticalalignment='center', horizontalalignment='right') + + plt.tight_layout() + fp_path = os.path.join(output_dir, f"{domain_name}_false_positives.png") + plt.savefig(fp_path, dpi=200, bbox_inches="tight") + plt.close() + print(f"False positive examples saved to: {fp_path}") + + if n_fn > 0: + fn_indices = np.where(false_negatives)[0] + n_show_fn = min(n_examples, len(fn_indices)) + selected_fn = np.random.choice(fn_indices, n_show_fn, replace=False) + + # Create figure for false negatives - show 3 channels separately + fig, axes = plt.subplots(n_show_fn, 3, figsize=(12, 3 * n_show_fn)) + fig.suptitle(f'{domain_name} Domain - False Negatives (Predicted Bogus, Actually Real)', fontsize=14) + + if n_show_fn == 1: + axes = axes.reshape(1, -1) + + channel_names = ['Science', 'Reference', 'Difference'] + + for idx, sample_idx in enumerate(selected_fn): + img = X[sample_idx] + + # Plot each channel separately + for ch in range(3): + axes[idx, ch].imshow(img[:, :, ch], cmap='gray', interpolation='nearest') + if idx == 0: + axes[idx, ch].set_title(f'{channel_names[ch]}', fontsize=11) + axes[idx, ch].axis('off') + + # Add prediction info on the left + axes[idx, 0].text(-0.15, 0.5, f'Pred: {y_pred_prob[sample_idx]:.3f}\nTrue: {y[sample_idx]}', + transform=axes[idx, 0].transAxes, fontsize=10, + verticalalignment='center', horizontalalignment='right') + + plt.tight_layout() + fn_path = os.path.join(output_dir, f"{domain_name}_false_negatives.png") + plt.savefig(fn_path, dpi=200, bbox_inches="tight") + plt.close() + print(f"False negative examples saved to: {fn_path}") + + + # Create magnitude histogram if metadata is available (only for Target domain) + if domain_name == "Target": + try: + # Apply scienceplots style for publication + plt.style.use(['science', 'ieee', 'no-latex']) + plt.rcParams.update({ + 'font.size': 10, + 'font.family': 'serif', + 'axes.labelsize': 11, + 'axes.titlesize': 11, + 'xtick.labelsize': 9, + 'ytick.labelsize': 9, + 'legend.fontsize': 9, + 'lines.linewidth': 1.5, + 'grid.linewidth': 0.5, + 'axes.linewidth': 0.8, + }) + + print(f"\nUsing threshold {threshold:.2f} for magnitude histogram") + + # Extract magnitude values using get_transient_magnitude function + # ONLY call the function for samples where y == 1 (ground truth positives) + mag_values = [] + pred_probs = [] # Store corresponding prediction probabilities + + for idx, m in enumerate(metadata): + # CRITICAL: Skip if not a ground truth positive (y != 1) + # get_transient_magnitude will fail for false positives + if y[idx] != 1: + continue + + try: + if isinstance(m, dict) and 'match_id' in m and 'jid_folder' in m: + mag = get_transient_magnitude(m['match_id'], m['jid_folder']) + # Only add if magnitude is valid (not NaN) + if not np.isnan(mag): + mag_values.append(mag) + pred_probs.append(y_pred_prob[idx]) + else: + print(f"Skipping sample {idx}: magnitude is NaN") + print(f"ID: {m['match_id']}, Folder: {m['jid_folder']}") + else: + print(f"Skipping sample {idx}: missing match_id or jid_folder in metadata") + except Exception as e: + print(f"Error getting magnitude for sample {idx}: {e}") + + if len(mag_values) == 0: + print("No valid magnitude values found. Skipping histogram.") + return { + "accuracy": accuracy, + "precision": precision, + "recall": recall, + "f1": f1, + "roc_auc": roc_auc, + "pr_auc": pr_auc, + "confusion_matrix": cm, + "best_f1_threshold": best_f1_threshold, + "best_f1_score": best_f1_score, + "best_f1_accuracy": accuracy_best_f1, + "best_f1_precision": precision_best_f1, + "best_f1_recall": recall_best_f1, + "best_f1_confusion_matrix": cm_best_f1, + "prec_60_threshold": prec_60_threshold, + "prec_60_precision": precision_prec_60, + "prec_60_recall": recall_prec_60, + "prec_60_f1": f1_prec_60, + "prec_60_accuracy": accuracy_prec_60, + "prec_60_confusion_matrix": cm_prec_60 + } + + mag_values = np.array(mag_values) + pred_probs = np.array(pred_probs) + + # All samples here are ground truth positives (y==1) + # Separate by model prediction: true positives vs false negatives + true_positive_mask = pred_probs > threshold + + mag_all_positives = mag_values # All ground truth positives + mag_true_positives = mag_values[true_positive_mask] # Correctly detected + + # Create single plot with both histograms overlaid + fig, ax = plt.subplots(figsize=(3.5, 2.8), dpi=600) + + # Use grayscale-friendly colors with different line styles + colors = { + 'all': '#000000', # Black + 'tp': 'blue', # Dark gray + } + + # Use common bins based on all positives range + bins = np.linspace(mag_all_positives.min(), mag_all_positives.max(), 30) + + # Plot all ground truth positives + ax.hist(mag_all_positives, + bins=bins, + label="Ground Truth Positives", + color=colors['all'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=2) + + # Plot true positives if available + if len(mag_true_positives) > 0: + ax.hist(mag_true_positives, + bins=bins, + label=f"True Positives", + color=colors['tp'], + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=3) + + # Set log scale + ax.set_yscale("log") + + # Labels with proper units + ax.set_xlabel("Magnitude (mag)", fontsize=10) + ax.set_ylabel("Count", fontsize=10) + + # Concise title + ax.set_title(f"{domain_name} Domain - Magnitude Distribution", fontsize=10, pad=8) + + # Legend with academic styling + ax.legend(loc='best', frameon=True, + fontsize=8, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.5, labelspacing=0.3) + + # Minimal grid on y-axis only (common for log-scale histograms) + ax.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax.set_axisbelow(True) + + # Clean white background + ax.set_facecolor('white') + fig.patch.set_facecolor('white') + + # Standard spine styling + for spine in ax.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + # Tight layout + plt.tight_layout(pad=0.3) + + # Save with high DPI for publication + histogram_plot_path = os.path.join(output_dir, f"{domain_name}_magnitude_histogram.png") + + # Save both PNG and PDF + plt.savefig(histogram_plot_path, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(histogram_plot_path.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + + plt.close() + + print(f"Magnitude histogram saved to: {histogram_plot_path}") + print(f"Magnitude histogram PDF saved to: {histogram_plot_path.replace('.png', '.pdf')}") + + # Print summary statistics + print(f"\nMagnitude Summary Statistics:") + print(f"All Ground Truth Positives: {len(mag_all_positives)}") + if len(mag_true_positives) > 0: + print(f"True Positives: {len(mag_true_positives)}") + print(f"Detection rate: {len(mag_true_positives)/len(mag_all_positives)*100:.1f}%") + print(f"True positives magnitude range: [{mag_true_positives.min():.2f}, {mag_true_positives.max():.2f}]") + print(f"All positives magnitude range: [{mag_all_positives.min():.2f}, {mag_all_positives.max():.2f}]") + + # Create separate plot for false negatives + false_negative_mask = pred_probs <= threshold + mag_false_negatives = mag_values[false_negative_mask] + + if len(mag_false_negatives) > 0: + print(f"\nCreating False Negatives magnitude histogram...") + print(f"False Negatives: {len(mag_false_negatives)}") + + fig_fn, ax_fn = plt.subplots(figsize=(3.5, 2.8), dpi=600) + + # Use red for false negatives to indicate missed detections + bins_fn = np.linspace(mag_false_negatives.min(), mag_false_negatives.max(), 30) + + ax_fn.hist(mag_false_negatives, + bins=bins_fn, + label=f"False Negatives (N={len(mag_false_negatives)})", + color='red', + histtype="step", + linewidth=1.5, + linestyle='-', + zorder=2) + + # Set log scale + ax_fn.set_yscale("log") + + # Labels + ax_fn.set_xlabel("Magnitude (mag)", fontsize=10) + ax_fn.set_ylabel("Count", fontsize=10) + ax_fn.set_title(f"{domain_name} Domain - False Negatives Magnitude Distribution", fontsize=10, pad=8) + + # Legend + ax_fn.legend(loc='best', frameon=True, + fontsize=8, framealpha=1, + edgecolor='black', fancybox=False, + borderpad=0.5, labelspacing=0.3) + + # Grid + ax_fn.grid(True, alpha=0.2, linestyle='-', linewidth=0.3, + color='gray', zorder=0, axis='y') + ax_fn.set_axisbelow(True) + + # Styling + ax_fn.set_facecolor('white') + fig_fn.patch.set_facecolor('white') + + for spine in ax_fn.spines.values(): + spine.set_edgecolor('black') + spine.set_linewidth(0.8) + + plt.tight_layout(pad=0.3) + + # Save + fn_histogram_path = os.path.join(output_dir, f"{domain_name}_false_negatives_magnitude_histogram.png") + plt.savefig(fn_histogram_path, dpi=600, bbox_inches='tight', + facecolor='white', edgecolor='none') + plt.savefig(fn_histogram_path.replace('.png', '.pdf'), + bbox_inches='tight', facecolor='white', edgecolor='none') + plt.close() + + print(f"False negatives magnitude histogram saved to: {fn_histogram_path}") + print(f"False negatives magnitude range: [{mag_false_negatives.min():.2f}, {mag_false_negatives.max():.2f}]") + + except Exception as e: + print(f"Could not create magnitude histogram: {e}") + + # Return metrics (include default, optimal F1, and 60% precision metrics) + return { + "accuracy": accuracy, + "precision": precision, + "recall": recall, + "f1": f1, + "roc_auc": roc_auc, + "pr_auc": pr_auc, + "confusion_matrix": cm, + "best_f1_threshold": best_f1_threshold, + "best_f1_score": best_f1_score, + "best_f1_accuracy": accuracy_best_f1, + "best_f1_precision": precision_best_f1, + "best_f1_recall": recall_best_f1, + "best_f1_confusion_matrix": cm_best_f1, + "prec_60_threshold": prec_60_threshold, + "prec_60_precision": precision_prec_60, + "prec_60_recall": recall_prec_60, + "prec_60_f1": f1_prec_60, + "prec_60_accuracy": accuracy_prec_60, + "prec_60_confusion_matrix": cm_prec_60 + } + + +def plot_comparison(source_metrics, target_metrics, output_dir): + """Plot comparison between source and target domain performance.""" + metrics = ["accuracy", "precision", "recall", "f1", "roc_auc", "pr_auc"] + source_values = [source_metrics[m] for m in metrics] + target_values = [target_metrics[m] for m in metrics] + + x = np.arange(len(metrics)) + width = 0.35 + + fig, ax = plt.subplots(figsize=(12, 6)) + bars1 = ax.bar(x - width/2, source_values, width, label="Source Domain", alpha=0.8) + bars2 = ax.bar(x + width/2, target_values, width, label="Target Domain", alpha=0.8) + + ax.set_xlabel("Metrics") + ax.set_ylabel("Score") + ax.set_title("Model Performance: Source vs Target Domain") + ax.set_xticks(x) + ax.set_xticklabels([m.replace("_", " ").title() for m in metrics]) + ax.legend() + ax.grid(True, alpha=0.3) + + # Add value labels on bars + for bars in [bars1, bars2]: + for bar in bars: + height = bar.get_height() + ax.annotate(f'{height:.3f}', + xy=(bar.get_x() + bar.get_width() / 2, height), + xytext=(0, 3), + textcoords="offset points", + ha='center', va='bottom', + fontsize=8) + + plt.tight_layout() + plt.savefig(os.path.join(output_dir, "domain_comparison.png"), dpi=300, bbox_inches="tight") + plt.close() + + print(f"\nComparison plot saved to: {os.path.join(output_dir, 'domain_comparison.png')}") + + +def main(): + parser = argparse.ArgumentParser( + description="Test Domain Adversarial model" + ) + + parser.add_argument( + "--model_path", + type=str, + default="./outputs/dann_train/best_model.h5", + # required=True, + help="Path to trained model (.h5 file)", + ) + parser.add_argument( + "--source_test", + type=str, + default="./data/source_test.npz", + # required=True, + help="Path to source domain test data (.npz file)", + ) + parser.add_argument( + "--target_test", + type=str, + default="./data/target_test.npz", + # required=True, + help="Path to target domain test data (.npz file)", + ) + parser.add_argument( + "--output_dir", + type=str, + default="./dann_test_results_tp_only", + help="Directory to save test results (default: ./dann_test_results)", + ) + parser.add_argument( + "--gpu", + type=int, + default=None, + help="GPU device ID to use (default: auto-select)", + ) + + args = parser.parse_args() + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Set GPU device if specified + if args.gpu is not None: + physical_devices = tf.config.experimental.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_visible_devices( + physical_devices[args.gpu], "GPU" + ) + tf.config.experimental.set_memory_growth(physical_devices[args.gpu], True) + + print("=" * 80) + print("DOMAIN ADVERSARIAL MODEL TESTING") + print("=" * 80) + print(f"Model: {args.model_path}") + print(f"Source test data: {args.source_test}") + print(f"Target test data: {args.target_test}") + print(f"Output directory: {args.output_dir}") + print("=" * 80) + + # Load model with custom objects + print("\nLoading model...") + model = tf.keras.models.load_model( + args.model_path, + custom_objects={ + "GradientReversalLayer": GradientReversalLayer, + "rot90_k1": rot90_k1, + "rot90_k2": rot90_k2, + "rot90_k3": rot90_k3 + } + ) + print("Model loaded successfully!") + + # Load test data + print("\n--- Loading Source Domain Test Data ---") + X_src, feats_src, y_src, metadata_src = load_data(args.source_test) + + print("\n--- Loading Target Domain Test Data ---") + X_tgt, feats_tgt, y_tgt, metadata_tgt = load_data(args.target_test) + + # Evaluate on source domain + source_metrics = evaluate_domain( + model, X_src, feats_src, y_src, + "Source", args.output_dir, metadata_src + ) + + # Evaluate on target domain + target_metrics = evaluate_domain( + model, X_tgt, feats_tgt, y_tgt, + "Target", args.output_dir, metadata_tgt + ) + + # Plot comparison + plot_comparison(source_metrics, target_metrics, args.output_dir) + + # Summary + print("\n" + "=" * 80) + print("SUMMARY") + print("=" * 80) + print("\nDefault Threshold (0.5):") + print(f" Source Domain F1: {source_metrics['f1']:.4f}") + print(f" Target Domain F1: {target_metrics['f1']:.4f}") + print(f" F1 Difference: {abs(source_metrics['f1'] - target_metrics['f1']):.4f}") + print(f" Source Domain ROC AUC: {source_metrics['roc_auc']:.4f}") + print(f" Target Domain ROC AUC: {target_metrics['roc_auc']:.4f}") + print(f" ROC AUC Difference: {abs(source_metrics['roc_auc'] - target_metrics['roc_auc']):.4f}") + + print("\nOptimal F1 Threshold:") + print(f" Source: threshold={source_metrics['best_f1_threshold']:.4f}, F1={source_metrics['best_f1_score']:.4f}") + print(f" Target: threshold={target_metrics['best_f1_threshold']:.4f}, F1={target_metrics['best_f1_score']:.4f}") + print(f" Best F1 Difference: {abs(source_metrics['best_f1_score'] - target_metrics['best_f1_score']):.4f}") + + print("\n60% Precision Threshold:") + print(f" Source: threshold={source_metrics['prec_60_threshold']:.4f}, Precision={source_metrics['prec_60_precision']:.4f}, Recall={source_metrics['prec_60_recall']:.4f}, F1={source_metrics['prec_60_f1']:.4f}") + print(f" Target: threshold={target_metrics['prec_60_threshold']:.4f}, Precision={target_metrics['prec_60_precision']:.4f}, Recall={target_metrics['prec_60_recall']:.4f}, F1={target_metrics['prec_60_f1']:.4f}") + print("=" * 80) + + # Save summary + summary_path = os.path.join(args.output_dir, "summary.txt") + with open(summary_path, "w") as f: + f.write("DOMAIN ADVERSARIAL MODEL - TEST SUMMARY\n") + f.write("=" * 80 + "\n\n") + f.write(f"Model: {args.model_path}\n\n") + + f.write("=" * 80 + "\n") + f.write("METRICS AT DEFAULT THRESHOLD (0.5)\n") + f.write("=" * 80 + "\n\n") + f.write("Source Domain Performance:\n") + for key in ["accuracy", "precision", "recall", "f1", "roc_auc", "pr_auc"]: + f.write(f" {key}: {source_metrics[key]:.4f}\n") + f.write("\nTarget Domain Performance:\n") + for key in ["accuracy", "precision", "recall", "f1", "roc_auc", "pr_auc"]: + f.write(f" {key}: {target_metrics[key]:.4f}\n") + f.write("\nDomain Differences:\n") + f.write(f" F1 Difference: {abs(source_metrics['f1'] - target_metrics['f1']):.4f}\n") + f.write(f" ROC AUC Difference: {abs(source_metrics['roc_auc'] - target_metrics['roc_auc']):.4f}\n") + + f.write("\n" + "=" * 80 + "\n") + f.write("METRICS AT OPTIMAL F1 THRESHOLD\n") + f.write("=" * 80 + "\n\n") + f.write("Source Domain (Optimal F1):\n") + f.write(f" Threshold: {source_metrics['best_f1_threshold']:.4f}\n") + f.write(f" F1 Score: {source_metrics['best_f1_score']:.4f}\n") + f.write(f" Accuracy: {source_metrics['best_f1_accuracy']:.4f}\n") + f.write(f" Precision: {source_metrics['best_f1_precision']:.4f}\n") + f.write(f" Recall: {source_metrics['best_f1_recall']:.4f}\n") + + f.write("\nTarget Domain (Optimal F1):\n") + f.write(f" Threshold: {target_metrics['best_f1_threshold']:.4f}\n") + f.write(f" F1 Score: {target_metrics['best_f1_score']:.4f}\n") + f.write(f" Accuracy: {target_metrics['best_f1_accuracy']:.4f}\n") + f.write(f" Precision: {target_metrics['best_f1_precision']:.4f}\n") + f.write(f" Recall: {target_metrics['best_f1_recall']:.4f}\n") + + f.write("\nDomain Differences (Optimal F1):\n") + f.write(f" Best F1 Difference: {abs(source_metrics['best_f1_score'] - target_metrics['best_f1_score']):.4f}\n") + f.write(f" Threshold Difference: {abs(source_metrics['best_f1_threshold'] - target_metrics['best_f1_threshold']):.4f}\n") + + f.write("\n" + "=" * 80 + "\n") + f.write("METRICS AT 60% PRECISION THRESHOLD\n") + f.write("=" * 80 + "\n\n") + f.write("Source Domain (60% Precision):\n") + f.write(f" Threshold: {source_metrics['prec_60_threshold']:.4f}\n") + f.write(f" Precision: {source_metrics['prec_60_precision']:.4f}\n") + f.write(f" Recall: {source_metrics['prec_60_recall']:.4f}\n") + f.write(f" F1 Score: {source_metrics['prec_60_f1']:.4f}\n") + f.write(f" Accuracy: {source_metrics['prec_60_accuracy']:.4f}\n") + + f.write("\nTarget Domain (60% Precision):\n") + f.write(f" Threshold: {target_metrics['prec_60_threshold']:.4f}\n") + f.write(f" Precision: {target_metrics['prec_60_precision']:.4f}\n") + f.write(f" Recall: {target_metrics['prec_60_recall']:.4f}\n") + f.write(f" F1 Score: {target_metrics['prec_60_f1']:.4f}\n") + f.write(f" Accuracy: {target_metrics['prec_60_accuracy']:.4f}\n") + + print(f"\nSummary saved to: {summary_path}") + print(f"\nTesting completed successfully!") + + +if __name__ == "__main__": + main() diff --git a/RuBR/experiments/domain_adaption/train_control_baseline.py b/RuBR/experiments/domain_adaption/train_control_baseline.py new file mode 100644 index 00000000..83acf219 --- /dev/null +++ b/RuBR/experiments/domain_adaption/train_control_baseline.py @@ -0,0 +1,396 @@ +""" +Control Training Script for Transient Detection + +This script trains a standard CNN model (without domain adversarial training) +on the source domain (e.g., Open Universe) to serve as a baseline for comparison +with domain adversarial approaches. + +This provides a control to measure the improvement gained from domain adaptation. +""" + +import gc +import argparse +import os +import tensorflow as tf +from tensorflow.keras import layers, Model +import numpy as np +from sklearn.metrics import confusion_matrix, classification_report +from matplotlib import pyplot as plt +from sklearn.model_selection import train_test_split +from model.data import load_dataset +from model.layers import image_encoder +from model.callbacks import F1EarlyStopping + +# Set random seeds for reproducibility +np.random.seed(42) +tf.random.set_seed(42) + +def load_data(data_path): + """ + Load the dataset from the specified path. + + Args: + data_path (str): Path to the dataset file (.npz file). + + Returns: + tuple: X, feats, y, metadata as numpy arrays. + """ + return load_dataset(data_path, mmap=False, allow_npy_dict=False) + + +def create_standard_model(img_shape, num_features): + """ + Create a standard CNN model for transient detection (no domain adaptation). + + Args: + img_shape (tuple): Shape of input images (H, W, C) + num_features (int): Number of tabular features + + Returns: + tf.keras.Model: Standard classification model + """ + # Inputs + img_input = layers.Input(shape=img_shape, name="image_input") + feat_input = layers.Input(shape=(num_features,), name="feature_input") + + # Feature extraction + img_features = image_encoder(img_input, img_shape, mode="mean") + feat_features = layers.Dense(32, activation="relu", name="feat_encoder")(feat_input) + + # Combine features + combined_features = layers.Concatenate(name="combined_features")([img_features, feat_features]) + + # Classification head + x = layers.Dense(128, activation="relu", name="fc1")(combined_features) + x = layers.Dropout(0.3)(x) + x = layers.Dense(64, activation="relu", name="fc2")(x) + output = layers.Dense(1, activation="sigmoid", name="output")(x) + + # Create model + model = Model( + inputs=[img_input, feat_input], + outputs=output, + name="StandardClassifier" + ) + + # Compile + model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), + loss="binary_crossentropy", + metrics=["accuracy", "precision", "recall"] + ) + + return model + + +def plot_training_history(history, output_dir): + """Plot and save training history.""" + fig, axes = plt.subplots(2, 2, figsize=(14, 10)) + + # Accuracy + axes[0, 0].plot(history.history["accuracy"], label="Train") + axes[0, 0].plot(history.history["val_accuracy"], label="Val") + axes[0, 0].set_title("Accuracy") + axes[0, 0].set_xlabel("Epoch") + axes[0, 0].set_ylabel("Accuracy") + axes[0, 0].legend() + axes[0, 0].grid(True) + + # Loss + axes[0, 1].plot(history.history["loss"], label="Train") + axes[0, 1].plot(history.history["val_loss"], label="Val") + axes[0, 1].set_title("Loss") + axes[0, 1].set_xlabel("Epoch") + axes[0, 1].set_ylabel("Loss") + axes[0, 1].legend() + axes[0, 1].grid(True) + + # Precision + axes[1, 0].plot(history.history["precision"], label="Train") + axes[1, 0].plot(history.history["val_precision"], label="Val") + axes[1, 0].set_title("Precision") + axes[1, 0].set_xlabel("Epoch") + axes[1, 0].set_ylabel("Precision") + axes[1, 0].legend() + axes[1, 0].grid(True) + + # Recall + axes[1, 1].plot(history.history["recall"], label="Train") + axes[1, 1].plot(history.history["val_recall"], label="Val") + axes[1, 1].set_title("Recall") + axes[1, 1].set_xlabel("Epoch") + axes[1, 1].set_ylabel("Recall") + axes[1, 1].legend() + axes[1, 1].grid(True) + + plt.tight_layout() + plt.savefig(os.path.join(output_dir, "training_history.png"), dpi=300, bbox_inches="tight") + plt.close() + + +def main(): + parser = argparse.ArgumentParser( + description="Train standard model (control) for transient detection" + ) + + # Data arguments + parser.add_argument( + "--train_data", + type=str, + required=True, + help="Path to training data (.npz file)", + ) + parser.add_argument( + "--max_samples", + type=int, + default=None, + help="Maximum number of training samples (default: use all)", + ) + + # Training arguments + parser.add_argument( + "--epochs", + type=int, + default=60, + help="Number of training epochs (default: 60)" + ) + parser.add_argument( + "--batch_size", + type=int, + default=256, + help="Batch size for training (default: 512)", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=0.005, + help="Learning rate for Adam optimizer (default: 0.001)", + ) + + # Data split arguments + parser.add_argument( + "--val_size", + type=float, + default=0.15, + help="Validation set size fraction (default: 0.15)", + ) + parser.add_argument( + "--random_state", + type=int, + default=42, + help="Random state for reproducible splits (default: 42)", + ) + + # Model arguments + parser.add_argument( + "--patience", + type=int, + default=10, + help="Early stopping patience (default: 10)" + ) + parser.add_argument( + "--class_weight_pos", + type=float, + default=2.0, + help="Class weight for positive class (default: 2.0)", + ) + + # Output arguments + parser.add_argument( + "--output_dir", + type=str, + default="./control_output_snr", + help="Directory to save model and results (default: ./control_output)", + ) + + # GPU arguments + parser.add_argument( + "--gpu", + type=int, + default=None, + help="GPU device ID to use (default: auto-select)", + ) + + args = parser.parse_args() + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Set GPU device if specified + if args.gpu is not None: + physical_devices = tf.config.experimental.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_visible_devices( + physical_devices[args.gpu], "GPU" + ) + tf.config.experimental.set_memory_growth(physical_devices[args.gpu], True) + + print("=" * 80) + print("STANDARD MODEL TRAINING (CONTROL)") + print("=" * 80) + print(f"Training data: {args.train_data}") + print(f"Output directory: {args.output_dir}") + print(f"Epochs: {args.epochs}") + print(f"Batch size: {args.batch_size}") + print(f"Learning rate: {args.learning_rate}") + print("=" * 80) + + # Load training data + print("\n--- Loading Training Data ---") + X, feats, y, metadata = load_data(args.train_data) + + # Limit samples if specified + if args.max_samples is not None and len(X) > args.max_samples: + print(f"Limiting to {args.max_samples} samples") + indices = np.random.choice(len(X), args.max_samples, replace=False) + X = X[indices] + feats = feats[indices] + y = y[indices] + metadata = metadata[indices] + + # Remove NaN values + mask = np.isnan(X).any(axis=(1, 2, 3)) | np.isnan(feats).any(axis=1) + if mask.any(): + print(f"Removing {mask.sum()} samples with NaN values") + X = X[~mask] + feats = feats[~mask] + y = y[~mask] + metadata = metadata[~mask] + + # Normalize data + print("\nNormalizing data...") + X = (X - X.mean(axis=(0, 1, 2), keepdims=True)) / (X.std(axis=(0, 1, 2), keepdims=True) + 1e-6) + feats = (feats - feats.mean(axis=0)) / (feats.std(axis=0) + 1e-6) + + print(f"\nProcessed data shapes:") + print(f"X: {X.shape}, feats: {feats.shape}, y: {y.shape}") + print(f"Class distribution: {np.bincount(y.astype(int))}") + + # Split into train and validation + print(f"\nSplitting data (val_size={args.val_size})...") + X_train, X_val, feats_train, feats_val, y_train, y_val = train_test_split( + X, feats, y, + test_size=args.val_size, + stratify=y, + random_state=args.random_state, + shuffle=True + ) + + print(f"Train set: {len(X_train)} samples") + print(f"Validation set: {len(X_val)} samples") + print(f"Train class distribution: {np.bincount(y_train.astype(int))}") + print(f"Val class distribution: {np.bincount(y_val.astype(int))}") + + # Clean up memory + del X, feats, y + gc.collect() + + # Create model + img_shape = X_train[0].shape + num_features = feats_train.shape[1] + + print(f"\nCreating standard model...") + print(f"Image shape: {img_shape}") + print(f"Number of features: {num_features}") + + model = create_standard_model(img_shape, num_features) + model.summary() + + # Setup callbacks + early_stopping = F1EarlyStopping( + precision_key="val_precision", + recall_key="val_recall", + patience=args.patience, + restore_best_weights=True, + ) + + callbacks = [ + early_stopping, + tf.keras.callbacks.ModelCheckpoint( + os.path.join(args.output_dir, "best_model.h5"), + save_best_only=True, + monitor="val_recall", + mode="max", + verbose=1 + ), + tf.keras.callbacks.ReduceLROnPlateau( + monitor="val_loss", factor=0.5, patience=5, min_lr=1e-6, verbose=1 + ), + ] + + # Sample weights (handle class imbalance) + sample_weights = np.ones(len(y_train)) + sample_weights[y_train == 1] = args.class_weight_pos + + print(f"\nClass weight for positive samples: {args.class_weight_pos}") + + # Train model + print(f"\nStarting training...") + history = model.fit( + [X_train, feats_train], + y_train, + epochs=args.epochs, + batch_size=args.batch_size, + validation_data=([X_val, feats_val], y_val), + callbacks=callbacks, + sample_weight=sample_weights, + verbose=1 + ) + + # Save final model + final_model_path = os.path.join(args.output_dir, "final_model.h5") + model.save(final_model_path) + print(f"\nFinal model saved to: {final_model_path}") + + # Plot training history + plot_training_history(history, args.output_dir) + + # Evaluate on validation set + print(f"\nEvaluating model on validation set...") + results = model.evaluate([X_val, feats_val], y_val, verbose=0) + + print(f"\nValidation Results:") + for name, value in zip(model.metrics_names, results): + print(f"{name}: {value:.4f}") + + # Make predictions + y_pred_prob = model.predict([X_val, feats_val], verbose=0).flatten() + y_pred = (y_pred_prob > 0.5).astype(int) + + # Classification report + print(f"\nClassification Report:") + print(classification_report(y_val, y_pred)) + + # Confusion matrix + cm = confusion_matrix(y_val, y_pred) + print(f"\nConfusion Matrix:") + print(cm) + + # Calculate F1 score + tn, fp, fn, tp = cm.ravel() + precision = tp / (tp + fp) if (tp + fp) > 0 else 0 + recall = tp / (tp + fn) if (tp + fn) > 0 else 0 + f1 = 2 * precision * recall / (precision + recall) if (precision + recall) > 0 else 0 + print(f"\nF1 Score: {f1:.4f}") + + # Save results + results_path = os.path.join(args.output_dir, "training_results.txt") + with open(results_path, "w") as f: + f.write("CONTROL MODEL TRAINING RESULTS\n") + f.write("=" * 80 + "\n\n") + f.write(f"Training data: {args.train_data}\n\n") + f.write("Validation Results:\n") + for name, value in zip(model.metrics_names, results): + f.write(f"{name}: {value:.4f}\n") + f.write(f"\nF1 Score: {f1:.4f}\n") + f.write(f"\nClassification Report:\n") + f.write(classification_report(y_val, y_pred)) + f.write(f"\nConfusion Matrix:\n{cm}\n") + + print(f"\nResults saved to: {results_path}") + print(f"\nTraining completed successfully!") + print(f"All outputs saved to: {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/RuBR/experiments/domain_adaption/train_dann.py b/RuBR/experiments/domain_adaption/train_dann.py new file mode 100644 index 00000000..a03f328e --- /dev/null +++ b/RuBR/experiments/domain_adaption/train_dann.py @@ -0,0 +1,514 @@ +""" +Domain Adversarial Training for Transient Detection + +This script implements Domain Adversarial Neural Networks (DANN) to learn +domain-invariant features for transient detection across different data domains. + +Architecture: +- Feature Extractor: Shared encoder for both domains +- Label Classifier: Predicts transient/non-transient +- Domain Classifier: Predicts source/target domain (with gradient reversal) +""" + +import gc +import argparse +import os +import tensorflow as tf +import numpy as np +from sklearn.metrics import confusion_matrix, classification_report +from matplotlib import pyplot as plt +from sklearn.model_selection import train_test_split +from model.data import load_dataset +from model.callbacks import F1EarlyStopping + +# Import model components from flat package module +from model.dann_model import ( + rot90_k1, + rot90_k2, + rot90_k3, + GradientReversalLayer, + create_dann_model +) + +# Set random seeds for reproducibility +np.random.seed(42) +tf.random.set_seed(42) + +only_tp = True + +def load_data(data_path): + return load_dataset(data_path, mmap=True, allow_npy_dict=False) + + +class BatchMetricsLogger(tf.keras.callbacks.Callback): + """Logs metrics at each batch for iteration-wise plotting.""" + + def __init__(self): + super(BatchMetricsLogger, self).__init__() + self.batch_metrics = { + 'iteration': [], + 'label_output_loss': [], + 'domain_output_loss': [], + 'loss': [] + } + self.current_iteration = 0 + + def on_train_batch_end(self, batch, logs=None): + self.current_iteration += 1 + self.batch_metrics['iteration'].append(self.current_iteration) + self.batch_metrics['label_output_loss'].append(logs.get('label_output_loss', 0)) + self.batch_metrics['domain_output_loss'].append(logs.get('domain_output_loss', 0)) + self.batch_metrics['loss'].append(logs.get('loss', 0)) + + +def plot_training_history(history, batch_metrics, output_dir): + """Plot and save training history (epoch-wise and iteration-wise).""" + + # Create figure with more subplots to include iteration-wise plots + fig = plt.figure(figsize=(20, 14)) + gs = fig.add_gridspec(3, 3, hspace=0.3, wspace=0.3) + + # Row 1: Label classifier metrics (epoch-wise) + ax1 = fig.add_subplot(gs[0, 0]) + ax1.plot(history.history["label_output_accuracy"], label="Train") + ax1.plot(history.history["val_label_output_accuracy"], label="Val") + ax1.set_title("Label Accuracy (Epoch-wise)") + ax1.set_xlabel("Epoch") + ax1.set_ylabel("Accuracy") + ax1.legend() + ax1.grid(True) + + ax2 = fig.add_subplot(gs[0, 1]) + ax2.plot(history.history["label_output_loss"], label="Train") + ax2.plot(history.history["val_label_output_loss"], label="Val") + ax2.set_title("Label Loss (Epoch-wise)") + ax2.set_xlabel("Epoch") + ax2.set_ylabel("Loss") + ax2.legend() + ax2.grid(True) + + ax3 = fig.add_subplot(gs[0, 2]) + ax3.plot(history.history["label_output_precision"], label="Train Prec") + ax3.plot(history.history["val_label_output_precision"], label="Val Prec") + ax3.plot(history.history["label_output_recall"], label="Train Rec") + ax3.plot(history.history["val_label_output_recall"], label="Val Rec") + ax3.set_title("Label Precision/Recall (Epoch-wise)") + ax3.set_xlabel("Epoch") + ax3.set_ylabel("Score") + ax3.legend() + ax3.grid(True) + + # Row 2: Domain classifier metrics (epoch-wise) + ax4 = fig.add_subplot(gs[1, 0]) + ax4.plot(history.history["domain_output_accuracy"], label="Train") + ax4.plot(history.history["val_domain_output_accuracy"], label="Val") + ax4.set_title("Domain Accuracy (Epoch-wise)") + ax4.set_xlabel("Epoch") + ax4.set_ylabel("Accuracy") + ax4.legend() + ax4.grid(True) + + ax5 = fig.add_subplot(gs[1, 1]) + ax5.plot(history.history["domain_output_loss"], label="Train") + ax5.plot(history.history["val_domain_output_loss"], label="Val") + ax5.set_title("Domain Loss (Epoch-wise)") + ax5.set_xlabel("Epoch") + ax5.set_ylabel("Loss") + ax5.legend() + ax5.grid(True) + + ax6 = fig.add_subplot(gs[1, 2]) + ax6.plot(history.history["loss"], label="Train Total") + ax6.plot(history.history["val_loss"], label="Val Total") + ax6.set_title("Total Loss (Epoch-wise)") + ax6.set_xlabel("Epoch") + ax6.set_ylabel("Loss") + ax6.legend() + ax6.grid(True) + + # Row 3: Iteration-wise metrics (from batch logger) + if batch_metrics is not None and len(batch_metrics['iteration']) > 0: + ax7 = fig.add_subplot(gs[2, 0]) + ax7.plot(batch_metrics['iteration'], batch_metrics['label_output_loss'], + linewidth=0.5, alpha=0.7) + ax7.set_title("Label Loss (Iteration-wise)") + ax7.set_xlabel("Iteration") + ax7.set_ylabel("Loss") + ax7.grid(True, alpha=0.3) + + ax8 = fig.add_subplot(gs[2, 1]) + ax8.plot(batch_metrics['iteration'], batch_metrics['domain_output_loss'], + linewidth=0.5, alpha=0.7) + ax8.set_title("Domain Loss (Iteration-wise)") + ax8.set_xlabel("Iteration") + ax8.set_ylabel("Loss") + ax8.grid(True, alpha=0.3) + + ax9 = fig.add_subplot(gs[2, 2]) + ax9.plot(batch_metrics['iteration'], batch_metrics['loss'], + linewidth=0.5, alpha=0.7) + ax9.set_title("Total Loss (Iteration-wise)") + ax9.set_xlabel("Iteration") + ax9.set_ylabel("Loss") + ax9.grid(True, alpha=0.3) + + plt.savefig(os.path.join(output_dir, "training_history.png"), dpi=300, bbox_inches="tight") + plt.close() + + # Save batch metrics to file for later analysis + if batch_metrics is not None and len(batch_metrics['iteration']) > 0: + batch_metrics_path = os.path.join(output_dir, "batch_metrics.npz") + np.savez(batch_metrics_path, **batch_metrics) + print(f"Batch metrics saved to: {batch_metrics_path}") + + +def main(): + parser = argparse.ArgumentParser( + description="Train Domain Adversarial model for transient detection" + ) + + # Data arguments + parser.add_argument( + "--source_data", + type=str, + required=True, + help="Path to source domain training data (.npz file)", + ) + parser.add_argument( + "--target_data", + type=str, + required=True, + help="Path to target domain data (.npz file)", + ) + parser.add_argument( + "--max_samples", + type=int, + default=100_000, + help="Maximum number of samples per domain (default: use all)", + ) + + # Training arguments + parser.add_argument( + "--epochs", + type=int, + default=60, + help="Number of training epochs (default: 60)" + ) + parser.add_argument( + "--batch_size", + type=int, + default=256, + help="Batch size for training (default: 512)", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=0.0025, + help="Learning rate for Adam optimizer (default: 0.0025)", + ) + parser.add_argument( + "--lambda_domain", + type=float, + default=1.0, + help="Weight for gradient reversal layer (default: 1.0)", + ) + + # Data split arguments + parser.add_argument( + "--val_size", + type=float, + default=0.15, + help="Validation set size fraction (default: 0.15)", + ) + parser.add_argument( + "--random_state", + type=int, + default=42, + help="Random state for reproducible splits (default: 42)", + ) + + # Model arguments + parser.add_argument( + "--patience", + type=int, + default=10, + help="Early stopping patience (default: 10)" + ) + parser.add_argument( + "--class_weight_pos", + type=float, + default=2.0, + help="Class weight for positive class in label classification (default: 8.0)", + ) + + # Output arguments + parser.add_argument( + "--output_dir", + type=str, + default="./dann_output_tp_only", + help="Directory to save model and results (default: ./dann_output_full)", + ) + + # GPU arguments + parser.add_argument( + "--gpu", + type=int, + default=None, + help="GPU device ID to use (default: auto-select)", + ) + + args = parser.parse_args() + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Set GPU device if specified + if args.gpu is not None: + physical_devices = tf.config.experimental.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_visible_devices( + physical_devices[args.gpu], "GPU" + ) + tf.config.experimental.set_memory_growth(physical_devices[args.gpu], True) + + print("=" * 80) + print("DOMAIN ADVERSARIAL NEURAL NETWORK TRAINING") + print("=" * 80) + print(f"Source domain data: {args.source_data}") + print(f"Target domain data: {args.target_data}") + print(f"Output directory: {args.output_dir}") + print(f"Epochs: {args.epochs}") + print(f"Batch size: {args.batch_size}") + print(f"Learning rate: {args.learning_rate}") + print(f"Lambda (domain): {args.lambda_domain}") + print("=" * 80) + + # Load source domain data + print("\n--- Loading Source Domain Data ---") + X_src, feats_src, y_src, metadata_src = load_data(args.source_data) + + # Load target domain data + print("\n--- Loading Target Domain Data ---") + X_tgt, feats_tgt, y_tgt, metadata_tgt = load_data(args.target_data) + + # Limit samples if specified + if args.max_samples is not None: + if len(X_src) > args.max_samples: + print(f"Limiting source domain to {args.max_samples} samples") + indices = np.random.choice(len(X_src), args.max_samples, replace=False) + X_src = X_src[indices] + feats_src = feats_src[indices] + y_src = y_src[indices] + metadata_src = metadata_src[indices] + + if len(X_tgt) > args.max_samples: + print(f"Limiting target domain to {args.max_samples} samples") + indices = np.random.choice(len(X_tgt), args.max_samples, replace=False) + X_tgt = X_tgt[indices] + feats_tgt = feats_tgt[indices] + y_tgt = y_tgt[indices] + metadata_tgt = metadata_tgt[indices] + + # Remove NaN values + mask_src = np.isnan(X_src).any(axis=(1, 2, 3)) | np.isnan(feats_src).any(axis=1) + if mask_src.any(): + print(f"Removing {mask_src.sum()} source samples with NaN values") + X_src = X_src[~mask_src] + feats_src = feats_src[~mask_src] + y_src = y_src[~mask_src] + metadata_src = metadata_src[~mask_src] + + mask_tgt = np.isnan(X_tgt).any(axis=(1, 2, 3)) | np.isnan(feats_tgt).any(axis=1) + if mask_tgt.any(): + print(f"Removing {mask_tgt.sum()} target samples with NaN values") + X_tgt = X_tgt[~mask_tgt] + feats_tgt = feats_tgt[~mask_tgt] + y_tgt = y_tgt[~mask_tgt] + metadata_tgt = metadata_tgt[~mask_tgt] + + # Normalize data + print("\nNormalizing data...") + X_src = (X_src - X_src.mean(axis=(0, 1, 2), keepdims=True)) / (X_src.std(axis=(0, 1, 2), keepdims=True) + 1e-6) + X_tgt = (X_tgt - X_tgt.mean(axis=(0, 1, 2), keepdims=True)) / (X_tgt.std(axis=(0, 1, 2), keepdims=True) + 1e-6) + + feats_src = (feats_src - feats_src.mean(axis=0)) / (feats_src.std(axis=0) + 1e-6) + feats_tgt = (feats_tgt - feats_tgt.mean(axis=0)) / (feats_tgt.std(axis=0) + 1e-6) + + # Create domain labels (0 = source, 1 = target) + domain_src = np.zeros(len(X_src)) + domain_tgt = np.ones(len(X_tgt)) + + # Combine data from both domains + X_combined = np.concatenate([X_src, X_tgt], axis=0) + feats_combined = np.concatenate([feats_src, feats_tgt], axis=0) + y_combined = np.concatenate([y_src, y_tgt], axis=0) + domain_combined = np.concatenate([domain_src, domain_tgt], axis=0) + metadata_combined = np.concatenate([metadata_src, metadata_tgt], axis=0) + + print(f"\nCombined data shapes:") + print(f"X: {X_combined.shape}, feats: {feats_combined.shape}") + print(f"Labels: {y_combined.shape}, Domains: {domain_combined.shape}") + print(f"Source samples: {len(X_src)}, Target samples: {len(X_tgt)}") + print(f"Label distribution: {np.bincount(y_combined.astype(int))}") + print(f"Domain distribution: {np.bincount(domain_combined.astype(int))}") + + # Split into train and validation + print(f"\nSplitting data (val_size={args.val_size})...") + ( + X_train, X_val, + feats_train, feats_val, + y_train, y_val, + domain_train, domain_val, + metadata_train, metadata_val + ) = train_test_split( + X_combined, feats_combined, y_combined, domain_combined, metadata_combined, + test_size=args.val_size, + stratify=domain_combined, # Stratify by domain to keep balance + random_state=args.random_state, + shuffle=True + ) + + print(f"Train set: {len(X_train)} samples") + print(f"Validation set: {len(X_val)} samples") + + # Clean up memory + del X_src, X_tgt, feats_src, feats_tgt, y_src, y_tgt + del X_combined, feats_combined, y_combined, domain_combined + gc.collect() + + # Create model + img_shape = X_train[0].shape + num_features = feats_train.shape[1] + + print(f"\nCreating DANN model...") + print(f"Image shape: {img_shape}") + print(f"Number of features: {num_features}") + print(f"Lambda (gradient reversal): {args.lambda_domain}") + + model = create_dann_model(img_shape, num_features, lambda_domain=args.lambda_domain) + model.summary() + + # Setup callbacks + batch_logger = BatchMetricsLogger() + early_stopping = F1EarlyStopping( + precision_key="val_label_output_precision", + recall_key="val_label_output_recall", + patience=args.patience, + restore_best_weights=True, + ) + + callbacks = [ + batch_logger, + early_stopping, + tf.keras.callbacks.ModelCheckpoint( + os.path.join(args.output_dir, "best_model.h5"), + save_best_only=True, + monitor="val_label_output_recall", + mode="max", + verbose=1 + ), + tf.keras.callbacks.ReduceLROnPlateau( + monitor="val_loss", factor=0.5, patience=5, min_lr=1e-6, verbose=1 + ), + ] + + # Calculate class weights for label classification (handle class imbalance) + # In unsupervised DANN: only train label classifier on SOURCE domain + # Train domain classifier on BOTH domains + + # Identify source vs target samples in training set + is_source_train = (domain_train == 0) + is_target_train = (domain_train == 1) + + n_source = is_source_train.sum() + n_target = is_target_train.sum() + n_neg_src = np.sum((y_train == 0) & is_source_train) + n_pos_src = np.sum((y_train == 1) & is_source_train) + + print(f"\nTraining set composition:") + print(f" Source samples: {n_source} (Negative={n_neg_src}, Positive={n_pos_src})") + print(f" Target samples: {n_target} (labels NOT used for training)") + print(f"Class weight for positive samples: {args.class_weight_pos}") + + # For multi-output models, sample_weight needs to be a list/tuple of arrays + # One array per output, matching the order of outputs in the model + + # Label classifier: Only train on SOURCE domain (set target weights to 0) + label_sample_weights = np.zeros(len(y_train)) # Initialize all to 0 + label_sample_weights[is_source_train] = 1.0 # Enable source samples + label_sample_weights[(y_train == 1) & is_source_train] = args.class_weight_pos # Weight positive class + + # Domain classifier: Train on BOTH domains equally + domain_sample_weights = np.ones(len(y_train)) if not only_tp else (y_train == 1) + + # Train model + print(f"\nStarting training...") + history = model.fit( + [X_train, feats_train], + [y_train, domain_train], # Use list instead of dict + epochs=args.epochs, + batch_size=args.batch_size, + validation_data=( + [X_val, feats_val], + [y_val, domain_val] # Use list instead of dict + ), + callbacks=callbacks, + sample_weight=[label_sample_weights, domain_sample_weights], # List of arrays + verbose=1 + ) + + # Save final model + final_model_path = os.path.join(args.output_dir, "final_model.h5") + model.save(final_model_path) + print(f"\nFinal model saved to: {final_model_path}") + + # Plot training history (epoch-wise and iteration-wise) + plot_training_history(history, batch_logger.batch_metrics, args.output_dir) + + # Evaluate on validation set + print(f"\nEvaluating model on validation set...") + results = model.evaluate( + [X_val, feats_val], + [y_val, domain_val], # Use list instead of dict + verbose=0 + ) + + print(f"\nValidation Results:") + for name, value in zip(model.metrics_names, results): + print(f"{name}: {value:.4f}") + + # Make predictions + y_pred_label, y_pred_domain = model.predict([X_val, feats_val], verbose=0) + y_pred_label_binary = (y_pred_label > 0.5).astype(int).flatten() + + # Classification report for labels + print(f"\nLabel Classification Report:") + print(classification_report(y_val, y_pred_label_binary)) + + # Confusion matrix + cm = confusion_matrix(y_val, y_pred_label_binary) + print(f"\nLabel Confusion Matrix:") + print(cm) + + # Save results + results_path = os.path.join(args.output_dir, "training_results.txt") + with open(results_path, "w") as f: + f.write("DOMAIN ADVERSARIAL TRAINING RESULTS\n") + f.write("=" * 80 + "\n\n") + f.write(f"Source domain: {args.source_data}\n") + f.write(f"Target domain: {args.target_data}\n") + f.write(f"Lambda (gradient reversal): {args.lambda_domain}\n\n") + f.write("Validation Results:\n") + for name, value in zip(model.metrics_names, results): + f.write(f"{name}: {value:.4f}\n") + f.write(f"\nLabel Classification Report:\n") + f.write(classification_report(y_val, y_pred_label_binary)) + f.write(f"\nLabel Confusion Matrix:\n{cm}\n") + + print(f"\nResults saved to: {results_path}") + print(f"\nTraining completed successfully!") + print(f"All outputs saved to: {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/RuBR/experiments/rotinv/evaluate_rotinv_with_features.py b/RuBR/experiments/rotinv/evaluate_rotinv_with_features.py new file mode 100644 index 00000000..a570f2e0 --- /dev/null +++ b/RuBR/experiments/rotinv/evaluate_rotinv_with_features.py @@ -0,0 +1,207 @@ +import os +import argparse +import numpy as np +import tensorflow as tf +import pandas as pd +from sklearn.metrics import classification_report +import matplotlib.pyplot as plt +import gc +from model.layers import rot90_k1, rot90_k2, rot90_k3 + +# Import load_data from the package train module +from experiments.rotinv.train_rotinv_with_features import load_data + +def load_model(path): + # Add custom Lambda functions to custom_objects for model loading + model = tf.keras.models.load_model( + path, + custom_objects={ + 'tf': tf, + 'rot90_k1': rot90_k1, + 'rot90_k2': rot90_k2, + 'rot90_k3': rot90_k3 + } + ) + return model + +def evaluate_model(model, X, feats, y): + predictions = model.predict([X, feats], verbose=0).flatten() + y_pred = predictions > 0.5 + assert y_pred.shape == y.shape + accuracy = np.mean(y_pred == y) + precision = ( + np.sum((y_pred == 1) & (y == 1)) / np.sum(y_pred == 1) + if np.sum(y_pred == 1) > 0 + else 0 + ) + recall = ( + np.sum((y_pred == 1) & (y == 1)) / np.sum(y == 1) if np.sum(y == 1) > 0 else 0 + ) + return accuracy, precision, recall, predictions + +def main(): + parser = argparse.ArgumentParser(description="Test a rot-inv-feat model on batched data") + parser.add_argument( + "--data_dir", + type=str, + required=True, + help="Directory containing the batched data files (.npz)", + ) + parser.add_argument( + "--model_path", + type=str, + required=True, + help="Path to the trained model file (.h5)", + ) + parser.add_argument( + "--output_dir", + type=str, + default="./inj_only_outs_sqrt_var", + help="Directory to save output plots (default: current directory)", + ) + parser.add_argument( + "--num_thresholds", + type=int, + default=100, + help="Number of thresholds to test for precision-recall curve (default: 100)", + ) + args = parser.parse_args() + + os.makedirs(args.output_dir, exist_ok=True) + all_predictions = [] + all_y = [] + all_metadata = [] + model = load_model(args.model_path) + + print(f"Using model: {args.model_path}") + print(f"Data directory: {args.data_dir}") + print(f"Output directory: {args.output_dir}") + print("=" * 50) + + for data_path in sorted(os.listdir(args.data_dir)): + if data_path.endswith(".npz") or data_path.endswith(".npy"): + print(f"Loading data from {data_path}") + X, feats, y, metadata = load_data(os.path.join(args.data_dir, data_path)) + # Ensure feats is a 2D float array + if feats.dtype == object: + if len(feats) > 0 and isinstance(feats[0], dict): + feat_keys = sorted(feats[0].keys()) + feats = np.array([[f[k] for k in feat_keys] for f in feats], dtype=np.float32) + else: + feats = np.array([list(f.values()) if hasattr(f, "values") else f for f in feats], dtype=np.float32) + else: + feats = feats.astype(np.float32) + X = X.astype(np.float32) + y = y.astype(np.int32) + + accuracy, precision, recall, predictions = evaluate_model(model, X, feats, y) + print(f"Model Accuracy: {accuracy:.4f}") + print(f"Model Precision: {precision:.4f}") + print(f"Model Recall: {recall:.4f}") + all_metadata.extend(metadata) + all_predictions.extend(predictions) + all_y.extend(y) + del X, feats + gc.collect() + + thresholds = np.linspace(0, 1, args.num_thresholds) + precisions = [] + recalls = [] + y = np.array(all_y) + predictions = np.array(all_predictions) + + for t in thresholds: + y_pred = predictions > t + precision = ( + np.sum((y_pred == 1) & (y == 1)) / np.sum(y_pred == 1) + if np.sum(y_pred == 1) > 0 + else 0 + ) + recall = ( + np.sum((y_pred == 1) & (y == 1)) / np.sum(y == 1) + if np.sum(y == 1) > 0 + else 0 + ) + precisions.append(precision) + recalls.append(recall) + + plt.figure(figsize=(4, 3)) + plt.plot(thresholds, precisions, label="Precision") + plt.plot(thresholds, recalls, label="Recall") + precisions = np.array(precisions) + recalls = np.array(recalls) + + # Annotate precision and recall at precision of 90, 95, 98 + for target_precision in [0.90, 0.95, 0.98]: + idx = (np.abs(precisions - target_precision)).argmin() + t = thresholds[idx] + p = precisions[idx] + r = recalls[idx] + plt.scatter([t], [p], color="red") + plt.scatter([t], [r], color="green") + plt.annotate( + f"P={p:.2f}, R={r:.2f}\nT={t:.2f}", + (t, p), + textcoords="offset points", + xytext=(0,10), + ha='center', + color="red", + fontsize=9, + arrowprops=dict(arrowstyle="->", color="red", lw=1) + ) + + plt.xlabel("Threshold") + plt.ylabel("Score") + plt.title("Precision and Recall vs Threshold") + plt.legend() + plt.grid(True) + precision_recall_plot_path = os.path.join( + args.output_dir, "precision_recall_vs_threshold.png" + ) + plt.savefig(precision_recall_plot_path) + print(f"Saved precision-recall plot to: {precision_recall_plot_path}") + + # Find the threshold where precision is closest to 0.9 + target_precision = 0.90 + precisions = np.array(precisions) + thresholds = np.array(thresholds) + idx = (np.abs(precisions - target_precision)).argmin() + best_threshold = thresholds[idx] + print( + f"Threshold where precision is closest to {target_precision*100}%: {best_threshold:.3f} (Precision: {precisions[idx]:.3f}, Recall: {recalls[idx]:.3f})" + ) + + # Plot histogram of number of elements with y==1 in different metadata['mag'] bins + mag_values = np.array([m["mag"] for m in all_metadata]) + y_positive = y == 1 + mag_positive = mag_values[y_positive] + mag_predict_positive = mag_values[(predictions > best_threshold) & y_positive] + plt.figure(figsize=(8, 6)) + plt.hist( + [mag_positive, mag_predict_positive], + label=["PSF Detections", "Model Detections"], + color=["black", "blue"], + histtype="step", + ) + plt.yscale("log") + plt.legend() + plt.xlabel("Magnitude (mag)") + plt.ylabel("Count") + plt.title("Magnitude Histogram") + plt.grid(True) + histogram_plot_path = os.path.join(args.output_dir, "histogram_y1_mag_bins.png") + plt.savefig(histogram_plot_path) + print(f"Saved magnitude histogram to: {histogram_plot_path}") + + filters = np.array([m["filter"] for m in all_metadata]) + tp_filters = filters[y_positive] + print("Unique filters:", np.unique(filters)) + filter_counts = pd.Series(filters).value_counts() + print("Filter counts:\n", filter_counts) + print("True Positive Filter counts:\n", pd.Series(tp_filters).value_counts()) + y_pred_final = predictions > best_threshold + print("Final Classification Report (threshold={:.3f}):".format(best_threshold)) + print(classification_report(y, y_pred_final, digits=4)) + +if __name__ == "__main__": + main() diff --git a/RuBR/experiments/rotinv/train_rotinv_with_features.py b/RuBR/experiments/rotinv/train_rotinv_with_features.py new file mode 100644 index 00000000..339f03aa --- /dev/null +++ b/RuBR/experiments/rotinv/train_rotinv_with_features.py @@ -0,0 +1,541 @@ +import gc +import argparse +import os +import random +import tensorflow as tf +from tensorflow.keras import layers, Model +import numpy as np +from sklearn.metrics import confusion_matrix, classification_report +from matplotlib import pyplot as plt +from sklearn.model_selection import train_test_split +from model.data import load_dataset +from model.layers import image_encoder +from model.callbacks import F1EarlyStopping + +means = { + "F184": [0.57817131, 0.58065492, 0.02057936], + "H158": [1.16567825, 1.16309479, 0.03536996], + "J129": [0.92620432, 0.92055472, 0.01911548], + "K213": [1.04127802, 1.04668246, 0.05275641], + "R062": [0.81320585, 0.81542744, 0.01229751], + "Y106": [0.70102083, 0.7024932, 0.01716464], + "Z087": [1.13010884, 1.12491324, 0.01309704], +} + +vars = { + "F184": [24.01642823, 28.37142752, 2.33357588], + "H158": [141.10233678, 166.69334902, 13.48953537], + "J129": [103.2038142, 122.14426572, 12.47955061], + "K213": [47.1821188, 55.16905395, 4.47816794], + "R062": [178.64370608, 212.49874184, 16.16102734], + "Y106": [74.89405617, 88.88895939, 7.68320053], + "Z087": [378.46329339, 451.93130827, 41.04080681], +} + +def set_seed(seed): + os.environ["PYTHONHASHSEED"] = str(seed) + random.seed(seed) + np.random.seed(seed) + tf.random.set_seed(seed) + +def load_data(data_path): + """ + Load the dataset from the specified path. + + Args: + data_path (str): Path to the dataset file (.npz). + + Returns: + tuple: X, feats, y, metadata as numpy arrays. + """ + X, feats, y, metadata = load_dataset(data_path, mmap=False, allow_npy_dict=True) + + # Vectorized normalization by filter + filters = np.array([m['filter'] for m in metadata]) + unique_filters = np.unique(filters) + print(unique_filters) + for f in unique_filters: + idx = filters == f + X[idx] = (X[idx] - means[str(f)]) / np.sqrt(vars[str(f)]) + + return X, feats, y, metadata + +def prepare_training_data(tp_train, fp_train): + X_tp, feats_tp, y_tp, metadata_tp = load_data(tp_train) + X_fp, feats_fp, y_fp, metadata_fp = load_data(fp_train) + + # Combine the data + X = np.concatenate((X_tp, X_fp), axis=0) + feats = np.concatenate((feats_tp, feats_fp), axis=0) + y = np.concatenate((y_tp, y_fp), axis=0) + metadata = np.concatenate((metadata_tp, metadata_fp), axis=0) + + indices = np.random.permutation(len(y)) + + X = X[indices] + feats = feats[indices] + y = y[indices] + metadata = metadata[indices] + + return X, feats, y, metadata + +def prepare_testing_data(tp_test, fp_test): + X_tp, feats_tp, y_tp, metadata_tp = load_data(tp_test) + X_fp, feats_fp, y_fp, metadata_fp = load_data(fp_test) + + # Combine the data + X = np.concatenate((X_tp, X_fp), axis=0) + feats = np.concatenate((feats_tp, feats_fp), axis=0) + y = np.concatenate((y_tp, y_fp), axis=0) + metadata = np.concatenate((metadata_tp, metadata_fp), axis=0) + + return X, feats, y, metadata + + +def create_hybrid_model(img_shape, num_features): + """ + Create a hybrid model that combines CNN for images and dense layers for features. + + Args: + img_shape (tuple): Shape of input images (H, W, C) + num_features (int): Number of tabular features + + Returns: + tf.keras.Model: Compiled hybrid model + """ + # Image input branch + img_input = layers.Input(shape=img_shape) + + x = image_encoder( + img_input, + img_shape, + mode="mean", + conv_kernel_initializer="he_normal", + dense_kernel_initializer="he_normal", + ) + + # Tabular features branch + feat_input = layers.Input(shape=(num_features,)) + y = layers.Dense(32, activation="relu", kernel_initializer="he_normal")(feat_input) + + # Combine both branches + combined = layers.Concatenate()([x, y]) + + # Output layer + output = layers.Dense(1, activation="sigmoid", kernel_initializer="he_normal")(combined) + + # Create and compile model + model = Model(inputs=[img_input, feat_input], outputs=output) + model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), + loss="binary_crossentropy", + metrics=["accuracy", "precision", "recall"], + ) + + return model + +def plot_training_history(history, output_dir): + """Plot and save training history.""" + plt.figure(figsize=(15, 5)) + + # Plot training & validation accuracy + plt.subplot(1, 3, 1) + plt.plot(history.history["accuracy"], label="Train") + plt.plot(history.history["val_accuracy"], label="Validation") + plt.title("Model Accuracy") + plt.ylabel("Accuracy") + plt.xlabel("Epoch") + plt.legend() + plt.grid(True) + + # Plot training & validation loss + plt.subplot(1, 3, 2) + plt.plot(history.history["loss"], label="Train") + plt.plot(history.history["val_loss"], label="Validation") + plt.title("Model Loss") + plt.ylabel("Loss") + plt.xlabel("Epoch") + plt.legend() + plt.grid(True) + + # Plot precision and recall + plt.subplot(1, 3, 3) + plt.plot(history.history["precision"], label="Train Precision") + plt.plot(history.history["val_precision"], label="Val Precision") + plt.plot(history.history["recall"], label="Train Recall") + plt.plot(history.history["val_recall"], label="Val Recall") + plt.title("Precision and Recall") + plt.ylabel("Score") + plt.xlabel("Epoch") + plt.legend() + plt.grid(True) + + plt.tight_layout() + plt.savefig( + os.path.join(output_dir, "training_history.png"), dpi=300, bbox_inches="tight" + ) + plt.show() + + +def evaluate_model(model, X_test, feats_test, y_test, output_dir): + """Evaluate the trained model and save results.""" + # Evaluate on test set + loss, accuracy, precision, recall = model.evaluate( + [X_test, feats_test], y_test, verbose=0 + ) + + print(f"\nTest Results:") + print(f"Test Loss: {loss:.4f}") + print(f"Test Accuracy: {accuracy:.4f}") + print(f"Test Precision: {precision:.4f}") + print(f"Test Recall: {recall:.4f}") + + # Make predictions + y_pred_prob = model.predict([X_test, feats_test], verbose=0) + y_pred = (y_pred_prob > 0.5).astype(int) + + # Calculate confusion matrix and classification report + cm = confusion_matrix(y_test, y_pred) + print(f"\nConfusion Matrix:") + print(cm) + print(f"\nClassification Report:") + print(classification_report(y_test, y_pred)) + + # Save results to file + results_path = os.path.join(output_dir, "evaluation_results.txt") + with open(results_path, "w") as f: + f.write(f"Test Results:\n") + f.write(f"Test Loss: {loss:.4f}\n") + f.write(f"Test Accuracy: {accuracy:.4f}\n") + f.write(f"Test Precision: {precision:.4f}\n") + f.write(f"Test Recall: {recall:.4f}\n\n") + f.write(f"Confusion Matrix:\n{cm}\n\n") + f.write(f"Classification Report:\n{classification_report(y_test, y_pred)}\n") + + print(f"Results saved to: {results_path}") + + return y_pred_prob, y_pred + + +def main(model=None): + parser = argparse.ArgumentParser( + description="Train hybrid CNN model for transient detection" + ) + + # Data arguments + parser.add_argument( + "--data_path", + type=str, + required=True, + help="Path to the training data (.npz file)", + ) + parser.add_argument( + "--max_samples", + type=int, + default=300000, + help="Maximum number of samples to use (default: 300000)", + ) + + # Training arguments + parser.add_argument( + "--epochs", type=int, default=60, help="Number of training epochs (default: 30)" + ) + parser.add_argument( + "--batch_size", + type=int, + default=512, + help="Batch size for training (default: 512)", + ) + parser.add_argument( + "--learning_rate", + type=float, + default=0.001, + help="Learning rate for Adam optimizer (default: 0.001)", + ) + + # Data split arguments + parser.add_argument( + "--test_size", + type=float, + default=0.15, + help="Test set size fraction (default: 0.15)", + ) + parser.add_argument( + "--val_size", + type=float, + default=0.15, + help="Validation set size fraction (default: 0.15)", + ) + parser.add_argument( + "--random_state", + type=int, + default=42, + help="Random state for reproducible splits (default: 42)", + ) + + # Model arguments + parser.add_argument( + "--patience", type=int, default=10, help="Early stopping patience (default: 10)" + ) + parser.add_argument( + "--class_weight_pos", + type=float, + default=2.0, + help="Class weight for positive class (default: 2.0)", + ) + + # Output arguments + parser.add_argument( + "--output_dir", + type=str, + default="./training_output_sqrt_var", + help="Directory to save model and results (default: ./training_output)", + ) + parser.add_argument( + "--model_name", + type=str, + default="hybrid_model_rot_inv_final_new_data", + help="Base name for saved models (default: hybrid_model)", + ) + + # GPU arguments + parser.add_argument( + "--gpu", + type=int, + default=None, + help="GPU device ID to use (default: auto-select)", + ) + + args = parser.parse_args() + + set_seed(args.random_state) + + # Create output directory + os.makedirs(args.output_dir, exist_ok=True) + + # Set GPU device if specified + if args.gpu is not None: + physical_devices = tf.config.experimental.list_physical_devices("GPU") + if physical_devices: + tf.config.experimental.set_visible_devices( + physical_devices[args.gpu], "GPU" + ) + tf.config.experimental.set_memory_growth(physical_devices[args.gpu], True) + + print("=" * 60) + print("HYBRID CNN TRAINING") + print("=" * 60) + print(f"Data path: {args.data_path}") + print(f"Max samples: {args.max_samples}") + print(f"Output directory: {args.output_dir}") + print(f"Epochs: {args.epochs}") + print(f"Batch size: {args.batch_size}") + print(f"Learning rate: {args.learning_rate}") + print("=" * 60) + + # Load data + X, feats, y, metadata = load_data(args.data_path) + + # Limit samples if specified + if len(X) > args.max_samples: + print(f"Limiting to {args.max_samples} samples") + X = X[: args.max_samples] + feats = feats[: args.max_samples] + y = y[: args.max_samples] + metadata = metadata[: args.max_samples] + + # Convert features to proper format + try: + if feats.dtype == object and len(feats) > 0 and isinstance(feats[0], dict): + feat_keys = sorted(feats[0].keys()) + feats_np = np.array([[f[k] for k in feat_keys] for f in feats]) + else: + feats_np = np.array([list(f.values()) for f in feats]) + except: + feats_np = feats + + X = X.astype("float32") + feats_np = feats_np.astype("float32") + y = y.astype("int32") + + # Remove samples with NaN values + print(feats_np.shape) + mask = np.isnan(X).any(axis=(1, 2, 3)) | np.isnan(feats_np).any(axis=1) + if mask.any(): + print(f"Removing {mask.sum()} samples with NaN values") + X = X[~mask] + feats_np = feats_np[~mask] + y = y[~mask] + metadata = metadata[~mask] + + print(f"Final data shapes - X: {X.shape}, feats: {feats_np.shape}, y: {y.shape}") + print(f"Final class distribution: {np.bincount(y)}") + + # Split data into train/val/test + print(f"\nSplitting data...") + print(f"Test size: {args.test_size}, Validation size: {args.val_size}") + + # First split: separate test set + if args.test_size > 0: + ( + X_train_val, + X_test, + feats_train_val, + feats_test, + y_train_val, + y_test, + metadata_train_val, + metadata_test, + ) = train_test_split( + X, + feats_np, + y, + metadata, + test_size=args.test_size, + stratify=y, + random_state=args.random_state, + shuffle=True, + ) + else: + X_train_val = X + feats_train_val = feats_np + y_train_val = y + metadata_train_val = metadata + X_test = np.empty((0,) + X.shape[1:]) + feats_test = np.empty((0, feats_np.shape[1])) + y_test = np.empty((0,)) + metadata_test = [] + + # Second split: separate train and validation + ( + X_train, + X_val, + feats_train, + feats_val, + y_train, + y_val, + metadata_train, + metadata_val, + ) = train_test_split( + X_train_val, + feats_train_val, + y_train_val, + metadata_train_val, + test_size=args.val_size, + stratify=y_train_val, + random_state=args.random_state, + shuffle=True, + ) + + print(f"Train set: {len(X_train)} samples") + print(f"Validation set: {len(X_val)} samples") + print(f"Test set: {len(X_test)} samples") + + # Clean up memory + del ( + X, + feats, + feats_np, + y, + metadata, + X_train_val, + feats_train_val, + y_train_val, + metadata_train_val, + ) + gc.collect() + + # Create model + img_shape = X_train[0].shape + num_features = feats_train.shape[1] + + print(f"\nCreating model...") + print(f"Image shape: {img_shape}") + print(f"Number of features: {num_features}") + + model = create_hybrid_model(img_shape, num_features) + model.summary() + + # Setup callbacks + early_stopping = F1EarlyStopping( + precision_key="val_precision", + recall_key="val_recall", + patience=args.patience, + restore_best_weights=True, + ) + + callbacks = [ + early_stopping, + tf.keras.callbacks.ModelCheckpoint( + os.path.join(args.output_dir, f"{args.model_name}_best_precision.h5"), + save_best_only=True, + monitor="val_precision", + mode="max", + verbose=1, + ), + tf.keras.callbacks.ModelCheckpoint( + os.path.join(args.output_dir, f"{args.model_name}_best_recall.h5"), + save_best_only=True, + monitor="val_recall", + mode="max", + verbose=1, + ), + tf.keras.callbacks.ReduceLROnPlateau( + monitor="val_loss", factor=0.5, patience=5, min_lr=1e-6, verbose=1 + ), + ] + + # Calculate class weights + # class_weights = compute_class_weight( + # "balanced", classes=np.unique(y_train), y=y_train + # ) + class_weight_dict = {0: 1.0, 1: args.class_weight_pos} + + print(f"\nClass weights: {class_weight_dict}") + + # Train model + print(f"\nStarting training...") + history = model.fit( + [X_train, feats_train], + y_train, + epochs=args.epochs, + batch_size=args.batch_size, + validation_data=([X_val, feats_val], y_val), + callbacks=callbacks, + class_weight=class_weight_dict, + verbose=1, + ) + + # Save final model + final_model_path = os.path.join(args.output_dir, f"{args.model_name}_final.h5") + model.save(final_model_path) + print(f"\nFinal model saved to: {final_model_path}") + + # Plot training history + plot_training_history(history, args.output_dir) + + # Evaluate model + print(f"\nEvaluating model on test set...") + y_pred_prob, y_pred = evaluate_model( + model, X_test, feats_test, y_test, args.output_dir + ) + + # Save test set predictions and metadata + test_results_path = os.path.join(args.output_dir, "test_predictions.npz") + np.savez( + test_results_path, + X=X_test, + feats=feats_test, + y_true=y_test, + y_pred_prob=y_pred_prob, + y_pred=y_pred, + metadata=metadata_test, + ) + print(f"Test predictions saved to: {test_results_path}") + + print(f"\nTraining completed successfully!") + print(f"All outputs saved to: {args.output_dir}") + + +if __name__ == "__main__": + main() diff --git a/RuBR/model/callbacks.py b/RuBR/model/callbacks.py new file mode 100644 index 00000000..f48c60ef --- /dev/null +++ b/RuBR/model/callbacks.py @@ -0,0 +1,44 @@ +import tensorflow as tf + + +class F1EarlyStopping(tf.keras.callbacks.Callback): + def __init__( + self, + precision_key, + recall_key, + patience=10, + restore_best_weights=True, + ): + super().__init__() + self.precision_key = precision_key + self.recall_key = recall_key + self.patience = patience + self.restore_best_weights = restore_best_weights + self.best_f1 = 0 + self.best_weights = None + self.wait = 0 + + def on_epoch_end(self, epoch, logs=None): + logs = logs or {} + current_precision = logs.get(self.precision_key, 0) + current_recall = logs.get(self.recall_key, 0) + if current_precision + current_recall == 0: + current_f1 = 0 + else: + current_f1 = 2 * current_precision * current_recall / (current_precision + current_recall) + + if current_f1 > self.best_f1: + self.best_f1 = current_f1 + self.wait = 0 + if self.restore_best_weights: + self.best_weights = self.model.get_weights() + return + + self.wait += 1 + if self.wait >= self.patience: + print(f"\nEarly stopping triggered after {epoch + 1} epochs") + print(f"Best F1: {self.best_f1:.4f}") + self.model.stop_training = True + if self.restore_best_weights and self.best_weights is not None: + print("Restoring best weights...") + self.model.set_weights(self.best_weights) diff --git a/RuBR/model/dann_model.py b/RuBR/model/dann_model.py new file mode 100644 index 00000000..441c72c2 --- /dev/null +++ b/RuBR/model/dann_model.py @@ -0,0 +1,80 @@ +""" +Domain Adversarial Neural Network (DANN) Model Architecture + +This module defines the DANN architecture for transient detection with domain adaptation. +""" + +import tensorflow as tf +from tensorflow.keras import layers, Model +from model.layers import ( + rot90_k1, + rot90_k2, + rot90_k3, + gradient_reversal, + GradientReversalLayer, + image_encoder, +) + + +def create_dann_model(img_shape, num_features, lambda_domain=1.0): + """ + Create a Domain Adversarial Neural Network model. + + Args: + img_shape (tuple): Shape of input images (H, W, C) + num_features (int): Number of tabular features + lambda_domain (float): Weight for gradient reversal layer + + Returns: + tf.keras.Model: DANN model with multiple outputs + """ + # Inputs + img_input = layers.Input(shape=img_shape, name="image_input") + feat_input = layers.Input(shape=(num_features,), name="feature_input") + + # Feature extractor (shared) + img_features = image_encoder(img_input, img_shape, mode="mean") + feat_features = layers.Dense(32, activation="relu", name="feat_encoder")(feat_input) + + # Combine features + combined_features = layers.Concatenate(name="combined_features")([img_features, feat_features]) + + # Label classifier (for transient detection) + label_classifier = layers.Dense(128, activation="relu", name="label_fc1")(combined_features) + label_classifier = layers.Dropout(0.3)(label_classifier) + label_classifier = layers.Dense(64, activation="relu", name="label_fc2")(label_classifier) + label_output = layers.Dense(1, activation="sigmoid", name="label_output")(label_classifier) + + # Domain classifier (for domain adaptation) + # Apply gradient reversal layer + domain_features = GradientReversalLayer(lambda_=lambda_domain, name="gradient_reversal")(combined_features) + domain_classifier = layers.Dense(128, activation="relu", name="domain_fc1")(domain_features) + domain_classifier = layers.Dropout(0.3)(domain_classifier) + domain_classifier = layers.Dense(64, activation="relu", name="domain_fc2")(domain_classifier) + domain_output = layers.Dense(1, activation="sigmoid", name="domain_output")(domain_classifier) + + # Create model with multiple outputs + model = Model( + inputs=[img_input, feat_input], + outputs=[label_output, domain_output], + name="DANN" + ) + + # Compile with multiple losses + model.compile( + optimizer=tf.keras.optimizers.Adam(learning_rate=0.001), + loss={ + "label_output": "binary_crossentropy", + "domain_output": "binary_crossentropy" + }, + loss_weights={ + "label_output": 1.0, + "domain_output": 1.0 + }, + metrics={ + "label_output": ["accuracy", "precision", "recall"], + "domain_output": ["accuracy"] + } + ) + + return model diff --git a/RuBR/model/data.py b/RuBR/model/data.py new file mode 100644 index 00000000..a3bbc300 --- /dev/null +++ b/RuBR/model/data.py @@ -0,0 +1,39 @@ +import numpy as np + + +def load_dataset(data_path, mmap=False, allow_npy_dict=True): + print(f"Loading data from {data_path}") + if allow_npy_dict and data_path.endswith(".npy"): + data = np.load(data_path, allow_pickle=True).item() + else: + mmap_mode = "r" if mmap else None + data = np.load(data_path, allow_pickle=True, mmap_mode=mmap_mode) + + if isinstance(data, np.lib.npyio.NpzFile): + X = data["X"] + feats = data["feats"] + y = data["y"] + metadata = data["metadata"] + else: + X = data["X"] + feats = data["feats"] + y = data["y"] + metadata = data["metadata"] + + if X.ndim == 4 and X.shape[1] == 3 and X.shape[-1] != 3: + X = X.transpose(0, 2, 3, 1) + + if getattr(feats, "dtype", None) == object: + feats = np.asarray([list(f.values()) if hasattr(f, "values") else f for f in feats]) + + y = np.asarray(y) + + print(f"Loaded data shapes - X: {X.shape}, feats: {feats.shape}, y: {y.shape}") + print(f"Class distribution: {np.bincount(y.astype(int))}") + return X, feats, y, metadata + + +def normalize_arrays(X, feats): + X_norm = (X - X.mean(axis=(0, 1, 2), keepdims=True)) / (X.std(axis=(0, 1, 2), keepdims=True) + 1e-6) + feats_norm = (feats - feats.mean(axis=0)) / (feats.std(axis=0) + 1e-6) + return X_norm, feats_norm diff --git a/RuBR/model/layers.py b/RuBR/model/layers.py new file mode 100644 index 00000000..2a263412 --- /dev/null +++ b/RuBR/model/layers.py @@ -0,0 +1,88 @@ +import tensorflow as tf +from tensorflow.keras import layers, Model + + +def rot90_k1(x): + return tf.image.rot90(x, k=1) + + +def rot90_k2(x): + return tf.image.rot90(x, k=2) + + +def rot90_k3(x): + return tf.image.rot90(x, k=3) + + +@tf.custom_gradient +def gradient_reversal(x, lambda_): + def grad(dy): + return -lambda_ * dy, None + + return x, grad + + +class GradientReversalLayer(layers.Layer): + def __init__(self, lambda_=1.0, **kwargs): + super().__init__(**kwargs) + self.lambda_ = lambda_ + + def call(self, x): + return gradient_reversal(x, self.lambda_) + + def get_config(self): + config = super().get_config() + config.update({"lambda_": self.lambda_}) + return config + + +def image_encoder( + inputs, + img_shape, + mode="mean", + conv_kernel_initializer=None, + dense_kernel_initializer=None, +): + x1 = inputs + x2 = layers.Lambda(rot90_k1, output_shape=img_shape)(inputs) + x3 = layers.Lambda(rot90_k2, output_shape=img_shape)(inputs) + x4 = layers.Lambda(rot90_k3, output_shape=img_shape)(inputs) + + enc_input = layers.Input(shape=img_shape) + + conv_kwargs = {"kernel_size": 3, "padding": "same"} + if conv_kernel_initializer is not None: + conv_kwargs["kernel_initializer"] = conv_kernel_initializer + + x = layers.Conv2D(32, **conv_kwargs)(enc_input) + x = layers.BatchNormalization()(x) + x = layers.ReLU()(x) + x = layers.MaxPooling2D(pool_size=2, strides=2)(x) + + x = layers.Conv2D(64, **conv_kwargs)(x) + x = layers.BatchNormalization()(x) + x = layers.ReLU()(x) + x = layers.MaxPooling2D(pool_size=2, strides=2)(x) + + x = layers.Conv2D(128, **conv_kwargs)(x) + x = layers.BatchNormalization()(x) + x = layers.ReLU()(x) + x = layers.MaxPooling2D(pool_size=2, strides=2)(x) + x = layers.Flatten()(x) + + dense_units = x.shape[-1] // 2 + dense_kwargs = {"activation": "relu"} + if dense_kernel_initializer is not None: + dense_kwargs["kernel_initializer"] = dense_kernel_initializer + x = layers.Dense(dense_units, **dense_kwargs)(x) + + encoder = Model(enc_input, x, name="shared_encoder") + + e1 = encoder(x1) + e2 = encoder(x2) + e3 = encoder(x3) + e4 = encoder(x4) + + if mode == "concat": + return layers.Concatenate(axis=1)([e1, e2, e3, e4]) + return layers.Average()([e1, e2, e3, e4]) diff --git a/RuBR/model/metadata.py b/RuBR/model/metadata.py new file mode 100644 index 00000000..93d005e3 --- /dev/null +++ b/RuBR/model/metadata.py @@ -0,0 +1,18 @@ +import os +import pandas as pd + + +def get_transient_magnitude(id, jid_folder): + file_name = os.path.join(jid_folder, "combined_truth_table.csv") + if not os.path.exists(file_name): + raise FileNotFoundError(f"File not found: {file_name}") + + df = pd.read_csv(file_name) + id = str(id) + if not id.endswith("_ou"): + id += "_ou" + row = df[df["id"] == id] + if row.empty: + raise ValueError(f"ID {id} not found in {file_name}") + mag = row.iloc[0]["mag"] + row.iloc[0]["zpt"] + return mag diff --git a/RuBR/requirements.txt b/RuBR/requirements.txt new file mode 100644 index 00000000..cacda9f7 --- /dev/null +++ b/RuBR/requirements.txt @@ -0,0 +1,6 @@ +numpy==1.26.4 +pandas==2.2.2 +scikit-learn==1.4.2 +matplotlib==3.8.4 +tensorflow==2.15.1 +scienceplots==2.1.1