Utility function to generate arrows and ellipsoids from physical data

This allows to visualize vectorial and tensorial properties for each atom.

docs/src/tutorial/input.rst

@@ -57,6 +57,10 @@ generate an output in chemiscope format.
 
 .. autofunction:: chemiscope.extract_lammps_shapes_from_ase
 
+.. autofunction:: chemiscope.ellipsoid_from_tensor
+
+.. autofunction:: chemiscope.arrow_from_vector
+
 .. _ase: https://wiki.fysik.dtu.dk/ase/index.html
 .. _ASAP: https://github.com/BingqingCheng/ASAP
 

python/chemiscope/__init__.py

@@ -2,8 +2,13 @@
 from .input import create_input, write_input  # noqa
 from .structures import (  # noqa
     all_atomic_environments,
+    arrow_from_vector,
+    center_shape,
     composition_properties,
+    ellipsoid_from_tensor,
     extract_lammps_shapes_from_ase,
+    extract_tensors_from_ase,
+    extract_vectors_from_ase,
     extract_properties,
     librascal_atomic_environments,
 )

python/chemiscope/jupyter.py

@@ -30,9 +30,10 @@ class ChemiscopeWidgetBase(ipywidgets.DOMWidget, ipywidgets.ValueWidget):
 
     def __init__(self, data, has_metadata):
         super().__init__()
+
         self.value = json.dumps(data)
         self.has_metadata = has_metadata
-        self.settings = {}
+        self.settings = data["settings"]
         self._settings_sync = True
 
     def save(self, path):
@@ -85,14 +86,15 @@ def show(
     properties=None,
     meta=None,
     environments=None,
+    shapes=None,
     settings=None,
     mode="default",
 ):
     """
     Show the dataset defined by the given ``frames`` and ``properties``
-    (optionally ``meta`` and ``environments`` as well) using a embedded chemiscope
-    visualizer inside a Jupyter notebook. These parameters have the same meaning
-    as in the :py:func:`chemiscope.create_input` function.
+    (optionally ``meta``, ``environments`` and ``shapes`` as well) using an embedded
+    chemiscope visualizer inside a Jupyter notebook. These parameters have the same
+    meaning as in the :py:func:`chemiscope.create_input` function.
 
     The ``mode`` keyword also allows overriding the default two-panels
     visualization to show only a structure panel (``mode = "structure"``) or the
@@ -167,6 +169,7 @@ def show(
         properties=properties,
         meta=meta,
         environments=environments,
+        shapes=shapes,
         settings=settings,
     )
 

python/chemiscope/structures/__init__.py

@@ -10,7 +10,17 @@
     _ase_valid_structures,
 )
 
-from ._ase import extract_lammps_shapes_from_ase  # noqa  isort: skip
+from ._ase import (  # noqa  isort: skip
+    extract_lammps_shapes_from_ase,
+    extract_tensors_from_ase,
+    extract_vectors_from_ase,
+)
+
+from ._shapes import (  # noqa
+    arrow_from_vector,
+    center_shape,
+    ellipsoid_from_tensor,
+)
 
 
 def _guess_adapter(frames):

python/chemiscope/structures/_ase.py

@@ -1,6 +1,7 @@
 import warnings
 from collections import Counter
 
+from ._shapes import ellipsoid_from_tensor, arrow_from_vector
 import numpy as np
 
 try:
@@ -304,6 +305,66 @@ def _is_convertible_to_property(value):
                 return False
 
 
+def extract_vectors_from_ase(frames, key="forces", **kwargs):
+    """
+    Extract a vectorial atom property from a list of ase.Atoms
+    objects, and returns a list of arrow shapes. Besides the specific
+    parameters it also accepts the same parameters as
+    `array_from_vector`, which are used to define the style of the
+    arrows.
+
+    :param frames: list of ASE Atoms objects
+    :param key: name of the ASE atom property. Should contain
+       three components corresponding to x,y,z
+    """
+
+    vectors = []
+
+    for f in frames:
+        if key not in f.arrays:
+            raise IndexError(f"Key {key} not found in `Atoms.arrays`")
+        values = f.arrays[key]
+        if len(values.shape) != 2 or values.shape[1] != 3:
+            raise ValueError(
+                f"Property array {key} has not the shape of a list of 3-vectors"
+            )
+
+        # makes a list of arrows to visualize the property
+        vectors.append([arrow_from_vector(v, **kwargs) for v in values])
+
+    return vectors
+
+
+def extract_tensors_from_ase(frames, key="tensor", **kwargs):
+    """
+    Extract a 3-tensor atom property from a list of ase.Atoms
+    objects, and returns a list of arrow shapes. Besides the specific
+    parameters it also accepts the same parameters as
+    `ellipsoid_from_tensor`, which are used to draw the shapes
+
+    :param frames: list of ASE Atoms objects
+    :param key: name of the ASE atom property. Should contain
+       nine components corresponding to xx,xy,xz,yx,yy,yz,zx,zy,zz or
+       six components corresponding to xx,yy,zz,xy,xz,yz
+    """
+
+    tensors = []
+
+    for f in frames:
+        if key not in f.arrays:
+            raise IndexError(f"Key {key} not found in `Atoms.arrays`")
+        values = f.arrays[key]
+        if len(values.shape) != 2 or (values.shape[1] != 6 and values.shape[1] != 9):
+            raise ValueError(
+                f"Property array {key} has not the shape of a list of 6 or 9-vectors"
+            )
+
+        # makes a list of arrows to visualize the property
+        tensors.append([ellipsoid_from_tensor(v, **kwargs) for v in values])
+
+    return tensors
+
+
 def extract_lammps_shapes_from_ase(frames, key="shape"):
     """
     Extract shapes from a LAMMPS data file read by ASE.

