"""Pinned CombiSolv-QM/COSMO preparation and frozen, full-precision MACE cache.""" 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,write,sha import concurrent.futures,fcntl,gc,hashlib,json,multiprocessing as mp,os,shutil,time,traceback import numpy as np import pandas as pd import torch from rdkit import Chem,RDLogger from rdkit.Chem import AllChem from huggingface_hub import HfApi,hf_hub_download,snapshot_download,CommitOperationAdd OUT=WORK/'multitask_combisolv' DATA=OUT/'data' MODEL_REPO='Queimo/rem3di-cosmo10-combisolv-multitask' DATA_REPO='Queimo/combisolv-qm-rem3di-multitask' RAW_REV='eacfee378eaafa49c98b4e12e9d76e2c26de1626' COSMO_REPO='Queimo/pubchemqc-cosmo-pm6' COSMO_REV='6aeb084c14fd757677d00665efac8d130611273c' MONROE_REV='de075eabe98e4adbbedbc86ca9e877f4bfd9b629' INITIAL_HASHES={ 'encoder.pth':'97e616a16e8317717784929495ba3a030bfcd821038be6e05324624b17c409c0', 'atomic_preprocessor.pth':'9683883267a902773b4f6777eedb366befd016f4f7083bed62a1bd35ae113e3d', 'geometric_preprocessor.pth':'4da0843646d06fcd2ff5b13e1fd4fb4b22c8c0456d0507c59df01e665c4d6d7b', 'cosmo_head.pth':'1a43270e60f00c9428099e0e37ab62f77a2c219ea611d4f05194bc1abf69b0a7'} DECODER_REPO='Queimo/rem3di-pcqm500k-corrected' DECODER_REV='2f5c3b40b36259a13adc8412689254054bb93288' DECODER_PATH='runs/corrected_20261027/evaluation/paper_audit/author_pcqm500k_5ep/epoch_02/decoder.pth' def status(phase,**kw): row=dict(phase=phase,pid=os.getpid(),updated_unix=time.time(),**kw) write(OUT/'status.json',row);print('MULTITASK_STATUS',json.dumps(row),flush=True) def connectivity(smiles): mol=Chem.MolFromSmiles(str(smiles)) assert mol is not None,smiles key=Chem.MolToInchiKey(mol) assert key,smiles return key.split('-')[0] def geometry_job(item): """One reproducible explicit-H 3D conformer, never substitute a 2D layout.""" index,smiles=item RDLogger.DisableLog('rdApp.warning') mol=Chem.AddHs(Chem.MolFromSmiles(smiles)) methods=[] for retry in range(3): params=AllChem.ETKDGv3();params.randomSeed=42026+index+100000*retry params.numThreads=1;params.useRandomCoords=retry>0;params.maxIterations=2000 result=AllChem.EmbedMolecule(mol,params) if result==0:break if result!=0: # Knowledge-based torsions can be inconsistent for strained macrocycles # or uncommon N=S=N valence. Generic 3D distance geometry avoids those # preferences; subsequent force-field optimization still applies. params=AllChem.EmbedParameters();params.randomSeed=342026+index params.useRandomCoords=True;params.ignoreSmoothingFailures=True params.useBasicKnowledge=False;params.useExpTorsionAnglePrefs=False params.maxIterations=5000;params.numThreads=1 result=AllChem.EmbedMolecule(mol,params) methods.append('generic-3D-distance-geometry') if result!=0:return dict(smiles=smiles,error='ETKDG and generic 3D distance geometry failed') charge=Chem.GetFormalCharge(mol);radicals=sum(a.GetNumRadicalElectrons() for a in mol.GetAtoms()) # MMFF can mutate formal-charge/aromaticity annotations (e.g. N#N=O). Keep # the source molecular graph intact; accept optimized coordinates only. optimized=Chem.Mol(mol) if AllChem.MMFFHasAllMoleculeParams(mol): code=AllChem.MMFFOptimizeMolecule(optimized,mmffVariant='MMFF94s',maxIters=500) methods.append('MMFF94s:'+str(code)) elif AllChem.UFFHasAllMoleculeParams(mol): code=AllChem.UFFOptimizeMolecule(optimized,maxIters=500);methods.append('UFF:'+str(code)) else:methods.append('ETKDG-only: no complete force-field parameters') # Some MMFF parameter checks themselves mutate the checked molecule. if Chem.GetFormalCharge(optimized)!=charge: methods=['ETKDG-only: rejected MMFF charge mutation'];optimized=None pos=np.asarray((optimized if optimized is not None else mol).GetConformer().GetPositions(),dtype=np.float32) nums=[a.GetAtomicNum() for a in mol.GetAtoms()] assert np.isfinite(pos).all() and pos.shape==(len(nums),3) assert radicals==0 and charge==0 and sum(nums)%2==0,(smiles,charge,radicals) return dict(smiles=smiles,molecule_group=connectivity(smiles),atomic_numbers=nums, positions=pos.tolist(),charge=charge,geometry_method=';'.join(methods)) def initialize(): OUT.mkdir(exist_ok=True);DATA.mkdir(exist_ok=True) initial=WORK/'checkpoints/epoch_10' bigsol=json.loads((WORK/'bigsoldb_evaluation/protocol.json').read_text()) for n,digest in INITIAL_HASHES.items(): assert sha(initial/n)==digest,(n,'Not the BigSolDB-evaluated PCQMep2COSMOep10 checkpoint') if n!='cosmo_head.pth':assert bigsol['checkpoint_sha256'][n]==digest decoder=Path(hf_hub_download(DECODER_REPO,DECODER_PATH,revision=DECODER_REV)) foundation=Path('/marimo/rem3di_training/MACE-POLAR-1-M.model') foundation_sha=sha(foundation) assert foundation_sha=='fab8b8713c832f31a2a853aaa22fd638be8a369cbf5095e6b3e982a18d10e93a' write(OUT/'initialization.json',dict(initial_checkpoint_repository='Queimo/rem3di-pcqm-epoch2-cosmo', initial_checkpoint_path='epoch_10',initial_checkpoint_revision='9266a41fbe3b36afe0615e0c3ba232ca74bdeef7',initial_hashes=INITIAL_HASHES, frozen_foundation_sha256=foundation_sha, verified_against_bigsoldb_protocol=True,bigsoldb_protocol_sha256=sha(WORK/'bigsoldb_evaluation/protocol.json'), denoising_decoder=dict(repository=DECODER_REPO,revision=DECODER_REV,path=DECODER_PATH,sha256=sha(decoder)), explanation='All encoder, atomic/geometric preprocessor and COSMO head weights come from the exact evaluated COSMO10 checkpoint. Only its omitted denoising decoder is restored from PCQM epoch2.')) return initial,decoder def prepare_tables(): if (DATA/'manifest.json').exists(): previous=json.loads((DATA/'manifest.json').read_text()) if previous.get('preparation_version')==2:return previous initialize();status('preparing_joint_splits') root=Path(snapshot_download(COSMO_REPO,repo_type='dataset',revision=COSMO_REV, allow_patterns=['data/*','manifest.json','label_transform.json'])) cosmo=pd.concat([pd.read_parquet(p).assign(original_split=split) for split in ['train','validation'] for p in sorted((root/'data').glob(split+'-*.parquet'))],ignore_index=True) assert len(cosmo)==360259 overlap=pd.concat([pd.read_parquet(p,columns=['molecule_group']) for p in sorted((root/'data').glob('benchmark_overlap-*.parquet'))],ignore_index=True) blocked=set(overlap.molecule_group.astype(str)) bigsol=pd.read_csv(WORK/'bigsoldb_evaluation/data/test.csv') bigsol_keys={connectivity(s) for s in bigsol.SMILES_Solute.unique()} assert len(bigsol_keys)<=126 raw_path=Path(hf_hub_download(DATA_REPO,'raw/combisolv_qm.parquet',repo_type='dataset',revision=RAW_REV)) raw_manifest=json.loads(Path(hf_hub_download(DATA_REPO,'raw/source_manifest.json',repo_type='dataset',revision=RAW_REV)).read_text()) assert sha(raw_path)==raw_manifest['parquet_sha256'] raw=pd.read_parquet(raw_path);assert len(raw)==1000000 # The supplied InChI avoids RDKit 2026's interpretation of the legacy # hypervalent N#N=O SMILES as a charged HN2O species. Preserve source strings. structures=pd.concat([raw[['smiles solute','inchi solute']].rename(columns={'smiles solute':'source_smiles','inchi solute':'source_inchi'}), raw[['smiles solvent','inchi solvent']].rename(columns={'smiles solvent':'source_smiles','inchi solvent':'source_inchi'})]).drop_duplicates() converted=[] for source_smiles,source_inchi in structures.itertuples(index=False,name=None): mol=Chem.MolFromInchi(source_inchi,removeHs=False) assert mol is not None,source_inchi converted.append(dict(source_smiles=source_smiles,source_inchi=source_inchi,smiles=Chem.MolToSmiles(mol,isomericSmiles=True))) structures=pd.DataFrame(converted) mapping=dict(zip(structures.source_smiles,structures.smiles)) molecules=[str(s) for s in pd.unique(pd.concat([raw['smiles solute'].map(mapping),raw['smiles solvent'].map(mapping)]))] lookup={s:i for i,s in enumerate(molecules)} keys=np.asarray([connectivity(s) for s in molecules]);key_by_smiles=dict(zip(molecules,keys)) solvent=np.asarray([lookup[mapping[s]] for s in raw['smiles solvent']],dtype=np.int32) solute=np.asarray([lookup[mapping[s]] for s in raw['smiles solute']],dtype=np.int32) solute_key=keys[solute];solvent_key=keys[solvent] heldout=bigsol_keys cosine_val=set(cosmo.loc[cosmo.original_split=='validation','molecule_group'].astype(str)) eligible=sorted(set(solute_key)-heldout-cosine_val) rng=np.random.default_rng(42026) valkeys=set(np.asarray(eligible)[rng.permutation(len(eligible))[:max(1,round(.02*len(eligible)))]]) # Validation solutes cannot occur in either input role in new supervised training. benchmark=np.isin(solute_key,list(heldout))|np.isin(solvent_key,list(heldout)) is_val=np.isin(solute_key,list(valkeys))&~benchmark auxiliary_holdout=np.isin(solute_key,list(cosine_val))|np.isin(solvent_key,list(cosine_val))|np.isin(solvent_key,list(valkeys)) split=np.full(len(raw),'train',dtype=object) split[auxiliary_holdout]='auxiliary_holdout';split[is_val]='validation';split[benchmark]='benchmark_overlap' train=split=='train' assert not set(solute_key[train])&valkeys and not set(solvent_key[train])&valkeys assert not set(solute_key[train])&heldout and not set(solvent_key[train])&heldout target=raw['dGsolv (298K) [kcal/mol]'].to_numpy(dtype=np.float64) assert np.isfinite(target).all() mean=float(target[train].mean());scale=float(target[train].std()) assert scale>0 pairs=pd.DataFrame(dict(source_row=np.arange(len(raw)),solute_index=solute,solvent_index=solvent, temperature_K=298.,dGsolv_kcal_mol=target,split=split)) pairs.to_parquet(DATA/'combisolv_pairs.parquet',index=False,compression='zstd') pd.DataFrame(dict(smiles=molecules,molecule_group=keys)).to_parquet(DATA/'combisolv_molecules.parquet',index=False,compression='zstd') structures.to_parquet(DATA/'source_structure_mapping.parquet',index=False,compression='zstd') cosmo['joint_split']=cosmo.original_split mask=np.isin(cosmo.molecule_group.astype(str),list(heldout|valkeys)) cosmo.loc[mask,'joint_split']='new_holdout' assert not set(cosmo.loc[cosmo.joint_split=='train','molecule_group'].astype(str))&(heldout|valkeys) # Retain every original row with explicit membership, rather than silently dropping it. for start in range(0,len(cosmo),50000): cosmo.iloc[start:start+50000].to_parquet(DATA/f'cosmo-{start//50000:03d}.parquet',index=False,compression='zstd') shutil.copyfile(root/'label_transform.json',DATA/'cosmo_label_transform.json') write(DATA/'combisolv_label_transform.json',dict(transform='z=(dGsolv-mean)/scale; no logarithm for signed free energies', target='dGsolv_kcal_mol',temperature_K=298.,mean=mean,scale=scale,training_rows=int(train.sum()), fit_split='CombiSolv-QM training only',temperature_feature='Temperature_K / 298; constant 1 for this source')) write(DATA/'heldout_groups.json',dict(bigsoldb_test_solute_connectivity=sorted(bigsol_keys), existing_esol_freesolv_quarantine=sorted(blocked),combisolv_validation_solutes=sorted(valkeys), cosmo_validation=sorted(cosine_val))) manifest=dict(preparation_version=2,raw_revision=RAW_REV,cosmo_revision=COSMO_REV,raw_combisolv_rows=len(raw), unique_qm_molecules=len(molecules),combisolv_splits={str(k):int(v) for k,v in pairs.split.value_counts().items()}, cosmo_splits={str(k):int(v) for k,v in cosmo.joint_split.value_counts().items()}, bigsoldb_test_solute_groups=len(bigsol_keys),new_training_benchmark_overlap=0, exclusion='Connectivity InChIKey, in both solute and solvent roles; validation solutes excluded across objectives', geometry_identity='Canonical isomeric SMILES from supplied source InChI; original SMILES/InChI mapping preserved. This avoids RDKit 2026 legacy hypervalent nitrogen reinterpretation.', esol_freesolv='Original COSMO quarantine preserved. CombiSolv training is protected against requested BigSolDB holdout; no claim of ESOL/FreeSolv quarantine for CombiSolv.', initial_checkpoint_exposure='New exclusions cannot remove any BigSolDB or validation molecule exposure already present in the starting COSMO10/PCQM2 checkpoint.', source_manifest=raw_manifest,files={p.name:sha(p) for p in DATA.iterdir() if p.is_file() and p.name!='manifest.json'}) write(DATA/'manifest.json',manifest);status('joint_splits_ready',**{k:v for k,v in manifest.items() if k not in ['files','source_manifest']}) return manifest def prepare_geometries(): path=DATA/'combisolv_geometries.parquet' if path.exists(): record_geometry_manifest(path);return molecules=pd.read_parquet(DATA/'combisolv_molecules.parquet') status('generating_combisolv_3d',total=len(molecules)) result=[] with concurrent.futures.ProcessPoolExecutor(max_workers=8,mp_context=mp.get_context('spawn')) as pool: for i,row in enumerate(pool.map(geometry_job,enumerate(molecules.smiles),chunksize=4),1): result.append(row) if i%256==0:status('generating_combisolv_3d',completed=i,total=len(molecules)) failures=[r for r in result if 'error' in r];write(DATA/'geometry_failures.json',failures) assert not failures,failures[:10] assert [r['smiles'] for r in result]==molecules.smiles.tolist() pd.DataFrame(result).to_parquet(path,index=False,compression='zstd') record_geometry_manifest(path) def record_geometry_manifest(path): from rdkit import rdBase frame=pd.read_parquet(path,columns=['smiles','geometry_method']) molecules=pd.read_parquet(DATA/'combisolv_molecules.parquet',columns=['smiles']) assert frame.smiles.tolist()==molecules.smiles.tolist() write(DATA/'geometry_manifest.json',dict(rows=len(frame),sha256=sha(path),seed=42026, rdkit_version=rdBase.rdkitVersion, algorithm='RDKit ETKDGv3 explicit H, deterministic random-coordinate retries, generic 3D distance geometry fallback; MMFF94s or UFF where parameterized', methods={str(k):int(v) for k,v in frame.geometry_method.value_counts().items()}, no_2d_fallback=True,missing=0)) def cosmo_frame(): return pd.concat([pd.read_parquet(p) for p in sorted(DATA.glob('cosmo-[0-9]*.parquet'))],ignore_index=True) class RawCache: """Flat float32 arrays, atom offsets and positions; learned preprocessors stay trainable.""" def __init__(self,name): self.root=OUT/'cache'/name self.ptr=np.load(self.root/'offsets.npy');self.lengths=np.diff(self.ptr) self.positions=np.load(self.root/'positions.npy',mmap_mode='r') self.features=np.load(self.root/'features.npy',mmap_mode='r') def sample(self,ids): from torch.nn.utils.rnn import pad_sequence from threedscriptors.data_handling.sample import Sample spans=[slice(int(self.ptr[i]),int(self.ptr[i+1])) for i in ids] f=pad_sequence([torch.from_numpy(np.array(self.features[s],copy=True)) for s in spans],batch_first=True).cuda() p=pad_sequence([torch.from_numpy(np.array(self.positions[s],copy=True)) for s in spans],batch_first=True).cuda() lens=torch.tensor(self.lengths[ids],device='cuda') mask=torch.arange(f.shape[1],device='cuda')[None,:]>=lens[:,None] return Sample(embeddings=f,atomic_positions=p,padding_mask=mask) def cache_mace(name,frame): root=OUT/'cache'/name;root.mkdir(parents=True,exist_ok=True) if (root/'complete.json').exists():return from remedi.data_handling.sample import Sample,yield_molecules_collate_fn from build_author_mace_runtime import build_author_mace_runtime from train_cosmo import batches numbers=[torch.tensor(np.asarray(z),dtype=torch.long) for z in frame.atomic_numbers] positions=[torch.tensor(np.stack(p),dtype=torch.float32) for p in frame.positions] lengths=np.asarray([len(z) for z in numbers]);ptr=np.r_[0,np.cumsum(lengths)] np.save(root/'offsets.npy',ptr) pos=np.lib.format.open_memmap(root/'positions.npy',mode='w+',dtype=np.float32,shape=(int(ptr[-1]),3)) for i,p in enumerate(positions):pos[ptr[i]:ptr[i+1]]=p.numpy() pos.flush();del pos mode='r+' if (root/'features.npy').exists() else 'w+' features=np.lib.format.open_memmap(root/'features.npy',mode=mode,dtype=np.float32,shape=(int(ptr[-1]),4096)) progress=json.loads((root/'progress.json').read_text()) if (root/'progress.json').exists() else {'batches':0} mace=build_author_mace_runtime().cuda().eval().requires_grad_(False) captures=[] hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(lambda m,a:captures.append(a[0].detach())) started=time.time() if name=='bigsoldb': plan=list(common.batch_indices(np.arange(len(frame)),lengths,max_size=24,atom_cap=2400)) else:plan=list(batches(np.arange(len(frame)),lengths,max_size=64,atom_cap=3600)) status('caching_frozen_mace',corpus=name,total=len(frame),atoms=int(ptr[-1]),bytes=int(ptr[-1])*4096*4,batches=len(plan)) with torch.inference_mode(): for bi,ids in enumerate(plan,1): if bi<=progress['batches']:continue batch=yield_molecules_collate_fn([Sample(atomic_positions=positions[i],atomic_numbers=numbers[i],total_charge=float(frame.charge.iloc[i]) if 'charge' in frame else 0.,multiplicity=1.) for i in ids]).to_('cuda') captures.clear();mace.preprocessor(batch);assert len(captures)==1 raw=captures.pop().float().cpu().numpy() for j,i in enumerate(ids): v=raw[j,:lengths[i]];assert v.shape==(lengths[i],4096) and np.isfinite(v).all() features[ptr[i]:ptr[i+1]]=v if bi%25==0 or bi==len(plan): features.flush();write(root/'progress.json',dict(batches=bi,total_batches=len(plan))) status('caching_frozen_mace',corpus=name,batch=bi,batches=len(plan),seconds=time.time()-started) features.flush();del features hook.remove();del mace,numbers,positions;gc.collect();torch.cuda.empty_cache() write(root/'complete.json',dict(rows=len(frame),atoms=int(ptr[-1]),width=4096,dtype='float32', mace_sha256='fab8b8713c832f31a2a853aaa22fd638be8a369cbf5095e6b3e982a18d10e93a', description='Frozen MACE only. Full-precision inputs to learned atomic/geometric preprocessors; no trainable activations cached.')) def publish_data(): marker=OUT/'prepared_data_hub.json' if marker.exists():return pairs=pd.read_parquet(DATA/'combisolv_pairs.parquet') for split,frame in pairs.groupby('split',sort=False): frame.to_parquet(DATA/f'combisolv_{split}.parquet',index=False,compression='zstd') manifest=json.loads((DATA/'manifest.json').read_text()) manifest['files']={p.name:sha(p) for p in DATA.iterdir() if p.is_file() and p.name!='manifest.json'} write(DATA/'manifest.json',manifest) header='''--- license: cc-by-4.0 task_categories: - tabular-regression configs: - config_name: raw_combisolv_qm default: true data_files: - split: records path: raw/combisolv_qm.parquet - config_name: prepared_pairs data_files: - split: train path: prepared/combisolv_train.parquet - split: validation path: prepared/combisolv_validation.parquet - split: benchmark_overlap path: prepared/combisolv_benchmark_overlap.parquet - split: auxiliary_holdout path: prepared/combisolv_auxiliary_holdout.parquet - config_name: combisolv_geometries data_files: - split: molecules path: prepared/combisolv_geometries.parquet - config_name: cosmo_pool data_files: - split: records path: prepared/cosmo-*.parquet --- ''' card=header+'\n# CombiSolv-QM and COSMO multi-task continuation data\n\n' card+='The unchanged original one-million-row CombiSolv-QM table is preserved under `raw/`, together with source archive/CSV/Parquet checksums. The target is signed solvation Gibbs free energy at 298 K, in kcal/mol, using 1 mol/L reference concentrations. It is distinct from experimental solubility.\n\n' card+=f"Prepared CombiSolv memberships: `{json.dumps(manifest['combisolv_splits'],sort_keys=True)}`. All {manifest['unique_qm_molecules']:,} unique structures have explicit-H coordinates; supplied InChI defines identity, and original SMILES/InChI mappings are retained. Solute/solvent indices refer to row order in the geometry and molecule tables.\n\n" card+='Geometries use seeded RDKit ETKDGv3, random-coordinate retries and generic 3D distance geometry for strained/unusual structures, followed by MMFF94s or UFF when parameterized. There is no 2D fallback. Geometry methods, software version, coverage and checksums are recorded.\n\n' card+=f"The COSMO pool retains the source training/validation records and adds an explicit `joint_split` membership: `{json.dumps(manifest['cosmo_splits'],sort_keys=True)}`. Filter that column for training. The complete original COSMO dataset and its existing ESOL/FreeSolv quarantine remain available in [Queimo/pubchemqc-cosmo-pm6](https://huggingface.co/datasets/{COSMO_REPO}) at revision `{COSMO_REV}`.\n\n" card+='BigSolDB2.0 held-out solute connectivity groups are excluded from all new objectives and both CombiSolv input roles. CombiSolv validation solutes are excluded across objectives. Existing COSMO validation groups are excluded from CombiSolv training. Raw rows remain available with explicit quarantine membership; earlier exposure in the starting checkpoint cannot be undone. ESOL/FreeSolv quarantine is preserved for COSMO, but is not claimed for CombiSolv.\n\n' card+='CombiSolv target mean/std is fitted on its retained training rows only; signed energies are not log transformed. The existing COSMO log1p/ln transform and normalization are preserved so the initialized output head retains its meaning.\n\n' card+='[SolProp v1.2 source](https://zenodo.org/records/5970538) · [Vermeire and Green (2021)](https://doi.org/10.1016/j.cej.2021.129307) · [Vermeire, Chung and Green (2022)](https://doi.org/10.1021/jacs.2c01768). Please cite the original data authors.\n' card_path=OUT/'dataset_README.md';card_path.write_text(card) files=[p for p in DATA.iterdir() if p.is_file()] commit=HfApi().create_commit(repo_id=DATA_REPO,repo_type='dataset', operations=[CommitOperationAdd(path_in_repo='prepared/'+p.name,path_or_fileobj=str(p)) for p in files]+[ CommitOperationAdd(path_in_repo='README.md',path_or_fileobj=str(card_path))], commit_message='Add reproducible joint splits, all geometries, label transforms and benchmark exclusions') remote={p.rfilename:p for p in HfApi().dataset_info(DATA_REPO,revision=commit.oid,files_metadata=True).siblings} for p in files: entry=remote['prepared/'+p.name] if entry.lfs:assert entry.lfs.sha256==sha(p) write(marker,dict(repository=DATA_REPO,revision=commit.oid,verified=True,files=len(files))) def main(): OUT.mkdir(exist_ok=True) lock=(OUT/'prepare.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB) prepare_tables();prepare_geometries() cache_mace('combisolv',pd.read_parquet(DATA/'combisolv_geometries.parquet')) cache_mace('cosmo',cosmo_frame()) # Evaluation geometries must be byte-identical to those used for the published baseline. import openadmet_evaluate as runtime runtime.restore_monroe() from monroe.model.constants import NODE_FEAT_LIST_ONE_HOT graphs=pd.read_parquet(WORK/'bigsoldb_evaluation/features/shared_graphs.parquet') vocab=NODE_FEAT_LIST_ONE_HOT['atomic_number'];col=list(NODE_FEAT_LIST_ONE_HOT).index('atomic_number') ptable=Chem.GetPeriodicTable() graph_atoms=pd.DataFrame(dict(smiles=graphs.smiles,positions=[np.stack(p).astype(np.float32).tolist() for p in graphs.pos_rdkit], atomic_numbers=[[ptable.GetAtomicNumber(vocab[int(i)]) for i in np.stack(codes)[:,col]] for codes in graphs.node_codes], charge=[int(round(float(np.stack(v)[:,5].sum()))) for v in graphs.node_float])) graph_atoms.to_parquet(DATA/'bigsoldb_fixed_geometries.parquet',index=False,compression='zstd') cache_mace('bigsoldb',graph_atoms) publish_data() write(OUT/'preparation_complete.json',dict(completed=True,time=time.time(),initialization=json.loads((OUT/'initialization.json').read_text()),data=json.loads((OUT/'prepared_data_hub.json').read_text()))) status('preparation_complete') if __name__=='__main__': try:main() except Exception:status('failed_preparation',error=traceback.format_exc());raise