source/eval_common.py

"""Shared protocol and I/O for matched random/scaffold checkpoint comparisons."""
from pathlib import Path
import hashlib, json, os, sys, time
AUDIT=Path('/marimo/rem3di_training/corrected_20261027/evaluation/paper_audit')
ROOT=AUDIT.parent.parent
WORK=Path(__file__).resolve().parent.parent
OUT=WORK/'evaluation'
ORIGINAL=AUDIT/'author_checkpoint/exp138/polar'
POST=AUDIT/'author_pcqm500k_5ep'
sys.path[:0]=[str(AUDIT),str(AUDIT/'author_legacy_source'),'/marimo/Rem3Di']
os.environ.update(CUDA_VISIBLE_DEVICES='0',OMP_NUM_THREADS='8',MKL_NUM_THREADS='8',WANDB_MODE='online')
from dotenv import load_dotenv
load_dotenv('/marimo/.env',override=False)
import numpy as np
import torch
torch.serialization.add_safe_globals([slice])
from torch import nn
from torch.nn.utils.rnn import pad_sequence
import yaml
from sklearn.model_selection import KFold,GroupKFold,GroupShuffleSplit,train_test_split
from rdkit import Chem
from rdkit.Chem.Scaffolds import MurckoScaffold
from threedscriptors.configuration.architecture_config import EncoderOnlyArchitectureConfig
from threedscriptors.data_handling.sample import Sample

torch.set_num_threads(8)
SEED=42026
DATASETS=['esol','freesolv']
CHECKPOINTS=list(range(1,11))
SPLITS=['random','scaffold']
PROTOCOL={
 'checkpoints':CHECKPOINTS,'initial_checkpoint':'Selected PCQM epoch 2, followed by 10 COSMO epochs',
 'datasets':DATASETS,'outer_cv':'3-fold; random KFold(shuffle=True, random_state=42026), scaffold GroupKFold by achiral Bemis-Murcko scaffold',
 'scaffold_empty':'All acyclic molecules share the empty scaffold; kept together, so scaffold folds can be unequal.',
 'ridge_probe':'StandardScaler + Ridge; alpha selected on 3 inner folds from 1e-6..1e6 (13 values); group-disjoint inner CV for scaffold',
 'kumo':'Kumo-Tabular large, installed sdm version; 4 estimators; all 640 descriptors; seed=16000+fold; default recipe with width limit raised to 640',
 'finetuning':{'max_epochs':100,'patience':15,'backbone_lr':1e-5,'head_lr':1e-3,'weight_decay':1e-4,'grad_clip':1.,
  'head':'640 -> 256 -> 1, ReLU, dropout 0.2','validation_fraction':0.2,'scaffold_validation':'Closest to 20% of rows among 512 seeded scaffold-group prefix partitions; no labels used',
  'trainable':'checkpoint encoder + learned atomic/chiral and geometric preprocessors + regression head; MACE frozen',
  'selection':'Best inner validation RMSE; outer test evaluated once after selection; no outer test-based tuning',
  'scaling':'Descriptor mean/std and target mean/std fit on inner training rows only; checkpoint atomic normalization buffers preserved',
  'seeds':'One matched initialization/shuffle seed per outer fold, identical across checkpoints',
  'refit':'No refit on validation rows after early stopping; report actual train/validation counts'},
 'metrics':['rmse','mae','r2'],'total_probe_fold_results':240,'total_finetuning_fold_results':120,
 'comparison':'Identical saved row indices across checkpoints and methods; probes use all outer training rows; finetuning reserves inner validation.',
 'dataset_scope':'Existing retained/geometry-valid subsets: ESOL 1108 and FreeSolv 625; no new filtering.',
}

def write(path,value):
 path.parent.mkdir(parents=True,exist_ok=True)
 tmp=path.with_suffix(path.suffix+'.partial')
 tmp.write_text(json.dumps(value,indent=2,allow_nan=False));tmp.replace(path)

def sha(path):
 h=hashlib.sha256()
 with path.open('rb') as f:
  for block in iter(lambda:f.read(8<<20),b''):h.update(block)
 return h.hexdigest()

def status(phase,**kw):
 write(OUT/'status.json',dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**kw))
 print('STATUS',phase,json.dumps(kw),flush=True)

def checkpoint(epoch):return WORK/'checkpoints'/f'epoch_{epoch:02d}'

def load_backbone(epoch):
 d=yaml.safe_load((ORIGINAL/'post_training_architecture_config.yaml').read_text())
 d.pop('decoder_config',None);d['kind']='encoder_only';d['mace_config']=None
 cfg=EncoderOnlyArchitectureConfig.model_validate(d)
 pre=cfg._build_preprocessor();enc=cfg._build_encoder()
 # Cast before loading to preserve the trained double-precision chiral weights.
 pre.atomic_preprocessor.double();pre.geometric_preprocessor.float();enc.float()
 for name,module in [('encoder.pth',enc),('atomic_preprocessor.pth',pre.atomic_preprocessor),('geometric_preprocessor.pth',pre.geometric_preprocessor)]:
  module.load_state_dict(torch.load(checkpoint(epoch)/name,map_location='cpu',weights_only=True),strict=True)
 return pre.cuda(),enc.cuda()

