sras-viewer/sras_render.py

"""Pure-matplotlib rendering of a displayed SRAS image: imshow + colorbar +
axis/title labeling, shared by the interactive Qt canvas
(sras_viewer.canvases.ImageCanvas, which supplies its own already-Qt-backed
Figure/Axes) and the headless batch image-export worker (which builds a
throwaway Agg Figure per file and never touches Qt) -- so an exported PNG
can never quietly start looking different from what the GUI actually shows.

Deliberately no PyQt6 import anywhere in this module: BatchExportImagesWorker
(sras_workers.py) may run export_view_image inside a spawned
ProcessPoolExecutor subprocess, exactly like sras_compute.cache_file, and
importing anything under the sras_viewer package would run its __init__.py
and pull in the whole Qt widget tree for no reason.
"""

from pathlib import Path

import matplotlib as mpl
import numpy as np
from matplotlib.backends.backend_agg import FigureCanvasAgg
from matplotlib.figure import Figure

from sras_compute import compute_rf_image, dc_image_mv
from sras_format import CH4_IDX, CH_NAMES, SrasFile, _axes_extent

# Figure size (inches) to fall back on when a caller has no live view to
# match: the size ImageCanvas is constructed at, so a headless render with
# no canvas behind it still looks like the viewer's starting layout.
DEFAULT_FIGSIZE = (7.0, 5.0)

# The canvas size travels from the GUI as-is, so it can be degenerate (a
# collapsed splitter pane, a window minimized mid-batch) or, on a very wide
# multi-monitor window, big enough that width * dpi approaches matplotlib's
# 2**16 pixel limit. Clamp rather than fail: an export's aspect ratio is a
# presentation detail and must never be the reason a batch loses an image.
_MIN_FIG_IN = 1.0
_MAX_FIG_IN = 100.0


def sanitize_figsize(figsize) -> tuple[float, float]:
    """(width, height) in inches, clamped to something renderable.

    *figsize* is None (use DEFAULT_FIGSIZE) or any 2-sequence of numbers --
    including the float64 pair Figure.get_size_inches returns, which is how
    the GUI hands over the live canvas's current size.
    """
    try:
        w, h = float(figsize[0]), float(figsize[1])
    except (TypeError, ValueError, IndexError, KeyError):
        return DEFAULT_FIGSIZE
    if not (np.isfinite(w) and np.isfinite(h)):
        return DEFAULT_FIGSIZE
    return (min(max(w, _MIN_FIG_IN), _MAX_FIG_IN),
            min(max(h, _MIN_FIG_IN), _MAX_FIG_IN))


def draw_view_image(ax, fig, img: np.ndarray, extent: list[float], cmap,
                    vmin: float, vmax: float, xlabel: str, ylabel: str,
                    title: str, colorbar_label: str = "", cb_ticks=None,
                    norm=None, bad_color=None):
    """imshow + colorbar + labels onto an already-created (ax, fig) pair.

    *cmap* may be a name or a Colormap instance. *norm* (which overrides
    vmin/vmax) and *cb_ticks* let a caller draw a discrete integer image
    with whole-number colorbar bands instead of a continuous shade.
    *bad_color*, if given, is the fill for NaN pixels -- a copy of *cmap* is
    made so a shared, registered instance is never mutated.

    Shared by ImageCanvas.show_image (Qt-backed ax/fig) and
    export_view_image (headless Agg ax/fig) so the two can never drift into
    showing different things for the same settings.
    """
    if bad_color is not None:
        cmap = (cmap if hasattr(cmap, "with_extremes")
               else mpl.colormaps[cmap]).with_extremes(bad=bad_color)

    kw = ({"norm": norm} if norm is not None
          else {"vmin": vmin, "vmax": vmax})
    im = ax.imshow(
        img, aspect="auto", origin="upper",
        extent=extent, cmap=cmap, interpolation="nearest", **kw,
    )
    cb = fig.colorbar(im, ax=ax, fraction=0.046, pad=0.04, ticks=cb_ticks)
    if colorbar_label:
        cb.set_label(colorbar_label)

    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)
    ax.set_title(title)
    return im


