diff --git a/data/poisson.npz b/data/poisson.npz
new file mode 100644
index 0000000000000000000000000000000000000000..456134a992071c2a6bcb91a8377cc9dae90a5a7f
Binary files /dev/null and b/data/poisson.npz differ
diff --git a/demo_poisson.ipynb b/demo_poisson.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..de9b7f8d22698c07e3efa15e5d4ab0146063664f
--- /dev/null
+++ b/demo_poisson.ipynb
@@ -0,0 +1,208 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Nifty tutorial for Poisson count data\n",
+ "====================================="
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Setup\n",
+ "-----"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import matplotlib.pyplot as plt\n",
+ "import nifty8 as ift\n",
+ "from utils import plot_2D, load_psf, geovi_sampling, plot_posterior\n",
+ "\n",
+ "ift.random.push_sseq_from_seed(42)\n",
+ "evidences = []"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Load data and visualize\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "position_space = ift.RGSpace([128, 128])\n",
+ "\n",
+ "# Homogeneous poisson process\n",
+ "\n",
+ "\n",
+ "print(model1)\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Set up likelihood & PSF\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Inference model 1\n",
+ "\n",
+ "\n",
+ "print(evidence)\n",
+ "evidences += [evidence, ]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Independent poisson process\n",
+ "\n",
+ "# Inference model 2\n",
+ "\n",
+ "\n",
+ "evidences += [evidence, ]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Compare evidence\n",
+ "print(evidences)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Diffuse poisson process\n",
+ "args = {\n",
+ " 'offset_mean': .5,\n",
+ " 'offset_std': (1., 1E-5),\n",
+ "\n",
+ " # Amplitude of field fluctuations\n",
+ " 'fluctuations': (1.5, 0.5), # 1.0, 1e-2\n",
+ "\n",
+ " # Exponent of power law power spectrum component\n",
+ " 'loglogavgslope': (-4., 1), # -6.0, 1\n",
+ "\n",
+ " # Amplitude of integrated Wiener process power spectrum component\n",
+ " 'flexibility': (1., 0.2), # 2.0, 1.0\n",
+ "\n",
+ " # How ragged the integrated Wiener process component is\n",
+ " 'asperity': (0.1, 0.01), # 0.1, 0.5\n",
+ "\n",
+ " # Name of the input keys\n",
+ " 'prefix' : 'diffuse'\n",
+ "}\n",
+ "correlated_field = ift.SimpleCorrelatedField(position_space, **args)\n",
+ "pspec = correlated_field.power_spectrum\n",
+ "\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Prior samples\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Inference model 3\n",
+ "\n",
+ "evidences += [evidence, ]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Compare evidence\n",
+ "print(evidences)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Posterior visualization\n",
+ "-----------------------"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plot_posterior(samples, data, model3, diffuse, model2, pspec)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.10.6"
+ },
+ "vscode": {
+ "interpreter": {
+ "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+ }
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/demo_poisson_solution.ipynb b/demo_poisson_solution.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..98a883a8f07d02cf17c71d4d885841ff1ac136be
--- /dev/null
+++ b/demo_poisson_solution.ipynb
@@ -0,0 +1,229 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Nifty tutorial for Poisson count data\n",
+ "====================================="
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Setup\n",
+ "-----"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "import numpy as np\n",
+ "import matplotlib.pyplot as plt\n",
+ "import nifty8 as ift\n",
+ "from utils import plot_2D, load_psf, geovi_sampling, plot_posterior\n",
+ "\n",
+ "ift.random.push_sseq_from_seed(42)\n",
+ "evidences = []"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Load data and visualize\n",
+ "data = np.load('data/poisson.npz')\n",
+ "print(data['data'].shape)\n",
+ "plot_2D(data['data'], 'Data')"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "position_space = ift.RGSpace([128, 128])\n",
+ "\n",
+ "# Homogeneous poisson process\n",
+ "projection = ift.VdotOperator(ift.full(position_space, 1.)).adjoint\n",
+ "model1 = ift.FieldAdapter(projection.domain, 'hom')\n",
+ "model1 = ift.exp(5. * model1)\n",
+ "model1 = projection @ model1\n",
+ "\n",
+ "print(model1)\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Set up likelihood & PSF\n",
+ "d = ift.makeField(position_space, data['data'])\n",
+ "likelihood = ift.PoissonianEnergy(d)\n",
+ "PSF_op, psf = load_psf(position_space)\n",
+ "plot_2D(psf, 'PSF')\n",
+ "likelihood = likelihood @ PSF_op"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Inference model 1\n",
+ "samples, evidence = geovi_sampling(likelihood @ model1)\n",
+ "plot_2D(samples.average(model1).val, 'model1 posterior mean')\n",
+ "print(evidence)\n",
+ "evidences += [evidence, ]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Independent poisson process\n",
+ "model2 = ift.InverseGammaOperator(position_space, 2., 3.).ducktape('independent')\n",
+ "\n",
+ "# Inference model 2\n",
+ "samples, evidence = geovi_sampling(likelihood @ model2)\n",
+ "plot_2D(samples.average(model2).val, 'model2 posterior mean')\n",
+ "print(evidence)\n",
+ "evidences += [evidence, ]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Compare evidence\n",
+ "print(evidences)"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Diffuse poisson process\n",
+ "args = {\n",
+ " 'offset_mean': .5,\n",
+ " 'offset_std': (1., 1E-5),\n",
+ "\n",
+ " # Amplitude of field fluctuations\n",
+ " 'fluctuations': (1.5, 0.5), # 1.0, 1e-2\n",
+ "\n",
+ " # Exponent of power law power spectrum component\n",
+ " 'loglogavgslope': (-4., 1), # -6.0, 1\n",
+ "\n",
+ " # Amplitude of integrated Wiener process power spectrum component\n",
+ " 'flexibility': (1., 0.2), # 2.0, 1.0\n",
+ "\n",
+ " # How ragged the integrated Wiener process component is\n",
+ " 'asperity': (0.1, 0.01), # 0.1, 0.5\n",
+ "\n",
+ " # Name of the input keys\n",
+ " 'prefix' : 'diffuse'\n",
+ "}\n",
+ "correlated_field = ift.SimpleCorrelatedField(position_space, **args)\n",
+ "pspec = correlated_field.power_spectrum\n",
+ "diffuse = correlated_field.exp()\n",
+ "\n",
+ "model3 = diffuse + model2"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Prior samples\n",
+ "pl = ift.Plot()\n",
+ "for _ in range(9):\n",
+ " pl.add(model3(ift.from_random(model3.domain)))\n",
+ "pl.output()"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Inference model 3\n",
+ "samples, evidence = geovi_sampling(likelihood @ model3)\n",
+ "plot_2D(samples.average(model3).val, 'model3 posterior mean')\n",
+ "print(evidence)\n",
+ "evidences += [evidence, ]"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Compare evidence\n",
+ "print(evidences)"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Posterior visualization\n",
+ "-----------------------"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "plot_posterior(samples, data, model3, diffuse, model2, pspec)"
+ ]
+ }
+ ],
+ "metadata": {
+ "kernelspec": {
+ "display_name": "Python 3 (ipykernel)",
+ "language": "python",
+ "name": "python3"
+ },
+ "language_info": {
+ "codemirror_mode": {
+ "name": "ipython",
+ "version": 3
+ },
+ "file_extension": ".py",
+ "mimetype": "text/x-python",
+ "name": "python",
+ "nbconvert_exporter": "python",
+ "pygments_lexer": "ipython3",
+ "version": "3.10.6"
+ },
+ "vscode": {
+ "interpreter": {
+ "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+ }
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}
diff --git a/utils.py b/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..e84fbc182f26809450d53ec8cb535ccd9a2dff6e
--- /dev/null
+++ b/utils.py
@@ -0,0 +1,120 @@
+import nifty8 as ift
+import numpy as np
+import matplotlib.pyplot as plt
+
+def load_psf(space):
+ hsp = space.get_default_codomain()
+ psf = lambda k: 1./(1.+(k/20.)**2)/space.size
+ PD = ift.PowerDistributor(hsp)
+ psf = ift.PS_field(PD.domain[0], psf)
+ psf = PD(psf)
+ ht = ift.HartleyOperator(hsp, space)
+ pos_psf = ht(psf).val
+ pos_psf = np.roll(pos_psf, pos_psf.shape[0]//2, axis = 0)
+ pos_psf = np.roll(pos_psf, pos_psf.shape[1]//2, axis = 1)
+ pos_psf = ift.makeField(space, pos_psf)
+ R = ht @ ift.DiagonalOperator(psf) @ ht.adjoint
+ return R, pos_psf.val
+
+def plot_2D(inp, label):
+ fig, ax = plt.subplots(nrows=1, ncols=1, figsize = (10, 8))
+ im = ax.imshow(inp.T, origin='lower', extent=[0,1,0,1])
+ ax.set_title(label)
+ fig.colorbar(im, ax=ax)
+ plt.show()
+
+def geovi_sampling(likelihood):
+ ic_samp = ift.AbsDeltaEnergyController(1E-3, iteration_limit = 30)
+ ic_sampnl = ift.AbsDeltaEnergyController(0.1, iteration_limit = 20,
+ convergence_level=2)
+ mini_samp = ift.NewtonCG(ic_sampnl)
+ ic_mini = ift.AbsDeltaEnergyController(0.1, iteration_limit=15,
+ name = 'Minimizer')
+ minimizer = ift.NewtonCG(ic_mini)
+
+ N_samples = 4
+ iteration_limit = 5
+ initial_mean = 0.1 * ift.from_random(likelihood.domain)
+ samples = ift.optimize_kl(likelihood, iteration_limit, N_samples,
+ minimizer, ic_samp, mini_samp,
+ initial_position=initial_mean)
+ ev, _ = ift.estimate_evidence_lower_bound(
+ ift.StandardHamiltonian(likelihood), samples,
+ min(100, likelihood.domain.size), verbose = False)
+ return samples, ev.average().val[()]
+
+def plot_posterior(samples, data, model3, diffuse, model2, pspec):
+ fig, ax = plt.subplots(nrows=5, ncols=3, figsize = (20, 28))
+
+ ax[0,0].set_visible(False)
+ ax[0,2].set_visible(False)
+ ax[4,0].set_visible(False)
+ ax[4,2].set_visible(False)
+
+ im = ax[0,1].imshow(data['data'].T, origin='lower', extent=[0,1,0,1])
+ ax[0,1].set_title('Data')
+ fig.colorbar(im, ax=ax[0,1])
+
+ im = ax[1,0].imshow(data['sky'].T, origin='lower', extent=[0,1,0,1])
+ ax[1,0].set_title('Ground truth')
+ fig.colorbar(im, ax=ax[1,0])
+
+ im = ax[1,1].imshow(samples.average(model3.force).val.T, origin='lower', extent=[0,1,0,1])
+ ax[1,1].set_title('Posterior mean')
+ fig.colorbar(im, ax=ax[1,1])
+ sm = tuple(s for s in samples.iterator(model3.force))
+ im = ax[1,2].imshow(sm[0].val.T, origin='lower', extent=[0,1,0,1])
+ ax[1,2].set_title('Posterior sample')
+ fig.colorbar(im, ax=ax[1,2])
+
+
+ im = ax[2,0].imshow(data['points'].T, origin='lower', extent=[0,1,0,1])
+ ax[2,0].set_title('Ground truth (sources)')
+ fig.colorbar(im, ax=ax[2,0])
+
+ im = ax[2,1].imshow(samples.average(model2.force).val.T, origin='lower', extent=[0,1,0,1])
+ ax[2,1].set_title('Posterior mean (sources)')
+ fig.colorbar(im, ax=ax[2,1])
+
+ sm = tuple(s for s in samples.iterator(model2.force))
+ im = ax[2,2].imshow(sm[0].val.T, origin='lower', extent=[0,1,0,1])
+ ax[2,2].set_title('Posterior sample (sources)')
+ fig.colorbar(im, ax=ax[2,2])
+
+
+ im = ax[3,0].imshow(data['diffuse'].T, origin='lower', extent=[0,1,0,1])
+ ax[3,0].set_title('Ground truth (diffuse)')
+ fig.colorbar(im, ax=ax[3,0])
+
+ im = ax[3,1].imshow(samples.average(diffuse.force).val.T, origin='lower', extent=[0,1,0,1])
+ ax[3,1].set_title('Posterior mean (diffuse)')
+ fig.colorbar(im, ax=ax[3,1])
+
+ sm = tuple(s for s in samples.iterator(diffuse.force))
+ im = ax[3,2].imshow(sm[0].val.T, origin='lower', extent=[0,1,0,1])
+ ax[3,2].set_title('Posterior sample (diffuse)')
+ fig.colorbar(im, ax=ax[3,2])
+
+ ax = ax[4,1]
+ ks = pspec.target[0].k_lengths
+ lbl = 'posterior samples'
+ for i,s in enumerate(samples.iterator()):
+ ss = pspec.force(s)
+ ax.plot(ks[1:], ss.val[1:], color='k', alpha = 0.4,
+ label=lbl)
+ lbl=None
+
+ ax.plot(ks[1:],data['pspec'][1:], color = 'g',label='ground truth')
+ mm = samples.average((pspec.log()).force)
+ ax.plot(ks[1:], mm.exp().val[1:], color='r',label='posterior mean')
+ #ax.set_ylim([1E-2*np.min(pp.val[1:]), 10*np.max(pp.val[1:])])
+ ax.set_xscale('log')
+ ax.set_yscale('log')
+ ax.set_xlabel(r'$|k|$')
+ ax.set_ylabel(r'$P_s\left(|k|\right)$')
+ ax.set_title('Power spectrum')
+ leg = ax.legend()
+ #for lh in leg.legendHandles:
+ # lh.set_alpha(1)
+ fig.tight_layout()
+ plt.show()