Rga proton energy loss corrections

src/iguana/algorithms/clas12/rga/ProtonEnergyLossCorrection/Action.yaml

@@ -0,0 +1,34 @@
+algorithm:
+  name: 'clas12::ProtonEnergyLossCorrection'
+
+actions:
+
+  - name: Transform
+    type: transformer
+    rank: scalar
+    inputs:
+      - name: pid
+        type: int
+      - name: status
+        type: int
+      - name: run
+        type: int
+      - name: px
+        type: double
+        cast: float
+      - name: py
+        type: double
+        cast: float
+      - name: pz
+        type: double
+        cast: float
+    outputs:
+      - name: px
+        type: double
+        cast: float
+      - name: py
+        type: double
+        cast: float
+      - name: pz
+        type: double
+        cast: float

src/iguana/algorithms/clas12/rga/ProtonEnergyLossCorrection/Algorithm.cc

@@ -0,0 +1,353 @@
+#include "Algorithm.h"
+
+#include <cmath>
+#include <stdexcept>
+
+#include <yaml-cpp/yaml.h>
+
+namespace iguana::clas12::rga {
+
+  REGISTER_IGUANA_ALGORITHM(ProtonEnergyLossCorrection);
+
+  namespace {
+    // Keep an angle in degrees within [0, 360).
+    double WrapDeg360(double x) {
+      while(x >= 360.0) { x -= 360.0; }
+      while(x < 0.0) { x += 360.0; }
+      return x;
+    }
+  } 
+
+  // -----------------------------------------------------------------------------
+  // Detector-region helpers (based on REC::Particle status)
+  // -----------------------------------------------------------------------------
+  //
+  // CLAS12 status definitions:
+  //   - FD tracks: abs(status) in [2000, 4000)
+  //   - CD tracks: abs(status) in [4000, 5000)
+  //
+  bool ProtonEnergyLossCorrection::IsFD(int status) {
+    int s = std::abs(status);
+    return (s >= 2000 && s < 4000);
+  }
+
+  bool ProtonEnergyLossCorrection::IsCD(int status) {
+    int s = std::abs(status);
+    return (s >= 4000 && s < 5000);
+  }
+
+  // -----------------------------------------------------------------------------
+  // math helpers (p, theta, phi)
+  // -----------------------------------------------------------------------------
+  double ProtonEnergyLossCorrection::Pmag(double px, double py, double pz) {
+    return std::sqrt(px * px + py * py + pz * pz);
+  }
+
+  // theta from momentum vector, in degrees.
+  // cos(theta) = pz / |p|
+  double ProtonEnergyLossCorrection::ThetaDeg(double px, double py, double pz) {
+    double r = Pmag(px, py, pz);
+    if(r <= 0.0) {
+      return 0.0;
+    }
+
+    double c = pz / r;
+    if(c > 1.0) {
+      c = 1.0;
+    }
+    if(c < -1.0) {
+      c = -1.0;
+    }
+
+    return (180.0 / M_PI) * std::acos(c);
+  }
+
+
+  // Returns phi in [0,360).
+  double ProtonEnergyLossCorrection::PhiDeg(double px, double py) {
+    double phi = (180.0 / M_PI) * std::atan2(px, py);
+    phi = phi - 90.0;
+    if(phi < 0.0) {
+      phi = 360.0 + phi;
+    }
+    phi = 360.0 - phi;
+
+    return WrapDeg360(phi);
+  }
+
+  // Convert from (p, theta_deg, phi_deg) back to Cartesian (px,py,pz).
+  void ProtonEnergyLossCorrection::SphericalToCartesian(double p, double theta_deg,
+    double phi_deg, double& px, double& py, double& pz) {
+    double theta = theta_deg * (M_PI / 180.0);
+    double phi   = phi_deg * (M_PI / 180.0);
+
+    px = p * std::sin(theta) * std::cos(phi);
+    py = p * std::sin(theta) * std::sin(phi);
+    pz = p * std::cos(theta);
+  }
+
+  // Evaluate polynomial c0 + c1*x + c2*x^2 + ...
+  double ProtonEnergyLossCorrection::EvalPoly(Poly const& p, double x) {
+    double sum = 0.0;
+    double xn  = 1.0;
+    for(double c : p.c) {
+      sum += c * xn;
+      xn *= x;
+    }
+    return sum;
+  }
+
+  // -----------------------------------------------------------------------------
+  // Period lookup
+  // -----------------------------------------------------------------------------
+  //
+  // Given a run number, find the first period def whose run_ranges contain it.
+  // If no match is found, return nullptr and the algorithm leaves the row unchanged.
+  ProtonEnergyLossCorrection::PeriodDef 
+    const* ProtonEnergyLossCorrection::FindPeriod(int run) const {
+    for(auto const& kv : m_periods) {
+      auto const& def = kv.second;
+      for(auto const& rr : def.run_ranges) {
+        if(run >= rr.first && run <= rr.second) {
+          return &def;
+        }
+      }
+    }
+    return nullptr;
+  }
+
+  // -----------------------------------------------------------------------------
+  // ConfigHook: read Config.yaml and populate m_periods
+  // -----------------------------------------------------------------------------
+  //
+  // YAML schema expected:
+  //
+  //   clas12:
+  //     ProtonEnergyLossCorrection:
+  //       periods:
+  //         <period_key>:
+  //           run_ranges:
+  //             - [min, max]
+  //             - [min, max]
+  //           FD:
+  //             A_p: [ ... ]
+  //             B_p: [ ... ]
+  //             ...
+  //           CD:
+  //             A_p: [ ... ]
+  //             ...
+  //
+  // For each period, load FD and CD RegionCoeffs. Each RegionCoeffs holds
+  // polynomial parameterizations in theta for A/B/C of p/theta/phi corrections.
+  // many A/B/C may be zero!!
+  void ProtonEnergyLossCorrection::ConfigHook() {
+    m_periods.clear();
+
+    // algorithm's installed Config.yaml
+    std::string const cfg_path =
+      GetConfig()->FindFile("algorithms/clas12/rga/ProtonEnergyLossCorrection/Config.yaml");
+
+    YAML::Node root = YAML::LoadFile(cfg_path);
+    YAML::Node node = root["clas12"]["ProtonEnergyLossCorrection"];
+    YAML::Node periods_node = node["periods"];
+
+    // Helper: load a polynomial array by name (e.g. "A_p") from a region node.
+    auto load_poly = [](YAML::Node const& region_node, char const* name) -> Poly {
+      YAML::Node a = region_node[name];
+      Poly p;
+      p.c.reserve((size_t)a.size());
+      for(auto v : a) {
+        p.c.push_back(v.as<double>());
+      }
+      return p;
+    };
+
+    // Helper: load all A/B/C for p/theta/phi for a given detector region node.
+    auto load_region = [&](YAML::Node const& rnode) -> RegionCoeffs {
+      RegionCoeffs rc;
+
+      rc.A_p = load_poly(rnode, "A_p");
+      rc.B_p = load_poly(rnode, "B_p");
+      rc.C_p = load_poly(rnode, "C_p");
+
+      rc.A_theta = load_poly(rnode, "A_theta");
+      rc.B_theta = load_poly(rnode, "B_theta");
+      rc.C_theta = load_poly(rnode, "C_theta");
+
+      rc.A_phi = load_poly(rnode, "A_phi");
+      rc.B_phi = load_poly(rnode, "B_phi");
+      rc.C_phi = load_poly(rnode, "C_phi");
+
+      return rc;
+    };
+
+    // Iterate periods dictionary: key -> data
+    for(auto it = periods_node.begin(); it != periods_node.end(); ++it) {
+      std::string key = it->first.as<std::string>();
+      YAML::Node pnode = it->second;
+
+      PeriodDef def;
+
+      for(auto rr : pnode["run_ranges"]) {
+        def.run_ranges.push_back({rr[0].as<int>(), rr[1].as<int>()});
+      }
+
+      def.fd = load_region(pnode["FD"]);
+      def.cd = load_region(pnode["CD"]);
+
+      m_periods[key] = def;
+    }
+  }
+
+  // -----------------------------------------------------------------------------
+  // StartHook: cache bank indices
+  // -----------------------------------------------------------------------------
+  void ProtonEnergyLossCorrection::StartHook(hipo::banklist& banks) {
+    m_b_rec_particle = GetBankIndex(banks, "REC::Particle");
+    m_b_run_config   = GetBankIndex(banks, "RUN::config");
+  }
+
+  // -----------------------------------------------------------------------------
+  // RunHook: apply correction to each matching REC::Particle row
+  // -----------------------------------------------------------------------------
+  bool ProtonEnergyLossCorrection::RunHook(hipo::banklist& banks) const {
+    auto& rec = GetBank(banks, m_b_rec_particle, "REC::Particle");
+    auto& run = GetBank(banks, m_b_run_config, "RUN::config");
+
+    int runnum = run.getInt("run", 0);
+
+    // Loop all rows and correct in place.
+    for(auto const& row : rec.getRowList()) {
+
+      // Select protons only.
+      int pid = rec.getInt("pid", row);
+      if(pid != 2212) {
+        continue;
+      }
+
+      // Select FD or CD only (based on status).
+      int status = rec.getInt("status", row);
+      if(!IsFD(status) && !IsCD(status)) {
+        continue;
+      }
+
+      // Read momentum components.
+      double px = rec.getFloat("px", row);
+      double py = rec.getFloat("py", row);
+      double pz = rec.getFloat("pz", row);
+
+      // Transform returns corrected values (or unchanged if not applicable).
+      auto [px_new, py_new, pz_new] = Transform(pid, status, runnum, px, py, pz);
+
+      // Write back to bank.
+      rec.putFloat("px", row, static_cast<float>(px_new));
+      rec.putFloat("py", row, static_cast<float>(py_new));
+      rec.putFloat("pz", row, static_cast<float>(pz_new));
+    }
+
+    return true;
+  }
+
+  void ProtonEnergyLossCorrection::StopHook() {
+  }
+
+  // -----------------------------------------------------------------------------
+  // Transform: core correction logic (FD vs CD functional forms)
+  // -----------------------------------------------------------------------------
+  std::tuple<vector_element_t, vector_element_t, vector_element_t>
+  ProtonEnergyLossCorrection::Transform(int const pid, int const status, int const run,
+    vector_element_t const px_in, vector_element_t const py_in,
+    vector_element_t const pz_in) const {
+    // only protons are corrected.
+    if(pid != 2212) {
+      return {px_in, py_in, pz_in};
+    }
+
+    // Detector region classification from status.
+    bool is_fd = IsFD(status);
+    bool is_cd = IsCD(status);
+    if(!is_fd && !is_cd) {
+      return {px_in, py_in, pz_in};
+    }
+
+    // Period lookup from run number. If run is not recognized, do nothing.
+    PeriodDef const* period = FindPeriod(run);
+    if(!period) {
+      return {px_in, py_in, pz_in};
+    }
+
+    // Compute kinematics from input momentum.
+    double px = (double)px_in;
+    double py = (double)py_in;
+    double pz = (double)pz_in;
+
+    double p = Pmag(px, py, pz);
+    if(p <= 0.0) {
+      return {px_in, py_in, pz_in};
+    }
+    double theta = ThetaDeg(px, py, pz); // degrees
+    double phi   = PhiDeg(px, py); // degrees in [0,360)
+
+    // Choose FD vs CD coefficients.
+    RegionCoeffs const& coeffs = is_fd ? period->fd : period->cd;
+
+    // Evaluate A/B/C polynomials at this theta.
+    double A_p = EvalPoly(coeffs.A_p, theta);
+    double B_p = EvalPoly(coeffs.B_p, theta);
+    double C_p = EvalPoly(coeffs.C_p, theta);
+
+    double A_theta = EvalPoly(coeffs.A_theta, theta);
+    double B_theta = EvalPoly(coeffs.B_theta, theta);
+    double C_theta = EvalPoly(coeffs.C_theta, theta);
+
+    double A_phi = EvalPoly(coeffs.A_phi, theta);
+    double B_phi = EvalPoly(coeffs.B_phi, theta);
+    double C_phi = EvalPoly(coeffs.C_phi, theta);
+
+    // Apply corrections.
+    //
+    // Forward Detector (FD):
+    //   p_new = p + A_p + B_p/p + C_p/p^2
+    //   theta_new = theta + A_theta + B_theta/theta + C_theta/theta^2
+    //   phi_new   = phi   + A_phi   + B_phi/phi     + C_phi/phi^2
+    //
+    // Central Detector (CD):
+    //   p_new = p + A_p + B_p*p + C_p*p^2
+    //   not all A_, B_ and C_ need be nonzero
+    double p_new     = p;
+    double theta_new = theta;
+    double phi_new   = phi;
+
+    if(is_fd) {
+      p_new += A_p + (B_p / p) + (C_p / (p * p));
+
+      if(theta != 0.0) {
+        theta_new += A_theta + (B_theta / theta) + (C_theta / (theta * theta));
+      }
+      if(phi != 0.0) {
+        phi_new += A_phi + (B_phi / phi) + (C_phi / (phi * phi));
+      }
+    } else {
+      // is_cd
+      p_new += A_p + (B_p * p) + (C_p * p * p);
+
+      if(theta != 0.0) {
+        theta_new += A_theta + (B_theta / theta) + (C_theta / (theta * theta));
+      }
+      if(phi != 0.0) {
+        phi_new += A_phi + (B_phi / phi) + (C_phi / (phi * phi));
+      }
+    }
+
+    phi_new = WrapDeg360(phi_new);
+
+    // Convert back to Cartesian using the corrected spherical variables.
+    double px_out = 0.0;
+    double py_out = 0.0;
+    double pz_out = 0.0;
+    SphericalToCartesian(p_new, theta_new, phi_new, px_out, py_out, pz_out);
+
+    return {px_out, py_out, pz_out};
+  }
+
+}