<commit_before>import unittest
import matplotlib
# matplotlib.use('Agg')
import flavio
import flavio.plots
import numpy as np

# NB, this test only runs with matplotlib>=1.5.3 due to a matplotlib bug

class TestPlots(unittest.TestCase):

    def test_band_plot(self):
        def dummy_loglikelihood(x):
            return -x[0]**2-x[1]**2
        # check that no error is raised and output dimensions match
        x, y, z = flavio.plots.band_plot(dummy_loglikelihood,
                                         -2, 2, -3, 3, steps=30)
        self.assertEqual(x.shape, (30, 30))
        self.assertEqual(y.shape, (30, 30))
        self.assertEqual(z.shape, (30, 30))
        # with interpolation_factor
        x, y, z = flavio.plots.band_plot(dummy_loglikelihood,
                                         -2, 2, -3, 3, steps=30,
                                         interpolation_factor=2)
        self.assertEqual(x.shape, (30, 30))
        self.assertEqual(y.shape, (30, 30))
        self.assertEqual(z.shape, (30, 30))
        # with pre_calculated_z
        x, y, z = flavio.plots.band_plot(None, -2, 2, -3, 3,
                                         pre_calculated_z=z,
                                         interpolation_factor=2)
        self.assertEqual(x.shape, (30, 30))
        self.assertEqual(y.shape, (30, 30))
        self.assertEqual(z.shape, (30, 30))<commit_msg>[plots] More unit tests for plot functions<commit_after>import unittest
import matplotlib
# matplotlib.use('Agg')
import flavio
from flavio.plots import *
import numpy as np
import scipy.stats
import warnings

def dummy_loglikelihood(x):
    return -x[0]**2-x[1]**2

class TestPlots(unittest.TestCase):
    def test_error_budget_pie(self):
        err_budget_bsmumu = {'DeltaGamma/Gamma_Bs': 0.0048858044281356464,
             'GF': 1.1082519567930236e-06,
             'Vcb': 0.039790267942911725,
             'Vub': 0.00039636600693301931,
             'Vus': 0.00040185451336679405,
             'alpha_e': 0.00017640953398991146,
             'alpha_s': 1.7980573055638197e-05,
             'f_Bs': 0.035073683628645283,
             'gamma': 0.0040950620820486205,
             'm_Bs': 4.395718017004336e-05,
             'm_mu': 3.89716785433222e-08,
             'tau_Bs': 0.003286868163723475}
        error_budget_pie(err_budget_bsmumu)

    def test_band_plot(self):
        # NB, this test only runs with matplotlib>=1.5.3 due to a matplotlib bug
        # check that no error is raised and output dimensions match
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            
            x, y, z = band_plot(dummy_loglikelihood,
                                             -2, 2, -3, 3, steps=30)
            self.assertEqual(x.shape, (30, 30))
            self.assertEqual(y.shape, (30, 30))
            self.assertEqual(z.shape, (30, 30))
            # with interpolation_factor
            x, y, z = band_plot(dummy_loglikelihood,
                                             -2, 2, -3, 3, steps=30,
                                             interpolation_factor=2)
            self.assertEqual(x.shape, (30, 30))
            self.assertEqual(y.shape, (30, 30))
            self.assertEqual(z.shape, (30, 30))
            # with pre_calculated_z
            x, y, z = band_plot(None, -2, 2, -3, 3,
                                             pre_calculated_z=z,
                                             interpolation_factor=2)
            self.assertEqual(x.shape, (30, 30))
            self.assertEqual(y.shape, (30, 30))
            self.assertEqual(z.shape, (30, 30))

    def test_density_contour_data(self):
        np.random.seed(42)
        xy = scipy.stats.multivariate_normal(mean=[2,3], cov=[[1,0.5],[0.5,1]]).rvs(size=100)
        data = density_contour_data(*xy.T)
        self.assertEqual(data['x'].shape, (100,100))
        self.assertEqual(data['y'].shape, (100,100))
        self.assertEqual(data['z'].shape, (100,100))
        self.assertEqual(len(data['levels']), 2) # by default 2 levels (1, 2sigma)
        self.assertTrue(min(data['levels']) > 0) # levels positive
        self.assertEqual(data['levels'], sorted(data['levels'])) # levels ascending
        self.assertEqual(np.min(data['z']), 0)
        # point in the middle should be close to maximum likelihood
        self.assertAlmostEqual(data['z'][50,50], 0, delta=0.1)
        # corners
        self.assertTrue(data['z'][0,0] < data['z'][-1,0])
        self.assertTrue(data['z'][-1,-1] < data['z'][0,-1])
        # symmetries
        self.assertAlmostEqual(data['z'][-1,-1], data['z'][0,0], delta=1.)
        self.assertAlmostEqual(data['z'][-1,0], data['z'][0,-1], delta=3.)

    def test_density_contour(self):
        # just check this works
        np.random.seed(42)
        xy = scipy.stats.multivariate_normal(mean=[2,3], cov=[[1,0.5],[0.5,1]]).rvs(size=100)
        density_contour(*xy.T)

    def test_likelihood_contour(self):
        # just check this works
        data = likelihood_countour_data(dummy_loglikelihood,
                                        -2, 2, -3, 3)
        self.assertEqual(data['x'].shape, (20,20))
        self.assertEqual(data['y'].shape, (20,20))
        self.assertEqual(data['z'].shape, (20,20))
        self.assertEqual(len(data['levels']), 1) # by default, plot 1 sigma contour
        self.assertAlmostEqual(data['levels'][0], 2.3, delta=0.01) #
        self.assertEqual(np.min(data['z']), 0)
{'fraction_non_alphanumeric': 0.1376029277218664, 'fraction_numerical': 0.08325709057639524, 'mean_word_length': 2.0553381777529345, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 6, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 22, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3378759', 'n_tokens_mistral': 2024, 'n_tokens_neox': 1726, 'n_words': 346}