trained on cryobife prior

flatironinstitute · Aug 9, 2023 · ec4d590 · ec4d590
1 parent 658ccdb
commit ec4d590
Show file tree

Hide file tree

Showing 8 changed files with 317 additions and 192 deletions.
diff --git a/Lukes_folder/cryoBIFE_test.ipynb b/Lukes_folder/cryoBIFE_test.ipynb
diff --git a/Lukes_folder/cryobife_loss b/Lukes_folder/cryobife_loss
diff --git a/Lukes_folder/image_params_train_cryobife.json b/Lukes_folder/image_params_train_cryobife.json
@@ -0,0 +1,14 @@
+{
+    "N_PIXELS": 128,
+    "PIXEL_SIZE": 1.5,
+    "SIGMA": [1.0, 5.0],
+    "MODEL_FILE": "../data/protein_models/hsp90_models.npy",
+    "SHIFT": 30,
+    "DEFOCUS": [0.5, 2],
+    "SNR": [0.01, 0.5],
+    "RADIUS_MASK": 64,
+    "AMP": 0.1,
+    "B_FACTOR": [1.0, 100.0],
+    "ELECWAVE": 0.019866,
+    "CRYOBIFE_INDEX_PRIOR": "True"
+}
diff --git a/Lukes_folder/train_BIFE_prior.ipynb b/Lukes_folder/train_BIFE_prior.ipynb
@@ -2,15 +2,15 @@
  "cells": [
   {
    "cell_type": "code",
-   "execution_count": 1,
+   "execution_count": 3,
    "metadata": {},
    "outputs": [
     {
-     "name": "stderr",
+     "name": "stdout",
      "output_type": "stream",
      "text": [
-      "/Data/Packages/Utilities/miniconda3/envs/cryosbi_env/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
-      "  from .autonotebook import tqdm as notebook_tqdm\n"
+      "The autoreload extension is already loaded. To reload it, use:\n",
+      "  %reload_ext autoreload\n"
      ]
     }
    ],
@@ -19,9 +19,10 @@
     "import matplotlib.pyplot as plt\n",
     "\n",
     "from cryo_sbi import CryoEmSimulator\n",
-    "from cryo_sbi import gen_training_set\n",
-    "from cryo_sbi.inference.NPE_train_from_disk import npe_train_from_disk\n",
-    "from cryo_sbi.inference.NPE_train_without_saving import npe_train_no_saving"
+    "#from cryo_sbi import gen_training_set\n",
+    "from cryo_sbi.inference.train_npe_model import npe_train_no_saving\n",
+    "%load_ext autoreload\n",
+    "%autoreload 2"
    ]
   },
   {
@@ -34,99 +35,43 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": 4,
    "metadata": {},
    "outputs": [
     {
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "hsp90_models.npy\n"
+      "Using CryoBife reference prior\n",
+      "Training neural network:\n"
      ]
     },
     {
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 10000/10000 [01:13<00:00, 136.60pair/s]\n",
-      "100%|██████████| 100/100 [00:00<00:00, 4779.94pair/s]\n"
+      "100%|██████████| 100/100 [56:24<00:00, 33.85s/epoch, loss=-.508]\n"
      ]
     }
    ],
    "source": [
-    "gen_training_set(\n",
-    "    config_file=\"config_file.json\",\n",
-    "    num_train_samples=10000,\n",
-    "    num_val_samples=100,\n",
-    "    file_name=\"tut_imgs\",\n",
-    "    save_as_tensor=False,\n",
-    "    n_workers=2,\n",
-    "    batch_size=100,\n",
-    ")"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": 6,
-   "metadata": {},
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "Training neural netowrk:\n"
-     ]
-    },
-    {
-     "name": "stderr",
-     "output_type": "stream",
-     "text": [
-      "100%|██████████| 30/30 [04:44<00:00,  9.48s/epoch, train_loss=-.981, val_loss=0.108]  \n"
-     ]
-    }
-   ],
-   "source": [
-    "npe_train_from_disk(\n",
-    "    train_config=\"resnet18_encoder.json\",\n",
-    "    epochs=30,\n",
-    "    train_data_dir=\"tut_imgs_train.h5\",\n",
-    "    val_data_dir=\"tut_imgs_valid.h5\",\n",
-    "    estimator_file=\"tut_estimator\",\n",
-    "    loss_file=\"tut_loss\",\n",
-    "    train_from_checkpoint=False,\n",
-    "    model_state_dict=None,\n",
-    "    n_workers=2,\n",
-    ")"
-   ]
-  },
-  {
-   "attachments": {},
-   "cell_type": "markdown",
-   "metadata": {},
-   "source": [
-    "## Training without saving images"
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
-   "metadata": {},
-   "outputs": [],
-   "source": [
+    "torch.cuda.empty_cache()\n",
     "npe_train_no_saving(\n",
-    "    image_config=\"image_params_snr01_128.json\",\n",
+    "    image_config=\"image_params_train_cryobife.json\",\n",
     "    train_config=\"resnet18_encoder.json\",\n",
-    "    epochs=350,\n",
-    "    estimator_file=\"resnet18_encoder.estimator\",\n",
-    "    loss_file=\"resnet18_encoder.estimator\",\n",
-    "    n_workers=2,  # CHANGE\n",
-    ")"
+    "    epochs=100,\n",
+    "    estimator_file=\"cryobife_estimator.estimator\",\n",
+    "    loss_file=\"cryobife_loss\",\n",
+    "    device=\"cuda\",\n",
+    "    simulation_batch_size=1024\n",
+    ")\n",
+    "\n"
    ]
   }
  ],
  "metadata": {
   "kernelspec": {
-   "display_name": "cryosbi_env",
+   "display_name": "Python 3.9.15 64-bit",
    "language": "python",
    "name": "python3"
   },
@@ -140,9 +85,14 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.9.16"
+   "version": "3.9.15"
   },
