<commit_before>from collections import defaultdict, OrderedDict

import numpy as np
import pylab as plt

from tracking import VortMax, VortMaxTrack
from utils.utils import dist
from utils.kalman import RTSSmoother, _plot_rts_smoother

CUM_DIST_CUTOFF = 100


class EnsembleMatch(object):
    '''Represents one match between many vorticity tracks'''
    def __init__(self, vort_track, store_all_tracks=False):
        self.av_vort_track = vort_track

        # May have to disable due to memory constraints.
        self.store_all_tracks = store_all_tracks
        if self.store_all_tracks:
            self.vort_tracks = []
            self.vort_tracks.append(vort_track)

        self.cum_dist = 0
        self.overlap = 0

        self.dates = vort_track.dates
        self.date_start = vort_track.dates[0]
        self.date_end = vort_track.dates[-1]

        index_start = np.where(vort_track.dates == self.date_start)[0][0]
        index_end = np.where(vort_track.dates == self.date_end)[0][0]

        self.tracks_added = np.ones(index_end - index_start + 1)

    def add_track(self, vort_track):
        '''Checks to see whether the track should be added, then adds it

        Compares the overlap and cumalitive distance between the track to the current average
        track, and if they are below a limit adds the track, which involves recalculating the
        average track for this match, and updating some fields.

        :param vort_track: vorticitiy track to be added
        :returns: True if it was added, otherwise False
        '''
        # Calculate some dates and offsets.
        overlap_start = max(self.date_start, vort_track.dates[0])
        overlap_end = min(self.date_end, vort_track.dates[-1])

        index1 = np.where(self.dates == overlap_start)[0][0]
        index2 = np.where(vort_track.dates == overlap_start)[0][0]

        end1 = np.where(self.dates == overlap_end)[0][0]
        end2 = np.where(vort_track.dates == overlap_end)[0][0]

        assert (end1 - index1) == (end2 - index2)

        overlap, cum_dist = self.__calc_distance(self.av_vort_track, vort_track,
                                                 index1, index2, end1, end2)

        # Check to see whether this vort_track should be added.
        if overlap < 6 or cum_dist / overlap > 6:
            return False

        if self.store_all_tracks:
            self.vort_tracks.append(vort_track)

        # import ipdb; ipdb.set_trace()
        new_tracks_added = np.zeros(len(self.dates) + len(vort_track.dates) - overlap)
        copy_start_date = self.dates[0]
        copy_length = len(self.dates)
        ones_start_date = vort_track.dates[0]
        ones_length = len(vort_track.dates)

        self.dates = np.array(sorted([d for d in set(self.dates) | set(vort_track.dates)]))
        self.date_start = self.dates[0]
        self.date_end = self.dates[-1]

        # Copy over the existing tracks added to the new tracks added.
        # new_tracks_added must be equal or greater in length than current.
        # index it using the greatest of index1, and index2.
        copy_index = np.where(self.dates == copy_start_date)[0][0]
        ones_index = np.where(self.dates == ones_start_date)[0][0]

        new_tracks_added[copy_index:copy_index + copy_length] = self.tracks_added
        new_tracks_added[ones_index:ones_index + ones_length] += np.ones(ones_length)

        self.tracks_added = new_tracks_added

        # Update some fields.
        self.overlap += overlap
        self.cum_dist += cum_dist

        vortmaxes = []
        # Update the average vort track by taking a weighted average of the new
        # track and the current average. (N.B. self.tracks_added not yet updated.)
        # import ipdb; ipdb.set_trace()
        for i, date in enumerate(self.dates):
            vm1, vm2 = None, None
            if date in self.av_vort_track.vortmax_by_date:
                vm1 = self.av_vort_track.vortmax_by_date[date]
            if date in vort_track.vortmax_by_date:
                vm2 = vort_track.vortmax_by_date[date]

            if vm1 and vm2:
                # index tracks_added using min of index1, index2
                tracks_added = self.tracks_added[i]

                # Calculate a weighted average of the averaged pos and the new pos.
                pos = tuple([((tracks_added - 1) * v1 + v2) / (tracks_added)
                             for v1, v2 in zip(vm1.pos, vm2.pos)])

                # Calculate a weighted average of the averaged vort and the new vort.
                vort = ((tracks_added - 1) * vm1.vort) / (tracks_added)

                new_vm = VortMax(date, pos, vort)
                index1 += 1
                index2 += 1
            elif vm1:
                new_vm = VortMax(vm1.date, vm1.pos, vm1.vort)
                index1 += 1
            elif vm2:
                new_vm = VortMax(vm2.date, vm2.pos, vm2.vort)
                index2 += 1
            else:
                import ipdb; ipdb.set_trace()
                assert False, 'Should never be reached'

            vortmaxes.append(new_vm)

        next_vm = None
        for vm in vortmaxes[::-1]:
            if next_vm:
                vm.next_vortmax.append(next_vm)
            next_vm = vm
        self.av_vort_track = VortMaxTrack(vm)
        return True

    def __calc_distance(self,
                        vort_track_0, vort_track_1,
                        index1, index2, end1, end2):
        overlap = end1 - index1 + 1

        cum_dist = 0

        while index1 <= end1 and index2 <= end2:
            pos1 = vort_track_0.vortmaxes[index1].pos
            pos2 = vort_track_1.vortmaxes[index2].pos

            cum_dist += dist(pos1, pos2)
            index1 += 1
            index2 += 1

        return overlap, cum_dist

    def av_dist(self):
        '''Returns the average distance between all vorticity tracks'''
        return self.cum_dist / self.overlap


