<commit_before>import pickle
import collections

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq, minimize
from scipy import stats
from tqdm import tqdm
from copy import deepcopy

from multihist import Histdd
from .source import Source


class Model(object):
    """Model for XENON1T dataset simulation and analysis
    """
    config = None       # type: dict
    no_wimp_strength = -10
    max_wimp_strength = 2
    exposure_factor = 1  # Increase this to quickly change the exposure of the model

    def __init__(self, config, ipyparallel_view=None, **kwargs):
        """
        :param config: Dictionary specifying detector parameters, source info, etc.
        :param ipyparallel_view: DirectView object to use for parallelizing pdf computation (optional)
        :param kwargs: Overrides for the config (optional)
        :return:
        """
        c = self.config = deepcopy(config)
        self.config.update(kwargs)

        self.space = collections.OrderedDict(c['analysis_space'])
        self.dims = list(self.space.keys())
        self.bins = list(self.space.values())

        self.sources = []
        for source_spec in tqdm(c['sources'], desc='Initializing sources'):
            source = Source(self.config, source_spec)

            # Has the PDF in the analysis space already been provided in the spec?
            # Usually not, do so now.
            if source.pdf_histogram is None or c['force_pdf_recalculation']:
                self.compute_source_pdf(source, ipyparallel_view=ipyparallel_view)
            self.sources.append(source)

    def compute_source_pdf(self, source, ipyparallel_view=None):
        """Computes the PDF for the source. Returns nothing but modifies source.
        """
        # Not a method of source, since it needs analysis space definition...
        # To be honest I'm not sure where this method (and source.simulate) actually would fit best.

        # Simulate batches of 1e6 events at a time (to avoid memory errors and show a meaningful progressbar)
        n_events = source.n_events_for_pdf
        batch_size = self.config['pdf_sampling_batch_size']
        n_events = int(batch_size * (int(n_events / batch_size) + 1))
        mh = Histdd(bins=self.bins)

        # Standalone counterpart of Model.to_space, needed for parallel simulation. Ugly!!
        def to_space(d, dims):
            return [d[dims[i]] for i in range(len(dims))]

        if ipyparallel_view is not None:
            # Get the necessary objects to the nodes
            ipyparallel_view.push(dict(source=source, bins=self.bins, batch_size=batch_size, to_space=to_space),
                                  block=True)

            # Run small simulations on each node. Histogram immediately so we don't have to pass gobs of data around.
            # Make a copy of the histogram first, so we don't have to
            def do_sim(dims):
                return Histdd(*to_space(source.simulate(batch_size), dims),
                              bins=bins)

            for q in ipyparallel_view.map_sync(do_sim, [self.dims for _ in range(int(n_events / batch_size))]):
                mh.histogram += q.histogram

        else:
            for _ in tqdm(range(int(n_events / batch_size)),
                          desc='Sampling PDF of %s' % source.name):
                mh.add(*self.to_space(source.simulate(batch_size)))

        source.fraction_in_range = mh.n / n_events

        # Convert the histogram to a PDF
        # This means we have to divide by
        #  - the number of events histogrammed
        #  - the bin sizes (particularly relevant for non-uniform bins!)
        mh.histogram = mh.histogram.astype(np.float) / mh.n
        mh.histogram /= np.outer(*[np.diff(self.bins[i]) for i in range(len(self.bins))])
        source.pdf_histogram = mh

    def range_cut(self, d):
        """Return events from dataset d which are in the analysis space"""
        mask = np.ones(len(d), dtype=np.bool)
        for dimension, bin_edges in self.space.items():
            mask = mask & (d[dimension] >= bin_edges[0]) & (d[dimension] <= bin_edges[-1])
        return d[mask]

    def simulate(self, wimp_strength=0, restrict=True):
        """Makes a toy dataset.
        if restrict=True, return only events inside analysis range
        """
        ds = []
        for s_i, source in enumerate(self.sources):<commit_msg>Better parallelization of PDF generation<commit_after>import pickle
import collections

import numpy as np
import matplotlib.pyplot as plt
from scipy.optimize import brentq, minimize
from scipy import stats
from tqdm import tqdm
from copy import deepcopy

from multihist import Histdd
from .source import Source


