RMGVertexConfinement.hh

// Copyright (C) 2022 Luigi Pertoldi <https://orcid.org/0000-0002-0467-2571>
//
// This program is free software: you can redistribute it and/or modify it under
// the terms of the GNU Lesser General Public License as published by the Free
// Software Foundation, either version 3 of the License, or (at your option) any
// later version.
//
// This program is distributed in the hope that it will be useful, but WITHOUT
// ANY WARRANTY; without even the implied warranty of MERCHANTABILITY or FITNESS
// FOR A PARTICULAR PURPOSE.  See the GNU Lesser General Public License for more
// details.
//
// You should have received a copy of the GNU Lesser General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
 
#ifndef _RMG_VERTEX_CONFINEMENT_HH_
#define _RMG_VERTEX_CONFINEMENT_HH_
 
#include <atomic>
#include <chrono>
#include <memory>
#include <optional>
#include <regex>
#include <string>
#include <vector>
 
#include "G4AutoLock.hh"
#include "G4GenericMessenger.hh"
#include "G4RotationMatrix.hh"
#include "G4ThreeVector.hh"
 
#include "RMGVVertexGenerator.hh"
 
class G4VPhysicalVolume;
class G4VSolid;

class RMGVertexConfinement : public RMGVVertexGenerator {
 
  public:

    enum class GeometricalSolidType {
      kSphere,
      kCylinder,
      kBox,
    };

    struct GenericGeometricalSolidData {
        GeometricalSolidType solid_type = GeometricalSolidType::kBox;
        G4ThreeVector volume_center = G4ThreeVector(0, 0, 0);
        double sphere_inner_radius = 0;
        double sphere_outer_radius = -1;
        double cylinder_inner_radius = 0;
        double cylinder_outer_radius = -1;
        double cylinder_height = -1;
        double cylinder_starting_angle = 0;
        double cylinder_spanning_angle = CLHEP::twopi;
        double box_x_length = -1;
        double box_y_length = -1;
        double box_z_length = -1;
    };

    enum class SamplingMode {
      kIntersectPhysicalWithGeometrical,
      kUnionAll,
      kSubtractGeometrical,
    };

    enum class VolumeType {
      kPhysical,
      kGeometrical,
      kUnset,
    };

    struct DepthProfile {

        enum class Type {
          kNone,              
          kExponential,       
          kTruncatedGaussian, 
          kUniform,           
        };

        Type type = Type::kNone;

        double mean = 0;

        double sigma = 0;

        double range_lo = 0;

        double range_hi = 0;

        [[nodiscard]] double Sample() const;

        void Validate() const;
    };
 
    RMGVertexConfinement();
 
    void BeginOfRunAction(const G4Run* run) override;
    void EndOfRunAction(const G4Run* run) override;

    bool GenerateVertex(G4ThreeVector& v) override;

    void AddPhysicalVolumeNameRegex(std::string name, std::string copy_nr = ".*");
 
    void AddGeometricalVolume(GenericGeometricalSolidData& data) {
      fGeomVolumeData.emplace_back(data);
    }
    void Reset();
 
    void SetSamplingMode(SamplingMode mode) { fSamplingMode = mode; }
    void SetFirstSamplingVolumeType(VolumeType type) { fFirstSamplingVolumeType = type; }
    void SetWeightByMass(bool mode) { fWeightByMass = mode; }
    void SetWeightByMassIsotope(int z, int n) {
      fWeightByMass = true;
      fWeightByMassIsotopeZ = z;
      fWeightByMassIsotopeN = n;
    }
 
    std::vector<GenericGeometricalSolidData>& GetGeometricalSolidDataList() {
      return fGeomVolumeData;
    }

    struct SampleableObject {
 
        SampleableObject() = default;
        SampleableObject(const SampleableObject&) = default;

        SampleableObject(
            G4VPhysicalVolume* physvol,
            G4RotationMatrix rot,
            G4ThreeVector trans,
            G4VSolid* solid,
            bool is_native_sampleable = false,
            bool on_surface = false
        );
 
