{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "286054ce",
   "metadata": {},
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "import cdtools\n",
    "import torch as t\n",
    "from matplotlib import pyplot as plt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "955bb242-e2ed-47c3-919c-1ea690681445",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Load and inspect a dataset\n",
    "\n",
    "filename = 'example_data/lab_ptycho_data.cxi'\n",
    "dataset = cdtools.datasets.Ptycho2DDataset.from_cxi(filename)\n",
    "\n",
    "dataset.inspect();"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "82f71fb4-f013-46bb-817b-1973ba23336a",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Initialize a model from the dataset and move it to the GPU.\n",
    "\n",
    "model = cdtools.models.FancyPtycho.from_dataset(\n",
    "    dataset,\n",
    "    n_modes=3, # Use 3 incoherently mixing probe modes\n",
    "    oversampling=2, # Simulate the probe on a 2xlarger real-space array\n",
    "    probe_support_radius=120, # Force the probe to 0 outside a radius of 120 pix\n",
    "    propagation_distance=5e-3, # Propagate the initial probe guess by 5 mm\n",
    "    units='mm', # Set the units for the live plots\n",
    "    obj_view_crop=-50, # Expands the field of view in the object plot by 50 pix,\n",
    ")\n",
    "\n",
    "if t.cuda.is_available():\n",
    "    model.to(device='cuda')\n",
    "    dataset.get_as(device='cuda')\n",
    "\n",
    "# Then, create a reconstructor object and view the initialized model\n",
    "\n",
    "recon = cdtools.reconstructors.AdamReconstructor(model, dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0e22a65e-280b-428d-a6f5-f3206d22110b",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Workaround reconstruction pattern for interactive plotting in jupyter:\n",
    "# First, a standalone cell to plot the current model state\n",
    "\n",
    "model.inspect(replot_all=True);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81a768b5",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Second, a cell for running the reconstruction. With this pattern, it is safe\n",
    "# to interrupt the kernel. Then, the cell above can be re-run to refresh the plots.\n",
    "while model.epoch < 50:\n",
    "    for loss in recon.optimize(1, lr=0.02, batch_size=10):\n",
    "        print(model.report())\n",
    "\n",
    "while model.epoch < 100:\n",
    "    for loss in recon.optimize(1, lr=0.005, batch_size=10):\n",
    "        print(model.report())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "237e9286-b6cf-41dc-aeb0-89abfe59b37c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Save out the results\n",
    "\n",
    "model.save_to_h5('lab_ptycho_reconstruction.h5');"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0f1565b6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Finalize the plotting and create the comparison plot\n",
    "\n",
    "# This orthogonalizes the recovered probe modes. It is best to do so\n",
    "# after saving the results, if you intend to initialize any further\n",
    "# reconstructions with the probe.\n",
    "model.tidy_probes()\n",
    "\n",
    "# Final plotting\n",
    "model.inspect()\n",
    "model.compare(dataset);"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "63ec3aa4-0d1c-4775-9fb2-3002d404faa4",
   "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.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}