diff --git a/demo_radio_blank.ipynb b/demo_radio_blank.ipynb
new file mode 100644
index 0000000000000000000000000000000000000000..3f153bf723b2240b68f932315f06446dd0503c94
--- /dev/null
+++ b/demo_radio_blank.ipynb
@@ -0,0 +1,530 @@
+{
+ "cells": [
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Nifty tutorial for radio interferometric imaging\n",
+ "================================================"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Setup\n",
+ "-----"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "#import packages\n",
+ "import numpy as np\n",
+ "import matplotlib.pyplot as plt\n",
+ "import nifty8 as ift\n",
+ "import resolve as rve"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Load mock data with EHT measurement configuration and visualize uv-plane"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# load the observations\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# plot uv\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Set up the FOV and number of pixels of the `nifty` space which defines the image plane."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# set the image parameters\n",
+ "# Number of pixels\n",
+ "\n",
+ "\n",
+ "# Field of view in x and y directon in myas\n",
+ "\n",
+ "\n",
+ "# the space\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Model of the sky brightness distribution\n",
+ "----------------------------------------\n",
+ "\n",
+ "The sky model is going to be a log-normal random process (the log-brightness is Gaussian distributed). As the prior correlation structure of the log-brightness is unkown, it will be generated using a `CorrelatedField` model and the power-spectrum is inferred along with the realization."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Hyperparameters (mean and std pairs) for prior models of parameters\n",
+ "args = {# Overall offset from zero of Gaussian process\n",
+ " 'offset_mean': -np.log(space.scalar_dvol) - 10.,\n",
+ "\n",
+ " # Variability of inferred offset amplitude\n",
+ " 'offset_std': (3., 1.),\n",
+ "\n",
+ " # Amplitude of field fluctuations\n",
+ " 'fluctuations': (1.5, 0.5),\n",
+ "\n",
+ " # Exponent of power law power spectrum component\n",
+ " 'loglogavgslope': (-4., .5),\n",
+ "\n",
+ " # Amplitude of integrated Wiener process power spectrum component\n",
+ " 'flexibility': (.3, .1),\n",
+ "\n",
+ " # How ragged the integrated Wiener process component is\n",
+ " 'asperity': None # Passing 'None' disables this part of the model\n",
+ " }"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# set the signal\n",
+ "\n",
+ "\n",
+ "# Exponentiate to recieve a log-normal distributed sky model\n",
+ "\n",
+ "\n",
+ "# Save the model of the power-spectrum for visualization\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# print some properties\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Prior samples\n",
+ "-------------\n",
+ "\n",
+ "To get a feeling for the prior variability of the sky model we generate several random realizations of the process and visualize them. The power-spectra, generated for each process realization, are depicted in the last panel."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# some prior samples\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Alternatively, one can define a log-normal process with a fixed, user-specified spectrum. Note that in case the spectrum does not match the true underlying data-generating process, this may yield suboptimal results!"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# fixed spectrum model\n",
+ "use_fixed_spectrum = False\n",
+ "\n",
+ "def fixed_spectrum_model(space):\n",
+ " hp = space.get_default_codomain()\n",
+ " a, k0 = 1e-36, 1E9\n",
+ " ker = ift.PS_field(ift.PowerSpace(hp), (lambda k: a / (1.+ (k/k0)**6)))\n",
+ " amplitude = ift.PowerDistributor(hp, power_space=ker.domain[0])(ker.sqrt())\n",
+ "\n",
+ " HT = ift.HarmonicTransformOperator(hp, target=space)\n",
+ " A = ift.DiagonalOperator(amplitude)\n",
+ "\n",
+ " signal = ift.FieldAdapter(hp, 'xi')\n",
+ " signal = ift.exp(HT @ A @ signal)\n",
+ " return (1e-4/signal.target.scalar_weight()) * signal\n",
+ "\n",
+ "if use_fixed_spectrum:\n",
+ " sky_model = fixed_spectrum_model(space)\n",
+ " pspec = None"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Setup a mock VLBI imaging task using the `InterferometryResponse` of `resolve`. The `mock_observation` contains all information relevant to set up the likelihood, including visibility data, uv-coordinates, and the noise levels of each measurement. The `measurement_sky` contains all relevant information regarding the prior model of the sky brightness distribution. The additional parameters passed to the `InterferometryResponse` control the accuracy and behaviour of the `wgridder` used within `resolve` which defines the response function."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# adapting to resolve standards\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Set up response, the additional response from sky brightness to visibilities, radio measurement equation\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Extract data and noise level from observation\n",
+ "\n",
+ "\n",
+ "# Build gaussian likelihood and apply it to the sky model\n"
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# just some properties\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Solve the inference problem\n",
+ "---------------------------\n",
+ "\n",
+ "The `likelihood` together with the `sky_model` fully specify a Bayesian inverse problem and imply a posterior probabiltiy distribution over the degrees of freedom (DOF) of the model. This distribution is in general a high-dimensional (number of pixels + DOF of power spectrum) and non-Gaussian distribution which prohibits analytical integration. To access its information and compute posterior expectation values numerical approximations have to be made.\n",
+ "\n",
+ "`nifty` provides multiple ways of psterior approximation, with Variational Inference (VI) being by far the most frequently used method. In VI the posterior distribution is approximated with another distribution by minimizing their respective forward Kullbach-Leibler divergence (KL). In the following, the Geometric VI method is employed which utilizes conceps of differential geometry to provide a local estimate of the distribution function.\n",
+ "\n",
+ "Its numerical implementation (`ift.optimize_kl`) consists of a repeated and successive re-approximation of the VI objective function (the KL) via a stochastic estimate. This estimate is generated using the at the time best available approximation of the posterior, and then the KL gets minimized to further improve it. The resulting algorithm cosists of a repeated re-generation of novel samples for the estimate and a successing optimization thereof until convergence is reached.\n",
+ "\n",
+ "The internal steps of `ift.optimize_kl` invoke the approximate solution of multiple interdependent optimization problems:\n",
+ "- For sample generation, a linear system of eqations is approximated using the `ConjugateGradient` (CG) method\n",
+ "- Furthermore, the sample generation invokes a non-linear optimization problem approximated using the `NewtonCG` method\n",
+ "- Finally, the approximative distribution is optimized by minimizing the KL between the true posterior and the approximation. This again invokes a non-linear optimization problem approximated with `NewtonCG`."
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Posterior visualization\n",
+ "-----------------------\n",
+ "\n",
+ "Before we set run the minimization routine, we set up a `plotting_callback` function for visualization. Note that additional information and plots regarding the reconstruction are generated during an `ift.optimize_kl` run and stored in the folder passed to the `output_directory` argument of `ift.optimize_kl`\n",
+ "The final output of `ift.optimize_kl` is a collection of approximate posterior samples and is provided via an instance of `ift.ResidualSampleList`. A `SampleList` provides a variety of convenience functions such as: \n",
+ "- `average`: to compute samples averages\n",
+ "- `sample_stat`: to get the approximate mean and variance of a model\n",
+ "- `iterator`: a python iterator over all samples\n",
+ "- ..."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "from IPython.display import clear_output\n",
+ "\n",
+ "def _imshow(figure, field, ax, title, vmin = 0, vmax = None, cmap='afmhot'):\n",
+ " im0 = ax.imshow(field.val.T, origin = 'lower', extent = [0, x_fov, 0, y_fov], cmap=cmap,\n",
+ " vmin = vmin, vmax = vmax)\n",
+ " figure.colorbar(im0, ax=ax)\n",
+ " ax.set_xlabel(r'x $\\left(\\mu as\\right)$')\n",
+ " ax.set_ylabel(r'y $\\left(\\mu as\\right)$')\n",
+ " ax.set_title(title)\n",
+ "\n",
+ "def _plot_histogram(nodes, hist, ax, title, ):\n",
+ " nodes = 0.5*(nodes[1:] + nodes[:-1])\n",
+ " ax.bar(nodes, hist)\n",
+ " rs = np.arange(nodes[0], nodes[-1], 0.1)\n",
+ " gauss = np.exp(-0.5*rs**2)/np.sqrt(2*np.pi)\n",
+ " ax.plot(rs, gauss, 'k--', label = r'standard Gauss')\n",
+ " ax.set_xlabel(r'$r$')\n",
+ " ax.set_ylabel(r'$P(r)$')\n",
+ " ax.set_title(title)\n",
+ " ax.legend()\n",
+ " ax.set_xlim([nodes[0], nodes[-1]])\n",
+ "\n",
+ "def plotting_callback(samples):\n",
+ " clear_output(wait=True) \n",
+ "\n",
+ " sky_mean, sky_var = samples.sample_stat(sky_model)\n",
+ "\n",
+ " pspec_mean = samples.average(pspec.log().force).exp()\n",
+ " pspecs = samples.iterator(pspec.force)\n",
+ "\n",
+ " residual = ift.Adder(mock_observation.vis, neg=True) @ R @ measurement_sky\n",
+ " residual = ift.makeOp(mock_observation.weight.sqrt()) @ residual\n",
+ " nbins = 50\n",
+ " hist = np.zeros(nbins)\n",
+ " for rr in samples.iterator(residual):\n",
+ " rr = rr.val.flatten()\n",
+ " rr = np.concatenate((rr.real, rr.imag))\n",
+ " wgt, nodes = np.histogram(rr, nbins, range=[-5, 5])\n",
+ " hist += wgt/wgt.sum()/(nodes[1]-nodes[0])\n",
+ " hist /= samples.n_samples\n",
+ "\n",
+ "\n",
+ "\n",
+ " fig, axs = plt.subplots(nrows=2, ncols=2, figsize=(15,13))\n",
+ " axs = axs.flatten()\n",
+ "\n",
+ " _imshow(fig, sky_mean, axs[0], 'Sky brightness mean')\n",
+ " _imshow(fig, sky_var.sqrt(), axs[1], 'Sky brightness std', cmap = 'afmhot')\n",
+ "\n",
+ " axs[1].yaxis.set_visible(False)\n",
+ " k_lengths = pspec_mean.domain[0].k_lengths[1:]\n",
+ " lbl = 'samples'\n",
+ " for ps in pspecs:\n",
+ " axs[2].plot(k_lengths, ps.val[1:], 'k-', alpha = 0.5, label = lbl)\n",
+ " lbl = None\n",
+ " axs[2].plot(k_lengths, pspec_mean.val[1:], 'r-', label = 'mean')\n",
+ " axs[2].set_xlim([k_lengths[0], k_lengths[-1]])\n",
+ " axs[2].set_xscale('log')\n",
+ " axs[2].set_yscale('log')\n",
+ " axs[2].set_xlabel(r'$|k|$')\n",
+ " axs[2].set_ylabel(r'$P_s\\left(|k|\\right)$')\n",
+ " axs[2].set_title(r'Power-spectrum of log-sky brightness')\n",
+ " axs[2].legend()\n",
+ "\n",
+ " _plot_histogram(nodes, hist, axs[3], \n",
+ " r'Inverse noise weighted data residual ($r$) distribution ($P(r)$)')\n",
+ "\n",
+ " fig.tight_layout()\n",
+ " plt.show();"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "WARNING: The entire reconstruction takes a few minutes to run."
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "# Minimization parameters and minimizers for the VI algorithm\n",
+ "ic_sampling = ift.AbsDeltaEnergyController(deltaE=0.05, iteration_limit=50)\n",
+ "\n",
+ "ic_sampling_nl = ift.AbsDeltaEnergyController(deltaE=0.5,\n",
+ " iteration_limit=5,\n",
+ " convergence_level=1)\n",
+ "\n",
+ "\n",
+ "ic_newton_early = ift.AbsDeltaEnergyController(name='Newton',\n",
+ " deltaE=0.1,\n",
+ " iteration_limit=5,\n",
+ " convergence_level=2)\n",
+ "\n",
+ "ic_newton_late = ift.AbsDeltaEnergyController(name='Newton',\n",
+ " deltaE=0.1,\n",
+ " iteration_limit=10,\n",
+ " convergence_level=2)\n",
+ "\n",
+ "minimizer_early = ift.NewtonCG(ic_newton_early)\n",
+ "minimizer_late = ift.NewtonCG(ic_newton_late)\n",
+ "\n",
+ "\n",
+ "\n",
+ "# Minimize KL between true posterior and approximation. Each iteration includes sample\n",
+ "# generation and optimization.\n",
+ "n_iterations = 15 # Total number of iterations. \n",
+ "n_samples = (lambda iiter: 2 if iiter < 10 else 5) # Number of samples used for KL approximation\n",
+ "\n",
+ "minimizer = (lambda iiter: minimizer_early if iiter < 10 else minimizer_late) # When to use which optimizer\n",
+ "minimizer_sampling = (lambda iiter: None if iiter < 10 else ift.NewtonCG(ic_sampling_nl)) # initially MGVI, \n",
+ " # later geoVI\n",
+ "samples = ift.optimize_kl(likelihood, n_iterations, n_samples, minimizer,\n",
+ " ic_sampling, minimizer_sampling,\n",
+ " inspect_callback=plotting_callback,\n",
+ " output_directory=\"mock_inference_5\")\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "Comparison of posterior to ground truth\n",
+ "======================================="
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "mock_sky = np.load('data/mock_signal.npz')['signal']\n",
+ "mock_sky = ift.makeField(space, mock_sky)\n",
+ "\n",
+ "sky_mean, sky_var = samples.sample_stat(sky_model)\n",
+ "sky_samples = list(s for s in samples.iterator(sky_model))\n",
+ "\n",
+ "fig, axs = plt.subplots(nrows=3, ncols=2, figsize = (15,18))\n",
+ "_imshow(fig, mock_sky, axs[0,0], 'Sky brightness ground truth', vmax = 2e19)\n",
+ "_imshow(fig, sky_mean, axs[0,1], 'Sky brightness mean', vmax = 2e19)\n",
+ "_imshow(fig, sky_var.sqrt()/sky_mean, axs[1,0], 'Sky brightness relative uncertainty')\n",
+ "_imshow(fig, sky_samples[0], axs[1,1], 'Sky brightness posterior sample (1)', vmax = 2e19)\n",
+ "_imshow(fig, sky_samples[1], axs[2,0], 'Sky brightness posterior sample (2)', vmax = 2e19)\n",
+ "_imshow(fig, sky_samples[2], axs[2,1], 'Sky brightness posterior sample (3)', vmax = 2e19)\n",
+ "fig.tight_layout()\n"
+ ]
+ },
+ {
+ "cell_type": "markdown",
+ "metadata": {},
+ "source": [
+ "UV - Uncertainty map\n",
+ "--------------------\n",
+ "\n",
+ "We can generate and study the posterior uncertainty maps for any kind on quantity we are interested in (i.E. not only for the sky brightness). In particular, we can also take a look at the uncertainty of the sky brightness in the UV plane. Comparing this to the measured UV-tracks yields information "
+ ]
+ },
+ {
+ "cell_type": "code",
+ "execution_count": null,
+ "metadata": {},
+ "outputs": [],
+ "source": [
+ "lim_uv = 1.2E7\n",
+ "u = np.linspace(-lim_uv, lim_uv, num=257)\n",
+ "v = np.linspace(-lim_uv, lim_uv, num=257)\n",
+ "uu, vv = np.meshgrid(u,v)\n",
+ "ww = np.zeros_like(uu)\n",
+ "uvw = np.stack(tuple(a.flatten() for a in (uu,vv,ww)), axis = -1)\n",
+ "grid_antennas = rve.AntennaPositions(uvw)\n",
+ "vis = np.zeros(uu.size, dtype=np.complex128).reshape((1,-1,1))\n",
+ "wgts = np.ones_like(vis)\n",
+ "grid_obs = rve.Observation(grid_antennas, vis, wgts, rve.Polarization.trivial(),\n",
+ " mock_observation.freq)\n",
+ "grid_R = rve.InterferometryResponse(grid_obs, measurement_sky.target, \n",
+ " do_wgridding=False, epsilon=1e-10)\n",
+ "def intensity_uv(inp):\n",
+ " r = (grid_R @ pre_response)(inp)\n",
+ " r = (r.real**2 + r.imag**2).sqrt()\n",
+ " return r.val.reshape(uu.shape)\n",
+ "\n",
+ "uv_mean, uv_var = samples.sample_stat(lambda xi: intensity_uv(sky_model(xi)))\n",
+ "uv_gt = intensity_uv(mock_sky)\n",
+ "\n",
+ "def uv_plot(inp, pre = \"\"):\n",
+ " f, ax = plt.subplots(ncols = 2, figsize = (15,10))\n",
+ " ax[0].imshow(inp, origin='lower', extent=(-lim_uv, lim_uv, -lim_uv, lim_uv))\n",
+ " ax[1].imshow(inp, origin='lower', extent=(-lim_uv, lim_uv, -lim_uv, lim_uv))\n",
+ " ax[1].scatter(mock_observation.uvw[:,0], mock_observation.uvw[:,1],\n",
+ " c='r', marker='.', s = 12, alpha = 0.5, label = 'measurements')\n",
+ " ax[1].scatter(-mock_observation.uvw[:,0], -mock_observation.uvw[:,1],\n",
+ " c='c', marker='.', s = 10, alpha = 0.5, label = 'measurements (c.c.)')\n",
+ " ax[1].legend()\n",
+ " for a in ax:\n",
+ " a.set_xlabel('u')\n",
+ " a.set_ylabel('v')\n",
+ " ax[0].set_title(pre+'UV posterior uncertainty')\n",
+ " ax[1].set_title(pre+'UV posterior uncertainty & measured UV-tracks')\n",
+ " fig.tight_layout()\n",
+ " plt.show()\n",
+ "\n",
+ "uv_plot(np.log(uv_var))\n",
+ "\n",
+ "f, ax = plt.subplots(ncols = 2, figsize = (15,10))\n",
+ "ax[0].imshow(np.log(uv_gt), origin='lower',\n",
+ " extent=(-lim_uv, lim_uv, -lim_uv, lim_uv))\n",
+ "ax[1].imshow(np.log(uv_mean), origin='lower',\n",
+ " extent=(-lim_uv, lim_uv, -lim_uv, lim_uv))\n",
+ "for a in ax:\n",
+ " a.set_xlabel('u')\n",
+ " a.set_ylabel('v')\n",
+ "ax[0].set_title('true UV intensity')\n",
+ "ax[1].set_title('rec. UV intensity')\n",
+ "fig.tight_layout()\n",
+ "plt.show()\n",
+ "\n",
+ "uv_plot(np.log(uv_var / uv_gt), pre = 'relative ')\n"
+ ]
+ }
+ ],
+ "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.12"
+ },
+ "vscode": {
+ "interpreter": {
+ "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+ }
+ }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}