{ "cells": [ { "cell_type": "code", "execution_count": 1, "id": "4d6d065a-be52-4462-9d33-fd12928aea6a", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:10:35.353150Z", "iopub.status.busy": "2026-06-25T14:10:35.352952Z", "iopub.status.idle": "2026-06-25T14:10:35.356293Z", "shell.execute_reply": "2026-06-25T14:10:35.355625Z" } }, "outputs": [], "source": [ "import os\n", "os.environ['CUDA_VISIBLE_DEVICES'] = '0'" ] }, { "cell_type": "markdown", "id": "036f012a-0623-494b-95f2-35ad8c4a604f", "metadata": {}, "source": [ "### Tutorial B6: Design\n", "\n", "> ersatz (substitute) + predict + ersatz (substitute) + predict + ..." ] }, { "cell_type": "markdown", "id": "1f33fe28-fe87-4f70-948a-8b18cfce5154", "metadata": {}, "source": [ "Most of the functions we've discussed so far center around extracting what a sequence-based machine learning model has learned after training. Marginalization experiments show the individual effect that each motif has on model predictions; ISM shows the individual effect that each mutation has on model predictions within an observed sequence. However, these models can be used for more than just identifying relevent sequence features and their syntax rules.\n", "\n", "For example, trained sequence-based models can be used in the design setting. Although these are numerous methods for designing sequences, each with their own set of important details, the overall approach is to train a model that predicts some sort of readout of interest and then perturb the sequence input until the predicted output matches what is desired by the user. Sometimes, this involves starting with an informative sequence and designing edits, other times this involves the de novo generation of sequences from scratch that exhibit the desired characters. Regardless, the trained model is considered to be an \"oracle\" that outputs how good the in-progress sequence is.\n", "\n", "#### Greedy Substitution\n", "\n", "The most conceptually simple design method is that of greedy substitution. Given an initial sequence, a predictive model, an output goal, and a set of motifs that one could substitute into the sequence, the task is to find which motifs and what exact positioning should be used to achieved a desired output from the predictive model. Because the aim here is to be conceptually simple, the method proceeds across several iterations where at each iteration, every motif is tried at every position in the sequence, and the motif+position pair that yields the largest improvement in terms of getting the model towards its goal. Naturally, this can be very time consuming but is the most straightforward approach.\n", "\n", "Let's start by loading up the Beluga model again." ] }, { "cell_type": "code", "execution_count": 2, "id": "4c78f4e4-4cb3-4013-89d6-a3d72a68ec4c", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:10:35.358182Z", "iopub.status.busy": "2026-06-25T14:10:35.358050Z", "iopub.status.idle": "2026-06-25T14:10:36.677351Z", "shell.execute_reply": "2026-06-25T14:10:36.676559Z" } }, "outputs": [ { "data": { "text/plain": [ "" ] }, "execution_count": 2, "metadata": {}, "output_type": "execute_result" } ], "source": [ "import torch\n", "from model import Beluga\n", "\n", "model = Beluga()\n", "model.load_state_dict(torch.load(\"deepsea.beluga.torch\"))" ] }, { "cell_type": "markdown", "id": "889446fe-2114-44c6-bec4-1aa809b6f6fb", "metadata": {}, "source": [ "Let's see if we can get this model to design a sequence that predicts high AP-1 binding. Yes, this task is unrealistically simple because AP-1 factors bind to a known motif and one could just insert that into the sequence, but it's a good initial task to show how to use the function. \n", "\n", "For the purpose of this example, we can randomly generate a single sequence and use a subset of predefined motifs from JASPAR. Here, we have inserted a motif that we known the Beluga AP-1 tasks respond to as the third motif in the list. " ] }, { "cell_type": "code", "execution_count": 3, "id": "61fd30eb-ff58-4dfe-965a-c40c8016024d", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:10:36.678955Z", "iopub.status.busy": "2026-06-25T14:10:36.678760Z", "iopub.status.idle": "2026-06-25T14:10:37.300960Z", "shell.execute_reply": "2026-06-25T14:10:37.296033Z" } }, "outputs": [], "source": [ "import numpy\n", "from tangermeme.utils import random_one_hot\n", "\n", "X = random_one_hot((1, 4, 2000), random_state=0).type(torch.float32)\n", "motifs = [\n", " 'GCTAATTAAC',\n", " 'ATGCCCACC',\n", " \"GTGACTCATC\",\n", " 'AGAACAGAATGTTCT',\n", " 'TGATGACGTCATCGC',\n", " 'ACATTCCA',\n", " 'GGGAGGAGGGAGAGGAGGAG',\n", " 'TAATCGATTA',\n", " 'CGTCTAGACA',\n", " 'TGCTATTTTTAG',\n", " 'CATTGTTTATTT',\n", " 'TTTCACACCTAGGTGTGAAA',\n", " 'GGCACGCGCC'\n", "]" ] }, { "cell_type": "markdown", "id": "521a90a2-a718-43ec-a84c-59a5897741d7", "metadata": {}, "source": [ "Next, we need a target that we are trying to get the model to predict with the designed sequence. In this case, we probably want all of the tasks for AP-1 factor proteins (e.g., c-/FOS, c-/JUN, etc). Since the predictions from Beluga are in logit space, we can set it to a relatively high positive value to indicate that we want binding of those specific tasks. Note that the first dimension is `1`. Although this method will accept a batch of sequences, the batch size must be `1` because, here, you are designing a single sequence." ] }, { "cell_type": "code", "execution_count": 4, "id": "e5114ea7-b91c-490a-b823-3650ab83d61c", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:10:37.306806Z", "iopub.status.busy": "2026-06-25T14:10:37.306458Z", "iopub.status.idle": "2026-06-25T14:10:37.311491Z", "shell.execute_reply": "2026-06-25T14:10:37.310961Z" } }, "outputs": [], "source": [ "names = numpy.loadtxt(\"beluga_target_names.txt\", delimiter=',', dtype=str)\n", "idxs = torch.tensor(['jun' in n.lower() or 'fos' in n.lower() for n in names])\n", "\n", "y = torch.zeros(1, 2002)\n", "y[:, idxs] = 3" ] }, { "cell_type": "markdown", "id": "ff622d6e-b217-41e9-8705-3840b24de721", "metadata": {}, "source": [ "Now, we can call the function. This function has a similar signature to other `tangermeme` functions in that it takes in the predictive model, the initial sequence, and a list of motifs. It differs from other functions in that it also takes in `y`, which is the target that we want the predictive model to output. Optionally, `greedy_substitution` can take in a mask that is applied to the predictions from the model and the target. Basically, if your model makes predictions for more outputs than you care about, this mask lets you remove some of those outputs from the loss. In our case, we can remove all predictions not related to the AP-1 factors and say that we want to design a sequence with the highest binding of AP-1 factors regardless of what happens to all other tasks." ] }, { "cell_type": "code", "execution_count": 5, "id": "f830db0f-c013-4962-be98-e6873b6df227", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:10:37.312595Z", "iopub.status.busy": "2026-06-25T14:10:37.312509Z", "iopub.status.idle": "2026-06-25T14:12:15.947447Z", "shell.execute_reply": "2026-06-25T14:12:15.946240Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 72.97, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/26 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "from matplotlib import pyplot as plt\n", "import seaborn; seaborn.set_style('whitegrid')\n", "\n", "from tangermeme.predict import predict\n", "\n", "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "774f60a5-0f09-4357-be16-d50c98ef1ae9", "metadata": {}, "source": [ "Not only do we see that the tasks related to AP-1 binding have increased logits, we can see that they seem to be close to the desired value of 3. Interestingly, not all of the AP-1 tasks are increased to that level, and there seems to be quite a bit of a jitter in all of the other tasks. " ] }, { "cell_type": "markdown", "id": "8623eb11-4db2-4256-a3b4-d001ced74a4f", "metadata": {}, "source": [ "### Greedy Saturation Mutagenesis / Evolutionary Algorithms\n", "\n", "Rather than implanting entire motifs into sequences, some approaches consider changing only one nucleotide at a time. This is sometimes done probabilistically, where only a subset of potential mutations are considered, and is sometimes comprehensive, e.g., saturation mutagenesis. Regardless, each approach proceeds similarly to the greedy approach described above: a nucleotide is changed, the change in score is recorded, and the best individual mutation is made at the end of each round.\n", "\n", "This approach of only changing a single nucleotide in each round can be performed by simply passing in the four nucleotides as options rather than entire motifs." ] }, { "cell_type": "code", "execution_count": 7, "id": "e59f16bb-46c5-46c8-85c8-2bd285e0716e", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:12:16.994758Z", "iopub.status.busy": "2026-06-25T14:12:16.994614Z", "iopub.status.idle": "2026-06-25T14:13:04.590883Z", "shell.execute_reply": "2026-06-25T14:13:04.590156Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 72.97, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/4 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "b4a5264b-3da6-4f0d-aee6-25389871ca02", "metadata": {}, "source": [ "Looks like we are starting to get there, but are not as close as implanting ten entire motifs.\n", "\n", "This method has strengths and weaknesses. The weaknesses are that it may take many more iterations to get good results and that it gives potentially too much flexiblility of the model, allowing it to overfit to spurious correlations (e.g., finding a random mutation that is not part of the model but happens to make the model predict closer to the desired output). The strengths are that it does not rely on a set of pre-identified set of motifs and so can make changes that fall outside the set of motifs, e.g., adding in a lower affinity version of a motif or an alternate version of the motif that is actually stronger than the one provided. " ] }, { "cell_type": "markdown", "id": "a369065f-723f-4021-b0cc-45a2035c1ddc", "metadata": {}, "source": [ "#### Balancing Losses\n", "\n", "Is it possible to design a sequence that keeps the other tasks as close to their original values as possible while still increasing the predictions for AP-1 binding? To do this, rather than using a mask to ignore the outputs from non-AP-1 tasks, we will set the target values for that task to the original predictions from the model." ] }, { "cell_type": "code", "execution_count": 9, "id": "f5e43a8f-137e-4d19-9f57-a3adda8f30d8", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:13:04.982322Z", "iopub.status.busy": "2026-06-25T14:13:04.982244Z", "iopub.status.idle": "2026-06-25T14:16:08.976800Z", "shell.execute_reply": "2026-06-25T14:16:08.975817Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 1.13, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/26 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat2).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "201fd87f-b7d1-4200-8267-c57c357a0fa1", "metadata": {}, "source": [ "And it seems like we've reached a new compromise! The AP-1 related tasks are nicely separated from the other tasks (except for two of them) and, although they do not reach the desired objective value of `3`, the flip side is that the other tasks seem to be much closer to their original values.\n", "\n", "\n", "#### Input Masks\n", "\n", "Sometimes, we may not want to consider substitutions at every position in the sequence. Perhaps we know that the most important part of the sequence is in the middle, or have known motifs in our sequence that we want to protect, or maybe we just do not need to be as precise as considering any possible possible. Regardless of the reason for having a mask, using one can speed up design significantly.\n", "\n", "Normally, the input to the model is 2000 bp. What happens if we only use the middle 200 bp for design?" ] }, { "cell_type": "code", "execution_count": 11, "id": "a54c6ad3-f0f4-416c-a095-0725aec55b8c", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:16:09.407196Z", "iopub.status.busy": "2026-06-25T14:16:09.406946Z", "iopub.status.idle": "2026-06-25T14:16:43.009684Z", "shell.execute_reply": "2026-06-25T14:16:43.008699Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 1.13, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/26 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat2).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "962201b0-af02-4360-944e-940e48064754", "metadata": {}, "source": [ "The results look quite similar to before. However, this is likely because the original sequence we used was completely randomly generated and the improvements are likely coming from simply adding AP-1 motifs nearby to each other, so limiting the possible space isn't a huge constraint. In settings with actual structured sequences as input or where running more iterations to get the maximum possible improvement is necessary, having more sequence can be helpful." ] }, { "cell_type": "markdown", "id": "0c20d7d3-84ea-4832-9e4a-bc29af0cabaf", "metadata": {}, "source": [ "#### Loss Functions\n", "\n", "By default, the loss used is `torch.nn.MSELoss` without reduction because it is a reasonable and general-purpose loss. However, you can use whatever loss you'd like by passing the loss function into `loss=...`. This can be any loss implemented in PyTorch that is reasonable for your model or any custom function with the signature `def loss(y, y_hat)` with `y.shape == (n, ...)` that returns a tensor with size `(n, ...)`. Basically, every example that gets passed into the loss function needs to have one or more losses associated with it. Everything except for the first dimension gets averaged such that you get one value per example. This is a critical detail because, by default, PyTorch losses will average over the batch dimension as well and only give you a single number regardless of the number of examples passed in. The design function needs to know the example with the minimal loss and so any function -- custom or not -- needs to return the loss for each example. \n", "\n", "To demonstrate, let's do the design but use a different loss function." ] }, { "cell_type": "code", "execution_count": 13, "id": "71bf8273-5963-4f2f-bad0-cc9710f23b36", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:16:43.435171Z", "iopub.status.busy": "2026-06-25T14:16:43.435080Z", "iopub.status.idle": "2026-06-25T14:17:21.856245Z", "shell.execute_reply": "2026-06-25T14:17:21.855246Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 0.1229, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/26 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat2).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "15c6d933-b6fd-4a9d-834f-ccc5f30401c4", "metadata": {}, "source": [ "#### Elimination of Motifs\n", "\n", "An important part of this greedy procedure is that it does not only involve inserting new motifs into uninformative sequence, but that it will eliminate motifs that yield activity that the model does not want. As an illustration, let's consider the simple case where a motif drives unwanted AP-1 binding activity and the goal is to design a sequence that does not exhibit AP-1 binding.\n", "\n", "Here, we will start with a sequence that has the AP-1 motif in it already." ] }, { "cell_type": "code", "execution_count": 15, "id": "660a4744-40bd-4b17-ae7b-111e98325c2a", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:17:22.289140Z", "iopub.status.busy": "2026-06-25T14:17:22.289060Z", "iopub.status.idle": "2026-06-25T14:17:22.291781Z", "shell.execute_reply": "2026-06-25T14:17:22.291247Z" } }, "outputs": [], "source": [ "from tangermeme.ersatz import substitute\n", "\n", "X = random_one_hot((1, 4, 2000), random_state=0).type(torch.float32)\n", "X = substitute(X, \"GTGACTCATC\")" ] }, { "cell_type": "markdown", "id": "b79852a6-5b03-492b-912e-e2ba18161c47", "metadata": {}, "source": [ "Now, we will specify that we want baseline levels of AP-1 binding while preserving everything else." ] }, { "cell_type": "code", "execution_count": 16, "id": "b20f862c-2c62-4a09-afd4-e9ae7c3aa9bb", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:17:22.292817Z", "iopub.status.busy": "2026-06-25T14:17:22.292750Z", "iopub.status.idle": "2026-06-25T14:17:42.647251Z", "shell.execute_reply": "2026-06-25T14:17:42.646247Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 0.1575, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/26 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat3).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "e12f20ca-bb3f-46fe-876e-0b49501005ff", "metadata": {}, "source": [ "And it looks like the predicted signal is lower after the elimination of the motif!\n", "\n", "We can do a more targetted approach by using the greedy saturated mutagenesis approach." ] }, { "cell_type": "code", "execution_count": 18, "id": "0413b121-b4f1-4c91-b33b-1e31facd17dd", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:17:43.067027Z", "iopub.status.busy": "2026-06-25T14:17:43.066944Z", "iopub.status.idle": "2026-06-25T14:17:45.881296Z", "shell.execute_reply": "2026-06-25T14:17:45.880607Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Iteration 0 -- Loss: 0.1575, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/4 [00:00" ] }, "metadata": {}, "output_type": "display_data" } ], "source": [ "X_hat3 = greedy_substitution(model, X, y, ['A', 'C', 'G', 'T'], reverse_complement=False, max_iter=1, verbose=True)\n", "\n", "y0 = predict(model, X).numpy(force=True)[0]\n", "y1 = predict(model, X_hat3).numpy(force=True)[0]\n", "\n", "plt.plot([-9, 5], [-9, 5], c='0.3')\n", "plt.scatter(y0, y1, s=1)\n", "plt.scatter(y0[idxs], y1[idxs], s=2, color='r', label='AP-1 Targets')\n", "plt.legend()\n", "plt.xlabel(\"Task Logits Before Design\")\n", "plt.ylabel(\"Task Logits After Design\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "156dac71-02d5-43f8-bc6c-3c437e6d837c", "metadata": {}, "source": [ "With only a single mutation knocking out a key position in the motif, we are able to achieve similar results. " ] }, { "cell_type": "markdown", "id": "994ea15b-c972-484d-b176-867e2a4dd214", "metadata": {}, "source": [ "#### Half Precision\n", "\n", "A strength of this prediction-based design is that it is fairly robust to operating at lower precision because we are simply looking for the best motif and position, rather than needing a very precise value, and if two choices happen to be similar and swapped at lower precision the consequences for design are minimal." ] }, { "cell_type": "code", "execution_count": 19, "id": "1a582771-5e8b-49a0-828c-fc53e114b4f9", "metadata": { "execution": { "iopub.execute_input": "2026-06-25T14:17:45.882629Z", "iopub.status.busy": "2026-06-25T14:17:45.882520Z", "iopub.status.idle": "2026-06-25T14:19:36.896742Z", "shell.execute_reply": "2026-06-25T14:19:36.895994Z" } }, "outputs": [ { "name": "stdout", "output_type": "stream", "text": [ "Full Precision\n", "Iteration 0 -- Loss: 28.48, Improvement: N/A, Idx: N/A, Time (s): 0s\n" ] }, { "name": "stderr", "output_type": "stream", "text": [ "\r", " 0%| | 0/26 [00:00