LIRA does not conserve flux

[Borlaff]

Dear pyLIRA developers,

I am running some experiments with LIRA and I am finding a problem related to the conservation of flux before and after deconvolution.

I found that the total flux before (sum of all initial counts) and after (sum of the posterior traces) the deconvolution in LIRA is not preserved. The output seems to be about 95% lower than the expected. I am unable to explain this difference. I tried using different number of simulations, setting fit_background_scale to True and False, different initial backgrounds, different initial fluxes, without any different result.

I made a minimal reproducible example in case that you can provide with assistance. The code is attached to the end of this issue. The fits file required can be accessed here:

https://drive.google.com/file/d/1eqvFbofKg6OgdctiXCBZWTjAeVdveiVD/view?usp=sharing

The example reads the attached image (obs_events.fits), generates the initial conditions for the pyLIRA deconvolver, adds the PSF from another example and runs the statistics to check the flux.

These are the input image and the PSF:

And this is the histogram of flux conservation, measured as:

Flux post-LIRA deconvolution / Flux pre-LIRA deconvolution (Technically it should be one).

The division between the total flux of the posterior_trace and the input image should be equal to 1, but rather it is closer to 0.95 (95%). I am already adding the background flux as recorded in the bkgScale column in result.parameter_trace. At first I thought that that was the problem, but after checking it, it seems that even after adding bkgScale to the sum of all the flux in the posterior, it is impossible to recover the 100% of the input flux. I must be missing something in the process.

I hope that you can point us in the right direction. Thank you so much, Alex

import numpy as np
import matplotlib.pyplot as plt
from astropy.io import fits
from pylira import LIRADeconvolver, LIRADeconvolverResult, LIRASignificanceEstimator
import matplotlib.colors as colors
import matplotlib.ticker as ticker
from astropy.io import ascii

random_state = np.random.RandomState(123)

# Load the example input data and the PSF
obs_events_fits = fits.open("obs_events.fits")
obs_events = obs_events_fits[0].data

from pylira.data import point_source_gauss_psf, gauss_and_point_sources_gauss_psf, lincoln

data = point_source_gauss_psf()
psf = data["psf"] # psf_fits[0].data / np.sum(psf_fits[0].data) # Normalize the PSF

# Generate a initial conditions for LIRA
background = np.zeros(obs_events.shape) + 1.0
exposure   = np.zeros(obs_events.shape) + 1.0
alpha_init = 0.1 * np.ones(np.log2(obs_events.shape[0]).astype(int))
flux = obs_events/exposure
flux_init = flux

# Plot the data
plt.imshow(obs_events, origin='lower', cmap="RdYlBu_r",
                     norm=colors.SymLogNorm(linthresh=1, linscale=1,
                                            vmin=0.01, vmax=100))
plt.title("Observed events (PSF-convolved)")
plt.colorbar()
plt.show()

plt.imshow(psf,origin='lower', cmap="RdYlBu_r",
                     norm=colors.SymLogNorm(linthresh=1E-4, linscale=1,
                                            vmin=1E-6, vmax=0.1))
plt.title("PSF")
plt.colorbar()
plt.show()

# Generate the LIRA Object
data = {"counts": obs_events.astype("float16"),
        "flux": flux.astype("float16"),
        "flux_init": flux_init.astype("float16"),
        "exposure": exposure.astype("float16"),
        "psf": psf.astype("float16"),
        "background": background.astype("float16")}

# Run LIRA
dec = LIRADeconvolver(
    alpha_init=alpha_init,
    n_iter_max=1000,
    n_burn_in=100,
    fit_background_scale=True,
    random_state=random_state
)
result = dec.run(data)

# Check the results
post_cube = result.image_trace
result.parameter_trace.write('values.dat', format='ascii', overwrite=True) 
data_table = ascii.read('values.dat')
ratio = []
for i in range(post_cube.shape[0]):
    post_flux = np.nansum(post_cube[i]+data_table["bkgScale"][i])
    pre_flux = np.nansum(data["counts"])
    ratio.append(post_flux/pre_flux)

plt.hist(ratio)
plt.xlim(0.90,1.10)
plt.axvline(1, linewidth=3, linestyle=":", color="black", label = "100% flux conservation")
plt.legend()
plt.title("Flux conservation histogram, taking into account bkgScale")
plt.xlim(0.85,1.15)
plt.legend()
plt.show()

plt.scatter(np.linspace(1,1000,1000), ratio)
plt.xlabel("LIRA iteration")
plt.title("Flux conservation as a function of the iteration")
plt.ylabel("Flux conservation ratio (Post-LIRA/Pre-LIRA)")
plt.show()

[adonath]

Thanks a lot @Borlaff for the issue! I have not yet looked into the example in detail, but my suspicion here is that this might be due to the boundary handling of the PSF convolution. I remember when I worked with the code, the flux of the point source in the center was recovered correctly.

Can you please add this check? Integrate the flux in a given circle around the test source. You can use for example photutils.aperture.CircularAperture for this.

[adonath]

You could also try to correct the exposure for the boundary effect along the lines of:

from pylira.utils.convolution import fftconvolve
from pylira.data import point_source_gauss_psf
import numpy as np

data = point_source_gauss_psf()
psf = data["psf"] / np.sum(data["psf"])

weights = fftconvolve(np.ones(data["exposure"].shape), psf)

data["exposure"] /= weights

Maybe this resolves the difference already.