"""Streak (scanline) assignment from the stage trajectory.
Reconstructs, for every frame, which scan row (streak) it belongs to and its
position within that row, by segmenting the stage motion on the sign of the
x-velocity. A "sweep" is a maximal run where the (windowed) horizontal speed is
above the measured noise floor and ``sign(vx)`` is constant: direction reversals
split rows regardless of x amplitude (so partial / mid-sweep first rows are
handled), while the parked lead-in and turnarounds (``vx ~ 0``) fall out between
sweeps. Only *complete* sweeps (spanning ~the full x width) are kept.
Writes three datasets under ``entry/data``:
``streak_id`` per (dense) frame: row index, or -1 where not scanning.
``streak_frame`` per (dense) frame: position within the row; resets to
``frame_id_start`` each streak and counts *missing* frames
so a dropped frame keeps its column slot.
``streak_endpoints`` compound table, one row per streak, with the start/end
stage coordinates and direction.
This is the replacement for the ``streakdirection`` produced by
:mod:`coseda.map.findscanlines`.
"""
from __future__ import annotations
from datetime import datetime
import h5py
import numpy as np
from scipy.ndimage import uniform_filter1d
from coseda.logging_utils import log_print
__all__ = ["assign_streaks", "find_sweeps", "rasterize_arrays", "rasterize_map",
"plot_rasterized_map", "add_scalebar", "write_atlas", "save_atlas",
"load_atlas"]
ATLAS_GROUP = "atlas"
# coseda stage positions are in metres; adjust if a dataset uses other units.
STAGE_UNIT_TO_M = 1.0
# "nice" 1-2-5 scalebar lengths, in nanometres.
_SCALEBAR_STEPS_NM = [10, 20, 50, 100, 200, 500, 1e3, 2e3, 5e3, 1e4, 2e4, 5e4,
1e5, 2e5, 5e5, 1e6, 2e6, 5e6, 1e7]
PLACEHOLDER = 999999.0
STREAK_ENDPOINT_DTYPE = np.dtype([
("streak_id", "<i4"),
("x_start", "<f8"), ("y_start", "<f8"),
("x_stop", "<f8"), ("y_stop", "<f8"),
("direction", "i1"),
])
def _longest_run(mask, gap=0):
"""Return (start, stop) of the longest True run in ``mask``, tolerating gaps
of up to ``gap`` False frames; ``None`` if there is no True frame."""
idx = np.flatnonzero(mask)
if len(idx) == 0:
return None
breaks = np.flatnonzero(np.diff(idx) > gap + 1)
starts = np.r_[idx[0], idx[breaks + 1]]
stops = np.r_[idx[breaks], idx[-1]] + 1
k = int(np.argmax(stops - starts))
return int(starts[k]), int(stops[k])
[docs]
def find_sweeps(x, y, window=17, noise_k=10.0, gap=3, min_frames=20,
min_span_frac=0.7):
"""Segment the stage trajectory into complete scan sweeps.
Parameters
----------
x, y : array_like
Stage x/y position per (valid) frame.
window : int
Smoothing width (frames) applied to positions before differentiating;
averages out encoder quantization. Forced odd.
noise_k : float
Motion threshold is ``noise_k * median(|vy|)`` -- the noise floor is
measured from the vertical speed, which rests at the shared x/y noise
level while a row is scanned in x.
gap : int
Bridge brief sub-threshold dips of up to this many frames within a sweep.
min_frames : int
Ignore runs shorter than this.
min_span_frac : float
A sweep is "complete" if its x-span is at least this fraction of the
median sweep x-span (drops partial / mid-sweep rows).
Returns
-------
kept : list of dict
Complete sweeps with ``streak_id``, ``start_i``, ``stop_i`` (valid-frame
index, half-open) and ``direction`` (+1/-1).
extras : dict
``raw`` (all sweeps before the completeness filter), ``dropped``,
``vx``/``vy`` (windowed speeds) and ``thr`` -- for diagnostics.
"""
x = np.asarray(x, dtype=float)
y = np.asarray(y, dtype=float)
win = window if window % 2 else window + 1
vx = np.gradient(uniform_filter1d(x, size=win, mode="nearest"))
vy = np.abs(np.gradient(uniform_filter1d(y, size=win, mode="nearest")))
thr = noise_k * np.median(vy)
d = np.where(np.abs(vx) > thr, np.sign(vx), 0).astype(np.int8)
raw, i, n = [], 0, len(d)
while i < n:
if d[i] == 0:
i += 1
continue
j = i
while j < n:
if d[j] == d[i]:
j += 1
elif d[j] == 0 and np.all(d[j:j + gap + 1] != -d[i]):
j += 1 # brief dip, not a real reversal
else:
break
stop = j
while stop > i and d[stop - 1] == 0: # trim trailing zeros the bridge added
stop -= 1
if stop - i >= min_frames:
raw.append({"start_i": i, "stop_i": stop, "direction": int(d[i]),
"span": abs(float(x[stop - 1] - x[i]))})
i = j
if raw:
med = np.median([s["span"] for s in raw])
kept = [s for s in raw if s["span"] >= min_span_frac * med]
else:
kept = []
for k, s in enumerate(kept):
s["streak_id"] = k
return kept, {"raw": raw, "dropped": [s for s in raw if s not in kept],
"vx": vx, "vy": vy, "thr": thr}
[docs]
def assign_streaks(input_path, window=17, noise_k=10.0, gap=3, min_frames=20,
min_span_frac=0.7, frame_id_start=1, write=True):
"""Assign streak ids / within-streak frame ids for a coseda HDF5 file.
Reads the stage trajectory (``stagepos_x_refined``/``stagepos_y_refined``,
falling back to the unrefined datasets), detects complete sweeps with
:func:`find_sweeps`, and builds the dense per-frame ``streak_id`` /
``streak_frame`` labels plus a ``streak_endpoints`` table. Frames whose
images were stripped (``index == -1``) are still labelled when their stage
positions are present, so full-length intensity maps can be regenerated.
Parameters
----------
input_path : str or pathlib.Path
Path to the coseda HDF5 file.
window, noise_k, gap, min_frames, min_span_frac
Passed through to :func:`find_sweeps`.
frame_id_start : int
Value of ``streak_frame`` at each streak's first frame.
write : bool
If True, write ``streak_id`` / ``streak_frame`` / ``streak_endpoints``
into ``entry/data`` (replacing any existing versions). If False, only
compute and return them.
Returns
-------
dict
``streak_id``, ``streak_frame`` (dense int32 arrays), ``endpoints``
(structured array), ``spans`` and ``extras`` (from :func:`find_sweeps`).
"""
with h5py.File(input_path, "r") as f:
g = f["entry/data"]
xf = g["stagepos_x_refined"][:] if "stagepos_x_refined" in g else g["stagepos_x"][:]
yf = g["stagepos_y_refined"][:] if "stagepos_y_refined" in g else g["stagepos_y"][:]
n_dense = len(xf)
index = g["index"][:] if "index" in g else np.arange(n_dense)
valid = (np.isfinite(xf) & np.isfinite(yf)
& (xf != PLACEHOLDER) & (yf != PLACEHOLDER))
x, y = xf[valid], yf[valid]
valid_pos = np.flatnonzero(valid) # dense positions with stage coordinates
spans, extras = find_sweeps(x, y, window=window, noise_k=noise_k, gap=gap,
min_frames=min_frames, min_span_frac=min_span_frac)
# per-streak start/end stage coordinates
endpoints = np.empty(len(spans), dtype=STREAK_ENDPOINT_DTYPE)
for k, s in enumerate(spans):
a, b = s["start_i"], s["stop_i"] - 1
endpoints[k] = (s["streak_id"], x[a], y[a], x[b], y[b], s["direction"])
# dense per-frame labelling; stripped-image frames keep their slot
streak_id = np.full(n_dense, -1, dtype=np.int32)
streak_frame = np.full(n_dense, -1, dtype=np.int32)
for s in spans:
p0 = valid_pos[s["start_i"]]
p1 = valid_pos[s["stop_i"] - 1]
streak_id[p0:p1 + 1] = s["streak_id"]
streak_frame[p0:p1 + 1] = np.arange(p1 - p0 + 1) + frame_id_start
missing = index == -1
log_print(f"Detected {len(spans)} complete streaks "
f"(dropped {len(extras['dropped'])} partial sweeps).")
log_print(f"Labelled {(streak_id != -1).sum()} frames "
f"(present {((streak_id != -1) & ~missing).sum()}, "
f"missing {((streak_id != -1) & missing).sum()}); "
f"{(streak_id == -1).sum()} unlabelled.")
if write:
with h5py.File(input_path, "r+") as f:
g = f["entry/data"]
for name, data in (("streak_id", streak_id),
("streak_frame", streak_frame),
("streak_endpoints", endpoints)):
if name in g:
del g[name]
g.create_dataset(name, data=data)
log_print(f"Wrote streak_id, streak_frame, streak_endpoints to {input_path}")
return {"streak_id": streak_id, "streak_frame": streak_frame,
"endpoints": endpoints, "spans": spans, "extras": extras}
def _resample_row(x_reg, xs, zs, max_gap_frac):
"""Resample one streak's values onto the common x-grid.
Duplicate x readings are averaged (keeps np.interp monotonic); pixels outside
the streak's x-range, or bridging a gap wider than ``max_gap_frac`` * the
median sample spacing (a run of missing frames), are left NaN.
"""
order = np.argsort(xs)
xs, zs = xs[order], zs[order]
ux, inv = np.unique(xs, return_inverse=True)
if ux.size < 2:
return None
acc = np.zeros(ux.size)
cnt = np.zeros(ux.size)
np.add.at(acc, inv, zs)
np.add.at(cnt, inv, 1.0)
uz = acc / np.maximum(cnt, 1.0)
row = np.interp(x_reg, ux, uz, left=np.nan, right=np.nan)
gap_thresh = max_gap_frac * float(np.median(np.diff(ux)))
j = np.clip(np.searchsorted(ux, x_reg), 1, ux.size - 1)
bracket = ux[j] - ux[j - 1] # width of the sample interval each pixel falls in
row[bracket > gap_thresh] = np.nan # keep missing-frame gaps empty
return row
def _streak_direction_map(streak_id, x, valid, streak_directions=None):
"""Return per-streak scan directions, using metadata with an x-start/end fallback."""
directions = dict(streak_directions or {})
for sid in np.unique(streak_id[valid]):
sid = int(sid)
if sid in directions:
continue
idx = np.flatnonzero(valid & (streak_id == sid))
if idx.size < 2:
continue
dx = float(x[idx[-1]] - x[idx[0]])
if dx != 0:
directions[sid] = 1 if dx > 0 else -1
return directions
[docs]
def rasterize_arrays(streak_id, x, y, z, valid=None, out_cols=None, max_gap_frac=3.0,
left_streak_offset_m=0.0, streak_directions=None):
"""Rasterize per-frame arrays into a regular physical grid (array-level core).
Rows are streaks (from ``streak_id``); within each streak ``z`` is resampled
onto a common physical-x grid using the measured ``x``. Because ``x``/``y`` are
absolute stage coordinates, both sweep directions share the same axis -- no
velocity model, flip or backlash needed. ``valid`` (bool, per frame) selects
the frames to use; it is always AND-ed with ``streak_id >= 0`` and finite
coords/values. Rows are ordered by measured y so the raster is upright.
Returns ``{raster, extent, row_y, streak_ids}`` with NaN where there is no data.
"""
streak_id = np.asarray(streak_id)
x = np.asarray(x, dtype=np.float64)
y = np.asarray(y, dtype=np.float64)
z = np.asarray(z, dtype=np.float64)
finite = (np.isfinite(x) & np.isfinite(y) & np.isfinite(z)
& (x != PLACEHOLDER) & (y != PLACEHOLDER))
valid = finite if valid is None else (np.asarray(valid, dtype=bool) & finite)
valid = valid & (streak_id >= 0)
left_streak_offset_m = float(left_streak_offset_m or 0.0)
if left_streak_offset_m != 0.0:
directions = _streak_direction_map(streak_id, x, valid, streak_directions)
left_ids = [sid for sid, direction in directions.items() if direction < 0]
if left_ids:
x = x.copy()
x[np.isin(streak_id, left_ids)] += left_streak_offset_m
sids = np.unique(streak_id[valid])
if sids.size == 0:
raise ValueError("No labelled frames to rasterize.")
x_min, x_max = float(np.min(x[valid])), float(np.max(x[valid]))
if out_cols is None: # ~ one pixel per frame: median frames per streak
counts = [int(np.count_nonzero((streak_id == s) & valid)) for s in sids]
out_cols = max(int(np.median(counts)), 2)
x_reg = np.linspace(x_min, x_max, out_cols)
rows = np.full((sids.size, out_cols), np.nan)
row_y = np.full(sids.size, np.nan)
for i, s in enumerate(sids):
m = (streak_id == s) & valid
row = _resample_row(x_reg, x[m], z[m], max_gap_frac)
if row is not None:
rows[i] = row
row_y[i] = float(np.median(y[m]))
order = np.argsort(row_y) # physically upright regardless of scan direction
rows, row_y, sids = rows[order], row_y[order], sids[order]
extent = [x_min, x_max, float(row_y[0]), float(row_y[-1])]
return {
"raster": rows,
"extent": extent,
"row_y": row_y,
"streak_ids": sids,
"left_streak_offset_m": left_streak_offset_m,
}
[docs]
def rasterize_map(input_path, zdim="frame_mean_intensities", out_cols=None,
max_gap_frac=3.0, left_streak_offset_um=0.0):
"""Rasterize a scan file onto a regular physical grid using measured coords.
Thin file wrapper over :func:`rasterize_arrays`: reads the stage trajectory,
``streak_id`` and the ``zdim`` value. Full-length ``zdim`` arrays are used
on the dense frame axis so atlases can be regenerated after stripping.
Compact per-image arrays are mapped onto the dense axis via ``index`` and
only available images contribute.
Requires ``assign_streaks`` to have been run (needs ``streak_id``).
"""
with h5py.File(input_path, "r") as f:
g = f["entry/data"]
x = g["stagepos_x_refined"][:] if "stagepos_x_refined" in g else g["stagepos_x"][:]
y = g["stagepos_y_refined"][:] if "stagepos_y_refined" in g else g["stagepos_y"][:]
if "streak_id" not in g:
raise KeyError("streak_id not found -- run assign_streaks first.")
streak_id = g["streak_id"][:]
streak_directions = {}
if "streak_endpoints" in g:
endpoints = g["streak_endpoints"][:]
if endpoints.dtype.names and {"streak_id", "direction"} <= set(endpoints.dtype.names):
streak_directions = {
int(row["streak_id"]): int(row["direction"])
for row in endpoints
if int(row["streak_id"]) >= 0 and int(row["direction"]) != 0
}
index = g["index"][:] if "index" in g else np.arange(len(x))
if zdim not in g:
raise KeyError(f"zdim '{zdim}' not found under entry/data")
z = g[zdim][:].astype(np.float64)
n = len(x)
available = index != -1
if len(z) == n:
z_dense = z
valid = np.ones(n, dtype=bool)
else:
# compact (per-image) zdim -> dense via index, like the GUI map view
z_dense = np.full(n, np.nan)
z_dense[available] = z[index[available]]
valid = available
out = rasterize_arrays(
streak_id, x, y, z_dense, valid=valid,
out_cols=out_cols, max_gap_frac=max_gap_frac,
left_streak_offset_m=float(left_streak_offset_um or 0.0) * 1e-6,
streak_directions=streak_directions,
)
out["zdim"] = zdim
log_print(f"Rasterized '{zdim}': {out['raster'].shape[0]} streaks x "
f"{out['raster'].shape[1]} px, "
f"{int(np.count_nonzero(np.isnan(out['raster'])))} empty px.")
return out
[docs]
def plot_rasterized_map(input_path, zdim="frame_mean_intensities", cmap="viridis",
vmin_pct=1.0, vmax_pct=99.0, ax=None, **kwargs):
"""Render :func:`rasterize_map` as an image (drop-in for the old scatter map)."""
import matplotlib.pyplot as plt
out = rasterize_map(input_path, zdim=zdim, **kwargs)
raster = out["raster"]
finite = raster[np.isfinite(raster)]
vmin = float(np.percentile(finite, vmin_pct)) if finite.size else 0.0
vmax = float(np.percentile(finite, vmax_pct)) if finite.size else 1.0
if not np.isfinite(vmax) or vmax <= vmin:
vmax = vmin + 1.0
if ax is None:
_fig, ax = plt.subplots(figsize=(9, 7))
im = ax.imshow(raster, extent=out["extent"], origin="lower", aspect="auto",
cmap=cmap, vmin=vmin, vmax=vmax, interpolation="nearest")
ax.set_xlabel("stage x")
ax.set_ylabel("stage y")
ax.set_title(f"Rasterized map ({zdim})")
ax.figure.colorbar(im, ax=ax, label=zdim)
add_scalebar(ax)
return ax, out
[docs]
def add_scalebar(ax, data_to_m=STAGE_UNIT_TO_M, frac=0.15, loc="lower left",
color="white"):
"""Draw a 1-2-5 'nice' scalebar on a map axis (labelled nm / µm / mm).
The bar is ~``frac`` of the x-range; ``data_to_m`` converts axis data units to
metres (coseda stage positions are metres -> 1.0). A translucent dark box keeps
it legible over any colormap. Returns the artist, or None if unavailable.
"""
try:
from mpl_toolkits.axes_grid1.anchored_artists import AnchoredSizeBar
from matplotlib.font_manager import FontProperties
except Exception:
return None
x0, x1 = ax.get_xlim()
y0, y1 = ax.get_ylim()
span = abs(x1 - x0)
yspan = abs(y1 - y0)
if not (span > 0 and np.isfinite(span)):
return None
target_nm = frac * span * data_to_m * 1e9
nice_nm = _SCALEBAR_STEPS_NM[0]
for step in _SCALEBAR_STEPS_NM:
if step <= target_nm:
nice_nm = step
length_data = (nice_nm * 1e-9) / data_to_m
if nice_nm < 1e3:
label = f"{nice_nm:g} nm"
elif nice_nm < 1e6:
label = f"{nice_nm / 1e3:g} µm"
else:
label = f"{nice_nm / 1e6:g} mm"
# Thickness from the y-range (not x), so the bar stays thin whatever the aspect.
bar = AnchoredSizeBar(
ax.transData, length_data, label, loc, color=color, frameon=True,
size_vertical=yspan * 0.012, pad=0.4, borderpad=0.6, sep=4,
fontproperties=FontProperties(size=10),
)
bar.patch.set(facecolor="black", alpha=0.5, edgecolor="none")
ax.add_artist(bar)
return bar
[docs]
def write_atlas(input_path, zdim, out, group=ATLAS_GROUP):
"""Persist an already-computed raster into ``entry/<group>/<zdim>``.
Stores the 2-D ``raster`` plus the physical ``x`` (per column) and ``y`` (per
row) coordinates, with ``extent`` and provenance as group attrs, so the atlas
can be re-rendered or exported later without recomputing. Overwrites any
existing atlas for this zdim. ``out`` is a dict from :func:`rasterize_arrays`
/ :func:`rasterize_map` (``raster``, ``extent``, ``row_y``, ``streak_ids``).
"""
raster = np.asarray(out["raster"], dtype=np.float32)
x_min, x_max = float(out["extent"][0]), float(out["extent"][1])
x_coords = (np.linspace(x_min, x_max, raster.shape[1]) if raster.shape[1] > 1
else np.array([x_min], dtype=np.float64))
with h5py.File(input_path, "r+") as f:
root = f.require_group(f"entry/{group}")
if zdim in root:
del root[zdim]
gz = root.create_group(zdim)
gz.create_dataset("raster", data=raster)
gz.create_dataset("x", data=x_coords.astype(np.float64))
gz.create_dataset("y", data=np.asarray(out["row_y"], dtype=np.float64))
gz.create_dataset("streak_ids", data=np.asarray(out["streak_ids"], dtype=np.int32))
gz.attrs["extent"] = np.asarray(out["extent"], dtype=np.float64)
gz.attrs["zdim"] = str(zdim)
gz.attrs["left_streak_offset_um"] = float(out.get("left_streak_offset_m", 0.0)) * 1e6
gz.attrs["created_at"] = datetime.utcnow().replace(microsecond=0).isoformat() + "Z"
log_print(f"Saved atlas 'entry/{group}/{zdim}' ({raster.shape[0]}x{raster.shape[1]} px)")
[docs]
def save_atlas(input_path, zdim="frame_mean_intensities", group=ATLAS_GROUP, **kwargs):
"""Rasterize ``zdim`` and persist it into the file (convenience wrapper)."""
out = rasterize_map(input_path, zdim=zdim, **kwargs)
write_atlas(input_path, zdim, out, group=group)
return out
[docs]
def load_atlas(input_path, zdim="frame_mean_intensities", group=ATLAS_GROUP):
"""Load a saved atlas from ``entry/<group>/<zdim>``, or None if absent."""
path = f"entry/{group}/{zdim}"
with h5py.File(input_path, "r") as f:
if path not in f:
return None
gz = f[path]
return {
"raster": gz["raster"][:],
"extent": [float(v) for v in gz.attrs["extent"]],
"x": gz["x"][:],
"row_y": gz["y"][:],
"streak_ids": gz["streak_ids"][:],
"zdim": zdim,
"left_streak_offset_um": float(gz.attrs.get("left_streak_offset_um", 0.0)),
"created_at": gz.attrs.get("created_at"),
}