From 8f10837bae7b7ca411468396d6069fa504557566 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 6 Oct 2026 11:33:10 -0600
Subject: [PATCH 01/11] Add display normalization module and test. Then hook up
 display normalization to visualization module to make scaling options
 function parameter. This supports different needs in visualization (e.g.
 directly comparing prediction to ground truth vs. displaying predictions and
 ground truth independently in optimal contrast)

---
 .../evaluation/display_normalization.py       |  97 ++++++++++
 .../evaluation/visualization.py               | 106 ++++++----
 .../evaluation/test_display_normalization.py  | 182 ++++++++++++++++++
 3 files changed, 352 insertions(+), 33 deletions(-)
 create mode 100644 src/virtual_stain_flow/evaluation/display_normalization.py
 create mode 100644 tests/evaluation/test_display_normalization.py

diff --git a/src/virtual_stain_flow/evaluation/display_normalization.py b/src/virtual_stain_flow/evaluation/display_normalization.py
new file mode 100644
index 0000000..0d3a0af
--- /dev/null
+++ b/src/virtual_stain_flow/evaluation/display_normalization.py
@@ -0,0 +1,97 @@
+"""Display-only intensity scaling; never modifies image or metric data."""
+
+from typing import List, Optional, Sequence, Tuple, Union
+
+import numpy as np
+from matplotlib.colors import Normalize
+
+DisplayLimits = Union[Tuple[float, float], Sequence[Tuple[float, float]]]
+
+
+class _ClippedNormalize(Normalize):
+    """Handle a constant reference without hiding brighter predictions.
+
+    For a zero-width range, equal values are middle gray, lower values are
+    black and higher values are white (with the gray colormap).
+    """
+
+    def __call__(self, value, clip=None):
+        if self.vmin != self.vmax:
+            return super().__call__(value, clip=clip)
+        result, is_scalar = self.process_value(value)
+        normalized = np.ma.array(
+            np.where(result.data < self.vmin, 0.0,
+                     np.where(result.data > self.vmax, 1.0, 0.5)),
+            mask=np.ma.getmaskarray(result) | np.isnan(result.data),
+        )
+        return normalized[0] if is_scalar else normalized
+
+
+def _fixed_limits(limits: Optional[DisplayLimits], channels: int, name: str) -> np.ndarray:
+    """Validate a common pair or one pair per *displayed* channel."""
+    try:
+        values = np.asarray(limits, dtype=float)
+    except (TypeError, ValueError) as error:
+        raise ValueError(f"{name} must be a (min, max) pair or one pair per displayed channel.") from error
+    if values.shape == (2,):
+        values = np.tile(values, (channels, 1))
+    if values.shape != (channels, 2) or not np.all(np.isfinite(values)):
+        raise ValueError(f"{name} must contain finite (min, max) pairs for {channels} displayed channels.")
+    if np.any(values[:, 0] >= values[:, 1]):
+        raise ValueError(f"{name} requires min < max.")
+    return values
+
+
+def _finite_range(images: np.ndarray) -> Tuple[float, float]:
+    finite = images[np.isfinite(images)]
+    if not finite.size:
+        raise ValueError("Cannot determine display limits from an image/channel with no finite values.")
+    return float(finite.min()), float(finite.max())
+
+
+def image_norms(
+    images: np.ndarray,
+    *,
+    scope: str = "image",
+    limits: Optional[DisplayLimits] = None,
+    other: Optional[np.ndarray] = None,
+    legacy: bool = False,
+) -> List[List[Normalize]]:
+    """Construct per-row/channel norms, optionally pooling another image stack."""
+    if scope not in ("image", "channel"):
+        raise ValueError("Display scale scope must be 'image' or 'channel'.")
+    fixed = _fixed_limits(limits, images.shape[1], "Display limits") if limits is not None else None
+    rows: List[List[Normalize]] = []
+    for row in range(images.shape[0]):
+        if row and (scope == "channel" or fixed is not None):
+            rows.append(rows[0])
+            continue
+        norms = []
+        for channel in range(images.shape[1]):
+            if fixed is not None:
+                lower, upper = fixed[channel]
+            else:
+                reference = images[:, channel] if scope == "channel" else images[row, channel]
+                lower, upper = _finite_range(reference)
+                if other is not None:
+                    comparison = other[:, channel] if scope == "channel" else other[row, channel]
+                    other_lower, other_upper = _finite_range(comparison)
+                    lower, upper = min(lower, other_lower), max(upper, other_upper)
+            norm_type = Normalize if legacy else _ClippedNormalize
+            norms.append(norm_type(vmin=lower, vmax=upper, clip=True))
+        rows.append(norms)
+    return rows
+
+
+def unpaired_norms(
+    images: np.ndarray, scaling: str, limits: Optional[DisplayLimits], name: str
+) -> List[List[Normalize]]:
+    """Input/raw panels have their own independent, channel-shared or fixed scale."""
+    if scaling not in ("independent", "channel", "fixed"):
+        raise ValueError(f"{name}_scaling must be 'independent', 'channel', or 'fixed'.")
+    if (scaling == "fixed") != (limits is not None):
+        raise ValueError(f"{name}_limits must be provided if and only if {name}_scaling='fixed'.")
+    return image_norms(
+        images, scope="channel" if scaling == "channel" else "image",
+        limits=limits, legacy=scaling == "independent",
+    )
diff --git a/src/virtual_stain_flow/evaluation/visualization.py b/src/virtual_stain_flow/evaluation/visualization.py
index 2628d07..491f9e6 100644
--- a/src/virtual_stain_flow/evaluation/visualization.py
+++ b/src/virtual_stain_flow/evaluation/visualization.py
@@ -12,6 +12,7 @@
 import torch
 import matplotlib.pyplot as plt
 from matplotlib.patches import Rectangle
+from torch.utils.data import TensorDataset
 
 from ..datasets.base_dataset import BaseImageDataset
 from ..datasets.crop_dataset import CropImageDataset
@@ -19,6 +20,7 @@
 from .evaluation_utils import evaluate_per_image_metric
 from .predict_utils import predict_image
 from .visualization_utils import extract_samples_from_dataset
