Commit Graph

187 Commits

Author SHA1 Message Date
0ea324ce74 Validate that a non-loop CIF key has only a single value
PiperOrigin-RevId: 802455341
Change-Id: I089c4e0a9a52862b5eba6d26023184a4a34efc96
2025-09-03 01:48:24 -07:00
7b816f4035 Add support for protein/DNA/RNA/ligand descriptions
Suggested in https://github.com/google-deepmind/alphafold3/issues/496.

PiperOrigin-RevId: 802081348
Change-Id: I666466fd6a770b6f4a891ed33e6a26651d600c4a
2025-09-02 04:11:53 -07:00
a2b03dab51 Validate the number of tokens in a CIF loop
PiperOrigin-RevId: 797771843
Change-Id: Ib34796c2e9f402cd2cad79d67127ddb32f0803ae
2025-08-21 08:17:40 -07:00
2e2ffc10ab Improve the gpu_device flag description
Based on the suggestion in https://github.com/google-deepmind/alphafold3/issues/493

PiperOrigin-RevId: 792089424
Change-Id: Iea8d0527725c2febbc5e09b619ca7ad1f778ddba
2025-08-07 03:53:37 -07:00
563896c75c Improve comments in CifDict
PiperOrigin-RevId: 791623617
Change-Id: I569b0c65d427f32bbb6a6e50a4791cb774347375
2025-08-06 04:05:16 -07:00
722437b4ff Fix handling of tokens that contain quotes surrounded by spaces in CifDict::ToString
* `A "A" A` must be quoted as `'A "A" A', B 'B' B` must be quoted as `"B 'B' B"`.
* `C "C" 'C' C` must be quoted as `\n;C "C" 'C' C\n;`.
* Previously, `CifDict` didn't quote these correctly and e.g. `A "A" A` would become three tokens because of the bad early return on the first space in the token.

This makes `CifDict::ToString` ~15% faster as the trivial check allows us to skip the extra newline check for most tokens.

Reported in https://github.com/google-deepmind/alphafold3/issues/471.

PiperOrigin-RevId: 791587762
Change-Id: If6616263e570dff534633bd0302afd23471f946f
2025-08-06 02:05:23 -07:00
f2edd59f87 Undo backwards incompatible Triton API change
PiperOrigin-RevId: 790639906
Change-Id: If98343e8a5601baeeb882229a4e4c5dcdc601881
2025-08-04 01:37:39 -07:00
903a931d36 Improve documentation on running data pipeline and inference separately
Prompted by https://github.com/google-deepmind/alphafold3/issues/485.

PiperOrigin-RevId: 789700619
Change-Id: I0e876a1130cc537421e8916fcad1230302196b64
2025-08-01 03:56:58 -07:00
3c4de6bb9b Update the output documentation with case sensitivity and force_output_dir
PiperOrigin-RevId: 789369238
Change-Id: I2ccab4e220300646dba2f8882a82ad1e21c1523f
2025-07-31 09:40:19 -07:00
ec4254ac5d Integrate Triton up to [6af1c4b5](6af1c4b507)
https://github.com/openxla/triton/tree/triton_integrate_branch-1.10

PiperOrigin-RevId: 789217115
Change-Id: Ib89b60601802c247b4dc96cfe6cbff4752a41402
2025-07-31 00:14:38 -07:00
a14376d249 Make the folding input sanitised name case sensitive
PiperOrigin-RevId: 788852391
Change-Id: Ic0c75d8bcae843bacc77b2369b40002aeab50a33
2025-07-30 04:46:12 -07:00
042202363b Suppress Jackhmmer stdout and stderr when checking whether seq_limit is supported
Inspired by https://github.com/google-deepmind/alphafold3/pull/470, thanks @owainkenwayucl.

PiperOrigin-RevId: 782928877
Change-Id: I8cb09f9bd75f174f78340ad72f80b805eb873be8
2025-07-14 09:11:20 -07:00
751a4b8612 Add support for --seq_limit in Jackhmmer - significantly reduces peak RAM use
* See https://github.com/EddyRivasLab/hmmer/issues/323 for more context.
* We observed 75.6 GB -> 13.6 GB peak RAM reduction for one of our queries.

PiperOrigin-RevId: 782808687
Change-Id: I4306dc6921015c88c5f8ced69a4ef46e10574a57
2025-07-14 02:00:15 -07:00
a209008f6f Add validation of template JSON keys
More user-friendly than the error reported in https://github.com/google-deepmind/alphafold3/issues/459.

PiperOrigin-RevId: 778401008
Change-Id: I80da973e2091f4cadee71f49caf2b8a37867556c
2025-07-02 01:53:24 -07:00
53d2acc5b7 Make query sequence logging less verbose in MSA tools
PiperOrigin-RevId: 775647191
Change-Id: I3804586da51525b58c98a91c400523cda160eaf7
2025-06-25 05:36:10 -07:00
6cf825575b Fix the documentation for RNA MSA
Fixes https://github.com/google-deepmind/alphafold3/issues/448

