[MATRIX-SIG] FancyArray

Timothy A. Hochberg hochberg@wwa.com
Fri, 29 Aug 1997 12:40:34 -0500 (CDT)


If anyone's interested in playing with it I've put together a FancyArray
class derived from UserArray that allows indexing with sequences. 

>>> a = FancyArray(zeros((10,10)))
>>> a[(1,4), (5,6)] = 5
>>> a[(2,9), 1:5] = 9
>>> a
FancyArray([[0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 5, 5, 0, 0, 0],
       [0, 9, 9, 9, 9, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 5, 5, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 0, 0, 0, 0, 0, 0, 0, 0, 0],
       [0, 9, 9, 9, 9, 0, 0, 0, 0, 0]])

Addition, etc, all work. However, there are deinitely still some bugs,
some of which are my fault, and some of which have to do with UserArray
and its interaction with the builtin functions. In particular, reshape
doesn't work right for FancyArray or UserArray. So, don't look at this as
a finished product, consider it a proof of concept type thang'. If enough
people are interested, this could be polished into a useful type for those
who yearn for S-like indexing (not that I've ever used S). 

It does this by keeping a data array (shared between different slices of
the same matrix) and a mask array. It should be not too slow since all the
loops are still in C. The one exception is that I needed to fake a 'give'
(opposite of take) function using map, and I don't know how fast that is,
but is could be easily be rewritten in C.

Anyway, enjoy. If this looks interesting let me know.
 ____   
  /im  

+------------------------------------------------+
|Tim Hochberg            Research Assistant      |
|hochberg <at> wwa.com   University of Illinois  |
|                        Electrical Engineering  |
+------------------------------------------------+


"""FancyArray is a class that provides fancy indexing."""
import sys, UserArray, operator
from Numeric import *
from types import *
from indexing import *


def give(fro, to, indices, axis=0):
    """Give 'to' the elements of 'fro' at 'indices'."""
    map(operator.setitem, [to]*len(fro), indices, fro)


class FancyArray(UserArray.UserArray):

    """An array-like object that supports "fancy" indexing.

    The index is either a tuple of subindinces or a single
    subindex. The subindexes can have one of four forms:
    
      1. An integer.
      2. A slice object (or i:j:k).
      3. An ellipses object (or ...).
      4. A NewAxis (None) object.
      5. A sequence of integers.

    The first four types of subindices are available with normal
    arrays, it is this last type which is new. If a single, sequence,
    subindex is used, the sequence must not be a tuple as that would be
    interpreted as a tuple of integer subindices.

    Examples:
      * A[(1,2,3), (0,2,4)] =>  A[1:4:2,0:5:2].
      * A[[1,4,7,0]         =>  The {A[1], A[4], A[7], A[0]}

    """

    def __init__(self, array, typecode=None, mask=None):
	"""Create a FancyArray."""
	try:
	    data = asarray(array.data, typecode)
	    mask = asarray(array.mask)
	except:
	    try:
		data = reshape(asarray(array, typecode), (-1,))
		if mask == None:
		    mask = reshape(arange(len(data)), shape(array))
	    except:
		raise ValueError, "Bad initializer."
	self.data = data
	self.mask = mask
	self.shape = shape(mask)
	self.typecode = self.data.typecode()
	self.name = string.split(str(self.__class__))[1]

    def __getattr__(self, key):
	"""Fake the presence of an array attribute."""
	if key == "array":
	    return take(self.data, self.mask)
	else:
	    raise AttributeError, key

    def __float__(self):
	"""Return the floating point representation of self."""
	# This could be put in UserArray instead - it would be better?!
	# __int__ (?), etc. should probably also be added to UserArray.
	return self.__class__(self, Float)
    
    def _return(self, mask):
	"""Return a new FancyArray from a mask."""
	if len(shape(mask)) == 0: self.data[mask]
	else: return self.__class__(self.data, mask=mask)

    def _rc(self, array):
	"""Return a new FancyArray from an array."""
	if len(shape(array)) == 0: return array
	else: return self.__class__(array)

    def _make_mask(self, indices):
	if type(indices) != TupleType: indices = (indices,)
	mask = self.mask
	skip = 0
	ellip = 0
	theslice = slice(None, None, None)
	for (axis, index) in indexed(indices):
	    itype = type(index)
	    if itype == SliceType:
		mask = mask[[theslice]*skip+[index]]
		skip = skip + 1
	    elif itype == IntType:
		mask = mask[[theslice]*skip+[index]]
	    elif itype is NoneType:
		mask = mask[[theslice]*skip+[index]]
		skip = skip + 2
	    elif itype == EllipsisType:
		skip = len(shape(mask)) - (len(indices) - axis) + 1
	    else:			# Try a sequence.
		try:
		    mask = take(mask, index, skip)
		    skip = skip + 1
		except:
		    raise IndexError, "Bad index of type %s" % itype
	return mask
		
    def __getitem__(self, indices): 
	"""Get an item from an array using generalized indexing."""
	return self._return(self._make_mask(indices))

    def __getslice__(self, i, j): 
	"""Get a slice from array."""
	return self._return(self.mask[i:j])

    def __setitem__(self, indices, value): 
	"""Set an item in array using generalized indexing."""
	mask = self._make_mask(indices) 
	if shape(value) != shape(mask):
	    raise IndexError, "Matrices not alligned for copy"
	give(reshape(value, (-1,)), self.data, reshape(mask, (-1,)))

    def __setslice__(self, i, j, value): 
	"""Set a slice in the array."""
	give(reshape(value, (-1,)), self.data, reshape(mask[i:j], (-1)))



def print_result(c,d):
    try:
	if alltrue(equal(c,d).flat) or \
	   (type(c) != InstanceType and c.__class__ != FancyArray):
	    print "OK"
	else: 
	    print "FAILED!"
    except:
	print
	print "ERROR IN TEST SUITE"
	print "Was processing:"
	print shape(c)
	print shape(d)
	


def test():
    """Test suite."""
    a0 = ones((5,5,5))
    a = FancyArray(a0)
    b0 = arange(5.)
    b = FancyArray(b0)
    for op in (add, subtract, multiply, divide):
 	print "Checking", `op`,
 	c = op(a0,b0)
 	d = op(a, b)
 	print_result(c,d)
    for expr, equiv in (("b[(range(5))]", "b0[0:5]"),
 			("b[(range(0,5,2))]", "b0[0:5:2]"),
 			("a[:,(range(0,5,2))]", "a0[:,0:5:2]"),
 			("a[:,:,(range(0,5,2))]", "a0[:,:,0:5:2]"),
 			("a[...,(range(0,5,2))]", "a0[...,0:5:2]"),
 			("a[NewAxis,...][...,(range(0,5,2))]", 
 			 "a0[NewAxis,...][...,0:5:2]"),
 			):
 	print "Checking", expr,
 	c = eval(expr)
 	d = eval(equiv)
 	print_result(c,d)
    for expr, equiv in \
	(("c[range(5)] = arange(5)[:,NewAxis,NewAxis] + zeros(shape(c[range(5)]))", 
	  "d[0:5] = arange(5)[:,NewAxis,NewAxis] + zeros(shape(d[0:5]))"),
	 ("c[(range(0,5,2))]=arange(0,5,2)[:,NewAxis,NewAxis]+zeros(shape(c[(range(0,5,2))])) ", 
	  "d[0:5:2] = arange(0,5,2)[:,NewAxis,NewAxis]+zeros(shape(c[(range(0,5,2))]))"),
	 ("c[:,(range(0,5,2))] = arange(0,5,2)[NewAxis,:,NewAxis]+zeros(shape(c[:,(range(0,5,2))]))", 
	  "d[:,0:5:2] = arange(0,5,2)[NewAxis,:,NewAxis]+zeros(shape(c[:,(range(0,5,2))]))"),
	 ("c[:,:,(range(0,5,2))] = arange(0,5,2)[NewAxis,NewAxis,:]+zeros(shape(c[:,:,(range(0,5,2))]))", 
	  "d[:,:,0:5:2] = arange(0,5,2)[NewAxis,NewAxis,:]+zeros(shape(c[:,:,(range(0,5,2))]))"),
	 ("c[...,(range(0,5,2))] = arange(0,5,2)[NewAxis,NewAxis,:]+zeros(shape(c[:,:,(range(0,5,2))]))", 
	  "d[...,0:5:2] = arange(0,5,2)[NewAxis,NewAxis,:]+zeros(shape(c[:,:,(range(0,5,2))]))"),
	 ("c[NewAxis,...][...,(range(0,5,2))] = arange(0,5,2)[NewAxis,NewAxis,NewAxis,:] + zeros(shape(c[NewAxis,...][...,(range(0,5,2))]))", 
	 "d[NewAxis,...][...,0:5:2] = arange(0,5,2)[NewAxis,NewAxis,NewAxis,:]+ zeros(shape(c[NewAxis,...][...,(range(0,5,2))]))"),
	  ):
	c = FancyArray(zeros((10,10,10), Float))
	d = zeros((10,10,10), Float)
	print "Checking", expr[:30], "...",
	exec (expr)
	exec (equiv)
	print_result(c,d)

    


if __name__ == "__main__": test()
~


_______________
MATRIX-SIG  - SIG on Matrix Math for Python

send messages to: matrix-sig@python.org
administrivia to: matrix-sig-request@python.org
_______________