Source code for landlab.utils.jaggedarray

"""The JaggedArray class to store arrays of variable-length arrays.

Examples
--------

Create a JaggedArray that stores link IDs for the links attached to the
nodes of a 3x3 grid.

>>> from landlab.utils.jaggedarray import JaggedArray
>>> links_at_node = JaggedArray([
...     [0, 6],
...     [1, 7, 0],
...     [8, 1],
...     [2, 9, 6],
...     [3, 10, 2, 7],
...     [11, 3, 8],
...     [4, 7],
...     [5, 10, 4],
...     [5, 11]])

Make up some data that provides values at each of the links.

>>> value_at_link = np.arange(12, dtype=float)

Create another JaggedArray. Here we store the values at each of the links
attached to nodes of the grid.

>>> values_at_node = JaggedArray.empty_like(links_at_node, dtype=float)
>>> values_at_node.array[:] = value_at_link[links_at_node.array]

Now operate on the link values for each node.

>>> values_at_node.foreach_row(sum)
array([  6.,   8.,   9.,  17.,  22.,  22.,  11.,  19.,  16.])
>>> values_at_node.foreach_row(min)
array([ 0.,  0.,  1.,  2.,  2.,  3.,  4.,  4.,  5.])
>>> values_at_node.foreach_row(np.ptp)
array([ 6.,  7.,  7.,  7.,  8.,  8.,  3.,  6.,  6.])
"""
import numpy as np