-  "orig_nbformat": 4
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "7b7fbdd20bcc2083504065e64dd68e11295ac29c39a09e225403f090756a3e6a"
+   }
+  }
  },
  "nbformat": 4,
  "nbformat_minor": 2

diff --git a/Lukes_folder/trying_it_out.ipynb b/Lukes_folder/trying_it_out.ipynb
diff --git a/src/cryo_sbi/inference/priors.py b/src/cryo_sbi/inference/priors.py
@@ -95,10 +95,26 @@ def get_image_priors(
         ndims=1,
     )
 
-    index_prior = zuko.distributions.BoxUniform(
-        lower=torch.tensor([0], dtype=torch.float32, device=device),
-        upper=torch.tensor([max_index], dtype=torch.float32, device=device),
-    )
+    # Giving an option in the JSON file to use a different prior on index, for now hardcoded from Cryo-Bife paper
+    if image_config.get("CRYOBIFE_INDEX_PRIOR"):
+        mean_a = 6.0
+        mean_b = 15.0
+        std_a = 2.0
+        std_b = 3.0
+        weight_a = 1.0
+        weight_b = 1.0 / 3.0
+        weights = torch.tensor([weight_a, weight_b], dtype=torch.float32, device=device)
+        gaussians = zuko.distributions.Normal(torch.tensor([[mean_a, mean_b]], dtype=torch.float32, device=device), torch.tensor([[std_a, std_b]], dtype=torch.float32, device=device))
+        lower = torch.tensor([0], dtype=torch.float32, device=device)
+        upper=torch.tensor([19.0], dtype=torch.float32, device=device)
+        index_prior = zuko.distributions.Mixture(zuko.distributions.Truncated(gaussians, lower=lower, upper=upper), logits = torch.log(weights))
+        print("Using CryoBife reference prior")
+    else:
+        index_prior = zuko.distributions.BoxUniform(
+            lower=torch.tensor([0], dtype=torch.float32, device=device),
+            upper=torch.tensor([max_index], dtype=torch.float32, device=device),
+        )
+
     quaternion_prior = QuaternionPrior(device)
     if image_config.get("ROTATIONS") and isinstance(image_config["ROTATIONS"], list) and len(image_config["ROTATIONS"]) == 4:
         test_quat = image_config["ROTATIONS"]

diff --git a/src/cryo_sbi/inference/train_npe_model.py b/src/cryo_sbi/inference/train_npe_model.py
@@ -88,13 +88,22 @@ def npe_train_no_saving(
     assert simulation_batch_size % train_config["BATCH_SIZE"] == 0
 
     if image_config["MODEL_FILE"].endswith("npy"):
-         models = (
-             torch.from_numpy(
-                 np.load(image_config["MODEL_FILE"]),
-             )
-             .to(device)
-             .to(torch.float32)
-         )
+
+        # Lukes hsp90 hack!
+        if "hsp90" in image_config["MODEL_FILE"]:
+            models = (
+                torch.from_numpy(
+                    np.load(image_config["MODEL_FILE"])[:, 0, :, :],
+                ).to(device).to(torch.float32)
+            )
+        else:
+            models = (
+                torch.from_numpy(
+                    np.load(image_config["MODEL_FILE"]),
+                )
+                .to(device)
+                .to(torch.float32)
+                    )
     else:
         models = torch.load(
         image_config["MODEL_FILE"], 
@@ -131,7 +140,7 @@ def npe_train_no_saving(
     step = GDStep(optimizer, clip=train_config["CLIP_GRADIENT"])
     mean_loss = []
 
-    print("Training neural netowrk:")
+    print("Training neural network:")
     estimator.train()
     with tqdm(range(epochs), unit="epoch") as tq:
         for epoch in tq:
@@ -153,6 +162,7 @@ def npe_train_no_saving(
                 )
                 for _indices, _images in zip(indices.split(train_config["BATCH_SIZE"]), images.split(train_config["BATCH_SIZE"])):
                     losses.append(step(loss(_indices.to(device, non_blocking=True), _images.to(device, non_blocking=True))))
+                #torch.cuda.empty_cache()
             losses = torch.stack(losses)
 
             tq.set_postfix(loss=losses.mean().item())

diff --git a/src/cryo_sbi/wpa_simulator/cryo_em_simulator.py b/src/cryo_sbi/wpa_simulator/cryo_em_simulator.py
@@ -116,6 +116,7 @@ def simulate(self, num_sim, indices=None, return_parameters=False):
             assert (
                 indices.ndim == 2
             ), "Indices are not a 2D tensor, converting to 2D tensor. With shape (batch_size, 1)."
+            parameters[0] = indices
         else:
             indices = parameters[0] 
         images = cryo_em_simulator(