Universal Bayesian Imaging Kit (UBIK)

Updated 21 December 2025

Universal Bayesian Imaging Kit (UBIK) is an open-source modular framework based on Information Field Theory, enabling the probabilistic reconstruction of astrophysical images.
It employs advanced inference techniques such as MAP, variational inference, MGVI, and geoVI to achieve precise source separation and uncertainty quantification.
Its modular, instrument-agnostic design with GPU/TPU acceleration enables scalable, multi-channel analysis across diverse electromagnetic datasets.

The Universal Bayesian Imaging Kit (UBIK) is an open-source, modular framework for end-to-end Bayesian imaging of astrophysical data across the electromagnetic spectrum, grounded in the principles of Information Field Theory (IFT) and implemented numerically via the NIFTy library, with accelerated backends such as JAX and NIFTy.re. UBIK provides instrument-agnostic generative modeling, supports joint multi-instrument and multi-wavelength analyses, and features advanced Bayesian inference schemes capable of quantifying uncertainties and separating astrophysical source components within a unified, scalable architecture. By encapsulating both domain-specific knowledge and flexible inference engines, it enables the generation of physically interpretable sky images, spatio-spectral cubes, and uncertainty maps with minimal bespoke coding, covering use cases from X-ray to radio astronomy (Enßlin, 24 Aug 2025, Enßlin et al., 13 Dec 2025, Eberle et al., 16 Sep 2024).

1. Theoretical Foundations and Mathematical Framework

UBIK is fundamentally underpinned by Information Field Theory (IFT), which treats astrophysical images and sky signals as continuous random fields. In this framework, inference is performed directly on functions, with uncertainties encoded as probability functionals over fields. The core elements are:

Prior $\mathcal{P}(s)$ : Encodes assumptions about spatial and spectral correlations, regularity, positivity, and non-Gaussianity in the sky signal $s(\mathbf{x}, E)$ .
Likelihood $\mathcal{P}(d|s)$ : Relates the physical sky to observed data $d$ via instrument response operators $R$ ; often compositional and able to capture Poisson (photon-counting) and Gaussian (read-noise) statistics.
Posterior $\mathcal{P}(s|d) \propto \mathcal{P}(d|s) \mathcal{P}(s)$ : Encapsulates all knowledge about the sky given the data, forming the target distribution for reconstruction and uncertainty quantification.

A latent Gaussian field $\xi(\mathbf{z})$ is typically introduced with

$\mathcal{P}(\xi) = \mathcal{G}(\xi,\,1\!\!1) = \frac{1}{\sqrt{\det(2\pi\,1\!\!1)}} \exp\left[-\frac12\,\xi^\dagger\,\xi\right],$

and the sky components are constructed via deterministic, differentiable maps $s = f(\xi)$ , enforcing all desired statistical and physical constraints a priori (Enßlin et al., 13 Dec 2025).

The field-Hamiltonian formulation is central: $\mathcal{H}(d,s) = -\ln \left[\mathcal{P}(d|s) \mathcal{P}(s)\right],$ where the Hamiltonian serves as the penalized misfit functional.

2. Modular Software Architecture

UBIK builds upon NIFTy/NIFTy.re field abstractions and operator mechanics, encapsulating its architecture into modular layers that align with the physical and observational workflow:

Data Ingestion: Instrument-specific loaders extract data from formats (e.g., FITS), perform metadata analysis, and output NIFTy Fields along with noise covariance representations.
Response Modules: Libraries of linear and nonlinear operators represent instrument actions such as PSF convolution, energy redistribution, Fourier-plane sampling (e.g., for ALMA), or coded-aperture imaging (e.g., Chandra, eROSITA).
Prior and Sky Models: UBIK offers additive/multiplicative prior modules (Gaussian fields, log-normal, Gaussian process, etc.), enabling mixture models that represent diffuse emission, point sources, and extended structures within a single generative framework.
Inference Engines: MAP solvers (e.g., L-BFGS), variational approximations (MGVI, geoVI), and, optionally, Monte Carlo sampling machinery, all powered by JAX for high-dimensional autodifferentiation and GPU/TPU acceleration.

Each subsystem is designed for composability and Pythonic interface exposure, with heavy numerical routines JIT-compiled to minimize overhead (Eberle et al., 16 Sep 2024).

3. Bayesian Inference Principles and Algorithms

UBIK integrates state-of-the-art Bayesian inference techniques:

MAP (Maximum a Posteriori) Estimation: Finds the minimizer of the information Hamiltonian using standard optimizers. Suitable for fast prototyping and initial diagnostics, but does not provide credible intervals.
Variational Inference (VI): Uses parametric surrogates (typically Gaussian with flexible low-rank or natural-metric covariance) to approximate the posterior. The Metric Gaussian VI (MGVI) and geometric VI (geoVI) are supported, with the latter leveraging Riemannian geometry through local normalizing flows to better capture posterior geometry in highly non-linear cases.
Latent-space Transformations: All model complexity is shifted to the deterministic map $s = f_\sigma(\xi)$ , enabling all priors to collapse to a simple quadratic form in $\xi$ : $\mathcal{P}(\xi) = \exp\left(-\frac{1}{2}\xi^\dagger\xi\right),$ thus standardizing the representation and simplifying automatic differentiation.

