Last active
May 31, 2022 16:20
-
-
Save rmrao/4c37ebfcbc88b02caf7c01f1c84a0596 to your computer and use it in GitHub Desktop.
My pymol rc script
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
from pymol import preset, cmd | |
from glob import glob | |
from os.path import sep, basename | |
from pathlib import Path | |
def bns(sel: str = "bb."): | |
preset.ball_and_stick(selection=sel, mode=1) | |
def plddt(sel: str = "all"): | |
seleobjs = cmd.get_object_list(sel) | |
for obj in seleobjs: | |
if obj.startswith("ranked") or "alphafold" in obj or "af2" in obj: | |
minimum: float = 50 | |
maximum: float = 90 | |
else: | |
minimum = 0.5 | |
maximum = 0.9 | |
cmd.spectrum("b", "red_green", obj, minimum, maximum) | |
def loadDir(dirName=".", suff="pdb", group=None): | |
""" | |
Loads all files with the suffix suff (the input parameter) from the directory dirName). | |
dirName: directory path | |
suff: file suffix. Should be simply "pdb" or "sdf" or similar. Will accept the | |
wildcard and dot in case the user doesn't read this. So, "*.pdb", ".pdb", | |
and "pdb" should work. The suffix can be anything valid that PyMOL knows | |
how to natively load. | |
group: groupName to add the files to. | |
example: | |
# load all the PDBs in the current directory | |
loadDir | |
# load all SD files from /tmp | |
loadDir /tmp, "sdf" | |
notes: | |
make sure you call this script w/o quotes around your parameters: | |
loadDir ., .pdb | |
as opposed to | |
loadDir ".", "*.pdb" | |
Use the former. | |
""" | |
g = dirName + sep + "*." + suff.split(".")[-1] | |
for c in glob( g ): | |
cmd.load(c) | |
if ( group != None ): | |
cmd.group( group, basename(c).split(".")[0], "add" ) | |
def alignto(sel: str, ref: str, ce: bool = False): | |
seleobjs = cmd.get_object_list(sel) | |
best_obj, best_rmsd = None, float("inf") | |
for obj in seleobjs: | |
if ce: | |
res = cmd.cealign(ref, obj) | |
rmsd = res["RMSD"] | |
else: | |
rmsd, *_ = cmd.align(obj, ref) | |
if rmsd < best_rmsd: | |
best_rmsd = rmsd | |
best_obj = obj | |
print(f"Best: {best_obj}, RMSD: {best_rmsd}") | |
def load_predictions(directory: str): | |
import pandas as pd | |
path = Path(directory).absolute() | |
target = path.parent.name | |
method = path.name | |
plddt = path / f"{target}-plddtmask.csv" | |
df = pd.read_csv(plddt) | |
max_plddt = df["plddt"].idxmax() | |
min_plddt = df["plddt"].idxmin() | |
base_path = path / f"{target}-base.pdb" | |
max_path = path / f"{target}-{max_plddt}.pdb" | |
min_path = path / f"{target}-{min_plddt}.pdb" | |
af2_path = path.parent.parent.parent / target / "ranked_0.pdb" | |
if base_path.exists(): | |
cmd.load(base_path) | |
if max_path.exists(): | |
cmd.load(max_path, f"{target}-max") | |
if min_path.exists(): | |
cmd.load(min_path, f"{target}-min") | |
if af2_path.exists(): | |
cmd.load(af2_path, f"af2-{target}") | |
# def load_af2(name) | |
cmd.extend("bns", bns) | |
cmd.extend("plddt", plddt) | |
cmd.extend("loadDir", loadDir) | |
cmd.extend("alignto", alignto) | |
cmd.extend("load_predictions", load_predictions) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
run ~/.pymol_functions.py |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment