Source code for chemtools.outputs.vmd

# -*- coding: utf-8 -*-
# ChemTools is a collection of interpretive chemical tools for
# analyzing outputs of the quantum chemistry calculations.
#
# Copyright (C) 2016-2019 The ChemTools Development Team
#
# This file is part of ChemTools.
#
# ChemTools is free software; you can redistribute it and/or
# modify it under the terms of the GNU General Public License
# as published by the Free Software Foundation; either version 3
# of the License, or (at your option) any later version.
#
# ChemTools is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, see <http://www.gnu.org/licenses/>
#
# --
"""Output Module.

This module contains functions for generating scripts for visualizing
purposes using VMD, GaussView, etc.
"""
import os
import numpy as np

__all__ = ['print_vmd_script_nci', 'print_vmd_script_isosurface',
           'print_vmd_script_multiple_cube', 'print_vmd_script_vector_field']


def _vmd_script_start():
    """Generate part of the beginning part of the VMD script.

    Returns
    -------
    VMD script responsible for beginning of VMD script
    """
    return ('#!/usr/local/bin/vmd\n'
            '# VMD script written by save_state $Revision: 1.41 $\n'
            '# VMD version: 1.8.6\n'
            'set viewplist\n'
            'set fixedlist\n'
            '#\n'
            '# Display settings\n'
            'display projection Perspective\n'
            'display nearclip set 0.000000\n'
            'display shadow off\n'
            'display rendermode GLSL\n'
            'color Element {C} gray\n'
            'color Element {Cl} green\n'
            'axes location Off\n'
            'light 2 on\n'
            'light 3 on\n'
            '#\n')


def _vmd_script_molecule(representation, *mol_files):
    """Generate part of the VMD script that loads the molecule information.

    Parameters
    ----------
    representation: str
        Representation of a molecule. Choose either 'CPK' or 'Line'.
    mol_files : str
        Names of the input files that represent the molecule
        .xyz and .cube files are supported
        Cube files can correspond to the density, reduced density gradient, isosurface, color of
        isosurface, etc

    Returns
    -------
    Part of the VMD script that constructs the molecule

    Raises
    ------
    ValueError
        If no files are given
    TypeError
        If unsupported file type (i.e. not xyz or cube)
    """
    output = '# load new molecule\n'
    if len(mol_files) == 0:
        raise ValueError('Need at least one molecule file')
    for i, mol in enumerate(mol_files):
        if i == 0:
            mol_type = 'new'
        else:
            mol_type = 'addfile'

        ext = os.path.splitext(mol)[1]
        if ext == '.xyz':
            file_type = '{xyz}'
        elif ext == '.cube':
            file_type = 'cube'
        else:
            raise TypeError('Unsupported file type, {0}'.format(ext))

        output += ('mol {0} {1} type {2} first 0 last -1 step 1 filebonds 1 autobonds 1 waitfor all'
                   '\n'.format(mol_type, mol, file_type))
    output += ('#\n'
               '# representation of the atoms\n')
    if representation.lower() == 'cpk':
        output += 'mol representation CPK 1.000000 0.300000 118.000000 131.000000\n'
    elif representation.lower() == 'line':
        output += 'mol representation Lines 4.000000\n'
    else:
        raise ValueError('Argument representation is not recognized.')
    output += ('mol delrep 0 top\n'
               'mol color Element\n'
               'mol selection {{all}}\n'
               'mol material Opaque\n'
               'mol addrep top\n'
               '#\n')
    return output