        // NOTE: G4 volume/solid pointers should be fully owned by G4, avoid trying to delete them
        ~SampleableObject() = default;

        [[nodiscard]] bool IsInside(const G4ThreeVector& vertex) const;

        [[nodiscard]] bool Sample(
            G4ThreeVector& vertex,
            size_t max_attempts,
            bool force_containment_check,
            size_t& n_trials
        ) const;

        [[nodiscard]] bool GenerateSurfacePoint(
            G4ThreeVector& vertex,
            size_t max_attempts,
            size_t max_intersections
        ) const;
 
        // methods for the generic surface sampling
        [[nodiscard]] std::vector<G4ThreeVector> GetIntersections(
            G4ThreeVector start,
            G4ThreeVector dir
        ) const;

        void GetDirection(G4ThreeVector& dir, G4ThreeVector& pos) const;

        void ApplyDepthProfile(G4ThreeVector& local_vertex, const G4VSolid* solid) const;
 
        void RecalcMass(int z, int n);
 
        G4VPhysicalVolume* physical_volume = nullptr;
        G4VSolid* sampling_solid = nullptr;
        G4RotationMatrix rotation;
        G4ThreeVector translation;
 
        double volume = -1;
        double mass = -1;
        double surface = -1;
 
        bool surface_sample = false;
        bool native_sample = false;
        size_t max_num_intersections = 0;

        DepthProfile depth_profile = {};
    };

    struct SampleableObjectCollection {
 
        SampleableObjectCollection() = default;
        ~SampleableObjectCollection() { data.clear(); }

        [[nodiscard]] const SampleableObject& SurfaceWeightedRand() const;

        [[nodiscard]] const SampleableObject& VolumeWeightedRand(bool weight_by_mass) const;
        [[nodiscard]] bool IsInside(const G4ThreeVector& vertex) const;
 
        // emulate @c std::vector
        [[nodiscard]] size_t size() const { return data.size(); }
        SampleableObject& at(size_t i) { return data.at(i); }
        template<typename... Args> void emplace_back(Args&&... args);
        [[nodiscard]] bool empty() const { return data.empty(); }
        SampleableObject& back() { return data.back(); }
        void clear() {
          data.clear();
          total_volume = 0;
          total_mass = 0;
          total_surface = 0;
        }
        void insert(SampleableObjectCollection& other) {
          data.reserve(this->size() + other.size());
          for (size_t i = 0; i < other.size(); ++i) this->emplace_back(other.at(i));
          this->total_volume += other.total_volume;
          this->total_mass += other.total_mass;
          this->total_surface += other.total_surface;
        }
 
        void recalc_total(bool weigh_by_mass, int mass_isotope_z, int mass_istotope_n);
 
        std::vector<SampleableObject> data;
        double total_volume = 0;
        double total_mass = 0;
        double total_surface = 0;
    };
 
  private:
 
    void InitializePhysicalVolumes();
    void InitializeGeometricalVolumes(bool use_excluded_volumes);
    bool ActualGenerateVertex(G4ThreeVector& v);
 
    std::vector<std::string> fPhysicalVolumeNameRegexes;
    std::vector<std::string> fPhysicalVolumeCopyNrRegexes;
 
    std::vector<GenericGeometricalSolidData> fGeomVolumeData;
    std::vector<GenericGeometricalSolidData> fExcludedGeomVolumeData;
 
    static G4Mutex fGeometryMutex;
    // the final geometry data is shared between all threads and protected by fGeometryMutex.
    // this is to prevent problems in G4SolidStore, which is apparently not safe to mutate from
    // multiple threads. note that some operations that just appear to read static data might also
    // mutate the G4SolidStore temporarily, e.g. G4SubstractionSolid::GetCubicVolume()
    static SampleableObjectCollection fPhysicalVolumes;
    static SampleableObjectCollection fGeomVolumeSolids;
    static SampleableObjectCollection fExcludedGeomVolumeSolids;
    // cached merged collection used by the default kUnionAll mode.
    static SampleableObjectCollection fUnionVolumes;
 
