|
| 1 | +# |
| 2 | +# PCMSolver, an API for the Polarizable Continuum Model |
| 3 | +# Copyright (C) 2017 Roberto Di Remigio, Luca Frediani and collaborators. |
| 4 | +# |
| 5 | +# This file is part of PCMSolver. |
| 6 | +# |
| 7 | +# PCMSolver is free software: you can redistribute it and/or modify |
| 8 | +# it under the terms of the GNU Lesser General Public License as published by |
| 9 | +# the Free Software Foundation, either version 3 of the License, or |
| 10 | +# (at your option) any later version. |
| 11 | +# |
| 12 | +# PCMSolver is distributed in the hope that it will be useful, |
| 13 | +# but WITHOUT ANY WARRANTY; without even the implied warranty of |
| 14 | +# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the |
| 15 | +# GNU Lesser General Public License for more details. |
| 16 | +# |
| 17 | +# You should have received a copy of the GNU Lesser General Public License |
| 18 | +# along with PCMSolver. If not, see <http://www.gnu.org/licenses/>. |
| 19 | +# |
| 20 | +# For information on the complete list of contributors to the |
| 21 | +# PCMSolver API, see: <http://pcmsolver.readthedocs.io/> |
| 22 | +# |
| 23 | + |
| 24 | +# -*- python -*- |
| 25 | +# -*- coding: utf-8 -*- |
| 26 | +# vim:filetype=python: |
| 27 | + |
| 28 | +# Written by Roberto Di Remigio <[email protected]> |
| 29 | +# University of Tromso, 2017 |
| 30 | + |
| 31 | +""" |
| 32 | +Plot molecular cavity from a compressed NumPy format file. |
| 33 | +Color map the finite elements according to a surface function, saved |
| 34 | +to NumPy format file. |
| 35 | +""" |
| 36 | + |
| 37 | +import os |
| 38 | +import sys |
| 39 | +sys.path.append(os.path.dirname(__file__)) |
| 40 | + |
| 41 | +try: |
| 42 | + import docopt |
| 43 | +except: |
| 44 | + sys.path.append('cmake/lib/docopt') |
| 45 | + import docopt |
| 46 | + |
| 47 | +import numpy as np |
| 48 | +import matplotlib as mpl |
| 49 | +import matplotlib.pyplot as plt |
| 50 | +from matplotlib.colors import Normalize |
| 51 | +from matplotlib import cm |
| 52 | +from mpl_toolkits.mplot3d.art3d import Poly3DCollection |
| 53 | + |
| 54 | +options = """ |
| 55 | +Usage: |
| 56 | + ./plot_cavity.py <cavity_npz> [--map-by <npy>] |
| 57 | + ./plot_cavity.py (-h | --help) |
| 58 | +
|
| 59 | +Options: |
| 60 | + <cavity_npz> Compressed NumPy file with cavity specifications. |
| 61 | + --map-by <npy> NumPy format file with surface function to color-map finite elements. |
| 62 | + -h --help Show this screen. |
| 63 | +""" |
| 64 | + |
| 65 | +def shiftedColorMap(cmap, start=0, midpoint=0.5, stop=1.0, name='shiftedcmap'): |
| 66 | + ''' |
| 67 | + Function to offset the "center" of a colormap. Useful for |
| 68 | + data with a negative min and positive max and you want the |
| 69 | + middle of the colormap's dynamic range to be at zero |
| 70 | +
|
| 71 | + Input |
| 72 | + ----- |
| 73 | + cmap : The matplotlib colormap to be altered |
| 74 | + start : Offset from lowest point in the colormap's range. |
| 75 | + Defaults to 0.0 (no lower ofset). Should be between |
| 76 | + 0.0 and `midpoint`. |
| 77 | + midpoint : The new center of the colormap. Defaults to |
| 78 | + 0.5 (no shift). Should be between 0.0 and 1.0. In |
| 79 | + general, this should be 1 - vmax/(vmax + abs(vmin)) |
| 80 | + For example if your data range from -15.0 to +5.0 and |
| 81 | + you want the center of the colormap at 0.0, `midpoint` |
| 82 | + should be set to 1 - 5/(5 + 15)) or 0.75 |
| 83 | + stop : Offset from highets point in the colormap's range. |
| 84 | + Defaults to 1.0 (no upper ofset). Should be between |
| 85 | + `midpoint` and 1.0. |
| 86 | + ''' |
| 87 | + cdict = { |
| 88 | + 'red': [], |
| 89 | + 'green': [], |
| 90 | + 'blue': [], |
| 91 | + 'alpha': [] |
| 92 | + } |
| 93 | + |
| 94 | + # regular index to compute the colors |
| 95 | + reg_index = np.linspace(start, stop, 257) |
| 96 | + |
| 97 | + # shifted index to match the data |
| 98 | + shift_index = np.hstack([ |
| 99 | + np.linspace(0.0, midpoint, 128, endpoint=False), |
| 100 | + np.linspace(midpoint, 1.0, 129, endpoint=True) |
| 101 | + ]) |
| 102 | + |
| 103 | + for ri, si in zip(reg_index, shift_index): |
| 104 | + r, g, b, a = cmap(ri) |
| 105 | + |
| 106 | + cdict['red'].append((si, r, r)) |
| 107 | + cdict['green'].append((si, g, g)) |
| 108 | + cdict['blue'].append((si, b, b)) |
| 109 | + cdict['alpha'].append((si, a, a)) |
| 110 | + |
| 111 | + newcmap = mpl.colors.LinearSegmentedColormap(name, cdict) |
| 112 | + plt.register_cmap(cmap=newcmap) |
| 113 | + |
| 114 | + return newcmap |
| 115 | + |
| 116 | + |
| 117 | +def plot(cavity_npz, surf_func_npy=None): |
| 118 | + fig = plt.figure() |
| 119 | + ax = fig.add_subplot(111, projection='3d') |
| 120 | + |
| 121 | + cavity = np.load(cavity_npz) |
| 122 | + |
| 123 | + nElements = cavity['elements'] |
| 124 | + centroids = cavity['centers'] |
| 125 | + |
| 126 | + # Plot collocation points |
| 127 | + ax.scatter(centroids[0, :], centroids[1,:], centroids[2,:], c='black', alpha=0.5) |
| 128 | + |
| 129 | + # Generate color mapping |
| 130 | + colors = (.5, .1, .3, 0.3) |
| 131 | + if surf_func_npy: |
| 132 | + surf_func = np.load(surf_func_npy) |
| 133 | + shifted_cmap = shiftedColorMap(cm.coolwarm, midpoint=0.75, name='shifted') |
| 134 | + mappable = cm.ScalarMappable(cmap=shifted_cmap) |
| 135 | + mappable.set_array(surf_func.flatten()) |
| 136 | + plt.colorbar(mappable) |
| 137 | + # Provide colors for Poly3DCollection |
| 138 | + colors = mappable.to_rgba(surf_func.flatten()) |
| 139 | + # Generate list of vertices |
| 140 | + vertices = [zip(cavity['vertices_' + str(i)][0, :], cavity['vertices_' + str(i)][1, :], cavity['vertices_' + str(i)][2, :]) for i in range(nElements)] |
| 141 | + elements = Poly3DCollection(vertices, facecolors=colors) |
| 142 | + ax.add_collection3d(elements) |
| 143 | + ax.set_axis_off() |
| 144 | + plt.show() |
| 145 | + |
| 146 | + |
| 147 | +def main(): |
| 148 | + try: |
| 149 | + arguments = docopt.docopt(options, argv=None) |
| 150 | + except docopt.DocoptExit: |
| 151 | + sys.stderr.write('ERROR: bad input to %s\n' % sys.argv[0]) |
| 152 | + sys.stderr.write(options) |
| 153 | + sys.exit(-1) |
| 154 | + cavity_npz = arguments['<cavity_npz>'] |
| 155 | + surf_func_npy = arguments['--map-by'] |
| 156 | + plot(cavity_npz, surf_func_npy) |
| 157 | + |
| 158 | + |
| 159 | +if __name__ == '__main__': |
| 160 | + main() |
0 commit comments