def export_view_image(path: str, *, out_dir: str, angle_idx: int, ch_idx: int,
                      is_fft_mode: bool, is_velocity: bool,
                      dc_threshold_mv: float, apply_bg_sub: bool,
                      pad_factor: int, min_freq_mhz: float, grating_um: float,
                      cmap: str, auto_scale: bool, vmin: float, vmax: float,
                      highlight_masked: bool, mode_str: str,
                      colorbar_label: str, mask_color: str = "magenta",
                      max_workers: int | None = None,
                      figsize: tuple[float, float] | None = None,
                      dpi: int = 150) -> tuple[str, str]:
    """One file's contribution to Batch Export View as Images: renders
    (angle_idx, ch_idx) at the given display settings to a PNG under
    *out_dir*, via draw_view_image -- so a batch export is a folder of what
    ImageCanvas.show_image would have put on screen for these settings, not
    a raw data dump.

    Module-level and picklable, like sras_compute.cache_file, so it can run
    in a ProcessPoolExecutor -- see BatchExportImagesWorker. Unlike
    cache_file this never writes to *path*: export is a read of the file's
    own data, not a cache conversion, so any version SrasFile can open
    works, with no v6/v7 precondition.

    *figsize* is the on-screen ImageCanvas's current size in inches, so the
    PNG carries the aspect ratio the view was actually being read at. It
    matters more here than it would for a plot with fixed data aspect:
    draw_view_image uses aspect="auto", so the image stretches to whatever
    box it is given -- rendering a view the user has sized wide into a
    hard-coded 7x5 squeezes the map into a different shape than the one
    they judged it by. Passing the full size, not just the ratio, also
    keeps titles, tick labels and the colorbar in the same proportion to
    the map as on screen; *dpi* alone then sets the output resolution.

    *pad_factor* (not n_fft) travels across files deliberately: n_fft
    depends on samples_per_frame, which can differ between files in the
    same batch, so n_fft is derived per file, here, from *this* file's own
    value -- the same reason sras_compute.cache_file does the same thing.

    row_avg_n is per file for the same reason and so is not a parameter at
    all: it comes from *this* file's stored cache, mirroring
    SrasViewerWindow._stored_fft_image and _start_compute, which both ask at
    the window size the store was written at. Leaving it at compute_rf_image's
    "raw per-pixel" default would make a row-averaged cache a mismatch, so a
    file the viewer displays from its store would export as a full raw
    recompute instead -- a different (and much slower) image than the one the
    batch was triggered to reproduce.

    Returns (error, out_name). error is "" on success. out_name is the
    filename this call targeted -- set as soon as it's known, even on most
    failures -- so the caller can flag same-stem collisions across the
    batch without any cross-process bookkeeping.
    """
    out_name = ""
    try:
        sras = SrasFile(path)
        if angle_idx >= sras.n_angles:
            return (f"angle {angle_idx} out of range "
                    f"(file has {sras.n_angles} angle(s))", out_name)

        out_name = f"{Path(path).stem}_angle{angle_idx}_{CH_NAMES[ch_idx]}.png"

        if is_fft_mode:
            n_fft = (sras.samples_per_frame * pad_factor
                     if pad_factor > 1 else None)
            freq = compute_rf_image(
                sras, angle_idx, dc_threshold_mv=dc_threshold_mv,
                apply_bg_sub=apply_bg_sub, n_fft=n_fft,
                row_avg_n=sras.precomputed_row_avg_n,
                min_freq_mhz=min_freq_mhz, max_workers=max_workers)
            img = freq * grating_um if is_velocity else freq
        else:
            img = dc_image_mv(sras, angle_idx, ch_idx, max_workers=max_workers)

        display_img, bad_color = img, None
        if highlight_masked and is_fft_mode:
            # Mirrors the viewer's _redraw_image rule: DC-masked pixels and
            # value-0 pixels (the "no valid peak" sentinel — DC-masked,
            # below the min-freq floor, or empty spectrum; the grating
            # multiply above preserves zeros, so this holds for Velocity
            # too) both render in the highlight color.
            dc4 = dc_image_mv(sras, angle_idx, CH4_IDX, max_workers=max_workers)
            valid = dc4 >= dc_threshold_mv
            if valid.shape == display_img.shape:
                valid &= display_img != 0.0
                display_img = display_img.astype(np.float32, copy=True)
                display_img[~valid] = np.nan
                bad_color = mask_color

        if auto_scale:
            v0, v1 = float(np.nanmin(display_img)), float(np.nanmax(display_img))
            if not np.isfinite(v0):
                v0, v1 = 0.0, 0.0   # every pixel masked out
        else:
            v0, v1 = vmin, vmax

        x_axis = sras.x_axis_mm(angle_idx)
        y_axis = sras.y_positions_mm(angle_idx)
        dx = x_axis[1] - x_axis[0] if len(x_axis) > 1 else sras.pixel_x_mm
        dy = float(y_axis[1] - y_axis[0]) if len(y_axis) > 1 else 1.0
        extent = _axes_extent(x_axis, y_axis, dx, dy)

        title = (f"{CH_NAMES[ch_idx]}  |  {mode_str}  |  "
                f"{sras.angles_deg[angle_idx]:.1f}°")

        fig = Figure(figsize=sanitize_figsize(figsize), tight_layout=True)
        FigureCanvasAgg(fig)   # Agg-only: never registered with pyplot
        ax = fig.add_subplot(111)
        draw_view_image(ax, fig, display_img, extent, cmap, v0, v1,
                        "X (mm)", "Y (mm)", title, colorbar_label,
                        bad_color=bad_color)
        fig.savefig(str(Path(out_dir) / out_name), dpi=dpi)
        return ("", out_name)
    except Exception as exc:
        return (str(exc), out_name)