Add a --write_distogram flag to save the distogram in .npz (AF3 mimic)

docs/prediction.md

@@ -170,6 +170,7 @@ Examples of common options include:
 | `--use_potentials`        | `FLAG`          | `False`                     | Whether to run the original Boltz-2 model using inference time potentials.                                                                                                        |
 | `--write_full_pae`       | `FLAG`          | `False`                     | Whether to save the full PAE matrix as a file.                                                                                                                                      |
 | `--write_full_pde`       | `FLAG`          | `False`                     | Whether to save the full PDE matrix as a file.                                                                                                                                      |
+| `--write_distogram`      | `FLAG`          | `False`                     | Whether to save the distogram logits as a file.                                                                                                                                     |
 
 ## Output
 
@@ -186,6 +187,7 @@ out_dir/
         ├── pae_[input_file1]_model_0.npz                      # The predicted PAE score for every pair of tokens
         ├── pde_[input_file1]_model_0.npz                      # The predicted PDE score for every pair of tokens
         ├── plddt_[input_file1]_model_0.npz                    # The predicted pLDDT score for every token
+        ├── distogram_[input_file1].npz                        # The predicted distogram logits for every pair of tokens
         ...
         └── [input_file1]_model_[diffusion_samples-1].cif      # The predicted structure in CIF format
         ...
@@ -248,6 +250,18 @@ The `affinity_pred_value` aims to measure the specific affinity of different bin
 
 You can convert the model's output to pIC50 in `kcal/mol` by using `y --> (6 - y) * 1.364` where `y` is the model's prediction.
 
+The distogram `.npz` file, written when `--write_distogram` is set, contains the raw logits of the distogram head together with the annotation of its token axes:
+
+| **Array**           | **Shape**              | **Description**                                                                                                              |
+|---------------------|------------------------|------------------------------------------------------------------------------------------------------------------------------|
+| `distogram_logits`  | `(num_tokens, num_tokens, num_bins)` | The unnormalised logits over distance bins, as `float16`. The bins are `num_bins - 1` boundaries evenly spaced over 2 to 22 Å, with a first and a last bin catching the distances below and above that range. Apply a softmax over the last axis to obtain distance probabilities. |
+| `asym_id`           | `(num_tokens,)`        | The chain each token belongs to.                                                                                             |
+| `residue_index`     | `(num_tokens,)`        | The residue index of each token.                                                                                             |
+| `entity_id`         | `(num_tokens,)`        | The entity each token belongs to.                                                                                            |
+| `mol_type`          | `(num_tokens,)`        | The molecule type of each token.                                                                                             |
+
+The distogram is a trunk output, shared by all diffusion samples of a prediction, so a single file is written per input rather than one per sample.
+
 
 ## Authentication to MSA Server
 

src/boltz/data/write/writer.py

