diff --git a/README.md b/README.md index 4e6f0d2..bc88afd 100644 --- a/README.md +++ b/README.md @@ -26,13 +26,34 @@ tiled server config. ### `dls_tiled.visr:visr_router` `GET /api/v1/binned/{uid}` bins a ViSR scan into a 2-D image on the server and -returns the mean `RedTotal`, `GreenTotal` and `BlueTotal` in each bin, plus the -bin edges as `x_limits` and `y_limits`. Positions come from the recorded -readbacks, or from the ScanSpec setpoints in the run's start document with -`setpoints=true`. Other query parameters: `x_dim_index`, `y_dim_index`, `width`, -`height`, `xmin`/`xmax`/`ymin`/`ymax` and the repeatable +returns the mean `RedTotal`, `GreenTotal` and `BlueTotal` in each bin as +matrices with one row per y bin and one column per x bin (`null` where no point +landed), the bin edges as `x_limits` and `y_limits`, and `n_points`/`n_total`, +the points used so far and the points the scan will have (`null` without a +spec). Positions come from the recorded readbacks, or from the ScanSpec +setpoints in the run's start document with `setpoints=true`. + +Unless `width`/`height` or a full `xmin`/`xmax`/`ymin`/`ymax` range are given, +the bins follow the scan's setpoint grid from its spec (one cell per setpoint, +fixed from the first poll), provided the spec is a grid of at least 3×3 and the +recorded positions lie on it. Positions that are off the grid by a constant +amount per axis (as fly scans currently record them) are first shifted onto it +by the median of each point's position minus its setpoint, and the shift is +returned as `position_offset`. Otherwise numpy's default 10×10 over the data. +Other query parameters: `x_dim_index`, `y_dim_index` and the repeatable `slice_dim=dim:center:thickness`. +A file still behind the catalog is retried briefly and then answered with 503 +and `Retry-After`. Identical requests share one result for a short while and +carry an `ETag`, so a repeated poll of an unchanged image gets a `304`. + +| Environment variable | Default | | +|---|---|---| +| `DLS_TILED_VISR_NOT_READY_RETRIES` | `3` | retries before answering 503 | +| `DLS_TILED_VISR_NOT_READY_RETRY_DELAY` | `0.2` | seconds between those retries | +| `DLS_TILED_VISR_CACHE_SECONDS` | `0.5` | how long a result is shared; `0` turns sharing off | +| `DLS_TILED_VISR_MAX_GRID_CELLS` | `100000` | largest scan grid used by default | + ```yaml routers: - dls_tiled.visr:visr_router diff --git a/src/dls_tiled/rendering.py b/src/dls_tiled/rendering.py new file mode 100644 index 0000000..ef3ad46 --- /dev/null +++ b/src/dls_tiled/rendering.py @@ -0,0 +1,303 @@ +"""Draw scan points into an image for a viewport (prototype). + +Each renderer takes the points of one channel (values and positions) and a +``Viewport`` (a rectangle in data coordinates and an output size in pixels), and +returns two ``(height, width)`` arrays: the image (NaN where nothing is drawn) +and the coverage (0 to 1, how much of the pixel is covered by the drawn points). +Row 0 is the viewport's ``y0`` edge and column 0 its ``x0`` edge. + +- ``cells``: the scan's grid cells, each with the mean of the points in it. +- ``rectangles``: a rectangle of the nominal cell size centred on each recorded + position, so the picture follows where the points really were. +- ``gaussian``: a beam-sized 2-D Gaussian per point, so any gaps the scan left + between beam footprints show up as low coverage. +- ovals (``oval_layout``): a beam-sized oval per point, filled with its value + and faintly outlined; overlaps are split between the ovals, not blended. +""" + +from __future__ import annotations + +import struct +import zlib +from dataclasses import dataclass + +import numpy + +FWHM_TO_SIGMA = 1 / (2 * numpy.sqrt(2 * numpy.log(2))) + + +@dataclass(frozen=True) +class Viewport: + """A rectangle in data coordinates drawn at ``width`` x ``height`` pixels.""" + + x0: float + x1: float + y0: float + y1: float + width: int + height: int + + @classmethod + def fit(cls, x0, x1, y0, y1, size): + """The viewport with ``size`` pixels along its longer side, square pixels.""" + span_x, span_y = abs(x1 - x0) or 1.0, abs(y1 - y0) or 1.0 + if span_x >= span_y: + width, height = size, max(1, round(size * span_y / span_x)) + else: + width, height = max(1, round(size * span_x / span_y)), size + return cls(float(x0), float(x1), float(y0), float(y1), int(width), int(height)) + + @property + def pixel_x(self) -> float: + return (self.x1 - self.x0) / self.width + + @property + def pixel_y(self) -> float: + return (self.y1 - self.y0) / self.height + + def centres(self): + """Pixel centres along x (columns) and y (rows).""" + xs = self.x0 + (numpy.arange(self.width) + 0.5) * self.pixel_x + ys = self.y0 + (numpy.arange(self.height) + 0.5) * self.pixel_y + return xs, ys + + def column(self, x): + """Fractional column of a data x coordinate.""" + return (numpy.asarray(x, dtype=float) - self.x0) / self.pixel_x + + def row(self, y): + return (numpy.asarray(y, dtype=float) - self.y0) / self.pixel_y + + +def _finish(total, weight, coverage): + image = numpy.full(total.shape, numpy.nan) + drawn = weight > 0 + image[drawn] = total[drawn] / weight[drawn] + return image, numpy.clip(coverage, 0.0, 1.0) + + +def render_cells(values, x, y, x_edges, y_edges, viewport: Viewport): + """The scan's grid cells (mean of the points in each), sampled at each pixel.""" + values = numpy.asarray(values, dtype=float) + counts, _, _ = numpy.histogram2d(x, y, bins=[x_edges, y_edges]) + sums, _, _ = numpy.histogram2d(x, y, bins=[x_edges, y_edges], weights=values) + xs, ys = viewport.centres() + col = numpy.searchsorted(x_edges, xs, side="right") - 1 + row = numpy.searchsorted(y_edges, ys, side="right") - 1 + inside_x = (col >= 0) & (col < len(x_edges) - 1) + inside_y = (row >= 0) & (row < len(y_edges) - 1) + total = numpy.zeros((viewport.height, viewport.width)) + weight = numpy.zeros_like(total) + rr, cc = numpy.meshgrid( + numpy.where(inside_y)[0], numpy.where(inside_x)[0], indexing="ij" + ) + cell_y, cell_x = row[rr], col[cc] + total[rr, cc] = sums[cell_x, cell_y] + weight[rr, cc] = counts[cell_x, cell_y] + return _finish(total, weight, (weight > 0).astype(float)) + + +def render_rectangles(values, x, y, cell_x, cell_y, viewport: Viewport): + """A rectangle of ``cell_x`` by ``cell_y`` centred on each point. + + Where rectangles overlap, the pixel takes the mean of their values, weighted + by how much of the pixel each one covers. + """ + values = numpy.asarray(values, dtype=float) + total = numpy.zeros((viewport.height, viewport.width)) + weight = numpy.zeros_like(total) + for v, px, py in zip(values, viewport.column(x), viewport.row(y), strict=True): + if not numpy.isfinite(v): + continue + cx0, cx1 = ( + px - cell_x / 2 / viewport.pixel_x, + px + cell_x / 2 / viewport.pixel_x, + ) + cy0, cy1 = ( + py - cell_y / 2 / viewport.pixel_y, + py + cell_y / 2 / viewport.pixel_y, + ) + c_lo, c_hi = ( + max(0, int(numpy.floor(cx0))), + min(viewport.width, int(numpy.ceil(cx1))), + ) + r_lo, r_hi = ( + max(0, int(numpy.floor(cy0))), + min(viewport.height, int(numpy.ceil(cy1))), + ) + if c_lo >= c_hi or r_lo >= r_hi: + continue + cols = numpy.arange(c_lo, c_hi) + rows = numpy.arange(r_lo, r_hi) + # fraction of each pixel the rectangle covers, along each axis + fx = numpy.clip(numpy.minimum(cols + 1, cx1) - numpy.maximum(cols, cx0), 0, 1) + fy = numpy.clip(numpy.minimum(rows + 1, cy1) - numpy.maximum(rows, cy0), 0, 1) + overlap = numpy.outer(fy, fx) + total[r_lo:r_hi, c_lo:c_hi] += overlap * v + weight[r_lo:r_hi, c_lo:c_hi] += overlap + return _finish(total, weight, weight) + + +def render_gaussian(values, x, y, fwhm_x, fwhm_y, viewport: Viewport, reach=3.0): + """A 2-D Gaussian of the beam's size (FWHM in data units) on each point. + + Each pixel takes the mean of the points' values weighted by their Gaussians. + Coverage is the summed Gaussian weight, 1 at the centre of an isolated beam, + so the space between beam footprints that the scan didn't cover shows up as + low coverage rather than being filled in. + """ + values = numpy.asarray(values, dtype=float) + sx = fwhm_x * FWHM_TO_SIGMA / viewport.pixel_x # sigma in pixels + sy = fwhm_y * FWHM_TO_SIGMA / viewport.pixel_y + total = numpy.zeros((viewport.height, viewport.width)) + weight = numpy.zeros_like(total) + for v, px, py in zip(values, viewport.column(x), viewport.row(y), strict=True): + if not numpy.isfinite(v): + continue + c_lo, c_hi = ( + max(0, int(px - reach * sx)), + min(viewport.width, int(px + reach * sx) + 1), + ) + r_lo, r_hi = ( + max(0, int(py - reach * sy)), + min(viewport.height, int(py + reach * sy) + 1), + ) + if c_lo >= c_hi or r_lo >= r_hi: + continue + gx = numpy.exp(-0.5 * ((numpy.arange(c_lo, c_hi) + 0.5 - px) / sx) ** 2) + gy = numpy.exp(-0.5 * ((numpy.arange(r_lo, r_hi) + 0.5 - py) / sy) ** 2) + g = numpy.outer(gy, gx) + total[r_lo:r_hi, c_lo:c_hi] += g * v + weight[r_lo:r_hi, c_lo:c_hi] += g + image, coverage = _finish(total, weight, weight) + image[coverage < 0.02] = numpy.nan # beyond the beams' reach: nothing drawn + return image, coverage + + +@dataclass(frozen=True) +class OvalLayout: + """Where each point's oval is drawn in a viewport, independent of any channel. + + ``owner`` is the index of the point each pixel takes its colour from (-1 + outside every oval) and ``outline`` the opacity (0 to 1) of the white + outlines drawn over the colours. + """ + + owner: numpy.ndarray + outline: numpy.ndarray + + def paint(self, values): + """One channel's image and coverage: each oval filled with its point's value.""" + values = numpy.asarray(values, dtype=float) + image = numpy.full(self.owner.shape, numpy.nan) + inside = self.owner >= 0 + image[inside] = values[self.owner[inside]] + return image, (~numpy.isnan(image)).astype(float) + + +def oval_layout( + x, y, width_x, width_y, viewport: Viewport, line_width=1.5, line_opacity=0.35 +) -> OvalLayout: + """An oval ``width_x`` by ``width_y`` (data units) centred on each point. + + Where ovals overlap, the overlap is split along the line through the two + points where their outlines cross, each side taking the colour of its own + oval. For ovals of one size that line is where the two centres are equally + far in units of the ovals' radii, so each pixel takes the oval whose + centre is nearest in those units; that also splits three-way overlaps + consistently. + + Every oval's whole outline is drawn as a faint white line, ``line_width`` + pixels wide at ``line_opacity`` (0 for none). + """ + a = width_x / 2 / viewport.pixel_x # semi-axes in pixels + b = width_y / 2 / viewport.pixel_y + px, py = viewport.column(x), viewport.row(y) + shape = (viewport.height, viewport.width) + best = numpy.full(shape, numpy.inf) # the nearest centre, in radii squared + owner = numpy.full(shape, -1) + outline = numpy.zeros(shape) + pad = line_width + 1 # reach past the oval far enough to draw its outline + for i in numpy.flatnonzero(numpy.isfinite(px) & numpy.isfinite(py)): + c_lo = max(0, int(numpy.floor(px[i] - a - pad))) + c_hi = min(viewport.width, int(numpy.ceil(px[i] + a + pad))) + r_lo = max(0, int(numpy.floor(py[i] - b - pad))) + r_hi = min(viewport.height, int(numpy.ceil(py[i] + b + pad))) + if c_lo >= c_hi or r_lo >= r_hi: + continue + window = (slice(r_lo, r_hi), slice(c_lo, c_hi)) + dx = (numpy.arange(c_lo, c_hi) + 0.5 - px[i]) / a # in radii + dy = (numpy.arange(r_lo, r_hi) + 0.5 - py[i]) / b + q = dy[:, None] ** 2 + dx[None, :] ** 2 # < 1 inside the oval + nearer = q < best[window] + best[window][nearer] = q[nearer] + owner[window][nearer] = i + if line_opacity > 0: + # distance to the outline in pixels, to first order + r = numpy.sqrt(q) + with numpy.errstate(divide="ignore", invalid="ignore"): + gradient = numpy.hypot(dx[None, :] / a, dy[:, None] / b) / r + distance = numpy.where(r > 0, (r - 1) / gradient, numpy.inf) + line = numpy.clip(line_width / 2 + 0.5 - numpy.abs(distance), 0.0, 1.0) + outline[window] = numpy.maximum(outline[window], line * line_opacity) + return OvalLayout(numpy.where(best < 1, owner, -1), outline) + + +def to_rows(image) -> list[list[float | None]]: + """An image as JSON rows, ``None`` where nothing was drawn or the value isn't + a number JSON can carry (NaN or infinite).""" + out = image.astype(object) + out[~numpy.isfinite(image)] = None + return out.tolist() + + +CHANNEL_COLOUR = {"RedTotal": 0, "GreenTotal": 1, "BlueTotal": 2} + + +def to_png(image, coverage, colour=0, vmin=None, vmax=None, outline=None) -> bytes: + """An RGBA PNG of one channel, tinted (0/1/2 = red/green/blue), transparent + where nothing was drawn, with an optional white ``outline`` layer (opacity + 0 to 1) drawn over it. Row 0 of the image is the top row of the PNG.""" + valid = image[~numpy.isnan(image)] + lo = float(valid.min()) if vmin is None and valid.size else (vmin or 0.0) + hi = float(valid.max()) if vmax is None and valid.size else (vmax or 1.0) + scaled = numpy.zeros(image.shape) + if hi > lo: + scaled = numpy.clip((numpy.nan_to_num(image, nan=lo) - lo) / (hi - lo), 0, 1) + rgb = numpy.zeros(image.shape + (3,)) + rgb[..., colour] = scaled + alpha = numpy.where(numpy.isnan(image), 0.0, coverage) + if outline is not None: + # white over the colour, where the colour itself may be transparent + total = outline + alpha * (1 - outline) + with numpy.errstate(divide="ignore", invalid="ignore"): + rgb = numpy.where( + total[..., None] > 0, + (outline[..., None] + rgb * (alpha * (1 - outline))[..., None]) + / total[..., None], + 0.0, + ) + alpha = total + rgba = numpy.zeros(image.shape + (4,), dtype=numpy.uint8) + rgba[..., :3] = numpy.clip(numpy.round(rgb * 255), 0, 255) + rgba[..., 3] = numpy.clip(numpy.round(alpha * 255), 0, 255) + return _encode_png(rgba) + + +def _encode_png(rgba) -> bytes: + height, width, _ = rgba.shape + raw = b"".join(b"\x00" + rgba[r].tobytes() for r in range(height)) + + def chunk(kind: bytes, payload: bytes) -> bytes: + body = kind + payload + return ( + struct.pack(">I", len(payload)) + body + struct.pack(">I", zlib.crc32(body)) + ) + + header = struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0) + return ( + b"\x89PNG\r\n\x1a\n" + + chunk(b"IHDR", header) + + chunk(b"IDAT", zlib.compress(raw, 6)) + + chunk(b"IEND", b"") + ) diff --git a/src/dls_tiled/visr.py b/src/dls_tiled/visr.py index a870f69..06817e7 100644 --- a/src/dls_tiled/visr.py +++ b/src/dls_tiled/visr.py @@ -3,18 +3,38 @@ Enable it in the tiled config with ``routers: [dls_tiled.visr:visr_router]``. """ +import asyncio import enum +import functools +import hashlib import inspect +import json import logging import os +import re +import time +from collections.abc import Callable +from typing import Any, NamedTuple import anyio.to_thread import numpy -from fastapi import APIRouter, Depends, HTTPException, Query, Request, Security +from fastapi import ( + APIRouter, + Depends, + HTTPException, + Query, + Request, + Response, + Security, +) from h5py._hl.dataset import Dataset as H5Dataset from scanspec.core import stack2dimension from scanspec.specs import Spec -from starlette.status import HTTP_422_UNPROCESSABLE_CONTENT +from starlette.status import ( + HTTP_422_UNPROCESSABLE_CONTENT, + HTTP_503_SERVICE_UNAVAILABLE, +) +from tiled.adapters.utils import DataNotReadyError from tiled.server.authentication import ( # type: ignore check_scopes, get_current_access_tags, @@ -27,12 +47,26 @@ from tiled.server.schemas import Principal from tiled.type_aliases import AccessTags, Scopes +from dls_tiled import rendering + logger = logging.getLogger(__name__) -CROSS_CHANNEL_RETRIES = int(os.getenv("DLS_TILED_VISR_CROSS_CHANNEL_RETRIES", "5")) + +def _env_count(name: str, default: int) -> int: + """A non-negative count from the environment (a negative one counts as 0).""" + return max(0, int(os.getenv(name, str(default)))) + + +CROSS_CHANNEL_RETRIES = _env_count("DLS_TILED_VISR_CROSS_CHANNEL_RETRIES", 5) CROSS_CHANNEL_RETRY_DELAY = float( os.getenv("DLS_TILED_VISR_CROSS_CHANNEL_RETRY_DELAY", "0.1") ) +# Tiled raises DataNotReadyError when the file on disk has fewer frames than the +# catalog advertises, i.e. a live write hasn't landed yet. It is retryable. +NOT_READY_RETRIES = _env_count("DLS_TILED_VISR_NOT_READY_RETRIES", 3) +NOT_READY_RETRY_DELAY = float(os.getenv("DLS_TILED_VISR_NOT_READY_RETRY_DELAY", "0.2")) +# The Retry-After Tiled's own array routes send for the same condition. +RETRY_AFTER_SECONDS = int(os.getenv("TILED_ARRAY_RETRY_AFTER", "1")) class ScanType(enum.Enum): @@ -43,6 +77,45 @@ class ScanType(enum.Enum): visr_router = APIRouter() +BINNED_CACHE_SECONDS = float(os.getenv("DLS_TILED_VISR_CACHE_SECONDS", "0.5")) +_recent: dict[tuple, tuple[float, tuple]] = {} +_inflight: dict[tuple, asyncio.Task] = {} + + +async def _shared(key: tuple, compute): + """Await ``compute()`` once for concurrent callers with the same key, and reuse + its result for ``BINNED_CACHE_SECONDS``. Errors are shared with the callers + waiting on that call but never kept. A cache time of 0 turns this off. + + The call runs as its own task, so a caller that goes away (a client + disconnecting cancels its request) doesn't cancel it for the others. + """ + if BINNED_CACHE_SECONDS <= 0: + return await compute() + hit = _recent.get(key) + if hit is not None and time.monotonic() - hit[0] < BINNED_CACHE_SECONDS: + return hit[1] + task = _inflight.get(key) + if task is None: + task = asyncio.ensure_future(_compute_and_keep(key, compute)) + # Mark any error as retrieved, so a call nobody waited on doesn't warn. + task.add_done_callback(lambda t: t.cancelled() or t.exception()) + _inflight[key] = task + return await asyncio.shield(task) + + +async def _compute_and_keep(key: tuple, compute): + try: + result = await compute() + finally: + _inflight.pop(key, None) + now = time.monotonic() + for stale in [k for k, (t, _) in _recent.items() if now - t > BINNED_CACHE_SECONDS]: + del _recent[stale] + _recent[key] = (now, result) + return result + + async def get_data(root, segments) -> H5Dataset | numpy.ndarray | dict: try: adapter = await root.lookup_adapter(segments) @@ -58,15 +131,30 @@ async def get_data(root, segments) -> H5Dataset | numpy.ndarray | dict: except Exception: raise else: - # Leaf node — read it - try: - if inspect.iscoroutinefunction(adapter.read): - data = await adapter.read() - else: - data = await anyio.to_thread.run_sync(adapter.read) - return data - except Exception: - raise + # Leaf node - read it. A file that is still behind the catalog is retried + # briefly, then reported as 503 with Retry-After (as Tiled's own routes + # do) rather than as an error the client can't tell from a bad request. + for attempt in range(NOT_READY_RETRIES + 1): + try: + if inspect.iscoroutinefunction(adapter.read): + return await adapter.read() + return await anyio.to_thread.run_sync(adapter.read) + except DataNotReadyError as err: + if attempt == NOT_READY_RETRIES: + raise HTTPException( + status_code=HTTP_503_SERVICE_UNAVAILABLE, + detail=str(err), + headers={"Retry-After": str(RETRY_AFTER_SECONDS)}, + ) from None + logger.info( + "Data for %s not ready (attempt %d/%d): %s", + "/".join(segments), + attempt + 1, + NOT_READY_RETRIES + 1, + err, + ) + await anyio.sleep(NOT_READY_RETRY_DELAY) + raise AssertionError("unreachable") # pragma: no cover async def fill_data(root, segments, shape=None, fill_value=numpy.nan): @@ -83,19 +171,279 @@ def _channel_length(array: H5Dataset | numpy.ndarray | dict) -> int: return array.shape[-1] +_AXIS_LETTERS = ("x", "y", "z") + + +def _axis_letter(name: object) -> str | None: + """Return "x", "y" or "z" if the axis name ends in one (e.g. "sample_stage-x").""" + if not isinstance(name, str): + return None + letter = re.split(r"[-_.]", name.lower())[-1] + return letter if letter in _AXIS_LETTERS else None + + +# The ``name=`` argument of a repr, not one that merely ends in it (``device_name=``) +_REPR_NAME_RE = re.compile(r"""(?:\(|,\s*)name=["']([^"']+)["']""") + + +def _axis_name(axis: object) -> str | None: + """The device name of a spec axis, or None if the spec doesn't carry one. + + A spec axis is whatever the plan passed in: a plain name, an ophyd-async + repr that carries the name (``Motor(name="sample_stage-x")``), or, in older + runs, an object repr with no name (``<... Motor object at 0x7f...>``). + """ + if not isinstance(axis, str): + return None + named = _REPR_NAME_RE.search(axis) + if named: + return named.group(1) + return None if axis.startswith("<") else axis + + +def _axis_letters(axes: list, motors: object) -> list[str | None]: + """The x/y/z letter of each spec axis, or None where it can't be told. + + The names come from the spec's axes. Where the spec has none (older runs + store an object repr) the run's ``motors`` list, which is in the same order + as the spec's axes, supplies them. + """ + names = [_axis_name(a) for a in axes] + if None in names and isinstance(motors, list) and len(motors) == len(axes): + names = list(motors) + return [_axis_letter(n) for n in names] + + +def _xyz_order(axes: list, motors: object) -> list[int]: + """Indices that put a spec's axes in x, y, z order. + + A spec lists its axes outermost dimension first, which is not necessarily x + first (a raster scan has y as the outer axis). If the axes cannot be told + apart as x, y and z, keep the spec's own order. + """ + letters = _axis_letters(axes, motors) + if None in letters or len(set(letters)) != len(letters): + return list(range(len(axes))) + return sorted(range(len(axes)), key=lambda i: _AXIS_LETTERS.index(letters[i])) + + +def _bin_edges(centres: numpy.ndarray) -> numpy.ndarray: + """Edges halfway between consecutive centres, extended half a step at each end.""" + if len(centres) == 1: + return numpy.array([centres[0] - 0.5, centres[0] + 0.5]) + mids = (centres[:-1] + centres[1:]) / 2 + first = centres[0] - (centres[1] - centres[0]) / 2 + last = centres[-1] + (centres[-1] - centres[-2]) / 2 + return numpy.concatenate([[first], mids, [last]]) + + +# Past this many cells the native-resolution image is too much to send on every +# poll (200 x 200 is about 2.3 MB of JSON), so the caller falls back to the +# coarse default binning. +MAX_GRID_CELLS = int(os.getenv("DLS_TILED_VISR_MAX_GRID_CELLS", "40000")) +# The ViSR plot needs at least three points on each axis, so a scan narrower than +# that keeps the previous default binning. +MIN_CELLS_PER_AXIS = 3 + + +class ScanGrid(NamedTuple): + """A scan's setpoint grid: bin edges on x and y, the point count, each + point's x and y setpoint in scan order, and whether the spec itself named + its axes (rather than the run's ``motors`` list, whose order isn't reliable: + see ``place_on_grid``).""" + + x_edges: numpy.ndarray + y_edges: numpy.ndarray + n_total: int + x_setpoints: numpy.ndarray + y_setpoints: numpy.ndarray + named: bool = True + + def transposed(self) -> "ScanGrid": + """The same grid with x and y swapped.""" + return ScanGrid( + self.y_edges, + self.x_edges, + self.n_total, + self.y_setpoints, + self.x_setpoints, + self.named, + ) + + +# A grid keeps every point's x and y setpoint, so a scan of a million points +# holds 16 MB; a few runs are enough, as the UI shows one or two at a time. +@functools.lru_cache(maxsize=8) +def _grid_from_spec(spec_json: str, motors_json: str) -> ScanGrid | None: + """Bin edges for x and y and the point count of the scan the spec describes. + + The grid is the set of distinct setpoint positions on each axis, so every + setpoint falls in its own cell and the cells' edges lie halfway between + neighbours. Returns None when the spec has no x and y axes that can be + identified, or doesn't lay points out on a grid (a spiral, say), where the + caller falls back to binning the positions as they come. Cached: the spec + of a run never changes, and this is asked for on every poll. + """ + spec = Spec.deserialize(json.loads(spec_json)) + by_axis = stack2dimension(spec.calculate()).midpoints + axes = list(by_axis) + letters = _axis_letters(axes, json.loads(motors_json)) + if "x" not in letters or "y" not in letters or len(set(letters)) != len(letters): + return None + n_total = len(by_axis[axes[0]]) + x_setpoints = numpy.asarray(by_axis[axes[letters.index("x")]], dtype=float) + y_setpoints = numpy.asarray(by_axis[axes[letters.index("y")]], dtype=float) + xs = numpy.unique(numpy.round(x_setpoints, 9)) + ys = numpy.unique(numpy.round(y_setpoints, 9)) + if min(len(xs), len(ys)) < MIN_CELLS_PER_AXIS: + return None # a line, or too few cells for the plot's axes + if len(xs) * len(ys) > min(n_total, MAX_GRID_CELLS): + return None # more cells than points: not a grid + named = None not in (_axis_name(a) for a in axes) + return ScanGrid( + _bin_edges(xs), _bin_edges(ys), n_total, x_setpoints, y_setpoints, named + ) + + +# The grid is used when at least this share of the recorded points lie within +# half a cell of their own setpoint (after any offset is taken out). +MIN_POINTS_ON_SETPOINT = 0.5 +# Offsets smaller than this share of a cell are taken as jitter and not applied. +MIN_OFFSET_CELLS = 0.1 +# Recorded points are paired with setpoints by index, and also one setpoint +# earlier and later, in case the capture gained or dropped a frame at the start. +SETPOINT_LAGS = (0, 1, -1) + + +def _first_row_length(grid: ScanGrid) -> int: + """Points in the scan's first row: up to the first move of its slower axis.""" + + def first_move(setpoints): + moved = numpy.flatnonzero(numpy.abs(setpoints - setpoints[0]) > 1e-9) + return int(moved[0]) if moved.size else len(setpoints) + + return max(first_move(grid.x_setpoints), first_move(grid.y_setpoints)) + + +def place_on_grid(x, y, grid: ScanGrid, uid=""): + """Bring recorded positions onto the spec's grid, or None if they don't fit it. + + Each recorded point is paired with its own setpoint, and the offset on each + axis is the median of their differences: positions captured in another + frame (fly scans currently record positions offset from their setpoints) + are shifted by it, offsets below ``MIN_OFFSET_CELLS`` of a cell are left + alone. The positions fit when at least ``MIN_POINTS_ON_SETPOINT`` of them + then lie within half a cell of their setpoint. + + Points are paired by index and also one place earlier and later, keeping + the pairing the most points fit: a frame gained or dropped at the start + would otherwise give an offset one cell out that fits within each row. Only + the points where a row wraps tell those apart, so a non-zero offset is + accepted only once the scan is past its first row; before that, positions + that need one don't fit yet. + + When the spec doesn't name its axes, which is x and which is y came from + the run's ``motors`` list, and older runs don't list it in the spec's axis + order (some list x first for a spec whose outer axis is y). Then the grid + with x and y swapped is tried as well, and the recorded points decide. + + Returns the positions, the offset ``(dx, dy)`` applied (or None) and the + grid they were placed on. + """ + x, y = numpy.asarray(x, dtype=float), numpy.asarray(y, dtype=float) + if x.size == 0: + # nothing received yet: the grid is all there is to show + return x, y, None, grid + best = None # (fit, grid), the spec's own orientation first on a tie + for candidate in [grid] if grid.named else [grid, grid.transposed()]: + fit = _fit(x, y, candidate) + if fit is not None and (best is None or fit[:2] > best[0][:2]): + best = (fit, candidate) + if best is None or best[0][0] < MIN_POINTS_ON_SETPOINT: + logger.info( + "Recorded positions for '%s' don't lie on the spec's grid " + "(%.0f%% on their setpoint at best); keeping the default binning", + uid, + 100 * (best[0][0] if best else 0.0), + ) + return None + (_, _, dx, dy), chosen = best + if chosen is not grid: + logger.info( + "Recorded positions of '%s' follow its spec with x and y the other " + "way round from its motors list", + uid, + ) + if dx == 0.0 and dy == 0.0: + return x, y, None, chosen + logger.info( + "Shifted the recorded positions of '%s' by (%s, %s) to match its spec", + uid, + dx, + dy, + ) + return x - dx, y - dy, (dx, dy), chosen + + +def _fit(x, y, grid: ScanGrid): + """How well positions fit a grid: ``(share on setpoint, by index, dx, dy)`` + for the best pairing of points with setpoints, or None if none works.""" + cell_x = float(numpy.min(numpy.diff(grid.x_edges))) + cell_y = float(numpy.min(numpy.diff(grid.y_edges))) + offsets_known = len(x) > _first_row_length(grid) + best = None + for lag in SETPOINT_LAGS: + # recorded point i is setpoint i + lag + recorded = numpy.arange(max(0, -lag), min(len(x), grid.n_total - lag)) + if recorded.size == 0: + continue + diff_x = x[recorded] - grid.x_setpoints[recorded + lag] + diff_y = y[recorded] - grid.y_setpoints[recorded + lag] + dx, dy = 0.0, 0.0 + if offsets_known: + dx, dy = float(numpy.nanmedian(diff_x)), float(numpy.nanmedian(diff_y)) + if not (numpy.isfinite(dx) and numpy.isfinite(dy)): + continue + dx = 0.0 if abs(dx) < MIN_OFFSET_CELLS * cell_x else dx + dy = 0.0 if abs(dy) < MIN_OFFSET_CELLS * cell_y else dy + on_setpoint = (numpy.abs(diff_x - dx) <= cell_x / 2) & ( + numpy.abs(diff_y - dy) <= cell_y / 2 + ) + candidate = (float(on_setpoint.mean()), lag == 0, dx, dy) + if best is None or candidate[:2] > best[:2]: + best = candidate + return best + + +async def get_scan_grid(root, uid): + """The run's setpoint grid ``(x_edges, y_edges, n_total)``, or None.""" + try: + adapter = await root.lookup_adapter([uid]) + start = adapter.metadata()["start"] + return _grid_from_spec( + json.dumps(start["spec"]), json.dumps(start.get("motors")) + ) + except Exception as e: + logger.debug("No usable scan grid for '%s': %s", uid, e) + return None + + async def get_setpoints(root, uid): """Return setpoints from the bluesky start document stored in a node's metadata.""" adapter = await root.lookup_adapter([uid]) metadata = adapter.metadata() try: - spec = Spec.deserialize(metadata["start"]["spec"]) + start = metadata["start"] + spec = Spec.deserialize(start["spec"]) except KeyError as e: raise HTTPException( status_code=HTTP_422_UNPROCESSABLE_CONTENT, detail=f"Could not find 'start.spec' in metadata for '{uid}': {e}", ) from None - midpoints = list(stack2dimension(spec.calculate()).midpoints.values()) + by_axis = stack2dimension(spec.calculate()).midpoints + axes = list(by_axis) + midpoints = [by_axis[axes[i]] for i in _xyz_order(axes, start.get("motors"))] x = midpoints[0] y = midpoints[1] if len(midpoints) > 1 else numpy.full(x.shape, numpy.nan) z = midpoints[2] if len(midpoints) > 2 else numpy.full(x.shape, numpy.nan) @@ -149,6 +497,8 @@ async def get_readbacks(root, uid, readback_x): """ try: readback_x, scan_type = await _fetch_readback_x(root, uid) + except HTTPException: + raise except Exception as e: raise HTTPException( status_code=HTTP_422_UNPROCESSABLE_CONTENT, @@ -251,6 +601,18 @@ async def binned( # type: ignore ): """Fetch a folded representation of an array dataset. + Returns ``RedTotal``/``GreenTotal``/``BlueTotal`` as matrices with one row per + y bin and one column per x bin (the layout image plots expect), ``null`` for + bins no point fell into, ``x_limits``/``y_limits`` as the bin edges (one more + than the columns/rows), ``n_points``, the number of scan points used, and + ``n_total``, the number the scan will have (``null`` if the run has no spec), + and ``position_offset``: ``{"x": dx, "y": dy}`` when the recorded positions + were shifted onto the spec's grid (see ``place_on_grid``), else ``null``. + + Unless ``width``/``height`` or a full range are given, bins follow the scan's + setpoint grid when its spec describes one: one cell per setpoint, a range + fixed from the first poll. + Args: x_dim_index: Index into the position array to use as the x axis (default 0). y_dim_index: Index into the position array to use as the y axis (default 1). @@ -274,7 +636,81 @@ async def binned( # type: ignore segments = [s for s in path.strip("/").split("/") if s] uid = segments[0] - # load data + # A UI polls this several times a second per client while points arrive only + # every couple of seconds, so identical requests share one computation (and + # one read of the files) and its result is reused for a moment. + # The key leaves out who is asking. That is safe only because the data is + # read through root_tree.lookup_adapter, which bypasses Tiled's access policy + # anyway; if per-run access control is added, the key needs the principal + # (or its access tags), or one user's result could be served to another. + key = (uid, tuple(sorted(request.query_params.multi_items()))) + + async def compute(): + output = await _binned_output( + root, + uid, + x_dim_index, + y_dim_index, + xmin, + xmax, + ymin, + ymax, + width, + height, + setpoints, + slice_dim, + ) + payload = await anyio.to_thread.run_sync(_encode_json, output) + return payload, f'"{hashlib.sha1(payload).hexdigest()[:20]}"' + + payload, etag = await _shared(key, compute) + return _conditional(request, payload, "application/json", etag) + + +def _encode_json(output: dict) -> bytes: + """Compact JSON, refusing NaN and infinities (which JSON can't carry).""" + return json.dumps(output, separators=(",", ":"), allow_nan=False).encode() + + +def _conditional(request: Request, payload: bytes, media: str, tag: str) -> Response: + """The response, or 304 if the client already has it (``If-None-Match``). + + The ETag is weak: Tiled's compression middleware may gzip the body and + keeps the header, so the tag can only promise the same content, not the + same bytes. Matching compares the quoted tag alone, so a client or proxy + sending it back in either form gets the 304. + """ + headers = {"ETag": f"W/{tag}", "Cache-Control": "no-cache"} + if tag in request.headers.get("if-none-match", ""): + return Response(status_code=304, headers=headers) + return Response(payload, media_type=media, headers=headers) + + +class ScanPoints(NamedTuple): + """The points of a run ready to draw: each channel's values, their positions, + how many there are, whether they were placed on the scan's grid, and the + offset taken out of the positions to do so (or None).""" + + data: dict + x: numpy.ndarray + y: numpy.ndarray + n_points: int + on_grid: bool + offset: tuple[float, float] | None + grid: "ScanGrid | None" + + +async def load_points( + root, uid, x_dim_index, y_dim_index, setpoints, slice_dim, grid=None +): + """Read a run's channels and positions, reconciled to one length and sliced. + + Positions are the recorded readbacks, or the spec's setpoints with + ``setpoints``. With a ``grid`` (for x and y as dimensions 0 and 1), the + positions are placed on it (see ``place_on_grid``) before ``slice_dim`` is + applied, since that pairs each point with its setpoint by its place in the + scan. ``n_points`` counts the points after slicing. + """ try: red_total = await get_data(root, [uid, "primary", "RedTotal"]) green_total = await get_data(root, [uid, "primary", "GreenTotal"]) @@ -291,6 +727,8 @@ async def binned( # type: ignore assert isinstance(blue_total, H5Dataset) or isinstance( blue_total, numpy.ndarray ) + except HTTPException: + raise except Exception as e: raise HTTPException( status_code=HTTP_422_UNPROCESSABLE_CONTENT, @@ -317,12 +755,20 @@ async def binned( # type: ignore # and retry re-fetching everything until they do. If they don't catch up # within the retry budget, fall back to truncating to whatever the shortest # currently is, so we always return a consistent (if possibly stale) image. - lengths = { - "RedTotal": _channel_length(red_total), - "GreenTotal": _channel_length(green_total), - "BlueTotal": _channel_length(blue_total), - "readbacks": readbacks.shape[-1], - } + # Setpoints come from the spec and are complete from the start, so they take + # no part in this: waiting for the data to reach them would only ever time + # out while a scan is running. + def current_lengths() -> dict[str, int]: + found = { + "RedTotal": _channel_length(red_total), + "GreenTotal": _channel_length(green_total), + "BlueTotal": _channel_length(blue_total), + } + if not setpoints: + found["readbacks"] = readbacks.shape[-1] + return found + + lengths = current_lengths() target_len = max(lengths.values()) for attempt in range(CROSS_CHANNEL_RETRIES): if min(lengths.values()) >= target_len: @@ -342,7 +788,7 @@ async def binned( # type: ignore new_green_total = await get_data(root, [uid, "primary", "GreenTotal"]) new_blue_total = await get_data(root, [uid, "primary", "BlueTotal"]) if setpoints: - new_readbacks = await get_setpoints(root, uid) + new_readbacks = readbacks else: new_readbacks, _ = await get_readbacks(root, uid, None) except Exception as e: @@ -361,12 +807,7 @@ async def binned( # type: ignore new_blue_total, new_readbacks, ) - lengths = { - "RedTotal": _channel_length(red_total), - "GreenTotal": _channel_length(green_total), - "BlueTotal": _channel_length(blue_total), - "readbacks": readbacks.shape[-1], - } + lengths = current_lengths() else: logger.info( "Giving up on cross-channel reconciliation for '%s' after %d attempt(s), " @@ -376,8 +817,10 @@ async def binned( # type: ignore lengths, ) - common_len = min(target_len, min(lengths.values())) - if any(length != common_len for length in lengths.values()): + common_len = min(target_len, min(lengths.values()), readbacks.shape[-1]) + if any(length != common_len for length in lengths.values()) or ( + readbacks.shape[-1] != common_len + ): red_total = red_total[:common_len] green_total = green_total[:common_len] blue_total = blue_total[:common_len] @@ -389,8 +832,17 @@ async def binned( # type: ignore "BlueTotal": blue_total, } + readbacks = numpy.asarray(readbacks, dtype=float) + on_grid, offset = False, None + if grid is not None and (x_dim_index, y_dim_index) == (0, 1): + placed = place_on_grid(readbacks[0], readbacks[1], grid, uid) + if placed is not None: + readbacks = readbacks.copy() + readbacks[0], readbacks[1], offset, grid = placed + on_grid = True + # mask out the points that lie outside the slice - mask = numpy.ones(readbacks.size, dtype=bool) + mask = numpy.ones(readbacks.shape[-1], dtype=bool) # one entry per point if slice_dim is not None: for slice_spec in slice_dim: try: @@ -432,30 +884,349 @@ async def binned( # type: ignore slice_dim is not None, ) - x_positions = readbacks[x_dim_index, :] - y_positions = readbacks[y_dim_index, :] + return ScanPoints( + data, + readbacks[x_dim_index, :], + readbacks[y_dim_index, :], + int(readbacks.shape[-1]), + on_grid, + offset, + grid, + ) + +async def _binned_output( + root, + uid, + x_dim_index, + y_dim_index, + xmin, + xmax, + ymin, + ymax, + width, + height, + setpoints, + slice_dim, +) -> dict: + """Compute the binned images for one run (the body of the /binned endpoint).""" # bundle the kwargs - histogram2d_kwargs = {} - if all(opt is not None for opt in (width, height)): + histogram2d_kwargs: dict = {} + explicit_bins = all(opt is not None for opt in (width, height)) + explicit_range = all(opt is not None for opt in (xmin, xmax, ymin, ymax)) + if explicit_bins: histogram2d_kwargs["bins"] = (width, height) - if all(opt is not None for opt in (xmin, xmax, ymin, ymax)): + if explicit_range: histogram2d_kwargs["range"] = ((xmin, xmax), (ymin, ymax)) - binned_output = {} - for channel in ("RedTotal", "GreenTotal", "BlueTotal"): - binned_channel = compute_binned_image( - data[channel], x_positions, y_positions, **histogram2d_kwargs - ) - binned_output[channel] = binned_channel["img"].tolist() - binned_output["x_limits"] = binned_channel["x"].tolist() - binned_output["y_limits"] = binned_channel["y"].tolist() + # Unless the caller chose the binning, use the scan's own setpoint grid: one + # cell per setpoint and a range fixed by the spec from the first poll, rather + # than a fixed 10x10 over whatever extent the points received so far happen + # to cover (which rescales the whole image as the scan fills in). + grid = await get_scan_grid(root, uid) + n_total = grid.n_total if grid is not None else None + use_grid = grid is not None and not explicit_bins and not explicit_range + points = await load_points( + root, + uid, + x_dim_index, + y_dim_index, + setpoints, + slice_dim, + grid if use_grid else None, + ) + data, x_positions, y_positions = points.data, points.x, points.y + n_points, position_offset = points.n_points, points.offset + if points.on_grid: + assert points.grid is not None + histogram2d_kwargs["bins"] = [points.grid.x_edges, points.grid.y_edges] + + def bin_channels() -> dict: + """The CPU-bound part, run off the event loop.""" + binned_output: dict = {} + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + binned_channel = compute_binned_image( + data[channel], x_positions, y_positions, **histogram2d_kwargs + ) + binned_output[channel] = image_rows_are_y( + binned_channel["img"], binned_channel["counts"] + ) + binned_output["x_limits"] = binned_channel["x"].tolist() + binned_output["y_limits"] = binned_channel["y"].tolist() + return binned_output + + binned_output = await anyio.to_thread.run_sync(bin_channels) + binned_output["n_points"] = n_points + binned_output["n_total"] = n_total + binned_output["position_offset"] = ( + None + if position_offset is None + else {"x": position_offset[0], "y": position_offset[1]} + ) return binned_output +RENDER_KINDS = ("cells", "rectangles", "gaussian", "ovals") +MAX_RENDER_SIZE = 2048 + + +def _nominal_step(grid: ScanGrid) -> tuple[float, float]: + """The scan's step on x and y: the median width of its grid cells.""" + return float(numpy.median(numpy.diff(grid.x_edges))), float( + numpy.median(numpy.diff(grid.y_edges)) + ) + + +def _beam_from_metadata(start: dict) -> tuple[float, float] | None: + """Beam size (FWHM) from the start document's ``beam_size``, if it has one: + a single number, or ``{"x": ..., "y": ...}``.""" + beam = start.get("beam_size") + if isinstance(beam, (int, float)): + return float(beam), float(beam) + if isinstance(beam, dict) and {"x", "y"} <= set(beam): + return float(beam["x"]), float(beam["y"]) + return None + + +@visr_router.get("/render/{path:path}") +async def render( # type: ignore + path: str, + request: Request, + kind: str = "cells", + x0: float | None = None, + x1: float | None = None, + y0: float | None = None, + y1: float | None = None, + size: int = 256, + setpoints: bool = False, + cell_x: float | None = None, + cell_y: float | None = None, + beam_x: float | None = None, + beam_y: float | None = None, + format: str = "json", + channel: str = "RedTotal", + vmin: float | None = None, + vmax: float | None = None, + line_width: float = 1.5, + line_opacity: float = 0.35, + slice_dim: list[str] | None = Query(None), # noqa: B008 + principal: Principal | None = Depends(get_current_principal), # type: ignore # noqa: B008 + root_tree=Depends(get_root_tree), # type: ignore # noqa: B008 + session_state: dict = Depends(get_session_state), # type: ignore # noqa: B008 + authn_access_tags: AccessTags | None = Depends(get_current_access_tags), # type: ignore # noqa: B008 + authn_scopes: Scopes = Depends(get_current_scopes), # type: ignore # noqa: B008 + _=Security(check_scopes, scopes=["read:data"]), # noqa: B008 +): + """Draw a run's points into an image for a viewport (prototype). + + ``kind`` picks how each point is drawn: ``cells`` (the scan's grid cells), + ``rectangles`` (a nominal-size cell centred on each recorded position), + ``gaussian`` (a beam-sized Gaussian per point) or ``ovals`` (a beam-sized + oval per point with a faint white outline, ``line_width`` pixels at + ``line_opacity``; overlaps are split between the ovals rather than blended, + see ``rendering.oval_layout``). The beam size (FWHM) comes from + ``beam_x``/``beam_y``, else the start document's ``beam_size``, else the + scan's step. The viewport + is ``x0``..``x1`` by ``y0``..``y1`` in data coordinates (default: the scan's + extent) at ``size`` pixels along its longer side. + + ``format=json`` returns every channel as rows (row 0 at ``y0``, ``null`` + where nothing is drawn) plus a ``coverage`` layer (0 to 1). ``format=png`` + returns one ``channel`` as a tinted RGBA image, transparent where nothing is + drawn, scaled to ``vmin``/``vmax`` if given (else to its own range). For + ``ovals``, JSON adds an ``outline`` layer (opacity of the white lines) and + the PNG has the lines drawn in. + """ + if kind not in RENDER_KINDS: + raise HTTPException( + status_code=400, detail=f"kind must be one of {RENDER_KINDS}" + ) + if format not in ("json", "png"): + raise HTTPException(status_code=400, detail="format must be json or png") + if format == "png" and channel not in rendering.CHANNEL_COLOUR: + raise HTTPException( + status_code=400, + detail=f"channel must be one of {tuple(rendering.CHANNEL_COLOUR)}", + ) + if not 1 <= size <= MAX_RENDER_SIZE: + raise HTTPException( + status_code=400, detail=f"size must be 1..{MAX_RENDER_SIZE}" + ) + if not (0 < line_width <= 20 and 0 <= line_opacity <= 1): + raise HTTPException( + status_code=400, + detail="line_width must be in (0, 20], line_opacity in [0, 1]", + ) + root = request.app.state.root_tree + uid = [s for s in path.strip("/").split("/") if s][0] + # Shared like /binned's results, and with the same caveat about access control. + key = ("render", uid, tuple(sorted(request.query_params.multi_items()))) + + async def compute(): + points = await load_points( + root, uid, 0, 1, setpoints, slice_dim, await get_scan_grid(root, uid) + ) + # the grid the points were placed on, which may have x and y swapped + x, y, offset, grid = points.x, points.y, points.offset, points.grid + + step = _nominal_step(grid) if grid is not None else None + if grid is not None: + extent = ( + grid.x_edges[0], + grid.x_edges[-1], + grid.y_edges[0], + grid.y_edges[-1], + ) + elif len(x): + extent = (float(x.min()), float(x.max()), float(y.min()), float(y.max())) + else: + extent = (0.0, 1.0, 0.0, 1.0) + bounds = [ + v if v is not None else e + for v, e in zip((x0, x1, y0, y1), extent, strict=True) + ] + if bounds[0] == bounds[1] or bounds[2] == bounds[3]: + raise HTTPException(status_code=400, detail="the viewport has no area") + viewport = rendering.Viewport.fit( + bounds[0], bounds[1], bounds[2], bounds[3], size + ) + + details: dict = {} + beam = None + draw: Callable[[Any], tuple[Any, Any]] + outline_of: Callable[[], Any] = lambda: None # noqa: E731 + if kind in ("gaussian", "ovals"): + if beam_x is not None and beam_y is not None: + beam, source = (beam_x, beam_y), "request" + else: + adapter = await root.lookup_adapter([uid]) + beam = _beam_from_metadata(adapter.metadata().get("start", {})) + source = "metadata" + if beam is None and step is not None: + beam, source = step, "scan step" + if beam is None: + raise HTTPException( + status_code=400, + detail=f"{kind} needs beam_x/beam_y without a beam size or grid", + ) + details["beam"] = {"x": beam[0], "y": beam[1], "source": source} + + if kind == "cells": + if grid is not None: + edges = (grid.x_edges, grid.y_edges) + else: # no spec: the previous default binning over the data + _, ex, ey = ( + numpy.histogram2d(x, y) if len(x) else (None, [0, 1], [0, 1]) + ) + edges = (numpy.asarray(ex), numpy.asarray(ey)) + + draw = lambda v: rendering.render_cells( # noqa: E731 + v, x, y, edges[0], edges[1], viewport + ) + elif kind == "rectangles": + size_x = cell_x if cell_x is not None else (step[0] if step else None) + size_y = cell_y if cell_y is not None else (step[1] if step else None) + if size_x is None or size_y is None: + raise HTTPException( + status_code=400, + detail="rectangles need cell_x/cell_y if the run has no scan grid", + ) + details["cell"] = {"x": size_x, "y": size_y} + + draw = lambda v: rendering.render_rectangles( # noqa: E731 + v, x, y, size_x, size_y, viewport + ) + elif kind == "gaussian": + assert beam is not None + draw = lambda v: rendering.render_gaussian( # noqa: E731 + v, x, y, beam[0], beam[1], viewport + ) + else: + assert beam is not None + ovals = beam + layout = functools.cache( + lambda: rendering.oval_layout( + x, y, ovals[0], ovals[1], viewport, line_width, line_opacity + ) + ) + outline_of = lambda: layout().outline # noqa: E731 + details["outline_style"] = {"width": line_width, "opacity": line_opacity} + draw = lambda v: layout().paint(v) # noqa: E731 + + def produce() -> tuple[bytes, str]: + """Draw and encode: the CPU-bound part, run off the event loop.""" + outline = outline_of() + if format == "png": + image, coverage = draw(points.data[channel]) + png = rendering.to_png( + image, + coverage, + rendering.CHANNEL_COLOUR[channel], + vmin, + vmax, + outline, + ) + return png, "image/png" + else: + images = { + name: draw(points.data[name]) + for name in ("RedTotal", "GreenTotal", "BlueTotal") + } + coverage = next(iter(images.values()))[1] + output = { + "kind": kind, + "viewport": { + "x0": viewport.x0, + "x1": viewport.x1, + "y0": viewport.y0, + "y1": viewport.y1, + "width": viewport.width, + "height": viewport.height, + }, + "channels": { + name: rendering.to_rows(img) + for name, (img, _) in images.items() + }, + "coverage": numpy.round(coverage, 3).tolist(), + **( + {"outline": numpy.round(outline, 3).tolist()} + if outline is not None + else {} + ), + "n_points": points.n_points, + "n_total": grid.n_total if grid is not None else None, + "position_offset": None + if offset is None + else {"x": offset[0], "y": offset[1]}, + **details, + } + return _encode_json(output), "application/json" + + payload, media = await anyio.to_thread.run_sync(produce) + return payload, f'"{hashlib.sha1(payload).hexdigest()[:20]}"', media + + payload, etag, media = await _shared(key, compute) + return _conditional(request, payload, media, etag) + + +def image_rows_are_y(img, counts) -> list[list[float | None]]: + """The binned image as plot-ready rows: one row per y bin, one column per x + bin, and ``None`` (JSON ``null``) where no point landed or the points' mean + isn't a number (NaN or infinite values, which JSON can't carry). + + ``numpy.histogram2d`` indexes its result ``[x_bin, y_bin]``, so it is + transposed here. An empty bin is reported as ``None`` rather than ``0`` so a + plot can tell "no data yet" from a measured zero. + """ + shown = (counts > 0) & numpy.isfinite(img) + rows = img.T.astype(object) + rows[~shown.T] = None + return rows.tolist() + + def compute_binned_image(data, readback_x, readback_y, **kwargs): counts, edges_x, edges_y = numpy.histogram2d(readback_x, readback_y, **kwargs) weights, _, _ = numpy.histogram2d(readback_x, readback_y, weights=data, **kwargs) img = numpy.divide(weights, counts, out=numpy.zeros_like(weights), where=counts > 0) - return {"img": img, "x": edges_x, "y": edges_y} + return {"img": img, "counts": counts, "x": edges_x, "y": edges_y} diff --git a/tests/test_render.py b/tests/test_render.py new file mode 100644 index 0000000..c93d88d --- /dev/null +++ b/tests/test_render.py @@ -0,0 +1,320 @@ +"""Tests for the prototype /render endpoint and its renderers.""" + +import struct + +import numpy +import pytest +from scanspec.specs import Line +from tiled.catalog import in_memory +from tiled.client import Context, from_context +from tiled.server.app import build_app + +from dls_tiled import rendering, visr +from dls_tiled.rendering import Viewport +from dls_tiled.visr import visr_router + + +@pytest.fixture(autouse=True) +def no_shared_results(monkeypatch): + visr._recent.clear() + visr._inflight.clear() + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.0) + + +@pytest.fixture +def client(tmp_path): + app = build_app( + in_memory(writable_storage=str(tmp_path)), include_routers=[visr_router] + ) + with Context.from_app(app) as context: + yield from_context(context) + + +# --- renderers --- + + +def test_viewport_fit_keeps_pixels_square(): + viewport = Viewport.fit(0, 4, 0, 2, 8) + + assert (viewport.width, viewport.height) == (8, 4) + xs, ys = viewport.centres() + assert xs[0] == 0.25 and ys[-1] == 1.75 + + +def test_cells_fill_each_grid_cell_with_its_mean(): + edges = numpy.array([-0.5, 0.5, 1.5, 2.5]) + x = numpy.array([0.0, 1.0, 2.0, 0.0]) + y = numpy.array([0.0, 0.0, 0.0, 2.0]) + values = numpy.array([1.0, 2.0, 3.0, 4.0]) + + image, coverage = rendering.render_cells( + values, x, y, edges, edges, Viewport.fit(-0.5, 2.5, -0.5, 2.5, 3) + ) + + assert numpy.array_equal(image[0], [1.0, 2.0, 3.0]) # row 0 is y = 0 + assert numpy.isnan(image[1]).all() + assert image[2, 0] == 4.0 and coverage[2, 0] == 1.0 and coverage[1, 1] == 0.0 + + +def test_cells_zoomed_into_one_cell_show_that_cell_everywhere(): + edges = numpy.array([-0.5, 0.5, 1.5]) + image, _ = rendering.render_cells( + numpy.array([1.0, 2.0]), + numpy.array([0.0, 1.0]), + numpy.array([0.0, 0.0]), + edges, + edges, + Viewport.fit(0.6, 1.4, -0.4, 0.4, 4), + ) + + assert numpy.all(image == 2.0) + + +def test_rectangles_cover_their_cell_around_the_recorded_position(): + image, coverage = rendering.render_rectangles( + [5.0], [2.0], [2.0], 2.0, 2.0, Viewport.fit(0, 4, 0, 4, 4) + ) + + assert numpy.all(image[1:3, 1:3] == 5.0) and numpy.all(coverage[1:3, 1:3] == 1.0) + assert numpy.isnan(image[0]).all() and coverage[0].sum() == 0.0 + + +def test_rectangles_cover_part_of_a_pixel_they_only_partly_overlap(): + # x: [1.75, 2.75] covers a quarter of column 1 and three quarters of column 2; + # y: [2, 3] covers row 2 exactly + _, coverage = rendering.render_rectangles( + [1.0], [2.25], [2.5], 1.0, 1.0, Viewport.fit(0, 4, 0, 4, 4) + ) + + assert coverage[2, 1] == pytest.approx(0.25) + assert coverage[2, 2] == pytest.approx(0.75) + assert coverage[1].sum() == 0.0 + + +def test_rectangles_average_where_they_overlap_weighted_by_coverage(): + # x: [0.5, 2.5] (value 2) and [1.5, 3.5] (value 4) + image, _ = rendering.render_rectangles( + [2.0, 4.0], [1.5, 2.5], [2.0, 2.0], 2.0, 2.0, Viewport.fit(0, 4, 0, 4, 4) + ) + + assert image[1, 0] == pytest.approx(2.0) # only the first: half the pixel + assert image[1, 1] == pytest.approx((2.0 * 1 + 4.0 * 0.5) / 1.5) + assert image[1, 2] == pytest.approx((2.0 * 0.5 + 4.0 * 1) / 1.5) + assert image[1, 3] == pytest.approx(4.0) # only the second + + +def test_gaussian_leaves_the_gap_between_far_apart_beams_empty(): + image, coverage = rendering.render_gaussian( + [1.0, 1.0], [1.0, 9.0], [1.0, 1.0], 1.0, 1.0, Viewport.fit(0, 10, 0, 2, 50) + ) + + middle = 25 # x = 5, four beam widths from either point + assert numpy.isnan(image[5, middle]) and coverage[5, middle] < 0.02 + assert coverage[5, 4] > 0.9 # x = 0.9, at a beam centre + + +def test_gaussian_weights_overlapping_beams_by_distance(): + image, _ = rendering.render_gaussian( + [0.0, 10.0], [1.0, 2.0], [1.0, 1.0], 2.0, 2.0, Viewport.fit(0, 3, 0, 2, 30) + ) + + assert ( + image[10, 10] < 5.0 < image[10, 20] + ) # nearer the first point, then the second + + +def test_ovals_split_an_overlap_along_the_line_through_the_crossings(): + # circles of radius 1 at x = 1 and x = 2.5 cross at x = 1.75 + layout = rendering.oval_layout( + [1.0, 2.5], [1.0, 1.0], 2.0, 2.0, Viewport.fit(0, 4, 0, 2, 80) + ) + image, coverage = layout.paint([10.0, 20.0]) + + middle_row = image[20] + assert middle_row[34] == 10.0 and middle_row[35] == 20.0 # x = 1.725 | 1.775 + assert numpy.isnan(middle_row[75]) and coverage[20, 75] == 0.0 # x = 3.775 + + +def test_ovals_are_scaled_on_each_axis(): + layout = rendering.oval_layout([2.0], [2.0], 2.0, 1.0, Viewport.fit(0, 4, 0, 4, 40)) + + assert layout.owner[20, 11] == 0 # x = 1.15: inside, 1 wide each side + assert layout.owner[14, 20] == -1 # y = 1.45: outside, 0.5 high each side + + +def test_oval_outlines_are_faint_and_whole_even_inside_another_oval(): + layout = rendering.oval_layout( + [1.0, 2.5], [1.0, 1.0], 2.0, 2.0, Viewport.fit(0, 4, 0, 2, 200) + ) + + # where each circle's outline runs inside the other: the second's leftmost + # point (x = 1.5) and the first's rightmost (x = 2), both at y = 1 + assert layout.outline[49:52, 74:77].max() > 0.25 + assert layout.outline[49:52, 99:102].max() > 0.25 + assert layout.outline.max() == pytest.approx(0.35) # never more than faint + assert layout.outline[50, 50] == 0.0 # inside, away from any outline + + +def test_oval_outlines_can_be_turned_off(): + layout = rendering.oval_layout( + [1.0], [1.0], 2.0, 2.0, Viewport.fit(0, 2, 0, 2, 40), line_opacity=0 + ) + + assert layout.outline.max() == 0.0 + + +def test_png_is_a_png_of_the_viewport_size_with_transparent_gaps(): + image = numpy.array([[1.0, numpy.nan], [2.0, 3.0]]) + payload = rendering.to_png(image, numpy.ones_like(image), colour=1) + + assert payload.startswith(b"\x89PNG\r\n\x1a\n") + width, height = struct.unpack(">II", payload[16:24]) + assert (width, height) == (2, 2) + + +# --- the endpoint --- + + +def write_scan(client, uid, n_points=9, x_shift=0.0, y_shift=0.0, start_extra=None): + """The first `n_points` of a 3x3 raster whose spec is in the start document.""" + spec = Line("Y", 0, 2, 3) * Line("X", 0, 2, 3) + start = {"spec": spec.serialize(), **(start_extra or {})} + primary = client.create_container(uid, metadata={"start": start}).create_container( + "primary" + ) + primary.write_array( + (numpy.array([0.0, 1.0, 2.0] * 3) + x_shift)[:n_points], key="X" + ) + primary.write_array( + (numpy.repeat([0.0, 1.0, 2.0], 3) + y_shift)[:n_points], key="Y" + ) + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(numpy.arange(9.0)[:n_points], key=channel) + + +def get(client, uid, **params): + return client.context.http_client.get(f"/api/v1/render/{uid}", params=params) + + +def test_render_cells_over_the_scan_extent_by_default(client): + write_scan(client, "grid") + + body = get(client, "grid", size=3).json() + + assert body["kind"] == "cells" + assert body["viewport"] == { + "x0": -0.5, + "x1": 2.5, + "y0": -0.5, + "y1": 2.5, + "width": 3, + "height": 3, + } + assert body["channels"]["RedTotal"] == [ + [0.0, 1.0, 2.0], + [3.0, 4.0, 5.0], + [6.0, 7.0, 8.0], + ] + assert body["n_points"] == 9 and body["n_total"] == 9 + + +def test_render_zooms_into_a_viewport(client): + write_scan(client, "zoom") + + body = get(client, "zoom", x0=0.5, x1=2.5, y0=-0.5, y1=0.5, size=4).json() + + assert (body["viewport"]["width"], body["viewport"]["height"]) == (4, 2) + assert body["channels"]["RedTotal"][0] == [1.0, 1.0, 2.0, 2.0] + + +def test_render_rectangles_use_the_scan_step_and_aligned_fly_positions(client): + write_scan(client, "fly", x_shift=1.0, y_shift=10.4) + + body = get(client, "fly", kind="rectangles", size=3).json() + + assert body["cell"] == {"x": 1.0, "y": 1.0} + assert body["position_offset"] == pytest.approx({"x": 1.0, "y": 10.4}) + assert body["channels"]["RedTotal"] == [ + [0.0, 1.0, 2.0], + [3.0, 4.0, 5.0], + [6.0, 7.0, 8.0], + ] + + +def test_render_gaussian_takes_the_beam_from_the_request_metadata_or_step(client): + write_scan(client, "beam", start_extra={"beam_size": {"x": 0.5, "y": 0.25}}) + write_scan(client, "nobeam") + + asked = get(client, "beam", kind="gaussian", beam_x=0.2, beam_y=0.2).json() + stored = get(client, "beam", kind="gaussian").json() + step = get(client, "nobeam", kind="gaussian").json() + + assert asked["beam"] == {"x": 0.2, "y": 0.2, "source": "request"} + assert stored["beam"] == {"x": 0.5, "y": 0.25, "source": "metadata"} + assert step["beam"] == {"x": 1.0, "y": 1.0, "source": "scan step"} + + +def test_render_gaussian_shows_gaps_a_small_beam_leaves(client): + write_scan(client, "gaps") + + body = get(client, "gaps", kind="gaussian", beam_x=0.2, beam_y=0.2, size=60).json() + + coverage = numpy.array(body["coverage"]) + assert coverage.max() > 0.9 # on a point + assert ( + coverage[:, 30].min() < 0.02 or coverage[30, :].min() < 0.02 + ) # between points + assert None in body["channels"]["RedTotal"][30] + + +def test_render_png(client): + write_scan(client, "png") + + response = get( + client, "png", format="png", channel="GreenTotal", size=16, vmin=0, vmax=8 + ) + + assert response.headers["content-type"] == "image/png" + assert response.content.startswith(b"\x89PNG") + assert struct.unpack(">II", response.content[16:24]) == (16, 16) + + +def test_render_shows_the_scan_extent_from_its_first_points(client): + write_scan(client, "early", n_points=2) + + body = get(client, "early", size=3).json() + + assert body["viewport"]["x0"] == -0.5 and body["viewport"]["y1"] == 2.5 + assert body["channels"]["RedTotal"][0] == [0.0, 1.0, None] + assert body["n_points"] == 2 + + +def test_render_ovals_returns_an_outline_layer(client): + write_scan(client, "ovals") + + body = get(client, "ovals", kind="ovals", beam_x=1.3, beam_y=1.3, size=30).json() + + assert body["beam"]["source"] == "request" + outline = numpy.array(body["outline"]) + assert outline.shape == (30, 30) and outline.max() == pytest.approx(0.35) + assert body["channels"]["RedTotal"][15][15] == 4.0 # the centre point + png = get(client, "ovals", kind="ovals", format="png", size=30) + assert png.headers["content-type"] == "image/png" + + +@pytest.mark.parametrize( + "params", + [ + {"kind": "contour"}, + {"format": "svg"}, + {"format": "png", "channel": "GrayTotal"}, + {"size": 0}, + {"x0": 1.0, "x1": 1.0}, + {"kind": "ovals", "line_width": 0}, + {"kind": "ovals", "line_opacity": 1.5}, + ], +) +def test_render_rejects_bad_requests(client, params): + write_scan(client, "bad") + + assert get(client, "bad", **params).status_code == 400 diff --git a/tests/test_visr.py b/tests/test_visr.py index 725e14f..21d1ed8 100644 --- a/tests/test_visr.py +++ b/tests/test_visr.py @@ -1,10 +1,24 @@ +import asyncio +import inspect +import json +from types import SimpleNamespace + import numpy import pytest +from fastapi import HTTPException +from scanspec.specs import Line, Spiral +from tiled.adapters.utils import DataNotReadyError from tiled.catalog import in_memory from tiled.client import Context, from_context from tiled.server.app import build_app -from dls_tiled.visr import compute_binned_image, visr_router +from dls_tiled import visr +from dls_tiled.visr import ( + compute_binned_image, + get_data, + get_setpoints, + visr_router, +) @pytest.fixture @@ -35,10 +49,50 @@ def test_binned_step_scan(client): assert response.status_code == 200, response.text body = response.json() - # mean of the two points in each (x, y) cell - assert body["RedTotal"] == [[0.5, 4.5], [2.5, 6.5]] - assert body["GreenTotal"] == [[5.0, 45.0], [25.0, 65.0]] + # mean of the two points in each cell; one row per y bin, one column per x bin + assert body["RedTotal"] == [[0.5, 2.5], [4.5, 6.5]] + assert body["GreenTotal"] == [[5.0, 25.0], [45.0, 65.0]] assert body["x_limits"] == [0.0, 0.55, 1.1] + assert body["y_limits"] == [0.0, 0.55, 1.1] + assert body["n_points"] == 8 + + +def test_binned_rows_are_y_and_columns_are_x(client): + # x only takes two values and y four, so the image must be 4 rows by 2 columns + primary = client.create_container("raster").create_container("primary") + x = numpy.array([0.0, 1.0] * 4) + y = numpy.repeat([0.0, 1.0, 2.0, 3.0], 2) + primary.write_array(x, key="X") + primary.write_array(y, key="Y") + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(numpy.arange(8.0), key=channel) + + response = client.context.http_client.get( + "/api/v1/binned/raster", params={"width": 2, "height": 4} + ) + + body = response.json() + assert body["RedTotal"] == [[0.0, 1.0], [2.0, 3.0], [4.0, 5.0], [6.0, 7.0]] + assert len(body["x_limits"]) == 3 and len(body["y_limits"]) == 5 + + +def test_binned_reports_empty_bins_as_null_and_keeps_a_measured_zero(client): + primary = client.create_container("sparse").create_container("primary") + primary.write_array(numpy.array([0.0, 1.0]), key="X") + primary.write_array(numpy.array([0.0, 1.0]), key="Y") + for channel, values in ( + ("RedTotal", [0.0, 5.0]), + ("GreenTotal", [1.0, 2.0]), + ("BlueTotal", [3.0, 4.0]), + ): + primary.write_array(numpy.array(values), key=channel) + + response = client.context.http_client.get( + "/api/v1/binned/sparse", params={"width": 2, "height": 2} + ) + + # a measured 0.0 stays 0.0; the two cells no point landed in are null + assert response.json()["RedTotal"] == [[0.0, None], [None, 5.0]] def test_binned_missing_run_is_422(client): @@ -57,3 +111,950 @@ def test_compute_binned_image_empty_bins_are_zero(): ) assert result["img"].tolist() == [[2.0, 0.0], [0.0, 0.0]] + assert result["counts"].tolist() == [[1.0, 0.0], [0.0, 0.0]] + + +def test_image_rows_are_y_transposes_and_marks_empty_bins(): + img = numpy.array([[1.0, 2.0, 0.0], [3.0, 0.0, 0.0]]) # [x_bin, y_bin] + counts = numpy.array([[1, 1, 0], [1, 0, 0]]) + + assert visr.image_rows_are_y(img, counts) == [ + [1.0, 3.0], + [2.0, None], + [None, None], + ] + + +def strict_json(response): + """The response body parsed as JSON proper: NaN and Infinity are refused.""" + + def refuse(constant): + raise ValueError(f"{constant} is not valid JSON") + + return json.loads(response.content, parse_constant=refuse) + + +def test_binned_reports_a_bin_whose_mean_is_not_a_number_as_null(client): + spec = Line("Y", 0, 2, 3) * Line("X", 0, 2, 3) + primary = client.create_container( + "nan", metadata={"start": {"spec": spec.serialize()}} + ).create_container("primary") + primary.write_array(numpy.array([0.0, 1.0, 2.0] * 3), key="X") + primary.write_array(numpy.repeat([0.0, 1.0, 2.0], 3), key="Y") + values = numpy.arange(9.0) + values[4] = numpy.nan # e.g. a detector frame that failed + values[5] = numpy.inf + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(values, key=channel) + + body = strict_json(client.context.http_client.get("/api/v1/binned/nan")) + + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, None, None], [6.0, 7.0, 8.0]] + + +def _setpoints(spec, **start): + """Run get_setpoints against a fake root whose run has this start document.""" + metadata = {"start": {"spec": spec.serialize(), **start}} + adapter = SimpleNamespace(metadata=lambda: metadata) + + async def lookup_adapter(path): + return adapter + + return asyncio.run( + get_setpoints(SimpleNamespace(lookup_adapter=lookup_adapter), "uid") + ) + + +def test_setpoints_follow_axis_names_not_dimension_order(): + # Raster scan: y is the outer dimension, x the inner one. + spec = Line("sample_stage-y", 0, 1, 2) * Line("sample_stage-x", 0, 10, 3) + + x, y, z = _setpoints(spec) + + assert x.tolist() == [0, 5, 10, 0, 5, 10] + assert y.tolist() == [0, 0, 0, 1, 1, 1] + assert numpy.isnan(z).all() + + +def test_setpoints_of_runs_with_repr_axes_use_the_motors_list(): + spec = Line("", 0, 1, 2) * Line( + "", 0, 10, 3 + ) + + x, y, _ = _setpoints(spec, motors=["sample_stage-y", "sample_stage-x"]) + + assert x.tolist() == [0, 5, 10, 0, 5, 10] + assert y.tolist() == [0, 0, 0, 1, 1, 1] + + +def test_setpoints_keep_dimension_order_when_axes_are_not_x_y_z(): + spec = Line("theta", 0, 1, 2) * Line("energy", 0, 10, 3) + + x, y, _ = _setpoints(spec) + + assert x.tolist() == [0, 0, 0, 1, 1, 1] + assert y.tolist() == [0, 5, 10, 0, 5, 10] + + +class _LaggingAdapter: + """A leaf whose file is behind the catalog for the first `lag` reads.""" + + def __init__(self, lag): + self.lag = lag + self.reads = 0 + + def read(self): + self.reads += 1 + if self.reads <= self.lag: + raise DataNotReadyError("advertises (25,), only (24,) available") + return numpy.arange(25.0) + + +def _root_with(adapter): + async def lookup_adapter(segments): + return adapter + + return SimpleNamespace(lookup_adapter=lookup_adapter) + + +@pytest.fixture +def fast_retries(monkeypatch): + monkeypatch.setattr(visr, "NOT_READY_RETRY_DELAY", 0.0) + + +def test_get_data_retries_a_file_that_is_behind_the_catalog(fast_retries): + adapter = _LaggingAdapter(lag=2) + + data = asyncio.run(get_data(_root_with(adapter), ["uid", "primary", "RedTotal"])) + + assert isinstance(data, numpy.ndarray) + assert data.shape == (25,) + assert adapter.reads == 3 + + +@pytest.mark.parametrize(("value", "count"), [("2", 2), ("0", 0), ("-1", 0)]) +def test_retry_counts_from_the_environment_are_never_negative( + monkeypatch, value, count +): + monkeypatch.setenv("DLS_TILED_VISR_TEST_COUNT", value) + + assert visr._env_count("DLS_TILED_VISR_TEST_COUNT", 3) == count + + +def test_retry_counts_default_when_not_set(): + assert visr._env_count("DLS_TILED_VISR_NOT_SET_ANYWHERE", 3) == 3 + + +def test_get_data_answers_503_with_retry_after_when_the_file_stays_behind(fast_retries): + adapter = _LaggingAdapter(lag=100) + + with pytest.raises(HTTPException) as excinfo: + asyncio.run(get_data(_root_with(adapter), ["uid", "primary", "RedTotal"])) + + assert excinfo.value.status_code == 503 + assert excinfo.value.headers == {"Retry-After": "1"} + assert adapter.reads == visr.NOT_READY_RETRIES + 1 + + +def test_binned_keeps_a_503_instead_of_turning_it_into_422(client, monkeypatch): + write_step_scan(client) + + async def not_ready(root, segments): + raise HTTPException( + status_code=503, detail="not ready", headers={"Retry-After": "1"} + ) + + monkeypatch.setattr(visr, "get_data", not_ready) + + response = client.context.http_client.get("/api/v1/binned/scan") + + assert response.status_code == 503 + assert response.headers["retry-after"] == "1" + + +# --- the scan's own grid ------------------------------------------------------------ + + +def _grid_or_none(spec, motors=None): + return visr._grid_from_spec(json.dumps(spec.serialize()), json.dumps(motors)) + + +def _grid(spec, motors=None): + grid = _grid_or_none(spec, motors) + assert grid is not None + return grid + + +def test_grid_edges_sit_halfway_between_setpoints(): + spec = Line("sample_stage-y", 0, 2, 3) * ~Line("sample_stage-x", 0, 10, 3) + + x_edges, y_edges, n_total, *_ = _grid(spec) + + assert x_edges.tolist() == [-2.5, 2.5, 7.5, 12.5] + assert y_edges.tolist() == [-0.5, 0.5, 1.5, 2.5] + assert n_total == 9 + + +def test_grid_of_runs_with_repr_axes_uses_the_motors_list(): + spec = Line("", 0, 2, 3) * Line( + "", 0, 10, 3 + ) + + x_edges, y_edges, *_ = _grid(spec, ["sample_stage-y", "sample_stage-x"]) + + assert len(x_edges) == 4 and len(y_edges) == 4 + + +def test_grid_is_none_when_the_axes_are_not_x_and_y(): + assert _grid_or_none(Line("theta", 0, 1, 2) * Line("energy", 0, 10, 3)) is None + + +def test_grid_is_none_when_the_points_do_not_lie_on_a_grid(): + assert ( + _grid_or_none( + Spiral( + x_axis="x", x_centre=0, x_diameter=5, x_step=1, y_axis="y", y_centre=0 + ) + ) + is None + ) + + +def test_grid_is_none_when_it_has_more_cells_than_the_cap(monkeypatch): + spec = Line("y", 0, 2, 4) * Line("x", 0, 2, 4) + visr._grid_from_spec.cache_clear() + assert _grid_or_none(spec) is not None + + monkeypatch.setattr(visr, "MAX_GRID_CELLS", 10) + visr._grid_from_spec.cache_clear() + + assert _grid_or_none(spec) is None + visr._grid_from_spec.cache_clear() + + +def test_grid_is_none_for_a_scan_narrower_than_three_setpoints(): + # a line, and a 2-wide grid: too few points on an axis for the UI's plot axes + assert _grid_or_none(Line("y", 0, 2, 3) * Line("x", 3, 3, 1)) is None + assert _grid_or_none(Line("y", 0, 2, 3) * Line("x", 0, 1, 2)) is None + + +def write_grid_scan(client, uid, n_points): + """The first `n_points` of a 3 (x) by 3 (y) raster; its spec is in the start doc.""" + spec = Line("Y", 0, 2, 3) * Line("X", 0, 2, 3) + metadata = {"start": {"spec": spec.serialize(), "motors": ["Y", "X"]}} + primary = client.create_container(uid, metadata=metadata).create_container( + "primary" + ) + x = numpy.array([0.0, 1.0, 2.0] * 3)[:n_points] + y = numpy.repeat([0.0, 1.0, 2.0], 3)[:n_points] + primary.write_array(x, key="X") + primary.write_array(y, key="Y") + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(numpy.arange(9.0)[:n_points], key=channel) + + +def test_binned_uses_the_scans_grid_without_being_asked(client): + write_grid_scan(client, "grid", 9) + + body = client.context.http_client.get("/api/v1/binned/grid").json() + + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]] + assert body["x_limits"] == [-0.5, 0.5, 1.5, 2.5] + assert body["y_limits"] == [-0.5, 0.5, 1.5, 2.5] + assert (body["n_points"], body["n_total"]) == (9, 9) + + +def test_binned_grid_is_the_same_shape_from_the_first_points_of_a_scan(client): + write_grid_scan(client, "partial", 7) # the last row is only partly measured + + body = client.context.http_client.get("/api/v1/binned/partial").json() + + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, None, None]] + assert body["y_limits"] == [-0.5, 0.5, 1.5, 2.5] + assert (body["n_points"], body["n_total"]) == (7, 9) + + +def test_binned_explicit_width_and_height_override_the_grid(client): + write_grid_scan(client, "chosen", 9) + + body = client.context.http_client.get( + "/api/v1/binned/chosen", params={"width": 1, "height": 1} + ).json() + + assert body["RedTotal"] == [[4.0]] + + +def test_binned_without_a_spec_still_reports_no_total(client): + write_step_scan(client) + + body = client.context.http_client.get("/api/v1/binned/scan").json() + + assert body["n_total"] is None + + +# --- sharing work between polls ------------------------------------------------------ + + +@pytest.fixture(autouse=True) +def isolated_binned_cache(monkeypatch): + """Start every test with an empty cache, switched off unless the test opts in.""" + visr._recent.clear() + visr._inflight.clear() + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.0) + + +class _Counting: + """An async compute() that records how often it ran.""" + + def __init__(self, result=(b"{}", '"e"'), delay=0.05, error=None): + self.calls = 0 + self.result, self.delay, self.error = result, delay, error + + async def __call__(self): + self.calls += 1 + await asyncio.sleep(self.delay) + if self.error: + raise self.error + return self.result + + +def test_concurrent_identical_requests_share_one_computation(monkeypatch): + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.5) + compute = _Counting() + + async def scenario(): + return await asyncio.gather( + *(visr._shared(("uid", ()), compute) for _ in range(8)) + ) + + results = asyncio.run(scenario()) + + assert compute.calls == 1 + assert results == [compute.result] * 8 + + +def test_a_result_is_reused_for_a_moment_and_then_recomputed(monkeypatch): + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.1) + compute = _Counting(delay=0) + + async def scenario(): + await visr._shared(("uid", ()), compute) + await visr._shared(("uid", ()), compute) + assert compute.calls == 1 + await asyncio.sleep(0.12) + await visr._shared(("uid", ()), compute) + + asyncio.run(scenario()) + + assert compute.calls == 2 + + +def test_different_requests_do_not_share(monkeypatch): + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.5) + compute = _Counting(delay=0) + + async def scenario(): + await visr._shared(("a", ()), compute) + await visr._shared(("b", ()), compute) + await visr._shared(("a", (("width", "3"),)), compute) + + asyncio.run(scenario()) + + assert compute.calls == 3 + + +def test_errors_reach_everyone_waiting_but_are_not_kept(monkeypatch): + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.5) + failing = _Counting(error=HTTPException(status_code=503, detail="not ready")) + working = _Counting(delay=0) + + async def scenario(): + outcomes = await asyncio.gather( + *(visr._shared(("uid", ()), failing) for _ in range(3)), + return_exceptions=True, + ) + assert all( + isinstance(o, HTTPException) and o.status_code == 503 for o in outcomes + ) + return await visr._shared(("uid", ()), working) # the error wasn't cached + + assert asyncio.run(scenario()) == working.result + assert failing.calls == 1 and working.calls == 1 + + +def test_a_caller_going_away_does_not_cancel_the_call_for_the_others(monkeypatch): + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.5) + compute = _Counting(delay=0.05) + + async def scenario(): + first = asyncio.ensure_future(visr._shared(("uid", ()), compute)) + await asyncio.sleep(0) # the first caller has started the call + others = [ + asyncio.ensure_future(visr._shared(("uid", ()), compute)) for _ in range(3) + ] + await asyncio.sleep(0.01) + first.cancel() # e.g. its client disconnected + results = await asyncio.gather(*others) + assert first.cancelled() + return results + + assert asyncio.run(scenario()) == [compute.result] * 3 + assert compute.calls == 1 + + +def test_a_cache_time_of_zero_turns_sharing_off(): + compute = _Counting(delay=0) + + async def scenario(): + await visr._shared(("uid", ()), compute) + await visr._shared(("uid", ()), compute) + + asyncio.run(scenario()) + + assert compute.calls == 2 + + +def test_binned_has_an_etag_and_answers_304_when_nothing_changed(client, monkeypatch): + monkeypatch.setattr(visr, "BINNED_CACHE_SECONDS", 0.5) + write_grid_scan(client, "grid", 9) + http = client.context.http_client + + first = http.get("/api/v1/binned/grid") + second = http.get( + "/api/v1/binned/grid", headers={"If-None-Match": first.headers["etag"]} + ) + + assert first.status_code == 200 and first.headers["cache-control"] == "no-cache" + assert first.headers["etag"].startswith('W/"') # the body may be gzipped + assert second.status_code == 304 and second.content == b"" + assert second.headers["etag"] == first.headers["etag"] + strong = http.get( + "/api/v1/binned/grid", + headers={"If-None-Match": first.headers["etag"].removeprefix("W/")}, + ) + assert strong.status_code == 304 + + +def test_binned_etag_changes_when_the_image_does(client): + write_grid_scan(client, "short", 7) + write_grid_scan(client, "full", 9) + http = client.context.http_client + + short = http.get("/api/v1/binned/short").headers["etag"] + full = http.get("/api/v1/binned/full").headers["etag"] + + assert short != full + stale = http.get("/api/v1/binned/full", headers={"If-None-Match": short}) + assert stale.status_code == 200 + + +# --- the axis forms real runs store --- + + +@pytest.mark.parametrize( + ("stored", "name"), + [ + ("sample_stage-x", "sample_stage-x"), + ('Motor(name="sample_stage-x")', "sample_stage-x"), + ("Motor(name='sample_stage-y')", "sample_stage-y"), + ('Motor(device_name="stage", name="sample_stage-x")', "sample_stage-x"), + ('Motor(prefix="BL01:", name="sample_stage-x")', "sample_stage-x"), + ("", None), + ], +) +def test_axis_name_reads_the_name_out_of_the_forms_runs_store(stored, name): + assert visr._axis_name(stored) == name + + +def test_setpoints_of_runs_with_named_reprs_need_no_motors_list(): + # recent runs store Motor(name=...) and fly scans have no motors list + spec = Line('Motor(name="sample_stage-y")', 0, 1, 2) * Line( + 'Motor(name="sample_stage-x")', 0, 10, 3 + ) + + x, y, _ = _setpoints(spec) + + assert x.tolist() == [0, 5, 10, 0, 5, 10] + assert y.tolist() == [0, 0, 0, 1, 1, 1] + + +def test_grid_of_runs_with_named_reprs_needs_no_motors_list(): + spec = Line('Motor(name="sample_stage-y")', 0, 2, 3) * ~Line( + 'Motor(name="sample_stage-x")', 0, 10, 3 + ) + + x_edges, y_edges, n_total, *_ = _grid(spec) + + assert x_edges.tolist() == [-2.5, 2.5, 7.5, 12.5] + assert y_edges.tolist() == [-0.5, 0.5, 1.5, 2.5] + assert n_total == 9 + + +def write_offset_scan(client, uid, x_shift, y_shift, x_scale=1.0, n_points=9): + """A 3x3 raster whose recorded positions are offset (and scaled) from the spec.""" + spec = Line("Y", 0, 2, 3) * Line("X", 0, 2, 3) + metadata = {"start": {"spec": spec.serialize()}} + primary = client.create_container(uid, metadata=metadata).create_container( + "primary" + ) + x = numpy.array([0.0, 1.0, 2.0] * 3) * x_scale + x_shift + y = numpy.repeat([0.0, 1.0, 2.0], 3) + y_shift + primary.write_array(x[:n_points], key="X") + primary.write_array(y[:n_points], key="Y") + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(numpy.arange(9.0)[:n_points], key=channel) + + +def test_binned_shifts_offset_positions_onto_the_scan_grid(client): + # fly scans currently record positions offset from their setpoints + write_offset_scan(client, "offset", x_shift=1.0, y_shift=10.4) + + body = client.context.http_client.get("/api/v1/binned/offset").json() + + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]] + assert body["x_limits"] == [-0.5, 0.5, 1.5, 2.5] # the spec's frame + assert body["position_offset"] == pytest.approx({"x": 1.0, "y": 10.4}) + + +def test_binned_keeps_the_default_binning_when_no_offset_explains_the_positions( + client, +): + write_offset_scan(client, "scaled", x_shift=1.0, y_shift=10.4, x_scale=3.0) + + body = client.context.http_client.get("/api/v1/binned/scaled").json() + + assert len(body["RedTotal"]) == 10 and len(body["RedTotal"][0]) == 10 + assert body["position_offset"] is None + assert body["n_total"] == 9 # still known from the spec + + +def test_positions_that_already_fit_are_not_shifted(client): + write_grid_scan(client, "fits", 9) + + body = client.context.http_client.get("/api/v1/binned/fits").json() + + assert body["position_offset"] is None + + +def test_binned_takes_an_offset_once_the_scan_is_past_its_first_row(client): + # until a row wraps, an offset can't be told from a frame gained or dropped + write_offset_scan(client, "row", x_shift=1.0, y_shift=10.4, n_points=3) + write_offset_scan(client, "wrapped", x_shift=1.0, y_shift=10.4, n_points=4) + http = client.context.http_client + + first_row = http.get("/api/v1/binned/row").json() + wrapped = http.get("/api/v1/binned/wrapped").json() + + assert first_row["position_offset"] is None + assert len(first_row["RedTotal"]) == 10 # the default binning, for now + assert wrapped["position_offset"] == pytest.approx({"x": 1.0, "y": 10.4}) + assert wrapped["RedTotal"] == [ + [0.0, 1.0, 2.0], + [3.0, None, None], + [None, None, None], + ] + + +def raster(n=20): + """A raster of n x n cells of width 1, x the fast axis.""" + return _grid(Line("y", 0, n - 1, n) * Line("x", 0, n - 1, n)) + + +def cells_hit(x, y, grid): + counts, _, _ = numpy.histogram2d(x, y, bins=[grid.x_edges, grid.y_edges]) + return counts + + +@pytest.mark.parametrize("shift", [0.5, 0.9]) +def test_place_on_grid_takes_out_an_offset_smaller_than_a_cell(shift): + # small enough that every position still lies inside the grid's edges + grid = raster() + + placed = visr.place_on_grid(grid.x_setpoints + shift, grid.y_setpoints, grid) + + assert placed is not None + assert placed[2] == pytest.approx((shift, 0.0)) + assert numpy.all(cells_hit(placed[0], placed[1], grid) == 1) + + +def test_place_on_grid_leaves_jitter_alone(): + grid = raster() + jitter = numpy.random.default_rng(0).uniform(-0.05, 0.05, grid.n_total) + + placed = visr.place_on_grid(grid.x_setpoints + jitter, grid.y_setpoints, grid) + + assert placed is not None and placed[2] is None + + +@pytest.mark.parametrize( + "frames", + ["dropped first", "extra run-up"], +) +def test_place_on_grid_finds_the_offset_when_a_frame_is_gained_or_lost(frames): + grid = raster() + x, y = grid.x_setpoints + 1.0, grid.y_setpoints + 10.4 + if frames == "dropped first": + x, y = x[1:], y[1:] + else: + x, y = numpy.r_[x[0] - 1, x], numpy.r_[y[0], y] + + placed = visr.place_on_grid(x, y, grid) + + assert placed is not None + assert placed[2] == pytest.approx((1.0, 10.4)) + assert cells_hit(placed[0], placed[1], grid).max() == 1 + + +def test_place_on_grid_finds_the_offset_despite_a_point_lost_late_in_the_scan(): + grid = raster(5) + x, y = grid.x_setpoints + 1.0, grid.y_setpoints + 10.4 + # one point lost near the end shifts the last few points by one place + x, y = numpy.delete(x, 21), numpy.delete(y, 21) + + placed = visr.place_on_grid(x, y, grid) + + assert placed is not None + assert placed[2] == pytest.approx((1.0, 10.4)) + + +def test_place_on_grid_without_points_is_the_bare_grid(): + grid = raster(5) + + placed = visr.place_on_grid(numpy.array([]), numpy.array([]), grid) + + assert placed is not None and placed[2] is None + + +@pytest.mark.parametrize( + "motors", + [["sample_stage-y", "sample_stage-x"], ["sample_stage-x", "sample_stage-y"]], +) +def test_binned_lets_the_points_decide_x_and_y_when_only_the_motors_list_names_them( + client, motors +): + # older runs store the axes as object reprs, and some list their motors in + # another order than the spec's axes: here y is the outer axis either way + spec = Line("", 0, 4, 3) * Line( + "", 0, 2, 3 + ) + write_points( + client, + "unnamed", + [0.0, 1.0, 2.0] * 3, + numpy.repeat([0.0, 2.0, 4.0], 3), + metadata={"start": {"spec": spec.serialize(), "motors": motors}}, + ) + + body = client.context.http_client.get("/api/v1/binned/unnamed").json() + + assert body["x_limits"] == [-0.5, 0.5, 1.5, 2.5] + assert body["y_limits"] == [-1.0, 1.0, 3.0, 5.0] + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]] + + +def test_place_on_grid_keeps_the_specs_own_names_over_a_better_fit(): + grid = raster(5) + + # positions that would fit the grid only with x and y swapped + placed = visr.place_on_grid(grid.y_setpoints, grid.x_setpoints, grid) + + assert placed is None + + +def test_binned_slice_dim_keeps_the_grid_of_an_offset_scan(client): + # z innermost: a slice in z leaves every other point, which must still be + # paired with its own setpoint to find the offset + spec = Line("y", 0, 2, 3) * Line("x", 0, 2, 3) * Line("z", 0, 1, 2) + grid = _grid(spec) + z_setpoints = numpy.tile([0.0, 1.0], 9) + write_points( + client, + "layers", + grid.x_setpoints + 1.0, + grid.y_setpoints + 10.4, + z_setpoints, + metadata={"start": {"spec": spec.serialize()}}, + ) + + body = client.context.http_client.get( + "/api/v1/binned/layers", params={"slice_dim": "2:1:0.25"} + ).json() + + assert body["position_offset"] == pytest.approx({"x": 1.0, "y": 10.4}) + assert body["RedTotal"] == [[1.0, 3.0, 5.0], [7.0, 9.0, 11.0], [13.0, 15.0, 17.0]] + assert body["n_points"] == 9 # the points in the slice + + +# --- the rest of the endpoint --- + + +@pytest.fixture +def app_and_client(tmp_path): + """Like `client`, with the app too, for tests that reach into its tree.""" + catalog = in_memory(writable_storage=str(tmp_path)) + app = build_app(catalog, include_routers=[visr_router]) + with Context.from_app(app) as context: + yield app, from_context(context) + + +@pytest.fixture +def fast_reconciliation(monkeypatch): + monkeypatch.setattr(visr, "CROSS_CHANNEL_RETRY_DELAY", 0.0) + + +def write_points( + client, uid, x, y, z=None, channels=None, where="primary", metadata=None +): + """Write positions and the three channels of a run, positions under `where`.""" + run = client.create_container(uid, metadata=metadata or {}) + primary = run.create_container("primary") + target = primary if where == "primary" else primary.create_container(where) + names = {"X": "X", "Y": "Y", "Z": "Z"} + if where != "primary" or channels == "named": + names = {"X": "sample_stage-x", "Y": "sample_stage-y", "Z": "sample_stage-z"} + target.write_array(numpy.asarray(x, dtype=float), key=names["X"]) + target.write_array(numpy.asarray(y, dtype=float), key=names["Y"]) + if z is not None: + target.write_array(numpy.asarray(z, dtype=float), key=names["Z"]) + values = numpy.arange(float(len(x))) + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(values, key=channel) + + +# setpoints=true over HTTP + + +def test_binned_setpoints_mode_bins_the_specs_positions(client): + # readbacks far off the spec: setpoints mode must not use them + spec = Line("Y", 0, 2, 3) * Line("X", 0, 2, 3) + write_points( + client, + "sp", + numpy.full(9, 50.0), + numpy.full(9, 50.0), + metadata={"start": {"spec": spec.serialize()}}, + ) + + body = client.context.http_client.get( + "/api/v1/binned/sp", params={"setpoints": True} + ).json() + + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, 7.0, 8.0]] + assert body["x_limits"] == [-0.5, 0.5, 1.5, 2.5] + assert body["position_offset"] is None + + +def test_binned_setpoints_mode_on_a_scan_still_running(client, monkeypatch): + spec = Line("Y", 0, 2, 3) * Line("X", 0, 2, 3) + write_points( + client, + "running", + numpy.zeros(7), + numpy.zeros(7), + metadata={"start": {"spec": spec.serialize()}}, + ) + calls = [] + real = visr.get_setpoints + + async def counting(root, uid): + calls.append(uid) + return await real(root, uid) + + monkeypatch.setattr(visr, "get_setpoints", counting) + + body = client.context.http_client.get( + "/api/v1/binned/running", params={"setpoints": True} + ).json() + + assert body["RedTotal"] == [[0.0, 1.0, 2.0], [3.0, 4.0, 5.0], [6.0, None, None]] + assert body["n_points"] == 7 + # the setpoints are complete from the start: nothing to wait for + assert len(calls) == 1 + + +def test_binned_setpoints_mode_without_a_spec_is_422(client): + write_step_scan(client) + + response = client.context.http_client.get( + "/api/v1/binned/scan", params={"setpoints": True} + ) + + assert response.status_code == 422 + assert "start.spec" in response.json()["detail"] + + +# slice_dim + + +def test_binned_slice_dim_keeps_only_points_in_the_slice(client): + # two layers in z; only the z = 0 layer is kept + x = [0.0, 1.0, 0.0, 1.0, 0.0, 1.0, 0.0, 1.0] + y = [0.0, 0.0, 1.0, 1.0, 0.0, 0.0, 1.0, 1.0] + z = [0.0, 0.0, 0.0, 0.0, 5.0, 5.0, 5.0, 5.0] + write_points(client, "layers", x, y, z) + + body = client.context.http_client.get( + "/api/v1/binned/layers", + params={"width": 2, "height": 2, "slice_dim": "2:0:0.5"}, + ).json() + + assert body["RedTotal"] == [[0.0, 1.0], [2.0, 3.0]] + + +@pytest.mark.parametrize( + ("slice_dim", "message"), + [ + ("not-a-slice", "Invalid slice_dim format"), + ("7:0:1", "out of range"), + ("0:0:1", "cannot contain x or y"), + ], +) +def test_binned_rejects_a_bad_slice_dim(client, slice_dim, message): + write_step_scan(client) + + response = client.context.http_client.get( + "/api/v1/binned/scan", params={"slice_dim": slice_dim} + ) + + assert response.status_code == 400 + assert message in response.json()["detail"] + + +# x_dim_index / y_dim_index + + +def test_binned_swapping_the_dimension_indices_transposes_the_image(client): + write_points(client, "swap", [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]) + http = client.context.http_client + size = {"width": 2, "height": 2} + + normal = http.get("/api/v1/binned/swap", params=size).json() + swapped = http.get( + "/api/v1/binned/swap", params={**size, "x_dim_index": 1, "y_dim_index": 0} + ).json() + + assert normal["RedTotal"] == [[0.0, 1.0], [2.0, 3.0]] + assert swapped["RedTotal"] == [[0.0, 2.0], [1.0, 3.0]] + + +# where the positions live + + +@pytest.mark.parametrize("where", ["internal", "named"]) +def test_binned_finds_positions_wherever_the_run_stores_them(client, where): + # older fly scans: primary/internal/sample_stage-x; others: primary/sample_stage-x + x, y = [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0] + if where == "internal": + write_points(client, where, x, y, where="internal") + else: + write_points(client, where, x, y, channels="named") + + body = client.context.http_client.get( + f"/api/v1/binned/{where}", params={"width": 2, "height": 2} + ).json() + + assert body["RedTotal"] == [[0.0, 1.0], [2.0, 3.0]] + assert body["n_points"] == 4 + + +# length reconciliation during a live write + + +def test_binned_truncates_to_the_shortest_channel_if_they_never_agree( + client, fast_reconciliation +): + primary = client.create_container("ragged").create_container("primary") + primary.write_array(numpy.array([0.0, 1.0, 0.0, 1.0]), key="X") + primary.write_array(numpy.array([0.0, 0.0, 1.0, 1.0]), key="Y") + primary.write_array(numpy.arange(3.0), key="RedTotal") # one frame behind + primary.write_array(numpy.arange(4.0), key="GreenTotal") + primary.write_array(numpy.arange(4.0), key="BlueTotal") + + body = client.context.http_client.get( + "/api/v1/binned/ragged", params={"width": 2, "height": 2} + ).json() + + assert body["n_points"] == 3 + assert body["GreenTotal"] == [[0.0, 1.0], [2.0, None]] + + +def test_binned_waits_for_a_channel_that_catches_up( + client, fast_reconciliation, monkeypatch +): + write_points(client, "catching", [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]) + real = visr.get_data + seen = {"RedTotal": 0} + + async def lagging(root, segments): + data = await real(root, segments) + if segments[-1] == "RedTotal": + seen["RedTotal"] += 1 + if seen["RedTotal"] == 1: + return data[:3] # the first read lands a frame behind + return data + + monkeypatch.setattr(visr, "get_data", lagging) + + body = client.context.http_client.get( + "/api/v1/binned/catching", params={"width": 2, "height": 2} + ).json() + + assert body["n_points"] == 4 + assert body["RedTotal"] == [[0.0, 1.0], [2.0, 3.0]] + assert seen["RedTotal"] >= 2 + + +def test_binned_truncates_positions_that_arrive_at_different_lengths( + client, fast_reconciliation +): + primary = client.create_container("axes").create_container("primary") + primary.write_array(numpy.array([0.0, 1.0, 0.0, 1.0]), key="X") + primary.write_array(numpy.array([0.0, 0.0, 1.0]), key="Y") # one point behind + for channel in ("RedTotal", "GreenTotal", "BlueTotal"): + primary.write_array(numpy.arange(4.0), key=channel) + + body = client.context.http_client.get( + "/api/v1/binned/axes", params={"width": 2, "height": 2} + ).json() + + assert body["n_points"] == 3 + + +# a file still behind the catalog, through the whole app + + +@pytest.mark.parametrize(("lag", "status"), [(2, 200), (100, 503)]) +def test_binned_retries_a_lagging_file_then_answers_503( + app_and_client, monkeypatch, fast_retries, lag, status +): + app, client = app_and_client + write_points(client, "lagging", [0.0, 1.0, 0.0, 1.0], [0.0, 0.0, 1.0, 1.0]) + tree = app.state.root_tree + real_lookup = tree.lookup_adapter + reads = {"n": 0} + + class Lagging: + def __init__(self, adapter): + self.adapter = adapter + + async def read(self, *args, **kwargs): + reads["n"] += 1 + if reads["n"] <= lag: + raise DataNotReadyError("advertises (4,), only (3,) available") + result = self.adapter.read(*args, **kwargs) + return await result if inspect.isawaitable(result) else result + + async def lookup(segments): + adapter = await real_lookup(segments) + # the catalog calls lookup_adapter recursively with shorter paths too + return Lagging(adapter) if segments and segments[-1] == "RedTotal" else adapter + + monkeypatch.setattr(tree, "lookup_adapter", lookup) + + response = client.context.http_client.get( + "/api/v1/binned/lagging", params={"width": 2, "height": 2} + ) + + assert response.status_code == status, response.text + if status == 503: + assert response.headers["retry-after"] == "1"