Add a flag to compress large output files using zstandard

Inspired by #544 (thank you @ntnn19), but with some differences:
1. Uses zstandard instead of gzip for faster compression times and better compression ratios.
2. Compresses only large files.
3. Doesn't compress already compressed files (distogram and embeddings).
4. Doesn't write large temporary files, compresses directly.

PiperOrigin-RevId: 856630504
Change-Id: I4413e26c9d2e8e4592c372dc675dcfa6044fc257

Author: Augustin-Zidek

run_alphafold.py

@@ -369,6 +369,13 @@
     ' and is non-empty. Useful to set this to True to run the data pipeline and'
     ' the inference separately, but use the same output directory.',
 )
+_COMPRESS_LARGE_OUTPUT_FILES = flags.DEFINE_bool(
+    'compress_large_output_files',
+    False,
+    'If True, compresses the output mmCIF and confidences JSON files (the two'
+    ' largest files) using zstandard. Note that embeddings and distogram, if'
+    ' saved, are already stored in a compressed format.',
+)
 
 
 def make_model_config(
@@ -594,6 +601,7 @@ def write_outputs(
     all_inference_results: Sequence[ResultsForSeed],
     output_dir: os.PathLike[str] | str,
     job_name: str,
+    compress_large_output_files: bool = False,
 ) -> None:
   """Writes outputs to the specified output directory."""
   ranking_scores = []
@@ -614,6 +622,7 @@ def write_outputs(
           inference_result=result,
           output_dir=sample_dir,
           name=f'{job_name}_seed-{seed}_sample-{sample_idx}',
+          compress=compress_large_output_files,
       )
       ranking_score = float(result.metadata['ranking_score'])
       ranking_scores.append((seed, sample_idx, ranking_score))
@@ -646,6 +655,7 @@ def write_outputs(
         # The output terms of use are the same for all seeds/samples.
         terms_of_use=output_terms,
         name=job_name,
+        compress=compress_large_output_files,
     )
     # Save csv of ranking scores with seeds and sample indices, to allow easier
     # comparison of ranking scores across different runs.
@@ -689,6 +699,7 @@ def process_fold_input(
     conformer_max_iterations: int | None = None,
     resolve_msa_overlaps: bool = True,
     force_output_dir: bool = False,
+    compress_large_output_files: bool = False,
 ) -> folding_input.Input:
   ...
 
@@ -705,6 +716,7 @@ def process_fold_input(
     conformer_max_iterations: int | None = None,
     resolve_msa_overlaps: bool = True,
     force_output_dir: bool = False,
+    compress_large_output_files: bool = False,
 ) -> Sequence[ResultsForSeed]:
   ...
 
@@ -720,6 +732,7 @@ def process_fold_input(
     conformer_max_iterations: int | None = None,
     resolve_msa_overlaps: bool = True,
     force_output_dir: bool = False,
+    compress_large_output_files: bool = False,
 ) -> folding_input.Input | Sequence[ResultsForSeed]:
   """Runs data pipeline and/or inference on a single fold input.
 
@@ -750,6 +763,8 @@ def process_fold_input(
       existing one is non-empty. Instead use the existing output directory and
       potentially overwrite existing files. If False, create a new timestamped
       output directory instead if the existing one is non-empty.
+    compress_large_output_files: If True, compress large output files (mmCIF and
+      confidences JSON) using zstandard.
 
   Returns:
     The processed fold input, or the inference results for each seed.
@@ -806,6 +821,7 @@ def process_fold_input(
         all_inference_results=all_inference_results,
         output_dir=output_dir,
         job_name=fold_input.sanitised_name(),
+        compress_large_output_files=compress_large_output_files,
     )
     output = all_inference_results
 
@@ -970,6 +986,7 @@ def main(_):
         conformer_max_iterations=_CONFORMER_MAX_ITERATIONS.value,
         resolve_msa_overlaps=_RESOLVE_MSA_OVERLAPS.value,
         force_output_dir=_FORCE_OUTPUT_DIR.value,
+        compress_large_output_files=_COMPRESS_LARGE_OUTPUT_FILES.value,
     )
     num_fold_inputs += 1
 

src/alphafold3/model/post_processing.py

@@ -19,6 +19,7 @@
 from alphafold3.model import mmcif_metadata
 from alphafold3.model import model
 import numpy as np
+import zstandard
 
 
 @dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
@@ -91,23 +92,34 @@ def write_output(
     output_dir: os.PathLike[str] | str,
     terms_of_use: str | None = None,
     name: str | None = None,
+    compress: bool = False,
 ) -> None:
   """Writes processed inference result to a directory."""
   processed_result = post_process_inference_result(inference_result)
 
   prefix = f'{name}_' if name is not None else ''
 
-  with open(os.path.join(output_dir, f'{prefix}model.cif'), 'wb') as f:
-    f.write(processed_result.cif)
+  if compress:
+    opener = zstandard.open
+    path_transform = lambda path: f'{path}.zst'
+  else:
+    opener = open
+    path_transform = lambda path: path
 
-  with open(
-      os.path.join(output_dir, f'{prefix}summary_confidences.json'), 'wb'
-  ) as f:
-    f.write(processed_result.structure_confidence_summary_json)
+  mmcif_path = os.path.join(output_dir, f'{prefix}model.cif')
+  with opener(path_transform(mmcif_path), 'wb') as f:
+    f.write(processed_result.cif)
 
-  with open(os.path.join(output_dir, f'{prefix}confidences.json'), 'wb') as f:
+  full_confidences_path = os.path.join(output_dir, f'{prefix}confidences.json')
+  with opener(path_transform(full_confidences_path), 'wb') as f:
     f.write(processed_result.structure_full_data_json)
 
+  summary_confidences_path = os.path.join(
+      output_dir, f'{prefix}summary_confidences.json'
+  )
+  with open(summary_confidences_path, 'wb') as f:
+    f.write(processed_result.structure_confidence_summary_json)
+
   if terms_of_use is not None:
     with open(os.path.join(output_dir, 'TERMS_OF_USE.md'), 'wt') as f:
       f.write(terms_of_use)