From a8d8b0f0a632affe2cf262faffa8f0216f2ea2fc Mon Sep 17 00:00:00 2001 From: Michele Ceriotti Date: Mon, 11 Dec 2023 15:19:53 -0800 Subject: [PATCH] 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 <47536110+rosecers@users.noreply.github.com> --- docs/src/manual/input-reference.rst | 26 ++- python/chemiscope/__init__.py | 7 +- python/chemiscope/input.py | 27 ++- python/chemiscope/structures/__init__.py | 11 +- python/chemiscope/structures/_ase.py | 81 +++++++++ src/dataset.ts | 4 +- src/structure/shapes.ts | 213 +++++++++++++++++++++-- src/structure/viewer.ts | 89 ++++++++-- 8 files changed, 408 insertions(+), 50 deletions(-) diff --git a/docs/src/manual/input-reference.rst b/docs/src/manual/input-reference.rst index dcf722602..e6778dada 100644 --- a/docs/src/manual/input-reference.rst +++ b/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": { : { - "kind" : <"sphere", "ellipsoid", "arrow", "custom">, + "kind" : <"sphere", "ellipsoid", "cylinder", "arrow", "custom">, "parameters" : { "global" : { }, "structure" : [ ], @@ -120,14 +120,6 @@ contains the following fields and values: "global" : { "radius" : 0.2 } } }, - // Arrow, with the given shape parameters, and `vector` direction - : { - "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 : { "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 + : { + "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 + : { + "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 diff --git a/python/chemiscope/__init__.py b/python/chemiscope/__init__.py index 78411b32d..610f3fa85 100644 --- a/python/chemiscope/__init__.py +++ b/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. diff --git a/python/chemiscope/input.py b/python/chemiscope/input.py index 67bba1cfd..217822edb 100644 --- a/python/chemiscope/input.py +++ b/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'])}" ) diff --git a/python/chemiscope/structures/__init__.py b/python/chemiscope/structures/__init__.py index 49d1abb67..3f6fa9991 100644 --- a/python/chemiscope/structures/__init__.py +++ b/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, ) diff --git a/python/chemiscope/structures/_ase.py b/python/chemiscope/structures/_ase.py index 548b3dcc8..1c496d209 100644 --- a/python/chemiscope/structures/_ase.py +++ b/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), + } diff --git a/src/dataset.ts b/src/dataset.ts index f62611259..dfc439636 100644 --- a/src/dataset.ts +++ b/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): 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); } diff --git a/src/structure/shapes.ts b/src/structure/shapes.ts index db105703f..71e3c12a2 100644 --- a/src/structure/shapes.ts +++ b/src/structure/shapes.ts @@ -61,6 +61,17 @@ export interface EllipsoidParameters extends BaseShapeParameters kind: 'ellipsoid'; } +// Interface for cylinder data (avoids orientation options, since it's redundant) +export interface CylinderData extends BaseShapeData { + vector: [number, number, number]; + radius?: number; +} + +/** Parameters for an arrow shape */ +export interface CylinderParameters extends BaseShapeParameters { + kind: 'cylinder'; +} + // Interface for arrow data (avoids orientation options, since it's redundant) export interface ArrowData extends BaseShapeData { vector: [number, number, number]; @@ -87,7 +98,7 @@ export interface CustomShapeParameters extends BaseShapeParameters) { + super(data); + assert(data.vector); + this.vector = [ + this.scale * data.vector[0], + this.scale * data.vector[1], + this.scale * data.vector[2], + ]; + this.radius = this.scale * (data.radius || 0.1); + } + + public static validateParameters(parameters: Record): string { + if (!('vector' in parameters)) { + return '"vector" is required for "arrow" shapes'; + } + + if (!Array.isArray(parameters.vector) || parameters.vector.length !== 3) { + return '"vector" must be an array with 3 elements for "vector" shapes'; + } + + const [ax, ay, az] = parameters.vector as unknown[]; + if (typeof ax !== 'number' || typeof ay !== 'number' || typeof az !== 'number') { + return '"vector" elements must be numbers for "vector" shapes'; + } + + if ('orientation' in parameters) { + return '"orientation" cannot be used on "cylinder" shapes. define "vector" instead'; + } + + return ''; + } + + public outputTo3Dmol(color: $3Dmol.ColorSpec, resolution: number = 20): $3Dmol.CustomShapeSpec { + const triangulation = triangulateCylinder(this.vector, this.radius, resolution); + const rawVertices = triangulation.vertices; + const indices = triangulation.indices; + const vertices: XYZ[] = []; + const simplices: [number, number, number][] = []; + + for (const v of rawVertices) { + const newVertex: XYZ = addXYZ(v, this.position); + vertices.push(newVertex); + } + + for (let i = 0; i < indices.length; i += 3) { + simplices.push(indices.slice(i, i + 3) as [number, number, number]); + } + + const normals = determineNormals(vertices, simplices); + + return { + vertexArr: vertices, + normalArr: normals, + faceArr: indices, + color: color, + }; + } +} + export class CustomShape extends Shape { public vertices: XYZ[]; public simplices: [number, number, number][]; @@ -663,3 +807,48 @@ export class CustomShape extends Shape { }; } } + +export function add_shapes( + shape_list: $3Dmol.CustomShapeSpec, + shape: $3Dmol.CustomShapeSpec, + viewer: $3Dmol.GLViewer, + max_vertices: number = 0 +): void { + // Dumps shapes if the number of vertices exceeds a set threshold + if (shape_list.vertexArr && shape.vertexArr && max_vertices > 0) { + if (shape_list.vertexArr.length + shape.vertexArr.length >= max_vertices) { + // adds to the viewer and resets the list + viewer.addCustom(shape_list); + shape_list.vertexArr.length = 0; + if (Array.isArray(shape_list.normalArr)) { + shape_list.normalArr.length = 0; + } + if (Array.isArray(shape_list.faceArr)) { + shape_list.faceArr.length = 0; + } + if (Array.isArray(shape_list.color)) { + shape_list.color.length = 0; + } + } + } + // Consolidates a list of shapes to add them all at once + if (shape_list.faceArr && shape.faceArr && shape_list.vertexArr) { + const shift = shape_list.vertexArr.length ?? 0; + const shiftedFaceArr = shape.faceArr.map((value) => value + shift); + shape_list.faceArr.push(...shiftedFaceArr); + } + if (shape_list.vertexArr && shape.vertexArr) { + shape_list.vertexArr?.push(...shape.vertexArr); + } + if (shape_list.normalArr && shape.normalArr) { + shape_list.normalArr?.push(...shape.normalArr); + } + if (shape.vertexArr && Array.isArray(shape_list.color)) { + const newcolor = shape.color && !Array.isArray(shape.color) ? shape.color : 0xffffff; + const newcolors: $3Dmol.ColorSpec[] = Array(shape.vertexArr.length ?? 0).fill( + newcolor + ) as $3Dmol.ColorSpec[]; + //eslint-disable-next-line @typescript-eslint/no-unsafe-assignment + shape_list.color.push(...newcolors); + } +} diff --git a/src/structure/viewer.ts b/src/structure/viewer.ts index c2283db3f..6b67c6e6d 100644 --- a/src/structure/viewer.ts +++ b/src/structure/viewer.ts @@ -12,7 +12,7 @@ import { arrayMaxMin, getElement, sendWarning, unreachable } from '../utils'; import { PositioningCallback } from '../utils'; import { Environment, Settings, Structure } from '../dataset'; -import { Arrow, CustomShape, Ellipsoid, ShapeData, Sphere } from './shapes'; +import { Arrow, CustomShape, Cylinder, Ellipsoid, ShapeData, Sphere, add_shapes } from './shapes'; import { StructureOptions } from './options'; @@ -1102,6 +1102,18 @@ export class MoleculeViewer { const active_shapes = this._options.shape.value.split(','); + // consolidates all shapes in a single CustomShapeSpec + const all_shapes: $3Dmol.CustomShapeSpec = { + vertexArr: [], + normalArr: [], + faceArr: [], + color: [], + }; + + // color function for atoms + const colorfunc = this._colorFunction(); + const atomspec: $3Dmol.AtomSpec = {}; + for (const shape of active_shapes) { if (shape === '') { continue; @@ -1154,31 +1166,54 @@ export class MoleculeViewer { shape_data.position = position; if (atom_pars.color) { shape_data.color = atom_pars.color; + } else if (colorfunc) { + atomspec.serial = i; + shape_data.color = colorfunc(atomspec); } else { - shape_data.color = - atom_pars.color || $3Dmol.elementColors.Jmol[name]; + shape_data.color = $3Dmol.elementColors.Jmol[name]; } if (current_shape.kind === 'sphere') { const shape = new Sphere(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + //this._viewer.addCustom( + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } else if (current_shape.kind === 'ellipsoid') { const shape = new Ellipsoid(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 + ); + } else if (current_shape.kind === 'cylinder') { + const shape = new Cylinder(shape_data); + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } else if (current_shape.kind === 'arrow') { const shape = new Arrow(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } else { assert(current_shape.kind === 'custom'); const shape = new CustomShape(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } } @@ -1186,24 +1221,36 @@ export class MoleculeViewer { // the shape is defined only at the structure level if (current_shape.kind === 'sphere') { const shape = new Sphere(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } else if (current_shape.kind === 'ellipsoid') { const shape = new Ellipsoid(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } else if (current_shape.kind === 'arrow') { const shape = new Arrow(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } else { assert(current_shape.kind === 'custom'); const shape = new CustomShape(shape_data); - this._viewer.addCustom( - shape.outputTo3Dmol(shape_data.color || 0xffffff) + add_shapes( + all_shapes, + shape.outputTo3Dmol(shape_data.color || 0xffffff), + this._viewer, + 32768 ); } } @@ -1212,6 +1259,10 @@ export class MoleculeViewer { } } + if (Array.isArray(all_shapes.faceArr) && all_shapes.faceArr.length > 0) { + // adds all shapes that have been accumulated + this._viewer.addCustom(all_shapes); + } this._viewer.render(); } /**