Overlap trace adjustment

.github/pull_request_template.md

@@ -21,10 +21,9 @@ Describe the tests that you ran to verify your changes.
 Provide instructions so that others can reproduce.
 -->
 
-- [ ] Ran unit tests
-- [ ] Ran integration tests
-- [ ] Performed manual testing
-- [ ] Updated existing tests
+- [ ] Unit tests
+- [ ] Integration tests
+- [ ] Existing tests update
 
 
 ## 🛠️ Dependencies
@@ -36,9 +35,5 @@ List any new dependencies added or existing ones updated.
 
 ## ✅ Checklist
 
-- [ ] My code follows the project's style guidelines
 - [ ] I have commented my code, particularly in hard-to-understand areas
 - [ ] I have made corresponding changes to the documentation
-- [ ] My changes generate no new warnings
-- [ ] I have added tests that prove my fix is effective or that my feature works
-- [ ] New and existing unit tests pass locally with my changes

src/cala/assets.py

@@ -86,6 +86,7 @@ class Footprints(Asset):
             member=Footprint.entity(),
             group_by=Dims.component,
             checks=[is_non_negative, has_no_nan],
+            allow_extra_coords=False,
         )
     )
 
@@ -134,6 +135,7 @@ def from_array(
             member=Trace.entity(),
             group_by=Dims.component,
             checks=[is_non_negative, has_no_nan],
+            allow_extra_coords=False,
         )
     )
 
@@ -145,6 +147,7 @@ class Movie(Asset):
             member=Frame.entity(),
             group_by=Dims.frame.value,
             checks=[is_non_negative, has_no_nan],
+            allow_extra_coords=False,
         )
     )
 
@@ -180,6 +183,7 @@ class CompStats(Asset):
             dims=comp_dims,
             dtype=float,
             checks=[is_non_negative, has_no_nan],
+            allow_extra_coords=False,
         )
     )
 
@@ -191,6 +195,7 @@ class PixStats(Asset):
             dims=(Dims.width.value, Dims.height.value, Dims.component.value),
             dtype=float,
             checks=[is_non_negative, has_no_nan],
+            allow_extra_coords=False,
         )
     )
 
@@ -202,6 +207,7 @@ class Overlaps(Asset):
             dims=comp_dims,
             dtype=bool,
             checks=[has_no_nan],
+            allow_extra_coords=False,
         )
     )
 
@@ -224,5 +230,6 @@ class Residual(Asset):
             member=Frame.entity(),
             group_by=Dims.frame.value,
             checks=[is_non_negative, has_no_nan],
+            allow_extra_coords=False,
         )
     )

src/cala/models/axis.py

@@ -3,7 +3,7 @@
 
 from pydantic import BaseModel, Field
 
-from cala.models.checks import is_unique, is_unit_interval
+from cala.models.checks import has_no_nan, is_unique
 
 
 class Axis:
@@ -17,7 +17,7 @@ class Axis:
 
     id_coord: str = "id_"
     timestamp_coord: str = "timestamp"
-    confidence_coord: str = "confidence"
+    detect_coord: str = "detected_on"
     frame_coord: str = "frame"
     width_coord: str = "width"
     height_coord: str = "height"
@@ -37,7 +37,7 @@ def component_rename(self) -> dict[str, str]:
         return {
             AXIS.component_dim: f"{AXIS.component_dim}'",
             AXIS.id_coord: f"{AXIS.id_coord}'",
-            AXIS.confidence_coord: f"{AXIS.confidence_coord}'",
+            AXIS.detect_coord: f"{AXIS.detect_coord}'",
         }
 
 
@@ -62,11 +62,11 @@ class Coords(Enum):
     width = Coord(name=AXIS.width_coord, dtype=int, checks=[is_unique])
     frame = Coord(name=AXIS.frame_coord, dtype=int, checks=[is_unique])
     timestamp = Coord(name=AXIS.timestamp_coord, dtype=str, checks=[is_unique])
-    confidence = Coord(name=AXIS.confidence_coord, dtype=float, checks=[is_unit_interval])
+    detected = Coord(name=AXIS.detect_coord, dtype=int, checks=[has_no_nan])
 
 
 class Dims(Enum):
     width = Dim(name=AXIS.width_dim, coords=[Coords.width.value])
     height = Dim(name=AXIS.height_dim, coords=[Coords.height.value])
     frame = Dim(name=AXIS.frames_dim, coords=[Coords.frame.value, Coords.timestamp.value])
-    component = Dim(name=AXIS.component_dim, coords=[Coords.id.value, Coords.confidence.value])
+    component = Dim(name=AXIS.component_dim, coords=[Coords.id.value, Coords.detected.value])

src/cala/models/entity.py

@@ -18,6 +18,7 @@ class Entity(BaseModel):
     coords: list[Coord] = Field(default_factory=list)
     dtype: type
     checks: list[Callable] = Field(default_factory=list)
+    allow_extra_coords: bool = True
 
     _model: DataArraySchema = PrivateAttr(DataArraySchema())
 
