Outsource PAE file into Zip archive

.gitignore

+# Don't have Emac's backup files
+*~
+
+# ignore some files used for testing
+1ake.1.pdb
+3lre.3.A.pdb
+6xne.pdb

convert_to_modelcif.py

@@ -10,6 +10,7 @@ import os
 import pickle
 import re
 import sys
+import zipfile
 
 from Bio import SeqIO
 from Bio.PDB import PDBParser, PPBuilder
@@ -18,6 +19,7 @@ from absl import app, flags, logging
 import numpy as np
 
 import modelcif
+import modelcif.associated
 import modelcif.dumper
 import modelcif.model
 
@@ -116,7 +118,7 @@ class _Biopython2ModelCIF(modelcif.model.AbInitioModel):
                 occupancy=atm.occupancy,
             )
 
-    def add_scores(self, scores_json, entry_id):
+    def add_scores(self, scores_json, entry_id, file_prefix):
         """Add QA metrics"""
         # global scores
         self.qa_metrics.extend(
@@ -171,6 +173,31 @@ class _Biopython2ModelCIF(modelcif.model.AbInitioModel):
                 i += 1
         self.qa_metrics.extend(lpae)
 
+        # outsource PAE to associated file
+        arc_files = [
+            modelcif.associated.LocalPairwiseQAScoresFile(
+                f"{file_prefix}_local_pairwise_qa.cif",
+                categories=["_ma_qa_metric_local_pairwise"],
+                copy_categories=["_ma_qa_metric"],
+                entry_id=entry_id,
+                entry_details="This file is an associated file consisting "
+                + "of local pairwise QA metrics. This is a partial mmCIF "
+                + "file and can be validated by merging with the main "
+                + "mmCIF file containing the model coordinates and other "
+                + "associated data.",
+                details="Predicted aligned error",
+            )
+        ]
+
+        return modelcif.associated.Repository(
+            "",
+            [
+                modelcif.associated.ZipFile(
+                    f"{file_prefix}.zip", files=arc_files
+                )
+            ],
+        )
+
 
 def _get_modelcif_entities(target_ents, asym_units, system):
     """Create ModelCIF entities and asymmetric units."""
@@ -225,9 +252,8 @@ def _store_as_modelcif(
     )
 
     # process scores
-    # system.repositories.append(
-    model.add_scores(data_json, system.id)
-    # )
+    mdl_file = os.path.splitext(os.path.basename(mdl_file))[0]
+    system.repositories.append(model.add_scores(data_json, system.id, mdl_file))
 
     system.model_groups.append(modelcif.model.ModelGroup([model]))
 
@@ -257,12 +283,19 @@ def _store_as_modelcif(
     os.chdir(out_dir)
     try:
         with open(
-            f"{os.path.splitext(os.path.basename(mdl_file))[0]}.cif",
+            f"{mdl_file}.cif",
             "w",
             encoding="ascii",
         ) as mmcif_fh:
             modelcif.dumper.write(mmcif_fh, [system])
-        # _package_associated_files(system.repositories[0])
+        # Create associated archive
+        for archive in system.repositories[0].files:
+            with zipfile.ZipFile(
+                archive.path, "w", zipfile.ZIP_BZIP2
+            ) as cif_zip:
+                for zfile in archive.files:
+                    cif_zip.write(zfile.path, arcname=zfile.path)
+                    os.remove(zfile.path)
         # if compress:
         #    _compress_cif_file(mdl_fle)
     finally: