File size: 4,140 Bytes
75b53c6 | 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 | """Render side-by-side C-alpha traces of an example input and the generated design (matplotlib, no extra viewer).
```bash
python render.py --inputs-dir inputs --designs-dir outputs --out images
```
"""
import argparse
from pathlib import Path
import matplotlib
matplotlib.use("Agg")
import matplotlib.pyplot as plt
import numpy as np
from biotite.structure import filter_amino_acids
from biotite.structure.io import pdbx
# Atoms that are not part of a C-alpha trace (ligands, motif side chains) are drawn as red dots.
GREY, COLORS = "#9aa0a6", ["#1a73e8", "#e8710a", "#188038", "#d93025", "#9334e6", "#b06000"]
def load(path: Path):
"""Return `(ca_traces, hetero, other, n_atoms)`: per-chain C-alpha coordinates, non-polymer atoms, other resolved
non-C-alpha atoms and the number of resolved atoms."""
atoms = pdbx.get_structure(pdbx.CIFFile.read(path), model=1)
atoms = atoms[np.isfinite(atoms.coord).all(axis=1)]
is_ca = filter_amino_acids(atoms) & (atoms.atom_name == "CA")
traces = {chain: atoms.coord[is_ca & (atoms.chain_id == chain)] for chain in dict.fromkeys(atoms.chain_id[is_ca])}
hetero = atoms.hetero & ~is_ca
return traces, atoms.coord[hetero], atoms.coord[~is_ca & ~hetero], len(atoms)
def draw(ax, traces, dots, title, grey, center):
for index, (chain, coords) in enumerate(traces.items()):
color = GREY if grey else COLORS[index % len(COLORS)]
if len(coords) > 1:
ax.plot(*(coords - center).T, color=color, lw=1.6)
else:
ax.scatter(*(coords - center).T, color=color, s=18)
if len(dots):
ax.scatter(*(dots - center).T, color="#d93025", s=22, alpha=0.95, depthshade=False)
ax.set_title(title, fontsize=10)
ax.set_axis_off()
def main():
parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
parser.add_argument("--inputs-dir", type=Path, required=True)
parser.add_argument("--designs-dir", type=Path, required=True)
parser.add_argument("--out", type=Path, required=True)
parser.add_argument("--names", nargs="*")
args = parser.parse_args()
args.out.mkdir(parents=True, exist_ok=True)
for design in sorted(args.designs_dir.glob("*.cif")):
if args.names and design.stem not in args.names:
continue
source = args.inputs_dir / ("inverse_folding.cif" if design.stem == "partial_diffusion" else f"{design.stem}.cif")
(given, given_hetero, given_other, given_atoms), (generated, generated_hetero, _, _) = load(source), load(design)
# Sparse conditioning (a motif or a ligand) is drawn atom by atom and zoomed; dense conditioning as a trace.
sparse = given_atoms <= 400
panels = [
(given, np.concatenate([given_hetero, given_other]) if sparse else given_hetero, True, sparse),
(generated, generated_hetero, False, False),
]
titles = ["input (conditioning, zoomed)" if sparse else "input (conditioning)", "generated design"]
figure = plt.figure(figsize=(8, 4), dpi=150)
centers, extents = [], []
for traces, dots, _, _ in panels:
points = np.concatenate([*traces.values(), dots]) if len(traces) or len(dots) else np.zeros((1, 3))
centers.append(points.mean(axis=0))
extents.append(max(1.0, np.abs(points - centers[-1]).max() * 1.1))
shared = max(extent for extent, (_, _, _, zoom) in zip(extents, panels) if not zoom)
for position, ((traces, dots, grey, zoom), center) in enumerate(zip(panels, centers)):
ax = figure.add_subplot(1, 2, position + 1, projection="3d")
draw(ax, traces, dots, titles[position], grey, center)
limit = extents[position] if zoom else shared
ax.set_xlim(-limit, limit), ax.set_ylim(-limit, limit), ax.set_zlim(-limit, limit)
figure.suptitle(design.stem.replace("_", " "), fontsize=11)
figure.tight_layout()
figure.savefig(args.out / f"{design.stem}.png")
plt.close(figure)
print("rendered", design.stem)
if __name__ == "__main__":
main()
|