Skip to content
Snippets Groups Projects
Commit 0341fb7f authored by Augustin Zidek's avatar Augustin Zidek
Browse files

Documentation fixes and silence irrelevant warning

PiperOrigin-RevId: 695740242
parent e2afc41e
No related branches found
No related tags found
1 merge request!1Cloned AlphaFold 3 repo into this one
...@@ -494,15 +494,15 @@ glycan to the protein residue, and all bonds that are within the glycan between ...@@ -494,15 +494,15 @@ glycan to the protein residue, and all bonds that are within the glycan between
its individual chemical components. its individual chemical components.
For example, to define the following glycan composed of 4 components (CMP1, 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 ALA CMP4
| | | |
ARG --- CMP1 --- CMP2 ASN ―― CMP1 ―― CMP2
| | | |
ALA CMP3 ALA CMP3
``` ```
...@@ -510,7 +510,7 @@ You will need to specify: ...@@ -510,7 +510,7 @@ You will need to specify:
1. Protein chain A. 1. Protein chain A.
2. Ligand chain B with the 4 components. 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 ## User-provided CCD
......
...@@ -70,8 +70,8 @@ them for numerical accuracy and throughput efficiency: ...@@ -70,8 +70,8 @@ them for numerical accuracy and throughput efficiency:
#### NVIDIA A100 (40 GB) #### NVIDIA A100 (40 GB)
AlphaFold 3 can run on a single NVIDIA A100 (40 GB) with the following AlphaFold 3 can run on inputs of size up to 4,352 tokens on a single NVIDIA A100
configuration changes: (40 GB) with the following configuration changes:
1. Enabling [unified memory](#unified-memory). 1. Enabling [unified memory](#unified-memory).
1. Adjusting `pair_transition_shard_spec` in `model_config.py`: 1. Adjusting `pair_transition_shard_spec` in `model_config.py`:
......
...@@ -9,6 +9,8 @@ ...@@ -9,6 +9,8 @@
# https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md # https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
"""Functions for extracting and processing confidences from model outputs.""" """Functions for extracting and processing confidences from model outputs."""
import warnings
from absl import logging from absl import logging
from alphafold3 import structure from alphafold3 import structure
from alphafold3.constants import residue_names from alphafold3.constants import residue_names
...@@ -321,7 +323,7 @@ def chain_pair_pde( ...@@ -321,7 +323,7 @@ def chain_pair_pde(
def weighted_nanmean( def weighted_nanmean(
value: np.ndarray, mask: np.ndarray, axis: int value: np.ndarray, mask: np.ndarray, axis: int
) -> np.ndarray: ) -> np.ndarray:
"""Nan-mean with weighting -- Empty slices return NaN.""" """Nan-mean with weighting -- empty slices return NaN."""
assert mask.shape == value.shape assert mask.shape == value.shape
assert not np.isnan(mask).all() assert not np.isnan(mask).all()
...@@ -329,9 +331,12 @@ def weighted_nanmean( ...@@ -329,9 +331,12 @@ def weighted_nanmean(
# Need to NaN the mask to get the correct denominator weighting. # Need to NaN the mask to get the correct denominator weighting.
mask_with_nan = mask.copy() mask_with_nan = mask.copy()
mask_with_nan[nan_idxs] = np.nan mask_with_nan[nan_idxs] = np.nan
return np.nanmean(value * mask_with_nan, axis=axis) / np.nanmean( with warnings.catch_warnings():
mask_with_nan, axis=axis # 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( def chain_pair_pae(
...@@ -453,10 +458,11 @@ def reduce_chain_pair( ...@@ -453,10 +458,11 @@ def reduce_chain_pair(
weight *= chain_weight[None] * chain_weight[:, None] weight *= chain_weight[None] * chain_weight[:, None]
xchain = weighted_nanmean(chain_pair_met, mask=weight, axis=agg_axis) xchain = weighted_nanmean(chain_pair_met, mask=weight, axis=agg_axis)
elif agg_type == 'min': elif agg_type == 'min':
is_self = np.eye( is_self = np.eye(num_chains)
num_chains, with warnings.catch_warnings():
) # Min over empty slice is ok and should return a NaN.
xchain = np.nanmin(chain_pair_met + 1e8 * is_self, axis=agg_axis) warnings.filterwarnings('ignore', message='All-NaN slice encountered')
xchain = np.nanmin(chain_pair_met + 1e8 * is_self, axis=agg_axis)
else: else:
raise ValueError(f'Unknown aggregation method: {agg_type}') raise ValueError(f'Unknown aggregation method: {agg_type}')
......
...@@ -8,15 +8,7 @@ ...@@ -8,15 +8,7 @@
# if received directly from Google. Use is subject to terms of use available at # 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 # https://github.com/google-deepmind/alphafold3/blob/main/WEIGHTS_TERMS_OF_USE.md
"""Post-processing utilities for AlphaFold inference results. """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.
"""
import dataclasses import dataclasses
import datetime import datetime
......
0% Loading or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment