recovar_coding_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "# Here is a short tutorial in case you want to use the code.\n",
    "\n",
    "# If you want to extend and use recovar, you should import this first\n",
    "import recovar.config\n",
    "\n",
    "from recovar import dataset\n",
    "from recovar.fourier_transform_utils import fourier_transform_utils\n",
    "import jax\n",
    "# Fourier transform utils version that runs on GPU by default (if you have one). If you pass numpy instead, it will run on cpu.\n",
    "ftu = fourier_transform_utils(jax.numpy) \n",
    "\n",
    "dataset_dict = dataset.get_default_dataset_option()\n",
    "\n",
    "# Fill these options with the path to preprocessed files\n",
    "experiment_directory = '/home/mg6942/mytigress/uniform/'\n",
    "dataset_dict['ctf_file'] = experiment_directory + 'ctf.pkl'\n",
    "dataset_dict['poses_file'] = experiment_directory + 'poses.pkl'\n",
    "dataset_dict['particles_file'] = experiment_directory + 'particles.128.mrcs'\n",
    "\n",
    "# Returns an object that knows everything about the dataset.\n",
    "cryo_dataset = dataset.load_dataset_from_dict(dataset_dict)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import jax.numpy as jnp\n",
    "# Run code on GPU (if you have one), in batch of 1000\n",
    "batch_size = 1000\n",
    "# Iterator used to send data in batch to GPU.\n",
    "dataset_iterator = cryo_dataset.get_dataset_generator(batch_size=batch_size) \n",
    "\n",
    "new_stack = np.empty([cryo_dataset.n_images, *cryo_dataset.image_shape])\n",
    "weiner_param = 1\n",
    "\n",
    "# Running code on GPU with JAX is very easy. By default, any jax.numpy.array will be allocated on GPU,\n",
    "# and can be used like a normal numpy.array to do operations on GPU. Once you want to send back to CPU,\n",
    "# either do jax.device_put(array, device = jax.devices(\"cpu\")[0]) or pass it to a numpy array.\n",
    "\n",
    "# Here is a simple example: CTF correcting an image stack \n",
    "for images, batch_image_ind in dataset_iterator:\n",
    "    # images are still on CPU at this point. you can do jnp.array(images) to send them to GPU explicitly, or they will be sent to GPU by any of the functions below.\n",
    "\n",
    "    # Do some computation on GPU    \n",
    "    CTFs = cryo_dataset.CTF_fun(cryo_dataset.CTF_params[batch_image_ind], cryo_dataset.image_shape, cryo_dataset.voxel_size) # Compute CTF\n",
    "    images = cryo_dataset.image_stack.process_images(images) # Compute DFT, masking\n",
    "    CTF_corrected_images = (CTFs / (CTFs**2 + weiner_param)) * images  # CTF correction\n",
    "    \n",
    "    CTF_corrected_images_real = ftu.get_idft2(CTF_corrected_images.reshape(CTF_corrected_images.shape[0], *cryo_dataset.image_shape )) # Back to real domain\n",
    "\n",
    "    # Send back to CPU\n",
    "    new_stack[batch_image_ind] = np.array(CTF_corrected_images_real.real)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# In recovar.core, you will find functions to do all the basic cryo-EM operations in batches on GPU: \n",
    "import recovar.core as core\n",
    "core.translate_images() \n",
    "core.slice_volume_by_nearest() # Slice a volume \n",
    "core.slice_volume_by_map() # For a second order discretization\n",
    "core.summed_adjoint_slice_by_nearest() # Summed adjoint of slicing: v = \\sum_i P_i^* im_i for some im_i\n",
    "core.batch_evaluate_ctf() # CTF function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Often, you can get a large speed up by ``jitting\" the inner loop. See https://jax.readthedocs.io/en/latest/jax-101/02-jitting.html\n",
    "# As an example, this is how you would jit the function above:\n",
    "# WIP:\n",
    "# (Unfinished)\n",
    "# @functools.partial(jax.jit, static_argnums = [3,4,6])\n",
    "# def weiner_CTF_correction(images, CTF_params, weiner_param,  image_shape voxel_size, CTF_fun ):\n",
    "#     CTFs = CTF_fun(CTF_params, image_shape, voxel_size) # Compute CTF\n",
    "#     # images = cryo_dataset.image_stack.process_images(images) # Compute DFT, masking\n",
    "#     CTF_corrected_images = (CTFs / (CTFs**2 + weiner_param)) * images  # CTF correction\n",
    "#     CTF_corrected_images_real = ftu.get_idft2(CTF_corrected_images.reshape(CTF_corrected_images.shape[0], *image_shape )) # Back to real domain\n",
    "#     return CTF_corrected_images_real\n"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "recovar2 [~/.conda/envs/recovar2/]",
   "language": "python",
   "name": "conda_recovar2"
  },
  "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.9.18"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}