{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Using nabu from python to reconstruct a dataset with GPU\n",
    "\n",
    "This notebook shows how to use the Nabu software for performing a basic reconstruction of a tomography dataset.  \n",
    "The computations are done on a local machine with a GPU and Cuda available.\n",
    "\n",
    "This tutorial goes a bit further than `nabu_basic_reconstruction.ipynb`:\n",
    "  - GPU implementation of each component is used\n",
    "  - We see how to start from a configuration file and devise a simple processing chain accordingly\n",
    "  \n",
    "The same dataset is used (binned scan of a bamboo stick, thanks Ludovic Broche, ESRF ID19)."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 1 - Load the dataset informations\n",
    "\n",
    "We must provide `nabu` with the the configuration file (`nabu.conf`), describing the path to the dataset and the processing steps. This is the equivalent of the `.par` file in PyHST2. In this file, no information is given on the detector size, energy, distance, etc: these informations are extracted from the dataset metadata.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from nabu.testutils import utilstest, get_file\n",
    "from nabu.pipeline.fullfield.processconfig import ProcessConfig"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Getting dataset (downloading if necessary) ...\")\n",
    "data_path = get_file(\"bamboo_reduced.nx\")\n",
    "print(\"... OK\")\n",
    "\n",
    "# Get the configuration file of this dataset\n",
    "conf_fname = get_file(\"bamboo_reduced.conf\")\n",
    "\n",
    "# Change directory to the path where the data is located (only useful for this tutorial)\n",
    "os.chdir(utilstest.data_home)\n",
    "\n",
    "# Parse this configuration file\n",
    "conf = ProcessConfig(conf_fname)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that `ProcessConfig` will do quite a few things under the hood:\n",
    "  - Parse the configuration file and check parameters correctness\n",
    "  - Browse the dataset\n",
    "  - Get or compute the reduced flats/darks\n",
    "  - Estimate the center of rotation\n",
    "  \n",
    "The resulting object contains all necessary information to process the dataset."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# We can easily get information on the processing steps.\n",
    "nabu_config = conf.nabu_config\n",
    "print(nabu_config)\n",
    "# The same can be done with the dataset structure\n",
    "dataset_info = conf.dataset_info\n",
    "# print([getattr(dataset_info, attr) for attr in [\"energy\", \"distance\", \"n_angles\", \"radio_dims\"]])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 2 - Chunk processing\n",
    "\n",
    "Nabu processes data by chunks of radios (see the [documentation](http://www.silx.org/pub/nabu/doc/definitions.html#radios-and-sinograms) for more explanations).  \n",
    "In a first step, we define how to read chunks of radios."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.io.reader import ChunkReader"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "What is the largest chunk size we can process ?  \n",
    "The answer is given by inspecting the current GPU memory, and the processing steps."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.cuda.utils import get_gpu_memory\n",
    "from nabu.pipeline.fullfield.computations import estimate_max_chunk_size"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "chunk_size = estimate_max_chunk_size(\n",
    "    get_gpu_memory(0), \n",
    "    conf\n",
    ")\n",
    "print(\"Chunk_size = %d\" % chunk_size)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the first 'chunk_size' lines of all the radios\n",
    "sub_region = (None, None, 0, chunk_size) # start_x, end_x, start_z, end_z\n",
    "chunk_reader = ChunkReader(\n",
    "    dataset_info.projections, \n",
    "    sub_region=sub_region, \n",
    "    convert_float=True\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load the current chunk\n",
    "chunk_reader.load_files() # takes some time"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(chunk_reader.files_data.shape)\n",
    "print(chunk_reader.files_data.dtype)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 3 - Initialize the GPU\n",
    "\n",
    "Most of the processing can be done on GPU (or many-core CPU if using OpenCL).  \n",
    "With `pycuda.gpuarray` (or its OpenCL counterpart `pyopencl.array`), we manipulate array objects with memory residing on device. This allows to avoid extraneous host <-> device copies."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pycuda.gpuarray as garray\n",
    "from nabu.cuda.utils import get_cuda_context\n",
    "import numpy as np"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a Cuda context on device ID 0\n",
    "# By default, all following GPU processings will be bound on this context\n",
    "ctx = get_cuda_context(device_id=0)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radios = chunk_reader.files_data\n",
    "n_angles, n_z, n_x = radios.shape\n",
    "# transfer the chunk on GPU\n",
    "d_radios = garray.to_gpu(radios)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 4 - Pre-processing\n",
    "\n",
    "Pre-processing utilities are available in the `nabu.preproc` module.  \n",
    "Utilities available with the cuda backend are implemented in a module with a `_cuda` suffix."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.1 - Flat-field"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.preproc.flatfield_cuda import CudaFlatFieldDataUrls"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "radios_indices = sorted(conf.dataset_info.projections.keys())\n",
    "# Configure the `FlatField` processor\n",
    "cuda_flatfield = CudaFlatFieldDataUrls(\n",
    "    d_radios.shape, \n",
    "    dataset_info.flats, \n",
    "    dataset_info.darks, \n",
    "    radios_indices=radios_indices,\n",
    "    sub_region=sub_region,\n",
    "    convert_float=True,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Perform the normalization on GPU\n",
    "if nabu_config[\"preproc\"][\"flatfield\"]:\n",
    "    print(\"Doing flat-field\")\n",
    "    cuda_flatfield.normalize_radios(d_radios)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.2 - Phase retrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.preproc.phase_cuda import CudaPaganinPhaseRetrieval"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "energy = dataset_info.energy\n",
    "# Phase retrieval is done on each radio individually, with the sub-region specified above\n",
    "if (nabu_config[\"phase\"][\"method\"] or \"\").lower() == \"paganin\":\n",
    "    print(\"Doing phase retrieval\")\n",
    "    cudapaganin = CudaPaganinPhaseRetrieval(\n",
    "        (n_z, n_x),\n",
    "        distance=dataset_info.distance,\n",
    "        energy=energy,\n",
    "        delta_beta=nabu_config[\"phase\"][\"delta_beta\"],\n",
    "        pixel_size=dataset_info.pixel_size * 1e6,\n",
    "    )\n",
    "    for i in range(n_angles):\n",
    "        cudapaganin.apply_filter(d_radios[i], output=d_radios[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 4.3 - Logarithm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.preproc.ccd_cuda import CudaLog"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "if nabu_config[\"preproc\"][\"take_logarithm\"]:\n",
    "    print(\"Taking logarithm\")\n",
    "    cuda_log = CudaLog(d_radios.shape, clip_min=0.01)\n",
    "    cuda_log.take_logarithm(d_radios)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 5 - Reconstruction\n",
    "\n",
    "We use the filtered backprojection with `nabu.reconstruction.fbp`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from nabu.reconstruction.fbp import Backprojector"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rec_options = conf.processing_options[\"reconstruction\"]\n",
    "B = Backprojector(\n",
    "    (n_angles, n_x), \n",
    "    angles=rec_options[\"angles\"], \n",
    "    rot_center=rec_options[\"rotation_axis_position\"],\n",
    "    padding_mode=\"edges\"\n",
    ")\n",
    "d_recs = garray.zeros((n_z, n_x, n_x), \"f\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Reconstructing...\", end=\"\")\n",
    "for i in range(n_z):\n",
    "    B.fbp(radios[:, i, :], output=d_recs[i])\n",
    "recs = d_recs.get()\n",
    "print(\" ... OK\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## 6 - Visualize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%pylab nbagg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "figure()\n",
    "imshow(recs[0], cmap=\"gray\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 2
}