<commit_before>import os
import unittest
from copy import deepcopy
from functools import partial

import six
from mock import MagicMock, patch, Mock

import ROOT
from PyAnalysisTools.AnalysisTools import SystematicsAnalyser as sa, MLHelper
from PyAnalysisTools.AnalysisTools.LimitHelpers import Yield
from PyAnalysisTools.AnalysisTools.SystematicsAnalyser import SystematicsCategory
from PyAnalysisTools.PlottingUtils.BasePlotter import BasePlotter
from PyAnalysisTools.PlottingUtils.PlotConfig import PlotConfig
from PyAnalysisTools.base.FileHandle import FileHandle
from PyAnalysisTools.base.ProcessConfig import Process
from PyAnalysisTools.base.YAMLHandle import YAMLLoader as yl
from .Mocks import hist
import numpy as np

cwd = os.path.dirname(__file__)


def entity_patch(*args, **kwargs):
    return args, kwargs


@classmethod
def mock_read_yaml(*args):
    return {'foo': 'bar'}


@classmethod
def mock_read_yaml2(*args):
    return {'foo': {'bar': 'foobar'}}


class TestSystematicsCategory(unittest.TestCase):
    def test_ctor(self):
        category = sa.SystematicsCategory()
        self.assertEqual("total", category.name)
        self.assertEqual("", category.systematics)
        self.assertEqual(ROOT.kBlue, category.color)

    def test_contain_check(self):
        category = sa.SystematicsCategory(systematics=['foo', 'bar'])
        self.assertTrue(category.contains_syst('total'))
        self.assertTrue(category.contains_syst('foo'))
        self.assertTrue(category.contains_syst('missing'))
        category = sa.SystematicsCategory(name='single', systematics=['foo', 'bar'])
        self.assertTrue(category.contains_syst('foo'))
        self.assertFalse(category.contains_syst('missing'))


class TestSystematics(unittest.TestCase):
    def setUp(self):
        self.syst = sa.Systematic('foo', variation='updown')
        self.scale_syst = sa.Systematic('foo', type='scale', variation='single:up')

    def tearDown(self):
        del self.syst
        del self.scale_syst

    def test_ctor_default(self):
        syst = sa.Systematic('foo')
        self.assertEqual('', syst.prefix)
        self.assertFalse(syst.symmetrise)
        self.assertIsNone(syst.symmetrise_option)
        self.assertIsNone(syst.call)
        self.assertIsNone(syst.affects)
        self.assertIsNone(syst.samples)
        self.assertIsNone(syst.group)
        self.assertIsNone(syst.title)
        self.assertIsNone(syst.type)
        self.assertIsNone(syst.expand)
        self.assertIsNone(syst.hist_name)
        self.assertIsNone(syst.variation)
        self.assertIsNone(syst.envelope)

    def test_ctor_scale(self):
        self.assertEqual('weight_', self.scale_syst.prefix)

    def test_str_op(self):
        self.assertEqual('Systematic uncertainty: foo', self.syst.__str__())

    def test_repr_op(self):
        self.assertEqual('Systematic uncertainty: foo\n', self.syst.__repr__())

    def test_get_variations_updown(self):
        self.assertEqual(['foo__1up', 'foo__1down'], self.syst.get_variations())

    def test_get_variations_scale(self):
        self.assertEqual(['weight_foo__1up'], self.scale_syst.get_variations())

    def test_get_variations_empty(self):
        self.assertEqual([], sa.Systematic('foo').get_variations())

    def test_get_variations_unsupported(self):
        syst = sa.Systematic('foo', type='scale', variation='foo', symmetrise=True)
        self.assertEqual([], syst.get_variations())

    def test_get_symmetrised_name_nosym(self):
        self.assertIsNone(self.syst.get_symmetrised_name())

    def test_get_symmetrised_name_invalid(self):
        syst = sa.Systematic('foo', type='scale', variation='single:foo', symmetrise=True)
        self.assertIsNone(syst.get_symmetrised_name())

    def test_get_symmetrised_name_up(self):
        scale_syst = sa.Systematic('foo', type='scale', variation='single:up', symmetrise=True)
        self.assertEqual(['weight_foo__1down'], scale_syst.get_symmetrised_name())

    def test_get_symmetrised_name_down(self):
        scale_syst = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        self.assertEqual(['weight_foo__1up'], scale_syst.get_symmetrised_name())

    def test_eq(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        syst2 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        self.assertEqual(syst1, syst2)

    def test_inequal(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        syst2 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True)
        self.assertNotEqual(syst1, syst2)

    def test_inequal_missing_key(self):
        syst1 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True, foo='bar')
        syst2 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True)
        self.assertNotEqual(syst1, syst2)

    def test_inequal_missing_key_swap(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        syst2 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True, foo='bar')
        self.assertNotEqual(syst1, syst2)

    def test_inequal_type(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        self.assertNotEqual(1., syst1)


syst = sa.Systematic('foo', bar='foobar', type='shape', variation='updown')
scale_syst = sa.Systematic('foo_scale', bar='foobar_scale', type='scale', variation='updown')


class TestSystematicsAnalyser(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_ctor_default(self):
        analyser = sa.SystematicsAnalyser(xs_handle='foo')
        self.assertEqual(1, analyser.ncpu)
        self.assertFalse(analyser.dump_hists)
        self.assertFalse(analyser.cluster_mode)
        self.assertFalse(analyser.store_abs_syst)
        self.assertIsNone(analyser.systematics)
        self.assertEqual({}, analyser.systematic_hists)
        self.assertEqual({}, analyser.systematic_variations)
        self.assertEqual({}, analyser.total_systematics)
        self.assertEqual([], analyser.scale_uncerts)
        self.assertEqual([], analyser.shape_uncerts)
        self.assertEqual([], analyser.file_handles)
        self.assertEqual([], analyser.custom_uncerts)
        self.assertEqual('foo', analyser.xs_handle)
        self.assertTrue(analyser.disable)

    @patch.object(yl, 'read_yaml', mock_read_yaml2)
    def test_parse_syst_cfg(self):
        systs = sa.SystematicsAnalyser.parse_syst_config('foo')
        self.assertEqual([sa.Systematic('foo', bar='foobar')], systs)

    @patch.object(FileHandle, 'get_object_by_name',
                  lambda *args: ROOT.TH1F('h_syst', '', 10, 0., 10.))
    def test_load_dumped_hists(self):
        fh = FileHandle(file_name='foo', process=Process('foo', dataset_info=None))
        pc, process, hist = sa.SystematicsAnalyser.load_dumped_hist((fh, PlotConfig()), 'foo')
        self.assertEqual(PlotConfig(), pc)
        self.assertEqual('foo', process.process_name)
        self.assertEqual('h_syst', hist.GetName())

    @patch.object(FileHandle, 'get_object_by_name',
                  lambda *args: None)
    def test_load_dumped_hists_tuple(self):
        fh = FileHandle(file_name='foo', process=Process('foo', dataset_info=None))
        self.assertEqual([(PlotConfig(), None, None)], sa.SystematicsAnalyser(xs_handle='foo').load_dumped_hists([fh],
                                                                                                         [PlotConfig()],
                                                                                                         'foo'))

    def test_load_dumped_hists_theory(self):
        fh = FileHandle(file_name='foo', process=Process('foo', dataset_info=None))
        pc, process, hist = sa.SystematicsAnalyser.load_dumped_hist((fh, PlotConfig()),
                                                                    'weight_pdf_uncert_MUR0p5_MUF0p5_PDF261000')
        self.assertIsNone(pc)
        self.assertIsNone(process)
        self.assertIsNone(hist)

    def test_retrieve_sys_hists_no_input_files(self):
        fh = FileHandle(file_name='foo', process=Process('data', dataset_info=None))
        analyser = sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo')
        self.assertEqual(1, len(analyser.file_handles))
        self.assertIsNone(analyser.retrieve_sys_hists())

    def test_retrieve_sys_hists_disable(self):
        self.assertIsNone(sa.SystematicsAnalyser(xs_handle='foo').retrieve_sys_hists())

    def test_retrieve_sys_hists_data_only(self):
        fh = FileHandle(file_name='foo', process=Process('foo', dataset_info=None))
        self.assertIsNone(sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo').retrieve_sys_hists())

    @patch.object(BasePlotter, 'read_histograms', lambda *args, **kwargs: [(PlotConfig(), 'foo',
                                                                            ROOT.TH1F('h_syst', '', 10, 0., 10.))])
    @patch.object(BasePlotter, 'apply_lumi_weights', lambda *args, **kwargs: None)
    @patch.object(sa.SystematicsAnalyser, 'parse_syst_config', lambda *args: [syst, scale_syst])
    @patch.object(BasePlotter, 'merge_histograms', lambda _: None)
    def test_retrieve_sys_hists(self):
        fh = FileHandle(file_name='foo', process=Process('foo_311011', dataset_info=None))
        print(sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo',
                                     plot_configs=[PlotConfig()]).retrieve_sys_hists())

    def test_get_variation_for_process_no_pc(self):
        dummy_process = Process('foo', dataset_info=None)
        fh = FileHandle(file_name='foo', process=dummy_process)
        analyser = sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo')
        analyser.systematic_hists = {'my_syst_foo': {PlotConfig(): [hist]}}<commit_msg>fix tests<commit_after>import os
import unittest
from copy import deepcopy
from functools import partial

import six
from mock import MagicMock, patch, Mock

import ROOT
from PyAnalysisTools.AnalysisTools import SystematicsAnalyser as sa, MLHelper
from PyAnalysisTools.AnalysisTools.LimitHelpers import Yield
from PyAnalysisTools.AnalysisTools.SystematicsAnalyser import SystematicsCategory
from PyAnalysisTools.PlottingUtils.BasePlotter import BasePlotter
from PyAnalysisTools.PlottingUtils.PlotConfig import PlotConfig
from PyAnalysisTools.base.FileHandle import FileHandle
from PyAnalysisTools.base.ProcessConfig import Process
from PyAnalysisTools.base.YAMLHandle import YAMLLoader as yl
from .Mocks import hist
import numpy as np

cwd = os.path.dirname(__file__)


def entity_patch(*args, **kwargs):
    return args, kwargs


@classmethod
def mock_read_yaml(*args):
    return {'foo': 'bar'}


@classmethod
def mock_read_yaml2(*args):
    return {'foo': {'bar': 'foobar'}}


class TestSystematicsCategory(unittest.TestCase):
    def test_ctor(self):
        category = sa.SystematicsCategory()
        self.assertEqual("total", category.name)
        self.assertEqual("", category.systematics)
        self.assertEqual(ROOT.kBlue, category.color)

    def test_contain_check(self):
        category = sa.SystematicsCategory(systematics=['foo', 'bar'])
        self.assertTrue(category.contains_syst('total'))
        self.assertTrue(category.contains_syst('foo'))
        self.assertTrue(category.contains_syst('missing'))
        category = sa.SystematicsCategory(name='single', systematics=['foo', 'bar'])
        self.assertTrue(category.contains_syst('foo'))
        self.assertFalse(category.contains_syst('missing'))


class TestSystematics(unittest.TestCase):
    def setUp(self):
        self.syst = sa.Systematic('foo', variation='updown')
        self.scale_syst = sa.Systematic('foo', type='scale', variation='single:up')

    def tearDown(self):
        del self.syst
        del self.scale_syst

    def test_ctor_default(self):
        syst = sa.Systematic('foo')
        self.assertEqual('', syst.prefix)
        self.assertFalse(syst.symmetrise)
        self.assertIsNone(syst.symmetrise_option)
        self.assertIsNone(syst.call)
        self.assertIsNone(syst.affects)
        self.assertIsNone(syst.samples)
        self.assertIsNone(syst.group)
        self.assertIsNone(syst.title)
        self.assertIsNone(syst.type)
        self.assertIsNone(syst.expand)
        self.assertIsNone(syst.hist_name)
        self.assertIsNone(syst.variation)
        self.assertIsNone(syst.envelope)

    def test_ctor_scale(self):
        self.assertEqual('weight_', self.scale_syst.prefix)

    def test_str_op(self):
        self.assertEqual('Systematic uncertainty: foo', self.syst.__str__())

    def test_repr_op(self):
        self.assertEqual('Systematic uncertainty: foo\n', self.syst.__repr__())

    def test_get_variations_updown(self):
        self.assertEqual(['foo__1up', 'foo__1down'], self.syst.get_variations())

    def test_get_variations_scale(self):
        self.assertEqual(['weight_foo__1up'], self.scale_syst.get_variations())

    def test_get_variations_empty(self):
        self.assertEqual([], sa.Systematic('foo').get_variations())

    def test_get_variations_unsupported(self):
        syst = sa.Systematic('foo', type='scale', variation='foo', symmetrise=True)
        self.assertEqual([], syst.get_variations())

    def test_get_symmetrised_name_nosym(self):
        self.assertIsNone(self.syst.get_symmetrised_name())

    def test_get_symmetrised_name_invalid(self):
        syst = sa.Systematic('foo', type='scale', variation='single:foo', symmetrise=True)
        self.assertIsNone(syst.get_symmetrised_name())

    def test_get_symmetrised_name_up(self):
        scale_syst = sa.Systematic('foo', type='scale', variation='single:up', symmetrise=True)
        self.assertEqual(['weight_foo__1down'], scale_syst.get_symmetrised_name())

    def test_get_symmetrised_name_down(self):
        scale_syst = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        self.assertEqual(['weight_foo__1up'], scale_syst.get_symmetrised_name())

    def test_eq(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        syst2 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        self.assertEqual(syst1, syst2)

    def test_inequal(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        syst2 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True)
        self.assertNotEqual(syst1, syst2)

    def test_inequal_missing_key(self):
        syst1 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True, foo='bar')
        syst2 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True)
        self.assertNotEqual(syst1, syst2)

    def test_inequal_missing_key_swap(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        syst2 = sa.Systematic('foo', type='shape', variation='single:down', symmetrise=True, foo='bar')
        self.assertNotEqual(syst1, syst2)

    def test_inequal_type(self):
        syst1 = sa.Systematic('foo', type='scale', variation='single:down', symmetrise=True)
        self.assertNotEqual(1., syst1)


syst = sa.Systematic('foo', bar='foobar', type='shape', variation='updown')
scale_syst = sa.Systematic('foo_scale', bar='foobar_scale', type='scale', variation='updown')


class TestSystematicsAnalyser(unittest.TestCase):
    def setUp(self):
        pass

    def tearDown(self):
        pass

    def test_ctor_default(self):
        analyser = sa.SystematicsAnalyser(xs_handle='foo')
        self.assertEqual(1, analyser.ncpu)
        self.assertFalse(analyser.dump_hists)
        self.assertFalse(analyser.cluster_mode)
        self.assertFalse(analyser.store_abs_syst)
        self.assertIsNone(analyser.systematics)
        self.assertEqual({}, analyser.systematic_hists)
        self.assertEqual({}, analyser.systematic_variations)
        self.assertEqual({}, analyser.total_systematics)
        self.assertEqual([], analyser.scale_uncerts)
        self.assertEqual([], analyser.shape_uncerts)
        self.assertEqual([], analyser.file_handles)
        self.assertEqual([], analyser.custom_uncerts)
        self.assertEqual('foo', analyser.xs_handle)
        self.assertTrue(analyser.disable)

    @patch.object(yl, 'read_yaml', mock_read_yaml2)
    def test_parse_syst_cfg(self):
        systs = sa.SystematicsAnalyser.parse_syst_config('foo')
        self.assertEqual([sa.Systematic('foo', bar='foobar')], systs)

    @patch.object(FileHandle, 'get_object_by_name',
                  lambda *args: ROOT.TH1F('h_syst', '', 10, 0., 10.))
    def test_load_dumped_hists(self):
        fh = FileHandle(file_name='foo', process=Process('foo', dataset_info=None))
        pc, process, hist = sa.SystematicsAnalyser.load_dumped_hist((fh, PlotConfig()), 'foo')
        self.assertEqual(PlotConfig(), pc)
        self.assertEqual('foo', process.process_name)
        self.assertEqual('h_syst', hist.GetName())

    @patch.object(FileHandle, 'get_object_by_name', lambda *args: None)
    def test_load_dumped_hists_tuple(self):
        process = Process('foo', dataset_info=None)
        fh = FileHandle(file_name='foo', process=process)
        self.assertEqual([(PlotConfig(), process, None)], sa.SystematicsAnalyser(xs_handle='foo').load_dumped_hists([fh],
                                                                                                         [PlotConfig()],
                                                                                                         'foo'))

    @patch.object(FileHandle, 'get_object_by_name', lambda *args: None)
    def test_load_dumped_hists_theory(self):
        process = Process('foo', dataset_info=None)
        fh = FileHandle(file_name='foo', process=process)
        pc, res_process, hist = sa.SystematicsAnalyser.load_dumped_hist((fh, PlotConfig()),
                                                                    'weight_pdf_uncert_MUR0p5_MUF0p5_PDF261000')
        self.assertEqual(pc, PlotConfig())
        self.assertEqual(process, res_process)
        self.assertIsNone(hist)

    def test_retrieve_sys_hists_no_input_files(self):
        fh = FileHandle(file_name='foo', process=Process('data', dataset_info=None))
        analyser = sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo')
        self.assertEqual(1, len(analyser.file_handles))
        self.assertIsNone(analyser.retrieve_sys_hists())

    def test_retrieve_sys_hists_disable(self):
        self.assertIsNone(sa.SystematicsAnalyser(xs_handle='foo').retrieve_sys_hists())

    def test_retrieve_sys_hists_data_only(self):
        fh = FileHandle(file_name='foo', process=Process('foo', dataset_info=None))
        self.assertIsNone(sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo').retrieve_sys_hists())

    @patch.object(BasePlotter, 'read_histograms', lambda *args, **kwargs: [(PlotConfig(), 'foo',
                                                                            ROOT.TH1F('h_syst', '', 10, 0., 10.))])
    @patch.object(BasePlotter, 'apply_lumi_weights', lambda *args, **kwargs: None)
    @patch.object(sa.SystematicsAnalyser, 'parse_syst_config', lambda *args: [syst, scale_syst])
    @patch.object(BasePlotter, 'merge_histograms', lambda _: None)
    def test_retrieve_sys_hists(self):
        fh = FileHandle(file_name='foo', process=Process('foo_311011', dataset_info=None))
        print(sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo',
                                     plot_configs=[PlotConfig()]).retrieve_sys_hists())

    def test_get_variation_for_process_no_pc(self):
        dummy_process = Process('foo', dataset_info=None)
        fh = FileHandle(file_name='foo', process=dummy_process)
        analyser = sa.SystematicsAnalyser(file_handles=[fh], xs_handle='foo')
        analyser.systematic_hists = {'my_syst_foo': {PlotConfig(): [hist]}}
{'fraction_non_alphanumeric': 0.12627210749925896, 'fraction_numerical': 0.005533050093864243, 'mean_word_length': 3.638634280476627, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 32, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1442219', 'n_tokens_mistral': 6349, 'n_tokens_neox': 6211, 'n_words': 951}