Merge pull request #41 from blinkseb/lines

Add ability to draw lines over plots
cp3-llbb · Jan 15, 2016 · 1abc17e · 1abc17e
2 parents c7a0bf0 + b2b5f34
commit 1abc17e
Show file tree

Hide file tree

Showing 8 changed files with 154 additions and 1 deletion.
diff --git a/include/types.h b/include/types.h
@@ -51,6 +51,12 @@ namespace plotIt {
       return node.as<bool>() ? True : False;
   }
 
+  enum Orientation {
+      UNSPECIFIED,
+      HORIZONTAL,
+      VERTICAL
+  };
+
   struct Summary {
     float n_events = 0;
     float n_events_error = 0;
@@ -200,6 +206,27 @@ namespace plotIt {
     Point position;
   };
 
+  struct Line {
+    Point start;
+    Point end;
+
+    bool operator==(const Line& other) {
+      return ((start == other.start) && (end == other.end));
+    }
+
+    bool valid() const {
+      return start.valid() && end.valid();
+    }
+
+    Line() = default;
+    Line(std::initializer_list<Point> c) {
+      assert(c.size() == 2);
+      start = *c.begin();
+      end = *(c.begin() + 1);
+    }
+    Line(const YAML::Node& node, Orientation);
+  };
+
   struct Plot {
     std::string name;
     std::string output_suffix;
@@ -231,7 +258,6 @@ namespace plotIt {
     std::string draw_string;  // Only used in tree mode
     std::string selection_string;  // Only used in tree mode
 
-
     std::vector<std::string> save_extensions = {"pdf"};
 
     bool show_ratio = false;
@@ -269,6 +295,8 @@ namespace plotIt {
     int yields_table_order = 0;
 
     bool is_rescaled = false;
+
+    std::vector<Line> lines;
 
     void print() {
       std::cout << "Plot '" << name << "'" << std::endl;
@@ -321,6 +349,10 @@ namespace plotIt {
     int16_t ratio_fit_error_fill_color = 42;
     int16_t ratio_fit_error_fill_style = 1001;
 
+    int16_t line_color = 0;
+    int16_t line_width = 1;
+    int16_t line_style = 1;
+
     std::vector<Label> labels;
 
     std::string experiment = "CMS";

diff --git a/include/utilities.h b/include/utilities.h
@@ -129,6 +129,17 @@ namespace plotIt {
 
   void setRange(TObject* object, const Range& x_range, const Range& y_range);
 
+  template<class T>
+    Range getXRange(T* object) {
+      Range range;
+      range.start = object->GetXaxis()->GetBinLowEdge(object->GetXaxis()->GetFirst());
+      range.end = object->GetXaxis()->GetBinUpEdge(object->GetXaxis()->GetLast());
+
+      return range;
+    }
+
+  Range getXRange(TObject* object);
+
   template<class T>
   float getPositiveMinimum(T* object) {
     return object->GetMinimum(0);

diff --git a/src/TH1Plotter.cc b/src/TH1Plotter.cc
@@ -4,6 +4,7 @@
 #include <TF1.h>
 #include <TFitResult.h>
 #include <TLatex.h>
+#include <TLine.h>
 #include <TObject.h>
 #include <TPave.h>
 #include <TVirtualFitter.h>
@@ -398,6 +399,39 @@ namespace plotIt {
         blinded_area->Draw("same");
     }
 
+    // Draw all the requested lines for this plot
+    auto resolveLine = [&](Line& line) {
+        Range x_range = getXRange(toDraw[0].first);
+
+        float y_range_start = gPad->GetUymin();
+        float y_range_end = gPad->GetUymax();
+
+        if (std::isnan(line.start.x))
+            line.start.x = x_range.start;
+
+        if (std::isnan(line.start.y))
+            line.start.y = y_range_start;
+
+        if (std::isnan(line.end.x))
+            line.end.x = x_range.end;
+
+        if (std::isnan(line.end.y))
+            line.end.y = y_range_end;
+    };
+
+    for (Line& line: plot.lines) {
+      resolveLine(line);
+
+      std::shared_ptr<TLine> l(new TLine(line.start.x, line.start.y, line.end.x, line.end.y));
+      TemporaryPool::get().add(l);
+
+      l->SetLineColor(m_plotIt.getConfiguration().line_color);
+      l->SetLineWidth(m_plotIt.getConfiguration().line_width);
+      l->SetLineStyle(m_plotIt.getConfiguration().line_style);
+
+      l->Draw("same");
+    }
+
     // Redraw only axis
     toDraw[0].first->Draw("axis same");
 

diff --git a/src/plotIt.cc b/src/plotIt.cc
@@ -211,6 +211,15 @@ namespace plotIt {
       if (node["blinded-range-fill-style"])
         m_config.blinded_range_fill_style = node["blinded-range-fill-style"].as<uint16_t>();
 
+      if (node["line-color"])
+        m_config.line_color = loadColor(node["line-color"]);
+
+      if (node["line-width"])
+        m_config.line_width = node["line-style"].as<int16_t>();
+
+      if (node["line-style"])
+        m_config.line_style = node["line-style"].as<int16_t>();
+
       if (node["labels"]) {
         YAML::Node labels = node["labels"];
         m_config.labels = parseLabelsNode(labels);
@@ -525,6 +534,24 @@ namespace plotIt {
 
       if (node["yields-table-order"])
         plot.yields_table_order = node["yields-table-order"].as<int>();
+
+      if (node["vertical-lines"]) {
+        for (const auto& line: node["vertical-lines"]) {
+          plot.lines.push_back(Line(line, VERTICAL));
+        }
+      }
+
+      if (node["horizontal-lines"]) {
+        for (const auto& line: node["horizontal-lines"]) {
+          plot.lines.push_back(Line(line, HORIZONTAL));
+        }
+      }
+
+      if (node["lines"]) {
+        for (const auto& line: node["lines"]) {
+          plot.lines.push_back(Line(line, UNSPECIFIED));
+        }
+      }
 
       // Handle log
       std::vector<bool> logs_x;

diff --git a/src/types.cc b/src/types.cc
@@ -81,4 +81,26 @@ namespace plotIt {
     if (node["marker-size"])
       marker_size = node["marker-size"].as<float>();
   }
+
+  Line::Line(const YAML::Node& node, Orientation orientation) {
+      auto NaN = std::numeric_limits<float>::quiet_NaN();
+
+      if (orientation == UNSPECIFIED) {
+          auto points = node.as<std::vector<Point>>();
+          if (points.size() != 2)
+              throw nullptr; // FIXME
+
+          start = points[0];
+          end = points[1];
+      } else {
+          float value = node.as<float>();
+          if (orientation == HORIZONTAL) {
+              start = {NaN, value};
+              end = {NaN, value};
+          } else {
+              start = {value, NaN};
+              end = {value, NaN};
+          }
+      }
+  }
 }
diff --git a/src/utilities.cc b/src/utilities.cc
@@ -163,6 +163,12 @@ namespace plotIt {
     CAST_AND_CALL(object, setRange, x_range, y_range);
   }
 
+  Range getXRange(TObject* object) {
+    CAST_AND_RETURN(object, getXRange);
+
+    return Range();
+  }
+
   float getPositiveMinimum(TObject* object) {
       if (dynamic_cast<TH1*>(object))
           return getPositiveMinimum(dynamic_cast<TH1*>(object));

diff --git a/test/golden/default_configuration_lines.pdf b/test/golden/default_configuration_lines.pdf
diff --git a/test/unit_tests/tests.py b/test/unit_tests/tests.py
@@ -227,3 +227,24 @@ def test_default_legend_columns(self):
                 os.path.join(self.output_folder.name, 'histo1.pdf'),
                 get_golden_file('default_configuration_2columns_samplesordering_legend.pdf')
                 )
+
+    def test_lines(self):
+        configuration = get_configuration()
+
+        configuration['configuration']['line-color'] = "#1685f2"
+        configuration['configuration']['line-width'] = 3
+        configuration['configuration']['line-style'] = 4
+
+        configuration['plots']['histo1']['vertical-lines'] = [3, 7]
+        configuration['plots']['histo1']['horizontal-lines'] = [80]
+        configuration['plots']['histo1']['lines'] = [
+                [[3, 20], [8, 90]],
+                [[5.5, float('nan')], [5.5, float('nan')]]
+                ]
+
+        self.run_plotit(configuration)
+
+        self.compare_images(
+                os.path.join(self.output_folder.name, 'histo1.pdf'),
+                get_golden_file('default_configuration_lines.pdf')
+                )