<commit_before>from collections import OrderedDict
import datetime as dt

import numpy as np
import pandas as pd

from logger import get_logger
from utils.kalman import KalmanFilter
from utils.utils import dist, geo_dist, pairwise, find_extrema
from tracking import VortMax, VortMaxTrack, CycloneTrack

log = get_logger('analysis.tracking', console_level_str='INFO')

NUM_ENSEMBLE_MEMBERS = 56


class FullVortmaxFinder(object):
    '''Finds all vortmaxes across ensemble members'''
    def __init__(self, fc20data):
        self.fc20data = fc20data

        # Some settings to document/consider playing with.
        self.use_vort_cuttoff = True
        self.use_dist_cuttoff = False
        self.use_range_cuttoff = True
        self.use_geo_dist = True

        if self.use_geo_dist:
            self.dist = geo_dist
            self.dist_cutoff = geo_dist((0, 0), (2, 0)) * 5
        else:
            self.dist = dist
            self.dist_cutoff = 5

        # self.vort_cutoff = 5e-5 # Old value with wrong vort calc.
        # self.vort_cutoff = 2.5e-5
        self.vort_cutoff = 1e-5

    def find_vort_maxima(self, start_date, end_date):
        '''Runs over the date range looking for all vorticity maxima'''
        if start_date < self.fc20data.dates[0]:
            raise Exception('Start date is out of date range, try setting the year appropriately')
        elif end_date > self.fc20data.dates[-1]:
            raise Exception('End date is out of date range, try setting the year appropriately')
        index = np.where(self.fc20data.dates == start_date)[0][0]
        end_index = np.where(self.fc20data.dates == end_date)[0][0]

        self.all_vortmax_time_series = []
        results = []

        for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
            self.all_vortmax_time_series.append(OrderedDict())

        while index <= end_index:
            date = self.fc20data.dates[index]
            self.fc20data.set_date(date)
            start = dt.datetime.now()

            log.info('Finding vortmaxima: {0}'.format(date))

            for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
                vortmax_time_series = self.all_vortmax_time_series[ensemble_member]
                vortmaxes = []
                vmaxs = self.fc20data.vmaxs[ensemble_member]

                for vmax in vmaxs:
                    if self.use_range_cuttoff and not (260 < vmax[1][0] < 340 and
                       0 < vmax[1][1] < 60):
                        continue

                    if self.use_vort_cuttoff and vmax[0] < self.vort_cutoff:
                        continue

                    vortmax = VortMax(date, vmax[1], vmax[0])
                    vortmaxes.append(vortmax)

                if self.use_dist_cuttoff:
                    secondary_vortmaxes = []
                    for i in range(len(vortmaxes)):
                        v1 = vortmaxes[i]
                        for j in range(i + 1, len(vortmaxes)):
                            v2 = vortmaxes[j]
                            if self.dist(v1.pos, v2.pos) < self.dist_cutoff:
                                if v1.vort > v2.vort:
                                    v1.secondary_vortmax.append(v2)
                                    secondary_vortmaxes.append(v2)
                                elif v1.vort <= v2.vort:
                                    v2.secondary_vortmax.append(v1)
                                    secondary_vortmaxes.append(v1)

                    for v in secondary_vortmaxes:
                        if v in vortmaxes:
                            vortmaxes.remove(v)

                vortmax_time_series[date] = vortmaxes

                for vortmax in vortmaxes:
                    row = {'date': date,
                           'em': ensemble_member,
                           'lon': vortmax.pos[0],
                           'lat': vortmax.pos[1],
                           'vort': vortmax.vort}
                    res = self.get_other_fields(ensemble_member, vortmax)
                    row.update(res)
                    results.append(row)

            print('  Found vortmaxima and fields in {}'.format(end - start))
            index += 1

        columns = ['date', 'em', 
                   'lon', 
                   'lat', 
                   'max_ws_lon',
                   'max_ws_lat',
                   'pmin_lon',
                   'pmin_lat',
                   'vort', 
                   'max_ws',
                   'pmin_dist',
                   'pmin',
                   'p_ambient_diff',
                   't850s',
                   't995s',
                   'capes',
                   'pwats']
        df = pd.DataFrame(results, columns=columns)
        end = dt.datetime.now()
        return df

    def get_other_fields(self, ensemble_member, vortmax):
        res = {}
        vmax_pos = vortmax.pos<commit_msg>Make sure end is set.<commit_after>from collections import OrderedDict
import datetime as dt

import numpy as np
import pandas as pd

from logger import get_logger
from utils.kalman import KalmanFilter
from utils.utils import dist, geo_dist, pairwise, find_extrema
from tracking import VortMax, VortMaxTrack, CycloneTrack

