| """ |
| Inference: rank all unknown spectra by EMR probability using the 5-model ensemble. |
| |
| Input : Data Parquet/unknown_data_after_noise_filter_cropped.parquet (sclk + spectrum_cropped) |
| Models: EMR Classifier/models/emr_model_seed_*.pth |
| Output: EMR Classifier/predictions/emr_candidates_unknown.csv (all rows, sorted desc) |
| EMR Classifier/predictions/emr_top100.csv (top 100 for manual review) |
| """ |
| import os |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| from torch.utils.data import DataLoader, TensorDataset |
|
|
| from train_emr_classifier import Spectra1DCNN, TARGET_LEN, minmax |
|
|
| ROOT = os.path.dirname(os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) |
| HERE = os.path.dirname(os.path.dirname(os.path.abspath(__file__))) |
| DATA_PATH = os.path.join(ROOT, 'Data Parquet', 'unknown_data_after_noise_filter_cropped.parquet') |
| MODELS_DIR = os.path.join(HERE, 'models') |
| OUT_DIR = os.path.join(HERE, 'predictions') |
| os.makedirs(OUT_DIR, exist_ok=True) |
|
|
| BATCH_SIZE = 256 |
|
|
|
|
| def main(): |
| device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') |
| print(f'Using device: {device}') |
|
|
| print('Loading unknown data (sclk + spectrum_cropped)...') |
| df = pd.read_parquet(DATA_PATH, columns=['sclk', 'spectrum_cropped']) |
| X = np.stack(df['spectrum_cropped'].to_numpy()).astype(np.float32) |
| X = minmax(X) |
| print(f' {X.shape}') |
|
|
| model_files = sorted(f for f in os.listdir(MODELS_DIR) if f.startswith('emr_model_seed_')) |
| models = [] |
| for mf in model_files: |
| m = Spectra1DCNN(TARGET_LEN) |
| m.load_state_dict(torch.load(os.path.join(MODELS_DIR, mf), map_location=device)) |
| m.to(device).eval() |
| models.append(m) |
| print(f'Loaded {len(models)} ensemble models: {model_files}') |
|
|
| loader = DataLoader(TensorDataset(torch.from_numpy(X).unsqueeze(1)), |
| batch_size=BATCH_SIZE, shuffle=False, num_workers=0) |
|
|
| preds = np.zeros(len(X), dtype=np.float32) |
| with torch.no_grad(): |
| for i, (bx,) in enumerate(loader): |
| bx = bx.to(device) |
| acc = torch.zeros(bx.size(0), device=device) |
| for m in models: |
| acc += torch.sigmoid(m(bx)).squeeze(-1) |
| acc /= len(models) |
| s = i * BATCH_SIZE |
| preds[s:s + bx.size(0)] = acc.cpu().numpy() |
| if (i + 1) % 40 == 0: |
| print(f' {(i + 1) * BATCH_SIZE:,} / {len(X):,}') |
|
|
| df['emr_probability'] = preds |
| df_sorted = (df.drop(columns=['spectrum_cropped']) |
| .sort_values('emr_probability', ascending=False) |
| .reset_index(names='original_parquet_index')) |
|
|
| out_all = os.path.join(OUT_DIR, 'emr_candidates_unknown.csv') |
| df_sorted.to_csv(out_all, index=False) |
| print(f'Wrote {out_all} ({len(df_sorted):,} rows)') |
|
|
| out_top = os.path.join(OUT_DIR, 'emr_top100.csv') |
| df_sorted.head(100).to_csv(out_top, index=False) |
| print(f'Wrote {out_top}') |
|
|
| print(f'\nProbability distribution: ' |
| f'mean={preds.mean():.4f} p50={np.percentile(preds, 50):.4f} ' |
| f'p99={np.percentile(preds, 99):.4f} p99.9={np.percentile(preds, 99.9):.4f} ' |
| f'max={preds.max():.4f}') |
| for thr in (0.5, 0.75, 0.9): |
| print(f' rows with emr_probability >= {thr}: {(preds >= thr).sum():,}') |
|
|
|
|
| if __name__ == '__main__': |
| main() |
