From bd6c23a9c5273283eb77d8685481ae98a1382b4b Mon Sep 17 00:00:00 2001
From: sss <sss@karma>
Date: Tue, 3 Dec 2024 17:16:28 -0500
Subject: [PATCH] Fix segmentation position calculation for Allegro ECal.

Fix the calculation of cell position for
FCCSWGridModuleThetaMerged_k4geo.
Also add a layer() method to that class to extract the layer from a cell ID.
---
 .../FCCSWGridModuleThetaMerged_k4geo.h        |  39 +++++-
 .../src/FCCSWGridModuleThetaMerged_k4geo.cpp  | 111 +++++++++++++++++-
 2 files changed, 142 insertions(+), 8 deletions(-)

diff --git a/detectorSegmentations/include/detectorSegmentations/FCCSWGridModuleThetaMerged_k4geo.h b/detectorSegmentations/include/detectorSegmentations/FCCSWGridModuleThetaMerged_k4geo.h
index d44f1c6c2..b53a291c5 100644
--- a/detectorSegmentations/include/detectorSegmentations/FCCSWGridModuleThetaMerged_k4geo.h
+++ b/detectorSegmentations/include/detectorSegmentations/FCCSWGridModuleThetaMerged_k4geo.h
@@ -3,6 +3,7 @@
 
 // FCCSW
 #include "detectorSegmentations/GridTheta_k4geo.h"
+#include <atomic>
 
 /** FCCSWGridModuleThetaMerged_k4geo Detector/DetSegmentation/DetSegmentation/FCCSWGridModuleThetaMerged_k4geo.h FCCSWGridModuleThetaMerged_k4geo.h
  *
@@ -21,7 +22,7 @@ class FCCSWGridModuleThetaMerged_k4geo : public GridTheta_k4geo {
   FCCSWGridModuleThetaMerged_k4geo(const BitFieldCoder* decoder);
 
   /// destructor
-  virtual ~FCCSWGridModuleThetaMerged_k4geo() = default;
+  virtual ~FCCSWGridModuleThetaMerged_k4geo();
 
   /// read n(modules) from detector metadata
   void GetNModulesFromGeom();
@@ -33,7 +34,7 @@ class FCCSWGridModuleThetaMerged_k4geo : public GridTheta_k4geo {
    *   @param[in] aCellId ID of a cell.
    *   return Position (relative to R, phi of Geant4 volume it belongs to, scaled for R=1).
    */
-  virtual Vector3D position(const CellID& aCellID) const;
+  virtual Vector3D position(const CellID& aCellID) const override;
   /**  Determine the cell ID based on the position.
    *   @param[in] aLocalPosition (not used).
    *   @param[in] aGlobalPosition
@@ -41,7 +42,7 @@ class FCCSWGridModuleThetaMerged_k4geo : public GridTheta_k4geo {
    *   return Cell ID.
    */
   virtual CellID cellID(const Vector3D& aLocalPosition, const Vector3D& aGlobalPosition,
-                        const VolumeID& aVolumeID) const;
+                        const VolumeID& aVolumeID) const override;
   /**  Determine the azimuthal angle (relative to the G4 volume) based on the cell ID.
    *   @param[in] aCellId ID of a cell.
    *   return Phi.
@@ -94,6 +95,20 @@ class FCCSWGridModuleThetaMerged_k4geo : public GridTheta_k4geo {
    */
   inline const std::string& fieldNameModule() const { return m_moduleID; }
 
+  /// Extract the layer index fom a cell ID.
+  int layer (const CellID& aCellID) const;
+
+  /// Determine the volume ID from the full cell ID by removing all local fields
+  virtual VolumeID volumeID(const CellID& cellID) const override;
+
+  /// Return true if this segmentation can have cells that span multiple
+  /// volumes.  That is, points from multiple distinct volumes may
+  /// be assigned to the same cell.
+  virtual bool cellsSpanVolumes() const override
+  {
+    return true;
+  }
+
 protected:
   /// the field name used for layer
   std::string m_layerID;
