MaterialScan.h

/**************************************************************************
 * basf2 (Belle II Analysis Software Framework)                           *
 * Author: The Belle II Collaboration                                     *
 *                                                                        *
 * See git log for contributors and copyright holders.                    *
 * This file is licensed under LGPL-3.0, see LICENSE.md.                  *
 **************************************************************************/
 
#pragma once
 
#include <framework/core/Module.h>
#include <framework/gearbox/Unit.h>
 
#include <optional>
#include <set>
#include <string>
#include <vector>
 
#include <G4UserSteppingAction.hh>
#include <G4ThreeVector.hh>
 
#include <TFile.h>
class TH1D;
class TH2D;
 
namespace Belle2 {
  class MaterialScanBase: public G4UserSteppingAction {
  public:
    MaterialScanBase(TFile* rootFile, const std::string& name, const std::string& axisLabel):
      m_rootFile(rootFile), m_name(name), m_axisLabel(axisLabel)
    {
      m_rootFile->mkdir(name.c_str());
    }
    std::string getName() const { return m_name; }
    virtual int getNRays() const = 0;
 
    virtual bool createNext(G4ThreeVector& origin, G4ThreeVector& direction) = 0;
 
  protected:
    bool checkStep(const G4Step* step);
    TFile* m_rootFile;
    std::string m_name;
    std::string m_axisLabel;
  private:
    static constexpr double c_zeroTolerance = 1e-6;
    static constexpr int c_maxZeroStepsNudge = 10;
    static constexpr int c_maxZeroStepsKill = 20;
    int m_zeroSteps {0};
  };
 
  class MaterialScan2D: public MaterialScanBase {
  public:
    struct ScanParams {
      ScanParams(): nU(0), nV(0), minU(0), maxU(0), minV(0), maxV(0), maxDepth(-1), splitByMaterials(false) {}
      int nU;
      int nV;
      double minU;
      double maxU;
      double minV;
      double maxV;
      double maxDepth;
      std::vector<std::string> ignoredMaterials;
      bool splitByMaterials;
    };
 
    MaterialScan2D(TFile* rootFile, const std::string& name, const std::string& axisLabel, const ScanParams& params);
 
    bool createNext(G4ThreeVector& origin, G4ThreeVector& direction) override;
 
    int getNRays() const override { return m_params.nU * m_params.nV; }
 
    void UserSteppingAction(const G4Step* step) override;
  protected:
    virtual void getRay(G4ThreeVector& origin, G4ThreeVector& direction) = 0;
 
    TH2D* getHistogram(const std::string& name);
 
    void fillValue(const std::string& name, double value);
 
    ScanParams m_params;
    double m_curU;
    double m_stepU;
    double m_curV;
    double m_stepV;
    double m_curDepth;
    std::map<std::string, std::unique_ptr<TH2D>> m_regions;
  };
 
  class MaterialScanSpherical: public MaterialScan2D {
  public:
    MaterialScanSpherical(TFile* rootFile, const G4ThreeVector& origin, const ScanParams& params, bool doCosTheta):
      MaterialScan2D(rootFile, "Spherical", doCosTheta ? "cos(#theta);#phi [deg]" : "#theta [deg];#phi [deg]", params), m_origin(origin),
      m_doCosTheta(doCosTheta)
    {
      if (doCosTheta) {
        m_params.minU = cos(m_params.minU * Unit::deg);
        m_params.maxU = cos(m_params.maxU * Unit::deg);
        if (m_params.minU > m_params.maxU) std::swap(m_params.minU, m_params.maxU);
        m_stepU = (m_params.maxU - m_params.minU) / m_params.nU;
        m_curU  = m_params.minU - m_stepU / 2.;
      }
    }
  protected:
    void getRay(G4ThreeVector& origin, G4ThreeVector& direction) override;
 
    G4ThreeVector m_origin;
 
    bool m_doCosTheta;
  };
 
  class MaterialScanPlanar: public MaterialScan2D {
  public:
    MaterialScanPlanar(TFile* rootFile, const G4ThreeVector& origin, const G4ThreeVector& dirU, const G4ThreeVector& dirV,
                       const ScanParams& params):
      MaterialScan2D(rootFile, "Planar", "u [cm];v [cm]", params), m_origin(origin), m_dirU(dirU.unit()), m_dirV(dirV.unit()),
      m_dirW(m_dirU.cross(m_dirV))
    {
    }
  protected:
    void getRay(G4ThreeVector& origin, G4ThreeVector& direction) override;
  protected:
    G4ThreeVector m_origin;
    G4ThreeVector m_dirU;
    G4ThreeVector m_dirV;
    G4ThreeVector m_dirW;
  };
 
  class MaterialScanRay: public MaterialScanBase {
  public:
    MaterialScanRay(TFile* rootFile, const G4ThreeVector& origin, const G4ThreeVector& dir, double opening,
                    int count, double sampleDepth, double maxDepth, bool splitByMaterials,
                    const std::vector<std::string>& ignoredMaterials):
      MaterialScanBase(rootFile, "Ray", "ray length [cm]; material budget [X_0/cm]"),
      m_ignoredMaterials(ignoredMaterials.begin(), ignoredMaterials.end()),
      m_origin(origin), m_dir(dir), m_opening(opening), m_sampleDepth(sampleDepth),
      m_maxDepth(maxDepth), m_count(std::max(1, count)), m_splitByMaterials(splitByMaterials)
    {
      if (m_opening <= 0) m_count = 1;
    }
 
    int getNRays() const override { return m_count + 1; }
    void UserSteppingAction(const G4Step* step) override;
    bool createNext(G4ThreeVector& origin, G4ThreeVector& direction) override;
  private:
    TH1D* getHistogram(const std::string& name);
    void fillValue(const std::string& name, double value, double steplength);
 
    std::set<std::string> m_ignoredMaterials;
    std::map<std::string, std::unique_ptr<TH1D>> m_regions;
    G4ThreeVector m_origin;
    G4ThreeVector m_dir;
    double m_opening;
    double m_sampleDepth;
    double m_maxDepth{0};
    double m_curDepth{0};
    double m_scanDepth{0};
    int m_count;
    int m_curRay{ -1};
    bool m_splitByMaterials;
  };
 
  class MaterialScanModule : public Module {
 
  public:
 
    /* The constructor of the module.
     * Sets the description and the parameters of the module.
     */
    MaterialScanModule();
 
    void initialize() override;
 
    void beginRun() override;
  private:
    G4ThreeVector getAxis(char name);
 
    TFile* m_rootFile;
    std::string m_filename;
    std::string m_planeName;
    MaterialScan2D::ScanParams m_spherical;
    MaterialScan2D::ScanParams m_planar;
    std::vector<double> m_customPlane;
    std::vector<double> m_sphericalOrigin;
    std::vector<double> m_rayOrigin{0, 0, 0};
    std::vector<std::string> m_rayIgnoredMaterials;
    std::optional<std::vector<double>> m_rayDirection;
    double m_rayTheta{0};
    double m_rayPhi{0};
    double m_rayOpening{0};
    double m_raySampleDepth{0.1};
    double m_rayMaxDepth{1000};
    unsigned int m_rayCount{0};
    bool m_doSpherical;
    bool m_doPlanar;
    bool m_doCosTheta;
    bool m_doRay;
    bool m_raySplitByMaterials{false};
  };
}