python/chemiscope/structures/_shapes.py

@@ -0,0 +1,167 @@
+import numpy as np
+
+
+def center_shape(shape):
+    """
+    Takes a dictionary that describes a custom shape, and centers it, subtracting the
+    mean of the vertex positions. Returns a shallow copy, with a new list for the
+    vertices.
+
+    :param shape: A dictionary describing a custom shape
+    """
+
+    if shape["kind"] != "custom":
+        raise ValueError("Only `custom` shapes can be centered")
+
+    new_shape = {k: shape[k] for k in shape}
+    points = np.array(shape["vertices"])
+    points -= points.mean(axis=0)
+    new_shape["vertices"] = points.tolist()
+
+    return new_shape
+
+
+def _oriented_circle(radius, vec, n_points=20):
+    """
+    Generates a circle in 3D centered at the origin ands oriented according to the
+    given vector
+    """
+    # makes sure the normal is unit
+    nvec = vec / np.linalg.norm(vec)
+
+    # Generate an arbitrary vector not collinear with n
+    if nvec[0] or nvec[1]:
+        vx = np.array([0, 0, 1])
+    else:
+        vx = np.array([0, 1, 0])
+
+    # generate orthogonal vectors in the plane defined by nvec
+    u = vx - np.dot(vx, nvec) * nvec
+    u = u / np.linalg.norm(u)
+    v = np.cross(nvec, u)
+
+    # generate n_points in the plane defined by nvec, centered at vec
+    angles = np.linspace(0, 2 * np.pi, n_points, endpoint=False)
+    circle_points = radius * np.outer(np.cos(angles), u) + np.outer(np.sin(angles), v)
+
+    return circle_points
+
+
+def arrow_from_vector(
+    vec, scale=1.0, radius=0.1, head_radius_scale=1.75, head_length_scale=4, n_points=16
+):
+    """
+    Draws an arrow from the origin to the specified 3D position. Returns a custom shape
+    in the form required by the chemiscope input.
+
+    :param scale: conversion from the units of the vector to the units of the atomic
+        positions (usually Å)
+    :param radius: radius of the stem of the arrow (same units as the atomic positions,
+        typically Å)
+    :param head_radius_scale: radius of the arrow tip, relative to the stem radius
+    :param head_length_scale: length of the arrow tip, relative to the stem radius
+    :param n_points: resolution of the discretization of the shape
+    """
+
+    tip = np.asarray(vec, dtype=float) * scale
+    head_tip = tip + tip / np.linalg.norm(tip) * radius * head_length_scale
+    circle = _oriented_circle(1.0, tip, n_points)
+    base_circle = radius * circle
+    stem_circle = radius * circle + tip
+    head_circle = radius * head_radius_scale * circle + tip
+
+    arrow = dict(
+        kind="custom",
+        vertices=[head_tip.tolist()]
+        + head_circle.tolist()
+        + stem_circle.tolist()
+        + base_circle.tolist()
+        + [[0, 0, 0]],
+        simplices=(
+            [[0, i, i + 1] for i in range(1, n_points)]
+            + [[0, n_points, 1]]
+            + [[i, i + 1, i + 1 + n_points] for i in range(1, n_points)]
+            + [[n_points, 1, n_points + 1]]
+            + [[i + n_points, i, i + 1 + n_points] for i in range(1, n_points)]
+            + [[2 * n_points, n_points, n_points + 1]]
+            + [[i, i + 1, i + 1 + n_points] for i in range(1 + n_points, 2 * n_points)]
+            + [[2 * n_points, 1 + n_points, 2 * n_points + 1]]
+            + [
+                [i + n_points, i, i + 1 + n_points]
+                for i in range(1 + n_points, 2 * n_points)
+            ]
+            + [[3 * n_points, 2 * n_points, 2 * n_points + 1]]
+            + [
+                [3 * n_points + 1, i, i + 1]
+                for i in range(1 + 2 * n_points, 3 * n_points)
+            ]
+            + [[3 * n_points + 1, 3 * n_points, 2 * n_points + 1]]
+        ),
+    )
+    return arrow
+
+
+def ellipsoid_from_tensor(tensor, scale=1.0, force_positive=False):
+    """
+    Returns an ellipsoid (semiaxes + quaternion) representation of a positive definite
+    tensor (e.g. a polarizability), in the form required by the chemiscope input.
+
+    :param tensor: a positive-definite tensor (3x3 or a 6-array [xx,yy,zz,xy,xz,yz])
+    :param scale: conversion from the units of the tensor to the units of the atomic
+        positions (usually Å)
+    :param force_positive: takes the absolute value of eigenvalues, to handle
+        non-positive tensors
+    """
+
+    tensor = np.asarray(tensor) * (scale**2)
+    if len(tensor.flatten()) == 9:
+        matrix = tensor.reshape((3, 3))
+    elif len(tensor) == 6:
+        matrix = np.array(
+            [
+                [tensor[0], tensor[3], tensor[4]],
+                [tensor[3], tensor[1], tensor[5]],
+                [tensor[4], tensor[5], tensor[2]],
+            ]
+        )
+    # Compute the eigenvalues and eigenvectors
+    eigenvalues, eigenvectors = np.linalg.eigh(matrix)
+
+    if force_positive:
+        eigenvalues = np.abs(eigenvalues)
+
+    # The lengths of the principal axes are the square roots of the eigenvalues
+    old_settings = np.seterr(invalid="raise")
+    try:
+        ax = np.sqrt(eigenvalues[0])
+        ay = np.sqrt(eigenvalues[1])
+        az = np.sqrt(eigenvalues[2])
+    except FloatingPointError as e:
+        raise ValueError(
+            f"Non-positive definite tensor found with eigenvalues {eigenvalues}.\n"
+            "If this is acceptable, set `force_positive=True` to take the "
+            "absolute values.",
+        ) from e
+    np.seterr(**old_settings)
+
+    # makes sure the rotation is proper
+    if np.linalg.det(eigenvectors) < 0:
+        eigenvectors[:, 0] *= -1
+
+    # Form the rotation matrix from eigenvectors
+    rotation = eigenvectors.T
+
+    # converts to quaternion
+    try:
+        from scipy.spatial.transform import Rotation
+    except ImportError as e:
+        raise RuntimeError(
+            "scipy is required to construct ellipsoids from tensors"
+        ) from e
+    quaternion = Rotation.from_matrix(rotation).as_quat()
+
+    return dict(
+        kind="ellipsoid",
+        semiaxes=[ax, ay, az],
+        orientation=list(quaternion),
+    )

