feat(event-display): visualize mc particles as arrows

packages/phoenix-event-display/src/loaders/edm4hep-json-loader.ts

@@ -29,6 +29,7 @@ export class Edm4hepJsonLoader extends PhoenixLoader {
         CaloClusters: {},
         Jets: {},
         MissingEnergy: {},
+        MCParticles: {},
         'event number': this.getEventNumber(event),
         'run number': this.getRunNumber(event),
       };
@@ -42,6 +43,7 @@ export class Edm4hepJsonLoader extends PhoenixLoader {
       oneEventData.CaloClusters = this.getCaloClusters(event);
       oneEventData.Jets = this.getJets(event);
       oneEventData.MissingEnergy = this.getMissingEnergy(event);
+      oneEventData.MCParticles = this.getMCParticles(event);
 
       this.eventData[eventName] = oneEventData;
     });
@@ -621,14 +623,106 @@ export class Edm4hepJsonLoader extends PhoenixLoader {
     return allMETs;
   }
 
-  /** Return a random colour */
+  /** Returns Monte Carlo particles */
+  private getMCParticles(event: any) {
+    const allParticles: any[] = [];
+
+    for (const collName in event) {
+      if (event[collName].constructor != Object) {
+        continue;
+      }
+
+      const collDict = event[collName];
+
+      if (!('collType' in collDict)) {
+        continue;
+      }
+
+      if (!collDict['collType'].includes('edm4hep::')) {
+        continue;
+      }
+
+      if (!collDict['collType'].includes('MCParticleCollection')) {
+        continue;
+      }
+
+      if (!('collection' in collDict)) {
+        continue;
+      }
+
+      const rawParticles = collDict['collection'];
+      const particles: any[] = [];
+
+      rawParticles.forEach((rawParticle: any) => {
+        const origin: number[] = [];
+        origin.push(rawParticle.vertex.x * 0.1);
+        origin.push(rawParticle.vertex.y * 0.1);
+        origin.push(rawParticle.vertex.z * 0.1);
+
+        const momentum: number[] = [];
+        momentum.push(rawParticle.momentum.x);
+        momentum.push(rawParticle.momentum.y);
+        momentum.push(rawParticle.momentum.z);
+
+        const particle = {
+          origin: origin,
+          momentum: momentum,
+          pdgid: rawParticle.PDG,
+          status: rawParticle.generatorStatus,
+          color: this.mcParticleColor(rawParticle.PDG),
+        };
+        particles.push(particle);
+      });
+
+      allParticles[collName] = particles;
+    }
+
+    return allParticles;
+  }
+
+  /** Return a random color */
   private randomColor() {
     return Math.floor(Math.random() * 16777215)
       .toString(16)
       .padStart(6, '0')
       .toUpperCase();
   }
 
