Merge pull request #19 from PytLab/feature/pyvista-migration

Migrate 3D visualization from mayavi to PyVista

Author: PytLab

.gitignore

@@ -82,3 +82,6 @@ Thumbs.db
 # Jupyter notebook
 .ipynb_checkpoints/
 
+# Claude Code
+.claude/
+

README.rst

@@ -46,20 +46,13 @@ Installation
 
     python setup.py install
 
-If you want to use **mayavi** to visualize VASP data, it is recommened to install `Canopy environment <https://store.enthought.com/downloads/#default>`_ on your device instead of installing it manually.
+3D visualization is done via **PyVista**, which is installed automatically as a dependency::
 
-After installing canopy, you can set corresponding aliases, for example:
-
-.. code-block:: shell
-
-    alias canopy='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/python'
-    alias canopy-pip='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/pip'
-    alias canopy-ipython='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/ipython'
-    alias canopy-jupyter='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/jupyter'
+    pip install vaspy
 
-Then you can install VASPy to canopy::
+Or install PyVista separately::
 
-    canopy-pip install vaspy
+    pip install pyvista
 
 Examples
 --------
@@ -98,7 +91,7 @@ Visualize ELFCAR
     >>> from vaspy.electro import ElfCar
     >>> a = ElfCar() 
     >>> a.plot_contour()   # Plot coutour
-    >>> a.plot_mcontour()  # Plot coutour using mlab(with Mayavi installed)
+    >>> a.plot_mcontour()  # Plot coutour using PyVista
     >>> a.plot_contour3d() # Plot 3D coutour
     >>> a.plot_field()     # Plot scalar field
 

requirements.txt

@@ -1,4 +1,4 @@
 matplotlib>=1.5.2
 numpy>=1.11.1
 scipy>=0.18.0
-
+pyvista>=0.42.0

scripts/md_viz/map_mayavi.py

@@ -4,7 +4,7 @@
 
 import numpy as np
 from scipy.interpolate import interp2d
-from mayavi import mlab
+import pyvista as pv
 
 from vaspy.iter import XdatCar, AniFile
 from traj import get_trajectories
@@ -43,11 +43,15 @@ def locate(x, y, position):
     newy = np.linspace(left, right, interp_resolution)
     newz = interp_func(newx, newy)
 
-    newy, newx = np.meshgrid(newx, newy)
-
-    face = mlab.surf(newx, newy, newz, warp_scale=40)
-    mlab.axes(xlabel="x", ylabel="y", zlabel="z")
-    mlab.outline(face)
-
-    mlab.show()
-
+    # PyVista surface plot
+    nx, ny = len(newx), len(newy)
+    X, Y = np.meshgrid(newx, newy, indexing='ij')
+    Z = newz.T
+    surface = pv.StructuredGrid(X[:, :, None], Y[:, :, None], Z[:, :, None])
+    surface.point_data['scalars'] = Z[:, :, None].flatten(order='F')
+
+    pl = pv.Plotter()
+    pl.add_mesh(surface, scalars='scalars', cmap='viridis',
+                show_scalar_bar=True)
+    pl.add_axes(xlabel='x', ylabel='y', zlabel='z')
+    pl.show()

scripts/plot_test.py

@@ -1,7 +1,7 @@
 from plotter import *
 import numpy as np
 from scipy.interpolate import interp2d
-from mayavi import mlab
+import pyvista as pv
 import mpl_toolkits.mplot3d
 import matplotlib.pyplot as plt
 
@@ -17,20 +17,14 @@
 newy = np.linspace(0, np.max(y), 1000)
 newz = interpfunc(newx, newy)
 
-#extent = [np.min(newx), np.max(newx), np.min(newy), np.max(newy)]
-#plt.contourf(newx.reshape(-1), newy.reshape(-1), newz, 20, extent=extent)
-#plt.colorbar()
-#
-##3d plot
-#fig3d = plt.figure()
-#ax3d = fig3d.add_subplot(111, projection='3d')
-#ax3d.plot_surface(newx, newy, newz, cmap=plt.cm.RdBu_r)
-#
-#plt.show()
+# PyVista surface plot
+nx, ny = len(newx), len(newy)
+X, Y = np.meshgrid(newx, newy, indexing='ij')
+Z = newz.T
+surface = pv.StructuredGrid(X[:, :, None], Y[:, :, None], Z[:, :, None])
+surface.point_data['scalars'] = Z[:, :, None].flatten(order='F')
 
-#mlab
-face = mlab.surf(newx, newy, newz, warp_scale=2)
-mlab.axes(xlabel='x', ylabel='y', zlabel='z')
-mlab.outline(face)
-
-mlab.show()
+pl = pv.Plotter()
+pl.add_mesh(surface, scalars='scalars', cmap='viridis', show_scalar_bar=True)
+pl.add_axes(xlabel='x', ylabel='y', zlabel='z')
+pl.show()

setup.py

@@ -58,27 +58,21 @@
 
     python setup.py install
 
-If you want to use **mayavi** to visualize VASP data, it is recommened to install `Canopy environment <https://store.enthought.com/downloads/#default>`_ on your device instead of installing it manually.
+3D visualization is done via **PyVista**, which is installed automatically as a dependency::
 
-After installing canopy, you can set corresponding aliases, for example:
-
-.. code-block:: shell
-
-    alias canopy='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/python'
-    alias canopy-pip='/Users/zjshao/Library/Enthought/Canopy/edm/envs/User/bin/pip'
-    alias canopy-ipython='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/ipython'
-    alias canopy-jupyter='/Users/<yourname>/Library/Enthought/Canopy/edm/envs/User/bin/jupyter'
+    pip install vaspy
 
-Then you can install VASPy to canopy::
+Or install PyVista separately::
 
-    canopy-pip install vaspy
+    pip install pyvista
 
 """
 
 install_requires = [
     'numpy>=1.11.1',
     'matplotlib>=1.5.2',
     'scipy>=0.18.0',
+    'pyvista>=0.42.0',
 ]
 
 license = 'LICENSE'