My pymol rc script

.pymol_functions.py

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)

.pymolrc

run ~/.pymol_functions.py