forked from sublee/trueskill
-
Notifications
You must be signed in to change notification settings - Fork 0
/
trueskillhelpers.py
155 lines (144 loc) · 5.73 KB
/
trueskillhelpers.py
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
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
# -*- coding: utf-8 -*-
from __future__ import with_statement
from contextlib import contextmanager
import functools
import inspect
import logging
import trueskill
from trueskill.backends import available_backends
from trueskill.factorgraph import (Variable, Factor, PriorFactor,
LikelihoodFactor, SumFactor)
from trueskill.mathematics import Gaussian
__all__ = ['substituted_trueskill', 'calc_dynamic_draw_probability',
'factor_graph_logging']
@contextmanager
def substituted_trueskill(*args, **kwargs):
"""Setup the global environment only within the context::
assert Rating().mu == 25
with substituted_trueskill(mu=0):
assert Rating().mu == 0
"""
env = trueskill.global_env()
params = [['mu', env.mu], ['sigma', env.sigma], ['beta', env.beta],
['tau', env.tau], ['draw_probability', env.draw_probability],
['backend', env.backend]]
# merge settings with previous TrueSkill object
for x, arg in enumerate(args):
params[x][1] = arg
params = dict(params)
for kw, arg in kwargs.items():
params[kw] = arg
try:
# setup the environment
yield trueskill.setup(**params)
finally:
# revert the environment
trueskill.setup(env=env)
def calc_dynamic_draw_probability(rating_group1, rating_group2, env=None):
from trueskill.factorgraph import (Variable, )
if env is None:
env = global_env()
team_perf_vars = []
for rating_group in [rating_group1, rating_group2]:
team_perf_var = Variable()
team_perf_vars.append(team_perf_var)
perf_vars = []
for rating in rating_group:
rating_var, perf_var = Variable(), Variable()
perf_vars.append(perf_var)
PriorFactor(rating_var, rating, env.tau).down()
LikelihoodFactor(rating_var, perf_var, env.beta ** 2).down()
SumFactor(team_perf_var, perf_vars, [1] * len(perf_vars)).down()
return env.draw_probability(trueskill.Rating(team_perf_vars[0]),
trueskill.Rating(team_perf_vars[1]), env)
@contextmanager
def factor_graph_logging(color=False):
"""In the context, a factor graph prints logs as DEBUG level. It will help
to follow factor graph running schedule::
with factor_graph_logging() as logger:
logger.setLevel(DEBUG)
logger.addHandler(StreamHandler(sys.stderr))
rate_1vs1(Rating(), Rating())
"""
import inspect
# color mode uses the termcolor module
if color:
try:
from termcolor import colored
except ImportError:
raise ImportError('To enable color mode, install termcolor')
else:
colored = lambda s, *a, **k: s
logger = logging.getLogger('TrueSkill')
orig_factor_init = Factor.__init__
orig_variable_set = Variable.set
def repr_factor(factor):
return '{}@{}'.format(type(factor).__name__, id(factor))
def repr_gauss(gauss):
return 'N(mu=%.3f, sigma=%.3f, pi=%r, tau=%r)' % \
(gauss.mu, gauss.sigma, gauss.pi, gauss.tau)
def r(val):
if isinstance(val, Factor):
return repr_factor(val)
elif isinstance(val, Gaussian):
return repr_gauss(val)
else:
return repr(val)
def factor_init(self, *args, **kwargs):
frames = inspect.getouterframes(inspect.currentframe())
layer_builder_name = frames[2][3]
assert (layer_builder_name.startswith('build_') and
layer_builder_name.endswith('_layer'))
self._layer_name = layer_builder_name[6:].replace('_', ' ').title()
return orig_factor_init(self, *args, **kwargs)
def variable_set(self, val):
old_value = Gaussian(pi=self.pi, tau=self.tau)
old_messages = {fac: Gaussian(pi=msg.pi, tau=msg.tau)
for fac, msg in self.messages.items()}
delta = orig_variable_set(self, val)
# inspect outer frames
frames = inspect.getouterframes(inspect.currentframe())
methods = [None, None]
for frame in frames:
method = frame[3]
if method.startswith('update_'):
methods[0] = method
elif method in ('up', 'down'):
methods[1] = method
break
factor = frame[0].f_locals['self']
before = Gaussian(pi=self.pi, tau=self.tau)
# helpers for logging
logs = []
l = logs.append
bullet = lambda changed: colored(' * ', 'red') if changed else ' '
# print layer
if getattr(logger, '_prev_layer_name', None) != factor._layer_name:
logger._prev_layer_name = factor._layer_name
l(colored('[{}]'.format(factor._layer_name), 'blue'))
# print factor
l(colored('<{}.{}>'.format(r(factor), methods[1]), 'cyan'))
# print value
if old_value == self:
line = '{}'.format(r(self))
else:
line = '{} -> {}'.format(r(old_value), r(self))
l(bullet(methods[0] == 'update_value') + line)
# print messages
fmt = '{}: {} -> {}'.format
for fac, msg in self.messages.items():
old_msg = old_messages[fac]
changed = fac is factor and methods[0] == 'update_message'
if old_msg == msg:
line = '{}: {}'.format(r(fac), r(msg))
else:
line = '{}: {} -> {}'.format(r(fac), r(old_msg), r(msg))
l(bullet(changed) + line)
# print buffered logs
map(logger.debug, logs)
return delta
try:
Factor.__init__, Variable.set = factor_init, variable_set
yield logger
finally:
Factor.__init__, Variable.set = orig_factor_init, orig_variable_set