EK-Thesis/data_generation.py at main · ekollrack/EK-Thesis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
"""
Generates data to be used for models. Play by play game data
and game result data will be used to predict attention
Last updated 1/15 by EK
"""

import pandas as pd
from datetime import timedelta
from data_mountain_query.query import get_ambient_tweets
from data_mountain_query.connection import get_connection
import time


def main():
    start_time = time.time()

    # Load games
    games = pd.read_csv("/Users/elisabethkollrack/Thesis/EK-thesis/games.csv")
    games['gameday'] = pd.to_datetime(games['gameday'], format='%m/%d/%y')

    # Load cumulative win percentages
    win_pct = pd.read_csv('/Users/elisabethkollrack/Thesis/EK-Thesis/R data/nfl_r_data.csv')

    # Merge win_pct into games on 'game_id'
    games = games.merge(
        win_pct[['game_id', 'home_win_pct', 'away_win_pct', 'lead_changes']],
        on='game_id',
        how='left'
    )

    p = 1.0
    collection, client = get_connection(user_loc=True)

    data_rows = []  # store data for final DataFrame

    # Loop over each season 2013–2017
    for season in range(2010, 2015):
        # Filter the games for that season
        games_season = games[games['season'] == season]

        # Unique matchups (order-insensitive)
        games_season['matchup_key'] = games_season.apply(
            lambda r: "_vs_".join(sorted([r['away_team'], r['home_team']])), axis=1
        )

        for _, row in games_season.iterrows():
            team1, team2 = row['away_team'], row['home_team']
            gameday = row['gameday']
            total_score = row['total']
            score_diff = row['result']

            anchors = [f"#{team1}vs{team2}", f"#{team2}vs{team1}"]
            start_date = gameday - timedelta(days=3)
            end_date = gameday + timedelta(days=3)
            dates = pd.date_range(start_date, end_date, freq='D')

            counts = {}
            for anchor in anchors:
                tweets_list = [t for t in get_ambient_tweets(anchor, dates, collection)]
                for tweet in tweets_list:
                    tweet_day = pd.to_datetime(tweet['tweet_created_at']).date()
                    counts[tweet_day] = counts.get(tweet_day, 0) + 1

            # total attention for this game (sum over the ±3 days)
            attention = sum(counts.get((gameday + timedelta(days=d)).date(), 0) for d in range(-3, 4))

            data_rows.append({
                'date': gameday.date(),
                'game_id': row['game_id'],
                'season': row['season'],
                'week': row['week'],
                'weekday': row['weekday'],
                'gametime': row['gametime'],
                'home_team': row['home_team'],
                'away_team': row['away_team'],
                'home_win_pct': row['home_win_pct'],
                'away_win_pct': row['away_win_pct'],
                'num_lead_changes': row['lead_changes'],
                'total_score': total_score,
                'score_differential': score_diff,
                'overtime': int(row['overtime']),
                'attention': attention
            })

    # Create final DataFrame
    df = pd.DataFrame(data_rows)

    # Save to CSV
    output_path = "/Users/elisabethkollrack/Thesis/EK-thesis/game_attention.csv"
    df.to_csv(output_path, index=False)

    print(df.head())  # preview
    print(f"\nSaved to: {output_path}")
    print(f"Total rows: {len(df)}")

    end_time = time.time()
    elapsed = end_time - start_time
    print(f"\nTotal runtime: {elapsed:.2f} seconds")

    return df

if __name__ == '__main__':
    main()