Posterior samples in latent space are pushed through $f(\xi)$ to obtain sky realizations, over which arbitrary statistics can be constructed, supporting pixelwise means, variances, or higher-order moments.

4. Spatio-Spectral Modeling and Instrument-Agnostic Design

UBIK represents astrophysical signals via spatio-spectral cubes $s(\mathbf{x}, E)$ , with physical modeling and instrument response factorized to enable joint, multi-instrument analysis:

Sky Model Decomposition: The sky is partitioned into diffuse, point source, and extended emission components, each with distinct latent fields and statistical structure: $s(\mathbf{x},E) = s_{\text{diffuse}}(\mathbf{x},E) + \sum_{p=1}^{N_{\rm ps}} A_p\,\delta(\mathbf{x} - \mathbf{x}_p)\,\phi_p(E) + s_{\text{extended}}(\mathbf{x},E)$
Instrument Response Formulation: Every instrument (e.g., Chandra, eROSITA, JWST, ALMA) is modeled by actions such as spatial convolution, spectral redistribution, mosaic grid tiling, or $(u, v)$ -plane sampling with baseline-dependent gain for interferometry: $R^\ell[s](i,\alpha) = \int d\mathbf{x}\, \mathrm{PSF}^\ell_i(\mathbf{x}, E_\alpha) A^\ell(E_\alpha)\,s(\mathbf{x}, E_\alpha)$
Likelihood Factorization Over Instruments: For multi-instrument fusion, the total likelihood is a product over all individual likelihoods: $\mathcal{P}(\{d^\ell\}|s) = \prod_\ell \mathcal{P}(d^\ell | R^\ell[s])$

This abstraction allows rapid extension to new instruments, provided their response operator can be programmatically specified (Enßlin et al., 13 Dec 2025).

5. Algorithmic Workflow and Example Pipelines

A typical reconstruction pipeline in UBIK proceeds as:

Model Instantiation: Latent Gaussian fields and their transforms for each emission component.
Forward Modeling: Evaluation of $s(\xi)$ through the instrument response stack for all supported instruments.
Objective Construction: Automatic assembly of the joint information Hamiltonian and its gradient via JAX autodiff.
Inference Execution: Running MGVI or geoVI to convergence, yielding a variational approximation to the posterior.
Uncertainty Quantification: Posterior sampling, extraction of mean and variance maps, or other functionals.

Several example codes demonstrate setup and execution for Chandra X-ray and JWST infrared imaging, highlighting the minimal code required to go from event data to Bayesian sky reconstruction with uncertainty maps (Enßlin, 24 Aug 2025, Eberle et al., 16 Sep 2024).

6. Performance, Scalability, and Benchmarking

UBIK leverages algorithmic and architectural strategies for handling high-dimensional problems:

FFT-accelerated Priors/Responses: Enables $\mathcal{O}(N\log N)$ scaling for dense convolutions in images with $N$ pixels or voxels.
Implicit Covariance Representation: Avoids explicit formation of large matrices, reducing memory and computational overhead, required only to implement matrix-vector (or action) routines.
JAX/Pmap Parallelization: Supports strong scaling across multiple GPUs or TPUs for multi-channel, multi-pointing, and high-resolution analysis.

Benchmarks report that full Bayesian analyses of $\sim10^6$ voxel cubes, with $\sim10^5$ latent parameters, are completed in under a day on a single GPU. A factor-of-2 improvement in dynamic range over CLEAN-based methods has been shown for ALMA continuum imaging, with robust separation of compact and diffuse sources and quantified posteriors. Similar efficiency is demonstrated for X-ray and IR applications (Enßlin et al., 13 Dec 2025, Eberle et al., 16 Sep 2024).

7. Limitations, Assumptions, and Ongoing Developments

UBIK faces a set of known limitations and targets for future improvement:

Latent Gaussianity: While deterministic transforms enable rich prior modeling, highly multimodal or non-Gaussian phenomena may require expressive flows or alternative latent representations.
Temporal Variability: The current focus is on static sky models; time-dependent (4D) imaging is a planned extension.
Instrument Calibration Complexity: The framework allows modeling calibration uncertainties and gain/PSF variations, but full integration with nonlinear calibration chains and metrology remains under development.
Polarization and Stokes Imaging: Extensions to polarimetric inference and comprehensive Stokes parameter imaging are in active progress.
Extreme Scalability and Distributed Inference: Handling SKA-scale survey data will need integration with distributed computing backends (e.g., Dask or MPI) and advanced preconditioning strategies.

Despite these areas of ongoing work, UBIK has established itself as a generalizable, scalable, and scientifically robust engine for probabilistic imaging in astrophysics (Enßlin et al., 13 Dec 2025, Enßlin, 24 Aug 2025, Eberle et al., 16 Sep 2024).

References: