Verify random initialization, real-data gradients and exact next-step resume before scratch training

04a3ade verified 2 days ago

13.2 kB

"""One full-context BigSolDB Kumo Large probe per saved fifth-epoch checkpoint."""
from pathlib import Path
import sys
sys.path.insert(0,str(Path(__file__).resolve().parent))
from scratch_model import *
from scratch_train import status
import bigsoldb_evaluate as reference
import random,gc,fcntl,argparse,traceback,shutil
from huggingface_hub import HfApi,CommitOperationAdd,hf_hub_download

EVAL=OUT/'evaluation'


def embeddings(point):
    checkpoint=OUT/'checkpoints'/f'probe_{point:03d}';path=EVAL/'features'/f'probe_{point:03d}.npz'
    hashes={n:sha(checkpoint/n) for n in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']}
    geometry=pd.read_parquet(DATA/'bigsoldb_fixed_geometries.parquet')
    if path.exists():
        saved=np.load(path)
        assert saved['smiles'].tolist()==geometry.smiles.tolist()
        assert json.loads(str(saved['checkpoint_sha256']))==hashes
        return saved['X'],hashes
    model=ScratchJoint(checkpoint).eval().requires_grad_(False)
    cache=RawCache('bigsoldb');z=np.empty((len(geometry),640),dtype=np.float32)
    with torch.inference_mode():
        for ids in batch_indices(np.arange(len(geometry)),cache.lengths,max_size=24,atom_cap=2400):
            _,_,values=model.encode(cache.sample(ids));z[ids]=values.float().cpu().numpy()
    assert np.isfinite(z).all() and z.shape==(1334,640)
    np.savez_compressed(path,X=z,smiles=geometry.smiles.to_numpy(dtype=str),checkpoint_sha256=json.dumps(hashes))
    del model,cache;gc.collect();torch.cuda.empty_cache()
    return z,hashes


def build_report(run):
    import matplotlib
    matplotlib.use('Agg')
    import matplotlib.pyplot as plt
    rows=pd.read_csv(EVAL/'results.csv').sort_values('point')
    prior_path=OUT/'prior_comparison.csv'
    if not prior_path.exists():
        p=hf_hub_download('Queimo/rem3di-cosmo10-combisolv-multitask','ablations/references/comparison.csv',
                         revision='ff86641eb441ffe038c8aca1df1ba32c02f74cb7')
        shutil.copyfile(p,prior_path)
    prior=pd.read_csv(prior_path)
    selected=prior.loc[prior.variant.isin(['initialization','published_rem3di','monroe','literature'])]
    fig,axes=plt.subplots(1,3,figsize=(13,3.8),layout='constrained')
    for ax,metric in zip(axes,['mae','rmse','r2']):
        ax.plot(rows.epoch_position,rows[metric],'-o',ms=3,label='Scratch REM3DI + COSMO + Kumo L')
        for variant,label,color in [('initialization','PCQM2→COSMO10','#888888'),('published_rem3di','Published Rem3Di','#332288'),('monroe','Monroe','#CC79A7')]:
            ax.axhline(float(prior.loc[prior.variant==variant,metric].iloc[0]),ls='--',color=color,label=label)
        ax.set_xlabel('Joint pretraining epoch');ax.set_ylabel(metric.upper());ax.grid(alpha=.2)
    axes[0].legend(fontsize=6);fig.savefig(EVAL/'analysis/learning_curve.png',dpi=180)
    fig.savefig(EVAL/'analysis/learning_curve.svg');plt.close(fig)
    lines=['# Random-initialized Rem3Di + COSMO: ten epochs, fifty BigSolDB probes','',
        f'Completed {len(rows)}/50 planned fifth-epoch frozen Kumo Large probes. Every learned Rem3Di preprocessor, encoder, denoising decoder and COSMO head started from its constructor initialization. No published Rem3Di or PCQM/COSMO checkpoint initializes this run. MACE-POLAR-1-M alone is pretrained and frozen. The encoder/decoder architecture and 640-dimensional pooled descriptor match the preceding experiments.','',
        'Two simultaneous streams: standard author atomic-embedding denoising on the retained PCQM500k structures, and 153 sigma-bin areas plus A/V reconstruction on the retained original PM6 COSMO structures. Each epoch covers both complete training streams; the shorter stream cycles. No CombiSolv head or covariance/decorrelation penalty is included.','',
        'The released denoising implementation, Gaussian sigma 0.3 and its original gradient paths are preserved. In particular the clean atomic reconstruction target remains differentiable, as in the author implementation. Raw and per-channel denoising losses, held-out descriptor variability and target variability are logged; low reconstruction loss alone is not treated as downstream quality.','',
        'Task uncertainty weighting uses sum(L/(2 exp(s)) + s/2). The unmodified raw denoising loss has 1088 channel terms, so its initial log variance is ln(1088)−0.5; COSMO starts at −0.5. This is mathematically equivalent, up to an additive constant, to starting both log variances at −0.5 with per-channel denoising units. All invariant feature normalization and COSMO log1p/ln target standardization were freshly fit using training rows only.','',
        'Held-out BigSolDB solute connectivity groups are excluded from both new training streams. PCQM and COSMO validation groups are also excluded across objectives. The pretrained MACE foundation can have prior exposure; this experiment makes no claim about removing its historical pretraining data. Exact memberships and all original coordinates are preserved in the prepared dataset.','',
        'Every fifth-epoch checkpoint is saved with all heads, task weights, AdamW state, OneCycle scheduler and RNG state, and verified against its Hugging Face manifest before probing. Probes run in separate processes and never change the training RNG, optimizer or schedule. There are exactly five probes per epoch, fifty total; the epoch-10 endpoint is the predeclared final checkpoint, and all intermediate test scores are reported without test-based tuning.','',
        '| Point | Epoch position | MAE | RMSE | R² |','|---:|---:|---:|---:|---:|']
    for r in rows.itertuples():lines.append(f'| {r.point} | {r.epoch_position:.5f} | {r.mae:.6f} | {r.rmse:.6f} | {r.r2:.6f} |')
    lines+=['','## Prior matched probes and literature values','','| Model | Origin | MAE | RMSE | R² |','|---|---|---:|---:|---:|']
    for r in selected.itertuples():lines.append(f'| {r.model} | {r.origin} | {r.mae:.6f} | {r.rmse:.6f} | {r.r2:.6f} |')
    lines+=['','All measured probes use exactly the released 9,686 test rows and the 75,731-row non-test complement as full Kumo context. Fixed shared baseline conformers, 1281 ordered solute640/solvent640/temperature-in-Kelvin features, Kumo Large v1.0.1, four estimators, seed 42, fp16, query chunks 128. No additional transformation of the supplied log10 mol/L labels. Monroe uses its full 1441 features. Literature rows are copied reported values, not rerun models. The author prediction file has 9,686 rows and reproduces Table 1 metrics; the prose says 9,250. The private neural validation subset is unavailable.','',
        f'[Prepared data](https://huggingface.co/datasets/{DATA_REPO}) · [Checkpoints](https://huggingface.co/{MODEL_REPO}) · [Probe W&B]({run.url}) · [Solvaformer Table 1]({reference.PAPER_URL})','',
        '![Fifty-point learning curve](learning_curve.png)','']
    text='\n'.join(lines);(EVAL/'analysis/REPORT.md').write_text(text)
    (OUT/'README.md').write_text('---\ntags:\n- chemistry\n- multitask\n---\n\n'+text.replace('](learning_curve.png)','](evaluation/analysis/learning_curve.png)'))
    return rows


def publish(run):
    files=[p for p in EVAL.rglob('*') if p.is_file() and p.suffix in ['.json','.csv','.npz','.png','.svg','.md']]
    files += [p for p in OUT.iterdir() if p.is_file() and p.suffix in ['.json','.csv','.yaml','.md']]
    sources=[Path(__file__).with_name(n) for n in ['scratch_prepare.py','scratch_data.py','scratch_model.py','scratch_train.py','scratch_probe.py','scratch_driver.py','scratch_cache_initial.py','scratch_fast_cache.py',
        'eval_common.py','train_cosmo.py','multitask_data.py','bigsoldb_evaluate.py']]
    ops=[CommitOperationAdd(path_in_repo=str(p.relative_to(OUT)),path_or_fileobj=str(p)) for p in files]
    ops += [CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources]
    commit=HfApi().create_commit(repo_id=MODEL_REPO,operations=ops,commit_message='Update scratch-training learning curve and exact fifth-epoch probe artifacts')
    write(OUT/'report_hub.json',dict(verified=True,revision=commit.oid,report=f'https://huggingface.co/{MODEL_REPO}/blob/{commit.oid}/evaluation/analysis/REPORT.md'))
    return commit.oid


def main(point):
    import sdm,wandb
    for name in ['features','predictions','analysis']:(EVAL/name).mkdir(parents=True,exist_ok=True)
    lock=(OUT/'probe.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX|fcntl.LOCK_NB)
    checkpoint=OUT/'checkpoints'/f'probe_{point:03d}';assert json.loads((checkpoint/'hub_verified.json').read_text())['verified']
    train=pd.read_csv(WORK/'bigsoldb_evaluation/data/train.csv');test=pd.read_csv(WORK/'bigsoldb_evaluation/data/test.csv')
    indices=np.load(WORK/'bigsoldb_evaluation/features/pair_indices.npz')
    assert np.array_equal(indices['test_clean_rows'],test['_clean_row']) and np.array_equal(indices['train_clean_rows'],train['_clean_row'])
    assert len(train)==75731 and len(test)==9686
    records=pd.read_csv(EVAL/'results.csv').to_dict('records') if (EVAL/'results.csv').exists() else []
    wb_path=OUT/'probe_wandb.json';wb=json.loads(wb_path.read_text()) if wb_path.exists() else dict(run_id=wandb.util.generate_id())
    write(wb_path,wb)
    run=wandb.init(entity='queimo',project='rem3di-reproduction',id=wb['run_id'],resume='allow',
        name='scratch-pcqm-cosmo-fifty-bigsoldb-kumo-large-probes',group='rem3di-scratch-pcqm-cosmo',job_type='frozen-probing',
        dir=str(OUT),config=CONFIG,settings=wandb.Settings(init_timeout=90))
    write(wb_path,dict(run_id=run.id,url=run.url));reference.status=status
    try:
        if not any(int(r['point'])==point for r in records):
            status('embedding_for_probe',point=point,wandb_url=run.url)
            z,hashes=embeddings(point)
            X=np.concatenate([z[indices['solute']],z[indices['solvent']],indices['temperature_K'][:,None]],axis=1).astype(np.float32)
            y=pd.concat([train,test])['LogS(mol/L)'].to_numpy(dtype=np.float64)
            assert X.shape==(len(y),1281) and np.isfinite(X).all()
            torch.manual_seed(42);torch.cuda.manual_seed_all(42);np.random.seed(42);random.seed(42)
            model=sdm.models.KumoTabular(task='regression',size='large',device='cuda');start=time.perf_counter()
            status('fitting_kumo',point=point,wandb_url=run.url)
            shapes=reference.fit_context(model,X,y,len(train),1);pred=reference.predict(model,X[len(train):])
            model.clear();del model,X,z;gc.collect();torch.cuda.empty_cache()
            metadata=json.loads((checkpoint/'metrics.json').read_text())
            row=dict(point=point,epoch_position=metadata['epoch_position'],nominal_epoch=metadata['nominal_epoch'],
                global_step=metadata['global_step'],n_train=len(train),n_test=len(test),raw_columns=1281,
                context_blocks=1,num_estimators=4,seed=42,seconds=time.perf_counter()-start,**reference.metrics(y[len(train):],pred))
            np.savez_compressed(EVAL/'predictions'/f'probe_{point:03d}.npz',y_true=y[len(train):],y_pred=pred,
                test_clean_rows=test['_clean_row'].to_numpy(),test_raw_rows=test['_raw_row'].to_numpy(),checkpoint_sha256=json.dumps(hashes))
            write(EVAL/'features'/f'probe_{point:03d}_context.json',dict(full_context=True,processed_shapes=shapes,checkpoint_sha256=hashes))
            records.append(row);pd.DataFrame(records).sort_values('point').to_csv(EVAL/'results.csv',index=False)
            run.log(row)
        for row in records:
            saved=np.load(EVAL/'predictions'/f'probe_{int(row["point"]):03d}.npz')
            assert np.array_equal(saved['y_true'],test['LogS(mol/L)']) and np.array_equal(saved['test_clean_rows'],test['_clean_row'])
            for key,value in reference.metrics(saved['y_true'],saved['y_pred']).items():assert np.isclose(value,row[key],rtol=1e-8,atol=1e-9)
        write(EVAL/'validation.json',dict(passed=True,probes=len(records),all_six_metrics_recomputed=True,exact_test_rows=True))
        write(EVAL/'protocol.json',dict(train_rows=len(train),test_rows=len(test),split_sha256={n:sha(WORK/'bigsoldb_evaluation/data'/n) for n in ['train.csv','test.csv']},
            pair_indices_sha256=sha(WORK/'bigsoldb_evaluation/features/pair_indices.npz'),kumo=CONFIG['probe']))
        build_report(run);revision=publish(run)
        artifact=wandb.Artifact(f'scratch-pcqm-cosmo-bigsoldb-point-{point:03d}',type='evaluation')
        for p in [EVAL/'predictions'/f'probe_{point:03d}.npz',EVAL/'results.csv',EVAL/'analysis/REPORT.md',EVAL/'analysis/learning_curve.png']:
            artifact.add_file(str(p),name=str(p.relative_to(EVAL)))
        run.log_artifact(artifact,aliases=['complete']).wait(timeout=60)
        write(OUT/'latest_probe.json',dict(point=point,total=50,revision=revision,metrics=next(r for r in records if int(r['point'])==point)))
        write(EVAL/f'probe_{point:03d}_complete.json',dict(completed=True,revision=revision))
        run.finish();status('probe_complete',point=point,probes=len(records))
    except Exception:run.finish(exit_code=1);raise


if __name__=='__main__':
    try:main(int(sys.argv[1]))
    except Exception:status('failed_probe',error=traceback.format_exc());raise

	"""One full-context BigSolDB Kumo Large probe per saved fifth-epoch checkpoint."""
	from pathlib import Path
	import sys
	sys.path.insert(0,str(Path(__file__).resolve().parent))
	from scratch_model import *
	from scratch_train import status
	import bigsoldb_evaluate as reference
	import random,gc,fcntl,argparse,traceback,shutil
	from huggingface_hub import HfApi,CommitOperationAdd,hf_hub_download

	EVAL=OUT/'evaluation'


	def embeddings(point):
	checkpoint=OUT/'checkpoints'/f'probe_{point:03d}';path=EVAL/'features'/f'probe_{point:03d}.npz'
	hashes={n:sha(checkpoint/n) for n in ['encoder.pth','atomic_preprocessor.pth','geometric_preprocessor.pth']}
	geometry=pd.read_parquet(DATA/'bigsoldb_fixed_geometries.parquet')
	if path.exists():
	saved=np.load(path)
	assert saved['smiles'].tolist()==geometry.smiles.tolist()
	assert json.loads(str(saved['checkpoint_sha256']))==hashes
	return saved['X'],hashes
	model=ScratchJoint(checkpoint).eval().requires_grad_(False)
	cache=RawCache('bigsoldb');z=np.empty((len(geometry),640),dtype=np.float32)
	with torch.inference_mode():
	for ids in batch_indices(np.arange(len(geometry)),cache.lengths,max_size=24,atom_cap=2400):
	_,_,values=model.encode(cache.sample(ids));z[ids]=values.float().cpu().numpy()
	assert np.isfinite(z).all() and z.shape==(1334,640)
	np.savez_compressed(path,X=z,smiles=geometry.smiles.to_numpy(dtype=str),checkpoint_sha256=json.dumps(hashes))
	del model,cache;gc.collect();torch.cuda.empty_cache()
	return z,hashes


	def build_report(run):
	import matplotlib
	matplotlib.use('Agg')
	import matplotlib.pyplot as plt
	rows=pd.read_csv(EVAL/'results.csv').sort_values('point')
	prior_path=OUT/'prior_comparison.csv'
	if not prior_path.exists():
	p=hf_hub_download('Queimo/rem3di-cosmo10-combisolv-multitask','ablations/references/comparison.csv',
	revision='ff86641eb441ffe038c8aca1df1ba32c02f74cb7')
	shutil.copyfile(p,prior_path)
	prior=pd.read_csv(prior_path)
	selected=prior.loc[prior.variant.isin(['initialization','published_rem3di','monroe','literature'])]
	fig,axes=plt.subplots(1,3,figsize=(13,3.8),layout='constrained')
	for ax,metric in zip(axes,['mae','rmse','r2']):
	ax.plot(rows.epoch_position,rows[metric],'-o',ms=3,label='Scratch REM3DI + COSMO + Kumo L')
	for variant,label,color in [('initialization','PCQM2→COSMO10','#888888'),('published_rem3di','Published Rem3Di','#332288'),('monroe','Monroe','#CC79A7')]:
	ax.axhline(float(prior.loc[prior.variant==variant,metric].iloc[0]),ls='--',color=color,label=label)
	ax.set_xlabel('Joint pretraining epoch');ax.set_ylabel(metric.upper());ax.grid(alpha=.2)
	axes[0].legend(fontsize=6);fig.savefig(EVAL/'analysis/learning_curve.png',dpi=180)
	fig.savefig(EVAL/'analysis/learning_curve.svg');plt.close(fig)
	lines=['# Random-initialized Rem3Di + COSMO: ten epochs, fifty BigSolDB probes','',
	f'Completed {len(rows)}/50 planned fifth-epoch frozen Kumo Large probes. Every learned Rem3Di preprocessor, encoder, denoising decoder and COSMO head started from its constructor initialization. No published Rem3Di or PCQM/COSMO checkpoint initializes this run. MACE-POLAR-1-M alone is pretrained and frozen. The encoder/decoder architecture and 640-dimensional pooled descriptor match the preceding experiments.','',
	'Two simultaneous streams: standard author atomic-embedding denoising on the retained PCQM500k structures, and 153 sigma-bin areas plus A/V reconstruction on the retained original PM6 COSMO structures. Each epoch covers both complete training streams; the shorter stream cycles. No CombiSolv head or covariance/decorrelation penalty is included.','',
	'The released denoising implementation, Gaussian sigma 0.3 and its original gradient paths are preserved. In particular the clean atomic reconstruction target remains differentiable, as in the author implementation. Raw and per-channel denoising losses, held-out descriptor variability and target variability are logged; low reconstruction loss alone is not treated as downstream quality.','',
	'Task uncertainty weighting uses sum(L/(2 exp(s)) + s/2). The unmodified raw denoising loss has 1088 channel terms, so its initial log variance is ln(1088)−0.5; COSMO starts at −0.5. This is mathematically equivalent, up to an additive constant, to starting both log variances at −0.5 with per-channel denoising units. All invariant feature normalization and COSMO log1p/ln target standardization were freshly fit using training rows only.','',
	'Held-out BigSolDB solute connectivity groups are excluded from both new training streams. PCQM and COSMO validation groups are also excluded across objectives. The pretrained MACE foundation can have prior exposure; this experiment makes no claim about removing its historical pretraining data. Exact memberships and all original coordinates are preserved in the prepared dataset.','',
	'Every fifth-epoch checkpoint is saved with all heads, task weights, AdamW state, OneCycle scheduler and RNG state, and verified against its Hugging Face manifest before probing. Probes run in separate processes and never change the training RNG, optimizer or schedule. There are exactly five probes per epoch, fifty total; the epoch-10 endpoint is the predeclared final checkpoint, and all intermediate test scores are reported without test-based tuning.','',
	'\| Point \| Epoch position \| MAE \| RMSE \| R² \|','\|---:\|---:\|---:\|---:\|---:\|']
	for r in rows.itertuples():lines.append(f'\| {r.point} \| {r.epoch_position:.5f} \| {r.mae:.6f} \| {r.rmse:.6f} \| {r.r2:.6f} \|')
	lines+=['','## Prior matched probes and literature values','','\| Model \| Origin \| MAE \| RMSE \| R² \|','\|---\|---\|---:\|---:\|---:\|']
	for r in selected.itertuples():lines.append(f'\| {r.model} \| {r.origin} \| {r.mae:.6f} \| {r.rmse:.6f} \| {r.r2:.6f} \|')
	lines+=['','All measured probes use exactly the released 9,686 test rows and the 75,731-row non-test complement as full Kumo context. Fixed shared baseline conformers, 1281 ordered solute640/solvent640/temperature-in-Kelvin features, Kumo Large v1.0.1, four estimators, seed 42, fp16, query chunks 128. No additional transformation of the supplied log10 mol/L labels. Monroe uses its full 1441 features. Literature rows are copied reported values, not rerun models. The author prediction file has 9,686 rows and reproduces Table 1 metrics; the prose says 9,250. The private neural validation subset is unavailable.','',
	f'[Prepared data](https://huggingface.co/datasets/{DATA_REPO}) · [Checkpoints](https://huggingface.co/{MODEL_REPO}) · [Probe W&B]({run.url}) · [Solvaformer Table 1]({reference.PAPER_URL})','',
	'![Fifty-point learning curve](learning_curve.png)','']
	text='\n'.join(lines);(EVAL/'analysis/REPORT.md').write_text(text)
	(OUT/'README.md').write_text('---\ntags:\n- chemistry\n- multitask\n---\n\n'+text.replace('](learning_curve.png)','](evaluation/analysis/learning_curve.png)'))
	return rows


	def publish(run):
	files=[p for p in EVAL.rglob('*') if p.is_file() and p.suffix in ['.json','.csv','.npz','.png','.svg','.md']]
	files += [p for p in OUT.iterdir() if p.is_file() and p.suffix in ['.json','.csv','.yaml','.md']]
	sources=[Path(__file__).with_name(n) for n in ['scratch_prepare.py','scratch_data.py','scratch_model.py','scratch_train.py','scratch_probe.py','scratch_driver.py','scratch_cache_initial.py','scratch_fast_cache.py',
	'eval_common.py','train_cosmo.py','multitask_data.py','bigsoldb_evaluate.py']]
	ops=[CommitOperationAdd(path_in_repo=str(p.relative_to(OUT)),path_or_fileobj=str(p)) for p in files]
	ops += [CommitOperationAdd(path_in_repo='source/'+p.name,path_or_fileobj=str(p)) for p in sources]
	commit=HfApi().create_commit(repo_id=MODEL_REPO,operations=ops,commit_message='Update scratch-training learning curve and exact fifth-epoch probe artifacts')
	write(OUT/'report_hub.json',dict(verified=True,revision=commit.oid,report=f'https://huggingface.co/{MODEL_REPO}/blob/{commit.oid}/evaluation/analysis/REPORT.md'))
	return commit.oid


	def main(point):
	import sdm,wandb
	for name in ['features','predictions','analysis']:(EVAL/name).mkdir(parents=True,exist_ok=True)
	lock=(OUT/'probe.lock').open('a');fcntl.flock(lock,fcntl.LOCK_EX\|fcntl.LOCK_NB)
	checkpoint=OUT/'checkpoints'/f'probe_{point:03d}';assert json.loads((checkpoint/'hub_verified.json').read_text())['verified']
	train=pd.read_csv(WORK/'bigsoldb_evaluation/data/train.csv');test=pd.read_csv(WORK/'bigsoldb_evaluation/data/test.csv')
	indices=np.load(WORK/'bigsoldb_evaluation/features/pair_indices.npz')
	assert np.array_equal(indices['test_clean_rows'],test['_clean_row']) and np.array_equal(indices['train_clean_rows'],train['_clean_row'])
	assert len(train)==75731 and len(test)==9686
	records=pd.read_csv(EVAL/'results.csv').to_dict('records') if (EVAL/'results.csv').exists() else []
	wb_path=OUT/'probe_wandb.json';wb=json.loads(wb_path.read_text()) if wb_path.exists() else dict(run_id=wandb.util.generate_id())
	write(wb_path,wb)
	run=wandb.init(entity='queimo',project='rem3di-reproduction',id=wb['run_id'],resume='allow',
	name='scratch-pcqm-cosmo-fifty-bigsoldb-kumo-large-probes',group='rem3di-scratch-pcqm-cosmo',job_type='frozen-probing',
	dir=str(OUT),config=CONFIG,settings=wandb.Settings(init_timeout=90))
	write(wb_path,dict(run_id=run.id,url=run.url));reference.status=status
	try:
	if not any(int(r['point'])==point for r in records):
	status('embedding_for_probe',point=point,wandb_url=run.url)
	z,hashes=embeddings(point)
	X=np.concatenate([z[indices['solute']],z[indices['solvent']],indices['temperature_K'][:,None]],axis=1).astype(np.float32)
	y=pd.concat([train,test])['LogS(mol/L)'].to_numpy(dtype=np.float64)
	assert X.shape==(len(y),1281) and np.isfinite(X).all()
	torch.manual_seed(42);torch.cuda.manual_seed_all(42);np.random.seed(42);random.seed(42)
	model=sdm.models.KumoTabular(task='regression',size='large',device='cuda');start=time.perf_counter()
	status('fitting_kumo',point=point,wandb_url=run.url)
	shapes=reference.fit_context(model,X,y,len(train),1);pred=reference.predict(model,X[len(train):])
	model.clear();del model,X,z;gc.collect();torch.cuda.empty_cache()
	metadata=json.loads((checkpoint/'metrics.json').read_text())
	row=dict(point=point,epoch_position=metadata['epoch_position'],nominal_epoch=metadata['nominal_epoch'],
	global_step=metadata['global_step'],n_train=len(train),n_test=len(test),raw_columns=1281,
	context_blocks=1,num_estimators=4,seed=42,seconds=time.perf_counter()-start,**reference.metrics(y[len(train):],pred))
	np.savez_compressed(EVAL/'predictions'/f'probe_{point:03d}.npz',y_true=y[len(train):],y_pred=pred,
	test_clean_rows=test['_clean_row'].to_numpy(),test_raw_rows=test['_raw_row'].to_numpy(),checkpoint_sha256=json.dumps(hashes))
	write(EVAL/'features'/f'probe_{point:03d}_context.json',dict(full_context=True,processed_shapes=shapes,checkpoint_sha256=hashes))
	records.append(row);pd.DataFrame(records).sort_values('point').to_csv(EVAL/'results.csv',index=False)
	run.log(row)
	for row in records:
	saved=np.load(EVAL/'predictions'/f'probe_{int(row["point"]):03d}.npz')
	assert np.array_equal(saved['y_true'],test['LogS(mol/L)']) and np.array_equal(saved['test_clean_rows'],test['_clean_row'])
	for key,value in reference.metrics(saved['y_true'],saved['y_pred']).items():assert np.isclose(value,row[key],rtol=1e-8,atol=1e-9)
	write(EVAL/'validation.json',dict(passed=True,probes=len(records),all_six_metrics_recomputed=True,exact_test_rows=True))
	write(EVAL/'protocol.json',dict(train_rows=len(train),test_rows=len(test),split_sha256={n:sha(WORK/'bigsoldb_evaluation/data'/n) for n in ['train.csv','test.csv']},
	pair_indices_sha256=sha(WORK/'bigsoldb_evaluation/features/pair_indices.npz'),kumo=CONFIG['probe']))
	build_report(run);revision=publish(run)
	artifact=wandb.Artifact(f'scratch-pcqm-cosmo-bigsoldb-point-{point:03d}',type='evaluation')
	for p in [EVAL/'predictions'/f'probe_{point:03d}.npz',EVAL/'results.csv',EVAL/'analysis/REPORT.md',EVAL/'analysis/learning_curve.png']:
	artifact.add_file(str(p),name=str(p.relative_to(EVAL)))
	run.log_artifact(artifact,aliases=['complete']).wait(timeout=60)
	write(OUT/'latest_probe.json',dict(point=point,total=50,revision=revision,metrics=next(r for r in records if int(r['point'])==point)))
	write(EVAL/f'probe_{point:03d}_complete.json',dict(completed=True,revision=revision))
	run.finish();status('probe_complete',point=point,probes=len(records))
	except Exception:run.finish(exit_code=1);raise


	if __name__=='__main__':
	try:main(int(sys.argv[1]))
	except Exception:status('failed_probe',error=traceback.format_exc());raise