From 5f8ead25a1ddce93fda760188289a5d166067453 Mon Sep 17 00:00:00 2001
From: Michele Ceriotti <michele.ceriotti@gmail.com>
Date: Fri, 8 Dec 2023 15:07:18 -0800
Subject: [PATCH] Added cylinder shape type

---
 python/chemiscope/input.py |  27 ++++++-
 src/dataset.ts             |   4 +-
 src/structure/shapes.ts    | 149 ++++++++++++++++++++++++++++++++++++-
 src/structure/viewer.ts    |   7 +-
 4 files changed, 181 insertions(+), 6 deletions(-)

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/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<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);
     }
diff --git a/src/structure/shapes.ts b/src/structure/shapes.ts
index db105703f..c0e219d56 100644
--- a/src/structure/shapes.ts
+++ b/src/structure/shapes.ts
@@ -61,6 +61,17 @@ export interface EllipsoidParameters extends BaseShapeParameters<EllipsoidData>
     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<CylinderData> {
+    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<CustomShapeDa
     kind: 'custom';
 }
 
-export type ShapeData = SphereData | EllipsoidData | ArrowData | CustomShapeData;
+export type ShapeData = SphereData | EllipsoidData | ArrowData | CylinderData | CustomShapeData;
 
 /**
  * Describes a shape, to be displayed alongside an atomic structure.
@@ -104,6 +115,7 @@ export type ShapeData = SphereData | EllipsoidData | ArrowData | CustomShapeData
 export type ShapeParameters =
     | SphereParameters
     | EllipsoidParameters
+    | CylinderParameters
     | ArrowParameters
     | CustomShapeParameters;
 
@@ -574,6 +586,141 @@ export class Arrow extends Shape {
     }
 }
 
+function triangulateCylinder(
+    vector: [number, number, number],
+    radius: number,
+    resolution: number = 20
+): { vertices: XYZ[]; indices: number[] } {
+    const [x, y, z] = vector;
+    const tip: XYZ = { x, y, z };
+    const v_len = Math.sqrt(x * x + y * y + z * z);
+
+    // generates a unit circle oriented in the right direction
+    const n_vec: XYZ = multXYZ(tip, 1.0 / v_len);
+
+    // Generate an arbitrary vector not collinear with n
+    let vx: XYZ;
+    if (n_vec.x !== 0.0 || n_vec.y !== 0.0) {
+        vx = { x: 0, y: 0, z: 1 };
+    } else {
+        vx = { x: 0, y: 1, z: 0 };
+    }
+
+    // generate orthogonal vectors in the plane defined by nvec
+    let u: XYZ = addXYZ(vx, multXYZ(n_vec, -dotXYZ(vx, n_vec)));
+    u = multXYZ(u, 1.0 / Math.sqrt(dotXYZ(u, u)));
+    const v: XYZ = crossXYZ(u, n_vec);
+
+    // generate n_points in the plane defined by nvec, centered at vec
+    const circle_points: XYZ[] = [];
+    for (let i = 0; i < resolution; i++) {
+        circle_points.push(
+            addXYZ(
+                multXYZ(u, Math.cos((i * 2 * Math.PI) / resolution)),
+                multXYZ(v, Math.sin((i * 2 * Math.PI) / resolution))
+            )
+        );
+    }
+
+    let indices: number[] = [];
+    const vertices: XYZ[] = [];
+
+    vertices.push({ x: 0, y: 0, z: 0 });
+    vertices.push(tip);
+    // the cylinder is built as a surface of revolution, by stacking |_| motifs
+    for (let i = 0; i < resolution; i++) {
+        // nb replicated points are needed to get sharp edges
+        vertices.push(multXYZ(circle_points[i], radius));
+        vertices.push(multXYZ(circle_points[i], radius));
+        vertices.push(addXYZ(multXYZ(circle_points[i], radius), tip));
+        vertices.push(addXYZ(multXYZ(circle_points[i], radius), tip));
+        const i_seg = 2 + i * 4;
+        const i_next = 2 + ((i + 1) % resolution) * 4;
+        indices = [
+            ...indices,
+            ...[
+                0,
+                i_seg,
+                i_next, // cylinder base
+                i_seg + 1,
+                i_seg + 2,
+                i_next + 1,
+                i_next + 1,
+                i_seg + 2,
+                i_next + 2, // cylinder side
+                i_seg + 3,
+                1,
+                i_next + 3, // cylinder top
+            ],
+        ];
+    }
+    return {
+        vertices: vertices,
+        indices: indices,
+    };
+}
+
+export class Cylinder extends Shape {
+    public vector: [number, number, number];
+    public radius: number;
+
+    constructor(data: Partial<CylinderData>) {
+        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, unknown>): 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 / 3; i++) {
+            simplices.push([indices[3 * i], indices[3 * i + 1], indices[3 * i + 2]]);
+        }
+
+        return {
+            vertexArr: vertices,
+            normalArr: determineNormals(vertices, simplices),
+            faceArr: indices,
+            color: color,
+        };
+    }
+}
+
 export class CustomShape extends Shape {
     public vertices: XYZ[];
     public simplices: [number, number, number][];
diff --git a/src/structure/viewer.ts b/src/structure/viewer.ts
index c2283db3f..32a13f2b2 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 } from './shapes';
 
 import { StructureOptions } from './options';
 
@@ -1169,6 +1169,11 @@ export class MoleculeViewer {
                                     this._viewer.addCustom(
                                         shape.outputTo3Dmol(shape_data.color || 0xffffff)
                                     );
+                                } else if (current_shape.kind === 'cylinder') {
+                                    const shape = new Cylinder(shape_data);
+                                    this._viewer.addCustom(
+                                        shape.outputTo3Dmol(shape_data.color || 0xffffff)
+                                    );
                                 } else if (current_shape.kind === 'arrow') {
                                     const shape = new Arrow(shape_data);
                                     this._viewer.addCustom(