Merge 797bdb2 into ba42122

trueskill/__init__.py

@@ -660,6 +660,40 @@ def quality_1vs1(rating1, rating2, env=None):
     return env.quality([(rating1,), (rating2,)])
 
 
+def probability_NvsM(team1_ratings, team2_ratings, env=None):
+    """Calculates the win probability of the first team over the second team.
+
+    :param team1_ratings: ratings of the first team participants.
+    :param team2_ratings: ratings of another team participants.
+    :param env: the :class:`TrueSkill` object.  Defaults to the global
+                environment.
+    """
+    if env is None:
+        env = global_env()
+
+    team1_mu = sum(r.mu for r in team1_ratings)
+    team1_sigma = sum((env.beta**2 + r.sigma**2) for r in team1_ratings)
+    team2_mu = sum(r.mu for r in team2_ratings)
+    team2_sigma = sum((env.beta**2 + r.sigma**2) for r in team2_ratings)
+
+    x = (team1_mu - team2_mu) / math.sqrt(team1_sigma + team2_sigma)
+    probability_win_team1 = env.cdf(x)
+    return probability_win_team1
+
+
+def probability_1vs1(rating1, rating2, env=None):
+    """A shortcut to calculate the win probability between just 2 players in
+    a head-to-head match
+
+    :param rating1: the rating.
+    :param rating2: the another rating.
+    :param env: the :class:`TrueSkill` object.  Defaults to the global
+                environment.
+    :return: probability the first player wins
+    """
+    return probability_NvsM((rating1,), (rating2,), env=env)
+
+
 def global_env():
     """Gets the :class:`TrueSkill` object which is the global environment."""
     try:

trueskilltest.py

@@ -9,7 +9,7 @@
 from conftest import various_backends
 import trueskill as t
 from trueskill import (
-    quality, quality_1vs1, rate, rate_1vs1, Rating, setup, TrueSkill)
+    quality, quality_1vs1, rate, rate_1vs1, Rating, setup, TrueSkill, probability_NvsM, probability_1vs1)
 
 
 warnings.simplefilter('always')
@@ -206,6 +206,22 @@ def test_backend():
         TrueSkill(backend='__not_defined__')
 
 
+def test_probability_NvsM():
+    def assert_almost_equal(value, correct):
+        eps = 1e-3
+        assert abs(value - correct) < eps
+
+    r1, r2 = Rating(), Rating()
+    assert_almost_equal(probability_1vs1(r1, r2), 0.5)
+
+    r1, r2 = rate_1vs1(r1, r2)
+    assert_almost_equal(probability_1vs1(r1, r2), 0.773)
+    assert_almost_equal(probability_1vs1(r2, r1) + probability_1vs1(r1, r2), 1.0)
+    assert_almost_equal(probability_NvsM((r1, r1, r1), (r2, r2, r2)), 0.903)
+    assert_almost_equal(probability_NvsM((r1, r1, r1), (r1, r2, r2)), 0.807)
+    assert_almost_equal(probability_NvsM((r2,), (r1, r2)), 0.02)
+
+
 # algorithm