From 7a4100946453f004bfadb0113f2a1dc9307de5da Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Thu, 1 Oct 2020 22:36:00 +0200
Subject: [PATCH 1/7] change TS signature to support different sample sizes

---
 pyrff/test_thompson.py | 11 ++++++-----
 pyrff/thompson.py      | 36 ++++++++++++++++++++++++------------
 2 files changed, 30 insertions(+), 17 deletions(-)

diff --git a/pyrff/test_thompson.py b/pyrff/test_thompson.py
index 9441710..23cb3bc 100644
--- a/pyrff/test_thompson.py
+++ b/pyrff/test_thompson.py
@@ -15,11 +15,11 @@ def test_sample_batch(self, batch_size, seed):
             low=[0, 0, -1],
             high=[0.2, 1, 0],
             size=(S, C)
-        )
-        numpy.testing.assert_array_equal(samples.shape, (S, C))
+        ).T
+        numpy.testing.assert_array_equal(samples.shape, (C, S))
 
         batch = thompson.sample_batch(
-            samples=samples, ids=ids,
+            candidate_samples=samples, ids=ids,
             batch_size=batch_size, seed=seed
         )
         assert len(batch) == batch_size
@@ -29,10 +29,11 @@ def test_sample_batch(self, batch_size, seed):
         pass
 
     def test_no_bias_on_sample_collisions(self):
-        samples = numpy.array([
+        samples = [
             [2, 2, 2],
+            [2, 2],
             [2, 2, 2],
-        ])
+        ]
         batch = thompson.sample_batch(samples, ids=('A', 'B', 'C'), batch_size=100, seed=1234)
         assert batch.count('A') != 100
         assert batch.count('C') != 0
diff --git a/pyrff/thompson.py b/pyrff/thompson.py
index 9fb3dc4..cecd634 100644
--- a/pyrff/thompson.py
+++ b/pyrff/thompson.py
@@ -4,31 +4,43 @@
 import numpy
 import typing
 
+# custom type shortcuts
+Sample = typing.Union[int, float]
+
 
 def sample_batch(
-    samples:numpy.ndarray, *, ids:typing.Sequence,
+    candidate_samples: typing.Sequence[typing.Sequence[Sample]], *,
+    ids:typing.Sequence,
     batch_size:int, seed:typing.Optional[int]=None
 ) -> tuple:
     """
-    Ranks all candidates by their samples.
+    Draws a batch of candidates by Thompson Sampling from posterior samples.
 
     Parameters
     ----------
-        samples : numpy.ndarray
-            (S, C) array of posterior samples (S) for each candidate (C)
-        ids : numpy.ndarray, list, tuple
-            (C,) candidate identifiers
-        batch_size : int
-            size of the next measurement batch (B)
+    candidate_samples : array-like
+        posterior samples for each candidate (C,)
+        (sample count may be different)
+    ids : numpy.ndarray, list, tuple
+        (C,) candidate identifiers
+    batch_size : int
+        number of candidates to draw (B)
 
     Returns
     -------
-        chosen_candidates : tuple
-            (B,) chosen candidate ids for the batch
+    chosen_candidates : tuple
+        (B,) chosen candidate ids for the batch
     """
-    n_samples, n_candidates = samples.shape
-    assert len(ids) == n_candidates
+    n_candidates = len(candidate_samples)
+    n_samples = tuple(map(len, candidate_samples))
+    if len(ids) != n_candidates:
+        raise ValueError(f"Number of candidate ids ({len(ids)}) does not match number of candidate_samples ({n_candidates}).")
     ids = numpy.atleast_1d(ids)
+    # work with matrix even if sample count is varies to get more efficient slicing
+    samples = numpy.zeros((max(n_samples), n_candidates))
+    samples[:] = numpy.nan
+    for c, (samps, s) in enumerate(zip(candidate_samples, n_samples)):
+        samples[:s, c] = samps
     chosen_candidates = []
     random = numpy.random.RandomState(seed)
     for i in range(batch_size):

From fa4e079e7662b295f2f836ffcb64ff77349362cd Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Thu, 1 Oct 2020 22:39:19 +0200
Subject: [PATCH 2/7] remove brute-force probability calculation

---
 pyrff/test_thompson.py | 14 --------------
 pyrff/thompson.py      | 25 -------------------------
 2 files changed, 39 deletions(-)

diff --git a/pyrff/test_thompson.py b/pyrff/test_thompson.py
index 23cb3bc..9133952 100644
--- a/pyrff/test_thompson.py
+++ b/pyrff/test_thompson.py
@@ -39,17 +39,3 @@ def test_no_bias_on_sample_collisions(self):
         assert batch.count('C') != 0
         pass
 
