Upload folder using huggingface_hub

.gitattributes

@@ -1,35 +1,54 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
 *.bin filter=lfs diff=lfs merge=lfs -text
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
 *.gz filter=lfs diff=lfs merge=lfs -text
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+*.pickle filter=lfs diff=lfs merge=lfs -text
+*.pkl filter=lfs diff=lfs merge=lfs -text
+*.tar filter=lfs diff=lfs merge=lfs -text
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+weight/release_data/ccd_cache/components.v20240608.cif filter=lfs diff=lfs merge=lfs -text
+model/PXDesignBench/pxdbench/metrics/TMalign filter=lfs diff=lfs merge=lfs -text
+model/Protenix-pxd/Protenix_Technical_Report.pdf filter=lfs diff=lfs merge=lfs -text
+model/Protenix-pxd/assets/constraint_metrics.png filter=lfs diff=lfs merge=lfs -text
+model/Protenix-pxd/assets/mini_performance.png filter=lfs diff=lfs merge=lfs -text
+model/Protenix-pxd/assets/protenix_predictions.gif filter=lfs diff=lfs merge=lfs -text
+model/Protenix-pxd/assets/v0.5.0_metrics.png filter=lfs diff=lfs merge=lfs -text

.gitignore

@@ -0,0 +1,171 @@
+# Byte-compiled / optimized / DLL files
+*__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+*.o
+*.obj
+*.d
+# ninjia relate
+*ninja*
+lock
+# Distribution / packaging
+.Python
+.vscode
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+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
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+*.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
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# data cache and checkpoints
+data_cache/
+checkpoints/
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__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/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+wandb/
+output*/
+release_data/
+tool_weights/

.gitmodules

@@ -0,0 +1,4 @@
+[submodule "eval_design"]
+	path = eval_design
+	url = https://github.com/bytedance/PXDesignBench.git
+	branch = master

CODE_OF_CONDUCT.md

@@ -0,0 +1,127 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

CONTRIBUTING.md

