types for RecoParticle relations including 1:1 and 1:many

key4hep · Jun 4, 2024 · 3e64c41 · 3e64c41
1 parent 83a9bbd
commit 3e64c41
Show file tree

Hide file tree

Showing 3 changed files with 102 additions and 54 deletions.
diff --git a/js/types/load.js b/js/types/load.js
@@ -18,10 +18,10 @@ const types = {
 
 const loadersConfig = [
   "ReconstructedParticle",
-  "ParticleID",
-  "Vertex",
-  "Track",
-  "Cluster",
+  // "ParticleID",
+  // "Vertex",
+  // "Track",
+  // "Cluster",
 ];
 
 export function buildLoader(config, version) {
@@ -61,11 +61,27 @@ export function loadParticles(jsonData, event, loadersConfig) {
     );
 
     particlesType.forEach(({ collection }) => {
-      const loadedParticles = loadFunction(collection);
+      const [particles, links] = loadFunction(collection);
     });
   }
 
   return particles;
 }
 
-loadParticles(jsonData, 4, loadersConfig);
+export function dynamicLoad(object, data, ignore = null) {
+  let filteredData = {};
+
+  if (ignore !== null) {
+    for (const key in data) {
+      if (!ignore.has(key)) filteredData[key] = data[key];
+    }
+  } else {
+    filteredData = data;
+  }
+
+  for (const [key, value] of Object.entries(data)) {
+    object[key] = value;
+  }
+}
+
+loadParticles(jsonData, 0, loadersConfig);
diff --git a/js/types/reconstruction.js b/js/types/reconstruction.js
@@ -1,14 +1,4 @@
-function dynamicLoad(object, data, ignore = null) {
-  if (ignore !== null) {
-    for (const key of ignore) {
-      delete data[key];
-    }
-  }
-
-  for (const [key, value] of Object.entries(data)) {
-    object[key] = value;
-  }
-}
+import { dynamicLoad } from "./load.js";
 
 export class Cluster {
   static MIN_VERSION = "0.7.0"; // may vary per type of particle
@@ -98,17 +88,34 @@ export class ReconstructedParticle {
 
   static load(collection) {
     const particles = [];
+    const links = createLinksManager([
+      "tracks",
+      "clusters",
+      "particles",
+      "startVertex",
+    ]);
 
     for (const [index, particle] of collection.entries()) {
       const reconstructedParticle = new ReconstructedParticle();
       reconstructedParticle.index = index;
 
-      dynamicLoad(reconstructedParticle, particle);
+      extractOneToManyLinks(
+        links,
+        ["tracks", "clusters", "particles"],
+        particle
+      );
+      extractOneToOneLinks(links, "startVertex", particle);
+
+      dynamicLoad(
+        reconstructedParticle,
+        particle,
+        new Set(["tracks", "clusters", "particles", "startVertex"])
+      );
 
       particles.push(reconstructedParticle);
     }
 
-    return particles;
+    return [particles, links];
   }
 }
 
@@ -165,3 +172,48 @@ export class Track {
     return particles;
   }
 }
+
+export class GenericLink {
+  // we may create a specific class for each type if needed
+  constructor(id, from, to) {
+    this.id = id;
+    this.from = from;
+    this.to = to;
+  }
+}
+
+function createLinksManager(types) {
+  const links = {};
+  types.forEach((type) => (links[type] = []));
+  return links;
+}
+
+export function createGenericLink(id, from, { collectionID, index }) {
+  const genericLink = new GenericLink(id, from, index);
+  genericLink.collectionID = collectionID;
+  return genericLink;
+}
+
+function extractOneToManyLinks(linksManager, keys, particle) {
+  for (const key of keys) {
+    particle[key].map((val) => {
+      const link = createGenericLink(
+        linksManager[key].length,
+        particle.index,
+        val
+      );
+      linksManager[key].push(link);
+      particle[key].push(link.id);
+    });
+  }
+}
+
+function extractOneToOneLinks(linksManager, key, particle) {
+  const link = createGenericLink(
+    linksManager[key].length,
+    particle.index,
+    particle[key]
+  );
+  linksManager[key].push(link);
+  particle[key] = link.id;
+}
diff --git a/test/load.test.js b/test/load.test.js
@@ -31,21 +31,21 @@ const data = {
         },
       ],
     },
-    "Jet": {
-      "collID": 13,
+    "ReconstructedParticles": {
+      "collID": 14,
       "collType": "edm4hep::ReconstructedParticleCollection",
       "collection": [
         {
           "charge": 1.0,
           "clusters": [],
           "covMatrix": [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
-          "energy": 12.062528610229492,
+          "energy": 0.358455091714859,
           "goodnessOfPID": 0.0,
-          "mass": 2.315242290496826,
+          "mass": 0.13957038521766663,
           "momentum": {
-            "x": 11.738886833190918,
-            "y": 1.2114704847335815,
-            "z": 0.9354811906814575,
+            "x": 0.09558331221342087,
+            "y": -0.11288221180438995,
+            "z": -0.2951807975769043,
           },
           "particleIDUsed": {
             "collectionID": -2,
@@ -54,35 +54,10 @@ const data = {
           "particleIDs": [
             {
               "collectionID": 4,
-              "index": 45,
-            },
-          ],
-          "particles": [
-            {
-              "collectionID": 14,
-              "index": 24,
-            },
-            {
-              "collectionID": 14,
-              "index": 22,
-            },
-            {
-              "collectionID": 14,
-              "index": 74,
-            },
-            {
-              "collectionID": 14,
-              "index": 23,
-            },
-            {
-              "collectionID": 14,
-              "index": 25,
-            },
-            {
-              "collectionID": 14,
-              "index": 26,
+              "index": 0,
             },
           ],
+          "particles": [],
           "referencePoint": {
             "x": 0.0,
             "y": 0.0,
@@ -92,7 +67,12 @@ const data = {
             "collectionID": -2,
             "index": -2,
           },
-          "tracks": [],
+          "tracks": [
+            {
+              "collectionID": 3,
+              "index": 0,
+            },
+          ],
           "type": 0,
         },
       ],