<commit_before>from __future__ import print_function, division
from copy import copy
from datetime import timedelta
import numpy as np
import pandas as pd
import nilmtk
from nilmtk.timeframegroup import TimeFrameGroup
from nilmtk.timeframe import TimeFrame
from neuralnilm.data.source import Sequence
from neuralnilm.utils import check_windows
from neuralnilm.data.activationssource import ActivationsSource

import logging
logger = logging.getLogger(__name__)


class RealAggregateSource(ActivationsSource):
    """
    Attributes
    ----------
    mains : dict
        Structure example:
        {<train | unseen_appliances | unseen_activations_of_seen_appliances>: {
             <building_name>: pd.Series of raw data
        }}
    mains_good_sections : dict
        Same structure as `mains`.
    sections_with_no_target : dict
        Same structure as `mains`.
    """
    def __init__(self, activations, target_appliance,
                 seq_length, filename, windows, sample_period,
                 target_inclusion_prob=0.5,
                 uniform_prob_of_selecting_each_building=True,
                 allow_incomplete_target=True,
                 include_incomplete_target_in_output=True,
                 allow_multiple_target_activations_in_aggregate=False,
                 include_multiple_targets_in_output=False,
                 rng_seed=None):
        self.activations = copy(activations)
        self.target_appliance = target_appliance
        self.seq_length = seq_length
        self.filename = filename
        check_windows(windows)
        self.windows = windows
        self.sample_period = sample_period
        self.target_inclusion_prob = target_inclusion_prob
        self.uniform_prob_of_selecting_each_building = (
            uniform_prob_of_selecting_each_building)
        self.allow_incomplete_target = allow_incomplete_target
        self.include_incomplete_target_in_output = (
            include_incomplete_target_in_output)
        self.allow_multiple_target_activations_in_aggregate = (
            allow_multiple_target_activations_in_aggregate)
        self.include_multiple_targets_in_output = (
            include_multiple_targets_in_output)
        super(RealAggregateSource, self).__init__(rng_seed=rng_seed)

        self._load_mains_into_memory()
        self._remove_activations_with_no_mains()
        self._find_sections_with_no_target()

    def _load_mains_into_memory(self):
        logger.info("Loading NILMTK mains...")

        # Load dataset
        dataset = nilmtk.DataSet(self.filename)

        self.mains = {}
        self.mains_good_sections = {}
        for fold, buildings_and_windows in self.windows.iteritems():
            for building_i, window in buildings_and_windows.iteritems():
                dataset.set_window(*window)
                elec = dataset.buildings[building_i].elec
                building_name = (
                    dataset.metadata['name'] +
                    '_building_{}'.format(building_i))

                logger.info(
                    "Loading mains for {}...".format(building_name))

                mains_meter = elec.mains()
                good_sections = mains_meter.good_sections()
                mains_data = mains_meter.power_series_all_data(
                    sample_period=self.sample_period,
                    sections=good_sections).dropna()

                def set_mains_data(dictionary, data):
                    dictionary.setdefault(fold, {})[building_name] = data

                if not mains_data.empty:
                    set_mains_data(self.mains, mains_data)
                    set_mains_data(self.mains_good_sections, good_sections)

                logger.info(
                    "Loaded mains data from building {} for fold {}"
                    " from {} to {}."
                    .format(building_name, fold,
                            mains_data.index[0], mains_data.index[-1]))

        dataset.store.close()
        logger.info("Done loading NILMTK mains data.")

    def _find_sections_with_no_target(self):
        """Finds the intersections of the mains good sections with the gaps
        between target appliance activations.
        """
        self.sections_with_no_target = {}
        seq_length_secs = self.seq_length * self.sample_period
        for fold, sects_per_building in self.mains_good_sections.iteritems():
            for building, good_sections in sects_per_building.iteritems():
                activations = (
                    self.activations[fold][self.target_appliance][building])
                mains = self.mains[fold][building]
                mains_good_sections = self.mains_good_sections[fold][building]
                gaps_between_activations = TimeFrameGroup()
                prev_end = mains.index[0]
                for activation in activations:
                    gap = TimeFrame(prev_end, activation.index[0])
                    gaps_between_activations.append(gap)
                    prev_end = activation.index[-1]
                gap = TimeFrame(prev_end, mains.index[-1])
                gaps_between_activations.append(gap)
                intersection = (
                    gaps_between_activations.intersection(mains_good_sections))
                intersection = intersection.remove_shorter_than(
                    seq_length_secs)
                self.sections_with_no_target.setdefault(fold, {})[building] = (
                    intersection)
                logger.info("Found {} sections without target for {} {}."
                            .format(len(intersection), fold, building))

    def _get_sequence_without_target(self, fold):
        # Choose a building and a gap
        all_gaps = []
        for building, gaps in self.sections_with_no_target[fold].iteritems():
            all_gaps.extend([(building, gap) for gap in gaps])
        lengths = np.array(
            [gap.timedelta.total_seconds() for (_, gap) in all_gaps])
        p = lengths / lengths.sum()
        gap_i = self.rng.choice(len(all_gaps), p=p)
        building, gap = all_gaps[gap_i]

        # Choose a start point in the gap
        latest_start_time = gap.end - timedelta(
            seconds=self.seq_length * self.sample_period)<commit_msg>speed up _get_sequence_without_target()<commit_after>from __future__ import print_function, division
