Python generator to produce lattice points on an n-simplex. Looking for a code review, and hopefully to learn something if someone has a nicer solution.  Here's what I wrote:



from __future__ import division, print_function
from future_builtins import *
import numpy as np

def _walk(num_dims, samples_per_dim, max_):
    if num_dims == 0:
        yield np.array([(max_ - 1) / (samples_per_dim - 1)])
    else:
        for i in range(max_):
            for rest in _walk(num_dims - 1, samples_per_dim, max_ - i):
                yield np.concatenate((np.array([i]) / (samples_per_dim - 1),
                                      rest))

def walk(num_dims, samples_per_dim):
    """
    A generator that returns lattice points on an n-simplex.
    """
    return _walk(num_dims, samples_per_dim, samples_per_dim)


So, list(walk(2, 3)) yields:

[array([ 0.,  0.,  1.]),
 array([ 0. ,  0.5,  0.5]),
 array([ 0.,  1.,  0.]),
 array([ 0.5,  0. ,  0.5]),
 array([ 0.5,  0.5,  0. ]),
 array([ 1.,  0.,  0.])]
Concatenating numpy arrays isn't a really good idea because that's not how they were designed to be used.

A better way might be this:

def walk(num_dims, samples_per_dim):
    """
    A generator that returns lattice points on an n-simplex.
    """
    values = np.arange(samples_per_dim) / (samples_per_dim - 1)
    for items in itertools.product(values, repeat = num_dims+1):
        if sum(items) == 1.0:
            yield items


Basically, we iterate over all possible combinations of those points, and filter out the invalid ones. This may seem wasteful, but since itertools is written in C, its probably actually faster then your solution. It produces tuples rather then numpy arrays.