Add support for polydisperse features

benchmarks/polydisperse_aniso_grid.py

@@ -0,0 +1,158 @@
+"""Benchmark: anisotropic (confocal) polydisperse ``locate`` vs. the naive
+``diameter=max`` baseline, in 3-D.
+
+Mirrors ``polydisperse_grid.py`` but in a confocal geometry: x and y equal, z
+one third as thick (``aspect=(1/3, 1, 1)``; trackpy axis order z, y, x). Runs a
+grid over size-range x density and reports, per cell, each method's recall
+(``r``), RMS position error over matched features (``e``) and precision (``p``,
+the fraction of detections that are true -- the rest are false positives), plus
+grid totals and the per-method wall-clock time.
+
+``characterize=False``: the uncertainty (ep) step is skipped. In 3-D it is
+dominated by ``measure_noise``'s large-radius morphological dilation -- a cost
+shared equally by both methods that would otherwise swamp the timing -- so
+excluding it isolates the detection + refinement work that the poly algorithm
+actually changes.
+
+Run:  python benchmarks/polydisperse_aniso_grid.py
+"""
+import time
+
+import numpy as np
+from scipy.spatial import cKDTree
+
+import trackpy as tp
+from trackpy import Polydisperse
+from trackpy.artificial import draw_feature
+from trackpy.utils import default_pos_columns
+
+MIN_DIAMETER = 5
+SHAPE = (48, 256, 256)  # z, y, x -- thin in z (larger than the unit test)
+ASPECT = (1 / 3, 1, 1)  # z one third of x/y
+MAX_DIAMETER_BY_RANGE = {'x5': 25, 'x10': 51}  # size-range label -> max x/y diameter
+GAP_BY_DENSITY = {'low': 20, 'med': 9, 'high': 3}  # density label -> edge gap (px); smaller = denser
+NOISE_STD = 12  # high background noise
+POS_COLUMNS = default_pos_columns(3)
+MINMASS = 50  # above the noise floor, below any real feature
+MATCH_TOLERANCE = 1.5  # max px between a detection and its true position
+TARGET = 150  # features per frame (packing target)
+
+
+def render(shape, centers, sizes, noise_std, seed):
+    """Render axis-aligned ellipsoidal Gaussians of given (center, reference Rg).
+
+    Each feature's per-axis radius of gyration is ``size * aspect[i]``.
+    """
+    image = np.zeros(shape, dtype=np.uint8)
+    for center, size in zip(centers, sizes):
+        draw_feature(image, center, tuple(size * a for a in ASPECT),
+                     max_value=200)
+    if noise_std:
+        rng = np.random.default_rng(seed)
+        image = np.clip(image.astype(float) + rng.normal(0, noise_std, shape),
+                        0, 255).astype(np.uint8)
+    return image
+
+
+def pack(max_diameter, gap, seed, target=TARGET, max_attempts=80000):
+    """Random non-overlapping ellipsoids spanning [~5, max_diameter] x/y diameters.
+
+    Rejection is done in aspect-normalised coordinates (where the ellipsoids are
+    spheres of the reference visible radius), so features never overlap; the
+    per-axis placement margin follows ``aspect`` (a thin z axis gets a thin
+    margin). Returns (centers, sizes) with size = reference radius of gyration.
+    """
+    rng = np.random.default_rng(seed)
+    max_rg = max_diameter / 3.5
+    asp = np.array(ASPECT, dtype=float)
+    margin = [int(max_diameter * a) // 2 + 3 for a in ASPECT]
+    centers = np.empty((0, len(SHAPE)))
+    visible_radii = np.empty(0)
+    sizes = []
+    for _ in range(max_attempts):
+        if len(sizes) >= target:
+            break
+        radius_of_gyration = rng.uniform(1.4, max_rg)
+        visible_radius = 2.0 * radius_of_gyration
+        candidate = np.array([rng.uniform(m, s - m)
+                              for s, m in zip(SHAPE, margin)])
+        normalized = np.sqrt(np.sum(((centers - candidate) / asp) ** 2, axis=1))
+        if len(centers) == 0 or np.all(
+                normalized >= visible_radii + visible_radius + gap):
+            centers = np.vstack([centers, candidate])
+            visible_radii = np.append(visible_radii, visible_radius)
+            sizes.append(radius_of_gyration)
+    return centers, np.array(sizes)
+
+
+def recall_and_error(features, true_centers):
+    """Recall (fraction of true features matched within tolerance), the RMS
+    position error over the matched features, and the number matched."""
+    if len(features) == 0:
+        return 0.0, float('nan'), 0
+    distance_to_detection, _ = cKDTree(
+        features[POS_COLUMNS].values).query(true_centers)
+    matched = distance_to_detection < MATCH_TOLERANCE
+    n_matched = int(matched.sum())
+    recall = float(np.mean(matched))
+    error = (np.sqrt(np.mean(distance_to_detection[matched] ** 2))
+             if matched.any() else float('nan'))
+    return recall, error, n_matched
+
+
+def main():
+    baseline_seconds = poly_seconds = 0.0
+    baseline_correct = poly_correct = total_true = 0
+    print("confocal aspect=%s  shape=%s  noise_std=%d" % (
+        ASPECT, SHAPE, NOISE_STD))
+    print("%-5s %-5s %4s | %-20s %-20s"
+          % ("range", "dens", "N", "baseline", "poly"))
+    print("-" * 70)
+    for range_i, (range_label, max_diameter) in enumerate(
+            MAX_DIAMETER_BY_RANGE.items()):
+        max_tuple = Polydisperse(MIN_DIAMETER, max_diameter,
+                                 aspect=ASPECT).for_ndim(3).max_diameter
+        poly = Polydisperse(MIN_DIAMETER, max_diameter, aspect=ASPECT)
+        for density_i, (density_label, gap) in enumerate(
+                GAP_BY_DENSITY.items()):
+            seed = 100 * range_i + 10 * density_i
+            true_centers, sizes = pack(max_diameter, gap, seed)
+            image = render(SHAPE, true_centers, sizes, NOISE_STD, seed)
+
+            start = time.perf_counter()
+            baseline = tp.locate(image, max_tuple, minmass=MINMASS,
+                                 characterize=False)
+            after_baseline = time.perf_counter()
+            found = tp.locate(image, poly, minmass=MINMASS, characterize=False)
+            after_poly = time.perf_counter()
+            baseline_seconds += after_baseline - start
+            poly_seconds += after_poly - after_baseline
+
+            base_recall, base_error, base_n = recall_and_error(
+                baseline, true_centers)
+            poly_recall, poly_error, poly_n = recall_and_error(
+                found, true_centers)
+            baseline_correct += base_n
+            poly_correct += poly_n
+            total_true += len(true_centers)
+            # Precision: fraction of a method's detections that are true (matched
+            # a distinct feature); the rest are false positives.
+            base_prec = base_n / len(baseline) if len(baseline) else float('nan')
+            poly_prec = poly_n / len(found) if len(found) else float('nan')
+            print("%-5s %-5s %4d | r%.2f e%.3f p%.2f  r%.2f e%.3f p%.2f"
+                  % (range_label, density_label, len(true_centers),
+                     base_recall, base_error, base_prec,
+                     poly_recall, poly_error, poly_prec))
+
+    print("-" * 70)
+    print("Total correct (within %.1fpx) of %d true:  baseline=%d  poly=%d  "
+          "(poly/baseline=%.2fx)"
+          % (MATCH_TOLERANCE, total_true, baseline_correct, poly_correct,
+             poly_correct / baseline_correct))
+    print("Total detect+refine time:  baseline=%.2fs  poly=%.2fs  "
+          "(poly/baseline=%.2fx)"
+          % (baseline_seconds, poly_seconds, poly_seconds / baseline_seconds))
+
+
+if __name__ == '__main__':
+    main()

benchmarks/polydisperse_grid.py

@@ -0,0 +1,142 @@
+"""Benchmark: polydisperse ``locate`` vs. the naive ``diameter=max`` baseline.
+
+Runs a grid of synthetic images over size range x density x noise and reports,
+per cell, each method's recall (``r``), RMS position error over matched features
+(``e``) and precision (``p``, the fraction of detections that are true -- the
+rest are false positives), plus grid totals and wall-clock time. SPIFF
+(``spiff='auto'``) is applied to both methods where a frame has enough features,
+and silently skipped otherwise.
+
+Run:  python benchmarks/polydisperse_grid.py
+"""
+import time
+
+import numpy as np
+from scipy.spatial import cKDTree
+
+import trackpy as tp
+from trackpy import Polydisperse
+from trackpy.artificial import draw_feature
+from trackpy.utils import default_pos_columns
+
+MIN_DIAMETER = 5
+SHAPE = (500, 500)
+MAX_DIAMETER_BY_RANGE = {'x3': 15, 'x5': 25, 'x10': 51}  # size-range label -> max diameter
+GAP_BY_DENSITY = {'low': 24, 'med': 10, 'high': 3}  # density label -> edge gap (px); smaller = denser
+STD_BY_NOISE = {'low': 1, 'med': 5, 'high': 12}  # noise label -> background std
+POS_COLUMNS = default_pos_columns(2)
+MINMASS = 500
+MATCH_TOLERANCE = 1.5  # max px between a detection and its true position
+
+
+def render(shape, centers, sizes, noise_std, seed):
+    """Render Gaussian features of given (center, radius-of-gyration)."""
+    image = np.zeros(shape, dtype=np.uint8)
+    for center, size in zip(centers, sizes):
+        draw_feature(image, center, size, max_value=200)
+    if noise_std:
+        rng = np.random.default_rng(seed)
+        image = np.clip(image.astype(float) + rng.normal(0, noise_std, shape),
+                        0, 255).astype(np.uint8)
+    return image
+
+
+def pack(max_diameter, gap, seed, target=300, max_attempts=15000):
+    """Random non-overlapping features spanning [~5, max_diameter] diameters.
+
+    Rejects a candidate whose centre is closer than the sum of the two features'
+    visible radii plus ``gap`` to any accepted feature, so a smaller ``gap``
+    packs more densely. Returns (centers, sizes) where size = radius of gyration.
+    """
+    rng = np.random.default_rng(seed)
+    max_rg = max_diameter / 3.5
+    centers = np.empty((0, 2))
+    radii = np.empty(0)
+    sizes = []
+    low, high = 28, SHAPE[0] - 28
+    for _ in range(max_attempts):
+        if len(sizes) >= target:
+            break
+        radius_of_gyration = rng.uniform(1.4, max_rg)
+        visible_radius = 2.0 * radius_of_gyration
+        candidate = rng.uniform(low, high, 2)
+        distances = np.hypot(*(centers - candidate).T)  # to every accepted feature
+        if len(centers) == 0 or np.all(distances >= radii + visible_radius + gap):
+            centers = np.vstack([centers, candidate])
+            radii = np.append(radii, visible_radius)
+            sizes.append(radius_of_gyration)
+    return centers, np.array(sizes)
+
+
+def recall_and_error(features, true_centers):
+    """Recall (fraction of true features matched within tolerance), the RMS
+    position error over the matched features, and the number matched."""
+    if len(features) == 0:
+        return 0.0, float('nan'), 0
+    distance_to_detection, _ = cKDTree(features[POS_COLUMNS].values).query(true_centers)
+    matched = distance_to_detection < MATCH_TOLERANCE
+    n_matched = int(matched.sum())
+    recall = float(np.mean(matched))
+    error = (np.sqrt(np.mean(distance_to_detection[matched] ** 2))
+             if matched.any() else float('nan'))
+    return recall, error, n_matched
+
+
+def main():
+    # Warm up the numba JIT so it doesn't skew the first timed cell.
+    warmup_centers, warmup_sizes = pack(25, 10, 0, target=20)
+    warmup_image = render(SHAPE, warmup_centers, warmup_sizes, 1, 0)
+    tp.locate(warmup_image, 25, minmass=MINMASS, spiff='auto', characterize=False)
+    tp.locate(warmup_image, Polydisperse(MIN_DIAMETER, 25), minmass=MINMASS,
+              spiff='auto', characterize=False)
+
+    baseline_seconds = poly_seconds = 0.0
+    baseline_correct = poly_correct = total_true = 0
+    print("%-5s %-5s %-5s %4s | %-20s %-20s"
+          % ("range", "dens", "noise", "N", "baseline", "poly"))
+    print("-" * 76)
+    for range_i, (range_label, max_diameter) in enumerate(MAX_DIAMETER_BY_RANGE.items()):
+        for density_i, (density_label, gap) in enumerate(GAP_BY_DENSITY.items()):
+            for noise_i, (noise_label, noise_std) in enumerate(STD_BY_NOISE.items()):
+                seed = 100 * range_i + 10 * density_i + noise_i
+                true_centers, sizes = pack(max_diameter, gap, seed)
+                image = render(SHAPE, true_centers, sizes, noise_std, seed)
+
+                start = time.perf_counter()
+                baseline = tp.locate(image, max_diameter, minmass=MINMASS,
+                                     spiff='auto', characterize=False)
+                after_baseline = time.perf_counter()
+                poly = tp.locate(image, Polydisperse(MIN_DIAMETER, max_diameter),
+                                 minmass=MINMASS, spiff='auto',
+                                 characterize=False)
+                after_poly = time.perf_counter()
+                baseline_seconds += after_baseline - start
+                poly_seconds += after_poly - after_baseline
+
+                baseline_recall, baseline_error, baseline_n = recall_and_error(
+                    baseline, true_centers)
+                poly_recall, poly_error, poly_n = recall_and_error(
+                    poly, true_centers)
+                baseline_correct += baseline_n
+                poly_correct += poly_n
+                total_true += len(true_centers)
+                # Precision: fraction of a method's detections that are true
+                # (matched a distinct feature); the rest are false positives.
+                baseline_prec = baseline_n / len(baseline) if len(baseline) else float('nan')
+                poly_prec = poly_n / len(poly) if len(poly) else float('nan')
+                print("%-5s %-5s %-5s %4d | r%.2f e%.3f p%.2f  r%.2f e%.3f p%.2f"
+                      % (range_label, density_label, noise_label, len(true_centers),
+                         baseline_recall, baseline_error, baseline_prec,
+                         poly_recall, poly_error, poly_prec))
+
+    print("-" * 76)
+    print("Total correct (within %.1fpx) of %d true:  baseline=%d  poly=%d  "
+          "(poly/baseline=%.2fx)"
+          % (MATCH_TOLERANCE, total_true, baseline_correct, poly_correct,
+             poly_correct / baseline_correct))
+    print("Total grid time:  baseline=%.2fs  poly=%.2fs  (poly/baseline=%.2fx)"
+          % (baseline_seconds, poly_seconds, poly_seconds / baseline_seconds))
+
+
+if __name__ == '__main__':
+    main()

trackpy/api.py

@@ -16,6 +16,7 @@
 from .filtering import filter_stubs, filter_clusters, filter
 from .feature import locate, batch, local_maxima, \
            estimate_mass, estimate_size, minmass_v03_change, minmass_v04_change
+from .polydisperse import Polydisperse
 from .preprocessing import bandpass, invert_image
 from .framewise_data import FramewiseData, PandasHDFStore, PandasHDFStoreBig, \
            PandasHDFStoreSingleNode