class Match(object):
    '''Represents one match between a best track and a vorticity track'''
    def __init__(self, best_track, vort_track):
        self.best_track = best_track
        self.vort_track = vort_track
        self.cum_dist = 0
        self.overlap = 1
        self.is_too_far_away = False
        self.overlap_start = max(best_track.dates[0], vort_track.dates[0])
        self.overlap_end = min(best_track.dates[-1], vort_track.dates[-1])

    def av_dist(self):
        '''Returns the average distance between the best and vorticity tracks'''
        return self.cum_dist / self.overlap


def match_vort_tracks(vort_tracks_by_dates):
    matches = OrderedDict()

    # import ipdb; ipdb.set_trace()
    vort_tracks_by_date_0 = vort_tracks_by_date[0]
    vort_tracks_by_date_1 = vort_tracks_by_date[1]

    all_dates = set(vort_tracks_by_date_0.keys()) | set(vort_tracks_by_date_1.keys())

    for date in all_dates:
        # This means that every vort track is represented **multiple times**, once
        # for each date where it appears. Hence the need to check for the key in matches dict.
        vort_tracks_0 = vort_tracks_by_date_0[date]
        vort_tracks_1 = vort_tracks_by_date_1[date]

        for vort_track_0 in vort_tracks_0:
            for vort_track_1 in vort_tracks_1:<commit_msg>Fix typeo.<commit_after>from collections import defaultdict, OrderedDict

import numpy as np
import pylab as plt

from tracking import VortMax, VortMaxTrack
from utils.utils import dist
from utils.kalman import RTSSmoother, _plot_rts_smoother

CUM_DIST_CUTOFF = 100


