website/src/data/models.ts

// Real HydroPD Step-4 models, loaded from generated JSON (see scripts/extract_data.py).
// 28 usable-tier hydrolysate models + 3 Arabidopsis transfer models (species + enzyme + MS protocol).
import modelsJson from './generated/models.json'

export interface RealModel {
  code: string
  hcode: string
  species: string
  speciesKey: string
  enzyme: string
  msProtocol: string
  taxId: string
  proteome: string
  proteomeApprox: boolean
  nPos: number
  bestModel: string
  bestFeatureSet: string
  gridAuroc: number | null
  tunedCvAuroc: number | null
  calibAuroc: number | null
  calibAuprc: number | null
  calibMcc: number | null
  calibF1: number | null
  brier: number | null
  ece: number | null
  headlineAuroc: number | null
  paperTitle: string
  paperLink: string
  rawMaterial: string
  flags: string[]
}

export const MODELS: RealModel[] = modelsJson as RealModel[]

// Display labels for the config vocabulary.
export const FEATURE_SET_LABEL: Record<string, string> = {
  phys: 'Physicochemical (14 descriptors)',
  pepbert: 'PepBERT embeddings',
  esm2: 'ESM2-t12-35M embeddings',
  pep_phys: 'PepBERT ⊕ physicochemical',
  esm2_phys: 'ESM2 ⊕ physicochemical',
}

export const CLASSIFIER_LABEL: Record<string, string> = {
  xgb: 'XGBoost',
  rf: 'Random forest',
  mlp: 'Neural net (MLP)',
}

// Shared training config (production defaults) shown in model details.
export const TRAINING_META = {
  cv: '5-fold StratifiedGroupKFold, grouped by lead protein',
  negatives: 'k-mer negatives (seed 24, 1:1 balance)',
  primaryMetric: 'AUROC',
}

export function featureSetLabel(id: string): string {
  return FEATURE_SET_LABEL[id] ?? id
}
export function classifierLabel(id: string): string {
  return CLASSIFIER_LABEL[id] ?? id
}

/** Headline AUROC as a number (calibrated, else grid); null when unavailable. */
export function displayAurocValue(m: RealModel): number | null {
  return m.calibAuroc ?? m.gridAuroc
}

/** Headline AUROC for display; null for degenerate models (e.g. H13A). */
export function displayAuroc(m: RealModel): string {
  const v = displayAurocValue(m)
  return v == null ? 'n/a' : v.toFixed(3)
}

// Out-of-the-box Pfly (wilhelm-lab dlomix DetectabilityModel), sequence-only and not
// calibrated to any single dataset. Selectable on the prediction page via modelCode
// 'pfly_oob'; the backend routes it to inference.predict_pfly (no joblib).
export const PFLY_OOB: RealModel = {
  code: 'pfly_oob',
  hcode: '',
  species: 'Any species (sequence-only)',
  speciesKey: 'pfly',
  enzyme: 'any',
  msProtocol: 'any',
  taxId: '',
  proteome: '',
  proteomeApprox: false,
  nPos: 0,
  bestModel: 'pfly',
  bestFeatureSet: 'sequence',
  gridAuroc: null,
  tunedCvAuroc: null,
  calibAuroc: null,
  calibAuprc: null,
  calibMcc: null,
  calibF1: null,
  brier: null,
  ece: null,
  headlineAuroc: null,
  paperTitle: 'Wilhelm-lab dlomix DetectabilityModel (Pfly), pretrained',
  paperLink: 'https://github.com/wilhelm-lab/dlomix',
  rawMaterial: '',
  flags: ['Sequence-only, pretrained; not calibrated to your dataset.'],
}

/** Look up a model by code across the 30 context models plus the Pfly option. */
export function findModel(code: string): RealModel {
  if (code === PFLY_OOB.code) return PFLY_OOB
  return MODELS.find((m) => m.code === code) ?? MODELS[0]
}