Adding explanations

Author: mirjagranfors

tutorials/2-examples/DTEx252_phase_mask_optimization.ipynb

@@ -5,23 +5,26 @@
    "id": "5e65db32",
    "metadata": {},
    "source": [
-    "# Particle localization and phase mask optimization"
+    "# Particle Localization and Phase Mask Optimization"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "d4aa8b6b",
    "metadata": {},
    "source": [
-    "This tutorial demonstrates how to optimize a phase mask to improve localization of closely spaced particles. The localization is performed by a convolutional neural network (CNN), trained jointly with the phase mask to improve both the optical setup and the network performance."
+    "This tutorial demonstrates how to jointly optimize an optical phase mask and a neural network to improve the localization of closely spaced particles in microscopy images.\n",
+    "\n",
+    "The phase mask, inserted in the optical pupil plane, shapes the microscope’s point spread function (PSF). Simultaneously, a convolutional neural network (CNN) is trained to reconstruct the 3D particle positions from simulated images. Through backpropagation, gradients flow not only through the network but also through the optical model. As a result, the optical system learns to produce images that are easier for the CNN to interpret, leading to improved localization accuracy."
    ]
   },
   {
    "cell_type": "markdown",
    "id": "82dfa30b",
    "metadata": {},
    "source": [
-    "# 1. Create a simulation pipeline"
+    "# 1. Create a simulation pipeline\n",
+    "The first step is to define a simulation pipeline that generates synthetic data consisting of fluorescent particles imaged through a fluorescence microscope. For each simulated image, the pipeline also provides the true 3D positions of the particles, which serve as the ground truth during training."
    ]
   },
   {
@@ -31,28 +34,24 @@
    "metadata": {},
    "outputs": [
     {
-     "name": "stdout",
+     "name": "stderr",
      "output_type": "stream",
      "text": [
-      "c:\\Users/xgrmir/Documents/DeepTrack2/DeepTrack2\\deeptrack\\__init__.py\n"
+      "WARNING:pint.util:Redefining '[magnetic_flux]' (<class 'pint.delegates.txt_defparser.plain.DerivedDimensionDefinition'>)\n"
      ]
     }
    ],
    "source": [
-    "import sys\n",
-    "\n",
-    "sys.path.insert(0, '/Users/xgrmir/Documents/DeepTrack2/DeepTrack2')\n",
-    "\n",
-    "import deeptrack as dt\n",
-    "print(dt.__file__)"
+    "import deeptrack as dt"
    ]
   },
   {
    "cell_type": "markdown",
    "id": "8a7638e5",
    "metadata": {},
    "source": [
-    "## 1.1 Set the backend"
+    "## 1.1 Set the backend\n",
+    "DeepTrack supports both NumPy and PyTorch backends. In this tutorial, the PyTorch backend is selected to enable backpropagation through the optical setup. Enabling CUDA (optional) allows the computations to run on the GPU, which accelerates both simulation and training."
    ]
   },
   {
@@ -86,13 +85,14 @@
    "metadata": {},
    "source": [
     "## 1.2 Define the trainable phase mask\n",
+    "A phase mask modifies the wavefront of light in the microscope pupil plane, shaping the resulting point spread function (PSF). By making the phase mask trainable, we allow the optical system itself to learn — optimizing its design through gradient descent to improve particle localization performance.\n",
     "\n",
-    "I don't know if inheriting from both nn.Module and dt.Aberration is correct? I inherit from dt.Aberration to get a dt feature that is compatible with the pipeline, and I inherit from nn.Module to make the layer trainable"
+    "The LearnablePhaseMask class inherits from both `nn.Module` and `dt.Aberration`. `nn.Module` makes the mask compatible with PyTorch’s autograd and optimizer system, allowing the phase parameters to be learned, while `dt.Aberration` ensures the class integrates seamlessly into DeepTrack’s feature graph, making it compatible with the rest of the optical pipeline."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "id": "3470f1bf",
    "metadata": {},
    "outputs": [],
@@ -107,19 +107,19 @@
     "        dt.Aberration.__init__(self, **kwargs)\n",
     "        super().__init__(**kwargs)\n",
     "        \n",
-    "        # Create a torch parameter for the phase\n",
+    "        # Create a trainable tensor representing the phase\n",
     "        self.phase = nn.Parameter(torch.zeros(shape, dtype=torch.float32))\n",
     "\n",
     "    def forward(self, pupil: torch.Tensor, **kwargs) -> torch.Tensor:\n",
-    "        \"\"\"PyTorch forward pass\"\"\"\n",
+    "        \"\"\"PyTorch forward pass for use in training\"\"\"\n",
     "        return self.apply_phase(pupil)\n",
     "\n",
     "    def get(self, pupil: torch.Tensor, **kwargs) -> torch.Tensor:\n",
     "        \"\"\"DeepTrack Feature graph call\"\"\"\n",
     "        return self.apply_phase(pupil)\n",
     "\n",
     "    def apply_phase(self, pupil: torch.Tensor) -> torch.Tensor:\n",
-    "        \"\"\"Shared implementation\"\"\"\n",
+    "        \"\"\"Shared implementation of phase modulation\"\"\"\n",
     "        phase = self.phase.to(device)\n",
     "        phase_mask = torch.cos(phase) + 1j * torch.sin(phase)\n",
     "        return pupil * phase_mask"
@@ -130,7 +130,8 @@
    "id": "e8a1553c",
    "metadata": {},
    "source": [
-    "## 1.3 Define the optical setup"
+    "## 1.3 Define the optical setup\n",
+    "The optical setup is simulated using `dt.Fluorescence`, incorporating the trainable phase mask as the pupil function."
    ]
   },
   {
@@ -166,25 +167,10 @@
    "id": "8fe88c92",
    "metadata": {},
    "source": [
-    "## 1.4 Simulate particles"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "id": "d5c1e312",
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "import matplotlib.pyplot as plt"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "id": "06cffcce",
-   "metadata": {},
-   "source": [
-    "### Random positions in 3D, and gt in 3D:"
+    "## 1.4 Simulate particles\n",
+    "A simulation pipeline is defined to generate images of randomly positioned 3D particles, together with their corresponding ground-truth localization targets.\n",
+    "\n",
+    "A custom DeepTrack feature, `Positions`, creates a binary 3D mask containing randomly placed particles. Each voxel corresponding to a particle is assigned a value of 1, while background voxels remain 0. The true particle coordinates are stored in self.points for later retrieval."
    ]
   },
   {
@@ -212,6 +198,14 @@
     "        return mask + image"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "38c4a606",
+   "metadata": {},
+   "source": [
+    "The number of points for each sample is randomized between 25 and 50 to provide diverse training examples."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 8,
@@ -223,14 +217,24 @@
     "xyz = Positions(num_points=num_points)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "446fee06",
+   "metadata": {},
+   "source": [
+    "To make the training data physically plausible, a combination of Poisson noise and Gaussian noise is added.\n",
+    "\n",
+    "Since true Poisson noise is not differentiable, a Gaussian approximation that preserves differentiability during backpropagation is used."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 9,
+   "execution_count": null,
    "id": "7d339aaf",
    "metadata": {},
    "outputs": [],
    "source": [
-    "class poisson_noise_approx(dt.Noise):       # Because the real poisson noise is not compatible with backpropagation\n",
+    "class poisson_noise_approx(dt.Noise):\n",
     "    def __init__(\n",
     "        self,\n",
     "        **kwargs,\n",
@@ -250,6 +254,14 @@
     "        return noisy_image"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "9ac539dc",
+   "metadata": {},
+   "source": [
+    "Next, the components, including the particle positions, the optics, and the noise, are combined to form the simulation pipeline."
+   ]
+  },
   {
    "cell_type": "code",
    "execution_count": 10,
@@ -275,9 +287,17 @@
     "pip = (im_pip & gt_pip) >> dt.MoveAxis(2, 0)"
    ]
   },
+  {
+   "cell_type": "markdown",
+   "id": "6bff75c1",
+   "metadata": {},
+   "source": [
+    "An example simulation is visualized below."
+   ]
+  },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": null,
    "id": "2f78ddd7",
    "metadata": {},
    "outputs": [
@@ -303,13 +323,23 @@
     }
    ],
    "source": [
+    "import matplotlib.pyplot as plt\n",
+    "\n",
     "pip.update()\n",
     "im, gt = pip.resolve()\n",
     "\n",
     "fig, axs = plt.subplots(1, 2, figsize=(8,4))\n",
     "\n",
     "axs[0].imshow(im[0].cpu().detach().numpy(), cmap='gray')\n",
-    "axs[1].imshow(gt.max(dim=0)[0])"
+    "axs[1].imshow(gt.max(dim=0)[0]);"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "id": "297f7ade",
+   "metadata": {},
+   "source": [
+    "The left image shows a synthetic microscopy image produced by the optical system with the current, untrained phase mask, while the right image displays the corresponding ground-truth particle positions projected onto the x–y plane."
    ]
   },
   {
@@ -2182,7 +2212,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 30,
+   "execution_count": null,
    "id": "edd1bcd0",
    "metadata": {},
    "outputs": [],
@@ -2193,7 +2223,7 @@
     "    output_region=(0,0, image_size, image_size),\n",
     ")\n",
     "\n",
-    "im_pip_no_phase_mask = optics_no_phase_mask(particle ^ num_points)\n",
+    "im_pip_no_phase_mask = optics_no_phase_mask(particle ^ num_points) ## Add noise!\n",
     "\n",
     "pip_no_phase_mask = (im_pip_no_phase_mask & gt_pip) >> dt.MoveAxis(2, 0)"
    ]