Add numpy to typing deps

.pre-commit-config.yaml

@@ -26,3 +26,6 @@ repos:
             additional_dependencies:
                 - types-setuptools
                 - types-requests
+                - dask
+                - numpy
+                - numba

benchmarks/filter_3d.py

@@ -27,7 +27,7 @@ def gen_signal_array(ny, nx):
 mp_3d_filter = VolumeFilter(
     soma_diameter=soma_diameter,
     setup_params=setup_params,
-    planes_paths_range=signal_array,
+    n_planes=len(signal_array),
     n_locks_release=1,
 )
 

src/cellfinder_core/detect/detect.py

@@ -17,14 +17,13 @@
 from datetime import datetime
 from queue import Queue
 from threading import Lock
-from typing import Callable, List, Optional, Sequence, Tuple, TypeVar, Union
+from typing import Callable, List, Optional, Sequence, Tuple, TypeVar
 
-import dask.array as da
 import numpy as np
 from imlib.cells.cells import Cell
 from imlib.general.system import get_num_processes
 
-from cellfinder_core import logger
+from cellfinder_core import logger, types
 from cellfinder_core.detect.filters.plane import TileProcessor
 from cellfinder_core.detect.filters.setup_filters import setup_tile_filtering
 from cellfinder_core.detect.filters.volume.volume_filter import VolumeFilter
@@ -56,7 +55,7 @@ def calculate_parameters_in_pixels(
 
 
 def main(
-    signal_array: Union[np.ndarray, da.Array],
+    signal_array: types.array,
     start_plane: int,
     end_plane: int,
     voxel_sizes: Tuple[float, float, float],
@@ -128,7 +127,7 @@ def main(
         soma_diameter=soma_diameter,
         setup_params=setup_params,
         soma_size_spread_factor=soma_spread_factor,
-        planes_paths_range=signal_array,
+        n_planes=len(signal_array),
         n_locks_release=n_ball_procs,
         save_planes=save_planes,
         plane_directory=plane_directory,
@@ -152,7 +151,9 @@ def main(
     mp_ctx = multiprocessing.get_context("spawn")
     with mp_ctx.Pool(n_ball_procs) as worker_pool:
         async_results, locks = _map_with_locks(
-            mp_tile_processor.get_tile_mask, signal_array, worker_pool
+            mp_tile_processor.get_tile_mask,
+            signal_array,  # type: ignore
+            worker_pool,
         )
 
         # Release the first set of locks for the 2D filtering

src/cellfinder_core/detect/filters/plane/plane_filter.py

@@ -5,6 +5,7 @@
 import dask.array as da
 import numpy as np
 
+from cellfinder_core import types
 from cellfinder_core.detect.filters.plane.classical_filter import enhance_peaks
 from cellfinder_core.detect.filters.plane.tile_walker import TileWalker
 
@@ -28,7 +29,7 @@ class TileProcessor:
     n_sds_above_mean_thresh: float
 
     def get_tile_mask(
-        self, plane: da.array, lock: Optional[Lock] = None
+        self, plane: types.array, lock: Optional[Lock] = None
     ) -> Tuple[np.ndarray, np.ndarray]:
         """
         This thresholds the input plane, and returns a mask indicating which

src/cellfinder_core/detect/filters/volume/ball_filter.py

@@ -141,7 +141,7 @@ def append(self, plane: np.ndarray, mask: np.ndarray) -> None:
                 [e for e in plane.shape[:2]],
             )
             assert [e for e in mask.shape[:2]] == [
-                e for e in self.inside_brain_tiles.shape[2]
+                e for e in self.inside_brain_tiles.shape[:2]
             ], 'mask shape mismatch, expected"{}", got {}"'.format(
                 [e for e in self.inside_brain_tiles.shape[:2]],
                 [e for e in mask.shape[:2]],

src/cellfinder_core/detect/filters/volume/structure_detection.py

@@ -1,8 +1,9 @@
 from dataclasses import dataclass
-from typing import Dict, List, Sequence, Tuple, TypeVar, Union
+from typing import Dict, List, Optional, Sequence, Tuple, TypeVar, Union
 
 import numba.typed
 import numpy as np
+import numpy.typing as npt
 from numba import njit
 from numba.core import types
 from numba.experimental import jitclass
@@ -139,7 +140,7 @@ def __init__(self, width: int, height: int, start_z: int):
         )
 
     def process(
-        self, plane: np.ndarray, previous_plane: np.ndarray
+        self, plane: np.ndarray, previous_plane: Optional[np.ndarray]
     ) -> np.ndarray:
         """
         Process a new plane.
@@ -152,7 +153,7 @@ def process(
         return plane
 
     def connect_four(
-        self, plane: np.ndarray, previous_plane: np.ndarray
+        self, plane: np.ndarray, previous_plane: Optional[np.ndarray]
     ) -> np.ndarray:
         """
         Perform structure labelling.
@@ -209,7 +210,9 @@ def add_point(self, sid: int, point: np.ndarray) -> None:
         """
         self.coords_maps[sid] = np.row_stack((self.coords_maps[sid], point))
 
-    def add(self, x: int, y: int, z: int, neighbour_ids: List[int]) -> int:
+    def add(
+        self, x: int, y: int, z: int, neighbour_ids: npt.NDArray[np.uint64]
+    ) -> int:
         """
         For the current coordinates takes all the neighbours and find the
         minimum structure including obsolete structures mapping to any of
@@ -232,7 +235,7 @@ def add(self, x: int, y: int, z: int, neighbour_ids: List[int]) -> int:
         self.add_point(updated_id, point)
         return updated_id
 
-    def sanitise_ids(self, neighbour_ids: List[int]) -> int:
+    def sanitise_ids(self, neighbour_ids: npt.NDArray[np.uint64]) -> int:
         """
         Get the smallest ID of all the structures that are connected to IDs
         in `neighbour_ids`.
@@ -253,7 +256,7 @@ def sanitise_ids(self, neighbour_ids: List[int]) -> int:
         return int(updated_id)
 
     def merge_structures(
-        self, updated_id: int, neighbour_ids: List[int]
+        self, updated_id: int, neighbour_ids: npt.NDArray[np.uint64]
     ) -> None:
         """
         For all the neighbours, reassign all the points of neighbour to

src/cellfinder_core/detect/filters/volume/volume_filter.py

@@ -2,7 +2,7 @@
 import os
 from queue import Queue
 from threading import Lock
-from typing import Any, Callable, List, Optional, Sequence, Tuple
+from typing import Any, Callable, List, Optional, Tuple
 
 import numpy as np
 from imlib.cells.cells import Cell
@@ -30,7 +30,7 @@ def __init__(
         soma_diameter: float,
         soma_size_spread_factor: float = 1.4,
         setup_params: Tuple[np.ndarray, Any, int, int, float, Any],
-        planes_paths_range: Sequence,
+        n_planes: int,
         n_locks_release: int,
         save_planes: bool = False,
         plane_directory: Optional[str] = None,
@@ -41,7 +41,7 @@ def __init__(
     ):
         self.soma_diameter = soma_diameter
         self.soma_size_spread_factor = soma_size_spread_factor
-        self.n_planes = len(planes_paths_range)
+        self.n_planes = n_planes
         self.z = start_plane
         self.save_planes = save_planes
         self.plane_directory = plane_directory
@@ -55,7 +55,7 @@ def __init__(
         self.threshold_value = None
         self.setup_params = setup_params
 
-        self.previous_plane = None
+        self.previous_plane: Optional[np.ndarray] = None
 
         self.ball_filter = get_ball_filter(
             plane=self.setup_params[0],
@@ -66,7 +66,7 @@ def __init__(
         )
 
         self.cell_detector = get_cell_detector(
-            plane_shape=self.setup_params[0].shape,
+            plane_shape=self.setup_params[0].shape,  # type: ignore
             ball_z_size=self.setup_params[3],
             z_offset=self.setup_params[5],
         )

src/cellfinder_core/tools/array_operations.py

@@ -110,5 +110,4 @@ def bin_mean_3d(
             for j in range(i, i + bin_depth)
         ]
         binned_arr.append(np.dstack(sub_stack).mean(axis=2))
-    binned_arr = np.dstack(binned_arr)
-    return binned_arr
+    return np.dstack(binned_arr)

src/cellfinder_core/tools/geometry.py

@@ -27,11 +27,11 @@ def make_sphere(
     # generate the grid for the support points
     # centered at the position indicated by position
     grid = [slice(-x0, dim - x0) for x0, dim in zip(position, ball_shape)]
-    position = np.ogrid[grid]
+    meshedgrid = np.ogrid[grid]
     # calculate the distance of all points from `position` center
     # scaled by the radius
     arr = np.zeros(ball_shape, dtype=float)
-    for x_i, half_size in zip(position, half_sizes):
+    for x_i, half_size in zip(meshedgrid, half_sizes):
         arr += np.abs(x_i / half_size) ** 2
     # the inner part of the sphere will have distance below 1
     return arr <= 1.0