File size: 4,214 Bytes
aa45001 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | 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}.')
|