source/multitask_data.py

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