Record a movie of a ChimeraX vseries with color interpolation

animate_vser.py

#!/usr/bin/env python
from coloraide import Color
import pyperclip
import argparse
from pathlib import Path

def main(args):
    vol_base, num_vols = args.num_vols.split(".")
    vol_base = "#" + vol_base
    num_vols = int(num_vols)

    colors = ["#" + x for x in args.color_seq.split(",")]
    color_interpolator = Color.interpolate(colors, space = "oklab")
    stepsize = 1 / (num_vols - 1)
    color_ramp = [color_interpolator(stepsize * i).convert("srgb").to_string(hex = True) for i in range(num_vols)]
    if args.reverse_colors:
        color_ramp = color_ramp[::-1]

    filename_iter = 0
    base_dir = Path(args.movie_dir).expanduser()
    movie_name = base_dir / f"vser-movie_{filename_iter:0>2}.mp4"
    while movie_name.exists():
        filename_iter += 1
        movie_name = base_dir / f"vser-movie_{filename_iter:0>2}.mp4"
        filename_iter += 1

    chimerax_commands = [
        "~show",
        f"show {vol_base}",
        "; ".join(f"color {vol_base}.{i + 1} {color_ramp[i]}" for i in range(num_vols)),
        "movie record super 3",
        f"vol morph {vol_base} slider false frames {2 * num_vols} playStep {1 / num_vols}",
        f"wait {2 * num_vols - 1}",
        f"movie encode {movie_name} framerate {args.framerate}"
    ]
    command = "; ".join(chimerax_commands)
    pyperclip.copy(command)
    print("Command copied to clipboard")

parser = argparse.ArgumentParser()
parser.add_argument(
    "num_vols",
    help = "Number of volumes over which to animate. Give as ChimeraX string (e.g., 1.30 for 30 subvolumes of #1)"
)
parser.add_argument(
    "color_seq",
    help = "Comma-separated hex colors over which to interpolate. Leave off the #"
)
parser.add_argument(
    "--reverse-colors",
    help = "Reverse the order of color interpolation",
    action = "store_true"
)
parser.add_argument(
    "--movie-dir",
    help = "Directory in which to save movie",
    default = "~/Desktop/"
)
parser.add_argument(
    "--framerate",
    help = "Framerate of the final movie",
    type = int,
    default = 30
)

if __name__ == "__main__":
    args = parser.parse_args()
    main(args)