From ff6fb4414c5966afa11f353dd397ac8b1b66fb61 Mon Sep 17 00:00:00 2001
From: Jammy2211 <JNightingale2211@gmail.com>
Date: Fri, 17 Jul 2026 08:16:25 +0100
Subject: [PATCH] feat: aggregate_images + aggregate_fits profiling stages

- mock templates carry the workflow payloads: panelled subplot_fit.png
  (4x3 grid) + multi-HDU fit.fits (4 EXTNAMEs)
- profile_aggregator.py: SubplotFit/FITSFit enums + png_make and
  fits_make pattern stages

These stages exposed the AggregateFITS fd leak fixed in Fit#1386.

Co-Authored-By: Claude Fable 5 <noreply@anthropic.com>
---
 scripts/profiling/aggregator/mock_results.py  | 33 ++++++++++--
 .../aggregator/profile_aggregator.py          | 51 +++++++++++++++++++
 2 files changed, 80 insertions(+), 4 deletions(-)

diff --git a/scripts/profiling/aggregator/mock_results.py b/scripts/profiling/aggregator/mock_results.py
index 6c75aa7..690f419 100644
--- a/scripts/profiling/aggregator/mock_results.py
+++ b/scripts/profiling/aggregator/mock_results.py
@@ -114,14 +114,39 @@ def write_template(
             paths.save_latent_samples(latent_samples=result.samples)
 
         if with_images:
-            image_path = paths.image_path
-            image = Image.new("RGB", (64, 64))
-            for name in ("subplot_fit", "corner_pdf", "search_internal"):
-                image.save(image_path / f"{name}.png")
+            _write_image_and_fits_payloads(paths)
 
         return Path(paths.output_path)
 
 
+def _write_image_and_fits_payloads(paths):
+    """
+    Write the payloads the workflow aggregators consume: a panelled subplot image
+    (AggregateImages extracts panels from a 4x3 grid, matching the real subplot_fit
+    convention) and a multi-HDU fits file (AggregateFITS extracts HDUs by EXTNAME).
+    """
+    import numpy as np
+    from astropy.io import fits
+
+    panel_w, panel_h = 120, 90
+    image = Image.new("RGB", (4 * panel_w, 3 * panel_h))
+    image.save(paths.image_path / "subplot_fit.png")
+    Image.new("RGB", (64, 64)).save(paths.image_path / "corner_pdf.png")
+
+    hdu_list = fits.HDUList([fits.PrimaryHDU()])
+    rng = np.random.default_rng(0)
+    for extname in (
+        "MODEL_IMAGE",
+        "RESIDUAL_MAP",
+        "NORMALIZED_RESIDUAL_MAP",
+        "CHI_SQUARED_MAP",
+    ):
+        hdu = fits.ImageHDU(data=rng.normal(size=(100, 100)).astype("float32"))
+        hdu.header["EXTNAME"] = extname
+        hdu_list.append(hdu)
+    hdu_list.writeto(paths._files_path / "fit.fits", overwrite=True)
+
+
 def _zip_directory(directory: Path):
     """
     Zip a result directory the way search output is zipped: contents at the archive
diff --git a/scripts/profiling/aggregator/profile_aggregator.py b/scripts/profiling/aggregator/profile_aggregator.py
index b657925..17a4565 100644
--- a/scripts/profiling/aggregator/profile_aggregator.py
+++ b/scripts/profiling/aggregator/profile_aggregator.py
@@ -39,10 +39,32 @@
 
 sys.path.insert(0, str(Path(__file__).parent))
 
+from enum import Enum
+
 import autofit as af
 from autofit.aggregator.aggregator import Aggregator
 from mock_results import generate_mock_results
 
+
+class SubplotFit(Enum):
+    """
+    Panel coordinates in the mock subplot_fit.png (4x3 grid) — the png_make pattern.
+    """
+
+    Data = (0, 0)
+    ModelData = (3, 0)
+    ResidualMap = (0, 2)
+    ChiSquaredMap = (2, 2)
+
+
+class FITSFit(Enum):
+    """
+    HDU EXTNAMEs in the mock fit.fits — the fits_make pattern.
+    """
+
+    ModelData = "MODEL_IMAGE"
+    ResidualMap = "RESIDUAL_MAP"
+
 RESULTS_PATH = Path("output") / "profiling_aggregator" / "results"
 
 # values("samples") holds every loaded Samples in memory (SearchOutput caches them),
@@ -148,6 +170,35 @@ def aggregate_csv():
 
     timings["aggregate_csv"] = timed(aggregate_csv)
 
+    agg = fresh_agg()
+
+    def aggregate_images():
+        agg_images = af.AggregateImages(agg)
+        agg_images.output_to_folder(
+            results_root.parent / "png",
+            name="dataset_name",
+            subplots=[
+                SubplotFit.Data,
+                SubplotFit.ModelData,
+                SubplotFit.ResidualMap,
+                SubplotFit.ChiSquaredMap,
+            ],
+        )
+
+    timings["aggregate_images"] = timed(aggregate_images)
+
+    agg = fresh_agg()
+
+    def aggregate_fits():
+        agg_fits = af.AggregateFITS(agg)
+        agg_fits.output_to_folder(
+            results_root.parent / "fits",
+            name="dataset_name",
+            hdus=[FITSFit.ModelData, FITSFit.ResidualMap],
+        )
+
+    timings["aggregate_fits"] = timed(aggregate_fits)
+
     if not keep:
         shutil.rmtree(results_root.parent)