diff --git a/src/plane.py b/src/plane.py index 86ecf68..4ece270 100644 --- a/src/plane.py +++ b/src/plane.py @@ -24,27 +24,22 @@ class Plane: numpy grid, without the same bloat as a straightforward N-dimensional grid of booleans for instance. """ - def __init__(self, size, N): + def __init__(self, shape): """ Construction of a plane. There are three cases: - If N == 0: We have an undefined plane. Nothing is in it. - If N == 1: We have a 1D plane. This is represented by a single number. - Otherwise: We have an N-D plane. Everything operates as expected. - - If N happens to be negative, an exception is thrown. + If shape is the empty tuple, we have an undefined plane. Nothing is in it. + If shape is length 1, we have a 1D plane. This is represented by a single number. + Otherwise, we have an N-D plane. Everything operates as expected. """ - if N < 0: - raise ValueError('Negative dimension nonsensical') - elif N == 0: - self.shape = () - self.grid = np.array([], dtype=np.object) - elif N == 1: - self.shape = (size,) - self.grid = np.array([0], dtype=np.object) + self.shape = tuple(shape) + + if len(shape) == 0: + self.grid = None + elif len(shape) == 1: + self.grid = 0 else: - self.shape = (size,) * N - self.grid = np.zeros((size**(N-1),), dtype=np.object) + self.grid = np.zeros((np.prod(shape[:-1]),), dtype=np.object) def __getitem__(self, idx): """ @@ -57,7 +52,7 @@ class Plane: # Passed in coordinates, access incrementally # Note this could be a tuple of slices or numbers - if type(idx) in [tuple]: + if type(idx) is tuple: tmp = self for i in idx: tmp = tmp[i] @@ -65,23 +60,27 @@ class Plane: # Reached last dimension, return bits instead elif len(self.shape) == 1: - bits = bm.bits_of(self.grid[0], self.shape[0])[idx] + bits = bm.bits_of(self.grid, self.shape[0])[idx] if isinstance(idx, slice): return list(map(int, bits)) else: return int(bits) # Simply relay to numpy methods - # Note this doesn't necessarily return a grid in the same notion but + # We check if we reach an actual number as opposed to a tuple # does still allow further indexing if desired. In addition, we can # be confident idx is either a list or a number so the final dimension # cannot be accessed from here - # TODO: Reconsider this... else: - full = np.reshape(self.grid, self.shape[:-1])[idx] - tmp = cls(1, len(self.shape) - 1) - tmp.grid = full.flat - return tmp + tmp = np.reshape(self.grid, self.shape[:-1])[idx] + try: + plane = Plane(tmp.shape + self.shape[-1:]) + plane.grid = tmp.flat + except AttributeError: + plane = Plane(self.shape[-1:]) + plane.grid = tmp + + return plane def _flatten(coordinates): """