analysis_v2.ipynb

{
 "cells": [
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Notebook showcasing the various options in Roodmus for analysis and visualisations\n",
    "In this notebook, the user can load metadata from one or several jobs from a processing pipeline done in RELION or cryoSPARC. This metadata, along with the ground-truth particle parameters are loaded into data frames, which allow for easy and conveniet plotting. We also provide several convenient functions to make plots.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### imports\n",
    "# general\n",
    "import os\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "\n",
    "# roodmus\n",
    "from roodmus.analysis.utils import load_data\n",
    "from roodmus.analysis.plot_ctf import plot_CTF, plot_defocus_scatter\n",
    "from roodmus.analysis.plot_picking import (\n",
    "    label_micrograph_picked, \n",
    "    label_micrograph_truth,\n",
    "    label_micrograph_truth_and_picked, \n",
    "    plot_precision, plot_recall, \n",
    "    plot_boundary_investigation,\n",
    "    plot_overlap_investigation, \n",
    "    plot_precision_and_recall, \n",
    "    plot_f1_score\n",
    ")\n",
    "from roodmus.analysis.plot_frames import plot_frame_distribution\n",
    "from roodmus.analysis.plot_classes import plot_2Dclass_precision, plot_2Dclasses_frames\n",
    "from roodmus.analysis.plot_alignment import plot_picked_pose_distribution, plot_true_pose_distribution\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### data loading\n",
    "config_dir = \"path/to/config/dir\"\n",
    "meta_files = [\n",
    "    \"path/to/metadata/file1.cs\",\n",
    "    \"path/to/metadata/file2.cs\",\n",
    "    [\"path/to/metadata/file3.cs\", \"path/to/metadata/file4.cs\"], # multiple metadata files are fused\n",
    "]\n",
    "\n",
    "jobtypes = {\n",
    "    \"path/to/metadata/file1.cs\": \"jobtype1\",\n",
    "    \"path/to/metadata/file2.cs\": \"jobtype2\",\n",
    "    \"path/to/metadata/file3.cs\": \"jobtype3\",\n",
    "}\n",
    "\n",
    "particle_diameter = 100 # approximate particle diameter in Angstroms\n",
    "ugraph_shape = (4000, 4000) # shape of the micrograph in pixels. Only needs to be given if the metadata file is a .star file\n",
    "verbose = True # prints out progress statements\n",
    "ignore_missing_files = True # if .mrc files are missing, the analysis will still be performed\n",
    "enable_tqdm = True # enables tqdm progress bars\n",
    "\n",
    "for i, meta_file in enumerate(meta_files):\n",
    "    if i == 0:\n",
    "        analysis = load_data(meta_file, config_dir, particle_diameter, ugraph_shape=ugraph_shape, verbose=verbose, enable_tqdm=enable_tqdm, ignore_missing_files=ignore_missing_files) # creates the class\n",
    "    else:\n",
    "        analysis.add_data(meta_file, config_dir, verbose=verbose) # updates the class with the next metadata file\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### turn the loaded data into a pandas dataframe\n",
    "df_picked = pd.DataFrame(analysis.results_picking)\n",
    "df_truth = pd.DataFrame(analysis.results_truth)\n",
    "df_picked.tail()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### saving the dataframes\n",
    "# it is recommended to save the dataframes after running the rest of the notebook, as they may be modified by downstream analysis\n",
    "\n",
    "df_picked.to_csv(\"picked_particles.csv\")\n",
    "df_truth.to_csv(\"truth_particles.csv\")"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### CTF estimation\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### scatter plot of the estimated vs. the true defocus values\n",
    "meta_index = 0 # index of the metadata file to plot\n",
    "\n",
    "palette = \"RdBu\"\n",
    "\n",
    "fig, ax = plot_defocus_scatter(df_picked,\n",
    "                                meta_files[meta_index],\n",
    "                                df_truth,\n",
    "                                palette=palette)\n",
    "\n",
    "fig.savefig(\"defocus_scatter.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"defocus_scatter.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the CTF estimation for a single micrograph\n",
    "meta_index = 0 # index of the metadata file to plot\n",
    "ugraph_index = 3 # which micrograph to plot\n",
    "\n",
    "fig, ax = plot_CTF(df_picked, meta_files[meta_index],\n",
    "                    df_truth, config_dir, ugraph_index)\n",
    "\n",
    "fig.savefig(\"CTF.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"CTF.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the CTF for the particle with the largest defocus error (should take no more than a few seconds)\n",
    "max_error_index = 0\n",
    "max_error = 0\n",
    "for i, groupname in enumerate(df_picked.groupby([\"ugraph_filename\"]).groups.keys()):\n",
    "    defocus_estimated = df_picked.groupby([\"ugraph_filename\"]).get_group(groupname)[\"defocusU\"].mean()\n",
    "    defcous_true = np.abs(df_truth.groupby([\"ugraph_filename\"]).get_group(groupname)[\"defocus\"].mean())\n",
    "    error = np.abs(defocus_estimated - defcous_true)\n",
    "    if error > max_error:\n",
    "        max_error = error\n",
    "        max_error_index = i\n",
    "\n",
    "fig, ax = plot_CTF(df_picked, None, df_truth, config_dir, max_error_index)\n",
    "\n",
    "# save high quality figure\n",
    "fig.savefig(\"CTF_max_error.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"CTF_max_error.pdf\", bbox_inches=\"tight\")\n",
    "\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Particle picking"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the picked particles\n",
    "ugraph_index = 0 # which micrograph to plot\n",
    "metadata_index = 0 # which metadata file to plot\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "\n",
    "fig, ax = label_micrograph_picked(df_picked, meta_files[meta_index],\n",
    "                                   ugraph_index, config_dir, box_width=48, box_height=48, verbose=verbose)\n",
    "\n",
    "fig.savefig(\"picked_particles.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"picked_particles.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the truth particles\n",
    "ugraph_index = 3 # which micrograph to plot\n",
    "\n",
    "fig, ax = label_micrograph_truth(df_truth, ugraph_index, config_dir, box_width=32, box_height=32, verbose=verbose)\n",
    "\n",
    "fig.savefig(\"truth_particles.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"truth_particles.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the truth and picked particles\n",
    "ugraph_index = 3 # which micrograph to plot\n",
    "metadata_index = 3 # which metadata file to plot\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "    \n",
    "fig, ax = label_micrograph_truth_and_picked(df_picked, meta_files[meta_index],\n",
    "                                             df_truth, ugraph_index, config_dir, box_width=48, box_height=48, verbose=verbose)\n",
    "\n",
    "fig.savefig(\"truth_and_picked_particles.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"truth_and_picked_particles.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### compute precision and recall (may take a few minutes)\n",
    "df_precision, df_picked = analysis.compute_precision(df_picked, df_truth, verbose=verbose)\n",
    "df_precision.tail()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot boxplot for precision and recall\n",
    "order = []\n",
    "for r in meta_files:\n",
    "    if type(r) == str:\n",
    "        order.append(r)\n",
    "    else:\n",
    "        order.append(r[0])    \n",
    "fig, ax = plot_precision(df_precision, jobtypes, order)\n",
    "xticklabels = ax.get_xticklabels()\n",
    "ax.set_xticklabels(xticklabels, fontsize=14)\n",
    "ax.set_title(\"\")\n",
    "fig.set_size_inches(7, 7)\n",
    "\n",
    "fig.savefig(\"precision.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"precision.pdf\", bbox_inches=\"tight\")\n",
    "\n",
    "fig, ax = plot_recall(df_precision, jobtypes, order)\n",
    "xticklabels = ax.get_xticklabels()\n",
    "ax.set_xticklabels(xticklabels, fontsize=14)\n",
    "ax.set_title(\"\")\n",
    "fig.set_size_inches(7, 7)\n",
    "\n",
    "fig.savefig(\"recall.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"recall.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### alternatively, plot the precision and recall in the same plot\n",
    "fig, ax = plot_precision_and_recall(df_precision, jobtypes)\n",
    "\n",
    "fig.savefig(\"precision_and_recall.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"precision_and_recall.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot f1-score\n",
    "fig, ax = plot_f1_score(df_precision, jobtypes)\n",
    "\n",
    "fig.savefig(\"f1_score.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"f1_score.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the picked particles, now with the TP and FP marked in green and red\n",
    "ugraph_index = 3 # which micrograph to plot\n",
    "metadata_index = 4 # which metadata file to plot\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "\n",
    "fig, ax = label_micrograph_picked(df_picked, meta_files[meta_index], ugraph_index, config_dir, box_width=48, box_height=48, verbose=verbose)\n",
    "\n",
    "fig.savefig(\"picked_particles_TP_FP.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"picked_particles_TP_FP.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the distribution of the particles in the ugraphs in x, y, and z directions\n",
    "metadata_index = 0 # which metadata file to plot\n",
    "bin_width = [100, 100, 10] # bin width for x, y, z\n",
    "axis = [\"x\", \"y\", \"z\"]\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "    \n",
    "for a, bnwdth in zip(axis, bin_width):\n",
    "    fig, ax = plot_boundary_investigation(df_truth, df_picked, metadata_filename, jobtypes, bnwdth, axis=a)\n",
    "\n",
    "    fig.savefig(\"boundary_investigation_{}.png\".format(a), dpi=600, bbox_inches=\"tight\")\n",
    "    fig.savefig(\"boundary_investigation_{}.pdf\".format(a), bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "df_overlap = analysis.compute_overlap(df_picked, df_truth, verbose=verbose)\n",
    "df_overlap.head()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the overlap between the picked and truth particles\n",
    "metadata_index = None # which metadata file to plot. If None, all metadata files are plotted\n",
    "\n",
    "if metadata_index is None:\n",
    "    metadata_filename = None\n",
    "else:\n",
    "    metadata_filename = meta_files[metadata_index]\n",
    "    if isinstance(metadata_filename, list):\n",
    "        metadata_filename = metadata_filename[0]\n",
    "\n",
    "fig, ax = plot_overlap_investigation(df_overlap, metadata_filename, job_types=jobtypes)\n",
    "\n",
    "fig.savefig(\"overlap_investigation.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"overlap_investigation.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the distribution of trajectory frames in a metadata file\n",
    "metadata_index = 6 # which metadata file to plot\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "print(metadata_filename)\n",
    "\n",
    "fig, ax = plot_frame_distribution(df_picked, metadata_filename, df_truth, particle_diameter, jobtypes)\n",
    "\n",
    "fig.savefig(\"frame_distribution.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"frame_distribution.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 2D classification"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the precision per class\n",
    "from importlib import reload\n",
    "import roodmus.analysis.plot_classes\n",
    "reload(roodmus.analysis.plot_classes)\n",
    "from roodmus.analysis.plot_classes import plot_2Dclass_precision\n",
    "metadata_index = 8 # which metadata file to plot. Must have a class2D column\n",
    "palette = \"coolwarm\"\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "\n",
    "fig, ax = plot_2Dclass_precision(df_picked, metadata_filename, jobtypes)\n",
    "\n",
    "fig.savefig(\"2Dclass_precision.png\", dpi=600, bbox_inches=\"tight\")\n",
    "fig.savefig(\"2Dclass_precision.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the distribution of frames over the 2D classes\n",
    "from importlib import reload\n",
    "import roodmus.analysis.plot_classes\n",
    "reload(roodmus.analysis.plot_classes)\n",
    "from roodmus.analysis.plot_classes import plot_2Dclasses_frames\n",
    "metadata_index = 1 # which metadata file to plot. Must have a class2D column\n",
    "palette = \"YlOrRd\"\n",
    "\n",
    "metadata_filename = meta_files[metadata_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "\n",
    "fig, ax = plot_2Dclasses_frames(df_picked, metadata_filename, bin_factor=100, palette=palette)\n",
    "ax.set_xlabel(\"trajectory frame\", fontsize=14)\n",
    "ax.set_xticks([])\n",
    "ax.set_ylabel(\"2D class\", fontsize=14)\n",
    "\n",
    "fig.savefig(\"2Dclass_frames.png\", dpi=600, bbox_inches=\"tight\")\t\n",
    "fig.savefig(\"2Dclass_frames.pdf\", bbox_inches=\"tight\")\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### 3D alignment"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "### plot the distribution of particle poses in the picked and truth particles\n",
    "meta_index = 0\n",
    "\n",
    "metadata_filename = meta_files[meta_index]\n",
    "if isinstance(metadata_filename, list):\n",
    "    metadata_filename = metadata_filename[0]\n",
    "print(metadata_filename)\n",
    "\n",
    "grid, vmin, vmax = plot_picked_pose_distribution(df_picked, metadata_filename)\n",
    "\n",
    "grid.fig.savefig(\"picked_pose_distribution.png\", dpi=600, bbox_inches=\"tight\")\n",
    "grid.fig.savefig(\"picked_pose_distribution.pdf\", bbox_inches=\"tight\")\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "grid = plot_true_pose_distribution(df_truth, vmin, vmax)\n",
    "\n",
    "grid.fig.savefig(\"true_pose_distribution.png\", dpi=600, bbox_inches=\"tight\")\n",
    "grid.fig.savefig(\"true_pose_distribution.pdf\", bbox_inches=\"tight\")\n"
   ]
  }
 ],
 "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.10.13"
  },
  "vscode": {
   "interpreter": {
    "hash": "e5112804f8795a0ac2f26cc90b499f055a6ce0f1dffbe6381875821b4c148204"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}