How to generate samples uniformly at random from multiple discrete variables subject to constraints?

Question

I would like to generate a Monte Carlo process to fill an urn with N balls of I colors, C[i]. Each color C[i] has a minimum and maximum number of balls which should be placed in the urn.

For example, I am trying to put 100 balls in the urn, and can fill it with four colors:

• Red - Minimum 0, Maximum 100 # NB, the actual maximum cannot be realized.
• Blue - Minimum 50, Maximum 100
• Yellow - Minimum 0, Maximum 50
• Green - Minimum 25, Maximum 75

How can I generate a N samples which is ensured to be uniformly distributed across possible outcomes?

I have seen solutions to this problem where the balls have no minimum or maximum, or alternatively, has the same implicit minima and maxima. See for example, this discussion on a slightly different subject:

Generate uniformly distributed weights that sum to unity?

But I'm having problems generalizing this solution.

Answer 1

Let $n_i$ denote the number of balls of color $C_i$. Also, let $m_i$ and $M_i$ denote the minimum and the maximum number of balls of color $C_i$, respectively. We want to sample $(n_1, \dots, n_I)$ uniformly at random subject to the following constraints:

1. $m_i \leq n_i \leq M_i$
2. $\sum_{i=1}^I n_i = N$

First of all, you can remove the lower bound constraints (i.e. $m_i \leq n_i$) by picking $m_i$ balls of color $C_i$ initially. This modifies the two constraints as follows:

1. $0 \leq n_i \leq b_i = M_i - m_i$
2. $\sum_{i=1}^I n_i = B = N - \sum_{i=1}^I m_i$

Let $P(n_1, \dots, n_I \mid B, b_{1:I})$ denote the uniform distribution that we are interested in. We can use chain rule and dynamic programming to sample from $P$ efficiently. First, we use chain rule to write $P$ as follows: $$ \begin{align} P(n_1, \dots, n_I \mid B, b_{1:I}) &= P(n_I \mid B, b_{1:I}) P(n_1, \dots, n_{I-1} \mid n_I, B, b_{1:I}) \\ &= P(n_I \mid B, b_{1:I}) P(n_1, \dots, n_{I-1} \mid B-n_I, b_{1:I-1}) \quad (1) \end{align} $$ where $P(n_I | B, b_{1:I}) = \sum_{n_1, \dots, n_{I-1}} P(n_1, \dots, n_I | B, b_{1:I})$ is the marginal distribution of $n_I$. Note that $P(n_I | B, b_{1:I})$ is a discrete distribution and can be computed efficiently using dynamic programming. Also, note that the second term in (1) can be computed recursively. We sample $n_I$ in the first round of the recursion, update the total number of balls to $B - n_I$ and recurse to sample $n_{I-1}$ in the next round.

The following is a Matlab implementation of the idea. The complexity of the algorithm is $O(I \times B \times K)$ where $K = \max_{i=1}^I b_i$. The code uses randomly generated $m_i$s in each run. As a result, some of the generated test cases may not have any valid solution, in which case the code prints out a warning message.

global dpm b

I = 5; % number of colors
N = 300; % total number of balls

m = randi(50, 1, I)-1; % minimum number of balls from each from each color
M = 99*ones(1, I); % maximum number of balls from each color

% print original constraints
print_constraints(I, N, m, M, 'original constraints');

% remove the lower bound constraints
b = M - m;
B = N - sum(m);
m = zeros(size(m));

% print transformed constraints
print_constraints(I, B, zeros(1, I), b, 'transformed constraints');