PiperOrigin-RevId: 774840546
Change-Id: Icd7e34c190052aa6da82d2edffa4050f36e00837
2025-06-23 10:43:45 -07:00
0f865a667c Configure tests to explicitly use jax_threefry_partitionable=False.
See https://github.com/jax-ml/jax/discussions/18480

PiperOrigin-RevId: 770183943
Change-Id: I2cea840e26e257b26e8605f759423e476755e10f
2025-06-11 09:26:55 -07:00
3f8871d982 Add empty paired MSA in the full example
Fixes https://github.com/google-deepmind/alphafold3/issues/444

PiperOrigin-RevId: 770122170
Change-Id: Ia82dec2f389cefcdc8c03afe6ec5ac2717b3ac1d
2025-06-11 06:33:43 -07:00
bc73f5652b Remove redundant aliases mmcif.tokenize and mmcif.split_line
PiperOrigin-RevId: 768055969
Change-Id: I7c6b4b47e9ccae15e0ffc220cfb15fa69457181b
2025-06-06 06:51:58 -07:00
344735c371 Check that ccdCodes is a list
Provides a more user-friendly error message than what a user reported in https://github.com/google-deepmind/alphafold3/issues/428.

PiperOrigin-RevId: 767193009
Change-Id: Iaba288b1d22173886967c1339bb0b7d6c07becab
2025-06-04 10:14:32 -07:00
06d28cad2e Update to Ubuntu 24.04 / CUDA 12.6.3 base image and use Python 3.12
PiperOrigin-RevId: 765229057
Change-Id: I7b1c8234fd1b067e62e308116d603f0e3d375972
2025-06-04 05:20:40 -07:00
6cc6364a8d Internal change
PiperOrigin-RevId: 765224124
Change-Id: I1aad4e6d7cee2e27e3f360c1845b4c4111b4eeec
2025-06-04 05:20:24 -07:00
64723739f5 Update to Ubuntu 24.04 / CUDA 12.6.3 base image and use Python 3.12
PiperOrigin-RevId: 765219906
Change-Id: I271f442012ee30356ef316eda87abd231319a673
2025-05-30 09:49:16 -07:00
935d47f8d4 Reduce the size of the Docker/Singularity container
PiperOrigin-RevId: 765209187
Change-Id: I075874b51b243a3e75bc8078ea81920f1984dbc5
2025-05-30 17:12:44 +01:00
5ecdfe883a Make register_dataclass compatible with JAX < v0.4.36
PiperOrigin-RevId: 765135084
Change-Id: Ifc89843c9706c38cbe6a9e2ed3c4f918b2c36954
2025-05-30 04:36:47 -07:00
6a34fea9a2 Remove unnecessary dependency on structure test_utils
PiperOrigin-RevId: 765132514
Change-Id: Ia1acdf55267083741c294864792f25009f5f4542
2025-05-30 04:25:22 -07:00
565f286892 Remove unnecessary chex dependency
PiperOrigin-RevId: 765086591
Change-Id: I34d7e7b83073d84fe083ee091a3b54e90dcdeba3
2025-05-30 01:44:10 -07:00
17afe151ea Improve AlphaFold 3 Dockerfile
* Use `apt-get` instead of `apt`. Should not make much difference since we set the `--quiet` flag though.
* Document better the rationale behind building HMMER from scratch.
* Check SHA-256 hash of the downloaded HMMER tar file.
* Addresses critique from https://www.linkedin.com/pulse/getting-alphafold-3-run-amd-gpus-owain-kenway-egite/.

PiperOrigin-RevId: 763703614
Change-Id: Ibc231b8c31d1d891f537e2765094c1710a76eaea
2025-05-27 02:54:58 -07:00
e64804c49f Add a flag to enable skipping unpaired MSA deduplication against paired MSA
Addresses an issue reported in https://github.com/google-deepmind/alphafold3/issues/404.

PiperOrigin-RevId: 761917258
Change-Id: I26592af0f40481a40ba611130739a40a62e4a6dc
2025-05-22 05:11:12 -07:00
fbabfe25b0 Name the deletion_matrix feature correctly when produced
PiperOrigin-RevId: 761487725
Change-Id: I99435be9cb0d1b043c3441b1f498ed31c7783f23
2025-05-21 05:45:06 -07:00
a20a29e021 Internal change.
PiperOrigin-RevId: 761425868
Change-Id: I1f24f047ad737a6509bc0866e7d86bcb8d725791
2025-05-21 01:58:37 -07:00
62074340cf Validate bonded atoms against the CCD
When user specifies a list of bonds in bondedAtomPairs, make sure the atoms that are specified are actually valid atoms within each of the components.

Addresses https://github.com/google-deepmind/alphafold3/issues/406.

PiperOrigin-RevId: 760603843
Change-Id: Ia1747e47c8ab45557173adff122cf32310c57c34
2025-05-19 06:41:00 -07:00
2d466dc791 Do not limit parallelism of make when building HMMER
PiperOrigin-RevId: 759516856
Change-Id: I1c926a30edf87f0d3e6feccb42e474ff96c06d6d
2025-05-16 02:07:09 -07:00
0cda346e40 Make template featurisation error messages clearer
Users hit this already multiple times:
https://github.com/google-deepmind/alphafold3/issues/317
https://github.com/google-deepmind/alphafold3/issues/416

