n dimensional grid in Python / numpy

Question

I have an unknown number n of variables that can range from 0 to 1 with some known step s, with the condition that they sum up to 1. I want to create a

Answer 1

Assuming that they always add up to 1, as you said:

import itertools

def make_grid(n):   
  # setup all possible values in one position
  p = [(float(1)/n)*i for i in range(n+1)]

  # combine values, filter by sum()==1
  return [x for x in itertools.product(p, repeat=n) if sum(x) == 1]

print(make_grid(n=3))

#[(0.0, 0.0, 1.0),
# (0.0, 0.3333333333333333, 0.6666666666666666),
# (0.0, 0.6666666666666666, 0.3333333333333333),
# (0.0, 1.0, 0.0),
# (0.3333333333333333, 0.0, 0.6666666666666666),
# (0.3333333333333333, 0.3333333333333333, 0.3333333333333333),
# (0.3333333333333333, 0.6666666666666666, 0.0),
# (0.6666666666666666, 0.0, 0.3333333333333333),
# (0.6666666666666666, 0.3333333333333333, 0.0),
# (1.0, 0.0, 0.0)]

Answer 2

EDIT

Here is a better solution. It basically partitions the number of steps into the amount of variables to generate all the valid combinations:

def partitions(n, k):
    if n < 0:
        return -partitions(-n, k)
    if k <= 0:
        raise ValueError('Number of partitions must be positive')
    if k == 1:
        return np.array([[n]])
    ranges = np.array([np.arange(i + 1) for i in range(n + 1)])
    parts = ranges[-1].reshape((-1, 1))
    s = ranges[-1]
    for _ in range(1, k - 1):
        d = n - s
        new_col = np.concatenate(ranges[d])
        parts = np.repeat(parts, d + 1, axis=0)
        s = np.repeat(s, d + 1) + new_col
        parts = np.append(parts, new_col.reshape((-1, 1)), axis=1)
    return np.append(parts, (n - s).reshape((-1, 1)), axis=1)

def make_grid_part(n, step):
    num_steps = round(1.0 / step)
    return partitions(num_steps, n) / float(num_steps)

print(make_grid_part(3, 0.33333))

Output:

array([[ 0.        ,  0.        ,  1.        ],
       [ 0.        ,  0.33333333,  0.66666667],
       [ 0.        ,  0.66666667,  0.33333333],
       [ 0.        ,  1.        ,  0.        ],
       [ 0.33333333,  0.        ,  0.66666667],
       [ 0.33333333,  0.33333333,  0.33333333],
       [ 0.33333333,  0.66666667,  0.        ],
       [ 0.66666667,  0.        ,  0.33333333],
       [ 0.66666667,  0.33333333,  0.        ],
       [ 1.        ,  0.        ,  0.        ]])

For comparison:

%timeit make_grid_part(5, .1)
>>> 338 µs ± 2.25 µs per loop (mean ± std. dev. of 7 runs, 1000 loops each)
%timeit make_grid_simple(5, .1)
>>> 26.4 ms ± 806 µs per loop (mean ± std. dev. of 7 runs, 10 loops each)

make_grid_simple actually runs out of memory if you push it just a bit further.

---

Here is one simple way:

def make_grid_simple(n, step):
    num_steps = round(1.0 / step)
    vs = np.meshgrid(*([np.linspace(0, 1, num_steps + 1)] * n))
    all_combs = np.stack([v.flatten() for v in vs], axis=1)
    return all_combs[np.isclose(all_combs.sum(axis=1), 1)]

print(make_grid_simple(3, 0.33333))

Output:

[[ 0.          0.          1.        ]
 [ 0.33333333  0.          0.66666667]
 [ 0.66666667  0.          0.33333333]
 [ 1.          0.          0.        ]
 [ 0.          0.33333333  0.66666667]
 [ 0.33333333  0.33333333  0.33333333]
 [ 0.66666667  0.33333333  0.        ]
 [ 0.          0.66666667  0.33333333]
 [ 0.33333333  0.66666667  0.        ]
 [ 0.          1.          0.        ]]

However, this is not the most efficient way to do it, since it is simply making all the possible combinations and then just picking the ones that add up to 1, instead of generating only the right ones in the first place. For small step sizes, it may incur in too high memory cost.

Answer 3

We can think of this as a problem of dividing some fixed number of things (1/s in this case and represented using sum_left parameter) between some given number of bins (n in this case). The most efficient way I can think of doing this is using a recursion:

In [31]: arr = []
In [32]: def fun(n, sum_left, arr_till_now):
    ...:     if n==1:
    ...:         n_arr = list(arr_till_now)
    ...:         n_arr.append(sum_left)
    ...:         arr.append(n_arr)
    ...:     else:
    ...:         for i in range(sum_left+1):
    ...:             n_arr = list(arr_till_now)
    ...:             n_arr.append(i)
    ...:             fun(n-1, sum_left-i, n_arr)

This would give an output like:

In [36]: fun(n, n, [])
In [37]: arr
Out[37]: 
[[0, 0, 3],
 [0, 1, 2],
 [0, 2, 1],
 [0, 3, 0],
 [1, 0, 2],
 [1, 1, 1],
 [1, 2, 0],
 [2, 0, 1],
 [2, 1, 0],
 [3, 0, 0]]

And now I can convert it to a numpy array to do an elementwise multiplication:

In [39]: s = 0.33
In [40]: arr_np = np.array(arr)
In [41]: arr_np * s
Out[41]: 
array([[ 0.        ,  0.        ,  0.99999999],
       [ 0.        ,  0.33333333,  0.66666666],
       [ 0.        ,  0.66666666,  0.33333333],
       [ 0.        ,  0.99999999,  0.        ],
       [ 0.33333333,  0.        ,  0.66666666],
       [ 0.33333333,  0.33333333,  0.33333333],
       [ 0.33333333,  0.66666666,  0.        ],
       [ 0.66666666,  0.        ,  0.33333333],
       [ 0.66666666,  0.33333333,  0.        ],
       [ 0.99999999,  0.        ,  0.        ]])

Answer 4

This method will also work for an arbitrary sum (total):

import numpy as np
import itertools as it
import scipy.special

n = 3
s = 1/3.
total = 1.00

interval = int(total/s)
n_combs = scipy.special.comb(n+interval-1, interval, exact=True)
counts = np.zeros((n_combs, n), dtype=int)

def count_elements(elements, n):
    count = np.zeros(n, dtype=int)
    for elem in elements:
        count[elem] += 1
    return count

for i, comb in enumerate(it.combinations_with_replacement(range(n), interval)):
    counts[i] = count_elements(comb, n)

ratios = counts*s
print(ratios)

Answer 5

Here is a direct method using itertools.combinations:

>>> import itertools as it
>>> import numpy as np
>>> 
>>> # k is 1/s
>>> n, k = 3, 3
>>> 
>>> combs = np.array((*it.combinations(range(n+k-1), n-1),), int)
>>> (np.diff(np.c_[np.full((len(combs),), -1), combs, np.full((len(combs),), n+k-1)]) - 1) / k
array([[0.        , 0.        , 1.        ],
       [0.        , 0.33333333, 0.66666667],
       [0.        , 0.66666667, 0.33333333],
       [0.        , 1.        , 0.        ],
       [0.33333333, 0.        , 0.66666667],
       [0.33333333, 0.33333333, 0.33333333],
       [0.33333333, 0.66666667, 0.        ],
       [0.66666667, 0.        , 0.33333333],
       [0.66666667, 0.33333333, 0.        ],
       [1.        , 0.        , 0.        ]])

If speed is a concern, itertools.combinations can be replaced by a numpy implementation.