@@ -24,6 +24,7 @@ def __init__(
         output_format: Literal["pdb", "mmcif"] = "mmcif",
         boltz2: bool = False,
         write_embeddings: bool = False,
+        write_distogram: bool = False,
     ) -> None:
         """Initialize the writer.
 
@@ -45,6 +46,7 @@ def __init__(
         self.boltz2 = boltz2
         self.output_dir.mkdir(parents=True, exist_ok=True)
         self.write_embeddings = write_embeddings
+        self.write_distogram = write_distogram
 
     def write_on_batch_end(
         self,
@@ -79,7 +81,9 @@ def write_on_batch_end(
             idx_to_rank = {i: i for i in range(len(records))}
 
         # Iterate over the records
-        for record, coord, pad_mask in zip(records, coords, pad_masks):
+        for record_idx, (record, coord, pad_mask) in enumerate(
+            zip(records, coords, pad_masks)
+        ):
             # Load the structure
             path = self.data_dir / f"{record.id}.npz"
             if self.boltz2:
@@ -246,6 +250,11 @@ def write_on_batch_end(
                     )
                     np.savez_compressed(path, pde=pde.cpu().numpy())
 
+            # Save distogram
+            if self.write_distogram and "pdistogram" in prediction:
+                path = struct_dir / f"distogram_{record.id}.npz"
+                self._write_distogram(path, prediction, batch, record_idx)
+
             # Save embeddings
             if self.write_embeddings and "s" in prediction and "z" in prediction:
                 s = prediction["s"].cpu().numpy()
@@ -257,6 +266,46 @@ def write_on_batch_end(
                 )
                 np.savez_compressed(path, s=s, z=z)
 
+    def _write_distogram(
+        self,
+        path: Path,
+        prediction: dict[str, Tensor],
+        batch: dict[str, Tensor],
+        record_idx: int,
+    ) -> None:
+        """Save the raw distogram logits of a record.
+
+        Parameters
+        ----------
+        path : Path
+            The file to write the distogram to.
+        prediction : dict[str, Tensor]
+            The predictions of the batch.
+        batch : dict[str, Tensor]
+            The features of the batch.
+        record_idx : int
+            The index of the record within the batch.
+
+        """
+        logits = prediction["pdistogram"][record_idx]
+
+        # Boltz-2 predicts a stack of distograms, Boltz-1 a single one
+        if logits.dim() == 4 and logits.shape[2] == 1:  # noqa: PLR2004
+            logits = logits[:, :, 0, :]
+
+        # Restrict to the tokens of the record
+        token_mask = batch["token_pad_mask"][record_idx].bool()
+        logits = logits[token_mask][:, token_mask].float()
+
+        # Token annotation of the distogram axes
+        arrays = {"distogram_logits": logits.cpu().numpy().astype(np.float16)}
+        for key in ("asym_id", "residue_index", "entity_id", "mol_type"):
+            if key in batch:
+                value = batch[key][record_idx][token_mask]
+                arrays[key] = value.cpu().numpy().astype(np.int32)
+
+        np.savez_compressed(path, **arrays)
+
     def on_predict_epoch_end(
         self,
         trainer: Trainer,  # noqa: ARG002

src/boltz/main.py

@@ -1039,6 +1039,14 @@ def cli() -> None:
     is_flag=True,
     help=" to dump the s and z embeddings into a npz file. Default is False.",
 )
+@click.option(
+    "--write_distogram",
+    is_flag=True,
+    help=(
+        "Whether to dump the distogram logits into a npz file. "
+        "Default is False."
+    ),
+)
 def predict(  # noqa: C901, PLR0915, PLR0912
     data: str,
     out_dir: str,
@@ -1077,6 +1085,7 @@ def predict(  # noqa: C901, PLR0915, PLR0912
     num_subsampled_msa: int = 1024,
     no_kernels: bool = False,
     write_embeddings: bool = False,
+    write_distogram: bool = False,
 ) -> None:
     """Run predictions with Boltz."""
     # If cpu, write a friendly warning
@@ -1250,6 +1259,7 @@ def predict(  # noqa: C901, PLR0915, PLR0912
         output_format=output_format,
         boltz2=model == "boltz2",
         write_embeddings=write_embeddings,
+        write_distogram=write_distogram,
     )
 
     # Set up trainer
@@ -1304,6 +1314,7 @@ def predict(  # noqa: C901, PLR0915, PLR0912
             "write_confidence_summary": True,
             "write_full_pae": write_full_pae,
             "write_full_pde": write_full_pde,
+            "write_distogram": write_distogram,
         }
 
         steering_args = BoltzSteeringParams()

src/boltz/model/models/boltz1.py

@@ -1165,6 +1165,8 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
             pred_dict["coords"] = out["sample_atom_coords"]
             pred_dict["s"] = out["s"]
             pred_dict["z"] = out["z"]
+            if self.predict_args.get("write_distogram", False):
+                pred_dict["pdistogram"] = out["pdistogram"]
             if self.predict_args.get("write_confidence_summary", True):
                 pred_dict["confidence_score"] = (
                     4 * out["complex_plddt"]

src/boltz/model/models/boltz2.py

@@ -1074,6 +1074,8 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> d
             pred_dict["token_masks"] = batch["token_pad_mask"]
             pred_dict["s"] = out["s"]
             pred_dict["z"] = out["z"]
+            if self.predict_args.get("write_distogram", False):
+                pred_dict["pdistogram"] = out["pdistogram"]
 
             if "keys_dict_out" in self.predict_args:
                 for key in self.predict_args["keys_dict_out"]:

tests/data/write/test_writer.py

@@ -0,0 +1,90 @@
+import tempfile
+import unittest
+from pathlib import Path
+
+import numpy as np
+import torch
+
+from boltz.data.write.writer import BoltzWriter
+
+
+class WriteDistogramTest(unittest.TestCase):
+    def setUp(self):
+        self.num_tokens = 7
+        self.num_padded = 10
+        self.num_bins = 64
+
+        self.writer = BoltzWriter(
+            data_dir=tempfile.mkdtemp(),
+            output_dir=tempfile.mkdtemp(),
+            write_distogram=True,
+        )
+
+        pad = self.num_padded - self.num_tokens
+        self.batch = {
+            "token_pad_mask": torch.tensor([[1.0] * self.num_tokens + [0.0] * pad]),
+            "asym_id": torch.tensor([[0] * 4 + [1] * 3 + [0] * pad]),
+            "residue_index": torch.arange(self.num_padded)[None],
+            "entity_id": torch.zeros(1, self.num_padded, dtype=torch.long),
+            "mol_type": torch.zeros(1, self.num_padded, dtype=torch.long),
+        }
+
+    def _write(self, logits):
+        with tempfile.TemporaryDirectory() as tmp_dir:
+            path = Path(tmp_dir) / "distogram.npz"
+            self.writer._write_distogram(  # noqa: SLF001
+                path, {"pdistogram": logits}, self.batch, 0
+            )
+            return dict(np.load(path))
+
+    def test_boltz1_shape(self):
+        """A single distogram is written unchanged."""
+        logits = torch.randn(1, self.num_padded, self.num_padded, self.num_bins)
+        arrays = self._write(logits)
+
+        expected = logits[0, : self.num_tokens, : self.num_tokens].numpy()
+        self.assertEqual(
+            arrays["distogram_logits"].shape,
+            (self.num_tokens, self.num_tokens, self.num_bins),
+        )
+        np.testing.assert_array_equal(
+            arrays["distogram_logits"], expected.astype(np.float16)
+        )
+
+    def test_boltz2_shape(self):
+        """A stack holding a single distogram loses its stacking axis."""
+        logits = torch.randn(1, self.num_padded, self.num_padded, 1, self.num_bins)
+        arrays = self._write(logits)
+
+        expected = logits[0, : self.num_tokens, : self.num_tokens, 0].numpy()
+        self.assertEqual(
+            arrays["distogram_logits"].shape,
+            (self.num_tokens, self.num_tokens, self.num_bins),
+        )
+        np.testing.assert_array_equal(
+            arrays["distogram_logits"], expected.astype(np.float16)
+        )
+
+    def test_distogram_stack_is_kept(self):
+        """A stack holding several distograms is written whole."""
+        logits = torch.randn(1, self.num_padded, self.num_padded, 3, self.num_bins)
+        arrays = self._write(logits)
+
+        self.assertEqual(
+            arrays["distogram_logits"].shape,
+            (self.num_tokens, self.num_tokens, 3, self.num_bins),
+        )
+
+    def test_token_annotation(self):
+        """The token annotation covers the unpadded tokens only."""
+        logits = torch.randn(1, self.num_padded, self.num_padded, 1, self.num_bins)
+        arrays = self._write(logits)
+
+        for key in ("asym_id", "residue_index", "entity_id", "mol_type"):
+            self.assertEqual(arrays[key].shape, (self.num_tokens,))
+        np.testing.assert_array_equal(arrays["asym_id"], [0, 0, 0, 0, 1, 1, 1])
+        np.testing.assert_array_equal(arrays["residue_index"], np.arange(7))
+
+
+if __name__ == "__main__":
+    unittest.main()