File size: 8,228 Bytes
d766458 | 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 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160 161 162 163 164 165 166 167 168 169 170 171 172 173 174 175 176 177 178 179 180 181 182 183 184 185 186 187 188 189 190 191 192 193 194 195 196 197 198 199 200 201 202 203 204 205 206 207 208 209 210 211 212 213 214 215 216 217 218 219 220 221 222 223 224 | 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)
-----------------------------------
'''
# backward compatibility
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 = []
# decide on shape
shape = (self._num, self._len, self._args.get("alphabet_size",20))
# initialize bias
if bias is None:
b = np.zeros(shape[1:])
else:
b = np.array(np.broadcast_to(bias, shape[1:]))
# disable certain amino acids
if rm_aa is not None:
for aa in rm_aa.split(","):
b[...,aa_order[aa]] -= 1e6
# use wildtype sequence
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
# initialize sequence
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 # add gap character
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)
# set seq/bias/state
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)
# make default
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}
# shuffle msa
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:
# randomly pick which sequence is query
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]]
# straight-through/reparameterization
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"]
# create pseudo sequence
seq["pseudo"] = opt["soft"] * seq["soft"] + (1-opt["soft"]) * seq["input"]
seq["pseudo"] = opt["hard"] * seq["hard"] + (1-opt["hard"]) * seq["pseudo"]
return seq
|