From 0341fb7f244779a634702d7cb6e37e540e2db26a Mon Sep 17 00:00:00 2001
From: Augustin Zidek <augustinzidek@google.com>
Date: Tue, 12 Nov 2024 16:14:32 +0000
Subject: [PATCH] Documentation fixes and silence irrelevant warning

PiperOrigin-RevId: 695740242
---
 docs/input.md                           | 14 +++++++-------
 docs/performance.md                     |  4 ++--
 src/alphafold3/model/confidences.py     | 22 ++++++++++++++--------
 src/alphafold3/model/post_processing.py | 10 +---------
 4 files changed, 24 insertions(+), 26 deletions(-)

diff --git a/docs/input.md b/docs/input.md
index f506a85..28a626c 100644
--- a/docs/input.md
+++ b/docs/input.md
@@ -494,15 +494,15 @@ glycan to the protein residue, and all bonds that are within the glycan between
 its individual chemical components.
 
 For example, to define the following glycan composed of 4 components (CMP1,
-CMP2, CMP3, CMP4) bound to an arginine in a protein chain A:
+CMP2, CMP3, CMP4) bound to an asparagine in a protein chain A:
 
 ```
  â‹®
-ALA              CMP4
- |                |
-ARG --- CMP1 --- CMP2
- |                |
-ALA              CMP3
+ALA            CMP4
+ |              |
+ASN ―― CMP1 ―― CMP2
+ |              |
+ALA            CMP3
  â‹®
 ```
 
@@ -510,7 +510,7 @@ You will need to specify:
 
 1.  Protein chain A.
 2.  Ligand chain B with the 4 components.
-3.  Bonds ARG-CMP1, CMP1-CMP2, CMP2-CMP3, CMP2-CMP4.
+3.  Bonds ASN-CMP1, CMP1-CMP2, CMP2-CMP3, CMP2-CMP4.
 
 ## User-provided CCD
 
diff --git a/docs/performance.md b/docs/performance.md
index 25d3cd8..c656c03 100644
--- a/docs/performance.md
+++ b/docs/performance.md
@@ -70,8 +70,8 @@ them for numerical accuracy and throughput efficiency:
 
 #### NVIDIA A100 (40 GB)
 
-AlphaFold 3 can run on a single NVIDIA A100 (40 GB) with the following
-configuration changes:
+AlphaFold 3 can run on inputs of size up to 4,352 tokens on a single NVIDIA A100
+(40 GB) with the following configuration changes:
 
 1.  Enabling [unified memory](#unified-memory).
 1.  Adjusting `pair_transition_shard_spec` in `model_config.py`:
diff --git a/src/alphafold3/model/confidences.py b/src/alphafold3/model/confidences.py
index 4a2a9f4..c7d7392 100644
--- a/src/alphafold3/model/confidences.py
+++ b/src/alphafold3/model/confidences.py
@@ -9,6 +9,8 @@
 # https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
 
 """Functions for extracting and processing confidences from model outputs."""
+import warnings
+
 from absl import logging
 from alphafold3 import structure
 from alphafold3.constants import residue_names
@@ -321,7 +323,7 @@ def chain_pair_pde(
 def weighted_nanmean(
     value: np.ndarray, mask: np.ndarray, axis: int
 ) -> np.ndarray:
-  """Nan-mean with weighting -- Empty slices return NaN."""
+  """Nan-mean with weighting -- empty slices return NaN."""
   assert mask.shape == value.shape
   assert not np.isnan(mask).all()
 
@@ -329,9 +331,12 @@ def weighted_nanmean(
   # Need to NaN the mask to get the correct denominator weighting.
   mask_with_nan = mask.copy()
   mask_with_nan[nan_idxs] = np.nan
-  return np.nanmean(value * mask_with_nan, axis=axis) / np.nanmean(
-      mask_with_nan, axis=axis
-  )
+  with warnings.catch_warnings():
+    # Mean of empty slice is ok and should return a NaN.
+    warnings.filterwarnings(action='ignore', message='Mean of empty slice')
+    return np.nanmean(value * mask_with_nan, axis=axis) / np.nanmean(
+        mask_with_nan, axis=axis
+    )
 
 
 def chain_pair_pae(
@@ -453,10 +458,11 @@ def reduce_chain_pair(
     weight *= chain_weight[None] * chain_weight[:, None]
     xchain = weighted_nanmean(chain_pair_met, mask=weight, axis=agg_axis)
   elif agg_type == 'min':
-    is_self = np.eye(
-        num_chains,
-    )
-    xchain = np.nanmin(chain_pair_met + 1e8 * is_self, axis=agg_axis)
+    is_self = np.eye(num_chains)
+    with warnings.catch_warnings():
+      # Min over empty slice is ok and should return a NaN.
+      warnings.filterwarnings('ignore', message='All-NaN slice encountered')
+      xchain = np.nanmin(chain_pair_met + 1e8 * is_self, axis=agg_axis)
   else:
     raise ValueError(f'Unknown aggregation method: {agg_type}')
 
diff --git a/src/alphafold3/model/post_processing.py b/src/alphafold3/model/post_processing.py
index ec34090..d8607b9 100644
--- a/src/alphafold3/model/post_processing.py
+++ b/src/alphafold3/model/post_processing.py
@@ -8,15 +8,7 @@
 # if received directly from Google. Use is subject to terms of use available at
 # https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
 
-"""Post-processing utilities for AlphaFold inference results.
-
-Post-processing includes:
-  - Adding watermark to the predicted structure.
-  - Adding mmCIF metadata fields.
-  - Calculating confidences.
-  - Compressing the resulting artifacts.
-  - Writing the artifacts to blobstore or returning raw result.
-"""
+"""Post-processing utilities for AlphaFold inference results."""
 
 import dataclasses
 import datetime
-- 
GitLab