python/jupyter/src/widget.css

@@ -85,6 +85,6 @@ html {
 
 /* hide the slider when showing the map only (no need for navigation)
   uses part() to access styling within the shadow DOM */
-.chemiscope-meta-and-map .chemiscope-info *::part(chsp-slider)  {
+.chemiscope-meta-and-map .chemiscope-info *::part(chsp-slider) {
     display: none;
 }

src/structure/shapes.ts

@@ -75,17 +75,6 @@ function rotateAndPlace(vertex: XYZ, quaternion: Quaternion, position: XYZ): XYZ
     );
 }
 
-function findCenter(a: XYZ[]): XYZ {
-    let center: XYZ = { x: 0, y: 0, z: 0 };
-    for (const v of a) {
-        center = addXYZ(center, v);
-    }
-    center.x /= a.length;
-    center.y /= a.length;
-    center.z /= a.length;
-    return center;
-}
-
 function determineNormals(vertices: XYZ[], simplices: [number, number, number][]): XYZ[] {
     const vertexNormals: XYZ[] = [];
     const nFaces: number[] = [];
@@ -407,10 +396,8 @@ export class CustomShape extends Shape {
         const indices = [];
         const vertices = [];
 
-        const center = findCenter(this.vertices);
-
         for (const v of this.vertices) {
-            vertices.push(rotateAndPlace(subXYZ(v, center), this.quaternion, this.position));
+            vertices.push(rotateAndPlace(v, this.quaternion, this.position));
         }
 
         for (const s of this.simplices) {