    // solids allocated by this class (bounding boxes and geometrical sampling shapes). They are
    // registered in the G4SolidStore and must be de-registered and freed on Reset().
    static std::vector<G4VSolid*> fOwnedSolids;
 
    static std::atomic<bool> fVolumesInitialized;
 
    SamplingMode fSamplingMode = SamplingMode::kUnionAll;
    VolumeType fFirstSamplingVolumeType = VolumeType::kUnset;
    bool fWeightByMass = false;
    int fWeightByMassIsotopeZ = 0;
    int fWeightByMassIsotopeN = 0;
 
    bool fOnSurface = false;
    bool fForceContainmentCheck = false;
    bool fLastSolidExcluded = false;
    size_t fSurfaceSampleMaxIntersections = 0;

    DepthProfile fDepthProfile = {};
 
    // counters used for the current run.
    size_t fTrials = 0;
    std::chrono::nanoseconds fVertexGenerationTime{};
 
    std::vector<std::unique_ptr<G4GenericMessenger>> fMessengers;
    void SetSamplingModeString(std::string mode);
    void SetFirstSamplingVolumeTypeString(std::string type);
 
    void AddGeometricalVolumeString(std::string solid);
    void AddExcludedGeometricalVolumeString(std::string solid);
 
    GenericGeometricalSolidData& SafeBack(
        std::optional<GeometricalSolidType> solid_type = std::nullopt
    );
 
    // FIXME: there is no easy way to set the position vector all at once with
    // G4GenericMessenger. Only ::DeclarePropertyWithUnit() accepts vectors and
    // one cannot call functions with more than 2 arguments (3 coordinated + 1
    // units needed). Ugly!
    void SetGeomVolumeCenter(const G4ThreeVector& v) { this->SafeBack().volume_center = v; }
    void SetGeomVolumeCenterX(double x) { this->SafeBack().volume_center.setX(x); }
    void SetGeomVolumeCenterY(double y) { this->SafeBack().volume_center.setY(y); }
    void SetGeomVolumeCenterZ(double z) { this->SafeBack().volume_center.setZ(z); }
 
    void SetGeomSphereInnerRadius(double r) {
      this->SafeBack(GeometricalSolidType::kSphere).sphere_inner_radius = r;
    }
    void SetGeomSphereOuterRadius(double r) {
      this->SafeBack(GeometricalSolidType::kSphere).sphere_outer_radius = r;
    }
 
    void SetGeomCylinderInnerRadius(double r) {
      this->SafeBack(GeometricalSolidType::kCylinder).cylinder_inner_radius = r;
    }
    void SetGeomCylinderOuterRadius(double r) {
      this->SafeBack(GeometricalSolidType::kCylinder).cylinder_outer_radius = r;
    }
    void SetGeomCylinderHeight(double h) {
      this->SafeBack(GeometricalSolidType::kCylinder).cylinder_height = h;
    }
    void SetGeomCylinderStartingAngle(double a) {
      this->SafeBack(GeometricalSolidType::kCylinder).cylinder_starting_angle = a;
    }
    void SetGeomCylinderSpanningAngle(double a) {
      this->SafeBack(GeometricalSolidType::kCylinder).cylinder_spanning_angle = a;
    }
 
    void SetGeomBoxXLength(double x) {
      this->SafeBack(GeometricalSolidType::kBox).box_x_length = x;
    }
    void SetGeomBoxYLength(double y) {
      this->SafeBack(GeometricalSolidType::kBox).box_y_length = y;
    }
    void SetGeomBoxZLength(double z) {
      this->SafeBack(GeometricalSolidType::kBox).box_z_length = z;
    }
 
    void DefineCommands();
 
    void SetDepthProfileTypeString(std::string type);
    void SetDepthProfileMean(double mean) { fDepthProfile.mean = mean; }
    void SetDepthProfileSigma(double sigma) { fDepthProfile.sigma = sigma; }
    void SetDepthProfileRangeLo(double lo) { fDepthProfile.range_lo = lo; }
    void SetDepthProfileRangeHi(double hi) { fDepthProfile.range_hi = hi; }
};
 
#endif
 
// vim: tabstop=2 shiftwidth=2 expandtab