add win probability feature to awpy2

awpy/stats/__init__.py

@@ -3,5 +3,6 @@
 from awpy.stats.adr import adr
 from awpy.stats.kast import calculate_trades, kast
 from awpy.stats.rating import impact, rating
+from awpy.stats.win_prob import win_probability 
 
-__all__ = ["adr", "calculate_trades", "kast", "impact", "rating"]
+__all__ = ["adr", "calculate_trades", "kast", "impact", "rating", "win_probability"]

awpy/stats/win_prob.py

@@ -0,0 +1,113 @@
+"""Calculates CT & T side likeliehood of winning a round at any given point"""
+
+from typing import List, Dict, Union
+import pandas as pd 
+import numpy as np
+
+from awpy import Demo
+
+def generate_random_weights(num_features):
+    return np.random.uniform(low=-1.0, high=1.0, size=num_features)
+
+
+def process_tick_data(tick_data: pd.DataFrame, demo: Demo) -> Dict[str, Union[int, str, bool]]:
+    """
+    Processes individual tick data to extract game state features.
+
+    Args:
+        tick_data (pd.DataFrame): DataFrame containing all entries for a single tick.
+        demo (Demo): A parsed demo object containing CS2 match data.
+
+    Returns:
+        Dict[str, Union[int, str, bool]]: Dictionary containing the game state features for the tick.
+    """
+    round_number = tick_data['round'].iloc[0]
+    map_name = demo.header.map_name
+    bomb_planted = tick_data['is_bomb_planted'].iloc[0]
+    players_alive_ct = tick_data[(tick_data['side'] == 'CT') & (tick_data['health'] > 0)]['steamid'].nunique()
+    players_alive_t = tick_data[(tick_data['side'] == 'TERRORIST') & (tick_data['health'] > 0)]['steamid'].nunique()
+    equipment_value_ct = tick_data[(tick_data['side'] == 'CT') & (tick_data['health'] > 0)]['current_equip_value'].sum()
+    equipment_value_t = tick_data[(tick_data['side'] == 'TERRORIST') & (tick_data['health'] > 0)]['current_equip_value'].sum()
+    hp_remaining_ct = tick_data[(tick_data['side'] == 'CT') & (tick_data['health'] > 0)]['health'].sum()
+    hp_remaining_t = tick_data[(tick_data['side'] == 'TERRORIST') & (tick_data['health'] > 0)]['health'].sum()
+    armor_ct = tick_data[(tick_data['side'] == 'CT') & (tick_data['health'] > 0) & (tick_data['armor_value'] > 0)].shape[0]
+    armor_t = tick_data[(tick_data['side'] == 'TERRORIST') & (tick_data['health'] > 0) & (tick_data['armor_value'] > 0)].shape[0]
+    has_helmet_ct = tick_data[(tick_data['side'] == 'CT') & (tick_data['health'] > 0) & (tick_data['has_helmet'] == True)].shape[0]
+    has_helmet_t = tick_data[(tick_data['side'] == 'TERRORIST') & (tick_data['health'] > 0) & (tick_data['has_helmet'] == True)].shape[0]
+
+
+    return {
+        "tick": tick_data['tick'].iloc[0],
+        "round": round_number,
+        "map_name": map_name,
+        "bomb_planted": bomb_planted,
+        "players_alive_ct": players_alive_ct,
+        "players_alive_t": players_alive_t,
+        "equipment_value_ct": equipment_value_ct,
+        "equipment_value_t": equipment_value_t,
+        "hp_remaining_ct": hp_remaining_ct,
+        "hp_remaining_t": hp_remaining_t,
+        "armor_ct": armor_ct,
+        "armor_t": armor_t,
+        "has_helmet_ct": has_helmet_ct,
+        "has_helmet_t": has_helmet_t
+    }
+
+def build_feature_matrix(demo: Demo, ticks: Union[int, List[int]]) -> List[Dict[str, Union[int, str, bool]]]:
+    """
+    Builds the game state feature matrix for specified ticks using a more efficient pandas apply method.
+
+    Args:
+        demo (Demo): A parsed demo object containing CS2 match data.
+        ticks (Union[int, List[int]]): A single tick or a list of ticks at which to calculate features.
+
+    Returns:
+        List[Dict[str, Union[int, str, bool]]]: A list of dictionaries where each dictionary contains the game state features at a given tick.
+
+    Raises:
+        ValueError: If specified ticks are not found within the demo object.
+    """
+    if demo.ticks.empty:
+        raise ValueError("Demo object does not contain tick data.")
+
+    if isinstance(ticks, int):
+        ticks = [ticks]
+
+    # Filtering ticks data for the specified ticks
+    filtered_ticks = demo.ticks[demo.ticks["tick"].isin(ticks)]
+
+    if filtered_ticks.empty:
+        raise ValueError("No data found for specified ticks.")
+
+    # Applying the external function to each group of tick data
+    game_state = filtered_ticks.groupby('tick').apply(lambda x: process_tick_data(x, demo)).tolist()
+    return game_state
+
+
+def win_probability(demo: Demo, ticks: Union[int, List[int]]) -> List[Dict[str, float]]:
+    """
+    Calculate win probabilities for specified ticks in a CS2 match demo.
+
+    Args:
+        demo (Demo): A parsed demo object containing CS2 match data.
+        ticks (Union[int, List[int]]): A single tick or a list of ticks at which to calculate win probabilities.
+
+    Returns:
+        List[Dict[str, float]]: A list of dictionaries with the calculated win probabilities for CT and T sides for each tick.
+
+    Raises:
+        NotImplementedError: This function has not yet been implemented.
+    """
+    feature_matrix = build_feature_matrix(demo, ticks)
+    mock_weights = generate_random_weights(12)
+    probabilities = []
+    for features in feature_matrix:
+        win_prob_ct = 0.50 
+        probabilities.append({
+            "tick": features["tick"],
+            "CT_win_probability": win_prob_ct,
+            "T_win_probability": 1 - win_prob_ct,
+        })
+    return probabilities
+
+

tests/test_win_prob.py

@@ -0,0 +1,49 @@
+"""Tests Win Probability module."""
+
+import pytest
+
+from awpy import Demo
+from awpy.stats import win_probability
+
+class TestWinProbability:
+    """Tests the win probability calculations.
+
+    https://www.hltv.org/matches/2369248/natus-vincere-vs-virtuspro-pgl-cs2-major-copenhagen-2024-europe-rmr-closed-qualifier-a
+    """
+
+    @classmethod
+    def setup_class(cls):
+        """
+        Setup method called before any tests are run.
+        Initializes the Demo object and random weights for testing.
+        """
+        cls.demo = Demo(file='tests/natus-vincere-vs-virtus-pro-m1-overpass.dem')  
+
+
+    @classmethod
+    def teardown_class(cls):
+        """
+        Teardown method called after all tests have run.
+        """
+        cls.demo = None
+
+
+    def test_win_probability_symmetry(self):
+        """Test to ensure P(CT Win) = 1 - P(T Win) at a given tick."""
+        probabilities = win_probability(self.demo, 168200)
+        for prob in probabilities:
+            assert prob["CT_win_probability"] == pytest.approx(1 - prob["T_win_probability"], 0.01)
+
+
+    def test_known_ct_sided_situation(self):
+        """Test for a known heavily CT-sided situation."""
+        probabilities = win_probability(self.demo, 168200)
+        assert probabilities[0]["tick"] == 168200
+        assert probabilities[0]["CT_win_probability"] > 0.7
+
+
+    def test_prediction_count_matches_ticks(self):
+        """Test that the number of predictions matches the number of ticks passed in."""
+        ticks = [166100, 167023, 168200]
+        probabilities = win_probability(self.demo, ticks)
+        assert len(probabilities) == len(ticks)