# Plotting Wannier functions ----
import numpy as np
from EPWpy.error_handling import error_handler
from EPWpy.default.default_dicts import *
try:
from mayavi import mlab
from mayavi.modules.iso_surface import IsoSurface
from tvtk.api import tvtk
except ModuleNotFoundError:
error_handler.error_mayavi()
#print('The mayavi package not found\nTo visualize polarons in EPWpy, install mayavi')
#print('Perform')
[docs]
def plot_isosurface_from_cube_file(Data):
"""
Take data read from .cube file,
and plot two isosurfaces: one positive at 10% of max value, and one negative at the same value but negative.
Plot atomic positions as well.
"""
# Take the data
scalar_data = Data['scalar_data']
grid_shape = Data['grid_shape']
spacing = Data['spacing']
origin = Data['origin']
atomic_positions = Data['atomic_positions']
axis_x = Data['axis_x']
axis_y = Data['axis_y']
axis_z = Data['axis_z']
connections = Data['connections']
if ('verbosity' in Data.keys()):
verbosity = Data['verbosity']
# Initializing mayavi with a backend
if ('in_notebook' in Data.keys()):
if('backend' in Data.keys()):
mlab.init_notebook(backend = Data['backend'])
else:
mlab.init_notebook()
# Initialize figure
mlab.figure(size=(800, 600))
# Extract grid dimensions and spacings
n_points_x, n_points_y, n_points_z = grid_shape
spacing_x, spacing_y, spacing_z = spacing
# Create a 3D grid for the coordinates
x_indices = np.arange(n_points_x)
y_indices = np.arange(n_points_y)
z_indices = np.arange(n_points_z)
# Compute the grid points based on axis vectors and spacings
X = origin[0] + x_indices[:, None, None] * axis_x[0] + y_indices[None, :, None] * axis_y[0] + z_indices[None, None, :] * axis_z[0]
Y = origin[1] + x_indices[:, None, None] * axis_x[1] + y_indices[None, :, None] * axis_y[1] + z_indices[None, None, :] * axis_z[1]
Z = origin[2] + x_indices[:, None, None] * axis_x[2] + y_indices[None, :, None] * axis_y[2] + z_indices[None, None, :] * axis_z[2]
# Flatten the coordinate arrays
coords = np.vstack([X.flatten(), Y.flatten(), Z.flatten()]).T
# Create a StructuredGrid in tvtk
structured_grid = tvtk.StructuredGrid(dimensions=(n_points_x, n_points_y, n_points_z))
structured_grid.points = coords # Assign the flattened coordinates to the grid
# Assign the scalar field to the grid as point data
structured_grid.point_data.scalars = scalar_data.flatten()
structured_grid.point_data.scalars.name = 'ScalarField'
# Find the maximum value in the scalar field data
max_value = np.max(scalar_data)
# Set the isosurface values for positive and negative contours
iso_value = 0.1 * max_value
print(f"Positive isosurface contour set to: {iso_value} (10% of max value: {max_value})")
# Create an iso_surface pipeline from the structured grid (positive)
iso_pos = mlab.pipeline.iso_surface(structured_grid, contours=[iso_value], colormap='Oranges', opacity=0.8)
# Create an iso_surface pipeline from the structured grid (positive)
iso_neg = mlab.pipeline.iso_surface(structured_grid, contours=[-iso_value], colormap='Blues', opacity=0.8)
# Plot atomic positions with jmol colors and atomic radius-based size
length = len(np.array(atomic_positions))
X = []
Y = []
Z = []
for i,atom in enumerate(atomic_positions):
atom_number, charge, x, y, z = atom
# Get jmol color for the atomic number
color = jmol_colors.get(atom_number, (0.5, 0.5, 0.5)) # Default to gray if not in map
# Get atomic radius (if available) and adjust the size accordingly
radius = atomic_radii.get(atom_number, 50*pm2bohr) # Default to 50 pm if not in the map
size = (radius / 2) # Scale size for visibility
# Print atomic info for the user
if (verbosity > 2):
print(f"Atom: {atom_number}, Position: ({x}, {y}, {z}), Radius: {radius} bohr")
X.append(x)
Y.append(y)
Z.append(z)
src = mlab.points3d(x, y, z, color=color, scale_factor=size)
src = mlab.points3d(X, Y, Z, color=(0,0,0), scale_factor=0.0)#size)
src.mlab_source.dataset.lines = np.array(connections)
tube = mlab.pipeline.tube(src, tube_radius=0.15)
tube.filter.radius_factor = 1.
mlab.pipeline.surface(tube, color=(0.8, 0.8, 0))
if ('in_notebook' in Data.keys()):
return(src)#mlab.points3d(x, y, z, color=color, scale_factor=size))
else:
mlab.show()
"""
for i,atom in enumerate(atomic_positions):
print(atom)
atom_number, charge, x, y, z = atom
# Get jmol color for the atomic number
color = jmol_colors.get(atom_number, (0.5, 0.5, 0.5)) # Default to gray if not in map
# Get atomic radius (if available) and adjust the size accordingly
radius = atomic_radii.get(atom_number, 50*pm2bohr) # Default to 50 pm if not in the map
size = radius / 2 # Scale size for visibility
# Print atomic info for the user
print(f"Atom: {atom_number}, Position: ({x}, {y}, {z}), Radius: {radius} bohr")
if (i == length - 1):
if ('in_notebook' in Data.keys()):
src = mlab.points3d(x, y, z, color=color, scale_factor=size)
src.mlab_source.dataset.lines = np.array(connections)
tube = mlab.pipeline.tube(src, tube_radius=0.15)
tube.filter.radius_factor = 1.
mlab.pipeline.surface(tube, color=(0.8, 0.8, 0))
return(src)#mlab.points3d(x, y, z, color=color, scale_factor=size))
else:
mlab.points3d(x, y, z, color=color, scale_factor=size)
# Show the plot
if ('in_notebook' in Data.keys()):
pass
else:
mlab.show()
"""
