| import random, os
|
| import jax
|
| import jax.numpy as jnp
|
| import numpy as np
|
| from colabdesign.af.alphafold.common import residue_constants
|
| from colabdesign.shared.utils import copy_dict, update_dict, Key, dict_to_str, to_float, softmax, categorical, to_list, copy_missing
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
|
| class _af_design:
|
|
|
| def restart(self, seed=None, opt=None, weights=None,
|
| seq=None, mode=None, keep_history=False, reset_opt=True, **kwargs):
|
| '''
|
| restart the optimization
|
| ------------
|
| note: model.restart() resets the [opt]ions and weights to their defaults
|
| use model.set_opt(..., set_defaults=True) and model.set_weights(..., set_defaults=True)
|
| or model.restart(reset_opt=False) to avoid this
|
| ------------
|
| seed=0 - set seed for reproducibility
|
| reset_opt=False - do NOT reset [opt]ions/weights to defaults
|
| keep_history=True - do NOT clear the trajectory/[opt]ions/weights
|
| '''
|
|
|
| if reset_opt and not keep_history:
|
| copy_missing(self.opt, self._opt)
|
| self.opt = copy_dict(self._opt)
|
| if hasattr(self,"aux"): del self.aux
|
|
|
| if not keep_history:
|
|
|
| self._tmp = {"traj":{"seq":[],"xyz":[],"plddt":[],"pae":[]},
|
| "log":[],"best":{}}
|
|
|
|
|
| self.set_opt(opt)
|
| self.set_weights(weights)
|
|
|
|
|
| self.set_seed(seed)
|
| self.set_seq(seq=seq, mode=mode, **kwargs)
|
|
|
|
|
| self._k = 0
|
| self.set_optimizer()
|
|
|
| def _get_model_nums(self, num_models=None, sample_models=None, models=None):
|
| '''decide which model params to use'''
|
| if num_models is None: num_models = self.opt["num_models"]
|
| if sample_models is None: sample_models = self.opt["sample_models"]
|
|
|
| ns_name = self._model_names
|
| ns = list(range(len(ns_name)))
|
| if models is not None:
|
| models = models if isinstance(models,list) else [models]
|
| ns = [ns[n if isinstance(n,int) else ns_name.index(n)] for n in models]
|
|
|
| m = min(num_models,len(ns))
|
| if sample_models and m != len(ns):
|
| model_nums = np.random.choice(ns,(m,),replace=False)
|
| else:
|
| model_nums = ns[:m]
|
| return model_nums
|
|
|
| def run(self, num_recycles=None, num_models=None, sample_models=None, models=None,
|
| backprop=True, callback=None, model_nums=None, return_aux=False):
|
| '''run model to get outputs, losses and gradients'''
|
|
|
|
|
| for fn in self._callbacks["design"]["pre"]: fn(self)
|
|
|
|
|
| if model_nums is None:
|
| model_nums = self._get_model_nums(num_models, sample_models, models)
|
| assert len(model_nums) > 0, "ERROR: no model params defined"
|
|
|
|
|
| auxs = []
|
| for n in model_nums:
|
| p = self._model_params[n]
|
| auxs.append(self._recycle(p, num_recycles=num_recycles, backprop=backprop))
|
| auxs = jax.tree_util.tree_map(lambda *x: np.stack(x), *auxs)
|
|
|
|
|
| def avg_or_first(x):
|
| if np.issubdtype(x.dtype, np.integer): return x[0]
|
| else: return x.mean(0)
|
|
|
| self.aux = jax.tree_util.tree_map(avg_or_first, auxs)
|
| self.aux["atom_positions"] = auxs["atom_positions"][0]
|
| self.aux["all"] = auxs
|
|
|
|
|
| for fn in (self._callbacks["design"]["post"] + to_list(callback)): fn(self)
|
|
|
|
|
| self.aux["log"] = {**self.aux["losses"]}
|
| self.aux["log"]["plddt"] = 1 - self.aux["log"]["plddt"]
|
| for k in ["loss","i_ptm","ptm"]: self.aux["log"][k] = self.aux[k]
|
| for k in ["hard","soft","temp"]: self.aux["log"][k] = self.opt[k]
|
|
|
|
|
| if self.protocol in ["fixbb","partial"] or (self.protocol == "binder" and self._args["redesign"]):
|
| if self.protocol == "partial":
|
| aatype = self.aux["aatype"][...,self.opt["pos"]]
|
| else:
|
| aatype = self.aux["seq"]["pseudo"].argmax(-1)
|
|
|
| mask = self._wt_aatype != -1
|
| true = self._wt_aatype[mask]
|
| pred = aatype[...,mask]
|
| self.aux["log"]["seqid"] = (true == pred).mean()
|
|
|
| self.aux["log"] = to_float(self.aux["log"])
|
| self.aux["log"].update({"recycles":int(self.aux["num_recycles"]),
|
| "models":model_nums})
|
|
|
| if return_aux: return self.aux
|
|
|
| def _single(self, model_params, backprop=True):
|
| '''single pass through the model'''
|
| self._inputs["opt"] = self.opt
|
| flags = [self._params, model_params, self._inputs, self.key()]
|
| if backprop:
|
| (loss, aux), grad = self._model["grad_fn"](*flags)
|
| else:
|
| loss, aux = self._model["fn"](*flags)
|
| grad = jax.tree_util.tree_map(np.zeros_like, self._params)
|
| aux.update({"loss":loss,"grad":grad})
|
| return aux
|
|
|
| def _recycle(self, model_params, num_recycles=None, backprop=True):
|
| '''multiple passes through the model (aka recycle)'''
|
| a = self._args
|
| mode = a["recycle_mode"]
|
| if num_recycles is None:
|
| num_recycles = self.opt["num_recycles"]
|
|
|
| if mode in ["backprop","add_prev"]:
|
|
|
| aux = self._single(model_params, backprop)
|
|
|
| else:
|
| L = self._inputs["residue_index"].shape[0]
|
|
|
|
|
| if "prev" not in self._inputs or a["clear_prev"]:
|
| prev = {'prev_msa_first_row': np.zeros([L,256]),
|
| 'prev_pair': np.zeros([L,L,128])}
|
|
|
| if a["use_initial_guess"] and "batch" in self._inputs:
|
| prev["prev_pos"] = self._inputs["batch"]["all_atom_positions"]
|
| else:
|
| prev["prev_pos"] = np.zeros([L,37,3])
|
|
|
| if a["use_dgram"]:
|
|
|
| prev["prev_dgram"] = np.zeros([L,L,64])
|
|
|
| if a["use_initial_atom_pos"]:
|
| if "batch" in self._inputs:
|
| self._inputs["initial_atom_pos"] = self._inputs["batch"]["all_atom_positions"]
|
| else:
|
| self._inputs["initial_atom_pos"] = np.zeros([L,37,3])
|
|
|
| self._inputs["prev"] = prev
|
|
|
| cycles = (num_recycles + 1)
|
| mask = [0] * cycles
|
|
|
| if mode == "sample": mask[np.random.randint(0,cycles)] = 1
|
| if mode == "average": mask = [1/cycles] * cycles
|
| if mode == "last": mask[-1] = 1
|
| if mode == "first": mask[0] = 1
|
|
|
|
|
| grad = []
|
| for m in mask:
|
| if m == 0:
|
| aux = self._single(model_params, backprop=False)
|
| else:
|
| aux = self._single(model_params, backprop)
|
| grad.append(jax.tree_util.tree_map(lambda x:x*m, aux["grad"]))
|
| self._inputs["prev"] = aux["prev"]
|
| if a["use_initial_atom_pos"]:
|
| self._inputs["initial_atom_pos"] = aux["prev"]["prev_pos"]
|
|
|
| aux["grad"] = jax.tree_util.tree_map(lambda *x: np.stack(x).sum(0), *grad)
|
|
|
| aux["num_recycles"] = num_recycles
|
| return aux
|
|
|
| def step(self, lr_scale=1.0, num_recycles=None,
|
| num_models=None, sample_models=None, models=None, backprop=True,
|
| callback=None, save_best=False, verbose=1):
|
| '''do one step of gradient descent'''
|
|
|
|
|
| self.run(num_recycles=num_recycles, num_models=num_models, sample_models=sample_models,
|
| models=models, backprop=backprop, callback=callback)
|
|
|
|
|
| if self.opt["norm_seq_grad"]: self._norm_seq_grad()
|
| self._state, self.aux["grad"] = self._optimizer(self._state, self.aux["grad"], self._params)
|
|
|
|
|
| lr = self.opt["learning_rate"] * lr_scale
|
| self._params = jax.tree_util.tree_map(lambda x,g:x-lr*g, self._params, self.aux["grad"])
|
|
|
|
|
| self._save_results(save_best=save_best, verbose=verbose)
|
|
|
|
|
| self._k += 1
|
|
|
| def _print_log(self, print_str=None, aux=None):
|
| if aux is None: aux = self.aux
|
| keys = ["models","recycles","hard","soft","temp","seqid","loss",
|
| "seq_ent","mlm","helix","pae","i_pae","exp_res","con","i_con",
|
| "sc_fape","sc_rmsd","dgram_cce","fape","plddt","ptm"]
|
|
|
| if "i_ptm" in aux["log"]:
|
| if len(self._lengths) > 1:
|
| keys.append("i_ptm")
|
| else:
|
| aux["log"].pop("i_ptm")
|
|
|
| print(dict_to_str(aux["log"], filt=self.opt["weights"],
|
| print_str=print_str, keys=keys+["rmsd"], ok=["plddt","rmsd"]))
|
|
|
| def _save_results(self, aux=None, save_best=False,
|
| best_metric=None, metric_higher_better=False,
|
| verbose=True):
|
| if aux is None: aux = self.aux
|
| self._tmp["log"].append(aux["log"])
|
| if (self._k % self._args["traj_iter"]) == 0:
|
|
|
| traj = {"seq": aux["seq"]["pseudo"],
|
| "xyz": aux["atom_positions"][:,1,:],
|
| "plddt": aux["plddt"],
|
| "pae": aux["pae"]}
|
| for k,v in traj.items():
|
| if len(self._tmp["traj"][k]) == self._args["traj_max"]:
|
| self._tmp["traj"][k].pop(0)
|
| self._tmp["traj"][k].append(v)
|
|
|
|
|
| if save_best:
|
| if best_metric is None:
|
| best_metric = self._args["best_metric"]
|
| metric = float(aux["log"][best_metric])
|
| if self._args["best_metric"] in ["plddt","ptm","i_ptm","seqid","composite"] or metric_higher_better:
|
| metric = -metric
|
| if "metric" not in self._tmp["best"] or metric < self._tmp["best"]["metric"]:
|
| self._tmp["best"]["aux"] = copy_dict(aux)
|
| self._tmp["best"]["metric"] = metric
|
|
|
| if verbose and ((self._k+1) % verbose) == 0:
|
| self._print_log(f"{self._k+1}", aux=aux)
|
|
|
| def predict(self, seq=None, bias=None,
|
| num_models=None, num_recycles=None, models=None, sample_models=False,
|
| dropout=False, hard=True, soft=False, temp=1,
|
| return_aux=False, verbose=True, seed=None, **kwargs):
|
| '''predict structure for input sequence (if provided)'''
|
|
|
| def load_settings():
|
| if "save" in self._tmp:
|
| [self.opt, self._args, self._params, self._inputs] = self._tmp.pop("save")
|
|
|
| def save_settings():
|
| load_settings()
|
| self._tmp["save"] = [copy_dict(x) for x in [self.opt, self._args, self._params, self._inputs]]
|
|
|
| save_settings()
|
|
|
|
|
| if seed is not None: self.set_seed(seed)
|
|
|
|
|
| if seq is not None: self.set_seq(seq=seq, bias=bias)
|
| self.set_opt(hard=hard, soft=soft, temp=temp, dropout=dropout, pssm_hard=True)
|
| self.set_args(shuffle_first=False)
|
|
|
|
|
| self.run(num_recycles=num_recycles, num_models=num_models,
|
| sample_models=sample_models, models=models, backprop=False, **kwargs)
|
| if verbose: self._print_log("predict")
|
|
|
| load_settings()
|
|
|
|
|
| if return_aux: return self.aux
|
|
|
|
|
|
|
|
|
| def design(self, iters=100,
|
| soft=0.0, e_soft=None,
|
| temp=1.0, e_temp=None,
|
| hard=0.0, e_hard=None,
|
| step=1.0, e_step=None,
|
| dropout=True, opt=None, weights=None,
|
| num_recycles=None, ramp_recycles=False,
|
| num_models=None, sample_models=None, models=None,
|
| backprop=True, callback=None, save_best=False, verbose=1):
|
|
|
|
|
| self.set_opt(opt, dropout=dropout)
|
| self.set_weights(weights)
|
| m = {"soft":[soft,e_soft],"temp":[temp,e_temp],
|
| "hard":[hard,e_hard],"step":[step,e_step]}
|
| m = {k:[s,(s if e is None else e)] for k,(s,e) in m.items()}
|
|
|
| if ramp_recycles:
|
| if num_recycles is None:
|
| num_recycles = self.opt["num_recycles"]
|
| m["num_recycles"] = [0,num_recycles]
|
|
|
| for i in range(iters):
|
| for k,(s,e) in m.items():
|
| if k == "temp":
|
| self.set_opt({k:(e+(s-e)*(1-(i+1)/iters)**2)})
|
| else:
|
| v = (s+(e-s)*((i+1)/iters))
|
| if k == "step": step = v
|
| elif k == "num_recycles": num_recycles = round(v)
|
| else: self.set_opt({k:v})
|
|
|
|
|
| lr_scale = step * ((1 - self.opt["soft"]) + (self.opt["soft"] * self.opt["temp"]))
|
|
|
| self.step(lr_scale=lr_scale, num_recycles=num_recycles,
|
| num_models=num_models, sample_models=sample_models, models=models,
|
| backprop=backprop, callback=callback, save_best=save_best, verbose=verbose)
|
|
|
| def design_logits(self, iters=100, **kwargs):
|
| ''' optimize logits '''
|
| self.design(iters, **kwargs)
|
|
|
| def design_soft(self, iters=100, temp=1, **kwargs):
|
| ''' optimize softmax(logits/temp)'''
|
| self.design(iters, soft=1, temp=temp, **kwargs)
|
|
|
| def design_hard(self, iters=100, **kwargs):
|
| ''' optimize argmax(logits) '''
|
| self.design(iters, soft=1, hard=1, **kwargs)
|
|
|
|
|
|
|
|
|
| def design_3stage(self, soft_iters=300, temp_iters=100, hard_iters=10,
|
| ramp_recycles=True, **kwargs):
|
| '''three stage design (logits→soft→hard)'''
|
|
|
| verbose = kwargs.get("verbose",1)
|
|
|
|
|
| if soft_iters > 0:
|
| if verbose: print("Stage 1: running (logits → soft)")
|
| self.design_logits(soft_iters, e_soft=1,
|
| ramp_recycles=ramp_recycles, **kwargs)
|
| self._tmp["seq_logits"] = self.aux["seq"]["logits"]
|
|
|
|
|
| if temp_iters > 0:
|
| if verbose: print("Stage 2: running (soft → hard)")
|
| self.design_soft(temp_iters, e_temp=1e-2, **kwargs)
|
|
|
|
|
| if hard_iters > 0:
|
| if verbose: print("Stage 3: running (hard)")
|
| kwargs["dropout"] = False
|
| kwargs["save_best"] = True
|
| kwargs["num_models"] = len(self._model_names)
|
| self.design_hard(hard_iters, temp=1e-2, **kwargs)
|
|
|
| def _mutate(self, seq, plddt=None, logits=None, mutation_rate=1):
|
| '''mutate random position'''
|
| seq = np.array(seq)
|
| N,L = seq.shape
|
|
|
|
|
| i_prob = np.ones(L) if plddt is None else np.maximum(1-plddt,0)
|
| i_prob[np.isnan(i_prob)] = 0
|
| if "fix_pos" in self.opt:
|
| if "pos" in self.opt:
|
| p = self.opt["pos"][self.opt["fix_pos"]]
|
| seq[...,p] = self._wt_aatype_sub
|
| else:
|
| p = self.opt["fix_pos"]
|
| seq[...,p] = self._wt_aatype[...,p]
|
| i_prob[p] = 0
|
|
|
| for m in range(mutation_rate):
|
|
|
|
|
| i = np.random.choice(np.arange(L),p=i_prob/i_prob.sum())
|
|
|
|
|
| logits = np.array(0 if logits is None else logits)
|
| if logits.ndim == 3: logits = logits[:,i]
|
| elif logits.ndim == 2: logits = logits[i]
|
| a_logits = logits - np.eye(self._args["alphabet_size"])[seq[:,i]] * 1e8
|
| a = categorical(softmax(a_logits))
|
|
|
|
|
| seq[:,i] = a
|
|
|
| return seq
|
|
|
| def design_semigreedy(self, iters=100, tries=10, dropout=False,
|
| save_best=True, seq_logits=None, e_tries=None, **kwargs):
|
|
|
| '''semigreedy search'''
|
| if e_tries is None: e_tries = tries
|
|
|
|
|
| if hasattr(self,"aux"):
|
| seq = self.aux["seq"]["logits"].argmax(-1)
|
| else:
|
| seq = (self._params["seq"] + self._inputs["bias"]).argmax(-1)
|
|
|
|
|
| if seq_logits is None: seq_logits = 0
|
|
|
| model_flags = {k:kwargs.pop(k,None) for k in ["num_models","sample_models","models"]}
|
| verbose = kwargs.pop("verbose",1)
|
|
|
|
|
| aux = self.predict(seq, return_aux=True, verbose=False, **model_flags, **kwargs)
|
| plddt = self.aux["plddt"]
|
| plddt = plddt[self._target_len:] if self.protocol == "binder" else plddt[:self._len]
|
|
|
|
|
| if verbose:
|
| print("Running semigreedy optimization...")
|
|
|
| for i in range(iters):
|
| buff = []
|
| model_nums = self._get_model_nums(**model_flags)
|
| num_tries = (tries+(e_tries-tries)*((i+1)/iters))
|
| for t in range(int(num_tries)):
|
| mut_seq = self._mutate(seq=seq, plddt=plddt,
|
| logits=seq_logits + self._inputs["bias"])
|
| aux = self.predict(seq=mut_seq, return_aux=True, model_nums=model_nums, verbose=False, **kwargs)
|
| buff.append({"aux":aux, "seq":np.array(mut_seq)})
|
|
|
|
|
| losses = [x["aux"]["loss"] for x in buff]
|
| best = buff[np.argmin(losses)]
|
| self.aux, seq = best["aux"], jnp.array(best["seq"])
|
| self.set_seq(seq=seq, bias=self._inputs["bias"])
|
| self._save_results(save_best=save_best, verbose=verbose)
|
|
|
|
|
| plddt = best["aux"]["plddt"]
|
| plddt = plddt[self._target_len:] if self.protocol == "binder" else plddt[:self._len]
|
| self._k += 1
|
|
|
| def design_pssm_semigreedy(self, soft_iters=300, hard_iters=32, tries=10, e_tries=None,
|
| ramp_recycles=True, ramp_models=True, **kwargs):
|
|
|
| verbose = kwargs.get("verbose",1)
|
|
|
|
|
| if soft_iters > 0:
|
| self.design_3stage(soft_iters, 0, 0, ramp_recycles=ramp_recycles, **kwargs)
|
| self._tmp["seq_logits"] = kwargs["seq_logits"] = self.aux["seq"]["logits"]
|
|
|
|
|
| if hard_iters > 0:
|
| kwargs["dropout"] = False
|
| if ramp_models:
|
| num_models = len(kwargs.get("models",self._model_names))
|
| iters = hard_iters
|
| for m in range(num_models):
|
| if verbose and m > 0: print(f'Increasing number of models to {m+1}.')
|
|
|
| kwargs["num_models"] = m + 1
|
| kwargs["save_best"] = (m + 1) == num_models
|
| self.design_semigreedy(iters, tries=tries, e_tries=e_tries, **kwargs)
|
| if m < 2: iters = iters // 2
|
| else:
|
| self.design_semigreedy(hard_iters, tries=tries, e_tries=e_tries, **kwargs)
|
|
|
|
|
|
|
|
|
|
|
| def _design_mcmc(self, steps=1000, half_life=200, T_init=0.01, mutation_rate=1,
|
| seq_logits=None, save_best=True, **kwargs):
|
| '''
|
| MCMC with simulated annealing
|
| ----------------------------------------
|
| steps = number for steps for the MCMC trajectory
|
| half_life = half-life for the temperature decay during simulated annealing
|
| T_init = starting temperature for simulated annealing. Temperature is decayed exponentially
|
| mutation_rate = number of mutations at each MCMC step
|
| '''
|
|
|
|
|
|
|
|
|
| verbose = kwargs.pop("verbose",1)
|
| model_flags = {k:kwargs.pop(k,None) for k in ["num_models","sample_models","models"]}
|
|
|
|
|
| plddt, best_loss, current_loss = None, np.inf, np.inf
|
| current_seq = (self._params["seq"] + self._inputs["bias"]).argmax(-1)
|
| if seq_logits is None: seq_logits = 0
|
|
|
|
|
| if verbose: print("Running MCMC with simulated annealing...")
|
| for i in range(steps):
|
|
|
|
|
| T = T_init * (np.exp(np.log(0.5) / half_life) ** i)
|
|
|
|
|
| if i == 0:
|
| mut_seq = current_seq
|
| else:
|
| mut_seq = self._mutate(seq=current_seq, plddt=plddt,
|
| logits=seq_logits + self._inputs["bias"],
|
| mutation_rate=mutation_rate)
|
|
|
|
|
| model_nums = self._get_model_nums(**model_flags)
|
| aux = self.predict(seq=mut_seq, return_aux=True, verbose=False, model_nums=model_nums, **kwargs)
|
| loss = aux["log"]["loss"]
|
|
|
|
|
| delta = loss - current_loss
|
| if i == 0 or delta < 0 or np.random.uniform() < np.exp( -delta / T):
|
|
|
|
|
| (current_seq,current_loss) = (mut_seq,loss)
|
|
|
| plddt = aux["all"]["plddt"].mean(0)
|
| plddt = plddt[self._target_len:] if self.protocol == "binder" else plddt[:self._len]
|
|
|
| if loss < best_loss:
|
| (best_loss, self._k) = (loss, i)
|
| self.set_seq(seq=current_seq, bias=self._inputs["bias"])
|
| self._save_results(save_best=save_best, verbose=verbose)
|
|
|
