Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+weight/wheels/antiberty-0.1.3-py3-none-any.whl filter=lfs diff=lfs merge=lfs -text
+weight/wheels/igfold-0.4.0-py3-none-any.whl filter=lfs diff=lfs merge=lfs -text
+weight/wheels/setuptools-81.0.0-py3-none-any.whl filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,136 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+pip-wheel-metadata/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+.python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# Other
+.DS_Store
+.vscode/
+trained_models/
+demo_data/
+scripts/

.ms_upload_cache

@@ -0,0 +1 @@
+{"version": 3, "repo_id": "OneScience/IgFold", "files": {".gitignore|1784526611.0|1862": {"hash": "5ed5ad2f1435210a66884db6ef33f08eac7b6fdfba7cf2b596e8e17bda93db90", "size": 1862, "status": "c"}, "README.md|1787713837.0|11346": {"hash": "cc24643909ffc55c453b64627fd24187513cc13eead8bf7bbbc9588df5d2b25b", "size": 11346, "status": "c"}, "MANIFEST.in|1787713268.0|55": {"hash": "4f09d2f77ba3b962461deeb420affa542b90585d86fce4d90691b445172e4845", "size": 55, "status": "c"}, "conf/config.json|1787713764.0|150": {"hash": "94e90467e86bbccb199440403341bcc42c6a865c7027b6b4d806e8646b5aff3d", "size": 150, "status": "c"}, "configuration.json|1787713766.0|335": {"hash": "65ca5d35ec4141fcfa197add0014d91d5879bf434b72dedf2b2fc9ad720a846a", "size": 335, "status": "c"}, "model/IgFoldRunner.py|1787713781.0|4907": {"hash": "3629ab4bf7789778cb6890b73bffb89ea9a049df04452b6ec1ab8ae70e1bfb18", "size": 4907, "status": "c"}, "LICENSE.md|1784526611.0|3330": {"hash": "02cd4b17669c8294425925009ccdf6a6a6c8bef1486d580d5288ee7837d11fca", "size": 3330, "status": "c"}, "model/model/__init__.py|1784526611.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}, "model/model/IgFold.py|1784526611.0|15121": {"hash": "1510dcc0d76cd149719843e930545b784621bdfbbeb4779a3e77ba8f44cb0027", "size": 15121, "status": "c"}, "model/__init__.py|1784526611.0|126": {"hash": "a284b2be34a18aa671e4fc984bde975832133d5fb41a2eafa45fc09952d1ffe7", "size": 126, "status": "c"}, "model/model/components/GraphTransformer.py|1784526611.0|5439": {"hash": "8f1d561ef520c0f04d634200261a449c0e5ef67979365457c2088d065717c104", "size": 5439, "status": "c"}, "model/model/components/TriangleGraphTransformer.py|1784526611.0|1763": {"hash": "d0b6f107416c2269764bf11a99a5b53ed1f21e4a64b36d4f935e062e47c8fa19", "size": 1763, "status": "c"}, "model/model/components/IPABlock.py|1784526611.0|8775": {"hash": "6f224f9cf09e450ed7ac24b379ce16ea4abbdc53909591a92b460cfc6f64ca5e", "size": 8775, "status": "c"}, "model/refine/__init__.py|1784526611.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}, "model/model/components/IPATransformer.py|1784526611.0|4603": {"hash": "4d3a943699c886c84eaeb64f83f17c627da91bb1b376bee90bf08cc2eb208af1", "size": 4603, "status": "c"}, "model/training/__init__.py|1784526611.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}, "model/model/components/TriangleMultiplicativeModule.py|1784526611.0|2648": {"hash": "85769957c68d4f7131d2d3d4df00e516efcc648c7b622a890be62f2a3b5ec349", "size": 2648, "status": "c"}, "model/model/interface.py|1784526611.0|1248": {"hash": "4b7b0f1c7b2a961915a3abacff5c134371f4c575fbb935683ae24f8aceededef", "size": 1248, "status": "c"}, "model/utils/__init__.py|1784526611.0|0": {"hash": "e3b0c44298fc1c149afbf4c8996fb92427ae41e4649b934ca495991b7852b855", "size": 0, "status": "c"}, "model/model/components/__init__.py|1784526611.0|146": {"hash": "90edde297923dc2de725e0ca765f99d8fe92e26b14ef26809d5ffeba995b4122", "size": 146, "status": "c"}, "model/refine/openmm_ref.py|1784526611.0|1653": {"hash": "083d98ec88c711cb3d96be786a7693a443ca62d3b22f6d9e1ebe407ce37a5657", "size": 1653, "status": "c"}, "model/refine/pyrosetta_ref.py|1784526611.0|3222": {"hash": "8cb3ec32244fb094749a002544a4f2e19b239fed929d021c2e3304adc1d5df03", "size": 3222, "status": "c"}, "model/training/utils.py|1784526611.0|6991": {"hash": "7eb3552824f976d7d8a80a869388b23fef363727ed97a8f20f580ee531f85d44", "size": 6991, "status": "c"}, "model/utils/abnumber_.py|1784526611.0|2334": {"hash": "1837911edccc67c1b1db5b4742e223ab46dfba146687279c577ace811efba571", "size": 2334, "status": "c"}, "model/utils/ab_metrics.py|1784526611.0|1724": {"hash": "0c3fd30611103df673a40d4aaa23db07364289cc1ea6cbaa57e2ca69498b624a", "size": 1724, "status": "c"}, "model/utils/fasta.py|1784526611.0|934": {"hash": "a447498b09bf75dc603cad25a86987f439b3a080fefddddcd2da0860d6c52524", "size": 934, "status": "c"}, "model/utils/coordinates.py|1784526611.0|1755": {"hash": "735f8e452600dc948e7249d60b1f6e391f8eced00f53a680375e70ec7960f2d1", "size": 1755, "status": "c"}, "model/utils/embed.py|1784526611.0|1183": {"hash": "55f95ec50afffebcb3f029bad68895b693f7154038bcfa01d4de6ee2107fd405", "size": 1183, "status": "c"}, "model/utils/constants.py|1784526611.0|368": {"hash": "34e058fc09c6b1976b188b870a8e8302fef5ebbdceef98497afe9ff32fec6415", "size": 368, "status": "c"}, "model/utils/general.py|1784526611.0|771": {"hash": "63de9d287bd8c54b30fca2db2e7bcf6e8ffbfdfce16b5ad575ff1480d3e75319", "size": 771, "status": "c"}, "model/utils/geometry.py|1784526611.0|1688": {"hash": "c7af1aed0310414620e1683bf22e5af9e84dc02c99633236491f2bc885e421fa", "size": 1688, "status": "c"}, "model/utils/pdb.py|1784526611.0|14536": {"hash": "e7be074946baac32b5561f8f8fa872586d6ffc299b0ce2dd1dd4b6ccabea9ff2", "size": 14536, "status": "c"}, "model/utils/folding.py|1784526611.0|6376": {"hash": "060f6f8a5f5049540dde81ddc954cec766f4c9a21e5a1d69e89469d60b829579", "size": 6376, "status": "c"}, "model/utils/tensor.py|1784526611.0|1127": {"hash": "7eb084b9c878e2cefa0d7519262d534430086142862d7b0feb343c2093af6c42", "size": 1127, "status": "c"}, "model/utils/visualize.py|1784526611.0|4867": {"hash": "052c0beae7f3ad9e1c375f42d2535d3d665c507713b839d4e0abbf315876ceea", "size": 4867, "status": "c"}, "model/utils/transforms.py|1784526611.0|18138": {"hash": "04573d0b32168a6e94de13e2580f0f412b121b3d54a6f8617fca358631de7558", "size": 18138, "status": "c"}, "pyproject.toml|1784526611.0|80": {"hash": "1c557fcb2d563b509abf60fdb4132d72373474dbf6d059d752b2b780788d43ed", "size": 80, "status": "c"}, "requirements.txt|1784526611.0|206": {"hash": "a696fb4a79d4a004d8277aa222f3c074c8bea524b7692f93a03edfda5c4a8551", "size": 206, "status": "c"}, "setup.cfg|1787713267.0|375": {"hash": "d18e548322261df23f3b45867e5173f7136744d36d137c027368755f8b062b3a", "size": 375, "status": "c"}, "scripts/inference.py|1787713264.0|3051": {"hash": "986639e5231e2e7462ccbeafe160d12faf480c759646653eee6f09b5ebccd553", "size": 3051, "status": "c"}, "scripts/IgFold.ipynb|1784526611.0|7614": {"hash": "5ed28dc41b9b50edda8f1624f0999f783c7af65139721acb25968fab48531053", "size": 7614, "status": "c"}, "weight/IgFold/LICENSE.md|1650791114.0|3330": {"hash": "02cd4b17669c8294425925009ccdf6a6a6c8bef1486d580d5288ee7837d11fca", "size": 3330, "status": "c"}, "weight/wheels/setuptools-81.0.0-py3-none-any.whl|1785393298.0|1062021": {"hash": "fdd925d5c5d9f62e4b74b30d6dd7828ce236fd6ed998a08d81de62ce5a6310d6", "size": 1062021, "status": "c"}, "weight/IgFold/igfold_1.ckpt|1676739842.0|6358637": {"hash": "66c2aa5a4fd4ef485e99f661e6c605037fa6bfefc02464395836c23263948249", "size": 6358637, "status": "c"}, "weight/IgFold/igfold_2.ckpt|1676739842.0|6358637": {"hash": "3917eb78a3b8f9a84d53e3dae1132e0096b6ded0868a081a60d0603d502028bc", "size": 6358637, "status": "c"}, "weight/IgFold/igfold_3.ckpt|1676739842.0|6358637": {"hash": "d9b60988f1400c12f715b440c7be0bebab3b3be73dc39657c35c2029c3174c52", "size": 6358637, "status": "c"}, "weight/IgFold/igfold_5.ckpt|1676739842.0|6358637": {"hash": "b75865aef284973d47af61598a2fd50cbcfb02c67015b636f7ecd23ed842b2b5", "size": 6358637, "status": "c"}, "weight/wheels/igfold-0.4.0-py3-none-any.whl|1693722179.0|23355225": {"hash": "ddf73636cf86ca2ef91eab25f79bf8052a3deaeae0ee4e056c8eecd152582c9c", "size": 23355225, "status": "c"}, "weight/wheels/antiberty-0.1.3-py3-none-any.whl|1689549108.0|96631471": {"hash": "30d910992b190013871bac49cdc032e01a19339f7d2b958ab99b0eb44638352a", "size": 96631471, "status": "c"}}}

LICENSE.md

@@ -0,0 +1,24 @@
+# JHU Academic Software License Agreement
+
+This license agreement ("License") is effective automatically between you ("Licensee") and The Johns Hopkins University (“JHU”) for use of the software with which this License is distributed (“Software”) so long as Licensee complies with the following terms and conditions:
+
+The requirement to acknowledge the copyright of JHU as follows: “© 2022 The Johns Hopkins University” and copyrights of any incorporated third party software as described in the associated documentation.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the source code must retain the above copyright notice, and these terms and conditions.
+
+Neither the name of JHU nor its affiliates may be used to endorse or promote products derived from this Software without specific prior written permission from an authorized JHU representative.
+
+THIS SOFTWARE IS PROVIDED BY JHU "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL JHU BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. LICENSEE AGREES TO DEFEND, INDEMNIFY AND HOLD HARMLESS JHU FOR ANY CLAIMS ARISING FROM LICENSEE’S USE OF THE SOFTWARE TO THE FULLEST EXTENT PERMITTED BY LAW.
+
+Only a non-exclusive, nontransferable license is granted to Licensee to use the Software for non-commercial purposes. Commercial use of the Software requires a separately executed written license agreement.
+
+Licensee agrees that it will use the Software, and any modifications, improvements, or derivatives to the Software that the Licensee may create (collectively, "Improvements") solely for non-commercial purposes and shall not distribute or transfer the Software or Improvements to any third party without requiring that such third parties adhere to the terms of this License. Licensee agrees that any Improvements made by Licensee shall be subject to the same terms and conditions as the Software.
+
+Licensee acknowledges that JHU holds copyright in the Software or portions of the Software, and that the Software may incorporate third party software which may be subject to additional terms and conditions. Licensee shall comply with any additional terms and conditions appliable to such third party software.
+
+Licensee agrees that any publication of results obtained with the Software will acknowledge its use by an appropriate citation.
+
+Licensee’s rights under this License terminate automatically without notice from JHU if Licensee fail to comply with any term(s) of this License.
+
+This License shall be governed by the laws of the State of Maryland, excluding the application of its conflicts of law rules. Licensee agrees that any dispute shall be appropriate only in the state and federal courts located within the State of Maryland.
+

MANIFEST.in

@@ -0,0 +1,2 @@
+include LICENSE.md
+recursive-include model/igfold *.py

README.md

@@ -0,0 +1,459 @@
+---
+frameworks:
+- PyTorch
+language:
+- en
+license: other
+tags:
+- OneScience
+- bioscience
+- antibody-structure-prediction
+- protein-structure-prediction
+- IgFold
+- AntiBERTy
+tasks: []
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">IgFold</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+IgFold is an open-source antibody structure prediction model developed by the Gray Lab. It rapidly predicts three-dimensional antibody structures directly from amino acid sequences.
+
+The model supports paired heavy/light-chain antibodies, single-chain antibodies, and nanobodies. It can also return residue-level predicted RMSD values and multiple levels of antibody sequence representations.
+
+Paper: [Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies](https://www.nature.com/articles/s41467-023-38063-x)
+
+# Model Description
+
+IgFold uses AntiBERTy to extract antibody sequence representations and predicts antibody structures through graph Transformer layers, template feature integration, and an invariant point attention-based structure module.
+
+The model provides the following main capabilities:
+
+- Predict paired antibody structures from heavy-chain and light-chain sequences;
+- Predict single-chain antibody or nanobody structures from a single heavy-chain or light-chain sequence;
+- Output predicted RMSD values for the N, CA, C, and CB atoms of each residue;
+- Output intermediate representations from AntiBERTy, the graph Transformer, and the structure module;
+- Support template structures;
+- Support structural refinement using PyRosetta or OpenMM;
+- Support conversion of predicted structures to Chothia numbering.
+
+# Use Cases
+
+| Use Case | Description |
+| :---: | :--- |
+| Paired antibody structure prediction | Generate antibody PDB structures from heavy-chain and light-chain sequences. |
+| Single-chain antibody prediction | Predict a structure from a single heavy-chain or light-chain sequence. |
+| Nanobody structure prediction | Generate a PDB structure from a nanobody heavy-chain sequence. |
+| Prediction error analysis | Obtain residue-level predicted RMSD values, which are also written to the B-factor column of the output PDB file. |
+| Antibody representation extraction | Extract AntiBERTy, graph Transformer, and structure-module embeddings. |
+
+# Usage
+
+## Manual Installation
+
+### Hardware Requirements
+
+- The model can run on CPU or accelerator devices supported by PyTorch;
+- GPU or other compatible accelerator devices are recommended for structure prediction;
+- PyRosetta refinement mainly uses CPU resources, and runtime depends on sequence length and CPU performance;
+- The actual accelerator configuration and installation procedure depend on the locally installed PyTorch version, drivers, and runtime environment.
+
+### Download the Model Package
+
+Install the Hugging Face command-line tool and download the model repository:
+
+```bash
+python -m pip install -U huggingface_hub
+
+hf download OneScience-Group/IgFold --local-dir ./IgFold
+cd IgFold
+```
+
+
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+
+python -m pip install "onescience[bio-dcu]" \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+**GPU Environment**
+
+```bash
+# Activate Conda first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+
+python -m pip install "onescience[bio-gpu]" \
+  -i http://mirrors.onescience.ai:3141/pypi/simple/ \
+  --trusted-host mirrors.onescience.ai
+```
+
+Install the additional dependencies required by EpHod:
+
+```bash
+python -m pip install --no-deps -r requirements.txt
+```
+
+The `weight/wheels/` directory also contains offline-installable IgFold and AntiBERTy wheel packages together with their official pretrained assets:
+
+```bash
+python -m pip install --no-deps \
+  weight/wheels/antiberty-0.1.3-py3-none-any.whl \
+  weight/wheels/igfold-0.4.0-py3-none-any.whl
+```
+
+PyTorch installation requirements vary across hardware platforms.
+
+To use GPU or another accelerator device, install a PyTorch build compatible with the corresponding hardware platform and ensure that the remaining dependencies satisfy the versions declared in `requirements.txt`.
+
+### Optional Dependencies
+
+#### PyRosetta Refinement
+
+IgFold supports structural refinement using PyRosetta.
+
+Install a PyRosetta version compatible with the current Python version and operating system according to the [official PyRosetta installation instructions](https://www.pyrosetta.org/downloads).
+
+#### OpenMM Refinement
+
+If PyRosetta is not used, OpenMM and PDBFixer can be installed instead:
+
+```bash
+conda install -c conda-forge openmm==7.7.0 pdbfixer
+```
+
+#### Chothia Numbering
+
+To convert predicted structures to Chothia numbering, install AbNumber:
+
+```bash
+conda install -c bioconda abnumber
+```
+
+### Quick Inference
+
+The model package provides a directly executable inference entry point.
+
+If neither sequences nor a FASTA file are specified, the script uses the paired heavy/light-chain example from the official IgFold README and writes the predicted structure to:
+
+```text
+output/inference/antibody.pdb
+```
+
+Run:
+
+```bash
+python scripts/inference.py
+```
+
+To use a custom FASTA file:
+
+```bash
+python scripts/inference.py \
+  --fasta /path/to/antibody.fasta \
+  --output output/inference/my_antibody.pdb
+```
+
+Chain identifiers in the FASTA file should be `H` and `L`.
+
+Example:
+
+```text
+>sample:H
+EVQLVQSGPEVKKPGTSVKVSCKAS...
+>sample:L
+DVVMTQTPFSLPVSLGDQASISCR...
+```
+
+When optional features are disabled, basic inference does not require:
+
+- SAbDab PDB data;
+- PyRosetta;
+- OpenMM;
+- AbNumber.
+
+After installing the corresponding optional dependencies, structural refinement and renumbering can be enabled through options such as:
+
+```text
+--refine
+--openmm
+--renum
+```
+
+### Paired Antibody Structure Prediction
+
+Heavy-chain and light-chain sequences are provided as a dictionary using `H` and `L` as keys:
+
+```python
+from igfold import IgFoldRunner
+from igfold.refine.pyrosetta_ref import init_pyrosetta
+
+init_pyrosetta()
+
+sequences = {
+    "H": "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS",
+    "L": "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK",
+}
+
+pred_pdb = "my_antibody.pdb"
+
+igfold = IgFoldRunner()
+
+igfold.fold(
+    pred_pdb,
+    sequences=sequences,
+    do_refine=True,
+    do_renum=True,
+)
+```
+
+After successful execution, the predicted structure is saved to:
+
+```text
+my_antibody.pdb
+```
+
+### Nanobody or Single-Chain Antibody Prediction
+
+For nanobody or single-chain heavy/light-chain prediction, only one sequence is required:
+
+```python
+from igfold import IgFoldRunner
+from igfold.refine.pyrosetta_ref import init_pyrosetta
+
+init_pyrosetta()
+
+sequences = {
+    "H": "QVQLQESGGGLVQAGGSLTLSCAVSGLTFSNYAMGWFRQAPGKEREFVAAITWDGGNTYYTDSVKGRFTISRDNAKNTVFLQMNSLKPEDTAVYYCAAKLLGSSRYELALAGYDYWGQGTQVTVS",
+}
+
+pred_pdb = "my_nanobody.pdb"
+
+igfold = IgFoldRunner()
+
+igfold.fold(
+    pred_pdb,
+    sequences=sequences,
+    do_refine=True,
+    do_renum=True,
+)
+```
+
+### Inference Without Structural Refinement
+
+If PyRosetta or OpenMM refinement is not required, set:
+
+```python
+do_refine=False
+```
+
+If Chothia renumbering is also unnecessary, set:
+
+```python
+do_renum=False
+```
+
+Example:
+
+```python
+from igfold import IgFoldRunner
+
+sequences = {
+    "H": "QVQLQESGGGLVQAGGSLTLSCAVSGLTFSNYAMGWFRQAPGKEREFVAAITWDGGNTYYTDSVKGRFTISRDNAKNTVFLQMNSLKPEDTAVYYCAAKLLGSSRYELALAGYDYWGQGTQVTVS",
+}
+
+pred_pdb = "my_nanobody.pdb"
+
+igfold = IgFoldRunner()
+
+igfold.fold(
+    pred_pdb,
+    sequences=sequences,
+    do_refine=False,
+    do_renum=False,
+)
+```
+
+In this configuration, PyRosetta, OpenMM, and AbNumber are not required.
+
+### Predicted RMSD
+
+IgFold predicts residue-level RMSD values and stores them in the B-factor column of the output PDB file.
+
+The same values are also returned by `fold()`:
+
+```python
+from igfold import IgFoldRunner
+
+sequences = {
+    "H": "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS",
+    "L": "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK",
+}
+
+pred_pdb = "my_antibody.pdb"
+
+igfold = IgFoldRunner()
+
+out = igfold.fold(
+    pred_pdb,
+    sequences=sequences,
+    do_refine=False,
+    do_renum=False,
+)
+
+print(out.prmsd)
+# Predicted RMSD for the N, CA, C, and CB atoms of each residue.
+# Shape: [1, L, 4]
+```
+
+`prmsd` is a model-predicted error estimate.
+
+It is not the actual RMSD obtained by aligning the predicted structure with an experimentally determined structure.
+
+### Antibody Sequence Embeddings
+
+The `embed()` method provides antibody representations from multiple stages of the IgFold pipeline:
+
+```python
+from igfold import IgFoldRunner
+
+sequences = {
+    "H": "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS",
+    "L": "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK",
+}
+
+igfold = IgFoldRunner()
+emb = igfold.embed(sequences=sequences)
+
+print(emb.bert_embs.shape)
+# AntiBERTy final-layer representations: [1, L, 512]
+
+print(emb.gt_embs.shape)
+# Graph Transformer representations: [1, L, 64]
+
+print(emb.structure_embs.shape)
+# Structure-module representations: [1, L, 64]
+```
+
+### Prefer OpenMM Refinement
+
+After installing OpenMM and PDBFixer, OpenMM refinement can be selected by setting:
+
+```python
+use_openmm=True
+```
+
+Example:
+
+```python
+from igfold import IgFoldRunner
+
+sequences = {
+    "H": "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS",
+    "L": "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK",
+}
+
+pred_pdb = "my_antibody.pdb"
+
+igfold = IgFoldRunner()
+
+igfold.fold(
+    pred_pdb,
+    sequences=sequences,
+    do_refine=True,
+    use_openmm=True,
+    do_renum=True,
+)
+```
+
+### Training
+
+The official IgFold repository does not provide a complete directly executable training entry point, dataset class, or end-to-end training script.
+
+Therefore, this Hugging Face model package does not provide a training command.
+
+The `model/training/` directory contains utility functions used for model loss computation but does not constitute a complete training program.
+
+The official structural training data used in the IgFold paper is available from Zenodo:
+
+https://doi.org/10.5281/zenodo.7820263
+
+The dataset includes:
+
+- Experimentally determined SAbDab antibody structures;
+- Predicted structures generated from paired OAS sequences;
+- Predicted structures generated from unpaired OAS sequences.
+
+These datasets are useful for reproducing the training methodology described in the paper or for structural evaluation.
+
+They are not required for sequence-to-structure inference.
+
+Because this model package does not provide a complete training entry point, the full SAbDab `.fasta` and `.pdb` datasets do not need to be included in the Hugging Face model repository.
+
+To independently reproduce IgFold training, users would need to implement or reconstruct the following components using the methodology described in the paper, upstream loss functions, and official training data:
+
+- Dataset;
+- DataLoader;
+- Optimizer;
+- Training loop;
+- Validation procedure;
+- Checkpoint management.
+
+### Predicted Antibody Structure Datasets
+
+The IgFold authors also released two large collections of predicted antibody structures:
+
+- 104K non-redundant paired antibody structures from OAS:  
+  https://data.graylab.jhu.edu/OAS_paired.tar.gz
+- 1.3M predicted human paired antibody structures from the Jaffe et al. dataset:  
+  https://data.graylab.jhu.edu/Jaffe2022.tar.gz
+
+# OneScience Official Resources
+
+| Platform | OneScience Main 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
+
+IgFold source code, pretrained models, and associated materials are distributed under the [JHU Academic Software License Agreement](https://github.com/Graylab/IgFold/blob/main/LICENSE.md).
+
+The license permits use under the academic and non-commercial terms specified by Johns Hopkins University.
+
+Commercial use may require a separate license obtained through Johns Hopkins Technology Ventures.
+
+This Hugging Face model package does not modify or extend the original licensing terms of IgFold, AntiBERTy, PyRosetta, SAbDab, OAS, pretrained models, datasets, or other third-party resources.
+
+```bibtex
+@article{ruffolo2023fast,
+  title={Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies},
+  author={Ruffolo, Jeffrey A and Chu, Lee-Shin and Mahajan, Sai Pooja and Gray, Jeffrey J},
+  journal={Nature Communications},
+  volume={14},
+  number={1},
+  pages={2389},
+  year={2023},
+  publisher={Nature Publishing Group UK London}
+}
+
+@article{ruffolo2021deciphering,
+  title={Deciphering antibody affinity maturation with language models and weakly supervised learning},
+  author={Ruffolo, Jeffrey A and Gray, Jeffrey J and Sulam, Jeremias},
+  journal={arXiv},
+  year={2021}
+}
+```

conf/config.json

@@ -0,0 +1,8 @@
+{
+  "num_models": 4,
+  "output": "output/inference/antibody.pdb",
+  "refine": false,
+  "use_openmm": false,
+  "renumber": false,
+  "device": "auto"
+}

configuration.json

@@ -0,0 +1,13 @@
+{
+  "model_name": "IgFold",
+  "framework": "PyTorch",
+  "task": "antibody-structure-prediction",
+  "entry_points": {
+    "inference": "scripts/inference.py",
+    "predict": "scripts/inference.py"
+  },
+  "source_package": "model",
+  "config": "conf/config.json",
+  "weight_dir": "weight",
+  "license": "JHU-Academic-Software-License"
+}

model/IgFoldRunner.py

@@ -0,0 +1,158 @@
+import os
+from time import time
+from glob import glob
+import torch
+
+from antiberty import AntiBERTyRunner
+
+import igfold
+from igfold.model.IgFold import IgFold
+from igfold.utils.folding import fold
+from igfold.utils.embed import embed
+from igfold.utils.general import exists
+
+
+def display_license():
+    license_url = "https://github.com/Graylab/IgFold/blob/main/LICENSE.md"
+    license_message = f"""
+    The code, data, and weights for this work are made available for non-commercial use 
+    (including at commercial entities) under the terms of the JHU Academic Software License 
+    Agreement. For commercial inquiries, please contact awichma2[at]jhu.edu.
+    License: {license_url}
+    """
+    print(license_message)
+
+
+class IgFoldRunner():
+    """
+    Wrapper for IgFold model predictions.
+    """
+    def __init__(self, num_models=4, model_ckpts=None, try_gpu=True):
+        """
+        Initialize IgFoldRunner.
+
+        :param num_models: Number of pre-trained IgFold models to use for prediction.
+        :param model_ckpts: List of model checkpoints to use (instead of pre-trained).
+        """
+
+        display_license()
+
+        if exists(model_ckpts):
+            num_models = len(model_ckpts)
+        else:
+            if num_models < 1 or num_models > 4:
+                raise ValueError("num_models must be between 1 and 4.")
+
+            if not exists(model_ckpts):
+                project_path = os.path.dirname(
+                    os.path.realpath(igfold.__file__))
+
+                ckpt_path = os.path.join(
+                    project_path,
+                    "trained_models/IgFold/*.ckpt",
+                )
+                model_ckpts = list(glob(ckpt_path))
+
+            model_ckpts = list(sorted(model_ckpts))[:num_models]
+
+        print(f"Loading {num_models} IgFold models...")
+
+        device = torch.device(
+            "cuda:0" if torch.cuda.is_available() and try_gpu else "cpu")
+        print(f"Using device: {device}")
+
+        self.models = []
+        for ckpt_file in model_ckpts:
+            print(f"Loading {ckpt_file}...")
+            self.models.append(
+                IgFold.load_from_checkpoint(ckpt_file).eval().to(device))
+
+        print(f"Successfully loaded {num_models} IgFold models.")
+
+        self.antiberty = AntiBERTyRunner()
+        self.antiberty.model.eval()
+        self.antiberty.model.to(device)
+        print("Loaded AntiBERTy model.")
+
+    def fold(
+        self,
+        pdb_file,
+        fasta_file=None,
+        sequences=None,
+        template_pdb=None,
+        ignore_cdrs=None,
+        ignore_chain=None,
+        skip_pdb=False,
+        do_refine=True,
+        use_openmm=False,
+        do_renum=True,
+        truncate_sequences=False,
+    ):
+        """
+        Predict antibody structure with IgFold.
+
+        :param pdb_file: PDB file to predict.
+        :param fasta_file: FASTA file containing sequences.
+        :param sequences: Dictionary of sequences.
+        :param template_pdb: PDB file containing template structure.
+        :param ignore_cdrs: List of CDRs to ignore.
+        :param ignore_chain: Chain to ignore.
+        :param skip_pdb: Skip PDB processing.
+        :param do_refine: Perform PyRosetta refinement.
+        :param do_renum: Renumber PDB to Chothia with AbNum.
+        :param truncate_sequences: Truncate sequences with AbNumber.
+        """
+        start_time = time()
+        model_out = fold(
+            self.antiberty,
+            self.models,
+            pdb_file=pdb_file,
+            fasta_file=fasta_file,
+            sequences=sequences,
+            template_pdb=template_pdb,
+            ignore_cdrs=ignore_cdrs,
+            ignore_chain=ignore_chain,
+            skip_pdb=skip_pdb,
+            do_refine=do_refine,
+            use_openmm=use_openmm,
+            do_renum=do_renum,
+            truncate_sequences=truncate_sequences,
+        )
+
+        print(f"Completed folding in {time() - start_time:.2f} seconds.")
+
+        return model_out
+
+    def embed(
+        self,
+        model_idx=0,
+        fasta_file=None,
+        sequences=None,
+        template_pdb=None,
+        ignore_cdrs=None,
+        ignore_chain=None,
+    ):
+        """
+        Embed antibody sequences with IgFold.
+
+        :param fasta_file: FASTA file containing sequences.
+        :param sequences: Dictionary of sequences.
+        :param template_pdb: PDB file containing template structure.
+        :param ignore_cdrs: List of CDRs to ignore.
+        :param ignore_chain: Chain to ignore.
+        """
+
+        start_time = time()
+        model_out = embed(
+            self.antiberty,
+            self.models[model_idx],
+            fasta_file=fasta_file,
+            sequences=sequences,
+            template_pdb=template_pdb,
+            ignore_cdrs=ignore_cdrs,
+            ignore_chain=ignore_chain,
+        )
+
+        print(f"Completed embedding in {time() - start_time:.2f} seconds.")
+
+        return model_out

model/__init__.py

@@ -0,0 +1,3 @@
+from .model.IgFold import IgFold
+from .model.interface import IgFoldInput, IgFoldOutput
+from .IgFoldRunner import IgFoldRunner

model/model/IgFold.py

@@ -0,0 +1,488 @@
+import os
+from einops import rearrange, repeat
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import pytorch_lightning as pl
+
+from igfold.model.interface import *
+from igfold.model.components import TriangleGraphTransformer, IPAEncoder, IPATransformer
+from igfold.utils.coordinates import get_ideal_coords, place_o_coords
+from igfold.utils.transforms import quaternion_to_matrix
+from igfold.training.utils import *
+from igfold.utils.general import exists
+
+ATOM_DIM = 3
+
+
+class IgFold(pl.LightningModule):
+    def __init__(
+        self,
+        config,
+        config_overwrite=None,
+    ):
+        super().__init__()
+
+        self.save_hyperparameters()
+        config = self.hparams.config
+        if exists(config_overwrite):
+            config.update(config_overwrite)
+
+        self.bert_feat_dim = 512
+        self.bert_attn_dim = 64
+
+        self.node_dim = config["node_dim"]
+
+        self.depth = config["depth"]
+        self.gt_depth = config["gt_depth"]
+        self.gt_heads = config["gt_heads"]
+
+        self.temp_ipa_depth = config["temp_ipa_depth"]
+        self.temp_ipa_heads = config["temp_ipa_heads"]
+
+        self.str_ipa_depth = config["str_ipa_depth"]
+        self.str_ipa_heads = config["str_ipa_heads"]
+
+        self.dev_ipa_depth = config["dev_ipa_depth"]
+        self.dev_ipa_heads = config["dev_ipa_heads"]
+
+        self.str_node_transform = nn.Sequential(
+            nn.Linear(
+                self.bert_feat_dim,
+                self.node_dim,
+            ),
+            nn.ReLU(),
+            nn.LayerNorm(self.node_dim),
+        )
+        self.str_edge_transform = nn.Sequential(
+            nn.Linear(
+                self.bert_attn_dim,
+                self.node_dim,
+            ),
+            nn.ReLU(),
+            nn.LayerNorm(self.node_dim),
+        )
+
+        self.main_block = TriangleGraphTransformer(
+            dim=self.node_dim,
+            edge_dim=self.node_dim,
+            depth=self.depth,
+            tri_dim_hidden=2 * self.node_dim,
+            gt_depth=self.gt_depth,
+            gt_heads=self.gt_heads,
+            gt_dim_head=self.node_dim // 2,
+        )
+        self.template_ipa = IPAEncoder(
+            dim=self.node_dim,
+            depth=self.temp_ipa_depth,
+            heads=self.temp_ipa_heads,
+            require_pairwise_repr=True,
+        )
+
+        self.structure_ipa = IPATransformer(
+            dim=self.node_dim,
+            depth=self.str_ipa_depth,
+            heads=self.str_ipa_heads,
+            require_pairwise_repr=True,
+        )
+
+        self.dev_node_transform = nn.Sequential(
+            nn.Linear(self.bert_feat_dim, self.node_dim),
+            nn.ReLU(),
+            nn.LayerNorm(self.node_dim),
+        )
+        self.dev_edge_transform = nn.Sequential(
+            nn.Linear(
+                self.bert_attn_dim,
+                self.node_dim,
+            ),
+            nn.ReLU(),
+            nn.LayerNorm(self.node_dim),
+        )
+        self.dev_ipa = IPAEncoder(
+            dim=self.node_dim,
+            depth=self.dev_ipa_depth,
+            heads=self.dev_ipa_heads,
+            require_pairwise_repr=True,
+        )
+        self.dev_linear = nn.Linear(
+            self.node_dim,
+            4,
+        )
+
+    def get_coords_tran_rot(
+        self,
+        temp_coords,
+        batch_size,
+        seq_len,
+    ):
+        res_coords = rearrange(
+            temp_coords,
+            "b (l a) d -> b l a d",
+            l=seq_len,
+        ).to(self.device)
+        ideal_coords = get_ideal_coords()
+        res_ideal_coords = repeat(
+            ideal_coords,
+            "a d -> b l a d",
+            b=batch_size,
+            l=seq_len,
+        ).to(self.device)
+        _, rotations, translations = kabsch(
+            res_ideal_coords,
+            res_coords,
+            return_translation_rotation=True,
+        )
+        translations = rearrange(
+            translations,
+            "b l () d -> b l d",
+        )
+
+        return translations, rotations
+
+    def clean_input(
+        self,
+        input: IgFoldInput,
+    ):
+        embeddings = input.embeddings
+        temp_coords = input.template_coords
+        temp_mask = input.template_mask
+        batch_mask = input.batch_mask
+        align_mask = input.align_mask
+
+        batch_size = embeddings[0].shape[0]
+        seq_lens = [max(e.shape[1], 0) for e in embeddings]
+        seq_len = sum(seq_lens)
+
+        if not exists(temp_coords):
+            temp_coords = torch.zeros(
+                batch_size,
+                4 * seq_len,
+                ATOM_DIM,
+                device=self.device,
+            ).float()
+        if not exists(temp_mask):
+            temp_mask = torch.zeros(
+                batch_size,
+                4 * seq_len,
+                device=self.device,
+            ).bool()
+        if not exists(batch_mask):
+            batch_mask = torch.ones(
+                batch_size,
+                4 * seq_len,
+                device=self.device,
+            ).bool()
+        if not exists(align_mask):
+            align_mask = torch.ones(
+                batch_size,
+                4 * seq_len,
+                device=self.device,
+            ).bool()
+
+        align_mask = align_mask & batch_mask  # Should already be masked by batch_mask anyway
+        temp_coords[~temp_mask] = 0.
+        for i, (tc, m) in enumerate(zip(temp_coords, temp_mask)):
+            temp_coords[i][m] -= tc[m].mean(-2)
+
+        input.template_coords = temp_coords
+        input.template_mask = temp_mask
+        input.batch_mask = batch_mask
+        input.align_mask = align_mask
+
+        return input, batch_size, seq_lens, seq_len
+
+    def forward(
+        self,
+        input: IgFoldInput,
+    ):
+        input, batch_size, seq_lens, seq_len = self.clean_input(input)
+        embeddings = input.embeddings
+        attentions = input.attentions
+        temp_coords = input.template_coords
+        temp_mask = input.template_mask
+        coords_label = input.coords_label
+        batch_mask = input.batch_mask
+        align_mask = input.align_mask
+        return_embeddings = input.return_embeddings
+
+        res_batch_mask = rearrange(
+            batch_mask,
+            "b (l a) -> b l a",
+            a=4,
+        ).all(-1).to(self.device)
+        res_temp_mask = rearrange(
+            temp_mask,
+            "b (l a) -> b l a",
+            a=4,
+        ).all(-1).to(self.device)
+
+        ### Model forward pass
+
+        bert_feats = torch.cat(embeddings, dim=1).to(self.device)
+        bert_attn = torch.zeros(
+            (batch_size, seq_len, seq_len, self.bert_attn_dim),
+            device=self.device,
+        )
+        for i, (a, l) in enumerate(zip(attentions, seq_lens)):
+            a = rearrange(a, "b n h l1 l2 -> b l1 l2 (n h)")
+            cum_l = sum(seq_lens[:i])
+            bert_attn[:, cum_l:cum_l + l, cum_l:cum_l + l, :] = a
+
+        temp_translations, temp_rotations = self.get_coords_tran_rot(
+            temp_coords,
+            batch_size,
+            seq_len,
+        )
+
+        str_nodes = self.str_node_transform(bert_feats)
+        str_edges = self.str_edge_transform(bert_attn)
+        str_nodes, str_edges = self.main_block(
+            str_nodes,
+            str_edges,
+            mask=res_batch_mask,
+        )
+        gt_embs = str_nodes
+        str_nodes = self.template_ipa(
+            str_nodes,
+            translations=temp_translations,
+            rotations=temp_rotations,
+            pairwise_repr=str_edges,
+            mask=res_temp_mask,
+        )
+        structure_embs = str_nodes
+
+        ipa_coords, ipa_translations, ipa_quaternions = self.structure_ipa(
+            str_nodes,
+            translations=None,
+            quaternions=None,
+            pairwise_repr=str_edges,
+            mask=res_batch_mask,
+        )
+        ipa_rotations = quaternion_to_matrix(ipa_quaternions)
+
+        dev_nodes = self.dev_node_transform(bert_feats)
+        dev_edges = self.dev_edge_transform(bert_attn)
+        dev_out_feats = self.dev_ipa(
+            dev_nodes,
+            translations=ipa_translations.detach(),
+            rotations=ipa_rotations.detach(),
+            pairwise_repr=dev_edges,
+            mask=res_batch_mask,
+        )
+        dev_pred = F.relu(self.dev_linear(dev_out_feats))
+        dev_pred = rearrange(dev_pred, "b l a -> b (l a)", a=4)
+
+        bb_coords = rearrange(
+            ipa_coords[:, :, :3],
+            "b l a d -> b (l a) d",
+        )
+        flat_coords = rearrange(
+            ipa_coords[:, :, :4],
+            "b l a d -> b (l a) d",
+        )
+
+        ### Calculate losses if given labels
+        loss = torch.zeros(
+            batch_size,
+            device=self.device,
+        )
+        if exists(coords_label):
+            rmsd_clamp = self.hparams.config["rmsd_clamp"]
+            coords_loss = kabsch_mse(
+                flat_coords,
+                coords_label,
+                align_mask=batch_mask,
+                mask=batch_mask,
+                clamp=rmsd_clamp,
+            )
+
+            bb_coords_label = rearrange(
+                rearrange(coords_label, "b (l a) d -> b l a d", a=4)[:, :, :3],
+                "b l a d -> b (l a) d")
+            bb_batch_mask = rearrange(
+                rearrange(batch_mask, "b (l a) -> b l a", a=4)[:, :, :3],
+                "b l a -> b (l a)")
+            bondlen_loss = bond_length_l1(
+                bb_coords,
+                bb_coords_label,
+                bb_batch_mask,
+            )
+
+            prmsd_loss = []
+            cum_seq_lens = np.cumsum([0] + seq_lens)
+            for sl_i, sl in enumerate(seq_lens):
+                align_mask_ = align_mask.clone()
+                align_mask_[:, :cum_seq_lens[sl_i]] = False
+                align_mask_[:, cum_seq_lens[sl_i + 1]:] = False
+                res_batch_mask_ = res_batch_mask.clone()
+                res_batch_mask_[:, :cum_seq_lens[sl_i]] = False
+                res_batch_mask_[:, cum_seq_lens[sl_i + 1]:] = False
+
+                if sl == 0 or align_mask_.sum() == 0 or res_batch_mask_.sum(
+                ) == 0:
+                    continue
+
+                prmsd_loss.append(
+                    bb_prmsd_l1(
+                        dev_pred,
+                        flat_coords.detach(),
+                        coords_label,
+                        align_mask=align_mask_,
+                        mask=res_batch_mask_,
+                    ))
+            prmsd_loss = sum(prmsd_loss)
+
+            coords_loss, bondlen_loss = list(
+                map(
+                    lambda l: rearrange(l, "(c b) -> b c", b=batch_size).mean(
+                        1),
+                    [coords_loss, bondlen_loss],
+                ))
+
+            loss += sum([coords_loss, bondlen_loss, prmsd_loss])
+        else:
+            prmsd_loss, coords_loss, bondlen_loss = None, None, None
+
+        if not exists(coords_label):
+            loss = None
+
+        bert_embs = bert_feats if return_embeddings else None
+        bert_attn = bert_attn if return_embeddings else None
+        gt_embs = gt_embs if return_embeddings else None
+        structure_embs = structure_embs if return_embeddings else None
+        output = IgFoldOutput(
+            coords=ipa_coords,
+            prmsd=dev_pred,
+            translations=ipa_translations,
+            rotations=ipa_rotations,
+            coords_loss=coords_loss,
+            bondlen_loss=bondlen_loss,
+            prmsd_loss=prmsd_loss,
+            loss=loss,
+            bert_embs=bert_embs,
+            bert_attn=bert_attn,
+            gt_embs=gt_embs,
+            structure_embs=structure_embs,
+        )
+
+        return output
+    
+    def score_coords(
+        self,
+        input: IgFoldInput,
+        output: IgFoldOutput,
+    ):
+        input, _, _, _ = self.clean_input(input)
+        batch_mask = input.batch_mask
+
+        res_batch_mask = rearrange(
+            batch_mask,
+            "b (l a) -> b l a",
+            a=4,
+        ).all(-1)
+
+        str_translations, str_rotations = output.translations, output.rotations
+
+        bert_feats = output.bert_embs
+        bert_attn = output.bert_attn
+
+        dev_nodes = self.dev_node_transform(bert_feats)
+        dev_edges = self.dev_edge_transform(bert_attn)
+        dev_out_feats = self.dev_ipa(
+            dev_nodes,
+            translations=str_translations.detach(),
+            rotations=str_rotations.detach(),
+            pairwise_repr=dev_edges,
+            mask=res_batch_mask,
+        )
+        dev_pred = F.relu(self.dev_linear(dev_out_feats)).squeeze(-1)
+        dev_pred = rearrange(dev_pred, "b l a -> b (l a)", a=4)
+
+        return dev_pred
+
+    def transform_ideal_coords(self, translations, rotations):
+        b, n, d = translations.shape
+        device = translations.device
+
+        ideal_coords = get_ideal_coords().to(device)
+        ideal_coords = repeat(
+            ideal_coords,
+            "a d -> b l a d",
+            b=b,
+            l=n,
+        )
+        points_global = torch.einsum(
+            'b n a c, b n c d -> b n a d',
+            ideal_coords,
+            rotations,
+        ) + rearrange(
+            translations,
+            "b l d -> b l () d",
+        )
+
+        return points_global
+
+    def gradient_refine(
+        self,
+        input: IgFoldInput,
+        output: IgFoldOutput,
+        num_steps: int = 80,
+    ):
+        input_, _, seq_lens, _ = self.clean_input(input)
+        batch_mask = input_.batch_mask
+        res_batch_mask = rearrange(
+            batch_mask,
+            "b (l a) -> b l a",
+            a=4,
+        ).all(-1)
+        translations, rotations = output.translations, output.rotations
+
+        in_coords = self.transform_ideal_coords(translations,
+                                                rotations).detach()
+        in_flat_coords = rearrange(
+            in_coords[:, :, :4],
+            "b l a d -> b (l a) d",
+        )
+
+        with torch.enable_grad():
+            translations.requires_grad = True
+            rotations.requires_grad = True
+
+            translations = nn.parameter.Parameter(translations)
+            rotations = nn.parameter.Parameter(rotations)
+
+            optimizer = torch.optim.Adam([translations, rotations], lr=2e-2)
+            for _ in range(num_steps):
+                optimizer.zero_grad()
+
+                coords = self.transform_ideal_coords(translations, rotations)
+                viol_loss = violation_loss(coords, seq_lens, res_batch_mask)
+
+                flat_coords = rearrange(
+                    coords[:, :, :4],
+                    "b l a d -> b (l a) d",
+                )
+                rmsd = kabsch_mse(
+                    flat_coords,
+                    in_flat_coords,
+                    align_mask=batch_mask,
+                    mask=batch_mask,
+                )
+
+                output.translations = translations
+                output.rotations = rotations
+
+                loss = rmsd + viol_loss
+
+                loss.backward()
+                optimizer.step()
+
+        prmsd = self.score_coords(input, output)
+
+        coords = place_o_coords(coords)
+        output.coords = coords
+        output.prmsd = prmsd
+
+        return output

model/model/components/GraphTransformer.py

@@ -0,0 +1,255 @@
+###
+#   Inspired by graph transformer implementation from https://github.com/lucidrains/graph-transformer-pytorch
+###
+
+import torch
+from torch import nn, einsum
+from einops import rearrange, repeat
+
+from igfold.utils.general import exists, default
+
+List = nn.ModuleList
+
+
+class PreNorm(nn.Module):
+    def __init__(
+        self,
+        dim,
+        fn,
+    ):
+        super().__init__()
+        self.fn = fn
+        self.norm = nn.LayerNorm(dim)
+
+    def forward(
+        self,
+        x,
+        *args,
+        **kwargs,
+    ):
+        x = self.norm(x)
+        return self.fn(
+            x,
+            *args,
+            **kwargs,
+        )
+
+
+# gated residual
+
+
+class Residual(nn.Module):
+    def forward(
+        self,
+        x,
+        res,
+    ):
+        return x + res
+
+
+class GatedResidual(nn.Module):
+    def __init__(
+        self,
+        dim,
+    ):
+        super().__init__()
+        self.proj = nn.Sequential(
+            nn.Linear(dim * 3, 1, bias=False),
+            nn.Sigmoid(),
+        )
+
+    def forward(self, x, res):
+        gate_input = torch.cat((x, res, x - res), dim=-1)
+        gate = self.proj(gate_input)
+        return x * gate + res * (1 - gate)
+
+
+# attention
+
+
+class Attention(nn.Module):
+    def __init__(
+        self,
+        dim,
+        dim_head=64,
+        heads=8,
+        edge_dim=None,
+    ):
+        super().__init__()
+        edge_dim = default(
+            edge_dim,
+            dim,
+        )
+
+        inner_dim = dim_head * heads
+        self.heads = heads
+        self.scale = dim_head**-0.5
+
+        self.to_q = nn.Linear(
+            dim,
+            inner_dim,
+        )
+        self.to_kv = nn.Linear(
+            dim,
+            inner_dim * 2,
+        )
+        self.edges_to_kv = nn.Linear(
+            edge_dim,
+            inner_dim,
+        )
+
+        self.to_out = nn.Linear(
+            inner_dim,
+            dim,
+        )
+
+    def forward(
+        self,
+        nodes,
+        edges,
+        mask=None,
+    ):
+        h = self.heads
+
+        q = self.to_q(nodes)
+        k, v = self.to_kv(nodes).chunk(
+            2,
+            dim=-1,
+        )
+
+        e_kv = self.edges_to_kv(edges)
+
+        q, k, v, e_kv = map(
+            lambda t: rearrange(
+                t,
+                'b ... (h d) -> (b h) ... d',
+                h=h,
+            ),
+            (q, k, v, e_kv),
+        )
+
+        ek, ev = e_kv, e_kv
+
+        k, v = map(
+            lambda t: rearrange(
+                t,
+                'b j d -> b () j d ',
+            ),
+            (k, v),
+        )
+        k = k + ek
+        v = v + ev
+
+        sim = einsum(
+            'b i d, b i j d -> b i j',
+            q,
+            k,
+        ) * self.scale
+
+        if exists(mask):
+            mask = rearrange(
+                mask,
+                'b i -> b i ()',
+            ) & rearrange(
+                mask,
+                'b j -> b () j',
+            )
+            mask = repeat(
+                mask,
+                "b ... -> (b h) ...",
+                h=self.heads,
+            )
+            max_neg_value = -torch.finfo(sim.dtype).max
+            sim.masked_fill_(~mask, max_neg_value)
+
+        attn = sim.softmax(dim=-1)
+        out = einsum(
+            'b i j, b i j d -> b i d',
+            attn,
+            v,
+        )
+        out = rearrange(
+            out,
+            '(b h) n d -> b n (h d)',
+            h=h,
+        )
+        return self.to_out(out)
+
+
+def FeedForward(dim, ff_mult=4):
+    return nn.Sequential(
+        nn.Linear(dim, dim * ff_mult),
+        nn.GELU(),
+        nn.Linear(dim * ff_mult, dim),
+    )
+
+
+class GraphTransformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        depth,
+        dim_head=64,
+        edge_dim=None,
+        heads=8,
+        with_feedforwards=False,
+        norm_edges=False,
+    ):
+        super().__init__()
+        self.layers = List([])
+        edge_dim = default(
+            edge_dim,
+            dim,
+        )
+        self.norm_edges = nn.LayerNorm(
+            edge_dim) if norm_edges else nn.Identity()
+
+        for _ in range(depth):
+            self.layers.append(
+                List([
+                    List([
+                        PreNorm(
+                            dim,
+                            Attention(
+                                dim,
+                                edge_dim=edge_dim,
+                                dim_head=dim_head,
+                                heads=heads,
+                            )),
+                        GatedResidual(dim)
+                    ]),
+                    List(
+                        [PreNorm(
+                            dim,
+                            FeedForward(dim),
+                        ),
+                         GatedResidual(dim)]) if with_feedforwards else None
+                ]))
+
+    def forward(
+        self,
+        nodes,
+        edges,
+        mask=None,
+    ):
+        edges = self.norm_edges(edges)
+
+        for attn_block, ff_block in self.layers:
+            attn, attn_residual = attn_block
+            nodes = attn_residual(
+                attn(
+                    nodes,
+                    edges,
+                    mask=mask,
+                ),
+                nodes,
+            )
+
+            if exists(ff_block):
+                ff, ff_residual = ff_block
+                nodes = ff_residual(
+                    ff(nodes),
+                    nodes,
+                )
+
+        return nodes, edges

model/model/components/IPABlock.py

@@ -0,0 +1,248 @@
+###
+#   Inspired by IPA implementation from https://github.com/lucidrains/invariant-point-attention
+###
+
+import torch
+import torch.nn.functional as F
+from torch.cuda.amp import autocast
+from contextlib import contextmanager
+from torch import nn, einsum
+
+from einops.layers.torch import Rearrange
+from einops import rearrange, repeat
+
+# helpers
+
+def exists(val):
+    return val is not None
+
+def default(val, d):
+    return val if exists(val) else d
+
+def max_neg_value(t):
+    return -torch.finfo(t.dtype).max
+
+@contextmanager
+def disable_tf32():
+    orig_value = torch.backends.cuda.matmul.allow_tf32
+    torch.backends.cuda.matmul.allow_tf32 = False
+    yield
+    torch.backends.cuda.matmul.allow_tf32 = orig_value
+
+# classes
+
+class InvariantPointAttention(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        heads = 8,
+        scalar_key_dim = 16,
+        scalar_value_dim = 16,
+        point_key_dim = 4,
+        point_value_dim = 4,
+        pairwise_repr_dim = None,
+        require_pairwise_repr = True,
+        eps = 1e-8
+    ):
+        super().__init__()
+        self.eps = eps
+        self.heads = heads
+        self.require_pairwise_repr = require_pairwise_repr
+
+        # num attention contributions
+
+        num_attn_logits = 3 if require_pairwise_repr else 2
+
+        # qkv projection for scalar attention (normal)
+
+        self.scalar_attn_logits_scale = (num_attn_logits * scalar_key_dim) ** -0.5
+
+        self.to_scalar_q = nn.Linear(dim, scalar_key_dim * heads, bias = False)
+        self.to_scalar_k = nn.Linear(dim, scalar_key_dim * heads, bias = False)
+        self.to_scalar_v = nn.Linear(dim, scalar_value_dim * heads, bias = False)
+
+        # qkv projection for point attention (coordinate and orientation aware)
+
+        point_weight_init_value = torch.log(torch.exp(torch.full((heads,), 1.)) - 1.)
+        self.point_weights = nn.Parameter(point_weight_init_value)
+
+        self.point_attn_logits_scale = ((num_attn_logits * point_key_dim) * (9 / 2)) ** -0.5
+
+        self.to_point_q = nn.Linear(dim, point_key_dim * heads * 3, bias = False)
+        self.to_point_k = nn.Linear(dim, point_key_dim * heads * 3, bias = False)
+        self.to_point_v = nn.Linear(dim, point_value_dim * heads * 3, bias = False)
+
+        # pairwise representation projection to attention bias
+
+        pairwise_repr_dim = default(pairwise_repr_dim, dim) if require_pairwise_repr else 0
+
+        if require_pairwise_repr:
+            self.pairwise_attn_logits_scale = num_attn_logits ** -0.5
+
+            self.to_pairwise_attn_bias = nn.Sequential(
+                nn.Linear(pairwise_repr_dim, heads),
+                Rearrange('b ... h -> (b h) ...')
+            )
+
+        # combine out - scalar dim + pairwise dim + point dim * (3 for coordinates in R3 and then 1 for norm)
+
+        self.to_out = nn.Linear(heads * (scalar_value_dim + pairwise_repr_dim + point_value_dim * (3 + 1)), dim)
+
+    def forward(
+        self,
+        single_repr,
+        pairwise_repr = None,
+        *,
+        rotations,
+        translations,
+        mask = None
+    ):
+        x, b, h, eps, require_pairwise_repr = single_repr, single_repr.shape[0], self.heads, self.eps, self.require_pairwise_repr
+        assert not (require_pairwise_repr and not exists(pairwise_repr)), 'pairwise representation must be given as second argument'
+
+        # get queries, keys, values for scalar and point (coordinate-aware) attention pathways
+
+        q_scalar, k_scalar, v_scalar = self.to_scalar_q(x), self.to_scalar_k(x), self.to_scalar_v(x)
+
+        q_point, k_point, v_point = self.to_point_q(x), self.to_point_k(x), self.to_point_v(x)
+
+        # split out heads
+
+        q_scalar, k_scalar, v_scalar = map(lambda t: rearrange(t, 'b n (h d) -> (b h) n d', h = h), (q_scalar, k_scalar, v_scalar))
+        q_point, k_point, v_point = map(lambda t: rearrange(t, 'b n (h d c) -> (b h) n d c', h = h, c = 3), (q_point, k_point, v_point))
+
+        rotations = repeat(rotations, 'b n r1 r2 -> (b h) n r1 r2', h = h)
+        translations = repeat(translations, 'b n c -> (b h) n () c', h = h)
+
+        # rotate qkv points into global frame
+
+        q_point = einsum('b n d c, b n c r -> b n d r', q_point, rotations) + translations
+        k_point = einsum('b n d c, b n c r -> b n d r', k_point, rotations) + translations
+        v_point = einsum('b n d c, b n c r -> b n d r', v_point, rotations) + translations
+
+        # derive attn logits for scalar and pairwise
+
+        attn_logits_scalar = einsum('b i d, b j d -> b i j', q_scalar, k_scalar) * self.scalar_attn_logits_scale
+
+        if require_pairwise_repr:
+            attn_logits_pairwise = self.to_pairwise_attn_bias(pairwise_repr) * self.pairwise_attn_logits_scale
+
+        # derive attn logits for point attention
+
+        point_qk_diff = rearrange(q_point, 'b i d c -> b i () d c') - rearrange(k_point, 'b j d c -> b () j d c')
+        point_dist = (point_qk_diff ** 2).sum(dim = -2)
+
+        point_weights = F.softplus(self.point_weights)
+        point_weights = repeat(point_weights, 'h -> (b h) () () ()', b = b)
+
+        attn_logits_points = -0.5 * (point_dist * point_weights * self.point_attn_logits_scale).sum(dim = -1)
+
+        # combine attn logits
+
+        attn_logits = attn_logits_scalar + attn_logits_points
+
+        if require_pairwise_repr:
+            attn_logits = attn_logits + attn_logits_pairwise
+
+        # mask
+
+        if exists(mask):
+            mask = rearrange(mask, 'b i -> b i ()') * rearrange(mask, 'b j -> b () j')
+            mask = repeat(mask, 'b i j -> (b h) i j', h = h)
+            mask_value = max_neg_value(attn_logits)
+            attn_logits = attn_logits.masked_fill(~mask, mask_value)
+
+        # attention
+
+        attn = attn_logits.softmax(dim = - 1)
+
+        with disable_tf32(), autocast(enabled = False):
+            # disable TF32 for precision
+
+            # aggregate values
+
+            results_scalar = einsum('b i j, b j d -> b i d', attn, v_scalar)
+
+            attn_with_heads = rearrange(attn, '(b h) i j -> b h i j', h = h)
+
+            if require_pairwise_repr:
+                results_pairwise = einsum('b h i j, b i j d -> b h i d', attn_with_heads, pairwise_repr)
+
+            # aggregate point values
+
+            results_points = einsum('b i j, b j d c -> b i d c', attn, v_point)
+
+            # rotate aggregated point values back into local frame
+
+            results_points = einsum('b n d c, b n c r -> b n d r', results_points - translations, rotations.transpose(-1, -2))
+            results_points_norm = torch.sqrt( torch.square(results_points).sum(dim=-1) + eps )
+
+        # merge back heads
+
+        results_scalar = rearrange(results_scalar, '(b h) n d -> b n (h d)', h = h)
+        results_points = rearrange(results_points, '(b h) n d c -> b n (h d c)', h = h)
+        results_points_norm = rearrange(results_points_norm, '(b h) n d -> b n (h d)', h = h)
+
+        results = (results_scalar, results_points, results_points_norm)
+
+        if require_pairwise_repr:
+            results_pairwise = rearrange(results_pairwise, 'b h n d -> b n (h d)', h = h)
+            results = (*results, results_pairwise)
+
+        # concat results and project out
+
+        results = torch.cat(results, dim = -1)
+        return self.to_out(results)
+
+# one transformer block based on IPA
+
+def FeedForward(dim, mult = 1., num_layers = 2, act = nn.ReLU):
+    layers = []
+    dim_hidden = dim * mult
+
+    for ind in range(num_layers):
+        is_first = ind == 0
+        is_last  = ind == (num_layers - 1)
+        dim_in   = dim if is_first else dim_hidden
+        dim_out  = dim if is_last else dim_hidden
+
+        layers.append(nn.Linear(dim_in, dim_out))
+
+        if is_last:
+            continue
+
+        layers.append(act())
+
+    return nn.Sequential(*layers)
+
+class IPABlock(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        ff_mult = 1,
+        ff_num_layers = 3,     # in the paper, they used 3 layer transition (feedforward) block
+        post_norm = True,      # in the paper, they used post-layernorm - offering pre-norm as well
+        **kwargs
+    ):
+        super().__init__()
+        self.post_norm = post_norm
+
+        self.attn_norm = nn.LayerNorm(dim)
+        self.attn = InvariantPointAttention(dim = dim, **kwargs)
+
+        self.ff_norm = nn.LayerNorm(dim)
+        self.ff = FeedForward(dim, mult = ff_mult, num_layers = ff_num_layers)
+
+    def forward(self, x, **kwargs):
+        post_norm = self.post_norm
+
+        attn_input = x if post_norm else self.attn_norm(x)
+        x = self.attn(attn_input, **kwargs) + x
+        x = self.attn_norm(x) if post_norm else x
+
+        ff_input = x if post_norm else self.ff_norm(x)
+        x = self.ff(ff_input) + x
+        x = self.ff_norm(x) if post_norm else x
+        return x

model/model/components/IPATransformer.py

@@ -0,0 +1,174 @@
+###
+#   Inspired by IPA implementation from https://github.com/lucidrains/invariant-point-attention
+###
+
+import torch
+import torch.nn.functional as F
+from torch import nn, einsum
+from einops import rearrange, repeat
+
+from igfold.model.components.IPABlock import IPABlock
+from igfold.utils.coordinates import get_ideal_coords, place_o_coords
+from igfold.utils.general import exists
+from igfold.utils.transforms import quaternion_multiply, quaternion_to_matrix
+
+
+class IPAEncoder(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        **kwargs,
+    ):
+        super().__init__()
+
+        # layers
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            self.layers.append(IPABlock(
+                dim=dim,
+                **kwargs,
+            ))
+
+    def forward(
+        self,
+        x,
+        *,
+        translations=None,
+        rotations=None,
+        pairwise_repr=None,
+        mask=None,
+    ):
+        for block in self.layers:
+            x = block(
+                x,
+                pairwise_repr=pairwise_repr,
+                rotations=rotations,
+                translations=translations,
+                mask=mask,
+            )
+
+        return x
+
+
+class IPATransformer(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        depth,
+        stop_rotation_grad=False,
+        **kwargs,
+    ):
+        super().__init__()
+
+        self.stop_rotation_grad = stop_rotation_grad
+
+        self.quaternion_to_matrix = quaternion_to_matrix
+        self.quaternion_multiply = quaternion_multiply
+
+        # layers
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            ipa_block = IPABlock(
+                dim=dim,
+                **kwargs,
+            )
+            linear = nn.Linear(dim, 6)
+            torch.nn.init.zeros_(linear.weight.data)
+            torch.nn.init.zeros_(linear.bias.data)
+            self.layers.append(nn.ModuleList([ipa_block, linear]))
+
+    def forward(
+        self,
+        single_repr,
+        *,
+        translations=None,
+        quaternions=None,
+        pairwise_repr=None,
+        mask=None,
+    ):
+        x, device, quaternion_multiply, quaternion_to_matrix = single_repr, single_repr.device, self.quaternion_multiply, self.quaternion_to_matrix
+        b, n, *_ = x.shape
+
+        # if no initial quaternions passed in, start from identity
+
+        if not exists(quaternions):
+            quaternions = torch.tensor(
+                [1., 0., 0., 0.],
+                device=device,
+            )  # initial rotations
+            quaternions = repeat(
+                quaternions,
+                'd -> b n d',
+                b=b,
+                n=n,
+            )
+
+        # if not translations passed in, start from identity
+
+        if not exists(translations):
+            translations = torch.zeros(
+                (b, n, 3),
+                device=device,
+            )
+
+        # go through the layers and apply invariant point attention and feedforward
+
+        for block, to_update in self.layers:
+            rotations = quaternion_to_matrix(quaternions)
+            if self.stop_rotation_grad:
+                rotations = rotations.detach()
+
+            x = block(
+                x,
+                pairwise_repr=pairwise_repr,
+                rotations=rotations,
+                translations=translations,
+                mask=mask,
+            )
+
+            # update quaternion and translation
+
+            quaternion_update, translation_update = to_update(x).chunk(
+                2,
+                dim=-1,
+            )
+            quaternion_update = F.pad(
+                quaternion_update,
+                (1, 0),
+                value=1.,
+            )
+
+            quaternions = quaternion_multiply(
+                quaternions,
+                quaternion_update,
+            )
+            translations = translations + einsum(
+                'b n c, b n c r -> b n r',
+                translation_update,
+                rotations,
+            )
+
+        ideal_coords = get_ideal_coords().to(device)
+        ideal_coords = repeat(
+            ideal_coords,
+            "a d -> b l a d",
+            b=b,
+            l=n,
+        )
+
+        rotations = quaternion_to_matrix(quaternions)
+        points_global = einsum(
+            'b n a c, b n c d -> b n a d',
+            ideal_coords,
+            rotations,
+        ) + rearrange(
+            translations,
+            "b l d -> b l () d",
+        )
+
+        points_global = place_o_coords(points_global)
+
+        return points_global, translations, quaternions

model/model/components/TriangleGraphTransformer.py

@@ -0,0 +1,62 @@
+from torch import nn
+
+from .GraphTransformer import GraphTransformer
+from .TriangleMultiplicativeModule import TriangleMultiplicativeModule
+from igfold.utils.general import exists
+
+
+class TriangleGraphTransformer(nn.Module):
+    def __init__(
+        self,
+        dim,
+        edge_dim,
+        depth,
+        gt_depth=1,
+        gt_dim_head=32,
+        gt_heads=8,
+        tri_dim_hidden=None,
+    ):
+        super().__init__()
+
+        self.layers = nn.ModuleList([])
+        for _ in range(depth):
+            graph_transformer = GraphTransformer(
+                dim=dim,
+                edge_dim=edge_dim,
+                depth=gt_depth,
+                heads=gt_heads,
+                dim_head=gt_dim_head,
+                with_feedforwards=True,
+            )
+            triangle_out = TriangleMultiplicativeModule(
+                dim=edge_dim,
+                hidden_dim=tri_dim_hidden,
+                mix='outgoing',
+            )
+            triangle_in = TriangleMultiplicativeModule(
+                dim=edge_dim,
+                hidden_dim=tri_dim_hidden,
+                mix='ingoing',
+            )
+
+            self.layers.append(
+                nn.ModuleList([graph_transformer, triangle_out, triangle_in]))
+
+    def forward(self, nodes, edges, mask=None):
+        for gt, tri_out, tri_in in self.layers:
+            if exists(mask):
+                tri_mask = mask.unsqueeze(-2) & mask.unsqueeze(-1)
+            else:
+                tri_mask = None
+
+            nodes, _ = gt(nodes, edges, mask=mask)
+            edges = edges + tri_out(
+                edges,
+                mask=tri_mask,
+            )
+            edges = edges + tri_in(
+                edges,
+                mask=tri_mask,
+            )
+
+        return nodes, edges

model/model/components/TriangleMultiplicativeModule.py

@@ -0,0 +1,109 @@
+###
+#   Inspired by triangle multiplicative update implementation from https://github.com/lucidrains/triangle-multiplicative-module
+###
+
+from torch import nn, einsum
+import torch.nn.functional as F
+from einops import rearrange
+
+from igfold.utils.general import exists, default
+
+
+class TriangleMultiplicativeModule(nn.Module):
+    def __init__(
+        self,
+        *,
+        dim,
+        hidden_dim=None,
+        mix='ingoing',
+    ):
+        super().__init__()
+        assert mix in {'ingoing',
+                       'outgoing'}, 'mix must be either ingoing or outgoing'
+
+        hidden_dim = default(
+            hidden_dim,
+            dim,
+        )
+        self.norm = nn.LayerNorm(dim)
+
+        self.left_proj = nn.Linear(
+            dim,
+            hidden_dim,
+        )
+        self.right_proj = nn.Linear(
+            dim,
+            hidden_dim,
+        )
+
+        self.left_gate = nn.Linear(
+            dim,
+            hidden_dim,
+        )
+        self.right_gate = nn.Linear(
+            dim,
+            hidden_dim,
+        )
+        self.out_gate = nn.Linear(dim, dim)
+
+        # initialize all gating to be identity
+
+        for gate in (
+                self.left_gate,
+                self.right_gate,
+                self.out_gate,
+        ):
+            nn.init.constant_(
+                gate.weight,
+                0.,
+            )
+            nn.init.constant_(
+                gate.bias,
+                1.,
+            )
+
+        if mix == 'outgoing':
+            self.mix_einsum_eq = '... i k d, ... j k d -> ... i j d'
+        elif mix == 'ingoing':
+            self.mix_einsum_eq = '... k j d, ... k i d -> ... i j d'
+
+        self.to_out_norm = nn.LayerNorm(hidden_dim)
+        self.to_out = nn.Linear(
+            hidden_dim,
+            dim,
+        )
+
+    def forward(self, x, mask=None):
+        assert x.shape[1] == x.shape[2], 'feature map must be symmetrical'
+        if exists(mask):
+            mask = rearrange(
+                mask,
+                'b i j -> b i j ()',
+            )
+
+        x = self.norm(x)
+
+        left = self.left_proj(x)
+        right = self.right_proj(x)
+
+        if exists(mask):
+            left = left * mask
+            right = right * mask
+
+        left_gate = self.left_gate(x).sigmoid()
+        right_gate = self.right_gate(x).sigmoid()
+        out_gate = self.out_gate(x).sigmoid()
+
+        left = left * left_gate
+        right = right * right_gate
+
+        out = einsum(
+            self.mix_einsum_eq,
+            left,
+            right,
+        )
+
+        out = self.to_out_norm(out)
+        out = self.to_out(out)
+        out = out * out_gate
+        return out

model/model/components/__init__.py

@@ -0,0 +1,4 @@
+from .GraphTransformer import *
+from .IPATransformer import *
+from .TriangleMultiplicativeModule import *
+from .TriangleGraphTransformer import *

model/model/interface.py

@@ -0,0 +1,40 @@
+from dataclasses import dataclass
+from typing import List, Optional, Union
+import torch
+
+
+@dataclass
+class IgFoldInput():
+    """
+    Input type of for IgFold model.
+    """
+
+    embeddings: List[List[torch.FloatTensor]]
+    attentions: List[List[torch.FloatTensor]]
+    template_coords: Optional[torch.FloatTensor] = None
+    template_mask: Optional[torch.BoolTensor] = None
+    batch_mask: Optional[torch.BoolTensor] = None
+    align_mask: Optional[torch.BoolTensor] = None
+    coords_label: Optional[torch.FloatTensor] = None
+    return_embeddings: Optional[bool] = False
+
+
+@dataclass
+class IgFoldOutput():
+    """
+    Output type of for IgFold model.
+    """
+
+    coords: torch.FloatTensor
+    prmsd: torch.FloatTensor
+    translations: torch.FloatTensor
+    rotations: torch.FloatTensor
+    coords_loss: Optional[torch.FloatTensor] = None
+    torsion_loss: Optional[torch.FloatTensor] = None
+    bondlen_loss: Optional[torch.FloatTensor] = None
+    prmsd_loss: Optional[torch.FloatTensor] = None
+    loss: Optional[torch.FloatTensor] = None
+    bert_embs: Optional[torch.FloatTensor] = None
+    bert_attn: Optional[torch.FloatTensor] = None
+    gt_embs: Optional[torch.FloatTensor] = None
+    structure_embs: Optional[torch.FloatTensor] = None

model/refine/openmm_ref.py

@@ -0,0 +1,44 @@
+import pdbfixer
+import openmm
+
+ENERGY = openmm.unit.kilocalories_per_mole
+LENGTH = openmm.unit.angstroms
+
+def refine(pdb_file, stiffness=10., tolerance=2.39, use_gpu=False):
+    tolerance = tolerance * ENERGY
+    stiffness = stiffness * ENERGY / (LENGTH**2)
+
+    fixer = pdbfixer.PDBFixer(pdb_file)
+    fixer.findMissingResidues()
+    fixer.findMissingAtoms()
+    fixer.addMissingAtoms()
+
+    force_field = openmm.app.ForceField("amber14/protein.ff14SB.xml")
+    modeller = openmm.app.Modeller(fixer.topology, fixer.positions)
+    modeller.addHydrogens(force_field)
+    system = force_field.createSystem(modeller.topology)
+
+    force = openmm.CustomExternalForce("0.5 * k * ((x-x0)^2 + (y-y0)^2 + (z-z0)^2)")
+    force.addGlobalParameter("k", stiffness)
+    for p in ["x0", "y0", "z0"]:
+        force.addPerParticleParameter(p)
+    for residue in modeller.topology.residues():
+        for atom in residue.atoms():
+            if atom.name in ["N", "CA", "C", "CB"]:
+                force.addParticle(atom.index,
+                                    modeller.positions[atom.index])
+    system.addForce(force)
+
+    integrator = openmm.LangevinIntegrator(0, 0.01, 1.0)
+    platform = openmm.Platform.getPlatformByName("CUDA" if use_gpu else "CPU")
+
+    simulation = openmm.app.Simulation(modeller.topology, system, integrator, platform)
+    simulation.context.setPositions(modeller.positions)
+    simulation.minimizeEnergy(tolerance)
+
+    with open(pdb_file, "w") as f:
+        openmm.app.PDBFile.writeFile(
+            simulation.topology,
+            simulation.context.getState(getPositions=True).getPositions(),
+            f,
+            keepIds=True,)

model/refine/pyrosetta_ref.py

@@ -0,0 +1,125 @@
+import pyrosetta
+
+from igfold.utils.general import exists
+
+
+def init_pyrosetta(init_string=None, silent=True):
+    if not exists(init_string):
+        init_string = "-mute all -ignore_zero_occupancy false -detect_disulf true -detect_disulf_tolerance 1.5 -check_cdr_chainbreaks false"
+    pyrosetta.init(init_string, silent=silent)
+
+
+def get_min_mover(
+    max_iter: int = 1000,
+    sf_name: str = "ref2015_cst",
+    coord_cst_weight: float = 1,
+    dih_cst_weight: float = 1,
+) -> pyrosetta.rosetta.protocols.moves.Mover:
+    """
+    Create full-atom minimization mover
+    """
+
+    sf = pyrosetta.create_score_function(sf_name)
+    sf.set_weight(
+        pyrosetta.rosetta.core.scoring.ScoreType.cart_bonded,
+        1,
+    )
+    sf.set_weight(
+        pyrosetta.rosetta.core.scoring.ScoreType.pro_close,
+        0,
+    )
+    sf.set_weight(
+        pyrosetta.rosetta.core.scoring.ScoreType.coordinate_constraint,
+        coord_cst_weight,
+    )
+    sf.set_weight(
+        pyrosetta.rosetta.core.scoring.ScoreType.dihedral_constraint,
+        dih_cst_weight,
+    )
+
+    mmap = pyrosetta.rosetta.core.kinematics.MoveMap()
+    mmap.set_bb(True)
+    mmap.set_chi(False)
+    mmap.set_jump(False)
+    min_mover = pyrosetta.rosetta.protocols.minimization_packing.MinMover(
+        mmap,
+        sf,
+        'lbfgs_armijo_nonmonotone',
+        0.0001,
+        True,
+    )
+    min_mover.max_iter(max_iter)
+    min_mover.cartesian(True)
+
+    return min_mover
+
+
+def get_fa_relax_mover(
+        max_iter: int = 100) -> pyrosetta.rosetta.protocols.moves.Mover:
+    """
+    Create full-atom relax mover
+    """
+
+    sf = pyrosetta.create_score_function('ref2015_cst')
+
+    mmap = pyrosetta.rosetta.core.kinematics.MoveMap()
+    mmap.set_bb(True)
+    mmap.set_chi(True)
+    mmap.set_jump(True)
+
+    relax = pyrosetta.rosetta.protocols.relax.FastRelax()
+    relax.set_scorefxn(sf)
+    relax.max_iter(max_iter)
+    relax.set_movemap(mmap)
+
+    return relax
+
+
+def get_repack_mover():
+    tf = pyrosetta.rosetta.core.pack.task.TaskFactory()
+    tf.push_back(
+        pyrosetta.rosetta.core.pack.task.operation.InitializeFromCommandline())
+    tf.push_back(
+        pyrosetta.rosetta.core.pack.task.operation.RestrictToRepacking())
+
+    packer = pyrosetta.rosetta.protocols.minimization_packing.PackRotamersMover(
+    )
+    packer.task_factory(tf)
+
+    return packer
+
+
+def refine(out_pdb_file,
+           pdb_string,
+           minimization_iter=100,
+           constrain=True,
+           idealize=False):
+    # create new pose
+    pose = pyrosetta.rosetta.core.pose.Pose()
+    pyrosetta.rosetta.core.import_pose.pose_from_pdbstring(
+        pose,
+        pdb_string,
+    )
+
+    if constrain:
+        cst_mover = pyrosetta.rosetta.protocols.relax.AtomCoordinateCstMover()
+        cst_mover.cst_sidechain(False)
+        cst_mover.apply(pose)
+
+    min_mover = get_min_mover(
+        max_iter=minimization_iter,
+        coord_cst_weight=1,
+        dih_cst_weight=0,
+    )
+    min_mover.apply(pose)
+
+    if idealize:
+        idealize_mover = pyrosetta.rosetta.protocols.idealize.IdealizeMover()
+        idealize_mover.apply(pose)
+
+    packer = get_repack_mover()
+    packer.apply(pose)
+
+    min_mover.apply(pose)
+
+    pose.dump_pdb(out_pdb_file)

model/training/utils.py

@@ -0,0 +1,287 @@
+from einops import rearrange, repeat
+import numpy as np
+import torch
+import torch.nn.functional as F
+
+from igfold.utils.constants import *
+from igfold.utils.general import exists
+from igfold.utils.geometry import dist, angle, dihedral
+
+
+def kabsch(
+    mobile,
+    stationary,
+    return_translation_rotation=False,
+):
+    X = rearrange(
+        mobile,
+        "... l d -> ... d l",
+    )
+    Y = rearrange(
+        stationary,
+        "... l d -> ... d l",
+    )
+
+    #  center X and Y to the origin
+    XT, YT = X.mean(dim=-1, keepdim=True), Y.mean(dim=-1, keepdim=True)
+    X_ = X - XT
+    Y_ = Y - YT
+
+    # calculate convariance matrix
+    C = torch.einsum("... x l, ... y l -> ... x y", X_, Y_)
+
+    # Optimal rotation matrix via SVD
+    if int(torch.__version__.split(".")[1]) < 8:
+        # warning! int torch 1.<8 : W must be transposed
+        V, S, W = torch.svd(C)
+        W = rearrange(W, "... a b -> ... b a")
+    else:
+        V, S, W = torch.linalg.svd(C)
+
+    # determinant sign for direction correction
+    v_det = torch.det(V.to("cpu")).to(X.device)
+    w_det = torch.det(W.to("cpu")).to(X.device)
+    d = (v_det * w_det) < 0.0
+    if d.any():
+        S[d] = S[d] * (-1)
+        V[d, :] = V[d, :] * (-1)
+
+    # Create Rotation matrix U
+    U = torch.matmul(V, W)  #.to(device)
+
+    U = rearrange(
+        U,
+        "... d x -> ... x d",
+    )
+    XT = rearrange(
+        XT,
+        "... d x -> ... x d",
+    )
+    YT = rearrange(
+        YT,
+        "... d x -> ... x d",
+    )
+
+    if return_translation_rotation:
+        return XT, U, YT
+
+    transform = lambda coords: torch.einsum(
+        "... l d, ... x d -> ... l x",
+        coords - XT,
+        U,
+    ) + YT
+    mobile = transform(mobile)
+
+    return mobile, transform
+
+
+def do_kabsch(
+    mobile,
+    stationary,
+    align_mask=None,
+):
+    mobile_, stationary_ = mobile.clone(), stationary.clone()
+    if exists(align_mask):
+        mobile_[~align_mask] = mobile_[align_mask].mean(dim=-2)
+        stationary_[~align_mask] = stationary_[align_mask].mean(dim=-2)
+        _, kabsch_xform = kabsch(
+            mobile_,
+            stationary_,
+        )
+    else:
+        _, kabsch_xform = kabsch(
+            mobile_,
+            stationary_,
+        )
+
+    return kabsch_xform(mobile)
+
+
+def kabsch_mse(
+    pred,
+    target,
+    align_mask=None,
+    mask=None,
+    clamp=0.,
+    sqrt=False,
+):
+    aligned_target = do_kabsch(
+        mobile=target,
+        stationary=pred.detach(),
+        align_mask=align_mask,
+    )
+    mse = F.mse_loss(
+        pred,
+        aligned_target,
+        reduction='none',
+    ).mean(-1)
+
+    if clamp > 0:
+        mse = torch.clamp(mse, max=clamp**2)
+
+    if exists(mask):
+        mse = torch.sum(
+            mse * mask,
+            dim=-1,
+        ) / torch.sum(
+            mask,
+            dim=-1,
+        )
+    else:
+        mse = mse.mean(-1)
+
+    if sqrt:
+        mse = mse.sqrt()
+
+    return mse
+
+
+def bond_length_l1(
+    pred,
+    target,
+    mask,
+    offsets=[1, 2],
+):
+    losses = []
+    for c in range(pred.shape[0]):
+        m, p, t = mask[c], pred[c], target[c]
+        for o in offsets:
+            m_ = (torch.stack([m[:-o], m[o:]])).all(0)
+            pred_lens = torch.norm(p[:-o] - p[o:], dim=-1)
+            target_lens = torch.norm(t[:-o] - t[o:], dim=-1)
+
+            losses.append(
+                torch.abs(pred_lens[m_] - target_lens[m_], ).mean() / o)
+
+    return torch.stack(losses)
+
+
+def bb_prmsd_l1(
+    pdev,
+    pred,
+    target,
+    align_mask=None,
+    mask=None,
+):
+    aligned_target = do_kabsch(
+        mobile=target,
+        stationary=pred,
+        align_mask=align_mask,
+    )
+    bb_dev = (pred - aligned_target).norm(dim=-1)
+    loss = F.l1_loss(
+        pdev,
+        bb_dev,
+        reduction='none',
+    )
+
+    if exists(mask):
+        mask = repeat(mask, "b l -> b (l 4)")
+        loss = torch.sum(
+            loss * mask,
+            dim=-1,
+        ) / torch.sum(
+            mask,
+            dim=-1,
+        )
+    else:
+        loss = loss.mean(-1)
+
+    loss = loss.mean(-1).unsqueeze(0)
+
+    return loss
+
+def bond_len_loss(pred, seq_lens, mask, eps=EPS):
+    b, l, a, d = pred.shape
+
+    pred_bb = pred[:, :, :3]
+    mask = repeat(mask, "b l -> b (l 3)")
+    for seq_len in seq_lens:
+        mask[:, 3 * seq_len - 1] = 0
+    mask_bb = mask[:, :-1] * mask[:, 1:]
+
+    pred_bond_lens = dist(
+        rearrange(pred_bb, "b l a d -> b (l a) d")[:, :-1],
+        rearrange(pred_bb, "b l a d -> b (l a) d")[:, 1:],
+    )
+    lit_bond_lens = repeat(
+        torch.tensor([BL_N_CA, BL_CA_C, BL_C_N]),
+        "bl -> b (l bl)",
+        b=b,
+        l=l,
+    )[:, :-1]
+    lit_bond_lens = lit_bond_lens.to(pred_bond_lens.device)
+
+    bl_loss = torch.abs(pred_bond_lens - lit_bond_lens) * mask_bb
+    bl_loss = bl_loss.sum(-1) / (mask.sum(-1) + eps)
+
+    return bl_loss
+
+
+def bond_angle_loss(pred, seq_lens, mask, eps=EPS):
+    b, l, a, d = pred.shape
+
+    for seq_len in seq_lens:
+        mask[:, seq_len - 1] = 0
+    mask_ = mask[:, 1:] * mask[:, :-1]
+
+    N, CA, C, CB = pred.unbind(-2)
+    ba_CA_C_N = angle(CA[:, :-1], C[:, :-1], N[:, 1:], eps=eps)
+    ba_CA_C_N_loss = 1 - torch.cos(ba_CA_C_N - BA_CA_C_N * np.pi / 180)
+    ba_CA_C_N_loss = ba_CA_C_N_loss * mask_
+
+    ba_C_N_CA = angle(C[:, :-1], N[:, 1:], CA[:, 1:], eps=eps)
+    ba_C_N_CA_loss = 1 - torch.cos(ba_C_N_CA - BA_C_N_CA * np.pi / 180)
+    ba_C_N_CA_loss = ba_C_N_CA_loss * mask_
+
+    loss = ba_CA_C_N_loss + ba_C_N_CA_loss
+    loss = loss.sum(-1) / (mask_.sum(-1) + eps)
+
+    return loss
+
+
+def vdw_clash_loss(pred, mask, tol=1.5, eps=EPS):
+    b, l, a, d = pred.shape
+
+    mask_ = repeat(mask, "b l -> b (l a)", a=a)
+    mask_ = (mask_.unsqueeze(-1) * mask_.unsqueeze(-2))
+
+    vdw_radii = torch.tensor([VDW_N, VDW_C, VDW_C, VDW_C])
+    vdw_radii = repeat(vdw_radii, "a -> b (l a)", b=b, l=l)
+    vdw_distances = (vdw_radii.unsqueeze(-2) + vdw_radii.unsqueeze(-3))
+    vdw_distances = vdw_distances.to(pred.device)
+
+    pred_ = rearrange(pred, "b l a d -> b (l a) d")
+    atomic_distances = (pred_.unsqueeze(-2) - pred_.unsqueeze(-3)).norm(dim=-1)
+
+    loss = (vdw_distances - tol - atomic_distances).clamp(min=0)
+    loss = loss.sum(dim=(-1, -2)) / (mask_.sum(dim=(-1, -2)) + eps)
+
+    return loss
+
+
+def cis_peptide_loss(pred, seq_lens, mask, eps=EPS):
+    for seq_len in seq_lens:
+        mask[:, seq_len - 1] = 0
+    mask_ = mask[:, 1:] * mask[:, :-1]
+
+    N, CA, C, _ = pred.unbind(-2)
+    dih = dihedral(CA[:, :-1], C[:, :-1], N[:, 1:], CA[:, 1:], eps=0)
+
+    loss = 1 - torch.cos(dih - np.pi)
+    loss = loss.sum(dim=(-1, -2)) / (mask_.sum(dim=(-1, -2)) + eps)
+
+    return loss
+
+
+def violation_loss(pred, seq_lens, mask, eps=EPS):
+    b, l, a, d = pred.shape
+
+    bl_loss = bond_len_loss(pred, seq_lens, mask, eps=eps)
+    ba_loss = bond_angle_loss(pred, seq_lens, mask, eps=eps)
+    vdw_loss = vdw_clash_loss(pred, mask)
+    cis_loss = cis_peptide_loss(pred, seq_lens, mask, eps=eps)
+
+    loss = bl_loss + ba_loss + vdw_loss + cis_loss
+
+    return loss

model/utils/ab_metrics.py

@@ -0,0 +1,60 @@
+import pyrosetta
+
+
+def get_vh_vl_orientation(pose_1):
+    pose_i1 = pyrosetta.rosetta.protocols.antibody.AntibodyInfo(pose_1)
+
+    #vl_vh_distance, opening_angle, opposite_angle, packing_angle
+    results = pyrosetta.protocols.antibody.vl_vh_orientation_coords(
+        pose_1,
+        pose_i1,
+    )
+
+    results_labels = [
+        'vl-vh_distance', 'opening_angle', 'opposite_angle', 'packing_angle'
+    ]
+    results_dict = {}
+    for i in range(4):
+        results_dict[results_labels[i]] = results[i + 1]
+    return results_dict
+
+
+def get_ab_metrics(
+    pose_1,
+    pose_2,
+):
+    pose_i1 = pyrosetta.rosetta.protocols.antibody.AntibodyInfo(pose_1)
+    pose_i2 = pyrosetta.rosetta.protocols.antibody.AntibodyInfo(pose_2)
+
+    results = pyrosetta.rosetta.protocols.antibody.cdr_backbone_rmsds(
+        pose_1,
+        pose_2,
+        pose_i1,
+        pose_i2,
+    )
+
+    results_labels = [
+        'ocd', 'frh_rms', 'h1_rms', 'h2_rms', 'h3_rms', 'frl_rms', 'l1_rms',
+        'l2_rms', 'l3_rms'
+    ]
+    results_dict = {}
+    for i in range(9):
+        results_dict[results_labels[i]] = results[i + 1]
+
+    return results_dict
+
+
+def get_pose_cdr_clusters(pose):
+    clus_name = {1: 'h1', 2: 'h2', 3: 'h3', 4: 'l1', 5: 'l2', 6: 'l3'}
+    ab_info = pyrosetta.rosetta.protocols.antibody.AntibodyInfo(
+        pose, pyrosetta.rosetta.protocols.antibody.CDRDefinitionEnum.North)
+    ab_info.setup_CDR_clusters(pose)
+
+    clusters = {}
+    for enum in range(1, ab_info.get_total_num_CDRs() + 1):
+        result = ab_info.get_CDR_cluster(
+            pyrosetta.rosetta.protocols.antibody.CDRNameEnum(enum))
+        clus = ab_info.get_cluster_name(result.cluster())
+        clusters[clus_name[enum]] = clus
+
+    return clusters

model/utils/abnumber_.py

@@ -0,0 +1,88 @@
+import warnings
+from abnumber import Chain
+from Bio.PDB import PDBParser, PDBIO
+from Bio.SeqUtils import seq1
+
+from igfold.utils.pdb import clean_pdb
+
+
+def is_heavy(seq):
+    chain = Chain(seq, scheme='chothia')
+
+    return chain.is_heavy_chain()
+
+
+def rechain_pdb(pdb_file):
+    parser = PDBParser()
+    with warnings.catch_warnings(record=True):
+        structure = parser.get_structure("_", pdb_file)
+
+    for chain in structure.get_chains():
+        seq = seq1(''.join([residue.resname for residue in chain]))
+        abnum_chain = Chain(seq, scheme='chothia')
+        chain_id = "H" if abnum_chain.is_heavy_chain() else "L"
+        try:
+            chain.id = chain_id
+        except ValueError:
+            chain.id = chain_id + "_"
+    for chain in structure.get_chains():
+        if "_" in chain.id:
+            chain.id = chain.id.replace("_", "")
+
+    io = PDBIO()
+    io.set_structure(structure)
+    io.save(pdb_file)
+
+
+def renumber_pdb(
+    in_pdb_file,
+    out_pdb_file=None,
+    scheme="chothia",
+):
+    """
+    Renumber the pdb file.
+    """
+    if out_pdb_file is None:
+        out_pdb_file = in_pdb_file
+
+    clean_pdb(in_pdb_file)
+
+    parser = PDBParser()
+    with warnings.catch_warnings(record=True):
+        structure = parser.get_structure(
+            "_",
+            in_pdb_file,
+        )
+
+    for chain in structure.get_chains():
+        seq = seq1(''.join([residue.resname for residue in chain]))
+        abnum_chain = Chain(seq, scheme=scheme)
+        numbering = abnum_chain.positions.items()
+
+        chain_res = list(chain.get_residues())
+        assert len(chain_res) == len(numbering)
+
+        for pdb_r, (pos, aa) in zip(chain_res, numbering):
+            if aa != seq1(pdb_r.get_resname()):
+                raise Exception(f"Failed to renumber PDB file {in_pdb_file}")
+            pos = str(pos)[1:]
+            if not pos[-1].isnumeric():
+                ins = pos[-1]
+                pos = int(pos[:-1])
+            else:
+                pos = int(pos)
+                ins = ' '
+
+            pdb_r._id = (' ', pos, ins)
+
+    io = PDBIO()
+    io.set_structure(structure)
+    io.save(out_pdb_file)
+
+
+def truncate_seq(seq, scheme="chothia"):
+    abnum_chain = Chain(seq, scheme=scheme)
+    numbering = abnum_chain.positions.items()
+    seq = "".join([r[1] for r in list(numbering)])
+
+    return seq

model/utils/constants.py

@@ -0,0 +1,22 @@
+EPS = 1e-6
+
+# Values from Structure quality and target parameters (Engh and Huber)
+
+# Backbone bond lengths
+BL_N_CA = 1.459
+BL_CA_C = 1.525
+BL_C_N = 1.336
+BL_C_O = 1.229
+
+# Backbone bond angles
+BA_N_CA_C = 111.0
+BA_N_CA_CB = 110.6
+BA_CA_C_N = 117.2
+BA_CA_C_O = 120.1
+BA_O_C_N = 122.7
+BA_C_CA_CB = 110.6
+BA_C_N_CA = 121.7
+
+# Van der Waals radii
+VDW_N = 1.55
+VDW_C = 1.7

model/utils/coordinates.py

@@ -0,0 +1,76 @@
+import torch
+
+
+def place_fourth_atom(
+    a_coord: torch.Tensor,
+    b_coord: torch.Tensor,
+    c_coord: torch.Tensor,
+    length: torch.Tensor,
+    planar: torch.Tensor,
+    dihedral: torch.Tensor,
+) -> torch.Tensor:
+    """
+    Given 3 coords + a length + a planar angle + a dihedral angle, compute a fourth coord
+    """
+    bc_vec = b_coord - c_coord
+    bc_vec = bc_vec / bc_vec.norm(dim=-1, keepdim=True)
+
+    n_vec = (b_coord - a_coord).expand(bc_vec.shape).cross(bc_vec)
+    n_vec = n_vec / n_vec.norm(dim=-1, keepdim=True)
+
+    m_vec = [bc_vec, n_vec.cross(bc_vec), n_vec]
+    d_vec = [
+        length * torch.cos(planar),
+        length * torch.sin(planar) * torch.cos(dihedral),
+        -length * torch.sin(planar) * torch.sin(dihedral)
+    ]
+
+    d_coord = c_coord + sum([m * d for m, d in zip(m_vec, d_vec)])
+
+    return d_coord
+
+
+def get_ideal_coords(center=False):
+    N = torch.tensor([[0, 0, -1.458]], dtype=float)
+    A = torch.tensor([[0, 0, 0]], dtype=float)
+    B = torch.tensor([[0, 1.426, 0.531]], dtype=float)
+    C = place_fourth_atom(
+        B,
+        A,
+        N,
+        torch.tensor(2.460),
+        torch.tensor(0.615),
+        torch.tensor(-2.143),
+    )
+
+    coords = torch.cat([N, A, C, B]).float()
+
+    if center:
+        coords -= coords.mean(
+            dim=0,
+            keepdim=True,
+        )
+
+    return coords
+
+
+def place_o_coords(coords):
+    N = coords[:, :, 0]
+    A = coords[:, :, 1]
+    C = coords[:, :, 2]
+
+    o_coords = place_fourth_atom(
+        torch.roll(N, shifts=-1, dims=1),
+        A,
+        C,
+        torch.tensor(1.231),
+        torch.tensor(2.108),
+        torch.tensor(-3.142),
+    ).unsqueeze(2)
+
+    coords = torch.cat(
+        [coords, o_coords],
+        dim=2,
+    )
+
+    return coords

model/utils/embed.py

@@ -0,0 +1,53 @@
+from einops import rearrange
+
+from igfold import IgFoldInput
+from igfold.utils.folding import get_sequence_dict, process_template
+
+
+def embed(
+    antiberty,
+    model,
+    fasta_file=None,
+    sequences=None,
+    template_pdb=None,
+    ignore_cdrs=None,
+    ignore_chain=None,
+    mask=None,
+):
+    seq_dict = get_sequence_dict(
+        sequences,
+        fasta_file,
+    )
+
+    embeddings, attentions = antiberty.embed(
+        seq_dict.values(),
+        return_attention=True,
+    )
+    embeddings = [e[1:-1].unsqueeze(0) for e in embeddings]
+    attentions = [a[:, :, 1:-1, 1:-1].unsqueeze(0) for a in attentions]
+
+    temp_coords, temp_mask = process_template(
+        template_pdb,
+        fasta_file,
+        ignore_cdrs=ignore_cdrs,
+        ignore_chain=ignore_chain,
+    )
+    model_in = IgFoldInput(
+        embeddings=embeddings,
+        attentions=attentions,
+        template_coords=temp_coords,
+        template_mask=temp_mask,
+        return_embeddings=True,
+        batch_mask=mask,
+    )
+
+    model_out = model(model_in)
+
+    prmsd = rearrange(
+        model_out.prmsd,
+        "b (l a) -> b l a",
+        a=4,
+    )
+    model_out.prmsd = prmsd
+
+    return model_out

model/utils/fasta.py

@@ -0,0 +1,43 @@
+from Bio import SeqIO, pairwise2
+
+
+def get_fasta_chain_seq(
+    fasta_file,
+    chain_id,
+):
+    for chain in SeqIO.parse(fasta_file, 'fasta'):
+        if ":{}".format(chain_id) in chain.id:
+            return str(chain.seq)
+
+
+def get_fasta_chain_dict(fasta_file):
+    seq_dict = {}
+    for chain in SeqIO.parse(fasta_file, 'fasta'):
+        seq_dict[chain.id] = str(chain.seq)
+
+    return seq_dict
+
+
+def pairwise_align(
+    seq1,
+    seq2,
+):
+    ###
+    # Aligns two sequences using the Needleman-Wunsch algorithm
+    # Returns alignment of seq2 into seq1
+    ###
+    ali = pairwise2.align.globalxx(
+        seq1,
+        seq2,
+    )[0]
+    ali_list = []
+    seq1_i, seq2_i = 0, 0
+    for ali_seq in ali.seqB.split("-"):
+        if len(ali_seq) == 0:
+            seq1_i += 1
+        else:
+            l = len(ali_seq)
+            ali_list.append((seq1_i, seq1_i + l, seq2_i, seq2_i + l))
+            seq1_i += l
+
+    return ali_list

model/utils/folding.py

@@ -0,0 +1,231 @@
+import os
+from typing import List
+from einops import rearrange
+import torch
+import numpy as np
+
+from igfold.model.interface import IgFoldInput
+from igfold.utils.fasta import get_fasta_chain_dict
+from igfold.utils.general import exists
+from igfold.utils.pdb import get_atom_coords, save_PDB, write_pdb_bfactor, cdr_indices
+
+
+def get_sequence_dict(
+    sequences,
+    fasta_file,
+):
+    if exists(sequences) and exists(fasta_file):
+        print("Both sequences and fasta file provided. Using fasta file.")
+        seq_dict = get_fasta_chain_dict(fasta_file)
+    elif not exists(sequences) and exists(fasta_file):
+        seq_dict = get_fasta_chain_dict(fasta_file)
+    elif exists(sequences):
+        seq_dict = sequences
+    else:
+        exit("Must provide sequences or fasta file.")
+
+    return seq_dict
+
+
+def process_template(
+    pdb_file,
+    fasta_file,
+    ignore_cdrs=None,
+    ignore_chain=None,
+):
+    temp_coords, temp_mask = None, None
+    if exists(pdb_file):
+        temp_coords = get_atom_coords(
+            pdb_file,
+            fasta_file=fasta_file,
+        )
+        temp_coords = torch.stack(
+            [
+                temp_coords['N'], temp_coords['CA'], temp_coords['C'],
+                temp_coords['CB']
+            ],
+            dim=1,
+        ).view(-1, 3).unsqueeze(0)
+
+        temp_mask = torch.ones(temp_coords.shape[:2]).bool()
+        temp_mask[temp_coords.isnan().any(-1)] = False
+        temp_mask[temp_coords.sum(-1) == 0] = False
+
+        if exists(ignore_cdrs):
+            cdr_names = ["h1", "h2", "h3", "l1", "l2", "l3"]
+            if ignore_cdrs == False:
+                cdr_names = []
+            elif isinstance(ignore_cdrs, list):
+                cdr_names = ignore_cdrs
+            elif isinstance(ignore_cdrs, str):
+                cdr_names = [ignore_cdrs]
+
+            for cdr in cdr_names:
+                cdr_range = cdr_indices(pdb_file, cdr)
+                temp_mask[:, (cdr_range[0] - 1) * 4:(cdr_range[1] + 2) *
+                          4] = False
+        if exists(ignore_chain) and ignore_chain in ["H", "L"]:
+            seq_dict = get_fasta_chain_dict(fasta_file)
+            hlen = len(seq_dict["H"])
+            if ignore_chain == "H":
+                temp_mask[:, :hlen * 4] = False
+            elif ignore_chain == "L":
+                temp_mask[:, hlen * 4:] = False
+
+    return temp_coords, temp_mask
+
+
+def process_prediction(
+    model_out,
+    pdb_file,
+    fasta_file,
+    skip_pdb=False,
+    do_refine=True,
+    use_openmm=False,
+    do_renum=False,
+):
+    prmsd = rearrange(
+        model_out.prmsd,
+        "b (l a) -> b l a",
+        a=4,
+    )
+    model_out.prmsd = prmsd
+
+    if skip_pdb:
+        return model_out
+
+    coords = model_out.coords.squeeze(0).detach()
+    res_rmsd = prmsd.square().mean(dim=-1).sqrt().squeeze(0)
+
+    seq_dict = get_fasta_chain_dict(fasta_file)
+    full_seq = "".join(list(seq_dict.values()))
+    chains = list(seq_dict.keys())
+    delims = np.cumsum([len(s) for s in seq_dict.values()]).tolist()
+
+    write_pdb = not do_refine or use_openmm
+    pdb_string = save_PDB(
+        pdb_file,
+        coords,
+        full_seq,
+        chains=chains,
+        atoms=['N', 'CA', 'C', 'CB', 'O'],
+        error=res_rmsd,
+        delim=delims,
+        write_pdb=write_pdb,
+    )
+
+    if do_refine:
+        if use_openmm:
+            try:
+                from igfold.refine.openmm_ref import refine
+                refine_input = [pdb_file]
+            except:
+                exit("OpenMM not installed. Please install OpenMM to use refinement.")
+        else:
+            try:
+                from igfold.refine.pyrosetta_ref import refine
+                refine_input = [pdb_file, pdb_string]
+            except:
+                exit("PyRosetta not installed. Please install PyRosetta to use refinement.")
+
+        refine(*refine_input)
+
+    if do_renum:
+        try:
+            from igfold.utils.abnumber_ import renumber_pdb
+        except:
+            exit("AbNumber not installed. Please install AbNumber to use renumbering.")
+        
+        renumber_pdb(
+            pdb_file,
+            pdb_file,
+        )
+
+    write_pdb_bfactor(
+        pdb_file,
+        pdb_file,
+        bfactor=res_rmsd,
+    )
+
+    return model_out
+
+
+def fold(
+    antiberty,
+    models,
+    pdb_file,
+    fasta_file=None,
+    sequences=None,
+    template_pdb=None,
+    ignore_cdrs=None,
+    ignore_chain=None,
+    skip_pdb=False,
+    do_refine=True,
+    use_openmm=False,
+    do_renum=True,
+    truncate_sequences=False,
+):
+    seq_dict = get_sequence_dict(
+        sequences,
+        fasta_file,
+    )
+
+    if truncate_sequences:
+        try:
+            from igfold.utils.abnumber_ import truncate_seq
+        except:
+            exit("AbNumber not installed. Please install AbNumber to use truncation.")
+
+        seq_dict = {k: truncate_seq(v) for k, v in seq_dict.items()}
+
+    if not exists(fasta_file):
+        fasta_file = pdb_file.replace(".pdb", ".fasta")
+        with open(fasta_file, "w") as f:
+            for chain, seq in seq_dict.items():
+                f.write(">{}\n{}\n".format(
+                    chain,
+                    seq,
+                ))
+
+    embeddings, attentions = antiberty.embed(
+        seq_dict.values(),
+        return_attention=True,
+    )
+    embeddings = [e[1:-1].unsqueeze(0) for e in embeddings]
+    attentions = [a[:, :, 1:-1, 1:-1].unsqueeze(0) for a in attentions]
+
+    temp_coords, temp_mask = process_template(
+        template_pdb,
+        fasta_file,
+        ignore_cdrs=ignore_cdrs,
+        ignore_chain=ignore_chain,
+    )
+    model_in = IgFoldInput(
+        embeddings=embeddings,
+        attentions=attentions,
+        template_coords=temp_coords,
+        template_mask=temp_mask,
+        return_embeddings=True,
+    )
+
+    model_outs, scores = [], []
+    with torch.no_grad():
+        for i, model in enumerate(models):
+            model_out = model(model_in)
+            model_out = model.gradient_refine(model_in, model_out)
+            scores.append(model_out.prmsd.quantile(0.9))
+            model_outs.append(model_out)
+
+    best_model_i = scores.index(min(scores))
+    model_out = model_outs[best_model_i]
+    process_prediction(
+        model_out,
+        pdb_file,
+        fasta_file,
+        skip_pdb=skip_pdb,
+        do_refine=do_refine,
+        use_openmm=use_openmm,
+        do_renum=do_renum,
+    )
+
+    return model_out

model/utils/general.py

@@ -0,0 +1,57 @@
+_aa_dict = {
+    'A': '0',
+    'C': '1',
+    'D': '2',
+    'E': '3',
+    'F': '4',
+    'G': '5',
+    'H': '6',
+    'I': '7',
+    'K': '8',
+    'L': '9',
+    'M': '10',
+    'N': '11',
+    'P': '12',
+    'Q': '13',
+    'R': '14',
+    'S': '15',
+    'T': '16',
+    'V': '17',
+    'W': '18',
+    'Y': '19'
+}
+
+_aa_1_3_dict = {
+    'A': 'ALA',
+    'C': 'CYS',
+    'D': 'ASP',
+    'E': 'GLU',
+    'F': 'PHE',
+    'G': 'GLY',
+    'H': 'HIS',
+    'I': 'ILE',
+    'K': 'LYS',
+    'L': 'LEU',
+    'M': 'MET',
+    'N': 'ASN',
+    'P': 'PRO',
+    'Q': 'GLN',
+    'R': 'ARG',
+    'S': 'SER',
+    'T': 'THR',
+    'V': 'VAL',
+    'W': 'TRP',
+    'Y': 'TYR',
+    '-': 'GAP'
+}
+
+
+def exists(x):
+    return x is not None
+
+
+def default(
+    val,
+    d,
+):
+    return val if exists(val) else d

model/utils/geometry.py

@@ -0,0 +1,77 @@
+import torch
+from einops import rearrange
+
+from igfold.utils.constants import EPS
+
+
+def normed_vec(vec, eps=EPS):
+    mag_sq = torch.sum(vec**2, dim=-1, keepdim=True)
+    mag = torch.sqrt(mag_sq + eps)
+    vec = vec / mag
+
+    return vec
+
+
+def normed_cross(vec1, vec2, eps=EPS):
+    vec1 = normed_vec(vec1, eps=eps)
+    vec2 = normed_vec(vec2, eps=eps)
+    cross = torch.cross(vec1, vec2, dim=-1)
+
+    return cross
+
+
+def dist(x_1, x_2, eps=EPS):
+    d_sq = (x_1 - x_2)**2
+    d = torch.sqrt(d_sq.sum(-1) + eps)
+
+    return d
+
+
+def dist_mat(c1, c2, dim=-3, eps=EPS):
+    c1 = c1.unsqueeze(dim)
+    c2 = c2.unsqueeze(dim - 1)
+    d = dist(c1, c2, eps=eps)
+
+    return d
+
+
+def angle(x_1, x_2, x_3, eps=EPS):
+    a = normed_vec(x_1 - x_2, eps=eps)
+    b = normed_vec(x_3 - x_2, eps=eps)
+    ang = torch.arccos((a * b).sum(-1))
+
+    return ang
+
+
+def dihedral(x_1, x_2, x_3, x_4, eps=EPS):
+    b1 = normed_vec(x_1 - x_2, eps=eps)
+    b2 = normed_vec(x_2 - x_3, eps=eps)
+    b3 = normed_vec(x_3 - x_4, eps=eps)
+    n1 = normed_cross(b1, b2, eps=eps)
+    n2 = normed_cross(b2, b3, eps=eps)
+    m1 = normed_cross(n1, b2, eps=eps)
+    x = (n1 * n2).sum(-1)
+    y = (m1 * n2).sum(-1)
+
+    dih = torch.atan2(y, x)
+
+    return dih
+
+
+def coords_to_frame(coords, eps=EPS):
+    if len(coords.shape) == 3:
+        coords = rearrange(
+            coords,
+            "b (l a) d -> b l a d",
+            l=coords.shape[-2] // 4,
+        )
+
+    N, CA, C, _ = coords.unbind(-2)
+    CA_N = normed_vec(N - CA, eps=eps)
+    CA_C = normed_vec(C - CA, eps=eps)
+    n1 = CA_N
+    n2 = normed_cross(n1, CA_C, eps=eps)
+    n3 = normed_cross(n1, n2, eps=eps)
+    rot = torch.stack([n1, n2, n3], -1)
+
+    return CA, rot

model/utils/pdb.py

@@ -0,0 +1,515 @@
+import time
+import sys
+import io
+from typing import Union, List
+import requests
+import warnings
+from os.path import splitext, basename
+from Bio import PDB
+from Bio.PDB import PDBParser, PDBIO
+from Bio.SeqUtils import seq1
+from Bio import SeqIO
+from bisect import bisect_left, bisect_right
+import torch
+import numpy as np
+
+from igfold.utils.coordinates import place_fourth_atom
+from igfold.utils.fasta import get_fasta_chain_seq
+from igfold.utils.general import _aa_1_3_dict, exists
+
+
+def renumber_pdb(old_pdb, renum_pdb=None):
+    if not exists(renum_pdb):
+        renum_pdb = old_pdb
+
+    success = False
+    time.sleep(5)
+    for i in range(10):
+        try:
+            with open(old_pdb, 'rb') as f:
+                response = requests.post(
+                    'http://www.bioinf.org.uk/abs/abnum/abnumpdb.cgi',
+                    params={
+                        "plain": "1",
+                        "output": "-HL",
+                        "scheme": "-c"
+                    },
+                    files={"pdb": f},
+                )
+
+            success = response.status_code == 200 and not ("<html>"
+                                                           in response.text)
+
+            if success:
+                break
+            else:
+                time.sleep((i + 1) * 5)
+        except requests.exceptions.ConnectionError:
+            time.sleep(60)
+
+    # if success:
+    if success:
+        new_pdb_data = response.text
+        with open(renum_pdb, "w") as f:
+            f.write(new_pdb_data)
+    else:
+        print(
+            "Failed to renumber PDB. This is likely due to a connection error or a timeout with the AbNum server."
+        )
+
+
+def count_pdb_chains(pdb_file):
+    parser = PDBParser()
+    with warnings.catch_warnings(record=True):
+        structure = parser.get_structure("_", pdb_file)
+
+    l = len(list(structure.get_chains()))
+
+    return l
+
+
+def reorder_pdb_chains(pdb_file, chain_order):
+    """Reorder the chains in a PDB file and update residue numbers"""
+
+    parser = PDBParser()
+    with warnings.catch_warnings(record=True):
+        structure = parser.get_structure("_", pdb_file)
+
+    chains = list(structure.get_chains())
+    if len(chains) != len(chain_order):
+        raise ValueError(
+            f"Number of chains in PDB file ({len(chains)}) does not match number of chains in chain order ({len(chain_order)})"
+        )
+
+    chain_order = [c.upper() for c in chain_order]
+    sorted_chains = sorted(chains, key=lambda c: chain_order.index(c.id))
+
+    new_structure = PDB.Structure.Structure("_")
+    new_model = PDB.Model.Model(0)
+    new_structure.add(new_model)
+    atom_num = 1
+    for chain in sorted_chains:
+        new_chain = PDB.Chain.Chain(chain.id)
+        new_model.add(new_chain)
+        for residue in chain.get_residues():
+            new_residue = PDB.Residue.Residue(
+                residue.id,
+                residue.resname,
+                residue.segid,
+            )
+            new_chain.add(new_residue)
+            for atom in residue:
+                new_atom = PDB.Atom.Atom(
+                    atom.name,
+                    atom.coord,
+                    atom.occupancy,
+                    atom.bfactor,
+                    atom.altloc,
+                    atom.fullname,
+                    atom_num,
+                    atom.element,
+                )
+                new_residue.add(new_atom)
+                atom_num += 1
+
+    io = PDBIO()
+    io.set_structure(new_structure)
+    io.save(pdb_file)
+
+
+def get_atom_coord(residue, atom_type):
+    if exists(residue) and atom_type in residue:
+        return residue[atom_type].get_coord()
+    else:
+        return [0, 0, 0]
+
+
+def get_cb_or_ca_coord(residue):
+    if not exists(residue):
+        return [0, 0, 0]
+
+    if 'CB' in residue:
+        return residue['CB'].get_coord()
+    elif 'CA' in residue:
+        return residue['CA'].get_coord()
+    else:
+        return [0, 0, 0]
+
+
+def get_continuous_ranges(residues):
+    """ Returns ranges of residues which are continuously connected (peptide bond length 1.2-1.45 Å) """
+    dists = []
+    for res_i in range(len(residues) - 1):
+        dists.append(
+            np.linalg.norm(
+                np.array(get_atom_coord(residues[res_i], "C")) -
+                np.array(get_atom_coord(residues[res_i + 1], "N"))))
+
+    ranges = []
+    start_i = 0
+    for d_i, d in enumerate(dists):
+        if d > 1.45 or d < 1.2:
+            ranges.append((start_i, d_i + 1))
+            start_i = d_i + 1
+        if d_i == len(dists) - 1:
+            ranges.append((start_i, None))
+
+    return ranges
+
+
+def place_missing_cb_o(atom_coords):
+    cb_coords = place_fourth_atom(
+        atom_coords['C'],
+        atom_coords['N'],
+        atom_coords['CA'],
+        torch.tensor(1.522),
+        torch.tensor(1.927),
+        torch.tensor(-2.143),
+    )
+    o_coords = place_fourth_atom(
+        torch.roll(atom_coords['N'], shifts=-1, dims=0),
+        atom_coords['CA'],
+        atom_coords['C'],
+        torch.tensor(1.231),
+        torch.tensor(2.108),
+        torch.tensor(-3.142),
+    )
+
+    bb_mask = get_atom_coords_mask(atom_coords['N']) & get_atom_coords_mask(
+        atom_coords['CA']) & get_atom_coords_mask(atom_coords['C'])
+    missing_cb = (get_atom_coords_mask(atom_coords['CB']) & bb_mask) == 0
+    atom_coords['CB'][missing_cb] = cb_coords[missing_cb]
+
+    bb_mask = get_atom_coords_mask(
+        torch.roll(
+            atom_coords['N'],
+            shifts=-1,
+            dims=0,
+        )) & get_atom_coords_mask(atom_coords['CA']) & get_atom_coords_mask(
+            atom_coords['C'])
+    missing_o = (get_atom_coords_mask(atom_coords['O']) & bb_mask) == 0
+    atom_coords['O'][missing_o] = o_coords[missing_o]
+
+
+def get_atom_coords(pdb_file, fasta_file=None):
+    p = PDBParser()
+    file_name = splitext(basename(pdb_file))[0]
+    structure = p.get_structure(
+        file_name,
+        pdb_file,
+    )
+
+    if fasta_file:
+        residues = []
+        for chain in structure.get_chains():
+            pdb_seq = get_pdb_chain_seq(
+                pdb_file,
+                chain.id,
+            )
+
+            chain_dict = {"A": "H", "B": "L", "H": "H", "L": "L"}
+            fasta_seq = get_fasta_chain_seq(
+                fasta_file,
+                chain_dict[chain.id],
+            )
+
+            chain_residues = list(chain.get_residues())
+            continuous_ranges = get_continuous_ranges(chain_residues)
+
+            fasta_residues = [None for _ in range(len(fasta_seq))]
+            fasta_r = (0, 0)
+            for pdb_r in continuous_ranges:
+                fasta_r_start = fasta_seq[fasta_r[1]:].index(
+                    pdb_seq[pdb_r[0]:pdb_r[1]]) + fasta_r[1]
+                fasta_r_end = (len(pdb_seq) if pdb_r[1] == None else
+                               pdb_r[1]) - pdb_r[0] + fasta_r_start
+                fasta_r = (fasta_r_start, fasta_r_end)
+                fasta_residues[fasta_r[0]:fasta_r[1]] = chain_residues[
+                    pdb_r[0]:pdb_r[1]]
+
+            residues += fasta_residues
+    else:
+        residues = list(structure.get_residues())
+
+    n_coords = torch.tensor([get_atom_coord(r, 'N') for r in residues])
+    ca_coords = torch.tensor([get_atom_coord(r, 'CA') for r in residues])
+    c_coords = torch.tensor([get_atom_coord(r, 'C') for r in residues])
+    cb_coords = torch.tensor([get_atom_coord(r, 'CB') for r in residues])
+    cb_ca_coords = torch.tensor([get_cb_or_ca_coord(r) for r in residues])
+    o_coords = torch.tensor([get_atom_coord(r, 'O') for r in residues])
+
+    atom_coords = {}
+    atom_coords['N'] = n_coords
+    atom_coords['CA'] = ca_coords
+    atom_coords['C'] = c_coords
+    atom_coords['CB'] = cb_coords
+    atom_coords['CBCA'] = cb_ca_coords
+    atom_coords['O'] = o_coords
+
+    place_missing_cb_o(atom_coords)
+
+    return atom_coords
+
+
+def get_atom_coords_mask(coords):
+    mask = torch.ByteTensor([1 if sum(_) != 0 else 0 for _ in coords])
+    mask = mask & (1 - torch.any(torch.isnan(coords), dim=1).byte())
+    return mask
+
+
+def get_atom_coords_mask_for_dict(atom_coords):
+    atom_coords_masks = {}
+    for atom, coords in atom_coords.items():
+        atom_coords_masks[atom] = get_atom_coords_mask(coords)
+
+    return atom_coords_masks
+
+
+def pdb2fasta(pdb_file, num_chains=None):
+    """Converts a PDB file to a fasta formatted string using its ATOM data"""
+    pdb_id = basename(pdb_file).split('.')[0]
+    parser = PDBParser()
+    structure = parser.get_structure(
+        pdb_id,
+        pdb_file,
+    )
+
+    real_num_chains = len([0 for _ in structure.get_chains()])
+    if num_chains is not None and num_chains != real_num_chains:
+        print('WARNING: Skipping {}. Expected {} chains, got {}'.format(
+            pdb_file, num_chains, real_num_chains))
+        return ''
+
+    fasta = ''
+    for chain in structure.get_chains():
+        id_ = chain.id
+        seq = seq1(''.join([residue.resname for residue in chain]))
+        fasta += '>{}:{}\t{}\n'.format(pdb_id, id_, len(seq))
+        max_line_length = 80
+        for i in range(0, len(seq), max_line_length):
+            fasta += f'{seq[i:i + max_line_length]}\n'
+    return fasta
+
+
+def get_pdb_chain_seq(
+    pdb_file,
+    chain_id,
+):
+    p = PDBParser()
+    file_name = splitext(basename(pdb_file))[0]
+    structure = p.get_structure(
+        file_name,
+        pdb_file,
+    )
+
+    pdb_seq = None
+    for chain in structure.get_chains():
+        if chain.id == chain_id:
+            pdb_seq = "".join(
+                [seq1(r.get_resname()) for r in chain.get_residues()])
+
+    return pdb_seq
+
+
+def cdr_indices(
+    chothia_pdb_file,
+    cdr,
+    offset_heavy=True,
+):
+    """Gets the index of a given CDR loop"""
+    cdr_chothia_range_dict = {
+        "h1": (26, 32),
+        "h2": (52, 56),
+        "h3": (95, 102),
+        "l1": (24, 34),
+        "l2": (50, 56),
+        "l3": (89, 97)
+    }
+
+    cdr = str.lower(cdr)
+    assert cdr in cdr_chothia_range_dict.keys()
+
+    chothia_range = cdr_chothia_range_dict[cdr]
+    chain_id = cdr[0].upper()
+
+    parser = PDBParser()
+    pdb_id = basename(chothia_pdb_file).split('.')[0]
+    structure = parser.get_structure(
+        pdb_id,
+        chothia_pdb_file,
+    )
+    cdr_chain_structure = None
+    for chain in structure.get_chains():
+        if chain.id == chain_id:
+            cdr_chain_structure = chain
+            break
+    if cdr_chain_structure is None:
+        print("PDB must have a chain with chain id \"[PBD ID]:{}\"".format(
+            chain_id))
+        sys.exit(-1)
+
+    residue_id_nums = [res.get_id()[1] for res in cdr_chain_structure]
+
+    # Binary search to find the start and end of the CDR loop
+    cdr_start = bisect_left(
+        residue_id_nums,
+        chothia_range[0],
+    )
+    cdr_end = bisect_right(
+        residue_id_nums,
+        chothia_range[1],
+    ) - 1
+
+    if len(get_pdb_chain_seq(
+            chothia_pdb_file,
+            chain_id=chain_id,
+    )) != len(residue_id_nums):
+        print('ERROR in PDB file ' + chothia_pdb_file)
+        print('residue id len', len(residue_id_nums))
+
+    if chain_id == "L" and offset_heavy:
+        heavy_seq_len = get_pdb_chain_seq(
+            chothia_pdb_file,
+            chain_id="H",
+        )
+        cdr_start += len(heavy_seq_len)
+        cdr_end += len(heavy_seq_len)
+
+    return cdr_start, cdr_end
+
+
+def get_cdr_range_dict(
+    chothia_pdb_file,
+    heavy_only=False,
+    light_only=False,
+    offset_heavy=True,
+):
+    cdr_names = ["h1", "h2", "h3", "l1", "l2", "l3"]
+    if heavy_only:
+        cdr_names = cdr_names[:3]
+    if light_only:
+        cdr_names = cdr_names[3:]
+
+    cdr_range_dict = {
+        cdr: cdr_indices(
+            chothia_pdb_file,
+            cdr,
+            offset_heavy=offset_heavy,
+        )
+        for cdr in cdr_names
+    }
+
+    return cdr_range_dict
+
+
+def h3_indices(chothia_pdb_file):
+    """Gets the index of the CDR H3 loop"""
+
+    return cdr_indices(chothia_pdb_file, cdr="h3")
+
+
+def get_chain_numbering(
+    pdb_file,
+    chain_id,
+):
+    seq = []
+    parser = PDBParser()
+    structure = parser.get_structure("_", pdb_file)
+    for chain in structure.get_chains():
+        if chain.id == chain_id:
+            for r in chain.get_residues():
+                res_num = str(r._id[1]) + r._id[2]
+                res_num = res_num.replace(" ", "")
+                seq.append(res_num)
+
+            return seq
+
+
+def save_PDB(
+    out_pdb: str,
+    coords: torch.Tensor,
+    seq: str,
+    chains: List[str] = None,
+    error: torch.Tensor = None,
+    delim: Union[int, List[int]] = None,
+    atoms=['N', 'CA', 'C', 'O', 'CB'],
+    write_pdb=True,
+) -> None:
+    """
+    Write set of N, CA, C, O, CB coords to PDB file
+    """
+
+    if not exists(chains):
+        chains = ["H", "L"]
+
+    if type(delim) == type(None):
+        delim = -1
+    elif type(delim) == int:
+        delim = [delim]
+
+    if not exists(error):
+        error = torch.zeros(len(seq))
+
+    pdb_string = ""
+    k = 0
+    for r, residue in enumerate(coords):
+        AA = _aa_1_3_dict[seq[r]]
+        for a, atom in enumerate(residue):
+            chain_num = np.where(np.array(delim) - r > 0)[0][0]
+            chain_id = chains[chain_num]
+
+            if AA == "GLY" and atoms[a] == "CB": continue
+            x, y, z = atom
+            pdb_string += "ATOM  %5d  %-2s  %3s %s%4d    %8.3f%8.3f%8.3f  %4.2f  %4.2f           %s  \n" % (
+                k + 1, atoms[a], AA, chain_id, r + 1, x, y, z, 1, error[r], atoms[a][0])
+            k += 1
+
+        if r + 1 == delim[chain_num]:
+            pdb_string += "TER   %5d      %3s %s%4d\n" % (
+                k + 1, AA, chain_id, r + 1)
+            k += 1
+                
+    pdb_string += "END\n"
+
+    if write_pdb:
+        with open(out_pdb, "w") as f:
+            f.write(pdb_string)
+
+    return pdb_string
+
+
+def write_pdb_bfactor(
+    in_pdb_file,
+    out_pdb_file,
+    bfactor,
+    b_chain=None,
+):
+    parser = PDBParser()
+    with warnings.catch_warnings(record=True):
+        structure = parser.get_structure(
+            "_",
+            in_pdb_file,
+        )
+
+    i = 0
+    for chain in structure.get_chains():
+        if exists(b_chain) and chain._id != b_chain:
+            continue
+
+        for r in chain.get_residues():
+            [a.set_bfactor(bfactor[i]) for a in r.get_atoms()]
+            i += 1
+
+    io = PDBIO()
+    io.set_structure(structure)
+    io.save(out_pdb_file)
+
+
+def clean_pdb(pdb_file):
+    with open(pdb_file, "r") as f:
+        lines = f.readlines()
+
+    with open(pdb_file, "w") as f:
+        for l in lines:
+            if "ATOM" in l:
+                f.write(l)

model/utils/tensor.py

@@ -0,0 +1,42 @@
+import torch
+import torch.nn.functional as F
+
+
+def max_shape(data):
+    """Gets the maximum length along all dimensions in a list of Tensors"""
+    shapes = torch.Tensor([_.shape for _ in data])
+    return torch.max(
+        shapes.transpose(0, 1),
+        dim=1,
+    )[0].int()
+
+
+def pad_data_to_same_shape(
+    tensor_list,
+    pad_value=0,
+):
+    target_shape = max_shape(tensor_list)
+
+    padded_dataset_shape = [len(tensor_list)] + list(target_shape)
+    padded_dataset = torch.Tensor(*padded_dataset_shape).type_as(
+        tensor_list[0])
+
+    for i, data in enumerate(tensor_list):
+        # Get how much padding is needed per dimension
+        padding = reversed(target_shape - torch.Tensor(list(data.shape)).int())
+
+        # Add 0 every other index to indicate only right padding
+        padding = F.pad(
+            padding.unsqueeze(0).t(),
+            (1, 0, 0, 0),
+        ).view(-1, 1)
+        padding = padding.view(1, -1)[0].tolist()
+
+        padded_data = F.pad(
+            data,
+            padding,
+            value=pad_value,
+        )
+        padded_dataset[i] = padded_data
+
+    return padded_dataset

model/utils/transforms.py

@@ -0,0 +1,520 @@
+# Copyright (c) Facebook, Inc. and its affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the BSD-style license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import torch
+import torch.nn.functional as F
+"""
+The transformation matrices returned from the functions in this file assume
+the points on which the transformation will be applied are column vectors.
+i.e. the R matrix is structured as
+
+    R = [
+            [Rxx, Rxy, Rxz],
+            [Ryx, Ryy, Ryz],
+            [Rzx, Rzy, Rzz],
+        ]  # (3, 3)
+
+This matrix can be applied to column vectors by post multiplication
+by the points e.g.
+
+    points = [[0], [1], [2]]  # (3 x 1) xyz coordinates of a point
+    transformed_points = R * points
+
+To apply the same matrix to points which are row vectors, the R matrix
+can be transposed and pre multiplied by the points:
+
+e.g.
+    points = [[0, 1, 2]]  # (1 x 3) xyz coordinates of a point
+    transformed_points = points * R.transpose(1, 0)
+"""
+
+
+def quaternion_to_matrix(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as quaternions to rotation matrices.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    r, i, j, k = torch.unbind(quaternions, -1)
+    two_s = 2.0 / (quaternions * quaternions).sum(-1)
+
+    o = torch.stack(
+        (
+            1 - two_s * (j * j + k * k),
+            two_s * (i * j - k * r),
+            two_s * (i * k + j * r),
+            two_s * (i * j + k * r),
+            1 - two_s * (i * i + k * k),
+            two_s * (j * k - i * r),
+            two_s * (i * k - j * r),
+            two_s * (j * k + i * r),
+            1 - two_s * (i * i + j * j),
+        ),
+        -1,
+    )
+    return o.reshape(quaternions.shape[:-1] + (3, 3))
+
+
+def _copysign(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Return a tensor where each element has the absolute value taken from the,
+    corresponding element of a, with sign taken from the corresponding
+    element of b. This is like the standard copysign floating-point operation,
+    but is not careful about negative 0 and NaN.
+
+    Args:
+        a: source tensor.
+        b: tensor whose signs will be used, of the same shape as a.
+
+    Returns:
+        Tensor of the same shape as a with the signs of b.
+    """
+    signs_differ = (a < 0) != (b < 0)
+    return torch.where(signs_differ, -a, a)
+
+
+def _sqrt_positive_part(x: torch.Tensor) -> torch.Tensor:
+    """
+    Returns torch.sqrt(torch.max(0, x))
+    but with a zero subgradient where x is 0.
+    """
+    ret = torch.zeros_like(x)
+    positive_mask = x > 0
+    ret[positive_mask] = torch.sqrt(x[positive_mask])
+    return ret
+
+
+def matrix_to_quaternion(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to quaternions.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+
+    batch_dim = matrix.shape[:-2]
+    m00, m01, m02, m10, m11, m12, m20, m21, m22 = torch.unbind(
+        matrix.reshape(batch_dim + (9, )), dim=-1)
+
+    q_abs = _sqrt_positive_part(
+        torch.stack(
+            [
+                1.0 + m00 + m11 + m22,
+                1.0 + m00 - m11 - m22,
+                1.0 - m00 + m11 - m22,
+                1.0 - m00 - m11 + m22,
+            ],
+            dim=-1,
+        ))
+
+    # we produce the desired quaternion multiplied by each of r, i, j, k
+    quat_by_rijk = torch.stack(
+        [
+            torch.stack([q_abs[..., 0]**2, m21 - m12, m02 - m20, m10 - m01],
+                        dim=-1),
+            torch.stack([m21 - m12, q_abs[..., 1]**2, m10 + m01, m02 + m20],
+                        dim=-1),
+            torch.stack([m02 - m20, m10 + m01, q_abs[..., 2]**2, m12 + m21],
+                        dim=-1),
+            torch.stack([m10 - m01, m20 + m02, m21 + m12, q_abs[..., 3]**2],
+                        dim=-1),
+        ],
+        dim=-2,
+    )
+
+    # We floor here at 0.1 but the exact level is not important; if q_abs is small,
+    # the candidate won't be picked.
+    flr = torch.tensor(0.1).to(dtype=q_abs.dtype, device=q_abs.device)
+    quat_candidates = quat_by_rijk / (2.0 * q_abs[..., None].max(flr))
+
+    # if not for numerical problems, quat_candidates[i] should be same (up to a sign),
+    # forall i; we pick the best-conditioned one (with the largest denominator)
+
+    return quat_candidates[F.one_hot(q_abs.argmax(
+        dim=-1), num_classes=4) > 0.5, :  # pyre-ignore[16]
+                           ].reshape(batch_dim + (4, ))
+
+
+def _axis_angle_rotation(axis: str, angle: torch.Tensor) -> torch.Tensor:
+    """
+    Return the rotation matrices for one of the rotations about an axis
+    of which Euler angles describe, for each value of the angle given.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z".
+        angle: any shape tensor of Euler angles in radians
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+
+    cos = torch.cos(angle)
+    sin = torch.sin(angle)
+    one = torch.ones_like(angle)
+    zero = torch.zeros_like(angle)
+
+    if axis == "X":
+        R_flat = (one, zero, zero, zero, cos, -sin, zero, sin, cos)
+    elif axis == "Y":
+        R_flat = (cos, zero, sin, zero, one, zero, -sin, zero, cos)
+    elif axis == "Z":
+        R_flat = (cos, -sin, zero, sin, cos, zero, zero, zero, one)
+    else:
+        raise ValueError("letter must be either X, Y or Z.")
+
+    return torch.stack(R_flat, -1).reshape(angle.shape + (3, 3))
+
+
+def euler_angles_to_matrix(euler_angles: torch.Tensor,
+                           convention: str) -> torch.Tensor:
+    """
+    Convert rotations given as Euler angles in radians to rotation matrices.
+
+    Args:
+        euler_angles: Euler angles in radians as tensor of shape (..., 3).
+        convention: Convention string of three uppercase letters from
+            {"X", "Y", and "Z"}.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    if euler_angles.dim() == 0 or euler_angles.shape[-1] != 3:
+        raise ValueError("Invalid input euler angles.")
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    matrices = [
+        _axis_angle_rotation(c, e)
+        for c, e in zip(convention, torch.unbind(euler_angles, -1))
+    ]
+    # return functools.reduce(torch.matmul, matrices)
+    return torch.matmul(torch.matmul(matrices[0], matrices[1]), matrices[2])
+
+
+def _angle_from_tan(axis: str, other_axis: str, data, horizontal: bool,
+                    tait_bryan: bool) -> torch.Tensor:
+    """
+    Extract the first or third Euler angle from the two members of
+    the matrix which are positive constant times its sine and cosine.
+
+    Args:
+        axis: Axis label "X" or "Y or "Z" for the angle we are finding.
+        other_axis: Axis label "X" or "Y or "Z" for the middle axis in the
+            convention.
+        data: Rotation matrices as tensor of shape (..., 3, 3).
+        horizontal: Whether we are looking for the angle for the third axis,
+            which means the relevant entries are in the same row of the
+            rotation matrix. If not, they are in the same column.
+        tait_bryan: Whether the first and third axes in the convention differ.
+
+    Returns:
+        Euler Angles in radians for each matrix in data as a tensor
+        of shape (...).
+    """
+
+    i1, i2 = {"X": (2, 1), "Y": (0, 2), "Z": (1, 0)}[axis]
+    if horizontal:
+        i2, i1 = i1, i2
+    even = (axis + other_axis) in ["XY", "YZ", "ZX"]
+    if horizontal == even:
+        return torch.atan2(data[..., i1], data[..., i2])
+    if tait_bryan:
+        return torch.atan2(-data[..., i2], data[..., i1])
+    return torch.atan2(data[..., i2], -data[..., i1])
+
+
+def _index_from_letter(letter: str) -> int:
+    if letter == "X":
+        return 0
+    if letter == "Y":
+        return 1
+    if letter == "Z":
+        return 2
+    raise ValueError("letter must be either X, Y or Z.")
+
+
+def matrix_to_euler_angles(matrix: torch.Tensor,
+                           convention: str) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to Euler angles in radians.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+        convention: Convention string of three uppercase letters.
+
+    Returns:
+        Euler angles in radians as tensor of shape (..., 3).
+    """
+    if len(convention) != 3:
+        raise ValueError("Convention must have 3 letters.")
+    if convention[1] in (convention[0], convention[2]):
+        raise ValueError(f"Invalid convention {convention}.")
+    for letter in convention:
+        if letter not in ("X", "Y", "Z"):
+            raise ValueError(f"Invalid letter {letter} in convention string.")
+    if matrix.size(-1) != 3 or matrix.size(-2) != 3:
+        raise ValueError(f"Invalid rotation matrix shape {matrix.shape}.")
+    i0 = _index_from_letter(convention[0])
+    i2 = _index_from_letter(convention[2])
+    tait_bryan = i0 != i2
+    if tait_bryan:
+        central_angle = torch.asin(matrix[..., i0, i2] *
+                                   (-1.0 if i0 - i2 in [-1, 2] else 1.0))
+    else:
+        central_angle = torch.acos(matrix[..., i0, i0])
+
+    o = (
+        _angle_from_tan(convention[0], convention[1], matrix[..., i2], False,
+                        tait_bryan),
+        central_angle,
+        _angle_from_tan(convention[2], convention[1], matrix[..., i0, :], True,
+                        tait_bryan),
+    )
+    return torch.stack(o, -1)
+
+
+def standardize_quaternion(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert a unit quaternion to a standard form: one in which the real
+    part is non negative.
+
+    Args:
+        quaternions: Quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Standardized quaternions as tensor of shape (..., 4).
+    """
+    return torch.where(quaternions[..., 0:1] < 0, -quaternions, quaternions)
+
+
+def quaternion_raw_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two quaternions.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions shape (..., 4).
+    """
+    aw, ax, ay, az = torch.unbind(a, -1)
+    bw, bx, by, bz = torch.unbind(b, -1)
+    ow = aw * bw - ax * bx - ay * by - az * bz
+    ox = aw * bx + ax * bw + ay * bz - az * by
+    oy = aw * by - ax * bz + ay * bw + az * bx
+    oz = aw * bz + ax * by - ay * bx + az * bw
+    return torch.stack((ow, ox, oy, oz), -1)
+
+
+def quaternion_multiply(a: torch.Tensor, b: torch.Tensor) -> torch.Tensor:
+    """
+    Multiply two quaternions representing rotations, returning the quaternion
+    representing their composition, i.e. the versor with nonnegative real part.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        a: Quaternions as tensor of shape (..., 4), real part first.
+        b: Quaternions as tensor of shape (..., 4), real part first.
+
+    Returns:
+        The product of a and b, a tensor of quaternions of shape (..., 4).
+    """
+    ab = quaternion_raw_multiply(a, b)
+    return standardize_quaternion(ab)
+
+
+def quaternion_invert(quaternion: torch.Tensor) -> torch.Tensor:
+    """
+    Given a quaternion representing rotation, get the quaternion representing
+    its inverse.
+
+    Args:
+        quaternion: Quaternions as tensor of shape (..., 4), with real part
+            first, which must be versors (unit quaternions).
+
+    Returns:
+        The inverse, a tensor of quaternions of shape (..., 4).
+    """
+
+    scaling = torch.tensor([1, -1, -1, -1], device=quaternion.device)
+    return quaternion * scaling
+
+
+def quaternion_apply(quaternion: torch.Tensor,
+                     point: torch.Tensor) -> torch.Tensor:
+    """
+    Apply the rotation given by a quaternion to a 3D point.
+    Usual torch rules for broadcasting apply.
+
+    Args:
+        quaternion: Tensor of quaternions, real part first, of shape (..., 4).
+        point: Tensor of 3D points of shape (..., 3).
+
+    Returns:
+        Tensor of rotated points of shape (..., 3).
+    """
+    if point.size(-1) != 3:
+        raise ValueError(f"Points are not in 3D, {point.shape}.")
+    real_parts = point.new_zeros(point.shape[:-1] + (1, ))
+    point_as_quaternion = torch.cat((real_parts, point), -1)
+    out = quaternion_raw_multiply(
+        quaternion_raw_multiply(quaternion, point_as_quaternion),
+        quaternion_invert(quaternion),
+    )
+    return out[..., 1:]
+
+
+def axis_angle_to_matrix(axis_angle: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as axis/angle to rotation matrices.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        Rotation matrices as tensor of shape (..., 3, 3).
+    """
+    return quaternion_to_matrix(axis_angle_to_quaternion(axis_angle))
+
+
+def matrix_to_axis_angle(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as rotation matrices to axis/angle.
+
+    Args:
+        matrix: Rotation matrices as tensor of shape (..., 3, 3).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    return quaternion_to_axis_angle(matrix_to_quaternion(matrix))
+
+
+def axis_angle_to_quaternion(axis_angle: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as axis/angle to quaternions.
+
+    Args:
+        axis_angle: Rotations given as a vector in axis angle form,
+            as a tensor of shape (..., 3), where the magnitude is
+            the angle turned anticlockwise in radians around the
+            vector's direction.
+
+    Returns:
+        quaternions with real part first, as tensor of shape (..., 4).
+    """
+    angles = torch.norm(axis_angle, p=2, dim=-1, keepdim=True)
+    half_angles = angles * 0.5
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles])
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48)
+    quaternions = torch.cat(
+        [torch.cos(half_angles), axis_angle * sin_half_angles_over_angles],
+        dim=-1)
+    return quaternions
+
+
+def quaternion_to_axis_angle(quaternions: torch.Tensor) -> torch.Tensor:
+    """
+    Convert rotations given as quaternions to axis/angle.
+
+    Args:
+        quaternions: quaternions with real part first,
+            as tensor of shape (..., 4).
+
+    Returns:
+        Rotations given as a vector in axis angle form, as a tensor
+            of shape (..., 3), where the magnitude is the angle
+            turned anticlockwise in radians around the vector's
+            direction.
+    """
+    norms = torch.norm(quaternions[..., 1:], p=2, dim=-1, keepdim=True)
+    half_angles = torch.atan2(norms, quaternions[..., :1])
+    angles = 2 * half_angles
+    eps = 1e-6
+    small_angles = angles.abs() < eps
+    sin_half_angles_over_angles = torch.empty_like(angles)
+    sin_half_angles_over_angles[~small_angles] = (
+        torch.sin(half_angles[~small_angles]) / angles[~small_angles])
+    # for x small, sin(x/2) is about x/2 - (x/2)^3/6
+    # so sin(x/2)/x is about 1/2 - (x*x)/48
+    sin_half_angles_over_angles[small_angles] = (
+        0.5 - (angles[small_angles] * angles[small_angles]) / 48)
+    return quaternions[..., 1:] / sin_half_angles_over_angles
+
+
+def rotation_6d_to_matrix(d6: torch.Tensor) -> torch.Tensor:
+    """
+    Converts 6D rotation representation by Zhou et al. [1] to rotation matrix
+    using Gram--Schmidt orthogonalization per Section B of [1].
+    Args:
+        d6: 6D rotation representation, of size (*, 6)
+
+    Returns:
+        batch of rotation matrices of size (*, 3, 3)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+
+    a1, a2 = d6[..., :3], d6[..., 3:]
+    b1 = F.normalize(a1, dim=-1)
+    b2 = a2 - (b1 * a2).sum(-1, keepdim=True) * b1
+    b2 = F.normalize(b2, dim=-1)
+    b3 = torch.cross(b1, b2, dim=-1)
+    return torch.stack((b1, b2, b3), dim=-2)
+
+
+def matrix_to_rotation_6d(matrix: torch.Tensor) -> torch.Tensor:
+    """
+    Converts rotation matrices to 6D rotation representation by Zhou et al. [1]
+    by dropping the last row. Note that 6D representation is not unique.
+    Args:
+        matrix: batch of rotation matrices of size (*, 3, 3)
+
+    Returns:
+        6D rotation representation, of size (*, 6)
+
+    [1] Zhou, Y., Barnes, C., Lu, J., Yang, J., & Li, H.
+    On the Continuity of Rotation Representations in Neural Networks.
+    IEEE Conference on Computer Vision and Pattern Recognition, 2019.
+    Retrieved from http://arxiv.org/abs/1812.07035
+    """
+    batch_dim = matrix.size()[:-2]
+    return matrix[..., :2, :].clone().reshape(batch_dim + (6, ))

model/utils/visualize.py

@@ -0,0 +1,183 @@
+import math
+import torch
+import numpy as np
+import py3Dmol
+import matplotlib.pyplot as plt
+import seaborn as sns
+
+from igfold.utils.folding import get_sequence_dict
+from igfold.utils.general import exists
+from igfold.utils.pdb import get_cdr_range_dict
+
+
+def show_pdb(
+        pdb_filename: str,
+        num_sequences,
+        bb_sticks=False,
+        sc_sticks=False,
+        color="b",
+        view_size=(500, 500),
+):
+    return show_pdbs(
+        [pdb_filename],
+        num_sequences,
+        bb_sticks=bb_sticks,
+        sc_sticks=sc_sticks,
+        color=color,
+        view_size=view_size,
+    )
+
+
+###
+#   Inspired by ColabFold visualization from https://github.com/sokrypton/ColabFold
+###
+
+
+def show_pdbs(
+        pdb_filenames,
+        num_sequences,
+        bb_sticks=False,
+        sc_sticks=False,
+        color="b",
+        view_size=(800, 800),
+):
+    grid_width = math.ceil(math.sqrt(len(pdb_filenames)))
+    grid_height = math.ceil(len(pdb_filenames) / grid_width)
+    view = py3Dmol.view(
+        js="https://3dmol.org/build/3Dmol.js",
+        viewergrid=(grid_height, grid_width),
+        width=view_size[0],
+        height=view_size[1],
+    )
+
+    for pdb_i, pdb_filename in enumerate(pdb_filenames):
+        grid_row, grid_col = pdb_i // grid_width, pdb_i % grid_width
+        view.addModel(
+            open(pdb_filename, "r").read(),
+            "pdb",
+            viewer=(grid_row, grid_col),
+        )
+
+    if color == "b":
+        view.setStyle({
+            "cartoon": {
+                "colorscheme": {
+                    "prop": "b",
+                    "gradient": "roygb",
+                    "min": 1.5,
+                    "max": 0.5,
+                }
+            }
+        })
+    elif color == "rainbow":
+        view.setStyle({"cartoon": {"color": "spectrum"}})
+    elif color == "chain":
+        for n, chain, color_ in zip(
+                range(num_sequences),
+                list("ABCDEFGH"),
+            [
+                "lime", "cyan", "magenta", "yellow", "salmon", "white", "blue",
+                "orange"
+            ],
+        ):
+            view.setStyle({"chain": chain}, {"cartoon": {"color": color_}})
+    if sc_sticks:
+        BB = ["C", "O", "N"]
+        view.addStyle(
+            {
+                "and": [
+                    {
+                        "resn": ["GLY", "PRO"],
+                        "invert": True
+                    },
+                    {
+                        "atom": BB,
+                        "invert": True
+                    },
+                ]
+            },
+            {"stick": {
+                "colorscheme": f"WhiteCarbon",
+                "radius": 0.2
+            }},
+        )
+        view.addStyle(
+            {"and": [{
+                "resn": "GLY"
+            }, {
+                "atom": "CA"
+            }]},
+            {"sphere": {
+                "colorscheme": f"WhiteCarbon",
+                "radius": 0.3
+            }},
+        )
+        view.addStyle(
+            {"and": [{
+                "resn": "PRO"
+            }, {
+                "atom": ["C", "O"],
+                "invert": True
+            }]},
+            {"stick": {
+                "colorscheme": f"WhiteCarbon",
+                "radius": 0.3
+            }},
+        )
+    if bb_sticks:
+        BB = ["C", "O", "N", "CA"]
+        view.addStyle(
+            {"atom": BB},
+            {"stick": {
+                "colorscheme": f"WhiteCarbon",
+                "radius": 0.3
+            }},
+        )
+
+    view.zoomTo()
+    return view
+
+
+def plot_prmsd(
+    sequences,
+    prmsd,
+    out_file=None,
+    shade_cdr=False,
+    pdb_file=None,
+):
+    seq_dict = get_sequence_dict(sequences, None)
+    delims = np.cumsum([len(s) for s in seq_dict.values()]).tolist()
+
+    res_rmsd = prmsd.cpu().square().mean(dim=-1).sqrt().squeeze(0)
+    chain_res_rmsd = np.split(res_rmsd, delims)
+
+    if shade_cdr and exists(pdb_file):
+        heavy_only = len(sequences) == 1 and "H" in sequences
+        light_only = len(sequences) == 1 and "L" in sequences
+        cdr_range_dict = get_cdr_range_dict(
+            pdb_file,
+            heavy_only=heavy_only,
+            light_only=light_only,
+            offset_heavy=False,
+        )
+        cdr_ranges = np.split(np.array(list(cdr_range_dict.values())), [3])
+
+    plt.figure(figsize=(8, 4))
+    for i, (chain, rmsd) in enumerate(zip(seq_dict.keys(), chain_res_rmsd)):
+        plt.subplot(1, len(seq_dict), i + 1)
+
+        res_nums = torch.arange(1, len(rmsd) + 1)
+        sns.lineplot(x=res_nums, y=rmsd)
+
+        if shade_cdr and exists(pdb_file):
+            for r in cdr_ranges[i]:
+                plt.axvspan(r[0], r[1], color="gray", alpha=0.5)
+
+        plt.xlabel("Residue Number")
+        plt.ylabel("Predicted RMSD (A)")
+        plt.title(f"Chain {chain}")
+
+    plt.tight_layout()
+
+    if exists(out_file):
+        plt.savefig(out_file, dpi=400)

pyproject.toml

@@ -0,0 +1,3 @@
+[build-system]
+requires = ["setuptools"]
+build-backend = "setuptools.build_meta"

requirements.txt

@@ -0,0 +1,16 @@
+# Dependencies already provided by the OneScience bio-dcu environment are kept
+# below as comments for traceability. Only uncommented packages need installing.
+
+antiberty>=0.1.1
+# biopython>=1.81  # provided by OneScience (1.84)
+# einops>=0.3.0  # provided by OneScience (>=0.7.0)
+future>=0.18.2
+# matplotlib>=3.7  # provided by OneScience
+# numpy==1.21.2  # provided by OneScience (1.26.3)
+# pytorch-lightning>=1.5.10,<1.9.0  # OneScience provides 2.0.6; keep its version
+# requests>=2.26.0  # provided by OneScience
+# seaborn>=0.12  # provided by OneScience
+# torch==1.7.1  # provided by the OneScience bio-dcu base environment
+# tqdm>=4.62.1  # provided by OneScience (>=4.60.0)
+yapf>=0.31.0
+# py3Dmol>=1.8.0  # provided by OneScience (2.5.2)

scripts/IgFold.ipynb

@@ -0,0 +1,199 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "5AEsv1z5KXXA"
+      },
+      "source": [
+        "# **IgFold**: Fast, accurate antibody structure prediction\n",
+        "\n",
+        "Official notebook for [IgFold](https://www.biorxiv.org/content/10.1101/2022.04.20.488972): Fast, accurate antibody structure prediction from deep learning on massive set of natural antibodies.  The code, data, and weights for this work are made available for non-commercial use. For commercial inquiries, please contact `jruffolo[at]jhu.edu`."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "0PsLNGK57LDq"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Input antibody Fv sequences then press `Runtime` -> `Run all`\n",
+        "\n",
+        "import os\n",
+        "import sys\n",
+        "\n",
+        "python_version = f\"{sys.version_info.major}.{sys.version_info.minor}\"\n",
+        "\n",
+        "name = \"my_antibody\" #@param {type:\"string\"}\n",
+        "pred_dir = name\n",
+        "os.makedirs(pred_dir, exist_ok=True)\n",
+        "\n",
+        "#@markdown Enter antibody sequences for structure prediction. To predict a nanobody structure (or an individual heavy or light chain), simply provide one sequence.\n",
+        "heavy_sequence = \"EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKFQERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS\" #@param {type:\"string\"}\n",
+        "light_sequence = \"DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGVPDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK\" #@param {type:\"string\"}\n",
+        "\n",
+        "sequences = {}\n",
+        "if len(heavy_sequence) > 0:\n",
+        "    sequences[\"H\"] = heavy_sequence\n",
+        "if len(light_sequence) > 0:\n",
+        "    sequences[\"L\"] = light_sequence\n",
+        "\n",
+        "#@markdown Perform structural refinement with OpenMM\n",
+        "do_refine = True #@param {type:\"boolean\"}\n",
+        "#@markdown Renumber predicted antibody structure (Chothia) with AbNumber\n",
+        "do_renum = False #@param {type:\"boolean\"}\n",
+        "#@markdown Use only a single model for predictions (instead of model ensemble)\n",
+        "single_model = False #@param {type:\"boolean\"}"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "LsJNdVE87Go2"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Install dependencies\n",
+        "\n",
+        "PYTHON_VERSION = python_version\n",
+        "\n",
+        "if not os.path.isfile(\"CONDA_READY\"):\n",
+        "  print(\"installing conda...\")\n",
+        "  os.system(\"wget -qnc https://github.com/jaimergp/miniforge/releases/latest/download/Mambaforge-colab-Linux-x86_64.sh\")\n",
+        "  os.system(\"bash Mambaforge-colab-Linux-x86_64.sh -bfp /usr/local\")\n",
+        "  os.system(\"mamba config --set auto_update_conda false\")\n",
+        "  os.system(\"touch CONDA_READY\")\n",
+        "\n",
+        "if not os.path.isfile(\"CODE_READY\"):\n",
+        "  print(\"installing igfold...\")\n",
+        "  torch_string = \"torch==1.11.0+cu113 torchvision==0.12.0+cu113 -f https://download.pytorch.org/whl/torch_stable.html\"\n",
+        "  os.system(f\"pip3 install {torch_string}\")\n",
+        "  os.system(f\"pip install 'igfold>=0.3.0' {torch_string}\")\n",
+        "  os.system(\"pip install -q --no-warn-conflicts 'py3Dmol>=2.0.1' matplotlib seaborn\")\n",
+        "  os.system(\"touch CODE_READY\")\n",
+        "\n",
+        "if do_refine and not os.path.isfile(\"AMBER_READY\"):\n",
+        "  print(\"installing amber...\")\n",
+        "  os.system(f\"mamba install -y -q -c conda-forge openmm=7.7.0 python='{PYTHON_VERSION}' pdbfixer 2>&1 1>/dev/null\")\n",
+        "  os.system(\"touch AMBER_READY\")\n",
+        "\n",
+        "if do_renum and not os.path.isfile(\"ABNUMBER_READY\"):\n",
+        "  print(\"installing abnumber...\")\n",
+        "  os.system(f\"mamba install -y -q -c bioconda abnumber python='{PYTHON_VERSION}' 2>&1 1>/dev/null\")\n",
+        "  os.system(\"pip install pandas --force-reinstall\")\n",
+        "  os.system(\"touch ABNUMBER_READY\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "a2a3BsiE9AXI"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Predict antibody structure with IgFold\n",
+        "\n",
+        "if f\"/usr/local/lib/python{python_version}/site-packages/\" not in sys.path:\n",
+        "    sys.path.insert(0, f\"/usr/local/lib/python{python_version}/site-packages/\")\n",
+        "\n",
+        "from igfold.utils.visualize import *\n",
+        "from igfold import IgFoldRunner\n",
+        "\n",
+        "num_models = 1 if single_model else 4\n",
+        "igfold = IgFoldRunner(num_models=num_models)\n",
+        "\n",
+        "pred_pdb = os.path.join(pred_dir, f\"{name}.pdb\")\n",
+        "pred = igfold.fold(\n",
+        "    pred_pdb,\n",
+        "    sequences=sequences,\n",
+        "    do_refine=do_refine,\n",
+        "    use_openmm=True,\n",
+        "    do_renum=do_renum,\n",
+        ")\n",
+        "show_pdb(pred_pdb, len(sequences), bb_sticks=False, sc_sticks=True, color=\"rainbow\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "xFOTYxsP9Cz1"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Plot per-residue predicted RMSD\n",
+        "\n",
+        "prmsd_fig_file = os.path.join(pred_dir, f\"{name}_prmsd.png\")\n",
+        "plot_prmsd(sequences, pred.prmsd.cpu(), prmsd_fig_file, shade_cdr=do_renum, pdb_file=pred_pdb)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "ajyElWbZ9EFF"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Show predicted structure with predicted RMSD\n",
+        "\n",
+        "#@markdown Structure is colored from low (blue) to high (red) pRMSD.\n",
+        "\n",
+        "show_pdb(pred_pdb, len(sequences), bb_sticks=False, sc_sticks=True, color=\"b\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "gZBzjpMdJ77q"
+      },
+      "outputs": [],
+      "source": [
+        "#@title Download results\n",
+        "\n",
+        "#@markdown Download zip file containing structure prediction and annotation results. If download fails, results are also accessible from file explorer on the left panel of the notebook.\n",
+        "\n",
+        "from google.colab import files\n",
+        "import locale\n",
+        "locale.getpreferredencoding = lambda: \"UTF-8\"\n",
+        "\n",
+        "!zip -FSr $name\".result.zip\" $pred_dir/ &> /dev/null\n",
+        "files.download(f\"{name}.result.zip\")"
+      ]
+    }
+  ],
+  "metadata": {
+    "colab": {
+      "collapsed_sections": [],
+      "name": "IgFold.ipynb",
+      "provenance": []
+    },
+    "kernelspec": {
+      "display_name": "Python 3.9.12 ('igfold_public')",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python",
+      "version": "3.9.12"
+    },
+    "vscode": {
+      "interpreter": {
+        "hash": "84181e5f1827f203c248bfcd3a60e7e3a4ffc08f0a7dd8a443bd855d4ab14b5d"
+      }
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

scripts/inference.py

@@ -0,0 +1,88 @@
+"""Run IgFold structure prediction from antibody sequences or a FASTA file."""
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+
+from igfold import IgFoldRunner
+
+
+DEFAULT_HEAVY = (
+    "EVQLVQSGPEVKKPGTSVKVSCKASGFTFMSSAVQWVRQARGQRLEWIGWIVIGSGNTNYAQKF"
+    "QERVTITRDMSTSTAYMELSSLRSEDTAVYYCAAPYCSSISCNDGFDIWGQGTMVTVS"
+)
+DEFAULT_LIGHT = (
+    "DVVMTQTPFSLPVSLGDQASISCRSSQSLVHSNGNTYLHWYLQKPGQSPKLLIYKVSNRFSGV"
+    "PDRFSGSGSGTDFTLKISRVEAEDLGVYFCSQSTHVPYTFGGGTKLEIK"
+)
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("--fasta", type=Path, help="Input FASTA containing H/L chains.")
+    parser.add_argument("--heavy", help="Heavy-chain amino-acid sequence.")
+    parser.add_argument("--light", help="Light-chain amino-acid sequence.")
+    parser.add_argument("--output", type=Path, default=Path("output/inference/antibody.pdb"))
+    parser.add_argument("--num-models", type=int, choices=range(1, 5), default=4)
+    parser.add_argument("--refine", action="store_true", help="Enable structure refinement.")
+    parser.add_argument("--openmm", action="store_true", help="Use OpenMM for refinement.")
+    parser.add_argument("--renum", action="store_true", help="Apply Chothia numbering.")
+    parser.add_argument("--cpu", action="store_true", help="Force CPU inference.")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    root = Path(__file__).resolve().parents[1]
+    checkpoints = sorted((root / "weight" / "IgFold").glob("*.ckpt"))[: args.num_models]
+    if len(checkpoints) != args.num_models:
+        raise FileNotFoundError(
+            f"Expected {args.num_models} checkpoints in {root / 'weight' / 'IgFold'}, "
+            f"found {len(checkpoints)}."
+        )
+
+    if args.fasta is not None:
+        fasta_file = str(args.fasta.resolve())
+        sequences = None
+    else:
+        fasta_file = None
+        sequences = {"H": args.heavy or DEFAULT_HEAVY}
+        light = args.light if args.light is not None else DEFAULT_LIGHT
+        if light:
+            sequences["L"] = light
+
+    if args.refine and not args.openmm:
+        from igfold.refine.pyrosetta_ref import init_pyrosetta
+
+        init_pyrosetta()
+
+    args.output.parent.mkdir(parents=True, exist_ok=True)
+    runner = IgFoldRunner(
+        model_ckpts=[str(path) for path in checkpoints],
+        try_gpu=not args.cpu,
+    )
+    output = runner.fold(
+        str(args.output),
+        fasta_file=fasta_file,
+        sequences=sequences,
+        do_refine=args.refine,
+        use_openmm=args.openmm,
+        do_renum=args.renum,
+    )
+
+    result = {
+        "status": "PASS",
+        "output_pdb": str(args.output.resolve()),
+        "models": len(runner.models),
+        "device": str(next(runner.models[0].parameters()).device),
+        "torch": torch.__version__,
+        "coords_shape": list(output.coords.shape),
+        "prmsd_shape": list(output.prmsd.shape),
+    }
+    print("IGFOLD_INFERENCE_RESULT=" + json.dumps(result, sort_keys=True))
+
+
+if __name__ == "__main__":
+    main()

setup.cfg

@@ -0,0 +1,20 @@
+[metadata]
+name = igfold
+version = 0.4.0
+long-description = file: README.md
+long_description_content_type = text/markdown
+
+[options]
+packages = find:
+package_dir =
+    = model
+install_requires =
+    antiberty>=0.1.1
+    biopython>=1.79
+    einops>=0.3.0
+    pytorch-lightning>=1.5.10,<1.9.0
+    torch>=1.7.1
+include_package_data = True
+
+[options.packages.find]
+where = model

weight/IgFold/LICENSE.md

@@ -0,0 +1,24 @@
+# JHU Academic Software License Agreement
+
+This license agreement ("License") is effective automatically between you ("Licensee") and The Johns Hopkins University (“JHU”) for use of the software with which this License is distributed (“Software”) so long as Licensee complies with the following terms and conditions:
+
+The requirement to acknowledge the copyright of JHU as follows: “© 2022 The Johns Hopkins University” and copyrights of any incorporated third party software as described in the associated documentation.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the source code must retain the above copyright notice, and these terms and conditions.
+
+Neither the name of JHU nor its affiliates may be used to endorse or promote products derived from this Software without specific prior written permission from an authorized JHU representative.
+
+THIS SOFTWARE IS PROVIDED BY JHU "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL JHU BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGES. LICENSEE AGREES TO DEFEND, INDEMNIFY AND HOLD HARMLESS JHU FOR ANY CLAIMS ARISING FROM LICENSEE’S USE OF THE SOFTWARE TO THE FULLEST EXTENT PERMITTED BY LAW.
+
+Only a non-exclusive, nontransferable license is granted to Licensee to use the Software for non-commercial purposes. Commercial use of the Software requires a separately executed written license agreement.
+
+Licensee agrees that it will use the Software, and any modifications, improvements, or derivatives to the Software that the Licensee may create (collectively, "Improvements") solely for non-commercial purposes and shall not distribute or transfer the Software or Improvements to any third party without requiring that such third parties adhere to the terms of this License. Licensee agrees that any Improvements made by Licensee shall be subject to the same terms and conditions as the Software.
+
+Licensee acknowledges that JHU holds copyright in the Software or portions of the Software, and that the Software may incorporate third party software which may be subject to additional terms and conditions. Licensee shall comply with any additional terms and conditions appliable to such third party software.
+
+Licensee agrees that any publication of results obtained with the Software will acknowledge its use by an appropriate citation.
+
+Licensee’s rights under this License terminate automatically without notice from JHU if Licensee fail to comply with any term(s) of this License.
+
+This License shall be governed by the laws of the State of Maryland, excluding the application of its conflicts of law rules. Licensee agrees that any dispute shall be appropriate only in the state and federal courts located within the State of Maryland.
+

weight/IgFold/igfold_1.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:66c2aa5a4fd4ef485e99f661e6c605037fa6bfefc02464395836c23263948249
+size 6358637

weight/IgFold/igfold_2.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:3917eb78a3b8f9a84d53e3dae1132e0096b6ded0868a081a60d0603d502028bc
+size 6358637

weight/IgFold/igfold_3.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d9b60988f1400c12f715b440c7be0bebab3b3be73dc39657c35c2029c3174c52
+size 6358637

weight/IgFold/igfold_5.ckpt

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b75865aef284973d47af61598a2fd50cbcfb02c67015b636f7ecd23ed842b2b5
+size 6358637

weight/wheels/antiberty-0.1.3-py3-none-any.whl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:30d910992b190013871bac49cdc032e01a19339f7d2b958ab99b0eb44638352a
+size 96631471

weight/wheels/igfold-0.4.0-py3-none-any.whl

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:ddf73636cf86ca2ef91eab25f79bf8052a3deaeae0ee4e056c8eecd152582c9c
+size 23355225