diff --git a/src/lgdo/types/vectorofvectors.py b/src/lgdo/types/vectorofvectors.py index 8d85276a..89a2bcb4 100644 --- a/src/lgdo/types/vectorofvectors.py +++ b/src/lgdo/types/vectorofvectors.py @@ -153,7 +153,10 @@ def __eq__(self, other: VectorOfVectors) -> bool: if isinstance(other, VectorOfVectors): return ( self.cumulative_length == other.cumulative_length - and np.all(self.flattened_data[:self.cumulative_length[-1]] == other.flattened_data[:other.cumulative_length[-1]]) + and np.all( + self.flattened_data[: self.cumulative_length[-1]] + == other.flattened_data[: other.cumulative_length[-1]] + ) and self.dtype == other.dtype and self.attrs == other.attrs ) @@ -337,7 +340,7 @@ def _set_vector_unsafe(self, i: int, vec: NDArray, lens: NDArray = None) -> None Vectors stored after index `i` can be overridden, producing unintended behavior. This method is typically used for fast sequential fill of a pre-allocated vector of vectors. - + If vec is 1D array and lens is None, set using full array. If vec is 2D, require lens to be included, and fill each array only up to lengths in lens. @@ -352,15 +355,15 @@ def _set_vector_unsafe(self, i: int, vec: NDArray, lens: NDArray = None) -> None append, replace, insert """ start = 0 if i == 0 else self.cumulative_length[i - 1] - if len(vec.shape)==1: + if len(vec.shape) == 1: vec = np.expand_dims(vec, axis=0) if lens is None: - lens = np.array([vec.shape[1]], dtype='u4') + lens = np.array([vec.shape[1]], dtype="u4") if not isinstance(lens, np.ndarray): - lens = np.array([lens], dtype='u4') + lens = np.array([lens], dtype="u4") cum_lens = start + lens.cumsum() - _nb_fill(vec, lens, self.flattened_data.nda[start:cum_lens[-1]]) - self.cumulative_length[i:i+len(lens)] = cum_lens + _nb_fill(vec, lens, self.flattened_data.nda[start : cum_lens[-1]]) + self.cumulative_length[i : i + len(lens)] = cum_lens def __iter__(self) -> Iterator[NDArray]: for j, stop in enumerate(self.cumulative_length): @@ -492,12 +495,34 @@ def _nb_build_cl(sorted_array_in: NDArray, cumulative_length_out: NDArray) -> ND ii += 1 return cumulative_length_out[:ii] -@numba.guvectorize([f"{t}[:,:],u4[:],{t}[:]" for t in ['b1', 'i1', 'i2', 'i4', 'i8', 'u1', 'u2', 'u4', 'u8', 'f4', 'f8', 'c8', 'c16']], - '(l,m),(l),(n)', cache=True, nopython=True) + +@numba.guvectorize( + [ + f"{t}[:,:],u4[:],{t}[:]" + for t in [ + "b1", + "i1", + "i2", + "i4", + "i8", + "u1", + "u2", + "u4", + "u8", + "f4", + "f8", + "c8", + "c16", + ] + ], + "(l,m),(l),(n)", + cache=True, + nopython=True, +) def _nb_fill(aoa_in: NDArray, len_in: NDArray, flattened_array_out: NDArray): """Vectorized function to fill flattened array from array of arrays and lengths. Values in aoa_in past lengths will not be copied. - + Parameters ---------- aoa_in @@ -508,16 +533,17 @@ def _nb_fill(aoa_in: NDArray, len_in: NDArray, flattened_array_out: NDArray): flattened array to copy values into. Must be longer than sum of lengths in len_in """ - + if len(flattened_array_out) < len_in.sum(): raise ValueError("flattened array not large enough to hold values") start = 0 for i, l in enumerate(len_in): - stop = start+l + stop = start + l flattened_array_out[start:stop] = aoa_in[i, :l] start = stop + def explode_cl(cumulative_length: NDArray, array_out: NDArray = None) -> NDArray: """Explode a `cumulative_length` array.