def _vmd_script_isosurface(isosurf=0.5, index=0, show_type='isosurface', draw_type='solid surface',
                           material='Opaque', scalemin=-0.05, scalemax=0.05, colorscheme='RGB'):
    """Generate part of the VMD script that configures the isosurface.

    Parameters
    ----------
    isosurf : float
        Iso-surface value to plot.
    index : int, optional
        Index of the file that contains the iso-surface data (in the order loaded by
        `_vmd_script_molecule`)
    show_type : str, optional
        Option that controls what will be shown
        One of 'isosurface', 'box', and 'box+isosurface'
    draw_type : str, optional
        Option that controls how the iso-surface will be drawn
        One of 'solid surface', 'wireframe', 'points', 'shaded points'
    material : str, optional
        The material setting of the iso-surface
        One of 'Opaque', 'Transparent', 'BrushedMetal', 'Diffuse', 'Ghost', 'Glass1', 'Glass2',
        'Glass3', 'Glossy', 'HardPlastic', 'MetallicPastel', 'Steel', 'Translucent', 'Edgy',
        'EdgyShiny', 'EdgyGlass', 'Goodsell', 'AOShiny', 'AOChalky', 'AOEdgy', 'BlownGlass',
        'GlassBubble', 'RTChrome'.
    scalemin : float, optional
        Smallest value to color on the iso-surface.
    scalemax : float, optional
        Largest value to color on the iso-surface.
    colorscheme : str or int, optional
        Color scheme used for the isosurface.
        If str, then appropriate gradient is used to color the iso-surface. It must be one of the
        following options
            =======  =====================================
            Options  Description
            =======  =====================================
            'RGB'    small=red, middle=green, large=blue
            'BGR'    small=blue, middle=green, large=red
            'RWB'    small=red, middle=white, large=blue
            'BWR'    small=blue, middle=white, large=red
            'RWG'    small=red, middle=white, large=green
            'GWR'    small=green, middle=white, large=red
            'GWB'    small=green, middle=white, large=blue
            'BWG'    small=blue, middle=white, large=green
            'BlkW'   small=black, large=white
            'WBlk'   small=white, large=black
            =======  =====================================
        If int, isosurface is colored with just one color. There are 1057 colors are available in
        VMD, check the program or website for more details.

    Returns
    -------
    Part of the VMD script that constructs the isosurface

    Raises
    ------
    TypeError
        If `isosurf` is not a float
        If `index` is not an integer
        If `show_type` is not one of 'isosurface', 'box', or 'box+isosurface'
        If `draw_type` is not one of 'solid surface', 'wireframe', 'points', or 'shaded points'
        If `material` is not one of 'Opaque', 'Transparent', 'BrushedMetal', 'Diffuse', 'Ghost',
        'Glass1', 'Glass2', 'Glass3', 'Glossy', 'HardPlastic', 'MetallicPastel', 'Steel',
        'Translucent', 'Edgy', 'EdgyShiny', 'EdgyGlass', 'Goodsell', 'AOShiny', 'AOChalky',
        'AOEdgy', 'BlownGlass', 'GlassBubble', 'RTChrome'
        If `scalemin` is not float
        If `scalemax` is not float
        If `colorscheme` is not one of 'RGB', 'BGR', 'RWB', 'BWR', 'RWG', 'GWR', 'GWB', 'BWG',
        'BlkW', 'WBlk' or int
    """
    if not isinstance(isosurf, float):
        raise TypeError('`isosurf` must be a float')

    if not isinstance(index, int):
        raise TypeError('`index` must be an integer')

    if show_type not in ['isosurface', 'box', 'box+isosurface']:
        raise TypeError('Unsupported `show_type`. Must be one of {0}, {1}, or {2}'
                        ''.format('box', 'isosurface', 'box+isosurface'))
    show_type = {'isosurface': 0, 'box': 1, 'box+isosurface': 2}[show_type]

    if draw_type not in ['solid surface', 'wireframe', 'points', 'shaded points']:
        raise TypeError('Unsupported `draw_type`. Must be one of {0}, {1}, {2} or {3}'
                        ''.format('solid surface', 'wireframe', 'points', 'shaded points'))
    draw_type = {'solid surface': 0, 'wireframe': 1, 'points': 2, 'shaded points': 3}[draw_type]

    allowed_materials = ['Opaque', 'Transparent', 'BrushedMetal', 'Diffuse', 'Ghost', 'Glass1',
                         'Glass2', 'Glass3', 'Glossy', 'HardPlastic', 'MetallicPastel', 'Steel',
                         'Translucent', 'Edgy', 'EdgyShiny', 'EdgyGlass', 'Goodsell', 'AOShiny',
                         'AOChalky', 'AOEdgy', 'BlownGlass', 'GlassBubble', 'RTChrome']
    if material not in allowed_materials:
        raise TypeError('Unsupported `material`. Must be one of {0}'.format(allowed_materials))

    if not isinstance(scalemin, float):
        raise TypeError('`scalemin` must be a float')
    if not isinstance(scalemax, float):
        raise TypeError('`scalemax` must be a float')

    if (not isinstance(colorscheme, (str, int)) or
            (isinstance(colorscheme, str) and
             colorscheme not in ['RGB', 'BGR', 'RWB', 'BWR', 'RWG', 'GWR', 'GWB', 'BWG', 'BlkW',
                                 'WBlk']) or
            (isinstance(colorscheme, int) and not 0 <= colorscheme < 1057)):
        raise TypeError('Unsupported colorscheme, {0}'.format(colorscheme))

    output = '# add representation of the surface\n'
    # mol representation Isosurface {0} {1} {2} {3} {4} {5}
    # {0} = isovalue
    # {1} = file index that will be plotted
    # {2} = what to show
    #       0 -> isosurface
    #       1 -> box
    #       2 -> box and isosurface
    # {3} = how to draw
    #       0 -> solid surface
    #       1 -> wireframe
    #       2 -> points
    #       3 -> shaded points
    # {4} = (integer) step size (distance between plot objects)
    # {5} = (integer) width of plot object (e.g. wire, point)
    output += ('mol representation Isosurface {isosurf:.5f} {index} {show} {draw} 1 1\n'
               ''.format(isosurf=isosurf, index=index, show=show_type, draw=draw_type))

    if isinstance(colorscheme, int):
        output += 'mol color ColorID {0}\n'.format(colorscheme)
    else:
        output += 'mol color Volume 0\n'

    output += ('mol selection {{all}}\n'
               'mol material {material}\n'
               'mol addrep top\n'
               'mol selupdate 1 top 0\n'
               'mol colupdate 1 top 0\n'
               'mol scaleminmax top 1 {scalemin:.6f} {scalemax:.6f}\n'
               'mol smoothrep top 1 0\n'
               'mol drawframes top 1 {{now}}\n'.format(material=material, scalemin=scalemin,
                                                       scalemax=scalemax))

    if isinstance(colorscheme, str):
        output += 'color scale method {0}\n'.format(colorscheme)
    else:
        output += 'color scale method RGB\n'
    output += 'color Display Background silver\n'
    output += '#\n'
    return output