class EnsembleMatch(object):
    '''Represents one match between many vorticity tracks'''
    def __init__(self, vort_track, store_all_tracks=False):
        self.av_vort_track = vort_track

        # May have to disable due to memory constraints.
        self.store_all_tracks = store_all_tracks
        if self.store_all_tracks:
            self.vort_tracks = []
            self.vort_tracks.append(vort_track)

        self.cum_dist = 0
        self.overlap = 0

        self.dates = vort_track.dates
        self.date_start = vort_track.dates[0]
        self.date_end = vort_track.dates[-1]

        index_start = np.where(vort_track.dates == self.date_start)[0][0]
        index_end = np.where(vort_track.dates == self.date_end)[0][0]

        self.tracks_added = np.ones(index_end - index_start + 1)

    def add_track(self, vort_track):
        '''Checks to see whether the track should be added, then adds it

        Compares the overlap and cumalitive distance between the track to the current average
        track, and if they are below a limit adds the track, which involves recalculating the
        average track for this match, and updating some fields.

        :param vort_track: vorticitiy track to be added
        :returns: True if it was added, otherwise False
        '''
        # Calculate some dates and offsets.
        overlap_start = max(self.date_start, vort_track.dates[0])
        overlap_end = min(self.date_end, vort_track.dates[-1])

        index1 = np.where(self.dates == overlap_start)[0][0]
        index2 = np.where(vort_track.dates == overlap_start)[0][0]

        end1 = np.where(self.dates == overlap_end)[0][0]
        end2 = np.where(vort_track.dates == overlap_end)[0][0]

        assert (end1 - index1) == (end2 - index2)

        overlap, cum_dist = self.__calc_distance(self.av_vort_track, vort_track,
                                                 index1, index2, end1, end2)

        # Check to see whether this vort_track should be added.
        if overlap < 6 or cum_dist / overlap > 6:
            return False

        if self.store_all_tracks:
            self.vort_tracks.append(vort_track)

        # import ipdb; ipdb.set_trace()
        new_tracks_added = np.zeros(len(self.dates) + len(vort_track.dates) - overlap)
        copy_start_date = self.dates[0]
        copy_length = len(self.dates)
        ones_start_date = vort_track.dates[0]
        ones_length = len(vort_track.dates)

        self.dates = np.array(sorted([d for d in set(self.dates) | set(vort_track.dates)]))
        self.date_start = self.dates[0]
        self.date_end = self.dates[-1]

        # Copy over the existing tracks added to the new tracks added.
        # new_tracks_added must be equal or greater in length than current.
        # index it using the greatest of index1, and index2.
        copy_index = np.where(self.dates == copy_start_date)[0][0]
        ones_index = np.where(self.dates == ones_start_date)[0][0]

        new_tracks_added[copy_index:copy_index + copy_length] = self.tracks_added
        new_tracks_added[ones_index:ones_index + ones_length] += np.ones(ones_length)

        self.tracks_added = new_tracks_added

        # Update some fields.
        self.overlap += overlap
        self.cum_dist += cum_dist

        vortmaxes = []
        # Update the average vort track by taking a weighted average of the new
        # track and the current average. (N.B. self.tracks_added not yet updated.)
        # import ipdb; ipdb.set_trace()
        for i, date in enumerate(self.dates):
            vm1, vm2 = None, None
            if date in self.av_vort_track.vortmax_by_date:
                vm1 = self.av_vort_track.vortmax_by_date[date]
            if date in vort_track.vortmax_by_date:
                vm2 = vort_track.vortmax_by_date[date]

            if vm1 and vm2:
                # index tracks_added using min of index1, index2
                tracks_added = self.tracks_added[i]

                # Calculate a weighted average of the averaged pos and the new pos.
                pos = tuple([((tracks_added - 1) * v1 + v2) / (tracks_added)
                             for v1, v2 in zip(vm1.pos, vm2.pos)])

                # Calculate a weighted average of the averaged vort and the new vort.
                vort = ((tracks_added - 1) * vm1.vort) / (tracks_added)

                new_vm = VortMax(date, pos, vort)
                index1 += 1
                index2 += 1
            elif vm1:
                new_vm = VortMax(vm1.date, vm1.pos, vm1.vort)
                index1 += 1
            elif vm2:
                new_vm = VortMax(vm2.date, vm2.pos, vm2.vort)
                index2 += 1
            else:
                import ipdb; ipdb.set_trace()
                assert False, 'Should never be reached'

            vortmaxes.append(new_vm)

        next_vm = None
        for vm in vortmaxes[::-1]:
            if next_vm:
                vm.next_vortmax.append(next_vm)
            next_vm = vm
        self.av_vort_track = VortMaxTrack(vm)
        return True

    def __calc_distance(self,
                        vort_track_0, vort_track_1,
                        index1, index2, end1, end2):
        overlap = end1 - index1 + 1

        cum_dist = 0

        while index1 <= end1 and index2 <= end2:
            pos1 = vort_track_0.vortmaxes[index1].pos
            pos2 = vort_track_1.vortmaxes[index2].pos

            cum_dist += dist(pos1, pos2)
            index1 += 1
            index2 += 1

        return overlap, cum_dist

    def av_dist(self):
        '''Returns the average distance between all vorticity tracks'''
        return self.cum_dist / self.overlap


class Match(object):
    '''Represents one match between a best track and a vorticity track'''
    def __init__(self, best_track, vort_track):
        self.best_track = best_track
        self.vort_track = vort_track
        self.cum_dist = 0
        self.overlap = 1
        self.is_too_far_away = False
        self.overlap_start = max(best_track.dates[0], vort_track.dates[0])
        self.overlap_end = min(best_track.dates[-1], vort_track.dates[-1])

    def av_dist(self):
        '''Returns the average distance between the best and vorticity tracks'''
        return self.cum_dist / self.overlap


def match_vort_tracks(vort_tracks_by_dates):
    matches = OrderedDict()

    # import ipdb; ipdb.set_trace()
    vort_tracks_by_date_0 = vort_tracks_by_dates[0]
    vort_tracks_by_date_1 = vort_tracks_by_dates[1]

    all_dates = set(vort_tracks_by_date_0.keys()) | set(vort_tracks_by_date_1.keys())

    for date in all_dates:
        # This means that every vort track is represented **multiple times**, once
        # for each date where it appears. Hence the need to check for the key in matches dict.
        vort_tracks_0 = vort_tracks_by_date_0[date]
        vort_tracks_1 = vort_tracks_by_date_1[date]

        for vort_track_0 in vort_tracks_0:
            for vort_track_1 in vort_tracks_1:
{'fraction_non_alphanumeric': 0.08492583649534322, 'fraction_numerical': 0.01738530527768196, 'mean_word_length': 2.216330153095185, 'pattern_counts': {'":': 0, '<': 9, '<?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': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4513426', 'n_tokens_mistral': 4765, 'n_tokens_neox': 4424, 'n_words': 1168}