diff --git a/structuretoolkit/__init__.py b/structuretoolkit/__init__.py index d50c50720..4829312f5 100644 --- a/structuretoolkit/__init__.py +++ b/structuretoolkit/__init__.py @@ -94,7 +94,7 @@ ) # Visualize -from structuretoolkit.visualize import plot3d +from structuretoolkit.visualize import plot3d, plot_isosurface from . import _version diff --git a/structuretoolkit/visualize.py b/structuretoolkit/visualize.py index 9a77c70c1..24cfb8b31 100644 --- a/structuretoolkit/visualize.py +++ b/structuretoolkit/visualize.py @@ -2,6 +2,8 @@ # 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 typing import Optional @@ -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: 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: float = 600.0, + 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