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


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,
                      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.

    *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.

    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,
                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=(7, 5), 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)