Multi dimensional array with automatic expansion

fitarray.py

#!python
# Multi dimensional array with automatic expansion
# compatible with masked arrays


import numpy as np

# convert a flat index to a dimension indexes
# 7 => (1,1,1)
def static_index_to_key(itemsize, strides, index):
  key = [0] * len(strides)
  index *= itemsize
  last_d = 0
  for d in range(len(strides)):
    key[d] = int((index - last_d) / strides[d])
    last_d += key[d] * strides[d]
  return(tuple(key))

# convert dimension indexes to flat index
# (1,1,1) => 7
def static_key_to_index(itemsize, strides, key):
  index = 0
  for d in range(len(strides)):
    index += (strides[d] / itemsize) * key[d]
  return(index)


# Multi dimensional array with automatic expansion
# You can assign values outside the current bounds and the array will fit to accomodate
# If you query a key outside the current bound it will return None
# v1.0 10/2016 pemn
class FitArray(np.ndarray):
  _default_value = np.nan
  # custom serializer that can handle multiple dimensions and None type
  # the default serializar will break when it sees a None instead of dtype
  def __str__(self):
    r = ""
    for i in range(self.size):
      r += "%d %s %s\n" % (i, self.index_to_key(i), self.flat[i])
    return(r)

  # get a value, or None if outside current bounds
  def __getitem__(self, key):
    # Ensure key has tuple type
    if type(key)==int: 
      key=key, 
    for i in range(self.ndim):
      # the requested key is outside the current bounds
      if i < len(key) and (isinstance(key[i], int) and key[i] >= self.shape[i]):
        break 
    else:
      # inside bounds
      return(super().__getitem__(key))
    # print("key def",key)
    # out of bounds
    return(self._default_value)

  # set a value, extending array dimensions as needed
  def __setitem__(self, key, value):
    # Ensure key has list type
    if type(key)==int: 
      key=[key, 0]
    new_shape = list(self.shape)
    for i in range(max(len(key),self.ndim)):
      if(i >= len(new_shape)):
        new_shape.append(1)
      # if we are trying to access beyond the current dimensions, we need to extend
      if key[i] >= new_shape[i]:
        new_shape[i] = key[i]+1
      # if key is smaller than the current data, adjust so it will access the first element of each dimension
      # Ex.: [3,3] => [3,3,0,0]
      if i >= len(key):
        key.append(0)
    if(self.shape != tuple(new_shape)):
      self.fit(new_shape)
    return super().__setitem__(key, value)

  # fit the array to a new (bigger) shape, while preserving the key addresses
  def fit(self, new_shape):
    old_size, old_strides = self.size, self.strides
    # resize will give the array a new shape, but data will fall on different keys addresses
    super().resize(new_shape, refcheck=False)
    # fixes the data positions according to old shape key addresses
    for i in range(self.size -1, -1, -1):
      # the current key on the new shape
      new_key = self.index_to_key(i)
      # the flat index of the current key in the old shape
      old_i = static_key_to_index(self.itemsize, old_strides, new_key)
      # the corresponding key on the old shape
      old_key = static_index_to_key(self.itemsize, old_strides, old_i)
      if(old_i < old_size and old_key == new_key):
        # put on this position the correct value restored directly from the position that inherited the data after the resize
        super().__setitem__(new_key, super().__getitem__(self.index_to_key(old_i)))
      else:
        super().__setitem__(new_key, self._default_value)

  def index_to_key(self, index):
    return(static_index_to_key(self.itemsize, self.strides, index))

  def key_to_index(self, index):
    return(static_key_to_index(self.itemsize, self.strides, index))


# entry point
if __name__ == "__main__":

  # 2 dimensions, fixed starting size
  a = FitArray((3,3), dtype = np.float_)
  for i in range(a.shape[0]):
    for j in range(a.shape[1]):
      a[i,j] = i * 10 + j
  # resize to a (4,4) shape preserving the current data structure
  a[3,3] = 33
  print(a)

  # 3 dimensions, free starting shape
  b = FitArray((0,0,0), 'float')
  b[1,1,1] = 111
  print(b)