Pymol functions for creating low depth cartoony images and generating color palettes programmatically.

pymol_utils.py

"""
To use these, open PyMOL separately and type:

    run <path/to/this_file.py>

Alternatively, place the above command in your ~/.pymolrc file.

"""

import math

from pymol import cmd


def interpolate_colors(hex_colors: list[str], n: int):
    """
    Interpolates colors between given hexcodes and returns n colors.

    :param hex_colors: List of 5 hexcodes (strings).
    :param n: The number of colors to return.
    :return: List of n interpolated hexcodes.
    """
    import matplotlib.colors as mcolors
    import numpy as np

    # Convert hex colors to RGB
    rgb_colors = [mcolors.hex2color(color) for color in hex_colors]

    # Prepare the interpolation
    x = np.linspace(0, 1, len(rgb_colors))
    xp = np.linspace(0, 1, n)

    # Interpolate each channel
    reds = np.interp(xp, x, [color[0] for color in rgb_colors])
    greens = np.interp(xp, x, [color[1] for color in rgb_colors])
    blues = np.interp(xp, x, [color[2] for color in rgb_colors])

    # Combine interpolated channels and convert back to hex
    interpolated_colors = [
        mcolors.to_hex((r, g, b)) for r, g, b in zip(reds, greens, blues)
    ]

    return interpolated_colors


@cmd.extend
def color_hcl_rainbow(
    selection: str = "all",
    hue: tuple[int, int] = (0, 270),
    chroma: int = 300,
    luminence: int = 90,
    rev=True,
):
    """
    Creates a pastel rainbow color palette. Default settings are usually good for HCL values.

    NOTE:
        To use this, you must install the `colorspace` library inside your *pymol* environment.

        If you installed PyMOL via homebrew, the only way to do this is via the following scary commands:

            /opt/homebrew/bin/python -m pip install colorspace --user --break-system-packages
            brew install python-tk  # installs tkinter

        This should not in fact break things on your system but it is messy.
    """
    from colorspace import rainbow_hcl

    chroma = int(chroma)
    luminence = int(luminence)
    if isinstance(hue, str):
        hue = eval(hue)  # type: ignore
    assert len(hue) == 2
    color_palette = rainbow_hcl(
        c=chroma, l=luminence, start=hue[0], end=hue[1], rev=rev
    )
    objects = cmd.get_object_list(selection)
    for obj in objects:
        residues = set()
        cmd.iterate(
            f"{obj} and {selection} and n. CA",
            "residues.add((chain, resv))",
            space={"residues": residues},
        )
        num_colors = len(residues)
        for color, (chain, resv) in zip(
            color_palette.colors(num_colors),  # pyright: ignore[reportArgumentType]
            sorted(residues, key=lambda x: (x[0], int(x[1]))),
        ):
            cmd.color(
                color.replace("#", "0x"),
                f"{obj} and {selection} and chain {chain} and i. {resv}",
            )


@cmd.extend
def color_seaborn(
    name: str,
    selection: str = "all",
    desat: float | None = None,
    start: float = 0,
    stop: float = 1,
):
    """
    Creates a seaborn color palette. Any valid sns.color_palette can be passed in.

    NOTE:
        To use this, you must install the `seaborn` library inside your *pymol* environment.

        If you installed PyMOL via homebrew, the only way to do this is via the following scary commands:

            /opt/homebrew/bin/python -m pip install seaborn --user --break-system-packages
            brew install python-tk  # installs tkinter

        This should not in fact break things on your system but it is messy.
    """
    import seaborn as sns

    assert 0 <= start < stop <= 1

    objects = cmd.get_object_list(selection)
    for obj in objects:
        residues = set()
        cmd.iterate(
            f"{obj} and {selection} and n. CA",
            "residues.add((chain, resv))",
            space={"residues": residues},
        )
        num_colors = len(residues)
        fraction = stop - start
        palette = sns.color_palette(
            name, desat=desat, n_colors=math.ceil(num_colors / fraction)
        ).as_hex()
        start_idx = math.floor(len(palette) * start)
        stop_idx = start_idx + num_colors
        palette = palette[start_idx:stop_idx]

        for color, (chain, resv) in zip(
            palette, sorted(residues, key=lambda x: (x[0], int(x[1])))
        ):
            cmd.color(
                color.replace("#", "0x"),
                f"{obj} and {selection} and chain {chain} and i. {resv}",
            )


@cmd.extend
def color_hex(hex_colors: str | list[str], selection: str = "all"):
    """
    Creates a color palette from a list of hex codes. These can be passed in as either a
        python list (["#ffffff", "#000000"]), or as a comma separated string ("#ffffff,#000000").

    NOTE:
        To use this, you must install the `numpy` library inside your *pymol* environment.

        If you installed PyMOL via homebrew, the only way to do this is via the following scary commands:

            /opt/homebrew/bin/python -m pip install numpy --user --break-system-packages

        This should not in fact break things on your system but it is messy.
    """
    import re

    if isinstance(hex_colors, str):
        hex_colors = re.split(",| ", hex_colors)
    objects = cmd.get_object_list(selection)
    for obj in objects:
        residues = set()
        cmd.iterate(
            f"{obj} and {selection} and n. CA",
            "residues.add((chain, resv))",
            space={"residues": residues},
        )
        num_colors = len(residues)
        palette = interpolate_colors(hex_colors, num_colors)

        for color, (chain, resv) in zip(
            palette, sorted(residues, key=lambda x: (x[0], int(x[1])))
        ):
            cmd.color(
                color.replace("#", "0x"),
                f"{obj} and {selection} and chain {chain} and i. {resv}",
            )


@cmd.extend
def low_depth_image():
    cmd.set("transparency", 0.8)  # pyright: ignore[reportArgumentType]
    cmd.set("ray_shadow", 0)
    cmd.set("ray_trace_color", "black")  # pyright: ignore[reportArgumentType]
    cmd.set("fog", 0)
    cmd.set("depth_cue", 0)
    cmd.set("antialias", 4)
    cmd.set("specular", 0)
    cmd.set("ambient", 1)
    cmd.set("orthoscopic", 1)
    cmd.set("light_count", 1)
    cmd.set("ray_trace_mode", 1)
    cmd.set("cartoon_discrete_colors", 1)
    cmd.set("ray_scatter", 0)
    cmd.set("surface_quality", 0)
    cmd.set("ambient_occlusion_mode", 0)