from copy import copy
from datetime import timedelta
import numpy as np
import pandas as pd
import nilmtk
from nilmtk.timeframegroup import TimeFrameGroup
from nilmtk.timeframe import TimeFrame
from neuralnilm.data.source import Sequence
from neuralnilm.utils import check_windows
from neuralnilm.data.activationssource import ActivationsSource
from neuralnilm.consts import DATA_FOLD_NAMES

import logging
logger = logging.getLogger(__name__)


class RealAggregateSource(ActivationsSource):
    """
    Attributes
    ----------
    mains : dict
        Structure example:
        {<train | unseen_appliances | unseen_activations_of_seen_appliances>: {
             <building_name>: pd.Series of raw data
        }}
    mains_good_sections : dict
        Same structure as `mains`.
    sections_with_no_target : dict
        Same structure as `mains`.
    all_gaps : dict of pd.DataFrames
        Each key is a fold name.
        Each DF has columns:
        building, gap, duration, p (short for 'probability')
        p is used by _get_sequence_without_target().
    """
    def __init__(self, activations, target_appliance,
                 seq_length, filename, windows, sample_period,
                 target_inclusion_prob=0.5,
                 uniform_prob_of_selecting_each_building=True,
                 allow_incomplete_target=True,
                 include_incomplete_target_in_output=True,
                 allow_multiple_target_activations_in_aggregate=False,
                 include_multiple_targets_in_output=False,
                 rng_seed=None):
        self.activations = copy(activations)
        self.target_appliance = target_appliance
        self.seq_length = seq_length
        self.filename = filename
        check_windows(windows)
        self.windows = windows
        self.sample_period = sample_period
        self.target_inclusion_prob = target_inclusion_prob
        self.uniform_prob_of_selecting_each_building = (
            uniform_prob_of_selecting_each_building)
        self.allow_incomplete_target = allow_incomplete_target
        self.include_incomplete_target_in_output = (
            include_incomplete_target_in_output)
        self.allow_multiple_target_activations_in_aggregate = (
            allow_multiple_target_activations_in_aggregate)
        self.include_multiple_targets_in_output = (
            include_multiple_targets_in_output)
        super(RealAggregateSource, self).__init__(rng_seed=rng_seed)

        self._load_mains_into_memory()
        self._remove_activations_with_no_mains()
        self._find_sections_with_no_target()
        self._compute_gap_probabilities()

    def _load_mains_into_memory(self):
        logger.info("Loading NILMTK mains...")

        # Load dataset
        dataset = nilmtk.DataSet(self.filename)

        self.mains = {}
        self.mains_good_sections = {}
        for fold, buildings_and_windows in self.windows.iteritems():
            for building_i, window in buildings_and_windows.iteritems():
                dataset.set_window(*window)
                elec = dataset.buildings[building_i].elec
                building_name = (
                    dataset.metadata['name'] +
                    '_building_{}'.format(building_i))

                logger.info(
                    "Loading mains for {}...".format(building_name))

                mains_meter = elec.mains()
                good_sections = mains_meter.good_sections()
                mains_data = mains_meter.power_series_all_data(
                    sample_period=self.sample_period,
                    sections=good_sections).dropna()

                def set_mains_data(dictionary, data):
                    dictionary.setdefault(fold, {})[building_name] = data

                if not mains_data.empty:
                    set_mains_data(self.mains, mains_data)
                    set_mains_data(self.mains_good_sections, good_sections)

                logger.info(
                    "Loaded mains data from building {} for fold {}"
                    " from {} to {}."
                    .format(building_name, fold,
                            mains_data.index[0], mains_data.index[-1]))

        dataset.store.close()
        logger.info("Done loading NILMTK mains data.")

    def _find_sections_with_no_target(self):
        """Finds the intersections of the mains good sections with the gaps
        between target appliance activations.
        """
        self.sections_with_no_target = {}
        seq_length_secs = self.seq_length * self.sample_period
        for fold, sects_per_building in self.mains_good_sections.iteritems():
            for building, good_sections in sects_per_building.iteritems():
                activations = (
                    self.activations[fold][self.target_appliance][building])
                mains = self.mains[fold][building]
                mains_good_sections = self.mains_good_sections[fold][building]
                gaps_between_activations = TimeFrameGroup()
                prev_end = mains.index[0]
                for activation in activations:
                    gap = TimeFrame(prev_end, activation.index[0])
                    gaps_between_activations.append(gap)
                    prev_end = activation.index[-1]
                gap = TimeFrame(prev_end, mains.index[-1])
                gaps_between_activations.append(gap)
                intersection = (
                    gaps_between_activations.intersection(mains_good_sections))
                intersection = intersection.remove_shorter_than(
                    seq_length_secs)
                self.sections_with_no_target.setdefault(fold, {})[building] = (
                    intersection)
                logger.info("Found {} sections without target for {} {}."
                            .format(len(intersection), fold, building))

    def _compute_gap_probabilities(self):
        # Choose a building and a gap
        self.all_gaps = {}
        for fold in DATA_FOLD_NAMES:
            all_gaps_for_fold = []
            for building, gaps in self.sections_with_no_target[fold].iteritems():
                gaps_for_building = [
                    (building, gap, gap.timedelta.total_seconds())
                    for gap in gaps]
                all_gaps_for_fold.extend(gaps_for_building)
            gaps_df = pd.DataFrame(
                all_gaps_for_fold, columns=['building', 'gap', 'duration'])
            gaps_df['p'] = gaps_df['duration'] / gaps_df['duration'].sum()
            self.all_gaps[fold] = gaps_df

    def _get_sequence_without_target(self, fold):
        # Choose a building and a gap
        all_gaps_for_fold = self.all_gaps[fold]
        n = len(all_gaps_for_fold)
        gap_i = self.rng.choice(n, p=all_gaps_for_fold['p'])
        row = all_gaps_for_fold.iloc[gap_i]
        building, gap = row['building'], row['gap']

        # Choose a start point in the gap
        latest_start_time = gap.end - timedelta(
            seconds=self.seq_length * self.sample_period)
{'fraction_non_alphanumeric': 0.07727374635677453, 'fraction_numerical': 0.0011957252821164336, 'mean_word_length': 2.2052694610778443, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7376431', 'n_tokens_mistral': 3766, 'n_tokens_neox': 3593, 'n_words': 680}