@@ -110,6 +125,24 @@ class FCCSWGridModuleThetaMerged_k4geo : public GridTheta_k4geo {
   /// number of layers (from the geometry)
   int m_nLayers;
   
+private:
+  /// Tabulate the cylindrical radii of all layers, as well as the
+  /// local x and z components needed for the proper phi offset.
+  struct LayerInfo {
+    LayerInfo (double the_rho, double the_xloc, double the_zloc)
+      : rho (the_rho), xloc (the_xloc), zloc (the_zloc)
+    {}
+    double rho;
+    double xloc;
+    double zloc;
+  };
+  std::vector<LayerInfo> initLayerInfo (const CellID& cID) const;
+
+  // The vector of tabulated values, indexed by layer number.
+  // We can't build this in the constructor --- the volumes won't have
+  // been created yet.  Instead, build it lazily the first time it's needed.
+  // Since that's in a const method, make it thread-safe.
+  mutable std::atomic<const std::vector<LayerInfo>*> m_layerInfo = nullptr;
 };
 }
 }
diff --git a/detectorSegmentations/src/FCCSWGridModuleThetaMerged_k4geo.cpp b/detectorSegmentations/src/FCCSWGridModuleThetaMerged_k4geo.cpp
index ac54d86a6..8153b4bcf 100644
--- a/detectorSegmentations/src/FCCSWGridModuleThetaMerged_k4geo.cpp
+++ b/detectorSegmentations/src/FCCSWGridModuleThetaMerged_k4geo.cpp
@@ -2,6 +2,7 @@
 
 #include <iostream>
 #include "DD4hep/Detector.h"