-    @pytest.mark.xfail(reason='Probabilities are currently computed by brute force and non-exact.')
-    def test_get_probabilities_exact_on_identical(self):
-        samples = numpy.array([
-            [1, 2, 3, 4, 5],
-            [5, 3, 4, 2, 1],
-            [1, 3, 4, 2, 5]
-        ]).T
-        S, C = samples.shape
-        assert S == 5
-        assert C == 3
-
-        probabilities = thompson.get_probabilities(samples)
-        numpy.testing.assert_array_equal(probabilities, [1/C]*C)
-        pass
diff --git a/pyrff/thompson.py b/pyrff/thompson.py
index cecd634..85c22ce 100644
--- a/pyrff/thompson.py
+++ b/pyrff/thompson.py
@@ -53,28 +53,3 @@ def sample_batch(
         chosen_candidates.append(best_candidate)
     random.seed(None)
     return tuple(chosen_candidates)
-
-
-def get_probabilities(samples:numpy.ndarray, nit:int=100_000):
-    """Get thompson sampling probabilities from posterior.
-
-    Parameters
-    ----------
-        samples : numpy.ndarray
-            (S, C) array of posterior samples (S) for each candidate (C)
-        nit : int
-            how many thompson draws samples to draw for the estimation
-
-    Returns
-    -------
-        probabilities : numpy.ndarray
-            (C,) probabilities that the candidates are sampled
-    """
-    n_samples, n_candidates = samples.shape
-
-    frequencies = numpy.zeros(n_candidates)
-    for _ in range(nit):
-        idx = numpy.random.randint(n_samples, size=n_candidates)
-        selected_samples = samples[idx, numpy.arange(n_candidates)]
-        frequencies[numpy.argmax(selected_samples)] += 1
-    return frequencies / frequencies.sum()

From 388eabccc9e14595594e1c63cf5dd19c69c41fac Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Sat, 3 Oct 2020 22:23:21 +0200
Subject: [PATCH 3/7] implement calculation of exact TS probabilities from
 independent candidate samples closes #3

---
 pyrff/__init__.py      |   4 +-
 pyrff/test_thompson.py |  55 +++++++++++++
 pyrff/thompson.py      | 178 ++++++++++++++++++++++++++++++++++++++++-
 requirements.txt       |   1 +
 4 files changed, 235 insertions(+), 3 deletions(-)

diff --git a/pyrff/__init__.py b/pyrff/__init__.py
index 1796ef9..d9904e8 100644
--- a/pyrff/__init__.py
+++ b/pyrff/__init__.py
@@ -1,6 +1,6 @@
 from . exceptions import DtypeError, ShapeError
 from . rff import sample_rff, save_rffs, load_rffs
-from . thompson import sample_batch, get_probabilities
+from . thompson import sample_batch, sampling_probabilities
 from . utils import multi_start_fmin
 
-__version__ = '1.0.1'
+__version__ = '2.0.0'
diff --git a/pyrff/test_thompson.py b/pyrff/test_thompson.py
index 9133952..a048bb1 100644
--- a/pyrff/test_thompson.py
+++ b/pyrff/test_thompson.py
@@ -39,3 +39,58 @@ def test_no_bias_on_sample_collisions(self):
         assert batch.count('C') != 0
         pass
 
+
+class TestThompsonProbabilities:
+    def test_sort_samples(self):
+        samples, sample_cols = thompson._sort_samples([
+            [3,1,2],
+            [4,-1],
+            [7],
+        ])
+        numpy.testing.assert_array_equal(samples, [-1, 1, 2, 3, 4, 7])
+        numpy.testing.assert_array_equal(sample_cols, [1, 0, 0, 0, 1, 2])
+        pass
+
+    def test_win_draw_prob(self):
+        assert thompson._win_draw_prob(numpy.array([
+            [1, 0, 0],
+            [0, 1, 1],
+            [0, 0, 0],
+        ])) == 0.0
+
+        assert thompson._win_draw_prob(numpy.array([
+            [0, 0, 0],
+            [0, 0, 0],
+            [1, 1, 1],
+        ])) == 0.25
+
+        numpy.testing.assert_allclose(thompson._win_draw_prob(numpy.array([
+            [0, 0],
+            [0.5, 0.75],
+            [0.5, 0.25],
+        ])), 0.041666666)
+        pass
+
+    def test_sampling_probability(self):
+        numpy.testing.assert_array_equal(thompson.sampling_probabilities([
+            [0, 1, 2],
+            [0, 1, 2],
+        ]), [0.5, 0.5])
+        
+        numpy.testing.assert_array_equal(thompson.sampling_probabilities([
+            [0, 1, 2],
+            [10],
+        ]), [0, 1])
+
+        numpy.testing.assert_array_equal(thompson.sampling_probabilities([
+            [0, 1, 2],
+            [3, 4, 5],
+            [5, 4, 3],
+        ]), [0, 0.5, 0.5])
+
+        numpy.testing.assert_array_equal(thompson.sampling_probabilities([
+            [5, 6],
+            [0, 0, 10, 20],
+            [5, 6],
+        ]), [0.25, 0.5, 0.25])
+        pass
\ No newline at end of file
diff --git a/pyrff/thompson.py b/pyrff/thompson.py
index 85c22ce..eee2e11 100644
--- a/pyrff/thompson.py
+++ b/pyrff/thompson.py
@@ -1,8 +1,11 @@
 """
 Convenience implementation of the standard Thompson Sampling algorithm.
 """
-import numpy
+import itertools
 import typing
+import fastprogress
+import numpy
+import pandas
 
 # custom type shortcuts
 Sample = typing.Union[int, float]
@@ -53,3 +56,176 @@ def sample_batch(
         chosen_candidates.append(best_candidate)
     random.seed(None)
     return tuple(chosen_candidates)
+
+
+def _sort_samples(
+    candidate_samples: typing.Sequence[typing.Sequence[Sample]],
+) -> typing.Tuple[numpy.ndarray, numpy.ndarray]:
+    """ Flattens samples into a sorted array and corresponding group numbers.
+
+    Parameters
+    ----------
+    candidate_samples : array-like
+        posterior samples for each candidate (C,)
+        (sample count may be different)
+
+    Returns
+    -------
+    samples : array
+        sorted array of all samples
+    sample_cols : array
+        the group numbers
+    """
+    flat_samples = numpy.vstack([
+        numpy.stack([samps, numpy.repeat(g, len(samps))]).T
+        for g, samps in enumerate(candidate_samples)
+    ])
+    flat_samples = flat_samples[numpy.argsort(flat_samples[:, 0]), :]
+    return flat_samples[:, 0], flat_samples[:, 1].astype(int)
+
+
+def _win_draw_prob(cprobs: numpy.ndarray) -> float:
+    """ Calculate the probability of winning by draw.
+
+    This function iterates over all possible combinations of draws.
+    The runtime complexity explodes exponentially with O(2^N - N - 1).
+
+    Parameters
+    ----------
+    cprobs : numpy.ndarray
+        (3, N) array of probabilities that a win/loose/draw occurs
+        between the candidate value and other candidates (candidate itself is not included in [cprobs])
+
+    Returns
+    -------
+    p_win_draw : float
+        probability of winning in a fair draw against the other candidates
+    """
+    C = cprobs.shape[1]
+    columns = numpy.arange(C)
+    # TODO: this can be further accelerated by not using a DataFrame
+    df_smart_combos = pandas.DataFrame(columns=["p_event", "p_win"])
+
+    drawable = tuple(columns[cprobs[2, :] > 0])
+    p_win_draw = 0
+    for n in range(1, len(drawable) + 1):
+        p_win = 1 / (n + 1)
+        for combination in itertools.combinations(drawable, r=n):
+            draw_probs = cprobs[2, list(combination)]
+            win_probs = cprobs[0, list(sorted(set(columns).difference(combination)))]
+            p_event = numpy.prod(draw_probs) * numpy.prod(win_probs)
+            p_win_draw += p_win * p_event
+            combo = ["W"]*C
+            for c in combination:
+                combo[c] = "D"
+            df_smart_combos.loc["".join(combo)] = (p_event, p_win)
+    return p_win_draw
+
+
+def _rolling_probs_calculation(
+    samples: numpy.ndarray, sample_cols: numpy.ndarray,
+    s_totals: numpy.ndarray,
+) -> pandas.DataFrame:
+    """ Calculates win, loose and win-by-draw probabilities for all samples.
+
+    Parameters
+    ----------
+    samples : numpy.ndarray
+        (S,) values of candidate samples
+    sample_cols : numpy.ndarray
+        (S,) corresponding candidate indices
+    s_totals : numpy.ndarray
+        (C,) numbers of samples per candidate
+
+    Returns
+    -------
+    df_probs : pandas.DataFrame
+        table with win/loose/win-by-draw probabilities for each sample
+    """
+    C = s_totals.shape[0]
+
+    p_win = numpy.zeros_like(samples, dtype=float)
+    p_loose = numpy.zeros_like(samples, dtype=float)
+    p_win_draw = numpy.zeros_like(samples, dtype=float)
+    
+    # s_smaller: number of samples IN EACH COLUMN that are smaller than [value]
+    s_smaller = numpy.zeros((C,))
+
+    # now iterate over groups with identical sample values in the DataFrame
+    # pandas.DataFrame.groupby is too slow for this -> DIY iterator using the unique idx & counts
+    unique_values, idx_from, counts = numpy.unique(samples, return_counts=True, return_index=True)
+    for value, ifrom, nsame in fastprogress.progress_bar(tuple(zip(unique_values, idx_from, counts))):
+        ito = ifrom + nsame
+        candidates_with_value = sample_cols[ifrom:ito]
+
+        # s_same: number of samples IN EACH COLUMN that have the same [value]
+        s_same = numpy.zeros(C)
+        same_cols, same_counts = numpy.unique(candidates_with_value, return_counts=True)
+        s_same[same_cols] = same_counts
+        # s_larger: number of samples IN EACH COLUMN that are larger than [value]
+        s_larger = s_totals - s_smaller - s_same
+
+        #numpy.testing.assert_array_equal(s_smaller + s_same + s_larger, s_totals)
+
+        # from the counts of smaller/same/larger values in other columns,
+        # calculate the probabilities of direct win, direct loss and draw
+        cprobs_all = numpy.array([
+            s_smaller,
+            s_larger,
+            s_same
+        ]) / s_totals
+        for s, fc in zip(range(ifrom, ito), candidates_with_value):
+            # do not look at probabilities w.r.t. the same column:
+            cprobs = numpy.hstack([cprobs_all[:, :fc], cprobs_all[:, fc+1:]])
+        
+            p_win[s] = numpy.prod(cprobs[0, :])
+            p_loose[s] = 1 - numpy.prod(1 - cprobs[1, :])
+            
+            if s_same[fc] != nsame:
+                # draws with other columns are possible -> calculate combinatorial event & win probabilities
+                p_win_draw[s] = _win_draw_prob(cprobs)
+            #else:
+            #    # no other candidate has a sample of this value
+            #    p_win_draw[s] was initialized to 0
+
+        # increment the column-wise count of samples that are smaller than [value]
+        # by this iterative counting, we avoid doing S*C </=/> comparisons, dramatically reducing complexity
+        s_smaller += s_same
+
+    # summarize results as DataFrame
+    df_probs = pandas.DataFrame(data={
+        "value": samples,
+        "candidate": sample_cols,
+        "p_win": p_win,
+        "p_loose": p_loose,
+        "p_win_draw": p_win_draw,
+    })
+    return df_probs
+
+
+def sampling_probabilities(
+    candidate_samples: typing.Sequence[typing.Sequence[Sample]],
+) -> numpy.ndarray:
+    """ Calculates the probability thompson sampling probability of each candidate.
+
+    Parameters
+    ----------
+    candidate_samples : array-like
+        posterior samples for each candidate (C,)
+        (sample count may be different)
+
+    Returns
+    -------
+    probabilities : numpy.ndarray
+        (C,) array of sampling probabilities
+    """
+    C = len(candidate_samples)
+    s_totals = numpy.array(tuple(map(len, candidate_samples)))
+    samples, sample_cols = _sort_samples(candidate_samples)
+
+    df_probs = _rolling_probs_calculation(samples, sample_cols, s_totals)
+        
+    probabilities = numpy.zeros(C)
+    for fc, df in df_probs.groupby("candidate"):
+        probabilities[fc] = sum(df.p_win + df.p_win_draw) / len(df)
+    return probabilities
diff --git a/requirements.txt b/requirements.txt
index 2273212..7afb8b4 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,6 +1,7 @@
 h5py
 numba
 numpy
+pandas
 pytest
 scipy
 typing

From 0f9ed973e52f598bcde87f8db16244df918155a8 Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Sun, 4 Oct 2020 00:29:02 +0200
Subject: [PATCH 4/7] support for both correlated and uncorrelated TS

---
 pyrff/test_thompson.py | 76 +++++++++++++++++++++++++++++++++++++-----
 pyrff/thompson.py      | 56 ++++++++++++++++++++++---------
 2 files changed, 109 insertions(+), 23 deletions(-)

diff --git a/pyrff/test_thompson.py b/pyrff/test_thompson.py
index a048bb1..b4c4915 100644
--- a/pyrff/test_thompson.py
+++ b/pyrff/test_thompson.py
@@ -1,6 +1,7 @@
 import numpy
 import pytest
 
+from . import exceptions
 from . import thompson
 
 
@@ -20,6 +21,7 @@ def test_sample_batch(self, batch_size, seed):
 
         batch = thompson.sample_batch(
             candidate_samples=samples, ids=ids,
+            correlated=False,
             batch_size=batch_size, seed=seed
         )
         assert len(batch) == batch_size
@@ -34,11 +36,56 @@ def test_no_bias_on_sample_collisions(self):
             [2, 2],
             [2, 2, 2],
         ]
-        batch = thompson.sample_batch(samples, ids=('A', 'B', 'C'), batch_size=100, seed=1234)
+        batch = thompson.sample_batch(samples, ids=('A', 'B', 'C'), correlated=False, batch_size=100, seed=1234)
         assert batch.count('A') != 100
         assert batch.count('C') != 0
         pass
 
+    def test_correlated_sampling(self):
+        samples = [
+            [1, 2, 3],
+            [1, 1, 1],
+            [0, 1, 2],
+        ]
+        batch = thompson.sample_batch(samples, ids=('A', 'B', 'C'), correlated=True, batch_size=100, seed=1234)
+        assert batch.count('A') < 100
+        assert batch.count('B') < 100 / 3
+        assert batch.count('C') == 0
+        pass
+
+
+class TestExceptions:
+    def test_id_count(self):
+        with pytest.raises(exceptions.ShapeError, match="candidate ids"):
+            thompson.sample_batch([
+                    [1,2,3],
+                    [1,2],
+                ],
+                ids=("A", "B", "C"),
+                correlated=False,
+                batch_size=30,
+            )
+
+    def test_correlated_sample_size_check(self):
+        with pytest.raises(exceptions.ShapeError, match="number of samples"):
+            thompson.sample_batch([
+                    [1,2,3],
+                    [1,2],
+                ],
+                ids=("A", "B"),
+                correlated=True,
+                batch_size=30,
+            )
+
+        with pytest.raises(exceptions.ShapeError):
+            thompson.sampling_probabilities([
+                    [1,2,3],
+                    [1,2],
+                ],
+                correlated=True,
+            )
+        pass
+
 
 class TestThompsonProbabilities:
     def test_sort_samples(self):
@@ -71,26 +118,39 @@ def test_win_draw_prob(self):
         ])), 0.041666666)
         pass
 
