| import argparse |
| import csv |
| import json |
| from pathlib import Path |
|
|
|
|
| def parser(mode): |
| p = argparse.ArgumentParser(description=f'Metalorian {mode} sampling') |
| p.add_argument('--target_idx', type=int, required=True, help='Metal index; Cu=3, Cd=1, Zn=13, Ni=8') |
| p.add_argument('--seq_len', type=int, default=30) |
| p.add_argument('--batch', type=int, default=1) |
| p.add_argument('--desired', type=int, default=1) |
| p.add_argument('--max_rounds', type=int, default=150) |
| p.add_argument('--grad_steps', type=int, default=4) |
| p.add_argument('--guidance_scale', type=float, default=6.0) |
| p.add_argument('--penalty', type=float, default=0.1) |
| p.add_argument('--tau', type=float, default=1.0) |
| p.add_argument('--label_clamp', type=float, default=0.8) |
| p.add_argument('--best_decode', type=int, default=16, help='Candidate decodes; 0 uses argmax') |
| p.add_argument('--decode_temp', type=float, default=1.0) |
| p.add_argument('--device', default='cuda:0') |
| p.add_argument('--seed', type=int, default=42) |
| p.add_argument('--output', type=Path, required=True) |
| if mode == 'prototype': |
| p.add_argument('--proto', required=True) |
| p.add_argument('--t_start', type=int, default=15) |
| p.set_defaults(guide_on='x0', no_accum=False, track_best=False, accept_all=False, |
| amp=False, out_prefix=None, cys_logit_bias=0.0, cys_target=0.0, |
| cys_weight=0.0, min_cys=0, max_aa_frac=1.0, relax_after=10**9, |
| init_noise=0.3) |
| return p |
|
|
|
|
| def collect(args, sampler_class, mode): |
| from .sampling import FLAGS, METAL_LABELS |
| if not 0 <= args.target_idx < 15 or not 1 <= args.seq_len <= 1022: |
| raise ValueError('target_idx must be 0–14; seq_len must be 1–1022') |
| if min(args.batch, args.desired, args.max_rounds) < 1: |
| raise ValueError('batch, desired and max_rounds must be positive') |
| if mode == 'prototype': |
| if not 0 < args.t_start < 50: |
| raise ValueError('t_start must be 1–49') |
| if not args.proto or set(args.proto)-set('ACDEFGHIKLMNPQRSTVWY'): |
| raise ValueError('proto must contain standard uppercase amino acids') |
| import random |
| import numpy as np |
| import torch |
| random.seed(args.seed) |
| np.random.seed(args.seed) |
| torch.manual_seed(args.seed) |
| if torch.cuda.is_available(): |
| torch.cuda.manual_seed_all(args.seed) |
| torch.backends.cuda.matmul.allow_tf32 = True |
| torch.backends.cudnn.allow_tf32 = True |
| sampler = sampler_class(args, device=args.device) |
| rows, seen, candidates = [], set(), 0 |
| for round_index in range(1, args.max_rounds+1): |
| sequences, _ = sampler.run_round() |
| sequences = [s for s in sequences if s] |
| candidates += len(sequences) |
| hits, _ = sampler.accept(sequences, None) if sequences else ([], None) |
| for hit in hits: |
| if hit['sequence'] in seen: |
| continue |
| seen.add(hit['sequence']) |
| rows.append(dict(metal=METAL_LABELS[args.target_idx], **hit)) |
| if len(rows) == args.desired: |
| break |
| print(f'round={round_index} candidates={candidates} unique_accepted={len(rows)}/{args.desired}', flush=True) |
| if len(rows) >= args.desired: |
| break |
| args.output.parent.mkdir(parents=True, exist_ok=True) |
| with args.output.open('w', newline='') as f: |
| writer = csv.DictWriter(f, fieldnames=['metal','sequence','score','n_cys','n_his']) |
| writer.writeheader() |
| writer.writerows(rows) |
| metadata = dict(mode=mode, arguments={k:str(v) if isinstance(v,Path) else v for k,v in vars(args).items()}, |
| candidates=candidates, unique_accepted=len(rows), rounds=round_index, |
| target_threshold=sampler.threshold, diffusion_steps=int(FLAGS.T), |
| beta_start=float(FLAGS.beta_1), beta_end=float(FLAGS.beta_T), |
| predictor='checkpoints/metalatte_predictor.ckpt', diffusion='checkpoints/metalorian_diffusion.pt') |
| args.output.with_suffix('.metadata.json').write_text(json.dumps(metadata, indent=2)+'\n') |
| print(f'Wrote {len(rows)} sequences to {args.output}; requested {args.desired}.') |
|
|
