feat: SCHWED-3119 Command line interface for QS-score

actions/ost-qs-score

-"""Calculate Quaternary Structure score (QS-score) between two complexes.
-
-"""
+"""Calculate Quaternary Structure score (QS-score) between two complexes."""
 
 import os
 import sys
+import json
 import argparse
 
 import ost
-from ost.io import LoadPDB, LoadMMCIF
+from ost.io import (LoadPDB, LoadMMCIF, MMCifInfoBioUnit, MMCifInfo,
+                    MMCifInfoTransOp)
 from ost import PushVerbosityLevel
 from ost.mol.alg import qsscoring
 
@@ -47,6 +47,11 @@ def _ParseArgs():
               "Each separate mapping consist of key:value pairs where key "
               "is the chain name in model and value is the chain name in "
               "reference."))
+    parser.add_argument(
+        "-o",
+        "--output",
+        dest="output",
+        help=("Output file name. The output will be saved as a JSON file."))
 
     opts = parser.parse_args()
     if opts.chain_mapping is not None:
@@ -69,17 +74,51 @@ def _ReadStructureFile(path):
     :returns: Entity
     :rtype: :class:`~ost.mol.EntityHandle`
     """
+    entities = list()
     if not os.path.isfile(path):
         raise IOError("%s is not a file" % path)
     try:
-        ent = LoadPDB(path, profile="SLOPPY")
-    except IOError:
+        entity = LoadPDB(path)
+        if not entity.IsValid():
+            raise IOError("Provided file does not contain valid entity.")
+        entity.SetName(os.path.basename(path))
+        entities.append(entity)
+    except Exception:
         try:
-            ent = LoadMMCIF(path, profile="SLOPPY")
-        except IOError as err:
-            raise err
-
-    return ent
+            tmp_entity, cif_info = LoadMMCIF(path, info=True)
+            if len(cif_info.biounits) == 0:
+                tbu = MMCifInfoBioUnit()
+                tbu.id = 'ASU of ' + entity.pdb_id
+                tbu.details = 'asymmetric unit'
+                for chain in tmp_entity.chains:
+                    tbu.AddChain(str(chain))
+                tinfo = MMCifInfo()
+                tops = MMCifInfoTransOp()
+                tinfo.AddOperation(tops)
+                tbu.AddOperations(tinfo.GetOperations())
+                entity = tbu.PDBize(tmp_entity, min_polymer_size=0)
+                entity.SetName(os.path.basename(path) + ".au")
+                entities.append(entity)
+            elif len(cif_info.biounits) > 1:
+                for i, biounit in enumerate(cif_info.biounits):
+                    entity = biounit.PDBize(tmp_entity, min_polymer_size=0)
+                    if not entity.IsValid():
+                        raise IOError(
+                            "Provided file does not contain valid entity.")
+                    entity.SetName(os.path.basename(path) + "." + str(i))
+                    entities.append(entity)
+            else:
+                biounit = cif_info.biounits[0]
+                entity = biounit.PDBize(tmp_entity, min_polymer_size=0)
+                if not entity.IsValid():
+                    raise IOError(
+                        "Provided file does not contain valid entity.")
+                entity.SetName(os.path.basename(path))
+                entities.append(entity)
+
+        except Exception as exc:
+            raise exc
+    return entities
 
 
 def _Main():
@@ -90,24 +129,56 @@ def _Main():
     #
     # Read the input files
     ost.LogInfo("Reading model from %s" % opts.model)
-    model = _ReadStructureFile(opts.model)
+    models = _ReadStructureFile(opts.model)
     ost.LogInfo("Reading reference from %s" % opts.reference)
-    reference = _ReadStructureFile(opts.reference)
+    references = _ReadStructureFile(opts.reference)
 
+    if len(models) > 1 or len(references) > 1:
+        ost.LogInfo(
+            "Multiple complexes detected. All combinations will be tried.")
+
+    result = {"result": []}
     #
     # Perform scoring
-    try:
-        qs_scorer = qsscoring.QSscorer(reference,
-                                       model)
-        if opts.chain_mapping is not None:
-            ost.LogInfo(
-                "Using custom chain mapping: %s" % str(opts.chain_mapping))
-            qs_scorer.chain_mapping = opts.chain_mapping
-        ost.LogScript('QSscore:', str(qs_scorer.global_score))
-        ost.LogScript('Chain mapping used:', str(qs_scorer.chain_mapping))
-    except qsscoring.QSscoreError as ex:
-        # default handling: report failure and set score to 0
-        ost.LogError('QSscore failed:', str(ex))
+    for model in models:
+        for reference in references:
+            ost.LogInfo("#\nComparing %s to %s" % (
+                model.GetName(),
+                reference.GetName()))
+            try:
+                qs_scorer = qsscoring.QSscorer(reference,
+                                               model)
+                if opts.chain_mapping is not None:
+                    ost.LogInfo(
+                        "Using custom chain mapping: %s" % str(
+                            opts.chain_mapping))
+                    qs_scorer.chain_mapping = opts.chain_mapping
+                result["result"].append({
+                    "status": "SUCCESS",
+                    "error": "",
+                    "model_name": model.GetName(),
+                    "reference_name": reference.GetName(),
+                    "global_score": qs_scorer.global_score,
+                    "lddt_score": qs_scorer.lddt_score,
+                    "best_score": qs_scorer.best_score,
+                    "chain_mapping": qs_scorer.chain_mapping
+                })
+            except qsscoring.QSscoreError as ex:
+                # default handling: report failure and set score to 0
+                ost.LogError('QSscore failed:', str(ex))
+                result["result"].append({
+                    "status": "FAILURE",
+                    "error": str(ex),
+                    "model_name": model.GetName(),
+                    "reference_name": reference.GetName(),
+                    "global_score": 0.0,
+                    "lddt_score": 0.0,
+                    "best_score": 0.0,
+                    "chain_mapping": None
+                })
+    if opts.output is not None:
+        with open(opts.output, "w") as outfile:
+            outfile.write(json.dumps(result, indent=4))
 
 
 if __name__ == '__main__':

[Gerardo Tauriello]

Two comments:

• the lddt scores should be in the other action for SCHWED-3120 (as we discussed we might drop the QS-only action completely but that's a discussion for later)
• there is a shortcut for "outfile.write(json.dumps(result, ...))" called "json.dump(result, outfile, ...)"