-    def test_sampling_probability(self):
+    def test_sampling_probability_uncorrelated(self):
         numpy.testing.assert_array_equal(thompson.sampling_probabilities([
             [0, 1, 2],
             [0, 1, 2],
-        ]), [0.5, 0.5])
-        
+        ], correlated=False), [0.5, 0.5])
+
         numpy.testing.assert_array_equal(thompson.sampling_probabilities([
             [0, 1, 2],
             [10],
-        ]), [0, 1])
+        ], correlated=False), [0, 1])
 
         numpy.testing.assert_array_equal(thompson.sampling_probabilities([
             [0, 1, 2],
             [3, 4, 5],
             [5, 4, 3],
-        ]), [0, 0.5, 0.5])
+        ], correlated=False), [0, 0.5, 0.5])
 
         numpy.testing.assert_array_equal(thompson.sampling_probabilities([
             [5, 6],
             [0, 0, 10, 20],
             [5, 6],
-        ]), [0.25, 0.5, 0.25])
-        pass
\ No newline at end of file
+        ], correlated=False), [0.25, 0.5, 0.25])
+        pass
+
+    def test_sampling_probability_correlated(self):
+        numpy.testing.assert_array_equal(thompson.sampling_probabilities([
+            [0, 1, 2],
+            [0, 1, 2],
+        ], correlated=True), [0.5, 0.5])
+
+        numpy.testing.assert_array_equal(thompson.sampling_probabilities([
+            [0, 4, 2],
+            [3, 4, 5],
+            [5, 1, 6],
+        ], correlated=True), [0.5/3, 0.5/3, 2/3])
+        pass
diff --git a/pyrff/thompson.py b/pyrff/thompson.py
index eee2e11..5284dd5 100644
--- a/pyrff/thompson.py
+++ b/pyrff/thompson.py
@@ -7,6 +7,8 @@
 import numpy
 import pandas
 
