<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Script that generates data.py - list of unit cells of contaminants
"""

import csv
import hashlib
import itertools
import json
import math
import os
import re
import sys
import urllib2
from collections import OrderedDict
import xml.etree.ElementTree as ET
import numpy

SCIPY_CLUSTERING = True

if SCIPY_CLUSTERING:
    import scipy.cluster
    import scipy.spatial

if __name__ == '__main__' and __package__ is None:
    sys.path.insert(1,
          os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import dimple.cell

WIKI_URL = ('https://raw.githubusercontent.com/wiki/'
            'ccp4/dimple/Crystallization-Contaminants.md')

CONTAMINER_XML = 'https://strube.cbrc.kaust.edu.sa/static/contabase.xml'

CLUSTER_CUTOFF = 0.04

OUTPUT_PY_FILE = 'data.py'
OUTPUT_JSON_FILE = 'data.json'
CACHE_DIR = 'cached'

# ContaBase has homologs of these
NOT_IN_CONTABASE = ['POLG_HRV2', 'THRB_HUMAN', 'FA10_BOVIN']

def cached_urlopen(url, cache_name=None):
    assert os.path.dirname(__file__) == '.', "Run it from script's directory"
    if not cache_name:
        cache_name = hashlib.sha1('abc').hexdigest()[:12]
    elif cache_name == -1:
        cache_name = url.split('/')[-1]
    path = os.path.join(CACHE_DIR, cache_name)
    if not os.path.exists(path):
        if not os.path.isdir(CACHE_DIR):
            os.mkdir(CACHE_DIR)
        print '--> %s' % path
        response = urllib2.urlopen(url)
        with open(path, 'wb') as f:
            f.write(response.read())
    return open(path)

def parse_wiki_page(page):
    pat = re.compile(r'`[A-Z0-9]+_[A-Z0-9]+`')
    data = OrderedDict()
    for line in page:
        obj = pat.search(line)
        if obj:
            name = obj.group(0).strip('`')
            if name == 'LIKE_THIS':
                continue
        elif '`' in line:
            sys.exit('No UniProt name? ' + line)
        else:
            continue
        sp = line[obj.span()[1]:].rsplit(' - ', 1)
        desc = sp[0].strip(' ①② ')
        pdb_ids = []
        if len(sp) == 2:
            pdb_ids = [i.strip().upper() for i in sp[1].split(',')]
            assert all(len(i) == 4 for i in pdb_ids), pdb_ids
            assert all(i[0].isdigit() for i in pdb_ids), pdb_ids
        data[name] = (desc, pdb_ids)
    return data

def fetch_contaminer_protein_acs():
    response = cached_urlopen(CONTAMINER_XML, -1)
    root = ET.parse(response).getroot()
    all_acs = []
    for contaminant in root:
        assert contaminant.tag == 'contaminant'
        acs = [t.text for t in contaminant.findall('uniprot_id')]
        assert len(acs) == 1
        all_acs += acs
    return all_acs

def fetch_pdb_info_from_ebi(pdb_id):
    pdb_id = pdb_id.lower()
    # We tried to remove heteromers when quering Uniprot, but we still have
    # things like 4TTD (Hetero 3-mer with a single link to UniProt AC).
    # An assembly that is a monomer can also be classified as "hetero", e.g.
    # http://www.ebi.ac.uk/pdbe/api/pdb/entry/summary/2war
    # From what I understand such forms are unlikely as contaminants?
    # The same about "homo" monomeric assemblies that contain 2 distinct
    # polypeptide molecules (e.g. 5AQ9) or polypeptide and DNA or RNA.
    entry_url = "http://www.ebi.ac.uk/pdbe/api/pdb/entry/"
    response = cached_urlopen(entry_url + 'summary/' + pdb_id,
                              'ebi-summary-%s.json' % pdb_id)
    summary = json.load(response)[pdb_id]
    assert isinstance(summary, list)
    assert len(summary) == 1
    summary = summary[0]
    exper_method = summary['experimental_method']
    if exper_method != ['X-ray diffraction']:
        if len(exper_method) != 1:
            print 'Hmmm: multiple methods %s in %s' % (exper_method, pdb_id)
        if exper_method[0] not in ('X-ray powder diffraction',
                                   'Solution NMR',
                                   'Electron Microscopy',
                                   'Electron crystallography',
                                   'Neutron Diffraction'):
            print 'WARNING: got %s in %s' % (exper_method, pdb_id)
        return None
    forms = set(a['form'] for a in summary['assemblies'])
    assert forms.issubset({'homo', 'hetero'}) # some are both (2EKS)
    if forms != {'homo'}:
        return None
    entities = summary['number_of_entities']
    big_molecules = ['polypeptide', 'dna', 'dna/rna', 'rna']
    if sum(entities[x] for x in big_molecules) > 1:
        return None
    # now http://www.ebi.ac.uk/pdbe/api/pdb/entry/experiment/:pdbid
    response = cached_urlopen(entry_url + 'experiment/' + pdb_id,
                              'ebi-exper-%s.json' % pdb_id)
    summary = json.load(response)[pdb_id]
    assert isinstance(summary, list)
    assert len(summary) == 1
    summary = summary[0]
    parameters = [summary['cell'][x] for x in 'a b c alpha beta gamma'.split()]
    sg = summary['spacegroup']
    if sg == 'A 1': # pesky 1LKS
        return None
    # Correct spacegroup, but in non-standard settings.
    # We could convert it to the reference/standard settings.
    if sg in ['P 1 1 21']: # only this one for now
        return None

    # Standarize unit cell settings (e.g. 3wai, 1y6e).
    # It'd be better to use a library function (cctbx?) for this.
    # to_standard() below makes a<b<c in primitive orthorhombic cells.
    if sg in ['P 1 2 1', 'P 1 21 1'] and parameters[0] > parameters[2]:
        parameters[0], parameters[2] = parameters[2], parameters[0]

    cell = dimple.cell.Cell(tuple(parameters), symmetry=sg).to_standard()

    cell.pdb_id = pdb_id
    # Estimate the model quality (cell.quality). Higher is better.
    # Use simplistic criterium, similar to the one from
    # http://www.rcsb.org/pdb/statistics/clusterStatistics.do
    try:
        cell.quality = 1.0 / summary['resolution'] - summary['r_free']
    except (ValueError, TypeError):
        cell.quality = -10
    if summary['experiment_data_available'] != 'Y':
        cell.quality = -20
    return cell

def dump_uniprot_tab(pdb_ids):
    query_url = ('http://www.uniprot.org/uniprot/'
                 '?query=database:(type:pdb%%20%s)&sort=score&format=tab')
    pdb_ids = [x.strip(',') for x in sys.stdin.read().split()]
    assert all(len(x) == 4 for x in pdb_ids)
    print >>sys.stderr, 'Read %d IDs' % len(pdb_ids)
    for pdb_id in pdb_ids:
        response = cached_urlopen(query_url % pdb_id, 'up-pdb-%s.tab' % pdb_id)
        text = response.read()
        print pdb_id
        print text
        print

def uniprot_names_to_acs(names):
    query_url = ('http://www.uniprot.org/uniprot/'
                 '?query=%s&columns=id,entry%%20name&format=tab')
    acs = OrderedDict()
    for name in names:
        response = cached_urlopen(query_url % name, 'up-%s-ie.tab' % name)
        lines = response.readlines()
        assert len(lines) >= 2, 'up-%s-ie.tab' % name
        ac, entry_name = lines[1].strip().split('\t')
        assert entry_name == name
        acs[ac] = name
    return acs

def fetch_uniref_clusters(acs):
    query_url = ('http://www.uniprot.org/uniref/'
                 '?query=%s&fil=identity:1.0&format=tab')
    clusters = OrderedDict()
    for ac in acs:
        response = cached_urlopen(query_url % ac, 'ur100-%s.tab' % ac)
        reader = csv.DictReader(response, delimiter='\t')
        for d in reader:
            key = d['Cluster ID']
            name = d['Cluster name']
            if name.startswith('Cluster: '):
                name = name[9:]
            members = [a.strip() for a in d['Cluster members'].split(';')]
            if ac not in members:  # querying P63165 also returns P63165-2
                continue
            #if key != 'UniRef100_P02931': continue
            if verbose:
                print '%s -> %s (%s) - %d members, %d residues' % (
                        ac, key, name, len(members), int(d['Length']))
            if key in clusters:
                sys.exit('Duplicated cluster: ' + key)
            clusters[key] = members
    return clusters

def read_current_pdb_entries():
    f = cached_urlopen('http://www.rcsb.org/pdb/rest/getCurrent',
                       'current-rcsb.xml')
    root = ET.parse(f).getroot()<commit_msg>contaminants/prepare.py: better error message<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Script that generates data.py - list of unit cells of contaminants
"""

import csv
import hashlib
import itertools
import json
import math
import os
import re
import sys
import urllib2
from collections import OrderedDict
import xml.etree.ElementTree as ET
import numpy

SCIPY_CLUSTERING = True

if SCIPY_CLUSTERING:
    import scipy.cluster
    import scipy.spatial

if __name__ == '__main__' and __package__ is None:
    sys.path.insert(1,
          os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
import dimple.cell

WIKI_URL = ('https://raw.githubusercontent.com/wiki/'
            'ccp4/dimple/Crystallization-Contaminants.md')

CONTAMINER_XML = 'https://strube.cbrc.kaust.edu.sa/static/contabase.xml'

CLUSTER_CUTOFF = 0.04

OUTPUT_PY_FILE = 'data.py'
OUTPUT_JSON_FILE = 'data.json'
CACHE_DIR = 'cached'

# ContaBase has homologs of these
NOT_IN_CONTABASE = ['POLG_HRV2', 'THRB_HUMAN', 'FA10_BOVIN']

def cached_urlopen(url, cache_name):
    assert os.path.dirname(__file__) == '.', "Run it from script's directory"
    if not cache_name:
        cache_name = hashlib.sha1('abc').hexdigest()[:12]
    elif cache_name == -1:
        cache_name = url.split('/')[-1]
    path = os.path.join(CACHE_DIR, cache_name)
    if not os.path.exists(path):
        if not os.path.isdir(CACHE_DIR):
            os.mkdir(CACHE_DIR)
        print '--> %s' % path
        response = urllib2.urlopen(url)
        with open(path, 'wb') as f:
            f.write(response.read())
    return open(path)

def parse_wiki_page(page):
    pat = re.compile(r'`[A-Z0-9]+_[A-Z0-9]+`')
    data = OrderedDict()
    for line in page:
        obj = pat.search(line)
        if obj:
            name = obj.group(0).strip('`')
            if name == 'LIKE_THIS':
                continue
        elif '`' in line:
            sys.exit('No UniProt name? ' + line)
        else:
            continue
        sp = line[obj.span()[1]:].rsplit(' - ', 1)
        desc = sp[0].strip(' ①② ')
        pdb_ids = []
        if len(sp) == 2:
            pdb_ids = [i.strip().upper() for i in sp[1].split(',')]
            assert all(len(i) == 4 for i in pdb_ids), pdb_ids
            assert all(i[0].isdigit() for i in pdb_ids), pdb_ids
        data[name] = (desc, pdb_ids)
    return data

def fetch_contaminer_protein_acs():
    response = cached_urlopen(CONTAMINER_XML, -1)
    root = ET.parse(response).getroot()
    all_acs = []
    for contaminant in root:
        assert contaminant.tag == 'contaminant'
        acs = [t.text for t in contaminant.findall('uniprot_id')]
        assert len(acs) == 1
        all_acs += acs
    return all_acs

def fetch_pdb_info_from_ebi(pdb_id):
    pdb_id = pdb_id.lower()
    # We tried to remove heteromers when quering Uniprot, but we still have
    # things like 4TTD (Hetero 3-mer with a single link to UniProt AC).
    # An assembly that is a monomer can also be classified as "hetero", e.g.
    # http://www.ebi.ac.uk/pdbe/api/pdb/entry/summary/2war
    # From what I understand such forms are unlikely as contaminants?
    # The same about "homo" monomeric assemblies that contain 2 distinct
    # polypeptide molecules (e.g. 5AQ9) or polypeptide and DNA or RNA.
    entry_url = "http://www.ebi.ac.uk/pdbe/api/pdb/entry/"
    response = cached_urlopen(entry_url + 'summary/' + pdb_id,
                              'ebi-summary-%s.json' % pdb_id)
    summary = json.load(response)[pdb_id]
    assert isinstance(summary, list)
    assert len(summary) == 1
    summary = summary[0]
    exper_method = summary['experimental_method']
    if exper_method != ['X-ray diffraction']:
        if len(exper_method) != 1:
            print 'Hmmm: multiple methods %s in %s' % (exper_method, pdb_id)
        if exper_method[0] not in ('X-ray powder diffraction',
                                   'Solution NMR',
                                   'Electron Microscopy',
                                   'Electron crystallography',
                                   'Neutron Diffraction'):
            print 'WARNING: got %s in %s' % (exper_method, pdb_id)
        return None
    forms = set(a['form'] for a in summary['assemblies'])
    assert forms.issubset({'homo', 'hetero'}) # some are both (2EKS)
    if forms != {'homo'}:
        return None
    entities = summary['number_of_entities']
    big_molecules = ['polypeptide', 'dna', 'dna/rna', 'rna']
    if sum(entities[x] for x in big_molecules) > 1:
        return None
    # now http://www.ebi.ac.uk/pdbe/api/pdb/entry/experiment/:pdbid
    response = cached_urlopen(entry_url + 'experiment/' + pdb_id,
                              'ebi-exper-%s.json' % pdb_id)
    summary = json.load(response)[pdb_id]
    assert isinstance(summary, list)
    assert len(summary) == 1
    summary = summary[0]
    parameters = [summary['cell'][x] for x in 'a b c alpha beta gamma'.split()]
    sg = summary['spacegroup']
    if sg == 'A 1': # pesky 1LKS
        return None
    # Correct spacegroup, but in non-standard settings.
    # We could convert it to the reference/standard settings.
    if sg in ['P 1 1 21']: # only this one for now
        return None

    # Standarize unit cell settings (e.g. 3wai, 1y6e).
    # It'd be better to use a library function (cctbx?) for this.
    # to_standard() below makes a<b<c in primitive orthorhombic cells.
    if sg in ['P 1 2 1', 'P 1 21 1'] and parameters[0] > parameters[2]:
        parameters[0], parameters[2] = parameters[2], parameters[0]

    cell = dimple.cell.Cell(tuple(parameters), symmetry=sg).to_standard()

    cell.pdb_id = pdb_id
    # Estimate the model quality (cell.quality). Higher is better.
    # Use simplistic criterium, similar to the one from
    # http://www.rcsb.org/pdb/statistics/clusterStatistics.do
    try:
        cell.quality = 1.0 / summary['resolution'] - summary['r_free']
    except (ValueError, TypeError):
        cell.quality = -10
    if summary['experiment_data_available'] != 'Y':
        cell.quality = -20
    return cell

def dump_uniprot_tab(pdb_ids):
    query_url = ('http://www.uniprot.org/uniprot/'
                 '?query=database:(type:pdb%%20%s)&sort=score&format=tab')
    pdb_ids = [x.strip(',') for x in sys.stdin.read().split()]
    assert all(len(x) == 4 for x in pdb_ids)
    print >>sys.stderr, 'Read %d IDs' % len(pdb_ids)
    for pdb_id in pdb_ids:
        response = cached_urlopen(query_url % pdb_id, 'up-pdb-%s.tab' % pdb_id)
        text = response.read()
        print pdb_id
        print text
        print

def uniprot_names_to_acs(names):
    query_url = ('http://www.uniprot.org/uniprot/'
                 '?query=%s&columns=id,entry%%20name&format=tab')
    acs = OrderedDict()
    for name in names:
        query = query_url % name
        response = cached_urlopen(query, 'up-%s-ie.tab' % name)
        lines = response.readlines()
        assert len(lines) >= 2, 'up-%s-ie.tab' % name
        ac, entry_name = lines[1].strip().split('\t')
        assert entry_name == name, '%s\n%s != %s' % (query, entry_name, name)
        acs[ac] = name
    return acs

def fetch_uniref_clusters(acs):
    query_url = ('http://www.uniprot.org/uniref/'
                 '?query=%s&fil=identity:1.0&format=tab')
    clusters = OrderedDict()
    for ac in acs:
        response = cached_urlopen(query_url % ac, 'ur100-%s.tab' % ac)
        reader = csv.DictReader(response, delimiter='\t')
        for d in reader:
            key = d['Cluster ID']
            name = d['Cluster name']
            if name.startswith('Cluster: '):
                name = name[9:]
            members = [a.strip() for a in d['Cluster members'].split(';')]
            if ac not in members:  # querying P63165 also returns P63165-2
                continue
            #if key != 'UniRef100_P02931': continue
            if verbose:
                print '%s -> %s (%s) - %d members, %d residues' % (
                        ac, key, name, len(members), int(d['Length']))
            if key in clusters:
                sys.exit('Duplicated cluster: ' + key)
            clusters[key] = members
    return clusters

def read_current_pdb_entries():
    f = cached_urlopen('http://www.rcsb.org/pdb/rest/getCurrent',
                       'current-rcsb.xml')
    root = ET.parse(f).getroot()
{'fraction_non_alphanumeric': 0.1322070582576805, 'fraction_numerical': 0.012625503517104552, 'mean_word_length': 2.416307249948655, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 17, 'https://': 4, 'lorem ipsum': 0, 'www.': 16, 'xml': 6}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 8, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '2229422', 'n_tokens_mistral': 5789, 'n_tokens_neox': 5530, 'n_words': 1403}