bayes-mef 0.2.0__tar.gz

bayes_mef-0.2.0/.gitignore

@@ -0,0 +1,23 @@
+# some extensions
+.vscode
+.DS_Store
+# jupyter notebook
+.ipynb_checkpoints
+*/.ipynb_checkpoints/*
+
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# folders
+data
+dist
+*.zip
+# Python / uv
+.venv/
+__pycache__/
+*.pyc
+.pytest_cache/
+.ruff_cache/
+dist/

bayes_mef-0.2.0/LICENSE

@@ -0,0 +1,28 @@
+BSD 3-Clause License
+
+Copyright (c) 2024, microscopic-image-analysis
+
+Redistribution and use in source and binary forms, with or without
+modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this
+   list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice,
+   this list of conditions and the following disclaimer in the documentation
+   and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its
+   contributors may be used to endorse or promote products derived from
+   this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE
+FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL
+DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
+SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER
+CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY,
+OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

bayes_mef-0.2.0/PKG-INFO

@@ -0,0 +1,182 @@
+Metadata-Version: 2.4
+Name: bayes-mef
+Version: 0.2.0
+Summary: Bayesian multi-exposure image fusion for robust HDR imaging
+Project-URL: Homepage, https://github.com/microscopic-image-analysis/bayes-mef
+Project-URL: Repository, https://github.com/microscopic-image-analysis/bayes-mef
+Project-URL: Issues, https://github.com/microscopic-image-analysis/bayes-mef/issues
+Author-email: Shantanu Kodgirwar <shantanu.kodgirwar@uni-jena.de>, Michael Habeck <michael.habeck@uni-jena.de>
+License-Expression: BSD-3-Clause
+License-File: LICENSE
+Keywords: HDR,expectation-maximization,image-fusion,jax,ptychography
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Science/Research
+Classifier: Programming Language :: Python :: 3
+Classifier: Topic :: Scientific/Engineering :: Image Recognition
+Classifier: Topic :: Scientific/Engineering :: Physics
+Requires-Python: >=3.11
+Requires-Dist: h5py>=3.9
+Requires-Dist: jax>=0.8
+Requires-Dist: numpy>=1.24
+Requires-Dist: scipy>=1.11
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]; extra == 'cuda12'
+Provides-Extra: cuda13
+Requires-Dist: jax[cuda13]; extra == 'cuda13'
+Description-Content-Type: text/markdown
+
+# Bayesian MEF
+[![PyPI](https://img.shields.io/pypi/v/bayes_mef)](https://pypi.org/project/bayes_mef/)
+![Python 3.11+](https://img.shields.io/badge/python-3.11+-green.svg)
+[![DOI](https://img.shields.io/badge/DOI-10.1364%2FOE.524284-blue.svg)](https://doi.org/10.1364/OE.524284)
+[![License](https://img.shields.io/badge/License-BSD_3--Clause-purple.svg)](https://opensource.org/licenses/BSD-3-Clause)
+
+Bayesian multi-exposure image fusion (MEF) is a general-purpose algorithm for robust
+high dynamic range (HDR) imaging under low SNR or varying illumination — in particular
+for phase retrieval in coherent diffractive imaging. The paper,
+["Bayesian multi-exposure image fusion for robust high dynamic range ptychography"](https://doi.org/10.1364/OE.524284),
+details the method and its benefits for ptychography ([reproducing the results](#reproducing-results)).
+
+![demo_mef](https://github.com/microscopic-image-analysis/bayes-mef/assets/64919085/d00a8c5e-5e53-4b7e-856b-381cc99523ba)
+
+This small library is implemented in [JAX](https://docs.jax.dev), so the same code runs on CPU/GPU/TPU. Inputs may be NumPy arrays, Python lists or JAX arrays; results come back as NumPy arrays.
+
+```bash
+pip install bayes_mef
+```
+
+Optionally for GPU (CUDA 12 or 13)
+
+```bash
+pip install "bayes_mef[cuda13]"  # GPU with CUDA 13; `cuda12` flag for older GPUs
+```
+
+After a GPU install, check the GPU is picked up with `bayes_mef check gpu`.
+
+## Usage
+
+A minimal example, simulating some data
+([Colab](https://colab.research.google.com/github/microscopic-image-analysis/bayes-mef/blob/main/demo.ipynb)):
+
+```python
+from bayes_mef import BayesianMEF, ConventionalMEF
+from skimage.data import camera
+import numpy as np
+
+truth = camera()
+background = 60
+times = np.array([0.1, 1, 10])  # exposure times / flux factors
+threshold = 1500                # detector limit
+
+# overexposed Poisson data from the image-formation model
+data = [np.random.poisson(t * truth + background) for t in times]
+data_saturated = np.clip(data, None, threshold, dtype="float")
+
+mef_em = BayesianMEF(data_saturated, threshold, times, background)
+mef_em.run(n_iter=100)
+fused_em = mef_em.fused_image.copy()
+
+# ConventionalMEF is the paper's MLE baseline, with the same interface
+mef_mle = ConventionalMEF(data_saturated, threshold, times, background)
+mef_mle.mle()
+fused_mle = mef_mle.fused_image.copy()
+```
+
+### Precision
+
+Single precision by default (accurate for censoring thresholds up to ~4096). For 16-bit detectors, switch to double precision *before* creating any array; a warning flags the risk otherwise:
+
+```python
+import bayes_mef
+bayes_mef.enable_x64()  # or set JAX_ENABLE_X64=1
+```
+
+### When the exposures are not known
+
+Omit `times` (and `threshold`) and they are estimated from the data, initialised from the non-saturated pixels (`flux_init="matched"`). The estimate is only a starting point, so **pass `update_fluxes=True` with it** and let EM refine it:
+
+```python
+mef_em = BayesianMEF(data_saturated, background=background, update_fluxes=True)
+mef_em.run(n_iter=200)
+```
+
+`update_fluxes` defaults to False — `times` are used exactly as given, estimated or not — and estimating them without it warns. On the simulation above, the unrefined estimate correlates 0.16 with the truth against 0.989 once refined. `flux_init="uncensored"` selects the v0.1.9 initialiser, whose ratios compress under heavy censoring.
+
+Once the background dominates the signal, i.e., the weak, low-SNR regime the [paper](https://doi.org/10.1364/OE.524284) targets: summed counts carry almost no information about the exposures, and it is EM's iterative background handling that recovers the range; a warning is raised when the estimate comes out nearly flat. Estimated fluxes are relative, so the fused image is on a relative scale. **Supply the real exposure times whenever you know them.**
+
+### Fusing a 4D ptychogram dataset
+
+For ptychography, we record multiple exposures per scan position. `LaunchMEF` fuses every scan position with a single vectorised program (CPU or GPU):
+
+```python
+from bayes_mef import LaunchMEF
+
+launch_mef = LaunchMEF(
+    ptychogram_stack,    # (n_exposures, n_scans, dp_x, dp_y)
+    background,          # a number, one dark frame, or one per exposure
+    times=None,          # None -> estimated from the data
+    threshold=None,      # None -> estimated automatically
+    update_fluxes=False, # True -> EM refines the fluxes; pair with estimated times
+    flux_init="matched",
+)
+
+# returns fused patterns (n_scans, dp_x, dp_y) and the flux factors
+fused_ptyem_stack, em_flux_factors = launch_mef.run_em(n_iter=150)
+
+# or the conventional MLE baseline over the whole stack (just the fused patterns)
+fused_ptymle_stack = launch_mef.run_mle()
+```
+
+Backgrounds are taken however they were recorded, on `LaunchMEF` and on the single stack classes alike: a scalar, one offset per exposure `(n_exposures,)`, a single dark frame `(dp_x, dp_y)`, one frame per exposure `(n_exposures, dp_x, dp_y)`, or one per image (the full stack shape). Each is given its exposure axis explicitly, so a per-exposure vector is never spread along the image columns, and a mismatched shape raises instead of broadcasting into something that quietly means the wrong thing.
+
+The old `n_cpus` argument is accepted but ignored (it warns), since there are no worker processes to size any more.
+
+Scans are fused in chunks sized to the device's free memory, so a stack larger than device memory still works. Set `batch_size` yourself if you hit a `MemoryError` or fuse alongside other work (results do not depend on it):
+
+```python
+fused, fluxes = launch_mef.run_em(n_iter, batch_size=8)
+```
+
+See [synthetic_mef.py](scripts/synthetic_mef.py) for detailed usage on synthetic
+ptychography data, and [benchmarks/FINDINGS.md](benchmarks/FINDINGS.md) for a study
+of when each method helps and for performance benchmarks.
+
+## Reproducing results
+
+To reproduce the ptychographic reconstructions from the paper:
+
+1. Clone the repo:
+   ```bash
+   git clone https://github.com/microscopic-image-analysis/bayes-mef.git
+   cd bayes-mef
+   ```
+2. Install the pinned dependencies with [uv](https://docs.astral.sh/uv/), then prefix
+   commands with `uv run` (e.g. `uv run python scripts/synthetic_mef.py`):
+   ```bash
+   uv sync --locked --group scripts
+   ```
+3. Download the data from [Zenodo](https://zenodo.org/doi/10.5281/zenodo.10964222):
+   ```bash
+   ./download_data.sh
+   ```
+4. Optional: install [`cupy`](https://docs.cupy.dev/en/stable/install.html) for faster
+   GPU reconstructions.
+5. Run files from [scripts/](scripts) to plot the results.
+
+## Citation
+If this algorithm or the publication was useful, please cite:
+```tex
+@article{Kodgirwar:24,
+author = {Shantanu Kodgirwar and Lars Loetgering and Chang Liu and Aleena Joseph and Leona Licht and Daniel S. Penagos Molina and Wilhelm Eschen and Jan Rothhardt and Michael Habeck},
+journal = {Opt. Express},
+number = {16},
+pages = {28090--28099},
+publisher = {Optica Publishing Group},
+title = {Bayesian multi-exposure image fusion for robust high dynamic range ptychography},
+volume = {32},
+month = {Jul},
+year = {2024},
+url = {https://opg.optica.org/oe/abstract.cfm?URI=oe-32-16-28090},
+doi = {10.1364/OE.524284},
+}
+```

bayes_mef-0.2.0/README.md

@@ -0,0 +1,155 @@
+# Bayesian MEF
+[![PyPI](https://img.shields.io/pypi/v/bayes_mef)](https://pypi.org/project/bayes_mef/)
+![Python 3.11+](https://img.shields.io/badge/python-3.11+-green.svg)
+[![DOI](https://img.shields.io/badge/DOI-10.1364%2FOE.524284-blue.svg)](https://doi.org/10.1364/OE.524284)
+[![License](https://img.shields.io/badge/License-BSD_3--Clause-purple.svg)](https://opensource.org/licenses/BSD-3-Clause)
+
+Bayesian multi-exposure image fusion (MEF) is a general-purpose algorithm for robust
+high dynamic range (HDR) imaging under low SNR or varying illumination — in particular
+for phase retrieval in coherent diffractive imaging. The paper,
+["Bayesian multi-exposure image fusion for robust high dynamic range ptychography"](https://doi.org/10.1364/OE.524284),
+details the method and its benefits for ptychography ([reproducing the results](#reproducing-results)).
+
+![demo_mef](https://github.com/microscopic-image-analysis/bayes-mef/assets/64919085/d00a8c5e-5e53-4b7e-856b-381cc99523ba)
+
+This small library is implemented in [JAX](https://docs.jax.dev), so the same code runs on CPU/GPU/TPU. Inputs may be NumPy arrays, Python lists or JAX arrays; results come back as NumPy arrays.
+
+```bash
+pip install bayes_mef
+```
+
+Optionally for GPU (CUDA 12 or 13)
+
+```bash
+pip install "bayes_mef[cuda13]"  # GPU with CUDA 13; `cuda12` flag for older GPUs
+```
+
+After a GPU install, check the GPU is picked up with `bayes_mef check gpu`.
+
+## Usage
+
+A minimal example, simulating some data
+([Colab](https://colab.research.google.com/github/microscopic-image-analysis/bayes-mef/blob/main/demo.ipynb)):
+
+```python
+from bayes_mef import BayesianMEF, ConventionalMEF
+from skimage.data import camera
+import numpy as np
+
+truth = camera()
+background = 60
+times = np.array([0.1, 1, 10])  # exposure times / flux factors
+threshold = 1500                # detector limit
+
+# overexposed Poisson data from the image-formation model
+data = [np.random.poisson(t * truth + background) for t in times]
+data_saturated = np.clip(data, None, threshold, dtype="float")
+
+mef_em = BayesianMEF(data_saturated, threshold, times, background)
+mef_em.run(n_iter=100)
+fused_em = mef_em.fused_image.copy()
+
+# ConventionalMEF is the paper's MLE baseline, with the same interface
+mef_mle = ConventionalMEF(data_saturated, threshold, times, background)
+mef_mle.mle()
+fused_mle = mef_mle.fused_image.copy()
+```
+
+### Precision
+
+Single precision by default (accurate for censoring thresholds up to ~4096). For 16-bit detectors, switch to double precision *before* creating any array; a warning flags the risk otherwise:
+
+```python
+import bayes_mef
+bayes_mef.enable_x64()  # or set JAX_ENABLE_X64=1
+```
+
+### When the exposures are not known
+
+Omit `times` (and `threshold`) and they are estimated from the data, initialised from the non-saturated pixels (`flux_init="matched"`). The estimate is only a starting point, so **pass `update_fluxes=True` with it** and let EM refine it:
+
+```python
+mef_em = BayesianMEF(data_saturated, background=background, update_fluxes=True)
+mef_em.run(n_iter=200)
+```
+
+`update_fluxes` defaults to False — `times` are used exactly as given, estimated or not — and estimating them without it warns. On the simulation above, the unrefined estimate correlates 0.16 with the truth against 0.989 once refined. `flux_init="uncensored"` selects the v0.1.9 initialiser, whose ratios compress under heavy censoring.
+
+Once the background dominates the signal, i.e., the weak, low-SNR regime the [paper](https://doi.org/10.1364/OE.524284) targets: summed counts carry almost no information about the exposures, and it is EM's iterative background handling that recovers the range; a warning is raised when the estimate comes out nearly flat. Estimated fluxes are relative, so the fused image is on a relative scale. **Supply the real exposure times whenever you know them.**
+
+### Fusing a 4D ptychogram dataset
+
+For ptychography, we record multiple exposures per scan position. `LaunchMEF` fuses every scan position with a single vectorised program (CPU or GPU):
+
+```python
+from bayes_mef import LaunchMEF
+
+launch_mef = LaunchMEF(
+    ptychogram_stack,    # (n_exposures, n_scans, dp_x, dp_y)
+    background,          # a number, one dark frame, or one per exposure
+    times=None,          # None -> estimated from the data
+    threshold=None,      # None -> estimated automatically
+    update_fluxes=False, # True -> EM refines the fluxes; pair with estimated times
+    flux_init="matched",
+)
+
+# returns fused patterns (n_scans, dp_x, dp_y) and the flux factors
+fused_ptyem_stack, em_flux_factors = launch_mef.run_em(n_iter=150)
+
+# or the conventional MLE baseline over the whole stack (just the fused patterns)
+fused_ptymle_stack = launch_mef.run_mle()
+```
+
+Backgrounds are taken however they were recorded, on `LaunchMEF` and on the single stack classes alike: a scalar, one offset per exposure `(n_exposures,)`, a single dark frame `(dp_x, dp_y)`, one frame per exposure `(n_exposures, dp_x, dp_y)`, or one per image (the full stack shape). Each is given its exposure axis explicitly, so a per-exposure vector is never spread along the image columns, and a mismatched shape raises instead of broadcasting into something that quietly means the wrong thing.
+
+The old `n_cpus` argument is accepted but ignored (it warns), since there are no worker processes to size any more.
+
+Scans are fused in chunks sized to the device's free memory, so a stack larger than device memory still works. Set `batch_size` yourself if you hit a `MemoryError` or fuse alongside other work (results do not depend on it):
+
+```python
+fused, fluxes = launch_mef.run_em(n_iter, batch_size=8)
+```
+
+See [synthetic_mef.py](scripts/synthetic_mef.py) for detailed usage on synthetic
+ptychography data, and [benchmarks/FINDINGS.md](benchmarks/FINDINGS.md) for a study
+of when each method helps and for performance benchmarks.
+
+## Reproducing results
+
+To reproduce the ptychographic reconstructions from the paper:
+
+1. Clone the repo:
+   ```bash
+   git clone https://github.com/microscopic-image-analysis/bayes-mef.git
+   cd bayes-mef
+   ```
+2. Install the pinned dependencies with [uv](https://docs.astral.sh/uv/), then prefix
+   commands with `uv run` (e.g. `uv run python scripts/synthetic_mef.py`):
+   ```bash
+   uv sync --locked --group scripts
+   ```
+3. Download the data from [Zenodo](https://zenodo.org/doi/10.5281/zenodo.10964222):
+   ```bash
+   ./download_data.sh
+   ```
+4. Optional: install [`cupy`](https://docs.cupy.dev/en/stable/install.html) for faster
+   GPU reconstructions.
+5. Run files from [scripts/](scripts) to plot the results.
+
+## Citation
+If this algorithm or the publication was useful, please cite:
+```tex
+@article{Kodgirwar:24,
+author = {Shantanu Kodgirwar and Lars Loetgering and Chang Liu and Aleena Joseph and Leona Licht and Daniel S. Penagos Molina and Wilhelm Eschen and Jan Rothhardt and Michael Habeck},
+journal = {Opt. Express},
+number = {16},
+pages = {28090--28099},
+publisher = {Optica Publishing Group},
+title = {Bayesian multi-exposure image fusion for robust high dynamic range ptychography},
+volume = {32},
+month = {Jul},
+year = {2024},
+url = {https://opg.optica.org/oe/abstract.cfm?URI=oe-32-16-28090},
+doi = {10.1364/OE.524284},
+}
+```

bayes_mef-0.2.0/bayes_mef/__init__.py

@@ -0,0 +1,54 @@
+"""Bayesian multi-exposure image fusion for robust HDR imaging."""
+
+import importlib
+from importlib.metadata import PackageNotFoundError, version
+from typing import TYPE_CHECKING
+
+if TYPE_CHECKING:
+    # Never executed. The public names are bound lazily below, which leaves type
+    # checkers and editors with nothing to resolve, so "go to definition" on
+    # `from bayes_mef import BayesianMEF` lands nowhere. These imports exist for
+    # them alone; at runtime the lazy path is still the only one.
+    from ._compat import enable_x64
+    from .algos import BayesianMEF, ConventionalMEF
+    from .mef_launcher import LaunchMEF
+    from .utils import set_fluxes, set_threshold
+
+__all__ = [
+    "BayesianMEF",
+    "ConventionalMEF",
+    "LaunchMEF",
+    "enable_x64",
+    "set_fluxes",
+    "set_threshold",
+]
+
+try:
+    __version__ = version("bayes-mef")
+except PackageNotFoundError:  # not installed, e.g. running from a source checkout
+    __version__ = "0.0.0"
+
+_LAZY = {
+    "BayesianMEF": "algos",
+    "ConventionalMEF": "algos",
+    "LaunchMEF": "mef_launcher",
+    "set_fluxes": "utils",
+    "set_threshold": "utils",
+    "enable_x64": "_compat",
+}
+
+
+def __getattr__(name):
+    """Import submodules on first use, so `import bayes_mef` does not load JAX.
+
+    Keeps the import cheap, and leaves a window for callers to configure JAX
+    (precision, platform) before the backend initialises.
+    """
+    if name in _LAZY:
+        module = importlib.import_module(f".{_LAZY[name]}", __name__)
+        return getattr(module, name)
+    raise AttributeError(f"module {__name__!r} has no attribute {name!r}")
+
+
+def __dir__():
+    return sorted(__all__)

bayes_mef-0.2.0/bayes_mef/_compat.py

@@ -0,0 +1,132 @@
+"""NumPy <-> JAX boundary: dtype policy, input coercion, precision control."""
+
+import warnings
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+_X64_WARNED = False
+
+
+def default_dtype():
+    """Working float dtype: float64 when x64 is enabled, float32 otherwise."""
+    return jnp.float64 if jax.config.jax_enable_x64 else jnp.float32
+
+
+def enable_x64():
+    """Run bayes_mef in double precision.
+
+    Must be called before any JAX array is created. Equivalent to setting
+    ``JAX_ENABLE_X64=1`` in the environment.
+    """
+    jax.config.update("jax_enable_x64", True)
+
+
+def as_array(x, dtype=None):
+    """Coerce numpy arrays, Python lists/scalars and JAX arrays to a JAX array."""
+    return jnp.asarray(x, dtype=default_dtype() if dtype is None else dtype)
+
+
+def as_background(x):
+    """Background may be None, a scalar, or a full array; all broadcast against data."""
+    return as_array(0.0 if x is None else x)
+
+
+def align_background(background, stack_shape, *, scan_axis=False):
+    """Lay a recorded background out so it broadcasts against the stack.
+
+    Backgrounds arrive in whatever form the experiment produced: a number, one
+    offset per exposure, a single dark frame, one frame per exposure, or one per
+    image. All of them describe the same thing, but NumPy lines shapes up from
+    the trailing axis, so only some happen to broadcast correctly against a
+    stack -- a per-exposure vector silently spreads along the image columns
+    instead. Each form is therefore given its exposure axis explicitly here,
+    once, so nothing downstream has to reason about the layout again.
+
+    `stack_shape` is the stack as the caller holds it, exposures first:
+    (n_exposures, *image), or (n_exposures, n_scans, *image) with
+    `scan_axis=True`. Accepted backgrounds, in that same layout:
+
+        None, scalar, (1,)      one offset for the whole stack
+        (n_exposures,)          one offset per exposure
+        image                   one dark frame, shared by every exposure
+        (n_exposures, *image)   one dark frame per exposure
+        stack_shape             one dark frame per image (scan_axis only)
+
+    Anything of the right rank whose axes are 1 or the stack's own length is
+    accepted too, so a background already written as (n_exposures, 1, 1) or
+    (1, X, Y) passes through untouched. Rank is what disambiguates the forms,
+    so a shape that is merely broadcastable by NumPy's trailing-axis rule is
+    rejected rather than guessed at.
+
+    Returns a NumPy array in the caller's layout. A per-image background that
+    does not actually vary across scan positions -- `np.broadcast_to(frame,
+    stack.shape)`, the usual way of saying "this frame everywhere" -- collapses
+    to the per-exposure form, so the launcher uploads it once rather than
+    re-sending a stack-sized copy with every chunk.
+    """
+    bg = np.asarray(0.0 if background is None else background)
+    stack_shape = tuple(stack_shape)
+    n_exposures = stack_shape[0]
+    image = stack_shape[2:] if scan_axis else stack_shape[1:]
+
+    def fits(target):
+        return len(bg.shape) == len(target) and all(
+            b in (1, t) for b, t in zip(bg.shape, target, strict=True)
+        )
+
+    if bg.ndim == 0 or bg.shape == (1,):
+        return bg.reshape(())
+
+    if bg.shape == (n_exposures,):
+        return bg.reshape((n_exposures,) + (1,) * len(image))
+
+    if fits(image) or fits((n_exposures, *image)):
+        return bg
+
+    if scan_axis and fits(stack_shape):
+        # A singleton or zero-strided scan axis means the frame was broadcast,
+        # not measured per position, so the axis carries nothing worth keeping.
+        if bg.shape[1] == 1 or bg.strides[1] == 0:
+            return bg[:, 0]
+        return bg
+
+    accepted = [f"a scalar, {(n_exposures,)}, {image}, {(n_exposures, *image)}"]
+    if scan_axis:
+        accepted.append(f"or {stack_shape}")
+    raise ValueError(
+        f"background shape {bg.shape} does not fit a stack of shape "
+        f"{stack_shape}. Expected {' '.join(accepted)}."
+    )
+
+
+def to_numpy(x):
+    """Public results are NumPy, so downstream `.flat`, h5py and in-place idioms work."""
+    return np.asarray(x)
+
+
+# Above this censoring threshold, XLA's float32 igamma loses accuracy near
+# rate ~ threshold -- exactly where the EM operates on censored pixels. Measured
+# relative error at rate = 0.9 * threshold: 1e-3 at 4096, 2e-2 at 16383, 1e-1 at
+# 65279. Below the bound the error stays at or under 2e-3.
+X32_SAFE_THRESHOLD = 4096.0
+
+
+def warn_if_precision_risk(threshold):
+    """Warn once when single precision is not accurate enough for this threshold."""
+    global _X64_WARNED
+    if _X64_WARNED or jax.config.jax_enable_x64:
+        return
+    if float(np.max(np.asarray(threshold))) <= X32_SAFE_THRESHOLD:
+        return
+    _X64_WARNED = True
+    warnings.warn(
+        f"censoring threshold exceeds {X32_SAFE_THRESHOLD:.0f}, where float32 "
+        "evaluation of the censored-Poisson expectation loses up to ~10% "
+        "accuracy for rates near the threshold. Call bayes_mef.enable_x64() "
+        "(or set JAX_ENABLE_X64=1) before creating any array to run in double "
+        "precision.",
+        UserWarning,
+        stacklevel=3,
+    )