public release 4.2.0 - see README.md and CHANGES.md for details

pmsco/graphics/cluster.py

@@ -0,0 +1,178 @@
+#!/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)