mcif to pdb conversion with gemmi

utils/pipelines.py

@@ -7,6 +7,7 @@ import spaces
 import subprocess
 from Bio.PDB import MMCIFParser, PDBIO
 import gzip
+import gemmi
 
 
 
@@ -72,7 +73,7 @@ def unconditional_generation(num_batches, num_designs_per_batch, length):
         for batch in range(num_batches):
             for design in range(num_designs_per_batch):
                 file_name = os.path.join(directory, f"_{batch}_model_{design}.cif.gz")
-                results.append({"batch": batch, "design": design, "file": file_name, "pdb": cif_gz_to_pdb(file_name)})
+                results.append({"batch": batch, "design": design, "file": file_name, "pdb": mcif_gz_to_pdb_string_gemmi(file_name)})
             
         print(results)
         return directory, results
@@ -89,18 +90,18 @@ def collect_outputs(gen_directory, num_batches, num_designs_per_batch):
         return f"Error: {str(e)}"
     
 
-def cif_gz_to_pdb(cif_gz_path):
-    """Convert .cif.gz to PDB string for viewer."""
-    # Decompress & parse
-    parser = MMCIFParser(QUIET=True)
-    with gzip.open(cif_gz_path, 'rt') as f:
-        struct = parser.get_structure('model', f)
+def mcif_gz_to_pdb_string_gemmi(file_path: str) -> str:
+    """
+    Converts a .mcif.gz file to a PDB-formatted string.
     
-    # Write to string
-    io = PDBIO()
-    io.set_structure(struct)
-    pdb_lines = []
-    class StringIO:
-        def write(self, s): pdb_lines.append(s)
-    io.save(StringIO())
-    return ''.join(pdb_lines)
+    Args:
+        file_path (str): Path to the .mcif.gz file.
+    
+    Returns:
+        str: PDB content as string.
+    
+    Requires: pip install gemmi
+    """
+    st = gemmi.read_structure(file_path)
+    st.setup_entities()  # Recommended for consistent entity handling [web:18]
+    return st.make_pdb_string(gemmi.PdbWriteOptions(minimal=True))