class Model(object):
    """Model for XENON1T dataset simulation and analysis
    """
    config = None       # type: dict
    no_wimp_strength = -10
    max_wimp_strength = 2
    exposure_factor = 1  # Increase this to quickly change the exposure of the model

    def __init__(self, config, ipp_client=None, **kwargs):
        """
        :param config: Dictionary specifying detector parameters, source info, etc.
        :param ipp_client: ipyparallel client to use for parallelizing pdf computation (optional)
        :param kwargs: Overrides for the config (optional)
        :return:
        """
        c = self.config = deepcopy(config)
        self.config.update(kwargs)

        self.space = collections.OrderedDict(c['analysis_space'])
        self.dims = list(self.space.keys())
        self.bins = list(self.space.values())

        self.sources = []
        for source_spec in tqdm(c['sources'], desc='Initializing sources'):
            source = Source(self.config, source_spec)

            # Has the PDF in the analysis space already been provided in the spec?
            # Usually not, do so now.
            if source.pdf_histogram is None or c['force_pdf_recalculation']:
                self.compute_source_pdf(source, ipp_client=ipp_client)
            self.sources.append(source)

    def compute_source_pdf(self, source, ipp_client=None):
        """Computes the PDF for the source. Returns nothing but modifies source.
        """
        # Not a method of source, since it needs analysis space definition...
        # To be honest I'm not sure where this method (and source.simulate) actually would fit best.

        # Simulate batches of events at a time (to avoid memory errors, show a progressbar, and split up among machines)
        # Number of events to simulate will be rounded up to the nearest batch size
        batch_size = self.config['pdf_sampling_batch_size']
        n_batches = int(source.n_events_for_pdf // batch_size + 1)
        n_events = n_batches * batch_size
        mh = Histdd(bins=self.bins)

        if ipp_client is not None:
            # We need both a directview and a load-balanced view: the latter doesn't have methods like push.
            directview = ipp_client[:]
            lbview = ipp_client.load_balanced_view()
        
            # Get the necessary objects to the engines
            # For some reason you can't directly .push(dict(bins=self.bins)),
            # it will fail with 'bins' is not defined error. When you first assign bins = self.bins it works.
            bins = self.bins
            dims = self.dims
            
            def to_space(d):
                """Standalone counterpart of Model.to_space, needed for parallel simulation. Ugly!!"""
                return [d[dims[i]] for i in range(len(dims))]
                
            def do_sim(_):
                """Run one simulation batch and histogram it immediately (so we don't pass gobs of data around)"""
                return Histdd(*to_space(source.simulate(batch_size)), bins=bins).histogram

            directview.push(dict(source=source, bins=bins, dims=dims, batch_size=batch_size, to_space=to_space),
                                  block=True)
                
            amap_result = lbview.map(do_sim, [None for _ in range(n_batches)], ordered=False, 
                                     block=self.config.get('block_during_simulation', False))
            for r in tqdm(amap_result, total=n_batches, desc='Sampling PDF of %s' % source.name):
                mh.histogram += r

        else:
            for _ in tqdm(range(n_batches),
                          desc='Sampling PDF of %s' % source.name):
                mh.add(*self.to_space(source.simulate(batch_size)))

        source.fraction_in_range = mh.n / n_events

        # Convert the histogram to a PDF
        # This means we have to divide by
        #  - the number of events histogrammed
        #  - the bin sizes (particularly relevant for non-uniform bins!)
        source.pdf_histogram = mh.similar_blank_hist()
        source.pdf_histogram.histogram = mh.histogram.astype(np.float) / mh.n
        source.pdf_histogram.histogram /= np.outer(*[np.diff(self.bins[i]) for i in range(len(self.bins))])

        # Estimate the MC statistical error. Not used for anything, but good to inspect.
        source.pdf_errors = source.pdf_histogram / np.sqrt(np.clip(mh.histogram, 1, float('inf')))
        source.pdf_errors[source.pdf_errors == 0] = float('nan')


    def range_cut(self, d):
        """Return events from dataset d which are in the analysis space"""
        mask = np.ones(len(d), dtype=np.bool)
        for dimension, bin_edges in self.space.items():
            mask = mask & (d[dimension] >= bin_edges[0]) & (d[dimension] <= bin_edges[-1])
        return d[mask]

    def simulate(self, wimp_strength=0, restrict=True):
        """Makes a toy dataset.
        if restrict=True, return only events inside analysis range
        """
        ds = []
        for s_i, source in enumerate(self.sources):
{'fraction_non_alphanumeric': 0.0905463474524248, 'fraction_numerical': 0.0022508696541845715, 'mean_word_length': 2.420223932820154, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 3, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5140225', 'n_tokens_mistral': 2846, 'n_tokens_neox': 2712, 'n_words': 882}