PiperOrigin-RevId: 759101683
Change-Id: I15c21208f5843e02a73579d4a10310ae26aa8494
2025-05-15 05:18:51 -07:00
38a9333b54 Fix flag descriptions for save_embeddings and save_distogram
PiperOrigin-RevId: 758180480
Change-Id: Iaccb91b1f0be4a6ea6d4b9b77c55d5593a4afd03
2025-05-13 05:52:26 -07:00
267ab05055 Save embeddings as float16 instead of float32 to make them 2x smaller
PiperOrigin-RevId: 758140305
Change-Id: I70cb7adc53440ce19d20fa89d64bbb36b84b13f7
2025-05-13 03:23:02 -07:00
424201003f Add an option to save the distogram
Based on the implementation in https://github.com/google-deepmind/alphafold3/pull/345 by @PasqM.

PiperOrigin-RevId: 758128292
Change-Id: I868d43018833f889d858e5c9a7d687a44cfd3e82
2025-05-13 02:41:03 -07:00
6a0e8b2afe [alphafold3] avoid use of deprecated jax.util APIs
jax.util was deprecated in JAX v0.6.0, and will be removed in JAX v0.7.0.

PiperOrigin-RevId: 756703396
Change-Id: I9e0b0267c4a403939aeb413f59d8ff7ef3ee390b
2025-05-09 04:09:41 -07:00
e274d27978 Internal change.
PiperOrigin-RevId: 748595613
Change-Id: I7e84c89c15c4887b51a66fc708d15c835aa8a0a8
2025-04-17 01:49:55 -07:00
7a4a2f7142 Mention unresolved residues in template documentation
Fixes https://github.com/google-deepmind/alphafold3/issues/373.

PiperOrigin-RevId: 747342015
Change-Id: I09586c5d00cda84db18764cdc735b057f2b4d749
2025-04-14 03:40:15 -07:00
5fbae81a9e Remove unused code
PiperOrigin-RevId: 744678177
Change-Id: Ib2444fa8ba6c9850f2331d133e038883fc32e681
2025-04-09 14:07:50 +01:00
dfd8b9cb63 Fix test numerics.
PiperOrigin-RevId: 742607438
Change-Id: I6670124aa30f3000a07efb753ccf75c959881bbe
2025-04-01 02:34:34 -07:00
69b749cb71 Use unsigned int for bias pipelining hack in Triton FlashAttention. A bitwise AND of 0xFFFFFFFF with a negative integer leads to overflow, which got caught by the latest Triton integration.
PiperOrigin-RevId: 739868199
Change-Id: I732e36901c18e4a07c8907c52412b7475d2ebb64
2025-03-24 03:01:53 -07:00
8151373a6b Do not deduplicate unpaired MSA in featurisation
If we generated the MSA ourselves, it is already deduplicated. If it is user-provided, keep it as is to prevent destroying desired pairing.

Reported in https://github.com/google-deepmind/alphafold3/issues/331.

PiperOrigin-RevId: 736489145
Change-Id: If0ffe4908e8a4dd93855962b299a7fe84a849119
2025-03-13 06:53:48 -07:00
3f706280fe Make pseudo_beta_fn work with both np and jnp.
Fix kindly suggested in https://github.com/google-deepmind/alphafold3/pull/350/files by @popfido.

PiperOrigin-RevId: 736443698
Change-Id: I4c3c20c2291d445976f88f8f44a7f1ac6d6b3cbb
2025-03-13 03:31:33 -07:00
9c0d146f32 Release the GIL when running CifDict::FromString
PiperOrigin-RevId: 735819832
Change-Id: Ibb4323ef474850bffa2a54ab227d9654091fd07a
2025-03-11 11:06:59 -07:00
b641827717 Make output file names unique for all seeds/samples
PiperOrigin-RevId: 735310700
Change-Id: I4f42ed52cd38a48d189eb6ba4d563117caff18b0
2025-03-10 03:30:01 -07:00
98e03d944f Internal change
PiperOrigin-RevId: 735308466
Change-Id: I9787f324d3357ebb7755987553fc17f988adc232
2025-03-10 03:20:20 -07:00
e9b2f0ce3b Set ref_max_modified_date when running featurisation
Addresses https://github.com/google-deepmind/alphafold3/issues/334.

PiperOrigin-RevId: 735299819
Change-Id: Ib16739fc33ea927b243e4e481582c0a9cdb4128a
2025-03-10 02:43:05 -07:00
2eb2a50b62 Add support for specifying the user-defined CCD using a file
Prompted by https://github.com/google-deepmind/alphafold3/issues/325.

PiperOrigin-RevId: 734476458
Change-Id: I0fa681f6e3cb77284b3a3b4b6f429cef14e1638b
2025-03-07 02:29:39 -08:00