% initialize the dynamic programming matrix (dpm)
% if dpm(i, n) <> -1, it denotes the value of the following marginal probability
% \sum_{k=1}^{i-1} P(n_1, ..., n_i |
dpm = -ones(I, B);

% sample the number of balls of each color, one at a time, using chain rule
running_B = B;  % we change the value of "running_B" on the fly, as we sample balls of different colors
for i = I : -1 : 1
    % compute marginal distribution P(n_i)

    % instead of P(n_i) we compute q(n_i) which is un-normalized.
    q_ni = zeros(1, b(i) + 1); % possibilities for ni are 0, 1, ..., b(i)
    for ni = 0 : b(i)
        q_ni(ni+1) = dpfunc(i-1, running_B-ni);
    end
    if(sum(q_ni) == 0)
        fprintf('Impossible!!! constraints can not be satisfied!\n');
        return;
    end
    P_ni = q_ni / sum(q_ni);
    ni = discretesample(P_ni, 1) - 1;
    fprintf('n_%d=%d\n', i, ni);
    running_B = running_B - ni;
end

where the function print_constraints is

function [] = print_constraints(I, N, m, M, note)
    fprintf('\n------ %s ------ \n', note);
    fprintf('%d <= n_%d <= %d\n', [m; [1:I]; M]);
    fprintf('========================\n');
    fprintf('sum_{i=1}^%d n_i = %d\n', I, N);
end

and the function dpfunc performs the dynamic programming computation as follows:

function res = dpfunc(i, n)
    global dpm b

    % check boundary cases
    if(n == 0)
        res = 1;
        return;
    end
    if(i == 0) % gets here only if n <> 0
        res = 0;
        return;
    end

    if(n < 0)
        res = 0;
        return;
    end

    if(dpm(i, n) == -1) % if <> -1, it has been compute before, so, just use it!
        % compute the value of dpm(i, n) = \sum_{n_1, ..., n_i} valid(n, n_1, ..., n_i)
        % where "valid" return 1 if \sum_{j=1}^i n_i = n and 0 <= n_i <= b_i, for all i
        % and 0 otherwise.
        dpm(i, n) = 0;
        for ni = 0 : b(i)
            dpm(i, n) = dpm(i, n) + dpfunc(i-1, n-ni);
        end
    end
    res = dpm(i, n);
end

and finally, the function discretesample(p, 1) draws a random sample from the discrete distribution $p$. You can find one implementation of this function here.

Answer 2

Let's consider a generalization of this problem. There are $m=4$ cans of paint of $m=4$ distinct colors and $n^{(0)}=100$ balls. Can $i$ can hold up to $a^{(0)}_i = (100, 100, 50, 75)$ balls. You wish to generate configurations of balls in the cans having at least $b_i = (0, 50, 0, 25)$ balls in can $i$ for each $i$, each configuration with equal probability.

Such configurations are in one-to-one correspondence with the configurations obtained after removing $b_i$ balls from can $i$, limiting the $n = n^{(0)} - \sum_i b_i = 100 - (0+50+0+25)=25$ remaining balls to at most $a_i = a^{(0)}_i - b_i = (100, 50, 50, 50)$ per can. I will therefore just generate these and let you adjust them afterwards (by putting $b_i$ balls back into can $i$ for every $i$).

To count these configurations up, fix all but two of the indices, say $i$ and $j$. Suppose there are $s_k$ balls already in can $k$ for each $k$ differing from $i$ and $j$. That leaves $s_i+s_j$ balls. Conditional on where the other $n - (s_i+s_j)$ balls are, these are distributed uniformly within cans $i$ and $j$. The possible configurations are $1 + \min(a_i + a_j - s_i - s_j, s_i+s_j)$ in number (see the comments), ranging from placing as many balls in can $i$ as possible all the way through placing as many balls in can $j$ as possible.

If you wish, you could count the total number of configurations by applying this argument recursively to the remaining $m-2$ cans. However, to obtain samples we don't even need to know this count. All we need to do is repeatedly visit all possible unordered pairs $\{i,j\}$ of cans and randomly (and uniformly) change the distribution of balls within those two cans. This is a Markov chain with a limiting probability distribution that is uniform over all possible states (as is readily shown using standard methods). Therefore it suffices to start in any state, run the chain long enough to reach the limiting distribution, and then keep track of the states visited by this procedure. As usual, to avoid serial correlation, this sequence of states should be "thinned" by skipping through them (or revisited randomly). Thinning by a factor of about half the number of cans tends to work well, because after that many steps on average each can has been affected, producing a genuinely new configuration.

This algorithm costs $O(m)$ effort to generate each random configuration on average. Although other $O(m)$ algorithms exist, this one has the advantage of not needing to do the combinatorial calculations beforehand.

---

As an example, let's work out a smaller situation manually. Let $a=(4,3,2,1)$ and $n=3$, for instance. There are 15 valid configurations, which may be written as strings of occupancy numbers. For example, 0201 places two balls into the second can and one ball in the fourth can. Emulating the argument, let's consider the total occupancy of the first two cans. When that is $s_1+s_2=3$ balls, no balls are left for the last two cans. That gives the states

30**, 21**, 12**, 03**

where ** represents all the possible occupancy numbers for the last two cans: namely, 00. When $s_1+s_2=2$, the states are

20**, 11**, 02**

where now ** can be either 10 or 01. That gives $3\times 2=6$ more states. When $s_1+s_2=1$, the states are

10**, 01**

where now ** can be 20, 11, but not 02 (due to the limit of one ball in the last can). That gives $2\times 2=4$ more states. Finally, when $s_1+s_2=0$, all balls are in the last two cans, which must be full to their limits of $2$ and $1$. The $4+6+4+1=15$ equally probable states therefore are

3000, 2100, 1200, 0300; 2010, 2001, 1110, 1101, 0210, 0201; 1020, 1011, 0120, 0111; 0021.

Using the code below, a sequence of $10,009$ such configurations was generated and thinned to every third one, creating $3337$ configurations of the $15$ states. Their frequencies were the following:

State: 3000 2100 1200 0300 2010 1110 0210 1020 0120 2001 1101 0201 1011 0111 0021 
Count:  202  227  232  218  216  208  238  227  237  209  239  222  243  211  208 

A $\chi^2$ test of uniformity gives a $\chi^2$ value of $11.2$, $p=0.67$ ($14$ degrees of freedom): that is beautiful agreement with the hypothesis that this procedure produces equally probable states.

---

This R code is set up to handle the situation in the question. Change a and n to work with other situations. Set N to be large enough to generate the number of realizations you need after thinning.

This code cheats a little bit by cycling systematically through all $(i,j)$ pairs. If you want to be strict about running the Markov chain, generate i, j, and ij randomly, as given in the commented code.

#
# Gibbs-like sampler.
#
# `a` is an array of maximum numbers of balls of each type.  Its values should
#     all be integers greater than zero.
# `n` is the total number of balls.
#------------------------------------------------------------------------------#
g <- function(j, state, a) {
  #
  # `state` contains the occupancy numbers.
  # `a`     is the array of maximum occupancy numbers.
  # `j`     is a pair of indexes into `a` to "rotate".
  #
  k <- sum(state[j]) # Total occupancy.
  x <- floor(runif(1, max(0, k - a[j[2]]), min(k, a[j[1]]) + 1))
  state[j] <- c(x, k-x)
  return(state)
}
#
# Set up the problem.
#
a <- c(100, 50, 50, 50)
n <- 25

# a <- 4:1
# n <- 3
#
# Initialize the state.
#
state <- round(n * a / sum(a))
i <- 1
while (sum(state) < n) {
  if (state[i] < a[i]) state[i] <- state[i] + 1
  i <- i+1
}
while (sum(state) > n) {
  i <- i-1
  if (state[i] > 0) state[i] <- state[i] - 1
}
#
# Conduct a sequence of random changes.
#
set.seed(17)
N <- 1e5
sim <- matrix(state, ncol=1)
u <- ceiling(N / choose(length(state), 2) / 2)
i <- rep(rep(1:length(state), each=length(state)-1), u)
j <- rep(rep(length(state):1, length(state)-1), u)
ij <- rbind(i, j)
#
# Alternatively, generate `ij` randomly:
#   i <- sample.int(length(state), N, replace=TRUE)
#   j <- sample.int(length(state)-1, N, replace=TRUE)
#   ij <- rbind(i, ((i+j-1) %% length(state))+1)
#
sim <- cbind(sim, apply(ij, 2, function(j) {state <<- g(j, state, a); state}))
rownames(sim) <- paste("Can", 1:nrow(sim))
#
# Thin them for use.  Each column is a state.
#
thin <- function(x, stride=1, start=1) {
  i <- round(seq(start, ncol(x), by=stride))
  x[, i]
}
#
# Make a scatterplot of the results, to illustrate.
#
par(mfrow=c(1,1))
s <- thin(sim, stride=max(1, N/1e4))
pairs(t(s) + runif(length(s), -1/2, 1/2), cex=1/2, col="#00000005", pch=16)