Publish the FLIP tutorial models
d4700c7 verified 8 days ago
27.9 kB
# Copyright (c) 2026 Guy's and St Thomas' NHS Foundation Trust & King's College London
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#     http://www.apache.org/licenses/LICENSE-2.0
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Figures for the image-synthesis tutorials, for local development only.
``save_grid`` and ``save_triplanar`` write debug images during a simulator run; run as a script, this
module plots each site's losses from a simulator workspace. It sits in ``image_synthesis/dev/``,
which is never uploaded: each tutorial's ``job.py`` stages it into simulator runs, and on a trust the
trainers' import of it fails and debug output is simply off.
Images are written only when ``config.json`` sets ``SAVE_DEBUG_SAMPLES`` **and** this is a
Makefile-launched simulator run (see :func:`samples_enabled`), into ``DEBUG_SAMPLES_DIR``. Nothing
here sends an image anywhere.
Usage, from a tutorial's directory after ``make run``::
    python ../dev/plot_utils.py                          # reads /tmp/nvflare/<this tutorial>/flip_fedavg
    python ../dev/plot_utils.py --workspace <dir> --out-dir <dir>
"""

from __future__ import annotations

import argparse
import logging
import os
import re
from pathlib import Path

import numpy as np
import torch
from flip.constants import FlipConstants
from torchvision.utils import make_grid, save_image

logger = logging.getLogger(__name__)

#: Set by the tutorial Makefiles' ``sim`` recipe only: where the images go, and half of the simulator
#: check in :func:`in_simulation`. Never a ``config.json`` key.
DEBUG_DIR_ENV = "DEBUG_SAMPLES_DIR"

#: How many images to write per grid when ``DEBUG_SAMPLES_MAX`` is absent from the config.
DEFAULT_MAX_IMAGES = 8

#: Training iterations between tri-planar figures when ``DEBUG_PLOT_EVERY`` is absent.
DEFAULT_PLOT_EVERY = 100

#: The three orthogonal views: display name, the array axis each one holds fixed, and whether the
#: slice needs a left-right flip after the shared rotation.
#:
#: Volumes reach here as ``(C, X, Y, Z)`` with the chain's ``Orientationd(axcodes="RAS")``
#: guaranteeing what the axes mean: +X is right, +Y anterior, +Z superior. So a sagittal view fixes
#: X, a coronal view fixes Y, and an axial view fixes Z. Without that guarantee these labels — and
#: the flip below — would be guesses.
#:
#: One ``np.rot90`` puts the vertical axis the right way up in all three (superior up for sagittal
#: and coronal, anterior up for axial), but it does not settle the horizontal one, and the sagittal
#: view is the odd one out. Coronal and axial both have **X** across the page, so rot90 leaves them
#: with the patient's left on the viewer's left — neurological convention. Sagittal has **Y** across
#: the page instead, which leaves anterior on the *right*: the head faces backwards compared with how
#: every other view here, and every standard neuroimaging viewer, shows it. Hence the flip.
_PLANES = (("sagittal", 0, True), ("coronal", 1, False), ("axial", 2, False))


def in_simulation() -> bool:
    """Whether this process is a local simulator run launched by a tutorial Makefile.
    Needs **both** signals, so that neither alone can open the gate on a trust:
    - ``FlipConstants.LOCAL_DEV`` is true. Every trust deployment (compose and Helm) sets it to
      ``false``, and the ``flip`` library already reads its data from the trust APIs on that value,
      so a client with it true cannot be running a real cohort. It is not enough on its own because
      it *defaults* to true when the variable is unset.
    - ``DEBUG_SAMPLES_DIR`` is set and non-blank. Only the tutorial Makefiles' ``sim`` recipe sets
      it; it is not a ``config.json`` key, so nothing uploaded with the app can supply it.
    Returns:
        bool: True only inside a Makefile-launched simulator run.
    """
    if not os.environ.get(DEBUG_DIR_ENV, "").strip():
        return False
    try:
        return FlipConstants.LOCAL_DEV is True
    except Exception:  # unreadable settings: fail closed
        return False


def samples_enabled(config: dict) -> bool:
    """Whether this job writes client-local debug images.
    True only when ``config.json`` sets ``SAVE_DEBUG_SAMPLES`` **and** :func:`in_simulation` holds.
    The config is the user's switch — false means no images, simulator included. It is never what
    separates simulation from production, because ``config.json`` travels with the app to the trust:
    in production (``LOCAL_DEV=false``, no ``DEBUG_SAMPLES_DIR``) this returns False whatever the
    config says, so patient-derived pixels are never written at a trust.
    This is the single gate for both writers. ``save_grid`` and ``save_triplanar`` consult it and
    no-op when it is false, so a caller must never add a second condition of its own: a sampling
    block that runs on a looser test than this one pays for its output and then discards it.
    Args:
        config (dict): The parsed ``config.json``.
    Returns:
        bool: True when ``SAVE_DEBUG_SAMPLES`` is truthy and this is a simulator run.
    """
    return bool(config.get("SAVE_DEBUG_SAMPLES", False)) and in_simulation()


def plot_every(config: dict) -> int:
    """Iterations between tri-planar figures, from ``DEBUG_PLOT_EVERY`` (default 100).
    Args:
        config (dict): The parsed ``config.json``.
    Returns:
        int: The interval. Zero or negative disables the figures while leaving the grids on.
    """
    return int(config.get("DEBUG_PLOT_EVERY", DEFAULT_PLOT_EVERY))


def due_for_plot(iteration: int, config: dict) -> bool:
    """Whether ``iteration`` (0-based) is one of the every-N steps that gets a tri-planar figure.
    True on iteration 0 as well as every Nth after it, so a run that crashes early still leaves one
    figure behind — the first is also the most diagnostic, being the only look at an untrained model.
    Args:
        iteration (int): Zero-based training step, counted across rounds and epochs (not within one
            epoch — an epoch shorter than the interval would otherwise fire once per epoch).
        config (dict): The parsed ``config.json``.
    Returns:
        bool: True when a figure is due and debug output is enabled at all.
    """
    if not samples_enabled(config):
        return False
    interval = plot_every(config)
    return interval > 0 and iteration % interval == 0


def debug_dir() -> Path:
    """``DEBUG_SAMPLES_DIR``, created if needed; the tutorial Makefiles point it under fl-tutorials/data/."""
    target = Path(os.environ[DEBUG_DIR_ENV].strip()).expanduser().resolve()
    target.mkdir(parents=True, exist_ok=True)
    return target


def _to_2d(batch: torch.Tensor) -> torch.Tensor:
    """Reduce a batch to something ``make_grid`` can tile, taking a mid-slice of any volume.
    Args:
        batch (torch.Tensor): ``(B, C, H, W)`` or ``(B, C, H, W, D)``.
    Returns:
        torch.Tensor: ``(B, C, H, W)`` — volumes sliced at the centre of the last axis, which for the
        RAS-oriented volumetric chain is the axial mid-slice.
    """
    if batch.ndim == 5:
        return batch[..., batch.shape[-1] // 2]
    return batch


def _name_parts(name: str, site_name: str, step: int | None) -> list[str]:
    """Filename parts shared by both writers: what, where, when — in that order."""
    parts = [name]
    if site_name:
        parts.append(site_name)
    if step is not None:
        parts.append(f"step{step:04d}")
    return parts


def _prepare(batch: torch.Tensor, max_images: int) -> torch.Tensor:
    """Detach, truncate and flatten a batch to greyscale 2-D images ready for tiling."""
    prepared = _to_2d(batch.detach().float().cpu()[:max_images])
    if prepared.shape[1] > 1:
        # Multi-channel here means a latent, not a colour image; show the first channel rather than
        # letting save_image interpret 3 channels as RGB.
        prepared = prepared[:, :1]
    return prepared


def save_grid(
    tensors: dict[str, torch.Tensor],
    name: str,
    config: dict,
    *,
    site_name: str = "",
    step: int | None = None,
) -> Path | None:
    """Write one PNG tiling the given batches, one row per entry. No-op unless enabled.
    Rows are normalised **jointly**, not per image: independent scaling would rescale a washed-out
    reconstruction back onto its input's range and hide precisely the intensity drift the comparison
    is meant to expose.
    Never raises. Debug output must not be able to fail a training round, so any error (a read-only
    workspace, a tensor of an unexpected rank) is logged and swallowed.
    Args:
        tensors (dict[str, torch.Tensor]): Batches to tile, one row each, in insertion order — e.g.
            ``{"input": images, "reconstruction": recon}``. Keys appear in the filename.
        name (str): Short label for what this grid shows, e.g. ``"reconstruction"``.
        config (dict): The parsed ``config.json``; read for ``SAVE_DEBUG_SAMPLES``/``DEBUG_SAMPLES_MAX``.
        site_name (str): Client name, included in the filename so a shared workspace cannot collide.
        step (int | None): Round or epoch counter, included in the filename to order the images.
    Returns:
        Path | None: The file written, or None when disabled or on any failure.
    """
    if not samples_enabled(config) or not tensors:
        return None

    try:
        max_images = int(config.get("DEBUG_SAMPLES_MAX", DEFAULT_MAX_IMAGES))
        rows = [_prepare(batch, max_images) for batch in tensors.values()]
        columns = min(row.shape[0] for row in rows)
        grid = make_grid(
            torch.cat([row[:columns] for row in rows]),
            nrow=columns,
            normalize=True,
            scale_each=False,
        )

        label = "_".join([name, *tensors.keys()]) if len(tensors) > 1 else name
        path = debug_dir() / ("_".join(_name_parts(label, site_name, step)) + ".png")

        save_image(grid, str(path))
        logger.info(f"[DEBUG] Wrote {columns} sample(s) to {path}")
        return path
    except Exception as err:
        logger.warning(f"[DEBUG] Could not write debug samples {name!r}: {err}")
        return None


def _central_slice(volume: torch.Tensor, axis: int, flip_lr: bool = False) -> np.ndarray:
    """The central slice of one 3-D volume along ``axis``, oriented for display.
    Args:
        volume (torch.Tensor): A single volume, ``(X, Y, Z)``, channel already removed.
        axis (int): Spatial axis to hold fixed — 0 sagittal, 1 coronal, 2 axial.
        flip_lr (bool): Mirror the slice horizontally after the rotation. True for the sagittal
            view only; see :data:`_PLANES` for why it alone needs it.
    Returns:
        np.ndarray: The 2-D slice, rotated so the axis that runs superior (or anterior, for the
        axial view) points up the page, and mirrored if this view needs it.
    """
    plane = volume.index_select(axis, torch.tensor([volume.shape[axis] // 2], device=volume.device)).squeeze(axis)
    rotated = np.rot90(plane.numpy())
    return np.fliplr(rotated) if flip_lr else rotated


def save_triplanar(
    volumes: dict[str, torch.Tensor],
    name: str,
    config: dict,
    *,
    site_name: str = "",
    step: int | None = None,
    sample: int = 0,
) -> Path | None:
    """Write a 3x2 tri-planar figure: one row per view, one column per batch. No-op unless enabled.
    Rows are sagittal / coronal / axial central cuts; columns are the entries of ``volumes`` in
    insertion order — so ``{"input": images, "reconstruction": recon}`` puts the input beside what
    the model made of it, in the same three planes.
    This is the figure that answers questions a mid-slice grid cannot. An autoencoder can look
    plausible on one axial slice while having learnt nothing about through-plane structure, and that
    shows up immediately in the sagittal and coronal views.
    **Each panel is scaled to its own range**, unlike :func:`save_grid`, which normalises jointly.
    The trade-off is deliberate and goes the other way here: a joint scale is the honest one for
    spotting intensity drift, but an early reconstruction spans a far wider range than the input
    (noise, including negatives), so a shared scale compresses the input to flat grey and neither
    column can be read. These figures exist to show *structure* — is there a brain, are the
    ventricles in the right place, does the tumour survive — and structure needs per-panel contrast.
    Read absolute intensity off the grids, not off these.
    Never raises: debug output must not be able to fail a training round.
    Args:
        volumes (dict[str, torch.Tensor]): Batches to compare, ``(B, C, X, Y, Z)``, one column each.
        name (str): Short label for what this figure shows, used in the filename.
        config (dict): The parsed ``config.json``.
        site_name (str): Client name, included in the filename so a shared workspace cannot collide.
        step (int | None): Iteration or epoch counter, included in the filename to order the figures.
        sample (int): Which item of the batch to show.
    Returns:
        Path | None: The file written, or None when disabled, when the tensors are not volumetric,
        or on any failure.
    """
    if not samples_enabled(config) or not volumes:
        return None

    try:
        import matplotlib

        # Agg before pyplot: an FL client has no display, and the default backend would try to pick
        # an interactive one and fail at import inside a training round.
        matplotlib.use("Agg")
        import matplotlib.pyplot as plt

        picked = {}
        for label, batch in volumes.items():
            if batch.ndim != 5:
                logger.info(f"[DEBUG] {name!r}: {label} is not volumetric ({tuple(batch.shape)}); skipping figure.")
                return None
            picked[label] = batch.detach().float().cpu()[min(sample, batch.shape[0] - 1), 0]

        figure, axes = plt.subplots(len(_PLANES), len(picked), figsize=(3.2 * len(picked), 3.2 * len(_PLANES)))
        axes = np.atleast_2d(axes).reshape(len(_PLANES), len(picked))
        figure.patch.set_facecolor("black")

        for row, (view, axis, flip_lr) in enumerate(_PLANES):
            for column, (label, volume) in enumerate(picked.items()):
                axis_handle = axes[row][column]
                # No vmin/vmax: matplotlib scales each panel to its own min/max, which is the point
                # (see the note on contrast above). A flat panel is left to matplotlib, which
                # renders a constant array as mid-grey rather than dividing by a zero range.
                axis_handle.imshow(_central_slice(volume, axis, flip_lr), cmap="gray")
                axis_handle.set_facecolor("black")
                axis_handle.axis("off")
                if row == 0:
                    axis_handle.set_title(label, color="0.85", fontsize=11)
                if column == 0:
                    # A y-label survives axis("off"), which is why the view name rides here rather
                    # than in a title that would collide with the column headings.
                    axis_handle.set_ylabel(view, color="0.85", fontsize=10)
                    axis_handle.axis("on")
                    axis_handle.set_xticks([])
                    axis_handle.set_yticks([])
                    for spine in axis_handle.spines.values():
                        spine.set_visible(False)

        figure.tight_layout()
        path = debug_dir() / ("_".join(_name_parts(name, site_name, step)) + ".png")
        figure.savefig(path, facecolor=figure.get_facecolor(), dpi=110)
        plt.close(figure)
        logger.info(f"[DEBUG] Wrote tri-planar figure to {path}")
        return path
    except Exception as err:
        logger.warning(f"[DEBUG] Could not write tri-planar figure {name!r}: {err}")
        return None


# --------------------------------------------------------------------------------------------
# Analysis half: loss curves from a simulator log. Host-side only — see the module docstring.
# --------------------------------------------------------------------------------------------

#: Log lines arrive as ``<timestamp> - <logger> - INFO - <message>``. Splitting the message off is
#: what stops the timestamp's own colons ("20:10:04") being read as data. A line that does *not*
#: match is a continuation of the message above it — the per-epoch summary is logged as one
#: multi-line string, so its numbers live on unprefixed lines — and is parsed whole.
_MESSAGE_RE = re.compile(r" - (?:DEBUG|INFO|WARNING|ERROR|CRITICAL) - (?P<message>.*)$")

#: One ``name: number`` pair inside a message. Deliberately generic: every tutorial here logs its
#: losses in that shape but with different names, so hard-coding one tutorial's names would leave
#: the other two with empty plots. The name charset excludes quotes,
#: commas and hyphens, which is what keeps a name from bridging across ``a: 1, b: 2`` or reaching
#: back into the logger prefix; ``nan``/``inf`` are accepted as values so a run that broke still
#: plots up to the point it broke.
_PAIR_RE = re.compile(r"(?P<name>[A-Za-z][\w ()]*?)\s*:\s*(?P<value>[-+]?(?:\d+\.?\d*(?:[eE][-+]?\d+)?|nan|inf))\b")

#: Fixed colours for the names this tree logs today, so two sites' figures can be read side by side
#: and successive runs stay comparable. Anything unrecognised is assigned from _PALETTE in sorted
#: order, which is deterministic for a given set of names.
COLOURS = {
    "l1": "#4ea1ff",
    "perceptual": "#ffd166",
    "gan": "#ff6b6b",
    "KL": "#9b7bff",
    "Total loss D": "#7bffb0",
    "Total loss DM": "#4ea1ff",
    "Validation loss (L1)": "#4ea1ff",
    "Validation SSIM (foreground)": "#ff9f4e",
    "Validation DM": "#ff9f4e",
    # The per-epoch aggregates, coloured to match the per-iteration term each one totals.
    "L1": "#4ea1ff",
    "Perceptual": "#ffd166",
    "GAN": "#ff6b6b",
    "KLD": "#9b7bff",
    "Total loss G": "#eeeeee",
}
_PALETTE = ("#4ea1ff", "#ffd166", "#ff6b6b", "#9b7bff", "#7bffb0", "#ff9f4e", "#59d4d4", "#d98cc4")

#: Smoothing for the per-iteration panel, which is unreadable raw at BATCH_SIZE 2.
EMA_ALPHA = 0.02

#: A series needs at least this many points to be worth an axis; it also filters out the stray
#: ``name: number`` pairs that appear once in a startup banner.
MIN_POINTS = 2

#: Series logged at the same cadence share a length. Anything within this fraction of the longest
#: series is treated as per-iteration; the rest is per-epoch. A tutorial that logs only per-epoch
#: (the two diffusion ones) therefore yields a single panel rather than an empty one.
#:
#: The same ratio is then applied *within* the per-epoch band, which is what keeps incidental
#: ``name: number`` pairs out of the figure. A log is full of them — ``localhost: 34953``,
#: ``PID: 1005143``, ``train size: 30`` — and they are not distinguishable from a metric by name or
#: by value. They are distinguishable by cadence: a metric is logged once per iteration or once per
#: epoch, so it appears as many times as there were iterations or epochs, while a startup line
#: appears once or twice however long the run goes on.
DENSE_RATIO = 0.9


def _ema(values: list[float], alpha: float = EMA_ALPHA) -> list[float]:
    """Exponential moving average, skipping non-finite points rather than propagating them."""
    out: list[float] = []
    acc = next((v for v in values if v == v and abs(v) != float("inf")), 0.0)
    for value in values:
        if value == value and abs(value) != float("inf"):
            acc = alpha * value + (1 - alpha) * acc
        out.append(acc)
    return out


def parse_log(path: Path) -> dict[str, list[float]]:
    """Return every ``name: number`` series found in one site's ``log.txt``, in order of appearance.
    Args:
        path (Path): The site's ``log.txt``.
    Returns:
        dict[str, list[float]]: Series name to its values, longest series first.
    """
    series: dict[str, list[float]] = {}
    with path.open(errors="replace") as handle:
        for line in handle:
            prefixed = _MESSAGE_RE.search(line)
            message = prefixed.group("message") if prefixed else line
            for match in _PAIR_RE.finditer(message):
                series.setdefault(match.group("name"), []).append(float(match.group("value")))
    return dict(sorted(series.items(), key=lambda item: -len(item[1])))


def split_by_cadence(series: dict[str, list[float]]) -> tuple[dict[str, list[float]], dict[str, list[float]]]:
    """Split series into per-iteration and per-epoch groups by how often each was logged.
    Args:
        series (dict[str, list[float]]): Output of :func:`parse_log`.
    Returns:
        tuple: ``(dense, sparse)``; ``sparse`` is empty when everything shares one cadence.
    """
    usable = {name: values for name, values in series.items() if len(values) >= MIN_POINTS}
    if not usable:
        return {}, {}
    longest = max(len(values) for values in usable.values())
    dense = {n: v for n, v in usable.items() if len(v) >= longest * DENSE_RATIO}

    remainder = {n: v for n, v in usable.items() if len(v) < longest * DENSE_RATIO}
    if not remainder:
        return dense, {}
    epoch_cadence = max(len(values) for values in remainder.values())
    sparse = {n: v for n, v in remainder.items() if len(v) >= epoch_cadence * DENSE_RATIO}
    return dense, sparse


def _colour_for(name: str, index: int) -> str:
    return COLOURS.get(name, _PALETTE[index % len(_PALETTE)])


def _style(axis, title: str) -> None:
    """Apply the shared dark styling to one axis."""
    axis.set_facecolor("#111111")
    axis.set_title(title, color="#eeeeee")
    axis.tick_params(colors="#bbbbbb")
    axis.grid(True, color="#333333", linewidth=0.5)
    for spine in axis.spines.values():
        spine.set_color("#444444")
    axis.xaxis.label.set_color("#bbbbbb")
    axis.yaxis.label.set_color("#bbbbbb")


def _legend(axis, handles=None, labels=None) -> None:
    """Legend outside the axes on the right, so it cannot sit on top of the curves.
    The generic parser finds however many series a tutorial logs — eight for the autoencoder's
    per-epoch panel — which is more than fits inside the plot area.
    """
    axis.legend(
        *([handles, labels] if handles is not None else []),
        loc="upper left",
        bbox_to_anchor=(1.12, 1.0),
        fontsize=8,
        facecolor="#1c1c1c",
        edgecolor="#444444",
        labelcolor="#eeeeee",
    )


def plot_site(site: str, series: dict[str, list[float]], out_dir: Path) -> Path:
    """Write the loss figure for one site and return its path.
    The per-iteration panel is log-scaled when every value is positive: the terms routinely span two
    decades, and on a linear axis the smallest of them — usually the one that is misweighted, which
    is exactly what you are looking for — is a flat line on the floor.
    Args:
        site (str): Site name, used in the title and filename.
        series (dict[str, list[float]]): Output of :func:`parse_log`.
        out_dir (Path): Directory to write into; created if absent.
    Returns:
        Path: The figure written.
    """
    import matplotlib

    matplotlib.use("Agg")
    import matplotlib.pyplot as plt

    dense, sparse = split_by_cadence(series)
    panels = [p for p in (dense, sparse) if p]
    figure, axes = plt.subplots(len(panels), 1, figsize=(13, 4 * len(panels)), facecolor="#111111", squeeze=False)

    for axis, panel in zip(axes.flatten(), panels, strict=True):
        is_dense = panel is dense
        # SSIM is bounded [0, 1] where the losses are not, so it gets its own axis or it flattens
        # every other series against the top of the plot.
        twin_names = [n for n in panel if "SSIM" in n] if not is_dense else []
        twin = axis.twinx() if twin_names else None

        for index, (name, values) in enumerate(panel.items()):
            colour = _colour_for(name, index)
            target = twin if name in twin_names else axis
            if is_dense:
                target.plot(values, color=colour, alpha=0.15, linewidth=0.6)
                target.plot(_ema(values), color=colour, linewidth=1.8, label=name)
            else:
                steps = range(1, len(values) + 1)
                target.plot(steps, values, color=colour, marker="o", markersize=3, label=name)

        # Log y whenever everything is positive. The terms routinely span two or three decades, and
        # on a linear axis the smallest of them — usually the misweighted one, which is exactly what
        # you are looking for — is a flat line on the floor. On the per-epoch panel the first epoch
        # is several times any later one, so a linear axis spends most of its height on it.
        plotted = [values for name, values in panel.items() if name not in twin_names]
        finite = [v for values in plotted for v in values if v == v and abs(v) != float("inf")]
        if finite and min(finite) > 0:
            axis.set_yscale("log")
        axis.set_xlabel("training iteration" if is_dense else "local epoch")
        axis.set_ylabel("weighted contribution to $\\mathcal{L}_G$" if is_dense else "loss")
        _style(axis, f"{site} — {'per-iteration loss terms' if is_dense else 'per-epoch summary'}")

        if twin is not None:
            twin.set_ylabel(" / ".join(twin_names), color=COLOURS.get(twin_names[0], _PALETTE[5]))
            twin.tick_params(axis="y", colors=COLOURS.get(twin_names[0], _PALETTE[5]))
            for spine in twin.spines.values():
                spine.set_color("#444444")
            handles, labels = axis.get_legend_handles_labels()
            extra_handles, extra_labels = twin.get_legend_handles_labels()
            _legend(axis, handles + extra_handles, labels + extra_labels)
        else:
            _legend(axis)

    figure.tight_layout()
    out_dir.mkdir(parents=True, exist_ok=True)
    destination = out_dir / f"losses_{site}.png"
    figure.savefig(destination, dpi=110, facecolor=figure.get_facecolor())
    plt.close(figure)
    return destination


def main() -> None:
    """Plot every site's losses from a simulator workspace. Host-side entry point."""
    parser = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    parser.add_argument(
        "--workspace",
        type=Path,
        default=Path("/tmp/nvflare") / Path.cwd().name / "flip_fedavg",
        help="simulator workspace root (default: this tutorial's, from the current directory)",
    )
    parser.add_argument("--out-dir", type=Path, default=None, help="defaults to $DEBUG_SAMPLES_DIR, else ./")
    args = parser.parse_args()

    out_dir = args.out_dir or Path(os.environ.get(DEBUG_DIR_ENV, "."))

    logs = sorted(args.workspace.glob("site-*/log.txt"))
    if not logs:
        raise SystemExit(f"no site logs under {args.workspace} — pass --workspace")

    for log in logs:
        site = log.parent.name
        series = parse_log(log)
        dense, sparse = split_by_cadence(series)
        if not dense and not sparse:
            print(f"{site}: no scalar lines yet, skipping")
            continue
        destination = plot_site(site, series, out_dir)
        counts = ", ".join(f"{name} x{len(values)}" for name, values in list(dense.items())[:4])
        print(f"{site}: {counts} -> {destination}")


if __name__ == "__main__":
    main()