micelle_plot.py

from mpl_toolkits.mplot3d.art3d import Poly3DCollection
from mpl_toolkits.mplot3d import axes3d, Axes3D
from collections.abc import Iterable
import matplotlib.pyplot as plt
import matplotlib as mpl
import numpy as np
import os, cmd

def link_axes(fig, ax):
    def on_move(event):
        if event.inaxes == ax[0]:
            if ax[0].button_pressed in ax[0]._rotate_btn:
                ax[1].view_init(elev=ax[0].elev, azim=ax[0].azim)
            elif ax[0].button_pressed in ax[0]._zoom_btn:
                ax[1].set_xlim3d(ax[0].get_xlim3d())
                ax[1].set_ylim3d(ax[0].get_ylim3d())
                ax[1].set_zlim3d(ax[0].get_zlim3d())
        elif event.inaxes == ax[1]:
            if ax[1].button_pressed in ax[1]._rotate_btn:
                ax[0].view_init(elev=ax[1].elev, azim=ax[1].azim)
            elif ax[1].button_pressed in ax[1]._zoom_btn:
                ax[0].set_xlim3d(ax[1].get_xlim3d())
                ax[0].set_ylim3d(ax[1].get_ylim3d())
                ax[0].set_zlim3d(ax[1].get_zlim3d())
        else:
            return
        fig.canvas.draw_idle()
    return fig.canvas.mpl_connect('motion_notify_event', on_move)

def plot_points(axis, p, *args, **kw):
    p = np.array(p)
    return axis.scatter3D(p[:,0], p[:,1], p[:,2], *args, **kw)

def plot_tri(axis, p, idx, **kw):
    verts = [p[i] for i in idx]
    poly = Poly3DCollection([verts])
    if 'alpha' in kw: poly.set_alpha(kw['alpha'])
    if 'c' in kw: poly.set_facecolor(kw['c'])
    poly.set_edgecolor('g')
    axis.add_collection3d(poly)
    return poly

def save_axis(fig, ax, fname):
    fpath = os.path.dirname(fname)
    if fpath and not os.path.exists(fpath):
        print(' ! creating folders %s' % fpath)
        os.makedirs(fpath)
    print(' | saving %s' % fname)
    extent = ax.get_window_extent().transformed(fig.dpi_scale_trans.inverted())
    fig.savefig(fname, bbox_inches=extent)

class PlotElement:
    def __init__(self, axis, data, default, **kw):
        self.axis, self.data, self.default = axis, data, default
        if 'size' in kw:
            kw['s'] = kw['size']
            del kw['size']
        if 'color' in kw:
            kw['c'] = kw['color']
            del kw['color']
        for k, v in default.items():
            kw[k] = kw[k] if k in kw else v
        if not 'alpha' in kw:
            kw['alpha'] = None
        self.kw = kw

class PointPlot(PlotElement):
    def __init__(self, axis, data, default, **kw):
        PlotElement.__init__(self, axis, data, default, **kw)
        self.obj = plot_points(self.axis, self.data, **self.kw)
    def hide(self): self.obj.set_alpha(0.)
    def show(self): self.obj.set_alpha(self.kw['alpha'])

class PolyPlot(PlotElement):
    def __init__(self, axis, data, indices, default, **kw):
        PlotElement.__init__(self, axis, data, default, **kw)
        self.obj = list(map(lambda idx: plot_tri(axis, data, idx, **self.kw), indices))
        self.indices = indices
    def set_alpha(self, i, alpha):
        self.obj[i].set_alpha(alpha)
        self.obj[i].set_facecolor(self.obj[i].get_facecolor())
        self.obj[i].set_edgecolor(self.obj[i].get_edgecolor())
    def hide(self): list(map(lambda i: self.set_alpha(i, 0.), range(len(self.obj))))
    def show(self): list(map(lambda i: self.set_alpha(i, self.kw['alpha']), range(len(self.obj))))

