add new class TblSMSNevt

Author: TaiSakuma

atheppy/heppyresult/TblSMSNevt.py

@@ -0,0 +1,48 @@
+# Tai Sakuma <[email protected]>
+from alphatwirl.misc import mkdir_p
+from alphatwirl.misc import list_to_aligned_text
+import os
+import numpy as np
+import ROOT
+
+##__________________________________________________________________||
+class TblSMSNevt(object):
+    def __init__(self, analyzerName, fileName, outPath):
+        self.analyzerName = analyzerName
+        self.fileName = fileName
+        self.outPath = outPath
+
+    def begin(self):
+        self.rows = [('component', 'histname', 'smsmass1', 'smsmass2', 'nevt')]
+
+    def read(self, component):
+        inputPath = os.path.join(getattr(component, self.analyzerName).path, self.fileName)
+        file_ = ROOT.TFile.Open(inputPath)
+
+        histnames = ['h_Gluino_LSP', 'h_Stop_LSP', 'h_Sbottom_LSP', 'h_Squark_LSP']
+        for histname in histnames:
+            h = file_.Get(histname)
+            xaxis_lastbin_lowedge = h.GetXaxis().GetBinLowEdge(h.GetNbinsX()+1)
+            mass1 = np.concatenate((np.array([1]), np.arange(50, xaxis_lastbin_lowedge+50, step=50, dtype=int)))
+            yaxis_lastbin_lowedge = h.GetYaxis().GetBinLowEdge(h.GetNbinsY()+1)
+            mass2 = np.concatenate((np.array([1]), np.arange(50, yaxis_lastbin_lowedge+50, step=50, dtype=int)))
+            for m1 in mass1:
+                binx = h.GetXaxis().FindBin(m1)
+                ## binlowedgex = h.GetXaxis().GetBinLowEdge(binx)
+                for m2 in mass2:
+                    biny = h.GetYaxis().FindBin(m2)
+                    ## binlowedgey = h.GetYaxis().GetBinLowEdge(biny)
+                    c = h.GetBinContent(binx, biny)
+                    if c == 0:
+                        continue
+                    self.rows.append(
+                        (component.name, histname,
+                        int(m1), int(m2), int(c))
+                    )
+
+    def end(self):
+        mkdir_p(os.path.dirname(self.outPath))
+        with open(self.outPath, 'w') as f:
+            f.write(list_to_aligned_text(self.rows))
+
+##__________________________________________________________________||

atheppy/heppyresult/__init__.py

@@ -24,3 +24,4 @@
     from .EventBuilderConfigMaker import EventBuilderConfigMaker
     from .TblBranch import TblBranch
     from .TblTree import TblTree
+    from .TblSMSNevt import TblSMSNevt

tests/heppyresult/test_TblSMSNevt.py

@@ -0,0 +1,70 @@
+# Tai Sakuma <[email protected]>
+import os
+import sys
+
+import pytest
+
+try:
+    import unittest.mock as mock
+except ImportError:
+    import mock
+
+from atheppy.heppyresult import TblSMSNevt
+
+##__________________________________________________________________||
+@pytest.fixture(autouse=True)
+def mock_mkdir_p(monkeypatch):
+    ret = mock.Mock()
+    module = sys.modules['atheppy.heppyresult.TblSMSNevt']
+    monkeypatch.setattr(module, 'mkdir_p', ret)
+    return ret
+
+## @pytest.fixture(autouse=True)
+## def mock_open(monkeypatch):
+##     ret = mock.MagicMock()
+##     module = sys.modules['atheppy.heppyresult.TblSMSNevt']
+##     monkeypatch.setitem(module, 'open', ret)
+##     return ret
+
+@pytest.fixture()
+def component_t1tttt():
+    ret = mock.Mock()
+    name = 'SMS_T1tttt_madgraphMLM'
+    thisdir = os.path.dirname(os.path.realpath(__file__))
+    analyzer_path = os.path.join(
+        thisdir, 'test_data', name, 'susyParameterScanAnalyzer')
+    ret.susyParameterScanAnalyzer.path = analyzer_path
+    ret.name = name
+    return ret
+
+@pytest.fixture()
+def component_t2bb():
+    ret = mock.Mock()
+    name = 'SMS_T2bb_madgraphMLM'
+    thisdir = os.path.dirname(os.path.realpath(__file__))
+    analyzer_path = os.path.join(
+        thisdir, 'test_data', name, 'susyParameterScanAnalyzer')
+    ret.susyParameterScanAnalyzer.path = analyzer_path
+    ret.name = name
+    return ret
+
+@pytest.fixture()
+def obj():
+    ret = TblSMSNevt(
+        analyzerName='susyParameterScanAnalyzer',
+        fileName='genEvtsPerMass.root',
+        outPath='tbt_nevt_sms.txt'
+    )
+    yield ret
+
+def test_read(obj, component_t1tttt, component_t2bb):
+    obj.begin()
+    obj.read(component_t1tttt)
+    obj.read(component_t2bb)
+
+    with mock.patch('__builtin__.open', mock.mock_open()) as mock_file:
+        obj.end()
+        open().write.call_args_list
+
+
+##__________________________________________________________________||