+from . import exceptions
+
 # custom type shortcuts
 Sample = typing.Union[int, float]
 
@@ -14,7 +16,9 @@
 def sample_batch(
     candidate_samples: typing.Sequence[typing.Sequence[Sample]], *,
     ids:typing.Sequence,
-    batch_size:int, seed:typing.Optional[int]=None
+    correlated:bool,
+    batch_size:int,
+    seed:typing.Optional[int] = None,
 ) -> tuple:
     """
     Draws a batch of candidates by Thompson Sampling from posterior samples.
@@ -26,8 +30,13 @@ def sample_batch(
         (sample count may be different)
     ids : numpy.ndarray, list, tuple
         (C,) candidate identifiers
+    correlated : bool
+        Switches between jointly (correlated=True) or independently (correlated=False) sampling the candidates.
+        When correlated=True, all candidates must have the same number of samples.
     batch_size : int
         number of candidates to draw (B)
+    seed : int
+        seed for the random number generator (will reset afterwards)
 
     Returns
     -------
@@ -36,8 +45,10 @@ def sample_batch(
     """
     n_candidates = len(candidate_samples)
     n_samples = tuple(map(len, candidate_samples))
+    if correlated and len(set(n_samples)) != 1:
+        raise exceptions.ShapeError("For correlated sampling, all candidates must have the same number of samples.")
     if len(ids) != n_candidates:
