Merge pull request #170 from yfukai/predefined_edges

Predefined edges

pyproject.toml

@@ -53,6 +53,8 @@ matplotlib = "^3.5.1"
 sphinxcontrib-napoleon = "^0.7"
 autodoc_pydantic = "^1.7.2"
 pre-commit = "^2.20.0"
+nox = "^2022.8.7"
+nox-poetry = "^1.0.1"
 
 [tool.poetry.scripts]
 laptrack = "laptrack.__main__:main"

src/laptrack/_cost_matrix.py

@@ -6,6 +6,7 @@
 
 from ._coo_matrix_builder import coo_matrix_builder
 from ._typing_utils import Float
+from ._typing_utils import FloatArray
 from ._typing_utils import Matrix
 
 EPSILON = 1e-6
@@ -14,18 +15,18 @@
 def build_frame_cost_matrix(
     dist_matrix: coo_matrix_builder,
     *,
-    track_start_cost: Optional[Float],
-    track_end_cost: Optional[Float],
+    track_start_cost: Optional[Union[Float, FloatArray]],
+    track_end_cost: Optional[Union[Float, FloatArray]],
 ) -> coo_matrix:
     """Build sparce array for frame-linking cost matrix.
 
     Parameters
     ----------
     dist_matrix : Matrix or `_utils.coo_matrix_builder`
         The distance matrix for points at time t and t+1.
-    track_start_cost : Float, optional
+    track_start_cost : Float or FloatArray, optional
         The cost for starting the track (b in Jaqaman et al 2008 NMeth)
-    track_end_cost : Float, optional
+    track_end_cost : Float or FloatArray, optional
         The cost for ending the track (d in Jaqaman et al 2008 NMeth)
 
     Returns
@@ -50,8 +51,11 @@ def build_frame_cost_matrix(
         else:
             track_end_cost = 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
+    track_end_costs = np.ones(N) * track_end_cost
+    track_start_costs = np.ones(M) * track_start_cost
+
+    C[np.arange(M, M + N), np.arange(N)] = track_end_costs
+    C[np.arange(M), np.arange(N, N + M)] = track_start_costs
     min_val = np.min(C.data) if len(C.data) > 0 else 0
     C[dist_matrix.col + M, dist_matrix.row + N] = min_val