+  /** Return color depending on the particle type */
+  private mcParticleColor(mcParticlePDGid) {
+    switch (Math.abs(mcParticlePDGid)) {
+      case 1:
+      case 2:
+      case 3:
+      case 4:
+      case 5:
+      case 6:
+        return '#aa0000';
+      case 11:
+      case 12:
+      case 13:
+      case 14:
+      case 15:
+      case 16:
+        return '#00aa00';
+      case 21:
+      case 22:
+      case 23:
+      case 24:
+        return '#0000aa';
+      case 25:
+        return '#00aaaa';
+      case 111:
+      case 211:
+      case 221:
+      case 311:
+      case 321:
+        return '#aa00aa';
+      default:
+        return '#aaaa00';
+    }
+  }
+
   /** Helper conversion of HSL to hexadecimal */
   private convHSLtoHEX(h: number, s: number, l: number): string {
     l /= 100;

packages/phoenix-event-display/src/loaders/objects/phoenix-objects.ts

@@ -7,6 +7,7 @@ import {
   MeshToonMaterial,
   Mesh,
   BufferGeometry,
+  ConeGeometry,
   LineBasicMaterial,
   Line,
   Group,
@@ -24,10 +25,13 @@ import {
   CanvasTexture,
 } from 'three';
 import { ConvexGeometry } from 'three/examples/jsm/geometries/ConvexGeometry.js';
+import { TextGeometry } from 'three/examples/jsm/geometries/TextGeometry.js';
+import { Font } from 'three/examples/jsm/loaders/FontLoader.js';
+import HelvetikerFont from '../../managers/three-manager/fonts/helvetiker_regular.typeface.json';
 import { EVENT_DATA_TYPE_COLORS } from '../../helpers/constants';
 import { RKHelper } from '../../helpers/rk-helper';
 import { CoordinateHelper } from '../../helpers/coordinate-helper';
-import { mergeBufferGeometries } from 'three/examples/jsm/utils/BufferGeometryUtils.js';
+import { mergeGeometries } from 'three/examples/jsm/utils/BufferGeometryUtils.js';
 import { TracksMaterial, TracksMesh } from './tracks';
 import * as BufferGeometryUtils from 'three/examples/jsm/utils/BufferGeometryUtils';
 
@@ -472,7 +476,7 @@ export class PhoenixObjects {
       boxGeometry.translate(pointPos[i], pointPos[i + 1], pointPos[i + 2]);
       geometries.push(boxGeometry);
     }
-    const geometry = mergeBufferGeometries(geometries);
+    const geometry = mergeGeometries(geometries);
     geometry.computeBoundingSphere();
     // material
     const material = new MeshPhongMaterial({
@@ -901,4 +905,186 @@ export class PhoenixObjects {
 
     return cell;
   }
+
+  /**
+   * Create and return a Monte Carlo particle arrow from the given parameters.
+   * @param mcParticleParams  MCParticle parameters.
+   * @returns  Calorimeter MCParticle object.
+   */
+  public static getMCParticle(mcParticleParams: {
+    origin: number[];
+    momentum: number[];
+    pdgid: number;
+    status?: number;
+    color?: string;
+    uuid: string;
+  }): Object3D {
+    const defaultColor: string = '#ffff00';
+
+    const origin = new Vector3(
+      mcParticleParams.origin[0],
+      mcParticleParams.origin[1],
+      mcParticleParams.origin[2],
+    );
+
+    const direction = new Vector3(
+      mcParticleParams.momentum[0],
+      mcParticleParams.momentum[1],
+      mcParticleParams.momentum[2],
+    );
+    const length = direction.length();
+    direction.normalize();
+
+    const lineLength = 0.85 * length;
+    let lineWidth = Math.log(length * 1000) / 100;
+    if (lineWidth < 0) {
+      lineWidth = 0.00001;
+    }
+    if (lineWidth > 0.1) {
+      lineWidth = 0.1;
+    }
+    const coneLength = 0.15 * length;
+    const coneWidth = 2.5 * lineWidth;
+
+    const lineGeometry = new CylinderGeometry(
+      lineWidth,
+      lineWidth,
+      lineLength,
+      16,
+    );
+    lineGeometry.rotateZ(Math.PI / 2);
+    lineGeometry.translate(lineLength / 2, 0, 0);
+
+    const coneGeometry = new ConeGeometry(coneWidth, coneLength, 16);
+    coneGeometry.rotateZ(-Math.PI / 2);
+    coneGeometry.translate(length - coneLength / 2, 0, 0);
+
+    const geometries = [lineGeometry, coneGeometry];
+    const mergedGeometry = mergeGeometries(geometries, false);
+
+    const buildDirection = new Vector3(1, 0, 0).normalize();
+
+    const quaternion = new Quaternion();
+    quaternion.setFromUnitVectors(buildDirection, direction);
+    mergedGeometry.applyQuaternion(quaternion);
+
+    const material = new MeshPhongMaterial({
+      color: mcParticleParams.color ?? defaultColor,
+    });
+
+    const arrowObject = new Mesh(mergedGeometry, material);
+
+    const font = new Font(HelvetikerFont);
+    const labelGeometry = new TextGeometry(
+      PhoenixObjects.getMCParticleName(mcParticleParams.pdgid),
+      {
+        font: font,
+        size: 100,
+        height: 15,
+      },
+    );
+    labelGeometry.scale(0.01, 0.01, 0.01);
+
+    labelGeometry.rotateX(Math.PI / 2);
+    labelGeometry.translate(length, 0, 0);
+
+    labelGeometry.applyQuaternion(quaternion);
+
+    const labelObject = new Mesh(labelGeometry, material);
+
+    const arrowContainer = new Group();
+    arrowContainer.position.x = origin.x;
+    arrowContainer.position.y = origin.y;
+    arrowContainer.position.z = origin.z;
+
+    arrowContainer.add(arrowObject);
+    arrowContainer.add(labelObject);
+
+    mcParticleParams.uuid = arrowContainer.uuid;
+    arrowContainer.userData = Object.assign({}, mcParticleParams);
+
+    arrowContainer.name = 'MCParticle';
+
+    return arrowContainer;
+  }
+
+  /**
+   * Return a Monte Carlo particle name from PDG ID.
+   * @param mcParticlePDGid  PDG ID.
+   * @returns  MCParticle name.
+   */
+  public static getMCParticleName(mcParticlePDGid: number): string {
+    switch (mcParticlePDGid) {
+      case 1:
+        return 'd';
+      case -1:
+        return '!d';
+      case 2:
+        return 'u';
+      case -2:
+        return '!u';
+      case 3:
+        return 's';
+      case -3:
+        return '!s';
+      case 4:
+        return 'c';
+      case -4:
+        return '!c';
+      case 5:
+        return 'b';
+      case -5:
+        return '!b';
+      case 6:
+        return 't';
+      case -6:
+        return '!t';
+      case 11:
+        return 'e-';
+      case -11:
+        return 'e+';
+      case 12:
+        return 'νe';
+      case 13:
+        return 'μ-';
+      case -13:
+        return 'μ+';
+      case 14:
+        return 'νμ';
+      case 15:
+        return 'τ';
+      case -15:
+        return 'τ-';
+      case 16:
+        return 'ντ';
+      case 21:
+        return 'g';
+      case 22:
+        return 'γ';
+      case 23:
+        return 'Z';
+      case 24:
+        return 'W+';
+      case -24:
+        return 'W-';
+      case 25:
+        return 'H';
+      case 111:
+        return 'π0';
+      case 211:
+        return 'π+';
+      case -211:
+        return 'π-';
+      case 221:
+        return 'η';
+      case 311:
+        return 'K0';
+      case 321:
+        return 'K+';
+      case -321:
+        return 'K-';
+      default:
+        return mcParticlePDGid.toString();
+    }
+  }
 }

packages/phoenix-event-display/src/loaders/phoenix-loader.ts

@@ -433,6 +433,47 @@ export class PhoenixLoader implements EventDataLoader {
         addMETSizeOption,
       );
     }
+
+    if (eventData.MCParticles) {
+      const cuts = [new Cut('status', 21, 29, 200)];
+
+      const addMCParticlesSizeOption = (
+        typeFolder: GUI,
+        typeFolderPM: PhoenixMenuNode,
+      ) => {
+        const scaleMCParticles = (value: number) => {
+          this.graphicsLibrary
+            .getSceneManager()
+            .scaleChildObjects('MCParticles', value / 100);
+        };
+        if (typeFolder) {
+          const sizeMenu = typeFolder
+            .add({ particleScale: 100 }, 'particleScale', 1, 400)
+            .name('Size (%)');
+          sizeMenu.onChange(scaleMCParticles);
+        }
+        // Phoenix menu
+        if (typeFolderPM) {
+          typeFolderPM.addConfig('slider', {
+            label: 'Size (%)',
+            value: 100,
+            min: 1,
+            max: 400,
+            allowCustomValue: true,
+            onChange: scaleMCParticles,
+          });
+        }
+      };
+
+      this.addObjectType(
+        eventData.MCParticles,
+        PhoenixObjects.getMCParticle,
+        'MCParticles',
+        false,
+        cuts,
+        addMCParticlesSizeOption,
+      );
+    }
   }
 
   /**