pmsco-public/pmsco/graphics/cluster.py

#!/usr/bin/env python

"""
@package pmsco.graphics.cluster
graphics rendering module for clusters.

this module is experimental.
interface and implementation may change without notice.

at the moment we are evaluating rendering solutions.

@author Matthias Muntwiler, matthias.muntwiler@psi.ch

@copyright (c) 2017 by Paul Scherrer Institut @n
Licensed under the Apache License, Version 2.0 (the "License"); @n
  you may not use this file except in compliance with the License.
  You may obtain a copy of the License at
  http://www.apache.org/licenses/LICENSE-2.0
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function
import sys
import os
import numpy as np
import argparse
import logging

logger = logging.getLogger(__name__)

try:
    import pymol2
except ImportError:
    logger.warning("error importing pymol2. cluster rendering using pymol2 disabled.")
    pymol2 = None

try:
    from mpl_toolkits.mplot3d import Axes3D
    from matplotlib.figure import Figure
    from matplotlib.backends.backend_agg import FigureCanvasAgg as FigureCanvas
    # from matplotlib.backends.backend_pdf import FigureCanvasPdf
    # from matplotlib.backends.backend_svg import FigureCanvasSVG
except ImportError:
    Axes3D = None
    Figure = None
    FigureCanvas = None
    logger.warning("error importing matplotlib. cluster rendering using matplotlib disabled.")


def render_file(spath, view):
    sname = "cluster"
    opath = spath + ".png"

    pm = pymol2.PyMOL()
    cmd = pm.cmd
    pm.start()
    try:
        cmd.reinitialize()
        cmd.load(spath, sname)
        cmd.disable("all")
        cmd.enable(sname)

        cmd.set("orthoscopic", 1)
        cmd.bg_color("white")
        cmd.show_as("spheres")
        cmd.alter("all", "vdw=0.8")
        #cmd.show("sticks")

        #zoom selection-expression # selection to fill the viewer
        #orient selection-expression # largest dim horizontal, second-largest vertical

        #cmd.orient() --- should stick to fixed orientation
        #cmd.turn("x", -90)
        #cmd.turn("x", 0)
        #cmd.turn("y", 0)

        #cmd.clip("slab", 5.0)
        cmd.viewport(640, 640)
        cmd.zoom(complete=1)
        #pymol.cmd.rebuild() #--- necessary?

        cmd.png(opath)
    finally:
        pm.stop()


def render_cluster(clu):
    pass


def set_axes_equal(ax):
    """
    Make axes of 3D plot have equal scale so that spheres appear as spheres,
    cubes as cubes, etc..  This is one possible solution to Matplotlib's
    ax.set_aspect('equal') and ax.axis('equal') not working for 3D.

    @author https://stackoverflow.com/a/31364297

    @param ax: a matplotlib axis, e.g., as output from plt.gca().
    """

    x_limits = ax.get_xlim3d()
    y_limits = ax.get_ylim3d()
    z_limits = ax.get_zlim3d()

    x_range = abs(x_limits[1] - x_limits[0])
    x_middle = np.mean(x_limits)
    y_range = abs(y_limits[1] - y_limits[0])
    y_middle = np.mean(y_limits)
    z_range = abs(z_limits[1] - z_limits[0])
    z_middle = np.mean(z_limits)

    # The plot bounding box is a sphere in the sense of the infinity
    # norm, hence I call half the max range the plot radius.
    plot_radius = 0.5*max([x_range, y_range, z_range])

    ax.set_xlim3d([x_middle - plot_radius, x_middle + plot_radius])
    ax.set_ylim3d([y_middle - plot_radius, y_middle + plot_radius])
    ax.set_zlim3d([z_middle - plot_radius, z_middle + plot_radius])


def render_xyz_matplotlib(filename, data, canvas=None):
    """
    produce a graphics file from an array of 3d coordinates in the matplotlib scatter style.

    the default file format is PNG.

    this function requires the matplotlib module.
    if it is not available, the function raises an error.

    @param filename: path and name of the scan file.
        this is used to derive the output file path by adding the extension of the graphics file format.
    @param data: numpy array of shape (N,3).
    @param canvas: a FigureCanvas class reference from a matplotlib backend.
        if None, the default FigureCanvasAgg is used which produces a bitmap file in PNG format.

    @return (str) path and name of the generated graphics file.
        empty string if an error occurred.

    @raise TypeError if matplotlib is not available.
    """

    if canvas is None:
        canvas = FigureCanvas
    fig = Figure()
    canvas(fig)

    ax = fig.add_subplot(111, projection='3d')
    # ax.set_aspect('equal')
    try:
        # method available in matplotlib 2.1 and later
        ax.set_proj_type('ortho')
    except AttributeError:
        pass
    ax.scatter(data[:, 0], data[:, 1], data[:, 2], c='r', marker='o')
    ax.set_xlabel('x')
    ax.set_ylabel('y')
    ax.set_zlabel('z')
    set_axes_equal(ax)

    out_filename = "{0}.{1}".format(filename, canvas.get_default_filetype())
    fig.savefig(out_filename)
    return out_filename


def exec_cli():
    parser = argparse.ArgumentParser()
    parser.add_argument('-v', '--view', default='z')
    parser.add_argument(dest='files', nargs='+')
    args = parser.parse_args()
    for fil in args.files:
        render_file(fil, args.view)


if __name__ == '__main__':
    exec_cli()
    sys.exit(0)