Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,47 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
-*tfevents* filter=lfs diff=lfs merge=lfs -text
+*tfevents* filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,6 @@
+data
+output
+*.pyc
+*.png
+*.log
+.tar

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2024 varun-shanker
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,316 @@
+---
+license: mit
+language:
+- en
+- zh
+tags:
+- OneScience
+- life-science
+- protein
+- antibody
+- mutation
+- inverse-folding
+- structural-evolution
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">Structural Evolution</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+Structural Evolution is an unsupervised protein and antibody mutation recommendation workflow based on a structure-informed protein language model. It takes a PDB/CIF structure of a protein or protein complex as input, scores candidate mutation sequences under structural conditions with ESM-IF1, and selects high-probability substitutions from deep mutational scanning candidates to assist protein and antibody sequence optimization.
+
+Paper:
+
+> **Unsupervised evolution of protein and antibody complexes with a structure-informed language model**  
+> https://doi.org/10.1126/science.adk8946
+
+# Model Description
+
+Structural Evolution uses the structure-conditioned language model ESM-IF1 to recommend mutations for a target chain in a given protein structure. The program first extracts the target chain's wild-type sequence from the input PDB/CIF and generates single-point deep mutational scanning candidates. It then calculates the log-likelihood of each candidate under a single-chain or multichain backbone condition and selects recommended mutations in descending score order.
+
+# Use Cases
+
+| Use case | Description |
+| --- | --- |
+| Single-point protein mutation recommendation | Select high-probability amino acid substitutions from structure information |
+| Antibody sequence optimization | Recommend structure-conditioned mutations for antibody heavy and light chains separately |
+| Protein complex optimization | Evaluate target-chain mutations in a multichain backbone context |
+| Deep mutational scanning candidate screening | Score all single-point mutations and output high-scoring candidates |
+
+# Usage
+
+## 1. Using OneCode
+
+Experience intelligent one-click AI4S programming in the OneCode online environment:
+
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+**Hardware Requirements**
+
+- Structural Evolution supports inference on CPUs and GPUs.
+- A GPU/DCU is recommended for ESM-IF1 inference to accelerate candidate sequence scoring.
+- Larger protein complexes or larger mutation candidate libraries require more GPU/DCU memory and host memory.
+- If no GPU/DCU is available, you can explicitly use the CPU with `--nogpu`, but execution will be significantly slower.
+
+### Set Up the Runtime Environment
+
+#### DCU Environment
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+
+# Install with uv support
+pip install onescience[bio] \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+#### Environment Notes
+
+- If you encounter missing dependencies or version incompatibilities during execution, refer to the dependency versions specified in `environment.yml` and install or adjust the relevant dependencies as needed.
+
+### Prepare the Weights
+
+Complete Structural Evolution inference only requires the additional **ESM-IF1 model weights**. Standard mutation recommendation does not require downloading any additional training dataset from the paper.
+
+#### 1) ESM-IF1 Model Weights
+
+Download the ESM-IF1 weights:
+
+```text
+https://zenodo.org/records/12631662
+```
+
+Command:
+
+```bash
+wget -P ~/.cache/torch/hub/checkpoints \
+  https://zenodo.org/records/12631662/files/esm_if1_20220410.zip
+
+unzip ~/.cache/torch/hub/checkpoints/esm_if1_20220410.zip \
+  -d ~/.cache/torch/hub/checkpoints/
+```
+
+After extraction, make sure that the following file exists:
+
+```text
+~/.cache/torch/hub/checkpoints/
+└── esm_if1_20220410.pt
+```
+
+The file must exist.
+
+`scripts/recommend.py` always loads the model from:
+
+```text
+~/.cache/torch/hub/checkpoints/esm_if1_20220410.pt
+```
+
+loads the model from this path. If you save the weights elsewhere, update the checkpoint path in the code.
+
+## 3. Quick Start
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/structural-evolution \
+  --local-dir ./structural-evolution
+
+cd structural-evolution
+```
+
+- Structural Evolution additionally depends on the **ESM-IF1 model weights**. Prepare `esm_if1_20220410.pt` as described in "Prepare the Weights" and make sure that the checkpoint path in the code matches its actual location.
+
+### Quick Verification
+
+View the inference options:
+
+```bash
+python scripts/recommend.py --help
+```
+
+Run a quick verification with the official example structure:
+
+```bash
+python scripts/recommend.py \
+  scripts/examples/7mmo_abc_fvar.pdb \
+  --chain A \
+  --n 10
+```
+
+# Example Data
+
+The official repository provides an example structure in:
+
+```text
+scripts/examples/
+```
+
+The directory contains:
+
+```text
+scripts/examples/
+└─�� 7mmo_abc_fvar.pdb
+```
+
+This structure contains the LYCoV-1404 antibody variable region and SARS-CoV-2 RBD, and demonstrates mutation recommendation for an antibody heavy chain.
+
+For your own tasks, prepare:
+
+```text
+PDB or CIF structure of a protein or protein complex
++
+Target chain ID
+```
+
+For example:
+
+```text
+structure.pdb
+chain A
+```
+
+# Inference Examples
+
+## Basic Mutation Recommendation
+
+The simplest way to run it is:
+
+```bash
+python scripts/recommend.py \
+  /path/to/structure.pdb \
+  --chain A
+```
+
+Default output:
+
+```text
+Top 10 mutations
+maxrep = 1
+multichain backbone = True
+```
+
+This recommends 10 mutations by default and allows each original residue position to appear at most once.
+
+## Antibody Example
+
+Official example:
+
+```bash
+python scripts/recommend.py \
+  scripts/examples/7mmo_abc_fvar.pdb \
+  --chain A \
+  --seqpath scripts/examples/7mmo_chainA_lib.fasta \
+  --outpath scripts/examples/7mmo_chainA_scores.csv \
+  --upperbound 109 \
+  --offset 1
+```
+
+Where:
+
+| Parameter | Description |
+| --- | --- |
+| `--chain` | Target chain ID |
+| `--seqpath` | Output path for the generated deep mutational scanning FASTA |
+| `--outpath` | Output path for the CSV containing scores for all mutation candidates |
+| `--n` | Number of mutations to recommend; default: 10 |
+| `--maxrep` | Maximum number of times the same position may appear in the recommendations; default: 1 |
+| `--upperbound` | During final recommendation, consider only positions below this residue number |
+| `--offset` | Offset correction for PDB residue numbering |
+| `--order` | Specify chain order in a multichain structure |
+| `--multichain-backbone` | Use all chains as the structural context |
+| `--singlechain-backbone` | Use only the target chain backbone |
+| `--nogpu` | Force CPU execution |
+
+In the antibody example above, `--chain A` specifies the heavy chain. `--upperbound 109` excludes mutations in the final framework region. Because the input structure lacks the first residue, `--offset 1` corrects the mutation numbering.
+
+## Customize the Number of Recommendations
+
+For example, output the top 20 candidates and allow the same position to appear at most twice:
+
+```bash
+python scripts/recommend.py \
+  /path/to/your_structure.pdb \
+  --chain A \
+  --n 20 \
+  --maxrep 2
+```
+
+## Single-Chain Backbone Condition
+
+To use only the backbone information of the target chain:
+
+```bash
+python scripts/recommend.py \
+  /path/to/your_structure.pdb \
+  --chain A \
+  --singlechain-backbone
+```
+
+## CPU Inference
+
+```bash
+python scripts/recommend.py \
+  /path/to/your_structure.pdb \
+  --chain A \
+  --nogpu
+```
+
+# Output Description
+
+Structural Evolution first generates a complete single-point deep mutational scanning sequence library for the target chain in the input structure, then uses ESM-IF1 to calculate the structure-conditioned log-likelihood of each candidate sequence.
+
+If the following options are not specified manually:
+
+```text
+--seqpath
+--outpath
+```
+
+the program automatically saves the results under `output/`.
+
+The main outputs include:
+
+```text
+*.fasta
+*.csv
+```
+
+The files are:
+
+| Output | Description |
+| --- | --- |
+| DMS FASTA | Contains the wild-type sequence and all single-point mutation candidates |
+| scores CSV | Stores candidate sequences and their corresponding `log_likelihood` scores |
+| Terminal output | Top-N recommended mutations sorted by `log_likelihood` |
+
+Candidate sequences in the CSV are sorted by:
+
+```text
+log_likelihood
+```
+
+in descending order, and recommended mutations are finally selected using parameters such as `n`, `maxrep`, and `upperbound`.
+
+# Official OneScience Information
+
+| Platform | Main OneScience repository | Skills repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation and License
+
+- Original Structural Evolution paper: [Unsupervised evolution of protein and antibody complexes with a structure-informed language model](https://doi.org/10.1126/science.adk8946).
+- Structural Evolution is released under the **MIT License**, which permits use, copying, modification, publication, distribution, sublicensing, and commercial use provided that the copyright notice and license text are retained. See `LICENSE` in the repository root.
+- Structural Evolution inference depends on ESM-IF1 and related ESM code. The MIT License of the Structural Evolution repository does not automatically cover third-party model weights or dependency resources. For commercial use, redistribution, or other purposes, also review the applicable ESM/ESM-IF1 licenses and model-weight terms of use.
+- If you use this repository in research, cite the original Structural Evolution paper and, as appropriate, the related work on ESM-IF1 and other components actually used.

config.json

@@ -0,0 +1 @@
+{}

environment.yml

@@ -0,0 +1,281 @@
+name: struct-evo
+channels:
+  - pytorch
+  - pyg
+  - conda-forge
+  - defaults
+dependencies:
+  - _libgcc_mutex=0.1=main
+  - _openmp_mutex=5.1=1_gnu
+  - abseil-cpp=20211102.0=hd4dd3e8_0
+  - anyio=4.2.0=py39h06a4308_0
+  - argon2-cffi=21.3.0=pyhd3eb1b0_0
+  - argon2-cffi-bindings=21.2.0=py39h7f8727e_0
+  - arrow-cpp=14.0.2=h374c478_1
+  - asttokens=2.0.5=pyhd3eb1b0_0
+  - attrs=23.1.0=py39h06a4308_0
+  - aws-c-auth=0.6.19=h5eee18b_0
+  - aws-c-cal=0.5.20=hdbd6064_0
+  - aws-c-common=0.8.5=h5eee18b_0
+  - aws-c-compression=0.2.16=h5eee18b_0
+  - aws-c-event-stream=0.2.15=h6a678d5_0
+  - aws-c-http=0.6.25=h5eee18b_0
+  - aws-c-io=0.13.10=h5eee18b_0
+  - aws-c-mqtt=0.7.13=h5eee18b_0
+  - aws-c-s3=0.1.51=hdbd6064_0
+  - aws-c-sdkutils=0.1.6=h5eee18b_0
+  - aws-checksums=0.1.13=h5eee18b_0
+  - aws-crt-cpp=0.18.16=h6a678d5_0
+  - aws-sdk-cpp=1.10.55=h721c034_0
+  - backcall=0.2.0=pyhd3eb1b0_0
+  - beautifulsoup4=4.12.3=py39h06a4308_0
+  - blas=1.0=mkl
+  - bleach=4.1.0=pyhd3eb1b0_0
+  - bokeh=2.4.3=py39h06a4308_0
+  - boost-cpp=1.82.0=hdb19cb5_2
+  - bottleneck=1.3.7=py39ha9d4c09_0
+  - brotli=1.0.9=h5eee18b_8
+  - brotli-bin=1.0.9=h5eee18b_8
+  - brotli-python=1.0.9=py39h5a03fae_7
+  - bzip2=1.0.8=h5eee18b_6
+  - c-ares=1.19.1=h5eee18b_0
+  - ca-certificates=2024.3.11=h06a4308_0
+  - certifi=2024.6.2=py39h06a4308_0
+  - cffi=1.16.0=py39h5eee18b_1
+  - charset-normalizer=3.3.2=pyhd8ed1ab_0
+  - click=8.1.7=py39h06a4308_0
+  - cloudpickle=2.2.1=py39h06a4308_0
+  - colorama=0.4.6=pyhd8ed1ab_0
+  - colorcet=3.0.0=py39h06a4308_0
+  - comm=0.2.1=py39h06a4308_0
+  - contourpy=1.2.0=py39hdb19cb5_0
+  - cudatoolkit=11.3.1=h2bc3f7f_2
+  - cycler=0.11.0=pyhd3eb1b0_0
+  - cyrus-sasl=2.1.28=h52b45da_1
+  - cytoolz=0.12.2=py39h5eee18b_0
+  - dask=2023.11.0=py39h06a4308_0
+  - dask-core=2023.11.0=py39h06a4308_0
+  - datashader=0.14.1=py39h06a4308_1
+  - datashape=0.5.4=py39h06a4308_1
+  - dbus=1.13.18=hb2f20db_0
+  - debugpy=1.6.7=py39h6a678d5_0
+  - decorator=5.1.1=pyhd3eb1b0_0
+  - defusedxml=0.7.1=pyhd3eb1b0_0
+  - distributed=2023.11.0=py39h06a4308_0
+  - entrypoints=0.4=py39h06a4308_0
+  - et_xmlfile=1.1.0=py39h06a4308_0
+  - exceptiongroup=1.2.0=py39h06a4308_0
+  - executing=0.8.3=pyhd3eb1b0_0
+  - expat=2.6.2=h6a678d5_0
+  - fontconfig=2.14.1=h4c34cd2_2
+  - fonttools=4.51.0=py39h5eee18b_0
+  - freetype=2.12.1=h4a9f257_0
+  - fsspec=2024.3.1=py39h06a4308_0
+  - gflags=2.2.2=h6a678d5_1
+  - glib=2.78.4=h6a678d5_0
+  - glib-tools=2.78.4=h6a678d5_0
+  - glog=0.5.0=h6a678d5_1
+  - grpc-cpp=1.48.2=he1ff14a_1
+  - gst-plugins-base=1.14.1=h6a678d5_1
+  - gstreamer=1.14.1=h5eee18b_1
+  - heapdict=1.0.1=pyhd3eb1b0_0
+  - holoviews=1.15.0=py39h06a4308_1
+  - hvplot=0.8.0=py39h06a4308_0
+  - icu=73.1=h6a678d5_0
+  - idna=3.7=pyhd8ed1ab_0
+  - importlib-metadata=7.0.1=py39h06a4308_0
+  - importlib_resources=6.1.1=py39h06a4308_1
+  - intel-openmp=2023.1.0=hdb19cb5_46306
+  - ipykernel=6.28.0=py39h06a4308_0
+  - ipython=8.15.0=py39h06a4308_0
+  - ipython_genutils=0.2.0=pyhd3eb1b0_1
+  - jedi=0.18.1=py39h06a4308_1
+  - jinja2=2.11.3=pyhd8ed1ab_2
+  - joblib=1.4.2=pyhd8ed1ab_0
+  - jpeg=9e=h5eee18b_1
+  - jsonschema=4.19.2=py39h06a4308_0
+  - jsonschema-specifications=2023.7.1=py39h06a4308_0
+  - jupyter_client=7.4.9=py39h06a4308_0
+  - jupyter_core=5.7.2=py39h06a4308_0
+  - jupyter_events=0.10.0=py39h06a4308_0
+  - jupyter_server=2.10.0=py39h06a4308_0
+  - jupyter_server_terminals=0.4.4=py39h06a4308_1
+  - jupyterlab_pygments=0.2.2=py39h06a4308_0
+  - kiwisolver=1.4.4=py39h6a678d5_0
+  - krb5=1.20.1=h143b758_1
+  - lcms2=2.12=h3be6417_0
+  - ld_impl_linux-64=2.38=h1181459_1
+  - lerc=3.0=h295c915_0
+  - libboost=1.82.0=h109eef0_2
+  - libbrotlicommon=1.0.9=h5eee18b_8
+  - libbrotlidec=1.0.9=h5eee18b_8
+  - libbrotlienc=1.0.9=h5eee18b_8
+  - libclang=14.0.6=default_hc6dbbc7_1
+  - libclang13=14.0.6=default_he11475f_1
+  - libcups=2.4.2=h2d74bed_1
+  - libcurl=8.7.1=h251f7ec_0
+  - libdeflate=1.17=h5eee18b_1
+  - libedit=3.1.20230828=h5eee18b_0
+  - libev=4.33=h7f8727e_1
+  - libevent=2.1.12=hdbd6064_1
+  - libffi=3.4.4=h6a678d5_1
+  - libgcc-ng=11.2.0=h1234567_1
+  - libgfortran-ng=11.2.0=h00389a5_1
+  - libgfortran5=11.2.0=h1234567_1
+  - libglib=2.78.4=hdc74915_0
+  - libgomp=11.2.0=h1234567_1
+  - libiconv=1.16=h5eee18b_3
+  - libllvm14=14.0.6=hdb19cb5_3
+  - libnghttp2=1.57.0=h2d74bed_0
+  - libpng=1.6.39=h5eee18b_0
+  - libpq=12.17=hdbd6064_0
+  - libprotobuf=3.20.3=he621ea3_0
+  - libsodium=1.0.18=h7b6447c_0
+  - libssh2=1.11.0=h251f7ec_0
+  - libstdcxx-ng=11.2.0=h1234567_1
+  - libthrift=0.15.0=h1795dd8_2
+  - libtiff=4.5.1=h6a678d5_0
+  - libuuid=1.41.5=h5eee18b_0
+  - libuv=1.44.2=h5eee18b_0
+  - libwebp-base=1.3.2=h5eee18b_0
+  - libxcb=1.15=h7f8727e_0
+  - libxkbcommon=1.0.1=h5eee18b_1
+  - libxml2=2.10.4=hfdd30dd_2
+  - llvmlite=0.42.0=py39h6a678d5_0
+  - locket=1.0.0=py39h06a4308_0
+  - lz4=4.3.2=py39h5eee18b_0
+  - lz4-c=1.9.4=h6a678d5_1
+  - markdown=3.4.1=py39h06a4308_0
+  - markupsafe=1.1.1=py39h3811e60_3
+  - matplotlib=3.8.0=py39h06a4308_0
+  - matplotlib-base=3.8.0=py39h1128e8f_0
+  - matplotlib-inline=0.1.6=py39h06a4308_0
+  - mistune=0.8.4=py39h27cfd23_1000
+  - mkl=2023.1.0=h213fc3f_46344
+  - mkl-service=2.4.0=py39h5eee18b_1
+  - mkl_fft=1.3.8=py39h5eee18b_0
+  - mkl_random=1.2.4=py39hdb19cb5_0
+  - msgpack-python=1.0.3=py39hd09550d_0
+  - multipledispatch=0.6.0=py39h06a4308_0
+  - mysql=5.7.24=h721c034_2
+  - nbclassic=1.1.0=py39h06a4308_0
+  - nbclient=0.5.13=py39h06a4308_0
+  - nbconvert=6.4.4=py39h06a4308_0
+  - nbformat=5.9.2=py39h06a4308_0
+  - ncurses=6.4=h6a678d5_0
+  - nest-asyncio=1.6.0=py39h06a4308_0
+  - notebook=6.5.7=py39h06a4308_0
+  - notebook-shim=0.2.3=py39h06a4308_0
+  - numba=0.59.1=py39h6a678d5_0
+  - numexpr=2.8.7=py39h85018f9_0
+  - numpy=1.22.3=py39hf6e8229_2
+  - numpy-base=1.22.3=py39h060ed82_2
+  - openjpeg=2.4.0=h9ca470c_1
+  - openpyxl=3.0.10=py39h5eee18b_0
+  - openssl=3.0.14=h5eee18b_0
+  - orc=1.7.4=hb3bc3d3_1
+  - overrides=7.4.0=py39h06a4308_0
+  - packaging=23.2=py39h06a4308_0
+  - pandocfilters=1.5.0=pyhd3eb1b0_0
+  - panel=0.13.1=py39h06a4308_0
+  - param=1.12.0=pyhd3eb1b0_0
+  - parso=0.8.3=pyhd3eb1b0_0
+  - partd=1.4.1=py39h06a4308_0
+  - pcre2=10.42=hebb0a14_1
+  - pexpect=4.8.0=pyhd3eb1b0_3
+  - pickleshare=0.7.5=pyhd3eb1b0_1003
+  - pillow=10.3.0=py39h5eee18b_0
+  - pip=24.0=py39h06a4308_0
+  - platformdirs=3.10.0=py39h06a4308_0
+  - ply=3.11=py39h06a4308_0
+  - prometheus_client=0.14.1=py39h06a4308_0
+  - prompt-toolkit=3.0.43=py39h06a4308_0
+  - psutil=5.9.0=py39h5eee18b_0
+  - ptyprocess=0.7.0=pyhd3eb1b0_2
+  - pure_eval=0.2.2=pyhd3eb1b0_0
+  - pyarrow=14.0.2=py39h1eedbd7_0
+  - pybind11-abi=4=hd8ed1ab_3
+  - pycparser=2.21=pyhd3eb1b0_0
+  - pyct=0.5.0=py39h06a4308_0
+  - pyg=2.1.0=py39_torch_1.11.0_cu113
+  - pygments=2.15.1=py39h06a4308_1
+  - pyparsing=3.0.9=py39h06a4308_0
+  - pyqt=5.15.10=py39h6a678d5_0
+  - pyqt5-sip=12.13.0=py39h5eee18b_0
+  - pysocks=1.7.1=pyha2e5f31_6
+  - python=3.9.19=h955ad1f_1
+  - python-dateutil=2.9.0post0=py39h06a4308_2
+  - python-fastjsonschema=2.16.2=py39h06a4308_0
+  - python-json-logger=2.0.7=py39h06a4308_0
+  - python-lmdb=1.4.1=py39h6a678d5_0
+  - python_abi=3.9=2_cp39
+  - pytorch=1.11.0=py3.9_cuda11.3_cudnn8.2.0_0
+  - pytorch-cluster=1.6.0=py39_torch_1.11.0_cu113
+  - pytorch-mutex=1.0=cuda
+  - pytorch-scatter=2.0.9=py39_torch_1.11.0_cu113
+  - pytorch-sparse=0.6.15=py39_torch_1.11.0_cu113
+  - pytz=2024.1=py39h06a4308_0
+  - pyviz_comms=3.0.2=py39h06a4308_0
+  - pyyaml=6.0.1=py39h5eee18b_0
+  - pyzmq=24.0.1=py39h5eee18b_0
+  - qt-main=5.15.2=h53bd1ea_10
+  - re2=2022.04.01=h295c915_0
+  - readline=8.2=h5eee18b_0
+  - referencing=0.30.2=py39h06a4308_0
+  - requests=2.32.3=pyhd8ed1ab_0
+  - rfc3339-validator=0.1.4=py39h06a4308_0
+  - rfc3986-validator=0.1.1=py39h06a4308_0
+  - rpds-py=0.10.6=py39hb02cf49_0
+  - s2n=1.3.27=hdbd6064_0
+  - scikit-learn=1.4.2=py39h1128e8f_1
+  - scipy=1.11.4=py39h5f9d8c6_0
+  - send2trash=1.8.2=py39h06a4308_0
+  - setuptools=69.5.1=py39h06a4308_0
+  - sip=6.7.12=py39h6a678d5_0
+  - six=1.16.0=pyhd3eb1b0_1
+  - snappy=1.1.10=h6a678d5_1
+  - sniffio=1.3.0=py39h06a4308_0
+  - sortedcontainers=2.4.0=pyhd3eb1b0_0
+  - soupsieve=2.5=py39h06a4308_0
+  - sqlite=3.45.3=h5eee18b_0
+  - stack_data=0.2.0=pyhd3eb1b0_0
+  - tbb=2021.8.0=hdb19cb5_0
+  - tblib=1.7.0=pyhd3eb1b0_0
+  - terminado=0.17.1=py39h06a4308_0
+  - testpath=0.6.0=py39h06a4308_0
+  - threadpoolctl=3.5.0=pyhc1e730c_0
+  - tk=8.6.14=h39e8969_0
+  - tomli=2.0.1=py39h06a4308_0
+  - toolz=0.12.0=py39h06a4308_0
+  - tornado=6.4.1=py39h5eee18b_0
+  - tqdm=4.66.4=pyhd8ed1ab_0
+  - traitlets=5.14.3=py39h06a4308_0
+  - typing-extensions=4.11.0=py39h06a4308_0
+  - typing_extensions=4.11.0=py39h06a4308_0
+  - unicodedata2=15.1.0=py39h5eee18b_0
+  - urllib3=2.2.2=pyhd8ed1ab_0
+  - utf8proc=2.6.1=h5eee18b_1
+  - wcwidth=0.2.5=pyhd3eb1b0_0
+  - webencodings=0.5.1=py39h06a4308_1
+  - websocket-client=1.8.0=py39h06a4308_0
+  - wheel=0.43.0=py39h06a4308_0
+  - xarray=2023.6.0=py39h06a4308_0
+  - xz=5.4.6=h5eee18b_1
+  - yaml=0.2.5=h7b6447c_0
+  - zeromq=4.3.5=h6a678d5_0
+  - zict=3.0.0=py39h06a4308_0
+  - zipp=3.17.0=py39h06a4308_0
+  - zlib=1.2.13=h5eee18b_1
+  - zstd=1.5.5=hc292b87_2
+  - pip:
+    - ablang==0.3.1
+    - biopython==1.78
+    - biotite==0.34.1
+    - git+https://github.com/facebookresearch/esm.git
+    - iqplot==0.3.3
+    - msgpack==1.0.8
+    - natsort==8.4.0
+    - networkx==3.2.1
+    - pandas==2.0.3
+    - tzdata==2024.1
+prefix: /data/home/vrs/miniconda3/envs/struct-evo

model/dms_enrichment.py

@@ -0,0 +1,382 @@
+import argparse
+from pathlib import Path
+from biotite.sequence.io.fasta import FastaFile, get_sequences
+from Bio import pairwise2
+import pandas as pd
+import shutil
+import os
+import numpy as np
+import pandas as pd
+import bokeh.io
+import bokeh.plotting
+import bokeh.palettes 
+from bokeh.transform import factor_cmap
+
+from bokeh.io import export_svg
+import iqplot 
+from colorcet import glasbey_category10
+
+import subprocess
+
+bla_dict = {
+    'protein':'bla', 
+    'pdbfile':'bla_1m40_a.pdb', 
+    'chain':'A', 
+    'dmsfile': 'dms_bla.csv', 
+    'fitness_col': 'DMS_amp_2500_(b)', 
+    'threshold' : 0.01
+}
+
+calm1_dict = {
+    'protein':'CALM1', 
+    'pdbfile':'calm1_5v03_r.pdb', 
+    'chain':'R', 
+    'dmsfile': 'dms_calm1.csv', 
+    'fitness_col':'DMS', 
+    'threshold': 1
+}
+
+haeiiim_dict = {
+    'protein':'haeIIIM', 
+    'pdbfile':'haeiiim_3ubt_b.pdb', 
+    'chain':'B', 
+    'dmsfile': 'dms_haeiiim.csv', 
+    'fitness_col':'DMS_G3', 
+    'threshold' : 0.01
+}
+
+gal4_dict = {
+    'protein':'GAL4', 
+    'pdbfile':'gal4_3coq_b.pdb', 
+    'chain':'B', 
+    'dmsfile': 'dms_gal4.csv', 
+    'fitness_col':'DMS_nonsel_24', 
+    'threshold': 1
+}
+
+hras_dict = {
+    'protein':'HRAS', 
+    'pdbfile':'hras_2ce2_x.pdb', 
+    'chain':'X', 
+    'dmsfile': 'dms_hras.csv', 
+    'fitness_col':'DMS_unregulated', 
+    'threshold': 1
+}
+
+mapk1_dict = {
+    'protein':'MAPK1', 
+    'pdbfile':'mapk1_4zzn_a.pdb', 
+    'chain':'A', 
+    'dmsfile': 'dms_mapk1.csv', 
+    'fitness_col':'DMS_VRT', 
+    'threshold': 1
+}
+
+tpk1_dict = {
+    'protein':'TPK1', 
+    'pdbfile':'tpk1_3s4y_a.pdb', 
+    'chain':'A', 
+    'dmsfile': 'dms_tpk1.csv', 
+    'fitness_col':'DMS', 
+    'threshold': 1
+}
+
+tpmt_dict = {
+    'protein':'TPMT', 
+    'pdbfile':'tpmt_2bzg_a.pdb', 
+    'chain':'A', 
+    'dmsfile': 'dms_tpmt.csv', 
+    'fitness_col':'DMS', 
+    'threshold': 1
+}
+
+ube2i_dict = {
+    'protein':'UBE2I', 
+    'pdbfile':'ube2i_5f6e_a.pdb', 
+    'chain':'A', 
+    'dmsfile': 'dms_tpmt.csv', 
+    'fitness_col':'DMS', 
+    'threshold': 1
+}
+
+ubi4_dict = {
+    'protein':'UBI4', 
+    'pdbfile':'ubi4_4q5e_b.pdb', 
+    'chain':'B', 
+    'dmsfile': 'dms_ubi4.csv', 
+    'fitness_col':'DMS_limiting_(b)', 
+    'threshold': 1
+}
+
+studies = [           
+        bla_dict,
+        calm1_dict,
+        gal4_dict, 
+        haeiiim_dict, 
+        hras_dict,
+        tpmt_dict,
+        tpk1_dict,
+        mapk1_dict,
+        ube2i_dict,
+        ubi4_dict,
+ ]
+
+def score_if(prot_dict):
+    prefix = 'data/dms/'
+    fasta_outpath = f"dms_libs/{prot_dict['pdbfile'][:-4]}_dmsLib.fasta"
+    score_if_cmd = [
+        'python', 
+        'model/score_log_likelihoods.py',
+        prefix + f"structures/{prot_dict['pdbfile']}", 
+        '--seqpath', prefix + fasta_outpath, 
+        '--outpath', f"output/dms/if_results/dms_{prot_dict['protein']}_if.csv"   
+        ,'--chain', prot_dict['chain']
+    ]
+    subprocess.run(score_if_cmd, check=True) 
+
+def get_stats(prot_dict, n, top_per):
+
+    top_per = top_per/100
+    prefix = 'output/dms/'
+    esm1v_outpath = 'dms_esm1vScored/' + f"dms_{prot_dict['protein']}" + '_maskMarginals.csv'
+    if_results_df = pd.read_csv(  prefix + f"if_results/dms_{prot_dict['protein']}_if.csv")
+    dms_df = pd.read_csv(prefix + esm1v_outpath)
+    
+    unprofiled_df = dms_df[dms_df[prot_dict['fitness_col']].isna()]
+    unprofiled_variants = unprofiled_df['variant'].to_list()
+    
+    dms_df = dms_df.dropna(subset = [prot_dict['fitness_col']])
+    pop_n = len(dms_df)
+    if_results_df = if_results_df[~if_results_df['seqid'].isin(unprofiled_variants)]
+
+    assert (len(if_results_df) == len(dms_df))
+
+    threshold = dms_df[prot_dict['fitness_col']].quantile(1-top_per)
+    highFit_subset_variants = dms_df[dms_df[prot_dict['fitness_col']] >= threshold]['variant']
+    n_highFit = len(highFit_subset_variants)
+
+
+    dms_df['esm1v'] = dms_df.loc[:, dms_df.columns.str.startswith('esm1v')].mean(axis=1)
+    top_esm1v_df = dms_df.sort_values(by = 'esm1v', ascending = False)[:n]
+    top_esm1v_variants = top_esm1v_df['variant'].to_list()
+
+    top_if_df = if_results_df.sort_values(by = 'log_likelihood', ascending = False)[:n]
+    top_if_variants = top_if_df['seqid'].to_list()
+    
+    esm1v_hits_df = top_esm1v_df[top_esm1v_df['variant'].isin(highFit_subset_variants)]
+    if_hits_df = top_if_df[top_if_df['seqid'].isin(highFit_subset_variants)]
+    
+    n_esm1v_hits = len(esm1v_hits_df)
+    n_if_hits = len(if_hits_df)
+
+    if_hit_enrich = (n_if_hits / n) / top_per
+    esm1v_hit_enrich = (n_esm1v_hits / n) / top_per
+
+    if_hit_rate = (n_if_hits / n)
+    esm1v_hit_rate = (n_esm1v_hits / n)
+
+    return (prot_dict['protein'], pop_n, n_highFit, n_if_hits, n_esm1v_hits, if_hit_enrich, esm1v_hit_enrich, if_hit_rate, esm1v_hit_rate )
+
+def plot_bars(melted_df):
+    p = bokeh.plotting.figure(
+    height=350,
+    width=1100,
+    y_axis_label="High Fitness \n Prediction Precision",
+    x_axis_label = "Functional Percentile Threshold \n for High Fitness Classification",
+    x_range=bokeh.models.FactorRange(*factors, group_padding = 1.2),
+    tools="save",
+    title=''
+    )
+    p.output_backend = "svg"
+
+    p.vbar(
+        source=melted_df,
+        x="cats",
+        top="Hit Rate",
+        width = 1,
+        line_color='black',
+        alpha = 'alpha',
+        legend_field='legend_label',
+        fill_color=bokeh.transform.factor_cmap(
+            'Method',
+            palette=['#999999', '#43a2ca'],
+            factors=list(melted_df['Method'].unique()),
+            start=1,
+            end=2
+        )
+    )
+
+    p.xgrid.grid_line_color = None
+    p.ygrid.grid_line_color = None
+    p.x_range.range_padding = 0.03
+    p.legend.location = "top_right"
+    p.legend.spacing = 20
+    p.legend.label_text_font_size = "11pt"
+    p.legend.label_text_color = "black"
+    p.legend.orientation = "horizontal"
+    p.xaxis.major_label_orientation = 1.2
+    p.xaxis.separator_line_alpha = 0
+    p.xaxis.group_text_font_size = '11pt'
+    p.xaxis.subgroup_text_font_size = '0pt'
+    p.xaxis.axis_label_text_font_size = '11pt'
+    p.xaxis.major_label_text_font_size = '10pt'
+    p.xaxis.axis_label_text_font_style = 'normal'
+    p.xaxis.major_label_text_color = 'black'
+    p.xaxis.axis_label_text_color = 'black'
+
+    p.yaxis.axis_label_text_font_size = '11pt'
+    p.yaxis.axis_label_text_font_style = 'normal'
+    p.yaxis.axis_label_text_color = 'black'
+    p.yaxis.major_label_text_font_size = '10pt'
+    p.xaxis.group_text_color = 'dimgrey'
+
+    return p
+
+def plot_ecdf_hits(prot_dict, n, top_per_lst):
+
+    #use least stringent threshold to plot all hits
+    top_per = max(top_per_lst)/100
+    prefix = 'output/dms/'
+
+    esm1v_outpath = 'dms_esm1vScored/' + f"dms_{prot_dict['protein']}" + '_maskMarginals.csv'
+    if_results_df = pd.read_csv(  prefix + f"if_results/dms_{prot_dict['protein']}_if.csv")
+    dms_df = pd.read_csv(prefix + esm1v_outpath)
+
+    
+    unprofiled_df = dms_df[dms_df[prot_dict['fitness_col']].isna()]
+    unprofiled_variants = unprofiled_df['variant'].to_list()
+    
+    dms_df = dms_df.dropna(subset = [prot_dict['fitness_col']])
+    dms_df['percentile'] = dms_df[prot_dict['fitness_col']].rank(pct=True, method='average')
+    dms_df['zscore'] = (dms_df[prot_dict['fitness_col']] - dms_df[prot_dict['fitness_col']].mean())/dms_df[prot_dict['fitness_col']].std()
+    if_results_df = if_results_df[~if_results_df['seqid'].isin(unprofiled_variants)]
+    
+    assert (len(if_results_df) == len(dms_df))
+
+    threshold = dms_df[prot_dict['fitness_col']].quantile(1-top_per)
+    highFit_subset_variants = dms_df[dms_df[prot_dict['fitness_col']] >= threshold]['variant']
+    n_highFit = len(highFit_subset_variants)
+
+    dms_df['esm1v'] = dms_df.loc[:, dms_df.columns.str.startswith('esm1v')].mean(axis=1)
+    top_esm1v_df = dms_df.sort_values(by = 'esm1v', ascending = False)[:n]
+    top_esm1v_variants = top_esm1v_df['variant'].to_list()
+
+    top_if_df = if_results_df.sort_values(by = 'log_likelihood', ascending = False)[:n]
+    top_if_variants = top_if_df['seqid'].to_list()
+    
+    esm1v_hits_df = top_esm1v_df[top_esm1v_df['variant'].isin(highFit_subset_variants)]
+    if_hits_df = top_if_df[top_if_df['seqid'].isin(highFit_subset_variants)]
+    if_hits_df = pd.merge(if_hits_df, dms_df[['variant', 'zscore', 'percentile', prot_dict['fitness_col']]], left_on='seqid', right_on='variant', how='left')
+
+    p_1v = iqplot.ecdf(
+            data=dms_df,
+            q=prot_dict['fitness_col'],
+            style="staircase",
+            palette = 'lightgrey',
+            line_kwargs= {'line_width':6},
+            x_axis_label = prot_dict['protein']+' Fitness Measure',
+            #title = prot_dict['protein']
+    )
+    p_if = iqplot.ecdf(
+            data=dms_df,
+            q=prot_dict['fitness_col'],
+            style="staircase",
+            palette = 'lightgrey',
+            line_kwargs= {'line_width':6},
+            x_axis_label = None, #prot_dict['protein']+'Fitness Measure',
+            title = prot_dict['protein']
+    )
+    colors = ['#1f78b4', '#33a02c', '#ff7f00',]
+    for i, top_per in enumerate(top_per_lst):
+        top_per = top_per/100
+        hline = bokeh.models.Span(location=1-top_per, dimension='width', line_color=colors[i], line_dash='dashed', line_width=4, )
+        p_1v.add_layout(hline)
+        p_if.add_layout(hline)
+
+
+    p_1v.circle(
+            source = esm1v_hits_df,
+            x = prot_dict['fitness_col'],
+            y = 'percentile',
+            color = 'dimgrey',
+            size = 11,
+            line_color = 'lightgrey',
+            alpha = 0.8,
+    )
+ 
+    p_if.circle(
+            source = if_hits_df,
+            x = prot_dict['fitness_col'],
+            y = 'percentile',
+            color = '#43a2ca',
+            size = 11,
+            line_color = 'white',
+            alpha = 0.8,
+    )
+
+    p_if.title.text_font_style = 'normal'
+    p_if.title.text_color = 'black'
+    p_if.title.align = 'center'
+    p_if.title.text_font_size = '16pt'
+    p_1v.xaxis.axis_label_text_font_size = '14pt'
+    
+
+    for p in [p_if, p_1v]:
+        p.xaxis.axis_label_text_font_style = 'normal'
+        p.yaxis.axis_label_text_font_style = 'normal'
+        p.xaxis.axis_label_text_font_size = '14pt'
+        p.yaxis.axis_label_text_font_size = '14pt'
+        p.xaxis.axis_label_text_color = 'black'
+        p.yaxis.axis_label_text_color = 'black'
+        p.xaxis.major_label_text_font_size = '14pt'
+        p.yaxis.major_label_text_font_size = '14pt'
+        p.xaxis.major_label_text_color = 'black'
+        p.yaxis.major_label_text_color = 'black'
+        p.output_backend = "svg"
+
+    return [p_if, p_1v]
+
+if __name__ == '__main__':
+    for s in studies:
+        score_if(s)
+    shutil.copytree("data/dms/dms_esm1vScored", "output/dms/dms_esm1vScored", dirs_exist_ok=True)
+
+    name, n_pop, n_popHits, n_if_hits, n_esm1v_hits, if_hit_enrich, esm1v_hit_enrich, if_hit_rate, esm1v_hit_rate, p_threshold = ([] for _ in range(10))
+    data_lists = [name, n_pop, n_popHits, n_if_hits, n_esm1v_hits, if_hit_enrich, esm1v_hit_enrich, if_hit_rate, esm1v_hit_rate, p_threshold]
+    for s in studies:
+        for p in [5, 10, 20 ]:
+            outputs = get_stats(s, 10, p)
+            for i in range(len(outputs)):
+                data_lists[i].append(outputs[i])
+            data_lists[len(data_lists)-1].append(p)
+    col_names = ['Protein', 'Total Library Size', 'Library Hits (Variants with Fitness >95th)', 'Inverse Folding Hits (in Top 10)', 'ESM1v Hits (in Top 10)', 'Inverse Folding Hit Enrichment', 'ESM1v Hit Enrichment','Inverse Folding Hit Rate', 'ESM1v Hit Rate','Percentile Threshold' ]
+    results_df = pd.DataFrame({col_names[i]: data_lists[i] for i in range(len(data_lists))})
+    # Melt the dataframe
+    melted_df = results_df.melt(id_vars=['Protein', 'Percentile Threshold'], value_vars=[ 'ESM1v Hit Rate', 'Inverse Folding Hit Rate',],
+                        var_name='Method', value_name='Hit Rate')
+    melted_df['Method'] = melted_df['Method'].replace({
+        'Inverse Folding Hit Rate': 'Inverse Folding',
+        'ESM1v Hit Rate': 'Language Model'})
+    melted_df['alpha'] = (1/melted_df['Percentile Threshold'].to_numpy())*3 + 0.4
+    melted_df['Percentile Threshold'] = melted_df['Percentile Threshold'].astype(str)
+    melted_df['legend_label'] = ['Structure-Informed Language Model' if x=='Inverse Folding' else 'Language Model' for x in melted_df['Method']]
+
+    melted_df['cats'] = melted_df.apply(lambda x: (x["Protein"],  x["Method"], x['Percentile Threshold'],), axis = 1)
+    factors = list(melted_df.cats)
+
+    p = plot_bars(melted_df)
+    comparePrecision_fname = 'output/dms/precision_comparison.html'
+    bokeh.plotting.output_file(comparePrecision_fname)
+    bokeh.io.show(p)
+
+    ecdf_plots = []
+    for s in studies:
+        ecdf_plots.extend(plot_ecdf_hits(s, 10, [20,10,5]))
+    mid = int(len(ecdf_plots)/2)
+    ecdf_fname = 'output/dms/fitness_ecdfs.html'
+    bokeh.plotting.output_file(ecdf_fname)
+    bokeh.io.show(bokeh.layouts.gridplot([ecdf_plots[:mid:2],
+                                        ecdf_plots[1:mid:2],
+                                        ecdf_plots[mid::2],
+                                        ecdf_plots[mid+1::2]])
+    )

model/dms_utils.py

@@ -0,0 +1,73 @@
+from collections import Counter
+import datetime
+from dateutil.parser import parse as dparse
+import errno
+import numpy as np
+import os
+import random
+import sys
+import time
+import warnings
+from Bio import pairwise2
+from Bio import BiopythonWarning
+warnings.simplefilter('ignore', BiopythonWarning)
+from Bio import Seq, SeqIO
+
+
+np.random.seed(1)
+random.seed(1)
+
+AAs = [
+    'A', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'K', 'L',
+    'M', 'N', 'P', 'Q', 'R', 'S', 'T', 'V', 'W', 'Y',
+]
+
+def tprint(string):
+    string = str(string)
+    sys.stdout.write(str(datetime.datetime.now()) + ' | ')
+    sys.stdout.write(string + '\n')
+    sys.stdout.flush()
+
+def mkdir_p(path):
+    try:
+        os.makedirs(path)
+    except OSError as exc:  # Python >2.5
+        if exc.errno == errno.EEXIST and os.path.isdir(path):
+            pass
+        else:
+            raise
+        
+def deep_mutational_scan(sequence, exclude_noop=True):
+    for pos, wt in enumerate(sequence):
+        for mt in AAs:
+            if exclude_noop and wt == mt:
+                continue
+            yield (pos, wt, mt)
+
+
+def make_mutations(seq, mutations):
+    mut_seq = [ char for char in seq ]
+    for mutation in mutations:
+        wt, pos, mt = mutation[0], int(mutation[1:-1]) - 1, mutation[-1]
+        assert(seq[pos] == wt)
+        mut_seq[pos] = mt
+    mut_seq = ''.join(mut_seq).replace('-', '')
+    return mut_seq
+
+def find_mutations(seq1, seq2):
+    alignment = pairwise2.align.globalms(
+        seq1, seq2, 5, -4, -6, -.1, one_alignment_only=True,
+    )[0]
+
+    mutation_set = []
+    pos1, pos2, pos_map = 0, 0, {}
+    for wt, mt in zip(alignment[0], alignment[1]):
+        if wt != '-':
+            pos1 += 1
+        if mt != '-':
+            pos2 += 1
+        if wt != mt and wt != '-' and mt != '-':
+            mut_str = f'{wt}{pos1}{mt}'
+            mutation_set.append(mut_str)
+
+    return mutation_set

model/eval_ablang.py

@@ -0,0 +1,65 @@
+import ablang
+import numpy as np
+import pandas as pd
+import scipy.special
+
+heavy_ablang = ablang.pretrained("heavy")
+heavy_ablang.freeze()
+
+light_ablang = ablang.pretrained("light")
+light_ablang.freeze()
+
+ab_dict = dict(
+    cr6261= ('EVQLVESGAEVKKPGSSVKVSCKASGGPFRSYAISWVRQAPGQGPEWMGGIIPIFGTTKYAPKFQGRVTITADDFAGTVYMELSSLRSEDTAMYYCAKHMGYQVRETMDVWGKGTTVTVSS', 
+                'QSVLTQPPSVSAAPGQKVTISCSGSSSNIGNDYVSWYQQLPGTAPKLLIYDNNKRPSGIPDRFSGSKSGTSATLGITGLQTGDEANYYCATWDRRPTAYVVFGGGTKLTVL'), 
+    cr9114=  ('QVQLVQSGAEVKKPGSSVKVSCKSSGGTSNNYAISWVRQAPGQGLDWMGGISPIFGSTAYAQKFQGRVTISADIFSNTAYMELNSLTSEDTAVYFCARHGNYYYYSGMDVWGQGTTVTVSS',
+                'QSALTQPPAVSGTPGQRVTISCSGSDSNIGRRSVNWYQQFPGTAPKLLIYSNDQRPSVVPDRFSGSKSGTSASLAISGLQSEDEAEYYCAAWDDSLKGAVFGGGTQLTVL'),
+    g6 =  ('EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTV',
+                'DIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIK')
+
+)
+
+def eval_ablang(s, ab, chain):
+    if chain == 'hc':
+        fname = 'output/ab_mutagenesis_expts/'+ab+'/' + ab +'_hc_ablangScores.csv'
+        log_likelihoods = heavy_ablang(s, mode = 'likelihood')[0][1:-1]
+        alphabet = heavy_ablang.tokenizer.vocab_to_aa
+    elif chain == 'lc':
+        fname = 'output/ab_mutagenesis_expts/'+ab+'/' + ab +'_lc_ablangScores.csv'
+        log_likelihoods = light_ablang(s, mode = 'likelihood')[0][1:-1]
+        alphabet = light_ablang.tokenizer.vocab_to_aa
+
+    assert (log_likelihoods.shape)[0] == len(s)
+
+    filt_alphabet = {key: value for key, value in alphabet.items() if value.isalpha()}
+    log_likelihood_ratio = []
+    for i,res_log_likelihoods in enumerate(log_likelihoods):
+            wt_res = s[i]
+            wt_index = list(filt_alphabet.values()).index(wt_res)
+            wt_log_likelihood = res_log_likelihoods[wt_index]
+            log_likelihood_ratio.extend(res_log_likelihoods-wt_log_likelihood)
+
+    res_order = [alphabet[key] for key in range(1, 21)] #extract order of residues in likelihood
+    mt = res_order * len(s)
+    wt = [char for char in s for i in range(len(res_order))]
+    pos = [i+1 for i in range(len(s)) for j in range(len(res_order))]
+    data = {
+        'pos': pos,
+        'wt': wt,
+        'mt': mt,
+        'log_likelihood' : log_likelihoods.flatten(),
+        'log_likelihood_ratio' : log_likelihood_ratio,
+    } 
+    df = pd.DataFrame(data)
+    df.to_csv(fname, index = False)
+
+
+def main():
+    
+    for ab in ab_dict:
+        for s,chain in zip(ab_dict[ab], ('hc','lc')):
+              eval_ablang(s, ab, chain)
+           
+
+if __name__ == '__main__':
+        main()

model/generate_dms.py

@@ -0,0 +1,60 @@
+import argparse
+from dms_utils import deep_mutational_scan
+from pathlib import Path
+import numpy as np
+
+import esm
+from util import load_structure, extract_coords_from_structure
+import biotite.structure
+from collections import defaultdict
+
+
+def get_native_seq(pdbfile, chain):
+    structure = load_structure(pdbfile, chain)
+    _ , native_seq = extract_coords_from_structure(structure)
+    return native_seq
+    
+def write_dms_lib(args):
+    '''Writes a deep mutational scanning library, including the native/wildtype (wt) of the 
+    indicated target chain in the structure to an output Fasta file'''
+
+    sequence = get_native_seq(args.pdbfile, args.chain)
+    Path(args.outpath).parent.mkdir(parents=True, exist_ok=True)
+    with open(args.dmspath, 'w') as f:
+        f.write('>wt\n')
+        f.write(sequence+'\n')
+        for pos, wt, mt in deep_mutational_scan(sequence):
+            assert(sequence[pos] == wt)
+            mut_seq = sequence[:pos] + mt + sequence[(pos + 1):]
+
+            f.write('>' + str(wt) + str(pos+1) + str(mt) + '\n')
+            f.write(mut_seq + '\n')
+
+
+def main():
+    parser = argparse.ArgumentParser(
+            description='Create a DMS library based on target chain in the structure.'
+    )
+    parser.add_argument(
+            'pdbfile', type=str,
+            help='input filepath, either .pdb or .cif',
+    )
+    parser.add_argument(
+            '--dmspath', type=str,
+            help='output filepath for dms library',
+    )
+    parser.add_argument(
+            '--chain', type=str,
+            help='chain id for the chain of interest', default='A',
+    )
+
+    args = parser.parse_args()
+
+    if args.dmspath is None:
+        args.dmspath = f'predictions/{args.pdbfile[:-4]}-{args.chain}_dms.fasta'
+
+    write_dms_lib(args)
+
+if __name__ == '__main__':
+    
+    main()

model/multichain_util.py

@@ -0,0 +1,216 @@
+import biotite.structure
+import numpy as np
+import torch
+from typing import Sequence, Tuple, List
+from util import (
+    load_structure,
+    extract_coords_from_structure,
+    load_coords,
+    get_sequence_loss,
+    get_encoder_output,
+)
+import util
+
+
+def extract_coords_from_complex(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: biotite AtomArray
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    coords = {}
+    seqs = {}
+    all_chains = biotite.structure.get_chains(structure)
+    for chain_id in all_chains:
+        chain = structure[structure.chain_id == chain_id]
+        coords[chain_id], seqs[chain_id] = extract_coords_from_structure(chain)
+    return coords, seqs
+
+
+def load_complex_coords(fpath, chains):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chains: the chain ids (the order matters for autoregressive model)
+    Returns:
+        Tuple (coords_list, seq_list)
+        - coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+          coordinates representing the backbone of each chain
+        - seqs: Dictionary mapping chain ids to native sequences of each chain
+    """
+    structure = load_structure(fpath, chains)
+    return extract_coords_from_complex(structure)
+
+#*
+def _concatenate_coords(
+        coords,
+        target_chain_id,
+        padding_length=10,
+        order=None
+):
+    """
+    Args:
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: Length of padding between concatenated chains
+    Returns:
+        Tuple (coords, seq)
+            - coords_concatenated is an L x 3 x 3 array for N, CA, C coordinates, a
+              concatenation of the chains with padding in between 
+              AND target chain placed first
+            - seq is the extracted sequence, with padding tokens inserted
+            between the concatenated chains
+    """
+    pad_coords = np.full((padding_length, 3, 3), np.nan, dtype=np.float32)
+    if order is None:
+        order = (
+            [ target_chain_id ] +
+            [ chain_id for chain_id in coords if chain_id != target_chain_id ]
+        )
+    coords_list, coords_chains = [], []
+    for idx, chain_id in enumerate(order):
+        if idx > 0:
+            coords_list.append(pad_coords)
+            coords_chains.append([ 'pad' ] * padding_length)
+        coords_list.append(list(coords[chain_id]))
+        coords_chains.append([ chain_id ] * coords[chain_id].shape[0])
+    coords_concatenated = np.concatenate(coords_list, axis=0)
+    coords_chains = np.concatenate(coords_chains, axis=0).ravel()
+    return coords_concatenated, coords_chains
+
+#*
+def _concatenate_seqs(
+        native_seqs,
+        target_seq,
+        target_chain_id,
+        padding_length=10,
+        order=None,
+):
+    """
+    Args:
+        native_seqs: Dictionary mapping chain ids to corresponding AA sequence
+        target_seq: The chain id to sample sequences for
+        padding_length: Length of padding between concatenated chains
+    Returns:
+        native_seqs_concatenated: Array of length L, concatenation of the chain 
+        sequences with padding in between
+    """
+    if order is None:
+        order = (
+            [ target_chain_id ] +
+            [ chain_id for chain_id in native_seqs if chain_id != target_chain_id ]
+        )
+    native_seqs_list = []
+    for idx, chain_id in enumerate(order):
+        if idx > 0:
+            native_seqs_list.append(['<mask>'] * (padding_length - 1) + ['<cath>'])
+        if chain_id == target_chain_id:
+            native_seqs_list.append(list(target_seq))
+        else:
+            native_seqs_list.append(list(native_seqs[chain_id]))
+    native_seqs_concatenated = ''.join(np.concatenate(native_seqs_list, axis=0))
+    return native_seqs_concatenated
+
+
+#*
+def sample_sequence_in_complex(model, coords, target_chain_id, temperature=1.,
+        padding_length=10):
+    """
+    Samples sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+        padding_length: padding length in between chains
+    Returns:
+        Sampled sequence for the target chain
+    """
+    target_chain_len = coords[target_chain_id].shape[0]
+    all_coords, coords_chains = _concatenate_coords(coords, target_chain_id)
+    device = next(model.parameters()).device
+
+    # Supply padding tokens for other chains to avoid unused sampling for speed
+    padding_pattern = ['<pad>'] * all_coords.shape[0]
+    for i in range(target_chain_len):
+        padding_pattern[i] = '<mask>'
+    sampled = model.sample(all_coords, partial_seq=padding_pattern,
+            temperature=temperature, device=device)
+    sampled = sampled[:target_chain_len]
+    return sampled
+
+
+#*
+def score_sequence_in_complex(
+        model,
+        alphabet,
+        coords,
+        native_seqs,
+        target_chain_id,
+        target_seq,
+        padding_length=10,
+        order=None,
+):
+    """
+    Scores sequence for one chain in a complex.
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        native_seqs: Dictionary mapping chain ids to sequence
+                extracted from each chain
+        target_chain_id: The chain id to sample sequences for
+        target_seq: Target sequence for the target chain for scoring.
+        padding_length: padding length in between chains
+    Returns:
+        Tuple (ll_fullseq, ll_withcoord)
+        - ll_fullseq: Average log-likelihood over the full target chain
+        - ll_targetseq Average log-likelihood in target chain excluding those
+            residues without coordinates
+    """
+
+    assert(len(target_seq) == len(native_seqs[target_chain_id]))
+
+    all_coords, coords_chains = _concatenate_coords(
+        coords,
+        target_chain_id,
+        order=order,
+    )
+    all_seqs = _concatenate_seqs(
+        native_seqs,
+        target_seq,
+        target_chain_id,
+        order=order,
+    )
+
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, all_coords,
+            all_seqs)
+    assert(all_coords.shape[0] == coords_chains.shape[0] == loss.shape[0])
+
+    ll_fullseq = -np.mean(loss[coords_chains != 'pad'])
+    ll_targetseq = -np.mean(loss[coords_chains == target_chain_id])
+
+    return ll_fullseq, ll_targetseq
+
+
+def get_encoder_output_for_complex(model, alphabet, coords, target_chain_id):
+    """
+    Args:
+        model: An instance of the GVPTransformer model
+        alphabet: Alphabet for the model
+        coords: Dictionary mapping chain ids to L x 3 x 3 array for N, CA, C
+            coordinates representing the backbone of each chain
+        target_chain_id: The chain id to sample sequences for
+    Returns:
+        Dictionary mapping chain id to encoder output for each chain
+    """
+    all_coords = _concatenate_coords(coords, target_chain_id)
+    all_rep = get_encoder_output(model, alphabet, all_coords)
+    target_chain_len = coords[target_chain_id].shape[0]
+    return all_rep[:target_chain_len]

model/parse_abysis.py

@@ -0,0 +1,65 @@
+import json
+import pandas as pd
+import os
+
+AAs = set('ACDEFGHIKLMNPQRSTVWY')
+
+seqs_abs = {
+    'cr6261_vh': 'EVQLVESGAEVKKPGSSVKVSCKASGGPFRSYAISWVRQAPGQGPEWMGGIIPIFGTTKYAPKFQGRVTITADDFAGTVYMELSSLRSEDTAMYYCAKHMGYQVRETMDVWGKGTTVTVSS',
+    'cr9114_vh': 'QVQLVQSGAEVKKPGSSVKVSCKSSGGTSNNYAISWVRQAPGQGLDWMGGISPIFGSTAYAQKFQGRVTISADIFSNTAYMELNSLTSEDTAVYFCARHGNYYYYSGMDVWGQGTTVTVSS',
+    'g6_vh':     'EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTV',  
+    'g6_vl':     'DIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIK',
+    'beb_vh':   'QITLKESGPTLVKPTQTLTLTCTFSGFSLSISGVGVGWLRQPPGKALEWLALIYWDDDKRYSPSLKSRLTISKDTSKNQVVLKMTNIDPVDTATYYCAHHSISTIFDHWGQGTLVTVSS',
+    'beb_vl' :   'QSALTQPASVSGSPGQSITISCTATSSDVGDYNYVSWYQQHPGKAPKLMIFEVSDRPSGISNRFSGSKSGNTASLTISGLQAEDEADYYCSSYTTSSAVFGGGTKLTVL',
+    'sa58_vh':  'QVQLAQSGSELRKPGASVKVSCDTSGHSFTSNAIHWVRQAPGQGLEWMGWINTDTGTPTYAQGFTGRFVFSLDTSARTAYLQISSLKADDTAVFYCARERDYSDYFFDYWGQGTLVTVSS',
+    'sa58_vl':   'EVVMTQSPASLSVSPGERATLSCRARASLGISTDLAWYQQRPGQAPRLLIYGASTRATGIPARFSGSGSGTEFTLTISSLQSEDSAVYYCQQYSNWPLTFGGGTKVEIK',
+}
+
+
+def parse_abysis(seq, seq_name):
+    fname = f'data/abysis/abysis_counts_{seq_name}.json'
+    with open(fname) as f:
+        json_data = json.load(f)
+
+    data = []
+    for idx, freq_table in enumerate(json_data['frequencies']):
+        wt = seq[idx]
+        total = freq_table['total']
+        for entry in freq_table['counts']:
+            if entry['aa'] == seq[idx]:
+                wt_frac = entry['c'] / total
+        aa_to_freq, seen = {}, set()
+        for entry in freq_table['counts']:
+            mt = entry['aa']
+            seen.add(mt)
+            counts = entry['c']
+            frac = entry['c'] / total
+            ratio = frac / wt_frac
+            data.append([ idx + 1, wt, mt, counts, frac, ratio ])
+        for mt in AAs - seen:
+            data.append([ idx + 1, wt, mt, 0, 0., 0. ])
+
+    df = pd.DataFrame(data, columns=[
+        'pos',
+        'wt',
+        'mt',
+        'counts',
+        'fraction',
+        'likelihood_ratio',
+    ])
+
+    if any(prefix in seq_name for prefix in ['cr6261', 'cr9114', 'g6']):
+        subdirectory = seq_name.split('_')[0]
+    else:
+        subdirectory = ''
+
+    output_dir = os.path.join('output', 'ab_mutagenesis_expts', subdirectory)
+    os.makedirs(output_dir, exist_ok=True)
+
+    output_path = os.path.join(output_dir, f'abysis_counts_{seq_name}.txt')
+    df.to_csv(output_path, sep='\t')
+
+if __name__ == '__main__':
+    for seq_name in seqs_abs:
+        seq = seqs_abs[seq_name]
+        parse_abysis(seq, seq_name)

model/plot_ab-binding_benchmarks.py

@@ -0,0 +1,755 @@
+import numpy as np
+import pandas as pd
+import bokeh.io
+import bokeh.plotting
+import bokeh.palettes 
+from bokeh.transform import factor_cmap
+import datashader
+import holoviews as hv
+import holoviews.operation.datashader
+hv.extension("bokeh")
+
+import warnings
+# Suppress FutureWarning messages
+warnings.simplefilter(action='ignore',)
+
+# Your Pandas code here
+
+if_conds = ['ag+ab', 'ab only', 'ab']
+
+cr9114_dict = {
+                'ab_name'   : 'CR9114',
+                'files'     : {
+                                'Ab-Ag':      'output/ab_mutagenesis_expts/cr9114/4fqi_ablh_scores.csv',
+                                'Ab only':    'output/ab_mutagenesis_expts/cr9114/4fqi_lh_scores.csv',
+                                'Ab VH only': 'output/ab_mutagenesis_expts/cr9114/4fqi_h_scores.csv',
+                                'ESM-1v':      'output/ab_mutagenesis_expts/cr9114/cr9114_exp_data_maskMargLabeled.csv',
+                                'AbLang':     'output/ab_mutagenesis_expts/cr9114/cr9114_hc_ablangScores.csv',
+                                'abYsis':     'output/ab_mutagenesis_expts/cr9114/abysis_counts_cr9114_vh.txt',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/cr9114/cr9114_exp_data.csv', 
+                                dtype = {'genotype': str},
+                                ).rename(columns={'h1_mean': 'H1', 'h3_mean' : 'H3'}),
+                'ag_columns': ['H1', 'H3'],
+                'expt_type':  'Combinatorial Mutagenesis for Affinity',
+                'palette':     bokeh.palettes.Spectral6
+}
+
+cr6261_dict = {
+                'ab_name'   : 'CR6261',
+                'files'     : {  
+                                'Ab-Ag':      'output/ab_mutagenesis_expts/cr6261/3gbn_ablh_scores.csv',
+                                'Ab only':    'output/ab_mutagenesis_expts/cr6261/3gbn_lh_scores.csv',
+                                'Ab VH only': 'output/ab_mutagenesis_expts/cr6261/3gbn_h_scores.csv',
+                                'ESM-1v':      'output/ab_mutagenesis_expts/cr6261/cr6261_exp_data_maskMargLabeled.csv',
+                                'AbLang':     'output/ab_mutagenesis_expts/cr6261/cr6261_hc_ablangScores.csv',
+                                'abYsis':     'output/ab_mutagenesis_expts/cr6261/abysis_counts_cr6261_vh.txt',
+                            },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/cr6261/cr6261_exp_data.csv',
+                            dtype={'genotype': str},
+                            ).rename(columns={'h1_mean': 'H1', 'h9_mean' : 'H9'}),
+                'ag_columns': ['H1', 'H9'],
+                'expt_type':  'Combinatorial Mutagenesis for Affinity',
+                'palette':     bokeh.palettes.Dark2_6,
+}
+
+
+g6LC_dict = {
+                'ab_name'   : 'g6',
+                'files'     : {
+                                'Ab-Ag':      'output/ab_mutagenesis_expts/g6/2fjg_vlh_lc_scores.csv',
+                                'Ab only':    'output/ab_mutagenesis_expts/g6/2fjg_lh_lc_scores.csv',
+                                'Ab VH/VL only': 'output/ab_mutagenesis_expts/g6/2fjg_l_lc_scores.csv',
+                                'ESM-1v':      'output/ab_mutagenesis_expts/g6/g6Lc_exp_data_maskMargLabeled.csv',
+                                'AbLang':     'output/ab_mutagenesis_expts/g6/g6_lc_ablangScores.csv',
+                                'abYsis':     'output/ab_mutagenesis_expts/g6/abysis_counts_g6_vl.txt',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv'),
+                'ag_columns': ['norm_binding'],
+                'expt_type':  'Deep Mutational Scan for Binding',
+                'palette':     bokeh.palettes.Pastel1_6,
+                'chain':      'VL'
+}
+
+g6HC_dict = {
+                'ab_name'   : 'g6',
+                'files'     : {
+                                'Ab-Ag':      'output/ab_mutagenesis_expts/g6/2fjg_vlh_hc_scores.csv',
+                                'Ab only':      'output/ab_mutagenesis_expts/g6/2fjg_lh_hc_scores.csv',
+                                'Ab VH/VL only': 'output/ab_mutagenesis_expts/g6/2fjg_h_hc_scores.csv',
+                                'ESM-1v':      'output/ab_mutagenesis_expts/g6/g6Hc_exp_data_maskMargLabeled.csv',
+                                'AbLang':     'output/ab_mutagenesis_expts/g6/g6_hc_ablangScores.csv',
+                                'abYsis':     'output/ab_mutagenesis_expts/g6/abysis_counts_g6_vh.txt',
+                            },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv'),
+                'ag_columns': ['norm_binding'],
+                'expt_type':  'Deep Mutational Scan for Binding',
+                'palette':     bokeh.palettes.Pastel1_4,
+                'chain':      'VH'
+}
+def apply_mask_and_average(scores, genotype):
+    #scores here should already be a in log space
+
+    if '1' in genotype:
+        masked_list = []
+        for s, mask_char in zip(scores, genotype):
+            if mask_char == '1':
+                masked_list.append(s)
+        multi_avg = np.mean(masked_list)
+    else:
+        #if genotype is all zeros = wt
+        multi_avg = 0
+    
+    return multi_avg
+
+
+def transform_single_to_multi( dms_df, singleMuts_df, condition, sort_col, ascending ):
+    multi_scores = []
+    #use input parameter as method bc key is kwarg in pd.sort_values
+    #method should be either esm1v or abysis or AbLang
+    
+    #for cr sorts the single mutation dataframe in residue order, ie binary '10000' before '00010'. ascending = False
+    #for g6 sorts the single mutation dataframe in position order, ie residue index. ascending = True
+
+    singleMuts_sorted = singleMuts_df.sort_values(sort_col, key=lambda x: x.astype(int), ascending=ascending)
+    singles_scores = singleMuts_sorted[condition].to_list()
+
+    for g in dms_df['genotype']:
+        multi_scores.append(apply_mask_and_average(singles_scores, g ))
+
+    return multi_scores
+
+def get_abysis_single_scores_df(singleMut_ids, abysis_df):
+    single_scores = []
+    all_pos = []
+
+    for id in singleMut_ids:
+        wt, pos, mt = id[0], int(id[1:-1]), id[-1]
+        likelihood_ratio = abysis_df.loc[(abysis_df['pos'] == pos) & (abysis_df['wt'] == wt) & (abysis_df['mt'] == mt)]['likelihood_ratio'].to_list()[0]
+        log_like_ratio = np.log10(likelihood_ratio)
+        single_scores.append(log_like_ratio)
+        all_pos.append(pos)
+
+    singleMuts_df = pd.DataFrame({'pos': all_pos, 'abYsis': single_scores})
+
+
+    return singleMuts_df
+
+def get_ablang_single_scores_df(singleMut_ids, ablang_df):
+    single_scores = []
+    all_pos = []
+
+    for id in singleMut_ids:
+        wt, pos, mt = id[0], int(id[1:-1]), id[-1]
+        log_likelihood_ratio = ablang_df.loc[(ablang_df['pos'] == pos) & (ablang_df['wt'] == wt) & (ablang_df['mt'] == mt)]['log_likelihood_ratio'].to_list()[0]
+        single_scores.append(log_likelihood_ratio)
+        all_pos.append(pos)
+
+    singleMuts_df = pd.DataFrame({'pos': all_pos, 'AbLang': single_scores})
+
+
+    return singleMuts_df
+
+
+#get melted correlations matrix of structure input x target antigen, for a given antibody 
+def get_corr(ab_name, files, dms_df, ag_columns, dropLOQ = False ):
+
+    inv_fold_files = {key: value for key, value in files.items() if key not in ['ESM-1v', 'AbLang', 'abYsis']}
+    other_files = {key: value for key, value in files.items() if key in ['ESM-1v', 'AbLang', 'abYsis']}
+
+    #retreive correlations from inverse fold scores
+    for key, filepath in inv_fold_files.items():
+        df = pd.read_csv(filepath)  # Read the CSV file into a DataFrame
+        column_label = key
+        log_likelihood_column = df['log_likelihood']  # Extract the 'log_likelihood' column
+        dms_df[column_label] = log_likelihood_column
+
+    #retreive correlations from abysis and esm1v scores
+    for key, filepath in other_files.items():        
+        if key == 'ESM-1v':
+            singleMuts_df = pd.read_csv(filepath)
+            # Get the columns for all esm1v models
+            esm_columns = [col for col in singleMuts_df.columns if col.startswith('esm')]      
+            # Calculate the average for each esm column
+            esm_avg_values = singleMuts_df[esm_columns].mean(axis=1)
+            singleMuts_df[key] = esm_avg_values
+            if ab_name == 'g6':
+                #g6 data set is only single mutations
+                dms_df[key] = esm_avg_values
+            else:
+                #cr datasets are multiple mutations
+                dms_df[key] = transform_single_to_multi(dms_df, singleMuts_df, key, 'genotype', ascending = False) 
+        elif key == 'abYsis':
+            abysis_df = pd.read_csv(filepath, sep = '\t', dtype = {'pos': int})
+            if ab_name == 'g6':
+                dms_df[key] = get_abysis_single_scores_df(dms_df['mutant'], abysis_df)[key]
+            else:
+                #cr data set are multiple mutations. use same scoring strategy as esm1v paper
+                singleMut_ids = pd.read_csv(files['ESM-1v'])['mutant'].to_list()
+                singleMuts_df = get_abysis_single_scores_df(singleMut_ids, abysis_df)
+                dms_df[key] = transform_single_to_multi(dms_df, singleMuts_df, key, 'pos', ascending = True)
+        elif key == 'AbLang':
+            ablang_df = pd.read_csv(filepath, dtype = {'pos': int})
+            if ab_name == 'g6':
+                dms_df[key] = get_ablang_single_scores_df(dms_df['mutant'], ablang_df)[key]
+            else:
+                #cr data set are multiple mutations. use same scoring strategy as esm1v paper
+                singleMut_ids = pd.read_csv(files['ESM-1v'])['mutant'].to_list()
+                singleMuts_df = get_ablang_single_scores_df(singleMut_ids, ablang_df)
+                dms_df[key] = transform_single_to_multi(dms_df, singleMuts_df, key, 'pos', ascending = True)
+
+    conditions = list(files.keys())
+
+    #
+    if not dropLOQ:
+        correlations = dms_df[ag_columns + conditions].corr(method='spearman')
+        correlations = correlations.drop(conditions, axis= 1)
+        correlations = correlations.drop(ag_columns, axis= 0)
+
+        # Melt the correlations DataFrame and rename the index column
+        melted_correlations = (
+            correlations
+            .reset_index()
+            .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            .rename(columns={'index': 'Input'})
+        )
+        melted_correlations['Method'] = ['LM' if (('ESM' in i) or ('AbLang' in i))
+                                         else 'IF' if 'Ab' in i 
+                                        else 'MSA'
+                                        for i in melted_correlations['Input']
+                                    ]
+
+        return melted_correlations
+    #if we wish to drop the samples on the lower limit of quantitation, drop samples from each ag group individually and compute correlation with IF log_like
+    else:
+        all_ag_melted = pd.DataFrame({})
+        for ag in ag_columns:
+            filt_df = dms_df[dms_df[ag] > min(dms_df[ag])]
+
+            correlations = filt_df[[ag] + conditions].corr(method='spearman')
+            correlations = correlations.drop(conditions, axis= 1)
+            correlations = correlations.drop(ag, axis= 0)
+
+            # Melt the correlations DataFrame and rename the index column
+            melted_correlations = (
+                correlations
+                .reset_index()
+                .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            ).rename(columns={'index': 'Input'})
+
+            all_ag_melted = pd.concat([all_ag_melted, melted_correlations], ignore_index= True)
+        all_ag_melted['Method'] =   ['LM' if (('ESM' in i) or ('AbLang' in i))
+                                         else 'IF' if 'Ab' in i 
+                                        else 'MSA'
+                                        for i in all_ag_melted['Input']
+                                    ]
+        
+        return all_ag_melted
+    
+
+#get melted correlations matrix of structure input x target antigen, for a given antibody 
+def get_g6_corr(g6Hc_dict, g6LC_dict, dropLOQ = False ):
+
+    ab_name = 'g6'
+    ag_columns = g6LC_dict['ag_columns']
+    vh_and_vl_dms_df = pd.DataFrame({})
+
+    for d in [g6HC_dict, g6LC_dict]:
+
+        files = d['files']
+        dms_df = d['dms_df']
+
+        inv_fold_files =  {key: value for key, value in files.items() if key not in ['ESM-1v', 'AbLang', 'abYsis']}
+        other_files =  {key: value for key, value in files.items() if key in ['ESM-1v', 'AbLang', 'abYsis']}
+
+        #retreive correlations from inverse fold scores
+        for key, filepath in inv_fold_files.items():
+            df = pd.read_csv(filepath)  # Read the CSV file into a DataFrame
+            column_label = key
+            log_likelihood_column = df['log_likelihood']  # Extract the 'log_likelihood' column
+            dms_df[column_label] = log_likelihood_column
+
+        #retreive correlations from abysis, ablang, and esm1v scores
+        for key, filepath in other_files.items():        
+            if key == 'ESM-1v':
+                singleMuts_df = pd.read_csv(filepath)
+                # Get the columns for all esm1v models
+                esm_columns = [col for col in singleMuts_df.columns if col.startswith('esm')]      
+                # Calculate the average for each esm column
+                esm_avg_values = singleMuts_df[esm_columns].mean(axis=1)
+                singleMuts_df[key] = esm_avg_values
+                if ab_name == 'g6':
+                    #g6 data set is only single mutations
+                    dms_df[key] = esm_avg_values
+            elif key == 'abYsis':
+                abysis_df = pd.read_csv(filepath, sep = '\t', dtype = {'pos': int})
+                if ab_name == 'g6':
+                    dms_df[key] = get_abysis_single_scores_df(dms_df['mutant'], abysis_df)[key]
+            elif key == 'AbLang':
+                ablang_df = pd.read_csv(filepath, dtype = {'pos': int})
+                if ab_name == 'g6':
+                    dms_df[key] = get_ablang_single_scores_df(dms_df['mutant'], ablang_df)[key]
+
+
+        vh_and_vl_dms_df = pd.concat([vh_and_vl_dms_df, dms_df], ignore_index= True)
+        
+        conditions = list(files.keys())
+
+    #
+    if not dropLOQ:
+        correlations = vh_and_vl_dms_df[ag_columns + conditions].corr(method='spearman')
+        correlations = correlations.drop(conditions, axis= 1)
+        correlations = correlations.drop(ag_columns, axis= 0)
+
+        # Melt the correlations DataFrame and rename the index column
+        melted_correlations = (
+            correlations
+            .reset_index()
+            .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            .rename(columns={'index': 'Input'})
+        )
+        melted_correlations['Method'] = ['LM' if (('ESM' in i) or ('AbLang' in i))
+                                         else 'IF' if 'Ab' in i 
+                                        else 'MSA'
+                                        for i in melted_correlations['Input']
+                                    ]
+        melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations)
+
+        return melted_correlations
+    #if we wish to drop the samples on the lower limit of quantitation, drop samples from each ag group individually and compute correlation with IF log_like
+    else:
+        all_ag_melted = pd.DataFrame({})
+        for ag in ag_columns:
+            filt_df = vh_and_vl_dms_df[vh_and_vl_dms_df[ag] > min(vh_and_vl_dms_df[ag])]
+
+            correlations = filt_df[[ag] + conditions].corr(method='spearman')
+            correlations = correlations.drop(conditions, axis= 1)
+            correlations = correlations.drop(ag, axis= 0)
+
+            # Melt the correlations DataFrame and rename the index column
+            melted_correlations = (
+                correlations
+                .reset_index()
+                .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            ).rename(columns={'index': 'Input'})
+
+            all_ag_melted = pd.concat([all_ag_melted, melted_correlations], ignore_index= True)
+        all_ag_melted['Method'] = ['LM' if (('ESM' in i) or ('AbLang' in i))
+                                         else 'IF' if 'Ab' in i 
+                                        else 'MSA'
+                                        for i in all_ag_melted['Input']
+                                    ]
+        melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations)
+
+        return all_ag_melted    
+
+
+#plot correlation bars for cr antibodies: cr6261  and cr9114
+def plot_hbar(title, melted_correlations, palette = bokeh.palettes.Spectral6, compare = 'model'):
+
+    if compare == 'model':
+        melted_correlations.replace('Ab-Ag','InverseFolding', inplace = True)
+        color_by = 'Input'
+    elif compare == 'IF':
+        color_by = 'Method'
+    else:
+        color_by = 'Target'
+
+
+    melted_correlations['cats'] = melted_correlations.apply(lambda x: (x["Target"], x["Input"]), axis = 1)
+    factors = list(melted_correlations.cats)[::-1]
+
+    if 'CR6261' in title:
+        x_min = -0.15
+    else:
+        x_min = 0
+    
+    p = bokeh.plotting.figure(
+    height=340,
+    width=440,
+    x_axis_label="Spearman Correlation",
+    x_range=[x_min, 1],
+    y_range=bokeh.models.FactorRange(*factors),
+    tools="save",
+    title = title
+    )
+        
+
+    p.hbar(
+        source=melted_correlations,
+        y="cats",
+        right="Correlation",
+        height=0.6,
+        line_color = 'black',
+        legend_field = color_by,
+        # use the palette to colormap based on the the x[1:2] values
+        fill_color=factor_cmap(
+                                color_by, 
+                                palette= palette, 
+                                factors = list(melted_correlations[color_by].unique()), 
+                                start=1, 
+                                end=3)
+    )
+
+    #labels_df = melted_correlations[melted_correlations['Input'].isin(['Ab-Ag','ESM-1v', 'abYsis', 'InverseFold'])]
+    labels_df = melted_correlations
+    labels_df['corr_str'] = labels_df['Correlation'].apply(lambda x: round(x, 2)).astype(str)
+    labels_source = bokeh.models.ColumnDataSource(labels_df)
+
+    labels = bokeh.models.LabelSet(x='Correlation', y='cats', text='corr_str',text_font_size = "10px",
+              x_offset=12, y_offset=-5, source=labels_source, render_mode='canvas')
+
+    p.ygrid.grid_line_color = None
+    p.y_range.range_padding = 0.1
+    p.add_layout(labels)
+
+    if compare == 'IF':
+        p.legend.orientation = "vertical"
+        p.legend.location = "bottom_right"
+        p.legend.label_text_font_size = "8pt"
+    else:
+        p.legend.visible = False
+
+    p.output_backend = "svg"
+    return p
+
+def plot_g6_scatter(g6Hc_df,  g6Lc_df, ag_columns, dropLOQ):
+    ag = ag_columns[0]
+    if dropLOQ:
+        g6Lc_df = g6Lc_df[g6Lc_df[ag] > np.min(g6Lc_df[ag])]
+        g6Hc_df = g6Hc_df[g6Hc_df[ag] > np.min(g6Hc_df[ag])]
+
+
+    hv.extension("bokeh")
+
+    plots = []
+    for (plot_df,title) in [(g6Hc_df, 'g6.31 VH'),(g6Lc_df, 'g6.31 VL')]:
+
+        # Generate HoloViews Points Element
+        points = hv.Points(
+            data=plot_df,
+            kdims=[ag, 'Ab-Ag'],
+        )
+
+        # Datashade with spreading of points
+        p = hv.operation.datashader.dynspread(
+            hv.operation.datashader.rasterize(
+                points
+            ).opts(
+                cmap='Magma',
+                cnorm='linear',
+            )
+            ).opts(
+                frame_width=350,
+                frame_height=300,
+                padding=0.05,
+                show_grid=False,
+                colorbar = True
+            )
+
+        p = hv.render(p)
+        p.output_backend = "svg"
+
+        plots.append(p)
+
+    return plots
+
+#plot scatter plot colored by n_muts for CR antibodies
+def single_cr_scatterMut(plot_df, ag_columns, x, y, title, ):
+    x_axis_label = x + '-logKd' if x in ag_columns else 'log likelihood'
+    y_axis_label = y + '-logKd' if y in ag_columns else 'log likelihood'
+
+
+    #set up figure
+    p = bokeh.plotting.figure(
+        width=500,
+        height=400,
+        x_axis_label= x_axis_label,
+        y_axis_label= y_axis_label,
+        title = title
+    )
+
+    #Set up color mapper to color by number of mutations from germline
+    mapper = bokeh.transform.linear_cmap(field_name='som_mut', palette=bokeh.palettes.inferno(16) ,low=min(plot_df.som_mut) ,high=max(plot_df.som_mut))
+    p.circle(
+        source=plot_df,
+        x= x,
+        y= y,
+        alpha = 0.2, 
+        line_color=mapper,
+        color=mapper,
+        size = 4
+    )
+
+    #
+    p.diamond(
+        #germline seqid is all 0's
+        source = plot_df[plot_df['genotype'].str.contains('1') == False],
+        x= x,
+        y= y,
+        color = 'dodgerblue',
+        size = 12,
+        legend_label = 'germline'
+    )
+    p.star(
+        #mature seqid is all 0's
+        source = plot_df[plot_df['genotype'].str.contains('0') == False],
+        x= x,
+        y= y,
+        color = 'green',
+        size = 12,
+        legend_label = 'mature'
+    )
+    p.legend.location = "top_left"
+    p.output_backend = "svg"
+
+    return p, mapper
+
+#plot single scatter plot for CR antibodies - with dynamic spreading and datashading
+def single_cr_scatter(plot_df, ag_columns, x, y, title, ):
+    x_axis_label = x + '-logKd' if x in ag_columns else 'log likelihood'
+    y_axis_label = y + '-logKd' if y in ag_columns else 'log likelihood'
+
+    hv.extension("bokeh")
+
+    # Generate HoloViews Points Element
+    points = hv.Points(
+        data=plot_df,
+        kdims=[x, y],
+    )
+
+    # Datashade with spreading of points
+    p = hv.operation.datashader.dynspread(
+        hv.operation.datashader.rasterize(
+            points
+        ).opts(
+            cmap='Magma',
+            cnorm='linear',
+        )
+        ).opts(
+            frame_width=350,
+            frame_height=300,
+            padding=0.05,
+            show_grid=False,
+            colorbar = True
+        )
+
+    p = hv.render(p)
+    
+    p.grid.visible = False
+    #
+    p.diamond(
+        #germline seqid is all 0's
+        source = plot_df[plot_df['genotype'].str.contains('1') == False],
+        x= x,
+        y= y,
+        color = 'dodgerblue',
+        size = 12,
+        #legend_label = 'germline'
+    )
+    p.star(
+        #mature seqid is all 0's
+        source = plot_df[plot_df['genotype'].str.contains('0') == False],
+        x= x,
+        y= y,
+        color = 'green',
+        size = 12,
+        #legend_label = 'mature'
+    )
+    #p.legend.location = "top_left"
+    p.output_backend = "svg"
+
+    return p
+
+def plot_CR_scatter_colorMuts(title, files, dms_df, ag_columns, dropLOQ = False):
+
+    agAb_files  = {key: value for key, value in files.items() if key in ['Ab-Ag',]}
+
+    for key, filepath in agAb_files.items():
+        df = pd.read_csv(filepath)  # Read the CSV file into a DataFrame
+        column_label = key
+        log_likelihood_column = df['log_likelihood']  # Extract the 'log_likelihood' column
+        dms_df[column_label] = log_likelihood_column
+
+    plots = []
+    #predictions vs experiment
+    for ag in ag_columns:
+
+        if dropLOQ:
+            plot_df = dms_df[dms_df[ag] > min(dms_df[ag])]
+            #add back wt which is on LOQ for visualization   
+            plot_df = pd.concat([plot_df, dms_df[dms_df['som_mut'] == 0]], ignore_index = True, axis = 0)
+        else:
+            plot_df = dms_df
+
+        p, _ = single_cr_scatterMut(plot_df, ag_columns, x = ag, y = 'Ab-Ag', title = f'{title} against {ag}')
+        plots.append(p)
+
+    #experimental vs experimental 
+    plot_df = dms_df
+    if dropLOQ: 
+        for ag in ag_columns:
+            plot_df = plot_df[plot_df[ag] > min(plot_df[ag])]
+        #add back wt which is on LOQ for visualization   
+        plot_df = pd.concat([plot_df, dms_df[dms_df['som_mut'] == 0]], ignore_index = True, axis = 0)
+
+    p_exp, mapper = single_cr_scatterMut(plot_df, ag_columns, x= ag_columns[0], y = ag_columns[1], title = 'Experimental Cross-Reactive Binding Landscape')
+    p_exp.legend.location = "top_left"
+    color_bar = bokeh.models.ColorBar(color_mapper=mapper['transform'], width=8, title = 'Amino Acid Mutations')
+    #p_exp.add_layout(color_bar, 'right')
+
+    plots.append(p_exp)
+    p.output_backend = "svg"
+
+    return plots
+
+def plot_CR_scatterplots(title, files, dms_df, ag_columns, dropLOQ = False):
+
+    agAb_files  = {key: value for key, value in files.items() if key in ['Ab-Ag',]}
+
+    for key, filepath in agAb_files.items():
+        df = pd.read_csv(filepath)  # Read the CSV file into a DataFrame
+        column_label = key
+        log_likelihood_column = df['log_likelihood']  # Extract the 'log_likelihood' column
+        dms_df[column_label] = log_likelihood_column
+
+    plots = []
+    #predictions vs experiment
+    for ag in ag_columns:
+
+        if dropLOQ:
+            plot_df = dms_df[dms_df[ag] > min(dms_df[ag])]
+            #add back wt which is on LOQ for visualization   
+            plot_df = pd.concat([plot_df, dms_df[dms_df['som_mut'] == 0]], ignore_index = True, axis = 0)
+        else:
+            plot_df = dms_df
+
+        p = single_cr_scatter(plot_df, ag_columns, x = ag, y = 'Ab-Ag', title = f'{title} against {ag}')
+        plots.append(p)
+
+    #experimental vs experimental 
+    plot_df = dms_df
+    if dropLOQ: 
+        for ag in ag_columns:
+            plot_df = plot_df[plot_df[ag] > min(plot_df[ag])]
+        #add back wt which is on LOQ for visualization   
+        plot_df = pd.concat([plot_df, dms_df[dms_df['som_mut'] == 0]], ignore_index = True, axis = 0)
+
+    p_exp = single_cr_scatter(plot_df, ag_columns, x= ag_columns[0], y = ag_columns[1], title = 'Experimental Cross-Reactive Binding Landscape')
+
+    plots.append(p_exp)
+    p_exp.output_backend = "svg"
+
+    return plots
+
+#plot correlation bars for cr antibodies: cr6261  and cr9114
+def plot_LOQ_comparison(ab_name, unfilt_corrs, filt_corrs, palette = bokeh.palettes.Spectral6,):
+    unfilt_corrs['LOQ'] = ['Included'] * len(unfilt_corrs)
+    filt_corrs['LOQ'] = ['Excluded'] * len(filt_corrs)
+    combined_melt = pd.concat([unfilt_corrs,filt_corrs], ignore_index= True)
+
+    combined_melt['cats'] = combined_melt.apply(lambda x: (x["Target"], x["Input"], x["LOQ"]), axis = 1)
+    factors = list(combined_melt.cats)[::-1]
+    
+    p = bokeh.plotting.figure(
+    height=850,
+    width=400,
+    x_axis_label="Spearman Correlation",
+    x_range=[-.25, 1],
+    y_range=bokeh.models.FactorRange(*factors),
+    tools="save",
+    title = f'{ab_name}, Impact of Ignoring Lower LOQ '
+    )
+
+    p.hbar(
+        source=combined_melt,
+        y="cats",
+        right="Correlation",
+        height=0.7,
+        line_color = 'black',
+        legend_field = 'LOQ',
+        # use the palette to colormap based on the the x[1:2] values
+        fill_color=factor_cmap(
+                                'LOQ', 
+                                palette= palette, 
+                                factors = list(combined_melt['LOQ'].unique()), 
+                                start=1, 
+                                end=2)
+    )
+
+    p.ygrid.grid_line_color = None
+    p.y_range.range_padding = 0.1
+    p.legend.location = "bottom_right"
+    p.legend.orientation = "vertical"
+    p.yaxis.axis_label = "Model Input"
+    p.output_backend = "svg"
+
+    return p
+
+
+
+if __name__ == '__main__':
+
+    datasets = [cr9114_dict, cr6261_dict, (g6HC_dict, g6LC_dict) ]
+    results = []
+    
+    #compute with and without points on lower limit of quantitation ignored
+    for dropLOQ in [False, True]:
+
+        g6_corrs = []
+        all_corr_plots = []
+
+        for d in datasets:
+            if type(d) is tuple:
+                g6Hc, g6Lc = d
+                title = g6Hc['ab_name'] + ', ' + g6Hc['expt_type']
+                g6_combined = get_g6_corr(g6Hc, g6Lc, dropLOQ= dropLOQ)
+
+                all_corr_plots.append(plot_hbar(title, g6_combined, g6Lc['palette'], 'IF'))
+
+                #also plot scatter plots
+                scatter_fname = f"output/ab_mutagenesis_expts/g6_scatterplot_colorMuts{'_dropLOQ' if dropLOQ else ''}.html"
+                bokeh.plotting.output_file(scatter_fname)
+                g6_scatterMuts_plots = plot_g6_scatter(g6Hc['dms_df'],  g6Lc['dms_df'], g6Hc['ag_columns'], dropLOQ= dropLOQ)
+                bokeh.io.show(bokeh.layouts.gridplot(g6_scatterMuts_plots, ncols = 2))
+
+                if (dropLOQ == False):
+                    results.append(g6_combined)
+
+            else:
+                corr_df = get_corr(d['ab_name'], d['files'], d['dms_df'], d['ag_columns'], dropLOQ= dropLOQ)
+                title = d['ab_name'] + ', ' + d['expt_type']
+
+                all_corr_plots.append(plot_hbar(title, corr_df, d['palette'], 'IF'))
+
+                #also plot scatter plots for CR9114, CR6261
+                scatter_fname = f"output/ab_mutagenesis_expts/{d['ab_name']}_scatter_colorMuts{'_dropLOQ' if dropLOQ else ''}.html"
+                bokeh.plotting.output_file(scatter_fname)
+                cr_scatterMuts_plots = plot_CR_scatter_colorMuts(title,  d['files'], d['dms_df'], d['ag_columns'], dropLOQ= dropLOQ)
+                bokeh.io.show(bokeh.layouts.gridplot(cr_scatterMuts_plots, ncols = 3))
+
+                #also scatter plots with datashading/dynamic spreading for CR9114, CR6261
+                scatter_fname = f"output/ab_mutagenesis_expts/{d['ab_name']}_scatter{'_dropLOQ' if dropLOQ else ''}.html"
+                bokeh.plotting.output_file(scatter_fname)
+                cr_scattter_plots = plot_CR_scatterplots(title,  d['files'], d['dms_df'], d['ag_columns'], dropLOQ= dropLOQ)
+                bokeh.io.show(bokeh.layouts.gridplot(cr_scattter_plots, ncols = 3))
+
+                if (dropLOQ == False):
+                    results.append(corr_df)
+     
+        all_corr_fname = f"output/ab_mutagenesis_expts/benchmarks{'_dropLOQ' if dropLOQ else ''}.html"
+        bokeh.plotting.output_file(all_corr_fname)
+        bokeh.io.show(bokeh.layouts.gridplot(all_corr_plots, ncols = len(all_corr_plots)))
+
+    compareLOQ_fname = 'output/ab_mutagenesis_expts/LOQ_comparison.html'
+    bokeh.plotting.output_file(compareLOQ_fname)
+    compareLOQ_plots = []
+    for d in datasets:
+        if (type(d) != tuple) and (d['ab_name'].startswith('CR')):  
+            corr_withLOQ =  get_corr(d['ab_name'], d['files'], d['dms_df'], d['ag_columns'], )
+            corr_withoutLOQ =  get_corr(d['ab_name'], d['files'], d['dms_df'], d['ag_columns'], dropLOQ= dropLOQ)
+
+            compareLOQ_plots.append(plot_LOQ_comparison(d['ab_name'], corr_withLOQ, corr_withoutLOQ , d['palette']))
+    bokeh.io.show(bokeh.layouts.gridplot(compareLOQ_plots, ncols = len(compareLOQ_plots)))
+
+    pd.concat(results, ignore_index= True).to_csv( 'output/ab_mutagenesis_expts/results.csv' , index = False)

model/plot_esm1v_benchmarks.py

@@ -0,0 +1,317 @@
+import numpy as np
+import pandas as pd
+import bokeh.io
+import bokeh.plotting
+import bokeh.palettes 
+from bokeh.transform import factor_cmap
+import datashader
+import holoviews as hv
+import holoviews.operation.datashader
+hv.extension("bokeh")
+
+import warnings
+# Suppress FutureWarning messages
+warnings.simplefilter(action='ignore',)
+
+cr9114_dict = {
+                'ab_name'   : 'CR9114',
+                'files'     : {
+                                'ESM-1v Ab-Ag':      'output/ab_mutagenesis_expts/cr9114/cr9114_esm1vAbAg_exp_data_maskMargLabeled.csv',
+                                'ESM-1v Ab only':    'output/ab_mutagenesis_expts/cr9114/cr9114_esm1vbothchains_exp_data_maskMargLabeled.csv',
+                                'ESM-1v Ab VH only': 'output/ab_mutagenesis_expts/cr9114/cr9114_exp_data_maskMargLabeled.csv',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/cr9114/cr9114_exp_data.csv', 
+                                dtype = {'genotype': str},
+                                ).rename(columns={'h1_mean': 'H1', 'h3_mean' : 'H3'}),
+                'ag_columns': ['H1', 'H3'],
+                'expt_type':  'Combinatorial Mutagenesis for Affinity',
+                'palette':     bokeh.palettes.Spectral6
+}
+
+cr6261_dict = {
+                'ab_name'   : 'CR6261',
+                'files'     : {
+                                'ESM-1v Ab-Ag':      'output/ab_mutagenesis_expts/cr6261/cr6261_esm1vAbAg_exp_data_maskMargLabeled.csv',
+                                'ESM-1v Ab only':    'output/ab_mutagenesis_expts/cr6261/cr6261_esm1vbothchains_exp_data_maskMargLabeled.csv',
+                                'ESM-1v Ab VH only': 'output/ab_mutagenesis_expts/cr6261/cr6261_exp_data_maskMargLabeled.csv',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/cr6261/cr6261_exp_data.csv',
+                            dtype={'genotype': str},
+                            ).rename(columns={'h1_mean': 'H1', 'h9_mean' : 'H9'}),
+                'ag_columns': ['H1', 'H9'],
+                'expt_type':  'Combinatorial Mutagenesis for Affinity',
+                'palette':     bokeh.palettes.Dark2_6,
+}
+
+g6LC_dict = {
+                'ab_name'   : 'g6',
+                'files'     : {
+                                'LM Ab-Ag':      'output/ab_mutagenesis_expts/g6/g6Lc_esm1vAbAg_exp_data_maskMargLabeled.csv',
+                                'LM Ab only':    'output/ab_mutagenesis_expts/g6/g6Lc_esm1vbothchains_exp_data_maskMargLabeled.csv',
+                                'LM Ab VH/VL only': 'output/ab_mutagenesis_expts/g6/g6Lc_exp_data_maskMargLabeled.csv',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv'),
+                'ag_columns': ['norm_binding'],
+                'expt_type':  'Deep Mutational Scan for Binding',
+                'palette':     bokeh.palettes.Pastel1_6,
+                'chain':      'VL'
+}
+
+g6HC_dict = {
+                'ab_name'   : 'g6',
+                'files'     : {
+                                'LM Ab-Ag':      'output/ab_mutagenesis_expts/g6/g6Hc_esm1vAbAg_exp_data_maskMargLabeled.csv',
+                                'LM Ab only':    'output/ab_mutagenesis_expts/g6/g6Hc_esm1vbothchains_exp_data_maskMargLabeled.csv',
+                                'LM Ab VH/VL only': 'output/ab_mutagenesis_expts/g6/g6Hc_exp_data_maskMargLabeled.csv',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv'),
+                'ag_columns': ['norm_binding'],
+                'expt_type':  'Deep Mutational Scan for Binding',
+                'palette':     bokeh.palettes.Pastel1_4,
+                'chain':      'VH'
+}
+
+def apply_mask_and_average(scores, genotype):
+    #scores here should already be a in log space
+
+    if '1' in genotype:
+        masked_list = []
+        for s, mask_char in zip(scores, genotype):
+            if mask_char == '1':
+                masked_list.append(s)
+        multi_avg = np.mean(masked_list)
+    else:
+        #if genotype is all zeros = wt
+        multi_avg = 0
+    
+    return multi_avg
+
+
+def transform_single_to_multi( dms_df, singleMuts_df, condition, sort_col, ascending ):
+    multi_scores = []
+    #use input parameter as method bc key is kwarg in pd.sort_values
+    #method should be either esm1v or abysis or AbLang
+    
+    #for cr sorts the single mutation dataframe in residue order, ie binary '10000' before '00010'. ascending = False
+    #for g6 sorts the single mutation dataframe in position order, ie residue index. ascending = True
+
+    singleMuts_sorted = singleMuts_df.sort_values(sort_col, key=lambda x: x.astype(int), ascending=ascending)
+    singles_scores = singleMuts_sorted[condition].to_list()
+
+    for g in dms_df['genotype']:
+        multi_scores.append(apply_mask_and_average(singles_scores, g ))
+
+    return multi_scores
+
+#get melted correlations matrix of structure input x target antigen, for a given antibody 
+def get_corr(ab_name, files, dms_df, ag_columns, dropLOQ = False ):
+
+    #retreive correlations from abysis and esm1v scores
+    for key, filepath in files.items():        
+        singleMuts_df = pd.read_csv(filepath)
+        # Get the columns for all esm1v models
+        esm_columns = [col for col in singleMuts_df.columns if col.startswith('esm')]      
+        # Calculate the average for each esm column
+        esm_avg_values = singleMuts_df[esm_columns].mean(axis=1)
+        singleMuts_df[key] = esm_avg_values
+        if ab_name == 'g6':
+            #g6 data set is only single mutations
+            dms_df[key] = esm_avg_values
+        else:
+            #cr datasets are multiple mutations
+            dms_df[key] = transform_single_to_multi(dms_df, singleMuts_df, key, 'genotype', ascending = False) 
+
+
+    conditions = list(files.keys())
+
+    if not dropLOQ:
+        correlations = dms_df[ag_columns + conditions].corr(method='spearman')
+        correlations = correlations.drop(conditions, axis= 1)
+        correlations = correlations.drop(ag_columns, axis= 0)
+
+        # Melt the correlations DataFrame and rename the index column
+        melted_correlations = (
+            correlations
+            .reset_index()
+            .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            .rename(columns={'index': 'Input'})
+        )
+        melted_correlations['Method'] = ['LM' for i in melted_correlations['Input']]
+
+        return melted_correlations
+    #if we wish to drop the samples on the lower limit of quantitation, drop samples from each ag group individually and compute correlation with IF log_like
+    else:
+        all_ag_melted = pd.DataFrame({})
+        for ag in ag_columns:
+            filt_df = dms_df[dms_df[ag] > min(dms_df[ag])]
+
+            correlations = filt_df[[ag] + conditions].corr(method='spearman')
+            correlations = correlations.drop(conditions, axis= 1)
+            correlations = correlations.drop(ag, axis= 0)
+
+            # Melt the correlations DataFrame and rename the index column
+            melted_correlations = (
+                correlations
+                .reset_index()
+                .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            ).rename(columns={'index': 'Input'})
+
+            all_ag_melted = pd.concat([all_ag_melted, melted_correlations], ignore_index= True)
+        all_ag_melted['Method'] =   ['LM' for i in all_ag_melted['Input']]
+        
+        return all_ag_melted
+    
+#get melted correlations matrix of structure input x target antigen, for a given antibody 
+def get_g6_corr(g6Hc_dict, g6LC_dict, dropLOQ = False ):
+
+    ab_name = 'g6'
+    ag_columns = g6LC_dict['ag_columns']
+    vh_and_vl_dms_df = pd.DataFrame({})
+
+    for d in [g6HC_dict, g6LC_dict]:
+
+        files = d['files']
+        dms_df = d['dms_df']
+
+        #retreive correlations from abysis, ablang, and esm1v scores
+        for key, filepath in files.items():        
+            singleMuts_df = pd.read_csv(filepath)
+            # Get the columns for all esm1v models
+            esm_columns = [col for col in singleMuts_df.columns if col.startswith('esm')]      
+            # Calculate the average for each esm column
+            esm_avg_values = singleMuts_df[esm_columns].mean(axis=1)
+            singleMuts_df[key] = esm_avg_values
+            if ab_name == 'g6':
+                #g6 data set is only single mutations
+                dms_df[key] = esm_avg_values
+
+
+        vh_and_vl_dms_df = pd.concat([vh_and_vl_dms_df, dms_df], ignore_index= True)
+
+        conditions = list(files.keys())
+
+    #
+    if not dropLOQ:
+        correlations = vh_and_vl_dms_df[ag_columns + conditions].corr(method='spearman')
+        correlations = correlations.drop(conditions, axis= 1)
+        correlations = correlations.drop(ag_columns, axis= 0)
+
+        # Melt the correlations DataFrame and rename the index column
+        melted_correlations = (
+            correlations
+            .reset_index()
+            .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            .rename(columns={'index': 'Input'})
+        )
+        melted_correlations['Method'] = ['LM' for i in melted_correlations['Input']]
+        melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations)
+
+        return melted_correlations
+    #if we wish to drop the samples on the lower limit of quantitation, drop samples from each ag group individually and compute correlation with IF log_like
+    else:
+        all_ag_melted = pd.DataFrame({})
+        for ag in ag_columns:
+            filt_df = vh_and_vl_dms_df[vh_and_vl_dms_df[ag] > min(vh_and_vl_dms_df[ag])]
+
+            correlations = filt_df[[ag] + conditions].corr(method='spearman')
+            correlations = correlations.drop(conditions, axis= 1)
+            correlations = correlations.drop(ag, axis= 0)
+
+            # Melt the correlations DataFrame and rename the index column
+            melted_correlations = (
+                correlations
+                .reset_index()
+                .melt(id_vars='index', var_name='Target', value_name='Correlation')
+            ).rename(columns={'index': 'Input'})
+
+            all_ag_melted = pd.concat([all_ag_melted, melted_correlations], ignore_index= True)
+        all_ag_melted['Method'] = ['LM' for i in melted_correlations['Input']]
+        melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations)
+
+        return all_ag_melted    
+    
+#plot correlation bars for cr antibodies: cr6261  and cr9114
+def plot_hbar(title, melted_correlations, palette = bokeh.palettes.Spectral6, compare = 'model'):
+
+    if compare == 'model':
+        melted_correlations.replace('Ab-Ag','InverseFolding', inplace = True)
+        color_by = 'Input'
+    elif compare == 'IF':
+        color_by = 'Method'
+    else:
+        color_by = 'Target'
+
+
+    melted_correlations['cats'] = melted_correlations.apply(lambda x: (x["Target"], x["Input"]), axis = 1)
+    factors = list(melted_correlations.cats)[::-1]
+    
+    p = bokeh.plotting.figure(
+    height=340,
+    width=440,
+    x_axis_label="Spearman Correlation",
+    x_range=[0, 1],
+    y_range=bokeh.models.FactorRange(*factors),
+    tools="save",
+    title = title
+    )
+        
+
+    p.hbar(
+        source=melted_correlations,
+        y="cats",
+        right="Correlation",
+        height=0.6,
+        line_color = 'black',
+        legend_field = color_by,
+        # use the palette to colormap based on the the x[1:2] values
+        # fill_color=factor_cmap(
+        #                         color_by, 
+        #                         palette= palette, 
+        #                         factors = list(melted_correlations[color_by].unique()), 
+        #                         start=1, 
+        #                         end=3)
+        fill_color=bokeh.palettes.Dark2_6[1]
+             
+    )
+
+    #labels_df = melted_correlations[melted_correlations['Input'].isin(['Ab-Ag','ESM-1v', 'abYsis', 'InverseFold'])]
+    labels_df = melted_correlations
+    labels_df['corr_str'] = labels_df['Correlation'].apply(lambda x: round(x, 2)).astype(str)
+    labels_source = bokeh.models.ColumnDataSource(labels_df)
+
+    labels = bokeh.models.LabelSet(x='Correlation', y='cats', text='corr_str',text_font_size = "10px",
+              x_offset=12, y_offset=-5, source=labels_source, render_mode='canvas')
+
+    p.ygrid.grid_line_color = None
+    p.y_range.range_padding = 0.1
+    p.add_layout(labels)
+    p.legend.visible = False
+
+    p.output_backend = "svg"
+    return p
+
+if __name__ == '__main__':
+    datasets = [cr9114_dict, cr6261_dict, (g6HC_dict, g6LC_dict) ]
+        
+    #compute with and without points on lower limit of quantitation ignored
+    for dropLOQ in [False]:
+
+        all_corr_plots = []
+
+        for d in datasets:
+            if type(d) is tuple:
+                g6Hc, g6Lc = d
+                title = g6Hc['ab_name'] + ', ' + g6Hc['expt_type']
+                g6_combined = get_g6_corr(g6Hc, g6Lc, dropLOQ= dropLOQ)
+            
+                all_corr_plots.append(plot_hbar(title, g6_combined, g6Lc['palette'], 'IF'))
+
+            else:
+                corr_df = get_corr(d['ab_name'], d['files'], d['dms_df'], d['ag_columns'], dropLOQ= dropLOQ)
+                title = d['ab_name'] + ', ' + d['expt_type']
+                all_corr_plots.append(plot_hbar(title, corr_df, d['palette'], 'Target'))
+
+    
+        all_corr_fname = f"output/ab_mutagenesis_expts/esm1v_benchmarks{'_dropLOQ' if dropLOQ else ''}.html"
+        bokeh.plotting.output_file(all_corr_fname)
+        bokeh.io.show(bokeh.layouts.gridplot(all_corr_plots, ncols = len(all_corr_plots)))

model/plot_mpnn_benchmarks.py

@@ -0,0 +1,232 @@
+import numpy as np
+import pandas as pd
+import bokeh.io
+import bokeh.plotting
+import bokeh.palettes 
+from bokeh.transform import factor_cmap
+import datashader
+import holoviews as hv
+import holoviews.operation.datashader
+import os
+from natsort import natsorted
+
+hv.extension("bokeh")
+
+import warnings
+# Suppress FutureWarning messages
+warnings.simplefilter(action='ignore',)
+
+cr9114_dict = {
+                'ab_name'   : 'CR9114',
+                'files'     : {
+                                'ESM-IF1':      'output/ab_mutagenesis_expts/cr9114/4fqi_ablh_scores.csv',
+                                'ProteinMPNN': 'output/ab_mutagenesis_expts/cr9114/mpnn/score_only/',
+
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/cr9114/cr9114_exp_data.csv', 
+                                dtype = {'genotype': str},
+                                ).rename(columns={'h1_mean': 'H1', 'h3_mean' : 'H3'}),
+                'ag_columns': ['H1', 'H3'],
+                'expt_type':  'Combinatorial Mutagenesis for Affinity',
+                'palette':     bokeh.palettes.Spectral6
+}
+
+cr6261_dict = {
+                'ab_name'   : 'CR6261',
+                'files'     : {
+                                'ESM-IF1': 'output/ab_mutagenesis_expts/cr6261/3gbn_ablh_scores.csv',
+                                'ProteinMPNN': 'output/ab_mutagenesis_expts/cr6261/mpnn/score_only/'
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/cr6261/cr6261_exp_data.csv',
+                            dtype={'genotype': str},
+                            ).rename(columns={'h1_mean': 'H1', 'h9_mean' : 'H9'}),
+                'ag_columns': ['H1', 'H9'],
+                'expt_type':  'Combinatorial Mutagenesis for Affinity',
+                'palette':     bokeh.palettes.Dark2_6,
+}
+
+
+g6LC_dict = {
+                'ab_name'   : 'g6',
+                'files'     : {
+                                'ESM-IF1':'output/ab_mutagenesis_expts/g6/2fjg_vlh_lc_scores.csv',
+                                'ProteinMPNN':'output/ab_mutagenesis_expts/g6/proteinMpnnLC/score_only/',
+                                
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv'),
+                'ag_columns': ['norm_binding'],
+                'expt_type':  'Deep Mutational Scan for Binding',
+                'palette':     bokeh.palettes.Pastel1_6,
+                'chain':      'VL'
+}
+
+g6HC_dict = {
+                'ab_name'   : 'g6',
+                'files'     : {
+                                'ESM-IF1': 'output/ab_mutagenesis_expts/g6/2fjg_vlh_hc_scores.csv',
+                                'ProteinMPNN':'output/ab_mutagenesis_expts/g6/proteinMpnnHC/score_only/',
+                                },
+                'dms_df'    : pd.read_csv('data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv'),
+                'ag_columns': ['norm_binding'],
+                'expt_type':  'Deep Mutational Scan for Binding',
+                'palette':     bokeh.palettes.Pastel1_4,
+                'chain':      'VH'
+}
+def combine_mpnn_scores(dir):
+    global_scores = []
+    npz_files = [f for f in os.listdir(dir) if f.endswith('.npz')]
+    npz_files = natsorted(npz_files)[:-1]
+
+    for fname in npz_files:
+        file_path = os.path.join(dir, fname)
+        data = np.load(file_path)
+        global_scores.append(-1* data['global_score'][0])
+        data.close()
+
+
+    return global_scores
+
+
+#get melted correlations matrix of structure input x target antigen, for a given antibody 
+def get_corr(ab_name, files, dms_df, ag_columns,):
+
+    #retreive correlations from inverse fold scores
+    for key, filepath in files.items():    
+        column_label = key
+        if key == 'ProteinMPNN':
+            scores = combine_mpnn_scores(filepath)
+            dms_df[column_label] = scores
+        elif key == 'ESM-IF1':
+            df = pd.read_csv(files['ESM-IF1'])  # Read the CSV file into a DataFrame
+            log_likelihood_column = df['log_likelihood']  # Extract the 'log_likelihood' column
+            dms_df[column_label] = log_likelihood_column
+
+
+    conditions = list(files.keys())
+
+    correlations = dms_df[ag_columns + conditions].corr(method='spearman')
+    correlations = correlations.drop(conditions, axis= 1)
+    correlations = correlations.drop(ag_columns, axis= 0)
+
+    # Melt the correlations DataFrame and rename the index column
+    melted_correlations = (
+        correlations
+        .reset_index()
+        .melt(id_vars='index', var_name='Target', value_name='Correlation')
+        .rename(columns={'index': 'Input'})
+    )
+
+    print(melted_correlations)
+    return melted_correlations
+    
+    return all_ag_melted
+    
+#get melted correlations matrix of structure input x target antigen, for g6 antibody 
+def get_g6_corr(g6Hc_dict, g6LC_dict, dropLOQ = False ):
+
+    ab_name = 'g6'
+    ag_columns = g6LC_dict['ag_columns']
+    vh_and_vl_dms_df = pd.DataFrame({})
+
+    for d in [g6HC_dict, g6LC_dict]:
+
+        files = d['files']
+        dms_df = d['dms_df']
+
+        #retreive correlations from abysis, ablang, and esm1v scores
+        for key, filepath in files.items():        
+            column_label = key
+            if key == 'ProteinMPNN':
+                scores = combine_mpnn_scores(filepath)
+                dms_df[column_label] = scores
+            elif key == 'ESM-IF1':
+                df = pd.read_csv(files['ESM-IF1'])  # Read the CSV file into a DataFrame
+                log_likelihood_column = df['log_likelihood']  # Extract the 'log_likelihood' column
+                dms_df[column_label] = log_likelihood_column
+
+
+        vh_and_vl_dms_df = pd.concat([vh_and_vl_dms_df, dms_df], ignore_index= True)
+
+    conditions = list(files.keys())
+
+    correlations = vh_and_vl_dms_df[ag_columns + conditions].corr(method='spearman')
+    correlations = correlations.drop(conditions, axis= 1)
+    correlations = correlations.drop(ag_columns, axis= 0)
+
+    # Melt the correlations DataFrame and rename the index column
+    melted_correlations = (
+        correlations
+        .reset_index()
+        .melt(id_vars='index', var_name='Target', value_name='Correlation')
+        .rename(columns={'index': 'Input'})
+    )
+    melted_correlations['Target'] = ['VEGF-A'] * len(melted_correlations)
+
+    print(melted_correlations)
+    return melted_correlations
+
+    
+#plot correlation bars for cr antibodies: cr6261  and cr9114
+def plot_hbar(title, melted_correlations, palette = bokeh.palettes.Spectral6, ):
+
+    melted_correlations['cats'] = melted_correlations.apply(lambda x: (x["Target"], x["Input"]), axis = 1)
+    factors = list(melted_correlations.cats)[::-1]
+    
+    p = bokeh.plotting.figure(
+    height=340,
+    width=440,
+    x_axis_label="Spearman Correlation",
+    x_range=[0, 1],
+    y_range=bokeh.models.FactorRange(*factors),
+    tools="save",
+    title = title
+    )
+        
+
+    p.hbar(
+        source=melted_correlations,
+        y="cats",
+        right="Correlation",
+        height=0.6,
+        line_color = 'black',
+        fill_color=bokeh.palettes.Dark2_6[0]
+             
+    )
+
+    labels_df = melted_correlations
+    labels_df['corr_str'] = labels_df['Correlation'].apply(lambda x: round(x, 2)).astype(str)
+    labels_source = bokeh.models.ColumnDataSource(labels_df)
+
+    labels = bokeh.models.LabelSet(x='Correlation', y='cats', text='corr_str',text_font_size = "10px",
+              x_offset=12, y_offset=-5, source=labels_source, render_mode='canvas')
+
+    p.ygrid.grid_line_color = None
+    p.y_range.range_padding = 0.1
+    p.add_layout(labels)
+    p.legend.visible = False
+
+    p.output_backend = "svg"
+    return p
+
+if __name__ == '__main__':
+    datasets = [cr9114_dict, cr6261_dict, (g6HC_dict, g6LC_dict) ]
+
+    all_corr_plots = []
+
+    for d in datasets:
+        if type(d) is tuple:
+            g6Hc, g6Lc = d
+            title = g6Hc['ab_name'] + ', ' + g6Hc['expt_type']
+            g6_combined = get_g6_corr(g6Hc, g6Lc)
+        
+            all_corr_plots.append(plot_hbar(title, g6_combined, g6Lc['palette']))
+
+        else:
+            corr_df = get_corr(d['ab_name'], d['files'], d['dms_df'], d['ag_columns'])
+            title = d['ab_name'] + ', ' + d['expt_type']
+            all_corr_plots.append(plot_hbar(title, corr_df, d['palette']))
+
+
+    all_corr_fname = f"output/ab_mutagenesis_expts/mpnn_benchmarks.html"
+    bokeh.plotting.output_file(all_corr_fname)
+    bokeh.io.show(bokeh.layouts.gridplot(all_corr_plots, ncols = len(all_corr_plots)))

model/score_log_likelihoods.py

@@ -0,0 +1,162 @@
+import argparse
+from biotite.sequence.io.fasta import FastaFile, get_sequences
+import numpy as np
+from pathlib import Path
+import torch
+import torch.nn.functional as F
+from tqdm import tqdm
+import warnings
+import os
+from dms_utils import deep_mutational_scan
+import esm
+import pandas as pd
+from multichain_util import extract_coords_from_complex, _concatenate_coords, _concatenate_seqs, score_sequence_in_complex
+from util import get_sequence_loss, load_structure, load_coords, score_sequence, extract_coords_from_structure
+
+def get_native_seq(pdbfile, chain):
+        structure = load_structure(pdbfile, chain)
+        _ , native_seq = extract_coords_from_structure(structure)
+        return native_seq
+
+def score_singlechain_backbone(model, alphabet, args):
+    if torch.cuda.is_available() and not args.nogpu:
+        model = model.cuda()
+        print("Transferred model to GPU")
+
+    coords, native_seq = load_coords(args.pdbfile, args.chain)
+    print('Native sequence loaded from structure file:')
+    print(native_seq)
+    print('\n')
+    ll, _ = score_sequence(
+            model, alphabet, coords, native_seq)
+    print('Native sequence')
+    print(f'Log likelihood: {ll:.2f}')
+    print(f'Perplexity: {np.exp(-ll):.2f}')
+    print('\nScoring variant sequences from sequence file..\n')
+    infile = FastaFile()
+    infile.read(args.seqpath)
+    seqs = get_sequences(infile)
+    Path(args.outpath).parent.mkdir(parents=True, exist_ok=True)
+    with open(args.outpath, 'w') as fout:
+        fout.write('seqid,log_likelihood\n')
+        for header, seq in tqdm(seqs.items()):
+            ll, _ = score_sequence(
+                    model, alphabet, coords, str(seq))
+            fout.write(header + ',' + str(ll) + '\n')
+    print(f'Results saved to {args.outpath}') 
+
+
+def score_multichain_backbone(model, alphabet, args):
+    if torch.cuda.is_available() and not args.nogpu:
+        model = model.cuda()
+        print("Transferred model to GPU")
+
+    structure = load_structure(args.pdbfile)
+    coords, native_seqs = extract_coords_from_complex(structure)
+    target_chain_id = args.chain
+    native_seq = native_seqs[target_chain_id]
+    order = args.order
+
+    print('Native sequence loaded from structure file:')
+    print(native_seq)
+    print('\n')
+
+    ll_complex, ll_targetchain = score_sequence_in_complex(
+        model,
+        alphabet,
+        coords,
+        native_seqs,
+        target_chain_id,
+        native_seq,
+        order=order,
+    )
+    print('Native sequence')
+    print(f'Log likelihood of complex: {ll_complex:.2f}')
+    print(f'Log likelihood of target chain: {ll_targetchain:.2f}')
+    print(f'Perplexity: {np.exp(ll_complex):.2f}')
+
+    print('\nScoring variant sequences from sequence file..\n')
+    infile = FastaFile()
+    infile.read(args.seqpath)
+    seqs = get_sequences(infile)
+    Path(args.outpath).parent.mkdir(parents=True, exist_ok=True)
+    with open(args.outpath, 'w') as fout:
+        fout.write('seqid,log_likelihood, log_likelihood_target\n')
+        for header, seq in tqdm(seqs.items()):
+            ll_complex, ll_targetchain = score_sequence_in_complex(
+                model,
+                alphabet,
+                coords,
+                native_seqs,
+                target_chain_id,
+                str(seq),
+                order=order,
+            )
+            fout.write(header + ',' + str(ll_complex) + ',' + str(ll_targetchain) + '\n')
+    print(f'Results saved to {args.outpath}') 
+
+def get_model_checkpoint_path(filename):
+    # Expanding the user's home directory
+    return os.path.expanduser(f"~/.cache/torch/hub/checkpoints/{filename}")
+
+def main():
+    parser = argparse.ArgumentParser(
+            description='Score sequences based on a given structure.'
+    )
+    parser.add_argument(
+            'pdbfile', type=str,
+            help='input filepath, either .pdb or .cif',
+    )
+    parser.add_argument(
+            '--seqpath', type=str,
+            help='input filepath for variant sequences in a .fasta file',
+    )
+    parser.add_argument(
+            '--outpath', type=str,
+            help='output filepath for scores of variant sequences',
+    )
+    parser.add_argument(
+            '--chain', type=str,
+            help='chain id for the chain of interest', default='A',
+    )
+    parser.set_defaults(multichain_backbone=True)
+    parser.add_argument(
+            '--multichain-backbone', action='store_true',
+            help='use the backbones of all chains in the input for conditioning'
+    )
+    parser.add_argument(
+            '--order', type=str, default=None,
+            help='for multichain, specify order'
+    )
+    parser.add_argument(
+            '--singlechain-backbone', dest='multichain_backbone',
+            action='store_false',
+            help='use the backbone of only target chain in the input for conditioning'
+    )
+
+    parser.add_argument(
+        "--nogpu", action="store_true", 
+        help="Do not use GPU even if available"
+    )
+    args = parser.parse_args()
+
+    if args.outpath is None:
+        args.outpath = f'output/{args.pdbfile[:-4]}-chain{args.chain}_scores.csv'
+
+    model_checkpoint_path = get_model_checkpoint_path('esm_if1_20220410.pt')
+    with warnings.catch_warnings():
+            warnings.simplefilter('ignore', UserWarning)
+            model, alphabet = esm.pretrained.load_model_and_alphabet( \
+                model_checkpoint_path \
+            )
+    model = model.eval()
+
+
+    if args.multichain_backbone:
+        score_multichain_backbone(model, alphabet, args)
+    else:
+        score_singlechain_backbone(model, alphabet, args)
+
+
+if __name__ == '__main__':
+    main()

model/util.py

@@ -0,0 +1,316 @@
+import json
+import math
+import biotite.structure
+from biotite.structure.io import pdbx, pdb
+from biotite.structure.residues import get_residues
+from biotite.structure import filter_backbone
+from biotite.structure import get_chains
+from biotite.sequence import ProteinSequence
+import numpy as np
+from scipy.spatial import transform
+from scipy.stats import special_ortho_group
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.utils.data as data
+from typing import Sequence, Tuple, List
+from esm.data import BatchConverter
+
+
+def load_structure(fpath, chain=None):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id or list of chain ids to load
+    Returns:
+        biotite.structure.AtomArray
+    """
+    if fpath.endswith('cif'):
+        with open(fpath) as fin:
+            pdbxf = pdbx.PDBxFile.read(fin)
+        structure = pdbx.get_structure(pdbxf, model=1)
+    elif fpath.endswith('pdb'):
+        with open(fpath) as fin:
+            pdbf = pdb.PDBFile.read(fin)
+        structure = pdb.get_structure(pdbf, model=1)
+    bbmask = filter_backbone(structure)
+    structure = structure[bbmask]
+    all_chains = get_chains(structure)
+    if len(all_chains) == 0:
+        raise ValueError('No chains found in the input file.')
+    if chain is None:
+        chain_ids = all_chains
+    elif isinstance(chain, list):
+        chain_ids = chain
+    else:
+        chain_ids = [chain] 
+    for chain in chain_ids:
+        if chain not in all_chains:
+            raise ValueError(f'Chain {chain} not found in input file')
+    chain_filter = [a.chain_id in chain_ids for a in structure]
+    structure = structure[chain_filter]
+    return structure
+
+
+def extract_coords_from_structure(structure: biotite.structure.AtomArray):
+    """
+    Args:
+        structure: An instance of biotite AtomArray
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    coords = get_atom_coords_residuewise(["N", "CA", "C"], structure)
+    residue_identities = get_residues(structure)[1]
+    seq = ''.join([ProteinSequence.convert_letter_3to1(r) for r in residue_identities])
+    return coords, seq
+
+
+def load_coords(fpath, chain):
+    """
+    Args:
+        fpath: filepath to either pdb or cif file
+        chain: the chain id
+    Returns:
+        Tuple (coords, seq)
+            - coords is an L x 3 x 3 array for N, CA, C coordinates
+            - seq is the extracted sequence
+    """
+    structure = load_structure(fpath, chain)
+    return extract_coords_from_structure(structure)
+
+
+def get_atom_coords_residuewise(atoms: List[str], struct: biotite.structure.AtomArray):
+    """
+    Example for atoms argument: ["N", "CA", "C"]
+    """
+    def filterfn(s, axis=None):
+        filters = np.stack([s.atom_name == name for name in atoms], axis=1)
+        sum = filters.sum(0)
+        if not np.all(sum <= np.ones(filters.shape[1])):
+            raise RuntimeError("structure has multiple atoms with same name")
+        index = filters.argmax(0)
+        coords = s[index].coord
+        coords[sum == 0] = float("nan")
+        return coords
+
+    return biotite.structure.apply_residue_wise(struct, struct, filterfn)
+
+#* now consistent with updated gpu
+def get_sequence_loss(model, alphabet, coords, seq):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, seq)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+
+    prev_output_tokens = tokens[:, :-1].to(device)
+    target = tokens[:, 1:]
+    target_padding_mask = (target == alphabet.padding_idx)
+    logits, _ = model.forward(coords, padding_mask, confidence, prev_output_tokens)
+    loss = F.cross_entropy(logits, target, reduction='none')
+    loss = loss[0].cpu().detach().numpy()
+    target_padding_mask = target_padding_mask[0].cpu().numpy()
+    return loss, target_padding_mask
+
+
+def score_sequence(model, alphabet, coords, seq):
+    loss, target_padding_mask = get_sequence_loss(model, alphabet, coords, seq)
+    ll_fullseq = -np.sum(loss * ~target_padding_mask) / np.sum(~target_padding_mask)
+    # Also calculate average when excluding masked portions
+    coord_mask = np.all(np.isfinite(coords), axis=(-1, -2))
+    ll_withcoord = -np.sum(loss * coord_mask) / np.sum(coord_mask)
+    return ll_fullseq, ll_withcoord
+
+
+def get_encoder_output(model, alphabet, coords):
+    device = next(model.parameters()).device
+    batch_converter = CoordBatchConverter(alphabet)
+    batch = [(coords, None, None)]
+    coords, confidence, strs, tokens, padding_mask = batch_converter(
+        batch, device=device)
+    encoder_out = model.encoder.forward(coords, padding_mask, confidence,
+            return_all_hiddens=False)
+    # remove beginning and end (bos and eos tokens)
+    return encoder_out['encoder_out'][0][1:-1, 0]
+
+
+def rotate(v, R):
+    """
+    Rotates a vector by a rotation matrix.
+    
+    Args:
+        v: 3D vector, tensor of shape (length x batch_size x channels x 3)
+        R: rotation matrix, tensor of shape (length x batch_size x 3 x 3)
+
+    Returns:
+        Rotated version of v by rotation matrix R.
+    """
+    R = R.unsqueeze(-3)
+    v = v.unsqueeze(-1)
+    return torch.sum(v * R, dim=-2)
+
+
+def get_rotation_frames(coords):
+    """
+    Returns a local rotation frame defined by N, CA, C positions.
+
+    Args:
+        coords: coordinates, tensor of shape (batch_size x length x 3 x 3)
+        where the third dimension is in order of N, CA, C
+
+    Returns:
+        Local relative rotation frames in shape (batch_size x length x 3 x 3)
+    """
+    v1 = coords[:, :, 2] - coords[:, :, 1]
+    v2 = coords[:, :, 0] - coords[:, :, 1]
+    e1 = normalize(v1, dim=-1)
+    u2 = v2 - e1 * torch.sum(e1 * v2, dim=-1, keepdim=True)
+    e2 = normalize(u2, dim=-1)
+    e3 = torch.cross(e1, e2, dim=-1)
+    R = torch.stack([e1, e2, e3], dim=-2)
+    return R
+
+
+def nan_to_num(ts, val=0.0):
+    """
+    Replaces nans in tensor with a fixed value.    
+    """
+    val = torch.tensor(val, dtype=ts.dtype, device=ts.device)
+    return torch.where(~torch.isfinite(ts), val, ts)
+
+
+def rbf(values, v_min, v_max, n_bins=16):
+    """
+    Returns RBF encodings in a new dimension at the end.
+    """
+    rbf_centers = torch.linspace(v_min, v_max, n_bins, device=values.device)
+    rbf_centers = rbf_centers.view([1] * len(values.shape) + [-1])
+    rbf_std = (v_max - v_min) / n_bins
+    v_expand = torch.unsqueeze(values, -1)
+    z = (values.unsqueeze(-1) - rbf_centers) / rbf_std
+    return torch.exp(-z ** 2)
+
+
+def norm(tensor, dim, eps=1e-8, keepdim=False):
+    """
+    Returns L2 norm along a dimension.
+    """
+    return torch.sqrt(
+            torch.sum(torch.square(tensor), dim=dim, keepdim=keepdim) + eps)
+
+
+def normalize(tensor, dim=-1):
+    """
+    Normalizes a tensor along a dimension after removing nans.
+    """
+    return nan_to_num(
+        torch.div(tensor, norm(tensor, dim=dim, keepdim=True))
+    )
+
+#*
+class CoordBatchConverter(BatchConverter):
+    def __call__(self, raw_batch: Sequence[Tuple[Sequence, str]], device=None):
+        """
+        Args:
+            raw_batch: List of tuples (coords, confidence, seq)
+            In each tuple,
+                coords: list of floats, shape L x 3 x 3
+                confidence: list of floats, shape L; or scalar float; or None
+                seq: string of length L
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        self.alphabet.cls_idx = self.alphabet.get_idx("<cath>") 
+        batch = []
+        for coords, confidence, seq in raw_batch:
+            if confidence is None:
+                confidence = 1.
+            if isinstance(confidence, float) or isinstance(confidence, int):
+                confidence = [float(confidence)] * len(coords)
+            if seq is None:
+                seq = 'X' * len(coords)
+            batch.append(((coords, confidence), seq))
+
+        coords_and_confidence, strs, tokens = super().__call__(batch)
+
+        # pad beginning and end of each protein due to legacy reasons
+        coords = [
+            F.pad(torch.tensor(cd), (0, 0, 0, 0, 1, 1), value=np.nan) #<---pad set to nan from np.inf
+            for cd, _ in coords_and_confidence
+        ]
+        confidence = [
+            F.pad(torch.tensor(cf), (1, 1), value=-1.)
+            for _, cf in coords_and_confidence
+        ]
+        coords = self.collate_dense_tensors(coords, pad_v=np.nan)
+        confidence = self.collate_dense_tensors(confidence, pad_v=-1.)
+        if device is not None:
+            coords = coords.to(device)
+            confidence = confidence.to(device)
+            tokens = tokens.to(device)
+        padding_mask = torch.isnan(coords[:,:,0,0])
+        coord_mask = torch.isfinite(coords.sum(-2).sum(-1))
+        confidence = confidence * coord_mask + (-1.) * padding_mask
+        return coords, confidence, strs, tokens, padding_mask
+
+    def from_lists(self, coords_list, confidence_list=None, seq_list=None, device=None):
+        """
+        Args:
+            coords_list: list of length batch_size, each item is a list of
+            floats in shape L x 3 x 3 to describe a backbone
+            confidence_list: one of
+                - None, default to highest confidence
+                - list of length batch_size, each item is a scalar
+                - list of length batch_size, each item is a list of floats of
+                    length L to describe the confidence scores for the backbone
+                    with values between 0. and 1.
+            seq_list: either None or a list of strings
+        Returns:
+            coords: Tensor of shape batch_size x L x 3 x 3
+            confidence: Tensor of shape batch_size x L
+            strs: list of strings
+            tokens: LongTensor of shape batch_size x L
+            padding_mask: ByteTensor of shape batch_size x L
+        """
+        batch_size = len(coords_list)
+        if confidence_list is None:
+            confidence_list = [None] * batch_size
+        if seq_list is None:
+            seq_list = [None] * batch_size
+        raw_batch = zip(coords_list, confidence_list, seq_list)
+        return self.__call__(raw_batch, device)
+
+    @staticmethod
+    def collate_dense_tensors(samples, pad_v):
+        """
+        Takes a list of tensors with the following dimensions:
+            [(d_11,       ...,           d_1K),
+             (d_21,       ...,           d_2K),
+             ...,
+             (d_N1,       ...,           d_NK)]
+        and stack + pads them into a single tensor of:
+        (N, max_i=1,N { d_i1 }, ..., max_i=1,N {diK})
+        """
+        if len(samples) == 0:
+            return torch.Tensor()
+        if len(set(x.dim() for x in samples)) != 1:
+            raise RuntimeError(
+                f"Samples has varying dimensions: {[x.dim() for x in samples]}"
+            )
+        (device,) = tuple(set(x.device for x in samples))  # assumes all on same device
+        max_shape = [max(lst) for lst in zip(*[x.shape for x in samples])]
+        result = torch.empty(
+            len(samples), *max_shape, dtype=samples[0].dtype, device=device
+        )
+        result.fill_(pad_v)
+        for i in range(len(samples)):
+            result_i = result[i]
+            t = samples[i]
+            result_i[tuple(slice(0, k) for k in t.shape)] = t
+        return result

scripts/esm1v_ab_benchmarking.sh

@@ -0,0 +1,84 @@
+#!/bin/bash
+
+run_prediction() {
+    python ../esm/examples/variant-prediction/predict.py \
+        --model-location esm1v_t33_650M_UR90S_1 esm1v_t33_650M_UR90S_2 esm1v_t33_650M_UR90S_3 esm1v_t33_650M_UR90S_4 esm1v_t33_650M_UR90S_5 \
+        --sequence "$1" \
+        --dms-input "$2" \
+        --mutation-col mutant \
+        --dms-output "$3" \
+        --scoring-strategy 'masked-marginals' \
+        --offset-idx "$4"
+}
+
+# CR6261
+run_prediction "EVQLVESGAEVKKPGSSVKVSCKASGGPFRSYAISWVRQAPGQGPEWMGGIIPIFGTTKYAPKFQGRVTITADDFAGTVYMELSSLRSEDTAMYYCAKHMGYQVRETMDVWGKGTTVTVSS" \
+    "data/ab_mutagenesis_expts/cr6261/cr6261_singleMuts_exp_data.csv" \
+    "output/ab_mutagenesis_expts/cr6261/cr6261_exp_data_maskMargLabeled.csv" \
+    1
+
+# CR9114
+run_prediction "QVQLVQSGAEVKKPGSSVKVSCKSSGGTSNNYAISWVRQAPGQGLDWMGGISPIFGSTAYAQKFQGRVTISADIFSNTAYMELNSLTSEDTAVYFCARHGNYYYYSGMDVWGQGTTVTVSS" \
+    "data/ab_mutagenesis_expts/cr9114/cr9114_singleMuts_exp_data.csv" \
+    "output/ab_mutagenesis_expts/cr9114/cr9114_exp_data_maskMargLabeled.csv" \
+    1
+
+# G6 HC
+run_prediction "EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTV" \
+    "data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv" \
+    "output/ab_mutagenesis_expts/g6/g6Hc_exp_data_maskMargLabeled.csv" \
+    1
+
+# G6 LC
+run_prediction "DIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIK" \
+    "data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv" \
+    "output/ab_mutagenesis_expts/g6/g6Lc_exp_data_maskMargLabeled.csv" \
+    1
+
+# G6 HC (both chains)
+run_prediction "EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTVDIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIK" \
+    "data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv" \
+    "output/ab_mutagenesis_expts/g6/g6Hc_esm1vbothchains_exp_data_maskMargLabeled.csv" \
+    1
+
+# G6 LC (both chains)
+run_prediction "EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTVDIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIK" \
+    "data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv" \
+    "output/ab_mutagenesis_expts/g6/g6Lc_esm1vbothchains_exp_data_maskMargLabeled.csv" \
+    -117
+
+# CR9114 (both chains)
+run_prediction "QVQLVQSGAEVKKPGSSVKVSCKSSGGTSNNYAISWVRQAPGQGLDWMGGISPIFGSTAYAQKFQGRVTISADIFSNTAYMELNSLTSEDTAVYFCARHGNYYYYSGMDVWGQGTTVTVSSQSALTQPPAVSGTPGQRVTISCSGSDSNIGRRSVNWYQQFPGTAPKLLIYSNDQRPSVVPDRFSGSKSGTSASLAISGLQSEDEAEYYCAAWDDSLKGAVFGGGTQLTVL" \
+    "data/ab_mutagenesis_expts/cr9114/cr9114_singleMuts_exp_data.csv" \
+    "output/ab_mutagenesis_expts/cr9114/cr9114_esm1vbothchains_exp_data_maskMargLabeled.csv" \
+    1
+
+# CR6261 (both chains)
+run_prediction "EVQLVESGAEVKKPGSSVKVSCKASGGPFRSYAISWVRQAPGQGPEWMGGIIPIFGTTKYAPKFQGRVTITADDFAGTVYMELSSLRSEDTAMYYCAKHMGYQVRETMDVWGKGTTVTVSSQSVLTQPPSVSAAPGQKVTISCSGSSSNIGNDYVSWYQQLPGTAPKLLIYDNNKRPSGIPDRFSGSKSGTSATLGITGLQTGDEANYYCATWDRRPTAYVVFGGGTKLTVL" \
+    "data/ab_mutagenesis_expts/cr6261/cr6261_singleMuts_exp_data.csv" \
+    "output/ab_mutagenesis_expts/cr6261/cr6261_esm1vbothchains_exp_data_maskMargLabeled.csv" \
+    1
+
+# CR9114 (antibody + antigen)
+run_prediction "QVQLVQSGAEVKKPGSSVKVSCKSSGGTSNNYAISWVRQAPGQGLDWMGGISPIFGSTAYAQKFQGRVTISADIFSNTAYMELNSLTSEDTAVYFCARHGNYYYYSGMDVWGQGTTVTVSSQSALTQPPAVSGTPGQRVTISCSGSDSNIGRRSVNWYQQFPGTAPKLLIYSNDQRPSVVPDRFSGSKSGTSASLAISGLQSEDEAEYYCAAWDDSLKGAVFGGGTQLTVLADPGDQICIGYHANNSTEQVDTIMEKNVTVTHAQDILEKKHNGKLCDLDGVKPLILRDCSVAGWLLGNPMCDEFINVPEWSYIVEKANPVNDLCYPGDFNDYEELKHLLSRINHFEKIQIIPKSSWSSHEASLGVSSACPYQGKSSFFRNVVWLIKKNSTYPTIKRSYNNTNQEDLLVLWGIHHPNDAAEQTKLYQNPTTYISVGTSTLNQRLVPRIATRSKVNGQSGRMEFFWTILKPNDAINFESNGNFIAPEYAYKIVKKGDSTIMKSELEYGNCNTKCQTPMGAINSSMPFHNIHPLTIGECPKYVKSNRLVLATGLRNSPQRERRRKKRGLFGAIAGFIEGGWQGMVDGWYGYHHSNEQGSGYAADKESTQKAIDGVTNKVNSIIDKMNTQFEAVGREFNNLERRIENLNKKMEDGFLDVWTYNAELLVLMENERTLDFHDSNVKNLYDKVRLQLRDNAKELGNGCFEFYHKCDNECMESVRNGTYDYPQYSEEARLKREEISSGR" \
+    "data/ab_mutagenesis_expts/cr9114/cr9114_singleMuts_exp_data.csv" \
+    "output/ab_mutagenesis_expts/cr9114/cr9114_esm1vAbAg_exp_data_maskMargLabeled.csv" \
+    1
+
+# CR6261 (antibody + antigen)
+run_prediction "EVQLVESGAEVKKPGSSVKVSCKASGGPFRSYAISWVRQAPGQGPEWMGGIIPIFGTTKYAPKFQGRVTITADDFAGTVYMELSSLRSEDTAMYYCAKHMGYQVRETMDVWGKGTTVTVSSQSVLTQPPSVSAAPGQKVTISCSGSSSNIGNDYVSWYQQLPGTAPKLLIYDNNKRPSGIPDRFSGSKSGTSATLGITGLQTGDEANYYCATWDRRPTAYVVFGGGTKLTVLADPGDTICIGYHANNSTDTVDTVLEKNVTVTHSVNLLEDSHNGKLCKLKGIAPLQLGKCNIAGWLLGNPECDLLLTASSWSYIVETSNSENGTCYPGDFIDYEELREQLSSVSSFEKFEIFPKTSSWPNHETTKGVTAACSYAGASSFYRNLLWLTKKGSSYPKLSKSYVNNKGKEVLVLWGVHHPPTGTDQQSLYQNADAYVSVGSSKYNRRFTPEIAARPKVRDQAGRMNYYWTLLEPGDTITFEATGNLIAPWYAFALNRGSGSGIITSDAPVHDCNTKCQTPHGAINSSLPFQNIHPVTIGECPKYVRSTKLRMATGLRNIPSIQSRGLFGAIAGFIEGGWTGMIDGWYGYHHQNEQGSGYAADQKSTQNAIDGITNKVNSVIEKMNTQFTAVGKEFNNLERRIENLNKKVDDGFLDIWTYNAELLVLLENERTLDFHDSNVRNLYEKVKSQLKNNAKEIGNGCFEFYHKCDDACMESVRNGTYDYPKYSEESKLNREEIDGVSGR" \
+    "data/ab_mutagenesis_expts/cr6261/cr6261_singleMuts_exp_data.csv" \
+    "output/ab_mutagenesis_expts/cr6261/cr6261_esm1vAbAg_exp_data_maskMargLabeled.csv" \
+    1
+
+# G6 HC (antibody + antigen)
+run_prediction "EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTVDIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIKGQNHHEVVKFMDVYQRSYCHPIETLVDIFQEYPDEIEYIFKPSCVPLMRCGGCCNDEGLECVPTEESNITMQIMRIKPHQGQHIGEMSFLQHNKCECRPKKD" \
+    "data/ab_mutagenesis_expts/g6/g6_hc_exp_data.csv" \
+    "output/ab_mutagenesis_expts/g6/g6Hc_esm1vAbAg_exp_data_maskMargLabeled.csv" \
+    1
+
+# G6 LC (antibody + antigen)
+run_prediction "EVQLVESGGGLVQPGGSLRLSCAASGFTISDYWIHWVRQAPGKGLEWVAGITPAGGYTYYADSVKGRFTISADTSKNTAYLQMNSLRAEDTAVYYCARFVFFLPYAMDYWGQGTLVTVDIQMTQSPSSLSASVGDRVTITCRASQDVSTAVAWYQQKPGKAPKLLIYSASFLYSGVPSRFSGSGSGTDFTLTISSLQPEDFATYYCQQSYTTPPTFGQGTKVEIKGQNHHEVVKFMDVYQRSYCHPIETLVDIFQEYPDEIEYIFKPSCVPLMRCGGCCNDEGLECVPTEESNITMQIMRIKPHQGQHIGEMSFLQHNKCECRPKKD" \
+    "data/ab_mutagenesis_expts/g6/g6_lc_exp_data.csv" \
+    "output/ab_mutagenesis_expts/g6/g6Lc_esm1vAbAg_exp_data_maskMargLabeled.csv" \
+    -117

scripts/if_ab_benchmarking.sh

@@ -0,0 +1,44 @@
+abs=("cr6261" "cr9114" "g6" "g6")
+ab_fastas=("cr6261_3gbn_hc_lib.fasta" "cr9114_4fqi_hc_lib.fasta" "g6_2fjg_hc_lib.fasta" "g6_2fjg_lc_lib.fasta")
+data_path="data/ab_mutagenesis_expts/"
+out_prefix="output/ab_mutagenesis_expts/"
+
+for ((i=0; i<${#abs[@]}; i++)); do
+    ab="${abs[i]}"
+    ab_fasta="${ab_fastas[i]}"
+    ab_dir_path="${data_path}${ab}/"
+    struc_list=("${ab_dir_path}"*.pdb)
+    ab_out_dir="${out_prefix}${ab}/"
+
+    # Set the default chain value
+    chain="H"
+    # Special handling for 'g6' antibody
+    if [[ "$ab" == "g6" ]]; then
+        [[ "$ab_fasta" == *"lc"* ]] && chain="L"
+    fi
+
+    # gather pdbs and filter the hc/lc only structure from being scored by library for the other chain
+    if [[ "$ab" == "g6" && "$chain" == "L" ]]; then
+        struc_list=($(echo "${struc_list[@]}" | tr ' ' '\n' | grep -v '_h_' | tr '\n' ' '))
+    elif [[ "$ab" == "g6" && "$chain" == "H" ]]; then
+        struc_list=($(echo "${struc_list[@]}" | tr ' ' '\n' | grep -v '_l_' | tr '\n' ' '))
+    fi
+    
+    mkdir -p "$ab_out_dir"
+
+    for struc in "${struc_list[@]}"; do
+        out_file="${ab_out_dir}${struc##*/}"
+        if [[ "$ab" == "g6" ]]; then
+            chain_modeled="$([ "$chain" == "H" ] && echo "hc" || echo "lc")"
+            out_file="${out_file%_fvar.pdb}_${chain_modeled}_scores.csv"
+        else
+            out_file="${out_file%_fvar.pdb}_scores.csv"
+        fi
+
+        if [[ ! -f "$out_file" ]]; then
+            python model/score_log_likelihoods.py "$struc" --chain "$chain" --seqpath "${ab_dir_path}${ab_fasta}" --outpath "$out_file"
+        else
+            echo "$out_file already exists. Skipping..."
+        fi
+    done
+done

scripts/mpnn_ab_benchmarking.sh

@@ -0,0 +1,39 @@
+#!/bin/bash
+
+run_protein_mpnn() {
+    python ../ProteinMPNN/protein_mpnn_run.py \
+        --path_to_fasta "$1" \
+        --score_only 1 \
+        --save_score 1 \
+        --pdb_path_chains "$2" \
+        --out_folder "$3" \
+        --pdb_path "$4"
+}
+
+# G6 Light Chain
+run_protein_mpnn \
+    "data/ab_mutagenesis_expts/g6/g6_2fjg_lc_lib.fasta" \
+    "L" \
+    "output/ab_mutagenesis_expts/g6/proteinMpnnLC" \
+    "data/ab_mutagenesis_expts/g6/2fjg_vlh_fvar.pdb"
+
+# G6 Heavy Chain
+run_protein_mpnn \
+    "data/ab_mutagenesis_expts/g6/g6_2fjg_hc_lib.fasta" \
+    "H" \
+    "output/ab_mutagenesis_expts/g6/proteinMpnnHC" \
+    "data/ab_mutagenesis_expts/g6/2fjg_vlh_fvar.pdb"
+
+# CR6261 Heavy Chain
+run_protein_mpnn \
+    "data/ab_mutagenesis_expts/cr6261/cr6261_3gbn_hc_lib.fasta" \
+    "H" \
+    "output/ab_mutagenesis_expts/cr6261/mpnn" \
+    "data/ab_mutagenesis_expts/cr6261/3gbn_ablh_fvar.pdb"
+
+# CR9114 Heavy Chain
+run_protein_mpnn \
+    "data/ab_mutagenesis_expts/cr9114/cr9114_4fqi_hc_lib.fasta" \
+    "H" \
+    "output/ab_mutagenesis_expts/cr9114/mpnn" \
+    "data/ab_mutagenesis_expts/cr9114/4fqi_ablh_fvar.pdb"

scripts/recommend.py

@@ -0,0 +1,152 @@
+import argparse
+import os
+import sys
+from pathlib import Path
+
+ROOT_DIR = Path(__file__).resolve().parents[1]
+MODEL_DIR = ROOT_DIR / "model"
+if str(MODEL_DIR) not in sys.path:
+    sys.path.insert(0, str(MODEL_DIR))
+
+
+def load_esm_model(checkpoint_path):
+    import esm
+    import warnings
+
+    with warnings.catch_warnings():
+        warnings.simplefilter('ignore', UserWarning)
+        return esm.pretrained.load_model_and_alphabet(checkpoint_path)
+
+def get_native_seq(pdbfile, chain):
+    import util
+
+    structure = util.load_structure(pdbfile, chain)
+    _ , native_seq = util.extract_coords_from_structure(structure)
+    return native_seq
+    
+def write_dms_lib(args):
+    '''Writes a deep mutational scanning library, including the native/wildtype (wt) of the 
+        indicated target chain in the structure to an output Fasta file'''
+    from dms_utils import deep_mutational_scan
+
+    sequence = get_native_seq(args.pdbfile, args.chain)
+    Path(args.seqpath).parent.mkdir(parents=True, exist_ok=True)
+    with open(args.seqpath, 'w') as f:
+        f.write('>wt\n')
+        f.write(sequence+'\n')
+        for pos, wt, mt in deep_mutational_scan(sequence):
+            assert(sequence[pos] == wt)
+            mut_seq = sequence[:pos] + mt + sequence[(pos + 1):]
+            f.write('>' + str(wt) + str(pos+1+args.offset) + str(mt) + '\n')
+            f.write(mut_seq + '\n')
+
+def get_top_n(args):
+    import pandas as pd
+
+    recs, rec_inds = [], []
+    scores_df = pd.read_csv(args.outpath).sort_values(by = 'log_likelihood', ascending = False)
+
+    for seqid in scores_df['seqid']:
+        res_ind = seqid[1:-1]
+        if (rec_inds.count(res_ind) < args.maxrep): 
+            if args.upperbound == None or (int(res_ind) < int(args.upperbound)):
+                recs.append(seqid)
+                rec_inds.append(res_ind)
+        if len(recs) == args.n:
+            break 
+
+    print(f'\n Chain {args.chain}')
+    print(*recs, sep='\n')
+
+def get_model_checkpoint_path(filename):
+    # Expanding the user's home directory
+    return os.path.expanduser(f"~/.cache/torch/hub/checkpoints/{filename}")
+        
+def main():
+    parser = argparse.ArgumentParser(
+        description='Score sequences based on a given structure.'
+    )
+    parser.add_argument(
+        'pdbfile', type=str,
+        help='input filepath, either .pdb or .cif',
+    )
+    parser.add_argument(
+        '--seqpath', type=str,
+        help='filepath where fasta of dms library should be saveda',
+    )
+    parser.add_argument(
+        '--outpath', type=str,
+        help='output filepath for scores of variant sequences',
+    )
+    parser.add_argument(
+        '--chain', type=str,
+        help='chain id for the chain of interest', default='A',
+    )
+    parser.set_defaults(multichain_backbone=True)
+    parser.add_argument(
+        '--multichain-backbone', action='store_true',
+        help='use the backbones of all chains in the input for conditioning'
+    )
+    parser.add_argument(
+        '--singlechain-backbone', dest='multichain_backbone',
+        action='store_false',
+        help='use the backbone of only target chain in the input for conditioning'
+    )
+    parser.add_argument(
+            '--order', type=str, default=None,
+            help='for multichain, option to specify order of chains'
+    )
+    parser.add_argument(
+        '--n', type=int,
+        help='number of desired predictions to be output', 
+        default=10,
+    )
+    parser.add_argument(
+        '--maxrep', type=int,
+        help='maximum representation of a single site in the top recommendations  \
+              (eg: maxrep = 1 is a unique set where no wildtype residue is mutated more than once)', 
+        default=1,
+    )
+    parser.add_argument(
+        '--offset', type=int,
+        help='integer offset for labeling of residue indices encoded in the structure',
+        default=0,
+    )
+    parser.add_argument(
+        '--upperbound', type=int,
+        help='only residue positions less than the user-defined upperbound are considered to be recommended for screening \
+                (but all positions are still conditioned for scoring)', 
+        default=None,
+    )
+    parser.add_argument(
+        "--nogpu", action="store_true", 
+        help="Do not use GPU even if available"
+    )
+
+    args = parser.parse_args()
+
+    if args.seqpath is None:
+        args.seqpath = f'output/{args.pdbfile[:-4]}-chain{args.chain}_dms.fasta'
+
+    if args.outpath is None:
+        args.outpath = f'output/{args.pdbfile[:-4]}-chain{args.chain}_scores.csv'
+
+    #write dms library for target chain 
+    write_dms_lib(args)
+
+    model_checkpoint_path = get_model_checkpoint_path('esm_if1_20220410.pt')
+    model, alphabet = load_esm_model(model_checkpoint_path)
+    model = model.eval()
+    
+    import score_log_likelihoods
+
+    if args.multichain_backbone:
+        score_log_likelihoods.score_multichain_backbone(model, alphabet, args)
+    else:
+        score_log_likelihoods.score_singlechain_backbone(model, alphabet, args)
+
+    get_top_n(args)
+
+
+if __name__ == '__main__':
+        main()