Triangular (square) matrix class for Python, using only half as much memory. Supports decent portions of what you'd expect for a numpy object

triangle.py

#!/usr/bin/env python -O
#
# Copyright (c) 2013 Kyle Gorman
#
# Permission is hereby granted, free of charge, to any person obtaining a
# copy of this software and associated documentation files (the
# "Software"), to deal in the Software without restriction, including
# without limitation the rights to use, copy, modify, merge, publish,
# distribute, sublicense, and/or sell copies of the Software, and to
# permit persons to whom the Software is furnished to do so, subject to
# the following conditions:
#
# The above copyright notice and this permission notice shall be included
# in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS
# OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT.
# IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY
# CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN ACTION OF CONTRACT,
# TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN CONNECTION WITH THE
# SOFTWARE OR THE USE OR OTHER DEALINGS IN THE SOFTWARE.
#
# triangle.py: class for storing numbers in a matrix "triangle"
# Kyle Gorman <[email protected]>
#
# Inspiration comes from:
#
# http://en.wikipedia.org/wiki/Triangular_matrix
# http://www.codeguru.com/cpp/cpp/algorithms/general/article.php/c11211/TIP-Half-Size-Triangular-Matrix.htm
#
# TODO make this act more like a numpy array


from __future__ import division

from numpy import asarray, asmatrix, ones, triu_indices_from, sqrt, zeros


