I am attempting to subclass a numpy array in order to create arrays that have periodic boundary conditions. I only need this to work for 1D arrays.
What I Want:
Suppose I have an object a=Periodic_Grid([0,1,2,3,4,5,6,7,8,9])
. I would like the following outputs
>>> a=Periodic_Grid(a)
>>> a
Periodic_Grid([0, 1, 2, 3, 4, 5, 6, 7, 8, 9])
>>> a[0]
0
>>> a[7]
7
>>> a[14]
4
>>> a[-5]
5
>>> a[-14]
6
>>> a[1:5]
array([1., 2., 3., 4.])
>>> a[8:13]
array([8., 9., 0., 1., 2.])
>>> a[-5:2]
array([5., 6., 7., 8., 9., 0., 1.])
In other words, indexing just wraps around. Slicing should also appropriately wrap around. Otherwise, it acts as normal.
Minimal Working Code
class Periodic_Grid(np.ndarray):
def __new__(cls, input_array):
obj = np.asarray(input_array).view(cls)
obj.grid_size = input_array.size
return obj
def __getitem__(self, index):
if isinstance(index, slice):
indicies=list(range(index.start, index.stop))
size=len(indicies)
a=np.empty(size)
print("Size: {0} indicies: {1}".format(size, indicies))
for i, ind in enumerate(indicies):
idx=self.wrapped_index(ind)
a[i]=self[idx]
return a
elif isinstance(index, tuple):
idx=index[0]
idx=self.wrapped_index(idx)
return super().__getitem__(idx)
elif isinstance(index, int):
index = self.wrapped_index(index)
return super().__getitem__(index)
else:
print("RAISING EXCEPTION")
raise ValueError
def __setitem__(self, index, item):
index = self.wrapped_index(index)
return super().__setitem__(index, item)
def __array_finalize__(self, obj):
if obj is None: return
self.grid_size = getattr(obj, 'grid_size', obj.size)
pass
def wrapped_index(self, index):
idx=None
#if hasattr(index, '__iter__'): idx=
if type(index) is tuple:
idx=index[0] #only support for 1D!
elif type(index) is int:
idx=index
else:
raise ValueError('Unexpected index: {}'.format(index))
My Question
The above code actually works well, so in that sense my problem is solved. However, I have a question about how I have implemented handling slices in __getitem__
. In the above code, I have resorted to using a loop and manually looping over the needing indices and filling the appropriate values. This occurs in a Python loop, whereas I would like to let numpy
handle this loop if possible. I think this should be possible using numpy.take like so (looking just at the block of __getitem__
that handles slice objects:
def __getitem__(self, index):
if isinstance(index, slice):
start=index.start
stop=index.stop
indices=list(range(start, stop))
return super().take(self, indices, mode='wrap')
However, now I run into an error:
>>> b[1:5]
Traceback (most recent call last):
File "<stdin>", line 1, in <module>
File "path/to/class/file", line 130, in __getitem__
return super().take(self, indices, mode='wrap')
TypeError: 'list' object cannot be interpreted as an integer
As far as I can tell however, passing a list of indices to take
is allowed, judging by the first example in the documentation. I have suspicions that I am using super()
wrong but I am not sure. Any help here is appreciated.