clean up tutorial

docs/examples/simulate-relion-dataset.ipynb

@@ -522,7 +522,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Checking how the image from the imaging pipeline looks"
+    "Checking how an image generated from the imaging pipeline looks"
    ]
   },
   {
@@ -555,7 +555,7 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "Now let's define a noise distribution"
+    "Define a noise distribution"
    ]
   },
   {
@@ -595,96 +595,6 @@
     ")"
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "Lastly, let's define our vmapping filter. This is needed, as some of the parameters loaded from the starfile have a batch dimension (i.e., everything that's not in the `optics` data block)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 15,
-   "metadata": {},
-   "outputs": [],
-   "source": [
-    "def vmap_filter(distribution: dist.AbstractDistribution):\n",
-    "    \"\"\"\n",
-    "    These are the only leaves that have a batch size in both\n",
-    "    the starfile and the imaging pipeline\n",
-    "    \"\"\"\n",
-    "    output = (\n",
-    "        distribution.imaging_pipeline.scattering_theory.structural_ensemble.pose.offset_x_in_angstroms,\n",
-    "        distribution.imaging_pipeline.scattering_theory.structural_ensemble.pose.offset_y_in_angstroms,\n",
-    "        distribution.imaging_pipeline.scattering_theory.structural_ensemble.pose.view_phi,\n",
-    "        distribution.imaging_pipeline.scattering_theory.structural_ensemble.pose.view_theta,\n",
-    "        distribution.imaging_pipeline.scattering_theory.structural_ensemble.pose.view_psi,\n",
-    "        distribution.imaging_pipeline.scattering_theory.transfer_theory.ctf.defocus_in_angstroms,\n",
-    "        distribution.imaging_pipeline.scattering_theory.transfer_theory.ctf.astigmatism_in_angstroms,\n",
-    "        distribution.imaging_pipeline.scattering_theory.transfer_theory.ctf.astigmatism_angle,\n",
-    "        distribution.imaging_pipeline.scattering_theory.transfer_theory.ctf.phase_shift,\n",
-    "        distribution.imaging_pipeline.scattering_theory.transfer_theory.envelope,\n",
-    "    )\n",
-    "    return output"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "You can also infer how to obtain this values by printing a RelionParticleStack directly. Look for values that have `[19]` (for this example I will load 19 images)"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 16,
-   "metadata": {},
-   "outputs": [
-    {
-     "data": {
-      "text/plain": [
-       "RelionParticleStack(\n",
-       "  instrument_config=InstrumentConfig(\n",
-       "    shape=(128, 128),\n",
-       "    pixel_size=f32[],\n",
-       "    voltage_in_kilovolts=f32[],\n",
-       "    electrons_per_angstrom_squared=f32[],\n",
-       "    padded_shape=(128, 128),\n",
-       "    pad_mode='constant'\n",
-       "  ),\n",
-       "  pose=EulerAnglePose(\n",
-       "    offset_x_in_angstroms=f32[19],\n",
-       "    offset_y_in_angstroms=f32[19],\n",
-       "    offset_z_in_angstroms=0.0,\n",
-       "    view_phi=f32[19],\n",
-       "    view_theta=f32[19],\n",
-       "    view_psi=f32[19]\n",
-       "  ),\n",
-       "  transfer_theory=ContrastTransferTheory(\n",
-       "    ctf=ContrastTransferFunction(\n",
-       "      defocus_in_angstroms=f32[19],\n",
-       "      astigmatism_in_angstroms=f32[19],\n",
-       "      astigmatism_angle=f32[19],\n",
-       "      voltage_in_kilovolts=300.0,\n",
-       "      spherical_aberration_in_mm=f32[19],\n",
-       "      amplitude_contrast_ratio=f32[19],\n",
-       "      phase_shift=f32[19]\n",
-       "    ),\n",
-       "    envelope=FourierGaussian(amplitude=f32[19], b_factor=f32[19])\n",
-       "  ),\n",
-       "  image_stack=None\n",
-       ")"
-      ]
-     },
-     "execution_count": 16,
-     "metadata": {},
-     "output_type": "execute_result"
-    }
-   ],
-   "source": [
-    "dataset[0:19]"
-   ]
-  },
   {
    "cell_type": "markdown",
    "metadata": {},
@@ -693,7 +603,9 @@
     "\n",
     "First we will define a function that generates images from our previously defined distribution. Notice that this function has three arguments: `dist_vmap`, `dist_novmap`, and `key`. As the name suggest, `dist_vmap` is an equinox filtered distribution object, which contains the leaves that we will vmap over (the ones we defined before with `vmap_filter`). Similarly, the `dist_novmap` contains the leaves we will not vmap over. We also need a different `key` per image, otherwise images would have the same noise, so we also vmap over that argument.\n",
     "\n",
-    "You could easily change this to a plain `imaging_pipeline` object. In that case you would not need the key, and you would to slightly change the definition of `vmap_filter`. Our `Simulate a batch of images` tutorial shows how to do so."
+    "    We perform this filtering inside our code, you only need to worry about defining your function properly.\n",
+    "\n",
+    "You could easily change this to a plain `imaging_pipeline` object. In that case you would not need the key, and you would to slightly change the definition of `vmap_filter`. Our `Simulate a batch of images` tutorial shows how to do so. There is also a short example of how to do it at the end of this tutorial."
    ]
   },
   {