| import jax
|
| import jax.numpy as jnp
|
| import numpy as np
|
| import matplotlib.pyplot as plt
|
| from matplotlib.gridspec import GridSpec
|
|
|
| from colabdesign.shared.protein import _np_kabsch
|
| from colabdesign.shared.utils import update_dict, Key
|
| from colabdesign.shared.plot import plot_pseudo_3D, make_animation, show_pdb
|
| from colabdesign.shared.protein import renum_pdb_str
|
| from colabdesign.af.alphafold.common import protein
|
|
|
|
|
|
|
|
|
| class _af_utils:
|
|
|
| def set_opt(self, *args, **kwargs):
|
| '''
|
| set [opt]ions
|
| -------------------
|
| note: model.restart() resets the [opt]ions to their defaults
|
| use model.set_opt(..., set_defaults=True)
|
| or model.restart(..., reset_opt=False) to avoid this
|
| -------------------
|
| model.set_opt(num_models=1, num_recycles=0)
|
| model.set_opt(con=dict(num=1)) or set_opt({"con":{"num":1}}) or set_opt("con",num=1)
|
| model.set_opt(lr=1, set_defaults=True)
|
| '''
|
| ks = list(kwargs.keys())
|
| self.set_args(**{k:kwargs.pop(k) for k in ks if k in self._args})
|
|
|
| if kwargs.pop("set_defaults", False):
|
| update_dict(self._opt, *args, **kwargs)
|
|
|
| update_dict(self.opt, *args, **kwargs)
|
|
|
| def set_args(self, **kwargs):
|
| '''
|
| set [arg]uments
|
| '''
|
| for k in ["best_metric", "traj_iter", "shuffle_first"]:
|
| if k in kwargs: self._args[k] = kwargs.pop(k)
|
|
|
| if "recycle_mode" in kwargs:
|
| ok_recycle_mode_swap = ["average","sample","first","last"]
|
| if kwargs["recycle_mode"] in ok_recycle_mode_swap and self._args["recycle_mode"] in ok_recycle_mode_swap:
|
| self._args["recycle_mode"] = kwargs.pop("recycle_mode")
|
| else:
|
| print(f"ERROR: use {self.__class__.__name__}(recycle_mode=...) to set the recycle_mode")
|
|
|
| if "optimizer" in kwargs:
|
| self.set_optimizer(kwargs.pop("optimizer"),
|
| learning_rate=kwargs.pop("learning_rate", None))
|
|
|
| ks = list(kwargs.keys())
|
| if len(ks) > 0:
|
| print(f"ERROR: the following args were not set: {ks}")
|
|
|
| def get_loss(self, x="loss"):
|
| '''output the loss (for entire trajectory)'''
|
| return np.array([loss[x] for loss in self._tmp["log"]])
|
|
|
| def save_pdb(self, filename=None, get_best=True, renum_pdb=True, aux=None):
|
| '''
|
| save pdb coordinates (if filename provided, otherwise return as string)
|
| - set get_best=False, to get the last sampled sequence
|
| '''
|
| if aux is None:
|
| aux = self._tmp["best"]["aux"] if (get_best and "aux" in self._tmp["best"]) else self.aux
|
| aux = aux["all"]
|
|
|
| p = {k:aux[k] for k in ["aatype","residue_index","atom_positions","atom_mask"]}
|
| p["b_factors"] = 100 * p["atom_mask"] * aux["plddt"][...,None]
|
|
|
| def to_pdb_str(x, n=None):
|
| p_str = protein.to_pdb(protein.Protein(**x))
|
| p_str = "\n".join(p_str.splitlines()[1:-2])
|
| if renum_pdb: p_str = renum_pdb_str(p_str, self._lengths)
|
| if n is not None:
|
| p_str = f"MODEL{n:8}\n{p_str}\nENDMDL\n"
|
| return p_str
|
|
|
| p_str = ""
|
| for n in range(p["atom_positions"].shape[0]):
|
| p_str += to_pdb_str(jax.tree_util.tree_map(lambda x:x[n],p), n+1)
|
| p_str += "END\n"
|
|
|
| if filename is None:
|
| return p_str
|
| else:
|
| with open(filename, 'w') as f:
|
| f.write(p_str)
|
|
|
|
|
|
|
|
|
| def animate(self, s=0, e=None, dpi=100, get_best=True, traj=None, aux=None, color_by="plddt"):
|
| '''
|
| animate the trajectory
|
| - use [s]tart and [e]nd to define range to be animated
|
| - use dpi to specify the resolution of animation
|
| - color_by = ["plddt","chain","rainbow"]
|
| '''
|
| if aux is None:
|
| aux = self._tmp["best"]["aux"] if (get_best and "aux" in self._tmp["best"]) else self.aux
|
| aux = aux["all"]
|
| if self.protocol in ["fixbb","binder"]:
|
| pos_ref = self._inputs["batch"]["all_atom_positions"][:,1].copy()
|
| pos_ref[(pos_ref == 0).any(-1)] = np.nan
|
| else:
|
| pos_ref = aux["atom_positions"][0,:,1,:]
|
|
|
| if traj is None: traj = self._tmp["traj"]
|
| sub_traj = {k:v[s:e] for k,v in traj.items()}
|
|
|
| align_xyz = self.protocol == "hallucination"
|
| return make_animation(**sub_traj, pos_ref=pos_ref, length=self._lengths,
|
| color_by=color_by, align_xyz=align_xyz, dpi=dpi)
|
|
|
| def plot_pdb(self, show_sidechains=False, show_mainchains=False,
|
| color="pLDDT", color_HP=False, size=(800,480), animate=False,
|
| get_best=True, aux=None, pdb_str=None):
|
| '''
|
| use py3Dmol to plot pdb coordinates
|
| - color=["pLDDT","chain","rainbow"]
|
| '''
|
| if pdb_str is None:
|
| pdb_str = self.save_pdb(get_best=get_best, aux=aux)
|
| view = show_pdb(pdb_str,
|
| show_sidechains=show_sidechains,
|
| show_mainchains=show_mainchains,
|
| color=color,
|
| Ls=self._lengths,
|
| color_HP=color_HP,
|
| size=size,
|
| animate=animate)
|
| view.show()
|
|
|
| def plot_traj(self, dpi=100):
|
| fig = plt.figure(figsize=(5,5), dpi=dpi)
|
| gs = GridSpec(4,1, figure=fig)
|
| ax1 = fig.add_subplot(gs[:3,:])
|
| ax2 = fig.add_subplot(gs[3:,:])
|
| ax1_ = ax1.twinx()
|
|
|
| if self.protocol in ["fixbb","partial"] or (self.protocol == "binder" and self._args["redesign"]):
|
| if self.protocol == "partial" and self._args["use_sidechains"]:
|
| rmsd = self.get_loss("sc_rmsd")
|
| else:
|
| rmsd = self.get_loss("rmsd")
|
| for k in [0.5,1,2,4,8,16,32]:
|
| ax1.plot([0,len(rmsd)],[k,k],color="lightgrey")
|
| ax1.plot(rmsd,color="black")
|
| seqid = self.get_loss("seqid")
|
| ax1_.plot(seqid,color="green",label="seqid")
|
|
|
| ax1.set_yscale("log")
|
| ticks = [0.25,0.5,1,2,4,8,16,32,64]
|
| ax1.set(xticks=[])
|
| ax1.set_yticks(ticks); ax1.set_yticklabels(ticks)
|
| ax1.set_ylabel("RMSD",color="black");ax1_.set_ylabel("seqid",color="green")
|
| ax1.set_ylim(0.25,64)
|
| ax1_.set_ylim(0,0.8)
|
|
|
| ax2.plot(self.get_loss("soft"),color="yellow",label="soft")
|
| ax2.plot(self.get_loss("temp"),color="orange",label="temp")
|
| ax2.plot(self.get_loss("hard"),color="red",label="hard")
|
| ax2.set_ylim(-0.1,1.1)
|
| ax2.set_xlabel("iterations")
|
| ax2.legend(loc='center left')
|
| else:
|
| print("TODO")
|
| plt.show()
|
|
|
| def clear_best(self):
|
| self._tmp["best"] = {}
|
|
|
| def save_current_pdb(self, filename=None):
|
| '''save pdb coordinates (if filename provided, otherwise return as string)'''
|
| self.save_pdb(filename=filename, get_best=False)
|
|
|
| def plot_current_pdb(self, show_sidechains=False, show_mainchains=False,
|
| color="pLDDT", color_HP=False, size=(800,480), animate=False):
|
| '''use py3Dmol to plot pdb coordinates
|
| - color=["pLDDT","chain","rainbow"]
|
| '''
|
| self.plot_pdb(show_sidechains=show_sidechains, show_mainchains=show_mainchains, color=color,
|
| color_HP=color_HP, size=size, animate=animate, get_best=False) |