+#include "DD4hep/VolumeManager.h"
 
 namespace dd4hep {
 namespace DDSegmentation {
@@ -35,6 +36,12 @@ FCCSWGridModuleThetaMerged_k4geo::FCCSWGridModuleThetaMerged_k4geo(const BitFiel
   GetNLayersFromGeom();
 }
 
+FCCSWGridModuleThetaMerged_k4geo::~FCCSWGridModuleThetaMerged_k4geo()
+{
+  delete m_layerInfo;
+}
+
+
 void FCCSWGridModuleThetaMerged_k4geo::GetNModulesFromGeom() {
   dd4hep::Detector* dd4hepgeo = &(dd4hep::Detector::getInstance());
   try {
@@ -59,14 +66,96 @@ void FCCSWGridModuleThetaMerged_k4geo::GetNLayersFromGeom() {
   std::cout << "Number of layers read from detector metadata and used in readout class: " << m_nLayers << std::endl;
 }
 
+/// Tabulate the cylindrical radii of all layers, as well as the
+/// local x and z components needed for the proper phi offset.
+std::vector<FCCSWGridModuleThetaMerged_k4geo::LayerInfo>
+FCCSWGridModuleThetaMerged_k4geo::initLayerInfo (const CellID& cID) const
+{
+  dd4hep::Detector* dd4hepgeo = &(dd4hep::Detector::getInstance());
+  VolumeManager vman = VolumeManager::getVolumeManager(*dd4hepgeo);
+
+  std::vector<LayerInfo> out;
+  out.reserve (m_nLayers);
+  VolumeID vID = cID;
+  _decoder->set(vID, m_thetaID, 0);
+  for (int l = 0; l < m_nLayers; l++) {
+
+    // Look up a volume in layer l in the volume manager, and find its radius
+    // by transforming the origin in the local coordinate system to global
+    // coordinates.
+    _decoder->set(vID, m_layerID, l);
+    VolumeManagerContext* vc = vman.lookupContext(vID);
+    Position wpos = vc->localToWorld ({0,0,0});
+    double rho = wpos.Rho();
+
+    // If different modules are ganged together, we want to put hits
+    // in the center of all the ganged modules.  This represents
+    // a rotation in phi (in global coordinates).  Convert this rotation
+    // to displacement in local coordinates.  This will be in the x-z plane,
+    // and will be constant for all modules in a layer, but will be different
+    // for different layers (even with identical ganging).
+    double xloc = 0;
+    double zloc = 0;
+    double phioff = phi (vID);
+    if (phioff > 0) {
+      // We need to apply a phi offset.  Calculate it by rotating
+      // the global position in phi and converting back to local
+      // coordinates.
+      //
+      // For the Allegro calorimeter, this can also be calculated as:
+      //
+      //   d = rho * 2 * sin(phioff/2)
+      //   beta = pi/2 - phioff/2 - asin(sin(InclinationAngle)*EMBarrel_rmin/rho)
+      //   xloc = - d * sin(beta)
+      //   yloc =   d * cos(beta)
+      //
+      // but we prefer to do the calculation via explicit rotations because
+      // it's easier to see that that was is correct, and it also avoids
+      // the explicit dependencies on the geometry parameters.
+      Position wpos2 = RotateZ (wpos, phioff);
+      Position lpos2 = vc->worldToLocal (wpos2);
+      xloc = lpos2.X();
+      zloc = lpos2.Z();
+    }
+
+    out.emplace_back (rho, xloc, zloc);
+  }
+  return out;
+}
+
+
 /// determine the local position based on the cell ID
 Vector3D FCCSWGridModuleThetaMerged_k4geo::position(const CellID& cID) const {
 
+  // Get the vector of layer info.  If it hasn't been made yet,
+  // calculate it now.
+  const std::vector<LayerInfo>* liv = m_layerInfo.load();
+  if (!liv) {
+    auto liv_new = new std::vector<LayerInfo> (initLayerInfo(cID));
+    if (m_layerInfo.compare_exchange_strong (liv, liv_new)) {
+      liv = liv_new;
+    }
+    else {
+      delete liv_new;
+    }
+  }
+
+  VolumeID vID = cID;
+  _decoder->set(vID, m_thetaID, 0);
+  int layer = this->layer (vID);
+
   // debug
   // std::cout << "cellID: " << cID << std::endl;
 
-  // cannot return for R=0 otherwise will lose phi info, return for R=1
-  return positionFromRThetaPhi(1.0, theta(cID), phi(cID));
+  // Calculate the position in local coordinates.
+  // The volume here has the cross-section of a cell in the x-z plane;
+  // it extends the length of the calorimeter along the y-axis.
+  // We set the y-coordinate based on the theta bin, and x and z
+  // based on the phi offset required for this layer.
+  const LayerInfo& li = liv->at(layer);
+  return Vector3D(li.xloc,
+                  -li.rho / tan(theta(cID)),
+                  li.zloc);
 }
 
 /// determine the cell ID based on the global position
@@ -75,7 +164,7 @@ CellID FCCSWGridModuleThetaMerged_k4geo::cellID(const Vector3D& /* localPosition
   CellID cID = vID;
 
   // retrieve layer (since merging depends on layer)
-  int layer = _decoder->get(vID, m_layerID);
+  int layer = this->layer (vID);
 
   // retrieve theta
   double lTheta = thetaFromXYZ(globalPosition);
@@ -112,7 +201,7 @@ CellID FCCSWGridModuleThetaMerged_k4geo::cellID(const Vector3D& /* localPosition
 double FCCSWGridModuleThetaMerged_k4geo::phi(const CellID& cID) const {
 
   // retrieve layer
-  int layer = _decoder->get(cID, m_layerID);
+  int layer = this->layer (cID);
 
   // calculate phi offset due to merging
   // assume that m_mergedModules[layer]>=1
@@ -131,7 +220,7 @@ double FCCSWGridModuleThetaMerged_k4geo::phi(const CellID& cID) const {
 double FCCSWGridModuleThetaMerged_k4geo::theta(const CellID& cID) const {
 
   // retrieve layer
-  int layer = _decoder->get(cID, m_layerID);
+  int layer = this->layer (cID);
 
   // retrieve theta bin from cellID and determine theta position
   CellID thetaValue = _decoder->get(cID, m_thetaID);
@@ -152,5 +241,17 @@ double FCCSWGridModuleThetaMerged_k4geo::theta(const CellID& cID) const {
   return _theta;
 }
 
+/// Extract the layer index fom a cell ID.
+int FCCSWGridModuleThetaMerged_k4geo::layer(const CellID& cID) const {
+  return _decoder->get(cID, m_layerID);
+}
+
+/// Determine the volume ID from the full cell ID by removing all local fields
+VolumeID FCCSWGridModuleThetaMerged_k4geo::volumeID(const CellID& cID) const {
+  VolumeID vID = cID;
+  _decoder->set(vID, m_thetaID, 0);
+  return vID;
+}
+
 }
 }