Added cylinder shape type (#321)

Very simple, self-contained PR to include a cylinder shape - started off
to switch to cylinder when headLength was zero, but it seems cleaner to
have a separate object even if it was a bit more coding.

Furthermore, there are a few improvements to the shapes functionality:

* changed the way shapes are displayed so it is much faster to add many - basically all the vertices are added at once rather than one shape at a time
* set things up so that when there are atom-based shapes with no explicit color definition the shapes will take the color of the corresponding atom.
* added a utility to make path integral visualizations.


---------

Co-authored-by: Rose K. Cersonsky <[email protected]>

docs/src/manual/input-reference.rst

@@ -106,7 +106,7 @@ contains the following fields and values:
         // at the structure level, or for individual atoms 
         "shapes": {
             <name>: {
-                "kind" : <"sphere", "ellipsoid", "arrow", "custom">,
+                "kind" : <"sphere", "ellipsoid", "cylinder", "arrow", "custom">,
                 "parameters" : {
                     "global" : { <global_parameters_dictionary> },
                     "structure" : [ <list_of_structure_parameter_dictionaries> ],
@@ -120,14 +120,6 @@ contains the following fields and values:
                     "global" : { "radius" : 0.2 }                    
                 }
             },
-            // Arrow, with the given shape parameters, and `vector` direction
-            <other_name>: {
-                "kind" : "sphere"
-                "parameters" : {
-                    "global" : { "baseRadius" : 0.2, 'headRadius': 0.3, 'headLength' : 0.4 },
-                    "atom" : [ {"vector" : [0,0,1]}, {"vector": [0,1,1]}, ... ]
-                }
-            },
             // Ellipsoid shapes, with the given `[ax, ay, az]` semi-axes
             <other_name>: {
                 "kind" : "ellipsoid"
@@ -136,6 +128,22 @@ contains the following fields and values:
                     "structure" : [ {"semiaxes": [1, 1, 2]}, ... ]
                 }
             },
+            // Cylinder, with the given radiys, and `vector` direction
+            <other_name>: {
+                "kind" : "cylinder"
+                "parameters" : {
+                    "global" : { "radius" : 0.2 },
+                    "atom" : [ {"vector" : [0,0,1]}, {"vector": [0,1,1]}, ... ]
+                }
+            },
+            // Arrow, with the given shape parameters, and `vector` direction
+            <other_name>: {
+                "kind" : "arrow"
+                "parameters" : {
+                    "global" : { "baseRadius" : 0.2, 'headRadius': 0.3, 'headLength' : 0.4 },
+                    "atom" : [ {"vector" : [0,0,1]}, {"vector": [0,1,1]}, ... ]
+                }
+            },            
             // Custom shapes. Must provide list of vertices, and the vertex
             // indices associated with simplices (the latter are autocalculated)
             // if omitted

python/chemiscope/__init__.py

@@ -1,8 +1,9 @@
 # -*- coding: utf-8 -*-
-from .input import create_input, write_input  # noqa
-from .structures import (  # noqa
+from .input import create_input, write_input  # noqa: F401
+from .structures import (  # noqa: F401
     all_atomic_environments,
     arrow_from_vector,
+    ase_merge_pi_frames,
     ase_tensors_to_ellipsoids,
     ase_vectors_to_arrows,
     center_shape,
@@ -11,7 +12,7 @@
     extract_properties,
     librascal_atomic_environments,
 )
-from .version import __version__  # noqa
+from .version import __version__  # noqa: F401
 
 try:
     # only import the chemiscope.show function if we have ipywidgets installed.

python/chemiscope/input.py

@@ -220,6 +220,13 @@ def create_input(
             "semiaxes": [float, float, float],
         }
 
+        # "kind" : "cylinder"
+        {   # "orientation" is redundant and hence ignored
+            "vector" : [float, float, float],  # orientation and shape of the cylinder
+            # the tip of the cylinder is at the end of the segment.
+            "radius" : float,
+        }
+
         # "kind" : "arrow"
         {   # "orientation" is redundant and hence ignored
             "vector" : [float, float, float],  # orientation and shape of the arrow
@@ -982,20 +989,34 @@ def _check_valid_shape(shape):
                 raise ValueError(
                     "'simplices' must be an Nx3 array values for 'custom' shape kind"
                 )
+    elif shape["kind"] == "cylinder":
+        if not isinstance(parameters["radius"], float):
+            raise TypeError(
+                "cylinder shape 'radius' must be a float, "
+                f"got {type(parameters['radius'])}"
+            )
+        vector_array = np.asarray(parameters["vector"]).astype(
+            np.float64, casting="safe", subok=False, copy=False
+        )
+
+        if not vector_array.shape == (3,):
+            raise ValueError(
+                "'vector' must be an array with 3 values for 'cylinder' shape kind"
+            )
     elif shape["kind"] == "arrow":
         if not isinstance(parameters["baseRadius"], float):
             raise TypeError(
-                "sphere shape 'baseRadius' must be a float, "
+                "arrow shape 'baseRadius' must be a float, "
                 f"got {type(parameters['baseRadius'])}"
             )
         if not isinstance(parameters["headRadius"], float):
             raise TypeError(
-                "sphere shape 'headRadius' must be a float, "
+                "arrow shape 'headRadius' must be a float, "
                 f"got {type(parameters['headRadius'])}"
             )
         if not isinstance(parameters["headLength"], float):
             raise TypeError(
-                "sphere shape 'headLength' must be a float, "
+                "arrow shape 'headLength' must be a float, "
                 f"got {type(parameters['headLength'])}"
             )
 

python/chemiscope/structures/__init__.py

@@ -9,11 +9,16 @@
     _ase_to_json,
     _ase_valid_structures,
 )
-from ._shapes import arrow_from_vector, center_shape, ellipsoid_from_tensor  # noqa
+from ._shapes import (  # noqa: F401
+    arrow_from_vector,
+    center_shape,
+    ellipsoid_from_tensor,
+)
 
-from ._ase import (  # noqa  isort: skip
-    ase_vectors_to_arrows,
+from ._ase import (  # noqa: F401
+    ase_merge_pi_frames,
     ase_tensors_to_ellipsoids,
+    ase_vectors_to_arrows,
 )
 
 

python/chemiscope/structures/_ase.py

@@ -499,3 +499,84 @@ def _get_shape_params_atom(prefix, shape_kind, dictionary, atom_i):
         )
 
     return shape
+
+
+def ase_merge_pi_frames(bead_frames):
+    """
+    Takes a list of lists of ase.Atoms objects, corresponding to
+    the trajectories of the beads of a path integral calculation.
+    Returns a consolidated trajectory in which all beads are in the
+    same frame, together with properties and shapes that help visualize
+    the path integral. The various pieces are returned as a dictionary
+    and can be passed directly to `create_input` and related functions.
+
+    :param bead_frames: list of lists of ASE Atoms objects. the outer
+       list corresponds to the bead index
+    """
+
+    nbeads, nframes, natoms = (
+        len(bead_frames),
+        len(bead_frames[0]),
+        len(bead_frames[0][0]),
+    )
+
+    # creates frames with all beads
+    full_frames = []
+    for i in range(nframes):
+        struc = bead_frames[0][i].copy()
+        for k in range(1, nbeads):
+            struc += bead_frames[k][i]
+        full_frames.append(struc)
+
+    atom_id = [
+        j
+        for i in range(nframes)  # frame
+        for k in range(nbeads)  # bead
+        for j in range(natoms)  # atom
+    ]
+    bead_id = [
+        k
+        for i in range(nframes)  # frame
+        for k in range(nbeads)  # bead
+        for j in range(natoms)  # atom
+    ]
+
+    paths_shapes = dict(
+        kind="cylinder",
+        parameters={
+            "global": {"radius": 0.2},
+            "atom": [
+                {
+                    "vector": full_frames[i]
+                    .get_distance(
+                        k * natoms + j,
+                        ((k + 1) % nbeads) * natoms + j,
+                        mic=True,
+                        vector=True,
+                    )
+                    .tolist(),
+                }
+                for i in range(nframes)  # frame
+                for k in range(nbeads)  # bead
+                for j in range(natoms)  # atom
+            ],
+        },
+    )
+    return {
+        "frames": full_frames,
+        "shapes": {"paths": paths_shapes},
+        "properties": {"atom_id": atom_id, "bead_id": bead_id},
+        "settings": {
+            "structure": [
+                {
+                    "atoms": True,
+                    "bonds": False,
+                    "shape": "paths",
+                    "unitCell": bool(full_frames[0].pbc[0]),
+                    "environments": {"activated": False},
+                    "color": {"property": "atom_id", "palette": "hsv (periodic)"},
+                }
+            ]
+        },
+        "environments": _ase_all_atomic_environments(full_frames, 4.0),
+    }

src/dataset.ts

@@ -3,7 +3,7 @@
  * @module main
  */
 
-import { Arrow, CustomShape, Ellipsoid, Sphere } from './structure/shapes';
+import { Arrow, CustomShape, Cylinder, Ellipsoid, Sphere } from './structure/shapes';
 import { ShapeParameters } from './structure/shapes';
 
 /** A dataset containing all the data to be displayed. */
@@ -403,6 +403,8 @@ function validateShape(kind: string, parameters: Record<string, unknown>): strin
         return Ellipsoid.validateParameters(parameters);
     } else if (kind === 'arrow') {
         return Arrow.validateParameters(parameters);
+    } else if (kind === 'cylinder') {
+        return Cylinder.validateParameters(parameters);
     } else if (kind === 'custom') {
         return CustomShape.validateParameters(parameters);
     }