class Triangle(object):
    """
    A sparse representation of the upper triangular portion of a square
    matrix (or "array" in the numpy sense), including the diagonal

    The following examples show two of the constructors and `as_array`
    and `as_matrix`, instance methods which convert the triangle to "dense"
    representations.

    >>> Triangle(xrange(21)).to_array()
    array([[  0.,   1.,   2.,   3.,   4.,   5.],
           [  0.,   6.,   7.,   8.,   9.,  10.],
           [  0.,   0.,  11.,  12.,  13.,  14.],
           [  0.,   0.,   0.,  15.,  16.,  17.],
           [  0.,   0.,   0.,   0.,  18.,  19.],
           [  0.,   0.,   0.,   0.,   0.,  20.]])
    >>> Triangle.ones(6).to_matrix()
    matrix([[ 1.,  1.,  1.,  1.,  1.,  1.],
            [ 0.,  1.,  1.,  1.,  1.,  1.],
            [ 0.,  0.,  1.,  1.,  1.,  1.],
            [ 0.,  0.,  0.,  1.,  1.,  1.],
            [ 0.,  0.,  0.,  0.,  1.,  1.],
            [ 0.,  0.,  0.,  0.,  0.,  1.]])
    """

    @staticmethod
    def is_square(n):
        """
        Determine whether a positive integer n is square using the 
        Babylonian algorithm. This is probably just good enough. See:

        http://stackoverflow.com/questions/2489435/

        >>> all(Triangle.is_square(i * i) for i in xrange(10000))
        True
        """
        # do the smallest ones by lookup
        if n <= 100 and n in {1, 4, 9, 16, 25, 36, 49, 64, 81, 100}:
            return True
        # normal case
        x = n // 2
        seen = {x}
        while x * x != n:
            x = (x + (n // x)) // 2
            if x in seen:
                return False
            seen.add(x)
        return True

    @staticmethod
    def edge_to_size(n):
        """
        Returns the size of the triangular portion of an n x n matrix

        >>> n = 6
        >>> Triangle.size_to_edge(Triangle.edge_to_size(n)) == n
        True
        """
        return (n * n + n) // 2

    @staticmethod
    def size_to_edge(n):
        """
        Returns the dimension of a square matrix whose diagonal is size n,
        the inverse function of `edge_to_size`

        >>> n = 21
        >>> Triangle.edge_to_size(Triangle.size_to_edge(n)) == n
        True
        """
        # Weird fact: an integer is "triangular" (fits into the "triangle"
        # of a square matrix) iff 8x + 1 is a square number. This also
        # holds when considering n x n triangular matrices whose diagonal
        # we are ignoring, (i.e., in the subclass TriangleNoDiagonal)
        # since that is equivalent to the triangle of a perfectly good
        # (n - 1) x (n - 1) matrix
        x = 8 * n + 1
        if not Triangle.is_square(x):
            raise ValueError('n ({}) is non-triangular'.format(n))
        return (int(sqrt(x)) - 1) // 2

    # primary constructor

    def __init__(self, iterable, dtype=float):
        L = len(iterable)
        self.n = self.size_to_edge(len(iterable))
        self.triangle = asarray(iterable, dtype)

    # alternative MATLAB-style constructors

    @classmethod
    def zeros(cls, n, dtype=float):
        """
        Construct the triangular component of an n x n matrix with all
        cells initialized to zero
        """
        return cls(zeros(cls.edge_to_size(n), dtype))

    @classmethod
    def ones(cls, n, dtype=float):
        """
        Construct the triangular component of an n x n matrix with 
        triangular cells initialized to one
        """
        return cls(ones(Triangle.edge_to_size(n), dtype))

    # important magic methods (with associated static methods)

    def __repr__(self):
        return '{0}({1} x {1})'.format(self.__class__.__name__, self.n)

    def _get_1d_index(self, indices):
        (row, col) = indices
        if col > row:
            raise ValueError('Point ({}, {}) is in lower triangle'.format(
                row, col))
        return row * self.n - (row - 1) * ((row - 1) + 1) / 2 + col - row

    def __getitem__(self, indices):
        return self.triangle[self._get_1d_index(indices)]

    def __setitem__(self, indices, value):
        self.triangle[self._get_1d_index(indices)] = value

    @staticmethod
    def tri_indices_from(the_array):
        return triu_indices_from(the_array)

    def to_array(self):
        """
        Return (non-sparse) array representation. Non-triangle cells are
        populated with zeros
        """
        # initialize the (full) array to be returned
        the_array = zeros((self.n, self.n))
        # write into new array; indices are selected by converting the x
        # and y vectors returned by the indexing function into a single
        # vector of (x, y) tuples
        the_array[self.tri_indices_from(the_array)] = self.triangle
        # and we're done
        return the_array

    def to_matrix(self):
        """
        The same as to_array, but returns a matrix instead
        """
        return asmatrix(self.to_array())


class TriangleNoDiagonal(Triangle):
    """
    A sparse representation of the upper triangular portion of a square
    matrix (or "array" in the numpy sense), excluding the diagonal

    The following examples show two of the constructors and `as_array`
    and `as_matrix`, instance methods which convert the triangle to "dense"
    representations.

    >>> x = TriangleNoDiagonal(xrange(15))
    >>> y = x.to_array()
    >>> y
    array([[  0.,   0.,   1.,   2.,   3.,   4.],
           [  0.,   0.,   5.,   6.,   7.,   8.],
           [  0.,   0.,   0.,   9.,  10.,  11.],
           [  0.,   0.,   0.,   0.,  12.,  13.],
           [  0.,   0.,   0.,   0.,   0.,  14.],
           [  0.,   0.,   0.,   0.,   0.,   0.]])
    >>> x[2, 3] == y[2, 3] == 9.
    True
    >>> TriangleNoDiagonal.ones(5).to_matrix()
    matrix([[ 0.,  1.,  1.,  1.,  1.,  1.],
            [ 0.,  0.,  1.,  1.,  1.,  1.],
            [ 0.,  0.,  0.,  1.,  1.,  1.],
            [ 0.,  0.,  0.,  0.,  1.,  1.],
            [ 0.,  0.,  0.,  0.,  0.,  1.],
            [ 0.,  0.,  0.,  0.,  0.,  0.]])
    """

    # all overloadings of the superclass

    @staticmethod
    def edge_to_size(n):
        """
        Returns the size of the triangular portion of an n x n matrix,
        excluding the diagonal
        """
        return (n * n - n) // 2

    @staticmethod
    def size_to_edge(n):
        """
        Returns the dimension of a square matrix whose diagonal is size n,
        the inverse function of `edge_to_size`
        """
        x = 8 * n + 1
        if not Triangle.is_square(x):
            raise ValueError('n ({}) is non-triangular'.format(n))
        return ((int(sqrt(x)) - 1) // 2) + 1

    @staticmethod
    def tri_indices_from(the_array):
        """
        The second argument to triu_indices_from (1) indicates an "offset"
        of 1; in this context, this means the diagonal indices are ignored
        """
        return triu_indices_from(the_array, 1)

    def _get_1d_index(self, indices):
        (row, col) = indices
        if col <= row:
            raise ValueError('Point ({}, {}) is in lower triangle'.format(
                row, col))
        return row * (self.n - 1) - (row - 1) * ((row - 1) + 1) / 2 + \
            col - row - 1


if __name__ == '__main__':
    import doctest
    doctest.testmod()

At line 153:

if col > row:
    raise ValueError(...

should the comparison be this instead:

if col < row:
    raise ValueError(...

because without this fix, a contradictory error is thrown:

ValueError: Point (0, 69) is in lower triangle

and also the following is already used with non-diagonals at line 249:

if col <= row:
    raise ValueError(...

Thanks man!