Merge pull request #30 from yfukai/fix_zero_distance_problem

Fix zero distance problem

Author: yfukai

codecov.yml

@@ -3,7 +3,7 @@ coverage:
   status:
     project:
       default:
-        target: "90"
+        target: "80"
     patch:
       default:
-        target: "100"
+        target: "80"

pyproject.toml

@@ -1,6 +1,6 @@
 [tool.poetry]
 name = "laptrack"
-version = "0.1.2"
+version = "0.1.3"
 description = "LapTrack"
 authors = ["Yohsuke Fukai <ysk@yfukai.net>"]
 license = "BSD-3-Clause"

src/laptrack/_cost_matrix.py

@@ -2,26 +2,25 @@
 from typing import Union
 
 import numpy as np
-from scipy.sparse import lil_matrix
+from scipy.sparse.coo import coo_matrix
 
 from ._typing_utils import Float
-from ._typing_utils import FloatArray
+from ._typing_utils import Matrix
+from ._utils import coo_matrix_builder
 
 
 def build_frame_cost_matrix(
-    dist_matrix: FloatArray,
-    track_cost_cutoff: Float,
+    dist_matrix: coo_matrix_builder,
+    *,
     track_start_cost: Optional[Float],
     track_end_cost: Optional[Float],
-) -> lil_matrix:
+) -> coo_matrix:
     """Build sparce array for frame-linking cost matrix.
 
     Parameters
     ----------
-    dist_matrix : FloatArray
+    dist_matrix : Matrix or `_utils.coo_matrix_builder`
         The distance matrix for points at time t and t+1.
-    track_cost_cutoff : Float, optional
-        The distance cutoff for the connected points in the track
     track_start_cost : Float, optional
         The cost for starting the track (b in Jaqaman et al 2008 NMeth)
     track_end_cost : Float, optional
@@ -35,43 +34,43 @@ def build_frame_cost_matrix(
     M = dist_matrix.shape[0]
     N = dist_matrix.shape[1]
 
-    C = lil_matrix((M + N, N + M), dtype=np.float32)
-    ind = np.where(dist_matrix < track_cost_cutoff)
-    C[(*ind,)] = dist_matrix[(*ind,)]
+    C = coo_matrix_builder((M + N, N + M), dtype=np.float32)
+    C.append_matrix(dist_matrix)
+
     if track_start_cost is None:
-        track_start_cost = np.concatenate(C.data).max() * 1.05
+        track_start_cost = np.max(C.data) * 1.05
     if track_end_cost is None:
-        track_end_cost = np.concatenate(C.data).max() * 1.05
+        track_end_cost = np.max(C.data) * 1.05
+
     C[np.arange(M, M + N), np.arange(N)] = np.ones(N) * track_end_cost
     C[np.arange(M), np.arange(N, N + M)] = np.ones(M) * track_start_cost
-    ind2 = [ind[1] + M, ind[0] + N]
-    min_val = np.concatenate(C.data).min() if len(C.data) > 0 else 0
-    C[(*ind2,)] = min_val
+    min_val = np.min(C.data) if len(C.data) > 0 else 0
+    C[dist_matrix.col + M, dist_matrix.row + N] = min_val
 
-    return C
+    return C.to_coo_matrix()
 
 
 def build_segment_cost_matrix(
-    gap_closing_dist_matrix: Union[lil_matrix, FloatArray],
-    splitting_dist_matrix: Union[lil_matrix, FloatArray],
-    merging_dist_matrix: Union[lil_matrix, FloatArray],
+    gap_closing_dist_matrix: Union[coo_matrix_builder, Matrix],
+    splitting_dist_matrix: Union[coo_matrix_builder, Matrix],
+    merging_dist_matrix: Union[coo_matrix_builder, Matrix],
     track_start_cost: Optional[Float],
     track_end_cost: Optional[Float],
     no_splitting_cost: Optional[Float],
     no_merging_cost: Optional[Float],
     alternative_cost_factor: Float = 1.05,
     alternative_cost_percentile: Float = 90,
     alternative_cost_percentile_interpolation: str = "lower",
-) -> lil_matrix:
+) -> coo_matrix:
     """Build sparce array for segment-linking cost matrix.
 
     Parameters
     ----------
-    gap_closing_dist_matrix : lil_matrix or FloatArray
+    gap_closing_dist_matrix : coo_matrix_builder or Matrix
         The distance matrix for closing gaps between segment i and j.
-    splitting_dist_matrix : lil_matrix or FloatArray
+    splitting_dist_matrix : coo_matrix_builder or Matrix
         The distance matrix for splitting between segment i and time/index j
-    merging_dist_matrix : lil_matrix or FloatArray
+    merging_dist_matrix : coo_matrix_builder or Matrix
         The distance matrix for merging between segment i and time/index j
     track_start_cost : Float, optional
         The cost for starting the track (b in Jaqaman et al 2008 NMeth)
@@ -92,7 +91,7 @@ def build_segment_cost_matrix(
 
     Returns
     -------
-    cost_matrix : Optional[FloatArray]
+    cost_matrix : Optional[coo_matrix]
         the cost matrix for frame linking, None if not appropriate
     """
     M = gap_closing_dist_matrix.shape[0]
@@ -104,14 +103,14 @@ def build_segment_cost_matrix(
 
     S = 2 * M + N1 + N2
 
-    C = lil_matrix((S, S), dtype=np.float32)
+    C = coo_matrix_builder((S, S), dtype=np.float32)
 
-    C[:M, :M] = gap_closing_dist_matrix
-    C[M : M + N1, :M] = splitting_dist_matrix.T
-    C[:M, M : M + N2] = merging_dist_matrix
+    C.append_matrix(gap_closing_dist_matrix)
+    C.append_matrix(splitting_dist_matrix.T, shift=(M, 0))
+    C.append_matrix(merging_dist_matrix, shift=(0, M))
 
-    all_data = np.concatenate(C.data)
-    if len(all_data) == 0:
+    upper_left_size = C.size()
+    if upper_left_size == 0:
         return None
 
     # Note:
@@ -133,7 +132,7 @@ def build_segment_cost_matrix(
     ):
         alternative_cost = (
             np.percentile(
-                all_data,
+                C.data,
                 alternative_cost_percentile,
                 interpolation=alternative_cost_percentile_interpolation,
             )
@@ -152,11 +151,12 @@ def build_segment_cost_matrix(
     C[np.arange(2 * M + N1, S), np.arange(M, M + N2)] = np.ones(N2) * no_merging_cost
     C[np.arange(M), np.arange(M + N2, 2 * M + N2)] = np.ones(M) * track_end_cost
     C[np.arange(M, M + N1), np.arange(2 * M + N2, S)] = np.ones(N1) * no_splitting_cost
-    subC = C[: M + N1, : M + N2]
-    inds = subC.nonzero()
     min_val = np.min(
         [track_start_cost, track_end_cost, no_splitting_cost, no_merging_cost]
     )
-    C[inds[1] + M + N1, inds[0] + M + N2] = min_val
 
-    return C
+    upper_left_rows = C.row[:upper_left_size]
+    upper_left_cols = C.col[:upper_left_size]
+    C[upper_left_cols + M + N1, upper_left_rows + M + N2] = min_val
+
+    return C.to_coo_matrix()

src/laptrack/_optimization.py

@@ -1,16 +1,16 @@
 from typing import Tuple
 
 import lap
-from scipy.sparse import lil_matrix
 from scipy.sparse.csr import csr_matrix
 
 from ._typing_utils import FloatArray
 from ._typing_utils import Int
 from ._typing_utils import IntArray
+from ._typing_utils import Matrix
 
 
 def __to_lap_sparse(
-    cost_matrix: lil_matrix,
+    cost_matrix: Matrix,
 ) -> Tuple[Int, FloatArray, IntArray, IntArray]:
     """Convert data for lap.lapmod."""
     n = cost_matrix.shape[0]
@@ -19,12 +19,12 @@ def __to_lap_sparse(
     return n, cost_matrix2.data, cost_matrix2.indptr, cost_matrix2.indices
 
 
-def lap_optimization(cost_matrix: lil_matrix) -> Tuple[float, IntArray, IntArray]:
+def lap_optimization(cost_matrix: Matrix) -> Tuple[float, IntArray, IntArray]:
     """Solves the linear assignment problem for a sparse matrix.
 
     Parameters
     ----------
-    cost_matrix : lil_matrix
+    cost_matrix : lil_matrix or coo_matrix
         the cost matrix in scipy.sparse.lil_matrix format
 
     Returns

src/laptrack/_tracking.py

@@ -18,13 +18,13 @@
 import numpy as np
 import pandas as pd
 from scipy.spatial.distance import cdist
-from scipy.sparse import lil_matrix
 
 from ._cost_matrix import build_frame_cost_matrix, build_segment_cost_matrix
 from ._optimization import lap_optimization
 from ._typing_utils import Float
 from ._typing_utils import FloatArray
 from ._typing_utils import Int
+from ._utils import coo_matrix_builder
 
 
 def laptrack(
@@ -137,9 +137,16 @@ def laptrack(
     # linking between frames
     for frame, (coord1, coord2) in enumerate(zip(coords[:-1], coords[1:])):
         dist_matrix = cdist(coord1, coord2, metric=track_dist_metric)
+        ind = np.where(dist_matrix < track_cost_cutoff)
+        dist_matrix = coo_matrix_builder(
+            dist_matrix.shape,
+            row=ind[0],
+            col=ind[1],
+            data=dist_matrix[(*ind,)],
+            dtype=dist_matrix.dtype,
+        )
         cost_matrix = build_frame_cost_matrix(
             dist_matrix,
-            track_cost_cutoff=track_cost_cutoff,
             track_start_cost=track_start_cost,
             track_end_cost=track_end_cost,
         )
@@ -203,7 +210,7 @@ def to_gap_closing_candidates(row):
                     indices2 = np.where(
                         target_dist_matrix[0] < gap_closing_cost_cutoff
                     )[0]
-                    return df.index[indices2], target_dist_matrix[0][indices2]
+                    return df.index[indices2].values, target_dist_matrix[0][indices2]
                 else:
                     return [], []
 
@@ -213,12 +220,14 @@ def to_gap_closing_candidates(row):
         else:
             segments_df["gap_closing_candidates"] = [[]] * len(segments_df)
 
-        gap_closing_dist_matrix = lil_matrix((N_segments, N_segments), dtype=np.float32)
+        gap_closing_dist_matrix = coo_matrix_builder(
+            (N_segments, N_segments), dtype=np.float32
+        )
         for ind, row in segments_df.iterrows():
             candidate_inds = row["gap_closing_candidates"][0]
             candidate_costs = row["gap_closing_candidates"][1]
             # row ... track end, col ... track start
-            gap_closing_dist_matrix[(ind, candidate_inds)] = candidate_costs
+            gap_closing_dist_matrix[(int(ind), candidate_inds)] = candidate_costs
 
         all_candidates = {}
         dist_matrices = {}
@@ -262,7 +271,9 @@ def to_candidates(row):
             )
 
             N_middle = len(all_candidates[prefix])
-            dist_matrices[prefix] = lil_matrix((N_segments, N_middle), dtype=np.float32)
+            dist_matrices[prefix] = coo_matrix_builder(
+                (N_segments, N_middle), dtype=np.float32
+            )
 
             all_candidates_dict = {
                 tuple(val): i for i, val in enumerate(all_candidates[prefix])
@@ -273,7 +284,7 @@ def to_candidates(row):
                     all_candidates_dict[tuple(fi)] for fi in candidate_frame_indices
                 ]
                 candidate_costs = row[f"{prefix}_candidates"][1]
-                dist_matrices[prefix][(ind, candidate_inds)] = candidate_costs
+                dist_matrices[prefix][(int(ind), candidate_inds)] = candidate_costs
 
         splitting_dist_matrix = dist_matrices["first"]
         merging_dist_matrix = dist_matrices["last"]

src/laptrack/_typing_utils.py

@@ -3,10 +3,14 @@
 
 import numpy as np
 import numpy.typing as npt
+from scipy.sparse import lil_matrix
+from scipy.sparse.coo import coo_matrix
 
 NumArray = npt.NDArray[Any]
 FloatArray = npt.NDArray[np.float_]
 IntArray = npt.NDArray[np.int_]
 
 Int = Union[int, np.int_]
 Float = Union[float, np.float_]
+
+Matrix = Union[FloatArray, coo_matrix, lil_matrix]