| """Extract frozen MACE inputs before unused learned preprocessing runs. |
| |
| The batch plan, MACE call, coordinates and cache layout match multitask_data. |
| Raising after its existing atomic-input capture avoids downstream preprocessing |
| that cannot affect the captured raw input. Existing batch-progress files resume. |
| """ |
| from pathlib import Path |
| import sys,gc |
| sys.path.insert(0,str(Path(__file__).resolve().parent)) |
| from scratch_data import * |
| common=shared.common |
|
|
|
|
| def cache_mace_fast(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=[] |
| class RawCaptured(Exception): |
| pass |
| def capture_raw(module, arguments): |
| captures.append(arguments[0].detach()) |
| raise RawCaptured |
| hook=mace.preprocessor.atomic_preprocessor.register_forward_pre_hook(capture_raw) |
| 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() |
| try: |
| mace.preprocessor(batch) |
| except RawCaptured: |
| pass |
| 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', |
| raw_capture_only=True,description='Frozen MACE only. Full-precision inputs to learned atomic/geometric preprocessors; no trainable activations cached.')) |
|
|