-        raise ValueError(f"Number of candidate ids ({len(ids)}) does not match number of candidate_samples ({n_candidates}).")
+        raise exceptions.ShapeError(f"Number of candidate ids ({len(ids)}) does not match number of candidate_samples ({n_candidates}).")
     ids = numpy.atleast_1d(ids)
     # work with matrix even if sample count is varies to get more efficient slicing
     samples = numpy.zeros((max(n_samples), n_candidates))
@@ -50,7 +61,10 @@ def sample_batch(
         # for every sample in the batch, randomize the column order
         # to prevent always selecting lower-numbered candidates when >=2 samples are equal
         col_order = random.permutation(n_candidates)
-        idx = random.randint(n_samples, size=n_candidates)
+        if correlated:
+            idx = numpy.repeat(numpy.random.randint(n_samples[0]), n_candidates)
+        else:
+            idx = random.randint(n_samples, size=n_candidates)
         selected_samples = samples[:, col_order][idx, numpy.arange(n_candidates)]
         best_candidate = ids[col_order][numpy.argmax(selected_samples)]
         chosen_candidates.append(best_candidate)
@@ -147,7 +161,7 @@ def _rolling_probs_calculation(
     p_win = numpy.zeros_like(samples, dtype=float)
     p_loose = numpy.zeros_like(samples, dtype=float)
     p_win_draw = numpy.zeros_like(samples, dtype=float)
-    
+
     # s_smaller: number of samples IN EACH COLUMN that are smaller than [value]
     s_smaller = numpy.zeros((C,))
 
@@ -165,8 +179,6 @@ def _rolling_probs_calculation(
         # s_larger: number of samples IN EACH COLUMN that are larger than [value]
         s_larger = s_totals - s_smaller - s_same
 
-        #numpy.testing.assert_array_equal(s_smaller + s_same + s_larger, s_totals)
-
         # from the counts of smaller/same/larger values in other columns,
         # calculate the probabilities of direct win, direct loss and draw
         cprobs_all = numpy.array([
@@ -177,14 +189,14 @@ def _rolling_probs_calculation(
         for s, fc in zip(range(ifrom, ito), candidates_with_value):
             # do not look at probabilities w.r.t. the same column:
             cprobs = numpy.hstack([cprobs_all[:, :fc], cprobs_all[:, fc+1:]])
-        
+
             p_win[s] = numpy.prod(cprobs[0, :])
             p_loose[s] = 1 - numpy.prod(1 - cprobs[1, :])
-            
+
             if s_same[fc] != nsame:
                 # draws with other columns are possible -> calculate combinatorial event & win probabilities
                 p_win_draw[s] = _win_draw_prob(cprobs)
-            #else:
+            # else:
             #    # no other candidate has a sample of this value
             #    p_win_draw[s] was initialized to 0
 
@@ -205,14 +217,18 @@ def _rolling_probs_calculation(
 
 def sampling_probabilities(
     candidate_samples: typing.Sequence[typing.Sequence[Sample]],
+    correlated:bool,
 ) -> numpy.ndarray:
-    """ Calculates the probability thompson sampling probability of each candidate.
+    """ Calculates the thompson sampling probability of each candidate.
 
     Parameters
     ----------
     candidate_samples : array-like
         posterior samples for each candidate (C,)
         (sample count may be different)
+    correlated : bool
+        Switches between jointly (correlated=True) or independently (correlated=False) sampling the candidates.
+        When correlated=True, all candidates must have the same number of samples.
 
     Returns
     -------
@@ -221,11 +237,21 @@ def sampling_probabilities(
     """
     C = len(candidate_samples)
     s_totals = numpy.array(tuple(map(len, candidate_samples)))
-    samples, sample_cols = _sort_samples(candidate_samples)
+    if correlated and len(set(s_totals)) != 1:
+        raise exceptions.ShapeError("For correlated sampling, all candidates must have the same number of samples.")
 
-    df_probs = _rolling_probs_calculation(samples, sample_cols, s_totals)
-        
     probabilities = numpy.zeros(C)
-    for fc, df in df_probs.groupby("candidate"):
-        probabilities[fc] = sum(df.p_win + df.p_win_draw) / len(df)
+    if correlated:
+        S = s_totals[0]
+        for s, samples in enumerate(numpy.array(candidate_samples).T):
+            vwin = numpy.max(samples)
+            i_winners = numpy.argwhere(samples == vwin)
+            n_winners = len(i_winners)
+            probabilities[i_winners] += 1 / n_winners
+        probabilities /= S
+    else:
+        samples, sample_cols = _sort_samples(candidate_samples)
+        df_probs = _rolling_probs_calculation(samples, sample_cols, s_totals)
+        for fc, df in df_probs.groupby("candidate"):
+            probabilities[fc] = sum(df.p_win + df.p_win_draw) / len(df)
     return probabilities

From 7bfce1dedd1f321848098d86c00e1de9aadc242a Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Sun, 4 Oct 2020 01:07:57 +0200
Subject: [PATCH 5/7] remove pandas dependency, comments

---
 pyrff/thompson.py | 50 +++++++++++++++++++++++++----------------------
 requirements.txt  |  1 -
 2 files changed, 27 insertions(+), 24 deletions(-)

diff --git a/pyrff/thompson.py b/pyrff/thompson.py
index 5284dd5..ae7e504 100644
--- a/pyrff/thompson.py
+++ b/pyrff/thompson.py
@@ -5,7 +5,6 @@
 import typing
 import fastprogress
 import numpy
-import pandas
 
 from . import exceptions
 
@@ -87,7 +86,7 @@ def _sort_samples(
     -------
     samples : array
         sorted array of all samples
-    sample_cols : array
+    sample_candidates : array
         the group numbers
     """
     flat_samples = numpy.vstack([
@@ -117,8 +116,6 @@ def _win_draw_prob(cprobs: numpy.ndarray) -> float:
     """
     C = cprobs.shape[1]
     columns = numpy.arange(C)
-    # TODO: this can be further accelerated by not using a DataFrame
-    df_smart_combos = pandas.DataFrame(columns=["p_event", "p_win"])
 
     drawable = tuple(columns[cprobs[2, :] > 0])
     p_win_draw = 0
@@ -132,21 +129,20 @@ def _win_draw_prob(cprobs: numpy.ndarray) -> float:
             combo = ["W"]*C
             for c in combination:
                 combo[c] = "D"
-            df_smart_combos.loc["".join(combo)] = (p_event, p_win)
     return p_win_draw
 
 
 def _rolling_probs_calculation(
-    samples: numpy.ndarray, sample_cols: numpy.ndarray,
+    samples: numpy.ndarray, sample_candidates: numpy.ndarray,
     s_totals: numpy.ndarray,
-) -> pandas.DataFrame:
+) -> typing.Tuple[numpy.ndarray, numpy.ndarray, numpy.ndarray]:
     """ Calculates win, loose and win-by-draw probabilities for all samples.
 
     Parameters
     ----------
     samples : numpy.ndarray
         (S,) values of candidate samples
-    sample_cols : numpy.ndarray
+    sample_candidates : numpy.ndarray
         (S,) corresponding candidate indices
     s_totals : numpy.ndarray
         (C,) numbers of samples per candidate
@@ -170,7 +166,7 @@ def _rolling_probs_calculation(
     unique_values, idx_from, counts = numpy.unique(samples, return_counts=True, return_index=True)
     for value, ifrom, nsame in fastprogress.progress_bar(tuple(zip(unique_values, idx_from, counts))):
         ito = ifrom + nsame
-        candidates_with_value = sample_cols[ifrom:ito]
+        candidates_with_value = sample_candidates[ifrom:ito]
 
         # s_same: number of samples IN EACH COLUMN that have the same [value]
         s_same = numpy.zeros(C)
@@ -204,15 +200,7 @@ def _rolling_probs_calculation(
         # by this iterative counting, we avoid doing S*C </=/> comparisons, dramatically reducing complexity
         s_smaller += s_same
 
-    # summarize results as DataFrame
-    df_probs = pandas.DataFrame(data={
-        "value": samples,
-        "candidate": sample_cols,
-        "p_win": p_win,
-        "p_loose": p_loose,
-        "p_win_draw": p_win_draw,
-    })
-    return df_probs
+    return p_win, p_loose, p_win_draw
 
 
 def sampling_probabilities(
@@ -221,6 +209,9 @@ def sampling_probabilities(
 ) -> numpy.ndarray:
     """ Calculates the thompson sampling probability of each candidate.
 
+    ATTENTION: When correlated=False is specified, the occurence of non-unique sample values can
+    increase the runtime complexity to worst-case O(2^total_samples).
+
     Parameters
     ----------
     candidate_samples : array-like
@@ -240,18 +231,31 @@ def sampling_probabilities(
     if correlated and len(set(s_totals)) != 1:
         raise exceptions.ShapeError("For correlated sampling, all candidates must have the same number of samples.")
 
-    probabilities = numpy.zeros(C)
+    probabilities = numpy.zeros(C, dtype=float)
     if correlated:
+        # this case is O(S^2 * C) because it does not need to account for combinations
         S = s_totals[0]
         for s, samples in enumerate(numpy.array(candidate_samples).T):
             vwin = numpy.max(samples)
+            # which candidates have the highest value?
             i_winners = numpy.argwhere(samples == vwin)
             n_winners = len(i_winners)
+            # attribute winning probability to the winners
             probabilities[i_winners] += 1 / n_winners
         probabilities /= S
     else:
-        samples, sample_cols = _sort_samples(candidate_samples)
-        df_probs = _rolling_probs_calculation(samples, sample_cols, s_totals)
-        for fc, df in df_probs.groupby("candidate"):
-            probabilities[fc] = sum(df.p_win + df.p_win_draw) / len(df)
+        # For uncorrelated TS, all possible combinations must be considered.
+        # Naively doing all combinations would be O(S^C), but this implementation simplifies it:
+        # 1. it's sufficient to categorize win/loose/draw
+        # 2. sorting the samples allows for an iteration that needs much fewer </=/> comparisons
+        # 3. combinatorics for draw win probabilities is only required for non-unique sample values
+
+        # first sort all samples into a vector
+        samples, sample_candidates = _sort_samples(candidate_samples)
+        # then calculate the win/loose/win-by-draw probabilities for each sample
+        p_win, p_loose, p_win_draw = _rolling_probs_calculation(samples, sample_candidates, s_totals)
+        # finally summarize the sample-wise probabilities by the corresponding candidate
+        for c in range(C):
+            mask = sample_candidates == c
+            probabilities[c] = numpy.sum(p_win[mask] + p_win_draw[mask]) / numpy.sum(mask)
     return probabilities
diff --git a/requirements.txt b/requirements.txt
index 7afb8b4..2273212 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,7 +1,6 @@
 h5py
 numba
 numpy
-pandas
 pytest
 scipy
 typing

From ce39cad9480c1e330b450c5563864da07a0d0799 Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Sun, 4 Oct 2020 01:21:38 +0200
Subject: [PATCH 6/7] include fastprogress for TS probability calculation
 progress bar

---
 requirements.txt | 1 +
 1 file changed, 1 insertion(+)

diff --git a/requirements.txt b/requirements.txt
index 2273212..fa0fe5b 100644
--- a/requirements.txt
+++ b/requirements.txt
@@ -1,3 +1,4 @@
+fastprogress
 h5py
 numba
 numpy

From e322ceb34ab5f3f6068811bf553e90bb3b098173 Mon Sep 17 00:00:00 2001
From: Michael Osthege <michael.osthege@outlook.com>
Date: Sun, 4 Oct 2020 02:03:54 +0200
Subject: [PATCH 7/7] pin only major version, trying to fix the CI

---
 .github/workflows/pipeline.yml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/.github/workflows/pipeline.yml b/.github/workflows/pipeline.yml
index 8333d49..7ce507a 100644
--- a/.github/workflows/pipeline.yml
+++ b/.github/workflows/pipeline.yml
@@ -33,7 +33,7 @@ jobs:
         export NUMBA_DISABLE_JIT=1
         pytest --cov=./pyrff --cov-report xml --cov-report term-missing pyrff/
     - name: Upload coverage
-      uses: codecov/codecov-action@v1.0.7
+      uses: codecov/codecov-action@v1
       if: matrix.python-version == 3.8
       with:
         file: ./coverage.xml