log = get_logger('analysis.tracking', console_level_str='INFO')

NUM_ENSEMBLE_MEMBERS = 56


class FullVortmaxFinder(object):
    '''Finds all vortmaxes across ensemble members'''
    def __init__(self, fc20data):
        self.fc20data = fc20data

        # Some settings to document/consider playing with.
        self.use_vort_cuttoff = True
        self.use_dist_cuttoff = False
        self.use_range_cuttoff = True
        self.use_geo_dist = True

        if self.use_geo_dist:
            self.dist = geo_dist
            self.dist_cutoff = geo_dist((0, 0), (2, 0)) * 5
        else:
            self.dist = dist
            self.dist_cutoff = 5

        # self.vort_cutoff = 5e-5 # Old value with wrong vort calc.
        # self.vort_cutoff = 2.5e-5
        self.vort_cutoff = 1e-5

    def find_vort_maxima(self, start_date, end_date):
        '''Runs over the date range looking for all vorticity maxima'''
        if start_date < self.fc20data.dates[0]:
            raise Exception('Start date is out of date range, try setting the year appropriately')
        elif end_date > self.fc20data.dates[-1]:
            raise Exception('End date is out of date range, try setting the year appropriately')
        index = np.where(self.fc20data.dates == start_date)[0][0]
        end_index = np.where(self.fc20data.dates == end_date)[0][0]

        self.all_vortmax_time_series = []
        results = []

        for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
            self.all_vortmax_time_series.append(OrderedDict())

        while index <= end_index:
            date = self.fc20data.dates[index]
            self.fc20data.set_date(date)
            start = dt.datetime.now()

            log.info('Finding vortmaxima: {0}'.format(date))

            for ensemble_member in range(NUM_ENSEMBLE_MEMBERS):
                vortmax_time_series = self.all_vortmax_time_series[ensemble_member]
                vortmaxes = []
                vmaxs = self.fc20data.vmaxs[ensemble_member]

                for vmax in vmaxs:
                    if self.use_range_cuttoff and not (260 < vmax[1][0] < 340 and
                       0 < vmax[1][1] < 60):
                        continue

                    if self.use_vort_cuttoff and vmax[0] < self.vort_cutoff:
                        continue

                    vortmax = VortMax(date, vmax[1], vmax[0])
                    vortmaxes.append(vortmax)

                if self.use_dist_cuttoff:
                    secondary_vortmaxes = []
                    for i in range(len(vortmaxes)):
                        v1 = vortmaxes[i]
                        for j in range(i + 1, len(vortmaxes)):
                            v2 = vortmaxes[j]
                            if self.dist(v1.pos, v2.pos) < self.dist_cutoff:
                                if v1.vort > v2.vort:
                                    v1.secondary_vortmax.append(v2)
                                    secondary_vortmaxes.append(v2)
                                elif v1.vort <= v2.vort:
                                    v2.secondary_vortmax.append(v1)
                                    secondary_vortmaxes.append(v1)

                    for v in secondary_vortmaxes:
                        if v in vortmaxes:
                            vortmaxes.remove(v)

                vortmax_time_series[date] = vortmaxes

                for vortmax in vortmaxes:
                    row = {'date': date,
                           'em': ensemble_member,
                           'lon': vortmax.pos[0],
                           'lat': vortmax.pos[1],
                           'vort': vortmax.vort}
                    res = self.get_other_fields(ensemble_member, vortmax)
                    row.update(res)
                    results.append(row)

            end = dt.datetime.now()
            print('  Found vortmaxima and fields in {}'.format(end - start))
            index += 1

        columns = ['date', 'em', 
                   'lon', 
                   'lat', 
                   'max_ws_lon',
                   'max_ws_lat',
                   'pmin_lon',
                   'pmin_lat',
                   'vort', 
                   'max_ws',
                   'pmin_dist',
                   'pmin',
                   'p_ambient_diff',
                   't850s',
                   't995s',
                   'capes',
                   'pwats']
        df = pd.DataFrame(results, columns=columns)
        end = dt.datetime.now()
        return df

    def get_other_fields(self, ensemble_member, vortmax):
        res = {}
        vmax_pos = vortmax.pos
{'fraction_non_alphanumeric': 0.0868239121594144, 'fraction_numerical': 0.01667344448962993, 'mean_word_length': 1.4116204952194165, 'pattern_counts': {'":': 0, '<': 21, '<?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': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1979085', 'n_tokens_mistral': 3047, 'n_tokens_neox': 2796, 'n_words': 617}