rcsb.exdb 1.31__py3-none-any.whl

rcsb/exdb/examples-seq/EntityPolymerExtractor.py

@@ -0,0 +1,544 @@
+##
+# File: EntityPolymerExtractor.py
+# Date: 19-Feb-2019  jdw
+#
+# Selected utilities to extract entity polymer mapping and feature data
+# from the exchange database schema.
+#
+# Updates:
+#
+#
+##
+__docformat__ = "google en"
+__author__ = "John Westbrook"
+__email__ = "jwest@rcsb.rutgers.edu"
+__license__ = "Apache 2.0"
+
+import copy
+import logging
+import os
+
+from rcsb.db.mongo.Connection import Connection
+from rcsb.db.mongo.MongoDbUtil import MongoDbUtil
+from rcsb.utils.io.MarshalUtil import MarshalUtil
+
+
+logger = logging.getLogger(__name__)
+
+
+class EntityPolymerExtractor(object):
+    """Utilities to extract polymer related data from entry and entity collections."""
+
+    def __init__(self, cfgOb, **kwargs):
+        self.__cfgOb = cfgOb
+        self.__resourceName = "MONGO_DB"
+        self.__mU = MarshalUtil()
+        self.__entryD, self.__authAsymIdIndex = self.__rebuildCache(**kwargs)
+        #
+
+    def __rebuildCache(self, **kwargs):
+        useCache = kwargs.get("useCache", True)
+        dirPath = kwargs.get("exdbDirPath", ".")
+        cacheKwargs = kwargs.get("cacheKwargs", {"fmt": "pickle"})
+        #
+        ext = "pic" if cacheKwargs["fmt"] == "pickle" else "json"
+        fn = "entity-polymer-extracted-data-cache" + "." + ext
+        cacheFilePath = os.path.join(dirPath, fn)
+        #
+        cD = {"entryD": {}, "authIdxD": {}}
+        try:
+            self.__mU.mkdir(dirPath)
+            if not useCache:
+                for fp in [cacheFilePath]:
+                    try:
+                        os.remove(fp)
+                    except Exception:
+                        pass
+
+            if useCache and cacheFilePath and os.access(cacheFilePath, os.R_OK):
+                cD = self.__mU.doImport(cacheFilePath, **cacheKwargs)
+            else:
+                entryD = self.__selectEntries(**kwargs)
+                entryD = self.__selectPolymerEntities(entryD, **kwargs)
+                authIdxD = self.__buildIndices(entryD)
+                cD["entryD"] = entryD
+                cD["authIdxD"] = authIdxD
+                if cacheFilePath:
+                    ok = self.__mU.doExport(cacheFilePath, cD, **cacheKwargs)
+                    logger.info("Saved entity-polymer extracted results (%d) status %r in %s", len(entryD), ok, cacheFilePath)
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+        return cD["entryD"], cD["authIdxD"]
+
+    def __buildIndices(self, entryD):
+        indD = {}
+        for entryId, eD in entryD.items():
+            entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+            for entityId, pD in entityD.items():
+                for authAsymId in pD["auth_asym_ids"]:
+                    # avoid tuples for json serialization
+                    # indD[(entryId, authAsymId)] = entityId
+                    indD[entryId + "_" + authAsymId] = entityId
+        return indD
+
+    def getEntryCount(self):
+        return len(self.__entryD)
+
+    def getRefSeqAccessions(self, dbName):
+        acL = []
+        try:
+            for _, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for _, pD in entityD.items():
+                    for dD in pD["struct_ref"]:
+                        if "pdbx_db_accession" in dD and dD["db_name"] == dbName:
+                            acL.append(dD["pdbx_db_accession"])
+            return list(set(acL))
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+
+        return acL
+
+    def countRefSeqAccessions(self, dbName):
+        cD = {}
+        try:
+            for _, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for _, pD in entityD.items():
+                    iCount = 0
+                    for dD in pD["struct_ref"]:
+                        if "pdbx_db_accession" in dD and dD["db_name"] == dbName:
+                            iCount += 1
+                    cD[iCount] = cD[iCount] + 1 if iCount in cD else 1
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+
+        return cD
+
+    def countRefSeqAccessionDbType(self):
+        cD = {}
+        try:
+            for _, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for _, pD in entityD.items():
+                    for dD in pD["struct_ref"]:
+                        if "pdbx_db_accession" in dD and "db_name" in dD:
+                            cD[dD["db_name"]] = cD[dD["db_name"]] + 1 if dD["db_name"] in cD else 1
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+
+        return cD
+
+    def countRefSeqAccessionAny(self):
+        cD = {}
+        try:
+            for _, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for _, pD in entityD.items():
+                    iCount = len(pD["struct_ref"])
+                    # if iCount == 0:
+                    #    logger.info("entryId %r " % (entryId, entityId))
+                    cD[iCount] = cD[iCount] + 1 if iCount in cD else 1
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+
+        return cD
+
+    def getUniqueTaxons(self):
+        #
+        tD = {}
+        try:
+            for _, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for _, pD in entityD.items():
+                    # logger.info("Entity dictionary %r", pD.keys())
+                    if "rcsb_entity_source_organism" in pD:
+                        for dd in pD["rcsb_entity_source_organism"]:
+                            if "ncbi_taxonomy_id" in dd:
+                                tD[dd["ncbi_taxonomy_id"]] = tD[dd["ncbi_taxonomy_id"]] + 1 if dd["ncbi_taxonomy_id"] in tD else 1
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+        logger.info("Taxon coverage %d", len(tD))
+        return tD
+
+    def getOrigTaxons(self):
+        #
+        tD = {}
+        try:
+            for entryId, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for entityId, pD in entityD.items():
+                    # logger.info("Entity dictionary %r", pD.keys())
+                    if "original_taxonomy_ids" in pD:
+                        for tV in pD["original_taxonomy_ids"]:
+                            tD.setdefault(entryId, []).append((entityId, tV))
+                if entryId not in tD:
+                    logger.debug("No taxonomy for %s", entryId)
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+        logger.info("Taxon coverage %d", len(tD))
+        return tD
+
+    def countRefSeqAccessionByTaxon(self, dbNameList=None):
+        #
+        tD = {}
+        iCount = 0
+        #
+        try:
+            for _, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for _, pD in entityD.items():
+                    # logger.info("Entity dictionary %r", pD.keys())
+                    if "rcsb_entity_source_organism" in pD:
+                        for dd in pD["rcsb_entity_source_organism"]:
+                            if "ncbi_taxonomy_id" in dd:
+                                tId = dd["ncbi_taxonomy_id"]
+                                for dD in pD["struct_ref"]:
+                                    if "pdbx_db_accession" in dD and "db_name" in dD:
+                                        if dD["db_name"] in dbNameList:
+                                            tD.setdefault(tId, []).append(dD["pdbx_db_accession"])
+                                        iCount += 1
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+
+        logger.info("Total observed accessions %d", iCount)
+        return tD
+
+    def checkRefSeqAlignRange(self, dbName):
+        ok = True
+        try:
+            eCount = 0
+            aCount = 0
+            tCount = 0
+            for entryId, eD in self.__entryD.items():
+                entityD = eD["selected_polymer_entities"] if "selected_polymer_entities" in eD else {}
+                for entityId, pD in entityD.items():
+                    for dD in pD["struct_ref"]:
+                        if "db_name" in dD and dD["db_name"] == dbName:
+                            if "pdbx_db_accession" in dD and "alignD" in dD and "pdbx_seq_one_letter_code" in dD and "pdbx_align_begin" in dD:
+                                seqLen = len(dD["pdbx_seq_one_letter_code"])
+                                dbBegin = 100000000
+                                dbEnd = -1
+                                refSeqDbBegin = dD["pdbx_align_begin"]
+                                for authAsymId, alDL in dD["alignD"].items():
+                                    tCount += 1
+                                    difL = []
+                                    for alD in alDL:
+                                        tBeg = alD["db_align_beg"]
+                                        tEnd = alD["db_align_end"]
+                                        tDif = tEnd - tBeg + 1
+                                        difL.append(tDif)
+                                        dbBegin = min(tBeg, dbBegin)
+                                        dbEnd = max(tEnd, dbEnd)
+
+                                        # range is calculate on off -
+                                        # if seqLen < dbEnd - dbBegin + 1:
+                                        if seqLen < dbEnd - dbBegin and not refSeqDbBegin == dbBegin:
+                                            fDif = sum(difL)
+                                            logger.debug(
+                                                "Bad alignment for %r %r %r %r (%d) seqLen %r (%d) dbBegin %r dbEnd %r difL %r tDif %r",
+                                                entryId,
+                                                entityId,
+                                                authAsymId,
+                                                alD["pdbx_strand_id"],
+                                                len(alDL),
+                                                seqLen,
+                                                dbEnd - dbBegin + 1,
+                                                dbBegin,
+                                                dbEnd,
+                                                difL,
+                                                fDif,
+                                            )
+                                            aCount += 1
+
+                            else:
+                                eCount += 1
+            logger.info("Incomplete %s struct_ref record count %d", dbName, eCount)
+            logger.info("Inconsistent %s db reference alignments %d/%d", dbName, aCount, tCount)
+
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+            ok = False
+
+        return ok
+
+    def getEntityRefSeqAccessions(self, dbName, entryId, entityId):
+        acL = []
+        try:
+            dL = self.__entryD[entryId]["selected_polymer_entities"][entityId]["struct_ref"]
+            acL = list(set([d["pdbx_db_accession"] for d in dL if d["db_name"] == dbName]))
+        except Exception as e:
+            logger.exception("Failing with %s %r %r %s", dbName, entryId, entityId, str(e))
+        return acL
+
+    def __selectEntries(self, **kwargs):
+        """Return a dictionary of PDB entries satifying the input conditions (e.g. method, resolution limit)"""
+
+        dbName = kwargs.get("dbName", "pdbx_core")
+        collectionName = kwargs.get("collectionName", "pdbx_core_entry")
+        selectionQueryD = kwargs.get("entrySelectionQuery", {})
+        #
+        entryD = {}
+        try:
+            with Connection(cfgOb=self.__cfgOb, resourceName=self.__resourceName) as client:
+                mg = MongoDbUtil(client)
+                if mg.collectionExists(dbName, collectionName):
+                    logger.info("%s %s document count is %d", dbName, collectionName, mg.count(dbName, collectionName))
+                    qD = {}
+                    if selectionQueryD:
+                        qD.update(qD)
+                    selectL = ["rcsb_entry_container_identifiers"]
+                    dL = mg.fetch(dbName, collectionName, selectL, queryD=qD)
+                    logger.info("Selection %r fetch result count %d", selectL, len(dL))
+                    #
+                    for dD in dL:
+                        #
+                        if (
+                            ("rcsb_entry_container_identifiers" in dD)
+                            and ("entry_id" in dD["rcsb_entry_container_identifiers"])
+                            and ("polymer_entity_ids" in dD["rcsb_entry_container_identifiers"])
+                            and dD["rcsb_entry_container_identifiers"]["polymer_entity_ids"]
+                        ):
+                            entryD[dD["rcsb_entry_container_identifiers"]["entry_id"]] = {"polymer_entity_ids": dD["rcsb_entry_container_identifiers"]["polymer_entity_ids"]}
+
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+        return entryD
+        #
+
+    def __selectPolymerEntities(self, entryD, **kwargs):
+        """Skeleton entity selector recovering essential biological sequence mapping features
+        for macromolecules (default type = protein).
+
+         "1CP9": {
+             "polymer_entity_ids": [
+                "1",
+                "2"
+             ],
+             "selected_polymer_entities": {
+                "1": {
+                   "rcsb_multiple_source_flag": "N",
+                   "asym_ids": [
+                      "A"
+                   ],
+                   "auth_asym_ids": [
+                      "A"
+                   ],
+                   "entity_id": "1",
+                   "type": "polypeptide(L)",
+                   "rcsb_entity_polymer_type": "Protein",
+                   "rcsb_entity_source_organism": [
+                      {
+                         "ncbi_taxonomy_id": 587,
+                         "beg_seq_num": 1,
+                         "end_seq_num": 205,
+                         "ncbi_scientific_name": "Providencia rettgeri"
+                      }
+                   ],
+                   "struct_ref": [
+                      {
+                         "id": "1",
+                         "db_name": "UNP",
+                         "pdbx_db_accession": "Q7WZI9",
+                         "entity_id": "1",
+                         "pdbx_seq_one_letter_code": "QSTQIKIERDNYGVPHIYANDTYSLFYGYGYA...",
+                         "alignD": {
+                            "A": [
+                               {
+                                  "align_id": "1",
+                                  "ref_id": "1",
+                                  "pdbx_PDB_id_code": "1CP9",
+                                  "pdbx_strand_id": "A",
+                                  "seq_align_beg": 1,
+                                  "seq_align_end": 205,
+                                  "pdbx_db_accession": "Q7WZI9",
+                                  "db_align_beg": 24,
+                                  "db_align_end": 228,
+                                  "pdbx_auth_seq_align_beg": "1",
+                                  "pdbx_auth_seq_align_end": "205",
+                                  "rcsb_entity_id": "1"
+                               }
+                            ]
+                         }
+                      }
+                   ]
+                },
+            "2": {
+                   "rcsb_multiple_source_flag": "N",
+                   "asym_ids": [
+                      "B"
+                   ],
+                   "auth_asym_ids": [
+                      "B"
+                   ],
+                   "entity_id": "2",
+                   "type": "polypeptide(L)",
+                   "rcsb_entity_polymer_type": "Protein",
+                   "rcsb_entity_source_organism": [
+                      {
+                         "ncbi_taxonomy_id": 587,
+                         "beg_seq_num": 1,
+                         "end_seq_num": 553,
+                         "ncbi_scientific_name": "Providencia rettgeri"
+                      }
+                   ],
+                   "struct_ref": [
+                      {
+                         "id": "2",
+                         "db_name": "UNP",
+                         "pdbx_db_accession": "Q7WZI9",
+                         "entity_id": "2",
+                         "pdbx_seq_one_letter_code": "SNVWLVGKTKASGAKAILLNGPQFGWFNPAYTYGIGLHG",
+                         "alignD": {
+                            "B": [
+                               {
+                                  "align_id": "2",
+                                  "ref_id": "2",
+                                  "pdbx_PDB_id_code": "1CP9",
+                                  "pdbx_strand_id": "B",
+                                  "seq_align_beg": 1,
+                                  "seq_align_end": 553,
+                                  "pdbx_db_accession": "Q7WZI9",
+                                  "db_align_beg": 285,
+                                  "db_align_end": 837,
+                                  "pdbx_auth_seq_align_beg": "1",
+                                  "pdbx_auth_seq_align_end": "553",
+                                  "rcsb_entity_id": "2"
+                               }
+                            ]
+                         }
+                      }
+                   ]
+                }
+             }
+           },
+
+        """
+        dbName = kwargs.get("dbName", "pdbx_core")
+        collectionName = kwargs.get("collectionName", "pdbx_core_polymer_entity")
+        resultKey = kwargs.get("resultKey", "selected_polymer_entities")
+
+        entryLimit = kwargs.get("entryLimit", None)
+        selectionQueryD = kwargs.get("entitySelectionQuery", {"entity_poly.rcsb_entity_polymer_type": "Protein"})
+        #
+        try:
+            with Connection(cfgOb=self.__cfgOb, resourceName=self.__resourceName) as client:
+                mg = MongoDbUtil(client)
+                if mg.collectionExists(dbName, collectionName):
+                    logger.info("%s %s document count is %d", dbName, collectionName, mg.count(dbName, collectionName))
+                    selectL = [
+                        "rcsb_polymer_entity_container_identifiers",
+                        "entity.rcsb_multiple_source_flag",
+                        "entity_poly.type",
+                        "entity_poly.rcsb_entity_polymer_type",
+                        "entity_poly.pdbx_seq_one_letter_code_can",
+                        "rcsb_entity_source_organism.ncbi_taxonomy_id",
+                        "rcsb_entity_source_organism.ncbi_scientific_name",
+                        "rcsb_entity_source_organism.beg_seq_num",
+                        "rcsb_entity_source_organism.end_seq_num",
+                        "struct_ref.id",
+                        "struct_ref.pdbx_db_accession",
+                        "struct_ref.db_name",
+                        "struct_ref.entity_id",
+                        "struct_ref.pdbx_seq_one_letter_code",
+                        "struct_ref.pdbx_align_begin",
+                        "struct_ref_seq",
+                        #
+                        "entity_src_nat.pdbx_ncbi_taxonomy_id",
+                        "entity_src_gen.pdbx_gene_src_ncbi_taxonomy_id",
+                        "entity_src_gen.pdbx_host_org_ncbi_taxonomy_id",
+                        "pdbx_entity_src_syn.ncbi_taxonomy_id",
+                    ]
+                    iCount = 0
+                    for entryId in entryD:
+                        #
+                        if resultKey in entryD[entryId]:
+                            continue
+                        #
+                        qD = {"rcsb_polymer_entity_container_identifiers.entry_id": entryId}
+                        qD.update(selectionQueryD)
+                        #
+                        dL = mg.fetch(dbName, collectionName, selectL, queryD=qD)
+                        logger.debug("%s query %r fetch result count %d", entryId, qD, len(dL))
+                        eD = {}
+                        for ii, dD in enumerate(dL, 1):
+                            rD = {}
+                            logger.debug("%s (%4d) d is %r", entryId, ii, dD)
+                            if "entity" in dD:
+                                rD["rcsb_multiple_source_flag"] = dD["entity"]["rcsb_multiple_source_flag"] if "rcsb_multiple_source_flag" in dD["entity"] else "N"
+                            #
+                            if "rcsb_polymer_entity_container_identifiers" in dD:
+                                rD["asym_ids"] = dD["rcsb_entity_container_identifiers"]["asym_ids"] if "asym_ids" in dD["rcsb_entity_container_identifiers"] else []
+                                rD["auth_asym_ids"] = dD["rcsb_entity_container_identifiers"]["auth_asym_ids"] if "auth_asym_ids" in dD["rcsb_entity_container_identifiers"] else []
+                                rD["entity_id"] = dD["rcsb_entity_container_identifiers"]["entity_id"]
+                            #
+                            if "entity_poly" in dD:
+                                rD["type"] = dD["entity_poly"]["type"] if "type" in dD["entity_poly"] else None
+                                rD["rcsb_entity_polymer_type"] = dD["entity_poly"]["rcsb_entity_polymer_type"] if "rcsb_entity_polymer_type" in dD["entity_poly"] else None
+                                rD["entity_polymer_length"] = len(dD["entity_poly"]["pdbx_seq_one_letter_code_can"]) if "pdbx_seq_one_letter_code_can" in dD["entity_poly"] else 0
+                            #
+                            tL = []
+                            if "rcsb_entity_source_organism" in dD:
+                                for tD in dD["rcsb_entity_source_organism"]:
+                                    tL.append(tD)
+                            rD["rcsb_entity_source_organism"] = copy.copy(tL)
+                            #
+                            qDL = []
+                            if "struct_ref" in dD:
+                                for tD in dD["struct_ref"]:
+                                    if "db_name" in tD:
+                                        tD["db_name"] = str(tD["db_name"]).upper().strip()
+                                        tD["db_name"] = "UNP" if tD["db_name"] in ["TREMBL"] else tD["db_name"]
+                                    qDL.append(tD)
+                                if "struct_ref_seq" in dD:
+                                    for qD in qDL:
+                                        refId = qD["id"]
+                                        alignL = []
+                                        for tD in dD["struct_ref_seq"]:
+                                            if refId == tD["ref_id"]:
+                                                alignL.append(tD)
+                                        # qD['align_list'] = copy.copy(aL)
+                                        for align in alignL:
+                                            authAsymId = align["pdbx_strand_id"]
+                                            qD.setdefault("alignD", {}).setdefault(authAsymId, []).append(align)
+
+                            rD["struct_ref"] = qDL
+                            #
+                            taxIdL = []
+                            if "entity_src_nat" in dD:
+                                for tD in dD["entity_src_nat"]:
+                                    if "pdbx_ncbi_taxonomy_id" in tD:
+                                        taxIdL.append(tD["pdbx_ncbi_taxonomy_id"])
+                            if "entity_src_gen" in dD:
+                                for tD in dD["entity_src_gen"]:
+                                    if "pdbx_gene_src_ncbi_taxonomy_id" in tD:
+                                        taxIdL.append(tD["pdbx_gene_src_ncbi_taxonomy_id"])
+                                    if "pdbx_host_org_ncbi_taxonomy_id" in tD:
+                                        taxIdL.append(tD["pdbx_host_org_ncbi_taxonomy_id"])
+                            if "pdbx_entity_src_syn" in dD:
+                                for tD in dD["pdbx_entity_src_syn"]:
+                                    if "ncbi_taxonomy_id" in tD:
+                                        taxIdL.append(tD["ncbi_taxonomy_id"])
+                            qL = []
+                            for taxId in taxIdL:
+                                ttL = [int(t.strip()) for t in taxId.split(",") if t.strip().isdigit()]
+                                qL.extend(ttL)
+                            logger.debug("TaxId list %r", qL)
+                            rD["original_taxonomy_ids"] = copy.copy(list(set(qL)))
+                            #
+                            if "entity_id" in rD:
+                                eD[rD["entity_id"]] = copy.copy(rD)
+
+                        entryD[entryId][resultKey] = copy.copy(eD)
+
+                        iCount += 1
+                        if iCount % 1000 == 0:
+                            logger.info("Completed fetch %d/%d entries", iCount, len(entryD))
+                        if entryLimit and iCount >= entryLimit:
+                            logger.info("Quitting after %d", iCount)
+                            break
+
+        except Exception as e:
+            logger.exception("Failing with %s", str(e))
+
+        return entryD