Create mesh visualization

structuretoolkit/__init__.py

@@ -56,7 +56,7 @@
 )
 
 # Visualize
-from structuretoolkit.visualize import plot3d
+from structuretoolkit.visualize import plot3d, plot_isosurface
 
 # Analyse - for backwards compatibility
 from structuretoolkit.analyse import (

structuretoolkit/visualize.py

@@ -2,12 +2,14 @@
 # Copyright (c) Max-Planck-Institut für Eisenforschung GmbH - Computational Materials Design (CM) Department
 # Distributed under the terms of "New BSD License", see the LICENSE file.
 
+from __future__ import annotations
 import warnings
 
 from ase.atoms import Atoms
 import numpy as np
-from typing import Optional
+from typing import Optional, Union
 from scipy.interpolate import interp1d
+from structuretoolkit.common.helper import get_cell
 
 from structuretoolkit.common.helper import get_cell
 
@@ -807,3 +809,78 @@ def _get_flattened_orientation(
     flattened_orientation[:3, :3] = _get_orientation(view_plane)
 
     return (distance_from_camera * flattened_orientation).ravel().tolist()
+
+
+def plot_isosurface(
+    mesh,
+    value,
+    cell: Optional[Union[Atoms, list, np.ndarray, float]] = None,
+    structure_plot: Optional["plotly.graph_objs._figure.Figure"] = None,
+    isomin: Optional[float] = None,
+    isomax: Optional[float] = None,
+    surface_fill: Optional[float] = None,
+    opacity: Optional[float] = None,
+    surface_count: Optional[int] = 5,
+    colorbar_nticks: Optional[int] = None,
+    caps: Optional[dict] = dict(x_show=False, y_show=False, z_show=False),
+    colorscale: Optional[str] = None,
+    height: Optional[float] = 600,
+    camera: Optional[str] = "orthographic",
+    **kwargs,
+):
+    """
+    Make a mesh plot
+
+    Args:
+        mesh (numpy.ndarray): Mesh grid. Must have a shape of (3, nx, ny, nz).
+            It can be generated from structuretoolkit.create_mesh
+        value: (numpy.ndarray): Value to plot. Must have a shape of (nx, ny, nz)
+        cell (Atoms|ndarray|list|float|tuple): Cell, ignored if
+            `structure_plot` is given
+        structure_plot (plotly.graph_objs._figure.Figure): Plot of the
+            structure to overlay. You should basically always use
+            structuretoolkit.plot3d(structure, mode="plotly")
+        isomin(float): Min color value
+        isomax(float): Max color value
+        surface_fill(float): Polygonal filling of the surface to choose between
+            0 and 1
+        opacity(float): Opacity
+        surface_count(int): Number of isosurfaces, 5 by default, which means
+            only min and max
+        colorbar_nticks(int): Colorbar ticks correspond to isosurface values
+        caps(dict): Whether to set cap on sides or not. Default is False. You
+            can set: caps=dict(x_show=True, y_show=True, z_show=True)
+        colorscale(str): Colorscale ("turbo", "gnbu" etc.)
+        height(float): Height of the figure. 600px by default
+        camera(str): Camera perspective to choose from "orthographic" and
+            "perspective". Default is "orthographic"
+    """
+    try:
+        import plotly.graph_objects as go
+    except ModuleNotFoundError:
+        raise ModuleNotFoundError("plotly not installed")
+    x_mesh = np.reshape(mesh, (3, -1))
+    data = go.Isosurface(
+        x=x_mesh[0],
+        y=x_mesh[1],
+        z=x_mesh[2],
+        value=np.array(value).flatten(),
+        isomin=isomin,
+        isomax=isomax,
+        surface_fill=surface_fill,
+        opacity=opacity,
+        surface_count=surface_count,
+        colorbar_nticks=colorbar_nticks,
+        caps=caps,
+        colorscale=colorscale,
+        **kwargs,
+    )
+    fig = go.Figure(data=data)
+    if structure_plot is not None:
+        fig = go.Figure(data=fig.data + structure_plot.data)
+    elif cell is not None:
+        fig = _draw_box_plotly(fig, cell)
+    fig.update_scenes(aspectmode="data")
+    fig.layout.scene.camera.projection.type = camera
+    fig.update_layout(autosize=True, height=height)
+    return fig