feat: background thread for on-the-fly visualization during sampling

[Jammy2211]

Overview

On-the-fly visualization via perform_quick_update currently blocks sampling while plots render and save to disk. This task adds a background thread worker so sampling continues uninterrupted during visualization. It also improves Jupyter notebook safety for display handling.

Plan

• Add a background visualization worker that runs perform_quick_update in a separate thread, so sampling is not blocked
• Use a "latest-only" queue pattern — if a new best-fit arrives while the previous visualization is still running, discard the stale one
• Conditionally convert JAX arrays to numpy before handing to the visualization thread (only when JAX is active and the visualizer doesn't yet support JAX natively), keeping the door open for full JAX visualization later
• Add Jupyter notebook safety — detect the notebook environment, save plots to file in the background thread, and display from the main thread on the next fitness call
• Make the feature opt-in via configuration so existing synchronous behavior is preserved by default
• Add tests for the background worker lifecycle

Detailed implementation plan

Affected Repositories

• PyAutoFit (primary)
• PyAutoGalaxy (minor)

Work Classification

Library

Branch Survey

Repository	Current Branch	Dirty?
PyAutoFit	main	clean
PyAutoGalaxy	main	scrap.py untracked

Suggested branch: feature/on-the-fly-modeling
Worktree root: ~/Code/PyAutoLabs-wt/on-the-fly-modeling/

Implementation Steps

1. Create autofit/non_linear/quick_update.py — new module containing a BackgroundQuickUpdate class:

   • Uses a single threading.Thread as a daemon worker
   • Accepts (analysis, paths, instance) tuples via a thread-safe queue
   • "Latest-only" semantics: uses threading.Event + atomic swap so only the most recent request is processed
   • Conditional JAX→numpy conversion: checks whether JAX is active and whether the visualizer supports JAX natively (supports_jax_visualization property). Only converts if JAX is active but the visualizer doesn't support it yet
   • Worker thread creates its own matplotlib figures, saves to disk, closes them — thread-safe for non-interactive backends
   • shutdown() method signals the thread to exit and joins it
   • Catches and logs exceptions in the worker thread (never crashes the sampler)

2. Modify autofit/non_linear/fitness.py — manage_quick_update method:

   • Accept a background_quick_update flag (from search config or Fitness.__init__)
   • When enabled, delegate to BackgroundQuickUpdate instead of calling analysis.perform_quick_update() directly
   • For Jupyter: after the background thread saves the plot, display the image file using IPython.display.Image + display() on the next manage_quick_update call (main thread, avoids Jupyter display issues)
   • On sampler completion, call shutdown() to cleanly join the worker thread

3. Modify autofit/non_linear/analysis/analysis.py — add two properties:

   • supports_background_update (default False) — downstream classes override to True when their visualizer is file-save safe
   • supports_jax_visualization (default False) — downstream classes override to True once their visualizer stack is JAX-native (future task)

4. Modify PyAutoGalaxy autogalaxy/analysis/analysis/analysis.py:

   • Override supports_background_update to return True
   • Leave supports_jax_visualization as False for now (future task)

5. Add configuration in autofit/config/ for quick_update_background (default: False)

6. Add tests in test_autofit/non_linear/test_quick_update.py:

   • Test latest-only queue semantics
   • Test conditional JAX array conversion
   • Test graceful shutdown
   • Test exception handling in worker thread

Key Files

• autofit/non_linear/quick_update.py — new background worker class
• autofit/non_linear/fitness.py — integration point (~lines 290-334)
• autofit/non_linear/analysis/analysis.py — new properties (~line 302)
• autogalaxy/analysis/analysis/analysis.py — override property (~line 154)

Original Prompt

Click to expand starting prompt

When I perform sampling, I am able to produce on-the-fly images and output, which currently overwrites
the method:

def perform_quick_update(self, paths, instance):
    raise NotImplementedError

In @PyAutoFit/autofit/non_linear/analysis.py

An example of how this is used in autolens to show lens modeling is:

def perform_quick_update(self, paths, instance):
    self.Visualizer().visualize(
        analysis=self,
        paths=paths,
        instance=instance,
        during_analysis=True,
        quick_update=True,
    )

Which is in @PyAutoGalaxy/autogalaxy/analysis

Where the visualizer object has been adapted to be reused but only call a subset of key functions and outputs for a quick update.

This works ok, but it is stop-start whereby a sampler has to stop to perform visualization. For certian users I have
heard matplotlib can cause display issues in Jupyter Notebooks.

My questions are, can this functionality be updated to output the viuslaizaiton on a separate process while sampling
contiues? Even with JAX and GPU? And can I put any safety checks in or improve it to perform more generally on Notebooks.

[Jammy2211]

Smoke Test Results — 2026-04-13

Workspace	Passed	Failed	Total
autofit_workspace	7	0	7
autogalaxy_workspace	4	2	6
autolens_workspace	6	1	7
autolens_workspace_test	6	2	8

All 5 failures are pre-existing missing .fits dataset files (FileNotFoundError) — unrelated to these code changes.

Failures (all pre-existing data issues)

autogalaxy_workspace/imaging/start_here.py

FileNotFoundError: 'dataset/imaging/extra_galaxies/mask_extra_galaxies.fits'

autogalaxy_workspace/interferometer/start_here.py

FileNotFoundError: 'dataset/interferometer/uv_wavelengths/sma.fits'

autolens_workspace/interferometer/modeling.py

FileNotFoundError: 'dataset/interferometer/uv_wavelengths/sma.fits'

autolens_workspace_test/database/scrape/general.py

FileNotFoundError: 'dataset/imaging/no_lens_light/data.fits'

autolens_workspace_test/jax_likelihood_functions/interferometer/rectangular.py

FileNotFoundError: 'dataset/interferometer/uv_wavelengths/sma.fits'

[Jammy2211]

Shipped — 2026-04-13

PRs

• feat: background thread for on-the-fly visualization #1212
• feat: opt in to background quick-update visualization PyAutoGalaxy#350

Summary

Added a background daemon thread for on-the-fly visualization during sampling. The perform_quick_update call now returns in ~0.004s (vs 4-14s synchronously) by submitting to a worker thread. Uses latest-only queue semantics and conditional JAX→numpy conversion. Opt-in via quick_update_background: true in general.yaml. Verified in both script and Jupyter notebook environments.

Session Notes

• The background worker drains pending work on shutdown to avoid losing the final visualization.
• PYAUTO_FAST_PLOTS=1 should be used alongside to avoid a pre-existing matplotlib LaTeX parsing crash in tight_layout().