def arrays(dataset):
 z=np.load(ROOT/'evaluation/cache'/f'{dataset}.npz')
 y=z['y'][:,0].astype(np.float64);smiles=z['smiles']
 assert np.isfinite(y).all()
 canonical=[Chem.MolToSmiles(Chem.MolFromSmiles(str(s)),isomericSmiles=True) for s in smiles]
 assert len(set(canonical))==len(canonical), 'Duplicate canonical molecules require grouped random splitting'
 groups=np.array([MurckoScaffold.MurckoScaffoldSmiles(mol=Chem.MolFromSmiles(str(s)),includeChirality=False) for s in smiles])
 return y,smiles,groups

def prepare_splits():
 for dataset in DATASETS:
  y,smiles,groups=arrays(dataset);indices=np.arange(len(y))
  for scheme in SPLITS:
   path=OUT/'splits'/f'{dataset}_{scheme}.json'
   folds=[]
   outer=KFold(3,shuffle=True,random_state=SEED).split(indices) if scheme=='random' else GroupKFold(3).split(indices,groups=groups)
   for fold,(tr,te) in enumerate(outer):
    if scheme=='random':fit,val=train_test_split(tr,test_size=.2,random_state=16000+fold)
    else:
     unique,counts=np.unique(groups[tr],return_counts=True)
     rng=np.random.default_rng(16000+fold);candidates=[]
     for _ in range(512):
      order=rng.permutation(len(unique));k=int(np.argmin(abs(np.cumsum(counts[order])-.2*len(tr))))+1
      mask=np.isin(groups[tr],unique[order[:k]])
      it,iv=np.flatnonzero(~mask),np.flatnonzero(mask)
      if len(iv)>=10 and len(it)>=20:candidates.append((it,iv))
     assert candidates, 'No useful scaffold validation partition'
     it,iv=min(candidates,key=lambda pair:abs(len(pair[1])/len(tr)-.2))
     fit,val=tr[it],tr[iv]
    assert not set(tr)&set(te) and set(fit)|set(val)==set(tr) and not set(fit)&set(val)
    if scheme=='scaffold':
     assert not set(groups[tr])&set(groups[te])
     assert not set(groups[fit])&set(groups[val])
    folds.append(dict(fold=fold,train=tr.tolist(),test=te.tolist(),fit=fit.tolist(),validation=val.tolist()))
   assert sorted(i for f in folds for i in f['test'])==indices.tolist()
   result=dict(dataset=dataset,scheme=scheme,smiles_sha256=hashlib.sha256('\n'.join(smiles).encode()).hexdigest(),
               n=len(y),n_scaffolds=len(set(groups)),folds=folds)
   if path.exists() and (OUT/'results.csv').exists():assert json.loads(path.read_text())==result, 'Saved split protocol changed'
   else:write(path,result)
   print('SPLITS',dataset,scheme,[(len(f['fit']),len(f['validation']),len(f['test'])) for f in folds],flush=True)

def get_splits(dataset,scheme):return json.loads((OUT/'splits'/f'{dataset}_{scheme}.json').read_text())['folds']

def batch_indices(indices,lengths,shuffle=None,max_size=48,atom_cap=4000):
 order=np.asarray(indices).copy()
 if shuffle is not None:shuffle.shuffle(order)
 # Sorting within buckets controls padding; random bucket membership changes each epoch.
 for start in range(0,len(order),256):
  chunk=sorted(order[start:start+256],key=lambda i:lengths[i]);batch=[];atoms=0
  for i in chunk:
   if batch and (atoms+lengths[i]>atom_cap or len(batch)>=max_size):
    yield batch;batch=[];atoms=0
   batch.append(int(i));atoms+=lengths[i]
  if batch:yield batch

def sample(ids,data):
 f=pad_sequence([data['features'][i] for i in ids],batch_first=True).cuda()
 p=pad_sequence([data['positions'][i] for i in ids],batch_first=True).cuda()
 lengths=torch.tensor([data['lengths'][i] for i in ids],device='cuda')
 mask=torch.arange(f.shape[1],device='cuda')[None]>=lengths[:,None]
 return Sample(embeddings=f,atomic_positions=p,padding_mask=mask)

def metric(y,p):
 from sklearn.metrics import mean_squared_error,mean_absolute_error,r2_score
 assert np.isfinite(p).all()
 return dict(rmse=float(np.sqrt(mean_squared_error(y,p))),mae=float(mean_absolute_error(y,p)),r2=float(r2_score(y,p)))