multidim reduction

#!/usr/bin/python2


import itertools

import numpy


def reduce(dataset):
    """Stand-in for the real function"""
    return dataset.sum()


def causal_graph(array):
    reduction = 2 # remaining dimensions
    ndim = array.ndim
    assert(ndim > reduction)
    allaxes = xrange(ndim)
    combinations = itertools.combinations
    axessets = combinations(allaxes, ndim - reduction)
    destinations = reversed(list(combinations(allaxes, reduction)))
    arraysum = array.sum
    datasets = (arraysum(axes) for axes in axessets)
    result = numpy.zeros((ndim, ) * reduction)
    # fill upper triangular with values
    for destination, dataset in itertools.izip(destinations, datasets):
        result[destination] = reduce(dataset)
    # copy upper triangular to lower triangular. Diagonal stays 0
    result[numpy.tril_indices(ndim, -1)] = result[numpy.triu_indices(ndim, 1)]
    return result