+from .display_normalization import DisplayLimits, image_norms, unpaired_norms
 
 
 def _select_channels(
@@ -88,6 +90,13 @@ def plot_predictions_grid(
     input_channel_indices: Optional[List[int]] = None,
     target_channel_indices: Optional[List[int]] = None,
     prediction_channel_indices: Optional[List[int]] = None,
+    target_scaling: str = "target",
+    target_scale_scope: str = "image",
+    target_limits: Optional[DisplayLimits] = None,
+    input_scaling: str = "independent",
+    input_limits: Optional[DisplayLimits] = None,
+    raw_scaling: str = "independent",
+    raw_limits: Optional[DisplayLimits] = None,
 ) -> plt.Figure:
     """
     Core visualization function for visualizing grid of input/target and predictions. 
@@ -127,6 +136,29 @@ def plot_predictions_grid(
     :param input_channel_indices: Optional list of channel indices to display for inputs.
     :param target_channel_indices: Optional list of channel indices to display for targets.
     :param prediction_channel_indices: Optional list of channel indices to display for predictions.
+    :param target_scaling: 'target' (default) uses only the target's range for
+        each target/prediction pair; 'joint' uses their combined range;
+        'independent' scales them separately (legacy); 'fixed' uses target_limits.
+        Values outside shared limits saturate at the colormap endpoints (black
+        and white for gray). Only display mapping changes, never data or metrics.
+    :param target_scale_scope: 'image' computes ranges per row/channel;
+        'channel' shares each displayed channel's range across all rows.
+        Independent mode still keeps target and prediction ranges separate.
+        Fixed mode is always shared across rows.
+    :param target_limits: Required only with target_scaling='fixed': a finite
+        (min, max) pair, or one pair per selected target channel in display order.
+        The corresponding prediction uses the same pair. Requires min < max.
+    :param input_scaling: 'independent' (default), 'channel' (shared across rows
+        per channel), or 'fixed'. Input scaling never shares target/raw limits.
+    :param input_limits: Fixed limits for inputs, in selected channel order.
+    :param raw_scaling: Like input_scaling, independently controls raw panels.
+    :param raw_limits: Fixed limits for raw images, in selected channel order.
+
+    A constant shared reference maps equal values to the colormap midpoint,
+    lower predictions to its low endpoint, and higher predictions to its high
+    endpoint. Independent image scaling retains Matplotlib's legacy constant
+    image behavior. Automatic ranges ignore nonfinite values; a reference with
+    no finite values raises ValueError.
     """
     if inputs.ndim != 4:
         raise ValueError(f"Inputs must have shape (N, C, H, W), received {inputs.shape}.")
@@ -188,6 +220,25 @@ def plot_predictions_grid(
             "Target and prediction channel counts must match for paired display."
         )
 
+    if target_scaling not in ("target", "joint", "independent", "fixed"):
+        raise ValueError("target_scaling must be 'target', 'joint', 'independent', or 'fixed'.")
+    if (target_scaling == "fixed") != (target_limits is not None):
+        raise ValueError("target_limits must be provided if and only if target_scaling='fixed'.")
+    target_norms = image_norms(
+        targets, scope=target_scale_scope, limits=target_limits,
+        other=predictions if target_scaling == "joint" else None,
+        legacy=target_scaling == "independent",
+    )
+    prediction_norms = (
+        image_norms(predictions, scope=target_scale_scope, legacy=True)
+        if has_predictions and target_scaling == "independent" else target_norms
+    )
+    input_norms = unpaired_norms(inputs, input_scaling, input_limits, "input")
+    raw_norms = (
+        unpaired_norms(raw_images, raw_scaling, raw_limits, "raw")
+        if has_raw_images else [[] for _ in range(num_samples)]
+    )
+
     raw_titles = _build_titles("Raw Input", raw_indices, raw_channel_names)
     input_titles = _build_titles("Input", input_indices, input_channel_names)
     target_titles = _build_titles("Target", target_indices, target_channel_names)
@@ -212,11 +263,7 @@ def plot_predictions_grid(
     # Create figure
     fig_width = panel_width * num_cols
     fig_height = panel_width * num_samples
-    fig, axes = plt.subplots(num_samples, num_cols, figsize=(fig_width, fig_height))
-
-    # Handle single-row case where axes is 1D
-    if num_samples == 1:
-        axes = axes.reshape(1, -1)
+    fig, axes = plt.subplots(num_samples, num_cols, figsize=(fig_width, fig_height), squeeze=False)
 
     for row_idx in range(num_samples):
         raw_row = list(raw_images[row_idx]) if has_raw_images else []
@@ -231,6 +278,12 @@ def plot_predictions_grid(
         ] if has_predictions else target_row
 
         img_set = raw_row + input_row + target_pred_row
+        paired_norms = [
+            norm
+            for target_norm, prediction_norm in zip(target_norms[row_idx], prediction_norms[row_idx])
+            for norm in (target_norm, prediction_norm)
+        ] if has_predictions else target_norms[row_idx]
+        norms = raw_norms[row_idx] + input_norms[row_idx] + paired_norms
 
         if len(img_set) != num_cols:
             raise ValueError(
@@ -242,7 +295,7 @@ def plot_predictions_grid(
 
             # Squeeze to 2D for display (handles (1, H, W) or (H, W))
             img_2d = np.squeeze(img)
-            ax.imshow(img_2d, cmap=cmap)
+            ax.imshow(img_2d, cmap=cmap, norm=norms[col_idx], interpolation="nearest")
 
             # Column title only on first row
             if row_idx == 0:
@@ -271,12 +324,8 @@ def plot_predictions_grid(
                 # Sample a small region (e.g., top-left 10% of image)
                 sample_size = max(1, int(min(img_2d.shape) * 0.1))
                 corner_region = img_2d[:sample_size, :sample_size]
-                # Normalize to 0-1 range for brightness check
-                img_min, img_max = img_2d.min(), img_2d.max()
-                if img_max > img_min:
-                    normalized_brightness = (corner_region.mean() - img_min) / (img_max - img_min)
-                else:
-                    normalized_brightness = 0.5
+                # Use the actual display mapping, including channel-shared limits.
+                normalized_brightness = norms[col_idx](corner_region).mean()
                 text_color = "black" if normalized_brightness > 0.5 else "white"
                 ax.text(
                     0.02, 0.98,  # Top-left corner in axes coordinates
@@ -332,7 +381,7 @@ def plot_dataset_grid(
     :param indices: List of dataset indices to display.
     :param save_path: Optional path to save the figure.
     :param kwargs: Additional arguments passed to `plot_predictions_grid`.
-        Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
+        Includes all display scaling/limits options, row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
         raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices,
         raw_channel_names, input_channel_names, target_channel_names, prediction_channel_names.
     """
@@ -371,9 +420,9 @@ def plot_predictions_grid_from_model(
     Plot predictions grid by running inference on a model.
 
     Performs the following steps:
-    1. Run inference on the specified dataset indices.
-    2. Compute per-image metrics.
-    3. Extract samples and plot using `plot_predictions_grid`.
+    1. Snapshot transformed samples and raw/crop metadata in one dataset pass.
+    2. Run inference and compute metrics on that same snapshot.
+    3. Plot the snapshot and predictions using `plot_predictions_grid`.
 
     :param model: PyTorch model for inference.
     :param dataset: BaseImageDataset, CropImageDataset, or BaseWrapperDataset to visualize.
@@ -385,32 +434,23 @@ def plot_predictions_grid_from_model(
     :param device: Device for inference ("cpu" or "cuda").
     :param save_path: Optional path to save the figure.
     :param kwargs: Additional arguments passed to `plot_predictions_grid`.
-        Supported: row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
+        Includes all display scaling/limits options, row_label_prefix, cmap, panel_width, show_plot, wspace, hspace,
         raw_channel_indices, input_channel_indices, target_channel_indices, prediction_channel_indices,
         raw_channel_names, input_channel_names, target_channel_names, prediction_channel_names.
     """
-    # Step 1: Run inference
-    targets_tensor, predictions_tensor, inputs_tensor = predict_image(
-        dataset, model, indices=indices, device=device
-    )
+    # Capture metadata immediately after each access, not in a second traversal
+    # that could resample stochastic transforms or replace mutable crop state.
+    (
+        inputs, targets, raw_images, patch_coords, input_channel_names, target_channel_names
+    ) = extract_samples_from_dataset(dataset, indices)
+    snapshot = TensorDataset(torch.from_numpy(inputs), torch.from_numpy(targets))
+    targets_tensor, predictions_tensor, _ = predict_image(snapshot, model, device=device)
 
     # Step 2: Compute metrics (if any)
     metrics_df = None
     if metrics:
         metrics_df = evaluate_per_image_metric(predictions_tensor, targets_tensor, metrics)
 
-    # Step 3: Re-access the dataset for CropImageDataset raw images and crop metadata.
-    (
-        _, _, raw_images, patch_coords, input_channel_names, target_channel_names
-    ) = extract_samples_from_dataset(dataset, indices)
-    if isinstance(inputs_tensor, list):
-        raise ValueError(
-            "Visualization requires a single batched input tensor with shape (N, C, H, W); "
-            "multi-input sequences are not supported."
-        )
-
-    inputs = inputs_tensor.detach().cpu().numpy()
-    targets = targets_tensor.detach().cpu().numpy()
     predictions = predictions_tensor.detach().cpu().numpy()
 
     # Step 4: Plot
diff --git a/tests/evaluation/test_display_normalization.py b/tests/evaluation/test_display_normalization.py
new file mode 100644
index 0000000..959d6bb
--- /dev/null
+++ b/tests/evaluation/test_display_normalization.py
@@ -0,0 +1,182 @@
+"""Regression tests for the actual Matplotlib image mappings, not just titles."""
+
+import matplotlib
+import numpy as np
+import pytest
+
+matplotlib.use("Agg")
+
+from virtual_stain_flow.evaluation.visualization import plot_predictions_grid
+
+
+def _images(figure):
+    return [axis.images[0] for axis in figure.axes]
+
+
+def _grid(**kwargs):
+    targets = np.array([0.2, 0.3, 0.5, 0.6]).reshape(1, 1, 2, 2)
+    predictions = np.array([-0.1, 0.3, 0.5, 1.2]).reshape(1, 1, 2, 2)
+    return plot_predictions_grid(
+        targets * 10, targets, predictions, show_plot=False, **kwargs
+    )
+
+
+def test_default_is_target_only_and_clips_prediction_display_not_data():
+    figure = _grid()
+    input_image, target, prediction = _images(figure)
+    assert target.norm is prediction.norm
+    assert target.get_clim() == (0.2, 0.6)
+    assert input_image.get_clim() == (2.0, 6.0)
+    assert prediction.norm(-0.1) == 0
+    assert prediction.norm(1.2) == 1
+    assert prediction.cmap(prediction.norm(-0.1))[:3] == (0, 0, 0)
+    assert prediction.cmap(prediction.norm(1.2))[:3] == (1, 1, 1)
+    assert prediction.get_array().min() == -0.1
+    assert prediction.get_array().max() == 1.2
+    # Force rendering too, including the shared norm's interaction with imshow.
+    figure.canvas.draw()
+    assert target.get_clim() == prediction.get_clim() == (0.2, 0.6)
+
+
+@pytest.mark.parametrize(
+    "mode,limits,target_range,prediction_range",
+    [
+        ("target", None, (0.2, 0.6), (0.2, 0.6)),
+        ("joint", None, (-0.1, 1.2), (-0.1, 1.2)),
+        ("independent", None, (0.2, 0.6), (-0.1, 1.2)),
+        ("fixed", (0, 1), (0, 1), (0, 1)),
+    ],
+)
+def test_selectable_target_modes(mode, limits, target_range, prediction_range):
+    _, target, prediction = _images(_grid(target_scaling=mode, target_limits=limits))
+    assert target.get_clim() == target_range
+    assert prediction.get_clim() == prediction_range
+
+
+@pytest.mark.parametrize("mode", ["target", "joint", "independent"])
+def test_channel_scope_shares_rows_but_not_channels(mode):
+    targets = np.array([
+        [[[1, 2], [3, 4]], [[10, 20], [30, 40]]],
+        [[[2, 3], [4, 5]], [[20, 30], [40, 50]]],
+    ], dtype=float)
+    predictions = targets * 2
+    figure = plot_predictions_grid(
+        targets[:, :1], targets, predictions, target_scaling=mode,
+        target_scale_scope="channel", show_plot=False,
+    )
+    images = _images(figure)
+    target_ranges = [(1, 5), (10, 50)]
+    prediction_ranges = [(2, 10), (20, 100)]
+    if mode == "joint":
+        target_ranges = [(1, 10), (10, 100)]
+    if mode != "independent":
+        prediction_ranges = target_ranges
+    for channel in range(2):
+        target_column = 1 + 2 * channel
+        assert images[target_column].norm is images[target_column + 5].norm
+        assert images[target_column + 1].norm is images[target_column + 6].norm
+        assert images[target_column].get_clim() == target_ranges[channel]
+        assert images[target_column + 1].get_clim() == prediction_ranges[channel]
+
+
+def test_image_scope_keeps_rows_independent_by_default():
+    targets = np.arange(8, dtype=float).reshape(2, 1, 2, 2)
+    images = _images(plot_predictions_grid(targets, targets, targets * 2, show_plot=False))
+    assert images[1].get_clim() == images[2].get_clim() == (0, 3)
+    assert images[4].get_clim() == images[5].get_clim() == (4, 7)
+
+
+def test_fixed_limits_follow_selected_channel_order():
+    images = np.arange(24, dtype=float).reshape(2, 3, 2, 2)
+    artists = _images(plot_predictions_grid(
+        images, images, images, raw_images=images,
+        input_channel_indices=[2, 0], raw_channel_indices=[1],
+        target_channel_indices=[2, 0], prediction_channel_indices=[0, 1],
+        target_scaling="fixed", target_limits=[(0, 100), (-1, 1)],
+        input_scaling="fixed", input_limits=[(0, 30), (0, 10)],
+        raw_scaling="fixed", raw_limits=(0, 65535), show_plot=False,
+    ))
+    expected = [(0, 65535), (0, 30), (0, 10), (0, 100), (0, 100), (-1, 1), (-1, 1)]
+    assert [image.get_clim() for image in artists] == expected * 2
+    np.testing.assert_array_equal(artists[3].get_array(), images[0, 2])
+    np.testing.assert_array_equal(artists[4].get_array(), images[0, 0])
+
+
+@pytest.mark.parametrize("scaling", ["independent", "channel"])
+def test_input_and_raw_scaling_are_separate_from_each_other_and_targets(scaling):
+    images = np.arange(16, dtype=float).reshape(2, 2, 2, 2)
+    artists = _images(plot_predictions_grid(
+        images, images[:, :1] / 100, raw_images=images * 100,
+        input_scaling=scaling, raw_scaling=scaling, show_plot=False,
+    ))
+    if scaling == "channel":
+        ranges = [(0, 1100), (400, 1500), (0, 11), (4, 15)]
+        for start in (0, 5):
+            assert [image.get_clim() for image in artists[start:start + 4]] == ranges
+    else:
+        assert artists[0].get_clim() == (0, 300)
+        assert artists[5].get_clim() == (800, 1100)
+        assert artists[2].get_clim() == (0, 3)
+        assert artists[7].get_clim() == (8, 11)
+
+
+def test_constant_target_still_exposes_under_and_over_prediction():
+    target = np.full((1, 1, 2, 2), 0.5)
+    prediction = np.array([0.4, 0.5, 0.6, np.nan]).reshape(1, 1, 2, 2)
+    figure = plot_predictions_grid(target, target, prediction, show_plot=False)
+    _, target_artist, pred_artist = _images(figure)
+    assert pred_artist.norm is target_artist.norm
+    assert pred_artist.get_clim() == (0.5, 0.5)
+    np.testing.assert_array_equal(pred_artist.norm([0.4, 0.5, 0.6]), [0, 0.5, 1])
+    assert np.ma.is_masked(pred_artist.norm(np.nan))
+    figure.canvas.draw()
+    np.testing.assert_array_equal(pred_artist.norm([0.4, 0.5, 0.6]), [0, 0.5, 1])
+    rgba, *_ = pred_artist.make_image(figure.canvas.get_renderer(), unsampled=True)
+    np.testing.assert_array_equal(rgba[0, 0], [0, 0, 0, 255])
+    np.testing.assert_allclose(rgba[0, 1, :3], [128, 128, 128], atol=1)
+    np.testing.assert_array_equal(rgba[1, 0], [255, 255, 255, 255])
+    assert rgba[1, 1, 3] == 0
+
+
+def test_legacy_constant_images_keep_matplotlib_behavior():
+    images = np.ones((1, 1, 2, 2))
+    artists = _images(plot_predictions_grid(
+        images, images, images, target_scaling="independent", show_plot=False,
+    ))
+    assert artists[1].norm(1.0) == 0
+
+
+def test_auto_limits_ignore_nonfinite_values():
+    targets = np.array([np.nan, np.inf, 0.2, 0.8]).reshape(1, 1, 2, 2)
+    artists = _images(plot_predictions_grid(np.zeros_like(targets), targets, show_plot=False))
+    assert artists[1].get_clim() == (0.2, 0.8)
+    with pytest.raises(ValueError, match="no finite values"):
+        plot_predictions_grid(np.zeros_like(targets), targets * np.nan, show_plot=False)
+
+
+@pytest.mark.parametrize("kwargs", [
+    {"target_scaling": "invalid"},
+    {"target_scale_scope": "invalid"},
+    {"target_scaling": "fixed"},
+    {"target_limits": (0, 1)},
+    {"target_scaling": "fixed", "target_limits": (1, 0)},
+    {"target_scaling": "fixed", "target_limits": (0, 0)},
+    {"target_scaling": "fixed", "target_limits": (0, np.inf)},
+    {"target_scaling": "fixed", "target_limits": [(0, 1), (0, 2)]},
+    {"input_scaling": "invalid"},
+    {"input_scaling": "fixed"},
+    {"input_limits": (0, 1)},
+    {"raw_scaling": "invalid", "raw_images": np.ones((1, 1, 2, 2))},
+])
+def test_invalid_scaling_settings_raise(kwargs):
+    with pytest.raises(ValueError):
+        _grid(**kwargs)
+
+
+def test_without_predictions_joint_uses_target_range_and_does_not_mutate_arrays():
+    inputs = np.arange(4, dtype=float).reshape(1, 1, 2, 2)
+    targets = inputs / 10
+    original = targets.copy()
+    artists = _images(plot_predictions_grid(inputs, targets, target_scaling="joint", show_plot=False))
+    assert artists[1].get_clim() == (0, 0.3)
+    np.testing.assert_array_equal(targets, original)

From 8952cfc797374c8bbc7e56391fb34ee5bd41732c Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 6 Oct 2026 11:43:26 -0600
Subject: [PATCH 02/11] Fix test relying on renderer backed canvas which may
 not be present in actions env

---
 tests/evaluation/test_display_normalization.py | 14 +++++++-------
 1 file changed, 7 insertions(+), 7 deletions(-)

diff --git a/tests/evaluation/test_display_normalization.py b/tests/evaluation/test_display_normalization.py
index 959d6bb..c391c7c 100644
--- a/tests/evaluation/test_display_normalization.py
+++ b/tests/evaluation/test_display_normalization.py
@@ -129,13 +129,13 @@ def test_constant_target_still_exposes_under_and_over_prediction():
     assert pred_artist.get_clim() == (0.5, 0.5)
     np.testing.assert_array_equal(pred_artist.norm([0.4, 0.5, 0.6]), [0, 0.5, 1])
     assert np.ma.is_masked(pred_artist.norm(np.nan))
-    figure.canvas.draw()
-    np.testing.assert_array_equal(pred_artist.norm([0.4, 0.5, 0.6]), [0, 0.5, 1])
-    rgba, *_ = pred_artist.make_image(figure.canvas.get_renderer(), unsampled=True)
-    np.testing.assert_array_equal(rgba[0, 0], [0, 0, 0, 255])
-    np.testing.assert_allclose(rgba[0, 1, :3], [128, 128, 128], atol=1)
-    np.testing.assert_array_equal(rgba[1, 0], [255, 255, 255, 255])
-    assert rgba[1, 1, 3] == 0
+    mapped = pred_artist.cmap(pred_artist.norm(prediction[0, 0]))
+    np.testing.assert_allclose(
+        mapped,
+        [[[0, 0, 0, 1], [0.5, 0.5, 0.5, 1]],
+         [[1, 1, 1, 1], [0, 0, 0, 0]]],
+        atol=1 / 255,
+    )
 
 
 def test_legacy_constant_images_keep_matplotlib_behavior():

From b71a45fb259a570a19836dbae28cb93183099a78 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 6 Oct 2026 11:47:54 -0600
Subject: [PATCH 03/11] Tighten visualization backend and add tests for copying
 behavior

---
 .../evaluation/visualization_utils.py         |  8 ++-
 tests/evaluation/test_visualization_utils.py  | 55 +++++++++++++++++++
 2 files changed, 61 insertions(+), 2 deletions(-)

diff --git a/src/virtual_stain_flow/evaluation/visualization_utils.py b/src/virtual_stain_flow/evaluation/visualization_utils.py
index 8083cc8..c500cb6 100644
--- a/src/virtual_stain_flow/evaluation/visualization_utils.py
+++ b/src/virtual_stain_flow/evaluation/visualization_utils.py
@@ -33,7 +33,8 @@ def _to_numpy_image(value: Union[np.ndarray, torch.Tensor], name: str) -> np.nda
             f"{name} must have shape (C, H, W), received shape {value.shape}."
         )
 
-    return value
+    # Freeze reused NumPy/tensor buffers before accessing the next sample.
+    return value.copy()
 
 
 def _stack_images(images: List[np.ndarray], name: str) -> np.ndarray:
@@ -73,7 +74,10 @@ def extract_samples_from_dataset(
     Primary function of this abstraction is to provide a consistent data
         access interface between Dataset objects and plotting functions by
         extracting input/target with __get_item__ and also accessing raw
-        images and crop annotations when available. 
+        images and crop annotations when available. Each requested sample is
+        accessed once and copied immediately, including its raw image metadata.
+        Wrappers around crop datasets must keep the underlying crop metadata
+        synchronized with the sample they return.
 
     :param dataset: A BaseImageDataset, CropImageDataset, or BaseWrapperDataset.
     :param indices: Dataset indices to extract, in the displayed order.
diff --git a/tests/evaluation/test_visualization_utils.py b/tests/evaluation/test_visualization_utils.py
index f842151..bf2b3eb 100644
--- a/tests/evaluation/test_visualization_utils.py
+++ b/tests/evaluation/test_visualization_utils.py
@@ -43,3 +43,58 @@ def __getitem__(self, index):
 
     with pytest.raises(ValueError, match="single channel-first image"):
         extract_samples_from_dataset(MultiInputDataset(basic_dataset), [0])
+
+
+def test_snapshot_copies_reused_buffers(basic_dataset):
+    import torch
+
+    class ReusingWrapper(BaseWrapperDataset):
+        def __init__(self, dataset):
+            super().__init__(dataset)
+            self.buffer = torch.empty((1, 2, 2))
+
+        def __getitem__(self, index):
+            self.buffer.fill_(index)
+            return self.buffer, self.buffer
+
+    dataset = ReusingWrapper(basic_dataset)
+    inputs, targets, *_ = extract_samples_from_dataset(dataset, [2, 0, 2])
+    dataset.buffer.fill_(-1)
+    for images in (inputs, targets):
+        np.testing.assert_array_equal(images[:, 0, 0, 0], [2, 0, 2])
+
+
+def test_snapshot_copies_reused_raw_buffers(crop_dataset, monkeypatch):
+    dataset_type = type(crop_dataset)
+    original_getitem = dataset_type.__getitem__
+    raw_buffer = np.empty((2, 10, 10))
+
+    def getitem(self, index):
+        result = original_getitem(self, index)
+        raw_buffer.fill(index)
+        return result
+
+    monkeypatch.setattr(dataset_type, "__getitem__", getitem)
+    monkeypatch.setattr(dataset_type, "original_input_image", property(lambda self: raw_buffer))
+    _, _, raw, coords, *_ = extract_samples_from_dataset(crop_dataset, [1, 2])
+    raw_buffer.fill(-1)
+    np.testing.assert_array_equal(raw[:, 0, 0, 0], [1, 2])
+    assert coords == [(5, 5, 4, 4), (0, 0, 4, 4)]
+
+
+@pytest.mark.parametrize("fixture_name", ["basic_dataset", "crop_dataset"])
+def test_snapshot_matches_dataloader_with_separate_transforms(request, fixture_name):
+    from torch.utils.data import DataLoader, Subset
+    from virtual_stain_flow.transforms.gamma import ContinuousGammaTransform
+    from virtual_stain_flow.transforms.normalizations import MaxScaleNormalize
+
+    dataset = request.getfixturevalue(fixture_name)
+    dataset.transforms = [MaxScaleNormalize(10)]
+    dataset.input_transforms = [MaxScaleNormalize(2)]
+    dataset.target_transforms = [ContinuousGammaTransform(0.3)]
+    inputs, targets, *_ = extract_samples_from_dataset(dataset, [2, 0])
+    loader_inputs, loader_targets = next(iter(DataLoader(Subset(dataset, [2, 0]), batch_size=2)))
+    np.testing.assert_array_equal(inputs, loader_inputs.numpy())
+    np.testing.assert_array_equal(targets, loader_targets.numpy())
+    np.testing.assert_allclose(inputs[1, 0], 1 / 10 / 2)
+    np.testing.assert_allclose(targets[1, 0], (1 / 10) ** 0.3)

From 7f4c3a619517f9d10f71dfaf423d0c7b52d9446c Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 6 Oct 2026 11:48:46 -0600
Subject: [PATCH 04/11] Add tests for new scaling options in visualization
 behavior.

---
 tests/evaluation/test_visualization.py | 99 ++++++++++++++++++++++++++
 1 file changed, 99 insertions(+)

diff --git a/tests/evaluation/test_visualization.py b/tests/evaluation/test_visualization.py
index 260f61d..193f387 100644
--- a/tests/evaluation/test_visualization.py
+++ b/tests/evaluation/test_visualization.py
@@ -79,3 +79,102 @@ def test_plot_predictions_grid_from_model_uses_tensor_batches(basic_dataset):
     )
 
     assert len(figure.axes) == 8
+
+
+@pytest.mark.parametrize("scaling,limits", [("target", None), ("fixed", (0, 0.01)), ("independent", None)])
+def test_model_plot_uses_one_transformed_snapshot_for_metrics_and_metadata(crop_dataset, scaling, limits):
+    from virtual_stain_flow.datasets.base_wrapper_dataset import BaseWrapperDataset
+    from virtual_stain_flow.transforms.gamma import ContinuousGammaTransform
+    from virtual_stain_flow.transforms.normalizations import MaxScaleNormalize
+
+    crop_dataset.input_transforms = [MaxScaleNormalize(10)]
+    crop_dataset.target_transforms = [MaxScaleNormalize(10), ContinuousGammaTransform(0.3)]
+
+    class StatefulWrapper(BaseWrapperDataset):
+        def __init__(self, dataset):
+            super().__init__(dataset)
+            self.calls = []
+
+        def __getitem__(self, index):
+            self.calls.append(index)
+            inputs, targets = self._dataset[index]
+            # Simulate augmentation that changes with every read.
+            return inputs + len(self.calls), targets
+
+    class RecordingModel(torch.nn.Module):
+        def __init__(self):
+            super().__init__()
+            self.inputs = []
+
+        def forward(self, inputs):
+            self.inputs.append(inputs.clone())
+            return inputs[:, :1] + 0.25
+
+    class RecordingL1(torch.nn.L1Loss):
+        def __init__(self):
+            super().__init__()
+            self.values = []
+
+        def forward(self, first, second):
+            result = super().forward(first, second)
+            self.values.append(result.item())
+            return result
+
+    dataset = StatefulWrapper(crop_dataset)
+    model, metric = RecordingModel(), RecordingL1()
+    figure = plot_predictions_grid_from_model(
+        model, dataset, indices=[3, 0], metrics=[metric], device="cpu",
+        target_scaling=scaling, target_limits=limits, show_plot=False,
+    )
+    assert dataset.calls == [3, 0]
+    for row, (raw_value, crop_xy) in enumerate([(3, (5, 5)), (1, (0, 0))]):
+        axes = figure.axes[row * 6:(row + 1) * 6]
+        displayed_input = axes[2].images[0].get_array()
+        displayed_target = axes[4].images[0].get_array()
+        displayed_prediction = axes[5].images[0].get_array()
+        np.testing.assert_allclose(displayed_input, model.inputs[row][0, 0].numpy())
+        np.testing.assert_allclose(displayed_input, raw_value / 10 + row + 1)
+        np.testing.assert_allclose(displayed_target, ((2 if row == 0 else 1) / 10) ** 0.3)
+        np.testing.assert_allclose(displayed_prediction, displayed_input + 0.25)
+        assert metric.values[row] == pytest.approx(np.abs(displayed_prediction - displayed_target).mean())
+        assert f"RecordingL1: {metric.values[row]:.3f}" in axes[5].get_title()
+        np.testing.assert_array_equal(axes[0].images[0].get_array(), raw_value)
+        assert axes[0].patches[0].get_xy() == crop_xy
+
+
+def test_dataset_wrapper_forwards_display_options(basic_dataset):
+    figure = plot_dataset_grid(
+        basic_dataset, [0, 1], target_scaling="fixed", target_limits=(0, 10),
+        input_scaling="channel", show_plot=False,
+    )
+    assert figure.axes[2].images[0].get_clim() == (0, 10)
+    assert figure.axes[5].images[0].get_clim() == (0, 10)
+    assert figure.axes[0].images[0].get_clim() == figure.axes[3].images[0].get_clim()
+
+
+def test_plot_callback_forwards_scaling_options(basic_dataset, tmp_path, monkeypatch):
+    from virtual_stain_flow.vsf_logging.callbacks.PlotCallback import PlotPredictionCallback
+
+    callback = PlotPredictionCallback(
+        "test", basic_dataset, save_path=tmp_path, indices=[0],
+        target_scaling="fixed", target_limits=(0, 1), input_scaling="channel",
+        show_plot=False, panel_width=0.5,
+    )
+    monkeypatch.setattr(callback, "get_model", lambda **kwargs: torch.nn.Conv2d(2, 1, 1))
+    monkeypatch.setattr(callback, "get_epoch", lambda: 0)
+    assert callback._plot().is_file()
+
+
+def test_inplace_model_does_not_change_displayed_snapshot(basic_dataset):
+    class InplaceModel(torch.nn.Module):
+        def forward(self, inputs):
+            inputs.add_(100)
+            return inputs[:, :1]
+
+    figure = plot_predictions_grid_from_model(
+        InplaceModel(), basic_dataset, [0], [], device="cpu", show_plot=False,
+    )
+    np.testing.assert_array_equal(figure.axes[0].images[0].get_array(), 1)
+    np.testing.assert_array_equal(figure.axes[1].images[0].get_array(), 2)
+    np.testing.assert_array_equal(figure.axes[2].images[0].get_array(), 1)
+    np.testing.assert_array_equal(figure.axes[3].images[0].get_array(), 101)

From 45b8166cf763a4676904cbdfe874a41226045a52 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 6 Oct 2026 13:44:05 -0600
Subject: [PATCH 05/11] Additionally tighten tensor creation to avoid modifying
 original inputs and targets and ensuring snapshot itself isn't modified
 unexpectedly

---
 src/virtual_stain_flow/evaluation/visualization.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/src/virtual_stain_flow/evaluation/visualization.py b/src/virtual_stain_flow/evaluation/visualization.py
index 491f9e6..d794a74 100644
--- a/src/virtual_stain_flow/evaluation/visualization.py
+++ b/src/virtual_stain_flow/evaluation/visualization.py
@@ -443,7 +443,7 @@ def plot_predictions_grid_from_model(
     (
         inputs, targets, raw_images, patch_coords, input_channel_names, target_channel_names
     ) = extract_samples_from_dataset(dataset, indices)
-    snapshot = TensorDataset(torch.from_numpy(inputs), torch.from_numpy(targets))
+    snapshot = TensorDataset(torch.from_numpy(inputs.copy()), torch.from_numpy(targets.copy()))
     targets_tensor, predictions_tensor, _ = predict_image(snapshot, model, device=device)
 
     # Step 2: Compute metrics (if any)

From f540bd58a0f5f7b4ce801b8c292496a06f226907 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 13 Oct 2026 15:40:26 -0600
Subject: [PATCH 06/11] Refactor type hints in display normalization and
 visualization modules to use Literal for better clarity and type safety

---
 .../evaluation/display_normalization.py                |  6 +++---
 src/virtual_stain_flow/evaluation/visualization.py     | 10 +++++-----
 2 files changed, 8 insertions(+), 8 deletions(-)

diff --git a/src/virtual_stain_flow/evaluation/display_normalization.py b/src/virtual_stain_flow/evaluation/display_normalization.py
index 0d3a0af..ab038bc 100644
--- a/src/virtual_stain_flow/evaluation/display_normalization.py
+++ b/src/virtual_stain_flow/evaluation/display_normalization.py
@@ -1,6 +1,6 @@
 """Display-only intensity scaling; never modifies image or metric data."""
 
-from typing import List, Optional, Sequence, Tuple, Union
+from typing import List, Optional, Sequence, Tuple, Union, Literal
 
 import numpy as np
 from matplotlib.colors import Normalize
@@ -52,7 +52,7 @@ def _finite_range(images: np.ndarray) -> Tuple[float, float]:
 def image_norms(
     images: np.ndarray,
     *,
-    scope: str = "image",
+    scope: Literal["image", "channel"] = "image",
     limits: Optional[DisplayLimits] = None,
     other: Optional[np.ndarray] = None,
     legacy: bool = False,
@@ -84,7 +84,7 @@ def image_norms(
 
 
 def unpaired_norms(
-    images: np.ndarray, scaling: str, limits: Optional[DisplayLimits], name: str
+    images: np.ndarray, scaling: Literal["independent", "channel", "fixed"], limits: Optional[DisplayLimits], name: str
 ) -> List[List[Normalize]]:
     """Input/raw panels have their own independent, channel-shared or fixed scale."""
     if scaling not in ("independent", "channel", "fixed"):
diff --git a/src/virtual_stain_flow/evaluation/visualization.py b/src/virtual_stain_flow/evaluation/visualization.py
index d794a74..080f921 100644
--- a/src/virtual_stain_flow/evaluation/visualization.py
+++ b/src/virtual_stain_flow/evaluation/visualization.py
@@ -5,7 +5,7 @@
 Supports both BaseImageDataset and CropImageDataset with optional metrics display.
 """
 
-from typing import List, Optional, Tuple, Union
+from typing import List, Literal, Optional, Tuple, Union
 
 import numpy as np
 import pandas as pd
@@ -90,12 +90,12 @@ def plot_predictions_grid(
     input_channel_indices: Optional[List[int]] = None,
     target_channel_indices: Optional[List[int]] = None,
     prediction_channel_indices: Optional[List[int]] = None,
-    target_scaling: str = "target",
-    target_scale_scope: str = "image",
+    target_scaling: Literal["target", "joint", "independent", "fixed"] = "target",
+    target_scale_scope: Literal["image", "channel"] = "image",
     target_limits: Optional[DisplayLimits] = None,
-    input_scaling: str = "independent",
+    input_scaling: Literal["independent", "channel", "fixed"] = "independent",
     input_limits: Optional[DisplayLimits] = None,
-    raw_scaling: str = "independent",
+    raw_scaling: Literal["independent", "channel", "fixed"] = "independent",
     raw_limits: Optional[DisplayLimits] = None,
 ) -> plt.Figure:
     """

From b3fafe0a3ab4512bc95c746cf838f5829f6129a9 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 13 Oct 2026 15:41:10 -0600
Subject: [PATCH 07/11] Specify matplotlib version range in dependencies for
 compatibility

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 628269a..807d08a 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -14,7 +14,7 @@ dependencies = [
   "numpy",
   "pandas>=2.1",
   "seaborn",
-  "matplotlib",
+  "matplotlib>=3.10,<3.12",
   "scikit-learn",
   "opencv-python",
   "typing-extensions",

From cf0f807f4b8819c4c77a2f70274a760ab46a3385 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 13 Oct 2026 16:02:44 -0600
Subject: [PATCH 08/11] Refactor image batch handling in
 extract_samples_from_dataset for dtype promotion and consistency checks

---
 .../evaluation/visualization_utils.py         | 61 +++++++++++++------
 tests/evaluation/test_visualization_utils.py  | 30 +++++++++
 2 files changed, 72 insertions(+), 19 deletions(-)

diff --git a/src/virtual_stain_flow/evaluation/visualization_utils.py b/src/virtual_stain_flow/evaluation/visualization_utils.py
index c500cb6..3cfa4be 100644
--- a/src/virtual_stain_flow/evaluation/visualization_utils.py
+++ b/src/virtual_stain_flow/evaluation/visualization_utils.py
@@ -33,19 +33,38 @@ def _to_numpy_image(value: Union[np.ndarray, torch.Tensor], name: str) -> np.nda
             f"{name} must have shape (C, H, W), received shape {value.shape}."
         )
 
-    # Freeze reused NumPy/tensor buffers before accessing the next sample.
-    return value.copy()
+    # The caller must snapshot this view before accessing the next sample.
+    return value
 
 
-def _stack_images(images: List[np.ndarray], name: str) -> np.ndarray:
-    """
-    Helper function to stack a list of images into a single numpy array 
-        with shape (N, C, H, W).
-    """
-    try:
-        return np.stack(images)
-    except ValueError as error:
-        raise ValueError(f"{name} must have a consistent (C, H, W) shape.") from error
+class _ImageBatch:
+    """Snapshot images directly into batch storage, promoting dtype if needed."""
+
+    def __init__(self, count: int, name: str):
+        self.count = count
+        self.name = name
+        self.batch: Optional[np.ndarray] = None
+        self._dtypes: List[np.dtype] = []
+        self._position = 0
+
+    def append(self, image: np.ndarray) -> None:
+        if self.batch is not None and image.shape != self.batch.shape[1:]:
+            raise ValueError(f"{self.name} must have a consistent (C, H, W) shape.")
+
+        # Use source dtypes rather than pairwise promotion to match np.stack.
+        if image.dtype not in self._dtypes:
+            self._dtypes.append(image.dtype)
+        dtype = np.result_type(*self._dtypes)
+        if self.batch is None:
+            self.batch = np.empty((self.count, *image.shape), dtype=dtype)
+        elif dtype != self.batch.dtype:
+            promoted = np.empty(self.batch.shape, dtype=dtype)
+            promoted[:self._position] = self.batch[:self._position]
+            self.batch = promoted
+
+        # Freeze reused NumPy/tensor buffers before accessing the next sample.
+        self.batch[self._position] = image
+        self._position += 1
 
 
 def _get_channel_names(
@@ -75,7 +94,9 @@ def extract_samples_from_dataset(
         access interface between Dataset objects and plotting functions by
         extracting input/target with __get_item__ and also accessing raw
         images and crop annotations when available. Each requested sample is
-        accessed once and copied immediately, including its raw image metadata.
+        accessed once and copied immediately into preallocated batch storage,
+        including its raw image metadata. Mixed dtypes may require reallocation
+        to preserve NumPy stacking promotion semantics.
         Wrappers around crop datasets must keep the underlying crop metadata
         synchronized with the sample they return.
 
@@ -110,9 +131,9 @@ def extract_samples_from_dataset(
             f"requested indices: {indices}"
         )
 
-    inputs: List[np.ndarray] = []
-    targets: List[np.ndarray] = []
-    raw_images: Optional[List[np.ndarray]] = [] if crop_dataset is not None else None
+    inputs = _ImageBatch(len(indices), "Dataset inputs")
+    targets = _ImageBatch(len(indices), "Dataset targets")
+    raw_images = _ImageBatch(len(indices), "Crop raw images") if crop_dataset is not None else None
     patch_coords: Optional[List[PatchCoords]] = [] if crop_dataset is not None else None
 
     for index in indices:
@@ -122,7 +143,7 @@ def extract_samples_from_dataset(
         inputs.append(_to_numpy_image(input_image, "Dataset input"))
         targets.append(_to_numpy_image(target_image, "Dataset target"))
 
-        if crop_dataset is not None:
+        if crop_dataset is not None and raw_images is not None and patch_coords is not None:
             raw_image = crop_dataset.original_input_image
             crop_info = crop_dataset.crop_info
             if raw_image is None or crop_info is None:
@@ -133,9 +154,11 @@ def extract_samples_from_dataset(
                 (crop_info.x, crop_info.y, crop_info.width, crop_info.height)
             )
 
-    input_batch = _stack_images(inputs, "Dataset inputs")
-    target_batch = _stack_images(targets, "Dataset targets")
-    raw_batch = _stack_images(raw_images, "Crop raw images") if raw_images is not None else None
+    input_batch = inputs.batch
+    target_batch = targets.batch
+    raw_batch = raw_images.batch if raw_images is not None else None
+    # Nonempty indices guarantee that both batches have been allocated.
+    assert input_batch is not None and target_batch is not None
 
     return (
         input_batch,
diff --git a/tests/evaluation/test_visualization_utils.py b/tests/evaluation/test_visualization_utils.py
index bf2b3eb..9d6d34d 100644
--- a/tests/evaluation/test_visualization_utils.py
+++ b/tests/evaluation/test_visualization_utils.py
@@ -62,6 +62,9 @@ def __getitem__(self, index):
     dataset.buffer.fill_(-1)
     for images in (inputs, targets):
         np.testing.assert_array_equal(images[:, 0, 0, 0], [2, 0, 2])
+    inputs.fill(10)
+    np.testing.assert_array_equal(targets[:, 0, 0, 0], [2, 0, 2])
+    assert np.all(dataset.buffer.numpy() == -1)
 
 
 def test_snapshot_copies_reused_raw_buffers(crop_dataset, monkeypatch):
@@ -80,6 +83,33 @@ def getitem(self, index):
     raw_buffer.fill(-1)
     np.testing.assert_array_equal(raw[:, 0, 0, 0], [1, 2])
     assert coords == [(5, 5, 4, 4), (0, 0, 4, 4)]
+    raw.fill(10)
+    assert np.all(raw_buffer == -1)
+
+
+@pytest.mark.parametrize("dtypes", [(np.uint8, np.int8, np.float16), (np.int16, np.float32, np.float64)])
+def test_snapshot_matches_stack_dtype_promotion(basic_dataset, dtypes):
+    images = [np.full((1, 2, 2), index + 0.5, dtype=dtype) for index, dtype in enumerate(dtypes)]
+
+    class MixedDtypeWrapper(BaseWrapperDataset):
+        def __getitem__(self, index):
+            return images[index], images[index]
+
+    inputs, targets, *_ = extract_samples_from_dataset(MixedDtypeWrapper(basic_dataset), [0, 1, 2])
+    expected = np.stack(images)
+    for batch in (inputs, targets):
+        assert batch.dtype == expected.dtype
+        np.testing.assert_array_equal(batch, expected)
+
+
+def test_snapshot_rejects_inconsistent_shapes(basic_dataset):
+    class VaryingShapeWrapper(BaseWrapperDataset):
+        def __getitem__(self, index):
+            image = np.zeros((1, index + 1, 2))
+            return image, image
+
+    with pytest.raises(ValueError, match=r"Dataset inputs must have a consistent \(C, H, W\) shape"):
+        extract_samples_from_dataset(VaryingShapeWrapper(basic_dataset), [0, 1])
 
 
 @pytest.mark.parametrize("fixture_name", ["basic_dataset", "crop_dataset"])

From 6759e0218244499d7acb92efb5315c37dabc52f6 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Tue, 13 Oct 2026 16:25:02 -0600
Subject: [PATCH 09/11] raise the project's declared Python minimum to 3.10

---
 pyproject.toml | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pyproject.toml b/pyproject.toml
index 807d08a..8b65cb7 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -8,7 +8,7 @@ build-backend = "setuptools.build_meta"
 name = "virtual_stain_flow"
 version = "0.4.10"
 description = "For developing virtual staining models"
-requires-python = ">=3.9"
+requires-python = ">=3.10"
 dependencies = [
   "pyarrow",
   "numpy",

From c60bf10bf486926e0d4a139d846030f2ff0064d4 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Wed, 14 Oct 2026 15:18:05 -0600
Subject: [PATCH 10/11] Enhance _finite_range to consider additional input for
 display limits and add tests for joint limits handling non-finite predictions

---
 .../evaluation/display_normalization.py       | 18 ++++++++++----
 .../evaluation/test_display_normalization.py  | 24 +++++++++++++++++++
 2 files changed, 37 insertions(+), 5 deletions(-)

diff --git a/src/virtual_stain_flow/evaluation/display_normalization.py b/src/virtual_stain_flow/evaluation/display_normalization.py
index ab038bc..b58a38a 100644
--- a/src/virtual_stain_flow/evaluation/display_normalization.py
+++ b/src/virtual_stain_flow/evaluation/display_normalization.py
@@ -42,11 +42,19 @@ def _fixed_limits(limits: Optional[DisplayLimits], channels: int, name: str) ->
     return values
 
 
-def _finite_range(images: np.ndarray) -> Tuple[float, float]:
+def _finite_range(images: np.ndarray, other: Optional[np.ndarray] = None) -> Tuple[float, float]:
     finite = images[np.isfinite(images)]
+    other_finite = other[np.isfinite(other)] if other is not None else np.array([])
+    if not finite.size:
+        finite = other_finite
+        other_finite = np.array([])
     if not finite.size:
         raise ValueError("Cannot determine display limits from an image/channel with no finite values.")
-    return float(finite.min()), float(finite.max())
+    lower, upper = float(finite.min()), float(finite.max())
+    if other_finite.size:
+        lower = min(lower, float(other_finite.min()))
+        upper = max(upper, float(other_finite.max()))
+    return lower, upper
 
 
 def image_norms(
@@ -72,11 +80,11 @@ def image_norms(
                 lower, upper = fixed[channel]
             else:
                 reference = images[:, channel] if scope == "channel" else images[row, channel]
-                lower, upper = _finite_range(reference)
                 if other is not None:
                     comparison = other[:, channel] if scope == "channel" else other[row, channel]
-                    other_lower, other_upper = _finite_range(comparison)
-                    lower, upper = min(lower, other_lower), max(upper, other_upper)
+                    lower, upper = _finite_range(reference, comparison)
+                else:
+                    lower, upper = _finite_range(reference)
             norm_type = Normalize if legacy else _ClippedNormalize
             norms.append(norm_type(vmin=lower, vmax=upper, clip=True))
         rows.append(norms)
diff --git a/tests/evaluation/test_display_normalization.py b/tests/evaluation/test_display_normalization.py
index c391c7c..7f32f5b 100644
--- a/tests/evaluation/test_display_normalization.py
+++ b/tests/evaluation/test_display_normalization.py
@@ -154,6 +154,30 @@ def test_auto_limits_ignore_nonfinite_values():
         plot_predictions_grid(np.zeros_like(targets), targets * np.nan, show_plot=False)
 
 
+def test_joint_limits_use_finite_target_when_prediction_channel_is_all_nonfinite():
+    targets = np.array([np.nan, np.inf, 0.2, 0.8]).reshape(1, 1, 2, 2)
+    predictions = np.full_like(targets, np.nan)
+    artists = _images(plot_predictions_grid(
+        np.zeros_like(targets), targets, predictions,
+        target_scaling="joint", show_plot=False,
+    ))
+    assert artists[1].get_clim() == artists[2].get_clim() == (0.2, 0.8)
+
+    predictions = np.array([-0.1, np.inf, 0.5, np.nan]).reshape(1, 1, 2, 2)
+    artists = _images(plot_predictions_grid(
+        np.zeros_like(targets), np.full_like(targets, np.nan), predictions,
+        target_scaling="joint", show_plot=False,
+    ))
+    assert artists[1].get_clim() == artists[2].get_clim() == (-0.1, 0.5)
+
+    predictions = np.full_like(targets, np.nan)
+    with pytest.raises(ValueError, match="no finite values"):
+        plot_predictions_grid(
+            np.zeros_like(targets), predictions, predictions,
+            target_scaling="joint", show_plot=False,
+        )
+
+
 @pytest.mark.parametrize("kwargs", [
     {"target_scaling": "invalid"},
     {"target_scale_scope": "invalid"},

From d448207735d5586e26abfc39c52aa1445e3ff108 Mon Sep 17 00:00:00 2001
From: Weishan Li <wli51@jh.edu>
Date: Thu, 15 Oct 2026 11:32:19 -0600
Subject: [PATCH 11/11] Update changelog for visualization normalization
 enhancements

---
 CHANGELOG.md   | 12 ++++++++++++
 pyproject.toml |  2 +-
 2 files changed, 13 insertions(+), 1 deletion(-)

diff --git a/CHANGELOG.md b/CHANGELOG.md
index 0101f68..375a0ea 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -5,6 +5,18 @@ All notable changes to this project will be documented in this file.
 The format is based on [Keep a Changelog](https://keepachangelog.com/en/1.1.0/),  
 and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0.html).
 
+## [0.4.11] - 2026-10-15
+
+### Refactored
+
+#### Visualization normalization/joint normalization modes and checks (`virtual_stain_flow/evluation/`)
+- Added configurable  image scaling for displaying the evaluation grids, including: 
+  1. target-referenced (prediction range is clamped in reference to target vmin and vmax to expose localized over-darkening and brightening)
+  2. joint (both target and reference are clamped by the pooled vim and vmax to reveal contrast differences)
+  3. independent (target and reference are independently optimized to display best effort contrast, for checking if texture and details and properly reconstructed correcting for illumination and contrast differences)
+  4. fixed ranges (manual free form)
+- Bumped minimal python version and pinned matplotlib version due to coupling with the normalization modes. 
+
 ## [0.4.10] - 2026-09-11
 
 ### Added
diff --git a/pyproject.toml b/pyproject.toml
index 8b65cb7..fe88f12 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -6,7 +6,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "virtual_stain_flow"
-version = "0.4.10"
+version = "0.4.11"
 description = "For developing virtual staining models"
 requires-python = ">=3.10"
 dependencies = [