<commit_before>import numbers
import logging

import yaml

from .generator import IndexGenerator
from .reader import RecordingsReader
from .buffer import BufferGenerator


class BatchProcessor(object):
    """
    Batch processing for large numpy matrices

    Parameters
    ----------
    path_to_recordings: str
        Path to recordings file

    dtype: str
        Numpy dtype

    n_channels: int
        Number of channels

    data_format: str
        Data format, it can be either 'long' (observations, channels) or
        'wide' (channels, observations)

    max_memory: int or str
        Max memory to use in each batch, interpreted as bytes if int,
        if string, it can be any of {N}KB, {N}MB or {N}GB

    buffer_size: int, optional
        Buffer size, defaults to 0

    Raises
    ------
    ValueError
        If dimensions do not match according to the file size, dtype and
        number of channels
    """
    def __init__(self, path_to_recordings, dtype, n_channels,
                 data_format, max_memory, buffer_size=0):
        self.data_format = data_format
        self.buffer_size = buffer_size
        self.reader = RecordingsReader(path_to_recordings, dtype, n_channels,
                                       data_format)
        self.indexer = IndexGenerator(self.reader.observations,
                                      self.reader.channels,
                                      dtype,
                                      max_memory)
        self.buffer_generator = BufferGenerator(self.reader.observations,
                                                self.data_format,
                                                buffer_size=buffer_size)

        self.logger = logging.getLogger(__name__)

    def single_channel(self, force_complete_channel_batch=True, from_time=None,
                       to_time=None, channels='all'):
        """
        Generate indexes where each index has observations from a single
        channel

        Returns
        -------
        A generator that yields indexes

        Examples
        --------

        .. literalinclude:: ../../examples/batch/single_channel.py
        """
        indexes = self.indexer.single_channel(force_complete_channel_batch,
                                              from_time, to_time,
                                              channels)
        if force_complete_channel_batch:
            for idx in indexes:
                yield self.reader[idx]
        else:
            for idx in indexes:
                channel_idx = idx[1]
                yield self.reader[idx], channel_idx

    def multi_channel(self, from_time=None, to_time=None, channels='all'):
        """
        Generate indexes where each index has observations from more than
        one channel

        Returns
        -------
        generator:
            A slice that yields indexes

        Examples
        --------
        .. literalinclude:: ../../examples/batch/multi_channel.py
        """
        indexes = self.indexer.multi_channel(from_time, to_time, channels)

        for idx in indexes:
            if self.buffer_size:
                (dx_new,
                 (buff_start, buff_end)) = (self.buffer_generator
                                            .update_key_with_buffer(idx))
                subset = self.reader[idx]
                yield self.buffer_generator.add_buffer(subset, buff_start,
                                                       buff_end)
            else:
                yield self.reader[idx]

    def single_channel_apply(self, function, output_path,
                             force_complete_channel_batch=True,
                             from_time=None, to_time=None, channels='all',
                             **kwargs):
        """
        Apply a transformation where each batch has observations from a
        single channel

        Parameters
        ----------
        function: callable
            Function to be applied, must accept a 1D numpy array as its first
            parameter

        output_path: str
            Where to save the output

        force_complete_channel_batch: bool, optional
            If True, every index generated will correspond to all the
            observations in a single channel, hence
            n_batches = n_selected_channels, defaults to True. If True
            from_time and to_time must be None

        from_time: int, optional
            Starting time, defaults to None<commit_msg>fixes bug in batch processor<commit_after>import numbers
import logging

import yaml

from .generator import IndexGenerator
from .reader import RecordingsReader
from .buffer import BufferGenerator


class BatchProcessor(object):
    """
    Batch processing for large numpy matrices

    Parameters
    ----------
    path_to_recordings: str
        Path to recordings file

    dtype: str
        Numpy dtype

    n_channels: int
        Number of channels

    data_format: str
        Data format, it can be either 'long' (observations, channels) or
        'wide' (channels, observations)

    max_memory: int or str
        Max memory to use in each batch, interpreted as bytes if int,
        if string, it can be any of {N}KB, {N}MB or {N}GB

    buffer_size: int, optional
        Buffer size, defaults to 0

    Raises
    ------
    ValueError
        If dimensions do not match according to the file size, dtype and
        number of channels
    """
    def __init__(self, path_to_recordings, dtype, n_channels,
                 data_format, max_memory, buffer_size=0):
        self.data_format = data_format
        self.buffer_size = buffer_size
        self.reader = RecordingsReader(path_to_recordings, dtype, n_channels,
                                       data_format)
        self.indexer = IndexGenerator(self.reader.observations,
                                      self.reader.channels,
                                      dtype,
                                      max_memory)
        self.buffer_generator = BufferGenerator(self.reader.observations,
                                                self.data_format,
                                                buffer_size=buffer_size)

        self.logger = logging.getLogger(__name__)

    def single_channel(self, force_complete_channel_batch=True, from_time=None,
                       to_time=None, channels='all'):
        """
        Generate indexes where each index has observations from a single
        channel

        Returns
        -------
        A generator that yields indexes

        Examples
        --------

        .. literalinclude:: ../../examples/batch/single_channel.py
        """
        indexes = self.indexer.single_channel(force_complete_channel_batch,
                                              from_time, to_time,
                                              channels)
        if force_complete_channel_batch:
            for idx in indexes:
                yield self.reader[idx]
        else:
            for idx in indexes:
                channel_idx = idx[1]
                yield self.reader[idx], channel_idx

    def multi_channel(self, from_time=None, to_time=None, channels='all'):
        """
        Generate indexes where each index has observations from more than
        one channel

        Returns
        -------
        generator:
            A slice that yields indexes

        Examples
        --------
        .. literalinclude:: ../../examples/batch/multi_channel.py
        """
        indexes = self.indexer.multi_channel(from_time, to_time, channels)

        for idx in indexes:
            if self.buffer_size:
                (idx_new,
                 (buff_start, buff_end)) = (self.buffer_generator
                                            .update_key_with_buffer(idx))
                subset = self.reader[idx_new]
                yield self.buffer_generator.add_buffer(subset, buff_start,
                                                       buff_end)
            else:
                yield self.reader[idx]

    def single_channel_apply(self, function, output_path,
                             force_complete_channel_batch=True,
                             from_time=None, to_time=None, channels='all',
                             **kwargs):
        """
        Apply a transformation where each batch has observations from a
        single channel

        Parameters
        ----------
        function: callable
            Function to be applied, must accept a 1D numpy array as its first
            parameter

        output_path: str
            Where to save the output

        force_complete_channel_batch: bool, optional
            If True, every index generated will correspond to all the
            observations in a single channel, hence
            n_batches = n_selected_channels, defaults to True. If True
            from_time and to_time must be None

        from_time: int, optional
            Starting time, defaults to None
{'fraction_non_alphanumeric': 0.06740576496674058, 'fraction_numerical': 0.0008869179600886918, 'mean_word_length': 1.4926775352307267, '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': '6683736', 'n_tokens_mistral': 2136, 'n_tokens_neox': 2041, 'n_words': 647}