{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "97fcdc49-7965-46a3-b7fd-4286855c31ee",
   "metadata": {},
   "source": [
    "# Nabu patch processing from a python script"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "05afb0fc-7d96-4fbe-aa99-c9ed12ebb07b",
   "metadata": {},
   "source": [
    "## connect to the slurm - cluster\n",
    "```bash\n",
    "ssh -XC slurm-cluster\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "62d869d4-2fb8-43b1-829a-ccaa79f09eca",
   "metadata": {},
   "source": [
    "## define the cluster node configuration you want to have"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f798fee0-4bd5-4909-b0d4-e600274eed10",
   "metadata": {},
   "outputs": [],
   "source": [
    "cluster_config = {  # slurm cluster configuration to be used for each reconstruction\n",
    "    \"cpu-per-task\": 4,\n",
    "    \"n_tasks\": 1,\n",
    "    \"memory\": 256,\n",
    "    \"partition\": \"p9gpu\",\n",
    "    \"n_gpus\": 1,\n",
    "    \"job_name\": \"nabu-batch-processing\",\n",
    "    \"python_venv\": \"/scisoft/tomotools/activate.sh dev\",\n",
    "    # \"modules\": \"tomotools\",  # either the module or the python virtual environment should be provided\n",
    "}"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9061c07d-3d93-474e-94a2-056caeb3e892",
   "metadata": {},
   "source": [
    "## define the nabu settings which are different from the default values\n",
    "\n",
    "for example here we change the default file format, the dataset location (mndatory for batch processing) and the algorithm to be used to compute the center of rotation"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9a7e69b-0849-4b0a-ad94-b3c8d5882bab",
   "metadata": {},
   "outputs": [],
   "source": [
    "def get_nabu_config(dataset_location, output_location):\n",
    "    \"\"\"\n",
    "    :return: tuned parameter from  the default configuration of nabu\n",
    "    You can tune need if needed by looking each key position directly in the generated nabu configuration\n",
    "    \"\"\"\n",
    "    return {\n",
    "        \"dataset\": {\n",
    "            \"location\": dataset_location,\n",
    "        },\n",
    "        \"reconstruction\": {\n",
    "            # Rotation axis position. It can be a number or the name of an estimation method (empty value means the middle of the detector).\n",
    "            # The following methods are available to find automatically the Center of Rotation (CoR):\n",
    "            #  - centered : a fast and simple auto-CoR method. It only works when the CoR is not far from the middle of the detector. It does not work for half-tomography.\n",
    "            #  - global : a slow but robust auto-CoR.\n",
    "            #  - sliding-window : semi-automatically find the CoR with a sliding window. You have to specify on which side the CoR is (left, center, right). Please see the 'cor_options' parameter.\n",
    "            #  - growing-window : automatically find the CoR with a sliding-and-growing window. You can tune the option with the parameter 'cor_options'.\n",
    "            #  - sino-coarse-to-fine: Estimate CoR from sinogram. Only works for 360 degrees scans.\n",
    "            #  - composite-coarse-to-fine: Estimate CoR from composite multi-angle images. Only works for 360 degrees scans.\n",
    "            \"rotation_axis_position\": \"sliding-window\",\n",
    "        },\n",
    "        \"output\": {\n",
    "            \"file_format\": \"tiff\",\n",
    "            \"location\": output_location,\n",
    "        },\n",
    "    }"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "641855a1-85a8-4c25-a855-c17f5f818c43",
   "metadata": {},
   "source": [
    "## define the remote processing to be done for each NXtomo - or scan"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "be3ca21c-2414-45cd-b224-e7d1f4136772",
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "from nabu.pipeline.config import generate_nabu_configfile\n",
    "from nabu.pipeline.fullfield.nabu_config import (\n",
    "    nabu_config as nabu_fullfield_default_config,\n",
    ")\n",
    "from sluurp.job import SBatchScriptJob\n",
    "from sluurp.executor import submit as submit_to_slurm_cluster\n",
    "import uuid\n",
    "from datetime import datetime\n",
    "\n",
    "\n",
    "def treat_single_nx_tomo(nx_tomo_file):\n",
    "    \"\"\"\n",
    "    treat a single .nx file\n",
    "    \"\"\"\n",
    "    print(f\"treat {nx_tomo_file}_file\")\n",
    "    # step1: save nabu configuration to file (near .nx file)\n",
    "    nabu_conf_file = nx_tomo_file.replace(\".nx\", \"_nabu.cfg\")\n",
    "    output_location = nx_tomo_file.replace(\".nx\", \"_rec\")\n",
    "    generate_nabu_configfile(\n",
    "        nabu_conf_file,\n",
    "        nabu_fullfield_default_config,\n",
    "        config=get_nabu_config(dataset_location=nx_tomo_file, output_location=output_location),\n",
    "        options_level=\"advanced\",\n",
    "    )\n",
    "    # step2 create SBatch job for slurm\n",
    "    now_str = str(datetime.now()).replace(\" \", \"_\")\n",
    "    script_name = os.path.splitext(os.path.basename(nx_tomo_file))[0] + f\"_{now_str}.sh\"\n",
    "    job = SBatchScriptJob(\n",
    "        slurm_config=cluster_config,\n",
    "        script=(\n",
    "            f\"python3 -m nabu.app.reconstruct {nabu_conf_file} --slice middle\",\n",
    "        ),  # TODO: remove '--slice middle' to reconstruct the full volume\n",
    "        script_path=os.path.join(\n",
    "            os.path.dirname(nx_tomo_file), \"slurm_scripts\", script_name\n",
    "        ),\n",
    "        clean_script=False,\n",
    "        working_directory=os.path.abspath(os.path.dirname(nx_tomo_file)),\n",
    "    )\n",
    "    future_slurm_job = submit_to_slurm_cluster(job)\n",
    "    # print(\"wait for {nx_tomo_file}\")\n",
    "    # future_slurm_job.result()\n",
    "    # print(\"{nx_tomo_file} finished\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "bdae4803-ad4e-45d9-9471-1fb744deae2c",
   "metadata": {},
   "source": [
    "## execute batch processing - do be done on a slurm-client (slurm cluster fron end)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7993646-51ef-44ee-a42c-20cc9f611d63",
   "metadata": {},
   "outputs": [],
   "source": [
    "from glob import glob\n",
    "\n",
    "# define nx_tomo to be processed\n",
    "nxtomo_s = glob(\"/data/visitor/esXXXX/YY/ZZZZ/PROCESSED_DATA/*/*/*.nx\")\n",
    "nxtomo_s = (\n",
    "    \"/data/visitor//esXXXX/YY/ZZZZ/PROCESSED_DATA/*/*/aaa.nx\",\n",
    "    \"/data/visitor//esXXXX/YY/ZZZZ/PROCESSED_DATA/*/*/aaa.nx\",\n",
    ")\n",
    "for nxtomo in nxtomo_s:\n",
    "    treat_single_nx_tomo(nxtomo)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76fbbd7c-c02e-40bb-8374-4502dda6706d",
   "metadata": {},
   "source": [
    "### note\n",
    "to simplify this tutorial consider that you already have the NXtomo (.nx) files. If necessary you can do some preprocessing to retrieve the bliss file on `RAW_DATA` then obtain NXtomo from [nxtomomill](https://tomotools.gitlab-pages.esrf.fr/nxtomomill/)\n",
    "either from [the command line interface](https://tomotools.gitlab-pages.esrf.fr/nxtomomill/_generated/nxtomomill.app.h52nx.html) or from the [python API](https://tomotools.gitlab-pages.esrf.fr/nxtomomill/_generated/nxtomomill.converter.hdf5.hdf5converter.from_h5_to_nx.html#nxtomomill.converter.hdf5.hdf5converter.from_h5_to_nx)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "25e772ad-e2b5-4a17-a358-0857cc295a53",
   "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
}