Use standard library instead of toolz for map, reduce, concat, accumulate (#890)

Author: tomwhite

cubed/array_api/manipulation_functions.py

@@ -1,10 +1,9 @@
 from bisect import bisect
+from functools import reduce
+from itertools import accumulate, chain
 from operator import add, mul
 from typing import Iterator
 
-import tlz
-from toolz import reduce
-
 from cubed.array_api.creation_functions import empty
 from cubed.backend_array_api import IS_IMMUTABLE_ARRAY
 from cubed.backend_array_api import namespace as nxp
@@ -39,7 +38,7 @@ def broadcast_arrays(*arrays):
 
     # Unify uneven chunking
     inds = [list(reversed(range(x.ndim))) for x in arrays]
-    uc_args = tlz.concat(zip(arrays, inds))
+    uc_args = chain.from_iterable(zip(arrays, inds))
     _, args = unify_chunks(*uc_args, warn=False)
 
     shape = broadcast_shapes(*(e.shape for e in args))
@@ -133,11 +132,11 @@ def concat(arrays, /, *, axis=0, chunks=None):
     inds = [list(range(x.ndim)) for x in arrays]
     for i, ind in enumerate(inds):
         ind[axis] = -(i + 1)
-    uc_args = tlz.concat(zip(arrays, inds))
+    uc_args = chain.from_iterable(zip(arrays, inds))
     chunkss, arrays = unify_chunks(*uc_args, warn=False)
 
     # offsets along axis for the start of each array
-    offsets = [0] + list(tlz.accumulate(add, [a.shape[axis] for a in arrays]))
+    offsets = [0] + list(accumulate([a.shape[axis] for a in arrays], add))
     in_shapes = tuple(array.shape for array in arrays)
 
     axis = validate_axis(axis, ndim)

cubed/core/array.py

@@ -1,8 +1,7 @@
+from functools import reduce
 from operator import mul
 from typing import Literal, Optional, TypeVar
 
-from toolz import map, reduce
-
 from cubed import config
 from cubed.backend_array_api import namespace as nxp
 from cubed.backend_array_api import numpy_array_to_backend_array

cubed/core/gufunc.py

@@ -1,5 +1,7 @@
+from itertools import chain
+
 import numpy as np
-from tlz import concat, merge, unique
+from tlz import merge, unique
 
 from cubed.vendor.dask.array.gufunc import _parse_gufunc_signature
 
@@ -151,7 +153,7 @@ def apply_gufunc(
             raise ValueError(f"Dimension `'{dim}'` with different chunksize present")
 
     # Apply function - use blockwise here
-    arginds = list(concat(zip(args, input_dimss)))
+    arginds = list(chain.from_iterable(zip(args, input_dimss)))
 
     from cubed.core.ops import blockwise
 

cubed/core/indexing.py

@@ -1,10 +1,10 @@
 import math
+from itertools import accumulate
 from operator import add
 from typing import TYPE_CHECKING
 
 import ndindex
 import numpy as np
-from toolz import accumulate, map
 
 from cubed.backend_array_api import backend_array_to_numpy_array
 from cubed.core.array import CoreArray
@@ -181,7 +181,9 @@ def _target_chunk_selection(target_chunks, idx, selection):
             offset = s.start or 0
             step = s.step if s.step is not None else 1
             start = tuple(
-                accumulate(add, tuple(x * step for x in target_chunks[i]), offset)
+                accumulate(
+                    tuple(x * step for x in target_chunks[i]), add, initial=offset
+                )
             )
             j = idx[i]
             sel.append(slice(start[j], start[j + 1], step))
@@ -190,7 +192,7 @@ def _target_chunk_selection(target_chunks, idx, selection):
         elif isinstance(s, np.ndarray):
             # find the cumulative chunk starts
             target_chunk_starts = [0] + list(
-                accumulate(add, [c for c in target_chunks[i]])
+                accumulate([c for c in target_chunks[i]], add)
             )
             # and use to slice the integer array
             j = idx[i]

cubed/core/ops.py

@@ -3,7 +3,7 @@
 import numbers
 from dataclasses import dataclass
 from functools import partial
-from itertools import product
+from itertools import chain, product
 from numbers import Integral, Number
 from typing import (
     TYPE_CHECKING,
@@ -20,8 +20,7 @@
 
 import numpy as np
 import zarr
-from tlz import concat, first, partition
-from toolz import map
+from tlz import first, partition
 
 from cubed import config
 from cubed.backend_array_api import IS_IMMUTABLE_ARRAY, numpy_array_to_backend_array
@@ -619,7 +618,7 @@ def elemwise(func, *args: "Array", dtype=None) -> "Array":
     return blockwise(
         func,
         expr_inds,
-        *concat((a, tuple(range(a.ndim)[::-1])) for a in args),
+        *chain.from_iterable((a, tuple(range(a.ndim)[::-1])) for a in args),
         dtype=dtype,
     )
 
@@ -887,7 +886,7 @@ def _map_blocks(
     return blockwise(
         func,
         out_ind,
-        *concat(argpairs),
+        *chain.from_iterable(argpairs),
         dtype=dtype,
         adjust_chunks=adjust_chunks,
         new_axes=new_axes,

cubed/diagnostics/tqdm.py

@@ -1,8 +1,6 @@
 import contextlib
 import sys
 
-from toolz import map
-
 from cubed.runtime.pipeline import visit_nodes
 from cubed.runtime.types import Callback
 

cubed/primitive/blockwise.py

@@ -24,7 +24,6 @@
 
 import toolz
 import zarr
-from toolz import map
 
 from cubed.backend_array_api import (
     backend_array_to_numpy_array,

cubed/storage/stores/tensorstore.py

@@ -1,12 +1,12 @@
 import dataclasses
 import math
+from functools import reduce
 from itertools import starmap
 from operator import mul
 from typing import Any, Dict, Optional, Union
 
 import numpy as np
 import tensorstore
-from toolz import reduce
 
 from cubed.types import T_DType, T_RegularChunks, T_Shape, T_Store
 from cubed.utils import itemsize, join_path

cubed/utils.py

@@ -7,7 +7,8 @@
 import traceback
 from collections.abc import Iterable
 from dataclasses import dataclass
-from functools import partial
+from functools import partial, reduce
+from itertools import accumulate
 from math import prod
 from operator import add, mul
 from pathlib import Path
@@ -17,8 +18,6 @@
 from urllib.parse import quote, unquote, urlsplit, urlunsplit
 
 import numpy as np
-import tlz as toolz
-from toolz import reduce
 
 from cubed.backend_array_api import backend_dtype_to_numpy_dtype
 from cubed.backend_array_api import namespace as nxp
@@ -76,9 +75,9 @@ def get_item(chunks: T_RectangularChunks, idx: Tuple[int, ...]) -> Tuple[slice,
 
 def _cumsum(seq, initial_zero=False):
     if initial_zero:
-        return tuple(toolz.accumulate(add, seq, 0))
+        return tuple(accumulate(seq, add, initial=0))
     else:
-        return tuple(toolz.accumulate(add, seq))
+        return tuple(accumulate(seq, add))
 
 
 def join_path(dir_url: PathType, child_path: str) -> str: