Optmized multidim reduction

#!/usr/bin/python2

#  Copyright 2015 Florian Philipp
#
#  Licensed under the Apache License, Version 2.0 (the "License");
#  you may not use this file except in compliance with the License.
#  You may obtain a copy of the License at
#
#      http://www.apache.org/licenses/LICENSE-2.0
#
#  Unless required by applicable law or agreed to in writing, software
#  distributed under the License is distributed on an "AS IS" BASIS,
#  WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#  See the License for the specific language governing permissions and
#  limitations under the License.


"""Demonstrates the combinatoric reduction of multidimensional datasets"""


import itertools

import numpy


def per_axis_combination(arr, reduce_comb=numpy.prod, eliminate_axis=numpy.sum,
                         odtype=None, remaindim=2):
    """Combines multidimensional datasets into lower dimensional datasets

    All combinations of dimensions in the input array are combined using
    repeated calls of eliminate_axis until remaindim dimensions are left.
    Then reduce_comb is called on the final result, once per combination.

    Example:
    arr.ndim = 5
    remaindim = 2
    result[0,1] = reduce_comb(combination of axis 2, 3, 4)
    result[0,2] = reduce_comb(combination of axis 1, 3, 4)
    ...
    result[3,4] = reduce_comb(combination of axis 0, 1, 2)

    Arguments:
    arr -- a multidimensional numpy array that serves as input
    reduce_comb -- functor that is called in the final reduction stage.
                   The input has remaindim dimensions. The output is expected
                   to be a scalar compatible with odtype
    eliminate_axis -- functor that is called in the first reduction stage.
                      eliminate_axis(N dimensions, axis) -> N-1 dimensions
    odtype -- output data type. Defaults to arr.dtype
    remaindim -- Final dimensionality

    Return value:
    An array with dtype=odtype, ndim=remaindim, shape=(arr.ndim,)*remaindim.
    Only the strictly upper triangle (or its higher-dimensional equivalent) is
    filled. The remaining entries are invalid. The indices correspond to the
    combined dimensions
    """
    ndim = arr.ndim
    if ndim < remaindim:
        raise IndexError('insufficient number of axes')
    if odtype is None:
        odtype = arr.dtype
    allaxes = xrange(ndim)
    destinations = reversed(list(itertools.combinations(allaxes, remaindim)))
    results = numpy.empty((ndim, ) * remaindim, odtype)
    finaldepth = ndim - remaindim
    def sum_recurse(outersum, outeraxis=-1):
        """Recursively applies eliminate_axis, then finally reduce_comb

        Relevant surrounding variables:
        destinations -- iterator that defines the indizes to which the results
                        are saved. The reduction happens in lexicographical
                        order of the axis numbers, e.g.
                        (0, 1, 2), (0, 1, 3) ... (2, 3, 4), so
                        the index order of destinations has to correspond, i.e.
                        (3, 4), (2, 4) ... (0, 1)
        results -- will be filled with return values of reduce_comb

        Arguments:
        outersum -- input array
        outeraxis -- last axis that was eliminated. -1 on initial call
        """
        depth = ndim - outersum.ndim
        if depth == finaldepth:
            results[next(destinations)] = reduce_comb(outersum)
            return
        startaxis = outeraxis + 1
        endaxis = depth + remaindim + 1
        for inneraxis in xrange(startaxis, endaxis):
            # Note that the actual axis numbers differ from the computed
            # numbers because some axes have been eliminated before
            innersum = eliminate_axis(outersum, inneraxis - depth)
            sum_recurse(innersum, inneraxis)
    sum_recurse(arr)
    return results


def triu_to_full(arr, diagonal=1):
    """Copies the transposition of the upper triangular matrix to the lower

    Arguments:
    arr -- array representing a symmetric matrix as a full matrix where only
           the upper triangular matrix is filled with valid values
    diagonal -- number of diagonal above which the values are valid.
                0 means upper triangular, 1 means strictly upper triangular

    Return value:
    A copy of the input array where all values are valid. With diagonal == 1,
    the main diagonal is filled with zeros
    """
    zeroed = numpy.triu(arr, diagonal)
    return zeroed + zeroed.T