| """Matched COSMO epoch 10 / Monroe + Kumo Large on official OpenADMET splits.""" |
| from pathlib import Path |
| import sys |
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| import eval_common as common |
| from eval_common import WORK, sha, write |
| import concurrent.futures |
| import fcntl |
| import gc |
| import hashlib |
| import importlib.metadata |
| import json |
| import multiprocessing as mp |
| import os |
| import shutil |
| import subprocess |
| import time |
| import traceback |
|
|
| import numpy as np |
| import pandas as pd |
| import torch |
| import wandb |
| from huggingface_hub import HfApi, hf_hub_download, CommitOperationAdd |
| from rdkit import Chem |
|
|
| OUT = WORK/'openadmet_evaluation' |
| REPOSITORY = 'Queimo/rem3di-pcqm-epoch2-cosmo' |
| DATASET = 'openadmet/openadmet-expansionrx-challenge-data' |
| DATA_REVISION = '6b898ccc43d10d25b230fb09e22a6e30c30022b5' |
| MONROE = WORK/'external/Monroe' |
| MONROE_REVISION = 'de075eabe98e4adbbedbc86ca9e877f4bfd9b629' |
| MONROE_SHA = 'df996e8e98b12a3cda5fe5acbef3e1e9a4ea31946d3736461cd07e53a5bf205a' |
| SEED = 42 |
| PROTOCOL = dict( |
| dataset=DATASET, dataset_revision=DATA_REVISION, |
| reference_notebook=f'https://github.com/blazejba/Monroe/blob/{MONROE_REVISION}/examples/openadmet.ipynb', |
| split='Original organizer train/test CSVs: 5326/2282; no new CV or test-based selection', |
| representations=['cosmo_epoch10','monroe'], methods=['kumo_large'], expected_endpoint_results=18, |
| labels='Raw numeric endpoint values, matching the reference notebook; only missing measurements excluded per endpoint', |
| geometry='One author Monroe ETKDG/MMFF94s/UFF conformer per unique molecule, cached and shared exactly by both encoders', |
| encoder='Frozen released Monroe encoder (720 dimensions) and frozen COSMO epoch 10 backbone (640 dimensions); strict state loading', |
| monroe_revision=MONROE_REVISION, monroe_weights_sha256=MONROE_SHA, |
| kumo=dict(size='large',num_estimators=4,seed=SEED,max_columns=640,estimator_batch_size=1, |
| precision='float16 autocast',query_chunk_size=256, |
| preprocessing='Default recipe fitted only on measured training rows; processed context shapes recorded'), |
| metrics=['rmse','mae','r2','spearman'], comparison='Same molecules, conformers, endpoint masks, seeds and Kumo recipe; only representation differs') |
|
|
|
|
| def status(phase, **kw): |
| write(OUT/'status.json',dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**kw)) |
| print('OPENADMET_STATUS',phase,json.dumps(kw),flush=True) |
|
|
|
|
| def prepare_data(): |
| status('preparing_official_dataset') |
| parts = [] |
| for split in ['train','test']: |
| name = f'expansion_data_{split}.csv' |
| source = Path(hf_hub_download(DATASET,name,repo_type='dataset',revision=DATA_REVISION)) |
| target = OUT/'data'/name |
| if not target.exists(): shutil.copyfile(source,target) |
| assert sha(source)==sha(target) |
| parts.append(pd.read_csv(target)) |
| train,test = parts |
| assert len(train)==5326 and len(test)==2282 |
| endpoints = [c for c in train if c not in ['Molecule Name','SMILES']] |
| assert len(endpoints)==9 and list(train.columns)==list(test.columns) |
| smiles = [str(s) for s in pd.unique(pd.concat([train.SMILES,test.SMILES]))] |
| assert len(smiles)==7608 |
| canonical = [{Chem.MolToSmiles(Chem.MolFromSmiles(s),isomericSmiles=True) for s in df.SMILES} for df in parts] |
| assert not canonical[0]&canonical[1] |
| assert all(np.isfinite(df[e].dropna().to_numpy(dtype=float)).all() for df in parts for e in endpoints) |
| PROTOCOL['endpoints'] = endpoints |
| PROTOCOL['dataset_files_sha256'] = {p.name:sha(p) for p in (OUT/'data').glob('*.csv')} |
| PROTOCOL['checkpoint_sha256'] = {n:sha(common.checkpoint(10)/n) for n in |
| ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']} |
| manifest=json.loads((common.checkpoint(10)/'manifest.json').read_text()) |
| for name,digest in PROTOCOL['checkpoint_sha256'].items():assert digest==manifest['files'][name]['sha256'] |
| write(OUT/'protocol.json',PROTOCOL) |
| write(OUT/'data/measurement_counts.json', {e:dict(train=int(train[e].notna().sum()),test=int(test[e].notna().sum())) for e in endpoints}) |
| return train,test,endpoints,smiles |
|
|
|
|
| def restore_monroe(): |
| env={**os.environ,'GIT_LFS_SKIP_SMUDGE':'1'} |
| MONROE.parent.mkdir(exist_ok=True) |
| if not (MONROE/'.git').exists(): |
| subprocess.run(['git','clone','--depth','1','https://github.com/blazejba/Monroe.git',str(MONROE)],check=True,env=env) |
| revision=subprocess.check_output(['git','-C',str(MONROE),'rev-parse','HEAD'],text=True).strip() |
| if revision!=MONROE_REVISION: |
| subprocess.run(['git','-C',str(MONROE),'fetch','--depth','1','origin',MONROE_REVISION],check=True,env=env) |
| subprocess.run(['git','-C',str(MONROE),'checkout',MONROE_REVISION],check=True,env=env) |
| weights=MONROE/'checkpoint/weights.pt' |
| if not weights.exists() or sha(weights)!=MONROE_SHA: |
| temp=weights.with_suffix('.download') |
| subprocess.run(['curl','-fL','--retry','4','-o',str(temp),f'https://media.githubusercontent.com/media/blazejba/Monroe/{MONROE_REVISION}/checkpoint/weights.pt'],check=True) |
| assert sha(temp)==MONROE_SHA |
| temp.replace(weights) |
| sys.path.insert(0,str(MONROE)) |
|
|
|
|
| def graph_job(item): |
| index,smiles=item |
| from monroe.eval.embed import _featurize_one |
| from rdkit import rdBase |
| rdBase.SeedRandomNumberGenerator(SEED+index) |
| return _featurize_one(smiles) |
|
|
|
|
| def prepare_graphs(smiles): |
| import pyarrow as pa |
| import pyarrow.parquet as pq |
| path=OUT/'features/shared_graphs.parquet' |
| if path.exists(): |
| records=pq.read_table(path).to_pylist() |
| assert [r['smiles'] for r in records]==smiles |
| return [(r.pop('smiles'),r) for r in records] |
| status('generating_shared_conformers',total=len(smiles)) |
| with concurrent.futures.ProcessPoolExecutor(max_workers=8,mp_context=mp.get_context('fork')) as pool: |
| built=list(pool.map(graph_job,enumerate(smiles),chunksize=8)) |
| missing=[smiles[i] for i,v in enumerate(built) if v is None] |
| write(OUT/'data/coverage.json',dict(requested=len(smiles),missing=missing)) |
| assert not missing, f'Do not silently shrink the official split: {missing[:8]}' |
| assert [s for s,_ in built]==smiles |
| records=[dict(smiles=s,**{k:np.asarray(g[k]).tolist() for k in |
| ['node_float','node_codes','edge_index','edge_codes','pos_rdkit']}) for s,g in built] |
| pq.write_table(pa.Table.from_pylist(records),path,compression='zstd') |
| write(OUT/'features/shared_geometry_manifest.json',dict(rows=len(smiles),sha256=sha(path), |
| all_z_zero_layouts=sum(bool(np.all(np.asarray(g['pos_rdkit'])[:,2]==0)) for _,g in built), |
| source_revision=MONROE_REVISION)) |
| return built |
|
|
|
|
| def monroe_embeddings(built,smiles): |
| path=OUT/'features/monroe.npz' |
| if path.exists(): |
| z=np.load(path);assert z['smiles'].tolist()==smiles;return z['X'] |
| status('embedding_monroe',total=len(built)) |
| from monroe.eval.embed import to_pyg |
| from monroe.model.ckpt import _build_encoder,_load_state_dict |
| from torch_geometric.data import Batch |
| hp=json.loads((MONROE/'checkpoint/config.json').read_text()) |
| encoder=_build_encoder(hp).eval() |
| state=_load_state_dict(MONROE/'checkpoint/weights.pt',device='cpu') |
| encoder.load_state_dict({k[len('encoder.'):]:v for k,v in state.items() if k.startswith('encoder.')},strict=True) |
| del state |
| encoder=encoder.cuda().requires_grad_(False) |
| X=np.empty((len(built),720),dtype=np.float32) |
| with torch.inference_mode(): |
| for start in range(0,len(built),8): |
| batch=Batch.from_data_list([to_pyg(g) for _,g in built[start:start+8]]).to('cuda') |
| emb,_=encoder(batch);X[start:start+len(emb)]=emb.float().cpu().numpy() |
| if start%512==0:status('embedding_monroe',completed=start+len(emb),total=len(built)) |
| assert np.isfinite(X).all() |
| np.savez_compressed(path,X=X,smiles=np.asarray(smiles),weights_sha256=MONROE_SHA) |
| del encoder;gc.collect();torch.cuda.empty_cache() |
| return X |
|
|
|
|
| def cosmo_embeddings(built,smiles): |
| path=OUT/'features/cosmo_epoch10.npz' |
| if path.exists(): |
| z=np.load(path);assert z['smiles'].tolist()==smiles |
| assert json.loads(str(z['checkpoint_sha256']))==PROTOCOL['checkpoint_sha256'] |
| return z['X'] |
| from collections import Counter |
| from monroe.model.constants import NODE_FEAT_LIST_ONE_HOT |
| from remedi.data_handling.sample import Sample,yield_molecules_collate_fn |
| from build_author_mace_runtime import build_author_mace_runtime |
| from threedscriptors.data_handling.sample import Sample as LegacySample |
| columns=list(NODE_FEAT_LIST_ONE_HOT);j=columns.index('atomic_number') |
| vocab=NODE_FEAT_LIST_ONE_HOT['atomic_number'];ptable=Chem.GetPeriodicTable() |
| atoms=[];charge_counts={} |
| for s,g in built: |
| codes=np.asarray(g['node_codes'])[:,j] |
| assert (codes<len(vocab)).all() |
| z=np.asarray([ptable.GetAtomicNumber(vocab[int(code)]) for code in codes],dtype=np.int64) |
| pos=np.asarray(g['pos_rdkit'],dtype=np.float32) |
| charge=int(round(float(np.asarray(g['node_float'])[:,5].sum()))) |
| original=Chem.AddHs(Chem.MolFromSmiles(s)) |
| assert Counter(z.tolist())==Counter(a.GetAtomicNum() for a in original.GetAtoms()),s |
| assert charge==Chem.GetFormalCharge(original),s |
| assert int(np.asarray(g['node_float'])[:,0].sum())==0 and (int(z.sum())-charge)%2==0,s |
| assert pos.shape==(len(z),3) and np.isfinite(pos).all() |
| atoms.append((torch.from_numpy(z),torch.from_numpy(pos),charge)) |
| charge_counts[str(charge)]=charge_counts.get(str(charge),0)+1 |
| status('embedding_cosmo_epoch10',completed=0,total=len(built)) |
| pre,encoder=common.load_backbone(10) |
| pre.eval().requires_grad_(False);encoder.eval().requires_grad_(False) |
| mace=build_author_mace_runtime().cuda().eval().requires_grad_(False) |
| captures=[] |
| hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(lambda module,args:captures.append(args[0].detach())) |
| lengths=[len(a[0]) for a in atoms] |
| X=np.empty((len(built),640),dtype=np.float32);completed=0 |
| with torch.inference_mode(): |
| for ids in common.batch_indices(np.arange(len(built)),lengths,max_size=24,atom_cap=2400): |
| samples=[Sample(atomic_numbers=atoms[i][0],atomic_positions=atoms[i][1],total_charge=float(atoms[i][2]),multiplicity=1.) for i in ids] |
| batch=yield_molecules_collate_fn(samples).to_('cuda') |
| captures.clear();mace.preprocessor(batch) |
| assert len(captures)==1 |
| result=encoder(pre(LegacySample(embeddings=captures.pop().float(), |
| atomic_positions=batch.atomic_positions,padding_mask=batch.padding_mask))) |
| if not torch.is_tensor(result):result=result.flat |
| X[ids]=result.float().cpu().numpy();completed+=len(ids) |
| if completed%512<24:status('embedding_cosmo_epoch10',completed=completed,total=len(built)) |
| hook.remove();assert np.isfinite(X).all() |
| np.savez_compressed(path,X=X,smiles=np.asarray(smiles),checkpoint_sha256=json.dumps(PROTOCOL['checkpoint_sha256'])) |
| write(OUT/'features/cosmo_input_validation.json',dict(passed=True,rows=len(built), |
| atomic_compositions_validated=True,formal_charges_validated=True,charge_counts=charge_counts, |
| geometry_shared_with_monroe=True,encoder_frozen=True)) |
| del pre,encoder,mace,atoms;gc.collect();torch.cuda.empty_cache() |
| return X |
|
|
|
|
| def metrics(y,p): |
| from scipy.stats import spearmanr |
| return dict(**common.metric(y,p),spearman=float(spearmanr(y,p).statistic)) |
|
|
|
|
| def fit_kumo(model,X,y,tr,te): |
| import sdm |
| import sdm.processing as sp |
| from sdm.models.callback import Callback |
| def table(x):return sdm.TableTensor.from_tensor(torch.as_tensor(np.ascontiguousarray(x),dtype=torch.float32,device='cuda')) |
| recipe=model.default_recipe() |
| selectors=[m for m in recipe.features.modules() if isinstance(m,sp.SelectColumns)] |
| assert len(selectors)==1;selectors[0].max_columns=640 |
| shapes=[] |
| class Audit(Callback): |
| def on_context_preprocessing_end(self,model,x,y,related_tables): |
| shapes.append(list(x.shape));return x,y,related_tables |
| values=[] |
| with torch.inference_mode(),torch.autocast('cuda',dtype=torch.float16): |
| model.fit(x=table(X[tr]),y=table(y[tr,None]),num_estimators=4,estimator_batch_size=1, |
| generator=torch.Generator(device='cuda').manual_seed(SEED),recipe=recipe,callbacks=[Audit()]) |
| for start in range(0,len(te),256): |
| q=model.predict(table(X[te[start:start+256]])).numerical.float().cpu().numpy() |
| assert q.shape[0]==len(te[start:start+256]) |
| values.append(q.mean(1,dtype=np.float64)) |
| pred=np.concatenate(values);assert np.isfinite(pred).all() |
| assert shapes and max(s[-1] for s in shapes)<=640 |
| model.clear();torch.cuda.empty_cache() |
| return pred,dict(processed_context_shapes=shapes,num_estimators=4,raw_columns=X.shape[1]) |
|
|
|
|
| def run_endpoints(run,features,train,test,endpoints,smiles): |
| import sdm |
| path=OUT/'results.csv' |
| records=pd.read_csv(path).to_dict('records') if path.exists() else [] |
| lookup={s:i for i,s in enumerate(smiles)} |
| model=sdm.models.KumoTabular(task='regression',size='large',device='cuda') |
| for representation in ['cosmo_epoch10','monroe']: |
| X=features[representation] |
| for ei,endpoint in enumerate(endpoints): |
| if any(r['representation']==representation and r['endpoint']==endpoint for r in records):continue |
| train_rows=np.flatnonzero(train[endpoint].notna().to_numpy()) |
| test_rows=np.flatnonzero(test[endpoint].notna().to_numpy()) |
| tr=np.asarray([lookup[s] for s in train.iloc[train_rows].SMILES]) |
| te=np.asarray([lookup[s] for s in test.iloc[test_rows].SMILES]) |
| y=np.full(len(smiles),np.nan,dtype=float) |
| y[tr]=train.iloc[train_rows][endpoint].to_numpy(dtype=float) |
| y[te]=test.iloc[test_rows][endpoint].to_numpy(dtype=float) |
| assert not set(tr)&set(te) |
| status('predicting_endpoints',representation=representation,endpoint=endpoint,completed=len(records),total=18,n_train=len(tr),n_test=len(te)) |
| started=time.perf_counter();pred,extras=fit_kumo(model,X,y,tr,te) |
| result=dict(representation=representation,endpoint=endpoint,method='kumo_large',seed=SEED, |
| n_train=len(tr),n_test=len(te),**metrics(y[te],pred),seconds=time.perf_counter()-started, |
| raw_columns=extras['raw_columns']) |
| p=OUT/'predictions'/f'{representation}_endpoint_{ei:02d}.npz' |
| np.savez_compressed(p,train_csv_rows=train_rows,test_csv_rows=test_rows, |
| train_feature_indices=tr,test_feature_indices=te,y_true=y[te],y_pred=pred, |
| test_smiles=test.iloc[test_rows].SMILES.to_numpy(dtype=str)) |
| write(p.with_suffix('.json'),{**result,**extras}) |
| records.append(result) |
| tmp=path.with_suffix('.partial.csv');pd.DataFrame(records).to_csv(tmp,index=False);tmp.replace(path) |
| write(OUT/'progress.json',dict(completed=len(records),total=18,latest=result)) |
| run.log({'endpoint/'+k:v for k,v in result.items()}) |
| artifact=wandb.Artifact('cosmo-openadmet-endpoint-backups',type='evaluation',metadata=result) |
| for file in [p,p.with_suffix('.json'),path,OUT/'progress.json']: |
| artifact.add_file(str(file),name=str(file.relative_to(OUT))) |
| run.log_artifact(artifact,aliases=['latest',f'endpoint-{len(records):02d}']).wait(timeout=60) |
| print('OPENADMET_RESULT',json.dumps(result),flush=True) |
| assert len(records)==18 |
| del model;gc.collect();torch.cuda.empty_cache() |
|
|
|
|
| def publish(run): |
| artifact=wandb.Artifact('cosmo-openadmet-evaluation-results',type='evaluation',metadata=dict(endpoint_results=18,protocol=PROTOCOL)) |
| files=[p for p in OUT.rglob('*') if p.is_file() and 'wandb' not in p.relative_to(OUT).parts and p.suffix in ['.json','.csv','.npz','.parquet','.png','.svg','.md']] |
| for p in files:artifact.add_file(str(p),name='openadmet_evaluation/'+str(p.relative_to(OUT))) |
| sources=[Path(__file__),Path(__file__).with_name('openadmet_report.py')] |
| for p in sources:artifact.add_file(str(p),name='source/'+p.name) |
| combined=WORK/'expanded_evaluation/analysis/REPORT.md' |
| artifact.add_file(str(combined),name='expanded_evaluation/analysis/REPORT.md') |
| run.log_artifact(artifact,aliases=['latest','complete']).wait(timeout=60) |
| operations=[CommitOperationAdd(path_in_repo='openadmet_evaluation/'+str(p.relative_to(OUT)),path_or_fileobj=str(p)) for p in files] |
| operations += [CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources] |
| operations.append(CommitOperationAdd(path_in_repo='expanded_evaluation/analysis/REPORT.md',path_or_fileobj=str(combined))) |
| revision=HfApi().create_commit(repo_id=REPOSITORY,operations=operations, |
| commit_message='Add matched COSMO epoch 10 and Monroe Kumo Large results on nine OpenADMET endpoints').oid |
| write(OUT/'hub_verified.json',dict(completed=True,repository=REPOSITORY,revision=revision)) |
| return revision |
|
|
|
|
| def main(): |
| OUT.mkdir(exist_ok=True) |
| with (OUT/'driver.lock').open('a') as lock: |
| try:fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) |
| except BlockingIOError:print('OPENADMET_ALREADY_RUNNING',flush=True);return |
| for folder in ['data','features','predictions','analysis']:(OUT/folder).mkdir(exist_ok=True) |
| PROTOCOL['versions']={n:importlib.metadata.version(n) for n in ['torch','numpy','rdkit','scikit-learn','structured-data-models','torch-geometric','datamol','wandb']} |
| train,test,endpoints,smiles=prepare_data() |
| restore_monroe() |
| built=prepare_graphs(smiles) |
| |
| features={'monroe':monroe_embeddings(built,smiles),'cosmo_epoch10':cosmo_embeddings(built,smiles)} |
| del built |
| write(OUT/'features/manifest.json',{n:dict(rows=len(X),columns=X.shape[1],sha256=sha(OUT/'features'/f'{n}.npz')) for n,X in features.items()}) |
| saved=json.loads((OUT/'wandb.json').read_text()) if (OUT/'wandb.json').exists() else {} |
| run=wandb.init(entity='queimo',project='rem3di-reproduction',id=saved.get('run_id') or wandb.util.generate_id(),resume='allow', |
| name='openadmet-cosmo-epoch10-vs-monroe-kumo-large',group='cosmo-openadmet',job_type='official-split-comparison', |
| config=PROTOCOL,dir=str(OUT),settings=wandb.Settings(init_timeout=90)) |
| write(OUT/'wandb.json',dict(run_id=run.id,url=run.url)) |
| try: |
| run_endpoints(run,features,train,test,endpoints,smiles) |
| from openadmet_report import build_report |
| build_report(OUT,train,test,endpoints,run) |
| write(OUT/'complete.json',dict(completed=True,endpoint_results=18,endpoints=9,representations=2,official_train_rows=5326,official_test_rows=2282,wandb_url=run.url)) |
| revision=publish(run) |
| run.summary.update(dict(all_evaluations_complete=True,endpoint_results=18,hub_revision=revision));run.finish() |
| status('complete',endpoint_results=18,hub_revision=revision) |
| HfApi().create_commit(repo_id=REPOSITORY,operations=[ |
| CommitOperationAdd(path_in_repo='openadmet_evaluation/'+p.name,path_or_fileobj=str(p)) for p in [OUT/'status.json',OUT/'hub_verified.json']], |
| commit_message='Verify completed OpenADMET evaluation publication') |
| except Exception: |
| status('failed',error=traceback.format_exc());run.finish(exit_code=1);raise |
|
|
|
|
| if __name__=='__main__': |
| try:main() |
| except Exception: |
| OUT.mkdir(exist_ok=True);status('failed',error=traceback.format_exc());raise |
|
|
