| import jax
|
| import jax.numpy as jnp
|
| import numpy as np
|
| import optax
|
|
|
| from colabdesign.shared.utils import copy_dict, update_dict, softmax, Key
|
| from colabdesign.shared.prep import rewire
|
| from colabdesign.af.alphafold.common import residue_constants
|
|
|
| aa_order = residue_constants.restype_order
|
| order_aa = {b:a for a,b in aa_order.items()}
|
|
|
| class design_model:
|
| def set_weights(self, *args, **kwargs):
|
| '''
|
| set weights
|
| -------------------
|
| note: model.restart() resets the weights to their defaults
|
| use model.set_weights(..., set_defaults=True) to avoid this
|
| -------------------
|
| model.set_weights(rmsd=1)
|
| '''
|
| if kwargs.pop("set_defaults", False):
|
| update_dict(self._opt["weights"], *args, **kwargs)
|
| update_dict(self.opt["weights"], *args, **kwargs)
|
|
|
| def set_seq(self, seq=None, mode=None, bias=None, rm_aa=None, set_state=True, **kwargs):
|
| '''
|
| set sequence params and bias
|
| -----------------------------------
|
| -seq=str or seq=[str,str] or seq=array(shape=(L,20) or shape=(?,L,20))
|
| -mode=
|
| -"wildtype"/"wt" = initialize sequence with sequence saved from input PDB
|
| -"gumbel" = initial sequence with gumbel distribution
|
| -"soft_???" = apply softmax-activation to initailized sequence (eg. "soft_gumbel")
|
| -bias=array(shape=(20,) or shape=(L,20)) - bias the sequence
|
| -rm_aa="C,W" = specify which amino acids to remove (aka. add a negative-infinity bias to these aa)
|
| -----------------------------------
|
| '''
|
|
|
|
|
| seq_init = kwargs.pop("seq_init",None)
|
| if seq_init is not None:
|
| modes = ["soft","gumbel","wildtype","wt"]
|
| if isinstance(seq_init,str):
|
| seq_init = seq_init.split("_")
|
| if isinstance(seq_init,list) and seq_init[0] in modes:
|
| mode = seq_init
|
| else:
|
| seq = seq_init
|
|
|
| if mode is None: mode = []
|
|
|
|
|
| shape = (self._num, self._len, self._args.get("alphabet_size",20))
|
|
|
|
|
| if bias is None:
|
| b = np.zeros(shape[1:])
|
| else:
|
| b = np.array(np.broadcast_to(bias, shape[1:]))
|
|
|
|
|
| if rm_aa is not None:
|
| for aa in rm_aa.split(","):
|
| b[...,aa_order[aa]] -= 1e6
|
|
|
|
|
| if ("wildtype" in mode or "wt" in mode) and hasattr(self,"_wt_aatype"):
|
| wt_seq = np.eye(shape[-1])[self._wt_aatype]
|
| wt_seq[self._wt_aatype == -1] = 0
|
| if "pos" in self.opt and self.opt["pos"].shape[0] == wt_seq.shape[0]:
|
| seq = np.zeros(shape)
|
| seq[:,self.opt["pos"],:] = wt_seq
|
| else:
|
| seq = wt_seq
|
|
|
|
|
| if seq is None:
|
| if hasattr(self,"key"):
|
| x = 0.01 * np.random.normal(size=shape)
|
| else:
|
| x = np.zeros(shape)
|
| else:
|
| if isinstance(seq, str):
|
| seq = [seq]
|
| if isinstance(seq, list):
|
| if isinstance(seq[0], str):
|
| aa_dict = copy_dict(aa_order)
|
| if shape[-1] > 21:
|
| aa_dict["-"] = 21
|
| seq = np.asarray([[aa_dict.get(aa,-1) for aa in s] for s in seq])
|
| else:
|
| seq = np.asarray(seq)
|
| else:
|
| seq = np.asarray(seq)
|
|
|
| if np.issubdtype(seq.dtype, np.integer):
|
| seq_ = np.eye(shape[-1])[seq]
|
| seq_[seq == -1] = 0
|
| seq = seq_
|
|
|
| if kwargs.pop("add_seq",False):
|
| b = b + seq * 1e7
|
|
|
| if seq.ndim == 2:
|
| x = np.pad(seq[None],[[0,shape[0]-1],[0,0],[0,0]])
|
| elif shape[0] > seq.shape[0]:
|
| x = np.pad(seq,[[0,shape[0]-seq.shape[0]],[0,0],[0,0]])
|
| else:
|
| x = seq
|
|
|
| if "gumbel" in mode:
|
| y_gumbel = jax.random.gumbel(self.key(),shape)
|
| if "soft" in mode:
|
| y = softmax(x + b + y_gumbel)
|
| elif "alpha" in self.opt:
|
| y = x + y_gumbel / self.opt["alpha"]
|
| else:
|
| y = x + y_gumbel
|
|
|
| x = np.where(x.sum(-1,keepdims=True) == 1, x, y)
|
|
|
|
|
| self._params["seq"] = x
|
| self._inputs["bias"] = b
|
|
|
| def _norm_seq_grad(self):
|
| g = self.aux["grad"]["seq"]
|
| eff_L = (np.square(g).sum(-1,keepdims=True) > 0).sum(-2,keepdims=True)
|
| gn = np.linalg.norm(g,axis=(-1,-2),keepdims=True)
|
| self.aux["grad"]["seq"] = g * np.sqrt(eff_L) / (gn + 1e-7)
|
|
|
| def set_optimizer(self, optimizer=None, learning_rate=None, norm_seq_grad=None, **kwargs):
|
| '''
|
| set/reset optimizer
|
| ----------------------------------
|
| supported optimizers include: [adabelief, adafactor, adagrad, adam, adamw,
|
| fromage, lamb, lars, noisy_sgd, dpsgd, radam, rmsprop, sgd, sm3, yogi]
|
| '''
|
| optimizers = {'adabelief':optax.adabelief,'adafactor':optax.adafactor,
|
| 'adagrad':optax.adagrad,'adam':optax.adam,
|
| 'adamw':optax.adamw,'fromage':optax.fromage,
|
| 'lamb':optax.lamb,'lars':optax.lars,
|
| 'noisy_sgd':optax.noisy_sgd,'dpsgd':optax.dpsgd,
|
| 'radam':optax.radam,'rmsprop':optax.rmsprop,
|
| 'sgd':optax.sgd,'sm3':optax.sm3,'yogi':optax.yogi}
|
|
|
| if optimizer is None: optimizer = self._args["optimizer"]
|
| if learning_rate is not None: self.opt["learning_rate"] = learning_rate
|
| if norm_seq_grad is not None: self.opt["norm_seq_grad"] = norm_seq_grad
|
|
|
| o = optimizers[optimizer](1.0, **kwargs)
|
| self._state = o.init(self._params)
|
|
|
| def update_grad(state, grad, params):
|
| updates, state = o.update(grad, state, params)
|
| grad = jax.tree_util.tree_map(lambda x:-x, updates)
|
| return state, grad
|
|
|
| self._optimizer = jax.jit(update_grad)
|
|
|
| def set_seed(self, seed=None):
|
| np.random.seed(seed=seed)
|
| self.key = Key(seed=seed).get
|
|
|
| def get_seq(self, get_best=True):
|
| '''
|
| get sequences as strings
|
| - set get_best=False, to get the last sampled sequence
|
| '''
|
| aux = self._tmp["best"]["aux"] if (get_best and "aux" in self._tmp["best"]) else self.aux
|
| x = aux["seq"]["hard"].argmax(-1)
|
| return ["".join([order_aa[a] for a in s]) for s in x]
|
|
|
| def get_seqs(self, get_best=True):
|
| return self.get_seq(get_best)
|
|
|
| def rewire(self, order=None, offset=0, loops=0):
|
| '''
|
| helper function for "partial" protocol
|
| -----------------------------------------
|
| -order=[0,1,2] - change order of specified segments
|
| -offset=0 - specify start position of the first segment
|
| -loops=[3,2] - specified loop lengths between segments
|
| -----------------------------------------
|
| '''
|
| self.opt["pos"] = rewire(length=self._pos_info["length"], order=order,
|
| offset=offset, loops=loops)
|
|
|
|
|
| if hasattr(self,"_opt"): self._opt["pos"] = self.opt["pos"]
|
|
|
| def soft_seq(x, bias, opt, key=None, num_seq=None, shuffle_first=True):
|
| seq = {"input":x}
|
|
|
| if x.ndim == 3 and x.shape[0] > 1 and key is not None:
|
| key, sub_key = jax.random.split(key)
|
| if num_seq is None or x.shape[0] == num_seq:
|
|
|
| if shuffle_first:
|
| n = jax.random.randint(sub_key,[],0,x.shape[0])
|
| seq["input"] = seq["input"].at[0].set(seq["input"][n]).at[n].set(seq["input"][0])
|
| else:
|
| n = jnp.arange(x.shape[0])
|
| if shuffle_first:
|
| n = jax.random.permutation(sub_key,n)
|
| else:
|
| n = jnp.append(0,jax.random.permutation(sub_key,n[1:]))
|
| seq["input"] = seq["input"][n[:num_seq]]
|
|
|
|
|
| seq["logits"] = seq["input"] * opt["alpha"]
|
| if bias is not None: seq["logits"] = seq["logits"] + bias
|
| seq["pssm"] = jax.nn.softmax(seq["logits"])
|
| seq["soft"] = jax.nn.softmax(seq["logits"] / opt["temp"])
|
| seq["hard"] = jax.nn.one_hot(seq["soft"].argmax(-1), seq["soft"].shape[-1])
|
| seq["hard"] = jax.lax.stop_gradient(seq["hard"] - seq["soft"]) + seq["soft"]
|
|
|
|
|
| seq["pseudo"] = opt["soft"] * seq["soft"] + (1-opt["soft"]) * seq["input"]
|
| seq["pseudo"] = opt["hard"] * seq["hard"] + (1-opt["hard"]) * seq["pseudo"]
|
| return seq
|
|
|
