Add the annotated-CIF preparation and export package for the Diffusers RFD4-Proteina pipeline
75b53c6 verified | """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() | |