# coding: utf-8
"""
mayavi_ toolkit.
WARNING: This code is still under development.
"""
import itertools
import numpy as np
DEFAULT_FIGURE_KWARGS = dict(size=(1024, 768), bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
[docs]
def get_fig_mlab(figure=None, **kwargs): # pragma: no cover
try:
from mayavi import mlab
except ImportError as exc:
print("mayavi is not installed. Use `conda install mayavi` or `pip install mayavi`")
raise exc
# To use the full envisage application
#mlab.options.backend = "envisage"
#mlab.options.backend = "test"
#mlab.options.offscreen = True
if figure is None:
# Add defaults
for k, v in DEFAULT_FIGURE_KWARGS.items():
if k not in kwargs: kwargs[k] = v
figure = mlab.figure(**kwargs)
return figure, mlab
[docs]
def plot_wigner_seitz(lattice, figure=None, **kwargs): # pragma: no cover
"""
Adds the skeleton of the Wigner-Seitz cell of the lattice to a mayavi_ figure
Args:
lattice: Reciprocal-space |Lattice| object
figure: mayavi figure, None to plot on the curretn figure
kwargs: kwargs passed to the mayavi function ``plot3d``. Color defaults to black
and line_width to 1.
Returns: mayavi figure
"""
figure, mlab = get_fig_mlab(figure=figure)
if "color" not in kwargs:
kwargs["color"] = (0, 0, 0)
if "line_width" not in kwargs:
kwargs["line_width"] = 1
if "tube_radius" not in kwargs:
kwargs["tube_radius"] = None
bz = lattice.get_wigner_seitz_cell()
for iface in range(len(bz)):
for line in itertools.combinations(bz[iface], 2):
for jface in range(len(bz)):
if iface < jface and any(np.all(line[0] == x) for x in bz[jface])\
and any(np.all(line[1] == x) for x in bz[jface]):
#do_plot = True
#if in_unit_cell:
# kred0 = lattice.get_fractional_coords(line[0])
# kred1 = lattice.get_fractional_coords(line[1])
# do_plot = np.alltrue((kred0 >= 0) & (kred0 <= 0.5) &
# (kred1 >= 0) & (kred1 <= 0.5))
# print(kred0, kred1, do_plot)
#if not do_plot: continue
mlab.plot3d(*zip(line[0], line[1]), figure=figure, **kwargs)
return figure
[docs]
def plot_unit_cell(lattice, figure=None, **kwargs): # pragma: no cover
"""
Adds the unit cell of the lattice to a mayavi_ figure.
Args:
lattice: Lattice object
figure: mayavi figure, None to plot on the curretn figure
kwargs: kwargs passed to the mayavi function ``plot3d``. Color defaults to black
and line_width to 1.
Returns: mayavi figure
"""
figure, mlab = get_fig_mlab(figure=figure)
if "color" not in kwargs:
kwargs["color"] = (0, 0, 0)
if "line_width" not in kwargs:
kwargs["line_width"] = 1
if "tube_radius" not in kwargs:
kwargs["tube_radius"] = None
v = 8 * [None]
v[0] = lattice.get_cartesian_coords([0.0, 0.0, 0.0])
v[1] = lattice.get_cartesian_coords([1.0, 0.0, 0.0])
v[2] = lattice.get_cartesian_coords([1.0, 1.0, 0.0])
v[3] = lattice.get_cartesian_coords([0.0, 1.0, 0.0])
v[4] = lattice.get_cartesian_coords([0.0, 1.0, 1.0])
v[5] = lattice.get_cartesian_coords([1.0, 1.0, 1.0])
v[6] = lattice.get_cartesian_coords([1.0, 0.0, 1.0])
v[7] = lattice.get_cartesian_coords([0.0, 0.0, 1.0])
for i, j in ((0, 1), (1, 2), (2, 3), (0, 3), (3, 4), (4, 5), (5, 6),
(6, 7), (7, 4), (0, 7), (1, 6), (2, 5), (3, 4)):
mlab.plot3d(*zip(v[i], v[j]), figure=figure, **kwargs)
#mlab.xlabel("x-axis")
#mlab.ylabel("y-axis")
#mlab.zlabel("z-axis")
return figure
[docs]
def plot_lattice_vectors(lattice, figure=None, **kwargs): # pragma: no cover
"""
Adds the basis vectors of the lattice provided to a mayavi_ figure.
Args:
lattice: |Lattice| object.
figure: mayavi figure, None if a new figure should be created.
kwargs: kwargs passed to the mayavi function ``plot3d``. Color defaults to black
and line_width to 1.
Returns: mayavi figure
"""
figure, mlab = get_fig_mlab(figure=figure)
if "color" not in kwargs:
kwargs["color"] = (0, 0, 0)
if "line_width" not in kwargs:
kwargs["line_width"] = 1
if "tube_radius" not in kwargs:
kwargs["tube_radius"] = None
vertex1 = lattice.get_cartesian_coords([0.0, 0.0, 0.0])
vertex2 = lattice.get_cartesian_coords([1.0, 0.0, 0.0])
mlab.plot3d(*zip(vertex1, vertex2), figure=figure, **kwargs)
vertex2 = lattice.get_cartesian_coords([0.0, 1.0, 0.0])
mlab.plot3d(*zip(vertex1, vertex2), figure=figure, **kwargs)
vertex2 = lattice.get_cartesian_coords([0.0, 0.0, 1.0])
mlab.plot3d(*zip(vertex1, vertex2), figure=figure, **kwargs)
return figure
[docs]
def plot_structure(structure, frac_coords=False, to_unit_cell=False, style="points+labels",
unit_cell_color=(0, 0, 0), color_scheme="VESTA", figure=None, show=False, **kwargs): # pragma: no cover
"""
Plot structure with mayavi.
Args:
structure: |Structure| object
frac_coords:
to_unit_cell: True if sites should be wrapped into the first unit cell.
style: "points+labels" to show atoms sites with labels.
unit_cell_color:
color_scheme: color scheme for atom types. Allowed values in ("Jmol", "VESTA")
figure:
kwargs:
Returns: mayavi figure
"""
figure, mlab = get_fig_mlab(figure=figure)
#if not frac_coords:
plot_unit_cell(structure.lattice, color=unit_cell_color, figure=figure)
from pymatgen.analysis.molecule_structure_comparator import CovalentRadius
from pymatgen.vis.structure_vtk import EL_COLORS
for site in structure:
symbol = site.specie.symbol
color = tuple(i / 255 for i in EL_COLORS[color_scheme][symbol])
radius = CovalentRadius.radius[symbol]
if to_unit_cell and hasattr(site, "to_unit_cell"): site = site.to_unit_cell
x, y, z = site.frac_coords if frac_coords else site.coords
if "points" in style:
mlab.points3d(x, y, z, figure=figure, scale_factor=radius,
resolution=20, color=color, scale_mode='none', **kwargs)
if "labels" in style:
mlab.text3d(x, y, z, symbol, figure=figure, color=(0, 0, 0), scale=0.2)
if show: mlab.show()
return figure
[docs]
def plot_labels(labels, lattice=None, coords_are_cartesian=False, figure=None, **kwargs): # pragma: no cover
"""
Adds labels to a mayavi_ figure.
Args:
labels: dict containing the label as a key and the coordinates as value.
lattice: |Lattice| object used to convert from reciprocal to cartesian coordinates
coords_are_cartesian: Set to True if you are providing.
coordinates in cartesian coordinates. Defaults to False.
Requires lattice if False.
figure: mayavi figure, None to plot on the curretn figure
kwargs: kwargs passed to the mayavi function `text3d`. Color defaults to blue and size to 25.
Returns: mayavi figure
"""
figure, mlab = get_fig_mlab(figure=figure)
#if "color" not in kwargs:
# kwargs["color"] = "b"
#if "size" not in kwargs:
# kwargs["size"] = 25
#if "width" not in kwargs:
# kwargs["width"] = 0.8
if "scale" not in kwargs:
kwargs["scale"] = 0.1
for k, coords in labels.items():
label = k
if k.startswith("\\") or k.find("_") != -1:
label = "$" + k + "$"
off = 0.01
if coords_are_cartesian:
coords = np.array(coords)
else:
if lattice is None:
raise ValueError("coords_are_cartesian False requires the lattice")
coords = lattice.get_cartesian_coords(coords)
x, y, z = coords + off
mlab.text3d(x, y, z, label, figure=figure, **kwargs)
return figure
[docs]
class MayaviFieldAnimator(object): # pragma: no cover
def __init__(self, filepaths):
self.filepaths = filepaths
self.num_files = len(filepaths)
[docs]
def volume_animate(self):
from abipy import abilab
with abilab.abiopen(self.filepaths[0]) as nc:
nsppol, nspden, nspinor = nc.nsppol, nc.nspden, nc.nspinor
structure = nc.structure
datar = nc.field.datar
# [nspden, nx, ny, nz] array
nx, ny, nz = datar.shape[1:]
s = datar[0]
print(s.dtype, s.shape)
#cart_coords = np.empty((nx*ny*nz, 3))
#cnt = 0
#for i in range(nx):
# for j in range(ny):
# for k in range(nz):
# cart_coords[ctn, :] = (i/nx, j/ny, k/nz)
# cnt += 1
#cart_coords = structure.lattice.get_cartesian_coords(cart_coords)
# We reorder the points, scalars and vectors so this is as per VTK's
# requirement of x first, y next and z last.
#pts = pts.transpose(2, 1, 0, 3).copy()
#pts.shape = pts.size / 3, 3
#scalars = scalars.T.copy()
#vectors = vectors.transpose(2, 1, 0, 3).copy()
#vectors.shape = vectors.size / 3, 3
#from tvtk.api import tvtk
#sgrid = tvtk.StructuredGrid(dimensions=(dims[1], dims[0], dims[2]))
#sgrid.points = pts
#s = random.random((dims[0]*dims[1]*dims[2]))
#sgrid.point_data.scalars = ravel(s.copy())
#sgrid.point_data.scalars.name = 'scalars'
figure, mlab = get_fig_mlab(figure=None)
source = mlab.pipeline.scalar_field(s)
data_min, data_max = s.min(), s.max()
print(data_min, data_max)
#mlab.pipeline.volume(source)
# #vmin=data_min + 0.65 * (data_max - data_min),
# #vmax=data_min + 0.9 * (data_max - data_min))
#mlab.pipeline.iso_surface(source)
mlab.pipeline.image_plane_widget(source, plane_orientation='x_axes', slice_index=0)
mlab.pipeline.image_plane_widget(source, plane_orientation='y_axes', slice_index=0)
mlab.pipeline.image_plane_widget(source, plane_orientation='z_axes', slice_index=0)
@mlab.show
@mlab.animate(delay=1000, ui=True)
def anim():
"""Animate."""
t = 1
while True:
#vmin, vmax = .1 * np.max(data[t]), .2 * np.max(data[t])
#print 'animation t = ',tax[t],', max = ',np.max(data[t])
with abilab.abiopen(self.filepaths[t]) as nc:
print("Animation step", t, "from file:", self.filepaths[t])
#nsppol, nspden, nspinor = nc.nsppol, nc.nspden, nc.nspinor
datar = nc.field.datar
# [nspden, nx, ny, nz] array
#nx, ny, nz = datar.shape[1:]
scalars = datar[0]
#data_min, data_max = scalars.min(), scalars.max(),
#mlab.pipeline.volume(source, vmin=data_min + 0.65 * (data_max - data_min),
# vmax=data_min + 0.9 * (data_max - data_min))
source.mlab_source.scalars = scalars
t = (t + 1) % self.num_files
yield
anim()