@@ -0,0 +1,61 @@
+# Contributing
+
+Thank you for investing your time in contributing to <NAME> project! 
+
+Read our [Code of Coduct](./CODE_OF_CONDUCT.md) to keep our community approachable and respectable.
+
+This guide details how to use issues and pull requests to improve <NAME> project.
+
+## General Guidelines
+
+### Pull Requests
+
+Make sure to keep Pull Requests small and functional to make them easier to review, understand, and look up in commit history. This repository uses "Squash and Commit" to keep our history clean and make it easier to revert changes based on PR.
+
+Adding the appropriate documentation, unit tests and e2e tests as part of a feature is the responsibility of the feature owner, whether it is done in the same Pull Request or not.
+
+Pull Requests should follow the "subject: message" format, where the subject describes what part of the code is being modified.
+
+Refer to the template for more information on what goes into a PR description.
+
+### Design Docs
+
+A contributor proposes a design with a PR on the repository to allow for revisions and discussions. If a design needs to be discussed before formulating a document for it, make use of Google doc and GitHub issue to involve the community on the discussion. 
+
+### GitHub Issues
+
+GitHub Issues are used to file bugs, work items, and feature requests with actionable items/issues (Please refer to the "Reporting Bugs/Feature Requests" section below for more information).
+
+### Reporting Bugs/Feature Requests
+
+We welcome you to use the GitHub issue tracker to report bugs or suggest features that have actionable items/issues (as opposed to introducing a feature request on GitHub Discussions).
+
+When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already reported the issue. Please try to include as much information as you can. Details like these are incredibly useful:
+
+- A reproducible test case or series of steps
+- The version of the code being used
+- Any modifications you've made relevant to the bug
+- Anything unusual about your environment or deployment
+
+## Contributing via Pull Requests
+
+### Find interesting issue
+
+If you spot a problem with the problem, [search if an issue already exists](https://github.com/bytedance/<NAME>/issues). If a related issue doesn't exist, you can open a new issue using [issue template](https://github.com/bytedance/<NAME>/issues/new/choose). 
+
+### Solve an issue
+
+Please check `DEVELOPMENT.md` in sub folder to get familar with running and testing codes.
+
+### Open a Pull request.
+
+When you're done making the changes, open a pull request and fill PR template so we can better review your PR. The template helps reviewers understand your changes and the purpose of your pull request. 
+
+Don't forget to link PR to issue if you are solving one.
+
+If you run into any merge issues, checkout this [git tutorial](https://lab.github.com/githubtraining/managing-merge-conflicts) to help you resolve merge conflicts and other issues.
+
+
+## Finding contributions to work on
+
+Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' and 'good first issue' issues are a great place to start.

LICENSE

@@ -0,0 +1,201 @@
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright [yyyy] [name of copyright owner]
+
+   Licensed under the Apache License, Version 2.0 (the "License");
+   you may not use this file except in compliance with the License.
+   You may obtain a copy of the License at
+
+       http://www.apache.org/licenses/LICENSE-2.0
+
+   Unless required by applicable law or agreed to in writing, software
+   distributed under the License is distributed on an "AS IS" BASIS,
+   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+   See the License for the specific language governing permissions and
+   limitations under the License.

README.md

@@ -0,0 +1,520 @@
+---
+license: apache-2.0
+language:
+- en
+- zh
+tags:
+- OneScience
+- life-science
+- protein-design
+- protein-binder
+- PXDesign
+- Protenix
+frameworks: PyTorch
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">PXDesign</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+PXDesign is an open-source suite from the ByteDance team for de novo protein binder design. Given a target protein structure, it generates candidate binders and further filters candidate structures through structure prediction and confidence evaluation workflows.
+
+The complete PXDesign workflow consists of the PXDesign diffusion generation model, ProteinMPNN sequence design, AF2-IG evaluation, and Protenix evaluation. The official implementation provides three primary modes, generation-only, preview, and extended, covering use cases from quick validation to complete candidate screening.
+
+Paper:
+> **PXDesign: Fast, Modular, and Accurate De Novo Design of Protein Binders**  
+> https://www.biorxiv.org/content/10.1101/2025.08.15.670647v1
+
+# Model Description
+
+The core task of PXDesign is to generate new protein binders from a target protein structure and specified design regions.
+
+The typical workflow is:
+```text
+Target protein structure and design constraints
+  -> PXDesign-d diffusion model
+        -> Binder Backbone Generation
+  -> ProteinMPNN sequence design
+  -> AF2-IG structure prediction and filtering
+  -> Protenix structure prediction and filtering (extended mode)
+        -> summary.csv
+  -> Filtered high-confidence binders
+```
+
+Where:
+- **PXDesign-d**: Generates candidate binder backbones from the target protein structure, hotspots, binder length, and other conditions.
+- **ProteinMPNN**: Designs amino acid sequences for the generated protein backbones.
+- **AF2-IG**: Predicts structures and applies quality filters to candidate binder-target complexes.
+- **Protenix**: Provides additional structure prediction and confidence evaluation in extended mode.
+- **summary.csv**: Summarizes AF2-IG, Protenix, and other evaluation metrics for candidate structures, along with the pass status of each filter.
+
+# Use Cases
+
+| Use case | Description |
+| --- | --- |
+| De novo protein binder design | Generate new candidate binders from a given target protein structure |
+| Interface-guided design | Use hotspots to specify target residues that the binder should preferentially bind |
+| Rapid validation of protein design workflows | Use preview mode to quickly evaluate whether the design task and parameters are reasonable |
+| High-quality candidate screening | Use extended mode with AF2-IG and Protenix for multistage filtering |
+| Structure generation research | Use `pxdesign infer` to run only the PXDesign generation stage |
+
+# Usage
+
+## 1. Using OneCode
+
+Experience intelligent one-click AI4S programming in the OneCode online environment:
+[Try intelligent one-click AI4S programming](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation and Usage
+
+### Hardware Requirements
+
+- A GPU/DCU is recommended for the PXDesign backbone generation stage; complete inference typically requires substantial GPU memory.
+- MSA generation and preparation primarily use the CPU. You can prepare MSAs in advance with `prepare-msa` or by precomputing them.
+- The ProteinMPNN, AF2-IG, and Protenix prediction and screening stages depend on deep learning frameworks such as PyTorch and JAX; a GPU/DCU is recommended.
+- If GPU/DCU resources are limited, prepare the MSA separately on the CPU first, then run the PXDesign generation, ProteinMPNN, AF2-IG, and Protenix evaluation stages.
+
+### Set Up the Runtime Environment
+
+#### DCU Environment
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+
+# Install with uv support
+pip install onescience[bio] -i http://mirrors.onescience.ai:3141/pypi/simple/ --trusted-host mirrors.onescience.ai
+```
+
+#### Environment Notes
+
+- If you encounter missing dependencies or version issues during execution, install additional dependencies according to the versions specified in `requirements.txt`.
+- Enter the project root and activate the environment:
+
+```bash
+cd /path/to/PXDesign-main
+conda activate your_env
+```
+
+Install PXDesign in the current environment:
+
+```bash
+python -m pip install -e model
+```
+
+Verify the installation:
+
+```bash
+which pxdesign
+pxdesign --help
+pxdesign pipeline --help
+```
+
+### Environment Variables
+
+After entering the PXDesign root directory, configure the following variables together:
+```bash
+export PXDESIGN_ROOT=$PWD
+export TOOL_WEIGHTS_ROOT=$PWD/weight/tool_weights
+export PROTENIX_DATA_ROOT_DIR=$PWD/weight/release_data/ccd_cache
+```
+
+Check them with:
+```bash
+echo $PXDESIGN_ROOT
+echo $TOOL_WEIGHTS_ROOT
+echo $PROTENIX_DATA_ROOT_DIR
+```
+
+to confirm the configuration.
+
+## Prepare Weights and Data
+
+The complete PXDesign workflow depends on PXDesign and Protenix model weights, as well as AlphaFold2, ProteinMPNN, and the CCD cache. This model repository already includes the CCD cache and PXDesign/Protenix checkpoints; users only need to prepare the `tool_weights/` portion separately. The complete preparation process is as follows:
+
+### 1) External Tool Weights and CCD Cache
+
+PXDesign provides an official download script:
+```bash
+bash scripts/download_tool_weights.sh
+```
+
+The script uses the official default directories and generates `tool_weights/` and `release_data/ccd_cache/` in the current directory when run directly. This project has been reorganized under a `weight/` directory, so it is recommended to organize or symlink the existing weights and cache to the locations below.
+
+In the current reorganized project structure, place the external tool weights as follows:
+```text
+weight/
+├── tool_weights/
+│   ├── af2/        # AlphaFold2 weights
+│   └── mpnn/       # ProteinMPNN weights
+└── release_data/
+    └── ccd_cache/  # Protenix CCD cache
+```
+
+- The recommended default location for the CCD cache is:
+
+```text
+weight/release_data/ccd_cache/
+```
+
+To specify another location, set:
+
+```bash
+export PROTENIX_DATA_ROOT_DIR=/path/to/ccd_cache
+```
+
+### 2) PXDesign and Protenix Checkpoints
+
+The following model weights are downloaded on demand during the first run, or can be downloaded to the corresponding locations in advance:
+```text
+PXDesign diffusion checkpoint
+
+Protenix checkpoints:
+├── base
+├── mini
+└── mini_tmpl
+```
+
+The recommended location in the reorganized structure is:
+
+```text
+weight/release_data/checkpoint/
+```
+
+The required files include:
+```text
+pxdesign_v0.1.0.pt
+protenix_base_default_v0.5.0.pt
+protenix_mini_default_v0.5.0.pt
+protenix_mini_tmpl_v0.5.0.pt
+```
+
+### 3) Check the Installation
+
+After preparation is complete, run:
+```bash
+ls weight/tool_weights/af2/
+ls weight/tool_weights/mpnn/
+ls weight/release_data/ccd_cache/
+ls weight/release_data/checkpoint/*.pt
+```
+to confirm that the required weights and data are ready.
+
+## 3. Quick Start
+
+### Download the Model Package
+
+```bash
+hf download OneScience-Group/PXDesign --local-dir ./PXDesign
+cd PXDesign
+```
+
+- PXDesign additionally depends on Protenix and PXDesignBench; the corresponding dependency source code is included in this model repository, so separate downloads are not required.
+- The complete PXDesign workflow also depends on AlphaFold2, ProteinMPNN, and the CCD cache required by Protenix. Prepare these resources as described in "Prepare Weights and Data" first.
+
+### Quick Verification
+
+First, verify that the command is available:
+```bash
+pxdesign --help
+```
+
+To save the results to `runs/` as in the examples below, first create the output directory:
+```bash
+mkdir -p runs
+```
+
+Then check the official example YAML:
+```bash
+pxdesign check-input \
+  --yaml conf/examples/PDL1_quick_start.yaml
+```
+
+On success, the output should be:
+```text
+YAML file is valid.
+```
+
+### Example Data
+
+The current project provides:
+```text
+conf/examples/
+├── PDL1_quick_start.yaml
+├── 5o45.cif
+└── msa/
+    └── PDL1/
+        └── 0/
+```
+
+`PDL1_quick_start.yaml` defines the PDL1 binder design task.
+
+The typical YAML format is:
+```yaml
+target:
+  file: "./conf/examples/5o45.cif"
+  chains:
+    A:
+      crop: ["1-116"]
+      hotspots: [40, 99, 107]
+      msa: "./conf/examples/msa/PDL1/0"
+
+binder_length: 80
+```
+
+Key fields:
+| Field | Description |
+| --- | --- |
+| `target.file` | Target protein structure file; mmCIF or PDB can be used |
+| `target.chains` | Target chains involved in the design |
+| `crop` | Residue range retained from the target chain |
+| `hotspots` | Target residues used to guide binder interface generation |
+| `msa` | Path to the precomputed MSA for the target chain |
+| `binder_length` | Amino acid length of the binder to be designed |
+
+PXDesign primarily uses the mmCIF `label_seq_id` as the standard residue index internally. For custom tasks, mmCIF files are recommended, and `parse-target` should be used to check that crop and hotspot specifications point to the intended positions.
+
+### Input Checking and Target Parsing
+
+#### 1) Check the YAML
+
+Run the following before formally executing a design task:
+
+```bash
+pxdesign check-input \
+  --yaml conf/examples/PDL1_quick_start.yaml
+```
+
+#### 2) Parse the Target and Generate Visualization Debug Files
+```bash
+pxdesign parse-target \
+  --yaml conf/examples/PDL1_quick_start.yaml \
+  -o runs/debug_target
+```
+
+This step is useful for checking the following before running a large-scale design:
+
+- whether the crop is correct;
+- whether the hotspots correspond to the intended residues;
+- whether the structure chains and residue numbering are correct.
+
+## Inference Examples
+
+PXDesign primarily provides three execution modes:
+```text
+Generation Only
+    -> Generate only the PXDesign binder backbone
+
+Preview Pipeline
+    -> PXDesign + ProteinMPNN + AF2-IG
+
+Extended Pipeline
+    -> PXDesign + ProteinMPNN + AF2-IG + Protenix
+```
+
+### 1. Generation Only: Run PXDesign Generation Only
+
+#### Quick Smoke Test
+
+To first verify that the model, weights, and GPU/DCU work correctly, use a smaller number of steps:
+```bash
+pxdesign infer \
+  -i conf/examples/PDL1_quick_start.yaml \
+  -o runs/test_infer \
+  --load_checkpoint_dir weight/release_data/checkpoint \
+  --N_sample 1 \
+  --N_step 20 \
+  --dtype bf16 \
+  --sample_diffusion_chunk_size 1
+```
+
+#### Full-Step Generation Test
+
+```bash
+pxdesign infer \
+  -i conf/examples/PDL1_quick_start.yaml \
+  -o runs/test_infer_full \
+  --load_checkpoint_dir weight/release_data/checkpoint \
+  --N_sample 10 \
+  --N_step 400 \
+  --dtype bf16
+```
+
+This mode only generates binders and does not provide complete AF2/Protenix filtering results.
+
+### 2. Preview Pipeline
+
+Preview mode runs:
+```text
+PXDesign generation
+        -> ProteinMPNN sequence design
+        -> AF2-IG filtering
+```
+
+```bash
+pxdesign pipeline \
+  --preset preview \
+  -i conf/examples/PDL1_quick_start.yaml \
+  -o runs/test_preview \
+  --load_checkpoint_dir weight/release_data/checkpoint \
+  --N_sample 2 \
+  --N_step 100 \
+  --dtype bf16 \
+  --use_fast_ln False \
+  --use_deepspeed_evo_attention False
+```
+
+Preview mode is suitable for:
+- initial validation of the complete pipeline;
+- checking whether the hotspot/crop settings are reasonable;
+- assessing the difficulty of the current design task;
+- running a small pilot experiment before a large-scale Extended task.
+
+### 3. Extended Pipeline
+
+Extended mode is the official PXDesign workflow for complete evaluation:
+
+```text
+PXDesign generation
+        -> ProteinMPNN
+        -> AF2-IG
+        -> Protenix
+        -> summary.csv
+```
+
+#### Small-Scale Validation
+```bash
+pxdesign pipeline \
+  --preset extended \
+  -i conf/examples/PDL1_quick_start.yaml \
+  -o runs/test_extended \
+  --load_checkpoint_dir weight/release_data/checkpoint \
+  --N_sample 2 \
+  --N_step 100 \
+  --dtype bf16 \
+  --use_fast_ln False \
+  --use_deepspeed_evo_attention False
+```
+
+#### Quick Start Scale
+
+The official Quick Start example uses:
+```text
+N_sample = 10
+N_step = 400
+```
+
+```bash
+pxdesign pipeline \
+  --preset extended \
+  -i conf/examples/PDL1_quick_start.yaml \
+  -o runs/test_extended_N10 \
+  --load_checkpoint_dir weight/release_data/checkpoint \
+  --N_sample 10 \
+  --N_step 400 \
+  --dtype bf16 \
+  --use_fast_ln False \
+  --use_deepspeed_evo_attention False
+```
+
+## Output Description
+
+The core results from Extended mode are typically located at:
+```text
+<OUT_DIR>/
+└── design_outputs/
+    └── <task_name>/
+        ├── summary.csv
+        ├── task_info.json
+        ├── server_extended_mode.png
+        ├── orig_designed/
+        ├── passing-AF2-IG-easy/
+        └── passing-Protenix-basic/
+```
+
+If no designs pass the corresponding filters, the `passing-AF2-IG-easy/` or `passing-Protenix-basic/` directories may not be generated. This is normal for small-sample tests.
+
+# Official OneScience Information
+
+| Platform | Main OneScience repository | Skills repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+
+# Citation and License
+
+If you use PXDesign in research, cite the official PXDesign paper:
+```bibtex
+@article{ren2025pxdesign,
+  title={PXDesign: Fast, Modular, and Accurate De Novo Design of Protein Binders},
+  author={Ren, Milong and Sun, Jinyuan and Guan, Jiaqi and Liu, Cong and
+          Gong, Chengyue and Wang, Yuzhe and Wang, Lan and Cai, Qixu and
+          Chen, Xinshi and Xiao, Wenzhi},
+  journal={bioRxiv},
+  pages={2025--08},
+  year={2025},
+  publisher={Cold Spring Harbor Laboratory}
+}
+```
+
+The complete PXDesign pipeline also depends on methods such as Protenix, ProteinMPNN, and AF2-IG. If you use these modules in research, also cite the corresponding original works as described in the official PXDesign README.
+
+Protenix:
+```bibtex
+@article{bytedance2025protenix,
+  title={Protenix - Advancing Structure Prediction Through a Comprehensive AlphaFold3 Reproduction},
+  author={ByteDance AML AI4Science Team and Chen, Xinshi and Zhang, Yuxuan
+          and Lu, Chan and Ma, Wenzhi and Guan, Jiaqi and Gong, Chengyue
+          and Yang, Jincai and Zhang, Hanyu and Zhang, Ke and Wu, Shenghao
+          and Zhou, Kuangqi and Yang, Yanping and Liu, Zhenyu and Wang, Lan
+          and Shi, Bo and Shi, Shaochen and Xiao, Wenzhi},
+  year={2025},
+  journal={bioRxiv},
+  publisher={Cold Spring Harbor Laboratory},
+  doi={10.1101/2025.01.08.631967}
+}
+```
+
+ProteinMPNN:
+```bibtex
+@article{dauparas2022robust,
+  title={Robust deep learning--based protein sequence design using ProteinMPNN},
+  author={Dauparas, Justas and Anishchenko, Ivan and Bennett, Nathaniel
+          and Bai, Hua and Ragotte, Robert J and Milles, Lukas and others},
+  journal={Science},
+  volume={378},
+  number={6615},
+  pages={49--56},
+  year={2022}
+}
+```
+
+AF2-IG:
+```bibtex
+@article{bennett2023improving,
+  title={Improving de novo protein binder design with deep learning},
+  author={Bennett, Nathaniel R and Coventry, Brian and Goreshnik, Inna
+          and Huang, Buwei and Allen, Aza and Vafeados, Dionne and others},
+  journal={Nature Communications},
+  volume={14},
+  number={1},
+  pages={2625},
+  year={2023}
+}
+```
+
+The official PXDesign repository is released under the **Apache License 2.0**. According to the official README, this license permits academic research and commercial use. Use, modification, and redistribution of the code must comply with the specific terms in this project's `LICENSE`.
+
+In addition:
+- AlphaFold2/AF2 weights and related resources are subject to their respective licenses and terms of use.
+- ProteinMPNN is subject to the license of its official repository.
+- Protenix is subject to the license of its official repository.
+- Models and data resources reused through SCNet shared directories remain subject to the authorization conditions of their respective original resources.
+
+For papers, reports, or public releases, also cite PXDesign, Protenix, ProteinMPNN, AF2-IG, and any other third-party models and data resources actually used.

conf/examples/PDL1_quick_start.yaml

@@ -0,0 +1,16 @@
+target:
+  # Path to the target structure (PDB or CIF)
+  file: "./conf/examples/5o45.cif"
+  # Specify the target chain(s)
+  chains:
+    # Configuration for Target Chain A
+    A:
+      # Crop regions (Optional)
+      crop: ["1-116"]
+      # Binding hotspot residues (Optional)
+      hotspots: [40, 99, 107]
+      # Path to pre-computed MSA (Optional but recommended)
+      msa: "./conf/examples/msa/PDL1/0"
+
+# Number of amino acids of the binder to be designed
+binder_length: 80

configuration.json

@@ -0,0 +1,39 @@
+{
+  "project": "PXDesign-main",
+  "layout": {
+    "conf": "Configuration files, example inputs, license files, and reproduction notes.",
+    "model": "PXDesign source code, bundled Protenix/PXDesignBench source trees, original README, and assets.",
+    "scripts": "Installation, dependency, Docker, and download scripts.",
+    "weight": "Model checkpoints, CCD cache, AF2/ProteinMPNN external tool weights, and MSA cache."
+  },
+  "paths": {
+    "pxdesign_package": "model/pxdesign",
+    "protenix_source": "model/Protenix-pxd",
+    "pxdesignbench_source": "model/PXDesignBench",
+    "example_yaml": "conf/examples/PDL1_quick_start.yaml",
+    "checkpoint_dir": "weight/release_data/checkpoint",
+    "ccd_cache_dir": "weight/release_data/ccd_cache",
+    "tool_weights_root": "weight/tool_weights",
+    "af2_weights": "weight/tool_weights/af2",
+    "mpnn_weights": "weight/tool_weights/mpnn",
+    "reproduction_notes": "conf/PXDESIGN_REPRODUCTION_NOTES.md",
+    "original_readme": "model/README.md"
+  },
+  "environment": {
+    "conda_env_used": "model_bio",
+    "tool_weights_root_env": "TOOL_WEIGHTS_ROOT=$PWD/weight/tool_weights",
+    "protenix_data_root_env": "PROTENIX_DATA_ROOT_DIR=$PWD/weight/release_data/ccd_cache"
+  },
+  "validated_tasks": [
+    "pxdesign check-input",
+    "pxdesign parse-target",
+    "pxdesign infer",
+    "pxdesign pipeline --preset preview",
+    "pxdesign pipeline --preset extended"
+  ],
+  "hpc_runtime_overrides": {
+    "use_fast_ln": false,
+    "use_deepspeed_evo_attention": false,
+    "reason": "The DTK/HIP environment failed to build fastfold_layer_norm_cuda with CUDA compiler flags."
+  }
+}

model/PXDesignBench/.gitignore

@@ -0,0 +1,168 @@
+# Byte-compiled / optimized / DLL files
+*__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+*.o
+*.obj
+*.d
+# ninjia relate
+*ninja*
+lock
+# Distribution / packaging
+.Python
+.vscode
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+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
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+*.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
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# data cache and checkpoints
+data_cache/
+checkpoints/
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .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
+
+# poetry
+#   Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+#   This is especially recommended for binary packages to ensure reproducibility, and is more
+#   commonly ignored for libraries.
+#   https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+#   Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+#   pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+#   in version control.
+#   https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__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/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+#  JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+#  be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+#  and can be added to the global gitignore or merged into this file.  For a more nuclear
+#  option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+output*/

model/PXDesignBench/CODE_OF_CONDUCT.md

@@ -0,0 +1,127 @@
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.

model/PXDesignBench/CONTRIBUTING.md

@@ -0,0 +1,61 @@
+# Contributing
+
+Thank you for investing your time in contributing to PXDesignBench project! 
+
+Read our [Code of Conduct](./CODE_OF_CONDUCT.md) to keep our community approachable and respectable.
+
+This guide details how to use issues and pull requests to improve PXDesignBench project.
+
+## General Guidelines
+
+### Pull Requests
+
+Make sure to keep Pull Requests small and functional to make them easier to review, understand, and look up in commit history. This repository uses "Squash and Commit" to keep our history clean and make it easier to revert changes based on PR.
+
+Adding the appropriate documentation, unit tests and e2e tests as part of a feature is the responsibility of the feature owner, whether it is done in the same Pull Request or not.
+
+Pull Requests should follow the "subject: message" format, where the subject describes what part of the code is being modified.
+
+Refer to the template for more information on what goes into a PR description.
+
+### Design Docs
+
+A contributor proposes a design with a PR on the repository to allow for revisions and discussions. If a design needs to be discussed before formulating a document for it, make use of Google doc and GitHub issue to involve the community on the discussion. 
+
+### GitHub Issues
+
+GitHub Issues are used to file bugs, work items, and feature requests with actionable items/issues (Please refer to the "Reporting Bugs/Feature Requests" section below for more information).
+
+### Reporting Bugs/Feature Requests
+
+We welcome you to use the GitHub issue tracker to report bugs or suggest features that have actionable items/issues (as opposed to introducing a feature request on GitHub Discussions).
+
+When filing an issue, please check existing open, or recently closed, issues to make sure somebody else hasn't already reported the issue. Please try to include as much information as you can. Details like these are incredibly useful:
+
+- A reproducible test case or series of steps
+- The version of the code being used
+- Any modifications you've made relevant to the bug
+- Anything unusual about your environment or deployment
+
+## Contributing via Pull Requests
+
+### Find interesting issue
+
+If you spot a problem with the problem, [search if an issue already exists](https://github.com/bytedance/PXDesignBench/issues). If a related issue doesn't exist, you can open a new issue using [issue template](https://github.com/bytedance/PXDesignBench/issues/new/choose). 
+
+### Solve an issue
+
+Please check `DEVELOPMENT.md` in sub folder to get familiar with running and testing codes.
+
+### Open a Pull request.
+
+When you're done making the changes, open a pull request and fill PR template so we can better review your PR. The template helps reviewers understand your changes and the purpose of your pull request. 
+
+Don't forget to link PR to the issue if you are solving one.
+
+If you run into any merge issues, checkout this [git tutorial](https://lab.github.com/githubtraining/managing-merge-conflicts) to help you resolve merge conflicts and other issues.
+
+
+## Finding contributions to work on
+
+Looking at the existing issues is a great way to find something to contribute on. As our projects, by default, use the default GitHub issue labels (enhancement/bug/duplicate/help wanted/invalid/question/wontfix), looking at any 'help wanted' and 'good first issue' issues are a great place to start.

model/PXDesignBench/ColabDesign/.github/workflows/python-publish.yml

@@ -0,0 +1,39 @@
+# This workflow will upload a Python Package using Twine when a release is created
+# For more information see: https://docs.github.com/en/actions/automating-builds-and-tests/building-and-testing-python#publishing-to-package-registries
+
+# This workflow uses actions that are not certified by GitHub.
+# They are provided by a third-party and are governed by
+# separate terms of service, privacy policy, and support
+# documentation.
+
+name: Upload Python Package
+
+on:
+  release:
+    types: [published]
+
+permissions:
+  contents: read
+
+jobs:
+  deploy:
+
+    runs-on: ubuntu-latest
+
+    steps:
+    - uses: actions/checkout@v3
+    - name: Set up Python
+      uses: actions/setup-python@v3
+      with:
+        python-version: '3.x'
+    - name: Install dependencies
+      run: |
+        python -m pip install --upgrade pip
+        pip install build
+    - name: Build package
+      run: python -m build
+    - name: Publish package
+      uses: pypa/gh-action-pypi-publish@release/v1
+      with:
+        user: __token__
+        password: ${{ secrets.PYPI_API_TOKEN }}

model/PXDesignBench/ColabDesign/.gitignore

@@ -0,0 +1,2 @@
+**/.DS_Store
+*.pyc

model/PXDesignBench/ColabDesign/LICENSE.txt

@@ -0,0 +1,8 @@
+/*
+ * ----------------------------------------------------------------------------
+ * "THE BEER-WARE LICENSE" (Revision 42):
+ * <so3@mit.edu> wrote this file.  As long as you retain this notice you
+ * can do whatever you want with this stuff. If we meet some day, and you think
+ * this stuff is worth it, you can buy me a beer in return.  Sergey Ovchinnikov
+ * ----------------------------------------------------------------------------
+ */

model/PXDesignBench/ColabDesign/MANIFEST.in

@@ -0,0 +1,4 @@
+include colabdesign/mpnn/weights/*.pkl
+include colabdesign/mpnn/weights_soluble/*.pkl
+include colabdesign/tr/weights/*.npy
+include colabdesign/af/weights/*.npy

model/PXDesignBench/ColabDesign/README.md

@@ -0,0 +1,29 @@
+# ColabDesign
+### Making Protein Design accessible to all via Google Colab! 
+- P(structure | sequence)
+  - [TrDesign](/tr) - using TrRosetta for design
+  - [AfDesign](/af) - using AlphaFold for design
+  - [WIP] [RfDesign](https://github.com/RosettaCommons/RFDesign) - using RoseTTAFold for design
+- P(sequence | structure)
+  - [ProteinMPNN](/mpnn)
+  - [WIP] TrMRF
+- P(sequence)
+  - [WIP] [MSA_transformer](/esm_msa)
+  - [WIP] [SEQ](/seq) - (GREMLIN, mfDCA, arDCA, plmDCA, bmDCA, etc)
+- P(structure)
+  - [Rfdiffusion](/rf)
+
+### Where can I chat with other ColabDesign users?
+  - See our [Discord](https://discord.gg/gna8maru7d) channel!
+
+
+### Presentations
+[Slides](https://docs.google.com/presentation/d/1Zy7lf_LBK0_G3e7YQLSPP5aj_-AR5I131fTsxJrLdg4/)
+[Talk](https://www.youtube.com/watch?v=2HmXwlKWMVs)
+
+### Contributors:
+- Sergey Ovchinnikov [@sokrypton](https://github.com/sokrypton)
+- Shihao Feng [@JeffSHF](https://github.com/JeffSHF)
+- Justas Dauparas [@dauparas](https://github.com/dauparas)
+- Weikun.Wu [@guyujun](https://github.com/guyujun) (from [Levinthal.bio](http://levinthal.bio/en/))
+- Christopher Frank [@chris-kafka](https://github.com/chris-kafka)

model/PXDesignBench/ColabDesign/af/LICENSE.txt

@@ -0,0 +1,8 @@
+/*
+ * ----------------------------------------------------------------------------
+ * "THE BEER-WARE LICENSE" (Revision 42):
+ * <so3@mit.edu> wrote this file.  As long as you retain this notice you
+ * can do whatever you want with this stuff. If we meet some day, and you think
+ * this stuff is worth it, you can buy me a beer in return.  Sergey Ovchinnikov
+ * ----------------------------------------------------------------------------
+ */

model/PXDesignBench/ColabDesign/af/README.md

@@ -0,0 +1,336 @@
+# AfDesign (v1.1.1)
+### Google Colab
+<a href="https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/design.ipynb">
+  <img src="https://colab.research.google.com/assets/colab-badge.svg" alt="Open In Colab"/>
+</a>
+
+# Updates
+- Jump to [Previous Updates](#previous-updates)
+- **15Oct2022** - v1.1.0
+  - integrating proteinMPNN!
+  - bugfix for sidechain loss
+- **17Nov2022**
+  - updating pae/plddt loss calculation to be consistent with pae/plddt outputs
+- **24Dec2022** - v1.1.1
+  - adding af_pseudo_diffusion examples
+  - updating to alphafold-multimer v2.3.0
+  - enabling fused_triangle_multiplication by default
+- **21Jan2023**
+  - add support for bfloat16 (enabled by default
+- **01Mar2023**
+  - adding support for [RfDiffusion](/rf)
+### setup
+first install jax (with GPU support)
+```bash
+pip install "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html
+```
+second install colabdesign
+```bash
+pip install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1
+
+# download alphafold weights
+mkdir params
+curl -fsSL https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar | tar x -C params
+```
+By default `mk_afdesign_model()` assumes alphafold "params" are saved in the run directory (`data_dir="."`). To override:
+```python
+model = mk_afdesign_model(..., data_dir="/location/of")
+```
+### import
+```python
+import numpy as np
+from IPython.display import HTML
+from colabdesign import mk_afdesign_model, clear_mem
+```
+### fixed backbone design
+For a given protein backbone, generate/design a new sequence that AlphaFold thinks folds into that conformation
+```python
+model = mk_afdesign_model(protocol="fixbb")
+model.prep_inputs(pdb_filename="1TEN.pdb", chain="A")
+model.design_3stage()
+```
+### hallucination
+For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured 
+protein (high plddt, low pae, many contacts).
+```python
+model = mk_afdesign_model(protocol="hallucination")
+model.prep_inputs(length=100)
+model.set_seq(mode="gumbel")
+model.design_soft(50)
+model.set_seq(model.aux["seq"]["pseudo"])
+model.design_3stage(50,50,10)
+```
+### binder hallucination
+For a given protein target and protein binder length, generate/hallucinate a protein binder sequence AlphaFold 
+thinks will bind to the target structure. To do this, we minimize PAE and maximize number of contacts at the 
+interface and within the binder, and we maximize pLDDT of the binder.
+```python
+model = mk_afdesign_model(protocol="binder")
+model.prep_inputs(pdb_filename="4MZK.pdb", chain="A", binder_len=19)
+model.design_3stage(100, 100, 10)
+```
+Instead of hallucination, you can redesign an existing binder:
+```python
+model.prep_inputs(pdb_filename="4MZK.pdb", chain="A", binder_chain="T")
+```
+
+### partial hallucination
+If you have a motif (binding motif, or functional motif) and you want to hallucinate a new scaffold around it,
+you can use partial hallucination. Or you have a protein and you want to extend one of the loops.
+```python
+af_model = mk_afdesign_model(protocol="partial")
+af_model.prep_inputs(pdb_filename="6MRR.pdb", chain="A", pos="3-30,33-68", length=100)
+af_model.rewire(loops=[36])
+```
+# FAQ
+
+#### Can I reuse the same model without needing to recompile?
+```python
+model.restart()
+```
+#### How do I change the loss weights?
+This can be done using the provided function:
+```python
+model.set_weights(pae=0.0, plddt=1.0)
+```
+or the dictionary directly:
+```python
+model.opt["weights"]["pae"] = 0.0
+```
+#### How do I control number of recycles used during design?
+```python 
+model = mk_afdesign_model(num_recycles=1, recycle_mode="average")
+# if recycle_mode in ["average",last","sample","first"] the number of recycles can change during optimization
+model.set_opt(num_recycles=1)
+```
+- `num_recycles` - number of recycles to use during design (for denovo proteins we find 0 is often enough)
+- `recycle_mode` - optimizing across all recycles can be tricky, we experiment with a couple of ways:
+  - *last* - use loss from last recycle. (Default)
+  - *average* - compute loss at each recycle and average gradients. (Previous default from v.1.0.5)
+  - *sample* - Same as *last* but each iteration a different number of recycles are used.
+  - *first* - use loss from first recycle.
+  - *add_prev* - average the outputs (dgram, plddt, pae) across all recycles before computing loss.
+  - *backprop* - use loss from last recycle, but backprop through all recycles.
+
+#### How do I control which model params are used during design?
+By default all five models are used during optimization. If `num_models` > 1, then multiple params are evaluated at each iteration and the gradients/losses are averaged. Each iteration a random set of model params are used unless `sample_models=False`.
+```python
+model = mk_afdesign_model(num_models=1, sample_models=True)
+# or
+model.set_opt(num_models=1, sample_models=True)
+```
+- `num_models` - number of model params to use at each iteration.
+- `sample_models`:
+  - *True* - randomly select models params to use. (Recommended)
+  - *False* - use the same model params each iteration.
+You can also specify exactly which models are used during any of the design protocols:
+```python
+model.design_(num_models=1, sample_models=True, models=[0,2,3])
+# or
+model.design_(num_models=2, sample_models=False, models=["model_1_ptm","model_3_ptm"])
+```
+#### Can I use OpenFold model params for design instead of AlphaFold?
+You may need to download them:
+```bash
+  for W in openfold_model_ptm_1 openfold_model_ptm_2 openfold_model_no_templ_ptm_1
+  do wget -qnc https://files.ipd.uw.edu/krypton/openfold/${W}.npz -P params; done
+```
+Once downloaded:
+```python
+model = mk_afdesign_model(use_openfold=True, use_alphafold=False)
+```
+#### For binder hallucination, can I specify the site I want to bind?
+```python
+model.prep_inputs(..., hotspot="1-10,15,3")
+```
+#### Can I input more than one chain?
+```python
+model.prep_inputs(..., chain="A,B")
+```
+#### For fixed backbone design, how do I force the sequence to be the same for homo-dimer optimization?
+```python
+model.prep_inputs(pdb_filename="6Q40.pdb", chain="A,B", copies=2, homooligomer=True)
+```
+WARNING, this functionality assumes the input chains are of equal length.
+#### How do I disable certain amino acids?
+```python
+model.restart(rm_aa="C,W")
+```
+#### How do I set the random seed for reproducibility?
+```python
+model.restart(seed=0)
+```
+#### What are all the different `design_???` methods?
+- For **design** we provide 5 different functions:
+  - `design_logits()` - optimize *logits* inputs (continious)
+  - `design_soft()` - optimize *softmax(logits)* inputs (probabilities)
+  - `design_hard()` - optimize *one_hot(logits)* inputs (discrete)
+
+- For complex topologies, we find directly optimizing one_hot encoded sequence `design_hard()` to be very challenging. 
+To get around this problem, we propose optimizing in 3 stages or first learning logits then switching to semigreedy optimization.
+  - `design_3stage()` - gradient based optimization (GD) (logits → soft → hard)
+  - `design_semigreedy(tries=X)` - tries X random mutations, accepts those that decrease loss
+  - `design_pssm_semigreey(tries=X)` - uses GD to get a sequence profile (PSSM), then uses the PSSM to bias semigreedy opt. (Recommended)
+
+#### What are all the different losses being optimized?
+- general losses
+  - *pae*       - minimizes the predicted alignment error
+  - *plddt*     - maximizes the predicted LDDT
+  - *pae* and *plddt* values are between 0 and 1 (where lower is better for both)
+
+- fixbb specific losses
+  - *dgram_cce* - minimizes the categorical-crossentropy between predicted distogram and one extracted from pdb.
+  - *fape*      - minimize difference between coordinates (see AlphaFold paper)
+  - we find *dgram_cce* loss to be more stable for design (compared to *fape*)
+
+- hallucination specific losses
+  - *con*       - maximize `1` contacts per position. `model.set_opt("con",num=1)`
+
+- binder specific losses
+  - *pae* - minimize PAE at interface and within binder
+  - *con* - - maximize `2` contacts per binder position, within binder. `model.set_opt("con",num=2)`
+  - *i_con* - maximize `1` contacts per binder position `model.set_opt("i_con",num=1)`
+
+- partial hallucination specific losses
+  - *sc_fape* - sidechain-specific fape
+
+#### How is contact defined? How do I change it?
+By default, 2 [con]tacts per positions are optimized to be within cβ-cβ < 14.0Å and sequence seperation ≥ 9. This can be changed with:
+```python
+model.set_opt(con=dict(cutoff=8, seqsep=5, num=1))
+```
+For interface:
+```python
+model.set_opt(i_con=dict(...))
+```
+
+#### Optax Optimizers
+By default, we use stochastic gradient descent `set_optimizer(optimizer="sgd", learning_rate=0.1, norm_seq_grad=True)` for optimization. This seems to work quite well for the default problems. But if you want to try other optimizers, ColabDesign is now fully integrated with all [Optax optimizers](https://optax.readthedocs.io/en/latest/api.html).
+Example how to change optimizer. Note the default learning_rate of 0.1 was calibrated for sgd with gradient normalization. A different learning_rate may be more optimal for other optimization settings.
+```python
+model = mk_afdesign_model(optimizer="adam", learning_rate=0.01)
+```
+Or for more control (or to change settings after model initialization), use:
+```python
+model.set_optimizer(optimizer="adam", learning_rate=0.01, b1=0.9, b2=0.999)
+```
+By default, the gradients for the sequence parameters are normalized. We find this helps with convergence. To disable:
+```python
+model.set_opt(norm_seq_grad=False)
+```
+
+# Advanced FAQ
+#### loss during Gradient descent is too jumpy, can I do some kind of greedy search towards the end?
+Gradient descent updates multiple positions each iteration, which can be a little too aggressive during hard (discrete) mode.
+Instead, one can try (`tries`) a few random mutations and accept one with lowest loss. If `use_plddt=True` the random mutations will be biased towards positions with low pLDDT.
+```python
+model.design_3stage(hard_iters=0)
+# set number of model params to evaluate at each iteration
+num_models = 2 if model.args["use_templates"] else 5
+model.design_semigreedy(iters=10, tries=20, num_models=num_models, use_plddt=True)
+```
+#### I was getting better results before the major update (19June2022), how do I revert back to the old settings?
+We are actively trying to find the best weights `model.opt["weights"]`, settings `model.opt` for each protocol.
+Please send us a note if you find something better! To revert back to old settings do this after prepping the model:
+- fixbb:
+```python
+model.set_weights(dgram_cce=1, pae=0.1, plddt=0.1)
+model.design_3stage()
+```
+- hallucination:
+```python
+model.set_seq(mode="gumbel")
+model.set_weights(pae=1, plddt=1, con=0.5)
+model.set_opt("con", binary=True, cutoff=21.6875, num=model._len, seqsep=0)
+model.design_2stage(100, 100, 10)
+```
+- binder hallucination:
+```python
+model.set_weights(plddt=0.1, pae=0.1, i_pae=1.0, con=0.1, i_con=0.5)
+model.set_opt("con", binary=True, cutoff=21.6875, num=model._binder_len, seqsep=0)
+model.set_opt("i_con", binary=True, cutoff=21.6875, num=model._target_len)
+model.design_3stage(100, 100, 10)
+```
+#### I don't like your design_??? function, can I write my own with more detailed control?
+```python
+def design_custom(self):
+  # set options
+  self.set_opt(dropout=True, soft=False, hard=False)
+  # set number of recycles
+  self.set_opt(num_recycles=0)
+  # take 100 steps
+  for _ in range(100): self.step()
+  # increase weight for plddt
+  self.set_weights(plddt=2.0)
+  # take another 100 steps
+  for _ in range(100): self.step()
+  # increase number of recycles
+  self.set_opt(num_recycles=1)
+  # take another 100 steps
+  for _ in range(100): self.step()
+  # etc...
+  
+model = mk_afdesign_model()
+design_custom(model)
+```
+
+#### custom callback examples
+Looking for more control over afdesign? The callback functions have gotten much smarter. Based on your input arguments, it will automatically fetch the variable of interest. You can now define your own custom losses, params to optimize, modify inputs before alphafold is run, and modify auxiliary outputs.
+```python
+def custom_pre_callback(inputs, aux, opt, key):
+  inputs["aatype"] = inputs["aatype"].at[:].set(0)
+  aux["pre"] = opt["pre"] + jax.random.randint(key,[],0,10)
+
+def custom_post_callback(outputs, aux):
+  aux["post"] = outputs["structure_module"]
+
+def custom_loss_callback(outputs, params):
+  loss = jnp.square(outputs["structure_module"]["final_atom14_positions"] + params["custom_param"]).mean()
+  return {"custom_loss":loss}
+
+af_model = mk_afdesign_model(protocol="fixbb",
+                             pre_callback=custom_pre_callback,
+                             post_callback=custom_post_callback,
+                             loss_callback=custom_loss_callback)
+af_model._params["custom_param"] = 1.0
+af_model.opt["weights"]["custom_loss"] = 0.1
+af_model.opt["pre"] = 100
+
+af_model.prep_inputs(pdb_filename=get_pdb("1TEN"), chain="A")
+```
+
+# Previous Updates
+- **24Feb2022** - "Beta" branch started. Refactoring code to allow homooligomeric hallucination/design and averaging gradients across recycles (which is now the default).
+Minor changes changes include renaming intra_pae/inter_con to pae/con and inter_pae/inter_con to i_pae/i_con for clarity.
+- **28Feb2022** - We find backprop through structure module to be unstable, all functions have been updated to only use distogram by default. The definition of contact has changed to minimize entropy within distance cutoff.
+- **02May2022** - The `design.py` code has been split up into multiple python files under `src/`
+- **14May2022** - Adding support for partial hallucination (if you want to constrain one part and generate structure/sequence for rest).
+- **19June2022** - "Beta" branch is now the "Main" branch. WARNING: Lots of default settings and weights were changed. [Click here](#i-was-getting-better-results-before-the-major-update-19june2022-how-do-i-revert-back-to-the-old-settings) for info on how to revert back to old settings. 
+- **28June2022** - v1.0.1 - Major code reorganization/refactoring to add support for callbacks (to allow integration w/ other tools during design) and to avoid clashes with existing trrosetta/alphafold installations. (eg. `af → colabdesign`, `af.src → colabdesign.af` and `alphafold → colabdesign.af.alphafold`).
+- **05July2022** - v1.0.2 - Major code cleanup, removing duplicate code. Adding support for custom loss functions.
+- **11July2022** - v1.0.3 - Improved homo-oligomeric support. RMSD and dgram losses have been refactored to automatically save aligned coordinates. Multimeric coordinates now saved with chain identifiers.
+- **23July2022** - v1.0.4 - Adding support for openfold weights. To enable set `mk_afdesign_model(..., use_openfold=True)`.
+- **31July2022** - v1.0.5 - Refactoring to add support for swapping batch features without recompile. Allowing for implementation of [AF2Rank](https://github.com/sokrypton/ColabDesign/blob/main/af/examples/AF2Rank.ipynb)!
+- **09Sept2022** - v1.0.6
+  - support for alphafold-multimer `model = mk_afdesign_model(..., use_multimer=True)`
+  - support for experimentally resolved loss `model.set_weights(exp_res=1)`
+  - support for multichain design/hallucination for fixbb, hallucination and partial protocols: `model.prep_inputs(..., copies=2)`
+  - support to fix the sequence for certain positions `model.prep_inputs(..., fix_pos="1-10")` (supported in protocols "fixbb" and "partial")
+  - binder protocol improved, prior protocol would try to optimize number of contacts per target, new default is to optimize number of contacts per binder position. Number of contacts per binder position can be controlled with `model.set_opt("i_con",num=1)` and number of positions that should be contact with `model.set_opt("i_con",num_pos=5)`
+  - implementing David Jones'-like protocol for semi-greedy optimization, where positions are selected based on plddt, and after 20 tries, the mutation that decreasing loss the most is accepted. `model.design_semigreedy()`
+  - WARNING: the returned pLDDT is now in the "correct" direction (higher is better)
+  - removing recycle dimension from the input features (to standardize with multimer inputs)
+  - removing all dependence on TensorFlow
+- **14Sept2022** - v1.0.7
+  - refactoring design.py to add `design_pssm_semigreedy()` protocol, which is a wrapper around `design_semigreedy(seq_logits=)`, and can be used to input/learn PSSM for biased optimization.
+  - adding example [peptide_binder_design.ipynb](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/peptide_binder_design.ipynb) targeted for peptide binder hallucination/design.
+  - adding [finer control](#how-do-i-control-which-model-params-are-used-during-design) over what models are used during optimization.
+  - fixing RAM memory leaks, `clear_mem()` now also does garbage collection
+  - fixing integration with TrDesign that got broken in v1.0.6
+- **22Sept2022** - v1.0.8
+  - [custom callback functions](#custom-callback-examples) (\[pre|loss|pos\]_callback) have been refactored to be more flexible.
+    - Supported input arguments include: ["inputs", "outputs", "params", "opt", "seq", "aux", "key"]. 
+    - The pre_callback function can be used to modify inputs before prediction, loss_callback to add cutstom loss.
+  - adding support for [Optax optimizers](#optax-optimizers)
+- **24Sept2022** - v1.0.9
+  - adding [contrib section](/af/contrib) where user contributed modifications and protocols will go.

model/PXDesignBench/ColabDesign/af/contrib/README.md

@@ -0,0 +1 @@
+WIP

model/PXDesignBench/ColabDesign/af/design.ipynb

@@ -0,0 +1,452 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "#AfDesign (v1.1.1)\n",
+        "Backprop through AlphaFold for protein design.\n",
+        "\n",
+        "**WARNING**\n",
+        "1.   This notebook is in active development and was designed for demonstration purposes only.\n",
+        "2.   Using AfDesign as the only \"loss\" function for design might be a bad idea, you may find adversarial sequences (aka. sequences that trick AlphaFold)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UUfKrOzT0gOS"
+      },
+      "source": [
+        "# fixed backbone design (fixbb)\n",
+        "For a given protein backbone, generate/design a new sequence that AlphaFold thinks folds into that conformation. "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qLd1DsnKzxBJ"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"fixbb\")\n",
+        "af_model.prep_inputs(pdb_filename=get_pdb(\"1TEN\"), chain=\"A\")\n",
+        "\n",
+        "print(\"length\",  af_model._len)\n",
+        "print(\"weights\", af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.restart()\n",
+        "af_model.design_3stage()"
+      ],
+      "metadata": {
+        "id": "u0AwskJ84NGx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.plot_traj()  "
+      ],
+      "metadata": {
+        "id": "8FB1v7dn1LL6"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YEApO8YzBoS0"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cW1KQiHKJpfp"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "YDrChASGVUUx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# hallucination\n",
+        "For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts)."
+      ],
+      "metadata": {
+        "id": "qLwS2s_xcjRI"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sZnYfCbfEvol"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "af_model.prep_inputs(length=100)\n",
+        "\n",
+        "print(\"length\",af_model._len)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# pre-design with gumbel initialization and softmax activation\n",
+        "af_model.restart(mode=\"gumbel\")\n",
+        "af_model.design_soft(50)\n",
+        "\n",
+        "# three stage design  \n",
+        "af_model.set_seq(af_model.aux[\"seq\"][\"pseudo\"])\n",
+        "af_model.design_3stage(50,50,10)"
+      ],
+      "metadata": {
+        "id": "f76xqCkw0vj9"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "A1GxeLZdTTya"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "L2E9Tn2Acchj"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YSKWYu0_GlUH"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.get_seqs()"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "dXfm4B8ISLuL"
+      },
+      "source": [
+        "# binder hallucination\n",
+        "For a given protein target and protein binder length, generate/hallucinate a protein binder sequence AlphaFold thinks will bind to the target structure.\n",
+        "To do this, we minimize PAE and maximize number of contacts at the interface and within the binder, and we maximize pLDDT of the binder.\n",
+        "By default, AlphaFold-ptm with residue index offset hack is used. To enable AlphaFold-multimer set: mk_afdesign_model(use_multimer=True).\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3XLJPiRKx5Mw"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"binder\")\n",
+        "af_model.prep_inputs(pdb_filename=get_pdb(\"4MZK\"), chain=\"A\", binder_len=19)\n",
+        "\n",
+        "print(\"target_length\",af_model._target_len)\n",
+        "print(\"binder_length\",af_model._binder_len)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "u6VxjuinyCZa"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.restart()\n",
+        "af_model.design_3stage(100,100,10)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ],
+      "metadata": {
+        "id": "sTlS7_L8Zfwf"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9cARoviGyIKb"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "RzE137NDZdZc"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#**ADVANCED**"
+      ],
+      "metadata": {
+        "id": "SK0IJLoen_EC"
+      }
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "## partial hallucination + custom Radius of Gyration (rg) loss\n",
+        "mix supervised (fixbb) and unsupervised (hallucination) losses to constrain the halluciation process."
+      ],
+      "metadata": {
+        "id": "zl6JGTUzXRnk"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "\n",
+        "# first off, let's implement a custom Radius of Gyration loss function\n",
+        "def rg_loss(inputs, outputs):\n",
+        "  positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n",
+        "  ca = positions[:,residue_constants.atom_order[\"CA\"]]\n",
+        "  center = ca.mean(0)\n",
+        "  rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n",
+        "  rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "  rg = jax.nn.elu(rg - rg_th)\n",
+        "  return {\"rg\":rg}"
+      ],
+      "metadata": {
+        "id": "spec3m8BlGer"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"partial\",\n",
+        "                             loss_callback=rg_loss, # add rg_loss\n",
+        "                             use_templates=False)   # set True to constrain positions using template input\n",
+        "\n",
+        "af_model.opt[\"weights\"][\"rg\"] = 0.1  # optional: specify weight for rg_loss\n",
+        "\n",
+        "af_model.prep_inputs(pdb_filename=get_pdb(\"6MRR\"),\n",
+        "                     chain=\"A\",\n",
+        "                     pos=\"3-30,33-68\",  # define positions to contrain\n",
+        "                     length=100)          # total length if different from input pdb\n",
+        "\n",
+        "af_model.rewire(loops=[36]) # set loop length between segments                     "
+      ],
+      "metadata": {
+        "id": "h_BvzwbAKo6V"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# initialize with wildtype seq, fill in the rest with soft_gumbel distribution\n",
+        "af_model.restart(mode=[\"soft\",\"gumbel\",\"wildtype\"])\n",
+        "af_model.design_3stage(100, 100, 10)"
+      ],
+      "metadata": {
+        "id": "5Unr9u2GYKRD"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ],
+      "metadata": {
+        "id": "BFweaqNWYuF0"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "HTML(af_model.animate())"
+      ],
+      "metadata": {
+        "id": "GSu2lB9HYw-t"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "2EG2t2_KY4Td"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "rTGKbhsI0t8k"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "design.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/RSO.ipynb

@@ -0,0 +1,1297 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/RSO.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "\n",
+        "\n",
+        "\n",
+        "#Protein Design using Relaxed Sequence Optimization\n",
+        "\n",
+        "\n",
+        "**Scalable protein design using optimization in a relaxed sequence space**\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "Christopher Frank, Ali Khoshouei, Lara Fuß, Lara Weber Dominik Schiewitz,Zhixuan Zhao, Motoyuki Hattori, Yosta de Stigter, Shihao Feng, Sergey Ovchinnikov and Hendrik Dietz\n",
+        "\n",
+        "\n",
+        "This notebook contains code to run relaxed sequence optimisation for de novo protein design as described in the manuscript. There are additional options to modify the pipeline according to ones needs\n",
+        "\n",
+        "We recommend using at least an L4 GPU to run this notebook, as the free T4 GPU struggles with larger proteins\n",
+        "\n",
+        "Alternativly a local installation of ColabDesign is strongly recommendet, especially for the design of larger proteins.\n",
+        "\n",
+        "For questions feel free to reach out to the authors\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install pyppeteer nest_asyncio\")\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "import requests, time\n",
+        "if not os.path.isfile(\"TMscore\"):\n",
+        "  os.system(\"wget -qnc https://zhanggroup.org/TM-score/TMscore.cpp\")\n",
+        "  os.system(\"g++ -static -O3 -ffast-math -lm -o TMscore TMscore.cpp\")\n",
+        "def tmscore(x,y):\n",
+        "  # pass to TMscore\n",
+        "  output = os.popen(f'./TMscore {x} {y}')\n",
+        "  # parse outputs\n",
+        "  parse_float = lambda x: float(x.split(\"=\")[1].split()[0])\n",
+        "  o = {}\n",
+        "  for line in output:\n",
+        "    line = line.rstrip()\n",
+        "    if line.startswith(\"RMSD\"): o[\"rms\"] = parse_float(line)\n",
+        "    if line.startswith(\"TM-score\"): o[\"tms\"] = parse_float(line)\n",
+        "    if line.startswith(\"GDT-TS-score\"): o[\"gdt\"] = parse_float(line)\n",
+        "  return o\n",
+        "\n",
+        "import asyncio\n",
+        "import nest_asyncio\n",
+        "from pyppeteer import launch\n",
+        "import base64\n",
+        "\n",
+        "# Apply nest_asyncio to enable nested event loops\n",
+        "nest_asyncio.apply()\n",
+        "\n",
+        "async def fetch_blob_content(page, blob_url):\n",
+        "  blob_to_base64 = \"\"\"\n",
+        "  async (blobUrl) => {\n",
+        "      const blob = await fetch(blobUrl).then(r => r.blob());\n",
+        "      return new Promise((resolve) => {\n",
+        "          const reader = new FileReader();\n",
+        "          reader.onloadend = () => resolve(reader.result);\n",
+        "          reader.readAsDataURL(blob);\n",
+        "      });\n",
+        "  }\n",
+        "  \"\"\"\n",
+        "  base64_data = await page.evaluate(blob_to_base64, blob_url)\n",
+        "  _, encoded = base64_data.split(',', 1)\n",
+        "  return base64.b64decode(encoded)\n",
+        "\n",
+        "async def extract_pdb_file_download_link_and_content(url):\n",
+        "  browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])\n",
+        "  page = await browser.newPage()\n",
+        "  await page.goto(url, {'waitUntil': 'networkidle0'})\n",
+        "  elements = await page.querySelectorAll('a.btn.bg-purple')\n",
+        "  for element in elements:\n",
+        "      href = await page.evaluate('(element) => element.getAttribute(\"href\")', element)\n",
+        "      if 'blob:https://esmatlas.com/' in href:\n",
+        "          content = await fetch_blob_content(page, href)\n",
+        "          await browser.close()\n",
+        "          return href, content\n",
+        "  await browser.close()\n",
+        "  return \"No PDB file link found.\", None\n",
+        "\n",
+        "def esmfold_api(sequence):\n",
+        "  url = f'https://esmatlas.com/resources/fold/result?fasta_header=%3Eunnamed&sequence={sequence}'\n",
+        "  result = asyncio.get_event_loop().run_until_complete(extract_pdb_file_download_link_and_content(url))\n",
+        "  if result[1]:\n",
+        "      pdb_str = result[1].decode('utf-8')\n",
+        "      return pdb_str\n",
+        "  else:\n",
+        "      return \"Failed to retrieve PDB content.\"\n",
+        "\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "from colabdesign.af.alphafold.common import residue_constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sZnYfCbfEvol",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title # Unconditional Generation (Custom)\n",
+        "#@markdown For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n",
+        "LENGTH = 100 #@param {type:\"integer\"}\n",
+        "#@markdown With copies you can specify the number of identical sequences design, resulting in homo oligomers. Copies = 1 is the standard, resulting in a monomer\n",
+        "\n",
+        "COPIES = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\"] {type:\"raw\"}\n",
+        "MODE = \"manuscript\"\n",
+        "\n",
+        "#@markdown Select the losses you want to use. For unconditional generation as reported in the manuscript use all the losses. To increase the diversity of designes remove confidence losses and/or increase the weight of the helix loss.\n",
+        "\n",
+        "use_rg_loss = True #@param {type:\"boolean\"}\n",
+        "#@markdown A too strong rg loss can lead to problems and clashes. Use 0.1 for backbones smaller then 600 AA and 0.01 for larger proteins (0.001 for 1000 AA).\n",
+        "rg_weight = 0.1 #@param {type:\"raw\"}\n",
+        "use_helix_loss = True #@param {type:\"boolean\"}\n",
+        "use_con_loss = True #@param {type:\"boolean\"}\n",
+        "use_confidence_loss = True #@param {type:\"boolean\"}\n",
+        "#@markdown How many halluicnation iteration you want to perform. The standard in the manuscript is 100.\n",
+        "\n",
+        "iters = 50 #@param [\"100\", \"50\", \"30\"] {type:\"raw\"}\n",
+        "\n",
+        "\n",
+        "#@markdown Select if you want to use the 'standard\" ProteinMPNN weights or the soluble ones. The soluble ones usually result in higher in silico as well as experimental sucess, but will increase the negative net charge of the protein which sould potentially interfer with certain protein design problems. The manuscript settings are soluble MPNN\n",
+        "\n",
+        "use_solubleMPNN = True #@param {type:\"boolean\"}\n",
+        "#@markdown Select this to use an experimental ProteinMPNN loss, also backpropagating through ProteinMPNN. This was not used in the manuscript\n",
+        "\n",
+        "use_mpnn_loss = False #@param {type:\"boolean\"}\n",
+        "#@markdown\n",
+        "\n",
+        "def add_rg_loss(self, weight=0.1):\n",
+        "  '''add radius of gyration loss'''\n",
+        "  def loss_fn(inputs, outputs):\n",
+        "    xyz = outputs[\"structure_module\"]\n",
+        "    ca = xyz[\"final_atom_positions\"][:,residue_constants.atom_order[\"CA\"]]\n",
+        "    if self.protocol == \"binder\":\n",
+        "      ca = ca[-self._binder_len:]\n",
+        "    #This uses a scaled version of the rg loss, only looking at every 5th residue\n",
+        "    if MODE == \"manuscript\":\n",
+        "      ca = ca[::5]\n",
+        "    rg = jnp.sqrt(jnp.square(ca - ca.mean(0)).sum(-1).mean() + 1e-8)\n",
+        "\n",
+        "    if MODE == \"original\":\n",
+        "      rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "      rg = jax.nn.elu(rg - rg_th)\n",
+        "    return {\"rg\":rg}\n",
+        "  self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n",
+        "  self.opt[\"weights\"][\"rg\"] = weight\n",
+        "\n",
+        "def add_mpnn_loss(self, mpnn=0.1, mpnn_seq=0.0):\n",
+        "  '''\n",
+        "  add mpnn loss\n",
+        "  mpnn = maximize confidence of proteinmpnn\n",
+        "  mpnn_seq = push designed sequence to match proteinmpnn logits\n",
+        "  '''\n",
+        "\n",
+        "  self._mpnn = mk_mpnn_model(weights = \"soluble\" if use_solubleMPNN else \"original\")\n",
+        "  def loss_fn(inputs, outputs, aux, key):\n",
+        "\n",
+        "    # get structure\n",
+        "    atom_idx = tuple(residue_constants.atom_order[k] for k in [\"N\",\"CA\",\"C\",\"O\"])\n",
+        "    I = {\"S\":           inputs[\"aatype\"],\n",
+        "         \"residue_idx\": inputs[\"residue_index\"],\n",
+        "         \"chain_idx\":   inputs[\"asym_id\"],\n",
+        "         \"X\":           outputs[\"structure_module\"][\"final_atom_positions\"][:,atom_idx],\n",
+        "         \"mask\":        outputs[\"structure_module\"][\"final_atom_mask\"][:,1],\n",
+        "         \"lengths\":     self._lengths,\n",
+        "         \"key\":         key}\n",
+        "\n",
+        "    if \"offset\" in inputs:\n",
+        "      I[\"offset\"] = inputs[\"offset\"]\n",
+        "\n",
+        "    # set autoregressive mask\n",
+        "    L = sum(self._lengths)\n",
+        "    if self.protocol == \"binder\":\n",
+        "      I[\"ar_mask\"] = 1 - np.eye(L)\n",
+        "      I[\"ar_mask\"][-self._len:,-self._len:] = 0\n",
+        "    else:\n",
+        "      I[\"ar_mask\"] = np.zeros((L,L))\n",
+        "\n",
+        "    # get logits\n",
+        "    logits = self._mpnn._score(**I)[\"logits\"][:,:20]\n",
+        "    if self.protocol == \"binder\":\n",
+        "      logits = logits[-self._len:]\n",
+        "    else:\n",
+        "      logits = logits[:self._len]\n",
+        "    aux[\"mpnn_logits\"] = logits\n",
+        "\n",
+        "    # compute loss\n",
+        "    log_q = jax.nn.log_softmax(logits)\n",
+        "    p = inputs[\"seq\"][\"hard\"]\n",
+        "    q = jax.nn.softmax(logits)\n",
+        "    losses = {}\n",
+        "    losses[\"mpnn\"] = -log_q.max(-1).mean()\n",
+        "    losses[\"mpnn_seq\"] = -(p * jax.lax.stop_gradient(log_q)).sum(-1).mean()\n",
+        "    return losses\n",
+        "\n",
+        "  self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n",
+        "  self.opt[\"weights\"][\"mpnn\"] = mpnn\n",
+        "  self.opt[\"weights\"][\"mpnn_seq\"] = mpnn_seq\n",
+        "\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "af_model.prep_inputs(length=LENGTH, copies=COPIES)\n",
+        "\n",
+        "# add extra losses\n",
+        "\n",
+        "if use_mpnn_loss: add_mpnn_loss(af_model)\n",
+        "\n",
+        "print(\"length\",af_model._lengths)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#This cell runs the design loop. Run this in a for loop for design of multiple proteins\n",
+        "\n",
+        "af_model.restart()\n",
+        "af_model.set_seq(mode=[\"gumbel\",\"soft\"])\n",
+        "if use_rg_loss:   add_rg_loss(af_model,rg_weight)\n",
+        "if use_helix_loss : af_model.set_weights(helix=-0.2)\n",
+        "if use_con_loss : af_model.set_weights(con=1.0)\n",
+        "if use_confidence_loss : af_model.set_weights(plddt=0.5, pae=0.5)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])\n",
+        "af_model.design_logits(iters-10)\n",
+        "af_model.design_logits(10, save_best=True)"
+      ],
+      "metadata": {
+        "id": "f76xqCkw0vj9"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "A1GxeLZdTTya"
+      },
+      "outputs": [],
+      "source": [
+        "#This cell plots and saves the results as a pdb file\n",
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "L2E9Tn2Acchj"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YSKWYu0_GlUH"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.get_seqs()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import pandas as pd\n",
+        "#@title # Designability test\n",
+        "#@markdown Test the designability of the backbone, taking in the backbone, generating sequences with solubleMPNN and predicting the sequence with AF2 in single sequence mode.\n",
+        "#@markdown Use Initial Guess (IG) and All Atom Initialisation (AA) for larger proteins\n",
+        "\n",
+        "AA = False #@param {type:\"boolean\"}\n",
+        "IG = False #@param {type:\"boolean\"}\n",
+        "#@markdown NOTE: we remove cysteines from all designed proteins. Additionally for large proteins we also exclude methions to reduce the number of internal start codons\n",
+        "\n",
+        "def designability_test(af_model_test, mpnn_model_test,\n",
+        "                       num_seqs=8, sampling_temp=0.1, num_recycles=3,\n",
+        "                       model_num=4, best_metric=\"rmsd\",\n",
+        "                       in_pdb=\"init.pdb\", out_pdb=\"final.pdb\",\n",
+        "                       verbose=False):\n",
+        "    alphafold_model = f\"model_{model_num}_ptm\"\n",
+        "\n",
+        "    af_model_test.prep_inputs(in_pdb)\n",
+        "    af_model_test.restart(rm_aa=\"C,M\")\n",
+        "    af_model_test._args[\"best_metric\"] = best_metric\n",
+        "    L = sum(af_model_test._lengths)\n",
+        "    mpnn_model_test.get_af_inputs(af_model_test)\n",
+        "    out = mpnn_model_test.sample(num=num_seqs // 8, batch=8,\n",
+        "                                 temperature=sampling_temp)\n",
+        "\n",
+        "    af_terms = [\"plddt\", \"ptm\", \"pae\", \"rmsd\", \"dgram_cce\"]\n",
+        "    for k in af_terms: out[k] = []\n",
+        "\n",
+        "    for n in range(num_seqs):\n",
+        "        seq = out[\"seq\"][n]\n",
+        "        af_model_test.predict(seq=seq,\n",
+        "                              num_recycles=num_recycles,\n",
+        "                              num_models=1,\n",
+        "                              verbose=False,\n",
+        "                              models=alphafold_model)\n",
+        "\n",
+        "        for k in af_terms: out[k].append(af_model_test.aux[\"log\"][k])\n",
+        "        out[\"pae\"][-1] = out[\"pae\"][-1] * 31\n",
+        "        af_model_test._save_results(save_best=True, verbose=verbose)\n",
+        "        af_model_test._k += 1\n",
+        "\n",
+        "    af_model_test.save_pdb(out_pdb)\n",
+        "    labels = [\"score\"] + af_terms + [\"seq\"]\n",
+        "    data = [[out[k][n] for k in labels] for n in range(num_seqs)]\n",
+        "    labels[0] = \"mpnn\"\n",
+        "    df = pd.DataFrame(data, columns=labels)\n",
+        "    return df\n",
+        "\n",
+        "af_model_test = mk_afdesign_model(protocol=\"fixbb\",best_metric=\"rmsd\",use_initial_guess=IG,use_initial_atom_pos=AA,use_templates=False)\n",
+        "mpnn_model_test = mk_mpnn_model(weights=\"soluble\")\n",
+        "\n",
+        "\n",
+        "lowest_rmsd = float('inf')\n",
+        "lowest_rmsd_data = None\n",
+        "\n",
+        "in_pdb = f\"{af_model.protocol}.pdb\"\n",
+        "out_pdb = f\"{af_model.protocol}_out.pdb\"\n",
+        "\n",
+        "\n",
+        "out = designability_test(af_model_test, mpnn_model_test,\n",
+        "                    num_seqs=8, sampling_temp=0.1, num_recycles=3,\n",
+        "                    model_num=4, best_metric=\"rmsd\",\n",
+        "                    in_pdb=in_pdb, out_pdb=out_pdb,\n",
+        "                    verbose=True)\n",
+        "\n"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "Y_-8_Ya4Q8oo"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "#@title # OPTIONAL Unconditional Generation (Manuscript Code)\n",
+        "\n",
+        "#@markdown This code generates a sample of 10 unconditional proteins for lengths between 100 and 800 AA exactly as in the manuscript. For larger proteins CUDA_UNIFIED_MEMORY is needed. This can be done by localy running the code on a CUDA capeable GPU with sufficient memory (A100 80GB e.g.) and running the code with the environment variables XLA_PYTHON_CLIENT_MEM_FRACTION=100.0 TF_FORCE_UNIFIED_MEMORY=1\n",
+        "def rg_loss(inputs, outputs):\n",
+        "  positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n",
+        "  ca = positions[::5,residue_constants.atom_order[\"CA\"]]\n",
+        "  center = ca.mean(0)\n",
+        "  rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n",
+        "  rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "  rg = jax.nn.elu(rg - rg_th)\n",
+        "  return {\"rg\":rg}\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "for length in [100,200,300,400,500,600,700,800]:\n",
+        "  model = mk_afdesign_model(protocol=\"hallucination\",loss_callback=rg_loss)\n",
+        "  model.prep_inputs(length=length)\n",
+        "  print(\"weights\",model.opt[\"weights\"])\n",
+        "  print('Starting up and compiling JAX model....')\n",
+        "\n",
+        "  for i in range(10):\n",
+        "      model.restart(mode=[\"gumbel\", \"soft\"],rm_aa=\"C\")\n",
+        "      model.opt[\"weights\"][\"rg\"] = 0.1\n",
+        "      if length > 600:\n",
+        "        model.opt[\"weights\"][\"rg\"] = 0.01\n",
+        "      #model.opt[\"weights\"]['helix'] = -0.1\n",
+        "      model.opt[\"weights\"]['plddt'] = 1.0\n",
+        "      model.opt[\"weights\"]['pae'] = 1.0\n",
+        "      model.opt[\"weights\"]['helix'] = -0.1\n",
+        "      print(\"weights\", model.opt[\"weights\"])\n",
+        "      model.design_logits(100)\n",
+        "\n",
+        "      #change the output path for local execution\n",
+        "      model.save_pdb(f\"Hallo_{i}.pdb\")"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "gPby7oP25b5i"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown #Redesign with ProteinMPNN for ESMFold prediction\n",
+        "#@markdown The standard manuscript settings were 8 sequences, 0.1 sampling temperature and the removal of cysteines\n",
+        "import pickle\n",
+        "num_seqs = 8 #@param [\"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n",
+        "mpnn_sampling_temp = 0.1 #@param [\"0.0001\", \"0.1\", \"0.15\", \"0.2\", \"0.25\", \"0.3\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
+        "rm_aa = \"C\" #@param {type:\"string\"}\n",
+        "use_solubleMPNN = False #@param {type:\"boolean\"}\n",
+        "#@markdown - `mpnn_sampling_temp` - control diversity of sampled sequences. (higher = more diverse).\n",
+        "#@markdown - `rm_aa='C'` - do not use [C]ysteines.\n",
+        "#@markdown - `use_solubleMPNN` - use weights trained only on soluble proteins.\n",
+        "#@markdown\n",
+        "\n",
+        "from colabdesign.shared.protein import alphabet_list as chain_list\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "mpnn_model.prep_inputs(pdb_filename=f\"{af_model.protocol}.pdb\",\n",
+        "                       chain=\",\".join(chain_list[:COPIES]),\n",
+        "                       homooligmer=COPIES>1,\n",
+        "                       rm_aa=rm_aa,\n",
+        "                       weights = \"soluble\" if use_solubleMPNN else\"original\")\n",
+        "out = mpnn_model.sample(num=num_seqs//8,\n",
+        "                        batch=8,\n",
+        "                        temperature=mpnn_sampling_temp)\n",
+        "for seq,score in zip(out[\"seq\"],out[\"score\"]):\n",
+        "  print(score,seq.split(\"/\")[0])\n",
+        "df = pd.DataFrame(out[\"seq\"])\n",
+        "\n",
+        "# Define the output path for saving the sequences as a .pkl file\n",
+        "output_pkl_file = \"redesigned_sequences.pkl\"\n",
+        "\n",
+        "# Save the DataFrame to a .pkl file\n",
+        "with open(output_pkl_file, 'wb') as f:\n",
+        "    pickle.dump(df, f)"
+      ],
+      "metadata": {
+        "id": "m2qAYsDsCfqJ",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown #Run ESMFold to test designability\n",
+        "#@markdown This cells runs ESMFold from huggingface and automatically calculates the RMSD to the designed backbone\n",
+        "#@markdown NOTE: GPU memory can be a big problem here. If you get memory errors please restart the runtime and run this cell again. It should be self contained. Additionally, after finish the ESMFold prediction rerun the setup cell\n",
+        "\n",
+        "import os\n",
+        "import pandas as pd\n",
+        "from Bio.PDB import PDBParser, Superimposer\n",
+        "import pickle\n",
+        "import torch\n",
+        "import numpy as np\n",
+        "from transformers import AutoTokenizer, EsmForProteinFolding\n",
+        "from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein\n",
+        "from transformers.models.esm.openfold_utils.feats import atom14_to_atom37\n",
+        "output_pkl_file = \"redesigned_sequences.pkl\"\n",
+        "\n",
+        "with open(output_pkl_file, 'rb') as f:\n",
+        "    seq = pickle.load(f)\n",
+        "seq_list = []\n",
+        "for i in np.asarray(seq):\n",
+        "  seq_list.append(i[0])\n",
+        "\n",
+        "pdb_file = \"hallucination.pdb\"\n",
+        "print(seq_list)\n",
+        "\n",
+        "\n",
+        "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esmfold_v1\")\n",
+        "model = EsmForProteinFolding.from_pretrained(\"facebook/esmfold_v1\", low_cpu_mem_usage=True)\n",
+        "\n",
+        "device = 'cuda:0'\n",
+        "model = model.cuda(device)\n",
+        "model.esm = model.esm.half()\n",
+        "model.trunk.set_chunk_size(64)\n",
+        "torch.backends.cuda.matmul.allow_tf32 = True\n",
+        "\n",
+        "def convert_outputs_to_pdb(outputs):\n",
+        "    final_atom_positions = atom14_to_atom37(outputs[\"positions\"][-1], outputs)\n",
+        "    outputs = {k: v.to(\"cpu\").numpy() for k, v in outputs.items()}\n",
+        "    final_atom_positions = final_atom_positions.cpu().numpy()\n",
+        "    final_atom_mask = outputs[\"atom37_atom_exists\"]\n",
+        "    pdbs = []\n",
+        "    for i in range(outputs[\"aatype\"].shape[0]):\n",
+        "        aa = outputs[\"aatype\"][i]\n",
+        "        pred_pos = final_atom_positions[i]\n",
+        "        mask = final_atom_mask[i]\n",
+        "        resid = outputs[\"residue_index\"][i] + 1\n",
+        "        pred = OFProtein(\n",
+        "            aatype=aa,\n",
+        "            atom_positions=pred_pos,\n",
+        "            atom_mask=mask,\n",
+        "            residue_index=resid,\n",
+        "            b_factors=outputs[\"plddt\"][i],\n",
+        "            chain_index=outputs[\"chain_index\"][i] if \"chain_index\" in outputs else None,\n",
+        "        )\n",
+        "        pdbs.append(to_pdb(pred))\n",
+        "    return pdbs\n",
+        "\n",
+        "def calculate_ca_rmsd(pdb_file1, pdb_file2):\n",
+        "    parser = PDBParser(QUIET=True)\n",
+        "\n",
+        "    structure1 = parser.get_structure(\"Protein1\", pdb_file1)\n",
+        "    structure2 = parser.get_structure(\"Protein2\", pdb_file2)\n",
+        "\n",
+        "    ca_atoms1 = [atom for atom in structure1.get_atoms() if atom.get_name() == \"CA\"]\n",
+        "    ca_atoms2 = [atom for atom in structure2.get_atoms() if atom.get_name() == \"CA\"]\n",
+        "\n",
+        "    super_imposer = Superimposer()\n",
+        "    super_imposer.set_atoms(ca_atoms1, ca_atoms2)\n",
+        "    super_imposer.apply(structure2.get_atoms())\n",
+        "    rmsd = super_imposer.rms\n",
+        "    return rmsd\n",
+        "\n",
+        "def process_sequences(seq_list, pdb_file):\n",
+        "    lowest_rmsd = float('inf')\n",
+        "    lowest_rmsd_data = None\n",
+        "    out_ss_path = \"./output\"\n",
+        "\n",
+        "    if not os.path.exists(out_ss_path):\n",
+        "        os.mkdir(out_ss_path)\n",
+        "\n",
+        "    for test_protein in seq_list:\n",
+        "        data = {}\n",
+        "        tokenized_input = tokenizer([test_protein], return_tensors=\"pt\", add_special_tokens=False)['input_ids']\n",
+        "        tokenized_input = tokenized_input.cuda(device)\n",
+        "\n",
+        "        with torch.no_grad():\n",
+        "            output = model(tokenized_input)\n",
+        "\n",
+        "        data['out'] = output\n",
+        "        data[\"plddt\"] = torch.mean(output['plddt']).item()\n",
+        "        data['pae'] = torch.mean(output['predicted_aligned_error']).item()\n",
+        "\n",
+        "        pdb_data = convert_outputs_to_pdb(output)\n",
+        "        tmp_pdb_file = os.path.join(out_ss_path, \"TMP.pdb\")\n",
+        "\n",
+        "        with open(tmp_pdb_file, 'w') as file:\n",
+        "            for line in pdb_data:\n",
+        "                file.write(line)\n",
+        "\n",
+        "        data['rmsd'] = calculate_ca_rmsd(tmp_pdb_file, pdb_file)\n",
+        "        print(f'Sequence: {test_protein}, plddt: {data[\"plddt\"]}, PAE: {data[\"pae\"]}, RMSD: {data[\"rmsd\"]}')\n",
+        "\n",
+        "        if data['rmsd'] < lowest_rmsd:\n",
+        "            lowest_rmsd = data['rmsd']\n",
+        "            lowest_rmsd_data = data\n",
+        "\n",
+        "    if lowest_rmsd_data is not None:\n",
+        "        print(f'Lowest RMSD: {lowest_rmsd}')\n",
+        "        best_pdb_data = convert_outputs_to_pdb(lowest_rmsd_data['out'])\n",
+        "        best_pdb_file = os.path.join(out_ss_path, \"best_structure.pdb\")\n",
+        "\n",
+        "        with open(best_pdb_file, 'w') as file:\n",
+        "            for line in best_pdb_data:\n",
+        "                file.write(line)\n",
+        "\n",
+        "        original_dict = lowest_rmsd_data\n",
+        "        key_to_exclude = 'out'\n",
+        "        data_out = {k: v for k, v in original_dict.items() if k != key_to_exclude}\n",
+        "\n",
+        "        with open(os.path.join(out_ss_path, \"best_structure_data.pkl\"), 'wb') as f:\n",
+        "            pickle.dump(data_out, f)\n",
+        "\n",
+        "        return lowest_rmsd, best_pdb_file, data_out\n",
+        "\n",
+        "    return None, None, None\n",
+        "\n",
+        "\n",
+        "\n",
+        "lowest_rmsd, best_pdb_file, best_data = process_sequences(seq_list, pdb_file)\n",
+        "if lowest_rmsd is not None:\n",
+        "    print(f\"Lowest RMSD: {lowest_rmsd}, Best PDB file: {best_pdb_file}\")\n",
+        "else:\n",
+        "    print(\"No valid result found.\")\n"
+      ],
+      "metadata": {
+        "id": "Ey29NmNAFtK0",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title # Heterodimer Design Prep\n",
+        "#@markdown Design a set of heterodimeric proteins with two chains making a complex. The settings are excatly the ones used in the manuscript to design the heterodimer binders.\n",
+        "LENGTH1 = 100 #@param {type:\"integer\"}\n",
+        "LENGTH2 = 100 #@param {type:\"integer\"}\n",
+        "\n",
+        "#@markdown ProteinMPNN Settings\n",
+        "use_solubleMPNN = True #@param {type:\"boolean\"}\n",
+        "#@markdown\n",
+        "\n",
+        "\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "\n",
+        "def hd_loss(inputs, outputs):\n",
+        "  positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n",
+        "  ca1 = positions[:LENGTH1, residue_constants.atom_order[\"CA\"]]\n",
+        "  center1 = ca1.mean(0)\n",
+        "  rg1 = jnp.sqrt(jnp.square(ca1 - center1).sum(-1).mean() + 1e-8)\n",
+        "  rg_th = 2.38 * ca1.shape[0] ** 0.365\n",
+        "  rg1 = jax.nn.elu(rg1 - rg_th)\n",
+        "\n",
+        "\n",
+        "  ca2 = positions[LENGTH2:, residue_constants.atom_order[\"CA\"]]\n",
+        "  center2 = ca2.mean(0)\n",
+        "  rg2 = jnp.sqrt(jnp.square(ca2 - center2).sum(-1).mean() + 1e-8)\n",
+        "  rg_th = 2.38 * ca2.shape[0] ** 0.365\n",
+        "  rg2 = jax.nn.elu(rg2 - rg_th)\n",
+        "\n",
+        "\n",
+        "\n",
+        "  return {\"hd\":rg1+rg2}\n",
+        "\n",
+        "total_length = LENGTH1 + LENGTH2\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\", loss_callback=hd_loss)\n",
+        "af_model.prep_inputs(length=total_length)\n",
+        "af_model._inputs['residue_index'][LENGTH1:] = np.arange(LENGTH2) + 50 + LENGTH1\n",
+        "# add extra losses\n",
+        "af_model.restart(mode=[\"gumbel\", \"soft\"])\n",
+        "af_model.opt[\"weights\"][\"hd\"] = 0.1\n",
+        "af_model.opt[\"weights\"]['plddt'] = 1.0\n",
+        "af_model.opt[\"weights\"]['pae'] = 1.0\n",
+        "af_model.opt[\"weights\"]['helix'] = -0.5\n",
+        "print(\"weights\", af_model.opt[\"weights\"])\n",
+        "print('Starting up and compiling JAX model....')\n"
+      ],
+      "metadata": {
+        "id": "88HCs0ID6CGc",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title # Run Design\n",
+        "af_model.design_logits(100)\n",
+        "af_model.save_pdb(\"Heterodimer.pdb\")"
+      ],
+      "metadata": {
+        "id": "5KVtj6LM7RQ4"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.save_pdb(\"Heterodimer.pdb\")\n",
+        "af_model.plot_pdb()"
+      ],
+      "metadata": {
+        "id": "qfOEMo4qpWfm"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title # Design Sequence using Homooligomer Filter\n",
+        "#@markdown We first test if the two protomers are predictd to fold into a high confidence protein on their own, removing proteins that are not likely to be expressed on their own. Then we predict the heterodimer using the AF multimer model. Generally the AF multimer model has a hard time predicting de novo designed proteins. This is why we use templates and remove any interchain information. Finally we predict each individual protomer with a copy of itself, testing for homooligomerisation.\n",
+        "\n",
+        "\n",
+        "file_path =\"Heterodimer.pdb\"\n",
+        "\n",
+        "folder_path = \"/content/\"\n",
+        "\n",
+        "######## make A - B chain file\n",
+        "\n",
+        "from Bio.PDB import PDBParser, PDBIO, Chain\n",
+        "\n",
+        "# Set the input and output PDB file names\n",
+        "input_pdb_file = file_path\n",
+        "if not os.path.exists(os.path.join(folder_path, 'AB')):\n",
+        "    os.mkdir(os.path.join(folder_path, 'AB'))\n",
+        "output_pdb_file = os.path.join(folder_path, 'AB',\"Heterodimer.pdb\")\n",
+        "\n",
+        "# Create a PDB parser and read the input PDB file\n",
+        "parser = PDBParser()\n",
+        "structure = parser.get_structure(\"input_structure\", input_pdb_file)\n",
+        "\n",
+        "# Find the initial chain id\n",
+        "initial_chain_id = None\n",
+        "for chain in structure[0]:\n",
+        "    initial_chain_id = chain.get_id()\n",
+        "    break\n",
+        "\n",
+        "# Create new chains A and B\n",
+        "chain_A = Chain.Chain(\"A\")\n",
+        "chain_B = Chain.Chain(\"B\")\n",
+        "\n",
+        "# Iterate over the residues in the original chain\n",
+        "for residue in structure[0][initial_chain_id]:\n",
+        "    res_id = residue.get_id()[1]\n",
+        "\n",
+        "    # Add residues 1-200 to chain A\n",
+        "    if 1 <= res_id <= 100:\n",
+        "        chain_A.add(residue.copy())\n",
+        "\n",
+        "    # Add residues 201-400 to chain B\n",
+        "    elif 151 <= res_id <= 450:\n",
+        "        chain_B.add(residue.copy())\n",
+        "\n",
+        "# Remove the existing chain\n",
+        "for model in structure:\n",
+        "    model.detach_child(initial_chain_id)\n",
+        "\n",
+        "# Add the new chains to the model\n",
+        "structure[0].add(chain_A)\n",
+        "structure[0].add(chain_B)\n",
+        "\n",
+        "# Save the modified structure to a new PDB file\n",
+        "io = PDBIO()\n",
+        "io.set_structure(structure)\n",
+        "io.save(output_pdb_file)\n",
+        "\n",
+        "\n",
+        "clear_mem()\n",
+        "he_model = mk_afdesign_model(protocol=\"fixbb\", use_templates=True, use_multimer=True)\n",
+        "\n",
+        "\n",
+        "ho_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "ho_model.prep_inputs(length=LENGTH1, copies=2)\n",
+        "\n",
+        "ho_model.set_weights(i_pae=1.0)\n",
+        "s_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "s_model.prep_inputs(length=LENGTH2)\n",
+        "mpnn_model = mk_mpnn_model(weights=\"soluble\")\n",
+        "\n",
+        "\n",
+        "mpnn_model.prep_inputs(pdb_filename=output_pdb_file, chain='A,B',rm_aa=\"C\")\n",
+        "samples = mpnn_model.sample_parallel(8)\n",
+        "\n",
+        "he_model.prep_inputs(pdb_filename=output_pdb_file, chain='A,B',rm_template_ic=True)\n",
+        "he_model._inputs['residue_index'][LENGTH1:] = np.arange(LENGTH2) + 50 + LENGTH1\n",
+        "\n",
+        "k = 0\n",
+        "for seq in samples['seq']:\n",
+        "    print('Predicting Protomer 1...')\n",
+        "    s_model.predict(seq=seq[:LENGTH1], num_recycles=3)\n",
+        "    plddt1 = s_model.aux['losses']['plddt']\n",
+        "    print('Predicting Protomer 2...')\n",
+        "    s_model.predict(seq=seq[LENGTH1+1:], num_recycles=3)\n",
+        "    plddt2 = s_model.aux['losses']['plddt']\n",
+        "    k = k + 1\n",
+        "    if plddt1 < 0.20 and plddt2 < 0.20:\n",
+        "        print('Passed Protomer Check! Predicting Heterodimer...')\n",
+        "        he_model.predict(seq=''.join([seq[:LENGTH1], seq[LENGTH1+1:]]), num_recycles=3)\n",
+        "\n",
+        "        if he_model.aux['losses']['plddt'] < 0.15 and he_model.aux['losses']['rmsd'] < 2.0:\n",
+        "            print('Passed Heterodimer Check! Predicting Homodimer 1...')\n",
+        "            ho_model.predict(seq=seq[:LENGTH1],num_recycles=3)\n",
+        "            print('Predicting Homodimer 2...')\n",
+        "            ipae1 = ho_model.aux['losses']['i_pae']\n",
+        "            ho_model.predict(seq=seq[LENGTH1+1:],num_recycles=3)\n",
+        "            ipae2 = ho_model.aux['losses']['i_pae']\n",
+        "            if ipae1 > 0.8 and ipae2 > 0.8:\n",
+        "              print('Passed Homodimer check!')\n",
+        "              he_model.save_pdb(f'Heterodimer_seq_{k}.pdb')\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "mII01sjwhAqX",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "def add_rg_loss(self, weight=0.1):\n",
+        "  '''add radius of gyration loss'''\n",
+        "  def loss_fn(inputs, outputs):\n",
+        "    xyz = outputs[\"structure_module\"]\n",
+        "    ca = xyz[\"final_atom_positions\"][:,residue_constants.atom_order[\"CA\"]]\n",
+        "\n",
+        "    ca = ca[-self._binder_len:]\n",
+        "\n",
+        "    rg = jnp.sqrt(jnp.square(ca - ca.mean(0)).sum(-1).mean() + 1e-8)\n",
+        "    rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "    rg = jax.nn.elu(rg - rg_th)\n",
+        "    return {\"rg\":rg}\n",
+        "  self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n",
+        "  self.opt[\"weights\"][\"rg\"] = weight\n",
+        "\n",
+        "\n",
+        "\n",
+        "#@title # Binder Design\n",
+        "#@markdown For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n",
+        "LENGTH = 100 #@param {type:\"integer\"}\n",
+        "binder_pdb = '5NGV' #@param {type:\"string\"}\n",
+        "binder_chain ='A' #@param {type:\"string\"}\n",
+        "hotspot ='' #@param {type:\"string\"}\n",
+        "if hotspot == \"\": hotspot = None\n",
+        "#@markdown ProteinMPNN Settings\n",
+        "use_solubleMPNN = True #@param {type:\"boolean\"}\n",
+        "#@markdown\n",
+        "\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"binder\")\n",
+        "add_rg_loss(af_model)\n",
+        "af_model.prep_inputs(pdb_filename=get_pdb(binder_pdb), chain=binder_chain,hotspot=hotspot, binder_len=LENGTH)\n",
+        "\n",
+        "\n",
+        "af_model.restart(mode=[\"gumbel\", \"soft\"])\n",
+        "\n",
+        "af_model.opt[\"weights\"][\"rg\"] = 0.5\n",
+        "\n",
+        "af_model.opt[\"weights\"]['helix'] = -0.2\n",
+        "af_model.opt[\"weights\"]['plddt'] = 0.1\n",
+        "af_model.opt[\"weights\"]['pae'] = 0.1\n",
+        "af_model.opt[\"weights\"]['i_pae'] = 0.1\n",
+        "af_model.opt[\"weights\"]['i_con'] = 2.0\n",
+        "\n",
+        "print(\"weights\", af_model.opt[\"weights\"])\n",
+        "print('Starting up and compiling JAX model....')\n"
+      ],
+      "metadata": {
+        "id": "H4WqduyJ785H",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.design_logits(100)\n",
+        "af_model.save_pdb(\"Binder.pdb\")"
+      ],
+      "metadata": {
+        "id": "bS83MeVd99FN"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.plot_pdb()"
+      ],
+      "metadata": {
+        "id": "UQMYKLCNZK2L"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "\n",
+        "#@title # Binder Sequence Design with AF Multimer filtering\n",
+        "#@markdown Use this to generate sequences for the binder candidate generated in the previous step\n",
+        "\n",
+        "#@markdown First we use the AF2 PTM model to predict the binder without receptor, acting as a fast pre filter. Then we use the AF Multimer model to predict the Receptor Binder complex. Again we use a template for the binder to help AF Multimer predicting the de novo designed protein\n",
+        "\n",
+        "binder_model = mk_afdesign_model(protocol=\"binder\",use_multimer=True,use_initial_guess=True)\n",
+        "hall_model = mk_afdesign_model(protocol=\"fixbb\")\n",
+        "\n",
+        "\n",
+        "binder_model.set_weights(i_pae=1.0)\n",
+        "\n",
+        "\n",
+        "mpnn_model = mk_mpnn_model(weights=\"soluble\")\n",
+        "mpnn_model.prep_inputs(pdb_filename=\"Binder.pdb\", chain='A,B', fix_pos='A',rm_aa=\"C\")\n",
+        "\n",
+        "samples = mpnn_model.sample_parallel(8,temperature=0.01)\n",
+        "hall_model.prep_inputs(pdb_filename=\"Binder.pdb\", chain='B')\n",
+        "binder_model.prep_inputs(pdb_filename=\"Binder.pdb\", chain='A', binder_chain='B',use_binder_template=True,rm_template_ic=True)\n",
+        "k=0\n",
+        "for seq in samples['seq']:\n",
+        "    print(\"Predicting binder only\")\n",
+        "    hall_model.predict(seq=seq[-LENGTH:], num_recycles=3)\n",
+        "    if hall_model.aux['losses']['rmsd'] < 2.0 :\n",
+        "        print(\"Passed! Predicting binder with receptor using AF Multimer\")\n",
+        "        binder_model.predict(seq=seq[-LENGTH:], num_recycles=3)\n",
+        "        plddt1 = binder_model.aux['losses']['plddt']\n",
+        "        i_pae = binder_model.aux['losses']['i_pae']\n",
+        "        if plddt1 < 0.15 and i_pae < 0.4:\n",
+        "           print(f\"Passed! Final I_PAE is {i_pae*31}\")\n",
+        "           binder_model.save_pdb(f'Binder_seq_{k}.pdb')\n",
+        "           binder_model.plot_pdb()\n",
+        "\n",
+        "    k = k + 1"
+      ],
+      "metadata": {
+        "id": "QfrRhvwTAyK-",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title # Site scaffolding example\n",
+        "#@markdown This cell provides the code to perform the site scaffolding in bulk.\n",
+        "#@markdown Just go to the commented section with names, contigs and length to insert the desired PDB identifier, contigs and final size and start designing.\n",
+        "#@markdown Num_designs controls how many backbones one designes per PDB file\n",
+        "\n",
+        "num_designs = 1 #@param {type:\"integer\"}\n",
+        "\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "\n",
+        "\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "import contextlib\n",
+        "\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "import pickle\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "\n",
+        "import re\n",
+        "import os\n",
+        "\n",
+        "\n",
+        "#Add the names of the PDB files for the scaffolding problem here\n",
+        "names = [\n",
+        "    \"1PRW\"\n",
+        "]\n",
+        "print(len(names))\n",
+        "#Add the design contigs here\n",
+        "\n",
+        "inputs = [\n",
+        "    \"5-20,A16-35,10-25,A52-71,5-20\"\n",
+        "]\n",
+        "#Add the total length here. We only use the maximum length specified\n",
+        "lengths = [\n",
+        "    \"60-105\"\n",
+        "]\n",
+        "\n",
+        "\n",
+        "def rg_loss(inputs, outputs):\n",
+        "    positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n",
+        "    ca = positions[::5, residue_constants.atom_order[\"CA\"]]\n",
+        "    center = ca.mean(0)\n",
+        "    rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n",
+        "    rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "    rg = jax.nn.elu(rg - rg_th)\n",
+        "    return {\"rg\": rg}\n",
+        "\n",
+        "\n",
+        "clear_mem()\n",
+        "\n",
+        "for _name, _input, _length in zip(\n",
+        "    names, inputs, lengths\n",
+        "):\n",
+        "    print(f\"Starting on {_name}\")\n",
+        "    _input = _input.replace(\" \", \"\")\n",
+        "    __name = _name.split(\"_\")[0]\n",
+        "    model = mk_afdesign_model(\n",
+        "        protocol=\"partial\"\n",
+        "    )\n",
+        "    wire_loop_repr = [\"l\" if re.search(\"[A-Z]\", x) else \"w\" for x in _input.split(\",\")]\n",
+        "\n",
+        "    _lengths = []\n",
+        "    for _id, rep in zip(wire_loop_repr, _input.split(\",\")):\n",
+        "        if \"-\" in rep:  # loop or range\n",
+        "            if _id == \"l\":  # loop\n",
+        "                rep = rep[1:]\n",
+        "                _len = int(rep.split(\"-\")[1]) - int(rep.split(\"-\")[0]) + 1\n",
+        "            else:  # range\n",
+        "                _len = int(rep.split(\"-\")[1])\n",
+        "        else:\n",
+        "            if _id == \"l\":\n",
+        "                rep = rep[1:]\n",
+        "            _len = int(1)\n",
+        "        _lengths.append(_len)\n",
+        "\n",
+        "    overall_length = sum(_lengths)\n",
+        "    print(overall_length)\n",
+        "    old_pos = list(filter(lambda x: re.search(\"[A-Z]\", x), _input.split(\",\")))\n",
+        "    order = list(range(len(old_pos)))\n",
+        "    old_pos = \",\".join(old_pos)\n",
+        "    wires = list(filter(lambda x: not re.search(\"[A-Z]\", x), _input.split(\",\")))\n",
+        "    wires = [\n",
+        "        int(wire) if \"-\" not in wire else int(wire.split(\"-\")[1]) for wire in wires\n",
+        "    ]\n",
+        "    offset = wires[0] if not wire_loop_repr[0] == \"l\" else 0\n",
+        "    if wire_loop_repr[0] == \"w\":\n",
+        "        wires = wires[1:]\n",
+        "    if wire_loop_repr[-1] == \"w\":\n",
+        "        wires = wires[:-1]\n",
+        "\n",
+        "    chain = re.findall(\"[A-Z]\", _input)\n",
+        "    chain = list(set(chain))\n",
+        "    assert len(chain) == 1\n",
+        "    chain = chain[0]\n",
+        "    if \"-\" in _length:\n",
+        "        _length = _length.split(\"-\")[1]\n",
+        "    _length = int(_length)\n",
+        "    if _length < overall_length:\n",
+        "        _length = overall_length\n",
+        "\n",
+        "    debug = False\n",
+        "    if debug:\n",
+        "        print(\"chain  \" + str(chain))\n",
+        "        print(\"old_pos  \" + str(old_pos))\n",
+        "        print(\"wires  \" + str(wires))\n",
+        "        print(\"offset  \" + str(offset))\n",
+        "        print(\"_length  \" + str(_length))\n",
+        "        print(\"order  \" + str(order))\n",
+        "    print(_name)\n",
+        "    pdb_file = get_pdb(_name)\n",
+        "    model.prep_inputs(\n",
+        "        pdb_file,\n",
+        "        chain=chain,\n",
+        "        pos=old_pos,\n",
+        "        length=_length,\n",
+        "        fix_seq=True,\n",
+        "    )\n",
+        "\n",
+        "\n",
+        "    model.rewire(\n",
+        "        order=order,  # set order of segments\n",
+        "        loops=wires,  # change loop length inbetween segments\n",
+        "        offset=offset,\n",
+        "    )  # essentially loop length at the N term\n",
+        "\n",
+        "    print(\"   Starting up and compiling JAX model....\")\n",
+        "\n",
+        "    for i in range(num_designs):\n",
+        "        print(f\"      Iteration {i} of 100\")\n",
+        "        model.restart(mode=[\"gumbel\", \"soft\"], rm_aa=\"C\")\n",
+        "        model.opt[\"weights\"][\"rg\"] = 0.1\n",
+        "        model.opt[\"weights\"][\"dgram_cce\"] = 2.0\n",
+        "        model.opt[\"weights\"][\"plddt\"] = 0.1\n",
+        "        model.opt[\"weights\"][\"pae\"] = 0.1\n",
+        "        model.opt[\"weights\"][\"rmsd\"] = 1.0\n",
+        "        model.opt[\"weights\"]['sc_rmsd'] = 1.0\n",
+        "        #           model.opt[\"weights\"]['fape'] = 1.0\n",
+        "\n",
+        "        model.design_logits(190)\n",
+        "        model.design_logits(10, save_best=True)\n",
+        "        outfile = f\"out_sc/{_name}_resesigned/{_name}_redesigned_{i}.pdb\"\n",
+        "        os.makedirs(os.path.dirname(outfile), exist_ok=True)\n",
+        "        model.save_pdb(outfile)\n",
+        "    mpnn_model = mk_mpnn_model()\n",
+        "\n",
+        "    p = (\n",
+        "        []\n",
+        "    )  # [homo if not n in _interfaceFixturesIndexSecChain else hetero for n, (homo, hetero) in enumerate(zip(list(ho2), list(he[-len(ho2):])))]\n",
+        "    for k in model.opt[\"pos\"]:\n",
+        "        p.append(str(k + 1))  # Might be wrong\n",
+        "        p.append(\",\")\n",
+        "    posf = \"\".join(p[:-1])\n",
+        "\n",
+        "    repredictionModel = mk_afdesign_model(\n",
+        "        protocol=\"fixbb\", use_templates=False\n",
+        "    )\n",
+        "    os.makedirs(os.path.dirname('out_sc_Redesigned/'), exist_ok=True)\n",
+        "\n",
+        "    for j in range(num_designs):\n",
+        "        print(f\"      Reprediction Iteration {j} of 100\")\n",
+        "        repredictionModel.prep_inputs(\n",
+        "            f\"out_sc/{_name}_resesigned/{_name}_redesigned_{j}.pdb\"\n",
+        "        )\n",
+        "\n",
+        "        mpnn_model.prep_inputs(\n",
+        "            pdb_filename=f\"out_sc/{_name}_resesigned/{_name}_redesigned_{j}.pdb\",\n",
+        "            chain=\"A\",\n",
+        "            fix_pos=posf,\n",
+        "            rm_aa=\"C\",\n",
+        "        )\n",
+        "        out = mpnn_model.sample(num=1, batch=8, temperature=0.1)\n",
+        "\n",
+        "        for n, i in enumerate(out[\"seq\"]):\n",
+        "\n",
+        "            repredictionModel.predict(seq=i, num_recycles=3)\n",
+        "            if (\n",
+        "                repredictionModel.aux[\"log\"][\"rmsd\"] < 2.0\n",
+        "                and repredictionModel.aux[\"log\"][\"plddt\"] > 0.85\n",
+        "            ):\n",
+        "                filename = f'out_sc_Redesigned/{_name}_resesigned/{_name}_redesigned-{j}_num-{n}_rmsd-{int(repredictionModel.aux[\"log\"][\"rmsd\"]*100)}.pdb'\n",
+        "                os.makedirs(os.path.dirname(filename), exist_ok=True)\n",
+        "                repredictionModel.save_pdb(filename)\n",
+        "\n",
+        "for _name, _input, _length in zip(\n",
+        "    names, inputs, lengths\n",
+        "):\n",
+        "    print(f\"Starting on {_name}\")\n",
+        "    clear_mem()\n",
+        "\n",
+        "    _input = _input.replace(\" \", \"\")\n",
+        "    __name = _name.split(\"_\")[0]\n",
+        "\n",
+        "    test_model = mk_afdesign_model(protocol='fixbb')\n",
+        "    model = mk_afdesign_model(\n",
+        "        protocol=\"partial\", use_templates=False\n",
+        "    )  # set True to constrain positions using template input\n",
+        "    # define positions we want to constrain (input PDB numbering)\n",
+        "\n",
+        "    wire_loop_repr = [\"l\" if re.search(\"[A-Z]\", x) else \"w\" for x in _input.split(\",\")]\n",
+        "\n",
+        "    _lengths = []\n",
+        "    for _id, rep in zip(wire_loop_repr, _input.split(\",\")):\n",
+        "        if \"-\" in rep:  # loop or range\n",
+        "            if _id == \"l\":  # loop\n",
+        "                rep = rep[1:]\n",
+        "                _len = int(rep.split(\"-\")[1]) - int(rep.split(\"-\")[0]) + 1\n",
+        "            else:  # range\n",
+        "                _len = int(rep.split(\"-\")[1])\n",
+        "        else:\n",
+        "            if _id == \"l\":\n",
+        "                rep = rep[1:]\n",
+        "            _len = 1\n",
+        "        _lengths.append(_len)\n",
+        "\n",
+        "    overall_length = sum(_lengths)\n",
+        "\n",
+        "    old_pos = list(filter(lambda x: re.search(\"[A-Z]\", x), _input.split(\",\")))\n",
+        "    order = list(range(len(old_pos)))\n",
+        "    old_pos = \",\".join(old_pos)\n",
+        "    wires = list(filter(lambda x: not re.search(\"[A-Z]\", x), _input.split(\",\")))\n",
+        "    wires = [\n",
+        "        int(wire) if \"-\" not in wire else int(wire.split(\"-\")[1]) for wire in wires\n",
+        "    ]\n",
+        "    offset = wires[0] if not wire_loop_repr[0] == \"l\" else 0\n",
+        "    if wire_loop_repr[0] == \"w\":\n",
+        "        wires = wires[1:]\n",
+        "    if wire_loop_repr[-1] == \"w\":\n",
+        "        wires = wires[:-1]\n",
+        "\n",
+        "    chain = re.findall(\"[A-Z]\", _input)\n",
+        "    chain = list(set(chain))\n",
+        "    assert len(chain) == 1\n",
+        "    chain = chain[0]\n",
+        "    if \"-\" in _length:\n",
+        "        _length = _length.split(\"-\")[1]\n",
+        "    _length = int(_length)\n",
+        "    if _length < overall_length:\n",
+        "        _length = overall_length\n",
+        "\n",
+        "\n",
+        "    print(_name)\n",
+        "    pdb_file = get_pdb(_name)\n",
+        "\n",
+        "\n",
+        "\n",
+        "    model.prep_inputs(\n",
+        "        pdb_file,\n",
+        "        chain=chain,\n",
+        "        pos=old_pos,  # define positions to contrain\n",
+        "        length=_length,  # define if the desired length is different from input PDB\n",
+        "        fix_seq=True,\n",
+        "    )  # set True to constrain the sequence\n",
+        "\n",
+        "    # set positions (if different from PDB)\n",
+        "    # reorder the segments,\n",
+        "    model.rewire(\n",
+        "        order=order,  # set order of segments\n",
+        "        loops=wires,  # change loop length inbetween segments\n",
+        "        offset=offset,\n",
+        "    )  # essentially loop length at the N term\n",
+        "\n",
+        "    in_files = os.listdir(f'out_sc_Redesigned/{_name}_resesigned/')\n",
+        "    if not os.path.exists(f'out_sc_Redesigned/{_name}_resesigned/out/'):\n",
+        "        os.mkdir(f'out_sc_Redesigned/{_name}_resesigned/out/')\n",
+        "    for ii in in_files:\n",
+        "        if ii[-1] == 'b':\n",
+        "            test_model.prep_inputs(pdb_filename=f'out_sc_Redesigned/{_name}_resesigned/{ii}')\n",
+        "            seq = test_model._inputs['batch'][\"aatype\"]\n",
+        "            #print(seq)\n",
+        "            model.predict(seq=seq, num_recycles=3)\n",
+        "            if model.aux[\"losses\"][\"rmsd\"] < 1.0:\n",
+        "                model.save_pdb(f'out_sc_Redesigned/{_name}_resesigned/out/{ii}')\n",
+        "                with open(f'out_sc_Redesigned/{_name}_resesigned/out/{ii[:-4]}_data.pkl', 'wb') as f:\n",
+        "                    pickle.dump(model.aux[\"losses\"][\"rmsd\"], f)\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "8sLixVvGGJ47",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "QHFiJeHjB4sM"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": [],
+      "machine_shape": "hm",
+      "gpuType": "L4",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/af2cycler.ipynb

@@ -0,0 +1,741 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "machine_shape": "hm",
+      "gpuType": "L4",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "accelerator": "GPU"
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/af2cycler.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# The AF2cycler\n",
+        "This notebook contains the code to run the af2cycler and use for improvement of suboptimal designed protein backbones.\n",
+        "Based on:\n",
+        "\n",
+        "**Alphafold2 refinement improves designability of large de novo proteins**\n",
+        "\n",
+        "Christopher Frank, Dominik Schiwietz, Lara Fuß, Sergey Ovchinnikov and Hendrik Dietz\n",
+        "\n",
+        "We recommend to run this Notebook with at leat a L4 or better a A100 GPU as the GPU memeory needed for ESMFold is quite significant\n",
+        "\n"
+      ],
+      "metadata": {
+        "id": "wsji9enfo8aM"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "yv2S1XouoxZF"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install pyppeteer nest_asyncio\")\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "import requests, time\n",
+        "if not os.path.isfile(\"TMscore\"):\n",
+        "  os.system(\"wget -qnc https://zhanggroup.org/TM-score/TMscore.cpp\")\n",
+        "  os.system(\"g++ -static -O3 -ffast-math -lm -o TMscore TMscore.cpp\")\n",
+        "def tmscore(x,y):\n",
+        "  # pass to TMscore\n",
+        "  output = os.popen(f'./TMscore {x} {y}')\n",
+        "  # parse outputs\n",
+        "  parse_float = lambda x: float(x.split(\"=\")[1].split()[0])\n",
+        "  o = {}\n",
+        "  for line in output:\n",
+        "    line = line.rstrip()\n",
+        "    if line.startswith(\"RMSD\"): o[\"rms\"] = parse_float(line)\n",
+        "    if line.startswith(\"TM-score\"): o[\"tms\"] = parse_float(line)\n",
+        "    if line.startswith(\"GDT-TS-score\"): o[\"gdt\"] = parse_float(line)\n",
+        "  return o\n",
+        "\n",
+        "import asyncio\n",
+        "import nest_asyncio\n",
+        "from pyppeteer import launch\n",
+        "import base64\n",
+        "\n",
+        "# Apply nest_asyncio to enable nested event loops\n",
+        "nest_asyncio.apply()\n",
+        "\n",
+        "async def fetch_blob_content(page, blob_url):\n",
+        "  blob_to_base64 = \"\"\"\n",
+        "  async (blobUrl) => {\n",
+        "      const blob = await fetch(blobUrl).then(r => r.blob());\n",
+        "      return new Promise((resolve) => {\n",
+        "          const reader = new FileReader();\n",
+        "          reader.onloadend = () => resolve(reader.result);\n",
+        "          reader.readAsDataURL(blob);\n",
+        "      });\n",
+        "  }\n",
+        "  \"\"\"\n",
+        "  base64_data = await page.evaluate(blob_to_base64, blob_url)\n",
+        "  _, encoded = base64_data.split(',', 1)\n",
+        "  return base64.b64decode(encoded)\n",
+        "\n",
+        "async def extract_pdb_file_download_link_and_content(url):\n",
+        "  browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])\n",
+        "  page = await browser.newPage()\n",
+        "  await page.goto(url, {'waitUntil': 'networkidle0'})\n",
+        "  elements = await page.querySelectorAll('a.btn.bg-purple')\n",
+        "  for element in elements:\n",
+        "      href = await page.evaluate('(element) => element.getAttribute(\"href\")', element)\n",
+        "      if 'blob:https://esmatlas.com/' in href:\n",
+        "          content = await fetch_blob_content(page, href)\n",
+        "          await browser.close()\n",
+        "          return href, content\n",
+        "  await browser.close()\n",
+        "  return \"No PDB file link found.\", None\n",
+        "\n",
+        "def esmfold_api(sequence):\n",
+        "  url = f'https://esmatlas.com/resources/fold/result?fasta_header=%3Eunnamed&sequence={sequence}'\n",
+        "  result = asyncio.get_event_loop().run_until_complete(extract_pdb_file_download_link_and_content(url))\n",
+        "  if result[1]:\n",
+        "      pdb_str = result[1].decode('utf-8')\n",
+        "      return pdb_str\n",
+        "  else:\n",
+        "      return \"Failed to retrieve PDB content.\"\n",
+        "\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "\n",
+        "if not os.path.exists('/content/in/'):\n",
+        "    os.mkdir('/content/in/')\n",
+        "\n",
+        "import py3Dmol\n",
+        "\n",
+        "def visualize_pdb_overlay(pdb1_path, pdb2_path):\n",
+        "    viewer = py3Dmol.view(width=800, height=600)\n",
+        "\n",
+        "    with open(pdb1_path, 'r') as f:\n",
+        "        pdb1_data = f.read()\n",
+        "    viewer.addModel(pdb1_data, 'pdb')\n",
+        "    viewer.setStyle({'model': 0}, {'cartoon': {'color': 'grey'}})\n",
+        "\n",
+        "    with open(pdb2_path, 'r') as f:\n",
+        "        pdb2_data = f.read()\n",
+        "    viewer.addModel(pdb2_data, 'pdb')\n",
+        "    viewer.setStyle({'model': 1}, {'cartoon': {'color' : 'red'}})\n",
+        "\n",
+        "    viewer.zoomTo()\n",
+        "    viewer.show()\n",
+        "\n",
+        "\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "bxL_tvmGqg-S"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Chroma Design\n",
+        "\n",
+        "On the start of this pipeline is the creation of draft backbones using Chroma or any other design method you want to use. We suggest you check out this two notebooks on Chroma to generate your desired proteins:\n",
+        "**Chroma Quickstart**\n",
+        "https://colab.research.google.com/github/generatebio/chroma/blob/main/notebooks/ChromaDemo.ipynb\n",
+        "\n",
+        "**Chroma API Tutorial**\n",
+        "https://colab.research.google.com/github/generatebio/chroma/blob/main/notebooks/ChromaAPI.ipynb\n",
+        "\n",
+        "If you have your PDB files please upload them into the **in/** folder and proceed\n"
+      ],
+      "metadata": {
+        "id": "WMpjGEyGqgSH"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Monomer af2cycling\n",
+        "#@markdown The af2cycler takes in the Chroma design and returns a new pdb file with improved structure\n",
+        "iterations = 10 #@param {type:\"integer\"}\n",
+        "#@markdown The af2cycled model is shown in red, while the chroma model is shown in grey\n",
+        "\n",
+        "\n",
+        "\n",
+        "import numpy as np\n",
+        "import warnings\n",
+        "\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os, re\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "\n",
+        "import os, re\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "\n",
+        "import numpy as np\n",
+        "\n",
+        "def sample_gumbel(shape, eps=1e-20):\n",
+        "  \"\"\"Sample from Gumbel(0, 1)\"\"\"\n",
+        "  U = np.random.uniform(size=shape)\n",
+        "  return -np.log(-np.log(U + eps) + eps)\n",
+        "\n",
+        "\n",
+        "clear_mem()\n",
+        "\n",
+        "iters = iterations\n",
+        "in_path = '/content/in/'\n",
+        "out_path = 'out/'\n",
+        "if not os.path.exists('/content/in/out/'):\n",
+        "    os.mkdir('/content/in/out/')\n",
+        "\n",
+        "\n",
+        "starting_seq = \"\"\n",
+        "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n",
+        "\n",
+        "\n",
+        "file_list = os.listdir(in_path)\n",
+        "\n",
+        "clear_mem()\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "af_model = mk_afdesign_model(protocol=\"fixbb\",use_templates=True,use_initial_atom_pos=True,use_initial_guess=True)\n",
+        "\n",
+        "for file_name in file_list:\n",
+        "    if file_name[-1] == 'b':\n",
+        "\n",
+        "        in_pdb = in_path + file_name\n",
+        "        out_pdb = in_path + out_path + 'Out_' + file_name\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "        af_model.prep_inputs(pdb_filename=in_pdb, chain='A')\n",
+        "\n",
+        "        starting_seq = af_model._inputs['batch'][\"aatype\"]\n",
+        "\n",
+        "        iterations = iters\n",
+        "\n",
+        "        use_dropout = True\n",
+        "        num_recycles = 0\n",
+        "\n",
+        "        mpnn_mode = \"conditional\"\n",
+        "        cmap_seqsep = 9\n",
+        "\n",
+        "        cmap_num = 2\n",
+        "        #cmap_cutoff = 14\n",
+        "        L = sum(af_model._lengths)\n",
+        "        af_model.restart(mode=\"gumbel\")\n",
+        "        af_model._args[\"clear_prev\"] = False\n",
+        "        #af_model.set_opt(cmap_cutoff=cmap_cutoff)\n",
+        "        af_model.set_weights(helix=1e-8)\n",
+        "        # gather info about inputs\n",
+        "        if \"offset\" in af_model._inputs:\n",
+        "            offset = af_model._inputs\n",
+        "        else:\n",
+        "            idx = af_model._inputs[\"residue_index\"]\n",
+        "            offset = idx[:, None] - idx[None, :]\n",
+        "        # initialize sequence\n",
+        "        if len(starting_seq) > 1:\n",
+        "            af_model.set_seq(seq=starting_seq)\n",
+        "        # initialize coordinates\n",
+        "        af_model._inputs.pop(\"prev\", None)\n",
+        "        init = af_model._inputs[\"batch\"]['all_atom_positions'].copy()\n",
+        "\n",
+        "        save_best = False\n",
+        "        for k in range(iterations):\n",
+        "\n",
+        "            if k > (iterations - 10):\n",
+        "                use_dropout = False\n",
+        "                save_best = True\n",
+        "\n",
+        "            # denoise\n",
+        "            aux = af_model.predict(return_aux=True, verbose=False,\n",
+        "                                   dropout=use_dropout,\n",
+        "                                   num_recycles=num_recycles)\n",
+        "\n",
+        "            #af_model._inputs[\"prev\"] = aux[\"prev\"]\n",
+        "            #af_model._inputs[\"prev\"][\"prev_msa_first_row\"] *= 0\n",
+        "            #af_model._inputs[\"prev\"][\"prev_pos\"] *= 0\n",
+        "\n",
+        "            cmap = aux[\"cmap\"] * (np.abs(offset) > cmap_seqsep)\n",
+        "            conf = np.sort(cmap)[:, -cmap_num:].mean(-1)\n",
+        "\n",
+        "            plddt = aux[\"plddt\"]\n",
+        "            seq = aux[\"seq\"][\"hard\"][0].argmax(-1)\n",
+        "            xyz = aux[\"atom_positions\"].copy()\n",
+        "            # update inputs\n",
+        "            af_model._inputs[\"batch\"][\"aatype\"] = seq\n",
+        "            af_model._inputs[\"batch\"][\"all_atom_positions\"] = xyz\n",
+        "\n",
+        "            if mpnn_mode != \"none\":\n",
+        "\n",
+        "                mpnn_model.get_af_inputs(af_model)\n",
+        "                opt = {\"mask\": np.sqrt(conf)}\n",
+        "                if mpnn_mode == \"unconditional\":\n",
+        "                    opt[\"ar_mask\"] = np.zeros((L, L))\n",
+        "                mpnn_out = mpnn_model.score(**opt)\n",
+        "                mpnn_logits = mpnn_out[\"logits\"][:, :20]\n",
+        "                aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"]\n",
+        "\n",
+        "                c = conf[:, None]\n",
+        "\n",
+        "                new_logits = (1 - c) * sample_gumbel(mpnn_logits.shape) + c * mpnn_logits\n",
+        "\n",
+        "                af_model._params[\"seq\"] = 0.9 * af_model._params[\"seq\"] + 0.1 * new_logits\n",
+        "\n",
+        "            # save results\n",
+        "            af_model._save_results(aux, save_best=save_best)\n",
+        "            af_model._k += 1\n",
+        "\n",
+        "        af_model.save_pdb(out_pdb)\n",
+        "        visualize_pdb_overlay(in_pdb, out_pdb)\n",
+        "\n"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "KeWmENFFqf1S"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "HTML(af_model.animate())"
+      ],
+      "metadata": {
+        "id": "HtpnTpvZtHi6"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown #Redesign with solubleMPNN for ESMFold prediction\n",
+        "#@markdown The standard manuscript settings were 8 sequences, 0.1 sampling temperature and the removal of cysteines\n",
+        "import pickle\n",
+        "import pandas as pd\n",
+        "in_path = '/content/in/out/'\n",
+        "out_path = 'out_sMPNN/'\n",
+        "if not os.path.exists(out_path):\n",
+        "    os.mkdir(out_path)\n",
+        "\n",
+        "file_list = os.listdir(in_path)\n",
+        "\n",
+        "\n",
+        "num_seqs = 8 #@param [\"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n",
+        "mpnn_sampling_temp = 0.1 #@param [\"0.0001\", \"0.1\", \"0.15\", \"0.2\", \"0.25\", \"0.3\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
+        "rm_aa = \"C\" #@param {type:\"string\"}\n",
+        "use_solubleMPNN = True #@param {type:\"boolean\"}\n",
+        "#@markdown - `mpnn_sampling_temp` - control diversity of sampled sequences. (higher = more diverse).\n",
+        "#@markdown - `rm_aa='C'` - do not use [C]ysteines.\n",
+        "#@markdown - `use_solubleMPNN` - use weights trained only on soluble proteins.\n",
+        "#@markdown\n",
+        "\n",
+        "from colabdesign.shared.protein import alphabet_list as chain_list\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "\n",
+        "\n",
+        "for file in file_list:\n",
+        "  if file[-4:] == '.pdb':\n",
+        "\n",
+        "    in_file1 = in_path + file\n",
+        "    mpnn_model.prep_inputs(pdb_filename=in_file1,\n",
+        "                          chain='A',\n",
+        "                          rm_aa=rm_aa,weights = \"soluble\")\n",
+        "    out = mpnn_model.sample(num=num_seqs//8,\n",
+        "                            batch=8,\n",
+        "                            temperature=mpnn_sampling_temp)\n",
+        "    for seq,score in zip(out[\"seq\"],out[\"score\"]):\n",
+        "      print(score,seq.split(\"/\")[0])\n",
+        "    df = pd.DataFrame(out[\"seq\"])\n",
+        "\n",
+        "    # Define the output path for saving the sequences as a .pkl file\n",
+        "    output_pkl_file = out_path + file[:-4] + \"_sequences.pkl\"\n",
+        "\n",
+        "    # Save the DataFrame to a .pkl file\n",
+        "    with open(output_pkl_file, 'wb') as f:\n",
+        "        pickle.dump(df, f)"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "KNcTvB5XzG-X"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown #Run ESMFold to test designability\n",
+        "\n",
+        "#@markdown **This cell is a little bit tricky. The problem is the compatibility between the JAX and PYTORCH frameworks between ColabDesign and ESMFold and GPU memory requirements**\n",
+        "\n",
+        "#@markdown Go to *Runtime* >> *Restart Session*\n",
+        "\n",
+        "#@markdown then Run this cell\n",
+        "\n",
+        "#@markdown Runtime 2-5 min\n",
+        "\n",
+        "#@markdown This cells runs ESMFold from huggingface and automatically calculates the RMSD to the designed backbone\n",
+        "#@markdown NOTE: GPU memory can be a big problem here. If you get memory errors please restart the runtime and run this cell again. It should be self contained. Additionally, after finish the ESMFold prediction rerun the setup cell\n",
+        "\n",
+        "import os\n",
+        "import pandas as pd\n",
+        "from Bio.PDB import PDBParser, Superimposer\n",
+        "import pickle\n",
+        "import torch\n",
+        "import numpy as np\n",
+        "from transformers import AutoTokenizer, EsmForProteinFolding\n",
+        "from transformers.models.esm.openfold_utils.protein import to_pdb, Protein as OFProtein\n",
+        "from transformers.models.esm.openfold_utils.feats import atom14_to_atom37\n",
+        "\n",
+        "import py3Dmol\n",
+        "\n",
+        "def visualize_pdb_overlay(pdb1_path, pdb2_path):\n",
+        "    viewer = py3Dmol.view(width=800, height=600)\n",
+        "\n",
+        "    with open(pdb1_path, 'r') as f:\n",
+        "        pdb1_data = f.read()\n",
+        "    viewer.addModel(pdb1_data, 'pdb')\n",
+        "    viewer.setStyle({'model': 0}, {'cartoon': {'color': 'grey'}})\n",
+        "\n",
+        "    with open(pdb2_path, 'r') as f:\n",
+        "        pdb2_data = f.read()\n",
+        "    viewer.addModel(pdb2_data, 'pdb')\n",
+        "    viewer.setStyle({'model': 1}, {'cartoon': {'color' : 'red'}})\n",
+        "\n",
+        "    viewer.zoomTo()\n",
+        "    viewer.show()\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "tokenizer = AutoTokenizer.from_pretrained(\"facebook/esmfold_v1\")\n",
+        "model = EsmForProteinFolding.from_pretrained(\"facebook/esmfold_v1\", low_cpu_mem_usage=True)\n",
+        "\n",
+        "device = 'cuda:0'\n",
+        "model = model.cuda(device)\n",
+        "model.esm = model.esm.half()\n",
+        "model.trunk.set_chunk_size(64)\n",
+        "torch.backends.cuda.matmul.allow_tf32 = True\n",
+        "\n",
+        "def convert_outputs_to_pdb(outputs):\n",
+        "    final_atom_positions = atom14_to_atom37(outputs[\"positions\"][-1], outputs)\n",
+        "    outputs = {k: v.to(\"cpu\").numpy() for k, v in outputs.items()}\n",
+        "    final_atom_positions = final_atom_positions.cpu().numpy()\n",
+        "    final_atom_mask = outputs[\"atom37_atom_exists\"]\n",
+        "    pdbs = []\n",
+        "    for i in range(outputs[\"aatype\"].shape[0]):\n",
+        "        aa = outputs[\"aatype\"][i]\n",
+        "        pred_pos = final_atom_positions[i]\n",
+        "        mask = final_atom_mask[i]\n",
+        "        resid = outputs[\"residue_index\"][i] + 1\n",
+        "        pred = OFProtein(\n",
+        "            aatype=aa,\n",
+        "            atom_positions=pred_pos,\n",
+        "            atom_mask=mask,\n",
+        "            residue_index=resid,\n",
+        "            b_factors=outputs[\"plddt\"][i],\n",
+        "            chain_index=outputs[\"chain_index\"][i] if \"chain_index\" in outputs else None,\n",
+        "        )\n",
+        "        pdbs.append(to_pdb(pred))\n",
+        "    return pdbs\n",
+        "\n",
+        "def calculate_ca_rmsd(pdb_file1, pdb_file2):\n",
+        "    parser = PDBParser(QUIET=True)\n",
+        "\n",
+        "    structure1 = parser.get_structure(\"Protein1\", pdb_file1)\n",
+        "    structure2 = parser.get_structure(\"Protein2\", pdb_file2)\n",
+        "\n",
+        "    ca_atoms1 = [atom for atom in structure1.get_atoms() if atom.get_name() == \"CA\"]\n",
+        "    ca_atoms2 = [atom for atom in structure2.get_atoms() if atom.get_name() == \"CA\"]\n",
+        "\n",
+        "    super_imposer = Superimposer()\n",
+        "    super_imposer.set_atoms(ca_atoms1, ca_atoms2)\n",
+        "    super_imposer.apply(structure2.get_atoms())\n",
+        "    rmsd = super_imposer.rms\n",
+        "    return rmsd\n",
+        "\n",
+        "def process_sequences(seq_list, pdb_file,in_path_pdb):\n",
+        "    lowest_rmsd = float('inf')\n",
+        "    lowest_rmsd_data = None\n",
+        "    pdb_id = pdb_file[16:-4]\n",
+        "    out_ss_path = in_path_pdb + \"output/\"\n",
+        "\n",
+        "    if not os.path.exists(out_ss_path):\n",
+        "        os.mkdir(out_ss_path)\n",
+        "\n",
+        "    for test_protein in seq_list:\n",
+        "        data = {}\n",
+        "        tokenized_input = tokenizer([test_protein], return_tensors=\"pt\", add_special_tokens=False)['input_ids']\n",
+        "        tokenized_input = tokenized_input.cuda(device)\n",
+        "\n",
+        "        with torch.no_grad():\n",
+        "            output = model(tokenized_input)\n",
+        "\n",
+        "        data['out'] = output\n",
+        "        data[\"plddt\"] = torch.mean(output['plddt']).item()\n",
+        "        data['pae'] = torch.mean(output['predicted_aligned_error']).item()\n",
+        "\n",
+        "        pdb_data = convert_outputs_to_pdb(output)\n",
+        "        tmp_pdb_file = os.path.join(out_ss_path, \"TMP.pdb\")\n",
+        "\n",
+        "        with open(tmp_pdb_file, 'w') as file:\n",
+        "            for line in pdb_data:\n",
+        "                file.write(line)\n",
+        "\n",
+        "        data['rmsd'] = calculate_ca_rmsd(tmp_pdb_file, pdb_file)\n",
+        "        print(f'Sequence: {test_protein}, plddt: {data[\"plddt\"]}, PAE: {data[\"pae\"]}, RMSD: {data[\"rmsd\"]}')\n",
+        "\n",
+        "        if data['rmsd'] < lowest_rmsd:\n",
+        "            lowest_rmsd = data['rmsd']\n",
+        "            lowest_rmsd_data = data\n",
+        "\n",
+        "    if lowest_rmsd_data is not None:\n",
+        "        print(f'Lowest RMSD: {lowest_rmsd}')\n",
+        "        best_pdb_data = convert_outputs_to_pdb(lowest_rmsd_data['out'])\n",
+        "        best_pdb_file = os.path.join(out_ss_path, f\"{pdb_id}_best_structure.pdb\")\n",
+        "\n",
+        "        with open(best_pdb_file, 'w') as file:\n",
+        "            for line in best_pdb_data:\n",
+        "                file.write(line)\n",
+        "\n",
+        "        original_dict = lowest_rmsd_data\n",
+        "        key_to_exclude = 'out'\n",
+        "        data_out = {k: v for k, v in original_dict.items() if k != key_to_exclude}\n",
+        "\n",
+        "        with open(os.path.join(out_ss_path, f\"{pdb_id}_best_structure_data.pkl\"), 'wb') as f:\n",
+        "            pickle.dump(data_out, f)\n",
+        "\n",
+        "        return lowest_rmsd, best_pdb_file, data_out\n",
+        "\n",
+        "    return None, None, None\n",
+        "\n",
+        "\n",
+        "in_path= '/content/out_sMPNN/'\n",
+        "file_list = os.listdir(in_path)\n",
+        "\n",
+        "for file in file_list:\n",
+        "  if file[-1] =='l':\n",
+        "\n",
+        "\n",
+        "\n",
+        "    output_pkl_file = in_path + file\n",
+        "\n",
+        "    with open(output_pkl_file, 'rb') as f:\n",
+        "        seq = pickle.load(f)\n",
+        "    seq_list = []\n",
+        "    for i in np.asarray(seq):\n",
+        "      seq_list.append(i[0])\n",
+        "\n",
+        "    pdb_file = \"/content/in/out/\" + file[:-14] + '.pdb'\n",
+        "    print(seq_list)\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "    lowest_rmsd, best_pdb_file, best_data = process_sequences(seq_list, pdb_file,in_path)\n",
+        "    if lowest_rmsd is not None:\n",
+        "        print(f\"Lowest RMSD: {lowest_rmsd}, Best PDB file: {best_pdb_file}\")\n",
+        "    else:\n",
+        "        print(\"No valid result found.\")\n",
+        "    #visualize_pdb_overlay(pdb_file, best_pdb_file)\n"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "collapsed": true,
+        "id": "Ctyijhvu0UVe"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown #Optional: Run AF2 based designability test\n",
+        "\n",
+        "\n",
+        "#@markdown Rerun **setup cell** if you tested ESMFold prediction before!\n",
+        "\n",
+        "#@markdown This cell predicts the solubleMPNN generated seuqneces with AF2 with Initial Guess & All atom initialisation\n",
+        "import re\n",
+        "clear_mem()\n",
+        "\n",
+        "in_path = '/content/in/'\n",
+        "out_path = 'out/'\n",
+        "if not os.path.exists('/content/in/out/'):\n",
+        "    os.mkdir('/content/in/out/')\n",
+        "\n",
+        "\n",
+        "starting_seq = \"\"\n",
+        "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n",
+        "\n",
+        "\n",
+        "file_list = os.listdir(in_path)\n",
+        "\n",
+        "clear_mem()\n",
+        "\n",
+        "af_model = mk_afdesign_model(protocol=\"fixbb\",use_initial_atom_pos=True,use_initial_guess=True)\n",
+        "\n",
+        "\n",
+        "def process_sequences(seq_list, pdb_file,in_path_pdb):\n",
+        "    lowest_rmsd = float('inf')\n",
+        "    lowest_rmsd_data = None\n",
+        "    pdb_id = pdb_file[16:-4]\n",
+        "    out_ss_path = in_path_pdb + \"output_AF2/\"\n",
+        "\n",
+        "    if not os.path.exists(out_ss_path):\n",
+        "        os.mkdir(out_ss_path)\n",
+        "    kk=0\n",
+        "    for test_protein in seq_list:\n",
+        "        data = {}\n",
+        "        af_model.prep_inputs(pdb_filename=pdb_file, chain='A')\n",
+        "        af_model.predict(seq=test_protein,num_recycles=3)\n",
+        "\n",
+        "\n",
+        "\n",
+        "        data[\"plddt\"] = af_model.aux['losses']['plddt']\n",
+        "        data['pae'] = af_model.aux['losses']['pae']*31\n",
+        "        data['rmsd'] = af_model.aux['losses']['rmsd']\n",
+        "        af_model.save_pdb(f'{out_ss_path}{pdb_id}_{kk}.pdb')\n",
+        "        print(f'Sequence: {test_protein}, plddt: {data[\"plddt\"]}, PAE: {data[\"pae\"]}, RMSD: {data[\"rmsd\"]}')\n",
+        "\n",
+        "        if data['rmsd'] < lowest_rmsd:\n",
+        "            lowest_rmsd = data['rmsd']\n",
+        "            lowest_rmsd_data = data\n",
+        "\n",
+        "    if lowest_rmsd_data is not None:\n",
+        "        print(f'Lowest RMSD: {lowest_rmsd}')\n",
+        "\n",
+        "        best_pdb_file = os.path.join(out_ss_path, f\"{pdb_id}_best_structure.pdb\")\n",
+        "\n",
+        "        return lowest_rmsd, best_pdb_file\n",
+        "\n",
+        "    return None, None, None\n",
+        "\n",
+        "\n",
+        "in_path= '/content/out_sMPNN/'\n",
+        "file_list = os.listdir(in_path)\n",
+        "\n",
+        "for file in file_list:\n",
+        "  if file[-1] =='l':\n",
+        "\n",
+        "\n",
+        "\n",
+        "    output_pkl_file = in_path + file\n",
+        "\n",
+        "    with open(output_pkl_file, 'rb') as f:\n",
+        "        seq = pickle.load(f)\n",
+        "    seq_list = []\n",
+        "    for i in np.asarray(seq):\n",
+        "      seq_list.append(i[0])\n",
+        "\n",
+        "    pdb_file = \"/content/in/out/\" + file[:-14] + '.pdb'\n",
+        "    print(seq_list)\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "\n",
+        "    lowest_rmsd, best_pdb_file = process_sequences(seq_list, pdb_file,in_path)\n",
+        "    if lowest_rmsd is not None:\n",
+        "        print(f\"Lowest RMSD: {lowest_rmsd}, Best PDB file: {best_pdb_file}\")\n",
+        "    else:\n",
+        "        print(\"No valid result found.\")\n",
+        "    #visualize_pdb_overlay(pdb_file, best_pdb_file)\n",
+        "\n",
+        "\n"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "s1EVl7Ee8ne-"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "ZQljfnQl8oTc"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "16XTY0u9YyxP"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "u6mCraHrdqM3"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ]
+}

model/PXDesignBench/ColabDesign/af/examples/af_cyc_design.ipynb

@@ -0,0 +1,320 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/af_cyc_design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "#af_cyc_design\n",
+        "\n",
+        "**Cyclic peptide structure prediction and design using AlphaFold**\n",
+        "\n",
+        "Stephen Rettie, Katelyn Campbell, Asim Bera, Alex Kang, Simon Kozlov, Joshmyn De La Cruz, Victor Adebomi, Guangfeng Zhou, Frank DiMaio, Sergey Ovchinnikov, Gaurav Bhardwaj\n",
+        "\n",
+        "doi: https://doi.org/10.1101/2023.02.25.529956\n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup (~2 minutes)\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "def add_cyclic_offset(self, offset_type=2):\n",
+        "  '''add cyclic offset to connect N and C term'''\n",
+        "  def cyclic_offset(L):\n",
+        "    i = np.arange(L)\n",
+        "    ij = np.stack([i,i+L],-1)\n",
+        "    offset = i[:,None] - i[None,:]\n",
+        "    c_offset = np.abs(ij[:,None,:,None] - ij[None,:,None,:]).min((2,3))\n",
+        "    if offset_type == 1:\n",
+        "      c_offset = c_offset\n",
+        "    elif offset_type >= 2:\n",
+        "      a = c_offset < np.abs(offset)\n",
+        "      c_offset[a] = -c_offset[a]\n",
+        "    if offset_type == 3:\n",
+        "      idx = np.abs(c_offset) > 2\n",
+        "      c_offset[idx] = (32 * c_offset[idx] )/  abs(c_offset[idx])\n",
+        "    return c_offset * np.sign(offset)\n",
+        "  idx = self._inputs[\"residue_index\"]\n",
+        "  offset = np.array(idx[:,None] - idx[None,:])\n",
+        "\n",
+        "  if self.protocol == \"binder\":\n",
+        "    c_offset = cyclic_offset(self._binder_len)\n",
+        "    offset[self._target_len:,self._target_len:] = c_offset\n",
+        "\n",
+        "  if self.protocol in [\"fixbb\",\"partial\",\"hallucination\"]:\n",
+        "    Ln = 0\n",
+        "    for L in self._lengths:\n",
+        "      offset[Ln:Ln+L,Ln:Ln+L] = cyclic_offset(L)\n",
+        "      Ln += L\n",
+        "  self._inputs[\"offset\"] = offset\n",
+        "\n",
+        "def add_rg_loss(self, weight=0.1):\n",
+        "  '''add radius of gyration loss'''\n",
+        "  def loss_fn(inputs, outputs):\n",
+        "    xyz = outputs[\"structure_module\"]\n",
+        "    ca = xyz[\"final_atom_positions\"][:,residue_constants.atom_order[\"CA\"]]\n",
+        "    rg = jnp.sqrt(jnp.square(ca - ca.mean(0)).sum(-1).mean() + 1e-8)\n",
+        "    rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "    rg = jax.nn.elu(rg - rg_th)\n",
+        "    return {\"rg\":rg}\n",
+        "  self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n",
+        "  self.opt[\"weights\"][\"rg\"] = weight"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UUfKrOzT0gOS"
+      },
+      "source": [
+        "# fixed backbone design (fixbb) (~2 minutes)\n",
+        "For a given protein backbone, generate/design a new sequence that AlphaFold thinks folds into that conformation."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qLd1DsnKzxBJ"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"fixbb\")\n",
+        "af_model.prep_inputs(pdb_filename=get_pdb(\"7m28\"), chain=\"A\")\n",
+        "add_cyclic_offset(af_model, offset_type=2)\n",
+        "\n",
+        "print(\"length\",  af_model._len)\n",
+        "print(\"weights\", af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "plt.imshow(af_model._inputs[\"offset\"],cmap=\"bwr\")"
+      ],
+      "metadata": {
+        "id": "xr6mw-CXl_o3"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.restart()\n",
+        "af_model.design_3stage()"
+      ],
+      "metadata": {
+        "id": "u0AwskJ84NGx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.plot_traj()"
+      ],
+      "metadata": {
+        "id": "8FB1v7dn1LL6"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YEApO8YzBoS0"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb(show_mainchains=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cW1KQiHKJpfp"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "YDrChASGVUUx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# hallucination (~1 minute)\n",
+        "For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts)."
+      ],
+      "metadata": {
+        "id": "qLwS2s_xcjRI"
+      }
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sZnYfCbfEvol"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "af_model.prep_inputs(length=13, rm_aa=\"C\")\n",
+        "add_cyclic_offset(af_model)\n",
+        "# add_rg_loss(af_model)\n",
+        "\n",
+        "print(\"length\",af_model._len)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# pre-design with gumbel initialization and softmax activation\n",
+        "af_model.restart()\n",
+        "af_model.set_seq(mode=\"gumbel\")\n",
+        "af_model.set_opt(\"con\", binary=True, cutoff=21.6875, num=af_model._len, seqsep=0)\n",
+        "af_model.set_weights(pae=1, plddt=1, con=0.5)\n",
+        "af_model.design_soft(50)\n",
+        "\n",
+        "# three stage design\n",
+        "af_model.set_seq(seq=af_model.aux[\"seq\"][\"pseudo\"])\n",
+        "af_model.design_3stage(50,50,10)"
+      ],
+      "metadata": {
+        "id": "f76xqCkw0vj9"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "A1GxeLZdTTya"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb(show_mainchains=True)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "L2E9Tn2Acchj"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YSKWYu0_GlUH"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.get_seqs()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/af_pseudo_diffusion.ipynb

@@ -0,0 +1,266 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_pseudo_diffusion.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VD9K5H2cnFxL"
+      },
+      "source": [
+        "#AF_pseudo_diffusion + proteinMPNN\n",
+        "Hacking AlphaFold to be a diffusion model (for backbone generation). At each step add logits from proteinMPNN.\n",
+        "\n",
+        "\n",
+        "**WARNING**: This notebook is experimental, designed as a control. Not intended for practical use at this stage.\n",
+        "\n",
+        "---\n",
+        "**NEW**\n",
+        "For latest version of notebook (reconfigured to work in distogram space) go here:\n",
+        "[af_pseudo_diffusion_dgram.ipynb](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_pseudo_diffusion_dgram.ipynb)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "YCRzPGdTZfEe"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os, re\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "def sample_gumbel(shape, eps=1e-20): \n",
+        "  \"\"\"Sample from Gumbel(0, 1)\"\"\"\n",
+        "  U = np.random.uniform(size=shape)\n",
+        "  return -np.log(-np.log(U + eps) + eps)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wxYMB4A9Zmrf",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title initialize the model\n",
+        "length = 100 #@param {type:\"integer\"}\n",
+        "#@markdown Provide a starting point (optional)\n",
+        "starting_seq = \"\" #@param {type:\"string\"}\n",
+        "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n",
+        "#@markdown - if `starting_seq` provided the `length` option will be overwritten by length of starting sequence.\n",
+        "use_starting_pdb = False #@param {type:\"boolean\"}\n",
+        "pdb = \"\" #@param {type:\"string\"}\n",
+        "#@markdown - specify PDB or Uniprot code or leave pdb blank for upload prompt\n",
+        "chains = \"A\" #@param {type:\"string\"}\n",
+        "#@markdown - (example `A` or `A,B` for complexes)\n",
+        "fix_pos = \"\" #@param {type:\"string\"}\n",
+        "#@markdown - specify which positions to keep fixed in the sequence (example: `1,2-10`)\n",
+        "#@markdown - you can also specify chain specific constraints (example: `A1-10,B1-20`)\n",
+        "#@markdown - you can also specify to fix entire chain(s) (example: `A`)\n",
+        "\n",
+        "if len(starting_seq) > 0:\n",
+        "  length = len(starting_seq)\n",
+        "\n",
+        "clear_mem()\n",
+        "if use_starting_pdb:\n",
+        "  af_model = mk_afdesign_model(protocol=\"fixbb\", use_templates=True)\n",
+        "  af_model.prep_inputs(get_pdb(pdb), chains, fix_pos=fix_pos) \n",
+        "else:\n",
+        "  af_model = mk_afdesign_model(protocol=\"hallucination\", use_templates=True)\n",
+        "  af_model.prep_inputs(length=length)\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "print(\"lengths\",af_model._lengths)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3Dt8i00UbxtW",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title run protocol\n",
+        "#@markdown Optimization options\n",
+        "iterations = 100 #@param [\"50\", \"100\"] {type:\"raw\"}\n",
+        "use_xyz_noise = True #@param {type:\"boolean\"}\n",
+        "use_seq_noise = True #@param {type:\"boolean\"}\n",
+        "use_dropout = True #@param {type:\"boolean\"}\n",
+        "use_plddt = True #@param {type:\"boolean\"}\n",
+        "store_denoised = True #@param {type:\"boolean\"}\n",
+        "#@markdown - this does not change results, but toggle between storing noised vs denoised coordinates for animation.\n",
+        "\n",
+        "#@markdown AlphaFold options\n",
+        "sample_models = False #@param {type:\"boolean\"}\n",
+        "rm_template_seq = True #@param {type:\"boolean\"}\n",
+        "\n",
+        "#@markdown proteinMPNN options (disable to keep sequence the same)\n",
+        "use_mpnn = True #@param {type:\"boolean\"}\n",
+        "mpnn_mode = \"conditional\" #@param [\"conditional\", \"unconditional\"]\n",
+        "\n",
+        "af_model.restart(mode=\"gumbel\")\n",
+        "\n",
+        "if len(starting_seq) > 1:\n",
+        "  af_model.set_seq(seq=starting_seq)\n",
+        "elif use_starting_pdb:\n",
+        "  af_model.set_seq(mode=\"wildtype\")\n",
+        "\n",
+        "af_model._inputs[\"rm_template_seq\"] = rm_template_seq\n",
+        "L = sum(af_model._lengths)\n",
+        "\n",
+        "af_model._inputs[\"bias\"] = np.zeros((L,20))\n",
+        "\n",
+        "if not use_starting_pdb:\n",
+        "  af_model._inputs[\"batch\"] = {\"aatype\":np.zeros(L).astype(int),\n",
+        "                               \"all_atom_mask\":np.zeros((L,37)),\n",
+        "                               \"all_atom_positions\":np.zeros((L,37,3))}\n",
+        "\n",
+        "for k in range(iterations):\n",
+        "  # add noise\n",
+        "  if use_seq_noise:\n",
+        "    af_model._inputs[\"bias\"] = 0.1 * sample_gumbel((L,20))\n",
+        "\n",
+        "  if use_xyz_noise:\n",
+        "    n = np.random.normal(size=(L,37,3)) * (1-k/iterations)\n",
+        "    af_model._inputs[\"batch\"][\"all_atom_positions\"] += n\n",
+        "    if not store_denoised and k > 0:\n",
+        "      af_model._tmp[\"traj\"][\"xyz\"][-1] += n[:,1]\n",
+        "\n",
+        "  # denoise\n",
+        "  aux = af_model.predict(return_aux=True, verbose=False,\n",
+        "                         sample_models=sample_models, dropout=use_dropout)\n",
+        "  plddt = af_model.aux[\"plddt\"]\n",
+        "  # update inputs\n",
+        "  af_model._inputs[\"batch\"][\"aatype\"] = af_model.aux[\"seq\"][\"hard\"].argmax(-1)[0]\n",
+        "  af_model._inputs[\"batch\"][\"all_atom_mask\"][:,:4] = 1\n",
+        "  if use_plddt:\n",
+        "    af_model._inputs[\"batch\"][\"all_atom_mask\"][:,(1,3)] = np.sqrt(plddt)[:,None]\n",
+        "  af_model._inputs[\"batch\"][\"all_atom_positions\"] = af_model.aux[\"atom_positions\"].copy()\n",
+        "\n",
+        "  # add logits from proteinmpnn at each stage\n",
+        "  if use_mpnn:\n",
+        "    mpnn_model.get_af_inputs(af_model)\n",
+        "    opt = {} if mpnn_mode == \"conditional\" else {\"ar_mask\":np.zeros((L,L))}\n",
+        "    mpnn_out = mpnn_model.score(**opt)\n",
+        "    aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"]\n",
+        "    mpnn_logits = mpnn_out[\"logits\"][:,:20]\n",
+        "    m = (k/iterations)\n",
+        "    af_model._params[\"seq\"] = (1-m) * af_model._params[\"seq\"] + m * mpnn_logits\n",
+        "\n",
+        "    # save results\n",
+        "  af_model._save_results(aux)\n",
+        "  af_model._k += 1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "v3dPQeEQcAcT"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.save_pdb(\"tmp.pdb\")\n",
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "6Ck-ME5_99qE"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6cJhizcYcUxm"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate(dpi=100))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "aVS6VC7wKyhD"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/af_pseudo_diffusion_dgram.ipynb

@@ -0,0 +1,376 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_pseudo_diffusion_dgram.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VD9K5H2cnFxL"
+      },
+      "source": [
+        "#AF_pseudo_diffusion + proteinMPNN\n",
+        "Hacking AlphaFold to be a diffusion model (for backbone generation) via distogram. At each step add logits from proteinMPNN.\n",
+        "\n",
+        "\n",
+        "**WARNING**: This notebook is experimental, designed as a control. Not intended for practical use at this stage."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "YCRzPGdTZfEe"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os, re\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "from colabdesign.shared.protein import _np_get_cb\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "import jax.numpy as jnp\n",
+        "import jax\n",
+        "from scipy.special import softmax\n",
+        "\n",
+        "import tqdm.notebook\n",
+        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "def sample_gumbel(shape, eps=1e-20): \n",
+        "  \"\"\"Sample from Gumbel(0, 1)\"\"\"\n",
+        "  U = np.random.uniform(size=shape)\n",
+        "  return -np.log(-np.log(U + eps) + eps)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "wxYMB4A9Zmrf",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title initialize the model\n",
+        "length = 100 #@param {type:\"integer\"}\n",
+        "#symmetry = \"C\" #@param [\"C\"]\n",
+        "#copies = 1 #@param {type:\"integer\"}\n",
+        "#@markdown Provide a starting point (optional)\n",
+        "starting_seq = \"\" #@param {type:\"string\"}\n",
+        "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n",
+        "#@markdown - if `starting_seq` provided the `length` option will be overwritten by length of starting sequence.\n",
+        "\n",
+        "#@markdown Experimental options\n",
+        "use_multimer = False \n",
+        "mode = \"dgram_retrain\" #@param [\"dgram\",\"dgram_retrain\"]\n",
+        "##@markdown - `xyz` - use structure output as template input\n",
+        "#@markdown - `dgram` - use distogram output as template input\n",
+        "#@markdown - `dgram_retrain` - replace distogram head from AlphaFold with one retrained to map output bins to template bins.\n",
+        "\n",
+        "if len(starting_seq) > 0:\n",
+        "  length = len(starting_seq)\n",
+        "\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\",\n",
+        "                             use_templates=True,\n",
+        "                             debug=True, \n",
+        "                             use_multimer=use_multimer)\n",
+        "af_model.prep_inputs(length=length)\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "print(\"lengths\",af_model._lengths)\n",
+        "\n",
+        "if \"dgram\" in mode:\n",
+        "  if \"retrain\" in mode and not use_multimer:\n",
+        "    # update distogram head to return all 39 bins\n",
+        "    af_model._cfg.model.heads.distogram.first_break = 3.25\n",
+        "    af_model._cfg.model.heads.distogram.last_break = 50.75\n",
+        "    af_model._cfg.model.heads.distogram.num_bins = 39\n",
+        "    af_model._model = af_model._get_model(af_model._cfg)\n",
+        "    from colabdesign.af.weights import __file__ as af_path\n",
+        "    template_dgram_head = np.load(os.path.join(os.path.dirname(af_path),'template_dgram_head.npy'))\n",
+        "    for k in range(len(af_model._model_params)):\n",
+        "      params = {\"weights\":jnp.array(template_dgram_head[k]),\"bias\":jnp.zeros(39)}\n",
+        "      af_model._model_params[k][\"alphafold/alphafold_iteration/distogram_head/half_logits\"] = params\n",
+        "  else:\n",
+        "    dgram_map = np.eye(39)[np.repeat(np.append(0,np.arange(15)),4)]\n",
+        "    dgram_map[-1,:] = 0 \n",
+        "\n",
+        "def get_dgram(positions, num_bins=39, min_bin=3.25, max_bin=50.75):\n",
+        "  atom_idx = residue_constants.atom_order\n",
+        "  atoms = {k:positions[...,atom_idx[k],:] for k in [\"N\",\"CA\",\"C\"]}\n",
+        "  cb = _np_get_cb(**atoms, use_jax=False)\n",
+        "  dist2 = np.square(cb[None,:] - cb[:,None]).sum(-1,keepdims=True)\n",
+        "  lower_breaks = np.linspace(min_bin, max_bin, num_bins)\n",
+        "  lower_breaks = np.square(lower_breaks)\n",
+        "  upper_breaks = np.concatenate([lower_breaks[1:],np.array([1e8], dtype=jnp.float32)], axis=-1)\n",
+        "  return ((dist2 > lower_breaks) * (dist2 < upper_breaks)).astype(float)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3Dt8i00UbxtW",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title run protocol\n",
+        "#@markdown Optimization options\n",
+        "iterations = 100 #@param [\"50\", \"100\"] {type:\"raw\"}\n",
+        "use_dgram_noise = True #@param {type:\"boolean\"}\n",
+        "use_seq_noise = True #@param {type:\"boolean\"}\n",
+        "use_dropout = True #@param {type:\"boolean\"}\n",
+        "seqsep_mask = 6 #@param {type:\"integer\"}\n",
+        "\n",
+        "#@markdown AlphaFold options\n",
+        "sample_models = True #@param {type:\"boolean\"}\n",
+        "num_recycles = 0 #@param [\"0\", \"1\", \"2\", \"3\"] {type:\"raw\"}\n",
+        "\n",
+        "#@markdown proteinMPNN options (set to `none` to disable)\n",
+        "mpnn_mode = \"conditional\" #@param [\"none\",\"sample\", \"conditional\", \"unconditional\"]\n",
+        "\n",
+        "L = sum(af_model._lengths)\n",
+        "af_model.restart(mode=\"gumbel\")\n",
+        "af_model._inputs[\"rm_template_seq\"] = True\n",
+        "\n",
+        "# gather info about inputs\n",
+        "if \"offset\" in af_model._inputs:\n",
+        "  offset = af_model._inputs\n",
+        "else:\n",
+        "  idx = af_model._inputs[\"residue_index\"]\n",
+        "  offset = idx[:,None] - idx[None,:]\n",
+        "\n",
+        "# initialize sequence\n",
+        "if len(starting_seq) > 1:\n",
+        "  af_model.set_seq(seq=starting_seq)\n",
+        "af_model._inputs[\"bias\"] = np.zeros((L,20))\n",
+        "\n",
+        "# initialize coordinates/dgram\n",
+        "af_model._inputs[\"batch\"] = {\"aatype\":np.zeros(L).astype(int),\n",
+        "                             \"all_atom_mask\":np.zeros((L,37)),\n",
+        "                             \"all_atom_positions\":np.zeros((L,37,3)),\n",
+        "                             \"dgram\":np.zeros((L,L,39))}\n",
+        "\n",
+        "for k in range(iterations):\n",
+        "\n",
+        "  # disable stochastic part for the last 10 steps\n",
+        "  if k > (iterations - 10):\n",
+        "    use_dropout = False\n",
+        "    sample_models = False\n",
+        "    use_seq_noise = False\n",
+        "    seqsep_mask = 0.0\n",
+        "\n",
+        "  # noise\n",
+        "  if k > 0:\n",
+        "    dgram_xyz = get_dgram(xyz)\n",
+        "    dgram_prob = softmax(dgram_logits,-1)\n",
+        "\n",
+        "    if use_seq_noise:\n",
+        "      af_model._inputs[\"bias\"] = 0.1 * sample_gumbel((L,20))\n",
+        "\n",
+        "    if mode == \"xyz\":\n",
+        "      dgram = dgram_xyz\n",
+        "    if mode == \"dgram\":\n",
+        "      dgram = dgram_prob @ dgram_map\n",
+        "      dgram[...,14:] = dgram_xyz[...,14:] * dgram_prob[...,-1:]\n",
+        "    if mode == \"dgram_retrain\":\n",
+        "      dgram = dgram_prob\n",
+        "    \n",
+        "    if use_dgram_noise:\n",
+        "      noise = sample_gumbel(dgram.shape) * (1 - k/iterations)\n",
+        "      dgram = softmax(np.log(dgram + 1e-8) + noise, -1)\n",
+        "\n",
+        "    # add mask to avoid local contacts being fixed (otherwise there is a bias toward helix)\n",
+        "    mask = np.abs(offset) > seqsep_mask\n",
+        "    af_model._inputs[\"batch\"][\"dgram\"] = dgram * mask[:,:,None]\n",
+        "\n",
+        "  # denoise\n",
+        "  aux = af_model.predict(return_aux=True, verbose=False,\n",
+        "                         sample_models=sample_models,\n",
+        "                         dropout=use_dropout, num_recycles=num_recycles)\n",
+        "  plddt = aux[\"plddt\"]\n",
+        "  seq = aux[\"seq\"][\"hard\"][0].argmax(-1)\n",
+        "  xyz = aux[\"atom_positions\"].copy()\n",
+        "  dgram_logits = aux[\"debug\"][\"outputs\"][\"distogram\"][\"logits\"] \n",
+        "  \n",
+        "  # update inputs    \n",
+        "  af_model._inputs[\"batch\"][\"aatype\"] = seq\n",
+        "  af_model._inputs[\"batch\"][\"all_atom_mask\"][:,:4] = np.sqrt(plddt)[:,None]\n",
+        "  af_model._inputs[\"batch\"][\"all_atom_positions\"] = xyz\n",
+        "\n",
+        "  # add logits from proteinmpnn at each stage\n",
+        "  if mpnn_mode != \"none\":\n",
+        "    mpnn_model.get_af_inputs(af_model)\n",
+        "    if mpnn_mode == \"sample\":\n",
+        "      mpnn_out = mpnn_model.sample(temp = 1-k/iterations)\n",
+        "      mpnn_logits = mpnn_out[\"logits\"][0,:,:20]\n",
+        "      aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"][0]\n",
+        "    else:\n",
+        "      opt = {} if mpnn_mode == \"conditional\" else {\"ar_mask\":np.zeros((L,L))}\n",
+        "      mpnn_out = mpnn_model.score(**opt)\n",
+        "      mpnn_logits = mpnn_out[\"logits\"][:,:20]\n",
+        "      aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"]\n",
+        "    \n",
+        "    beta = np.square(k/iterations) * plddt[:,None]\n",
+        "    af_model._params[\"seq\"] = (1-beta) * af_model._params[\"seq\"] + beta * mpnn_logits\n",
+        "\n",
+        "  # save results\n",
+        "  af_model._save_results(aux)\n",
+        "  af_model._k += 1"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.save_pdb(\"0.pdb\")\n",
+        "af_model.plot_pdb()\n",
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "YeYD4KF8MUA_"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6cJhizcYcUxm"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate(dpi=100))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title sample new sequences using proteinMPNN and rescore with alphafold (w/o template)\n",
+        "#@markdown #### Design Options\n",
+        "num_seqs = 32 #@param [\"32\", \"64\", \"128\", \"256\", \"512\", \"1024\"] {type:\"raw\"}\n",
+        "sampling_temp = 0.1 \n",
+        "import pandas as pd\n",
+        "from google.colab import data_table\n",
+        "data_table.enable_dataframe_formatter()\n",
+        "\n",
+        "# zero out template inputs\n",
+        "out = mpnn_model.sample(num=num_seqs//32, batch=32,\n",
+        "                        temperature=sampling_temp)\n",
+        "af_terms = [\"plddt\",\"ptm\",\"pae\"]\n",
+        "for k in af_terms: out[k] = []\n",
+        "os.system(\"mkdir -p output/all_pdb\")\n",
+        "\n",
+        "af_model._inputs[\"batch\"][\"dgram\"] = np.zeros((L,L,39))\n",
+        "with tqdm.notebook.tqdm(total=out[\"S\"].shape[0], bar_format=TQDM_BAR_FORMAT) as pbar:\n",
+        "  with open(\"design.fasta\",\"w\") as fasta:\n",
+        "    for n in range(num_seqs):\n",
+        "      seq = out[\"seq\"][n]\n",
+        "      af_model.predict(seq=seq,\n",
+        "                      num_recycles=1,\n",
+        "                      num_models=1,\n",
+        "                      verbose=False)\n",
+        "\n",
+        "      for t in af_terms:\n",
+        "        out[t].append(af_model.aux[\"log\"][t])\n",
+        "      out[\"pae\"][-1] = out[\"pae\"][-1] * 31\n",
+        "      af_model._save_results(save_best=True, verbose=False)\n",
+        "      af_model.save_current_pdb(f\"output/all_pdb/n{n}.pdb\")\n",
+        "      af_model._k += 1\n",
+        "\n",
+        "      line = f'>mpnn:{out[\"score\"][n]:.3f}_plddt:{out[\"plddt\"][n]:.3f}_ptm:{out[\"ptm\"][n]:.3f}_pae:{out[\"pae\"][n]:.3f}\\n{out[\"seq\"][n]}'\n",
+        "      fasta.write(line+\"\\n\")\n",
+        "      pbar.update(1)\n",
+        "\n",
+        "labels = [\"score\"] + af_terms + [\"seq\"]\n",
+        "data = [[out[k][n] for k in labels] for n in range(num_seqs)]\n",
+        "labels[0] = \"mpnn\"\n",
+        "\n",
+        "df = pd.DataFrame(data, columns=labels)\n",
+        "df.to_csv('output/mpnn_results.csv')\n",
+        "data_table.DataTable(df.round(3).sort_values(\"pae\"))"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "3H_B0AoYIiWH"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "v3dPQeEQcAcT"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(\"1.pdb\")\n",
+        "af_model.plot_pdb()\n",
+        "af_model.get_seqs()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/af_pseudo_diffusion_recycle.ipynb

@@ -0,0 +1,329 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_pseudo_diffusion_recycle.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "VD9K5H2cnFxL"
+      },
+      "source": [
+        "#AF_pseudo_diffusion + proteinMPNN\n",
+        "Hacking AlphaFold to be a diffusion model (for backbone generation) via recycling mechanism. At each step add logits from proteinMPNN.\n",
+        "\n",
+        "\n",
+        "**WARNING**: This notebook is experimental, designed as a control. Not intended for practical use at this stage."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "YCRzPGdTZfEe"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os, re\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "from colabdesign.shared.protein import _np_get_cb\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "import jax.numpy as jnp\n",
+        "import jax\n",
+        "from scipy.special import softmax, log_softmax\n",
+        "\n",
+        "import tqdm.notebook\n",
+        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "def sample_gumbel(shape, eps=1e-20): \n",
+        "  \"\"\"Sample from Gumbel(0, 1)\"\"\"\n",
+        "  U = np.random.uniform(size=shape)\n",
+        "  return -np.log(-np.log(U + eps) + eps)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "wxYMB4A9Zmrf"
+      },
+      "outputs": [],
+      "source": [
+        "#@title initialize the model\n",
+        "length = 100 #@param {type:\"integer\"}\n",
+        "#@markdown Provide a starting point (optional)\n",
+        "starting_seq = \"\" #@param {type:\"string\"}\n",
+        "starting_seq = re.sub(\"[^A-Z]\", \"\", starting_seq.upper())\n",
+        "#@markdown - if `starting_seq` provided the `length` option will be overwritten by length of starting sequence.\n",
+        "\n",
+        "if len(starting_seq) > 0:\n",
+        "  length = len(starting_seq)\n",
+        "\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "af_model.prep_inputs(length=length)\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "\n",
+        "af_model_test = mk_afdesign_model(protocol=\"fixbb\", best_metric=\"rmsd\")\n",
+        "mpnn_model_test = mk_mpnn_model()\n",
+        "\n",
+        "print(\"lengths\",af_model._lengths)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3Dt8i00UbxtW",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title run protocol\n",
+        "#@markdown Optimization options\n",
+        "iterations = 100 #@param [\"50\", \"100\", \"200\"] {type:\"raw\"}\n",
+        "\n",
+        "#AlphaFold options\n",
+        "use_dropout = True\n",
+        "num_recycles = 0\n",
+        "\n",
+        "#@markdown proteinMPNN options (set to `none` to disable)\n",
+        "mpnn_mode = \"conditional\" #@param [\"none\",\"conditional\", \"unconditional\"]\n",
+        "\n",
+        "#@markdown proteinMPNN contact map masking options\n",
+        "cmap_seqsep = 6 #@param {type:\"raw\"}\n",
+        "cmap_num = 1 #@param {type:\"integer\"}\n",
+        "cmap_cutoff = 8 #@param {type:\"raw\"}\n",
+        "\n",
+        "L = sum(af_model._lengths)\n",
+        "af_model.restart(mode=\"gumbel\")\n",
+        "af_model._args[\"clear_prev\"] = False\n",
+        "af_model.set_opt(cmap_cutoff=cmap_cutoff)\n",
+        "af_model.set_weights(helix=1e-8)\n",
+        "\n",
+        "# gather info about inputs\n",
+        "if \"offset\" in af_model._inputs:\n",
+        "  offset = af_model._inputs\n",
+        "else:\n",
+        "  idx = af_model._inputs[\"residue_index\"]\n",
+        "  offset = idx[:,None] - idx[None,:]\n",
+        "\n",
+        "# initialize sequence\n",
+        "if len(starting_seq) > 1:\n",
+        "  af_model.set_seq(seq=starting_seq)\n",
+        "\n",
+        "# initialize coordinates\n",
+        "af_model._inputs.pop(\"prev\",None)\n",
+        "af_model._inputs[\"batch\"] = {\"aatype\":np.zeros(L).astype(int),\n",
+        "                             \"all_atom_mask\":np.zeros((L,37)),\n",
+        "                             \"all_atom_positions\":np.zeros((L,37,3))}\n",
+        "\n",
+        "save_best = False\n",
+        "for k in range(iterations):\n",
+        "\n",
+        "  if k > (iterations - 10):\n",
+        "    use_dropout = False\n",
+        "    save_best = True\n",
+        "\n",
+        "\n",
+        "  # denoise\n",
+        "  aux = af_model.predict(return_aux=True, verbose=False,\n",
+        "                         dropout=use_dropout,\n",
+        "                         num_recycles=num_recycles)\n",
+        "  af_model._inputs[\"prev\"] = aux[\"prev\"]\n",
+        "  af_model._inputs[\"prev\"][\"prev_msa_first_row\"] *= 0\n",
+        "  af_model._inputs[\"prev\"][\"prev_pos\"] *= 0\n",
+        "\n",
+        "  # per position confidence\n",
+        "  cmap = aux[\"cmap\"] * (np.abs(offset) > cmap_seqsep)\n",
+        "  conf = np.sort(cmap)[:,-cmap_num:].mean(-1)\n",
+        "\n",
+        "  # gather features\n",
+        "  plddt = aux[\"plddt\"]\n",
+        "  seq = aux[\"seq\"][\"hard\"][0].argmax(-1)\n",
+        "  xyz = aux[\"atom_positions\"].copy()\n",
+        "      \n",
+        "  # update inputs    \n",
+        "  af_model._inputs[\"batch\"][\"aatype\"] = seq\n",
+        "  af_model._inputs[\"batch\"][\"all_atom_positions\"] = xyz\n",
+        "\n",
+        "  # add logits from proteinmpnn at each stage\n",
+        "  if mpnn_mode != \"none\":    \n",
+        "    mpnn_model.get_af_inputs(af_model)\n",
+        "    opt = {\"mask\":np.sqrt(conf)}    \n",
+        "    if mpnn_mode == \"unconditional\":\n",
+        "      opt[\"ar_mask\"] = np.zeros((L,L))\n",
+        "    mpnn_out = mpnn_model.score(**opt)\n",
+        "    mpnn_logits = mpnn_out[\"logits\"][:,:20]\n",
+        "    aux[\"log\"][\"mpnn\"] = mpnn_out[\"score\"]\n",
+        "    \n",
+        "    c = conf[:,None]\n",
+        "    new_logits = (1 - c) * sample_gumbel(mpnn_logits.shape) + c * mpnn_logits\n",
+        "    af_model._params[\"seq\"] = 0.9 * af_model._params[\"seq\"] + 0.1 * new_logits\n",
+        "\n",
+        "  # save results\n",
+        "  af_model._save_results(aux, save_best=save_best)\n",
+        "  af_model._k += 1\n",
+        "\n",
+        "af_model.save_pdb(\"init.pdb\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YeYD4KF8MUA_"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.plot_pdb()\n",
+        "af_model.get_seqs()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "6cJhizcYcUxm"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate(dpi=100))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "3H_B0AoYIiWH"
+      },
+      "outputs": [],
+      "source": [
+        "#@title sample new sequences using proteinMPNN and rescore with alphafold (w/o template)\n",
+        "#@markdown #### MPNN Options\n",
+        "num_seqs = 16 #@param [\"8\", \"16\", \"32\", \"64\", \"128\", \"256\", \"512\", \"1024\"] {type:\"raw\"}\n",
+        "sampling_temp = 0.1 \n",
+        "#@markdown #### AlphaFold Options\n",
+        "alphafold_model = \"model_4_ptm\" #@param [\"model_1_ptm\", \"model_2_ptm\", \"model_3_ptm\", \"model_4_ptm\", \"model_5_ptm\"]\n",
+        "num_recycles = 3 #@param [\"0\", \"1\", \"2\", \"3\"] {type:\"raw\"}\n",
+        "import pandas as pd\n",
+        "\n",
+        "# zero out template inputs\n",
+        "af_model_test.prep_inputs(\"init.pdb\")\n",
+        "mpnn_model_test.get_af_inputs(af_model_test)\n",
+        "out = mpnn_model_test.sample(num=num_seqs//8, batch=8,\n",
+        "                             temperature=sampling_temp)\n",
+        "af_terms = [\"plddt\",\"ptm\",\"pae\",\"rmsd\",\"dgram_cce\"]\n",
+        "for k in af_terms: out[k] = []\n",
+        "os.system(\"mkdir -p output/all_pdb\")\n",
+        "\n",
+        "with tqdm.notebook.tqdm(total=out[\"S\"].shape[0], bar_format=TQDM_BAR_FORMAT) as pbar:\n",
+        "  with open(\"design.fasta\",\"w\") as fasta:\n",
+        "    for n in range(num_seqs):\n",
+        "      seq = out[\"seq\"][n]\n",
+        "      af_model_test.predict(seq=seq,\n",
+        "                            num_recycles=num_recycles,\n",
+        "                            num_models=1,\n",
+        "                            verbose=False,\n",
+        "                            models=alphafold_model)\n",
+        "\n",
+        "      for t in af_terms:\n",
+        "        out[t].append(af_model_test.aux[\"log\"][t])\n",
+        "      out[\"pae\"][-1] = out[\"pae\"][-1] * 31\n",
+        "      af_model_test._save_results(save_best=True, verbose=False)\n",
+        "      af_model_test.save_current_pdb(f\"output/all_pdb/n{n}.pdb\")\n",
+        "      af_model_test._k += 1\n",
+        "\n",
+        "      line = f'>mpnn:{out[\"score\"][n]:.3f}_plddt:{out[\"plddt\"][n]:.3f}_ptm:{out[\"ptm\"][n]:.3f}_pae:{out[\"pae\"][n]:.3f}\\n{out[\"seq\"][n]}'\n",
+        "      fasta.write(line+\"\\n\")\n",
+        "      pbar.update(1)\n",
+        "\n",
+        "af_model_test.save_pdb(\"final.pdb\")\n",
+        "\n",
+        "labels = [\"score\"] + af_terms + [\"seq\"]\n",
+        "data = [[out[k][n] for k in labels] for n in range(num_seqs)]\n",
+        "labels[0] = \"mpnn\"\n",
+        "\n",
+        "df = pd.DataFrame(data, columns=labels)\n",
+        "df.to_csv('output/mpnn_results.csv')\n",
+        "df.round(3).sort_values(\"rmsd\")"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "v3dPQeEQcAcT"
+      },
+      "outputs": [],
+      "source": [
+        "af_model_test.plot_pdb()\n",
+        "af_model_test.get_seqs()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "gpuClass": "standard",
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/af_relax_design.ipynb

@@ -0,0 +1,400 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/af_relax_design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "#af_relax_design (WIP)\n",
+        "\n",
+        "\n",
+        "**Efficient and scalable de novo protein design using a relaxed sequence space**\n",
+        "\n",
+        "Christopher Josef Frank, Ali Khoshouei, Yosta de Stigter, Dominik Schiewitz, Shihao Feng, Sergey Ovchinnikov, Hendrik Dietz\n",
+        "\n",
+        "doi: https://doi.org/10.1101/2023.02.24.529906\n",
+        "\n",
+        "**<font color=\"red\">WARNING</font>** This notebook is in development, we are still working on adding all the options from the manuscript above."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install pyppeteer nest_asyncio\")\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.mpnn import mk_mpnn_model\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "import requests, time\n",
+        "if not os.path.isfile(\"TMscore\"):\n",
+        "  os.system(\"wget -qnc https://zhanggroup.org/TM-score/TMscore.cpp\")\n",
+        "  os.system(\"g++ -static -O3 -ffast-math -lm -o TMscore TMscore.cpp\")\n",
+        "def tmscore(x,y):\n",
+        "  # pass to TMscore\n",
+        "  output = os.popen(f'./TMscore {x} {y}')\n",
+        "  # parse outputs\n",
+        "  parse_float = lambda x: float(x.split(\"=\")[1].split()[0])\n",
+        "  o = {}\n",
+        "  for line in output:\n",
+        "    line = line.rstrip()\n",
+        "    if line.startswith(\"RMSD\"): o[\"rms\"] = parse_float(line)\n",
+        "    if line.startswith(\"TM-score\"): o[\"tms\"] = parse_float(line)\n",
+        "    if line.startswith(\"GDT-TS-score\"): o[\"gdt\"] = parse_float(line)\n",
+        "  return o\n",
+        "\n",
+        "import asyncio\n",
+        "import nest_asyncio\n",
+        "from pyppeteer import launch\n",
+        "import base64\n",
+        "\n",
+        "# Apply nest_asyncio to enable nested event loops\n",
+        "nest_asyncio.apply()\n",
+        "\n",
+        "async def fetch_blob_content(page, blob_url):\n",
+        "  blob_to_base64 = \"\"\"\n",
+        "  async (blobUrl) => {\n",
+        "      const blob = await fetch(blobUrl).then(r => r.blob());\n",
+        "      return new Promise((resolve) => {\n",
+        "          const reader = new FileReader();\n",
+        "          reader.onloadend = () => resolve(reader.result);\n",
+        "          reader.readAsDataURL(blob);\n",
+        "      });\n",
+        "  }\n",
+        "  \"\"\"\n",
+        "  base64_data = await page.evaluate(blob_to_base64, blob_url)\n",
+        "  _, encoded = base64_data.split(',', 1)\n",
+        "  return base64.b64decode(encoded)\n",
+        "\n",
+        "async def extract_pdb_file_download_link_and_content(url):\n",
+        "  browser = await launch(headless=True, args=['--no-sandbox', '--disable-setuid-sandbox'])\n",
+        "  page = await browser.newPage()\n",
+        "  await page.goto(url, {'waitUntil': 'networkidle0'})\n",
+        "  elements = await page.querySelectorAll('a.btn.bg-purple')\n",
+        "  for element in elements:\n",
+        "      href = await page.evaluate('(element) => element.getAttribute(\"href\")', element)\n",
+        "      if 'blob:https://esmatlas.com/' in href:\n",
+        "          content = await fetch_blob_content(page, href)\n",
+        "          await browser.close()\n",
+        "          return href, content\n",
+        "  await browser.close()\n",
+        "  return \"No PDB file link found.\", None\n",
+        "\n",
+        "def esmfold_api(sequence):\n",
+        "  url = f'https://esmatlas.com/resources/fold/result?fasta_header=%3Eunnamed&sequence={sequence}'\n",
+        "  result = asyncio.get_event_loop().run_until_complete(extract_pdb_file_download_link_and_content(url))\n",
+        "  if result[1]:\n",
+        "      pdb_str = result[1].decode('utf-8')\n",
+        "      return pdb_str\n",
+        "  else:\n",
+        "      return \"Failed to retrieve PDB content.\"\n",
+        "\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "from colabdesign.af.alphafold.common import residue_constants"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "sZnYfCbfEvol",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title # hallucination\n",
+        "#@markdown For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n",
+        "LENGTH = 100 #@param {type:\"integer\"}\n",
+        "COPIES = 1 #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"6\", \"7\", \"8\"] {type:\"raw\"}\n",
+        "MODE = \"manuscript\" #@param [\"original\", \"manuscript\"]\n",
+        "use_rg_loss = True #@param {type:\"boolean\"}\n",
+        "\n",
+        "#@markdown ProteinMPNN Settings\n",
+        "use_mpnn_loss = False #@param {type:\"boolean\"}\n",
+        "use_solubleMPNN = False #@param {type:\"boolean\"}\n",
+        "#@markdown\n",
+        "\n",
+        "def add_rg_loss(self, weight=0.1):\n",
+        "  '''add radius of gyration loss'''\n",
+        "  def loss_fn(inputs, outputs):\n",
+        "    xyz = outputs[\"structure_module\"]\n",
+        "    ca = xyz[\"final_atom_positions\"][:,residue_constants.atom_order[\"CA\"]]\n",
+        "    if self.protocol == \"binder\":\n",
+        "      ca = ca[-self._binder_len:]\n",
+        "    if MODE == \"manuscript\":\n",
+        "      ca = ca[::5]\n",
+        "    rg = jnp.sqrt(jnp.square(ca - ca.mean(0)).sum(-1).mean() + 1e-8)\n",
+        "    if MODE == \"original\":\n",
+        "      rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "      rg = jax.nn.elu(rg - rg_th)\n",
+        "    return {\"rg\":rg}\n",
+        "  self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n",
+        "  self.opt[\"weights\"][\"rg\"] = weight\n",
+        "\n",
+        "def add_mpnn_loss(self, mpnn=0.1, mpnn_seq=0.0):\n",
+        "  '''\n",
+        "  add mpnn loss\n",
+        "  mpnn = maximize confidence of proteinmpnn\n",
+        "  mpnn_seq = push designed sequence to match proteinmpnn logits\n",
+        "  '''\n",
+        "\n",
+        "  self._mpnn = mk_mpnn_model(weights = \"soluble\" if use_solubleMPNN else \"original\")\n",
+        "  def loss_fn(inputs, outputs, aux, key):\n",
+        "\n",
+        "    # get structure\n",
+        "    atom_idx = tuple(residue_constants.atom_order[k] for k in [\"N\",\"CA\",\"C\",\"O\"])\n",
+        "    I = {\"S\":           inputs[\"aatype\"],\n",
+        "         \"residue_idx\": inputs[\"residue_index\"],\n",
+        "         \"chain_idx\":   inputs[\"asym_id\"],\n",
+        "         \"X\":           outputs[\"structure_module\"][\"final_atom_positions\"][:,atom_idx],\n",
+        "         \"mask\":        outputs[\"structure_module\"][\"final_atom_mask\"][:,1],\n",
+        "         \"lengths\":     self._lengths,\n",
+        "         \"key\":         key}\n",
+        "\n",
+        "    if \"offset\" in inputs:\n",
+        "      I[\"offset\"] = inputs[\"offset\"]\n",
+        "\n",
+        "    # set autoregressive mask\n",
+        "    L = sum(self._lengths)\n",
+        "    if self.protocol == \"binder\":\n",
+        "      I[\"ar_mask\"] = 1 - np.eye(L)\n",
+        "      I[\"ar_mask\"][-self._len:,-self._len:] = 0\n",
+        "    else:\n",
+        "      I[\"ar_mask\"] = np.zeros((L,L))\n",
+        "\n",
+        "    # get logits\n",
+        "    logits = self._mpnn._score(**I)[\"logits\"][:,:20]\n",
+        "    if self.protocol == \"binder\":\n",
+        "      logits = logits[-self._len:]\n",
+        "    else:\n",
+        "      logits = logits[:self._len]\n",
+        "    aux[\"mpnn_logits\"] = logits\n",
+        "\n",
+        "    # compute loss\n",
+        "    log_q = jax.nn.log_softmax(logits)\n",
+        "    p = inputs[\"seq\"][\"hard\"]\n",
+        "    q = jax.nn.softmax(logits)\n",
+        "    losses = {}\n",
+        "    losses[\"mpnn\"] = -log_q.max(-1).mean()\n",
+        "    losses[\"mpnn_seq\"] = -(p * jax.lax.stop_gradient(log_q)).sum(-1).mean()\n",
+        "    return losses\n",
+        "\n",
+        "  self._callbacks[\"model\"][\"loss\"].append(loss_fn)\n",
+        "  self.opt[\"weights\"][\"mpnn\"] = mpnn\n",
+        "  self.opt[\"weights\"][\"mpnn_seq\"] = mpnn_seq\n",
+        "\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "af_model.prep_inputs(length=LENGTH, copies=COPIES)\n",
+        "\n",
+        "# add extra losses\n",
+        "if use_rg_loss:   add_rg_loss(af_model)\n",
+        "if use_mpnn_loss: add_mpnn_loss(af_model)\n",
+        "\n",
+        "print(\"length\",af_model._lengths)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.restart()\n",
+        "if MODE == \"original\":\n",
+        "  # pre-design with gumbel initialization and softmax activation\n",
+        "  af_model.set_weights(plddt=0.0, pae=0.0)\n",
+        "  af_model.set_seq(mode=[\"gumbel\"])\n",
+        "  af_model.design_soft(50)\n",
+        "  af_model.set_seq(af_model.aux[\"seq\"][\"pseudo\"])\n",
+        "\n",
+        "if MODE == \"manuscript\":\n",
+        "  af_model.set_seq(mode=[\"gumbel\",\"soft\"])\n",
+        "\n",
+        "af_model.set_weights(plddt=1.0, pae=1.0)\n",
+        "af_model.design_logits(40)\n",
+        "af_model.design_logits(10, save_best=True)"
+      ],
+      "metadata": {
+        "id": "f76xqCkw0vj9"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "A1GxeLZdTTya"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "L2E9Tn2Acchj"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YSKWYu0_GlUH"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.get_seqs()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown #Redesign with ProteinMPNN\n",
+        "num_seqs = 8 #@param [\"8\", \"16\", \"32\", \"64\"] {type:\"raw\"}\n",
+        "mpnn_sampling_temp = 0.1 #@param [\"0.0001\", \"0.1\", \"0.15\", \"0.2\", \"0.25\", \"0.3\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
+        "rm_aa = \"C\" #@param {type:\"string\"}\n",
+        "use_solubleMPNN = False #@param {type:\"boolean\"}\n",
+        "#@markdown - `mpnn_sampling_temp` - control diversity of sampled sequences. (higher = more diverse).\n",
+        "#@markdown - `rm_aa='C'` - do not use [C]ysteines.\n",
+        "#@markdown - `use_solubleMPNN` - use weights trained only on soluble proteins. See [preprint](https://www.biorxiv.org/content/10.1101/2023.05.09.540044v2).\n",
+        "#@markdown"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "m2qAYsDsCfqJ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from colabdesign.shared.protein import alphabet_list as chain_list\n",
+        "mpnn_model = mk_mpnn_model()\n",
+        "mpnn_model.prep_inputs(pdb_filename=f\"{af_model.protocol}.pdb\",\n",
+        "                       chain=\",\".join(chain_list[:COPIES]),\n",
+        "                       homooligmer=COPIES>1,\n",
+        "                       rm_aa=rm_aa,\n",
+        "                       weights = \"soluble\" if use_solubleMPNN else\"original\")\n",
+        "out = mpnn_model.sample(num=num_seqs//8,\n",
+        "                        batch=8,\n",
+        "                        temperature=mpnn_sampling_temp)\n",
+        "for seq,score in zip(out[\"seq\"],out[\"score\"]):\n",
+        "  print(score,seq.split(\"/\")[0])"
+      ],
+      "metadata": {
+        "id": "uQa0FAp7bGQo"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Run ESMfold"
+      ],
+      "metadata": {
+        "id": "eDvyemgjNbX4"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "print(\"# rmsd tmscore sequence\")\n",
+        "best = {}\n",
+        "best_rmsd = None\n",
+        "for n,seq in enumerate(out[\"seq\"]):\n",
+        "  x = seq.split(\"/\")[0]\n",
+        "  with open(f\"{af_model.protocol}.esmfold.{n}.pdb\",\"w\") as handle:\n",
+        "    pdb_str = esmfold_api(x)\n",
+        "    handle.write(pdb_str)\n",
+        "  o = tmscore(f\"{af_model.protocol}.pdb\",\n",
+        "              f\"{af_model.protocol}.esmfold.{n}.pdb\")\n",
+        "  print(n,o[\"rms\"],o[\"tms\"],x)\n",
+        "  if best_rmsd is None or o[\"rms\"] < best_rmsd:\n",
+        "    best_rmsd = o[\"rms\"]\n",
+        "    best = {**o,\"seq\":x}"
+      ],
+      "metadata": {
+        "id": "Ey29NmNAFtK0"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "best"
+      ],
+      "metadata": {
+        "id": "ltH6cLw5NhuX"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/af_single_sequential.ipynb

@@ -0,0 +1,164 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/af_single_sequential.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title af_single_sequential\n",
+        "#@markdown implementing [@MartinPacesa](https://twitter.com/MartinPacesa/status/1640005299124707328)'s idea\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  print(\"installing AlphaFold...\")\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  print(\"downloading AlphaFold params...\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-03-02.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-03-02.tar -C params\")\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "import tqdm.notebook\n",
+        "TQDM_BAR_FORMAT = '{l_bar}{bar}| {n_fmt}/{total_fmt} [elapsed: {elapsed} remaining: {remaining}]'\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\"\n",
+        "\n",
+        "#@markdown Specify either pdb/pdb_chain or sequence:\n",
+        "pdb = \"6D0T\" #@param {type:\"string\"}\n",
+        "pdb_chain = \"A\" #@param {type:\"string\"}\n",
+        "sequence = \"\" #@param {type:\"string\"}\n",
+        "num_recycles = 0 #@param [\"0\", \"1\", \"2\", \"4\", \"8\", \"16\"] {type:\"raw\"}\n",
+        "clear_mem()\n",
+        "if sequence == \"\":\n",
+        "  pdb_filename = get_pdb(pdb)\n",
+        "  af_model = mk_afdesign_model(protocol=\"fixbb\")\n",
+        "  af_model.prep_inputs(pdb_filename=pdb_filename, chain=pdb_chain)\n",
+        "  af_model.set_seq(mode=\"wildtype\")\n",
+        "  mode = \"pdb\"\n",
+        "else:\n",
+        "  af_model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "  af_model.prep_inputs(length=len(sequence))\n",
+        "  af_model.set_seq(sequence)\n",
+        "  mode = \"free\"\n",
+        "\n",
+        "import random, string\n",
+        "path = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))\n",
+        "while os.path.isdir(path):\n",
+        "  path = ''.join(random.choices(string.ascii_lowercase + string.digits, k=5))\n",
+        "\n",
+        "seq = af_model._params[\"seq\"].copy()\n",
+        "os.makedirs(path, exist_ok=True)\n",
+        "with tqdm.notebook.tqdm(total=af_model._len, bar_format=TQDM_BAR_FORMAT) as pbar:\n",
+        "  for n in range(1,af_model._len+1):\n",
+        "  \n",
+        "    # replace sequence (this is just for visualization)\n",
+        "    af_model._params[\"seq\"] = seq.copy()\n",
+        "    af_model._params[\"seq\"][:,n:] = 0\n",
+        "    \n",
+        "    # mask part of the input features\n",
+        "    af_model._inputs[\"seq_mask\"] = (np.arange(af_model._len) < n).astype(float)\n",
+        "    af_model.predict(verbose=False, hard=False, num_recycles=num_recycles)\n",
+        "    af_model._save_results(verbose=False)\n",
+        "    af_model.save_current_pdb(f\"{path}/{n:04d}.pdb\")\n",
+        "    pbar.update(1)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "HTML(af_model.animate(dpi=100))"
+      ],
+      "metadata": {
+        "id": "lphWx8wkEpwP"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.plot_pdb()"
+      ],
+      "metadata": {
+        "id": "9h1uXK44TygM"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Package and download results\n",
+        "#@markdown If you are having issues downloading the result archive, \n",
+        "#@markdown try disabling your adblocker and run this cell again. \n",
+        "#@markdown  If that fails click on the little folder icon to the \n",
+        "#@markdown  left, navigate to file: `name.result.zip`, \n",
+        "#@markdown  right-click and select \\\"Download\\\" \n",
+        "#@markdown (see [screenshot](https://pbs.twimg.com/media/E6wRW2lWUAEOuoe?format=jpg&name=small)).\n",
+        "from google.colab import files\n",
+        "!zip -r {path}.result.zip {path}/*\n",
+        "files.download(f\"{path}.result.zip\")"
+      ],
+      "metadata": {
+        "id": "zGE1CJxuTzyQ",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    },
+    "gpuClass": "standard"
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/afdesign_hotspot_test.ipynb

@@ -0,0 +1,247 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/afdesign_hotspot_test.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "# AfDesign - binder + hotspot test"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "-AXy0s_4cKaK",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]      
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "%%bash\n",
+        "wget -qnc https://raw.githubusercontent.com/casperg92/MaSIF_colab/main/example/monomerexample.pdb"
+      ],
+      "metadata": {
+        "id": "TIGDV-NqGMAw"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "3XLJPiRKx5Mw"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "model = mk_afdesign_model(protocol=\"binder\")\n",
+        "\n",
+        "# pos = define positions (hotspot) on target you want to target for binder\n",
+        "model.prep_inputs(pdb_filename=\"monomerexample.pdb\", chain=\"A\",\n",
+        "                  binder_len=20,\n",
+        "                  hotspot=\"33,17,6,34,30,19,37,15\") \n",
+        "\n",
+        "print(\"target_length\",model._target_len)\n",
+        "print(\"binder_length\",model._binder_len)\n",
+        "print(\"weights\",model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "u6VxjuinyCZa"
+      },
+      "outputs": [],
+      "source": [
+        "model.restart(mode=\"soft_gumbel\")\n",
+        "model.design_3stage(100, 100, 10)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "9cARoviGyIKb"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "RzE137NDZdZc"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model.plot_pdb()\n",
+        "model.save_pdb(f\"{model.protocol}.pdb\")"
+      ],
+      "metadata": {
+        "id": "sTlS7_L8Zfwf"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "as a control, let's try hallucinate a binder without defining a hotspot"
+      ],
+      "metadata": {
+        "id": "aGW9s5t3Xlk4"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "clear_mem()\n",
+        "model = mk_afdesign_model(protocol=\"binder\")\n",
+        "model.prep_inputs(pdb_filename=\"monomerexample.pdb\", chain=\"A\", binder_len=20)"
+      ],
+      "metadata": {
+        "id": "KzkNzuS6Xpqj"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model.restart(mode=\"soft_gumbel\")\n",
+        "model.design_3stage(100, 100, 10)"
+      ],
+      "metadata": {
+        "id": "wjFi2uhwXtPz"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "HTML(model.animate())"
+      ],
+      "metadata": {
+        "id": "pBRNpspoX6pR"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "J43oqomrYA8v"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model.plot_pdb()\n",
+        "model.save_pdb(f\"{model.protocol}.control.pdb\")"
+      ],
+      "metadata": {
+        "id": "r2d86kMVX_SL"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "xw_5hblQmRqn"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "afdesign_hotspot_test.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/binder_hallucination.ipynb

@@ -0,0 +1,42 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/peptide_binder_design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "this notebook moved to: [peptide_binder_design.ipynb](https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/peptide_binder_design.ipynb)"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/disulfide_design.ipynb

@@ -0,0 +1,348 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/disulfide_design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "YRk_XS2NILil"
+      },
+      "source": [
+        "# disulfide-hallucination\n",
+        "For a given length, generate/hallucinate a disulfide-rich sequence that AlphaFold thinks folds into a well structured protein"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "7JbLyVu4ILil"
+      },
+      "outputs": [],
+      "source": [
+        "import random\n",
+        "from jax.lax import dynamic_slice\n",
+        "import jax.numpy as jnp\n",
+        "\n",
+        "from colabdesign.af.loss import _get_con_loss\n",
+        "def get_con_loss(dgram, dgram_bins, cutoff=None, binary=True,\n",
+        "                 num=1, seqsep=0, offset=None):\n",
+        "  '''convert distogram into contact loss'''  \n",
+        "  x = _get_con_loss(dgram, dgram_bins, cutoff, binary)  \n",
+        "  a,b = x.shape\n",
+        "  if offset is None:\n",
+        "    mask = jnp.abs(jnp.arange(a)[:,None] - jnp.arange(b)[None,:]) >= seqsep\n",
+        "  else:\n",
+        "    mask = jnp.abs(offset) >= seqsep\n",
+        "  x = jnp.sort(jnp.where(mask,x,jnp.nan))\n",
+        "  k_mask = (jnp.arange(b) < num) * (jnp.isnan(x) == False)    \n",
+        "  return jnp.where(k_mask,x,0.0).sum(-1) / (k_mask.sum(-1) + 1e-8)\n",
+        "\n",
+        "def generate_disulfide_pattern(L, disulfide_num, min_sep=5):\n",
+        "    disulfide_pattern = []\n",
+        "    positions = list(range(L))\n",
+        "    for n in range(disulfide_num):\n",
+        "        for _ in range(100): # try 100 time per postion.\n",
+        "            i,j = random.sample(positions, k=2)\n",
+        "            if abs(i-j)<=min_sep: continue # set min loop len.\n",
+        "            positions.remove(i)\n",
+        "            positions.remove(j)\n",
+        "            disulfide_pattern.append((i,j))\n",
+        "            # check\n",
+        "            if _ > 99:\n",
+        "                print('Not find good disulfide_pos! exit....')\n",
+        "                return 0  # not good pose!\n",
+        "            else:\n",
+        "                break\n",
+        "    sequence_pattern = list('X'*L)\n",
+        "    for pair in disulfide_pattern:\n",
+        "        for i in pair: sequence_pattern[i] = 'C'\n",
+        "\n",
+        "    return disulfide_pattern, ''.join(sequence_pattern), L\n",
+        "\n",
+        "def disulfide_loss(inputs, outputs):\n",
+        "  def get_disulfide_loss(dgram, dgram_bins, disulfide_pattern):\n",
+        "    '''\n",
+        "    Func: simple disulfide loss, make the contacts < 7.0/7.5A.\n",
+        "    # see: https://www.ncbi.nlm.nih.gov/pmc/articles/PMC7316719/\n",
+        "    params: disulfide_pattern: List[(pos1, pos2), (pos3, pos4)...]\n",
+        "    '''\n",
+        "    disulfide_loss = 0.0\n",
+        "    for pair in disulfide_pattern:\n",
+        "      i,j = pair\n",
+        "      pair_dgram = dynamic_slice(dgram, (i,j,0), (1,1,len(dgram_bins))) + dynamic_slice(dgram, (j,i,0), (1,1,len(dgram_bins)))\n",
+        "      disulfide_loss += get_con_loss(pair_dgram, dgram_bins, cutoff=7.0, binary=False, num=1)\n",
+        "    return disulfide_loss.mean()\n",
+        "\n",
+        "  # add disulfide loss here:\n",
+        "  dgram_logits = outputs['distogram']['logits']\n",
+        "  dgram_bins = jnp.append(0, outputs[\"distogram\"][\"bin_edges\"])\n",
+        "  return {\"disulfide\":get_disulfide_loss(dgram_logits, dgram_bins, inputs['opt']['disulfide_pattern'])}"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "aU4e-ZkfILim"
+      },
+      "source": [
+        "0. generate_cys_pattern"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "vrSatV2sILim"
+      },
+      "outputs": [],
+      "source": [
+        "disulfide_pattern, sequence_pattern, L = generate_disulfide_pattern(35, 3)\n",
+        "print(disulfide_pattern)\n",
+        "print(sequence_pattern)\n",
+        "print(L)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "jGNGjHeyILim"
+      },
+      "source": [
+        "1. hallucination with disulfide-pattern"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "HmsBDNhmILin"
+      },
+      "outputs": [],
+      "source": [
+        "# using some disulfide sequence as input:\n",
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\", loss_callback=disulfide_loss)\n",
+        "af_model.opt[\"weights\"][\"disulfide\"] = 1.0\n",
+        "af_model.prep_inputs(length=L)\n",
+        "\n",
+        "print(\"length\",af_model._len)\n",
+        "print(\"weights\",af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "MmBtBMnBILin"
+      },
+      "outputs": [],
+      "source": [
+        "# set disulfide_pattern sequence.\n",
+        "af_model.restart(seq=sequence_pattern, add_seq=True, rm_aa='C')\n",
+        "# set disulfide_pattern:\n",
+        "af_model.opt['disulfide_pattern'] = disulfide_pattern\n",
+        "# reweight con:\n",
+        "af_model.opt[\"weights\"]['con'] = 0.5"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.design_3stage(50,50,10)"
+      ],
+      "metadata": {
+        "id": "wstI0n4zMydi"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ewlj8671ILio"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb(show_sidechains=True)"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "iVEqJKaKILio"
+      },
+      "source": [
+        "2. fix-sidechain with pyRosetta"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "VGwyvaodILio"
+      },
+      "outputs": [],
+      "source": [
+        "from pyrosetta import init, pose_from_pdb, Pose, create_score_function\n",
+        "from pyrosetta.rosetta.protocols.relax import FastRelax\n",
+        "from pyrosetta.rosetta.core.pack.task import TaskFactory\n",
+        "from pyrosetta.rosetta.core.kinematics import MoveMap\n",
+        "from pyrosetta.rosetta.core.pack.task.operation import InitializeFromCommandline\n",
+        "from pyrosetta.rosetta.protocols.denovo_design import DisulfidizeMover\n",
+        "from pyrosetta.rosetta.core.pack.task.operation import RestrictToRepacking\n",
+        "from pyrosetta.rosetta.core.select.residue_selector import ChainSelector\n",
+        "\n",
+        "def fastrelax(pose):\n",
+        "    full_score = create_score_function('ref2015')\n",
+        "    tf = TaskFactory()\n",
+        "    # tf.push_back(InitializeFromCommandline())\n",
+        "    tf.push_back(RestrictToRepacking())\n",
+        "    mmap = MoveMap()\n",
+        "    mmap.set_bb(True)\n",
+        "    mmap.set_chi(True)\n",
+        "\n",
+        "    # FastRelax:\n",
+        "    fast_design = FastRelax(full_score, 4)  # repeat:6\n",
+        "    fast_design.set_task_factory(tf)\n",
+        "    fast_design.set_movemap(mmap)\n",
+        "    fast_design.apply(pose)\n",
+        "    \n",
+        "def build_stapled_pose(pose, cys_pattern):\n",
+        "    # rebuild \n",
+        "    disulfidizer = DisulfidizeMover()\n",
+        "    disulfidizer.set_match_rt_limit(999.0)\n",
+        "    disulfidizer.set_max_disulf_score(999.0)\n",
+        "    for pair in cys_pattern:\n",
+        "        full_score = create_score_function('ref2015')\n",
+        "        stapled_pose = Pose().assign(pose)  # pose for disulfides stapled\n",
+        "        residue1, residue2 = pair\n",
+        "        disulfidizer.make_disulfide(stapled_pose, residue1+1, residue2+1, False, full_score)\n",
+        "        # fastdesign mover\n",
+        "        pdbnum1 = stapled_pose.pdb_info().pose2pdb(residue1)\n",
+        "        pdbnum2 = stapled_pose.pdb_info().pose2pdb(residue2)\n",
+        "        print('# OPTIMIZE:  Around These Residues: %s %s' % (pdbnum1, pdbnum2))\n",
+        "\n",
+        "    # save pdb:\n",
+        "    fastrelax(stapled_pose)\n",
+        "    \n",
+        "    # filter: all stapled?\n",
+        "    v = ChainSelector(1).apply(stapled_pose)\n",
+        "    l = disulfidizer.find_current_disulfides(stapled_pose, v, v)\n",
+        "    print(len([i for i in l]), len(cys_pattern))\n",
+        "    if len([i for i in l]) == len(cys_pattern):\n",
+        "        return stapled_pose\n",
+        "    else:\n",
+        "        return 0\n",
+        "    "
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "zhMey-4lILip"
+      },
+      "outputs": [],
+      "source": [
+        "# optimze the side-chain\n",
+        "init('-mute all')\n",
+        "pose = pose_from_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "stapled_pose = build_stapled_pose(pose, disulfide_pattern)"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "disulfide_design.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3.8.10 64-bit",
+      "language": "python",
+      "name": "python3"
+    },
+    "language_info": {
+      "codemirror_mode": {
+        "name": "ipython",
+        "version": 3
+      },
+      "file_extension": ".py",
+      "mimetype": "text/x-python",
+      "name": "python",
+      "nbconvert_exporter": "python",
+      "pygments_lexer": "ipython3",
+      "version": "3.8.10"
+    },
+    "vscode": {
+      "interpreter": {
+        "hash": "31f2aee4e71d21fbe5cf8b01ff0e069b9275f58929596ceb00d14d90e3e16cd6"
+      }
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/hallucination.ipynb

@@ -0,0 +1,193 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/hallucination.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "# AfDesign - hallucination\n",
+        "For a given length, generate/hallucinate a protein sequence that AlphaFold thinks folds into a well structured protein (high plddt, low pae, many contacts).\n",
+        "\n",
+        "**WARNING**\n",
+        "1.   This notebook is in active development and was designed for demonstration purposes only.\n",
+        "2.   Using AfDesign as the only \"loss\" function for design might be a bad idea, you may find adversarial sequences (aka. sequences that trick AlphaFold)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title Hallucination Options\n",
+        "length =  100#@param {type:\"integer\"}\n",
+        "copies =  1#@param {type:\"integer\"}\n",
+        "#@markdown ####Weights\n",
+        "#@markdown - Minimizing `pae` or maximizing `plddt` often results in a single helix.\n",
+        "#@markdown To avoid this, we start with a random sequence and instead try to optimize \n",
+        "#@markdown defined `num`ber of `con`tacts per position. \n",
+        "pae = 0.1 #@param [\"0.01\", \"0.1\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
+        "plddt = 0.1 #@param [\"0.01\", \"0.1\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
+        "con = 1.0 #@param [\"0.01\", \"0.1\", \"0.5\", \"1.0\"] {type:\"raw\"}\n",
+        "#@markdown ####Contact Definition\n",
+        "#@markdown - The contact definition is based on Cb-Cb diststance `cutoff`. To avoid \n",
+        "#@markdown biasing towards helical contact, only contacts with sequence seperation > \n",
+        "#@markdown `seqsep` are considered.\n",
+        "\n",
+        "\n",
+        "seqsep = 9 #@param [\"0\",\"5\",\"9\"] {type:\"raw\"}\n",
+        "cutoff = \"14\" #@param [\"8\", \"14\", \"max\"]\n",
+        "num = \"2\" #@param [\"1\", \"2\", \"4\", \"8\", \"max\"]\n",
+        "binary = False #@param {type:\"boolean\"}\n",
+        "if cutoff == \"max\": cutoff = 21.6875\n",
+        "if num == \"max\": num = length\n",
+        "\n",
+        "opt = {\"con\":{\"seqsep\":int(seqsep),\"cutoff\":float(cutoff),\"num\":int(num),\n",
+        "              \"binary\":binary}}\n",
+        "weights = {\"con\":float(con), \"pae\":float(pae),\"plddt\":float(plddt)}\n",
+        "\n",
+        "if \"model\" not in dir() or model._len != length or model._args[\"copies\"] != copies:\n",
+        "  clear_mem()\n",
+        "  model = mk_afdesign_model(protocol=\"hallucination\")\n",
+        "  model.prep_inputs(length=length, copies=copies)\n",
+        "\n",
+        "#@markdown ####Optimizer settings\n",
+        "pre_iters =  100 #@param {type:\"integer\"}\n",
+        "soft_iters = 50 #@param {type:\"integer\"}\n",
+        "temp_iters = 50 #@param {type:\"integer\"}\n",
+        "hard_iters = 10 #@param {type:\"integer\"}\n",
+        "\n",
+        "# pre-design with gumbel initialization and softmax activation\n",
+        "model.restart(mode=\"gumbel\",opt=opt,weights=weights)\n",
+        "model.design_soft(pre_iters)\n",
+        "\n",
+        "# three stage design  \n",
+        "model.set_seq(model.aux[\"seq\"][\"pseudo\"])\n",
+        "model.design_3stage(soft_iters,temp_iters,hard_iters)"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "eCGc3J663NGz"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "A1GxeLZdTTya",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@markdown ## display hallucinated protein {run: \"auto\"}\n",
+        "color = \"pLDDT\" #@param [\"chain\", \"pLDDT\", \"rainbow\"]\n",
+        "show_sidechains = False #@param {type:\"boolean\"}\n",
+        "show_mainchains = False #@param {type:\"boolean\"}\n",
+        "model.save_pdb(f\"{model.protocol}.pdb\")\n",
+        "model.plot_pdb(show_sidechains=show_sidechains,\n",
+        "                   show_mainchains=show_mainchains,\n",
+        "                   color=color)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "L2E9Tn2Acchj"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YSKWYu0_GlUH"
+      },
+      "outputs": [],
+      "source": [
+        "model.get_seqs()"
+      ]
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "hallucination.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/hallucination_custom_loss.ipynb

@@ -0,0 +1,260 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/hallucination_custom_loss.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "# AfDesign - hallucination custom loss example\n",
+        "Backprop through AlphaFold for protein design."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "UUfKrOzT0gOS"
+      },
+      "source": [
+        "# Custom loss"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qLd1DsnKzxBJ"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\", debug=True)\n",
+        "af_model.prep_inputs(length=100)\n",
+        "\n",
+        "print(\"length\",  af_model._len)\n",
+        "print(\"weights\", af_model.opt[\"weights\"])"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.restart(mode=\"gumbel\",seed=0)\n",
+        "af_model.set_opt(soft=True)\n",
+        "af_model.run(backprop=False)"
+      ],
+      "metadata": {
+        "id": "u0AwskJ84NGx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "import jax\n",
+        "import jax.numpy as jnp\n",
+        "\n",
+        "# first off, let's implement a custom Radius of Gyration loss function\n",
+        "def rg_loss(inputs, outputs):\n",
+        "  positions = outputs[\"structure_module\"][\"final_atom_positions\"]\n",
+        "  ca = positions[:,residue_constants.atom_order[\"CA\"]]\n",
+        "  center = ca.mean(0)\n",
+        "  rg = jnp.sqrt(jnp.square(ca - center).sum(-1).mean() + 1e-8)\n",
+        "  rg_th = 2.38 * ca.shape[0] ** 0.365\n",
+        "  rg = jax.nn.elu(rg - rg_th)\n",
+        "  return {\"rg\":rg}"
+      ],
+      "metadata": {
+        "id": "SGxkLR_4VsQI"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.aux[\"debug\"].keys()"
+      ],
+      "metadata": {
+        "id": "LAgsoVLhcJdr"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "rg_loss(**af_model.aux[\"debug\"])"
+      ],
+      "metadata": {
+        "id": "enMbbgFnWLwF"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "#Let's add it to the model!"
+      ],
+      "metadata": {
+        "id": "jmHGCynter0p"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "clear_mem()\n",
+        "af_model = mk_afdesign_model(protocol=\"hallucination\",\n",
+        "                             debug=False,\n",
+        "                             loss_callback=rg_loss) # add our custom loss\n",
+        "af_model.opt[\"weights\"][\"rg\"] = 0.1 # add our loss to weights (so we can later control it)\n",
+        "af_model.prep_inputs(length=100)\n",
+        "print(\"weights\", af_model.opt[\"weights\"])"
+      ],
+      "metadata": {
+        "id": "0D7Z0U6aVD1V"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.restart(mode=\"gumbel\", seed=0)\n",
+        "af_model.design_soft(50)\n",
+        "\n",
+        "# three stage design  \n",
+        "af_model.set_seq(af_model.aux[\"seq\"][\"pseudo\"])\n",
+        "af_model.design_3stage(50,50,10)"
+      ],
+      "metadata": {
+        "id": "Wsc9IWsqXreX"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YEApO8YzBoS0"
+      },
+      "outputs": [],
+      "source": [
+        "af_model.save_pdb(f\"{af_model.protocol}.pdb\")\n",
+        "af_model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cW1KQiHKJpfp"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(af_model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "af_model.get_seqs()"
+      ],
+      "metadata": {
+        "id": "YDrChASGVUUx"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "-js6TX9ZytR9"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "hallucination_custom_loss.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/partial_hallucination_rewire.ipynb

@@ -0,0 +1,161 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/partial_hallucination_rewire.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "# AfDesign - rewire (partial hallucination)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# setup the model\n",
+        "clear_mem()\n",
+        "model = mk_afdesign_model(protocol=\"partial\",\n",
+        "                          use_templates=False) # set True to constrain positions using template input\n",
+        "                        \n",
+        "# define positions we want to constrain (input PDB numbering)\n",
+        "old_pos = \"4-9,13-26,30-33,37-50,54-59,63-76,80-84,112-125\"\n",
+        "\n",
+        "model.prep_inputs(get_pdb(\"2MR5\"), chain=\"A\",\n",
+        "                  pos=old_pos,               # define positions to contrain\n",
+        "                  length=100,                # define if the desired length is different from input PDB\n",
+        "                  fix_seq=False)             # set True to constrain the sequence\n",
+        "\n",
+        "# set positions (if different from PDB)\n",
+        "# reorder the segments, \n",
+        "model.rewire(order = [1,2,3,0,7,4,5,6], # set order of segments\n",
+        "             loops =  [3,3,3,4,3,3,4],  # change loop length inbetween segments\n",
+        "             offset=0)                  # essentially loop length at the N term\n",
+        "\n",
+        "print(model.opt[\"pos\"])"
+      ],
+      "metadata": {
+        "id": "_m7ZaaP4NkII"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# restart optimization\n",
+        "model.restart()\n",
+        "# model.set_seq(mode=\"wildtype\") # initialize with wildtype sequence\n",
+        "\n",
+        "# balance weights [dgram_cce = restraint weight], [con = hallucination weight]\n",
+        "model.set_weights(dgram_cce=1, con=0)\n",
+        "model.design_3stage(300,100,10)"
+      ],
+      "metadata": {
+        "id": "bDA8cPqMT_eY"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "model.save_pdb(f\"{model.protocol}.pdb\")\n",
+        "model.plot_pdb()"
+      ],
+      "metadata": {
+        "id": "qFegtgcnHdqr"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "HTML(model.animate())"
+      ],
+      "metadata": {
+        "id": "9sjG7kvoHaVi"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "partial_hallucination_rewire.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/peptide_binder_design.ipynb

@@ -0,0 +1,319 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/peptide_binder_design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "# AfDesign - peptide binder design\n",
+        "For a given protein target and protein binder length, generate/hallucinate a protein binder sequence AlphaFold thinks will bind to the target structure. To do this, we maximize number of contacts at the interface and maximize pLDDT of the binder.\n",
+        "\n",
+        "**WARNING**\n",
+        "1.   This notebook is in active development and was designed for demonstration purposes only.\n",
+        "2.   Using AfDesign as the only \"loss\" function for design might be a bad idea, you may find adversarial sequences (aka. sequences that trick AlphaFold)."
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title **setup**\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.shared.utils import copy_dict\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "\n",
+        "from IPython.display import HTML\n",
+        "from google.colab import files\n",
+        "import numpy as np\n",
+        "\n",
+        "#########################\n",
+        "def get_pdb(pdb_code=\"\"):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  elif os.path.isfile(pdb_code):\n",
+        "    return pdb_code\n",
+        "  elif len(pdb_code) == 4:\n",
+        "    os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "    return f\"{pdb_code}.pdb\"\n",
+        "  else:\n",
+        "    os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "    return f\"AF-{pdb_code}-F1-model_v3.pdb\""
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title **prep inputs**\n",
+        "import re\n",
+        "#@markdown ---\n",
+        "#@markdown **target info**\n",
+        "pdb = \"4N5T\" #@param {type:\"string\"}\n",
+        "#@markdown - enter PDB code or UniProt code (to fetch AlphaFoldDB model) or leave blink to upload your own\n",
+        "target_chain = \"A\" #@param {type:\"string\"}\n",
+        "target_hotspot = \"\" #@param {type:\"string\"}\n",
+        "if target_hotspot == \"\": target_hotspot = None\n",
+        "#@markdown - restrict loss to predefined positions on target (eg. \"1-10,12,15\")\n",
+        "target_flexible = False #@param {type:\"boolean\"}\n",
+        "#@markdown - allow backbone of target structure to be flexible\n",
+        "\n",
+        "#@markdown ---\n",
+        "#@markdown **binder info**\n",
+        "binder_len = 14 #@param {type:\"integer\"}\n",
+        "#@markdown - length of binder to hallucination\n",
+        "binder_seq = \"\" #@param {type:\"string\"}\n",
+        "binder_seq = re.sub(\"[^A-Z]\", \"\", binder_seq.upper())\n",
+        "if len(binder_seq) > 0:\n",
+        "  binder_len = len(binder_seq)\n",
+        "else:\n",
+        "  binder_seq = None\n",
+        "#@markdown - if defined, will initialize design with this sequence\n",
+        "\n",
+        "binder_chain = \"\" #@param {type:\"string\"}\n",
+        "if binder_chain == \"\": binder_chain = None\n",
+        "#@markdown - if defined, supervised loss is used (binder_len is ignored)\n",
+        "\n",
+        "#@markdown ---\n",
+        "#@markdown **model config**\n",
+        "use_multimer = False #@param {type:\"boolean\"}\n",
+        "#@markdown - use alphafold-multimer for design\n",
+        "num_recycles = 0 #@param [\"0\", \"1\", \"3\", \"6\"] {type:\"raw\"}\n",
+        "num_models = \"2\" #@param [\"1\", \"2\", \"3\", \"4\", \"5\", \"all\"]\n",
+        "num_models = 5 if num_models == \"all\" else int(num_models)\n",
+        "#@markdown - number of trained models to use during optimization\n",
+        "\n",
+        "\n",
+        "x = {\"pdb_filename\":pdb,\n",
+        "     \"chain\":target_chain,\n",
+        "     \"binder_len\":binder_len,\n",
+        "     \"binder_chain\":binder_chain,\n",
+        "     \"hotspot\":target_hotspot,\n",
+        "     \"use_multimer\":use_multimer,\n",
+        "     \"rm_target_seq\":target_flexible}\n",
+        "     \n",
+        "x[\"pdb_filename\"] = get_pdb(x[\"pdb_filename\"])     \n",
+        "\n",
+        "if \"x_prev\" not in dir() or x != x_prev:\n",
+        "  clear_mem()\n",
+        "  model = mk_afdesign_model(protocol=\"binder\",\n",
+        "                            use_multimer=x[\"use_multimer\"],\n",
+        "                            num_recycles=num_recycles,\n",
+        "                            recycle_mode=\"sample\")\n",
+        "  model.prep_inputs(**x,\n",
+        "                    ignore_missing=False)\n",
+        "  x_prev = copy_dict(x)\n",
+        "  print(\"target length:\", model._target_len)\n",
+        "  print(\"binder length:\", model._binder_len)\n",
+        "  binder_len = model._binder_len"
+      ],
+      "metadata": {
+        "id": "HSgE99WALOE-",
+        "cellView": "form"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@title **run AfDesign**\n",
+        "from scipy.special import softmax\n",
+        "\n",
+        "optimizer = \"pssm_semigreedy\" #@param [\"pssm_semigreedy\", \"3stage\", \"semigreedy\", \"pssm\", \"logits\", \"soft\", \"hard\"]\n",
+        "#@markdown - `pssm_semigreedy` - uses the designed PSSM to bias semigreedy opt. (Recommended)\n",
+        "#@markdown - `3stage` - gradient based optimization (GD) (logits → soft → hard)\n",
+        "#@markdown - `pssm` - GD optimize (logits → soft) to get a sequence profile (PSSM).\n",
+        "#@markdown - `semigreedy` - tries X random mutations, accepts those that decrease loss\n",
+        "#@markdown - `logits` - GD optimize logits inputs (continious)\n",
+        "#@markdown - `soft` - GD optimize softmax(logits) inputs (probabilities)\n",
+        "#@markdown - `hard` - GD optimize one_hot(logits) inputs (discrete)\n",
+        "\n",
+        "#@markdown WARNING: The output sequence from `pssm`,`logits`,`soft` is not one_hot. To get a valid sequence use the other optimizers, or redesign the output backbone with another protocol like ProteinMPNN.\n",
+        "\n",
+        "#@markdown ----\n",
+        "#@markdown #### advanced GD settings\n",
+        "GD_method = \"sgd\" #@param [\"adabelief\", \"adafactor\", \"adagrad\", \"adam\", \"adamw\", \"fromage\", \"lamb\", \"lars\", \"noisy_sgd\", \"dpsgd\", \"radam\", \"rmsprop\", \"sgd\", \"sm3\", \"yogi\"]\n",
+        "learning_rate = 0.1 #@param {type:\"raw\"}\n",
+        "norm_seq_grad = True #@param {type:\"boolean\"}\n",
+        "dropout = True #@param {type:\"boolean\"}\n",
+        "\n",
+        "model.restart(seq=binder_seq)\n",
+        "model.set_optimizer(optimizer=GD_method,\n",
+        "                    learning_rate=learning_rate,\n",
+        "                    norm_seq_grad=norm_seq_grad)\n",
+        "models = model._model_names[:num_models]\n",
+        "\n",
+        "flags = {\"num_recycles\":num_recycles,\n",
+        "         \"models\":models,\n",
+        "         \"dropout\":dropout}\n",
+        "\n",
+        "if optimizer == \"3stage\":\n",
+        "  model.design_3stage(120, 60, 10, **flags)\n",
+        "  pssm = softmax(model._tmp[\"seq_logits\"],-1)\n",
+        "\n",
+        "if optimizer == \"pssm_semigreedy\":\n",
+        "  model.design_pssm_semigreedy(120, 32, **flags)\n",
+        "  pssm = softmax(model._tmp[\"seq_logits\"],1)\n",
+        "\n",
+        "if optimizer == \"semigreedy\":\n",
+        "  model.design_pssm_semigreedy(0, 32, **flags)\n",
+        "  pssm = None\n",
+        "\n",
+        "if optimizer == \"pssm\":\n",
+        "  model.design_logits(120, e_soft=1.0, num_models=1, ramp_recycles=True, **flags)\n",
+        "  model.design_soft(32, num_models=1, **flags)\n",
+        "  flags.update({\"dropout\":False,\"save_best\":True})\n",
+        "  model.design_soft(10, num_models=num_models, **flags)\n",
+        "  pssm = softmax(model.aux[\"seq\"][\"logits\"],-1)\n",
+        "\n",
+        "O = {\"logits\":model.design_logits,\n",
+        "     \"soft\":model.design_soft,\n",
+        "     \"hard\":model.design_hard}\n",
+        "\n",
+        "if optimizer in O:\n",
+        "  O[optimizer](120, num_models=1, ramp_recycles=True, **flags)\n",
+        "  flags.update({\"dropout\":False,\"save_best\":True})\n",
+        "  O[optimizer](10, num_models=num_models, **flags)\n",
+        "  pssm = softmax(model.aux[\"seq\"][\"logits\"],-1)\n",
+        "\n",
+        "model.save_pdb(f\"{model.protocol}.pdb\")"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "60qmxpzno0yV"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "A1GxeLZdTTya",
+        "cellView": "form"
+      },
+      "outputs": [],
+      "source": [
+        "#@title display hallucinated protein {run: \"auto\"}\n",
+        "color = \"pLDDT\" #@param [\"chain\", \"pLDDT\", \"rainbow\"]\n",
+        "show_sidechains = False #@param {type:\"boolean\"}\n",
+        "show_mainchains = False #@param {type:\"boolean\"}\n",
+        "color_HP = False #@param {type:\"boolean\"}\n",
+        "animate = True #@param {type:\"boolean\"}\n",
+        "model.plot_pdb(show_sidechains=show_sidechains,\n",
+        "               show_mainchains=show_mainchains,\n",
+        "               color=color, color_HP=color_HP, animate=animate)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "L2E9Tn2Acchj"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(model.animate(dpi=100))"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YSKWYu0_GlUH"
+      },
+      "outputs": [],
+      "source": [
+        "model.save_pdb(f\"{model.protocol}.pdb\")\n",
+        "model.get_seqs()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "#@markdown ### Amino acid probabilties\n",
+        "import plotly.express as px\n",
+        "alphabet = \"ACDEFGHIKLMNPQRSTVWY\"\n",
+        "if \"pssm\" in dir() and pssm is not None:\n",
+        "  fig = px.imshow(pssm.mean(0).T,\n",
+        "                  labels=dict(x=\"positions\", y=\"amino acids\", color=\"probability\"),\n",
+        "                  y=residue_constants.restypes,\n",
+        "                  zmin=0,\n",
+        "                  zmax=1,\n",
+        "                  template=\"simple_white\",\n",
+        "                )\n",
+        "  fig.update_xaxes(side=\"top\")\n",
+        "  fig.show()"
+      ],
+      "metadata": {
+        "cellView": "form",
+        "id": "3pvptBM55GHU"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# log\n",
+        "model._tmp[\"best\"][\"aux\"][\"log\"]"
+      ],
+      "metadata": {
+        "id": "1SGmdJKLNKvb"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [],
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af/examples/use_esm_1b_bias.ipynb

@@ -0,0 +1,329 @@
+{
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/v1.1.1/af/examples/use_esm_1b_bias.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "# Gather inputs"
+      ],
+      "metadata": {
+        "id": "M9xamJDvx1pd"
+      }
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "import requests, re\n",
+        "from google.colab import files\n",
+        "\n",
+        "def get_uniprot_seq(uid):\n",
+        "  url = f'https://rest.uniprot.org/uniprotkb/stream?compressed=false&format=fasta&query={uid}'\n",
+        "  sequence = \"\".join(re.split(r'\\n(?=>)', requests.get(url).text)[0].split(\"\\n\")[1:])\n",
+        "  return sequence\n",
+        "\n",
+        "def get_pdb(pdb_code=\"\",alphafold_model=False):\n",
+        "  if pdb_code is None or pdb_code == \"\":\n",
+        "    upload_dict = files.upload()\n",
+        "    pdb_string = upload_dict[list(upload_dict.keys())[0]]\n",
+        "    with open(\"tmp.pdb\",\"wb\") as out: out.write(pdb_string)\n",
+        "    return \"tmp.pdb\"\n",
+        "  else:\n",
+        "    if alphafold_model:\n",
+        "      os.system(f\"wget -qnc https://alphafold.ebi.ac.uk/files/AF-{pdb_code}-F1-model_v3.pdb\")\n",
+        "      return f\"AF-{pdb_code}-F1-model_v3.pdb\"      \n",
+        "    else:\n",
+        "      os.system(f\"wget -qnc https://files.rcsb.org/view/{pdb_code}.pdb\")\n",
+        "      return f\"{pdb_code}.pdb\""
+      ],
+      "metadata": {
+        "id": "lev6Nc_-vx2D"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "UNIPROT = \"P0A6A8\"\n",
+        "SEQUENCE = get_uniprot_seq(UNIPROT)\n",
+        "PDB_FILENAME = get_pdb(UNIPROT, alphafold_model=True)"
+      ],
+      "metadata": {
+        "id": "VkPdQj-Fv1LQ"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "H5LVw2BwijkB"
+      },
+      "source": [
+        "#ESM_1b\n",
+        "use logits from ESM_1b as prior to AfDesign"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "8j12AU-SikvN"
+      },
+      "outputs": [],
+      "source": [
+        "!pip -q install fair-esm\n",
+        "import esm\n",
+        "import gc\n",
+        "import torch\n",
+        "import numpy as np"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "63LNj3dHilRW"
+      },
+      "outputs": [],
+      "source": [
+        "model, alphabet = esm.pretrained.esm1b_t33_650M_UR50S()\n",
+        "\n",
+        "# run model on GPU if available\n",
+        "device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')\n",
+        "model = model.to(device)\n",
+        "model = model.eval()\n",
+        "model.args.token_dropout = False"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "ql8K460bjQtv"
+      },
+      "outputs": [],
+      "source": [
+        "def get_bias_from_esm(seq, p=None):\n",
+        "  '''p=None; number of calculation done in parallel (increase if you have more gpu-memory)'''\n",
+        "\n",
+        "  # map esm-alphabet to standard-alphabet\n",
+        "  tmp_a2n = {a:n for n,a in enumerate(alphabet.all_toks[4:24])}\n",
+        "  tmp_aa_map = np.array([tmp_a2n[a] for a in \"ARNDCQEGHILKMFPSTWYV\"])\n",
+        "\n",
+        "  x,ln = alphabet.get_batch_converter()([(None,seq)])[-1],len(seq)\n",
+        "  if p is None: p = ln\n",
+        "  with torch.no_grad():\n",
+        "    f = lambda x: model(x)[\"logits\"][:,1:(ln+1),4:24]\n",
+        "    logits = np.zeros((ln,20))\n",
+        "    for n in range(0,ln,p):\n",
+        "      m = min(n+p,ln)\n",
+        "      x_h = torch.tile(torch.clone(x),[m-n,1])\n",
+        "      for i in range(m-n):\n",
+        "        x_h[i,n+i+1] = alphabet.mask_idx\n",
+        "      fx_h = f(x_h.to(device))\n",
+        "      for i in range(m-n):\n",
+        "        logits[n+i] = fx_h[i,n+i].cpu().numpy()\n",
+        "  \n",
+        "    return logits[:,tmp_aa_map]"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "mrD1PRRMkZt3"
+      },
+      "outputs": [],
+      "source": [
+        "# get bias\n",
+        "seq = SEQUENCE\n",
+        "bias = get_bias_from_esm(seq)\n",
+        "np.savetxt(\"bias.txt\",bias)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "K_8S9rXOkwZ8"
+      },
+      "outputs": [],
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "plt.imshow(bias.T,cmap=\"bwr_r\",vmin=-10,vmax=10)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [
+        "# clear GPU memory\n",
+        "del model\n",
+        "gc.collect()\n",
+        "torch.cuda.empty_cache()"
+      ],
+      "metadata": {
+        "id": "10_kZfNQJpsz"
+      },
+      "execution_count": null,
+      "outputs": []
+    },
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "OA2k3sAYuiXe"
+      },
+      "source": [
+        "#AfDesign \n"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "cellView": "form",
+        "id": "-AXy0s_4cKaK"
+      },
+      "outputs": [],
+      "source": [
+        "#@title setup afdesign\n",
+        "%%time\n",
+        "import os\n",
+        "if not os.path.isdir(\"params\"):\n",
+        "  # get code\n",
+        "  os.system(\"pip -q install git+https://github.com/sokrypton/ColabDesign.git@v1.1.1\")\n",
+        "  # for debugging\n",
+        "  os.system(\"ln -s /usr/local/lib/python3.*/dist-packages/colabdesign colabdesign\")\n",
+        "  # download params\n",
+        "  os.system(\"mkdir params\")\n",
+        "  os.system(\"apt-get install aria2 -qq\")\n",
+        "  os.system(\"aria2c -q -x 16 https://storage.googleapis.com/alphafold/alphafold_params_2022-12-06.tar\")\n",
+        "  os.system(\"tar -xf alphafold_params_2022-12-06.tar -C params\")\n",
+        "\n",
+        "import warnings\n",
+        "warnings.simplefilter(action='ignore', category=FutureWarning)\n",
+        "\n",
+        "import os\n",
+        "from colabdesign import mk_afdesign_model, clear_mem\n",
+        "from colabdesign.af.alphafold.common import residue_constants\n",
+        "from IPython.display import HTML\n",
+        "import numpy as np\n",
+        "\n"
+      ]
+
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "qLd1DsnKzxBJ"
+      },
+      "outputs": [],
+      "source": [
+        "clear_mem()\n",
+        "model = mk_afdesign_model(protocol=\"fixbb\",\n",
+        "                          use_templates=False) # set True to constrain structure\n",
+        "model.prep_inputs(PDB_FILENAME, chain=\"A\")\n",
+        "print(\"length\",  model._len)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "WZGyGfMTmdXE"
+      },
+      "outputs": [],
+      "source": [
+        "import matplotlib.pyplot as plt\n",
+        "bias = np.loadtxt(\"bias.txt\")\n",
+        "plt.imshow(bias.T,cmap=\"bwr_r\",vmin=-10,vmax=10)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "xeVuwTYWmQmp"
+      },
+      "outputs": [],
+      "source": [
+        "model.restart()\n",
+        "model.set_seq(bias=bias)\n",
+        "model.design_3stage(50,50,10)"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "cW1KQiHKJpfp"
+      },
+      "outputs": [],
+      "source": [
+        "HTML(model.animate())"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YDrChASGVUUx"
+      },
+      "outputs": [],
+      "source": [
+        "model.get_seqs()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "execution_count": null,
+      "metadata": {
+        "id": "YEApO8YzBoS0"
+      },
+      "outputs": [],
+      "source": [
+        "model.save_pdb(f\"{model.protocol}.pdb\")\n",
+        "model.plot_pdb()"
+      ]
+    },
+    {
+      "cell_type": "code",
+      "source": [],
+      "metadata": {
+        "id": "0Xd7zQ-PqTly"
+      },
+      "execution_count": null,
+      "outputs": []
+    }
+  ],
+  "metadata": {
+    "accelerator": "GPU",
+    "colab": {
+      "collapsed_sections": [
+        "q4qiU9I0QHSz"
+      ],
+      "name": "use_esm_1b_bias.ipynb",
+      "provenance": [],
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "display_name": "Python 3",
+      "name": "python3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "nbformat": 4,
+  "nbformat_minor": 0
+}

model/PXDesignBench/ColabDesign/af_design_tutorial.ipynb

@@ -0,0 +1,43 @@
+{
+  "nbformat": 4,
+  "nbformat_minor": 0,
+  "metadata": {
+    "colab": {
+      "provenance": [],
+      "authorship_tag": "ABX9TyM3Mazv53Et9F6cqnRak08f",
+      "include_colab_link": true
+    },
+    "kernelspec": {
+      "name": "python3",
+      "display_name": "Python 3"
+    },
+    "language_info": {
+      "name": "python"
+    }
+  },
+  "cells": [
+    {
+      "cell_type": "markdown",
+      "metadata": {
+        "id": "view-in-github",
+        "colab_type": "text"
+      },
+      "source": [
+        "<a href=\"https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/design.ipynb\" target=\"_parent\"><img src=\"https://colab.research.google.com/assets/colab-badge.svg\" alt=\"Open In Colab\"/></a>"
+      ]
+    },
+    {
+      "cell_type": "markdown",
+      "source": [
+        "The general tutorial has moved to:\n",
+        "https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/design.ipynb\n",
+        "\n",
+        "For peptide design tutorial see:\n",
+        "https://colab.research.google.com/github/sokrypton/ColabDesign/blob/main/af/examples/peptide_binder_design.ipynb"
+      ],
+      "metadata": {
+        "id": "IabCAgHR6Xa3"
+      }
+    }
+  ]
+}

model/PXDesignBench/ColabDesign/colabdesign/__init__.py

@@ -0,0 +1,16 @@
+import os,jax
+# disable triton_gemm for jax versions > 0.3
+if int(jax.__version__.split(".")[1]) > 3:
+  os.environ["XLA_FLAGS"] = "--xla_gpu_enable_triton_gemm=false"
+
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
+from colabdesign.shared.utils import clear_mem
+from colabdesign.af.model import mk_af_model
+from colabdesign.tr.model import mk_tr_model
+from colabdesign.mpnn.model import mk_mpnn_model
+
+# backward compatability
+mk_design_model = mk_afdesign_model = mk_af_model
+mk_trdesign_model = mk_tr_model

model/PXDesignBench/ColabDesign/colabdesign/af/README.md

@@ -0,0 +1,6 @@
+- model.py  - configure the model
+- inputs.py - configure the inputs
+- loss.py   - configure the loss
+- prep.py   - prep features
+- design.py - gradient update loop
+- utils.py  - various tools for saving/plotting

model/PXDesignBench/ColabDesign/colabdesign/af/__init__.py

@@ -0,0 +1,13 @@
+import os,jax
+# disable triton_gemm for jax versions > 0.3
+if int(jax.__version__.split(".")[1]) > 3:
+  os.environ["XLA_FLAGS"] = "--xla_gpu_enable_triton_gemm=false"
+
+import warnings
+warnings.simplefilter(action='ignore', category=FutureWarning)
+
+from colabdesign.shared.utils import clear_mem
+from colabdesign.af.model import mk_af_model
+
+# backward compatability
+mk_design_model = mk_afdesign_model = mk_af_model

model/PXDesignBench/ColabDesign/colabdesign/af/alphafold/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""An implementation of the inference pipeline of AlphaFold v2.0."""

model/PXDesignBench/ColabDesign/colabdesign/af/alphafold/common/__init__.py

@@ -0,0 +1,14 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Common data types and constants used within Alphafold."""

model/PXDesignBench/ColabDesign/colabdesign/af/alphafold/common/confidence.py

@@ -0,0 +1,169 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Functions for processing confidence metrics."""
+
+import jax.numpy as jnp
+import jax
+import numpy as np
+from colabdesign.af.alphafold.common import residue_constants
+import scipy.special
+
+def compute_tol(prev_pos, current_pos, mask, use_jnp=False):
+    # Early stopping criteria based on criteria used in
+    # AF2Complex: https://www.nature.com/articles/s41467-022-29394-2    
+    _np = jnp if use_jnp else np
+    dist = lambda x:_np.sqrt(((x[:,None] - x[None,:])**2).sum(-1))
+    ca_idx = residue_constants.atom_order['CA']
+    sq_diff = _np.square(dist(prev_pos[:,ca_idx])-dist(current_pos[:,ca_idx]))
+    mask_2d = mask[:,None] * mask[None,:]
+    return _np.sqrt((sq_diff * mask_2d).sum()/mask_2d.sum() + 1e-8)
+
+
+def compute_plddt(logits, use_jnp=False):
+  """Computes per-residue pLDDT from logits.
+  Args:
+    logits: [num_res, num_bins] output from the PredictedLDDTHead.
+  Returns:
+    plddt: [num_res] per-residue pLDDT.
+  """
+  if use_jnp:
+    _np, _softmax = jnp, jax.nn.softmax
+  else:
+    _np, _softmax = np, scipy.special.softmax
+  
+  num_bins = logits.shape[-1]
+  bin_width = 1.0 / num_bins
+  bin_centers = _np.arange(start=0.5 * bin_width, stop=1.0, step=bin_width)
+  probs = _softmax(logits, axis=-1)
+  predicted_lddt_ca = (probs * bin_centers[None, :]).sum(-1)
+  return predicted_lddt_ca * 100
+
+def _calculate_bin_centers(breaks, use_jnp=False):
+  """Gets the bin centers from the bin edges.
+  Args:
+    breaks: [num_bins - 1] the error bin edges.
+  Returns:
+    bin_centers: [num_bins] the error bin centers.
+  """
+  _np = jnp if use_jnp else np
+  step = breaks[1] - breaks[0]
+
+  # Add half-step to get the center
+  bin_centers = breaks + step / 2
+
+  # Add a catch-all bin at the end.
+  return _np.append(bin_centers, bin_centers[-1] + step)
+
+def _calculate_expected_aligned_error(
+  alignment_confidence_breaks,
+  aligned_distance_error_probs,
+  use_jnp=False):
+  """Calculates expected aligned distance errors for every pair of residues.
+  Args:
+    alignment_confidence_breaks: [num_bins - 1] the error bin edges.
+    aligned_distance_error_probs: [num_res, num_res, num_bins] the predicted
+      probs for each error bin, for each pair of residues.
+  Returns:
+    predicted_aligned_error: [num_res, num_res] the expected aligned distance
+      error for each pair of residues.
+    max_predicted_aligned_error: The maximum predicted error possible.
+  """
+  bin_centers = _calculate_bin_centers(alignment_confidence_breaks, use_jnp=use_jnp)
+  # Tuple of expected aligned distance error and max possible error.
+  pae = (aligned_distance_error_probs * bin_centers).sum(-1)
+  return (pae, bin_centers[-1])
+
+def compute_predicted_aligned_error(logits, breaks, use_jnp=False):
+  """Computes aligned confidence metrics from logits.
+  Args:
+    logits: [num_res, num_res, num_bins] the logits output from
+      PredictedAlignedErrorHead.
+    breaks: [num_bins - 1] the error bin edges.
+
+  Returns:
+    aligned_confidence_probs: [num_res, num_res, num_bins] the predicted
+      aligned error probabilities over bins for each residue pair.
+    predicted_aligned_error: [num_res, num_res] the expected aligned distance
+      error for each pair of residues.
+    max_predicted_aligned_error: The maximum predicted error possible.
+  """
+  _softmax = jax.nn.softmax if use_jnp else scipy.special.softmax
+  aligned_confidence_probs = _softmax(logits,axis=-1)
+  predicted_aligned_error, max_predicted_aligned_error = \
+  _calculate_expected_aligned_error(breaks, aligned_confidence_probs, use_jnp=use_jnp)
+
+  return {
+      'aligned_confidence_probs': aligned_confidence_probs,
+      'predicted_aligned_error': predicted_aligned_error,
+      'max_predicted_aligned_error': max_predicted_aligned_error,
+  }
+
+def predicted_tm_score(logits, breaks, residue_weights = None,
+    asym_id = None, use_jnp=False):
+  """Computes predicted TM alignment or predicted interface TM alignment score.
+
+  Args:
+    logits: [num_res, num_res, num_bins] the logits output from
+      PredictedAlignedErrorHead.
+    breaks: [num_bins] the error bins.
+    residue_weights: [num_res] the per residue weights to use for the
+      expectation.
+    asym_id: [num_res] the asymmetric unit ID - the chain ID. Only needed for
+      ipTM calculation.
+
+  Returns:
+    ptm_score: The predicted TM alignment or the predicted iTM score.
+  """
+  if use_jnp:
+    _np, _softmax = jnp, jax.nn.softmax
+  else:
+    _np, _softmax = np, scipy.special.softmax
+
+  # residue_weights has to be in [0, 1], but can be floating-point, i.e. the
+  # exp. resolved head's probability.
+  if residue_weights is None:
+    residue_weights = _np.ones(logits.shape[0])
+
+  bin_centers = _calculate_bin_centers(breaks, use_jnp=use_jnp)
+  num_res = residue_weights.shape[0]
+
+  # Clip num_res to avoid negative/undefined d0.
+  clipped_num_res = _np.maximum(residue_weights.sum(), 19)
+
+  # Compute d_0(num_res) as defined by TM-score, eqn. (5) in Yang & Skolnick
+  # "Scoring function for automated assessment of protein structure template
+  # quality", 2004: http://zhanglab.ccmb.med.umich.edu/papers/2004_3.pdf
+  d0 = 1.24 * (clipped_num_res - 15) ** (1./3) - 1.8
+
+  # Convert logits to probs.
+  probs = _softmax(logits, axis=-1)
+
+  # TM-Score term for every bin.
+  tm_per_bin = 1. / (1 + _np.square(bin_centers) / _np.square(d0))
+  # E_distances tm(distance).
+  predicted_tm_term = (probs * tm_per_bin).sum(-1)
+
+  if asym_id is None:
+    pair_mask = _np.full((num_res,num_res),True)
+  else:
+    pair_mask = asym_id[:, None] != asym_id[None, :]
+
+  predicted_tm_term *= pair_mask
+
+  pair_residue_weights = pair_mask * (residue_weights[None, :] * residue_weights[:, None])
+  normed_residue_mask = pair_residue_weights / (1e-8 + pair_residue_weights.sum(-1, keepdims=True))
+  per_alignment = (predicted_tm_term * normed_residue_mask).sum(-1)
+
+  return (per_alignment * residue_weights).max()

model/PXDesignBench/ColabDesign/colabdesign/af/alphafold/common/protein.py

@@ -0,0 +1,229 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Protein data type."""
+import dataclasses
+import io
+from typing import Any, Mapping, Optional
+from colabdesign.af.alphafold.common import residue_constants
+from Bio.PDB import PDBParser
+import numpy as np
+
+FeatureDict = Mapping[str, np.ndarray]
+ModelOutput = Mapping[str, Any]  # Is a nested dict.
+
+
+@dataclasses.dataclass(frozen=True)
+class Protein:
+  """Protein structure representation."""
+
+  # Cartesian coordinates of atoms in angstroms. The atom types correspond to
+  # residue_constants.atom_types, i.e. the first three are N, CA, CB.
+  atom_positions: np.ndarray  # [num_res, num_atom_type, 3]
+
+  # Amino-acid type for each residue represented as an integer between 0 and
+  # 20, where 20 is 'X'.
+  aatype: np.ndarray  # [num_res]
+
+  # Binary float mask to indicate presence of a particular atom. 1.0 if an atom
+  # is present and 0.0 if not. This should be used for loss masking.
+  atom_mask: np.ndarray  # [num_res, num_atom_type]
+
+  # Residue index as used in PDB. It is not necessarily continuous or 0-indexed.
+  residue_index: np.ndarray  # [num_res]
+
+  # B-factors, or temperature factors, of each residue (in sq. angstroms units),
+  # representing the displacement of the residue from its ground truth mean
+  # value.
+  b_factors: np.ndarray  # [num_res, num_atom_type]
+
+
+def from_pdb_string(pdb_str: str, chain_id: Optional[str] = None) -> Protein:
+  """Takes a PDB string and constructs a Protein object.
+
+  WARNING: All non-standard residue types will be converted into UNK. All
+    non-standard atoms will be ignored.
+
+  Args:
+    pdb_str: The contents of the pdb file
+    chain_id: If None, then the pdb file must contain a single chain (which
+      will be parsed). If chain_id is specified (e.g. A), then only that chain
+      is parsed.
+
+  Returns:
+    A new `Protein` parsed from the pdb contents.
+  """
+  pdb_fh = io.StringIO(pdb_str)
+  parser = PDBParser(QUIET=True)
+  structure = parser.get_structure('none', pdb_fh)
+  models = list(structure.get_models())
+  if len(models) != 1:
+    raise ValueError(
+        f'Only single model PDBs are supported. Found {len(models)} models.')
+  model = models[0]
+
+  if chain_id is not None:
+    chain = model[chain_id]
+  else:
+    chains = list(model.get_chains())
+    if len(chains) != 1:
+      raise ValueError(
+          'Only single chain PDBs are supported when chain_id not specified. '
+          f'Found {len(chains)} chains.')
+    else:
+      chain = chains[0]
+
+  atom_positions = []
+  aatype = []
+  atom_mask = []
+  residue_index = []
+  b_factors = []
+
+  for res in chain:
+    if res.id[2] != ' ':
+      raise ValueError(
+          f'PDB contains an insertion code at chain {chain.id} and residue '
+          f'index {res.id[1]}. These are not supported.')
+    res_shortname = residue_constants.restype_3to1.get(res.resname, 'X')
+    restype_idx = residue_constants.restype_order.get(
+        res_shortname, residue_constants.restype_num)
+    pos = np.zeros((residue_constants.atom_type_num, 3))
+    mask = np.zeros((residue_constants.atom_type_num,))
+    res_b_factors = np.zeros((residue_constants.atom_type_num,))
+    for atom in res:
+      if atom.name not in residue_constants.atom_types:
+        continue
+      pos[residue_constants.atom_order[atom.name]] = atom.coord
+      mask[residue_constants.atom_order[atom.name]] = 1.
+      res_b_factors[residue_constants.atom_order[atom.name]] = atom.bfactor
+    if np.sum(mask) < 0.5:
+      # If no known atom positions are reported for the residue then skip it.
+      continue
+    aatype.append(restype_idx)
+    atom_positions.append(pos)
+    atom_mask.append(mask)
+    residue_index.append(res.id[1])
+    b_factors.append(res_b_factors)
+
+  return Protein(
+      atom_positions=np.array(atom_positions),
+      atom_mask=np.array(atom_mask),
+      aatype=np.array(aatype),
+      residue_index=np.array(residue_index),
+      b_factors=np.array(b_factors))
+
+
+def to_pdb(prot: Protein) -> str:
+  """Converts a `Protein` instance to a PDB string.
+
+  Args:
+    prot: The protein to convert to PDB.
+
+  Returns:
+    PDB string.
+  """
+  restypes = residue_constants.restypes + ['X']
+  res_1to3 = lambda r: residue_constants.restype_1to3.get(restypes[r], 'UNK')
+  atom_types = residue_constants.atom_types
+
+  pdb_lines = []
+
+  atom_mask = prot.atom_mask
+  aatype = prot.aatype
+  atom_positions = prot.atom_positions
+  residue_index = prot.residue_index.astype(np.int32)
+  b_factors = prot.b_factors
+
+  if np.any(aatype > residue_constants.restype_num):
+    raise ValueError('Invalid aatypes.')
+
+  pdb_lines.append('MODEL     1')
+  atom_index = 1
+  chain_id = 'A'
+  # Add all atom sites.
+  for i in range(aatype.shape[0]):
+    res_name_3 = res_1to3(aatype[i])
+    for atom_name, pos, mask, b_factor in zip(
+        atom_types, atom_positions[i], atom_mask[i], b_factors[i]):
+      if mask < 0.5:
+        continue
+
+      record_type = 'ATOM'
+      name = atom_name if len(atom_name) == 4 else f' {atom_name}'
+      alt_loc = ''
+      insertion_code = ''
+      occupancy = 1.00
+      element = atom_name[0]  # Protein supports only C, N, O, S, this works.
+      charge = ''
+      # PDB is a columnar format, every space matters here!
+      atom_line = (f'{record_type:<6}{atom_index:>5} {name:<4}{alt_loc:>1}'
+                   f'{res_name_3:>3} {chain_id:>1}'
+                   f'{residue_index[i]:>4}{insertion_code:>1}   '
+                   f'{pos[0]:>8.3f}{pos[1]:>8.3f}{pos[2]:>8.3f}'
+                   f'{occupancy:>6.2f}{b_factor:>6.2f}          '
+                   f'{element:>2}{charge:>2}')
+      pdb_lines.append(atom_line)
+      atom_index += 1
+
+  # Close the chain.
+  chain_end = 'TER'
+  chain_termination_line = (
+      f'{chain_end:<6}{atom_index:>5}      {res_1to3(aatype[-1]):>3} '
+      f'{chain_id:>1}{residue_index[-1]:>4}')
+  pdb_lines.append(chain_termination_line)
+  pdb_lines.append('ENDMDL')
+
+  pdb_lines.append('END')
+  pdb_lines.append('')
+  return '\n'.join(pdb_lines)
+
+
+def ideal_atom_mask(prot: Protein) -> np.ndarray:
+  """Computes an ideal atom mask.
+
+  `Protein.atom_mask` typically is defined according to the atoms that are
+  reported in the PDB. This function computes a mask according to heavy atoms
+  that should be present in the given sequence of amino acids.
+
+  Args:
+    prot: `Protein` whose fields are `numpy.ndarray` objects.
+
+  Returns:
+    An ideal atom mask.
+  """
+  return residue_constants.STANDARD_ATOM_MASK[prot.aatype]
+
+
+def from_prediction(features: FeatureDict, result: ModelOutput,
+                    b_factors: Optional[np.ndarray] = None) -> Protein:
+  """Assembles a protein from a prediction.
+
+  Args:
+    features: Dictionary holding model inputs.
+    result: Dictionary holding model outputs.
+    b_factors: (Optional) B-factors to use for the protein.
+
+  Returns:
+    A protein instance.
+  """
+  fold_output = result['structure_module']
+  if b_factors is None:
+    b_factors = np.zeros_like(fold_output['final_atom_mask'])
+
+  return Protein(
+      aatype=features['aatype'][0],
+      atom_positions=fold_output['final_atom_positions'],
+      atom_mask=fold_output['final_atom_mask'],
+      residue_index=features['residue_index'][0] + 1,
+      b_factors=b_factors)

model/PXDesignBench/ColabDesign/colabdesign/af/alphafold/common/residue_constants.py

@@ -0,0 +1,911 @@
+# Copyright 2021 DeepMind Technologies Limited
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#      http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Constants used in AlphaFold."""
+
+import collections
+import functools
+from typing import List, Mapping, Tuple
+
+import numpy as np
+import tree
+
+# Internal import (35fd).
+
+
+# Distance from one CA to next CA [trans configuration: omega = 180].
+ca_ca = 3.80209737096
+
+# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
+# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
+# chi angles so their chi angle lists are empty.
+chi_angles_atoms = {
+    'ALA': [],
+    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
+    'ARG': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'NE'], ['CG', 'CD', 'NE', 'CZ']],
+    'ASN': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'OD1']],
+    'ASP': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'OD1']],
+    'CYS': [['N', 'CA', 'CB', 'SG']],
+    'GLN': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'OE1']],
+    'GLU': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'OE1']],
+    'GLY': [],
+    'HIS': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'ND1']],
+    'ILE': [['N', 'CA', 'CB', 'CG1'], ['CA', 'CB', 'CG1', 'CD1']],
+    'LEU': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'LYS': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD'],
+            ['CB', 'CG', 'CD', 'CE'], ['CG', 'CD', 'CE', 'NZ']],
+    'MET': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'SD'],
+            ['CB', 'CG', 'SD', 'CE']],
+    'PHE': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'PRO': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD']],
+    'SER': [['N', 'CA', 'CB', 'OG']],
+    'THR': [['N', 'CA', 'CB', 'OG1']],
+    'TRP': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'TYR': [['N', 'CA', 'CB', 'CG'], ['CA', 'CB', 'CG', 'CD1']],
+    'VAL': [['N', 'CA', 'CB', 'CG1']],
+}
+
+# If chi angles given in fixed-length array, this matrix determines how to mask
+# them for each AA type. The order is as per restype_order (see below).
+chi_angles_mask = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [1.0, 1.0, 1.0, 1.0],  # ARG
+    [1.0, 1.0, 0.0, 0.0],  # ASN
+    [1.0, 1.0, 0.0, 0.0],  # ASP
+    [1.0, 0.0, 0.0, 0.0],  # CYS
+    [1.0, 1.0, 1.0, 0.0],  # GLN
+    [1.0, 1.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [1.0, 1.0, 0.0, 0.0],  # HIS
+    [1.0, 1.0, 0.0, 0.0],  # ILE
+    [1.0, 1.0, 0.0, 0.0],  # LEU
+    [1.0, 1.0, 1.0, 1.0],  # LYS
+    [1.0, 1.0, 1.0, 0.0],  # MET
+    [1.0, 1.0, 0.0, 0.0],  # PHE
+    [1.0, 1.0, 0.0, 0.0],  # PRO
+    [1.0, 0.0, 0.0, 0.0],  # SER
+    [1.0, 0.0, 0.0, 0.0],  # THR
+    [1.0, 1.0, 0.0, 0.0],  # TRP
+    [1.0, 1.0, 0.0, 0.0],  # TYR
+    [1.0, 0.0, 0.0, 0.0],  # VAL
+]
+
+# The following chi angles are pi periodic: they can be rotated by a multiple
+# of pi without affecting the structure.
+chi_pi_periodic = [
+    [0.0, 0.0, 0.0, 0.0],  # ALA
+    [0.0, 0.0, 0.0, 0.0],  # ARG
+    [0.0, 0.0, 0.0, 0.0],  # ASN
+    [0.0, 1.0, 0.0, 0.0],  # ASP
+    [0.0, 0.0, 0.0, 0.0],  # CYS
+    [0.0, 0.0, 0.0, 0.0],  # GLN
+    [0.0, 0.0, 1.0, 0.0],  # GLU
+    [0.0, 0.0, 0.0, 0.0],  # GLY
+    [0.0, 0.0, 0.0, 0.0],  # HIS
+    [0.0, 0.0, 0.0, 0.0],  # ILE
+    [0.0, 0.0, 0.0, 0.0],  # LEU
+    [0.0, 0.0, 0.0, 0.0],  # LYS
+    [0.0, 0.0, 0.0, 0.0],  # MET
+    [0.0, 1.0, 0.0, 0.0],  # PHE
+    [0.0, 0.0, 0.0, 0.0],  # PRO
+    [0.0, 0.0, 0.0, 0.0],  # SER
+    [0.0, 0.0, 0.0, 0.0],  # THR
+    [0.0, 0.0, 0.0, 0.0],  # TRP
+    [0.0, 1.0, 0.0, 0.0],  # TYR
+    [0.0, 0.0, 0.0, 0.0],  # VAL
+    [0.0, 0.0, 0.0, 0.0],  # UNK
+]
+
+# Atoms positions relative to the 8 rigid groups, defined by the pre-omega, phi,
+# psi and chi angles:
+# 0: 'backbone group',
+# 1: 'pre-omega-group', (empty)
+# 2: 'phi-group', (currently empty, because it defines only hydrogens)
+# 3: 'psi-group',
+# 4,5,6,7: 'chi1,2,3,4-group'
+# The atom positions are relative to the axis-end-atom of the corresponding
+# rotation axis. The x-axis is in direction of the rotation axis, and the y-axis
+# is defined such that the dihedral-angle-definiting atom (the last entry in
+# chi_angles_atoms above) is in the xy-plane (with a positive y-coordinate).
+# format: [atomname, group_idx, rel_position]
+rigid_group_atom_positions = {
+    'ALA': [
+        ['N', 0, (-0.525, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, -0.000, -0.000)],
+        ['CB', 0, (-0.529, -0.774, -1.205)],
+        ['O', 3, (0.627, 1.062, 0.000)],
+    ],
+    'ARG': [
+        ['N', 0, (-0.524, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, -0.000)],
+        ['CB', 0, (-0.524, -0.778, -1.209)],
+        ['O', 3, (0.626, 1.062, 0.000)],
+        ['CG', 4, (0.616, 1.390, -0.000)],
+        ['CD', 5, (0.564, 1.414, 0.000)],
+        ['NE', 6, (0.539, 1.357, -0.000)],
+        ['NH1', 7, (0.206, 2.301, 0.000)],
+        ['NH2', 7, (2.078, 0.978, -0.000)],
+        ['CZ', 7, (0.758, 1.093, -0.000)],
+    ],
+    'ASN': [
+        ['N', 0, (-0.536, 1.357, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, -0.000, -0.000)],
+        ['CB', 0, (-0.531, -0.787, -1.200)],
+        ['O', 3, (0.625, 1.062, 0.000)],
+        ['CG', 4, (0.584, 1.399, 0.000)],
+        ['ND2', 5, (0.593, -1.188, 0.001)],
+        ['OD1', 5, (0.633, 1.059, 0.000)],
+    ],
+    'ASP': [
+        ['N', 0, (-0.525, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, 0.000, -0.000)],
+        ['CB', 0, (-0.526, -0.778, -1.208)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.593, 1.398, -0.000)],
+        ['OD1', 5, (0.610, 1.091, 0.000)],
+        ['OD2', 5, (0.592, -1.101, -0.003)],
+    ],
+    'CYS': [
+        ['N', 0, (-0.522, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.524, 0.000, 0.000)],
+        ['CB', 0, (-0.519, -0.773, -1.212)],
+        ['O', 3, (0.625, 1.062, -0.000)],
+        ['SG', 4, (0.728, 1.653, 0.000)],
+    ],
+    'GLN': [
+        ['N', 0, (-0.526, 1.361, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, 0.000, 0.000)],
+        ['CB', 0, (-0.525, -0.779, -1.207)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.615, 1.393, 0.000)],
+        ['CD', 5, (0.587, 1.399, -0.000)],
+        ['NE2', 6, (0.593, -1.189, -0.001)],
+        ['OE1', 6, (0.634, 1.060, 0.000)],
+    ],
+    'GLU': [
+        ['N', 0, (-0.528, 1.361, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, -0.000, -0.000)],
+        ['CB', 0, (-0.526, -0.781, -1.207)],
+        ['O', 3, (0.626, 1.062, 0.000)],
+        ['CG', 4, (0.615, 1.392, 0.000)],
+        ['CD', 5, (0.600, 1.397, 0.000)],
+        ['OE1', 6, (0.607, 1.095, -0.000)],
+        ['OE2', 6, (0.589, -1.104, -0.001)],
+    ],
+    'GLY': [
+        ['N', 0, (-0.572, 1.337, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.517, -0.000, -0.000)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+    ],
+    'HIS': [
+        ['N', 0, (-0.527, 1.360, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, 0.000, 0.000)],
+        ['CB', 0, (-0.525, -0.778, -1.208)],
+        ['O', 3, (0.625, 1.063, 0.000)],
+        ['CG', 4, (0.600, 1.370, -0.000)],
+        ['CD2', 5, (0.889, -1.021, 0.003)],
+        ['ND1', 5, (0.744, 1.160, -0.000)],
+        ['CE1', 5, (2.030, 0.851, 0.002)],
+        ['NE2', 5, (2.145, -0.466, 0.004)],
+    ],
+    'ILE': [
+        ['N', 0, (-0.493, 1.373, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, -0.000, -0.000)],
+        ['CB', 0, (-0.536, -0.793, -1.213)],
+        ['O', 3, (0.627, 1.062, -0.000)],
+        ['CG1', 4, (0.534, 1.437, -0.000)],
+        ['CG2', 4, (0.540, -0.785, -1.199)],
+        ['CD1', 5, (0.619, 1.391, 0.000)],
+    ],
+    'LEU': [
+        ['N', 0, (-0.520, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, -0.000)],
+        ['CB', 0, (-0.522, -0.773, -1.214)],
+        ['O', 3, (0.625, 1.063, -0.000)],
+        ['CG', 4, (0.678, 1.371, 0.000)],
+        ['CD1', 5, (0.530, 1.430, -0.000)],
+        ['CD2', 5, (0.535, -0.774, 1.200)],
+    ],
+    'LYS': [
+        ['N', 0, (-0.526, 1.362, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, 0.000, 0.000)],
+        ['CB', 0, (-0.524, -0.778, -1.208)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.619, 1.390, 0.000)],
+        ['CD', 5, (0.559, 1.417, 0.000)],
+        ['CE', 6, (0.560, 1.416, 0.000)],
+        ['NZ', 7, (0.554, 1.387, 0.000)],
+    ],
+    'MET': [
+        ['N', 0, (-0.521, 1.364, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, 0.000, 0.000)],
+        ['CB', 0, (-0.523, -0.776, -1.210)],
+        ['O', 3, (0.625, 1.062, -0.000)],
+        ['CG', 4, (0.613, 1.391, -0.000)],
+        ['SD', 5, (0.703, 1.695, 0.000)],
+        ['CE', 6, (0.320, 1.786, -0.000)],
+    ],
+    'PHE': [
+        ['N', 0, (-0.518, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.524, 0.000, -0.000)],
+        ['CB', 0, (-0.525, -0.776, -1.212)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['CG', 4, (0.607, 1.377, 0.000)],
+        ['CD1', 5, (0.709, 1.195, -0.000)],
+        ['CD2', 5, (0.706, -1.196, 0.000)],
+        ['CE1', 5, (2.102, 1.198, -0.000)],
+        ['CE2', 5, (2.098, -1.201, -0.000)],
+        ['CZ', 5, (2.794, -0.003, -0.001)],
+    ],
+    'PRO': [
+        ['N', 0, (-0.566, 1.351, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, -0.000, 0.000)],
+        ['CB', 0, (-0.546, -0.611, -1.293)],
+        ['O', 3, (0.621, 1.066, 0.000)],
+        ['CG', 4, (0.382, 1.445, 0.0)],
+        # ['CD', 5, (0.427, 1.440, 0.0)],
+        ['CD', 5, (0.477, 1.424, 0.0)],  # manually made angle 2 degrees larger
+    ],
+    'SER': [
+        ['N', 0, (-0.529, 1.360, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, -0.000)],
+        ['CB', 0, (-0.518, -0.777, -1.211)],
+        ['O', 3, (0.626, 1.062, -0.000)],
+        ['OG', 4, (0.503, 1.325, 0.000)],
+    ],
+    'THR': [
+        ['N', 0, (-0.517, 1.364, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.526, 0.000, -0.000)],
+        ['CB', 0, (-0.516, -0.793, -1.215)],
+        ['O', 3, (0.626, 1.062, 0.000)],
+        ['CG2', 4, (0.550, -0.718, -1.228)],
+        ['OG1', 4, (0.472, 1.353, 0.000)],
+    ],
+    'TRP': [
+        ['N', 0, (-0.521, 1.363, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.525, -0.000, 0.000)],
+        ['CB', 0, (-0.523, -0.776, -1.212)],
+        ['O', 3, (0.627, 1.062, 0.000)],
+        ['CG', 4, (0.609, 1.370, -0.000)],
+        ['CD1', 5, (0.824, 1.091, 0.000)],
+        ['CD2', 5, (0.854, -1.148, -0.005)],
+        ['CE2', 5, (2.186, -0.678, -0.007)],
+        ['CE3', 5, (0.622, -2.530, -0.007)],
+        ['NE1', 5, (2.140, 0.690, -0.004)],
+        ['CH2', 5, (3.028, -2.890, -0.013)],
+        ['CZ2', 5, (3.283, -1.543, -0.011)],
+        ['CZ3', 5, (1.715, -3.389, -0.011)],
+    ],
+    'TYR': [
+        ['N', 0, (-0.522, 1.362, 0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.524, -0.000, -0.000)],
+        ['CB', 0, (-0.522, -0.776, -1.213)],
+        ['O', 3, (0.627, 1.062, -0.000)],
+        ['CG', 4, (0.607, 1.382, -0.000)],
+        ['CD1', 5, (0.716, 1.195, -0.000)],
+        ['CD2', 5, (0.713, -1.194, -0.001)],
+        ['CE1', 5, (2.107, 1.200, -0.002)],
+        ['CE2', 5, (2.104, -1.201, -0.003)],
+        ['OH', 5, (4.168, -0.002, -0.005)],
+        ['CZ', 5, (2.791, -0.001, -0.003)],
+    ],
+    'VAL': [
+        ['N', 0, (-0.494, 1.373, -0.000)],
+        ['CA', 0, (0.000, 0.000, 0.000)],
+        ['C', 0, (1.527, -0.000, -0.000)],
+        ['CB', 0, (-0.533, -0.795, -1.213)],
+        ['O', 3, (0.627, 1.062, -0.000)],
+        ['CG1', 4, (0.540, 1.429, -0.000)],
+        ['CG2', 4, (0.533, -0.776, 1.203)],
+    ],
+}
+
+# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
+residue_atoms = {
+    'ALA': ['C', 'CA', 'CB', 'N', 'O'],
+    'ARG': ['C', 'CA', 'CB', 'CG', 'CD', 'CZ', 'N', 'NE', 'O', 'NH1', 'NH2'],
+    'ASP': ['C', 'CA', 'CB', 'CG', 'N', 'O', 'OD1', 'OD2'],
+    'ASN': ['C', 'CA', 'CB', 'CG', 'N', 'ND2', 'O', 'OD1'],
+    'CYS': ['C', 'CA', 'CB', 'N', 'O', 'SG'],
+    'GLU': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'O', 'OE1', 'OE2'],
+    'GLN': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'NE2', 'O', 'OE1'],
+    'GLY': ['C', 'CA', 'N', 'O'],
+    'HIS': ['C', 'CA', 'CB', 'CG', 'CD2', 'CE1', 'N', 'ND1', 'NE2', 'O'],
+    'ILE': ['C', 'CA', 'CB', 'CG1', 'CG2', 'CD1', 'N', 'O'],
+    'LEU': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'N', 'O'],
+    'LYS': ['C', 'CA', 'CB', 'CG', 'CD', 'CE', 'N', 'NZ', 'O'],
+    'MET': ['C', 'CA', 'CB', 'CG', 'CE', 'N', 'O', 'SD'],
+    'PHE': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'N', 'O'],
+    'PRO': ['C', 'CA', 'CB', 'CG', 'CD', 'N', 'O'],
+    'SER': ['C', 'CA', 'CB', 'N', 'O', 'OG'],
+    'THR': ['C', 'CA', 'CB', 'CG2', 'N', 'O', 'OG1'],
+    'TRP': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE2', 'CE3', 'CZ2', 'CZ3',
+            'CH2', 'N', 'NE1', 'O'],
+    'TYR': ['C', 'CA', 'CB', 'CG', 'CD1', 'CD2', 'CE1', 'CE2', 'CZ', 'N', 'O',
+            'OH'],
+    'VAL': ['C', 'CA', 'CB', 'CG1', 'CG2', 'N', 'O']
+}
+
+# Naming swaps for ambiguous atom names.
+# Due to symmetries in the amino acids the naming of atoms is ambiguous in
+# 4 of the 20 amino acids.
+# (The LDDT paper lists 7 amino acids as ambiguous, but the naming ambiguities
+# in LEU, VAL and ARG can be resolved by using the 3d constellations of
+# the 'ambiguous' atoms and their neighbours)
+residue_atom_renaming_swaps = {
+    'ASP': {'OD1': 'OD2'},
+    'GLU': {'OE1': 'OE2'},
+    'PHE': {'CD1': 'CD2', 'CE1': 'CE2'},
+    'TYR': {'CD1': 'CD2', 'CE1': 'CE2'},
+}
+
+# Van der Waals radii [Angstroem] of the atoms (from Wikipedia)
+van_der_waals_radius = {
+    'C': 1.7,
+    'N': 1.55,
+    'O': 1.52,
+    'S': 1.8,
+}
+
+Bond = collections.namedtuple(
+    'Bond', ['atom1_name', 'atom2_name', 'length', 'stddev'])
+BondAngle = collections.namedtuple(
+    'BondAngle',
+    ['atom1_name', 'atom2_name', 'atom3name', 'angle_rad', 'stddev'])
+
+
+@functools.lru_cache(maxsize=None)
+def load_stereo_chemical_props() -> Tuple[Mapping[str, List[Bond]],
+                                          Mapping[str, List[Bond]],
+                                          Mapping[str, List[BondAngle]]]:
+  """Load stereo_chemical_props.txt into a nice structure.
+
+  Load literature values for bond lengths and bond angles and translate
+  bond angles into the length of the opposite edge of the triangle
+  ("residue_virtual_bonds").
+
+  Returns:
+    residue_bonds:  dict that maps resname --> list of Bond tuples
+    residue_virtual_bonds: dict that maps resname --> list of Bond tuples
+    residue_bond_angles: dict that maps resname --> list of BondAngle tuples
+  """
+  stereo_chemical_props_path = (
+      'alphafold/common/stereo_chemical_props.txt')
+  with open(stereo_chemical_props_path, 'rt') as f:
+    stereo_chemical_props = f.read()
+  lines_iter = iter(stereo_chemical_props.splitlines())
+  # Load bond lengths.
+  residue_bonds = {}
+  next(lines_iter)  # Skip header line.
+  for line in lines_iter:
+    if line.strip() == '-':
+      break
+    bond, resname, length, stddev = line.split()
+    atom1, atom2 = bond.split('-')
+    if resname not in residue_bonds:
+      residue_bonds[resname] = []
+    residue_bonds[resname].append(
+        Bond(atom1, atom2, float(length), float(stddev)))
+  residue_bonds['UNK'] = []
+
+  # Load bond angles.
+  residue_bond_angles = {}
+  next(lines_iter)  # Skip empty line.
+  next(lines_iter)  # Skip header line.
+  for line in lines_iter:
+    if line.strip() == '-':
+      break
+    bond, resname, angle_degree, stddev_degree = line.split()
+    atom1, atom2, atom3 = bond.split('-')
+    if resname not in residue_bond_angles:
+      residue_bond_angles[resname] = []
+    residue_bond_angles[resname].append(
+        BondAngle(atom1, atom2, atom3,
+                  float(angle_degree) / 180. * np.pi,
+                  float(stddev_degree) / 180. * np.pi))
+  residue_bond_angles['UNK'] = []
+
+  def make_bond_key(atom1_name, atom2_name):
+    """Unique key to lookup bonds."""
+    return '-'.join(sorted([atom1_name, atom2_name]))
+
+  # Translate bond angles into distances ("virtual bonds").
+  residue_virtual_bonds = {}
+  for resname, bond_angles in residue_bond_angles.items():
+    # Create a fast lookup dict for bond lengths.
+    bond_cache = {}
+    for b in residue_bonds[resname]:
+      bond_cache[make_bond_key(b.atom1_name, b.atom2_name)] = b
+    residue_virtual_bonds[resname] = []
+    for ba in bond_angles:
+      bond1 = bond_cache[make_bond_key(ba.atom1_name, ba.atom2_name)]
+      bond2 = bond_cache[make_bond_key(ba.atom2_name, ba.atom3name)]
+
+      # Compute distance between atom1 and atom3 using the law of cosines
+      # c^2 = a^2 + b^2 - 2ab*cos(gamma).
+      gamma = ba.angle_rad
+      length = np.sqrt(bond1.length**2 + bond2.length**2
+                       - 2 * bond1.length * bond2.length * np.cos(gamma))
+
+      # Propagation of uncertainty assuming uncorrelated errors.
+      dl_outer = 0.5 / length
+      dl_dgamma = (2 * bond1.length * bond2.length * np.sin(gamma)) * dl_outer
+      dl_db1 = (2 * bond1.length - 2 * bond2.length * np.cos(gamma)) * dl_outer
+      dl_db2 = (2 * bond2.length - 2 * bond1.length * np.cos(gamma)) * dl_outer
+      stddev = np.sqrt((dl_dgamma * ba.stddev)**2 +
+                       (dl_db1 * bond1.stddev)**2 +
+                       (dl_db2 * bond2.stddev)**2)
+      residue_virtual_bonds[resname].append(
+          Bond(ba.atom1_name, ba.atom3name, length, stddev))
+
+  return (residue_bonds,
+          residue_virtual_bonds,
+          residue_bond_angles)
+
+
+# Between-residue bond lengths for general bonds (first element) and for Proline
+# (second element).
+between_res_bond_length_c_n = [1.329, 1.341]
+between_res_bond_length_stddev_c_n = [0.014, 0.016]
+
+# Between-residue cos_angles.
+between_res_cos_angles_c_n_ca = [-0.5203, 0.0353]  # degrees: 121.352 +- 2.315
+between_res_cos_angles_ca_c_n = [-0.4473, 0.0311]  # degrees: 116.568 +- 1.995
+
+# This mapping is used when we need to store atom data in a format that requires
+# fixed atom data size for every residue (e.g. a numpy array).
+atom_types = [
+    'N', 'CA', 'C', 'CB', 'O', 'CG', 'CG1', 'CG2', 'OG', 'OG1', 'SG', 'CD',
+    'CD1', 'CD2', 'ND1', 'ND2', 'OD1', 'OD2', 'SD', 'CE', 'CE1', 'CE2', 'CE3',
+    'NE', 'NE1', 'NE2', 'OE1', 'OE2', 'CH2', 'NH1', 'NH2', 'OH', 'CZ', 'CZ2',
+    'CZ3', 'NZ', 'OXT'
+]
+atom_order = {atom_type: i for i, atom_type in enumerate(atom_types)}
+atom_type_num = len(atom_types)  # := 37.
+
+# A compact atom encoding with 14 columns
+# pylint: disable=line-too-long
+# pylint: disable=bad-whitespace
+restype_name_to_atom14_names = {
+    'ALA': ['N', 'CA', 'C', 'O', 'CB', '',    '',    '',    '',    '',    '',    '',    '',    ''],
+    'ARG': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'NE',  'CZ',  'NH1', 'NH2', '',    '',    ''],
+    'ASN': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'ND2', '',    '',    '',    '',    '',    ''],
+    'ASP': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'OD1', 'OD2', '',    '',    '',    '',    '',    ''],
+    'CYS': ['N', 'CA', 'C', 'O', 'CB', 'SG',  '',    '',    '',    '',    '',    '',    '',    ''],
+    'GLN': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'NE2', '',    '',    '',    '',    ''],
+    'GLU': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'OE1', 'OE2', '',    '',    '',    '',    ''],
+    'GLY': ['N', 'CA', 'C', 'O', '',   '',    '',    '',    '',    '',    '',    '',    '',    ''],
+    'HIS': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'ND1', 'CD2', 'CE1', 'NE2', '',    '',    '',    ''],
+    'ILE': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', 'CD1', '',    '',    '',    '',    '',    ''],
+    'LEU': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', '',    '',    '',    '',    '',    ''],
+    'LYS': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  'CE',  'NZ',  '',    '',    '',    '',    ''],
+    'MET': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'SD',  'CE',  '',    '',    '',    '',    '',    ''],
+    'PHE': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  '',    '',    ''],
+    'PRO': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD',  '',    '',    '',    '',    '',    '',    ''],
+    'SER': ['N', 'CA', 'C', 'O', 'CB', 'OG',  '',    '',    '',    '',    '',    '',    '',    ''],
+    'THR': ['N', 'CA', 'C', 'O', 'CB', 'OG1', 'CG2', '',    '',    '',    '',    '',    '',    ''],
+    'TRP': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'NE1', 'CE2', 'CE3', 'CZ2', 'CZ3', 'CH2'],
+    'TYR': ['N', 'CA', 'C', 'O', 'CB', 'CG',  'CD1', 'CD2', 'CE1', 'CE2', 'CZ',  'OH',  '',    ''],
+    'VAL': ['N', 'CA', 'C', 'O', 'CB', 'CG1', 'CG2', '',    '',    '',    '',    '',    '',    ''],
+    'UNK': ['',  '',   '',  '',  '',   '',    '',    '',    '',    '',    '',    '',    '',    ''],
+
+}
+# pylint: enable=line-too-long
+# pylint: enable=bad-whitespace
+
+
+# This is the standard residue order when coding AA type as a number.
+# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
+restypes = [
+    'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P',
+    'S', 'T', 'W', 'Y', 'V'
+]
+restype_order = {restype: i for i, restype in enumerate(restypes)}
+restype_num = len(restypes)  # := 20.
+unk_restype_index = restype_num  # Catch-all index for unknown restypes.
+
+restypes_with_x = restypes + ['X']
+restype_order_with_x = {restype: i for i, restype in enumerate(restypes_with_x)}
+
+
+def sequence_to_onehot(
+    sequence: str,
+    mapping: Mapping[str, int],
+    map_unknown_to_x: bool = False) -> np.ndarray:
+  """Maps the given sequence into a one-hot encoded matrix.
+
+  Args:
+    sequence: An amino acid sequence.
+    mapping: A dictionary mapping amino acids to integers.
+    map_unknown_to_x: If True, any amino acid that is not in the mapping will be
+      mapped to the unknown amino acid 'X'. If the mapping doesn't contain
+      amino acid 'X', an error will be thrown. If False, any amino acid not in
+      the mapping will throw an error.
+
+  Returns:
+    A numpy array of shape (seq_len, num_unique_aas) with one-hot encoding of
+    the sequence.
+
+  Raises:
+    ValueError: If the mapping doesn't contain values from 0 to
+      num_unique_aas - 1 without any gaps.
+  """
+  num_entries = max(mapping.values()) + 1
+
+  if sorted(set(mapping.values())) != list(range(num_entries)):
+    raise ValueError('The mapping must have values from 0 to num_unique_aas-1 '
+                     'without any gaps. Got: %s' % sorted(mapping.values()))
+
+  one_hot_arr = np.zeros((len(sequence), num_entries), dtype=np.int32)
+
+  for aa_index, aa_type in enumerate(sequence):
+    if map_unknown_to_x:
+      if aa_type.isalpha() and aa_type.isupper():
+        aa_id = mapping.get(aa_type, mapping['X'])
+      else:
+        raise ValueError(f'Invalid character in the sequence: {aa_type}')
+    else:
+      aa_id = mapping[aa_type]
+    one_hot_arr[aa_index, aa_id] = 1
+
+  return one_hot_arr
+
+
+restype_1to3 = {
+    'A': 'ALA',
+    'R': 'ARG',
+    'N': 'ASN',
+    'D': 'ASP',
+    'C': 'CYS',
+    'Q': 'GLN',
+    'E': 'GLU',
+    'G': 'GLY',
+    'H': 'HIS',
+    'I': 'ILE',
+    'L': 'LEU',
+    'K': 'LYS',
+    'M': 'MET',
+    'F': 'PHE',
+    'P': 'PRO',
+    'S': 'SER',
+    'T': 'THR',
+    'W': 'TRP',
+    'Y': 'TYR',
+    'V': 'VAL',
+}
+
+
+# NB: restype_3to1 differs from Bio.PDB.protein_letters_3to1 by being a simple
+# 1-to-1 mapping of 3 letter names to one letter names. The latter contains
+# many more, and less common, three letter names as keys and maps many of these
+# to the same one letter name (including 'X' and 'U' which we don't use here).
+restype_3to1 = {v: k for k, v in restype_1to3.items()}
+
+# Define a restype name for all unknown residues.
+unk_restype = 'UNK'
+
+resnames = [restype_1to3[r] for r in restypes] + [unk_restype]
+resname_to_idx = {resname: i for i, resname in enumerate(resnames)}
+
+
+# The mapping here uses hhblits convention, so that B is mapped to D, J and O
+# are mapped to X, U is mapped to C, and Z is mapped to E. Other than that the
+# remaining 20 amino acids are kept in alphabetical order.
+# There are 2 non-amino acid codes, X (representing any amino acid) and
+# "-" representing a missing amino acid in an alignment.  The id for these
+# codes is put at the end (20 and 21) so that they can easily be ignored if
+# desired.
+HHBLITS_AA_TO_ID = {
+    'A': 0,
+    'B': 2,
+    'C': 1,
+    'D': 2,
+    'E': 3,
+    'F': 4,
+    'G': 5,
+    'H': 6,
+    'I': 7,
+    'J': 20,
+    'K': 8,
+    'L': 9,
+    'M': 10,
+    'N': 11,
+    'O': 20,
+    'P': 12,
+    'Q': 13,
+    'R': 14,
+    'S': 15,
+    'T': 16,
+    'U': 1,
+    'V': 17,
+    'W': 18,
+    'X': 20,
+    'Y': 19,
+    'Z': 3,
+    '-': 21,
+}
+
+# Partial inversion of HHBLITS_AA_TO_ID.
+ID_TO_HHBLITS_AA = {
+    0: 'A',
+    1: 'C',  # Also U.
+    2: 'D',  # Also B.
+    3: 'E',  # Also Z.
+    4: 'F',
+    5: 'G',
+    6: 'H',
+    7: 'I',
+    8: 'K',
+    9: 'L',
+    10: 'M',
+    11: 'N',
+    12: 'P',
+    13: 'Q',
+    14: 'R',
+    15: 'S',
+    16: 'T',
+    17: 'V',
+    18: 'W',
+    19: 'Y',
+    20: 'X',  # Includes J and O.
+    21: '-',
+}
+
+restypes_with_x_and_gap = restypes + ['X', '-']
+MAP_HHBLITS_AATYPE_TO_OUR_AATYPE = tuple(
+    restypes_with_x_and_gap.index(ID_TO_HHBLITS_AA[i])
+    for i in range(len(restypes_with_x_and_gap)))
+
+
+def _make_standard_atom_mask() -> np.ndarray:
+  """Returns [num_res_types, num_atom_types] mask array."""
+  # +1 to account for unknown (all 0s).
+  mask = np.zeros([restype_num + 1, atom_type_num], dtype=np.int32)
+  for restype, restype_letter in enumerate(restypes):
+    restype_name = restype_1to3[restype_letter]
+    atom_names = residue_atoms[restype_name]
+    for atom_name in atom_names:
+      atom_type = atom_order[atom_name]
+      mask[restype, atom_type] = 1
+  return mask
+
+
+STANDARD_ATOM_MASK = _make_standard_atom_mask()
+
+
+# A one hot representation for the first and second atoms defining the axis
+# of rotation for each chi-angle in each residue.
+def chi_angle_atom(atom_index: int) -> np.ndarray:
+  """Define chi-angle rigid groups via one-hot representations."""
+  chi_angles_index = {}
+  one_hots = []
+
+  for k, v in chi_angles_atoms.items():
+    indices = [atom_types.index(s[atom_index]) for s in v]
+    indices.extend([-1]*(4-len(indices)))
+    chi_angles_index[k] = indices
+
+  for r in restypes:
+    res3 = restype_1to3[r]
+    one_hot = np.eye(atom_type_num)[chi_angles_index[res3]]
+    one_hots.append(one_hot)
+
+  one_hots.append(np.zeros([4, atom_type_num]))  # Add zeros for residue `X`.
+  one_hot = np.stack(one_hots, axis=0)
+  one_hot = np.transpose(one_hot, [0, 2, 1])
+
+  return one_hot
+
+chi_atom_1_one_hot = chi_angle_atom(1)
+chi_atom_2_one_hot = chi_angle_atom(2)
+
+# An array like chi_angles_atoms but using indices rather than names.
+chi_angles_atom_indices = [chi_angles_atoms[restype_1to3[r]] for r in restypes]
+chi_angles_atom_indices = tree.map_structure(
+    lambda atom_name: atom_order[atom_name], chi_angles_atom_indices)
+chi_angles_atom_indices = np.array([
+    chi_atoms + ([[0, 0, 0, 0]] * (4 - len(chi_atoms)))
+    for chi_atoms in chi_angles_atom_indices])
+
+# Mapping from (res_name, atom_name) pairs to the atom's chi group index
+# and atom index within that group.
+chi_groups_for_atom = collections.defaultdict(list)
+for res_name, chi_angle_atoms_for_res in chi_angles_atoms.items():
+  for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
+    for atom_i, atom in enumerate(chi_group):
+      chi_groups_for_atom[(res_name, atom)].append((chi_group_i, atom_i))
+chi_groups_for_atom = dict(chi_groups_for_atom)
+
+
+def _make_rigid_transformation_4x4(ex, ey, translation):
+  """Create a rigid 4x4 transformation matrix from two axes and transl."""
+  # Normalize ex.
+  ex_normalized = ex / np.linalg.norm(ex)
+
+  # make ey perpendicular to ex
+  ey_normalized = ey - np.dot(ey, ex_normalized) * ex_normalized
+  ey_normalized /= np.linalg.norm(ey_normalized)
+
+  # compute ez as cross product
+  eznorm = np.cross(ex_normalized, ey_normalized)
+  m = np.stack([ex_normalized, ey_normalized, eznorm, translation]).transpose()
+  m = np.concatenate([m, [[0., 0., 0., 1.]]], axis=0)
+  return m
+
+
+# create an array with (restype, atomtype) --> rigid_group_idx
+# and an array with (restype, atomtype, coord) for the atom positions
+# and compute affine transformation matrices (4,4) from one rigid group to the
+# previous group
+restype_atom37_to_rigid_group = np.zeros([21, 37], dtype=np.int32)
+restype_atom37_mask = np.zeros([21, 37], dtype=np.float32)
+restype_atom37_rigid_group_positions = np.zeros([21, 37, 3], dtype=np.float32)
+restype_atom14_to_rigid_group = np.zeros([21, 14], dtype=np.int32)
+restype_atom14_mask = np.zeros([21, 14], dtype=np.float32)
+restype_atom14_rigid_group_positions = np.zeros([21, 14, 3], dtype=np.float32)
+restype_rigid_group_default_frame = np.zeros([21, 8, 4, 4], dtype=np.float32)
+
+###############################################
+restype_atom14_to_atom37 = []
+restype_atom37_to_atom14 = []
+for rt in restypes:
+  atom_names = restype_name_to_atom14_names[restype_1to3[rt]]
+  restype_atom14_to_atom37.append([(atom_order[name] if name else 0) for name in atom_names])
+  atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
+  restype_atom37_to_atom14.append([(atom_name_to_idx14[name] if name in atom_name_to_idx14 else 0) for name in atom_types])
+restype_atom14_to_atom37.append([0] * 14)
+restype_atom37_to_atom14.append([0] * 37)
+restype_atom14_to_atom37 = np.array(restype_atom14_to_atom37, dtype=np.int32)
+restype_atom37_to_atom14 = np.array(restype_atom37_to_atom14, dtype=np.int32)
+################################################
+
+def _make_rigid_group_constants():
+  """Fill the arrays above."""
+  
+  
+  for restype, restype_letter in enumerate(restypes):
+    resname = restype_1to3[restype_letter]
+    for atomname, group_idx, atom_position in rigid_group_atom_positions[resname]:
+      atomtype = atom_order[atomname]
+      restype_atom37_to_rigid_group[restype, atomtype] = group_idx
+      restype_atom37_mask[restype, atomtype] = 1
+      restype_atom37_rigid_group_positions[restype, atomtype, :] = atom_position
+
+      atom14idx = restype_name_to_atom14_names[resname].index(atomname)
+      restype_atom14_to_rigid_group[restype, atom14idx] = group_idx
+      restype_atom14_mask[restype, atom14idx] = 1
+      restype_atom14_rigid_group_positions[restype, atom14idx, :] = atom_position
+    
+    atom_names = residue_atoms[resname]
+    atom_name_to_idx14 = {name: i for i, name in enumerate(atom_names)}
+      
+  for restype, restype_letter in enumerate(restypes):
+    resname = restype_1to3[restype_letter]
+    atom_positions = {name: np.array(pos) for name, _, pos
+                      in rigid_group_atom_positions[resname]}
+
+    # backbone to backbone is the identity transform
+    restype_rigid_group_default_frame[restype, 0, :, :] = np.eye(4)
+
+    # pre-omega-frame to backbone (currently dummy identity matrix)
+    restype_rigid_group_default_frame[restype, 1, :, :] = np.eye(4)
+
+    # phi-frame to backbone
+    mat = _make_rigid_transformation_4x4(
+        ex=atom_positions['N'] - atom_positions['CA'],
+        ey=np.array([1., 0., 0.]),
+        translation=atom_positions['N'])
+    restype_rigid_group_default_frame[restype, 2, :, :] = mat
+
+    # psi-frame to backbone
+    mat = _make_rigid_transformation_4x4(
+        ex=atom_positions['C'] - atom_positions['CA'],
+        ey=atom_positions['CA'] - atom_positions['N'],
+        translation=atom_positions['C'])
+    restype_rigid_group_default_frame[restype, 3, :, :] = mat
+
+    # chi1-frame to backbone
+    if chi_angles_mask[restype][0]:
+      base_atom_names = chi_angles_atoms[resname][0]
+      base_atom_positions = [atom_positions[name] for name in base_atom_names]
+      mat = _make_rigid_transformation_4x4(
+          ex=base_atom_positions[2] - base_atom_positions[1],
+          ey=base_atom_positions[0] - base_atom_positions[1],
+          translation=base_atom_positions[2])
+      restype_rigid_group_default_frame[restype, 4, :, :] = mat
+
+    # chi2-frame to chi1-frame
+    # chi3-frame to chi2-frame
+    # chi4-frame to chi3-frame
+    # luckily all rotation axes for the next frame start at (0,0,0) of the
+    # previous frame
+    for chi_idx in range(1, 4):
+      if chi_angles_mask[restype][chi_idx]:
+        axis_end_atom_name = chi_angles_atoms[resname][chi_idx][2]
+        axis_end_atom_position = atom_positions[axis_end_atom_name]
+        mat = _make_rigid_transformation_4x4(
+            ex=axis_end_atom_position,
+            ey=np.array([-1., 0., 0.]),
+            translation=axis_end_atom_position)
+        restype_rigid_group_default_frame[restype, 4 + chi_idx, :, :] = mat
+
+
+_make_rigid_group_constants()
+
+
+def make_atom14_dists_bounds(overlap_tolerance=1.5,
+                             bond_length_tolerance_factor=15):
+  """compute upper and lower bounds for bonds to assess violations."""
+  restype_atom14_bond_lower_bound = np.zeros([21, 14, 14], np.float32)
+  restype_atom14_bond_upper_bound = np.zeros([21, 14, 14], np.float32)
+  restype_atom14_bond_stddev = np.zeros([21, 14, 14], np.float32)
+  residue_bonds, residue_virtual_bonds, _ = load_stereo_chemical_props()
+  for restype, restype_letter in enumerate(restypes):
+    resname = restype_1to3[restype_letter]
+    atom_list = restype_name_to_atom14_names[resname]
+
+    # create lower and upper bounds for clashes
+    for atom1_idx, atom1_name in enumerate(atom_list):
+      if not atom1_name:
+        continue
+      atom1_radius = van_der_waals_radius[atom1_name[0]]
+      for atom2_idx, atom2_name in enumerate(atom_list):
+        if (not atom2_name) or atom1_idx == atom2_idx:
+          continue
+        atom2_radius = van_der_waals_radius[atom2_name[0]]
+        lower = atom1_radius + atom2_radius - overlap_tolerance
+        upper = 1e10
+        restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+        restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+        restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+        restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+
+    # overwrite lower and upper bounds for bonds and angles
+    for b in residue_bonds[resname] + residue_virtual_bonds[resname]:
+      atom1_idx = atom_list.index(b.atom1_name)
+      atom2_idx = atom_list.index(b.atom2_name)
+      lower = b.length - bond_length_tolerance_factor * b.stddev
+      upper = b.length + bond_length_tolerance_factor * b.stddev
+      restype_atom14_bond_lower_bound[restype, atom1_idx, atom2_idx] = lower
+      restype_atom14_bond_lower_bound[restype, atom2_idx, atom1_idx] = lower
+      restype_atom14_bond_upper_bound[restype, atom1_idx, atom2_idx] = upper
+      restype_atom14_bond_upper_bound[restype, atom2_idx, atom1_idx] = upper
+      restype_atom14_bond_stddev[restype, atom1_idx, atom2_idx] = b.stddev
+      restype_atom14_bond_stddev[restype, atom2_idx, atom1_idx] = b.stddev
+  return {'lower_bound': restype_atom14_bond_lower_bound,  # shape (21,14,14)
+          'upper_bound': restype_atom14_bond_upper_bound,  # shape (21,14,14)
+          'stddev': restype_atom14_bond_stddev,  # shape (21,14,14)
+         }