| """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() |
| |
| 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) |
| |
| 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))) |
|
|
