{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "01ce690b",
   "metadata": {},
   "source": [
    "# Basic reconstruction with nabu\n",
    "\n",
    "In this notebook, we see how to use the different building blocks of nabu to parse a dataset and perform a basic reconstruction.\n",
    "\n",
    "This notebook uses a dataset of a bamboo stick (acquired at ESRF ID19, courtesy Ludovic Broche). The projections were binned by a factor of 4 in each dimension, and the angular range was also subsampled by 4."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c7cd7abd",
   "metadata": {},
   "source": [
    "## Browse the dataset\n",
    "\n",
    "We use nabu utility `analyze_dataset` which builds a data structure with all information on the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ff271ef6",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "from tempfile import mkdtemp\n",
    "from nabu.resources.dataset_analyzer import analyze_dataset\n",
    "from nabu.resources.nxflatfield import update_dataset_info_flats_darks\n",
    "from nabu.testutils import get_file"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8ecf7c13",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Getting dataset (downloading if necessary) ...\")\n",
    "data_path = get_file(\"bamboo_reduced.nx\")\n",
    "print(\"... OK\")\n",
    "output_dir = mkdtemp(prefix=\"nabu_reconstruction\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b74f55fd",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Parse the \".nx\" file. This NX file is our entry point when it comes to data,\n",
    "# as it's only the format which is remotely stable\n",
    "# From this .nx file, build a data structure with readily available information\n",
    "dataset_info = analyze_dataset(data_path)\n",
    "\n",
    "# A tomography scan usually contains a series of darks/flats.\n",
    "# These darks/flats are then averaged (the legacy pipeline called them 'refHST' and 'darkend')\n",
    "# The flat-field normalization should be done only with reduced flats/darks. \n",
    "update_dataset_info_flats_darks(dataset_info, True, darks_flats_dir=output_dir)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "24b8e0a5",
   "metadata": {},
   "source": [
    "## Estimate the center of rotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3c5057ec",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.pipeline.estimators import CORFinder"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "61422676",
   "metadata": {},
   "outputs": [],
   "source": [
    "cor_finder = CORFinder(\n",
    "    \"sliding-window\",\n",
    "    dataset_info,\n",
    "    do_flatfield=True,\n",
    ")\n",
    "cor = cor_finder.find_cor()\n",
    "print(\"Estimated center of rotation: %.2f\" % cor)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6846b426",
   "metadata": {},
   "source": [
    "## Define how the data should be processed\n",
    "\n",
    "Instantiate the various components of the pipeline."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dbca90a4",
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.io.reader import ChunkReader\n",
    "from nabu.preproc.flatfield import FlatFieldDataUrls\n",
    "from nabu.preproc.phase import PaganinPhaseRetrieval\n",
    "from nabu.reconstruction.fbp import Backprojector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "42e5367b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the sub-region to read (and reconstruct)\n",
    "# In each radio: will read (start_x, end_x, start_y, end_y)\n",
    "sub_region = (None, None, 270, 289)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c223d65",
   "metadata": {},
   "outputs": [],
   "source": [
    "reader = ChunkReader(\n",
    "    dataset_info.projections,\n",
    "    sub_region=sub_region,\n",
    "    convert_float=True,\n",
    ")\n",
    "radios = reader.data\n",
    "n_angles, n_z, n_x = radios.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2313e5ab",
   "metadata": {},
   "outputs": [],
   "source": [
    "flatfield = FlatFieldDataUrls(\n",
    "    radios.shape,\n",
    "    dataset_info.flats,\n",
    "    dataset_info.darks,\n",
    "    radios_indices=sorted(list(dataset_info.projections.keys())),\n",
    "    sub_region=sub_region\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6fe3916e",
   "metadata": {},
   "outputs": [],
   "source": [
    "paganin = PaganinPhaseRetrieval(\n",
    "    radios[0].shape,\n",
    "    distance=dataset_info.distance,\n",
    "    energy=dataset_info.energy,\n",
    "    delta_beta=250., # should be tuned\n",
    "    pixel_size=dataset_info.pixel_size * 1e-6,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a740140a",
   "metadata": {},
   "outputs": [],
   "source": [
    "reconstruction = Backprojector(\n",
    "    (n_angles, n_x),\n",
    "    angles=dataset_info.rotation_angles,\n",
    "    rot_center=cor,\n",
    "    halftomo=dataset_info.is_halftomo,\n",
    "    padding_mode=\"edges\",\n",
    "    scale_factor=1/(dataset_info.pixel_size * 1e-4), # mu/cm\n",
    "    extra_options={\"centered_axis\": True, \"clip_outer_circle\": True}\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ab49b848",
   "metadata": {},
   "source": [
    "## Read data\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b74a22e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "reader.load_data(overwrite=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "3c9d1915",
   "metadata": {},
   "source": [
    "## Pre-processing"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e2739e59",
   "metadata": {},
   "outputs": [],
   "source": [
    "flatfield.normalize_radios(radios) # in-place\n",
    "print()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a59d4133",
   "metadata": {},
   "outputs": [],
   "source": [
    "radios_phase = np.zeros_like(radios)\n",
    "for i in range(radios.shape[0]):\n",
    "    paganin.retrieve_phase(radios[i], output=radios_phase[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "89cf8a73",
   "metadata": {},
   "source": [
    "## Reconstruction\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "55e3c0af",
   "metadata": {},
   "outputs": [],
   "source": [
    "volume = np.zeros((n_z, n_x, n_x), \"f\")\n",
    "\n",
    "for i in range(n_z):\n",
    "    volume[i] = reconstruction.fbp(radios[:, i, :])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9a86229c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b3c61dab",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure()\n",
    "plt.imshow(volume[0], cmap=\"gray\")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "94dec68a",
   "metadata": {},
   "source": [
    "## Going further: GPU processing\n",
    "\n",
    "All the above components have a Cuda backend: `SinoBuilder` becomes `CudaSinoBuilder`, `PaganinPhaseRetrieval` becomes `CudaPaganinPhaseRetrieval`, and so on.\n",
    "Importantly, the cuda counterpart of these classes have the same API."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4ab4eac8",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "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.11.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}