File size: 3,347 Bytes
d561d1e | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | """
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__))) # EMR Classifier/
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() |