class MicellePlot:
    def __init__(self, micelle_points, boundary_points=None, convex_hull_points=None,
                    convex_hull_triangles_indices=None, non_convex_points=list()):
        if micelle_points is None:
            raise RuntimeError("The Micelle points cannot be None if a plot is desired.")
        self.micelle_points = micelle_points
        self.boundary_points = boundary_points
        self.convex_hull_points = convex_hull_points
        self.convex_hull_triangles_indices = convex_hull_triangles_indices
        self.non_convex_points = non_convex_points
        self.defs = {'undecimated' : {'s' : 10, 'c' : "#000000"},
                    'boundary' : {'s' : 10, 'c' : "blue"},
                    'convex' : {'s' : 10, 'c' : "green"},
                    'non_convex' : {'s' : 10, 'c' : "#FF00FF"},
                    'hull' : {'c' : 'black', 'alpha' : 0.1}}
        self.init_plot()
        self.plot()
    def init_plot(self, axis_on=True):
        plt.ion()
        self.fig = plt.figure(figsize=(11,5))
        self.ax = list(map(lambda i: self.fig.add_subplot(1, 2, i, projection="3d"), ((1,2))))
        plt.subplots_adjust(wspace=0, left=0, bottom=0, right=1., top=1.)
        if not axis_on: list(map(lambda x: x.axis('off'), self.ax))
        link_axes(self.fig, self.ax)
        self.cache = {}
    def plot_undecimated(self, axis, **kw):
        return PointPlot(axis, self.micelle_points, self.defs['undecimated'], **kw)
    def plot_boundary(self, axis, **kw):
        return PointPlot(axis, self.boundary_points, self.defs['boundary'], **kw)
    def plot_convex(self, axis, **kw):
        return PointPlot(axis, self.convex_hull_points, self.defs['convex'], **kw)
    def plot_non_convex(self, axis, **kw):
        return PointPlot(axis, self.non_convex_points, self.defs['non_convex'], **kw)
    def plot_hull(self, axis, **kw):
        return PolyPlot(axis, self.convex_hull_points,
                        self.convex_hull_triangles_indices,
                        self.defs['hull'], **kw)
    def save_default(self):
        save_axis(self.fig, self.ax[0], "undecimated.png")
        save_axis(self.fig, self.ax[1], "decimated.png")
    def plot(self):
        self.cache['undecimated'] = self.plot_undecimated(self.ax[0])
        self.cache['boundary'] = self.plot_boundary(self.ax[1])
        if self.convex_hull_points:
            self.cache['convex'] = self.plot_convex(self.ax[1])
            if self.convex_hull_triangles_indices:
                self.cache['hull'] = self.plot_hull(self.ax[1])
        self.cache['non_convex'] = self.plot_non_convex(self.ax[1])
        self.save_default()

class InteractiveMicellePlot(cmd.Cmd, MicellePlot):
    def __init__(self, *args, **kw):
        MicellePlot.__init__(self, *args, **kw)
        cmd.Cmd.__init__(self)
        self.prompt = '[ '
    def do_show(self, key):
        keys = key.split() + ['undecimated']
        for k,v in self.cache.items():
            if not k in keys: self.cache[k].hide()
            else: self.cache[k].show()
    def do_all(self, args):
        for k,v in self.cache.items():
            self.cache[k].show()
    def do_save(self, args):
        try:
            i, fname = args.split()
            save_axis(self.fig, self.ax[int(i)], fname)
        except:
            print(' ! invalid args %s' % args)
    def do_exit(self, args): return True
    def do_EOF(self, line): return True


''''''''''''
''' UTIL '''
''''''''''''
# plt.ion()
# fig = plt.figure(figsize=(11,5))
# ax = list(map(lambda i: fig.add_subplot(1, 2, i, projection="3d"), ((1,2))))
# plt.subplots_adjust(wspace=0, left=0, bottom=0, right=1., top=1.)
# list(map(lambda x: x.axis('off'), ax))
#
# def plot_micelle(micelle_points, boundary_points, convex_hull_points,
#         convex_hull_triangles_indices, non_convex_points=[], size=5, alpha=0.1):
#     plot_points(ax[0], micelle_points, c="#000000", s=size)
#     plot_points(ax[1], boundary_points, s=size)
#     if convex_hull_points:
#         plot_points(ax[1], convex_hull_points, c='g', s=size)
#         if convex_hull_triangles_indices:
#             for idx in convex_hull_triangles_indices:
#                 plot_tri(ax[1], convex_hull_points, idx, c='black', alpha=alpha)
#     plot_points(ax[1], non_convex_points, c="#FF00FF", s=size)

''''''''''''''''''
''' DEPRECATED '''
''''''''''''''''''
# def plot_tri_edges(axis, p, idx, *args, **kw):
#     (first_pt_x, first_pt_y, first_pt_z) = p[idx[0]]
#     (second_pt_x, second_pt_y, second_pt_z) = p[idx[1]]
#     (third_pt_x, third_pt_y, third_pt_z) = p[idx[2]]
#     axis.plot([first_pt_x, second_pt_x], [first_pt_y, second_pt_y], [first_pt_z, second_pt_z], *args, **kw)
#     axis.plot([second_pt_x, third_pt_x], [second_pt_y, third_pt_y], [second_pt_z, third_pt_z], *args, **kw)
#     axis.plot([third_pt_x, first_pt_x], [third_pt_y, first_pt_y], [third_pt_z, first_pt_z], *args, **kw)