{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "5d94300e",
   "metadata": {
    "tags": []
   },
   "source": [
    "# Image Deconvolution for 511 keV Extended Sources\n",
    "\n",
    "updated on 2025-04-11\n",
    "\n",
    "This notebook introduces COSI's image deconvolution with the Compton data space (CDS) in the Galactic coordinate system. An example of the image analysis will be presented using the 511keV thin disk 3-month simulation data created for DC2.\n",
    "\n",
    "We have two options on the coordinate system to describe the Compton scattering direction ($\\psi\\chi$) in the image deconvolution, namely Galactic coordinates or detector coordinates. Using Galactic coordinates is intuitive, and the spectral fitting adopts this convention. However, image deconvolution in Galactic coordinates take significant computation time since the detector response needs to be converted into Galactic coordinates for each sky pixel. To reduce the burden on the end-user, a pre-computed converted response is provided in DC2 and DC3 for several main sources (511 keV, Al-26, Ti-44, continuum). These pre-computed responses assume that we analyze 3-month data without extracting time intervals, and the pixel resolution of the model map is already fixed. While there is less flexibility in binning/modeling, it is relatively fast to perform the image deconvolution since the most computationally heavy part, the coordinate conversion of the response, can be skipped.\n",
    "\n",
    "You can also check out the related notebook [511keV-ScAtt-DataReduction.ipynb](https://github.com/cositools/cosipy/blob/main/docs/tutorials/image_deconvolution/optional/scatt_binned_data_w_local_coordinate/511keV-ScAtt-DataReduction.ipynb) for imaging in detector coordinates."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb396370",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "import sys\n",
    "logger = logging.getLogger('cosipy')\n",
    "logger.setLevel(logging.INFO)\n",
    "logger.addHandler(logging.StreamHandler(sys.stdout))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7c55c698",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "\n",
    "# cosipy uses astropy units\n",
    "import astropy.units as u\n",
    "\n",
    "import healpy as hp\n",
    "from mhealpy import HealpixMap\n",
    "\n",
    "from histpy import Histogram\n",
    "\n",
    "from cosipy.data_io import BinnedData\n",
    "from cosipy.image_deconvolution import  DataIF_COSI_DC2, ImageDeconvolution\n",
    "from cosipy.util import fetch_wasabi_file\n",
    "\n",
    "%matplotlib inline"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "15576714",
   "metadata": {
    "tags": []
   },
   "source": [
    "## 0. Files needed for this notebook\n",
    "\n",
    "From wasabi\n",
    "- cosi-pipeline-public/COSI-SMEX/develop/data/Responses/PointSourceResponse/psr_gal_511_DC2.h5\n",
    "    - a pre-computed 511 keV line response file converted into the Galactic coordinate system\n",
    "- cosi-pipeline-public/COSI-SMEX/DC2/Data/Sources/511_thin_disk_3months_unbinned_data.fits.gz\n",
    "- cosi-pipeline-public/COSI-SMEX/DC2/Data/Backgrounds/albedo_photons_3months_unbinned_data.fits.gz\n",
    "    - In this notebook, only the albedo gamma-ray background is considered for a tutorial.\n",
    "    - If you want to consider all of the background components, please replace it with cosi-pipeline-public/COSI-SMEX/DC2/Data/Backgrounds/total_bg_3months_unbinned_data.fits.gz\n",
    "    - Note that total_bg_3months_unbinned_data.fits.gz is 14.15 GB.\n",
    "\n",
    "From docs/tutorials/image_deconvolution/511keV/GalacticCDS\n",
    "- inputs_511keV_DC2.yaml\n",
    "- imagedeconvolution_parfile_gal_511keV.yml"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4a9f31fe",
   "metadata": {},
   "source": [
    "You can download the data and detector response from wasabi. You can skip this cell if you already have downloaded the files."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "559b3250",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Response file:\n",
    "# wasabi path: COSI-SMEX/develop/Data/Responses/PointSourceResponse/psr_gal_511_DC2.h5\n",
    "# File size: 3.82 GB\n",
    "fetch_wasabi_file('COSI-SMEX/develop/Data/Responses/PointSourceResponse/psr_gal_511_DC2.h5', checksum='50eb36c2bc0089ba230af457ec768fa0')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3d2f6d62",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Source file (511 keV thin disk model):\n",
    "# wasabi path: COSI-SMEX/DC2/Data/Sources/511_thin_disk_3months_unbinned_data.fits.gz\n",
    "# File size: 202.45 MB\n",
    "fetch_wasabi_file('COSI-SMEX/DC2/Data/Sources/511_thin_disk_3months_unbinned_data.fits.gz', checksum='e0db59ba35aabd0f3b9371b85986841d')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a49e6f32",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Background file (albedo gamma):\n",
    "# wasabi path: COSI-SMEX/DC2/Data/Backgrounds/albedo_photons_3months_unbinned_data.fits.gz\n",
    "# File size: 2.69 GB\n",
    "fetch_wasabi_file('COSI-SMEX/DC2/Data/Backgrounds/albedo_photons_3months_unbinned_data.fits.gz', checksum='4c125410d8f127d7e12682f008d5651d')"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee377e6d",
   "metadata": {},
   "source": [
    "## 1. Create binned event/background files in the Galactic coordinate system\n",
    "\n",
    "<span style=\"color:red\"> please modify \"path_data\" corresponding to your environment.</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "87405a24",
   "metadata": {},
   "outputs": [],
   "source": [
    "path_data = \"./\"\n",
    "# example\n",
    "# path_data = \"/Users/ckierans/Software/COSItools/COSItools/cosipy/docs/tutorials/image_deconvolution/511keV/GalacticCDS/\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c50857da",
   "metadata": {},
   "source": [
    "**Source**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "791a88e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "signal_filepath = path_data + \"511_thin_disk_3months_unbinned_data.fits.gz\"\n",
    "\n",
    "binned_signal = BinnedData(input_yaml = \"inputs_511keV_DC2.yaml\")\n",
    "\n",
    "binned_signal.get_binned_data(unbinned_data = signal_filepath, psichi_binning=\"galactic\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1db6990e",
   "metadata": {},
   "source": [
    "**Background**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5eb86e9c",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "bkg_filepath = path_data + \"albedo_photons_3months_unbinned_data.fits.gz\"\n",
    "\n",
    "binned_bkg = BinnedData(input_yaml = \"inputs_511keV_DC2.yaml\")\n",
    "\n",
    "binned_bkg.get_binned_data(unbinned_data = bkg_filepath, psichi_binning=\"galactic\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e95224a",
   "metadata": {},
   "source": [
    "Convert the data into sparse matrices & add the signal to the background"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d94e1354",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = binned_signal.binned_data.to_dense()\n",
    "bkg = binned_bkg.binned_data.to_dense()\n",
    "event = signal + bkg"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "90b9a1e6",
   "metadata": {},
   "source": [
    "Save the binned histograms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ae234042",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal.write(\"511keV_dc2_galactic_signal.hdf5\", overwrite = True)\n",
    "bkg.write(\"511keV_dc2_galactic_bkg.hdf5\", overwrite = True)\n",
    "event.write(\"511keV_dc2_galactic_event.hdf5\", overwrite = True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "027e73c3",
   "metadata": {},
   "source": [
    "Load the saved files"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "60093ee2",
   "metadata": {},
   "outputs": [],
   "source": [
    "signal = Histogram.open(\"511keV_dc2_galactic_signal.hdf5\")\n",
    "bkg = Histogram.open(\"511keV_dc2_galactic_bkg.hdf5\")\n",
    "event = Histogram.open(\"511keV_dc2_galactic_event.hdf5\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "129a8156",
   "metadata": {},
   "source": [
    "In DC2, the number of time bins should be 1 when you perform the image deconvolution using the Galactic CDS.\n",
    "It is because the pre-computed response files in the galactic coordinate have no time axis, and all of the events are assumed to be projected into a single Galactic CDS.\n",
    "In the future, we plan to introduce more flexible binning."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7857047",
   "metadata": {},
   "outputs": [],
   "source": [
    "bkg.axes['Time'].edges"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ac17634",
   "metadata": {},
   "source": [
    "## 2. Load the response matrix"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8e536819",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "response_path = path_data + \"psr_gal_511_DC2.h5\"\n",
    "\n",
    "image_response = Histogram.open(response_path)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "676ba432",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_response.axes.labels"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5c4b744e",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_response.contents.shape"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6396be1e",
   "metadata": {},
   "source": [
    "## 3. Imaging deconvolution"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7c5da67f",
   "metadata": {
    "jp-MarkdownHeadingCollapsed": true
   },
   "source": [
    "### Brief overview of the image deconvolution\n",
    "\n",
    "Basically, we have to maximize the following likelihood function\n",
    "\n",
    "$$\n",
    "\\log L = \\sum_i X_i \\log \\epsilon_i - \\sum_i \\epsilon_i\n",
    "$$\n",
    "\n",
    "$X_i$: detected counts at $i$-th bin ( $i$ : index of the Compton Data Space)\n",
    "\n",
    "$\\epsilon_i = \\sum_j R_{ij} \\lambda_j + b_i$ : expected counts ( $j$ : index of the model space)\n",
    "\n",
    "$\\lambda_j$ : the model map (basically gamma-ray flux at $j$-th pixel)\n",
    "\n",
    "$b_i$ : the background at $i$-th bin\n",
    "\n",
    "$R_{ij}$ : the response matrix\n",
    "\n",
    "Since we have to optimize the flux in each pixel, and the number of parameters is large, we adopt an iterative approach to find a solution of the above equation. The simplest one is the ML-EM (Maximum Likelihood Expectation Maximization) algorithm. It is also known as the Richardson-Lucy algorithm.\n",
    "\n",
    "$$\n",
    "\\lambda_{j}^{k+1} = \\lambda_{j}^{k} + \\delta \\lambda_{j}^{k}\n",
    "$$\n",
    "$$\n",
    "\\delta \\lambda_{j}^{k} = \\frac{\\lambda_{j}^{k}}{\\sum_{i} R_{ij}} \\sum_{i} \\left(\\frac{ X_{i} }{\\epsilon_{i}} - 1 \\right) R_{ij} \n",
    "$$\n",
    "\n",
    "We refer to $\\delta \\lambda_{j}^{k}$ as the delta map.\n",
    "\n",
    "As for now, the two improved algorithms are implemented in COSIpy.\n",
    "\n",
    "- Accelerated ML-EM algorithm (Knoedlseder+99)\n",
    "\n",
    "$$\n",
    "\\lambda_{j}^{k+1} = \\lambda_{j}^{k} + \\alpha^{k} \\delta \\lambda_{j}^{k}\n",
    "$$\n",
    "$$\n",
    "\\alpha^{k} < \\mathrm{max}(- \\lambda_{j}^{k} / \\delta \\lambda_{j}^{k})\n",
    "$$\n",
    "\n",
    "Practically, in order not to accelerate the algorithm excessively, we set the maximum value of $\\alpha$ ($\\alpha_{\\mathrm{max}}$). Then, $\\alpha$ is calculated as:\n",
    "\n",
    "$$\n",
    "\\alpha^{k} = \\mathrm{min}(\\mathrm{max}(- \\lambda_{j}^{k} / \\delta \\lambda_{j}^{k}), \\alpha_{\\mathrm{max}})\n",
    "$$\n",
    "\n",
    "- Noise damping using gaussian smoothing (Knoedlseder+05, Siegert+20)\n",
    "\n",
    "$$\n",
    "\\lambda_{j}^{k+1} = \\lambda_{j}^{k} + \\alpha^{k} \\left[ w_j \\delta \\lambda_{j}^{k} \\right]_{\\mathrm{gauss}}\n",
    "$$\n",
    "$$\n",
    "w_j = \\left(\\sum_{i} R_{ij}\\right)^\\beta\n",
    "$$\n",
    "\n",
    "$\\left[ ... \\right]_{\\mathrm{gauss}}$ means that the differential image is smoothed by a gaussian filter."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "dd2adde9",
   "metadata": {},
   "source": [
    "### 3-1. Prepare DataInterface containing all necessary datasets"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6af931c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "data_interface = DataIF_COSI_DC2.load(name = \"511keV\",\n",
    "                                      event_binned_data = event.project(['Em', 'Phi', 'PsiChi']),\n",
    "                                      dict_bkg_binned_data = {\"albedo\": bkg.project(['Em', 'Phi', 'PsiChi'])},\n",
    "                                      rsp = image_response,\n",
    "                                      coordsys_conv_matrix=None)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f9d2395e",
   "metadata": {},
   "source": [
    "### 3-2. Initialize the instance of the image deconvolution class\n",
    "\n",
    "First, we prepare an instance of the ImageDeconvolution class and then register the dataset and parameters for the deconvolution. After that, you can start the calculation."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "47778146",
   "metadata": {},
   "source": [
    "<span style=\"color:red\"> please modify this parameter_filepath corresponding to your environment.</span>"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3f533ea1",
   "metadata": {},
   "outputs": [],
   "source": [
    "parameter_filepath = \"imagedeconvolution_parfile_gal_511keV.yml\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e97c5b9a",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_deconvolution = ImageDeconvolution()\n",
    "\n",
    "# set data_interface to image_deconvolution\n",
    "image_deconvolution.set_dataset([data_interface])\n",
    "\n",
    "# set a parameter file for the image deconvolution\n",
    "image_deconvolution.read_parameterfile(parameter_filepath)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "32a7aae2",
   "metadata": {},
   "source": [
    "### 3-3. Initialize image_deconvolution\n",
    "\n",
    "In this process, a model map is defined following the input parameters, and it is initialized. Also, it prepares ancillary data for the image deconvolution, e.g., the expected counts with the initial model map, gaussian smoothing filter etc.\n",
    "\n",
    "I describe parameters in the parameter file.\n",
    "\n",
    "#### model_definition:class\n",
    "\n",
    "Specify the name of the image class that will be used in the image deconvolution. As for now, it should be \"AllSkyImage\".\n",
    "\n",
    "#### model_definition:property\n",
    "\n",
    "| Name | Unit | Description | Notes |\n",
    "| :---: | :---: | :---: | :---: |\n",
    "| coordinate | str | the coordinate system of the model map | As for now, it must be 'galactic' |\n",
    "| nside | int | NSIDE of the model map | it must be the same as NSIDE of 'lb' axis of the coordinate conversion matrix|\n",
    "| scheme | str | SCHEME of the model map | As for now, it must be 'ring' |\n",
    "| energy_edges:value | list of float | The definition of the energy bins of the model map | As for now, it must be the same as that of the response matrix |\n",
    "| energy_edges:unit | str | The physical unit of the above list | In most of cases, it is keV. |\n",
    "| unit | str | The physical unit of the image | As for now, it should be \"\"cm-2 s-1 sr-1\" |\n",
    "\n",
    "#### model_definition:initialization\n",
    "\n",
    "| Name | Unit | Description | Notes |\n",
    "| :---: | :---: | :---: | :---: |\n",
    "| algorithm | str | the method name to initialize the model map | As for now, only 'flat' can be used |\n",
    "| parameter_flat:value | list of float | the list of photon fluxes for each energy band | the length of the list should be the same as the length of \"energy_edges\" - 1 |\n",
    "| parameter_flat:unit | str | The physical unit of the above list | As for now, it should be \"\"cm-2 s-1 sr-1\" |\n",
    "\n",
    "#### deconvolution\n",
    "\n",
    "| Name | Unit | Description | Notes |\n",
    "| :---: | :---: | :---: | :---: |\n",
    "|algorithm | str | the name of the image deconvolution algorithm| As for now, only 'RL' is supported |\n",
    "|||||\n",
    "|parameter_RL:iteration_max | int | The maximum number of the iteration | |\n",
    "|parameter_RL:acceleration:activate | bool | whether the accelerated ML-EM algorithm (Knoedlseder+99) is used | |\n",
    "|parameter_RL:acceleration:accel_factor_max | float | the maximum value for the acceleration parameter | |\n",
    "|parameter_RL:response_weighting:activate | bool | whether a delta map is renormalized based on the exposure time on each pixel, namely $w_j = (\\sum_{i} R_{ij})^{\\beta}$ (see Knoedlseder+05, Siegert+20) | |\n",
    "|parameter_RL:response_weighting:index | float | $\\beta$ in the above equation | |\n",
    "|parameter_RL:smoothing:activate | bool | whether a Gaussian filter is used (see Knoedlseder+05, Siegert+20)  | |\n",
    "|parameter_RL:smoothing:FWHM:value | float | the FWHM of the Gaussian in the filter | |\n",
    "|parameter_RL:smoothing:FWHM:unit | str | the unit of the above value. deg. or rad. | |\n",
    "|parameter_RL:background_normalization:activate | bool | whether the background normalization factor is optimized at each iteration | As for now, the single background normalization factor is used in all of the bins |\n",
    "|parameter_RL:background_normalization:range | list of float | the range of the normalization factor | should be positive |\n",
    "|parameter_RL:stopping_criteria:statistics | str | the statistics to be checked | In most of cases, it is \"log-likelihood\" |\n",
    "|parameter_RL:stopping_criteria:threshold | float | the iteration is terminated when the increase of the selected statistics is below this value |  |\n",
    "|parameter_RL:save_results:activate | bool | whether an updated model map, detal map, likelihood etc. are saved after the image deconvolution | |\n",
    "|parameter_RL:save_results:directory | bool | the directory where the files are saved | |\n",
    "|parameter_RL:save_results:only_final_result | bool | If True, only the result after the last iteration is saved | |"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ba85bc8",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_deconvolution.initialize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a27c72f1",
   "metadata": {},
   "source": [
    "### (You can change the parameters as follows)\n",
    "\n",
    "Note that when you modify the parameters, do not forget to run \"initialize\" again!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "47043a33",
   "metadata": {},
   "outputs": [],
   "source": [
    "image_deconvolution.override_parameter(\"deconvolution:parameter:iteration_max = 50\")\n",
    "image_deconvolution.override_parameter(\"deconvolution:parameter:background_normalization_optimization:activate = True\")\n",
    "image_deconvolution.override_parameter(\"deconvolution:parameter:acceleration:accel_factor_max = 2.0\")\n",
    "image_deconvolution.override_parameter(\"deconvolution:parameter:smoothing:activate = False\")\n",
    "#image_deconvolution.override_parameter(\"deconvolution:parameter:smoothing:activate = True\")\n",
    "image_deconvolution.override_parameter(\"deconvolution:parameter:response_weighting:activate = False\")\n",
    "\n",
    "image_deconvolution.initialize()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2f491fb",
   "metadata": {},
   "source": [
    "### 3-4. Start the image deconvolution\n",
    "\n",
    "**With MacBook Pro with M1 Max and 64 GB memory, it takes about 1.5 minutes for 50 iterations.**"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "183d38d5",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "image_deconvolution.run_deconvolution()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c7c6121e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "image_deconvolution.results"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "111e7d81",
   "metadata": {},
   "source": [
    "## 4. Analyze the results\n",
    "Examples to see/analyze the results are shown below."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1291b338",
   "metadata": {},
   "source": [
    "### Log-likelihood\n",
    "\n",
    "Plotting the log-likelihood vs the number of iterations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "94db94dd",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = [], []\n",
    "\n",
    "for result in image_deconvolution.results:\n",
    "    x.append(result['iteration'])\n",
    "    y.append(result['log-likelihood'])\n",
    "    \n",
    "plt.plot(x, y)\n",
    "plt.grid()\n",
    "plt.xlabel(\"iteration\")\n",
    "plt.ylabel(\"log-likelihood\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8dbf9492",
   "metadata": {},
   "source": [
    "### Alpha (the factor used for the acceleration)\n",
    "\n",
    "Plotting $\\alpha$ vs the number of iterations. $\\alpha$ is a parameter to accelerate the EM algorithm (see the beginning of Section 4). If it is too large, reconstructed images may have artifacts."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42e413ac",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = [], []\n",
    "\n",
    "for result in image_deconvolution.results:\n",
    "    x.append(result['iteration'])\n",
    "    y.append(result['accel_factor'])\n",
    "    \n",
    "plt.plot(x, y)\n",
    "plt.grid()\n",
    "plt.xlabel(\"iteration\")\n",
    "plt.ylabel(\"accel_factor\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99f14485",
   "metadata": {},
   "source": [
    "### Background normalization\n",
    "\n",
    "Plotting the background nomalization factor vs the number of iterations. If the backgroud model is accurate and the image is reconstructed perfectly, this factor should be close to 1."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1f1e6bbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "x, y = [], []\n",
    "\n",
    "for result in image_deconvolution.results:\n",
    "    x.append(result['iteration'])\n",
    "    y.append(result['background_normalization']['albedo'])\n",
    "    \n",
    "plt.plot(x, y)\n",
    "plt.grid()\n",
    "plt.xlabel(\"iteration\")\n",
    "plt.ylabel(\"background_normalization\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2ba1169",
   "metadata": {},
   "source": [
    "### The reconstructed images"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4d9ced90",
   "metadata": {
    "scrolled": true
   },
   "outputs": [],
   "source": [
    "def plot_reconstructed_image(result, source_position = None): # source_position should be (l,b) in degrees\n",
    "    iteration = result['iteration']\n",
    "    image = result['model']\n",
    "\n",
    "    for energy_index in range(image.axes['Ei'].nbins):\n",
    "        map_healpxmap = HealpixMap(data = image[:,energy_index], unit = image.unit)\n",
    "\n",
    "        _, ax = map_healpxmap.plot('mollview')        \n",
    "        \n",
    "        _.colorbar.set_label(str(image.unit))\n",
    "        \n",
    "        if source_position is not None:\n",
    "            ax.scatter(source_position[0]*u.deg, source_position[1]*u.deg, transform=ax.get_transform('world'), color = 'red')\n",
    "\n",
    "        plt.title(label = f\"iteration = {iteration}, energy_index = {energy_index} ({image.axes['Ei'].bounds[energy_index][0]}-{image.axes['Ei'].bounds[energy_index][1]})\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "abeaed95",
   "metadata": {},
   "source": [
    "Plotting the reconstructed images in all of the energy bands at the 50th iteration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6481d04c",
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration = 49\n",
    "\n",
    "plot_reconstructed_image(image_deconvolution.results[iteration])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29681a14",
   "metadata": {},
   "source": [
    "An example to plot the image in the log scale"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "95260440",
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration_idx = 49\n",
    "\n",
    "result = image_deconvolution.results[iteration_idx]\n",
    "\n",
    "iteration = result['iteration']\n",
    "image = result['model']\n",
    "\n",
    "data = image[:,0]\n",
    "data[data <= 0 * data.unit] = 1e-12 * data.unit\n",
    "\n",
    "hp.mollview(data, min = 1e-5, norm ='log', unit = str(data.unit), title = f'511 keV image at {iteration}th iteration', cmap = 'magma')\n",
    "\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a56f58a7",
   "metadata": {},
   "source": [
    "**Note: when you use a smoothing filter, you may see arfacts like many rings around the galactic center as seen in the next image. This is likely caused by the calculation errors of hp.smoothing and the high contrast of the 511 keV image. It is expected that these artifacts will be mitigated by using finer pixel resolution.**\n",
    "\n",
    "For reference, the parameters used for the image below is:\n",
    "```\n",
    "algorithm: \"RL\"\n",
    "parameter:\n",
    "  iteration_max: 50\n",
    "  acceleration:\n",
    "    activate: True\n",
    "    accel_factor_max: 2.0\n",
    "  response_weighting:\n",
    "    activate: False\n",
    "    index: 0.5\n",
    "  smoothing:\n",
    "    activate: False\n",
    "    FWHM:\n",
    "      value: 2.0\n",
    "      unit: \"deg\"\n",
    "  stopping_criteria:\n",
    "    statistics: \"log-likelihood\"\n",
    "    threshold: 0.01\n",
    "  background_normalization_optimization:\n",
    "    activate: True\n",
    "    range: {\"albedo\": [0.01, 10.0]}\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "64f885e6",
   "metadata": {},
   "outputs": [],
   "source": [
    "iteration_idx = 49\n",
    "\n",
    "result = image_deconvolution.results[iteration_idx]\n",
    "\n",
    "iteration = result['iteration']\n",
    "image = result['model']\n",
    "\n",
    "data = image[:,0]\n",
    "data[data <= 0 * data.unit] = 1e-12 * data.unit\n",
    "\n",
    "hp.mollview(data, min = 1e-5, norm ='log', unit = str(data.unit), title = f'511 keV image at {iteration}th iteration', cmap = 'magma')\n",
    "\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "cosipy-312",
   "language": "python",
   "name": "cosipy-312"
  },
  "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.12.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}