[docs]def flatten_jagged_array(jagged, dtype=None): """Flatten a list of lists. Parameters ---------- jagged : array_like of array_like An array of arrays of unequal length. Returns ------- (data, offset) : (ndarray, ndarray of int) A tuple the data, as a flat numpy array, and offsets into that array for every item of the original list. Examples -------- >>> from landlab.utils.jaggedarray import flatten_jagged_array >>> data, offset = flatten_jagged_array([[1, 2], [], [3, 4, 5]], dtype=int) >>> data array([1, 2, 3, 4, 5]) >>> offset array([0, 2, 2, 5]) """ data = np.concatenate(jagged).astype(dtype=dtype) # if len(jagged) > 1: # data = np.concatenate(jagged).astype(dtype=dtype) # else: # data = np.array(jagged[0]).astype(dtype=dtype) items_per_block = np.array([len(block) for block in jagged], dtype=int) offset = np.empty(len(items_per_block) + 1, dtype=int) offset[0] = 0 offset[1:] = np.cumsum(items_per_block) return data, offset
[docs]def unravel(data, offset, out=None, pad=None): """Unravel a jagged array. Parameters ---------- data : ndarray Flattened-array of the data. offset : ndarray of int Offsets to the start of rows of the jagged array. out : ndarray Buffer into which to place the unravelled data. pad : number Value to use to pad rows of the jagged array. Returns ------- ndarray Matrix that holds the unravelled jagged array. """ from .ext.jaggedarray import unravel n_cols = np.diff(offset).max() if out is None: if pad is None: out = np.empty((len(offset) - 1, n_cols), dtype=data.dtype) else: out = np.full((len(offset) - 1, n_cols), pad, dtype=data.dtype) else: if pad is not None: out.fill(pad) unravel(data, offset, out) return out
[docs]class JaggedArray: """A container for an array of variable-length arrays. JaggedArray([row0, row1, ...]) JaggedArray(values, values_per_row) Examples -------- Create a JaggedArray with an array of arrays. >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.array array([0, 1, 2, 3, 4]) Create a JaggedArray as a 1D array and a list or row lengths. >>> x = JaggedArray([0, 1, 2, 3, 4], (3, 2)) >>> x.array array([0, 1, 2, 3, 4]) """
[docs] def __init__(self, *args): """JaggedArray([row0, row1, ...]) JaggedArray(values, values_per_row) Examples -------- Create a JaggedArray with an array of arrays. >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.array array([0, 1, 2, 3, 4]) >>> x = JaggedArray([[0, 1, 2]]) >>> x.array array([0, 1, 2]) >>> x.offset array([0, 3]) Create a JaggedArray as a 1D array and a list or row lengths. >>> x = JaggedArray([0, 1, 2, 3, 4], (3, 2)) >>> x.array array([0, 1, 2, 3, 4]) """ if len(args) == 1: values, values_per_row = ( np.concatenate(args[0]), [len(row) for row in args[0]], ) # if len(args[0]) > 1: # values, values_per_row = (np.concatenate(args[0]), # [len(row) for row in args[0]]) # else: # values, values_per_row = (np.array(args[0]), [len(args[0])]) else: values, values_per_row = (np.array(args[0]), args[1]) self._values = values self._number_of_rows = len(values_per_row) self._offsets = JaggedArray._offsets_from_values_per_row(values_per_row) self._offsets.flags["WRITEABLE"] = False
@property def array(self): """The jagged array as a 1D array. Returns ------- array : A view of the underlying 1D array. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.array array([0, 1, 2, 3, 4]) >>> x.array[0] = 1 >>> x.array array([1, 1, 2, 3, 4]) """ return self._values @property def offset(self): """The offsets to rows of a 1D array. Returns ------- array : Offsets into the underlying 1D array. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.offset array([0, 3, 5]) From the offsets you can get values for rows of the jagged array. >>> x.array[x.offset[0]:x.offset[1]] array([0, 1, 2]) Once the array is created, you can't change the offsets. >>> x.offset[0] = 1 Traceback (most recent call last): ValueError: assignment destination is read-only """ return self._offsets @property def size(self): """Number of array elements. Returns ------- int : Number of values in the array. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.size 5 """ return len(self._values) @property def number_of_rows(self): """Number of array rows. Returns ------- int : Number of rows in the array. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.number_of_rows 2 """ return self._number_of_rows @staticmethod def _offsets_from_values_per_row(values_per_row): """Get offsets into the base array from array lengths. Parameters ---------- values_per_row : array of int The number of values in each row of the JaggedArray. Returns ------- ndarray An array of offsets. """ offset = np.empty(len(values_per_row) + 1, dtype=int) np.cumsum(values_per_row, out=offset[1:]) offset[0] = 0 return offset
[docs] @staticmethod def empty_like(jagged, dtype=None): """Create a new JaggedArray that is like another one. Parameters ---------- jagged : JaggedArray A JaggedArray to copy. dtype : np.dtype The data type of the new JaggedArray. Returns ------- JaggedArray A new JaggedArray. """ return JaggedArray( np.empty_like(jagged.array, dtype=dtype), np.diff(jagged.offset) )
[docs] def length_of_row(self, row): """Number of values in a given row. Parameters ---------- row : int Index to a row. Returns ------- int : Number of values in the row. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.length_of_row(0) 3 >>> x.length_of_row(1) 2 """ return self._offsets[row + 1] - self._offsets[row]
[docs] def row(self, row): """Get the values of a row as an array. Parameters ---------- row : int Index to a row. Returns ------- array : Values in the row as a slice of the underlying array. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.row(0) array([0, 1, 2]) >>> x.row(1) array([3, 4]) >>> y = x.row(0) >>> y[0] = 1 >>> x.row(0) array([1, 1, 2]) """ return self._values[self._offsets[row] : self._offsets[row + 1]]
[docs] def __iter__(self): """Iterate over the rows of the array. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> for row in x: row array([0, 1, 2]) array([3, 4]) """ for row_number in range(self._number_of_rows): yield self.row(row_number)
[docs] def foreach_row(self, func, out=None): """Apply an operator row-by-row. Examples -------- >>> from landlab.utils.jaggedarray import JaggedArray >>> x = JaggedArray([[0, 1, 2], [3, 4]]) >>> x.foreach_row(sum) array([3, 7]) >>> out = np.empty(2, dtype=int) >>> x.foreach_row(sum, out=out) is out True >>> out array([3, 7]) """ if out is None: out = np.empty(self.number_of_rows, dtype=self._values.dtype) for row_number, row in enumerate(self): out[row_number] = func(row) return out