Add get_structure

poetry.lock

@@ -137,6 +137,17 @@ soupsieve = ">1.2"
 html5lib = ["html5lib"]
 lxml = ["lxml"]
 
+[[package]]
+name = "biopython"
+version = "1.81"
+description = "Freely available tools for computational molecular biology."
+category = "main"
+optional = false
+python-versions = ">=3.7"
+
+[package.dependencies]
+numpy = "*"
+
 [[package]]
 name = "bleach"
 version = "6.0.0"
@@ -271,6 +282,17 @@ category = "main"
 optional = false
 python-versions = ">=3.6"
 
+[[package]]
+name = "exceptiongroup"
+version = "1.1.1"
+description = "Backport of PEP 654 (exception groups)"
+category = "dev"
+optional = false
+python-versions = ">=3.7"
+
+[package.extras]
+test = ["pytest (>=6)"]
+
 [[package]]
 name = "executing"
 version = "1.2.0"
@@ -347,6 +369,14 @@ docs = ["sphinx (>=3.5)", "jaraco.packaging (>=9)", "rst.linker (>=1.9)", "furo"
 perf = ["ipython"]
 testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "flake8 (<5)", "pytest-cov", "pytest-enabler (>=1.3)", "packaging", "pyfakefs", "flufl.flake8", "pytest-perf (>=0.9.2)", "pytest-black (>=0.3.7)", "pytest-mypy (>=0.9.1)", "pytest-flake8", "importlib-resources (>=1.3)"]
 
+[[package]]
+name = "iniconfig"
+version = "2.0.0"
+description = "brain-dead simple config-ini parsing"
+category = "dev"
+optional = false
+python-versions = ">=3.7"
+
 [[package]]
 name = "ipykernel"
 version = "6.21.3"
@@ -971,6 +1001,18 @@ python-versions = ">=3.7"
 docs = ["furo (>=2022.12.7)", "proselint (>=0.13)", "sphinx-autodoc-typehints (>=1.22,!=1.23.4)", "sphinx (>=6.1.3)"]
 test = ["appdirs (==1.4.4)", "covdefaults (>=2.2.2)", "pytest-cov (>=4)", "pytest-mock (>=3.10)", "pytest (>=7.2.1)"]
 
+[[package]]
+name = "pluggy"
+version = "1.0.0"
+description = "plugin and hook calling mechanisms for python"
+category = "dev"
+optional = false
+python-versions = ">=3.6"
+
+[package.extras]
+dev = ["pre-commit", "tox"]
+testing = ["pytest", "pytest-benchmark"]
+
 [[package]]
 name = "prometheus-client"
 version = "0.16.0"
@@ -1104,6 +1146,26 @@ category = "main"
 optional = false
 python-versions = ">=3.7"
 
+[[package]]
+name = "pytest"
+version = "7.2.2"
+description = "pytest: simple powerful testing with Python"
+category = "dev"
+optional = false
+python-versions = ">=3.7"
+
+[package.dependencies]
+attrs = ">=19.2.0"
+colorama = {version = "*", markers = "sys_platform == \"win32\""}
+exceptiongroup = {version = ">=1.0.0rc8", markers = "python_version < \"3.11\""}
+iniconfig = "*"
+packaging = "*"
+pluggy = ">=0.12,<2.0"
+tomli = {version = ">=1.0.0", markers = "python_version < \"3.11\""}
+
+[package.extras]
+testing = ["argcomplete", "hypothesis (>=3.56)", "mock", "nose", "pygments (>=2.7.2)", "requests", "xmlschema"]
+
 [[package]]
 name = "python-dateutil"
 version = "2.8.2"
@@ -1389,6 +1451,14 @@ category = "main"
 optional = false
 python-versions = ">=2.6, !=3.0.*, !=3.1.*, !=3.2.*"
 
+[[package]]
+name = "tomli"
+version = "2.0.1"
+description = "A lil' TOML parser"
+category = "dev"
+optional = false
+python-versions = ">=3.7"
+
 [[package]]
 name = "tomlkit"
 version = "0.11.6"
