"""Extensible numpy arrays."""
__all__ = ['earray']
import sys
import numpy
[docs]class earray(numpy.ndarray):
"""Extensible array.
Warning: Arrays may only be resized if there are no references to
them to avoid memory corruption. One potential case where this
occurs is at the command prompt where displaying an Earray sets
the reference special variable "_". Using print will avoid this
problem.
>>> a = earray([1,2,3],maxrefcount=8+2) # Extra 2 for doctesting
>>> a[3] = 4
>>> a # Don't do this!
earray([1, 2, 3, 4])
>>> a[4] = 5 # doctest: +SKIP
ValueError: cannot resize an array that has been referenced or is referencing
another array in this way. Use the resize function
>>> a = earray([1,2,3],maxrefcount=8+2)
>>> a[2,3] = 1
>>> print a
[[1 0 0 0]
[2 0 0 0]
[3 0 0 1]]
>>> print a[3,3]
0
>>> print a
[[1 0 0 0]
[2 0 0 0]
[3 0 0 1]
[0 0 0 0]]
"""
_extension_factor = 1
def __new__(cls,*argv,**kwargs):
"""Make array a an Earray.
The argument maxrefcount sets the maximum number of references
allowed before resizing is prohibited. Resize should not be
allowed when there are references to the object as this could
cause segmentation faults. Unfortunately, there is no better
way than reference counting known to allow for this. For
example, self references this in methods. Thus, maxrefcount
should be 8. When doctesting, however, an additional
reference is kept that must be counted. Nose tests seem to
make this work.
"""
if 'maxrefcount' in kwargs:
maxrefcount = kwargs.pop('maxrefcount')
else:
maxrefcount = 8
self = numpy.array(*argv,**kwargs).view(cls).copy()
self._maxrefcount = maxrefcount
return self
def __setitem__(self,index,value):
try:
numpy.ndarray.__setitem__(self.view(numpy.ndarray),index,value)
except IndexError,index_err:
self._resize_to_index(index)
numpy.ndarray.__setitem__(self.view(numpy.ndarray),index,value)
def __getitem__(self,index):
try:
return numpy.ndarray.__getitem__(self.view(numpy.ndarray),index)
except IndexError,index_err:
self._resize_to_index(index)
return numpy.ndarray.__getitem__(self.view(numpy.ndarray),index)
def _resize_to_index(self,index):
"""Resize the array to ensure that the specified index
exists."""
# Try to extend array
shape = list(self.shape)
if isinstance(index,int):
if 1 < len(shape):
raise IndexError(
"Insufficient number of indices for a resize.")
else:
new_shape = max(int(self._extension_factor*shape[0]),
index + 1)
elif isinstance(index,tuple):
if len(index) == len(shape):
# Index has same number of dimensions as array
new_shape = [int(max(x+1,self._extension_factor*y))
for (x,y)
in zip(index,shape)]
elif len(index) > len(shape):
# Extend the dimension of the array
shape0 = shape + [0]*(len(index)-len(shape))
new_shape = [int(max(x+1,self._extension_factor*y))
for (x,y)
in zip(index,shape0)]
else:
# insufficient number of indices, re-raise original
# exception
raise IndexError(
"Insufficient number of indices for a resize.")
#elif isinstance(index,slice):
# raise ValueError(
# "Unable to extend array with index of type %s"%
# type(index).__name__)
else:
raise ValueError(
"Unable to extend array with index of type %s"%
type(index).__name__)
# Resize
if sys.getrefcount(self) <= self._maxrefcount:
# Force a resize if there are only six references or less.
self.resize(new_shape,refcheck=False)
else:
self.resize(new_shape,refcheck=True)
if not isinstance(index,int): # No need to copy 1d data
n_data = numpy.prod(shape)
n_new = numpy.prod(new_shape)
self1d = self.reshape(n_new)
old_data = self1d[:n_data].copy().reshape(shape)
self1d[:n_data] = 0 # Zero out old data
# Copy data
block = [slice(x) for x in shape] +\
[0]*(len(new_shape)-len(shape))
numpy.ndarray.__setitem__(self,block,old_data)
del self1d # Remove reference