import optparse
import numpy as np
from ase.calculators.calculator import get_calculator_class
from ase.data import covalent_radii
from ase.data.colors import cpk_colors
from ase.io.cube 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 = cp.actor.actors[0].mapper.input
pts = np.array(polydata.points) - 1
# Transform the points to the unit cell:
polydata.points = np.dot(pts, 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:
mlab.show()
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()