@@ -1569,7 +1639,7 @@ testing = ["pytest (>=6)", "pytest-checkdocs (>=2.4)", "flake8 (<5)", "pytest-co
 [metadata]
 lock-version = "1.1"
 python-versions = "^3.10"
-content-hash = "23a858c56f35b320a1eb54d146b830db79180c7be4a22c0e9afb55c0a7163f40"
+content-hash = "1e79d688b56335b1eafcb169572e0b8983eff0cb2da5ece8807ae02316f25f12"
 
 [metadata.files]
 altair = []
@@ -1610,6 +1680,7 @@ backcall = [
     {file = "backcall-0.2.0.tar.gz", hash = "sha256:5cbdbf27be5e7cfadb448baf0aa95508f91f2bbc6c6437cd9cd06e2a4c215e1e"},
 ]
 beautifulsoup4 = []
+biopython = []
 bleach = []
 blinker = []
 cachetools = []
@@ -1633,6 +1704,7 @@ entrypoints = [
     {file = "entrypoints-0.4-py3-none-any.whl", hash = "sha256:f174b5ff827504fd3cd97cc3f8649f3693f51538c7e4bdf3ef002c8429d42f9f"},
     {file = "entrypoints-0.4.tar.gz", hash = "sha256:b706eddaa9218a19ebcd67b56818f05bb27589b1ca9e8d797b74affad4ccacd4"},
 ]
+exceptiongroup = []
 executing = []
 fastjsonschema = []
 fqdn = []
@@ -1640,6 +1712,7 @@ gitdb = []
 gitpython = []
 idna = []
 importlib-metadata = []
+iniconfig = []
 ipykernel = []
 ipyspeck = []
 ipython = []
@@ -1694,6 +1767,10 @@ pickleshare = [
 ]
 pillow = []
 platformdirs = []
+pluggy = [
+    {file = "pluggy-1.0.0-py2.py3-none-any.whl", hash = "sha256:74134bbf457f031a36d68416e1509f34bd5ccc019f0bcc952c7b909d06b37bd3"},
+    {file = "pluggy-1.0.0.tar.gz", hash = "sha256:4224373bacce55f955a878bf9cfa763c1e360858e330072059e10bad68531159"},
+]
 prometheus-client = []
 prompt-toolkit = []
 protobuf = []
@@ -1716,6 +1793,7 @@ pydeck = []
 pygments = []
 pympler = []
 pyrsistent = []
+pytest = []
 python-dateutil = [
     {file = "python-dateutil-2.8.2.tar.gz", hash = "sha256:0123cacc1627ae19ddf3c27a5de5bd67ee4586fbdd6440d9748f8abb483d3e86"},
     {file = "python_dateutil-2.8.2-py2.py3-none-any.whl", hash = "sha256:961d03dc3453ebbc59dbdea9e4e11c5651520a876d0f4db161e8674aae935da9"},
@@ -1789,6 +1867,10 @@ toml = [
     {file = "toml-0.10.2-py2.py3-none-any.whl", hash = "sha256:806143ae5bfb6a3c6e736a764057db0e6a0e05e338b5630894a5f779cabb4f9b"},
     {file = "toml-0.10.2.tar.gz", hash = "sha256:b3bda1d108d5dd99f4a20d24d9c348e91c4db7ab1b749200bded2f839ccbe68f"},
 ]
+tomli = [
+    {file = "tomli-2.0.1-py3-none-any.whl", hash = "sha256:939de3e7a6161af0c887ef91b7d41a53e7c5a1ca976325f429cb46ea9bc30ecc"},
+    {file = "tomli-2.0.1.tar.gz", hash = "sha256:de526c12914f0c550d15924c62d72abc48d6fe7364aa87328337a31007fe8a4f"},
+]
 tomlkit = []
 toolz = []
 tornado = []

protention/attention.py

@@ -0,0 +1,34 @@
+from io import StringIO
+from urllib import request
+
+from Bio.PDB import PDBParser, Structure
+
+
+def get_structure(pdb_code: str) -> Structure:
+    """
+    Get structure from PDB
+    """
+    pdb_url = f"https://files.rcsb.org/download/{pdb_code}.pdb"
+    pdb_data = request.urlopen(pdb_url).read().decode("utf-8")
+    file = StringIO(pdb_data)
+    parser = PDBParser()
+    structure = parser.get_structure(pdb_code, file)
+    return structure
+
+
+def get_attention(
+    pdb_code: str, chain_ids: list[str], layer: int, head: int, min_attn: float = 0.2
+):
+    """
+    Get attention from T5
+    """
+    # fetch structure
+    structure = get_structure(pdb_code)
+
+    # get model
+
+    # call model
+
+    # get attention
+
+    # extract attention

pyproject.toml

@@ -8,8 +8,10 @@ authors = ["Aksel Lenes <[email protected]>"]
 python = "^3.10"
 streamlit = "^1.20.0"
 stmol = "^0.0.9"
+biopython = "^1.81"
 
 [tool.poetry.dev-dependencies]
+pytest = "^7.2.2"
 
 [build-system]
 requires = ["poetry-core>=1.0.0"]

tests/test_attention.py

@@ -0,0 +1,10 @@
+from Bio.PDB.Structure import Structure
+
+from protention.attention import get_structure
+
+
+def test_get_structure():
+    pdb_id = "1AKE"
+    structure = get_structure(pdb_id)
+    assert structure is not None
+    assert isinstance(structure, Structure)