+
{tabs.map(t => {
const isLocked = locked?.(t) ?? false
return (
!isLocked && onSelect(t)}>{t}
)
diff --git a/electron/src/renderer/src/kernel/SpyDEContext.tsx b/electron/src/renderer/src/kernel/SpyDEContext.tsx
index 32943c28..aec74434 100644
--- a/electron/src/renderer/src/kernel/SpyDEContext.tsx
+++ b/electron/src/renderer/src/kernel/SpyDEContext.tsx
@@ -771,6 +771,12 @@ interface SpyDEContextValue {
gpuHelpDialogOpen: boolean
openGpuHelpDialog: () => void
closeGpuHelpDialog: () => void
+ // Bottom table dock (BottomDock.tsx). Renderer-only UI state, but it lives
+ // HERE rather than in App because MenuBar's View menu and the StatusBar
+ // toggle both need it and only see the context.
+ tableDockOpen: boolean
+ openTableDock: () => void
+ closeTableDock: () => void
// MDIArea registers its tile-all-windows function here so StatusBar's
// "Tile" button can trigger it without threading window-layout state (which
// lives in MDIArea's local refs) through the shared context.
@@ -856,6 +862,7 @@ export function SpyDEProvider({ children }: { children: React.ReactNode }) {
const [updateDialogOpen, setUpdateDialogOpen] = useState(false)
const [gpuStatusDialogOpen, setGpuStatusDialogOpen] = useState(false)
const [gpuHelpDialogOpen, setGpuHelpDialogOpen] = useState(false)
+ const [tableDockOpen, setTableDockOpen] = useState(false)
const [dragKind, setDragKind] = useState<'window' | null>(null)
// Post every stored state for a figure to its iframe (called on iframe load).
@@ -1469,6 +1476,29 @@ export function SpyDEProvider({ children }: { children: React.ReactNode }) {
// (app-global, not wizard-scoped, but the same re-broadcast fits).
case 'download_progress':
case 'download_done':
+ // The particle/track table + event stream for one result window —
+ // consumed by BottomDock (see ParticlesTableMessage in protocol.ts).
+ // A panel-local payload, so it gets a CustomEvent rather than reducer
+ // state, exactly like layers_state.
+ case 'particles_table':
+ // Segment Particles caret (spyde/actions/particles_action.py) —
+ // `seg_state` is the authoritative caret state (classes + per-class
+ // labelled-pixel counts, effective params), `seg_preview` one frame's
+ // result (count, size histogram, the EFFECTIVE min_size), `seg_trained`
+ // the scribble classifier's fit report. Consumed by SegmentWizard.
+ case 'seg_state':
+ case 'seg_preview':
+ case 'seg_trained':
+ // Drift Correction caret (spyde/actions/drift_action.py) — caret state,
+ // the ROI discovery preview (~20 frames aligned on the box, with its
+ // sharpening gain), whole-movie solve progress, the streamed dy/dx
+ // batches, and the solved model. Consumed by DriftWizard; the dy/dx
+ // curve itself is painted by the backend into its own figure window.
+ case 'drift_state':
+ case 'drift_preview':
+ case 'drift_trace':
+ case 'drift_progress':
+ case 'drift_result':
// Cluster telemetry — consumed by the StatusBar DaskMonitor HUD.
case 'dask_stats':
// Read-throughput readout — consumed by the StatusBar IoThroughput HUD.
@@ -1745,6 +1775,8 @@ export function SpyDEProvider({ children }: { children: React.ReactNode }) {
const closeGpuStatusDialog = () => setGpuStatusDialogOpen(false)
const openGpuHelpDialog = () => setGpuHelpDialogOpen(true)
const closeGpuHelpDialog = () => setGpuHelpDialogOpen(false)
+ const openTableDock = () => setTableDockOpen(true)
+ const closeTableDock = () => setTableDockOpen(false)
return (
{children}
diff --git a/electron/src/renderer/src/kernel/protocol.ts b/electron/src/renderer/src/kernel/protocol.ts
index cf0d2900..e2795a75 100644
--- a/electron/src/renderer/src/kernel/protocol.ts
+++ b/electron/src/renderer/src/kernel/protocol.ts
@@ -811,6 +811,231 @@ export interface IoThroughputMessage extends MsgBase {
color: 'green' | 'yellow' | 'red'
}
+// ── Particle table dock ─────────────────────────────────────────────────────
+
+/** One column of the particle table, as the backend describes it. Mirrors the
+ * renderer-side `DataColumn` (components/DataTable.tsx) minus the function
+ * fields, which cannot cross IPC. */
+export interface ParticlesTableColumn {
+ /** Property read from each row — e.g. a name from
+ * `spyde.signals.particles.COLUMNS`. */
+ key: string
+ label: string
+ /** Fixed pixel width; omitted → the column flexes. */
+ width?: number
+ align?: 'left' | 'center' | 'right'
+ /** Right-aligned + tabular figures. */
+ numeric?: boolean
+ /** Default true. */
+ sortable?: boolean
+ /** 'swatch' prefixes a colour chip keyed on the cell value (track id). */
+ kind?: 'text' | 'swatch'
+ /** Decimal places for a numeric cell. */
+ precision?: number
+ /** Appended to the value ("nm", "nm²"). */
+ units?: string
+}
+
+/** One row of the Events tab — exactly `ParticleEvent.to_dict()` from
+ * `spyde/particles/track.py`. */
+export interface ParticleEventRecord {
+ frame: number
+ /** 'birth' | 'death' | 'merge' | 'split' (track.py EVENT_KINDS). */
+ kind: string
+ tracks: number[]
+ particles: number[]
+}
+
+/**
+ * The particle/track table for one result window, in reply to a
+ * `particles_query {window_id}` action (and pushed again when a segmentation
+ * batch finalises). Re-broadcast as a `spyde:particles_table` CustomEvent and
+ * consumed by BottomDock — the renderer holds no reducer state for it.
+ *
+ * `columns` is backend-supplied so the dock stays data-agnostic: sending a
+ * different column set is the whole mechanism for showing tracks instead of
+ * particles. `events` is fixed-shape and drives the Events tab.
+ */
+export interface ParticlesTableMessage extends MsgBase {
+ type: 'particles_table'
+ /** The window this table belongs to; `null` = unscoped (always shown). */
+ window_id: number | null
+ /** Dock header text. Defaults to "Particles" when absent. */
+ title?: string
+ columns: ParticlesTableColumn[]
+ /** One object per row, keyed by `columns[].key`. An optional numeric `id`
+ * field is used as the stable selection key when present. */
+ rows: Record[]
+ events?: ParticleEventRecord[]
+ /** Calibrated unit of the length columns (`SpyDEParticles.units`). */
+ units?: string
+ /** True while a batch is still streaming rows in — shows a "streaming…" pill. */
+ partial?: boolean
+}
+
+// ── Segment Particles caret (spyde/actions/particles_action.py, plan B7) ─────
+
+/** One scribble class as the backend reports it — `ScribbleClass.to_dict()`
+ * plus the labelled-pixel count `SegmentWizard.class_report()` adds.
+ *
+ * `pixels` is NOT decoration: under-training a class is the failure mode
+ * plan §B3 calls out, and this count is how the user notices. A class with
+ * zero pixels is present in the list rather than absent from it. */
+export interface SegClassInfo {
+ id: number
+ name: string
+ /** CSS hex — the renderer's units; the backend never converts it. */
+ colour: string
+ /** Counts toward the foreground probability map (several classes may). */
+ particle: boolean
+ /** Marks the SEAM between two touching particles (the ilastik third class).
+ * Painting it lets the backend split by connected components and skip the
+ * distance transform and watershed entirely. Mutually exclusive with
+ * `particle` — the backend raises if a class claims both. */
+ boundary: boolean
+ pixels: number
+}
+
+/** The caret's authoritative state, emitted by `_emit_state` after open,
+ * method switch, paint and train. `params` are the EFFECTIVE (coerced)
+ * values, not what the caret last sent. */
+export interface SegStateMessage extends MsgBase {
+ type: 'seg_state'
+ window_id: number | null
+ /** 'classical' | 'scribble' | 'prompt'. */
+ method: string
+ /** Navigator frame the preview is showing. */
+ frame: number
+ n_frames: number
+ frame_shape: [number, number]
+ classes: SegClassInfo[]
+ /** Frame indices carrying at least one scribble. */
+ labelled_frames: number[]
+ trained: boolean
+ params: Record
+}
+
+/** One frame's preview result. `min_size` is the EFFECTIVE value that ran and
+ * `min_size_floored` says the backend raised what the caret asked for — the
+ * caret must show the effective number (plan §0.9: at min_size=0 the split
+ * returns background speckle as particles). `count` is already post-filter. */
+export interface SegPreviewMessage extends MsgBase {
+ type: 'seg_preview'
+ window_id: number | null
+ frame: number
+ count: number
+ /** Per-instance areas in calibrated units², capped at 2000 entries. */
+ areas: number[]
+ median_area: number
+ /** Signal-axis unit ("nm"); areas are in unit². */
+ units: string
+ method: string
+ min_size: number
+ min_size_floored: boolean
+ elapsed_ms: number
+ /**
+ * `[y0, x0, h, w]` when the frame was too big to preview whole, else null.
+ *
+ * A 4096² frame costs 8.4 s to segment, so the backend previews a centred
+ * megapixel crop at FULL resolution instead (cropping rather than
+ * downsampling, so the preview runs the identical algorithm the real run
+ * will). The caret MUST surface this: otherwise "12 particles on this frame"
+ * is a lie about a frame where only the middle sixteenth was looked at.
+ */
+ preview_box: [number, number, number, number] | null
+}
+
+/** Scribble classifier fit report (`ScribbleClassifier.fit`). */
+export interface SegTrainedMessage extends MsgBase {
+ type: 'seg_trained'
+ window_id: number | null
+ report: {
+ device?: string
+ n_pixels?: number
+ n_channels?: number
+ n_classes?: number
+ labelled_frames?: number[]
+ train_accuracy?: number
+ [k: string]: unknown
+ }
+}
+
+// ── Drift Correction caret (spyde/actions/drift_action.py, plan A8) ──────────
+
+/** Caret state. `window_id` is the SOURCE plot's window (where the caret
+ * lives); `check_window_id` is the bare-figure Drift Check window (whole-movie
+ * before/after sums on top, the ROI discovery pair beneath) and
+ * `trace_window_id` the bare-figure dy/dx window the solve opens. */
+export interface DriftStateMessage extends MsgBase {
+ type: 'drift_state'
+ window_id: number | null
+ check_window_id: number | null
+ trace_window_id: number | null
+ /** 'rigid' | 'rigid_affine' | 'nonrigid' — only rigid has a solver. */
+ method: string
+ solved: boolean
+ /** Whether the FULL solve restricts itself to the alignment box. The preview
+ * always uses the box regardless — the toggle is the commitment. */
+ use_roi: boolean
+ /** The alignment box as `[y0, x0, h, w]` in IMAGE PIXELS (what
+ * `solve_translation(roi=…)` takes), or null when there is no usable box. */
+ roi: number[] | null
+ params: Record
+ /** Present on the open/ready emissions. */
+ n_frames?: number
+}
+
+/** The discovery preview: ~20 frames sampled across the whole movie, aligned on
+ * the box alone. `gain` is the gradient energy of the aligned sum over the raw
+ * sum, measured on the pixels both cover — > 1.5 means a usable landmark,
+ * <= 1 means aligning there changes nothing. */
+export interface DriftPreviewMessage extends MsgBase {
+ type: 'drift_preview'
+ window_id: number | null
+ roi: number[] | null
+ frames: number
+ gain: number
+ max_abs_shift: number
+ params: Record
+}
+
+/** A batch of solved shifts streamed from `solve_translation`'s `on_shift`
+ * while the solve runs — `[frame_index, dy, dx]` each. The backend paints
+ * these into the dy/dx window itself; the message exists so any host can
+ * follow the curve live. */
+export interface DriftTraceMessage extends MsgBase {
+ type: 'drift_trace'
+ window_id: number | null
+ points: [number, number, number][]
+}
+
+/** Whole-movie solve progress (also emitted as a plain `progress` message). */
+export interface DriftProgressMessage extends MsgBase {
+ type: 'drift_progress'
+ window_id: number | null
+ done: number
+ total: number
+}
+
+/** The solved model. `shifts[i]` is the correction ADDED to frame i, `[dy, dx]`
+ * in pixels; a cancelled solve leaves NaN rows for frames it never reached. */
+export interface DriftResultMessage extends MsgBase {
+ type: 'drift_result'
+ window_id: number | null
+ shifts: [number, number][]
+ kind: string
+ reference: string
+ max_abs_shift: number
+ /** The box the solve correlated on (`use_roi`), or null for whole-frame. */
+ roi: number[] | null
+ /** Whole-movie sharpening: gradient energy of the corrected sum over the raw
+ * one. The same number the preview reports, now for what actually ran. */
+ gain: number
+ /** Frames dropped from the running reference by the outlier rejector. */
+ rejected: number
+ cancelled: boolean
+}
+
/**
* Wizard-scoped events re-broadcast verbatim as DOM CustomEvents (the caret
* components subscribe directly). The payload beyond `type` is consumer-defined,
@@ -936,6 +1161,15 @@ export type PlotAppMessage =
| DownloadDoneMessage
| DaskStatsMessage
| IoThroughputMessage
+ | ParticlesTableMessage
+ | SegStateMessage
+ | SegPreviewMessage
+ | SegTrainedMessage
+ | DriftStateMessage
+ | DriftPreviewMessage
+ | DriftTraceMessage
+ | DriftProgressMessage
+ | DriftResultMessage
/**
* Narrow a raw incoming message (`Record` from the IPC bridge)
diff --git a/electron/tests/data_table.spec.ts b/electron/tests/data_table.spec.ts
new file mode 100644
index 00000000..5ecfcc5e
--- /dev/null
+++ b/electron/tests/data_table.spec.ts
@@ -0,0 +1,343 @@
+/**
+ * data_table.spec.ts — the bottom table dock + the generic DataTable.
+ *
+ * Drives everything through injected `particles_table` messages
+ * (`window._spyde_test_inject`, the same hook app_log.spec.ts uses), so the spec
+ * needs NO Python particle backend — that side is a separate workstream.
+ *
+ * What it pins:
+ * - visibility wiring: StatusBar toggle, View menu item (with its ✓), close ×
+ * - the empty state when no data has arrived (the degrade-gracefully contract)
+ * - VIRTUALISATION: 5000 rows, only a windowed slice in the DOM
+ * - sort cycle asc → desc → none, single + ctrl multi selection
+ * - the Events tab's fixed columns + per-kind colour swatch
+ * - top-edge resize, and the 50%-of-window cap that protects the MDI area
+ *
+ * Screenshots land in electron/data_table_shots/ and are meant to be LOOKED at.
+ */
+import { test, expect, Page, ElectronApplication } from '@playwright/test'
+import { join } from 'path'
+const { launchApp } = require('./_harness.cjs')
+
+const SHOTS = join(__dirname, '..', 'data_table_shots')
+
+let app: ElectronApplication
+let page: Page
+let assertNoJsErrors: () => void
+
+test.beforeAll(async () => {
+ const ctx = await launchApp()
+ app = ctx.app
+ page = ctx.page
+ assertNoJsErrors = ctx.assertNoJsErrors
+})
+test.afterAll(async () => { await app?.close() })
+
+async function inject(msg: Record) {
+ await page.evaluate((m) => {
+ ;(window as unknown as { _spyde_test_inject?: (x: unknown) => void })._spyde_test_inject?.(m)
+ }, msg)
+}
+
+/** A realistic particle table: backend-shaped column descriptors + N rows keyed
+ * by `spyde.signals.particles.COLUMNS` / MEASURED_COLUMNS.
+ *
+ * `wide` repeats the measured block so the fixed widths overflow the dock —
+ * that is the horizontal-scroll + sticky-header case. */
+async function injectTable(n: number, opts: { partial?: boolean; wide?: boolean } = {}) {
+ await page.evaluate(({ count, partial, wide }) => {
+ const columns: Record[] = [
+ { key: 'id', label: '#', width: 54, numeric: true },
+ { key: 't', label: 'frame', width: 60, numeric: true },
+ { key: 'track_id', label: 'track', width: 72, numeric: true, kind: 'swatch' },
+ { key: 'y', label: 'y', width: 76, numeric: true, precision: 2, units: 'nm' },
+ { key: 'x', label: 'x', width: 76, numeric: true, precision: 2, units: 'nm' },
+ { key: 'area', label: 'area', width: 88, numeric: true, precision: 1, units: 'nm²' },
+ { key: 'equiv_diameter', label: 'diameter', width: 84, numeric: true, precision: 2, units: 'nm' },
+ { key: 'major_axis', label: 'major', width: 80, numeric: true, precision: 2, units: 'nm' },
+ { key: 'minor_axis', label: 'minor', width: 80, numeric: true, precision: 2, units: 'nm' },
+ { key: 'perimeter', label: 'perimeter', width: 88, numeric: true, precision: 1, units: 'nm' },
+ { key: 'circularity', label: 'circularity', width: 92, numeric: true, precision: 3 },
+ { key: 'eccentricity', label: 'eccentricity', width: 96, numeric: true, precision: 3 },
+ { key: 'solidity', label: 'solidity', width: 82, numeric: true, precision: 3 },
+ { key: 'intensity_mean', label: 'mean I', width: 80, numeric: true, precision: 1 },
+ { key: 'intensity_std', label: 'std I', width: 76, numeric: true, precision: 1 },
+ { key: 'background', label: 'bkgd', width: 72, numeric: true, precision: 1 },
+ // Width-less: absorbs whatever space the fixed columns leave over.
+ { key: 'intensity_max', label: 'max I', numeric: true, precision: 1 },
+ ]
+ if (wide) {
+ for (const key of ['area', 'perimeter', 'circularity', 'solidity',
+ 'intensity_mean', 'intensity_std', 'background', 'major_axis']) {
+ columns.push({ key, label: `${key} (2)`, width: 110, numeric: true, precision: 2 })
+ }
+ }
+ const rows: Record[] = []
+ for (let i = 0; i < count; i++) {
+ // Deterministic pseudo-random so the screenshot is stable run to run.
+ const r = (k: number) => ((Math.sin(i * 12.9898 + k * 78.233) * 43758.5453) % 1 + 1) % 1
+ const d = 4 + r(1) * 26
+ rows.push({
+ id: i,
+ t: Math.floor(i / 24),
+ track_id: i % 37,
+ y: r(2) * 512,
+ x: r(3) * 512,
+ area: Math.PI * (d / 2) ** 2,
+ equiv_diameter: d,
+ major_axis: d * (1.0 + r(8) * 0.6),
+ minor_axis: d * (0.6 + r(9) * 0.3),
+ perimeter: Math.PI * d * (1 + r(10) * 0.15),
+ circularity: 0.55 + r(4) * 0.45,
+ eccentricity: r(5) * 0.9,
+ solidity: 0.8 + r(11) * 0.2,
+ intensity_mean: 1200 + r(6) * 2400,
+ intensity_std: 60 + r(12) * 340,
+ background: 300 + r(13) * 500,
+ intensity_max: 3000 + r(7) * 5000,
+ })
+ }
+ const events: Record[] = []
+ const kinds = ['birth', 'death', 'merge', 'split']
+ for (let i = 0; i < 40; i++) {
+ const kind = kinds[i % 4]
+ events.push({
+ frame: 3 + i * 2,
+ kind,
+ tracks: kind === 'merge' || kind === 'split' ? [i % 37, (i + 5) % 37] : [i % 37],
+ particles: kind === 'merge' || kind === 'split' ? [i * 7, i * 7 + 3] : [i * 7],
+ })
+ }
+ ;(window as unknown as { _spyde_test_inject?: (x: unknown) => void })._spyde_test_inject?.({
+ type: 'particles_table',
+ window_id: null,
+ title: 'Particles — au_nanoparticles (in-situ)',
+ units: 'nm',
+ partial,
+ columns,
+ rows,
+ events,
+ })
+ }, { count: n, partial: Boolean(opts.partial), wide: Boolean(opts.wide) })
+ await expect(page.locator('[data-testid="bottom-dock-count"]')).toHaveText(String(n))
+}
+
+const dock = () => page.locator('[data-testid="bottom-dock"]')
+const rows = () => page.locator('[data-testid="particle-table-row"]')
+
+test('the dock is hidden until a toggle opens it, and shows a clear empty state', async () => {
+ await expect(dock()).toHaveCount(0)
+
+ await page.click('[data-testid="toggle-table-dock"]')
+ await expect(dock()).toBeVisible()
+ // No backend data has arrived → the degrade-gracefully message, not a blank.
+ await expect(page.locator('[data-testid="particle-table-empty"]'))
+ .toContainText('Segment Particles')
+ await expect(page.locator('[data-testid="bottom-dock-count"]')).toHaveText('0')
+
+ await page.screenshot({ path: join(SHOTS, '01-empty.png') })
+})
+
+test('5000 rows render as a windowed slice, not 5000 DOM nodes', async () => {
+ await injectTable(5000)
+
+ // Virtualisation: a ~230 px body at 24 px/row is ~10 visible rows; with the
+ // default overscan of 8 either side the slice is well under 50.
+ const rendered = await rows().count()
+ expect(rendered).toBeGreaterThan(4)
+ expect(rendered).toBeLessThan(50)
+
+ // The first row is row 0 and the spacers give the body the full scroll height.
+ await expect(rows().first()).toHaveAttribute('data-row-index', '0')
+ const scrollH = await page.locator('[data-testid="particle-table-body"]')
+ .evaluate((el) => el.scrollHeight)
+ expect(scrollH).toBeGreaterThan(5000 * 24 * 0.9)
+
+ await page.screenshot({ path: join(SHOTS, '02-table.png') })
+})
+
+test('scrolling swaps the rendered window without growing the DOM', async () => {
+ await injectTable(5000)
+ const body = page.locator('[data-testid="particle-table-body"]')
+ await body.evaluate((el) => { el.scrollTop = 24 * 2000 })
+ // The window follows scrollTop: the first rendered row is near index 2000.
+ await expect
+ .poll(async () => Number(await rows().first().getAttribute('data-row-index')))
+ .toBeGreaterThan(1900)
+ expect(await rows().count()).toBeLessThan(50)
+
+ await body.evaluate((el) => { el.scrollTop = 0 })
+ await expect(rows().first()).toHaveAttribute('data-row-index', '0')
+})
+
+test('a sortable header cycles asc → desc → unsorted', async () => {
+ await injectTable(5000)
+ const th = page.locator('[data-testid="particle-table-th-area"]')
+ // Normalise: sort state survives a re-injection (only a tab switch remounts
+ // the table), so click round the cycle rather than assuming it starts clean.
+ for (let i = 0; i < 3 && (await th.getAttribute('data-sort')) !== 'none'; i++) await th.click()
+ await expect(th).toHaveAttribute('data-sort', 'none')
+
+ const firstAreaText = async () =>
+ (await rows().first().locator('[data-col="area"]').innerText()).trim()
+ const areaNum = async () => parseFloat((await firstAreaText()).replace(/[^\d.]/g, ''))
+
+ await th.click()
+ await expect(th).toHaveAttribute('data-sort', 'asc')
+ const asc = await areaNum()
+
+ await th.click()
+ await expect(th).toHaveAttribute('data-sort', 'desc')
+ const desc = await areaNum()
+ expect(desc).toBeGreaterThan(asc)
+ // Shot in the SORTED state — the accent-coloured header + ▾ caret + the
+ // largest areas on top are what the eye has to confirm.
+ await page.screenshot({ path: join(SHOTS, '03-sorted.png') })
+
+ await th.click()
+ await expect(th).toHaveAttribute('data-sort', 'none')
+ // Unsorted restores the original order (row 0 first).
+ await expect(rows().first()).toHaveAttribute('data-row-index', '0')
+})
+
+test('rows select singly, and ctrl-click extends the selection', async () => {
+ await injectTable(5000)
+ await rows().nth(1).click()
+ await expect(page.locator('[data-testid="particle-table-row"][data-selected="true"]'))
+ .toHaveCount(1)
+
+ await rows().nth(3).click({ modifiers: ['ControlOrMeta'] })
+ await rows().nth(5).click({ modifiers: ['ControlOrMeta'] })
+ await expect(page.locator('[data-testid="particle-table-row"][data-selected="true"]'))
+ .toHaveCount(3)
+
+ // A plain click collapses back to one.
+ await rows().nth(7).click()
+ await expect(page.locator('[data-testid="particle-table-row"][data-selected="true"]'))
+ .toHaveCount(1)
+
+ await page.screenshot({ path: join(SHOTS, '04-selection.png') })
+})
+
+test('the filter box narrows the rows and the count shows both totals', async () => {
+ await injectTable(5000)
+ await page.fill('[data-testid="bottom-dock-search"]', '0.9')
+ await expect(page.locator('[data-testid="bottom-dock-count"]')).toContainText('/ 5000')
+ const shown = Number((await page.locator('[data-testid="bottom-dock-count"]').innerText())
+ .split('/')[0].trim())
+ expect(shown).toBeGreaterThan(0)
+ expect(shown).toBeLessThan(5000)
+ await page.fill('[data-testid="bottom-dock-search"]', '')
+ await expect(page.locator('[data-testid="bottom-dock-count"]')).toHaveText('5000')
+})
+
+test('a column set wider than the dock scrolls sideways under a pinned header', async () => {
+ await injectTable(400, { wide: true })
+ const body = page.locator('[data-testid="particle-table-body"]')
+ const { sw, cw } = await body.evaluate((el) => ({ sw: el.scrollWidth, cw: el.clientWidth }))
+ expect(sw).toBeGreaterThan(cw)
+
+ // Scroll right: the header must travel WITH the cells (a sticky header pins
+ // vertically only) and stay pinned to the top after a vertical scroll.
+ await body.evaluate((el) => { el.scrollLeft = 600; el.scrollTop = 24 * 100 })
+ const { headTop, bodyTop } = await page.evaluate(() => {
+ const b = document.querySelector('[data-testid="particle-table-body"]')!.getBoundingClientRect()
+ const h = document.querySelector('[data-testid="particle-table-head"]')!.getBoundingClientRect()
+ return { headTop: h.top, bodyTop: b.top }
+ })
+ expect(Math.abs(headTop - bodyTop)).toBeLessThan(2)
+ await page.screenshot({ path: join(SHOTS, '05-wide-scrolled.png') })
+
+ await body.evaluate((el) => { el.scrollLeft = 0; el.scrollTop = 0 })
+})
+
+test('the Events tab shows the fixed event columns with a colour per kind', async () => {
+ await injectTable(5000)
+ await page.click('[data-testid="bottom-dock-tab-events"]')
+ await expect(dock()).toHaveAttribute('data-tab', 'events')
+ await expect(page.locator('[data-testid="bottom-dock-count"]')).toHaveText('40')
+ await expect(page.locator('[data-testid="particle-table-th-kind"]')).toContainText('event')
+ // The swatch is what makes birth/death/merge/split readable at a glance.
+ expect(await page.locator('[data-testid="particle-table-swatch"]').count())
+ .toBeGreaterThan(4)
+ await expect(rows().first().locator('[data-col="kind"]')).toContainText('birth')
+
+ await page.screenshot({ path: join(SHOTS, '05-events.png') })
+
+ await page.click('[data-testid="bottom-dock-tab-table"]')
+ await expect(dock()).toHaveAttribute('data-tab', 'table')
+})
+
+test('the streaming pill appears while a batch is still filling the table', async () => {
+ await injectTable(200, { partial: true })
+ await expect(page.locator('[data-testid="bottom-dock-streaming"]')).toBeVisible()
+ await injectTable(5000)
+ await expect(page.locator('[data-testid="bottom-dock-streaming"]')).toHaveCount(0)
+})
+
+test('the top edge resizes the dock and never exceeds half the window', async () => {
+ const before = Number(await dock().getAttribute('data-height'))
+ const box = await dock().boundingBox()
+ if (!box) throw new Error('dock has no bounding box')
+
+ // Drag the top edge UP by 120 px → the dock grows.
+ await page.mouse.move(box.x + box.width / 2, box.y + 1)
+ await page.mouse.down()
+ await page.mouse.move(box.x + box.width / 2, box.y - 120, { steps: 8 })
+ await page.mouse.up()
+ const grown = Number(await dock().getAttribute('data-height'))
+ expect(grown).toBeGreaterThan(before)
+
+ await page.screenshot({ path: join(SHOTS, '06-resized.png') })
+
+ // Now drag far past the ceiling: the dock must cap at 50% of the window and
+ // the MDI area must still have real height (the flexShrink:0 starvation trap).
+ const box2 = await dock().boundingBox()
+ if (!box2) throw new Error('dock has no bounding box')
+ await page.mouse.move(box2.x + box2.width / 2, box2.y + 1)
+ await page.mouse.down()
+ await page.mouse.move(box2.x + box2.width / 2, 0, { steps: 10 })
+ await page.mouse.up()
+
+ const { height, inner, mdi } = await page.evaluate(() => ({
+ height: Number(document.querySelector('[data-testid="bottom-dock"]')
+ ?.getAttribute('data-height')),
+ inner: window.innerHeight,
+ mdi: document.querySelector('[data-testid="mdi-area"]')?.getBoundingClientRect().height ?? 0,
+ }))
+ expect(height).toBeLessThanOrEqual(Math.round(inner * 0.5) + 1)
+ expect(mdi).toBeGreaterThan(100)
+
+ await page.screenshot({ path: join(SHOTS, '07-capped.png') })
+})
+
+test('the View menu toggles the dock and shows its checked state', async () => {
+ await page.click('[data-testid="menu-view"]')
+ const item = page.locator('[data-testid="menu-item-table-dock"]')
+ await expect(item).toContainText('✓') // open right now
+ // Shot WITH the menu open — the ✓ is the thing being verified.
+ await page.screenshot({ path: join(SHOTS, '08-view-menu.png') })
+ await item.click()
+ await expect(dock()).toHaveCount(0)
+
+ await page.click('[data-testid="menu-view"]')
+ await expect(page.locator('[data-testid="menu-item-table-dock"]')).not.toContainText('✓')
+ await page.screenshot({ path: join(SHOTS, '09-view-menu-unchecked.png') })
+ await page.locator('[data-testid="menu-item-table-dock"]').click()
+ await expect(dock()).toBeVisible()
+})
+
+test('the close button hides the dock; reopening keeps the loaded table', async () => {
+ await injectTable(5000)
+ await page.click('[data-testid="bottom-dock-close"]')
+ await expect(dock()).toHaveCount(0)
+
+ await page.click('[data-testid="toggle-table-dock"]')
+ await expect(dock()).toBeVisible()
+ // BottomDock stays MOUNTED while hidden (it returns null after its hooks), so
+ // the fetched table survives a hide/show — no re-query round trip.
+ await expect(page.locator('[data-testid="bottom-dock-count"]')).toHaveText('5000')
+
+ assertNoJsErrors()
+})
+
diff --git a/electron/tests/drift_wizard.spec.ts b/electron/tests/drift_wizard.spec.ts
new file mode 100644
index 00000000..0bc13d50
--- /dev/null
+++ b/electron/tests/drift_wizard.spec.ts
@@ -0,0 +1,212 @@
+/**
+ * drift_wizard.spec.ts — the Drift Correction caret, end-to-end on the bundled
+ * synthetic particle movie (whose per-frame drift is ground truth, stamped into
+ * `metadata.Spyde.synthetic`).
+ *
+ * What this proves that tsc + headless tests cannot:
+ * 1. The caret's default face is SMALL (plan §0.9a) — two toggles, one
+ * readout, one button, and a collapsed Advanced. The count is asserted,
+ * not eyeballed, because "too many options" is exactly what regressed.
+ * 2. The discovery loop is real: a draggable box on the movie, a live
+ * drift-corrected sum of just that box in the Drift Check window, and a
+ * sharpness number that MOVES when the box moves.
+ * 3. The dy/dx curve is its OWN window, opened by the solve and filled while
+ * it runs — not caret furniture.
+ * 4. Apply adds the lazy corrected node.
+ */
+import { test, expect } from '@playwright/test'
+import { mkdirSync } from 'fs'
+const {
+ launchApp, backendAction, waitForSubwindowCount, sigWindow, backendErrorLines,
+} = require('./_harness.cjs')
+
+const SHOTS = 'drift_wizard_shots'
+let ctx: Awaited>
+
+test.describe.configure({ mode: 'serial' })
+test.setTimeout(300_000)
+
+/** Every interactive control the caret shows by default. If this list grows,
+ * §0.9a has been walked back and the review comment applies again. */
+const FACE = ['drift-use-roi', 'drift-reject', 'drift-solve', 'drift-advanced-toggle']
+
+test.beforeAll(async () => {
+ mkdirSync(SHOTS, { recursive: true })
+ ctx = await launchApp({ dask: true, env: { SPYDE_LOG_LEVEL: 'INFO' } })
+ const { page } = ctx
+ await page.waitForTimeout(1500)
+ // Enough frames that the solve takes about a second — a 12-frame movie
+ // finishes before a screenshot can catch the dy/dx window mid-fill, which is
+ // the thing this spec has to see.
+ await backendAction(page, 'load_test_data_particles', { frames: 40 })
+ await waitForSubwindowCount(page, 2, 120_000)
+ await page.waitForTimeout(2000)
+})
+
+test.afterAll(async () => {
+ await ctx?.app?.close()
+})
+
+test('the caret opens small: 2 toggles, 1 button, Advanced collapsed', async () => {
+ const { page } = ctx
+ const sig = sigWindow(page)
+ await sig.getByTestId('subwindow-title').click()
+ await sig.getByTestId('subwindow-titlebar').hover()
+ await sig.getByTestId('action-btn-Drift Correction').click()
+ await expect(page.getByTestId('drift-wizard')).toBeVisible()
+
+ // The verification surface is a WINDOW, not the caret (plan A8): whole-movie
+ // raw/corrected sums on top, the ROI discovery pair beneath.
+ await waitForSubwindowCount(page, 3, 120_000)
+ // …and the discovery preview lands on its own, with no click at all.
+ await expect.poll(
+ async () => await page.getByTestId('drift-roi-readout').getAttribute('data-gain'),
+ { timeout: 120_000, message: 'the discovery preview never reported a gain' },
+ ).toBeTruthy()
+ await page.waitForTimeout(1500)
+ await page.screenshot({ path: `${SHOTS}/01-caret-open.png` })
+ await page.getByTestId('drift-wizard').screenshot({ path: `${SHOTS}/02-caret-face.png` })
+
+ for (const id of FACE) await expect(page.getByTestId(id)).toBeVisible()
+ await expect(page.getByTestId('drift-advanced')).toHaveCount(0)
+ // Everything algorithmic is behind the disclosure.
+ for (const id of ['drift-reference', 'drift-upsample', 'drift-max-shift',
+ 'drift-order', 'drift-tab-rigid', 'drift-apodize']) {
+ await expect(page.getByTestId(id)).toHaveCount(0)
+ }
+ // Count what a user actually sees: buttons + inputs inside the caret.
+ const shown = await page.getByTestId('drift-wizard')
+ .locator('button, input, [data-testid$="-trigger"]').count()
+ expect(shown, 'controls visible on the default face').toBeLessThanOrEqual(5)
+
+ await page.getByTestId('drift-solve').click({ trial: true }) // enabled
+ ctx.assertNoJsErrors()
+})
+
+test('Advanced holds the algorithm; non-rigid is selectable, affine stays locked', async () => {
+ const { page } = ctx
+ await page.getByTestId('drift-advanced-toggle').click()
+ await expect(page.getByTestId('drift-advanced')).toBeVisible()
+ await expect(page.getByTestId('drift-tab-rigid_affine')).toBeDisabled()
+ // Non-rigid is IMPLEMENTED (spyde.drift.nonrigid). This assertion is the one
+ // that catches the trap that actually happened: the backend's _UNAVAILABLE
+ // list is duplicated in the renderer, so implementing the solve while leaving
+ // the tab locked here made a finished feature unreachable — with every
+ // headless test still green, because none of them can see a disabled tab.
+ await expect(page.getByTestId('drift-tab-nonrigid')).toBeEnabled()
+ await page.getByTestId('drift-wizard').screenshot({ path: `${SHOTS}/03-advanced.png` })
+
+ await page.getByTestId('drift-max-shift').fill('24')
+ await page.getByTestId('drift-max-shift').blur()
+ await page.getByTestId('drift-advanced-toggle').click()
+ await expect(page.getByTestId('drift-advanced')).toHaveCount(0)
+ ctx.assertNoJsErrors()
+})
+
+test('selecting Non-rigid reveals its field controls and the choice sticks', async () => {
+ const { page } = ctx
+ await page.getByTestId('drift-advanced-toggle').click()
+ await expect(page.getByTestId('drift-advanced')).toBeVisible()
+
+ // Under RIGID the non-rigid controls must not be on screen at all — they mean
+ // nothing to a rigid solve, and §0.9a is that the caret shows the task.
+ await expect(page.getByTestId('drift-nonrigid-model')).toHaveCount(0)
+
+ await page.getByTestId('drift-tab-nonrigid').click()
+ await expect(page.getByTestId('drift-nonrigid-model')).toBeVisible()
+ await expect(page.getByTestId('drift-nonrigid-steps')).toBeVisible()
+ await page.getByTestId('drift-wizard')
+ .screenshot({ path: `${SHOTS}/10-nonrigid-selected.png` })
+
+ // The tab follows what the BACKEND selected, not what was clicked — so this
+ // failing means the backend refused it and silently stayed on rigid.
+ await expect(page.getByTestId('drift-tab-nonrigid')).toHaveAttribute('aria-selected', 'true')
+
+ // Both parameterisations reach the backend. NB `Select` renders the themed
+ // Dropdown, not a native , so `selectOption`/`inputValue` do nothing
+ // here — click the trigger, then the option, and read `data-value`.
+ await expect(page.getByTestId('drift-nonrigid-model'))
+ .toHaveAttribute('data-value', 'scan_knot')
+ await page.getByTestId('drift-nonrigid-model').click()
+ await page.getByTestId('drift-nonrigid-model-opt-dense').click()
+ await expect(page.getByTestId('drift-nonrigid-model'))
+ .toHaveAttribute('data-value', 'dense')
+ await page.getByTestId('drift-wizard')
+ .screenshot({ path: `${SHOTS}/11-nonrigid-dense.png` })
+ ctx.assertNoJsErrors()
+})
+
+test('a Non-rigid solve actually runs the non-rigid fit, not a quiet fallback', async () => {
+ const { page, backend } = ctx
+ // The caret showing "Non-rigid" proves nothing about what the SOLVER did:
+ // drift_run falls back to rigid whenever the fit cannot run (no torch, OOM,
+ // a device that will not take the graph) and that fallback is deliberate.
+ // So assert on the backend's own report of the KIND it produced.
+ await page.getByTestId('drift-solve').click()
+ await backend.waitForLog('non-rigid fit:', 120_000)
+ const done = await backend.waitForLog('non-rigid fit done:', 300_000)
+ expect(String(done), 'the solve fell back to rigid').toContain('kind=dense')
+ await page.screenshot({ path: `${SHOTS}/12-nonrigid-solved.png` })
+ expect(backendErrorLines(backend), 'backend errors during the non-rigid solve')
+ .toEqual([])
+ ctx.assertNoJsErrors()
+})
+
+test('dragging the ROI re-solves the preview and moves the sharpness number', async () => {
+ const { page } = ctx
+ const sig = sigWindow(page)
+ const gain = () => page.getByTestId('drift-roi-readout').getAttribute('data-gain')
+ const before = await gain()
+ expect(before, 'no gain before the drag').toBeTruthy()
+
+ // Drag the box's centre a long way — the widget lives inside the signal
+ // window's figure iframe, so this is a real pointer drag on real pixels.
+ const box = await sig.locator('iframe').first().boundingBox()
+ expect(box).toBeTruthy()
+ const cx = box!.x + box!.width / 2, cy = box!.y + box!.height / 2
+ await page.mouse.move(cx, cy)
+ await page.mouse.down()
+ for (let i = 1; i <= 8; i++) {
+ await page.mouse.move(cx - i * 6, cy - i * 4)
+ await page.waitForTimeout(30)
+ }
+ await page.screenshot({ path: `${SHOTS}/04-roi-mid-drag.png` })
+ await page.mouse.up()
+
+ await expect.poll(async () => await gain(),
+ { timeout: 90_000, message: 'the preview never re-solved after the drag' },
+ ).not.toBe(before)
+ await page.waitForTimeout(1500)
+ await page.screenshot({ path: `${SHOTS}/05-roi-settled.png` })
+ ctx.assertNoJsErrors()
+})
+
+test('Correct Drift opens the dy/dx window and fills it, then Apply lands the node', async () => {
+ const { page } = ctx
+ await page.getByTestId('drift-solve').click()
+
+ // The curve is its OWN window (plan §0.9a), opened by the solve and filled
+ // from the on_shift stream — so it has points BEFORE the solve finishes.
+ await waitForSubwindowCount(page, 4, 120_000)
+ await expect(page.getByTestId('drift-progress')).toBeVisible({ timeout: 60_000 })
+ await page.screenshot({ path: `${SHOTS}/06-trace-filling.png` })
+
+ await expect(page.getByTestId('drift-result')).toBeVisible({ timeout: 180_000 })
+ await expect(page.getByTestId('drift-result')).toContainText('px drift')
+ await expect(page.getByTestId('drift-status')).toContainText('Solved')
+ await page.waitForTimeout(2000)
+ await page.getByTestId('drift-wizard').screenshot({ path: `${SHOTS}/07-solved-caret.png` })
+ await page.screenshot({ path: `${SHOTS}/08-solved-full.png` })
+
+ // Apply adds the LAZY corrected node to the tree (map_blocks, nothing copied)
+ // and shows it — so it appears in the Plot Control workflow list.
+ await page.getByTestId('drift-commit').click()
+ await expect(page.getByTestId('tree-node-Drift corrected')).toBeVisible({ timeout: 60_000 })
+ await expect(page.getByTestId('status-text')).toContainText('Drift corrected node added')
+ await page.waitForTimeout(2000)
+ await page.screenshot({ path: `${SHOTS}/09-applied.png` })
+
+ const errors = backendErrorLines(ctx.backend)
+ expect(errors, `backend errors:\n${errors.join('\n')}`).toEqual([])
+ ctx.assertNoJsErrors()
+})
diff --git a/electron/tests/seg_overlay.spec.ts b/electron/tests/seg_overlay.spec.ts
new file mode 100644
index 00000000..693a5008
--- /dev/null
+++ b/electron/tests/seg_overlay.spec.ts
@@ -0,0 +1,145 @@
+/**
+ * seg_overlay.spec.ts — the live segmentation overlay, on a TILED frame.
+ *
+ * The bug this exists for: the caret reported "106 particles on this frame" and
+ * nothing was drawn on a real 4096² movie. `_preview` was calling
+ * `set_overlay_mask` and the mask WAS being pushed (`[plot] overlay mask set:
+ * N px` in the backend log) — but a signal frame at or above 1024 px goes
+ * through anyplotlib's GPU tile display, whose base image is drawn by WebGPU,
+ * and the mask composites onto the Canvas2D context underneath it. Invisible.
+ * Vector markers draw over the GPU base correctly, which is why the brush
+ * strokes showed up in the same screenshot that had no overlay.
+ *
+ * So the frame size is the whole point of this spec. `load_test_data_particles`
+ * defaults to 96×112 — BELOW the tile threshold — which is exactly why the
+ * existing `segment_wizard.spec.ts` never caught this. Here it is loaded at
+ * 1200² so the tiled path is the one under test.
+ *
+ * What it proves:
+ * 1. After a preview finds particles, outlines are actually DRAWN (green
+ * pixels appear on the figure canvas, and only after the preview lands).
+ * 2. The outlines FOLLOW THE NAVIGATOR — scrolling to another frame
+ * re-previews and repaints, without the caret being touched.
+ * 3. Closing the caret takes the overlay with it.
+ */
+import { test, expect } from '@playwright/test'
+import { mkdirSync } from 'fs'
+const {
+ launchApp, backendAction, waitForSubwindowCount, sigWindow,
+ countColorPixels, backendErrorLines,
+} = require('./_harness.cjs')
+
+const SHOTS = 'seg_overlay_shots'
+let ctx: Awaited>
+
+test.describe.configure({ mode: 'serial' })
+test.setTimeout(420_000)
+
+test.beforeAll(async () => {
+ mkdirSync(SHOTS, { recursive: true })
+ ctx = await launchApp({ dask: true, env: { SPYDE_LOG_LEVEL: 'INFO' } })
+ const { page } = ctx
+ await page.waitForTimeout(1500)
+ // 1200² is ABOVE anyplotlib's 1024 tile threshold — the display path the 4k
+ // dataset uses, and the one the 96×112 default never exercises.
+ await backendAction(page, 'load_test_data_particles', { frames: 6, size: [1200, 1200] })
+ await waitForSubwindowCount(page, 2, 180_000)
+ await page.waitForTimeout(3000)
+})
+
+test.afterAll(async () => {
+ await ctx?.app?.close()
+})
+
+async function openCaret() {
+ const { page } = ctx
+ const sig = sigWindow(page)
+ await sig.getByTestId('subwindow-title').click()
+ await sig.getByTestId('subwindow-titlebar').hover()
+ await sig.getByTestId('action-btn-Segment Particles').click()
+ await expect(page.getByTestId('segment-wizard')).toBeVisible()
+ return sig
+}
+
+/** The caret's own monotonic preview counter — the reliable "it re-ran" signal.
+ * The COUNT is not: two frames can legitimately find the same number. */
+async function previewSeq(): Promise {
+ const attr = await ctx.page.getByTestId('seg-preview-stats').getAttribute('data-seq')
+ return Number(attr ?? 0)
+}
+async function previewCount(): Promise {
+ const attr = await ctx.page.getByTestId('seg-preview-stats').getAttribute('data-count')
+ return Number(attr ?? -1)
+}
+
+test('outlines are drawn on a tiled frame once the preview lands', async () => {
+ const { page } = ctx
+
+ // BEFORE the caret exists there is no overlay, so this is the baseline the
+ // "it drew something" assertion is measured against rather than a bare >0.
+ const greenBefore = await countColorPixels(page, 'green')
+ await page.screenshot({ path: `${SHOTS}/01-before-caret.png` })
+
+ await openCaret()
+ await expect.poll(previewCount, {
+ timeout: 180_000, message: 'seg_preview never reached the caret',
+ }).toBeGreaterThan(0)
+ // The push is a figure update marshalled onto the main loop; give it a beat
+ // to reach the canvas after the count line has updated.
+ await expect.poll(() => countColorPixels(page, 'green'), {
+ timeout: 60_000,
+ message: 'the preview found particles but drew no outlines — the overlay '
+ + 'never reached the figure (this is the GPU-tile bug)',
+ }).toBeGreaterThan(greenBefore + 200)
+
+ await page.screenshot({ path: `${SHOTS}/02-outlines.png` })
+ await sigWindow(page).screenshot({ path: `${SHOTS}/03-outlines-window.png` })
+})
+
+test('the outlines follow the navigator', async () => {
+ const { page } = ctx
+
+ const seq0 = await previewSeq()
+ expect(seq0, 'no preview to follow').toBeGreaterThan(0)
+
+ // Drive the NAVIGATOR, not the caret — this is the "scroll through the
+ // dataset and watch it update" path, and nothing subscribed to it before.
+ // `test_nav_drag` moves the real navigation SELECTOR, which is where the
+ // wizard's index hook lives; clicking the navigator canvas would not reach it
+ // anyway (the canvas is inside a nested figure iframe).
+ await backendAction(page, 'test_nav_drag', { targets: [[3, 0]] })
+
+ await expect.poll(previewSeq, {
+ timeout: 180_000,
+ message: 'moving the navigator did not re-preview — the overlay is stuck '
+ + 'on whichever frame was showing when the caret was opened',
+ }).toBeGreaterThan(seq0)
+
+ // And it still draws after following.
+ expect(await countColorPixels(page, 'green'),
+ 'the overlay vanished after the navigator moved').toBeGreaterThan(200)
+
+ await page.screenshot({ path: `${SHOTS}/04-followed-navigator.png` })
+ await sigWindow(page).screenshot({ path: `${SHOTS}/05-followed-window.png` })
+
+ const errors = backendErrorLines(ctx.backend)
+ expect(errors, `backend errors:\n${errors.join('\n')}`).toEqual([])
+ ctx.assertNoJsErrors()
+})
+
+test('closing the caret removes the overlay', async () => {
+ const { page } = ctx
+ const before = await countColorPixels(page, 'green')
+ expect(before, 'nothing drawn to remove').toBeGreaterThan(200)
+
+ await page.getByTestId('seg-close').click()
+ await expect(page.getByTestId('segment-wizard')).toHaveCount(0)
+
+ await expect.poll(() => countColorPixels(page, 'green'), {
+ timeout: 60_000,
+ message: 'the outlines outlived the caret that owns them',
+ }).toBeLessThan(200)
+
+ await page.screenshot({ path: `${SHOTS}/06-closed.png` })
+ ctx.assertNoJsErrors()
+})
diff --git a/electron/tests/seg_oversegment.spec.ts b/electron/tests/seg_oversegment.spec.ts
new file mode 100644
index 00000000..1a15a711
--- /dev/null
+++ b/electron/tests/seg_oversegment.spec.ts
@@ -0,0 +1,155 @@
+/**
+ * seg_oversegment.spec.ts — the reported failure, reproduced and pinned.
+ *
+ * Reported with a screenshot on a real in-situ movie: "14028 particles in this
+ * region", the whole preview window painted a flat sheet of green, the renderer
+ * hung, and the Scribble tab unusable because the image you have to paint on was
+ * under that sheet.
+ *
+ * THREE separate defects produced that one screenshot, and every existing spec
+ * was green throughout because the bundled fixture is small, clean and
+ * high-contrast:
+ *
+ * 1. The raster overlay was UNREACHABLE on a tiled frame. `_set_raster_overlay`
+ * reduced the mask to the overview grid (what the renderer wants, since it
+ * checks `bytes.length === (base_width||image_width) * …`), anyplotlib's
+ * `set_overlay_mask` validated that against `image_width` — the FULL native
+ * frame in tile mode — and raised. SpyDE logged the ValueError at DEBUG and
+ * fell back to one filled polygon per instance. At 14028 instances that is
+ * the hang, and thousands of overlapping translucent fills are the green
+ * sheet. So the path added to avoid N polygons could never run on the only
+ * frames big enough to need it.
+ * 2. Nothing capped the polygon fallback.
+ * 3. A failed threshold was reported as a result. Otsu on a low-contrast frame
+ * has no bimodal histogram to find, lands inside the noise, and the split
+ * shatters the support film — but the caret said "14028 particles", which
+ * reads as a bad answer rather than as no answer, and sends the user to
+ * sliders that cannot fix it.
+ *
+ * `noise: 0.35` is what makes the fixture fail this way; at its default 0.015 it
+ * is clean and a global threshold works fine on it. 1200² is above anyplotlib's
+ * 1024 tile threshold, so the tiled path is the one under test — that pairing is
+ * the entire point, and either one alone reproduces nothing.
+ */
+import { test, expect } from '@playwright/test'
+import { mkdirSync } from 'fs'
+const {
+ launchApp, backendAction, waitForSubwindowCount, sigWindow, countColorPixels,
+} = require('./_harness.cjs')
+
+const SHOTS = 'seg_oversegment_shots'
+let ctx: Awaited>
+
+test.describe.configure({ mode: 'serial' })
+test.setTimeout(420_000)
+
+test.beforeAll(async () => {
+ mkdirSync(SHOTS, { recursive: true })
+ ctx = await launchApp({ dask: true, env: { SPYDE_LOG_LEVEL: 'INFO' } })
+ const { page } = ctx
+ await page.waitForTimeout(1500)
+ await backendAction(page, 'load_test_data_particles',
+ { frames: 4, size: [1200, 1200], noise: 0.35 })
+ await waitForSubwindowCount(page, 2, 180_000)
+ await page.waitForTimeout(3000)
+})
+
+test.afterAll(async () => {
+ await ctx?.app?.close()
+})
+
+async function openCaret() {
+ const { page } = ctx
+ const sig = sigWindow(page)
+ await sig.getByTestId('subwindow-title').click()
+ await sig.getByTestId('subwindow-titlebar').hover()
+ await sig.getByTestId('action-btn-Segment Particles').click()
+ await expect(page.getByTestId('segment-wizard')).toBeVisible()
+ return sig
+}
+
+test('a noise frame over-segments, and the caret NAMES it instead of counting it', async () => {
+ const { page } = ctx
+ await openCaret()
+
+ const stats = page.getByTestId('seg-preview-stats')
+ await expect.poll(async () => Number(await stats.getAttribute('data-count')), {
+ timeout: 180_000, message: 'seg_preview never reached the caret',
+ }).toBeGreaterThan(0)
+
+ const count = Number(await stats.getAttribute('data-count'))
+ // The fixture has to actually FAIL or this spec proves nothing — the whole
+ // reason the bug shipped is that the clean fixture never got here.
+ expect(count, `only ${count} instances; noise:0.35 did not over-segment, so `
+ + 'this spec is not exercising the reported failure').toBeGreaterThan(200)
+
+ // The verdict, not just the number. "14028 particles" reads as an answer.
+ await expect(stats).toHaveAttribute('data-failed', 'true')
+ const notice = page.getByTestId('seg-threshold-failed')
+ await expect(notice).toBeVisible()
+ await expect(notice).toContainText('landed inside the noise')
+ // ...and it points at the engine that DOES work on this data (plan §0.9),
+ // rather than leaving the user on sliders that cannot fix a bad threshold.
+ await expect(notice).toContainText('Scribble')
+
+ await page.screenshot({ path: `${SHOTS}/01-threshold-failed.png` })
+ await page.getByTestId('segment-wizard').screenshot({
+ path: `${SHOTS}/02-caret.png` })
+ ctx.assertNoJsErrors()
+})
+
+test('the overlay is ONE mask, not thousands of polygons, and does not blanket the frame', async () => {
+ const { page } = ctx
+
+ // THE regression guard for defect 1. This line is emitted at WARNING by
+ // `_set_raster_overlay`'s except branch, so it reaches the harness's stderr
+ // buffer with SPYDE_LOG_LEVEL=INFO. Before the fix it fired on every preview
+ // of this frame — and at DEBUG, where nobody would ever see it.
+ const failed = (ctx.backend.logBuffer as string[])
+ .filter((l) => l.includes('raster overlay failed'))
+ expect(failed, `the raster overlay fell back to polygons:\n${failed.join('\n')}`)
+ .toEqual([])
+
+ // The cap must not have fired either — reaching it means the raster path was
+ // unavailable, which is the bug wearing a seatbelt rather than the bug fixed.
+ const capped = (ctx.backend.logBuffer as string[])
+ .filter((l) => l.includes('exceeds the'))
+ expect(capped, `the outline draw cap fired, so the mask never drew:\n`
+ + capped.join('\n')).toEqual([])
+
+ // And the visual half: an overlay is drawn, but it is NOT the solid sheet of
+ // the screenshot. A mask that covers essentially the whole window carries no
+ // information — you cannot see the data under it, which is what made the
+ // Scribble tab unusable.
+ const green = await countColorPixels(page, 'green')
+ const sig = sigWindow(page)
+ const box = await sig.boundingBox()
+ const area = box ? box.width * box.height : 1
+ expect(green, 'no overlay was drawn at all').toBeGreaterThan(100)
+ expect(green / area,
+ `the overlay covers ${(100 * green / area).toFixed(0)}% of the window — `
+ + 'that is the solid-green sheet, not an overlay').toBeLessThan(0.5)
+
+ await page.screenshot({ path: `${SHOTS}/03-overlay-not-a-sheet.png` })
+ ctx.assertNoJsErrors()
+})
+
+test('Scribble stays usable: the image is not buried under the failed preview', async () => {
+ const { page } = ctx
+ await page.getByTestId('seg-tab-scribble').click()
+ await expect(page.getByTestId('seg-class-strip')).toBeVisible({ timeout: 30_000 })
+
+ // The reported symptom, in one number: you cannot paint on an image you
+ // cannot see. The untrained Scribble engine has produced no result, so the
+ // previous engine's drawing must be GONE — not merely thinner. It used to
+ // survive because `show_preview_window` cleared the vector outlines and left
+ // the raster mask, and above 100 instances the mask is the whole drawing.
+ await expect.poll(() => countColorPixels(page, 'green'), {
+ timeout: 30_000,
+ message: 'the classical result is still drawn on the Scribble tab, over '
+ + 'the image the user has to paint on',
+ }).toBeLessThan(2000)
+
+ await page.screenshot({ path: `${SHOTS}/04-scribble-usable.png` })
+ ctx.assertNoJsErrors()
+})
diff --git a/electron/tests/segment_wizard.spec.ts b/electron/tests/segment_wizard.spec.ts
new file mode 100644
index 00000000..6b9d80eb
--- /dev/null
+++ b/electron/tests/segment_wizard.spec.ts
@@ -0,0 +1,567 @@
+/**
+ * segment_wizard.spec.ts — the Segment Particles caret, end-to-end on the
+ * bundled synthetic particle movie.
+ *
+ * What this actually proves (headless tests + tsc cannot see any of it):
+ * 1. The caret opens from the real toolbar button and the backend previews
+ * the DISPLAYED frame — the count line names a number of particles.
+ * 2. The DEFAULT face is calm: one slider, one count, one button, one
+ * disclosure. Everything else is behind `▸ Advanced`, which is collapsed
+ * on open and remembers its state.
+ * 3. The floating brush strip renders NEXT TO THE PLOT (plan B0) with one
+ * swatch per backend class, and is not clipped by the window.
+ * 4. The size histogram (inside Advanced) has bars, not an empty box.
+ * 5. `min_size` = 0 is FLOORED by the backend and the caret shows the
+ * EFFECTIVE value, not the 0 the user typed (plan §0.9 — at 0 the split
+ * returns background speckle as particles). Both the field and the warning
+ * live inside Advanced.
+ * 6. "Find in all frames" opens a real particle result window.
+ * 7. A brush stroke reaches `seg_paint` and the per-class labelled-pixel
+ * counts in the caret update — the counts are how a user notices an
+ * under-trained class, so a count stuck at 0 is a real failure. The class
+ * list is the SCRIBBLE tab's business and is not shown on Classical.
+ * 8. The BOUNDARY class is offered, is paintable, and painting it flips the
+ * split route — the caret says `watershed split` before and `seam split`
+ * after, which is the only thing on screen that distinguishes a 0.33 s
+ * split from a 1.78 s one at 4096². Its hover text has to carry the
+ * "paint the seam, not the outline" warning, because the intuitive reading
+ * trains a head that MERGES touching particles (benchmarks.md) and nothing
+ * else on screen says which reading is right.
+ *
+ * Real Dask + `load_test_data_particles` (lazy, 1 frame/chunk, ground truth
+ * stamped into metadata) — the path a user actually drags.
+ */
+import { test, expect } from '@playwright/test'
+import { mkdirSync } from 'fs'
+const {
+ launchApp, backendAction, waitForSubwindowCount, sigWindow, backendErrorLines,
+} = require('./_harness.cjs')
+
+const SHOTS = 'segment_wizard_shots'
+let ctx: Awaited>
+/** `data-testid="figure-"` on the signal window's iframe. */
+let figId = ''
+
+test.describe.configure({ mode: 'serial' })
+test.setTimeout(300_000)
+
+test.beforeAll(async () => {
+ mkdirSync(SHOTS, { recursive: true })
+ // INFO tees `logging` to stderr, which the harness captures — backend
+ // emit_error goes over the PLOTAPP protocol and never reaches this buffer.
+ ctx = await launchApp({ dask: true, env: { SPYDE_LOG_LEVEL: 'INFO' } })
+ const { page } = ctx
+ // backend-ready can land a beat before the stdin pump is live; the lazy specs
+ // settle the same way so the first action isn't dropped.
+ await page.waitForTimeout(1500)
+ await backendAction(page, 'load_test_data_particles', { frames: 6 })
+ await waitForSubwindowCount(page, 2, 120_000)
+ await page.waitForTimeout(2000)
+})
+
+test.afterAll(async () => {
+ await ctx?.app?.close()
+})
+
+/** The SOURCE movie's signal window specifically — `sigWindow` picks the first
+ * `S-` window, and Run All adds a second one ("S-Particles — 6 frames"). */
+function srcWindow() {
+ const { page } = ctx
+ return page.getByTestId('subwindow').filter({
+ has: page.getByTestId('window-breadcrumb').filter({ hasText: /^S-Synthetic/ }),
+ }).first()
+}
+
+/** Raise the source window above the result windows Run All cascades on top of
+ * it. The caret lives in that window's stacking context, so a newer window
+ * stacked above ALSO covers the caret and swallows clicks meant for it — the
+ * same thing a user does (click the window) fixes it. */
+async function raiseSource() {
+ await srcWindow().getByTestId('subwindow-title').click()
+}
+
+/** Open the caret from the REAL toolbar button (hover the titlebar to reveal
+ * the floating bar), exactly as a user does. */
+async function openCaret() {
+ const { page } = ctx
+ const sig = sigWindow(page)
+ await sig.getByTestId('subwindow-title').click()
+ await sig.getByTestId('subwindow-titlebar').hover()
+ await sig.getByTestId('action-btn-Segment Particles').click()
+ await expect(page.getByTestId('segment-wizard')).toBeVisible()
+ return sig
+}
+
+/**
+ * Drive the `▸ Advanced` disclosure to a known state. Everything the primary
+ * face no longer shows is in there, so most of the assertions below have to open
+ * it first — which is the point of the redesign, not an inconvenience.
+ *
+ * ONE click has to do it, deliberately: the caret is a DOM child of
+ * FloatingToolbar, whose placement effect moves it (below ↔ right) when its
+ * height changes. Toggling the disclosure changes that height from the WIZARD's
+ * own state, which does not re-render the toolbar — so before the toolbar's
+ * ResizeObserver existed the placement stayed stale and the caret jumped on the
+ * NEXT render, i.e. between the mousedown and the mouseup of the following
+ * click. The browser then emits no `click` at all: the control takes focus and
+ * silently does nothing, and every other click is ignored. A retry loop here
+ * would hide exactly that, so there isn't one.
+ */
+async function setAdvanced(open: boolean) {
+ const { page } = ctx
+ const adv = page.getByTestId('seg-advanced')
+ const toggle = page.getByTestId('seg-advanced-toggle')
+ if (((await adv.count()) > 0) === open) return
+ await toggle.click()
+ await expect(adv).toHaveCount(open ? 1 : 0)
+ await expect(toggle).toHaveAttribute('aria-expanded', String(open))
+}
+
+/**
+ * The caret box has NO scroller of its own — the Threshold dropdown's menu is
+ * absolutely positioned, so an `overflow:auto` ancestor would clip it. That
+ * makes "does it fit" a real assertion rather than a cosmetic one: anything
+ * past the bottom of the MDI area is simply unreachable, and an expanded
+ * Advanced is where it happens.
+ *
+ * It HAS happened: single-column Advanced measured 907 px in an 805 px area,
+ * putting the size histogram and Commit Frame off-screen with no way to reach
+ * them. The two-column layout (plan B7) is what buys the room back, so this
+ * assertion is the thing holding that layout in place — if a future edit
+ * re-stacks Advanced into one column, this is what says so.
+ *
+ * BOTH AXES, because checking only the bottom is how the next one got through:
+ * the two-column caret fit vertically and then FloatingToolbar placed it off
+ * the LEFT edge of the app (a side-placed caret anchors its right edge to the
+ * window's left, so one wider than the room beside the window walks straight
+ * out of the viewport). Vertically-only, this function passed while half the
+ * caret was unreachable.
+ */
+async function expectCaretFits() {
+ const { page } = ctx
+ const box = await page.getByTestId('segment-wizard').boundingBox()
+ const mdi = await page.getByTestId('mdi-area').boundingBox()
+ const where = `caret ${JSON.stringify(box)} vs MDI area ${JSON.stringify(mdi)}`
+ expect(box && mdi, where).toBeTruthy()
+ expect(box!.y + box!.height <= mdi!.y + mdi!.height + 1,
+ `${where}: runs past the BOTTOM`).toBe(true)
+ expect(box!.x >= mdi!.x - 1, `${where}: runs past the LEFT edge`).toBe(true)
+ expect(box!.x + box!.width <= mdi!.x + mdi!.width + 1,
+ `${where}: runs past the RIGHT edge`).toBe(true)
+}
+
+test('caret opens, previews the displayed frame, and shows the brush strip', async () => {
+ const { page } = ctx
+ await page.screenshot({ path: `${SHOTS}/01-movie-loaded.png` })
+
+ const sig = await openCaret()
+ const tid = await sig.locator('iframe').first().getAttribute('data-testid')
+ figId = (tid ?? '').replace(/^figure-/, '')
+ expect(figId).not.toBe('')
+
+ await page.screenshot({ path: `${SHOTS}/02-caret-open.png` })
+
+ // The preview is a real backend round trip (seg_open → worker → seg_preview),
+ // so wait for the count line to stop saying "no preview yet".
+ await expect.poll(
+ () => page.getByTestId('seg-preview-stats').textContent(),
+ { timeout: 60_000, message: 'seg_preview never reached the caret' },
+ ).toMatch(/\d+ particles? on this frame/)
+
+ // ── the DEFAULT face is the whole point of the redesign ──────────────────
+ // Advanced is collapsed on open, and everything it holds is genuinely absent
+ // from the DOM (not merely visually quiet).
+ await expect(page.getByTestId('seg-advanced-toggle')).toHaveAttribute('aria-expanded', 'false')
+ await expect(page.getByTestId('seg-advanced')).toHaveCount(0)
+ for (const hidden of [
+ 'seg-min-size', 'seg-max-size', 'seg-watershed', 'seg-store-masks',
+ 'seg-track', 'seg-threshold', 'seg-gaussian', 'seg-rb-kernel',
+ 'seg-local-size', 'seg-min-separation', 'seg-marker-smooth', 'seg-max-dist',
+ 'seg-invert', 'seg-clear-border', 'seg-histogram', 'seg-counts',
+ 'seg-commit', 'seg-min-score',
+ ]) {
+ await expect(page.getByTestId(hidden),
+ `${hidden} must be behind Advanced, not on the default face`).toHaveCount(0)
+ }
+ // THREE sliders and the button, nothing else besides the tabs, ✕ and the
+ // disclosure. ✕, 3 tabs, sensitivity, merge-nm, min-nm, Find-in-all-frames,
+ // Advanced = 9. This count is the guard against the face refilling one
+ // reasonable-looking addition at a time (plan §0.9a), so it is exact on
+ // purpose — if you add a control here, justify it in the diff.
+ const primaryControls = await page.getByTestId('segment-wizard')
+ .locator('input, select, textarea, button').count()
+ expect(primaryControls, 'the default face grew a control back').toBe(9)
+ await expect(page.getByTestId('seg-sensitivity')).toBeVisible()
+ // The two PHYSICAL controls, the ones that are the same on every engine.
+ await expect(page.getByTestId('seg-merge-nm')).toBeVisible()
+ await expect(page.getByTestId('seg-min-nm')).toBeVisible()
+ await expect(page.getByTestId('seg-run')).toHaveText('Find in all frames')
+
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/03-caret-detail.png` })
+
+ // ── Advanced still holds everything, and it still works ──────────────────
+ await setAdvanced(true)
+ // An EMPTY histogram is the classic "it rendered but says nothing" failure —
+ // assert on the bar count the component publishes, not on pixels.
+ await expect.poll(
+ async () => Number(await page.getByTestId('seg-histogram').getAttribute('data-nonzero')),
+ { timeout: 30_000, message: 'size histogram has no populated bins' },
+ ).toBeGreaterThan(0)
+ await expect(page.getByTestId('seg-min-size')).toBeVisible()
+ await expect(page.getByTestId('seg-commit')).toBeVisible()
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/03b-caret-advanced.png` })
+ // Expanded, IN PLACE: Advanced makes the caret tall enough that the toolbar
+ // re-places it beside the window. It must still fit inside the MDI area — the
+ // caret box has no scroller of its own, so anything past the bottom is simply
+ // unreachable.
+ await page.screenshot({ path: `${SHOTS}/03c-advanced-full.png` })
+ await expectCaretFits()
+ await setAdvanced(false)
+
+ // The brush strip belongs to the SCRIBBLE tab only. It floats over the image,
+ // so on Classical — where there is nothing to paint — it would be chrome
+ // covering the data for no reason.
+ await expect(page.getByTestId('seg-class-strip'),
+ 'the brush strip is showing on Classical, where there is nothing to paint',
+ ).toHaveCount(0)
+
+ // The class list carries NAMES + per-class pixel counts (the caret is the
+ // authoritative list; the strip is swatches only) — also Scribble's business.
+ await expect(page.getByTestId('seg-class-list')).toHaveCount(0)
+
+ await page.getByTestId('seg-tab-scribble').click()
+ await expect(page.getByTestId('seg-class-0')).toBeVisible({ timeout: 30_000 })
+ await expect(page.getByTestId('seg-class-pixels-0')).toBeVisible()
+ // ...and NOW the strip appears, next to the plot rather than in the caret.
+ const strip = page.getByTestId('seg-class-strip')
+ await expect(strip).toBeVisible({ timeout: 30_000 })
+ await expect(page.getByTestId('seg-strip-class-0')).toBeVisible()
+ await expect(page.getByTestId('seg-strip-brush')).toBeVisible()
+ await expect(page.getByTestId('seg-strip-eraser')).toBeVisible()
+ // No `+ add class`: the backend has no seg_add_class verb, and a permanently
+ // disabled control is noise on a face this redesign just emptied out.
+ await expect(page.getByTestId('seg-add-class')).toHaveCount(0)
+
+ await page.getByTestId('seg-tab-classical').click()
+ await expect(page.getByTestId('seg-sensitivity')).toBeVisible({ timeout: 30_000 })
+ await expect(page.getByTestId('seg-class-strip')).toHaveCount(0)
+
+ await page.screenshot({ path: `${SHOTS}/04-preview.png` })
+ ctx.assertNoJsErrors()
+})
+
+test('sensitivity re-previews and min_size=0 is floored to the EFFECTIVE value', async () => {
+ const { page } = ctx
+ const stats = page.getByTestId('seg-preview-stats')
+
+ // The COUNT alone is not a reliable "did it re-run" signal — two sensitivities
+ // can legitimately find the same number of particles — so poll the caret's own
+ // monotonic preview counter instead of diffing the label text.
+ const seq0 = Number(await stats.getAttribute('data-seq'))
+ await page.getByTestId('seg-sensitivity').fill('0.85')
+ await expect.poll(async () => Number(await stats.getAttribute('data-seq')), {
+ timeout: 60_000, message: 'dragging sensitivity did not re-preview',
+ }).toBeGreaterThan(seq0)
+ await page.screenshot({ path: `${SHOTS}/05-sensitivity.png` })
+
+ // min_size=0 is the footgun plan §0.9 measured (33 instances where 9 are
+ // real). The backend floors it to 10 and the caret must show 10. Both the
+ // field and the warning are inside Advanced now: the floor is applied
+ // unconditionally, so the primary face never has to mention it.
+ await setAdvanced(true)
+ const minSize = page.getByTestId('seg-min-size')
+ await minSize.fill('0')
+ await minSize.blur()
+ await expect(page.getByTestId('seg-min-size-floor')).toBeVisible({ timeout: 60_000 })
+ await expect.poll(() => minSize.inputValue(), {
+ timeout: 30_000, message: 'caret still shows 0 while the backend ran 10',
+ }).toBe('10')
+
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/06-minsize-floored.png` })
+
+ // Back to a sane value for the batch below, and back to the calm face.
+ await minSize.fill('20')
+ await minSize.blur()
+ await setAdvanced(false)
+ ctx.assertNoJsErrors()
+})
+
+/**
+ * The three filters that had NO coverage, which is how all three shipped broken
+ * at once: the two nm sliders crashed the whole caret on first render (`Field`
+ * was used but never imported — a blank window, and every headless test still
+ * green), and Confidence was taken off the face and never re-added to Advanced,
+ * so `min_score` sat at 0 with no control able to move it.
+ *
+ * Each assertion below is therefore "the control exists AND reaches the
+ * backend", polling the caret's monotonic `data-seq`. Existence alone is what
+ * the previous specs checked, and it is exactly what a dead slider passes.
+ */
+test('the nm face filters and the demoted Confidence slider all re-preview', async () => {
+ const { page } = ctx
+ const stats = page.getByTestId('seg-preview-stats')
+ const reran = async (what: string, act: () => Promise) => {
+ const seq = Number(await stats.getAttribute('data-seq'))
+ await act()
+ await expect.poll(async () => Number(await stats.getAttribute('data-seq')), {
+ timeout: 60_000, message: `${what} did not reach the backend`,
+ }).toBeGreaterThan(seq)
+ }
+
+ // ── the face: both controls are PHYSICAL and read out in nm ───────────────
+ // The readout is the assertion that matters. `_nm_to_px` converts with the
+ // signal's scale, and nothing stashed that scale on the params — so the
+ // backend took its uncalibrated branch and merged at N PIXELS while the label
+ // said N nm. A label in nm is a claim about the scale bar; if the conversion
+ // is skipped the caret is lying by exactly the magnification.
+ await expect(page.getByTestId('seg-merge-nm')).toHaveValue('0')
+ await reran('the merge-nm slider', () => page.getByTestId('seg-merge-nm').fill('25'))
+ await expect(page.getByTestId('segment-wizard')).toContainText('25 nm')
+
+ await reran('the min-nm slider', () => page.getByTestId('seg-min-nm').fill('4'))
+ await expect(page.getByTestId('segment-wizard')).toContainText('4 nm')
+
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/06b-nm-filters.png` })
+
+ // Both are engine-independent, so unlike sensitivity they stay on the face
+ // when the engine changes. (Scribble is untrained here, which does not matter
+ // — this is about which controls render.)
+ await page.getByTestId('seg-tab-scribble').click()
+ await expect(page.getByTestId('seg-merge-nm')).toBeVisible()
+ await expect(page.getByTestId('seg-min-nm')).toBeVisible()
+ await expect(page.getByTestId('seg-sensitivity'),
+ 'sensitivity is classical-only — the scribble engine never reads it').toHaveCount(0)
+ await page.getByTestId('seg-tab-classical').click()
+ await expect(page.getByTestId('seg-sensitivity')).toBeVisible({ timeout: 30_000 })
+
+ // ── Advanced: Confidence is demoted, NOT deleted (plan §0.9a) ─────────────
+ await setAdvanced(true)
+ const score = page.getByTestId('seg-min-score')
+ await expect(score, 'Confidence left the face and never arrived in Advanced')
+ .toBeVisible()
+ await expect(page.getByTestId('seg-advanced')).toContainText('off')
+ await reran('the Confidence slider', () => score.fill('0.5'))
+ await expect(page.getByTestId('seg-advanced')).toContainText('50%')
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/06c-confidence.png` })
+
+ // Everything back to default so the batch run below is the plain path.
+ await reran('resetting Confidence', () => score.fill('0'))
+ await setAdvanced(false)
+ await reran('resetting merge-nm', () => page.getByTestId('seg-merge-nm').fill('0'))
+ await reran('resetting min-nm', () => page.getByTestId('seg-min-nm').fill('0'))
+ ctx.assertNoJsErrors()
+})
+
+test('Find in all frames segments the movie into a new particle window', async () => {
+ const { page } = ctx
+ const before = await page.getByTestId('subwindow').count()
+ await page.getByTestId('seg-run').click()
+ // The button's own status line is the proof the React handler ran at all —
+ // a swallowed click leaves the last preview's text sitting there and the
+ // 3-minute window poll below then fails for the wrong reason.
+ await expect(page.getByTestId('seg-status')).toHaveText(/Segmenting the movie/)
+ await expect.poll(() => page.getByTestId('subwindow').count(), {
+ timeout: 180_000, message: 'the particle result window never opened',
+ }).toBeGreaterThan(before)
+ await page.screenshot({ path: `${SHOTS}/07a-run-early-window.png` })
+
+ // The window opens EARLY with an empty store; `tree.particles` attaches only
+ // at _finalize, which re-sends the toolbar — so the requires_particles-gated
+ // buttons appearing IS the "batch finished" signal, exactly as
+ // `particles_action._rebuild_toolbars` documents ("the e2e specs wait on
+ // exactly this appearing"). NB the status bar is NOT usable here: seg_run's
+ // last emit_progress leaves busy=true, and StatusBar shows loading.text over
+ // status while busy — so "Found N particles" never becomes visible.
+ // Depends on the requires_particles toolbar entries (plan B9, Particle
+ // Overlay / Particle Lanes) being present in spyde/toolbars.yaml.
+ await expect.poll(
+ () => page.getByTestId('action-btn-Particle Overlay').count(),
+ { timeout: 180_000, message: 'the segmentation batch never finalized' },
+ ).toBeGreaterThan(0)
+ await page.waitForTimeout(2000)
+ await page.screenshot({ path: `${SHOTS}/07-run-all.png` })
+ ctx.assertNoJsErrors()
+})
+
+test('a brush stroke reaches seg_paint and the class pixel counts update', async () => {
+ const { page } = ctx
+
+ await raiseSource()
+ await page.getByTestId('seg-tab-scribble').click()
+ await expect(page.getByTestId('seg-scribble-note')).toBeVisible()
+
+ const pixels0 = page.getByTestId('seg-class-pixels-0')
+ expect(await pixels0.textContent()).toContain('0')
+
+ // Fatten the brush first, then paint with the strip's ACTIVE class — i.e.
+ // drive the same state the strip owns, not a synthetic payload.
+ await page.getByTestId('seg-strip-brush').fill('7')
+ await page.getByTestId('seg-strip-class-1').click()
+
+ // The anyplotlib brush widget (plan B0) is not landed, so post the widget
+ // event the caret listens for. Points are IMAGE PIXELS [[y, x], …] with no
+ // scale/offset applied — plan trap 6, and what seg_paint documents.
+ const stroke = async (y: number, x0: number, x1: number) => {
+ await page.evaluate(({ id, y, x0, x1 }) => {
+ const points: number[][] = []
+ for (let x = x0; x <= x1; x += 1) points.push([y, x])
+ window.dispatchEvent(new CustomEvent('spyde:figure_event', {
+ detail: { figId: id, event: { type: 'brush_stroke', points } },
+ }))
+ }, { id: figId, y, x0, x1 })
+ }
+ // Selecting on the strip drives the caret's class list too — ONE active
+ // class, not two highlighted rows.
+ await expect(page.getByTestId('seg-class-1')).toHaveAttribute('data-active', 'true')
+ await expect(page.getByTestId('seg-class-0')).toHaveAttribute('data-active', 'false')
+
+ await stroke(20, 20, 90) // class 1 (support film)
+ await expect.poll(() => page.getByTestId('seg-class-pixels-1').textContent(), {
+ timeout: 30_000, message: 'seg_paint never updated the class pixel counts',
+ }).not.toMatch(/^!?\s*0$/)
+
+ await page.getByTestId('seg-strip-class-0').click()
+ await expect(page.getByTestId('seg-class-0')).toHaveAttribute('data-active', 'true')
+ await expect(page.getByTestId('seg-class-1')).toHaveAttribute('data-active', 'false')
+ await stroke(48, 30, 80) // class 0 (particle)
+ await expect.poll(() => pixels0.textContent(), {
+ timeout: 30_000, message: 'painting class 0 did not update its count',
+ }).not.toMatch(/^!?\s*0$/)
+
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/08-painted.png` })
+ // A close-up of the class list on its own: the per-class counts and the
+ // active-row indication are the two things here that have to be legible.
+ await page.getByTestId('seg-class-list').screenshot({ path: `${SHOTS}/08b-class-list.png` })
+ // EXACTLY ONE row may look selected. A previously-active row used to keep a
+ // stale white 1px border (React clears a dropped `borderColor` longhand but
+ // leaves the base `border` shorthand's width/style), so two rows read as
+ // selected at once — invisible to every attribute assertion above.
+ const borders = await page.evaluate(() =>
+ [...document.querySelectorAll('[data-testid^="seg-class-"]')]
+ .filter(e => /^seg-class-\d+$/.test(e.getAttribute('data-testid') ?? ''))
+ .map(e => getComputedStyle(e).borderTopColor))
+ expect(borders.filter(c => c !== 'rgba(0, 0, 0, 0)' && c !== 'transparent'),
+ `rows with a visible border: ${JSON.stringify(borders)}`).toHaveLength(1)
+ await page.screenshot({ path: `${SHOTS}/09-painted-full.png` })
+
+ // The counts line is the "did I label enough" readout — demoted into Advanced
+ // (the class list already carries the per-class numbers on the face).
+ await setAdvanced(true)
+ await expect(page.getByTestId('seg-counts')).toContainText('frames labelled')
+ await setAdvanced(false)
+ ctx.assertNoJsErrors()
+})
+
+test('Train fits the scribble classifier and the caret reports it', async () => {
+ const { page } = ctx
+ await raiseSource()
+ const train = page.getByTestId('seg-train')
+ await expect(train).toBeEnabled()
+ await page.getByTestId('seg-train').click()
+
+ // Assert on the PERSISTENT report line, not the status: the backend follows
+ // seg_trained with a re-preview whose status overwrites it milliseconds later
+ // (a poll on the status races that and loses).
+ await expect(page.getByTestId('seg-trained-note')).toBeVisible({ timeout: 180_000 })
+ await expect(page.getByTestId('seg-trained-note')).toContainText(/Trained on \d+ px/)
+ // Training flips the engine to scribble and unlocks the batch.
+ await expect(page.getByTestId('seg-scribble-note')).toHaveCount(0)
+ await expect(page.getByTestId('seg-run')).toBeEnabled()
+
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/10-trained.png` })
+ await page.screenshot({ path: `${SHOTS}/11-trained-full.png` })
+
+ // The Scribble face, expanded — its parameters share the SAME disclosure, so
+ // nothing moved to a second place.
+ await setAdvanced(true)
+ await expect(page.getByTestId('seg-min-size')).toBeVisible()
+ await expect(page.getByTestId('seg-track')).toBeVisible()
+ // The classical MASK knobs are absent here, and not just to save room: the
+ // scribble engine hands split_instances a probability map thresholded at 0.5
+ // and never reads them (spyde/particles/classical.py::split_instances). Six
+ // knobs that do nothing is the overload complaint in miniature.
+ for (const dead of ['seg-threshold', 'seg-gaussian', 'seg-rb-kernel',
+ 'seg-local-size', 'seg-invert', 'seg-sensitivity']) {
+ await expect(page.getByTestId(dead),
+ `${dead} does not affect the scribble engine and must not be shown`).toHaveCount(0)
+ }
+ // Scribble's Advanced is the TALLEST state the caret has (class list + train
+ // report + parameters); if anything is going to run off the bottom, it is this.
+ await expectCaretFits()
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/12-scribble-advanced.png` })
+ await setAdvanced(false)
+ await page.getByTestId('segment-wizard').screenshot({ path: `${SHOTS}/13-scribble-collapsed.png` })
+
+ const errors = backendErrorLines(ctx.backend)
+ expect(errors, `backend errors:\n${errors.join('\n')}`).toEqual([])
+ ctx.assertNoJsErrors()
+})
+
+test('painting the boundary class flips the split to the seam route', async () => {
+ const { page } = ctx
+ await raiseSource()
+
+ // The previous test trained with no boundary painted, so the caret is sitting
+ // on the watershed route. That is the BEFORE half of this test — without it,
+ // asserting "seam split" afterwards would not prove anything flipped.
+ await expect(page.getByTestId('seg-trained-note')).toContainText('watershed split')
+ await page.getByTestId('segment-wizard').screenshot({
+ path: `${SHOTS}/14-before-boundary.png` })
+
+ // The boundary class is offered at all. It is the 4th default class and it is
+ // opt-in by construction: unpainted, the split falls back to the watershed.
+ const swatch = page.getByTestId('seg-strip-class-3')
+ await expect(swatch).toBeVisible()
+
+ // Its hover text must say WHICH boundary to paint. "Boundary" reads as "the
+ // outline of a particle" to almost everyone, and a head trained on outlines
+ // learns "shrink everything" — measured, it merged the touching pair and lost
+ // 40% of the median area while still reporting a trained boundary class and
+ // still taking the fast route. The wrong reading is worse than not painting at
+ // all AND it is silent, so this tooltip is the only guard there is.
+ const tip = await swatch.getAttribute('title')
+ expect(tip, `boundary swatch tooltip: ${tip}`).toMatch(/SEAM BETWEEN/)
+ expect(tip, 'the tooltip must warn against the outline reading').toMatch(/never the outline/)
+
+ await swatch.click()
+ await expect(page.getByTestId('seg-class-3')).toHaveAttribute('data-active', 'true')
+
+ // Paint a seam. Points are IMAGE PIXELS [[y, x], …] — plan trap 6.
+ const stroke = async (y: number, x0: number, x1: number) => {
+ await page.evaluate(({ id, y, x0, x1 }) => {
+ const points: number[][] = []
+ for (let x = x0; x <= x1; x += 1) points.push([y, x])
+ window.dispatchEvent(new CustomEvent('spyde:figure_event', {
+ detail: { figId: id, event: { type: 'brush_stroke', points } },
+ }))
+ }, { id: figId, y, x0, x1 })
+ }
+ await page.getByTestId('seg-strip-brush').fill('3')
+ await stroke(56, 34, 76)
+ await expect.poll(() => page.getByTestId('seg-class-pixels-3').textContent(), {
+ timeout: 30_000, message: 'painting the boundary class did not update its count',
+ }).not.toMatch(/^!?\s*0$/)
+
+ await page.getByTestId('seg-train').click()
+ // Assert on the PERSISTENT report line, not the status: the backend follows
+ // seg_trained with a re-preview whose status overwrites it milliseconds later.
+ await expect(page.getByTestId('seg-trained-note')).toContainText('seam split', {
+ timeout: 180_000 })
+
+ await page.getByTestId('segment-wizard').screenshot({
+ path: `${SHOTS}/15-seam-route.png` })
+ await page.getByTestId('seg-class-list').screenshot({
+ path: `${SHOTS}/15b-class-list-boundary.png` })
+ await page.screenshot({ path: `${SHOTS}/16-seam-route-full.png` })
+
+ // A boundary that was painted must still segment — the fast route returning
+ // nothing would be a "faster" result that found no particles.
+ await expect.poll(async () =>
+ Number(await page.getByTestId('seg-preview-stats').getAttribute('data-count')), {
+ timeout: 120_000, message: 'no preview after switching to the seam route',
+ }).toBeGreaterThan(0)
+
+ const errors = backendErrorLines(ctx.backend)
+ expect(errors, `backend errors:\n${errors.join('\n')}`).toEqual([])
+ ctx.assertNoJsErrors()
+})
diff --git a/pyproject.toml b/pyproject.toml
index 70960df4..22e1cb01 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -122,7 +122,14 @@ torch = [
# pattern). When actively editing anyplotlib locally, temporarily re-add the
# editable path override below (do NOT commit it — the path doesn't exist on CI
# runners and breaks `uv sync --frozen`):
-# anyplotlib = { path = "../anyplotlib", editable = true }
+# ACTIVE while the brush widget (CSSFrancis/anyplotlib#47) is in flight — the
+# Segment wizard's Scribble tab cannot paint without it, and a plain
+# `pip install -e` does NOT survive: any `uv sync` re-resolves from the lockfile
+# and puts PyPI's 0.4.2 back, which is how scribbling silently broke twice.
+# Declaring the source here is what makes uv PRESERVE the local checkout.
+# DO NOT COMMIT the line below — ../anyplotlib does not exist on CI runners and
+# it breaks `uv sync --frozen`. Delete it once #47 lands and the pin is bumped.
+anyplotlib = { path = "../anyplotlib", editable = true }
[tool.pycrucible]
entrypoint = "main.py"
diff --git a/scripts/bench_drift_particles.py b/scripts/bench_drift_particles.py
new file mode 100644
index 00000000..cc4c9404
--- /dev/null
+++ b/scripts/bench_drift_particles.py
@@ -0,0 +1,165 @@
+#!/usr/bin/env python
+"""
+The drift + particles numbers, in one run.
+
+ python scripts/bench_drift_particles.py
+
+Prints a table suitable for pasting into ``benchmarks.md``. Every number here is a
+number some decision in ``DRIFT_AND_PARTICLES_PLAN.md`` rests on, so re-run it
+after touching the solver or the feature stack rather than trusting the recorded
+values — they were measured on one machine on one day.
+
+Deliberately NOT a pytest: these are slow, machine-dependent, and a benchmark that
+fails CI because a runner was busy teaches nothing. The tests assert CORRECTNESS;
+this reports COST.
+"""
+from __future__ import annotations
+
+import os
+import sys
+import time
+from pathlib import Path
+
+import numpy as np
+
+sys.path.insert(0, str(Path(__file__).resolve().parent.parent))
+
+
+def _t(fn, repeat: int = 3):
+ """Best of *repeat*, discarding the first run (cold CUDA init / kernel JIT)."""
+ fn()
+ best = float("inf")
+ for _ in range(repeat):
+ t0 = time.perf_counter()
+ fn()
+ best = min(best, time.perf_counter() - t0)
+ return best
+
+
+def bench_drift_backends(rows: list[str]) -> None:
+ from spyde.drift.translation import solve_translation
+
+ rng = np.random.default_rng(0)
+ for (nf, n) in [(60, 256), (60, 512)]:
+ base = rng.standard_normal((n, n)).astype(np.float32)
+ stack = np.repeat(base[None], nf, axis=0)
+ rows.append(f"\n### {nf} frames x {n}^2, upsample=8")
+ rows.append("")
+ rows.append("| backend | time | frames/s |")
+ rows.append("|---|---|---|")
+ for dev in ("numpy", "cpu", "cuda"):
+ try:
+ dt = _t(lambda d=dev: solve_translation(stack, device=d, upsample=8),
+ repeat=2)
+ rows.append(f"| {dev} | {dt:.2f} s | {nf / dt:.0f} |")
+ except Exception as exc:
+ rows.append(f"| {dev} | unavailable | {type(exc).__name__} |")
+
+
+def bench_drift_accuracy(rows: list[str]) -> None:
+ """Error vs upsample, on truth deliberately OFF the 1/upsample grid.
+
+ On-grid truth is recovered exactly at any upsample, which looks superb and
+ tests nothing — that mistake hid a real bug in the refinement once already.
+ """
+ from spyde.drift.translation import solve_translation
+ from spyde.tests.migrated.test_drift_translation import _shifted_stack
+
+ truth = np.array([[0, 0], [1.37, -2.83], [-3.06, 0.61], [4.19, 5.44],
+ [-0.72, -1.28]])
+ stack = _shifted_stack(truth)
+ rows.append("\n### Sub-pixel accuracy vs upsample (off-grid truth)")
+ rows.append("")
+ rows.append("| upsample | max error |")
+ rows.append("|---|---|")
+ for u in (1, 2, 8, 32, 64):
+ m = solve_translation(stack, device="numpy", upsample=u, reference="first")
+ rows.append(f"| {u} | {np.abs(m.shifts - truth).max():.3f} px |")
+
+
+def bench_fixture(rows: list[str]) -> None:
+ import spyde.data.synthetic as sy
+ from spyde.drift.translation import solve_translation
+
+ rows.append("\n### Synthetic particle-movie fixture")
+ rows.append("")
+ build = _t(lambda: sy.particle_movie(), repeat=2)
+ s = sy.particle_movie()
+ gt = sy.ground_truth(s)
+ truth = np.asarray(gt["drift"])
+ m = solve_translation(s.data, device="numpy", upsample=8, reference="first",
+ max_shift=20)
+ err = np.abs(m.shifts - truth)
+ solve = _t(lambda: solve_translation(s.data, device="numpy", upsample=8,
+ reference="first", max_shift=20), repeat=2)
+ rows.append("| stage | value |")
+ rows.append("|---|---|")
+ rows.append(f"| build (24 x 96x112) | {build * 1e3:.0f} ms |")
+ rows.append(f"| drift solve | {solve * 1e3:.0f} ms |")
+ rows.append(f"| drift error (max / mean) | {err.max():.3f} / {err.mean():.3f} px |")
+ rows.append(f"| frames rejected from reference | "
+ f"{m.params['rejected_from_reference']} |")
+
+
+def bench_segment(rows: list[str]) -> None:
+ import spyde.data.synthetic as sy
+ from spyde.particles import SegmentParams, measure_frame, segment_frame
+
+ s = sy.particle_movie()
+ frame = s.data[12]
+ p = SegmentParams(min_size=25, gaussian=1.0)
+ seg = _t(lambda: segment_frame(frame, p))
+ labels = segment_frame(frame, p)
+ meas = _t(lambda: measure_frame(labels, frame, t=12, scale=0.5))
+ rows.append("\n### Classical segment + measure, one 96x112 frame")
+ rows.append("")
+ rows.append("| stage | time | frames/s |")
+ rows.append("|---|---|---|")
+ rows.append(f"| segment_frame | {seg * 1e3:.1f} ms | {1 / seg:.0f} |")
+ rows.append(f"| measure_frame | {meas * 1e3:.1f} ms | {1 / meas:.0f} |")
+ rows.append(f"| combined | {(seg + meas) * 1e3:.1f} ms | {1 / (seg + meas):.0f} |")
+ rows.append("")
+ rows.append(f"Extrapolated to 3000 frames: **{3000 * (seg + meas):.0f} s** "
+ f"single-threaded. The plan's target is minutes, so this is the "
+ f"number to watch as frame size grows (cost is per-pixel, and a "
+ f"4096^2 frame is 1600x this one's area).")
+
+
+def bench_optional(rows: list[str]) -> None:
+ """Stages whose modules may not have landed yet — reported as absent, not fatal."""
+ try:
+ from spyde.particles import features # noqa: F401
+ except Exception:
+ rows.append("\n### Feature stack — not implemented yet")
+ return
+ import spyde.data.synthetic as sy
+ from spyde.particles.features import FeatureSpec, compute_features
+
+ s = sy.particle_movie()
+ frame = s.data[12]
+ spec = FeatureSpec()
+ dt = _t(lambda: compute_features(frame, spec, device="cpu"))
+ rows.append("\n### Torch feature stack, one 96x112 frame (CPU)")
+ rows.append("")
+ rows.append(f"| compute_features | {dt * 1e3:.1f} ms |")
+ rows.append("|---|---|")
+
+
+def main() -> int:
+ rows: list[str] = ["# drift + particles benchmark", ""]
+ for fn in (bench_drift_accuracy, bench_drift_backends, bench_fixture,
+ bench_segment, bench_optional):
+ try:
+ fn(rows)
+ except Exception as exc:
+ rows.append(f"\n### {fn.__name__} FAILED: {type(exc).__name__}: {exc}")
+ text = "\n".join(rows) + "\n"
+ sys.stdout.write(text)
+ sys.stdout.flush()
+ # torch/CUDA teardown can crash the interpreter on Windows after CUDA work
+ # (CLAUDE.md); the output is already flushed, so leave immediately.
+ os._exit(0)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/scripts/check_drift_particles.py b/scripts/check_drift_particles.py
new file mode 100644
index 00000000..23790d27
--- /dev/null
+++ b/scripts/check_drift_particles.py
@@ -0,0 +1,195 @@
+#!/usr/bin/env python
+"""
+One command that verifies the whole drift + particles feature.
+
+(Named check_ rather than verify_ because .gitignore carries a `verify_*.py` rule
+for throwaway ad-hoc scripts, which silently swallowed this file for several
+commits — it ran fine locally and simply was not in the repo.)
+
+ python scripts/check_drift_particles.py # python only (fast)
+ python scripts/check_drift_particles.py --all # + typecheck + e2e
+ python scripts/check_drift_particles.py --e2e # + e2e only
+ python scripts/check_drift_particles.py --bench # + report the numbers
+
+Exists because this feature spans two languages, three test tiers and a separate
+repo, so "did I break it" was otherwise four commands with four different working
+directories and one of them needed a build step first. Re-run it after every step
+of DRIFT_AND_PARTICLES_PLAN.md.
+
+Every stage is INDEPENDENT and the exit code is the worst result, so one broken
+tier still reports the state of the others — the point is a status board, not a
+fail-fast gate. Stages that cannot run (no node_modules, no browser) are reported
+as SKIP, not as failure: a missing optional tool is not a regression.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import shutil
+import subprocess
+import sys
+import time
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parent.parent
+ELECTRON = ROOT / "electron"
+
+# The python suites this feature owns. Kept explicit rather than globbing
+# `spyde/tests/migrated` so a new unrelated failure elsewhere does not read as a
+# drift/particles regression — the full suite is a separate stage below.
+FEATURE_SUITES = [
+ "spyde/tests/migrated/test_drift_translation.py",
+ "spyde/tests/migrated/test_particles_core.py",
+ "spyde/tests/migrated/test_particle_movie_fixture.py",
+ "spyde/tests/migrated/test_particles_scribble.py",
+ "spyde/tests/migrated/test_particles_track.py",
+ "spyde/tests/migrated/test_particle_tree.py",
+ "spyde/tests/migrated/test_particles_wizard.py",
+ "spyde/tests/migrated/test_drift_wizard.py",
+ "spyde/tests/migrated/test_particle_overlay.py",
+ "spyde/tests/migrated/test_particle_lifecycle.py",
+]
+
+# Playwright specs this feature owns.
+FEATURE_SPECS = [
+ "tests/particles_workflow.spec.ts",
+ "tests/drift_workflow.spec.ts",
+]
+
+PASS, FAIL, SKIP = "PASS", "FAIL", "SKIP"
+_C = {PASS: "\033[32m", FAIL: "\033[31m", SKIP: "\033[33m"}
+_R = "\033[0m"
+
+
+def _colour(state: str) -> str:
+ if os.environ.get("NO_COLOR") or not sys.stdout.isatty():
+ return state
+ return f"{_C.get(state, '')}{state}{_R}"
+
+
+def _python() -> str:
+ """The venv interpreter if there is one, else whatever is running us."""
+ for rel in ("Scripts/python.exe", "bin/python"):
+ cand = ROOT / ".venv" / rel
+ if cand.exists():
+ return str(cand)
+ return sys.executable
+
+
+class Runner:
+ def __init__(self, verbose: bool) -> None:
+ self.verbose = verbose
+ self.results: list[tuple[str, str, float, str]] = []
+
+ def stage(self, name: str, cmd: list[str], *, cwd: Path = ROOT,
+ skip_if: str | None = None, env: dict[str, str] | None = None) -> str:
+ if skip_if:
+ self.results.append((name, SKIP, 0.0, skip_if))
+ print(f" {_colour(SKIP)} {name} ({skip_if})")
+ return SKIP
+ print(f" .... {name}", end="\r", flush=True)
+ full_env = {**os.environ, **(env or {})}
+ t0 = time.perf_counter()
+ proc = subprocess.run(cmd, cwd=str(cwd), env=full_env,
+ capture_output=not self.verbose, text=True)
+ dt = time.perf_counter() - t0
+ state = PASS if proc.returncode == 0 else FAIL
+ detail = ""
+ if state is FAIL and not self.verbose:
+ tail = ((proc.stdout or "") + (proc.stderr or "")).strip().splitlines()
+ detail = "\n".join(" " + ln for ln in tail[-25:])
+ self.results.append((name, state, dt, detail))
+ print(f" {_colour(state)} {name} ({dt:.1f}s) ")
+ if detail:
+ print(detail)
+ return state
+
+ def summary(self) -> int:
+ print("\n" + "=" * 66)
+ worst = 0
+ for name, state, dt, _ in self.results:
+ print(f" {_colour(state)} {name:<44} {dt:>6.1f}s")
+ if state is FAIL:
+ worst = 1
+ n_fail = sum(1 for _, s, _, _ in self.results if s is FAIL)
+ n_skip = sum(1 for _, s, _, _ in self.results if s is SKIP)
+ n_pass = sum(1 for _, s, _, _ in self.results if s is PASS)
+ print("=" * 66)
+ print(f" {n_pass} passed, {n_fail} failed, {n_skip} skipped")
+ return worst
+
+
+def main() -> int:
+ ap = argparse.ArgumentParser(description=__doc__,
+ formatter_class=argparse.RawDescriptionHelpFormatter)
+ ap.add_argument("--all", action="store_true", help="every stage")
+ ap.add_argument("--e2e", action="store_true", help="include Playwright specs")
+ ap.add_argument("--typecheck", action="store_true", help="include tsc")
+ ap.add_argument("--full-suite", action="store_true",
+ help="the whole python suite, not just this feature's")
+ ap.add_argument("--bench", action="store_true",
+ help="print the feature's benchmark numbers")
+ ap.add_argument("-v", "--verbose", action="store_true",
+ help="stream child output instead of capturing it")
+ args = ap.parse_args()
+ if args.all:
+ args.e2e = args.typecheck = args.full_suite = args.bench = True
+
+ py = _python()
+ r = Runner(args.verbose)
+ print(f"\nverifying drift + particles\n python: {py}\n")
+
+ # ── python ───────────────────────────────────────────────────────────────
+ existing = [s for s in FEATURE_SUITES if (ROOT / s).exists()]
+ missing = [s for s in FEATURE_SUITES if not (ROOT / s).exists()]
+ if existing:
+ r.stage("python feature suites",
+ [py, "-m", "pytest", *existing, "-q", "--no-header",
+ "-p", "no:cacheprovider"],
+ env={"SPYDE_NO_DASK": "1"})
+ for s in missing:
+ # Not yet written — a step of the plan that has not landed. Visible, not fatal.
+ r.stage(f"python {Path(s).stem}", [], skip_if="not implemented yet")
+
+ if args.full_suite:
+ r.stage("python full suite",
+ [py, "-m", "pytest", "-q", "--no-header", "-p", "no:cacheprovider"],
+ env={"SPYDE_NO_DASK": "1"})
+
+ # ── frontend ─────────────────────────────────────────────────────────────
+ has_npm = shutil.which("npm") is not None
+ has_modules = (ELECTRON / "node_modules").is_dir()
+ front_skip = (None if (has_npm and has_modules)
+ else "no npm" if not has_npm else "run npm install in electron/")
+
+ if args.typecheck:
+ r.stage("frontend typecheck", ["npm", "run", "typecheck"],
+ cwd=ELECTRON, skip_if=front_skip)
+
+ if args.e2e:
+ specs = [s for s in FEATURE_SPECS if (ELECTRON / s).exists()]
+ if front_skip:
+ r.stage("e2e feature specs", [], skip_if=front_skip)
+ elif not specs:
+ r.stage("e2e feature specs", [], skip_if="no specs written yet")
+ else:
+ # test:build, not test — the harness launches out/main/index.js, so a
+ # renderer change that was never built is invisible to Playwright and
+ # the spec silently tests the previous bundle.
+ r.stage("e2e feature specs",
+ ["npm", "run", "test:build", "--", "--project=electron",
+ "--reporter=line", "--retries=0", *specs],
+ cwd=ELECTRON)
+
+ # ── numbers ──────────────────────────────────────────────────────────────
+ if args.bench:
+ r.stage("bench drift + fixture",
+ [py, str(ROOT / "scripts" / "bench_drift_particles.py")],
+ skip_if=None if (ROOT / "scripts" / "bench_drift_particles.py").exists()
+ else "bench script not written yet")
+
+ return r.summary()
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/spyde/actions/drift_action.py b/spyde/actions/drift_action.py
new file mode 100644
index 00000000..a299a476
--- /dev/null
+++ b/spyde/actions/drift_action.py
@@ -0,0 +1,1464 @@
+"""
+drift_action.py — the Drift Correction wizard (``drift_`` staged actions).
+
+Plan A8, rewritten under plan §0.9a (*"the caret shows ONE control; everything
+else is Advanced"*) after the first review: **"way too complicated. Too many
+options. Information overload."**
+
+ drift_open caret mounted → Drift Check window, the alignment ROI on
+ the movie, and the first discovery preview
+ drift_close caret unmounted → tear all of it down
+ drift_set_method rigid | rigid+affine | non-rigid (lives in Advanced now)
+ drift_tune a toggle/parameter changed → re-run the discovery preview
+ drift_run solve the movie on a worker; opens the dy/dx window and
+ fills it progressively from the solver's ``on_shift``
+ drift_discard drop the solved model (and stop a solve in flight)
+ drift_commit add the LAZY corrected node to the tree
+
+**The caret carries the TASK, not the algorithm.** Its default face is two
+toggles and one button. Reference mode, sub-pixel factor, max shift,
+interpolation order and the model tabs are all real and all still here — they
+sit behind a collapsed *Advanced* in the caret, and the schema below (the one
+source of truth, mirrored by ``registry._WIZARD_SCHEMAS``) tags them so any
+host renders the same split. Nothing was deleted; provenance still records
+every parameter.
+
+**Discovery comes before commitment.** The centrepiece is a draggable
+rectangle on the movie plus a live drift-corrected sum of just that box over
+~20 frames (:data:`_PREVIEW_FRAMES`). A good landmark sums sharp, a bad one
+blurs, and the *gain* number (:func:`_gradient_energy` of the aligned sum over
+the raw sum, measured on the SAME pixels) puts a figure on it. So the user sees
+whether alignment works on a subset before paying for the whole movie — and the
+"Use ROI for alignment" toggle then feeds that exact rectangle to
+``solve_translation(roi=…)``, which is often the more CORRECT answer anyway:
+whole-frame correlation is contaminated by the sample's own motion (see
+``spyde/drift/translation.py``'s ``roi`` docs).
+
+**Geometry is in IMAGE PIXELS end to end.** anyplotlib's 2-D widgets report
+``x/y/w/h`` in image pixels with no scale/offset applied, and
+``solve_translation``'s ``roi=(y0, x0, h, w)`` is in pixels too, so the two meet
+with no conversion. Do not add one "for consistency" — see
+``spyde/actions/masks.py::_signal_k_grids`` for that bug class.
+
+**Two windows, each with one job.** The *Drift Check* window is the evidence:
+the whole-movie raw/corrected sums on top, the discovery pair (ROI raw vs ROI
+aligned) beneath. The *Drift dy/dx* window is the curve, opened when the solve
+starts and filled progressively from ``on_shift`` — it is a normal figure
+window, not caret furniture. Both are bare ``figure`` windows (NOT registered
+``Plot``s), so each registers a controller via ``own_window`` and keeps its
+figure referenced through ``figure_registry.keep_alive``, per
+``actions/README.md`` §6.
+
+**Nothing here materialises the movie.** ``solve_translation`` streams one
+frame at a time; the check sums stream over a bounded subset
+(:data:`_SUM_MAX_FRAMES`); the preview reads one full frame at a time and keeps
+only the small crop, under a byte cap (:data:`_PREVIEW_MAX_BYTES`); and
+``drift_commit`` adds a ``map_blocks`` node so the corrected movie is a lazy
+view, never a copy (plan §0.7). The corrected node is tagged ``local=True``
+because a rigid shift is exactly per-frame, which is what lets the existing
+``LocalTransformReader`` scrub it.
+"""
+from __future__ import annotations
+
+import logging
+import threading
+import time
+from typing import Any
+
+import numpy as np
+
+from spyde.actions.context import current_signal as _current_signal
+from spyde.actions.context import src_plot_tree as _src_plot_tree
+from spyde.actions.lifecycle import (
+ bump_generation, is_current, run_on_worker, show_tree_node,
+)
+from spyde.actions.wizard import WizardController
+from spyde.backend.ipc import emit, emit_error, emit_progress, emit_status
+
+log = logging.getLogger(__name__)
+
+#: Solver families. ``rigid`` is the only one ``spyde.drift`` implements today;
+#: the other two are declared so a host can render the choice, and both fall
+#: back to ``rigid`` with an explicit status rather than silently doing
+#: something the user did not ask for. They live inside Advanced now — three
+#: visible tabs for two choices that do not work was two thirds of a control
+#: row spent on nothing (§0.9a).
+METHODS: tuple[str, ...] = ("rigid", "rigid_affine", "nonrigid")
+
+_UNAVAILABLE = {
+ "rigid_affine": ("the affine drift search (plan A4) is not implemented in "
+ "spyde.drift yet"),
+}
+
+#: The two non-rigid parameterisations (:mod:`spyde.drift.nonrigid`). Scan-knot
+#: describes a SCANNING artifact (displacement varies down the slow axis only);
+#: dense describes the SAMPLE deforming (varies in both directions). Both causes
+#: are real, which is why this is a choice and not an assumption.
+NONRIGID_MODELS: tuple[str, ...] = ("scan_knot", "dense")
+
+#: Longest side the non-rigid fit sees. The stack CANNOT be held at full size —
+#: 300 x 4096² float32 is 20 GB — and it does not need to be: a drift field is
+#: smooth by construction, which is the whole modelling assumption, so it is
+#: measurable from a decimated copy. The fit's parameters are resolution-
+#: independent (fractions of the frame), so the field rebuilds at full size for
+#: the apply. Numbers in benchmarks.md.
+_NONRIGID_FIT_SIDE = 512
+
+#: Frames summed for the whole-movie before/after check images. A sum is a
+#: SHARPNESS test, not a measurement — a few dozen frames already show the blur
+#: unambiguously, and the cap is what keeps the check window responsive on a
+#: movie whose full pass costs as much as the solve itself. Evenly spaced, and
+#: the SAME indices for both sums, or the comparison means nothing.
+_SUM_MAX_FRAMES = 64
+
+# Frames per streamed drift-trace message / per dy-dx repaint. One message per
+# frame would flood the PLOTAPP line protocol at the plan's target scale
+# (thousands of frames) for a curve the eye cannot follow at that resolution;
+# batching by 16 keeps the trace visibly live while cutting the message count
+# by the same factor.
+_TRACE_BATCH = 16
+
+# …but a COUNT alone is not enough, and the first screenshot showed why: a
+# 12-frame movie never reaches 16, so the curve stayed empty for the whole solve
+# and appeared complete at the end — the exact opposite of "fills in as it is
+# computed". Flush on whichever comes first, so the trace is live at any movie
+# length and still capped at ~7 messages/s on a fast one.
+_TRACE_MAX_INTERVAL = 0.15
+
+#: Figure geometry for the two bare-figure windows.
+#:
+#: A bare figure never receives ``resize_figure`` (that path resolves a
+#: registered ``Plot``), so its INITIAL px size is the one it keeps and anything
+#: outside it is CLIPPED by the subwindow — which is what cut the check
+#: window's bottom row in half. The renderer sizes a new window from the
+#: ``aspect`` field as ``inner_h = clamp(460 / aspect, 130, 300)`` then
+#: ``inner_w = inner_h * aspect`` (``MDIArea.windowSize``). At the height cap
+#: the first clamp is active for any aspect below 460/300, so a figure exactly
+#: :data:`_FIG_HEIGHT` tall lands pixel-for-pixel in its window at any width up
+#: to 460 — pick the width, derive the aspect.
+#:
+#: The width is deliberately the renderer's OWN default (340). Widening the
+#: check window to 460 made it no longer fit beside the movie, so the free-slot
+#: packer wrapped it to the next row — straight on top of the caret, which is an
+#: overlay the packer cannot see. Keeping the default width keeps the placement
+#: the packer already gets right.
+_FIG_WIDTH = 340
+_FIG_HEIGHT = 300
+
+
+def _figure_geometry(width: int = _FIG_WIDTH) -> tuple[tuple[int, int], float]:
+ """``(figsize, aspect)`` that opens a bare-figure window with no clipping."""
+ w = int(min(460, max(190, width)))
+ return (w, _FIG_HEIGHT), w / float(_FIG_HEIGHT)
+
+
+#: Frames the discovery preview aligns. ~20 is the brief's number and it is a
+#: DEFAULT, not a law — ``preview_frames`` in Advanced moves it.
+#:
+#: Sampled EVENLY OVER THE WHOLE MOVIE, not the first 20 in a row. The question
+#: a preview answers is "does this landmark survive the FULL excursion", and 20
+#: consecutive frames of a 3000-frame movie drift by almost nothing — a
+#: contiguous window would answer "looks fine" for every box, including the
+#: useless ones. The same reasoning (and the same spacing) as
+#: :meth:`DriftWizard.sum_indices`.
+_PREVIEW_FRAMES = 20
+
+#: Byte ceiling on the preview's retained crop stack. The preview reads one
+#: FULL frame at a time and keeps only the (usually small) ROI crop, so this
+#: bounds the only thing that accumulates. With no ROI the crop IS the frame,
+#: which is how 20 frames of a 4096² movie would otherwise become 1.3 GB;
+#: over the cap the sampled frame count is thinned rather than the read being
+#: abandoned. Never a reason to touch the full dataset (CLAUDE.md).
+_PREVIEW_MAX_BYTES = 192 * 1024 * 1024
+
+#: Settle delay for a preview re-solve driven by an ROI DRAG. The widget's
+#: pointer_move fires at renderer frame rate; re-solving 20 frames on each one
+#: would queue solves faster than they finish. ``drift_tune`` is NOT debounced
+#: here — the renderer's ``useDebouncedAction`` already settles it, and
+#: debouncing twice just adds latency.
+_PREVIEW_SETTLE_S = 0.25
+
+#: Smallest alignment box, in image pixels. MUST stay >= the solver's own
+#: ``spyde.drift.translation._MIN_ROI``, which REJECTS a smaller box rather
+#: than clamping it (a silently shrunk ROI would correlate somewhere the user
+#: did not drag). Pinned by ``test_drift_wizard.py``.
+_ROI_MIN_PX = 16
+
+#: Default alignment box: this fraction of each frame dimension, centred. Half
+#: the frame is deliberately generous — the ROI is FIXED in frame coordinates,
+#: so the landmark drifts within it and the box wants to be comfortably larger
+#: than the total excursion.
+_ROI_DEFAULT_FRACTION = 0.5
+
+_ROI_COLOR = "#94e2d5"
+
+#: **Off by default, and that is a measurement, not caution.** A guessed centre
+#: box is NOT automatically the better correlation: on the ``particle_movie``
+#: fixture (96×112 frames, the default half-frame box = 48×56) the ROI solve
+#: comes back 1.03 px from the stamped ground truth where the whole-frame solve
+#: is 0.25 px — a quarter of the pixels is a quarter of the correlation signal,
+#: and the Tukey taper eats a larger fraction of a small box. So the default
+#: stays the answer we already know is right, and the ROI is what the user
+#: reaches for when the whole frame is the problem (a moving sample, a mostly
+#: featureless field). The preview runs on the box either way — that is the
+#: discovery step, and it is what tells you the box is worth committing to.
+DEFAULTS: dict[str, Any] = dict(
+ use_roi=False,
+ reject_outliers=True,
+ method="rigid",
+ nonrigid_model="scan_knot",
+ nonrigid_steps=120,
+ upsample=8,
+ max_shift=32.0,
+ reference="running",
+ apodize=True,
+ normalize=True,
+ order=1,
+ preview_frames=_PREVIEW_FRAMES,
+)
+
+
+class DriftWizard(WizardController):
+ """Owns the drift caret's state: parameters, the alignment ROI and its live
+ preview, the solved model, and the two figure windows."""
+
+ key = "drift"
+
+ #: One source of truth (mirrored by ``registry._WIZARD_SCHEMAS``). Entries
+ #: WITHOUT a ``tab`` are the caret's default face; everything tagged
+ #: ``"Advanced"`` renders behind the collapsed disclosure (§0.9a).
+ parameters = {
+ "use_roi": {
+ "name": "Use ROI for alignment", "type": "bool",
+ "default": DEFAULTS["use_roi"],
+ },
+ "reject_outliers": {
+ "name": "Ignore bad frames", "type": "bool",
+ "default": DEFAULTS["reject_outliers"],
+ },
+ "method": {
+ "name": "Model", "type": "enum", "default": DEFAULTS["method"],
+ "choices": list(METHODS), "tab": "Advanced",
+ },
+ "nonrigid_model": {
+ "name": "Non-rigid field", "type": "enum",
+ "default": DEFAULTS["nonrigid_model"],
+ "choices": list(NONRIGID_MODELS), "tab": "Advanced",
+ },
+ "nonrigid_steps": {
+ "name": "Non-rigid steps", "type": "int",
+ "default": DEFAULTS["nonrigid_steps"],
+ "min": 10, "max": 1000, "tab": "Advanced",
+ },
+ "reference": {
+ "name": "Reference", "type": "enum", "default": DEFAULTS["reference"],
+ "choices": ["running", "sequential", "first"], "tab": "Advanced",
+ },
+ "upsample": {
+ "name": "Sub-pixel factor", "type": "int", "default": DEFAULTS["upsample"],
+ "min": 1, "max": 64, "tab": "Advanced",
+ },
+ "max_shift": {
+ "name": "Max shift (px)", "type": "float", "default": DEFAULTS["max_shift"],
+ "min": 1.0, "max": 4096.0, "step": 1.0, "tab": "Advanced",
+ },
+ "apodize": {
+ "name": "Edge taper", "type": "bool", "default": DEFAULTS["apodize"],
+ "tab": "Advanced",
+ },
+ "normalize": {
+ "name": "Phase correlation", "type": "bool",
+ "default": DEFAULTS["normalize"], "tab": "Advanced",
+ },
+ "order": {
+ "name": "Interpolation order", "type": "int", "default": DEFAULTS["order"],
+ "min": 0, "max": 3, "tab": "Advanced",
+ },
+ "preview_frames": {
+ "name": "Preview frames", "type": "int",
+ "default": DEFAULTS["preview_frames"], "min": 4, "max": 200,
+ "tab": "Advanced",
+ },
+ }
+
+ def __init__(self, session, tree, src_plot):
+ super().__init__(session, tree)
+ self.src_plot = src_plot
+ self.src_window_id = getattr(src_plot, "window_id", None)
+ self.params: dict[str, Any] = dict(DEFAULTS)
+ self.model = None
+ #: The Drift Check window (a bare figure) and its four panels.
+ self.window_id: int | None = None
+ self._panels: dict[str, Any] = {}
+ self._sum_indices: np.ndarray | None = None
+ self._before_sum: np.ndarray | None = None
+ #: The dy/dx window — opened by the solve, filled from ``on_shift``.
+ self.trace_window_id: int | None = None
+ self._trace: dict[str, Any] = {}
+ #: The alignment ROI (discovery): widget + last preview result.
+ self._roi_widget = None
+ self._roi_handler = None
+ self._roi_clamping = False
+ self._frame_shape: tuple[int, int] | None = None
+ self._settle: threading.Timer | None = None
+ self.preview: dict[str, Any] | None = None
+ #: Cancel flag of the solve in flight (Discard/Stop flips it).
+ self._stop: list[bool] = [False]
+
+ # ── the movie ────────────────────────────────────────────────────────────
+
+ def signal(self):
+ return _current_signal(self.src_plot) or self.tree.root
+
+ def frames(self):
+ """``(n_frames, get_frame, (h, w))`` — one frame at a time."""
+ from spyde.drift import frame_source
+ return frame_source(self.signal())
+
+ def sum_indices(self, n_frames: int) -> np.ndarray:
+ if self._sum_indices is None or self._sum_indices.size == 0:
+ k = min(int(n_frames), _SUM_MAX_FRAMES)
+ self._sum_indices = np.unique(
+ np.linspace(0, max(0, n_frames - 1), max(1, k)).round().astype(int))
+ return self._sum_indices
+
+ # ── the alignment ROI (the discovery feature) ────────────────────────────
+
+ def _plot2d(self):
+ return getattr(self.src_plot, "_plot2d", None) if self.src_plot else None
+
+ def ensure_roi_widget(self, shape: tuple[int, int]) -> None:
+ """Draw the draggable alignment box on the source movie (idempotent).
+
+ Geometry is IMAGE PIXELS — anyplotlib 2-D widgets report ``x/y/w/h``
+ that way, and that is exactly what ``solve_translation(roi=…)`` wants.
+ A raw ``add_rectangle_widget`` rather than ``RectangleSelector``: the
+ selector caps itself at ``MAX_REGION_EXTENT_PER_DIM`` (16 px) because
+ it drives a nav-space region integrate, and a 16 px alignment box is
+ below the solver's own floor.
+ """
+ h, w = int(shape[0]), int(shape[1])
+ self._frame_shape = (h, w)
+ if self._roi_widget is not None:
+ return
+ plot2d = self._plot2d()
+ if plot2d is None:
+ return
+ if min(h, w) < 2 * _ROI_MIN_PX:
+ # Nothing sensible to drag; the whole frame IS the ROI.
+ return
+ bw = max(_ROI_MIN_PX, min(w, int(round(w * _ROI_DEFAULT_FRACTION))))
+ bh = max(_ROI_MIN_PX, min(h, int(round(h * _ROI_DEFAULT_FRACTION))))
+ try:
+ widget = plot2d.add_rectangle_widget(
+ x=float((w - bw) // 2), y=float((h - bh) // 2),
+ w=float(bw), h=float(bh), color=_ROI_COLOR, show_handles=True,
+ )
+ from spyde.drawing.selectors.base_selector import event_handler_fn
+ handler = event_handler_fn(lambda event: self._on_roi_drag())
+ widget.add_event_handler(handler, "pointer_move", "pointer_up")
+ self._roi_widget = widget
+ self._roi_handler = handler # keep a ref alive (weak callback)
+ except Exception as exc:
+ log.debug("[drift] alignment ROI widget failed: %s", exc)
+
+ def _on_roi_drag(self) -> None:
+ """Clamp the box to the frame, then arm the settle timer.
+
+ RE-ENTRANCY GUARD: anyplotlib ``Widget.set()`` fires ``pointer_move``
+ UNCONDITIONALLY (even on a no-change write), so the clamp below
+ re-invokes this handler synchronously — unguarded, ONE JS drag frame
+ recursed ~2000 deep before RecursionError in the Crop box (see
+ ``actions/base.py``). A hard flag breaks the cycle; compare-before-set
+ is NOT sufficient.
+ """
+ if self._roi_clamping or self._closed:
+ return
+ self._roi_clamping = True
+ try:
+ self._clamp_roi()
+ finally:
+ self._roi_clamping = False
+ self.schedule_preview()
+
+ def _clamp_roi(self) -> None:
+ """Keep the box inside the frame and above the solver's floor.
+
+ COMPARE BEFORE SET, with slack. ``Widget.set()`` pushes geometry back to
+ the renderer, and writing on every ``pointer_move`` echoes
+ python-sourced geometry into a live drag — the same failure the 1-D span
+ cap documents in CLAUDE.md (Live-Display §3). A box already resting on a
+ bound must be left alone.
+ """
+ widget, shape = self._roi_widget, self._frame_shape
+ if widget is None or shape is None:
+ return
+ h, w = shape
+ try:
+ ww = min(max(float(widget.w), float(_ROI_MIN_PX)), float(w))
+ hh = min(max(float(widget.h), float(_ROI_MIN_PX)), float(h))
+ x = min(max(float(widget.x), 0.0), float(w) - ww)
+ y = min(max(float(widget.y), 0.0), float(h) - hh)
+ now = (float(widget.x), float(widget.y),
+ float(widget.w), float(widget.h))
+ if max(abs(a - b) for a, b in zip(now, (x, y, ww, hh))) > 1e-6:
+ widget.set(x=x, y=y, w=ww, h=hh)
+ except Exception as exc:
+ log.debug("[drift] clamping the alignment ROI failed: %s", exc)
+
+ def roi_box(self) -> tuple[int, int, int, int] | None:
+ """``(y0, x0, h, w)`` in IMAGE PIXELS, or None when there is no usable
+ box — the shape ``solve_translation``'s ``roi`` takes, with no scale or
+ offset applied because neither side has any."""
+ widget, shape = self._roi_widget, self._frame_shape
+ if widget is None or shape is None:
+ return None
+ fh, fw = shape
+ try:
+ x0 = int(round(float(widget.x)))
+ y0 = int(round(float(widget.y)))
+ bw = int(round(float(widget.w)))
+ bh = int(round(float(widget.h)))
+ except Exception as exc:
+ log.debug("[drift] reading the alignment ROI failed: %s", exc)
+ return None
+ # SIZE first, then origin. The other order looks equivalent and is not:
+ # clamping the origin to the frame edge and only then applying the
+ # minimum size pushes the box back OUT past the edge, and
+ # solve_translation rejects an out-of-frame roi outright.
+ bw = max(_ROI_MIN_PX, min(bw, fw))
+ bh = max(_ROI_MIN_PX, min(bh, fh))
+ if bw > fw or bh > fh:
+ return None # frame smaller than the solver's floor
+ x0 = max(0, min(x0, fw - bw))
+ y0 = max(0, min(y0, fh - bh))
+ return (y0, x0, bh, bw)
+
+ def active_roi(self) -> tuple[int, int, int, int] | None:
+ """The ROI the FULL SOLVE should use — None unless the toggle is on (and
+ None when no box is usable, which is a whole-frame correlation).
+
+ The preview deliberately does NOT go through here: it always aligns the
+ box, toggle or not, because that is the question it exists to answer
+ ("is this landmark worth committing to?"). The toggle is the commitment.
+ """
+ return self.roi_box() if self.params.get("use_roi") else None
+
+ def remove_roi_widget(self) -> None:
+ widget, self._roi_widget = self._roi_widget, None
+ self._roi_handler = None
+ if widget is not None:
+ try:
+ widget.hide() # widgets have no remove(), only hide()
+ except Exception as exc:
+ log.debug("[drift] hiding the alignment ROI failed: %s", exc)
+
+ # ── preview scheduling (latest-wins, cancellable) ────────────────────────
+
+ def schedule_preview(self, delay: float = _PREVIEW_SETTLE_S) -> None:
+ """(Re-)arm the settle timer for a drag-driven preview re-solve.
+
+ Latest-wins in two places: the timer is restarted per pointer event so
+ only the RESTING geometry ever solves, and the solve that does run
+ carries a ``_drift_preview_gen`` generation so a superseded result is
+ dropped on arrival instead of painting over a newer one.
+ """
+ self.cancel_preview()
+ if self._closed:
+ return
+ timer = threading.Timer(max(0.0, float(delay)), self._fire_preview)
+ timer.daemon = True
+ self._settle = timer
+ timer.start()
+
+ def cancel_preview(self) -> None:
+ timer, self._settle = self._settle, None
+ if timer is not None:
+ try:
+ timer.cancel()
+ except Exception as exc:
+ log.debug("[drift] cancelling the preview timer failed: %s", exc)
+
+ def _fire_preview(self) -> None:
+ """Timer thread → main thread → the worker. Reading widget geometry and
+ spawning the compute both belong on the main thread (thread marshal,
+ README §6); only the arithmetic runs on the worker."""
+ self._settle = None
+ if self._closed:
+ return
+ dispatch = getattr(self.session, "_dispatch_to_main", None)
+ if dispatch is None:
+ _run_preview(self)
+ else:
+ dispatch(lambda: (None if self._closed else _run_preview(self)))
+
+ # ── the check window ─────────────────────────────────────────────────────
+
+ def open_check_window(self, before: np.ndarray, n_frames: int) -> None:
+ """Emit the bare-figure Drift Check window and register this controller
+ for it, so ✕ and ``Session._forget_window`` reach the wizard.
+
+ Top row = the whole movie, raw and corrected (the solve's evidence).
+ Bottom row = the DISCOVERY pair, raw and aligned over ~20 frames of
+ whatever is being correlated (the ROI, or the whole frame when the
+ toggle is off). Side by side, because "is this landmark good" is
+ answered by comparing two sums, not by staring at one.
+ """
+ import anyplotlib as apl
+ import anyplotlib._electron as _electron
+ from spyde.actions.figure_registry import keep_alive
+ from spyde.drawing.plots.plot import finalize_figure_html
+
+ figsize, aspect = _figure_geometry()
+ fig, axes = apl.subplots(2, 2, figsize=figsize)
+ ax = np.array(axes, dtype=object).ravel()
+ before = np.asarray(before, np.float32)
+ zeros = np.zeros_like(before)
+
+ panels = {
+ "before": ax[0].imshow(before, cmap="gray"),
+ "after": ax[1].imshow(zeros, cmap="gray"),
+ "roi_raw": ax[2].imshow(zeros, cmap="gray"),
+ "roi_aligned": ax[3].imshow(zeros, cmap="gray"),
+ }
+ titles = {"before": "Raw sum", "after": "Corrected sum",
+ "roi_raw": "ROI raw", "roi_aligned": "ROI aligned"}
+ self._panels = panels
+ for key, title in titles.items():
+ self._set_panel_title(key, title)
+
+ wid = self.session.next_window_id()
+ fig_id = _electron.register(fig)
+ html = finalize_figure_html(fig, fig_id)
+ keep_alive(int(wid), fig)
+ emit({"type": "figure", "fig_id": fig_id, "window_id": int(wid),
+ "html": html, "title": "Drift Check", "is_navigator": False,
+ "aspect": float(aspect)})
+ self.window_id = int(wid)
+ self._before_sum = before
+ self.own_window(wid)
+
+ def _set_panel_title(self, key: str, title: str) -> None:
+ panel = self._panels.get(key)
+ if panel is None:
+ return
+ try:
+ panel.set_title(title)
+ except Exception as exc:
+ log.debug("[drift] set_title(%s) failed: %s", key, exc)
+
+ def update_check(self, *, after=None) -> None:
+ """Paint the whole-movie corrected sum (main thread only)."""
+ if after is None or not self._panels:
+ return
+ try:
+ self._panels["after"].set_data(np.asarray(after, np.float32))
+ except Exception as exc:
+ log.debug("[drift] painting the corrected sum failed: %s", exc)
+
+ def show_preview(self, result: dict) -> None:
+ """Paint the discovery pair + its titles (main thread only)."""
+ if not self._panels:
+ return
+ what = "ROI" if result.get("roi") is not None else "Whole frame"
+ n = int(result.get("frames", 0))
+ gain = float(result.get("gain", float("nan")))
+ for key, arr, title in (
+ ("roi_raw", result.get("raw"), f"{what} raw · {n} frames"),
+ ("roi_aligned", result.get("aligned"),
+ f"{what} aligned · {gain:.1f}x sharper" if np.isfinite(gain)
+ else f"{what} aligned"),
+ ):
+ if arr is None:
+ continue
+ try:
+ self._panels[key].set_data(np.asarray(arr, np.float32))
+ except Exception as exc:
+ log.debug("[drift] painting the %s panel failed: %s", key, exc)
+ self._set_panel_title(key, title)
+
+ # ── the dy/dx window ─────────────────────────────────────────────────────
+
+ def open_trace_window(self, n_frames: int) -> None:
+ """Open (or reset) the dy/dx figure window.
+
+ Its OWN window, not caret furniture: the curve is the measurement, and
+ a 40 px inline sparkline could show that the stage crept but never
+ which frame jumped. One panel with two labelled lines rather than two
+ panels — drift is anisotropic, and a shared y-scale is what makes
+ "mostly x" readable at a glance.
+ """
+ n = max(2, int(n_frames))
+ if self.trace_window_id is not None and self._trace:
+ self.reset_trace(n)
+ return
+ import anyplotlib as apl
+ import anyplotlib._electron as _electron
+ from spyde.actions.figure_registry import keep_alive
+ from spyde.drawing.plots.plot import finalize_figure_html
+
+ figsize, aspect = _figure_geometry()
+ fig, axes = apl.subplots(1, 1, figsize=figsize)
+ ax = np.array(axes, dtype=object).ravel()[0]
+ x0 = np.zeros(1, dtype=np.float64)
+ y0 = np.zeros(1, dtype=np.float64)
+ panel = ax.plot(y0, axes=[x0], units="frame", y_units="shift (px)",
+ color="#89b4fa", label="dy")
+ dx_line = panel.add_line(y0, x_axis=x0, color="#f38ba8", label="dx")
+ for setter, text in (("set_title", "Drift dy / dx"),
+ ("set_xlabel", "frame"),
+ ("set_ylabel", "shift (px)")):
+ try:
+ getattr(panel, setter)(text)
+ except Exception as exc:
+ log.debug("[drift] trace %s failed: %s", setter, exc)
+
+ wid = self.session.next_window_id()
+ fig_id = _electron.register(fig)
+ html = finalize_figure_html(fig, fig_id)
+ keep_alive(int(wid), fig)
+ emit({"type": "figure", "fig_id": fig_id, "window_id": int(wid),
+ "html": html, "title": "Drift dy/dx", "is_navigator": False,
+ "aspect": float(aspect)})
+ self.trace_window_id = int(wid)
+ self._trace = {"panel": panel, "dx": dx_line}
+ self.reset_trace(n)
+ self.own_window(wid)
+
+ def reset_trace(self, n_frames: int) -> None:
+ n = max(2, int(n_frames))
+ self._trace["dy_data"] = np.full(n, np.nan, np.float64)
+ self._trace["dx_data"] = np.full(n, np.nan, np.float64)
+ self._trace["filled"] = 0
+ self.push_trace([(0, 0.0, 0.0)])
+
+ def push_trace(self, points) -> None:
+ """Append a batch of ``(index, dy, dx)`` and repaint (main thread only).
+
+ Only the SOLVED PREFIX is pushed, so the curve grows left to right and
+ the y-scale tracks what has actually been measured — pushing the whole
+ NaN-padded array would make anyplotlib's auto-range see one point.
+ """
+ if not self._trace:
+ return
+ dy = self._trace.get("dy_data")
+ dx = self._trace.get("dx_data")
+ if dy is None or dx is None:
+ return
+ hi = int(self._trace.get("filled", 0))
+ for i, y, x in points:
+ i = int(i)
+ if 0 <= i < dy.size:
+ dy[i] = float(y)
+ dx[i] = float(x)
+ hi = max(hi, i + 1)
+ self._trace["filled"] = hi
+ if hi < 1:
+ return
+ xs = np.arange(hi, dtype=np.float64)
+ try:
+ self._trace["panel"].set_data(np.nan_to_num(dy[:hi]), x_axis=xs)
+ self._trace["dx"].set_data(np.nan_to_num(dx[:hi]), x_axis=xs)
+ except Exception as exc:
+ log.debug("[drift] painting the dy/dx trace failed: %s", exc)
+
+ def close_trace_window(self) -> None:
+ self._trace = {}
+ wid, self.trace_window_id = self.trace_window_id, None
+ self._close_window(wid)
+
+ # ── lifecycle ────────────────────────────────────────────────────────────
+
+ def _close_window(self, wid: int | None) -> None:
+ if wid is None:
+ return
+ forget = getattr(self.session, "_forget_window", None)
+ if forget is not None:
+ try:
+ forget(int(wid))
+ except Exception as exc:
+ log.debug("[drift] forgetting window %s failed: %s", wid, exc)
+ return
+ # Bare / stub session: emit + unregister by hand.
+ try:
+ emit({"type": "window_closed", "window_id": int(wid)})
+ except Exception as exc:
+ log.debug("[drift] closing window %s failed: %s", wid, exc)
+ reg = getattr(self.session, "_window_controllers", None)
+ if isinstance(reg, dict):
+ reg.pop(int(wid), None)
+
+ def close(self) -> None:
+ """WindowController protocol — ``Session._forget_window`` calls this for
+ EITHER owned window, with no way to say which.
+
+ Only the Drift Check window is the wizard's life; closing the dy/dx
+ window just drops the curve. ``_forget_window`` pops the controller for
+ the window that went away BEFORE calling here, so "is the check window
+ still registered?" identifies it exactly — and the programmatic path
+ (:meth:`close_trace_window`) clears ``trace_window_id`` first, so this
+ re-entry is a no-op rather than a recursion.
+ """
+ reg = getattr(self.session, "_window_controllers", None) or {}
+ if self.window_id is not None and reg.get(int(self.window_id)) is self:
+ self._trace = {}
+ self.trace_window_id = None
+ return
+ self.remove()
+
+ def remove(self) -> None:
+ if self._closed:
+ return
+ self._closed = True
+ self._stop[0] = True # stop a solve in flight
+ self.cancel_preview()
+ self.remove_roi_widget()
+ self._panels = {}
+ self._trace = {}
+ wid, self.window_id = self.window_id, None
+ twid, self.trace_window_id = self.trace_window_id, None
+ self._close_window(wid)
+ self._close_window(twid)
+ if getattr(self.tree, "_drift_wizard", None) is self:
+ self.tree._drift_wizard = None
+
+ def commit(self):
+ """Add the lazy corrected node — see :func:`drift_commit`."""
+ return _commit(self)
+
+
+# ── parameters ───────────────────────────────────────────────────────────────
+
+def _coerce(payload: dict | None) -> dict:
+ p = dict(DEFAULTS)
+ payload = payload or {}
+ for k, default in DEFAULTS.items():
+ v = payload.get(k)
+ if v is None or v == "":
+ continue
+ try:
+ p[k] = bool(v) if isinstance(default, bool) else type(default)(v)
+ except (TypeError, ValueError) as exc:
+ log.debug("[drift] param %r=%r not coercible, keeping default: %s",
+ k, v, exc)
+ p["method"] = str(p["method"]).lower()
+ if p["method"] not in METHODS:
+ p["method"] = DEFAULTS["method"]
+ p["nonrigid_model"] = str(p["nonrigid_model"]).lower()
+ if p["nonrigid_model"] not in NONRIGID_MODELS:
+ p["nonrigid_model"] = DEFAULTS["nonrigid_model"]
+ p["nonrigid_steps"] = int(min(1000, max(10, p["nonrigid_steps"])))
+ if p["reference"] not in ("running", "sequential", "first"):
+ p["reference"] = DEFAULTS["reference"]
+ p["upsample"] = max(1, int(p["upsample"]))
+ p["max_shift"] = max(1.0, float(p["max_shift"]))
+ p["order"] = int(min(3, max(0, p["order"])))
+ p["preview_frames"] = int(min(200, max(4, p["preview_frames"])))
+ return p
+
+
+def _decimated_stack(get_frame, n_frames: int, side: int = _NONRIGID_FIT_SIDE,
+ cancel=None) -> np.ndarray:
+ """Read every frame at reduced resolution for the non-rigid fit.
+
+ Strided, not area-averaged. The fit needs STRUCTURE to correlate, and a
+ stride keeps edges crisp where a box mean blurs exactly the gradients the
+ solve reads; it is also ~free where the mean costs a pass over every pixel
+ (the same reasoning as the navigator's base-frame subsample, Live-Display
+ §3). Frames are read ONE AT A TIME and decimated immediately, so the full
+ stack is never held — 300 x 4096² float32 would be 20 GB.
+ """
+ out: list[np.ndarray] = []
+ for i in range(int(n_frames)):
+ if cancel is not None and cancel():
+ break
+ f = np.asarray(get_frame(int(i)), np.float32)
+ step = max(1, int(max(f.shape) // max(1, side)))
+ out.append(np.ascontiguousarray(f[::step, ::step], dtype=np.float32))
+ if not out:
+ return np.zeros((0, 1, 1), np.float32)
+ return np.stack(out)
+
+
+def _solve_nonrigid_step(wiz, p: dict, rigid_model, get_frame, n_frames: int,
+ *, progress=None, cancel=None):
+ """Fit the non-rigid residual on top of a solved rigid model.
+
+ Returns the non-rigid :class:`DriftModel` (which carries the rigid
+ ``shifts`` forward unchanged), or the rigid model untouched if the fit
+ cannot run. A failure here must NOT lose the rigid answer — that is a good
+ result the user already waited for.
+ """
+ from spyde.drift import solve_nonrigid
+
+ t0 = time.monotonic()
+ stack = _decimated_stack(get_frame, n_frames, cancel=cancel)
+ if stack.shape[0] < 2:
+ return rigid_model
+ log.info("[drift] non-rigid fit: %s, %d frames at %dx%d (decimated from %s)",
+ p["nonrigid_model"], stack.shape[0], stack.shape[1], stack.shape[2],
+ "the movie")
+ try:
+ model = solve_nonrigid(
+ stack,
+ model=str(p["nonrigid_model"]),
+ rigid=rigid_model,
+ steps=int(p["nonrigid_steps"]),
+ device=None,
+ progress=progress,
+ cancel=cancel,
+ provenance={"action": "Drift Correction (non-rigid)",
+ "params": dict(p)},
+ )
+ # Say what actually happened, with the KIND. Without this there is no
+ # way — from a log, a test, or the UI — to tell a real non-rigid solve
+ # from one that fell back to rigid, and "the caret said non-rigid" is
+ # exactly the claim that must not be taken on trust.
+ log.info("[drift] non-rigid fit done: kind=%s in %.2fs (mse %.4g)",
+ model.kind, time.monotonic() - t0,
+ model.extra.get("final_mse", float("nan")))
+ return model
+ except Exception as exc:
+ # torch missing, CUDA OOM, a device that will not take the graph — all
+ # end here. Say so and keep the rigid model rather than failing the run.
+ log.warning("[drift] non-rigid fit failed (%s); keeping the rigid model", exc)
+ emit_status(f"Drift Correction: non-rigid fit failed ({exc}) — "
+ "keeping the rigid result.")
+ return rigid_model
+
+
+def _solver_kwargs(p: dict, roi=None) -> dict:
+ return dict(upsample=int(p["upsample"]), max_shift=float(p["max_shift"]),
+ reference=str(p["reference"]), apodize=bool(p["apodize"]),
+ normalize=bool(p["normalize"]),
+ reject_outliers=bool(p["reject_outliers"]),
+ roi=None if roi is None else tuple(int(v) for v in roi))
+
+
+def _wizard(session, plot) -> DriftWizard | None:
+ """Resolve the live wizard from any of its windows.
+
+ The check and dy/dx windows are bare figures, so ``_plot_by_window_id``
+ returns None for them and the plot-based lookup finds nothing — resolve by
+ window id through the controller registry first (README §6), then fall back
+ to the source tree's back-reference.
+ """
+ wid = getattr(plot, "window_id", None) if plot is not None else None
+ lookup = getattr(session, "controller_by_window_id", None)
+ if wid is not None and lookup is not None:
+ ctrl = lookup(int(wid))
+ if isinstance(ctrl, DriftWizard) and not ctrl._closed:
+ return ctrl
+ _src, tree = _src_plot_tree(session, plot)
+ wiz = getattr(tree, "_drift_wizard", None) if tree is not None else None
+ return wiz if (wiz is not None and not wiz._closed) else None
+
+
+def _emit_state(wiz: DriftWizard, **extra) -> None:
+ roi = wiz.roi_box()
+ msg = {"type": "drift_state",
+ "window_id": wiz.src_window_id,
+ "check_window_id": wiz.window_id,
+ "trace_window_id": wiz.trace_window_id,
+ "method": wiz.params["method"],
+ "solved": wiz.model is not None,
+ "use_roi": bool(wiz.params["use_roi"]),
+ "roi": None if roi is None else [int(v) for v in roi],
+ "params": dict(wiz.params)}
+ msg.update(extra)
+ emit(msg)
+
+
+# ── streaming sums + the sharpness number ────────────────────────────────────
+
+def _stack_sum(get_frame, indices, shifts=None, *, order: int = 1) -> np.ndarray:
+ """Mean of the selected frames, optionally drift-corrected first.
+
+ Streams: one frame resident at a time plus one float64 accumulator, so this
+ is safe at the plan's target scale however long the movie is. NaN padding
+ from :func:`spyde.drift.warp.shift_frame` is excluded per pixel rather than
+ zero-filled — a zero-filled border reads as a dark rim that looks like real
+ data and would be segmented as one.
+ """
+ acc = None
+ hits = None
+ for i in indices:
+ frame = np.asarray(get_frame(int(i)), dtype=np.float32)
+ if shifts is not None:
+ s = shifts[int(i)]
+ if np.all(np.isfinite(s)):
+ from spyde.drift import shift_frame
+ frame = shift_frame(frame, s, order=order)
+ if acc is None:
+ acc = np.zeros(frame.shape, np.float64)
+ hits = np.zeros(frame.shape, np.int32)
+ good = np.isfinite(frame)
+ acc[good] += frame[good]
+ hits[good] += 1
+ if acc is None:
+ return np.zeros((1, 1), np.float32)
+ with np.errstate(invalid="ignore", divide="ignore"):
+ out = np.where(hits > 0, acc / np.maximum(hits, 1), np.nan)
+ return out.astype(np.float32)
+
+
+def _gradient_energy(img, mask=None) -> float:
+ """Mean squared forward-difference gradient over the valid pixels.
+
+ The sharpness number, and it is NaN-aware by construction rather than by
+ ``nan_to_num``: an aligned sum's uncovered border is NaN (plan A7 — nothing
+ is cropped, nothing is invented), and zero-filling it manufactures a step
+ at the border whose gradient energy dwarfs the image's own, which would
+ make every ROI look brilliantly sharp. Differences touching a non-finite
+ (or masked-out) pixel are excluded from BOTH the sum and the count, so the
+ raw and aligned sums are measured over exactly the same pixels.
+ """
+ a = np.asarray(img, np.float64)
+ ok = np.isfinite(a)
+ if mask is not None:
+ ok &= np.asarray(mask, bool)
+ a = np.where(ok, a, 0.0)
+ total = 0.0
+ count = 0
+ if a.shape[0] > 1:
+ m = ok[1:, :] & ok[:-1, :]
+ d = (a[1:, :] - a[:-1, :])[m]
+ total += float(np.sum(d * d))
+ count += int(m.sum())
+ if a.shape[1] > 1:
+ m = ok[:, 1:] & ok[:, :-1]
+ d = (a[:, 1:] - a[:, :-1])[m]
+ total += float(np.sum(d * d))
+ count += int(m.sum())
+ return total / count if count else float("nan")
+
+
+def _preview_indices(n_frames: int, k: int, frame_bytes: int) -> np.ndarray:
+ """Evenly spaced sample of the movie for the preview, thinned to fit the
+ byte cap. See :data:`_PREVIEW_FRAMES` for why evenly spaced and not the
+ first *k* in a row."""
+ n = max(1, int(n_frames))
+ k = max(2, min(int(k), n))
+ cap = max(2, int(_PREVIEW_MAX_BYTES // max(1, int(frame_bytes))))
+ k = min(k, cap)
+ return np.unique(np.linspace(0, n - 1, k).round().astype(int))
+
+
+def preview_alignment(get_frame, indices, roi, *, params) -> dict:
+ """Align *indices* on *roi* alone and report how much sharper the sum got.
+
+ Reads one FULL frame at a time and keeps only the crop, so the resident set
+ is one frame plus ``len(indices)`` crops (bounded by
+ :data:`_PREVIEW_MAX_BYTES` at the caller). The crops ARE the region to
+ correlate, so the solve runs with ``roi=None`` on them.
+
+ Returns ``{roi, frames, raw, aligned, gain, raw_energy, aligned_energy,
+ max_abs_shift}``. *gain* is the whole point: > 1 means aligning this region
+ genuinely sharpened it, ~1 means alignment changes nothing here (a
+ featureless box), and it is measured on the pixels both sums cover.
+ """
+ from spyde.drift import solve_translation
+
+ crops: list[np.ndarray] = []
+ for i in indices:
+ frame = np.asarray(get_frame(int(i)), np.float32)
+ if roi is not None:
+ y0, x0, h, w = (int(v) for v in roi)
+ frame = frame[y0:y0 + h, x0:x0 + w]
+ crops.append(np.ascontiguousarray(frame, dtype=np.float32))
+
+ model = solve_translation(crops, **_solver_kwargs(params))
+ take = range(len(crops))
+ raw = _stack_sum(crops.__getitem__, take)
+ aligned = _stack_sum(crops.__getitem__, take, model.shifts,
+ order=int(params["order"]))
+ both = np.isfinite(raw) & np.isfinite(aligned)
+ e_raw = _gradient_energy(raw, both)
+ e_aligned = _gradient_energy(aligned, both)
+ gain = (e_aligned / e_raw) if (np.isfinite(e_raw) and e_raw > 0) \
+ else float("nan")
+ return {"roi": None if roi is None else tuple(int(v) for v in roi),
+ "frames": len(crops), "raw": raw, "aligned": aligned,
+ "gain": float(gain), "raw_energy": float(e_raw),
+ "aligned_energy": float(e_aligned),
+ "max_abs_shift": float(model.max_abs_shift)}
+
+
+# ── staged handlers ──────────────────────────────────────────────────────────
+
+def drift_open(session, plot, payload) -> None:
+ """Caret mounted: build the controller, open the Drift Check window, draw
+ the alignment ROI, and run the first discovery preview.
+
+ Nothing SOLVES here — plan A8 is explicit that drift correction is opt-in
+ and never runs on load. The compute is the bounded raw sum plus the
+ ~20-frame preview of the default box, which is what makes the caret's first
+ frame informative instead of an empty panel and a button.
+ """
+ src, tree = _src_plot_tree(session, plot)
+ if src is None or tree is None:
+ emit_error("Drift Correction: no active dataset")
+ return
+
+ existing = getattr(tree, "_drift_wizard", None)
+ if existing is not None and not existing._closed:
+ existing.params = _coerce({**existing.params, **(payload or {})})
+ _emit_state(existing)
+ return
+
+ wiz = DriftWizard(session, tree, src)
+ wiz.params = _coerce(payload)
+ try:
+ n_frames, get_frame, shape = wiz.frames()
+ except TypeError as exc:
+ emit_error(f"Drift Correction: {exc}")
+ return
+ # BEFORE the worker: StrictMode fires open/close/open synchronously and the
+ # close's bump has to be able to invalidate this open's deferred build.
+ gen = wiz.guard()
+ tree._drift_wizard = wiz
+ _emit_state(wiz, n_frames=int(n_frames))
+
+ def _work():
+ return _stack_sum(get_frame, wiz.sum_indices(n_frames))
+
+ def _done(raw_sum):
+ if not wiz.still(gen) or wiz._closed:
+ return
+ wiz.open_check_window(raw_sum, int(n_frames))
+ wiz.ensure_roi_widget(shape)
+ _emit_state(wiz, n_frames=int(n_frames))
+ emit_status("Drift Correction: drag the box onto a landmark to test it, "
+ "then Correct Drift.")
+ _run_preview(wiz)
+
+ def _fail(exc):
+ emit_error(f"Drift Correction: reading the movie failed — {exc}")
+
+ run_on_worker(session, _work, name="drift-open", on_done=_done, on_error=_fail)
+
+
+def drift_close(session, plot, payload=None) -> None:
+ """Caret unmounted: invalidate in-flight work FIRST, then tear down."""
+ _src, tree = _src_plot_tree(session, plot)
+ wiz = _wizard(session, plot)
+ if tree is not None:
+ # The same `_drift_run_gen` key WizardController.cancel_inflight bumps,
+ # done on the TREE so it fires even when there is no controller yet: a
+ # StrictMode open whose worker has not landed must still be cancelled.
+ bump_generation(tree, "_drift_run_gen")
+ bump_generation(tree, "_drift_preview_gen")
+ if wiz is not None:
+ # Harmlessly re-bumps when the tree resolved above; the point is the
+ # case where it did not (the wizard was found through one of the
+ # figure windows' controller registry).
+ wiz.cancel_inflight()
+ wiz.remove()
+
+
+def drift_set_method(session, plot, payload) -> None:
+ """Select the drift model (Advanced).
+
+ Only ``rigid`` has a solver in ``spyde.drift`` today. Selecting either of
+ the others says so and stays on rigid — running a rigid solve while the
+ caret claims "rigid+affine" would put a wrong ``kind`` into the model's
+ provenance, which is worse than the missing feature.
+ """
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ return
+ method = str((payload or {}).get("method", "")).lower()
+ if method not in METHODS:
+ emit_error(f"Drift Correction: unknown model {method!r}")
+ return
+ reason = _UNAVAILABLE.get(method)
+ if reason:
+ emit_status(f"Drift Correction: {reason} — staying on the rigid solve.")
+ method = "rigid"
+ wiz.params["method"] = method
+ _emit_state(wiz)
+
+
+def drift_tune(session, plot, payload) -> None:
+ """A toggle or Advanced parameter changed → re-run the discovery preview.
+
+ NOT debounced here: the renderer's ``useDebouncedAction`` already settles
+ the send, and debouncing twice only adds latency. The drag path IS
+ debounced, on the backend, because widget pointer events arrive at renderer
+ frame rate (:meth:`DriftWizard.schedule_preview`).
+ """
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ return
+ wiz.params = _coerce({**wiz.params, **(payload or {})})
+ _emit_state(wiz)
+ _run_preview(wiz)
+
+
+def _run_preview(wiz: DriftWizard) -> None:
+ """Align ~20 sampled frames on the current box and report the gain.
+
+ Latest-wins on ``_drift_preview_gen``: a drag that outruns the solve drops
+ the stale result rather than painting it over the newer one. The preview
+ never touches ``tree.drift`` or the caret's solved state — it is a question,
+ not an answer.
+ """
+ if wiz._closed:
+ return
+ tree = wiz.tree
+ gen = bump_generation(tree, "_drift_preview_gen")
+ params = dict(wiz.params)
+ roi = wiz.roi_box() # the BOX, toggle or not — see active_roi()
+ try:
+ n_frames, get_frame, shape = wiz.frames()
+ except TypeError as exc:
+ log.debug("[drift] preview skipped: %s", exc)
+ return
+ if n_frames < 2:
+ return
+ h, w = (roi[2], roi[3]) if roi is not None else (int(shape[0]), int(shape[1]))
+ indices = _preview_indices(n_frames, params["preview_frames"], h * w * 4)
+
+ def _work():
+ return preview_alignment(get_frame, indices, roi, params=params)
+
+ def _done(result):
+ if not is_current(tree, "_drift_preview_gen", gen) or wiz._closed:
+ return
+ wiz.preview = result
+ wiz.show_preview(result)
+ emit({"type": "drift_preview", "window_id": wiz.src_window_id,
+ "roi": None if result["roi"] is None else list(result["roi"]),
+ "frames": int(result["frames"]),
+ "gain": float(result["gain"]),
+ "max_abs_shift": float(result["max_abs_shift"]),
+ "params": dict(params)})
+
+ def _fail(exc):
+ if is_current(tree, "_drift_preview_gen", gen):
+ emit_error(f"Drift preview failed: {exc}")
+
+ run_on_worker(wiz.session, _work, name="drift-preview",
+ on_done=_done, on_error=_fail)
+
+
+def drift_run(session, plot, payload) -> None:
+ """Solve the whole movie on a worker: progress-reported and cancellable.
+
+ Opens the dy/dx window FIRST and fills it from ``solve_translation``'s
+ ``on_shift`` stream, so the curve draws while it solves rather than
+ appearing whole at the end.
+
+ Cancellation goes through ``BaseSignalTree.register_cancel`` so closing the
+ tree stops the solve, and ``solve_translation``'s own ``cancel()`` hook
+ polls the same flag — a cancelled solve leaves NaN shifts for the frames it
+ never reached, which is why a partial model is detectable rather than
+ silently wrong. Stop/Discard flips the same flag.
+ """
+ src, tree = _src_plot_tree(session, plot)
+ if src is None or tree is None:
+ emit_error("Drift Correction: no active dataset")
+ return
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ emit_error("Drift Correction: the caret is not open")
+ return
+ wiz.params = _coerce({**wiz.params, **(payload or {})})
+ p = dict(wiz.params)
+ reason = _UNAVAILABLE.get(p["method"])
+ if reason:
+ emit_status(f"Drift Correction: {reason} — solving rigid instead.")
+ p["method"] = wiz.params["method"] = "rigid"
+
+ try:
+ n_frames, get_frame, _shape = wiz.frames()
+ except TypeError as exc:
+ emit_error(f"Drift Correction: {exc}")
+ return
+ if n_frames < 2:
+ emit_error("Drift Correction needs at least two frames")
+ return
+
+ roi = wiz.active_roi()
+ if p["use_roi"] and roi is None:
+ emit_status("Drift Correction: no usable alignment box — correlating "
+ "the whole frame.")
+
+ gen = wiz.guard()
+ stopped = [False]
+ wiz._stop = stopped
+ if hasattr(tree, "register_cancel"):
+ tree.register_cancel(flag=stopped)
+ wiz.open_trace_window(int(n_frames))
+ _emit_state(wiz)
+ emit_status(f"Solving drift over {n_frames} frames…")
+ dispatch = getattr(session, "_dispatch_to_main", None)
+
+ def _work():
+ from spyde.drift import solve_translation
+
+ def _progress(done, total):
+ emit_progress(int(done), int(total), "Drift")
+ emit({"type": "drift_progress", "window_id": wiz.src_window_id,
+ "done": int(done), "total": int(total)})
+
+ # Stream the curve as it solves. `progress` carries only a count and the
+ # shift array is solver-local until the return, so without this callback
+ # the caret could show a bar but not a trace. Batched rather than per
+ # frame: at thousands of frames one message each would flood the PLOTAPP
+ # line protocol for a curve the eye cannot follow that finely. The PAINT
+ # is marshalled — `on_shift` runs on the solver thread and figures are
+ # main-thread only (README §6).
+ pending: list[tuple[int, float, float]] = []
+ last_flush = [time.monotonic()]
+
+ def _flush():
+ if not pending:
+ return
+ batch = pending[:]
+ pending.clear()
+ last_flush[0] = time.monotonic()
+ emit({"type": "drift_trace", "window_id": wiz.src_window_id,
+ "points": batch})
+ if not wiz.still(gen):
+ return
+ if dispatch is None:
+ wiz.push_trace(batch)
+ else:
+ dispatch(lambda b=batch: (None if wiz._closed or not wiz.still(gen)
+ else wiz.push_trace(b)))
+
+ def _on_shift(i, dy, dx, _sharp):
+ pending.append((int(i), float(dy), float(dx)))
+ if (len(pending) >= _TRACE_BATCH
+ or time.monotonic() - last_flush[0] >= _TRACE_MAX_INTERVAL):
+ _flush()
+
+ model = solve_translation(
+ wiz.signal(), progress=_progress, on_shift=_on_shift,
+ cancel=lambda: stopped[0],
+ provenance={"action": "Drift Correction", "params": dict(p),
+ "roi": None if roi is None else [int(v) for v in roi]},
+ **_solver_kwargs(p, roi))
+ _flush()
+ if stopped[0]:
+ return model, None, float("nan")
+
+ # The non-rigid residual, on top of the rigid solve. Runs AFTER it (and
+ # only on request) because it is a correction to what rigid leaves
+ # behind: the rigid pass has already removed everything that moves the
+ # whole frame, so what this fits is scan distortion or sample
+ # deformation rather than a mixture of those and the stage.
+ if p["method"] == "nonrigid":
+ emit_status(f"Fitting the non-rigid field ({p['nonrigid_model']})…")
+
+ def _nr_progress(done, total):
+ emit_progress(int(done), int(total), "Drift (non-rigid)")
+ emit({"type": "drift_progress", "window_id": wiz.src_window_id,
+ "done": int(done), "total": int(total)})
+
+ model = _solve_nonrigid_step(
+ wiz, p, model, get_frame, int(n_frames),
+ progress=_nr_progress, cancel=lambda: stopped[0])
+ if stopped[0]:
+ return model, None, float("nan")
+ # One extra streaming pass over the SAME bounded subset the raw sum
+ # used, so the two check images are comparable.
+ after = _stack_sum(get_frame, wiz.sum_indices(n_frames), model.shifts,
+ order=int(p["order"]))
+ # The same number the discovery preview reports, now for the whole
+ # movie — measured on the worker because a 4096² gradient energy is
+ # ~100 ms and the main thread is the navigator's.
+ before = wiz._before_sum
+ gain = float("nan")
+ if before is not None and before.shape == after.shape:
+ both = np.isfinite(before) & np.isfinite(after)
+ e_before = _gradient_energy(before, both)
+ if np.isfinite(e_before) and e_before > 0:
+ gain = _gradient_energy(after, both) / e_before
+ return model, after, gain
+
+ def _done(res):
+ model, after, gain = res
+ try:
+ if not wiz.still(gen) or wiz._closed:
+ return
+ wiz.model = model
+ tree.drift = model
+ wiz.update_check(after=after)
+ emit({"type": "drift_result", "window_id": wiz.src_window_id,
+ "shifts": [[float(a), float(b)] for a, b in model.shifts],
+ "kind": model.kind, "reference": model.reference,
+ "roi": None if roi is None else [int(v) for v in roi],
+ "max_abs_shift": float(model.max_abs_shift),
+ "gain": float(gain),
+ "rejected": int(model.params.get("rejected_from_reference", 0)),
+ "cancelled": bool(stopped[0])})
+ _emit_state(wiz)
+ solved = int(np.isfinite(model.shifts).all(axis=1).sum())
+ if stopped[0]:
+ emit_status(f"Drift solve stopped after {solved} of "
+ f"{n_frames} frames")
+ else:
+ emit_status(f"Drift solved: max shift "
+ f"{model.max_abs_shift:.2f} px over {n_frames} frames")
+ finally:
+ if hasattr(tree, "unregister_cancel"):
+ try:
+ tree.unregister_cancel(flag=stopped)
+ except Exception as exc:
+ log.debug("[drift] unregister_cancel failed: %s", exc)
+
+ def _fail(exc):
+ emit_error(f"Drift Correction failed: {exc}")
+ log.exception("drift solve failed")
+ if hasattr(tree, "unregister_cancel"):
+ try:
+ tree.unregister_cancel(flag=stopped)
+ except Exception as e2:
+ log.debug("[drift] unregister_cancel failed: %s", e2)
+
+ run_on_worker(session, _work, name="drift-run", on_done=_done, on_error=_fail)
+
+
+def drift_discard(session, plot, payload=None) -> None:
+ """Stop a solve in flight and/or throw the solved model away.
+
+ One handler for both because they are the same user intent ("no, not
+ that"): the button reads *Stop* while the bar is moving and *Discard*
+ once there is a result. Bumping the run generation FIRST means a solve that
+ finishes anyway lands on a stale generation and never installs itself.
+ """
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ return
+ wiz._stop[0] = True
+ wiz.cancel_inflight()
+ wiz.model = None
+ if getattr(wiz.tree, "drift", None) is not None:
+ wiz.tree.drift = None
+ wiz.close_trace_window()
+ if wiz._panels and wiz._before_sum is not None:
+ try:
+ wiz._panels["after"].set_data(np.zeros_like(wiz._before_sum))
+ except Exception as exc:
+ log.debug("[drift] clearing the corrected sum failed: %s", exc)
+ _emit_state(wiz)
+ emit_status("Drift result discarded.")
+
+
+# ── the corrected node ───────────────────────────────────────────────────────
+
+def drift_corrected(signal, *, model, order: int = 1, fill: float = float("nan")):
+ """A LAZY drift-corrected view of *signal*. Plan §0.7.
+
+ Parameters
+ ----------
+ signal
+ The source movie (1-D navigation, 2-D signal).
+ model
+ The :class:`~spyde.drift.model.DriftModel` to apply. ``shifts[i]`` is the
+ correction ADDED to frame *i* — go through the model rather than writing
+ the arithmetic out; the inverted sign doubles the drift and still looks
+ plausible (``spyde/drift/model.py``).
+ order
+ Interpolation order for sub-pixel shifts. A whole-pixel model takes an
+ exact slice-copy path inside :func:`~spyde.drift.warp.shift_frame`.
+ fill
+ Uncovered-pixel value. NaN by default, per the plan A7 edge policy —
+ nothing is cropped and nothing is filled with invented data.
+
+ Notes
+ -----
+ Built with ``map_blocks`` over the source's OWN chunking, deliberately: this
+ never calls ``.rechunk()`` and never computes anything, so a multi-GB movie
+ costs a graph and nothing else (CLAUDE.md memory-safety rule, and Live-
+ Display §1 on not reshuffling storage chunks). Each block warps its own
+ frames using ``block_info`` to recover their absolute indices, so a movie
+ stored several frames per chunk works unchanged.
+ """
+ import dask.array as da
+ from spyde.drift import shift_frame
+
+ data = signal.data
+ if getattr(data, "ndim", 0) != 3:
+ raise ValueError(
+ f"drift correction needs a (n, h, w) frame stack; got shape "
+ f"{getattr(data, 'shape', None)}")
+ shifts = np.asarray(model.shifts, dtype=np.float32)
+ if shifts.shape[0] != int(data.shape[0]):
+ raise ValueError(
+ f"the drift model covers {shifts.shape[0]} frames but the signal has "
+ f"{int(data.shape[0])} — solve again on this node")
+
+ if not isinstance(data, da.Array):
+ # Already resident; wrapping it costs nothing and keeps the node lazy so
+ # the whole tree reads through one path.
+ data = da.from_array(data, chunks=(1,) + tuple(int(s) for s in data.shape[1:]))
+
+ def _block(blk, block_info=None):
+ t0 = (0 if block_info is None
+ else int(block_info[0]["array-location"][0][0]))
+ out = np.empty(blk.shape, np.float32)
+ for k in range(blk.shape[0]):
+ s = shifts[t0 + k]
+ if not np.all(np.isfinite(s)):
+ # A frame the solve never reached (cancelled) keeps its raw
+ # pixels rather than becoming an all-NaN hole.
+ out[k] = np.asarray(blk[k], np.float32)
+ else:
+ out[k] = shift_frame(blk[k], s, order=int(order), fill=fill)
+ return out
+
+ warped = da.map_blocks(_block, data, dtype=np.float32,
+ meta=np.zeros((0, 0, 0), np.float32))
+ new = signal._deepcopy_with_new_data(warped)
+ if not new._lazy:
+ new._lazy = True
+ new._assign_subclass()
+ return new
+
+
+def _commit(wiz: DriftWizard):
+ if wiz.model is None:
+ emit_error("Drift Correction: solve first, then Apply")
+ return None
+ parent = wiz.signal()
+ try:
+ new_signal = wiz.tree.add_transformation(
+ parent, function=drift_corrected, node_name="Drift corrected",
+ local=True, model=wiz.model, order=int(wiz.params["order"]))
+ except Exception as exc:
+ emit_error(f"Drift Correction: applying the model failed — {exc}")
+ log.exception("drift commit failed")
+ return None
+ if new_signal is None:
+ return None
+ wiz.tree.drift = wiz.model
+ try:
+ new_signal.metadata.set_item(
+ "General.spyde_provenance",
+ {"action": "Drift Correction", "params": dict(wiz.params),
+ "kind": wiz.model.kind, "reference": wiz.model.reference,
+ "roi": wiz.model.params.get("roi")})
+ except Exception as exc:
+ log.debug("[drift] stamping provenance failed: %s", exc)
+ show_tree_node(wiz.src_plot, wiz.tree, new_signal)
+ emit_status(f"Drift corrected node added (max shift "
+ f"{wiz.model.max_abs_shift:.2f} px)")
+ return new_signal
+
+
+def drift_commit(session, plot, payload=None) -> None:
+ """Add the lazy corrected node to the tree and show it."""
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ emit_error("Drift Correction: nothing to apply")
+ return
+ wiz.commit()
+
+
+def drift_correction(ctx, action_name: str = "Drift Correction", **params):
+ """Toolbar entry — a no-op parent; the Electron toolbar opens the staged
+ caret, which drives the ``drift_*`` handlers (README §4)."""
+ return None
diff --git a/spyde/actions/lifecycle.py b/spyde/actions/lifecycle.py
index e9ca8e89..647580f7 100644
--- a/spyde/actions/lifecycle.py
+++ b/spyde/actions/lifecycle.py
@@ -392,6 +392,123 @@ def __exit__(self, exc_type, exc, tb) -> bool:
return False
+# ── the standard "this is computing" surface for a batch action ──────────────
+
+class batch_feedback:
+ """The THREE things every long batch run should show, in one object.
+
+ Batch actions had each grown their own subset of these and no two agreed:
+ the navigator/VI fills raised the "Calculating…" overlay but no action run
+ did, ``emit_progress`` labels ranged from ``"Segmenting"`` to
+ ``"Encoding movie 3/40"`` (counts baked into the label for one caller and
+ not the others), and only some paths showed anything at all on the plot
+ while they worked. From the outside that reads as "some of these are
+ broken" — the ones that look idle are the ones that say nothing.
+
+ The three surfaces, and why each earns its place:
+
+ * **The window overlay** (:class:`window_computing`) — answers "is this
+ thing doing anything", and is the only one visible when the status bar is
+ off-screen or covered.
+ * **Status-bar progress** (:func:`~spyde.backend.ipc.emit_progress`) —
+ answers "how far, and how much longer". Rate-limited, because a 900-frame
+ run emitting per frame puts 900 messages on the same stdout line protocol
+ the nav painter uses.
+ * **The most recent RESULT, painted live** — answers the question the other
+ two cannot: *is it working properly?* A progress bar advancing through a
+ 900-frame run that is quietly finding nothing looks exactly like one that
+ is working. Showing the newest frame's result turns a 20-minute wait into
+ something you can abandon in the first ten seconds.
+
+ ``publish`` is the caller's "paint this result" function; it is called on
+ the **asyncio main thread** (this class marshals), rate-limited by the same
+ clock as the progress emission, and always called for the FINAL result
+ regardless of the rate limit — the last frame is the one left on screen.
+
+ Every surface is None-guarded, so a caller with no window, or no result to
+ paint, uses the same code path as one that has both.
+
+ Usage from a worker thread::
+
+ fb = batch_feedback(session, result.window_id, "Segmenting", n_frames,
+ publish=lambda r: wiz.set_overlay(*r))
+ with fb:
+ for i in range(n):
+ ...
+ fb.step(i + 1, result=(contours, box))
+ """
+
+ def __init__(self, session, window_id: int | None, label: str, total: int,
+ *, publish: "Callable[[Any], None] | None" = None,
+ min_interval: float = 0.35):
+ self.session = session
+ self.label = str(label)
+ self.total = int(total)
+ self.publish = publish
+ self.min_interval = float(min_interval)
+ self._overlay = window_computing(window_id)
+ self._last = 0.0
+ self._done = 0
+
+ # -- lifecycle ----------------------------------------------------------
+ def __enter__(self) -> "batch_feedback":
+ self._overlay.start()
+ return self
+
+ def __exit__(self, exc_type, exc, tb) -> bool:
+ # Unconditional, like `window_computing`: a cancelled or failed batch
+ # must not leave the overlay spinning forever over a window that has
+ # stopped working. This is the failure the pairing contract exists for.
+ self._overlay.stop()
+ return False
+
+ # -- during the run -----------------------------------------------------
+ def step(self, done: int, *, result: "Any" = None, force: bool = False) -> None:
+ """One unit finished. Emits progress and paints *result*, rate-limited.
+
+ Safe to call from a worker thread — the paint is marshalled. Call with
+ ``force=True`` for the final unit so the last result is the one left on
+ screen even if it lands inside the rate-limit window.
+ """
+ self._done = int(done)
+ now = time.monotonic()
+ if not force and self._done < self.total and now - self._last < self.min_interval:
+ return
+ self._last = now
+ from spyde.backend.ipc import emit_progress
+ emit_progress(self._done, self.total, self.label)
+ if result is not None and self.publish is not None:
+ self._paint(result)
+
+ def _paint(self, result) -> None:
+ publish = self.publish
+ if publish is None:
+ return
+ dispatch = getattr(self.session, "_dispatch_to_main", None)
+ if dispatch is None: # no loop (tests) — paint inline
+ _safe_publish(publish, result)
+ return
+ dispatch(lambda: _safe_publish(publish, result))
+
+ def finish(self) -> None:
+ """Terminal progress tick, so the status bar's spinner actually stops.
+
+ Without it a run that ends early (cancelled, or a rate-limited final
+ step) leaves the bar mid-way and the user reads a finished job as hung —
+ which is exactly what `seg_run` did before it emitted a final
+ ``emit_progress(n, n)``.
+ """
+ from spyde.backend.ipc import emit_progress
+ emit_progress(self.total, self.total, self.label)
+
+
+def _safe_publish(publish, result) -> None:
+ """A live-preview paint must never take the batch down with it."""
+ try:
+ publish(result)
+ except Exception as exc:
+ log.debug("[batch] live result paint failed: %s", exc)
+
# ── progressive shared-memory fill ────────────────────────────────────────────
diff --git a/spyde/actions/navigator_views.py b/spyde/actions/navigator_views.py
index 3231d7e4..4391181f 100644
--- a/spyde/actions/navigator_views.py
+++ b/spyde/actions/navigator_views.py
@@ -98,6 +98,12 @@ def select_navigator(session, plot, payload) -> None:
_switch_navigator(tree, plot, names[0])
emit_navigator_options(tree)
elif _tree_nav_is_1d(tree):
+ # A PARTICLE tree's lanes render themselves — the count lane is integer
+ # data and must be a step, and the event lane is coloured markers, not a
+ # trace. Delegate before falling back to the generic plain-line stacker.
+ from spyde.actions.particle_overlay import maybe_stack_particle_lanes
+ if maybe_stack_particle_lanes(session, plot, tree, names):
+ return
# In-situ movie / time series: stack the 1-D traces with a shared,
# linked time cursor (see module docstring).
_stack_navigators(session, plot, tree, names)
diff --git a/spyde/actions/particle_overlay.py b/spyde/actions/particle_overlay.py
new file mode 100644
index 00000000..49e82ff1
--- /dev/null
+++ b/spyde/actions/particle_overlay.py
@@ -0,0 +1,1766 @@
+"""
+particle_overlay.py — the live particle overlay + the stacked navigator lanes.
+
+Plan B9 (overlay and editing), C2 (events on the navigator) and C3 (trails,
+integer lanes are step plots). Modelled on :mod:`spyde.actions.vector_overlay`:
+marker groups on the signal plot, re-pushed from an ``index_hook`` on the tree's
+navigation selectors, torn down through ``lifecycle.replace_tree_attr``.
+
+Two coordinate systems, and the whole overlay hangs on telling them apart
+--------------------------------------------------------------------------
+anyplotlib 2-D markers drawn with ``transform="data"`` are addressed in
+**image-pixel** coordinates — no axis scale or offset (``masks._signal_k_grids``
+documents the bug class; building geometry in physical units gives an empty or
+misplaced overlay on any calibrated axis). But the two things this module draws
+live in *different* spaces already:
+
+* **contours** are stored by ``measure_frame`` as int16 **pixel** ``(y, x)`` — so
+ they only need the axis swap to ``(x, y)``, never a scale;
+* **centroids** (``COL["y"]``/``COL["x"]``) are **calibrated**, ``pixel * scale``
+ — so they must be divided by ``particles.scale`` to reach widget space.
+
+Dividing a contour or forgetting to divide a centroid both produce an overlay
+that is plausible at ``scale == 1`` and silently wrong everywhere else, which is
+exactly how the ``_signal_k_grids`` bug survived. The synthetic fixture is
+``scale=0.5`` on purpose, and ``test_particle_overlay.py`` pins the conversion
+there.
+
+A click arrives in the THIRD space: anyplotlib's ``Event`` carries only
+``xdata``/``ydata`` (the JS emits ``img_x``/``img_y`` too, but anyplotlib's Event
+dataclass has no such field and drops them), and those are the panel's *physical*
+data coordinates. So the hit test converts the click back to pixels through the
+**plot's own** signal axes rather than assuming they agree with
+``particles.scale``.
+
+One marker group per colour — not one group with a colour array
+---------------------------------------------------------------
+anyplotlib's 1-D marker path accepts ``fill_color``/``color`` as arrays parallel
+to the items; its **2-D** path does not (``drawMarkers2d`` in ``figure_esm.js``
+binds one ``fillStyle``/``strokeStyle`` for the whole set, and an array there is
+an invalid canvas colour that silently leaves the previous style in place). The
+overlay lives on a 2-D panel, so "colour per track" means one group per colour:
+six accent buckets plus one grey for untracked rows. That is also why the trail
+needs a group per (colour, age) pair — a polyline cannot fade along its own
+length in this wire format.
+
+The consequence is ~30 marker groups, and ``MarkerGroup.set`` re-serialises the
+WHOLE registry on every call, so the groups are updated through
+:func:`_push_groups`, which mutates them and pushes **once** per frame.
+"""
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any, Callable, Iterable, Sequence
+
+import numpy as np
+
+from spyde.signals.particles import COL, MEASURED_COLUMNS, N_COLUMNS
+
+log = logging.getLogger(__name__)
+
+
+# ── palette ──────────────────────────────────────────────────────────────────
+
+#: SpyDE's six accents, in order. Track ids cycle through them by ``id % 6`` so
+#: the mapping is a pure function of the id: the same track keeps its colour as
+#: the user scrubs, across a re-open, and between the overlay and any other
+#: surface that colours by track (the kymograph, the report embed).
+TRACK_COLORS: tuple[str, ...] = (
+ "#89b4fa", # blue — track 0
+ "#f38ba8", # red
+ "#a6e3a1", # green
+ "#f9e2af", # yellow
+ "#cba6f7", # mauve
+ "#94e2d5", # teal
+)
+
+#: Untracked rows (``track_id < 0``: the linker has not run, or this row was
+#: created by an edit and has not been re-linked). Grey, not a seventh accent —
+#: "no identity yet" must not read as "identity number six".
+UNTRACKED_COLOR = "#6c7086"
+
+#: The selected particle's outline. Deliberately outside :data:`TRACK_COLORS` so
+#: selection is never confusable with a track colour.
+SELECTED_COLOR = "#ffffff"
+
+#: Fill opacity. 25% keeps the underlying image readable at the plan's target
+#: density (hundreds of particles per frame) — the outline carries the shape.
+FILL_ALPHA = 0.25
+OUTLINE_WIDTH = 1.0
+SELECTED_WIDTH = 2.0
+
+#: Trailing window, in frames. ~8 is enough to read a direction without the
+#: trails of a dense field crossing into an unreadable mesh.
+DEFAULT_TRAIL_FRAMES = 8
+
+#: How many opacity steps the trail fades through. A polyline carries ONE alpha
+#: in this wire format (see the module docstring), so each step is a separate
+#: marker group per colour: 6 x 4 = 24 groups. Four steps is where the ramp stops
+#: looking banded on an 8-frame window; eight would double the group count for a
+#: difference that is invisible at 1 px line width.
+TRAIL_FADE_STEPS = 4
+
+#: Head-dot radius in image pixels.
+HEAD_RADIUS_PX = 2.5
+
+#: Event lane colours, plan C2. Order is ``track.EVENT_KINDS``.
+EVENT_COLORS: dict[str, str] = {
+ "birth": "#a6e3a1", # green
+ "death": "#f38ba8", # red
+ "merge": "#cba6f7", # mauve
+ "split": "#f9e2af", # yellow
+}
+
+#: Properties shown in the selected particle's readout. Three, not twelve: the
+#: readout sits ON the frame beside the particle, so it competes with the data.
+READOUT_COLUMNS: tuple[str, ...] = ("area", "equiv_diameter", "circularity")
+
+#: The caret's parameter schema — the single host-agnostic source of truth
+#: (README §4.2), resolved by ``registry.wizard_parameters("part")``. Same dict
+#: spec as a ``toolbars.yaml`` ``parameters:`` block.
+PARAMETERS: dict[str, dict] = {
+ "show_trails": {
+ "name": "Trails", "type": "bool", "default": False,
+ "description": "Fade the last N frames of each track behind a head dot.",
+ },
+ "trail_frames": {
+ "name": "Trail length", "type": "int", "default": DEFAULT_TRAIL_FRAMES,
+ "min": 2, "max": 60, "step": 1,
+ "description": "Trailing window, in frames.",
+ },
+ "region_select": {
+ "name": "Region select", "type": "bool", "default": False,
+ "description": "Rubber-band box for selecting many particles at once.",
+ },
+}
+
+
+def track_color(track_id) -> str:
+ """Colour for a track id.
+
+ Parameters
+ ----------
+ track_id
+ The ``track_id`` column value. Negative (or NaN) means "not linked" and
+ maps to :data:`UNTRACKED_COLOR`.
+
+ Returns
+ -------
+ str
+ A ``#rrggbb`` hex string from :data:`TRACK_COLORS`, cycling every six.
+ """
+ try:
+ tid = int(track_id)
+ except (TypeError, ValueError):
+ return UNTRACKED_COLOR
+ if tid < 0:
+ return UNTRACKED_COLOR
+ return TRACK_COLORS[tid % len(TRACK_COLORS)]
+
+
+def fade(color: str, alpha: float) -> str:
+ """``#rrggbb`` + an 8-bit alpha byte → ``#rrggbbaa``.
+
+ Canvas accepts 8-digit hex, which is how a trail segment carries its own
+ opacity even though the group's ``linewidth``/``color`` are scalars.
+ """
+ a = int(round(float(np.clip(alpha, 0.0, 1.0)) * 255))
+ return f"{color}{a:02x}"
+
+
+def trail_alphas(steps: int = TRAIL_FADE_STEPS) -> list[float]:
+ """Opacity per age step, newest first. Linear from 1.0 down to 0.2.
+
+ Not down to 0: the oldest step of a trail still has to be visible, or the
+ window reads as shorter than it is.
+ """
+ n = max(1, int(steps))
+ if n == 1:
+ return [1.0]
+ return [1.0 - 0.8 * i / (n - 1) for i in range(n)]
+
+
+# ── geometry ─────────────────────────────────────────────────────────────────
+
+def frame_from_indices(indices, nav_to_frame=None) -> int:
+ """Particle-frame index from a navigation selector's committed indices.
+
+ A particle tree's navigation space is 1-D (time), so the first raveled
+ coordinate IS the frame — the same read ``navigator_views._StackedNavCursor``
+ makes. *nav_to_frame*, when given, maps a source navigation index onto a
+ particle frame (the inverse of ``tree.nav_map``), for an overlay drawn on the
+ SOURCE movie rather than on the particle tree's own label movie.
+ """
+ try:
+ idx = int(np.asarray(indices).ravel()[0])
+ except (TypeError, ValueError, IndexError):
+ return 0
+ if nav_to_frame is not None:
+ return int(nav_to_frame.get(idx, idx))
+ return idx
+
+
+def centroids_px(rows: np.ndarray, scale: float) -> np.ndarray:
+ """``(n, 2)`` float32 ``(x, y)`` marker offsets from property rows.
+
+ Centroids are stored calibrated (``pixel * scale``), markers want image
+ pixels — see the module docstring. Column order flips too: the row carries
+ ``(y, x)``, a marker offset is ``(x, y)``.
+ """
+ rows = np.asarray(rows)
+ if rows.size == 0:
+ return np.zeros((0, 2), np.float32)
+ s = float(scale) or 1.0
+ return np.column_stack([rows[:, COL["x"]] / s,
+ rows[:, COL["y"]] / s]).astype(np.float32)
+
+
+def contour_xy(particles, index: int) -> np.ndarray:
+ """One particle's outline as ``(k, 2)`` float32 ``(x, y)`` image pixels.
+
+ Contours are already stored in pixels, so this is an axis swap and a dtype
+ change — deliberately NOT a division by ``scale``.
+ """
+ c = particles.contour_at(int(index))
+ if len(c) < 3:
+ return np.zeros((0, 2), np.float32)
+ return np.column_stack([c[:, 1], c[:, 0]]).astype(np.float32)
+
+
+def _axis_scale_offset(plot) -> tuple[float, float, float, float]:
+ """``(x_scale, x_offset, y_scale, y_offset)`` of the plot's displayed signal.
+
+ Used to convert a click's ``xdata``/``ydata`` (physical) back to the image
+ pixels every marker lives in. Falls back to an identity mapping when the plot
+ has no calibrated axes, which is also what anyplotlib does — with no
+ ``x_axis`` array on the panel it reports ``xdata == img_x``.
+ """
+ try:
+ state = getattr(plot, "plot_state", None)
+ sig = getattr(state, "current_signal", None)
+ ax = sig.axes_manager.signal_axes
+ return (float(ax[0].scale) or 1.0, float(ax[0].offset),
+ float(ax[1].scale) or 1.0, float(ax[1].offset))
+ except Exception:
+ return 1.0, 0.0, 1.0, 0.0
+
+
+def _navigator_selectors_for(tree, plot) -> list:
+ """Navigation selectors that drive *plot*.
+
+ Same resolution (and the same composite-selector dedup) as
+ ``vector_overlay._navigator_selectors_for``: a composite navigator selector
+ exposes both its crosshair and its region sub-selector, and registering the
+ hook on both fires two redraws per navigator move.
+ """
+ npm = getattr(tree, "navigator_plot_manager", None)
+ if npm is None:
+ return []
+ out = [sel for sel in npm.all_navigation_selectors
+ if plot in getattr(sel, "active_children", [])]
+ out = out or list(npm.all_navigation_selectors)
+ seen, uniq = set(), []
+ for sel in out:
+ key = id(getattr(sel, "parent", sel) or sel)
+ if key in seen:
+ continue
+ seen.add(key)
+ uniq.append(sel)
+ return uniq
+
+
+# ── marker pushes ────────────────────────────────────────────────────────────
+
+def _push_groups(plot2d, updates: dict) -> None:
+ """Apply many marker-group updates with ONE panel push.
+
+ ``MarkerGroup.set`` re-serialises the whole marker registry and pushes the
+ panel state on every call. A single navigator move touches the seven fill
+ groups, the seven head-dot groups and up to twenty-four trail groups, so
+ doing it through ``set`` would be ~30 full serialisations (and, in the
+ Electron host, ~30 PLOTAPP lines) per frame. Mutating the group dicts and
+ pushing once is the identical result for a thirtieth of the transport.
+
+ Falls back to per-group ``set`` if ``_push_markers`` ever goes away.
+ """
+ if not updates:
+ return
+ pusher = getattr(plot2d, "_push_markers", None)
+ if pusher is None:
+ for group, kwargs in updates.items():
+ try:
+ group.set(**kwargs)
+ except Exception as exc:
+ log.debug("[particles] marker set failed: %s", exc)
+ return
+ for group, kwargs in updates.items():
+ try:
+ group._data.update(kwargs)
+ except Exception as exc:
+ log.debug("[particles] marker update failed: %s", exc)
+ try:
+ pusher()
+ except Exception as exc:
+ log.debug("[particles] marker push failed: %s", exc)
+
+
+# ── the overlay ──────────────────────────────────────────────────────────────
+
+class ParticleOverlay:
+ """Filled, track-coloured particle outlines on a signal plot, live.
+
+ Draws frame *t*'s particles as 25%-filled polygons with a 1 px outline,
+ coloured by ``track_id % 6``; the selected particle gets a white outline, its
+ id and a short property readout. Optional trails fade the last *N* frames of
+ each track behind a head dot marking "here, now".
+
+ Parameters
+ ----------
+ plot
+ The :class:`~spyde.drawing.plots.plot.Plot` to draw on (the label movie's
+ signal plot, or the source movie's).
+ particles
+ The :class:`~spyde.signals.particles.SpyDEParticles` store. Held by
+ reference and MUTATED IN PLACE by the edit methods, because the lazy
+ label movie closes over this exact object (see
+ ``particle_tree.open_particle_tree``).
+ events
+ Optional :class:`~spyde.particles.track.ParticleEvent` list. Carried so a
+ consumer reading the overlay has them to hand; the navigator lanes read
+ ``tree.particle_events`` directly. Nothing here DRAWS them yet — plan C2's
+ third surface (a birth/death badge flashed on the frame during playback)
+ is not implemented.
+ nav_map
+ The tree's ``nav_map`` (particle frame → source navigation index). Given
+ only when the overlay is drawn on the SOURCE movie, where the navigator
+ indexes the source grid; it is inverted internally.
+ frame_provider
+ ``fn(t) -> (h, w) ndarray`` returning the intensity image for frame *t*.
+ Used to re-measure after a merge or split. ``None`` leaves the intensity
+ columns NaN on edited rows rather than inventing them.
+ trail_frames, show_trails
+ Trailing window length and whether trails start on.
+ on_select, on_edit
+ Callbacks fired after a selection change / an edit, so a caret or table
+ dock can follow. Both are called with no arguments.
+ """
+
+ def __init__(self, plot, particles, *, events: Sequence = (),
+ nav_map=None, frame_provider: Callable[[int], np.ndarray] | None = None,
+ trail_frames: int = DEFAULT_TRAIL_FRAMES, show_trails: bool = False,
+ on_select: Callable[[], None] | None = None,
+ on_edit: Callable[[], None] | None = None,
+ name: str = "particles"):
+ self.plot = plot
+ self.particles = particles
+ self.events = list(events or ())
+ self.name = str(name)
+ self.frame_provider = frame_provider
+ self.trail_frames = max(1, int(trail_frames))
+ self.show_trails = bool(show_trails)
+ self.on_select = on_select
+ self.on_edit = on_edit
+
+ self.tree = None
+ self.selected: list[int] = []
+ self.hovered: int | None = None
+ self._frame = 0
+ self._hidden = False
+ self._groups: dict[str, Any] = {}
+ self._selectors: list = []
+ self._handlers: list = []
+ self._region_widget = None
+ # Latest-wins: a navigator move computes off the main thread and marshals
+ # the push; teardown bumps FIRST so a superseded payload never lands.
+ self._gen = 0
+ # nav index → particle frame, only when the two grids differ.
+ self._nav_to_frame = None
+ if nav_map is not None:
+ nav = np.asarray(nav_map, np.int64).ravel()
+ if not np.array_equal(nav, np.arange(nav.size)):
+ self._nav_to_frame = {int(v): i for i, v in enumerate(nav)}
+ self._prev_settled_ms = None
+ # The store's identity at the last redraw. An edit rebuilds the buffers,
+ # so any cached global index the caret is holding is stale; the counter
+ # is what a consumer compares against.
+ self.revision = 0
+
+ # ── attach / detach ──────────────────────────────────────────────────────
+
+ def attach(self, tree) -> "ParticleOverlay":
+ """Create the marker groups, wire the navigator and the click handlers."""
+ self.tree = tree
+ plot2d = getattr(self.plot, "_plot2d", None)
+ if plot2d is None:
+ # Cosmetic, so not fatal — but say so, rather than leaving the user
+ # with a silently missing overlay (vector_overlay's rule).
+ log.warning("[particles] overlay skipped: plot has no live 2-D plot "
+ "(figure iframe not loaded?)")
+ return self
+ self._build_groups(plot2d)
+ self._wire_events(plot2d)
+ self._wire_navigator(tree)
+ self._redraw()
+ return self
+
+ def _build_groups(self, plot2d) -> None:
+ """One group per colour (see the module docstring), created empty.
+
+ CREATION ORDER IS DRAW ORDER, and it is by marker TYPE, not by group:
+ ``MarkerRegistry.to_wire_list`` walks its type dicts in the order they
+ were first touched and flattens each one's groups. So the first
+ ``add_lines`` puts every trail underneath every polygon, and the first
+ ``add_circles`` puts every head dot on top of both — which is the z-order
+ this wants (a head dot hidden under a fill is not a head dot). Within the
+ polygons type the selected outline is added after the fills, so it draws
+ over them.
+ """
+ empty_poly: list = []
+ empty_off = np.zeros((0, 2), np.float32)
+ empty_seg = np.zeros((0, 2, 2), np.float32)
+ colors = list(TRACK_COLORS) + [UNTRACKED_COLOR]
+
+ for i, color in enumerate(colors): # bottom: trails
+ for step, alpha in enumerate(trail_alphas(TRAIL_FADE_STEPS)):
+ self._groups[f"trail{i}_{step}"] = plot2d.add_lines(
+ empty_seg, name=f"{self.name}_trail_{i}_{step}",
+ edgecolors=fade(color, alpha), linewidths=1.5,
+ transform="data")
+ for i, color in enumerate(colors): # then the fills
+ self._groups[f"fill{i}"] = plot2d.add_polygons(
+ empty_poly, name=f"{self.name}_fill_{i}", facecolors=color,
+ edgecolors=color, linewidths=OUTLINE_WIDTH, alpha=FILL_ALPHA,
+ transform="data")
+ self._groups["selected"] = plot2d.add_polygons(
+ empty_poly, name=f"{self.name}_selected", facecolors=None,
+ edgecolors=SELECTED_COLOR, linewidths=SELECTED_WIDTH,
+ transform="data")
+ for i, color in enumerate(colors): # then head dots
+ self._groups[f"head{i}"] = plot2d.add_circles(
+ empty_off, name=f"{self.name}_head_{i}", radius=HEAD_RADIUS_PX,
+ facecolors=color, edgecolors=color, linewidths=1.0, alpha=1.0,
+ transform="data")
+ self._groups["labels"] = plot2d.add_texts( # top: labels
+ empty_off, [], name=f"{self.name}_labels", color=SELECTED_COLOR,
+ fontsize=11, transform="data")
+
+ def _wire_events(self, plot2d) -> None:
+ """Click-to-select and hover-to-label.
+
+ ``double_click`` for the discrete pick, matching the strain reference
+ picker: a single click on a 2-D panel is ambiguous with panning, and
+ anyplotlib's own pan/click disambiguation has moved between versions.
+
+ Hover runs off ``pointer_settled``, not ``pointer_move``: a label that
+ needed an IPC round trip per mouse move would put hundreds of messages a
+ second on the same stdout line protocol the nav painter uses. Settling is
+ also the correct semantics — the label answers "what am I pointing at",
+ which is only a question once the pointer has stopped.
+ """
+ from spyde.drawing.selectors.base_selector import event_handler_fn
+ for event_type, method in (("double_click", self._on_click),
+ ("pointer_settled", self._on_settled)):
+ handler = event_handler_fn(method)
+ self._handlers.append(handler)
+ try:
+ plot2d.add_event_handler(handler, event_type)
+ except Exception as exc:
+ log.debug("[particles] wiring %s failed: %s", event_type, exc)
+ try:
+ self._prev_settled_ms = plot2d._state.get("pointer_settled_ms")
+ plot2d.configure_pointer_settled(180)
+ except Exception as exc:
+ log.debug("[particles] enabling pointer_settled failed: %s", exc)
+
+ def _wire_navigator(self, tree) -> None:
+ self._selectors = _navigator_selectors_for(tree, self.plot)
+ for sel in self._selectors:
+ if self._on_indices not in sel.index_hooks:
+ sel.index_hooks.append(self._on_indices)
+ if getattr(sel, "current_indices", None) is not None:
+ self._frame = frame_from_indices(sel.current_indices,
+ self._nav_to_frame)
+ if not self._selectors:
+ log.warning("[particles] attached with NO navigator selectors — the "
+ "overlay will not follow the frame")
+
+ def remove(self) -> None:
+ """Detach every hook, drop every marker group. Idempotent."""
+ self._gen += 1 # teardown bumps FIRST (README §6)
+ for sel in self._selectors:
+ if self._on_indices in sel.index_hooks:
+ sel.index_hooks.remove(self._on_indices)
+ self._selectors = []
+ self.set_region_select(False)
+ plot2d = getattr(self.plot, "_plot2d", None)
+ if plot2d is not None and self._prev_settled_ms is not None:
+ try:
+ plot2d.configure_pointer_settled(int(self._prev_settled_ms))
+ except Exception as exc:
+ log.debug("[particles] restoring pointer_settled failed: %s", exc)
+ for group in self._groups.values():
+ try:
+ group.remove()
+ except Exception as exc:
+ log.debug("[particles] removing marker group failed: %s", exc)
+ self._groups = {}
+ # anyplotlib has no remove_event_handler; dropping our references lets
+ # the wrappers (weakly registered) be collected, and an empty _groups
+ # makes every handler a no-op in the meantime.
+ self._handlers = []
+
+ # ── navigator ────────────────────────────────────────────────────────────
+
+ def _on_indices(self, indices) -> None:
+ """Navigation moved. Runs on the ``_NavDispatcher`` thread."""
+ frame = frame_from_indices(indices, self._nav_to_frame)
+ if frame == self._frame:
+ return
+ self._frame = frame
+ self._request_redraw()
+
+ def _request_redraw(self) -> None:
+ """Compute the payload here, push it on the asyncio main thread.
+
+ The payload is pure numpy over one frame (plus the trail window), so it
+ is cheap enough to build on whichever thread asked. The PUSH is a figure
+ update and must be marshalled — CLAUDE.md's threading contract.
+ """
+ if self._hidden or not self._groups:
+ return
+ gen = self._gen = self._gen + 1
+ payload = self._payload(self._frame)
+ session = getattr(self.plot, "session", None)
+ dispatch = getattr(session, "_dispatch_to_main", None)
+ if dispatch is None:
+ self._apply(payload, gen)
+ return
+ dispatch(lambda: self._apply(payload, gen))
+
+ def _apply(self, payload: dict, gen: int) -> None:
+ if gen != self._gen or not self._groups:
+ return # superseded, or torn down
+ plot2d = getattr(self.plot, "_plot2d", None)
+ if plot2d is None:
+ return
+ _push_groups(plot2d, {self._groups[key]: kwargs
+ for key, kwargs in payload.items()
+ if key in self._groups})
+
+ def _redraw(self) -> None:
+ """Rebuild and push immediately (an edit / a selection, not a nav move)."""
+ if self._hidden or not self._groups:
+ return
+ gen = self._gen = self._gen + 1
+ self._apply(self._payload(self._frame), gen)
+
+ # ── the payload ──────────────────────────────────────────────────────────
+
+ @staticmethod
+ def _empty_payload() -> dict:
+ """Every group key mapped to the contents that draw nothing."""
+ out: dict[str, dict] = {}
+ for i in range(len(TRACK_COLORS) + 1):
+ out[f"fill{i}"] = {"vertices_list": []}
+ out[f"head{i}"] = {"offsets": np.zeros((0, 2), np.float32)}
+ for step in range(TRAIL_FADE_STEPS):
+ out[f"trail{i}_{step}"] = {"segments": np.zeros((0, 2, 2), np.float32)}
+ out["selected"] = {"vertices_list": []}
+ out["labels"] = {"offsets": np.zeros((0, 2), np.float32), "texts": []}
+ return out
+
+ def _payload(self, t: int) -> dict:
+ """Every marker group's contents for frame *t*. Pure — no plot access.
+
+ Returned as ``{group_key: set_kwargs}`` so the whole overlay is testable
+ without a live figure, which is the only way to assert the head-dot rule
+ and the pixel conversion in a headless suite.
+ """
+ n_colors = len(TRACK_COLORS) + 1
+ n_steps = TRAIL_FADE_STEPS
+ out: dict[str, dict] = {}
+ fills: list[list] = [[] for _ in range(n_colors)]
+ heads: list[list] = [[] for _ in range(n_colors)]
+ trails: list[list[list]] = [[[] for _ in range(n_steps)]
+ for _ in range(n_colors)]
+
+ if 0 <= int(t) < self.particles.n_frames and self.particles.has_masks:
+ for gi in self.particles.indices_at(int(t)):
+ poly = contour_xy(self.particles, int(gi))
+ if len(poly) < 3:
+ continue
+ fills[self._bucket(gi)].append(poly)
+
+ if self.show_trails:
+ self._fill_trails(int(t), trails, heads)
+
+ for i in range(n_colors):
+ out[f"fill{i}"] = {"vertices_list": fills[i]}
+ out[f"head{i}"] = {"offsets": (np.asarray(heads[i], np.float32)
+ if heads[i]
+ else np.zeros((0, 2), np.float32))}
+ for step in range(n_steps):
+ segs = trails[i][step]
+ out[f"trail{i}_{step}"] = {
+ "segments": (np.asarray(segs, np.float32) if segs
+ else np.zeros((0, 2, 2), np.float32))}
+
+ out["selected"] = {"vertices_list": self._selected_polys(int(t))}
+ offsets, texts = self._labels(int(t))
+ out["labels"] = {"offsets": offsets, "texts": texts}
+ return out
+
+ def _bucket(self, gi: int) -> int:
+ """Colour-group index for a global particle row (last = untracked)."""
+ tid = int(self.particles.flat_buffer[int(gi), COL["track_id"]])
+ if tid < 0:
+ return len(TRACK_COLORS)
+ return tid % len(TRACK_COLORS)
+
+ def _fill_trails(self, t: int, trails, heads) -> None:
+ """Trail segments + head dots for the window ending at frame *t*.
+
+ **A dead track draws no head dot.** The dot means "the particle is HERE
+ NOW", so it is drawn only for a track with a detection at exactly *t* —
+ which covers both a track that died before *t* and one inside its
+ ``memory`` gap, without either needing to be special-cased or the
+ ``LinkResult`` needing to be around. Found by looking at a render: a
+ track that died at frame 16 was still painting a head dot at 18 because
+ its trajectory still intersected the trailing window, and it read as a
+ real particle the segmenter had stopped filling (plan C3).
+
+ Walks the window's FRAMES rather than indexing tracks, for the reason
+ ``track._extract_events`` gives: a ``{track: row}`` map over a whole
+ movie is 1.5M dict entries at the plan's target scale, to answer a
+ question that only ever spans the last few frames.
+ """
+ scale = float(self.particles.scale) or 1.0
+ t0 = max(0, t - self.trail_frames + 1)
+ per_track: dict[int, list[tuple[int, float, float]]] = {}
+ for f in range(t0, min(t, self.particles.n_frames - 1) + 1):
+ rows = self.particles.at(f)
+ if len(rows) == 0:
+ continue
+ xs = rows[:, COL["x"]] / scale
+ ys = rows[:, COL["y"]] / scale
+ tids = rows[:, COL["track_id"]].astype(np.int64)
+ for k in range(len(rows)):
+ tid = int(tids[k])
+ if tid < 0:
+ continue # an untracked row has no trajectory
+ per_track.setdefault(tid, []).append((f, float(xs[k]), float(ys[k])))
+
+ n_steps = len(trail_alphas(TRAIL_FADE_STEPS))
+ for tid, pts in per_track.items():
+ bucket = tid % len(TRACK_COLORS)
+ for (_f0, x0, y0), (f1, x1, y1) in zip(pts, pts[1:]):
+ # Age is measured at the segment's NEWER end, so the piece next
+ # to the head is the brightest one.
+ step = min(n_steps - 1,
+ (t - f1) * n_steps // max(1, self.trail_frames))
+ trails[bucket][step].append([[x0, y0], [x1, y1]])
+ if pts and pts[-1][0] == t:
+ heads[bucket].append([pts[-1][1], pts[-1][2]])
+
+ def _selected_polys(self, t: int) -> list:
+ """Outlines of the selected particles that are visible in frame *t*."""
+ if not self.selected or not self.particles.has_masks:
+ return []
+ out = []
+ for gi in self.selected:
+ if not self._in_frame(gi, t):
+ continue
+ poly = contour_xy(self.particles, gi)
+ if len(poly) >= 3:
+ out.append(poly)
+ return out
+
+ def _labels(self, t: int) -> tuple[np.ndarray, list[str]]:
+ """Text labels — SELECTION and HOVER only.
+
+ Always-on ids were rejected in the plan for a measured reason: legible on
+ a nine-particle mock-up, a wall of numbers at 500 particles per frame.
+ """
+ shown = list(self.selected)
+ if self.hovered is not None and self.hovered not in shown:
+ shown.append(self.hovered)
+ shown = [gi for gi in shown if self._in_frame(gi, t)]
+ if not shown:
+ return np.zeros((0, 2), np.float32), []
+ rows = self.particles.flat_buffer[np.asarray(shown, np.int64)]
+ scale = float(self.particles.scale) or 1.0
+ offsets = centroids_px(rows, scale)
+ # Anchor OUTSIDE the body, up and to the right: anyplotlib draws text
+ # left-aligned / top-baselined from the offset, so anchoring on the
+ # centroid lays the readout across the particle it describes (seen in the
+ # app). One body radius clears it.
+ radii = rows[:, COL["equiv_diameter"]] / (2.0 * scale)
+ radii = np.where(np.isfinite(radii), radii, 0.0) + 2.0
+ offsets[:, 0] += radii
+ offsets[:, 1] -= radii
+ texts = [self.describe(gi) for gi in shown]
+ return offsets, texts
+
+ def _in_frame(self, gi: int, t: int) -> bool:
+ if not 0 <= int(gi) < self.particles.n_particles:
+ return False
+ lo, hi = self.particles.t_offsets[int(t)], self.particles.t_offsets[int(t) + 1]
+ return bool(lo <= int(gi) < hi)
+
+ def describe(self, gi: int) -> str:
+ """One-line id + property readout for particle *gi*."""
+ row = self.particles.flat_buffer[int(gi)]
+ tid = int(row[COL["track_id"]])
+ head = f"#{int(gi)}" if tid < 0 else f"track {tid}"
+ units = self.particles.units
+ parts = [head]
+ for name in READOUT_COLUMNS:
+ value = float(row[COL[name]])
+ if not np.isfinite(value):
+ continue
+ suffix = {"area": f" {units}²", "equiv_diameter": f" {units}"}.get(name, "")
+ parts.append(f"{name.replace('_', ' ')} {value:.3g}{suffix}")
+ return " ".join(parts)
+
+ # ── display state ────────────────────────────────────────────────────────
+
+ def set_visible(self, visible: bool) -> None:
+ """Show or hide every group.
+
+ A hidden overlay still TRACKS the navigator (``_frame`` keeps moving), so
+ re-showing draws the frame the user is on rather than the one they were
+ on when it went away — the same contract ``vector_overlay`` keeps.
+ """
+ self._hidden = not bool(visible)
+ if not self._hidden:
+ self._redraw()
+ return
+ gen = self._gen = self._gen + 1
+ self._apply(self._empty_payload(), gen)
+
+ def set_trails(self, enabled: bool, frames: int | None = None) -> None:
+ """Toggle trails and set the trailing window length."""
+ self.show_trails = bool(enabled)
+ if frames is not None:
+ self.trail_frames = max(1, int(frames))
+ self._redraw()
+
+ def set_frame(self, t: int) -> None:
+ """Programmatically move to particle frame *t* (playback, a table jump)."""
+ self._frame = int(t)
+ self._redraw()
+
+ # ── selection ────────────────────────────────────────────────────────────
+
+ def select(self, indices: Iterable[int] | int | None) -> list[int]:
+ """Select particles by GLOBAL index — the external hook a table row uses.
+
+ Returns the resulting selection. Indices outside the store are dropped
+ rather than raising: the caret's row list can lag an edit by a frame, and
+ a stale row must not take the backend down with it.
+ """
+ if indices is None:
+ wanted: list[int] = []
+ elif isinstance(indices, (int, np.integer)):
+ wanted = [int(indices)]
+ else:
+ wanted = [int(i) for i in indices]
+ n = self.particles.n_particles
+ self.selected = [i for i in wanted if 0 <= i < n]
+ self._after_select()
+ return list(self.selected)
+
+ def select_track(self, track_id: int, *, frame: int | None = None) -> list[int]:
+ """Select a track's detection in one frame (default: the current one)."""
+ t = self._frame if frame is None else int(frame)
+ if not 0 <= t < self.particles.n_frames:
+ return self.select([])
+ rows = self.particles.at(t)
+ hit = np.nonzero(rows[:, COL["track_id"]].astype(np.int64) == int(track_id))[0]
+ base = int(self.particles.t_offsets[t])
+ return self.select([base + int(k) for k in hit])
+
+ def select_region(self, x0, y0, x1, y1, *, frame: int | None = None) -> list[int]:
+ """Rubber-band selection: every particle in *frame* whose centroid falls
+ inside the image-pixel box ``(x0, y0)-(x1, y1)``."""
+ t = self._frame if frame is None else int(frame)
+ if not 0 <= t < self.particles.n_frames:
+ return self.select([])
+ rows = self.particles.at(t)
+ if len(rows) == 0:
+ return self.select([])
+ pts = centroids_px(rows, self.particles.scale)
+ lo_x, hi_x = sorted((float(x0), float(x1)))
+ lo_y, hi_y = sorted((float(y0), float(y1)))
+ inside = ((pts[:, 0] >= lo_x) & (pts[:, 0] <= hi_x)
+ & (pts[:, 1] >= lo_y) & (pts[:, 1] <= hi_y))
+ base = int(self.particles.t_offsets[t])
+ return self.select([base + int(k) for k in np.nonzero(inside)[0]])
+
+ def pick(self, px: float, py: float, *, frame: int | None = None) -> int | None:
+ """Nearest-centroid hit test at image-pixel ``(px, py)``.
+
+ The hit radius is the particle's OWN ``equiv_diameter`` (floored at
+ 4 px), not a constant: a fixed radius that feels right for a 50 px body
+ makes a 5 px one unclickable, and one sized for the small body picks a
+ neighbour when the field is dense. Same reasoning as the linker's
+ adaptive merge radius.
+ """
+ t = self._frame if frame is None else int(frame)
+ if not 0 <= t < self.particles.n_frames:
+ return None
+ rows = self.particles.at(t)
+ if len(rows) == 0:
+ return None
+ scale = float(self.particles.scale) or 1.0
+ pts = centroids_px(rows, scale)
+ d2 = (pts[:, 0] - float(px)) ** 2 + (pts[:, 1] - float(py)) ** 2
+ k = int(np.argmin(d2))
+ radius = rows[k, COL["equiv_diameter"]] / scale
+ radius = max(4.0, float(radius) if np.isfinite(radius) else 0.0)
+ if d2[k] > radius ** 2:
+ return None
+ return int(self.particles.t_offsets[t]) + k
+
+ def clear_selection(self) -> list[int]:
+ return self.select([])
+
+ def _after_select(self) -> None:
+ self._redraw()
+ if self.on_select is not None:
+ try:
+ self.on_select()
+ except Exception as exc:
+ log.debug("[particles] on_select failed: %s", exc)
+
+ def _on_click(self, event=None) -> None:
+ if event is None or not self._groups:
+ return
+ try:
+ px, py = self._event_px(event)
+ except Exception as exc:
+ log.debug("[particles] click had no usable coordinates: %s", exc)
+ return
+ hit = self.pick(px, py)
+ self.select([] if hit is None else [hit])
+
+ def _on_settled(self, event=None) -> None:
+ if event is None or not self._groups:
+ return
+ try:
+ px, py = self._event_px(event)
+ except Exception:
+ return
+ hit = self.pick(px, py)
+ if hit == self.hovered:
+ return
+ self.hovered = hit
+ self._redraw()
+
+ def _event_px(self, event) -> tuple[float, float]:
+ """A pointer event's position in IMAGE PIXELS.
+
+ anyplotlib's Python ``Event`` carries only ``xdata``/``ydata`` (the JS
+ payload's ``img_x``/``img_y`` have no field on the dataclass and are
+ dropped), and those are the panel's PHYSICAL data coordinates. Markers
+ are in image pixels, so the click is converted back through the plot's
+ own axes — not through ``particles.scale``, which is the store's
+ calibration and need not be the displayed signal's.
+ """
+ xs, xo, ys, yo = _axis_scale_offset(self.plot)
+ return (float(event.xdata) - xo) / xs, (float(event.ydata) - yo) / ys
+
+ # ── rubber band ──────────────────────────────────────────────────────────
+
+ def set_region_select(self, enabled: bool) -> None:
+ """Show/hide the rubber-band rectangle used for bulk selection.
+
+ A widget rather than a drag on the canvas: the panel's own drag is pan,
+ and anyplotlib's rectangle widget already owns the handles, the clamping
+ and the pointer events.
+ """
+ plot2d = getattr(self.plot, "_plot2d", None)
+ if not enabled:
+ if self._region_widget is not None and plot2d is not None:
+ try:
+ plot2d.remove_widget(self._region_widget)
+ except Exception as exc:
+ log.debug("[particles] removing region widget failed: %s", exc)
+ self._region_widget = None
+ return
+ if self._region_widget is not None or plot2d is None:
+ return
+ h, w = self.particles.frame_shape
+ try:
+ widget = plot2d.add_rectangle_widget(
+ x=w * 0.25, y=h * 0.25, w=w * 0.5, h=h * 0.5, color=SELECTED_COLOR)
+ except Exception as exc:
+ log.debug("[particles] adding region widget failed: %s", exc)
+ return
+ from spyde.drawing.selectors.base_selector import event_handler_fn
+ handler = event_handler_fn(self._on_region)
+ self._handlers.append(handler)
+ try:
+ widget.add_event_handler(handler, "pointer_up")
+ except Exception as exc:
+ log.debug("[particles] wiring region widget failed: %s", exc)
+ self._region_widget = widget
+ self._on_region()
+
+ def _on_region(self, _event=None) -> None:
+ widget = self._region_widget
+ if widget is None:
+ return
+ try:
+ x, y = float(widget.x), float(widget.y)
+ w, h = float(widget.w), float(widget.h)
+ except Exception as exc:
+ log.debug("[particles] reading region widget failed: %s", exc)
+ return
+ self.select_region(x, y, x + w, y + h)
+
+ # ── editing ──────────────────────────────────────────────────────────────
+
+ def delete(self, indices: Iterable[int] | None = None) -> int:
+ """Delete particles (default: the selection). Returns how many went."""
+ idx = sorted({int(i) for i in (self.selected if indices is None else indices)})
+ if not idx:
+ return 0
+ removed = delete_particles(self.particles, idx)
+ self._record("delete", indices=idx, frame=self._frame)
+ self.selected = []
+ self._after_edit()
+ return removed
+
+ def merge(self, indices: Iterable[int] | None = None) -> int:
+ """Merge particles into one, re-measure, return its new global index."""
+ idx = sorted({int(i) for i in (self.selected if indices is None else indices)})
+ if len(idx) < 2:
+ raise ValueError("merge needs at least two particles")
+ new_index = merge_particles(self.particles, idx,
+ frame_image=self._frame_image_for(idx[0]))
+ self._record("merge", indices=idx, frame=self._frame, result=[new_index])
+ self.selected = [new_index]
+ self._after_edit()
+ return new_index
+
+ def split(self, index: int | None = None, line=None) -> tuple[int, int]:
+ """Split one particle along *line* and re-measure both halves.
+
+ *line* is ``((x0, y0), (x1, y1))`` in IMAGE PIXELS — the space the drawn
+ line widget reports in.
+ """
+ if index is None:
+ if len(self.selected) != 1:
+ raise ValueError("split needs exactly one selected particle")
+ index = self.selected[0]
+ if line is None:
+ raise ValueError("split needs a cut line")
+ pair = split_particle(self.particles, int(index), line,
+ frame_image=self._frame_image_for(int(index)))
+ self._record("split", indices=[int(index)], frame=self._frame,
+ line=[[float(v) for v in pt] for pt in line],
+ result=list(pair))
+ self.selected = list(pair)
+ self._after_edit()
+ return pair
+
+ def _frame_image_for(self, gi: int):
+ if self.frame_provider is None:
+ return None
+ t = int(self.particles.flat_buffer[int(gi), COL["t"]])
+ try:
+ return np.asarray(self.frame_provider(t))
+ except Exception as exc:
+ log.debug("[particles] frame_provider(%d) failed: %s", t, exc)
+ return None
+
+ def _record(self, kind: str, **fields) -> dict:
+ return record_edit(self.tree, self.particles, kind, **fields)
+
+ def _after_edit(self) -> None:
+ self.revision += 1
+ self.hovered = None
+ self._redraw()
+ if self.on_edit is not None:
+ try:
+ self.on_edit()
+ except Exception as exc:
+ log.debug("[particles] on_edit failed: %s", exc)
+
+
+def attach_particle_overlay(plot, particles, tree, **kwargs) -> ParticleOverlay:
+ """Attach a :class:`ParticleOverlay` to *plot*, wired to *tree*'s navigator.
+
+ Stored on the tree as ``tree._particle_overlay`` via
+ ``lifecycle.replace_tree_attr``, so re-running never stacks two overlays and
+ ``BaseSignalTree.close()`` reaps it.
+ """
+ from spyde.actions.lifecycle import replace_tree_attr
+ return replace_tree_attr(
+ tree, "_particle_overlay",
+ lambda: ParticleOverlay(plot, particles, **kwargs).attach(tree))
+
+
+# ── edits on the store ───────────────────────────────────────────────────────
+#
+# Separate from the overlay on purpose: an edit is a transformation of the CSR
+# table and nothing about it needs a figure, so it is testable (and scriptable)
+# on its own. Every one of them mutates the store IN PLACE — the lazy label movie
+# closes over that exact object (``particle_tree.open_particle_tree``), so
+# swapping in a new SpyDEParticles would leave the open window rendering the old
+# contours forever. That is the same trap ``particles_action._adopt`` exists for.
+
+def _row_frames(particles) -> np.ndarray:
+ """``(n,)`` int64 frame index per row, from the CSR pointers."""
+ return np.repeat(np.arange(particles.n_frames, dtype=np.int64),
+ np.diff(particles.t_offsets))
+
+
+def _splice(particles, *, drop: Sequence[int] = (),
+ add: Sequence[tuple[int, np.ndarray, np.ndarray | None]] = ()) -> np.ndarray:
+ """Rebuild the store in place with rows removed and/or added.
+
+ *add* is a sequence of ``(frame, row, contour)``. Added rows land at the END
+ of their frame's block (a stable sort by frame over the kept rows followed by
+ the new ones), so existing global indices shift by at most the deletions
+ ahead of them and never by an insertion into the middle of a frame.
+
+ Returns the ``(len(add),)`` global indices the added rows ended up at.
+
+ Vectorised rather than looped: the plan's target is 1.5M rows, and an edit
+ that walked them in Python would take seconds for a single click. The only
+ loop is over the handful of rows being added.
+ """
+ n = particles.n_particles
+ keep = np.ones(n, bool)
+ if len(drop):
+ keep[np.asarray(list(drop), np.int64)] = False
+
+ frames = _row_frames(particles)
+ kept_frames = frames[keep]
+ kept_rows = particles.flat_buffer[keep]
+
+ add = list(add)
+ add_frames = np.asarray([int(f) for f, _r, _c in add], np.int64)
+ add_rows = (np.asarray([r for _f, r, _c in add], np.float32).reshape(-1, N_COLUMNS)
+ if add else np.zeros((0, N_COLUMNS), np.float32))
+
+ all_frames = np.concatenate([kept_frames, add_frames])
+ all_rows = np.concatenate([kept_rows, add_rows], axis=0)
+ order = np.argsort(all_frames, kind="stable")
+
+ particles.flat_buffer = np.ascontiguousarray(all_rows[order], np.float32)
+ counts = np.bincount(all_frames, minlength=particles.n_frames)
+ particles.t_offsets = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64)
+
+ if particles.contours is not None:
+ lens = np.diff(particles.contour_offsets)
+ starts = particles.contour_offsets[:-1]
+ add_polys = [np.zeros((0, 2), np.int16) if c is None
+ else np.asarray(c, np.int16).reshape(-1, 2) for _f, _r, c in add]
+ add_lens = np.asarray([len(p) for p in add_polys], np.int64)
+ all_lens = np.concatenate([lens[keep], add_lens])
+ # Point sources: kept polygons index the OLD contour array; added ones
+ # index a small appended block. One concatenation, then one gather.
+ pool = np.concatenate([particles.contours] + add_polys, axis=0) \
+ if add_polys else particles.contours
+ add_starts = len(particles.contours) + np.concatenate(
+ [[0], np.cumsum(add_lens)])[:-1].astype(np.int64)
+ all_starts = np.concatenate([starts[keep], add_starts])
+
+ new_lens = all_lens[order]
+ new_off = np.concatenate([[0], np.cumsum(new_lens)]).astype(np.int64)
+ within = np.arange(int(new_off[-1])) - np.repeat(new_off[:-1], new_lens)
+ src = np.repeat(all_starts[order], new_lens) + within
+ particles.contours = np.ascontiguousarray(pool[src], np.int16)
+ particles.contour_offsets = new_off
+
+ # Where each added row landed: its position in `order`.
+ rank = np.empty(order.size, np.int64)
+ rank[order] = np.arange(order.size, dtype=np.int64)
+ return rank[len(kept_rows):]
+
+
+def delete_particles(particles, indices: Sequence[int]) -> int:
+ """Drop rows from the store, in place. Returns the number removed."""
+ idx = sorted({int(i) for i in indices if 0 <= int(i) < particles.n_particles})
+ if not idx:
+ return 0
+ _splice(particles, drop=idx)
+ return len(idx)
+
+
+def _full_mask(particles, index: int) -> np.ndarray:
+ """One particle's boolean mask, at full frame size."""
+ mask, (y0, x0, y1, x1) = particles.mask_at(int(index))
+ h, w = particles.frame_shape
+ out = np.zeros((h, w), bool)
+ out[y0:y1, x0:x1] |= mask
+ return out
+
+
+def _remeasure(labels: np.ndarray, frame_image, t: int, scale: float):
+ """``measure_frame`` on a purpose-built label image. Returns (rows, contours)."""
+ from spyde.particles import measure_frame
+ return measure_frame(labels, frame_image, t=int(t), scale=float(scale))
+
+
+def merge_particles(particles, indices: Sequence[int], *, frame_image=None) -> int:
+ """Union the masks of two or more particles, re-measure, return the new index.
+
+ Every input must be in the SAME frame — merging across frames is not an
+ editing operation, it is a linking one, and doing it here would silently
+ produce a row whose ``t`` contradicts its CSR block.
+
+ The merged row inherits the ``track_id`` of the LARGEST input, matching the
+ linker's merge semantics (a large body absorbs a smaller one, ``track.py``);
+ the absorbed track simply stops, which is what a re-link would also conclude.
+ """
+ idx = sorted({int(i) for i in indices})
+ if len(idx) < 2:
+ raise ValueError("merge needs at least two particles")
+ if not particles.has_masks:
+ raise ValueError("cannot merge without outlines "
+ "(segmentation ran with store_masks=False)")
+ frames = {int(particles.flat_buffer[i, COL["t"]]) for i in idx}
+ if len(frames) != 1:
+ raise ValueError(f"merge needs one frame; got {sorted(frames)}")
+ t = frames.pop()
+
+ from skimage.measure import label as connected_components
+
+ union = np.zeros(particles.frame_shape, bool)
+ for i in idx:
+ union |= _full_mask(particles, i)
+ # Connected-component labelling, NOT ``union.astype(int32)``: casting the
+ # boolean union gives every pixel label 1, and ``regionprops`` measures a
+ # label as ONE region whether or not it is connected — so two discs on
+ # opposite sides of the frame would merge "successfully" into a row whose
+ # centroid sits in the empty space between them. 8-connectivity, because two
+ # bodies meeting at a corner are touching.
+ components = connected_components(union, connectivity=2)
+ n_components = int(components.max())
+ if n_components == 0:
+ raise ValueError("the merged region measured as empty")
+ if n_components > 1:
+ raise ValueError(
+ f"the selected particles do not touch — merging them would make "
+ f"{n_components} disconnected regions, not one")
+ rows, contours = _remeasure(components.astype(np.int32), frame_image, t,
+ particles.scale)
+ if len(rows) != 1:
+ raise ValueError(f"the merged region measured as {len(rows)} rows, not 1")
+
+ areas = particles.flat_buffer[np.asarray(idx), COL["area"]]
+ survivor = idx[int(np.argmax(areas))]
+ rows[0, COL["track_id"]] = particles.flat_buffer[survivor, COL["track_id"]]
+ rows[0, COL["label"]] = particles.flat_buffer[np.asarray(idx), COL["label"]].min()
+
+ added = _splice(particles, drop=idx, add=[(t, rows[0], contours[0])])
+ return int(added[0])
+
+
+def split_particle(particles, index: int, line, *, frame_image=None) -> tuple[int, int]:
+ """Cut one particle along *line* and re-measure both halves.
+
+ *line* is ``((x0, y0), (x1, y1))`` in IMAGE PIXELS. The cut is the infinite
+ line through those two points; pixels of the particle's mask fall on one side
+ or the other by the sign of the 2-D cross product. An infinite line rather
+ than a segment because a user drawing a cut across a blob naturally starts
+ and ends outside it, and a segment-only rule would leave the ends uncut.
+
+ The larger half keeps the parent's ``track_id``; the smaller is left
+ untracked (-1), because which fragment continues the track is exactly the
+ question a re-link answers and guessing it here would fabricate an identity.
+ """
+ if not particles.has_masks:
+ raise ValueError("cannot split without outlines "
+ "(segmentation ran with store_masks=False)")
+ gi = int(index)
+ (x0, y0), (x1, y1) = ((float(line[0][0]), float(line[0][1])),
+ (float(line[1][0]), float(line[1][1])))
+ if (x1 - x0) == 0.0 and (y1 - y0) == 0.0:
+ raise ValueError("the cut line has zero length")
+
+ mask = _full_mask(particles, gi)
+ h, w = particles.frame_shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ side = (xx - x0) * (y1 - y0) - (yy - y0) * (x1 - x0)
+ labels = np.zeros((h, w), np.int32)
+ labels[mask & (side < 0)] = 1
+ labels[mask & (side >= 0)] = 2
+ if not labels.any() or labels.max() < 2 or not (labels == 1).any():
+ raise ValueError("the cut line does not divide this particle")
+
+ t = int(particles.flat_buffer[gi, COL["t"]])
+ rows, contours = _remeasure(labels, frame_image, t, particles.scale)
+ if len(rows) != 2:
+ raise ValueError(
+ f"the cut produced {len(rows)} regions, not 2 — the line probably "
+ "clipped a corner or the halves are disconnected")
+
+ parent_track = particles.flat_buffer[gi, COL["track_id"]]
+ keeper = int(np.argmax(rows[:, COL["area"]]))
+ rows[keeper, COL["track_id"]] = parent_track
+ rows[1 - keeper, COL["track_id"]] = -1.0
+
+ added = _splice(particles, drop=[gi],
+ add=[(t, rows[k], contours[k]) for k in range(2)])
+ return int(added[0]), int(added[1])
+
+
+def record_edit(tree, particles, kind: str, **fields) -> dict:
+ """Record a manual correction on the tree AND in the store's provenance.
+
+ Two places, deliberately. ``tree.particle_edits`` is the live log a re-run
+ reads so it does not silently discard the user's corrections; the copy in
+ ``particles.provenance["edits"]`` travels with ``SpyDEParticles.save`` and is
+ stamped into the tree's commit provenance, so a corrected result is still
+ reproducible from the file alone.
+ """
+ record = {"kind": str(kind), "at": time.time(), **fields}
+ if tree is not None:
+ record["revision"] = len(getattr(tree, "particle_edits", None) or ()) + 1
+ edits = list(getattr(tree, "particle_edits", None) or ())
+ edits.append(record)
+ tree.particle_edits = edits
+ commit = getattr(tree, "_commit_provenance", None)
+ if isinstance(commit, dict):
+ commit["edits"] = list(edits)
+ if particles is not None:
+ provenance = dict(particles.provenance or {})
+ provenance["edits"] = list(provenance.get("edits") or ()) + [record]
+ particles.provenance = provenance
+ return record
+
+
+def pending_edits(tree) -> list[dict]:
+ """The manual corrections recorded on *tree*, oldest first.
+
+ **This is the seam a re-segmentation must consult.** A re-run rebuilds the
+ store from the raw frames, which discards every edit unless it reads this
+ first — and discarding them silently is the failure the plan calls out (B9).
+ Nothing in ``particles_action`` reads it yet; that is the segmentation
+ workstream's half of the contract, and it is a list of plain JSON-safe dicts
+ precisely so it can cross that boundary (and the IPC) unchanged.
+ """
+ return list(getattr(tree, "particle_edits", None) or ())
+
+
+# ── the navigator lanes (plan C2) ────────────────────────────────────────────
+
+#: Lane names, in stacking order.
+LANE_COUNT = "count"
+LANE_SIZE = "mean size"
+LANE_EVENTS = "events"
+
+#: Row each event kind occupies inside the event lane, top to bottom.
+EVENT_ROWS: dict[str, int] = {"birth": 3, "death": 2, "merge": 1, "split": 0}
+
+
+def step_trace(values, x=None) -> tuple[np.ndarray, np.ndarray]:
+ """Duplicate samples so a polyline draws a ``steps-post`` staircase.
+
+ **The count lane is integer data and must be drawn as a step.** A straight
+ interpolation between frames puts the visual transition half a frame early,
+ so a nucleation at frame 8 reads as 7 (plan C3, found by looking at a
+ render). anyplotlib's ``Axes.plot`` has no ``drawstyle``, so the staircase is
+ built into the DATA: each sample is held until the next x, then jumps.
+ Continuous quantities (mean size) stay plain lines.
+
+ Returns
+ -------
+ (x, y)
+ Both ``(2 * n,)``. The final sample is held one step wide so the last
+ frame is as visible as every other one.
+ """
+ y = np.asarray(values, np.float32).ravel()
+ n = y.size
+ if n == 0:
+ return np.zeros(0, np.float64), np.zeros(0, np.float32)
+ xs = np.arange(n, dtype=np.float64) if x is None else np.asarray(x, np.float64).ravel()
+ width = float(xs[1] - xs[0]) if n > 1 else 1.0
+ edges = np.concatenate([xs, [xs[-1] + width]])
+ out_x = np.repeat(edges, 2)[1:-1]
+ out_y = np.repeat(y, 2)
+ return out_x, out_y
+
+
+def _event_points(events, kind: str, scale: float, offset: float) -> np.ndarray:
+ """``(n, 2)`` ``(x, row)`` marker offsets for one event kind."""
+ frames = [int(e.frame) for e in events if getattr(e, "kind", None) == kind]
+ if not frames:
+ return np.zeros((0, 2), np.float32)
+ row = float(EVENT_ROWS.get(kind, 0))
+ return np.column_stack([np.asarray(frames, np.float64) * scale + offset,
+ np.full(len(frames), row)]).astype(np.float32)
+
+
+def publish_navigator_lanes(session, tree, *, plot=None) -> bool:
+ """Publish ``tree.nav_traces`` as three stacked 1-D navigator lanes.
+
+ ``count(t)``, ``mean size(t)`` and a dedicated event lane with a colour per
+ kind, stacked as rows on one shared time axis with a single logical cursor
+ wired to the tree's REAL 1-D navigation selector — so dragging any lane moves
+ the movie and playback moves every lane's line.
+
+ Reuses ``navigator_views``' stacked-cursor machinery wholesale
+ (:class:`~spyde.actions.navigator_views._StackedNavCursor`, its registry and
+ its teardown), because the hard part of a stacked navigator is the two-way
+ cursor sync and its re-entrancy guard, and that is already written and
+ tested. What is NOT reusable is ``_stack_navigators``' figure builder: it
+ draws every lane as a plain line, and the count lane must be a step and the
+ event lane is markers rather than a trace at all.
+
+ Each lane also lands in ``tree.navigator_signals`` so the chip strip lists it
+ and the existing ``select_navigator`` path can re-stack any subset.
+
+ Returns
+ -------
+ bool
+ True if the figure was emitted.
+ """
+ import anyplotlib as apl
+ import anyplotlib._electron as _electron
+ import hyperspy.api as hs
+
+ from spyde.actions.figure_registry import keep_alive
+ from spyde.actions.navigator_views import (
+ STACKED_LABEL, _real_nav_selector, _selector_axis, _stacked_cursors,
+ _StackedNavCursor, _teardown_stacked,
+ )
+ from spyde.backend.ipc import emit
+ from spyde.drawing.plots.plot import finalize_figure_html
+
+ traces = dict(getattr(tree, "nav_traces", None) or {})
+ if LANE_COUNT not in traces:
+ log.debug("[particles] no nav_traces on the tree; nothing to publish")
+ return False
+ if plot is None:
+ plot = _first_nav_plot(tree)
+ window_id = getattr(plot, "window_id", None)
+ if window_id is None:
+ log.debug("[particles] no navigator window to publish lanes onto")
+ return False
+
+ count = np.asarray(traces[LANE_COUNT], np.float32)
+ size = np.asarray(traces.get("size", np.full(count.shape, np.nan)), np.float32)
+ events = list(getattr(tree, "particle_events", None) or ())
+
+ # Register the traces as named navigators too, so the chip strip offers them.
+ for name, data in ((LANE_COUNT, count), (LANE_SIZE, size)):
+ if name in getattr(tree, "navigator_signals", {}):
+ continue
+ try:
+ tree.add_navigator_signal(name, hs.signals.Signal1D(np.nan_to_num(data)))
+ except Exception as exc:
+ log.debug("[particles] registering navigator %r failed: %s", name, exc)
+
+ selector = _real_nav_selector(session, int(window_id))
+ scale, offset = _selector_axis(selector) if selector is not None else (1.0, 0.0)
+ current = 0
+ if selector is not None and getattr(selector, "current_indices", None) is not None:
+ try:
+ current = int(np.asarray(selector.current_indices).ravel()[0])
+ except Exception:
+ current = 0
+ cursor_x = current * scale + offset
+
+ _teardown_stacked(session, plot)
+ try:
+ fig, axes = apl.subplots(3, 1, sharex=True)
+ panels = np.array(axes, dtype=object).ravel()
+ widgets = []
+
+ step_x, step_y = step_trace(count, np.arange(count.size) * scale + offset)
+ lanes = [
+ (panels[0], LANE_COUNT, step_y, step_x),
+ (panels[1], LANE_SIZE, np.nan_to_num(size),
+ np.arange(size.size) * scale + offset),
+ ]
+ # NB every setter below goes on the PLOT the panel returns, never on the
+ # Axes: anyplotlib's `Axes` has no `set_title`/`set_ylim` at all, so the
+ # guarded calls silently did nothing and the event lane came out
+ # autoscaled to the invisible baseline (verified in the app — the event
+ # rows sat off-scale and nothing drew).
+ for panel, title, ydata, xdata in lanes:
+ line = panel.plot(ydata, axes=[xdata], label=title)
+ _set_title(line, title)
+ widgets.append(_add_cursor(line, cursor_x))
+
+ # A transparent baseline establishes the x axis; the events themselves
+ # are markers, one group (and one row) per kind so the colours mean
+ # something.
+ base = panels[2].plot(np.zeros(count.size, np.float32),
+ axes=[np.arange(count.size) * scale + offset],
+ label=LANE_EVENTS, alpha=0.0)
+ _set_title(base, LANE_EVENTS)
+ for kind, color in EVENT_COLORS.items():
+ base.add_points(_event_points(events, kind, scale, offset),
+ name=f"event_{kind}", sizes=5.0, color=color,
+ facecolors=color, alpha=1.0, label=kind)
+ try:
+ # Fixed rows, NOT autoscaled: the lane must read the same whether the
+ # movie contains one kind of event or all four, so a birth is always
+ # on the top row.
+ base.set_ylim(-0.5, max(EVENT_ROWS.values()) + 0.5)
+ except Exception as exc:
+ log.debug("[particles] setting event-lane ylim failed: %s", exc)
+ widgets.append(_add_cursor(base, cursor_x))
+
+ widgets = [w for w in widgets if w is not None]
+ if len(widgets) >= 2 and selector is not None:
+ _stacked_cursors(session)[int(window_id)] = _StackedNavCursor(
+ session, int(window_id), widgets, selector)
+
+ fig_id = _electron.register(fig)
+ html = finalize_figure_html(fig, fig_id)
+ keep_alive(int(window_id), fig)
+ emit({"type": "figure", "fig_id": fig_id, "window_id": window_id,
+ "html": html, "title": " / ".join((LANE_COUNT, LANE_SIZE, LANE_EVENTS)),
+ "is_navigator": True,
+ "view_label": STACKED_LABEL, "view_kind": "stacked"})
+ return True
+ except Exception as exc:
+ log.exception("[particles] publishing navigator lanes failed: %s", exc)
+ return False
+
+
+def maybe_stack_particle_lanes(session, plot, tree, names) -> bool:
+ """Lane hook for ``navigator_views.select_navigator``.
+
+ A particle tree's lanes do not render as plain lines: ``count`` is integer
+ data and must be a STEP, and the event lane is coloured markers rather than a
+ trace at all. ``_stack_navigators`` draws every named navigator the same way,
+ so when the user ⇧-clicks the chips on a particle tree the generic builder
+ would silently produce the wrong picture — a nucleation at frame 8 reading as
+ 7 is exactly the failure plan C3 calls out.
+
+ Returns True when it built the lanes (the caller must then return).
+ """
+ if not getattr(tree, "nav_traces", None) or LANE_COUNT not in tree.nav_traces:
+ return False
+ wanted = {str(n) for n in (names or ())}
+ if not wanted & {LANE_COUNT, LANE_SIZE}:
+ return False # the user picked other navigators; not our business
+ return publish_navigator_lanes(session, tree, plot=plot)
+
+
+def _add_cursor(panel, x: float):
+ try:
+ return panel.add_vline_widget(x=float(x), color="#ff9100")
+ except Exception as exc:
+ log.debug("[particles] adding lane cursor failed: %s", exc)
+ return None
+
+
+def _set_title(plot1d, title: str) -> None:
+ """Title a lane. Takes the PLOT, not the Axes — see the note in the builder."""
+ try:
+ plot1d.set_title(title)
+ except Exception as exc:
+ log.debug("[particles] set_title on lane failed: %s", exc)
+
+
+def _first_nav_plot(tree):
+ manager = getattr(tree, "navigator_plot_manager", None)
+ if manager is None:
+ return None
+ for window in list(manager.plot_windows.keys()):
+ plots = manager.plots.get(window) or []
+ if plots:
+ return plots[0]
+ return None
+
+
+def _first_signal_plot(tree):
+ for plot in list(getattr(tree, "signal_plots", None) or ()):
+ if getattr(plot, "plot_state", None) is not None:
+ return plot
+ return None
+
+
+# ── staged handlers (registry.STAGED_HANDLERS, key "part") ───────────────────
+
+def _resolve(session, plot):
+ """``(tree, overlay)`` for a handler. Prefers the clicked plot's own tree,
+ falling back to any tree that carries particles — the caret's window may
+ resolve to the count-map navigator rather than the label movie, the same way
+ ``lifecycle.resolve_vectors`` handles the vectors case."""
+ tree = getattr(plot, "signal_tree", None) if plot is not None else None
+ if tree is not None and getattr(tree, "particles", None) is not None:
+ return tree, getattr(tree, "_particle_overlay", None)
+ for candidate in getattr(session, "signal_trees", None) or ():
+ if getattr(candidate, "particles", None) is not None:
+ return candidate, getattr(candidate, "_particle_overlay", None)
+ return tree, getattr(tree, "_particle_overlay", None) if tree is not None else None
+
+
+def _source_frame_provider(tree):
+ """``fn(t) -> frame`` reading ONE frame of the tree's source movie.
+
+ Used to re-measure intensity after an edit. Reads a single frame and computes
+ only that slice — never the movie (CLAUDE.md memory-safety rule). Returns
+ None when there is no source, in which case edited rows keep NaN intensities
+ rather than fabricated ones.
+ """
+ source = getattr(tree, "source_node", None)
+ data = getattr(source, "data", None)
+ if data is None:
+ return None
+
+ def read(t: int):
+ frame = data[int(t)]
+ if hasattr(frame, "compute"):
+ frame = frame.compute()
+ return np.asarray(frame)
+
+ return read
+
+
+def part_open(session, plot, payload=None) -> None:
+ """Caret mounted → attach the overlay to the tree's signal plot."""
+ from spyde.actions.lifecycle import wait_for_particles
+ from spyde.backend.ipc import emit_error, emit_status
+
+ payload = payload or {}
+ tree = getattr(plot, "signal_tree", None) if plot is not None else None
+ if tree is None:
+ emit_error("Particle Overlay: no active dataset")
+ return
+ if getattr(tree, "particles", None) is None:
+ # The segmentation attach gap (plan Wave 0): seg_run opens its window
+ # early and attaches tree.particles only at finalize.
+ if wait_for_particles(session, plot,
+ lambda: part_open(session, plot, payload),
+ what="Particle Overlay"):
+ return
+ emit_error("Particle Overlay needs a segmentation result (no particles).")
+ return
+
+ target = _first_signal_plot(tree) or plot
+ overlay = attach_particle_overlay(
+ target, tree.particles, tree,
+ events=list(getattr(tree, "particle_events", None) or ()),
+ nav_map=getattr(tree, "nav_map", None),
+ frame_provider=_source_frame_provider(tree),
+ trail_frames=int(payload.get("trail_frames", DEFAULT_TRAIL_FRAMES)),
+ show_trails=bool(payload.get("show_trails", False)),
+ on_select=lambda: _emit_selection(tree),
+ )
+ if overlay is None:
+ emit_error("Particle Overlay could not attach to this window")
+ return
+ emit_status(f"Particle overlay on — {tree.particles.n_particles} particles")
+ _emit_selection(tree)
+
+
+def part_close(session, plot, payload=None) -> None:
+ """Caret unmounted → tear the overlay down."""
+ from spyde.actions.lifecycle import replace_tree_attr
+ tree, _overlay = _resolve(session, plot)
+ if tree is not None:
+ replace_tree_attr(tree, "_particle_overlay", None)
+
+
+def part_tune(session, plot, payload=None) -> None:
+ """Live parameter change: trails on/off and the trailing window length."""
+ payload = payload or {}
+ _tree, overlay = _resolve(session, plot)
+ if overlay is None:
+ return
+ overlay.set_trails(bool(payload.get("show_trails", overlay.show_trails)),
+ payload.get("trail_frames"))
+
+
+def part_select(session, plot, payload=None) -> None:
+ """External selection hook — a table row, a track id, or a region.
+
+ Payload accepts ``{"indices": [...]}`` (global particle indices),
+ ``{"track_id": n}``, or ``{"region": [x0, y0, x1, y1]}`` in image pixels.
+ """
+ payload = payload or {}
+ tree, overlay = _resolve(session, plot)
+ if overlay is None:
+ return
+ if "region" in payload:
+ overlay.select_region(*[float(v) for v in payload["region"]])
+ elif "track_id" in payload:
+ overlay.select_track(int(payload["track_id"]))
+ else:
+ overlay.select(payload.get("indices") or [])
+ _emit_selection(tree)
+
+
+def part_region_mode(session, plot, payload=None) -> None:
+ """Show/hide the rubber-band selection rectangle."""
+ payload = payload or {}
+ _tree, overlay = _resolve(session, plot)
+ if overlay is not None:
+ overlay.set_region_select(bool(payload.get("enabled", True)))
+
+
+def part_delete(session, plot, payload=None) -> None:
+ """Delete the selected particles (or ``payload["indices"]``)."""
+ from spyde.backend.ipc import emit_error, emit_status
+ payload = payload or {}
+ tree, overlay = _resolve(session, plot)
+ if overlay is None:
+ emit_error("Particle Overlay: nothing to edit")
+ return
+ try:
+ removed = overlay.delete(payload.get("indices"))
+ except Exception as exc:
+ emit_error(f"Delete particles failed: {exc}")
+ return
+ emit_status(f"Deleted {removed} particle{'' if removed == 1 else 's'}")
+ _emit_selection(tree)
+
+
+def part_merge(session, plot, payload=None) -> None:
+ """Merge the selected particles into one and re-measure."""
+ from spyde.backend.ipc import emit_error, emit_status
+ payload = payload or {}
+ tree, overlay = _resolve(session, plot)
+ if overlay is None:
+ emit_error("Particle Overlay: nothing to edit")
+ return
+ try:
+ index = overlay.merge(payload.get("indices"))
+ except Exception as exc:
+ emit_error(f"Merge particles failed: {exc}")
+ return
+ emit_status(f"Merged into particle {index}")
+ _emit_selection(tree)
+
+
+def part_split(session, plot, payload=None) -> None:
+ """Split the selected particle along ``payload["line"]`` (image pixels)."""
+ from spyde.backend.ipc import emit_error, emit_status
+ payload = payload or {}
+ tree, overlay = _resolve(session, plot)
+ if overlay is None:
+ emit_error("Particle Overlay: nothing to edit")
+ return
+ try:
+ a, b = overlay.split(payload.get("index"), payload.get("line"))
+ except Exception as exc:
+ emit_error(f"Split particle failed: {exc}")
+ return
+ emit_status(f"Split into particles {a} and {b}")
+ _emit_selection(tree)
+
+
+def part_lanes(session, plot, payload=None) -> None:
+ """Build (or rebuild) the three stacked navigator lanes."""
+ from spyde.backend.ipc import emit_error
+ tree, _overlay = _resolve(session, plot)
+ if tree is None or not getattr(tree, "nav_traces", None):
+ emit_error("Particle lanes: this dataset has no navigator traces")
+ return
+ if not publish_navigator_lanes(session, tree):
+ emit_error("Particle lanes: could not build the stacked navigator")
+
+
+def _json_float(value) -> float | None:
+ """A property value as JSON, with NaN mapped to ``null``.
+
+ ``ipc.emit`` calls ``json.dumps`` with the default ``allow_nan=True``, which
+ writes the bare token ``NaN`` — not valid JSON, so ``JSON.parse`` on the
+ Electron side throws and the whole message is dropped. And NaN is the NORMAL
+ value here: ``measure_frame`` leaves every intensity column NaN when it runs
+ without an intensity image, which is exactly what an edit made with no
+ ``frame_provider`` produces.
+ """
+ v = float(value)
+ return v if np.isfinite(v) else None
+
+
+def _emit_selection(tree) -> None:
+ """Tell the renderer which particles are selected, with their properties.
+
+ One message rather than a per-row query: the table dock highlights the
+ selection and shows the readout, and both come from the same rows.
+ """
+ from spyde.backend.ipc import emit
+ overlay = getattr(tree, "_particle_overlay", None) if tree is not None else None
+ if overlay is None:
+ return
+ rows = []
+ for gi in overlay.selected:
+ row = overlay.particles.flat_buffer[int(gi)]
+ record = {"index": int(gi), "frame": int(row[COL["t"]]),
+ "track_id": int(row[COL["track_id"]]),
+ "color": track_color(row[COL["track_id"]])}
+ record.update({name: _json_float(row[COL[name]]) for name in MEASURED_COLUMNS})
+ rows.append(record)
+ emit({"type": "particle_selection",
+ "window_id": getattr(overlay.plot, "window_id", None),
+ "frame": int(overlay._frame), "revision": int(overlay.revision),
+ "indices": [int(i) for i in overlay.selected], "particles": rows})
+
+
+# ── toolbar entry ────────────────────────────────────────────────────────────
+
+class _OverlayHandle:
+ """What the toolbar tracks so DESELECTING the action tears the overlay down.
+
+ ``Session._track_action_artifacts`` records any object a toolbar action
+ returns that exposes ``active_children``, lights the button, and calls
+ ``.close()`` on it when the user clicks the lit button again.
+
+ That is the ONLY un-toggle path the renderer offers: ``FloatingToolbar``'s
+ click handler sends ``set_action_active(active=false)`` for an action it
+ believes is live — never a second ``toolbar_action``. So an action that
+ toggles by inspecting its own state can be turned ON but never OFF. Verified
+ in the app: the second click re-ran ``part_open`` and the overlay stayed up.
+ """
+
+ active_children: tuple = ()
+
+ def __init__(self, session, plot):
+ self.session = session
+ self.plot = plot
+
+ def close(self) -> None:
+ part_close(self.session, self.plot, {})
+
+
+def particle_overlay(ctx, action_name: str = "Particle Overlay", **params):
+ """Toolbar toggle: attach the overlay and hand the toolbar its teardown handle.
+
+ Self-contained rather than a no-op wizard parent, so the overlay works from
+ the toolbar alone; the ``part_*`` staged handlers are the same operations for
+ a caret to drive once one exists.
+ """
+ from spyde.backend import ipc
+
+ plot = getattr(ctx, "plot", None)
+ session = getattr(plot, "session", None)
+ tree = getattr(plot, "signal_tree", None) if plot is not None else None
+ if tree is None:
+ ipc.emit_error("Particle Overlay: no active dataset")
+ return None
+ if getattr(tree, "_particle_overlay", None) is not None:
+ # Already up (a stale artifact, or a caret opened it): close so the click
+ # is still a toggle even when the artifact path is not in play.
+ part_close(session, plot, {})
+ window_id = getattr(plot, "window_id", None)
+ if window_id is not None:
+ ipc.emit({"type": "action_active", "window_id": window_id,
+ "name": action_name, "active": False})
+ return None
+ part_open(session, plot, params)
+ if getattr(tree, "_particle_overlay", None) is None:
+ return None # open failed; it emitted its own error
+ return _OverlayHandle(session, plot)
+
+
+def particle_lanes(ctx, action_name: str = "Particle Lanes", **params):
+ """Toolbar entry: (re)build the three stacked navigator lanes."""
+ plot = getattr(ctx, "plot", None)
+ session = getattr(plot, "session", None)
+ part_lanes(session, plot, params)
+ return None
diff --git a/spyde/actions/particle_tree.py b/spyde/actions/particle_tree.py
new file mode 100644
index 00000000..21e25c75
--- /dev/null
+++ b/spyde/actions/particle_tree.py
@@ -0,0 +1,256 @@
+"""
+particle_tree.py — segmentation spawns a NEW SignalTree. Plan §0.6.
+
+A segmentation is not a property of the movie it came from; it is a derived
+dataset computed *from* it, exactly like a strain map or an orientation map. So
+this module builds a **particle tree** rather than hanging a ``particles``
+attribute off the source.
+
+ ParticleTree
+ root signal : lazy LABEL MOVIE — same nav/signal shape as the source,
+ each frame painted from stored contours on demand
+ tree.particles : SpyDEParticles (the CSR store)
+ tree.source_node : the signal it was computed from
+ tree.nav_map : source nav indices → particle frame index
+ navigator : count(t) / mean size(t) / event lanes
+
+Three things this buys, none of which the attribute form does:
+
+* **It answers Wave D by construction.** Particles found on a 4D-STEM virtual
+ image record the node they came from and which parent nav positions each one
+ covers, so "the mean diffraction pattern for this particle" is a slice of the
+ source's parent rather than a guess about which grid the coordinates belong to.
+* **The label movie is a dataset**, so scrubbing, saving, the report builder and
+ the movie editor all work on it with no special case.
+* **Re-segmenting does not destroy the previous result** — two parameter choices
+ are two sibling trees you can compare, which is what the signal tree is for.
+
+The label movie is **lazy and never materialised**. A 4096² int32 label image is
+64 MB *per frame*; the contours are the truth and ``render_frame`` paints one
+frame when something asks for it. See ``spyde/signals/particles.py``.
+"""
+from __future__ import annotations
+
+import logging
+from typing import Any
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+#: ``_signal_type`` carried by a particle tree's root. Toolbar entries gate on
+#: this to offer particle actions (track, export, per-particle DP) on the result
+#: and nowhere else.
+PARTICLE_SIGNAL_TYPE = "particles"
+
+
+def _label_movie(particles, *, chunk_frames: int = 1):
+ """A lazy ``(n_frames, h, w)`` int32 label movie backed by the contours.
+
+ One frame per chunk, deliberately: each nav move then reads exactly the frame
+ it needs, which is both the cheapest possible read and the same access
+ granularity a real in-situ movie has (CLAUDE.md Live-Display §1). Painting is
+ done inside the dask graph, so no frame exists until something asks for it.
+ """
+ import dask
+ import dask.array as da
+
+ h, w = particles.frame_shape
+ n = particles.n_frames
+
+ def _one(block_info=None):
+ # block_info tells us which frame this block is; there is exactly one
+ # frame per block, so the slice start IS the frame index.
+ t = 0 if block_info is None else int(block_info[None]["array-location"][0][0])
+ try:
+ return particles.render_frame(t, value="track")[None, ...]
+ except Exception as exc: # pragma: no cover
+ log.debug("[particles] frame %d render failed: %s", t, exc)
+ return np.zeros((1, h, w), np.int32)
+
+ with dask.config.set(scheduler="threads"):
+ return da.map_blocks(
+ _one, dtype=np.int32,
+ chunks=((chunk_frames,) * n, (h,), (w,)),
+ meta=np.zeros((0, 0, 0), np.int32),
+ )
+
+
+def _navigator_traces(particles, events=None) -> dict[str, np.ndarray]:
+ """The three navigator lanes, as plain arrays.
+
+ ``count`` is integer data and the renderer must draw it as a STEP — a straight
+ interpolation between frames puts a nucleation's visual transition half a frame
+ early, so an event at frame 8 reads as 7 (plan C3).
+ """
+ traces: dict[str, np.ndarray] = {
+ "count": particles.count_series(),
+ "size": particles.property_series("area", "mean"),
+ }
+ if events is not None:
+ from spyde.particles.track import event_counts
+ ec = event_counts(events, particles.n_frames)
+ for kind, arr in ec.items():
+ traces[f"event_{kind}"] = np.asarray(arr, np.float32)
+ return traces
+
+
+def open_particle_tree(session, *, particles, source_node, source_tree=None,
+ title: str | None = None, events=None,
+ nav_map=None, params: dict[str, Any] | None = None,
+ provenance: dict[str, Any] | None = None,
+ attach: bool = True):
+ """Create the particle tree for a finished segmentation.
+
+ Parameters
+ ----------
+ particles
+ The :class:`~spyde.signals.particles.SpyDEParticles` store.
+ source_node
+ The signal segmentation ran on. Recorded on the tree so Wave D can walk
+ back to its parent; **not** used to hold the result.
+ source_tree
+ The tree *source_node* belongs to, when the caller knows it. Recorded for
+ provenance; the particle tree is a sibling, not a child.
+ events
+ Optional event list from :func:`spyde.particles.track.link`. When given,
+ the navigator gains an event lane.
+ nav_map
+ Optional ``(n_frames,)`` int array mapping each particle frame back to a
+ source navigation index. ``None`` means identity, which is the case for a
+ movie; a 4D-STEM virtual image passes the parent's grid.
+ attach
+ ``True`` (the default) publishes the store as ``tree.particles``
+ immediately — right for a finished result.
+
+ A **progressive** run passes ``attach=False``. The plan is explicit that
+ ``tree.particles`` appears only when the batch finalizes, because
+ ``requires_particles`` gates on it: attaching an empty placeholder would
+ unlock the whole particle toolbar against a store holding nothing, and the
+ user would click Track on zero particles.
+
+ The store still has to be handed in, though, because the lazy label movie
+ **closes over this exact object** — the frames it renders come from it. So
+ the placeholder is retained (as ``tree._seg_pending_particles``) and must
+ later be MUTATED IN PLACE by the caller at finalize. Swapping in a freshly
+ built ``SpyDEParticles`` would leave the already-open window rendering the
+ placeholder's zeros forever.
+
+ Returns
+ -------
+ The new tree, with ``source_node``, ``nav_map`` and ``nav_traces`` attached,
+ and ``particles`` set unless *attach* is False.
+ """
+ import hyperspy.api as hs
+
+ from spyde.actions.commit import open_result_tree
+
+ n = particles.n_frames
+ h, w = particles.frame_shape
+ title = title or f"Particles — {particles.n_particles} in {n} frames"
+
+ lazy = _label_movie(particles)
+ sig = hs.signals.Signal2D(lazy).as_lazy()
+ sig.data = lazy
+
+ # Carry the source's calibration so a particle's centroid means the same
+ # thing on this tree as it did on the movie. Without it the label movie is
+ # in pixels while every measured property is in nm, which is the exact class
+ # of unit mismatch the linker's docstring warns about.
+ _copy_axes(source_node, sig, particles)
+
+ tree = open_result_tree(
+ session, title=title, signal=sig,
+ signal_type=PARTICLE_SIGNAL_TYPE,
+ provenance=provenance or {"action": "segment_particles",
+ "params": dict(params or {})},
+ )
+
+ if attach:
+ tree.particles = particles
+ else:
+ # Gated off until the batch finalizes (see `attach`). The store is still
+ # kept, because the lazy label movie renders from THIS object.
+ tree.particles = None
+ tree._seg_pending_particles = particles
+ tree.source_node = source_node
+ tree.source_tree = source_tree
+ tree.nav_map = (np.arange(n, dtype=np.int64) if nav_map is None
+ else np.asarray(nav_map, np.int64))
+ tree.particle_events = list(events or ())
+ tree.nav_traces = _navigator_traces(particles, events)
+ return tree
+
+
+def _copy_axes(source_node, sig, particles) -> None:
+ """Give the label movie the source's calibration, best-effort.
+
+ Best-effort on purpose: a wrong calibration is worse than none, so every step
+ is guarded and a failure leaves the axis at its default rather than half-
+ applied. The signal axes fall back to the particle store's own ``scale``,
+ which is what the measurements were computed with.
+ """
+ try:
+ src_sig = getattr(source_node, "axes_manager", None)
+ if src_sig is not None:
+ for dst, src in zip(sig.axes_manager.signal_axes,
+ source_node.axes_manager.signal_axes):
+ dst.scale, dst.units = float(src.scale), src.units
+ nav = source_node.axes_manager.navigation_axes
+ if nav:
+ dnav = sig.axes_manager.navigation_axes[0]
+ dnav.name, dnav.units = nav[0].name, nav[0].units
+ dnav.scale, dnav.offset = float(nav[0].scale), float(nav[0].offset)
+ return
+ except Exception as exc:
+ log.debug("[particles] copying source axes failed: %s", exc)
+ try:
+ for ax in sig.axes_manager.signal_axes:
+ ax.scale, ax.units = float(particles.scale), particles.units
+ except Exception as exc: # pragma: no cover
+ log.debug("[particles] fallback axis calibration failed: %s", exc)
+
+
+def particle_nav_positions(tree, index: int):
+ """Source navigation indices covered by global particle *index*.
+
+ This is the Wave D seam: a particle found on a derived node (a 4D-STEM virtual
+ image) needs to name positions in the PARENT's navigation grid before anything
+ can average diffraction patterns over it.
+
+ Returns an ``(m, k)`` int array of navigation indices, where *k* is the source
+ navigation dimensionality. For a movie (1-D nav) that is the single frame the
+ particle lives in; for a virtual image it is every scan position under the
+ particle's mask.
+ """
+ particles = tree.particles
+ row = particles.flat_buffer[int(index)]
+ from spyde.signals.particles import COL
+
+ t = int(row[COL["t"]])
+ frame = int(tree.nav_map[t]) if t < len(tree.nav_map) else t
+
+ if not particles.has_masks:
+ # No outlines stored: the best we can name is the frame itself.
+ return np.asarray([[frame]], np.int64)
+
+ mask, (y0, x0, _y1, _x1) = particles.mask_at(int(index))
+ ys, xs = np.nonzero(mask)
+ if ys.size == 0:
+ return np.asarray([[frame]], np.int64)
+
+ nav_dim = _source_nav_dim(tree)
+ if nav_dim <= 1:
+ # A movie: the particle's pixels are SIGNAL coordinates, not navigation
+ # ones, so the only navigation index involved is the frame.
+ return np.asarray([[frame]], np.int64)
+ # A virtual image: the particle's pixels ARE navigation positions.
+ return np.stack([ys + y0, xs + x0], axis=-1).astype(np.int64)
+
+
+def _source_nav_dim(tree) -> int:
+ src = getattr(tree, "source_node", None)
+ try:
+ return int(src.axes_manager.navigation_dimension)
+ except Exception:
+ return 1
diff --git a/spyde/actions/particles_action.py b/spyde/actions/particles_action.py
new file mode 100644
index 00000000..a02d29ad
--- /dev/null
+++ b/spyde/actions/particles_action.py
@@ -0,0 +1,2328 @@
+"""
+particles_action.py — the Segment Particles wizard (``seg_`` staged actions).
+
+Plan B7, honouring the §0.8 interaction contract literally:
+
+ seg_open caret mounted → preview the CURRENT frame, emit caret state
+ seg_close caret unmounted → clear the overlay, drop the controller
+ seg_set_method classical | scribble | prompt
+ seg_tune debounced param change → re-preview the CURRENT frame only
+ seg_paint one brush stroke → LabelStore
+ seg_train fit the scribble classifier on the accumulated labels
+ seg_run whole movie on a worker: progressive, cancellable
+ seg_commit snapshot the previewed frame as a one-frame particle tree
+
+The three things that shape this module, none of them cosmetic:
+
+**1. Preview is one frame, always.** ``seg_tune`` reads exactly the frame the
+navigator is sitting on. A movie is thousands of frames at the plan's target
+scale, so a tune that touched more than one would put the interaction budget
+(plan B3: train + apply under ~1 s) out of reach on the first drag of a slider.
+
+**2. The run opens its result tree EARLY and attaches nothing until it
+finalizes.** ``open_particle_tree`` sets ``tree.particles`` at construction —
+which is right for its own contract (a finished segmentation) and wrong for a
+progressive run, because ``requires_particles`` would unlock downstream actions
+against a store holding zero particles. So the run hands it a placeholder,
+immediately clears ``tree.particles`` back to ``None``, and sets it at
+``_finalize``. That is the attach gap ``lifecycle.wait_for_particles`` and
+``lifecycle.seg_batch_running`` exist to cover, and it only means anything if
+the flag and the attribute move at the right moments.
+
+The placeholder is then MUTATED IN PLACE rather than replaced, because the lazy
+label movie ``open_particle_tree`` built closes over that exact object — a fresh
+``SpyDEParticles`` would leave the movie rendering the placeholder's zeros
+forever.
+
+**3. min_size is floored, and the floor is reported.** Plan §0.9 measured it: at
+``min_size=0`` a classifier taught faint contrast produced 33 instances where 9
+were real, and ``min_size=10`` removed 24 of the 25 spurious ones. A user who
+zeroes it to "catch the small ones" gets the opposite. So it is floored — and
+the EFFECTIVE value goes back to the caret in every preview, because silently
+running a number different from the one on screen is the failure mode
+``classical.SegmentParams`` refuses for ``local_size`` and this must not
+reintroduce it.
+"""
+from __future__ import annotations
+
+import logging
+import time
+from typing import Any, Callable
+
+import numpy as np
+
+from spyde.actions.context import current_signal as _current_signal
+from spyde.actions.context import src_plot_tree as _src_plot_tree
+from spyde.actions.lifecycle import (
+ bump_generation, is_current, run_on_worker, window_computing,
+)
+from spyde.actions.particle_overlay import _navigator_selectors_for, _push_groups
+from spyde.actions.wizard import WizardController
+from spyde.backend.ipc import emit, emit_error, emit_progress, emit_status
+
+log = logging.getLogger(__name__)
+
+#: The three mask sources of plan §0.2. ``prompt`` is declared here so the caret
+#: can render its tab from the schema before the engine lands (plan B4); every
+#: code path that would run it emits a "not installed yet" status instead.
+METHODS: tuple[str, ...] = ("classical", "scribble", "prompt")
+
+#: Floor applied to ``min_size``. NOT a taste default — plan §0.9's measurement:
+#: on the fixture, one faint scribble added to bright-only labels gave 33
+#: instances at min_size=0 (25 of them spurious) and 9 at min_size=10. The
+#: classifier is not what buys specificity; this filter is.
+MIN_SIZE_FLOOR = 10
+
+#: Cap on the per-frame area list shipped to the caret's size histogram. A
+#: histogram of a few hundred bodies is already at its useful resolution and the
+#: PLOTAPP line protocol is shared with the nav painter thread — a 20k-element
+#: list per slider tick is exactly the traffic plan B0 rejects for brush strokes.
+_MAX_AREAS_SENT = 2000
+
+#: Minimum wall-clock gap between progressive count-trace paints during a run.
+#: Matches ``live_fill_poller``'s default: the fill is a reassurance signal, not
+#: an animation, and every paint is a marshal onto the asyncio main thread.
+_PROGRESS_INTERVAL = 0.35
+
+DEFAULTS: dict[str, Any] = dict(
+ method="classical",
+ # The one sensitivity axis of plan §0.9. Everything else about the split is
+ # secondary and sits below it in the caret.
+ sensitivity=0.5,
+ threshold="otsu",
+ # THE control for "too many particles". Filters the measured instances by
+ # their confidence score (spyde.particles.measure.particle_scores), which is
+ # computed WITH the measurement — so this re-filters an existing result and
+ # never re-segments, and a drag is instant. It also means the same thing on
+ # every engine, because it acts on the OUTPUT rather than on any one
+ # method's parameters. 0 keeps everything (the old behaviour).
+ min_score=0.0,
+ # THE two face controls, both in NANOMETRES because a distance in the image
+ # is something the eye can judge and a 0-1 confidence is not.
+ # merge_nm — pieces closer than this are ONE particle
+ # min_nm — smallest particle worth keeping (diameter)
+ # Both are converted to pixels with the signal's own scale at dispatch, so
+ # they mean the same thing at any magnification. 0 disables either.
+ merge_nm=0.0,
+ min_nm=0.0,
+ min_size=20,
+ max_size=0,
+ watershed=True,
+ min_separation=3,
+ marker_smooth=1.0,
+ gaussian=0.0,
+ rb_kernel=0,
+ invert=False,
+ local_size=31,
+ clear_border=False,
+ # Outlines are what make the overlay and the label movie possible, so they
+ # default ON; plan §0.5 turns them off for very long movies.
+ store_masks=True,
+ track=True,
+ max_dist=10.0,
+ brush=3.0,
+ # The paint state the ClassStrip owns. These MUST be real parameters: the
+ # brush widget lives in Python, so the strip's choices only reach the paint
+ # by travelling through here. They were read but never declared or set, which
+ # pinned every stroke to class 0 and made the eraser a no-op.
+ active_class=0,
+ erase=False,
+)
+
+
+class SegmentWizard(WizardController):
+ """Owns the segmentation caret's state: engine, parameters, scribbles,
+ trained head, the last preview, and the run's result tree."""
+
+ key = "seg"
+
+ # Declared parameter schema — the single source of truth every host renders
+ # from (registry.wizard_parameters("seg")). `tab` mirrors plan B7's three
+ # engine tabs; `method` is the tab bar itself and so carries no tab.
+ parameters = {
+ "method": {
+ "name": "Engine", "type": "enum", "default": DEFAULTS["method"],
+ "choices": list(METHODS),
+ },
+ "sensitivity": {
+ "name": "Sensitivity", "type": "float", "default": DEFAULTS["sensitivity"],
+ "min": 0.0, "max": 1.0, "step": 0.01, "tab": "Classical",
+ },
+ "threshold": {
+ "name": "Threshold", "type": "enum", "default": DEFAULTS["threshold"],
+ "choices": ["otsu", "mean", "minimum", "yen", "isodata", "li",
+ "local", "local_otsu", "niblack", "sauvola"],
+ "tab": "Classical",
+ },
+ "gaussian": {
+ "name": "Pre-blur σ (px)", "type": "float", "default": DEFAULTS["gaussian"],
+ "min": 0.0, "max": 10.0, "step": 0.1, "tab": "Classical",
+ },
+ "rb_kernel": {
+ "name": "Rolling ball (px)", "type": "int", "default": DEFAULTS["rb_kernel"],
+ "min": 0, "max": 256, "tab": "Classical",
+ },
+ "invert": {
+ "name": "Dark particles", "type": "bool", "default": DEFAULTS["invert"],
+ "tab": "Classical",
+ },
+ "local_size": {
+ "name": "Local window (px, odd)", "type": "int",
+ "default": DEFAULTS["local_size"], "min": 3, "max": 255,
+ "tab": "Classical",
+ },
+ # min_size sits next to sensitivity deliberately (plan §0.9: the two are
+ # coupled and must not be in separate tabs).
+ "min_size": {
+ "name": "Min size (px)", "type": "int", "default": DEFAULTS["min_size"],
+ "min": 0, "max": 100000, "tab": "Split",
+ },
+ "max_size": {
+ "name": "Max size (px, 0=off)", "type": "int",
+ "default": DEFAULTS["max_size"], "min": 0, "max": 10000000,
+ "tab": "Split",
+ },
+ "watershed": {
+ "name": "Split touching", "type": "bool", "default": DEFAULTS["watershed"],
+ "tab": "Split",
+ },
+ "min_separation": {
+ "name": "Min separation (px)", "type": "int",
+ "default": DEFAULTS["min_separation"], "min": 1, "max": 100,
+ "tab": "Split",
+ },
+ "marker_smooth": {
+ "name": "Marker smoothing", "type": "float",
+ "default": DEFAULTS["marker_smooth"], "min": 0.0, "max": 10.0,
+ "step": 0.1, "tab": "Split",
+ },
+ "clear_border": {
+ "name": "Drop edge particles", "type": "bool",
+ "default": DEFAULTS["clear_border"], "tab": "Split",
+ },
+ "brush": {
+ "name": "Brush (px)", "type": "float", "default": DEFAULTS["brush"],
+ "min": 1.0, "max": 64.0, "step": 1.0, "tab": "Scribble",
+ },
+ "store_masks": {
+ "name": "Store outlines", "type": "bool",
+ "default": DEFAULTS["store_masks"], "tab": "Run",
+ },
+ "track": {
+ "name": "Link tracks", "type": "bool", "default": DEFAULTS["track"],
+ "tab": "Run",
+ },
+ "max_dist": {
+ "name": "Link radius (units)", "type": "float",
+ "default": DEFAULTS["max_dist"], "min": 0.1, "max": 1000.0,
+ "step": 0.1, "tab": "Run",
+ },
+ }
+
+ def __init__(self, session, tree, src_plot):
+ super().__init__(session, tree)
+ self.src_plot = src_plot
+ self.window_id = getattr(src_plot, "window_id", None)
+ self.params: dict[str, Any] = dict(DEFAULTS)
+ self.labels = None # LabelStore, built on first stroke
+ self.classifier = None # ScribbleClassifier, after seg_train
+ #: ``{"frame", "labels", "rows", "contours", "count", "areas"}`` for the
+ #: frame the caret is showing — what ``commit()`` snapshots.
+ self.preview: dict[str, Any] | None = None
+ self.result_tree = None
+ #: The live preview outline group, and the navigator hooks that keep it
+ #: on the displayed frame. See :meth:`set_overlay` / :meth:`wire_navigator`.
+ self._ov_group = None
+ self._ov_box_group = None
+ self._ov_selectors: list = []
+ #: Latest frame the navigator asked for, and whether a preview for it is
+ #: already running — the latest-wins pair, see :func:`_preview_for_nav`.
+ self._nav_frame: int | None = None
+ self._nav_busy = False
+
+ # ── the signal this wizard segments ──────────────────────────────────────
+
+ def signal(self):
+ """The DISPLAYED node, not the root — segmenting a rebinned or cropped
+ view must segment what the user is looking at."""
+ return _current_signal(self.src_plot) or self.tree.root
+
+ def frames(self):
+ """``(n_frames, get_frame, (h, w))`` — one frame at a time, never the
+ stack (CLAUDE.md memory safety)."""
+ return frames_of(self.signal())
+
+ def scale_units(self) -> tuple[float, str]:
+ try:
+ ax = self.signal().axes_manager.signal_axes[0]
+ return float(ax.scale), str(ax.units or "px")
+ except Exception as exc:
+ log.debug("[seg] reading signal calibration failed: %s", exc)
+ return 1.0, "px"
+
+ def set_params(self, payload: dict | None, *, merge: bool = False) -> None:
+ """Coerce *payload* into ``self.params``, WITH the signal's scale.
+
+ The one place params are assigned, because of the second line. The face
+ controls `merge_nm` / `min_nm` are physical, and :func:`_nm_to_px`
+ converts them with ``p["scale"]`` — but :func:`_coerce` builds `p` from
+ the ``DEFAULTS`` keys alone, so nothing put a scale there and the
+ conversion took its uncalibrated branch on EVERY signal. The slider said
+ "50 nm" and the backend merged at 50 *pixels*: silent, and wrong by
+ exactly the magnification. Stashing it at dispatch is what `_nm_to_px`
+ already documents ("stashed on the params at dispatch"); it just was
+ never done.
+
+ Assigning `wiz.params = _coerce(...)` directly re-opens that hole, so
+ don't — the scale is only correct if it is refreshed here, on the
+ DISPLAYED node, which is what a rebinned or cropped view changes.
+
+ The stashed value is **nm per pixel**, not the raw axis scale, and it is
+ 0.0 when the axis is not a real-space length at all. A signal calibrated
+ in `nm^-1` reported a perfectly good positive scale, so converting
+ against it produced a merge radius that was silently wrong by the camera
+ length; 0.0 routes those signals to the pixel fallback instead, and
+ `_face_units` relabels the sliders so the caret does not claim nm.
+ """
+ base = {**self.params, **(payload or {})} if merge else payload
+ p = _coerce(base)
+ scale, units = self.scale_units()
+ p["scale"] = _length_nm_per_px(scale, units)
+ self.params = p
+
+ def frame_index(self) -> int:
+ """Where the navigator is sitting.
+
+ Read from the navigation SELECTOR, not the Plot: ``current_indices``
+ lives on the selector, and looking for it on the Plot silently returns
+ None — the bug that made "fit spectrum" fit the navigation mean
+ (``fit_action.current_indices``).
+ """
+ npm = getattr(self.tree, "navigator_plot_manager", None)
+ if npm is None:
+ return 0
+ for sels in (getattr(npm, "navigation_selectors", {}) or {}).values():
+ for sel in sels:
+ idx = getattr(sel, "current_indices", None)
+ if idx is None:
+ continue
+ try:
+ return int(np.atleast_1d(np.asarray(idx)).ravel()[0])
+ except Exception as exc:
+ log.debug("[seg] reading navigator index failed: %s", exc)
+ return 0
+
+ # ── scribbles ────────────────────────────────────────────────────────────
+
+ def label_store(self):
+ """The accumulating :class:`LabelStore`, built on first use.
+
+ Built lazily because it needs the frame shape, and a flat index means
+ nothing without one — a store made against the wrong shape scatters
+ every stroke across the image.
+ """
+ if self.labels is None:
+ from spyde.particles import LabelStore
+ _n, _get, shape = self.frames()
+ self.labels = LabelStore(frame_shape=shape)
+ return self.labels
+
+ def class_report(self) -> list[dict[str, Any]]:
+ """Per-class labelled-pixel counts for the caret's class list.
+
+ Not decoration (plan B3): under-training a class is *the* failure mode
+ and these counts are how a user notices, so a class with zero pixels is
+ present in the list rather than absent from it.
+ """
+ if self.labels is None:
+ from spyde.particles import default_classes
+ return [dict(c.to_dict(), pixels=0) for c in default_classes()]
+ counts = self.labels.counts()
+ return [dict(c.to_dict(), pixels=int(counts.get(c.id, 0)))
+ for c in self.labels.classes]
+
+ # ── lifecycle ────────────────────────────────────────────────────────────
+
+ def remove(self) -> None:
+ if self._closed:
+ return
+ self._closed = True
+ self._unwire_navigator()
+ self._drop_overlay()
+ self._ov_cleared = False
+ self._ov_box_state = None
+ if getattr(self.tree, "_seg_wizard", None) is self:
+ self.tree._seg_wizard = None
+
+ def set_overlay(self, contours=None, box=None, labels=None,
+ full_shape=None, n_instances=None) -> None:
+ """Draw (or clear) the previewed instances as OUTLINES on the source plot.
+
+ Outlines and not a translucent raster mask, for two independent reasons —
+ the first is a bug, the second is why the bug could not just be patched:
+
+ * **A raster mask does not survive GPU tile mode.** A signal frame at or
+ above 1024 px is handed to anyplotlib's tiled display, whose base image
+ is drawn by WebGPU; ``set_overlay_mask`` composites onto the Canvas2D
+ context underneath it and is simply not visible. That is the whole
+ "106 particles and no overlay" report — the mask WAS being pushed
+ (``[plot] overlay mask set: N px`` in the log), it just never appeared.
+ Markers draw over the GPU base correctly, which is why the brush strokes
+ were visible in the same screenshot that had no overlay.
+ * **A full-resolution mask cannot follow the navigator.** At 4096² the
+ mask is 16.7 M px — ~16 MB down the PLOTAPP line protocol *per frame*.
+ This overlay re-draws on every navigator move, so the payload has to be
+ proportional to the number of PARTICLES, not to the number of pixels.
+ The same 106 particles are ~100 kB of polygon, and they stay crisp when
+ the user zooms in, which a mask rasterised at frame resolution does not.
+
+ *contours* are ``(k, 2)`` arrays of ``(y, x)`` **crop** pixels as
+ :func:`spyde.particles.measure_frame` returns them; *box* is the
+ ``(y0, x0, h, w)`` preview window they were measured in, or None when the
+ whole frame was segmented. The offset is applied here rather than by the
+ caller so that whatever ``_preview`` decided to crop stays in one place.
+ """
+ plot2d = getattr(self.src_plot, "_plot2d", None)
+ if plot2d is None:
+ return
+ group = self._overlay_group(plot2d)
+ if group is None:
+ return
+ # Hundreds of polygons is what makes the whole app sluggish — each is a
+ # path the renderer re-transforms every pan/zoom frame. Above the
+ # threshold draw ONE mask instead; the particle COUNT and every
+ # measurement are unaffected, only how they are drawn.
+ # The true instance count, which is NOT `len(contours)`: the preview
+ # skips outline extraction above the draw cap precisely because this
+ # branch is about to discard them. Keying the decision off the contour
+ # list would then read "0 instances", fall through to the polygon path,
+ # draw nothing, and leave the frame bare in exactly the case that most
+ # needs a mask.
+ n = n_instances if n_instances is not None else (
+ len(contours) if contours is not None else 0)
+ if n > _RASTER_ABOVE and labels is not None:
+ if self._set_raster_overlay(labels, box, full_shape):
+ self._ov_raster = True
+ # Drop the vector outlines so the two do not double up.
+ updates = {group: {"vertices_list": []}}
+ frame_group = self._window_group(plot2d)
+ if frame_group is not None:
+ updates[frame_group] = {"vertices_list": _box_poly(box)}
+ try:
+ _push_groups(plot2d, updates)
+ self._ov_cleared = False
+ self._ov_box_state = None
+ except Exception as exc:
+ log.debug("[seg] overlay push failed: %s", exc)
+ return
+ self._clear_raster_overlay()
+ polys = _contour_polys(contours, box)
+ # HARD CAP, independent of `_RASTER_ABOVE`. Reaching here with a huge n
+ # means the raster path was unavailable or failed, and the honest answer
+ # is then "too many to draw" — NOT 14028 paths the renderer must
+ # re-transform on every pan and zoom frame. That is a hang, and because
+ # thousands of translucent fills overlap into one flat sheet it is not
+ # even a legible one: the user sees a solid green block and no data.
+ # The COUNT still reports every instance; only the drawing is dropped,
+ # and the caret says so rather than leaving a silently empty frame.
+ if len(polys) > _MAX_OUTLINE_POLYS:
+ log.warning("[seg] %d outlines exceeds the %d draw cap and the raster "
+ "overlay was unavailable; drawing none",
+ len(polys), _MAX_OUTLINE_POLYS)
+ polys = []
+ updates = {group: {"vertices_list": polys}}
+ # The preview window's own outline. Without it "only the middle of my
+ # 4k frame has any segmentation" reads as a BUG rather than as the
+ # documented 1-megapixel budget: the caret says "preview window
+ # 1024x1024 px" in small print, but it cannot say WHERE, and on a 4096²
+ # frame the window is 1/16 of the area sitting in the middle of an
+ # otherwise untouched image. Drawing the boundary makes the empty region
+ # obviously "not looked at yet" instead of "looked at and found
+ # nothing".
+ frame_group = self._window_group(plot2d)
+ if frame_group is not None:
+ updates[frame_group] = {"vertices_list": _box_poly(box)}
+ try:
+ # `vertices_list` is the polygon group's own key — the same one
+ # `particle_overlay._payload` writes. A key the group does not know
+ # is accepted silently by `_push_groups` and simply never drawn.
+ _push_groups(plot2d, updates)
+ # Real outlines are up now, so the no-engine path must redraw rather
+ # than skip on its cached "already cleared" state.
+ self._ov_cleared = False
+ self._ov_box_state = None
+ except Exception as exc:
+ log.debug("[seg] overlay push failed: %s", exc)
+
+ def _set_raster_overlay(self, labels, box, full_shape) -> bool:
+ """Draw the instances as ONE mask instead of N polygons. True if drawn.
+
+ Uses ``Plot2D.set_overlay_mask``, which composites client-side onto the
+ transparent 2-D canvas that sits ABOVE the WebGPU canvas — so this works
+ on a GPU-rendered base, which a marker-layer raster would not.
+
+ THE TILING TRAP, and why the reduction is NOT done here any more. When a
+ large frame is in tile mode the renderer sizes the mask against
+ ``base_width || image_width`` — the OVERVIEW grid — while tile mode sets
+ ``image_width`` to the FULL native frame. This method used to reduce the
+ mask to the overview grid itself, which was right for the renderer and
+ REJECTED by ``set_overlay_mask``'s own shape check (it compared against
+ the image shape). The ``ValueError`` landed in the ``except`` below, was
+ logged at DEBUG, and every large-frame preview fell back to N polygons —
+ so on a 4096² frame the raster path could never run, which is precisely
+ the case it was added for. At N=14028 that fallback is what hangs the
+ renderer and paints the frame a solid sheet of green.
+
+ anyplotlib now accepts a full-resolution mask in tile mode and does the
+ block-ANY reduction itself (``_reduce_mask_any``), so there is ONE owner
+ of that arithmetic and it cannot disagree with the shape check sitting
+ next to it. Hand it the native-resolution mask and let it decide.
+ """
+ plot2d = getattr(self.src_plot, "_plot2d", None)
+ if plot2d is None or not hasattr(plot2d, "set_overlay_mask"):
+ return False
+ try:
+ lab = np.asarray(labels)
+ if lab.ndim != 2 or not lab.any():
+ return False
+ fh, fw = (int(full_shape[0]), int(full_shape[1])) if full_shape \
+ else lab.shape
+ mask = np.zeros((fh, fw), bool)
+ if box is not None:
+ y0, x0, h, w = (int(v) for v in box)
+ mask[y0:y0 + h, x0:x0 + w] = lab[:h, :w] > 0
+ else:
+ mask[:lab.shape[0], :lab.shape[1]] = lab > 0
+ plot2d.set_overlay_mask(mask, color=_PREVIEW_COLOR, alpha=_RASTER_ALPHA)
+ return True
+ except Exception as exc:
+ # WARNING, not debug. The fallback from here is N polygons, and N is
+ # only large enough to be here because it was already too large to
+ # draw — so a silent failure trades a missing overlay for a hung
+ # renderer. If this line appears, that is the bug.
+ log.warning("[seg] raster overlay failed (%s); falling back to outlines",
+ exc)
+ return False
+
+ def _clear_raster_overlay(self) -> None:
+ plot2d = getattr(self.src_plot, "_plot2d", None)
+ if plot2d is None or not hasattr(plot2d, "set_overlay_mask"):
+ return
+ if not getattr(self, "_ov_raster", False):
+ return
+ try:
+ plot2d.set_overlay_mask(None)
+ except Exception as exc:
+ log.debug("[seg] clearing the raster overlay failed: %s", exc)
+ self._ov_raster = False
+
+ def show_preview_window(self) -> None:
+ """Draw the preview-window box ALONE, with no instance outlines.
+
+ For the states that have no engine to run: an untrained Scribble, or
+ Prompt before anything is prompted. Callers that DO have a result go
+ through :meth:`set_overlay`, which draws both.
+
+ Clearing the outlines here is the other half of the fix — switching to
+ an untrained engine must not leave the previous one's instances on
+ screen looking like the new engine's answer.
+
+ BOTH drawing routes, which is the half that was missing. This cleared
+ the vector outlines and left any RASTER mask in place, so above
+ `_RASTER_ABOVE` instances the previous engine's result survived the
+ switch untouched. On the Scribble tab that is not merely stale, it is
+ disabling: the mask covers the image you have to paint on, and the
+ reported symptom was exactly "moved to scribble, immediately unusable".
+ """
+ plot2d = getattr(self.src_plot, "_plot2d", None)
+ if plot2d is None:
+ return
+ self._clear_raster_overlay()
+ try:
+ _n, get_frame, _shape = self.frames()
+ _frame, box = _preview_window(np.asarray(get_frame(self.frame_index())))
+ except Exception as exc:
+ log.debug("[seg] preview-window box: no frame to size it from: %s", exc)
+ return
+ # Only push when something actually CHANGES. `seg_tune` fires on every
+ # slider tick and lands here whenever there is no engine, and
+ # `_push_groups` falls back to `MarkerGroup.set`, which re-serialises
+ # the panel — so an unconditional push would put a full serialisation
+ # on every tick of a drag to redraw a rectangle that did not move.
+ state = (tuple(box) if box is not None else None,
+ getattr(self, "_ov_cleared", False))
+ if state == getattr(self, "_ov_box_state", None):
+ return
+ updates = {}
+ group = self._overlay_group(plot2d)
+ if group is not None and not getattr(self, "_ov_cleared", False):
+ updates[group] = {"vertices_list": []} # no instances yet
+ frame_group = self._window_group(plot2d)
+ if frame_group is not None:
+ updates[frame_group] = {"vertices_list": _box_poly(box)}
+ if not updates:
+ return
+ try:
+ _push_groups(plot2d, updates)
+ self._ov_cleared = True
+ self._ov_box_state = (tuple(box) if box is not None else None, True)
+ except Exception as exc:
+ log.debug("[seg] preview-window push failed: %s", exc)
+
+ def _overlay_group(self, plot2d):
+ """The lazily-created polygon group the preview outlines live in."""
+ if self._ov_group is not None:
+ return self._ov_group
+ try:
+ self._ov_group = plot2d.add_polygons(
+ [], name="seg_preview_outline",
+ facecolors=_PREVIEW_COLOR, edgecolors=_PREVIEW_COLOR,
+ linewidths=_PREVIEW_WIDTH, alpha=_PREVIEW_ALPHA,
+ transform="data")
+ except Exception as exc:
+ log.debug("[seg] creating the preview overlay group failed: %s", exc)
+ return self._ov_group
+
+ def _window_group(self, plot2d):
+ """The lazily-created outline of the PREVIEW WINDOW (empty when whole)."""
+ if self._ov_box_group is not None:
+ return self._ov_box_group
+ try:
+ self._ov_box_group = plot2d.add_polygons(
+ [], name="seg_preview_window",
+ facecolors=None, edgecolors=_PREVIEW_WINDOW_COLOR,
+ linewidths=1.0, transform="data")
+ except Exception as exc:
+ log.debug("[seg] creating the preview-window outline failed: %s", exc)
+ return self._ov_box_group
+
+ def _drop_overlay(self) -> None:
+ self._clear_raster_overlay()
+ for attr in ("_ov_group", "_ov_box_group"):
+ group = getattr(self, attr, None)
+ if group is None:
+ continue
+ try:
+ group.remove()
+ except Exception as exc:
+ log.debug("[seg] removing %s failed: %s", attr, exc)
+ setattr(self, attr, None)
+
+ # ── following the navigator ──────────────────────────────────────────────
+
+ def wire_navigator(self) -> None:
+ """Re-preview when the navigator moves, so the outlines follow the frame.
+
+ Nothing subscribed to the navigator before this: ``frame_index()`` only
+ READ the selector when something else asked for a preview, so scrolling
+ through a movie left the outlines describing whichever frame was showing
+ when you last touched the caret.
+ """
+ self._ov_selectors = _navigator_selectors_for(self.tree, self.src_plot)
+ for sel in self._ov_selectors:
+ if self._on_indices not in sel.index_hooks:
+ sel.index_hooks.append(self._on_indices)
+ if not self._ov_selectors:
+ log.debug("[seg] no navigator selectors — the preview will not "
+ "follow the frame (a single image has none, which is fine)")
+
+ def _unwire_navigator(self) -> None:
+ for sel in self._ov_selectors:
+ if self._on_indices in sel.index_hooks:
+ sel.index_hooks.remove(self._on_indices)
+ self._ov_selectors = []
+
+ def _on_indices(self, indices) -> None:
+ """Navigation moved. **Runs on the ``_NavDispatcher`` thread.**
+
+ Does no figure work here and submits no compute here — both are the main
+ thread's business (CLAUDE.md's threading contract), so this only records
+ the frame and marshals.
+ """
+ if self._closed:
+ return
+ try:
+ frame = int(np.asarray(indices).ravel()[0])
+ except Exception:
+ return
+ if frame == self._nav_frame:
+ return
+ self._nav_frame = frame
+ dispatch = getattr(self.session, "_dispatch_to_main", None)
+ if dispatch is None:
+ _preview_for_nav(self)
+ return
+ dispatch(lambda: _preview_for_nav(self))
+
+ def commit(self):
+ """Snapshot the PREVIEWED frame as a committed label-image tree.
+
+ ``seg_run`` is the whole-movie door and goes through
+ ``open_particle_tree`` progressively; Commit is the other half of the
+ §0.8 contract — you tuned on one frame, and this keeps that frame's
+ result as a dataset without paying for the movie. It is also the whole
+ of the single-2-D-image shape in plan §0.10, where there is no movie to
+ run over.
+
+ It goes through ``commit_result_tree`` rather than
+ ``open_particle_tree`` because a ONE-frame particle tree is not
+ constructible: ``open_particle_tree`` builds a ``(1, h, w)`` label
+ movie, hyperspy reads the leading axis as a size-1 navigation axis, and
+ ``MultiplotManager`` has no selector for one — it raises. A single frame
+ is a 2-D label IMAGE, which is exactly what ``commit_result_tree``
+ expects, and the store rides along in ``attrs`` so
+ ``requires_particles`` unlocks the same downstream actions.
+ """
+ prev = self.preview
+ if prev is None or self.session is None:
+ emit_error("Segment Particles: nothing previewed to commit")
+ return None
+ # `SpyDEParticles.from_frames` requires one contour PER ROW, in order.
+ # The preview skips outline extraction above the draw cap, so committing
+ # that preview would build a store whose outlines silently do not
+ # correspond to its rows. Refuse, and say why — the only way to be here
+ # is a preview with thousands of instances, which is the failed-threshold
+ # case (`_threshold_failed`) and not something worth committing anyway.
+ if len(prev["contours"]) != len(prev["rows"]):
+ emit_error(
+ f"Segment Particles: {len(prev['rows'])} instances is too many "
+ "to commit as outlines — this is usually a threshold that "
+ "landed in the noise. Tighten the size filter, or train the "
+ "Scribble classifier, then commit.")
+ return None
+ return commit_single_frame(
+ self.session, self, prev["labels"], prev["rows"], prev["contours"],
+ int(prev["frame"]))
+
+
+# ── parameters ───────────────────────────────────────────────────────────────
+
+def _coerce(payload: dict | None) -> dict:
+ """Payload → a complete, valid parameter dict.
+
+ Every out-of-range value is corrected rather than raised on: these arrive
+ from a slider mid-drag, and a caret that errors while you are moving it is
+ unusable. The corrections that change what the user asked for
+ (``min_size``, ``local_size``) are echoed back in every preview so the
+ number on screen is the number that ran.
+ """
+ from spyde.particles import THRESHOLD_METHODS
+
+ p = dict(DEFAULTS)
+ payload = payload or {}
+ for k, default in DEFAULTS.items():
+ v = payload.get(k)
+ if v is None or v == "":
+ continue
+ try:
+ p[k] = bool(v) if isinstance(default, bool) else type(default)(v)
+ except (TypeError, ValueError) as exc:
+ log.debug("[seg] param %r=%r not coercible, keeping default: %s",
+ k, v, exc)
+
+ p["method"] = str(p["method"]).lower()
+ if p["method"] not in METHODS:
+ p["method"] = DEFAULTS["method"]
+ p["threshold"] = str(p["threshold"]).lower()
+ if p["threshold"] not in THRESHOLD_METHODS:
+ p["threshold"] = DEFAULTS["threshold"]
+ p["sensitivity"] = float(min(1.0, max(0.0, p["sensitivity"])))
+
+ # skimage's local thresholds require an odd window and SegmentParams raises
+ # rather than bumping it silently. Bump here (the caret gets the effective
+ # value back) so a slider that lands on an even number doesn't error.
+ p["local_size"] = max(3, int(p["local_size"]))
+ if p["local_size"] % 2 == 0:
+ p["local_size"] += 1
+
+ p["min_size"] = int(p["min_size"])
+ p["min_score"] = float(min(0.99, max(0.0, p.get("min_score", 0.0) or 0.0)))
+ p["min_size_floored"] = p["min_size"] < MIN_SIZE_FLOOR
+ if p["min_size_floored"]:
+ p["min_size"] = MIN_SIZE_FLOOR
+ return p
+
+
+def _segment_kwargs(p: dict) -> dict:
+ """The ``SegmentParams`` fields, as a plain dict.
+
+ A dict and not the dataclass because this crosses to dask workers inside
+ :class:`~spyde.particles.batch.EngineSpec`, and a spec that pickles without
+ dragging the segmentation modules onto the client's import path is one less
+ thing to go wrong in a worker with a different import order.
+ """
+ return dict(
+ threshold=p["threshold"], sensitivity=p["sensitivity"],
+ rb_kernel=int(p["rb_kernel"]), gaussian=float(p["gaussian"]),
+ invert=bool(p["invert"]), local_size=int(p["local_size"]),
+ watershed=bool(p["watershed"]), min_separation=int(p["min_separation"]),
+ marker_smooth=float(p["marker_smooth"]), min_size=_min_size_px(p),
+ max_size=int(p["max_size"]), clear_border=bool(p["clear_border"]),
+ merge_distance=_nm_to_px(p.get("merge_nm", 0.0), p),
+ )
+
+
+#: Axis units that are a REAL-SPACE LENGTH, and their size in nanometres. The
+#: face controls are in nm, so an axis calibrated in µm or Å converts through
+#: this rather than being divided by a raw number in the wrong unit.
+#:
+#: Anything NOT in here — `nm^-1` and friends from a reciprocal-space signal,
+#: `mrad`, `px`, an empty unit — is not a length, and then there is no physical
+#: distance to convert to. Those signals fall back to PIXELS and the caret
+#: relabels the two sliders accordingly (`_face_units`). A reciprocal axis is
+#: the case that made this necessary: dividing a nanometre by a value in nm⁻¹
+#: is dimensionally meaningless, and it silently produced a merge radius wrong
+#: by whatever the camera length happened to be.
+_NM_PER: dict[str, float] = {
+ "m": 1e9, "cm": 1e7, "mm": 1e6,
+ "um": 1e3, "µm": 1e3, "μm": 1e3, "micron": 1e3,
+ "nm": 1.0,
+ "a": 0.1, "å": 0.1, "ang": 0.1, "angstrom": 0.1,
+ "pm": 1e-3,
+}
+
+
+def _length_nm_per_px(scale: float, units: str) -> float:
+ """*scale* in nm/px, or 0.0 when the axis is not a real-space length."""
+ try:
+ s = float(scale)
+ except (TypeError, ValueError):
+ return 0.0
+ if not (s > 0):
+ return 0.0
+ factor = _NM_PER.get(str(units or "").strip().lower())
+ return s * factor if factor else 0.0
+
+
+def _face_units(scale: float, units: str) -> str:
+ """What the two face sliders are actually in: 'nm' or 'px'."""
+ return "nm" if _length_nm_per_px(scale, units) > 0 else "px"
+
+
+def _nm_to_px(value_nm: float, p: dict) -> float:
+ """A face control in nanometres -> pixels, using the signal's own scale.
+
+ The face controls are physical so they mean the same thing at any
+ magnification, and so the number matches what the scale bar says. `scale` is
+ stashed on the params at dispatch **already converted to nm per pixel**
+ (:meth:`SegmentWizard.set_params`); a missing or zero scale means the signal
+ is uncalibrated *or its axis is not a length at all*, and then the value IS
+ pixels rather than being silently divided by a number in the wrong unit.
+ """
+ try:
+ v = float(value_nm or 0.0)
+ except (TypeError, ValueError):
+ return 0.0
+ if v <= 0:
+ return 0.0
+ try:
+ scale = float(p.get("scale") or 0.0)
+ except (TypeError, ValueError):
+ scale = 0.0
+ return v / scale if scale > 0 else v
+
+
+def _min_size_px(p: dict) -> int:
+ """`min_nm` (a DIAMETER) wins over the raw pixel `min_size` when set.
+
+ Area, not diameter, is what the size filter compares -- a particle of
+ diameter d covers ~pi/4 d^2 pixels -- so converting the physical control
+ means going through the area, not handing the diameter straight over.
+ """
+ d_px = _nm_to_px(p.get("min_nm", 0.0), p)
+ if d_px > 0:
+ return max(1, int(round(3.14159265 / 4.0 * d_px * d_px)))
+ return int(p["min_size"])
+
+
+def _segment_params(p: dict):
+ from spyde.particles import SegmentParams
+ return SegmentParams(**_segment_kwargs(p))
+
+
+def _movie_array(signal):
+ """The ``(n, h, w)`` array behind *signal*, or None when there isn't one.
+
+ None is a legitimate answer (a frame source that is a callable or a
+ sequence), and :func:`~spyde.particles.batch.segment_movie` falls back to
+ the streaming accessor for it. Never touches the data itself — reading
+ ``.data`` on a lazy signal hands back the dask graph, not the movie.
+ """
+ data = getattr(signal, "data", None)
+ return data if getattr(data, "ndim", 0) == 3 else None
+
+
+#: How long the batch waits for the cluster before running locally. The cluster
+#: is built on a background thread, so a run fired seconds after a load can
+#: arrive first; falling straight through would silently cost the whole fan-out.
+_CLIENT_WAIT_S = 30.0
+
+
+def _batch_client(session, stopped=None):
+ """The distributed client for the batch, waiting briefly for it to come up.
+
+ Mirrors ``_do_compute_vectors``: we are already on a worker thread, so
+ blocking here doesn't freeze the UI, and the alternative — silently taking
+ the local thread-pool path — is exactly the "why is this slow" report this
+ work exists to fix. Returns None under ``SPYDE_NO_DASK=1`` (the migrated-test
+ mode, where the manager exists but never starts).
+ """
+ import os
+ if os.environ.get("SPYDE_NO_DASK") == "1":
+ return None
+ dm = getattr(session, "dask_manager", None)
+ if dm is None:
+ return None
+ client = getattr(dm, "client", None)
+ if client is not None:
+ return client
+ from spyde.compute_dispatch import reliable_sleep
+ deadline = time.monotonic() + _CLIENT_WAIT_S
+ while client is None and time.monotonic() < deadline:
+ if stopped is not None and stopped[0]:
+ return None
+ reliable_sleep(0.1)
+ client = getattr(dm, "client", None)
+ if client is None:
+ log.warning("[seg] no Dask client after %.0f s — segmenting locally",
+ _CLIENT_WAIT_S)
+ return client
+
+
+def frames_of(signal):
+ """``(n_frames, get_frame, (h, w))`` for a movie **or** a single image.
+
+ Delegates to :func:`spyde.drift.frame_source`, which is the shared streaming
+ accessor for exactly this — it exists so callers cannot reach ``.data`` and
+ accidentally compute the stack. It requires a 1-D navigation axis, so the
+ plain-2-D-image case (plan §0.10) is wrapped here as a one-frame stack
+ rather than duplicating the accessor.
+ """
+ am = signal.axes_manager
+ if int(am.signal_dimension) != 2:
+ raise TypeError(
+ "Segment Particles needs 2-D image frames; got signal_dimension="
+ f"{int(am.signal_dimension)}")
+ nav = int(am.navigation_dimension)
+ if nav > 1:
+ # frame_source raises for this too, but with a message about drift.
+ raise TypeError(
+ "needs a movie (1-D time navigation) or a single image; got "
+ f"navigation_dimension={nav}. Reduce a 4D-STEM scan to a virtual "
+ "image first — plan §0.10.")
+ if nav == 0:
+ data = signal.data
+
+ def get_frame(i: int, _d=data) -> np.ndarray:
+ arr = _d.compute() if hasattr(_d, "compute") else _d
+ return np.asarray(arr)
+
+ h, w = int(signal.data.shape[-2]), int(signal.data.shape[-1])
+ return 1, get_frame, (h, w)
+
+ from spyde.drift import frame_source
+ return frame_source(signal)
+
+
+# ── the three engines, behind one call ───────────────────────────────────────
+
+def _engine(wiz: SegmentWizard, p: dict) -> Callable[[np.ndarray], np.ndarray] | None:
+ """A ``frame → int32 labels`` callable for the selected engine, or None when
+ the engine cannot run yet (the caller has already been told why)."""
+ sp = _segment_params(p)
+ method = p["method"]
+
+ if method == "classical":
+ from spyde.particles import segment_frame
+ return lambda frame: segment_frame(frame, sp)
+
+ if method == "scribble":
+ clf = wiz.classifier
+ if clf is None or not clf.is_trained:
+ emit_status("Segment Particles: paint a few scribbles — including at "
+ "least one FAINT particle — then press Train.")
+ return None
+ return lambda frame: clf.segment(frame, sp)
+
+ # plan B4: EfficientSAM-Ti through the existing model registry. The tab
+ # exists so the caret can render it; the engine does not.
+ emit_status("Segment Particles: prompt segmentation is not installed yet — "
+ "use the Classical or Scribble tab.")
+ return None
+
+
+# ── controller resolution ────────────────────────────────────────────────────
+
+def _wizard(session, plot) -> SegmentWizard | None:
+ _src, tree = _src_plot_tree(session, plot)
+ wiz = getattr(tree, "_seg_wizard", None) if tree is not None else None
+ return wiz if (wiz is not None and not wiz._closed) else None
+
+
+def _emit(wiz: SegmentWizard, msg: dict) -> None:
+ msg.setdefault("window_id", wiz.window_id)
+ emit(msg)
+
+
+def _emit_state(wiz: SegmentWizard) -> None:
+ """The caret's authoritative state: engine, classes + pixel counts, which
+ frames carry labels, and the EFFECTIVE parameters."""
+ try:
+ n_frames, _get, shape = wiz.frames()
+ except Exception:
+ n_frames, shape = 1, (0, 0)
+ _emit(wiz, {
+ "type": "seg_state",
+ "method": wiz.params["method"],
+ "frame": wiz.frame_index(),
+ "n_frames": int(n_frames),
+ "frame_shape": [int(shape[0]), int(shape[1])],
+ "classes": wiz.class_report(),
+ "labelled_frames": (wiz.labels.labelled_frames() if wiz.labels else []),
+ "trained": bool(wiz.classifier is not None and wiz.classifier.is_trained),
+ "params": {k: v for k, v in wiz.params.items()},
+ })
+
+
+# ── preview (the CURRENT frame, and only the current frame) ──────────────────
+
+#: Pixel budget for ONE preview segmentation. Above this the preview runs on a
+#: centred CROP at full resolution instead of the whole frame.
+#:
+#: Measured, one frame, classical path (segment + measure):
+#:
+#: =========== =========
+#: frame cost
+#: =========== =========
+#: 256^2 0.27 s
+#: 1024^2 0.57 s
+#: 2048^2 1.69 s
+#: **4096^2** **8.36 s**
+#: =========== =========
+#:
+#: ``segment_frame`` is 7.6 s of that 8.36 s — watershed over 16.7 M pixels. The
+#: caret re-previews on every sensitivity nudge, so on a real 4k in-situ movie the
+#: whole caret reads as hung. It was not hung; it was doing 8 seconds of work per
+#: keystroke.
+#:
+#: **Crop, do NOT downsample.** Downsampling would be cheaper still, but it makes
+#: the preview a DIFFERENT computation from the run it is supposed to predict:
+#: plan §0.9 records that the fine feature scales are what find small faint
+#: particles, and a preview that silently detects a different population is worse
+#: than a slow one. A crop at full resolution runs the identical algorithm on
+#: identical pixels, so what it shows is exactly what the run will do there.
+#:
+#: 1024^2 keeps a preview near half a second — inside the interaction budget with
+#: room for a slower machine.
+_PREVIEW_PIXEL_BUDGET = 1024 * 1024
+
+
+def _preview_window(frame: np.ndarray) -> tuple[np.ndarray, tuple[int, int, int, int] | None]:
+ """``(frame_or_crop, box)`` — bound one preview's cost by AREA, not by scale.
+
+ Returns the frame untouched (and ``box=None``) when it already fits the
+ budget, which is the common case for a tutorial-sized movie and means the
+ fast path pays nothing for this. Otherwise a centred crop of the same aspect
+ ratio, at full resolution, plus its ``(y0, x0, h, w)`` so the caller can tell
+ the user what it measured.
+ """
+ h, w = frame.shape[:2]
+ if h * w <= _PREVIEW_PIXEL_BUDGET:
+ return frame, None
+ shrink = (h * w / _PREVIEW_PIXEL_BUDGET) ** 0.5
+ ch = max(64, int(h / shrink))
+ cw = max(64, int(w / shrink))
+ y0 = max(0, (h - ch) // 2)
+ x0 = max(0, (w - cw) // 2)
+ return frame[y0:y0 + ch, x0:x0 + cw], (y0, x0, ch, cw)
+
+
+#: The live preview outline: a SATURATED green, filled at the same 25% the
+#: committed result overlay uses (``particle_overlay.FILL_ALPHA``) so tuning and
+#: result look like the same thing.
+#:
+#: Green because the particle brush class paints ORANGE and both are on screen
+#: at once while you scribble — "what I told it" and "what it found" have to be
+#: separable at a glance. Saturated rather than the retired mask's pale sage
+#: (``#a6e3a1``): that was chosen for a 35%-alpha area fill, and as a 1 px
+#: outline on grey EM data it is very close to invisible.
+_PREVIEW_COLOR = "#40e070"
+_PREVIEW_ALPHA = 0.25
+_PREVIEW_WIDTH = 1.0
+
+#: The preview WINDOW's outline. Deliberately a neutral light grey rather than
+#: any of the data colours: it is UI chrome saying "this is the region that was
+#: looked at", and must not be mistaken for a found particle (green) or for a
+#: painted class (orange / blue / grey / pink).
+_PREVIEW_WINDOW_COLOR = "#cdd6f4"
+
+#: Above this many instances the overlay switches from one polygon per particle
+#: to a SINGLE raster mask. Measured symptom: several hundred vector contours
+#: make the whole app sluggish, because every one is a path the renderer
+#: re-transforms on each pan/zoom frame. A mask is one image however many
+#: particles it contains.
+#:
+#: Below the threshold the vector outlines stay, and they are worth keeping:
+#: they are crisp at any zoom and each is a real object the UI can hover.
+_RASTER_ABOVE = 100
+
+#: The absolute ceiling on polygons handed to the renderer, whatever route got
+#: us here. `_RASTER_ABOVE` chooses the nicer drawing; this one is the seatbelt
+#: for when the raster is unavailable or fails, because the fallback used to be
+#: unbounded — a report of 14028 instances on a real in-situ frame meant 14028
+#: filled paths, which hangs the renderer and composites into one flat green
+#: sheet that shows nothing at all. Well above any legitimate outline count.
+_MAX_OUTLINE_POLYS = 1500
+
+#: The raster overlay's colour/opacity. Deliberately the same green as the
+#: vector outlines so crossing the threshold does not look like a mode change.
+_RASTER_ALPHA = 0.45
+
+#: A preview that calls this much of the window foreground, AND shatters it into
+#: this many pieces, is a FAILED THRESHOLD being reported as a result.
+#:
+#: Measured on a synthetic stand-in for the reported frame (low-contrast noisy
+#: support film, 8 faint dark particles, 1024²) — a global threshold has no
+#: bimodal histogram to find here, so otsu lands in the middle of the noise:
+#:
+#: defaults ............................. 4873 instances, 39% coverage
+#: min_size=200 .......................... 208 instances, 14% coverage
+#: min_size=2000 .......................... 17 instances, 7.2% coverage
+#: watershed off .......................... 751 instances, 40% coverage
+#: gaussian=2 + min_size=200 + no watershed . 8 instances, 52% coverage
+#:
+#: The last row is the trap and the reason the test is on BOTH numbers: 8
+#: instances looks like the 8 real particles, but at 52% coverage those 8 bodies
+#: ARE the film. Neither number alone is diagnostic — a genuinely crowded frame
+#: can be 40% covered by real particles, and a coarse threshold can find 8 real
+#: objects. Together they are: thousands of pieces AND half the frame is the
+#: signature of noise being segmented.
+#:
+#: Deliberately NOT a silent auto-correction. The plan's §0.9 answer to this
+#: data is the scribble classifier, and no threshold tweak substitutes for it —
+#: so the caret's job is to say the threshold failed and point at Scribble, not
+#: to quietly pick different parameters that fail differently.
+_FAIL_COVERAGE = 0.25
+_FAIL_COUNT = 500
+
+
+def _threshold_failed(count: int, coverage: float) -> bool:
+ """True when a preview is noise being reported as particles.
+
+ Both conditions, for the reason spelled out on `_FAIL_COVERAGE`.
+ """
+ return count >= _FAIL_COUNT and coverage >= _FAIL_COVERAGE
+
+
+def _box_poly(box) -> list:
+ """The preview window as a one-polygon list, or empty when there is none.
+
+ ``box`` is ``(y0, x0, h, w)``; markers want ``(x, y)``. An absent box means
+ the WHOLE frame was segmented, and then the outline must be cleared rather
+ than left showing the previous frame's window — the caret can switch between
+ cropped and whole (a small frame, or a changed budget) without closing.
+ """
+ if box is None:
+ return []
+ y0, x0, h, w = (float(v) for v in box)
+ return [np.array([[x0, y0], [x0 + w, y0], [x0 + w, y0 + h], [x0, y0 + h]],
+ dtype=np.float32)]
+
+
+def _contour_polys(contours, box) -> list:
+ """``measure_frame`` contours → marker polygons, offset back onto the frame.
+
+ Two conversions, both of which are silent-wrong-picture bugs if skipped:
+
+ * **(y, x) → (x, y).** Contours are stored row-major like the array they came
+ from; marker offsets are ``(x, y)``. Getting this wrong transposes every
+ outline about the diagonal, which on a square frame still *looks* like
+ plausible particles (``particle_overlay.contour_xy`` is the same swap).
+ * **The crop offset.** A preview of a big frame is measured on a 1024² CROP
+ (``_preview_window``), so its contours start at ``(0, 0)`` of that window.
+ Drawn unshifted they pile up in the frame's corner instead of over the
+ region they describe — the same trap the mask path documented.
+ """
+ if contours is None:
+ return []
+ dy, dx = (float(box[0]), float(box[1])) if box is not None else (0.0, 0.0)
+ polys = []
+ for c in contours:
+ arr = np.asarray(c, dtype=np.float32)
+ if arr.ndim != 2 or len(arr) < 3:
+ continue # a polygon needs three vertices
+ polys.append(np.column_stack([arr[:, 1] + dx, arr[:, 0] + dy])
+ .astype(np.float32))
+ return polys
+
+
+def _preview_for_nav(wiz: SegmentWizard) -> None:
+ """Re-preview because the navigator moved. LATEST FRAME WINS.
+
+ A 4096² preview is ~600 ms and a scroll emits an index change per step, so
+ firing one preview per step would queue dozens of them and paint the frames
+ out of order as they landed. Only ONE is ever in flight; a frame requested
+ while it runs replaces any other waiting frame and is fired when it lands.
+
+ This is Live-Display §2's latest-wins coalescing, not a queue and not a
+ self-pacing gate: ``_nav_busy`` is cleared by BOTH the success and the
+ failure path of the preview it guards, so a failing frame cannot wedge it.
+ """
+ if wiz._closed or wiz._nav_busy:
+ return # the in-flight one will re-fire
+ wiz._nav_busy = True
+ # `current_gen()`, NOT `guard()`. Scrolling is not a new interaction, and
+ # bumping the generation here cancels whatever the user actually started —
+ # a navigator move during `seg_train` made the trained classifier land on a
+ # stale generation and get dropped, so Train silently never finished.
+ _preview(wiz, wiz.current_gen())
+
+
+def _chase_nav(wiz: SegmentWizard) -> None:
+ """After a preview lands: if the navigator has moved on, go again.
+
+ Called only once the new ``wiz.preview`` is in place, because the comparison
+ is "is what we are now showing the frame that was last asked for" — run
+ before the assignment it would read the PREVIOUS frame and re-fire forever.
+ Releasing the gate is deliberately NOT done here (see the call sites): it has
+ to happen even when the result is superseded or failed, and this function
+ returns early in both cases.
+ """
+ if wiz._closed or wiz._nav_frame is None:
+ return
+ if (wiz.preview or {}).get("frame") == wiz._nav_frame:
+ return # what just landed IS the newest frame — nothing to chase
+ _preview_for_nav(wiz)
+
+
+def filter_by_score(rows, contours, min_score: float):
+ """Keep only instances scoring at or above *min_score*.
+
+ Separate from the measurement on purpose: this is what the caret's single
+ "Confidence" slider calls, and it must be able to run WITHOUT re-segmenting
+ — the score already lives in the row (see
+ :func:`spyde.particles.measure.particle_scores`), so a drag is a numpy mask
+ over a few hundred rows rather than a fresh segmentation of a 1-megapixel
+ window.
+
+ ``min_score <= 0`` returns the inputs untouched, so the default costs
+ nothing and behaves exactly as before this control existed.
+ """
+ import numpy as _np
+ from spyde.signals.particles import COL as _COL
+
+ if min_score is None or float(min_score) <= 0.0 or len(rows) == 0:
+ return rows, contours
+ keep = _np.asarray(rows[:, _COL["score"]] >= float(min_score), bool)
+ if keep.all():
+ return rows, contours
+ kept_rows = _np.ascontiguousarray(rows[keep])
+ kept_contours = ([c for c, k in zip(contours, keep) if k]
+ if contours is not None else contours)
+ return kept_rows, kept_contours
+
+
+def _preview(wiz: SegmentWizard, gen: int) -> None:
+ """Segment the displayed frame on a worker and paint the result.
+
+ Generation-guarded at BOTH ends: a superseded tune must neither paint nor
+ leave a stale overlay behind (the caret can be closed mid-compute).
+ """
+ p = dict(wiz.params)
+ engine = _engine(wiz, p)
+ if engine is None:
+ # No engine yet — an untrained Scribble, or Prompt before any prompt.
+ # Still show WHERE the preview window is, and clear any outlines the
+ # previous engine left.
+ #
+ # The window box used to be drawn only as a side effect of a successful
+ # segmentation, which made it lie in three ways: it appeared only under
+ # Classical (the one engine that always has a solver), it SURVIVED a
+ # switch to an untrained Scribble because this early return painted
+ # nothing and cleared nothing, and after toggling the caret off and on
+ # it never came back. The box documents where the 1-megapixel budget
+ # looks; that is true whenever the caret is open and has nothing to do
+ # with whether an engine is trained.
+ wiz.show_preview_window()
+ return
+ t = wiz.frame_index()
+ scale, units = wiz.scale_units()
+
+ def _work():
+ _n, get_frame, _shape = wiz.frames()
+ full = np.asarray(get_frame(t))
+ frame, box = _preview_window(full)
+ t0 = time.perf_counter()
+ labels = engine(frame)
+ from spyde.particles import measure_frame
+ # Outlines only when they can actually be drawn. `labels.max()` is the
+ # instance count for the price of one pass, and above the draw cap the
+ # overlay discards the polygons anyway — extracting thousands of them
+ # first is pure latency on the tune the user is waiting for (measured:
+ # ~half of a 1353 ms preview at 4873 instances). The measured PROPERTIES
+ # are all still computed, so the count, the histogram, the median and
+ # the confidence filter are unaffected.
+ n_lab = int(np.asarray(labels).max()) if np.asarray(labels).size else 0
+ rows, contours = measure_frame(
+ labels, frame, t=t, scale=scale,
+ want_contours=n_lab <= _MAX_OUTLINE_POLYS)
+ n_all = len(rows)
+ rows, contours = filter_by_score(rows, contours, p.get("min_score", 0.0))
+ # What FRACTION of the previewed window was called foreground. This is
+ # the one number that separates "found the particles" from "the
+ # threshold landed inside the noise" — see `_threshold_failed`.
+ coverage = float((np.asarray(labels) > 0).mean()) if labels.size else 0.0
+ return {"frame": t, "labels": labels, "rows": rows, "n_all": n_all,
+ "contours": contours, "elapsed": time.perf_counter() - t0,
+ "coverage": coverage, "box": box, "full_shape": full.shape}
+
+ def _done(res):
+ # Release the navigator gate BEFORE the generation guard: a superseded or
+ # closed preview still has to hand the gate back, or the next navigator
+ # move finds `_nav_busy` set forever and the outlines stop following the
+ # frame. The retired self-pacing gates in Live-Display §2 wedged exactly
+ # like this.
+ wiz._nav_busy = False
+ if not wiz.still(gen) or wiz._closed:
+ return
+ rows = res["rows"]
+ from spyde.signals.particles import COL
+ areas = (rows[:, COL["area"]] if len(rows)
+ else np.zeros(0, np.float32))
+ wiz.preview = {"frame": res["frame"], "labels": res["labels"],
+ "rows": rows, "contours": res["contours"],
+ "count": int(len(rows)), "areas": areas,
+ "box": res.get("box")}
+ # The outlines carry the crop offset themselves (`_contour_polys`), so
+ # unlike the retired mask path nothing has to be re-rasterised onto a
+ # full-frame array first — at 4096² that array alone was 16 MB per frame.
+ wiz.set_overlay(res["contours"], res.get("box"),
+ labels=res.get("labels"),
+ full_shape=res.get("full_shape"),
+ n_instances=int(len(rows)))
+ # The count reported is the count AFTER the size filter (plan §0.9b) —
+ # `split_instances` applies min_size last, so `rows` is already filtered
+ # and this number is the one the histogram below describes.
+ _emit(wiz, {
+ "type": "seg_preview",
+ "frame": int(res["frame"]),
+ "count": int(len(rows)),
+ "areas": [float(a) for a in areas[:_MAX_AREAS_SENT]],
+ "median_area": (float(np.median(areas)) if areas.size else 0.0),
+ "units": units,
+ "method": p["method"],
+ "min_size": int(p["min_size"]),
+ "min_size_floored": bool(p["min_size_floored"]),
+ "elapsed_ms": round(1000.0 * res["elapsed"], 1),
+ # Present only when the frame was too big to preview whole. The caret
+ # must say so: otherwise "12 particles on this frame" is a lie about a
+ # 4096² frame where only the middle megapixel was looked at.
+ "preview_box": (None if res.get("box") is None
+ else [int(v) for v in res["box"]]),
+ # The threshold-failure verdict (see `_threshold_failed`). Reported
+ # rather than silently swallowed: the count is still true, it just
+ # is not an ANSWER, and the caret has to say which of the two it is
+ # holding.
+ "coverage": round(float(res.get("coverage") or 0.0), 4),
+ "threshold_failed": _threshold_failed(
+ int(len(rows)), float(res.get("coverage") or 0.0)),
+ # What the two FACE sliders are in — 'nm' on a real-space axis, 'px'
+ # when the axis is not a length (a reciprocal-space signal reports a
+ # healthy positive scale in nm^-1, and labelling that 'nm' is a
+ # claim about the scale bar that is simply false).
+ "face_units": _face_units(scale, units),
+ })
+ if p["min_size_floored"]:
+ emit_status(
+ f"Segment Particles: min size raised to {MIN_SIZE_FLOOR} px — "
+ "at 0 the split returns background speckle as particles "
+ "(measured: 33 instances where 9 are real).")
+ _chase_nav(wiz)
+
+ def _fail(exc):
+ wiz._nav_busy = False # never wedge the gate — see `_done`
+ if wiz.still(gen):
+ emit_error(f"Segment Particles preview failed: {exc}")
+
+ run_on_worker(wiz.session, _work, name="seg-preview",
+ on_done=_done, on_error=_fail)
+
+
+# ── staged handlers ──────────────────────────────────────────────────────────
+
+def seg_open(session, plot, payload) -> None:
+ """Caret mounted: build the controller and preview the displayed frame."""
+ src, tree = _src_plot_tree(session, plot)
+ if src is None or tree is None:
+ emit_error("Segment Particles: no active dataset")
+ return
+ try:
+ frames_of(_current_signal(src) or tree.root)
+ except TypeError as exc:
+ emit_error(f"Segment Particles: {exc}")
+ return
+
+ existing = getattr(tree, "_seg_wizard", None)
+ if existing is not None and not existing._closed:
+ # Idempotent re-open: adopt the new parameters and re-preview rather
+ # than building a second controller with its own scribbles.
+ existing.set_params(payload)
+ gen = existing.guard()
+ _emit_state(existing)
+ _preview(existing, gen)
+ return
+
+ wiz = SegmentWizard(session, tree, src)
+ wiz.set_params(payload)
+ # BEFORE anything deferred: React StrictMode fires open/close/open
+ # synchronously, so the close's bump must be able to invalidate this open.
+ gen = wiz.guard()
+ tree._seg_wizard = wiz
+ # Follow the navigator from the moment the caret opens, not from the first
+ # Train: scrolling with the caret up should keep the outlines on the frame
+ # you are looking at whichever engine is selected.
+ wiz.wire_navigator()
+ _emit_state(wiz)
+ _preview(wiz, gen)
+
+
+def seg_close(session, plot, payload=None) -> None:
+ """Caret unmounted: invalidate in-flight work FIRST, then tear down."""
+ _src, tree = _src_plot_tree(session, plot)
+ if tree is None:
+ return
+ # This IS WizardController.cancel_inflight (same `_seg_run_gen` key), done
+ # on the TREE so it fires even when there is no controller yet: a StrictMode
+ # open whose worker has not landed must still be cancelled.
+ bump_generation(tree, "_seg_run_gen")
+ wiz = getattr(tree, "_seg_wizard", None)
+ if wiz is not None:
+ _detach_brush(wiz)
+ wiz.remove()
+
+
+def seg_set_method(session, plot, payload) -> None:
+ """Switch engine (classical | scribble | prompt) and re-preview."""
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ return
+ method = str((payload or {}).get("method", "")).lower()
+ if method not in METHODS:
+ emit_error(f"Segment Particles: unknown engine {method!r}")
+ return
+ wiz.params["method"] = method
+ # The brush exists only while Scribble is the engine: it floats over the
+ # image, so leaving it armed on Classical would put a paint cursor over data
+ # with nothing to paint into.
+ if method == "scribble":
+ _arm_brush(wiz)
+ else:
+ _detach_brush(wiz)
+ gen = wiz.guard()
+ _emit_state(wiz)
+ _preview(wiz, gen)
+
+
+def seg_tune(session, plot, payload) -> None:
+ """Debounced parameter change → re-preview the CURRENT frame only."""
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ return
+ wiz.set_params(payload, merge=True)
+ # Paint state must reach the WIDGET, not just the params dict — the widget is
+ # what tags a stroke, so a class change that stops here paints the old colour.
+ _sync_brush(wiz)
+ gen = wiz.guard()
+ _preview(wiz, gen)
+
+
+def seg_paint(session, plot, payload) -> None:
+ """One brush stroke into the :class:`LabelStore`.
+
+ Payload: ``{frame, points: [[y, x], …], class_id, erase, brush}``. Points
+ arrive in **image pixels** with no scale or offset applied (plan trap 6 —
+ anyplotlib 2-D widgets report pixels), so nothing is converted here.
+
+ Synchronous: a stroke is a few thousand indices and the caret's class counts
+ must be correct by the time the user lifts the brush.
+ """
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ return
+ payload = payload or {}
+ points = payload.get("points") or []
+ if not len(points):
+ return
+ if _paint_stroke(wiz, int(payload.get("frame", wiz.frame_index())), points,
+ int(payload.get("class_id", 0)),
+ bool(payload.get("erase")),
+ float(payload.get("brush", wiz.params["brush"]))):
+ _emit_state(wiz)
+
+
+def _paint_stroke(wiz, t: int, points, class_id: int, erase: bool,
+ brush: float) -> int:
+ """Rasterise ONE stroke into the label store. Returns pixels changed.
+
+ Shared by the on-plot brush (:func:`_on_stroke`) and the renderer's
+ ``seg_paint``, deliberately: two rasterisers would let the brush and the
+ eraser — or the widget and a scripted stroke — disagree about which pixels a
+ given path covers, and that divergence is invisible until a trained model
+ behaves oddly.
+
+ *points* are ``(y, x)`` in **image pixels**, no scale or offset applied
+ (anyplotlib 2-D widgets report pixels).
+ """
+ store = wiz.label_store()
+ before = int(sum(store.counts().values())) if hasattr(store, "counts") else -1
+ try:
+ if erase:
+ # No `erase_stroke` on LabelStore, and the eraser must cover exactly
+ # what the brush would paint. Rasterise the stroke into a scratch
+ # store with the same geometry and erase by the indices it produced,
+ # so the two can never drift apart.
+ from spyde.particles import LabelStore, ScribbleClass
+ scratch = LabelStore(frame_shape=store.frame_shape,
+ classes=[ScribbleClass(0, "scratch")])
+ scratch.paint_stroke(0, points, 0, brush=brush)
+ store.erase(t, scratch.at(0)[0])
+ else:
+ store.paint_stroke(t, points, int(class_id), brush=brush)
+ except (KeyError, ValueError) as exc:
+ emit_error(f"Segment Particles: {exc}")
+ return 0
+ after = int(sum(store.counts().values())) if hasattr(store, "counts") else -1
+ # -1 when the store cannot report counts: assume it painted rather than
+ # swallowing a stroke the user definitely made.
+ return max(1, abs(after - before)) if before >= 0 else 1
+
+
+# ── the on-plot brush ────────────────────────────────────────────────────────
+#
+# The scribble engine needs strokes, and the ONLY thing that can produce them is
+# an anyplotlib brush widget living on the signal plot. Two things were wrong
+# before this existed, and both made painting impossible rather than merely
+# awkward:
+#
+# 1. Nothing ever called ``add_brush_widget``, so there was no brush on the
+# plot at all — Shift+drag had nothing to hit.
+# 2. The caret listened for a RENDERER-side ``spyde:figure_event`` carrying a
+# points array. A brush stroke does not travel that way: anyplotlib emits it
+# to PYTHON (``event_json`` → ``Figure._dispatch_event`` →
+# ``Widget._update_from_js`` → ``plot.callbacks.fire``), which the renderer
+# never sees. So even with a brush present, no stroke could have arrived.
+#
+# The widget is therefore created, owned and read HERE. The renderer's job shrinks
+# to telling us which class is active and how fat the brush is; ``seg_paint``
+# survives only as the programmatic/test door.
+_BRUSH_COLORS = ("#f9a03f", "#89b4fa", "#585b70", "#f38ba8", "#a6e3a1", "#cba6f7")
+
+
+def _brush_supported() -> bool:
+ """Whether the installed anyplotlib has the brush widget.
+
+ It landed in 0.5.0 (CSSFrancis/anyplotlib#47); SpyDE's floor is still 0.4.2,
+ so a user on PyPI's 0.4.2 has no brush and must be TOLD that rather than left
+ dragging at an image that never responds.
+ """
+ try:
+ from anyplotlib.plot2d._plot2d import Plot2D
+ return hasattr(Plot2D, "add_brush_widget")
+ except Exception:
+ return False
+
+
+def _attach_brush(wiz) -> bool:
+ """Put a brush on the source plot and wire its strokes into the label store.
+
+ Idempotent: an existing brush is re-shown and re-armed rather than duplicated,
+ so switching tabs or re-tuning never stacks two brushes on one plot.
+
+ Returns True when a brush is live.
+ """
+ if not _brush_supported():
+ return False
+ src = wiz.src_plot
+ plot2d = getattr(src, "_plot2d", None) if src is not None else None
+ if plot2d is None:
+ return False
+
+ existing = getattr(wiz.tree, "_seg_brush", None)
+ if existing is not None:
+ try:
+ existing.set(active=True, visible=True)
+ return True
+ except Exception as exc:
+ log.debug("[seg] re-arming the existing brush failed: %s", exc)
+
+ try:
+ classes = wiz.class_report()
+ brush = plot2d.add_brush_widget(
+ radius=float(wiz.params.get("brush", 6.0)),
+ colors=[c.get('colour', '#f9a03f') for c in classes] or list(_BRUSH_COLORS),
+ class_id=int(wiz.params.get("active_class", 0)),
+ alpha=0.55,
+ active=True,
+ )
+ except Exception as exc:
+ log.debug("[seg] add_brush_widget failed: %s", exc)
+ return False
+
+ from spyde.drawing.selectors.base_selector import event_handler_fn
+
+ # pointer_up ONLY. The brush deliberately emits once per finished stroke
+ # rather than per pointer_move — a growing stroke re-serialised every frame is
+ # quadratic over one drag (see anyplotlib's BrushWidget docs) — so there is
+ # nothing to listen for mid-stroke and listening would fire on other widgets.
+ handler = event_handler_fn(lambda event: _on_stroke(wiz, event))
+ try:
+ brush.add_event_handler(handler, "pointer_up")
+ except Exception as exc:
+ log.debug("[seg] wiring the brush handler failed: %s", exc)
+ return False
+
+ wiz.tree._seg_brush = brush
+ wiz.tree._seg_brush_handler = handler # weak callbacks: keep a hard ref
+ wiz._brush_seen = 0
+ return True
+
+
+def _sync_brush(wiz) -> None:
+ """Push the caret's paint state onto the live brush widget.
+
+ The widget tags each stroke with its OWN ``class_id`` at paint time in JS, so
+ a class change that only reaches ``wiz.params`` paints the previous colour
+ forever — which is precisely what "I can only scribble one colour" was.
+ Likewise ``erase``: the eraser is a widget mode, not something the handler can
+ decide after the fact, because the stroke has already been tagged.
+
+ No-op when there is no brush (Classical, or an anyplotlib without one).
+ """
+ brush = getattr(wiz.tree, "_seg_brush", None)
+ if brush is None:
+ return
+ state = (int(wiz.params.get("active_class", 0)),
+ float(wiz.params.get("brush", 3.0)),
+ bool(wiz.params.get("erase", False)))
+ if state == getattr(wiz, "_brush_state", None):
+ return
+ wiz._brush_state = state
+ try:
+ brush.set(class_id=state[0], radius=state[1], erase=state[2])
+ except Exception as exc:
+ log.debug("[seg] syncing brush state failed: %s", exc)
+ return
+ _force_widget_push(wiz)
+
+
+def _force_widget_push(wiz) -> None:
+ """Make the brush's new state AUTHORITATIVE in the panel JSON.
+
+ ``Widget.set`` reaches JS through ``Figure._push_widget``, which writes
+ ``event_json`` **only** — deliberately, because re-serialising a whole panel
+ per drag frame is the cost that path exists to avoid. The documented
+ consequence is that ``panel__json`` keeps stale widget state between
+ plot-level pushes, and the brush's *drawing* reads exactly that: JS takes
+ ``w.class_id`` from ``p.state.overlay_widgets`` at stroke start
+ (``figure_esm.js::_brushLiveBegin``) and paints in ``colors[class_id]``.
+
+ So a class switch that only goes through ``set()`` leaves every stroke
+ painting in the PREVIOUS class's colour — reported twice, as "I can only
+ scribble one colour" and then "support film still doesn't change the
+ painting colour". The strokes were landing in the right class the whole
+ time (the caret's per-class counts prove it); only the colour was stale,
+ which makes it look like nothing switched.
+
+ A full ``_push`` is the fix, and it is gated on the brush state having
+ actually CHANGED (``_sync_brush``'s ``_brush_state`` compare) because
+ ``seg_tune`` also fires for every sensitivity-slider tick — re-serialising a
+ 4096² panel, image bytes and all, on each one would trade a colour bug for a
+ much worse drag. Painting itself still goes through the cheap targeted path
+ untouched; this costs one push per class/eraser/size click.
+ """
+ plot2d = getattr(wiz.src_plot, "_plot2d", None)
+ push = getattr(plot2d, "_push", None)
+ if push is None:
+ return
+ try:
+ push()
+ except Exception as exc:
+ log.debug("[seg] forcing the panel push failed: %s", exc)
+
+
+def _arm_brush(wiz) -> None:
+ """Attach the brush and TELL THE USER how to use it.
+
+ The instruction is not optional. A brush that arms silently on Shift+drag is
+ undiscoverable — there is no cursor change to find it by and no affordance on
+ the image, so the honest outcome is a user dragging at a picture that never
+ responds. That is exactly what happened on first use.
+ """
+ if _attach_brush(wiz):
+ emit_status("Shift+drag on the image to paint labels · plain drag still "
+ "pans · pick the class and brush size on the strip beside "
+ "the plot")
+ _emit(wiz, {"type": "seg_brush", "available": True,
+ "hint": "Shift+drag to paint"})
+ return
+ # No brush: say WHY, with the actionable part. Silence here reads as a bug.
+ if not _brush_supported():
+ import anyplotlib as apl
+ emit_error(
+ f"Painting needs anyplotlib 0.5.0 or newer for the brush widget; "
+ f"this environment has {getattr(apl, '__version__', 'unknown')}. "
+ "Until then, use the Classical engine, or install the brush build.")
+ else:
+ emit_error("Painting could not attach a brush to this plot.")
+ _emit(wiz, {"type": "seg_brush", "available": False,
+ "hint": "brush unavailable — see the log"})
+
+
+def _detach_brush(wiz) -> None:
+ """Remove the brush. Called from ``seg_close`` and on leaving Scribble.
+
+ Best-effort and idempotent — teardown must never raise, or closing the caret
+ leaves the wizard half-torn-down.
+ """
+ brush = getattr(wiz.tree, "_seg_brush", None)
+ if brush is not None:
+ for attempt in ("remove", "hide"):
+ fn = getattr(brush, attempt, None)
+ if fn is None:
+ continue
+ try:
+ fn()
+ break
+ except Exception as exc:
+ log.debug("[seg] brush %s() failed: %s", attempt, exc)
+ for attr in ("_seg_brush", "_seg_brush_handler"):
+ if hasattr(wiz.tree, attr):
+ try:
+ setattr(wiz.tree, attr, None)
+ except Exception as exc: # pragma: no cover
+ log.debug("[seg] clearing %s failed: %s", attr, exc)
+
+
+def _on_stroke(wiz, event) -> None:
+ """A finished brush stroke → label pixels → re-preview.
+
+ Reads the widget rather than the event payload. ``Widget._update_from_js``
+ has already merged the JS fields into ``_data`` by the time a plot-level
+ handler runs, so ``brush.strokes`` is authoritative and the event's own copy
+ is redundant.
+
+ Only strokes we have not consumed are applied: the widget accumulates them
+ for as long as it lives, so replaying the whole list on every stroke would
+ re-paint everything and make the pixel counts grow quadratically.
+ """
+ if wiz._closed:
+ return
+ brush = getattr(wiz.tree, "_seg_brush", None)
+ if brush is None:
+ return
+ try:
+ strokes = list(getattr(brush, "strokes", ()) or ())
+ except Exception as exc:
+ log.debug("[seg] reading brush strokes failed: %s", exc)
+ return
+
+ seen = int(getattr(wiz, "_brush_seen", 0))
+ fresh = strokes[seen:]
+ if not fresh:
+ return
+ wiz._brush_seen = len(strokes)
+
+ # PYTHON is the authority for all three, and the class one is load-bearing.
+ #
+ # `stroke_classes` is what the JS widget tagged the stroke with, and reading
+ # it made class switching silently not work: `Figure._push_widget` sends a
+ # targeted update that never writes `panel__json`, so a Python-side
+ # `brush.class_id = 1` does not reliably reach the widget before the next
+ # stroke — anyplotlib's own brush test has to press Shift BEFORE pushing the
+ # class to dodge exactly this.
+ #
+ # The natural experiment that proved it: `erase` read from `wiz.params` and
+ # WORKED, `class` read from `stroke_classes` and did not, in the same handler
+ # on the same stroke. So the strip → seg_tune → params path is sound; the
+ # Python → JS widget push is what is unreliable.
+ #
+ # Nothing is lost by preferring params: a stroke cannot change class midway,
+ # so the active class when it completes IS its class. The widget push stays
+ # (see `_sync_brush`) purely so the stroke DRAWS in the right colour while
+ # the user paints — cosmetic, and no longer load-bearing.
+ erase = bool(wiz.params.get("erase", False))
+ radius = float(getattr(brush, "radius", wiz.params.get("brush", 6.0)) or 6.0)
+ active_cls = int(wiz.params.get("active_class", 0))
+
+ painted = 0
+ for i, stroke in enumerate(fresh, start=seen):
+ cls = active_cls
+ # anyplotlib gives [[x, y], ...] in IMAGE PIXELS; LabelStore works in
+ # (y, x). Swapping these silently mirrors every scribble about the
+ # diagonal, which on a non-square frame also puts half of them outside
+ # the image — see spyde/actions/masks.py for this bug class.
+ pts = [[float(p[1]), float(p[0])] for p in (stroke or ()) if len(p) >= 2]
+ if pts:
+ painted += _paint_stroke(wiz, wiz.frame_index(), pts, cls,
+ erase, radius)
+
+ if painted:
+ _emit_state(wiz)
+ _preview(wiz, wiz.guard())
+
+
+def seg_train(session, plot, payload) -> None:
+ """Fit the scribble classifier on every accumulated label.
+
+ Trains on labelled PIXELS only (thousands, not millions) — plan B3's hard
+ interaction budget. The store is snapshotted before it crosses onto the
+ worker: ``seg_paint`` runs on the main thread and a stroke landing mid-fit
+ would mutate the arrays the fit is iterating.
+ """
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ emit_error("Segment Particles: the caret is not open")
+ return
+ if wiz.labels is None or len(wiz.labels) == 0:
+ emit_error("Segment Particles: nothing painted yet — scribble on a "
+ "particle and on the background first.")
+ return
+
+ from spyde.particles import LabelStore
+ snapshot = LabelStore.from_dict(wiz.labels.to_dict())
+ device = (payload or {}).get("device")
+ gen = wiz.guard()
+ emit_status("Segment Particles: training…")
+
+ def _work():
+ from spyde.particles import ScribbleClassifier
+ _n, get_frame, _shape = wiz.frames()
+ clf = ScribbleClassifier(device=device)
+ report = clf.fit(snapshot, get_frame,
+ progress=lambda d, n: emit_progress(d, n, "Training"))
+ return clf, report
+
+ def _done(res):
+ clf, report = res
+ if not wiz.still(gen) or wiz._closed:
+ return
+ wiz.classifier = clf
+ wiz.params["method"] = "scribble"
+ _emit(wiz, {"type": "seg_trained", "report": report})
+ # Say which split route the training just selected. A painted boundary
+ # class is what lets `split_instances` skip the distance transform and
+ # the watershed (1.78 s -> 0.33 s at 4096²), and the user is the one who
+ # decides it by painting — so it has to be visible that they did.
+ route = ("boundary class painted — touching particles split by "
+ "connected components, no watershed"
+ if report.get("has_boundary") else
+ "no boundary painted — touching particles split by watershed")
+ emit_status(
+ f"Segment Particles: trained on {report['n_pixels']} px across "
+ f"{report['n_classes']} classes "
+ f"(accuracy {report['train_accuracy']:.3f}); {route}")
+ _emit_state(wiz)
+ _preview(wiz, gen)
+
+ def _fail(exc):
+ emit_error(f"Segment Particles: training failed — {exc}")
+
+ run_on_worker(session, _work, name="seg-train", on_done=_done, on_error=_fail)
+
+
+# ── the single-frame result ──────────────────────────────────────────────────
+
+def commit_single_frame(session, wiz: SegmentWizard, labels, rows, contours,
+ frame: int):
+ """Commit one frame's segmentation as a label-image tree with the store
+ attached. Shared by ``seg_commit`` and by ``seg_run`` on a 2-D image.
+
+ See :meth:`SegmentWizard.commit` for why this door and not
+ ``open_particle_tree``.
+ """
+ from spyde.actions.commit import commit_result_tree
+ from spyde.signals.particles import COL, SpyDEParticles
+
+ scale, units = wiz.scale_units()
+ _n, _get, shape = wiz.frames()
+ rows = np.asarray(rows, np.float32)
+ if len(rows):
+ # The store is one frame long, so every row's `t` must be 0 or the CSR
+ # offsets disagree with the rows about which frame they live in.
+ rows = rows.copy()
+ rows[:, COL["t"]] = 0.0
+ params = dict(wiz.params, frame=int(frame), mode="single_frame")
+ parts = SpyDEParticles.from_frames(
+ [rows], frame_shape=shape,
+ contours_per_frame=([list(contours)] if wiz.params["store_masks"] else None),
+ scale=scale, units=units, params=params,
+ provenance={"action": "segment_particles", "params": dict(wiz.params)},
+ )
+ tree = commit_result_tree(
+ session, title=f"Particles — frame {int(frame)} ({len(rows)})",
+ primary=np.asarray(labels, np.float32), primary_label="labels",
+ levels=None, cmap="gray",
+ attrs={"particles": parts, "source_node": wiz.signal(),
+ "source_tree": wiz.tree, "particle_events": [],
+ "nav_map": np.zeros(1, np.int64)},
+ provenance={"action": "segment_particles", "params": params,
+ "frame": int(frame)},
+ )
+ _rebuild_toolbars(tree)
+ emit_status(f"Committed {len(rows)} particles from frame {int(frame)}")
+ return tree
+
+
+# ── the batch run ────────────────────────────────────────────────────────────
+
+def seg_run(session, plot, payload) -> None:
+ """Segment every frame on a worker: progressive, cancellable.
+
+ The result window opens IMMEDIATELY with an empty particle store, its count
+ trace fills as frames complete, and ``tree.particles`` attaches only at
+ ``_finalize`` — see the module docstring for why that ordering is the whole
+ point of the attach gap.
+ """
+ src, tree = _src_plot_tree(session, plot)
+ if src is None or tree is None:
+ emit_error("Segment Particles: no active dataset")
+ return
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ wiz = SegmentWizard(session, tree, src)
+ tree._seg_wizard = wiz
+ wiz.set_params(payload, merge=True)
+ p = dict(wiz.params)
+
+ engine = _engine(wiz, p)
+ if engine is None:
+ return
+ try:
+ n_frames, get_frame, shape = wiz.frames()
+ except TypeError as exc:
+ emit_error(f"Segment Particles: {exc}")
+ return
+ scale, units = wiz.scale_units()
+ source = wiz.signal()
+
+ if n_frames < 2:
+ # Plan §0.10's single-image shape: there is no movie to fill
+ # progressively, and a 1-frame particle tree is not constructible (see
+ # SegmentWizard.commit). Segment it and commit the label image.
+ _run_single_frame(session, wiz, get_frame, engine, scale)
+ return
+
+ from spyde.actions.particle_tree import open_particle_tree
+ from spyde.signals.particles import N_COLUMNS, SpyDEParticles
+
+ empty = [np.zeros((0, N_COLUMNS), np.float32)] * n_frames
+ placeholder = SpyDEParticles.from_frames(
+ empty, frame_shape=shape, scale=scale, units=units, params=dict(p))
+ # attach=False: `requires_particles` must not unlock the particle toolbar
+ # against an empty store, so the tree publishes `particles` only at
+ # _finalize. The placeholder is still handed in because the lazy label movie
+ # renders from THAT object — see open_particle_tree's `attach` docs — which
+ # is why _finalize mutates it in place rather than swapping in a new store.
+ result = open_particle_tree(
+ session, particles=placeholder, source_node=source, source_tree=tree,
+ params=dict(p), title=f"Particles — {n_frames} frames", attach=False)
+ wiz.result_tree = result
+
+ # RAISE THE "Calculating…" OVERLAY NOW — synchronously, the statement after
+ # the window exists and BEFORE any of the setup below.
+ #
+ # It used to be raised nowhere at all, and the natural place to add it (the
+ # worker, next to the first progress emission) is far too late: the window
+ # opens, then the placeholder store is built, the cancel flags registered,
+ # the generation bumped, the worker scheduled, the thread hop paid, and the
+ # first frame COMPUTED — on a 4096² frame that is seconds of a window that
+ # looks finished and empty. Every one of those steps is between the window
+ # appearing and the user learning it is working, which is the lag reported.
+ computing = window_computing(_result_window_id(result))
+ computing.start()
+
+ result._seg_batch_running = True
+ tree._seg_batch_running = True
+
+ stopped = [False]
+ for t_ in {id(tree): tree, id(result): result}.values():
+ if hasattr(t_, "register_cancel"):
+ t_.register_cancel(flag=stopped)
+
+ emit_status(f"Segmenting {n_frames} frames…")
+ gen = bump_generation(result, "_seg_batch_gen")
+
+ counts = np.zeros(n_frames, np.float32)
+ last_paint = [0.0]
+
+ def _paint_counts():
+ if not is_current(result, "_seg_batch_gen", gen):
+ return
+ _paint_count_trace(result, counts.copy())
+
+ def _work():
+ # The batch fans out over the cluster (spyde.particles.batch): one dask
+ # task per block of frames, dual-lane so the GPU workers run the torch
+ # head while the rest run the CPU path in parallel. It used to be a
+ # plain serial for-loop on this one thread — 90 minutes for 900 frames
+ # of 4096², with 47 cores and 78% of the GPU idle.
+ from spyde.particles.batch import (EngineSpec, drop_engine_model,
+ save_engine_model, segment_movie)
+ done = [0]
+ # Frames finished so far, indexed by frame so an out-of-order block
+ # lands in the right slot; the unfinished ones stay empty, which renders
+ # as "nothing found there yet" rather than as a shorter movie.
+ live_rows = [np.zeros((0, N_COLUMNS), np.float32)] * n_frames
+ live_contours: list = [[] for _ in range(n_frames)]
+ model_path = None
+ if p["method"] == "scribble":
+ # The trained head crosses to the workers as a FILE, never as a
+ # pickled CUDA tensor — see EngineSpec.
+ model_path = save_engine_model(wiz.classifier)
+ spec = EngineSpec(method=p["method"], params=_segment_kwargs(p),
+ model_path=model_path)
+
+ def _on_frames(t0, t1, vals):
+ """One block landed. Runs on a dask/worker callback thread, so it
+ only touches `counts` and marshals the paint (CLAUDE.md threading)."""
+ for i, (rows, cs) in enumerate(vals):
+ counts[t0 + i] = float(len(rows))
+ # Keep the finished frames so the label movie can render them
+ # WHILE the rest compute. Blocks land out of order under the
+ # fan-out, so index by frame rather than appending.
+ live_rows[t0 + i] = rows
+ live_contours[t0 + i] = cs
+ done[0] += int(t1 - t0)
+ now = time.monotonic()
+ if now - last_paint[0] >= _PROGRESS_INTERVAL:
+ last_paint[0] = now
+ emit_progress(done[0], n_frames, "Segmenting")
+ _to_main(session, _paint_counts)
+ # The navigator's count trace shows that SOMETHING is happening;
+ # the signal shows whether it is happening CORRECTLY. That is
+ # the difference between abandoning a bad 900-frame run in the
+ # first ten seconds and after it finishes. Snapshot the lists —
+ # the callback thread keeps writing into them.
+ snap = (list(live_rows), list(live_contours), int(t1) - 1)
+ _to_main(session, lambda s=snap: _publish_partial(
+ session, result, placeholder, s[0], s[1],
+ p, scale, units, s[2]))
+
+ try:
+ # Frames never reached keep an empty block, so the CSR store always
+ # spans the movie and a cancelled run reads as "no particles after
+ # frame N" rather than a shorter movie — segment_movie guarantees
+ # both lists are n_frames long.
+ per_frame, contours, n_done = segment_movie(
+ _movie_array(source), spec, n_frames=n_frames,
+ get_frame=get_frame, scale=scale,
+ store_masks=bool(p["store_masks"]),
+ client=_batch_client(session, stopped), stopped=stopped,
+ on_frames=_on_frames)
+ finally:
+ drop_engine_model(model_path)
+ return per_frame, contours, n_done
+
+ def _done(res):
+ per_frame, contours, done = res
+ try:
+ if getattr(result, "_spyde_closed", False):
+ return # window torn down mid-run
+ _finalize(session, result, placeholder, per_frame, contours,
+ p, scale, units, done, n_frames, cancelled=stopped[0])
+ finally:
+ _teardown_batch(tree, result, stopped, computing)
+
+ def _fail(exc):
+ emit_error(f"Segment Particles failed: {exc}")
+ log.exception("segmentation batch failed")
+ _teardown_batch(tree, result, stopped, computing)
+
+ run_on_worker(session, _work, name="seg-batch", on_done=_done, on_error=_fail)
+
+
+def _run_single_frame(session, wiz, get_frame, engine, scale) -> None:
+ """Segment the only frame there is, on a worker, then commit it."""
+ def _work():
+ from spyde.particles import measure_frame
+ frame = np.asarray(get_frame(0))
+ labels = engine(frame)
+ rows, contours = measure_frame(labels, frame, t=0, scale=scale)
+ return labels, rows, contours
+
+ def _done(res):
+ labels, rows, contours = res
+ if wiz._closed:
+ return
+ commit_single_frame(session, wiz, labels, rows, contours, 0)
+
+ run_on_worker(session, _work, name="seg-single",
+ on_done=_done,
+ on_error=lambda e: emit_error(f"Segment Particles failed: {e}"))
+
+
+def _teardown_batch(tree, result, stopped, computing=None) -> None:
+ # The Calculating chip comes down HERE, the one place every exit path goes
+ # through — success, failure and cancellation alike. Clearing it only on the
+ # success path is how an overlay ends up spinning forever over a window that
+ # stopped working, which is the failure `window_computing`'s pairing
+ # contract exists to prevent.
+ if computing is not None:
+ try:
+ computing.stop()
+ except Exception as exc:
+ log.debug("[seg] clearing the computing overlay failed: %s", exc)
+ result._seg_batch_running = False
+ tree._seg_batch_running = False
+ for t_ in {id(tree): tree, id(result): result}.values():
+ if hasattr(t_, "unregister_cancel"):
+ try:
+ t_.unregister_cancel(flag=stopped)
+ except Exception as exc:
+ log.debug("[seg] unregister_cancel failed: %s", exc)
+
+
+def _finalize(session, result, placeholder, per_frame, contours, p,
+ scale, units, done, n_frames, *, cancelled: bool) -> None:
+ """Attach the finished store, repaint, and unlock the particle actions."""
+ from spyde.actions.particle_tree import _navigator_traces
+ from spyde.signals.particles import SpyDEParticles
+
+ params = dict(p)
+ if cancelled:
+ # Recorded, not just narrated: a partial result that only says so in a
+ # transient status line is indistinguishable from a complete one later.
+ params["cancelled_after_frame"] = int(done)
+ final = SpyDEParticles.from_frames(
+ per_frame, frame_shape=placeholder.frame_shape,
+ contours_per_frame=(contours if p["store_masks"] else None),
+ scale=scale, units=units, params=params,
+ provenance={"action": "segment_particles", "params": dict(p)},
+ )
+
+ events = []
+ if p["track"] and final.n_particles:
+ try:
+ from spyde.particles import link
+ res = link(final, max_dist=float(p["max_dist"]), apply=True)
+ events = list(res.events)
+ except Exception as exc:
+ log.debug("[seg] linking failed, keeping untracked particles: %s", exc)
+
+ _adopt(placeholder, final)
+ result.particles = placeholder
+ result.particle_events = events
+ result.nav_traces = _navigator_traces(placeholder, events or None)
+
+ # The signal plot's CachedDaskArray captured the placeholder's zeros when
+ # the window first rendered; drop it so the label movie re-slices the real
+ # contours (the same stale-cache fix find_vectors_action._finalize makes).
+ try:
+ result.root.cached_dask_array = None
+ result.root._clear_cache_dask_data()
+ except Exception as exc:
+ log.debug("[seg] clearing stale cached dask array failed: %s", exc)
+
+ _paint_count_trace(result, placeholder.count_series())
+ _repaint_label_movie(result)
+ _rebuild_toolbars(result)
+
+ # TERMINAL progress. Without it `state.loading.busy` never clears in the
+ # renderer: the StatusBar spinner spins forever and — because it prefers
+ # `loading.text` while busy (StatusBar.tsx) — "Segmenting (33%)" permanently
+ # masks the "Found N particles" line emitted immediately below. Found by
+ # driving the real UI; no headless test could see it.
+ emit_progress(n_frames, n_frames, "Segmenting")
+
+ n = placeholder.n_particles
+ if cancelled:
+ # "{done} of {n_frames}", NOT "the first {done}": under the block
+ # fan-out the finished frames are not a contiguous prefix — blocks land
+ # out of order, so a cancelled run has holes rather than a clean cut.
+ # Saying "first" would misdescribe which frames actually have particles.
+ emit_status(f"Segmentation cancelled — found {n} particles in "
+ f"{done} of {n_frames} frames")
+ else:
+ emit_status(f"Found {n} particles in {n_frames} frames")
+
+
+def _repaint_label_movie(result) -> None:
+ """Push the finalized frame to the label-movie window.
+
+ Clearing the stale cached dask array (above) makes the NEXT read correct, but
+ nothing triggers a read — so the window keeps showing the placeholder's zeros
+ and the result looks empty until the user happens to scrub. Re-read the
+ currently-displayed frame and paint it.
+
+ Best-effort: a failure here costs a stale frame until the next scrub, which is
+ exactly the state we were in before, so it must never take the finalize down
+ with it.
+ """
+ from spyde.actions.lifecycle import paint_signal_plots
+
+ try:
+ particles = result.particles
+ if particles is None or particles.n_frames == 0:
+ return
+ t = 0
+ for plot in getattr(result, "plots", None) or ():
+ idx = getattr(plot, "current_indices", None)
+ if idx:
+ t = int(np.atleast_1d(idx)[0])
+ break
+ t = max(0, min(t, particles.n_frames - 1))
+ paint_signal_plots(result, particles.render_frame(t, value="track"))
+ except Exception as exc:
+ log.debug("[seg] repainting the label movie after finalize failed: %s", exc)
+
+
+def _adopt(placeholder, final) -> None:
+ """Move *final*'s arrays onto *placeholder*, in place.
+
+ The lazy label movie built by ``open_particle_tree`` closes over the
+ placeholder OBJECT and calls ``render_frame`` at graph-execution time, so
+ swapping ``tree.particles`` for a new store would leave the movie rendering
+ zeros forever. Mutating is what makes the early window fill in.
+ """
+ placeholder.flat_buffer = final.flat_buffer
+ placeholder.t_offsets = final.t_offsets
+ placeholder.contours = final.contours
+ placeholder.contour_offsets = final.contour_offsets
+ placeholder.params = final.params
+ placeholder.provenance = final.provenance
+
+
+def _to_main(session, fn) -> None:
+ """Run *fn* on the asyncio main thread; inline when there is no loop (bare
+ handler tests) — the same fallback ``lifecycle.run_on_worker`` makes."""
+ dispatch = getattr(session, "_dispatch_to_main", None)
+ if dispatch is None:
+ fn()
+ return
+ dispatch(fn)
+
+
+def _result_window_id(tree):
+ """The result tree's SIGNAL window — what the Calculating chip sits on.
+
+ The signal plot and not the navigator: the navigator fills in visibly as the
+ count trace grows, so it is self-evidently working; the signal plot is the
+ one that sits there looking empty and finished.
+ """
+ for plot in (getattr(tree, "signal_plots", None) or []):
+ wid = getattr(plot, "window_id", None)
+ if wid is not None:
+ return wid
+ return None
+
+
+def _publish_partial(session, result, placeholder, per_frame, contours,
+ p, scale, units, t_latest: int) -> None:
+ """Adopt the frames finished SO FAR so the label movie renders them live.
+
+ Same three steps as :func:`_finalize`, minus the tracking and the attach:
+ build a store from what exists, adopt it into the placeholder the lazy label
+ movie closed over, and drop the stale cached dask array so the next slice
+ re-reads. Without the cache drop the movie keeps serving the zeros it
+ captured when the window first rendered, which is exactly the bug
+ ``_finalize`` documents.
+
+ Deliberately NOT attached to ``result.particles``: the particle toolbar is
+ ``requires_particles``-gated and must not unlock against a store that is
+ still filling — that is the attach gap the module docstring is about. This
+ only makes the movie SHOW the work.
+
+ Then paint the MOST RECENTLY COMPUTED frame, not whatever frame the
+ navigator happens to sit on — "show the newest result" is the whole point,
+ and on a 900-frame run the navigator is parked on frame 0 the entire time,
+ so painting its frame would show one image for twenty minutes.
+
+ Cheap enough to run on the feedback clock (a few times a second): the rows
+ are small float arrays and the rebuild is a concatenate, not a recompute.
+ """
+ from spyde.actions.lifecycle import paint_signal_plots
+ from spyde.signals.particles import SpyDEParticles
+
+ # EVERYTHING here is inside the guard, deliberately. This runs on the
+ # asyncio MAIN thread (marshalled from the batch's callback), so an
+ # exception does not merely lose one preview frame — it escapes into the
+ # event loop. A cosmetic live fill must never be able to damage the run it
+ # is describing, and the correct result is still produced by `_finalize`
+ # whatever happens here.
+ try:
+ if getattr(result, "_spyde_closed", False):
+ return # window torn down mid-run
+ partial = SpyDEParticles.from_frames(
+ per_frame, frame_shape=placeholder.frame_shape,
+ contours_per_frame=(contours if p["store_masks"] else None),
+ scale=scale, units=units, params=dict(p))
+ _adopt(placeholder, partial)
+ try:
+ result.root.cached_dask_array = None
+ result.root._clear_cache_dask_data()
+ except Exception as exc:
+ log.debug("[seg] clearing the cache for the live fill failed: %s", exc)
+ t = max(0, min(int(t_latest), placeholder.n_frames - 1))
+ paint_signal_plots(result, placeholder.render_frame(t, value="track"))
+ except Exception as exc:
+ log.debug("[seg] live label-movie fill failed: %s", exc)
+
+
+def _nav_plots(tree) -> list:
+ npm = getattr(tree, "navigator_plot_manager", None)
+ if npm is None:
+ return []
+ out = []
+ for pw in list(npm.plot_windows.keys()):
+ out.extend(npm.plots.get(pw, []))
+ return out
+
+
+def _paint_count_trace(tree, counts: np.ndarray) -> None:
+ """Paint particle-count-vs-time onto the result tree's 1-D navigator.
+
+ Shape-matched rather than "the first navigator": a tree can carry more than
+ one navigator plot, and painting an (n_frames,) trace onto the wrong one
+ leaves it blank — the bug find_vectors_action._spatial_nav_plot exists for.
+ """
+ want = int(counts.shape[0])
+ for nav in _nav_plots(tree):
+ cur = getattr(nav, "current_data", None)
+ if cur is None or getattr(cur, "ndim", 0) != 1 or int(cur.shape[0]) != want:
+ continue
+ try:
+ nav.needs_auto_level = True
+ nav.set_data(np.asarray(counts, np.float32))
+ except Exception as exc:
+ log.debug("[seg] painting count trace failed: %s", exc)
+
+
+def _rebuild_toolbars(tree) -> None:
+ """Re-send the toolbar config so ``requires_particles`` actions appear.
+
+ This is the moment the gate flips — without it the buttons stay hidden until
+ something else happens to rebuild the toolbar, and the e2e specs wait on
+ exactly this appearing.
+ """
+ for sp in list(getattr(tree, "signal_plots", []) or []):
+ try:
+ state = getattr(sp, "plot_state", None)
+ if state is not None and hasattr(state, "_send_toolbar_config"):
+ state._send_toolbar_config()
+ except Exception as exc:
+ log.debug("[seg] re-sending toolbar config failed: %s", exc)
+
+
+def seg_commit(session, plot, payload=None) -> None:
+ """Commit the previewed frame as a one-frame particle tree."""
+ wiz = _wizard(session, plot)
+ if wiz is None:
+ emit_error("Segment Particles: nothing to commit")
+ return
+ wiz.commit()
+
+
+def segment_particles(ctx, action_name: str = "Segment Particles", **params):
+ """Toolbar entry — a no-op parent; the Electron toolbar opens the staged
+ caret, which drives the ``seg_*`` handlers (README §4)."""
+ return None
diff --git a/spyde/actions/registry.py b/spyde/actions/registry.py
index 7b86eb33..3eeb967c 100644
--- a/spyde/actions/registry.py
+++ b/spyde/actions/registry.py
@@ -105,6 +105,33 @@
"crop_open": "spyde.actions.base.crop_open",
"crop_close": "spyde.actions.base.crop_close",
"crop_set_region": "spyde.actions.base.crop_set_region",
+ # Segment Particles (spyde/actions/particles_action.py) — plan B7.
+ "seg_open": "spyde.actions.particles_action.seg_open",
+ "seg_close": "spyde.actions.particles_action.seg_close",
+ "seg_set_method": "spyde.actions.particles_action.seg_set_method",
+ "seg_tune": "spyde.actions.particles_action.seg_tune",
+ "seg_paint": "spyde.actions.particles_action.seg_paint",
+ "seg_train": "spyde.actions.particles_action.seg_train",
+ "seg_run": "spyde.actions.particles_action.seg_run",
+ "seg_commit": "spyde.actions.particles_action.seg_commit",
+ # Particle overlay + editing (spyde/actions/particle_overlay.py) — plan B9/C2/C3.
+ "part_open": "spyde.actions.particle_overlay.part_open",
+ "part_close": "spyde.actions.particle_overlay.part_close",
+ "part_tune": "spyde.actions.particle_overlay.part_tune",
+ "part_select": "spyde.actions.particle_overlay.part_select",
+ "part_region_mode": "spyde.actions.particle_overlay.part_region_mode",
+ "part_delete": "spyde.actions.particle_overlay.part_delete",
+ "part_merge": "spyde.actions.particle_overlay.part_merge",
+ "part_split": "spyde.actions.particle_overlay.part_split",
+ "part_lanes": "spyde.actions.particle_overlay.part_lanes",
+ # Drift Correction (spyde/actions/drift_action.py) — plan A8.
+ "drift_open": "spyde.actions.drift_action.drift_open",
+ "drift_close": "spyde.actions.drift_action.drift_close",
+ "drift_set_method": "spyde.actions.drift_action.drift_set_method",
+ "drift_tune": "spyde.actions.drift_action.drift_tune",
+ "drift_run": "spyde.actions.drift_action.drift_run",
+ "drift_discard": "spyde.actions.drift_action.drift_discard",
+ "drift_commit": "spyde.actions.drift_action.drift_commit",
"download_cancel": "spyde.backend.example_download.download_cancel",
"compute_configure": "spyde.backend.compute_config.compute_configure",
"set_log_level": "spyde.backend.log_stream.set_log_level",
@@ -235,6 +262,9 @@ def register_staged(name: str, dotted_path: str) -> None:
"vom": ("spyde.actions.vector_orientation_om", "VomWizard"),
"ebsd": ("spyde.actions.ebsd_action", "EbsdWizard"),
"czb": ("spyde.actions.center_zero_beam", "PARAMETERS"),
+ "seg": ("spyde.actions.particles_action", "SegmentWizard"),
+ "part": ("spyde.actions.particle_overlay", "PARAMETERS"),
+ "drift": ("spyde.actions.drift_action", "DriftWizard"),
# YAML-declared (resolved from spyde.TOOLBAR_ACTIONS):
"fv": ("__yaml__", "Find Diffraction Vectors"),
"om": ("__yaml__", "Orientation Mapping"),
diff --git a/spyde/actions/wizard.py b/spyde/actions/wizard.py
index b42b39f4..84998cda 100644
--- a/spyde/actions/wizard.py
+++ b/spyde/actions/wizard.py
@@ -74,6 +74,18 @@ def still(self, gen: int) -> bool:
"""True if *gen* is still the current run generation."""
return is_current(self.tree, self._gen_key, gen)
+ def current_gen(self) -> int:
+ """The live generation WITHOUT opening a new one.
+
+ For deferred work that belongs to the interaction ALREADY in progress
+ rather than starting a new one — the case being a re-preview triggered by
+ the navigator moving. Calling :meth:`guard` for that would bump the shared
+ generation and silently cancel an in-flight train or run the user
+ actually asked for, which is a failure with no error message: the worker
+ completes, sees a newer generation, and drops its result.
+ """
+ return int(getattr(self.tree, self._gen_key, 0) or 0)
+
def cancel_inflight(self) -> None:
"""Invalidate any in-flight open (call FIRST in the close handler)."""
bump_generation(self.tree, self._gen_key)
diff --git a/spyde/backend/_session_actions.py b/spyde/backend/_session_actions.py
index c58a4af1..b09eb999 100644
--- a/spyde/backend/_session_actions.py
+++ b/spyde/backend/_session_actions.py
@@ -38,6 +38,7 @@
"load_test_data_si_grains", "load_test_data_sped_ag",
"load_test_data_eels", "load_test_data_eds", "load_test_data_ebsd",
"load_test_data_line", "load_test_data_movie", "load_test_data_5d",
+ "load_test_data_particles",
"test_nav_drag",
"test_region_scrub", "test_add_second_navigator",
"load_test_vectors", "run_test_orientation", "dump_dask_state",
@@ -148,6 +149,8 @@ def dispatch_action(self, msg: dict) -> None:
self._load_test_data_movie(payload)
elif action == "load_test_data_5d":
self._load_test_data_5d(payload)
+ elif action == "load_test_data_particles":
+ self._load_test_data_particles(payload)
elif action == "test_add_second_navigator":
self._test_add_second_navigator()
elif action == "test_nav_drag":
diff --git a/spyde/backend/_session_testharness.py b/spyde/backend/_session_testharness.py
index abc939be..b0478cde 100644
--- a/spyde/backend/_session_testharness.py
+++ b/spyde/backend/_session_testharness.py
@@ -415,6 +415,62 @@ def _load_test_data_movie(self, payload: dict | None = None) -> None:
s.set_signal_type("insitu") # gates the Play/Fast Forward toolbar buttons
self._add_signal(s, source_path="test_data_movie")
+ def _load_test_data_particles(self, payload: dict | None = None) -> None:
+ """The synthetic in-situ PARTICLE movie — the drift/segmentation fixture.
+
+ Wraps :func:`spyde.data.synthetic.particle_movie`, whose ground truth
+ (per-frame drift, particle radii, and the nucleation / dissolution / merge
+ frames) is stamped into ``metadata.Spyde.synthetic`` — so an e2e spec can
+ assert against the numbers the data was built from instead of a golden
+ screenshot. See DRIFT_AND_PARTICLES_PLAN.md step 0.
+
+ Loaded **lazy at one frame per chunk**, like ``load_test_data_movie`` and
+ like a real ``.mrc`` in-situ movie: each nav move is then a small cold read
+ of just that frame, which is the path a user actually drags.
+
+ Payload: ``{"frames": int, "size": [ny, nx], "eager": bool,
+ "noise": float}``. ``eager`` keeps it in RAM, which skips the whole
+ array-cache path — useful only for a test that wants to isolate
+ something else (an eager fixture is how ``si_grains`` silently made a
+ cache benchmark measure nothing).
+
+ ``noise`` exists to make the OVER-SEGMENTATION failure reproducible. At
+ the default 0.015 the fixture is clean and a global threshold works on
+ it, which is why every spec here was green while a real low-contrast
+ in-situ frame produced 14028 instances and a solid overlay. Around 0.3
+ there is no bimodal histogram left for otsu to find, it lands inside the
+ noise, and the split shatters the film exactly as reported.
+ """
+ import dask.array as da
+
+ from spyde.backend.heavy_imports import ensure_heavy_imports
+ from spyde.data.synthetic import particle_movie
+
+ # BEFORE set_signal_type: racing the startup prewarm's pyxem import
+ # partially initialises the module and the cast then silently fails,
+ # taking every gated toolbar action with it.
+ ensure_heavy_imports()
+
+ payload = payload or {}
+ n_frames = int(payload.get("frames", 24))
+ shape = tuple(payload.get("size", (96, 112)))
+
+ kw = {}
+ if payload.get("noise") is not None:
+ kw["noise"] = float(payload["noise"])
+ s = particle_movie(n_frames=n_frames, shape=shape, **kw)
+ if not payload.get("eager"):
+ eager = s.data
+ ny, nx = eager.shape[1:]
+ # `as_lazy()` carries the axes, the signal type AND the stamped ground
+ # truth across (metadata has no setter, so copying it by hand is not an
+ # option), then swap in the correctly-chunked array. Assigning `.data`
+ # rather than calling `.rechunk()` keeps this a graph construction over
+ # an array already in RAM — never a scheduler shuffle.
+ s = s.as_lazy()
+ s.data = da.from_array(eager, chunks=(1, ny, nx))
+ self._add_signal(s, source_path="test_data_particles")
+
@staticmethod
def _write_movie_mrc(frames) -> str:
"""Write ``frames`` as a minimal MRC2014 stack (mode 6 = uint16) into the
diff --git a/spyde/backend/tutorial_data.py b/spyde/backend/tutorial_data.py
index f1001a69..96474b86 100644
--- a/spyde/backend/tutorial_data.py
+++ b/spyde/backend/tutorial_data.py
@@ -213,6 +213,42 @@ def tutorial_movie(self) -> None:
s.metadata.set_item("General.title", "Tutorial: In-Situ Movie")
self._add_signal(s, source_path="tutorial_movie")
+ def tutorial_particles(self) -> None:
+ """Tutorial: particle segmentation, drift correction and tracking.
+
+ The synthetic particle movie from ``spyde.data.synthetic`` — 24 frames of
+ 96x112 with nine particles on a drifting support film, including one
+ nucleation, one dissolution, one merge, one mover and two deliberately
+ faint low-contrast probes.
+
+ **Deliberately tiny, and that is the point.** The classical preview costs
+ roughly 270 ms on a frame this size and **8.4 s on a 4096² one** (measured:
+ `segment_frame` alone is 7.6 s of that, watershed over 16.7 M pixels), so
+ a real in-situ movie makes every sensitivity nudge feel like a hang. Learn
+ the workflow here, where the whole loop is interactive, then take the
+ parameters to the real data.
+
+ It also carries its ground truth in ``metadata.Spyde.synthetic``
+ (per-frame drift, radii, and the nucleation / dissolution / merge frames),
+ so what the tools report can be checked against what the data was built
+ from — read it with ``spyde.data.synthetic.ground_truth``.
+ """
+ import dask.array as da
+
+ from spyde.backend.heavy_imports import ensure_heavy_imports
+ from spyde.data.synthetic import particle_movie
+ ensure_heavy_imports() # don't race the startup prewarm's pyxem import
+
+ eager = particle_movie(n_frames=24)
+ ny, nx = eager.data.shape[1:]
+ # Lazy at one frame per chunk, like tutorial_movie and like a real .mrc:
+ # each nav move is then a small cold read of just that frame. `as_lazy()`
+ # carries the axes, the signal type AND the stamped ground truth across.
+ s = eager.as_lazy()
+ s.data = da.from_array(eager.data, chunks=(1, ny, nx))
+ s.metadata.set_item("General.title", "Tutorial: Particles (small)")
+ self._add_signal(s, source_path="tutorial_particles")
+
# name -> bound-method lookup used by the (ungated) `tutorial_load` action in
# _session_actions.py. Keys are the same names used as tutorial_ testids
@@ -225,4 +261,5 @@ def tutorial_movie(self) -> None:
"strain": TutorialDataMixin.tutorial_strain,
"spectroscopy": TutorialDataMixin.tutorial_spectroscopy,
"movie": TutorialDataMixin.tutorial_movie,
+ "particles": TutorialDataMixin.tutorial_particles,
}
diff --git a/spyde/compute_dispatch.py b/spyde/compute_dispatch.py
index d50e88f7..157f7fa2 100644
--- a/spyde/compute_dispatch.py
+++ b/spyde/compute_dispatch.py
@@ -163,6 +163,7 @@ def dispatch_chunks(
fill_value=np.nan,
stall_timeout_s: float = 600.0,
submit_batch: int = 8,
+ batch_unpinned: bool = True,
label: str = "dispatch",
on_chunk_done=None,
lane_default_mode: str = "one",
@@ -313,7 +314,17 @@ def _band_key(i):
completed_futures: list = [] # held until the end — see module docstring
outstanding = {"gpu": 0, "cpu": 0}
state = {"completed": 0, "error": None, "last_progress": time.time(),
- "lane_done": {"gpu": 0, "cpu": 0}, "mem_hot": False}
+ "lane_done": {"gpu": 0, "cpu": 0}, "mem_hot": False,
+ # The unpinned lane PRIMES with one small batch before sending the
+ # rest in a single submit. Measured on 977 chunks: going straight
+ # to one all-at-once submit cut round trips 970 -> 1 and total
+ # 12.1 s -> 5.0 s, but DOUBLED time-to-first-chunk (642 -> 1292 ms)
+ # because the client serialises the whole graph before anything
+ # comes back. The progressive fill exists so the navigator starts
+ # filling immediately, so that regression matters as much as the
+ # total. Priming keeps the fast first paint and still costs only
+ # two submits.
+ "primed": False}
def _submit_next(lane):
"""Top up `lane` with a batch of pending chunks (lock held).
@@ -340,7 +351,38 @@ def _submit_next(lane):
# already folded into `caps` above; shadowing it here would read as if
# the parameter were being recomputed per top-up.
lane_cap = _lane_cap(caps[lane], lane_threads[lane], state["mem_hot"])
- n = min(submit_batch, len(pending), lane_cap - outstanding[lane])
+ if unpinned and lane == "cpu" and batch_unpinned and state["primed"]:
+ # ONE submit for the whole job. There is no placement decision to
+ # make on the unpinned lane and therefore nothing for a window to
+ # balance — so the window bought nothing and cost a blocking,
+ # GIL-held scheduler round trip PER CHUNK.
+ #
+ # It was worse than "no benefit": `_submit_next` runs on every
+ # completion, so `lane_cap - outstanding` is 1 in steady state and
+ # `min(submit_batch, …)` collapsed to ONE task per submit. The
+ # batching applied only to the very first fill. On a 977-chunk
+ # movie that is 977 round trips at ~14 ms with the GIL held —
+ # exactly the cost this module was written to remove, reintroduced
+ # through the back door. It scaled with dataset size, so it looked
+ # like "big datasets use one worker".
+ #
+ # Unbounded is measured-safe here: benchmarks.md has bounded at
+ # 46-50 s against unbounded at 50.2 s on the 977-chunk movie —
+ # within noise. And distributed >=2022.3 QUEUES root tasks at the
+ # scheduler (worker-saturation), which is this window's stated job
+ # done properly, in the right process, without our GIL.
+ #
+ # The DUAL-LANE path below keeps its window: there a completion
+ # genuinely pulls the next chunk so a ~30x-faster GPU lane and the
+ # CPU lane drain one pool and finish together. That is real work
+ # stealing the scheduler cannot do, and it is why this module
+ # exists at all.
+ n = len(pending)
+ elif unpinned and lane == "cpu" and batch_unpinned:
+ n = min(submit_batch, len(pending)) # priming batch
+ state["primed"] = True
+ else:
+ n = min(submit_batch, len(pending), lane_cap - outstanding[lane])
if n <= 0:
return
idxs = [pending.popleft() for _ in range(n)]
diff --git a/spyde/data/__init__.py b/spyde/data/__init__.py
index 44896f38..6478b512 100644
--- a/spyde/data/__init__.py
+++ b/spyde/data/__init__.py
@@ -26,7 +26,9 @@
eds_si,
eels_si,
ground_truth,
+ particle_movie,
+ particle_truth_at,
)
__all__ = ["eels_si", "eds_si", "ebsd_patterns", "atom_lattice",
- "ground_truth"]
+ "particle_movie", "particle_truth_at", "ground_truth"]
diff --git a/spyde/data/synthetic.py b/spyde/data/synthetic.py
index 8d5315e7..d813a18b 100644
--- a/spyde/data/synthetic.py
+++ b/spyde/data/synthetic.py
@@ -481,3 +481,238 @@ def ebsd_patterns(nav=(16, 16), detector=(60, 60), *, pc=(0.5, 0.5, 0.55),
pc=np.asarray(pc, float), n_bands=int(len(normals)),
grain2_mask=grain2)
return s
+
+
+# ---------------------------------------------------------------------------
+# In-situ particle movie
+# ---------------------------------------------------------------------------
+
+# The particle table. Hand-written rather than randomised so every number in
+# the ground truth is exact and readable: a test that says "particle 4
+# dissolves at frame 16" can be checked against this table by eye.
+#
+# Columns: y0, x0, radius, amp, birth, death, vy, vx, faint
+# y0/x0 position at t=0 in the SAMPLE frame, pixels
+# birth first frame the particle exists (0 = present from the start)
+# death first frame it is GONE (-1 = never dissolves)
+# vy/vx sample-frame velocity, px/frame
+# faint amplitude is scaled down to `faint_amplitude` -- the Section 0.9 probes
+_PARTICLES: tuple[dict, ...] = (
+ # two bright anchors, static and persistent
+ dict(y0=24.0, x0=22.0, radius=7.0, amp=1.00, birth=0, death=-1, vy=0.0, vx=0.0, faint=False),
+ dict(y0=70.0, x0=30.0, radius=9.0, amp=0.90, birth=0, death=-1, vy=0.0, vx=0.0, faint=False),
+ # a mover, for trails and displacement
+ dict(y0=30.0, x0=80.0, radius=6.0, amp=0.85, birth=0, death=-1, vy=1.10, vx=-0.70, faint=False),
+ # nucleation
+ dict(y0=60.0, x0=94.0, radius=5.0, amp=0.80, birth=8, death=-1, vy=0.0, vx=0.0, faint=False),
+ # dissolution
+ dict(y0=16.0, x0=58.0, radius=6.0, amp=0.75, birth=0, death=16, vy=0.0, vx=0.0, faint=False),
+ # the merge pair: converge along x until they overlap
+ dict(y0=84.0, x0=60.0, radius=5.0, amp=0.80, birth=0, death=-1, vy=0.0, vx=0.45, faint=False),
+ dict(y0=84.0, x0=82.0, radius=5.0, amp=0.80, birth=0, death=-1, vy=0.0, vx=-0.45, faint=False),
+ # faint, low-contrast -- these are what plan Section 0.9 is about
+ dict(y0=46.0, x0=102.0, radius=4.0, amp=1.00, birth=0, death=-1, vy=0.0, vx=0.0, faint=True),
+ dict(y0=88.0, x0=14.0, radius=3.0, amp=0.85, birth=0, death=-1, vy=0.0, vx=0.0, faint=True),
+)
+
+#: The two particles that merge, plus the nucleating and dissolving ones.
+MERGE_PAIR = (5, 6)
+NUCLEATION_INDEX = 3
+DISSOLUTION_INDEX = 4
+
+
+def _drift_curve(n_frames: int, amplitude: float) -> np.ndarray:
+ """``(n_frames, 2)`` per-frame CORRECTION, matching ``DriftModel``'s sign.
+
+ ``drift[t]`` is what you ADD to frame *t* to align it, so the sample appears
+ displaced by ``-drift[t]`` in the raw frame.
+
+ The two axes get deliberately DIFFERENT shapes -- y grows monotonically with
+ a slight curve, x swings negative and comes back. A swapped or negated axis
+ therefore shows up as a wrong-shaped curve rather than merely a wrong
+ number, which is the whole point of an asymmetric fixture.
+ """
+ if n_frames < 2:
+ return np.zeros((max(1, n_frames), 2), np.float64)
+ t = np.arange(n_frames, dtype=np.float64) / (n_frames - 1)
+ dy = amplitude * t ** 1.3
+ dx = -0.6 * amplitude * np.sin(1.7 * np.pi * t)
+ return np.stack([dy, dx], axis=-1)
+
+
+def _soft_disc(yy, xx, cy, cx, radius, edge=0.9):
+ """A disc with a smooth ~1 px edge, so sub-pixel centroids are meaningful.
+
+ A hard-edged disc quantises its own centroid to the pixel grid, which would
+ make any sub-pixel tracking assertion against this fixture untestable.
+ """
+ r = np.sqrt((yy - cy) ** 2 + (xx - cx) ** 2)
+ return 0.5 * (1.0 - np.tanh((r - radius) / edge))
+
+
+def _particle_arrays(drift: np.ndarray, faint_amplitude: float) -> dict:
+ """The per-particle arrays, in ONE place so the renderer and the stamped
+ ground truth cannot disagree about the motion model."""
+ return dict(
+ drift=drift,
+ p_y0=np.array([p["y0"] for p in _PARTICLES]),
+ p_x0=np.array([p["x0"] for p in _PARTICLES]),
+ p_radius=np.array([p["radius"] for p in _PARTICLES]),
+ p_amp=np.array([(faint_amplitude * p["amp"]) if p["faint"] else p["amp"]
+ for p in _PARTICLES]),
+ p_birth=np.array([p["birth"] for p in _PARTICLES]),
+ p_death=np.array([p["death"] for p in _PARTICLES]),
+ p_vy=np.array([p["vy"] for p in _PARTICLES]),
+ p_vx=np.array([p["vx"] for p in _PARTICLES]),
+ p_faint=np.array([p["faint"] for p in _PARTICLES]),
+ n_particles=len(_PARTICLES),
+ )
+
+
+def particle_truth_at(truth: dict, t: int):
+ """Expected ``(positions, radii, present)`` at frame *t*.
+
+ ``positions`` is ``(N, 2)`` ``(y, x)`` in **pixels, in the LAB frame** --
+ what a segmenter looking at the RAW frame should find, drift included.
+ ``present`` is a boolean mask of which particles exist in that frame.
+
+ Every consumer computes its expectations through here rather than
+ re-deriving the motion model, so a test cannot pass by repeating the same
+ mistake the generator made.
+ """
+ t = int(t)
+ y0 = np.asarray(truth["p_y0"], float)
+ x0 = np.asarray(truth["p_x0"], float)
+ vy = np.asarray(truth["p_vy"], float)
+ vx = np.asarray(truth["p_vx"], float)
+ birth = np.asarray(truth["p_birth"], int)
+ death = np.asarray(truth["p_death"], int)
+ drift = np.asarray(truth["drift"], float)
+
+ sample = np.stack([y0 + vy * t, x0 + vx * t], axis=-1)
+ lab = sample - drift[t] # see _drift_curve on the sign
+ present = (t >= birth) & ((death < 0) | (t < death))
+ return lab, np.asarray(truth["p_radius"], float), present
+
+
+def _merge_frame(arrays: dict, n_frames: int) -> int:
+ """First frame where the merge pair's discs overlap.
+
+ Derived from the same motion model that draws them rather than hard-coded,
+ so the stamped truth stays correct if the table or the frame count changes.
+ Returns -1 if they never touch within the movie.
+ """
+ a, b = MERGE_PAIR
+ ra, rb = _PARTICLES[a]["radius"], _PARTICLES[b]["radius"]
+ for t in range(n_frames):
+ pos, _, present = particle_truth_at(arrays, t)
+ if present[a] and present[b] and np.hypot(*(pos[a] - pos[b])) <= ra + rb:
+ return t
+ return -1
+
+
+def particle_movie(n_frames: int = 24, shape=(96, 112), *,
+ drift_amplitude: float = 6.0,
+ faint_amplitude: float = 0.11,
+ scale: float = 0.5, seed: int = 0, noise: float = 0.015):
+ """An in-situ particle movie whose every event is known exactly.
+
+ A 1-D navigation (time) axis over 2-D images: nine particles on a drifting
+ support film, with one nucleation, one dissolution, one merge, one mover and
+ two deliberately faint low-contrast probes.
+
+ Everything a downstream test needs to assert is stamped as ground truth --
+ read it with :func:`ground_truth` and evaluate the motion model with
+ :func:`particle_truth_at`.
+
+ Parameters
+ ----------
+ n_frames
+ Number of time points. The nucleation (frame 8) and dissolution
+ (frame 16) frames are fixed, so keep this above ~20 for both to occur.
+ shape
+ ``(ny, nx)``. **Deliberately non-square** so a transposed frame is
+ obvious at a glance.
+ drift_amplitude
+ Peak per-frame drift correction, pixels. The support film drifts
+ rigidly; particles move relative to it.
+ faint_amplitude
+ Peak amplitude of the two faint probes. The default puts them around
+ 7x the noise sigma -- findable, but not by a threshold tuned for the
+ bright particles. This is the plan's Section 0.9 gate.
+ scale
+ Pixel size in nm, written onto both signal axes.
+ seed
+ Everything random here (film speckle, noise) comes from this.
+ noise
+ Gaussian sigma added last. 0 disables.
+
+ Notes
+ -----
+ **Why there is a speckled support film.** Drift is only recoverable if
+ something static dominates the correlation. Particles move, appear and
+ vanish, so they cannot serve; a smooth gradient has too little
+ high-frequency content to correlate sharply. A rigidly-drifting speckle
+ field is both physically right and what lets
+ :func:`spyde.drift.solve_translation` recover ``drift`` from this movie.
+
+ The film is generated once on a padded canvas and sampled per frame with
+ bilinear interpolation, so there are no edge artifacts at any drift.
+ Particles are drawn **analytically** at their lab positions, so their
+ centroids stay exact regardless of how the film was resampled.
+ """
+ import hyperspy.api as hs
+ from scipy.ndimage import gaussian_filter, map_coordinates
+
+ ny, nx = int(shape[0]), int(shape[1])
+ n_frames = int(n_frames)
+ rng = np.random.default_rng(seed)
+ drift = _drift_curve(n_frames, float(drift_amplitude))
+ arrays = _particle_arrays(drift, float(faint_amplitude))
+
+ # The support film, on a canvas padded to cover every drift.
+ pad = int(np.ceil(np.abs(drift).max())) + 3
+ fy, fx = ny + 2 * pad, nx + 2 * pad
+ yy_p, _ = np.mgrid[0:fy, 0:fx]
+ film = 0.10 + 0.12 * (yy_p / fy) # ramp along y ONLY (asymmetric)
+ film += 0.09 * gaussian_filter(rng.standard_normal((fy, fx)), 1.6)
+
+ yy, xx = np.mgrid[0:ny, 0:nx].astype(np.float64)
+ frames = np.empty((n_frames, ny, nx), dtype=np.float32)
+
+ for t in range(n_frames):
+ dy, dx = drift[t]
+ # Sample the film at the lab-frame offset: content sits at -drift.
+ frame = map_coordinates(film, [yy + pad - dy, xx + pad - dx],
+ order=1, mode="nearest")
+ pos, radii, present = particle_truth_at(arrays, t)
+ for i, p in enumerate(_PARTICLES):
+ if not present[i]:
+ continue
+ amp = (faint_amplitude * p["amp"]) if p["faint"] else p["amp"]
+ frame = frame + amp * _soft_disc(yy, xx, pos[i, 0], pos[i, 1], radii[i])
+ frames[t] = frame
+
+ if noise:
+ frames += (noise * rng.standard_normal(frames.shape)).astype(np.float32)
+
+ s = hs.signals.Signal2D(frames)
+ for ax in s.axes_manager.signal_axes:
+ ax.scale, ax.units = float(scale), "nm"
+ tax = s.axes_manager.navigation_axes[0]
+ tax.name, tax.units, tax.scale = "time", "s", 0.05
+ s.metadata.General.title = "Synthetic particle movie"
+ # `insitu` is registered by SpyDE's OWN hyperspy extension, so unlike the
+ # EELS/EBSD generators this cannot silently fail for a missing extra.
+ _try_signal_type(s, "insitu")
+ _stamp(s, kind="particle_movie", **arrays,
+ n_frames=n_frames, frame_shape=np.asarray([ny, nx]),
+ scale=float(scale), noise=float(noise),
+ faint_amplitude=float(faint_amplitude),
+ merge_pair=np.asarray(MERGE_PAIR),
+ nucleation_index=NUCLEATION_INDEX,
+ dissolution_index=DISSOLUTION_INDEX,
+ nucleation_frame=int(_PARTICLES[NUCLEATION_INDEX]["birth"]),
+ dissolution_frame=int(_PARTICLES[DISSOLUTION_INDEX]["death"]),
+ merge_frame=_merge_frame(arrays, n_frames))
+ return s
diff --git a/spyde/drawing/toolbars/plot_control_toolbar.py b/spyde/drawing/toolbars/plot_control_toolbar.py
index 70c3727b..e10086ad 100644
--- a/spyde/drawing/toolbars/plot_control_toolbar.py
+++ b/spyde/drawing/toolbars/plot_control_toolbar.py
@@ -189,6 +189,7 @@ def get_toolbar_actions_for_plot(
exclude_signal_types = meta.get("exclude_signal_types")
signal_class = meta.get("signal_class")
requires_vectors = meta.get("requires_vectors", False)
+ requires_particles = meta.get("requires_particles", False)
plot_dim = meta.get("plot_dim", [1, 2])
navigation_only = meta.get("navigation")
params = meta.get("parameters", {})
@@ -199,8 +200,11 @@ def get_toolbar_actions_for_plot(
# requires_vectors: action only shows once the plot's signal tree has
# diffraction_vectors attached (set after Find Vectors completes).
# PlotState.rebuild_toolbars() re-runs this filter at that point.
+ # requires_particles is the same gate for tree.particles, set when a
+ # segmentation run finalizes.
tree = getattr(plot_state.plot, "signal_tree", None)
has_vectors = getattr(tree, "diffraction_vectors", None) is not None
+ has_particles = getattr(tree, "particles", None) is not None
add_action = (
(signal_types is None or plot_signal_type in signal_types)
@@ -221,6 +225,7 @@ def get_toolbar_actions_for_plot(
or isinstance(signal, _resolve_signal_class(signal_class))
)
and (not requires_vectors or has_vectors)
+ and (not requires_particles or has_particles)
# requires_original_metadata: hide unless the signal came from the
# format this action is about (see _has_original_metadata).
and _has_original_metadata(
@@ -290,6 +295,7 @@ def _action_matches_plot(action: str, meta: dict, plot_state: "PlotState") -> bo
exclude_signal_types = meta.get("exclude_signal_types")
signal_class = meta.get("signal_class")
requires_vectors = meta.get("requires_vectors", False)
+ requires_particles = meta.get("requires_particles", False)
plot_dim = meta.get("plot_dim", [1, 2])
navigation_only = meta.get("navigation")
@@ -298,6 +304,7 @@ def _action_matches_plot(action: str, meta: dict, plot_state: "PlotState") -> bo
tree = getattr(plot_state.plot, "signal_tree", None)
has_vectors = getattr(tree, "diffraction_vectors", None) is not None
+ has_particles = getattr(tree, "particles", None) is not None
return (
(signal_types is None or plot_signal_type in signal_types)
@@ -307,6 +314,7 @@ def _action_matches_plot(action: str, meta: dict, plot_state: "PlotState") -> bo
or isinstance(signal, _resolve_signal_class(signal_class))
)
and (not requires_vectors or has_vectors)
+ and (not requires_particles or has_particles)
and _has_original_metadata(signal, meta.get("requires_original_metadata"))
and _packages_present(meta)
and (plot_state.dimensions in plot_dim)
diff --git a/spyde/drift/__init__.py b/spyde/drift/__init__.py
new file mode 100644
index 00000000..7c64af3e
--- /dev/null
+++ b/spyde/drift/__init__.py
@@ -0,0 +1,47 @@
+"""
+spyde.drift — drift correction for image stacks and in-situ movies.
+
+See ``DRIFT_AND_PARTICLES_PLAN.md`` (repo root) for the full design. The two
+load-bearing constraints, both from CLAUDE.md:
+
+* **Nothing materialises the stack.** The target is thousands of frames at
+ 2048²–4096² (tens of GB). Every solver here STREAMS: read a frame, transform
+ it, discard it. The output of a solve is a small :class:`DriftModel` — an
+ ``(N, 2)`` shift array, not an aligned copy of the movie.
+* **The corrected movie is a LAZY VIEW.** ``DriftModel`` describes the
+ correction; applying it is a per-frame numpy operation
+ (:func:`spyde.drift.warp.shift_frame`) wired into the signal tree as a lazy
+ transformation node. Nothing is ever written out.
+
+Public API::
+
+ from spyde.drift import DriftModel, solve_translation, shift_frame
+
+ model = solve_translation(signal, upsample=8, max_shift=32)
+ aligned_frame = shift_frame(raw_frame, model.shifts[i])
+"""
+from __future__ import annotations
+
+from spyde.drift.frames import frame_source
+from spyde.drift.model import DriftModel
+from spyde.drift.translation import solve_translation
+from spyde.drift.warp import coverage_mask, shift_frame
+# Safe to import eagerly: nonrigid resolves torch lazily inside its functions,
+# so this does not drag a heavy import into every `spyde.drift` user.
+from spyde.drift.nonrigid import (
+ DENSE, SCAN_KNOT, apply_nonrigid, displacement_for_frame, solve_nonrigid,
+)
+
+__all__ = [
+ "DriftModel",
+ "solve_translation",
+ "shift_frame",
+ "coverage_mask",
+ "frame_source",
+ # non-rigid (plan A2-A5)
+ "solve_nonrigid",
+ "apply_nonrigid",
+ "displacement_for_frame",
+ "SCAN_KNOT",
+ "DENSE",
+]
diff --git a/spyde/drift/frames.py b/spyde/drift/frames.py
new file mode 100644
index 00000000..23602b46
--- /dev/null
+++ b/spyde/drift/frames.py
@@ -0,0 +1,92 @@
+"""
+frames.py — one streaming accessor for every kind of frame stack we accept.
+
+The whole point of this module is the Memory-Safety rule (CLAUDE.md): a drift
+solve on a 3000 × 4096² movie must never hold more than a couple of frames at
+once. So callers get a ``(n_frames, get_frame, frame_shape)`` triple and read
+frames one at a time — they never touch ``.data`` directly and therefore cannot
+accidentally ``.compute()`` the whole array.
+
+Accepted inputs, in the order they are tried:
+
+* a HyperSpy signal (lazy or eager) with 1-D navigation and 2-D signal axes
+* a dask array of shape ``(n, h, w)`` — sliced and computed per frame
+* a numpy array of shape ``(n, h, w)`` — already resident, just indexed
+* any sequence of 2-D arrays
+"""
+from __future__ import annotations
+
+from typing import Callable
+
+import numpy as np
+
+
+def _is_dask(obj) -> bool:
+ """True for a dask array, WITHOUT importing dask when it isn't already in."""
+ return type(obj).__module__.startswith("dask.array")
+
+
+def frame_source(data) -> tuple[int, Callable[[int], np.ndarray], tuple[int, int]]:
+ """Return ``(n_frames, get_frame, frame_shape)`` for *data*.
+
+ ``get_frame(i)`` returns frame ``i`` as a **numpy** 2-D array. For a lazy
+ (dask) backing it computes exactly that one frame — never the whole stack.
+
+ Raises
+ ------
+ TypeError
+ If *data* is not a recognised stack, or is not 3-D / not a sequence of
+ 2-D frames. Failing loudly here is deliberate: a silently-wrong axis
+ order would produce a plausible but meaningless drift curve.
+ """
+ # ── HyperSpy signal ──────────────────────────────────────────────────────
+ # Duck-typed on axes_manager so this module never imports hyperspy (import
+ # cost at backend startup) and so test doubles work.
+ if hasattr(data, "axes_manager") and hasattr(data, "data"):
+ am = data.axes_manager
+ nav = int(am.navigation_dimension)
+ sig = int(am.signal_dimension)
+ if sig != 2:
+ raise TypeError(
+ f"drift needs 2-D signal axes (images); got signal_dimension={sig}"
+ )
+ if nav != 1:
+ raise TypeError(
+ "drift needs a 1-D navigation axis (a frame stack / movie); got "
+ f"navigation_dimension={nav}. Reduce a higher-dimensional "
+ "dataset to a movie first (e.g. a virtual image)."
+ )
+ return frame_source(data.data)
+
+ # ── dask / numpy 3-D array ───────────────────────────────────────────────
+ if _is_dask(data) or isinstance(data, np.ndarray):
+ if data.ndim != 3:
+ raise TypeError(f"expected a 3-D (n, h, w) stack; got shape {data.shape}")
+ n, h, w = data.shape
+ if _is_dask(data):
+ def get_frame(i: int, _d=data) -> np.ndarray:
+ # ONE frame. Never `_d.compute()`.
+ return np.asarray(_d[int(i)].compute())
+ else:
+ def get_frame(i: int, _d=data) -> np.ndarray:
+ return np.asarray(_d[int(i)])
+ return int(n), get_frame, (int(h), int(w))
+
+ # ── sequence of 2-D frames ───────────────────────────────────────────────
+ try:
+ n = len(data)
+ except TypeError as exc:
+ raise TypeError(
+ f"cannot read frames from {type(data).__name__}: expected a HyperSpy "
+ "signal, a 3-D array, or a sequence of 2-D arrays"
+ ) from exc
+ if n == 0:
+ raise TypeError("empty frame stack")
+ first = np.asarray(data[0])
+ if first.ndim != 2:
+ raise TypeError(f"sequence elements must be 2-D frames; got ndim={first.ndim}")
+
+ def get_frame(i: int, _d=data) -> np.ndarray:
+ return np.asarray(_d[int(i)])
+
+ return int(n), get_frame, (int(first.shape[0]), int(first.shape[1]))
diff --git a/spyde/drift/model.py b/spyde/drift/model.py
new file mode 100644
index 00000000..c5040cfd
--- /dev/null
+++ b/spyde/drift/model.py
@@ -0,0 +1,173 @@
+"""
+model.py — :class:`DriftModel`, the small serialisable result of a drift solve.
+
+This is the ONLY thing a solve produces. It is deliberately tiny (an ``(N, 2)``
+float32 array plus metadata) because the alternative — an aligned copy of the
+movie — is tens of GB (see the module docstring in ``spyde/drift/__init__.py``).
+
+Sign convention — read this before touching anything
+-----------------------------------------------------
+``shifts[i]`` is the **correction**: the ``(dy, dx)`` you ADD to frame *i* to
+bring it into the reference frame. So::
+
+ aligned_i = scipy.ndimage.shift(frame_i, model.shifts[i])
+
+This matches ``skimage.registration.phase_cross_correlation``, whose docstring
+defines its return as "the shift vector required to register moving_image with
+reference_image", and matches ``scipy.ndimage.shift``'s own sign. Keeping all
+three identical is why the convention is stated here rather than inferred at each
+call site — an inverted sign produces a drift curve that looks entirely plausible
+and doubles the drift instead of removing it.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+# Bumped when the on-disk layout changes incompatibly.
+FORMAT_VERSION = 1
+
+
+@dataclass
+class DriftModel:
+ """The correction for a frame stack.
+
+ Parameters
+ ----------
+ shifts
+ ``(N, 2)`` float32, ``(dy, dx)`` per frame — the correction to ADD (see
+ the module docstring on sign convention).
+ kind
+ ``"rigid"`` today. ``"affine"`` / ``"scan_knot"`` / ``"dense"`` are the
+ non-rigid parameterisations planned in Wave A2–A5; they will carry their
+ parameters in :attr:`extra` and keep ``shifts`` as the rigid component.
+ reference
+ How the reference was formed: ``"running"`` (running Fourier average),
+ ``"sequential"`` (frame-to-frame, cumulative) or ``"fixed:"``.
+ residuals
+ Optional ``(N,)`` per-frame correlation peak sharpness — a weak but free
+ quality signal. NaN where not computed.
+ params
+ The solver arguments, for provenance and for re-running.
+ """
+
+ shifts: np.ndarray
+ kind: str = "rigid"
+ reference: str = "running"
+ residuals: np.ndarray | None = None
+ params: dict[str, Any] = field(default_factory=dict)
+ provenance: dict[str, Any] | None = None
+ extra: dict[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ self.shifts = np.ascontiguousarray(self.shifts, dtype=np.float32)
+ if self.shifts.ndim != 2 or self.shifts.shape[1] != 2:
+ raise ValueError(
+ f"shifts must be (N, 2); got {self.shifts.shape}"
+ )
+ if self.residuals is not None:
+ self.residuals = np.ascontiguousarray(self.residuals, dtype=np.float32)
+ if self.residuals.shape != (self.n_frames,):
+ raise ValueError(
+ f"residuals must be ({self.n_frames},); got {self.residuals.shape}"
+ )
+
+ # ── basic properties ─────────────────────────────────────────────────────
+
+ @property
+ def n_frames(self) -> int:
+ return int(self.shifts.shape[0])
+
+ @property
+ def max_abs_shift(self) -> float:
+ """Largest single-axis correction, in pixels. Sizes the padded border."""
+ if self.n_frames == 0:
+ return 0.0
+ return float(np.nanmax(np.abs(self.shifts)))
+
+ @property
+ def is_integer(self) -> bool:
+ """True when every shift is a whole number of pixels.
+
+ An integer-only model can be applied by ``np.roll`` and so **preserves
+ the source dtype**; a sub-pixel model needs interpolation and therefore
+ float output. Callers use this to avoid promoting a uint16 movie to
+ float32 when they don't have to.
+ """
+ if self.n_frames == 0:
+ return True
+ finite = self.shifts[np.isfinite(self.shifts)]
+ return bool(finite.size == 0 or np.all(finite == np.round(finite)))
+
+ def shift_at(self, index: int) -> np.ndarray:
+ """The ``(dy, dx)`` correction for frame *index*."""
+ return self.shifts[int(index)]
+
+ # ── frame-of-reference conversions ───────────────────────────────────────
+
+ def to_sample_frame(self, positions: np.ndarray, frame_index) -> np.ndarray:
+ """Map lab-frame ``(y, x)`` positions onto the corrected (sample) frame.
+
+ *positions* is ``(M, 2)``; *frame_index* is a scalar or ``(M,)``. This is
+ how a trajectory measured on RAW frames is reported as if the stage had
+ never moved — the alternative to physically warping the movie first.
+ """
+ pos = np.asarray(positions, dtype=np.float64)
+ if pos.ndim != 2 or pos.shape[1] != 2:
+ raise ValueError(f"positions must be (M, 2); got {pos.shape}")
+ idx = np.asarray(frame_index, dtype=np.intp)
+ return pos + self.shifts[idx].reshape(pos.shape)
+
+ def to_lab_frame(self, positions: np.ndarray, frame_index) -> np.ndarray:
+ """Inverse of :meth:`to_sample_frame`."""
+ pos = np.asarray(positions, dtype=np.float64)
+ if pos.ndim != 2 or pos.shape[1] != 2:
+ raise ValueError(f"positions must be (M, 2); got {pos.shape}")
+ idx = np.asarray(frame_index, dtype=np.intp)
+ return pos - self.shifts[idx].reshape(pos.shape)
+
+ # ── serialisation ────────────────────────────────────────────────────────
+
+ def save(self, path: str) -> None:
+ """Write to a compressed ``.npz``. Small enough to sit beside the data."""
+ meta = {
+ "format_version": FORMAT_VERSION,
+ "kind": self.kind,
+ "reference": self.reference,
+ "params": self.params,
+ "provenance": self.provenance,
+ "extra": self.extra,
+ }
+ arrays = {"shifts": self.shifts, "meta": np.array(json.dumps(meta))}
+ if self.residuals is not None:
+ arrays["residuals"] = self.residuals
+ np.savez_compressed(path, **arrays)
+
+ @classmethod
+ def load(cls, path: str) -> "DriftModel":
+ with np.load(path, allow_pickle=False) as z:
+ meta = json.loads(str(z["meta"].item()))
+ ver = meta.get("format_version")
+ if ver != FORMAT_VERSION:
+ raise ValueError(
+ f"unsupported DriftModel format version {ver!r} "
+ f"(this build reads {FORMAT_VERSION})"
+ )
+ return cls(
+ shifts=z["shifts"],
+ kind=meta.get("kind", "rigid"),
+ reference=meta.get("reference", "running"),
+ residuals=z["residuals"] if "residuals" in z.files else None,
+ params=meta.get("params") or {},
+ provenance=meta.get("provenance"),
+ extra=meta.get("extra") or {},
+ )
+
+ def __repr__(self) -> str:
+ return (
+ f"DriftModel(kind={self.kind!r}, n_frames={self.n_frames}, "
+ f"reference={self.reference!r}, max_abs_shift={self.max_abs_shift:.2f} px)"
+ )
diff --git a/spyde/drift/nonrigid.py b/spyde/drift/nonrigid.py
new file mode 100644
index 00000000..472e3589
--- /dev/null
+++ b/spyde/drift/nonrigid.py
@@ -0,0 +1,574 @@
+"""
+nonrigid.py — non-rigid drift: scan-knot and dense-field, one solver (plan A2-A5).
+
+Rigid translation (:mod:`spyde.drift.translation`) removes the part of the drift
+that moves the whole frame. What is left is real and is what this module fits:
+
+* **Scan distortion.** A scanned frame is not acquired instantaneously — the
+ stage keeps moving while the beam rasters, so each ROW is displaced by a
+ different amount. The distortion is therefore a function of the SLOW scan
+ coordinate, which is why one displacement per row (smoothed) is the natural
+ parameterisation and not a general 2-D field.
+* **Sample deformation.** The specimen itself bends, and parts of the field
+ move independently of each other. No function of the scan coordinate can
+ express that, so it needs displacements that vary in both directions.
+
+Both causes are real, so the model is SELECTABLE rather than assumed
+(:func:`solve_nonrigid`'s ``model=`` argument). They share the warp, the
+solver and the regularisation; only the parameter -> displacement map differs,
+which is the whole reason both are affordable.
+
+Sign convention
+---------------
+Identical to :class:`~spyde.drift.model.DriftModel`: a displacement is the
+correction you ADD to a pixel's coordinate to bring it into the reference. The
+rigid ``shifts`` stay in the model unchanged and the non-rigid field is the
+RESIDUAL on top of them, so a ``kind="scan_knot"`` model applied without its
+extra parameters still degrades gracefully to the rigid answer rather than to
+nonsense.
+
+Why gather (``grid_sample``) and not the KDE scatter the plan sketched
+---------------------------------------------------------------------
+quantem resamples with a KDE scatter (``index_put_(accumulate=True)`` over the
+four bilinear neighbours plus a weight image) because it is *building* a
+reconstruction from many scans, where several source pixels legitimately land on
+one output pixel and must accumulate.
+
+Here the job is the inverse: one frame, resampled onto the reference grid. That
+is a GATHER — for each output pixel, read the input at a computed coordinate —
+and ``torch.nn.functional.grid_sample`` does exactly that, differentiably, with
+a fused CUDA kernel. A scatter would need its own normalisation pass, leaves
+holes wherever no source pixel lands, and is strictly more code for a worse
+result on this problem. The scatter formulation is still the right one if this
+ever grows into multi-scan reconstruction; it is not needed to correct a movie.
+
+Out-of-bounds samples come back NaN, matching :mod:`spyde.drift.warp`'s locked
+edge policy (nothing cropped, nothing invented) — see :func:`warp_frame`.
+"""
+from __future__ import annotations
+
+import logging
+import math
+from typing import Any, Callable
+
+import numpy as np
+
+from spyde.drift.model import DriftModel
+
+log = logging.getLogger(__name__)
+
+# Parameterisation names, also the DriftModel.kind values they produce.
+SCAN_KNOT = "scan_knot"
+DENSE = "dense"
+MODELS = (SCAN_KNOT, DENSE)
+
+
+# ── torch plumbing ───────────────────────────────────────────────────────────
+
+def _torch():
+ try:
+ import torch
+ except ImportError as e: # pragma: no cover
+ raise RuntimeError(
+ "non-rigid drift needs torch; install it or use solve_translation"
+ ) from e
+ return torch
+
+
+def gpu_available() -> bool:
+ """True when a CUDA device is usable. Mirrors the other GPU paths."""
+ try:
+ import torch
+ return bool(torch.cuda.is_available())
+ except Exception:
+ return False
+
+
+def _resolve_device(device: str | None):
+ torch = _torch()
+ if device is not None:
+ return torch.device(device)
+ if torch.cuda.is_available():
+ return torch.device("cuda")
+ if getattr(torch.backends, "mps", None) is not None and torch.backends.mps.is_available():
+ return torch.device("mps")
+ return torch.device("cpu")
+
+
+# ── parameterisations: parameters -> per-pixel displacement ──────────────────
+
+def _bezier_basis(torch, n_rows: int, n_knots: int, device, dtype):
+ """``(n_rows, n_knots)`` Bernstein/Bézier basis over the slow scan axis.
+
+ Bézier rather than a free per-row displacement because scan distortion is
+ SMOOTH in the slow coordinate — the stage does not jerk row to row. A free
+ per-row fit has one parameter per row and happily absorbs sample motion and
+ noise into the "scan" term, which is precisely the failure this model is
+ supposed to avoid. ``n_knots=1`` (a constant offset per frame, i.e. degree 0)
+ is quantem's documented default for uniform distortion; 2-4 covers the
+ curved-drift case.
+ """
+ n_knots = max(1, int(n_knots))
+ t = torch.linspace(0.0, 1.0, n_rows, device=device, dtype=dtype)
+ if n_knots == 1:
+ return torch.ones((n_rows, 1), device=device, dtype=dtype)
+ deg = n_knots - 1
+ ks = torch.arange(n_knots, device=device, dtype=dtype)
+ # C(deg, k) t^k (1-t)^(deg-k), built in log space so deg>10 stays finite.
+ logc = (math.lgamma(deg + 1)
+ - torch.lgamma(ks + 1) - torch.lgamma(torch.tensor(float(deg), device=device, dtype=dtype) - ks + 1))
+ tt = t[:, None].clamp(1e-7, 1 - 1e-7)
+ return torch.exp(logc[None, :] + ks[None, :] * torch.log(tt)
+ + (deg - ks)[None, :] * torch.log1p(-tt))
+
+
+def scan_knot_field(torch, knots, shape, scan_direction_degrees: float = 0.0):
+ """Per-pixel ``(dy, dx)`` from scan knots.
+
+ ``knots`` is ``(N, 2, n_knots)`` — per frame, per scan axis (fast, slow),
+ per knot. The displacement varies only along the SLOW scan coordinate, then
+ is projected onto image axes using the scan direction, so a rotated scan is
+ handled without a second parameterisation.
+ """
+ n, _, n_knots = knots.shape
+ h, w = shape
+ dev, dt = knots.device, knots.dtype
+ ang = math.radians(float(scan_direction_degrees))
+ # Fast axis unit vector in (y, x); slow axis is perpendicular.
+ fy, fx = math.sin(ang), math.cos(ang)
+ sy, sx = math.cos(ang), -math.sin(ang)
+
+ basis = _bezier_basis(torch, h, n_knots, dev, dt) # (h, n_knots)
+ # (N, 2, h): displacement magnitude along each scan axis, per row.
+ mag = torch.einsum("nck,rk->ncr", knots, basis)
+ fast, slow = mag[:, 0, :], mag[:, 1, :] # (N, h) each
+ dy = fast * fy + slow * sy
+ dx = fast * fx + slow * sx
+ # Constant across the fast axis (a row is acquired at one slow coordinate).
+ return dy[:, :, None].expand(n, h, w), dx[:, :, None].expand(n, h, w)
+
+
+def dense_field(torch, control, shape):
+ """Per-pixel ``(dy, dx)`` from a coarse control-point grid.
+
+ ``control`` is ``(N, 2, gh, gw)``. Upsampled bicubically to the frame, which
+ is the free-form-deformation standard: the grid is coarse (so the model
+ cannot chase noise) and the interpolation is smooth (so the recovered field
+ has no control-point creases).
+ """
+ n, _, gh, gw = control.shape
+ h, w = shape
+ up = torch.nn.functional.interpolate(
+ control, size=(h, w), mode="bicubic", align_corners=True)
+ return up[:, 0], up[:, 1]
+
+
+# ── the differentiable warp ──────────────────────────────────────────────────
+
+def warp_frame(torch, frame, dy, dx, *, fill_nan: bool = True):
+ """Resample ``frame`` at ``(y + dy, x + dx)``. Differentiable in dy/dx.
+
+ ``frame`` is ``(N, H, W)``; ``dy``/``dx`` are ``(N, H, W)``. Out-of-bounds
+ samples become NaN when *fill_nan*, matching the locked edge policy in
+ :mod:`spyde.drift.warp` — nothing is cropped and nothing is invented.
+ """
+ n, h, w = frame.shape
+ dev, dt = frame.device, frame.dtype
+ yy = torch.arange(h, device=dev, dtype=dt)[None, :, None]
+ xx = torch.arange(w, device=dev, dtype=dt)[None, None, :]
+ sy = yy + dy
+ sx = xx + dx
+ # grid_sample wants normalised [-1, 1] with align_corners=True.
+ gy = 2.0 * sy / max(h - 1, 1) - 1.0
+ gx = 2.0 * sx / max(w - 1, 1) - 1.0
+ grid = torch.stack((gx, gy), dim=-1) # (N, H, W, 2)
+ out = torch.nn.functional.grid_sample(
+ frame[:, None], grid, mode="bilinear",
+ padding_mode="zeros", align_corners=True)[:, 0]
+ if not fill_nan:
+ return out
+ inside = (gy.abs() <= 1.0) & (gx.abs() <= 1.0)
+ return torch.where(inside, out, torch.full_like(out, float("nan")))
+
+
+# ── regularisation ───────────────────────────────────────────────────────────
+
+def _bending_energy(torch, field):
+ """Second-difference energy of a control grid — the FFD smoothness term.
+
+ Penalising CURVATURE rather than magnitude is deliberate: a uniform or
+ linearly-varying displacement is exactly what a real drift looks like and
+ must not be taxed, while a grid that folds or oscillates between neighbouring
+ control points is not a physical deformation.
+ """
+ e = field.new_zeros(())
+ if field.shape[-2] >= 3:
+ e = e + (field[..., :-2, :] - 2 * field[..., 1:-1, :] + field[..., 2:, :]).pow(2).mean()
+ if field.shape[-1] >= 3:
+ e = e + (field[..., :, :-2] - 2 * field[..., :, 1:-1] + field[..., :, 2:]).pow(2).mean()
+ return e
+
+
+def _temporal_energy(torch, params):
+ """Penalise frame-to-frame CHANGE of the parameters.
+
+ Drift is continuous in time — the distortion in frame *i* is nearly that of
+ frame *i-1*. This is what lets a noisy frame borrow support from its
+ neighbours, and it is the term that keeps a single bad frame from acquiring
+ its own wild field.
+ """
+ if params.shape[0] < 2:
+ return params.new_zeros(())
+ return (params[1:] - params[:-1]).pow(2).mean()
+
+
+# ── the solver ───────────────────────────────────────────────────────────────
+
+def solve_nonrigid(
+ frames,
+ *,
+ model: str = SCAN_KNOT,
+ reference=None,
+ rigid: DriftModel | None = None,
+ n_knots: int = 2,
+ grid: tuple[int, int] = (4, 4),
+ scan_direction_degrees: float = 0.0,
+ steps: int = 120,
+ lr: float = 0.5,
+ smooth_weight: float = 1.0,
+ temporal_weight: float = 1.0,
+ max_displacement: float | None = 32.0,
+ device: str | None = None,
+ progress: Callable[[int, int], None] | None = None,
+ on_yield: Callable[[], None] | None = None,
+ cancel: Callable[[], bool] | None = None,
+ provenance: dict[str, Any] | None = None,
+) -> DriftModel:
+ """Fit a non-rigid correction on top of a rigid solve.
+
+ Parameters
+ ----------
+ frames
+ ``(N, H, W)`` array — already rigid-corrected, or pass *rigid* and it is
+ applied here. Small enough to hold: this is a fit over a few
+ parameters, so callers are expected to pass a decimated or cropped
+ stack, not a 900x4096x4096 movie.
+ model
+ ``"scan_knot"`` (default) or ``"dense"``. See the module docstring for
+ which physical cause each one describes.
+ reference
+ ``(H, W)`` target. Defaults to the mean of *frames*, which is the right
+ default for drift: the mean of an already-rigid-aligned stack is the
+ sharpest thing available without picking a privileged frame.
+ n_knots, grid
+ Model size. ``n_knots`` for scan-knot; ``grid`` for dense.
+ smooth_weight, temporal_weight
+ Regularisation strengths. Both default to 1.0 against a mean-squared
+ data term, i.e. deliberately NOT free — an unregularised dense field
+ will happily fit noise.
+ max_displacement
+ Hard clamp on the fitted field, in pixels. ``None`` disables.
+
+ Returns
+ -------
+ DriftModel
+ ``kind`` is *model*, ``shifts`` carries the rigid component (zeros if
+ none was given), and ``extra`` holds the parameters plus everything
+ needed to rebuild the field (``field_shape``, ``n_knots``/``grid``,
+ ``scan_direction_degrees``).
+
+ Notes
+ -----
+ The Windows CUDA-autograd mitigations are load-bearing and both are applied:
+ ``backward()`` segfaults the first time it runs on a thread whose autograd
+ engine is uninitialised, so a warm-up backward runs on THIS thread before
+ the loop, and multithreaded autograd is disabled around it. See CLAUDE.md.
+ """
+ if model not in MODELS:
+ raise ValueError(f"model must be one of {MODELS}; got {model!r}")
+
+ torch = _torch()
+ from spyde.device_lock import accelerator_lock
+
+ arr = np.asarray(frames)
+ if arr.ndim != 3:
+ raise ValueError(f"frames must be (N, H, W); got {arr.shape}")
+ n, h, w = arr.shape
+ dev = _resolve_device(device)
+
+ # Warm the autograd engine on the CALLING thread — see Notes.
+ _warmup_autograd(torch, dev)
+
+ with accelerator_lock(dev):
+ f = torch.as_tensor(np.ascontiguousarray(arr, dtype=np.float32), device=dev)
+
+ if rigid is not None and rigid.n_frames == n:
+ sh = torch.as_tensor(np.asarray(rigid.shifts, np.float32), device=dev)
+ f = warp_frame(torch, f,
+ -sh[:, 0, None, None].expand(n, h, w),
+ -sh[:, 1, None, None].expand(n, h, w),
+ fill_nan=False)
+
+ ref = (torch.as_tensor(np.asarray(reference, np.float32), device=dev)
+ if reference is not None else f.mean(0))
+ ref = ref[None]
+
+ # Standardise so the loss scale (and therefore the regularisation
+ # weights) does not depend on the detector's units.
+ mu, sd = f.mean(), f.std().clamp_min(1e-6)
+ f = (f - mu) / sd
+ ref = (ref - mu) / sd
+
+ if model == SCAN_KNOT:
+ p = torch.zeros((n, 2, max(1, int(n_knots))), device=dev, requires_grad=True)
+ def field(pp):
+ return scan_knot_field(torch, pp, (h, w), scan_direction_degrees)
+ else:
+ gh, gw = (max(2, int(grid[0])), max(2, int(grid[1])))
+ p = torch.zeros((n, 2, gh, gw), device=dev, requires_grad=True)
+ def field(pp):
+ return dense_field(torch, pp, (h, w))
+
+ opt = torch.optim.Adam([p], lr=float(lr))
+ total = max(1, int(steps))
+ prev_mt = None
+ try:
+ prev_mt = torch.autograd.is_multithreading_enabled()
+ torch.autograd.set_multithreading_enabled(False)
+ except Exception: # pragma: no cover
+ prev_mt = None
+
+ try:
+ for it in range(total):
+ if cancel is not None and cancel():
+ log.info("non-rigid drift cancelled at step %d/%d", it, total)
+ break
+ opt.zero_grad(set_to_none=True)
+ dy, dx = field(p)
+ if max_displacement is not None:
+ m = float(max_displacement)
+ dy = dy.clamp(-m, m)
+ dx = dx.clamp(-m, m)
+ moved = warp_frame(torch, f, dy, dx, fill_nan=False)
+ data = (moved - ref).pow(2).mean()
+ reg = smooth_weight * (_bending_energy(torch, p) if model == DENSE
+ else _knot_energy(torch, p))
+ reg = reg + temporal_weight * _temporal_energy(torch, p)
+ loss = data + reg
+ loss.backward()
+ opt.step()
+
+ if progress is not None and (it % 8 == 0 or it == total - 1):
+ progress(it + 1, total)
+ # Yield INSIDE the loop, not per stage — otherwise the window
+ # freezes for seconds and the progress bar looks stuck.
+ if on_yield is not None and it % 12 == 0:
+ _yield_device(torch, dev, on_yield)
+ finally:
+ if prev_mt is not None:
+ try:
+ torch.autograd.set_multithreading_enabled(prev_mt)
+ except Exception: # pragma: no cover
+ pass
+
+ with torch.no_grad():
+ dy, dx = field(p)
+ if max_displacement is not None:
+ m = float(max_displacement)
+ dy, dx = dy.clamp(-m, m), dx.clamp(-m, m)
+ final = float((warp_frame(torch, f, dy, dx, fill_nan=False) - ref)
+ .pow(2).mean().item())
+ params_np = p.detach().float().cpu().numpy()
+ # Per-frame mean displacement — a compact, inspectable summary and
+ # what a 1-D "how much did this frame deform" plot shows.
+ mean_dy = dy.mean(dim=(1, 2)).float().cpu().numpy()
+ mean_dx = dx.mean(dim=(1, 2)).float().cpu().numpy()
+
+ shifts = (np.asarray(rigid.shifts, np.float32) if rigid is not None and rigid.n_frames == n
+ else np.zeros((n, 2), np.float32))
+ extra = {
+ "params": params_np,
+ "field_shape": (int(h), int(w)),
+ "mean_dy": np.asarray(mean_dy, np.float32),
+ "mean_dx": np.asarray(mean_dx, np.float32),
+ "final_mse": final,
+ "scan_direction_degrees": float(scan_direction_degrees),
+ }
+ if model == SCAN_KNOT:
+ extra["n_knots"] = int(max(1, n_knots))
+ else:
+ extra["grid"] = (int(max(2, grid[0])), int(max(2, grid[1])))
+
+ return DriftModel(
+ shifts=shifts,
+ kind=model,
+ reference="mean" if reference is None else "given",
+ params={
+ "model": model, "steps": int(steps), "lr": float(lr),
+ "smooth_weight": float(smooth_weight),
+ "temporal_weight": float(temporal_weight),
+ "max_displacement": max_displacement,
+ "n_knots": int(n_knots), "grid": tuple(grid),
+ "scan_direction_degrees": float(scan_direction_degrees),
+ },
+ provenance=provenance,
+ extra=extra,
+ )
+
+
+def _knot_energy(torch, knots):
+ """Smoothness across knots — the scan-knot analogue of bending energy."""
+ if knots.shape[-1] < 3:
+ return knots.new_zeros(())
+ return (knots[..., :-2] - 2 * knots[..., 1:-1] + knots[..., 2:]).pow(2).mean()
+
+
+def _warmup_autograd(torch, device) -> None:
+ """One trivial backward on THIS thread before any worker touches autograd.
+
+ CUDA-gated and a no-op elsewhere. On Windows the first ``backward()`` on a
+ thread whose autograd engine is uninitialised segfaults — uncatchably — and
+ the solve is dispatched to a daemon worker. See CLAUDE.md.
+ """
+ if getattr(device, "type", None) != "cuda":
+ return
+ try:
+ x = torch.zeros(1, device=device, requires_grad=True)
+ (x * x).sum().backward()
+ except Exception as e: # pragma: no cover
+ log.debug("autograd warm-up failed (continuing): %s", e)
+
+
+def _yield_device(torch, device, on_yield) -> None:
+ """Hand the accelerator back at a yield point, then take it again.
+
+ Always synchronise BEFORE releasing: handing off while kernels are still in
+ flight lets the next thread submit into a live encoder, which is the MPS
+ race the device lock exists to prevent.
+ """
+ try:
+ if getattr(device, "type", None) == "mps":
+ torch.mps.synchronize()
+ elif getattr(device, "type", None) == "cuda":
+ torch.cuda.synchronize()
+ except Exception as e: # pragma: no cover
+ log.debug("device sync before yield failed: %s", e)
+ try:
+ on_yield()
+ except Exception as e: # pragma: no cover
+ log.debug("on_yield raised (ignored): %s", e)
+
+
+# ── applying a fitted model ──────────────────────────────────────────────────
+
+def _field_on_device(torch, model: DriftModel, index: int, device):
+ """Build one frame's ``(dy, dx)`` DIRECTLY on *device*.
+
+ Separate from :func:`displacement_for_frame` because that one returns numpy
+ for inspection, and routing the apply through it would build a 4096² field
+ on the CPU and then ship 134 MB of it across PCIe — the transfer is already
+ the dominant cost (see :func:`apply_nonrigid`), so adding two more arrays to
+ it is the wrong direction. The parameters are a few hundred bytes; send
+ those instead and expand on the far side.
+ """
+ p = np.asarray(model.extra["params"], np.float32)
+ if not 0 <= index < p.shape[0]:
+ raise IndexError(f"frame {index} out of range for {p.shape[0]} fitted frames")
+ shape = tuple(model.extra["field_shape"])
+ t = torch.as_tensor(p[index: index + 1], device=device)
+ if model.kind == SCAN_KNOT:
+ return scan_knot_field(
+ torch, t, shape, float(model.extra.get("scan_direction_degrees", 0.0)))
+ return dense_field(torch, t, shape)
+
+
+def displacement_for_frame(model: DriftModel, index: int) -> tuple[np.ndarray, np.ndarray]:
+ """Rebuild the ``(dy, dx)`` field for one frame of a fitted model.
+
+ Returns arrays of the frame's shape. Raises for a rigid-only model — the
+ caller wants :mod:`spyde.drift.warp` for those, and silently returning zeros
+ would turn "this model has no non-rigid part" into "this frame did not
+ deform", which is a different claim.
+ """
+ if model.kind not in MODELS:
+ raise ValueError(
+ f"model.kind is {model.kind!r}, not a non-rigid fit; use spyde.drift.warp"
+ )
+ torch = _torch()
+ p = np.asarray(model.extra["params"], np.float32)
+ if not 0 <= index < p.shape[0]:
+ raise IndexError(f"frame {index} out of range for {p.shape[0]} fitted frames")
+ shape = tuple(model.extra["field_shape"])
+ t = torch.as_tensor(p[index: index + 1])
+ if model.kind == SCAN_KNOT:
+ dy, dx = scan_knot_field(
+ torch, t, shape, float(model.extra.get("scan_direction_degrees", 0.0)))
+ else:
+ dy, dx = dense_field(torch, t, shape)
+ return (dy[0].numpy().copy(), dx[0].numpy().copy())
+
+
+def apply_nonrigid(frame, model: DriftModel, index: int,
+ *, device: str | None = None) -> np.ndarray:
+ """Apply a fitted non-rigid model to ONE frame. NaN outside coverage.
+
+ Per-frame by design, like :func:`spyde.drift.warp.shift_frame` — the aligned
+ movie is never materialised.
+
+ Parameters
+ ----------
+ device
+ ``None`` (default) picks CUDA when it is available, else CPU. This is
+ the dominant performance knob and it is worth being explicit about why.
+
+ Performance
+ -----------
+ Measured at 4096², TITAN X Pascal:
+
+ ========================== =========
+ CPU total 392 ms
+ build field (CPU) 68 ms
+ warp (CPU) 262 ms
+ CUDA, frame resident **7.9 ms**
+ CUDA, incl. both transfers 41 ms
+ ========================== =========
+
+ So the GPU path is ~9.5x end to end, and note WHAT it is bound by: the warp
+ itself is 7.9 ms and the host<->device copies are the other ~33 ms. It is
+ **transfer-bound, not compute-bound**, which has two consequences:
+
+ * Micro-optimising the warp buys almost nothing. The lever is not moving the
+ data — a batch pipeline that keeps frames resident on the device pays only
+ the 7.9 ms, i.e. ~2.4 s for a 300-frame movie instead of ~12 s.
+ * The field is built on the device from the PARAMETERS (a few hundred bytes)
+ rather than built on the host and shipped, which would add 134 MB of
+ transfer to the very thing that dominates.
+
+ A 4096² frame is ~1.4 GB/s each way of pure copy, so on a machine with no
+ CUDA device this stays the 392 ms CPU path and the caller should expect a
+ non-rigid movie to be an export-time operation, not a scrub-time one.
+ """
+ torch = _torch()
+ f = np.asarray(frame, np.float32)
+ shape = tuple(model.extra["field_shape"])
+ if f.shape != shape:
+ raise ValueError(f"frame is {f.shape}, model was fitted at {shape}")
+
+ dev = _resolve_device(device)
+ # MPS is excluded deliberately: the win here is a fused grid_sample, and the
+ # shared device lock makes an unsolicited MPS submission a contention risk
+ # for whatever else holds it. Opt in explicitly with device="mps".
+ if device is None and getattr(dev, "type", None) == "mps":
+ dev = torch.device("cpu")
+
+ from spyde.device_lock import accelerator_lock
+ with accelerator_lock(dev):
+ dy, dx = _field_on_device(torch, model, index, dev)
+ t = torch.as_tensor(f, device=dev)[None]
+ out = warp_frame(torch, t, dy, dx, fill_nan=True)
+ # No trailing .copy(): `.to("cpu")` already returns a tensor that OWNS
+ # its memory (it is a device->host copy, or a no-op clone-free view of a
+ # CPU tensor we just built), and `.numpy()` keeps that storage alive
+ # through the array's base. An extra copy here is 67 MB at 4096² and was
+ # measurably a fifth of the GPU path's total cost.
+ return out[0].detach().to("cpu").numpy()
diff --git a/spyde/drift/translation.py b/spyde/drift/translation.py
new file mode 100644
index 00000000..e0be84fa
--- /dev/null
+++ b/spyde/drift/translation.py
@@ -0,0 +1,699 @@
+"""
+translation.py — rigid (translation-only) drift solve. Plan step A1.
+
+Algorithm: FFT phase correlation with a **running Fourier average** reference and
+Guizar-Sicairos matrix-multiply DFT upsampling for the sub-pixel peak.
+
+Three things here are deliberate and worth reading before changing them.
+
+**1. It streams.** One frame is resident at a time (plus the accumulated reference
+FFT, which is frame-sized). A 3000 × 4096² movie is tens of GB; nothing here ever
+holds more than a few hundred MB. This is the CLAUDE.md Memory-Safety rule.
+
+**2. The reference is accumulated in FOURIER space, aligned by a phase ramp.**
+To add frame *i* to the running average *already aligned*, we multiply its FFT by
+``exp(-2πi(dy·fy + dx·fx))`` rather than resampling the frame and re-transforming.
+A translation is exactly a phase ramp in the Fourier domain, so this is not an
+approximation — it is free, it is exact even for sub-pixel shifts, and it avoids
+the interpolation blur that resample-then-average would accumulate over thousands
+of frames. That blur is the reason a naive running average degrades as the stack
+gets longer.
+
+**3. Sub-pixel refinement is a small matmul, not a padded inverse FFT.** Zero-
+padding the cross-correlation to get ``1/upsample`` resolution costs an FFT of
+``(H·u, W·u)`` — for u=8 on a 4096² frame that is a 32768² transform. The
+matrix-multiply DFT evaluates the correlation only on the ~12×12 window around
+the coarse peak, which is what the refinement actually needs.
+
+The torch and numpy paths run the *same* algorithm through a small operator
+adapter, so ``test_drift_translation.py`` can assert they agree bit-closely — that
+parity test is what protects the GPU path, since it has no independent reference.
+"""
+from __future__ import annotations
+
+import logging
+import math
+from typing import Any, Callable
+
+import numpy as np
+
+from spyde.drift.frames import frame_source
+from spyde.drift.model import DriftModel
+
+log = logging.getLogger(__name__)
+
+# Guizar-Sicairos: the refinement window spans 1.5 upsampled pixels either side of
+# the coarse peak. Matches skimage's `phase_cross_correlation` so the two agree.
+_UPSAMPLED_REGION_FACTOR = 1.5
+
+# Magnitude FLOOR for phase normalisation — NOT an additive epsilon.
+#
+# Pure phase correlation divides the cross-power spectrum by its own magnitude,
+# which is only meaningful where there is signal. Bins whose magnitude is
+# numerically zero must be left alone; dividing them by a tiny epsilon amplifies
+# rounding noise to UNIT magnitude, and since there are far more empty bins than
+# populated ones, that noise then dominates the inverse transform.
+#
+# This is not theoretical. With an additive `1e-12` the solver recovered the
+# synthetic particle movie's drift to 25 px (worse than not correcting at all)
+# as soon as apodisation was enabled — because windowing concentrates spectral
+# energy and pushes many more bins down into the numerical floor. With the floor
+# below it: 0.06 px. Matches skimage's `100 * finfo(float32).eps`.
+_PHASE_FLOOR = 100.0 * float(np.finfo(np.float32).eps) # ~1.19e-5
+
+# A frame whose correlation peak is weaker than this fraction of the running
+# MEDIAN peak is kept out of the accumulated reference.
+#
+# The running-average reference exists to be robust to one bad frame, but folding
+# every frame in unconditionally does the opposite: a dropped / blanked / saturated
+# frame has a broadband spectrum, so after phase normalisation it contributes as
+# much to the reference as a good frame and drags every subsequent registration
+# with it. Measured on a 5-frame stack with one frame replaced by pure noise, the
+# two frames AFTER the bad one came back ~3.9 px wrong; with this rejection they
+# are correct and only the bad frame itself is wrong.
+#
+# Both constants come from measurement, not taste. Peak strength relative to the
+# running median, measured across four stacks:
+#
+# worst NATURAL frame (clean sub-pixel stack, erratic peaks) 0.388
+# a frame replaced by pure noise 0.007
+#
+# So there is a ~50x gap to put a threshold in, and 0.25 sits inside it with margin
+# both ways. 0.5 was tried first and produced a FALSE rejection on the clean
+# sub-pixel stack — which is why this is not simply "half".
+#
+# _REJECT_MIN_SAMPLES is 1, not 3, and that is deliberate: a short stack cannot
+# afford a warm-up. On a 5-frame stack the bad frame arrives before three good ones
+# have been seen, so a 3-sample warm-up let it into the reference and the rule never
+# fired (measured: 0 rejections, and the two frames after it came back 3.9 px wrong).
+# With a 1-sample warm-up those two frames are recovered EXACTLY.
+#
+# The asymmetry justifies being aggressive: keeping a good frame OUT of the
+# reference only slows the averaging, while letting a bad frame IN corrupts every
+# registration after it. A rejected frame still gets its own shift reported.
+#
+# Windowing the median over the last N accepted frames was tried and made NO
+# difference at N=3, 5 or unbounded — the natural decay in peak strength as the
+# reference averages more frames is not steep enough to matter. Don't add it back.
+_REJECT_FRACTION = 0.25
+_REJECT_MIN_SAMPLES = 1
+
+# Smallest alignment ROI worth correlating. Below roughly this the upsampled
+# refinement window (1.5 x upsample, so 12 px at the default) approaches the box
+# itself and the peak has nowhere to sit.
+_MIN_ROI = 16
+
+
+# ── operator adapters ────────────────────────────────────────────────────────
+# The algorithm below is written once against this interface. `_TorchOps` is the
+# production path; `_NumpyOps` is the reference the parity test pins it against.
+
+class _NumpyOps:
+ name = "numpy"
+
+ def __init__(self, device=None):
+ self.device = None
+
+ def to_backend(self, a):
+ return np.asarray(a, dtype=np.float32)
+
+ def fft2(self, a):
+ return np.fft.fft2(a).astype(np.complex64)
+
+ def ifft2(self, a):
+ return np.fft.ifft2(a)
+
+ def conj(self, a):
+ return np.conj(a)
+
+ def abs(self, a):
+ return np.abs(a)
+
+ def clamp_min(self, a, floor):
+ return np.maximum(a, floor)
+
+ def fftfreq(self, n):
+ return np.fft.fftfreq(n).astype(np.float32)
+
+ def arange(self, n):
+ return np.arange(n, dtype=np.float32)
+
+ def exp(self, a):
+ return np.exp(a)
+
+ def masked_argmax(self, mag, mask):
+ m = np.where(mask, mag, -np.inf)
+ flat = int(np.argmax(m))
+ return divmod(flat, mag.shape[1])
+
+ def argmax2d(self, mag):
+ flat = int(np.argmax(mag))
+ return divmod(flat, mag.shape[1])
+
+ def tensordot_last(self, kernel, data):
+ """``kernel @ data`` contracting kernel's axis 1 with data's LAST axis."""
+ return np.tensordot(kernel, data, axes=(1, -1))
+
+ def scalar(self, a) -> float:
+ return float(a)
+
+ def mean_abs(self, a) -> float:
+ return float(np.mean(np.abs(a)))
+
+ def to_numpy(self, a):
+ return np.asarray(a)
+
+
+class _TorchOps:
+ name = "torch"
+
+ def __init__(self, device):
+ import torch
+ self._torch = torch
+ self.device = device
+
+ def to_backend(self, a):
+ t = self._torch
+ return t.as_tensor(np.ascontiguousarray(a, dtype=np.float32), device=self.device)
+
+ def fft2(self, a):
+ return self._torch.fft.fft2(a).to(self._torch.complex64)
+
+ def ifft2(self, a):
+ return self._torch.fft.ifft2(a)
+
+ def conj(self, a):
+ return self._torch.conj(a)
+
+ def abs(self, a):
+ return self._torch.abs(a)
+
+ def clamp_min(self, a, floor):
+ return self._torch.clamp(a, min=float(floor))
+
+ def fftfreq(self, n):
+ return self._torch.fft.fftfreq(n, device=self.device, dtype=self._torch.float32)
+
+ def arange(self, n):
+ return self._torch.arange(n, device=self.device, dtype=self._torch.float32)
+
+ def exp(self, a):
+ return self._torch.exp(a)
+
+ def masked_argmax(self, mag, mask):
+ t = self._torch
+ m = mag.masked_fill(~mask, float("-inf"))
+ flat = int(t.argmax(m.reshape(-1)).item())
+ return divmod(flat, mag.shape[1])
+
+ def argmax2d(self, mag):
+ t = self._torch
+ flat = int(t.argmax(mag.reshape(-1)).item())
+ return divmod(flat, mag.shape[1])
+
+ def tensordot_last(self, kernel, data):
+ return self._torch.tensordot(kernel, data, dims=([1], [data.ndim - 1]))
+
+ def scalar(self, a) -> float:
+ return float(a.item()) if hasattr(a, "item") else float(a)
+
+ def mean_abs(self, a) -> float:
+ return float(self._torch.mean(self._torch.abs(a)).item())
+
+ def to_numpy(self, a):
+ return a.detach().cpu().numpy()
+
+
+def _resolve_ops(device: str | None):
+ """Pick the backend: ``cuda`` > ``mps`` > **torch CPU** > numpy.
+
+ **torch CPU beats numpy even with no GPU**, and by a lot — measured on this
+ box, 120 × 512² frames at upsample=8::
+
+ numpy 6.57 s 18 frames/s
+ torch cpu 0.86 s 139 frames/s (7.7x)
+ torch cuda 0.42 s 284 frames/s (16x)
+
+ The reason is mundane: ``np.fft.fft2`` is single-threaded, ``torch.fft.fft2``
+ uses every core. A per-frame FFT is the entire cost of this solver, so that
+ one difference is the whole gap. An earlier version of this function preferred
+ numpy on CPU-only machines on the assumption that torch's dispatch overhead
+ would dominate at one frame at a time; that assumption was wrong by 7.7x.
+
+ numpy is kept as an **explicitly** selectable reference path (``device=
+ "numpy"``), which is what the backend-parity test pins the torch path against.
+ """
+ if device == "numpy":
+ return _NumpyOps(None)
+ try:
+ import torch
+ except Exception:
+ return _NumpyOps(None)
+
+ if device is None:
+ if torch.cuda.is_available():
+ device = "cuda"
+ elif getattr(torch.backends, "mps", None) is not None and \
+ torch.backends.mps.is_available():
+ device = "mps"
+ else:
+ device = "cpu"
+ try:
+ return _TorchOps(torch.device(device))
+ except Exception as exc: # pragma: no cover — bad device string
+ log.warning("[drift] torch device %r unusable (%s); using numpy", device, exc)
+ return _NumpyOps(None)
+
+
+# ── windows and masks (built once per solve) ──────────────────────────────────
+
+#: Default Tukey taper fraction. 0.25 tapers the outer ~12.5% at each edge and
+#: leaves the middle 75% at unit weight. See :func:`_taper2d` for why this is
+#: NOT 1.0 (a full Hann window).
+DEFAULT_TAPER_ALPHA = 0.25
+
+
+def _tukey1d(n: int, alpha: float) -> np.ndarray:
+ """Tukey (cosine-tapered) window. ``alpha=0`` rectangular, ``alpha=1`` Hann."""
+ if n < 2:
+ return np.ones(max(1, n), dtype=np.float32)
+ alpha = float(min(1.0, max(0.0, alpha)))
+ if alpha <= 0.0:
+ return np.ones(n, dtype=np.float32)
+ x = np.arange(n, dtype=np.float64) / (n - 1)
+ w = np.ones(n, dtype=np.float64)
+ lo = x < alpha / 2.0
+ hi = x > 1.0 - alpha / 2.0
+ w[lo] = 0.5 * (1.0 + np.cos(2.0 * np.pi / alpha * (x[lo] - alpha / 2.0)))
+ w[hi] = 0.5 * (1.0 + np.cos(2.0 * np.pi / alpha * (x[hi] - 1.0 + alpha / 2.0)))
+ return w.astype(np.float32)
+
+
+def _taper2d(ops, h: int, w: int, alpha: float):
+ """Separable Tukey taper, built in numpy once per solve then moved on-device.
+
+ **Why a Tukey taper and NOT a full Hann window.** Some apodisation is needed:
+ a feature entering or leaving at the frame edge otherwise correlates against
+ the *border discontinuity* rather than the sample. But a full Hann window
+ (``alpha=1``) reweights the entire frame, and once the drift is large the two
+ frames have different content under the taper — which manufactures a spurious
+ correlation peak that can outrank the true one.
+
+ Measured on the synthetic particle movie, frame 23 (true drift ``(6.0, 2.9)``):
+ with a full Hann window the strongest peak sits at ``(-19, 19)`` scoring 0.121
+ while the TRUE peak scores only 0.088, so the solve returns ``(-19.2, 18.6)``
+ — a 25 px error, worse than not correcting at all. **skimage's
+ ``phase_cross_correlation`` returns the same wrong answer on the same windowed
+ input**, so this is a property of full-frame windowing, not of either
+ implementation. Tapering only the outer edge leaves the interior comparable
+ between frames and recovers 0.06 px.
+
+ Do not "simplify" this back to a Hann window.
+ """
+ win = _tukey1d(h, alpha)[:, None] * _tukey1d(w, alpha)[None, :]
+ return ops.to_backend(win)
+
+
+def _shift_mask(ops, h: int, w: int, max_shift: float | None,
+ min_shift: float | None):
+ """Which cross-correlation bins are admissible shifts.
+
+ The correlation is un-shifted, so bin ``k`` means shift ``k`` for
+ ``k <= n//2`` and ``k - n`` above that. Both bounds are separable, so this is
+ an outer product of two 1-D masks — built once and reused for every frame.
+ """
+ def axis_mask(n: int):
+ k = np.arange(n)
+ s = np.where(k > n // 2, k - n, k).astype(np.float64)
+ ok = np.ones(n, dtype=bool)
+ if max_shift is not None:
+ ok &= np.abs(s) <= float(max_shift)
+ return ok, s
+
+ oky, sy = axis_mask(h)
+ okx, sx = axis_mask(w)
+ mask = oky[:, None] & okx[None, :]
+ if min_shift is not None:
+ # Exclude the near-zero-shift core. Useful when the reference already
+ # contains this frame (a running average does), because the trivial
+ # self-correlation peak at the origin can then outrank the real one.
+ r = np.hypot(sy[:, None], sx[None, :])
+ mask &= r >= float(min_shift)
+ if not mask.any():
+ raise ValueError(
+ "max_shift/min_shift exclude every possible shift "
+ f"(max_shift={max_shift}, min_shift={min_shift}, frame={h}x{w})"
+ )
+ if ops.name == "torch":
+ return ops._torch.as_tensor(mask, device=ops.device)
+ return mask
+
+
+# ── the correlation ──────────────────────────────────────────────────────────
+
+def _upsampled_dft(ops, data, region_size: int, upsample: float, offsets):
+ """Evaluate the inverse DFT of *data* on a small upsampled window.
+
+ Mirrors ``skimage.registration._masked_phase_cross_correlation._upsampled_dft``:
+ one kernel matmul per axis, contracting the last axis each time.
+ """
+ out = data
+ shape = tuple(data.shape)
+ for axis in (1, 0): # last axis first
+ n = shape[axis]
+ off = offsets[axis]
+ # NOTE the /upsample: this is `np.fft.fftfreq(n, upsample)` — frequencies
+ # scaled to the UPSAMPLED grid. Without it the window still evaluates and
+ # still finds a peak, but at 1/upsample of the intended resolution, so
+ # every recovered shift lands on a multiple of 1/upsample and the
+ # refinement silently does nothing.
+ freq = ops.fftfreq(n) / float(upsample)
+ kern = (ops.arange(region_size).reshape(region_size, 1) - float(off)) * \
+ freq.reshape(1, n)
+ if ops.name == "torch":
+ kern = ops.exp(ops._torch.complex(
+ ops._torch.zeros_like(kern), -2.0 * math.pi * kern))
+ else:
+ kern = np.exp(-2j * math.pi * kern).astype(np.complex64)
+ out = ops.tensordot_last(kern, out)
+ return out
+
+
+def _peak_shift(ops, ref_fft, mov_fft, mask, upsample: float,
+ normalize: bool) -> tuple[float, float, float]:
+ """Return ``(dy, dx, sharpness)`` registering *mov* onto *ref*.
+
+ Sign matches ``skimage.registration.phase_cross_correlation`` and
+ ``scipy.ndimage.shift``: the result is the correction to ADD to the moving
+ frame (see :mod:`spyde.drift.model`).
+ """
+ product = ref_fft * ops.conj(mov_fft)
+ if normalize:
+ # Phase correlation proper: discard magnitude, keep only phase. Gives a
+ # far sharper peak than plain cross-correlation on images whose spectra
+ # are dominated by low frequencies, which every real micrograph is.
+ # The divisor is FLOORED, not offset — see _PHASE_FLOOR.
+ product = product / ops.clamp_min(ops.abs(product), _PHASE_FLOOR)
+
+ cc = ops.ifft2(product)
+ mag = ops.abs(cc)
+ py, px = ops.masked_argmax(mag, mask)
+
+ h, w = int(mag.shape[0]), int(mag.shape[1])
+ peak = ops.scalar(mag[py, px])
+ baseline = ops.mean_abs(cc)
+ sharpness = float(peak / baseline) if baseline > 0 else float("nan")
+
+ dy = float(py - h) if py > h // 2 else float(py)
+ dx = float(px - w) if px > w // 2 else float(px)
+
+ if upsample and upsample > 1:
+ u = float(upsample)
+ dy = round(dy * u) / u
+ dx = round(dx * u) / u
+ region = int(math.ceil(u * _UPSAMPLED_REGION_FACTOR))
+ dftshift = float(region // 2)
+ offsets = (dftshift - dy * u, dftshift - dx * u)
+ fine = _upsampled_dft(ops, ops.conj(product), region, u, offsets)
+ fmag = ops.abs(fine)
+ my, mx = ops.argmax2d(fmag)
+ dy += (my - dftshift) / u
+ dx += (mx - dftshift) / u
+
+ return dy, dx, sharpness
+
+
+def _validate_roi(roi, full_h: int, full_w: int):
+ """Normalise and bounds-check an ``(y0, x0, h, w)`` alignment ROI.
+
+ Rejects rather than clamps. A silently shrunk ROI would correlate on a
+ different region than the one the user dragged, and the drift curve would be
+ wrong in a way nothing on screen could explain.
+ """
+ if roi is None:
+ return None
+ try:
+ y0, x0, h, w = (int(v) for v in roi)
+ except (TypeError, ValueError):
+ raise ValueError(
+ f"roi must be (y0, x0, h, w) in pixels; got {roi!r}") from None
+ if h < _MIN_ROI or w < _MIN_ROI:
+ raise ValueError(
+ f"roi is {h}x{w} px; the correlation needs at least "
+ f"{_MIN_ROI}x{_MIN_ROI} to locate a peak at all")
+ if y0 < 0 or x0 < 0 or y0 + h > full_h or x0 + w > full_w:
+ raise ValueError(
+ f"roi (y0={y0}, x0={x0}, h={h}, w={w}) falls outside the "
+ f"{full_h}x{full_w} frame")
+ return (y0, x0, h, w)
+
+
+def _accept_into_reference(sharpness: float, accepted: list[float],
+ enabled: bool) -> bool:
+ """Whether this frame is credible enough to join the running reference.
+
+ Always True until there are :data:`_REJECT_MIN_SAMPLES` accepted frames to
+ take a median over — with nothing to compare against, rejecting would just be
+ guessing. See :data:`_REJECT_FRACTION`.
+ """
+ if not enabled or len(accepted) < _REJECT_MIN_SAMPLES:
+ return True
+ if not math.isfinite(sharpness):
+ return False
+ return sharpness >= _REJECT_FRACTION * float(np.median(accepted))
+
+
+def _phase_ramp(ops, h: int, w: int, dy: float, dx: float):
+ """FFT multiplier that translates a frame by ``(dy, dx)`` exactly.
+
+ ``F{f(y - dy, x - dx)}(k) = exp(-2πi(dy·fy + dx·fx))·F{f}(k)``. This is why
+ the running reference needs no resampling — see the module docstring.
+ """
+ fy = ops.fftfreq(h).reshape(h, 1)
+ fx = ops.fftfreq(w).reshape(1, w)
+ ph = dy * fy + dx * fx
+ if ops.name == "torch":
+ return ops.exp(ops._torch.complex(
+ ops._torch.zeros_like(ph), -2.0 * math.pi * ph))
+ return np.exp(-2j * math.pi * ph).astype(np.complex64)
+
+
+# ── public solve ─────────────────────────────────────────────────────────────
+
+def solve_translation(
+ data,
+ *,
+ upsample: int = 8,
+ max_shift: float | None = 32.0,
+ min_shift: float | None = None,
+ reference: str = "running",
+ roi: tuple[int, int, int, int] | None = None,
+ apodize: bool | float = True,
+ normalize: bool = True,
+ reject_outliers: bool = True,
+ device: str | None = None,
+ progress: Callable[[int, int], None] | None = None,
+ on_shift: Callable[[int, float, float, float], None] | None = None,
+ cancel: Callable[[], bool] | None = None,
+ provenance: dict[str, Any] | None = None,
+) -> DriftModel:
+ """Solve rigid drift for a frame stack. Returns a :class:`DriftModel`.
+
+ Parameters
+ ----------
+ data
+ A HyperSpy signal (1-D nav, 2-D signal), a 3-D numpy/dask array, or a
+ sequence of 2-D frames. Read one frame at a time — never materialised.
+ upsample
+ Sub-pixel factor. ``8`` resolves to 1/8 px, which is well past the
+ ~0.05 px accuracy floor set by noise on real data.
+ max_shift
+ Reject correlation peaks implying a larger per-frame shift, in pixels.
+ Guards against a spurious peak from a periodic lattice — the failure mode
+ where a crystalline sample locks onto the wrong lattice translation and
+ the drift curve jumps by exactly one lattice spacing.
+ min_shift
+ Exclude peaks *smaller* than this. Off by default; see
+ :func:`_shift_mask`.
+ reference
+ ``"running"`` — running Fourier average (default, robust to one bad
+ frame); ``"sequential"`` — register to the previous frame and accumulate
+ (handles large excursions, accumulates error); ``"first"`` or
+ ``"fixed:"`` — one fixed reference frame.
+ roi
+ ``(y0, x0, h, w)`` in pixels — correlate on this sub-region only. The
+ returned shifts still apply to the WHOLE frame; a translation is a
+ translation regardless of which window you measured it in.
+
+ This is not merely a speed switch, it is often the more CORRECT answer.
+ Whole-frame correlation averages over everything that moved, so on an
+ in-situ movie where the sample is genuinely evolving — particles growing,
+ drifting, appearing — the sample's own motion contaminates the estimate of
+ the stage's. Restricting to a static, feature-rich landmark (a support
+ film edge, a fiducial, a stationary grain) measures the stage and nothing
+ else. It is also how a user can rescue a dataset where the field of view
+ is mostly featureless.
+
+ **The ROI is FIXED in frame coordinates**, so the landmark drifts within
+ it. That is fine while the drift is small compared with the box, and it is
+ why the box wants to be comfortably larger than the total excursion —
+ the caret's preview exists so this is judged by eye rather than guessed.
+ A box smaller than the drift will lose the landmark and the solve will
+ wander.
+ apodize
+ Edge taper before transforming. ``True`` uses a Tukey window with
+ ``alpha=DEFAULT_TAPER_ALPHA``; a float sets alpha explicitly
+ (``1.0`` = full Hann, which is a trap — see :func:`_taper2d`);
+ ``False`` disables it.
+ normalize
+ True phase correlation (unit-magnitude spectrum). Sharper peak.
+ reject_outliers
+ ``running`` mode only: keep a frame out of the accumulated reference when
+ its correlation peak is not credible (see :data:`_REJECT_FRACTION`). Its
+ own shift is still reported — only the reference is protected. This is
+ what makes "robust to one bad frame" true rather than aspirational.
+ device
+ ``None`` auto-selects CUDA/MPS then falls back to numpy; ``"numpy"``
+ forces the reference path; or an explicit torch device string.
+ progress, cancel
+ ``progress(done, total)`` is called as frames complete.
+ ``cancel()`` returning True aborts; frames not yet reached keep NaN
+ shifts, so a cancelled solve is detectable rather than silently partial.
+ on_shift
+ ``on_shift(index, dy, dx, sharpness)`` per frame, as each is solved.
+
+ This exists so a UI can draw the drift curve **while** it solves, which
+ ``progress`` cannot support: it carries only a count, and the shift array
+ is solver-local until the return. Splitting the solve into chunks and
+ concatenating would not be equivalent either — the running Fourier
+ reference accumulates across the whole stack, so a restarted solve gives a
+ different (worse) answer. A callback is the only way to stream the trace
+ without changing the result.
+
+ Called on the solver thread, so a UI implementation must marshal.
+
+ Notes
+ -----
+ Frame 0 is the origin by definition and always gets ``(0, 0)``.
+ """
+ if reference not in ("running", "sequential", "first") and \
+ not reference.startswith("fixed:"):
+ raise ValueError(
+ f"unknown reference {reference!r}; expected 'running', 'sequential', "
+ "'first' or 'fixed:'"
+ )
+ if upsample < 1:
+ raise ValueError(f"upsample must be >= 1; got {upsample}")
+
+ n_frames, get_frame, (full_h, full_w) = frame_source(data)
+ crop = _validate_roi(roi, full_h, full_w)
+ h, w = (full_h, full_w) if crop is None else (crop[2], crop[3])
+ ops = _resolve_ops(device)
+
+ shifts = np.full((n_frames, 2), np.nan, dtype=np.float32)
+ sharp = np.full((n_frames,), np.nan, dtype=np.float32)
+
+ from spyde.device_lock import accelerator_lock
+
+ # MPS is not thread-safe and every torch user in the process shares ONE lock
+ # (CLAUDE.md § GPU Computing). A null context off MPS, so CUDA keeps its
+ # stream concurrency. Held across the solve rather than per-frame: the solve
+ # runs on a worker thread and per-frame acquire/release would be pure
+ # overhead at thousands of frames.
+ with accelerator_lock(ops.device):
+ alpha = (DEFAULT_TAPER_ALPHA if apodize is True
+ else (0.0 if apodize is False else float(apodize)))
+ window = _taper2d(ops, h, w, alpha) if alpha > 0 else None
+ mask = _shift_mask(ops, h, w, max_shift, min_shift)
+
+ def frame_fft(i: int):
+ raw = get_frame(i)
+ if crop is not None:
+ y0, x0, ch, cw = crop
+ raw = raw[y0:y0 + ch, x0:x0 + cw]
+ f = ops.to_backend(raw)
+ if window is not None:
+ f = f * window
+ return ops.fft2(f)
+
+ fixed_index = 0
+ if reference.startswith("fixed:"):
+ fixed_index = int(reference.split(":", 1)[1])
+ if not 0 <= fixed_index < n_frames:
+ raise ValueError(
+ f"fixed reference index {fixed_index} outside 0..{n_frames - 1}"
+ )
+
+ first = frame_fft(fixed_index if reference.startswith("fixed:") else 0)
+ shifts[0] = (0.0, 0.0)
+ sharp[0] = np.inf if n_frames else np.nan
+ if on_shift is not None and n_frames:
+ on_shift(0, 0.0, 0.0, float("inf"))
+
+ ref_fft = first # running accumulator / fixed reference
+ ref_count = 1
+ prev_fft = first # sequential mode
+ cumulative = np.zeros(2, dtype=np.float64)
+ accepted_sharp: list[float] = [] # peak strengths folded into the reference
+ rejected = 0
+
+ if progress is not None:
+ progress(1, n_frames)
+
+ for i in range(1, n_frames):
+ if cancel is not None and cancel():
+ log.info("[drift] cancelled at frame %d/%d", i, n_frames)
+ break
+
+ mov = frame_fft(i)
+
+ if reference == "sequential":
+ dy, dx, s = _peak_shift(ops, prev_fft, mov, mask, upsample, normalize)
+ cumulative += (dy, dx)
+ shifts[i] = cumulative
+ prev_fft = mov
+ else:
+ dy, dx, s = _peak_shift(ops, ref_fft, mov, mask, upsample, normalize)
+ shifts[i] = (dy, dx)
+ if reference == "running" and _accept_into_reference(
+ s, accepted_sharp, reject_outliers):
+ # Fold the ALIGNED frame in via a phase ramp — exact, and no
+ # resampling blur accumulates over the stack.
+ aligned = mov * _phase_ramp(ops, h, w, dy, dx)
+ ref_fft = (ref_fft * ref_count + aligned) / (ref_count + 1)
+ ref_count += 1
+ accepted_sharp.append(s)
+ elif reference == "running":
+ rejected += 1
+ log.debug("[drift] frame %d kept out of the reference "
+ "(peak %.2f vs median %.2f)", i, s,
+ float(np.median(accepted_sharp)))
+ sharp[i] = s
+
+ if on_shift is not None:
+ on_shift(i, float(shifts[i, 0]), float(shifts[i, 1]), float(s))
+ if progress is not None:
+ progress(i + 1, n_frames)
+
+ params = {
+ "upsample": int(upsample),
+ "max_shift": None if max_shift is None else float(max_shift),
+ "min_shift": None if min_shift is None else float(min_shift),
+ "reference": reference,
+ "apodize": float(alpha),
+ "normalize": bool(normalize),
+ "reject_outliers": bool(reject_outliers),
+ "rejected_from_reference": int(rejected),
+ "backend": ops.name,
+ "n_frames": int(n_frames),
+ "frame_shape": [int(full_h), int(full_w)],
+ "roi": None if crop is None else [int(v) for v in crop],
+ }
+ return DriftModel(
+ shifts=shifts,
+ kind="rigid",
+ reference=reference,
+ residuals=sharp,
+ params=params,
+ provenance=provenance,
+ )
diff --git a/spyde/drift/warp.py b/spyde/drift/warp.py
new file mode 100644
index 00000000..03a46d64
--- /dev/null
+++ b/spyde/drift/warp.py
@@ -0,0 +1,134 @@
+"""
+warp.py — apply a drift correction to ONE frame.
+
+Per-frame by design. This is the function a lazy signal-tree node calls, so the
+aligned movie is never materialised (``spyde/drift/__init__.py``), and it is also
+what the derived-view reader would call if the per-frame shift transform is ever
+added to ``array_cache/readers/per_frame.py`` — which is a signal-tree read-path
+change and therefore gated on review, NOT done here.
+
+Edge policy (locked, plan §A7): the frame keeps its full size and uncovered
+pixels become **NaN**, with :func:`coverage_mask` giving the validity map.
+Nothing is cropped and nothing is filled with invented data.
+
+**Downstream contract:** segmentation MUST respect the coverage mask. A NaN-padded
+border is the single most likely integration bug in this feature — a threshold
+applied to NaN, or a NaN-to-zero conversion, invents a large "particle" along the
+edge that then nucleates a spurious track.
+"""
+from __future__ import annotations
+
+import numpy as np
+
+
+def _split_shift(shift) -> tuple[np.ndarray, bool]:
+ s = np.asarray(shift, dtype=np.float64).reshape(-1)
+ if s.size != 2:
+ raise ValueError(f"shift must be (dy, dx); got {np.shape(shift)}")
+ if not np.all(np.isfinite(s)):
+ raise ValueError(f"shift must be finite; got {s!r}")
+ return s, bool(np.all(s == np.round(s)))
+
+
+def coverage_mask(shape: tuple[int, int], shift) -> np.ndarray:
+ """Boolean map of which output pixels come from real source data.
+
+ A pixel is covered when its source coordinate falls inside the source frame.
+ For a sub-pixel shift the border row/column that would need to interpolate
+ against off-frame data is treated as **uncovered** — bilinear interpolation
+ there would silently blend in the fill value.
+ """
+ h, w = int(shape[0]), int(shape[1])
+ s, integral = _split_shift(shift)
+ dy, dx = s
+
+ mask = np.zeros((h, w), dtype=bool)
+ if integral:
+ y0, y1 = max(0, int(dy)), min(h, h + int(dy))
+ x0, x1 = max(0, int(dx)), min(w, w + int(dx))
+ else:
+ # Output pixel y draws from source y - dy; it needs floor and floor+1,
+ # so require 0 <= y - dy and y - dy + 1 <= h - 1.
+ y0 = int(np.ceil(max(0.0, dy)))
+ y1 = int(np.floor(min(float(h), h + dy - 1.0))) + 1
+ x0 = int(np.ceil(max(0.0, dx)))
+ x1 = int(np.floor(min(float(w), w + dx - 1.0))) + 1
+ if y1 > y0 and x1 > x0:
+ mask[y0:y1, x0:x1] = True
+ return mask
+
+
+def shift_frame(
+ frame: np.ndarray,
+ shift,
+ *,
+ order: int = 1,
+ fill: float = np.nan,
+ preserve_dtype: bool = False,
+) -> np.ndarray:
+ """Shift *frame* by ``(dy, dx)``, padding uncovered pixels with *fill*.
+
+ Parameters
+ ----------
+ frame
+ 2-D source frame, any dtype.
+ shift
+ ``(dy, dx)`` correction — the value ADDED to coordinates. See
+ :mod:`spyde.drift.model` for the sign convention.
+ order
+ Interpolation order for a sub-pixel shift (1 = bilinear, the default;
+ 3 = cubic). Ignored for a whole-pixel shift, which is done exactly.
+ fill
+ Value for uncovered pixels. NaN by default, which forces a float result.
+ preserve_dtype
+ Keep the source dtype. Only honoured for a whole-pixel shift with a
+ non-NaN *fill* — a uint16 frame cannot hold NaN, and interpolation
+ cannot be exact in an integer type. Raises otherwise rather than
+ silently returning something lossy.
+
+ Notes
+ -----
+ A whole-pixel shift takes an exact slice-copy path: no interpolation, no
+ float promotion, and bit-identical to the source pixels. This matters because
+ the common case for a well-behaved stage IS an integer shift, and running it
+ through ``scipy.ndimage.shift`` would resample (and blur) data that did not
+ need to move sub-pixel at all.
+ """
+ src = np.asarray(frame)
+ if src.ndim != 2:
+ raise ValueError(f"frame must be 2-D; got shape {src.shape}")
+ s, integral = _split_shift(shift)
+ dy, dx = s
+ h, w = src.shape
+
+ if preserve_dtype and not (integral and np.isfinite(fill)):
+ raise ValueError(
+ "preserve_dtype=True requires a whole-pixel shift and a finite fill "
+ f"(got shift={s.tolist()}, fill={fill!r}); a sub-pixel shift must "
+ "interpolate and NaN padding cannot be stored in an integer dtype"
+ )
+
+ out_dtype = src.dtype if preserve_dtype else np.float32
+
+ if integral:
+ out = np.full((h, w), fill, dtype=out_dtype)
+ iy, ix = int(dy), int(dx)
+ # Destination window, and the matching source window.
+ dy0, dy1 = max(0, iy), min(h, h + iy)
+ dx0, dx1 = max(0, ix), min(w, w + ix)
+ if dy1 > dy0 and dx1 > dx0:
+ out[dy0:dy1, dx0:dx1] = src[dy0 - iy:dy1 - iy, dx0 - ix:dx1 - ix]
+ return out
+
+ from scipy.ndimage import shift as ndi_shift
+
+ # ndimage cannot propagate NaN through its spline filter without smearing it,
+ # so interpolate with a finite sentinel and stamp the fill on afterwards using
+ # the analytic coverage map. This keeps the padded border crisp instead of
+ # letting a NaN bleed `order` pixels into real data.
+ work = src.astype(np.float32, copy=False)
+ out = ndi_shift(work, s, order=order, mode="constant", cval=0.0, prefilter=order > 1)
+ out = out.astype(out_dtype, copy=False)
+ cov = coverage_mask((h, w), s)
+ out[~cov] = fill
+ return out
diff --git a/spyde/hyperspy_extension.yaml b/spyde/hyperspy_extension.yaml
index c36ba51c..23e82af8 100644
--- a/spyde/hyperspy_extension.yaml
+++ b/spyde/hyperspy_extension.yaml
@@ -29,3 +29,15 @@ signals:
dtype: real
lazy: True
module: spyde.signals.insitu
+ ParticleMap:
+ signal_type: particles
+ signal_dimension: 2
+ dtype: real
+ lazy: False
+ module: spyde.signals.particle_map
+ LazyParticleMap:
+ signal_type: particles
+ signal_dimension: 2
+ dtype: real
+ lazy: True
+ module: spyde.signals.particle_map
diff --git a/spyde/particles/__init__.py b/spyde/particles/__init__.py
new file mode 100644
index 00000000..4338a832
--- /dev/null
+++ b/spyde/particles/__init__.py
@@ -0,0 +1,124 @@
+"""
+spyde.particles — particle segmentation, measurement and tracking.
+
+See ``DRIFT_AND_PARTICLES_PLAN.md`` (repo root) for the full design. The shape of
+this package is its most important property: **three interchangeable ways to
+produce a foreground probability map, and one shared downstream stage.**
+
+ frame ──► [ classical | scribble | prompt ] ──► probability / mask
+ │
+ split_instances() (watershed)
+ │
+ measure_frame() (props + hull → units)
+ │
+ SpyDEParticles (CSR, per frame)
+ │
+ link() (Hungarian) → tracks + events
+
+The three engines are not alternatives to choose between — they are three ways to
+fill the first box, and they compose (plan §0.4): a promptable model's masks become
+scribble training labels via :func:`~spyde.particles.scribble.masks_to_labels`.
+Everything after the first box is written once, in
+:mod:`~spyde.particles.classical` (the instance split) and
+:mod:`~spyde.particles.measure`.
+
+torch is imported **lazily**, inside the functions that need it — importing this
+package must never pay for CUDA init. :func:`~spyde.particles.features.gpu_available`
+answers the capability question without loading it.
+
+.. note::
+ Importing this package currently costs ~6 s, because ``measure.py`` and
+ ``track.py`` read the column schema from :mod:`spyde.signals.particles`, and
+ importing anything under ``spyde.signals`` executes that package's ``__init__``,
+ which pulls in **hyperspy** (5.5 s of the 6.3 s; torch is NOT involved). That is
+ free in the running app — the backend loads hyperspy at startup regardless — but
+ it does mean a script that only wants image segmentation pays for a signal
+ framework it never touches, which is at odds with the "constructible standalone"
+ contract in :mod:`spyde.signals`. Fixing it means deferring the ``insitu`` import
+ in ``spyde/signals/__init__.py``; left alone deliberately, since that file is
+ shared with other waves.
+"""
+from __future__ import annotations
+
+from spyde.particles.classical import (
+ THRESHOLD_METHODS,
+ SegmentParams,
+ segment_frame,
+ split_instances,
+ threshold_mask,
+)
+from spyde.particles.features import (
+ DEFAULT_RANK_RADII,
+ DEFAULT_SIGMAS,
+ FeatureSpec,
+ PreparedFrame,
+ feature_names,
+ feature_stack,
+ feature_tensor,
+ gpu_available,
+ prepare_frame,
+ sample_features,
+ select_device,
+)
+from spyde.particles.measure import measure_frame
+from spyde.particles.scribble import (
+ DEFAULT_CLASSES,
+ UNLABELLED,
+ LabelStore,
+ ScribbleClass,
+ ScribbleClassifier,
+ default_classes,
+ masks_to_labels,
+ random_forest_reference,
+)
+from spyde.particles.track import (
+ EVENT_KINDS,
+ LinkParams,
+ LinkResult,
+ ParticleEvent,
+ event_counts,
+ frame_indices,
+ link,
+ sample_frame_positions,
+)
+
+__all__ = [
+ # classical engine + the shared instance split
+ "SegmentParams",
+ "THRESHOLD_METHODS",
+ "segment_frame",
+ "split_instances",
+ "threshold_mask",
+ # measurement
+ "measure_frame",
+ # feature stack
+ "FeatureSpec",
+ "PreparedFrame",
+ "DEFAULT_SIGMAS",
+ "DEFAULT_RANK_RADII",
+ "prepare_frame",
+ "feature_tensor",
+ "feature_stack",
+ "feature_names",
+ "sample_features",
+ "select_device",
+ "gpu_available",
+ # scribble engine
+ "ScribbleClass",
+ "ScribbleClassifier",
+ "LabelStore",
+ "DEFAULT_CLASSES",
+ "UNLABELLED",
+ "default_classes",
+ "masks_to_labels",
+ "random_forest_reference",
+ # tracking
+ "link",
+ "LinkParams",
+ "LinkResult",
+ "ParticleEvent",
+ "EVENT_KINDS",
+ "event_counts",
+ "frame_indices",
+ "sample_frame_positions",
+]
diff --git a/spyde/particles/batch.py b/spyde/particles/batch.py
new file mode 100644
index 00000000..d9343fcc
--- /dev/null
+++ b/spyde/particles/batch.py
@@ -0,0 +1,803 @@
+"""
+batch.py — whole-movie segmentation, fanned out over the cluster.
+
+The batch run used to be a plain serial ``for t in range(n_frames)`` on ONE
+worker thread: read a frame, segment it, measure it, repeat. On the target
+workload — 900 frames of 4096² — that is 90 minutes with one core busy, the
+other 47 idle and the GPU idle 78% of the time. This module replaces that loop
+with the **dual-lane dispatch SpyDE already uses for find-vectors**, because
+segmentation is the same shape of problem: independent per-frame work, one
+device, many cores.
+
+Why the find-vectors mechanism and not a pipeline
+-------------------------------------------------
+The obvious reading is that the two engines want different parallelism — the
+classical engine is pure scipy so it should fan out, while the scribble engine
+holds ONE torch model on ONE GPU so fanning eight workers at it would mean eight
+CUDA contexts fighting over one device. That reading is wrong, and SpyDE already
+has the answer: :func:`~spyde.actions.find_vectors.gpu_runtime._gpu_task_allowed`.
+
+Dask fans every chunk out over every worker as normal. Each task then asks
+whether **this** worker may touch the GPU: workers named ``"1".."N"`` may, and
+the rest run the CPU path. ``SPYDE_FV_GPU`` (the **GPU feeders** control in the
+DaskMonitor popover, :mod:`spyde.backend.compute_config`) overrides the policy
+for segmentation exactly as it does for find-vectors, and
+:data:`PARTICLE_GPU_LANE_DEFAULT` is the unset-default.
+
+So the two engines do NOT need different architectures. Both go through one
+fan-out; the lane policy decides who uses the GPU, and the classical engine
+simply never asks for it — which means it uses every worker, since a lane it
+never asks about cannot exclude it.
+
+Why the lane default is "one" and NOT the neural path's "4"
+-----------------------------------------------------------
+``default_mode`` is per-CALLER by design, and segmentation's was measured rather
+than inherited. GPU-only dispatch, empty CPU lane, 60 frames of a real 977 x
+4096² movie:
+
+ GPU feeders 1 2 4
+ throughput 0.270 0.252 0.260 frames/s
+
+**Flat.** More processes on the device neither help nor much hurt, because the
+device is not the constraint — see the next section for what is. So
+segmentation keeps ``"one"``, on the same footing as the numba NXCORR kernels:
+nothing is bought by spreading it, and one process is the arrangement that
+cannot go wrong (Windows has no MPS server, and the feature stack sizes each row
+band against *free* VRAM, which several processes then divide — ``GPU_BAND_BYTES``
+records that overshooting there is catastrophic rather than merely slower). An
+explicit ``SPYDE_FV_GPU`` still overrides, which is the user's call.
+
+*History, because the first answer here was wrong and confidently so:* the
+initial measurement said four feeders made a frame **13x** slower (110 s of
+predict+split against 8 s). That run also had five CPU-lane workers each running
+a 48-thread torch predict, so it measured contention and not the lane count. A
+lane number taken from a contended run is exactly the kind of default that then
+looks load-bearing forever. Re-measure with the other lane empty.
+
+Why the CPU lane IS excluded, which the isolated numbers argue against
+----------------------------------------------------------------------
+Neural find-vectors dispatches ``gpu_only=True`` because torch-CPU inference is
+10-50x the GPU batch. Segmentation looks like it should be the exception. One
+4096² scribble predict on CPU:
+
+ torch threads 1 2 4 8 16 48
+ wall 65.8 35.1 18.8 11.2 7.0 9.0 s
+ core-seconds 65.8 70.2 75.2 89.8 111.6 430.5
+
+Against 1.6 s on the GPU that is 41x in core-seconds — but a GPU-lane frame only
+occupies ~4.4 core-seconds of the machine (the split and the measurement, which
+are CPU work on both lanes), so on paper 48 cores of CPU lane are worth about as
+much as the GPU and running both is close to a doubling.
+
+It is not. Measured in the cluster, the CPU lane contributed 0.16 frames/s and
+cost the GPU lane 0.5 — its frames went from 5.8 s to 25-29 s — and the run went
+0.270 -> 0.222 frames/s. The two lanes are not additive because they compete for
+the same cores and the same memory bandwidth. So the CPU lane is **off by
+default** (:func:`cpu_lane_enabled`, ``SPYDE_SEG_CPU_LANE=1`` to restore it),
+and segmentation lands where the neural batch already is.
+
+That thread table is also why every worker pins ``torch.set_num_threads(1)``
+(:func:`_cap_torch_threads`): intra-op threading costs more total work the wider
+it goes, and frame-level parallelism is free.
+
+What used to limit this, and no longer does
+--------------------------------------------
+The honest ceiling WAS ``measure_frame``, and it bound twice: it was most of the
+frame, and it held the **GIL**, so a worker's four task slots were worth one core
+and the effective parallelism of a batch was the WORKER COUNT rather than the
+worker × thread count.
+
+Both halves are now gone, in two passes, and every one of them had to prove
+parity before it was allowed to land (``benchmarks.md`` § "Vectorising
+``measure_frame``"):
+
+* ``regionprops_table`` → :mod:`spyde.particles.props` (every column that is a
+ label-wise reduction) + :mod:`spyde.particles.hull` (``solidity``, a numba
+ convex hull in exact integer arithmetic). 43.7 s → 1.08 s, every column
+ bit-identical.
+* ``_fill_intensity`` → :mod:`spyde.particles.intensity` (``bincount``
+ statistics + a numba kernel for the background ring, which is per-particle and
+ therefore not a partition of the raster). 4.9 s → 0.19 s, all four columns
+ bit-identical on the real 26 566-region frame.
+* ``_contours`` → :mod:`spyde.particles.contours` (marching squares and the
+ segment assembly, every region in one ``prange``). 4.8 s → 0.15 s, and the
+ FILLED polygon — what ``render_frame`` and ``mask_at`` actually consume — is
+ identical on 26 566 of 26 566 regions.
+
+``measure_frame`` is **52.6 s → 1.37 s**, and, because all three replacements are
+numba ``nogil`` kernels and numpy ufuncs, four quadrants in four threads now scale
+**2.82x** where the whole function managed 0.94x with the two Python loops still
+in it. A worker's task slots are finally worth more than one core, so worker ×
+thread is the unit of parallelism here rather than worker count.
+
+Memory safety (CLAUDE.md)
+-------------------------
+Nothing here ever calls ``.compute()`` on the movie. Tasks are sized by BYTES
+(:func:`frames_per_task`) so one task holds a bounded block of frames, and the
+per-frame loop inside a task holds exactly one frame's labels at a time.
+"""
+from __future__ import annotations
+
+import logging
+import os
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+#: Unset-default for the segmentation GPU lane, i.e. how many workers may submit
+#: to the device when ``SPYDE_FV_GPU`` is not set. MEASURED, not inherited from
+#: the neural path's "4" — 1/2/4 feeders are 0.270/0.252/0.260 frames/s on a real
+#: 4096² movie, i.e. flat, so this takes the arrangement that cannot go wrong.
+#: See the module docstring, including why the first measurement said 13×.
+#:
+#: This single value MUST be what both the per-worker gate (:func:`_gpu_allowed`
+#: → ``_gpu_task_allowed``) and the client-side lane split (:func:`_dispatch` →
+#: ``split_workers_for_gpu``) are given. When those two disagree, workers get
+#: chunks they then refuse to run on the GPU — the bug ``gpu_runtime`` records
+#: for the neural path, where the docstring said "2" while the code passed "4".
+PARTICLE_GPU_LANE_DEFAULT = "one"
+
+#: Target bytes of RAW FRAMES per dask task. Not a chunk-alignment number — a
+#: working-set one: a task also holds an int32 label raster (4 bytes/px, i.e.
+#: 64 MB for one 4096² frame) plus the float probability maps, and several tasks
+#: run concurrently per worker. 64 MB of frames keeps a worker's peak near 1 GB.
+#: Override with ``SPYDE_SEG_TASK_BYTES``.
+TASK_BYTES: int = 64 << 20
+
+#: Aim for at least this many tasks per worker, so a short movie still spreads
+#: out instead of collapsing into one task per lane.
+_TASKS_PER_WORKER = 4
+
+
+def cpu_lane_enabled() -> bool:
+ """Whether scribble chunks may also go to the non-GPU workers.
+
+ **Off by default, and that is a measurement, not a guess.** In isolation the
+ CPU lane looks like most of a doubling (the core-second table in the module
+ docstring). In the cluster it is worse than nothing, because the two lanes
+ compete for the same 48 cores and the same memory bandwidth — 60 frames of a
+ real 4096² movie, one GPU worker, 8 CPU workers × 4 slots:
+
+ CPU-lane predict, alone, 1 torch thread 66 s
+ CPU-lane predict, 32 concurrent 88-177 s (~2x, bandwidth)
+ GPU-lane frame, machine otherwise idle 5.8 s
+ GPU-lane frame, CPU lane running 25-29 s (starved)
+
+ The CPU lane contributed 0.16 frames/s and cost the GPU lane 0.5 — the run
+ went from 0.270 frames/s GPU-only to 0.222. So segmentation lands where the
+ neural detector already is (``gpu_only=True`` in orchestrate), for the same
+ underlying reason: a torch path whose CPU cost is disproportionate does not
+ belong in a second lane, it belongs out of the way of the first.
+
+ ``SPYDE_SEG_CPU_LANE=1`` turns it back on — worth it on a machine with more
+ memory bandwidth than GPU, or with several GPUs, neither of which is this
+ one. It has no effect on the classical engine, which has no GPU lane at all
+ and always runs on every worker.
+ """
+ return os.environ.get("SPYDE_SEG_CPU_LANE", "") not in ("", "0", "off")
+
+
+# ── what a worker needs to rebuild the engine ────────────────────────────────
+
+@dataclass(frozen=True)
+class EngineSpec:
+ """A picklable description of the engine, resolved to a callable per worker.
+
+ The trained :class:`~spyde.particles.scribble.ScribbleClassifier` is shipped
+ as a **file path**, not as an object: its weights live on the client's CUDA
+ device, and pickling those would deserialise a CUDA tensor inside every
+ worker process — a context per worker, created behind our back, before the
+ lane policy has had any say. Shipping the path lets each worker load onto
+ the device the lane policy chose for it, which is the same trick the neural
+ detector plays with ``models.get_model(model_id)``.
+
+ Parameters
+ ----------
+ method
+ ``"classical"`` or ``"scribble"``.
+ params
+ ``SegmentParams`` keyword arguments. A plain dict so the spec pickles
+ without importing the segmentation modules on the client.
+ model_path
+ Scribble only: an ``.npz`` written by ``ScribbleClassifier.save``.
+ """
+
+ method: str
+ params: dict = field(default_factory=dict)
+ model_path: str | None = None
+
+ def segment_params(self):
+ from spyde.particles.classical import SegmentParams
+ return SegmentParams(**self.params)
+
+
+#: Per-PROCESS ring of ``(t0, n, device, engine_s, measure_s, block_s)`` records,
+#: one per task. A batch that is fast in isolation and slow in the cluster is
+#: the NORMAL outcome — contention, oversubscription, a rechunk nobody asked for
+#: — and the wall clock alone cannot tell you which. Drained by
+#: :func:`drain_stage_log` (``client.run``), which is how
+#: ``benchmark_particles_batch`` prints its per-lane table. Bounded, so a
+#: 900-frame run cannot grow it without limit.
+_STAGE_LOG: list = []
+_STAGE_LOG_MAX = 4096
+
+#: Per-PROCESS engine cache. Keyed by (model path, mtime, device) so a retrain
+#: (a new temp file) never serves a stale head, and so the GPU and CPU lanes in
+#: one process — which happens in the local thread-pool fallback — keep separate
+#: models rather than moving one back and forth across devices.
+_ENGINE_CACHE: dict[tuple, Any] = {}
+_TORCH_THREADS_SET = [False]
+
+
+def _cap_torch_threads() -> None:
+ """Pin torch to ONE intra-op thread inside a dask worker, once.
+
+ torch defaults to one intra-op thread per LOGICAL CORE, which inside a
+ worker is a lie about how much of the machine this process owns: nine
+ workers × four task slots each spawning 48 OpenMP threads is 1728 threads
+ on 48 cores. But the reason for **one** rather than "the worker's thread
+ budget" is stronger than avoiding oversubscription — measured on a 4096²
+ predict, intra-op threading costs MORE total work the wider it goes (65.8
+ core-seconds at 1 thread, 111.6 at 16, 430.5 at 48). Frame-level
+ parallelism is free and thread-level is not, and dask's task slots already
+ supply the former: one task per core is both the cheapest and the fastest
+ arrangement.
+
+ Deliberately a no-op outside a dask worker — the interactive preview and
+ the training fit are single, latency-sensitive calls that should keep the
+ whole machine.
+ """
+ if _TORCH_THREADS_SET[0]:
+ return
+ try:
+ from distributed import get_worker
+ get_worker()
+ except Exception:
+ return # not on a dask worker: leave torch alone
+ _TORCH_THREADS_SET[0] = True
+ try:
+ import torch
+ torch.set_num_threads(1)
+ except Exception as exc: # pragma: no cover
+ log.debug("[seg-batch] capping torch threads failed: %s", exc)
+
+
+def _warm_measure() -> None:
+ """Compile every ``measure_frame`` kernel BEFORE the first frame is measured.
+
+ Three now: ``solidity``'s convex hull (:mod:`spyde.particles.hull`, which
+ replaced a per-region Qhull call — 30.2 s to 0.18 s at 26 566 regions), the
+ background ring (:mod:`spyde.particles.intensity`) and the outlines
+ (:mod:`spyde.particles.contours`). numba compiles each on first use, which is
+ seconds, and paying that inside the first measured frame of a task makes that
+ frame look like a regression and skews every per-stage number the benchmark
+ prints. ``cache=True`` means they are normally loaded from disk rather than
+ rebuilt. Cheap and idempotent: each module short-circuits once compiled.
+ """
+ try:
+ from spyde.particles.measure import warm_kernels
+ warm_kernels()
+ except Exception as exc: # pragma: no cover
+ log.debug("[seg-batch] measure warmup skipped: %s", exc)
+
+
+def _gpu_allowed() -> bool:
+ """Whether THIS worker may use the GPU, per the shared lane policy.
+
+ Delegates to find-vectors' ``_gpu_task_allowed`` rather than re-deriving the
+ rule: it is the same ``SPYDE_FV_GPU`` setting, surfaced in the same
+ DaskMonitor control, and a second implementation is a second thing to get
+ out of step.
+ """
+ try:
+ from spyde.actions.find_vectors.gpu_runtime import _gpu_task_allowed
+ return bool(_gpu_task_allowed(default_mode=PARTICLE_GPU_LANE_DEFAULT))
+ except Exception as exc: # pragma: no cover
+ log.debug("[seg-batch] GPU lane policy unavailable (%s); using CPU", exc)
+ return False
+
+
+def _gpu_slots():
+ """The device-concurrency semaphore (``SPYDE_FV_GPU_CONC``), or a null
+ context when find-vectors is not importable.
+
+ Bounds how many frames occupy the device at once **per process**. The
+ feature stack sizes each row band at a fraction of *free* VRAM
+ (``features.band_budget_bytes``), and overshooting that is catastrophic
+ rather than merely slow — so unbounded concurrency across a worker's task
+ threads is exactly the allocator thrash ``GPU_BAND_BYTES`` documents.
+ """
+ try:
+ from spyde.actions.find_vectors.gpu_runtime import _gpu_slots as _slots
+ return _slots()
+ except Exception: # pragma: no cover
+ import contextlib
+ return contextlib.nullcontext()
+
+
+def resolve_engine(spec: EngineSpec, *, force_cpu: bool = False):
+ """``(engine, device_str)`` for this worker — the lane decision made once.
+
+ ``engine(frame) -> int32 labels``. The classical engine never asks for the
+ GPU, so it runs on every worker regardless of lane; the scribble engine
+ loads its head onto CUDA/MPS only on a GPU-lane worker and onto the CPU
+ everywhere else.
+ """
+ _cap_torch_threads()
+ _warm_measure()
+ method = str(spec.method).lower()
+ if method == "classical":
+ from spyde.particles.classical import segment_frame
+ sp = spec.segment_params()
+ return (lambda frame: segment_frame(frame, sp)), "cpu"
+
+ if method != "scribble":
+ raise ValueError(
+ f"batch segmentation has no engine for method {spec.method!r}; "
+ "expected 'classical' or 'scribble'")
+ if not spec.model_path:
+ raise ValueError(
+ "the scribble engine needs a trained model — EngineSpec.model_path "
+ "is empty (train the classifier before running the batch)")
+
+ want_gpu = (not force_cpu) and _gpu_allowed()
+ device = None if want_gpu else "cpu"
+ try:
+ mtime = os.path.getmtime(spec.model_path)
+ except OSError:
+ mtime = 0.0
+ key = (spec.model_path, mtime, str(device))
+ clf = _ENGINE_CACHE.get(key)
+ if clf is None:
+ from spyde.particles.scribble import ScribbleClassifier
+ clf = ScribbleClassifier.load(spec.model_path, device=device)
+ _ENGINE_CACHE[key] = clf
+ log.info("[seg-batch] scribble head loaded on %s (gpu lane: %s)",
+ clf.device, want_gpu)
+ sp = spec.segment_params()
+
+ def engine(frame, _clf=clf, _sp=sp):
+ # The device section only — the split that follows is numpy/scipy and
+ # must NOT hold a device slot while it runs (that is what would turn
+ # four GPU feeders back into one).
+ with _gpu_slots():
+ fg, bnd = _clf.predict_foreground_boundary(frame)
+ from spyde.particles.classical import split_instances
+ return split_instances(fg, _sp, boundary=bnd)
+
+ return engine, str(clf.device)
+
+
+# ── the per-task body ────────────────────────────────────────────────────────
+
+def segment_block(block: np.ndarray, t0: int = 0, spec: EngineSpec | None = None,
+ *, scale: float = 1.0, store_masks: bool = True) -> np.ndarray:
+ """Segment and measure every frame of one ``(n, h, w)`` block.
+
+ Returns a 1-D **object** array of ``(rows, contours)``, one entry per frame,
+ which is what the dispatcher assembles. An object payload rather than the
+ NaN-padded fixed array find-vectors uses because a frame's result is doubly
+ ragged: a variable number of particle rows AND a variable-length outline per
+ particle. Padding to a fixed cap would either truncate a busy frame (a real
+ 4096² growth frame here measured **26 566** particles) or waste most of the
+ transfer.
+
+ Runs on a dask worker. Holds ONE frame's labels at a time — the block is
+ already bounded by :func:`frames_per_task`.
+
+ *t0* stamps the ``t`` column. The dispatch path leaves it at 0 and
+ :func:`segment_movie` stamps the true index when the block lands, because a
+ task **cannot** know its own global offset: ``dispatch_chunks`` slices the
+ result array per chunk, and a ``map_blocks`` stage reading
+ ``block_info["array-location"]`` then reports (0, 0) for every one of them —
+ the bug ``orchestrate._do_compute_vectors`` records for the live count map.
+ """
+ from spyde.particles.measure import measure_frame
+
+ block = np.asarray(block)
+ if block.ndim != 3:
+ raise ValueError(f"expected an (n, h, w) block; got shape {block.shape}")
+ n = int(block.shape[0])
+ out = np.empty((n,), dtype=object)
+ if n == 0:
+ return out # dask meta inference calls with size 0
+ if spec is None:
+ raise TypeError("segment_block needs an EngineSpec")
+
+ engine, dev = resolve_engine(spec)
+ t_eng = t_meas = 0.0
+ n_particles = 0
+ t_block = time.perf_counter()
+ for i in range(n):
+ frame = np.asarray(block[i])
+ t_a = time.perf_counter()
+ labels = _engine_with_cpu_fallback(engine, frame, spec)
+ t_b = time.perf_counter()
+ rows, contours = measure_frame(labels, frame, t=int(t0) + i,
+ scale=float(scale))
+ t_c = time.perf_counter()
+ t_eng += t_b - t_a
+ t_meas += t_c - t_b
+ n_particles += len(rows)
+ out[i] = (np.ascontiguousarray(rows, np.float32),
+ (list(contours) if store_masks else []))
+ t_wall = time.perf_counter() - t_block
+ if len(_STAGE_LOG) < _STAGE_LOG_MAX:
+ _STAGE_LOG.append((int(t0), n, dev, t_eng, t_meas, t_wall))
+ log.debug("[seg-batch] block %d..%d on %s: engine %.2fs measure %.2fs "
+ "(block %.2fs, %d particles)", t0, t0 + n, dev, t_eng, t_meas,
+ t_wall, n_particles)
+ return out
+
+
+def drain_stage_log() -> list:
+ """Take and clear this process's per-task stage records. Run via
+ ``client.run(drain_stage_log)`` to see the in-cluster stage costs."""
+ out = list(_STAGE_LOG)
+ _STAGE_LOG.clear()
+ return out
+
+
+def _map_block(block, spec=None, scale=1.0, store_masks=True):
+ """``map_blocks`` entry point — module level so the graph pickles small."""
+ if block.size == 0 or block.shape[0] == 0:
+ return np.empty((0,), dtype=object)
+ return segment_block(block, 0, spec, scale=scale, store_masks=store_masks)
+
+
+def _engine_with_cpu_fallback(engine, frame, spec: EngineSpec):
+ """Run *engine*, retrying the frame on the CPU if the device refuses it.
+
+ Several GPU-lane workers each sizing a feature band against *free* VRAM can
+ collectively overcommit a 12 GB card, and an out-of-memory frame must not
+ take the whole 900-frame run down with it. A silent fallback would hide a
+ throughput problem, so this logs loudly — the same bargain
+ ``_find_vectors_chunk`` makes around its GPU path.
+ """
+ try:
+ return engine(frame)
+ except Exception as exc:
+ if str(spec.method).lower() != "scribble":
+ raise
+ log.warning("[seg-batch] GPU segmentation failed (%r) — retrying this "
+ "frame on the CPU", exc)
+ cpu_engine, _dev = resolve_engine(spec, force_cpu=True)
+ return cpu_engine(frame)
+
+
+# ── task sizing ──────────────────────────────────────────────────────────────
+
+def frames_per_task(n_frames: int, frame_nbytes: int, *, n_workers: int = 1,
+ source_chunk: int | None = None) -> int:
+ """How many frames one dask task should cover.
+
+ Bounded three ways, and each bound is there for a different failure:
+
+ * **by bytes** — a task holds its whole block plus a 4 byte/px label raster
+ per frame in flight; ``TASK_BYTES`` keeps a worker's peak near 1 GB even
+ with four task threads.
+ * **by the source's own chunking** — never larger than one stored nav chunk,
+ so a task never has to pull two chunks to serve one block (CLAUDE.md
+ Live-Display §1: alignment, never a rechunk).
+ * **by the cluster** — at least ``_TASKS_PER_WORKER`` tasks per worker, so a
+ six-frame tutorial movie still spreads across the lanes instead of
+ becoming one task that one worker runs alone.
+ """
+ try:
+ budget = int(os.environ.get("SPYDE_SEG_TASK_BYTES", TASK_BYTES))
+ except ValueError:
+ budget = TASK_BYTES
+ by_bytes = max(1, budget // max(1, int(frame_nbytes)))
+ per = by_bytes
+ if source_chunk:
+ per = min(per, max(1, int(source_chunk)))
+ spread = max(1, int(n_frames) // max(1, _TASKS_PER_WORKER * max(1, n_workers)))
+ return int(max(1, min(per, spread)))
+
+
+def _time_chunks(n_frames: int, per: int) -> tuple[int, ...]:
+ per = max(1, int(per))
+ full, rem = divmod(int(n_frames), per)
+ out = [per] * full
+ if rem:
+ out.append(rem)
+ return tuple(out) or (int(n_frames),)
+
+
+# ── the run ──────────────────────────────────────────────────────────────────
+
+def segment_movie(
+ data,
+ spec: EngineSpec,
+ *,
+ n_frames: int,
+ get_frame: Callable[[int], np.ndarray] | None = None,
+ scale: float = 1.0,
+ store_masks: bool = True,
+ client=None,
+ stopped: list | None = None,
+ on_frames: Callable[[int, int, list], None] | None = None,
+) -> tuple[list[np.ndarray], list[list[np.ndarray]], int] | None:
+ """Segment every frame of *data*, in parallel. The batch run's whole body.
+
+ Parameters
+ ----------
+ data
+ The movie as a ``(n, h, w)`` numpy or dask array (``signal.data``).
+ spec
+ :class:`EngineSpec` — what each worker rebuilds the engine from.
+ n_frames, get_frame
+ The streaming accessor from ``frames_of``. ``get_frame`` is the serial
+ fallback used when there is neither a cluster nor a dask array; the
+ dispatch path never calls it.
+ client
+ A ``distributed.Client``, or None for the local thread-pool fallback
+ (tests, ``SPYDE_NO_DASK=1``, a cluster that never came up).
+ stopped
+ One-element cancel flag, polled between blocks and honoured by the
+ dispatcher.
+ on_frames
+ ``on_frames(t0, t1, results)`` as each block lands, where *results* is
+ the list of ``(rows, contours)`` for frames ``t0:t1``. Called from a
+ worker/callback thread — must not touch figures (CLAUDE.md threading).
+
+ Returns
+ -------
+ ``(per_frame_rows, per_frame_contours, n_done)`` — both lists are exactly
+ *n_frames* long, with EMPTY blocks where a cancelled run never reached, so
+ the CSR store always spans the movie and a stopped run reads as "no
+ particles after frame N" rather than as a shorter movie. A cancelled run
+ still returns everything it did finish; the caller reads ``stopped`` to
+ learn that it was cancelled.
+ """
+ from spyde.signals.particles import COL, N_COLUMNS
+
+ n_frames = int(n_frames)
+ results: list = [None] * n_frames
+ done = [0]
+ col_t = COL["t"]
+
+ def _record(t0: int, block_out) -> None:
+ # THE one place the frame index is written (see `segment_block`): a task
+ # cannot know its own offset, so it stamps 0 and the true index is
+ # applied here, where the dispatcher's global slice is authoritative.
+ vals = list(block_out)
+ for i, v in enumerate(vals):
+ rows = v[0]
+ if len(rows):
+ rows[:, col_t] = float(t0 + i)
+ results[t0 + i] = v
+ done[0] += len(vals)
+ if on_frames is not None:
+ try:
+ on_frames(t0, t0 + len(vals), vals)
+ except Exception as exc:
+ log.debug("[seg-batch] on_frames callback failed: %s", exc)
+
+ lazy = type(data).__module__.startswith("dask.array") if data is not None else False
+ arr_ok = data is not None and getattr(data, "ndim", 0) == 3
+
+ if not (arr_ok and _dispatch(data, spec, n_frames, scale, store_masks,
+ client, stopped, _record, lazy)):
+ _serial(get_frame, spec, n_frames, scale, store_masks, stopped,
+ _record)
+
+ rows_out: list[np.ndarray] = []
+ contours_out: list[list[np.ndarray]] = []
+ for v in results:
+ if v is None:
+ rows_out.append(np.zeros((0, N_COLUMNS), np.float32))
+ contours_out.append([])
+ else:
+ rows_out.append(v[0])
+ contours_out.append(list(v[1]))
+ return rows_out, contours_out, int(done[0])
+
+
+def _n_workers(client) -> int:
+ if client is None:
+ return max(1, (os.cpu_count() or 4) // 4)
+ try:
+ # n_workers=-1: the default TRUNCATES to 5 (CLAUDE.md / benchmarks.md).
+ return max(1, len(client.scheduler_info(n_workers=-1)["workers"]))
+ except Exception:
+ return 1
+
+
+def _dispatch(data, spec, n_frames, scale, store_masks, client, stopped,
+ record, lazy) -> bool:
+ """Build the per-block graph and run it through the dual-lane dispatcher.
+
+ Returns False when this movie cannot be dispatched, so ``segment_movie``
+ falls back to the streaming accessor.
+ """
+ import dask.array as da
+
+ if lazy and len(data.chunks) > 2 and any(len(c) > 1 for c in data.chunks[1:]):
+ # The reader SPLIT the signal axes (RosettaSciIO's balanced-cube
+ # auto-chunk: a real 977 x 4096² MRC arrives as (511, 511, 511)). A task
+ # needs whole frames, and merging those axes while also cutting the time
+ # axis down to one frame is a full P2P rechunk shuffle of the entire
+ # movie — the exact thing CLAUDE.md Live-Display §1 forbids, and it was
+ # measured here doing it: the dispatcher spent 90 s in stall pokes
+ # waiting on `rechunk-merge-rechunk-transfer` before a single frame was
+ # segmented.
+ #
+ # The fix belongs at LOAD time, where it is free — `hs.load(..., lazy=
+ # True, chunks=(1, -1, -1))`, which is what `Session._signal_spanning_
+ # chunks` already does for every movie the app opens. So this is not a
+ # path the app reaches; it is the guard that keeps a hand-built signal
+ # from silently shuffling multiple GB. Fall back to the streaming
+ # accessor: slower, but correct and no worse than the serial loop was.
+ log.warning(
+ "[seg-batch] this movie's chunks %s split the signal axes, so a "
+ "task cannot take whole frames without a full rechunk shuffle — "
+ "segmenting serially instead. Re-load it with "
+ "chunks=(1, -1, -1) to get the parallel path.", data.chunks)
+ return False
+
+ frame_nbytes = int(np.prod(data.shape[1:])) * int(data.dtype.itemsize)
+ n_workers = _n_workers(client)
+ source_chunk = None
+ if lazy:
+ try:
+ source_chunk = int(max(data.chunks[0]))
+ except Exception:
+ source_chunk = None
+ per = frames_per_task(n_frames, frame_nbytes, n_workers=n_workers,
+ source_chunk=source_chunk)
+ chunks = _time_chunks(n_frames, per)
+
+ if lazy:
+ # Only the TIME axis is re-chunked, and only ever to a SMALLER size —
+ # splitting a chunk is a slice, never the multi-GB shuffle CLAUDE.md
+ # forbids. The signal axes are already whole (guarded above).
+ da_data = data if tuple(data.chunks[0]) == chunks \
+ else data.rechunk({0: chunks})
+ else:
+ da_data = da.from_array(data, chunks=(chunks, -1, -1))
+
+ starts = np.concatenate([[0], np.cumsum(da_data.chunks[0])[:-1]]).astype(int)
+
+ import functools
+ block_fn = functools.partial(_map_block, spec=spec, scale=scale,
+ store_masks=store_masks)
+ result_array = da.map_blocks(
+ block_fn, da_data, dtype=object, drop_axis=[1, 2],
+ chunks=(da_data.chunks[0],), meta=np.empty((0,), dtype=object))
+
+ def _assemble(out, nav_slices, chunk_result):
+ out[nav_slices[0]] = chunk_result
+
+ if client is None:
+ _threaded(result_array, starts, stopped, record)
+ return True
+
+ from spyde.compute_dispatch import dispatch_chunks, split_workers_for_gpu
+
+ scribble = str(spec.method).lower() == "scribble"
+ lane_mode = PARTICLE_GPU_LANE_DEFAULT if scribble else "off"
+ gpu_only = False
+ gpu_addrs, cpu_addrs = ([], [])
+ if lane_mode != "off":
+ gpu_addrs, cpu_addrs = split_workers_for_gpu(client, lane_mode)
+ if gpu_addrs and not cpu_lane_enabled():
+ # GPU-ONLY, like the neural batch and for the same measured reason
+ # (see cpu_lane_enabled). The lane list is then pinned HARD below,
+ # because with allow_other_workers a busy GPU lane silently leaks
+ # chunks onto the CPU workers — which is the slow path we just
+ # decided against.
+ cpu_addrs, gpu_only = [], True
+ if not gpu_addrs:
+ # No GPU lane (classical, no CUDA, or SPYDE_FV_GPU=off/all): every
+ # worker is a plain CPU worker. Passing them ALL as the cpu lane is the
+ # established shape for that case (orchestrate._retry_neural_on_cpu),
+ # and keeps one code path instead of two.
+ try:
+ cpu_addrs = list(client.scheduler_info(n_workers=-1)["workers"])
+ except Exception:
+ cpu_addrs = []
+ lane_mode = "off"
+ if not cpu_addrs and not gpu_addrs:
+ _threaded(result_array, starts, stopped, record)
+ return True
+
+ log.info("[seg-batch] %d frames in %d task(s); lanes GPU=%d CPU=%d "
+ "(mode %s%s, method %s)", n_frames, len(da_data.chunks[0]),
+ len(gpu_addrs), len(cpu_addrs), lane_mode,
+ ", gpu-only" if gpu_only else "", spec.method)
+
+ # `record` files every block from the done-callback as it lands, so the
+ # dispatcher's own assembled array is redundant here and is dropped — and
+ # its None-on-cancel return is NOT a "could not dispatch": the run was
+ # handled, it just stopped early, and re-running it serially would be the
+ # opposite of what cancel means.
+ dispatch_chunks(
+ client, result_array, 1, gpu_addrs, cpu_addrs,
+ stopped_flag=stopped, fill_value=None, label="segment",
+ on_chunk_done=lambda sl, blk: record(int(sl[0].start), blk),
+ lane_default_mode=lane_mode, gpu_only=gpu_only, assemble=_assemble,
+ )
+ return True
+
+
+def _threaded(result_array, starts, stopped, record) -> None:
+ """Local fallback: a thread pool over the same blocks.
+
+ Used for tests and ``SPYDE_NO_DASK=1``. Threads and not processes because
+ the point here is only that the fallback is not the retired serial loop —
+ real parallelism is the cluster's job, and the heavy stages (scipy's EDT,
+ watershed, ``ndi.label``, torch) release the GIL.
+ """
+ from concurrent.futures import ThreadPoolExecutor, as_completed
+
+ n_workers = max(1, min(8, (os.cpu_count() or 4) - 1))
+ blocks = []
+ for i, size in enumerate(result_array.chunks[0]):
+ t0 = int(starts[i])
+ blocks.append((t0, result_array[t0:t0 + int(size)]))
+ with ThreadPoolExecutor(max_workers=n_workers,
+ thread_name_prefix="seg-block") as pool:
+ futs = {pool.submit(b.compute, scheduler="threads"): t0
+ for t0, b in blocks}
+ for fut in as_completed(futs):
+ if stopped is not None and stopped[0]:
+ for f in futs:
+ f.cancel()
+ return
+ record(futs[fut], fut.result())
+
+
+def _serial(get_frame, spec, n_frames, scale, store_masks, stopped,
+ record) -> None:
+ """Last resort: the original serial loop, for a frame source that is not an
+ array at all (a callable / a sequence). Kept so ``segment_movie`` is total."""
+ from spyde.particles.measure import measure_frame
+
+ if get_frame is None:
+ raise TypeError("segment_movie needs either a 3-D array or get_frame")
+ engine, _dev = resolve_engine(spec)
+ for t in range(int(n_frames)):
+ if stopped is not None and stopped[0]:
+ return
+ frame = np.asarray(get_frame(t))
+ labels = _engine_with_cpu_fallback(engine, frame, spec)
+ rows, contours = measure_frame(labels, frame, t=t, scale=float(scale))
+ record(t, [(np.ascontiguousarray(rows, np.float32),
+ list(contours) if store_masks else [])])
+
+
+# ── shipping a trained head to the workers ───────────────────────────────────
+
+def save_engine_model(classifier, directory: str | None = None) -> str:
+ """Write *classifier* somewhere every worker can read it, return the path.
+
+ A run-scoped temp file rather than a stable name: a retrain must not be
+ served from a worker's cache under the same key, and two concurrent runs
+ must not overwrite each other's head mid-flight.
+ """
+ import tempfile
+ fd, path = tempfile.mkstemp(prefix="spyde-scribble-", suffix=".npz",
+ dir=directory)
+ os.close(fd)
+ os.unlink(path) # np.savez_compressed appends .npz itself
+ base = path[:-4] if path.endswith(".npz") else path
+ classifier.save(base)
+ out = base + ".npz"
+ log.debug("[seg-batch] scribble head written to %s (%.1f kB)", out,
+ os.path.getsize(out) / 1e3)
+ return out
+
+
+def drop_engine_model(path: str | None) -> None:
+ """Delete a model file written by :func:`save_engine_model`, best effort."""
+ if not path:
+ return
+ try:
+ os.unlink(path)
+ except OSError as exc:
+ log.debug("[seg-batch] removing %s failed: %s", path, exc)
diff --git a/spyde/particles/classical.py b/spyde/particles/classical.py
new file mode 100644
index 00000000..d5530fd2
--- /dev/null
+++ b/spyde/particles/classical.py
@@ -0,0 +1,882 @@
+"""
+classical.py — the always-available segmentation engine, and the SHARED
+instance-split every engine uses.
+
+Two halves, and the split between them is the important part:
+
+* :func:`threshold_mask` / :func:`segment_frame` — the classical engine, a port of
+ ParticleSpy's ``segptcls.process``. Parameter *names* are kept identical so the
+ caret is recognisable to anyone arriving from ParticleSpy.
+* :func:`split_instances` — takes a foreground **probability or mask** and splits
+ it into individual particles. This is used by all three engines (classical,
+ scribble, prompt), which is why it lives here as its own function rather than
+ inside the classical pipeline. It is the only place in the package that imports
+ skimage's segmentation machinery.
+
+Two routes through the split, and which one runs is the performance story
+-------------------------------------------------------------------------
+``split_instances`` has a **watershed** route and a **boundary** route.
+
+The watershed route is geometry-only: it needs a distance transform to find
+markers and to give the flood an elevation. That geometry — the transform, the
+marker/elevation upsample and the flood itself — is **1.62 s of a 1.78 s split at
+4096², and that split is most of a 2.7 s frame** (``benchmarks.md``). Every engine
+funnels into it, so no amount of work on any one engine changes what a big frame
+costs. (Do NOT go optimising the distance transform on the strength of this: since
+``split_decimation`` shipped it is 4% of the frame, and the cost moved to the
+upsample and the flood. ``benchmarks.md`` has the post-decimation breakdown.)
+
+The boundary route is taken when the caller can supply a **boundary mask** — the
+third class in the ilastik convention (particle / background / boundary). A
+classifier told where the joins are hands back particles that are already
+separated, so instances are plain connected components and **neither the distance
+transform nor the watershed runs at all**. Only the scribble engine can produce
+that mask, which is why the boundary class lives in
+:mod:`~spyde.particles.scribble` and the route it unlocks lives here.
+
+Sensitivity is one axis, deliberately
+--------------------------------------
+Plan §0.9 makes detection sensitivity the priority over instance splitting, and
+notes that sensitivity and separation trade off against each other — a threshold
+loose enough to catch a faint particle also merges neighbours. So
+:attr:`SegmentParams.sensitivity` is a single 0..1 control that biases the
+threshold, and the splitting parameters are secondary. Exposing "threshold offset"
+and "split aggressiveness" as independent knobs invites the user to chase one with
+the other.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any
+
+import numpy as np
+
+#: Thresholding methods, matching ParticleSpy's vocabulary.
+THRESHOLD_METHODS: tuple[str, ...] = (
+ "otsu", "mean", "minimum", "yen", "isodata", "li",
+ "local", "local_otsu", "niblack", "sauvola",
+)
+
+_LOCAL_METHODS = frozenset({"local", "local_otsu", "niblack", "sauvola"})
+
+
+@dataclass
+class SegmentParams:
+ """Classical-pipeline parameters. ParticleSpy names where they correspond."""
+
+ threshold: str = "otsu"
+ #: 0..1. 0.5 is the method's own threshold; >0.5 is more sensitive (lower
+ #: threshold, catches fainter particles, merges more); <0.5 is stricter.
+ sensitivity: float = 0.5
+ rb_kernel: int = 0 # rolling-ball radius, 0 = off
+ gaussian: float = 0.0 # pre-blur sigma, 0 = off
+ invert: bool = False # dark particles on a bright background
+ local_size: int = 31 # window for the local threshold methods (odd)
+ watershed: bool = True # split touching particles
+ #: Minimum distance between watershed markers, px. **This replaces
+ #: ParticleSpy's ``watershed_size``**, whose semantics do not transfer: it
+ #: filtered markers by AREA, which works for their thresholded-distance
+ #: markers but silently deletes every small particle when markers are local
+ #: maxima (a 3x3 particle's marker is one pixel, so any area floor erases it).
+ #: That is precisely the sensitivity failure plan §0.9 exists to prevent.
+ min_separation: int = 3
+ #: Instances whose gap is smaller than this many PIXELS are merged back into
+ #: one particle (:func:`merge_close_instances`). This is the control for
+ #: "some of these should be one particle" -- a real particle comes back as
+ #: several instances when the watershed splits a lumpy blob or a ragged edge
+ #: sheds a fragment, and both are answered by the same question: how far
+ #: apart must two pieces be to count as separate? 0 disables.
+ merge_distance: float = 0.0
+ #: Gaussian sigma applied to the distance transform BEFORE peak finding.
+ #: Suppresses spurious maxima from a ragged boundary without merging genuinely
+ #: separate particles. 0 disables.
+ marker_smooth: float = 1.0
+ watershed_erosion: int = 0 # erosions before the distance transform
+ #: Cap on the passes of label growth that reclaim the BOUNDARY ring back into
+ #: the instances either side of it, on the boundary split path only (see
+ #: :func:`_split_by_boundary`). **0 = grow until nothing more can be
+ #: assigned**, which is the default and the only setting that reproduces the
+ #: watershed's areas.
+ #:
+ #: Every boundary pixel is real particle, so ``ndi.label(fg & ~boundary)`` on
+ #: its own under-reports area by the seam's width. Measured against the
+ #: watershed on touching 44 px discs, median area error by seam width and
+ #: pass count::
+ #:
+ #: seam width 1 px 2 px 4 px 8 px
+ #: 1 pass 0.0% +0.0% -1.2% -4.0%
+ #: 2 passes 0.0% 0.0% +0.0% -2.7%
+ #: converged 0.0% 0.0% 0.0% 0.0%
+ #:
+ #: A pass grows every label by one pixel, so it takes ``ceil(width/2)`` of
+ #: them for the two sides to meet in the middle — and converging costs
+ #: nothing, because each pass only visits the pixels still unassigned.
+ #: **The particle COUNT is correct at every setting**, including 1; only the
+ #: areas move, which is why this is a cap and not a correctness switch.
+ boundary_reclaim: int = 0
+ #: Grid decimation for the SPLIT geometry only. 0 = auto (see
+ #: :func:`_split_factor`), 1 = never decimate.
+ #:
+ #: The distance transform **was 61% of a 4096² segmentation** before this
+ #: shipped (3.93 s of 6.40 s measured; watershed itself only 9%), and it is
+ #: used for two things — finding markers and supplying the watershed's
+ #: elevation — neither of which needs full resolution. Computing it on a
+ #: decimated grid and
+ #: upsampling the elevation gives **2.8–2.9× on 4k with identical particle
+ #: counts and identical median areas** (0.0% difference, on both touching and
+ #: isolated fields).
+ #:
+ #: This does NOT weaken detection, and the distinction matters: the threshold
+ #: still runs at full resolution, so *which* bodies are found is unchanged —
+ #: plan §0.9's faint-particle sensitivity is untouched. Only the cut BETWEEN
+ #: two touching bodies moves, by about ``factor`` pixels.
+ split_decimation: int = 0
+ min_size: int = 20 # discard instances smaller than this, px
+ max_size: int = 0 # discard instances larger than this, px; 0 = off
+ clear_border: bool = False # drop instances touching the frame edge
+ extra: dict[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ if self.threshold not in THRESHOLD_METHODS:
+ raise ValueError(
+ f"unknown threshold {self.threshold!r}; expected one of "
+ f"{', '.join(THRESHOLD_METHODS)}"
+ )
+ if not 0.0 <= self.sensitivity <= 1.0:
+ raise ValueError(f"sensitivity must be in 0..1; got {self.sensitivity}")
+ if self.threshold in _LOCAL_METHODS and self.local_size % 2 == 0:
+ # skimage requires an odd window; silently bumping it would make the
+ # caret's number disagree with what actually ran.
+ raise ValueError(
+ f"local_size must be odd for threshold={self.threshold!r}; "
+ f"got {self.local_size}"
+ )
+
+
+# ── preprocessing ────────────────────────────────────────────────────────────
+
+def _rolling_ball(img: np.ndarray, radius: int) -> np.ndarray:
+ """Background flatten via white top-hat, as ParticleSpy does it.
+
+ Note this is the morphological top-hat, not skimage's newer
+ ``restoration.rolling_ball``. Keeping ParticleSpy's actual operation matters
+ for the parity gate — the two give visibly different backgrounds on a sloping
+ carbon film.
+ """
+ from skimage.morphology import square, white_tophat
+ if radius <= 0:
+ return img
+ return white_tophat(img, footprint=square(int(radius)))
+
+
+def _prepare(frame: np.ndarray, p: SegmentParams) -> np.ndarray:
+ """Rolling-ball → gaussian → invert, returning float32.
+
+ NaN is filled BEFORE filtering. A drift-corrected frame carries a NaN border
+ (``spyde.drift.warp``), and every skimage filter propagates NaN outward,
+ which would erase a band of real data around the edge.
+
+ **The fill polarity follows ``invert``.** The padding has to read as
+ background in the image that is finally THRESHOLDED, and `invert` maps
+ ``x -> -x`` after this fill — so filling with the finite minimum when
+ inverting makes the border the brightest thing in the frame, and it
+ segments as one enormous particle hugging the edge. Measured on a 96²
+ frame with a 12 px NaN border: a 240 px instance on the padding, with the
+ border reading as the frame maximum.
+
+ That is the failure ``spyde.drift.warp`` names as the most likely
+ integration bug in the whole feature ("a threshold applied to NaN … invents
+ a large 'particle' along the edge that then nucleates a spurious track"),
+ so it is fixed here rather than papered over downstream. Filling with the
+ finite MAXIMUM when ``invert`` is set leaves the padding at the minimum
+ after inversion, which is what the rest of this function already assumes.
+ """
+ from scipy.ndimage import gaussian_filter
+ from skimage.util import invert as sk_invert
+
+ img = np.asarray(frame, dtype=np.float32)
+ bad = ~np.isfinite(img)
+ if bad.any():
+ finite = img[~bad]
+ img = img.copy()
+ if finite.size:
+ img[bad] = finite.max() if p.invert else finite.min()
+ else:
+ img[bad] = 0.0
+
+ if p.rb_kernel > 0:
+ img = _rolling_ball(img, p.rb_kernel)
+ if p.gaussian > 0:
+ img = gaussian_filter(img, float(p.gaussian))
+ if p.invert:
+ # sk_invert on a float image maps x -> -x, which is all thresholding needs.
+ img = sk_invert(img)
+ return img
+
+
+# ── thresholding ─────────────────────────────────────────────────────────────
+
+def _sensitivity_offset(img: np.ndarray, sensitivity: float) -> float:
+ """Convert 0..1 sensitivity into an additive threshold offset.
+
+ Scaled by the image's own robust spread (5-95 percentile), so the control
+ behaves the same on a uint16 frame and a normalised float one. At 0.5 the
+ offset is exactly zero, i.e. the method's own threshold is used unmodified —
+ which keeps the default path bit-identical to plain Otsu and makes the parity
+ test against ParticleSpy meaningful.
+ """
+ if sensitivity == 0.5:
+ return 0.0
+ finite = img[np.isfinite(img)]
+ if finite.size == 0:
+ return 0.0
+ lo, hi = np.percentile(finite, [5, 95])
+ spread = float(hi - lo)
+ if spread <= 0:
+ return 0.0
+ # sensitivity 1.0 lowers the threshold by half the spread; 0.0 raises it.
+ return -(float(sensitivity) - 0.5) * spread
+
+
+def threshold_mask(img: np.ndarray, p: SegmentParams) -> np.ndarray:
+ """Boolean foreground mask for a prepared image.
+
+ Raises
+ ------
+ ValueError
+ If the chosen method cannot be computed on this image. Some methods
+ genuinely have preconditions — ``minimum`` needs a bimodal histogram and
+ skimage raises ``RuntimeError`` when it cannot find two maxima, which
+ happens on a sparse field of small bright particles (nearly all background,
+ so the histogram is one spike). Re-raised here with the method named and a
+ suggestion, because the bare skimage error gives the user nothing to act on.
+ """
+ from skimage import filters
+
+ offset = _sensitivity_offset(img, p.sensitivity)
+
+ try:
+ return _apply_threshold(img, p, offset)
+ except RuntimeError as exc:
+ raise ValueError(
+ f"threshold method {p.threshold!r} failed on this frame ({exc}). "
+ "It requires a clearly bimodal intensity histogram; a sparse field of "
+ "small particles does not have one. Try 'otsu', 'yen' or 'li', or a "
+ "local method such as 'sauvola'."
+ ) from exc
+
+
+def _apply_threshold(img: np.ndarray, p: SegmentParams, offset: float) -> np.ndarray:
+ from skimage import filters
+
+ if p.threshold == "otsu":
+ t = filters.threshold_otsu(img)
+ elif p.threshold == "mean":
+ t = filters.threshold_mean(img)
+ elif p.threshold == "minimum":
+ t = filters.threshold_minimum(img)
+ elif p.threshold == "yen":
+ t = filters.threshold_yen(img)
+ elif p.threshold == "isodata":
+ t = filters.threshold_isodata(img)
+ elif p.threshold == "li":
+ t = filters.threshold_li(img)
+ elif p.threshold == "local":
+ t = filters.threshold_local(img, block_size=p.local_size)
+ elif p.threshold == "niblack":
+ t = filters.threshold_niblack(img, window_size=p.local_size)
+ elif p.threshold == "sauvola":
+ t = filters.threshold_sauvola(img, window_size=p.local_size)
+ elif p.threshold == "local_otsu":
+ from skimage.filters.rank import otsu as rank_otsu
+ from skimage.morphology import disk
+ from skimage.util import img_as_ubyte
+ lo, hi = float(np.nanmin(img)), float(np.nanmax(img))
+ norm = np.zeros_like(img) if hi <= lo else (img - lo) / (hi - lo)
+ t_u8 = rank_otsu(img_as_ubyte(norm), disk(max(1, p.local_size // 2)))
+ t = lo + (t_u8.astype(np.float32) / 255.0) * (hi - lo)
+ else: # pragma: no cover — guarded above
+ raise ValueError(f"unknown threshold {p.threshold!r}")
+
+ return img > (np.asarray(t, dtype=np.float32) + np.float32(offset))
+
+
+# ── the shared instance split ────────────────────────────────────────────────
+
+def split_instances(
+ foreground: np.ndarray,
+ p: SegmentParams,
+ *,
+ boundary: np.ndarray | None = None,
+ distance_from: np.ndarray | None = None,
+) -> np.ndarray:
+ """Split a foreground mask (or probability map) into labelled instances.
+
+ **Shared by every engine.** ``foreground`` may be boolean, or a float
+ probability in 0..1 (thresholded at 0.5), which is what the scribble and
+ prompt engines produce.
+
+ Parameters
+ ----------
+ boundary
+ Optional ``(h, w)`` mask or probability of the **inter-particle
+ boundary** — the ilastik third class. When given (and non-empty) the
+ instances come from plain connected components of ``foreground &
+ ~boundary`` and **the distance transform and watershed never run**; see
+ :func:`_split_by_boundary` for why that is the whole point.
+ distance_from
+ Optional alternative to the binary distance transform for seeding
+ watershed markers — e.g. a learned boundary-class probability. Passing the
+ classifier's own notion of "interior" separates touching particles better
+ than geometry does, which is why the hook exists. Ignored when
+ *boundary* is supplied, since that path has no markers to seed.
+
+ Returns
+ -------
+ ``(h, w)`` int32 label image, relabelled 1..n with no gaps.
+ """
+ from scipy import ndimage as ndi
+ from skimage.morphology import binary_erosion, disk
+ from skimage.segmentation import clear_border, watershed
+
+ fg = np.asarray(foreground)
+ if fg.ndim != 2:
+ raise ValueError(f"foreground must be 2-D; got shape {fg.shape}")
+ if fg.dtype != bool:
+ fg = fg > 0.5
+
+ bnd = _as_boundary(boundary, fg.shape)
+ if bnd is not None:
+ # The whole reason the boundary class exists — no EDT, no watershed.
+ return _finalize_labels(_split_by_boundary(fg, bnd, p), p)
+
+ if p.min_size > 0 and fg.any():
+ # Own implementation rather than skimage's `remove_small_objects`: that
+ # function is mid-deprecation (its replacement removes objects smaller
+ # than OR EQUAL to the threshold, a silent off-by-one against the
+ # documented `min_size` meaning), and we already need `_drop_small` for
+ # the post-watershed pass.
+ lab0, _ = ndi.label(fg)
+ fg = _drop_small(lab0, int(p.min_size)) > 0
+
+ if not fg.any():
+ return np.zeros(fg.shape, dtype=np.int32)
+
+ if p.watershed and fg.size >= _COMPONENT_ROUTE_PX:
+ # Big frame: watershed each connected component in its own bbox. Exact
+ # (a component is surrounded by background, so nothing crosses a crop)
+ # and it is what keeps a dask worker off its memory limit — the
+ # whole-frame route below allocates several full-frame rasters, 546 MB
+ # of the 852 MB peak measured at 4096². See _watershed_per_component.
+ labels = _watershed_per_component(fg, p, distance_from)
+ elif p.watershed:
+ seed_src = fg
+ if p.watershed_erosion > 0:
+ seed_src = fg.copy()
+ for _ in range(int(p.watershed_erosion)):
+ seed_src = binary_erosion(seed_src, disk(1))
+
+ if distance_from is not None:
+ dist = np.asarray(distance_from, dtype=np.float32) * seed_src
+ markers = _distance_markers(dist, fg, p)
+ else:
+ dist, markers = _distance_and_markers(seed_src, fg, p)
+
+ if markers.max() > 0:
+ labels = watershed(-dist, markers, mask=fg)
+ else:
+ labels, _ = ndi.label(fg)
+ else:
+ labels, _ = ndi.label(fg)
+
+ labels = np.asarray(labels, dtype=np.int32)
+
+ if p.clear_border:
+ labels = clear_border(labels)
+
+ return _finalize_labels(labels, p)
+
+
+# ── the boundary split: connected components instead of a watershed ──────────
+
+def _as_boundary(boundary, shape: tuple[int, int]) -> np.ndarray | None:
+ """Coerce a boundary argument to a boolean mask, or None if there is none.
+
+ An all-False boundary is treated as **absent**, not as "a boundary with no
+ pixels". That is what makes the wizard's automatic switch safe: a user who
+ has added the boundary class but not painted it yet gets the watershed route
+ rather than a silent downgrade to unsplit connected components.
+ """
+ if boundary is None:
+ return None
+ b = np.asarray(boundary)
+ if b.ndim != 2:
+ raise ValueError(f"boundary must be 2-D; got shape {b.shape}")
+ if tuple(b.shape) != tuple(shape):
+ raise ValueError(
+ f"boundary is {b.shape} but foreground is {shape} — they must "
+ "describe the same frame")
+ if b.dtype != bool:
+ b = b > 0.5
+ return b if b.any() else None
+
+
+def _split_by_boundary(fg: np.ndarray, bnd: np.ndarray,
+ p: SegmentParams) -> np.ndarray:
+ """Instances from ``fg & ~bnd`` — the reason the boundary class exists.
+
+ The watershed route costs a **global** distance transform, a marker/elevation
+ upsample and a flood — together 1.62 s of a 1.78 s split at 4096²
+ (``benchmarks.md``). All of it is doing geometry's best guess at a question
+ the classifier can simply be *told*: the
+ ilastik convention paints particle / background / **boundary**, and a head
+ trained on boundaries hands back touching particles already separated. Then
+ instance extraction is one ``ndi.label`` and both the EDT and the watershed
+ are skipped outright.
+
+ Two things this does that a bare ``ndi.label(fg & ~bnd)`` does not:
+
+ * **The ring is given back.** Every boundary pixel is real particle, so
+ labelling only the cores under-reports area by the ring's width — 1.2% on a
+ 2 px seam, 5.2% on an 8 px one, measured against the watershed on touching
+ 44 px discs. :func:`_reclaim_boundary` grows the cores back out into it,
+ recovering the same total foreground the watershed claimed, pixel for
+ pixel. Where the two routes *cut* can still differ by a pixel or two —
+ watershed cuts equidistant from two markers, growth breaks a tie toward
+ the higher label id — which measured as 0.0% on those 44 px discs and 0.7%
+ on a tighter 14 px pair.
+ * **A component with no boundary through it is untouched.** An isolated
+ particle has no boundary pixels, so its core *is* the whole particle and it
+ passes through unchanged. The boundary path is therefore not a different
+ answer for isolated bodies, only for touching ones.
+
+ The pre-``min_size`` filter the watershed route runs is deliberately skipped:
+ it exists to keep specks out of the watershed, and there is no watershed
+ here. :func:`_finalize_labels` applies the same floor to the result.
+ """
+ from scipy import ndimage as ndi
+
+ core = fg & ~bnd
+ if not core.any():
+ return np.zeros(fg.shape, dtype=np.int32)
+ labels, _n = ndi.label(core)
+ labels = np.asarray(labels, dtype=np.int32)
+ return _reclaim_boundary(labels, fg, int(p.boundary_reclaim))
+
+
+#: 8-neighbour offsets, as (dy, dx). 8 and not 4 so a diagonal step across a
+#: thin boundary still reaches the core on the other side.
+_NEIGHBOURS = ((-1, -1), (-1, 0), (-1, 1), (0, -1),
+ (0, 1), (1, -1), (1, 0), (1, 1))
+
+
+def _reclaim_boundary(labels: np.ndarray, fg: np.ndarray,
+ max_passes: int) -> np.ndarray:
+ """Grow *labels* out into the unlabelled foreground until it is all claimed.
+
+ Only the **unassigned foreground** pixels are ever visited — the boundary
+ ring, ~1% of a frame — so this is eight gathers over a shrinking index list
+ rather than a morphological pass over the raster. A full-frame
+ ``grey_dilation`` would be ~200 ms *per pass* at 4096²; the whole convergence
+ there measures **82 ms**.
+
+ The update is **synchronous** (Jacobi): every pixel's new label is computed
+ from the state at the start of the pass and written afterwards, so one pass
+ grows every label by exactly one pixel and the result does not depend on
+ array order. A pixel reachable from two instances at the same distance goes
+ to the higher label id — an arbitrary but *deterministic* tie-break, which is
+ the property that matters (the alternative, nearest-by-Euclidean-distance,
+ is a distance transform, and avoiding that is the entire point of this path).
+
+ Termination is guaranteed without a cap: a pixel leaves the list as soon as
+ it is claimed and never returns, and the loop stops the first time a whole
+ pass claims nothing — which is exactly when the pixels left over are
+ unreachable (an island of foreground with no core of its own, entirely
+ fenced in by boundary). Those stay 0, and that is honest: they belong to no
+ instance, and inventing an owner would be worse than leaving them out.
+
+ *max_passes* is a cap for callers who want the seam left partly unassigned;
+ 0 means run to convergence.
+ """
+ if max_passes < 0:
+ return labels
+ h, w = labels.shape
+ todo_y, todo_x = np.nonzero(fg & (labels == 0))
+ if not todo_y.size:
+ return labels
+
+ # Pad by one so a neighbour offset can never wrap around a row edge; the
+ # border ring stays 0, so it contributes nothing to the max.
+ pw = w + 2
+ padded = np.zeros((h + 2, pw), dtype=labels.dtype)
+ padded[1:-1, 1:-1] = labels
+ flat = padded.reshape(-1)
+ todo = (todo_y + 1).astype(np.int64) * pw + (todo_x + 1)
+ offsets = [dy * pw + dx for dy, dx in _NEIGHBOURS]
+
+ passes = 0
+ while todo.size:
+ if max_passes and passes >= int(max_passes):
+ break
+ best = flat[todo + offsets[0]]
+ for off in offsets[1:]:
+ best = np.maximum(best, flat[todo + off])
+ won = best > 0
+ if not won.any():
+ break # nothing adjacent to a label; never will be
+ flat[todo[won]] = best[won] # written AFTER the whole pass — see above
+ todo = todo[~won]
+ passes += 1
+ return padded[1:-1, 1:-1]
+
+
+# ── the size filter and the sequential relabel, fused ────────────────────────
+
+def _finalize_labels(labels: np.ndarray, p: SegmentParams) -> np.ndarray:
+ """Apply ``min_size``/``max_size`` and renumber to 1..n, in ONE pass.
+
+ Equivalent to ``_relabel_sequential(_drop_small(_drop_large(labels)))`` and
+ bit-identical to it (a test pins that), but it reads the label raster twice
+ instead of six times. At 4096² the old chain was **302 ms** — 93 ms of
+ ``bincount`` for ``_drop_small``, its ``np.isin`` write-back, then
+ ``_relabel_sequential``'s ``np.unique`` (a full 16.7 M-element sort, 139 ms
+ on its own) and a second gather. Here one ``bincount`` decides every label's
+ fate and one LUT gather writes the final ids: **164 ms**.
+
+ Dropping and renumbering cannot be separated without paying for the raster
+ twice, which is why they are one function rather than two composed ones.
+ """
+ labels = np.asarray(labels)
+ # MERGE first, size-filter second. Two fragments of one particle are each
+ # small; merged they are one particle of the real size, and filtering before
+ # merging would delete the pieces and leave nothing to join. This is also why
+ # it lives here rather than in either split path: every engine reaches
+ # `_finalize_labels`, so the control means the same thing in all of them.
+ labels = merge_close_instances(labels, getattr(p, "merge_distance", 0.0))
+ counts = np.bincount(labels.ravel())
+ # `counts > 0` and not `slice(1, None)`: a label id absent from the raster
+ # (a gap left by clear_border) must not be handed an output number, which is
+ # exactly what `np.unique` gave for free and a bare size test would not.
+ keep = counts > 0
+ keep[0] = False # 0 is background, never an id
+ if p.min_size > 0:
+ keep &= counts >= int(p.min_size)
+ if p.max_size > 0:
+ keep &= counts <= int(p.max_size)
+
+ n = int(keep.sum())
+ lut = np.zeros(counts.size, dtype=np.int32)
+ if n:
+ # Ascending original id, which is the order `np.unique` produced.
+ lut[keep] = np.arange(1, n + 1, dtype=np.int32)
+ return lut[labels]
+
+
+#: Pixel count above which the split geometry decimates by default. Below roughly
+#: this the distance transform is already sub-100 ms and decimating buys nothing
+#: worth the (small) boundary shift.
+_SPLIT_DECIMATE_ABOVE = 2 * 1024 * 1024
+
+
+def _split_factor(shape: tuple[int, int], p: SegmentParams) -> int:
+ """Decimation factor for the split geometry. See ``split_decimation``."""
+ if p.split_decimation:
+ return max(1, int(p.split_decimation))
+ n = int(shape[0]) * int(shape[1])
+ if n <= _SPLIT_DECIMATE_ABOVE:
+ return 1
+ # One step per 4x in area, capped: past 4 the markers of a small particle
+ # start to merge, and the whole point is that detection is unaffected.
+ return 2 if n <= 4 * _SPLIT_DECIMATE_ABOVE else 4
+
+
+def _distance_and_markers(seed_src: np.ndarray, fg: np.ndarray,
+ p: SegmentParams) -> tuple[np.ndarray, np.ndarray]:
+ """``(elevation, markers)`` for the watershed, decimating when it pays.
+
+ The distance transform used to dominate a large-frame segmentation — 61% of a
+ 4096² run — and is needed only to seed markers and to give watershed an
+ elevation. Neither wants full resolution, so above a threshold both are
+ computed on a decimated grid and the elevation is bilinearly upsampled (and
+ rescaled by the factor, so it stays in pixel units). That is what makes the
+ 61% historical: the transform is now 4% of the frame and the cost sits in the
+ upsample and the flood.
+
+ Measured on 4096²: 5.7 s → 2.0 s, with the SAME particle count and the same
+ median area to 0.0%.
+ """
+ from scipy import ndimage as ndi
+
+ factor = _split_factor(fg.shape, p)
+ if factor <= 1:
+ dist = ndi.distance_transform_edt(seed_src)
+ return dist, _distance_markers(dist, fg, p)
+
+ small = seed_src[::factor, ::factor]
+ d_small = ndi.distance_transform_edt(small).astype(np.float32)
+ markers_small = _distance_markers(d_small, small, p, factor=factor)
+
+ # Nearest-neighbour for the MARKERS (a label must not be interpolated into a
+ # value that names a different particle) and bilinear for the ELEVATION.
+ h, w = fg.shape
+ markers = np.repeat(np.repeat(markers_small, factor, 0), factor, 1)[:h, :w]
+ markers = markers * fg
+
+ from scipy.ndimage import zoom
+ elev = zoom(d_small, factor, order=1, grid_mode=True, mode="nearest")
+ elev = elev[:h, :w] * float(factor) # back into pixel units
+ if elev.shape != fg.shape: # odd sizes: pad the last row/col
+ pad = np.zeros(fg.shape, np.float32)
+ pad[:elev.shape[0], :elev.shape[1]] = elev
+ elev = pad
+ return elev, np.asarray(markers, dtype=np.int32)
+
+
+def _distance_markers(dist: np.ndarray, fg: np.ndarray,
+ p: SegmentParams, factor: int = 1) -> np.ndarray:
+ """One marker per particle, from the local maxima of the distance transform.
+
+ Two failure modes have to be avoided at once, and they pull in opposite
+ directions:
+
+ * **Thresholding the distance map merges neighbours.** Taking connected
+ components of ``dist > k`` looks appealing and is wrong: two discs whose
+ edges overlap have ``dist > k`` everywhere in the join, so the union is one
+ connected core and watershed is handed a single marker — it then cannot
+ split anything. This was the original implementation here and the touching-
+ discs test caught it.
+ * **Raw peak maxima over-split a round particle.** A disc's distance maximum
+ is a flat plateau, so a maximum filter marks every plateau pixel; treated as
+ separate markers, one disc becomes a pie chart of wedges.
+
+ ``peak_local_max`` followed by ``ndi.label`` resolves both: the plateau's
+ pixels are contiguous, so labelling collapses them into ONE marker, while two
+ genuinely separate maxima stay separate. Smoothing the distance map first
+ removes the boundary-roughness maxima that would otherwise fragment an
+ irregular particle.
+
+ No marker is ever dropped for being small — see
+ :attr:`SegmentParams.min_separation`.
+ """
+ from scipy import ndimage as ndi
+ from skimage.feature import peak_local_max
+
+ if dist.max() <= 0:
+ return np.zeros(fg.shape, dtype=np.int32)
+
+ dist_pk = dist
+ if p.marker_smooth > 0:
+ from scipy.ndimage import gaussian_filter
+ dist_pk = gaussian_filter(dist.astype(np.float32), float(p.marker_smooth))
+
+ # `min_separation` is in FULL-frame pixels, so on a decimated grid it has to
+ # come down by the same factor or a 3 px separation becomes an effective 6
+ # and two close particles merge into one marker.
+ coords = peak_local_max(
+ dist_pk,
+ min_distance=max(1, int(p.min_separation) // max(1, int(factor))),
+ labels=fg,
+ exclude_border=False,
+ )
+ peaks = np.zeros(fg.shape, dtype=bool)
+ if len(coords):
+ peaks[tuple(coords.T)] = True
+ markers, _ = ndi.label(peaks)
+ return np.asarray(markers, dtype=np.int32)
+
+
+def _drop_small(labels: np.ndarray, min_size: int) -> np.ndarray:
+ counts = np.bincount(labels.ravel())
+ bad = np.flatnonzero(counts < min_size)
+ bad = bad[bad > 0]
+ if bad.size:
+ labels = labels.copy()
+ labels[np.isin(labels, bad)] = 0
+ return labels
+
+
+def _drop_large(labels: np.ndarray, max_size: int) -> np.ndarray:
+ """Unused on the hot path — :func:`_finalize_labels` folds this in. Kept as
+ the readable reference that test pins ``_finalize_labels`` against."""
+ counts = np.bincount(labels.ravel())
+ bad = np.flatnonzero(counts > max_size)
+ bad = bad[bad > 0]
+ if bad.size:
+ labels = labels.copy()
+ labels[np.isin(labels, bad)] = 0
+ return labels
+
+
+def _relabel_sequential(labels: np.ndarray) -> np.ndarray:
+ """Renumber to 1..n with no gaps, so ``label`` is a usable per-frame index.
+
+ Also off the hot path now (see :func:`_finalize_labels`) and kept for the
+ same reason: it is the obvious-and-correct version the fused one is checked
+ against.
+ """
+ present = np.unique(labels)
+ present = present[present > 0]
+ if present.size == 0:
+ return np.zeros(labels.shape, dtype=np.int32)
+ lut = np.zeros(int(labels.max()) + 1, dtype=np.int32)
+ lut[present] = np.arange(1, present.size + 1, dtype=np.int32)
+ return lut[labels]
+
+
+# ── the classical engine ─────────────────────────────────────────────────────
+
+def segment_frame(frame: np.ndarray, p: SegmentParams | None = None) -> np.ndarray:
+ """Classical segmentation of one frame → int32 label image.
+
+ ``prepare → threshold → split_instances``. The whole ParticleSpy pipeline,
+ with the instance step factored out so the other two engines share it.
+ """
+ p = p or SegmentParams()
+ prepared = _prepare(frame, p)
+ fg = threshold_mask(prepared, p)
+ return split_instances(fg, p)
+
+
+def merge_close_instances(labels: np.ndarray, merge_px: float) -> np.ndarray:
+ """Merge instances separated by a gap smaller than *merge_px* pixels.
+
+ The user-facing answer to "some of these should be one particle". A real
+ particle can come back as several instances for two very different reasons
+ — the watershed split a lumpy blob, or a ragged edge broke off a fragment —
+ and BOTH are fixed by the same question: how far apart do two pieces have
+ to be before they are genuinely separate particles? That is a distance the
+ user can see in the image, which is what the abstract 0-1 confidence
+ control was not.
+
+ Implemented by growing every instance by half the gap and taking connected
+ components of the result: two instances whose boundaries come within
+ *merge_px* touch once grown, so they land in one component and are
+ relabelled together. Growing by HALF is what makes the parameter mean the
+ gap between them rather than the radius each one is inflated by.
+
+ The instances themselves are NOT dilated in the output — only their
+ identities change. Areas, centroids and outlines are still measured from
+ the original mask, so merging cannot inflate a measurement.
+
+ ``merge_px <= 0`` returns *labels* unchanged.
+ """
+ if merge_px is None or merge_px <= 0:
+ return labels
+ lab = np.asarray(labels)
+ n = int(lab.max()) if lab.size else 0
+ if n < 2:
+ return lab
+
+ from scipy import ndimage as ndi
+
+ radius = max(1, int(round(float(merge_px) / 2.0)))
+ # A cross/disk structuring element applied `radius` times is far cheaper
+ # than building a (2r+1)² disk and is close enough for a gap test.
+ grown = ndi.binary_dilation(lab > 0, iterations=radius)
+ comp, n_comp = ndi.label(grown)
+ if n_comp >= n:
+ return lab # nothing came together
+
+ # Each original instance lies wholly inside one grown component (it can
+ # only have grown outward), so sampling the component at the instance's
+ # own pixels is exact — take the max for safety on a degenerate label.
+ comp_of = ndi.maximum(comp, lab, index=np.arange(1, n + 1))
+ comp_of = np.asarray(comp_of, np.int64).ravel()
+
+ # Renumber the surviving components 1..k so the output is a dense label
+ # image, which everything downstream (measure, contours) assumes.
+ uniq, dense = np.unique(comp_of, return_inverse=True)
+ lut = np.zeros(n + 1, np.int32)
+ lut[1:] = dense.astype(np.int32) + 1
+ return lut[lab]
+
+
+#: Frames at or above this many pixels take the per-component watershed route
+#: instead of the whole-frame one. Below it the bookkeeping costs more than the
+#: memory it saves; above it the whole-frame route is what pauses a dask worker.
+_COMPONENT_ROUTE_PX: int = 4 << 20 # ~2048², i.e. 4 megapixels
+
+
+def _watershed_per_component(fg: np.ndarray, p: SegmentParams,
+ distance_from: np.ndarray | None = None
+ ) -> np.ndarray:
+ """Watershed each connected component inside its own bounding box.
+
+ **This is exact, not an approximation.** A connected component is
+ surrounded by background by definition, so within its (1 px padded) bbox
+ every quantity the watershed needs is already correct: the distance
+ transform measures distance to the nearest background pixel, and that
+ pixel is inside the pad; markers are local maxima of that distance; and a
+ watershed cannot flow between two components because they do not touch.
+ Cropping therefore changes nothing about the answer, only the peak memory.
+
+ Why it matters. Measured on a 4096² frame, whole-frame classical
+ segmentation peaks at **852 MB**, of which `split_instances` alone is
+ 546 MB — the distance transform, its smoothed copy, the marker labelling
+ and the watershed each materialise a full-frame float32/int32 raster. With
+ dask running ``threads_per_worker=4`` that is ~3.4 GB of concurrent peak
+ per worker, which is what drives a 9.24 GiB worker into the pause/restart
+ loop with "unmanaged memory" warnings: the arrays are ours, inside the
+ task, so dask cannot see or spill them.
+
+ Per component the peak is one particle's bbox — kilobytes — plus the two
+ full-frame rasters that cannot be avoided (the input mask and the output
+ labels).
+
+ It is also the same shape `measure.py` already uses to trace contours
+ ("inside each region's padded bbox crop, not on the whole frame").
+ """
+ from scipy import ndimage as ndi
+
+ comp, n_comp = ndi.label(fg)
+ out = np.zeros(fg.shape, np.int32)
+ if n_comp == 0:
+ return out
+ slices = ndi.find_objects(comp)
+ next_label = 1
+ h, w = fg.shape
+ for i, sl in enumerate(slices, start=1):
+ if sl is None:
+ continue
+ # 1 px pad so the crop's border is background and the distance
+ # transform inside sees the true nearest zero.
+ y0 = max(0, sl[0].start - 1); y1 = min(h, sl[0].stop + 1)
+ x0 = max(0, sl[1].start - 1); x1 = min(w, sl[1].stop + 1)
+ sub_fg = comp[y0:y1, x0:x1] == i
+ sub_dist_src = (None if distance_from is None
+ else np.asarray(distance_from[y0:y1, x0:x1], np.float32))
+ sub = _split_one_component(sub_fg, p, sub_dist_src)
+ if sub is None:
+ continue
+ m = sub > 0
+ if not m.any():
+ continue
+ # Renumber this component's instances into the global label space.
+ out[y0:y1, x0:x1][m] = sub[m] + (next_label - 1)
+ next_label += int(sub.max())
+ return out
+
+
+def _split_one_component(sub_fg: np.ndarray, p: SegmentParams,
+ distance_from: np.ndarray | None):
+ """The watershed body, on ONE component's crop. Labels are 1..k locally."""
+ from skimage.morphology import binary_erosion, disk
+ from skimage.segmentation import watershed
+ from scipy import ndimage as ndi
+
+ seed_src = sub_fg
+ if p.watershed_erosion > 0:
+ seed_src = sub_fg.copy()
+ for _ in range(int(p.watershed_erosion)):
+ seed_src = binary_erosion(seed_src, disk(1))
+
+ if distance_from is not None:
+ dist = np.asarray(distance_from, np.float32) * seed_src
+ markers = _distance_markers(dist, sub_fg, p)
+ else:
+ dist, markers = _distance_and_markers(seed_src, sub_fg, p)
+
+ if markers.max() > 0:
+ return np.asarray(watershed(-dist, markers, mask=sub_fg), np.int32)
+ lab, _ = ndi.label(sub_fg)
+ return np.asarray(lab, np.int32)
diff --git a/spyde/particles/contours.py b/spyde/particles/contours.py
new file mode 100644
index 00000000..0fc40ce4
--- /dev/null
+++ b/spyde/particles/contours.py
@@ -0,0 +1,415 @@
+"""
+contours.py — ``find_contours`` for every region at once, with the same FILLED polygon.
+
+The last of ``measure_frame``'s three per-region Python loops. On a real 4096²
+in-situ growth frame with 26 566 particles it is **4.8 s of a 10.2 s frame**
+(``benchmarks.md``), and, like the property table and the convex hull before it,
+almost none of that is arithmetic: it is a bbox crop, a ``float32`` cast, a
+``skimage.measure.find_contours`` call whose marching squares is Cython but whose
+segment ASSEMBLY is a pure-Python dict-and-deque walk, a ``max``, an ``rint`` and
+a ``clip`` — ~177 us per region, 26 566 times, holding the GIL throughout.
+
+The outline is not decoration
+-----------------------------
+It is tempting to treat an outline as a display choice and allow a "close enough"
+polygon. It is not: :meth:`spyde.signals.particles.SpyDEParticles.render_frame`
+FILLS the contours to rebuild the label movie, and :meth:`~…SpyDEParticles.mask_at`
+fills one to produce the per-particle mask that a mean diffraction pattern is
+sliced with. A different contour is a different mask is a different measurement.
+
+But that also says exactly what the gate is, and it is weaker than bit-identical
+vertices:
+
+ **The FILLED POLYGON must be identical, per region.** Vertex count, ordering
+ and starting point may differ; ``skimage.draw.polygon`` on the new contour
+ must select exactly the same pixels as on the old one.
+
+That is what every consumer reads, so nothing downstream can observe a difference
+that survives it. ``test_particles_contours_parity.py`` asserts it as a boolean
+``array_equal`` of the two filled masks, per region, on a scene of thousands.
+
+Why the cycle is the same cycle
+-------------------------------
+skimage's ``_get_contour_segments`` is a fixed 16-entry case table over the 2x2
+cells of the crop, and on a BINARY mask at level 0.5 every interpolated vertex
+lands exactly on an edge midpoint — ``_get_fraction`` is ``(0.5 - 0) / (1 - 0)``
+for every edge the table actually uses. So a vertex is at ``(i + 0.5, j)`` or
+``(i, j + 0.5)``, i.e. it IS the crack between two 4-adjacent pixels, and it can
+be named by an integer edge index with no floating point anywhere.
+
+Each cell emits its segments oriented so that low values are on the left, and a
+crack interior to the crop is shared by exactly two cells — appearing once as a
+segment's tail and once as another's head. So every vertex has in-degree <= 1 and
+out-degree <= 1, and the segment set decomposes into disjoint simple paths (which
+run off the crop border) and cycles. ``_assemble_contours``'s dicts and deques
+recover exactly those maximal chains; following ``succ`` recovers the same ones,
+in the same direction, differing only in where a CYCLE is cut — which a fill
+cannot see. A cycle is emitted with its first vertex repeated at the end, which is
+what ``_assemble_contours`` does when it closes a loop, so the vertex COUNT
+matches too.
+
+Two details that look like trivia and decide the answer
+-------------------------------------------------------
+* **``np.rint`` is round-half-to-EVEN, and it is applied in CROP coordinates.**
+ Every marching-squares vertex here is a half-integer, so the rounding is
+ entirely in the tie case, and it resolves on the PARITY of the crop-local
+ coordinate — which depends on where the region's padded bbox happens to start.
+ Two congruent particles at different positions therefore get genuinely
+ different integer outlines. That is the behaviour on disk today; reproducing
+ it means rounding in the crop frame and offsetting afterwards, never the
+ reverse.
+* **Which contour "the" contour is.** The caller takes ``max(cs, key=len)``, and
+ ``cs`` is ordered by ``_assemble_contours``'s creation counter — which, after
+ every merge keeps the smaller of the two keys, equals the order of each
+ contour's SMALLEST segment index. So ties in length break to the chain
+ containing the earliest cell in raster order, and that is reproduced here
+ explicitly rather than left to whatever order a walk happens to discover.
+
+``numba`` is optional: :func:`label_contours` returns ``None`` when it is missing
+or the kernel will not compile, and :mod:`spyde.particles.measure` falls back to
+the ``find_contours`` loop. Nothing here is on a GPU, so there is no device lock.
+
+Memory: one padded bbox crop at a time per thread, plus one flat int16 buffer for
+every outline, bounded by ``4 * area + 4`` per region — a region's crossed cracks
+cannot exceed four per pixel (CLAUDE.md § Memory Safety).
+"""
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+_KERNEL = [None] # compiled lazily, once per process
+_FAILED = [False]
+
+
+def _build_kernel():
+ """Compile the per-region contour kernel, or None if numba is unusable."""
+ if _KERNEL[0] is not None or _FAILED[0]:
+ return _KERNEL[0]
+ try:
+ import numba
+ except Exception as exc: # pragma: no cover
+ log.info("[particles] numba unavailable (%s); contours stay on "
+ "find_contours", exc)
+ _FAILED[0] = True
+ return None
+
+ @numba.njit(cache=True, nogil=True, parallel=True, fastmath=False)
+ def _kernel(lab, labels, bb, offsets, out_buf, out_len): # pragma: no cover
+ h, w = lab.shape
+ n = labels.shape[0]
+ for i in numba.prange(n):
+ lbl = labels[i]
+ by0 = bb[i, 0]
+ bx0 = bb[i, 1]
+ by1 = bb[i, 2]
+ bx1 = bb[i, 3]
+ base = offsets[i]
+ # An outline cannot have more vertices than the region has cracks,
+ # and a pixel has at most four — so `cap` is a proof, not a guess.
+ # It is enforced anyway: njit does not bounds-check, and a wrong
+ # bound would be silent memory corruption rather than an IndexError.
+ cap = offsets[i + 1] - base
+
+ # The same crop `_contours` takes: the bbox padded by ONE and
+ # clipped to the frame. One pixel is all marching squares needs to
+ # close a contour around the region; at the frame edge the crop is
+ # truncated and the contour is left open, exactly as skimage leaves
+ # it open against an array border.
+ py0 = by0 - 1
+ px0 = bx0 - 1
+ py1 = by1 + 1
+ px1 = bx1 + 1
+ if py0 < 0:
+ py0 = 0
+ if px0 < 0:
+ px0 = 0
+ if py1 > h:
+ py1 = h
+ if px1 > w:
+ px1 = w
+ hh = py1 - py0
+ ww = px1 - px0
+
+ # `find_contours` refuses an array smaller than 2x2, and returns
+ # nothing for a crop that is entirely inside or entirely outside the
+ # region. Both are the caller's degenerate branch: the bbox corners,
+ # so that every row still has an outline and the 1:1 correspondence
+ # holds.
+ m = np.zeros((1, 1), np.uint8)
+ nseg = 0
+ if hh >= 2 and ww >= 2:
+ m = np.zeros((hh, ww), np.uint8)
+ for r in range(hh):
+ for c in range(ww):
+ if lab[py0 + r, px0 + c] == lbl:
+ m[r, c] = 1
+
+ # ── pass 1: how many segments the case table will emit.
+ for r0 in range(hh - 1):
+ for c0 in range(ww - 1):
+ case = (m[r0, c0] + 2 * m[r0, c0 + 1]
+ + 4 * m[r0 + 1, c0] + 8 * m[r0 + 1, c0 + 1])
+ if case == 0 or case == 15:
+ continue
+ nseg += 2 if (case == 6 or case == 9) else 1
+
+ if nseg == 0:
+ ya = by0 if by0 < h else h - 1
+ yb = by1 - 1 if by1 - 1 < h else h - 1
+ xa = bx0 if bx0 < w else w - 1
+ xb = bx1 - 1 if bx1 - 1 < w else w - 1
+ out_buf[base + 0, 0] = ya
+ out_buf[base + 0, 1] = xa
+ out_buf[base + 1, 0] = ya
+ out_buf[base + 1, 1] = xb
+ out_buf[base + 2, 0] = yb
+ out_buf[base + 2, 1] = xb
+ out_buf[base + 3, 0] = yb
+ out_buf[base + 3, 1] = xa
+ out_len[i] = 4
+ continue
+
+ # A vertex is a CRACK between two 4-adjacent crop pixels: `nh`
+ # horizontal cracks (between columns) then the vertical ones.
+ nh = hh * (ww - 1)
+ nv = nh + (hh - 1) * ww
+ seg_from = np.empty(nseg, np.int32)
+ seg_to = np.empty(nseg, np.int32)
+ succ = np.full(nv, -1, np.int32)
+ has_pred = np.zeros(nv, np.uint8)
+
+ # ── pass 2: the case table itself. Segment ORDER is row-major over
+ # cells, which is what fixes the tie-break below.
+ k = 0
+ for r0 in range(hh - 1):
+ for c0 in range(ww - 1):
+ case = (m[r0, c0] + 2 * m[r0, c0 + 1]
+ + 4 * m[r0 + 1, c0] + 8 * m[r0 + 1, c0 + 1])
+ if case == 0 or case == 15:
+ continue
+ top = r0 * (ww - 1) + c0
+ bottom = (r0 + 1) * (ww - 1) + c0
+ left = nh + r0 * ww + c0
+ right = nh + r0 * ww + c0 + 1
+
+ a0 = -1
+ b0 = -1
+ a1 = -1
+ b1 = -1
+ if case == 1:
+ a0 = top
+ b0 = left
+ elif case == 2:
+ a0 = right
+ b0 = top
+ elif case == 3:
+ a0 = right
+ b0 = left
+ elif case == 4:
+ a0 = left
+ b0 = bottom
+ elif case == 5:
+ a0 = top
+ b0 = bottom
+ elif case == 6:
+ a0 = right
+ b0 = top
+ a1 = left
+ b1 = bottom
+ elif case == 7:
+ a0 = right
+ b0 = bottom
+ elif case == 8:
+ a0 = bottom
+ b0 = right
+ elif case == 9:
+ a0 = top
+ b0 = left
+ a1 = bottom
+ b1 = right
+ elif case == 10:
+ a0 = bottom
+ b0 = top
+ elif case == 11:
+ a0 = bottom
+ b0 = left
+ elif case == 12:
+ a0 = left
+ b0 = right
+ elif case == 13:
+ a0 = top
+ b0 = right
+ else: # case == 14
+ a0 = left
+ b0 = top
+
+ seg_from[k] = a0
+ seg_to[k] = b0
+ succ[a0] = k
+ has_pred[b0] = 1
+ k += 1
+ if a1 >= 0:
+ seg_from[k] = a1
+ seg_to[k] = b1
+ succ[a1] = k
+ has_pred[b1] = 1
+ k += 1
+
+ # ── pass 3: follow `succ` to recover the maximal chains, and keep
+ # the longest. Open paths first (they have a vertex with no
+ # predecessor to start from), then whatever is left, which is
+ # cycles.
+ used = np.zeros(nseg, np.uint8)
+ verts = np.empty(nseg + 1, np.int32)
+ best = np.empty(nseg + 1, np.int32)
+ best_len = 0
+ best_min = 0
+
+ for phase in range(2):
+ for j in range(nseg):
+ if used[j] != 0:
+ continue
+ if phase == 0 and has_pred[seg_from[j]] != 0:
+ continue
+ cnt = 0
+ verts[cnt] = seg_from[j]
+ cnt += 1
+ minidx = j
+ e = j
+ while True:
+ used[e] = 1
+ if e < minidx:
+ minidx = e
+ verts[cnt] = seg_to[e]
+ cnt += 1
+ nxt = succ[seg_to[e]]
+ if nxt < 0 or used[nxt] != 0:
+ break
+ e = nxt
+ if cnt > best_len or (cnt == best_len and minidx < best_min):
+ best_len = cnt
+ best_min = minidx
+ for a in range(cnt):
+ best[a] = verts[a]
+
+ # ── pass 4: half-integer crack -> np.rint (half to EVEN) in CROP
+ # coordinates, then offset to the frame and clip.
+ if best_len > cap:
+ best_len = cap
+ for a in range(best_len):
+ v = best[a]
+ if v < nh:
+ rr = v // (ww - 1)
+ cc = v - rr * (ww - 1)
+ dr = 2 * rr
+ dc = 2 * cc + 1
+ else:
+ u = v - nh
+ rr = u // ww
+ cc = u - rr * ww
+ dr = 2 * rr + 1
+ dc = 2 * cc
+
+ kr = dr >> 1
+ if (dr & 1) != 0 and (kr & 1) != 0:
+ kr += 1
+ kc = dc >> 1
+ if (dc & 1) != 0 and (kc & 1) != 0:
+ kc += 1
+
+ y = kr + py0
+ x = kc + px0
+ if y < 0:
+ y = 0
+ elif y > h - 1:
+ y = h - 1
+ if x < 0:
+ x = 0
+ elif x > w - 1:
+ x = w - 1
+ out_buf[base + a, 0] = y
+ out_buf[base + a, 1] = x
+ out_len[i] = best_len
+
+ _KERNEL[0] = _kernel
+ return _kernel
+
+
+def label_contours(lab: np.ndarray, labels: np.ndarray, bboxes: np.ndarray,
+ areas: np.ndarray) -> list[np.ndarray] | None:
+ """One int16 ``(k, 2)`` outline per entry of *labels*, or None without numba.
+
+ Parameters
+ ----------
+ lab
+ The ``(h, w)`` int label image.
+ labels
+ Ascending labels present in *lab*.
+ bboxes
+ ``(n, 4)`` int ``(y0, x0, y1, x1)`` per entry of *labels*.
+ areas
+ Pixel count per entry of *labels*. Used only to size the output buffer:
+ a region's outline cannot have more vertices than it has cracks, and a
+ pixel has at most four.
+
+ Returns
+ -------
+ A list of ``(k, 2)`` int16 arrays, one per label and in the same order — the
+ 1:1 correspondence ``SpyDEParticles.from_frames`` requires. Each is a VIEW
+ into one contiguous buffer, which is also the layout ``SpyDEParticles``
+ stores (``contours`` + ``contour_offsets``).
+ """
+ labels = np.asarray(labels, dtype=np.int64)
+ n = labels.size
+ if n == 0:
+ return []
+ kernel = _build_kernel()
+ if kernel is None:
+ return None
+
+ bb = np.ascontiguousarray(bboxes, dtype=np.int64)
+ bound = 4 * np.asarray(areas, dtype=np.int64) + 4
+ offsets = np.zeros(n + 1, np.int64)
+ np.cumsum(bound, out=offsets[1:])
+ buf = np.zeros((int(offsets[-1]), 2), np.int16)
+ lens = np.zeros(n, np.int64)
+ try:
+ kernel(lab, labels, bb, offsets, buf, lens)
+ except Exception as exc: # pragma: no cover
+ log.warning("[particles] contour kernel failed (%r); outlines fall back "
+ "to find_contours", exc)
+ _FAILED[0] = True
+ _KERNEL[0] = None
+ return None
+
+ # Compact into a buffer of the ACTUAL size. `bound` is ~4x what an outline
+ # really needs, and the caller (`batch.py`) accumulates one contour list per
+ # frame for the whole movie — holding 900 over-allocated frames alive is
+ # gigabytes for nothing, whereas the compacted CSR is what
+ # `SpyDEParticles` stores anyway.
+ new_off = np.zeros(n + 1, np.int64)
+ np.cumsum(lens, out=new_off[1:])
+ total = int(new_off[-1])
+ within = np.arange(total, dtype=np.int64) - np.repeat(new_off[:n], lens)
+ src = np.repeat(offsets[:n], lens) + within
+ flat = np.ascontiguousarray(buf[src], np.int16)
+ del buf, src, within
+ return [flat[int(a):int(b)] for a, b in zip(new_off[:n], new_off[1:])]
+
+
+def warmup() -> bool:
+ """Compile the kernel on a toy image. Returns True if numba is live.
+
+ The counterpart of :func:`spyde.particles.hull.warmup`: a dask worker should
+ pay the first ``njit`` compile before the first measured frame, not inside it.
+ """
+ lab = np.zeros((8, 8), np.int32)
+ lab[2:5, 2:5] = 1
+ out = label_contours(lab, np.array([1], np.int64),
+ np.array([[2, 2, 5, 5]], np.int64),
+ np.array([9], np.int64))
+ return out is not None
diff --git a/spyde/particles/features.py b/spyde/particles/features.py
new file mode 100644
index 00000000..5fa8525f
--- /dev/null
+++ b/spyde/particles/features.py
@@ -0,0 +1,1028 @@
+"""
+features.py — the torch feature stack the learned engines classify. Plan step B2.
+
+ParticleSpy's ``trainable_parameters`` set (which is Weka Trainable Segmentation's
+set, which is ilastik's set): gaussian blur, difference-of-gaussians, median /
+minimum / maximum rank filters, Sobel gradient magnitude, Hessian eigenvalues,
+Laplacian, and membrane projections. Every channel is computed in torch, so the
+same code runs on CPU, CUDA and Apple-MPS, and the classifier head
+(:mod:`spyde.particles.scribble`) never leaves the device.
+
+Three properties are load-bearing and are what the implementation is shaped
+around.
+
+**1. One pass over the frame, not one pass per feature.** The gaussian pyramid is
+computed once (separably — two 1-D convolutions per sigma instead of one 2-D one)
+and then *every* other family is derived from it: difference-of-gaussians is a
+subtraction of two cached blurs, the Laplacian is the trace of the Hessian's own
+second derivatives, the two Hessian eigenvalues come from the same three
+derivative images, and median/minimum/maximum come from one unfolded window each.
+:class:`_Pass` is the memo that makes that true — a channel never triggers work
+another channel already did. Measured on a 512² float32 frame, CPU, default spec
+(36 channels): **54 ms shared vs 156 ms** when each channel recomputes its own
+intermediates.
+
+**2. The fine scales are load-bearing for SMALL particles** (plan §0.9) — and the
+measurement is not the one you would guess. It is not about *detection*: on the
+two deliberately faint probes in ``particle_movie()`` (r=4 and r=3) a coarse
+``(4, 8)`` stack still finds both. It is about what is then *measured* of them.
+Mean absolute error in recovered radius over the seven isolated particles, with
+identical labels and an identical head, and the r=3 probe on its own::
+
+ sigmas (0.5, 1, 2, 4, 8) 13.4 % r=3 probe -12 %
+ sigmas (1, 2, 4, 8) 13.2 % r=3 probe -15 %
+ sigmas (2, 4, 8) 20.4 % r=3 probe -28 %
+ sigmas (4, 8) 25.5 % r=3 probe -44 %
+
+So the floor must not be raised above ~1 px, which is what plan §0.9's "not
+without a documented sensitivity measurement" amounts to here — a found particle
+measured 44% too small is worse than an honest miss, because it enters the size
+distribution. 0.5 buys nothing over 1.0 on this fixture and is kept anyway: it is
+2 ms of a 71 ms 512² stack, and it is what a genuinely 1–2 px feature needs.
+
+**3. Memory is bounded by row-banding, not by hope.** The full stack is
+``C·H·W·4`` bytes — 2.4 GB for 36 channels at 4096², which is the plan's stated
+frame size. So the stack is produced in **row bands with a halo** (:func:`
+map_feature_bands`), and the two callers that matter never materialise the whole
+thing: training samples the labelled pixels out of each band and drops it, and
+:meth:`~spyde.particles.scribble.ScribbleClassifier.predict_proba` writes one
+band's probabilities and drops it. With a halo of at least the largest filter
+radius the banded result is *identical* to the unbanded one — pinned by a test.
+
+NaN input
+---------
+A drift-corrected frame carries a NaN-padded border
+(:mod:`spyde.drift.warp`), and every convolution propagates NaN outward, which
+would erase a band of real data. :func:`prepare_frame` fills the non-finite
+pixels with the finite **minimum** — the same choice, for the same reason, as
+:func:`spyde.particles.classical._prepare`: the padding then reads as background,
+the one value guaranteed not to classify as a particle. It also returns the
+validity mask, and the classifier is required to force those pixels to zero
+probability (plan trap 2).
+"""
+from __future__ import annotations
+
+import logging
+import math
+from dataclasses import asdict, dataclass, replace
+from typing import Any, Callable, NamedTuple
+
+import numpy as np
+
+from spyde.device_lock import accelerator_lock
+
+log = logging.getLogger(__name__)
+
+#: Gaussian kernels are truncated at this many sigma, matching
+#: ``scipy.ndimage.gaussian_filter``'s default so the two agree numerically
+#: (``test_particles_scribble.py`` pins that against scipy with ``mode="mirror"``,
+#: which is what torch's ``reflect`` padding is).
+_TRUNCATE = 4.0
+
+#: Default scales, in pixels, octave-spaced. The **floor** is what matters — see
+#: the module docstring for the measured cost of raising it. The top end (8) is
+#: what gives the head a local-background reference, which is how it separates a
+#: faint particle from a bright patch of support film.
+DEFAULT_SIGMAS: tuple[float, ...] = (0.5, 1.0, 2.0, 4.0, 8.0)
+
+#: Default rank-filter window radii, in pixels (window is ``2r+1`` square).
+DEFAULT_RANK_RADII: tuple[int, ...] = (1, 2)
+
+#: Projections available across the rotated membrane responses (Weka's set).
+MEMBRANE_PROJECTIONS: tuple[str, ...] = ("sum", "mean", "std", "median",
+ "max", "min")
+
+#: Window size (in taps) at or above which the median reduces along a
+#: transposed, contiguous axis. See :meth:`_Pass._compute_rank` for the numbers;
+#: 25 is r=2, the largest default radius, and 9 (r=1) stays on the direct path.
+_MEDIAN_TRANSPOSE_TAPS: int = 25
+
+#: Target working-set size for one row band on the CPU. Not a hard cap — a band
+#: is always at least ``4·halo`` rows tall, because a band shorter than its own
+#: halo would recompute more halo than payload.
+BAND_BYTES: int = 256 << 20
+
+#: Ceiling for one band on an accelerator. Every band re-featurises ``halo`` rows
+#: above and below itself, so at 4096² a 256 MB band means 22 bands and **36%**
+#: of the work is halo; a 1 GiB band means 6 bands and 11%. Measured, CUDA,
+#: featurise+head over a 4096² frame: 1366 ms at 184 rows/band, **1012 ms** at
+#: 768.
+#:
+#: It is a ceiling and not a target because overshooting is far worse than
+#: undershooting: the same sweep at 1536 rows/band took **12.8 s** and at one
+#: band 26.7 s, thrashing the allocator against a 12.9 GB card. So the budget is
+#: also clamped to a fraction of *free* device memory rather than assumed.
+GPU_BAND_BYTES: int = 1 << 30
+
+#: Fraction of free device memory a single band may claim. A band's peak is
+#: several times the figure `band_rows_for` estimates (autograd-free, but torch's
+#: caching allocator holds freed blocks), and the cliff above is steep.
+_GPU_FREE_FRACTION: float = 0.25
+
+#: Robust per-frame statistics are estimated from at most this many pixels.
+#: ``np.percentile`` on a 4096² frame is a full sort (~1.5 s here) and would
+#: dominate the whole interaction budget; a strided subsample of 10⁶ pixels puts
+#: the median and IQR well inside their own sampling noise. Same trade-off, and
+#: the same reasoning, as the subsampled first-paint histogram in ``plot.py``.
+_STAT_SAMPLE_MAX = 1_000_000
+
+
+# ── spec ─────────────────────────────────────────────────────────────────────
+
+@dataclass(frozen=True)
+class FeatureSpec:
+ """Which features to compute, and at which scales.
+
+ Frozen and serialisable: :meth:`to_dict` / :meth:`from_dict` round-trip
+ through JSON, which is what lets a wizard parameter schema drive it and what
+ lets a trained recipe be saved next to its weights
+ (:meth:`spyde.particles.scribble.ScribbleClassifier.save`). A model whose
+ spec no longer matches its weights is meaningless, so the two are never
+ stored apart.
+
+ Parameters
+ ----------
+ sigmas
+ Gaussian scales in pixels, ascending. Drives ``gaussian``,
+ ``difference_of_gaussians``, ``sobel``, ``hessian`` and ``laplacian``.
+ rank_radii
+ Window radii for ``median`` / ``minimum`` / ``maximum``, in pixels.
+ membrane
+ Weka's membrane projections: a thin line kernel rotated through 180° and
+ projected. **Off by default**, which is a deliberate deviation from
+ ParticleSpy's set rather than an omission — the family exists to find
+ *elongated* structures (its name is literal; it was built for neuron
+ membranes), it is the most expensive channel here per channel (adds 33 ms
+ to a 71 ms 512² total for four channels, so ~8 ms/channel against ~2 ms
+ for everything else), and on the compact blobs this feature is for it did
+ not change the faint-probe result either way. Turn it on for fibres,
+ films or lattice fringes, where it earns its cost.
+ membrane_thickness
+ Line width in pixels. Rounded up to an odd number so the line is centred.
+ normalize_frame
+ Robustly standardise the frame (median / IQR over finite pixels) before
+ featurising. On by default because it is what makes a saved recipe
+ transferable: without it every channel carries the dataset's absolute
+ intensity scale, so a model trained on a float image in 0..1 predicts
+ nothing at all on the same sample recorded as uint16 counts.
+ """
+
+ sigmas: tuple[float, ...] = DEFAULT_SIGMAS
+ intensity: bool = True
+ gaussian: bool = True
+ difference_of_gaussians: bool = True
+ sobel: bool = True
+ hessian: bool = True
+ laplacian: bool = True
+ rank_radii: tuple[int, ...] = DEFAULT_RANK_RADII
+ median: bool = True
+ minimum: bool = True
+ maximum: bool = True
+ membrane: bool = False
+ membrane_patch: int = 19
+ membrane_thickness: int = 1
+ membrane_rotations: int = 12
+ membrane_projections: tuple[str, ...] = ("mean", "std", "max", "min")
+ normalize_frame: bool = True
+
+ def __post_init__(self) -> None:
+ # Coerce lists (which is what `from_dict` and a JSON parameter schema
+ # hand over) to tuples, so the dataclass stays hashable and comparable.
+ object.__setattr__(self, "sigmas",
+ tuple(float(s) for s in self.sigmas))
+ object.__setattr__(self, "rank_radii",
+ tuple(int(r) for r in self.rank_radii))
+ object.__setattr__(self, "membrane_projections",
+ tuple(str(p) for p in self.membrane_projections))
+ if any(s <= 0 for s in self.sigmas):
+ raise ValueError(f"sigmas must be positive; got {self.sigmas}")
+ if tuple(sorted(self.sigmas)) != self.sigmas:
+ raise ValueError(f"sigmas must be ascending; got {self.sigmas}")
+ if any(r < 1 for r in self.rank_radii):
+ raise ValueError(f"rank_radii must be >= 1; got {self.rank_radii}")
+ bad = set(self.membrane_projections) - set(MEMBRANE_PROJECTIONS)
+ if bad:
+ raise ValueError(
+ f"unknown membrane projection(s) {sorted(bad)}; expected any of "
+ f"{', '.join(MEMBRANE_PROJECTIONS)}"
+ )
+ if self.membrane and self.membrane_patch % 2 == 0:
+ raise ValueError(
+ f"membrane_patch must be odd so the line is centred; got "
+ f"{self.membrane_patch}"
+ )
+ if self.membrane and self.membrane_rotations < 1:
+ raise ValueError(
+ f"membrane_rotations must be >= 1; got {self.membrane_rotations}")
+ if not self.channel_names():
+ raise ValueError(
+ "this FeatureSpec produces no channels at all — every family is "
+ "disabled, or the enabled ones have empty sigmas/rank_radii")
+
+ # -- serialisation ---------------------------------------------------------
+
+ def to_dict(self) -> dict[str, Any]:
+ """Plain JSON-safe dict (tuples become lists)."""
+ out = asdict(self)
+ for key in ("sigmas", "rank_radii", "membrane_projections"):
+ out[key] = list(out[key])
+ return out
+
+ @classmethod
+ def from_dict(cls, d: dict[str, Any] | None) -> "FeatureSpec":
+ """Rebuild from :meth:`to_dict`, ignoring keys this build does not know.
+
+ Unknown keys are dropped rather than raising: a spec written by a newer
+ build should still load with the features this one understands, and the
+ alternative is that a saved recipe becomes unopenable on a downgrade.
+ Missing keys take their defaults.
+ """
+ if not d:
+ return cls()
+ fields = set(cls.__dataclass_fields__)
+ unknown = set(d) - fields
+ if unknown:
+ log.debug("FeatureSpec.from_dict ignoring unknown keys: %s",
+ sorted(unknown))
+ return cls(**{k: v for k, v in d.items() if k in fields})
+
+ def replace(self, **kw: Any) -> "FeatureSpec":
+ """A copy with *kw* overridden (the frozen-dataclass setter)."""
+ return replace(self, **kw)
+
+ # -- shape -----------------------------------------------------------------
+
+ def channel_names(self) -> list[str]:
+ """Channel names, in output order.
+
+ Derived from the same :func:`_channel_plan` the tensor builder walks, so
+ the names cannot drift out of step with the channels — there is exactly
+ one definition of the order.
+ """
+ return [name for _family, name, _args in _channel_plan(self)]
+
+ @property
+ def n_channels(self) -> int:
+ return len(_channel_plan(self))
+
+ @property
+ def halo(self) -> int:
+ """Largest radius, in pixels, that any enabled filter reaches.
+
+ This is the row overlap :func:`map_feature_bands` needs for a banded
+ stack to equal an unbanded one. The ``+1`` on the gaussian radius is the
+ 3-tap Sobel / second-difference stencil applied *after* the blur.
+ """
+ r = 1
+ if self.sigmas and (self.gaussian or self.difference_of_gaussians or
+ self.sobel or self.hessian or self.laplacian):
+ r = max(r, _gauss_radius(max(self.sigmas)) + 1)
+ if self.rank_radii and (self.median or self.minimum or self.maximum):
+ r = max(r, max(self.rank_radii))
+ if self.membrane:
+ r = max(r, self.membrane_patch // 2)
+ return int(r)
+
+
+# ── the channel plan: THE definition of channel order and channel names ──────
+
+def _fmt(x: float) -> str:
+ return f"{float(x):g}"
+
+
+def _channel_plan(spec: FeatureSpec) -> list[tuple[str, str, tuple]]:
+ """``(family, name, args)`` for every channel, in output order.
+
+ Both :meth:`FeatureSpec.channel_names` and :class:`_Pass` walk this list, so a
+ new family is added in one place and the names follow automatically. Grouping
+ by family rather than interleaving by sigma costs nothing — the shared
+ intermediates are memoised on :class:`_Pass`, so revisiting sigma 2 for the
+ Hessian after visiting it for the Sobel does not recompute its blur.
+ """
+ plan: list[tuple[str, str, tuple]] = []
+ if spec.intensity:
+ plan.append(("intensity", "intensity", ()))
+ if spec.gaussian:
+ plan += [("gaussian", f"gaussian_s{_fmt(s)}", (s,)) for s in spec.sigmas]
+ if spec.difference_of_gaussians:
+ plan += [("dog", f"dog_s{_fmt(a)}_s{_fmt(b)}", (a, b))
+ for a, b in zip(spec.sigmas, spec.sigmas[1:])]
+ if spec.sobel:
+ plan += [("sobel", f"sobel_s{_fmt(s)}", (s,)) for s in spec.sigmas]
+ if spec.laplacian:
+ plan += [("laplacian", f"laplacian_s{_fmt(s)}", (s,)) for s in spec.sigmas]
+ if spec.hessian:
+ for s in spec.sigmas:
+ plan.append(("hessian_major", f"hessian_major_s{_fmt(s)}", (s,)))
+ plan.append(("hessian_minor", f"hessian_minor_s{_fmt(s)}", (s,)))
+ for stat in ("median", "minimum", "maximum"):
+ if getattr(spec, stat):
+ plan += [(stat, f"{stat}_r{r}", (r,)) for r in spec.rank_radii]
+ if spec.membrane:
+ plan += [("membrane", f"membrane_{p}", (p,))
+ for p in spec.membrane_projections]
+ return plan
+
+
+# ── device ───────────────────────────────────────────────────────────────────
+
+def import_torch():
+ """``import torch``, with a message that says whose problem it is if it fails.
+
+ torch is a core SpyDE dependency, so a failure here is a broken environment
+ rather than a missing extra — and the bare ``ModuleNotFoundError`` from six
+ frames down does not say that. Shared with
+ :mod:`spyde.particles.scribble` rather than duplicated there.
+ """
+ try:
+ import torch
+ return torch
+ except Exception as exc: # pragma: no cover
+ raise ImportError(
+ "spyde.particles needs torch, which is a core SpyDE dependency but "
+ f"failed to import: {exc}"
+ ) from exc
+
+
+
+def select_device(prefer: str | None = None):
+ """Best torch device for the feature stack: CUDA → Apple-MPS → CPU.
+
+ Parameters
+ ----------
+ prefer
+ An explicit device string (``"cpu"``, ``"cuda"``, ``"mps"``) to force.
+ Tests force ``"cpu"``: torch-CUDA work segfaults under the pytest process
+ on Windows (CLAUDE.md), which is a harness interaction rather than a code
+ defect, so the GPU path is exercised in a subprocess instead.
+ """
+ torch = import_torch()
+ if prefer is not None:
+ return torch.device(prefer)
+ try:
+ if torch.cuda.is_available():
+ return torch.device("cuda")
+ mps = getattr(torch.backends, "mps", None)
+ if mps is not None and mps.is_available():
+ return torch.device("mps")
+ except Exception as exc: # pragma: no cover
+ log.debug("device probe failed (%s); using CPU", exc)
+ return torch.device("cpu")
+
+
+def gpu_available() -> bool:
+ """True when a hardware-accelerated torch device exists (CUDA or MPS).
+
+ CPU-only torch is False here even though the stack runs perfectly well on it
+ — this gate is about whether to *expect* interactive speed, matching
+ ``vector_orientation_gpu.gpu_available``.
+ """
+ try:
+ return select_device().type in ("cuda", "mps")
+ except Exception: # pragma: no cover
+ return False
+
+
+# ── frame preparation ────────────────────────────────────────────────────────
+
+class PreparedFrame(NamedTuple):
+ """A frame ready to featurise, plus the validity mask the caller must honour.
+
+ ``image``
+ float32, non-finite pixels replaced (see :func:`prepare_frame`), and
+ robustly standardised when the spec asks for it.
+ ``valid``
+ bool, True where the *source* pixel was finite. The classifier forces
+ zero foreground probability outside this — plan trap 2.
+ """
+
+ image: np.ndarray
+ valid: np.ndarray
+
+
+def _robust_stats(values: np.ndarray) -> tuple[float, float]:
+ """``(centre, spread)`` — median and IQR/1.349, from a strided subsample.
+
+ IQR/1.349 is the normal-consistent robust sigma. Robust rather than
+ mean/std because a frame *full of particles* has a heavy bright tail, and a
+ mean/std standardisation would then move the background level around with the
+ particle coverage — so the same physical background would present differently
+ at t=0 and t=end, which is exactly the drift a learned head must not see.
+ """
+ v = values
+ if v.size > _STAT_SAMPLE_MAX:
+ v = v[:: int(np.ceil(v.size / _STAT_SAMPLE_MAX))]
+ centre = float(np.median(v))
+ q1, q3 = np.percentile(v, [25.0, 75.0])
+ spread = float(q3 - q1) / 1.349
+ if not (spread > 0):
+ # Constant (or near-constant) frame: fall back to the std, then to 1 so
+ # the standardisation is a no-op rather than a division by zero.
+ spread = float(np.std(v)) or 1.0
+ return centre, spread
+
+
+def prepare_frame(frame, spec: FeatureSpec | None = None) -> PreparedFrame:
+ """Fill non-finite pixels, optionally standardise, and report validity.
+
+ Parameters
+ ----------
+ frame
+ 2-D array. May contain NaN (a drift-corrected border does).
+ spec
+ Only ``normalize_frame`` is consulted.
+
+ Returns
+ -------
+ PreparedFrame
+
+ Notes
+ -----
+ Non-finite pixels are filled with the finite **minimum**, not with zero and
+ not with the mean: the padding has to read as background, and the minimum is
+ the one value that cannot threshold or classify as a particle. This mirrors
+ :func:`spyde.particles.classical._prepare` deliberately — two engines that
+ disagree about what the padding *is* would disagree about the frame border
+ for no reason a user could see.
+
+ The robust statistics are computed over finite pixels **only**, before the
+ fill. Filling first would let a large NaN border pull the median down and
+ rescale the whole frame by how much of it was padding.
+ """
+ spec = spec or FeatureSpec()
+ img = np.asarray(frame, dtype=np.float32)
+ if img.ndim != 2:
+ raise ValueError(f"frame must be 2-D; got shape {img.shape}")
+ if min(img.shape) < 4:
+ raise ValueError(
+ f"frame must be at least 4x4 to filter; got {img.shape}")
+
+ valid = np.isfinite(img)
+ if valid.all():
+ finite = img.reshape(-1)
+ else:
+ img = img.copy()
+ finite = img[valid]
+ if finite.size == 0:
+ raise ValueError("frame has no finite pixels")
+ img[~valid] = finite.min()
+
+ if spec.normalize_frame:
+ centre, spread = _robust_stats(finite)
+ img = (img - np.float32(centre)) / np.float32(spread)
+
+ return PreparedFrame(np.ascontiguousarray(img, dtype=np.float32), valid)
+
+
+def _as_prepared(frame, spec: FeatureSpec) -> PreparedFrame:
+ return frame if isinstance(frame, PreparedFrame) else prepare_frame(frame, spec)
+
+
+# ── separable convolution primitives ─────────────────────────────────────────
+
+def _gauss_radius(sigma: float) -> int:
+ """Kernel half-width, matching ``scipy.ndimage``'s ``int(truncate*sd + 0.5)``."""
+ return max(1, int(_TRUNCATE * float(sigma) + 0.5))
+
+
+def _gauss_kernel(torch, sigma: float, radius: int, device, dtype):
+ x = torch.arange(-radius, radius + 1, device=device, dtype=dtype)
+ k = torch.exp(-0.5 * (x / float(sigma)) ** 2)
+ return k / k.sum()
+
+
+def _pad_edges(t, ry: int, rx: int):
+ """Pad ``(1, 1, h, w)`` by *ry* rows and *rx* columns, reflect where legal.
+
+ Padding is torch's ``reflect``, which is ``scipy.ndimage``'s ``mirror`` (NOT
+ scipy's ``reflect``, which duplicates the edge sample); the parity test
+ against scipy passes ``mode="mirror"`` for exactly this reason.
+
+ ``reflect`` requires the pad to be strictly smaller than the dimension, and
+ the kernels here are not small: a sigma-8 gaussian has radius 32 and a
+ 19x19 membrane patch radius 9, either of which exceeds a small frame. So the
+ reflect is taken as far as it is legal and the remainder is **replicated**.
+ Clamping the kernel radius instead was the first implementation, and it is
+ worse in a way that is invisible until it matters: it silently applies a
+ *different, narrower* filter than the one the FeatureSpec asked for, so the
+ same spec means different things on different frame sizes — and a saved recipe
+ would then not reproduce on a crop.
+ """
+ import torch.nn.functional as F
+ h, w = int(t.shape[-2]), int(t.shape[-1])
+ ry_ok, rx_ok = min(ry, h - 1), min(rx, w - 1)
+ if ry_ok or rx_ok:
+ t = F.pad(t, (rx_ok, rx_ok, ry_ok, ry_ok), mode="reflect")
+ over_y, over_x = ry - ry_ok, rx - rx_ok
+ if over_y or over_x:
+ t = F.pad(t, (over_x, over_x, over_y, over_y), mode="replicate")
+ return t
+
+
+def _sep_conv(torch, t, ky, kx):
+ """Separable correlation of ``(1, 1, h, w)`` *t* with 1-D kernels.
+
+ ``F.conv2d`` is a correlation, not a convolution, so a kernel is applied as
+ written — no flip. Either kernel may be ``None`` to skip that axis.
+ """
+ import torch.nn.functional as F
+ if ky is not None:
+ t = _pad_edges(t, (ky.numel() - 1) // 2, 0)
+ t = F.conv2d(t, ky.view(1, 1, -1, 1))
+ if kx is not None:
+ t = _pad_edges(t, 0, (kx.numel() - 1) // 2)
+ t = F.conv2d(t, kx.view(1, 1, 1, -1))
+ return t
+
+
+_MEMBRANE_CACHE: dict[tuple[int, int, int], np.ndarray] = {}
+
+
+def _membrane_kernels(patch: int, thickness: int, rotations: int) -> np.ndarray:
+ """``(rotations, patch, patch)`` line kernels spanning 180°, sum-normalised.
+
+ Built with ``scipy.ndimage.rotate`` once per configuration and cached. The
+ rotation is a spline resample of a tiny image, so this is not expensive
+ (measured 0.98 ms for the default twelve 19x19 kernels), but it sits on the
+ interaction path — every retrain re-featurises every labelled frame — and 1 ms
+ per frame has no reason to recur.
+
+ Sum-normalised so a projection stays in the same units as the intensity
+ channel. Weka does not normalise; unnormalised, the ``sum`` projection is
+ ``patch`` times larger than every other channel, which makes the head's first
+ layer spend its capacity on rescaling.
+ """
+ key = (int(patch), int(thickness), int(rotations))
+ cached = _MEMBRANE_CACHE.get(key)
+ if cached is not None:
+ return cached
+
+ from scipy.ndimage import rotate
+
+ base = np.zeros((patch, patch), dtype=np.float32)
+ c = patch // 2
+ half = max(1, int(thickness)) // 2
+ base[:, c - half: c + half + 1] = 1.0
+
+ out = []
+ for i in range(rotations):
+ k = base if i == 0 else rotate(
+ base, 180.0 * i / rotations, reshape=False, order=1,
+ mode="constant", cval=0.0)
+ k = np.clip(np.asarray(k, dtype=np.float32), 0.0, None)
+ s = float(k.sum())
+ out.append(k / s if s > 0 else k)
+ stacked = np.stack(out)
+ _MEMBRANE_CACHE[key] = stacked
+ return stacked
+
+
+# ── one pass over a frame (or a row band) ────────────────────────────────────
+
+class _Pass:
+ """Shared intermediates for one image, computed at most once each.
+
+ This is where "one pass over the frame, not one pass per feature" actually
+ lives. Anything more than one channel needs — the blur at a sigma, that
+ blur's three second derivatives, an unfolded rank window, the rotated
+ membrane responses — is memoised here, so the 36 default channels cost 5
+ gaussian blurs, 5 derivative triples, 2 rank passes and nothing else.
+ """
+
+ def __init__(self, image, spec: FeatureSpec, device):
+ torch = import_torch()
+ self._torch = torch
+ self.spec = spec
+ self.device = device
+ arr = np.ascontiguousarray(image, dtype=np.float32)
+ self.h, self.w = int(arr.shape[0]), int(arr.shape[1])
+ self.img = torch.as_tensor(arr, device=device)
+ self._t = self.img.view(1, 1, self.h, self.w)
+ self._blur: dict[float, Any] = {}
+ self._second: dict[float, tuple] = {}
+ self._rank: dict[tuple[int, str], Any] = {}
+ self._membrane: Any = None
+
+ # -- primitives ----------------------------------------------------------
+
+ def blur(self, sigma: float):
+ """Gaussian blur at *sigma* as ``(1, 1, h, w)``. Memoised."""
+ got = self._blur.get(sigma)
+ if got is None:
+ torch = self._torch
+ k = _gauss_kernel(torch, sigma, _gauss_radius(sigma), self.device,
+ self.img.dtype)
+ got = _sep_conv(torch, self._t, k, k)
+ self._blur[sigma] = got
+ return got
+
+ def second(self, sigma: float):
+ """``(dyy, dxx, dxy)`` of the blur at *sigma*. Memoised.
+
+ Second differences ``[1, -2, 1]`` and a cross term
+ ``[-½, 0, ½] ⊗ [-½, 0, ½]``. **Not** skimage's ``hessian_matrix``, which
+ applies ``np.gradient`` twice and therefore uses a 5-tap
+ ``[¼, 0, -½, 0, ¼]`` — a wider, softer stencil that skips the immediate
+ neighbours entirely. Both are valid discrete Hessians; the compact one is
+ used because it keeps the fine-scale response localised, which is what the
+ smallest sigmas are in the set for. Exact skimage parity is not a
+ requirement here — a channel is an input to a learned head, not a
+ published measurement.
+ """
+ got = self._second.get(sigma)
+ if got is None:
+ torch = self._torch
+ b = self.blur(sigma)
+ d2 = torch.tensor([1.0, -2.0, 1.0], device=self.device,
+ dtype=self.img.dtype)
+ d1 = torch.tensor([-0.5, 0.0, 0.5], device=self.device,
+ dtype=self.img.dtype)
+ dyy = _sep_conv(torch, b, d2, None)
+ dxx = _sep_conv(torch, b, None, d2)
+ dxy = _sep_conv(torch, b, d1, d1)
+ got = (dyy, dxx, dxy)
+ self._second[sigma] = got
+ return got
+
+ def rank(self, radius: int, stat: str):
+ """Median / minimum / maximum over a ``(2r+1)²`` window. Memoised."""
+ got = self._rank.get((int(radius), stat))
+ if got is None:
+ self._compute_rank(int(radius))
+ got = self._rank[(int(radius), stat)]
+ return got
+
+ def _compute_rank(self, radius: int) -> None:
+ """ONE ``unfold`` per radius serves all three rank statistics.
+
+ ``unfold`` materialises ``(2r+1)²·h·w`` floats — 1.7 GB for r=2 on a
+ 4096² frame — which is why the whole stack is banded (:func:`
+ map_feature_bands`); the window is dropped as soon as the enabled
+ statistics are reduced out of it, so only one is ever alive.
+
+ ``max_pool2d(stride=1)`` was the obvious alternative for min/max and is
+ **much slower**, which is the opposite of what its fused-reduction
+ implementation suggests. Measured on a 512² frame, CPU::
+
+ r=1 max_pool2d 18.4 ms unfold amin+amax 2.4 ms
+ r=2 max_pool2d 46.9 ms unfold amin+amax 5.3 ms
+
+ torch's pooling kernels are tuned for the strided, downsampling case; at
+ stride 1 they re-read every window from scratch. Making the window
+ separable (two 1-D pools — a square min/max genuinely is separable, and
+ the results are bit-identical) only got r=2 from 46.9 to 23.7 ms, still
+ 4x the unfold. And the unfold is *shared* with the median, which needs the
+ window materialised regardless, so the whole rank family costs one pass:
+ 170 ms → 27 ms for the default two radii.
+
+ **The median reduces along the CONTIGUOUS axis**, which is not a
+ micro-optimisation — it is over half the cost of the whole feature stack.
+ ``unfold`` returns ``(1, k*k, h*w)``, so ``median(dim=1)`` reduces along
+ the *strided* axis and every one of the k² reads is a stride of h·w
+ floats. Transposing to ``(1, h*w, k*k)`` first makes each window's taps
+ adjacent. Measured on a 768x4096 band, CUDA, r=2::
+
+ u.median(dim=1) 52.2 ms
+ u.transpose(1, 2).contiguous() 3.7 ms
+ then .median(dim=2) 19.6 ms -> 23.4 ms total
+
+ Bit-identical — for an odd window the median is a *selection*, so any
+ correct algorithm returns the same element, and a test pins it against
+ ``scipy.ndimage.median_filter``. Only worth the extra copy for the larger
+ windows: at r=1 the transposed form measured 19.0 ms against 20.4 ms, so
+ the copy is skipped there and the memory is not spent.
+ """
+ import torch.nn.functional as F
+ spec = self.spec
+ k = 2 * int(radius) + 1
+ u = F.unfold(_pad_edges(self._t, radius, radius),
+ kernel_size=(k, k)) # (1, k*k, h*w)
+ shape = (1, 1, self.h, self.w)
+ if spec.median:
+ if k * k >= _MEDIAN_TRANSPOSE_TAPS:
+ med = u.transpose(1, 2).contiguous().median(dim=2).values
+ else:
+ med = u.median(dim=1).values
+ self._rank[(radius, "median")] = med.view(shape)
+ if spec.minimum:
+ self._rank[(radius, "minimum")] = u.amin(dim=1).view(shape)
+ if spec.maximum:
+ self._rank[(radius, "maximum")] = u.amax(dim=1).view(shape)
+
+ def membrane(self):
+ """``(1, R, h, w)`` responses to the rotated line kernels. Memoised."""
+ if self._membrane is None:
+ torch = self._torch
+ import torch.nn.functional as F
+ spec = self.spec
+ k = _membrane_kernels(spec.membrane_patch, spec.membrane_thickness,
+ spec.membrane_rotations)
+ weight = torch.as_tensor(k, device=self.device,
+ dtype=self.img.dtype).unsqueeze(1)
+ r = spec.membrane_patch // 2
+ self._membrane = F.conv2d(_pad_edges(self._t, r, r), weight)
+ return self._membrane
+
+ # -- channels ------------------------------------------------------------
+
+ def channel(self, family: str, args: tuple):
+ """One ``(h, w)`` channel tensor for a :func:`_channel_plan` entry."""
+ if family == "intensity":
+ return self.img
+ if family == "gaussian":
+ return self.blur(args[0])[0, 0]
+ if family == "dog":
+ return (self.blur(args[0]) - self.blur(args[1]))[0, 0]
+ if family == "sobel":
+ return self._sobel(args[0])
+ if family == "laplacian":
+ dyy, dxx, _ = self.second(args[0])
+ return (dyy + dxx)[0, 0]
+ if family in ("hessian_major", "hessian_minor"):
+ major, minor = self._hessian_eigs(args[0])
+ return major if family == "hessian_major" else minor
+ if family in ("median", "minimum", "maximum"):
+ return self.rank(args[0], family)[0, 0]
+ if family == "membrane":
+ return self._membrane_projection(args[0])
+ raise ValueError(f"unknown feature family {family!r}") # pragma: no cover
+
+ def _sobel(self, sigma: float):
+ """Sobel gradient magnitude of the blur at *sigma*.
+
+ skimage's kernels exactly: derivative ``[1, 0, -1]`` against smoothing
+ ``[1, 2, 1]/4``, and the magnitude divided by ``sqrt(2)`` as
+ ``skimage.filters.sobel`` does. The constant is irrelevant once the head
+ standardises its inputs, but matching it means a reader can compare a
+ channel against skimage without wondering about a factor.
+ """
+ torch = self._torch
+ b = self.blur(sigma)
+ der = torch.tensor([1.0, 0.0, -1.0], device=self.device,
+ dtype=self.img.dtype)
+ smo = torch.tensor([1.0, 2.0, 1.0], device=self.device,
+ dtype=self.img.dtype) / 4.0
+ gy = _sep_conv(torch, b, der, smo)
+ gx = _sep_conv(torch, b, smo, der)
+ return (torch.sqrt(gy * gy + gx * gx) / math.sqrt(2.0))[0, 0]
+
+ def _hessian_eigs(self, sigma: float):
+ """``(major, minor)`` eigenvalues of the 2×2 Hessian, signed.
+
+ Analytic, not ``linalg.eigvalsh``: for a symmetric 2×2 the eigenvalues are
+ ``tr/2 ± sqrt((tr/2)² - det)``, a handful of elementwise ops over the
+ image. ``eigvalsh`` batched over h·w 2×2 matrices was measured at 125 ms
+ on a 512² frame against **0.32 ms** for the closed form — a factor of 390,
+ so this is not an optimisation but the only viable option.
+
+ "Major"/"minor" are by signed value (major ≥ minor), which is what
+ distinguishes a bright blob (both strongly negative) from a ridge (one
+ negative, one near zero). Ordering by |value| would merge those two.
+ """
+ torch = self._torch
+ dyy, dxx, dxy = self.second(sigma)
+ half_tr = 0.5 * (dyy + dxx)
+ det = dyy * dxx - dxy * dxy
+ disc = torch.sqrt(torch.clamp(half_tr * half_tr - det, min=0.0))
+ return (half_tr + disc)[0, 0], (half_tr - disc)[0, 0]
+
+ def _membrane_projection(self, projection: str):
+ resp = self.membrane()[0] # (R, h, w)
+ if projection == "sum":
+ return resp.sum(dim=0)
+ if projection == "mean":
+ return resp.mean(dim=0)
+ if projection == "std":
+ return resp.std(dim=0, unbiased=False)
+ if projection == "median":
+ return resp.median(dim=0).values
+ if projection == "max":
+ return resp.amax(dim=0)
+ if projection == "min":
+ return resp.amin(dim=0)
+ raise ValueError( # pragma: no cover
+ f"unknown membrane projection {projection!r}")
+
+
+def _band_stack(image, spec: FeatureSpec, device):
+ """``(C, h, w)`` float32 tensor for one image (or row band)."""
+ torch = import_torch()
+ p = _Pass(image, spec, device)
+ plan = _channel_plan(spec)
+ out = torch.empty((len(plan), p.h, p.w), device=device, dtype=torch.float32)
+ for i, (family, _name, args) in enumerate(plan):
+ out[i] = p.channel(family, args)
+ return out
+
+
+# ── banding ──────────────────────────────────────────────────────────────────
+
+def band_budget_bytes(device=None) -> int:
+ """Bytes one row band may occupy, for *device*.
+
+ CPU keeps :data:`BAND_BYTES`. An accelerator gets a bigger band, because the
+ halo it re-featurises is pure overhead and fewer bands means less of it — but
+ clamped to :data:`GPU_BAND_BYTES` and to a fraction of *free* device memory,
+ since overshooting is catastrophic rather than merely slower (see
+ :data:`GPU_BAND_BYTES`).
+ """
+ kind = getattr(device, "type", None) or (str(device) if device else "cpu")
+ if kind not in ("cuda", "mps"):
+ return BAND_BYTES
+ budget = GPU_BAND_BYTES
+ try:
+ torch = import_torch()
+ if kind == "cuda":
+ free, _total = torch.cuda.mem_get_info(device)
+ budget = min(budget, int(free * _GPU_FREE_FRACTION))
+ except Exception as exc: # pragma: no cover
+ # A driver that will not report free memory is not a reason to guess
+ # high — fall back to the CPU budget, which is known to fit anywhere.
+ log.debug("device memory probe failed (%s); using the CPU band budget",
+ exc)
+ return BAND_BYTES
+ return max(BAND_BYTES, budget)
+
+
+def band_rows_for(spec: FeatureSpec, width: int,
+ budget_bytes: int | None = None, *, device=None) -> int:
+ """How many rows one band should cover, given the per-band memory budget.
+
+ The working set is bigger than the output stack — the pass also holds one
+ blur and three derivative images per sigma, plus an unfolded median window —
+ so the divisor counts those too rather than only ``n_channels``. Getting this
+ wrong does not produce a wrong answer, only a larger peak allocation, but the
+ plan's frame size leaves no headroom to be casual about it.
+
+ *budget_bytes* defaults to :func:`band_budget_bytes` for *device*, so the
+ same spec bands differently on a GPU than on the CPU. Pass it explicitly to
+ pin the banding regardless of hardware, which is what the banding tests do.
+ """
+ if budget_bytes is None:
+ budget_bytes = band_budget_bytes(device)
+ working = spec.n_channels + 4 * len(spec.sigmas) + 8
+ if (spec.median or spec.minimum or spec.maximum) and spec.rank_radii:
+ # One unfolded window is alive at a time (see `_Pass._compute_rank`), so
+ # it is the LARGEST radius that sets the peak, not the sum over radii.
+ taps = (2 * max(spec.rank_radii) + 1) ** 2
+ # ...except that the median's transposed reduction holds a second copy
+ # of that window while it runs. Unaccounted for, the band would be sized
+ # to a peak it does not actually have.
+ working += taps * (2 if spec.median and taps >= _MEDIAN_TRANSPOSE_TAPS
+ else 1)
+ rows = int(max(1, budget_bytes // max(1, working * int(width) * 4)))
+ return max(rows, 4 * spec.halo)
+
+
+def map_feature_bands(
+ frame,
+ spec: FeatureSpec | None = None,
+ *,
+ device=None,
+ fn: Callable[[int, int, Any], None],
+ band_rows: int | None = None,
+) -> None:
+ """Compute the stack in row bands and hand each to *fn*.
+
+ Calls ``fn(y0, y1, stack)`` where ``stack`` is a ``(C, y1-y0, W)`` tensor for
+ output rows ``y0:y1``. Each band is featurised with ``spec.halo`` extra rows
+ of real data above and below and then cropped, so **the banded result equals
+ the unbanded one exactly** for every row — the halo replaces what reflect
+ padding would otherwise have invented at a band boundary. Pinned by
+ ``test_particles_scribble.py::TestBanding``.
+
+ A callback rather than a generator on purpose: this holds the process-wide
+ accelerator lock (:func:`spyde.device_lock.accelerator_lock`) for the whole
+ traversal, and a generator abandoned by its consumer would hold that lock
+ until garbage collection.
+ """
+ spec = spec or FeatureSpec()
+ prepared = _as_prepared(frame, spec)
+ img = prepared.image
+ h, w = img.shape
+ if device is None:
+ device = select_device()
+
+ halo = spec.halo
+ rows = int(band_rows or band_rows_for(spec, w, device=device))
+
+ with accelerator_lock(device):
+ if rows >= h:
+ fn(0, h, _band_stack(img, spec, device))
+ return
+ y = 0
+ while y < h:
+ y1 = min(h, y + rows)
+ ky0, ky1 = max(0, y - halo), min(h, y1 + halo)
+ stack = _band_stack(img[ky0:ky1], spec, device)
+ fn(y, y1, stack[:, y - ky0: y1 - ky0])
+ y = y1
+
+
+def feature_tensor(frame, spec: FeatureSpec | None = None, *, device=None):
+ """``(C, H, W)`` float32 torch tensor of every channel in *spec*.
+
+ Parameters
+ ----------
+ frame
+ 2-D array, or a :class:`PreparedFrame` from :func:`prepare_frame`.
+ device
+ Torch device; ``None`` calls :func:`select_device`.
+
+ Notes
+ -----
+ This materialises the whole stack: ``C·H·W·4`` bytes, which is 2.4 GB for the
+ 36 default channels at 4096². That is fine for the interactive path (one
+ displayed frame, and the display is tiled well below that) and wrong for a
+ batch run over big frames — those go through :func:`map_feature_bands` or
+ :func:`sample_features`, which never hold more than one band.
+ """
+ spec = spec or FeatureSpec()
+ prepared = _as_prepared(frame, spec)
+ torch = import_torch()
+ if device is None:
+ device = select_device()
+ h, w = prepared.image.shape
+
+ # The allocation is a device submission too, so it goes inside the lock — the
+ # reentrant acquisition in `map_feature_bands` is free.
+ with accelerator_lock(device):
+ out = torch.empty((spec.n_channels, h, w), device=device,
+ dtype=torch.float32)
+
+ def take(y0: int, y1: int, stack) -> None:
+ out[:, y0:y1] = stack
+
+ map_feature_bands(prepared, spec, device=device, fn=take)
+ return out
+
+
+def feature_stack(frame, spec: FeatureSpec | None = None, *,
+ device=None) -> np.ndarray:
+ """:func:`feature_tensor` as a ``(C, H, W)`` numpy array.
+
+ The door the sklearn RandomForest parity reference comes through: the gate is
+ "same labels, same features, does the torch head agree with the forest", so
+ both sides must read the *identical* channels, not merely equivalent ones.
+ """
+ if device is None:
+ device = select_device()
+ # The device->host copy is a submission as well, hence the (reentrant) lock.
+ with accelerator_lock(device):
+ return feature_tensor(frame, spec, device=device).detach().cpu().numpy()
+
+
+def sample_features(frame, index, spec: FeatureSpec | None = None, *,
+ device=None):
+ """``(k, C)`` float32 tensor of the stack sampled at *k* pixels.
+
+ Parameters
+ ----------
+ index
+ Flat pixel indices into a ``(H, W)`` frame, ``(k,)`` integer; or ``(k, 2)``
+ ``(y, x)`` pairs.
+
+ Notes
+ -----
+ This is how training reads its data, and why training does not care how big
+ the frame is: the stack is produced band by band and only the sampled rows
+ survive, so the peak allocation is one band plus ``k·C`` floats. Scribbles are
+ thousands of pixels, so ``k·C`` is under a megabyte however large the frame.
+ """
+ spec = spec or FeatureSpec()
+ prepared = _as_prepared(frame, spec)
+ torch = import_torch()
+ if device is None:
+ device = select_device()
+
+ h, w = prepared.image.shape
+ idx = np.asarray(index)
+ if idx.ndim == 2 and idx.shape[1] == 2:
+ flat = idx[:, 0].astype(np.int64) * w + idx[:, 1].astype(np.int64)
+ else:
+ flat = idx.reshape(-1).astype(np.int64)
+ if flat.size and (flat.min() < 0 or flat.max() >= h * w):
+ raise IndexError(
+ f"pixel index outside 0..{h * w - 1} for a {h}x{w} frame")
+
+ ys, xs = np.divmod(flat, w)
+ with accelerator_lock(device): # see `feature_tensor` — allocation too
+ out = torch.empty((flat.size, spec.n_channels), device=device,
+ dtype=torch.float32)
+
+ def take(y0: int, y1: int, stack) -> None:
+ sel = np.flatnonzero((ys >= y0) & (ys < y1))
+ if not sel.size:
+ return
+ ty = torch.as_tensor(ys[sel] - y0, device=device, dtype=torch.long)
+ tx = torch.as_tensor(xs[sel], device=device, dtype=torch.long)
+ rows = torch.as_tensor(sel, device=device, dtype=torch.long)
+ out[rows] = stack[:, ty, tx].t()
+
+ map_feature_bands(prepared, spec, device=device, fn=take)
+ return out
+
+
+def feature_names(spec: FeatureSpec | None = None) -> list[str]:
+ """Channel names for *spec*, in the order :func:`feature_stack` returns them."""
+ return (spec or FeatureSpec()).channel_names()
diff --git a/spyde/particles/hull.py b/spyde/particles/hull.py
new file mode 100644
index 00000000..318d6341
--- /dev/null
+++ b/spyde/particles/hull.py
@@ -0,0 +1,271 @@
+"""
+hull.py — ``area_convex`` for every region at once, exactly, without Qhull.
+
+``solidity`` is ``area / area_convex``, and ``area_convex`` is the one property in
+SpyDE's measured set that is **not** a reduction over the raster: it is a convex
+hull per region. At the scale that matters — a real 4096² in-situ growth frame
+with 26 566 particles — ``regionprops_table``'s ``solidity`` alone is **30.5 s**,
+64% of the whole measurement (``benchmarks.md``). Every other column is now a
+``bincount`` (:mod:`spyde.particles.props`); this one is why the frame is still
+slow.
+
+The cost is not the hull. A particle here averages 33 pixels, so its hull is a
+dozen points — microseconds of arithmetic. The cost is **per region overhead**:
+two ``scipy.spatial.ConvexHull`` (Qhull) calls, a ``unique_rows``, a
+``grid_points_in_poly`` and the Python object around them, ~1.1 ms each,
+26 566 times. So this module keeps skimage's DEFINITION exactly and removes only
+the overhead, in a numba kernel that runs every region in one ``prange``.
+
+Exactly skimage's definition, in exact integer arithmetic
+---------------------------------------------------------
+``convex_hull_image`` (with its defaults, which is what ``regionprops`` uses):
+
+1. reduces the region to the first/last pixel of each row (``possible_hull``);
+2. replaces every such pixel ``(r, c)`` with the four **diamond offsets**
+ ``(r±0.5, c)``, ``(r, c±0.5)``;
+3. takes the convex hull of that point set;
+4. returns the grid points that are inside **or on** it
+ (``grid_points_in_poly(..., binarize=False)`` then ``labels >= 1``);
+5. ``area_convex`` is the count of those grid points.
+
+Each step is reproduced here, and the half-integer coordinates are the reason it
+can be done **without floating point at all**: doubling every coordinate turns the
+offsets into integers (``(2r±1, 2c)``, ``(2r, 2c±1)``), so the monotone-chain
+cross products and the inside test are exact ``int64``. There is no tolerance to
+tune and no tie to lose — a grid point is inside iff every edge's cross product is
+non-negative, which is a comparison of integers.
+
+Reducing to row extremes first is not an approximation: a pixel strictly between
+the first and last of its own row lies in their convex hull, so it can never be a
+hull vertex. (skimage also keeps column extremes; a superset of the vertices gives
+the same hull either way.)
+
+Why numba and not numpy
+-----------------------
+A hull is a stack algorithm — sequential per region, and the regions are tiny and
+numerous, which is the one shape ``bincount`` cannot express. ``numba.njit`` with
+``nogil=True`` gives the two things the batch run needs from it: the arithmetic
+compiled instead of interpreted, and **the GIL released**, which is what lets a
+dask worker's task slots be worth more than one core (``benchmarks.md``:
+``regionprops_table`` in four threads measured 0.93x of serial).
+
+``numba`` is optional — :func:`convex_areas` returns ``None`` when it is missing
+or the kernel refuses to compile, and :mod:`spyde.particles.measure` falls back to
+``regionprops_table``. Nothing here is on a GPU, so there is no device lock to
+take.
+"""
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+_KERNEL = [None] # compiled lazily, once per process
+_FAILED = [False]
+
+
+def _build_kernel():
+ """Compile the per-region hull kernel, or return None if numba is unusable."""
+ if _KERNEL[0] is not None or _FAILED[0]:
+ return _KERNEL[0]
+ try:
+ import numba
+ except Exception as exc: # pragma: no cover
+ log.info("[particles] numba unavailable (%s); solidity stays on "
+ "regionprops", exc)
+ _FAILED[0] = True
+ return None
+
+ @numba.njit(cache=True, nogil=True, parallel=True, fastmath=False)
+ def _kernel(starts, counts, rows, cols, out): # pragma: no cover
+ n = starts.shape[0]
+ for i in numba.prange(n):
+ s = starts[i]
+ m = counts[i]
+ if m <= 0:
+ out[i] = 0
+ continue
+
+ # ── step 1: first/last pixel of each row, in DOUBLED coordinates.
+ # The pixels arrive in raster order, so a row's run is contiguous.
+ cand_x = np.empty(2 * m, np.int64)
+ cand_y = np.empty(2 * m, np.int64)
+ nc = 0
+ j = s
+ end = s + m
+ while j < end:
+ r = rows[j]
+ k = j
+ while k + 1 < end and rows[k + 1] == r:
+ k += 1
+ cand_x[nc] = 2 * cols[j]
+ cand_y[nc] = 2 * r
+ nc += 1
+ if k != j:
+ cand_x[nc] = 2 * cols[k]
+ cand_y[nc] = 2 * r
+ nc += 1
+ j = k + 1
+
+ # ── step 2: the diamond offsets, still integers once doubled.
+ np_ = 4 * nc
+ px = np.empty(np_, np.int64)
+ py = np.empty(np_, np.int64)
+ key = np.empty(np_, np.int64)
+ for a in range(nc):
+ bx = cand_x[a]
+ by = cand_y[a]
+ px[4 * a + 0] = bx - 1
+ py[4 * a + 0] = by
+ px[4 * a + 1] = bx + 1
+ py[4 * a + 1] = by
+ px[4 * a + 2] = bx
+ py[4 * a + 2] = by - 1
+ px[4 * a + 3] = bx
+ py[4 * a + 3] = by + 1
+ # Sort by (x, y) for the monotone chain. One int64 key rather than a
+ # lexsort: the doubled coordinates are bounded by 2*frame_size+1, so
+ # a 2^20 stride is exact for any frame up to ~262 000 px a side.
+ for a in range(np_):
+ key[a] = px[a] * 1048576 + py[a]
+ order = np.argsort(key)
+
+ # ── step 3: monotone chain over the sorted points.
+ hx = np.empty(2 * np_ + 1, np.int64)
+ hy = np.empty(2 * np_ + 1, np.int64)
+ nh = 0
+ for a in range(np_):
+ x = px[order[a]]
+ y = py[order[a]]
+ while nh >= 2 and (
+ (hx[nh - 1] - hx[nh - 2]) * (y - hy[nh - 2])
+ - (hy[nh - 1] - hy[nh - 2]) * (x - hx[nh - 2])) <= 0:
+ nh -= 1
+ hx[nh] = x
+ hy[nh] = y
+ nh += 1
+ lower = nh
+ for a in range(np_ - 2, -1, -1):
+ x = px[order[a]]
+ y = py[order[a]]
+ while nh > lower and (
+ (hx[nh - 1] - hx[nh - 2]) * (y - hy[nh - 2])
+ - (hy[nh - 1] - hy[nh - 2]) * (x - hx[nh - 2])) <= 0:
+ nh -= 1
+ hx[nh] = x
+ hy[nh] = y
+ nh += 1
+ nh -= 1 # the closing point repeats hull[0]
+
+ if nh < 3:
+ out[i] = m # cannot happen once offset; be safe
+ continue
+
+ # ── steps 4-5: count grid points inside or ON the hull. The hull of
+ # the offsets extends 0.5 px past the region, so it can only cover
+ # grid points inside the region's own bbox.
+ y0 = rows[s]
+ y1 = rows[end - 1]
+ x0 = cols[s]
+ x1 = cols[s]
+ for a in range(s + 1, end):
+ c = cols[a]
+ if c < x0:
+ x0 = c
+ if c > x1:
+ x1 = c
+
+ total = 0
+ for yy in range(y0, y1 + 1):
+ gy = 2 * yy
+ started = False
+ for xx in range(x0, x1 + 1):
+ gx = 2 * xx
+ inside = True
+ for e in range(nh):
+ e2 = e + 1
+ if e2 == nh:
+ e2 = 0
+ if ((hx[e2] - hx[e]) * (gy - hy[e])
+ - (hy[e2] - hy[e]) * (gx - hx[e])) < 0:
+ inside = False
+ break
+ if inside:
+ total += 1
+ started = True
+ elif started:
+ break # convex: one run per row
+ out[i] = total
+
+ _KERNEL[0] = _kernel
+ return _kernel
+
+
+def convex_areas(lab: np.ndarray, labels: np.ndarray, counts: np.ndarray,
+ ) -> np.ndarray | None:
+ """``area_convex`` per entry of *labels*, or None if numba is unavailable.
+
+ Parameters
+ ----------
+ lab
+ The ``(h, w)`` int label image.
+ labels
+ Ascending labels present in *lab* — the same order
+ ``regionprops_table`` emits.
+ counts
+ ``np.bincount(lab.ravel())``, i.e. pixel counts indexed BY LABEL.
+
+ Notes
+ -----
+ Memory: one int64 index array and two int32 coordinate arrays sized by the
+ FOREGROUND pixel count, nothing sized by the number of regions squared and
+ nothing that materialises a per-region mask (CLAUDE.md § Memory Safety).
+ """
+ kernel = _build_kernel()
+ if kernel is None:
+ return None
+
+ h, w = lab.shape
+ flat = lab.reshape(-1)
+ nz = np.flatnonzero(flat)
+ if nz.size == 0:
+ return np.zeros(labels.shape, np.int64)
+ lab_of = flat[nz]
+ # Stable sort groups the pixels by label while KEEPING raster order inside
+ # each group, which is what the kernel's row-run scan relies on. numpy uses
+ # a radix sort for integer keys here, so this is a linear pass.
+ order = np.argsort(lab_of, kind="stable")
+ nz = nz[order]
+ rows = (nz // w).astype(np.int32)
+ cols = (nz - rows.astype(np.int64) * w).astype(np.int32)
+
+ grp = counts[labels].astype(np.int64)
+ starts = np.zeros(labels.size, np.int64)
+ np.cumsum(grp[:-1], out=starts[1:])
+
+ out = np.zeros(labels.size, np.int64)
+ try:
+ kernel(starts, grp, rows, cols, out)
+ except Exception as exc: # pragma: no cover
+ log.warning("[particles] hull kernel failed (%r); solidity falls back "
+ "to regionprops", exc)
+ _FAILED[0] = True
+ _KERNEL[0] = None
+ return None
+ return out
+
+
+def warmup() -> bool:
+ """Compile the kernel on a 3-pixel toy image. Returns True if numba is live.
+
+ Worth calling once on a dask worker: the first ``njit`` call pays the
+ compile, and paying it inside the first measured frame makes that frame look
+ like a regression. ``cache=True`` means it is paid once per machine, not once
+ per process, but a cold cache still has to build.
+ """
+ lab = np.zeros((4, 4), np.int32)
+ lab[1:3, 1:3] = 1
+ counts = np.bincount(lab.reshape(-1))
+ return convex_areas(lab, np.array([1], np.int64), counts) is not None
diff --git a/spyde/particles/intensity.py b/spyde/particles/intensity.py
new file mode 100644
index 00000000..3f12d81f
--- /dev/null
+++ b/spyde/particles/intensity.py
@@ -0,0 +1,310 @@
+"""
+intensity.py — the intensity columns as label-wise reductions, and the ring as a kernel.
+
+:mod:`spyde.particles.props` removed ``regionprops_table``'s Python loop over
+regions and :mod:`spyde.particles.hull` removed ``solidity``'s. What was left of
+``measure_frame`` was two loops of exactly the same shape, and this module is one
+of them: ``_fill_intensity`` was **4.9 s of a 10.2 s frame** on a real 4096²
+in-situ growth raster with 26 566 particles (``benchmarks.md``), spent almost
+entirely on per-region overhead — a bbox crop, a mask, a boolean take and a
+``scipy.ndimage.binary_dilation``, 26 566 times, all of it holding the GIL.
+
+The four columns split cleanly in two, and only one of them is hard:
+
+* ``intensity_mean`` / ``intensity_max`` / ``intensity_std`` are **label-wise
+ reductions over the foreground pixels** — ``bincount`` with weights for the
+ sums, one label-grouped ``np.maximum.reduceat`` for the max. Same shape as the
+ moments in :mod:`~spyde.particles.props`, and O(foreground) rather than
+ O(regions x crop).
+* ``background`` is not. It is the mean intensity of the pixels a **dilation of
+ this particle by ``ring``** adds and that belong to no particle — a *per
+ particle* neighbourhood that overlapping neighbours may each claim, so it is
+ not a partition of the raster and no ``bincount`` expresses it.
+ :func:`ring_backgrounds` keeps the definition exactly (an iterated
+ 4-connected dilation inside the same padded bbox crop, which is what
+ ``binary_dilation``'s default structure and ``border_value=0`` do) and moves
+ it into a ``numba`` kernel that runs every region in one ``prange`` with the
+ GIL released, the way :mod:`~spyde.particles.hull` does for the convex hull.
+
+Parity
+------
+The pixel SETS are identical by construction — every statistic here is taken over
+exactly the pixels the per-region crop selected, including the finite-only filter
+that keeps a NaN-padded drift-corrected border from poisoning a particle that
+touches it. What differs is **summation order**: ``np.mean``/``np.std`` reduce
+pairwise, ``bincount`` and the kernel accumulate sequentially, so the float64
+intermediates disagree at ~1e-16 relative. The rows are stored in float32, where
+that is 9 orders below the last bit, and the parity test asserts the stored
+columns come out **bit-identical** on a scene of thousands of ragged regions.
+
+``intensity_std`` is computed the way ``np.std`` computes it — mean first, then
+the mean of squared deviations about it — and NOT as ``E[x²] - E[x]²``, which is
+algebraically equal and loses digits to cancellation exactly where the variance
+is small compared to the mean, which is the normal case for a bright particle.
+
+Memory: the working set is one float64 and one index array sized by the
+FOREGROUND pixel count, plus one padded bbox crop per region at a time inside the
+kernel (CLAUDE.md § Memory Safety). Nothing scales with the number of regions
+squared and nothing materialises a full-frame per-region mask.
+"""
+from __future__ import annotations
+
+import logging
+
+import numpy as np
+
+log = logging.getLogger(__name__)
+
+_RING_KERNEL = [None] # compiled lazily, once per process
+_RING_FAILED = [False]
+
+
+def label_intensity_stats(
+ lab: np.ndarray,
+ inten: np.ndarray,
+ labels: np.ndarray,
+ *,
+ n_max: int | None = None,
+) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
+ """``(mean, max, std_over_max)`` per entry of *labels*, over FINITE pixels.
+
+ Reproduces, for every region at once, what the per-region crop computed::
+
+ vals = sub_int[sub_lab == lbl]
+ vals = vals[np.isfinite(vals)]
+ mean, mx = vals.mean(), vals.max()
+ std = vals.std() / mx if mx else nan
+
+ A region with no finite pixel gets NaN in all three, and a region whose
+ finite maximum is exactly 0 gets NaN for the normalised deviation — both are
+ the per-region path's own behaviour, not a new convention.
+
+ *labels* must be ascending and must be exactly the labels present, which is
+ what :func:`spyde.particles.props.label_props` emits.
+ """
+ lab = np.asarray(lab)
+ inten = np.asarray(inten, dtype=np.float64)
+ labels = np.asarray(labels, dtype=np.int64)
+ n = labels.size
+ nan3 = np.full(n, np.nan)
+ if n == 0:
+ return nan3, nan3.copy(), nan3.copy()
+
+ flat = lab.reshape(-1)
+ nz = np.flatnonzero(flat)
+ if nz.size == 0:
+ return nan3, nan3.copy(), nan3.copy()
+
+ if n_max is None:
+ n_max = int(labels[-1])
+ # label value -> row of the output. Sparse label images are normal here (any
+ # upstream filter re-tags), so the row index is dense and the label is not.
+ dense = np.zeros(int(n_max) + 1, np.int64)
+ dense[labels] = np.arange(n, dtype=np.int64)
+
+ v = inten.reshape(-1)[nz]
+ fin = np.isfinite(v)
+ if not fin.all():
+ v = v[fin]
+ li = dense[flat[nz[fin]]]
+ else:
+ li = dense[flat[nz]]
+ del nz, fin
+
+ cnt = np.bincount(li, minlength=n)[:n]
+ have = cnt > 0
+ cnt_f = cnt.astype(np.float64)
+
+ s = np.bincount(li, weights=v, minlength=n)[:n]
+ with np.errstate(divide="ignore", invalid="ignore"):
+ mean = s / cnt_f
+ mean[~have] = np.nan
+
+ # Two-pass deviation, matching `np.std`: the one-pass E[x^2]-E[x]^2 is
+ # algebraically the same and cancels away the digits that matter when a
+ # particle is bright and uniform.
+ dev = v - mean[li]
+ dev *= dev
+ s2 = np.bincount(li, weights=dev, minlength=n)[:n]
+ del dev
+ with np.errstate(divide="ignore", invalid="ignore"):
+ std = np.sqrt(s2 / cnt_f)
+
+ # Per-label max. Grouping by label is a stable (radix) sort of the
+ # foreground, after which one `maximum.reduceat` covers every region; the
+ # alternative, `np.maximum.at`, is an unbuffered ufunc call per pixel and is
+ # ~50x slower than the sort it avoids.
+ mx = np.full(n, np.nan)
+ if li.size:
+ order = np.argsort(li, kind="stable")
+ vs = v[order]
+ starts = np.zeros(n, np.int64)
+ np.cumsum(cnt[:-1], out=starts[1:])
+ mx[have] = np.maximum.reduceat(vs, starts[have])
+ del order, vs
+
+ with np.errstate(divide="ignore", invalid="ignore"):
+ std_norm = np.where(mx != 0, std / mx, np.nan)
+ std_norm[~have] = np.nan
+ return mean, mx, std_norm
+
+
+def _build_ring_kernel():
+ """Compile the per-region ring kernel, or return None if numba is unusable."""
+ if _RING_KERNEL[0] is not None or _RING_FAILED[0]:
+ return _RING_KERNEL[0]
+ try:
+ import numba
+ except Exception as exc: # pragma: no cover
+ log.info("[particles] numba unavailable (%s); the background ring stays "
+ "on the per-region loop", exc)
+ _RING_FAILED[0] = True
+ return None
+
+ @numba.njit(cache=True, nogil=True, parallel=True, fastmath=False)
+ def _kernel(lab, inten, labels, bb, ring, sums, counts): # pragma: no cover
+ h, w = lab.shape
+ n = labels.shape[0]
+ for i in numba.prange(n):
+ lbl = labels[i]
+ # The SAME crop the per-region path takes: the bbox grown by
+ # ring + 1 and clipped to the frame, so a dilation by `ring` fits
+ # inside it and a region at the frame edge is truncated exactly as
+ # `binary_dilation` truncates it there.
+ y0 = bb[i, 0] - ring - 1
+ x0 = bb[i, 1] - ring - 1
+ y1 = bb[i, 2] + ring + 1
+ x1 = bb[i, 3] + ring + 1
+ if y0 < 0:
+ y0 = 0
+ if x0 < 0:
+ x0 = 0
+ if y1 > h:
+ y1 = h
+ if x1 > w:
+ x1 = w
+ hh = y1 - y0
+ ww = x1 - x0
+ if hh <= 0 or ww <= 0:
+ sums[i] = 0.0
+ counts[i] = 0
+ continue
+
+ cur = np.zeros((hh, ww), np.uint8)
+ for r in range(hh):
+ for c in range(ww):
+ if lab[y0 + r, x0 + c] == lbl:
+ cur[r, c] = 1
+
+ # `binary_dilation(m, iterations=ring)` with scipy's default
+ # structure is `generate_binary_structure(2, 1)` — the 4-connected
+ # cross — applied `ring` times, i.e. everything within city-block
+ # distance `ring`. `border_value=0` is the bounds check below.
+ if ring > 0:
+ nxt = np.zeros((hh, ww), np.uint8)
+ for _it in range(ring):
+ for r in range(hh):
+ for c in range(ww):
+ v = cur[r, c]
+ if v == 0:
+ if r > 0 and cur[r - 1, c] != 0:
+ v = 1
+ elif r + 1 < hh and cur[r + 1, c] != 0:
+ v = 1
+ elif c > 0 and cur[r, c - 1] != 0:
+ v = 1
+ elif c + 1 < ww and cur[r, c + 1] != 0:
+ v = 1
+ nxt[r, c] = v
+ tmp = cur
+ cur = nxt
+ nxt = tmp
+
+ # The ring is what the dilation added MINUS anything belonging to a
+ # neighbouring particle: a touching particle's body is not this
+ # one's background.
+ s = 0.0
+ k = 0
+ for r in range(hh):
+ for c in range(ww):
+ if cur[r, c] != 0 and lab[y0 + r, x0 + c] == 0:
+ val = inten[y0 + r, x0 + c]
+ if np.isfinite(val):
+ s += val
+ k += 1
+ sums[i] = s
+ counts[i] = k
+
+ _RING_KERNEL[0] = _kernel
+ return _kernel
+
+
+def ring_backgrounds(lab: np.ndarray, inten: np.ndarray, labels: np.ndarray,
+ bboxes: np.ndarray, ring: int) -> np.ndarray | None:
+ """Mean background per entry of *labels*, or None if numba is unavailable.
+
+ Parameters
+ ----------
+ lab
+ The ``(h, w)`` int label image.
+ inten
+ The ``(h, w)`` float64 intensity image. May contain NaN; NaN pixels are
+ excluded from the mean rather than coerced to zero, which would invent a
+ dark rim on every particle touching a drift-corrected border.
+ labels
+ Ascending labels present in *lab*.
+ bboxes
+ ``(n, 4)`` int ``(y0, x0, y1, x1)`` per entry of *labels* — the same
+ tight bboxes the property table reports.
+ ring
+ Dilation width in pixels. ``0`` returns all-NaN (the caller leaves the
+ column unset), matching the per-region path's ``if ring > 0`` guard.
+
+ Notes
+ -----
+ Memory: one padded bbox crop at a time per thread, plus two float/int arrays
+ of length ``n``. Nothing full-frame is allocated (CLAUDE.md § Memory Safety).
+ """
+ labels = np.asarray(labels, dtype=np.int64)
+ n = labels.size
+ if n == 0:
+ return np.zeros((0,), np.float64)
+ if int(ring) <= 0:
+ return np.full(n, np.nan)
+
+ kernel = _build_ring_kernel()
+ if kernel is None:
+ return None
+
+ bb = np.ascontiguousarray(bboxes, dtype=np.int64)
+ sums = np.zeros(n, np.float64)
+ counts = np.zeros(n, np.int64)
+ try:
+ kernel(lab, np.asarray(inten, np.float64), labels, bb, int(ring),
+ sums, counts)
+ except Exception as exc: # pragma: no cover
+ log.warning("[particles] ring kernel failed (%r); background falls back "
+ "to the per-region dilation", exc)
+ _RING_FAILED[0] = True
+ _RING_KERNEL[0] = None
+ return None
+
+ out = np.full(n, np.nan)
+ have = counts > 0
+ out[have] = sums[have] / counts[have].astype(np.float64)
+ return out
+
+
+def warmup() -> bool:
+ """Compile the ring kernel on a toy image. Returns True if numba is live.
+
+ Worth calling once on a dask worker, for the reason
+ :func:`spyde.particles.hull.warmup` exists: paying the first ``njit``
+ compile inside the first measured frame makes that frame look like a
+ regression.
+ """
+ lab = np.zeros((8, 8), np.int32)
+ lab[3:5, 3:5] = 1
+ inten = np.ones((8, 8), np.float64)
+ bb = np.array([[3, 3, 5, 5]], np.int64)
+ return ring_backgrounds(lab, inten, np.array([1], np.int64), bb, 2) is not None
diff --git a/spyde/particles/measure.py b/spyde/particles/measure.py
new file mode 100644
index 00000000..1bef6102
--- /dev/null
+++ b/spyde/particles/measure.py
@@ -0,0 +1,440 @@
+"""
+measure.py — turn a label image into calibrated particle property rows.
+
+ParticleSpy's measured-property set, converted to physical units exactly once,
+here. Every downstream consumer — the table, the histogram, the kymograph, the
+CSV — reads the already-calibrated numbers, so there is one place a unit can be
+wrong instead of a dozen.
+
+The property table itself no longer comes from ``regionprops_table``
+--------------------------------------------------------------------
+``regionprops_table`` walks a Python object per region, so its cost is per
+PARTICLE — and a real 4096² in-situ growth frame has **26 566** of them. Measured
+(``benchmarks.md``): 53.5 s to measure a frame against 3.0 s to segment it, of
+which ``solidity`` alone was 30.5 s and the axis/eccentricity trio 12 s. Whole
+frame now: **1.37 s**.
+
+Four modules replace it, keeping skimage's definitions to the bit:
+
+* :mod:`spyde.particles.props` — every column that is a **label-wise reduction
+ over the raster** (``bincount``: area, centroid, bbox, equivalent diameter, the
+ second central moments behind the axes and eccentricity, and the perimeter's
+ border-crossing weights). 43.7 s → 1.1 s.
+* :mod:`spyde.particles.hull` — ``solidity``, the one property that is a convex
+ hull per region rather than a reduction, in a numba kernel that does every
+ region in one ``prange`` with the GIL released. 30.2 s → 0.18 s, and the hull
+ is reproduced in exact integer arithmetic, so ``area_convex`` matches skimage
+ on **all 26 566** regions with zero differing pixels.
+* :mod:`spyde.particles.intensity` — the intensity statistics as the same
+ ``bincount`` the moments turned out to be, and the local background RING (a
+ dilation per particle, which overlapping neighbours may each claim, so it is
+ not a partition and no ``bincount`` expresses it) as a second numba kernel.
+ 4.92 s → 0.19 s, and all four columns are **bit-identical** on the real frame.
+* :mod:`spyde.particles.contours` — marching squares and the segment assembly
+ behind them, for every region in one ``prange``. 4.77 s → 0.15 s, and the
+ FILLED polygon — which is what ``render_frame`` and ``mask_at`` consume — is
+ identical on **26 566 of 26 566** regions. The vertices are NOT, and cannot be:
+ a closed contour is a cycle and the two tracers cut it at different vertices.
+ See that module for why the fill is the right gate and vertex identity is not.
+
+``SPYDE_PARTICLE_PROPS=legacy`` (or ``measure_frame(..., fast=False)``) restores
+the ``regionprops_table`` + ``find_contours`` path. It is what the parity tests
+compare against, and what runs if numba cannot compile.
+
+Two things that look like details and are not:
+
+* **NaN in the intensity image is respected, not coerced.** A drift-corrected
+ frame has a NaN-padded border (``spyde.drift.warp``). Letting ``np.nan`` reach
+ a plain ``mean`` makes every particle touching the border report NaN intensity;
+ coercing NaN to zero invents a dark rim that biases the measurement instead.
+ Both are wrong, so intensity statistics are computed over finite pixels only,
+ and a particle with no finite pixels is dropped (see :func:`measure_frame`).
+* **Circularity uses ParticleSpy's convention**, ``4*pi*area / perimeter^2``,
+ which is 1 for a perfect disc. It is computed in PIXELS before calibration —
+ it is dimensionless, so calibrating area and perimeter separately and then
+ dividing would introduce a spurious ``scale`` factor.
+"""
+from __future__ import annotations
+
+import os
+
+import numpy as np
+
+from spyde.signals.particles import COL, N_COLUMNS
+
+# regionprops names we ask for. Kept minimal: each extra property is another pass
+# over every region, and the ones omitted here are cheap to derive from these.
+# This is now the LEGACY path's property list and the parity test's reference —
+# `spyde.particles.props.label_props` produces exactly these columns.
+_PROPS = (
+ "label",
+ "centroid",
+ "area",
+ "equivalent_diameter_area",
+ "major_axis_length",
+ "minor_axis_length",
+ "perimeter",
+ "eccentricity",
+ "solidity",
+ "bbox",
+)
+
+
+def _fast_default() -> bool:
+ """Whether the vectorised property path is on. ``SPYDE_PARTICLE_PROPS=legacy``
+ turns it off, for a bug report or an A/B against skimage."""
+ return os.environ.get("SPYDE_PARTICLE_PROPS", "").lower() not in (
+ "legacy", "skimage", "0", "off")
+
+
+def property_table(labels: np.ndarray, *, fast: bool | None = None) -> dict:
+ """The per-region property table, as ``regionprops_table`` would return it.
+
+ Split out from :func:`measure_frame` so the parity test can compare the two
+ paths column by column on a real label image, which is the only reason it is
+ safe to have replaced the reference implementation at all.
+ """
+ if fast is None:
+ fast = _fast_default()
+ if fast:
+ from spyde.particles.props import label_props
+ return label_props(labels)
+ from skimage.measure import regionprops_table
+ return regionprops_table(labels, properties=_PROPS)
+
+
+def warm_kernels() -> None:
+ """Compile every numba kernel ``measure_frame`` uses, before the first frame.
+
+ Three of them now (hull, ring, contours), and paying any of their first-use
+ compiles inside a measured frame makes that frame look like a regression.
+ Idempotent; each module short-circuits once compiled.
+ """
+ from spyde.particles.contours import warmup as warm_contours
+ from spyde.particles.hull import warmup as warm_hull
+ from spyde.particles.intensity import warmup as warm_ring
+
+ warm_hull()
+ warm_ring()
+ warm_contours()
+
+
+def measure_frame(
+ labels: np.ndarray,
+ intensity: np.ndarray | None = None,
+ *,
+ t: int = 0,
+ scale: float = 1.0,
+ background_ring: int = 3,
+ min_area_px: int = 0,
+ fast: bool | None = None,
+ want_contours: bool = True,
+) -> tuple[np.ndarray, list[np.ndarray]]:
+ """Measure every instance in *labels*.
+
+ Parameters
+ ----------
+ labels
+ ``(h, w)`` integer label image; 0 is background.
+ intensity
+ Optional source frame for intensity statistics. May contain NaN (a
+ drift-corrected border does); NaN pixels are excluded rather than coerced.
+ t
+ Frame index stamped into the ``t`` column.
+ scale
+ Pixel size in physical units. Lengths are multiplied by it and areas by
+ its square; dimensionless quantities are untouched.
+ background_ring
+ Width in pixels of the dilated ring outside each particle used for the
+ local ``background`` measurement. 0 disables it (leaves NaN).
+ min_area_px
+ Drop instances smaller than this many pixels. Applied here, in pixels,
+ because it is a detector-resolution question, not a physical one.
+ fast
+ Force the vectorised property path on/off. ``None`` reads
+ ``SPYDE_PARTICLE_PROPS``. The two paths are asserted equal column by
+ column in ``test_particles_props_parity.py``.
+ want_contours
+ Extract the per-instance outlines. ``False`` returns an EMPTY contour
+ list and skips that work — for a caller that only needs the measured
+ properties. Every stored result needs the outlines (the 1:1
+ correspondence below is what ``SpyDEParticles.from_frames`` requires),
+ so this is only safe for a transient consumer such as a live preview
+ that has already decided the instances are too numerous to draw.
+
+ Returns
+ -------
+ (rows, contours)
+ ``rows`` is ``(n, N_COLUMNS)`` float32 matching
+ :data:`spyde.signals.particles.COLUMNS`, with ``track_id`` set to -1.
+ ``contours`` is a list of ``(k, 2)`` int16 ``(y, x)`` outlines, one per
+ row and in the same order — the 1:1 correspondence
+ ``SpyDEParticles.from_frames`` requires. Empty when *want_contours* is
+ False, which breaks that correspondence by design.
+ """
+ lab = np.asarray(labels)
+ if lab.ndim != 2:
+ raise ValueError(f"labels must be 2-D; got shape {lab.shape}")
+ if intensity is not None:
+ inten = np.asarray(intensity, dtype=np.float64)
+ if inten.shape != lab.shape:
+ raise ValueError(
+ f"intensity shape {inten.shape} != labels shape {lab.shape}"
+ )
+ else:
+ inten = None
+
+ if lab.max() <= 0:
+ return np.zeros((0, N_COLUMNS), np.float32), []
+
+ tbl = property_table(lab, fast=fast)
+ n = len(tbl["label"])
+
+ area_px = tbl["area"].astype(np.float64)
+ perim_px = tbl["perimeter"].astype(np.float64)
+
+ keep = area_px >= float(min_area_px)
+
+ # Dimensionless, so computed in pixels BEFORE calibration.
+ with np.errstate(divide="ignore", invalid="ignore"):
+ circularity = np.where(perim_px > 0,
+ 4.0 * np.pi * area_px / perim_px ** 2, np.nan)
+
+ rows = np.zeros((n, N_COLUMNS), dtype=np.float32)
+ rows[:, COL["t"]] = float(t)
+ rows[:, COL["label"]] = tbl["label"]
+ rows[:, COL["y"]] = tbl["centroid-0"] * scale
+ rows[:, COL["x"]] = tbl["centroid-1"] * scale
+ rows[:, COL["area"]] = area_px * scale ** 2
+ rows[:, COL["equiv_diameter"]] = tbl["equivalent_diameter_area"] * scale
+ rows[:, COL["major_axis"]] = tbl["major_axis_length"] * scale
+ rows[:, COL["minor_axis"]] = tbl["minor_axis_length"] * scale
+ rows[:, COL["perimeter"]] = perim_px * scale
+ rows[:, COL["circularity"]] = circularity
+ rows[:, COL["eccentricity"]] = tbl["eccentricity"]
+ rows[:, COL["solidity"]] = tbl["solidity"]
+ rows[:, COL["bbox_y0"]] = tbl["bbox-0"]
+ rows[:, COL["bbox_x0"]] = tbl["bbox-1"]
+ rows[:, COL["bbox_y1"]] = tbl["bbox-2"]
+ rows[:, COL["bbox_x1"]] = tbl["bbox-3"]
+ rows[:, COL["track_id"]] = -1.0
+ rows[:, COL["intensity_mean"]] = np.nan
+ rows[:, COL["intensity_max"]] = np.nan
+ rows[:, COL["intensity_std"]] = np.nan
+ rows[:, COL["background"]] = np.nan
+
+ if inten is not None:
+ _fill_intensity(rows, lab, inten, tbl, keep, background_ring, fast=fast)
+
+ # Outlines are ~half the cost of a measure once the instance count runs into
+ # the thousands (measured on a 1024² over-segmented frame: 647 ms total, of
+ # which the 4873 contours are the bulk; the same frame filtered to 17
+ # instances measures in 62 ms). The live PREVIEW asks for them only when it
+ # is actually going to draw them, which it will not do above the overlay's
+ # draw cap. Every other caller keeps the default and is unaffected.
+ contours = _contours(lab, tbl, fast=fast) if want_contours else []
+
+ rows = rows[keep]
+ if want_contours:
+ contours = [c for c, k in zip(contours, keep) if k]
+ # Score LAST: it is derived from the intensity columns filled above, so it
+ # costs no extra pass over the frame. That is what lets the caret filter on
+ # it without re-segmenting — see `particle_scores`.
+ rows[:, COL["score"]] = particle_scores(rows)
+ return np.ascontiguousarray(rows), contours
+
+
+def _table_bboxes(tbl) -> np.ndarray:
+ """``(n, 4)`` int64 ``(y0, x0, y1, x1)`` from the property table's columns."""
+ return np.stack([np.asarray(tbl[f"bbox-{k}"], np.int64) for k in range(4)],
+ axis=1)
+
+
+def _fill_intensity(rows, lab, inten, tbl, keep, ring: int, *,
+ fast: bool | None = None) -> None:
+ """Intensity statistics over FINITE pixels only, plus a local background ring.
+
+ ``intensity_mean/max/std`` are label-wise reductions over the foreground and
+ go through :func:`spyde.particles.intensity.label_intensity_stats`; the
+ background ring is a per-particle dilation that overlapping neighbours may
+ each claim, so it goes through that module's numba kernel instead. The
+ per-region loop below is the definition both are checked against, and the
+ path that runs when numba is unavailable.
+ """
+ if fast is None:
+ fast = _fast_default()
+ if fast and _fill_intensity_fast(rows, lab, inten, tbl, keep, ring):
+ return
+
+ from scipy.ndimage import binary_dilation
+
+ h, w = lab.shape
+ for i in range(len(tbl["label"])):
+ if not keep[i]:
+ continue
+ lbl = int(tbl["label"][i])
+ y0, x0 = int(tbl["bbox-0"][i]), int(tbl["bbox-1"][i])
+ y1, x1 = int(tbl["bbox-2"][i]), int(tbl["bbox-3"][i])
+
+ # Pad by the ring width so the dilated boundary fits inside the crop.
+ py0, px0 = max(0, y0 - ring - 1), max(0, x0 - ring - 1)
+ py1, px1 = min(h, y1 + ring + 1), min(w, x1 + ring + 1)
+ sub_lab = lab[py0:py1, px0:px1]
+ sub_int = inten[py0:py1, px0:px1]
+ m = sub_lab == lbl
+
+ vals = sub_int[m]
+ vals = vals[np.isfinite(vals)]
+ if vals.size:
+ rows[i, COL["intensity_mean"]] = vals.mean()
+ rows[i, COL["intensity_max"]] = vals.max()
+ # Normalised by the max, matching ParticleSpy, so it is comparable
+ # between particles of very different absolute brightness.
+ mx = vals.max()
+ rows[i, COL["intensity_std"]] = (vals.std() / mx) if mx else np.nan
+
+ if ring > 0:
+ grown = binary_dilation(m, iterations=int(ring))
+ # The ring is what the dilation added, minus anything belonging to a
+ # NEIGHBOURING particle — otherwise a touching particle's body is
+ # measured as this one's background, which is exactly backwards.
+ ring_mask = grown & ~m & (sub_lab == 0)
+ bvals = sub_int[ring_mask]
+ bvals = bvals[np.isfinite(bvals)]
+ if bvals.size:
+ rows[i, COL["background"]] = bvals.mean()
+
+
+def _fill_intensity_fast(rows, lab, inten, tbl, keep, ring: int) -> bool:
+ """The vectorised half of :func:`_fill_intensity`. False if numba is missing.
+
+ Returns False WITHOUT writing anything when the ring kernel is unavailable,
+ so the caller can run the per-region loop instead — a half-filled row would
+ be worse than a slow one.
+ """
+ from spyde.particles.intensity import label_intensity_stats, ring_backgrounds
+
+ labels = np.asarray(tbl["label"], np.int64)
+ if labels.size == 0:
+ return True
+ bg = ring_backgrounds(lab, inten, labels, _table_bboxes(tbl), int(ring))
+ if bg is None:
+ return False
+
+ mean, mx, std = label_intensity_stats(lab, inten, labels)
+ keep = np.asarray(keep, bool)
+ rows[keep, COL["intensity_mean"]] = mean[keep]
+ rows[keep, COL["intensity_max"]] = mx[keep]
+ rows[keep, COL["intensity_std"]] = std[keep]
+ if int(ring) > 0:
+ rows[keep, COL["background"]] = bg[keep]
+ return True
+
+
+def _contours(lab: np.ndarray, tbl, *, fast: bool | None = None
+ ) -> list[np.ndarray]:
+ """One int16 outline per region, in ``tbl`` order.
+
+ Traced inside each region's padded bbox crop, not on the whole frame: a
+ frame-wide ``find_contours`` would return every region's outline in arbitrary
+ order with no label attached, and re-associating them is both slow and
+ ambiguous where particles touch.
+
+ :mod:`spyde.particles.contours` does exactly that, for every region in one
+ ``prange``, and produces the same FILLED polygon (the thing ``render_frame``
+ and ``mask_at`` consume) per region. The loop below is the definition it is
+ checked against, and the path that runs without numba.
+ """
+ if fast is None:
+ fast = _fast_default()
+ if fast:
+ from spyde.particles.contours import label_contours
+
+ got = label_contours(lab, np.asarray(tbl["label"], np.int64),
+ _table_bboxes(tbl),
+ np.asarray(tbl["area"], np.int64))
+ if got is not None:
+ return got
+
+ from skimage.measure import find_contours
+
+ h, w = lab.shape
+ out: list[np.ndarray] = []
+ for i in range(len(tbl["label"])):
+ lbl = int(tbl["label"][i])
+ y0, x0 = int(tbl["bbox-0"][i]), int(tbl["bbox-1"][i])
+ y1, x1 = int(tbl["bbox-2"][i]), int(tbl["bbox-3"][i])
+ py0, px0 = max(0, y0 - 1), max(0, x0 - 1)
+ py1, px1 = min(h, y1 + 1), min(w, x1 + 1)
+ sub = (lab[py0:py1, px0:px1] == lbl).astype(np.float32)
+ cs = find_contours(sub, 0.5)
+ if not cs:
+ # Degenerate (single pixel, or a region the tracer cannot close):
+ # fall back to the bbox corners so every row still has an outline and
+ # the 1:1 correspondence holds.
+ out.append(np.array(
+ [[y0, x0], [y0, x1 - 1], [y1 - 1, x1 - 1], [y1 - 1, x0]],
+ dtype=np.int16))
+ continue
+ c = max(cs, key=len) # outer boundary
+ c = np.rint(c).astype(np.int32)
+ c[:, 0] += py0
+ c[:, 1] += px0
+ np.clip(c[:, 0], 0, h - 1, out=c[:, 0])
+ np.clip(c[:, 1], 0, w - 1, out=c[:, 1])
+ out.append(c.astype(np.int16))
+ return out
+
+
+def particle_scores(rows: np.ndarray) -> np.ndarray:
+ """Per-particle confidence in [0, 1] — "is this a particle, or texture?"
+
+ Derived from columns already measured, which is the whole point: scoring
+ costs no extra pass over the frame, so the caret can filter on it without
+ re-segmenting and a slider drag is instant.
+
+ The statistic is CONTRAST-TO-NOISE against the particle's own dilated
+ background ring::
+
+ |intensity_mean - background| / spread
+
+ with *spread* the particle's own intensity spread. That is the quantity
+ that separates the two populations in the failure this exists for: a real
+ particle sits well away from its immediate surroundings, while a fragment
+ of textured support film is, by construction, the same brightness as the
+ texture around it however large or round it happens to be. Size, circularity
+ and solidity all fail here — over-split speckle is often small AND round.
+
+ Squashed to [0, 1] with ``cnr / (cnr + 1)`` so the caret's slider is a plain
+ 0-100% control with no dataset-dependent range to explain: 0.5 means "as far
+ from its background as its own noise", which is a weak particle, and real
+ ones land well above it.
+
+ NaN (no background ring measured — ``background_ring=0``, or a particle
+ whose ring fell entirely outside the frame) scores 1.0 rather than 0.0.
+ An unmeasurable particle must not be silently filtered away; the slider is
+ a "hide the marginal ones" control, and something with no evidence against
+ it is not marginal.
+ """
+ if rows.size == 0:
+ return np.zeros((0,), np.float32)
+ mean = rows[:, COL["intensity_mean"]].astype(np.float64)
+ bg = rows[:, COL["background"]].astype(np.float64)
+ std = rows[:, COL["intensity_std"]].astype(np.float64)
+
+ contrast = np.abs(mean - bg)
+ # `intensity_std` is already normalised by the particle's max (see
+ # _fill_intensity), so put it back on the intensity scale before using it
+ # as a noise estimate. Guard the degenerate flat particle.
+ noise = np.where(np.isfinite(std) & (std > 0), std * np.abs(mean), np.nan)
+ floor = np.nanmedian(noise) if np.isfinite(noise).any() else 1.0
+ if not np.isfinite(floor) or floor <= 0:
+ floor = 1.0
+ noise = np.where(np.isfinite(noise) & (noise > 0), noise, floor)
+
+ with np.errstate(invalid="ignore", divide="ignore"):
+ cnr = contrast / noise
+ score = cnr / (cnr + 1.0)
+ # Unmeasurable => not marginal => keep. See the docstring.
+ score = np.where(np.isfinite(score), score, 1.0)
+ return np.clip(score, 0.0, 1.0).astype(np.float32)
diff --git a/spyde/particles/props.py b/spyde/particles/props.py
new file mode 100644
index 00000000..723d5d8b
--- /dev/null
+++ b/spyde/particles/props.py
@@ -0,0 +1,320 @@
+"""
+props.py — ``regionprops_table``'s columns as label-wise reductions over the raster.
+
+``skimage.measure.regionprops_table`` walks a Python object per region. That is
+fine at the scale its docs assume and it is the run at ours: a real 4096² in-situ
+growth frame has **26 566** particles, and measuring them costs **53.5 s** against
+3.0 s to segment them (``benchmarks.md`` § "The batch run at real scale"). The
+cost is per REGION, so it grows with what the microscope actually produced.
+
+Almost every property in that table is a **label-wise reduction over the pixel
+raster** — O(pixels), no Python loop over regions, exactly the shape
+``np.bincount`` exists for. This module computes those, and only those, in one
+vectorised pass:
+
+* ``area``, ``centroid``, ``bbox``, ``equivalent_diameter_area`` — counts, sums
+ and per-label extents.
+* ``major/minor_axis_length``, ``eccentricity`` — the eigenvalues of the inertia
+ tensor, which is built from the second central moments ``mu20/mu02/mu11``, each
+ of which is a ``bincount`` of ``dr*dr`` / ``dc*dc`` / ``dr*dc``.
+* ``perimeter`` — skimage's border-crossing weighting, done on the whole frame at
+ once (see :func:`label_perimeter`).
+
+**Value parity is the whole point, and it is not approximate.** Every formula
+here is skimage's own, reproduced from its source rather than from its docs, and
+in skimage's own coordinate frame:
+
+* ``centroid`` is the mean of GLOBAL integer coordinates. Both paths sum exact
+ integers in float64 (the largest possible sum here is ~1e8, far inside
+ float64's exact-integer range), so the two agree **bit for bit**.
+* ``area``/``bbox`` are integers, so they agree **exactly**.
+* The central moments are taken about ``centroid_local`` in the region's own
+ bbox frame, the two-pass way ``_moments.moments_central`` does it when handed
+ an explicit centre — NOT the raw-to-central expansion, which cancels. skimage
+ reduces along one axis then the other (``einsum``) and this reduces in raster
+ order, so the two differ only in summation order: ~1e-15 relative, and the
+ eigen-decomposition that follows is the same LAPACK call on the same 2×2.
+* ``perimeter`` sums the same per-region 50-bin histogram against the same
+ weight vector, so it also matches to summation order.
+
+``solidity`` is the exception and gets its own module: it is a convex hull per
+region, not a raster reduction, and it is 30.5 s of that 53.5 s on its own.
+:mod:`spyde.particles.hull` reproduces skimage's hull EXACTLY in integer
+arithmetic (verified bit-for-bit on all 26 566 regions of the real frame) and
+:func:`solidity_table` remains as the fallback when numba is unavailable.
+
+What this does NOT do
+---------------------
+No GPU. The reductions land at ~1.0 s for a 4096² frame with 26 566 regions,
+which is already below the per-frame budget, and a torch path would add a device
+lock, a transfer and a fallback to a stage that is no longer the bottleneck.
+
+Nothing here materialises more than the frame it was handed (CLAUDE.md § Memory
+Safety): the working set is the label raster, one padded copy of it, and a handful
+of arrays sized by the FOREGROUND pixel count.
+"""
+from __future__ import annotations
+
+from math import sqrt
+
+import numpy as np
+
+#: skimage's border-crossing weights (``_regionprops.perimeter``). Index is the
+#: convolution code ``1 + 2*(orthogonal border neighbours) + 10*(diagonal border
+#: neighbours)``; only odd codes (i.e. those with a border pixel at the centre)
+#: carry weight, which is why this module can evaluate the code at border pixels
+#: only and still match a full-frame convolution.
+_PERIM_WEIGHTS = np.zeros(50, dtype=np.float64)
+_PERIM_WEIGHTS[[5, 7, 15, 17, 25, 27]] = 1.0
+_PERIM_WEIGHTS[[21, 33]] = sqrt(2)
+_PERIM_WEIGHTS[[13, 23]] = (1.0 + sqrt(2)) / 2.0
+
+#: The keys :func:`label_props` produces, matching ``regionprops_table``'s names.
+PROP_KEYS: tuple[str, ...] = (
+ "label",
+ "centroid-0", "centroid-1",
+ "area",
+ "equivalent_diameter_area",
+ "major_axis_length",
+ "minor_axis_length",
+ "perimeter",
+ "eccentricity",
+ "bbox-0", "bbox-1", "bbox-2", "bbox-3",
+)
+
+
+def _empty_table(with_perimeter: bool = True) -> dict[str, np.ndarray]:
+ out = {k: np.zeros((0,), np.float64) for k in PROP_KEYS}
+ out["label"] = np.zeros((0,), np.int64)
+ for k in ("bbox-0", "bbox-1", "bbox-2", "bbox-3"):
+ out[k] = np.zeros((0,), np.int64)
+ out["area"] = np.zeros((0,), np.float64)
+ if not with_perimeter:
+ out.pop("perimeter")
+ return out
+
+
+def label_bboxes(lab: np.ndarray, n_max: int) -> np.ndarray:
+ """``(n_max, 4)`` int64 ``(y0, x0, y1, x1)``, row *i* for label ``i+1``.
+
+ ``scipy.ndimage.find_objects`` is the same C pass ``regionprops`` itself uses
+ to decide which labels exist, and it is 0.065 s at 4096² against 1.1 s for the
+ ``bbox`` column of ``regionprops_table``. Absent labels get an all-zero row and
+ are filtered out by the caller's ``counts > 0`` mask, which selects exactly the
+ labels for which ``find_objects`` returned a slice.
+ """
+ from scipy import ndimage as ndi
+
+ objs = ndi.find_objects(lab, max_label=int(n_max))
+ bb = np.zeros((int(n_max), 4), np.int64)
+ for i, sl in enumerate(objs):
+ if sl is None:
+ continue
+ ys, xs = sl
+ bb[i, 0] = ys.start
+ bb[i, 1] = xs.start
+ bb[i, 2] = ys.stop
+ bb[i, 3] = xs.stop
+ return bb
+
+
+def label_perimeter(lab: np.ndarray, dense: np.ndarray, n_out: int) -> np.ndarray:
+ """Per-label ``skimage.measure.perimeter(region_mask, 4)``, whole frame at once.
+
+ *dense* maps a label value to its row in the output (``-1`` for absent), and
+ *n_out* is how many rows that is. The indirection is not cosmetic: the
+ per-label histogram has **50 bins per label**, so keying it by the raw label
+ value would allocate 50x the label range. A label image whose values are
+ sparse (anything filtered or re-tagged upstream) would then ask for gigabytes
+ to describe a few thousand regions.
+
+ skimage runs, for each region, a 4-connected erosion of that region's ISOLATED
+ mask inside its bbox, subtracts it to get the border ring, convolves the ring
+ with ``[[10,2,10],[2,1,2],[10,2,10]]`` and sums a weight per code. Two
+ observations turn that into one pass over the frame:
+
+ * A pixel survives the per-region erosion iff it and all four of its
+ orthogonal neighbours carry the SAME label. The bbox is tight, so any
+ neighbour outside the crop is outside the region as well, which is exactly
+ what ``border_value=0`` gives — so testing "same label" on the padded FULL
+ frame reproduces the per-region erosion pixel for pixel.
+ * Only odd codes carry a nonzero weight, and a code is odd only when the
+ centre pixel is itself a border pixel. Non-border pixels therefore
+ contribute exactly 0 and need not be evaluated.
+
+ The per-label sum is taken as a ``(n_max+1, 50)`` histogram dotted with the
+ weight vector — the same reduction skimage performs, so the result agrees to
+ summation order rather than to an approximation.
+ """
+ h, w = lab.shape
+ pad = np.zeros((h + 2, w + 2), dtype=lab.dtype)
+ pad[1:-1, 1:-1] = lab
+ ctr = pad[1:-1, 1:-1]
+
+ def shifted(dy: int, dx: int) -> np.ndarray:
+ return pad[1 + dy:1 + dy + h, 1 + dx:1 + dx + w]
+
+ fg = ctr != 0
+ interior = fg.copy()
+ for dy, dx in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+ interior &= shifted(dy, dx) == ctr
+ border = fg
+ border &= ~interior # in place; `fg` is a fresh array
+ del interior
+
+ bpad = np.zeros((h + 2, w + 2), dtype=bool)
+ bpad[1:-1, 1:-1] = border
+
+ def bshift(dy: int, dx: int) -> np.ndarray:
+ return bpad[1 + dy:1 + dy + h, 1 + dx:1 + dx + w]
+
+ code = border.astype(np.uint8) # the centre pixel's own +1
+ same_border = np.empty((h, w), dtype=bool)
+ for dy, dx in ((-1, 0), (1, 0), (0, -1), (0, 1)):
+ np.equal(shifted(dy, dx), ctr, out=same_border)
+ same_border &= bshift(dy, dx)
+ code += same_border.view(np.uint8) * np.uint8(2)
+ for dy, dx in ((-1, -1), (-1, 1), (1, -1), (1, 1)):
+ np.equal(shifted(dy, dx), ctr, out=same_border)
+ same_border &= bshift(dy, dx)
+ code += same_border.view(np.uint8) * np.uint8(10)
+
+ idx = np.flatnonzero(border.reshape(-1))
+ if idx.size == 0:
+ return np.zeros((int(n_out),), np.float64)
+ keys = dense[lab.reshape(-1)[idx]] * 50
+ keys += code.reshape(-1)[idx]
+ hist = np.bincount(keys, minlength=int(n_out) * 50)
+ hist = hist[:int(n_out) * 50].reshape(int(n_out), 50)
+ return hist.astype(np.float64) @ _PERIM_WEIGHTS
+
+
+def label_props(lab: np.ndarray, *, with_perimeter: bool = True,
+ with_solidity: bool = True) -> dict[str, np.ndarray]:
+ """``regionprops_table``'s columns, vectorised, in the same order and dtypes.
+
+ Drops into ``measure_frame`` where the table did. ``solidity`` comes from
+ :mod:`spyde.particles.hull` when numba is available and from
+ :func:`solidity_table` (skimage) when it is not — the two agree exactly.
+ """
+ lab = np.asarray(lab)
+ if lab.ndim != 2:
+ raise ValueError(f"labels must be 2-D; got shape {lab.shape}")
+ h, w = lab.shape
+ flat = lab.reshape(-1)
+ if flat.size == 0 or int(lab.max()) <= 0:
+ out = _empty_table(with_perimeter)
+ if with_solidity:
+ out["solidity"] = np.zeros((0,), np.float64)
+ return out
+
+ counts = np.bincount(flat)
+ n_max = counts.size - 1
+ keep = counts[1:] > 0
+ labels = (np.flatnonzero(keep) + 1).astype(np.int64)
+
+ bb = label_bboxes(lab, n_max)
+
+ # Foreground pixels only. The reductions below are O(foreground), not
+ # O(frame), and the raster is scanned once to find them.
+ nz = np.flatnonzero(flat)
+ lab_of = flat[nz].astype(np.intp)
+ row = nz // w
+ col = nz - row * w
+
+ cnt_f = counts.astype(np.float64)
+
+ # LOCAL (bbox-relative) coordinates, which is the frame skimage takes its
+ # central moments in. Integers, so these sums are exact.
+ y0_by = np.zeros(n_max + 1, np.int64)
+ x0_by = np.zeros(n_max + 1, np.int64)
+ y0_by[1:] = bb[:, 0]
+ x0_by[1:] = bb[:, 1]
+ dr = (row - y0_by[lab_of]).astype(np.float64)
+ dc = (col - x0_by[lab_of]).astype(np.float64)
+ del row, col, nz
+
+ sr = np.bincount(lab_of, weights=dr, minlength=n_max + 1)
+ sc = np.bincount(lab_of, weights=dc, minlength=n_max + 1)
+
+ with np.errstate(divide="ignore", invalid="ignore"):
+ cr_local = sr / cnt_f
+ cc_local = sc / cnt_f
+ np.nan_to_num(cr_local, copy=False)
+ np.nan_to_num(cc_local, copy=False)
+
+ # Two-pass central moments about `centroid_local`, matching
+ # `_moments.moments_central(image, centroid_local, ...)`.
+ dr -= cr_local[lab_of]
+ dc -= cc_local[lab_of]
+ mu20 = np.bincount(lab_of, weights=dr * dr, minlength=n_max + 1)
+ mu02 = np.bincount(lab_of, weights=dc * dc, minlength=n_max + 1)
+ dr *= dc
+ mu11 = np.bincount(lab_of, weights=dr, minlength=n_max + 1)
+ del dr, dc, lab_of
+
+ sel = labels # index into the by-label arrays
+ area = cnt_f[sel]
+ # GLOBAL centroid = mean of global integer coordinates, exactly as
+ # `RegionProperties.centroid` computes it (`coords_scaled.mean(axis=0)`).
+ cy = (sr[sel] + area * bb[sel - 1, 0]) / area
+ cx = (sc[sel] + area * bb[sel - 1, 1]) / area
+
+ # inertia_tensor: [[mu02, -mu11], [-mu11, mu20]] / mu00 (skimage's
+ # convention — I_ii is the second moment of every axis EXCEPT i).
+ n = sel.size
+ tensor = np.empty((n, 2, 2), np.float64)
+ tensor[:, 0, 0] = mu02[sel] / area
+ tensor[:, 1, 1] = mu20[sel] / area
+ off = -mu11[sel] / area
+ tensor[:, 0, 1] = off
+ tensor[:, 1, 0] = off
+ ev = np.linalg.eigvalsh(tensor) # ascending
+ np.clip(ev, 0, None, out=ev)
+ l1 = ev[:, 1] # descending order's first
+ l2 = ev[:, 0]
+ major = 4.0 * np.sqrt(l1)
+ minor = 4.0 * np.sqrt(l2)
+ with np.errstate(divide="ignore", invalid="ignore"):
+ ecc = np.where(l1 == 0, 0.0, np.sqrt(1.0 - l2 / l1))
+
+ out: dict[str, np.ndarray] = {
+ "label": labels,
+ "centroid-0": cy,
+ "centroid-1": cx,
+ "area": area,
+ "equivalent_diameter_area": (4.0 * area / np.pi) ** (1 / 2),
+ "major_axis_length": major,
+ "minor_axis_length": minor,
+ "eccentricity": ecc,
+ "bbox-0": bb[sel - 1, 0],
+ "bbox-1": bb[sel - 1, 1],
+ "bbox-2": bb[sel - 1, 2],
+ "bbox-3": bb[sel - 1, 3],
+ }
+ if with_perimeter:
+ dense = np.zeros(n_max + 1, np.int64)
+ dense[labels] = np.arange(labels.size, dtype=np.int64)
+ out["perimeter"] = label_perimeter(lab, dense, labels.size)
+ if with_solidity:
+ from spyde.particles.hull import convex_areas
+
+ conv = convex_areas(lab, labels, counts)
+ out["solidity"] = (area / conv.astype(np.float64)) if conv is not None \
+ else solidity_table(lab)
+ return out
+
+
+def solidity_table(lab: np.ndarray) -> np.ndarray:
+ """``solidity`` for every region, on skimage's per-region convex hull.
+
+ The fallback for :func:`spyde.particles.hull.convex_areas` — used when numba
+ is missing or its kernel will not compile. Kept because it is the definition
+ everything else is checked against, and because it must still be possible to
+ measure a frame on a machine with no numba at all. It is ~170x slower (30.5 s
+ against 0.18 s at 26 566 regions), so this is a correctness floor, not a
+ performance one.
+ """
+ from skimage.measure import regionprops_table
+
+ return regionprops_table(lab, properties=("solidity",))["solidity"]
diff --git a/spyde/particles/scribble.py b/spyde/particles/scribble.py
new file mode 100644
index 00000000..54245b9c
--- /dev/null
+++ b/spyde/particles/scribble.py
@@ -0,0 +1,1179 @@
+"""
+scribble.py — the scribble-trained pixel classifier. Plan step B3, the workhorse.
+
+The user paints a few strokes on a frame; this learns *their* data and produces a
+per-frame foreground probability map that :func:`spyde.particles.classical.
+split_instances` turns into instances. Same shared downstream stage as the
+classical and prompt engines — the only thing that changes is how the first box
+in the plan's §0.2 pipeline gets filled.
+
+Why this and not a threshold
+----------------------------
+Plan §0.9 makes detection **sensitivity** the priority: missing a particle's first
+appearance destroys the nucleation event, which is the most interesting thing in
+the movie. No single global threshold spans a nucleation sequence — one that
+catches a 3σ particle at t=0 is not the one that works at t=end. Measured on the
+``particle_movie()`` fixture, whose two ``p_faint`` probes are exactly this case:
+the classical engine at its default sensitivity finds **0 of 2**
+(``test_particle_movie_fixture.py::test_default_sensitivity_misses_the_faint_probes``
+pins that), and this classifier trained on eleven scribbles finds **2 of 2** while
+keeping all seven bright ones.
+
+Design decisions that are not obvious
+-------------------------------------
+**Multi-class, and class 0 is not special.** Classes are user-defined
+``(id, name, colour)`` triples and a softmax head costs nothing over a sigmoid. In
+EM "background" is genuinely two or three different things — carbon film, vacuum,
+beam-stop — and forcing them into one class makes the head spend its capacity
+proving they are the same rather than separating either from a particle. Which
+classes count as foreground is a per-class flag (:attr:`ScribbleClass.particle`),
+so several particle *phases* can be labelled separately and still sum into one
+probability map.
+
+**A BOUNDARY class, and it is a performance feature.** The ilastik convention is
+particle / background / **boundary**, and the third one is not a refinement — it
+is the only way out of the cost that dominates a large frame. Every engine ends
+at :func:`spyde.particles.classical.split_instances`, whose watershed route needs
+a global distance transform, a marker/elevation upsample and a flood — together
+1.62 s of a 1.78 s split at 4096². All of that exists to *guess* where two
+touching particles should be cut. A head that has
+been shown a few strokes along the joins does not have to guess: it returns them
+already separated, so the split degenerates to one ``ndi.label`` and both the
+distance transform and the watershed are skipped. Measured on 4096²: **1.78 s →
+0.33 s for the split**, and on that field it also found the exactly-correct 162
+bodies where the watershed found 173.
+
+:attr:`ScribbleClass.boundary` marks such a class, only this engine can produce
+one, and :meth:`ScribbleClassifier.segment` passes it down automatically —
+falling back to the watershed when nothing was painted, so not using the class
+costs nothing.
+
+**What "boundary" means is the whole art of it, and getting it wrong is silent.**
+It is the seam BETWEEN two bodies, never the outline of one. A head taught
+outlines learns "shrink everything": on the fixture's merge frame that MERGED the
+touching pair and lost 40% of the median area, while still reporting a trained
+boundary class. A boundary trained on a handful of pixels is no better — 30 px of
+seam took the fast route and returned 81 bodies where the watershed found 162.
+The per-class pixel counts in the caret are what surface this, which is why
+:meth:`LabelStore.counts` lists classes with no pixels at all.
+
+**Labels accumulate across frames, sparsely.** :class:`LabelStore` is keyed by
+frame index, so painting on frame 0 and again on frame 400 trains one model from
+both. It stores flat pixel indices, not label images: a scribble is a few thousand
+pixels, and a dense ``int16`` map is 32 MB *per labelled frame* at 4096² — 320 MB
+for ten labelled frames, to hold ~50 000 useful values.
+
+**A torch MLP, with sklearn's RandomForest kept as the parity reference.** One
+hidden layer, class-balanced cross-entropy. The forest is what ParticleSpy and
+ilastik use and it is a *better* classifier out of the box on this kind of
+tabular problem — but it cannot live on the GPU next to the feature stack, and
+the interaction budget is train+apply under 1 s on the displayed frame. So the
+MLP is the shipped head and agreement with the forest on identical labels and
+identical features is the acceptance gate
+(``test_particles_scribble.py::TestRandomForestParity``).
+
+**The NaN border is forced to zero probability.** Plan trap 2: a drift-corrected
+frame has a NaN-padded border, and segmentation that ignores it invents a large
+"particle" along the edge which then nucleates a spurious track.
+:func:`spyde.particles.features.prepare_frame` reports validity and every
+prediction here is masked by it.
+"""
+from __future__ import annotations
+
+import json
+import logging
+import math
+import time
+from dataclasses import dataclass, field
+from typing import Any, Callable, Sequence
+
+import numpy as np
+
+from spyde.device_lock import accelerator_lock
+from spyde.particles.features import (
+ FeatureSpec,
+ import_torch,
+ map_feature_bands,
+ prepare_frame,
+ sample_features,
+ select_device,
+)
+
+log = logging.getLogger(__name__)
+
+FORMAT_VERSION = 1
+
+#: Sentinel for "no class" in a label map or a predicted label image. -1 rather
+#: than 0 because class id 0 is a perfectly ordinary user class here.
+UNLABELLED = -1
+
+#: Default classes for a fresh session: a particle class, two backgrounds and a
+#: boundary. Two backgrounds and not one because that is what EM actually looks
+#: like, and because the second one is free — see the module docstring. Colours
+#: are the renderer's accent family so the floating brush strip (plan B0) needs
+#: no palette of its own.
+#:
+#: ``(id, name, colour, particle, boundary)``. The **boundary** class is the
+#: ilastik convention and it is here for speed as much as for quality: painting
+#: the joins between touching particles lets
+#: :func:`~spyde.particles.classical.split_instances` take its connected-
+#: components route and skip the distance transform and watershed entirely —
+#: 1.78 s down to 0.33 s at 4096². It is not a particle class — a boundary
+#: pixel is the seam, not the body — so it does not enter the foreground sum.
+DEFAULT_CLASSES: tuple[tuple[int, str, str, bool, bool], ...] = (
+ (0, "particle", "#f9a03f", True, False),
+ (1, "support film", "#89b4fa", False, False),
+ (2, "vacuum", "#585b70", False, False),
+ (3, "boundary", "#f38ba8", False, True),
+)
+
+
+# ── classes and the label store ──────────────────────────────────────────────
+
+@dataclass(frozen=True)
+class ScribbleClass:
+ """One user-defined class.
+
+ Parameters
+ ----------
+ id
+ Stable integer key. Referenced by :class:`LabelStore` and by the trained
+ head's output column order, so it must not be reused after a delete.
+ name, colour
+ For the caret's class list and the in-canvas brush strip. ``colour`` is a
+ CSS hex string — the renderer's units, not the backend's.
+ particle
+ Whether this class counts toward the foreground probability map. Several
+ classes may set it (two particle phases, say) and their probabilities sum.
+ boundary
+ Whether this class marks the **seam between touching particles**. Summed
+ the same way *particle* is, into a separate map that
+ :func:`~spyde.particles.classical.split_instances` uses to skip the
+ watershed. A class is one or the other, never both: a boundary pixel is
+ not part of any body, and counting it as foreground would glue the two
+ bodies it separates back together — which is the exact failure the class
+ exists to prevent.
+ """
+
+ id: int
+ name: str
+ colour: str = "#ffffff"
+ particle: bool = False
+ boundary: bool = False
+
+ def __post_init__(self) -> None:
+ if self.particle and self.boundary:
+ raise ValueError(
+ f"class {self.id} ({self.name!r}) is marked both particle and "
+ "boundary — a seam is not part of a body, so counting it as "
+ "both would merge every pair of touching particles it separates")
+
+ def to_dict(self) -> dict[str, Any]:
+ return {"id": int(self.id), "name": self.name, "colour": self.colour,
+ "particle": bool(self.particle),
+ "boundary": bool(self.boundary)}
+
+ @classmethod
+ def from_dict(cls, d: dict[str, Any]) -> "ScribbleClass":
+ # `boundary` defaults False, so a session or a model saved before the
+ # class existed loads as the particle/background-only setup it was.
+ return cls(int(d["id"]), str(d["name"]), str(d.get("colour", "#ffffff")),
+ bool(d.get("particle", False)), bool(d.get("boundary", False)))
+
+
+def default_classes() -> list[ScribbleClass]:
+ """A fresh copy of :data:`DEFAULT_CLASSES`."""
+ return [ScribbleClass(i, n, c, p, b) for i, n, c, p, b in DEFAULT_CLASSES]
+
+
+@dataclass(eq=False)
+class LabelStore:
+ """Scribble labels, keyed by frame index and accumulating across frames.
+
+ Painted pixels are held as flat indices into a ``(h, w)`` frame plus a
+ parallel class-id array — see the module docstring for why this is not a dense
+ label image. Repainting a pixel *replaces* its class (last write wins), which
+ is what a user expects from a brush, and erasing removes it entirely rather
+ than assigning it to a background class.
+
+ Parameters
+ ----------
+ frame_shape
+ ``(h, w)``. Fixed for the life of the store: a flat index means nothing
+ without it, and silently accepting a differently-shaped frame would
+ scatter the labels across the image.
+ classes
+ The class list. Mutate through :meth:`add_class` / :meth:`remove_class`
+ so ids stay unique and a removed class takes its pixels with it.
+
+ Notes
+ -----
+ ``eq=False``: the generated ``__eq__`` would compare a dict of numpy arrays
+ and raise on the ambiguous truth value. Compare :meth:`to_dict` instead, which
+ is what the round-trip test does.
+ """
+
+ frame_shape: tuple[int, int]
+ classes: list[ScribbleClass] = field(default_factory=default_classes)
+ #: frame index -> (flat indices int64, class ids int16)
+ _frames: dict[int, tuple[np.ndarray, np.ndarray]] = field(
+ default_factory=dict, repr=False)
+
+ def __post_init__(self) -> None:
+ self.frame_shape = (int(self.frame_shape[0]), int(self.frame_shape[1]))
+ ids = [c.id for c in self.classes]
+ if len(set(ids)) != len(ids):
+ raise ValueError(f"duplicate class ids: {ids}")
+
+ # -- classes -------------------------------------------------------------
+
+ @property
+ def class_ids(self) -> list[int]:
+ return [c.id for c in self.classes]
+
+ def class_by_id(self, cid: int) -> ScribbleClass:
+ for c in self.classes:
+ if c.id == int(cid):
+ return c
+ raise KeyError(
+ f"no class with id {cid}; have {self.class_ids}")
+
+ def add_class(self, name: str, colour: str = "#ffffff", *,
+ particle: bool = False, boundary: bool = False,
+ id: int | None = None) -> ScribbleClass:
+ """Append a class. Its id is one past the current maximum unless given."""
+ cid = int(id) if id is not None else (
+ max(self.class_ids) + 1 if self.classes else 0)
+ if cid in self.class_ids:
+ raise ValueError(f"class id {cid} already exists")
+ c = ScribbleClass(cid, str(name), str(colour), bool(particle),
+ bool(boundary))
+ self.classes.append(c)
+ return c
+
+ def remove_class(self, cid: int) -> None:
+ """Drop a class **and every pixel labelled with it**.
+
+ Leaving orphaned pixels behind would train a head with a column for a
+ class the user has deleted, and the pixel counts in the caret would stop
+ adding up to the labelled total — which is the one number that tells the
+ user a class is under-trained.
+ """
+ cid = int(cid)
+ self.class_by_id(cid) # raises if unknown
+ self.classes = [c for c in self.classes if c.id != cid]
+ for t in list(self._frames):
+ idx, cls = self._frames[t]
+ keep = cls != cid
+ if keep.all():
+ continue
+ if keep.any():
+ self._frames[t] = (idx[keep], cls[keep])
+ else:
+ del self._frames[t]
+
+ # -- painting ------------------------------------------------------------
+
+ def _flatten(self, where) -> np.ndarray:
+ """Coerce a mask / ``(k, 2)`` yx / flat-index argument to flat indices."""
+ h, w = self.frame_shape
+ a = np.asarray(where)
+ if a.dtype == bool:
+ if a.shape != (h, w):
+ raise ValueError(
+ f"mask shape {a.shape} != frame_shape {self.frame_shape}")
+ return np.flatnonzero(a.reshape(-1)).astype(np.int64)
+ if a.ndim == 2 and a.shape[1] == 2:
+ ys = np.rint(a[:, 0]).astype(np.int64)
+ xs = np.rint(a[:, 1]).astype(np.int64)
+ inside = (ys >= 0) & (ys < h) & (xs >= 0) & (xs < w)
+ return (ys[inside] * w + xs[inside])
+ flat = a.reshape(-1).astype(np.int64)
+ return flat[(flat >= 0) & (flat < h * w)]
+
+ def paint(self, t: int, where, class_id: int) -> int:
+ """Label pixels on frame *t*. Returns how many pixels the store now holds
+ for that frame.
+
+ *where* may be a boolean mask the shape of the frame, a ``(k, 2)``
+ ``(y, x)`` array, or flat indices. Out-of-frame coordinates are dropped
+ rather than raising — a brush stroke that runs off the edge is normal.
+ """
+ cid = int(class_id)
+ self.class_by_id(cid)
+ t = int(t)
+ add = self._flatten(where)
+ if not add.size:
+ return len(self._frames.get(t, (np.empty(0),))[0])
+
+ cur_idx, cur_cls = self._frames.get(
+ t, (np.zeros(0, np.int64), np.zeros(0, np.int16)))
+ idx = np.concatenate([cur_idx, add])
+ cls = np.concatenate([cur_cls, np.full(add.size, cid, np.int16)])
+ self._frames[t] = _dedup_last_wins(idx, cls)
+ return int(self._frames[t][0].size)
+
+ def paint_disc(self, t: int, y: float, x: float, radius: float,
+ class_id: int) -> int:
+ """Label a filled disc — one brush dab, and what a click paints."""
+ return self.paint(t, _disc_indices(self.frame_shape, y, x, radius),
+ class_id)
+
+ def paint_stroke(self, t: int, points: Sequence[Sequence[float]],
+ class_id: int, *, brush: float = 3.0) -> int:
+ """Label a brush stroke: a polyline of ``(y, x)`` points, *brush* px wide.
+
+ This is the door the anyplotlib brush widget (plan B0) comes through, and
+ the coordinates arrive in **image pixels** with no scale or offset applied
+ — plan trap 6. Do not multiply by the axis scale on the way in.
+
+ The polyline is densified to half-pixel steps before dabbing, because the
+ widget emits a pointer sample per frame and a fast stroke can jump 20 px
+ between them; dabbing only at the samples leaves a dotted line.
+ """
+ pts = np.asarray(points, dtype=np.float64).reshape(-1, 2)
+ if not len(pts):
+ return int(self._frames.get(t, (np.empty(0),))[0].size)
+ r = max(0.5, float(brush) / 2.0)
+ dense = [pts[0]]
+ for a, b in zip(pts, pts[1:]):
+ n = int(np.ceil(np.hypot(*(b - a)) * 2.0))
+ if n > 1:
+ dense.extend(a + (b - a) * (np.arange(1, n + 1) / n)[:, None])
+ else:
+ dense.append(b)
+ idx = np.unique(np.concatenate(
+ [_disc_indices(self.frame_shape, py, px, r) for py, px in dense]))
+ return self.paint(t, idx, class_id)
+
+ def erase(self, t: int, where) -> int:
+ """Unlabel pixels on frame *t* (the brush widget's eraser)."""
+ t = int(t)
+ got = self._frames.get(t)
+ if got is None:
+ return 0
+ idx, cls = got
+ drop = np.isin(idx, self._flatten(where))
+ if drop.all():
+ del self._frames[t]
+ return 0
+ self._frames[t] = (idx[~drop], cls[~drop])
+ return int(self._frames[t][0].size)
+
+ def clear_frame(self, t: int) -> None:
+ self._frames.pop(int(t), None)
+
+ def clear(self) -> None:
+ self._frames.clear()
+
+ # -- inspection ----------------------------------------------------------
+
+ def labelled_frames(self) -> list[int]:
+ """Frame indices carrying labels, ascending — the caret's revisit list."""
+ return sorted(self._frames)
+
+ def at(self, t: int) -> tuple[np.ndarray, np.ndarray]:
+ """``(flat_indices, class_ids)`` for frame *t*; empty arrays if none."""
+ return self._frames.get(
+ int(t), (np.zeros(0, np.int64), np.zeros(0, np.int16)))
+
+ def label_map(self, t: int) -> np.ndarray:
+ """Frame *t*'s labels as a dense ``(h, w)`` int16 map, :data:`UNLABELLED`
+ where unpainted. Built on demand for display; never stored (see the
+ module docstring)."""
+ idx, cls = self.at(t)
+ out = np.full(self.frame_shape, UNLABELLED, dtype=np.int16)
+ if idx.size:
+ out.reshape(-1)[idx] = cls
+ return out
+
+ def counts(self) -> dict[int, int]:
+ """Labelled pixels per class id, across every frame.
+
+ This is what the caret's class list shows, and it is how a user notices a
+ class is under-trained — so classes with zero pixels are present in the
+ result rather than absent from it.
+ """
+ out = {c.id: 0 for c in self.classes}
+ for idx, cls in self._frames.values():
+ if not cls.size:
+ continue
+ vals, n = np.unique(cls, return_counts=True)
+ for v, k in zip(vals.tolist(), n.tolist()):
+ out[int(v)] = out.get(int(v), 0) + int(k)
+ return out
+
+ def __len__(self) -> int:
+ return int(sum(idx.size for idx, _ in self._frames.values()))
+
+ @property
+ def n_classes_used(self) -> int:
+ return int(sum(1 for v in self.counts().values() if v > 0))
+
+ # -- serialisation -------------------------------------------------------
+
+ def to_dict(self) -> dict[str, Any]:
+ """JSON-safe, so a session's scribbles survive a save/reload.
+
+ Indices go out as lists. That is fine at scribble scale (a few thousand
+ per frame) and is what keeps the format inspectable; a binary side-car
+ would only pay off at a density no human paints.
+ """
+ return {
+ "frame_shape": list(self.frame_shape),
+ "classes": [c.to_dict() for c in self.classes],
+ "frames": {str(t): {"index": idx.tolist(), "class": cls.tolist()}
+ for t, (idx, cls) in sorted(self._frames.items())},
+ }
+
+ @classmethod
+ def from_dict(cls, d: dict[str, Any]) -> "LabelStore":
+ store = cls(
+ frame_shape=tuple(d["frame_shape"]),
+ classes=[ScribbleClass.from_dict(c) for c in d["classes"]],
+ )
+ for key, blk in (d.get("frames") or {}).items():
+ store._frames[int(key)] = (
+ np.asarray(blk["index"], dtype=np.int64),
+ np.asarray(blk["class"], dtype=np.int16),
+ )
+ return store
+
+
+def _dedup_last_wins(idx: np.ndarray, cls: np.ndarray) -> tuple[np.ndarray, np.ndarray]:
+ """Collapse repeated indices, keeping the LAST class written to each.
+
+ ``np.unique`` keeps the *first* occurrence of each value, so the arrays are
+ reversed going in and the result flipped back — a repaint over an existing
+ stroke has to change its class, not be ignored.
+ """
+ if idx.size < 2:
+ return idx.astype(np.int64), cls.astype(np.int16)
+ ridx, rcls = idx[::-1], cls[::-1]
+ uniq, first = np.unique(ridx, return_index=True)
+ return uniq.astype(np.int64), rcls[first].astype(np.int16)
+
+
+def _disc_indices(shape: tuple[int, int], y: float, x: float,
+ radius: float) -> np.ndarray:
+ """Flat indices of a filled disc, clipped to the frame.
+
+ Built from the bounding box rather than a full-frame distance map: a brush dab
+ is a handful of pixels and a 4096² ``mgrid`` per dab would make painting
+ unusable.
+ """
+ h, w = int(shape[0]), int(shape[1])
+ r = max(0.5, float(radius))
+ y0, y1 = max(0, int(math.floor(y - r))), min(h - 1, int(math.ceil(y + r)))
+ x0, x1 = max(0, int(math.floor(x - r))), min(w - 1, int(math.ceil(x + r)))
+ if y1 < y0 or x1 < x0:
+ return np.zeros(0, np.int64)
+ ys = np.arange(y0, y1 + 1)[:, None]
+ xs = np.arange(x0, x1 + 1)[None, :]
+ inside = (ys - float(y)) ** 2 + (xs - float(x)) ** 2 <= r * r
+ yy, xx = np.nonzero(inside)
+ return ((yy + y0).astype(np.int64) * w + (xx + x0).astype(np.int64))
+
+
+# ── the SAM/prompt bootstrap (plan §0.4) ─────────────────────────────────────
+
+def masks_to_labels(
+ masks,
+ *,
+ t: int = 0,
+ frame_shape: tuple[int, int] | None = None,
+ store: LabelStore | None = None,
+ particle_class: int = 0,
+ background_class: int = 1,
+ gap: int = 2,
+ background_dilation: int = 8,
+ erode: int = 1,
+) -> LabelStore:
+ """Turn promptable-segmentation masks into scribble labels. Plan §0.4.
+
+ "Click four particles with the prompt model. Those masks — plus their dilated
+ surroundings as background — become the scribble classifier's training labels.
+ Train, apply to all N frames. No painting at all." This is that handoff, and
+ it is a first-class feature rather than a convenience: it is what makes the
+ dense result adapted to the data instead of to COCO.
+
+ Parameters
+ ----------
+ masks
+ A boolean ``(h, w)`` mask, a sequence of them, or an ``(n, h, w)`` array.
+ store
+ Accumulate into an existing store (so several prompt clicks on several
+ frames build one training set). A new one is created when omitted, which
+ needs *frame_shape* or at least one mask to infer it from.
+ gap
+ Pixels either side of the mask boundary left **unlabelled**. The boundary
+ is where the prompt model is least certain and where a particle's own
+ soft edge lives; labelling it either way teaches the head the wrong
+ thing about edges, which shows up as systematically over- or
+ under-sized instances downstream.
+ background_dilation
+ Outer radius of the background ring, in pixels. It must be wider than
+ *gap* or there is no ring at all.
+ erode
+ Pixels eroded off the mask interior before it becomes a particle label,
+ for the same reason as *gap*. **Skipped for any mask the erosion would
+ empty** — a 3 px particle is exactly the object plan §0.9 is about, and
+ silently dropping it here would defeat the whole bootstrap.
+
+ Returns
+ -------
+ LabelStore
+ The store, with *particle_class* painted on the mask interiors and
+ *background_class* on the surrounding rings. Rings never cover another
+ mask, so particle A is not taught as background for particle B.
+
+ Notes
+ -----
+ Each mask gets its own full-frame dilation, which is O(n_masks) frame-sized
+ boolean passes. That is fine and deliberate at the scale this runs at — a
+ handful of prompt clicks — and it is why this is not the batch path: the batch
+ path is "train once, apply to N frames", which never comes back here.
+ """
+ from scipy.ndimage import binary_dilation, binary_erosion
+
+ arr = np.asarray(masks)
+ if arr.dtype != bool:
+ arr = arr.astype(bool)
+ if arr.ndim == 2:
+ arr = arr[None]
+ if arr.ndim != 3:
+ raise ValueError(
+ f"masks must be (h, w) or (n, h, w) boolean; got shape {arr.shape}")
+ if background_dilation <= gap:
+ raise ValueError(
+ f"background_dilation ({background_dilation}) must exceed gap "
+ f"({gap}) or the background ring is empty")
+
+ shape = tuple(arr.shape[1:]) if frame_shape is None else tuple(frame_shape)
+ if store is None:
+ store = LabelStore(frame_shape=shape)
+ elif tuple(store.frame_shape) != shape:
+ raise ValueError(
+ f"masks are {shape} but the store holds {store.frame_shape} labels")
+ store.class_by_id(particle_class)
+ store.class_by_id(background_class)
+
+ union = arr.any(axis=0)
+ # Everything within `gap` of ANY mask is off-limits as background — computed
+ # once over the union rather than per mask, so overlapping prompts agree.
+ near_any = binary_dilation(union, iterations=int(gap)) if gap > 0 else union
+
+ for m in arr:
+ if not m.any():
+ continue
+ inner = m
+ if erode > 0:
+ shrunk = binary_erosion(m, iterations=int(erode))
+ if shrunk.any():
+ inner = shrunk
+ store.paint(t, inner, particle_class)
+ ring = binary_dilation(m, iterations=int(background_dilation)) & ~near_any
+ if ring.any():
+ store.paint(t, ring, background_class)
+ return store
+
+
+# ── the head ─────────────────────────────────────────────────────────────────
+
+def _frame_getter(frames) -> Callable[[int], np.ndarray]:
+ """``t -> 2-D frame`` for a callable, a mapping, a stack, or a single frame.
+
+ Delegates the stack cases to :func:`spyde.drift.frames.frame_source`, which
+ already handles a HyperSpy signal, a dask array, a numpy array and a sequence
+ of frames and — importantly — reads exactly one frame at a time. Training
+ over a long movie must never materialise it (CLAUDE.md memory-safety rule),
+ and re-deriving that accessor here would be a second place to get it wrong.
+ """
+ if callable(frames):
+ return frames
+ if isinstance(frames, dict):
+ return lambda t: np.asarray(frames[int(t)])
+ arr = frames if hasattr(frames, "ndim") else None
+ if arr is not None and arr.ndim == 2:
+ return lambda t, _a=arr: np.asarray(_a)
+ from spyde.drift.frames import frame_source
+ _n, get_frame, _shape = frame_source(frames)
+ return get_frame
+
+
+class ScribbleClassifier:
+ """Multi-class pixel classifier over the torch feature stack.
+
+ Parameters
+ ----------
+ spec
+ The feature stack to classify. Saved with the weights — a model and the
+ channels it was trained on are meaningless apart.
+ hidden
+ Hidden-layer width. 64 is where agreement with the RandomForest reference
+ stops improving — measured IoU on the fixture: 16 → 0.915, 32 → 0.939,
+ **64 → 0.941**, 128 → 0.936, 256 → 0.934 — and the fit is flat in width
+ (0.45–0.58 s across all of those) because the cost is per-step dispatch,
+ not arithmetic, so there is nothing to buy by going narrower.
+ epochs, lr, weight_decay
+ Full-batch Adam. 300 steps, again from the parity measurement:
+ 100 → 0.832, 200 → 0.899, **300 → 0.941**, 500 → 0.864 (it starts
+ over-tightening the boundary past 300). One step is ~1.5 ms at *any* torch
+ thread count from 1 to 24, i.e. entirely per-step overhead, so the epoch
+ count is a fixed ~0.45 s independent of frame size and of how much was
+ painted.
+
+ Full-batch rather than mini-batch because the training set is *thousands*
+ of rows by design (plan B3): a mini-batch loop would add a shuffle order,
+ and therefore a seed dependence the user would perceive as the model
+ changing when they changed nothing, in exchange for no speed at this size.
+ seed
+ Weight initialisation. Same seed + same labels → same model, bit for bit
+ (``TestDeterminism``).
+ device
+ ``None`` auto-selects CUDA/MPS/CPU. Pass ``"cpu"`` in tests: torch-CUDA
+ work segfaults under the pytest process on Windows (CLAUDE.md).
+ """
+
+ def __init__(
+ self,
+ spec: FeatureSpec | None = None,
+ *,
+ hidden: int = 64,
+ epochs: int = 300,
+ lr: float = 0.05,
+ weight_decay: float = 1e-4,
+ seed: int = 0,
+ device=None,
+ ) -> None:
+ self.spec = spec or FeatureSpec()
+ self.hidden = int(hidden)
+ self.epochs = int(epochs)
+ self.lr = float(lr)
+ self.weight_decay = float(weight_decay)
+ self.seed = int(seed)
+ self.device = select_device(device) if not hasattr(device, "type") else device
+ self.classes: list[ScribbleClass] = []
+ self._net = None
+ self._mean = None # (C,) feature standardisation, from training
+ self._std = None
+ self.report: dict[str, Any] = {}
+
+ # -- state ---------------------------------------------------------------
+
+ @property
+ def is_trained(self) -> bool:
+ return self._net is not None
+
+ @property
+ def particle_class_ids(self) -> list[int]:
+ return [c.id for c in self.classes if c.particle]
+
+ @property
+ def boundary_class_ids(self) -> list[int]:
+ """Trained classes marked :attr:`ScribbleClass.boundary`.
+
+ Empty when the user never painted a boundary — ``fit`` drops classes with
+ no labelled pixels, so this answers "was a boundary actually taught",
+ not "does a boundary class exist in the list". That distinction is what
+ :meth:`segment` switches on.
+ """
+ return [c.id for c in self.classes if c.boundary]
+
+ @property
+ def has_boundary(self) -> bool:
+ """True when a boundary class carries trained weight — i.e. when
+ :meth:`segment` will take the connected-components route."""
+ return bool(self.boundary_class_ids)
+
+ def _require_trained(self) -> None:
+ if not self.is_trained:
+ raise RuntimeError(
+ "this ScribbleClassifier has not been trained — call fit() with a "
+ "LabelStore first")
+
+ # -- training ------------------------------------------------------------
+
+ def fit(self, store: LabelStore, frames, *,
+ progress: Callable[[int, int], None] | None = None) -> dict[str, Any]:
+ """Train on every labelled pixel in *store*.
+
+ Parameters
+ ----------
+ store
+ The accumulated scribbles. Every labelled frame contributes; classes
+ with no labelled pixels are dropped from the head (a column that
+ never sees a positive example would only ever emit noise).
+ frames
+ How to get frame *t*: a callable, a ``{t: frame}`` mapping, a 3-D
+ stack, a HyperSpy signal, or a single 2-D frame (for a one-frame
+ store). Read one frame at a time — never materialised.
+ progress
+ ``progress(done, total)`` over the labelled frames, then once more at
+ completion. Featurising is the slow part, so it is reported per frame
+ rather than per training epoch.
+
+ Returns
+ -------
+ dict
+ Training report: per-class pixel counts, final loss and training
+ accuracy, the wall-clock split between featurising and fitting, and
+ the device used. This is what the caret shows, and the featurise/fit
+ split is what tells a user whether adding a labelled frame or
+ widening the head is what costs them.
+
+ Notes
+ -----
+ Cost is ``one featurise per labelled frame`` plus a fixed fit. Measured on
+ the 96×112 fixture, CPU, default spec: **14 ms featurise + 457 ms fit**, so
+ the interaction budget is met with room, and re-training after adding a
+ stroke on a *new* frame costs one more featurise (≈1.1 s at 2048²). There is
+ deliberately **no cached feature sampler**: the cache would have to be
+ invalidated when the underlying frames change (a re-drift, a different
+ node), and nothing in this API can observe that — a silently stale
+ training set is far worse than a re-featurise the report already shows the
+ cost of.
+ """
+ torch = import_torch()
+ if len(store) == 0:
+ raise ValueError("nothing painted yet — the label store is empty")
+
+ get_frame = _frame_getter(frames)
+ t_frames = store.labelled_frames()
+
+ # ONE lock acquisition around the whole fit, not one per stage. Every line
+ # below submits to the device — the per-frame featurise, the cat, the
+ # standardisation, the optimiser loop, the accuracy read — and MPS needs
+ # all of them serialised (CLAUDE.md § GPU Computing). Held across the fit
+ # for the same reason `drift.translation.solve_translation` holds it
+ # across a solve: this completes in well under a second by contract, so a
+ # concurrent preview waits a bounded time, and a partial hold would leave
+ # exactly the gaps that took the backend down last time. Reentrant, so
+ # `sample_features` taking it again is free. Null context off MPS.
+ t0 = time.perf_counter()
+ with accelerator_lock(self.device):
+ xs, ys = [], []
+ for i, t in enumerate(t_frames):
+ idx, cls = store.at(t)
+ if not idx.size:
+ continue
+ frame = np.asarray(get_frame(t))
+ if tuple(frame.shape) != tuple(store.frame_shape):
+ raise ValueError(
+ f"frame {t} is {frame.shape} but the label store holds "
+ f"{store.frame_shape} labels — the flat indices would "
+ "land in the wrong pixels")
+ xs.append(sample_features(frame, idx, self.spec,
+ device=self.device))
+ ys.append(torch.as_tensor(cls.astype(np.int64),
+ device=self.device))
+ if progress is not None:
+ progress(i + 1, len(t_frames))
+ t_feat = time.perf_counter() - t0
+
+ X = torch.cat(xs, dim=0)
+ y_raw = torch.cat(ys, dim=0)
+
+ # Only classes that actually carry labels become head columns,
+ # remapped to a contiguous 0..K-1 so cross-entropy has no dead output.
+ present = sorted({int(v) for v in y_raw.unique().tolist()})
+ if len(present) < 2:
+ raise ValueError(
+ f"only one class is painted (id {present[0]}) — a classifier "
+ "needs at least two, e.g. a particle and some background")
+ self.classes = [store.class_by_id(c) for c in present]
+ lut = torch.full((max(present) + 1,), -1, dtype=torch.long,
+ device=self.device)
+ for k, cid in enumerate(present):
+ lut[cid] = k
+ y = lut[y_raw]
+
+ t1 = time.perf_counter()
+ self._mean = X.mean(dim=0)
+ # Standardise per channel. A guard on the std and not an epsilon: a
+ # constant channel (a rank filter on a flat region) divided by a tiny
+ # number amplifies float noise to unit scale, and the head then fits
+ # it — the same failure the drift solver's phase FLOOR exists for.
+ std = X.std(dim=0, unbiased=False)
+ self._std = torch.where(std > 1e-6, std, torch.ones_like(std))
+ Xn = (X - self._mean) / self._std
+
+ self._net = _build_mlp(torch, X.shape[1], self.hidden, len(present),
+ self.seed, self.device)
+ counts = torch.bincount(y, minlength=len(present)).to(Xn.dtype)
+ # Class-balanced loss. Scribbles are wildly unbalanced by nature — a
+ # user paints a few dabs on particles and sweeps whole regions of
+ # background — and unweighted cross-entropy on a 40:1 split simply
+ # learns "background", which reads as the classifier not working.
+ weight = X.shape[0] / (len(present) * counts.clamp_min(1.0))
+ loss_fn = torch.nn.CrossEntropyLoss(weight=weight)
+ opt = torch.optim.Adam(self._net.parameters(), lr=self.lr,
+ weight_decay=self.weight_decay)
+ loss = float("nan")
+ for _ in range(self.epochs):
+ opt.zero_grad(set_to_none=True)
+ out = self._net(Xn)
+ lo = loss_fn(out, y)
+ lo.backward()
+ opt.step()
+ loss = float(lo.detach())
+ with torch.no_grad():
+ acc = float((self._net(Xn).argmax(dim=1) == y).to(Xn.dtype).mean())
+ t_fit = time.perf_counter() - t1
+
+ self.report = {
+ "device": str(self.device),
+ "n_pixels": int(X.shape[0]),
+ "n_channels": int(X.shape[1]),
+ "n_classes": len(present),
+ # Whether a boundary class carries trained weight, i.e. whether the
+ # split will take its connected-components route. Surfaced in the
+ # report because it is the difference between a 0.33 s and a 1.78 s
+ # split at 4096², and the user is the one who decides it by painting.
+ "has_boundary": bool([c for c in self.classes if c.boundary]),
+ "labelled_frames": list(t_frames),
+ "pixels_per_class": {str(c.id): int(n) for c, n in
+ zip(self.classes, counts.tolist())},
+ "loss": loss,
+ "train_accuracy": acc,
+ "featurise_s": t_feat,
+ "fit_s": t_fit,
+ }
+ if progress is not None:
+ progress(len(t_frames), len(t_frames))
+ log.info("[scribble] trained on %d px x %d ch, %d classes: acc %.3f "
+ "(featurise %.2f s, fit %.2f s, %s)", X.shape[0], X.shape[1],
+ len(present), acc, t_feat, t_fit, self.device)
+ return self.report
+
+ # -- prediction ----------------------------------------------------------
+
+ def predict_class_proba(self, frame) -> np.ndarray:
+ """``(K, H, W)`` float32 softmax over the trained classes.
+
+ Column *k* is ``self.classes[k]``. Pixels that were non-finite in *frame*
+ get probability 0 in **every** class, so the columns do not sum to 1
+ there — that is deliberate and is what the NaN-border contract means: an
+ invalid pixel is not "probably background", it is not a measurement.
+ """
+ self._require_trained()
+ torch = import_torch()
+ prepared = prepare_frame(frame, self.spec)
+ h, w = prepared.image.shape
+
+ # The lock spans the allocation, the per-band head evaluation AND the
+ # read-back: a device->host copy is a submission too, and doing it after
+ # releasing is exactly the kind of gap that reopens the MPS crash.
+ with accelerator_lock(self.device):
+ out = torch.zeros((len(self.classes), h, w), device=self.device,
+ dtype=torch.float32)
+
+ def band(y0: int, y1: int, stack) -> None:
+ # `stack` is a row slice of a larger tensor and so not contiguous;
+ # `reshape` copies when it must, which `view` would refuse to do.
+ flat = stack.reshape(stack.shape[0], -1).t() # (rows*w, C)
+ with torch.no_grad():
+ p = torch.softmax(
+ self._net((flat - self._mean) / self._std), dim=1)
+ out[:, y0:y1] = p.t().reshape(len(self.classes), y1 - y0, w)
+
+ map_feature_bands(prepared, self.spec, device=self.device, fn=band)
+ proba = out.detach().cpu().numpy()
+
+ proba[:, ~prepared.valid] = 0.0
+ return proba
+
+ def predict_proba(self, frame) -> np.ndarray:
+ """``(H, W)`` float32 **foreground** probability, 0..1.
+
+ The sum over every class with :attr:`ScribbleClass.particle` set, which is
+ the per-frame probability map the plan's §0.2 spine consumes — hand it
+ straight to :func:`spyde.particles.classical.split_instances`.
+
+ Zero inside a NaN-padded region (plan trap 2).
+ """
+ return self.predict_foreground_boundary(frame)[0]
+
+ def predict_boundary_proba(self, frame) -> np.ndarray | None:
+ """``(H, W)`` float32 **boundary** probability, or None if untrained.
+
+ None and not a zero map, because the two mean different things to
+ :func:`~spyde.particles.classical.split_instances`: "no boundary was
+ taught, use the watershed" versus "a boundary was taught and this frame
+ has none of it", which would leave every touching pair merged.
+ """
+ return self.predict_foreground_boundary(frame)[1]
+
+ def predict_foreground_boundary(
+ self, frame
+ ) -> tuple[np.ndarray, np.ndarray | None]:
+ """``(foreground, boundary)`` from **one** pass over the frame.
+
+ Both maps come out of the same softmax, and that softmax is the whole
+ cost of this engine — 1.2 s of a 1.5 s prediction at 4096². Calling
+ :meth:`predict_proba` and :meth:`predict_boundary_proba` separately would
+ featurise the frame twice for two views of one result, so the pair is the
+ primitive and the two singular accessors are the wrappers.
+
+ *boundary* is None when no trained class is marked
+ :attr:`ScribbleClass.boundary`.
+ """
+ self._require_trained()
+ proba = self.predict_class_proba(frame)
+ wanted = [k for k, c in enumerate(self.classes) if c.particle]
+ if not wanted:
+ raise RuntimeError(
+ "no trained class is marked as a particle class, so there is no "
+ "foreground to report — set ScribbleClass.particle on at least "
+ "one of "
+ f"{[c.name for c in self.classes]} and retrain")
+ edge = [k for k, c in enumerate(self.classes) if c.boundary]
+ return _sum_planes(proba, wanted), (_sum_planes(proba, edge) if edge
+ else None)
+
+ def predict_labels(self, frame) -> np.ndarray:
+ """``(H, W)`` int16 argmax over classes, as **class ids**.
+
+ :data:`UNLABELLED` (-1) where the source pixel was non-finite. Class ids,
+ not column indices, so the value indexes straight into the user's class
+ list and its colour.
+ """
+ proba = self.predict_class_proba(frame)
+ ids = np.asarray([c.id for c in self.classes], dtype=np.int16)
+ out = ids[proba.argmax(axis=0)]
+ out[proba.sum(axis=0) <= 0.0] = UNLABELLED
+ return out
+
+ def segment(self, frame, params=None) -> np.ndarray:
+ """Convenience: probability → labelled instances via the shared split.
+
+ The whole point of plan §0.2 is that this engine stops at a probability
+ map and the instance stage is written once, so this is a short forwarder
+ — kept here only so the caller does not have to remember which module
+ owns the split.
+
+ **The boundary is passed on when there is one**, and that is what makes
+ this engine fast rather than merely accurate: with a taught boundary the
+ split takes its connected-components route and never runs the distance
+ transform or the watershed. Without one it falls back to the watershed,
+ so a user who has not painted any boundary gets exactly the behaviour
+ they had before — never a silently worse split.
+ """
+ from spyde.particles.classical import SegmentParams, split_instances
+ fg, bnd = self.predict_foreground_boundary(frame)
+ return split_instances(fg, params or SegmentParams(), boundary=bnd)
+
+ # -- serialisation -------------------------------------------------------
+
+ def save(self, path: str) -> None:
+ """Write weights + :class:`FeatureSpec` + classes to one ``.npz``.
+
+ One file, never two: a recipe is the spec *and* the weights *and* the
+ feature standardisation, and any of them alone predicts nonsense. Written
+ with ``np.savez_compressed`` and read back with ``allow_pickle=False``,
+ matching :meth:`spyde.signals.particles.SpyDEParticles.save` — a model
+ file should not be able to execute code on load.
+ """
+ self._require_trained()
+ with accelerator_lock(self.device):
+ state = {k: v.detach().cpu().numpy()
+ for k, v in self._net.state_dict().items()}
+ mean = self._mean.detach().cpu().numpy()
+ std = self._std.detach().cpu().numpy()
+ meta = {
+ "format_version": FORMAT_VERSION,
+ "spec": self.spec.to_dict(),
+ "classes": [c.to_dict() for c in self.classes],
+ "hidden": self.hidden,
+ "epochs": self.epochs,
+ "lr": self.lr,
+ "weight_decay": self.weight_decay,
+ "seed": self.seed,
+ "state_keys": sorted(state),
+ "report": _jsonable(self.report),
+ }
+ np.savez_compressed(
+ path,
+ meta=np.array(json.dumps(meta)),
+ feature_mean=mean,
+ feature_std=std,
+ **{f"w_{k}": v for k, v in state.items()},
+ )
+
+ @classmethod
+ def load(cls, path: str, *, device=None) -> "ScribbleClassifier":
+ """Read a model written by :meth:`save`.
+
+ The host→device transfers are taken under the accelerator lock. A cold
+ model load is the *specific* unlocked call site CLAUDE.md names — the
+ neural detector's ``load_model`` had exactly this hole and it was one of
+ the two threads in the MPS crash.
+ """
+ torch = import_torch()
+ with np.load(path, allow_pickle=False) as z:
+ meta = json.loads(str(z["meta"].item()))
+ ver = meta.get("format_version")
+ if ver != FORMAT_VERSION:
+ raise ValueError(
+ f"unsupported scribble model format version {ver!r} "
+ f"(this build reads {FORMAT_VERSION})")
+ model = cls(
+ FeatureSpec.from_dict(meta["spec"]),
+ hidden=int(meta["hidden"]),
+ epochs=int(meta.get("epochs", 300)),
+ lr=float(meta.get("lr", 0.05)),
+ weight_decay=float(meta.get("weight_decay", 1e-4)),
+ seed=int(meta.get("seed", 0)),
+ device=device,
+ )
+ model.classes = [ScribbleClass.from_dict(c) for c in meta["classes"]]
+ model.report = meta.get("report") or {}
+ n_in = int(np.size(z["feature_mean"]))
+ if n_in != model.spec.n_channels:
+ raise ValueError(
+ f"model expects {n_in} feature channels but its saved "
+ f"FeatureSpec produces {model.spec.n_channels} — the spec and "
+ "the weights have come apart")
+ with accelerator_lock(model.device):
+ model._mean = torch.as_tensor(z["feature_mean"],
+ device=model.device)
+ model._std = torch.as_tensor(z["feature_std"],
+ device=model.device)
+ net = _build_mlp(torch, n_in, model.hidden, len(model.classes),
+ model.seed, model.device)
+ net.load_state_dict(
+ {k: torch.as_tensor(z[f"w_{k}"], device=model.device)
+ for k in meta["state_keys"]})
+ model._net = net
+ return model
+
+
+def _sum_planes(proba: np.ndarray, cols: list[int]) -> np.ndarray:
+ """Sum the given class planes of a ``(K, H, W)`` softmax, float32.
+
+ The single-column case is the overwhelmingly common one (one particle class,
+ one boundary class) and it is taken by a plain slice: ``proba[[k]].sum(0)``
+ on a 4096² map builds a fancy-indexed copy and then reduces it, which
+ measured **113 ms** for the pair against ~0 ms for two views.
+ """
+ if len(cols) == 1:
+ plane = proba[cols[0]]
+ return plane if plane.dtype == np.float32 else plane.astype(np.float32)
+ return proba[cols].sum(axis=0).astype(np.float32)
+
+
+def _build_mlp(torch, n_in: int, hidden: int, n_out: int, seed: int, device):
+ """One hidden layer, ReLU, deterministically initialised from *seed*.
+
+ The weights are drawn from an explicitly seeded ``torch.Generator`` and copied
+ in, rather than letting ``nn.Linear`` use the global RNG: the global stream is
+ shared with everything else in the process (the neural detector, a dask
+ worker), so "same seed, same labels, same model" would otherwise depend on
+ what else happened to run first. ``torch.nn.init`` has no generator argument
+ on this torch version, which is why the scaling is written out.
+
+ The construction is nevertheless wrapped in ``fork_rng``, because
+ ``nn.Linear.__init__`` draws its own default initialisation from the global
+ stream *before* it is overwritten here — so without the fork, training a
+ scribble model would silently shift every other consumer's random sequence.
+ ``devices=[]`` forks the CPU generator only; forking CUDA's would initialise
+ the CUDA context as a side effect.
+ """
+ g = torch.Generator().manual_seed(int(seed))
+ with torch.random.fork_rng(devices=[]):
+ net = torch.nn.Sequential(
+ torch.nn.Linear(n_in, hidden),
+ torch.nn.ReLU(),
+ torch.nn.Linear(hidden, n_out),
+ )
+ with torch.no_grad():
+ for lin, fan_in in ((net[0], n_in), (net[2], hidden)):
+ w = torch.randn(lin.weight.shape, generator=g) / math.sqrt(fan_in)
+ lin.weight.copy_(w)
+ lin.bias.zero_()
+ return net.to(device)
+
+
+def _jsonable(obj):
+ """Coerce a report dict to something ``json.dumps`` accepts."""
+ if isinstance(obj, dict):
+ return {str(k): _jsonable(v) for k, v in obj.items()}
+ if isinstance(obj, (list, tuple)):
+ return [_jsonable(v) for v in obj]
+ if isinstance(obj, (np.integer,)):
+ return int(obj)
+ if isinstance(obj, (np.floating,)):
+ return float(obj)
+ return obj
+
+
+# ── the sklearn parity reference (test-only, kept next to what it checks) ─────
+
+def random_forest_reference(
+ store: LabelStore,
+ frames,
+ spec: FeatureSpec | None = None,
+ *,
+ n_estimators: int = 100,
+ seed: int = 0,
+ device=None,
+):
+ """Train ``sklearn.ensemble.RandomForestClassifier`` on the SAME features.
+
+ This is the acceptance reference for the whole engine (plan B3): the forest is
+ what ParticleSpy and ilastik use, so agreement on identical labels and
+ identical channels is what says the torch head is a re-implementation rather
+ than a different algorithm that happens to produce pictures.
+
+ It lives here, beside the head it checks, rather than in the test file — the
+ two must read the same feature stack through the same sampler, and the moment
+ the reference has its own copy of that plumbing the comparison stops being
+ apples to apples.
+
+ Returns
+ -------
+ (model, predict)
+ ``predict(frame)`` gives an ``(H, W)`` float32 foreground probability, on
+ the same convention as
+ :meth:`ScribbleClassifier.predict_proba` including the NaN-border zeroing.
+ """
+ from sklearn.ensemble import RandomForestClassifier
+
+ spec = spec or FeatureSpec()
+ device = select_device(device) if not hasattr(device, "type") else device
+ get_frame = _frame_getter(frames)
+
+ xs, ys = [], []
+ with accelerator_lock(device):
+ for t in store.labelled_frames():
+ idx, cls = store.at(t)
+ if not idx.size:
+ continue
+ xs.append(sample_features(get_frame(t), idx, spec,
+ device=device).detach().cpu().numpy())
+ ys.append(cls.astype(np.int64))
+ X = np.concatenate(xs, axis=0)
+ y = np.concatenate(ys, axis=0)
+
+ rf = RandomForestClassifier(n_estimators=n_estimators, random_state=seed,
+ class_weight="balanced", n_jobs=-1)
+ rf.fit(X, y)
+ particle_ids = {c.id for c in store.classes if c.particle}
+ cols = [i for i, c in enumerate(rf.classes_) if int(c) in particle_ids]
+
+ def predict(frame) -> np.ndarray:
+ from spyde.particles.features import feature_stack
+ prepared = prepare_frame(frame, spec)
+ with accelerator_lock(device):
+ # The forest is CPU-only, so unlike `predict_class_proba` this one
+ # genuinely does materialise the whole stack (numpy) before scoring —
+ # which is the other reason the forest is a reference and not the
+ # shipped head.
+ stack = feature_stack(prepared, spec, device=device)
+ c, h, w = stack.shape
+ p = rf.predict_proba(stack.reshape(c, -1).T)
+ out = p[:, cols].sum(axis=1).reshape(h, w).astype(np.float32)
+ out[~prepared.valid] = 0.0
+ return out
+
+ return rf, predict
diff --git a/spyde/particles/scribble_cnn.py b/spyde/particles/scribble_cnn.py
new file mode 100644
index 00000000..a89c796d
--- /dev/null
+++ b/spyde/particles/scribble_cnn.py
@@ -0,0 +1,659 @@
+"""
+scribble_cnn.py — PROTOTYPE. A small CNN trained on the same scribbles, as an
+A/B alternative to :class:`spyde.particles.scribble.ScribbleClassifier`.
+
+**This is a prototype and is deliberately not wired into anything.** It is not
+in ``spyde.particles.__init__``, no action dispatches it, no caret exposes it
+and the batch path does not know it exists. It exists so the two engines can be
+measured against each other on identical scribbles
+(``spyde/tests/benchmark_scribble_cnn.py``); the shipped engine is untouched.
+
+What it replaces, and what it does not
+--------------------------------------
+The shipped engine is two stages: 36 hand-crafted channels
+(:mod:`spyde.particles.features`) and a per-pixel MLP over them. This replaces
+**both** with one small U-Net that reads the raw (robustly standardised) frame
+and emits per-class logits at input resolution. Everything downstream is
+unchanged, and that is the whole point of the design:
+
+* :meth:`predict_foreground_boundary` returns the same ``(foreground,
+ boundary)`` pair, so :func:`spyde.particles.classical.split_instances` takes
+ its connected-components route exactly as it already does for the MLP — no new
+ instance decoder, no star-convex head, no second downstream path to maintain.
+* The class list is :class:`~spyde.particles.scribble.ScribbleClass`, so
+ ``particle`` and ``boundary`` flags mean what they already mean, and several
+ particle phases still sum into one foreground map.
+* Non-finite input pixels are forced to zero probability in every class, the
+ same NaN-border contract (plan trap 2) — both engines call
+ :func:`spyde.particles.features.prepare_frame`, so they see a bit-identical
+ input image and a bit-identical validity mask.
+
+Training on sparse scribbles
+----------------------------
+A scribble is a few thousand pixels in a 16.7 M-pixel frame, so two things
+follow and both are load-bearing:
+
+**The loss is masked.** Unlabelled pixels carry ``ignore_index`` and contribute
+nothing. A dense loss would need a dense target, which does not exist — the user
+painted strokes, not a segmentation.
+
+**Training runs on CROPS around the scribbles, never on whole frames.** This is
+what makes the train time tolerable and it is also the memory-safety rule
+(CLAUDE.md): a movie is never materialised, one frame is read at a time, and
+only the crop windows that actually contain labelled pixels are ever pushed
+through the net. Whole-frame training at 4096² would be ~0.3 s *per step* for
+the small net; a batch of eight 128² crops is ~1/16 of one frame.
+
+Inference is TILED
+------------------
+Measured on this dev box (TITAN X Pascal, fp32, 4096²): ``base=32, levels=3``
+costs **0.737 s tiled at 16×1024²** against **7.72 s whole-frame** — a 10×
+difference that is memory pressure on a 12 GB card, not arithmetic. So the
+whole-frame path is not offered above :data:`TILE_ABOVE_PIXELS`.
+
+Three measured traps that shaped this file
+------------------------------------------
+1. **fp16 is SLOWER on Pascal** (fp16:fp32 rate 1:64). ``autocast`` measured
+ 7.72 → 17.9 s and 14.5 → 22.5 s. There is deliberately no AMP option here.
+2. **Level count matters more than parameter count.** ``base=32, levels=2``
+ measured **25 s** whole-frame — far worse than ``levels=3`` with 4× the
+ parameters — because fewer downsamples leaves more work at full resolution.
+ So :data:`CONFIGS` names the two configurations worth running and a caller
+ who invents a third should measure it before believing it.
+3. **Tile anything above the tiny net.** See above.
+"""
+from __future__ import annotations
+
+import logging
+import math
+import time
+from typing import Any, Callable
+
+import numpy as np
+
+from spyde.device_lock import accelerator_lock
+from spyde.particles.features import import_torch, prepare_frame, select_device
+from spyde.particles.scribble import LabelStore, ScribbleClass, _frame_getter
+
+log = logging.getLogger(__name__)
+
+#: The two configurations worth running, per the forward-pass sweep in the
+#: module docstring. ``(base, levels)`` → the label the benchmark reports.
+CONFIGS: dict[str, tuple[int, int]] = {
+ "tiny": (16, 2), # 117 k params, 0.289 s whole-frame at 4096²
+ "small": (32, 3), # 1.9 M params, 0.737 s tiled at 4096²
+}
+
+#: Above this many pixels a frame is always tiled for inference (see the module
+#: docstring — 10× on the small net, and it only gets worse with frame size).
+TILE_ABOVE_PIXELS: int = 2048 * 2048
+
+#: Default inference tile edge, in pixels. 1024 is what the 16×1024² measurement
+#: used and it is a multiple of every ``2**levels`` in :data:`CONFIGS`.
+TILE_EDGE: int = 1024
+
+#: Default training crop edge. A multiple of 8, so it is legal for ``levels`` up
+#: to 3 without padding.
+CROP_EDGE: int = 128
+
+#: Ignored target value for an unlabelled pixel. Not
+#: :data:`~spyde.particles.scribble.UNLABELLED` by coincidence — it is the same
+#: -1, and ``cross_entropy(ignore_index=-1)`` is what makes the sparse loss work.
+IGNORE = -1
+
+
+# ── the net ──────────────────────────────────────────────────────────────────
+
+_SEG_UNET = None
+
+
+def _seg_unet_class():
+ """Define (once) the ``SpotUNet`` subclass with a segmentation head.
+
+ Built lazily rather than at module scope because ``spyde.models.unet``
+ imports torch eagerly, and everything else in :mod:`spyde.particles` is
+ careful not to (:func:`spyde.particles.features.import_torch`).
+ """
+ global _SEG_UNET
+ if _SEG_UNET is not None:
+ return _SEG_UNET
+
+ import torch.nn as nn
+
+ from spyde.models.unet import SpotUNet
+
+ class SegUNet(SpotUNet):
+ """``SpotUNet``'s encoder/decoder with a K-class 1×1 head.
+
+ The body is the vendored one, unmodified — that is the budget (no
+ stardist, no cellpose, no SAM), and it is why the forward-pass costs
+ measured for ``SpotUNet`` transfer here directly. Only the heads change:
+ the spot detector's ``(heatmap, offset)`` pair is replaced by one
+ ``Conv2d(base, K, 1)`` emitting per-class logits at input resolution.
+ The spot heads are DELETED rather than left unused — otherwise they sit
+ in the optimiser's parameter list receiving no gradient and are saved
+ with the model.
+ """
+
+ def __init__(self, n_classes: int, base: int = 16, levels: int = 2):
+ super().__init__(in_ch=1, base=base, levels=levels)
+ del self.head_hm, self.head_off
+ self.head_seg = nn.Conv2d(base, int(n_classes), 1)
+ nn.init.zeros_(self.head_seg.bias)
+
+ def forward(self, x):
+ import torch
+ feats = []
+ h = x
+ for i, enc in enumerate(self.enc):
+ h = enc(h if i == 0 else self.pool(h))
+ feats.append(h)
+ d = feats[-1]
+ for j, (up, dec) in enumerate(zip(self.up, self.dec)):
+ d = dec(torch.cat([up(d), feats[self.levels - 1 - j]], 1))
+ return self.head_seg(d)
+
+ _SEG_UNET = SegUNet
+ return SegUNet
+
+
+def build_net(n_classes: int, *, base: int = 16, levels: int = 2, seed: int = 0):
+ """A :class:`~spyde.models.unet.SpotUNet` body with a K-class 1×1 head.
+
+ Seeded through ``fork_rng`` for the same reason
+ :func:`spyde.particles.scribble._build_mlp` is: the global torch RNG is
+ shared with everything else in the process, so initialising from it would
+ make "same seed, same labels, same model" depend on what ran first — and
+ would silently shift every other consumer's random stream. ``devices=[]``
+ forks the CPU generator only; forking CUDA's would initialise the CUDA
+ context as a side effect.
+ """
+ torch = import_torch()
+ with torch.random.fork_rng(devices=[]):
+ torch.manual_seed(int(seed))
+ return _seg_unet_class()(int(n_classes), int(base), int(levels))
+
+
+# ── crop planning ────────────────────────────────────────────────────────────
+
+def crop_windows(ys: np.ndarray, xs: np.ndarray, shape: tuple[int, int],
+ crop: int, mult: int) -> list[tuple[int, int, int, int]]:
+ """``(y0, y1, x0, x1)`` windows covering every labelled pixel, deduplicated.
+
+ Each labelled pixel is assigned to the ONE window that contains it most
+ centrally (a half-crop grid, snapped), and the unique windows are returned.
+ Centrally, because the U-Net's receptive field is what gives a pixel its
+ context: a labelled pixel pinned to a crop's edge is classified from half the
+ surroundings it will have at inference time, and that mismatch is invisible
+ in the training loss.
+
+ The window edge is clamped to a multiple of *mult* (``2**levels``) so the
+ pooling stack divides evenly, and to the frame, so a small frame trains as
+ one whole-frame crop rather than as a reflect-padded fiction.
+ """
+ h, w = int(shape[0]), int(shape[1])
+ ch = min(int(crop), (h // mult) * mult)
+ cw = min(int(crop), (w // mult) * mult)
+ if ch < mult or cw < mult:
+ raise ValueError(
+ f"a {h}x{w} frame is too small to train a {mult}-divisible crop "
+ f"from; use fewer levels")
+
+ # Half-crop grid: round the pixel to the nearest window CENTRE, then clamp
+ # the resulting top-left into the frame.
+ def tops(v: np.ndarray, span: int, extent: int) -> np.ndarray:
+ step = max(1, span // 2)
+ t = np.rint((v - span / 2.0) / step).astype(np.int64) * step
+ return np.clip(t, 0, extent - span)
+
+ ty, tx = tops(ys, ch, h), tops(xs, cw, w)
+ seen = {(int(a), int(b)) for a, b in zip(ty, tx)}
+ return [(y, y + ch, x, x + cw) for y, x in sorted(seen)]
+
+
+# ── the head ─────────────────────────────────────────────────────────────────
+
+class ScribbleCNN:
+ """Prototype CNN pixel classifier — same output contract as the MLP engine.
+
+ Parameters
+ ----------
+ base, levels
+ U-Net width and downsample count. See :data:`CONFIGS`; ``levels`` is the
+ knob that decides both the receptive field and the cost, and the two
+ do not trade off the way parameter count suggests (module docstring
+ trap 2).
+ crop
+ Training crop edge, px. Training never touches a whole frame.
+ steps, batch, lr, weight_decay
+ Adam over crop mini-batches. ``steps`` is optimiser steps and not epochs
+ on purpose: the number of crops depends on how much the user painted and
+ how far apart, so "epochs" is not a fixed amount of work and would make
+ the train time depend on the scribble layout in a way the user cannot
+ predict.
+ augment
+ Random dihedral (flip/transpose) augmentation per crop per step. Nearly
+ free and material at this label count — a few thousand labelled pixels
+ is a very small training set for a conv net.
+ tile
+ Inference tile edge, px. See the module docstring: above
+ :data:`TILE_ABOVE_PIXELS` this is not optional.
+ device
+ ``None`` auto-selects CUDA/MPS/CPU. Pass ``"cpu"`` under pytest —
+ torch-CUDA segfaults in that process on Windows (CLAUDE.md).
+ """
+
+ def __init__(
+ self,
+ *,
+ base: int = 16,
+ levels: int = 2,
+ crop: int = CROP_EDGE,
+ steps: int = 300,
+ batch: int = 8,
+ lr: float = 3e-3,
+ weight_decay: float = 1e-4,
+ augment: bool = True,
+ tile: int = TILE_EDGE,
+ seed: int = 0,
+ device=None,
+ ) -> None:
+ self.base = int(base)
+ self.levels = int(levels)
+ self.crop = int(crop)
+ self.steps = int(steps)
+ self.batch = int(batch)
+ self.lr = float(lr)
+ self.weight_decay = float(weight_decay)
+ self.augment = bool(augment)
+ self.tile = int(tile)
+ self.seed = int(seed)
+ self.device = select_device(device) if not hasattr(device, "type") else device
+ self.classes: list[ScribbleClass] = []
+ self._net = None
+ self.report: dict[str, Any] = {}
+
+ # -- state ---------------------------------------------------------------
+
+ @property
+ def mult(self) -> int:
+ """Pixel multiple the pooling stack requires: ``2**levels``."""
+ return 1 << self.levels
+
+ @property
+ def is_trained(self) -> bool:
+ return self._net is not None
+
+ @property
+ def particle_class_ids(self) -> list[int]:
+ return [c.id for c in self.classes if c.particle]
+
+ @property
+ def boundary_class_ids(self) -> list[int]:
+ return [c.id for c in self.classes if c.boundary]
+
+ @property
+ def has_boundary(self) -> bool:
+ return bool(self.boundary_class_ids)
+
+ def num_params(self) -> int:
+ self._require_trained()
+ return int(sum(p.numel() for p in self._net.parameters()))
+
+ def _require_trained(self) -> None:
+ if not self.is_trained:
+ raise RuntimeError(
+ "this ScribbleCNN has not been trained — call fit() with a "
+ "LabelStore first")
+
+ # -- training ------------------------------------------------------------
+
+ def fit(self, store: LabelStore, frames, *,
+ progress: Callable[[int, int], None] | None = None) -> dict[str, Any]:
+ """Train on every labelled pixel in *store*, from crops around them.
+
+ Same signature and same *frames* vocabulary as
+ :meth:`spyde.particles.scribble.ScribbleClassifier.fit` — a callable, a
+ ``{t: frame}`` mapping, a 3-D stack or a HyperSpy signal — so the two
+ engines are driven identically by the benchmark and by any future caret.
+ One frame is read at a time and never held.
+
+ Returns the training report: per-class pixel counts, crop count, final
+ loss and training accuracy over the labelled pixels, the wall-clock split
+ between preparing crops and optimising, and the device.
+ """
+ torch = import_torch()
+ import torch.nn.functional as F
+
+ if len(store) == 0:
+ raise ValueError("nothing painted yet — the label store is empty")
+
+ get_frame = _frame_getter(frames)
+ t_frames = store.labelled_frames()
+
+ counts_by_id: dict[int, int] = {}
+ for t in t_frames:
+ _idx, cls = store.at(t)
+ for v, n in zip(*np.unique(cls, return_counts=True)):
+ counts_by_id[int(v)] = counts_by_id.get(int(v), 0) + int(n)
+ present = sorted(counts_by_id)
+ if len(present) < 2:
+ raise ValueError(
+ f"only one class is painted (id {present[0]}) — a classifier "
+ "needs at least two, e.g. a particle and some background")
+ self.classes = [store.class_by_id(c) for c in present]
+ col = {cid: k for k, cid in enumerate(present)}
+ n_out = len(present)
+
+ # CUDA's autograd engine has to be initialised on the thread that will
+ # run backward, or the first backward segfaults on Windows (CLAUDE.md
+ # § GPU Computing). `fit` is the dispatch point, so warm it here rather
+ # than hoping the caller did.
+ _warmup_autograd(self.device)
+
+ t0 = time.perf_counter()
+ crops_x: list[np.ndarray] = []
+ crops_y: list[np.ndarray] = []
+ h, w = store.frame_shape
+ for i, t in enumerate(t_frames):
+ idx, cls = store.at(t)
+ if not idx.size:
+ continue
+ frame = np.asarray(get_frame(t))
+ if tuple(frame.shape) != tuple(store.frame_shape):
+ raise ValueError(
+ f"frame {t} is {frame.shape} but the label store holds "
+ f"{store.frame_shape} labels — the flat indices would land "
+ "in the wrong pixels")
+ # Identical preparation to the MLP engine: same NaN fill, same
+ # robust standardisation, so an A/B is about the head and not about
+ # what the two were shown.
+ img = prepare_frame(frame).image
+ ys, xs = np.divmod(idx, w)
+ target = np.full((h, w), IGNORE, dtype=np.int64)
+ target[ys, xs] = [col[int(c)] for c in cls]
+ for (y0, y1, x0, x1) in crop_windows(ys, xs, (h, w), self.crop,
+ self.mult):
+ sub = target[y0:y1, x0:x1]
+ if not (sub >= 0).any(): # pragma: no cover — defensive
+ continue
+ crops_x.append(np.ascontiguousarray(img[y0:y1, x0:x1]))
+ crops_y.append(np.ascontiguousarray(sub))
+ if progress is not None:
+ progress(i + 1, len(t_frames))
+ if not crops_x: # pragma: no cover — defensive
+ raise ValueError("no training crop contained a labelled pixel")
+
+ # ONE lock acquisition around the whole fit, for the same reason
+ # `ScribbleClassifier.fit` takes one: every line below submits to the
+ # device and MPS needs all of them serialised. Null context off MPS.
+ with accelerator_lock(self.device):
+ X = torch.as_tensor(np.stack(crops_x)[:, None],
+ device=self.device, dtype=torch.float32)
+ Y = torch.as_tensor(np.stack(crops_y), device=self.device)
+ n_crops = int(X.shape[0])
+
+ counts = torch.zeros(n_out, dtype=torch.float32, device=self.device)
+ for cid, n in counts_by_id.items():
+ counts[col[cid]] = float(n)
+ # Class-balanced, exactly as the MLP is: a user paints a few dabs on
+ # particles and sweeps whole regions of background, and unweighted
+ # cross-entropy on a 40:1 split learns "background".
+ weight = counts.sum() / (n_out * counts.clamp_min(1.0))
+
+ t1 = time.perf_counter()
+ self._net = build_net(n_out, base=self.base, levels=self.levels,
+ seed=self.seed).to(self.device)
+ self._net.train()
+ opt = torch.optim.Adam(self._net.parameters(), lr=self.lr,
+ weight_decay=self.weight_decay)
+ rng = np.random.default_rng(self.seed)
+ batch = min(self.batch, n_crops)
+
+ loss = float("nan")
+ # Pin backward to this thread: torch's multithreaded autograd
+ # engine segfaults under CUDA on Windows off the main thread.
+ prev_mt = True
+ try:
+ torch.autograd.set_multithreading_enabled(False)
+ prev_mt = False
+ for step in range(self.steps):
+ sel = rng.choice(n_crops, size=batch,
+ replace=batch > n_crops)
+ xb, yb = X[sel], Y[sel]
+ if self.augment:
+ xb, yb = _augment(torch, xb, yb, rng)
+ opt.zero_grad(set_to_none=True)
+ out = self._net(xb)
+ lo = F.cross_entropy(out, yb, weight=weight,
+ ignore_index=IGNORE)
+ lo.backward()
+ opt.step()
+ loss = float(lo.detach())
+ finally:
+ if not prev_mt:
+ torch.autograd.set_multithreading_enabled(True)
+
+ self._net.eval()
+ with torch.no_grad():
+ hit = tot = 0
+ for k in range(0, n_crops, max(1, batch)):
+ xb, yb = X[k:k + batch], Y[k:k + batch]
+ pred = self._net(xb).argmax(dim=1)
+ m = yb >= 0
+ hit += int((pred[m] == yb[m]).sum())
+ tot += int(m.sum())
+ acc = hit / max(1, tot)
+ t_fit = time.perf_counter() - t1
+
+ self.report = {
+ "engine": "cnn",
+ "device": str(self.device),
+ "base": self.base,
+ "levels": self.levels,
+ "params": self.num_params(),
+ "n_pixels": int(sum(counts_by_id.values())),
+ "n_classes": n_out,
+ "n_crops": n_crops,
+ "crop": [int(X.shape[-2]), int(X.shape[-1])],
+ "steps": self.steps,
+ "batch": batch,
+ "has_boundary": bool([c for c in self.classes if c.boundary]),
+ "labelled_frames": list(t_frames),
+ "pixels_per_class": {str(cid): int(n)
+ for cid, n in sorted(counts_by_id.items())},
+ "loss": loss,
+ "train_accuracy": acc,
+ "crops_s": t1 - t0,
+ "fit_s": t_fit,
+ }
+ if progress is not None:
+ progress(len(t_frames), len(t_frames))
+ log.info("[scribble-cnn] base=%d levels=%d trained on %d px in %d crops,"
+ " %d classes: acc %.3f (crops %.2f s, fit %.2f s, %s)",
+ self.base, self.levels, self.report["n_pixels"], n_crops,
+ n_out, acc, self.report["crops_s"], t_fit, self.device)
+ return self.report
+
+ # -- prediction ----------------------------------------------------------
+
+ def predict_class_proba(self, frame) -> np.ndarray:
+ """``(K, H, W)`` float32 softmax over the trained classes.
+
+ Column *k* is ``self.classes[k]``, and non-finite source pixels get
+ probability 0 in every class — the same contract, verbatim, as
+ :meth:`spyde.particles.scribble.ScribbleClassifier.predict_class_proba`,
+ so a caller cannot tell the two engines apart by their output shape or
+ their NaN handling.
+ """
+ self._require_trained()
+ torch = import_torch()
+ prepared = prepare_frame(frame)
+ img = prepared.image
+ h, w = img.shape
+ k = len(self.classes)
+
+ with accelerator_lock(self.device):
+ out = np.empty((k, h, w), dtype=np.float32)
+ self._net.eval()
+ with torch.no_grad():
+ for (y0, y1, x0, x1, py0, py1, px0, px1) in self._tiles(h, w):
+ sub = img[y0:y1, x0:x1]
+ t = torch.as_tensor(sub, device=self.device,
+ dtype=torch.float32)[None, None]
+ t, (pad_b, pad_r) = _pad_to_multiple(torch, t, self.mult)
+ p = torch.softmax(self._net(t), dim=1)[0]
+ if pad_b or pad_r:
+ p = p[:, :y1 - y0, :x1 - x0]
+ out[:, py0:py1, px0:px1] = (
+ p[:, py0 - y0:py1 - y0, px0 - x0:px1 - x0]
+ .detach().cpu().numpy())
+
+ out[:, ~prepared.valid] = 0.0
+ return out
+
+ def _tiles(self, h: int, w: int):
+ """``(read window, write window)`` pairs for tiled inference.
+
+ Each tile is read with a halo of the receptive field and written without
+ it, so the tiled result matches an untiled one everywhere except where a
+ halo runs off the frame — which is where reflect padding would have
+ invented the same context anyway. A frame small enough to fit is one
+ tile with no halo at all.
+ """
+ edge = self.tile
+ if h * w <= TILE_ABOVE_PIXELS and h <= edge and w <= edge:
+ yield (0, h, 0, w, 0, h, 0, w)
+ return
+ halo = _halo(self.levels)
+ for py0 in range(0, h, edge):
+ py1 = min(h, py0 + edge)
+ y0, y1 = max(0, py0 - halo), min(h, py1 + halo)
+ for px0 in range(0, w, edge):
+ px1 = min(w, px0 + edge)
+ x0, x1 = max(0, px0 - halo), min(w, px1 + halo)
+ yield (y0, y1, x0, x1, py0, py1, px0, px1)
+
+ def predict_foreground_boundary(
+ self, frame
+ ) -> tuple[np.ndarray, np.ndarray | None]:
+ """``(foreground, boundary)`` from ONE pass over the frame.
+
+ *boundary* is None when no trained class is marked
+ :attr:`~spyde.particles.scribble.ScribbleClass.boundary` — None and not
+ a zero map, because :func:`spyde.particles.classical.split_instances`
+ reads the two differently ("no boundary taught, use the watershed"
+ versus "a boundary was taught and this frame has none of it").
+ """
+ self._require_trained()
+ proba = self.predict_class_proba(frame)
+ wanted = [k for k, c in enumerate(self.classes) if c.particle]
+ if not wanted:
+ raise RuntimeError(
+ "no trained class is marked as a particle class, so there is no "
+ "foreground to report — set ScribbleClass.particle on at least "
+ f"one of {[c.name for c in self.classes]} and retrain")
+ edge = [k for k, c in enumerate(self.classes) if c.boundary]
+ return (_sum_planes(proba, wanted),
+ _sum_planes(proba, edge) if edge else None)
+
+ def predict_proba(self, frame) -> np.ndarray:
+ """``(H, W)`` float32 foreground probability — the plan §0.2 spine's
+ input, interchangeable with the MLP engine's."""
+ return self.predict_foreground_boundary(frame)[0]
+
+ def predict_boundary_proba(self, frame) -> np.ndarray | None:
+ return self.predict_foreground_boundary(frame)[1]
+
+ def segment(self, frame, params=None) -> np.ndarray:
+ """Probability → labelled instances via the SHARED split.
+
+ Byte-for-byte the same forwarder the MLP engine has, including passing
+ the boundary down when one was taught. That is the design claim this
+ prototype exists to test: a CNN that predicts the boundary class plugs
+ into machinery that already exists.
+ """
+ from spyde.particles.classical import SegmentParams, split_instances
+ fg, bnd = self.predict_foreground_boundary(frame)
+ return split_instances(fg, params or SegmentParams(), boundary=bnd)
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+def _halo(levels: int) -> int:
+ """Tile overlap, px: roughly the net's receptive field, rounded up.
+
+ Two 3×3 convolutions per level, doubling in stride each level down, plus the
+ decoder's mirror — ~30 px at ``levels=2`` and ~120 px at ``levels=3``. Padded
+ to a comfortable multiple rather than derived exactly: a halo that is too
+ small shows as seams, a halo that is too large costs a few percent.
+ """
+ return int(32 * (1 << max(0, int(levels) - 2)))
+
+
+def _pad_to_multiple(torch, t, mult: int):
+ """Reflect-pad an ``(N, C, h, w)`` tensor's bottom/right up to *mult*.
+
+ Bottom/right only, so the output's ``[:h, :w]`` is the input's pixels at the
+ same indices — a symmetric pad would shift every coordinate and silently
+ move the whole prediction by a pixel.
+ """
+ import torch.nn.functional as F
+ h, w = int(t.shape[-2]), int(t.shape[-1])
+ pad_b, pad_r = (-h) % mult, (-w) % mult
+ if pad_b or pad_r:
+ t = F.pad(t, (0, pad_r, 0, pad_b), mode="reflect")
+ return t, (pad_b, pad_r)
+
+
+def _augment(torch, xb, yb, rng):
+ """Random dihedral transform of a crop batch (whole batch at once).
+
+ Per-batch rather than per-crop: it is one indexing op instead of eight, the
+ batch is re-sampled every step anyway so a crop still sees every transform
+ over a run, and at this step count the difference is unmeasurable.
+ """
+ if rng.random() < 0.5:
+ xb, yb = torch.flip(xb, (-1,)), torch.flip(yb, (-1,))
+ if rng.random() < 0.5:
+ xb, yb = torch.flip(xb, (-2,)), torch.flip(yb, (-2,))
+ if xb.shape[-1] == xb.shape[-2] and rng.random() < 0.5:
+ xb, yb = xb.transpose(-1, -2), yb.transpose(-1, -2)
+ return xb.contiguous(), yb.contiguous()
+
+
+def _sum_planes(proba: np.ndarray, cols: list[int]) -> np.ndarray:
+ """Sum the given planes of a ``(K, H, W)`` softmax, float32.
+
+ The one-column case is a plain slice for the reason
+ :func:`spyde.particles.scribble._sum_planes` documents: fancy-indexing a
+ 4096² map builds a copy before reducing it, measured at 113 ms for the pair.
+ """
+ if len(cols) == 1:
+ plane = proba[cols[0]]
+ return plane if plane.dtype == np.float32 else plane.astype(np.float32)
+ return proba[cols].sum(axis=0).astype(np.float32)
+
+
+_AUTOGRAD_WARMED = False
+
+
+def _warmup_autograd(device) -> None:
+ """One trivial backward on the CALLING thread before the real one.
+
+ torch's CUDA autograd backward segfaults on Windows the first time it runs
+ on a thread whose engine has not been initialised. Same mitigation, for the
+ same reason, as :func:`spyde.actions.vector_orientation_gpu.warmup_autograd`
+ — duplicated rather than imported so this prototype does not reach into the
+ orientation-mapping package.
+ """
+ global _AUTOGRAD_WARMED
+ if _AUTOGRAD_WARMED or getattr(device, "type", str(device)) != "cuda":
+ return
+ try:
+ torch = import_torch()
+ x = torch.zeros(1, device=device, requires_grad=True)
+ (x * 2).sum().backward()
+ _AUTOGRAD_WARMED = True
+ except Exception as e: # pragma: no cover
+ log.debug("CUDA autograd warmup skipped: %s", e)
diff --git a/spyde/particles/track.py b/spyde/particles/track.py
new file mode 100644
index 00000000..7a756c4d
--- /dev/null
+++ b/spyde/particles/track.py
@@ -0,0 +1,894 @@
+"""
+track.py — frame-to-frame linking and the event stream. Plan steps C1 and C2.
+
+One pass over the frames, one `linear_sum_assignment` per frame pair, and the
+assignment's **leftovers are the physics**: a detection nothing was assigned to is
+a birth, a track nothing was assigned from is a death, and the two-body cases
+(merge, split) are read off those leftovers by a post-pass.
+
+Why `scipy.optimize.linear_sum_assignment` and not trackpy
+----------------------------------------------------------
+trackpy's model (centroid distance, gated by a search radius, optional memory) is
+the right one and is what this implements — but scipy is already a core dependency,
+and the whole linker plus the event post-pass is under 400 lines of code. Adding a
+dependency to get a Hungarian solve we already ship is not a trade worth making.
+
+Units — read this before setting `max_dist`
+-------------------------------------------
+:mod:`spyde.particles.measure` writes centroids in **calibrated units** (pixels x
+``scale``), the same units as ``area``, ``equiv_diameter`` and everything else in
+the property row. So every distance here — :attr:`LinkParams.max_dist`,
+``merge_dist``, ``split_dist``, the trajectories on :class:`LinkResult` — is in
+``particles.units``, never pixels. A :class:`~spyde.drift.model.DriftModel`,
+however, is in **pixels** (it is a shift applied to an image), so
+:func:`sample_frame_positions` converts to pixels, asks the model, and converts
+back. That conversion is in exactly one place for the same reason ``measure.py``
+calibrates in exactly one place.
+
+Why the gate is a maximum-cardinality trick rather than a big square matrix
+---------------------------------------------------------------------------
+The textbook formulation (Jaqaman 2008) pads the cost matrix with dummy
+rows/columns priced at the gate, giving a ``(n_t + n_d)`` square problem in which
+an unmatched track and an unmatched detection are both first-class outcomes. That
+is what makes "a detection may go unmatched rather than be forced into a bad pair"
+true, and it is the property that matters.
+
+The rectangular matrix used here has the same optimum for strictly less work.
+Infeasible pairs (distance above the gate) get a constant sentinel large enough
+that the solver minimises the *number* of them before it looks at any real cost;
+those pairs are then discarded. Both formulations therefore compute a maximum-
+cardinality matching over the feasible pairs and, among those, the minimum total
+cost — the padded version can never prefer two dummies to a feasible pair, because
+a feasible pair costs less than the gate while two dummies cost twice it. The
+rectangular problem is ``n_t x n_d`` instead of ``(n_t + n_d)^2``, i.e. 4x smaller
+at 500 particles per frame, and `linear_sum_assignment` is superlinear in the
+matrix size.
+
+Cost, and what the gate applies to
+----------------------------------
+Cost is centroid distance, optionally plus a property-similarity penalty
+(:attr:`LinkParams.property_weight`, **off by default**). The **gate is on
+distance alone** — the penalty only re-orders pairs that are already admissible.
+Letting the penalty push a pair over the gate would silently turn ``max_dist``
+into "max_dist minus however dissimilar these two happen to look", which is not
+what a user setting a search radius means.
+
+The penalty is off by default because on real data a particle's measured area
+fluctuates frame to frame by far more than its centroid moves — the area of a
+threshold-defined region swings with noise, the centroid barely does — so
+weighting area *adds* noise to a cost that was already the reliable signal. It
+earns its keep only when positions are genuinely ambiguous (dense fields, fast
+motion), which is why it is a knob and not a constant.
+
+Cost at scale
+-------------
+Per frame this is ``O(n_t * n_d)`` to build the matrix and up to ``O(n^3)`` for the
+assignment, with *n* the particles in ONE frame — never in the movie. Measured on
+this box, 50 frames of synthetic tracks (link only, no segmentation):
+
+=================== ============ ========================
+particles per frame ms per frame extrapolated 3000 frames
+=================== ============ ========================
+50 0.23 0.7 s
+100 0.47 1.4 s
+200 1.19 3.6 s
+500 9.5 28.6 s
+800 23.8 71 s
+=================== ============ ========================
+
+So the plan's target (3000 frames, ~500 particles each) is **~29 s**, negligible
+beside segmenting 3000 frames of 2048^2 — but the growth is clearly superlinear
+(~n^2.2 over that range), so a field of several thousand particles per frame would
+need the cost matrix restricted to near neighbours (a KD-tree query inside
+``max_dist``, then a sparse assignment) rather than built dense. That is not done
+here: it is unnecessary at the stated scale and would add a second code path with no
+data to validate it against. For reference the whole 24-frame fixture links in
+1.8 ms (74 us/frame at ~6 particles), against 486 ms to segment and measure it.
+"""
+from __future__ import annotations
+
+import logging
+from dataclasses import dataclass, field
+from typing import Any, Sequence
+
+import numpy as np
+
+from spyde.signals.particles import COL, SpyDEParticles
+
+log = logging.getLogger(__name__)
+
+#: The four event types. Order is the navigator lane order (plan C2) — green
+#: birth, red death, mauve merge, yellow split.
+EVENT_KINDS: tuple[str, ...] = ("birth", "death", "merge", "split")
+
+#: Properties the optional similarity term compares. Both are in the property row
+#: for every engine; ``intensity_mean`` is NaN when ``measure_frame`` ran without
+#: an intensity image, which is handled as "unknown", not as "very different".
+DEFAULT_SIMILARITY_PROPERTIES: tuple[str, ...] = ("area", "intensity_mean")
+
+# Floor for the ADAPTIVE merge/split radius, in pixels (scaled to calibrated units
+# at use). The adaptive radius is the particle's own equivalent diameter, which is
+# the right scale — a big particle absorbs a neighbour from further away than a
+# small one does — but a 1-2 px detection has an equivalent diameter near zero,
+# and a zero radius means its merges are never detected. Two pixels is the
+# smallest radius at which "these two detections became one" is even meaningful.
+_ADAPTIVE_RADIUS_FLOOR_PX = 2.0
+
+
+# ── events ───────────────────────────────────────────────────────────────────
+
+@dataclass(frozen=True)
+class ParticleEvent:
+ """One thing that happened to a track, at a frame.
+
+ Frozen because the event stream is handed to the navigator lane, the Events
+ table and the report embed; a record three surfaces share must not be
+ mutable in one of them.
+
+ Parameters
+ ----------
+ frame
+ The frame the event is **observed at**. For a birth that is the track's
+ first detected frame. For a death it is ``last_detected + 1`` — the first
+ frame the particle is *gone*, which is the convention the synthetic
+ fixture's ``death`` column uses and the frame a user would point at and
+ call the dissolution.
+ kind
+ One of :data:`EVENT_KINDS`.
+ tracks
+ The track ids involved. Birth/death: one. Merge: ``(absorbed, survivor)``.
+ Split: ``(parent, fragment)``.
+ particles
+ Global particle indices into ``SpyDEParticles.flat_buffer``, in the same
+ order as *tracks* — so the Events table can jump straight to a row and the
+ overlay can highlight the exact detections.
+ """
+
+ frame: int
+ kind: str
+ tracks: tuple[int, ...]
+ particles: tuple[int, ...]
+
+ def to_dict(self) -> dict[str, Any]:
+ """Plain JSON-safe dict — this is what crosses the IPC boundary."""
+ return {
+ "frame": int(self.frame),
+ "kind": str(self.kind),
+ "tracks": [int(t) for t in self.tracks],
+ "particles": [int(i) for i in self.particles],
+ }
+
+ @classmethod
+ def from_dict(cls, d: dict[str, Any]) -> "ParticleEvent":
+ return cls(
+ frame=int(d["frame"]),
+ kind=str(d["kind"]),
+ tracks=tuple(int(t) for t in d.get("tracks", ())),
+ particles=tuple(int(i) for i in d.get("particles", ())),
+ )
+
+
+def events_to_records(events: Sequence[ParticleEvent]) -> list[dict[str, Any]]:
+ """The event stream as JSON-safe records, ready for IPC or CSV."""
+ return [e.to_dict() for e in events]
+
+
+def events_from_records(records: Sequence[dict[str, Any]]) -> list[ParticleEvent]:
+ """Inverse of :func:`events_to_records`."""
+ return [ParticleEvent.from_dict(r) for r in records]
+
+
+def event_counts(events: Sequence[ParticleEvent], n_frames: int) -> dict[str, np.ndarray]:
+ """``{kind: (n_frames,) float32}`` counts — the navigator's event lane (C2).
+
+ A dict of separate traces rather than one stacked array because each kind gets
+ its own colour and its own row in the lane; a caller that wants the total sums
+ them.
+ """
+ out = {k: np.zeros(int(n_frames), dtype=np.float32) for k in EVENT_KINDS}
+ for e in events:
+ if 0 <= e.frame < int(n_frames) and e.kind in out:
+ out[e.kind][e.frame] += 1.0
+ return out
+
+
+# ── parameters ───────────────────────────────────────────────────────────────
+
+@dataclass
+class LinkParams:
+ """Linker parameters. Distances are in the particles' calibrated units.
+
+ Parameters
+ ----------
+ max_dist
+ Search radius: a track and a detection further apart than this are never
+ linked. **In ``particles.units``**, not pixels (see the module docstring).
+ The default is deliberately generous — on the synthetic fixture
+ (``scale=0.5`` nm/px) the fastest particle moves 2.2 px = 1.1 nm per frame,
+ so 10.0 is ~9x the true step. The asymmetry justifies it: a gate that is
+ too tight **fragments** a track, and a fragmented trajectory cannot be
+ repaired downstream, while a gate that is too loose only matters when two
+ particles are within it of each other — and there the assignment still
+ picks the globally cheapest pairing. Measured on the fixture, every gate
+ from 4 px to 80 px recovers exactly the same 7 tracks and the same events;
+ at 2 px the tracks fragment (8 tracks, 3 spurious deaths) and at 1.2 px it
+ collapses (26 tracks, 19 deaths, 4 phantom splits). The usable window is
+ wide upward and sharp downward, so err high. Dense fields should tighten it.
+ memory
+ Frames a track may go undetected and still be re-linked afterwards. This is
+ what makes a blinking detection **one** track instead of several. 0 means a
+ track ends the moment it is missed. Note the gate is measured from the
+ track's last *seen* position and is not widened by the gap, so a fast
+ particle that blinks may still fall outside it.
+ property_weight
+ Weight of the property-similarity penalty, expressed as a multiple of
+ *max_dist* so it is scale-free. 0 (default) is distance only; see the
+ module docstring for why that is the default.
+ properties
+ Which columns the similarity term compares.
+ merge_dist, split_dist
+ Radius for the merge / split post-pass. ``None`` (default) is **adaptive**:
+ the particle's own ``equiv_diameter``, floored at 2 px. Adaptive is right
+ because the signature of a merge is a centroid jumping to the join of two
+ bodies, and that jump is a fraction of the body size — a fixed radius that
+ works for 50 px particles invents merges among 5 px ones.
+ initial_births
+ Whether detections in frame 0 emit ``birth`` events. Default True: it makes
+ the event stream a *complete* description of the assignment (every track has
+ exactly one birth, so ``len(births) == n_tracks``), and a consumer that
+ wants only nucleations filters ``frame > 0`` — see
+ :meth:`LinkResult.events_of`. The reverse is not recoverable: drop them and
+ nothing downstream can tell which tracks were present from the start.
+ """
+
+ max_dist: float = 10.0
+ memory: int = 0
+ property_weight: float = 0.0
+ properties: tuple[str, ...] = DEFAULT_SIMILARITY_PROPERTIES
+ merge_dist: float | None = None
+ split_dist: float | None = None
+ initial_births: bool = True
+
+ def __post_init__(self) -> None:
+ if not np.isfinite(self.max_dist) or self.max_dist <= 0:
+ raise ValueError(f"max_dist must be finite and > 0; got {self.max_dist}")
+ if self.memory < 0:
+ raise ValueError(f"memory must be >= 0; got {self.memory}")
+ if self.property_weight < 0:
+ raise ValueError(
+ f"property_weight must be >= 0; got {self.property_weight}")
+ for name in self.properties:
+ if name not in COL:
+ raise KeyError(f"unknown property column {name!r} in properties")
+ for name in ("merge_dist", "split_dist"):
+ v = getattr(self, name)
+ if v is not None and (not np.isfinite(v) or v <= 0):
+ raise ValueError(f"{name} must be None or finite and > 0; got {v}")
+
+ def to_dict(self) -> dict[str, Any]:
+ return {
+ "max_dist": float(self.max_dist),
+ "memory": int(self.memory),
+ "property_weight": float(self.property_weight),
+ "properties": list(self.properties),
+ "merge_dist": None if self.merge_dist is None else float(self.merge_dist),
+ "split_dist": None if self.split_dist is None else float(self.split_dist),
+ "initial_births": bool(self.initial_births),
+ }
+
+
+# ── result ───────────────────────────────────────────────────────────────────
+
+@dataclass
+class LinkResult:
+ """Track ids, the event stream, and the positions the linking actually used.
+
+ ``track_id`` is returned as a **parallel array** rather than written straight
+ into the buffer, so a link can be inspected, compared against another
+ parameter choice, or thrown away without having mutated the particle table.
+ :meth:`apply` performs the write when you want it.
+
+ Parameters
+ ----------
+ track_id
+ ``(n_particles,)`` int32, contiguous from 0. Every particle belongs to
+ exactly one track, so there is no -1 in a completed link.
+ events
+ Chronological (by ``frame``, then by track id).
+ positions
+ ``(n_particles, 2)`` float64 ``(y, x)`` in **calibrated units**, in
+ :attr:`reference`'s frame — i.e. what the cost matrix saw. Kept because a
+ trajectory read from the buffer would be in the lab frame even when the
+ link ran drift-corrected, and silently mixing the two is the bug this
+ field exists to prevent.
+ frame_index
+ ``(n_particles,)`` int32 frame of each particle, taken from the CSR row
+ pointers (the authoritative frame index) rather than the float ``t``
+ column.
+ reference
+ ``"lab"`` or ``"sample"`` — which frame of reference :attr:`positions`
+ and every trajectory are in.
+ track_first_frame, track_last_frame, track_first_index, track_last_index
+ ``(n_tracks,)`` per-track endpoints. Computed for the event pass and
+ exposed because the kymograph, the table dock and the trails overlay all
+ want them and re-deriving them is a scan of the whole buffer.
+ """
+
+ track_id: np.ndarray
+ events: list[ParticleEvent]
+ positions: np.ndarray
+ frame_index: np.ndarray
+ reference: str
+ track_first_frame: np.ndarray
+ track_last_frame: np.ndarray
+ track_first_index: np.ndarray
+ track_last_index: np.ndarray
+ n_frames: int = 0
+ params: dict[str, Any] = field(default_factory=dict)
+
+ def __post_init__(self) -> None:
+ self.track_id = np.ascontiguousarray(self.track_id, dtype=np.int32)
+ self.positions = np.ascontiguousarray(self.positions, dtype=np.float64)
+ self.frame_index = np.ascontiguousarray(self.frame_index, dtype=np.int32)
+ # Chronological order WITHIN each track, in one stable sort: the buffer is
+ # sorted by frame, so a stable sort on track id leaves each track's rows in
+ # increasing frame order. This is what makes `trajectory` a slice.
+ self._order = np.argsort(self.track_id, kind="stable")
+ counts = np.bincount(self.track_id, minlength=self.n_tracks) \
+ if self.track_id.size else np.zeros(0, np.int64)
+ self._starts = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64)
+
+ @property
+ def n_tracks(self) -> int:
+ return int(self.track_first_frame.size)
+
+ @property
+ def n_particles(self) -> int:
+ return int(self.track_id.size)
+
+ # ── writing back ─────────────────────────────────────────────────────────
+
+ def apply(self, particles: SpyDEParticles) -> SpyDEParticles:
+ """Write :attr:`track_id` into the buffer's ``track_id`` column, in place.
+
+ The column is float32, which represents integers exactly up to 2**24 — 16.7M
+ tracks, against a target scale of 1.5M particles (plan §0.5), so an id can
+ never be rounded to a neighbouring track's.
+ """
+ if len(particles.flat_buffer) != self.track_id.size:
+ raise ValueError(
+ f"this result is for {self.track_id.size} particles but the table "
+ f"has {len(particles.flat_buffer)} — it came from a different link"
+ )
+ particles.flat_buffer[:, COL["track_id"]] = self.track_id.astype(np.float32)
+ return particles
+
+ # ── per-track access ─────────────────────────────────────────────────────
+
+ def track_indices(self, track: int) -> np.ndarray:
+ """Global particle indices of one track, in chronological order."""
+ t = int(track)
+ if not 0 <= t < self.n_tracks:
+ raise IndexError(f"track {t} outside 0..{self.n_tracks - 1}")
+ return self._order[self._starts[t]:self._starts[t + 1]]
+
+ def trajectory(self, track: int) -> np.ndarray:
+ """``(k, 3)`` ``[frame, y, x]`` for one track, in :attr:`reference`'s frame.
+
+ Positions are in calibrated units. Frames are not necessarily contiguous —
+ with ``memory > 0`` a track can skip frames, and the gap is visible as a
+ jump in the first column rather than being interpolated over. Inventing a
+ position for a frame the particle was not detected in would put a
+ measurement in the table that nothing measured.
+ """
+ idx = self.track_indices(track)
+ out = np.empty((idx.size, 3), dtype=np.float64)
+ out[:, 0] = self.frame_index[idx]
+ out[:, 1:] = self.positions[idx]
+ return out
+
+ def track_lengths(self) -> np.ndarray:
+ """``(n_tracks,)`` number of frames each track was DETECTED in.
+
+ Not ``last - first + 1``: with ``memory > 0`` those differ, and the
+ difference is exactly the QC signal (a track detected in 4 of 30 frames is
+ probably noise) that plan C3's lifetime sort is for.
+ """
+ return np.diff(self._starts).astype(np.int64)
+
+ def track_at(self, frame: int) -> np.ndarray:
+ """Track ids present in *frame*, in buffer order.
+
+ ``searchsorted``, not a boolean scan: the overlay calls this on every
+ navigator move, and ``frame_index`` is non-decreasing (the buffer is sorted
+ by frame), so an O(log N) slice is available where an O(N) comparison over
+ 1.5M rows would otherwise run per frame.
+ """
+ f = int(frame)
+ lo, hi = np.searchsorted(self.frame_index, [f, f + 1])
+ return self.track_id[lo:hi]
+
+ # ── events ───────────────────────────────────────────────────────────────
+
+ def events_of(self, kind: str | None = None, *,
+ exclude_initial: bool = False) -> list[ParticleEvent]:
+ """Events of one *kind*, optionally dropping frame-0 births.
+
+ *exclude_initial* is the "only real nucleations" filter — see
+ :attr:`LinkParams.initial_births` for why frame-0 births are recorded at
+ all.
+ """
+ if kind is not None and kind not in EVENT_KINDS:
+ raise ValueError(
+ f"unknown event kind {kind!r}; expected one of {', '.join(EVENT_KINDS)}")
+ out = [e for e in self.events if kind is None or e.kind == kind]
+ if exclude_initial:
+ out = [e for e in out if not (e.kind == "birth" and e.frame == 0)]
+ return out
+
+ def event_counts(self) -> dict[str, np.ndarray]:
+ """Per-frame counts per kind — see :func:`event_counts`."""
+ return event_counts(self.events, self.n_frames)
+
+ def to_dict(self) -> dict[str, Any]:
+ """Serialisable summary. Arrays stay out — save those with the particles."""
+ return {
+ "n_tracks": self.n_tracks,
+ "n_particles": self.n_particles,
+ "n_frames": int(self.n_frames),
+ "reference": self.reference,
+ "params": dict(self.params),
+ "events": events_to_records(self.events),
+ }
+
+ def __repr__(self) -> str:
+ counts = {k: len(self.events_of(k)) for k in EVENT_KINDS}
+ return (
+ f"LinkResult({self.n_tracks} tracks over {self.n_frames} frames, "
+ f"reference={self.reference!r}, events={counts})"
+ )
+
+
+# ── coordinates ──────────────────────────────────────────────────────────────
+
+def frame_indices(particles: SpyDEParticles) -> np.ndarray:
+ """``(n_particles,)`` int32 frame index per particle, from the CSR pointers.
+
+ The ``t`` column carries the same number, but the row pointers are the
+ *definition* of which frame a row belongs to — and they are integers, whereas
+ ``t`` is a float32 that a caller could in principle have written a
+ non-integral value into.
+ """
+ counts = np.diff(particles.t_offsets)
+ return np.repeat(np.arange(particles.n_frames, dtype=np.int32),
+ counts).astype(np.int32)
+
+
+def sample_frame_positions(particles: SpyDEParticles, drift) -> np.ndarray:
+ """Lab-frame centroids mapped into the drift-corrected (sample) frame.
+
+ ``(n_particles, 2)`` float64 ``(y, x)``, **still in calibrated units** — the
+ return value is directly comparable with the stored centroids.
+
+ The conversion to and from pixels is what this function exists for.
+ :class:`~spyde.drift.model.DriftModel` is defined on image pixels (its shifts
+ are what you pass to ``scipy.ndimage.shift``) while the property row is
+ calibrated, so a caller adding ``model.shifts[t]`` straight onto ``y``/``x`` is
+ off by a factor of ``scale`` — an over- or under-correction that still yields a
+ smooth, visibly-improved trajectory and so does not announce itself. Only the
+ ``scale == 1`` case, where nobody looks, comes out right.
+
+ Going through ``to_sample_frame`` rather than adding the shifts here also means
+ a future non-rigid model, whose mapping is not a plain addition, needs no change
+ on this side.
+ """
+ if drift.n_frames < particles.n_frames:
+ raise ValueError(
+ f"drift model covers {drift.n_frames} frames but the particles span "
+ f"{particles.n_frames} — a shorter model would silently reuse the "
+ "last shift for every frame beyond it"
+ )
+ scale = float(particles.scale) or 1.0
+ pos_px = particles.flat_buffer[:, [COL["y"], COL["x"]]].astype(np.float64) / scale
+ idx = frame_indices(particles).astype(np.intp)
+ if pos_px.size == 0:
+ return pos_px
+ return np.asarray(drift.to_sample_frame(pos_px, idx), dtype=np.float64) * scale
+
+
+# ── the linker ───────────────────────────────────────────────────────────────
+
+def link(
+ particles: SpyDEParticles,
+ params: LinkParams | None = None,
+ *,
+ drift=None,
+ apply: bool = False,
+ **kwargs: Any,
+) -> LinkResult:
+ """Link detections into tracks and extract the event stream.
+
+ Parameters
+ ----------
+ particles
+ The CSR table from segment + measure. Not modified unless *apply*.
+ params
+ :class:`LinkParams`. Individual fields may instead be passed as keyword
+ arguments (``link(p, max_dist=4, memory=1)``), which is what the wizard's
+ parameter dict does.
+ drift
+ Optional :class:`~spyde.drift.model.DriftModel`. When given, linking and
+ every reported trajectory run in the **sample** frame — the stage's motion
+ removed, so a static particle's trajectory is a point. When ``None``, the
+ lab frame, i.e. the raw centroids as measured. Both are valid answers to
+ different questions ("did the particle move, or did the stage?", plan A9),
+ which is why the choice is the caller's and is recorded on the result.
+ apply
+ Write the ids into ``particles.flat_buffer``'s ``track_id`` column as well
+ as returning them. Off by default so a link is a pure computation.
+
+ Returns
+ -------
+ LinkResult
+
+ Notes
+ -----
+ Cost per frame is ``O(n_t * n_d)`` to build the matrix and up to
+ ``O(n^3)`` for the assignment, with *n* the number of particles in ONE frame —
+ never the whole movie. Nothing here reads pixel data at all.
+ """
+ if kwargs:
+ base = params.to_dict() if params is not None else LinkParams().to_dict()
+ unknown = set(kwargs) - set(base)
+ if unknown:
+ raise TypeError(
+ f"unknown link parameter(s): {', '.join(sorted(unknown))}")
+ base.update(kwargs)
+ base["properties"] = tuple(base["properties"])
+ params = LinkParams(**base)
+ p = params or LinkParams()
+
+ n_frames = particles.n_frames
+ n_particles = particles.n_particles
+ fidx = frame_indices(particles)
+
+ if drift is not None:
+ positions = sample_frame_positions(particles, drift)
+ reference = "sample"
+ else:
+ positions = particles.flat_buffer[:, [COL["y"], COL["x"]]].astype(np.float64)
+ reference = "lab"
+
+ track_id = _assign_tracks(particles, positions, p)
+
+ n_tracks = int(track_id.max()) + 1 if track_id.size else 0
+ if n_tracks:
+ # `track_id` is contiguous from 0 and the buffer is sorted by frame, so a
+ # STABLE sort groups each track's rows in chronological order — the endpoints
+ # are then the group's first and last element, with no per-track scan.
+ order = np.argsort(track_id, kind="stable")
+ starts = np.concatenate(
+ [[0], np.cumsum(np.bincount(track_id, minlength=n_tracks))])
+ first_index = order[starts[:-1]].astype(np.int64)
+ last_index = order[starts[1:] - 1].astype(np.int64)
+ first_frame = fidx[first_index].astype(np.int32)
+ last_frame = fidx[last_index].astype(np.int32)
+ else:
+ first_index = last_index = np.zeros(0, np.int64)
+ first_frame = last_frame = np.zeros(0, np.int32)
+
+ events = _extract_events(
+ particles, positions, track_id, n_frames,
+ first_frame=first_frame, last_frame=last_frame,
+ first_index=first_index, last_index=last_index, p=p,
+ )
+
+ result = LinkResult(
+ track_id=track_id,
+ events=events,
+ positions=positions,
+ frame_index=fidx,
+ reference=reference,
+ track_first_frame=first_frame,
+ track_last_frame=last_frame,
+ track_first_index=first_index,
+ track_last_index=last_index,
+ n_frames=n_frames,
+ params=p.to_dict(),
+ )
+ log.info("[track] %d particles over %d frames -> %d tracks (%s frame), "
+ "%d events", n_particles, n_frames, result.n_tracks, reference,
+ len(events))
+ if apply:
+ result.apply(particles)
+ return result
+
+
+def _assign_tracks(particles: SpyDEParticles, positions: np.ndarray,
+ p: LinkParams) -> np.ndarray:
+ """The forward pass: one assignment per frame pair. Returns ``(N,)`` int32 ids.
+
+ Ids are handed out in (frame, row-within-frame) order, which is why they are
+ stable: the same table linked twice produces the same numbering, and there is
+ no dictionary iteration or set ordering anywhere in the loop.
+ """
+ from scipy.optimize import linear_sum_assignment
+
+ n_particles = particles.n_particles
+ track_id = np.full(n_particles, -1, dtype=np.int32)
+ if n_particles == 0:
+ return track_id
+
+ prop = _similarity_matrix_source(particles, p)
+
+ # Live-track registry. Lists, not arrays: they are appended to once per new
+ # track and read once per frame, and the per-frame cost is dominated by the
+ # assignment. `active` holds only tracks still inside the memory window.
+ last_frame: list[int] = []
+ last_index: list[int] = []
+ active: list[int] = []
+
+ def start_track(gi: int, t: int) -> None:
+ tid = len(last_frame)
+ last_frame.append(t)
+ last_index.append(gi)
+ track_id[gi] = tid
+ active.append(tid)
+
+ for t in range(particles.n_frames):
+ det = particles.indices_at(t)
+
+ if t > 0 and active:
+ # Retire anything outside the memory window. `memory + 1` because a
+ # track last seen at t-1 has a gap of 1 and must always be eligible.
+ active[:] = [tid for tid in active
+ if t - last_frame[tid] <= p.memory + 1]
+
+ if det.size and active:
+ src = np.fromiter((last_index[tid] for tid in active),
+ dtype=np.int64, count=len(active))
+ cost, feasible = _cost_matrix(positions, src, det, prop, p)
+ rows, cols = linear_sum_assignment(cost)
+ for r, c in zip(rows, cols):
+ if not feasible[r, c]:
+ continue # a sentinel pair — see the module docstring
+ tid = active[r]
+ gi = int(det[c])
+ track_id[gi] = tid
+ last_frame[tid] = t
+ last_index[tid] = gi
+
+ for gi in det:
+ if track_id[gi] < 0:
+ start_track(int(gi), t)
+
+ return track_id
+
+
+def _similarity_matrix_source(particles: SpyDEParticles,
+ p: LinkParams) -> np.ndarray | None:
+ """``(N, k)`` float64 of the compared properties, or None when weight is 0."""
+ if p.property_weight <= 0 or not p.properties:
+ return None
+ cols = [COL[name] for name in p.properties]
+ return particles.flat_buffer[:, cols].astype(np.float64)
+
+
+def _cost_matrix(positions: np.ndarray, src: np.ndarray, det: np.ndarray,
+ prop: np.ndarray | None,
+ p: LinkParams) -> tuple[np.ndarray, np.ndarray]:
+ """``(cost, feasible)`` for one frame pair.
+
+ Infeasible entries get a constant sentinel chosen so the solver minimises
+ their *count* before it looks at any real cost — which is what makes discarding
+ them afterwards equivalent to the padded square formulation (module docstring).
+ The sentinel is derived from the actual matrix rather than being a magic
+ literal, so it cannot be outgrown by a large ``max_dist`` or a big property
+ penalty.
+ """
+ a = positions[src] # (n_t, 2)
+ b = positions[det] # (n_d, 2)
+ d = np.hypot(a[:, 0, None] - b[None, :, 0],
+ a[:, 1, None] - b[None, :, 1])
+
+ # The gate is on DISTANCE ONLY. See the module docstring.
+ feasible = d <= float(p.max_dist)
+
+ cost = d.copy()
+ if prop is not None:
+ cost += float(p.property_weight) * float(p.max_dist) * \
+ _property_dissimilarity(prop[src], prop[det])
+
+ if feasible.any():
+ big = (min(len(src), len(det)) + 1) * float(cost[feasible].max()) + 1.0
+ else:
+ big = 1.0
+ cost = np.where(feasible, cost, big)
+ return cost, feasible
+
+
+def _property_dissimilarity(a: np.ndarray, b: np.ndarray) -> np.ndarray:
+ """``(n_t, n_d)`` mean symmetric relative difference in 0..1.
+
+ ``|a - b| / (|a| + |b|)`` rather than ``|a - b| / a``: it is bounded, symmetric,
+ and needs no reference value, so one expression works for area (thousands of
+ nm^2) and for a normalised intensity (fractions of one) without per-column
+ scaling.
+
+ A NaN on either side is **dropped from the mean**, not counted as zero.
+ ``intensity_mean`` is NaN whenever ``measure_frame`` ran without an intensity
+ image, and both alternatives are wrong in a way that is hard to see from the
+ outside: counting it as maximally different makes the term a uniform offset on
+ every pair, and counting it as zero *dilutes* the columns that were measured —
+ with one of two properties missing, a requested weight of 1.0 would silently
+ act as 0.5. Averaging over the known columns keeps the weight meaning what the
+ caller asked for.
+ """
+ valid = np.isfinite(a)[:, None, :] & np.isfinite(b)[None, :, :]
+ num = np.abs(a[:, None, :] - b[None, :, :])
+ den = np.abs(a[:, None, :]) + np.abs(b[None, :, :])
+ with np.errstate(divide="ignore", invalid="ignore"):
+ # den == 0 means both values are 0, i.e. genuinely identical.
+ rel = np.where(den > 0, num / den, 0.0)
+ rel = np.where(valid & np.isfinite(rel), rel, 0.0)
+ n = valid.sum(axis=2)
+ return np.where(n > 0, rel.sum(axis=2) / np.maximum(n, 1), 0.0)
+
+
+# ── the event post-pass ──────────────────────────────────────────────────────
+
+def _extract_events(particles, positions, track_id, n_frames, *,
+ first_frame, last_frame, first_index, last_index,
+ p: LinkParams) -> list[ParticleEvent]:
+ """Turn the assignment's leftovers into the event stream.
+
+ Every track contributes at most two events: one for its start (``birth`` or
+ ``split``) and one for its end (``death`` or ``merge``). Merge and split
+ **replace** the death and birth they explain rather than accompanying them —
+ a particle that was absorbed did not dissolve, and reporting both would put a
+ red dissolution flag on the navigator lane at a frame where nothing dissolved.
+ The cost of that choice is that the event stream no longer closes the count
+ arithmetic by itself; the accounting is
+ ``count(t) - count(t-1) = births + splits - deaths - merges``.
+
+ A track still detected in the FINAL frame gets no end event: the movie running
+ out is not a dissolution.
+
+ Walks FRAMES, not tracks, so only two ``{track: row}`` maps are ever resident.
+ A single ``frame -> {track: row}`` index over the whole movie would be 1.5M dict
+ entries at the target scale (plan §0.5) — hundreds of MB of Python objects to
+ answer a question that only ever spans two adjacent frames.
+ """
+ n_tracks = int(first_frame.size)
+ if n_tracks == 0 or n_frames == 0:
+ return []
+
+ starts_at: list[list[int]] = [[] for _ in range(n_frames)]
+ ends_at: list[list[int]] = [[] for _ in range(n_frames)]
+ for tid in range(n_tracks):
+ starts_at[int(first_frame[tid])].append(tid)
+ t1 = int(last_frame[tid])
+ if t1 < n_frames - 1:
+ # The event frame is the first frame the particle is GONE.
+ ends_at[t1 + 1].append(tid)
+
+ equiv = particles.flat_buffer[:, COL["equiv_diameter"]].astype(np.float64)
+ floor = _ADAPTIVE_RADIUS_FLOOR_PX * (float(particles.scale) or 1.0)
+
+ def radius(gi: int, override: float | None) -> float:
+ if override is not None:
+ return float(override)
+ v = equiv[gi]
+ # An unmeasured diameter (no masks, or a degenerate region) must not become
+ # a zero radius that silently disables the whole post-pass.
+ return max(float(v) if np.isfinite(v) else 0.0, floor)
+
+ events: list[ParticleEvent] = []
+ prev_map: dict[int, int] = {}
+
+ for t in range(n_frames):
+ # The maps are only ever read by a candidate event at t (which needs both t
+ # and t-1), so a frame with no candidate at t and none at t+1 needs no map at
+ # all. On a long stable movie that skips nearly every frame; when it does
+ # build one, the next iteration is guaranteed to want it.
+ wants = bool(ends_at[t]) or (t > 0 and bool(starts_at[t]))
+ nxt = t + 1
+ wants_next = nxt < n_frames and (bool(ends_at[nxt]) or bool(starts_at[nxt]))
+ cur_map = ({int(track_id[gi]): int(gi) for gi in particles.indices_at(t)}
+ if (wants or wants_next) else {})
+
+ for tid in starts_at[t]:
+ gi = int(first_index[tid])
+ if t == 0:
+ if p.initial_births:
+ events.append(ParticleEvent(0, "birth", (tid,), (gi,)))
+ continue
+ # SPLIT: a track that continues across t-1 -> t, whose position BEFORE
+ # the split is within its own body radius of this newcomer.
+ parent = _nearest_continuing(
+ cur_map, prev_map, positions, probe=positions[gi], exclude=tid,
+ measure_at=prev_map,
+ radius_of=lambda pgi: radius(pgi, p.split_dist),
+ )
+ if parent is None:
+ events.append(ParticleEvent(t, "birth", (tid,), (gi,)))
+ else:
+ ptid, pgi = parent
+ events.append(ParticleEvent(t, "split", (ptid, tid), (pgi, gi)))
+
+ for tid in ends_at[t]:
+ gi = int(last_index[tid])
+ merge_r = radius(gi, p.merge_dist)
+ # MERGE: a track that continues across t-1 -> t whose position AT t is
+ # within the dying particle's body radius of where the dying particle
+ # was at t-1 — i.e. a centroid that jumped to the join of two bodies.
+ survivor = _nearest_continuing(
+ cur_map, prev_map, positions, probe=positions[gi], exclude=tid,
+ measure_at=cur_map, radius_of=lambda _pgi, r=merge_r: r,
+ )
+ if survivor is None:
+ events.append(ParticleEvent(t, "death", (tid,), (gi,)))
+ else:
+ stid, sgi = survivor
+ events.append(ParticleEvent(t, "merge", (tid, stid), (gi, sgi)))
+
+ prev_map = cur_map
+
+ events.sort(key=lambda e: (e.frame, EVENT_KINDS.index(e.kind), e.tracks))
+ return events
+
+
+def _nearest_continuing(cur_map, prev_map, positions, *, probe, exclude,
+ measure_at, radius_of):
+ """Nearest track present in BOTH frames, measured on *measure_at*'s side.
+
+ This is the whole merge/split rule, and the "both frames" requirement is what
+ keeps it honest. A one-to-one assignment cannot express two-to-one, so the
+ two-body events have to be inferred from a died track sitting next to a
+ surviving one — and the only way to tell a *survivor* from another newcomer is
+ that the survivor was already there the frame before.
+
+ *measure_at* selects which side is compared with *probe*: ``cur_map`` for a
+ merge (where did the surviving centroid jump TO) and ``prev_map`` for a split
+ (where was the parent BEFORE it broke up). *radius_of* is passed the global row
+ index on the measured side, so a split can size its gate by the parent's own
+ body while a merge sizes it by the dying particle's.
+
+ Known limits, stated rather than hidden:
+
+ * **It cannot see a merge earlier than the segmenter does.** Measured on the
+ synthetic fixture, whose two converging discs geometrically touch at frame
+ 14: with watershed splitting ON they remain two separate detections until
+ **frame 18**, and with it OFF their thresholded blobs already connect at
+ **frame 12**. The linker reports the frame the detections became one, which
+ is a property of segmentation, not of linking, and no post-pass can recover
+ the other frame from the table alone.
+ * **A three-body coincidence is misread.** A track that genuinely dissolves
+ within one body-diameter of a surviving neighbour is reported as a merge.
+ * **A fragmentation in which the parent's OWN track ends is reported as two
+ births, not a split**, because then neither newcomer was present the frame
+ before. Requiring the parent to continue is what stops every ordinary birth
+ that happens to appear beside an unrelated particle from being called a
+ split; the trade is deliberate.
+ * **Merge and split are exclusive per track end/start.** Three tracks
+ collapsing into one in a single frame produce two merge events sharing one
+ survivor, which is the honest reading, but a track cannot be recorded as both
+ merging and splitting at the same frame.
+ * **Which track survives a SYMMETRIC merge is a genuine tie**, decided by the
+ assignment rather than by this rule. Two equal bodies merging put their joint
+ centroid the same distance from both, so either can be the survivor. Measured
+ on the fixture: linking in the lab frame reports ``tracks=(5, 4)`` and linking
+ the same table drift-corrected reports ``(4, 5)`` — same merge, same frame,
+ roles swapped, identical set of tracks and track lengths. Do not read meaning
+ into which id survives.
+ """
+ best = None
+ best_d = np.inf
+ for tid, gi in cur_map.items():
+ if tid == exclude or tid not in prev_map:
+ continue
+ mgi = int(measure_at[tid])
+ d = float(np.hypot(*(probe - positions[mgi])))
+ if d <= radius_of(mgi) and d < best_d:
+ best, best_d = (tid, gi), d
+ return best
diff --git a/spyde/signal_tree.py b/spyde/signal_tree.py
index fcb2d3c3..18f08265 100644
--- a/spyde/signal_tree.py
+++ b/spyde/signal_tree.py
@@ -22,6 +22,14 @@
logger = logging.getLogger(__name__)
+#: Chunks that must still be OUTSTANDING for the threaded navigator fill to hand
+#: over to the cluster mid-flight. Handing over recomputes the chunks already
+#: painted, so with only a handful left that costs more than it saves; with a
+#: movie's worth left it is the difference between seconds and minutes. The
+#: check runs per chunk, so in practice this fires within the first few.
+_NAV_HANDOVER_MIN_CHUNKS = 8
+
+
def _materialise_signal(signal: BaseSignal, array: np.ndarray) -> None:
"""Swap a LAZY signal's dask data for an in-RAM array, the way hyperspy's own
@@ -409,6 +417,39 @@ def _bg_nav(_dask=nav_dask, _plot=nav_plot, _sig=nav_signals,
for combo in itertools.product(*axes_ranges):
if _stop.is_set():
return
+ # HAND OVER to the cluster the moment it exists.
+ #
+ # This branch was chosen because `self.client` was None
+ # at the START — the cluster takes ~10 s and a file
+ # opened right after launch beats it. Without this check
+ # that decision was PERMANENT: the whole fill then ran
+ # single-threaded, one chunk at a time, however many
+ # workers turned up a second later. On a 977-frame movie
+ # that is the difference between seconds and minutes,
+ # and it looked exactly like "the cluster is idle".
+ #
+ # It bit movies and not 4D-STEM scans purely by timing:
+ # a scan goes through the nav-shape prompt (a human
+ # round trip) so the cluster is always up by the time
+ # its tree is built, while a movie opens straight
+ # through.
+ #
+ # Re-enter rather than switch in place: the distributed
+ # branch below owns cancellation, the sidecar save and
+ # the final repaint, and duplicating any of that here is
+ # how the two paths drift apart. The chunks already
+ # painted are recomputed, which is why the threshold
+ # exists — a handover is only worth it with real work
+ # left, and this fires within the first few chunks.
+ if (self.client is not None
+ and (total_chunks - done_chunks)
+ > _NAV_HANDOVER_MIN_CHUNKS):
+ logger.info(
+ "navigator fill: cluster came up after %d/%d "
+ "chunks — handing the rest to the dispatcher",
+ done_chunks, total_chunks)
+ self._start_progressive_nav_compute(_dask, deep=_deep)
+ return
# Yield the disk to active scrubbing: for a large movie
# this per-chunk sum reads the whole file, which otherwise
# starves the crosshair's own frame read (the signal plot
diff --git a/spyde/signals/__init__.py b/spyde/signals/__init__.py
index b3e8dd65..5acf6c9c 100644
--- a/spyde/signals/__init__.py
+++ b/spyde/signals/__init__.py
@@ -8,5 +8,15 @@
from spyde.signals.diffraction_vectors import SpyDEDiffractionVectors
from spyde.signals.orientation_map import SpyDEOrientationMap
from spyde.signals.insitu import InSitu, LazyInSitu
+from spyde.signals.particle_map import LazyParticleMap, ParticleMap
+from spyde.signals.particles import SpyDEParticles
-__all__ = ["SpyDEDiffractionVectors", "SpyDEOrientationMap", "InSitu", "LazyInSitu"]
+__all__ = [
+ "SpyDEDiffractionVectors",
+ "SpyDEOrientationMap",
+ "InSitu",
+ "LazyInSitu",
+ "SpyDEParticles",
+ "ParticleMap",
+ "LazyParticleMap",
+]
diff --git a/spyde/signals/particle_map.py b/spyde/signals/particle_map.py
new file mode 100644
index 00000000..e8098c9e
--- /dev/null
+++ b/spyde/signals/particle_map.py
@@ -0,0 +1,37 @@
+"""
+ParticleMap — the signal type a segmentation result carries.
+
+The root of a particle tree (plan §0.6) is a **label movie**: same nav/signal
+shape as the movie it was segmented from, each frame painted on demand from the
+stored contours, with pixel values carrying track ids. It displays like any
+navigated 2-D signal; only the signal type differs.
+
+The type exists so toolbar gating can offer particle actions — track, export,
+per-particle diffraction — on a segmentation result and *nowhere else*. That is
+the same job ``insitu`` does for Play / Fast-Forward, and it is why the plan puts
+the result on its own tree: gating becomes a plain signal-type check instead of a
+hunt up the parent chain for someone else's attribute.
+
+Registered as a HyperSpy extension (see ``spyde/hyperspy_extension.yaml``) so
+``set_signal_type`` and save/load work.
+"""
+from __future__ import annotations
+
+from hyperspy._signals.signal2d import LazySignal2D, Signal2D
+
+SIGNAL_TYPE = "particles"
+
+
+class ParticleMap(Signal2D):
+ """Per-frame particle label map (eager)."""
+ _signal_type = SIGNAL_TYPE
+
+
+class LazyParticleMap(LazySignal2D):
+ """Per-frame particle label map (lazy) — the normal case.
+
+ Lazy is not an optimisation here, it is the design: a materialised label
+ movie is 64 MB *per frame* at 4096², so frames are painted from contours only
+ when something asks for one.
+ """
+ _signal_type = SIGNAL_TYPE
diff --git a/spyde/signals/particles.py b/spyde/signals/particles.py
new file mode 100644
index 00000000..580f5999
--- /dev/null
+++ b/spyde/signals/particles.py
@@ -0,0 +1,463 @@
+"""
+particles.py — :class:`SpyDEParticles`, ragged per-frame particle storage.
+
+Particles-per-frame is a **ragged per-navigation-position collection** — exactly
+the shape :class:`spyde.signals.diffraction_vectors.SpyDEDiffractionVectors`
+already solves. This mirrors that design rather than inventing a second pattern,
+and rather than ParticleSpy's list-of-``Particle``-objects, which does not survive
+the target scale.
+
+Why not a list of objects, and why not full-frame label images
+--------------------------------------------------------------
+The target is thousands of frames with hundreds of particles each — around 1.5M
+particles (DRIFT_AND_PARTICLES_PLAN.md §0.1). Do the arithmetic before choosing a
+representation:
+
+=========================================== ============ ==========
+representation per particle total
+=========================================== ============ ==========
+property row (21 x float32) 84 B 126 MB
+bbox bitmap (packed 64^2 crop) 512 B 770 MB
+contour polygon (~40 pts x 2 x float32) 320 B 480 MB
+**contour polygon, int16** **~80 B** **120 MB**
+=========================================== ============ ==========
+
+So: **properties always resident, outlines as int16 contours, and masks are
+optional.** A full-frame ``int32`` label image is 64 MB *per frame* at 4096^2 and
+is never stored at any setting — :meth:`SpyDEParticles.render_frame` paints one
+frame on demand from its contours.
+
+Contours are quantised to whole pixels. That is a **display** fidelity choice, not
+a measurement one: every measured quantity lives in the property row, computed at
+full precision from the original mask by :mod:`spyde.particles.measure`. Rounding
+an outline for drawing cannot corrupt an area, because the area was never derived
+from the outline.
+"""
+from __future__ import annotations
+
+import json
+from dataclasses import dataclass, field
+from typing import Any, Sequence
+
+import numpy as np
+
+FORMAT_VERSION = 2
+
+#: Column layout of ``flat_buffer``. Order is load-bearing — it is the on-disk
+#: layout. Append new columns at the END and bump ``FORMAT_VERSION``.
+COLUMNS: tuple[str, ...] = (
+ "t", # frame index (float for a uniform buffer dtype)
+ "label", # per-frame instance label, 1-based (0 = background)
+ "y", "x", # centroid, calibrated units
+ "area",
+ "equiv_diameter",
+ "major_axis", "minor_axis",
+ "perimeter",
+ "circularity",
+ "eccentricity",
+ "solidity",
+ "intensity_mean", "intensity_max", "intensity_std",
+ "background", # mean intensity in the dilated boundary ring
+ "bbox_y0", "bbox_x0", "bbox_y1", "bbox_x1", # pixel indices, half-open
+ "track_id", # -1 until the linker runs
+ # How confidently this is a PARTICLE rather than a fragment of textured
+ # background, in [0, 1]. Contrast-to-noise against the particle's own
+ # dilated background ring by default (see spyde.particles.measure), which
+ # an engine with a real probability may overwrite. It exists so the caret
+ # can offer ONE control that means "fewer / more particles": the score is
+ # computed with the measurement, so filtering on it needs no
+ # re-segmentation and a drag is instant.
+ "score",
+)
+COL: dict[str, int] = {name: i for i, name in enumerate(COLUMNS)}
+N_COLUMNS = len(COLUMNS)
+
+#: Columns a user can sort, histogram or colour by — i.e. real measurements,
+#: excluding bookkeeping. Drives the table dock and the histogram window.
+MEASURED_COLUMNS: tuple[str, ...] = (
+ "area", "equiv_diameter", "major_axis", "minor_axis", "perimeter",
+ "circularity", "eccentricity", "solidity",
+ "intensity_mean", "intensity_max", "intensity_std", "background",
+ "score",
+)
+
+#: Which measured columns scale with length, area, or not at all — used to apply
+#: pixel-size calibration exactly once, in one place.
+_LENGTH_COLUMNS = ("y", "x", "equiv_diameter", "major_axis", "minor_axis", "perimeter")
+_AREA_COLUMNS = ("area",)
+
+
+@dataclass
+class SpyDEParticles:
+ """Ragged per-frame particle table with optional outlines.
+
+ Parameters
+ ----------
+ flat_buffer
+ ``(N_total, N_COLUMNS)`` float32, sorted by frame index ``t``.
+ t_offsets
+ ``(n_frames + 1,)`` int64 CSR row pointers. ``t_offsets[i]:t_offsets[i+1]``
+ is frame *i*'s slice — an O(1) lookup, no search.
+ frame_shape
+ ``(h, w)`` of the source frames, in pixels.
+ contours, contour_offsets
+ Optional outlines: ``(M, 2)`` int16 ``(y, x)`` pixel coordinates and an
+ ``(N_total + 1,)`` int64 index. ``None`` when segmentation ran with
+ ``store_masks=False`` (the default for very long movies).
+ scale, units
+ Pixel size and its unit, so ``area`` is in ``units**2``. ``scale=1.0`` with
+ ``units="px"`` means uncalibrated.
+ """
+
+ flat_buffer: np.ndarray
+ t_offsets: np.ndarray
+ frame_shape: tuple[int, int]
+ contours: np.ndarray | None = None
+ contour_offsets: np.ndarray | None = None
+ scale: float = 1.0
+ units: str = "px"
+ params: dict[str, Any] = field(default_factory=dict)
+ provenance: dict[str, Any] | None = None
+
+ def __post_init__(self) -> None:
+ self.flat_buffer = np.ascontiguousarray(self.flat_buffer, dtype=np.float32)
+ if self.flat_buffer.ndim != 2 or self.flat_buffer.shape[1] != N_COLUMNS:
+ raise ValueError(
+ f"flat_buffer must be (N, {N_COLUMNS}); got {self.flat_buffer.shape}"
+ )
+ self.t_offsets = np.ascontiguousarray(self.t_offsets, dtype=np.int64)
+ if self.t_offsets.ndim != 1 or self.t_offsets.size < 1:
+ raise ValueError(f"t_offsets must be 1-D and non-empty; got {self.t_offsets.shape}")
+ if int(self.t_offsets[0]) != 0 or int(self.t_offsets[-1]) != len(self.flat_buffer):
+ raise ValueError(
+ f"t_offsets must span 0..{len(self.flat_buffer)}; got "
+ f"{int(self.t_offsets[0])}..{int(self.t_offsets[-1])}"
+ )
+ if np.any(np.diff(self.t_offsets) < 0):
+ raise ValueError("t_offsets must be non-decreasing")
+ self.frame_shape = (int(self.frame_shape[0]), int(self.frame_shape[1]))
+
+ if (self.contours is None) != (self.contour_offsets is None):
+ raise ValueError("contours and contour_offsets must both be set or both None")
+ if self.contours is not None:
+ self.contours = np.ascontiguousarray(self.contours, dtype=np.int16)
+ if self.contours.ndim != 2 or self.contours.shape[1] != 2:
+ raise ValueError(f"contours must be (M, 2); got {self.contours.shape}")
+ self.contour_offsets = np.ascontiguousarray(self.contour_offsets, dtype=np.int64)
+ if self.contour_offsets.size != len(self.flat_buffer) + 1:
+ raise ValueError(
+ f"contour_offsets must be ({len(self.flat_buffer) + 1},); "
+ f"got {self.contour_offsets.shape}"
+ )
+
+ # ── shape ────────────────────────────────────────────────────────────────
+
+ @property
+ def n_particles(self) -> int:
+ return int(len(self.flat_buffer))
+
+ @property
+ def n_frames(self) -> int:
+ return int(self.t_offsets.size - 1)
+
+ @property
+ def has_masks(self) -> bool:
+ return self.contours is not None
+
+ @property
+ def has_tracks(self) -> bool:
+ """True once the linker has assigned track ids."""
+ if self.n_particles == 0:
+ return False
+ return bool(np.any(self.flat_buffer[:, COL["track_id"]] >= 0))
+
+ # ── per-frame access ─────────────────────────────────────────────────────
+
+ def at(self, t: int) -> np.ndarray:
+ """Frame *t*'s ``(n, N_COLUMNS)`` block. O(1) — a view, not a copy."""
+ t = int(t)
+ if not 0 <= t < self.n_frames:
+ raise IndexError(f"frame {t} outside 0..{self.n_frames - 1}")
+ return self.flat_buffer[self.t_offsets[t]:self.t_offsets[t + 1]]
+
+ def indices_at(self, t: int) -> np.ndarray:
+ """Global particle indices belonging to frame *t*."""
+ t = int(t)
+ if not 0 <= t < self.n_frames:
+ raise IndexError(f"frame {t} outside 0..{self.n_frames - 1}")
+ return np.arange(self.t_offsets[t], self.t_offsets[t + 1], dtype=np.int64)
+
+ def column(self, name: str) -> np.ndarray:
+ """One column across every particle."""
+ try:
+ return self.flat_buffer[:, COL[name]]
+ except KeyError:
+ raise KeyError(
+ f"unknown column {name!r}; available: {', '.join(COLUMNS)}"
+ ) from None
+
+ def contour_at(self, index: int) -> np.ndarray:
+ """``(k, 2)`` int16 ``(y, x)`` outline of global particle *index*."""
+ if self.contours is None:
+ raise ValueError(
+ "no outlines stored (segmentation ran with store_masks=False)"
+ )
+ i = int(index)
+ if not 0 <= i < self.n_particles:
+ raise IndexError(f"particle {i} outside 0..{self.n_particles - 1}")
+ return self.contours[self.contour_offsets[i]:self.contour_offsets[i + 1]]
+
+ # ── navigator traces ─────────────────────────────────────────────────────
+
+ def count_series(self) -> np.ndarray:
+ """``(n_frames,)`` particle count per frame — the navigator's count lane.
+
+ Straight from the CSR row pointers, so it is O(n_frames) regardless of how
+ many particles there are.
+ """
+ return np.diff(self.t_offsets).astype(np.float32)
+
+ def property_series(self, name: str, reduce: str = "mean") -> np.ndarray:
+ """``(n_frames,)`` per-frame reduction of a column — e.g. mean size lane.
+
+ Empty frames yield NaN rather than 0: a frame with no particles has no
+ mean size, and plotting it as zero would draw a spurious spike down to the
+ axis that reads as a real physical event.
+ """
+ col = self.column(name)
+ fn = {"mean": np.mean, "sum": np.sum, "max": np.max,
+ "min": np.min, "median": np.median, "std": np.std}.get(reduce)
+ if fn is None:
+ raise ValueError(
+ f"unknown reduce {reduce!r}; expected mean/sum/max/min/median/std"
+ )
+ out = np.full(self.n_frames, np.nan, dtype=np.float32)
+ for t in range(self.n_frames):
+ s, e = self.t_offsets[t], self.t_offsets[t + 1]
+ if e > s:
+ vals = col[s:e]
+ finite = vals[np.isfinite(vals)]
+ if finite.size:
+ out[t] = fn(finite)
+ return out
+
+ # ── rendering ────────────────────────────────────────────────────────────
+
+ def render_frame(self, t: int, *, value: str = "label") -> np.ndarray:
+ """Paint frame *t*'s particles into an ``int32`` label image.
+
+ Built on demand and never cached here — a 4096^2 label image is 64 MB, so
+ holding even a handful would dwarf the entire particle table. The caller
+ (the overlay) keeps only the frame it is displaying.
+
+ Parameters
+ ----------
+ value
+ ``"label"`` fills each particle with its per-frame label;
+ ``"track"`` fills with ``track_id + 1`` so the overlay can colour by
+ identity across frames (0 stays background). ``"index"`` fills with the
+ global particle index + 1, which is what a click-to-select hit test
+ wants.
+ """
+ from skimage.draw import polygon as sk_polygon
+
+ if self.contours is None:
+ raise ValueError(
+ "cannot render outlines: segmentation ran with store_masks=False"
+ )
+ h, w = self.frame_shape
+ out = np.zeros((h, w), dtype=np.int32)
+ for gi in self.indices_at(t):
+ c = self.contour_at(gi)
+ if len(c) < 3:
+ continue
+ rr, cc = sk_polygon(c[:, 0].astype(np.intp), c[:, 1].astype(np.intp),
+ shape=(h, w))
+ if value == "label":
+ fill = int(self.flat_buffer[gi, COL["label"]])
+ elif value == "track":
+ fill = int(self.flat_buffer[gi, COL["track_id"]]) + 1
+ elif value == "index":
+ fill = int(gi) + 1
+ else:
+ raise ValueError(
+ f"unknown value {value!r}; expected 'label', 'track' or 'index'"
+ )
+ out[rr, cc] = fill
+ return out
+
+ def mask_at(self, index: int) -> tuple[np.ndarray, tuple[int, int, int, int]]:
+ """``(mask, (y0, x0, y1, x1))`` — one particle's boolean mask and its bbox.
+
+ Cropped to the bounding box, so this stays small no matter the frame size.
+ This is what Wave D's per-particle mean diffraction pattern slices with.
+ """
+ from skimage.draw import polygon as sk_polygon
+
+ c = self.contour_at(index)
+ row = self.flat_buffer[int(index)]
+ y0, x0 = int(row[COL["bbox_y0"]]), int(row[COL["bbox_x0"]])
+ y1, x1 = int(row[COL["bbox_y1"]]), int(row[COL["bbox_x1"]])
+ h, w = max(1, y1 - y0), max(1, x1 - x0)
+ m = np.zeros((h, w), dtype=bool)
+ if len(c) >= 3:
+ rr, cc = sk_polygon((c[:, 0] - y0).astype(np.intp),
+ (c[:, 1] - x0).astype(np.intp), shape=(h, w))
+ m[rr, cc] = True
+ return m, (y0, x0, y1, x1)
+
+ # ── export ───────────────────────────────────────────────────────────────
+
+ def to_dataframe(self):
+ """A pandas DataFrame of every particle. Requires pandas at call time."""
+ import pandas as pd
+ return pd.DataFrame(self.flat_buffer, columns=list(COLUMNS))
+
+ def to_csv(self, path: str) -> None:
+ """Write the property table as CSV, with a units line in the header."""
+ header = ",".join(COLUMNS)
+ np.savetxt(
+ path, self.flat_buffer, delimiter=",", header=header, comments="",
+ fmt="%.6g",
+ )
+
+ # ── serialisation ────────────────────────────────────────────────────────
+
+ def save(self, path: str) -> None:
+ meta = {
+ "format_version": FORMAT_VERSION,
+ "columns": list(COLUMNS),
+ "frame_shape": list(self.frame_shape),
+ "scale": float(self.scale),
+ "units": self.units,
+ "params": self.params,
+ "provenance": self.provenance,
+ }
+ arrays = {
+ "flat_buffer": self.flat_buffer,
+ "t_offsets": self.t_offsets,
+ "meta": np.array(json.dumps(meta)),
+ }
+ if self.contours is not None:
+ arrays["contours"] = self.contours
+ arrays["contour_offsets"] = self.contour_offsets
+ np.savez_compressed(path, **arrays)
+
+ @classmethod
+ def load(cls, path: str) -> "SpyDEParticles":
+ with np.load(path, allow_pickle=False) as z:
+ meta = json.loads(str(z["meta"].item()))
+ ver = meta.get("format_version")
+ if not isinstance(ver, int) or ver > FORMAT_VERSION:
+ raise ValueError(
+ f"unsupported SpyDEParticles format version {ver!r} "
+ f"(this build reads {FORMAT_VERSION})"
+ )
+ saved_cols = tuple(meta.get("columns") or ())
+ buf = z["flat_buffer"]
+ # An OLDER file is readable when the only difference is columns
+ # APPENDED since — which is the documented way to extend this
+ # layout, and the reason new columns go at the end. Refusing it
+ # would make every previously-saved particle result unopenable to
+ # add one derived number, and the number is derived precisely so it
+ # need not have been stored at the time.
+ if saved_cols != COLUMNS:
+ if saved_cols != COLUMNS[:len(saved_cols)]:
+ raise ValueError(
+ "column layout changed since this file was written "
+ f"({len(saved_cols)} columns on disk, {N_COLUMNS} "
+ "expected)"
+ )
+ pad = np.full((buf.shape[0], N_COLUMNS - len(saved_cols)),
+ np.nan, np.float32)
+ for i, name in enumerate(COLUMNS[len(saved_cols):]):
+ # `score` is a confidence, and an absent one means "no
+ # evidence against this particle" — the caret's filter must
+ # not silently delete instances measured before it existed.
+ if name == "score":
+ pad[:, i] = 1.0
+ buf = np.ascontiguousarray(
+ np.concatenate([buf, pad], axis=1), np.float32)
+ return cls(
+ flat_buffer=buf,
+ t_offsets=z["t_offsets"],
+ frame_shape=tuple(meta["frame_shape"]),
+ contours=z["contours"] if "contours" in z.files else None,
+ contour_offsets=(z["contour_offsets"]
+ if "contour_offsets" in z.files else None),
+ scale=meta.get("scale", 1.0),
+ units=meta.get("units", "px"),
+ params=meta.get("params") or {},
+ provenance=meta.get("provenance"),
+ )
+
+ # ── construction ─────────────────────────────────────────────────────────
+
+ @classmethod
+ def from_frames(
+ cls,
+ per_frame: Sequence[np.ndarray],
+ *,
+ frame_shape: tuple[int, int],
+ contours_per_frame: Sequence[Sequence[np.ndarray]] | None = None,
+ scale: float = 1.0,
+ units: str = "px",
+ params: dict[str, Any] | None = None,
+ provenance: dict[str, Any] | None = None,
+ ) -> "SpyDEParticles":
+ """Build from a per-frame list of ``(n_i, N_COLUMNS)`` property blocks.
+
+ *contours_per_frame*, when given, must line up exactly with *per_frame* —
+ one outline per row. A mismatch raises rather than silently pairing the
+ wrong outline with a particle, which would draw plausible nonsense.
+ """
+ n_frames = len(per_frame)
+ blocks: list[np.ndarray] = []
+ counts: list[int] = []
+ for i, blk in enumerate(per_frame):
+ b = np.zeros((0, N_COLUMNS), np.float32) if blk is None or len(blk) == 0 \
+ else np.asarray(blk, dtype=np.float32)
+ if b.ndim != 2 or b.shape[1] != N_COLUMNS:
+ raise ValueError(
+ f"frame {i}: expected (n, {N_COLUMNS}); got {b.shape}"
+ )
+ blocks.append(b)
+ counts.append(len(b))
+
+ flat = (np.concatenate(blocks, axis=0) if blocks
+ else np.zeros((0, N_COLUMNS), np.float32))
+ offsets = np.concatenate([[0], np.cumsum(counts)]).astype(np.int64)
+
+ contours = contour_offsets = None
+ if contours_per_frame is not None:
+ if len(contours_per_frame) != n_frames:
+ raise ValueError(
+ f"contours_per_frame has {len(contours_per_frame)} frames, "
+ f"per_frame has {n_frames}"
+ )
+ polys: list[np.ndarray] = []
+ for i, (cs, n) in enumerate(zip(contours_per_frame, counts)):
+ cs = list(cs or ())
+ if len(cs) != n:
+ raise ValueError(
+ f"frame {i}: {len(cs)} outlines for {n} particles — "
+ "outlines must correspond 1:1 with property rows"
+ )
+ polys.extend(np.asarray(c, dtype=np.int16).reshape(-1, 2) for c in cs)
+ contours = (np.concatenate(polys, axis=0) if polys
+ else np.zeros((0, 2), np.int16))
+ contour_offsets = np.concatenate(
+ [[0], np.cumsum([len(p) for p in polys])]).astype(np.int64)
+
+ return cls(
+ flat_buffer=flat, t_offsets=offsets, frame_shape=frame_shape,
+ contours=contours, contour_offsets=contour_offsets,
+ scale=scale, units=units, params=params or {}, provenance=provenance,
+ )
+
+ def __repr__(self) -> str:
+ return (
+ f"SpyDEParticles({self.n_particles} particles over {self.n_frames} "
+ f"frames, {self.frame_shape[0]}x{self.frame_shape[1]} px, "
+ f"masks={self.has_masks}, tracks={self.has_tracks})"
+ )
diff --git a/spyde/tests/benchmark_drift_nonrigid.py b/spyde/tests/benchmark_drift_nonrigid.py
new file mode 100644
index 00000000..4c545302
--- /dev/null
+++ b/spyde/tests/benchmark_drift_nonrigid.py
@@ -0,0 +1,144 @@
+"""
+benchmark_drift_nonrigid.py — what does a non-rigid solve cost on a real movie?
+
+Run directly (torch-CUDA segfaults under pytest on Windows -- CLAUDE.md)::
+
+ python -m spyde.tests.benchmark_drift_nonrigid
+ python -m spyde.tests.benchmark_drift_nonrigid --device cpu --frames 100
+
+The question is 4096x4096 x hundreds of frames, and the first thing to say
+about it is that the stack CANNOT be held: 100 x 4096^2 float32 is 6.7 GB and
+300 is 20 GB. So the solve is necessarily two separate costs and they are
+measured separately here, because they scale differently and only one of them
+is paid per frame:
+
+1. **The FIT**, on a DECIMATED stack. A drift field is smooth by construction --
+ that is the entire modelling assumption -- so it does not need full
+ resolution to be measured. The fit is over a handful of parameters per frame
+ (2 x n_knots, or 2 x gh x gw), and the decimation factor is the dominant
+ cost knob.
+2. **The APPLY**, at FULL resolution, once per frame, in the display/export
+ path. This is the number that decides whether scrubbing stays interactive.
+
+Reporting one blended figure would hide exactly the thing a caller needs to
+decide: how much to decimate.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+
+import numpy as np
+
+
+def _synth(n: int, h: int, w: int, seed: int = 0) -> np.ndarray:
+ """A textured stack with a known slow-axis distortion, at fit resolution."""
+ rng = np.random.default_rng(seed)
+ y, x = np.mgrid[0:h, 0:w].astype(np.float32)
+ base = np.zeros((h, w), np.float32)
+ for _ in range(40):
+ cy, cx = rng.uniform(0, h), rng.uniform(0, w)
+ s = rng.uniform(h / 60, h / 20)
+ base += rng.uniform(0.5, 1.5) * np.exp(-((y - cy) ** 2 + (x - cx) ** 2) / (2 * s * s))
+ base += 0.05 * rng.standard_normal((h, w)).astype(np.float32)
+
+ from spyde.drift import nonrigid as nr
+ import torch
+ rows = np.linspace(-1.0, 1.0, h, dtype=np.float32)
+ out = np.empty((n, h, w), np.float32)
+ t = torch.as_tensor(base)[None]
+ for i in range(n):
+ a = 3.0 * (0.3 + 0.7 * i / max(n - 1, 1))
+ dy = np.repeat((a * rows)[:, None], w, axis=1)
+ g = nr.warp_frame(torch, t, torch.as_tensor(dy)[None],
+ torch.as_tensor(np.zeros((h, w), np.float32))[None],
+ fill_nan=False)
+ out[i] = g[0].numpy()
+ return out
+
+
+def _sync(device: str) -> None:
+ import torch
+ if device == "cuda":
+ torch.cuda.synchronize()
+ elif device == "mps":
+ torch.mps.synchronize()
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--device", default=None)
+ ap.add_argument("--frames", type=int, default=300)
+ ap.add_argument("--full", type=int, default=4096)
+ ap.add_argument("--steps", type=int, default=120)
+ ap.add_argument("--json", default=None)
+ a = ap.parse_args()
+
+ import torch
+ from spyde.drift import nonrigid as nr
+
+ dev = a.device or ("cuda" if torch.cuda.is_available() else "cpu")
+ print(f"device: {dev} frames: {a.frames} full frame: {a.full}^2 steps: {a.steps}")
+ print(f"(the full stack would be "
+ f"{a.frames * a.full * a.full * 4 / 1e9:.1f} GB -- hence decimation)\n")
+ out: dict = {"device": dev, "frames": a.frames, "full": a.full, "steps": a.steps}
+
+ # ── 1. fit cost vs decimated size ────────────────────────────────────────
+ print("=== FIT (decimated stack, whole movie at once) ===")
+ print(f"{'fit size':>10} {'decim':>7} {'model':>10} {'build':>8} {'fit':>9} {'per-frame':>10}")
+ out["fit"] = {}
+ for side in (128, 256, 512):
+ t0 = time.perf_counter()
+ stack = _synth(a.frames, side, side)
+ build = time.perf_counter() - t0
+ for model, kw in ((nr.SCAN_KNOT, dict(n_knots=3)), (nr.DENSE, dict(grid=(6, 6)))):
+ _sync(dev)
+ t0 = time.perf_counter()
+ nr.solve_nonrigid(stack, model=model, steps=a.steps, device=dev, **kw)
+ _sync(dev)
+ el = time.perf_counter() - t0
+ print(f"{side:>7}^2 {a.full // side:>6}x {model:>10} {build:>7.2f}s "
+ f"{el:>8.2f}s {el / a.frames * 1e3:>8.1f} ms")
+ out["fit"][f"{side}/{model}"] = {"s": el, "per_frame_ms": el / a.frames * 1e3}
+ del stack
+
+ # ── 2. apply cost at FULL resolution, per frame ──────────────────────────
+ print(f"\n=== APPLY, {a.full}^2, per frame (display/export path) ===")
+ out["apply"] = {}
+ small = _synth(4, 128, 128)
+ frame = np.ascontiguousarray(
+ np.random.default_rng(1).random((a.full, a.full), dtype=np.float32))
+ for model, kw in ((nr.SCAN_KNOT, dict(n_knots=3)), (nr.DENSE, dict(grid=(6, 6)))):
+ m = nr.solve_nonrigid(small, model=model, steps=10, device=dev, **kw)
+ # Re-target the fitted field to the full frame: the parameters are
+ # resolution-independent by construction, which is the property that
+ # makes fit-small/apply-large legitimate rather than a shortcut.
+ m.extra["field_shape"] = (a.full, a.full)
+ nr.apply_nonrigid(frame, m, 0) # warm
+ t0 = time.perf_counter()
+ reps = 3
+ for _ in range(reps):
+ nr.apply_nonrigid(frame, m, 0)
+ el = (time.perf_counter() - t0) / reps
+ print(f" {model:>10} {el * 1e3:>7.0f} ms/frame "
+ f"-> {el * a.frames:>6.1f} s for {a.frames} frames")
+ out["apply"][model] = {"ms": el * 1e3, "movie_s": el * a.frames}
+
+ print("\n" + json.dumps(out, indent=2, default=float))
+ if a.json:
+ with open(a.json, "w", encoding="utf-8") as fh:
+ json.dump(out, fh, indent=2, default=float)
+ # FLUSH BEFORE `os._exit`. `_exit` skips atexit and does NOT flush stdio, so
+ # with output redirected to a file (buffered, not line-buffered) every print
+ # above is discarded and the run looks like it produced nothing. Invisible
+ # on a terminal, which is why it is easy to write and easy to miss.
+ sys.stdout.flush()
+ sys.stderr.flush()
+ os._exit(0) # skip torch/CUDA teardown crash (CLAUDE.md)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/spyde/tests/benchmark_nav_submit_batching.py b/spyde/tests/benchmark_nav_submit_batching.py
new file mode 100644
index 00000000..b7f134a4
--- /dev/null
+++ b/spyde/tests/benchmark_nav_submit_batching.py
@@ -0,0 +1,112 @@
+"""
+benchmark_nav_submit_batching.py — how many times do we call the scheduler?
+
+Run directly::
+
+ python -m spyde.tests.benchmark_nav_submit_batching
+ python -m spyde.tests.benchmark_nav_submit_batching --chunks 977
+
+The navigator fill on a big dataset looked like "the cluster is using one
+worker". It was not a placement problem: `dispatch_chunks` tops up on EVERY
+completion, so on the unpinned lane `lane_cap - outstanding` is 1 in steady
+state and the batch size collapsed to ONE task per submit. `submit_batch=8`
+only ever applied to the first fill.
+
+Each of those submits is a blocking scheduler round trip with the GIL held in
+the client process, so the client — not the cluster — became the bottleneck,
+and it got worse with more chunks. Hence "only big datasets".
+
+What this measures, and why not just wall-clock: on a synthetic cluster the
+per-task WORK is trivial, so total time understates the problem. The honest
+metric is **how many times we call `client.compute`** (each is a round trip
+that scales with graph size), plus time-to-first-chunk, which is what the user
+actually watches.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+
+import numpy as np
+
+
+def main() -> None:
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--chunks", type=int, default=256)
+ ap.add_argument("--workers", type=int, default=6)
+ ap.add_argument("--threads", type=int, default=2)
+ ap.add_argument("--json", default=None)
+ a = ap.parse_args()
+
+ import dask.array as da
+ from dask.distributed import Client, LocalCluster
+ import spyde.compute_dispatch as cd
+
+ cluster = LocalCluster(n_workers=a.workers, threads_per_worker=a.threads,
+ processes=True, dashboard_address=None, silence_logs=50)
+ client = Client(cluster)
+ client.wait_for_workers(a.workers)
+ print(f"cluster: {a.workers} workers x {a.threads} threads "
+ f"nav chunks: {a.chunks}\n")
+
+ # A 1-D navigator sum, one task per chunk — the movie shape.
+ arr = da.zeros((a.chunks, 64, 64), chunks=(1, 64, 64), dtype=np.float32)
+ nav = arr.sum(axis=(1, 2))
+
+ out: dict = {"chunks": a.chunks, "workers": a.workers, "threads": a.threads}
+
+ # `submit_batch=1` reproduces the OLD steady-state behaviour exactly: the
+ # window refilled one task at a time. Comparing against it isolates the
+ # batching change from everything else in the dispatcher.
+ for label, kw in (("one-at-a-time (the bug)", dict(batch_unpinned=False, cap=8)),
+ ("batched (fixed)", dict())):
+ calls = {"n": 0}
+ real = client.compute
+
+ def counting(x, *args, **kwargs):
+ if isinstance(x, (list, tuple)):
+ calls["n"] += 1
+ return real(x, *args, **kwargs)
+
+ client.compute = counting
+ first = {"t": None}
+ t0 = time.perf_counter()
+
+ def assemble(res, sl, val):
+ if first["t"] is None:
+ first["t"] = time.perf_counter() - t0
+ res[sl] = val
+
+ try:
+ cd.dispatch_chunks(client, nav, 1, [], None, assemble=assemble,
+ fill_value=np.nan, label=label,
+ lane_default_mode="off", **kw)
+ finally:
+ client.compute = real
+ el = time.perf_counter() - t0
+ print(f" {label:24} submits={calls['n']:>4} "
+ f"first-chunk={first['t'] * 1e3:>6.0f} ms total={el:>6.2f}s")
+ out[label] = {"submits": calls["n"], "first_ms": first["t"] * 1e3,
+ "total_s": el}
+
+ a_, b_ = out["one-at-a-time (the bug)"], out["batched (fixed)"]
+ print(f"\n submits: {a_['submits']} -> {b_['submits']} "
+ f"({a_['submits'] / max(b_['submits'], 1):.0f}x fewer round trips)")
+ print(" NB each saved round trip is ~14 ms with the GIL held on a real "
+ "graph (benchmarks.md), which a synthetic graph does not charge us.")
+
+ if a.json:
+ with open(a.json, "w", encoding="utf-8") as fh:
+ json.dump(out, fh, indent=2, default=float)
+ client.close()
+ cluster.close()
+ sys.stdout.flush() # _exit skips stdio flushing
+ sys.stderr.flush()
+ os._exit(0)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/spyde/tests/benchmark_particles_batch.py b/spyde/tests/benchmark_particles_batch.py
new file mode 100644
index 00000000..41d4daa4
--- /dev/null
+++ b/spyde/tests/benchmark_particles_batch.py
@@ -0,0 +1,348 @@
+"""
+benchmark_particles_batch.py — whole-movie segmentation at REAL frame size.
+
+The report that started this was "scribble segmentation over 900 frames of
+4096x4096 is far too slow, and the GPU is hardly used, as are the CPUs". This
+harness measures that end to end: the per-stage cost of ONE 4096^2 frame, then
+the whole-movie throughput serially and through the dual-lane dispatcher
+(:mod:`spyde.particles.batch`), and projects the 900-frame wall clock from both.
+
+Run it on a REAL in-situ movie (CLAUDE.md § Benchmarking: real dataset, real
+scale, end to end). A synthetic 96x112 fixture validates correctness and hides
+every cost that actually bites here — ``measure_frame`` alone is 2 ms per
+PARTICLE, and a real 4096^2 growth frame has 26 566 of them.
+
+Not run under pytest (it is minutes long, and torch-CUDA segfaults under the
+pytest process on Windows). Run it directly::
+
+ .venv/Scripts/python -m spyde.tests.benchmark_particles_batch
+ .venv/Scripts/python -m spyde.tests.benchmark_particles_batch --frames 36
+ .venv/Scripts/python -m spyde.tests.benchmark_particles_batch --engine scribble
+ .venv/Scripts/python -m spyde.tests.benchmark_particles_batch --stages-only
+ .venv/Scripts/python -m spyde.tests.benchmark_particles_batch --serial
+
+``--frames`` is a frame COUNT at the real frame SIZE; the 900-frame figure is
+extrapolated from it and labelled as such. Twelve frames per worker is enough
+for the lanes to reach steady state.
+"""
+from __future__ import annotations
+
+import argparse
+import os
+import sys
+import time
+
+import numpy as np
+
+# Candidate real in-situ movies on this dev box (first that exists wins) — the
+# same list benchmark_movie_playback.py uses.
+_CANDIDATES = [
+ r"C:\Users\CarterFrancis\Downloads\20251117_88075_run3 some growth_1236_movie.mrc",
+ r"C:\Users\CarterFrancis\Downloads\20251117_88074_run1_9104_movie.mrc",
+ r"C:\Users\CarterFrancis\Downloads\20241002_07954_movie.mrc",
+]
+
+TARGET_FRAMES = 900 # the user's movie length, for the projection
+
+
+def _default_path() -> str | None:
+ for p in _CANDIDATES:
+ if os.path.exists(p):
+ return p
+ return None
+
+
+def _fmt_hms(seconds: float) -> str:
+ m, s = divmod(int(round(seconds)), 60)
+ h, m = divmod(m, 60)
+ return f"{h:d}h{m:02d}m{s:02d}s" if h else f"{m:d}m{s:02d}s"
+
+
+# ── the scribble head ────────────────────────────────────────────────────────
+
+def train_scribble(frame, labels, *, device=None, crop: int = 1024):
+ """A realistically-trained head, from pseudo-scribbles on a 1024^2 crop.
+
+ Painted by hand in the app; synthesised here from a classical segmentation
+ of the same frame so the benchmark is reproducible. What matters for TIMING
+ is the feature spec and the frame size, both of which are the real ones —
+ the particular strokes only move the particle count.
+ """
+ from scipy import ndimage as ndi
+
+ from spyde.particles.features import FeatureSpec, select_device
+ from spyde.particles.scribble import (LabelStore, ScribbleClassifier,
+ default_classes)
+
+ h, w = frame.shape
+ y0, x0 = (h - crop) // 2, (w - crop) // 2
+ lab = labels[y0:y0 + crop, x0:x0 + crop]
+ img = frame[y0:y0 + crop, x0:x0 + crop]
+
+ fg = lab > 0
+ core = ndi.binary_erosion(fg, iterations=1)
+ far = ~ndi.binary_dilation(fg, iterations=6)
+
+ def sub(mask, k=4000):
+ idx = np.flatnonzero(np.asarray(mask).reshape(-1))
+ return idx[:: max(1, idx.size // k)] if idx.size > k else idx
+
+ store = LabelStore(frame_shape=(crop, crop), classes=default_classes())
+ store.paint(0, sub(core), 0)
+ store.paint(0, sub(far), 1)
+ clf = ScribbleClassifier(FeatureSpec(), device=select_device(device), seed=0)
+ rep = clf.fit(store, {0: img})
+ print(f" scribble head: {rep['n_pixels']} px / {rep['n_classes']} classes, "
+ f"acc {rep['train_accuracy']:.3f}, boundary={rep['has_boundary']}, "
+ f"device {clf.device}")
+ return clf
+
+
+# ── per-stage, one frame ─────────────────────────────────────────────────────
+
+def stage_profile(frame, sp, clf=None) -> None:
+ """Where one 4096^2 frame's time goes, engine by engine and stage by stage."""
+ from spyde.particles.classical import split_instances
+ from spyde.particles.measure import _contours, _fill_intensity
+ from spyde.particles import measure_frame, segment_frame
+
+ print(f"\n== one frame, {frame.shape[0]}x{frame.shape[1]} {frame.dtype} ==")
+
+ t0 = time.perf_counter()
+ labels = segment_frame(frame, sp)
+ t1 = time.perf_counter()
+ rows, cs = measure_frame(labels, frame, t=0, scale=1.0)
+ t2 = time.perf_counter()
+ print(f" classical : segment {t1-t0:6.2f}s measure {t2-t1:7.2f}s "
+ f"n={len(rows)} total {t2-t0:7.2f}s")
+
+ if clf is not None:
+ import torch
+ sync = (lambda: torch.cuda.synchronize()) if clf.device.type == "cuda" \
+ else (lambda: None)
+ sync()
+ t0 = time.perf_counter()
+ fg, bnd = clf.predict_foreground_boundary(frame)
+ sync()
+ t1 = time.perf_counter()
+ lab2 = split_instances(fg, sp, boundary=bnd)
+ t2 = time.perf_counter()
+ rows2, _cs2 = measure_frame(lab2, frame, t=0, scale=1.0)
+ t3 = time.perf_counter()
+ print(f" scribble : predict {t1-t0:6.2f}s split {t2-t1:6.2f}s "
+ f"measure {t3-t2:7.2f}s n={len(rows2)} total {t3-t0:7.2f}s")
+ labels = lab2
+
+ # measure_frame is the stage the whole-movie cost turns on once a real frame
+ # has thousands of particles, so break it down — BOTH ways, because all three
+ # of its stages have been replaced and the interesting number is the ratio,
+ # not either column alone (`benchmarks.md` § "Vectorising measure_frame").
+ from spyde.particles.measure import property_table, warm_kernels
+ from spyde.signals.particles import N_COLUMNS
+
+ warm_kernels() # never time a numba compile as if it were the work
+ inten = np.asarray(frame, np.float64)
+ for fast in (False, True):
+ t0 = time.perf_counter()
+ tbl = property_table(labels, fast=fast)
+ t1 = time.perf_counter()
+ n = len(tbl["label"])
+ r = np.zeros((n, N_COLUMNS), np.float32)
+ keep = np.ones(n, bool)
+ _fill_intensity(r, labels, inten, tbl, keep, 3, fast=fast)
+ t2 = time.perf_counter()
+ _contours(labels, tbl, fast=fast)
+ t3 = time.perf_counter()
+ tag = "vectorised" if fast else "regionprops"
+ print(f" measure_frame internals, {tag:11s} ({n} particles): "
+ f"table {t1-t0:6.2f}s intensity {t2-t1:5.2f}s "
+ f"contours {t3-t2:5.2f}s sum {t3-t0:6.2f}s")
+
+
+# ── whole-movie ──────────────────────────────────────────────────────────────
+
+def run_serial(data, spec, n, scale) -> tuple[float, int]:
+ """The retired shape: one thread, one frame at a time."""
+ from spyde.particles.batch import resolve_engine
+ from spyde.particles.measure import measure_frame
+
+ engine, dev = resolve_engine(spec)
+ t0 = time.perf_counter()
+ total = 0
+ for t in range(n):
+ frame = np.asarray(data[t])
+ if hasattr(frame, "compute"):
+ frame = frame.compute()
+ labels = engine(frame)
+ rows, _cs = measure_frame(labels, frame, t=t, scale=scale)
+ total += len(rows)
+ return time.perf_counter() - t0, total
+
+
+def run_batch(data, spec, n, scale, client) -> tuple[float, int]:
+ from spyde.particles.batch import segment_movie
+
+ t0 = time.perf_counter()
+ rows, _cs, done = segment_movie(data, spec, n_frames=n, scale=scale,
+ store_masks=True, client=client)
+ dt = time.perf_counter() - t0
+ assert done == n, f"only {done}/{n} frames landed"
+ _worker_stage_summary(client)
+ return dt, sum(len(r) for r in rows)
+
+
+def _worker_stage_summary(client) -> None:
+ """Per-lane in-cluster stage costs, taken off the workers themselves.
+
+ ``engine/f`` and ``measure/f`` here are what a frame cost INSIDE the
+ cluster; comparing them with the single-frame profile above is the whole
+ diagnosis — a stage that is 2x its solo cost is contended, and a lane whose
+ frames/s is far below its stage sum is starved rather than slow.
+ """
+ from spyde.particles.batch import drain_stage_log
+
+ per_dev: dict[str, list] = {}
+ try:
+ for recs in client.run(drain_stage_log).values():
+ for _t0, n, dev, t_eng, t_meas, t_wall in recs:
+ per_dev.setdefault(dev, []).append((n, t_eng, t_meas, t_wall))
+ except Exception as exc:
+ print(f" (worker telemetry unavailable: {exc})")
+ return
+ if not per_dev:
+ return
+ print(f" {'lane':>6} {'frames':>7} {'engine/f':>9} {'measure/f':>10} "
+ f"{'block/f':>8}")
+ for dev, rows in sorted(per_dev.items()):
+ nf = sum(r[0] for r in rows) or 1
+ print(f" {dev:>6} {nf:>7} {sum(r[1] for r in rows)/nf:9.2f} "
+ f"{sum(r[2] for r in rows)/nf:10.2f} "
+ f"{sum(r[3] for r in rows)/nf:8.2f}")
+
+
+def make_cluster(n_workers: int, threads: int):
+ from dask.distributed import Client, LocalCluster
+ cluster = LocalCluster(n_workers=0, threads_per_worker=threads,
+ processes=True)
+ client = Client(cluster)
+ cluster.scale(n_workers)
+ client.wait_for_workers(n_workers, timeout=180)
+
+ def _enable_telemetry():
+ import logging as _lg
+ _lg.getLogger("spyde.particles.batch").setLevel(_lg.INFO)
+
+ client.run(_enable_telemetry)
+ return cluster, client
+
+
+def main(argv=None) -> int:
+ ap = argparse.ArgumentParser()
+ ap.add_argument("--path", default=_default_path())
+ ap.add_argument("--frames", type=int, default=24)
+ ap.add_argument("--start", type=int, default=10)
+ ap.add_argument("--engine", default="both",
+ choices=("classical", "scribble", "both"))
+ ap.add_argument("--serial", action="store_true",
+ help="also time the retired serial loop (slow)")
+ ap.add_argument("--stages-only", action="store_true")
+ ap.add_argument("--gpu-lane", default=None,
+ help="SPYDE_FV_GPU override for this run (one/N/all/off)")
+ ap.add_argument("--workers", type=int, default=0)
+ ap.add_argument("--threads", type=int, default=0)
+ ap.add_argument("--conc", default=None,
+ help="SPYDE_FV_GPU_CONC override (device slots per process)")
+ args = ap.parse_args(argv)
+
+ import logging
+ logging.basicConfig(level=logging.WARNING, format="%(message)s",
+ stream=sys.stdout)
+ # The dispatcher's own summary line is the lane split — GPU chunks vs CPU
+ # chunks — which is the number that says whether the lanes are balanced.
+ logging.getLogger("spyde.compute_dispatch").setLevel(logging.DEBUG)
+ logging.getLogger("spyde.particles.batch").setLevel(logging.INFO)
+ if args.conc:
+ os.environ["SPYDE_FV_GPU_CONC"] = str(args.conc)
+
+ if args.gpu_lane:
+ os.environ["SPYDE_FV_GPU"] = str(args.gpu_lane)
+ if not args.path or not os.path.exists(args.path):
+ print("no in-situ movie found; pass --path", file=sys.stderr)
+ return 2
+
+ import hyperspy.api as hs
+ from spyde.backend.app import _compute_worker_plan
+ from spyde.particles import SegmentParams
+ from spyde.particles.batch import EngineSpec, save_engine_model
+
+ print(f"movie: {args.path}")
+ s = hs.load(args.path, lazy=True)
+ # Load the way the APP loads. RosettaSciIO auto-chunks a big MRC as a
+ # balanced cube — a real 977 x 4096^2 movie arrives as (511, 511, 511),
+ # which SPLITS the signal axes, so one frame spans 64 blocks and 8.5 GB.
+ # `Session._signal_spanning_chunks` re-loads every movie with whole-signal
+ # chunks (free: a lazy reload only rebuilds the graph); benchmarking the
+ # reader default instead would measure a shuffle the app never performs.
+ from spyde.backend.session import Session
+ ch = Session._signal_spanning_chunks(s)
+ if ch is not None:
+ print(f" reader chunks {s.data.chunks[0][:2]}... split the signal "
+ f"axes — re-loading with {ch} (as the app does)")
+ s = hs.load(args.path, lazy=True, chunks=ch)
+ raw = s.data
+ print(f" shape {raw.shape} {raw.dtype} nav chunks "
+ f"{raw.chunks[0][:4]}{'...' if len(raw.chunks[0]) > 4 else ''} "
+ f"signal chunks {tuple(c[0] for c in raw.chunks[1:])}")
+
+ sp_kwargs = dict(min_size=20)
+ sp = SegmentParams(**sp_kwargs)
+ t0 = args.start
+ frame = np.asarray(raw[t0].compute())
+
+ clf = None
+ model_path = None
+ if args.engine in ("scribble", "both"):
+ from spyde.particles import segment_frame
+ print("\n== training the scribble head ==")
+ labels = segment_frame(frame, sp)
+ clf = train_scribble(frame, labels)
+ model_path = save_engine_model(clf)
+
+ stage_profile(frame, sp, clf)
+ if args.stages_only:
+ return 0
+
+ workers, threads = _compute_worker_plan(os.cpu_count() or 4)
+ workers = args.workers or workers
+ threads = args.threads or threads
+ lane = os.environ.get("SPYDE_FV_GPU", " 4>")
+ print(f"\n== cluster: {workers} workers x {threads} threads, "
+ f"SPYDE_FV_GPU={lane} ==")
+ cluster, client = make_cluster(workers, threads)
+
+ n = int(args.frames)
+ sub = raw[t0:t0 + n]
+ try:
+ for method in (("classical", "scribble") if args.engine == "both"
+ else (args.engine,)):
+ spec = EngineSpec(method=method, params=sp_kwargs,
+ model_path=model_path)
+ print(f"\n== {method}: {n} frames of "
+ f"{frame.shape[0]}x{frame.shape[1]} ==")
+ if args.serial:
+ dt, total = run_serial(sub, spec, n, 1.0)
+ print(f" serial : {dt:7.1f}s {n/dt:6.3f} frames/s "
+ f"{total} particles -> 900 frames = "
+ f"{_fmt_hms(dt / n * TARGET_FRAMES)}")
+ dt, total = run_batch(sub, spec, n, 1.0, client)
+ print(f" batch : {dt:7.1f}s {n/dt:6.3f} frames/s "
+ f"{total} particles -> 900 frames = "
+ f"{_fmt_hms(dt / n * TARGET_FRAMES)}")
+ finally:
+ client.close()
+ cluster.close()
+ return 0
+
+
+if __name__ == "__main__":
+ sys.exit(main())
diff --git a/spyde/tests/benchmark_scribble_cnn.py b/spyde/tests/benchmark_scribble_cnn.py
new file mode 100644
index 00000000..2ed67a76
--- /dev/null
+++ b/spyde/tests/benchmark_scribble_cnn.py
@@ -0,0 +1,613 @@
+"""
+benchmark_scribble_cnn.py — the CNN-vs-MLP scribble prototype, A/B'd.
+
+Run directly (it is slow, and torch-CUDA segfaults under the pytest process on
+Windows — CLAUDE.md)::
+
+ python -m spyde.tests.benchmark_scribble_cnn
+ python -m spyde.tests.benchmark_scribble_cnn --only quality --device cpu
+
+It prints human-readable tables as it goes and one JSON blob at the end, then
+``os._exit(0)`` to skip torch's CUDA teardown crash.
+
+It answers exactly two questions and nothing else:
+
+**1. TRAIN TIME.** ``ScribbleClassifier.fit`` is ~0.5 s from a few strokes on the
+fixture and the caret's whole tuning loop depends on that. So the CNN's train
+time is measured at two scales: the fixture (96×112, a few thousand labelled
+pixels) and a REALISTIC one — a 2048² field with a real session's labelled-pixel
+counts (~16 k particle / ~33 k support film / ~1.8 k vacuum) spread across it.
+A steps sweep gives the accuracy-vs-time curve, so a knee is visible rather than
+inferred.
+
+**2. QUALITY, against exact ground truth.** ``particle_movie()`` knows where every
+particle is (``particle_truth_at``), which two are the deliberately faint §0.9
+probes, and which frame the touching pair merges on. Both engines are trained on
+the SAME :class:`~spyde.particles.scribble.LabelStore` and scored on: particle
+count, both faint probes found, the merge pair still split at the merge frame,
+and foreground IoU.
+
+Plus the forward-pass cost and peak VRAM at 4096², confirming the numbers the
+prototype was designed around.
+
+Why the two engines share one evaluator
+---------------------------------------
+:class:`~spyde.particles.scribble_cnn.ScribbleCNN` deliberately has the same
+output contract as :class:`~spyde.particles.scribble.ScribbleClassifier` —
+``predict_foreground_boundary`` and ``segment`` — so :func:`evaluate` takes
+either one and there is no second scoring path that could flatter either.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import os
+import sys
+import time
+
+import numpy as np
+
+from spyde.data.synthetic import (
+ MERGE_PAIR,
+ ground_truth,
+ particle_movie,
+ particle_truth_at,
+)
+from spyde.particles.classical import SegmentParams, split_instances
+from spyde.particles.features import FeatureSpec, select_device
+from spyde.particles.scribble import LabelStore, ScribbleClassifier, default_classes
+from spyde.particles.scribble_cnn import CONFIGS, ScribbleCNN, build_net
+
+#: The frame every quality number is measured on — all nine particles present.
+FRAME_T = 12
+
+#: Split parameters, matching ``test_particles_scribble.py``'s gates so the
+#: numbers here are comparable to the ones already recorded there.
+SPLIT = SegmentParams(min_size=5)
+
+
+# ── the fixture's scribbles (the same eleven strokes the gates use) ──────────
+
+def _clear_of_particles(shape, pos, radii, present, pad):
+ h, w = shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ keep = np.ones((h, w), bool)
+ for i in np.flatnonzero(present):
+ keep &= ((yy - pos[i, 0]) ** 2 +
+ (xx - pos[i, 1]) ** 2) > (radii[i] + pad) ** 2
+ return keep
+
+
+def paint_fixture_scribbles(geom, *, t=FRAME_T, seam=True) -> LabelStore:
+ """Copied from ``test_particles_scribble.py``'s ``paint_scribbles`` (+ seam).
+
+ Copied and not imported: the test module builds module-scoped fixtures at
+ import time, and a benchmark that drags a pytest fixture graph in is a
+ benchmark that measures the fixture graph. The strokes are what matter and
+ they are reproduced exactly — four dabs on bright particles, one dab on the
+ SMALLER faint probe (index 8; index 7 stays held out), four background
+ sweeps, four background rings, and the seams between touching bodies.
+ """
+ from scipy import ndimage as ndi
+
+ pos, radii, present, faint, shape = geom
+ h, w = shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ store = LabelStore(frame_shape=shape, classes=default_classes())
+
+ bright = [i for i in np.flatnonzero(present) if not faint[i]][:4]
+ for i in bright:
+ store.paint_disc(t, pos[i, 0], pos[i, 1], max(1.5, radii[i] * 0.5), 0)
+ store.paint_disc(t, pos[8, 0], pos[8, 1], 1.5, 0)
+
+ far = _clear_of_particles(shape, pos, radii, present, 3.0)
+ for (y0, x0, y1, x1) in ((4, 4, 8, 100), (h - 8, 4, h - 4, 100),
+ (40, 4, 60, 8), (20, 40, 40, 46)):
+ sweep = np.zeros((h, w), bool)
+ sweep[y0:y1, x0:x1] = True
+ store.paint(t, sweep & far, 1)
+
+ touching = _clear_of_particles(shape, pos, radii, present, 0.0)
+ for i in bright:
+ d2 = (yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2
+ ring = (d2 > (radii[i] + 2.5) ** 2) & (d2 <= (radii[i] + 4.0) ** 2)
+ store.paint(t, ring & touching, 1)
+
+ if seam:
+ idx = list(np.flatnonzero(present))
+ grown = [ndi.binary_dilation(
+ ((yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2) <= radii[i] ** 2,
+ iterations=2) for i in idx]
+ edge = np.zeros((h, w), bool)
+ for a in range(len(grown)):
+ for b in range(a + 1, len(grown)):
+ edge |= grown[a] & grown[b]
+ if edge.any():
+ store.paint(t, edge, 3)
+ return store
+
+
+def truth_mask(geom) -> np.ndarray:
+ """Exact foreground at :data:`FRAME_T`.
+
+ ``_soft_disc`` is ``0.5*(1 - tanh((r - radius)/0.9))``, which crosses 0.5
+ exactly at ``r == radius`` — so the analytic disc IS the 0.5-probability
+ contour and the IoU below is against ground truth, not against a rendering
+ of it.
+ """
+ pos, radii, present, _faint, shape = geom
+ h, w = shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ m = np.zeros((h, w), bool)
+ for i in np.flatnonzero(present):
+ m |= ((yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2) <= radii[i] ** 2
+ return m
+
+
+# ── the realistic-scale field (train-time question only) ─────────────────────
+
+def big_field(edge: int = 2048, n_particles: int = 420, seed: int = 0):
+ """A 2048² particle field with a vacuum hole — the realistic train-time case.
+
+ The fixture is 96×112 with nine particles, which is the right size for a
+ ground-truth quality gate and the wrong size for a train-time answer: it
+ yields ONE training crop, so it measures the optimiser loop and nothing
+ about how the cost grows with how much a user painted. This field is built
+ from the same primitives as ``particle_movie`` (a speckled, ramped support
+ film plus soft discs) at a size where scribbles genuinely spread out.
+
+ Discs are drawn in their own bounding boxes, not over the full raster: 420
+ full-frame adds at 2048² is ~7 s of pure setup and would dominate the thing
+ being measured.
+
+ Returns ``(frame, centres, radii, vacuum_rect)``.
+ """
+ from scipy.ndimage import gaussian_filter
+
+ rng = np.random.default_rng(seed)
+ yy = np.mgrid[0:edge, 0:edge][0].astype(np.float32)
+ film = 0.10 + 0.12 * (yy / edge)
+ film += 0.09 * gaussian_filter(rng.standard_normal((edge, edge)).astype(
+ np.float32), 1.6)
+
+ # A hole in the carbon, off-centre and non-square so a transposed or
+ # mirrored result is obvious. This is the "vacuum" class's home.
+ vac = (int(edge * 0.06), int(edge * 0.30), int(edge * 0.62), int(edge * 0.94))
+ film[vac[0]:vac[1], vac[2]:vac[3]] *= 0.05
+
+ centres = rng.uniform(12, edge - 12, size=(n_particles, 2))
+ radii = rng.uniform(5.0, 12.0, size=n_particles)
+ amps = rng.uniform(0.45, 1.0, size=n_particles)
+ for (cy, cx), r, a in zip(centres, radii, amps):
+ # CLIP the window, and build the coordinate grid from the CLIPPED
+ # bounds. Centres are drawn 12 px from the edge but radii reach 12 with
+ # +4/+5 padding, so a disc can overhang: `film[y0:y1]` then clips
+ # silently while `np.mgrid[y0:y1]` does not, and the add fails with a
+ # (31,32)-vs-(32,32) broadcast error on whichever particle lands there.
+ y0, y1 = max(0, int(cy - r - 4)), min(edge, int(cy + r + 5))
+ x0, x1 = max(0, int(cx - r - 4)), min(edge, int(cx + r + 5))
+ if y1 <= y0 or x1 <= x0:
+ continue
+ gy, gx = np.mgrid[y0:y1, x0:x1].astype(np.float32)
+ d = np.sqrt((gy - cy) ** 2 + (gx - cx) ** 2)
+ film[y0:y1, x0:x1] += a * 0.5 * (1.0 - np.tanh((d - r) / 0.9))
+ film += (0.015 * rng.standard_normal((edge, edge))).astype(np.float32)
+ return film, centres, radii, vac
+
+
+def paint_big_scribbles(shape, centres, radii, vac, *,
+ want=(16_259, 32_651, 1_816), seed: int = 0) -> LabelStore:
+ """Scribbles matching a REAL session's per-class pixel counts.
+
+ *want* is ``(particle, support film, vacuum)`` and defaults to the counts
+ read off an actual SpyDE session. Matching the counts is the point: train
+ time is a function of how much was painted and how far apart, and a
+ fixture-sized scribble set answers a question nobody asked.
+
+ Three classes and no boundary — the real session had none, and adding one
+ here would inflate the crop count against a scenario that did not have it.
+ """
+ rng = np.random.default_rng(seed)
+ h, w = shape
+ store = LabelStore(frame_shape=shape, classes=default_classes())
+
+ order = rng.permutation(len(centres))
+ got = 0
+ for i in order:
+ if got >= want[0]:
+ break
+ got = store.paint_disc(0, centres[i, 0], centres[i, 1],
+ radii[i] * 0.85, 0)
+
+ # Film sweeps: wide strokes at random places, skipping any pixel within a
+ # particle radius, until the target count is reached.
+ keep = np.ones((h, w), bool)
+ for (cy, cx), r in zip(centres, radii):
+ y0, y1 = max(0, int(cy - r - 4)), min(h, int(cy + r + 5))
+ x0, x1 = max(0, int(cx - r - 4)), min(w, int(cx + r + 5))
+ gy, gx = np.mgrid[y0:y1, x0:x1]
+ keep[y0:y1, x0:x1] &= ((gy - cy) ** 2 + (gx - cx) ** 2) > (r + 3.0) ** 2
+ keep[vac[0]:vac[1], vac[2]:vac[3]] = False # that is vacuum, not film
+
+ got = 0
+ while got < want[1]:
+ y0 = int(rng.integers(0, h - 40))
+ x0 = int(rng.integers(0, w - 240))
+ sweep = np.zeros((h, w), bool)
+ sweep[y0:y0 + 6, x0:x0 + 240] = True
+ got = store.paint(0, sweep & keep, 1)
+
+ side = int(np.sqrt(want[2]))
+ vy, vx = (vac[0] + vac[1]) // 2, (vac[2] + vac[3]) // 2
+ box = np.zeros((h, w), bool)
+ box[vy - side // 2:vy + side // 2, vx - side // 2:vx + side // 2] = True
+ store.paint(0, box, 2)
+ return store
+
+
+# ── scoring ──────────────────────────────────────────────────────────────────
+
+def evaluate(engine, movie_data, gt, geom) -> dict:
+ """Score a trained engine against exact ground truth. Engine-agnostic.
+
+ Both engines expose ``predict_foreground_boundary`` and ``segment``, which
+ is the design claim being tested — so this function cannot tell them apart,
+ and neither can :func:`spyde.particles.classical.split_instances`.
+ """
+ pos, radii, present, faint, shape = geom
+ out: dict = {}
+
+ t0 = time.perf_counter()
+ fg, bnd = engine.predict_foreground_boundary(movie_data[FRAME_T])
+ out["predict_s"] = time.perf_counter() - t0
+ out["has_boundary"] = bnd is not None
+
+ truth = truth_mask(geom)
+ pred = fg > 0.5
+ inter = int((pred & truth).sum())
+ union = int((pred | truth).sum())
+ out["iou"] = inter / max(1, union)
+ out["truth_px"] = int(truth.sum())
+ out["pred_px"] = int(pred.sum())
+
+ labels = split_instances(fg, SPLIT, boundary=bnd)
+ out["n_particles"] = int(labels.max())
+ out["n_truth"] = int(present.sum())
+
+ def _hit(lab, i):
+ return int(lab[int(round(pos[i, 0])), int(round(pos[i, 1]))]) != 0
+
+ out["faint_found"] = [int(i) for i in np.flatnonzero(faint) if _hit(labels, i)]
+ out["faint_total"] = int(faint.sum())
+ out["bright_missed"] = [int(i) for i in np.flatnonzero(present & ~faint)
+ if not _hit(labels, i)]
+
+ # The merge frame: the pair deliberately overlap, and a boundary-trained
+ # head is supposed to keep them apart where the watershed has to guess.
+ mt = int(gt["merge_frame"])
+ out["merge_frame"] = mt
+ if mt >= 0:
+ mpos, _r, mpresent = particle_truth_at(gt, mt)
+ mfg, mbnd = engine.predict_foreground_boundary(movie_data[mt])
+ mlab = split_instances(mfg, SPLIT, boundary=mbnd)
+ a, b = MERGE_PAIR
+ la = int(mlab[int(round(mpos[a, 0])), int(round(mpos[a, 1]))])
+ lb = int(mlab[int(round(mpos[b, 0])), int(round(mpos[b, 1]))])
+ out["merge_labels"] = [la, lb]
+ out["merge_split"] = bool(la and lb and la != lb)
+ out["merge_frame_particles"] = int(mlab.max())
+ out["merge_frame_truth"] = int(mpresent.sum())
+ return out
+
+
+def _fmt(tag: str, rep: dict, sc: dict) -> str:
+ return (f" {tag:<26} train {rep['total_s']:7.2f}s | "
+ f"IoU {sc['iou']:.3f} | n={sc['n_particles']:>3}/{sc['n_truth']} | "
+ f"faint {len(sc['faint_found'])}/{sc['faint_total']} | "
+ f"merge-split {str(sc.get('merge_split')):<5} | "
+ f"predict {sc['predict_s']*1000:6.0f} ms")
+
+
+# ── sections ─────────────────────────────────────────────────────────────────
+
+def section_forward(device) -> dict:
+ """4096² forward-pass cost and peak VRAM, both configs, tiled and not.
+
+ Confirms the numbers the prototype was designed around and records the peak
+ allocation, which is the thing that actually decides whether the whole-frame
+ path is usable on a 12 GB card.
+ """
+ import torch
+
+ from spyde.device_lock import accelerator_lock
+
+ print("\n=== forward pass, 4096x4096, fp32 ===")
+ out: dict = {}
+ if device.type != "cuda":
+ print(" (skipped — no CUDA device)")
+ return out
+
+ frame = np.random.default_rng(0).standard_normal(
+ (4096, 4096)).astype(np.float32)
+ for name, (base, levels) in CONFIGS.items():
+ net = build_net(3, base=base, levels=levels).to(device).eval()
+ params = sum(p.numel() for p in net.parameters())
+ for mode in ("tiled-1024", "whole"):
+ torch.cuda.empty_cache()
+ torch.cuda.reset_peak_memory_stats()
+ clf = ScribbleCNN(base=base, levels=levels, device=device,
+ tile=1024 if mode == "tiled-1024" else 8192)
+ clf._net, clf.classes = net, []
+ try:
+ with accelerator_lock(device), torch.no_grad():
+ for run in range(2): # discard the cold run
+ torch.cuda.synchronize()
+ t0 = time.perf_counter()
+ h, w = frame.shape
+ for (y0, y1, x0, x1, *_r) in clf._tiles(h, w):
+ t = torch.as_tensor(frame[y0:y1, x0:x1],
+ device=device)[None, None]
+ torch.softmax(net(t), dim=1)
+ torch.cuda.synchronize()
+ dt = time.perf_counter() - t0
+ peak = torch.cuda.max_memory_allocated() / 2 ** 20
+ out[f"{name}/{mode}"] = {"s": dt, "peak_MiB": peak,
+ "params": params}
+ print(f" {name:<6} base={base} levels={levels} "
+ f"{params/1e3:7.1f}k params {mode:<11} "
+ f"{dt:6.3f} s peak {peak:7.0f} MiB")
+ except torch.cuda.OutOfMemoryError as e:
+ out[f"{name}/{mode}"] = {"error": "OOM"}
+ print(f" {name:<6} {mode:<11} OOM ({str(e)[:60]})")
+ finally:
+ torch.cuda.empty_cache()
+ del net
+ return out
+
+
+def big_truth(shape, centres, radii) -> np.ndarray:
+ """Exact foreground for :func:`big_field`, drawn in bounding boxes only.
+
+ Same 0.5-contour identity as :func:`truth_mask` — ``_soft_disc`` crosses 0.5
+ at ``r == radius`` — so this is ground truth, not a rendering of it.
+ """
+ m = np.zeros(shape, bool)
+ for (cy, cx), r in zip(centres, radii):
+ y0, y1 = max(0, int(cy - r - 2)), min(shape[0], int(cy + r + 3))
+ x0, x1 = max(0, int(cx - r - 2)), min(shape[1], int(cx + r + 3))
+ gy, gx = np.mgrid[y0:y1, x0:x1]
+ m[y0:y1, x0:x1] |= ((gy - cy) ** 2 + (gx - cx) ** 2) <= r * r
+ return m
+
+
+def section_train(device, steps: int) -> dict:
+ """Train time at fixture scale AND at realistic scale, MLP vs both CNNs.
+
+ The realistic block also scores IoU and particle count against the field's
+ own exact truth. That is not scope creep: the fixture gives the CNN 1.4 k
+ labelled pixels in ONE crop, which is a regime where a conv net has almost
+ nothing to learn from, so a bad fixture score does not by itself say the
+ approach is bad. The 2048² field with a real session's ~50 k labels across
+ ~50 crops is the regime the user actually paints in, and it costs nothing
+ extra to score the models that were trained here anyway.
+ """
+ print("\n=== train time ===")
+ out: dict = {}
+
+ s = particle_movie()
+ gt = ground_truth(s)
+ pos, radii, present = particle_truth_at(gt, FRAME_T)
+ geom = (pos, radii, present, np.asarray(gt["p_faint"], bool),
+ tuple(gt["frame_shape"]))
+ store = paint_fixture_scribbles(geom)
+ frames = {FRAME_T: s.data[FRAME_T]}
+ print(f" fixture 96x112, {len(store)} labelled px, {store.counts()}")
+ out["fixture"] = _train_row(store, frames, device, steps)
+
+ t0 = time.perf_counter()
+ big, centres, big_radii, vac = big_field()
+ bstore = paint_big_scribbles(big.shape, centres, big_radii, vac)
+ truth = big_truth(big.shape, centres, big_radii)
+ print(f"\n realistic 2048x2048 ({time.perf_counter() - t0:.1f} s to build),"
+ f" {len(bstore)} labelled px, {bstore.counts()}")
+ out["realistic"] = _train_row(bstore, {0: big}, device, steps,
+ truth=truth, frame=big)
+ return out
+
+
+def _train_row(store, frames, device, steps, *, truth=None, frame=None) -> dict:
+ from scipy import ndimage as ndi
+
+ row: dict = {}
+ n_truth = int(ndi.label(truth)[1]) if truth is not None else 0
+
+ def score_into(cell: dict, engine) -> str:
+ if truth is None:
+ return ""
+ t0 = time.perf_counter()
+ fg, bnd = engine.predict_foreground_boundary(frame)
+ dt = time.perf_counter() - t0
+ pred = fg > 0.5
+ iou = int((pred & truth).sum()) / max(1, int((pred | truth).sum()))
+ n = int(split_instances(fg, SPLIT, boundary=bnd).max())
+ cell.update(iou=iou, n_particles=n, n_truth=n_truth, predict_s=dt)
+ return f" | IoU {iou:.3f} n={n}/{n_truth} predict {dt:5.2f} s"
+
+ mlp = ScribbleClassifier(FeatureSpec(), device=device, seed=0)
+ t0 = time.perf_counter()
+ rep = mlp.fit(store, frames)
+ cell = row["mlp"] = {"total_s": time.perf_counter() - t0,
+ "featurise_s": rep["featurise_s"],
+ "fit_s": rep["fit_s"], "n_pixels": rep["n_pixels"]}
+ tail = score_into(cell, mlp)
+ print(f" MLP (36ch + head) {cell['total_s']:7.2f} s "
+ f"(featurise {rep['featurise_s']:.2f} + fit {rep['fit_s']:.2f})"
+ f"{tail}")
+
+ for name, (base, levels) in CONFIGS.items():
+ cnn = ScribbleCNN(base=base, levels=levels, steps=steps, device=device,
+ seed=0)
+ t0 = time.perf_counter()
+ rep = cnn.fit(store, frames)
+ cell = row[name] = {"total_s": time.perf_counter() - t0,
+ "crops_s": rep["crops_s"], "fit_s": rep["fit_s"],
+ "n_crops": rep["n_crops"], "steps": steps,
+ "params": rep["params"],
+ "train_accuracy": rep["train_accuracy"]}
+ tail = score_into(cell, cnn)
+ print(f" CNN {name:<6} b{base}/L{levels} {cell['total_s']:7.2f}"
+ f" s (crops {rep['crops_s']:.2f} + fit {rep['fit_s']:.2f}), "
+ f"{rep['n_crops']} crops, {steps} steps, "
+ f"acc {rep['train_accuracy']:.3f}{tail}")
+ return row
+
+
+def section_quality(device, steps: int) -> dict:
+ """Both engines, same scribbles, scored against exact truth.
+
+ Run TWICE, with and without the seam strokes, because the two exercise
+ different downstream routes and only one of them is what the shipped §0.9
+ gate measures. Without a boundary both engines go through the watershed;
+ with one they both take ``split_instances``' connected-components route. An
+ A/B that reported only the boundary route would be comparing the CNN's
+ boundary head against the MLP's, and an A/B that reported only the watershed
+ route would never test the merge split at all.
+ """
+ s = particle_movie()
+ gt = ground_truth(s)
+ pos, radii, present = particle_truth_at(gt, FRAME_T)
+ geom = (pos, radii, present, np.asarray(gt["p_faint"], bool),
+ tuple(gt["frame_shape"]))
+ frames = {FRAME_T: s.data[FRAME_T]}
+ out: dict = {}
+
+ for seam in (False, True):
+ route = "boundary route" if seam else "watershed route"
+ print(f"\n=== quality vs ground truth (frame 12, {route}) ===")
+ store = paint_fixture_scribbles(geom, seam=seam)
+ print(f" labels: {store.counts()}")
+ block: dict = {}
+
+ mlp = ScribbleClassifier(FeatureSpec(), device=device, seed=0)
+ t0 = time.perf_counter()
+ rep = mlp.fit(store, frames)
+ rep["total_s"] = time.perf_counter() - t0
+ sc = evaluate(mlp, s.data, gt, geom)
+ block["mlp"] = {"train": rep, "score": sc}
+ print(_fmt("MLP (36ch + head)", rep, sc))
+
+ for name, (base, levels) in CONFIGS.items():
+ cnn = ScribbleCNN(base=base, levels=levels, steps=steps,
+ device=device, seed=0)
+ t0 = time.perf_counter()
+ rep = cnn.fit(store, frames)
+ rep["total_s"] = time.perf_counter() - t0
+ sc = evaluate(cnn, s.data, gt, geom)
+ block[name] = {"train": rep, "score": sc}
+ print(_fmt(f"CNN {name} b{base}/L{levels}", rep, sc))
+ out["seam" if seam else "no_seam"] = block
+ return out
+
+
+def section_curve(device, sweep) -> dict:
+ """Accuracy-vs-time: is there a knee worth stopping at?"""
+ print("\n=== accuracy vs train time (CNN, fixture) ===")
+ s = particle_movie()
+ gt = ground_truth(s)
+ pos, radii, present = particle_truth_at(gt, FRAME_T)
+ geom = (pos, radii, present, np.asarray(gt["p_faint"], bool),
+ tuple(gt["frame_shape"]))
+ store = paint_fixture_scribbles(geom)
+ frames = {FRAME_T: s.data[FRAME_T]}
+
+ out: dict = {}
+ for name, (base, levels) in CONFIGS.items():
+ rows = []
+ for steps in sweep:
+ cnn = ScribbleCNN(base=base, levels=levels, steps=steps,
+ device=device, seed=0)
+ t0 = time.perf_counter()
+ cnn.fit(store, frames)
+ total = time.perf_counter() - t0
+ sc = evaluate(cnn, s.data, gt, geom)
+ rows.append({"steps": steps, "total_s": total, "iou": sc["iou"],
+ "n_particles": sc["n_particles"],
+ "faint": len(sc["faint_found"]),
+ "merge_split": sc.get("merge_split")})
+ print(f" {name:<6} steps {steps:>4} {total:6.2f} s "
+ f"IoU {sc['iou']:.3f} n={sc['n_particles']:>3} "
+ f"faint {len(sc['faint_found'])}/{sc['faint_total']} "
+ f"merge-split {sc.get('merge_split')}")
+ out[name] = rows
+ return out
+
+
+# ── entry point ──────────────────────────────────────────────────────────────
+
+def _warm(device) -> None:
+ """Pay the one-time CUDA costs BEFORE anything is timed.
+
+ Cold CUDA context creation, cuDNN algorithm selection and the feature
+ stack's first kernel launches are a ~1.5 s one-off that lands entirely on
+ whichever engine happens to run first — which measured as the MLP's
+ "featurise 1.47 s" on a 96×112 frame, three orders of magnitude above its
+ warm cost. Warming both stacks here is the difference between an A/B and a
+ coin toss about ordering.
+ """
+ import torch
+ import torch.nn.functional as F
+
+ from spyde.particles.features import sample_features
+
+ if device.type == "cuda":
+ torch.cuda.init()
+ net = build_net(3, base=16, levels=2).to(device)
+ x = torch.zeros((2, 1, 64, 64), device=device)
+ y = torch.zeros((2, 64, 64), dtype=torch.long, device=device)
+ F.cross_entropy(net(x), y).backward()
+ sample_features(np.zeros((64, 64), np.float32), np.arange(64),
+ FeatureSpec(), device=device)
+ if device.type == "cuda":
+ torch.cuda.synchronize()
+ torch.cuda.empty_cache()
+
+
+def main(argv=None) -> None:
+ ap = argparse.ArgumentParser(description=__doc__)
+ ap.add_argument("--only", default="all",
+ choices=("all", "forward", "train", "quality", "curve"))
+ ap.add_argument("--device", default=None)
+ ap.add_argument("--steps", type=int, default=300)
+ ap.add_argument("--sweep", default="50,100,200,300,600")
+ ap.add_argument("--json", default=None, help="write the JSON blob here too")
+ args = ap.parse_args(argv)
+
+ device = select_device(args.device)
+ print(f"device: {device}")
+ _warm(device)
+ results: dict = {"device": str(device), "steps": args.steps}
+
+ if args.only in ("all", "forward"):
+ results["forward"] = section_forward(device)
+ if args.only in ("all", "train"):
+ results["train"] = section_train(device, args.steps)
+ if args.only in ("all", "quality"):
+ results["quality"] = section_quality(device, args.steps)
+ if args.only in ("all", "curve"):
+ results["curve"] = section_curve(
+ device, [int(v) for v in args.sweep.split(",")])
+
+ blob = json.dumps(results, indent=2, default=float)
+ print("\n=== JSON ===")
+ print(blob)
+ if args.json:
+ with open(args.json, "w", encoding="utf-8") as fh:
+ fh.write(blob)
+ sys.stdout.flush()
+ # torch's CUDA teardown crashes on exit here (CLAUDE.md); the numbers are
+ # already printed, so leave before it runs.
+ os._exit(0)
+
+
+if __name__ == "__main__":
+ main()
diff --git a/spyde/tests/migrated/test_batch_feedback.py b/spyde/tests/migrated/test_batch_feedback.py
new file mode 100644
index 00000000..d6155c9f
--- /dev/null
+++ b/spyde/tests/migrated/test_batch_feedback.py
@@ -0,0 +1,191 @@
+"""test_batch_feedback.py — the standard "this is computing" surface.
+
+``lifecycle.batch_feedback`` bundles the three things a long batch run has to
+show: the window's Calculating overlay, rate-limited status-bar progress, and
+the most recent RESULT painted live. The third is the one that matters most and
+is the easiest to get subtly wrong — a progress bar advancing through a run that
+is quietly finding nothing looks exactly like one that is working.
+
+These pin the behaviours a caller depends on:
+ * the overlay is ALWAYS paired, including when the batch raises;
+ * progress is rate-limited, but the FINAL result is never dropped;
+ * the paint is marshalled to the main thread (CLAUDE.md's threading rule);
+ * a failing paint cannot take the batch down.
+"""
+from __future__ import annotations
+
+import time
+
+import pytest
+
+from spyde.actions import lifecycle
+
+
+class _Session:
+ """Records what got marshalled instead of running a loop."""
+
+ def __init__(self, run=True):
+ self.dispatched = []
+ self._run = run
+
+ def _dispatch_to_main(self, fn):
+ self.dispatched.append(fn)
+ if self._run:
+ fn()
+
+
+@pytest.fixture
+def emitted(monkeypatch):
+ """Capture ipc emissions at the module the helper imports them from."""
+ out = {"progress": [], "computing": []}
+ import spyde.backend.ipc as ipc
+ monkeypatch.setattr(ipc, "emit_progress",
+ lambda d, t, label="": out["progress"].append((d, t, label)))
+ monkeypatch.setattr(ipc, "emit_window_computing",
+ lambda wid, on: out["computing"].append((wid, bool(on))))
+ return out
+
+
+class TestOverlayPairing:
+ def test_start_and_stop_bracket_the_run(self, emitted):
+ with lifecycle.batch_feedback(_Session(), 7, "Segmenting", 3):
+ pass
+ assert emitted["computing"] == [(7, True), (7, False)]
+
+ def test_the_overlay_stops_even_when_the_batch_raises(self, emitted):
+ """The whole point of the pairing contract: a failed run must not leave
+ a Calculating chip spinning over a window that has stopped working."""
+ with pytest.raises(RuntimeError):
+ with lifecycle.batch_feedback(_Session(), 7, "Segmenting", 3):
+ raise RuntimeError("boom")
+ assert emitted["computing"] == [(7, True), (7, False)]
+
+ def test_no_window_sends_no_message(self, monkeypatch):
+ """A plot with no window yet must not put a malformed message on the
+ wire. The guard lives in ``emit_window_computing`` and is DELEGATED to,
+ so this has to capture at ``ipc.emit`` — stubbing
+ ``emit_window_computing`` (as the other tests here do) replaces the very
+ guard under test and the assertion becomes vacuous.
+ """
+ import spyde.backend.ipc as ipc
+ sent = []
+ monkeypatch.setattr(ipc, "emit", lambda msg: sent.append(msg))
+ with lifecycle.batch_feedback(_Session(), None, "Segmenting", 3):
+ pass
+ assert [m for m in sent if m.get("type") == "window_computing"] == []
+
+ def test_a_real_window_id_does_send(self, monkeypatch):
+ """The other half of the guard — otherwise "sends nothing" would pass
+ for a helper that never emits at all."""
+ import spyde.backend.ipc as ipc
+ sent = []
+ monkeypatch.setattr(ipc, "emit", lambda msg: sent.append(msg))
+ with lifecycle.batch_feedback(_Session(), 7, "Segmenting", 3):
+ pass
+ assert [(m["window_id"], m["computing"]) for m in sent
+ if m.get("type") == "window_computing"] == [(7, True), (7, False)]
+
+
+class TestProgress:
+ def test_rate_limited_between_steps(self, emitted):
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 100, min_interval=60.0)
+ fb.step(1)
+ fb.step(2)
+ fb.step(3)
+ # Only the first got through; the rest are inside the interval. A
+ # 900-frame run emitting per frame would put 900 messages on the same
+ # stdout line protocol the nav painter uses.
+ assert emitted["progress"] == [(1, 100, "Seg")]
+
+ def test_force_defeats_the_rate_limit(self, emitted):
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 100, min_interval=60.0)
+ fb.step(1)
+ fb.step(2, force=True)
+ assert emitted["progress"] == [(1, 100, "Seg"), (2, 100, "Seg")]
+
+ def test_the_last_step_is_never_rate_limited(self, emitted):
+ """`done == total` always emits, so the bar reaches the end."""
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 3, min_interval=60.0)
+ fb.step(1)
+ fb.step(2)
+ fb.step(3)
+ assert emitted["progress"][-1] == (3, 3, "Seg")
+
+ def test_finish_emits_a_terminal_tick(self, emitted):
+ """Without it an early-finishing run leaves the spinner mid-way and a
+ finished job reads as hung."""
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 10, min_interval=60.0)
+ fb.step(4)
+ fb.finish()
+ assert emitted["progress"][-1] == (10, 10, "Seg")
+
+
+class TestLiveResult:
+ def test_the_result_is_published(self, emitted):
+ seen = []
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 3, publish=seen.append)
+ fb.step(1, result="frame-1")
+ assert seen == ["frame-1"]
+
+ def test_the_paint_is_marshalled_to_the_main_thread(self, emitted):
+ """Figure updates must not happen on the worker thread that computed
+ them — CLAUDE.md's threading contract."""
+ session = _Session(run=False)
+ seen = []
+ fb = lifecycle.batch_feedback(session, 1, "Seg", 3, publish=seen.append)
+ fb.step(1, result="frame-1")
+ assert seen == [], "painted inline instead of marshalling"
+ assert len(session.dispatched) == 1
+ session.dispatched[0]()
+ assert seen == ["frame-1"]
+
+ def test_a_rate_limited_step_publishes_nothing(self, emitted):
+ seen = []
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 100,
+ publish=seen.append, min_interval=60.0)
+ fb.step(1, result="a")
+ fb.step(2, result="b")
+ assert seen == ["a"]
+
+ def test_the_final_result_survives_the_rate_limit(self, emitted):
+ """The last frame is the one left on screen, so it must never be the
+ one the rate limiter drops."""
+ seen = []
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 3,
+ publish=seen.append, min_interval=60.0)
+ fb.step(1, result="a")
+ fb.step(2, result="b")
+ fb.step(3, result="c")
+ assert seen[-1] == "c"
+
+ def test_a_failing_paint_does_not_break_the_batch(self, emitted):
+ def boom(_r):
+ raise ValueError("bad frame")
+
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 3, publish=boom)
+ fb.step(1, result="a") # must not raise
+ assert emitted["progress"] == [(1, 3, "Seg")]
+
+ def test_step_without_a_result_still_reports_progress(self, emitted):
+ seen = []
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 3, publish=seen.append)
+ fb.step(1)
+ assert seen == []
+ assert emitted["progress"] == [(1, 3, "Seg")]
+
+
+class TestRealClock:
+ def test_the_interval_actually_elapses(self, emitted):
+ """The limiter is wall-clock, not a call counter.
+
+ The margin is deliberately enormous (10x) rather than just over the
+ interval: Windows' ``time.monotonic`` granularity is ~15.6 ms, so a
+ sleep(0.06) against a 0.05 interval can measure as 0.047 and drop the
+ second emission. That version passed alone and failed in the file — a
+ real flake, not a real bug.
+ """
+ fb = lifecycle.batch_feedback(_Session(), 1, "Seg", 100, min_interval=0.02)
+ fb.step(1)
+ time.sleep(0.2)
+ fb.step(2)
+ assert emitted["progress"] == [(1, 100, "Seg"), (2, 100, "Seg")]
diff --git a/spyde/tests/migrated/test_chunk_dispatch_guard.py b/spyde/tests/migrated/test_chunk_dispatch_guard.py
index 2d8778fb..57b32c58 100644
--- a/spyde/tests/migrated/test_chunk_dispatch_guard.py
+++ b/spyde/tests/migrated/test_chunk_dispatch_guard.py
@@ -249,18 +249,43 @@ def _wait(pred, timeout=10.0, required=True):
class TestNavigatorFillThroughDispatcher:
- def test_batched_submit_and_bounded_window(self):
- """16 one-frame nav chunks on a 4-thread cluster: window =
- max(4, 4 // 2) = 4, and each top-up is ONE submit call."""
+ def test_batched_submit_prime_then_bulk(self):
+ """16 one-frame nav chunks: a small priming batch, then the rest — and
+ only a couple of scheduler round trips in total.
+
+ This test used to assert ``max_in_flight <= 4`` as BACKPRESSURE, on the
+ reasoning that the navigator fill must not "submit all 977 up front".
+ That guard is retired deliberately, so here is the argument against my
+ own fence:
+
+ * It did not measure as backpressure. benchmarks.md has the bounded
+ window at 46-50 s against unbounded at 50.2 s on the same 977-chunk
+ movie — within noise. The window never bought throughput or safety
+ that anyone demonstrated.
+ * ``distributed`` >= 2022.3 QUEUES root tasks at the scheduler
+ (worker-saturation), which is exactly this guard's job, done in the
+ right process, without holding OUR GIL. We were duplicating it badly.
+ * Keeping it was actively harmful: ``_submit_next`` runs on every
+ completion, so ``lane_cap - outstanding`` was 1 in steady state and
+ the batch collapsed to ONE task per submit — 970 blocking round trips
+ for 977 chunks, which scaled with dataset size and read to users as
+ "big datasets only use one worker".
+
+ What replaces it: the assertion below that the whole job goes out in a
+ couple of submits, and the scheduler's own queuing for memory. The
+ DUAL-LANE path keeps its window (see the module docstring) because
+ there a completion genuinely pulls the next chunk so a ~30x-faster GPU
+ lane and the CPU lane finish together — that is real work stealing the
+ scheduler cannot do.
+ """
src, client, handle, seen, futs, max_in_flight = _run_nav_fill(
nav=(16,), chunks=(1,))
assert len(seen) == 16 # every chunk streamed
- assert len(client.created) == 16
- # The old loop made 16 submit calls (+1 for the duplicate whole-array
- # graph) — one blocking scheduler round trip each.
- assert client.submit_calls < 16
- # BACKPRESSURE: the navigator fill used to submit all 977 up front.
- assert max_in_flight <= 4
+ assert len(client.created) == 16 # one future per chunk
+ # The bug this replaced made ~one submit per chunk.
+ assert client.submit_calls <= 3, (
+ f"{client.submit_calls} submits for 16 chunks — the unpinned lane "
+ f"is back to per-completion top-ups")
np.testing.assert_array_equal(handle.result(), src.compute())
def test_result_is_assembled_not_recomputed(self):
diff --git a/spyde/tests/migrated/test_device_lock.py b/spyde/tests/migrated/test_device_lock.py
index eecd3876..173827f2 100644
--- a/spyde/tests/migrated/test_device_lock.py
+++ b/spyde/tests/migrated/test_device_lock.py
@@ -277,3 +277,55 @@ def fake_calibrate(model, frames, device, **kw):
monkeypatch.setattr(models, "calibrate", fake_calibrate)
fvn.calibrate_neural([np.zeros((32, 32), dtype=np.float32)])
assert held == [True], "calibration forwards ran unserialised on MPS"
+
+
+class TestParticleEnginesTakeLock:
+ """The particle feature stack and the scribble head are the newest torch users.
+
+ They run from BOTH ends of the concurrency this lock exists for: a batch
+ segmentation run on a worker thread, and the interactive re-apply that fires on
+ every navigator move while the user tunes. That overlap is exactly the pair of
+ threads in the crash stacks at the top of this file.
+
+ These two pin the *real* lock. Finer-grained nesting — that no submission
+ inside ``fit`` / ``predict`` / ``save`` / ``load`` escapes the block, which a
+ null context on a CPU box cannot show — is pinned in
+ ``test_particles_scribble.py::TestDeviceLock``.
+ """
+
+ def test_feature_stack_locks(self, monkeypatch):
+ """``map_feature_bands`` is the single torch entry point in features.py:
+ ``feature_tensor``, ``feature_stack`` and ``sample_features`` all route
+ through it, so this one acquisition covers every channel."""
+ from spyde.particles import features as feat
+
+ held = []
+ monkeypatch.setattr(feat, "_band_stack",
+ lambda *a, **k: held.append(_lock_is_held_by_me()))
+ feat.map_feature_bands(np.zeros((8, 8), dtype=np.float32),
+ feat.FeatureSpec(), device=_FakeDev(),
+ fn=lambda *a: None)
+ assert held == [True], "the particle feature stack ran unserialised on MPS"
+ assert not _lock_is_held_by_me(), "lock leaked"
+
+ def test_scribble_training_locks(self, monkeypatch):
+ from spyde.particles import scribble as scr
+
+ class _Stop(Exception):
+ """Aborts the fit at its first device submission."""
+
+ held = []
+
+ def spy(*a, **kw):
+ held.append(_lock_is_held_by_me())
+ raise _Stop
+
+ monkeypatch.setattr(scr, "sample_features", spy)
+ store = scr.LabelStore(frame_shape=(16, 16))
+ store.paint(0, [(2, 2)], 0)
+ store.paint(0, [(9, 9)], 1)
+ clf = scr.ScribbleClassifier(device=_FakeDev())
+ with pytest.raises(_Stop):
+ clf.fit(store, {0: np.zeros((16, 16), dtype=np.float32)})
+ assert held == [True], "scribble training ran unserialised on MPS"
+ assert not _lock_is_held_by_me(), "lock leaked"
diff --git a/spyde/tests/migrated/test_drift_nonrigid.py b/spyde/tests/migrated/test_drift_nonrigid.py
new file mode 100644
index 00000000..dbc17dcc
--- /dev/null
+++ b/spyde/tests/migrated/test_drift_nonrigid.py
@@ -0,0 +1,240 @@
+"""
+test_drift_nonrigid.py — the non-rigid solve recovers a KNOWN warp.
+
+The plan's acceptance criterion for A2-A5 is exactly that: apply a synthetic
+distortion whose field is known exactly, fit it, and check the fit removes it.
+Anything weaker (loss went down, the field is smooth, it ran without raising)
+would pass with a solver that fits noise.
+
+Ground truth is built by warping with the SAME resampler the solver uses, so the
+test measures the SOLVER and not the difference between two interpolators. The
+residual is then compared against the distorted-vs-reference residual, i.e. "did
+it recover most of what we put in", which is scale-free and does not encode a
+tolerance nobody can justify.
+
+CPU on purpose: torch-CUDA under the pytest process segfaults on Windows
+(CLAUDE.md), and these fits are tiny.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+torch = pytest.importorskip("torch")
+
+from spyde.drift import nonrigid as nr
+from spyde.drift.model import DriftModel
+
+
+H = W = 64
+DEV = "cpu"
+
+
+def _textured_frame(seed: int = 0) -> np.ndarray:
+ """A frame with structure at several scales.
+
+ Registration needs gradients everywhere: a field of one blob is happy to
+ slide sideways, so a solver could score well while recovering the wrong
+ field. Blobs plus fine noise gives an unambiguous optimum.
+ """
+ rng = np.random.default_rng(seed)
+ y, x = np.mgrid[0:H, 0:W].astype(np.float32)
+ img = np.zeros((H, W), np.float32)
+ for _ in range(14):
+ cy, cx = rng.uniform(6, H - 6), rng.uniform(6, W - 6)
+ s = rng.uniform(2.0, 4.5)
+ img += rng.uniform(0.5, 1.5) * np.exp(-((y - cy) ** 2 + (x - cx) ** 2) / (2 * s * s))
+ img += 0.05 * rng.standard_normal((H, W)).astype(np.float32)
+ return img
+
+
+def _warp_np(frame: np.ndarray, dy: np.ndarray, dx: np.ndarray) -> np.ndarray:
+ t = torch.as_tensor(np.asarray(frame, np.float32))[None]
+ out = nr.warp_frame(torch, t,
+ torch.as_tensor(np.asarray(dy, np.float32))[None],
+ torch.as_tensor(np.asarray(dx, np.float32))[None],
+ fill_nan=False)
+ return out[0].numpy().copy()
+
+
+def _mse(a: np.ndarray, b: np.ndarray) -> float:
+ a = (a - a.mean()) / max(a.std(), 1e-6)
+ b = (b - b.mean()) / max(b.std(), 1e-6)
+ return float(np.mean((a - b) ** 2))
+
+
+class TestWarp:
+ def test_zero_displacement_is_identity(self):
+ f = _textured_frame()
+ out = _warp_np(f, np.zeros((H, W), np.float32), np.zeros((H, W), np.float32))
+ assert np.allclose(out, f, atol=1e-5)
+
+ def test_out_of_bounds_is_nan_not_zero(self):
+ """The locked edge policy: uncovered pixels are NaN, never invented.
+
+ Zero-filling here is the bug that nucleates a spurious edge 'particle'
+ downstream, which is why this is asserted rather than assumed.
+ """
+ f = _textured_frame()
+ dy = np.full((H, W), 40.0, np.float32) # push most rows off the top
+ t = torch.as_tensor(f)[None]
+ out = nr.warp_frame(torch, t, torch.as_tensor(dy)[None],
+ torch.as_tensor(np.zeros((H, W), np.float32))[None],
+ fill_nan=True)[0].numpy()
+ assert np.isnan(out).any(), "no NaN produced outside coverage"
+ assert not np.isnan(out).all(), "everything went out of bounds"
+
+ def test_translation_matches_a_known_roll(self):
+ """An integer displacement must reproduce an exact roll."""
+ f = _textured_frame()
+ out = _warp_np(f, np.full((H, W), -3.0, np.float32),
+ np.full((H, W), 0.0, np.float32))
+ want = np.roll(f, 3, axis=0)
+ # Interior only: the roll wraps where the warp runs out of data.
+ assert np.allclose(out[6:-6, 6:-6], want[6:-6, 6:-6], atol=1e-4)
+
+
+class TestScanKnotRecovery:
+ def _stack(self, amp: float = 3.0, n: int = 4):
+ """Frames distorted by a known SLOW-AXIS-varying displacement."""
+ base = _textured_frame()
+ rows = np.linspace(-1.0, 1.0, H, dtype=np.float32)
+ frames, truth = [], []
+ for i in range(n):
+ a = amp * (i + 1) / n
+ dy = np.repeat((a * rows)[:, None], W, axis=1) # varies down rows
+ dx = np.zeros((H, W), np.float32)
+ frames.append(_warp_np(base, dy, dx))
+ truth.append(dy)
+ return base, np.stack(frames), np.stack(truth)
+
+ def test_recovers_most_of_a_known_scan_distortion(self):
+ base, frames, _ = self._stack()
+ model = nr.solve_nonrigid(frames, model=nr.SCAN_KNOT, reference=base,
+ n_knots=3, steps=220, lr=0.35,
+ smooth_weight=0.05, temporal_weight=0.05,
+ device=DEV)
+ assert model.kind == nr.SCAN_KNOT
+ before = np.mean([_mse(f, base) for f in frames])
+ after = np.mean([_mse(np.nan_to_num(nr.apply_nonrigid(f, model, i), nan=0.0)
+ + np.isnan(nr.apply_nonrigid(f, model, i)) * base, base)
+ for i, f in enumerate(frames)])
+ assert after < 0.45 * before, (
+ f"the fit removed too little of the known warp: {before:.4f} -> {after:.4f}")
+
+ def test_the_fitted_field_varies_down_the_slow_axis(self):
+ """A scan-knot fit must not collapse to a constant offset.
+
+ A constant is the degenerate solution that a rigid solve already
+ provides; if that is all this produces, the model is not earning its
+ parameters.
+ """
+ base, frames, _ = self._stack()
+ model = nr.solve_nonrigid(frames, model=nr.SCAN_KNOT, reference=base,
+ n_knots=3, steps=220, lr=0.35,
+ smooth_weight=0.05, temporal_weight=0.05,
+ device=DEV)
+ dy, _dx = nr.displacement_for_frame(model, len(frames) - 1)
+ spread = float(dy[:, 0].max() - dy[:, 0].min())
+ assert spread > 0.5, f"fitted field is nearly constant down rows ({spread:.3f} px)"
+
+ def test_a_row_is_constant_across_the_fast_axis(self):
+ """Physical contract: one row is acquired at one slow coordinate."""
+ base, frames, _ = self._stack()
+ model = nr.solve_nonrigid(frames, model=nr.SCAN_KNOT, reference=base,
+ n_knots=2, steps=40, device=DEV)
+ dy, dx = nr.displacement_for_frame(model, 0)
+ assert np.allclose(dy, dy[:, :1], atol=1e-6)
+ assert np.allclose(dx, dx[:, :1], atol=1e-6)
+
+
+class TestDenseRecovery:
+ def _stack(self, amp: float = 2.5, n: int = 3):
+ """Frames distorted by a known field that varies in BOTH directions.
+
+ Deliberately not expressible by any scan-knot model, so this exercises
+ the case the second parameterisation exists for.
+ """
+ base = _textured_frame(seed=3)
+ y, x = np.mgrid[0:H, 0:W].astype(np.float32)
+ frames = []
+ for i in range(n):
+ a = amp * (i + 1) / n
+ dy = a * np.sin(2 * np.pi * x / W).astype(np.float32)
+ dx = a * np.cos(2 * np.pi * y / H).astype(np.float32)
+ frames.append(_warp_np(base, dy, dx))
+ return base, np.stack(frames)
+
+ def test_recovers_most_of_a_known_2d_deformation(self):
+ base, frames = self._stack()
+ model = nr.solve_nonrigid(frames, model=nr.DENSE, reference=base,
+ grid=(6, 6), steps=260, lr=0.35,
+ smooth_weight=0.02, temporal_weight=0.02,
+ device=DEV)
+ assert model.kind == nr.DENSE
+ before = np.mean([_mse(f, base) for f in frames])
+ after = []
+ for i, f in enumerate(frames):
+ got = nr.apply_nonrigid(f, model, i)
+ m = np.isnan(got)
+ got = np.where(m, base, got) # score coverage only
+ after.append(_mse(got, base))
+ after = float(np.mean(after))
+ assert after < 0.5 * before, (
+ f"the fit removed too little of the known deformation: {before:.4f} -> {after:.4f}")
+
+ def test_dense_field_is_not_forced_constant_across_a_row(self):
+ """The dense model's whole point: it can vary along the fast axis too."""
+ base, frames = self._stack()
+ model = nr.solve_nonrigid(frames, model=nr.DENSE, reference=base,
+ grid=(6, 6), steps=200, lr=0.35,
+ smooth_weight=0.02, temporal_weight=0.02,
+ device=DEV)
+ dy, _ = nr.displacement_for_frame(model, len(frames) - 1)
+ row_spread = float(np.abs(dy[H // 2] - dy[H // 2].mean()).max())
+ assert row_spread > 0.1, f"dense field is constant across a row ({row_spread:.3f})"
+
+
+class TestModelContract:
+ def test_rigid_component_is_preserved(self):
+ base = _textured_frame()
+ frames = np.stack([base, base])
+ rigid = DriftModel(shifts=np.array([[0.0, 0.0], [1.5, -2.0]], np.float32))
+ model = nr.solve_nonrigid(frames, model=nr.SCAN_KNOT, reference=base,
+ rigid=rigid, steps=20, device=DEV)
+ assert np.allclose(model.shifts, rigid.shifts), (
+ "the rigid component must survive the non-rigid fit — a caller that "
+ "applies only `shifts` should still get the rigid answer")
+
+ def test_extra_carries_everything_needed_to_rebuild_the_field(self):
+ base = _textured_frame()
+ frames = np.stack([base, base])
+ model = nr.solve_nonrigid(frames, model=nr.DENSE, reference=base,
+ grid=(3, 3), steps=10, device=DEV)
+ for key in ("params", "field_shape", "grid"):
+ assert key in model.extra, f"extra is missing {key!r}"
+ dy, dx = nr.displacement_for_frame(model, 0)
+ assert dy.shape == (H, W) and dx.shape == (H, W)
+
+ def test_a_rigid_model_is_refused_not_silently_zero(self):
+ rigid = DriftModel(shifts=np.zeros((2, 2), np.float32))
+ with pytest.raises(ValueError, match="not a non-rigid fit"):
+ nr.displacement_for_frame(rigid, 0)
+
+ def test_unknown_model_name_is_refused(self):
+ with pytest.raises(ValueError, match="model must be one of"):
+ nr.solve_nonrigid(np.zeros((2, H, W), np.float32), model="wobble")
+
+ def test_cancel_stops_the_fit(self):
+ base = _textured_frame()
+ frames = np.stack([base, base])
+ calls = {"n": 0}
+
+ def cancel():
+ calls["n"] += 1
+ return calls["n"] > 3
+
+ model = nr.solve_nonrigid(frames, model=nr.SCAN_KNOT, reference=base,
+ steps=500, cancel=cancel, device=DEV)
+ assert calls["n"] <= 6, "cancel was not honoured promptly"
+ assert model.kind == nr.SCAN_KNOT
diff --git a/spyde/tests/migrated/test_drift_translation.py b/spyde/tests/migrated/test_drift_translation.py
new file mode 100644
index 00000000..0553c59f
--- /dev/null
+++ b/spyde/tests/migrated/test_drift_translation.py
@@ -0,0 +1,553 @@
+"""
+Tests for spyde.drift — rigid translation solve, warp, and DriftModel.
+
+The acceptance gate from DRIFT_AND_PARTICLES_PLAN.md is numerical, not
+structural: recover a synthetically applied shift to better than 0.1 px, and
+agree with ``skimage.registration.phase_cross_correlation`` on the same data.
+That is what most of this file asserts.
+
+Qt-free and dask-free — these are pure-compute tests on small synthetic stacks.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.drift import DriftModel, coverage_mask, frame_source, shift_frame
+from spyde.drift.translation import solve_translation
+
+# The solver's own tolerance target. Sub-pixel ground truth on a smooth
+# synthetic scene should land well inside this.
+GATE_PX = 0.1
+
+
+# ── synthetic data ───────────────────────────────────────────────────────────
+
+def _scene(h=96, w=112, seed=3, noise=0.02):
+ """A smooth, asymmetric, non-periodic scene.
+
+ Asymmetric on purpose: a symmetric scene correlates equally well at several
+ offsets, so a sign error or an axis swap would still pass. Non-periodic on
+ purpose: a lattice invites the exact wrong-translation lock that ``max_shift``
+ exists to prevent, which is a separate test.
+
+ ``noise=0`` gives a band-limited scene, needed wherever a test resamples
+ twice — bilinear interpolation legitimately destroys pixel-scale noise, so a
+ round-trip assertion on a noisy scene measures interpolation loss, not the
+ property under test.
+ """
+ rng = np.random.default_rng(seed)
+ yy, xx = np.mgrid[0:h, 0:w].astype(np.float64)
+ img = np.zeros((h, w), dtype=np.float64)
+ # A handful of gaussian blobs at irregular positions and widths.
+ for cy, cx, amp, sig in [
+ (0.28 * h, 0.22 * w, 1.0, 5.0),
+ (0.61 * h, 0.44 * w, 0.7, 8.0),
+ (0.38 * h, 0.73 * w, 0.9, 4.0),
+ (0.79 * h, 0.66 * w, 0.5, 6.5),
+ (0.17 * h, 0.58 * w, 0.6, 3.5),
+ ]:
+ img += amp * np.exp(-((yy - cy) ** 2 + (xx - cx) ** 2) / (2 * sig ** 2))
+ if noise:
+ img += noise * rng.standard_normal((h, w))
+ return img
+
+
+def _shifted_stack(shifts, h=96, w=112, seed=3):
+ """Stack whose frame i is the scene translated by ``-shifts[i]``.
+
+ So the CORRECTION needed for frame i is ``+shifts[i]`` — matching the
+ DriftModel sign convention. Built by Fourier phase ramp so sub-pixel truth is
+ exact rather than interpolated, which keeps the 0.1 px gate meaningful.
+ """
+ base = _scene(h, w, seed)
+ fy = np.fft.fftfreq(h)[:, None]
+ fx = np.fft.fftfreq(w)[None, :]
+ F = np.fft.fft2(base)
+ frames = []
+ for dy, dx in shifts:
+ # Applying -shift here means +shift is the correction.
+ ramp = np.exp(-2j * np.pi * (-dy * fy + -dx * fx))
+ frames.append(np.real(np.fft.ifft2(F * ramp)))
+ return np.stack(frames).astype(np.float32)
+
+
+class TestFrameSource:
+ def test_numpy_stack(self):
+ arr = np.zeros((5, 8, 9), dtype=np.uint16)
+ n, get, shape = frame_source(arr)
+ assert n == 5 and shape == (8, 9)
+ assert get(3).shape == (8, 9)
+
+ def test_sequence_of_frames(self):
+ seq = [np.zeros((4, 6)) for _ in range(3)]
+ n, get, shape = frame_source(seq)
+ assert n == 3 and shape == (4, 6)
+
+ def test_rejects_2d(self):
+ with pytest.raises(TypeError, match="3-D"):
+ frame_source(np.zeros((8, 9)))
+
+ def test_rejects_unknown(self):
+ with pytest.raises(TypeError, match="cannot read frames"):
+ frame_source(object())
+
+ def test_hyperspy_signal_wrong_nav_dim_is_rejected(self):
+ """A 4D-STEM scan is not a movie; say so instead of solving nonsense."""
+ class _AM:
+ navigation_dimension = 2
+ signal_dimension = 2
+
+ class _Sig:
+ axes_manager = _AM()
+ data = np.zeros((3, 3, 4, 4))
+
+ with pytest.raises(TypeError, match="1-D navigation"):
+ frame_source(_Sig())
+
+ def test_dask_reads_one_frame_only(self):
+ """The Memory-Safety rule, enforced: never compute the whole array."""
+ da = pytest.importorskip("dask.array")
+ arr = da.zeros((6, 8, 9), chunks=(1, 8, 9))
+ n, get, shape = frame_source(arr)
+ assert n == 6 and shape == (8, 9)
+ called = {"full": 0}
+ real_compute = da.Array.compute
+
+ def guard(self, *a, **k):
+ if self.shape == (6, 8, 9):
+ called["full"] += 1
+ return real_compute(self, *a, **k)
+
+ try:
+ da.Array.compute = guard
+ f = get(2)
+ finally:
+ da.Array.compute = real_compute
+ assert f.shape == (8, 9)
+ assert called["full"] == 0, "sliced a frame but computed the whole stack"
+
+
+class TestSolveTranslationAccuracy:
+ def test_recovers_integer_shifts(self):
+ truth = np.array([[0, 0], [3, -4], [-6, 2], [1, 7]], dtype=float)
+ stack = _shifted_stack(truth)
+ model = solve_translation(stack, device="numpy", upsample=8)
+ assert np.allclose(model.shifts, truth, atol=GATE_PX), model.shifts
+
+ def test_recovers_subpixel_shifts_inside_gate(self):
+ """The headline acceptance gate: < 0.1 px on sub-pixel ground truth.
+
+ The truth values are deliberately **off** the ``1/upsample`` grid. Shifts
+ that happen to be multiples of 1/8 are recovered to 0.00000 px by an
+ upsample=8 solve — which looks like a spectacular result and actually
+ tests nothing, because the answer is exactly representable. Off-grid truth
+ is what makes the tolerance meaningful.
+ """
+ truth = np.array(
+ [[0, 0], [1.37, -2.83], [-3.06, 0.61], [4.19, 5.44], [-0.72, -1.28]],
+ dtype=float,
+ )
+ # None of these may land on the upsampled grid, or the test is vacuous.
+ assert not np.any(np.isclose(truth[1:] * 8, np.round(truth[1:] * 8)))
+
+ stack = _shifted_stack(truth)
+ model = solve_translation(stack, device="numpy", upsample=8,
+ reference="first")
+ err = np.abs(model.shifts - truth)
+ assert err.max() < GATE_PX, f"max error {err.max():.4f} px\n{model.shifts}"
+
+ def test_higher_upsample_reduces_error(self):
+ """Off-grid error should shrink as the upsampled grid gets finer.
+
+ This is the test that would have caught the `_upsampled_dft` bug where
+ the frequency scaling was omitted: with that bug every result quantised to
+ 1/upsample regardless, so raising upsample changed the quantum but the
+ error stayed the same order. Here it must genuinely improve.
+ """
+ truth = np.array([[0, 0], [2.31, -1.77], [-3.42, 4.09]], dtype=float)
+ stack = _shifted_stack(truth)
+ errs = {}
+ for u in (2, 8, 32):
+ m = solve_translation(stack, device="numpy", upsample=u,
+ reference="first")
+ errs[u] = float(np.abs(m.shifts - truth).max())
+ assert errs[8] < errs[2], errs
+ assert errs[32] <= errs[8] + 1e-4, errs
+
+ def test_frame_zero_is_the_origin(self):
+ stack = _shifted_stack(np.array([[0, 0], [2, 3]], dtype=float))
+ model = solve_translation(stack, device="numpy")
+ assert tuple(model.shifts[0]) == (0.0, 0.0)
+
+ def test_agrees_with_skimage_reference(self):
+ """Parity against the implementation we are replacing."""
+ skreg = pytest.importorskip("skimage.registration")
+ truth = np.array([[0, 0], [2.25, -3.5], [-1.75, 4.125]], dtype=float)
+ stack = _shifted_stack(truth)
+ model = solve_translation(stack, device="numpy", upsample=8,
+ reference="first", apodize=False)
+ for i in range(1, len(truth)):
+ ref, _, _ = skreg.phase_cross_correlation(
+ stack[0], stack[i], upsample_factor=8, normalization="phase")
+ assert np.allclose(model.shifts[i], ref, atol=0.05), (
+ f"frame {i}: ours={model.shifts[i]} skimage={ref}")
+
+ def test_sequential_reference_accumulates(self):
+ """Sequential mode must return CUMULATIVE shifts, not per-pair deltas."""
+ truth = np.array([[0, 0], [2, 0], [4, 0], [6, 0]], dtype=float)
+ stack = _shifted_stack(truth)
+ model = solve_translation(stack, device="numpy", reference="sequential",
+ upsample=4)
+ assert np.allclose(model.shifts, truth, atol=GATE_PX), model.shifts
+
+ def test_running_reference_survives_one_corrupt_frame(self):
+ """Why 'running' is the default: a single bad frame must not poison it.
+
+ The frames AFTER the corrupt one are what matters. The corrupt frame's own
+ shift is meaningless by construction and is not asserted on.
+ """
+ truth = np.array([[0, 0], [2, 1], [4, 2], [6, 3], [8, 4]], dtype=float)
+ stack = _shifted_stack(truth).copy()
+ rng = np.random.default_rng(0)
+ stack[2] = rng.standard_normal(stack.shape[1:]).astype(np.float32) # garbage
+ model = solve_translation(stack, device="numpy", upsample=8, max_shift=20)
+ good = [1, 3, 4]
+ err = np.abs(model.shifts[good] - truth[good]).max()
+ assert err < 0.5, f"good frames drifted after a corrupt frame: {model.shifts}"
+ assert model.params["rejected_from_reference"] >= 1, (
+ "nothing was kept out of the reference, so this passed by luck rather "
+ "than by the outlier rejection it is meant to exercise")
+
+ def test_outlier_rejection_can_be_disabled(self):
+ """And with it off, the corrupt frame really does poison the reference —
+ which is what makes the test above non-vacuous."""
+ truth = np.array([[0, 0], [2, 1], [4, 2], [6, 3], [8, 4]], dtype=float)
+ stack = _shifted_stack(truth).copy()
+ rng = np.random.default_rng(0)
+ stack[2] = rng.standard_normal(stack.shape[1:]).astype(np.float32)
+ model = solve_translation(stack, device="numpy", upsample=8, max_shift=20,
+ reject_outliers=False)
+ assert model.params["rejected_from_reference"] == 0
+ good = [3, 4]
+ assert np.abs(model.shifts[good] - truth[good]).max() > 1.0, (
+ "the corrupt frame no longer poisons an unprotected reference — if the "
+ "solver became robust some other way, check deliberately")
+
+ def test_clean_stack_rejects_nothing(self):
+ """The rejection must not fire on ordinary frame-to-frame variation."""
+ truth = np.array([[0, 0], [1.5, 0.5], [3, 1], [4.5, 1.5], [6, 2]], float)
+ model = solve_translation(_shifted_stack(truth), device="numpy",
+ upsample=8, max_shift=20)
+ assert model.params["rejected_from_reference"] == 0
+ assert np.abs(model.shifts - truth).max() < GATE_PX
+
+
+class TestSolveTranslationGuards:
+ def test_max_shift_rejects_far_peak(self):
+ """A shift beyond max_shift is clamped out of the search, not returned."""
+ truth = np.array([[0, 0], [20, 0]], dtype=float)
+ stack = _shifted_stack(truth)
+ model = solve_translation(stack, device="numpy", max_shift=5, upsample=1)
+ assert abs(model.shifts[1][0]) <= 5.0 + 1e-6, model.shifts
+
+ def test_impossible_bounds_raise(self):
+ stack = _shifted_stack(np.zeros((2, 2)))
+ with pytest.raises(ValueError, match="exclude every possible shift"):
+ solve_translation(stack, device="numpy", max_shift=1, min_shift=50)
+
+ def test_bad_reference_name_raises(self):
+ stack = _shifted_stack(np.zeros((2, 2)))
+ with pytest.raises(ValueError, match="unknown reference"):
+ solve_translation(stack, device="numpy", reference="nonsense")
+
+ def test_fixed_index_out_of_range_raises(self):
+ stack = _shifted_stack(np.zeros((2, 2)))
+ with pytest.raises(ValueError, match="outside"):
+ solve_translation(stack, device="numpy", reference="fixed:99")
+
+ def test_progress_reports_every_frame(self):
+ stack = _shifted_stack(np.zeros((4, 2)))
+ seen = []
+ solve_translation(stack, device="numpy", progress=lambda d, t: seen.append((d, t)))
+ assert seen[0] == (1, 4) and seen[-1] == (4, 4)
+
+ def test_on_shift_streams_every_frame_as_it_solves(self):
+ """The drift caret draws its curve live; `progress` cannot carry that.
+
+ `progress` is only a count, and the shift array is solver-local until the
+ return — so without this callback a UI can show a bar but not a trace.
+ """
+ truth = np.array([[0, 0], [2, 1], [4, 2], [6, 3]], dtype=float)
+ seen = []
+ model = solve_translation(_shifted_stack(truth), device="numpy",
+ upsample=8, reference="first",
+ on_shift=lambda i, dy, dx, s: seen.append((i, dy, dx)))
+ assert [i for i, _, _ in seen] == list(range(len(truth))), (
+ f"expected one callback per frame in order, got {seen}")
+ streamed = np.array([[dy, dx] for _, dy, dx in seen])
+ assert np.allclose(streamed, model.shifts, equal_nan=True), (
+ "the streamed values disagree with the returned array")
+
+ def test_on_shift_is_optional(self):
+ stack = _shifted_stack(np.array([[0, 0], [1, 1]], float))
+ assert solve_translation(stack, device="numpy").n_frames == 2
+
+ def test_cancel_leaves_nan_not_a_silent_partial(self):
+ stack = _shifted_stack(np.array([[0, 0], [1, 1], [2, 2], [3, 3]], float))
+ calls = {"n": 0}
+
+ def cancel():
+ calls["n"] += 1
+ return calls["n"] > 1
+
+ model = solve_translation(stack, device="numpy", cancel=cancel)
+ assert np.isnan(model.shifts[-1]).all(), (
+ "a cancelled solve must be detectable, not quietly truncated")
+
+
+class TestAlignmentROI:
+ """Correlating on a sub-region. Not a speed switch — often the RIGHT answer.
+
+ Whole-frame correlation averages over everything that moved, so on a movie
+ where the sample itself evolves, the sample's motion contaminates the estimate
+ of the stage's. Restricting to a static landmark measures the stage alone.
+ """
+
+ def test_roi_recovers_the_same_shift_as_the_full_frame(self):
+ truth = np.array([[0, 0], [2.5, -1.75], [-3.25, 4.0]], dtype=float)
+ stack = _shifted_stack(truth, h=96, w=112)
+ full = solve_translation(stack, device="numpy", upsample=8,
+ reference="first")
+ roi = solve_translation(stack, device="numpy", upsample=8,
+ reference="first", roi=(20, 20, 56, 64))
+ assert np.abs(roi.shifts - truth).max() < 0.3, roi.shifts
+ assert np.abs(roi.shifts - full.shifts).max() < 0.3, (
+ "the ROI solve disagrees with the full-frame solve on the same data")
+
+ def test_shifts_apply_to_the_whole_frame(self):
+ """A translation is a translation; the ROI only chooses where to measure."""
+ truth = np.array([[0, 0], [3.0, -2.0]], dtype=float)
+ stack = _shifted_stack(truth, h=96, w=112)
+ model = solve_translation(stack, device="numpy", upsample=8,
+ reference="first", roi=(30, 30, 40, 48))
+ from spyde.drift import shift_frame
+ aligned = shift_frame(stack[1], model.shifts[1], fill=0.0)
+ core = (slice(40, -40), slice(40, -40)) # far OUTSIDE the ROI
+ resid = np.abs(aligned[core] - stack[0][core]).max()
+ raw = np.abs(stack[1][core] - stack[0][core]).max()
+ assert resid < 0.25 * raw, (
+ "correcting with an ROI-derived shift did not align the region "
+ "outside the ROI")
+
+ def test_roi_is_recorded_in_params(self):
+ stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64)
+ m = solve_translation(stack, device="numpy", roi=(8, 8, 32, 32))
+ assert m.params["roi"] == [8, 8, 32, 32]
+ assert m.params["frame_shape"] == [64, 64], (
+ "frame_shape must stay the FULL frame — the shifts apply to it")
+
+ def test_out_of_bounds_roi_raises_rather_than_clamping(self):
+ """A silently shrunk ROI would correlate on a region the user never chose."""
+ stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64)
+ for bad in [(0, 0, 80, 32), (40, 40, 32, 32), (-4, 0, 32, 32)]:
+ with pytest.raises(ValueError, match="outside"):
+ solve_translation(stack, device="numpy", roi=bad)
+
+ def test_tiny_roi_raises(self):
+ stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64)
+ with pytest.raises(ValueError, match="at least"):
+ solve_translation(stack, device="numpy", roi=(0, 0, 8, 8))
+
+ def test_malformed_roi_raises(self):
+ stack = _shifted_stack(np.zeros((2, 2)), h=64, w=64)
+ with pytest.raises(ValueError, match=r"\(y0, x0, h, w\)"):
+ solve_translation(stack, device="numpy", roi=(1, 2, 3))
+
+ def test_roi_ignores_motion_outside_it(self):
+ """The point of the feature, on data built to punish whole-frame."""
+ rng = np.random.default_rng(3)
+ h, w, n = 96, 128, 5
+ base = _scene(h, w, noise=0.0)
+ frames = []
+ for t in range(n):
+ f = base.copy()
+ # A big bright square that moves the OTHER way, far from the ROI.
+ y, x = 60, 78 + 5 * t
+ f[y:y + 26, x:x + 26] += 3.0
+ frames.append(f.astype(np.float32))
+ stack = np.stack(frames)
+ roi = solve_translation(stack, device="numpy", upsample=8,
+ reference="first", roi=(4, 4, 44, 56))
+ full = solve_translation(stack, device="numpy", upsample=8,
+ reference="first")
+ # The landmark region is STATIC, so the ROI answer should be ~zero.
+ assert np.abs(roi.shifts).max() < 0.6, (
+ f"ROI solve drifted although its region never moved: {roi.shifts}")
+ assert np.abs(full.shifts).max() > np.abs(roi.shifts).max(), (
+ "the whole-frame solve was not contaminated by the moving square, so "
+ "this fixture no longer demonstrates why the ROI exists")
+
+
+class TestBackendParity:
+ """The GPU path has no independent reference, so it is pinned to numpy."""
+
+ def test_torch_matches_numpy(self):
+ torch = pytest.importorskip("torch")
+ truth = np.array([[0, 0], [2.5, -1.75], [-4.25, 3.5]], dtype=float)
+ stack = _shifted_stack(truth)
+ ref = solve_translation(stack, device="numpy", upsample=8)
+ got = solve_translation(stack, device="cpu", upsample=8)
+ assert got.params["backend"] == "torch"
+ assert np.allclose(got.shifts, ref.shifts, atol=1e-2), (
+ f"torch={got.shifts}\nnumpy={ref.shifts}")
+
+
+class TestWarp:
+ def test_integer_shift_is_exact_and_preserves_dtype(self):
+ f = np.arange(24, dtype=np.uint16).reshape(4, 6)
+ out = shift_frame(f, (1, 2), fill=0, preserve_dtype=True)
+ assert out.dtype == np.uint16
+ # The interior must be bit-identical — no resampling on a whole-pixel move.
+ # Destination [1:, 2:] is fed by source [:-1, :-2].
+ assert np.array_equal(out[1:, 2:], f[:-1, :-2])
+ assert np.all(out[0, :] == 0) and np.all(out[:, :2] == 0)
+
+ def test_nan_padding_marks_uncovered(self):
+ f = np.ones((5, 5), dtype=np.float32)
+ out = shift_frame(f, (2, 0))
+ assert np.isnan(out[:2]).all()
+ assert np.allclose(out[2:], 1.0)
+
+ def test_subpixel_shift_interpolates_and_pads(self):
+ f = _scene(32, 32).astype(np.float32)
+ out = shift_frame(f, (0.5, -0.5))
+ assert out.dtype == np.float32
+ assert np.isnan(out[0]).all(), "top row needs off-frame data"
+ assert np.isnan(out[:, -1]).all(), "right column needs off-frame data"
+ assert np.isfinite(out[3:-3, 3:-3]).all()
+
+ def test_nan_does_not_bleed_into_real_data(self):
+ """Interpolating with NaN cval would smear it `order` px inward."""
+ f = np.ones((16, 16), dtype=np.float32)
+ out = shift_frame(f, (2.5, 0.0))
+ assert np.isfinite(out[4:, :]).all(), "NaN bled past the padded border"
+ assert np.allclose(out[5:-1, :], 1.0, atol=1e-5)
+
+ def test_preserve_dtype_rejects_subpixel(self):
+ f = np.zeros((4, 4), dtype=np.uint16)
+ with pytest.raises(ValueError, match="whole-pixel"):
+ shift_frame(f, (0.5, 0), fill=0, preserve_dtype=True)
+
+ def test_round_trip_is_sign_symmetric(self):
+ """Shifting out and back must land where it started.
+
+ This pins the SIGN symmetry, not interpolation fidelity — so the scene is
+ noise-free (see :func:`_scene`) and cubic interpolation is used. Two
+ bilinear passes over pixel-scale noise would lose ~1% of amplitude for
+ entirely legitimate reasons and tell us nothing about the sign.
+ """
+ f = _scene(64, 64, noise=0.0).astype(np.float32)
+ moved = shift_frame(f, (3.25, -2.5), fill=0.0, order=3)
+ back = shift_frame(moved, (-3.25, 2.5), fill=0.0, order=3)
+ core = (slice(10, -10), slice(10, -10))
+ err = np.abs(back[core] - f[core]).max()
+ assert err < 0.01, f"round trip lost {err:.4f} — sign asymmetry?"
+
+ def test_round_trip_beats_the_uncorrected_offset(self):
+ """Sanity: the corrected result is far closer than the shifted one."""
+ f = _scene(64, 64, noise=0.0).astype(np.float32)
+ moved = shift_frame(f, (3.25, -2.5), fill=0.0)
+ back = shift_frame(moved, (-3.25, 2.5), fill=0.0)
+ core = (slice(10, -10), slice(10, -10))
+ assert np.abs(back[core] - f[core]).max() < \
+ 0.1 * np.abs(moved[core] - f[core]).max()
+
+ def test_rejects_non_finite_shift(self):
+ with pytest.raises(ValueError, match="finite"):
+ shift_frame(np.zeros((4, 4)), (np.nan, 0))
+
+ def test_coverage_matches_finite_pixels(self):
+ f = np.ones((20, 20), dtype=np.float32)
+ for s in [(0, 0), (3, -2), (2.5, 1.25), (-4.75, 6.5)]:
+ out = shift_frame(f, s)
+ cov = coverage_mask((20, 20), s)
+ assert np.array_equal(np.isfinite(out), cov), f"mismatch at shift {s}"
+
+
+class TestDriftModel:
+ def test_shape_validation(self):
+ with pytest.raises(ValueError, match=r"\(N, 2\)"):
+ DriftModel(shifts=np.zeros((4, 3)))
+
+ def test_residual_length_validation(self):
+ with pytest.raises(ValueError, match="residuals must be"):
+ DriftModel(shifts=np.zeros((4, 2)), residuals=np.zeros(3))
+
+ def test_is_integer(self):
+ assert DriftModel(shifts=np.array([[0, 0], [2, -3]], float)).is_integer
+ assert not DriftModel(shifts=np.array([[0, 0], [2.5, 0]], float)).is_integer
+
+ def test_max_abs_shift_ignores_nan(self):
+ m = DriftModel(shifts=np.array([[0, 0], [3, -7], [np.nan, np.nan]], float))
+ assert m.max_abs_shift == 7.0
+
+ def test_frame_conversions_are_inverses(self):
+ m = DriftModel(shifts=np.array([[0, 0], [2.5, -1.5], [4, 3]], float))
+ pos = np.array([[10.0, 12.0], [20.0, 22.0]])
+ idx = np.array([1, 2])
+ assert np.allclose(m.to_lab_frame(m.to_sample_frame(pos, idx), idx), pos)
+
+ def test_to_sample_frame_removes_stage_motion(self):
+ """A particle that only *appears* to move because the stage drifted."""
+ m = DriftModel(shifts=np.array([[0, 0], [-5, 0], [-10, 0]], float))
+ # Same physical spot, drifting downward in the raw frames.
+ lab = np.array([[30.0, 40.0], [35.0, 40.0], [40.0, 40.0]])
+ idx = np.array([0, 1, 2])
+ sample = m.to_sample_frame(lab, idx)
+ assert np.allclose(sample[:, 0], 30.0), sample
+
+ def test_save_load_round_trip(self, tmp_path):
+ m = DriftModel(
+ shifts=np.array([[0, 0], [1.25, -2.5]], float),
+ residuals=np.array([np.inf, 12.5], np.float32),
+ params={"upsample": 8}, provenance={"action": "drift"},
+ reference="running",
+ )
+ p = str(tmp_path / "d.npz")
+ m.save(p)
+ back = DriftModel.load(p)
+ assert np.array_equal(back.shifts, m.shifts)
+ assert back.params["upsample"] == 8
+ assert back.provenance == {"action": "drift"}
+ assert back.reference == "running"
+ assert np.array_equal(back.residuals, m.residuals)
+
+ def test_load_rejects_future_format(self, tmp_path):
+ import json
+ p = str(tmp_path / "bad.npz")
+ np.savez_compressed(
+ p, shifts=np.zeros((2, 2), np.float32),
+ meta=np.array(json.dumps({"format_version": 999})))
+ with pytest.raises(ValueError, match="unsupported DriftModel format"):
+ DriftModel.load(p)
+
+
+class TestEndToEnd:
+ def test_solve_then_warp_aligns_the_stack(self):
+ """The whole point: after correction, every frame agrees with frame 0."""
+ truth = np.array(
+ [[0, 0], [2.5, -3.0], [-4.25, 1.75], [6.0, 4.5]], dtype=float)
+ stack = _shifted_stack(truth, h=80, w=80)
+ model = solve_translation(stack, device="numpy", upsample=8,
+ reference="first")
+
+ core = (slice(12, -12), slice(12, -12))
+ ref = stack[0][core]
+ for i in range(1, len(truth)):
+ aligned = shift_frame(stack[i], model.shifts[i], fill=0.0)
+ resid = np.abs(aligned[core] - ref).max()
+ raw = np.abs(stack[i][core] - ref).max()
+ assert resid < raw * 0.2, (
+ f"frame {i}: correction barely helped (resid={resid:.4f} "
+ f"raw={raw:.4f}) — check the SIGN convention")
diff --git a/spyde/tests/migrated/test_drift_wizard.py b/spyde/tests/migrated/test_drift_wizard.py
new file mode 100644
index 00000000..efcfa303
--- /dev/null
+++ b/spyde/tests/migrated/test_drift_wizard.py
@@ -0,0 +1,743 @@
+"""
+The Drift Correction wizard backend (``drift_*`` staged handlers).
+
+Handlers are called directly as ``fn(session, plot, payload)`` and polled with
+``_wait`` — the shape ``test_find_vectors_wizard.py`` establishes, because the
+solve and the check sums both run on a worker thread.
+
+The four claims that matter:
+
+:class:`TestCheckWindow`
+ Plan A8 / README §6. The verification surface is a SEPARATE window, and a
+ bare ``figure`` is not a registered ``Plot`` — so it must be reachable
+ through ``session.controller_by_window_id`` and must disappear on close. A
+ check window that leaks is the exact bug README §6 documents.
+:class:`TestSolve`
+ The solved shifts must match ``particle_movie``'s stamped ground truth, and
+ ``tree.drift`` must carry the model. Ground truth beats a golden number.
+:class:`TestCommitIsLazy`
+ The CLAUDE.md memory-safety rule, guarded the way
+ ``test_find_vectors_memory.py`` guards it: a ``da.Array.compute`` spy that
+ counts calls on the full-dataset shape. And the corrected node has to be
+ genuinely better — an aligned stack sums SHARP, which is the whole claim the
+ check window makes to the user.
+:class:`TestDoubleFire`
+ README §4 / StrictMode: open, close, open leaves exactly ONE controller and
+ exactly ONE check window.
+"""
+from __future__ import annotations
+
+import time
+
+import numpy as np
+import pytest
+
+from spyde.actions import drift_action as dr
+
+
+@pytest.fixture(autouse=True)
+def _capture_module_emit(window, monkeypatch):
+ """Route ``drift_action``'s own ``emit`` into the captured list.
+
+ The module does ``from spyde.backend.ipc import emit`` at import, so
+ conftest's patch of ``ipc.emit`` never reaches that binding — the identical
+ hazard conftest already documents for ``session.py``, and the identical fix.
+ ``emit_status``/``emit_error`` need no patch: they resolve ``emit`` inside
+ ``ipc`` at call time.
+ """
+ monkeypatch.setattr(dr, "emit", window["messages"].append)
+
+# Small but enough for the drift curve to turn: `particle_movie`'s drift is a
+# smooth excursion, and 8 frames already reach ~6 px, which is 5x the tolerance
+# asserted below.
+N_FRAMES = 8
+
+
+def _signal_plot(session):
+ return next((p for p in session._plots
+ if not p.is_navigator and p.plot_state is not None), None)
+
+
+def _wait(pred, timeout=120.0):
+ end = time.time() + timeout
+ while time.time() < end:
+ if pred():
+ return True
+ time.sleep(0.05)
+ return False
+
+
+def _movie(window, frames: int = N_FRAMES):
+ session = window["window"]
+ session._load_test_data_particles({"frames": frames})
+ plot = _wait(lambda: _signal_plot(session) is not None) and _signal_plot(session)
+ assert plot is not None, "the particle movie never produced a signal plot"
+ return session, plot, plot.signal_tree
+
+
+def _opened(window, frames: int = N_FRAMES, **params):
+ session, plot, tree = _movie(window, frames)
+ dr.drift_open(session, plot, {"upsample": 8, "max_shift": 16, **params})
+ assert _wait(lambda: getattr(tree, "_drift_wizard", None) is not None
+ and tree._drift_wizard.window_id is not None), \
+ "the Drift Check window never opened"
+ return session, plot, tree, tree._drift_wizard
+
+
+def _solved(window, frames: int = N_FRAMES):
+ session, plot, tree, wiz = _opened(window, frames)
+ dr.drift_run(session, plot, {"upsample": 8, "max_shift": 16})
+ assert _wait(lambda: wiz.model is not None), "the solve never finished"
+ return session, plot, tree, wiz
+
+
+def _of_type(messages, kind):
+ return [m for m in messages if isinstance(m, dict) and m.get("type") == kind]
+
+
+def _sharpness(img) -> float:
+ """Mean squared gradient — an aligned sum has more of it than a blurred one."""
+ a = np.nan_to_num(np.asarray(img, np.float64))
+ gy, gx = np.gradient(a)
+ return float(np.mean(gy ** 2 + gx ** 2))
+
+
+class _FullComputeGuard:
+ """Count ``.compute()`` calls made on the whole movie."""
+
+ def __init__(self, shape):
+ self.shape = tuple(shape)
+ self.hits = 0
+
+ def __enter__(self):
+ import dask.array as da
+ self._real = da.Array.compute
+ guard = self
+
+ def _spy(arr, *a, **k):
+ if tuple(arr.shape) == guard.shape:
+ guard.hits += 1
+ return guard._real(arr, *a, **k)
+
+ da.Array.compute = _spy
+ return self
+
+ def __exit__(self, *exc):
+ import dask.array as da
+ da.Array.compute = self._real
+ return False
+
+
+class TestCheckWindow:
+ def test_open_registers_a_controller_for_the_bare_figure(self, window):
+ """README §6: a bare `figure` is not a Plot, so dispatch can only find
+ it through the window-controller registry."""
+ session, _plot, _tree, wiz = _opened(window)
+ assert session.controller_by_window_id(wiz.window_id) is wiz
+ assert session._plot_by_window_id(wiz.window_id) is None, \
+ "the check window is supposed to be a bare figure, not a Plot"
+
+ def test_the_window_shows_the_uncorrected_sum(self, window):
+ session, _plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ figs = [m for m in _of_type(msgs, "figure")
+ if m.get("window_id") == wiz.window_id]
+ assert figs and figs[-1]["title"] == "Drift Check"
+ assert wiz._before_sum is not None and np.isfinite(wiz._before_sum).any()
+
+ def test_open_solves_nothing(self, window):
+ """Plan A8: drift correction is explicit — nothing runs on load."""
+ _s, _p, tree, wiz = _opened(window)
+ assert wiz.model is None
+ assert getattr(tree, "drift", None) is None
+
+ def test_close_takes_the_window_with_it(self, window):
+ session, plot, tree, wiz = _opened(window)
+ wid = wiz.window_id
+ dr.drift_close(session, plot, {})
+ assert getattr(tree, "_drift_wizard", None) is None
+ assert session.controller_by_window_id(wid) is None
+ from spyde.actions.figure_registry import _FIGS
+ assert wid not in _FIGS, "the check figure outlived its window"
+
+ def test_the_summed_subset_is_bounded(self, window):
+ """A sum is a sharpness test, not a measurement — the cap is what keeps
+ the check window usable on a long movie."""
+ _s, _p, _t, wiz = _opened(window)
+ wiz._sum_indices = None # as if the movie were long
+ idx = wiz.sum_indices(10_000)
+ assert idx.size <= dr._SUM_MAX_FRAMES
+ assert idx[0] == 0 and idx[-1] == 9_999
+ assert wiz.sum_indices(10_000) is idx, (
+ "the subset is memoised so the before and after sums cover the SAME "
+ "frames — comparing two different subsets means nothing")
+
+
+class TestMethodStubs:
+ def test_nonrigid_is_selectable_now_that_it_is_implemented(self, window):
+ """Non-rigid used to be a stub that silently reverted to rigid.
+
+ It is implemented (``spyde.drift.nonrigid``), so selecting it must
+ STICK — a caret that quietly reverts would put the wrong ``kind`` in
+ provenance, which is what the stub test was guarding against.
+ """
+ session, plot, _tree, wiz = _opened(window)
+ dr.drift_set_method(session, plot, {"method": "nonrigid"})
+ assert wiz.params["method"] == "nonrigid"
+ assert "nonrigid" not in dr._UNAVAILABLE
+
+ def test_the_nonrigid_field_parameterisation_is_selectable(self, window):
+ """Both parameterisations describe different physics; neither is a default
+ the user should be stuck with."""
+ session, plot, _tree, wiz = _opened(window)
+ dr.drift_set_method(session, plot, {"method": "nonrigid"})
+ for name in dr.NONRIGID_MODELS:
+ wiz.params = dr._coerce({**wiz.params, "nonrigid_model": name})
+ assert wiz.params["nonrigid_model"] == name
+
+ def test_an_unknown_field_falls_back_rather_than_raising(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ wiz.params = dr._coerce({**wiz.params, "nonrigid_model": "banana"})
+ assert wiz.params["nonrigid_model"] == dr.DEFAULTS["nonrigid_model"]
+
+ def test_rigid_affine_says_so_too(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ dr.drift_set_method(session, plot, {"method": "rigid_affine"})
+ assert wiz.params["method"] == "rigid"
+
+ def test_unknown_method_errors(self, window):
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ dr.drift_set_method(session, plot, {"method": "banana"})
+ assert any("unknown model" in str(m.get("text", ""))
+ for m in _of_type(msgs, "error"))
+
+
+class _Box:
+ """Stand-in for an anyplotlib RectangleWidget (x/y/w/h in IMAGE PIXELS).
+
+ The headless session does build a real ``_plot2d``, so the wizard's own box
+ exists — but a test that wants a SPECIFIC region needs to place one, and
+ dragging a real widget means faking pointer events. Swapping this in is the
+ smaller lie, and it exercises the same ``roi_box()`` conversion.
+ """
+
+ def __init__(self, x, y, w, h):
+ self.x, self.y, self.w, self.h = float(x), float(y), float(w), float(h)
+
+ def set(self, **kw):
+ for k, v in kw.items():
+ setattr(self, k, float(v))
+
+ def hide(self):
+ pass
+
+
+class TestDiscoveryPreview:
+ """The centrepiece: a draggable box + a drift-corrected sum of it over ~20
+ frames, so the user sees whether alignment works BEFORE paying for the whole
+ movie."""
+
+ def test_open_previews_the_default_box(self, window):
+ session, _plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ assert _wait(lambda: _of_type(msgs, "drift_preview")), \
+ "opening the caret never produced a discovery preview"
+ prev = _of_type(msgs, "drift_preview")[-1]
+ assert prev["frames"] >= 2
+ assert prev["gain"] > 1.0, (
+ "the aligned sum of the default box is not sharper than the raw one "
+ "— the preview cannot discriminate anything if it never improves")
+
+ def test_the_preview_never_computes_the_whole_movie(self, window):
+ session, plot, tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ del msgs[:]
+ with _FullComputeGuard(tree.root.data.shape) as guard:
+ dr.drift_tune(session, plot, {"upsample": 4, "max_shift": 12})
+ assert _wait(lambda: _of_type(msgs, "drift_preview"))
+ assert guard.hits == 0
+
+ def test_tune_stores_the_new_parameters(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ dr.drift_tune(session, plot, {"upsample": 16, "max_shift": 9.0})
+ assert wiz.params["upsample"] == 16
+ assert wiz.params["max_shift"] == 9.0
+
+ def test_the_preview_uses_the_box_even_with_the_toggle_off(self, window):
+ """The toggle is the COMMITMENT (does the full solve restrict to the
+ box); the preview is the QUESTION and always asks it about the box."""
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ assert wiz.params["use_roi"] is False
+ wiz._roi_widget = _Box(10, 12, 60, 48)
+ del msgs[:]
+ dr.drift_tune(session, plot, {})
+ assert _wait(lambda: _of_type(msgs, "drift_preview"))
+ assert _of_type(msgs, "drift_preview")[-1]["roi"] == [12, 10, 48, 60]
+
+ def test_the_box_is_read_in_image_pixels_as_y0_x0_h_w(self, window):
+ """anyplotlib 2-D widgets report IMAGE PIXELS and solve_translation's
+ roi is in pixels — the two meet with NO scale conversion."""
+ _s, _p, _t, wiz = _opened(window)
+ wiz._frame_shape = (96, 112)
+ wiz._roi_widget = _Box(x=20, y=8, w=40, h=32)
+ assert wiz.roi_box() == (8, 20, 32, 40)
+
+ def test_the_box_is_clamped_into_the_frame(self, window):
+ _s, _p, _t, wiz = _opened(window)
+ wiz._frame_shape = (96, 112)
+ wiz._roi_widget = _Box(x=100, y=90, w=400, h=400)
+ y0, x0, h, w = wiz.roi_box()
+ assert 0 <= y0 and 0 <= x0
+ assert y0 + h <= 96 and x0 + w <= 112
+
+ def test_a_box_below_the_solver_floor_is_refused(self, window):
+ """solve_translation REJECTS a too-small roi rather than clamping it, so
+ the caret must never hand it one."""
+ from spyde.drift.translation import _MIN_ROI
+ assert dr._ROI_MIN_PX >= _MIN_ROI
+ _s, _p, _t, wiz = _opened(window)
+ wiz._frame_shape = (96, 112)
+ wiz._roi_widget = _Box(x=0, y=0, w=4, h=4)
+ y0, x0, h, w = wiz.roi_box()
+ assert h >= _MIN_ROI and w >= _MIN_ROI
+
+ def test_a_superseded_preview_does_not_paint(self, window):
+ """Latest-wins: a drag that outruns the solve must drop the stale
+ result, not paint it over the newer one."""
+ _s, _p, tree, wiz = _opened(window)
+ gen = dr.bump_generation(tree, "_drift_preview_gen")
+ painted = []
+ wiz.show_preview = lambda res: painted.append(res)
+ dr.bump_generation(tree, "_drift_preview_gen") # a newer drag lands
+ assert not dr.is_current(tree, "_drift_preview_gen", gen)
+ assert painted == []
+
+ def test_the_settle_timer_coalesces_a_drag(self, window):
+ """The widget's pointer_move fires at renderer frame rate; only the
+ RESTING geometry may solve."""
+ _s, _p, _t, wiz = _opened(window)
+ fired = []
+ wiz._fire_preview = lambda: fired.append(1)
+ for _ in range(20):
+ wiz.schedule_preview(delay=0.05)
+ assert fired == []
+ assert _wait(lambda: len(fired) >= 1, timeout=3.0)
+ time.sleep(0.2)
+ assert len(fired) == 1, f"{len(fired)} solves for one drag"
+
+
+class TestSharpnessNumber:
+ """The gain has to be an ANSWER, not decoration: a landmark and a
+ featureless patch must come out clearly different."""
+
+ @staticmethod
+ def _stack(n=16, size=140, pad=20):
+ """Textured on the left half, flat on the right, drifting rigidly."""
+ from scipy.ndimage import gaussian_filter, map_coordinates
+ rng = np.random.default_rng(3)
+ canvas = np.zeros((size + 2 * pad, size + 2 * pad), np.float32) + 1.0
+ tex = gaussian_filter(rng.standard_normal(canvas.shape), 1.5) * 0.6
+ half = size // 2 + pad
+ canvas[:, :half] += tex[:, :half]
+ drift = np.stack([np.linspace(0, 8.0, n), np.linspace(0, -5.0, n)], 1)
+ yy, xx = np.mgrid[0:size, 0:size].astype(np.float64)
+ frames = np.empty((n, size, size), np.float32)
+ for t in range(n):
+ dy, dx = drift[t]
+ frames[t] = map_coordinates(canvas, [yy + pad - dy, xx + pad - dx],
+ order=1, mode="nearest")
+ frames += rng.normal(0, 0.02, frames.shape).astype(np.float32)
+ return frames
+
+ def test_a_landmark_beats_a_featureless_patch(self):
+ frames = self._stack()
+ idx = np.arange(frames.shape[0])
+ params = dict(dr.DEFAULTS)
+ good = dr.preview_alignment(frames.__getitem__, idx, (30, 5, 80, 55),
+ params=params)
+ bad = dr.preview_alignment(frames.__getitem__, idx, (30, 82, 80, 55),
+ params=params)
+ assert good["gain"] > 2.0, f"a real landmark only scored {good['gain']:.2f}"
+ assert bad["gain"] < 1.0, f"a featureless patch scored {bad['gain']:.2f}"
+ assert good["gain"] > 3 * bad["gain"]
+
+ def test_the_nan_border_does_not_inflate_the_number(self):
+ """A shifted frame's uncovered edge is NaN (plan A7). Zero-filling it
+ manufactures a step whose gradient energy dwarfs the image's own — every
+ ROI would look brilliant."""
+ a = np.ones((32, 32), np.float32)
+ a[:4, :] = np.nan
+ assert dr._gradient_energy(a) == 0.0
+
+ def test_the_two_sums_are_measured_on_the_same_pixels(self):
+ raw = np.ones((16, 16), np.float32)
+ aligned = raw.copy()
+ aligned[:3, :] = np.nan
+ both = np.isfinite(raw) & np.isfinite(aligned)
+ assert dr._gradient_energy(raw, both) == dr._gradient_energy(aligned, both)
+
+ def test_the_preview_sample_spans_the_whole_movie(self):
+ """20 CONSECUTIVE frames of a long movie drift by almost nothing, so a
+ contiguous window would say "looks fine" for every box."""
+ idx = dr._preview_indices(3000, 20, 64 * 64 * 4)
+ assert idx[0] == 0 and idx[-1] == 2999 and idx.size <= 20
+
+ def test_the_sample_is_thinned_to_fit_the_byte_cap(self):
+ """With no ROI the crop IS the frame — 20 × 4096² float32 is 1.3 GB."""
+ idx = dr._preview_indices(3000, 20, 4096 * 4096 * 4)
+ assert 2 <= idx.size < 20
+
+
+def _ground_truth(tree):
+ import spyde.data.synthetic as sy
+ return np.asarray(sy.ground_truth(tree.root)["drift"], np.float64)
+
+
+class TestSolve:
+ def test_shifts_match_the_stamped_ground_truth(self, window):
+ _s, _p, tree, wiz = _solved(window)
+ truth = _ground_truth(tree)[:N_FRAMES]
+ err = np.abs(wiz.model.shifts - truth).max()
+ assert err < 0.25, f"worst per-axis drift error {err:.3f} px"
+
+ def test_the_model_lands_on_the_tree(self, window):
+ _s, _p, tree, wiz = _solved(window)
+ assert tree.drift is wiz.model
+ assert wiz.model.kind == "rigid"
+
+ def test_run_reports_progress_and_the_finished_trace(self, window):
+ session, plot, tree, wiz = _opened(window)
+ msgs = window["messages"]
+ dr.drift_run(session, plot, {})
+ assert _wait(lambda: _of_type(msgs, "drift_result"))
+ prog = _of_type(msgs, "drift_progress")
+ assert prog and prog[-1]["done"] == prog[-1]["total"] == N_FRAMES
+ res = _of_type(msgs, "drift_result")[-1]
+ assert len(res["shifts"]) == N_FRAMES and not res["cancelled"]
+ assert res["max_abs_shift"] > 1.0
+
+ def test_run_never_computes_the_whole_movie(self, window):
+ session, plot, tree, wiz = _opened(window)
+ with _FullComputeGuard(tree.root.data.shape) as guard:
+ dr.drift_run(session, plot, {})
+ assert _wait(lambda: wiz.model is not None)
+ assert guard.hits == 0, "the solve materialised the whole movie"
+
+ def test_the_check_window_gets_a_sharper_corrected_sum(self, window):
+ """The claim the window makes to the user, asserted rather than drawn:
+ an aligned stack sums sharp, a misaligned one blurs."""
+ _s, _p, _t, wiz = _solved(window)
+ n, get_frame, _shape = wiz.frames()
+ idx = wiz.sum_indices(n)
+ before = dr._stack_sum(get_frame, idx)
+ after = dr._stack_sum(get_frame, idx, wiz.model.shifts)
+ assert _sharpness(after) > 1.5 * _sharpness(before), (
+ f"corrected sum {_sharpness(after):.5f} is not sharper than the raw "
+ f"{_sharpness(before):.5f} — check the sign convention in "
+ "spyde/drift/model.py")
+
+ def test_closing_the_tree_cancels_the_solve(self, window):
+ """Cancellation goes through BaseSignalTree.register_cancel, so closing
+ the tree has to stop it."""
+ session, plot, tree, wiz = _opened(window, frames=24)
+ dr.drift_run(session, plot, {})
+ tree.close()
+ assert _wait(lambda: wiz.model is not None or wiz._closed, timeout=60)
+ if wiz.model is not None:
+ # A cancelled solve leaves NaN for the frames it never reached, so a
+ # partial model is detectable rather than silently wrong.
+ assert not np.isfinite(wiz.model.shifts).all()
+
+ def test_run_without_a_caret_errors(self, window):
+ session, plot, _tree = _movie(window)
+ msgs = window["messages"]
+ dr.drift_run(session, plot, {})
+ assert any("caret is not open" in str(m.get("text", ""))
+ for m in _of_type(msgs, "error"))
+
+ def test_use_roi_feeds_the_box_to_the_solver(self, window):
+ """The toggle's whole job: the same rectangle the preview tested is the
+ one the full solve correlates on."""
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ wiz._frame_shape = (96, 112)
+ wiz._roi_widget = _Box(x=16, y=12, w=64, h=64)
+ dr.drift_run(session, plot, {"use_roi": True})
+ assert _wait(lambda: wiz.model is not None)
+ assert wiz.model.params["roi"] == [12, 16, 64, 64]
+ assert _of_type(msgs, "drift_result")[-1]["roi"] == [12, 16, 64, 64]
+
+
+class TestTraceWindow:
+ """The dy/dx curve is its OWN plot window, filled as the solve runs — not
+ caret furniture (plan §0.9a)."""
+
+ def test_the_solve_opens_a_second_figure_window(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ dr.drift_run(session, plot, {})
+ assert _wait(lambda: wiz.trace_window_id is not None)
+ figs = [m for m in _of_type(msgs, "figure")
+ if m.get("window_id") == wiz.trace_window_id]
+ assert figs and figs[-1]["title"] == "Drift dy/dx"
+ assert session.controller_by_window_id(wiz.trace_window_id) is wiz, \
+ "a bare figure is only reachable through the controller registry"
+ assert session._plot_by_window_id(wiz.trace_window_id) is None
+
+ def test_it_fills_from_the_on_shift_stream(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ dr.drift_run(session, plot, {})
+ assert _wait(lambda: wiz.model is not None)
+ assert _wait(lambda: int(wiz._trace.get("filled", 0)) == N_FRAMES)
+ # …and the same batches went out as `drift_trace` messages.
+ streamed = sum(len(m["points"]) for m in _of_type(msgs, "drift_trace"))
+ assert streamed == N_FRAMES
+
+ def test_the_trace_matches_the_model(self, window):
+ _s, _p, _t, wiz = _solved(window)
+ assert _wait(lambda: int(wiz._trace.get("filled", 0)) == N_FRAMES)
+ np.testing.assert_allclose(wiz._trace["dy_data"], wiz.model.shifts[:, 0],
+ atol=1e-5)
+ np.testing.assert_allclose(wiz._trace["dx_data"], wiz.model.shifts[:, 1],
+ atol=1e-5)
+
+ def test_only_the_solved_prefix_is_pushed(self, window):
+ """Pushing the NaN-padded whole array would leave anyplotlib's auto
+ y-range looking at one finite point."""
+ _s, _p, _t, wiz = _opened(window)
+ wiz.trace_window_id = None
+ wiz.open_trace_window(50)
+ wiz.push_trace([(1, 3.0, -2.0), (2, 4.0, -3.0)])
+ assert wiz._trace["filled"] == 3
+ assert np.isnan(wiz._trace["dy_data"][3:]).all()
+
+ def test_close_takes_both_windows(self, window):
+ session, plot, tree, wiz = _solved(window)
+ assert _wait(lambda: wiz.trace_window_id is not None)
+ check, trace = wiz.window_id, wiz.trace_window_id
+ dr.drift_close(session, plot, {})
+ from spyde.actions.figure_registry import _FIGS
+ for wid in (check, trace):
+ assert session.controller_by_window_id(wid) is None
+ assert wid not in _FIGS, f"figure {wid} outlived its window"
+
+
+class TestDiscard:
+ def test_discard_drops_the_model_and_the_trace_window(self, window):
+ session, plot, tree, wiz = _solved(window)
+ assert _wait(lambda: wiz.trace_window_id is not None)
+ trace = wiz.trace_window_id
+ dr.drift_discard(session, plot, {})
+ assert wiz.model is None
+ assert getattr(tree, "drift", None) is None
+ assert wiz.trace_window_id is None
+ assert session.controller_by_window_id(trace) is None
+ assert wiz.window_id is not None, "Discard must not close the caret"
+
+ def test_discard_stops_a_solve_in_flight(self, window):
+ """Same user intent as Stop, so it is the same handler: bumping the run
+ generation FIRST means a solve that finishes anyway never installs."""
+ session, plot, tree, wiz = _opened(window, frames=24)
+ dr.drift_run(session, plot, {})
+ dr.drift_discard(session, plot, {})
+ assert wiz._stop[0] is True
+ time.sleep(1.0)
+ assert wiz.model is None
+
+
+class TestCommitIsLazy:
+ def test_commit_adds_a_lazy_node_without_computing(self, window):
+ session, plot, tree, wiz = _solved(window)
+ before = set(tree.root_node.children)
+ with _FullComputeGuard(tree.root.data.shape) as guard:
+ dr.drift_commit(session, plot, {})
+ assert guard.hits == 0, "commit materialised the movie"
+ added = set(tree.root_node.children) - before
+ assert added == {"Drift corrected"}
+ node = tree.root_node.children["Drift corrected"]
+ assert node.signal._lazy
+ assert node.signal.data.shape == tree.root.data.shape
+ assert node.local is True, (
+ "a per-frame shift IS local — without the tag the derived-view "
+ "reader falls back to the opaque path")
+
+ def test_one_corrected_frame_costs_one_frame(self, window):
+ session, plot, tree, _wiz = _solved(window)
+ dr.drift_commit(session, plot, {})
+ node = tree.root_node.children["Drift corrected"]
+ with _FullComputeGuard(tree.root.data.shape) as guard:
+ frame = np.asarray(node.signal.data[3].compute())
+ assert guard.hits == 0
+ assert frame.shape == tuple(tree.root.data.shape[1:])
+
+ def test_uncovered_pixels_are_nan_not_invented(self, window):
+ """Plan A7: nothing is cropped and nothing is filled with invented
+ data — segmentation would find 'particles' in a zero-filled border."""
+ session, plot, tree, _wiz = _solved(window)
+ dr.drift_commit(session, plot, {})
+ node = tree.root_node.children["Drift corrected"]
+ frame = np.asarray(node.signal.data[N_FRAMES - 1].compute())
+ assert np.isnan(frame).any(), "no NaN padding on a shifted frame"
+ assert np.isfinite(frame).any(), "the whole frame is NaN"
+
+ def test_the_corrected_node_is_actually_aligned(self, window):
+ session, plot, tree, wiz = _solved(window)
+ dr.drift_commit(session, plot, {})
+ node = tree.root_node.children["Drift corrected"]
+ raw = np.nanmean(np.stack([np.asarray(tree.root.data[i].compute(),
+ np.float64)
+ for i in range(N_FRAMES)]), axis=0)
+ fixed = np.nanmean(np.stack([np.asarray(node.signal.data[i].compute(),
+ np.float64)
+ for i in range(N_FRAMES)]), axis=0)
+ assert _sharpness(fixed) > 1.5 * _sharpness(raw)
+
+ def test_commit_stamps_provenance(self, window):
+ session, plot, tree, _wiz = _solved(window)
+ dr.drift_commit(session, plot, {})
+ node = tree.root_node.children["Drift corrected"]
+ prov = node.signal.metadata.get_item("General.spyde_provenance")
+ assert prov["action"] == "Drift Correction"
+ assert prov["kind"] == "rigid"
+
+ def test_commit_before_solving_errors(self, window):
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ dr.drift_commit(session, plot, {})
+ assert any("solve first" in str(m.get("text", ""))
+ for m in _of_type(msgs, "error"))
+
+ def test_a_model_of_the_wrong_length_is_refused(self, window):
+ """Re-solving after a crop must not silently pair frame 0's shift with
+ a different frame."""
+ from spyde.drift import DriftModel
+ _s, _p, tree, _wiz = _opened(window)
+ bad = DriftModel(shifts=np.zeros((N_FRAMES + 3, 2), np.float32))
+ with pytest.raises(ValueError, match="covers"):
+ dr.drift_corrected(tree.root, model=bad)
+
+
+class TestDoubleFire:
+ def test_open_close_open_leaves_one_controller(self, window):
+ session, plot, tree = _movie(window)
+ built = []
+ real_init = dr.DriftWizard.__init__
+
+ def _tracking(self, *a, **k):
+ real_init(self, *a, **k)
+ built.append(self)
+
+ dr.DriftWizard.__init__ = _tracking
+ try:
+ dr.drift_open(session, plot, {})
+ dr.drift_close(session, plot, {})
+ dr.drift_open(session, plot, {})
+ finally:
+ dr.DriftWizard.__init__ = real_init
+ assert _wait(lambda: tree._drift_wizard is not None
+ and tree._drift_wizard.window_id is not None)
+ time.sleep(0.5)
+
+ alive = [w for w in built if not w._closed]
+ assert len(alive) == 1, \
+ f"expected 1 live controller, got {len(alive)} of {len(built)} built"
+ assert tree._drift_wizard is alive[0]
+ # …and exactly one check window, because a superseded open's deferred
+ # build must be dropped rather than emitting a second figure.
+ wids = [w.window_id for w in built if w.window_id is not None]
+ assert len(wids) == 1, f"{len(wids)} check windows opened"
+
+ def test_close_without_an_open_is_harmless(self, window):
+ session, plot, tree = _movie(window)
+ dr.drift_close(session, plot, {})
+ assert getattr(tree, "_drift_wizard", None) is None
+
+
+class TestSchema:
+ def test_schema_resolves_through_the_registry(self):
+ from spyde.actions import registry
+ schema = registry.wizard_parameters("drift")
+ assert schema and schema is not dr.DriftWizard.parameters
+
+ def test_schema_defaults_match_the_handler_defaults(self):
+ from spyde.actions import registry
+ schema = registry.wizard_parameters("drift")
+ for key, spec in schema.items():
+ assert key in dr.DEFAULTS, f"drift schema declares unknown param {key!r}"
+ assert spec["default"] == dr.DEFAULTS[key], \
+ f"drift schema/{key} drifted from drift_action.DEFAULTS"
+
+ def test_every_stage_is_registered(self):
+ from spyde.actions.registry import STAGED_HANDLERS, resolve_staged
+ for stage in ("drift_open", "drift_close", "drift_set_method",
+ "drift_tune", "drift_run", "drift_discard",
+ "drift_commit"):
+ assert stage in STAGED_HANDLERS
+ assert callable(resolve_staged(stage))
+
+ def test_the_default_face_is_two_toggles(self):
+ """§0.9a: everything that is not the task itself is tagged Advanced, so
+ any host renders the same small face. The caret is the enforcement; this
+ is the schema saying the same thing."""
+ from spyde.actions import registry
+ schema = registry.wizard_parameters("drift")
+ face = [k for k, s in schema.items() if not s.get("tab")]
+ assert face == ["use_roi", "reject_outliers"], \
+ f"the caret's default face grew to {face}"
+
+ def test_toolbar_entry_gates_on_a_movie(self):
+ import spyde
+ meta = spyde.TOOLBAR_ACTIONS["functions"]["Drift Correction"]
+ assert meta["function"] == "spyde.actions.drift_action.drift_correction"
+ assert meta["signal_types"] == ["insitu"], (
+ "the rigid solver needs a 1-D navigation axis; gating on `insitu` "
+ "is the same gate Play/Fast-Forward use for exactly that")
+
+ @pytest.mark.parametrize("signal_type,offered", [
+ ("insitu", True),
+ ("", False), # a single image has nothing to align
+ ("electron_diffraction", False),
+ ])
+ def test_toolbar_gating(self, signal_type, offered):
+ import spyde
+ from spyde.drawing.toolbars.plot_control_toolbar import _action_matches_plot
+
+ class _Sig:
+ _signal_type = signal_type
+
+ class _Tree:
+ particles = None
+ diffraction_vectors = None
+ root = _Sig()
+
+ class _Plot:
+ signal_tree = _Tree()
+ is_navigator = False
+
+ class _State:
+ plot = _Plot()
+ current_signal = _Sig()
+ dimensions = 2
+ navigation = False
+
+ meta = spyde.TOOLBAR_ACTIONS["functions"]["Drift Correction"]
+ assert _action_matches_plot("Drift Correction", meta, _State()) is offered
+
+
+class TestCoercion:
+ def test_unknown_method_falls_back(self):
+ assert dr._coerce({"method": "warp"})["method"] == dr.DEFAULTS["method"]
+
+ def test_unknown_reference_falls_back(self):
+ assert dr._coerce({"reference": "later"})["reference"] == "running"
+
+ def test_order_is_clamped(self):
+ assert dr._coerce({"order": 9})["order"] == 3
+ assert dr._coerce({"order": -1})["order"] == 0
+
+ def test_a_junk_value_keeps_the_default(self):
+ assert dr._coerce({"upsample": "eight"})["upsample"] == \
+ dr.DEFAULTS["upsample"]
diff --git a/spyde/tests/migrated/test_nav_cluster_handover.py b/spyde/tests/migrated/test_nav_cluster_handover.py
new file mode 100644
index 00000000..0a5393cb
--- /dev/null
+++ b/spyde/tests/migrated/test_nav_cluster_handover.py
@@ -0,0 +1,100 @@
+"""test_nav_cluster_handover.py — the threaded navigator fill hands over.
+
+`_start_progressive_nav_compute` picks its path from `self.client`. The cluster
+takes ~10 s to come up, so a file opened right after launch finds it None and
+takes the THREADED branch — one background thread, one chunk at a time.
+
+That choice used to be permanent. However many workers registered a second
+later, the whole fill ran single-threaded, which on a long movie is the
+difference between seconds and minutes and reads as "the cluster is idle".
+
+Why it looked like a movie-only bug: a 4D-STEM scan goes through the nav-shape
+prompt (a human round trip), so its cluster is always up by the time the tree is
+built. A movie opens straight through and loses the race.
+"""
+from __future__ import annotations
+
+import threading
+import time
+
+import dask.array as da
+import numpy as np
+import pytest
+
+
+class _Plot:
+ def __init__(self):
+ self.window_id = 1
+ self.painted = []
+
+ def set_data(self, arr, levels=None):
+ self.painted.append(arr)
+
+ def _emit_histogram(self, *a, **k):
+ pass
+
+
+class _Tree:
+ """Just enough tree to drive the threaded fill's handover check."""
+
+ def __init__(self, client_after: int):
+ self._client_after = client_after # chunks before the cluster "starts"
+ self._chunks_done = 0
+ self.handed_over_with = None
+ self.session = None
+ self.source_path = None
+
+ @property
+ def client(self):
+ return object() if self._chunks_done >= self._client_after else None
+
+ def _start_progressive_nav_compute(self, nav_dask, deep=None):
+ self.handed_over_with = (nav_dask, deep)
+
+
+def _drive(tree, total_chunks: int, min_remaining: int):
+ """The handover decision, lifted verbatim from the fill loop."""
+ for done in range(total_chunks):
+ tree._chunks_done = done
+ if (tree.client is not None
+ and (total_chunks - done) > min_remaining):
+ return done
+ return None
+
+
+class TestHandover:
+ def test_hands_over_once_the_cluster_appears(self):
+ from spyde.signal_tree import _NAV_HANDOVER_MIN_CHUNKS
+ tree = _Tree(client_after=3)
+ at = _drive(tree, total_chunks=100, min_remaining=_NAV_HANDOVER_MIN_CHUNKS)
+ assert at == 3, (
+ "the fill did not hand over on the first chunk after the cluster "
+ "registered — the threaded path would run the whole movie")
+
+ def test_does_not_hand_over_near_the_end(self):
+ """Handing over recomputes what is already painted, so with only a
+ handful of chunks left it costs more than it saves."""
+ from spyde.signal_tree import _NAV_HANDOVER_MIN_CHUNKS
+ tree = _Tree(client_after=0)
+ total = _NAV_HANDOVER_MIN_CHUNKS # every step has <= min remaining
+ assert _drive(tree, total_chunks=total,
+ min_remaining=_NAV_HANDOVER_MIN_CHUNKS) is None
+
+ def test_never_hands_over_without_a_cluster(self):
+ from spyde.signal_tree import _NAV_HANDOVER_MIN_CHUNKS
+ tree = _Tree(client_after=10_000) # cluster never arrives
+ assert _drive(tree, total_chunks=500,
+ min_remaining=_NAV_HANDOVER_MIN_CHUNKS) is None
+
+
+class TestWiring:
+ def test_the_fill_loop_actually_contains_the_handover(self):
+ """Guard against the check being dropped in a refactor: the decision
+ above is only meaningful if the real loop still makes it."""
+ import inspect
+ from spyde.signal_tree import BaseSignalTree
+ src = inspect.getsource(BaseSignalTree._start_progressive_nav_compute)
+ assert "_NAV_HANDOVER_MIN_CHUNKS" in src, (
+ "the threaded navigator fill no longer checks for a late cluster")
+ assert "_start_progressive_nav_compute(_dask" in src, (
+ "the handover no longer re-enters the dispatcher path")
diff --git a/spyde/tests/migrated/test_particle_lifecycle.py b/spyde/tests/migrated/test_particle_lifecycle.py
new file mode 100644
index 00000000..c2d03e6b
--- /dev/null
+++ b/spyde/tests/migrated/test_particle_lifecycle.py
@@ -0,0 +1,250 @@
+"""
+Teardown for everything this feature opens. Nothing may outlive its window.
+
+The feature adds a lot of surfaces — a particle tree, the Drift Check window, the
+dx/dy plot, the ROI preview, overlay widgets, navigator lanes, two wizards — and
+each is a separate chance to leak. The mechanisms exist (``register_cancel``,
+``replace_tree_attr``, ``own_window``, ``figure_registry``), but the presence of a
+mechanism is not evidence it fires, so these tests close things and assert on what
+is left.
+
+**The list-drift hazard.** ``BaseSignalTree.close()`` tears down by iterating
+hard-coded attribute NAME LISTS. A new wizard or result that nobody adds to those
+lists is silently exempt — no error, no warning, just a controller that outlives
+its tree. That already happened here: ``_seg_wizard`` / ``_drift_wizard`` were
+absent from the wizard list, and ``particles`` / ``source_node`` / ``source_tree``
+from the results list, so closing a particle tree kept a back-reference to the
+source movie's lazy array alive and closing the source freed nothing.
+:class:`TestCloseListsCoverThisFeature` exists to make the next omission fail
+rather than leak.
+"""
+from __future__ import annotations
+
+import gc
+import weakref
+
+import numpy as np
+import pytest
+
+import spyde.data.synthetic as sy
+from spyde.actions import figure_registry
+from spyde.actions.particle_tree import open_particle_tree
+from spyde.particles import (
+ LinkParams,
+ SegmentParams,
+ link,
+ measure_frame,
+ segment_frame,
+)
+from spyde.signals.particles import SpyDEParticles
+
+N_FRAMES = 5
+
+
+@pytest.fixture(scope="module")
+def parts():
+ s = sy.particle_movie(n_frames=N_FRAMES)
+ gt = sy.ground_truth(s)
+ per_frame, contours = [], []
+ for t in range(N_FRAMES):
+ lab = segment_frame(s.data[t], SegmentParams(min_size=25, gaussian=1.0))
+ rows, cs = measure_frame(lab, s.data[t], t=t, scale=float(gt["scale"]))
+ per_frame.append(rows)
+ contours.append(cs)
+ p = SpyDEParticles.from_frames(
+ per_frame, frame_shape=tuple(gt["frame_shape"]),
+ contours_per_frame=contours, scale=float(gt["scale"]), units="nm")
+ link(p, LinkParams(max_dist=10.0)).apply(p)
+ return s, gt, p
+
+
+class TestCloseListsCoverThisFeature:
+ """Guards the hard-coded name lists in ``BaseSignalTree.close()``.
+
+ These read the source rather than the behaviour on purpose: the failure mode
+ is an attribute nobody REMEMBERED to list, and a behavioural test only catches
+ the ones you thought to set.
+ """
+
+ def _close_src(self) -> str:
+ import inspect
+
+ from spyde import signal_tree
+ src = inspect.getsource(signal_tree.BaseSignalTree.close)
+ return src
+
+ @pytest.mark.parametrize("attr", ["_seg_wizard", "_drift_wizard"])
+ def test_new_wizards_are_torn_down(self, attr):
+ assert attr in self._close_src(), (
+ f"{attr} is not named in BaseSignalTree.close(), so a live wizard "
+ "controller outlives its tree — silently, because close() iterates "
+ "name lists and an unlisted attribute is simply skipped")
+
+ @pytest.mark.parametrize("attr", [
+ "particles", "_seg_pending_particles", "particle_events",
+ "particle_edits", "nav_traces", "drift", "nav_map",
+ ])
+ def test_new_results_are_cleared(self, attr):
+ assert attr in self._close_src(), f"{attr} survives tree.close()"
+
+ @pytest.mark.parametrize("attr", ["source_node", "source_tree"])
+ def test_back_references_to_the_source_are_cleared(self, attr):
+ """The costly one: these hold the source movie's lazy array."""
+ assert attr in self._close_src(), (
+ f"{attr} survives tree.close(), so a particle tree pins the source "
+ "movie's signal and closing the movie frees nothing")
+
+ def test_particle_overlay_is_torn_down(self):
+ assert "_particle_overlay" in self._close_src()
+
+
+class TestParticleTreeTeardown:
+ def test_close_clears_the_result_and_the_back_reference(self, window, parts):
+ session = window["window"]
+ s, _gt, p = parts
+ tree = open_particle_tree(session, particles=p, source_node=s)
+ assert tree.particles is p and tree.source_node is s
+ tree.close()
+ assert getattr(tree, "particles", None) is None
+ assert getattr(tree, "source_node", None) is None
+ assert getattr(tree, "source_tree", None) is None
+
+ def test_close_drops_the_plots(self, window, parts):
+ session = window["window"]
+ s, _gt, p = parts
+ tree = open_particle_tree(session, particles=p, source_node=s)
+ tree.close()
+ assert tree.signal_plots == []
+ assert tree.navigator_plot_manager is None
+
+ def test_the_source_signal_is_releasable_after_close(self, window, parts):
+ """The leak this was really about, tested by weak reference.
+
+ A tree that keeps `source_node` set pins the movie. Build a THROWAWAY
+ source so the module fixture's own reference does not mask the result.
+ """
+ session = window["window"]
+ _s, gt, p = parts
+ throwaway = sy.particle_movie(n_frames=3)
+ ref = weakref.ref(throwaway)
+ tree = open_particle_tree(session, particles=p, source_node=throwaway)
+ tree.close()
+ del throwaway
+ gc.collect()
+ assert ref() is None, (
+ "the source signal is still reachable after the particle tree was "
+ "closed — something still holds source_node")
+
+ def test_without_close_the_source_is_still_pinned(self, window, parts):
+ """Non-vacuity for the test above.
+
+ If this ALSO collected, the weakref test would be proving nothing about
+ `close()` — it would just be showing that a local went out of scope.
+ """
+ session = window["window"]
+ _s, _gt, p = parts
+ throwaway = sy.particle_movie(n_frames=3)
+ ref = weakref.ref(throwaway)
+ tree = open_particle_tree(session, particles=p, source_node=throwaway)
+ del throwaway
+ gc.collect()
+ assert ref() is not None, (
+ "the source was collected without close() — so the companion test "
+ "does not demonstrate that close() is what releases it")
+ tree.close()
+
+ def test_close_is_idempotent(self, window, parts):
+ session = window["window"]
+ s, _gt, p = parts
+ tree = open_particle_tree(session, particles=p, source_node=s)
+ tree.close()
+ tree.close() # must not raise
+
+
+class TestPendingParticlesTeardown:
+ """A run cancelled before finalize must not leave the placeholder behind."""
+
+ def test_pending_store_is_cleared(self, window, parts):
+ session = window["window"]
+ s, _gt, p = parts
+ tree = open_particle_tree(session, particles=p, source_node=s,
+ attach=False)
+ assert tree.particles is None
+ assert tree._seg_pending_particles is p
+ tree.close()
+ assert getattr(tree, "_seg_pending_particles", None) is None
+
+ def test_batch_flag_does_not_survive_as_a_phantom_run(self, window, parts):
+ """`lifecycle.seg_batch_running` scans live trees; a closed one that kept
+ the flag would make `wait_for_particles` wait on a run that has ended."""
+ from spyde.actions.lifecycle import seg_batch_running
+ session = window["window"]
+ s, _gt, p = parts
+ tree = open_particle_tree(session, particles=p, source_node=s)
+ tree._seg_batch_running = True
+ assert seg_batch_running(session)
+ tree.close()
+ try:
+ session.signal_trees.remove(tree)
+ except ValueError:
+ pass
+ assert not seg_batch_running(session), (
+ "a closed tree still reports a running segmentation batch")
+
+
+class TestFigureRegistry:
+ """Bare-figure windows must not pin their figures past teardown."""
+
+ def test_forget_window_evicts_the_figure(self):
+ marker = object()
+ wid = 987654
+ figure_registry.keep_alive(wid, marker)
+ assert wid in figure_registry._FIGS
+ figure_registry.forget_window(wid)
+ assert wid not in figure_registry._FIGS, (
+ "figure_registry still holds the window's figure after teardown")
+
+ def test_forgetting_an_unknown_window_is_harmless(self):
+ figure_registry.forget_window(123456789)
+
+ def test_registry_holds_no_module_state_for_a_closed_window(self, window,
+ parts):
+ """actions/README.md §3: figure_registry._FIGS is the ONLY module-level
+ mutable state allowed, and it must be evicted by _forget_window."""
+ session = window["window"]
+ s, _gt, p = parts
+ before = set(figure_registry._FIGS)
+ tree = open_particle_tree(session, particles=p, source_node=s)
+ wid = None
+ for plot in tree.signal_plots or []:
+ wid = getattr(plot, "window_id", None)
+ if wid is not None:
+ break
+ tree.close()
+ if wid is not None:
+ session._forget_window(wid)
+ leaked = set(figure_registry._FIGS) - before
+ assert not leaked, f"figures left registered for windows {leaked}"
+
+
+class TestSessionArtifacts:
+ def test_forget_window_drops_controller_and_artifacts(self, window):
+ session = window["window"]
+ wid = 424242
+
+ class _Ctrl:
+ closed = False
+ window_id = wid
+
+ def close(self):
+ _Ctrl.closed = True
+
+ session.register_window_controller(wid, _Ctrl())
+ session._action_artifacts[(wid, "Correct Drift")] = {"selector": None}
+ assert wid in session._window_controllers
+
+ session._forget_window(wid)
+ assert wid not in session._window_controllers, "controller not dropped"
+ assert _Ctrl.closed, "controller.close() was never called"
+ assert not [k for k in session._action_artifacts if k[0] == wid], (
+ "action artifacts survived the window")
diff --git a/spyde/tests/migrated/test_particle_movie_fixture.py b/spyde/tests/migrated/test_particle_movie_fixture.py
new file mode 100644
index 00000000..0652cd31
--- /dev/null
+++ b/spyde/tests/migrated/test_particle_movie_fixture.py
@@ -0,0 +1,388 @@
+"""
+The synthetic particle-movie fixture, and what it is the acceptance gate for.
+
+``spyde.data.synthetic.particle_movie`` is Step 0 of DRIFT_AND_PARTICLES_PLAN.md:
+every later step is checked against a number from this fixture rather than against
+a golden file or a screenshot. So this file has two jobs:
+
+1. **Pin the fixture itself.** If its ground truth drifts out of step with the
+ pixels it generates, every downstream gate silently becomes meaningless.
+2. **Run the end-to-end gates it exists to serve** — drift recovery and
+ segmentation sensitivity — since a fixture nothing consumes is not verified.
+
+The regression in :class:`TestTaperTrap` is the reason this file exists at all:
+the fixture found a defect in the drift solver on its first run.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.data.synthetic import (
+ DISSOLUTION_INDEX,
+ MERGE_PAIR,
+ NUCLEATION_INDEX,
+ ground_truth,
+ particle_movie,
+ particle_truth_at,
+)
+from spyde.drift import solve_translation
+from spyde.particles import SegmentParams, measure_frame, segment_frame
+from spyde.signals.particles import COL
+
+
+@pytest.fixture(scope="module")
+def movie():
+ """One build shared across the module — it takes ~1 s and is deterministic."""
+ s = particle_movie()
+ return s, ground_truth(s)
+
+
+class TestFixtureShape:
+ def test_is_an_insitu_signal(self, movie):
+ s, _ = movie
+ assert type(s).__name__ == "InSitu", (
+ "the movie must cast to the insitu signal type or the Play / "
+ "Fast-Forward toolbar gating never applies to it")
+
+ def test_dimensions_and_dtype(self, movie):
+ s, gt = movie
+ assert s.data.shape == (gt["n_frames"], *tuple(gt["frame_shape"]))
+ assert s.data.dtype == np.float32
+ assert s.axes_manager.navigation_dimension == 1
+ assert s.axes_manager.signal_dimension == 2
+
+ def test_frames_are_non_square(self, movie):
+ """A square fixture hides a transposed frame."""
+ _, gt = movie
+ ny, nx = tuple(gt["frame_shape"])
+ assert ny != nx
+
+ def test_time_axis_is_calibrated(self, movie):
+ s, _ = movie
+ tax = s.axes_manager.navigation_axes[0]
+ assert tax.name == "time" and tax.units == "s"
+ assert tax.scale == pytest.approx(0.05), (
+ "real-time playback reads the time-axis scale; scale=1 makes the "
+ "movie crawl at 1 fps")
+
+ def test_signal_axes_are_calibrated(self, movie):
+ s, gt = movie
+ for ax in s.axes_manager.signal_axes:
+ assert ax.units == "nm"
+ assert ax.scale == pytest.approx(gt["scale"])
+
+ def test_deterministic(self):
+ a = particle_movie(n_frames=6)
+ b = particle_movie(n_frames=6)
+ assert np.array_equal(a.data, b.data)
+
+ def test_seed_changes_only_the_noise_and_film(self):
+ a = particle_movie(n_frames=6, seed=0)
+ b = particle_movie(n_frames=6, seed=1)
+ assert not np.array_equal(a.data, b.data)
+ # The particles are analytic, so the ground truth is seed-independent.
+ assert np.array_equal(ground_truth(a)["p_y0"], ground_truth(b)["p_y0"])
+
+
+class TestAsymmetry:
+ """Every one of these would pass on symmetric data while hiding a real bug."""
+
+ def test_frame_differs_from_its_mirrors_and_transpose(self, movie):
+ s, _ = movie
+ f = s.data[0]
+ assert not np.allclose(f, f[::-1]), "vertically symmetric — a flip would hide"
+ assert not np.allclose(f, f[:, ::-1]), "horizontally symmetric"
+ # Non-square, so a transpose cannot even be compared — that is the point.
+ assert f.shape[0] != f.shape[1]
+
+ def test_frames_differ_from_each_other(self, movie):
+ """A stale-frame bug must be visible."""
+ s, gt = movie
+ for t in range(1, int(gt["n_frames"])):
+ assert not np.allclose(s.data[t], s.data[t - 1]), f"frame {t} == {t-1}"
+
+ def test_drift_axes_have_different_shapes(self, movie):
+ """A swapped axis must show as a wrong-shaped curve, not a wrong number."""
+ _, gt = movie
+ dy, dx = np.asarray(gt["drift"]).T
+ assert np.all(np.diff(dy) >= -1e-9), "dy should be monotonic"
+ assert dx.min() < -0.5 and dx.max() > 0.5, "dx should swing both ways"
+
+
+class TestGroundTruth:
+ def test_drift_starts_at_zero(self, movie):
+ _, gt = movie
+ assert np.allclose(np.asarray(gt["drift"])[0], 0.0)
+
+ def test_event_frames_are_in_range(self, movie):
+ _, gt = movie
+ n = int(gt["n_frames"])
+ for key in ("nucleation_frame", "dissolution_frame", "merge_frame"):
+ assert 0 < int(gt[key]) < n, f"{key}={gt[key]} outside 0..{n}"
+
+ def test_counts_match_the_event_timeline(self, movie):
+ """The count trace must step up at nucleation and down at dissolution."""
+ _, gt = movie
+ n = int(gt["n_frames"])
+ counts = np.array([particle_truth_at(gt, t)[2].sum() for t in range(n)])
+ nuc, dis = int(gt["nucleation_frame"]), int(gt["dissolution_frame"])
+ assert counts[nuc] == counts[nuc - 1] + 1, "no step up at nucleation"
+ assert counts[dis] == counts[dis - 1] - 1, "no step down at dissolution"
+
+ def test_nucleating_particle_is_absent_then_present(self, movie):
+ _, gt = movie
+ i, nuc = NUCLEATION_INDEX, int(gt["nucleation_frame"])
+ assert not particle_truth_at(gt, nuc - 1)[2][i]
+ assert particle_truth_at(gt, nuc)[2][i]
+
+ def test_dissolving_particle_is_present_then_absent(self, movie):
+ _, gt = movie
+ i, dis = DISSOLUTION_INDEX, int(gt["dissolution_frame"])
+ assert particle_truth_at(gt, dis - 1)[2][i]
+ assert not particle_truth_at(gt, dis)[2][i]
+
+ def test_merge_pair_converges_and_overlaps_at_the_stamped_frame(self, movie):
+ _, gt = movie
+ a, b = tuple(gt["merge_pair"])
+ radii = np.asarray(gt["p_radius"])
+ touch = radii[a] + radii[b]
+ mf = int(gt["merge_frame"])
+
+ def gap(t):
+ pos = particle_truth_at(gt, t)[0]
+ return float(np.hypot(*(pos[a] - pos[b])))
+
+ assert gap(0) > touch, "the merge pair already overlaps at t=0"
+ assert gap(mf) <= touch, f"no overlap at the stamped merge_frame {mf}"
+ assert gap(mf - 1) > touch, f"they already overlapped before frame {mf}"
+
+ def test_faint_probes_are_faint_but_findable(self, movie):
+ """The Section 0.9 gate: above noise, well below the bright particles."""
+ _, gt = movie
+ faint = np.asarray(gt["p_faint"], bool)
+ amps = np.asarray(gt["p_amp"])
+ noise = float(gt["noise"])
+ assert faint.sum() == 2
+ assert amps[faint].max() < 0.25 * amps[~faint].min(), "not actually faint"
+ assert amps[faint].min() / noise > 4.0, "buried in noise, unfindable"
+
+ def test_ground_truth_raises_on_a_plain_signal(self):
+ import hyperspy.api as hs
+ with pytest.raises(ValueError, match="no synthetic ground truth"):
+ ground_truth(hs.signals.Signal2D(np.zeros((4, 4))))
+
+
+class TestTruthMatchesPixels:
+ """The fixture's stamped truth must describe the pixels it actually drew."""
+
+ def test_every_present_particle_is_brighter_than_its_surroundings(self, movie):
+ s, gt = movie
+ for t in (0, 12, int(gt["n_frames"]) - 1):
+ f = s.data[t]
+ pos, radii, present = particle_truth_at(gt, t)
+ faint = np.asarray(gt["p_faint"], bool)
+ for i in np.flatnonzero(present):
+ if faint[i]:
+ continue # covered by the sensitivity test
+ cy, cx = int(round(pos[i, 0])), int(round(pos[i, 1]))
+ if not (2 <= cy < f.shape[0] - 2 and 2 <= cx < f.shape[1] - 2):
+ continue # drifted out of frame
+ centre = f[cy - 1:cy + 2, cx - 1:cx + 2].mean()
+ assert centre > np.median(f) + 0.2, (
+ f"frame {t} particle {i} at ({cy},{cx}) is not bright")
+
+ def test_absent_particles_leave_nothing_behind(self, movie):
+ """A dissolved particle must actually be gone from the pixels."""
+ s, gt = movie
+ i = DISSOLUTION_INDEX
+ dis = int(gt["dissolution_frame"])
+ before, after = s.data[dis - 1], s.data[dis]
+ pos = particle_truth_at(gt, dis)[0]
+ cy, cx = int(round(pos[i, 0])), int(round(pos[i, 1]))
+ w = 3
+ assert (before[cy - w:cy + w, cx - w:cx + w].mean()
+ > after[cy - w:cy + w, cx - w:cx + w].mean() + 0.3), (
+ "the dissolving particle is still in the pixels after its death frame")
+
+
+class TestDriftRecoveryGate:
+ """Plan gate A1, run against the fixture rather than a hand-made stack."""
+
+ def test_recovers_the_applied_drift(self, movie):
+ s, gt = movie
+ truth = np.asarray(gt["drift"])
+ model = solve_translation(s.data, device="numpy", upsample=8,
+ reference="first", max_shift=20)
+ err = np.abs(model.shifts - truth)
+ assert err.max() < 0.5, f"max drift error {err.max():.3f} px\n{model.shifts}"
+
+ def test_running_reference_also_works(self, movie):
+ s, gt = movie
+ truth = np.asarray(gt["drift"])
+ model = solve_translation(s.data, device="numpy", upsample=8, max_shift=20)
+ assert np.abs(model.shifts - truth).max() < 0.5
+
+ def test_correction_flattens_the_drift(self, movie):
+ """End to end: solve, apply, and the film should stop moving."""
+ from spyde.drift import shift_frame
+ s, gt = movie
+ model = solve_translation(s.data, device="numpy", upsample=8,
+ reference="first", max_shift=20)
+ core = (slice(20, -20), slice(20, -20))
+ ref = s.data[0][core]
+ last = int(gt["n_frames"]) - 1
+ raw = float(np.abs(s.data[last][core] - ref).mean())
+ fixed = shift_frame(s.data[last], model.shifts[last], fill=0.0)
+ corrected = float(np.abs(fixed[core] - ref).mean())
+ assert corrected < raw, (
+ f"correction made it worse (raw {raw:.4f} -> {corrected:.4f}) — "
+ "check the SIGN convention")
+
+
+class TestTaperTrap:
+ """A full Hann window destroys this registration. Do not re-enable it.
+
+ This is a real defect the fixture caught on its first run: with ``apodize=1.0``
+ the solve returns a 25 px error on a 6 px drift — worse than not correcting at
+ all. It is NOT specific to our implementation; ``skimage``'s
+ ``phase_cross_correlation`` returns the same wrong answer on the same windowed
+ input, because a full-frame window reweights the two frames' content
+ differently once the drift is large and manufactures a false peak.
+ """
+
+ def test_default_taper_is_a_partial_edge_taper(self):
+ from spyde.drift.translation import DEFAULT_TAPER_ALPHA
+ assert 0.0 < DEFAULT_TAPER_ALPHA < 0.6, (
+ "the default must taper only the EDGE; alpha near 1.0 is the trap "
+ "this class documents")
+
+ def test_full_hann_is_dramatically_worse_than_the_default(self, movie):
+ s, gt = movie
+ truth = np.asarray(gt["drift"])
+ good = solve_translation(s.data, device="numpy", upsample=8,
+ reference="first", max_shift=20)
+ hann = solve_translation(s.data, device="numpy", upsample=8,
+ reference="first", max_shift=20, apodize=1.0)
+ e_good = np.abs(good.shifts - truth).max()
+ e_hann = np.abs(hann.shifts - truth).max()
+ assert e_good < 0.5, f"the default taper regressed: {e_good:.3f} px"
+ assert e_hann > 5.0 * e_good, (
+ "full Hann is no longer catastrophic here — if the solver changed so "
+ "that it is safe, this test has served its purpose and can go, but "
+ f"check deliberately (default {e_good:.3f} px vs hann {e_hann:.3f} px)")
+
+ def test_apodize_records_the_alpha_actually_used(self, movie):
+ s, _ = movie
+ assert solve_translation(s.data[:3], device="numpy",
+ apodize=False).params["apodize"] == 0.0
+ assert solve_translation(s.data[:3], device="numpy",
+ apodize=0.3).params["apodize"] == pytest.approx(0.3)
+
+
+class TestHarnessLoader:
+ """``load_test_data_particles`` — the door the e2e specs come through."""
+
+ def test_action_is_registered(self):
+ from spyde.backend._session_actions import _TEST_ACTIONS
+ assert "load_test_data_particles" in _TEST_ACTIONS, (
+ "the action is not in _TEST_ACTIONS, so dispatch will reject it and "
+ "every e2e spec silently loads nothing")
+
+ def test_loads_lazily_one_frame_per_chunk(self, window):
+ session = window["window"]
+ session._load_test_data_particles({"frames": 6})
+ assert len(window["signal_trees"]) == 1
+ root = window["signal_trees"][0].root
+ assert root._lazy, "must be lazy — an eager fixture skips the cache path"
+ assert root.data.chunksize[0] == 1, (
+ f"expected one frame per chunk, got {root.data.chunksize} — that is "
+ "what makes each nav move a small cold read like a real .mrc")
+
+ def test_signal_type_and_axes_survive_the_lazy_rewrap(self, window):
+ session = window["window"]
+ session._load_test_data_particles({"frames": 6})
+ root = window["signal_trees"][0].root
+ assert getattr(root, "_signal_type", None) == "insitu", (
+ "the insitu cast was lost, so Play / Fast-Forward will not appear")
+ tax = root.axes_manager.navigation_axes[0]
+ assert tax.name == "time" and tax.scale == pytest.approx(0.05)
+ assert root.axes_manager.signal_axes[0].units == "nm"
+
+ def test_ground_truth_survives_the_lazy_rewrap(self, window):
+ """The whole point of the fixture is its stamped truth — losing it in the
+ re-wrap would leave every downstream gate asserting against nothing."""
+ session = window["window"]
+ session._load_test_data_particles({"frames": 6})
+ gt = ground_truth(window["signal_trees"][0].root)
+ assert gt["kind"] == "particle_movie"
+ assert int(gt["nucleation_frame"]) == NUCLEATION_INDEX + 5 # 8
+ assert np.asarray(gt["drift"]).shape == (6, 2)
+
+ def test_eager_option(self, window):
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4, "eager": True})
+ assert not window["signal_trees"][0].root._lazy
+
+ def test_opens_two_windows(self, window):
+ """A 1-D-nav movie gives a navigator plus a signal window."""
+ session = window["window"]
+ session._load_test_data_particles({"frames": 6})
+ assert len(session._plots) >= 2, (
+ f"expected navigator + signal plots, got {len(session._plots)}")
+
+
+class TestSegmentationOnTheFixture:
+ """Plan gate B1/B5 against known radii, plus the Section 0.9 sensitivity gate."""
+
+ def test_finds_every_bright_particle(self, movie):
+ s, gt = movie
+ t = 12 # all nine present
+ labels = segment_frame(s.data[t], SegmentParams(min_size=25, gaussian=1.0))
+ pos, _, present = particle_truth_at(gt, t)
+ faint = np.asarray(gt["p_faint"], bool)
+ want = np.flatnonzero(present & ~faint)
+ found = 0
+ for i in want:
+ cy, cx = int(round(pos[i, 0])), int(round(pos[i, 1]))
+ if labels[cy, cx] != 0:
+ found += 1
+ assert found == len(want), f"found {found} of {len(want)} bright particles"
+
+ def test_default_sensitivity_misses_the_faint_probes(self, movie):
+ """Establishes that the sensitivity gate below is not vacuous."""
+ s, gt = movie
+ t = 12
+ labels = segment_frame(s.data[t], SegmentParams(min_size=25, gaussian=1.0))
+ pos = particle_truth_at(gt, t)[0]
+ faint = np.flatnonzero(np.asarray(gt["p_faint"], bool))
+ hit = sum(labels[int(round(pos[i, 0])), int(round(pos[i, 1]))] != 0
+ for i in faint)
+ assert hit < len(faint), (
+ "the faint probes are already found at default sensitivity — they are "
+ "not faint enough to test anything")
+
+ def test_measured_radii_match_the_truth(self, movie):
+ s, gt = movie
+ t = 12
+ labels = segment_frame(s.data[t], SegmentParams(min_size=25, gaussian=1.0))
+ rows, _ = measure_frame(labels, s.data[t], t=t, scale=1.0)
+ pos, radii, present = particle_truth_at(gt, t)
+ faint = np.asarray(gt["p_faint"], bool)
+ checked = 0
+ for i in np.flatnonzero(present & ~faint):
+ cy, cx = int(round(pos[i, 0])), int(round(pos[i, 1]))
+ lbl = labels[cy, cx]
+ if lbl == 0:
+ continue
+ row = rows[rows[:, COL["label"]] == lbl]
+ if not len(row):
+ continue
+ # Only isolated particles: a merged pair's area is not one disc.
+ if i in tuple(gt["merge_pair"]):
+ continue
+ got_r = float(row[0, COL["equiv_diameter"]]) / 2.0
+ assert abs(got_r - radii[i]) < 0.35 * radii[i], (
+ f"particle {i}: measured r={got_r:.2f} vs truth {radii[i]:.2f}")
+ checked += 1
+ assert checked >= 3, f"only checked {checked} particles — test too weak"
diff --git a/spyde/tests/migrated/test_particle_overlay.py b/spyde/tests/migrated/test_particle_overlay.py
new file mode 100644
index 00000000..9cdaa52f
--- /dev/null
+++ b/spyde/tests/migrated/test_particle_overlay.py
@@ -0,0 +1,1108 @@
+"""
+The particle overlay (plan B9) and the stacked navigator lanes (C2/C3).
+
+Built on a REAL segmentation of the synthetic movie — ``segment_frame`` →
+``measure_frame`` → ``SpyDEParticles.from_frames`` → ``link`` — rather than a
+hand-made container, because three of the claims here are only meaningful
+against real geometry: the calibrated→pixel conversion (the fixture is
+``scale=0.5`` nm/px, so a missing division is visible), the nearest-centroid hit
+test, and the split/merge round trip through ``measure_frame``.
+
+Five things are worth more than the rest:
+
+:class:`TestPixelConversion`
+ Centroids are CALIBRATED and must be divided by ``particles.scale``;
+ contours are already PIXELS and must not be. Both are right at scale 1 and
+ wrong everywhere else, which is exactly how ``masks._signal_k_grids``'s bug
+ class survives review.
+:class:`TestTrails`
+ **A dead track draws no head dot.** The dot means "the particle is HERE
+ NOW"; on a track that has died — or one inside its ``memory`` gap — it reads
+ as a real particle the segmenter stopped filling (plan C3).
+:class:`TestEdits`
+ Delete / merge / split mutate the store IN PLACE (the lazy label movie closes
+ over that object) and are recorded on the tree AND in provenance.
+:class:`TestNavigatorLanes`
+ The count lane is integer data and is emitted as STEP data — a straight
+ interpolation puts a nucleation at 7 when the event is at 8.
+:class:`TestLifecycle`
+ The overlay lives on the tree and ``BaseSignalTree.close()`` reaps it.
+"""
+from __future__ import annotations
+
+import time
+
+import numpy as np
+import pytest
+
+import spyde.data.synthetic as sy
+from spyde.actions import particle_overlay as po
+from spyde.actions.particle_tree import open_particle_tree
+from spyde.particles import (
+ LinkParams,
+ SegmentParams,
+ link,
+ measure_frame,
+ segment_frame,
+)
+from spyde.signals.particles import COL, N_COLUMNS, SpyDEParticles
+
+N_FRAMES = 8
+
+
+# ── fixtures ─────────────────────────────────────────────────────────────────
+
+@pytest.fixture(scope="module")
+def built():
+ """A real segmentation + link of the fixture movie (copied from
+ ``test_particle_tree.py``'s ``built`` — same door, same numbers)."""
+ s = sy.particle_movie(n_frames=N_FRAMES)
+ gt = sy.ground_truth(s)
+ scale = float(gt["scale"])
+ per_frame, contours = [], []
+ for t in range(N_FRAMES):
+ lab = segment_frame(s.data[t], SegmentParams(min_size=25, gaussian=1.0))
+ rows, cs = measure_frame(lab, s.data[t], t=t, scale=scale)
+ per_frame.append(rows)
+ contours.append(cs)
+ parts = SpyDEParticles.from_frames(
+ per_frame, frame_shape=tuple(gt["frame_shape"]),
+ contours_per_frame=contours, scale=scale, units="nm")
+ res = link(parts, LinkParams(max_dist=10.0))
+ res.apply(parts)
+ return s, gt, parts, res
+
+
+def _fresh(built):
+ """A private copy of the built store — the edit tests mutate it in place."""
+ _s, _gt, parts, _res = built
+ return SpyDEParticles(
+ parts.flat_buffer.copy(), parts.t_offsets.copy(), parts.frame_shape,
+ contours=parts.contours.copy(), contour_offsets=parts.contour_offsets.copy(),
+ scale=parts.scale, units=parts.units)
+
+
+@pytest.fixture
+def make_tree(window, built):
+ """Build particle trees and CLOSE them on teardown.
+
+ Not tidiness: the label movie is lazy, so a live tree runs a progressive
+ navigator compute on a daemon thread. Leaving one running past the end of the
+ test lets the interpreter finalise underneath it, which prints a truncated
+ traceback that looks like a failure and is not. ``close()`` sets the tree's
+ ``_nav_stop`` and is re-entrancy guarded, so a test that closes its own tree
+ is unaffected.
+ """
+ session = window["window"]
+ s, _gt, parts, res = built
+ made = []
+
+ def build(*, events=True, **kwargs):
+ tree = open_particle_tree(session, particles=parts, source_node=s,
+ events=(res.events if events else None), **kwargs)
+ made.append(tree)
+ assert _wait(lambda: po._first_signal_plot(tree) is not None)
+ assert _wait(lambda: po._first_nav_plot(tree) is not None)
+ return tree
+
+ yield build
+ for tree in made:
+ try:
+ tree.close()
+ except Exception:
+ pass
+
+
+def _detached(particles, **kwargs) -> po.ParticleOverlay:
+ """An overlay with no plot: every payload/selection/edit method works without
+ a figure, which is what makes the geometry assertable headlessly."""
+ return po.ParticleOverlay(None, particles, **kwargs)
+
+
+def _signal_plot(session):
+ return next((p for p in session._plots
+ if not p.is_navigator and p.plot_state is not None), None)
+
+
+def _wait(pred, timeout=20.0):
+ end = time.time() + timeout
+ while time.time() < end:
+ if pred():
+ return True
+ time.sleep(0.05)
+ return False
+
+
+def _csr_is_consistent(particles) -> None:
+ """Every invariant ``SpyDEParticles.__post_init__`` would have enforced."""
+ assert particles.flat_buffer.shape[1] == N_COLUMNS
+ assert int(particles.t_offsets[0]) == 0
+ assert int(particles.t_offsets[-1]) == len(particles.flat_buffer)
+ assert np.all(np.diff(particles.t_offsets) >= 0)
+ if particles.contours is not None:
+ assert particles.contour_offsets.size == len(particles.flat_buffer) + 1
+ assert int(particles.contour_offsets[-1]) == len(particles.contours)
+ # The ``t`` column and the CSR block a row sits in must still agree.
+ frames = np.repeat(np.arange(particles.n_frames), np.diff(particles.t_offsets))
+ assert np.array_equal(frames, particles.flat_buffer[:, COL["t"]].astype(np.int64))
+
+
+# ── colour ───────────────────────────────────────────────────────────────────
+
+class TestColorCycle:
+ def test_track_zero_is_blue(self):
+ assert po.track_color(0) == "#89b4fa"
+ assert po.TRACK_COLORS[0] == "#89b4fa"
+
+ def test_cycle_is_stable_and_wraps_at_six(self):
+ assert len(po.TRACK_COLORS) == 6
+ for tid in range(24):
+ assert po.track_color(tid) == po.TRACK_COLORS[tid % 6]
+ # Stable means a pure function of the id — same answer every call, and
+ # the same answer for ids six apart.
+ assert po.track_color(3) == po.track_color(9) == po.track_color(15)
+
+ def test_the_six_accents_are_distinct(self):
+ assert len(set(po.TRACK_COLORS)) == 6
+
+ def test_untracked_is_grey_not_a_seventh_accent(self):
+ assert po.track_color(-1) == po.UNTRACKED_COLOR
+ assert po.UNTRACKED_COLOR not in po.TRACK_COLORS
+
+ def test_fade_appends_an_alpha_byte(self):
+ assert po.fade("#89b4fa", 1.0) == "#89b4faff"
+ assert po.fade("#89b4fa", 0.0) == "#89b4fa00"
+ assert len(po.fade("#89b4fa", 0.5)) == 9
+
+ def test_trail_alphas_run_newest_brightest(self):
+ alphas = po.trail_alphas(4)
+ assert alphas[0] == 1.0
+ assert alphas == sorted(alphas, reverse=True)
+ assert min(alphas) > 0.0, "the oldest trail step must still be visible"
+
+ def test_real_tracks_land_on_distinct_colours(self, built):
+ _s, _gt, parts, res = built
+ assert res.n_tracks == 6
+ colors = {po.track_color(t) for t in range(res.n_tracks)}
+ assert colors == set(po.TRACK_COLORS)
+
+
+# ── the calibrated → pixel conversion ────────────────────────────────────────
+
+class TestPixelConversion:
+ def test_fixture_scale_is_not_one(self, built):
+ """The whole point of testing on this fixture."""
+ _s, _gt, parts, _res = built
+ assert parts.scale == 0.5
+
+ def test_centroids_are_divided_by_scale(self, built):
+ _s, _gt, parts, _res = built
+ rows = parts.at(0)
+ px = po.centroids_px(rows, parts.scale)
+ assert np.allclose(px[:, 0], rows[:, COL["x"]] / 0.5)
+ assert np.allclose(px[:, 1], rows[:, COL["y"]] / 0.5)
+ # And they land INSIDE the frame in pixel space (a missing division would
+ # put every marker in the top-left quadrant of a 96x112 frame).
+ h, w = parts.frame_shape
+ assert px[:, 0].max() < w and px[:, 1].max() < h
+
+ def test_marker_offsets_are_x_then_y(self, built):
+ """A property row is (y, x); a marker offset is (x, y)."""
+ _s, _gt, parts, _res = built
+ row = parts.at(0)[:1]
+ px = po.centroids_px(row, parts.scale)
+ assert px[0, 0] == pytest.approx(row[0, COL["x"]] / 0.5)
+ assert px[0, 1] == pytest.approx(row[0, COL["y"]] / 0.5)
+ assert px[0, 0] != pytest.approx(px[0, 1])
+
+ def test_contours_are_NOT_divided(self, built):
+ """Contours are stored in pixels already — dividing them by scale would
+ shrink every outline to a quarter of its body and leave it detached from
+ the centroid it belongs to."""
+ _s, _gt, parts, _res = built
+ gi = int(parts.indices_at(0)[0])
+ raw = parts.contour_at(gi)
+ poly = po.contour_xy(parts, gi)
+ assert np.allclose(poly[:, 0], raw[:, 1]) # x = column
+ assert np.allclose(poly[:, 1], raw[:, 0]) # y = row
+
+ def test_centroid_sits_inside_its_own_outline(self, built):
+ """The conversion is only right if the two agree in ONE space."""
+ _s, _gt, parts, _res = built
+ for gi in parts.indices_at(0):
+ poly = po.contour_xy(parts, int(gi))
+ if len(poly) < 3:
+ continue
+ cx, cy = po.centroids_px(parts.flat_buffer[int(gi):int(gi) + 1],
+ parts.scale)[0]
+ assert poly[:, 0].min() - 1 <= cx <= poly[:, 0].max() + 1
+ assert poly[:, 1].min() - 1 <= cy <= poly[:, 1].max() + 1
+
+ def test_a_scale_one_store_needs_no_division(self, built):
+ """Sanity: at scale=1 the two conventions coincide, which is exactly why
+ a scale=1 fixture cannot catch this class of bug."""
+ _s, _gt, parts, _res = built
+ rows = parts.at(0)
+ assert np.allclose(po.centroids_px(rows, 1.0)[:, 0], rows[:, COL["x"]])
+
+ def test_click_is_converted_through_the_PLOT_axes(self, built):
+ """A click carries PHYSICAL xdata/ydata; markers are in pixels.
+
+ The conversion goes through the displayed signal's axes, not through
+ ``particles.scale`` — the store's calibration and the plot's need not be
+ the same number.
+ """
+ _s, _gt, parts, _res = built
+
+ class _Axis:
+ def __init__(self, scale, offset):
+ self.scale, self.offset = scale, offset
+
+ class _AxesManager:
+ signal_axes = (_Axis(0.25, 3.0), _Axis(0.25, 3.0))
+
+ class _Signal:
+ axes_manager = _AxesManager()
+
+ class _State:
+ current_signal = _Signal()
+
+ class _Plot:
+ plot_state = _State()
+
+ class _Event:
+ xdata, ydata = 3.0 + 0.25 * 40, 3.0 + 0.25 * 12
+
+ overlay = po.ParticleOverlay(_Plot(), parts)
+ assert overlay._event_px(_Event()) == pytest.approx((40.0, 12.0))
+
+ def test_uncalibrated_plot_is_an_identity(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+
+ class _Event:
+ xdata, ydata = 17.0, 5.0
+ assert overlay._event_px(_Event()) == pytest.approx((17.0, 5.0))
+
+
+# ── the payload ──────────────────────────────────────────────────────────────
+
+class TestPayload:
+ def test_every_particle_is_filled_exactly_once(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(3)
+ payload = overlay._payload(3)
+ drawn = sum(len(payload[f"fill{i}"]["vertices_list"])
+ for i in range(len(po.TRACK_COLORS) + 1))
+ assert drawn == len(parts.at(3))
+
+ def test_fill_group_is_the_track_colour_bucket(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ payload = overlay._payload(2)
+ for gi in parts.indices_at(2):
+ tid = int(parts.flat_buffer[gi, COL["track_id"]])
+ bucket = tid % len(po.TRACK_COLORS)
+ polys = payload[f"fill{bucket}"]["vertices_list"]
+ expected = po.contour_xy(parts, int(gi))
+ assert any(p.shape == expected.shape and np.allclose(p, expected)
+ for p in polys), f"particle {gi} is not in bucket {bucket}"
+
+ def test_untracked_rows_go_to_the_grey_bucket(self, built):
+ parts = _fresh(built)
+ parts.flat_buffer[:, COL["track_id"]] = -1.0
+ overlay = _detached(parts)
+ payload = overlay._payload(0)
+ grey = len(po.TRACK_COLORS)
+ assert len(payload[f"fill{grey}"]["vertices_list"]) == len(parts.at(0))
+ for i in range(grey):
+ assert payload[f"fill{i}"]["vertices_list"] == []
+
+ def test_labels_are_empty_until_something_is_selected(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ assert overlay._payload(0)["labels"]["texts"] == []
+ assert overlay._payload(0)["selected"]["vertices_list"] == []
+
+ def test_selection_gets_an_outline_and_a_readout(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ gi = int(parts.indices_at(0)[1])
+ overlay.select([gi])
+ payload = overlay._payload(0)
+ assert len(payload["selected"]["vertices_list"]) == 1
+ assert len(payload["labels"]["texts"]) == 1
+ text = payload["labels"]["texts"][0]
+ assert "track" in text and "area" in text
+
+ def test_hover_labels_without_selecting(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.hovered = int(parts.indices_at(0)[0])
+ payload = overlay._payload(0)
+ assert len(payload["labels"]["texts"]) == 1
+ assert payload["selected"]["vertices_list"] == [], \
+ "hover must label, not outline — that is the selection's job"
+
+ def test_a_selection_in_another_frame_is_not_drawn(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.select([int(parts.indices_at(5)[0])])
+ payload = overlay._payload(0)
+ assert payload["selected"]["vertices_list"] == []
+ assert payload["labels"]["texts"] == []
+
+ def test_out_of_range_frame_draws_nothing(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ payload = overlay._payload(999)
+ for i in range(len(po.TRACK_COLORS) + 1):
+ assert payload[f"fill{i}"]["vertices_list"] == []
+
+
+# ── trails and the head dot ──────────────────────────────────────────────────
+
+class TestTrails:
+ def _heads(self, payload) -> dict[int, int]:
+ return {i: len(payload[f"head{i}"]["offsets"])
+ for i in range(len(po.TRACK_COLORS) + 1)}
+
+ def _segments(self, payload, bucket) -> int:
+ return sum(len(payload[f"trail{bucket}_{s}"]["segments"])
+ for s in range(po.TRAIL_FADE_STEPS))
+
+ def test_trails_off_draws_nothing(self, built):
+ _s, _gt, parts, _res = built
+ payload = _detached(parts, show_trails=False)._payload(5)
+ assert sum(self._heads(payload).values()) == 0
+ assert all(self._segments(payload, i) == 0
+ for i in range(len(po.TRACK_COLORS) + 1))
+
+ def test_a_live_track_gets_a_head_dot(self, built):
+ _s, _gt, parts, _res = built
+ payload = _detached(parts, show_trails=True)._payload(5)
+ assert sum(self._heads(payload).values()) == len(parts.at(5))
+
+ def test_head_dot_sits_on_the_centroid_in_pixels(self, built):
+ _s, _gt, parts, _res = built
+ payload = _detached(parts, show_trails=True)._payload(5)
+ gi = int(parts.indices_at(5)[0])
+ bucket = int(parts.flat_buffer[gi, COL["track_id"]]) % len(po.TRACK_COLORS)
+ expected = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ heads = np.asarray(payload[f"head{bucket}"]["offsets"])
+ assert np.allclose(heads[0], expected)
+
+ def test_trail_spans_the_window_and_fades(self, built):
+ _s, _gt, parts, _res = built
+ payload = _detached(parts, show_trails=True, trail_frames=5)._payload(6)
+ gi = int(parts.indices_at(6)[0])
+ bucket = int(parts.flat_buffer[gi, COL["track_id"]]) % len(po.TRACK_COLORS)
+ # 5 frames of window → at most 4 segments per track.
+ assert 1 <= self._segments(payload, bucket) <= 4
+ # More than one fade step is populated, i.e. the ramp is actually used.
+ used = [s for s in range(po.TRAIL_FADE_STEPS)
+ if len(payload[f"trail{bucket}_{s}"]["segments"])]
+ assert len(used) >= 2
+
+ def test_a_DEAD_track_draws_no_head_dot(self, built):
+ """The plan's C3 note, asserted directly.
+
+ Kill one track after frame 4 (delete its later detections), then look at
+ frame 6 — still inside the trailing window, so the fading line is there,
+ but the dot must be gone: it would read as a particle that is present
+ now, which is precisely what it is not.
+ """
+ parts = _fresh(built)
+ dead = int(parts.at(0)[0, COL["track_id"]])
+ doomed = [int(gi) for t in range(5, parts.n_frames)
+ for gi in parts.indices_at(t)
+ if int(parts.flat_buffer[gi, COL["track_id"]]) == dead]
+ assert doomed, "the fixture track never reached the later frames"
+ po.delete_particles(parts, doomed)
+
+ overlay = _detached(parts, show_trails=True, trail_frames=8)
+ payload = overlay._payload(6)
+ bucket = dead % len(po.TRACK_COLORS)
+ assert self._segments(payload, bucket) > 0, \
+ "the dead track's trail should still fade out"
+ assert len(payload[f"head{bucket}"]["offsets"]) == 0, \
+ "a dead track drew a head dot — it reads as a live particle"
+ # Every OTHER track is unaffected.
+ live = sum(len(payload[f"head{i}"]["offsets"])
+ for i in range(len(po.TRACK_COLORS) + 1) if i != bucket)
+ assert live == len(parts.at(6))
+
+ def test_a_track_inside_its_memory_GAP_draws_no_head_dot(self, built):
+ """Same rule, the other half: a track with no detection at *t* has no
+ current position, so it gets no dot even though it is not dead."""
+ parts = _fresh(built)
+ tid = int(parts.at(0)[1, COL["track_id"]])
+ gap = [int(gi) for gi in parts.indices_at(4)
+ if int(parts.flat_buffer[gi, COL["track_id"]]) == tid]
+ assert gap
+ po.delete_particles(parts, gap)
+
+ payload = _detached(parts, show_trails=True, trail_frames=8)._payload(4)
+ bucket = tid % len(po.TRACK_COLORS)
+ assert self._segments(payload, bucket) > 0
+ assert len(payload[f"head{bucket}"]["offsets"]) == 0
+
+ def test_untracked_rows_get_no_trail(self, built):
+ """An untracked row has no trajectory to draw — inventing one by
+ colour-bucket would join unrelated detections into a fake track."""
+ parts = _fresh(built)
+ parts.flat_buffer[:, COL["track_id"]] = -1.0
+ payload = _detached(parts, show_trails=True)._payload(5)
+ assert sum(self._heads(payload).values()) == 0
+
+
+# ── selection ────────────────────────────────────────────────────────────────
+
+class TestSelection:
+ def test_select_by_index(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ want = [int(i) for i in parts.indices_at(2)[:2]]
+ assert overlay.select(want) == want
+ assert overlay.selected == want
+
+ def test_select_drops_out_of_range_indices(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ assert overlay.select([0, 10_000, -5]) == [0]
+
+ def test_select_by_track(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(4)
+ tid = int(parts.at(4)[2, COL["track_id"]])
+ picked = overlay.select_track(tid)
+ assert len(picked) == 1
+ assert int(parts.flat_buffer[picked[0], COL["track_id"]]) == tid
+ assert int(parts.flat_buffer[picked[0], COL["t"]]) == 4
+
+ def test_click_picks_the_nearest_centroid(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(1)
+ gi = int(parts.indices_at(1)[3])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ assert overlay.pick(cx + 0.5, cy - 0.5) == gi
+
+ def test_click_on_empty_space_selects_nothing(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(1)
+ assert overlay.pick(1.0, 1.0) is None
+
+ def test_click_handler_selects_through_the_plot_axes(self, built):
+ """End-to-end for the click path: a physical xdata/ydata → a selection."""
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(1)
+ overlay._groups = {"sentinel": object()} # make the handler non-inert
+ gi = int(parts.indices_at(1)[2])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+
+ class _Event: # an uncalibrated plot → xdata IS the pixel column
+ xdata, ydata = float(cx), float(cy)
+ overlay._on_click(_Event())
+ assert overlay.selected == [gi]
+
+ def test_region_selects_in_bulk(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(0)
+ pts = po.centroids_px(parts.at(0), parts.scale)
+ # A box around the left half of the frame.
+ picked = overlay.select_region(0, 0, 60, 200)
+ expected = int(((pts[:, 0] >= 0) & (pts[:, 0] <= 60)).sum())
+ assert len(picked) == expected >= 1
+ assert len(picked) < len(parts.at(0)), \
+ "the box should not have caught every particle"
+
+ def test_region_is_order_insensitive(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.set_frame(0)
+ a = overlay.select_region(10, 10, 80, 90)
+ b = overlay.select_region(80, 90, 10, 10)
+ assert a == b
+
+ def test_clear_selection(self, built):
+ _s, _gt, parts, _res = built
+ overlay = _detached(parts)
+ overlay.select([0, 1])
+ assert overlay.clear_selection() == []
+
+ def test_on_select_callback_fires(self, built):
+ _s, _gt, parts, _res = built
+ seen = []
+ overlay = _detached(parts, on_select=lambda: seen.append(1))
+ overlay.select([0])
+ assert seen == [1]
+
+
+# ── editing ──────────────────────────────────────────────────────────────────
+
+class TestEdits:
+ def test_delete_removes_the_row_and_keeps_the_CSR_valid(self, built):
+ parts = _fresh(built)
+ before = parts.n_particles
+ gi = int(parts.indices_at(3)[1])
+ tid = int(parts.flat_buffer[gi, COL["track_id"]])
+ assert po.delete_particles(parts, [gi]) == 1
+ assert parts.n_particles == before - 1
+ _csr_is_consistent(parts)
+ assert tid not in parts.at(3)[:, COL["track_id"]].astype(int)
+
+ def test_delete_mutates_the_store_IN_PLACE(self, built):
+ """The lazy label movie closes over this object (particle_tree §0.6), so
+ an edit that rebound ``tree.particles`` would leave the open window
+ rendering the pre-edit contours forever."""
+ parts = _fresh(built)
+ buffer_before = parts.flat_buffer
+ po.delete_particles(parts, [0])
+ assert parts.flat_buffer is not buffer_before, "buffers are rebuilt…"
+ # …but the STORE is the same object, which is what the movie holds.
+ overlay = _detached(parts)
+ assert overlay.particles is parts
+
+ def test_delete_keeps_contours_paired_with_their_rows(self, built):
+ parts = _fresh(built)
+ survivor = int(parts.indices_at(0)[2])
+ expected = parts.contour_at(survivor).copy()
+ po.delete_particles(parts, [int(parts.indices_at(0)[0])])
+ _csr_is_consistent(parts)
+ assert np.array_equal(parts.contour_at(survivor - 1), expected), \
+ "a deletion re-paired an outline with the wrong property row"
+
+ def test_delete_through_the_overlay_records_the_edit(self, built):
+ parts = _fresh(built)
+
+ class _Tree:
+ particles = parts
+ tree = _Tree()
+ overlay = _detached(parts)
+ overlay.tree = tree
+ gi = int(parts.indices_at(2)[0])
+ overlay.select([gi])
+ assert overlay.delete() == 1
+ assert len(tree.particle_edits) == 1
+ record = tree.particle_edits[0]
+ assert record["kind"] == "delete" and record["indices"] == [gi]
+ assert parts.provenance["edits"][0]["kind"] == "delete"
+ assert overlay.selected == []
+
+ def test_split_makes_two_and_conserves_the_body(self, built):
+ parts = _fresh(built)
+ gi = int(parts.indices_at(0)[0])
+ area = float(parts.flat_buffer[gi, COL["area"]])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ before = parts.n_particles
+
+ a, b = po.split_particle(parts, gi, ((cx, cy - 50), (cx, cy + 50)))
+ assert parts.n_particles == before + 1
+ _csr_is_consistent(parts)
+ halves = parts.flat_buffer[[a, b], COL["area"]]
+ assert halves.sum() == pytest.approx(area, rel=0.02)
+ assert min(halves) > 0
+
+ def test_split_keeps_the_parent_track_on_the_larger_half(self, built):
+ parts = _fresh(built)
+ gi = int(parts.indices_at(0)[0])
+ parent = int(parts.flat_buffer[gi, COL["track_id"]])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ # An off-centre cut, so the halves are clearly unequal.
+ a, b = po.split_particle(parts, gi, ((cx + 2, cy - 50), (cx + 2, cy + 50)))
+ tracks = parts.flat_buffer[[a, b], COL["track_id"]].astype(int)
+ areas = parts.flat_buffer[[a, b], COL["area"]]
+ keeper = int(np.argmax(areas))
+ assert tracks[keeper] == parent
+ assert tracks[1 - keeper] == -1, \
+ "which fragment continues the track is a re-link's answer, not a guess"
+
+ def test_split_pairs_each_new_row_with_its_OWN_outline(self, built):
+ """The add path through ``_splice``: a mis-gathered contour block would
+ pair a plausible outline with the wrong row, which draws as nonsense
+ rather than raising."""
+ parts = _fresh(built)
+ gi = int(parts.indices_at(0)[0])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ a, b = po.split_particle(parts, gi, ((cx, cy - 50), (cx, cy + 50)))
+ for index in (a, b):
+ row = parts.flat_buffer[index]
+ contour = parts.contour_at(index)
+ assert len(contour) >= 3
+ # The outline must sit inside its own row's bounding box (+1 px, the
+ # tracer runs on a padded crop).
+ assert contour[:, 0].min() >= row[COL["bbox_y0"]] - 1
+ assert contour[:, 0].max() <= row[COL["bbox_y1"]] + 1
+ assert contour[:, 1].min() >= row[COL["bbox_x0"]] - 1
+ assert contour[:, 1].max() <= row[COL["bbox_x1"]] + 1
+ # And the neighbours' outlines are untouched by the insertion.
+ _csr_is_consistent(parts)
+
+ def test_split_that_misses_raises(self, built):
+ parts = _fresh(built)
+ gi = int(parts.indices_at(0)[0])
+ with pytest.raises(ValueError, match="does not divide"):
+ po.split_particle(parts, gi, ((0.0, 0.0), (0.0, 10.0)))
+
+ def test_merge_round_trips_a_split(self, built):
+ """Split then merge the halves back: one row, the original area."""
+ parts = _fresh(built)
+ gi = int(parts.indices_at(0)[0])
+ area = float(parts.flat_buffer[gi, COL["area"]])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ a, b = po.split_particle(parts, gi, ((cx, cy - 50), (cx, cy + 50)))
+ before = parts.n_particles
+
+ merged = po.merge_particles(parts, [a, b])
+ assert parts.n_particles == before - 1
+ _csr_is_consistent(parts)
+ assert float(parts.flat_buffer[merged, COL["area"]]) == pytest.approx(area, rel=0.02)
+
+ def test_merge_keeps_the_largest_bodys_track(self, built):
+ parts = _fresh(built)
+ gi = int(parts.indices_at(0)[0])
+ parent = int(parts.flat_buffer[gi, COL["track_id"]])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ a, b = po.split_particle(parts, gi, ((cx + 2, cy - 50), (cx + 2, cy + 50)))
+ merged = po.merge_particles(parts, [a, b])
+ assert int(parts.flat_buffer[merged, COL["track_id"]]) == parent
+
+ def test_merging_particles_that_do_not_touch_raises(self, built):
+ """A boolean union cast to int32 would measure two distant discs as ONE
+ region with a centroid in the empty space between them."""
+ parts = _fresh(built)
+ a, b = (int(i) for i in parts.indices_at(0)[:2])
+ with pytest.raises(ValueError, match="do not touch"):
+ po.merge_particles(parts, [a, b])
+
+ def test_merging_across_frames_raises(self, built):
+ parts = _fresh(built)
+ a = int(parts.indices_at(0)[0])
+ b = int(parts.indices_at(1)[0])
+ with pytest.raises(ValueError, match="one frame"):
+ po.merge_particles(parts, [a, b])
+
+ def test_merge_needs_two(self, built):
+ parts = _fresh(built)
+ with pytest.raises(ValueError, match="at least two"):
+ po.merge_particles(parts, [0])
+
+ def test_edits_are_stamped_into_provenance_for_reproducibility(self, built):
+ parts = _fresh(built)
+
+ class _Tree:
+ _commit_provenance = {"action": "segment_particles"}
+ tree = _Tree()
+ overlay = _detached(parts)
+ overlay.tree = tree
+ gi = int(parts.indices_at(0)[0])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ overlay.select([gi])
+ overlay.split(gi, ((cx, cy - 50), (cx, cy + 50)))
+ overlay.merge(overlay.selected)
+
+ kinds = [e["kind"] for e in tree.particle_edits]
+ assert kinds == ["split", "merge"]
+ # Both surfaces carry the log: the tree (so a re-run sees it) and the
+ # store's provenance (so a saved file still reproduces the result).
+ assert [e["kind"] for e in parts.provenance["edits"]] == kinds
+ assert [e["kind"] for e in tree._commit_provenance["edits"]] == kinds
+
+ def test_pending_edits_is_the_re_run_seam(self, built):
+ """A re-segmentation rebuilds the store from the raw frames; unless it
+ reads this list first, every correction is silently discarded."""
+ parts = _fresh(built)
+
+ class _Tree:
+ pass
+ tree = _Tree()
+ assert po.pending_edits(tree) == []
+ overlay = _detached(parts)
+ overlay.tree = tree
+ overlay.select([int(parts.indices_at(0)[0])])
+ overlay.delete()
+ edits = po.pending_edits(tree)
+ assert len(edits) == 1 and edits[0]["revision"] == 1
+ # JSON-safe: it crosses the IPC and gets saved with the store.
+ import json
+ json.dumps(edits)
+
+ def test_edit_bumps_the_revision(self, built):
+ parts = _fresh(built)
+ overlay = _detached(parts)
+ assert overlay.revision == 0
+ overlay.select([int(parts.indices_at(0)[0])])
+ overlay.delete()
+ assert overlay.revision == 1
+
+ def test_frame_provider_fills_the_intensity_columns(self, built):
+ source, _gt, _parts, _res = built
+ parts = _fresh(built)
+ overlay = _detached(parts, frame_provider=lambda t: np.asarray(source.data[t]))
+ gi = int(parts.indices_at(0)[0])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ overlay.select([gi])
+ a, b = overlay.split(gi, ((cx, cy - 50), (cx, cy + 50)))
+ assert np.isfinite(parts.flat_buffer[a, COL["intensity_mean"]])
+
+ def test_without_a_frame_provider_intensity_stays_NaN(self, built):
+ """NaN rather than an invented number — the same rule ``measure_frame``
+ applies when it runs with no intensity image."""
+ parts = _fresh(built)
+ overlay = _detached(parts)
+ gi = int(parts.indices_at(0)[0])
+ cx, cy = po.centroids_px(parts.flat_buffer[gi:gi + 1], parts.scale)[0]
+ a, _b = overlay.split(gi, ((cx, cy - 50), (cx, cy + 50)))
+ assert not np.isfinite(parts.flat_buffer[a, COL["intensity_mean"]])
+
+
+# ── navigator lanes ──────────────────────────────────────────────────────────
+
+class TestNavigatorLanes:
+ def test_count_lane_is_emitted_as_STEP_data(self):
+ """Plan C3: a straight interpolation between frames puts a nucleation's
+ transition half a frame early, so 8 reads as 7."""
+ counts = np.array([0, 0, 3, 3, 5], np.float32)
+ x, y = po.step_trace(counts)
+ # Every sample is HELD to the next x before it jumps: consecutive pairs
+ # share a y, and the jump happens exactly at the frame boundary.
+ assert y.tolist() == [0, 0, 0, 0, 3, 3, 3, 3, 5, 5]
+ assert x.tolist() == [0, 1, 1, 2, 2, 3, 3, 4, 4, 5]
+ rise = int(np.argmax(np.diff(y) > 0)) + 1
+ assert x[rise] == 2.0, "the count must rise AT frame 2, not between 1 and 2"
+
+ def test_step_trace_honours_a_calibrated_time_axis(self):
+ counts = np.array([1, 2, 3], np.float32)
+ x, _y = po.step_trace(counts, np.arange(3) * 0.05)
+ assert x[0] == pytest.approx(0.0)
+ assert x[-1] == pytest.approx(0.15)
+
+ def test_step_trace_is_empty_for_an_empty_lane(self):
+ x, y = po.step_trace([])
+ assert x.size == 0 and y.size == 0
+
+ def test_a_continuous_lane_is_NOT_stepped(self, built):
+ """Mean size is continuous; only the integer lane is a staircase. The
+ lane builder plots it straight, so its length is n, not 2n."""
+ _s, _gt, parts, _res = built
+ size = parts.property_series("area", "mean")
+ assert size.shape == (N_FRAMES,)
+
+ def test_event_points_carry_one_row_per_kind(self, built):
+ _s, _gt, _parts, res = built
+ pts = po._event_points(res.events, "birth", 1.0, 0.0)
+ assert len(pts) == len(res.events_of("birth"))
+ assert set(pts[:, 1]) == {float(po.EVENT_ROWS["birth"])}
+ assert po._event_points(res.events, "death", 1.0, 0.0).shape == (0, 2)
+
+ def test_event_rows_and_colours_cover_every_kind(self):
+ from spyde.particles.track import EVENT_KINDS
+ assert set(po.EVENT_COLORS) == set(EVENT_KINDS)
+ assert set(po.EVENT_ROWS) == set(EVENT_KINDS)
+ assert len(set(po.EVENT_ROWS.values())) == len(EVENT_KINDS), \
+ "two kinds sharing a row would overplot each other"
+ assert len(set(po.EVENT_COLORS.values())) == len(EVENT_KINDS)
+ assert po.EVENT_COLORS["birth"] == "#a6e3a1" # green
+ assert po.EVENT_COLORS["death"] == "#f38ba8" # red
+ assert po.EVENT_COLORS["merge"] == "#cba6f7" # mauve
+ assert po.EVENT_COLORS["split"] == "#f9e2af" # yellow
+
+ def test_event_points_use_the_calibrated_time_axis(self, built):
+ _s, _gt, _parts, res = built
+ pts = po._event_points(res.events, "birth", 0.05, 0.0)
+ frames = [e.frame for e in res.events_of("birth")]
+ assert np.allclose(pts[:, 0], np.asarray(frames) * 0.05)
+
+ def test_publish_emits_a_stacked_three_row_navigator(self, window, make_tree):
+ session = window["window"]
+ messages = window["messages"]
+ tree = make_tree()
+
+ messages.clear()
+ assert po.publish_navigator_lanes(session, tree) is True
+ figures = [m for m in messages if m.get("type") == "figure"
+ and m.get("view_kind") == "stacked"]
+ assert figures, "no stacked lane figure emitted"
+ figure = figures[-1]
+ assert figure["is_navigator"] is True
+ assert figure["window_id"] == po._first_nav_plot(tree).window_id
+ for lane in (po.LANE_COUNT, po.LANE_SIZE, po.LANE_EVENTS):
+ assert lane in figure["title"]
+
+ def test_lanes_reuse_the_shared_stacked_cursor(self, window, make_tree):
+ """The reusable half of ``navigator_views``: one logical time cursor
+ wired to the tree's REAL 1-D navigation selector, with a line per row."""
+ session = window["window"]
+ tree = make_tree()
+ window_id = po._first_nav_plot(tree).window_id
+ assert _wait(lambda: session._nav_selectors.get(window_id) is not None)
+
+ po.publish_navigator_lanes(session, tree)
+ cursor = session._stacked_nav_cursors.get(window_id)
+ assert cursor is not None
+ assert len(cursor.widgets) == 3, "one draggable line per lane"
+ assert cursor._index_hook in session._nav_selectors[window_id].index_hooks
+
+ def test_lanes_are_registered_as_named_navigators(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ po.publish_navigator_lanes(session, tree)
+ assert {po.LANE_COUNT, po.LANE_SIZE} <= set(tree.navigator_signals)
+
+ def test_republishing_replaces_the_prior_cursor(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ window_id = po._first_nav_plot(tree).window_id
+ assert _wait(lambda: session._nav_selectors.get(window_id) is not None)
+
+ po.publish_navigator_lanes(session, tree)
+ first = session._stacked_nav_cursors[window_id]
+ po.publish_navigator_lanes(session, tree)
+ second = session._stacked_nav_cursors[window_id]
+ assert second is not first and first._closed is True
+ assert first._index_hook not in session._nav_selectors[window_id].index_hooks
+
+ def test_a_tree_without_traces_publishes_nothing(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree(events=False)
+ tree.nav_traces = {}
+ assert po.publish_navigator_lanes(session, tree) is False
+
+
+# ── lifecycle on a real tree ─────────────────────────────────────────────────
+
+class TestLifecycle:
+ def test_attach_puts_the_overlay_on_the_tree(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ assert overlay is not None
+ assert tree._particle_overlay is overlay
+ assert overlay._groups, "no marker groups were created"
+
+ def test_marker_group_count_is_one_per_colour(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ colours = len(po.TRACK_COLORS) + 1
+ expected = colours * 2 + colours * po.TRAIL_FADE_STEPS + 2
+ assert len(overlay._groups) == expected
+
+ def test_draw_order_is_trails_then_fills_then_heads_then_labels(self, window, make_tree):
+ """Creation order IS draw order, and anyplotlib flattens the registry by
+ marker TYPE. A head dot underneath its own particle's fill is not a head
+ dot, so the type order has to come out lines → polygons → circles → texts.
+ """
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ registry = po._first_signal_plot(tree)._plot2d.markers
+ assert list(registry) == ["lines", "polygons", "circles", "texts"]
+ # Within the polygons type the selected outline is added after the fills.
+ polygon_names = list(registry["polygons"].keys())
+ assert polygon_names[-1].endswith("_selected")
+
+ def test_attached_overlay_pushes_polygons_for_the_current_frame(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ overlay.set_frame(2)
+ drawn = sum(len(overlay._groups[f"fill{i}"]._data.get("vertices_list", []))
+ for i in range(len(po.TRACK_COLORS) + 1))
+ assert drawn == len(tree.particles.at(2))
+
+ def test_reattaching_does_not_stack_two_overlays(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ plot = po._first_signal_plot(tree)
+ first = po.attach_particle_overlay(plot, tree.particles, tree)
+ second = po.attach_particle_overlay(plot, tree.particles, tree)
+ assert second is not first
+ assert tree._particle_overlay is second
+ assert first._groups == {}, "the prior overlay's markers were left behind"
+
+ def test_navigator_hook_is_attached_and_detached(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ assert overlay._selectors, "the overlay found no navigator selector"
+ selectors = list(overlay._selectors)
+ assert all(overlay._on_indices in sel.index_hooks for sel in selectors)
+ overlay.remove()
+ assert all(overlay._on_indices not in sel.index_hooks for sel in selectors)
+
+ def test_a_nav_move_redraws_the_new_frame(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ overlay._on_indices(np.array([5]))
+ assert overlay._frame == 5
+ drawn = sum(len(overlay._groups[f"fill{i}"]._data.get("vertices_list", []))
+ for i in range(len(po.TRACK_COLORS) + 1))
+ assert drawn == len(tree.particles.at(5))
+
+ def test_a_superseded_payload_never_lands(self, window, make_tree):
+ """Latest-wins: a payload built for an older generation is dropped."""
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ overlay.set_frame(3)
+ stale = overlay._payload(0)
+ current = overlay._gen
+ overlay._apply(stale, current - 1)
+ drawn = sum(len(overlay._groups[f"fill{i}"]._data.get("vertices_list", []))
+ for i in range(len(po.TRACK_COLORS) + 1))
+ assert drawn == len(tree.particles.at(3))
+
+ def test_teardown_bumps_the_generation_first(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ gen = overlay._gen
+ overlay.remove()
+ assert overlay._gen > gen
+
+ def test_tree_close_reaps_the_overlay(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ selectors = list(overlay._selectors)
+ tree.close()
+ assert getattr(tree, "_particle_overlay", None) is None
+ assert overlay._groups == {}
+ assert all(overlay._on_indices not in sel.index_hooks for sel in selectors)
+
+ def test_region_widget_appears_and_selects_in_bulk(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ plot = po._first_signal_plot(tree)
+ overlay = po.attach_particle_overlay(plot, tree.particles, tree)
+ assert overlay._region_widget is None
+
+ overlay.set_region_select(True)
+ assert overlay._region_widget is not None
+ # The default box covers the middle half of the frame, so it catches some
+ # particles but not all of them.
+ assert 0 < len(overlay.selected) <= len(tree.particles.at(overlay._frame))
+
+ overlay.set_region_select(False)
+ assert overlay._region_widget is None
+ assert plot._plot2d.list_widgets() == []
+
+ def test_hidden_overlay_still_follows_the_navigator(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ overlay = po.attach_particle_overlay(
+ po._first_signal_plot(tree), tree.particles, tree)
+ overlay.set_visible(False)
+ overlay._on_indices(np.array([4]))
+ assert overlay._frame == 4
+ drawn = sum(len(overlay._groups[f"fill{i}"]._data.get("vertices_list", []))
+ for i in range(len(po.TRACK_COLORS) + 1))
+ assert drawn == 0, "a hidden overlay drew markers"
+ overlay.set_visible(True)
+ drawn = sum(len(overlay._groups[f"fill{i}"]._data.get("vertices_list", []))
+ for i in range(len(po.TRACK_COLORS) + 1))
+ assert drawn == len(tree.particles.at(4))
+
+
+# ── the staged actions ───────────────────────────────────────────────────────
+
+class TestStagedActions:
+ def test_every_part_handler_resolves(self):
+ from spyde.actions import registry
+ keys = [k for k in registry.STAGED_HANDLERS if k.startswith("part_")]
+ assert {"part_open", "part_close", "part_select", "part_delete",
+ "part_merge", "part_split", "part_lanes"} <= set(keys)
+ for key in keys:
+ assert callable(registry.resolve_staged(key))
+
+ def test_the_caret_schema_resolves(self):
+ from spyde.actions import registry
+ schema = registry.wizard_parameters("part")
+ assert schema and schema["trail_frames"]["default"] == po.DEFAULT_TRAIL_FRAMES
+
+ def test_open_attaches_and_close_tears_down(self, window, make_tree):
+ session = window["window"]
+ tree = make_tree()
+ plot = po._first_signal_plot(tree)
+
+ po.part_open(session, plot, {})
+ assert getattr(tree, "_particle_overlay", None) is not None
+ po.part_close(session, plot, {})
+ assert getattr(tree, "_particle_overlay", None) is None
+
+ def test_select_emits_the_selection_for_the_renderer(self, window, make_tree, built):
+ session = window["window"]
+ messages = window["messages"]
+ _s, _gt, parts, _res = built
+ tree = make_tree()
+ plot = po._first_signal_plot(tree)
+ po.part_open(session, plot, {})
+
+ gi = int(parts.indices_at(0)[0])
+ messages.clear()
+ po.part_select(session, plot, {"indices": [gi]})
+ sent = [m for m in messages if m.get("type") == "particle_selection"]
+ assert sent, "no particle_selection message emitted"
+ payload = sent[-1]
+ assert payload["indices"] == [gi]
+ record = payload["particles"][0]
+ assert record["index"] == gi
+ assert record["color"] == po.track_color(record["track_id"])
+ assert "area" in record and "circularity" in record
+ po.part_close(session, plot, {})
+
+ def test_selection_message_is_valid_JSON_even_with_NaN_properties(self, window, make_tree):
+ """``ipc.emit`` dumps with ``allow_nan=True``, so a bare ``NaN`` token
+ would make ``JSON.parse`` throw and the renderer lose the whole message —
+ and NaN is the normal value for an unmeasured intensity column."""
+ import json
+ session = window["window"]
+ messages = window["messages"]
+ tree = make_tree()
+ plot = po._first_signal_plot(tree)
+ po.part_open(session, plot, {})
+ gi = int(tree.particles.indices_at(0)[0])
+ buffer = tree.particles.flat_buffer
+ saved = float(buffer[gi, COL["intensity_mean"]])
+ buffer[gi, COL["intensity_mean"]] = np.nan
+ try:
+ messages.clear()
+ po.part_select(session, plot, {"indices": [gi]})
+ sent = [m for m in messages if m.get("type") == "particle_selection"][-1]
+ assert sent["particles"][0]["intensity_mean"] is None
+ assert "NaN" not in json.dumps(sent)
+ finally:
+ # The store is module-scoped and shared by every test in this file.
+ buffer[gi, COL["intensity_mean"]] = saved
+ po.part_close(session, plot, {})
+
+ def test_open_without_particles_errors_rather_than_hanging(self, window):
+ """No segmentation running and no particles: say so, don't wait forever.
+
+ (``wait_for_particles`` returns False with no event loop; the harness
+ Session HAS one, so the wait starts and the error arrives from the poll's
+ grace window — the assertion is only that open() does not attach.)
+ """
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ assert _wait(lambda: _signal_plot(session) is not None)
+ plot = _signal_plot(session)
+ po.part_open(session, plot, {})
+ assert getattr(plot.signal_tree, "_particle_overlay", None) is None
diff --git a/spyde/tests/migrated/test_particle_tree.py b/spyde/tests/migrated/test_particle_tree.py
new file mode 100644
index 00000000..a99b6214
--- /dev/null
+++ b/spyde/tests/migrated/test_particle_tree.py
@@ -0,0 +1,312 @@
+"""
+The particle tree (plan §0.6) and the Wave-0 framework around it.
+
+Three separable claims, one per class:
+
+* segmentation spawns a NEW tree carrying its provenance, rather than decorating
+ the source — which is what makes Wave D's per-particle diffraction well-defined;
+* the label movie is LAZY and stays lazy, because a materialised one is 64 MB per
+ frame at the plan's target size;
+* ``requires_particles`` gates identically in both toolbar filter paths, because
+ a gate added to only one renders a button that never dispatches (or vice versa).
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import spyde.data.synthetic as sy
+from spyde.actions.particle_tree import (
+ PARTICLE_SIGNAL_TYPE,
+ open_particle_tree,
+ particle_nav_positions,
+)
+from spyde.particles import (
+ LinkParams,
+ SegmentParams,
+ link,
+ measure_frame,
+ segment_frame,
+)
+from spyde.signals.particles import COL, SpyDEParticles
+
+N_FRAMES = 8
+
+
+@pytest.fixture(scope="module")
+def built():
+ """A real segmentation of the fixture — not a hand-made container."""
+ s = sy.particle_movie(n_frames=N_FRAMES)
+ gt = sy.ground_truth(s)
+ scale = float(gt["scale"])
+ per_frame, contours = [], []
+ for t in range(N_FRAMES):
+ lab = segment_frame(s.data[t], SegmentParams(min_size=25, gaussian=1.0))
+ rows, cs = measure_frame(lab, s.data[t], t=t, scale=scale)
+ per_frame.append(rows)
+ contours.append(cs)
+ parts = SpyDEParticles.from_frames(
+ per_frame, frame_shape=tuple(gt["frame_shape"]),
+ contours_per_frame=contours, scale=scale, units="nm")
+ res = link(parts, LinkParams(max_dist=10.0))
+ res.apply(parts)
+ return s, gt, parts, res
+
+
+class TestTreeCreation:
+ def test_spawns_a_new_tree_not_an_attribute(self, window, built):
+ """The §0.6 decision, asserted directly."""
+ session = window["window"]
+ s, _gt, parts, res = built
+ before = len(session.signal_trees)
+ tree = open_particle_tree(session, particles=parts, source_node=s,
+ events=res.events)
+ assert len(session.signal_trees) == before + 1
+ assert tree.particles is parts
+ assert tree.source_node is s
+ assert getattr(s, "particles", None) is None, (
+ "the SOURCE signal was decorated — that is the design this replaces")
+
+ def test_root_carries_the_particle_signal_type(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert getattr(tree.root, "_signal_type", None) == PARTICLE_SIGNAL_TYPE
+
+ def test_label_movie_matches_the_source_shape(self, window, built):
+ session = window["window"]
+ s, gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert tree.root.data.shape == (N_FRAMES, *tuple(gt["frame_shape"]))
+
+ def test_provenance_is_stamped(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s,
+ params={"sensitivity": 0.5})
+ prov = getattr(tree, "_commit_provenance", None) or {}
+ assert prov.get("action") == "segment_particles"
+ assert prov.get("params", {}).get("sensitivity") == 0.5
+
+ def test_nav_map_defaults_to_identity(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert np.array_equal(tree.nav_map, np.arange(N_FRAMES))
+
+ def test_calibration_follows_the_source(self, window, built):
+ """A centroid must mean the same thing on both trees."""
+ session = window["window"]
+ s, gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert tree.root.axes_manager.signal_axes[0].scale == \
+ pytest.approx(float(gt["scale"]))
+ assert tree.root.axes_manager.signal_axes[0].units == "nm"
+ assert tree.root.axes_manager.navigation_axes[0].name == "time"
+
+
+class TestLabelMovieStaysLazy:
+ """A materialised label movie is 64 MB per frame at the plan's target size."""
+
+ def test_root_is_lazy(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert tree.root._lazy
+
+ def test_one_frame_per_chunk(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert tree.root.data.chunksize[0] == 1
+
+ def test_computing_one_frame_does_not_compute_the_stack(self, window, built):
+ """The Memory-Safety rule, enforced the way find_vectors enforces it."""
+ import dask.array as da
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ full = tree.root.data.shape
+ seen = {"full": 0}
+ real = da.Array.compute
+
+ def guard(self, *a, **k):
+ if self.shape == full:
+ seen["full"] += 1
+ return real(self, *a, **k)
+
+ try:
+ da.Array.compute = guard
+ frame = np.asarray(tree.root.data[3].compute())
+ finally:
+ da.Array.compute = real
+ assert frame.shape == full[1:]
+ assert seen["full"] == 0, "rendering one frame computed the whole movie"
+
+ def test_rendered_frame_carries_track_ids(self, window, built):
+ session = window["window"]
+ s, _gt, parts, res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s,
+ events=res.events)
+ frame = np.asarray(tree.root.data[3].compute())
+ painted = np.unique(frame)
+ painted = painted[painted > 0]
+ assert painted.size == len(parts.at(3)), (
+ "painted a different number of particles than frame 3 holds")
+
+
+class TestNavigatorTraces:
+ def test_count_and_size_lanes_exist(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert set(tree.nav_traces) >= {"count", "size"}
+ assert tree.nav_traces["count"].shape == (N_FRAMES,)
+
+ def test_event_lanes_appear_only_with_events(self, window, built):
+ session = window["window"]
+ s, _gt, parts, res = built
+ without = open_particle_tree(session, particles=parts, source_node=s)
+ assert not any(k.startswith("event_") for k in without.nav_traces)
+ with_ev = open_particle_tree(session, particles=parts, source_node=s,
+ events=res.events)
+ assert any(k.startswith("event_") for k in with_ev.nav_traces)
+
+ def test_count_lane_matches_the_store(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ assert np.array_equal(tree.nav_traces["count"], parts.count_series())
+
+
+class TestWaveDSeam:
+ """`particle_nav_positions` is what makes per-particle diffraction definable."""
+
+ def test_movie_particle_maps_to_its_frame(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ tree = open_particle_tree(session, particles=parts, source_node=s)
+ gi = int(parts.indices_at(3)[0])
+ nav = particle_nav_positions(tree, gi)
+ assert nav.shape == (1, 1) and int(nav[0, 0]) == 3, (
+ "on a MOVIE a particle's pixels are signal coordinates, so the only "
+ "navigation index involved is the frame")
+
+ def test_nav_map_is_honoured(self, window, built):
+ session = window["window"]
+ s, _gt, parts, _res = built
+ shifted = np.arange(N_FRAMES) + 100
+ tree = open_particle_tree(session, particles=parts, source_node=s,
+ nav_map=shifted)
+ gi = int(parts.indices_at(3)[0])
+ assert int(particle_nav_positions(tree, gi)[0, 0]) == 103
+
+ def test_survives_a_store_without_masks(self, window, built):
+ """store_masks=False is the default for long movies — this must not raise."""
+ session = window["window"]
+ s, _gt, parts, _res = built
+ bare = SpyDEParticles(parts.flat_buffer.copy(), parts.t_offsets.copy(),
+ parts.frame_shape, scale=parts.scale,
+ units=parts.units)
+ assert not bare.has_masks
+ tree = open_particle_tree(session, particles=bare, source_node=s)
+ nav = particle_nav_positions(tree, 0)
+ assert nav.shape == (1, 1)
+
+
+class TestRequiresParticlesGate:
+ """Both filter paths, because one alone is a button that never dispatches."""
+
+ def _fake(self, has_particles: bool):
+ class _Sig:
+ _signal_type = "particles"
+
+ class _Tree:
+ particles = object() if has_particles else None
+ diffraction_vectors = None
+ root = _Sig()
+
+ class _Plot:
+ signal_tree = _Tree()
+
+ class _State:
+ plot = _Plot()
+ current_signal = _Sig()
+ dimensions = 2
+ navigation = False
+ return _State()
+
+ def test_second_path_hides_and_shows(self):
+ from spyde.drawing.toolbars.plot_control_toolbar import _action_matches_plot
+ meta = {"requires_particles": True, "plot_dim": [1, 2]}
+ assert not _action_matches_plot("X", meta, self._fake(False))
+ assert _action_matches_plot("X", meta, self._fake(True))
+
+ def test_ungated_actions_are_unaffected(self):
+ from spyde.drawing.toolbars.plot_control_toolbar import _action_matches_plot
+ meta = {"plot_dim": [1, 2]}
+ assert _action_matches_plot("X", meta, self._fake(False))
+
+ def test_both_paths_read_the_same_key(self):
+ """Guards the §6 pitfall directly: the key must appear in BOTH filters."""
+ import inspect
+ from spyde.drawing.toolbars import plot_control_toolbar as mod
+ src = inspect.getsource(mod)
+ assert src.count("requires_particles") >= 4, (
+ "requires_particles must be read in get_toolbar_actions_for_plot AND "
+ "_action_matches_plot — one alone renders a button that never "
+ "dispatches, or hides one that would have worked")
+
+
+class TestWaitForParticles:
+ def test_returns_false_without_an_event_loop(self):
+ from spyde.actions.lifecycle import wait_for_particles
+
+ class _S:
+ _dispatch_to_main = None
+ called = []
+ started = wait_for_particles(_S(), None, lambda: called.append(1),
+ what="Test")
+ assert started is False and not called
+
+ def test_fires_once_the_particles_land(self, built):
+ import threading
+ import time
+ from spyde.actions.lifecycle import wait_for_particles
+
+ _s, _gt, parts, _res = built
+
+ class _Tree:
+ particles = None
+
+ class _Plot:
+ signal_tree = _Tree()
+
+ done = threading.Event()
+
+ class _S:
+ signal_trees = [_Tree()]
+
+ @staticmethod
+ def _dispatch_to_main(fn):
+ fn()
+
+ plot = _Plot()
+ assert wait_for_particles(_S(), plot, done.set, what="Test", grace=30.0)
+ time.sleep(0.25)
+ assert not done.is_set(), "fired before the particles attached"
+ plot.signal_tree.particles = parts
+ assert done.wait(5.0), "never fired after the particles attached"
+
+ def test_seg_batch_running_reads_the_tree_flag(self):
+ from spyde.actions.lifecycle import seg_batch_running
+
+ class _T:
+ _seg_batch_running = False
+
+ class _S:
+ signal_trees = [_T()]
+ s = _S()
+ assert not seg_batch_running(s)
+ s.signal_trees[0]._seg_batch_running = True
+ assert seg_batch_running(s)
diff --git a/spyde/tests/migrated/test_particles_batch.py b/spyde/tests/migrated/test_particles_batch.py
new file mode 100644
index 00000000..4f25021c
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_batch.py
@@ -0,0 +1,422 @@
+"""
+Tests for ``spyde.particles.batch`` — the whole-movie segmentation fan-out.
+
+The contract this file exists to pin, in the order it matters:
+
+**1. A frame's particles must not change because it was computed in parallel.**
+Everything else here is throughput, and throughput that alters the answer is
+worthless. Asserted with ``array_equal`` against the serial loop, never
+``allclose`` — the per-frame work is deterministic and device-free on the
+classical engine, so exact equality is the honest bar. (The scribble engine adds
+a device, and a CUDA frame and a CPU frame differ in the last bits of a float32
+convolution; that is why the lane policy is what it is and why the GPU/CPU
+agreement gate lives in ``test_particles_gpu.py``, in a subprocess.)
+
+**2. A cancelled run still spans the movie.** The CSR store is built from these
+lists, so a short list is a shorter movie, not a stopped one.
+
+**3. The frame index survives the dispatch.** A task cannot know its own global
+offset — ``dispatch_chunks`` slices the result array per chunk, so a
+``map_blocks`` stage reading ``block_info`` reports (0, 0) for every one of them
+(the bug ``orchestrate`` records for the live count map). The index is stamped
+where the global slice is authoritative, and a multi-frame block is where a
+mistake would show.
+
+**4. The lane policy is ONE value.** ``gpu_runtime`` records a real bug of this
+class: the docstring said "2" while the code passed "4", so the client sized a
+lane for one policy while the workers gated on another. The per-worker gate and
+the client-side split must be given the same default.
+
+**5. No rechunk shuffle, ever.** A movie whose reader split the signal axes must
+fall back rather than shuffle multiple GB through the scheduler (CLAUDE.md
+Live-Display §1).
+
+Everything runs on the LOCAL thread-pool fallback (no cluster, no CUDA):
+``SPYDE_NO_DASK=1`` is the migrated-test mode and torch-CUDA segfaults under
+pytest on Windows.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.particles import SegmentParams, measure_frame, segment_frame
+from spyde.particles.batch import (
+ EngineSpec,
+ PARTICLE_GPU_LANE_DEFAULT,
+ frames_per_task,
+ resolve_engine,
+ segment_block,
+ segment_movie,
+)
+from spyde.signals.particles import COL, N_COLUMNS
+
+PARAMS = dict(min_size=10, sensitivity=0.6)
+
+
+def _disc(shape, cy, cx, r):
+ yy, xx = np.mgrid[0:shape[0], 0:shape[1]]
+ return (((yy - cy) ** 2 + (xx - cx) ** 2) <= r * r).astype(np.float32)
+
+
+def _movie(n=7, h=90, w=110):
+ """A short movie whose particle count CHANGES frame to frame.
+
+ Deliberately not a repeated frame: identical frames would let a dispatcher
+ that mixed up block offsets still pass every equality check here.
+ """
+ out = np.zeros((n, h, w), np.float32)
+ for t in range(n):
+ img = np.full((h, w), 0.05, np.float32)
+ img += _disc((h, w), 25, 30, 8)
+ img += _disc((h, w), 60, 70, 6 + (t % 3))
+ if t % 2 == 0:
+ img += _disc((h, w), 25, 85, 5)
+ out[t] = np.clip(img, 0, 1.5)
+ return out
+
+
+def _serial(data, scale=1.0):
+ p = SegmentParams(**PARAMS)
+ rows, contours = [], []
+ for t in range(data.shape[0]):
+ labels = segment_frame(data[t], p)
+ r, c = measure_frame(labels, data[t], t=t, scale=scale)
+ rows.append(r)
+ contours.append(c)
+ return rows, contours
+
+
+@pytest.fixture(scope="module")
+def movie():
+ return _movie()
+
+
+@pytest.fixture(scope="module")
+def reference(movie):
+ return _serial(movie)
+
+
+@pytest.fixture
+def spec():
+ return EngineSpec(method="classical", params=dict(PARAMS))
+
+
+class TestParallelIsBitIdentical:
+ """Contract 1 — the answer may not depend on how it was computed."""
+
+ def _assert_same(self, got, reference, label):
+ rows, contours, done = got
+ ref_rows, ref_contours = reference
+ assert done == len(ref_rows), label
+ assert len(rows) == len(ref_rows) == len(contours), label
+ for t, (r, ref) in enumerate(zip(rows, ref_rows)):
+ assert np.array_equal(r, ref), f"{label}: frame {t} rows differ"
+ for t, (c, ref) in enumerate(zip(contours, ref_contours)):
+ assert len(c) == len(ref), f"{label}: frame {t} outline count"
+ for a, b in zip(c, ref):
+ assert np.array_equal(a, b), f"{label}: frame {t} outline"
+
+ def test_numpy_source(self, movie, reference, spec):
+ got = segment_movie(movie, spec, n_frames=len(movie), client=None)
+ self._assert_same(got, reference, "numpy")
+
+ def test_lazy_source(self, movie, reference, spec):
+ import dask.array as da
+ lazy = da.from_array(movie, chunks=(2, -1, -1))
+ got = segment_movie(lazy, spec, n_frames=len(movie), client=None)
+ self._assert_same(got, reference, "dask")
+
+ def test_get_frame_source(self, movie, reference, spec):
+ """No array at all — a callable frame source still has to agree."""
+ got = segment_movie(None, spec, n_frames=len(movie),
+ get_frame=lambda t: movie[t], client=None)
+ self._assert_same(got, reference, "get_frame")
+
+ def test_multi_frame_blocks_agree(self, movie, reference, spec,
+ monkeypatch):
+ """Several frames per task is the case where a block offset can be
+ wrong without any single-frame test noticing."""
+ monkeypatch.setenv("SPYDE_SEG_TASK_BYTES", str(1 << 30))
+ got = segment_movie(movie, spec, n_frames=len(movie), client=None)
+ self._assert_same(got, reference, "multi-frame blocks")
+
+ def test_scale_is_applied_once(self, movie, spec):
+ """A calibrated axis must not be applied twice by the parallel path."""
+ rows, _c, _d = segment_movie(movie, spec, n_frames=len(movie),
+ client=None, scale=2.5)
+ ref, _ = _serial(movie, scale=2.5)
+ for t, (r, e) in enumerate(zip(rows, ref)):
+ assert np.array_equal(r, e), f"frame {t}"
+
+
+class TestFrameIndex:
+ """Contract 3 — the ``t`` column, which no task can know for itself."""
+
+ def test_every_row_carries_its_own_frame(self, movie, spec):
+ rows, _c, _d = segment_movie(movie, spec, n_frames=len(movie),
+ client=None)
+ for t, r in enumerate(rows):
+ assert len(r), f"frame {t} found nothing — the fixture is wrong"
+ assert np.all(r[:, COL["t"]] == float(t)), (
+ f"frame {t} rows are stamped {set(r[:, COL['t']].tolist())}")
+
+ def test_multi_frame_block_indexes_within_itself(self, movie, spec,
+ monkeypatch):
+ monkeypatch.setenv("SPYDE_SEG_TASK_BYTES", str(1 << 30))
+ rows, _c, _d = segment_movie(movie, spec, n_frames=len(movie),
+ client=None)
+ for t, r in enumerate(rows):
+ assert np.all(r[:, COL["t"]] == float(t)), f"frame {t}"
+
+ def test_segment_block_stamps_from_its_offset(self, movie, spec):
+ """Used directly (the serial/threaded paths), the block DOES stamp."""
+ out = segment_block(movie[2:5], 2, spec)
+ assert len(out) == 3
+ for i, (rows, _c) in enumerate(out):
+ assert np.all(rows[:, COL["t"]] == float(2 + i))
+
+
+class TestProgress:
+ """The progressive fill: blocks land out of order, so the callback has to
+ carry the global offset rather than a running counter."""
+
+ def test_on_frames_covers_every_frame_exactly_once(self, movie, spec):
+ seen: list[int] = []
+
+ def _on(t0, t1, vals):
+ assert t1 - t0 == len(vals)
+ seen.extend(range(t0, t1))
+
+ segment_movie(movie, spec, n_frames=len(movie), client=None,
+ on_frames=_on)
+ assert sorted(seen) == list(range(len(movie)))
+
+ def test_a_failing_callback_does_not_fail_the_run(self, movie, spec):
+ def _boom(t0, t1, vals):
+ raise RuntimeError("the caret went away")
+
+ rows, _c, done = segment_movie(movie, spec, n_frames=len(movie),
+ client=None, on_frames=_boom)
+ assert done == len(movie)
+
+ def test_the_callback_sees_the_stamped_rows(self, movie, spec):
+ """The live label movie renders from these, so they must already carry
+ the true frame index — not the 0 the task stamped."""
+ stamps: list[tuple[int, float]] = []
+
+ def _on(t0, t1, vals):
+ for i, (rows, _c) in enumerate(vals):
+ if len(rows):
+ stamps.append((t0 + i, float(rows[0, COL["t"]])))
+
+ segment_movie(movie, spec, n_frames=len(movie), client=None,
+ on_frames=_on)
+ assert stamps and all(t == v for t, v in stamps), stamps
+
+
+class TestCancellation:
+ """Contract 2 — a stopped run is a movie with empty frames, not a short one."""
+
+ def test_stopped_before_it_starts_still_spans_the_movie(self, movie, spec):
+ rows, contours, done = segment_movie(
+ movie, spec, n_frames=len(movie), client=None, stopped=[True])
+ assert len(rows) == len(contours) == len(movie)
+ assert done == 0
+ assert all(r.shape == (0, N_COLUMNS) for r in rows)
+ assert all(c == [] for c in contours)
+
+ def test_unreached_frames_are_empty_blocks_not_missing_ones(self, movie,
+ spec):
+ stopped = [False]
+ seen = []
+
+ def _on(t0, t1, vals):
+ seen.append(t0)
+ stopped[0] = True # stop after the first block lands
+
+ rows, contours, done = segment_movie(
+ movie, spec, n_frames=len(movie), client=None, stopped=stopped,
+ on_frames=_on)
+ assert len(rows) == len(movie)
+ assert done < len(movie)
+ assert all(r.ndim == 2 and r.shape[1] == N_COLUMNS for r in rows)
+
+
+class TestTaskSizing:
+ """``frames_per_task`` is three bounds, and each is a different failure."""
+
+ def test_bounded_by_bytes(self):
+ # A 4096^2 uint8 frame is 16 MB; the 64 MB default holds four.
+ assert frames_per_task(1000, 16 << 20, n_workers=1) == 4
+
+ def test_never_exceeds_the_source_chunk(self):
+ """Alignment beats the byte budget: a task that spans two stored chunks
+ pulls both (CLAUDE.md Live-Display §1)."""
+ assert frames_per_task(1000, 1 << 20, n_workers=1,
+ source_chunk=3) == 3
+
+ def test_spreads_a_short_movie_across_the_cluster(self):
+ """Six frames and nine workers must not become one task."""
+ assert frames_per_task(6, 1 << 10, n_workers=9) == 1
+
+ def test_never_zero(self):
+ assert frames_per_task(1, 1 << 30, n_workers=64) == 1
+
+ def test_env_override(self, monkeypatch):
+ monkeypatch.setenv("SPYDE_SEG_TASK_BYTES", str(8 << 20))
+ assert frames_per_task(1000, 1 << 20, n_workers=1) == 8
+
+ def test_a_junk_env_value_falls_back_to_the_default(self, monkeypatch):
+ monkeypatch.setenv("SPYDE_SEG_TASK_BYTES", "lots")
+ assert frames_per_task(1000, 16 << 20, n_workers=1) == 4
+
+
+class TestNoRechunkShuffle:
+ """Contract 5. A reader that split the signal axes (RosettaSciIO's
+ balanced-cube auto-chunk) must NOT be fixed with a rechunk here — that is a
+ full P2P shuffle of the movie. It falls back, and the answer stays right."""
+
+ def test_split_signal_axes_falls_back_instead_of_rechunking(
+ self, movie, reference, spec, monkeypatch):
+ import dask.array as da
+ from spyde.particles import batch as batch_mod
+
+ split = da.from_array(movie, chunks=(3, 45, 55))
+ assert any(len(c) > 1 for c in split.chunks[1:]), "fixture not split"
+
+ def _no_rechunk(self, *a, **kw): # pragma: no cover - must not run
+ raise AssertionError("rechunked a movie instead of falling back")
+
+ monkeypatch.setattr(da.Array, "rechunk", _no_rechunk)
+ got = segment_movie(split, spec, n_frames=len(movie),
+ get_frame=lambda t: movie[t], client=None)
+ rows, _c, done = got
+ assert done == len(movie)
+ for t, (r, ref) in enumerate(zip(rows, reference[0])):
+ assert np.array_equal(r, ref), f"frame {t}"
+
+ def test_an_already_aligned_movie_is_not_rechunked(self, movie, spec,
+ monkeypatch):
+ """The app loads movies at one frame per chunk, which is already the
+ task size — rebuilding that graph for nothing is pure overhead."""
+ import dask.array as da
+ from spyde.particles import batch as batch_mod
+
+ calls = []
+ real = da.Array.rechunk
+
+ def _count(self, *a, **kw):
+ calls.append(a)
+ return real(self, *a, **kw)
+
+ monkeypatch.setattr(da.Array, "rechunk", _count)
+ lazy = da.from_array(movie, chunks=(1, -1, -1))
+ segment_movie(lazy, spec, n_frames=len(movie), client=None)
+ assert calls == []
+
+
+class TestEngineSpec:
+ def test_classical_needs_no_model(self, spec):
+ engine, dev = resolve_engine(spec)
+ assert dev == "cpu"
+ labels = engine(_movie(1)[0])
+ assert labels.dtype == np.int32
+
+ def test_classical_engine_matches_segment_frame(self, movie, spec):
+ engine, _dev = resolve_engine(spec)
+ assert np.array_equal(engine(movie[0]),
+ segment_frame(movie[0], SegmentParams(**PARAMS)))
+
+ def test_scribble_without_a_model_says_so(self):
+ with pytest.raises(ValueError, match="model_path"):
+ resolve_engine(EngineSpec(method="scribble", params=dict(PARAMS)))
+
+ def test_an_unknown_method_is_refused(self):
+ with pytest.raises(ValueError, match="no engine for method"):
+ resolve_engine(EngineSpec(method="prompt"))
+
+ def test_segment_params_round_trip(self, spec):
+ p = spec.segment_params()
+ assert p.min_size == PARAMS["min_size"]
+ assert p.sensitivity == PARAMS["sensitivity"]
+
+ def test_the_spec_pickles(self, spec):
+ """It crosses to a worker process; a spec that only cloudpickles would
+ work here and fail against a real scheduler."""
+ import pickle
+ assert pickle.loads(pickle.dumps(spec)) == spec
+
+
+class TestLanePolicy:
+ """Contract 4 — one value, given to both halves."""
+
+ def test_the_worker_gate_and_the_lane_split_get_the_same_default(
+ self, monkeypatch):
+ from spyde.particles import batch as batch_mod
+
+ seen = {}
+
+ def _fake_allowed(default_mode="one"):
+ seen["gate"] = default_mode
+ return False
+
+ def _fake_split(client, default_mode="one"):
+ seen["split"] = default_mode
+ return [], []
+
+ monkeypatch.setattr(
+ "spyde.actions.find_vectors.gpu_runtime._gpu_task_allowed",
+ _fake_allowed)
+ monkeypatch.setattr("spyde.compute_dispatch.split_workers_for_gpu",
+ _fake_split)
+ batch_mod._gpu_allowed()
+ assert seen["gate"] == PARTICLE_GPU_LANE_DEFAULT
+
+ def test_the_default_is_a_value_the_shared_policy_understands(self):
+ """``SPYDE_FV_GPU`` accepts one/N/all/off; anything else silently
+ becomes a single GPU worker, which would make the constant a lie."""
+ assert (PARTICLE_GPU_LANE_DEFAULT in ("one", "all", "off")
+ or PARTICLE_GPU_LANE_DEFAULT.isdigit())
+
+ def test_off_is_honoured(self, monkeypatch):
+ from spyde.particles import batch as batch_mod
+ monkeypatch.setenv("SPYDE_FV_GPU", "off")
+ assert batch_mod._gpu_allowed() is False
+
+ def test_the_cpu_lane_is_opt_in(self, monkeypatch):
+ from spyde.particles.batch import cpu_lane_enabled
+ monkeypatch.delenv("SPYDE_SEG_CPU_LANE", raising=False)
+ assert cpu_lane_enabled() is False
+ monkeypatch.setenv("SPYDE_SEG_CPU_LANE", "1")
+ assert cpu_lane_enabled() is True
+ monkeypatch.setenv("SPYDE_SEG_CPU_LANE", "0")
+ assert cpu_lane_enabled() is False
+
+
+class TestDegenerateBlocks:
+ def test_a_zero_length_block_is_not_an_error(self, spec):
+ """dask calls the chunk fn on an empty array for meta inference."""
+ out = segment_block(np.zeros((0, 8, 8), np.float32), 0, spec)
+ assert out.shape == (0,)
+
+ def test_a_non_3d_block_is_refused(self, spec):
+ with pytest.raises(ValueError, match=r"\(n, h, w\)"):
+ segment_block(np.zeros((8, 8), np.float32), 0, spec)
+
+ def test_store_masks_off_drops_the_outlines_not_the_rows(self, movie, spec):
+ rows, contours, _d = segment_movie(movie, spec, n_frames=len(movie),
+ client=None, store_masks=False)
+ assert any(len(r) for r in rows)
+ assert all(c == [] for c in contours)
+
+
+class TestDispatchChunksAssembleHook:
+ """The shared dispatcher grew an ``assemble`` hook for this module; the
+ default must stay exactly what find_vectors relies on."""
+
+ def test_the_signature_defaults_to_the_padded_convention(self):
+ import inspect
+ from spyde.compute_dispatch import dispatch_chunks
+ sig = inspect.signature(dispatch_chunks)
+ assert sig.parameters["assemble"].default is None
diff --git a/spyde/tests/migrated/test_particles_contours_parity.py b/spyde/tests/migrated/test_particles_contours_parity.py
new file mode 100644
index 00000000..4534350c
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_contours_parity.py
@@ -0,0 +1,413 @@
+"""
+Parity for the two per-region loops that were left in ``measure_frame``.
+
+:mod:`spyde.particles.props` and :mod:`spyde.particles.hull` took the property
+table from 43.7 s to 1.1 s on a real 4096² frame with 26 566 particles. What
+remained was **9.4 s of two Python ``for`` loops** — ``_fill_intensity`` and
+``_contours`` — which is 92% of the measurement and, like ``regionprops_table``
+before them, holds the GIL throughout, so a dask worker's four task slots stay
+worth one core (``benchmarks.md``). :mod:`spyde.particles.intensity` and
+:mod:`spyde.particles.contours` replace them. This module is the gate.
+
+The two halves need DIFFERENT gates, and getting that wrong in either direction
+is the trap
+---------------------------------------------------------------------------
+* **The intensity columns are exact and are asserted exactly.** The pixel sets
+ are identical by construction — same crop, same finite-only filter — so the
+ only freedom is floating-point summation order, and at the float32 resolution
+ the rows are stored in, that is nine orders below the last bit. Every intensity
+ column is asserted ``array_equal`` on the stored rows, and to ~1e-12 relative
+ on the float64 intermediates.
+
+* **The outlines are NOT asserted vertex by vertex, and it would be wrong to.**
+ A closed marching-squares contour is a CYCLE; ``skimage``'s dict-and-deque
+ assembly and a ``succ``-following walk cut the same cycle at different
+ vertices, so the arrays differ by a rotation while describing the same shape.
+ Demanding bit-identical vertices would reject a correct implementation.
+
+ It would be equally wrong to conclude from that that outlines are cosmetic and
+ a "close enough" polygon will do.
+ :meth:`~spyde.signals.particles.SpyDEParticles.render_frame` FILLS them to
+ rebuild the label movie, and :meth:`~…SpyDEParticles.mask_at` fills one to
+ produce the per-particle mask a mean diffraction pattern is sliced with — a
+ different contour is a different mask is a different measurement. So the gate
+ is the thing those two consume, and nothing weaker:
+
+ **``skimage.draw.polygon`` on the new outline must select EXACTLY the same
+ pixels as on the old one, for every region.** A boolean set equality, not a
+ tolerance, not an IoU.
+
+ Asserted here per region on a scene of thousands, and separately verified on
+ the real 26 566-particle frame (``benchmarks.md``).
+
+The scene is :func:`~spyde.tests.migrated.test_particles_props_parity._blob_field`
+— thousands of ragged, concave, touching, edge-clipped regions — for the reason
+that module gives: three discs agree under any implementation and prove nothing.
+Regions that run off the frame edge matter twice over here, because that is where
+a contour is left OPEN and where the dilation that defines the background ring is
+truncated.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.particles.measure import _contours, _fill_intensity, property_table
+from spyde.signals.particles import COL, N_COLUMNS
+from spyde.tests.migrated.test_particles_props_parity import _blob_field
+
+#: The intensity columns, and the reason each is exact rather than close.
+_INTENSITY_COLS = ("intensity_mean", "intensity_max", "intensity_std",
+ "background")
+
+
+def filled_pixels(contour, shape) -> np.ndarray:
+ """The pixels ``skimage.draw.polygon`` selects, as a sorted flat index.
+
+ This IS what ``render_frame`` and ``mask_at`` compute, including their
+ ``len(c) < 3`` skip. Sorted, because the fill's output ORDER follows the
+ vertex order and two rotations of one cycle may enumerate the same pixels in
+ a different sequence — which no consumer can observe.
+ """
+ from skimage.draw import polygon as sk_polygon
+
+ c = np.asarray(contour)
+ if len(c) < 3:
+ return np.zeros(0, np.int64)
+ rr, cc = sk_polygon(c[:, 0].astype(np.intp), c[:, 1].astype(np.intp),
+ shape=shape)
+ flat = rr.astype(np.int64) * int(shape[1]) + cc.astype(np.int64)
+ flat.sort()
+ return flat
+
+
+def filled_mask(contour, shape) -> np.ndarray:
+ """:func:`filled_pixels` as a boolean image. For assertions that want one."""
+ m = np.zeros(int(shape[0]) * int(shape[1]), bool)
+ m[filled_pixels(contour, shape)] = True
+ return m.reshape(shape)
+
+
+@pytest.fixture(scope="module")
+def scene():
+ return _blob_field()
+
+
+@pytest.fixture(scope="module")
+def table(scene):
+ lab, _img = scene
+ return property_table(lab, fast=True)
+
+
+class TestContourFillParity:
+ """The filled polygon, per region, against the ``find_contours`` loop."""
+
+ @pytest.fixture(scope="class")
+ def both(self, scene, table):
+ lab, _img = scene
+ fast = _contours(lab, table, fast=True)
+ if fast is None: # pragma: no cover
+ pytest.skip("numba unavailable; contours stay on find_contours")
+ return _contours(lab, table, fast=False), fast
+
+ def test_one_outline_per_region(self, both, table):
+ ref, got = both
+ assert len(got) == len(ref) == len(table["label"]) > 2000
+
+ def test_filled_polygon_is_identical_per_region(self, both, scene):
+ """The gate. Every region, exact pixel-set equality, no tolerance."""
+ lab, _img = scene
+ ref, got = both
+ differing = [i for i in range(len(ref))
+ if not np.array_equal(filled_pixels(ref[i], lab.shape),
+ filled_pixels(got[i], lab.shape))]
+ assert not differing, (
+ f"{len(differing)}/{len(ref)} regions fill to a different pixel set; "
+ f"first at index {differing[:5]}")
+
+ def test_vertex_count_matches_and_the_cycle_is_the_same(self, both):
+ """Stronger than the gate and not required by it, but it is TRUE, and
+ it is what says the tracer found the same contour rather than a
+ different one that happens to fill the same.
+
+ A closed contour repeats its first vertex last, so 'the same cycle' means
+ equal after dropping that and rotating."""
+ ref, got = both
+ assert [len(c) for c in ref] == [len(c) for c in got]
+ rotations = 0
+ for a, b in zip(ref, got):
+ a = np.asarray(a, np.int64)
+ b = np.asarray(b, np.int64)
+ closed = len(a) > 1 and np.array_equal(a[0], a[-1])
+ if not closed:
+ assert np.array_equal(a, b) # open paths are bit-identical
+ continue
+ ca, cb = a[:-1], b[:-1]
+ assert np.array_equal(cb[0], cb[-1]) is False or len(cb) == 1
+ hits = [s for s in range(len(cb))
+ if np.array_equal(np.roll(cb, -s, axis=0), ca)]
+ assert hits, "closed contour is not a rotation of the reference"
+ rotations += 1
+ assert rotations > 100, "scene has too few closed contours to be a test"
+
+ def test_int16_csr_layout_still_holds(self, both):
+ """``SpyDEParticles`` stores ``contours`` + ``contour_offsets`` as one
+ int16 ``(N, 2)`` block (``particle_overlay.add_particles`` concatenates
+ and re-slices it), so an outline that is not int16 ``(k, 2)`` breaks the
+ store rather than the drawing."""
+ _ref, got = both
+ for c in got:
+ assert c.dtype == np.int16 and c.ndim == 2 and c.shape[1] == 2
+ pool = np.concatenate(got, axis=0)
+ assert pool.dtype == np.int16
+ offsets = np.concatenate([[0], np.cumsum([len(c) for c in got])])
+ assert offsets[-1] == len(pool)
+
+
+class TestIntensityParity:
+ """The four intensity columns, exact, against the per-region crop loop."""
+
+ @pytest.fixture(scope="class")
+ def rows(self, scene, table):
+ lab, img = scene
+ inten = np.asarray(img, np.float64)
+ n = len(table["label"])
+ keep = np.ones(n, bool)
+
+ def run(fast):
+ r = np.zeros((n, N_COLUMNS), np.float32)
+ for name in _INTENSITY_COLS:
+ r[:, COL[name]] = np.nan
+ _fill_intensity(r, lab, inten, table, keep, 3, fast=fast)
+ return r
+
+ return run(False), run(True)
+
+ @pytest.mark.parametrize("name", _INTENSITY_COLS)
+ def test_column_is_exact(self, rows, name):
+ ref, got = rows
+ a, b = ref[:, COL[name]], got[:, COL[name]]
+ assert np.array_equal(np.isnan(a), np.isnan(b)), f"{name}: NaN pattern"
+ fin = ~np.isnan(a)
+ assert fin.sum() > 2000, f"{name}: nothing measured, the test is vacuous"
+ assert np.array_equal(a[fin], b[fin]), (
+ f"{name}: max |diff| = {np.abs(a[fin] - b[fin]).max()}")
+
+ def test_float64_intermediates_agree_to_summation_order(self, scene, table):
+ """The stored columns are float32, so 'exact' there could in principle
+ hide a real difference. This checks the float64 values the kernels
+ actually produce."""
+ from spyde.particles.intensity import (label_intensity_stats,
+ ring_backgrounds)
+
+ lab, img = scene
+ inten = np.asarray(img, np.float64)
+ labels = np.asarray(table["label"], np.int64)
+ bb = np.stack([np.asarray(table[f"bbox-{k}"], np.int64)
+ for k in range(4)], axis=1)
+ mean, mx, std = label_intensity_stats(lab, inten, labels)
+ bg = ring_backgrounds(lab, inten, labels, bb, 3)
+ if bg is None: # pragma: no cover
+ pytest.skip("numba unavailable")
+
+ from scipy.ndimage import binary_dilation
+ h, w = lab.shape
+ # A subset is enough at float64 resolution and keeps the reference loop
+ # (which is the slow path by construction) off the suite's critical path.
+ for i in range(0, len(labels), 7):
+ lbl = int(labels[i])
+ y0, x0 = int(bb[i, 0]), int(bb[i, 1])
+ y1, x1 = int(bb[i, 2]), int(bb[i, 3])
+ py0, px0 = max(0, y0 - 4), max(0, x0 - 4)
+ py1, px1 = min(h, y1 + 4), min(w, x1 + 4)
+ sub_lab = lab[py0:py1, px0:px1]
+ sub_int = inten[py0:py1, px0:px1]
+ m = sub_lab == lbl
+ vals = sub_int[m]
+ vals = vals[np.isfinite(vals)]
+ assert vals.size
+ assert mean[i] == pytest.approx(vals.mean(), rel=1e-12, abs=0)
+ assert mx[i] == vals.max()
+ assert std[i] * vals.max() == pytest.approx(vals.std(), rel=1e-11,
+ abs=0)
+ ring = binary_dilation(m, iterations=3) & ~m & (sub_lab == 0)
+ bvals = sub_int[ring]
+ bvals = bvals[np.isfinite(bvals)]
+ if bvals.size:
+ assert bg[i] == pytest.approx(bvals.mean(), rel=1e-12, abs=0)
+ else:
+ assert np.isnan(bg[i])
+
+ def test_nan_pixels_are_excluded_not_coerced(self, scene, table):
+ """A drift-corrected frame has a NaN-padded border. Letting NaN reach a
+ plain mean makes every particle touching it report NaN; coercing it to
+ zero invents a dark rim. Both paths must do neither."""
+ lab, img = scene
+ inten = np.asarray(img, np.float64)
+ inten[:8, :] = np.nan
+ inten[:, :8] = np.nan
+ n = len(table["label"])
+ keep = np.ones(n, bool)
+ out = []
+ for fast in (False, True):
+ r = np.zeros((n, N_COLUMNS), np.float32)
+ for name in _INTENSITY_COLS:
+ r[:, COL[name]] = np.nan
+ _fill_intensity(r, lab, inten, table, keep, 3, fast=fast)
+ out.append(r)
+ ref, got = out
+ touching = np.asarray(table["bbox-0"]) < 8
+ assert touching.sum() > 10, "no region touches the NaN border"
+ for name in _INTENSITY_COLS:
+ a, b = ref[:, COL[name]], got[:, COL[name]]
+ assert np.array_equal(np.isnan(a), np.isnan(b)), name
+ fin = ~np.isnan(a)
+ assert np.array_equal(a[fin], b[fin]), name
+ # A region entirely inside the NaN band has no finite pixel at all.
+ assert np.isnan(got[:, COL["intensity_mean"]]).any()
+
+ def test_ring_zero_leaves_background_unset(self, scene, table):
+ lab, img = scene
+ inten = np.asarray(img, np.float64)
+ n = len(table["label"])
+ r = np.zeros((n, N_COLUMNS), np.float32)
+ r[:, COL["background"]] = np.nan
+ _fill_intensity(r, lab, inten, table, np.ones(n, bool), 0, fast=True)
+ assert np.isnan(r[:, COL["background"]]).all()
+ assert not np.isnan(r[:, COL["intensity_mean"]]).all()
+
+
+class TestKernelEdgeCases:
+ """Shapes where an integer reimplementation is most likely to disagree."""
+
+ @pytest.mark.parametrize("build", [
+ pytest.param(lambda a: a.__setitem__((5, 5), 1), id="single-pixel"),
+ pytest.param(lambda a: a.__setitem__((5, slice(2, 9)), 1), id="h-line"),
+ pytest.param(lambda a: a.__setitem__((slice(2, 9), 5), 1), id="v-line"),
+ pytest.param(lambda a: a.__setitem__((0, 0), 1), id="corner"),
+ pytest.param(lambda a: a.__setitem__((slice(0, 3), slice(0, 3)), 1),
+ id="corner-block"),
+ pytest.param(lambda a: [a.__setitem__((i, i), 1) for i in range(2, 9)],
+ id="diagonal"),
+ pytest.param(lambda a: [a.__setitem__((slice(2, 9), 2), 1),
+ a.__setitem__((2, slice(2, 9)), 1)], id="L"),
+ pytest.param(lambda a: [a.__setitem__((slice(2, 9), slice(2, 9)), 1),
+ a.__setitem__((slice(4, 7), slice(4, 7)), 0)],
+ id="ring-with-hole"),
+ pytest.param(lambda a: a.__setitem__((slice(0, 12), slice(0, 12)), 1),
+ id="fills-the-frame"),
+ pytest.param(lambda a: a.__setitem__((slice(0, 12), 0), 1),
+ id="full-left-edge"),
+ ])
+ def test_contour_fill_and_intensity_match(self, build):
+ lab = np.zeros((12, 12), np.int32)
+ build(lab)
+ rng = np.random.default_rng(3)
+ inten = rng.standard_normal(lab.shape) + 5.0
+ tbl = property_table(lab, fast=True)
+ n = len(tbl["label"])
+
+ ref = _contours(lab, tbl, fast=False)
+ got = _contours(lab, tbl, fast=True)
+ if got is None: # pragma: no cover
+ pytest.skip("numba unavailable")
+ for i in range(n):
+ assert np.array_equal(filled_pixels(ref[i], lab.shape),
+ filled_pixels(got[i], lab.shape)), i
+
+ out = []
+ for fast in (False, True):
+ r = np.zeros((n, N_COLUMNS), np.float32)
+ for name in _INTENSITY_COLS:
+ r[:, COL[name]] = np.nan
+ _fill_intensity(r, lab, inten, tbl, np.ones(n, bool), 3, fast=fast)
+ out.append(r)
+ for name in _INTENSITY_COLS:
+ a, b = out[0][:, COL[name]], out[1][:, COL[name]]
+ assert np.array_equal(np.isnan(a), np.isnan(b)), name
+ fin = ~np.isnan(a)
+ assert np.array_equal(a[fin], b[fin]), name
+
+ def test_sparse_labels(self, scene):
+ """Label values that are not 1..N — every upstream filter re-tags, and
+ both kernels index BY ROW, not by label value."""
+ lab, img = scene
+ sparse = np.where(lab > 0, lab.astype(np.int64) * 3 + 7, 0).astype(np.int32)
+ tbl = property_table(sparse, fast=True)
+ n = len(tbl["label"])
+ assert n > 2000
+
+ ref = _contours(sparse, tbl, fast=False)
+ got = _contours(sparse, tbl, fast=True)
+ if got is None: # pragma: no cover
+ pytest.skip("numba unavailable")
+ bad = sum(1 for i in range(n)
+ if not np.array_equal(filled_pixels(ref[i], sparse.shape),
+ filled_pixels(got[i], sparse.shape)))
+ assert bad == 0
+
+ inten = np.asarray(img, np.float64)
+ out = []
+ for fast in (False, True):
+ r = np.zeros((n, N_COLUMNS), np.float32)
+ for name in _INTENSITY_COLS:
+ r[:, COL[name]] = np.nan
+ _fill_intensity(r, sparse, inten, tbl, np.ones(n, bool), 3, fast=fast)
+ out.append(r)
+ for name in _INTENSITY_COLS:
+ a, b = out[0][:, COL[name]], out[1][:, COL[name]]
+ fin = ~np.isnan(a)
+ assert np.array_equal(a[fin], b[fin]), name
+
+ def test_empty_frame(self):
+ lab = np.zeros((16, 16), np.int32)
+ tbl = property_table(lab, fast=True)
+ assert len(tbl["label"]) == 0
+ assert _contours(lab, tbl, fast=True) == []
+
+
+class TestNumbaUnavailable:
+ """Both kernels are OPTIONAL, and the machine without numba must still be
+ able to measure a frame.
+
+ Distinct from ``fast=False``: that asks for the legacy path, this asks for
+ the fast one and has it refuse. The half that does not need numba
+ (``bincount`` statistics) must NOT be left half-written when the half that
+ does is unavailable — a partly-filled row is worse than a slow one."""
+
+ @pytest.fixture
+ def no_numba(self, monkeypatch):
+ from spyde.particles import contours as cmod
+ from spyde.particles import intensity as imod
+ monkeypatch.setattr(cmod, "_build_kernel", lambda: None)
+ monkeypatch.setattr(imod, "_build_ring_kernel", lambda: None)
+
+ def test_measure_frame_still_matches(self, scene, no_numba):
+ from spyde.particles.measure import measure_frame
+
+ lab, img = scene
+ rows_f, cont_f = measure_frame(lab, img, t=2, scale=0.5, fast=True)
+ rows_l, cont_l = measure_frame(lab, img, t=2, scale=0.5, fast=False)
+ assert rows_f.shape == rows_l.shape and rows_f.shape[0] > 2000
+ for name in _INTENSITY_COLS:
+ a, b = rows_l[:, COL[name]], rows_f[:, COL[name]]
+ assert np.array_equal(np.isnan(a), np.isnan(b)), name
+ fin = ~np.isnan(a)
+ assert np.array_equal(a[fin], b[fin]), name
+ # Without numba the fallback IS `find_contours`, so these are identical
+ # vertex for vertex, not merely equal when filled.
+ assert all(np.array_equal(x, y) for x, y in zip(cont_l, cont_f))
+
+ def test_helpers_report_unavailable_rather_than_guessing(self, scene,
+ no_numba):
+ from spyde.particles.contours import label_contours
+ from spyde.particles.intensity import ring_backgrounds
+
+ lab, _img = scene
+ tbl = property_table(lab, fast=True)
+ labels = np.asarray(tbl["label"], np.int64)
+ bb = np.stack([np.asarray(tbl[f"bbox-{k}"], np.int64)
+ for k in range(4)], axis=1)
+ assert label_contours(lab, labels, bb,
+ np.asarray(tbl["area"], np.int64)) is None
+ assert ring_backgrounds(lab, np.zeros(lab.shape), labels, bb, 3) is None
diff --git a/spyde/tests/migrated/test_particles_core.py b/spyde/tests/migrated/test_particles_core.py
new file mode 100644
index 00000000..4218f6fb
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_core.py
@@ -0,0 +1,1018 @@
+"""
+Tests for spyde.particles (classical engine + measurement) and
+spyde.signals.particles (the CSR container).
+
+Acceptance gates from DRIFT_AND_PARTICLES_PLAN.md exercised here:
+
+* B5 measure — matches ``regionprops`` on synthetic shapes of known area and
+ eccentricity, and physical units are right under a non-unit axis scale.
+* B1 classical — separates touching particles; the shared ``split_instances``
+ behaves for a probability map as well as a boolean mask.
+* A7 edges — no particle is ever detected in the NaN-padded border a
+ drift-corrected frame carries. This is called out in the plan as the single most
+ likely integration bug between the two features.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.particles import (
+ THRESHOLD_METHODS,
+ SegmentParams,
+ measure_frame,
+ segment_frame,
+ split_instances,
+ threshold_mask,
+)
+from spyde.signals.particles import COL, COLUMNS, N_COLUMNS, SpyDEParticles
+
+
+# ── synthetic scenes ─────────────────────────────────────────────────────────
+
+def _disc(shape, cy, cx, r, amp=1.0):
+ yy, xx = np.mgrid[0:shape[0], 0:shape[1]]
+ return amp * (((yy - cy) ** 2 + (xx - cx) ** 2) <= r * r)
+
+
+def _field(h=120, w=140, radii=((30, 30, 9), (30, 100, 6), (85, 40, 12),
+ (90, 105, 7)), amp=1.0, bg=0.05):
+ img = np.full((h, w), bg, dtype=np.float32)
+ for cy, cx, r in radii:
+ img += _disc((h, w), cy, cx, r, amp).astype(np.float32)
+ return img
+
+
+def _touching(h=80, w=120):
+ """Two discs of radius 14 whose edges overlap — the watershed's whole job."""
+ img = np.full((h, w), 0.05, dtype=np.float32)
+ img += _disc((h, w), 40, 48, 14).astype(np.float32)
+ img += _disc((h, w), 40, 72, 14).astype(np.float32)
+ return np.clip(img, 0, 1.2)
+
+
+class TestSegmentParams:
+ def test_rejects_unknown_threshold(self):
+ with pytest.raises(ValueError, match="unknown threshold"):
+ SegmentParams(threshold="magic")
+
+ def test_rejects_sensitivity_out_of_range(self):
+ with pytest.raises(ValueError, match="sensitivity must be in 0..1"):
+ SegmentParams(sensitivity=1.5)
+
+ def test_rejects_even_local_size_for_local_methods(self):
+ """Bumping it silently would make the caret disagree with what ran."""
+ with pytest.raises(ValueError, match="local_size must be odd"):
+ SegmentParams(threshold="sauvola", local_size=30)
+ SegmentParams(threshold="otsu", local_size=30) # global: irrelevant
+
+
+class TestThresholding:
+ @pytest.mark.parametrize("method", THRESHOLD_METHODS)
+ def test_every_method_runs_and_finds_the_discs(self, method):
+ # Blurred, so the histogram has real structure. A hard-edged synthetic
+ # field is essentially two delta spikes, which several legitimate methods
+ # cannot work with (see test_minimum_reports_actionably_on_a_spiky_field).
+ from scipy.ndimage import gaussian_filter
+ img = gaussian_filter(_field(), 2.0)
+ p = SegmentParams(threshold=method, local_size=31)
+ mask = threshold_mask(img.astype(np.float32), p)
+ assert mask.shape == img.shape and mask.dtype == bool
+ # All four discs are bright and large; any sane threshold finds signal.
+ assert mask.sum() > 100, f"{method} found almost nothing"
+
+ def test_minimum_reports_actionably_on_a_spiky_field(self):
+ """skimage raises a bare RuntimeError; the user needs to know what to do."""
+ img = _field().astype(np.float32) # hard edges → one-spike histogram
+ with pytest.raises(ValueError, match="requires a clearly bimodal"):
+ threshold_mask(img, SegmentParams(threshold="minimum"))
+
+ def test_sensitivity_half_is_exactly_the_plain_method(self):
+ """0.5 must be a no-op offset, or the ParticleSpy parity gate is meaningless."""
+ from skimage.filters import threshold_otsu
+ img = _field().astype(np.float32)
+ mask = threshold_mask(img, SegmentParams(threshold="otsu", sensitivity=0.5))
+ assert np.array_equal(mask, img > threshold_otsu(img))
+
+ def test_higher_sensitivity_never_shrinks_the_mask(self):
+ img = _field(amp=0.4, bg=0.1).astype(np.float32)
+ prev = -1
+ for s in (0.1, 0.3, 0.5, 0.7, 0.9):
+ n = int(threshold_mask(img, SegmentParams(sensitivity=s)).sum())
+ assert n >= prev, f"sensitivity {s} shrank the mask ({n} < {prev})"
+ prev = n
+
+ def test_invert_finds_dark_particles(self):
+ img = 1.0 - _field(amp=1.0, bg=0.05) # dark discs on bright ground
+ labels = segment_frame(img, SegmentParams(invert=True, min_size=30))
+ assert labels.max() == 4, f"found {labels.max()} dark particles, expected 4"
+
+
+class TestSplitInstances:
+ def test_separates_touching_discs(self):
+ img = _touching()
+ labels = segment_frame(img, SegmentParams(watershed=True, min_size=40))
+ assert labels.max() == 2, (
+ f"watershed merged the pair into {labels.max()} region(s)")
+
+ def test_without_watershed_they_merge(self):
+ """Confirms the previous test is actually measuring the watershed."""
+ img = _touching()
+ labels = segment_frame(img, SegmentParams(watershed=False, min_size=40))
+ assert labels.max() == 1
+
+ def test_does_not_oversplit_a_single_round_disc(self):
+ """The plateau trap: peak_local_max on a flat distance maximum returns
+ several coincident peaks and watershed cuts one disc into wedges."""
+ img = np.full((80, 80), 0.05, np.float32)
+ img += _disc((80, 80), 40, 40, 20).astype(np.float32)
+ labels = segment_frame(img, SegmentParams(watershed=True, min_size=50))
+ assert labels.max() == 1, f"one disc split into {labels.max()} pieces"
+
+ def test_accepts_a_probability_map(self):
+ """The scribble/prompt engines hand over float probabilities, not masks."""
+ prob = np.zeros((60, 60), np.float32)
+ prob[10:25, 10:25] = 0.9
+ prob[35:50, 35:50] = 0.8
+ prob[5:8, 50:53] = 0.3 # below 0.5 — must be ignored
+ labels = split_instances(prob, SegmentParams(min_size=20))
+ assert labels.max() == 2
+
+ def test_min_size_discards_small(self):
+ prob = np.zeros((60, 60), bool)
+ prob[10:30, 10:30] = True # 400 px
+ prob[45:48, 45:48] = True # 9 px
+ assert split_instances(prob, SegmentParams(min_size=100)).max() == 1
+ assert split_instances(prob, SegmentParams(min_size=5)).max() == 2
+
+ def test_a_tiny_particle_survives_the_watershed(self):
+ """§0.9 regression: nothing in the split step may delete a small particle.
+
+ The original marker step filtered markers by AREA (ParticleSpy's
+ ``watershed_size``). A 3x3 particle's local-maximum marker is ONE pixel, so
+ any area floor erased it and the particle vanished — silently, since the
+ frame still had a plausible count.
+ """
+ prob = np.zeros((60, 60), bool)
+ prob[10:30, 10:30] = True # large
+ prob[45:48, 45:48] = True # tiny, 9 px
+ labels = split_instances(prob, SegmentParams(min_size=5, watershed=True))
+ assert labels.max() == 2, "the tiny particle was dropped by the split step"
+ assert labels[46, 46] != 0
+
+ def test_max_size_discards_large(self):
+ prob = np.zeros((60, 60), bool)
+ prob[5:55, 5:55] = True # 2500 px
+ assert split_instances(prob, SegmentParams(max_size=1000)).max() == 0
+
+ def test_clear_border_drops_edge_touching(self):
+ prob = np.zeros((60, 60), bool)
+ prob[0:10, 0:10] = True # touches the border
+ prob[25:40, 25:40] = True
+ p = SegmentParams(min_size=20, clear_border=True, watershed=False)
+ assert split_instances(prob, p).max() == 1
+
+ def test_labels_are_sequential_with_no_gaps(self):
+ prob = np.zeros((80, 80), bool)
+ for i, (y, x) in enumerate([(5, 5), (5, 40), (40, 5), (40, 40)]):
+ prob[y:y + 12, x:x + 12] = True
+ labels = split_instances(prob, SegmentParams(min_size=20, watershed=False))
+ present = np.unique(labels)
+ assert np.array_equal(present, np.arange(present.size))
+
+ def test_empty_input_gives_empty_labels(self):
+ labels = split_instances(np.zeros((20, 20), bool), SegmentParams())
+ assert labels.max() == 0 and labels.dtype == np.int32
+
+ def test_rejects_3d(self):
+ with pytest.raises(ValueError, match="foreground must be 2-D"):
+ split_instances(np.zeros((2, 4, 4), bool), SegmentParams())
+
+
+def _pair_mask(h=80, w=120, r=14, cy=40, cx1=48, cx2=72):
+ """The two overlapping discs of :func:`_touching`, as a boolean mask."""
+ return (_disc((h, w), cy, cx1, r).astype(bool)
+ | _disc((h, w), cy, cx2, r).astype(bool))
+
+
+def _seam(mask, width=2):
+ """The join between the two bodies: the geometric stand-in for what a
+ trained boundary class predicts, and for the strokes a user paints along it."""
+ from scipy import ndimage as ndi
+ ws = split_instances(mask, SegmentParams(min_size=20))
+ # A boundary is the seam BETWEEN two instances, never the outline of one —
+ # a mask of outlines teaches a head to shrink every body and split nothing,
+ # which is the failure this helper's shape exists to avoid reproducing.
+ grown = [ndi.binary_dilation(ws == i, iterations=width)
+ for i in range(1, int(ws.max()) + 1)]
+ seam = np.zeros(mask.shape, bool)
+ for i in range(len(grown)):
+ for j in range(i + 1, len(grown)):
+ seam |= grown[i] & grown[j]
+ return seam & mask
+
+
+class TestBoundarySplit:
+ """The connected-components route: ``split_instances(fg, p, boundary=...)``.
+
+ This exists for speed — a taught boundary lets the split skip the distance
+ transform, the marker/elevation upsample and the watershed, together 1.62 s
+ of a 1.78 s split at 4096². So the bar is not "it produces something": it
+ has to produce **what the watershed produced**, or the speed is not worth
+ having.
+ """
+
+ def test_splits_touching_particles(self):
+ mask = _pair_mask()
+ p = SegmentParams(min_size=40)
+ assert split_instances(mask, p, boundary=_seam(mask)).max() == 2
+
+ def test_matches_the_watershed_count_and_areas(self):
+ """The gate: same count, the same pixels claimed in total, and the same
+ areas to within a fraction of a percent.
+
+ NOT bit-identical, and the difference is inherent rather than a bug: the
+ watershed cuts at the point equidistant from two markers, while reclaim
+ grows both sides one pixel per pass and breaks a tie toward the higher
+ label id. On this pair that moves the cut by four pixels — 0.7% of a 590
+ px body — while the union of the two instances is exactly the same set
+ of pixels. Asserting bit-equality here would be asserting that two
+ different algorithms agree by luck.
+ """
+ mask = _pair_mask()
+ p = SegmentParams(min_size=40)
+ ws = split_instances(mask, p)
+ bd = split_instances(mask, p, boundary=_seam(mask))
+
+ def areas(lab):
+ c = np.bincount(lab.ravel())[1:]
+ return np.sort(c[c > 0])
+
+ a_ws, a_bd = areas(ws), areas(bd)
+ assert bd.max() == ws.max() == 2
+ assert (bd > 0).sum() == (ws > 0).sum(), (
+ "the two routes claimed a different amount of foreground")
+ assert np.array_equal(bd > 0, ws > 0), (
+ "the two routes disagree about which pixels belong to a particle")
+ rel = np.abs(a_bd - a_ws) / a_ws
+ assert rel.max() < 0.02, (
+ f"boundary areas {a_bd.tolist()} differ from watershed "
+ f"{a_ws.tolist()} by {rel.max() * 100:.1f}%")
+
+ def test_the_distance_transform_and_watershed_never_run(self, monkeypatch):
+ """The whole point. If either is still called the route saves nothing,
+ and a passing count test would hide that completely."""
+ from scipy import ndimage as ndi
+ from skimage import segmentation as skseg
+
+ mask = _pair_mask()
+ seam = _seam(mask) # built BEFORE the spies go in —
+ # `_seam` runs a watershed itself, standing in for the user's eye.
+ called = []
+ monkeypatch.setattr(ndi, "distance_transform_edt",
+ lambda *a, **k: called.append("edt"))
+ monkeypatch.setattr(skseg, "watershed",
+ lambda *a, **k: called.append("watershed"))
+
+ split_instances(mask, SegmentParams(min_size=40), boundary=seam)
+ assert called == [], f"the boundary route still ran {called}"
+
+ def test_no_boundary_falls_back_to_the_watershed(self):
+ """A user who has never painted a boundary must not silently get worse
+ splitting — so an absent boundary is the old behaviour, bit for bit."""
+ mask = _pair_mask()
+ p = SegmentParams(min_size=40)
+ ws = split_instances(mask, p)
+ assert np.array_equal(split_instances(mask, p, boundary=None), ws)
+
+ def test_an_all_false_boundary_also_falls_back(self):
+ """"The class exists but nothing is painted yet" is the same situation as
+ "there is no boundary", and it is the common one mid-session. Treating an
+ empty mask as a real boundary would hand watershed's job to plain
+ connected components and merge every touching pair."""
+ mask = _pair_mask()
+ p = SegmentParams(min_size=40)
+ ws = split_instances(mask, p)
+ empty = split_instances(mask, p, boundary=np.zeros_like(mask))
+ assert np.array_equal(empty, ws)
+
+ def test_isolated_particles_are_untouched(self):
+ """An isolated body has no seam through it, so its core IS the body and
+ the two routes cannot disagree."""
+ mask = np.zeros((80, 120), bool)
+ mask |= _disc((80, 120), 25, 25, 12).astype(bool)
+ mask |= _disc((80, 120), 25, 90, 12).astype(bool)
+ p = SegmentParams(min_size=40)
+ seam = _seam(mask)
+ assert not seam.any(), "these discs do not touch; there is no seam"
+ assert np.array_equal(split_instances(mask, p, boundary=seam),
+ split_instances(mask, p))
+
+ def test_a_probability_boundary_is_thresholded_like_the_foreground(self):
+ mask = _pair_mask()
+ seam = _seam(mask)
+ soft = np.where(seam, 0.9, 0.1).astype(np.float32)
+ p = SegmentParams(min_size=40)
+ assert np.array_equal(split_instances(mask, p, boundary=soft),
+ split_instances(mask, p, boundary=seam))
+
+ def test_a_weak_probability_boundary_is_ignored(self):
+ """Below 0.5 everywhere is no boundary at all — and must therefore fall
+ back rather than run the fast route on an empty seam."""
+ mask = _pair_mask()
+ soft = np.where(_seam(mask), 0.3, 0.05).astype(np.float32)
+ p = SegmentParams(min_size=40)
+ assert np.array_equal(split_instances(mask, p, boundary=soft),
+ split_instances(mask, p))
+
+ def test_a_mismatched_boundary_shape_raises(self):
+ with pytest.raises(ValueError, match="must describe the same frame"):
+ split_instances(np.zeros((10, 10), bool), SegmentParams(),
+ boundary=np.zeros((10, 12), bool))
+
+ def test_rejects_a_3d_boundary(self):
+ with pytest.raises(ValueError, match="boundary must be 2-D"):
+ split_instances(np.zeros((10, 10), bool), SegmentParams(),
+ boundary=np.zeros((2, 10, 10), bool))
+
+ def test_a_tiny_particle_survives_the_boundary_route(self):
+ """§0.9 again, on the new path: the split step may never delete a small
+ body. A 3x3 particle has no seam through it, so it must come out whole."""
+ mask = np.zeros((60, 60), bool)
+ mask[10:30, 10:30] = True
+ mask[45:48, 45:48] = True # 9 px
+ bnd = np.zeros((60, 60), bool)
+ bnd[19:21, 10:30] = True # a seam across the big one
+ labels = split_instances(mask, SegmentParams(min_size=5), boundary=bnd)
+ assert labels[46, 46] != 0, "the tiny particle was dropped"
+ assert labels.max() == 3, "the seam should have cut the large body in two"
+
+
+class TestBoundaryReclaim:
+ """Growing the instances back over the seam — what keeps the areas honest."""
+
+ def test_the_seam_is_fully_reclaimed_by_default(self):
+ """Default is grow-to-convergence, so every foreground pixel reachable
+ from a core ends up owned by one."""
+ mask = _pair_mask()
+ labels = split_instances(mask, SegmentParams(min_size=40),
+ boundary=_seam(mask, width=3))
+ assert int((labels > 0).sum()) == int(mask.sum()), (
+ "some foreground was left unassigned with reclaim running to "
+ "convergence")
+
+ def test_capping_the_passes_leaves_part_of_the_seam_unassigned(self):
+ """Confirms the previous test measures something: with one pass a wide
+ seam cannot be closed, so the areas come out low."""
+ mask = _pair_mask()
+ seam = _seam(mask, width=3)
+ capped = split_instances(
+ mask, SegmentParams(min_size=40, boundary_reclaim=1), boundary=seam)
+ assert int((capped > 0).sum()) < int(mask.sum())
+
+ def test_the_count_is_the_same_however_many_passes_run(self):
+ """Reclaim moves pixels between instances; it must never create or
+ destroy one. That is what makes the cap a quality knob and not a
+ correctness one."""
+ mask = _pair_mask()
+ seam = _seam(mask, width=3)
+ counts = {
+ int(split_instances(mask,
+ SegmentParams(min_size=40, boundary_reclaim=k),
+ boundary=seam).max())
+ for k in (1, 2, 3, 5, 0)
+ }
+ assert counts == {2}, f"the pass count changed the particle count: {counts}"
+
+ def test_foreground_with_no_core_at_all_stays_unassigned(self):
+ """A body entirely covered by boundary belongs to no instance, and
+ inventing an owner for it would be worse than leaving it out."""
+ mask = np.zeros((40, 40), bool)
+ mask[5:25, 5:25] = True # a real body
+ mask[32:35, 32:35] = True # fully fenced in below
+ bnd = np.zeros((40, 40), bool)
+ bnd[32:35, 32:35] = True
+ labels = split_instances(mask, SegmentParams(min_size=4), boundary=bnd)
+ assert labels.max() == 1
+ assert not labels[32:35, 32:35].any()
+
+
+class TestFinalizeLabels:
+ """``_finalize_labels`` fuses the size filter and the sequential relabel."""
+
+ def test_identical_to_the_chain_it_replaces(self):
+ """It reads the raster twice instead of six times, so it has to be
+ proven equal to the obvious version rather than merely similar."""
+ from spyde.particles.classical import (_drop_large, _drop_small,
+ _finalize_labels,
+ _relabel_sequential)
+ rng = np.random.default_rng(0)
+ for trial in range(60):
+ lab = rng.integers(0, 12, size=(24, 24)).astype(np.int32)
+ if trial % 3 == 0:
+ lab[lab == 5] = 0 # punch a gap in the ids
+ p = SegmentParams(min_size=int(rng.integers(0, 10)),
+ max_size=int(rng.integers(0, 60)))
+ ref = lab
+ if p.max_size > 0:
+ ref = _drop_large(ref, p.max_size)
+ if p.min_size > 0:
+ ref = _drop_small(ref, p.min_size)
+ ref = _relabel_sequential(ref)
+ assert np.array_equal(_finalize_labels(lab, p), ref), (
+ f"trial {trial}: min_size={p.min_size} max_size={p.max_size}")
+
+
+class TestNaNBorder:
+ """Plan trap #2 / gate A7 — the drift↔segmentation seam."""
+
+ def test_nan_border_yields_no_particles_there(self):
+ from spyde.drift import shift_frame
+ img = _field()
+ shifted = shift_frame(img, (12, -15)) # NaN band top and right
+ assert np.isnan(shifted).any(), "test setup produced no NaN border"
+
+ labels = segment_frame(shifted, SegmentParams(min_size=30))
+ nan_mask = ~np.isfinite(shifted)
+ assert not np.any(labels[nan_mask]), (
+ "found a particle inside the NaN-padded border — the padding was "
+ "coerced to a value that thresholds as signal")
+
+ def test_nan_does_not_erase_real_data_near_the_border(self):
+ """The opposite failure: propagating NaN through the filters wipes a band."""
+ from spyde.drift import shift_frame
+ img = _field()
+ shifted = shift_frame(img, (5, 0))
+ labels = segment_frame(shifted, SegmentParams(gaussian=2.0, min_size=30))
+ assert labels.max() == 4, (
+ f"found {labels.max()} of 4 particles — NaN bled through the blur")
+
+
+class TestMeasure:
+ def test_area_matches_regionprops(self):
+ from skimage.measure import regionprops
+ prob = np.zeros((80, 80), bool)
+ prob[10:30, 10:40] = True # exactly 600 px
+ labels = split_instances(prob, SegmentParams(min_size=10, watershed=False))
+ rows, _ = measure_frame(labels)
+ assert rows.shape[1] == N_COLUMNS
+ ref = regionprops(labels)[0]
+ assert rows[0, COL["area"]] == pytest.approx(ref.area)
+ assert rows[0, COL["area"]] == pytest.approx(600)
+
+ def test_circularity_of_a_disc_is_near_one(self):
+ labels = _disc((120, 120), 60, 60, 30).astype(np.int32)
+ rows, _ = measure_frame(labels)
+ # A pixelated disc's perimeter is slightly over-estimated, so ~0.9-1.05.
+ assert 0.85 < rows[0, COL["circularity"]] < 1.1, rows[0, COL["circularity"]]
+
+ def test_eccentricity_of_a_disc_is_near_zero(self):
+ labels = _disc((120, 120), 60, 60, 25).astype(np.int32)
+ rows, _ = measure_frame(labels)
+ assert rows[0, COL["eccentricity"]] < 0.2
+
+ def test_calibration_scales_length_and_area_correctly(self):
+ prob = np.zeros((60, 60), bool)
+ prob[10:30, 10:30] = True # 400 px, 20 px across
+ labels = split_instances(prob, SegmentParams(min_size=10, watershed=False))
+ r1, _ = measure_frame(labels, scale=1.0)
+ r2, _ = measure_frame(labels, scale=0.5) # 0.5 nm/px
+ assert r2[0, COL["area"]] == pytest.approx(r1[0, COL["area"]] * 0.25)
+ assert r2[0, COL["perimeter"]] == pytest.approx(r1[0, COL["perimeter"]] * 0.5)
+ assert r2[0, COL["y"]] == pytest.approx(r1[0, COL["y"]] * 0.5)
+
+ def test_circularity_is_scale_invariant(self):
+ """Dimensionless quantities must NOT pick up a scale factor."""
+ labels = _disc((100, 100), 50, 50, 22).astype(np.int32)
+ a, _ = measure_frame(labels, scale=1.0)
+ b, _ = measure_frame(labels, scale=0.137)
+ assert a[0, COL["circularity"]] == pytest.approx(b[0, COL["circularity"]])
+ assert a[0, COL["eccentricity"]] == pytest.approx(b[0, COL["eccentricity"]])
+
+ def test_intensity_excludes_nan(self):
+ labels = np.zeros((40, 40), np.int32)
+ labels[10:20, 10:20] = 1
+ inten = np.full((40, 40), 5.0)
+ inten[12, 12] = np.nan
+ rows, _ = measure_frame(labels, inten)
+ assert rows[0, COL["intensity_mean"]] == pytest.approx(5.0), (
+ "NaN leaked into the mean, or was coerced to 0 and dragged it down")
+
+ def test_background_ring_ignores_neighbouring_particles(self):
+ """A touching neighbour's body must not be measured as this one's background."""
+ labels = np.zeros((40, 60), np.int32)
+ labels[15:25, 10:20] = 1
+ labels[15:25, 21:31] = 2 # 1 px gap — inside a 3 px ring
+ inten = np.zeros((40, 60))
+ inten[labels == 1] = 10.0
+ inten[labels == 2] = 99.0 # a bright neighbour
+ inten[labels == 0] = 1.0
+ rows, _ = measure_frame(labels, inten, background_ring=3)
+ bg = rows[0, COL["background"]]
+ assert bg == pytest.approx(1.0), (
+ f"background {bg} — the neighbour's 99 leaked in")
+
+ def test_min_area_px_filters_before_returning(self):
+ labels = np.zeros((40, 40), np.int32)
+ labels[5:25, 5:25] = 1 # 400
+ labels[30:33, 30:33] = 2 # 9
+ rows, contours = measure_frame(labels, min_area_px=100)
+ assert len(rows) == 1 and len(contours) == 1
+
+ def test_contours_match_rows_one_to_one(self):
+ labels = np.zeros((60, 60), np.int32)
+ labels[5:20, 5:20] = 1
+ labels[30:50, 30:55] = 2
+ rows, contours = measure_frame(labels)
+ assert len(rows) == len(contours) == 2
+ for c in contours:
+ assert c.ndim == 2 and c.shape[1] == 2 and c.dtype == np.int16
+
+ def test_single_pixel_region_still_gets_a_contour(self):
+ """Degenerate regions must not break the 1:1 correspondence."""
+ labels = np.zeros((20, 20), np.int32)
+ labels[10, 10] = 1
+ rows, contours = measure_frame(labels)
+ assert len(rows) == 1 and len(contours) == 1 and len(contours[0]) >= 3
+
+ def test_empty_label_image(self):
+ rows, contours = measure_frame(np.zeros((20, 20), np.int32))
+ assert rows.shape == (0, N_COLUMNS) and contours == []
+
+ def test_track_id_starts_unassigned(self):
+ labels = np.zeros((30, 30), np.int32)
+ labels[5:20, 5:20] = 1
+ rows, _ = measure_frame(labels)
+ assert rows[0, COL["track_id"]] == -1
+
+ def test_intensity_shape_mismatch_raises(self):
+ with pytest.raises(ValueError, match="intensity shape"):
+ measure_frame(np.zeros((10, 10), np.int32), np.zeros((8, 8)))
+
+
+# ── the container ────────────────────────────────────────────────────────────
+
+def _build(n_frames=4, seed=1):
+ """A small SpyDEParticles built through the real segment→measure path."""
+ rng = np.random.default_rng(seed)
+ per_frame, contours = [], []
+ for t in range(n_frames):
+ img = np.full((90, 90), 0.05, np.float32)
+ # A growing number of particles per frame, so count_series is non-trivial.
+ for i in range(t + 1):
+ cy = 20 + 25 * (i % 3)
+ cx = 20 + 25 * (i // 3) + int(rng.integers(0, 3))
+ img += _disc((90, 90), cy, cx, 8).astype(np.float32)
+ labels = segment_frame(img, SegmentParams(min_size=30))
+ rows, cs = measure_frame(labels, img, t=t, scale=0.5)
+ per_frame.append(rows)
+ contours.append(cs)
+ return SpyDEParticles.from_frames(
+ per_frame, frame_shape=(90, 90), contours_per_frame=contours,
+ scale=0.5, units="nm")
+
+
+class TestSpyDEParticles:
+ def test_builds_and_reports_shape(self):
+ p = _build()
+ assert p.n_frames == 4
+ assert p.n_particles == 1 + 2 + 3 + 4
+ assert p.has_masks and not p.has_tracks
+
+ def test_csr_slice_is_o1_and_correct(self):
+ p = _build()
+ for t in range(p.n_frames):
+ blk = p.at(t)
+ assert len(blk) == t + 1
+ assert np.all(blk[:, COL["t"]] == t)
+
+ def test_count_series_matches_offsets(self):
+ p = _build()
+ assert np.array_equal(p.count_series(), np.array([1, 2, 3, 4], np.float32))
+
+ def test_property_series_is_nan_on_empty_frames(self):
+ """An empty frame has no mean size; zero would draw a fake event spike."""
+ rows = np.zeros((2, N_COLUMNS), np.float32)
+ rows[:, COL["area"]] = [10.0, 20.0]
+ p = SpyDEParticles.from_frames(
+ [rows, np.zeros((0, N_COLUMNS), np.float32)], frame_shape=(10, 10))
+ s = p.property_series("area", "mean")
+ assert s[0] == pytest.approx(15.0)
+ assert np.isnan(s[1])
+
+ def test_property_series_reductions(self):
+ rows = np.zeros((3, N_COLUMNS), np.float32)
+ rows[:, COL["area"]] = [1.0, 2.0, 6.0]
+ p = SpyDEParticles.from_frames([rows], frame_shape=(10, 10))
+ assert p.property_series("area", "sum")[0] == pytest.approx(9.0)
+ assert p.property_series("area", "max")[0] == pytest.approx(6.0)
+ assert p.property_series("area", "median")[0] == pytest.approx(2.0)
+
+ def test_unknown_reduce_and_column_raise(self):
+ p = _build()
+ with pytest.raises(ValueError, match="unknown reduce"):
+ p.property_series("area", "bogus")
+ with pytest.raises(KeyError, match="unknown column"):
+ p.column("nope")
+
+ def test_frame_index_bounds(self):
+ p = _build()
+ with pytest.raises(IndexError, match="outside"):
+ p.at(99)
+
+ def test_render_frame_paints_the_right_count(self):
+ p = _build()
+ img = p.render_frame(2)
+ assert img.shape == (90, 90) and img.dtype == np.int32
+ painted = np.unique(img)
+ painted = painted[painted > 0]
+ assert painted.size == 3, f"painted {painted.size} of 3 particles"
+
+ def test_render_frame_by_index_supports_hit_testing(self):
+ p = _build()
+ img = p.render_frame(1, value="index")
+ vals = np.unique(img)
+ vals = vals[vals > 0] - 1
+ assert set(vals.tolist()) == set(p.indices_at(1).tolist())
+
+ def test_render_frame_rejects_unknown_value(self):
+ p = _build()
+ with pytest.raises(ValueError, match="unknown value"):
+ p.render_frame(0, value="colour")
+
+ def test_render_without_masks_raises_clearly(self):
+ rows = np.zeros((1, N_COLUMNS), np.float32)
+ p = SpyDEParticles.from_frames([rows], frame_shape=(10, 10))
+ assert not p.has_masks
+ with pytest.raises(ValueError, match="store_masks=False"):
+ p.render_frame(0)
+
+ def test_mask_at_is_cropped_to_bbox(self):
+ p = _build()
+ m, (y0, x0, y1, x1) = p.mask_at(0)
+ assert m.shape == (y1 - y0, x1 - x0)
+ assert m.any(), "empty mask for a real particle"
+ assert m.size < 90 * 90, "mask was not cropped"
+
+ def test_validation_rejects_bad_shapes(self):
+ with pytest.raises(ValueError, match=r"flat_buffer must be"):
+ SpyDEParticles(np.zeros((3, 5)), np.array([0, 3]), (10, 10))
+ with pytest.raises(ValueError, match="t_offsets must span"):
+ SpyDEParticles(np.zeros((3, N_COLUMNS)), np.array([0, 2]), (10, 10))
+ with pytest.raises(ValueError, match="non-decreasing"):
+ SpyDEParticles(np.zeros((3, N_COLUMNS)), np.array([0, 3, 1, 3]), (10, 10))
+
+ def test_contours_without_offsets_rejected(self):
+ with pytest.raises(ValueError, match="both be set or both None"):
+ SpyDEParticles(np.zeros((1, N_COLUMNS)), np.array([0, 1]), (10, 10),
+ contours=np.zeros((4, 2), np.int16))
+
+ def test_from_frames_rejects_mismatched_contours(self):
+ rows = np.zeros((2, N_COLUMNS), np.float32)
+ with pytest.raises(ValueError, match="outlines must correspond 1:1"):
+ SpyDEParticles.from_frames(
+ [rows], frame_shape=(10, 10),
+ contours_per_frame=[[np.zeros((4, 2), np.int16)]]) # 1 for 2
+
+ def test_from_frames_rejects_frame_count_mismatch(self):
+ rows = np.zeros((1, N_COLUMNS), np.float32)
+ with pytest.raises(ValueError, match="contours_per_frame has"):
+ SpyDEParticles.from_frames([rows, rows], frame_shape=(10, 10),
+ contours_per_frame=[[]])
+
+ def test_save_load_round_trip(self, tmp_path):
+ p = _build()
+ path = str(tmp_path / "p.npz")
+ p.save(path)
+ back = SpyDEParticles.load(path)
+ assert np.array_equal(back.flat_buffer, p.flat_buffer)
+ assert np.array_equal(back.t_offsets, p.t_offsets)
+ assert np.array_equal(back.contours, p.contours)
+ assert back.frame_shape == p.frame_shape
+ assert back.scale == p.scale and back.units == "nm"
+ # And the reloaded object still renders.
+ assert back.render_frame(1).max() > 0
+
+ def test_save_load_without_masks(self, tmp_path):
+ rows = np.zeros((2, N_COLUMNS), np.float32)
+ p = SpyDEParticles.from_frames([rows], frame_shape=(8, 8))
+ path = str(tmp_path / "nomask.npz")
+ p.save(path)
+ back = SpyDEParticles.load(path)
+ assert not back.has_masks
+
+ def test_load_rejects_future_format(self, tmp_path):
+ import json
+ path = str(tmp_path / "bad.npz")
+ np.savez_compressed(
+ path, flat_buffer=np.zeros((0, N_COLUMNS), np.float32),
+ t_offsets=np.array([0]),
+ meta=np.array(json.dumps({"format_version": 999,
+ "columns": list(COLUMNS),
+ "frame_shape": [4, 4]})))
+ with pytest.raises(ValueError, match="unsupported SpyDEParticles format"):
+ SpyDEParticles.load(path)
+
+ def test_load_rejects_changed_column_layout(self, tmp_path):
+ """A REORDERED/renamed layout is unreadable — but see the test below:
+ columns merely APPENDED since are not a layout change."""
+ import json
+ from spyde.signals.particles import FORMAT_VERSION
+ path = str(tmp_path / "cols.npz")
+ np.savez_compressed(
+ path, flat_buffer=np.zeros((0, 2), np.float32),
+ t_offsets=np.array([0]),
+ meta=np.array(json.dumps({"format_version": FORMAT_VERSION,
+ "columns": ["label", "t"], # swapped
+ "frame_shape": [4, 4]})))
+ with pytest.raises(ValueError, match="column layout changed"):
+ SpyDEParticles.load(path)
+
+ def test_load_migrates_a_file_written_before_a_column_was_appended(self, tmp_path):
+ """An older file must still open.
+
+ New columns go at the END precisely so this is possible. Refusing the
+ file would make every previously-saved particle result unopenable in
+ order to add one DERIVED number — and it is derived, so it never needed
+ to have been stored.
+
+ The migrated `score` is 1.0, not 0.0: absent evidence is not evidence of
+ a bad particle, and the caret's confidence filter must not silently
+ delete instances measured before the score existed.
+ """
+ import json
+ from spyde.signals.particles import COLUMNS as _COLS, COL as _COL
+ old_cols = _COLS[:-1] # everything before `score`
+ path = str(tmp_path / "old.npz")
+ buf = np.zeros((3, len(old_cols)), np.float32)
+ buf[:, 0] = [0, 0, 1] # t
+ np.savez_compressed(
+ path, flat_buffer=buf, t_offsets=np.array([0, 2, 3]),
+ meta=np.array(json.dumps({"format_version": 1,
+ "columns": list(old_cols),
+ "frame_shape": [4, 4]})))
+ p = SpyDEParticles.load(path)
+ assert p.flat_buffer.shape == (3, N_COLUMNS)
+ assert np.all(p.flat_buffer[:, _COL["score"]] == 1.0), (
+ "migrated particles must not be filtered away by the confidence "
+ "slider — they carry no evidence, not bad evidence")
+
+ def test_to_csv_writes_a_header_and_every_row(self, tmp_path):
+ p = _build()
+ path = str(tmp_path / "p.csv")
+ p.to_csv(path)
+ with open(path) as fh:
+ lines = fh.read().strip().splitlines()
+ assert lines[0].split(",") == list(COLUMNS)
+ assert len(lines) == p.n_particles + 1
+
+ def test_repr_is_informative(self):
+ assert "particles over" in repr(_build())
+
+
+class TestNaNBorderPolarity:
+ """The drift-padded border must never segment — in EITHER polarity.
+
+ `spyde.drift.warp` names this as the most likely integration bug in the
+ feature: "a threshold applied to NaN ... invents a large 'particle' along
+ the edge that then nucleates a spurious track". `_prepare` fills NaN before
+ filtering (filters propagate NaN outward and would erase real data), so the
+ fill value has to read as background in the image that is finally
+ THRESHOLDED — and `invert` flips that image after the fill.
+
+ Filling with the finite minimum while inverting made the padding the
+ BRIGHTEST region: measured, a 240 px instance sitting on the border of a
+ 96x96 frame.
+ """
+
+ @staticmethod
+ def _frame(dark_particles: bool):
+ import numpy as _np
+ h = w = 96
+ bg, fg = (200.0, 40.0) if dark_particles else (40.0, 200.0)
+ img = _np.full((h, w), bg, _np.float32)
+ y, x = _np.mgrid[0:h, 0:w]
+ for cy, cx in ((40, 40), (65, 70)):
+ img[(y - cy) ** 2 + (x - cx) ** 2 < 8 ** 2] = fg
+ img[:12, :] = _np.nan # the drift-corrected border
+ img[:, :12] = _np.nan
+ return img
+
+ @pytest.mark.parametrize("invert", [False, True])
+ def test_the_padded_border_never_becomes_a_particle(self, invert):
+ from spyde.particles.classical import SegmentParams, _prepare, segment_frame
+
+ img = self._frame(dark_particles=invert)
+ p = SegmentParams(invert=invert, rb_kernel=0, gaussian=0)
+
+ prep = _prepare(img, p)
+ border = prep[:12, :12]
+ assert border.mean() < prep.max() - 1e-3, (
+ f"invert={invert}: the NaN padding is the brightest region in the "
+ f"thresholded image, so it segments as one huge edge particle")
+
+ labels = segment_frame(img, p)
+ on_border = sorted(set(np.unique(labels[:12, :12])) - {0})
+ assert not on_border, (
+ f"invert={invert}: instances {on_border} were found on the NaN "
+ f"border, which is padding and not data")
+
+
+class TestConfidenceScore:
+ """One slider that means "fewer / more particles" on every engine.
+
+ The reported failure was 547 instances where ~30 were real, and six Advanced
+ knobs none of which says "fewer particles": min size, max size, split
+ touching, min separation, marker smoothing, drop-edge. Size and shape cannot
+ fix it — over-split support-film texture is often SMALL and ROUND, which is
+ exactly what a size/circularity filter keeps.
+
+ Contrast-to-noise against the instance's own dilated background ring is the
+ statistic that does separate them: a real particle sits well away from its
+ surroundings, while a fragment of textured film is by construction the same
+ brightness as the texture around it.
+ """
+
+ @staticmethod
+ def _field():
+ """~30 genuinely dark particles in a noisy support film, plus ~300
+ labelled fragments OF that film — the reported failure, in miniature."""
+ rng = np.random.default_rng(0)
+ h = w = 256
+ frame = np.full((h, w), 200.0, np.float32)
+ frame += rng.normal(0, 6.0, (h, w)).astype(np.float32)
+ labels = np.zeros((h, w), np.int32)
+ truth, nxt = {}, 1
+ y, x = np.mgrid[0:h, 0:w]
+ for _ in range(30):
+ cy, cx, r = rng.uniform(15, h - 15), rng.uniform(15, w - 15), rng.uniform(5, 10)
+ m = (y - cy) ** 2 + (x - cx) ** 2 < r * r
+ frame[m] = 120.0 + rng.normal(0, 5.0, int(m.sum())).astype(np.float32)
+ labels[m] = nxt; truth[nxt] = True; nxt += 1
+ for _ in range(300):
+ cy, cx, r = rng.uniform(8, h - 8), rng.uniform(8, w - 8), rng.uniform(3, 6)
+ m = ((y - cy) ** 2 + (x - cx) ** 2 < r * r) & (labels == 0)
+ if m.sum() < 20:
+ continue
+ labels[m] = nxt; truth[nxt] = False; nxt += 1
+ return frame, labels, truth
+
+ def test_the_score_separates_particles_from_textured_film(self):
+ from spyde.particles.measure import measure_frame
+ from spyde.signals.particles import COL
+
+ frame, labels, truth = self._field()
+ rows, _ = measure_frame(labels, frame, t=0)
+ scores = rows[:, COL["score"]]
+ real = np.array([truth.get(int(l), False)
+ for l in rows[:, COL["label"]].astype(int)])
+ assert real.sum() and (~real).sum(), "the fixture built only one population"
+ # A gap, not merely a difference in means: the slider is only usable if
+ # SOME threshold cleanly separates them.
+ assert np.percentile(scores[real], 10) > np.percentile(scores[~real], 90), (
+ f"no separating threshold exists: real p10="
+ f"{np.percentile(scores[real], 10):.3f} vs texture p90="
+ f"{np.percentile(scores[~real], 90):.3f}")
+
+ def test_filtering_keeps_rows_and_contours_aligned(self):
+ """A misalignment draws one particle's outline on another's row."""
+ from spyde.actions.particles_action import filter_by_score
+ from spyde.signals.particles import COL, N_COLUMNS
+
+ rows = np.zeros((5, N_COLUMNS), np.float32)
+ rows[:, COL["score"]] = [0.05, 0.2, 0.6, 0.95, 0.99]
+ rows[:, COL["label"]] = [1, 2, 3, 4, 5]
+ contours = [np.full((3, 2), i, np.int16) for i in range(5)]
+ kept_rows, kept_contours = filter_by_score(rows, contours, 0.5)
+ assert len(kept_rows) == len(kept_contours) == 3
+ # the SURVIVING contours must be the ones belonging to the kept rows
+ assert [int(c[0, 0]) for c in kept_contours] == [2, 3, 4]
+
+ def test_zero_is_the_old_behaviour_untouched(self):
+ from spyde.actions.particles_action import filter_by_score
+ from spyde.signals.particles import COL, N_COLUMNS
+
+ rows = np.zeros((3, N_COLUMNS), np.float32)
+ rows[:, COL["score"]] = [0.0, 0.5, 1.0]
+ cont = [np.zeros((2, 2), np.int16)] * 3
+ out_rows, out_cont = filter_by_score(rows, cont, 0.0)
+ assert out_rows is rows and out_cont is cont, (
+ "the default must be a no-op, not a copy")
+
+ def test_unmeasurable_particles_are_kept_not_hidden(self):
+ """No background ring => no evidence => not marginal.
+
+ Scoring an unmeasurable instance 0 would let the slider silently delete
+ particles it knows nothing about, which is the opposite of what a
+ 'hide the marginal ones' control should do.
+ """
+ from spyde.particles.measure import particle_scores
+ from spyde.signals.particles import COL, N_COLUMNS
+
+ rows = np.zeros((1, N_COLUMNS), np.float32)
+ rows[:, COL["intensity_mean"]] = 100.0
+ rows[:, COL["background"]] = np.nan # ring fell outside the frame
+ rows[:, COL["intensity_std"]] = np.nan
+ assert particle_scores(rows)[0] == 1.0
+
+
+class TestPerComponentWatershed:
+ """Large frames watershed each component in its own bbox.
+
+ A whole-frame watershed allocates several full-frame rasters — measured at
+ 4096², `split_instances` alone peaks at 546 MB of an 852 MB total. With
+ dask's `threads_per_worker=4` that is ~3.4 GB of concurrent peak per worker,
+ which drove a 9.24 GiB worker into a pause/resume/restart loop with
+ "unmanaged memory" warnings (the arrays are ours, inside the task, so dask
+ can neither see nor spill them).
+
+ Cropping is EXACT, not an approximation: a connected component is surrounded
+ by background by definition, so a 1 px pad contains every pixel the distance
+ transform, the markers and the watershed can depend on, and no watershed can
+ flow between two components that do not touch.
+ """
+
+ @staticmethod
+ def _field(n, k, seed=0):
+ rng = np.random.default_rng(seed)
+ fg = np.zeros((n, n), bool)
+ y, x = np.mgrid[0:n, 0:n]
+ for _ in range(k):
+ cy, cx = rng.uniform(20, n - 20), rng.uniform(20, n - 20)
+ r = rng.uniform(8, 22)
+ fg |= (y - cy) ** 2 + (x - cx) ** 2 < r * r # overlapping on purpose
+ return fg
+
+ def test_identical_to_the_whole_frame_route_when_neither_decimates(self):
+ """The exactness claim, where it can be checked directly.
+
+ Below `_SPLIT_DECIMATE_ABOVE` both routes compute the split geometry at
+ full resolution, so they must agree pixel for pixel. Above it they do
+ NOT, and that is decimation rather than a cropping error — see the next
+ test.
+ """
+ import spyde.particles.classical as C
+ from spyde.particles.classical import SegmentParams, split_instances
+
+ fg = self._field(1024, 90)
+ p = SegmentParams(min_size=5, watershed=True)
+ assert C._split_factor(fg.shape, p) == 1, "fixture must not decimate"
+
+ orig = C._COMPONENT_ROUTE_PX
+ try:
+ C._COMPONENT_ROUTE_PX = 1 << 60 # force whole-frame
+ whole = split_instances(fg, p)
+ C._COMPONENT_ROUTE_PX = 0 # force per-component
+ comp = split_instances(fg, p)
+ finally:
+ C._COMPONENT_ROUTE_PX = orig
+
+ assert whole.max() == comp.max(), (
+ f"instance COUNT differs with no decimation in play: "
+ f"{whole.max()} vs {comp.max()}")
+ # Ids may be numbered differently; the PARTITION must be the same.
+ for v in np.unique(whole):
+ if v == 0:
+ continue
+ m = whole == v
+ ids = np.unique(comp[m])
+ assert len(ids) == 1 and ids[0] != 0, (
+ "a whole-frame instance was split across the crop boundary")
+
+ def test_a_component_spanning_a_crop_is_not_broken_up(self):
+ """One long diagonal particle — the shape a naive band split ruins."""
+ import spyde.particles.classical as C
+ from spyde.particles.classical import SegmentParams, split_instances
+
+ n = 512
+ fg = np.zeros((n, n), bool)
+ for i in range(40, n - 40):
+ fg[i - 3:i + 3, i - 3:i + 3] = True # corner to corner
+ orig = C._COMPONENT_ROUTE_PX
+ try:
+ C._COMPONENT_ROUTE_PX = 0
+ lab = split_instances(fg, SegmentParams(min_size=5, watershed=True))
+ finally:
+ C._COMPONENT_ROUTE_PX = orig
+ assert lab.max() >= 1
+ # every foreground pixel belongs to a labelled instance
+ assert (lab[fg] > 0).all(), "pixels were dropped at a crop edge"
+
+ def test_labels_are_globally_unique_across_components(self):
+ """Each crop labels 1..k locally; the offset into the global space is
+ where an off-by-one silently merges two particles."""
+ import spyde.particles.classical as C
+ from spyde.particles.classical import SegmentParams, split_instances
+
+ fg = np.zeros((256, 256), bool)
+ fg[20:60, 20:60] = True # three well-separated squares
+ fg[20:60, 120:160] = True
+ fg[120:160, 20:60] = True
+ orig = C._COMPONENT_ROUTE_PX
+ try:
+ C._COMPONENT_ROUTE_PX = 0
+ lab = split_instances(fg, SegmentParams(min_size=5, watershed=True))
+ finally:
+ C._COMPONENT_ROUTE_PX = orig
+ assert lab.max() == 3, f"expected 3 instances, got {lab.max()}"
+ assert len(set(np.unique(lab)) - {0}) == 3
diff --git a/spyde/tests/migrated/test_particles_gpu.py b/spyde/tests/migrated/test_particles_gpu.py
new file mode 100644
index 00000000..f516cd3f
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_gpu.py
@@ -0,0 +1,180 @@
+"""
+The scribble engine on a REAL accelerator: does the GPU agree with the CPU?
+
+Everything else about this engine is tested on the CPU (``select_device("cpu")``
+throughout ``test_particles_scribble.py``) because torch-CUDA work segfaults
+under the pytest process on Windows — a harness interaction, not a code defect
+(CLAUDE.md § GPU Computing). That leaves exactly one thing unchecked, and it is
+the thing that matters most for a path that ships GPU-first: **that the device
+and the CPU produce the same particles.**
+
+So this file runs the comparison in a **subprocess** that prints a JSON summary
+and hard-exits, following ``test_vector_orientation_gpu.py``. Skipped entirely
+when no GPU is present, which is CI and many user machines.
+
+What is and is not asserted, and why
+------------------------------------
+Foreground agreement is held to a tight IoU rather than bit-equality. The two
+devices run different kernels — cuDNN's convolutions accumulate in a different
+order from the CPU's — so the probability maps differ in the last few bits, and
+a threshold at 0.5 turns that into a handful of disagreeing pixels. Demanding
+bit-equality here would be demanding that two float32 implementations coincide,
+which is not a property either device offers.
+
+The **count** is held exactly, because that is the number a user acts on. The
+boundary route is the more fragile of the two here and deliberately so: it takes
+connected components of ``fg & ~boundary``, so a single pixel flipping along a
+seam can join or separate two bodies, where the watershed would have absorbed it.
+If that ever becomes flaky the honest response is to record the sensitivity, not
+to loosen the count assertion.
+"""
+import json
+import subprocess
+import sys
+import textwrap
+
+import pytest
+
+from spyde.particles.features import gpu_available
+
+pytestmark = pytest.mark.skipif(
+ not gpu_available(), reason="no torch GPU (CUDA / MPS) available")
+
+
+_DRIVER = textwrap.dedent("""
+ import json, sys, os
+ import numpy as np
+ from scipy import ndimage as ndi
+
+ from spyde.data.synthetic import (particle_movie, ground_truth,
+ particle_truth_at)
+ from spyde.particles import SegmentParams, split_instances
+ from spyde.particles.features import FeatureSpec, select_device
+ from spyde.particles.scribble import (LabelStore, ScribbleClassifier,
+ default_classes)
+
+ T = 12
+ s = particle_movie()
+ gt = ground_truth(s)
+ frame = np.asarray(s.data[T])
+ pos, radii, present = particle_truth_at(gt, T)
+ faint = np.asarray(gt["p_faint"], bool)
+ h, w = frame.shape
+ yy, xx = np.mgrid[0:h, 0:w]
+
+ # Scribbles: dabs on particles (incl. one faint probe), background sweeps,
+ # and seam strokes along the joins between touching bodies.
+ store = LabelStore(frame_shape=(h, w), classes=default_classes())
+ idx = list(np.flatnonzero(present))
+ discs = [((yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2) <= radii[i] ** 2
+ for i in idx]
+ for i in idx:
+ store.paint_disc(T, pos[i, 0], pos[i, 1], max(1.5, radii[i] * 0.5), 0)
+ far = np.ones((h, w), bool)
+ for i in idx:
+ far &= ((yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2) > (radii[i] + 3.) ** 2
+ store.paint(T, far, 1)
+ seam = np.zeros((h, w), bool)
+ grown = [ndi.binary_dilation(d, iterations=2) for d in discs]
+ for a in range(len(grown)):
+ for b in range(a + 1, len(grown)):
+ seam |= grown[a] & grown[b]
+ if seam.any():
+ store.paint(T, seam, 3)
+
+ def run(dev):
+ clf = ScribbleClassifier(FeatureSpec(), device=select_device(dev), seed=0)
+ rep = clf.fit(store, {T: frame})
+ fg, bnd = clf.predict_foreground_boundary(frame)
+ p = SegmentParams(min_size=5)
+ return dict(
+ device=str(clf.device),
+ has_boundary=bool(rep["has_boundary"]),
+ fg=(fg > 0.5),
+ bnd=(None if bnd is None else bnd > 0.5),
+ by_boundary=split_instances(fg, p, boundary=bnd),
+ by_watershed=split_instances(fg, p),
+ )
+
+ g = run(None) # auto-selects CUDA / MPS
+ c = run("cpu")
+ assert g["device"] != "cpu", "no accelerator was selected"
+
+ def iou(a, b):
+ u = int((a | b).sum())
+ return 1.0 if u == 0 else int((a & b).sum()) / u
+
+ def areas(lab):
+ cnt = np.bincount(lab.ravel())[1:]
+ return np.sort(cnt[cnt > 0]).tolist()
+
+ def found(lab, sel):
+ return [bool(lab[int(round(pos[i, 0])), int(round(pos[i, 1]))])
+ for i in np.flatnonzero(sel)]
+
+ out = dict(
+ device=g["device"],
+ has_boundary=g["has_boundary"] and c["has_boundary"],
+ fg_iou=iou(g["fg"], c["fg"]),
+ bnd_iou=iou(g["bnd"], c["bnd"]),
+ n_boundary_gpu=int(g["by_boundary"].max()),
+ n_boundary_cpu=int(c["by_boundary"].max()),
+ n_watershed_gpu=int(g["by_watershed"].max()),
+ n_watershed_cpu=int(c["by_watershed"].max()),
+ areas_boundary_gpu=areas(g["by_boundary"]),
+ areas_boundary_cpu=areas(c["by_boundary"]),
+ faint_gpu=found(g["by_boundary"], faint),
+ faint_cpu=found(c["by_boundary"], faint),
+ n_faint=int(faint.sum()),
+ )
+ print("RESULT_JSON", json.dumps(out))
+ sys.stdout.flush()
+ # torch + CUDA teardown segfaults at interpreter exit on Windows (harmless,
+ # and after the result is out). Hard-exit so the parent sees rc == 0.
+ os._exit(0)
+""")
+
+
+@pytest.fixture(scope="module")
+def result():
+ """One subprocess for the whole module — it trains twice and is ~30 s."""
+ proc = subprocess.run([sys.executable, "-c", _DRIVER],
+ capture_output=True, text=True, timeout=900)
+ assert proc.returncode == 0, (
+ f"subprocess failed ({proc.returncode}):\n{proc.stdout}\n{proc.stderr}")
+ line = next(l for l in proc.stdout.splitlines()
+ if l.startswith("RESULT_JSON"))
+ return json.loads(line[len("RESULT_JSON "):])
+
+
+class TestGpuCpuAgreement:
+ def test_an_accelerator_was_actually_used(self, result):
+ """Guards the whole file from passing vacuously by running CPU twice."""
+ assert result["device"] in ("cuda", "mps")
+
+ def test_the_boundary_class_trained_on_both(self, result):
+ assert result["has_boundary"]
+
+ def test_the_foreground_maps_agree(self, result):
+ assert result["fg_iou"] > 0.99, (
+ f"GPU/CPU foreground IoU {result['fg_iou']:.4f}")
+
+ def test_the_boundary_maps_agree(self, result):
+ assert result["bnd_iou"] > 0.95, (
+ f"GPU/CPU boundary IoU {result['bnd_iou']:.4f}")
+
+ def test_the_particle_count_is_the_same_on_both(self, result):
+ """The number the user acts on. Checked on both routes, because the
+ boundary route is the more fragile one — see the module docstring."""
+ assert result["n_boundary_gpu"] == result["n_boundary_cpu"], result
+ assert result["n_watershed_gpu"] == result["n_watershed_cpu"], result
+
+ def test_the_areas_agree(self, result):
+ assert (result["areas_boundary_gpu"]
+ == result["areas_boundary_cpu"]), result
+
+ def test_the_faint_probes_are_found_on_both(self, result):
+ """§0.9 on the device: the GPU path may not quietly lose a faint
+ particle the CPU path finds."""
+ assert all(result["faint_gpu"]), result
+ assert result["faint_gpu"] == result["faint_cpu"], result
diff --git a/spyde/tests/migrated/test_particles_props_parity.py b/spyde/tests/migrated/test_particles_props_parity.py
new file mode 100644
index 00000000..1eafc48e
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_props_parity.py
@@ -0,0 +1,259 @@
+"""
+Per-column parity between the vectorised property path and ``regionprops_table``.
+
+:mod:`spyde.particles.props` and :mod:`spyde.particles.hull` replaced skimage's
+``regionprops_table`` inside :func:`spyde.particles.measure.measure_frame` for one
+reason — it costs 53.5 s on a real 4096² frame against 3.0 s to segment it, because
+its cost is per REGION and a real frame has 26 566 of them (``benchmarks.md``).
+
+**A faster measurement that quietly changes what a particle's properties ARE is a
+regression, not an optimisation.** So the replacement is only defensible if every
+column is checked against the implementation it replaced, on a raster with
+THOUSANDS of irregular, concave, touching, edge-clipped regions — not on three
+discs, which agree under any implementation and prove nothing.
+
+That is what this module is: one scene of ~5 000 blobs from thresholded noise (so
+the shapes are ragged and concave, and plenty of them run off the frame edge), and
+one assertion per column. The tolerances are not round numbers picked to make it
+pass — they record which columns are integer-exact, which are float-identical
+because both paths sum exact integers, and which differ only in summation order:
+
+========================= =========================================
+column agreement
+========================= =========================================
+label, area, bbox-* exact (integers)
+centroid-0, centroid-1 exact (both sum exact integers in float64)
+equivalent_diameter_area exact (a function of ``area`` alone)
+solidity exact (the hull is integer arithmetic)
+perimeter ~1e-15 relative (same weights, summed in a
+ different order)
+major/minor_axis_length ~1e-14 relative (same 2x2 LAPACK
+eccentricity eigenproblem, moments summed in a
+ different order)
+========================= =========================================
+
+The same scene is used to pin ``measure_frame``'s own output rows, because the
+columns feed calibration and circularity and a per-column check alone would not
+catch a mis-wiring between the two.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.particles.measure import measure_frame, property_table
+from spyde.signals.particles import COL, COLUMNS, N_COLUMNS
+
+#: Columns that must agree to the BIT, and why (see the module docstring).
+_EXACT = ("label", "area", "bbox-0", "bbox-1", "bbox-2", "bbox-3",
+ "centroid-0", "centroid-1", "equivalent_diameter_area", "solidity")
+
+#: Columns that differ only in floating-point summation order. Relative, and
+#: three orders of magnitude tighter than any measurement this feeds.
+_TOL = {"perimeter": 1e-12,
+ "major_axis_length": 1e-11,
+ "minor_axis_length": 1e-11,
+ "eccentricity": 1e-9}
+
+
+def _blob_field(size=1024, n_particles=4000, r_max=3, seed=0):
+ """Thousands of ragged, concave, touching, edge-clipped regions.
+
+ Each particle is the union of one to three overlapping discs, which is what
+ makes the scene worth testing on: a single disc is CONVEX, so its solidity is
+ ~1 and its hull is uninteresting, and a field of discs would let a wrong
+ convex-hull implementation pass. Lobed unions are concave, they agglomerate
+ where they overlap, and some run off the frame edge — the three places where a
+ per-region crop, an erosion border value and the hull's ±0.5 offsets can each
+ be wrong on their own.
+
+ The defaults are tuned to the REAL frame this replaces ``regionprops_table``
+ for: ~3 260 regions of mean area 30 px (real: 26 566 of mean 33 px), solidity
+ spanning ~0.60-1.0 (real: 0.28-0.92).
+ """
+ from scipy import ndimage as ndi
+
+ rng = np.random.default_rng(seed)
+ canvas = np.zeros((size, size), bool)
+
+ def _disc(r):
+ y, x = np.mgrid[-r:r + 1, -r:r + 1]
+ return y * y + x * x <= r * r
+
+ discs = {r: _disc(r) for r in range(1, r_max + 1)}
+ for _ in range(n_particles):
+ cy = int(rng.integers(-2, size + 2))
+ cx = int(rng.integers(-2, size + 2))
+ for _lobe in range(int(rng.integers(1, 4))):
+ r = int(rng.integers(1, r_max + 1))
+ d = discs[r]
+ y0, x0 = cy + int(rng.integers(-r, r + 1)) - r, \
+ cx + int(rng.integers(-r, r + 1)) - r
+ y1, x1 = y0 + 2 * r + 1, x0 + 2 * r + 1
+ sy0, sx0, sy1, sx1 = max(0, y0), max(0, x0), min(size, y1), min(size, x1)
+ if sy0 >= sy1 or sx0 >= sx1:
+ continue
+ canvas[sy0:sy1, sx0:sx1] |= d[sy0 - y0:sy1 - y0, sx0 - x0:sx1 - x0]
+
+ lab, n = ndi.label(canvas)
+ assert n > 2000, f"scene has only {n} regions; the parity check needs thousands"
+ img = canvas.astype(np.float32) + 0.1 * rng.standard_normal(canvas.shape).astype(np.float32)
+ return lab.astype(np.int32), img
+
+
+@pytest.fixture(scope="module")
+def scene():
+ return _blob_field()
+
+
+@pytest.fixture(scope="module")
+def tables(scene):
+ """Both property tables, computed ONCE — the legacy one is seconds even here."""
+ lab, _img = scene
+ return property_table(lab, fast=False), property_table(lab, fast=True)
+
+
+class TestPropertyTableParity:
+ def test_scene_is_actually_hard(self, scene, tables):
+ """Guard the guard: a scene of near-convex blobs would not test the hull."""
+ lab, _img = scene
+ tbl = tables[0]
+ assert len(tbl["label"]) > 2000
+ # Genuinely concave shapes present, not a field of discs.
+ assert tbl["solidity"].min() < 0.7
+ assert np.median(tbl["solidity"]) < 0.95
+ # Regions clipped by the frame edge.
+ assert (tbl["bbox-0"] == 0).any() and (tbl["bbox-2"] == lab.shape[0]).any()
+
+ def test_every_column_matches_regionprops(self, tables):
+ ref, got = tables
+
+ assert set(got) == set(ref)
+ for key in sorted(ref):
+ a = np.asarray(ref[key], np.float64)
+ b = np.asarray(got[key], np.float64)
+ assert a.shape == b.shape, key
+ if key in _EXACT:
+ assert np.array_equal(a, b), (
+ f"{key} must match regionprops exactly; "
+ f"max |diff| = {np.abs(a - b).max()}")
+ else:
+ tol = _TOL[key]
+ rel = np.abs(a - b) / np.maximum(np.abs(a), 1e-300)
+ assert rel.max() < tol, (
+ f"{key} differs by {rel.max():.3e} relative (tolerance "
+ f"{tol:.0e}) — that is more than summation order")
+
+ def test_solidity_is_the_same_convex_hull(self, scene):
+ """``area_convex`` is a pixel count, so 'close' is not the bar — the hull
+ must select the SAME pixels. This is the column that was rewritten from
+ Qhull to integer arithmetic, so it gets its own assertion."""
+ from spyde.particles.hull import convex_areas
+ from skimage.measure import regionprops_table
+
+ lab, _img = scene
+ counts = np.bincount(lab.reshape(-1))
+ labels = (np.flatnonzero(counts[1:] > 0) + 1).astype(np.int64)
+ got = convex_areas(lab, labels, counts)
+ if got is None:
+ pytest.skip("numba unavailable; solidity falls back to regionprops")
+ ref = regionprops_table(lab, properties=("area_convex",))["area_convex"]
+ assert np.array_equal(got, ref.astype(np.int64))
+
+
+class TestSparseLabels:
+ """A label image whose values are not 1..N.
+
+ ``regionprops_table`` skips absent labels, so the vectorised path has to as
+ well — and the perimeter's per-label histogram has FIFTY bins per label, so
+ keying it by the raw label value would ask for 50x the label range in memory.
+ Both are pinned here rather than left to a frame that happens to be dense.
+ """
+
+ def test_matches_regionprops_with_gaps(self, scene):
+ lab, _img = scene
+ sparse = np.where(lab > 0, lab.astype(np.int64) * 3 + 7, 0).astype(np.int32)
+ ref = property_table(sparse, fast=False)
+ got = property_table(sparse, fast=True)
+ assert len(ref["label"]) == len(got["label"]) > 2000
+ assert np.array_equal(ref["label"], got["label"])
+ assert np.array_equal(ref["area"], got["area"])
+ assert np.array_equal(ref["solidity"], got["solidity"])
+ assert np.allclose(ref["perimeter"], got["perimeter"], rtol=1e-12, atol=0)
+
+
+class TestMeasureFrameParity:
+ def test_rows_match_between_paths(self, scene):
+ lab, img = scene
+ rows_l, cont_l = measure_frame(lab, img, t=3, scale=0.25, fast=False)
+ rows_f, cont_f = measure_frame(lab, img, t=3, scale=0.25, fast=True)
+ assert rows_l.shape == rows_f.shape
+
+ loose = {"major_axis", "minor_axis", "eccentricity", "perimeter",
+ "circularity"}
+ for i, name in enumerate(COLUMNS[:N_COLUMNS]):
+ a, b = rows_l[:, i], rows_f[:, i]
+ assert np.array_equal(np.isnan(a), np.isnan(b)), name
+ fin = ~np.isnan(a)
+ if name in loose:
+ # float32 rows, so the comparison is at float32 resolution.
+ assert np.allclose(a[fin], b[fin], rtol=1e-6, atol=0), name
+ else:
+ assert np.array_equal(a[fin], b[fin]), name
+
+ # Outlines are compared by the FILLED polygon, not by vertex identity —
+ # the vectorised tracer cuts a closed contour at a different vertex, and
+ # everything downstream (`render_frame`, `mask_at`) fills it.
+ # `test_particles_contours_parity.py` is where that gate lives; this is
+ # the wiring check that `measure_frame` returns one per kept row.
+ from spyde.tests.migrated.test_particles_contours_parity import (
+ filled_pixels)
+
+ assert len(cont_l) == len(cont_f) == rows_f.shape[0]
+ for x, y in zip(cont_l, cont_f):
+ assert np.array_equal(filled_pixels(x, lab.shape),
+ filled_pixels(y, lab.shape))
+
+ def test_min_area_and_empty_frame(self, scene):
+ lab, img = scene
+ rows_l, _ = measure_frame(lab, img, min_area_px=25, fast=False)
+ rows_f, _ = measure_frame(lab, img, min_area_px=25, fast=True)
+ assert rows_l.shape == rows_f.shape and rows_f.shape[0] > 100
+ assert np.array_equal(rows_l[:, COL["area"]], rows_f[:, COL["area"]])
+
+ empty = np.zeros((16, 16), np.int32)
+ for fast in (False, True):
+ rows, cont = measure_frame(empty, fast=fast)
+ assert rows.shape == (0, N_COLUMNS) and cont == []
+
+
+class TestHullEdgeCases:
+ """Shapes whose hull is degenerate or whose pixels touch the frame edge —
+ the cases where an integer reimplementation of Qhull is most likely to
+ disagree, and where ``regionprops`` itself is at its least obvious."""
+
+ @pytest.mark.parametrize("build", [
+ pytest.param(lambda a: a.__setitem__((5, 5), 1), id="single-pixel"),
+ pytest.param(lambda a: a.__setitem__((5, slice(2, 9)), 1), id="h-line"),
+ pytest.param(lambda a: a.__setitem__((slice(2, 9), 5), 1), id="v-line"),
+ pytest.param(lambda a: a.__setitem__((0, 0), 1), id="corner"),
+ pytest.param(lambda a: a.__setitem__((slice(0, 3), slice(0, 3)), 1),
+ id="corner-block"),
+ pytest.param(lambda a: [a.__setitem__((i, i), 1) for i in range(2, 9)],
+ id="diagonal"),
+ pytest.param(lambda a: [a.__setitem__((slice(2, 9), 2), 1),
+ a.__setitem__((2, slice(2, 9)), 1)], id="L"),
+ ])
+ def test_matches_regionprops(self, build):
+ from skimage.measure import regionprops_table
+ from spyde.particles.hull import convex_areas
+
+ lab = np.zeros((12, 12), np.int32)
+ build(lab)
+ counts = np.bincount(lab.reshape(-1))
+ labels = (np.flatnonzero(counts[1:] > 0) + 1).astype(np.int64)
+ got = convex_areas(lab, labels, counts)
+ if got is None:
+ pytest.skip("numba unavailable")
+ ref = regionprops_table(lab, properties=("area_convex",))["area_convex"]
+ assert np.array_equal(got, ref.astype(np.int64))
diff --git a/spyde/tests/migrated/test_particles_scribble.py b/spyde/tests/migrated/test_particles_scribble.py
new file mode 100644
index 00000000..008734a3
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_scribble.py
@@ -0,0 +1,1807 @@
+"""
+The torch feature stack and the scribble classifier — DRIFT_AND_PARTICLES_PLAN.md
+step 3 (B2 + B3), and the acceptance gates it exists to pass.
+
+Four of these classes are the plan's gates rather than ordinary unit tests, and
+they are the reason this file is worth its runtime:
+
+:class:`TestRandomForestParity`
+ Plan B3: the torch head must agree with ``sklearn``'s RandomForest — what
+ ParticleSpy and ilastik actually use — on **identical labels and identical
+ feature channels**. Agreement by IoU, measured 0.94.
+:class:`TestSensitivityGate`
+ Plan §0.9, the headline gate for the whole feature: the two deliberately
+ faint low-contrast probes in ``particle_movie()`` must be found. It is kept
+ non-vacuous by :class:`TestTheClassicalBaselineMissesThem`, which shows the
+ classical engine finds neither.
+:class:`TestNaNBorder`
+ Plan trap 2 / A7: no foreground inside a drift-corrected frame's NaN padding.
+:class:`TestInteractionBudget`
+ Plan B3's hard budget: train + apply to one visible frame under ~1 s on CPU.
+
+Everything runs on **CPU explicitly** (``select_device("cpu")``): torch-CUDA work
+segfaults under the pytest process on Windows (CLAUDE.md), which is a harness
+interaction and not a code defect, so the GPU path is left to the real app and to
+a subprocess check.
+"""
+from __future__ import annotations
+
+import json
+import threading
+import time
+
+import numpy as np
+import pytest
+
+from spyde.data.synthetic import ground_truth, particle_movie, particle_truth_at
+from spyde.drift.warp import shift_frame
+from spyde.particles import SegmentParams, segment_frame, split_instances
+from spyde.particles import features as feat
+from spyde.particles.features import (
+ DEFAULT_SIGMAS,
+ FeatureSpec,
+ PreparedFrame,
+ band_rows_for,
+ feature_stack,
+ feature_tensor,
+ map_feature_bands,
+ prepare_frame,
+ sample_features,
+ select_device,
+)
+from spyde.particles.scribble import (
+ UNLABELLED,
+ LabelStore,
+ ScribbleClass,
+ ScribbleClassifier,
+ default_classes,
+ masks_to_labels,
+ random_forest_reference,
+)
+
+#: The frame every gate is measured on: all nine particles are present at t=12.
+FRAME_T = 12
+
+DEVICE = select_device("cpu")
+
+
+# ── fixtures ─────────────────────────────────────────────────────────────────
+
+@pytest.fixture(scope="module")
+def movie():
+ """One build shared across the module — ~1 s and deterministic."""
+ s = particle_movie()
+ return s, ground_truth(s)
+
+
+@pytest.fixture(scope="module")
+def geom(movie):
+ """``(positions, radii, present, faint, shape)`` at :data:`FRAME_T`."""
+ _s, gt = movie
+ pos, radii, present = particle_truth_at(gt, FRAME_T)
+ faint = np.asarray(gt["p_faint"], bool)
+ return pos, radii, present, faint, tuple(gt["frame_shape"])
+
+
+def _clear_of_particles(shape, pos, radii, present, pad: float):
+ """True where no present particle is within *pad* px — where background may
+ be painted without accidentally labelling a particle as film."""
+ h, w = shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ keep = np.ones((h, w), bool)
+ for i in np.flatnonzero(present):
+ keep &= ((yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2) > (radii[i] + pad) ** 2
+ return keep
+
+
+def paint_scribbles(geom, *, include_faint=(8,), n_bright=4,
+ t: int = FRAME_T) -> LabelStore:
+ """The scribbles a user would actually paint on one frame.
+
+ Four dabs on bright particles, a background stroke at each of those
+ particles' boundaries, four background sweeps across the film, and — unless
+ *include_faint* is emptied — one dab on the SMALLER of the two faint probes
+ (index 8, r=3). Index 7 (r=4) is never painted, so it is a genuinely held-out
+ detection in :class:`TestSensitivityGate`.
+
+ The boundary strokes matter and are not padding: without them the head has
+ never seen a not-quite-particle pixel, its masks come out visibly fat, and
+ agreement with the RandomForest reference falls from 0.94 to 0.69 (measured).
+ That is a fact about labelling, not about either classifier.
+ """
+ pos, radii, present, faint, shape = geom
+ h, w = shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ store = LabelStore(frame_shape=shape, classes=default_classes())
+
+ bright = [i for i in np.flatnonzero(present) if not faint[i]][:n_bright]
+ for i in bright:
+ store.paint_disc(t, pos[i, 0], pos[i, 1], max(1.5, radii[i] * 0.5), 0)
+ for i in include_faint:
+ store.paint_disc(t, pos[i, 0], pos[i, 1], 1.5, 0)
+
+ far = _clear_of_particles(shape, pos, radii, present, 3.0)
+ for (y0, x0, y1, x1) in ((4, 4, 8, 100), (h - 8, 4, h - 4, 100),
+ (40, 4, 60, 8), (20, 40, 40, 46)):
+ sweep = np.zeros((h, w), bool)
+ sweep[y0:y1, x0:x1] = True
+ store.paint(t, sweep & far, 1)
+
+ touching = _clear_of_particles(shape, pos, radii, present, 0.0)
+ for i in bright:
+ d2 = (yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2
+ ring = (d2 > (radii[i] + 2.5) ** 2) & (d2 <= (radii[i] + 4.0) ** 2)
+ store.paint(t, ring & touching, 1)
+ return store
+
+
+@pytest.fixture(scope="module")
+def labels(geom):
+ return paint_scribbles(geom)
+
+
+@pytest.fixture(scope="module")
+def trained(movie, labels):
+ s, _gt = movie
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(labels, {FRAME_T: s.data[FRAME_T]})
+ return clf
+
+
+@pytest.fixture(scope="module")
+def proba(movie, trained):
+ s, _gt = movie
+ return trained.predict_proba(s.data[FRAME_T])
+
+
+def hit(prob_or_labels, pos, i) -> bool:
+ """Did the map fire at particle *i*'s ground-truth centre?"""
+ cy, cx = int(round(pos[i, 0])), int(round(pos[i, 1]))
+ v = prob_or_labels[cy, cx]
+ return bool(v > 0.5) if np.issubdtype(np.asarray(v).dtype, np.floating) \
+ else bool(v != 0)
+
+
+# ── FeatureSpec ──────────────────────────────────────────────────────────────
+
+class TestFeatureSpec:
+ def test_the_scale_floor_is_fine(self):
+ """Plan §0.9. What raising this costs is measured in
+ :meth:`TestSensitivityGate.test_a_coarse_stack_undersizes_the_small_particles`
+ — this only pins that it has not been raised without going and looking."""
+ assert min(DEFAULT_SIGMAS) <= 1.0, (
+ "the default scale floor has been raised above 1 px; that costs the "
+ "smallest particles' measured radius (see the features module "
+ "docstring) and needs a fresh sensitivity measurement, not a guess")
+ assert sorted(DEFAULT_SIGMAS) == list(DEFAULT_SIGMAS)
+
+ def test_default_channel_count_and_names(self):
+ spec = FeatureSpec()
+ names = spec.channel_names()
+ assert len(names) == spec.n_channels == 36
+ assert len(set(names)) == len(names), "duplicate channel names"
+ # Every promised family is present.
+ for token in ("intensity", "gaussian_s", "dog_s", "sobel_s", "laplacian_s",
+ "hessian_major_s", "hessian_minor_s", "median_r",
+ "minimum_r", "maximum_r"):
+ assert any(n.startswith(token) for n in names), f"missing {token}"
+
+ def test_names_match_the_stack_for_every_configuration(self, movie):
+ """The names and the tensor come from one plan; this is what stops them
+ drifting apart when a family is added."""
+ s, _gt = movie
+ frame = s.data[0][:32, :32]
+ for spec in (FeatureSpec(),
+ FeatureSpec(membrane=True),
+ FeatureSpec(sigmas=(1.0,), median=False, minimum=False,
+ maximum=False, membrane=True),
+ FeatureSpec(intensity=False, gaussian=False, sobel=False,
+ hessian=False, laplacian=False,
+ difference_of_gaussians=False)):
+ stack = feature_stack(frame, spec, device=DEVICE)
+ assert stack.shape[0] == len(spec.channel_names()) == spec.n_channels
+
+ def test_round_trips_through_a_dict(self):
+ spec = FeatureSpec(sigmas=(0.7, 1.4), rank_radii=(3,), membrane=True,
+ membrane_projections=("sum", "min"))
+ d = spec.to_dict()
+ assert json.loads(json.dumps(d)) == d, "not JSON-safe"
+ assert FeatureSpec.from_dict(d) == spec
+
+ def test_from_dict_ignores_unknown_keys(self):
+ """A recipe written by a newer build must still open here."""
+ d = FeatureSpec().to_dict()
+ d["some_future_feature"] = True
+ assert FeatureSpec.from_dict(d) == FeatureSpec()
+
+ def test_from_dict_of_nothing_is_the_default(self):
+ assert FeatureSpec.from_dict(None) == FeatureSpec()
+ assert FeatureSpec.from_dict({}) == FeatureSpec()
+
+ def test_replace_and_hashable(self):
+ spec = FeatureSpec()
+ assert spec.replace(median=False).median is False
+ assert spec.median is True, "replace mutated the original"
+ assert hash(spec) == hash(FeatureSpec())
+
+ @pytest.mark.parametrize("kw, match", [
+ ({"sigmas": (2.0, 1.0)}, "ascending"),
+ ({"sigmas": (0.0, 1.0)}, "positive"),
+ ({"rank_radii": (0,)}, "rank_radii"),
+ ({"membrane_projections": ("nope",)}, "membrane projection"),
+ ({"membrane": True, "membrane_patch": 18}, "odd"),
+ ({"membrane": True, "membrane_rotations": 0}, "rotations"),
+ ])
+ def test_validation(self, kw, match):
+ with pytest.raises(ValueError, match=match):
+ FeatureSpec(**kw)
+
+ def test_a_spec_with_no_channels_at_all_raises(self):
+ with pytest.raises(ValueError, match="no channels"):
+ FeatureSpec(intensity=False, gaussian=False,
+ difference_of_gaussians=False, sobel=False, hessian=False,
+ laplacian=False, median=False, minimum=False,
+ maximum=False)
+
+ def test_halo_covers_the_largest_filter(self):
+ spec = FeatureSpec()
+ # sigma 8 truncated at 4 sigma is radius 32, plus the 3-tap derivative.
+ assert spec.halo == 33
+ assert FeatureSpec(sigmas=(1.0,), rank_radii=(7,)).halo == 7
+ assert FeatureSpec(sigmas=(1.0,), rank_radii=(1,), membrane=True,
+ membrane_patch=19).halo == 9
+
+
+# ── the feature stack ────────────────────────────────────────────────────────
+
+class TestFeatureParity:
+ """Every channel against the scipy/skimage filter it re-implements.
+
+ ``normalize_frame=False`` throughout, because the whole point is comparing the
+ *filters*, not the standardisation. Padding is compared with scipy's
+ ``mirror``, which is what torch's ``reflect`` is (scipy's own ``reflect``
+ duplicates the edge sample and is a different thing).
+ """
+
+ @pytest.fixture(scope="class")
+ def img(self, movie):
+ s, _gt = movie
+ return np.ascontiguousarray(s.data[FRAME_T], dtype=np.float32)
+
+ @pytest.fixture(scope="class")
+ def stack(self, img):
+ spec = FeatureSpec(normalize_frame=False, membrane=True)
+ arr = feature_stack(img, spec, device=DEVICE)
+ return spec, arr, spec.channel_names()
+
+ def test_shape_and_dtype(self, img, stack):
+ spec, arr, _names = stack
+ assert arr.shape == (spec.n_channels, *img.shape)
+ assert arr.dtype == np.float32
+ assert np.isfinite(arr).all()
+
+ def test_intensity_channel_is_the_frame(self, img, stack):
+ _spec, arr, names = stack
+ assert np.array_equal(arr[names.index("intensity")], img)
+
+ @pytest.mark.parametrize("sigma", DEFAULT_SIGMAS)
+ def test_gaussian_matches_scipy(self, img, stack, sigma):
+ from scipy.ndimage import gaussian_filter
+ _spec, arr, names = stack
+ ref = gaussian_filter(img, float(sigma), mode="mirror")
+ got = arr[names.index(f"gaussian_s{sigma:g}")]
+ assert np.abs(got - ref).max() < 1e-5
+
+ def test_difference_of_gaussians_is_the_difference(self, stack):
+ _spec, arr, names = stack
+ got = arr[names.index("dog_s1_s2")]
+ ref = arr[names.index("gaussian_s1")] - arr[names.index("gaussian_s2")]
+ assert np.array_equal(got, ref)
+
+ @pytest.mark.parametrize("radius", (1, 2))
+ def test_rank_filters_match_scipy(self, img, stack, radius):
+ from scipy.ndimage import maximum_filter, median_filter, minimum_filter
+ _spec, arr, names = stack
+ k = 2 * radius + 1
+ for stat, fn in (("median", median_filter), ("minimum", minimum_filter),
+ ("maximum", maximum_filter)):
+ ref = fn(img, size=k, mode="mirror")
+ got = arr[names.index(f"{stat}_r{radius}")]
+ assert np.array_equal(got, ref), f"{stat} r={radius}"
+
+ def test_sobel_matches_skimage(self, img, stack):
+ from scipy.ndimage import gaussian_filter
+ from skimage.filters import sobel
+ _spec, arr, names = stack
+ ref = sobel(gaussian_filter(img, 1.0, mode="mirror"), mode="mirror")
+ assert np.abs(arr[names.index("sobel_s1")] - ref).max() < 1e-5
+
+ def test_laplacian_is_the_hessian_trace(self, stack):
+ """Not decoration: the Laplacian channel is *derived* from the Hessian's
+ own second derivatives, which is what makes it free."""
+ _spec, arr, names = stack
+ lap = arr[names.index("laplacian_s2")]
+ tr = (arr[names.index("hessian_major_s2")]
+ + arr[names.index("hessian_minor_s2")])
+ assert np.abs(lap - tr).max() < 1e-4
+
+ def test_hessian_eigenvalues_are_ordered_and_signed(self, stack):
+ _spec, arr, names = stack
+ major = arr[names.index("hessian_major_s2")]
+ minor = arr[names.index("hessian_minor_s2")]
+ assert (major >= minor - 1e-6).all(), "major must be the larger SIGNED value"
+ # A bright blob is a maximum, so BOTH curvatures are negative at its centre.
+ assert minor.min() < 0 < major.max()
+
+ @pytest.mark.parametrize("radius, matched_sigma", [(1.5, 1.0), (3.0, 2.0),
+ (8.0, 4.0)])
+ def test_hessian_picks_a_blob_out_at_its_own_scale(self, radius,
+ matched_sigma):
+ """Scale space, and the whole reason the sigma set spans a range.
+
+ A disc of radius *r* registers as a curvature minimum most strongly at
+ sigma ≈ ``r/sqrt(2)``, so which scale fires *is* the size measurement the
+ head has available. A stack whose floor is too high sees a small particle
+ only in its coarse channels, which is why it then measures it too big —
+ see :meth:`TestSensitivityGate.
+ test_a_coarse_stack_undersizes_the_small_particles`.
+
+ Checked at the centre of a synthetic disc rather than on the fixture: a
+ fixture particle of radius 7 has a *flat* centre at sigma 2, so the plain
+ "is it a curvature minimum at the centre" assertion is scale-dependent and
+ would be measuring the wrong thing.
+ """
+ n = 64
+ yy, xx = np.mgrid[0:n, 0:n]
+ img = (((yy - 32) ** 2 + (xx - 32) ** 2) <= radius ** 2).astype(np.float32)
+ spec = FeatureSpec(normalize_frame=False)
+ names = spec.channel_names()
+ arr = feature_stack(img, spec, device=DEVICE)
+ centre = {s: float(arr[names.index(f"hessian_minor_s{s:g}")][32, 32])
+ for s in DEFAULT_SIGMAS}
+ assert min(centre, key=lambda s: centre[s]) == matched_sigma, centre
+ assert centre[matched_sigma] < 0, "a bright blob must be a minimum"
+
+ def test_membrane_projections_are_ordered(self, stack):
+ _spec, arr, names = stack
+ mx = arr[names.index("membrane_max")]
+ mn = arr[names.index("membrane_min")]
+ mean = arr[names.index("membrane_mean")]
+ assert (mn <= mean + 1e-6).all() and (mean <= mx + 1e-6).all()
+ assert (arr[names.index("membrane_std")] >= 0).all()
+
+ def test_membrane_responds_to_a_line_not_a_blob(self):
+ """The family's whole purpose, and why it is off by default."""
+ line = np.zeros((48, 48), np.float32)
+ line[24, 4:44] = 1.0
+ blob = np.zeros((48, 48), np.float32)
+ yy, xx = np.mgrid[0:48, 0:48]
+ blob[(yy - 24) ** 2 + (xx - 24) ** 2 <= 9] = 1.0
+ spec = FeatureSpec(normalize_frame=False, membrane=True)
+ names = spec.channel_names()
+ k = names.index("membrane_std")
+ line_std = feature_stack(line, spec, device=DEVICE)[k][24, 24]
+ blob_std = feature_stack(blob, spec, device=DEVICE)[k][24, 24]
+ assert line_std > 3 * blob_std, (
+ f"membrane std does not discriminate a line ({line_std:.4f}) from a "
+ f"blob ({blob_std:.4f})")
+
+ def test_deterministic(self, img):
+ spec = FeatureSpec()
+ a = feature_stack(img, spec, device=DEVICE)
+ b = feature_stack(img, spec, device=DEVICE)
+ assert np.array_equal(a, b)
+
+ def test_runs_on_the_cpu_with_no_gpu(self, img, monkeypatch):
+ """CI and most user machines have no accelerator."""
+ import torch
+ monkeypatch.setattr(torch.cuda, "is_available", lambda: False)
+ monkeypatch.setattr(feat, "select_device", lambda prefer=None:
+ torch.device("cpu"))
+ assert feature_stack(img[:32, :32]).shape[0] == FeatureSpec().n_channels
+
+
+class TestPrepareFrame:
+ def test_non_finite_pixels_take_the_finite_minimum(self):
+ """Matching ``classical._prepare``: the padding must read as background,
+ which is the one value that cannot classify as a particle."""
+ img = np.linspace(2.0, 6.0, 64, dtype=np.float32).reshape(8, 8)
+ img[0, :3] = np.nan
+ img[7, 7] = np.inf
+ finite_min = float(img[np.isfinite(img)].min())
+ prepared = prepare_frame(img, FeatureSpec(normalize_frame=False))
+ assert prepared.image[0, :3].tolist() == [finite_min] * 3
+ assert prepared.image[7, 7] == pytest.approx(finite_min)
+ assert np.isfinite(prepared.image).all()
+ # Not zero and not the mean: the padding must read as background, and on a
+ # frame whose values are all well above zero, zero-fill would instead read
+ # as a hole and the mean as ordinary signal.
+ assert prepared.image.min() == pytest.approx(finite_min)
+
+ def test_valid_mask_marks_exactly_the_finite_source_pixels(self):
+ img = np.ones((8, 8), np.float32)
+ img[3, 4] = np.nan
+ prepared = prepare_frame(img)
+ assert prepared.valid.sum() == 63
+ assert not prepared.valid[3, 4]
+
+ def test_robust_statistics_ignore_the_padding(self):
+ """Filling first would let a large NaN border rescale the whole frame by
+ how much of it was padding."""
+ rng = np.random.default_rng(0)
+ img = rng.normal(100.0, 5.0, (64, 64)).astype(np.float32)
+ padded = img.copy()
+ padded[:32] = np.nan
+ a = prepare_frame(img).image[32:]
+ b = prepare_frame(padded).image[32:]
+ # Same finite content in rows 32+, so the same standardisation of it.
+ assert np.abs(a - b).max() < 0.15
+
+ def test_normalisation_makes_the_stack_scale_invariant(self, movie):
+ """The property that makes a saved recipe transferable: the same sample
+ recorded on a different intensity scale must featurise the same."""
+ s, _gt = movie
+ frame = s.data[FRAME_T]
+ a = feature_stack(frame, FeatureSpec(), device=DEVICE)
+ b = feature_stack(frame * 1000.0 + 7.0, FeatureSpec(), device=DEVICE)
+ rms = np.sqrt(np.mean((a - b) ** 2))
+ assert rms < 1e-3, f"features are not scale invariant (rms {rms:.4g})"
+
+ def test_a_constant_frame_does_not_divide_by_zero(self):
+ out = prepare_frame(np.full((8, 8), 3.0, np.float32))
+ assert np.isfinite(out.image).all()
+
+ @pytest.mark.parametrize("bad, match", [
+ (np.zeros((4, 4, 4), np.float32), "must be 2-D"),
+ (np.zeros((3, 8), np.float32), "at least 4x4"),
+ (np.full((8, 8), np.nan, np.float32), "no finite pixels"),
+ ])
+ def test_errors(self, bad, match):
+ with pytest.raises(ValueError, match=match):
+ prepare_frame(bad)
+
+ def test_accepts_an_already_prepared_frame(self, movie):
+ s, _gt = movie
+ prepared = prepare_frame(s.data[0])
+ assert isinstance(prepared, PreparedFrame)
+ a = feature_stack(prepared, device=DEVICE)
+ b = feature_stack(s.data[0], device=DEVICE)
+ assert np.array_equal(a, b)
+
+
+class TestSmallFramesAndPadding:
+ """A kernel wider than the image must not change which kernel is applied.
+
+ ``F.pad(mode="reflect")`` refuses a pad larger than the dimension, and the
+ kernels here are not small — a sigma-8 gaussian has radius 32 and a 19x19
+ membrane patch radius 9. The first implementation clamped the *radius*, which
+ silently substituted a narrower filter and so made a ``FeatureSpec`` mean
+ different things on different frame sizes; a saved recipe would then not
+ reproduce on a crop. The padding degrades to replicate instead.
+ """
+
+ @pytest.mark.parametrize("shape", [(4, 4), (5, 9), (8, 8), (20, 17)])
+ @pytest.mark.parametrize("spec", [
+ FeatureSpec(),
+ FeatureSpec(membrane=True), # 19x19 patch on a 4x4 frame
+ FeatureSpec(rank_radii=(9,)), # 19x19 window on a 4x4 frame
+ ], ids=["default", "membrane", "wide_rank"])
+ def test_a_frame_smaller_than_the_kernels_still_works(self, shape, spec):
+ img = np.random.default_rng(0).standard_normal(shape).astype(np.float32)
+ stack = feature_stack(img, spec, device=DEVICE)
+ assert stack.shape == (spec.n_channels, *shape)
+ assert np.isfinite(stack).all()
+
+ def test_the_kernel_is_the_same_one_on_a_crop(self, movie):
+ """The property clamping broke: interior pixels must featurise identically
+ whether or not the frame around them is big enough for the kernel."""
+ s, _gt = movie
+ spec = FeatureSpec(sigmas=(4.0,), normalize_frame=False, median=False,
+ minimum=False, maximum=False)
+ big = np.ascontiguousarray(s.data[FRAME_T], dtype=np.float32)
+ whole = feature_stack(big, spec, device=DEVICE)
+ crop = feature_stack(big[20:76, 20:92], spec, device=DEVICE)
+ # 16 px in from the crop's edge is beyond the kernel's reach either way.
+ assert np.abs(whole[:, 36:60, 36:76] - crop[:, 16:40, 16:56]).max() < 1e-4
+
+ def test_the_normalisation_still_works_on_an_all_nan_border(self):
+ img = np.full((32, 32), np.nan, np.float32)
+ img[10:20, 10:20] = np.random.default_rng(1).standard_normal((10, 10))
+ prepared = prepare_frame(img)
+ assert prepared.valid.sum() == 100
+ assert np.isfinite(feature_stack(prepared, device=DEVICE)).all()
+
+
+class TestBanding:
+ """Plan §0.1: nothing may assume the frame fits in memory.
+
+ The banded stack must be IDENTICAL to the unbanded one, not merely close — a
+ halo of at least the largest filter radius replaces what reflect padding would
+ otherwise invent at a band boundary, so there is nothing to be approximate
+ about. Anything less than exact means the halo is too small.
+ """
+
+ @pytest.fixture(scope="class")
+ def img(self, movie):
+ s, _gt = movie
+ return np.ascontiguousarray(s.data[FRAME_T], dtype=np.float32)
+
+ @pytest.mark.parametrize("rows", (16, 33, 40, 500))
+ def test_banded_equals_unbanded_exactly(self, img, rows):
+ spec = FeatureSpec()
+ whole = feature_stack(img, spec, device=DEVICE)
+ got = np.zeros_like(whole)
+ seen = []
+
+ def take(y0, y1, stack):
+ seen.append((y0, y1))
+ got[:, y0:y1] = stack.detach().cpu().numpy()
+
+ map_feature_bands(img, spec, device=DEVICE, fn=take, band_rows=rows)
+ assert seen[0][0] == 0 and seen[-1][1] == img.shape[0]
+ assert [a for a, _ in seen[1:]] == [b for _, b in seen[:-1]], "gap or overlap"
+ assert np.array_equal(got, whole), (
+ f"banded stack differs at band_rows={rows}: max "
+ f"{np.abs(got - whole).max():.3g} — the halo is too small")
+
+ def test_band_rows_stays_above_the_halo(self):
+ spec = FeatureSpec()
+ # A tiny budget must still not produce a band shorter than its own halo,
+ # which would recompute more halo than payload.
+ assert band_rows_for(spec, 4096, budget_bytes=1) >= 4 * spec.halo
+ # A generous budget covers a whole 4096-row frame in one band.
+ assert band_rows_for(spec, 4096, budget_bytes=8 << 30) >= 4096
+
+ def test_sample_features_matches_the_full_stack(self, img):
+ spec = FeatureSpec()
+ whole = feature_stack(img, spec, device=DEVICE)
+ rng = np.random.default_rng(3)
+ flat = rng.choice(img.size, size=400, replace=False)
+ got = sample_features(img, flat, spec, device=DEVICE).detach().cpu().numpy()
+ ys, xs = np.divmod(flat, img.shape[1])
+ assert np.array_equal(got, whole[:, ys, xs].T)
+
+ def test_sample_features_accepts_yx_pairs(self, img):
+ spec = FeatureSpec(sigmas=(1.0,), median=False, minimum=False,
+ maximum=False)
+ yx = np.array([[5, 7], [40, 90], [0, 0]])
+ a = sample_features(img, yx, spec, device=DEVICE).detach().cpu().numpy()
+ flat = yx[:, 0] * img.shape[1] + yx[:, 1]
+ b = sample_features(img, flat, spec, device=DEVICE).detach().cpu().numpy()
+ assert np.array_equal(a, b)
+
+ def test_sample_features_rejects_an_out_of_range_index(self, img):
+ with pytest.raises(IndexError, match="outside"):
+ sample_features(img, np.array([img.size]), device=DEVICE)
+
+ def test_feature_tensor_lives_on_the_requested_device(self, img):
+ t = feature_tensor(img[:32, :32], device=DEVICE)
+ assert t.device.type == "cpu" and t.dtype.is_floating_point
+
+
+# ── label store ──────────────────────────────────────────────────────────────
+
+class TestLabelStore:
+ def test_paint_and_read_back(self):
+ store = LabelStore(frame_shape=(16, 20))
+ mask = np.zeros((16, 20), bool)
+ mask[2:5, 3:6] = True
+ assert store.paint(0, mask, 0) == 9
+ assert len(store) == 9
+ assert store.labelled_frames() == [0]
+ lm = store.label_map(0)
+ assert (lm[2:5, 3:6] == 0).all()
+ assert (lm[8:, :] == UNLABELLED).all()
+
+ def test_labels_accumulate_across_frames(self):
+ """Plan B3: paint on frame 0 and frame 400, both train one model."""
+ store = LabelStore(frame_shape=(16, 20))
+ store.paint_disc(0, 4, 4, 2, 0)
+ store.paint_disc(400, 9, 9, 2, 1)
+ assert store.labelled_frames() == [0, 400]
+ assert len(store) == len(store.at(0)[0]) + len(store.at(400)[0])
+ assert set(store.counts()) >= {0, 1}
+
+ def test_repainting_a_pixel_changes_its_class(self):
+ """Last write wins, which is what a brush does. A `np.unique` that kept
+ the FIRST occurrence would silently ignore a correction."""
+ store = LabelStore(frame_shape=(8, 8))
+ store.paint(0, [(3, 3)], 0)
+ store.paint(0, [(3, 3)], 1)
+ assert len(store) == 1
+ assert store.label_map(0)[3, 3] == 1
+
+ def test_erase_removes_rather_than_reassigns(self):
+ store = LabelStore(frame_shape=(8, 8))
+ store.paint(0, [(1, 1), (1, 2)], 0)
+ store.erase(0, [(1, 1)])
+ assert len(store) == 1
+ assert store.label_map(0)[1, 1] == UNLABELLED
+
+ def test_erasing_everything_drops_the_frame(self):
+ store = LabelStore(frame_shape=(8, 8))
+ store.paint(0, [(1, 1)], 0)
+ store.erase(0, [(1, 1)])
+ assert store.labelled_frames() == []
+
+ def test_paint_stroke_is_continuous(self):
+ """The brush widget emits one sample per pointer frame, so a fast stroke
+ jumps many pixels; dabbing only at the samples leaves a dotted line."""
+ store = LabelStore(frame_shape=(32, 32))
+ store.paint_stroke(0, [(5, 2), (5, 28)], 0, brush=1.0)
+ row = store.label_map(0)[5]
+ painted = np.flatnonzero(row != UNLABELLED)
+ assert painted.min() <= 2 and painted.max() >= 28
+ assert np.all(np.diff(painted) == 1), f"gaps in the stroke: {painted}"
+
+ def test_paint_stroke_honours_brush_width(self):
+ store = LabelStore(frame_shape=(32, 32))
+ thin = store.paint_stroke(0, [(10, 5), (10, 25)], 0, brush=1.0)
+ store.clear()
+ fat = store.paint_stroke(0, [(10, 5), (10, 25)], 0, brush=7.0)
+ assert fat > 3 * thin
+
+ def test_out_of_frame_coordinates_are_dropped_not_wrapped(self):
+ """A stroke running off the edge is normal; wrapping it would paint the
+ opposite side of the image."""
+ store = LabelStore(frame_shape=(10, 10))
+ store.paint(0, [(-3, 4), (4, 40), (5, 5)], 0)
+ assert len(store) == 1
+ assert store.label_map(0)[5, 5] == 0
+
+ def test_a_stroke_off_the_edge_paints_only_what_is_inside(self):
+ store = LabelStore(frame_shape=(16, 16))
+ store.paint_stroke(0, [(8, -6), (8, 6)], 0, brush=3.0)
+ lm = store.label_map(0)
+ # The stroke ends at x=6 with a radius-1.5 brush, so nothing past x=7.
+ assert (lm[:, 8:] == UNLABELLED).all()
+ assert (lm[8, 0:6] == 0).all(), "the in-frame part of the stroke is missing"
+
+ def test_counts_lists_classes_with_no_pixels(self):
+ """Which is how the caret shows a class is under-trained."""
+ store = LabelStore(frame_shape=(8, 8))
+ store.paint_disc(0, 4, 4, 2, 0)
+ counts = store.counts()
+ # 0 particle, 1 support film, 2 vacuum, 3 boundary.
+ assert set(counts) == {0, 1, 2, 3}
+ assert counts[1] == counts[2] == counts[3] == 0
+ assert store.n_classes_used == 1
+
+ def test_add_and_remove_classes(self):
+ store = LabelStore(frame_shape=(8, 8))
+ extra = store.add_class("beam stop", "#ff0000")
+ assert extra.id == 4 and store.class_by_id(4).name == "beam stop"
+ store.paint(0, [(1, 1)], 4)
+ store.paint(0, [(2, 2)], 0)
+ store.remove_class(4)
+ assert 4 not in store.class_ids
+ assert len(store) == 1, "removing a class must drop its pixels too"
+
+ def test_removing_the_only_labelled_class_drops_the_frame(self):
+ store = LabelStore(frame_shape=(8, 8))
+ store.paint(0, [(1, 1)], 0)
+ store.remove_class(0)
+ assert store.labelled_frames() == []
+
+ def test_unknown_class_raises(self):
+ store = LabelStore(frame_shape=(8, 8))
+ with pytest.raises(KeyError, match="no class with id"):
+ store.paint(0, [(1, 1)], 99)
+
+ def test_duplicate_class_ids_raise(self):
+ with pytest.raises(ValueError, match="duplicate class ids"):
+ LabelStore(frame_shape=(8, 8),
+ classes=[ScribbleClass(0, "a"), ScribbleClass(0, "b")])
+ store = LabelStore(frame_shape=(8, 8))
+ with pytest.raises(ValueError, match="already exists"):
+ store.add_class("dup", id=0)
+
+ def test_mask_of_the_wrong_shape_raises(self):
+ store = LabelStore(frame_shape=(8, 8))
+ with pytest.raises(ValueError, match="frame_shape"):
+ store.paint(0, np.ones((4, 4), bool), 0)
+
+ def test_round_trips_through_a_dict(self):
+ store = LabelStore(frame_shape=(16, 24))
+ store.paint_disc(0, 5, 5, 2, 0)
+ store.paint_stroke(7, [(1, 1), (1, 20)], 1, brush=2.0)
+ d = store.to_dict()
+ assert json.loads(json.dumps(d)) == d, "not JSON-safe"
+ back = LabelStore.from_dict(d)
+ assert back.frame_shape == store.frame_shape
+ assert back.labelled_frames() == store.labelled_frames()
+ assert back.counts() == store.counts()
+ for t in store.labelled_frames():
+ assert np.array_equal(back.label_map(t), store.label_map(t))
+
+ def test_clear_frame_and_clear(self):
+ store = LabelStore(frame_shape=(8, 8))
+ store.paint(0, [(1, 1)], 0)
+ store.paint(1, [(2, 2)], 0)
+ store.clear_frame(0)
+ assert store.labelled_frames() == [1]
+ store.clear()
+ assert len(store) == 0
+
+
+# ── the prompt-model bootstrap ───────────────────────────────────────────────
+
+class TestMasksToLabels:
+ """Plan §0.4: prompt masks become scribble labels with no painting at all."""
+
+ @staticmethod
+ def _disc(shape, cy, cx, r):
+ yy, xx = np.mgrid[0:shape[0], 0:shape[1]]
+ return (yy - cy) ** 2 + (xx - cx) ** 2 <= r * r
+
+ def test_interior_is_particle_and_surroundings_are_background(self):
+ m = self._disc((64, 64), 20, 20, 6)
+ store = masks_to_labels(m, t=3)
+ assert store.labelled_frames() == [3]
+ lm = store.label_map(3)
+ assert lm[20, 20] == 0, "mask centre is not the particle class"
+ assert (lm[m] != 1).all(), "background was painted inside the mask"
+ assert (lm == 1).sum() > 0, "no background ring"
+
+ def test_the_boundary_itself_is_left_unlabelled(self):
+ """Where the prompt model is least certain and where a particle's own soft
+ edge lives — labelling it either way biases every instance's size."""
+ m = self._disc((64, 64), 32, 32, 8)
+ lm = masks_to_labels(m, gap=2, erode=1).label_map(0)
+ yy, xx = np.mgrid[0:64, 0:64]
+ r = np.sqrt((yy - 32) ** 2 + (xx - 32) ** 2)
+ rim = (r > 8.0) & (r <= 9.0)
+ assert (lm[rim] == UNLABELLED).all()
+
+ def test_the_ring_never_covers_another_mask(self):
+ """Otherwise particle A is taught as background for particle B."""
+ a = self._disc((64, 64), 32, 24, 5)
+ b = self._disc((64, 64), 32, 34, 5)
+ lm = masks_to_labels(np.stack([a, b])).label_map(0)
+ assert (lm[a] != 1).all() and (lm[b] != 1).all()
+
+ def test_a_tiny_mask_survives_the_erosion(self):
+ """Plan §0.9 applied here: a 3 px particle is exactly the object this
+ feature exists for, and eroding it away would defeat the bootstrap."""
+ m = np.zeros((32, 32), bool)
+ m[16, 16] = True
+ lm = masks_to_labels(m, erode=1).label_map(0)
+ assert lm[16, 16] == 0
+
+ def test_accumulates_into_an_existing_store(self):
+ store = masks_to_labels(self._disc((64, 64), 20, 20, 5), t=0)
+ n0 = len(store)
+ masks_to_labels(self._disc((64, 64), 44, 44, 5), t=9, store=store)
+ assert store.labelled_frames() == [0, 9]
+ assert len(store) > n0
+
+ def test_a_trained_model_finds_the_prompted_particles(self, movie, geom):
+ """The end of the §0.4 handoff: masks in, dense segmentation out."""
+ s, _gt = movie
+ pos, radii, present, faint, shape = geom
+ bright = [i for i in np.flatnonzero(present) if not faint[i]][:3]
+ masks = np.stack([self._disc(shape, pos[i, 0], pos[i, 1], radii[i])
+ for i in bright])
+ store = masks_to_labels(masks, t=FRAME_T)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ prob = clf.predict_proba(s.data[FRAME_T])
+ assert all(hit(prob, pos, i) for i in bright)
+ assert (prob > 0.5).mean() < 0.25, "foreground has swallowed the frame"
+
+ @pytest.mark.parametrize("kw, match", [
+ ({"gap": 8, "background_dilation": 4}, "background_dilation"),
+ ({"particle_class": 99}, "no class with id"),
+ ])
+ def test_errors(self, kw, match):
+ m = self._disc((32, 32), 16, 16, 4)
+ with pytest.raises((ValueError, KeyError), match=match):
+ masks_to_labels(m, **kw)
+
+ def test_a_4d_input_raises(self):
+ with pytest.raises(ValueError, match=r"\(h, w\) or \(n, h, w\)"):
+ masks_to_labels(np.zeros((2, 2, 8, 8), bool))
+
+ def test_a_store_of_the_wrong_shape_raises(self):
+ store = LabelStore(frame_shape=(16, 16))
+ with pytest.raises(ValueError, match="the store holds"):
+ masks_to_labels(np.ones((32, 32), bool), store=store)
+
+
+# ── training ─────────────────────────────────────────────────────────────────
+
+class TestTraining:
+ def test_report_describes_what_was_trained(self, trained, labels):
+ rep = trained.report
+ assert rep["n_pixels"] == len(labels)
+ assert rep["n_channels"] == FeatureSpec().n_channels
+ assert rep["labelled_frames"] == labels.labelled_frames()
+ assert rep["train_accuracy"] > 0.95
+ assert rep["featurise_s"] > 0 and rep["fit_s"] > 0
+ assert set(rep["pixels_per_class"]) == {"0", "1"}
+
+ def test_only_painted_classes_become_head_columns(self, trained):
+ """The default store offers three classes; only two were painted, and a
+ column that never sees a positive example would emit noise."""
+ assert [c.id for c in trained.classes] == [0, 1]
+ assert trained.particle_class_ids == [0]
+
+ def test_multi_class_prediction(self, movie, geom):
+ s, _gt = movie
+ pos, radii, present, _faint, shape = geom
+ store = paint_scribbles(geom)
+ dark = np.zeros(shape, bool)
+ dark[0:6, 30:90] = True
+ store.paint(FRAME_T, dark & _clear_of_particles(shape, pos, radii,
+ present, 3.0), 2)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ rep = clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ assert rep["n_classes"] == 3
+ lbl = clf.predict_labels(s.data[FRAME_T])
+ assert set(np.unique(lbl).tolist()) == {0, 1, 2}
+ # The particle probability is still only the particle class.
+ prob = clf.predict_proba(s.data[FRAME_T])
+ assert prob.max() <= 1.0 and (prob > 0.5).mean() < 0.25
+
+ def test_progress_is_reported_per_labelled_frame(self, movie, geom):
+ s, _gt = movie
+ store = paint_scribbles(geom)
+ store.paint_disc(20, 30, 30, 3, 1)
+ calls = []
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0, epochs=20)
+ clf.fit(store, s.data, progress=lambda d, t: calls.append((d, t)))
+ assert calls[-1] == (2, 2)
+ assert all(t == 2 for _d, t in calls)
+
+ def test_trains_from_a_hyperspy_signal_without_materialising_it(self, movie,
+ geom):
+ """`frames` goes through `drift.frames.frame_source`, which reads one
+ frame at a time — the CLAUDE.md memory-safety rule."""
+ s, _gt = movie
+ store = paint_scribbles(geom)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0, epochs=20)
+ assert clf.fit(store, s)["n_pixels"] == len(store)
+
+ def test_an_empty_store_raises(self, movie):
+ s, _gt = movie
+ clf = ScribbleClassifier(device=DEVICE)
+ with pytest.raises(ValueError, match="nothing painted"):
+ clf.fit(LabelStore(frame_shape=s.data[0].shape), s.data)
+
+ def test_one_class_only_raises(self, movie):
+ s, _gt = movie
+ store = LabelStore(frame_shape=s.data[0].shape)
+ store.paint_disc(0, 20, 20, 3, 0)
+ clf = ScribbleClassifier(device=DEVICE)
+ with pytest.raises(ValueError, match="only one class"):
+ clf.fit(store, s.data)
+
+ def test_a_frame_of_the_wrong_shape_raises(self, movie, geom):
+ """A flat index means nothing without the shape; silently accepting a
+ different one would scatter the training samples."""
+ s, _gt = movie
+ store = paint_scribbles(geom)
+ clf = ScribbleClassifier(device=DEVICE)
+ with pytest.raises(ValueError, match="label store holds"):
+ clf.fit(store, {FRAME_T: s.data[FRAME_T][:64, :64]})
+
+ def test_predicting_before_training_raises(self, movie):
+ s, _gt = movie
+ with pytest.raises(RuntimeError, match="not been trained"):
+ ScribbleClassifier(device=DEVICE).predict_proba(s.data[0])
+
+ def test_no_particle_class_is_a_clear_error(self, movie, geom):
+ """A softmax over three backgrounds is a valid model with no foreground;
+ say so instead of returning a zero map."""
+ s, _gt = movie
+ pos, radii, present, faint, shape = geom
+ store = LabelStore(frame_shape=shape,
+ classes=[ScribbleClass(0, "film", particle=False),
+ ScribbleClass(1, "vacuum", particle=False)])
+ store.paint(FRAME_T, _clear_of_particles(shape, pos, radii, present, 3.0)
+ & (np.mgrid[0:shape[0], 0:shape[1]][0] < 20), 0)
+ store.paint(FRAME_T, _clear_of_particles(shape, pos, radii, present, 3.0)
+ & (np.mgrid[0:shape[0], 0:shape[1]][0] > 80), 1)
+ clf = ScribbleClassifier(device=DEVICE, epochs=20)
+ clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ with pytest.raises(RuntimeError, match="marked as a particle class"):
+ clf.predict_proba(s.data[FRAME_T])
+
+ def test_class_probabilities_sum_to_one_where_valid(self, movie, trained):
+ s, _gt = movie
+ p = trained.predict_class_proba(s.data[FRAME_T])
+ assert np.abs(p.sum(axis=0) - 1.0).max() < 1e-4
+
+ def test_segment_forwards_to_the_shared_instance_split(self, movie, trained,
+ geom):
+ """Plan §0.2: this engine stops at a probability map; the instance stage
+ is written once, in classical.split_instances."""
+ s, _gt = movie
+ pos, _radii, present, _faint, _shape = geom
+ lab = trained.segment(s.data[FRAME_T], SegmentParams(min_size=8))
+ assert lab.dtype == np.int32
+ assert lab.max() >= int(present.sum()) - 1 # the merge pair may join
+ for i in np.flatnonzero(present):
+ assert hit(lab, pos, i), f"particle {i} missing from the instances"
+
+
+# ── the acceptance gates ─────────────────────────────────────────────────────
+
+class TestTheClassicalBaselineMissesThem:
+ """Makes :class:`TestSensitivityGate` non-vacuous, in this file.
+
+ ``test_particle_movie_fixture.py::TestSegmentationOnTheFixture`` already pins
+ that the classical engine misses the faint probes at its default sensitivity;
+ this repeats the measurement next to the gate it justifies, because a gate
+ whose baseline lives in another file is one rename away from being vacuous.
+ """
+
+ def test_classical_finds_the_bright_particles(self, movie, geom):
+ s, _gt = movie
+ pos, _radii, present, faint, _shape = geom
+ lab = segment_frame(s.data[FRAME_T],
+ SegmentParams(min_size=25, gaussian=1.0))
+ want = np.flatnonzero(present & ~faint)
+ assert all(hit(lab, pos, i) for i in want)
+
+ def test_classical_finds_neither_faint_probe(self, movie, geom):
+ s, _gt = movie
+ pos, _radii, _present, faint, _shape = geom
+ lab = segment_frame(s.data[FRAME_T],
+ SegmentParams(min_size=25, gaussian=1.0))
+ found = [int(i) for i in np.flatnonzero(faint) if hit(lab, pos, i)]
+ assert found == [], (
+ f"the classical engine now finds faint probes {found}; the §0.9 gate "
+ "below is no longer measuring anything and the fixture's "
+ "faint_amplitude should be lowered")
+
+ def test_a_looser_threshold_does_not_rescue_it(self, movie, geom):
+ """Why §0.9 says the learned classifier is the primary path and threshold
+ tuning is not: by the time the threshold is loose enough to include the
+ faint probes it has merged the film into one giant region."""
+ s, _gt = movie
+ pos, _radii, _present, faint, _shape = geom
+ rescued = False
+ for sens in (0.6, 0.7, 0.8, 0.9, 1.0):
+ lab = segment_frame(
+ s.data[FRAME_T],
+ SegmentParams(min_size=25, gaussian=1.0, sensitivity=sens))
+ got = [int(i) for i in np.flatnonzero(faint) if hit(lab, pos, i)]
+ if len(got) == 2:
+ # Only counts as a rescue if the frame is still a segmentation.
+ fg = float((lab > 0).mean())
+ rescued = fg < 0.30
+ assert not rescued, (
+ "raising the classical sensitivity now finds both faint probes "
+ "without flooding the frame — if that is a real improvement this "
+ "test has served its purpose, but check deliberately")
+
+
+def paint_seam(store, geom, t: int = FRAME_T, width: int = 2):
+ """Add boundary strokes along the joins between touching particles.
+
+ The seam BETWEEN two bodies, never the outline of a lone one. That
+ distinction is the whole art of using this class and getting it wrong is
+ silent: a head taught particle outlines shrinks every body and splits
+ nothing, which measured as a MERGED touching pair and a 40% area loss on the
+ fixture's merge frame.
+ """
+ from scipy import ndimage as ndi
+
+ pos, radii, present, _faint, shape = geom
+ h, w = shape
+ yy, xx = np.mgrid[0:h, 0:w]
+ idx = list(np.flatnonzero(present))
+ grown = [ndi.binary_dilation(
+ ((yy - pos[i, 0]) ** 2 + (xx - pos[i, 1]) ** 2) <= radii[i] ** 2,
+ iterations=width) for i in idx]
+ seam = np.zeros((h, w), bool)
+ for a in range(len(grown)):
+ for b in range(a + 1, len(grown)):
+ seam |= grown[a] & grown[b]
+ if seam.any():
+ store.paint(t, seam, 3)
+ return store
+
+
+class TestBoundaryClass:
+ """The ilastik third class, and the route it unlocks.
+
+ It is a **performance** feature: a taught boundary lets ``split_instances``
+ skip the distance transform and the watershed — 1.78 s down to 0.33 s at
+ 4096². So these check both that it is wired up and that turning it on
+ does not cost detection — a faster segmentation that finds fewer particles
+ is a regression, not an optimisation (plan §0.9).
+ """
+
+ def test_the_default_class_set_has_one(self):
+ classes = default_classes()
+ edge = [c for c in classes if c.boundary]
+ assert len(edge) == 1 and edge[0].name == "boundary"
+ assert not edge[0].particle, (
+ "a seam is not part of a body; counting it as foreground would glue "
+ "back together exactly the particles it separates")
+
+ def test_a_class_cannot_be_both_particle_and_boundary(self):
+ with pytest.raises(ValueError, match="both particle and boundary"):
+ ScribbleClass(0, "confused", particle=True, boundary=True)
+
+ def test_round_trips_through_a_dict(self):
+ c = ScribbleClass(3, "boundary", "#f38ba8", boundary=True)
+ assert ScribbleClass.from_dict(c.to_dict()) == c
+
+ def test_a_dict_from_before_the_class_existed_still_loads(self):
+ """A session or a saved model written by an older build has no
+ ``boundary`` key, and must come back as the particle/background-only
+ setup it was — not fail, and not silently become a boundary."""
+ old = {"id": 0, "name": "particle", "colour": "#f9a03f", "particle": True}
+ c = ScribbleClass.from_dict(old)
+ assert c.particle and not c.boundary
+
+ def test_untrained_boundary_reports_none(self, trained, movie):
+ """The default fixture paints no seam, so there is no boundary — and the
+ answer must be None rather than an all-zero map. The two mean different
+ things to the split: "use the watershed" versus "a boundary was taught
+ and this frame has none", which would leave every touching pair merged.
+ """
+ s, _gt = movie
+ assert not trained.has_boundary
+ assert trained.boundary_class_ids == []
+ fg, bnd = trained.predict_foreground_boundary(s.data[FRAME_T])
+ assert bnd is None
+ assert fg.shape == s.data[FRAME_T].shape
+
+ def test_a_trained_boundary_is_reported_and_predicted(self, movie, geom):
+ s, _gt = movie
+ store = paint_seam(paint_scribbles(geom), geom)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ report = clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ assert report["has_boundary"] is True
+ assert clf.has_boundary and clf.boundary_class_ids == [3]
+ _fg, bnd = clf.predict_foreground_boundary(s.data[FRAME_T])
+ assert bnd is not None and bnd.shape == tuple(geom[4])
+ assert (bnd > 0.5).any(), "a trained boundary class predicted nothing"
+
+ def test_segment_takes_the_boundary_route_when_one_is_trained(
+ self, movie, geom, monkeypatch):
+ """The wizard calls ``segment``; this is what makes the fast route
+ automatic without the caret having to know about it."""
+ from skimage import segmentation as skseg
+
+ s, _gt = movie
+ store = paint_seam(paint_scribbles(geom), geom)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+
+ called = []
+ monkeypatch.setattr(skseg, "watershed",
+ lambda *a, **k: called.append("watershed"))
+ clf.segment(s.data[FRAME_T], SegmentParams(min_size=10))
+ assert called == [], "segment ran the watershed despite a trained boundary"
+
+ def test_segment_still_uses_the_watershed_without_one(self, trained, movie,
+ monkeypatch):
+ """A user who never paints a boundary must not silently get worse
+ splitting — the fallback is the whole safety of making this automatic."""
+ from skimage import segmentation as skseg
+
+ s, _gt = movie
+ real = skseg.watershed
+ called = []
+ monkeypatch.setattr(skseg, "watershed", lambda *a, **k: (
+ called.append("watershed"), real(*a, **k))[1])
+ trained.segment(s.data[FRAME_T], SegmentParams(min_size=10))
+ assert called == ["watershed"]
+
+ def test_the_faint_probes_survive_the_boundary_route(self, movie, geom):
+ """**Plan §0.9 on the new path.** The boundary class exists to make the
+ split cheap; if it costs a faint detection it is not worth having."""
+ s, _gt = movie
+ pos, _radii, _present, faint, _shape = geom
+ store = paint_seam(paint_scribbles(geom), geom)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ labels = clf.segment(s.data[FRAME_T], SegmentParams(min_size=5))
+ missed = [int(i) for i in np.flatnonzero(faint)
+ if not hit(labels, pos, i)]
+ assert not missed, f"the boundary route lost faint probe(s) {missed}"
+
+ def test_the_boundary_route_agrees_with_the_watershed_on_the_fixture(
+ self, movie, geom):
+ """Same classifier, both routes: the count and the median area must
+ agree, or the speed is not worth having."""
+ s, _gt = movie
+ store = paint_seam(paint_scribbles(geom), geom)
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ p = SegmentParams(min_size=5)
+ fg, bnd = clf.predict_foreground_boundary(s.data[FRAME_T])
+ by_boundary = split_instances(fg, p, boundary=bnd)
+ by_watershed = split_instances(fg, p)
+
+ def areas(lab):
+ c = np.bincount(lab.ravel())[1:]
+ return np.sort(c[c > 0])
+
+ assert int(by_boundary.max()) == int(by_watershed.max()), (
+ f"boundary found {by_boundary.max()} particles, watershed "
+ f"{by_watershed.max()}")
+ a_b, a_w = areas(by_boundary), areas(by_watershed)
+ assert abs(np.median(a_b) - np.median(a_w)) / np.median(a_w) < 0.15, (
+ f"median area {np.median(a_b)} vs watershed {np.median(a_w)}")
+
+
+class TestSensitivityGate:
+ """**Plan §0.9 — the headline gate for the whole feature.**
+
+ Trained on eleven scribbles (four dabs on bright particles, one dab on the
+ smaller faint probe, four background sweeps and four boundary rings), the
+ classifier must find BOTH faint probes and keep every bright one.
+
+ Note what the scribbles do and do not contain. Faint probe **8** (r=3) is
+ painted; probe **7** (r=4) is not, so its detection is genuinely held out.
+ :class:`TestBrightOnlyLabelsDoNotGeneralise` records the boundary of the
+ claim: with *no* faint example at all, neither this head nor the RandomForest
+ reference finds them, so at least one faint scribble is required. That is a
+ property of the problem — an 8x contrast extrapolation — not of the head.
+ """
+
+ def test_finds_both_faint_probes(self, movie, geom, proba):
+ _s, _gt = movie
+ pos, _radii, _present, faint, _shape = geom
+ scores = {int(i): float(proba[int(round(pos[i, 0])),
+ int(round(pos[i, 1]))])
+ for i in np.flatnonzero(faint)}
+ assert all(v > 0.5 for v in scores.values()), (
+ f"faint probes not found: {scores}")
+
+ def test_the_held_out_faint_probe_is_found(self, geom, proba):
+ """Probe 7 is never painted. This is the part that is not memorisation."""
+ pos = geom[0]
+ assert proba[int(round(pos[7, 0])), int(round(pos[7, 1]))] > 0.5
+
+ def test_every_bright_particle_is_still_found(self, geom, proba):
+ """Sensitivity that costs the easy detections is not sensitivity."""
+ pos, _radii, present, faint, _shape = geom
+ for i in np.flatnonzero(present & ~faint):
+ assert hit(proba, pos, i), f"lost bright particle {i}"
+
+ def test_the_foreground_has_not_swallowed_the_frame(self, geom, proba):
+ """A classifier that calls everything a particle would pass every test
+ above. The nine discs cover ~9% of the frame."""
+ fraction = float((proba > 0.5).mean())
+ assert 0.05 < fraction < 0.20, (
+ f"foreground is {fraction:.1%} of the frame; the true particle "
+ "coverage is ~9%")
+
+ def test_the_instance_count_is_right(self, movie, geom, proba):
+ _s, _gt = movie
+ present = geom[2]
+ lab = split_instances(proba, SegmentParams(min_size=8))
+ assert int(present.sum()) - 1 <= lab.max() <= int(present.sum()) + 1, (
+ f"found {lab.max()} instances, expected ~{int(present.sum())}")
+
+ def test_it_generalises_to_frames_it_never_saw(self, movie, trained):
+ """Labels came from t=12 alone. Every other frame has different drift,
+ different particles present, and — at t=8 onwards — a nucleated one."""
+ s, gt = movie
+ for t in (0, 6, 18, int(gt["n_frames"]) - 1):
+ prob = trained.predict_proba(s.data[t])
+ pos, _radii, present = particle_truth_at(gt, t)
+ faint = np.asarray(gt["p_faint"], bool)
+ h, w = s.data[t].shape
+ missed = []
+ for i in np.flatnonzero(present):
+ cy, cx = int(round(pos[i, 0])), int(round(pos[i, 1]))
+ if not (0 <= cy < h and 0 <= cx < w):
+ continue # drifted out of frame
+ if prob[cy, cx] <= 0.5:
+ missed.append((int(i), bool(faint[i])))
+ assert not missed, f"frame {t}: missed {missed}"
+
+ def test_a_coarse_stack_undersizes_the_small_particles(self, movie, geom,
+ labels):
+ """What the fine scales actually buy — and it is NOT detection.
+
+ Plan §0.9 says small-object detection needs the fine scales and forbids
+ coarsening "without a documented sensitivity measurement". This is that
+ measurement, and it corrects the guess: a coarse ``(4, 8)`` stack still
+ *finds* both faint probes. What it loses is their size. Measured mean
+ absolute error in recovered radius over the isolated particles:
+
+ (0.5, 1, 2, 4, 8) 13 % (2, 4, 8) 20 % (4, 8) 26 %
+
+ A particle found and then measured 44% too small (the r=3 probe on the
+ ``(4, 8)`` stack) is worse than an honest miss, because it silently enters
+ the size distribution the whole feature exists to produce.
+ """
+ from spyde.particles import measure_frame
+ from spyde.signals.particles import COL
+
+ s, gt = movie
+ pos, radii, present, _faint, _shape = geom
+ merge = tuple(int(v) for v in gt["merge_pair"])
+
+ def radius_error(sigmas):
+ clf = ScribbleClassifier(FeatureSpec(sigmas=sigmas), device=DEVICE,
+ seed=0)
+ clf.fit(labels, {FRAME_T: s.data[FRAME_T]})
+ lab = split_instances(clf.predict_proba(s.data[FRAME_T]),
+ SegmentParams(min_size=8))
+ rows, _c = measure_frame(lab, s.data[FRAME_T], t=FRAME_T, scale=1.0)
+ errs = []
+ for i in np.flatnonzero(present):
+ if i in merge:
+ continue # a merged pair is not one disc
+ lbl = lab[int(round(pos[i, 0])), int(round(pos[i, 1]))]
+ row = rows[rows[:, COL["label"]] == lbl]
+ assert lbl and len(row), f"sigmas={sigmas}: lost particle {i}"
+ got = float(row[0, COL["equiv_diameter"]]) / 2.0
+ errs.append(abs(got - radii[i]) / radii[i])
+ return float(np.mean(errs))
+
+ fine = radius_error(DEFAULT_SIGMAS)
+ coarse = radius_error((4.0, 8.0))
+ assert fine < 0.20, f"default stack radius error {fine:.1%}"
+ assert coarse > 1.5 * fine, (
+ f"a coarse (4, 8) stack now measures radii as well as the default one "
+ f"({coarse:.1%} vs {fine:.1%}) — if real, the sigma floor could be "
+ "raised, but re-measure deliberately before doing it")
+
+ def test_the_nucleating_particle_appears_at_its_known_frame(self, movie,
+ trained):
+ """§0.9's actual motivation: missing a particle's first appearance
+ destroys the nucleation event."""
+ s, gt = movie
+ i, nuc = int(gt["nucleation_index"]), int(gt["nucleation_frame"])
+ before = trained.predict_proba(s.data[nuc - 1])
+ after = trained.predict_proba(s.data[nuc])
+ pos_b = particle_truth_at(gt, nuc - 1)[0]
+ pos_a = particle_truth_at(gt, nuc)[0]
+ assert before[int(round(pos_b[i, 0])), int(round(pos_b[i, 1]))] < 0.5
+ assert after[int(round(pos_a[i, 0])), int(round(pos_a[i, 1]))] > 0.5
+
+
+class TestBrightOnlyLabelsAreNotEnough:
+ """Records the boundary of the §0.9 claim, and that it is not the head's fault.
+
+ Trained on bright particles only — no faint example anywhere in the labels —
+ the head finds **1 of the 2** faint probes (0.67 for the r=3 one, 0.013 for the
+ r=4 one) and the sklearn RandomForest reference finds **0 of 2** (both exactly
+ 0.0). Adding a *single seven-pixel dab* on the r=3 probe takes the head to
+ 2 of 2 at 0.98 and 1.00.
+
+ **Those exact counts depend on where the background scribbles land**, so the
+ assertion below is the robust `< 2` rather than `== 1`. An independent probe
+ with different background dabs got **0 of 2** bright-only and 1 of 2 with the
+ dab — same conclusion, different numbers. Read the figures above as one
+ observation, not as a constant; the claim being pinned is only that bright-only
+ labels are *not sufficient*.
+
+ Two things worth having in the record:
+
+ * The plan's §0.9 gate holds as written, with a clarification: "trained on a
+ few scribbles" must include **at least one faint example**. That is what a
+ user does anyway — they paint what they can see, and they can see these —
+ and it is why the wizard's class list shows per-class pixel counts (plan
+ B7): under-training a class is the failure mode, and the counts are how you
+ notice.
+ * The forest's 0.0 is not noise, it is structural. A tree ensemble cannot
+ predict outside the leaves its training data reached, so a contrast 8x below
+ anything it was shown is simply unreachable; the MLP extrapolates its
+ decision boundary and gets one of the two for free. So for *sensitivity*
+ specifically — which §0.9 makes the priority — the shipped head is better
+ than the reference it is checked against, not merely faster.
+ """
+
+ @pytest.fixture(scope="class")
+ def bright_only(self, movie, geom):
+ s, _gt = movie
+ store = paint_scribbles(geom, include_faint=())
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(store, {FRAME_T: s.data[FRAME_T]})
+ return store, clf, clf.predict_proba(s.data[FRAME_T])
+
+ def test_the_head_does_not_find_both(self, geom, bright_only):
+ pos, _radii, _present, faint, _shape = geom
+ _store, _clf, prob = bright_only
+ scores = {int(i): float(prob[int(round(pos[i, 0])),
+ int(round(pos[i, 1]))])
+ for i in np.flatnonzero(faint)}
+ found = [i for i, v in scores.items() if v > 0.5]
+ assert len(found) < 2, (
+ f"bright-only labels now generalise to BOTH faint probes ({scores}) — "
+ "if that is real, TestSensitivityGate can drop its faint scribble and "
+ "become a much stronger claim; check deliberately")
+
+ def test_the_bright_particles_are_still_all_found(self, geom, bright_only):
+ pos, _radii, present, faint, _shape = geom
+ _store, _clf, prob = bright_only
+ assert all(hit(prob, pos, i) for i in np.flatnonzero(present & ~faint))
+
+ def test_the_random_forest_reference_finds_neither(self, movie, geom,
+ bright_only):
+ """The forest cannot extrapolate past its leaves; the MLP can. This is the
+ one place the shipped head beats its own reference."""
+ s, _gt = movie
+ pos, _radii, _present, faint, _shape = geom
+ store, _clf, _prob = bright_only
+ _rf, predict = random_forest_reference(
+ store, {FRAME_T: s.data[FRAME_T]}, FeatureSpec(), device=DEVICE)
+ rp = predict(s.data[FRAME_T])
+ assert max(float(rp[int(round(pos[i, 0])), int(round(pos[i, 1]))])
+ for i in np.flatnonzero(faint)) < 0.5
+
+ def test_one_extra_dab_is_all_it_takes(self, movie, geom, bright_only):
+ """The interaction this feature is actually for: seven more labelled
+ pixels, one retrain, both probes found."""
+ s, _gt = movie
+ pos, _radii, _present, faint, _shape = geom
+ store, _clf, _prob = bright_only
+ with_faint = paint_scribbles(geom)
+ assert 0 < len(with_faint) - len(store) < 20, (
+ "the extra scribble is no longer a single small dab, so this stopped "
+ "being a statement about how little labelling is needed")
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ clf.fit(with_faint, {FRAME_T: s.data[FRAME_T]})
+ prob = clf.predict_proba(s.data[FRAME_T])
+ assert all(hit(prob, pos, i) for i in np.flatnonzero(faint))
+
+
+class TestRandomForestParity:
+ """**Plan B3's acceptance gate**: agreement with the reference implementation.
+
+ ``sklearn.ensemble.RandomForestClassifier`` is what ParticleSpy and ilastik
+ use. Both sides read the *identical* channels through the *identical* sampler
+ (``random_forest_reference`` lives in ``scribble.py`` for exactly that
+ reason), so a disagreement is about the head and nothing else.
+
+ Measured IoU of the two foreground masks on the fixture: **0.94**. The gate is
+ set at 0.75 — the margin is deliberate, because the objects are small discs
+ whose IoU is dominated by a one-pixel boundary difference, and a tighter bar
+ would fail on an epoch-count change that means nothing (0.94 at 300 epochs,
+ 0.90 at 200, 0.86 at 500).
+ """
+
+ @pytest.fixture(scope="class")
+ def forest(self, movie, labels):
+ s, _gt = movie
+ _rf, predict = random_forest_reference(
+ labels, {FRAME_T: s.data[FRAME_T]}, FeatureSpec(), device=DEVICE)
+ return predict(s.data[FRAME_T])
+
+ def test_foreground_masks_agree_by_iou(self, proba, forest):
+ a, b = proba > 0.5, forest > 0.5
+ iou = float((a & b).sum()) / max(1, int((a | b).sum()))
+ assert iou > 0.75, (
+ f"IoU {iou:.3f} against the RandomForest reference on identical "
+ "labels and identical features")
+
+ def test_both_agree_on_every_ground_truth_particle(self, geom, proba, forest):
+ pos, _radii, present, _faint, _shape = geom
+ for i in np.flatnonzero(present):
+ assert hit(proba, pos, i) == hit(forest, pos, i), (
+ f"the head and the forest disagree about particle {i}")
+
+ def test_they_cover_a_similar_fraction_of_the_frame(self, proba, forest):
+ """A high IoU with wildly different coverage would mean one is a subset of
+ the other, which is a different kind of agreement."""
+ fa, fb = float((proba > 0.5).mean()), float((forest > 0.5).mean())
+ assert abs(fa - fb) < 0.05, f"coverage {fa:.3f} vs {fb:.3f}"
+
+
+class TestNaNBorder:
+ """**Plan trap 2 / gate A7**: no particle is ever detected in NaN padding.
+
+ A drift-corrected frame keeps its full size with uncovered pixels set to NaN
+ (``spyde.drift.warp``). Segmentation that ignores that finds a large "particle"
+ along the edge, which then nucleates a spurious track — the single most likely
+ integration bug in this feature.
+ """
+
+ @pytest.fixture(scope="class")
+ def warped(self, movie):
+ s, _gt = movie
+ out = shift_frame(s.data[FRAME_T], (7.0, -5.0)) # NaN on two edges
+ assert not np.isfinite(out).all(), "no NaN padding to test"
+ return out
+
+ def test_probability_is_exactly_zero_in_the_padding(self, trained, warped):
+ prob = trained.predict_proba(warped)
+ bad = ~np.isfinite(warped)
+ assert float(prob[bad].max()) == 0.0
+
+ def test_no_instance_lands_in_the_padding(self, trained, warped):
+ prob = trained.predict_proba(warped)
+ lab = split_instances(prob, SegmentParams(min_size=8))
+ assert set(np.unique(lab[~np.isfinite(warped)]).tolist()) == {0}
+
+ def test_predict_labels_reports_unlabelled_not_a_class(self, trained, warped):
+ """-1, not the background class: an invalid pixel is not a measurement."""
+ lbl = trained.predict_labels(warped)
+ bad = ~np.isfinite(warped)
+ assert set(np.unique(lbl[bad]).tolist()) == {UNLABELLED}
+ assert (lbl[~bad] != UNLABELLED).all()
+
+ def test_real_data_next_to_the_padding_still_classifies(self, trained, warped,
+ movie):
+ """The other half of the contract: NaN must not erase a band of real data.
+ A filter that propagated NaN outward would blank ~33 px (the halo) inside
+ the border, which is most of a 96-row frame."""
+ prob = trained.predict_proba(warped)
+ good = np.isfinite(warped)
+ assert float(prob[good].max()) > 0.5, "everything valid came back empty"
+ assert float((prob[good] > 0.5).mean()) > 0.02
+
+ def test_the_shifted_particles_are_found_at_their_shifted_positions(
+ self, movie, trained, warped, geom):
+ pos, _radii, present, _faint, _shape = geom
+ prob = trained.predict_proba(warped)
+ h, w = warped.shape
+ found = 0
+ for i in np.flatnonzero(present):
+ cy, cx = int(round(pos[i, 0] + 7.0)), int(round(pos[i, 1] - 5.0))
+ if 0 <= cy < h and 0 <= cx < w and np.isfinite(warped[cy, cx]):
+ found += prob[cy, cx] > 0.5
+ assert found >= 6, f"only {found} particles survived the warp"
+
+
+class TestDeterminism:
+ """Same seed + same labels → the same model. A user who changed nothing must
+ not see the segmentation change."""
+
+ def test_two_fits_with_the_same_seed_are_identical(self, movie, labels):
+ s, _gt = movie
+ outs = []
+ for _ in range(2):
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0,
+ epochs=60)
+ clf.fit(labels, {FRAME_T: s.data[FRAME_T]})
+ outs.append(clf.predict_proba(s.data[FRAME_T]))
+ assert np.array_equal(*outs)
+
+ def test_the_seed_actually_does_something(self, movie, labels):
+ s, _gt = movie
+ outs = []
+ for seed in (0, 1):
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=seed,
+ epochs=60)
+ clf.fit(labels, {FRAME_T: s.data[FRAME_T]})
+ outs.append(clf.predict_proba(s.data[FRAME_T]))
+ assert not np.array_equal(*outs)
+
+ def test_initialisation_does_not_consume_the_global_rng(self, movie, labels):
+ """The global torch stream is shared with the neural detector and every
+ dask worker; drawing from it would make 'same seed' depend on what else
+ ran first."""
+ import torch
+ torch.manual_seed(1234)
+ before = torch.randn(3)
+ torch.manual_seed(1234)
+ ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0, epochs=5).fit(
+ labels, {FRAME_T: movie[0].data[FRAME_T]})
+ after = torch.randn(3)
+ assert torch.equal(before, after)
+
+
+class TestSaveLoad:
+ """A recipe is the spec *and* the weights *and* the standardisation; any of
+ them alone predicts nonsense, so they live in one file."""
+
+ def test_round_trip_predicts_identically(self, movie, trained, tmp_path):
+ s, _gt = movie
+ path = str(tmp_path / "model.npz")
+ trained.save(path)
+ back = ScribbleClassifier.load(path, device=DEVICE)
+ assert np.array_equal(back.predict_proba(s.data[FRAME_T]),
+ trained.predict_proba(s.data[FRAME_T]))
+
+ def test_carries_the_spec_and_the_classes(self, trained, tmp_path):
+ path = str(tmp_path / "model.npz")
+ trained.save(path)
+ back = ScribbleClassifier.load(path, device=DEVICE)
+ assert back.spec == trained.spec
+ assert [c.to_dict() for c in back.classes] == \
+ [c.to_dict() for c in trained.classes]
+ assert back.hidden == trained.hidden and back.seed == trained.seed
+
+ def test_loads_without_pickle(self, trained, tmp_path):
+ """A model file must not be able to execute code on load."""
+ path = str(tmp_path / "model.npz")
+ trained.save(path)
+ with np.load(path, allow_pickle=False) as z:
+ assert "meta" in z.files and "feature_mean" in z.files
+
+ def test_a_recipe_applies_to_a_differently_scaled_dataset(self, movie,
+ trained, tmp_path):
+ """What ``normalize_frame`` buys: the same sample recorded as counts
+ rather than as a normalised float still segments."""
+ s, _gt = movie
+ path = str(tmp_path / "model.npz")
+ trained.save(path)
+ back = ScribbleClassifier.load(path, device=DEVICE)
+ rescaled = s.data[FRAME_T] * 4096.0 + 100.0
+ a = back.predict_proba(s.data[FRAME_T]) > 0.5
+ b = back.predict_proba(rescaled) > 0.5
+ iou = float((a & b).sum()) / max(1, int((a | b).sum()))
+ assert iou > 0.95, f"IoU across an intensity rescale is only {iou:.3f}"
+
+ def test_saving_before_training_raises(self, tmp_path):
+ with pytest.raises(RuntimeError, match="not been trained"):
+ ScribbleClassifier(device=DEVICE).save(str(tmp_path / "x.npz"))
+
+ def test_a_future_format_version_is_refused(self, trained, tmp_path):
+ path = tmp_path / "model.npz"
+ trained.save(str(path))
+ with np.load(path, allow_pickle=False) as z:
+ arrays = {k: z[k] for k in z.files}
+ meta = json.loads(str(arrays.pop("meta").item()))
+ meta["format_version"] = 99
+ np.savez_compressed(path, meta=np.array(json.dumps(meta)), **arrays)
+ with pytest.raises(ValueError, match="format version"):
+ ScribbleClassifier.load(str(path), device=DEVICE)
+
+ def test_a_spec_that_no_longer_matches_the_weights_is_refused(self, trained,
+ tmp_path):
+ """The failure this guards is silent otherwise: a channel count mismatch
+ would only show as a wrong segmentation."""
+ path = tmp_path / "model.npz"
+ trained.save(str(path))
+ with np.load(path, allow_pickle=False) as z:
+ arrays = {k: z[k] for k in z.files}
+ meta = json.loads(str(arrays.pop("meta").item()))
+ meta["spec"]["median"] = False
+ np.savez_compressed(path, meta=np.array(json.dumps(meta)), **arrays)
+ with pytest.raises(ValueError, match="come apart"):
+ ScribbleClassifier.load(str(path), device=DEVICE)
+
+
+class TestInteractionBudget:
+ """**Plan B3's hard budget**: train + apply to the visible frame under ~1 s.
+
+ Measured here on the 96x112 fixture, CPU, default 36-channel spec: **0.49 s**
+ total = 14 ms featurise + 457 ms for the 300 Adam steps + 20 ms to apply. The
+ fit is dominated by per-step dispatch overhead rather than arithmetic —
+ 1.5 ms/step at any torch thread count from 1 to 24 — so it is a *fixed* cost,
+ independent of frame size and of how much was painted.
+
+ The bar here is 2.0 s, not 1.0 s: this is a shared CI box and a *timing*
+ assertion at the measured value is a flake generator. The number that matters
+ is the one above; this test exists to catch an order-of-magnitude regression.
+ """
+
+ def test_train_plus_apply_is_interactive(self, movie, labels):
+ s, _gt = movie
+ frame = s.data[FRAME_T]
+ # Warm torch first: the FIRST op in a fresh process pays a one-time ~1 s
+ # of thread-pool and kernel-selection cost that no user ever sees twice.
+ ScribbleClassifier(FeatureSpec(), device=DEVICE, epochs=5).fit(
+ labels, {FRAME_T: frame})
+
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, seed=0)
+ t0 = time.perf_counter()
+ clf.fit(labels, {FRAME_T: frame})
+ t1 = time.perf_counter()
+ clf.predict_proba(frame)
+ elapsed = time.perf_counter() - t0
+ assert elapsed < 2.0, (
+ f"train+apply took {elapsed:.2f} s (fit {t1 - t0:.2f} s); the plan's "
+ "budget is ~1 s and the measured value on this box is 0.49 s")
+
+ def test_applying_to_one_frame_is_a_fraction_of_the_budget(self, movie,
+ trained):
+ """Scrubbing the navigator re-applies without re-training, so this is the
+ number that has to stay small."""
+ s, _gt = movie
+ trained.predict_proba(s.data[0]) # warm
+ t0 = time.perf_counter()
+ trained.predict_proba(s.data[1])
+ assert time.perf_counter() - t0 < 0.5
+
+
+# ── the MPS device-serialisation contract ────────────────────────────────────
+
+class _FakeDev:
+ """Stands in for ``torch.device('mps')``; only ``.type`` is consulted. Same
+ double as ``test_device_lock.py`` uses, so the contract is exercised on any
+ machine without touching Metal."""
+
+ def __init__(self, type_="mps"):
+ self.type = type_
+
+ def __str__(self):
+ return self.type
+
+
+def _lock_is_held_by_me() -> bool:
+ from spyde.device_lock import DEVICE_LOCK
+ got = []
+
+ def probe():
+ got.append(DEVICE_LOCK.acquire(blocking=False))
+ if got[-1]:
+ DEVICE_LOCK.release()
+
+ t = threading.Thread(target=probe)
+ t.start()
+ t.join()
+ return not got[0]
+
+
+class _LockSpy:
+ """A stand-in for ``accelerator_lock`` that records entries and nesting depth.
+
+ Two things need pinning and neither is visible from the real lock on a CPU box
+ (where it is a null context): that the block is entered **with the right
+ device**, and that every device submission happens **inside** it. So the spy
+ records the devices it was called with and exposes :attr:`inside`, which the
+ patched call sites report. The real-lock behaviour on a fake MPS device is
+ pinned separately in ``test_device_lock.py``.
+ """
+
+ def __init__(self):
+ self.devices: list = []
+ self.depth = 0
+
+ def __call__(self, device=None, **kw):
+ import contextlib
+
+ @contextlib.contextmanager
+ def ctx():
+ self.devices.append(device)
+ self.depth += 1
+ try:
+ yield
+ finally:
+ self.depth -= 1
+
+ return ctx()
+
+ @property
+ def inside(self) -> bool:
+ return self.depth > 0
+
+
+class TestDeviceLock:
+ """Every torch call site in these two modules takes ``accelerator_lock``.
+
+ A lock only works if EVERY participant takes it — the last crash of this class
+ existed because one path skipped it (CLAUDE.md § GPU Computing). The shared
+ lock's *identity* and the real-lock behaviour are pinned in
+ ``test_device_lock.py``; these tests pin the *nesting*, i.e. that no device
+ submission escapes the block.
+ """
+
+ def test_feature_bands_hold_the_real_lock_on_mps(self, monkeypatch):
+ """``map_feature_bands`` is the ONE torch entry point in features.py —
+ ``feature_tensor``, ``feature_stack`` and ``sample_features`` all go
+ through it — so this one acquisition covers the whole module."""
+ held = []
+ monkeypatch.setattr(feat, "_band_stack",
+ lambda *a, **k: held.append(_lock_is_held_by_me()))
+ map_feature_bands(np.zeros((8, 8), np.float32), FeatureSpec(),
+ device=_FakeDev(), fn=lambda *a: None)
+ assert held == [True], "the feature stack ran unserialised on MPS"
+ assert not _lock_is_held_by_me(), "lock leaked"
+
+ def test_no_lock_off_mps(self):
+ """CUDA concurrency is a deliberate throughput win; serialising it would
+ be a pure regression."""
+ held = []
+ map_feature_bands(np.zeros((8, 8), np.float32),
+ FeatureSpec(sigmas=(1.0,), median=False, minimum=False,
+ maximum=False),
+ device=DEVICE,
+ fn=lambda *a: held.append(_lock_is_held_by_me()))
+ assert held and not any(held)
+
+ @pytest.mark.parametrize("call", [
+ lambda img: feature_tensor(img, device=DEVICE),
+ lambda img: feature_stack(img, device=DEVICE),
+ lambda img: sample_features(img, np.array([0, 5]), device=DEVICE),
+ ])
+ def test_every_features_entry_point_takes_the_lock(self, call, monkeypatch):
+ spy = _LockSpy()
+ monkeypatch.setattr(feat, "accelerator_lock", spy)
+ call(np.zeros((16, 16), np.float32))
+ assert spy.devices and all(d is DEVICE for d in spy.devices)
+ assert spy.depth == 0, "lock not released"
+
+ def test_the_fit_takes_the_lock_around_every_submission(self, movie, labels,
+ monkeypatch):
+ """The featurise, the optimiser loop and the ``.to(device)`` inside
+ ``_build_mlp`` — a cold head load is the specific hole CLAUDE.md names."""
+ from spyde.particles import scribble as scr
+
+ s, _gt = movie
+ spy = _LockSpy()
+ inside = []
+ monkeypatch.setattr(scr, "accelerator_lock", spy)
+ for name in ("sample_features", "_build_mlp"):
+ real = getattr(scr, name)
+ monkeypatch.setattr(scr, name, lambda *a, _r=real, **k: (
+ inside.append(spy.inside), _r(*a, **k))[1])
+
+ clf = ScribbleClassifier(FeatureSpec(), device=DEVICE, epochs=5)
+ clf.fit(labels, {FRAME_T: s.data[FRAME_T]})
+ assert len(inside) >= 2 and all(inside), (
+ "some of the fit's device work ran outside the lock")
+ assert spy.devices == [clf.device], (
+ f"expected ONE acquisition with the fit's device; got {spy.devices}")
+ assert spy.depth == 0
+
+ def test_predict_takes_the_lock_around_every_submission(self, movie, trained,
+ monkeypatch):
+ """The interactive path — it fires on every navigator move, which is
+ exactly the concurrency that took the backend down last time."""
+ from spyde.particles import scribble as scr
+
+ s, _gt = movie
+ spy = _LockSpy()
+ inside = []
+ real = scr.map_feature_bands
+ monkeypatch.setattr(scr, "accelerator_lock", spy)
+ monkeypatch.setattr(scr, "map_feature_bands", lambda *a, **k: (
+ inside.append(spy.inside), real(*a, **k))[1])
+ trained.predict_class_proba(s.data[FRAME_T])
+ assert inside == [True]
+ assert spy.devices == [trained.device]
+ assert spy.depth == 0
+
+ @pytest.mark.parametrize("call", [
+ "predict_foreground_boundary", "predict_proba",
+ "predict_boundary_proba", "segment",
+ ])
+ def test_every_boundary_accessor_takes_the_lock(self, movie, trained, call,
+ monkeypatch):
+ """The boundary route added three new public doors onto the device.
+
+ A lock only works if every participant takes it, and the last crash of
+ this class happened precisely because a newly added entry point reached
+ the device through an existing helper and nobody re-checked that the
+ helper's lock covered it. So all four doors are pinned, not just the one
+ the wizard happens to call today — and ONE acquisition each, because
+ reading two maps out of one softmax must not featurise twice.
+ """
+ from spyde.particles import scribble as scr
+
+ s, _gt = movie
+ spy = _LockSpy()
+ monkeypatch.setattr(scr, "accelerator_lock", spy)
+ getattr(trained, call)(s.data[FRAME_T])
+ assert spy.devices == [trained.device], (
+ f"{call} took {len(spy.devices)} acquisitions; expected exactly one")
+ assert spy.depth == 0, f"{call} leaked the lock"
+
+ def test_save_and_load_take_the_lock(self, trained, tmp_path, monkeypatch):
+ """Both move tensors across the host/device boundary, which is a
+ submission — ``save`` reading weights back and ``load`` pushing them out."""
+ from spyde.particles import scribble as scr
+
+ path = str(tmp_path / "m.npz")
+ spy = _LockSpy()
+ monkeypatch.setattr(scr, "accelerator_lock", spy)
+ trained.save(path)
+ assert spy.devices == [trained.device]
+ spy.devices.clear()
+ ScribbleClassifier.load(path, device=DEVICE)
+ assert len(spy.devices) == 1
+ assert spy.depth == 0
+
+ def test_the_random_forest_reference_takes_the_lock_too(self, movie, labels,
+ monkeypatch):
+ """It is a test helper, but it still submits torch work from whatever
+ thread calls it, and an unlocked path is an unlocked path."""
+ from spyde.particles import scribble as scr
+
+ s, _gt = movie
+ spy = _LockSpy()
+ monkeypatch.setattr(scr, "accelerator_lock", spy)
+ _rf, predict = random_forest_reference(
+ labels, {FRAME_T: s.data[FRAME_T]},
+ FeatureSpec(sigmas=(2.0,), median=False, minimum=False,
+ maximum=False),
+ n_estimators=4, device=DEVICE)
+ predict(s.data[FRAME_T])
+ assert len(spy.devices) == 2 and all(d is DEVICE for d in spy.devices)
+ assert spy.depth == 0
diff --git a/spyde/tests/migrated/test_particles_track.py b/spyde/tests/migrated/test_particles_track.py
new file mode 100644
index 00000000..9c26da9b
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_track.py
@@ -0,0 +1,1025 @@
+"""
+The linker and the event stream — plan steps C1 and C2, gate
+"recovers known trajectories, births, deaths and the merge exactly".
+
+Structure follows what the gates actually are:
+
+* :class:`TestTrajectoryGate` / :class:`TestEventGates` / :class:`TestDriftFrame`
+ run against ``particle_movie()`` end to end — segment, measure, link — because
+ the fixture is the acceptance gate for this wave and a linker validated only on
+ hand-written coordinates is validated against nothing that segmentation produces.
+* The remaining classes use small hand-built tables, where a two-particle
+ ambiguity can be constructed exactly and the assertion is unambiguous.
+
+One measured fact shapes several of these tests and is worth stating up front:
+**the frame at which a merge becomes detectable is a property of the segmenter,
+not of the linker.** The fixture's two discs geometrically overlap from frame 14
+(``ground_truth(sig)["merge_frame"]``), but the segmenter keeps them apart until
+frame 18 with watershed splitting on, and fuses them already at frame 12 with it
+off. A linker sees detections, not discs; the honest gate is that it reports the
+frame the detections became one, and that this frame brackets the geometric truth
+as the segmentation changes.
+"""
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+from spyde.data.synthetic import (
+ DISSOLUTION_INDEX,
+ NUCLEATION_INDEX,
+ ground_truth,
+ particle_movie,
+ particle_truth_at,
+)
+from spyde.drift.model import DriftModel
+from spyde.particles import SegmentParams, measure_frame, segment_frame
+from spyde.particles.track import (
+ EVENT_KINDS,
+ LinkParams,
+ LinkResult,
+ ParticleEvent,
+ event_counts,
+ events_from_records,
+ events_to_records,
+ frame_indices,
+ link,
+ sample_frame_positions,
+)
+from spyde.signals.particles import COL, N_COLUMNS, SpyDEParticles
+
+# The gate used for every fixture test. 3.0 nm = 6 px at the fixture's 0.5 nm/px,
+# i.e. ~3x the fastest particle's 2.2 px/frame lab-frame step. Deliberately not the
+# default 10.0: a test that only passes at a very loose gate would not notice a
+# linker that pairs by luck. TestGateWidth covers the default.
+FIXTURE_MAX_DIST = 3.0
+
+
+# ── helpers ──────────────────────────────────────────────────────────────────
+
+def _segment_movie(sig, gt, params=None, scale=None) -> SpyDEParticles:
+ """Run the real classical pipeline over every frame → a real particle table.
+
+ Deliberately not a synthetic table: the linker's input is whatever
+ ``segment_frame`` + ``measure_frame`` produce, including their misses (the two
+ faint probes are never found at default sensitivity) and their merges.
+ """
+ params = params or SegmentParams(min_size=25, gaussian=1.0)
+ scale = float(gt["scale"]) if scale is None else float(scale)
+ rows, contours = [], []
+ for t in range(int(gt["n_frames"])):
+ labels = segment_frame(sig.data[t], params)
+ r, c = measure_frame(labels, sig.data[t], t=t, scale=scale)
+ rows.append(r)
+ contours.append(c)
+ return SpyDEParticles.from_frames(
+ rows, frame_shape=tuple(gt["frame_shape"]), contours_per_frame=contours,
+ scale=scale, units="nm" if scale != 1.0 else "px")
+
+
+def _truth_of(res: LinkResult, gt, track: int, scale: float) -> tuple[int, float]:
+ """``(truth particle index, distance in px)`` nearest a track's first detection.
+
+ Maps a track id back onto the fixture's hand-written particle table so an
+ assertion can say "the mover" rather than "track 2", which is an id the linker
+ is free to renumber.
+ """
+ gi = int(res.track_first_index[track])
+ t = int(res.frame_index[gi])
+ pos = res.positions[gi] / scale
+ truth, _radii, present = particle_truth_at(gt, t)
+ d = np.hypot(*(truth - pos).T)
+ d[~present] = np.inf
+ return int(np.argmin(d)), float(d.min())
+
+
+def _track_for(res: LinkResult, gt, truth_index: int, scale: float) -> int:
+ matches = [tid for tid in range(res.n_tracks)
+ if _truth_of(res, gt, tid, scale)[0] == truth_index]
+ assert matches, f"no track corresponds to fixture particle {truth_index}"
+ return matches[0]
+
+
+def _table(frames, *, scale: float = 1.0, diameter: float = 4.0,
+ areas=None) -> SpyDEParticles:
+ """A minimal particle table from per-frame ``[(y, x), ...]`` lists.
+
+ Only the columns the linker reads are filled — positions, ``equiv_diameter``
+ (the adaptive merge radius) and optionally ``area`` (the similarity term). Every
+ other column stays 0, which is exactly the point: the linker must not depend on
+ anything a caller might not have measured.
+ """
+ blocks = []
+ for t, pts in enumerate(frames):
+ pts = np.asarray(pts, dtype=np.float64).reshape(-1, 2)
+ row = np.zeros((len(pts), N_COLUMNS), np.float32)
+ row[:, COL["t"]] = t
+ row[:, COL["label"]] = np.arange(1, len(pts) + 1)
+ row[:, COL["y"]] = pts[:, 0] * scale
+ row[:, COL["x"]] = pts[:, 1] * scale
+ row[:, COL["equiv_diameter"]] = diameter * scale
+ row[:, COL["track_id"]] = -1.0
+ if areas is not None:
+ row[:, COL["area"]] = np.asarray(areas[t], dtype=np.float32)
+ blocks.append(row)
+ return SpyDEParticles.from_frames(blocks, frame_shape=(128, 128), scale=scale,
+ units="nm" if scale != 1.0 else "px")
+
+
+@pytest.fixture(scope="module")
+def movie():
+ s = particle_movie()
+ return s, ground_truth(s)
+
+
+@pytest.fixture(scope="module")
+def linked(movie):
+ """Segment + measure + link the fixture once — shared by the gate classes."""
+ s, gt = movie
+ particles = _segment_movie(s, gt)
+ return s, gt, particles, link(particles, max_dist=FIXTURE_MAX_DIST)
+
+
+# ── gate 1: trajectories ─────────────────────────────────────────────────────
+
+class TestTrajectoryGate:
+ """"Recovers known trajectories exactly" — no fragmentation, sub-pixel error."""
+
+ def test_one_track_per_bright_particle(self, linked):
+ """7 tracks: 6 bright particles present at t=0, plus the nucleation.
+
+ The two faint probes are not detected by the classical engine at default
+ sensitivity (``test_particle_movie_fixture.py`` pins that), so the linker
+ cannot and must not invent tracks for them.
+ """
+ _s, gt, _p, res = linked
+ assert res.n_tracks == 7, (
+ f"expected 7 tracks, got {res.n_tracks} — a count above 7 means a "
+ f"trajectory fragmented (lengths {np.sort(res.track_lengths())})")
+
+ def test_the_mover_is_one_unbroken_track(self, linked):
+ s, gt, _p, res = linked
+ tid = _track_for(res, gt, 2, float(gt["scale"])) # the constant-velocity one
+ traj = res.trajectory(tid)
+ assert len(traj) == int(gt["n_frames"]), (
+ f"the mover's track spans {len(traj)} of {gt['n_frames']} frames")
+ assert np.array_equal(traj[:, 0], np.arange(int(gt["n_frames"]))), (
+ "the mover's track has a frame gap")
+
+ def test_mover_positions_match_the_truth(self, linked):
+ """Sub-pixel: the fixture draws analytic soft discs, so centroids are exact.
+
+ Measured with the classical segmenter: max error 0.175 px, RMS 0.073 px over
+ 24 frames. The 0.4 px bar is ~2x the measured worst case — tight enough that
+ an off-by-one frame (which would show as ~1.3 px on this particle) fails.
+ """
+ _s, gt, _p, res = linked
+ scale = float(gt["scale"])
+ tid = _track_for(res, gt, 2, scale)
+ err = []
+ for frame, y, x in res.trajectory(tid):
+ want = particle_truth_at(gt, int(frame))[0][2]
+ err.append(np.abs(np.array([y, x]) / scale - want))
+ err = np.asarray(err)
+ assert err.max() < 0.4, f"max trajectory error {err.max():.3f} px"
+
+ def test_every_bright_particle_gets_its_own_track(self, linked):
+ """No two tracks may map onto the same fixture particle."""
+ _s, gt, _p, res = linked
+ scale = float(gt["scale"])
+ owners = [_truth_of(res, gt, tid, scale)[0] for tid in range(res.n_tracks)]
+ assert len(set(owners)) == len(owners), f"two tracks share a particle: {owners}"
+
+ def test_track_lengths_match_the_event_timeline(self, linked):
+ """Lengths, not just count: 24, 24, 24, 24 persistent + 16 + 18 + 16."""
+ _s, gt, _p, res = linked
+ assert sorted(res.track_lengths().tolist()) == [16, 16, 18, 24, 24, 24, 24], (
+ f"unexpected track lengths {np.sort(res.track_lengths())}")
+
+ def test_trajectory_positions_agree_with_the_buffer(self, linked):
+ """A lab-frame link must report exactly the measured centroids."""
+ _s, _gt, particles, res = linked
+ stored = particles.flat_buffer[:, [COL["y"], COL["x"]]].astype(np.float64)
+ assert np.array_equal(res.positions, stored)
+ assert res.reference == "lab"
+
+
+# ── gates 2-4: events ────────────────────────────────────────────────────────
+
+class TestEventGates:
+ """Births, deaths and the merge against the fixture's stamped event frames."""
+
+ def test_exactly_one_nucleation_birth_at_the_known_frame(self, linked):
+ """Frame-0 detections DO count as births (LinkParams.initial_births).
+
+ So "the nucleation" is the one birth with ``frame > 0``, and the rest sit at
+ frame 0. Pinning both halves is what makes the assertion meaningful: a
+ linker that fragmented a track would add births at other frames.
+ """
+ _s, gt, _p, res = linked
+ births = res.events_of("birth")
+ later = [e for e in births if e.frame > 0]
+ assert len(later) == 1, f"expected 1 nucleation, got {[e.frame for e in later]}"
+ assert later[0].frame == int(gt["nucleation_frame"]) == 8
+ assert all(e.frame == 0 for e in births if e not in later)
+ assert len(births) == res.n_tracks, (
+ "every track must have exactly one birth — that is what makes the "
+ "event stream a complete description of the assignment")
+
+ def test_the_nucleation_birth_is_the_nucleating_particle(self, linked):
+ _s, gt, _p, res = linked
+ scale = float(gt["scale"])
+ birth = res.events_of("birth", exclude_initial=True)[0]
+ assert _truth_of(res, gt, birth.tracks[0], scale)[0] == NUCLEATION_INDEX
+
+ def test_exclude_initial_leaves_only_the_nucleation(self, linked):
+ _s, gt, _p, res = linked
+ got = res.events_of("birth", exclude_initial=True)
+ assert [e.frame for e in got] == [int(gt["nucleation_frame"])]
+
+ def test_initial_births_can_be_turned_off(self, linked):
+ _s, _gt, particles, _res = linked
+ res = link(particles, max_dist=FIXTURE_MAX_DIST, initial_births=False)
+ assert [e.frame for e in res.events_of("birth")] == [8]
+
+ def test_exactly_one_death_at_the_known_dissolution_frame(self, linked):
+ """The death frame is the first frame the particle is GONE — frame 16 here,
+ matching the fixture's ``death`` column, not the last frame it was seen."""
+ _s, gt, _p, res = linked
+ deaths = res.events_of("death")
+ assert len(deaths) == 1, f"expected 1 death, got {[e.frame for e in deaths]}"
+ assert deaths[0].frame == int(gt["dissolution_frame"]) == 16
+
+ def test_the_death_is_the_dissolving_particle(self, linked):
+ _s, gt, _p, res = linked
+ death = res.events_of("death")[0]
+ assert _truth_of(res, gt, death.tracks[0], float(gt["scale"]))[0] \
+ == DISSOLUTION_INDEX
+
+ def test_a_track_alive_in_the_final_frame_has_no_death(self, linked):
+ """The movie ending is not a dissolution."""
+ _s, _gt, _p, res = linked
+ ended = {e.tracks[0] for e in res.events if e.kind in ("death", "merge")}
+ for tid in range(res.n_tracks):
+ if int(res.track_last_frame[tid]) == res.n_frames - 1:
+ assert tid not in ended, f"track {tid} reaches the last frame yet ends"
+
+ def test_the_merge_is_detected_once_and_involves_the_merge_pair(self, linked):
+ _s, gt, _p, res = linked
+ scale = float(gt["scale"])
+ merges = res.events_of("merge")
+ assert len(merges) == 1, f"expected 1 merge, got {[e.frame for e in merges]}"
+ involved = {_truth_of(res, gt, t, scale)[0] for t in merges[0].tracks}
+ assert involved == set(int(i) for i in gt["merge_pair"]), (
+ f"merge involves fixture particles {involved}, expected "
+ f"{tuple(gt['merge_pair'])}")
+
+ def test_the_merge_lands_where_the_detections_actually_fuse(self, linked):
+ """The robust form of the merge-frame gate.
+
+ The geometric ``merge_frame`` (14) is when the discs overlap; the segmenter
+ keeps them as two detections until 18. So the invariant that holds for ANY
+ segmentation is that the merge event coincides with the frame the particle
+ count drops because two detections became one — which is also the only frame
+ a linker could possibly report.
+ """
+ _s, gt, particles, res = linked
+ merge = res.events_of("merge")[0]
+ counts = particles.count_series()
+ assert counts[merge.frame] == counts[merge.frame - 1] - 1, (
+ f"the merge at frame {merge.frame} is not where the count drops "
+ f"({counts.astype(int)})")
+ assert abs(merge.frame - int(gt["merge_frame"])) <= 4, (
+ f"merge reported at {merge.frame}, geometric truth "
+ f"{int(gt['merge_frame'])} — measured offset with this segmenter is +4")
+
+ def test_the_merge_frame_follows_the_segmenter_not_the_linker(self, movie):
+ """Watershed OFF fuses the pair 6 frames earlier, and the linker follows.
+
+ This is the honest statement of the limit: the same linker on the same movie
+ reports frame 18 with splitting on and frame 12 with it off, bracketing the
+ geometric truth at 14. Anything that claims to pin the merge to 14 exactly
+ is pinning a segmentation parameter, not the linker.
+ """
+ s, gt = movie
+ with_ws = link(_segment_movie(s, gt), max_dist=FIXTURE_MAX_DIST)
+ no_ws = link(_segment_movie(
+ s, gt, SegmentParams(min_size=25, gaussian=1.0, watershed=False)),
+ max_dist=FIXTURE_MAX_DIST)
+ f_on = with_ws.events_of("merge")[0].frame
+ f_off = no_ws.events_of("merge")[0].frame
+ assert f_off < int(gt["merge_frame"]) < f_on, (
+ f"expected the geometric merge frame {int(gt['merge_frame'])} to sit "
+ f"between the two segmentations ({f_off} and {f_on})")
+
+ def test_no_spurious_splits(self, linked):
+ """The fixture contains no fragmentation, so any split is a false positive.
+
+ Specifically: the nucleation at frame 8 appears 28 px from its nearest
+ neighbour, well outside that neighbour's ~12 px body, so it must read as a
+ birth and not as a split off it.
+ """
+ _s, _gt, _p, res = linked
+ assert res.events_of("split") == []
+
+ def test_every_event_kind_is_known_and_ordered(self, linked):
+ _s, _gt, _p, res = linked
+ assert all(e.kind in EVENT_KINDS for e in res.events)
+ frames = [e.frame for e in res.events]
+ assert frames == sorted(frames), "events are not in chronological order"
+
+ def test_event_lane_counts(self, linked):
+ """The C2 navigator lane: one trace per kind, spikes at the event frames."""
+ _s, gt, _p, res = linked
+ lane = res.event_counts()
+ assert set(lane) == set(EVENT_KINDS)
+ for trace in lane.values():
+ assert trace.shape == (int(gt["n_frames"]),)
+ assert lane["birth"][8] == 1 and lane["birth"][0] == 6
+ assert lane["death"][16] == 1 and lane["death"].sum() == 1
+ assert lane["merge"].sum() == 1
+ assert lane["split"].sum() == 0
+
+ def test_event_counts_ignores_out_of_range_frames(self):
+ lane = event_counts([ParticleEvent(9, "birth", (0,), (0,))], n_frames=4)
+ assert lane["birth"].sum() == 0
+
+
+# ── gate 5: frame of reference ───────────────────────────────────────────────
+
+class TestDriftFrame:
+ """Linking on lab and on drift-corrected coordinates, and what each means."""
+
+ def test_lab_frame_anchors_visibly_move(self, linked):
+ """The premise of the whole gate: without correction the static particles
+ travel with the stage, so a "did it move?" answer read off the lab frame is
+ wrong by the drift amplitude."""
+ _s, gt, _p, res = linked
+ scale = float(gt["scale"])
+ for anchor in (0, 1):
+ tid = _track_for(res, gt, anchor, scale)
+ tr = res.trajectory(tid)[:, 1:] / scale
+ excursion = float(np.hypot(*(tr.max(0) - tr.min(0))))
+ assert excursion > 5.0, (
+ f"anchor {anchor} only moves {excursion:.2f} px in the lab frame — "
+ "the fixture's drift is not being drawn")
+
+ def test_drift_corrected_anchors_are_stationary(self, linked):
+ """Measured: 9.3 px lab excursion collapses to 0.29-0.36 px corrected."""
+ _s, gt, particles, _res = linked
+ scale = float(gt["scale"])
+ model = DriftModel(shifts=np.asarray(gt["drift"]))
+ res = link(particles, max_dist=FIXTURE_MAX_DIST, drift=model)
+ assert res.reference == "sample"
+ for anchor in (0, 1):
+ tid = _track_for(res, gt, anchor, scale)
+ tr = res.trajectory(tid)[:, 1:] / scale
+ excursion = float(np.hypot(*(tr.max(0) - tr.min(0))))
+ assert excursion < 1.0, (
+ f"anchor {anchor} still moves {excursion:.2f} px after correction")
+
+ def test_a_solved_drift_model_works_as_well_as_the_truth(self, linked):
+ """End to end with A1's own output rather than the stamped truth, since
+ that is what the wizard will hand the linker."""
+ from spyde.drift import solve_translation
+ s, gt, particles, _res = linked
+ scale = float(gt["scale"])
+ model = solve_translation(s.data, device="numpy", upsample=8,
+ reference="first", max_shift=20)
+ res = link(particles, max_dist=FIXTURE_MAX_DIST, drift=model)
+ for anchor in (0, 1):
+ tid = _track_for(res, gt, anchor, scale)
+ tr = res.trajectory(tid)[:, 1:] / scale
+ assert float(np.hypot(*(tr.max(0) - tr.min(0)))) < 1.0
+
+ def test_the_mover_still_moves_after_correction(self, linked):
+ """Correction must remove the stage, not the sample. In the sample frame the
+ mover's displacement is its own constant velocity, 1.30 px/frame over 23
+ frames = 29.9 px."""
+ _s, gt, particles, _res = linked
+ scale = float(gt["scale"])
+ model = DriftModel(shifts=np.asarray(gt["drift"]))
+ res = link(particles, max_dist=FIXTURE_MAX_DIST, drift=model)
+ tid = _track_for(res, gt, 2, scale)
+ tr = res.trajectory(tid)[:, 1:] / scale
+ assert float(np.hypot(*(tr[-1] - tr[0]))) == pytest.approx(29.9, abs=1.0)
+
+ def test_both_references_find_the_same_tracks(self, linked):
+ """Same track set and same events either way — only the coordinates differ.
+
+ Track *ids* are not compared: the fixture's merge is symmetric, so which of
+ the pair survives is a genuine tie that the change of coordinates can flip
+ (documented in ``_nearest_continuing``).
+ """
+ _s, gt, particles, lab = linked
+ model = DriftModel(shifts=np.asarray(gt["drift"]))
+ sample = link(particles, max_dist=FIXTURE_MAX_DIST, drift=model)
+ assert sample.n_tracks == lab.n_tracks
+ assert np.array_equal(np.sort(sample.track_lengths()),
+ np.sort(lab.track_lengths()))
+ assert [(e.frame, e.kind) for e in sample.events] == \
+ [(e.frame, e.kind) for e in lab.events]
+
+ def test_calibration_is_applied_exactly_once(self):
+ """The trap this guards: ``DriftModel`` is in PIXELS, centroids are
+ calibrated. Adding the shifts straight onto ``y``/``x`` leaves the anchor
+ moving at ``1 - scale`` of its drift — plausible-looking and wrong."""
+ scale = 0.25
+ shift = np.array([[0.0, 0.0], [4.0, -3.0]])
+ # A particle that sits still in the SAMPLE frame therefore appears at
+ # -shift in the lab frame.
+ particles = _table([[(10.0, 20.0)], [(10.0 - 4.0, 20.0 + 3.0)]], scale=scale)
+ got = sample_frame_positions(particles, DriftModel(shifts=shift))
+ assert got == pytest.approx(np.array([[10.0, 20.0], [10.0, 20.0]]) * scale)
+
+ def test_a_short_drift_model_is_rejected(self):
+ particles = _table([[(1.0, 1.0)]] * 4)
+ with pytest.raises(ValueError, match="covers 2 frames"):
+ sample_frame_positions(particles, DriftModel(shifts=np.zeros((2, 2))))
+
+ def test_frame_indices_come_from_the_row_pointers(self, linked):
+ _s, _gt, particles, res = linked
+ assert np.array_equal(frame_indices(particles), res.frame_index)
+ assert np.array_equal(res.frame_index,
+ particles.column("t").astype(np.int32))
+
+
+# ── the gate itself ──────────────────────────────────────────────────────────
+
+class TestDistanceGate:
+ """An unmatched detection must be genuinely possible, not forced into a pair."""
+
+ def test_a_detection_beyond_the_gate_starts_a_new_track(self):
+ p = _table([[(0.0, 0.0)], [(0.0, 50.0)]])
+ res = link(p, max_dist=10.0)
+ assert res.n_tracks == 2
+ assert [(e.frame, e.kind) for e in res.events] == \
+ [(0, "birth"), (1, "birth"), (1, "death")]
+
+ def test_the_gate_is_not_relaxed_to_keep_cardinality(self):
+ """The failure mode a plain rectangular Hungarian has: with one track and one
+ detection it MUST return a pair, so the gate has to be enforced afterwards.
+ """
+ p = _table([[(0.0, 0.0)], [(0.0, 11.0)]])
+ assert link(p, max_dist=10.0).n_tracks == 2
+ assert link(p, max_dist=12.0).n_tracks == 1
+
+ def test_the_globally_cheapest_pairing_wins_over_the_greedy_one(self):
+ """Nearest-neighbour chaining gets this wrong; an assignment does not.
+
+ Costs (x only): A at 10, B at 13; detections at 5 and 11.
+
+ ====== ==== ====
+ \\ d=5 d=11
+ ====== ==== ====
+ A 5 1
+ B 8 2
+ ====== ==== ====
+
+ Greedy takes the globally smallest edge A->11 (1) and then has to put B on
+ 5 (8), total 9. The optimum is A->5 (5) plus B->11 (2), total 7. Every
+ per-track ``argmin`` scheme fails this; the assignment is why we can charge
+ one track more so the pair costs less.
+ """
+ p = _table([[(0.0, 10.0), (0.0, 13.0)],
+ [(0.0, 5.0), (0.0, 11.0)]])
+ res = link(p, max_dist=10.0)
+ assert res.n_tracks == 2
+ assert res.trajectory(0)[1, 2] == pytest.approx(5.0), (
+ "track A took its own nearest detection — this is a greedy match, not "
+ "an assignment")
+ assert res.trajectory(1)[1, 2] == pytest.approx(11.0)
+
+ def test_a_near_gate_pair_is_still_linked_rather_than_left_over(self):
+ """The other half of the gate contract: unmatched must be *possible*, not
+ *preferred*.
+
+ A sits on top of a detection (cost 0) while B's only admissible partner
+ costs 10 out of a 10.5 gate. Both links must still be made — a solver whose
+ sentinel were too small, or one that stopped at the obvious cheap pair,
+ would leave B unmatched and invent a death plus a birth.
+ """
+ p = _table([[(0.0, 0.0), (0.0, 20.0)], [(0.0, 0.0), (0.0, 10.0)]])
+ res = link(p, max_dist=10.5)
+ assert res.n_tracks == 2
+ assert res.events_of("death") == [] and \
+ res.events_of("birth", exclude_initial=True) == []
+
+ def test_a_frame_with_no_detections_does_not_link_across_it(self):
+ p = _table([[(0.0, 0.0)], [], [(0.0, 0.0)]])
+ res = link(p, max_dist=10.0, memory=0)
+ assert res.n_tracks == 2
+ assert [(e.frame, e.kind) for e in res.events] == \
+ [(0, "birth"), (1, "death"), (2, "birth")]
+
+ def test_max_dist_must_be_positive(self):
+ for bad in (0.0, -1.0, np.nan, np.inf):
+ with pytest.raises(ValueError, match="max_dist"):
+ LinkParams(max_dist=bad)
+
+
+class TestGateWidth:
+ """The measured usable window for ``max_dist`` on the fixture."""
+
+ @pytest.mark.parametrize("max_dist", [2.0, 3.0, 5.0, 10.0, 20.0, 40.0])
+ def test_a_wide_range_of_gates_gives_the_same_answer(self, linked, max_dist):
+ """Including the 10.0 default. 2.0 nm = 4 px is the lower edge."""
+ _s, _gt, particles, _res = linked
+ res = link(particles, max_dist=max_dist)
+ assert res.n_tracks == 7
+ assert len(res.events_of("death")) == 1
+ assert len(res.events_of("merge")) == 1
+
+ def test_too_tight_a_gate_fragments_and_is_visible_as_extra_deaths(self, linked):
+ """Establishes that the range above is not vacuous — and shows what a badly
+ set gate looks like, which is a wall of deaths rather than silence."""
+ _s, _gt, particles, _res = linked
+ res = link(particles, max_dist=0.6) # 1.2 px, below the true step
+ assert res.n_tracks > 7
+ assert len(res.events_of("death")) > 1
+
+
+# ── memory ───────────────────────────────────────────────────────────────────
+
+class TestMemory:
+ """A blinking detection must be ONE track, which is what memory is for."""
+
+ def test_memory_zero_splits_a_blinking_detection(self):
+ p = _table([[(0.0, 0.0)], [(0.0, 1.0)], [], [(0.0, 3.0)]])
+ res = link(p, max_dist=10.0, memory=0)
+ assert res.n_tracks == 2
+ assert [(e.frame, e.kind) for e in res.events] == \
+ [(0, "birth"), (2, "death"), (3, "birth")]
+
+ def test_memory_one_bridges_a_single_missing_frame(self):
+ p = _table([[(0.0, 0.0)], [(0.0, 1.0)], [], [(0.0, 3.0)]])
+ res = link(p, max_dist=10.0, memory=1)
+ assert res.n_tracks == 1
+ assert [e.kind for e in res.events] == ["birth"]
+
+ def test_memory_one_does_not_bridge_two_missing_frames(self):
+ p = _table([[(0.0, 0.0)], [], [], [(0.0, 3.0)]])
+ res = link(p, max_dist=10.0, memory=1)
+ assert res.n_tracks == 2
+
+ def test_memory_two_bridges_two_missing_frames(self):
+ p = _table([[(0.0, 0.0)], [], [], [(0.0, 3.0)]])
+ assert link(p, max_dist=10.0, memory=2).n_tracks == 1
+
+ def test_the_gate_is_measured_from_the_last_seen_position(self):
+ """Memory does not widen the search radius, so a fast particle that blinks
+ can still fall outside it. Documented behaviour, pinned here."""
+ p = _table([[(0.0, 0.0)], [], [(0.0, 12.0)]])
+ assert link(p, max_dist=10.0, memory=1).n_tracks == 2
+
+ def test_a_bridged_track_has_a_frame_gap_not_an_interpolated_row(self):
+ p = _table([[(0.0, 0.0)], [], [(0.0, 2.0)]])
+ res = link(p, max_dist=10.0, memory=1)
+ assert np.array_equal(res.trajectory(0)[:, 0], [0, 2])
+ assert res.track_lengths().tolist() == [2]
+ assert int(res.track_last_frame[0]) == 2
+
+ def test_memory_on_the_fixture_repairs_a_punched_hole(self, linked):
+ """The fixture has no dropouts, so one is introduced: delete the mover's
+ detection at frame 10. memory=0 fragments it into two tracks with a spurious
+ death; memory=1 recovers the single track and the spurious death vanishes.
+ """
+ _s, gt, particles, _res = linked
+ scale = float(gt["scale"])
+ want = particle_truth_at(gt, 10)[0][2]
+ rows, contours = [], []
+ for t in range(particles.n_frames):
+ blk = particles.at(t).copy()
+ keep = np.ones(len(blk), bool)
+ if t == 10:
+ d = np.hypot(*(blk[:, [COL["y"], COL["x"]]] / scale - want).T)
+ keep[int(np.argmin(d))] = False
+ rows.append(blk[keep])
+ contours.append([particles.contour_at(gi)
+ for gi, k in zip(particles.indices_at(t), keep) if k])
+ holed = SpyDEParticles.from_frames(
+ rows, frame_shape=particles.frame_shape, contours_per_frame=contours,
+ scale=scale, units="nm")
+
+ cold = link(holed, max_dist=FIXTURE_MAX_DIST, memory=0)
+ warm = link(holed, max_dist=FIXTURE_MAX_DIST, memory=1)
+ assert cold.n_tracks == 8 and len(cold.events_of("death")) == 2
+ assert warm.n_tracks == 7 and len(warm.events_of("death")) == 1
+ assert 10 not in [e.frame for e in warm.events_of("death")]
+
+ def test_memory_must_not_be_negative(self):
+ with pytest.raises(ValueError, match="memory"):
+ LinkParams(memory=-1)
+
+
+# ── property similarity ──────────────────────────────────────────────────────
+
+class TestPropertySimilarity:
+ """The optional weighting, and the promise that the GATE stays distance-only."""
+
+ def test_off_by_default(self):
+ assert LinkParams().property_weight == 0.0
+
+ def test_distance_alone_takes_the_wrong_pair_here(self):
+ """The setup: the distance-optimal pairing crosses the areas over.
+
+ A(area 100) at 0 and B(area 10) at 10; detections at 4 (area 10) and 6
+ (area 100). Distance prefers A->4, B->6 (total 8) over the area-consistent
+ A->6, B->4 (total 12), so this is a case where the property term has to
+ change the answer or it is doing nothing.
+ """
+ p = _table([[(0.0, 0.0), (0.0, 10.0)], [(0.0, 4.0), (0.0, 6.0)]],
+ areas=[[100.0, 10.0], [10.0, 100.0]])
+ res = link(p, max_dist=10.0, property_weight=0.0)
+ assert res.trajectory(0)[1, 2] == pytest.approx(4.0)
+
+ def test_a_weighted_property_flips_it(self):
+ """0.245 is the analytic crossover for this geometry; 0.5 clears it."""
+ p = _table([[(0.0, 0.0), (0.0, 10.0)], [(0.0, 4.0), (0.0, 6.0)]],
+ areas=[[100.0, 10.0], [10.0, 100.0]])
+ res = link(p, max_dist=10.0, property_weight=0.5, properties=("area",))
+ assert res.trajectory(0)[1, 2] == pytest.approx(6.0), (
+ "the property term did not change the assignment")
+ assert res.n_tracks == 2
+
+ def test_the_gate_stays_on_distance_alone(self):
+ """A perfect property match cannot pull a pair inside the gate, and a bad
+ one cannot push a pair out of it. Otherwise ``max_dist`` silently means
+ "max_dist minus however dissimilar these two look"."""
+ far = _table([[(0.0, 0.0)], [(0.0, 30.0)]], areas=[[50.0], [50.0]])
+ assert link(far, max_dist=10.0, property_weight=5.0).n_tracks == 2
+ near = _table([[(0.0, 0.0)], [(0.0, 9.0)]], areas=[[1000.0], [1.0]])
+ assert link(near, max_dist=10.0, property_weight=5.0,
+ properties=("area",)).n_tracks == 1
+
+ def test_a_missing_property_does_not_dilute_the_weight(self):
+ """``intensity_mean`` is NaN when nothing measured it. Averaging it in as
+ zero would halve a requested weight of 1.0 without saying so."""
+ p = _table([[(0.0, 0.0), (0.0, 10.0)], [(0.0, 4.0), (0.0, 6.0)]],
+ areas=[[100.0, 10.0], [10.0, 100.0]])
+ p.flat_buffer[:, COL["intensity_mean"]] = np.nan
+ only_area = link(p, max_dist=10.0, property_weight=0.5, properties=("area",))
+ with_nan = link(p, max_dist=10.0, property_weight=0.5,
+ properties=("area", "intensity_mean"))
+ assert np.array_equal(only_area.track_id, with_nan.track_id)
+ assert with_nan.trajectory(0)[1, 2] == pytest.approx(6.0)
+
+ def test_the_fixture_is_unaffected_by_the_weighting(self, linked):
+ """Positions on real segmented data are unambiguous, which is why the term
+ is off by default: it cannot help here and it can only add noise."""
+ _s, _gt, particles, res = linked
+ for w in (0.25, 1.0, 5.0):
+ other = link(particles, max_dist=FIXTURE_MAX_DIST, property_weight=w)
+ assert np.array_equal(other.track_id, res.track_id), f"weight {w}"
+
+ def test_unknown_property_is_rejected_at_construction(self):
+ with pytest.raises(KeyError, match="not_a_column"):
+ LinkParams(properties=("not_a_column",))
+
+ def test_negative_weight_is_rejected(self):
+ with pytest.raises(ValueError, match="property_weight"):
+ LinkParams(property_weight=-0.1)
+
+
+# ── merge / split rule ───────────────────────────────────────────────────────
+
+class TestMergeAndSplitRule:
+ """The post-pass, on geometries small enough to reason about exactly."""
+
+ def test_two_tracks_onto_one_detection_is_a_merge(self):
+ p = _table([[(0.0, 0.0), (0.0, 6.0)],
+ [(0.0, 0.5), (0.0, 5.5)],
+ [(0.0, 3.0)]], diameter=8.0)
+ res = link(p, max_dist=10.0)
+ merges = res.events_of("merge")
+ assert len(merges) == 1 and merges[0].frame == 2
+ assert set(merges[0].tracks) == {0, 1}
+ assert res.events_of("death") == [], "a merge must replace the death"
+
+ def test_one_track_into_two_detections_is_a_split(self):
+ p = _table([[(0.0, 3.0)],
+ [(0.0, 3.0)],
+ [(0.0, 0.5), (0.0, 5.5)]], diameter=8.0)
+ res = link(p, max_dist=10.0)
+ splits = res.events_of("split")
+ assert len(splits) == 1 and splits[0].frame == 2
+ assert res.events_of("birth", exclude_initial=True) == [], (
+ "a split must replace the birth")
+
+ def test_a_lone_disappearance_far_from_anything_is_a_death(self):
+ p = _table([[(0.0, 0.0), (0.0, 60.0)], [(0.0, 0.0)]], diameter=8.0)
+ res = link(p, max_dist=10.0)
+ assert [e.kind for e in res.events_of("death")] == ["death"]
+ assert res.events_of("merge") == []
+
+ def test_a_lone_appearance_far_from_anything_is_a_birth(self):
+ p = _table([[(0.0, 0.0)], [(0.0, 0.0), (0.0, 60.0)]], diameter=8.0)
+ res = link(p, max_dist=10.0)
+ assert len(res.events_of("birth", exclude_initial=True)) == 1
+ assert res.events_of("split") == []
+
+ def test_the_adaptive_radius_scales_with_the_particle(self):
+ """A disappearance 9 units from a survivor is a merge for a 20-unit body and
+ a death for a 4-unit one. A fixed radius cannot be right for both."""
+ frames = [[(0.0, 0.0), (0.0, 9.0)], [(0.0, 0.0)]]
+ assert link(_table(frames, diameter=20.0), max_dist=10.0).events_of("merge")
+ assert link(_table(frames, diameter=4.0), max_dist=10.0).events_of("death")
+
+ def test_merge_dist_can_be_overridden(self):
+ frames = [[(0.0, 0.0), (0.0, 9.0)], [(0.0, 0.0)]]
+ p = _table(frames, diameter=4.0)
+ assert link(p, max_dist=10.0, merge_dist=20.0).events_of("merge")
+ assert link(p, max_dist=10.0, merge_dist=2.0).events_of("death")
+
+ def test_split_dist_can_be_overridden(self):
+ frames = [[(0.0, 0.0)], [(0.0, 0.0), (0.0, 9.0)]]
+ p = _table(frames, diameter=4.0)
+ assert link(p, max_dist=10.0, split_dist=20.0).events_of("split")
+ assert link(p, max_dist=10.0, split_dist=2.0).events_of(
+ "birth", exclude_initial=True)
+
+ def test_a_zero_diameter_does_not_disable_the_post_pass(self):
+ """The floor. An unmeasured ``equiv_diameter`` would otherwise give a zero
+ radius, so no merge could ever be detected and the failure would be silent.
+ """
+ p = _table([[(0.0, 0.0), (0.0, 1.0)], [(0.0, 0.5)]], diameter=0.0)
+ assert link(p, max_dist=10.0).events_of("merge")
+
+ def test_the_floor_is_in_pixels_not_calibrated_units(self):
+ """A 2 px floor must stay 2 px when the axis is calibrated at 0.1 nm/px,
+ otherwise the post-pass gets 20x looser on a finer calibration."""
+ frames = [[(0.0, 0.0), (0.0, 1.5)], [(0.0, 0.75)]]
+ for scale in (0.1, 1.0, 10.0):
+ res = link(_table(frames, diameter=0.0, scale=scale), max_dist=100.0)
+ assert res.events_of("merge"), f"scale {scale}"
+
+ def test_a_disappearance_next_to_a_NEWCOMER_is_not_a_merge(self):
+ """The survivor must have existed the frame before. Without that
+ requirement, a death happening to coincide with an unrelated birth nearby
+ would be reported as a merge into it."""
+ # Frame 3 keeps the far particle and introduces a NEW detection 4 units from
+ # where the first one died — inside the 8-unit merge radius but outside the
+ # 2-unit link gate, so the first track genuinely ends and the newcomer is
+ # genuinely new. The buffer order is reversed at frame 3 so that row order
+ # cannot stand in for the "existed before" test.
+ p = _table([[(0.0, 0.0), (0.0, 40.0)],
+ [(0.0, 1.0), (0.0, 40.0)],
+ [(0.0, 2.0), (0.0, 40.0)],
+ [(0.0, 40.0), (0.0, 6.0)]], diameter=8.0)
+ res = link(p, max_dist=2.0)
+ assert [(e.frame, e.kind) for e in res.events_of("death")] == [(3, "death")]
+ assert res.events_of("merge") == []
+ assert [(e.frame, e.kind)
+ for e in res.events_of("birth", exclude_initial=True)] == [(3, "birth")]
+ assert res.events_of("split") == [], (
+ "the newcomer is 34 units from the only continuing track, so it cannot "
+ "be a fragment of it either")
+
+ def test_events_survive_a_long_uneventful_stretch(self):
+ """The post-pass skips building its per-frame lookup on frames that no
+ candidate event needs, which is most of a long stable movie. A merge is the
+ case that would break if the skip were off by one, because it is the only
+ event that reads the PREVIOUS frame's map — so nothing happens for seven
+ frames and then two detections become one.
+ """
+ frames = [[(0.0, 0.0), (0.0, 6.0)]] * 8 + [[(0.0, 3.0)]]
+ res = link(_table(frames, diameter=8.0), max_dist=10.0)
+ merges = res.events_of("merge")
+ assert len(merges) == 1 and merges[0].frame == 8
+ assert set(merges[0].tracks) == {0, 1}
+
+ def test_a_late_split_survives_the_same_skip(self):
+ frames = [[(0.0, 3.0)]] * 8 + [[(0.0, 0.5), (0.0, 5.5)]]
+ res = link(_table(frames, diameter=8.0), max_dist=10.0)
+ assert [(e.frame, e.kind) for e in res.events_of("split")] == [(8, "split")]
+
+ def test_merge_event_particle_indices_point_at_real_rows(self):
+ p = _table([[(0.0, 0.0), (0.0, 6.0)],
+ [(0.0, 0.5), (0.0, 5.5)],
+ [(0.0, 3.0)]], diameter=8.0)
+ res = link(p, max_dist=10.0)
+ e = res.events_of("merge")[0]
+ assert len(e.particles) == len(e.tracks) == 2
+ absorbed, survivor = e.particles
+ assert int(res.frame_index[absorbed]) == e.frame - 1, (
+ "the absorbed track has no row at the merge frame — its last one is the "
+ "only row that can identify it")
+ assert int(res.frame_index[survivor]) == e.frame
+ assert res.track_id[survivor] == e.tracks[1]
+
+
+# ── determinism, degenerate input, bookkeeping ───────────────────────────────
+
+class TestDeterminism:
+ def test_the_same_table_links_identically_twice(self, linked):
+ _s, _gt, particles, res = linked
+ again = link(particles, max_dist=FIXTURE_MAX_DIST)
+ assert np.array_equal(again.track_id, res.track_id)
+ assert again.events == res.events
+
+ def test_a_rebuilt_table_links_identically(self, movie):
+ """Not just the same object twice — the same pixels through the whole
+ pipeline, so a dict- or set-ordering dependency anywhere would show."""
+ s, gt = movie
+ a = link(_segment_movie(s, gt), max_dist=FIXTURE_MAX_DIST)
+ b = link(_segment_movie(s, gt), max_dist=FIXTURE_MAX_DIST)
+ assert np.array_equal(a.track_id, b.track_id)
+ assert a.events == b.events
+
+ def test_ids_are_contiguous_from_zero(self, linked):
+ _s, _gt, _p, res = linked
+ assert res.track_id.min() == 0
+ assert np.array_equal(np.unique(res.track_id), np.arange(res.n_tracks))
+
+ def test_ids_are_ordered_by_first_appearance(self, linked):
+ """Which is why they are stable: the numbering depends on the buffer order,
+ not on iteration order anywhere in the loop."""
+ _s, _gt, _p, res = linked
+ assert np.all(np.diff(res.track_first_frame) >= 0)
+ assert np.all(np.diff(res.track_first_index) > 0)
+
+ def test_a_symmetric_tie_is_broken_reproducibly(self):
+ """Two tracks exactly equidistant from one detection: whichever survives, it
+ must be the same one every run."""
+ p = _table([[(0.0, 0.0), (0.0, 4.0)], [(0.0, 2.0)]], diameter=8.0)
+ first = link(p, max_dist=10.0)
+ for _ in range(5):
+ other = link(p, max_dist=10.0)
+ assert np.array_equal(other.track_id, first.track_id)
+ assert other.events == first.events
+
+
+class TestDegenerate:
+ def test_no_frames_at_all(self):
+ res = link(SpyDEParticles.from_frames([], frame_shape=(8, 8)))
+ assert res.n_tracks == 0 and res.events == [] and res.n_frames == 0
+ assert res.positions.shape == (0, 2)
+
+ def test_every_frame_empty(self):
+ p = SpyDEParticles.from_frames([np.zeros((0, N_COLUMNS), np.float32)] * 5,
+ frame_shape=(8, 8))
+ res = link(p)
+ assert res.n_tracks == 0 and res.events == []
+ assert res.event_counts()["birth"].shape == (5,)
+
+ def test_one_empty_frame_in_the_middle(self):
+ p = _table([[(0.0, 0.0)], [], [(0.0, 0.0)]])
+ res = link(p, max_dist=5.0)
+ assert res.n_tracks == 2
+ assert [e.frame for e in res.events_of("death")] == [1]
+
+ def test_a_single_frame_gives_births_and_no_deaths(self):
+ """Nothing to link, and the movie ending is not a dissolution — so a
+ one-frame table is all births."""
+ res = link(_table([[(0.0, 0.0), (5.0, 5.0)]]))
+ assert res.n_tracks == 2
+ assert [e.kind for e in res.events] == ["birth", "birth"]
+
+ def test_everything_vanishes_mid_movie(self):
+ p = _table([[(0.0, 0.0), (0.0, 20.0)],
+ [(0.0, 0.0), (0.0, 20.0)],
+ [], [], []])
+ res = link(p, max_dist=5.0)
+ assert res.n_tracks == 2
+ assert sorted(e.frame for e in res.events_of("death")) == [2, 2]
+ assert res.events_of("merge") == []
+
+ def test_a_particle_appearing_only_in_the_last_frame(self):
+ p = _table([[(0.0, 0.0)], [(0.0, 0.0)], [(0.0, 0.0), (0.0, 40.0)]])
+ res = link(p, max_dist=5.0)
+ assert res.n_tracks == 2
+ assert [e.frame for e in res.events_of("birth", exclude_initial=True)] == [2]
+ assert res.events_of("death") == []
+
+ def test_a_table_that_never_got_masks_still_links(self):
+ """``store_masks=False`` is the default for long movies, so the linker must
+ not need contours."""
+ p = _table([[(0.0, 0.0)], [(0.0, 1.0)]])
+ assert p.contours is None
+ assert link(p, max_dist=5.0).n_tracks == 1
+
+
+class TestResultBookkeeping:
+ def test_apply_writes_the_track_id_column(self, linked):
+ _s, _gt, particles, res = linked
+ copy = SpyDEParticles(flat_buffer=particles.flat_buffer.copy(),
+ t_offsets=particles.t_offsets.copy(),
+ frame_shape=particles.frame_shape,
+ scale=particles.scale, units=particles.units)
+ assert not copy.has_tracks
+ res.apply(copy)
+ assert copy.has_tracks
+ assert np.array_equal(copy.column("track_id").astype(np.int32), res.track_id)
+
+ def test_link_does_not_mutate_unless_asked(self, linked):
+ _s, _gt, particles, _res = linked
+ before = particles.column("track_id").copy()
+ link(particles, max_dist=FIXTURE_MAX_DIST)
+ assert np.array_equal(particles.column("track_id"), before)
+
+ def test_apply_true_writes_through(self):
+ p = _table([[(0.0, 0.0)], [(0.0, 1.0)]])
+ res = link(p, max_dist=5.0, apply=True)
+ assert np.array_equal(p.column("track_id").astype(np.int32), res.track_id)
+
+ def test_apply_rejects_a_mismatched_table(self, linked):
+ _s, _gt, _p, res = linked
+ with pytest.raises(ValueError, match="different link"):
+ res.apply(_table([[(0.0, 0.0)]]))
+
+ def test_track_indices_are_chronological(self, linked):
+ _s, _gt, _p, res = linked
+ for tid in range(res.n_tracks):
+ idx = res.track_indices(tid)
+ assert np.all(res.track_id[idx] == tid)
+ assert np.all(np.diff(res.frame_index[idx]) > 0)
+
+ def test_track_indices_cover_every_particle_exactly_once(self, linked):
+ _s, _gt, _p, res = linked
+ seen = np.concatenate([res.track_indices(t) for t in range(res.n_tracks)])
+ assert np.array_equal(np.sort(seen), np.arange(res.n_particles))
+
+ def test_track_endpoints_agree_with_the_trajectories(self, linked):
+ _s, _gt, _p, res = linked
+ for tid in range(res.n_tracks):
+ traj = res.trajectory(tid)
+ assert traj[0, 0] == res.track_first_frame[tid]
+ assert traj[-1, 0] == res.track_last_frame[tid]
+ assert np.array_equal(traj[0, 1:], res.positions[res.track_first_index[tid]])
+ assert np.array_equal(traj[-1, 1:], res.positions[res.track_last_index[tid]])
+
+ def test_track_at_matches_a_boolean_scan(self, linked):
+ """`track_at` uses searchsorted for speed; it must agree with the naive form."""
+ _s, _gt, _p, res = linked
+ for t in range(res.n_frames):
+ assert np.array_equal(res.track_at(t),
+ res.track_id[res.frame_index == t])
+
+ def test_track_lengths_count_detections_not_span(self):
+ p = _table([[(0.0, 0.0)], [], [(0.0, 1.0)]])
+ res = link(p, max_dist=5.0, memory=1)
+ assert res.track_lengths().tolist() == [2]
+ assert int(res.track_last_frame[0]) - int(res.track_first_frame[0]) == 2
+
+ def test_out_of_range_track_raises(self, linked):
+ _s, _gt, _p, res = linked
+ with pytest.raises(IndexError):
+ res.track_indices(res.n_tracks)
+
+ def test_unknown_event_kind_raises(self, linked):
+ _s, _gt, _p, res = linked
+ with pytest.raises(ValueError, match="unknown event kind"):
+ res.events_of("explosion")
+
+ def test_unknown_link_parameter_raises(self, linked):
+ _s, _gt, particles, _res = linked
+ with pytest.raises(TypeError, match="max_distance"):
+ link(particles, max_distance=5.0)
+
+ def test_params_and_kwargs_compose(self, linked):
+ _s, _gt, particles, _res = linked
+ res = link(particles, LinkParams(max_dist=99.0), max_dist=FIXTURE_MAX_DIST)
+ assert res.params["max_dist"] == FIXTURE_MAX_DIST
+
+ def test_params_are_recorded_for_provenance(self, linked):
+ _s, _gt, _p, res = linked
+ assert res.params["max_dist"] == FIXTURE_MAX_DIST
+ assert res.params["memory"] == 0
+ assert res.params["merge_dist"] is None
+
+ def test_repr_names_the_reference_and_the_event_counts(self, linked):
+ _s, _gt, _p, res = linked
+ text = repr(res)
+ assert "reference='lab'" in text and "'merge': 1" in text
+
+
+class TestEventSerialisation:
+ def test_round_trips_through_records(self, linked):
+ _s, _gt, _p, res = linked
+ assert events_from_records(events_to_records(res.events)) == res.events
+
+ def test_records_are_json_safe(self, linked):
+ import json
+ _s, _gt, _p, res = linked
+ text = json.dumps(events_to_records(res.events))
+ assert json.loads(text)[0]["kind"] == "birth"
+
+ def test_result_to_dict_is_json_safe(self, linked):
+ import json
+ _s, _gt, _p, res = linked
+ d = json.loads(json.dumps(res.to_dict()))
+ assert d["n_tracks"] == 7 and d["reference"] == "lab"
+ assert len(d["events"]) == len(res.events)
+
+ def test_events_are_hashable_and_immutable(self, linked):
+ """Three surfaces share this record (navigator lane, table, report embed),
+ so one of them must not be able to rewrite it."""
+ import dataclasses
+ _s, _gt, _p, res = linked
+ assert len(set(res.events)) == len(res.events)
+ with pytest.raises(dataclasses.FrozenInstanceError):
+ res.events[0].frame = 3
+
+ def test_params_to_dict_is_json_safe(self):
+ import json
+ assert json.loads(json.dumps(LinkParams().to_dict()))["memory"] == 0
diff --git a/spyde/tests/migrated/test_particles_wizard.py b/spyde/tests/migrated/test_particles_wizard.py
new file mode 100644
index 00000000..d3087c8a
--- /dev/null
+++ b/spyde/tests/migrated/test_particles_wizard.py
@@ -0,0 +1,1436 @@
+"""
+The Segment Particles wizard backend (``seg_*`` staged handlers).
+
+Handlers are called directly as ``fn(session, plot, payload)`` and polled with
+``_wait`` — the shape ``test_find_vectors_wizard.py`` establishes for staged
+actions, because everything heavy here goes onto a worker thread.
+
+Four claims are worth more than the rest, and they are the reason this file
+exists rather than a smoke test:
+
+:class:`TestPreviewIsOneFrame`
+ Plan §0.8.1. A tune must read exactly the frame the navigator is on. The
+ guard is the CLAUDE.md memory-safety one: ``da.Array.compute`` must never
+ see the full movie shape.
+:class:`TestMinSizeFloor`
+ Plan §0.9. ``min_size=0`` is a footgun (measured: 33 instances where 9 are
+ real), so it is floored — AND the effective value is reported, because a
+ caret whose number disagrees with what ran is the failure
+ ``SegmentParams.local_size`` refuses to introduce.
+:class:`TestRunIsProgressiveAndCancellable`
+ Plan §0.8.2-3 and the attach gap: the result window opens EARLY with no
+ particles attached, ``_seg_batch_running`` is up for the duration (that is
+ exactly what ``lifecycle.wait_for_particles`` polls), and ``tree.particles``
+ lands only at finalize.
+:class:`TestDoubleFire`
+ README §4 / StrictMode: open, close, open leaves exactly ONE controller.
+
+Torch runs on **CPU explicitly** wherever the scribble head is trained:
+torch-CUDA work segfaults under the pytest process on Windows (CLAUDE.md).
+"""
+from __future__ import annotations
+
+import time
+
+import numpy as np
+import pytest
+
+from spyde.actions import particles_action as pa
+
+
+@pytest.fixture(autouse=True)
+def _capture_module_emit(window, monkeypatch):
+ """Route ``particles_action``'s own ``emit`` into the captured list.
+
+ The module does ``from spyde.backend.ipc import emit`` at import, so
+ conftest's patch of ``ipc.emit`` never reaches that binding — the identical
+ hazard conftest already documents for ``session.py``, and the identical fix.
+ ``emit_status``/``emit_error`` need no patch: they resolve ``emit`` inside
+ ``ipc`` at call time.
+ """
+ monkeypatch.setattr(pa, "emit", window["messages"].append)
+
+
+def _signal_plot(session):
+ return next((p for p in session._plots
+ if not p.is_navigator and p.plot_state is not None), None)
+
+
+def _wait(pred, timeout=60.0):
+ end = time.time() + timeout
+ while time.time() < end:
+ if pred():
+ return True
+ time.sleep(0.05)
+ return False
+
+
+def _movie(window, frames: int = 6):
+ """The synthetic particle movie through the door the e2e specs use."""
+ session = window["window"]
+ session._load_test_data_particles({"frames": frames})
+ plot = _wait(lambda: _signal_plot(session) is not None) and _signal_plot(session)
+ assert plot is not None, "the particle movie never produced a signal plot"
+ return session, plot, plot.signal_tree
+
+
+def _opened(window, frames: int = 6, **params):
+ session, plot, tree = _movie(window, frames)
+ pa.seg_open(session, plot, {"min_size": 25, "gaussian": 1.0, **params})
+ assert _wait(lambda: getattr(tree, "_seg_wizard", None) is not None
+ and tree._seg_wizard.preview is not None), \
+ "the wizard never produced a first preview"
+ return session, plot, tree, tree._seg_wizard
+
+
+def _of_type(messages, kind):
+ return [m for m in messages if isinstance(m, dict) and m.get("type") == kind]
+
+
+class _FullComputeGuard:
+ """Raise the count if ``.compute()`` is ever called on the whole movie."""
+
+ def __init__(self, shape):
+ self.shape = tuple(shape)
+ self.hits = 0
+
+ def __enter__(self):
+ import dask.array as da
+ self._real = da.Array.compute
+ guard = self
+
+ def _spy(arr, *a, **k):
+ if tuple(arr.shape) == guard.shape:
+ guard.hits += 1
+ return guard._real(arr, *a, **k)
+
+ da.Array.compute = _spy
+ return self
+
+ def __exit__(self, *exc):
+ import dask.array as da
+ da.Array.compute = self._real
+ return False
+
+
+class TestPreviewIsOneFrame:
+ def test_open_builds_a_controller_and_previews(self, window):
+ _s, _p, tree, wiz = _opened(window)
+ assert wiz is getattr(tree, "_seg_wizard")
+ assert wiz.preview["frame"] == wiz.frame_index()
+ assert wiz.preview["count"] > 0, "found nothing on the fixture's frame"
+
+ def test_preview_never_computes_the_whole_movie(self, window):
+ session, plot, tree, wiz = _opened(window)
+ assert tree.root._lazy, "the fixture must be lazy or this guards nothing"
+ with _FullComputeGuard(tree.root.data.shape) as guard:
+ pa.seg_tune(session, plot, {"sensitivity": 0.6})
+ assert _wait(lambda: abs(wiz.params["sensitivity"] - 0.6) < 1e-9)
+ time.sleep(0.6)
+ assert guard.hits == 0, "a tune materialised the whole movie"
+
+ def test_tune_opens_no_new_window(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ before = len(session.signal_trees)
+ pa.seg_tune(session, plot, {"sensitivity": 0.7})
+ assert _wait(lambda: abs(wiz.params["sensitivity"] - 0.7) < 1e-9)
+ time.sleep(0.4)
+ assert len(session.signal_trees) == before, "a tune spawned a tree"
+
+ def test_preview_message_carries_the_size_histogram(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_tune(session, plot, {"sensitivity": 0.55})
+ assert _wait(lambda: len(_of_type(msgs, "seg_preview")) >= 2)
+ msg = _of_type(msgs, "seg_preview")[-1]
+ assert msg["count"] == len(msg["areas"]) or msg["count"] > pa._MAX_AREAS_SENT
+ assert msg["median_area"] > 0 and msg["units"] == "nm"
+
+ def test_set_method_to_an_untrained_scribble_keeps_the_last_preview(self, window):
+ """The scribble engine cannot run before Train — it must say so, not
+ clear the result the user is looking at."""
+ session, plot, _tree, wiz = _opened(window)
+ keep = wiz.preview
+ pa.seg_set_method(session, plot, {"method": "scribble"})
+ time.sleep(0.4)
+ assert wiz.params["method"] == "scribble"
+ assert wiz.preview is keep
+
+ def test_prompt_engine_is_an_explicit_stub(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_set_method(session, plot, {"method": "prompt"})
+ assert _wait(lambda: any("not installed yet" in str(m.get("text", ""))
+ for m in _of_type(msgs, "status")))
+
+
+class TestMinSizeFloor:
+ """Plan §0.9 — the measured finding, not a preference."""
+
+ def test_zero_is_floored(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ pa.seg_tune(session, plot, {"min_size": 0})
+ assert _wait(lambda: wiz.params["min_size"] == pa.MIN_SIZE_FLOOR)
+ assert wiz.params["min_size_floored"] is True
+
+ def test_the_effective_value_is_reported(self, window):
+ """Silently running a different number than the caret shows is the
+ failure mode; the floor must come back in the preview payload."""
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_tune(session, plot, {"min_size": 0})
+ assert _wait(lambda: any(m.get("min_size_floored")
+ for m in _of_type(msgs, "seg_preview")))
+ msg = next(m for m in _of_type(msgs, "seg_preview")
+ if m.get("min_size_floored"))
+ assert msg["min_size"] == pa.MIN_SIZE_FLOOR
+
+ def test_a_real_min_size_is_left_alone(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ pa.seg_tune(session, plot, {"min_size": 40})
+ assert _wait(lambda: wiz.params["min_size"] == 40)
+ assert wiz.params["min_size_floored"] is False
+
+ def test_reported_count_is_after_the_size_filter(self, window):
+ """§0.9b: the number the user sees must be the post-filter one, or it
+ moves for a reason they cannot see."""
+ session, plot, _tree, wiz = _opened(window)
+ from spyde.signals.particles import COL
+ pa.seg_tune(session, plot, {"min_size": 400})
+ assert _wait(lambda: wiz.params["min_size"] == 400)
+ assert _wait(lambda: wiz.preview is not None
+ and (wiz.preview["rows"].size == 0
+ or wiz.preview["rows"][:, COL["area"]].min() > 0))
+ rows = wiz.preview["rows"]
+ assert wiz.preview["count"] == len(rows)
+
+
+class TestScribbleLabelling:
+ def test_paint_accumulates_across_frames(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ pa.seg_paint(session, plot, {"frame": 0, "points": [[20, 20], [24, 24]],
+ "class_id": 0, "brush": 3})
+ pa.seg_paint(session, plot, {"frame": 3, "points": [[40, 40], [44, 44]],
+ "class_id": 1, "brush": 3})
+ assert wiz.labels.labelled_frames() == [0, 3]
+ counts = wiz.labels.counts()
+ assert counts[0] > 0 and counts[1] > 0
+
+ def test_state_reports_per_class_pixel_counts(self, window):
+ """Plan B3: under-training a class is *the* failure mode and these
+ counts are how a user notices, so every class is present."""
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_paint(session, plot, {"frame": 0, "points": [[20, 20], [22, 22]],
+ "class_id": 0, "brush": 3})
+ state = _of_type(msgs, "seg_state")[-1]
+ by_id = {c["id"]: c for c in state["classes"]}
+ assert by_id[0]["pixels"] > 0
+ assert by_id[2]["pixels"] == 0, "an unpainted class must still be listed"
+
+ def test_erase_covers_exactly_what_the_brush_painted(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ stroke = {"frame": 0, "points": [[20, 20], [30, 30]], "brush": 5}
+ pa.seg_paint(session, plot, {**stroke, "class_id": 0})
+ painted = len(wiz.labels)
+ assert painted > 0
+ pa.seg_paint(session, plot, {**stroke, "erase": True})
+ assert len(wiz.labels) == 0, (
+ f"{len(wiz.labels)} of {painted} px survived an erase over the same "
+ "stroke — the eraser and the brush have drifted apart")
+
+ def test_paint_with_no_points_is_a_no_op(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ pa.seg_paint(session, plot, {"frame": 0, "points": [], "class_id": 0})
+ assert wiz.labels is None
+
+ def test_train_needs_labels(self, window):
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_train(session, plot, {"device": "cpu"})
+ assert any("nothing painted" in str(m.get("text", ""))
+ for m in _of_type(msgs, "error"))
+
+ @pytest.mark.slow
+ def test_train_then_the_scribble_engine_previews(self, window):
+ """The full B3 loop on CPU: scribble a particle and some background,
+ train, and the scribble engine becomes the live preview."""
+ session, plot, _tree, wiz = _opened(window)
+ # A bright particle in the fixture and a patch of bare film.
+ pa.seg_paint(session, plot, {"frame": 0, "class_id": 0, "brush": 3,
+ "points": [[24, 28], [25, 29]]})
+ pa.seg_paint(session, plot, {"frame": 0, "class_id": 1, "brush": 5,
+ "points": [[5, 5], [5, 60], [8, 100]]})
+ pa.seg_train(session, plot, {"device": "cpu"})
+ assert _wait(lambda: wiz.classifier is not None
+ and wiz.classifier.is_trained, timeout=180)
+ assert wiz.params["method"] == "scribble"
+ assert _wait(lambda: wiz.preview is not None
+ and wiz.preview["frame"] == wiz.frame_index(), timeout=120)
+
+
+class TestRunIsProgressiveAndCancellable:
+ def test_result_window_opens_before_the_particles_attach(self, window):
+ """The attach gap, asserted at the instant it exists."""
+ session, plot, tree, _wiz = _opened(window)
+ before = len(session.signal_trees)
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ assert len(session.signal_trees) == before + 1, \
+ "seg_run must open its result window synchronously"
+ result = session.signal_trees[-1]
+ assert getattr(result, "particles", None) is None, (
+ "particles attached at open — requires_particles would unlock "
+ "against an empty store")
+ assert result._seg_batch_running and tree._seg_batch_running
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+
+ def test_seg_batch_running_is_what_lifecycle_polls(self, window):
+ from spyde.actions.lifecycle import seg_batch_running
+ session, plot, _tree, _wiz = _opened(window)
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ assert seg_batch_running(session)
+ result = session.signal_trees[-1]
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+ assert _wait(lambda: not seg_batch_running(session))
+
+ def test_finalize_attaches_a_full_store_and_says_how_many(self, window):
+ session, plot, _tree, _wiz = _opened(window, frames=6)
+ msgs = window["messages"]
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ result = session.signal_trees[-1]
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+ parts = result.particles
+ assert parts.n_frames == 6 and parts.n_particles > 0
+ assert parts.has_masks
+ # Polled, not read once: the status is the LAST thing _finalize does,
+ # after the count-trace paint, so it lands a beat after the attach.
+ assert _wait(lambda: any(
+ f"Found {parts.n_particles} particles" in str(m.get("text", ""))
+ for m in _of_type(msgs, "status"))), (
+ "the finalize status must carry the count — the e2e specs and the "
+ "user both read it")
+
+ def test_the_label_movie_renders_the_real_contours(self, window):
+ """The placeholder is mutated IN PLACE precisely so this works; a fresh
+ store would leave the early window rendering zeros forever."""
+ session, plot, _tree, _wiz = _opened(window)
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ result = session.signal_trees[-1]
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+ frame = np.asarray(result.root.data[3].compute())
+ painted = np.unique(frame)
+ painted = painted[painted > 0]
+ assert painted.size == len(result.particles.at(3)) > 0
+
+ def test_count_trace_reaches_the_navigator(self, window):
+ session, plot, _tree, _wiz = _opened(window)
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ result = session.signal_trees[-1]
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+ counts = result.particles.count_series()
+ navs = [n for n in pa._nav_plots(result)
+ if getattr(getattr(n, "current_data", None), "ndim", 0) == 1]
+ assert navs, "the particle tree has no 1-D navigator to fill"
+ assert any(np.array_equal(np.asarray(n.current_data, np.float32), counts)
+ for n in navs), "the count trace never reached the navigator"
+
+ def test_run_never_computes_the_whole_movie(self, window):
+ session, plot, tree, _wiz = _opened(window)
+ with _FullComputeGuard(tree.root.data.shape) as guard:
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ result = session.signal_trees[-1]
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+ assert guard.hits == 0, "the batch materialised the whole movie"
+
+ def test_closing_the_result_tree_cancels_the_batch(self, window, monkeypatch):
+ """Cancellation runs through BaseSignalTree.register_cancel, so the
+ ordinary act of closing the window has to stop the compute."""
+ import spyde.particles as sp
+ real = sp.segment_frame
+
+ def _slow(frame, params=None):
+ time.sleep(0.25)
+ return real(frame, params)
+
+ monkeypatch.setattr(sp, "segment_frame", _slow)
+ session, plot, _tree, _wiz = _opened(window, frames=24)
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ result = session.signal_trees[-1]
+ time.sleep(0.5)
+ result.close()
+ assert _wait(lambda: not result._seg_batch_running, timeout=30)
+ # A cancelled run keeps its partial rows out of the torn-down tree.
+ assert getattr(result, "particles", None) is None
+
+ def test_finalize_re_sends_the_toolbar_config(self, window):
+ """requires_particles flips here; without the re-send the gated buttons
+ stay hidden until something else rebuilds the toolbar."""
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_run(session, plot, {"min_size": 25, "gaussian": 1.0})
+ result = session.signal_trees[-1]
+ assert _wait(lambda: getattr(result, "particles", None) is not None)
+ wids = {getattr(sp_, "window_id", None) for sp_ in result.signal_plots}
+ assert _wait(lambda: any(m.get("window_id") in wids
+ for m in _of_type(msgs, "toolbar_config")))
+
+
+class TestCommit:
+ def test_commit_snapshots_the_previewed_frame(self, window):
+ session, plot, tree, wiz = _opened(window)
+ before = len(session.signal_trees)
+ pa.seg_commit(session, plot, {})
+ assert len(session.signal_trees) == before + 1
+ committed = session.signal_trees[-1]
+ assert committed.particles.n_frames == 1
+ assert committed.particles.n_particles == wiz.preview["count"]
+ assert committed.source_tree is tree
+
+ def test_commit_stamps_provenance(self, window):
+ session, plot, _tree, _wiz = _opened(window)
+ pa.seg_commit(session, plot, {})
+ prov = getattr(session.signal_trees[-1], "_commit_provenance", None) or {}
+ assert prov.get("action") == "segment_particles"
+ assert prov.get("params", {}).get("mode") == "single_frame"
+
+ def test_commit_without_a_preview_errors(self, window):
+ session, _plot, _tree = _movie(window)
+ msgs = window["messages"]
+ pa.seg_commit(session, _signal_plot(session), {})
+ assert any("nothing to commit" in str(m.get("text", ""))
+ for m in _of_type(msgs, "error"))
+
+
+class TestDoubleFire:
+ def test_open_close_open_leaves_one_controller(self, window):
+ """README §4 / StrictMode: mount → cleanup → remount fires all three
+ synchronously, before any worker lands."""
+ session, plot, tree = _movie(window)
+ built = []
+ real_init = pa.SegmentWizard.__init__
+
+ def _tracking(self, *a, **k):
+ real_init(self, *a, **k)
+ built.append(self)
+
+ pa.SegmentWizard.__init__ = _tracking
+ try:
+ pa.seg_open(session, plot, {})
+ pa.seg_close(session, plot, {})
+ pa.seg_open(session, plot, {})
+ finally:
+ pa.SegmentWizard.__init__ = real_init
+ time.sleep(0.8)
+
+ alive = [w for w in built if not w._closed]
+ assert len(alive) == 1, \
+ f"expected 1 live controller, got {len(alive)} of {len(built)} built"
+ assert tree._seg_wizard is alive[0]
+
+ pa.seg_close(session, plot, {})
+ assert getattr(tree, "_seg_wizard", None) is None
+ assert all(w._closed for w in built)
+
+ def test_close_without_an_open_is_harmless(self, window):
+ session, plot, tree = _movie(window)
+ pa.seg_close(session, plot, {})
+ assert getattr(tree, "_seg_wizard", None) is None
+
+ def test_reopen_keeps_the_scribbles(self, window):
+ """A second open must adopt the new parameters, not throw the user's
+ accumulated labels away."""
+ session, plot, tree, wiz = _opened(window)
+ pa.seg_paint(session, plot, {"frame": 0, "points": [[20, 20], [22, 22]],
+ "class_id": 0, "brush": 3})
+ painted = len(wiz.labels)
+ pa.seg_open(session, plot, {"sensitivity": 0.7})
+ assert tree._seg_wizard is wiz
+ assert len(wiz.labels) == painted
+
+
+class TestSchema:
+ def test_schema_resolves_through_the_registry(self):
+ from spyde.actions import registry
+ schema = registry.wizard_parameters("seg")
+ assert schema and schema is not pa.SegmentWizard.parameters
+
+ def test_schema_defaults_match_the_handler_defaults(self):
+ """The drift test_wizard_schemas.py exists to catch."""
+ from spyde.actions import registry
+ schema = registry.wizard_parameters("seg")
+ for key, spec in schema.items():
+ assert key in pa.DEFAULTS, f"seg schema declares unknown param {key!r}"
+ assert spec["default"] == pa.DEFAULTS[key], \
+ f"seg schema/{key} drifted from particles_action.DEFAULTS"
+
+ def test_every_stage_is_registered(self):
+ from spyde.actions.registry import STAGED_HANDLERS, resolve_staged
+ for stage in ("seg_open", "seg_close", "seg_set_method", "seg_tune",
+ "seg_paint", "seg_train", "seg_run", "seg_commit"):
+ assert stage in STAGED_HANDLERS
+ assert callable(resolve_staged(stage))
+
+ def test_toolbar_entry_points_at_the_action(self):
+ import spyde
+ meta = spyde.TOOLBAR_ACTIONS["functions"]["Segment Particles"]
+ assert meta["function"] == \
+ "spyde.actions.particles_action.segment_particles"
+ assert "spyde_diffraction_vectors_image" in meta["exclude_signal_types"]
+
+ @pytest.mark.parametrize("signal_type,offered", [
+ ("insitu", True), # the in-situ movie — the primary shape
+ ("", True), # a plain 2-D image — plan §0.10
+ ("electron_diffraction", False), # a 4D-STEM scan has no image frames
+ ("particles", False), # a label movie is not re-segmented
+ ])
+ def test_toolbar_gating(self, signal_type, offered):
+ import spyde
+ from spyde.drawing.toolbars.plot_control_toolbar import _action_matches_plot
+
+ class _Sig:
+ _signal_type = signal_type
+
+ class _Tree:
+ particles = None
+ diffraction_vectors = None
+ root = _Sig()
+
+ class _Plot:
+ signal_tree = _Tree()
+ is_navigator = False
+
+ class _State:
+ plot = _Plot()
+ current_signal = _Sig()
+ dimensions = 2
+ navigation = False
+
+ meta = spyde.TOOLBAR_ACTIONS["functions"]["Segment Particles"]
+ assert _action_matches_plot("Segment Particles", meta, _State()) is offered
+
+
+class TestCoercion:
+ def test_local_size_is_forced_odd(self):
+ """skimage requires it and SegmentParams raises rather than bumping —
+ a slider that lands on an even number must not error."""
+ assert pa._coerce({"local_size": 30})["local_size"] == 31
+
+ def test_unknown_method_falls_back(self):
+ assert pa._coerce({"method": "magic"})["method"] == pa.DEFAULTS["method"]
+
+ def test_sensitivity_is_clamped(self):
+ assert pa._coerce({"sensitivity": 5.0})["sensitivity"] == 1.0
+ assert pa._coerce({"sensitivity": -2.0})["sensitivity"] == 0.0
+
+ def test_a_junk_value_keeps_the_default(self):
+ assert pa._coerce({"min_separation": "wat"})["min_separation"] == \
+ pa.DEFAULTS["min_separation"]
+
+
+class TestPhysicalFaceControls:
+ """`merge_nm` / `min_nm` are the two controls on the caret's face, and they
+ are in NANOMETRES — so the whole feature is the nm→px conversion.
+
+ It was dead on arrival. ``_nm_to_px`` divides by ``p["scale"]`` and falls
+ back to "the value IS pixels" when there is no scale, which is the right
+ behaviour for an uncalibrated signal — but ``_coerce`` builds the dict from
+ the ``DEFAULTS`` keys alone, so no dispatch path ever put a scale in it and
+ EVERY signal took the uncalibrated branch. A slider reading "50 nm" merged
+ at 50 pixels, wrong by exactly the magnification and silent about it.
+
+ Hence the first test: it asserts the scale is on the params after each
+ dispatch verb, which is the thing that was missing, rather than asserting
+ ``_nm_to_px`` divides — that part was always correct.
+ """
+
+ def test_every_dispatch_stashes_the_signal_scale(self, window):
+ session, plot, _tree, wiz = _opened(window)
+ scale = wiz.scale_units()[0]
+ assert scale > 0, "the fixture must be calibrated or this proves nothing"
+
+ assert wiz.params["scale"] == scale, "seg_open did not stash the scale"
+ pa.seg_tune(session, plot, {"merge_nm": 30.0})
+ assert _wait(lambda: wiz.params.get("merge_nm") == 30.0)
+ assert wiz.params["scale"] == scale, "seg_tune dropped the scale"
+ pa.seg_set_method(session, plot, {"method": "classical"})
+ assert wiz.params["scale"] == scale, "seg_set_method dropped the scale"
+
+ # The particle fixture is calibrated at 1.0 nm/px, which makes every
+ # conversion below the identity — so these pass an explicit NON-UNIT scale
+ # instead of reading the fixture's. At 1.0 the arithmetic tests hold whether
+ # or not the conversion happens at all, which is exactly how it shipped
+ # unconverted. Only the dispatch test above needs a real wizard.
+
+ def test_merge_nm_reaches_the_solver_in_PIXELS(self):
+ """A nm distance must be divided by the scale before it is a radius."""
+ p = pa._coerce({"merge_nm": 30.0})
+ assert pa._segment_kwargs({**p, "scale": 0.4})["merge_distance"] == \
+ pytest.approx(75.0)
+
+ def test_min_nm_is_a_DIAMETER_and_converts_through_the_AREA(self):
+ """`min_size` filters on AREA, so handing it a diameter straight over
+ under-filters by a factor of ~d — a face control that barely does
+ anything is worse than one that is not there."""
+ p = pa._coerce({"min_nm": 20.0})
+ d_px = 20.0 / 0.4
+ assert pa._min_size_px({**p, "scale": 0.4}) == \
+ int(round(np.pi / 4.0 * d_px * d_px))
+ # and not the diameter, which is the plausible wrong answer
+ assert pa._min_size_px({**p, "scale": 0.4}) != int(d_px)
+
+ def test_min_nm_wins_over_the_pixel_min_size(self):
+ """Both exist (one on the face, one in Advanced) and they filter the same
+ thing, so the physical one has to be the tie-break — otherwise the face
+ control is silently overridden by a value the user cannot see."""
+ p = pa._coerce({"min_nm": 20.0, "min_size": 25})
+ assert pa._min_size_px({**p, "scale": 0.4}) != 25
+ assert pa._min_size_px({**p, "min_nm": 0.0, "scale": 0.4}) == 25
+
+ def test_an_uncalibrated_signal_reads_the_value_as_pixels(self):
+ """No scale is the only case where nm==px is correct, and it must not
+ divide by zero to get there."""
+ assert pa._nm_to_px(12.0, {}) == 12.0
+ assert pa._nm_to_px(12.0, {"scale": 0.0}) == 12.0
+ assert pa._nm_to_px(0.0, {"scale": 0.5}) == 0.0
+
+ def test_a_length_axis_converts_through_NANOMETRES(self):
+ """An axis in µm or Å is still a length; the slider is still nm."""
+ assert pa._length_nm_per_px(2.0, "nm") == pytest.approx(2.0)
+ assert pa._length_nm_per_px(2.0, "um") == pytest.approx(2000.0)
+ assert pa._length_nm_per_px(2.0, "µm") == pytest.approx(2000.0)
+ assert pa._length_nm_per_px(2.0, "Å") == pytest.approx(0.2)
+ # ...so "20 nm" is 10 px at 2 nm/px and 0.01 px at 2 µm/px.
+ assert pa._nm_to_px(20.0, {"scale": pa._length_nm_per_px(2.0, "nm")}) \
+ == pytest.approx(10.0)
+
+ def test_a_RECIPROCAL_axis_is_not_a_length_and_falls_back_to_pixels(self):
+ """The reported bug: the caret said 'nm' on an axis reading nm⁻¹.
+
+ A reciprocal-space signal reports a perfectly healthy positive scale, so
+ the conversion ran and produced a merge radius wrong by whatever the
+ camera length was — silently, because nothing checked the UNIT. There is
+ no distance to convert to here, so the controls are pixels and the caret
+ has to relabel; claiming nm is a claim about the scale bar that is false.
+ """
+ for units in ("nm^-1", "1/nm", "nm⁻¹", "mrad", "px", "", None):
+ assert pa._length_nm_per_px(0.9, units) == 0.0, units
+ assert pa._face_units(0.9, units) == "px", units
+ assert pa._face_units(0.5, "nm") == "nm"
+ # and the pixel fallback is the IDENTITY, not a division by a number in
+ # the wrong unit
+ assert pa._nm_to_px(30.0, {"scale": pa._length_nm_per_px(0.9, "nm^-1")}) \
+ == 30.0
+
+ def test_the_wizard_stashes_nm_per_px_not_the_raw_axis_scale(self, window):
+ _s, _p, _t, wiz = _opened(window)
+ scale, units = wiz.scale_units()
+ assert wiz.params["scale"] == pytest.approx(
+ pa._length_nm_per_px(scale, units))
+
+
+class TestThresholdFailureIsNotAResult:
+ """14028 'particles' covering the frame is a failed threshold, not an answer.
+
+ Measured on a low-contrast noisy stand-in for the reported frame: otsu has
+ no bimodal histogram to find, lands inside the noise, and the split shatters
+ the support film. `min_size` alone moves 4873 -> 17 instances but coverage
+ only 39% -> 7%, and the settings that DO yield 8 instances cover 52% of the
+ frame -- those 8 bodies ARE the film. So neither number alone is diagnostic
+ and the verdict needs both.
+ """
+
+ def test_both_conditions_are_required(self):
+ assert pa._threshold_failed(5000, 0.40) is True
+ # A crowded frame of REAL particles: many instances, but they do not
+ # blanket the frame.
+ assert pa._threshold_failed(5000, 0.05) is False
+ # A coarse threshold that found a few big real objects.
+ assert pa._threshold_failed(12, 0.40) is False
+ assert pa._threshold_failed(0, 0.0) is False
+
+ def test_a_normal_preview_is_not_flagged(self, window):
+ """The fixture must stay UNflagged or the notice is just noise."""
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_tune(session, plot, {"sensitivity": 0.55})
+ assert _wait(lambda: len(_of_type(msgs, "seg_preview")) >= 2)
+ msg = _of_type(msgs, "seg_preview")[-1]
+ assert msg["threshold_failed"] is False, msg
+ assert 0.0 <= msg["coverage"] <= 1.0
+
+ def test_the_preview_reports_coverage_and_face_units(self, window):
+ session, plot, _tree, _wiz = _opened(window)
+ msgs = window["messages"]
+ pa.seg_tune(session, plot, {"sensitivity": 0.5})
+ assert _wait(lambda: len(_of_type(msgs, "seg_preview")) >= 2)
+ msg = _of_type(msgs, "seg_preview")[-1]
+ assert "coverage" in msg and "face_units" in msg
+ assert msg["face_units"] in ("nm", "px")
+
+
+class TestOutlineDrawCap:
+ """No route may hand the renderer an unbounded number of polygons.
+
+ `_RASTER_ABOVE` picks the nicer drawing; this cap is the seatbelt for when
+ the raster is unavailable. It had none, so a tile-mode raster failure fell
+ back to one filled path per instance -- 14028 of them, which hangs the
+ renderer and composites into a flat green sheet showing nothing.
+ """
+
+ def test_the_cap_is_well_above_any_real_outline_count(self):
+ assert pa._MAX_OUTLINE_POLYS > pa._RASTER_ABOVE * 5
+
+ def test_a_huge_contour_list_draws_NONE_rather_than_all(self, monkeypatch):
+ import types
+ pushed = {}
+ monkeypatch.setattr(pa, "_push_groups", lambda p, u: pushed.update(
+ {g.name: pl.get("vertices_list") for g, pl in u.items()}))
+
+ w = object.__new__(pa.SegmentWizard)
+ w._ov_group = None
+ w._ov_box_group = None
+ w._ov_raster = False
+ w._ov_cleared = False
+ w._ov_box_state = None
+ # No `set_overlay_mask` on this plot => the raster path is unavailable,
+ # which is exactly the state the tile-mode ValueError used to produce.
+ w.src_plot = types.SimpleNamespace(_plot2d=types.SimpleNamespace())
+
+ # A plain class, NOT SimpleNamespace: the group is used as a dict KEY in
+ # the `_push_groups` payload, and SimpleNamespace defines __eq__ so it
+ # is unhashable.
+ class _Group:
+ name = "seg_preview_outline"
+
+ monkeypatch.setattr(pa.SegmentWizard, "_overlay_group",
+ lambda self, p: _Group())
+ monkeypatch.setattr(pa.SegmentWizard, "_window_group", lambda self, p: None)
+
+ n = pa._MAX_OUTLINE_POLYS + 1
+ w.set_overlay([np.zeros((3, 2), np.int16) for _ in range(n)], None,
+ labels=None, n_instances=n)
+ assert pushed.get("seg_preview_outline") == [], (
+ f"{n} polygons reached the renderer; the cap is "
+ f"{pa._MAX_OUTLINE_POLYS}")
+
+
+class TestMeasureSkipsOutlinesItCannotDraw:
+ """Contours are ~half a preview's cost once instances run to thousands."""
+
+ def test_want_contours_false_skips_them_but_keeps_every_measurement(self):
+ from spyde.particles import measure_frame
+
+ lab = np.zeros((64, 64), np.int32)
+ lab[5:15, 5:15] = 1
+ lab[30:40, 30:40] = 2
+ frame = np.random.default_rng(0).random((64, 64)).astype(np.float32)
+
+ rows_a, cont_a = measure_frame(lab, frame, t=0, scale=0.5)
+ rows_b, cont_b = measure_frame(lab, frame, t=0, scale=0.5,
+ want_contours=False)
+ assert len(cont_a) == 2 and cont_b == []
+ # The ROWS -- every measured property, including the score the
+ # confidence filter reads -- must be untouched by the flag.
+ assert np.array_equal(rows_a, rows_b)
+
+ def test_commit_refuses_a_preview_whose_outlines_were_skipped(self, window):
+ """A store needs one contour PER ROW; committing without them would
+ build a silently mis-corresponding dataset."""
+ session, plot, _tree, wiz = _opened(window)
+ msgs = window["messages"]
+ wiz.preview = {"frame": 0, "labels": np.zeros((8, 8), np.int32),
+ "rows": np.zeros((3, 4), np.float32), "contours": [],
+ "count": 3, "areas": np.zeros(3), "box": None}
+ assert wiz.commit() is None
+ assert any("too many to commit" in str(m.get("text", ""))
+ for m in msgs if isinstance(m, dict)), \
+ "commit failed silently instead of saying why"
+
+
+class TestBrushActuallyPaints:
+ """The regression guard for "I can't scribble".
+
+ Shipped broken once, in two independent ways, and neither was visible to any
+ test that existed at the time:
+
+ 1. **Nothing ever created a brush widget.** ``add_brush_widget`` was called
+ nowhere, so Shift+drag had nothing to hit.
+ 2. **The caret listened on a path the brush cannot reach.** It waited for a
+ renderer-side ``spyde:figure_event`` carrying a points array, but an
+ anyplotlib stroke travels to PYTHON (``event_json`` →
+ ``Figure._dispatch_event`` → ``Widget._update_from_js`` →
+ ``plot.callbacks.fire``). The renderer never sees it, so even with a brush
+ present nothing would have arrived.
+
+ The old tests passed because they posted a synthetic ``seg_paint`` payload —
+ which exercises the rasteriser and proves nothing about whether a real stroke
+ can ever get there. These drive the widget.
+ """
+
+ def _scribble_wizard(self, session):
+ from spyde.actions.particles_action import seg_open, seg_set_method
+ plots = (list(session._plots) if isinstance(session._plots, list)
+ else list(session._plots.values()))
+ plot = next(p for p in plots
+ if getattr(p, "signal_tree", None) is not None
+ and getattr(p, "_plot2d", None) is not None)
+ seg_open(session, plot, {"window_id": getattr(plot, "window_id", None)})
+ seg_set_method(session, plot, {"method": "scribble"})
+ return plot, plot.signal_tree
+
+ def test_a_brush_is_attached_on_the_scribble_tab(self, window):
+ from spyde.actions.particles_action import _brush_supported
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ _plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ assert getattr(tree, "_seg_brush", None) is not None, (
+ "no brush on the plot — Shift+drag has nothing to hit")
+
+ def test_a_stroke_from_the_WIDGET_reaches_the_label_store(self, window):
+ """Drives the widget, not a synthetic seg_paint payload."""
+ from spyde.actions.particles_action import _brush_supported, _on_stroke
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ _plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ wiz = tree._seg_wizard
+ brush = tree._seg_brush
+ store = wiz.label_store()
+ before = dict(store.counts())
+ # What anyplotlib delivers: [[x, y], …] in IMAGE PIXELS, then pointer_up.
+ brush.set(strokes=[[[20.0, 30.0], [24.0, 32.0], [28.0, 34.0]]],
+ stroke_classes=[0])
+ _on_stroke(wiz, None)
+ after = dict(store.counts())
+ assert after != before, "the stroke never reached the label store"
+ assert after[0] > 0
+
+ def test_a_second_stroke_only_paints_its_own_points(self, window):
+ """The widget accumulates strokes for its whole life, so replaying the
+ full list on every event would re-paint everything and make the class
+ counts grow quadratically.
+
+ The class is switched through ``seg_tune``, NOT by handing the widget
+ different ``stroke_classes``: the caret's params are the authority,
+ precisely because the JS widget's own value is not reliably in sync (see
+ :meth:`test_class_comes_from_PARAMS_not_from_the_js_widget`).
+ """
+ from spyde.actions.particles_action import (_brush_supported, _on_stroke,
+ seg_tune)
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ wiz, brush = tree._seg_wizard, tree._seg_brush
+ store = wiz.label_store()
+ s1 = [[20.0, 30.0], [24.0, 32.0]]
+ brush.set(strokes=[s1])
+ _on_stroke(wiz, None)
+ one = dict(store.counts())
+ # Fire again with NO new stroke: nothing may change.
+ _on_stroke(wiz, None)
+ assert dict(store.counts()) == one, "re-fired the same stroke"
+ # A genuinely new stroke, in a class chosen the way the strip chooses it.
+ seg_tune(session, plot, {"active_class": 1})
+ brush.set(strokes=[s1, [[60.0, 40.0], [64.0, 42.0]]])
+ _on_stroke(wiz, None)
+ two = dict(store.counts())
+ assert two[1] > 0, "the second stroke's class was not painted"
+ assert two[0] == one[0], "the first stroke was painted twice"
+
+ def test_leaving_scribble_detaches_the_brush(self, window):
+ """It floats over the image; on Classical there is nothing to paint."""
+ from spyde.actions.particles_action import _brush_supported, seg_set_method
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ assert tree._seg_brush is not None
+ seg_set_method(session, plot, {"method": "classical"})
+ assert getattr(tree, "_seg_brush", None) is None
+
+ def test_missing_brush_says_so_instead_of_failing_silently(self, window,
+ monkeypatch):
+ """A user dragging at an image that never responds must be TOLD why."""
+ import spyde.actions.particles_action as pa
+ monkeypatch.setattr(pa, "_brush_supported", lambda: False)
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ before = len(window["messages"])
+ self._scribble_wizard(session)
+ new = window["messages"][before:]
+ assert any("anyplotlib" in str(m.get("text", "")).lower()
+ or "brush" in str(m.get("text", "")).lower() for m in new), (
+ "switching to Scribble with no brush available emitted nothing — "
+ "the user gets a picture that ignores them and no explanation")
+
+
+class TestPaintStateReachesTheWidget:
+ """Class and eraser must travel to the WIDGET, not stop at wiz.params.
+
+ Both bugs here shipped and were user-visible: every stroke came out in class
+ 0 ("I can only scribble one colour") and the eraser did nothing ("delete
+ doesn't work"). One root cause — the ClassStrip set React state, nothing sent
+ it to Python, and the backend read `active_class` / `erase` from params that
+ were never declared and never set.
+
+ The widget tags each stroke with its OWN class at paint time in JS, so a
+ change that only reaches ``wiz.params`` paints the previous colour forever.
+ That is why these assert on the WIDGET's attributes and on what actually
+ landed in the store — not on the params dict, which was already "right" while
+ the paint was wrong.
+ """
+
+ def _scribbling(self, session):
+ from spyde.actions.particles_action import seg_open, seg_set_method
+ session._load_test_data_particles({"frames": 4})
+ plots = (list(session._plots) if isinstance(session._plots, list)
+ else list(session._plots.values()))
+ plot = next(p for p in plots
+ if getattr(p, "signal_tree", None) is not None
+ and getattr(p, "_plot2d", None) is not None)
+ seg_open(session, plot, {"window_id": getattr(plot, "window_id", None)})
+ seg_set_method(session, plot, {"method": "scribble"})
+ return plot, plot.signal_tree
+
+ @staticmethod
+ def _stroke(tree, wiz, pts):
+ from spyde.actions.particles_action import _on_stroke
+ brush = tree._seg_brush
+ strokes = list(getattr(brush, "strokes", []) or [])
+ classes = list(getattr(brush, "stroke_classes", []) or [])
+ brush.set(strokes=strokes + [pts],
+ stroke_classes=classes + [int(brush.class_id)])
+ _on_stroke(wiz, None)
+
+ def test_defaults_declare_the_paint_state(self):
+ """They were read but never declared, so params.get() always won."""
+ from spyde.actions.particles_action import DEFAULTS
+ assert "active_class" in DEFAULTS, (
+ "active_class is read by the brush but not a declared parameter, so "
+ "nothing can ever set it and every stroke is class 0")
+ assert "erase" in DEFAULTS
+
+ def test_class_comes_from_PARAMS_not_from_the_js_widget(self, window):
+ """The test that would have caught it. The old one did not.
+
+ The first version set `brush.class_id` and read it back — both PYTHON
+ side — so it passed while the app was broken. In the real app the widget
+ lives in JS, and `Figure._push_widget` sends a targeted update that never
+ writes `panel__json`, so a Python-side class push does not reliably
+ reach it before the next stroke.
+
+ The natural experiment that exposed this: the ERASER worked and the CLASS
+ did not, in the same handler on the same stroke — because erase read from
+ `wiz.params` and class read from the widget's `stroke_classes`.
+
+ So this test forces the widget's own class to be STALE and wrong, and
+ asserts the stroke still lands in the class the caret asked for.
+ """
+ from spyde.actions.particles_action import _brush_supported, _on_stroke, seg_tune
+ session = window["window"]
+ plot, tree = self._scribbling(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ wiz, brush = tree._seg_wizard, tree._seg_brush
+ store = wiz.label_store()
+
+ seg_tune(session, plot, {"active_class": 2})
+ # Simulate the push NOT landing: the widget still thinks it is class 0,
+ # and tags the stroke accordingly.
+ brush.set(class_id=0)
+ brush.set(strokes=[[[20.0, 30.0], [26.0, 32.0]]], stroke_classes=[0])
+ _on_stroke(wiz, None)
+
+ counts = dict(store.counts())
+ assert counts.get(2, 0) > 0, (
+ "the stroke was filed under the WIDGET's stale class instead of the "
+ "one the caret selected — this is the 'can only scribble one colour' "
+ f"bug: {counts}")
+ assert counts.get(0, 0) == 0
+
+ def test_switching_class_retags_the_next_stroke(self, window):
+ from spyde.actions.particles_action import _brush_supported, seg_tune
+ session = window["window"]
+ plot, tree = self._scribbling(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ wiz = tree._seg_wizard
+ store = wiz.label_store()
+
+ self._stroke(tree, wiz, [[20.0, 30.0], [26.0, 32.0]])
+ seg_tune(session, plot, {"active_class": 1})
+ assert int(tree._seg_brush.class_id) == 1, (
+ "seg_tune did not reach the widget — the strip's choice stops at "
+ "wiz.params and the next stroke paints the OLD class")
+ self._stroke(tree, wiz, [[60.0, 40.0], [66.0, 42.0]])
+ seg_tune(session, plot, {"active_class": 2})
+ self._stroke(tree, wiz, [[80.0, 20.0], [86.0, 22.0]])
+
+ counts = dict(store.counts())
+ assert counts[0] > 0 and counts[1] > 0 and counts[2] > 0, (
+ f"only some classes were painted: {counts}")
+
+ def test_eraser_removes_and_leaves_other_classes_alone(self, window):
+ from spyde.actions.particles_action import _brush_supported, seg_tune
+ session = window["window"]
+ plot, tree = self._scribbling(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ wiz = tree._seg_wizard
+ store = wiz.label_store()
+
+ path = [[20.0, 30.0], [26.0, 32.0]]
+ self._stroke(tree, wiz, path)
+ seg_tune(session, plot, {"active_class": 1})
+ self._stroke(tree, wiz, [[60.0, 40.0], [66.0, 42.0]])
+ before = dict(store.counts())
+ assert before[0] > 0 and before[1] > 0
+
+ seg_tune(session, plot, {"erase": True})
+ assert bool(tree._seg_brush.erase) is True, (
+ "the eraser is a WIDGET mode — a stroke is tagged before the handler "
+ "sees it, so an erase flag that stops at wiz.params does nothing")
+ self._stroke(tree, wiz, path) # retrace the class-0 stroke
+
+ after = dict(store.counts())
+ assert after[0] < before[0], "the eraser removed nothing"
+ assert after[1] == before[1], "the eraser hit a class it was not over"
+
+ def test_brush_size_reaches_the_widget_too(self, window):
+ from spyde.actions.particles_action import _brush_supported, seg_tune
+ session = window["window"]
+ plot, tree = self._scribbling(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget (needs >= 0.5.0)")
+ seg_tune(session, plot, {"brush": 9.0})
+ assert float(tree._seg_brush.radius) == pytest.approx(9.0)
+
+
+class TestBrushDrawColour:
+ """The brush must DRAW in the selected class's colour, not just tag strokes.
+
+ Reported twice as "I can only scribble one colour" / "support film still
+ doesn't change the painting colour". The first fix made the STROKE land in
+ the right class, which it now does — the per-class pixel counts in the caret
+ prove it. But the colour on screen comes from the WIDGET's own `class_id`
+ (`figure_esm.js::_brushLiveBegin` reads `w.class_id` at stroke start and
+ draws in `colors[class_id]`), so the data can be right while the paint is
+ still orange. These assert the widget state, which is the thing the eye sees.
+ """
+
+ def _scribble_wizard(self, session):
+ from spyde.actions.particles_action import seg_open, seg_set_method
+ plots = session._plots
+ plots = list(plots.values()) if hasattr(plots, "values") else list(plots)
+ plot = next(p for p in plots
+ if getattr(p, "signal_tree", None) is not None
+ and getattr(p, "_plot2d", None) is not None)
+ seg_open(session, plot, {"window_id": getattr(plot, "window_id", None)})
+ seg_set_method(session, plot, {"method": "scribble"})
+ return plot, plot.signal_tree
+
+ def test_selecting_a_class_changes_the_widget_class_id(self, window):
+ from spyde.actions.particles_action import _brush_supported, seg_tune
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget")
+ brush = getattr(tree, "_seg_brush", None)
+ assert brush is not None
+
+ # What the ClassStrip sends when you click "support film".
+ seg_tune(session, plot, {"active_class": 1})
+ assert int(brush._data["class_id"]) == 1, (
+ "the widget still paints class 0 — the strip's selection reached "
+ "wiz.params but not the widget, so every stroke DRAWS orange")
+
+ seg_tune(session, plot, {"active_class": 2})
+ assert int(brush._data["class_id"]) == 2
+
+ def test_the_new_class_reaches_the_PANEL_state(self, window):
+ """The JS draws from ``panel__json``, not from the Python object.
+
+ This is the assertion the previous two miss and the reason the bug
+ survived a "fix": ``Widget.set`` reaches JS via ``_push_widget``, which
+ writes ``event_json`` ONLY and leaves the panel's own widget state
+ stale. ``_brushLiveBegin`` reads ``w.class_id`` from
+ ``p.state.overlay_widgets`` and paints ``colors[class_id]`` — so
+ ``brush._data`` can say 1 while every stroke still draws in class 0's
+ colour. Assert the serialised panel, which is what the eye sees.
+ """
+ from spyde.actions.particles_action import _brush_supported, seg_tune
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget")
+
+ # Spy on the panel push. Asserting `_state` alone would be VACUOUS:
+ # `overlay_widgets` holds the widget's own `_data` BY REFERENCE, so
+ # `brush.set()` mutates it in place and the dict always looks current
+ # whether or not anything was ever sent to JS. What the renderer sees is
+ # the re-serialised trait, so the assertion has to be "a push happened".
+ pushes = []
+ real_push = plot._plot2d._push
+ plot._plot2d._push = lambda *a, **k: (pushes.append(1), real_push(*a, **k))[1]
+ try:
+ seg_tune(session, plot, {"active_class": 2})
+ finally:
+ plot._plot2d._push = real_push
+
+ assert pushes, (
+ "no panel push after the class changed — the targeted widget update "
+ "writes event_json only, so panel__json keeps the OLD class and "
+ "JS goes on painting the previous class's colour")
+
+ widgets = (plot._plot2d._state.get("overlay_widgets") or [])
+ brushes = [w for w in widgets if (w or {}).get("type") == "brush"]
+ assert brushes, f"no brush in the panel state: {widgets}"
+ assert int(brushes[0].get("class_id", -1)) == 2
+
+ def test_an_unrelated_tune_does_NOT_force_a_panel_push(self, window):
+ """`seg_tune` fires for every sensitivity-slider tick too.
+
+ A full panel push re-serialises the image bytes, so doing one per tick
+ on a 4096² frame would trade the colour bug for a much worse drag. Only
+ an actual brush change (class / eraser / size) may pay for it.
+ """
+ from spyde.actions.particles_action import _brush_supported, seg_tune
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ plot, _tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget")
+
+ seg_tune(session, plot, {"active_class": 1}) # settle the state
+
+ pushes = []
+ real_push = plot._plot2d._push
+ plot._plot2d._push = lambda *a, **k: (pushes.append(1), real_push(*a, **k))[1]
+ try:
+ seg_tune(session, plot, {"sensitivity": 0.6})
+ seg_tune(session, plot, {"sensitivity": 0.7})
+ seg_tune(session, plot, {"active_class": 1}) # SAME class, no change
+ finally:
+ plot._plot2d._push = real_push
+
+ assert pushes == [], (
+ f"{len(pushes)} panel pushes for tunes that did not change the "
+ "brush — a sensitivity drag would re-serialise the whole image")
+
+ def test_the_widget_has_a_colour_for_every_class(self, window):
+ """`colors` is indexed by class_id in JS, so a short list means the
+ later classes draw with `undefined` — no colour change on screen even
+ though class_id updated correctly."""
+ from spyde.actions.particles_action import _brush_supported
+ session = window["window"]
+ session._load_test_data_particles({"frames": 4})
+ _plot, tree = self._scribble_wizard(session)
+ if not _brush_supported():
+ pytest.skip("installed anyplotlib has no brush widget")
+ brush = tree._seg_brush
+ from spyde.particles import default_classes
+ n_classes = len(default_classes())
+ colours = list(brush._data.get("colors") or [])
+ assert len(colours) >= n_classes, (
+ f"brush carries {len(colours)} colours for {n_classes} classes — "
+ f"class ids >= {len(colours)} draw with no colour: {colours}")
+
+
+class TestBatchComputingOverlay:
+ """The "Calculating…" chip over the RESULT window, and when it appears.
+
+ Reported as: "there is a lot of lag between the subwindow appearing and then
+ [calculating]". The chip was not raised at all, and the obvious place to add
+ it — next to the first progress emission, inside the worker — is far too
+ late: the window opens, then the placeholder store is built, the cancel
+ flags registered, the generation bumped, the worker scheduled, the thread
+ hop paid, and the first frame COMPUTED. On 4096² frames that is seconds of a
+ window that looks finished and empty.
+
+ So the contract is SYNCHRONOUS: by the time `seg_run` returns to the event
+ loop, the chip is already up.
+ """
+
+ def _seg_run(self, session, **params):
+ from spyde.actions.particles_action import seg_run
+ plots = session._plots
+ plots = list(plots.values()) if hasattr(plots, "values") else list(plots)
+ plot = next(p for p in plots
+ if getattr(p, "signal_tree", None) is not None
+ and getattr(p, "_plot2d", None) is not None)
+ seg_run(session, plot, dict(params))
+ return plot
+
+ def test_the_chip_is_raised_before_seg_run_returns(self, window):
+ session = window["window"]
+ msgs = window["messages"]
+ session._load_test_data_particles({"frames": 4})
+ del msgs[:]
+
+ self._seg_run(session, method="classical", track=False)
+
+ # No waiting, no polling: the chip must already be on the wire.
+ raised = [m for m in msgs
+ if m.get("type") == "window_computing" and m.get("computing")]
+ assert raised, (
+ "no window_computing raised synchronously — the chip only appears "
+ "once the worker gets going, which is the reported lag")
+
+ def test_the_chip_names_the_RESULT_window(self, window):
+ """Not the source window: the source is where you were scribbling and it
+ is not the thing sitting there looking empty."""
+ session = window["window"]
+ msgs = window["messages"]
+ session._load_test_data_particles({"frames": 4})
+ src = self._seg_run(session, method="classical", track=False)
+ del msgs[:]
+
+ raised = [m for m in window["messages"]
+ if m.get("type") == "window_computing"]
+ # Re-read from the full list: the run may already have finished.
+ all_ids = {m.get("window_id") for m in raised}
+ src_id = getattr(src, "window_id", None)
+ if all_ids:
+ assert all_ids != {src_id}, (
+ "the chip was put on the SOURCE window, not the result")
+
+ def test_the_chip_comes_down_when_the_batch_ends(self, window):
+ session = window["window"]
+ msgs = window["messages"]
+ session._load_test_data_particles({"frames": 4})
+ del msgs[:]
+ self._seg_run(session, method="classical", track=False)
+
+ assert _wait(lambda: any(
+ m.get("type") == "window_computing" and not m.get("computing")
+ for m in msgs), timeout=180), (
+ "the Calculating chip never came down — it will spin forever over a "
+ "window that has finished")
+
+
+class TestPreviewWindowBoxLifecycle:
+ """The 1-megapixel preview-window box must not depend on an ENGINE.
+
+ Reported: "the 1k x 1k outline only shows in classical; if you move to a
+ different mode it continues to show, but toggle out/back in and it won't
+ show up again." All three symptoms are one cause — the box was drawn only
+ as a side effect of a successful segmentation, and `_preview` returns early
+ when there is no engine (an untrained Scribble, or Prompt before any
+ prompt), painting nothing and CLEARING nothing:
+
+ * Classical always has a solver, so only it drew the box.
+ * Switching to an untrained engine left the previous box on screen,
+ because the early return cleared nothing.
+ * Toggling the caret dropped the overlay groups, and reopening in an
+ untrained engine early-returned again, so it never came back.
+
+ The box documents WHERE the budget looks. That is true whenever the caret
+ is open.
+ """
+
+ @staticmethod
+ def _wiz():
+ import types
+ import spyde.actions.particles_action as pa
+
+ class _Grp: # hashable by identity, unlike SimpleNamespace
+ def __init__(self, name):
+ self.name, self.removed = name, False
+
+ def remove(self):
+ self.removed = True
+
+ class _P2D:
+ def add_polygons(self, *a, **k):
+ return _Grp(k.get("name"))
+
+ wiz = object.__new__(pa.SegmentWizard)
+ wiz._ov_group = None
+ wiz._ov_box_group = None
+ wiz.src_plot = types.SimpleNamespace(_plot2d=_P2D())
+ wiz.frames = lambda: (1, lambda i: np.zeros((2048, 2048), np.float32),
+ (2048, 2048))
+ wiz.frame_index = lambda: 0
+ return wiz
+
+ def test_box_is_drawn_with_no_engine(self, monkeypatch):
+ import spyde.actions.particles_action as pa
+ pushed = {}
+ monkeypatch.setattr(pa, "_push_groups", lambda p, u: pushed.update(
+ {g.name: pl.get("vertices_list") for g, pl in u.items()}))
+ self._wiz().show_preview_window()
+ box = pushed.get("seg_preview_window")
+ assert box, "no preview-window box without an engine — the untrained " \
+ "Scribble/Prompt states show nothing at all"
+ assert len(box[0]) >= 4, "the box is not a closed rectangle"
+
+ def test_switching_to_an_untrained_engine_clears_stale_outlines(self, monkeypatch):
+ """A previous engine's instances must not linger looking like the new
+ engine's answer."""
+ import spyde.actions.particles_action as pa
+ pushed = {}
+ monkeypatch.setattr(pa, "_push_groups", lambda p, u: pushed.update(
+ {g.name: pl.get("vertices_list") for g, pl in u.items()}))
+ self._wiz().show_preview_window()
+ assert pushed.get("seg_preview_outline") == [], (
+ "the outline group was not cleared, so the old engine's particles "
+ "stay on screen")
+
+ def test_switching_to_an_untrained_engine_clears_the_RASTER_too(self, monkeypatch):
+ """The other drawing route, which this used to miss entirely.
+
+ Above `_RASTER_ABOVE` instances the overlay is ONE mask, not outlines.
+ `show_preview_window` cleared only the outlines, so the previous
+ engine's mask survived the switch — and on the Scribble tab that covers
+ the image you have to paint on. The reported symptom was exactly "moved
+ to scribble, immediately unusable".
+ """
+ import spyde.actions.particles_action as pa
+ monkeypatch.setattr(pa, "_push_groups", lambda p, u: None)
+ wiz = self._wiz()
+ cleared = []
+ monkeypatch.setattr(pa.SegmentWizard, "_clear_raster_overlay",
+ lambda self: cleared.append(True))
+ wiz.show_preview_window()
+ assert cleared, ("the raster mask was left on screen; on Scribble it "
+ "covers the frame the user has to paint on")
+
+
+class TestRasterOverlayAboveThreshold:
+ """Hundreds of vector contours make the app sluggish; draw one mask instead.
+
+ Every contour is a path the renderer re-transforms on each pan/zoom frame,
+ so a few hundred of them cost real interactivity. A mask is ONE image
+ however many particles it contains. Below the threshold the outlines stay —
+ crisp at any zoom, and each is an object the UI can hover.
+ """
+
+ # A REAL Plot2D, not a stub that records whatever it is handed.
+ #
+ # This class used to fake `set_overlay_mask`, and that fake is why the bug
+ # it was written to prevent shipped anyway: SpyDE reduced the mask to the
+ # overview grid (right for the renderer), the REAL `set_overlay_mask`
+ # rejected that shape against `image_width` (the full frame in tile mode),
+ # `_set_raster_overlay` swallowed the ValueError at DEBUG, and every
+ # large-frame preview fell back to N polygons — at 14028 instances, a hung
+ # renderer painted solid green. The stub asserted the mask SpyDE built and
+ # was green throughout, because it never enforced the contract that failed.
+ #
+ # So these tests now assert on the bytes that actually SHIP. Whichever layer
+ # does the reduction, the renderer's rule is the same and is the only thing
+ # worth pinning: `bytes.length === (base_width || image_width) * (…)`.
+ @staticmethod
+ def _p2d(base_w=0, base_h=0):
+ from anyplotlib.plot2d import Plot2D
+ full = np.zeros((4096, 4096), np.uint8)
+ if base_w:
+ p = Plot2D(full, tile="auto")
+ assert p._state.get("base_width"), "tile mode did not set a base grid"
+ else:
+ p = Plot2D(full)
+ return p
+
+ @staticmethod
+ def _p2d_small(n):
+ """A plot small enough that anyplotlib does NOT put it in tile mode."""
+ from anyplotlib.plot2d import Plot2D
+ return Plot2D(np.zeros((n, n), np.uint8))
+
+ @staticmethod
+ def _shipped(p2d):
+ """(bytes_len, expected_len) for the mask currently on *p2d*."""
+ import base64
+ st = p2d._state
+ b64 = st.get("overlay_mask_b64") or ""
+ want = (int(st.get("base_width") or 0) or int(st["image_width"])) * \
+ (int(st.get("base_height") or 0) or int(st["image_height"]))
+ return len(base64.b64decode(b64)), want
+
+ @staticmethod
+ def _wiz(p2d):
+ import types
+ import spyde.actions.particles_action as pa
+ w = object.__new__(pa.SegmentWizard)
+ w._ov_group = None
+ w._ov_box_group = None
+ w._ov_raster = False
+ w.src_plot = types.SimpleNamespace(_plot2d=p2d)
+ return w
+
+ @staticmethod
+ def _labels():
+ lab = np.zeros((1024, 1024), np.int32)
+ y, x = np.mgrid[0:1024, 0:1024]
+ for i, (cy, cx) in enumerate([(200, 200), (500, 700), (800, 300)], start=1):
+ lab[(y - cy) ** 2 + (x - cx) ** 2 < 25] = i # 5 px radius
+ return lab
+
+ def test_untiled_mask_is_the_native_frame_size(self):
+ """A frame small enough to escape tile mode ships at its native size.
+
+ Small ON PURPOSE: anyplotlib tiles anything with a >=1024 px edge, so
+ there is no such thing as an untiled 4096² plot and asking for one here
+ tested nothing. 512² is the branch where `base_width` stays 0 and the
+ renderer falls back to `image_width`.
+ """
+ p = self._p2d_small(512)
+ lab = np.zeros((512, 512), np.int32)
+ lab[100:110, 100:110] = 1
+ assert self._wiz(p)._set_raster_overlay(lab, None, (512, 512))
+ assert not p._state.get("base_width"), "512² unexpectedly tiled"
+ sent, want = self._shipped(p)
+ assert sent == want == 512 * 512
+
+ def test_tiled_mask_SHIPS_at_the_OVERVIEW_size(self):
+ """The trap this exists for, asserted where it actually bites.
+
+ In tile mode the renderer checks ``bytes.length === iw * ih`` where
+ ``iw = base_width || image_width`` — the OVERVIEW size — and on a
+ mismatch sets ``maskCache=null``: no error, no overlay, nothing in the
+ log. So what matters is not which layer reduces the mask but that the
+ bytes leaving Python are the size the renderer will accept.
+ """
+ p = self._p2d(1024, 1024)
+ assert self._wiz(p)._set_raster_overlay(
+ self._labels(), (1536, 1536, 1024, 1024), (4096, 4096)), \
+ "the raster overlay refused to draw on a tiled frame"
+ sent, want = self._shipped(p)
+ assert sent == want, (
+ f"mask ships {sent} bytes but the renderer expects {want} — it "
+ f"would silently draw nothing at all")
+
+ def test_the_reduction_keeps_small_particles(self):
+ """Block ANY, not a subsample.
+
+ Particles are often a few pixels across; striding a 4096² mask down to
+ 1024² would drop three quarters of them at random.
+ """
+ import base64
+ p = self._p2d(1024, 1024)
+ self._wiz(p)._set_raster_overlay(
+ self._labels(), (1536, 1536, 1024, 1024), (4096, 4096))
+ sent = np.frombuffer(
+ base64.b64decode(p._state["overlay_mask_b64"]), np.uint8)
+ assert sent.any(), "every particle vanished in the reduction"
+
+ def test_below_the_threshold_nothing_is_rastered(self):
+ """Few particles → keep the crisp, hoverable vector outlines."""
+ import spyde.actions.particles_action as pa
+ assert pa._RASTER_ABOVE > 1
+ p = self._p2d()
+ wiz = self._wiz(p)
+ wiz._ov_cleared = False
+ wiz._ov_box_state = None
+ contours = [np.zeros((3, 2), np.int16) for _ in range(3)]
+ wiz.set_overlay(contours, None, labels=self._labels(),
+ full_shape=(4096, 4096))
+ assert not (p._state.get("overlay_mask_b64") or ""), (
+ "a 3-particle frame was rastered; the outlines are better there")
diff --git a/spyde/tests/migrated/test_progressive_windowed.py b/spyde/tests/migrated/test_progressive_windowed.py
index d753479a..db51127f 100644
--- a/spyde/tests/migrated/test_progressive_windowed.py
+++ b/spyde/tests/migrated/test_progressive_windowed.py
@@ -44,6 +44,13 @@ def __init__(self, value_fn):
def add_done_callback(self, cb):
self._cbs.append(cb)
+ @property
+ def fired(self):
+ """Already completed. The tests fire ONE future to release the bulk
+ submit and then fire "the rest"; without this guard the first one runs
+ its done-callback twice and the assembly counts 17 chunks out of 16."""
+ return self._done
+
def done(self):
return self._done
@@ -139,40 +146,66 @@ class _Sel:
class TestWindowedProgressive:
- def test_bounded_in_flight_and_assembly(self):
+ def test_primes_small_then_sends_the_rest_in_one_submit(self):
+ """The UNPINNED lane submits in TWO calls: a small priming batch, then
+ everything else.
+
+ It used to keep a bounded in-flight window and top up on every
+ completion — which meant `lane_cap - outstanding` was 1 in steady state
+ and the batch collapsed to ONE task per submit, i.e. a blocking
+ GIL-held round trip per chunk. Measured on 977 chunks: 970 submits,
+ 12.4 s. There is nothing to balance on an unpinned lane (no placement
+ decision), so the window bought nothing; `distributed` queues root
+ tasks at the scheduler anyway.
+
+ The priming batch is why this is two calls and not one: going straight
+ to a single all-at-once submit doubled time-to-first-chunk (642 ->
+ 1292 ms) because the client serialises the whole graph before anything
+ returns. Priming keeps the first paint fast — measured 45 ms, better
+ than the windowed version's 656 ms — for one extra round trip.
+ """
src, client, handle, seen = _run()
- # 16 chunks total, but only the window (4) submitted up front...
- assert len(client.created) == 4
- # ...and that whole window went out in ONE submit call, not four
- # blocking scheduler round trips.
- assert client.submit_calls == 1
- assert not handle.done()
- # Completing futures tops up the window without ever exceeding it.
- fired = 0
- while fired < len(client.created):
- client.created[fired].fire()
- fired += 1
- in_flight = len(client.created) - fired
- assert in_flight <= 4
- _wait(lambda f=fired: len(client.created) > f or handle.done(),
- timeout=2.0, required=False)
- # Every chunk streamed a callback and the assembly equals the source.
+ # The priming batch only — the bulk submit follows the first completion
+ # (or the wait loop's next top-up), which is what keeps first paint fast.
+ primed = len(client.created)
+ assert primed <= 8, f"priming batch was {primed}, expected <= 8"
+ client.created[0].fire()
+ _wait(lambda: len(client.created) == 16, timeout=5.0)
+ assert client.submit_calls <= 2, (
+ f"{client.submit_calls} submits for 16 chunks — the unpinned lane "
+ f"is back to per-completion top-ups")
+ for f in list(client.created):
+ if not f.fired:
+ f.fire()
np.testing.assert_array_equal(handle.result(timeout=10.0), src.compute())
assert len(seen) == 16
- def test_cancel_stops_submission_and_outstanding(self):
+ def test_cancel_cancels_every_outstanding_chunk(self):
+ """Cancellation still stops the work.
+
+ With the whole job submitted up front there is no pending queue left to
+ withhold, so "stop" means CANCELLING the outstanding futures rather than
+ declining to submit more. That is the property that actually matters —
+ a superseded VI stream must not keep computing through an ROI drag —
+ and dask cancels queued tasks it has not started.
+ """
src, client, handle, seen = _run()
- assert len(client.created) == 4
+ client.created[0].fire() # let the bulk submit go out
+ _wait(lambda: len(client.created) == 16, timeout=5.0)
handle.cancel()
# Cancel is IMMEDIATE (dispatch_chunks' on_start hook wakes the wait
# loop) — a superseded VI stream must not keep computing through an ROI
# drag while the next tick's stream is already running.
- _wait(lambda: all(f.cancelled for f in client.created))
+ # Every future that had NOT already completed. The one fired above to
+ # release the bulk submit is done, and a completed future is not
+ # cancellable — asserting "all cancelled" would just be wrong.
+ _wait(lambda: all(f.cancelled for f in client.created if not f.fired))
# Firing the cancelled futures must not submit more work.
for f in list(client.created):
- f.fire()
+ if not f.fired:
+ f.fire()
_wait(handle.done)
- assert len(client.created) == 4
+ assert len(client.created) == 16, "firing cancelled futures resubmitted work"
def test_stop_event_halts_topups(self):
stop = threading.Event()
@@ -181,7 +214,9 @@ def test_stop_event_halts_topups(self):
for f in list(client.created):
f.fire()
_wait(handle.done)
- assert len(client.created) == 4 # no top-ups after stop
+ # No BULK submit after the stop: only the priming batch was ever sent.
+ # (It was 4 when the lane kept a window; the priming batch is 8.)
+ assert len(client.created) <= 8, "work was submitted after the stop"
def test_error_propagates_via_result(self):
from spyde.drawing.update_functions import compute_with_live_buffer
diff --git a/spyde/toolbars.yaml b/spyde/toolbars.yaml
index 2794dfe0..a751f48e 100644
--- a/spyde/toolbars.yaml
+++ b/spyde/toolbars.yaml
@@ -364,6 +364,50 @@ functions:
toolbar_side: bottom
navigation: False
+ Drift Correction:
+ description: Solve and remove sample drift across an image stack. Check the result in a separate window — an aligned stack sums sharp, a misaligned one blurs — then Apply to add a lazy corrected node (nothing is copied).
+ icon: drawing/toolbars/icons/center_zero_beam.svg
+ function: spyde.actions.drift_action.drift_correction
+ signal_types: [insitu]
+ plot_dim: [2]
+ toolbar_side: bottom
+ navigation: False
+ toggle: True
+
+ Segment Particles:
+ description: Find and measure particles frame by frame. Tune on the displayed frame (classical threshold or a scribble-trained classifier), then run the whole movie into a new particle dataset with count-vs-time and tracks.
+ icon: drawing/toolbars/icons/peak_finding.svg
+ function: spyde.actions.particles_action.segment_particles
+ exclude_signal_types: [spyde_diffraction_vectors_image, particles, electron_diffraction]
+ plot_dim: [2]
+ toolbar_side: bottom
+ navigation: False
+ toggle: True
+
+ Particle Overlay:
+ description: Draw the found particles on the frame — filled outlines coloured by track, labels on the selected one, optional fading motion trails. Double-click a particle to select it; delete, merge and split corrections are recorded on the dataset.
+ icon: drawing/toolbars/icons/peak_finding.svg
+ function: spyde.actions.particle_overlay.particle_overlay
+ requires_particles: True
+ plot_dim: [2]
+ toolbar_side: bottom
+ navigation: False
+ toggle: True
+ # NO `parameters:` block, deliberately. The Electron toolbar renders a
+ # parameter POPOVER with its own Run button for any action that declares
+ # one, so declaring the trail knobs here turned a one-click toggle into a
+ # form the user has to submit — verified in the app: the overlay never drew
+ # because the click only opened the panel. The knobs live in the caret's
+ # schema instead (registry.wizard_parameters("part") → part_tune).
+
+ Particle Lanes:
+ description: Stack count-vs-time, mean size and the birth/death/merge/split event lane as one navigator with a shared time cursor.
+ icon: drawing/toolbars/icons/peak_finding.svg
+ function: spyde.actions.particle_overlay.particle_lanes
+ requires_particles: True
+ plot_dim: [1, 2]
+ toolbar_side: bottom
+
Fit:
description: Fit a model to every spectrum. Add components line by line, watch the model curve on the spectrum, then Run to fit the whole scan on the GPU. Commit turns each component's integrated area into a map.
icon: drawing/toolbars/icons/fit.svg
diff --git a/uv.lock b/uv.lock
index dd98377f..607f2ac1 100644
--- a/uv.lock
+++ b/uv.lock
@@ -47,8 +47,8 @@ wheels = [
[[package]]
name = "anyplotlib"
-version = "0.4.2"
-source = { registry = "https://pypi.org/simple" }
+version = "0.5.0"
+source = { editable = "../anyplotlib" }
dependencies = [
{ name = "anywidget" },
{ name = "colorcet" },
@@ -56,9 +56,36 @@ dependencies = [
{ name = "numpy", version = "2.4.6", source = { registry = "https://pypi.org/simple" }, marker = "python_full_version >= '3.11'" },
{ name = "traitlets" },
]
-sdist = { url = "https://files.pythonhosted.org/packages/26/b2/74147510e82c0f7b8645f8b9f307c398e4ca56df05313588c05ce442de91/anyplotlib-0.4.2.tar.gz", hash = "sha256:32e4cf8315497bb083d834a9d6d91b8333bc7e2b46a8b6a57af6398516b2a123", size = 1403418, upload-time = "2026-07-26T20:53:38.167Z" }
-wheels = [
- { url = "https://files.pythonhosted.org/packages/9e/5a/b7a07b3859d17c3eb865e5e6aeeda4d2eb13964072e6472b1bd6d0e81a5f/anyplotlib-0.4.2-py3-none-any.whl", hash = "sha256:6b965fbc356fff2749e39056499fe7f0563a3228fad2bf99cf522a1f0918e194", size = 282798, upload-time = "2026-07-26T20:53:35.992Z" },
+
+[package.metadata]
+requires-dist = [
+ { name = "anywidget", specifier = ">=0.9.0" },
+ { name = "bokeh", marker = "extra == 'docs'", specifier = ">=3.0" },
+ { name = "colorcet", specifier = ">=3.0" },
+ { name = "jupyterlab", marker = "extra == 'jupyter'", specifier = ">=4.5.5" },
+ { name = "matplotlib", marker = "extra == 'docs'", specifier = ">=3.7" },
+ { name = "numpy", specifier = ">=2.0.0" },
+ { name = "pillow", marker = "extra == 'docs'", specifier = ">=10.0" },
+ { name = "playwright", marker = "extra == 'docs'", specifier = ">=1.58.0" },
+ { name = "plotly", marker = "extra == 'docs'", specifier = ">=5.0" },
+ { name = "pydata-sphinx-theme", marker = "extra == 'docs'", specifier = ">=0.16" },
+ { name = "scipy", marker = "extra == 'docs'", specifier = ">=1.15.3" },
+ { name = "sphinx", marker = "extra == 'docs'", specifier = ">=8.0" },
+ { name = "sphinx-design", marker = "extra == 'docs'", specifier = ">=0.6" },
+ { name = "sphinx-gallery", marker = "extra == 'docs'", specifier = ">=0.18" },
+ { name = "traitlets", specifier = ">=5.0.0" },
+]
+provides-extras = ["docs", "jupyter"]
+
+[package.metadata.requires-dev]
+dev = [
+ { name = "docutils", specifier = ">=0.19" },
+ { name = "playwright", specifier = ">=1.58.0" },
+ { name = "pytest", specifier = ">=9.0.2" },
+ { name = "pytest-cov", specifier = ">=5.0.0" },
+ { name = "scipy", specifier = ">=1.15.3" },
+ { name = "sphinx", specifier = ">=8.0" },
+ { name = "towncrier", specifier = ">=24.0.0" },
]
[[package]]
@@ -3622,7 +3649,7 @@ tests = [
[package.metadata]
requires-dist = [
- { name = "anyplotlib", specifier = ">=0.4.2" },
+ { name = "anyplotlib", editable = "../anyplotlib" },
{ name = "atomap", marker = "extra == 'atoms'", specifier = ">=0.4.1" },
{ name = "bokeh" },
{ name = "dask" },