<jupyter_start><jupyter_text>Prepare protein structure dataset- Parse the PDB files into json documents- Parse the CATH labels from RCSB API, merge into the json documents- Create index on PDB ID- Insert json documents into DocumentDB<jupyter_code>stack_name = 'gnn-proteins' # name of CloudFormation stack
import os
import json

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
plt.style.use('seaborn')
import seaborn as sns
%matplotlib inline
!mkdir -p data<jupyter_output><empty_output><jupyter_text>Download data from [AlphaFold](https://alphafold.ebi.ac.uk/download)<jupyter_code>!wget -O data/UP000000805_243232_METJA.tar \
    https://ftp.ebi.ac.uk/pub/databases/alphafold/UP000000805_243232_METJA.tar 
!cd data && tar -xvf UP000000805_243232_METJA.tar
# count number of pdb files
!ls data/*pdb.gz | wc -l
## get uniprot ids from pdb filenames
files = os.listdir('data/')

uniprot_ids = []
for file in files:
    if file.startswith('AF-') and file.endswith('pdb.gz'):
        uniprot_id = file.split('-')[1]
        uniprot_ids.append(uniprot_id)
print(len(uniprot_ids))
# https://www.uniprot.org/help/api_idmapping
# mapping uniprot id to pdb id
import urllib.parse
import urllib.request

url = 'https://www.uniprot.org/uploadlists/'

params = {
    'from': 'ACC+ID',
    'to': 'PDB_ID',
    'format': 'tab',
    'query': ' '.join(uniprot_ids)
}

data = urllib.parse.urlencode(params)
data = data.encode('utf-8')
req = urllib.request.Request(url, data)
with urllib.request.urlopen(req) as f:
    response = f.read()

res = response.decode('utf-8')
import io
id_mappings = pd.read_csv(io.StringIO(res), sep='\t')
id_mappings.columns = ['uniprot_id', 'pdb_id']
print(id_mappings.shape)
id_mappings.head()
id_mappings.nunique()
id_mappings.drop_duplicates().shape
pdb_ids = id_mappings['pdb_id'].unique()
with open('pdb_ids.txt', 'w') as out:
    out.write(','.join(list(pdb_ids)))
%%bash
# download from RCSB PDB

cd data
bash ../batch_download.sh -f ../pdb_ids.txt -p
!ls data/*.pdb.gz | wc -l<jupyter_output>2233<jupyter_text>Connect to DocumentDB<jupyter_code>from pymongo import MongoClient
from utils import get_secret
secrets = get_secret(stack_name)
uri = 'mongodb://{}:{}@{}:{}/?tls=true&tlsCAFile=rds-combined-ca-bundle.pem&replicaSet=rs0&readPreference=secondaryPreferred&retryWrites=false'\
    .format(secrets['username'], secrets['password'], secrets['host'], secrets['port'])

client = MongoClient(uri)
db = client['proteins'] # create a database
collection = db['proteins'] # create a collection
# collection.delete_many({}) # to delete all documents in this collection
collection.create_index('id', unique=True) # unique index in the collection<jupyter_output><empty_output><jupyter_text>Parse PDB files<jupyter_code>from tqdm import tqdm
from Bio.PDB import PDBParser

import xpdb
from pdb_parse import parse_pdb_file_to_json_record
# PDB parser
pdb_parser = PDBParser(
    QUIET=True,
    PERMISSIVE=True,
    structure_builder=xpdb.SloppyStructureBuilder(),
)<jupyter_output><empty_output><jupyter_text>Parse and save data to DB- for rcsb, use PDB-chain as id, add identifiers {uniprot_ids: [], pdb_ids: []}- for AF, use `AF-Q58321` as id, add identifiers {uniprot_ids: [], pdb_ids: []}<jupyter_code>import glob
# Parse PDB files from AlphaFold DB
for pdb_file in tqdm(glob.glob('data/AF-*.pdb.gz')):
    # AlphaFold structure
    id_ = '-'.join(os.path.basename(pdb_file).split('-')[:2])
    rec = parse_pdb_file_to_json_record(
        pdb_parser,
        pdb_file,
        id_
    )[0]
    
    uniprot_id = os.path.basename(pdb_file).split('-')[1]
    # look up pdb id
    pdb_ids = list(id_mappings.loc[id_mappings['uniprot_id']==uniprot_id, 'pdb_id'])
    identifiers = {
        'uniprot_ids': [uniprot_id],
        'pdb_ids': pdb_ids
    }
    rec['id'] = id_
    rec['identifiers'] = identifiers
    rec['is_AF'] = 1

    collection.insert_one(rec)
# Parse PDB files from RCSB
for pdb_file in tqdm(glob.glob('data/*.pdb.gz')):
    if not os.path.basename(pdb_file).startswith('AF-'):
        # RCSB structure
        pdb_id = os.path.basename(pdb_file).split('.')[0]
        recs = parse_pdb_file_to_json_record(
            pdb_parser,
            pdb_file,
            pdb_id
        )
        
        # look up uniprot_ids
        uniprot_ids = list(id_mappings.loc[id_mappings['pdb_id']==pdb_id, 'uniprot_id'])
        identifiers = {
            'uniprot_ids': uniprot_ids,
            'pdb_ids': [pdb_id]
        }
        for rec in recs:
            rec['id'] = rec['name']
            rec['identifiers'] = identifiers
            rec['is_AF'] = 0
        
        collection.insert_many(recs)
print('Number of protein chains in the DocumentDB:', 
      collection.count_documents({}))<jupyter_output>Number of protein chains in the DocumentDB: 3151<jupyter_text>Load protein metadata from DocumentDB to split train/valid/test<jupyter_code>from sklearn.model_selection import train_test_split
match = {"is_AF": {"$exists": True}}
project = {"y": "$is_AF"}

pipeline = [
    {"$match": match},
    {"$project": project},
]
# aggregation pipeline
cur = collection.aggregate(pipeline)
# retrieve documents from the DB cursor
docs = [doc for doc in cur]
# split train/valid/test
df = pd.DataFrame(docs)
print(df.shape)
df.head()
# stratified split: full -> train/test
df_train, df_test = train_test_split(
    df, 
    test_size=0.2,
    stratify=df['y'], 
    random_state=42
)
print(df_train.shape, df_test.shape)
df_train['y'].mean(), df_test['y'].mean()
# stratified split: train -> train/valid
df_train, df_valid = train_test_split(
    df_train, 
    test_size=0.2,
    stratify=df_train['y'], 
    random_state=42
)
print(df_train.shape, df_valid.shape, df_test.shape)
# Add split information into Documents
for split, df_split in zip(
    ['train', 'valid', 'test'],
    [df_train, df_valid, df_test]
):
    result = collection.update_many(
        {'_id': {'$in': df_split['_id'].tolist()}}, 
        {'$set': {'split': split}}
    )
    print('Number of documents modified:',result.modified_count)<jupyter_output>Number of documents modified: 2016
Number of documents modified: 504
Number of documents modified: 631
{'fraction_non_alphanumeric': 0.13627546071774976, 'fraction_numerical': 0.013902360168121564, 'mean_word_length': 4.125932062966031, 'pattern_counts': {'":': 5, '<': 22, '<?xml version=': 0, '>': 24, 'https://': 4, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 4, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '6', 'dir': 'jupyter-structured-clean-dedup', 'id': '178309', 'license': 'Apache-2.0', 'n_tokens_mistral': 2325, 'n_tokens_neox': 2236, 'n_words': 486, 'path': 'Prepare_data.ipynb', 'repo_name': 'aws-samples/sagemaker-documentdb-train-gnn-for-millions-proteins', 'size': '6186'}