@@ -47,15 +48,14 @@ def to_schema(self) -> DataArraySchema:
             checks=self.checks,
         )
 
-    @staticmethod
-    def _build_coord_schema(coords: list[Coord]) -> CoordsSchema:
+    def _build_coord_schema(self, coords: list[Coord]) -> CoordsSchema:
         spec = dict()
 
         for c in coords:
             dim = DimsSchema((c.dim,)) if c.dim else None
             spec[c.name] = DataArraySchema(dims=dim, dtype=DTypeSchema(c.dtype), checks=c.checks)
 
-        return CoordsSchema(spec)
+        return CoordsSchema(spec, allow_extra_keys=self.allow_extra_coords)
 
 
 class Group(Entity):

src/cala/nodes/cleanup.py

@@ -5,10 +5,31 @@
 import xarray as xr
 from noob import Name
 
-from cala.assets import CompStats, Footprints, Overlaps, PixStats, Traces
+from cala.assets import CompStats, Footprints, Overlaps, PixStats, Residual, Traces
 from cala.models import AXIS
 
 
+def clear_overestimates(
+    footprints: Footprints, residuals: Residual, nmf_error: float
+) -> A[Footprints, Name("footprints")]:
+    """
+    Remove all sections of the footprints that cause negative residuals.
+
+    This occurs by:
+    1. find "significant" negative residual spots that is more than a noise level, and thus
+    cannot be clipped to zero. !!!! (only of the latest frame, and then go back to trace update..?)
+    2. all footprint values at these spots go to zero.
+    """
+    if residuals.array is None:
+        return footprints
+    R_min = residuals.array.isel({AXIS.frames_dim: -1}).reset_coords(
+        [AXIS.frame_coord, AXIS.timestamp_coord], drop=True
+    )
+    tuned_fp = footprints.array.where(R_min > -nmf_error, 0, drop=False)
+
+    return tuned_fp
+
+
 def purge_razed_components(
     footprints: Footprints,
     traces: Traces,
@@ -101,13 +122,13 @@ def filter_components(
             comp_stats.array.set_xindex(AXIS.id_coord)
             .set_xindex(f"{AXIS.id_coord}'")
             .sel({AXIS.id_coord: keep_ids, f"{AXIS.id_coord}'": keep_ids.values.tolist()})
-            .reset_index(AXIS.id_coord)
+            .reset_index([AXIS.id_coord, f"{AXIS.id_coord}'"])
         )
         overlaps.array = (
             overlaps.array.set_xindex(AXIS.id_coord)
             .set_xindex(f"{AXIS.id_coord}'")
             .sel({AXIS.id_coord: keep_ids, f"{AXIS.id_coord}'": keep_ids.values.tolist()})
-            .reset_index(AXIS.id_coord)
+            .reset_index([AXIS.id_coord, f"{AXIS.id_coord}'"])
         )
 
     return footprints, traces, pix_stats, comp_stats, overlaps

src/cala/nodes/detect/catalog.py

@@ -1,6 +1,7 @@
 from collections.abc import Hashable, Iterable
 from typing import Annotated as A
 
+import cv2
 import numpy as np
 import xarray as xr
 from noob import Name
@@ -72,21 +73,19 @@ def process(
         footprints = xr.concat(
             footprints,
             dim=AXIS.component_dim,
-            coords=[AXIS.id_coord, AXIS.confidence_coord],
+            coords=[AXIS.id_coord, AXIS.detect_coord],
             combine_attrs=combine_attr_replaces,
         )
         traces = xr.concat(
             traces,
             dim=AXIS.component_dim,
-            coords=[AXIS.id_coord, AXIS.confidence_coord],
+            coords=[AXIS.id_coord, AXIS.detect_coord],
             combine_attrs=combine_attr_replaces,
         )
 
         return Footprints.from_array(footprints), Traces.from_array(traces)
 
-    def _register(
-        self, new_fp: xr.DataArray, new_tr: xr.DataArray, confidence: float = 0.0
-    ) -> tuple[Footprint, Trace]:
+    def _register(self, new_fp: xr.DataArray, new_tr: xr.DataArray) -> tuple[Footprint, Trace]:
 
         new_id = create_id()
 
@@ -95,7 +94,10 @@ def _register(
             .assign_coords(
                 {
                     AXIS.id_coord: (AXIS.component_dim, [new_id]),
-                    AXIS.confidence_coord: (AXIS.component_dim, [confidence]),
+                    AXIS.detect_coord: (
+                        AXIS.component_dim,
+                        [new_tr[AXIS.frame_coord].max().item()],
+                    ),
                 }
             )
             .isel({AXIS.component_dim: 0})
@@ -105,7 +107,10 @@ def _register(
             .assign_coords(
                 {
                     AXIS.id_coord: (AXIS.component_dim, [new_id]),
-                    AXIS.confidence_coord: (AXIS.component_dim, [confidence]),
+                    AXIS.detect_coord: (
+                        AXIS.component_dim,
+                        [new_tr[AXIS.frame_coord].max().item()],
+                    ),
                 }
             )
             .isel({AXIS.component_dim: 0})
@@ -211,8 +216,21 @@ def _merge_matrix(
             fps_base = fps_base.rename(AXIS.component_rename)
             trs_base = trs_base.rename(AXIS.component_rename)
 
+        fps = self._expand_boundary(fps > 0)
+
         overlaps = fps @ fps_base > 0
         # this should later reflect confidence
         corrs = xr.corr(trs, trs_base, dim=AXIS.frames_dim) > self.merge_threshold
 
         return overlaps * corrs
+
+    def _expand_boundary(self, mask: xr.DataArray) -> xr.DataArray:
+        kernel = cv2.getStructuringElement(cv2.MORPH_CROSS, (3, 3))
+        return xr.apply_ufunc(
+            lambda x: cv2.morphologyEx(x, cv2.MORPH_DILATE, kernel, iterations=1),
+            mask.astype(np.uint8),
+            input_core_dims=[AXIS.spatial_dims],
+            output_core_dims=[AXIS.spatial_dims],
+            vectorize=True,
+            dask="parallelized",
+        )

src/cala/nodes/detect/slice_nmf.py

@@ -19,6 +19,9 @@ class SliceNMF(Node):
     """Wait until this number of frames to begin detecting."""
     detect_thresh: float
     """Minimum detection threshold for brightness fluctuation."""
+    reprod_tol: float
+    """Mean pixel value error tolerance for reproduced slice from the new component"""
+
     nmf_kwargs: dict[str, Any] = Field(default_factory=dict)
 
     error_: float = Field(None)
@@ -43,7 +46,7 @@ def process(
         fps = []
         trs = []
 
-        while np.sqrt(energy.max()).item() > self.detect_thresh:  # or use res directly
+        while energy.max().item() >= self.detect_thresh:  # or use res directly
             # Find and analyze neighborhood of maximum variance
             slice_ = self._get_max_energy_slice(
                 arr=res, energy_landscape=energy, radius=detect_radius
@@ -55,25 +58,25 @@ def process(
             )
 
             l1_norm = slice_.sum().item()
-            # l0_norm = np.prod(slice_.shape)  # this fails when the residuals are tiny
             comp_recon = a_new @ c_new
 
             energy.loc[{ax: slice_.coords[ax] for ax in AXIS.spatial_dims}] = 0
 
-            if (self.error_ / l1_norm) <= self._model.tol:
+            if (self.error_ / l1_norm) <= self.reprod_tol:
                 fps.append(Footprint.from_array(a_new))
                 trs.append(Trace.from_array(c_new))
                 res = (res - comp_recon).clip(0)
             else:
-                res.loc[{ax: slice_.coords[ax] for ax in AXIS.spatial_dims}] = 0
+                l0_norm = np.prod(slice_.shape)
+                res.loc[{ax: slice_.coords[ax] for ax in AXIS.spatial_dims}] = self.error_ / l0_norm
 
         return fps, trs
 
     def _get_energy(self, res: xr.DataArray) -> xr.DataArray:
         pixels_median = res.median(dim=AXIS.frames_dim)
         V = res - pixels_median
 
-        return (V**2).sum(dim=AXIS.frames_dim)
+        return np.sqrt((V**2).mean(dim=AXIS.frames_dim))
 
     def _get_max_energy_slice(
         self,

src/cala/nodes/footprints.py

@@ -6,10 +6,12 @@
 from noob import Name, process_method
 
 from cala.assets import CompStats, Footprints, PixStats
+from cala.logging import init_logger
 from cala.models import AXIS
 
 
 class Footprinter:
+    logger = init_logger(__name__)
 
     def __init__(self, tol: float, max_iter: int | None = None, bep: int | None = None):
         self.bep = bep
@@ -55,7 +57,7 @@ def ingest_frame(
             kernel = self._expansion_kernel()
             mask = self._expand_boundary(kernel, mask)
 
-        # for _ in range(self.max_iter):
+        cnt = 0
         while True:
             AM = A.rename(AXIS.component_rename) @ M
             numerator = W - AM
@@ -72,8 +74,15 @@ def ingest_frame(
             update = numerator / (M_diag + np.finfo(float).tiny)
             A_new = (mask * (A + update)).clip(min=0)
 
-            if (np.abs(A - A_new).sum() / np.prod(A.shape)).item() < self.tol:
+            step = (np.abs(A - A_new).sum() / np.prod(A.shape)).item()
+
+            cnt += 1
+            maxed = self.max_iter and (cnt == self.max_iter)
+
+            if step < self.tol or maxed:
                 footprints.array = A_new
+                if maxed:
+                    self.logger.debug(msg="max_iter reached before converging.")
                 return footprints
             else:
                 A = A_new
@@ -93,7 +102,9 @@ def _expand_boundary(self, kernel: np.ndarray, mask: xr.DataArray) -> xr.DataArr
         )
 
 
-def ingest_component(footprints: Footprints, new_footprints: Footprints) -> Footprints:
+def ingest_component(
+    footprints: Footprints, new_footprints: Footprints
+) -> A[Footprints, Name("footprints")]:
     if new_footprints.array is None:
         return footprints