How to sample without replacement and reweigh each time (conditional sampling)?

Viewed 315

Consider a dataset of N rows with weights. This is the basic algorithm:

  1. Normalize the weights so that they sum to 1.
  2. Backup the weights into another column to record sample probabilities
  3. Randomly choose 1 row (without replacement), given the sample probabilities, and add it to the sample dataset
  4. Remove the drawn weight from the original dataset, and recompute the sample probabilities by normalizing the weights of the remaining rows
  5. Repeat steps 3 and 4 till sum of weights in sample reaches or exceeds threshold (assume 0.6)

Here is a toy example:

import pandas as pd
import numpy as np

def sampler(n):
    df = pd.DataFrame(np.random.rand(n), columns=['weight'])
    df['weight'] = df['weight']/df['weight'].sum()
    df['samp_prob'] = df['weight']

    samps = pd.DataFrame(columns=['weight'])

    while True:
        choice = np.random.choice(df.index, 1, replace=False, p=df['samp_prob'])[0]
        samps.loc[choice, 'weight'] = df.loc[choice, 'weight']
        df.drop(choice, axis=0, inplace=True)
        df['samp_prob'] = df['weight']/df['weight'].sum()
        if samps['weight'].sum() >= 0.6:
            break
    return samps

The problem with the toy example is the exponential growth in run times with increasing size of n:

Exponential Growth

2 Answers

Starting off approach

Few observations :

  • The dropping of rows per iteration that results in creation of new dataframes isn't helping with the performance.
  • Doesn't look like easy to vectorize, BUT should be easy to work with the underlying array data for performance. The idea would be to use masks and avoid re-creating dataframes or arrays. Starting off, we would be using two columns array, corresponding to the columns named : 'weights' and 'samp_prob'.

So, with those in mind, the starting approach would be something like this -

def sampler2(n):
    a = np.random.rand(n,2)
    a[:,0] /= a[:,0].sum()
    a[:,1] = a[:,0]
    N = len(a)

    idx = np.arange(N)
    mask = np.ones(N,dtype=bool)
    while True:
        choice = np.random.choice(idx[mask], 1, replace=False, p=a[mask,1])[0]
        mask[choice] = 0
        a_masked = a[mask,0]
        a[mask,1] = a_masked/a_masked.sum()
        if a[~mask,0].sum() >= 0.6:
            break
    out = a[~mask,0]
    return out

Improvement #1

A later observation revealed that the first column of the array isn't changing across iterations. So, we could optimize for the masked summations for the first column, by pre-computing the total summation and then at each iteration, a[~mask,0].sum() would be simply the total summation minus a_masked.sum(). Thsi leads us to the first improvement, listed below -

def sampler3(n):
    a = np.random.rand(n,2)
    a[:,0] /= a[:,0].sum()
    a[:,1] = a[:,0]
    N = len(a)

    idx = np.arange(N)
    mask = np.ones(N,dtype=bool)
    a0_sum = a[:,0].sum()
    while True:
        choice = np.random.choice(idx[mask], 1, replace=False, p=a[mask,1])[0]
        mask[choice] = 0
        a_masked = a[mask,0]
        a_masked_sum = a_masked.sum()
        a[mask,1] = a_masked/a_masked_sum
        if a0_sum - a_masked_sum >= 0.6:
            break
    out = a[~mask,0]
    return out

Improvement #2

Now, slicing and masking into the columns of a 2D array could be improved by using two separate arrays instead, given that the first column wasn't changing between iterations. That gives us a modified version, like so -

def sampler4(n):
    a = np.random.rand(n)
    a /= a.sum()
    b = a.copy()
    N = len(a)

    idx = np.arange(N)
    mask = np.ones(N,dtype=bool)
    a_sum = a.sum()
    while True:
        choice = np.random.choice(idx[mask], 1, replace=False, p=b[mask])[0]
        mask[choice] = 0
        a_masked = a[mask]
        a_masked_sum = a_masked.sum()
        b[mask] = a_masked/a_masked_sum
        if a_sum - a_masked_sum >= 0.6:
            break
    out = a[~mask]
    return out

Runtime test -

In [250]: n = 1000

In [251]: %timeit sampler(n) # original app
     ...: %timeit sampler2(n)
     ...: %timeit sampler3(n)
     ...: %timeit sampler4(n)
1 loop, best of 3: 655 ms per loop
10 loops, best of 3: 50 ms per loop
10 loops, best of 3: 44.9 ms per loop
10 loops, best of 3: 38.4 ms per loop

In [252]: n = 2000

In [253]: %timeit sampler(n) # original app
     ...: %timeit sampler2(n)
     ...: %timeit sampler3(n)
     ...: %timeit sampler4(n)
1 loop, best of 3: 1.32 s per loop
10 loops, best of 3: 134 ms per loop
10 loops, best of 3: 119 ms per loop
10 loops, best of 3: 100 ms per loop

Thus, we are getting 17x+ and 13x+ speedups with the final version over the original method for n=1000 and n=2000 sizes!

I think you can rewrite this while loop to do it in a single pass:

while True:
    choice = np.random.choice(df.index, 1, replace=False, p=df['samp_prob'])[0]
    samps.loc[choice, 'weight'] = df.loc[choice, 'weight']
    df.drop(choice, axis=0, inplace=True)
    df['samp_prob'] = df['weight']/df['weight'].sum()
    if samps['weight'].sum() >= 0.6:
        break

to something more like:

n = len(df.index)
ind = np.random.choice(n, n, replace=False, p=df["samp_prob"])
res = df.iloc[ind]
i = (res.cumsum() >= 0.6).idxmax()  # first index that satisfies .sum() >= 0.6
samps = res.iloc[:i+1]

The key parts are that choice can take multiple elements (indeed the entire array) whilst still respecting the probabilities. The cumsum allows you to cut off after passing the 0.6 threshold.


In this example you can see that the array is randomly chosen, but that 4 is most likely chosen nearer the top.

In [11]: np.random.choice(5, 5, replace=False, p=[0.05, 0.05, 0.1, 0.2, 0.6])
Out[11]: array([0, 4, 3, 2, 1])

In [12]: np.random.choice(5, 5, replace=False, p=[0.05, 0.05, 0.1, 0.2, 0.6])
Out[12]: array([3, 4, 1, 2, 0])

In [13]: np.random.choice(5, 5, replace=False, p=[0.05, 0.05, 0.1, 0.2, 0.6])
Out[13]: array([0, 4, 3, 1, 2])

In [14]: np.random.choice(5, 5, replace=False, p=[0.05, 0.05, 0.1, 0.2, 0.6])
Out[14]: array([4, 3, 0, 2, 1])

In [15]: np.random.choice(5, 5, replace=False, p=[0.05, 0.05, 0.1, 0.2, 0.6])
Out[15]: array([4, 2, 3, 0, 1])

In [16]: np.random.choice(5, 5, replace=False, p=[0.05, 0.05, 0.1, 0.2, 0.6])
Out[16]: array([3, 4, 2, 0, 1])

Note: The replace=False, ensures the probabilities are "reweighed" in the sense that it can't be picked again.

Related