draw.py

#!/usr/bin/env python

import sys
import typing as ty
import copy
import argparse
from pathlib import Path

from ase.io import read
from ase.utils import rotate
import numpy as np

from rdkit import Chem
from xyz2mol import xyz2mol

from cinemol.geometry import CylinderCapType
from cinemol.style import Color
from cinemol.api import Look, Style, Atom, Bond, filter_atoms_and_bonds, Scene, \
    draw_atoms_in_spacefilling_style, draw_bonds_in_tube_style
from cinemol.parsers import parse_sdf
from cinemol.svg import Svg, ViewBox


COLORS = {
    'O': (222, 64, 51),
    'C': (100,) * 3,
    'H': (210,) * 3,
    'N': (54, 147, 202),
}


def draw_molecule(
    atoms: list[Atom],
    bonds: list[Bond],
    style: Style,
    look: Look,
    resolution: int,
    window: ty.Optional[tuple[float, float]] = None,
    view_box: ty.Optional[tuple[float, float, float, float]] = None,
    rotation_over_x_axis: float = 0.0,
    rotation_over_y_axis: float = 0.0,
    rotation_over_z_axis: float = 0.0,
    scale: float = 1.0,
    focal_length: ty.Optional[float] = None,
    exclude_atoms: ty.Optional[list[str]] = None,
) -> Svg:
    atoms, bonds = filter_atoms_and_bonds(atoms, bonds, exclude_atoms)

    # override colors with better defaults
    for atom in atoms:
        atom.color = COLORS[atom.symbol]

    scene = Scene(nodes=[])  # Define empty scene explicitly to prevent memory leak.

    # Default settings for drawing.
    include_spheres = False
    include_cylinders = False
    include_wires = False
    calculate_sphere_sphere_intersections = False
    calculate_sphere_cylinder_intersections = False
    calculate_cylinder_sphere_intersections = False
    calculate_cylinder_cylinder_intersections = False

    # Default settings for cartoon look.
    dark_color = (50,) * 3
    stroke_color = Color(*dark_color)
    stroke_width = 0.03

    include_spheres = True
    include_cylinders = True
    calculate_cylinder_sphere_intersections = True
    cap_type = CylinderCapType.NO_CAP
    draw_atoms_in_spacefilling_style(
        scene,
        atoms,
        look,
        stroke_color,
        stroke_width,
        radius_scale=1.0 / 3.0,  # Lower the radius of the atoms for ball-and-stick style.
    )

    # bonds should be drawn without taking into account the color of the atoms;
    atoms_dark = [copy.deepcopy(atom) for atom in atoms]
    for atom in atoms_dark:
        atom.color = dark_color
    for bond in bonds:
        bond.radius = 0.2
    draw_bonds_in_tube_style(
        scene, atoms_dark, bonds, style, look, cap_type, stroke_color, stroke_width
    )

    svg = scene.draw(
        resolution=resolution,
        window=window,
        view_box=ViewBox(*view_box) if view_box is not None else None,
        rotation_over_x_axis=rotation_over_x_axis,
        rotation_over_y_axis=rotation_over_y_axis,
        rotation_over_z_axis=rotation_over_z_axis,
        include_spheres=include_spheres,
        include_cylinders=include_cylinders,
        include_wires=include_wires,
        calculate_sphere_sphere_intersections=calculate_sphere_sphere_intersections,
        calculate_sphere_cylinder_intersections=calculate_sphere_cylinder_intersections,
        calculate_cylinder_sphere_intersections=calculate_cylinder_sphere_intersections,
        calculate_cylinder_cylinder_intersections=calculate_cylinder_cylinder_intersections,
        filter_nodes_for_intersecting=True,
        scale=scale,
        focal_length=focal_length,
    )

    return svg


if __name__ == '__main__':
    parser = argparse.ArgumentParser()

    parser.add_argument(
        '--xyz',
        type=str,
        required=True,
    )
    parser.add_argument(
        '--index',
        type=str,
        default='-1',
    )
    parser.add_argument(
        '--rotation',
        type=str,
        required=False,
        default='0x,0y,0z',
    )
    parser.add_argument(
        '--resolution',
        type=int,
        default=300,
    )
    parser.add_argument(
        '--scale',
        type=int,
        default=400,
    )
    parser.add_argument(
        '--output',
        type=str,
        default=None,
        required=None,
    )
    args = parser.parse_args()
    atoms = read(args.xyz, index=int(args.index))

    # switch x and y axes to calibrate with convention in cinemol
    positions = atoms.get_positions()
    positions -= np.mean(positions, axis=0, keepdims=True)
    #rotation = rotate('-4.6x,-29y,2.4z')
    rotation = rotate(args.rotation)
    positions = positions @ rotation
    x = np.copy(positions[:, 0])
    y = np.copy(positions[:, 1])
    positions[:, 0] = x
    positions[:, 1] = -y
    atoms.set_positions(positions)

    # generate sdf string
    atomic_numbers = [int(n) for n in atoms.numbers]
    coordinates = np.array(atoms.positions)
    molecules = xyz2mol(atomic_numbers, coordinates)
    sdf_str = Chem.MolToMolBlock(molecules[0])  # only one molecule

    # load into cinemol and draw
    atoms, bonds = parse_sdf(sdf_str, include_hs=True)
    svg = draw_molecule(
        atoms=atoms,
        bonds=bonds,
        style=Style.BALL_AND_STICK,
        look=Look.CARTOON,
        resolution=args.resolution,
        rotation_over_x_axis=0.0,
        rotation_over_y_axis=0.0,
        rotation_over_z_axis=0.0,
        scale=args.scale,
        focal_length=None,
        exclude_atoms=None,
    )
    svg_str = svg.to_svg()

    if args.output is not None:
        path_svg = Path(args.xyz).parent / (args.output + '.svg')
    else:
        path_svg = Path(args.xyz).stem + '.svg'
    with open(path_svg, 'w') as f:
        f.write(svg_str)