Create mesh visualization - just for CI

structuretoolkit/__init__.py

@@ -37,6 +37,7 @@
     grainboundary,
     high_index_surface,
     sqs_structures,
+    create_mesh,
 )
 
 # Common
@@ -51,10 +52,11 @@
     get_wrapped_coordinates,
     pymatgen_to_ase,
     select_index,
+    get_cell,
 )
 
 # Visualize
-from structuretoolkit.visualize import plot3d
+from structuretoolkit.visualize import plot3d, plot_isosurface
 
 # Analyse - for backwards compatibility
 from structuretoolkit.analyse import (

structuretoolkit/build/__init__.py

@@ -8,4 +8,5 @@
 from structuretoolkit.build.surface import (
     high_index_surface,
     get_high_index_surface_info
-)
+)
+from structuretoolkit.build.mesh import create_mesh

structuretoolkit/build/mesh.py

@@ -0,0 +1,53 @@
+from __future__ import annotations
+
+import numpy as np
+import warnings
+import typing
+from structuretoolkit.common.helper import get_cell
+
+
+class MeshInputError(ValueError):
+    """ Raised when mesh input format is wrong """
+
+
+def create_mesh(
+    cell: typing.Union["ase.atoms.Atoms", np.ndarray, list, float],
+    n_mesh: typing.Union[int, list[int, int, int]] = 10,
+    density: typing.Optional[float] = None,
+    endpoint: bool = False
+):
+    """
+    Create a mesh based on a structure cell
+
+    Args:
+        cell (ase.atoms.Atoms|np.ndarray|list|float): ASE Atoms or cell
+        n_mesh (int): Number of grid points in each direction. If one number
+            is given, it will be repeated in every direction (i.e. n_mesh = 3
+            is the same as n_mesh = [3, 3, 3])
+        density (float): Density of grid points. Ignored when n_mesh is not
+            None
+        endpoint (bool): Whether both the edges get separate points or not.
+            cf. endpoint in numpy.linspace
+
+    Returns:
+        (3, n, n, n)-array: mesh
+    """
+    cell = get_cell(cell)
+    if n_mesh is None:
+        if density is None:
+            raise MeshInputError("either n_mesh or density must be specified")
+        n_mesh = np.rint(np.linalg.norm(cell, axis=-1) / density).astype(int)
+    elif density is not None:
+        raise MeshInputError(
+            "You cannot set n_mesh at density at the same time. Set one of"
+            " them to None"
+        )
+    n_mesh = np.atleast_1d(n_mesh).astype(int)
+    if len(n_mesh) == 1:
+        n_mesh = np.repeat(n_mesh, 3)
+    elif len(n_mesh) != 3:
+        raise MeshInputError("n_mesh must be a 3-dim vector")
+    linspace = [np.linspace(0, 1, nn, endpoint=endpoint) for nn in n_mesh]
+    x_mesh = np.meshgrid(*linspace, indexing='ij')
+    return np.einsum("ixyz,ij->jxyz", x_mesh, cell)
+

structuretoolkit/common/__init__.py

@@ -6,6 +6,7 @@
     get_vertical_length,
     get_wrapped_coordinates,
     select_index,
+    get_cell,
 )
 from structuretoolkit.common.pymatgen import (
     ase_to_pymatgen,

structuretoolkit/common/helper.py

@@ -2,7 +2,7 @@
 from ase.atoms import Atoms
 from ase.data import atomic_numbers
 from scipy.sparse import coo_matrix
-from typing import Optional
+from typing import Optional, Union
 
 
 def get_extended_positions(
@@ -226,3 +226,38 @@ def apply_strain(
     structure_copy.set_cell(cell, scale_atoms=True)
     if return_box:
         return structure_copy
+
+
+def get_cell(cell: Union[Atoms, list, np.ndarray, float]):
+    """
+    Get cell of an ase structure, or convert a float or a (3,)-array into a
+    orthogonal cell.
+
+    Args:
+        cell (Atoms|ndarray|list|float|tuple): Cell
+
+    Returns:
+        (3, 3)-array: Cell
+    """
+    if isinstance(cell, Atoms):
+        return cell.cell
+    # Convert float into (3,)-array. No effect if it is (3,3)-array or
+    # (3,)-array. Raises error if the shape is not correct
+    try:
+        cell = cell * np.ones(3)
+    except ValueError:
+        raise ValueError(
+            "cell must be a float, (3,)-ndarray/list/tuple or"
+            " (3,3)-ndarray/list/tuple"
+        )
+
+    if np.shape(cell) == (3, 3):
+        return cell
+    # Convert (3,)-array into (3,3)-array. Raises error if the shape is wrong
+    try:
+        return cell * np.eye(3)
+    except ValueError:
+        raise ValueError(
+            "cell must be a float, (3,)-ndarray/list/tuple or"
+            " (3,3)-ndarray/list/tuple"
+        )

structuretoolkit/visualize.py

@@ -2,12 +2,16 @@
 # 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
 
 __author__ = "Joerg Neugebauer, Sudarsan Surendralal"
 __copyright__ = (
@@ -166,7 +170,7 @@ def _get_box_skeleton(cell: np.ndarray):
 
 
 def _draw_box_plotly(fig, structure, px, go):
-    cell = structure.cell
+    cell = get_cell(structure)
     data = fig.data
     for lines in _get_box_skeleton(cell):
         fig = px.line_3d(**{xx: vv for xx, vv in zip(["x", "y", "z"], lines.T)})
@@ -805,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

tests/test_helpers.py

@@ -0,0 +1,28 @@
+# coding: utf-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.
+
+import unittest
+import numpy as np
+from ase.build import bulk
+import structuretoolkit as stk
+
+
+class TestHelpers(unittest.TestCase):
+    def test_get_cell(self):
+        self.assertEqual((3 * np.eye(3)).tolist(), stk.get_cell(3).tolist())
+        self.assertEqual(
+            ([1, 2, 3] * np.eye(3)).tolist(), stk.get_cell([1, 2, 3]).tolist()
+        )
+        atoms = bulk("Fe")
+        self.assertEqual(
+            atoms.cell.tolist(), stk.get_cell(atoms).tolist()
+        )
+        with self.assertRaises(ValueError):
+            stk.get_cell(np.arange(4))
+        with self.assertRaises(ValueError):
+            stk.get_cell(np.ones((4, 3)))
+
+
+if __name__ == "__main__":
+    unittest.main()

tests/test_mesh.py

@@ -0,0 +1,31 @@
+# coding: utf-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.
+
+import unittest
+from ase.build import bulk
+import structuretoolkit as stk
+
+
+class TestMesh(unittest.TestCase):
+    def test_mesh(self):
+        structure = bulk("Al", cubic=True)
+        self.assertEqual(stk.create_mesh(structure, n_mesh=4).shape, (3, 4, 4, 4))
+        with self.assertRaises(stk.build.mesh.MeshInputError):
+            stk.create_mesh(structure, n_mesh=None, density=None)
+        with self.assertRaises(stk.build.mesh.MeshInputError):
+            stk.create_mesh(
+                structure, n_mesh=10, density=structure.cell[0, 0] / 4
+            )
+        self.assertEqual(
+            stk.create_mesh(
+                structure, n_mesh=None, density=structure.cell[0, 0] / 4
+            ).shape,
+            (3, 4, 4, 4),
+        )
+        with self.assertRaises(stk.build.mesh.MeshInputError):
+            _ = stk.create_mesh(structure, n_mesh=[1, 2, 3, 4])
+
+
+if __name__ == "__main__":
+    unittest.main()