def _vmd_script_vector_field(centers, unit_vecs, weights, weight_threshold=1e-1,
                             has_shadow=True, material='Transparent', color=0):
    """Generate part of the VMD script that constructs the vector field.

    Parameters
    ----------
    centers : np.ndarray(N, 3)
        Coordinates of the centers of each vector
    unit_vecs : np.ndarray(N, 3)
        Unit direction of each vector
    weights : np.ndarray(N)
        Weights that determine the size (length and/or thickness) of each vector
    weight_threshold : float
        Threshold of the weight of the vector
        Vectors with norms smaller than this threshold are not plotted
        Default is 1e-1
    material : str
        The material setting of each vector
        One of 'Opaque', 'Transparent', 'BrushedMetal', 'Diffuse', 'Ghost', 'Glass1', 'Glass2',
        'Glass3', 'Glossy', 'HardPlastic', 'MetallicPastel', 'Steel', 'Translucent', 'Edgy',
        'EdgyShiny', 'EdgyGlass', 'Goodsell', 'AOShiny', 'AOChalky', 'AOEdgy', 'BlownGlass',
        'GlassBubble', 'RTChrome'.
        Default is 'Transparent'
    color : int
        Color of each vector
        Integer between 0 and 1057. See VMD program or manual for details.
        Default color is Blue.

    Returns
    -------
    VMD script for generating vector field

    Raises
    ------
    ValueError
        If unit_vecs are not unit vectors
        If centers, unit_vecs, and weights do not have the same number of rows (vectors)
    TypeError
        If centers is not a numpy array with 3 columns
        If unit_vecs is not a numpy array with 3 columns
        If weights is not a one dimensional numpy array
        If has_shadow is not boolean
        If material is not supported
        If color is not supported

    Note
    ----
    If there are too many vectors, you many need to turn off has_shadow (or materials in VMD)
    """
    # check unit directions
    if not (isinstance(centers, np.ndarray) and len(centers.shape) == 2 and centers.shape[1] == 3):
        raise TypeError('Given centers must be a two-dimensional numpy array with 3 columns')
    if not (isinstance(unit_vecs, np.ndarray) and len(unit_vecs.shape) == 2 and
            unit_vecs.shape[1] == 3):
        raise TypeError('Given unit_vecs must be a two-dimensional numpy array with 3 '
                        'columns')
    if not (isinstance(weights, np.ndarray) and len(weights.shape) == 1):
        raise TypeError('Given weights must be a one-dimensional numpy array')
    if not centers.size/3 == unit_vecs.size/3 == weights.size:
        raise ValueError('The number of vectors must match up with centers, unit_vecs, and '
                         'weights')

    if not np.allclose(np.linalg.norm(unit_vecs, axis=1), 1):
        raise ValueError('Given direction vectors are not unit vectors')

    if not isinstance(has_shadow, bool):
        raise TypeError('Given has_shadow must be boolean')
    if has_shadow:
        has_shadow = 'on'
    else:
        has_shadow = 'off'

    allowed_materials = ['Opaque', 'Transparent', 'BrushedMetal', 'Diffuse', 'Ghost', 'Glass1',
                         'Glass2', 'Glass3', 'Glossy', 'HardPlastic', 'MetallicPastel', 'Steel',
                         'Translucent', 'Edgy', 'EdgyShiny', 'EdgyGlass', 'Goodsell', 'AOShiny',
                         'AOChalky', 'AOEdgy', 'BlownGlass', 'GlassBubble', 'RTChrome']
    if material not in allowed_materials:
        raise TypeError('Unsupported `material`. Must be one of {0}'.format(allowed_materials))

    if isinstance(color, int) and not 0 <= color < 1057:
        raise TypeError('Unsupported color, {0}'.format(color))

    # vmd/tcl function for constructing arrow
    output = '# Add function for vector field\n'
    output += ('proc vmd_draw_arrow {mol center unit_dir cyl_radius cone_radius length} {\n'
               'set start [vecsub $center [vecscale [vecscale 0.5 $length] $unit_dir]]\n'
               'set end [vecadd $start [vecscale $length $unit_dir]]\n'
               'set middle [vecsub $end [vecscale [vecscale 1.732050808 $cone_radius] $unit_dir]]\n'
               'graphics $mol cylinder $start $middle radius $cyl_radius\n'
               'graphics $mol cone $middle $end radius $cone_radius\n'
               '}\n'
               '#\n')

    output += 'draw materials {0}\n'.format(has_shadow)
    output += 'draw material {0}\n'.format(material)
    output += 'draw color {0}\n'.format(color)

    def decompose_weight(weight):
        """Decompose a weight to the corresponding cylinder radius, cone radius and length.

        Parameters
        ----------
        weight : float
            Weight of a vector

        Returns
        -------
        cyl_radius : float
            Radius of cylinder in vector
        cone_radius : float
            Radius of cone in vector
        length : float
            Length of vector
        """
        return np.array([0.08, 0.15, 0.7])*(np.array(weight)**0.333)

    for (center_x, center_y, center_z), (unit_x, unit_y, unit_z), weight in zip(centers,
                                                                                unit_vecs,
                                                                                weights):
        if weight < weight_threshold:
            continue
        output += ('draw arrow {{{0} {1} {2}}} {{{3} {4} {5}}} {6} {7} {8}\n'
                   ''.format(center_x, center_y, center_z, unit_x, unit_y, unit_z,
                             *decompose_weight(weight)))
    output += '#\n'
    return output














def print_vmd_script_topology(fname, points, radius=0.2):
    """Generate VMD script for visualizing critical points & gradient path.

    Parameters
    ----------
    fname : str
        The name of the VMD script file.
    points : dict
        Dictionary of color (key) and coordinates (values) of critical points.
    radius : float, optional
        Radius of spheres representing the critical points.

    """
    output = ('#!/usr/local/bin/vmd\n'
              '#\n'
              '# Display settings\n'
              'display shadow off\n'
              'axes location Off\n'
              'light 2 on\n'
              'light 3 on\n'
              'color Display Background white\n'
              '#\n'
              '# Critical Points\n')
    for color, coordinates in points.items():
        output += 'draw color {}\n'.format(color)
        for coord in coordinates:
            coord = '% .6f % .6f % .6f' % tuple(coord)
            # resolution determines how many polygons are used in the approximation of a sphere
            output += 'draw sphere {%s} radius %f resolution 36\n' % (coord, radius)
    # write output
    if not fname.endswith('.vmd'):
        fname += '.vmd'
    with open(fname, 'w') as f:
        f.write(output)