Source code for ase.visualize.mlab

import optparse

import numpy as np

from ase.calculators.calculator import get_calculator_class
from import covalent_radii
from import cpk_colors
from import read_cube_data

[docs]def plot(atoms, data, contours): """Plot atoms, unit-cell and iso-surfaces using Mayavi. Parameters: atoms: Atoms object Positions, atomiz numbers and unit-cell. data: 3-d ndarray of float Data for iso-surfaces. countours: list of float Contour values. """ # Delay slow imports: import os from mayavi import mlab # mayavi GUI bug fix for remote access via ssh (X11 forwarding) if "SSH_CONNECTION" in os.environ: f = mlab.gcf() f.scene._lift() mlab.figure(1, bgcolor=(1, 1, 1)) # make a white figure # Plot the atoms as spheres: for pos, Z in zip(atoms.positions, atoms.numbers): mlab.points3d(*pos, scale_factor=covalent_radii[Z], resolution=20, color=tuple(cpk_colors[Z])) # Draw the unit cell: A = atoms.cell for i1, a in enumerate(A): i2 = (i1 + 1) % 3 i3 = (i1 + 2) % 3 for b in [np.zeros(3), A[i2]]: for c in [np.zeros(3), A[i3]]: p1 = b + c p2 = p1 + a mlab.plot3d([p1[0], p2[0]], [p1[1], p2[1]], [p1[2], p2[2]], tube_radius=0.1) cp = mlab.contour3d(data, contours=contours, transparent=True, opacity=0.5, colormap='hot') # Do some tvtk magic in order to allow for non-orthogonal unit cells: polydata =[0].mapper.input pts = np.array(polydata.points) - 1 # Transform the points to the unit cell: polydata.points =, A / np.array(data.shape)[:, np.newaxis]) # Apparently we need this to redraw the figure, maybe it can be done in # another way? mlab.view(azimuth=155, elevation=70, distance='auto') # Show the 3d plot:
def view_mlab(atoms, *args, **kwargs): return plot(atoms, *args, **kwargs) description = """\ Plot iso-surfaces from a cube-file or a wave function or an electron density from a calculator-restart file.""" def main(args=None): parser = optparse.OptionParser(usage='%prog [options] filename', description=description) add = parser.add_option add('-n', '--band-index', type=int, metavar='INDEX', help='Band index counting from zero.') add('-s', '--spin-index', type=int, metavar='SPIN', help='Spin index: zero or one.') add('-e', '--electrostatic-potential', action='store_true', help='Plot the electrostatic potential.') add('-c', '--contours', default='4', help='Use "-c 3" for 3 contours or "-c -0.5,0.5" for specific ' + 'values. Default is four contours.') add('-r', '--repeat', help='Example: "-r 2,2,2".') add('-C', '--calculator-name', metavar='NAME', help='Name of calculator.') opts, args = parser.parse_args(args) if len(args) != 1: parser.error('Incorrect number of arguments') arg = args[0] if arg.endswith('.cube'): data, atoms = read_cube_data(arg) else: calc = get_calculator_class(opts.calculator_name)(arg, txt=None) atoms = calc.get_atoms() if opts.band_index is None: if opts.electrostatic_potential: data = calc.get_electrostatic_potential() else: data = calc.get_pseudo_density(opts.spin_index) else: data = calc.get_pseudo_wave_function(opts.band_index, opts.spin_index or 0) if data.dtype == complex: data = abs(data) mn = data.min() mx = data.max() print('Min: %16.6f' % mn) print('Max: %16.6f' % mx) if opts.contours.isdigit(): n = int(opts.contours) d = (mx - mn) / n contours = np.linspace(mn + d / 2, mx - d / 2, n).tolist() else: contours = [float(x) for x in opts.contours.rstrip(',').split(',')] if len(contours) == 1: print('1 contour:', contours[0]) else: print('%d contours: %.6f, ..., %.6f' % (len(contours), contours[0], contours[-1])) if opts.repeat: repeat = [int(r) for r in opts.repeat.split(',')] data = np.tile(data, repeat) atoms *= repeat plot(atoms, data, contours) if __name__ == '__main__': main()