From 2346740b749cdef0968879327c33c378d70ff50b Mon Sep 17 00:00:00 2001
From: Alvaro Ezquerro <aezquerro@unizar.es>
Date: Sat, 21 Oct 2023 16:22:53 +0200
Subject: [PATCH] Adding methods to set any split distribution 'by hand'.

---
 .../analysis/inc/TRestDataSetGainMap.h        | 23 +++++----
 .../analysis/src/TRestDataSetGainMap.cxx      | 49 +++++++++++++++++--
 2 files changed, 56 insertions(+), 16 deletions(-)

diff --git a/source/framework/analysis/inc/TRestDataSetGainMap.h b/source/framework/analysis/inc/TRestDataSetGainMap.h
index e2ac6ef2e..3f7dbe652 100644
--- a/source/framework/analysis/inc/TRestDataSetGainMap.h
+++ b/source/framework/analysis/inc/TRestDataSetGainMap.h
@@ -132,10 +132,6 @@ class TRestDataSetGainMap : public TRestMetadata {
         std::vector<std::vector<TGraph*>> fSegLinearFit = {};
 
        public:
-        void SetSplitX();
-        void SetSplitY();
-        void SetSplits();
-
         void AddPeak(const double& energyPeak, const TVector2& rangePeak = TVector2(0, 0)) {
             fEnergyPeaks.push_back(energyPeak);
             fRangePeaks.push_back(rangePeak);
@@ -184,16 +180,19 @@ class TRestDataSetGainMap : public TRestMetadata {
         }
         void SetCalibrationRange(const TVector2& calibrationRange) { fCalibRange = calibrationRange; }
         void SetNBins(const Int_t& nBins) { fNBins = nBins; }
-
-        void SetNumberOfSegmentsX(const Int_t& numberOfSegmentsX) {
-            fNumberOfSegmentsX = numberOfSegmentsX;
-            SetSplitX();
-        }
-        void SetNumberOfSegmentsY(const Int_t& numberOfSegmentsY) {
-            fNumberOfSegmentsY = numberOfSegmentsY;
-            SetSplitY();
+        void SetSplitX();
+        void SetSplitY();
+        void SetSplitX(const std::set<double>& splitX);
+        void SetSplitY(const std::set<double>& splitY);
+        void SetSplits();
+        void SetSplits(const std::set<double>& splitXandY) {
+            SetSplitX(splitXandY);
+            SetSplitY(splitXandY);
         }
 
+        void SetNumberOfSegmentsX(const Int_t& numberOfSegmentsX) { fNumberOfSegmentsX = numberOfSegmentsX; }
+        void SetNumberOfSegmentsY(const Int_t& numberOfSegmentsY) { fNumberOfSegmentsY = numberOfSegmentsY; }
+
         void SetDataSetFileName(const std::string& dataSetFileName) { fDataSetFileName = dataSetFileName; }
         void SetReadoutRange(const TVector2& readoutRange) { fReadoutRange = readoutRange; }
         void SetZeroPoint(const bool& ZeroPoint) { fZeroPoint = ZeroPoint; }
diff --git a/source/framework/analysis/src/TRestDataSetGainMap.cxx b/source/framework/analysis/src/TRestDataSetGainMap.cxx
index 6872c0f13..0212e179b 100644
--- a/source/framework/analysis/src/TRestDataSetGainMap.cxx
+++ b/source/framework/analysis/src/TRestDataSetGainMap.cxx
@@ -527,12 +527,31 @@ void TRestDataSetGainMap::Module::SetSplits() {
 /// It uses the number of segments and the readout range to define the
 /// edges of the segments.
 void TRestDataSetGainMap::Module::SetSplitX() {
-    fSplitX.clear();
+    if (fNumberOfSegmentsX < 1) {
+        RESTError << "SetSplitX: fNumberOfSegmentsX must be >=1." << p->RESTendl;
+        return;
+    }
+    std::set<double> split;
     for (int i = 0; i <= fNumberOfSegmentsX; i++) {  // <= so that the last segment is included
         double x =
             fReadoutRange.X() + ((fReadoutRange.Y() - fReadoutRange.X()) / (float)fNumberOfSegmentsX) * i;
-        fSplitX.insert(x);
+        split.insert(x);
+    }
+    SetSplitX(split);
+}
+
+void TRestDataSetGainMap::Module::SetSplitX(const std::set<double>& splitX) {
+    if (splitX.size() < 2) {
+        RESTError << "SetSplitX: split size must be >=2 (start and end of range must be included)."
+                  << p->RESTendl;
+        return;
     }
+    if (!fSlope.empty())
+        RESTWarning << "SetSplitX: changing split but current gain map and calibration paremeters correspond "
+                       "to previous splitting. Use GenerateGainMap() to update them."
+                    << p->RESTendl;
+    fSplitX = splitX;
+    fNumberOfSegmentsX = fSplitX.size() - 1;
 }
 /////////////////////////////////////////////
 /// \brief Function to set the class members for segmentation of
@@ -541,12 +560,31 @@ void TRestDataSetGainMap::Module::SetSplitX() {
 /// It uses the number of segments and the readout range to define the
 /// edges of the segments.
 void TRestDataSetGainMap::Module::SetSplitY() {
-    fSplitY.clear();
+    if (fNumberOfSegmentsY < 1) {
+        RESTError << "SetSplitY: fNumberOfSegmentsY must be >=1." << p->RESTendl;
+        return;
+    }
+    std::set<double> split;
     for (int i = 0; i <= fNumberOfSegmentsY; i++) {  // <= so that the last segment is included
         double y =
             fReadoutRange.X() + ((fReadoutRange.Y() - fReadoutRange.X()) / (float)fNumberOfSegmentsY) * i;
-        fSplitY.insert(y);
+        split.insert(y);
     }
+    SetSplitY(split);
+}
+
+void TRestDataSetGainMap::Module::SetSplitY(const std::set<double>& splitY) {
+    if (splitY.size() < 2) {
+        RESTError << "SetSplitY: split size must be >=2 (start and end of range must be included)."
+                  << p->RESTendl;
+        return;
+    }
+    if (!fSlope.empty())
+        RESTWarning << "SetSplitY: changing split but current gain map and calibration paremeters correspond "
+                       "to previous splitting. Use GenerateGainMap() to update them."
+                    << p->RESTendl;
+    fSplitY = splitY;
+    fNumberOfSegmentsY = fSplitY.size() - 1;
 }
 
 /////////////////////////////////////////////
@@ -585,6 +623,9 @@ void TRestDataSetGainMap::Module::GenerateGainMap() {
 
     SetSplits();
 
+    if (fSplitX.empty()) SetSplitX();
+    if (fSplitY.empty()) SetSplitY();
+
     //--- Get the calibration range if not provided (default is 0,0) ---
     if (fCalibRange.X() >= fCalibRange.Y()) {
         // Get spectrum for this file