Add files using upload-large-folder tool

.gitattributes

@@ -1,35 +1,85 @@
 *.7z filter=lfs diff=lfs merge=lfs -text
 *.arrow filter=lfs diff=lfs merge=lfs -text
-*.bin filter=lfs diff=lfs merge=lfs -text
+ 
+*.bin.* filter=lfs diff=lfs merge=lfs -text
 *.bz2 filter=lfs diff=lfs merge=lfs -text
-*.ckpt filter=lfs diff=lfs merge=lfs -text
 *.ftz filter=lfs diff=lfs merge=lfs -text
-*.gz filter=lfs diff=lfs merge=lfs -text
+ 
 *.h5 filter=lfs diff=lfs merge=lfs -text
 *.joblib filter=lfs diff=lfs merge=lfs -text
 *.lfs.* filter=lfs diff=lfs merge=lfs -text
-*.mlmodel filter=lfs diff=lfs merge=lfs -text
 *.model filter=lfs diff=lfs merge=lfs -text
 *.msgpack filter=lfs diff=lfs merge=lfs -text
-*.npy filter=lfs diff=lfs merge=lfs -text
-*.npz filter=lfs diff=lfs merge=lfs -text
 *.onnx filter=lfs diff=lfs merge=lfs -text
 *.ot filter=lfs diff=lfs merge=lfs -text
 *.parquet filter=lfs diff=lfs merge=lfs -text
 *.pb filter=lfs diff=lfs merge=lfs -text
-*.pickle filter=lfs diff=lfs merge=lfs -text
-*.pkl filter=lfs diff=lfs merge=lfs -text
 *.pt filter=lfs diff=lfs merge=lfs -text
 *.pth filter=lfs diff=lfs merge=lfs -text
 *.rar filter=lfs diff=lfs merge=lfs -text
-*.safetensors filter=lfs diff=lfs merge=lfs -text
 saved_model/**/* filter=lfs diff=lfs merge=lfs -text
-*.tar.* filter=lfs diff=lfs merge=lfs -text
-*.tar filter=lfs diff=lfs merge=lfs -text
+ 
 *.tflite filter=lfs diff=lfs merge=lfs -text
 *.tgz filter=lfs diff=lfs merge=lfs -text
-*.wasm filter=lfs diff=lfs merge=lfs -text
 *.xz filter=lfs diff=lfs merge=lfs -text
-*.zip filter=lfs diff=lfs merge=lfs -text
+ 
+*.zstandard filter=lfs diff=lfs merge=lfs -text
+*.tfevents* filter=lfs diff=lfs merge=lfs -text
+*.db* filter=lfs diff=lfs merge=lfs -text
+*.ark* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*data* filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.meta filter=lfs diff=lfs merge=lfs -text
+**/*ckpt*.index filter=lfs diff=lfs merge=lfs -text
+*.safetensors filter=lfs diff=lfs merge=lfs -text
+*.ckpt filter=lfs diff=lfs merge=lfs -text
+*.gguf* filter=lfs diff=lfs merge=lfs -text
+*.ggml filter=lfs diff=lfs merge=lfs -text
+*.llamafile* filter=lfs diff=lfs merge=lfs -text
+*.pt2 filter=lfs diff=lfs merge=lfs -text
+*.mlmodel filter=lfs diff=lfs merge=lfs -text
+*.npy filter=lfs diff=lfs merge=lfs -text
+*.npz filter=lfs diff=lfs merge=lfs -text
+ 
+*.pkl filter=lfs diff=lfs merge=lfs -text
+ 
+*.wasm filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+
+inputs/7r6r_data.json filter=lfs diff=lfs merge=lfs -text
+
+flax_model/alphafold3/_data/libcifpp/components.cif filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_data/libcifpp/mmcif_ma.dic filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_data/libcifpp/mmcif_pdbx.dic filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_tools/hmmer/bin/hmmpgmd filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_tools/hmmer/bin/hmmpgmd_shard filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_tools/hmmer/bin/jackhmmer filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_tools/hmmer/bin/nhmmer filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/_tools/hmmer/bin/phmmer filter=lfs diff=lfs merge=lfs -text
+
+*.pickle filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/cpp.cpython-311-x86_64-linux-gnu.so filter=lfs diff=lfs merge=lfs -text
+
+ 
+ 
+*.tar filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/test_data/miniature_databases/nt_rna_2023_02_23_clust_seq_id_90_cov_80_rep_seq__subsampled_1000.fasta filter=lfs diff=lfs merge=lfs -text
+flax_model/alphafold3/test_data/miniature_databases/rnacentral_active_seq_id_90_cov_80_linclust__subsampled_1000.fasta filter=lfs diff=lfs merge=lfs -text
+
+ 
+*.zip filter=lfs diff=lfs merge=lfs -text
+
+weight/_dep/eigen/Eigen/src/misc/lapacke.h filter=lfs diff=lfs merge=lfs -text
+
+ 
+weight/_dep/libcifpp/rsrc/components.cif filter=lfs diff=lfs merge=lfs -text
+weight/_dep/libcifpp/rsrc/mmcif_ma.dic filter=lfs diff=lfs merge=lfs -text
+weight/_dep/libcifpp/rsrc/mmcif_pdbx.dic filter=lfs diff=lfs merge=lfs -text
+
+*.gz filter=lfs diff=lfs merge=lfs -text
+*.tar.* filter=lfs diff=lfs merge=lfs -text
+
+*.bin filter=lfs diff=lfs merge=lfs -text
+weight/libflash_atten_c.so filter=lfs diff=lfs merge=lfs -text
+weight/mmseqs/bin/mmseqs filter=lfs diff=lfs merge=lfs -text
+weight/mmseqs/lib/libmarv.so filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,10 @@
+AlphaFold3 source files in this split retain the license headers from the
+upstream AlphaFold 3 implementation integrated in OneScience.
+
+The copied AlphaFold3 source headers state that AlphaFold 3 source code is
+licensed under CC BY-NC-SA 4.0 and refer to:
+
+https://creativecommons.org/licenses/by-nc-sa/4.0/
+
+Model parameters are not included in this repository and are subject to their
+own terms of use.

README.md

@@ -0,0 +1,277 @@
+---
+license: cc-by-nc-sa-4.0
+tasks:
+  - protein-complex-structure-prediction
+frameworks:
+  - jax
+language:
+  - en
+  - zh
+tags:
+  - OneScience
+  - Life Sciences
+  - Protein Structure Prediction
+  - Biomolecular Interactions
+  - Complex Structure Prediction
+  - AlphaFold3
+datasets:
+  - OneScience-Sugon/AlphaFold3_dataset
+---
+
+<p align="center">
+  <strong>
+    <span style="font-size: 30px;">AlphaFold3</span>
+  </strong>
+</p>
+
+# Model Introduction
+
+AlphaFold 3 is a biomolecular structure prediction model developed by Google DeepMind and Isomorphic Labs. It predicts the three-dimensional structures and interactions of proteins, DNA, RNA, small-molecule ligands, and their complexes.
+
+Paper: Accurate structure prediction of biomolecular interactions with AlphaFold 3  
+https://www.nature.com/articles/s41586-024-07487-w
+
+# Model Description
+
+AlphaFold 3 uses a Pairformer and a diffusion model to predict biomolecular complex structures. This model package provides a JAX / Flax inference implementation and database-search scripts, together with the accompanying Hugging Face dataset `OneScience-Sugon/AlphaFold3_dataset`.
+
+# Use Cases
+
+| Scenario | Description |
+| :---: | :--- |
+| Direct inference from existing features | Takes an AlphaFold 3 JSON file containing precomputed features such as MSAs and templates as input and produces structure predictions |
+| Protein structure prediction | Takes a protein sequence as input, generates features by searching databases, and predicts its structure |
+| Biomolecular complex modeling | Takes multicomponent systems comprising proteins, DNA, RNA, ligands, and other molecules as input and predicts their 3D structures |
+| Database-search pipeline validation | Uses Jackhmmer / Nhmmer or MMseqs pipelines to verify database paths and the availability of search tools |
+| Hugging Face / OneCode execution | After downloading the model project and complete dataset, quickly verifies that the scripts run correctly in a life-sciences runtime environment |
+
+
+
+# Usage Instructions
+
+## 1. OneCode
+
+Try one-click AI4S development in the OneCode online environment:
+
+[Try one-click AI4S development](https://web-2069360198568017922-iaaj.ksai.scnet.cn:58043/home)
+
+## 2. Manual Installation
+
+**Hardware Requirements**
+
+- GPU or DCU is recommended.
+- A CPU can be used for import checks and lightweight configuration tests; full training and inference will be slow.
+- DCU users must install DTK in advance. DTK 25.04.2 or later is recommended, or a OneScience-recommended version matching the current cluster.
+
+
+
+
+
+**Environment Check**
+
+- NVIDIA GPU:
+
+```bash
+nvidia-smi
+```
+
+- Hygon DCU:
+
+```bash
+hy-smi
+```
+
+### Download the Model Package
+
+```bash
+hf download --model OneScience-Sugon/AlphaFold3 --local-dir ./AlphaFold3
+cd AlphaFold3
+```
+
+### Install the Runtime Environment
+
+**DCU Environment**
+
+```bash
+# Activate DTK and CONDA first
+conda create -n onescience311 python=3.11 -y
+conda activate onescience311
+# uv installation supported
+pip install onescience[bio-dcu] -i http://mirrors.onescience.ai:3141/pypi/simple/  --trusted-host mirrors.onescience.ai
+```
+
+After installation, return to the model package directory:
+
+```bash
+cd ./AlphaFold3
+```
+
+If the AlphaFold 3 C++ extension and runtime data files have not yet been built in the current environment, run:
+
+```bash
+python -m onescience.flax_model.alphafold3.build_extension
+python -m onescience.flax_models.alphafold3.build_data
+```
+
+### Training and Inference Data Overview
+
+The OneScience community has uploaded the complete data required for AlphaFold 3 inference and database searches to Hugging Face: [OneScience-Sugon/AlphaFold3_dataset](https://huggingface.co/datasets/OneScience-Sugon/AlphaFold3_dataset). This model package does not include a training entry point; the dataset is primarily used to construct MSA / template features and perform database searches before inference.
+
+```bash
+hf download --dataset OneScience-Sugon/AlphaFold3_dataset --local-dir ./data/alphafold3
+```
+### Model Weights
+
+Model weights will be available soon.
+
+### Preparing Weights
+
+Place the AlphaFold 3 model weights in the following directory, or specify them via environment variables:
+
+```text
+weight/
+  AlphaFold3/
+    ...
+```
+
+Default lookup order:
+
+- `ALPHAFOLD3_MODEL_DIR`
+- `${ONESCIENCE_MODELS_DIR}/AlphaFold3`
+- `weight/AlphaFold3`
+
+Example:
+
+```bash
+export ALPHAFOLD3_MODEL_DIR=/path/to/AlphaFold3
+```
+
+### Direct Inference
+
+When the input JSON already contains features such as MSA and templates, you can run inference directly:
+
+```bash
+bash scripts/infer.sh
+```
+
+Equivalent Python command example:
+
+```bash
+python scripts/run_alphafold.py \
+  --json_path inputs/7r6r_data.json \
+  --model_dir weight/AlphaFold3 \
+  --output_dir outputs \
+  --run_data_pipeline=false \
+  --flash_attention_implementation=triton
+```
+
+Output is written to `outputs/` and includes the top-ranked structure, structures generated from different seed / sample combinations, a CSV file containing ranking scores, and a copy of the input JSON.
+
+### Jackhmmer / Nhmmer Database Search
+
+When the input JSON contains only sequences and requires a local database search, use:
+
+```bash
+bash scripts/infer_jackhmmer.sh
+```
+
+Common environment variables:
+
+```bash
+export ALPHAFOLD3_DATASET_ROOT=/path/to/alphafold3
+export ALPHAFOLD3_MODEL_DIR=/path/to/AlphaFold3
+export ALPHAFOLD3_JSON_PATH=inputs/t1119_search.json
+export ALPHAFOLD3_OUTPUT_DIR=outputs
+export ALPHAFOLD3_RUN_INFERENCE=false
+```
+
+`ALPHAFOLD3_DATASET_ROOT` is expected to contain directories such as `public_databases/`, `jackhmmer_split/`, and `mmseqsDB/`.
+
+### MMseqs Database Search
+
+If the runtime environment provides the MMseqs executable and the required MMseqs databases, use:
+
+```bash
+bash scripts/infer_mmseqs.sh
+```
+
+Common environment variables:
+
+```bash
+export ALPHAFOLD3_MMSEQS_HOME=/path/to/mmseqs
+export ALPHAFOLD3_DATASET_ROOT=/path/to/alphafold3
+export ALPHAFOLD3_MMSEQS_DB_DIR=/path/to/alphafold3/mmseqsDB
+export ALPHAFOLD3_RUN_INFERENCE=false
+```
+
+To proceed with inference after the database search, set `ALPHAFOLD3_RUN_INFERENCE` to `true` and ensure that the weights directory is available.
+
+# Data Format
+
+AlphaFold 3 inputs are provided in JSON format with the following basic structure:
+
+```json
+{
+  "dialect": "alphafold3",
+  "version": 1,
+  "name": "example",
+  "sequences": [
+    {
+      "protein": {
+        "id": "A",
+        "sequence": "..."
+      }
+    }
+  ],
+  "modelSeeds": [100],
+  "bondedAtomPairs": null,
+  "userCCD": null
+}
+```
+
+This repository provides two examples:
+
+- `inputs/7r6r_data.json`: contains sequence, MSA, and template information; suitable for direct inference.
+- `inputs/t1119_search.json`: contains only sequences; suitable for database-search pipeline validation.
+
+It is recommended to download the full Hugging Face dataset `OneScience-Sugon/AlphaFold3_dataset` to `data/alphafold3/` under the model package. The relative directory structure expected by the database-search pipeline is as follows:
+
+```text
+data/
+  alphafold3/
+    public_databases/
+      mmcif_files/
+      pdb_seqres_2022_09_28.fasta
+      ...
+    jackhmmer_split/
+      bfd-first_non_consensus_sequences.fasta@64
+      mgy_clusters_2022_05.fa@512
+      uniprot_cluster_annot_2021_04.fa@256
+      uniref90_2022_05.fa@128
+    mmseqsDB/
+      small_bfd_db
+      mgnify_db
+      uniprot_cluster_annot_db
+      uniref90_db
+```
+
+# Verification
+
+Static import check:
+
+```bash
+python tests/check_import_boundaries.py
+```
+
+# OneScience Official Information
+
+| Platform | OneScience Main Repository | Skills Repository |
+| --- | --- | --- |
+| Gitee | https://gitee.com/onescience-ai/onescience | https://gitee.com/onescience-ai/oneskills |
+| GitHub | https://github.com/onescience-ai/OneScience | https://github.com/onescience-ai/oneskills |
+
+# Citation & License
+
+- This repository is adapted from the open-source AlphaFold 3 model to support DCUs.
+- The AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0; model parameters are subject to separate usage terms.
+- For scientific use, please cite the original paper: [Accurate structure prediction of biomolecular interactions with AlphaFold 3](https://www.nature.com/articles/s41586-024-07487-w).

configuration.json

@@ -0,0 +1 @@
+{"framework":"JAX","task":"protein-structure-prediction"}

flax_model/__init__.py

@@ -0,0 +1 @@
+"""Local flax model namespace for the split AlphaFold3 project."""

flax_model/alphafold3/CMakeLists.txt

@@ -0,0 +1,161 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+#
+# To request access to the AlphaFold 3 model parameters, follow the process set
+# out at https://github.com/google-deepmind/alphafold3. You may only use these
+# 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
+
+cmake_minimum_required(VERSION 3.24)
+project(
+  "${SKBUILD_PROJECT_NAME}"
+  LANGUAGES CXX
+  VERSION "${SKBUILD_PROJECT_VERSION}")
+
+include(FetchContent)
+set(CMAKE_CXX_STANDARD 20)
+set(CMAKE_CXX_STANDARD_REQUIRED ON)
+set(CMAKE_POSITION_INDEPENDENT_CODE TRUE)
+set(ABSL_PROPAGATE_CXX_STD ON)
+
+# Remove support for scan deps, which is only useful when using C++ modules.
+unset(CMAKE_CXX_SCANDEP_SOURCE)
+
+set(ENV{GIT_CONFIG_PARAMETERS} "'init.defaultRefFormat=files'")
+
+if(DEFINED ENV{ALPHAFOLD3_DEP_DIR})
+   set(DEP_BASE_DIR "$ENV{ALPHAFOLD3_DEP_DIR}")
+endif()
+
+function(declare_af3_dependency dep_name local_path_env repo_env tag_env)
+  if(DEFINED ENV{${local_path_env}} AND EXISTS "$ENV{${local_path_env}}")
+    message(STATUS "Using local dependency for ${dep_name}: $ENV{${local_path_env}}")
+    FetchContent_Declare(
+      ${dep_name}
+      SOURCE_DIR $ENV{${local_path_env}}
+      ${ARGN})
+  else()
+    message(STATUS "Fetching remote dependency for ${dep_name}: $ENV{${repo_env}} @ $ENV{${tag_env}}")
+    FetchContent_Declare(
+      ${dep_name}
+      GIT_REPOSITORY $ENV{${repo_env}}
+      GIT_TAG $ENV{${tag_env}}
+      ${ARGN})
+  endif()
+endfunction()
+
+
+declare_af3_dependency(
+  abseil-cpp
+  ABSEIL_LOCAL_PATH
+  ABSEIL_REPO
+  ABSEIL_TAG
+  EXCLUDE_FROM_ALL)
+
+declare_af3_dependency(
+  pybind11
+  PYBIND11_LOCAL_PATH
+  PYBIND11_REPO
+  PYBIND11_TAG
+  EXCLUDE_FROM_ALL)
+
+declare_af3_dependency(
+  pybind11_abseil
+  PYBIND11_ABSEIL_LOCAL_PATH
+  PYBIND11_ABSEIL_REPO
+  PYBIND11_ABSEIL_TAG
+  EXCLUDE_FROM_ALL)
+
+declare_af3_dependency(
+  cifpp
+  CIFPP_LOCAL_PATH
+  CIFPP_REPO
+  CIFPP_TAG)
+
+declare_af3_dependency(
+  dssp
+  DSSP_LOCAL_PATH
+  DSSP_REPO
+  DSSP_TAG
+  EXCLUDE_FROM_ALL)
+
+if(DEFINED ENV{ALPHAFOLD3_INSTALL_PREFIX})
+    set(CMAKE_INSTALL_PREFIX $ENV{ALPHAFOLD3_INSTALL_PREFIX} CACHE PATH "Install prefix" FORCE)
+    message(STATUS "Setting install prefix to AlphaFold3 package directory: $ENV{ALPHAFOLD3_INSTALL_PREFIX}")
+elseif(DEFINED ENV{CONDA_PREFIX})
+    set(CMAKE_INSTALL_PREFIX $ENV{CONDA_PREFIX} CACHE PATH "Install prefix" FORCE)
+    message(STATUS "Setting install prefix to conda environment: $ENV{CONDA_PREFIX}")
+endif()
+
+if(DEFINED ENV{ALPHAFOLD3_CIFPP_DATA_DIR})
+    set(CIFPP_DATA_DIR "$ENV{ALPHAFOLD3_CIFPP_DATA_DIR}" CACHE PATH "CIFPP data directory" FORCE)
+    message(STATUS "Setting CIFPP_DATA_DIR to ${CIFPP_DATA_DIR}")
+elseif(DEFINED ENV{CONDA_PREFIX})
+    execute_process(
+        COMMAND python -c "import site; print(site.getsitepackages()[0])"
+        OUTPUT_VARIABLE PYTHON_SITE_PACKAGES
+        OUTPUT_STRIP_TRAILING_WHITESPACE
+        ERROR_QUIET
+    )
+
+    if(PYTHON_SITE_PACKAGES)
+        set(CIFPP_DATA_DIR "${PYTHON_SITE_PACKAGES}/share/libcifpp" CACHE PATH "CIFPP data directory" FORCE)
+        message(STATUS "Setting CIFPP_DATA_DIR to Python site-packages: ${CIFPP_DATA_DIR}")
+    endif()
+endif()
+
+FetchContent_MakeAvailable(pybind11 abseil-cpp pybind11_abseil cifpp dssp)
+
+find_package(
+  Python3
+  COMPONENTS Interpreter Development NumPy
+  REQUIRED)
+
+message(STATUS "AF3 Python3_EXECUTABLE=${Python3_EXECUTABLE}")
+message(STATUS "AF3 Python3_VERSION=${Python3_VERSION}")
+message(STATUS "AF3 Python3_INCLUDE_DIRS=${Python3_INCLUDE_DIRS}")
+message(STATUS "AF3 Python3_SITELIB=${Python3_SITELIB}")
+message(STATUS "AF3 Python3_SITEARCH=${Python3_SITEARCH}")
+message(STATUS "AF3 Python3_NumPy_INCLUDE_DIRS=${Python3_NumPy_INCLUDE_DIRS}")
+
+include_directories(${Python3_INCLUDE_DIRS})
+include_directories(${Python3_NumPy_INCLUDE_DIRS})
+include_directories(${CMAKE_CURRENT_SOURCE_DIR}/..)
+
+# Only compile AF3's own C++ sources.
+file(GLOB_RECURSE cpp_srcs ${CMAKE_CURRENT_SOURCE_DIR}/*.cc)
+# Drop any third-party / build-time sources such as benchmarks/tests/examples.
+list(FILTER cpp_srcs EXCLUDE REGEX ".*/_deps/.*")
+list(FILTER cpp_srcs EXCLUDE REGEX ".*/mirror_deps/.*")
+list(FILTER cpp_srcs EXCLUDE REGEX ".*/CMake/install_test_project/.*")
+list(FILTER cpp_srcs EXCLUDE REGEX "${CMAKE_CURRENT_SOURCE_DIR}/test/.*")
+#list(FILTER cpp_srcs EXCLUDE REGEX ".*\\(_test\\|_main\\|_benchmark\\).cc$")
+list(FILTER cpp_srcs EXCLUDE REGEX ".*(_test|_main|_benchmark)\\.cc$")
+message(STATUS "AF3 cpp_srcs after filter: ${cpp_srcs}")
+
+add_compile_definitions(NPY_NO_DEPRECATED_API=NPY_1_7_API_VERSION)
+
+pybind11_add_module(cpp ${cpp_srcs})
+
+
+
+target_link_libraries(
+  cpp
+  PRIVATE absl::check
+          absl::flat_hash_map
+          absl::node_hash_map
+          absl::strings
+          absl::status
+          absl::statusor
+          absl::log
+          pybind11_abseil::absl_casters
+          Python3::NumPy
+          dssp::dssp
+          cifpp::cifpp)
+
+target_compile_definitions(cpp PRIVATE VERSION_INFO=${PROJECT_VERSION})
+install(TARGETS cpp LIBRARY DESTINATION .)
+
+

flax_model/alphafold3/README.md

@@ -0,0 +1,81 @@
+# AlphaFold 3 Integration for OneScience
+
+This is the AlphaFold 3 implementation integrated into the OneScience framework as a submodule.
+
+## Overview
+
+AlphaFold 3 is a state-of-the-art machine learning model for predicting protein structure, developed by DeepMind. This submodule integrates AlphaFold 3 into the OneScience framework, allowing it to be used as part of larger scientific computing workflows.
+
+## Installation
+
+### As Part of OneScience
+
+The recommended way to install alphafold3 is as part of the complete OneScience package:
+
+```bash
+# Install OneScience with alphafold3 support
+
+# install jackhmmer
+# mkdir ~/hmmer_build ~/hmmer
+# wget http://eddylab.org/software/hmmer/hmmer-3.4.tar.gz --directory-prefix ~/hmmer_build
+# cd ~/hmmer_build &&  tar zxf hmmer-3.4.tar.gz && rm hmmer-3.4.tar.gz
+# patch -p0 < jackhmmer_seq_limit.patch
+# cd ~/hmmer-3.4
+# ./configure --prefix ~/hmmer
+# make -j && make install && cd ./easel && make install
+# rm -R ~/hmmer_build
+
+# # install extension
+# pip install .[bio] -c constraints.txt
+# cp -r /public/onestore/onedatasets/alphafold3/_dep xxx/
+# export ALPHAFOLD3_DEP_DIR=/public/onestore/onedatasets/alphafold3/_dep
+# cd src/onescience/flax_models/alphafold3/
+# python build_extension.py
+
+# optional create mmseqs2 database (please contact ai4s@sugon.com for mmseqs2 program)
+export mmfasta=/root/public_databases
+cd /root/public_databases && mkdir mmseqsDB
+export mmdb=/root/public_databases/mmseqsDB
+export CUDA_VISIBLE_DEVICES=0
+mmseqs createdb $mmfasta/bfd-first_non_consensus_sequences.fasta $mmdb/small_bfd_db --gpu 1 --threads 32 --createdb-mode 2
+mmseqs createdb $mmfasta/mgy_clusters_2022_05.fa $mmdb/mgnify_db --gpu 1 --threads 32 --createdb-mode 2
+mmseqs createdb $mmfasta/uniprot_all_2021_04.fa $mmdb/uniprot_cluster_annot_db --gpu 1 --threads 32 --createdb-mode 2
+mmseqs createdb $mmfasta/uniref90_2022_05.fa $mmdb/uniref90_db --gpu 1 --threads 32 --createdb-mode 2
+```
+
+## Usage
+
+```python
+# Import alphafold3 as part of onescience
+import flax_model.alphafold3 as af3
+
+# Access alphafold3 components
+from flax_model.alphafold3 import structure, model, data
+
+# Use alphafold3 functionality
+print(f"AlphaFold3 version: {af3.__version__}")
+```
+
+## Requirements
+
+- Python 3.11+
+- JAX with CUDA support (optional, for GPU acceleration)
+- CMake 3.28+ (for building C++ extensions)
+- Additional dependencies listed in pyproject.toml
+
+## License
+
+This code is licensed under CC BY-NC-SA 4.0. See the original AlphaFold 3 repository for more details on usage restrictions and licensing terms.
+
+## Citation
+
+If you use this code in your research, please cite the AlphaFold 3 paper:
+
+```
+Abramson, J., Adler, J., Dunger, J. et al. Accurate structure prediction of biomolecular interactions with AlphaFold 3. Nature 630, 493�?00 (2024).
+```
+
+## Links
+
+- [Original AlphaFold 3 Repository](https://github.com/google-deepmind/alphafold3)
+- [AlphaFold 3 Paper](https://www.nature.com/articles/s41586-024-07487-w) 

flax_model/alphafold3/__init__.py

@@ -0,0 +1,26 @@
+
+
+"""An implementation of the inference pipeline of AlphaFold 3."""
+from importlib import resources
+from pathlib import Path
+import os
+import warnings
+
+
+def _data_artifacts_exist() -> bool:
+    try:
+        root = Path(
+            resources.files(__name__ + ".constants.converters")
+        )
+    except Exception:
+        return False
+    return (root / "ccd.pickle").exists() and (root / "chemical_component_sets.pickle").exists()
+
+
+if not _data_artifacts_exist():
+    warnings.warn(
+        "AlphaFold3 data files (ccd.pickle, chemical_component_sets.pickle) are missing.\n"
+        "Please run the following command once to build local artifacts:\n\n"
+        "    python -m flax_model.alphafold3.build_extension\n",
+        stacklevel=1,
+    )

flax_model/alphafold3/_build.py

@@ -0,0 +1,649 @@
+import json
+import os
+import shutil
+import subprocess
+import sys
+import sysconfig
+import tarfile
+import tempfile
+import urllib.request
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+MIRROR_CONFIG_PATH = ROOT / "docs" / "af3_dependency_mirrors.template.json"
+AF3_DIR = Path(__file__).resolve().parent
+CONVERTERS_DIR = AF3_DIR / "constants" / "converters"
+TMP_ROOT = Path(os.environ.get("ALPHAFOLD3_TMP_ROOT", tempfile.gettempdir())).resolve() / "alphafold3_split_af3"
+BUILD_DIR = TMP_ROOT / "build"
+REMOTE_DEPS_DIR = TMP_ROOT / "mirror_deps"
+HMMER_INSTALL_DIR = AF3_DIR / "_tools" / "hmmer"
+HMMER_BIN_DIR = HMMER_INSTALL_DIR / "bin"
+HMMER_PATCH_PATH = AF3_DIR / "jackhmmer_seq_limit.patch"
+HMMER_SOURCE_URL = "https://gitee.com/zhang-yuqi-sudo/hmmer-github/releases/download/hmmer-3.4/hmmer-3.4.tar.gz"
+HMMER_ARCHIVE_NAME = "hmmer-3.4.tar.gz"
+HMMER_SOURCE_DIRNAME = "hmmer-3.4"
+HMMER_BINARIES = ("jackhmmer", "nhmmer", "hmmalign", "hmmsearch", "hmmbuild")
+HMMER_BUILD_ROOT = TMP_ROOT / "hmmer_build"
+HMMER_BUILD_DIR = HMMER_BUILD_ROOT / HMMER_SOURCE_DIRNAME
+HMMER_ARCHIVE_PATH = HMMER_BUILD_ROOT / HMMER_ARCHIVE_NAME
+HMMER_PATCH_MARKER = HMMER_BUILD_ROOT / ".patch_applied"
+HMMER_BUILD_MARKER = HMMER_INSTALL_DIR / ".built"
+DATA_FILES = (
+    CONVERTERS_DIR / "ccd.pickle",
+    CONVERTERS_DIR / "chemical_component_sets.pickle",
+)
+LIB_PATTERNS = ("cpp*.so", "cpp*.pyd", "cpp*.dll", "cpp*.dylib")
+
+TOP_LEVEL_MIRROR_MAP = {
+    "abseil-cpp": "ABSEIL",
+    "pybind11": "PYBIND11",
+    "pybind11_abseil": "PYBIND11_ABSEIL",
+    "libcifpp": "CIFPP",
+    "dssp": "DSSP",
+}
+
+DEFAULT_DOWNSTREAM_TAGS = {
+    "boost-regex": "boost-1.87.0",
+    "libmcfp": "v1.4.2",
+    "catch2": "v3.4.0",
+}
+
+PREPARED_ENV: dict[str, str] | None = None
+PREPARED_DEP_DIR: Path | None = None
+PREPARED_REMOTE = False
+
+
+class AF3BuildError(RuntimeError):
+    pass
+
+
+def _env_flag(name: str, default: str = "auto") -> str:
+    return os.environ.get(name, default).strip().lower()
+
+
+def load_mirror_config() -> dict:
+    if not MIRROR_CONFIG_PATH.exists():
+        return {}
+    return json.loads(MIRROR_CONFIG_PATH.read_text())
+
+
+def mirror_value(name: str, field: str, default: str | None = None) -> str | None:
+    entry = load_mirror_config().get(name, {})
+    value = entry.get(field)
+    if isinstance(value, str) and value.strip() and "<fill-" not in value and "<your-org>" not in value:
+        return value.strip()
+    return default
+
+
+def reset_prepared_state() -> None:
+    global PREPARED_ENV, PREPARED_DEP_DIR, PREPARED_REMOTE
+    PREPARED_ENV = None
+    PREPARED_DEP_DIR = None
+    PREPARED_REMOTE = False
+
+
+def should_build() -> bool:
+    if _env_flag("ALPHAFOLD3_SKIP_BUILD", "0") in {"1", "true", "on", "yes"}:
+        return False
+    if _env_flag("ONESCIENCE_SKIP_AF3_BUILD", "0") in {"1", "true", "on", "yes"}:
+        return False
+    mode = _env_flag("ALPHAFOLD3_BUILD", _env_flag("ONESCIENCE_BUILD_AF3", "auto"))
+    if mode in {"0", "false", "off", "no"}:
+        return False
+    if mode in {"1", "true", "on", "yes", "force"}:
+        return True
+    return AF3_DIR.exists()
+
+
+def is_strict() -> bool:
+    return _env_flag("ALPHAFOLD3_STRICT", _env_flag("ONESCIENCE_AF3_STRICT", "0")) in {"1", "true", "on", "yes"}
+
+
+def force_rebuild() -> bool:
+    return _env_flag("ALPHAFOLD3_FORCE_REBUILD", _env_flag("ONESCIENCE_AF3_FORCE_REBUILD", "0")) in {"1", "true", "on", "yes"}
+
+
+def resolve_dep_dir() -> Path | None:
+    dep_dir = os.environ.get("ALPHAFOLD3_DEP_DIR")
+    if dep_dir:
+        return Path(dep_dir).expanduser().resolve()
+    candidate = ROOT / "third_party" / "alphafold3_deps"
+    if candidate.exists():
+        return candidate
+    return None
+
+
+def using_remote_dependencies() -> bool:
+    return resolve_dep_dir() is None
+
+
+def install_destination() -> Path:
+    return AF3_DIR
+
+
+def build_data_output_dir() -> Path:
+    return CONVERTERS_DIR
+
+
+def cifpp_data_dir() -> Path | None:
+    env_path = os.environ.get("ALPHAFOLD3_CIFPP_DATA_DIR")
+    if env_path:
+        return Path(env_path)
+    return AF3_DIR / "_data" / "libcifpp"
+
+
+def cifpp_components_path() -> Path | None:
+    data_dir = cifpp_data_dir()
+    if data_dir is None:
+        return None
+    return data_dir / "components.cif"
+
+
+def current_python_include_dir() -> str:
+    paths = sysconfig.get_paths()
+    return paths.get("include", "")
+
+
+def current_numpy_include_dir() -> str | None:
+    try:
+        import numpy
+        return numpy.get_include()
+    except Exception:
+        return None
+
+
+def _lib_exists() -> bool:
+    return any(AF3_DIR.glob(pattern) for pattern in LIB_PATTERNS)
+
+
+def _data_exists() -> bool:
+    return all(path.exists() for path in DATA_FILES)
+
+
+def artifacts_exist() -> bool:
+    return _lib_exists() and _data_exists()
+
+
+def hmmer_binaries_exist() -> bool:
+    return all((HMMER_BIN_DIR / name).exists() for name in HMMER_BINARIES)
+
+
+def download_hmmer_source() -> None:
+    if HMMER_ARCHIVE_PATH.exists():
+        return
+    HMMER_BUILD_ROOT.mkdir(parents=True, exist_ok=True)
+    print(f"[AF3] downloading HMMER source from {HMMER_SOURCE_URL}")
+    urllib.request.urlretrieve(HMMER_SOURCE_URL, HMMER_ARCHIVE_PATH)
+
+
+def extract_hmmer_source() -> None:
+    if HMMER_BUILD_DIR.exists():
+        return
+    HMMER_BUILD_ROOT.mkdir(parents=True, exist_ok=True)
+    with tarfile.open(HMMER_ARCHIVE_PATH, "r:gz") as tar:
+        tar.extractall(HMMER_BUILD_ROOT)
+
+
+def apply_hmmer_patch() -> None:
+    if HMMER_PATCH_MARKER.exists():
+        return
+    subprocess.run(
+        ["patch", "-p0", "-i", str(HMMER_PATCH_PATH.resolve())],
+        cwd=HMMER_BUILD_ROOT,
+        check=True,
+    )
+    HMMER_PATCH_MARKER.write_text("patched\n")
+
+
+def build_hmmer_tools() -> None:
+    if hmmer_binaries_exist() and HMMER_BUILD_MARKER.exists():
+        return
+    if shutil.which("make") is None and shutil.which("gmake") is None:
+        raise AF3BuildError("make or gmake is required to build HMMER tools.")
+    make_program = shutil.which("gmake") or shutil.which("make")
+    download_hmmer_source()
+    extract_hmmer_source()
+    apply_hmmer_patch()
+    HMMER_INSTALL_DIR.mkdir(parents=True, exist_ok=True)
+    print(f"[AF3] building HMMER tools into {HMMER_INSTALL_DIR}")
+    try:
+        subprocess.run(["./configure", f"--prefix={HMMER_INSTALL_DIR}"], cwd=HMMER_BUILD_DIR, check=True)
+        build_parallelism = min(os.cpu_count() or 8, 64)
+        subprocess.run([make_program, "-j", str(build_parallelism)], cwd=HMMER_BUILD_DIR, check=True)
+        subprocess.run([make_program, "install"], cwd=HMMER_BUILD_DIR, check=True)
+        subprocess.run([make_program, "install"], cwd=HMMER_BUILD_DIR / "easel", check=True)
+    except subprocess.CalledProcessError as exc:
+        raise AF3BuildError(
+            f"HMMER tools build failed with exit code {exc.returncode}. "
+            f"Check HMMER build logs under {HMMER_BUILD_DIR} for details."
+        ) from exc
+    HMMER_BUILD_MARKER.write_text("built\n")
+    print(f"[AF3] HMMER tools ready at {HMMER_BIN_DIR}")
+
+def ensure_hmmer_tools_ready() -> None:
+    if hmmer_binaries_exist():
+        return
+    if shutil.which("patch") is None:
+        raise AF3BuildError("patch is required to build HMMER tools.")
+    build_hmmer_tools()
+
+
+def ensure_dependencies_ready() -> None:
+    dep_dir = resolve_dep_dir()
+    if dep_dir is None:
+        return
+    required = ("abseil-cpp", "pybind11", "pybind11_abseil", "libcifpp", "dssp")
+    missing = [name for name in required if not (dep_dir / name).exists()]
+    if missing:
+        raise AF3BuildError(
+            f"AlphaFold3 dependency directories are missing under {dep_dir}: {', '.join(missing)}"
+        )
+
+
+def ensure_build_prerequisites() -> None:
+    if shutil.which("cmake") is None:
+        raise AF3BuildError("CMake is required to build AlphaFold3.")
+    if using_remote_dependencies() and shutil.which("git") is None:
+        raise AF3BuildError(
+            "Git is required for remote AlphaFold3 dependency fetching when no local dependency directory is configured."
+        )
+    ensure_dependencies_ready()
+    ensure_hmmer_tools_ready()
+
+
+def resolve_fetch_base() -> str:
+    return os.environ.get("ALPHAFOLD3_FETCH_BASE", "https://gitee.com/zhang-yuqi-sudo").rstrip("/")
+
+
+def dependency_specs() -> dict[str, dict[str, str]]:
+    base = resolve_fetch_base()
+    return {
+        "ABSEIL": {"repo": f"{base}/abseil-cpp", "tag": "d7aaad83b488fd62bd51c81ecf16cd938532cc0a", "local_path": "abseil-cpp"},
+        "PYBIND11": {"repo": f"{base}/pybind11", "tag": "2e0815278cb899b20870a67ca8205996ef47e70f", "local_path": "pybind11"},
+        "PYBIND11_ABSEIL": {"repo": f"{base}/pybind11_abseil", "tag": "bddf30141f9fec8e577f515313caec45f559d319", "local_path": "pybind11_abseil"},
+        "CIFPP": {"repo": f"{base}/libcifpp", "tag": "ac98531a2fc8daf21131faa0c3d73766efa46180", "local_path": "libcifpp"},
+        "DSSP": {"repo": f"{base}/dssp", "tag": "57560472b4260dc41f457706bc45fc6ef0bc0f10", "local_path": "dssp"},
+    }
+
+
+def announce_build_mode() -> None:
+    if PREPARED_DEP_DIR is not None and PREPARED_REMOTE:
+        print(f"Using prepared mirrored dependencies from {PREPARED_DEP_DIR}.")
+    elif PREPARED_DEP_DIR is not None:
+        print(f"Using local AlphaFold3 dependencies from {PREPARED_DEP_DIR}.")
+    elif using_remote_dependencies():
+        print(f"Using remote AlphaFold3 dependencies from {resolve_fetch_base()}.")
+    else:
+        print(f"Using local AlphaFold3 dependencies from {resolve_dep_dir()}.")
+
+
+def announce_install_destination() -> None:
+    print(f"Installing AlphaFold3 artifacts into {install_destination()}.")
+
+
+def announce_data_destination() -> None:
+    print(f"Writing AlphaFold3 data files into {build_data_output_dir()}.")
+
+
+def announce_cifpp_data_location() -> None:
+    data_dir = cifpp_data_dir()
+    if data_dir is not None:
+        print(f"Using libcifpp data directory {data_dir}.")
+
+
+def announce_mirror_overrides() -> None:
+    config = load_mirror_config()
+    active = []
+    for name, entry in config.items():
+        mirror = entry.get("mirror")
+        if isinstance(mirror, str) and mirror.strip() and "<your-org>" not in mirror:
+            active.append(f"{name} -> {mirror}")
+    if active:
+        print("Using dependency mirror config:")
+        for item in active:
+            print(f"  {item}")
+
+
+def announce_hmmer_location() -> None:
+    print(f"Using HMMER tool directory {HMMER_BIN_DIR}.")
+
+
+def ensure_install_dirs() -> None:
+    install_destination().mkdir(parents=True, exist_ok=True)
+    build_data_output_dir().mkdir(parents=True, exist_ok=True)
+
+
+def prepare_install_layout() -> None:
+    ensure_install_dirs()
+    announce_build_mode()
+    announce_install_destination()
+    announce_data_destination()
+    announce_cifpp_data_location()
+    announce_hmmer_location()
+    announce_mirror_overrides()
+
+
+def ensure_local_components_file() -> None:
+    components_path = cifpp_components_path()
+    if components_path is None or components_path.exists():
+        return
+    if shutil.which("wget") is None:
+        return
+    components_path.parent.mkdir(parents=True, exist_ok=True)
+    try:
+        subprocess.run(
+            [
+                "wget",
+                "-O",
+                str(components_path),
+                "https://files.wwpdb.org/pub/pdb/data/monomers/components.cif",
+            ],
+            check=True,
+        )
+    except subprocess.CalledProcessError as exc:
+        raise AF3BuildError(f"Failed to download components.cif with wget: {exc}") from exc
+
+
+def ensure_components_file_available() -> None:
+    components_path = cifpp_components_path()
+    if components_path is None:
+        return
+    if components_path.exists():
+        return
+    raise AF3BuildError(
+        f"Expected libcifpp components file at {components_path}, but it was not created during dependency install."
+    )
+
+
+def finalise_install_layout() -> None:
+    ensure_components_file_available()
+
+
+def patch_text(path: Path, old: str, new: str) -> None:
+    if not path.exists() or not new:
+        return
+    text = path.read_text()
+    if old in text:
+        path.write_text(text.replace(old, new))
+
+
+def clone_or_update_repo(dep_name: str, repo: str, tag: str) -> Path:
+    target = REMOTE_DEPS_DIR / dep_name
+    if target.exists():
+        shutil.rmtree(target, ignore_errors=True)
+    target.parent.mkdir(parents=True, exist_ok=True)
+    subprocess.run(["git", "clone", repo, str(target)], check=True)
+    subprocess.run(["git", "-C", str(target), "checkout", tag], check=True)
+    return target
+
+
+def disable_tests(source_dir: Path) -> None:
+    patch_text(
+        source_dir / "CMakeLists.txt",
+        "if(BUILD_TESTING AND PROJECT_IS_TOP_LEVEL)",
+        "if(FALSE AND BUILD_TESTING AND PROJECT_IS_TOP_LEVEL)",
+    )
+    patch_text(
+        source_dir / "CMakeLists.txt",
+        "if(BUILD_PYTHON_MODULE)",
+        "if(FALSE AND BUILD_PYTHON_MODULE)",
+    )
+
+
+def patch_catch2_sources(test_cmake: Path, env: dict[str, str]) -> None:
+    patch_text(test_cmake, "https://github.com/catchorg/Catch2.git", env.get("ALPHAFOLD3_CATCH2_REPO", ""))
+    patch_text(test_cmake, "v3.4.0", env.get("ALPHAFOLD3_CATCH2_TAG", ""))
+
+
+def patch_regex_sources(source_dir: Path, env: dict[str, str]) -> None:
+    patch_text(source_dir / "CMakeLists.txt", "https://github.com/boostorg/regex", env.get("ALPHAFOLD3_BOOST_REGEX_REPO", ""))
+    patch_text(source_dir / "CMakeLists.txt", "boost-1.87.0", env.get("ALPHAFOLD3_BOOST_REGEX_TAG", ""))
+
+
+# def patch_archive_urls(source_dir: Path, env: dict[str, str]) -> None:
+#     patch_text(
+#         source_dir / "pcre2-simple" / "CMakeLists.txt",
+#         "https://github.com/PCRE2Project/pcre2/releases/download/pcre2-10.46/pcre2-10.46.tar.gz",
+#         env.get("ALPHAFOLD3_PCRE2_URL", ""),
+#     )
+
+
+def patch_eigen_sources(source_dir: Path, env: dict[str, str]) -> None:
+    cmake_path = source_dir / "CMakeLists.txt"
+    patch_text(
+        cmake_path,
+        "https://gitlab.com/libeigen/eigen.git",
+        env.get("ALPHAFOLD3_EIGEN_REPO", ""),
+    )
+    patch_text(
+        cmake_path,
+        "GIT_TAG 3.4.0",
+        f"GIT_TAG {env.get('ALPHAFOLD3_EIGEN_TAG', '')}",
+    )
+    patch_text(
+        cmake_path,
+        "\t# Create a private copy of eigen3 and populate it only, no need to build\n",
+        "\t# Create a private copy of eigen3 and populate it only, no need to build\n\tmessage(STATUS \"AF3 libcifpp: Eigen3 not found locally, populating my-eigen3 from GIT repository\")\n\tmessage(STATUS \"AF3 libcifpp: my-eigen3 repo=$ENV{ALPHAFOLD3_EIGEN_REPO} tag=$ENV{ALPHAFOLD3_EIGEN_TAG}\")\n",
+    )
+    patch_text(
+        cmake_path,
+        "\tFetchContent_GetProperties(my-eigen3)\n",
+        "\tFetchContent_GetProperties(my-eigen3)\n\tmessage(STATUS \"AF3 libcifpp: my-eigen3 populated=${my-eigen3_POPULATED}\")\n",
+    )
+    patch_text(
+        cmake_path,
+        "\tif(NOT my-eigen3_POPULATED)\n\t\tFetchContent_Populate(my-eigen3)\n\tendif()\n",
+        "\tif(NOT my-eigen3_POPULATED)\n\t\tmessage(STATUS \"AF3 libcifpp: starting FetchContent_Populate(my-eigen3)\")\n\t\tFetchContent_Populate(my-eigen3)\n\t\tmessage(STATUS \"AF3 libcifpp: finished FetchContent_Populate(my-eigen3), source=${my-eigen3_SOURCE_DIR}\")\n\tendif()\n",
+    )
+    patch_text(
+        cmake_path,
+        "\tset(EIGEN_INCLUDE_DIR ${my-eigen3_SOURCE_DIR})\n",
+        "\tset(EIGEN_INCLUDE_DIR ${my-eigen3_SOURCE_DIR})\n\tmessage(STATUS \"AF3 libcifpp: EIGEN_INCLUDE_DIR=${EIGEN_INCLUDE_DIR}\")\n",
+    )
+
+
+def patch_all_known_downstream_sources(source_dir: Path, env: dict[str, str]) -> None:
+    disable_tests(source_dir)
+    patch_regex_sources(source_dir, env)
+    # patch_archive_urls(source_dir, env)
+    patch_eigen_sources(source_dir, env)
+    test_cmake = source_dir / "test" / "CMakeLists.txt"
+    if test_cmake.exists():
+        patch_catch2_sources(test_cmake, env)
+
+
+def patch_libcifpp_sources(source_dir: Path, env: dict[str, str]) -> None:
+    patch_all_known_downstream_sources(source_dir, env)
+
+
+def patch_dssp_sources(source_dir: Path, env: dict[str, str]) -> None:
+    patch_all_known_downstream_sources(source_dir, env)
+    patch_text(source_dir / "CMakeLists.txt", "https://github.com/mhekkel/libmcfp", env.get("ALPHAFOLD3_LIBMCFP_REPO", ""))
+    patch_text(source_dir / "CMakeLists.txt", "v1.4.2", env.get("ALPHAFOLD3_LIBMCFP_TAG", ""))
+    patch_text(source_dir / "CMakeLists.txt", "https://github.com/PDB-REDO/libcifpp", env.get("CIFPP_REPO", ""))
+    patch_text(source_dir / "CMakeLists.txt", "v10.0.1", env.get("CIFPP_TAG", ""))
+
+
+def prepare_remote_dependency_sources(env: dict[str, str]) -> Path:
+    if REMOTE_DEPS_DIR.exists():
+        shutil.rmtree(REMOTE_DEPS_DIR, ignore_errors=True)
+    REMOTE_DEPS_DIR.mkdir(parents=True, exist_ok=True)
+
+    specs = dependency_specs()
+    prepared = {}
+    clone_order = ["ABSEIL", "PYBIND11", "PYBIND11_ABSEIL", "CIFPP", "DSSP"]
+    for key_prefix in clone_order:
+        spec = specs[key_prefix]
+        dep_name = spec["local_path"]
+        repo = env[f"{key_prefix}_REPO"]
+        tag = env[f"{key_prefix}_TAG"]
+        prepared[dep_name] = clone_or_update_repo(dep_name, repo, tag)
+        if dep_name == "libcifpp":
+            patch_libcifpp_sources(prepared[dep_name], env)
+        elif dep_name == "dssp":
+            patch_dssp_sources(prepared[dep_name], env)
+
+    return REMOTE_DEPS_DIR
+
+
+def prepare_build_env_once() -> dict[str, str]:
+    global PREPARED_ENV, PREPARED_DEP_DIR, PREPARED_REMOTE
+    if PREPARED_ENV is not None:
+        return PREPARED_ENV
+
+    env = os.environ.copy()
+    env["PATH"] = str(HMMER_BIN_DIR) + os.pathsep + env.get("PATH", "")
+    specs = dependency_specs()
+    for prefix, spec in specs.items():
+        env[f"{prefix}_REPO"] = mirror_value(spec["local_path"], "mirror", spec["repo"]) or spec["repo"]
+        env[f"{prefix}_TAG"] = mirror_value(spec["local_path"], "tag", spec["tag"]) or spec["tag"]
+
+    env["ALPHAFOLD3_BOOST_REGEX_REPO"] = mirror_value("boost-regex", "mirror", "https://github.com/boostorg/regex") or "https://github.com/boostorg/regex"
+    env["ALPHAFOLD3_BOOST_REGEX_TAG"] = mirror_value("boost-regex", "tag", DEFAULT_DOWNSTREAM_TAGS["boost-regex"]) or DEFAULT_DOWNSTREAM_TAGS["boost-regex"]
+    env["ALPHAFOLD3_LIBMCFP_REPO"] = mirror_value("libmcfp", "mirror", "https://github.com/mhekkel/libmcfp") or "https://github.com/mhekkel/libmcfp"
+    env["ALPHAFOLD3_LIBMCFP_TAG"] = mirror_value("libmcfp", "tag", DEFAULT_DOWNSTREAM_TAGS["libmcfp"]) or DEFAULT_DOWNSTREAM_TAGS["libmcfp"]
+    env["ALPHAFOLD3_CATCH2_REPO"] = mirror_value("catch2", "mirror", "https://github.com/catchorg/Catch2.git") or "https://github.com/catchorg/Catch2.git"
+    env["ALPHAFOLD3_CATCH2_TAG"] = mirror_value("catch2", "tag", DEFAULT_DOWNSTREAM_TAGS["catch2"]) or DEFAULT_DOWNSTREAM_TAGS["catch2"]
+    pcre2_url = mirror_value("pcre2", "mirror")
+    if pcre2_url:
+        env["ALPHAFOLD3_PCRE2_URL"] = pcre2_url
+    eigen_repo = mirror_value("eigen", "mirror", "https://gitlab.com/libeigen/eigen.git") or "https://gitlab.com/libeigen/eigen.git"
+    eigen_tag = mirror_value("eigen", "tag", "3.4.0") or "3.4.0"
+    env["ALPHAFOLD3_EIGEN_REPO"] = eigen_repo
+    env["ALPHAFOLD3_EIGEN_TAG"] = eigen_tag
+    dep_dir = resolve_dep_dir()
+    if dep_dir is None and using_remote_dependencies():
+        dep_dir = prepare_remote_dependency_sources(env)
+        PREPARED_REMOTE = True
+    if dep_dir is not None:
+        PREPARED_DEP_DIR = dep_dir
+        env["ALPHAFOLD3_DEP_DIR"] = str(dep_dir)
+        for prefix, spec in specs.items():
+            env[f"{prefix}_LOCAL_PATH"] = str(dep_dir / spec["local_path"])
+
+    env["ALPHAFOLD3_INSTALL_PREFIX"] = str(install_destination())
+    env["ALPHAFOLD3_DATA_OUTPUT_DIR"] = str(build_data_output_dir())
+    data_dir = cifpp_data_dir()
+    components_path = cifpp_components_path()
+    if data_dir is not None:
+        env["ALPHAFOLD3_CIFPP_DATA_DIR"] = str(data_dir)
+    if components_path is not None:
+        env["ALPHAFOLD3_CIFPP_COMPONENTS"] = str(components_path)
+    env["ALPHAFOLD3_PYTHON_EXECUTABLE"] = sys.executable
+    python_include_dir = current_python_include_dir()
+    if python_include_dir:
+        env["ALPHAFOLD3_PYTHON_INCLUDE_DIR"] = python_include_dir
+    numpy_include_dir = current_numpy_include_dir()
+    if numpy_include_dir:
+        env["ALPHAFOLD3_NUMPY_INCLUDE_DIR"] = numpy_include_dir
+    pythonpath_entries = [str(ROOT)]
+    existing_pythonpath = env.get("PYTHONPATH")
+    if existing_pythonpath:
+        pythonpath_entries.append(existing_pythonpath)
+    env["PYTHONPATH"] = os.pathsep.join(pythonpath_entries)
+
+    PREPARED_ENV = env
+    return env
+
+
+def _build_env() -> dict[str, str]:
+    return prepare_build_env_once().copy()
+
+
+def _run(command: list[str], cwd: Path) -> None:
+    env = _build_env()
+    if command and command[0] == "cmake":
+        resolved_cmake = shutil.which("cmake", path=env.get("PATH"))
+        print(f"[AF3] build python executable: {sys.executable}")
+        print(f"[AF3] build cwd: {cwd}")
+        print(f"[AF3] build PATH: {env.get('PATH', '')}")
+        print(f"[AF3] resolved cmake: {resolved_cmake}")
+        print(f"[AF3] build PYTHONPATH: {env.get('PYTHONPATH', '')}")
+        print(f"[AF3] build CONDA_PREFIX: {env.get('CONDA_PREFIX', '')}")
+        print(f"[AF3] build VIRTUAL_ENV: {env.get('VIRTUAL_ENV', '')}")
+        print(f"[AF3] ALPHAFOLD3_PYTHON_EXECUTABLE: {env.get('ALPHAFOLD3_PYTHON_EXECUTABLE', '')}")
+        print(f"[AF3] ALPHAFOLD3_PYTHON_INCLUDE_DIR: {env.get('ALPHAFOLD3_PYTHON_INCLUDE_DIR', '')}")
+        print(f"[AF3] ALPHAFOLD3_NUMPY_INCLUDE_DIR: {env.get('ALPHAFOLD3_NUMPY_INCLUDE_DIR', '')}")
+    subprocess.run(command, cwd=cwd, env=env, check=True)
+
+
+def build_cpp_extension() -> None:
+    BUILD_DIR.mkdir(exist_ok=True)
+    try:
+        cmake_configure_command = [
+            "cmake",
+            str(AF3_DIR),
+            "-DCMAKE_BUILD_TYPE=Release",
+            "-DCMAKE_CXX_STANDARD=20",
+            "-DCMAKE_POSITION_INDEPENDENT_CODE=ON",
+            "-DBUILD_TESTING=OFF",
+            "-DCMAKE_CXX_SCAN_FOR_MODULES=OFF",
+            "-DSKBUILD_PROJECT_NAME=cpp",
+            "-DSKBUILD_PROJECT_VERSION=0.3.0",
+            f"-DPython3_EXECUTABLE={sys.executable}",
+            "-DPython3_FIND_STRATEGY=LOCATION",
+        ]
+        python_include_dir = current_python_include_dir()
+        if python_include_dir:
+            cmake_configure_command.append(f"-DPython3_INCLUDE_DIR={python_include_dir}")
+        numpy_include_dir = current_numpy_include_dir()
+        if numpy_include_dir:
+            cmake_configure_command.append(f"-DPython3_NumPy_INCLUDE_DIR={numpy_include_dir}")
+        _run(cmake_configure_command, cwd=BUILD_DIR)
+        build_parallelism = min(os.cpu_count() or 8, 64)
+        _run(["cmake", "--build", ".", "--parallel", str(build_parallelism)], cwd=BUILD_DIR)
+        _run(["cmake", "--install", ".", "--prefix", str(install_destination())], cwd=BUILD_DIR)
+    except FileNotFoundError as exc:
+        raise AF3BuildError(f"CMake not found: {exc}") from exc
+    except subprocess.CalledProcessError as exc:
+        raise AF3BuildError(f"Failed to build AlphaFold3 C++ extension: {exc}") from exc
+    finally:
+        cleanup_build_artifacts()
+
+
+def build_data_files() -> None:
+    try:
+        from flax_model.alphafold3.build_data import build_data
+        build_data()
+    except Exception as exc:
+        raise AF3BuildError(f"Failed to build AlphaFold3 data files: {exc}") from exc
+
+
+def cleanup_build_artifacts() -> None:
+    for path in (
+        BUILD_DIR,
+        AF3_DIR / "include",
+        AF3_DIR / "lib",
+        AF3_DIR / "lib64",
+        AF3_DIR / "var",
+        AF3_DIR / "etc",
+    ):
+        if path.exists():
+            shutil.rmtree(path, ignore_errors=True)
+
+
+def cleanup_prepared_remote_deps() -> None:
+    if PREPARED_REMOTE and REMOTE_DEPS_DIR.exists():
+        shutil.rmtree(REMOTE_DEPS_DIR, ignore_errors=True)
+
+
+def build_all() -> None:
+    reset_prepared_state()
+    ensure_build_prerequisites()
+    prepare_build_env_once()
+    ensure_local_components_file()
+    prepare_install_layout()
+    print("Building AlphaFold3 C++ extension...")
+    build_cpp_extension()
+    finalise_install_layout()
+    print("Building AlphaFold3 data files...")
+    build_data_files()
+    cleanup_prepared_remote_deps()
+
+
+def build_if_needed() -> None:
+    if not should_build():
+        return
+    if artifacts_exist() and not force_rebuild():
+        return
+    build_all()

flax_model/alphafold3/build_data.py

@@ -0,0 +1,51 @@
+
+
+"""Script for building intermediate data."""
+
+import os
+from importlib import resources
+import pathlib
+import site
+
+import flax_model.alphafold3.constants.converters
+from flax_model.alphafold3.constants.converters import ccd_pickle_gen
+from flax_model.alphafold3.constants.converters import chemical_component_sets_gen
+
+
+def build_data():
+  """Builds intermediate data."""
+  explicit_cif_path = os.environ.get('ALPHAFOLD3_CIFPP_COMPONENTS')
+  if explicit_cif_path:
+    cif_path = pathlib.Path(explicit_cif_path)
+    if not cif_path.exists():
+      raise ValueError(f'Configured components.cif does not exist: {cif_path}')
+  else:
+    for site_path in site.getsitepackages():
+      path = pathlib.Path(site_path) / 'share/libcifpp/components.cif'
+      if path.exists():
+        cif_path = path
+        break
+    else:
+      raise ValueError('Could not find components.cif')
+
+  output_root = os.environ.get('ALPHAFOLD3_DATA_OUTPUT_DIR')
+  if output_root:
+    out_root = pathlib.Path(output_root)
+    out_root.mkdir(parents=True, exist_ok=True)
+  else:
+    out_root = pathlib.Path(resources.files(flax_model.alphafold3.constants.converters))
+  ccd_pickle_path = out_root / 'ccd.pickle'
+  chemical_component_sets_pickle_path = out_root / 'chemical_component_sets.pickle'
+  ccd_pickle_gen.main(['', str(cif_path), str(ccd_pickle_path)])
+  
+  chemical_component_sets_gen.main(
+      ['', str(chemical_component_sets_pickle_path)]
+  )
+
+
+def main() -> None:
+    build_data()
+
+
+if __name__ == "__main__":
+    main()

flax_model/alphafold3/build_extension.py

@@ -0,0 +1,19 @@
+#!/usr/bin/env python3
+"""Compatibility entry point for local AlphaFold3 build helpers."""
+
+import sys
+from pathlib import Path
+
+ROOT = Path(__file__).resolve().parents[2]
+sys.path.insert(0, str(ROOT))
+
+from flax_model.alphafold3._build import AF3BuildError, build_all
+
+
+if __name__ == "__main__":
+    try:
+        build_all()
+    except AF3BuildError as exc:
+        print(f"Error: {exc}")
+        sys.exit(1)
+    sys.exit(0)

flax_model/alphafold3/common/base_config.py

@@ -0,0 +1,141 @@
+
+
+"""Config for the protein folding model and experiment."""
+
+from collections.abc import Mapping
+import copy
+import dataclasses
+import types
+import typing
+from typing import Any, ClassVar, TypeVar
+
+
+_T = TypeVar('_T')
+_ConfigT = TypeVar('_ConfigT', bound='BaseConfig')
+
+
+def _strip_optional(t: type[Any]) -> type[Any]:
+  """Transforms type annotations of the form `T | None` to `T`."""
+  if typing.get_origin(t) in (typing.Union, types.UnionType):
+    args = set(typing.get_args(t)) - {types.NoneType}
+    if len(args) == 1:
+      return args.pop()
+  return t
+
+
+_NO_UPDATE = object()
+
+
+class _Autocreate:
+
+  def __init__(self, **defaults: Any):
+    self.defaults = defaults
+
+
+def autocreate(**defaults: Any) -> Any:
+  """Marks a field as having a default factory derived from its type."""
+  return _Autocreate(**defaults)
+
+
+def _clone_field(
+    field: dataclasses.Field[_T], new_default: _T
+) -> dataclasses.Field[_T]:
+  if new_default is _NO_UPDATE:
+    return copy.copy(field)
+  return dataclasses.field(
+      default=new_default,
+      init=True,
+      kw_only=True,
+      repr=field.repr,
+      hash=field.hash,
+      compare=field.compare,
+      metadata=field.metadata,
+  )
+
+
+@typing.dataclass_transform()
+class ConfigMeta(type):
+  """Metaclass that synthesizes a __post_init__ that coerces dicts to Config subclass instances."""
+
+  def __new__(mcs, name, bases, classdict):
+    cls = super().__new__(mcs, name, bases, classdict)
+
+    def _coercable_fields(self) -> Mapping[str, tuple[ConfigMeta, Any]]:
+      type_hints = typing.get_type_hints(self.__class__)
+      fields = dataclasses.fields(self.__class__)
+      field_to_type_and_default = {
+          field.name: (_strip_optional(type_hints[field.name]), field.default)
+          for field in fields
+      }
+      coercable_fields = {
+          f: t
+          for f, t in field_to_type_and_default.items()
+          if issubclass(type(t[0]), ConfigMeta)
+      }
+      return coercable_fields
+
+    cls._coercable_fields = property(_coercable_fields)
+
+    old_post_init = getattr(cls, '__post_init__', None)
+
+    def _post_init(self) -> None:
+      # Use get_type_hints instead of Field.type to ensure that forward
+      # references are resolved.
+      for field_name, (
+          field_type,
+          field_default,
+      ) in self._coercable_fields.items():  # pylint: disable=protected-access
+        field_value = getattr(self, field_name)
+        if field_value is None:
+          continue
+        try:
+          match field_value:
+            case _Autocreate():
+              # Construct from field defaults.
+              setattr(self, field_name, field_type(**field_value.defaults))
+            case Mapping():
+              # Field value is not yet a `Config` instance; Assume we can create
+              # one by splatting keys and values.
+              args = {}
+              # Apply default args first, if present.
+              if isinstance(field_default, _Autocreate):
+                args.update(field_default.defaults)
+              args.update(field_value)
+              setattr(self, field_name, field_type(**args))
+            case _:
+              pass
+        except TypeError as e:
+          raise TypeError(
+              f'Failure while coercing field {field_name!r} of'
+              f' {self.__class__.__qualname__}'
+          ) from e
+      if old_post_init:
+        old_post_init(self)
+
+    cls.__post_init__ = _post_init
+
+    return dataclasses.dataclass(kw_only=True)(cls)
+
+
+class BaseConfig(metaclass=ConfigMeta):
+  """Config base class.
+
+  Subclassing Config automatically makes the subclass a kw_only dataclass with
+  a `__post_init__` that coerces Config-subclass field values from mappings to
+  instances of the right type.
+  """
+  # Provided by dataclasses.make_dataclass
+  __dataclass_fields__: ClassVar[dict[str, dataclasses.Field[Any]]]
+
+  # Overridden by metaclass
+  @property
+  def _coercable_fields(self) -> Mapping[str, tuple[type['BaseConfig'], Any]]:
+    return {}
+
+  def as_dict(self) -> Mapping[str, Any]:
+    result = dataclasses.asdict(self)
+    for field_name in self._coercable_fields:
+      field_value = getattr(self, field_name, None)
+      if isinstance(field_value, BaseConfig):
+        result[field_name] = field_value.as_dict()
+    return result

flax_model/alphafold3/common/folding_input.py

@@ -0,0 +1,1528 @@
+
+
+"""Model input dataclass."""
+
+from collections.abc import Collection, Iterator, Mapping, Sequence
+import dataclasses
+import gzip
+import json
+import logging
+import lzma
+import os
+import pathlib
+import random
+import re
+import string
+from typing import Any, Final, Self, TypeAlias, cast
+
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import cif_dict
+from flax_model.alphafold3.structure import mmcif as mmcif_lib
+import rdkit.Chem as rd_chem
+import zstandard as zstd
+
+
+BondAtomId: TypeAlias = tuple[str, int, str]
+
+JSON_DIALECT: Final[str] = 'alphafold3'
+JSON_VERSIONS: Final[tuple[int, ...]] = (1, 2, 3, 4)
+JSON_VERSION: Final[int] = JSON_VERSIONS[-1]
+
+ALPHAFOLDSERVER_JSON_DIALECT: Final[str] = 'alphafoldserver'
+ALPHAFOLDSERVER_JSON_VERSION: Final[int] = 1
+
+
+def _validate_keys(actual: Collection[str], expected: Collection[str]):
+  """Validates that the JSON doesn't contain any extra unwanted keys."""
+  if bad_keys := set(actual) - set(expected):
+    raise ValueError(f'Unexpected JSON keys in: {", ".join(sorted(bad_keys))}')
+
+
+def _read_file(path: pathlib.Path, json_path: pathlib.Path | None) -> str:
+  """Reads a maybe compressed (gzip, xz, zstd) file from the given path.
+
+  Args:
+    path: The path to the file to read. This can be either absolute path, or a
+      path relative to the JSON file path.
+    json_path: The path to the JSON file. If None, the path must be absolute.
+
+  Returns:
+    The contents of the file.
+  """
+  if not path.is_absolute():
+    if json_path is None:
+      raise ValueError('json_path must be specified if path is not absolute.')
+    path = (json_path.parent / path).resolve()
+
+  with open(path, 'rb') as f:
+    first_six_bytes = f.read(6)
+    f.seek(0)
+
+    # Detect the compression type using the magic number in the header.
+    if first_six_bytes[:2] == b'\x1f\x8b':
+      with gzip.open(f, 'rt') as gzip_f:
+        return cast(str, gzip_f.read())
+    elif first_six_bytes == b'\xfd\x37\x7a\x58\x5a\x00':
+      with lzma.open(f, 'rt') as xz_f:
+        return cast(str, xz_f.read())
+    elif first_six_bytes[:4] == b'\x28\xb5\x2f\xfd':
+      with zstd.open(f, 'rt') as zstd_f:
+        return cast(str, zstd_f.read())
+    else:
+      return f.read().decode('utf-8')
+
+
+class Template:
+  """Structural template input."""
+
+  __slots__ = ('_mmcif', '_query_to_template')
+
+  def __init__(self, *, mmcif: str, query_to_template_map: Mapping[int, int]):
+    """Initializes the template.
+
+    Args:
+      mmcif: The structural template in mmCIF format. The mmCIF should have only
+        one protein chain.
+      query_to_template_map: A mapping from query residue index to template
+        residue index.
+    """
+    self._mmcif = mmcif
+    # Needed to make the Template class hashable.
+    self._query_to_template = tuple(query_to_template_map.items())
+
+  @property
+  def query_to_template_map(self) -> Mapping[int, int]:
+    return dict(self._query_to_template)
+
+  @property
+  def mmcif(self) -> str:
+    return self._mmcif
+
+  def __hash__(self) -> int:
+    return hash((self._mmcif, tuple(sorted(self._query_to_template))))
+
+  def __eq__(self, other: Self) -> bool:
+    mmcifs_equal = self._mmcif == other._mmcif
+    maps_equal = sorted(self._query_to_template) == sorted(
+        other._query_to_template
+    )
+    return mmcifs_equal and maps_equal
+
+
+class ProteinChain:
+  """Protein chain input."""
+
+  __slots__ = (
+      '_id',
+      '_sequence',
+      '_ptms',
+      '_description',
+      '_paired_msa',
+      '_unpaired_msa',
+      '_templates',
+  )
+
+  def __init__(
+      self,
+      *,
+      id: str,  # pylint: disable=redefined-builtin
+      sequence: str,
+      ptms: Sequence[tuple[str, int]],
+      description: str | None = None,
+      paired_msa: str | None = None,
+      unpaired_msa: str | None = None,
+      templates: Sequence[Template] | None = None,
+  ):
+    """Initializes a single protein chain input.
+
+    Args:
+      id: Unique protein chain identifier.
+      sequence: The amino acid sequence of the chain.
+      ptms: A list of tuples containing the post-translational modification type
+        and the (1-based) residue index where the modification is applied.
+      description: An optional textual description of the protein chain.
+      paired_msa: Paired A3M-formatted MSA for this chain. This MSA is not
+        deduplicated and will be used to compute paired features. If None, this
+        field is unset and must be filled in by the data pipeline before
+        featurisation. If set to an empty string, it will be treated as a custom
+        MSA with no sequences.
+      unpaired_msa: Unpaired A3M-formatted MSA for this chain. This will be
+        deduplicated and used to compute unpaired features. If None, this field
+        is unset and must be filled in by the data pipeline before
+        featurisation. If set to an empty string, it will be treated as a custom
+        MSA with no sequences.
+      templates: A list of structural templates for this chain. If None, this
+        field is unset and must be filled in by the data pipeline before
+        featurisation. The list can be empty or contain up to 20 templates.
+    """
+    if not all(res.isalpha() for res in sequence):
+      raise ValueError(f'Protein must contain only letters, got "{sequence}"')
+    if any(not 0 < mod[1] <= len(sequence) for mod in ptms):
+      raise ValueError(f'Invalid protein modification index: {ptms}')
+    if any(mod[0].startswith('CCD_') for mod in ptms):
+      raise ValueError(
+          f'Protein ptms must not contain the "CCD_" prefix, got {ptms}'
+      )
+    # Use hashable containers for ptms and templates.
+    self._id = id
+    self._sequence = sequence
+    self._ptms = tuple(ptms)
+    self._description = description
+    self._paired_msa = paired_msa
+    self._unpaired_msa = unpaired_msa
+    self._templates = tuple(templates) if templates is not None else None
+
+  @property
+  def id(self) -> str:
+    return self._id
+
+  @property
+  def sequence(self) -> str:
+    """Returns a single-letter sequence, taking modifications into account.
+
+    Uses 'X' for all unknown residues.
+    """
+    return ''.join([
+        residue_names.letters_three_to_one(r, default='X')
+        for r in self.to_ccd_sequence()
+    ])
+
+  @property
+  def ptms(self) -> Sequence[tuple[str, int]]:
+    return self._ptms
+
+  @property
+  def description(self) -> str | None:
+    return self._description
+
+  @property
+  def paired_msa(self) -> str | None:
+    return self._paired_msa
+
+  @property
+  def unpaired_msa(self) -> str | None:
+    return self._unpaired_msa
+
+  @property
+  def templates(self) -> Sequence[Template] | None:
+    return self._templates
+
+  def __len__(self) -> int:
+    return len(self._sequence)
+
+  def __eq__(self, other: Self) -> bool:
+    return (
+        self._id == other._id
+        and self._sequence == other._sequence
+        and self._ptms == other._ptms
+        and self._description == other._description
+        and self._paired_msa == other._paired_msa
+        and self._unpaired_msa == other._unpaired_msa
+        and self._templates == other._templates
+    )
+
+  def __hash__(self) -> int:
+    return hash((
+        self._id,
+        self._sequence,
+        self._ptms,
+        self._description,
+        self._paired_msa,
+        self._unpaired_msa,
+        self._templates,
+    ))
+
+  def hash_without_id(self) -> int:
+    """Returns a hash ignoring the ID - useful for deduplication."""
+    return hash((
+        self._sequence,
+        self._ptms,
+        self._description,
+        self._paired_msa,
+        self._unpaired_msa,
+        self._templates,
+    ))
+
+  @classmethod
+  def from_alphafoldserver_dict(
+      cls, json_dict: Mapping[str, Any], seq_id: str
+  ) -> Self:
+    """Constructs ProteinChain from the AlphaFoldServer JSON dict."""
+    _validate_keys(
+        json_dict.keys(),
+        {
+            'sequence',
+            'glycans',
+            'modifications',
+            'count',
+            'maxTemplateDate',
+            'useStructureTemplate',
+        },
+    )
+    sequence = json_dict['sequence']
+
+    if 'glycans' in json_dict:
+      raise ValueError(
+          f'Specifying glycans in the `{ALPHAFOLDSERVER_JSON_DIALECT}` format'
+          ' is not supported.'
+      )
+
+    if 'maxTemplateDate' in json_dict:
+      raise ValueError(
+          f'Specifying maxTemplateDate in the `{ALPHAFOLDSERVER_JSON_DIALECT}`'
+          ' format is not supported, use the --max_template_date flag instead.'
+      )
+
+    templates = None  # Search for templates unless explicitly disabled.
+    if not json_dict.get('useStructureTemplate', True):
+      templates = []  # Do not use any templates.
+
+    ptms = [
+        (mod['ptmType'].removeprefix('CCD_'), mod['ptmPosition'])
+        for mod in json_dict.get('modifications', [])
+    ]
+    return cls(id=seq_id, sequence=sequence, ptms=ptms, templates=templates)
+
+  @classmethod
+  def from_dict(
+      cls,
+      json_dict: Mapping[str, Any],
+      json_path: pathlib.Path | None = None,
+      seq_id: str | None = None,
+  ) -> Self:
+    """Constructs ProteinChain from the AlphaFold JSON dict."""
+    json_dict = json_dict['protein']
+    _validate_keys(
+        json_dict.keys(),
+        {
+            'id',
+            'sequence',
+            'modifications',
+            'description',
+            'unpairedMsa',
+            'unpairedMsaPath',
+            'pairedMsa',
+            'pairedMsaPath',
+            'templates',
+        },
+    )
+
+    sequence = json_dict['sequence']
+    ptms = [
+        (mod['ptmType'], mod['ptmPosition'])
+        for mod in json_dict.get('modifications', [])
+    ]
+
+    unpaired_msa = json_dict.get('unpairedMsa', None)
+    unpaired_msa_path = json_dict.get('unpairedMsaPath', None)
+    if unpaired_msa and unpaired_msa_path:
+      raise ValueError('Only one of unpairedMsa/unpairedMsaPath can be set.')
+    if (
+        unpaired_msa
+        and len(unpaired_msa) < 256
+        and os.path.exists(unpaired_msa)
+    ):
+      raise ValueError(
+          'Set the unpaired MSA path using the "unpairedMsaPath" field.'
+      )
+    elif unpaired_msa_path:
+      unpaired_msa = _read_file(pathlib.Path(unpaired_msa_path), json_path)
+
+    paired_msa = json_dict.get('pairedMsa', None)
+    paired_msa_path = json_dict.get('pairedMsaPath', None)
+    if paired_msa and paired_msa_path:
+      raise ValueError('Only one of pairedMsa/pairedMsaPath can be set.')
+    if paired_msa and len(paired_msa) < 256 and os.path.exists(paired_msa):
+      raise ValueError(
+          'Set the paired MSA path using the "pairedMsaPath" field.'
+      )
+    elif paired_msa_path:
+      paired_msa = _read_file(pathlib.Path(paired_msa_path), json_path)
+
+    raw_templates = json_dict.get('templates', None)
+
+    if raw_templates is None:
+      templates = None
+    else:
+      templates = []
+      for raw_template in raw_templates:
+        _validate_keys(
+            raw_template.keys(),
+            {'mmcif', 'mmcifPath', 'queryIndices', 'templateIndices'},
+        )
+        mmcif = raw_template.get('mmcif', None)
+        mmcif_path = raw_template.get('mmcifPath', None)
+        if mmcif and mmcif_path:
+          raise ValueError('Only one of mmcif/mmcifPath can be set.')
+        if mmcif and len(mmcif) < 256 and os.path.exists(mmcif):
+          raise ValueError('Set the template path using the "mmcifPath" field.')
+        if mmcif_path:
+          mmcif = _read_file(pathlib.Path(mmcif_path), json_path)
+        query_to_template_map = dict(
+            zip(raw_template['queryIndices'], raw_template['templateIndices'])
+        )
+        templates.append(
+            Template(mmcif=mmcif, query_to_template_map=query_to_template_map)
+        )
+
+    return cls(
+        id=seq_id or json_dict['id'],
+        sequence=sequence,
+        ptms=ptms,
+        description=json_dict.get('description', None),
+        paired_msa=paired_msa,
+        unpaired_msa=unpaired_msa,
+        templates=templates,
+    )
+
+  def to_dict(
+      self, seq_id: str | Sequence[str] | None = None
+  ) -> Mapping[str, Mapping[str, Any]]:
+    """Converts ProteinChain to an AlphaFold JSON dict."""
+    if self._templates is None:
+      templates = None
+    else:
+      templates = [
+          {
+              'mmcif': template.mmcif,
+              'queryIndices': list(template.query_to_template_map.keys()),
+              'templateIndices': (
+                  list(template.query_to_template_map.values()) or None
+              ),
+          }
+          for template in self._templates
+      ]
+    contents = {
+        'id': seq_id or self._id,
+        'sequence': self._sequence,
+        'modifications': [
+            {'ptmType': ptm[0], 'ptmPosition': ptm[1]} for ptm in self._ptms
+        ],
+        'unpairedMsa': self._unpaired_msa,
+        'pairedMsa': self._paired_msa,
+        'templates': templates,
+    }
+    if self._description is not None:
+      contents['description'] = self._description
+    return {'protein': contents}
+
+  def to_ccd_sequence(self) -> Sequence[str]:
+    """Converts to a sequence of CCD codes."""
+    ccd_coded_seq = [
+        residue_names.PROTEIN_COMMON_ONE_TO_THREE.get(res, residue_names.UNK)
+        for res in self._sequence
+    ]
+    for ptm_code, ptm_index in self._ptms:
+      ccd_coded_seq[ptm_index - 1] = ptm_code
+    return ccd_coded_seq
+
+  def fill_missing_fields(self) -> Self:
+    """Fill missing MSA and template fields with default values."""
+    return ProteinChain(
+        id=self.id,
+        sequence=self._sequence,
+        ptms=self._ptms,
+        description=self._description,
+        unpaired_msa=self._unpaired_msa or '',
+        paired_msa=self._paired_msa or '',
+        templates=self._templates or [],
+    )
+
+
+class RnaChain:
+  """RNA chain input."""
+
+  __slots__ = (
+      '_id',
+      '_sequence',
+      '_modifications',
+      '_description',
+      '_unpaired_msa',
+  )
+
+  def __init__(
+      self,
+      *,
+      id: str,  # pylint: disable=redefined-builtin
+      sequence: str,
+      modifications: Sequence[tuple[str, int]],
+      description: str | None = None,
+      unpaired_msa: str | None = None,
+  ):
+    """Initializes a single strand RNA chain input.
+
+    Args:
+      id: Unique RNA chain identifier.
+      sequence: The RNA sequence of the chain.
+      modifications: A list of tuples containing the modification type and the
+        (1-based) residue index where the modification is applied.
+      description: An optional textual description of the RNA chain.
+      unpaired_msa: Unpaired A3M-formatted MSA for this chain. This will be
+        deduplicated and used to compute unpaired features. If None, this field
+        is unset and must be filled in by the data pipeline before
+        featurisation. If set to an empty string, it will be treated as a custom
+        MSA with no sequences.
+    """
+    if not all(res.isalpha() for res in sequence):
+      raise ValueError(f'RNA must contain only letters, got "{sequence}"')
+    if any(not 0 < mod[1] <= len(sequence) for mod in modifications):
+      raise ValueError(f'Invalid RNA modification index: {modifications}')
+    if any(mod[0].startswith('CCD_') for mod in modifications):
+      raise ValueError(
+          'RNA modifications must not contain the "CCD_" prefix, got'
+          f' {modifications}'
+      )
+    self._id = id
+    self._sequence = sequence
+    # Use hashable container for modifications.
+    self._modifications = tuple(modifications)
+    self._description = description
+    self._unpaired_msa = unpaired_msa
+
+  @property
+  def id(self) -> str:
+    return self._id
+
+  @property
+  def sequence(self) -> str:
+    """Returns a single-letter sequence, taking modifications into account.
+
+    Uses 'N' for all unknown residues.
+    """
+    return ''.join([
+        residue_names.letters_three_to_one(r, default='N')
+        for r in self.to_ccd_sequence()
+    ])
+
+  @property
+  def modifications(self) -> Sequence[tuple[str, int]]:
+    return self._modifications
+
+  @property
+  def description(self) -> str | None:
+    return self._description
+
+  @property
+  def unpaired_msa(self) -> str | None:
+    return self._unpaired_msa
+
+  def __len__(self) -> int:
+    return len(self._sequence)
+
+  def __eq__(self, other: Self) -> bool:
+    return (
+        self._id == other._id
+        and self._sequence == other._sequence
+        and self._modifications == other._modifications
+        and self._description == other._description
+        and self._unpaired_msa == other._unpaired_msa
+    )
+
+  def __hash__(self) -> int:
+    return hash((
+        self._id,
+        self._sequence,
+        self._modifications,
+        self._description,
+        self._unpaired_msa,
+    ))
+
+  def hash_without_id(self) -> int:
+    """Returns a hash ignoring the ID - useful for deduplication."""
+    return hash((
+        self._sequence,
+        self._modifications,
+        self._description,
+        self._unpaired_msa,
+    ))
+
+  @classmethod
+  def from_alphafoldserver_dict(
+      cls, json_dict: Mapping[str, Any], seq_id: str
+  ) -> Self:
+    """Constructs RnaChain from the AlphaFoldServer JSON dict."""
+    _validate_keys(json_dict.keys(), {'sequence', 'modifications', 'count'})
+    sequence = json_dict['sequence']
+    modifications = [
+        (mod['modificationType'].removeprefix('CCD_'), mod['basePosition'])
+        for mod in json_dict.get('modifications', [])
+    ]
+    return cls(id=seq_id, sequence=sequence, modifications=modifications)
+
+  @classmethod
+  def from_dict(
+      cls,
+      json_dict: Mapping[str, Any],
+      json_path: pathlib.Path | None = None,
+      seq_id: str | None = None,
+  ) -> Self:
+    """Constructs RnaChain from the AlphaFold JSON dict."""
+    json_dict = json_dict['rna']
+    _validate_keys(
+        json_dict.keys(),
+        {
+            'id',
+            'sequence',
+            'modifications',
+            'description',
+            'unpairedMsa',
+            'unpairedMsaPath',
+        },
+    )
+    sequence = json_dict['sequence']
+    modifications = [
+        (mod['modificationType'], mod['basePosition'])
+        for mod in json_dict.get('modifications', [])
+    ]
+
+    unpaired_msa = json_dict.get('unpairedMsa', None)
+    unpaired_msa_path = json_dict.get('unpairedMsaPath', None)
+    if unpaired_msa and unpaired_msa_path:
+      raise ValueError('Only one of unpairedMsa/unpairedMsaPath can be set.')
+    if (
+        unpaired_msa
+        and len(unpaired_msa) < 256
+        and os.path.exists(unpaired_msa)
+    ):
+      raise ValueError(
+          'Set the unpaired MSA path using the "unpairedMsaPath" field.'
+      )
+    elif unpaired_msa_path:
+      unpaired_msa = _read_file(pathlib.Path(unpaired_msa_path), json_path)
+
+    return cls(
+        id=seq_id or json_dict['id'],
+        sequence=sequence,
+        modifications=modifications,
+        description=json_dict.get('description', None),
+        unpaired_msa=unpaired_msa,
+    )
+
+  def to_dict(
+      self, seq_id: str | Sequence[str] | None = None
+  ) -> Mapping[str, Mapping[str, Any]]:
+    """Converts RnaChain to an AlphaFold JSON dict."""
+    contents = {
+        'id': seq_id or self._id,
+        'sequence': self._sequence,
+        'modifications': [
+            {'modificationType': mod[0], 'basePosition': mod[1]}
+            for mod in self._modifications
+        ],
+        'unpairedMsa': self._unpaired_msa,
+    }
+    if self._description is not None:
+      contents['description'] = self._description
+    return {'rna': contents}
+
+  def to_ccd_sequence(self) -> Sequence[str]:
+    """Converts to a sequence of CCD codes."""
+    mapping = {r: r for r in residue_names.RNA_TYPES}  # Same 1-letter and CCD.
+    ccd_coded_seq = [
+        mapping.get(res, residue_names.UNK_RNA) for res in self._sequence
+    ]
+    for ccd_code, modification_index in self._modifications:
+      ccd_coded_seq[modification_index - 1] = ccd_code
+    return ccd_coded_seq
+
+  def fill_missing_fields(self) -> Self:
+    """Fill missing MSA fields with default values."""
+    return RnaChain(
+        id=self.id,
+        sequence=self.sequence,
+        modifications=self.modifications,
+        unpaired_msa=self._unpaired_msa or '',
+    )
+
+
+class DnaChain:
+  """Single strand DNA chain input."""
+
+  __slots__ = ('_id', '_sequence', '_modifications', '_description')
+
+  def __init__(
+      self,
+      *,
+      id: str,  # pylint: disable=redefined-builtin
+      sequence: str,
+      modifications: Sequence[tuple[str, int]],
+      description: str | None = None,
+  ):
+    """Initializes a single strand DNA chain input.
+
+    Args:
+      id: Unique DNA chain identifier.
+      sequence: The DNA sequence of the chain.
+      modifications: A list of tuples containing the modification type and the
+        (1-based) residue index where the modification is applied.
+      description: An optional textual description of the DNA chain.
+    """
+    if not all(res.isalpha() for res in sequence):
+      raise ValueError(f'DNA must contain only letters, got "{sequence}"')
+    if any(not 0 < mod[1] <= len(sequence) for mod in modifications):
+      raise ValueError(f'Invalid DNA modification index: {modifications}')
+    if any(mod[0].startswith('CCD_') for mod in modifications):
+      raise ValueError(
+          'DNA modifications must not contain the "CCD_" prefix, got'
+          f' {modifications}'
+      )
+    self._id = id
+    self._sequence = sequence
+    # Use hashable container for modifications.
+    self._modifications = tuple(modifications)
+    self._description = description
+
+  @property
+  def id(self) -> str:
+    return self._id
+
+  @property
+  def sequence(self) -> str:
+    """Returns a single-letter sequence, taking modifications into account.
+
+    Uses 'N' for all unknown residues.
+    """
+    return ''.join([
+        residue_names.letters_three_to_one(r, default='N')
+        for r in self.to_ccd_sequence()
+    ])
+
+  @property
+  def description(self) -> str | None:
+    return self._description
+
+  def __len__(self) -> int:
+    return len(self._sequence)
+
+  def __eq__(self, other: Self) -> bool:
+    return (
+        self._id == other._id
+        and self._sequence == other._sequence
+        and self._modifications == other._modifications
+        and self._description == other._description
+    )
+
+  def __hash__(self) -> int:
+    return hash(
+        (self._id, self._sequence, self._modifications, self._description)
+    )
+
+  def modifications(self) -> Sequence[tuple[str, int]]:
+    return self._modifications
+
+  def hash_without_id(self) -> int:
+    """Returns a hash ignoring the ID - useful for deduplication."""
+    return hash((self._sequence, self._modifications, self._description))
+
+  @classmethod
+  def from_alphafoldserver_dict(
+      cls, json_dict: Mapping[str, Any], seq_id: str
+  ) -> Self:
+    """Constructs DnaChain from the AlphaFoldServer JSON dict."""
+    _validate_keys(json_dict.keys(), {'sequence', 'modifications', 'count'})
+    sequence = json_dict['sequence']
+    modifications = [
+        (mod['modificationType'].removeprefix('CCD_'), mod['basePosition'])
+        for mod in json_dict.get('modifications', [])
+    ]
+    return cls(id=seq_id, sequence=sequence, modifications=modifications)
+
+  @classmethod
+  def from_dict(
+      cls, json_dict: Mapping[str, Any], seq_id: str | None = None
+  ) -> Self:
+    """Constructs DnaChain from the AlphaFold JSON dict."""
+    json_dict = json_dict['dna']
+    _validate_keys(
+        json_dict.keys(), {'id', 'sequence', 'modifications', 'description'}
+    )
+    sequence = json_dict['sequence']
+    modifications = [
+        (mod['modificationType'], mod['basePosition'])
+        for mod in json_dict.get('modifications', [])
+    ]
+    return cls(
+        id=seq_id or json_dict['id'],
+        sequence=sequence,
+        modifications=modifications,
+        description=json_dict.get('description', None),
+    )
+
+  def to_dict(
+      self, seq_id: str | Sequence[str] | None = None
+  ) -> Mapping[str, Mapping[str, Any]]:
+    """Converts DnaChain to an AlphaFold JSON dict."""
+    contents = {
+        'id': seq_id or self._id,
+        'sequence': self._sequence,
+        'modifications': [
+            {'modificationType': mod[0], 'basePosition': mod[1]}
+            for mod in self._modifications
+        ],
+    }
+    if self._description is not None:
+      contents['description'] = self._description
+    return {'dna': contents}
+
+  def to_ccd_sequence(self) -> Sequence[str]:
+    """Converts to a sequence of CCD codes."""
+    ccd_coded_seq = [
+        residue_names.DNA_COMMON_ONE_TO_TWO.get(res, residue_names.UNK_DNA)
+        for res in self._sequence
+    ]
+    for ccd_code, modification_index in self._modifications:
+      ccd_coded_seq[modification_index - 1] = ccd_code
+    return ccd_coded_seq
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Ligand:
+  """Ligand input.
+
+  Attributes:
+    id: Unique ligand "chain" identifier.
+    ccd_ids: The Chemical Component Dictionary or user-defined CCD IDs of the
+      chemical components of the ligand. Typically, this is just a single ID,
+      but some ligands are composed of multiple components. If that is the case,
+      a bond linking these components should be added to the bonded_atom_pairs
+      Input field.
+    smiles: The SMILES representation of the ligand.
+    description: An optional textual description of the ligand.
+  """
+
+  id: str
+  ccd_ids: Sequence[str] | None = None
+  smiles: str | None = None
+  description: str | None = None
+
+  def __post_init__(self):
+    if (self.ccd_ids is None) == (self.smiles is None):
+      raise ValueError('Ligand must have one of CCD ID or SMILES set.')
+
+    if self.smiles is not None:
+      mol = rd_chem.MolFromSmiles(self.smiles)
+      if not mol:
+        raise ValueError(f'Unable to make RDKit Mol from SMILES: {self.smiles}')
+
+    # Use hashable types for ccd_ids.
+    if self.ccd_ids is not None:
+      object.__setattr__(self, 'ccd_ids', tuple(self.ccd_ids))
+
+  def __len__(self) -> int:
+    if self.ccd_ids is not None:
+      return len(self.ccd_ids)
+    else:
+      return 1
+
+  def hash_without_id(self) -> int:
+    """Returns a hash ignoring the ID - useful for deduplication."""
+    return hash((self.ccd_ids, self.smiles, self.description))
+
+  @classmethod
+  def from_alphafoldserver_dict(
+      cls, json_dict: Mapping[str, Any], seq_id: str
+  ) -> Self:
+    """Constructs Ligand from the AlphaFoldServer JSON dict."""
+    # Ligand can be specified either as a ligand, or ion (special-case).
+    _validate_keys(json_dict.keys(), {'ligand', 'ion', 'count'})
+    if 'ligand' in json_dict:
+      return cls(id=seq_id, ccd_ids=[json_dict['ligand'].removeprefix('CCD_')])
+    elif 'ion' in json_dict:
+      return cls(id=seq_id, ccd_ids=[json_dict['ion']])
+    else:
+      raise ValueError(f'Unknown ligand type: {json_dict}')
+
+  @classmethod
+  def from_dict(
+      cls, json_dict: Mapping[str, Any], seq_id: str | None = None
+  ) -> Self:
+    """Constructs Ligand from the AlphaFold JSON dict."""
+    json_dict = json_dict['ligand']
+    _validate_keys(
+        json_dict.keys(), {'id', 'ccdCodes', 'smiles', 'description'}
+    )
+    if json_dict.get('ccdCodes') and json_dict.get('smiles'):
+      raise ValueError(
+          'Ligand cannot have both CCD code and SMILES set at the same time, '
+          f'got CCD: {json_dict["ccdCodes"]} and SMILES: {json_dict["smiles"]}'
+      )
+
+    if 'ccdCodes' in json_dict:
+      ccd_codes = json_dict['ccdCodes']
+      if not isinstance(ccd_codes, (list, tuple)):
+        raise ValueError(
+            'CCD codes must be a list of strings, got '
+            f'{type(ccd_codes).__name__} instead: {ccd_codes}'
+        )
+      return cls(
+          id=seq_id or json_dict['id'],
+          ccd_ids=ccd_codes,
+          description=json_dict.get('description', None),
+      )
+    elif 'smiles' in json_dict:
+      return cls(
+          id=seq_id or json_dict['id'],
+          smiles=json_dict['smiles'],
+          description=json_dict.get('description', None),
+      )
+    else:
+      raise ValueError(f'Unknown ligand type: {json_dict}')
+
+  def to_dict(
+      self, seq_id: str | Sequence[str] | None = None
+  ) -> Mapping[str, Mapping[str, Any]]:
+    """Converts Ligand to an AlphaFold JSON dict."""
+    contents = {'id': seq_id or self.id}
+    if self.ccd_ids is not None:
+      contents['ccdCodes'] = self.ccd_ids
+    if self.smiles is not None:
+      contents['smiles'] = self.smiles
+    if self.description is not None:
+      contents['description'] = self.description
+    return {'ligand': contents}
+
+
+def _sample_rng_seed() -> int:
+  """Sample a random seed for AlphaFoldServer job."""
+  # See https://alphafoldserver.com/faq#what-are-seeds-and-how-are-they-set.
+  return random.randint(0, 2**32 - 1)
+
+
+def _validate_user_ccd_keys(keys: Sequence[str], component_name: str) -> None:
+  """Validates the keys of the user-defined CCD dictionary."""
+  mandatory_keys = (
+      '_chem_comp.id',
+      '_chem_comp.name',
+      '_chem_comp.type',
+      '_chem_comp.formula',
+      '_chem_comp.mon_nstd_parent_comp_id',
+      '_chem_comp.pdbx_synonyms',
+      '_chem_comp.formula_weight',
+      '_chem_comp_atom.comp_id',
+      '_chem_comp_atom.atom_id',
+      '_chem_comp_atom.type_symbol',
+      '_chem_comp_atom.charge',
+      '_chem_comp_atom.pdbx_model_Cartn_x_ideal',
+      '_chem_comp_atom.pdbx_model_Cartn_y_ideal',
+      '_chem_comp_atom.pdbx_model_Cartn_z_ideal',
+      '_chem_comp_bond.atom_id_1',
+      '_chem_comp_bond.atom_id_2',
+      '_chem_comp_bond.value_order',
+      '_chem_comp_bond.pdbx_aromatic_flag',
+  )
+  if missing_keys := set(mandatory_keys) - set(keys):
+    raise ValueError(
+        f'Component {component_name} in the user-defined CCD is missing these'
+        f' keys: {missing_keys}'
+    )
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Input:
+  """AlphaFold input.
+
+  Attributes:
+    name: The name of the target.
+    chains: Protein chains, RNA chains, DNA chains, or ligands.
+    protein_chains: Protein chains.
+    rna_chains: RNA chains.
+    dna_chains: Single strand DNA chains.
+    ligands: Ligand (including ion) inputs.
+    rng_seeds: Random number generator seeds, one for each model execution.
+    bonded_atom_pairs: A list of tuples of atoms that are bonded to each other.
+      Each atom is defined by a tuple of (chain_id, res_id, atom_name). Chain
+      IDs must be set if there are any bonded atoms. Residue IDs are 1-indexed.
+      Atoms in ligands defined by SMILES can't be bonded since SMILES doesn't
+      define unique atom names.
+    user_ccd: Optional user-defined chemical component dictionary in the CIF
+      format. This can be used to provide additional CCD entries that are not
+      present in the default CCD and thus define arbitrary new ligands. This is
+      more expressive than SMILES since it allows to name all atoms within the
+      ligand which in turn makes it possible to define bonds using those atoms.
+  """
+
+  name: str
+  chains: Sequence[ProteinChain | RnaChain | DnaChain | Ligand]
+  rng_seeds: Sequence[int]
+  bonded_atom_pairs: Sequence[tuple[BondAtomId, BondAtomId]] | None = None
+  user_ccd: str | None = None
+
+  def __post_init__(self):
+    if not self.rng_seeds:
+      raise ValueError('Input must have at least one RNG seed.')
+
+    if not self.name.strip() or not self.sanitised_name():
+      raise ValueError(
+          'Input name must be non-empty and contain at least one valid'
+          ' character (letters, numbers, dots, dashes, underscores).'
+      )
+
+    chain_ids = [c.id for c in self.chains]
+    if any(not c.id.isalpha() or c.id.islower() for c in self.chains):
+      raise ValueError(f'IDs must be upper case letters, got: {chain_ids}')
+    if len(set(chain_ids)) != len(chain_ids):
+      raise ValueError('Input JSON contains sequences with duplicate IDs.')
+
+    # Use hashable types for chains, rng_seeds, and bonded_atom_pairs.
+    object.__setattr__(self, 'chains', tuple(self.chains))
+    object.__setattr__(self, 'rng_seeds', tuple(self.rng_seeds))
+    if self.bonded_atom_pairs is not None:
+      object.__setattr__(
+          self, 'bonded_atom_pairs', tuple(self.bonded_atom_pairs)
+      )
+
+    if self.user_ccd is not None:
+      for component_name, component_cif in cif_dict.parse_multi_data_cif(
+          self.user_ccd
+      ).items():
+        _validate_user_ccd_keys(component_cif.keys(), component_name)
+
+  @property
+  def protein_chains(self) -> Sequence[ProteinChain]:
+    return [chain for chain in self.chains if isinstance(chain, ProteinChain)]
+
+  @property
+  def rna_chains(self) -> Sequence[RnaChain]:
+    return [chain for chain in self.chains if isinstance(chain, RnaChain)]
+
+  @property
+  def dna_chains(self) -> Sequence[DnaChain]:
+    return [chain for chain in self.chains if isinstance(chain, DnaChain)]
+
+  @property
+  def ligands(self) -> Sequence[Ligand]:
+    return [chain for chain in self.chains if isinstance(chain, Ligand)]
+
+  def sanitised_name(self) -> str:
+    """Returns sanitised version of the name that can be used as a filename."""
+    spaceless_name = self.name.replace(' ', '_')
+    allowed_chars = set(string.ascii_letters + string.digits + '_-.')
+    return ''.join(l for l in spaceless_name if l in allowed_chars)
+
+  @classmethod
+  def from_alphafoldserver_fold_job(cls, fold_job: Mapping[str, Any]) -> Self:
+    """Constructs Input from an AlphaFoldServer fold job."""
+
+    # Validate the fold job has the correct format.
+    _validate_keys(
+        fold_job.keys(),
+        {'name', 'modelSeeds', 'sequences', 'dialect', 'version'},
+    )
+    if 'dialect' not in fold_job and 'version' not in fold_job:
+      dialect = ALPHAFOLDSERVER_JSON_DIALECT
+      version = ALPHAFOLDSERVER_JSON_VERSION
+    elif 'dialect' in fold_job and 'version' in fold_job:
+      dialect = fold_job['dialect']
+      version = fold_job['version']
+    else:
+      raise ValueError(
+          'AlphaFold Server input JSON must either contain both `dialect` and'
+          ' `version` fields, or neither. If neither is specified, it is'
+          f' assumed that `dialect="{ALPHAFOLDSERVER_JSON_DIALECT}"` and'
+          f' `version="{ALPHAFOLDSERVER_JSON_VERSION}"`.'
+      )
+
+    if dialect != ALPHAFOLDSERVER_JSON_DIALECT:
+      raise ValueError(
+          f'AlphaFold Server input JSON has unsupported dialect: {dialect}, '
+          f'expected {ALPHAFOLDSERVER_JSON_DIALECT}.'
+      )
+
+    # For now, there is only one AlphaFold Server JSON version.
+    if version != ALPHAFOLDSERVER_JSON_VERSION:
+      raise ValueError(
+          f'AlphaFold Server input JSON has unsupported version: {version}, '
+          f'expected {ALPHAFOLDSERVER_JSON_VERSION}.'
+      )
+
+    # Parse the chains.
+    chains = []
+    for sequence in fold_job['sequences']:
+      if 'proteinChain' in sequence:
+        for _ in range(sequence['proteinChain'].get('count', 1)):
+          chains.append(
+              ProteinChain.from_alphafoldserver_dict(
+                  sequence['proteinChain'],
+                  seq_id=mmcif_lib.int_id_to_str_id(len(chains) + 1),
+              )
+          )
+      elif 'rnaSequence' in sequence:
+        for _ in range(sequence['rnaSequence'].get('count', 1)):
+          chains.append(
+              RnaChain.from_alphafoldserver_dict(
+                  sequence['rnaSequence'],
+                  seq_id=mmcif_lib.int_id_to_str_id(len(chains) + 1),
+              )
+          )
+      elif 'dnaSequence' in sequence:
+        for _ in range(sequence['dnaSequence'].get('count', 1)):
+          chains.append(
+              DnaChain.from_alphafoldserver_dict(
+                  sequence['dnaSequence'],
+                  seq_id=mmcif_lib.int_id_to_str_id(len(chains) + 1),
+              )
+          )
+      elif 'ion' in sequence:
+        for _ in range(sequence['ion'].get('count', 1)):
+          chains.append(
+              Ligand.from_alphafoldserver_dict(
+                  sequence['ion'],
+                  seq_id=mmcif_lib.int_id_to_str_id(len(chains) + 1),
+              )
+          )
+      elif 'ligand' in sequence:
+        for _ in range(sequence['ligand'].get('count', 1)):
+          chains.append(
+              Ligand.from_alphafoldserver_dict(
+                  sequence['ligand'],
+                  seq_id=mmcif_lib.int_id_to_str_id(len(chains) + 1),
+              )
+          )
+      else:
+        raise ValueError(f'Unknown sequence type: {sequence}')
+
+    if 'modelSeeds' in fold_job and fold_job['modelSeeds']:
+      rng_seeds = [int(seed) for seed in fold_job['modelSeeds']]
+    else:
+      rng_seeds = [_sample_rng_seed()]
+
+    return cls(name=fold_job['name'], chains=chains, rng_seeds=rng_seeds)
+
+  @classmethod
+  def from_json(
+      cls, json_str: str, json_path: pathlib.Path | None = None
+  ) -> Self:
+    """Loads the input from the AlphaFold JSON string."""
+    raw_json = json.loads(json_str)
+
+    _validate_keys(
+        raw_json.keys(),
+        {
+            'dialect',
+            'version',
+            'name',
+            'modelSeeds',
+            'sequences',
+            'bondedAtomPairs',
+            'userCCD',
+            'userCCDPath',
+        },
+    )
+
+    if 'dialect' not in raw_json or 'version' not in raw_json:
+      raise ValueError(
+          'AlphaFold 3 input JSON must contain `dialect` and `version` fields.'
+      )
+
+    if raw_json['dialect'] != JSON_DIALECT:
+      raise ValueError(
+          'AlphaFold 3 input JSON has unsupported dialect:'
+          f' {raw_json["dialect"]}, expected {JSON_DIALECT}.'
+      )
+
+    if raw_json['version'] not in JSON_VERSIONS:
+      raise ValueError(
+          'AlphaFold 3 input JSON has unsupported version:'
+          f' {raw_json["version"]}, expected one of {JSON_VERSIONS}.'
+      )
+
+    if 'sequences' not in raw_json:
+      raise ValueError('AlphaFold 3 input JSON does not contain any sequences.')
+
+    if 'modelSeeds' not in raw_json or not raw_json['modelSeeds']:
+      raise ValueError(
+          'AlphaFold 3 input JSON must specify at least one rng seed in'
+          ' `modelSeeds`.'
+      )
+
+    sequences = raw_json['sequences']
+
+    # Make sure sequence IDs are all set.
+    raw_sequence_ids = [next(iter(s.values())).get('id') for s in sequences]
+    if all(raw_sequence_ids):
+      sequence_ids = []
+      for sequence_id in raw_sequence_ids:
+        if isinstance(sequence_id, list):
+          sequence_ids.append(sequence_id)
+        else:
+          sequence_ids.append([sequence_id])
+    else:
+      raise ValueError(
+          'AlphaFold 3 input JSON contains sequences with unset IDs.'
+      )
+
+    flat_seq_ids = []
+    for seq_ids in sequence_ids:
+      flat_seq_ids.extend(seq_ids)
+
+    chains = []
+    for seq_ids, sequence in zip(sequence_ids, sequences, strict=True):
+      if len(sequence) != 1:
+        raise ValueError(f'Chain {seq_ids} has more than 1 sequence.')
+      for seq_id in seq_ids:
+        if 'protein' in sequence:
+          chains.append(ProteinChain.from_dict(sequence, json_path, seq_id))
+        elif 'rna' in sequence:
+          chains.append(RnaChain.from_dict(sequence, json_path, seq_id))
+        elif 'dna' in sequence:
+          chains.append(DnaChain.from_dict(sequence, seq_id=seq_id))
+        elif 'ligand' in sequence:
+          chains.append(Ligand.from_dict(sequence, seq_id=seq_id))
+        else:
+          raise ValueError(f'Unknown sequence type: {sequence}')
+
+    smiles_ligand_ids = set(
+        c.id for c in chains if isinstance(c, Ligand) and c.smiles is not None
+    )
+    chain_lengths = {chain.id: len(chain) for chain in chains}
+    bonded_atom_pairs = None
+    if bonds := raw_json.get('bondedAtomPairs'):
+      bonded_atom_pairs = []
+      for bond in bonds:
+        if len(bond) != 2:
+          raise ValueError(f'Bond {bond} must have 2 atoms, got {len(bond)}.')
+        bond_beg, bond_end = bond
+        if (
+            len(bond_beg) != 3
+            or not isinstance(bond_beg[0], str)
+            or not isinstance(bond_beg[1], int)
+            or not isinstance(bond_beg[2], str)
+        ):
+          raise ValueError(
+              f'Atom {bond_beg} in bond {bond} must have 3 components: '
+              '(chain_id: str, res_id: int, atom_name: str).'
+          )
+        if (
+            len(bond_end) != 3
+            or not isinstance(bond_end[0], str)
+            or not isinstance(bond_end[1], int)
+            or not isinstance(bond_end[2], str)
+        ):
+          raise ValueError(
+              f'Atom {bond_end} in bond {bond} must have 3 components: '
+              '(chain_id: str, res_id: int, atom_name: str).'
+          )
+        if bond_beg[0] not in flat_seq_ids or bond_end[0] not in flat_seq_ids:
+          raise ValueError(f'Invalid chain ID(s) in bond {bond}')
+        if (
+            not 0 < bond_beg[1] <= chain_lengths[bond_beg[0]]
+            or not 0 < bond_end[1] <= chain_lengths[bond_end[0]]
+        ):
+          raise ValueError(f'Invalid residue ID(s) in bond {bond}')
+        if bond_beg[0] in smiles_ligand_ids:
+          raise ValueError(
+              f'Bond {bond} involves an unsupported SMILES ligand {bond_beg[0]}'
+          )
+        if bond_end[0] in smiles_ligand_ids:
+          raise ValueError(
+              f'Bond {bond} involves an unsupported SMILES ligand {bond_end[0]}'
+          )
+        bonded_atom_pairs.append((tuple(bond_beg), tuple(bond_end)))
+
+      if len(bonded_atom_pairs) != len(set(bonded_atom_pairs)):
+        raise ValueError(f'Bonds are not unique: {bonded_atom_pairs}')
+
+    user_ccd = raw_json.get('userCCD')
+    user_ccd_path = raw_json.get('userCCDPath')
+    if user_ccd and user_ccd_path:
+      raise ValueError('Only one of userCCD/userCCDPath can be set.')
+    if user_ccd and len(user_ccd) < 256 and os.path.exists(user_ccd):
+      raise ValueError('Set the user CCD path using the "userCCDPath" field.')
+    elif user_ccd_path:
+      user_ccd = _read_file(pathlib.Path(user_ccd_path), json_path)
+
+    return cls(
+        name=raw_json['name'],
+        chains=chains,
+        rng_seeds=[int(seed) for seed in raw_json['modelSeeds']],
+        bonded_atom_pairs=bonded_atom_pairs,
+        user_ccd=user_ccd,
+    )
+
+  @classmethod
+  def from_mmcif(cls, mmcif_str: str, ccd: chemical_components.Ccd) -> Self:
+    """Loads the input from an mmCIF string.
+
+    WARNING: Since rng seeds are not stored in mmCIFs, an rng seed is sampled
+    in the returned `Input`.
+
+    Args:
+      mmcif_str: The mmCIF string.
+      ccd: The chemical components dictionary.
+
+    Returns:
+      The input in an Input format.
+    """
+
+    struc = structure.from_mmcif(
+        mmcif_str,
+        # Change MSE residues to MET residues.
+        fix_mse_residues=True,
+        # Fix arginine atom names. This is not needed since the input discards
+        # any atom-level data, but kept for consistency with the paper.
+        fix_arginines=True,
+        # Fix unknown DNA residues to the correct unknown DNA residue type.
+        fix_unknown_dna=True,
+        # Do not include water molecules.
+        include_water=False,
+        # Do not include things like DNA/RNA hybrids. This will be changed once
+        # we have a way of handling these in the AlphaFold 3 input format.
+        include_other=False,
+        # Include the specific bonds defined in the mmCIF bond table, e.g.
+        # covalent bonds for PTMs.
+        include_bonds=True,
+    )
+
+    # Create default bioassembly, expanding structures implied by stoichiometry.
+    struc = struc.generate_bioassembly(None)
+
+    sequences = struc.chain_single_letter_sequence(
+        include_missing_residues=True
+    )
+
+    chains = []
+    for chain_id, chain_type in zip(
+        struc.group_by_chain.chain_id, struc.group_by_chain.chain_type
+    ):
+      sequence = sequences[chain_id]
+
+      if chain_type in mmcif_names.NON_POLYMER_CHAIN_TYPES:
+        residues = list(struc.chain_res_name_sequence()[chain_id])
+        if all(ccd.get(res) is not None for res in residues):
+          chains.append(Ligand(id=chain_id, ccd_ids=residues))
+        elif len(residues) == 1:
+          comp_name = residues[0]
+          comps = struc.chemical_components_data
+          if comps is None:
+            raise ValueError(
+                'Missing mmCIF chemical components data - this is required for '
+                f'a non-CCD ligand {comp_name} defined using SMILES string.'
+            )
+          chains.append(
+              Ligand(id=chain_id, smiles=comps.chem_comp[comp_name].pdbx_smiles)
+          )
+        else:
+          raise ValueError(
+              'Multi-component ligand must be defined using CCD IDs, defining'
+              ' using SMILES is supported only for single-component ligands. '
+              f'Got {residues}'
+          )
+      else:
+        residues = struc.chain_res_name_sequence()[chain_id]
+        fixed = struc.chain_res_name_sequence(
+            fix_non_standard_polymer_res=True
+        )[chain_id]
+        modifications = [
+            (orig, i + 1)
+            for i, (orig, fixed) in enumerate(zip(residues, fixed, strict=True))
+            if orig != fixed
+        ]
+
+        if chain_type == mmcif_names.PROTEIN_CHAIN:
+          chains.append(
+              ProteinChain(id=chain_id, sequence=sequence, ptms=modifications)
+          )
+        elif chain_type == mmcif_names.RNA_CHAIN:
+          chains.append(
+              RnaChain(
+                  id=chain_id, sequence=sequence, modifications=modifications
+              )
+          )
+        elif chain_type == mmcif_names.DNA_CHAIN:
+          chains.append(
+              DnaChain(
+                  id=chain_id, sequence=sequence, modifications=modifications
+              )
+          )
+
+    bonded_atom_pairs = []
+    chain_ids = set(c.id for c in chains)
+    for atom_a, atom_b, _ in struc.iter_bonds():
+      if atom_a['chain_id'] in chain_ids and atom_b['chain_id'] in chain_ids:
+        beg = (atom_a['chain_id'], int(atom_a['res_id']), atom_a['atom_name'])
+        end = (atom_b['chain_id'], int(atom_b['res_id']), atom_b['atom_name'])
+        bonded_atom_pairs.append((beg, end))
+
+    return cls(
+        name=struc.name,
+        chains=chains,
+        # mmCIFs don't store rng seeds, so we need to sample one here.
+        rng_seeds=[_sample_rng_seed()],
+        bonded_atom_pairs=bonded_atom_pairs or None,
+    )
+
+  def to_structure(self, ccd: chemical_components.Ccd) -> structure.Structure:
+    """Converts Input to a Structure.
+
+    WARNING: This method does not preserve the rng seeds.
+
+    Args:
+      ccd: The chemical components dictionary.
+
+    Returns:
+      The input in a structure.Structure format.
+    """
+    ids: list[str] = []
+    sequences: list[str] = []
+    poly_types: list[str] = []
+    formats: list[structure.SequenceFormat] = []
+
+    for chain in self.chains:
+      ids.append(chain.id)
+      match chain:
+        case ProteinChain():
+          sequences.append('(' + ')('.join(chain.to_ccd_sequence()) + ')')
+          poly_types.append(mmcif_names.PROTEIN_CHAIN)
+          formats.append(structure.SequenceFormat.CCD_CODES)
+        case RnaChain():
+          sequences.append('(' + ')('.join(chain.to_ccd_sequence()) + ')')
+          poly_types.append(mmcif_names.RNA_CHAIN)
+          formats.append(structure.SequenceFormat.CCD_CODES)
+        case DnaChain():
+          sequences.append('(' + ')('.join(chain.to_ccd_sequence()) + ')')
+          poly_types.append(mmcif_names.DNA_CHAIN)
+          formats.append(structure.SequenceFormat.CCD_CODES)
+        case Ligand():
+          if chain.ccd_ids is not None:
+            sequences.append('(' + ')('.join(chain.ccd_ids) + ')')
+            if len(chain.ccd_ids) == 1:
+              poly_types.append(mmcif_names.NON_POLYMER_CHAIN)
+            else:
+              poly_types.append(mmcif_names.BRANCHED_CHAIN)
+            formats.append(structure.SequenceFormat.CCD_CODES)
+          elif chain.smiles is not None:
+            # Convert to `<unique ligand ID>:<smiles>` format that is expected
+            # by structure.from_sequences_and_bonds.
+            sequences.append(f'LIG_{chain.id}:{chain.smiles}')
+            poly_types.append(mmcif_names.NON_POLYMER_CHAIN)
+            formats.append(structure.SequenceFormat.LIGAND_SMILES)
+          else:
+            raise ValueError('Ligand must have one of CCD ID or SMILES set.')
+
+    # Remap bond chain IDs from chain IDs to chain indices and convert to
+    # 0-based residue indexing.
+    bonded_atom_pairs = []
+    chain_indices = {cid: i for i, cid in enumerate(ids)}
+    if self.bonded_atom_pairs is not None:
+      for bond_beg, bond_end in self.bonded_atom_pairs:
+        bonded_atom_pairs.append((
+            (chain_indices[bond_beg[0]], bond_beg[1] - 1, bond_beg[2]),
+            (chain_indices[bond_end[0]], bond_end[1] - 1, bond_end[2]),
+        ))
+
+    return structure.from_sequences_and_bonds(
+        sequences=sequences,
+        chain_types=poly_types,
+        sequence_formats=formats,
+        chain_ids=ids,
+        bonded_atom_pairs=bonded_atom_pairs,
+        ccd=ccd,
+        name=self.sanitised_name(),
+        bond_type=mmcif_names.COVALENT_BOND,
+        release_date=None,
+    )
+
+  def to_json(self) -> str:
+    """Converts Input to an AlphaFold JSON."""
+    deduped_chains = {}
+    deduped_chain_ids = {}
+    for chain in self.chains:
+      deduped_chains[chain.hash_without_id()] = chain
+      deduped_chain_ids.setdefault(chain.hash_without_id(), []).append(chain.id)
+
+    sequences = []
+    for chain_content_hash, ids in deduped_chain_ids.items():
+      chain = deduped_chains[chain_content_hash]
+      sequences.append(chain.to_dict(seq_id=ids if len(ids) > 1 else ids[0]))
+
+    alphafold_json = json.dumps(
+        {
+            'dialect': JSON_DIALECT,
+            'version': JSON_VERSION,
+            'name': self.name,
+            'sequences': sequences,
+            'modelSeeds': self.rng_seeds,
+            'bondedAtomPairs': self.bonded_atom_pairs,
+            'userCCD': self.user_ccd,
+        },
+        indent=2,
+    )
+    # Remove newlines from the query/template indices arrays. We match the
+    # queryIndices/templatesIndices with a non-capturing group. We then match
+    # the entire region between the square brackets by looking for lines
+    # containing only whitespace, number, or a comma.
+    return re.sub(
+        r'("(?:queryIndices|templateIndices)": \[)([\s\n\d,]+)(\],?)',
+        lambda mtch: mtch[1] + re.sub(r'\n\s+', ' ', mtch[2].strip()) + mtch[3],
+        alphafold_json,
+    )
+
+  def fill_missing_fields(self) -> Self:
+    """Fill missing MSA and template fields with default values."""
+    with_missing_fields = [
+        c.fill_missing_fields()
+        if isinstance(c, (ProteinChain, RnaChain))
+        else c
+        for c in self.chains
+    ]
+    return dataclasses.replace(self, chains=with_missing_fields)
+
+  def with_multiple_seeds(self, num_seeds: int) -> Self:
+    """Returns a copy of the input with num_seeds rng seeds."""
+    if num_seeds <= 1:
+      raise ValueError('Number of seeds must be greater than 1.')
+    if len(self.rng_seeds) != 1:
+      raise ValueError('Input must have one rng seed to set multiple seeds.')
+
+    return dataclasses.replace(
+        self,
+        rng_seeds=list(range(self.rng_seeds[0], self.rng_seeds[0] + num_seeds)),
+    )
+
+
+def load_fold_inputs_from_path(json_path: pathlib.Path) -> Iterator[Input]:
+  """Loads multiple fold inputs from a JSON string."""
+  with open(json_path, 'r') as f:
+    json_str = f.read()
+
+  # Parse the JSON string, so we can detect its format.
+  raw_json = json.loads(json_str)
+
+  if isinstance(raw_json, list):
+    # AlphaFold Server JSON.
+    logging.info('Loading %d fold jobs from %s', len(raw_json), json_path)
+    for fold_job_idx, fold_job in enumerate(raw_json):
+      try:
+        yield Input.from_alphafoldserver_fold_job(fold_job)
+      except ValueError as e:
+        raise ValueError(
+            f'Failed to load fold job {fold_job_idx} from {json_path}'
+            f' (AlphaFold Server dialect): {e}'
+        ) from e
+  else:
+    # AlphaFold 3 JSON.
+    try:
+      yield Input.from_json(json_str, json_path)
+    except ValueError as e:
+      raise ValueError(
+          f'Failed to load input from {json_path} (AlphaFold 3 dialect): {e}'
+      ) from e
+
+
+def load_fold_inputs_from_dir(input_dir: pathlib.Path) -> Iterator[Input]:
+  """Loads multiple fold inputs from all JSON files in a given input_dir.
+
+  Args:
+    input_dir: The directory containing the JSON files.
+
+  Yields:
+    The fold inputs from all JSON files in the input directory.
+  """
+  for file_path in sorted(input_dir.glob('*.json')):
+    if not file_path.is_file():
+      continue
+
+    yield from load_fold_inputs_from_path(file_path)

flax_model/alphafold3/common/resources.py

@@ -0,0 +1,68 @@
+
+
+"""Load external resources, such as external tools or data resources."""
+
+from collections.abc import Iterator
+import os
+import pathlib
+import typing
+from typing import BinaryIO, Final, Literal, TextIO
+
+from importlib import resources
+import flax_model.alphafold3.common
+
+
+_DATA_ROOT:  Final[pathlib.Path] = (
+    resources.files(flax_model.alphafold3.common).joinpath('..').resolve()
+)
+ROOT = _DATA_ROOT
+
+
+def filename(name: str | os.PathLike[str]) -> str:
+  """Returns the absolute path to an external resource.
+
+  Note that this calls resources.GetResourceFilename under the hood and hence
+  causes par file unpacking, which might be unfriendly on diskless machines.
+
+
+  Args:
+    name: the name of the resource corresponding to its path relative to the
+      root of the repository.
+  """
+  return (_DATA_ROOT / name).as_posix()
+
+
+@typing.overload
+def open_resource(
+    name: str | os.PathLike[str], mode: Literal['r', 'rt'] = 'rt'
+) -> TextIO:
+  ...
+
+
+@typing.overload
+def open_resource(
+    name: str | os.PathLike[str], mode: Literal['rb']
+) -> BinaryIO:
+  ...
+
+
+def open_resource(
+    name: str | os.PathLike[str], mode: str = 'rb'
+) -> TextIO | BinaryIO:
+  """Returns an open file object for the named resource.
+
+  Args:
+    name: the name of the resource corresponding to its path relative to the
+      root of the repository.
+    mode: the mode to use when opening the file.
+  """
+  return (_DATA_ROOT / name).open(mode)
+
+
+def get_resource_dir(path: str | os.PathLike[str]) -> os.PathLike[str]:
+  return _DATA_ROOT / path
+
+
+def walk(path: str) -> Iterator[tuple[str, list[str], list[str]]]:
+  """Walks the directory tree of resources similar to os.walk."""
+  return os.walk((_DATA_ROOT / path).as_posix())

flax_model/alphafold3/common/testing/data.py

@@ -0,0 +1,61 @@
+
+
+"""Module that provides an abstraction for accessing test data."""
+
+import os
+import pathlib
+from typing import Literal, overload
+
+from absl.testing import absltest
+
+
+class Data:
+  """Provides an abstraction for accessing test data."""
+
+  def __init__(self, data_dir: os.PathLike[str] | str):
+    """Initiailizes data wrapper, providing users with high level data access.
+
+    Args:
+      data_dir: Directory containing test data.
+    """
+    self._data_dir = pathlib.Path(data_dir)
+
+  def path(self, data_name: str | os.PathLike[str] | None = None) -> str:
+    """Returns the path to a given test data.
+
+    Args:
+      data_name: the name of the test data file relative to data_dir. If not
+        set, this will return the absolute path to the data directory.
+    """
+    data_dir_path = (
+        pathlib.Path(absltest.get_default_test_srcdir()) / self._data_dir
+    )
+
+    if data_name:
+      return str(data_dir_path / data_name)
+
+    return str(data_dir_path)
+
+  @overload
+  def load(
+      self, data_name: str | os.PathLike[str], mode: Literal['rt'] = 'rt'
+  ) -> str:
+    ...
+
+  @overload
+  def load(
+      self, data_name: str | os.PathLike[str], mode: Literal['rb'] = 'rb'
+  ) -> bytes:
+    ...
+
+  def load(
+      self, data_name: str | os.PathLike[str], mode: str = 'rt'
+  ) -> str | bytes:
+    """Returns the contents of a given test data.
+
+    Args:
+      data_name: the name of the test data file relative to data_dir.
+      mode: the mode in which to read the data file. Defaults to text ('rt').
+    """
+    with open(self.path(data_name), mode=mode) as f:
+      return f.read()

flax_model/alphafold3/constants/atom_types.py

@@ -0,0 +1,253 @@
+
+
+"""List of atom types with reverse look-up."""
+
+from collections.abc import Mapping, Sequence, Set
+import itertools
+import sys
+from typing import Final
+from flax_model.alphafold3.constants import residue_names
+
+# Note:
+# `sys.intern` places the values in the Python internal db for fast lookup.
+
+# 37 common residue atoms.
+N = sys.intern('N')
+CA = sys.intern('CA')
+C = sys.intern('C')
+CB = sys.intern('CB')
+O = sys.intern('O')
+CG = sys.intern('CG')
+CG1 = sys.intern('CG1')
+CG2 = sys.intern('CG2')
+OG = sys.intern('OG')
+OG1 = sys.intern('OG1')
+SG = sys.intern('SG')
+CD = sys.intern('CD')
+CD1 = sys.intern('CD1')
+CD2 = sys.intern('CD2')
+ND1 = sys.intern('ND1')
+ND2 = sys.intern('ND2')
+OD1 = sys.intern('OD1')
+OD2 = sys.intern('OD2')
+SD = sys.intern('SD')
+CE = sys.intern('CE')
+CE1 = sys.intern('CE1')
+CE2 = sys.intern('CE2')
+CE3 = sys.intern('CE3')
+NE = sys.intern('NE')
+NE1 = sys.intern('NE1')
+NE2 = sys.intern('NE2')
+OE1 = sys.intern('OE1')
+OE2 = sys.intern('OE2')
+CH2 = sys.intern('CH2')
+NH1 = sys.intern('NH1')
+NH2 = sys.intern('NH2')
+OH = sys.intern('OH')
+CZ = sys.intern('CZ')
+CZ2 = sys.intern('CZ2')
+CZ3 = sys.intern('CZ3')
+NZ = sys.intern('NZ')
+OXT = sys.intern('OXT')
+
+# 29 common nucleic acid atoms.
+C1PRIME = sys.intern("C1'")
+C2 = sys.intern('C2')
+C2PRIME = sys.intern("C2'")
+C3PRIME = sys.intern("C3'")
+C4 = sys.intern('C4')
+C4PRIME = sys.intern("C4'")
+C5 = sys.intern('C5')
+C5PRIME = sys.intern("C5'")
+C6 = sys.intern('C6')
+C7 = sys.intern('C7')
+C8 = sys.intern('C8')
+N1 = sys.intern('N1')
+N2 = sys.intern('N2')
+N3 = sys.intern('N3')
+N4 = sys.intern('N4')
+N6 = sys.intern('N6')
+N7 = sys.intern('N7')
+N9 = sys.intern('N9')
+O2 = sys.intern('O2')
+O2PRIME = sys.intern("O2'")
+O3PRIME = sys.intern("O3'")
+O4 = sys.intern('O4')
+O4PRIME = sys.intern("O4'")
+O5PRIME = sys.intern("O5'")
+O6 = sys.intern('O6')
+OP1 = sys.intern('OP1')
+OP2 = sys.intern('OP2')
+OP3 = sys.intern('OP3')
+P = sys.intern('P')
+
+# A list of atoms (excluding hydrogen) for each AA type. PDB naming convention.
+RESIDUE_ATOMS: Mapping[str, tuple[str, ...]] = {
+    residue_names.ALA: (C, CA, CB, N, O),
+    residue_names.ARG: (C, CA, CB, CG, CD, CZ, N, NE, O, NH1, NH2),
+    residue_names.ASN: (C, CA, CB, CG, N, ND2, O, OD1),
+    residue_names.ASP: (C, CA, CB, CG, N, O, OD1, OD2),
+    residue_names.CYS: (C, CA, CB, N, O, SG),
+    residue_names.GLN: (C, CA, CB, CG, CD, N, NE2, O, OE1),
+    residue_names.GLU: (C, CA, CB, CG, CD, N, O, OE1, OE2),
+    residue_names.GLY: (C, CA, N, O),
+    residue_names.HIS: (C, CA, CB, CG, CD2, CE1, N, ND1, NE2, O),
+    residue_names.ILE: (C, CA, CB, CG1, CG2, CD1, N, O),
+    residue_names.LEU: (C, CA, CB, CG, CD1, CD2, N, O),
+    residue_names.LYS: (C, CA, CB, CG, CD, CE, N, NZ, O),
+    residue_names.MET: (C, CA, CB, CG, CE, N, O, SD),
+    residue_names.PHE: (C, CA, CB, CG, CD1, CD2, CE1, CE2, CZ, N, O),
+    residue_names.PRO: (C, CA, CB, CG, CD, N, O),
+    residue_names.SER: (C, CA, CB, N, O, OG),
+    residue_names.THR: (C, CA, CB, CG2, N, O, OG1),
+    residue_names.TRP:
+        (C, CA, CB, CG, CD1, CD2, CE2, CE3, CZ2, CZ3, CH2, N, NE1, O),
+    residue_names.TYR: (C, CA, CB, CG, CD1, CD2, CE1, CE2, CZ, N, O, OH),
+    residue_names.VAL: (C, CA, CB, CG1, CG2, N, O),
+}  # pyformat: disable
+
+# Used to identify backbone for alignment and distance calculation for sterics.
+PROTEIN_BACKBONE_ATOMS: tuple[str, ...] = (N, CA, C)
+
+# Naming swaps for ambiguous atom names. Due to symmetries in the amino acids
+# the naming of atoms is ambiguous in 4 of the 20 amino acids. (The LDDT paper
+# lists 7 amino acids as ambiguous, but the naming ambiguities in LEU, VAL and
+# ARG can be resolved by using the 3D constellations of the 'ambiguous' atoms
+# and their neighbours)
+AMBIGUOUS_ATOM_NAMES: Mapping[str, Mapping[str, str]] = {
+    residue_names.ASP: {OD1: OD2},
+    residue_names.GLU: {OE1: OE2},
+    residue_names.PHE: {CD1: CD2, CE1: CE2},
+    residue_names.TYR: {CD1: CD2, CE1: CE2},
+}
+
+# Used when we need to store atom data in a format that requires fixed atom data
+# size for every protein residue (e.g. a numpy array).
+ATOM37: tuple[str, ...] = (
+    N, CA, C, CB, O, CG, CG1, CG2, OG, OG1, SG, CD, CD1, CD2, ND1, ND2, OD1,
+    OD2, SD, CE, CE1, CE2, CE3, NE, NE1, NE2, OE1, OE2, CH2, NH1, NH2, OH, CZ,
+    CZ2, CZ3, NZ, OXT)  # pyformat: disable
+ATOM37_ORDER: Mapping[str, int] = {name: i for i, name in enumerate(ATOM37)}
+ATOM37_NUM: Final[int] = len(ATOM37)  # := 37.
+
+# Used when we need to store protein atom data in a format that requires fixed
+# atom data size for any residue but takes less space than ATOM37 by having 14
+# fields, which is sufficient for storing atoms of all protein residues (e.g. a
+# numpy array).
+ATOM14: Mapping[str, tuple[str, ...]] = {
+    residue_names.ALA: (N, CA, C, O, CB),
+    residue_names.ARG: (N, CA, C, O, CB, CG, CD, NE, CZ, NH1, NH2),
+    residue_names.ASN: (N, CA, C, O, CB, CG, OD1, ND2),
+    residue_names.ASP: (N, CA, C, O, CB, CG, OD1, OD2),
+    residue_names.CYS: (N, CA, C, O, CB, SG),
+    residue_names.GLN: (N, CA, C, O, CB, CG, CD, OE1, NE2),
+    residue_names.GLU: (N, CA, C, O, CB, CG, CD, OE1, OE2),
+    residue_names.GLY: (N, CA, C, O),
+    residue_names.HIS: (N, CA, C, O, CB, CG, ND1, CD2, CE1, NE2),
+    residue_names.ILE: (N, CA, C, O, CB, CG1, CG2, CD1),
+    residue_names.LEU: (N, CA, C, O, CB, CG, CD1, CD2),
+    residue_names.LYS: (N, CA, C, O, CB, CG, CD, CE, NZ),
+    residue_names.MET: (N, CA, C, O, CB, CG, SD, CE),
+    residue_names.PHE: (N, CA, C, O, CB, CG, CD1, CD2, CE1, CE2, CZ),
+    residue_names.PRO: (N, CA, C, O, CB, CG, CD),
+    residue_names.SER: (N, CA, C, O, CB, OG),
+    residue_names.THR: (N, CA, C, O, CB, OG1, CG2),
+    residue_names.TRP:
+        (N, CA, C, O, CB, CG, CD1, CD2, NE1, CE2, CE3, CZ2, CZ3, CH2),
+    residue_names.TYR: (N, CA, C, O, CB, CG, CD1, CD2, CE1, CE2, CZ, OH),
+    residue_names.VAL: (N, CA, C, O, CB, CG1, CG2),
+    residue_names.UNK: (),
+}  # pyformat: disable
+
+# A compact atom encoding with 14 columns, padded with '' in empty slots.
+ATOM14_PADDED: Mapping[str, Sequence[str]] = {
+    k: [v for _, v in itertools.zip_longest(range(14), values, fillvalue='')]
+    for k, values in ATOM14.items()
+}
+
+ATOM14_ORDER: Mapping[str, Mapping[str, int]] = {
+    k: {name: i for i, name in enumerate(v)} for k, v in ATOM14.items()
+}
+ATOM14_NUM: Final[int] = max(len(v) for v in ATOM14.values())
+
+# Used when we need to store protein and nucleic atom library.
+DENSE_ATOM: Mapping[str, tuple[str, ...]] = {
+    # Protein.
+    residue_names.ALA: (N, CA, C, O, CB),
+    residue_names.ARG: (N, CA, C, O, CB, CG, CD, NE, CZ, NH1, NH2),
+    residue_names.ASN: (N, CA, C, O, CB, CG, OD1, ND2),
+    residue_names.ASP: (N, CA, C, O, CB, CG, OD1, OD2),
+    residue_names.CYS: (N, CA, C, O, CB, SG),
+    residue_names.GLN: (N, CA, C, O, CB, CG, CD, OE1, NE2),
+    residue_names.GLU: (N, CA, C, O, CB, CG, CD, OE1, OE2),
+    residue_names.GLY: (N, CA, C, O),
+    residue_names.HIS: (N, CA, C, O, CB, CG, ND1, CD2, CE1, NE2),
+    residue_names.ILE: (N, CA, C, O, CB, CG1, CG2, CD1),
+    residue_names.LEU: (N, CA, C, O, CB, CG, CD1, CD2),
+    residue_names.LYS: (N, CA, C, O, CB, CG, CD, CE, NZ),
+    residue_names.MET: (N, CA, C, O, CB, CG, SD, CE),
+    residue_names.PHE: (N, CA, C, O, CB, CG, CD1, CD2, CE1, CE2, CZ),
+    residue_names.PRO: (N, CA, C, O, CB, CG, CD),
+    residue_names.SER: (N, CA, C, O, CB, OG),
+    residue_names.THR: (N, CA, C, O, CB, OG1, CG2),
+    residue_names.TRP:
+        (N, CA, C, O, CB, CG, CD1, CD2, NE1, CE2, CE3, CZ2, CZ3, CH2),
+    residue_names.TYR: (N, CA, C, O, CB, CG, CD1, CD2, CE1, CE2, CZ, OH),
+    residue_names.VAL: (N, CA, C, O, CB, CG1, CG2),
+    residue_names.UNK: (),
+    # RNA.
+    residue_names.A:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, O2PRIME, C1PRIME, N9, C8, N7, C5, C6, N6, N1, C2, N3, C4),
+    residue_names.C:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, O2PRIME, C1PRIME, N1, C2, O2, N3, C4, N4, C5, C6),
+    residue_names.G:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, O2PRIME, C1PRIME, N9, C8, N7, C5, C6, O6, N1, C2, N2, N3, C4),
+    residue_names.U:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, O2PRIME, C1PRIME, N1, C2, O2, N3, C4, O4, C5, C6),
+    residue_names.UNK_RNA: (),
+    # DNA.
+    residue_names.DA:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, C1PRIME, N9, C8, N7, C5, C6, N6, N1, C2, N3, C4),
+    residue_names.DC:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, C1PRIME, N1, C2, O2, N3, C4, N4, C5, C6),
+    residue_names.DG:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, C1PRIME, N9, C8, N7, C5, C6, O6, N1, C2, N2, N3, C4),
+    residue_names.DT:
+        (OP3, P, OP1, OP2, O5PRIME, C5PRIME, C4PRIME, O4PRIME, C3PRIME, O3PRIME,
+         C2PRIME, C1PRIME, N1, C2, O2, N3, C4, O4, C5, C7, C6),
+    # Unknown nucleic.
+    residue_names.UNK_DNA: (),
+}  # pyformat: disable
+
+DENSE_ATOM_ORDER: Mapping[str, Mapping[str, int]] = {
+    k: {name: i for i, name in enumerate(v)} for k, v in DENSE_ATOM.items()
+}
+DENSE_ATOM_NUM: Final[int] = max(len(v) for v in DENSE_ATOM.values())
+
+# Used when we need to store atom data in a format that requires fixed atom data
+# size for every nucleic molecule (e.g. a numpy array).
+ATOM29: tuple[str, ...] = (
+    "C1'", 'C2', "C2'", "C3'", 'C4', "C4'", 'C5', "C5'", 'C6', 'C7', 'C8', 'N1',
+    'N2', 'N3', 'N4', 'N6', 'N7', 'N9', 'OP3', 'O2', "O2'", "O3'", 'O4', "O4'",
+    "O5'", 'O6', 'OP1', 'OP2', 'P')  # pyformat: disable
+ATOM29_ORDER: Mapping[str, int] = {
+    atom_type: i for i, atom_type in enumerate(ATOM29)
+}
+ATOM29_NUM: Final[int] = len(ATOM29)  # := 29
+
+# Hydrogens that exist depending on the protonation state of the residue.
+# Extracted from third_party/py/openmm/app/data/hydrogens.xml
+PROTONATION_HYDROGENS: Mapping[str, Set[str]] = {
+    'ASP': {'HD2'},
+    'CYS': {'HG'},
+    'GLU': {'HE2'},
+    'HIS': {'HD1', 'HE2'},
+    'LYS': {'HZ3'},
+}

flax_model/alphafold3/constants/chemical_component_sets.py

@@ -0,0 +1,29 @@
+
+
+"""Sets of chemical components."""
+
+import pickle
+from typing import Final
+
+from flax_model.alphafold3.common import resources
+
+
+_CCD_SETS_CCD_PICKLE_FILE = resources.filename(
+    resources.ROOT / 'constants/converters/chemical_component_sets.pickle'
+)
+
+_CCD_SET = pickle.load(open(_CCD_SETS_CCD_PICKLE_FILE, 'rb'))
+
+# Glycan (or 'Saccharide') ligands.
+# _chem_comp.type containing 'saccharide' and 'linking' (when lower-case).
+GLYCAN_LINKING_LIGANDS: Final[frozenset[str]] = _CCD_SET['glycans_linking']
+
+# _chem_comp.type containing 'saccharide' and not 'linking' (when lower-case).
+GLYCAN_OTHER_LIGANDS: Final[frozenset[str]] = _CCD_SET['glycans_other']
+
+# Each of these molecules appears in over 1k PDB structures, are used to
+# facilitate crystallization conditions, but do not have biological relevance.
+COMMON_CRYSTALLIZATION_AIDS: Final[frozenset[str]] = frozenset({
+    'SO4', 'GOL', 'EDO', 'PO4', 'ACT', 'PEG', 'DMS', 'TRS', 'PGE', 'PG4', 'FMT',
+    'EPE', 'MPD', 'MES', 'CD', 'IOD',
+})  # pyformat: disable

flax_model/alphafold3/constants/chemical_components.py

@@ -0,0 +1,190 @@
+
+
+"""Chemical Components found in PDB (CCD) constants."""
+
+from collections.abc import ItemsView, Iterator, KeysView, Mapping, Sequence, ValuesView
+import dataclasses
+import functools
+import os
+import pickle
+
+from flax_model.alphafold3.common import resources
+from flax_model.alphafold3.cpp import cif_dict
+
+
+_CCD_PICKLE_FILE = resources.filename(
+    resources.ROOT / 'constants/converters/ccd.pickle'
+)
+
+
+@functools.cache
+def _load_ccd_pickle_cached(
+    path: os.PathLike[str],
+) -> dict[str, Mapping[str, Sequence[str]]]:
+  """Loads the CCD pickle file and caches it so that it is only loaded once."""
+  with open(path, 'rb') as f:
+    return pickle.loads(f.read())
+
+
+class Ccd(Mapping[str, Mapping[str, Sequence[str]]]):
+  """Chemical Components found in PDB (CCD) constants.
+
+  See https://academic.oup.com/bioinformatics/article/31/8/1274/212200 for CCD
+  CIF format documentation.
+
+  Wraps the dict to prevent accidental mutation.
+  """
+
+  __slots__ = ('_dict', '_ccd_pickle_path')
+
+  def __init__(
+      self,
+      ccd_pickle_path: os.PathLike[str] | None = None,
+      user_ccd: str | None = None,
+  ):
+    """Initialises the chemical components dictionary.
+
+    Args:
+      ccd_pickle_path: Path to the CCD pickle file. If None, uses the default
+        CCD pickle file included in the source code.
+      user_ccd: A string containing the user-provided CCD. This has to conform
+        to the same format as the CCD, see https://www.wwpdb.org/data/ccd. If
+        provided, takes precedence over the CCD for the the same key. This can
+        be used to override specific entries in the CCD if desired.
+    """
+    self._ccd_pickle_path = ccd_pickle_path or _CCD_PICKLE_FILE
+    self._dict = _load_ccd_pickle_cached(self._ccd_pickle_path)
+
+    if user_ccd is not None:
+      if not user_ccd:
+        raise ValueError('User CCD cannot be an empty string.')
+      user_ccd_cifs = {
+          key: {k: tuple(v) for k, v in value.items()}
+          for key, value in cif_dict.parse_multi_data_cif(user_ccd).items()
+      }
+      self._dict.update(user_ccd_cifs)
+
+  def __getitem__(self, key: str) -> Mapping[str, Sequence[str]]:
+    return self._dict[key]
+
+  def __contains__(self, key: str) -> bool:
+    return key in self._dict
+
+  def __iter__(self) -> Iterator[str]:
+    return self._dict.__iter__()
+
+  def __len__(self) -> int:
+    return len(self._dict)
+
+  def __hash__(self) -> int:
+    return id(self)  # Ok since this is immutable.
+
+  def get(
+      self, key: str, default: None | Mapping[str, Sequence[str]] = None
+  ) -> Mapping[str, Sequence[str]] | None:
+    return self._dict.get(key, default)
+
+  def items(self) -> ItemsView[str, Mapping[str, Sequence[str]]]:
+    return self._dict.items()
+
+  def values(self) -> ValuesView[Mapping[str, Sequence[str]]]:
+    return self._dict.values()
+
+  def keys(self) -> KeysView[str]:
+    return self._dict.keys()
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class ComponentInfo:
+  name: str
+  type: str
+  pdbx_synonyms: str
+  formula: str
+  formula_weight: str
+  mon_nstd_parent_comp_id: str
+  mon_nstd_flag: str
+  pdbx_smiles: str
+
+
+def mmcif_to_info(mmcif: Mapping[str, Sequence[str]]) -> ComponentInfo:
+  """Converts CCD mmCIFs to component info. Missing fields are left empty."""
+  names = mmcif['_chem_comp.name']
+  types = mmcif['_chem_comp.type']
+  mon_nstd_parent_comp_ids = mmcif['_chem_comp.mon_nstd_parent_comp_id']
+  pdbx_synonyms = mmcif['_chem_comp.pdbx_synonyms']
+  formulas = mmcif['_chem_comp.formula']
+  formula_weights = mmcif['_chem_comp.formula_weight']
+
+  def front_or_empty(values: Sequence[str]) -> str:
+    return values[0] if values else ''
+
+  type_ = front_or_empty(types)
+  mon_nstd_parent_comp_id = front_or_empty(mon_nstd_parent_comp_ids)
+  if type_.lower() == 'non-polymer':
+    # Unset for non-polymers, e.g. water or ions.
+    mon_nstd_flag = '.'
+  elif mon_nstd_parent_comp_id == '?':
+    # A standard component - it doesn't have a standard parent, e.g. MET.
+    mon_nstd_flag = 'y'
+  else:
+    # A non-standard component, e.g. MSE.
+    mon_nstd_flag = 'n'
+
+  # Default SMILES is the canonical SMILES, but we fall back to the SMILES if a
+  # canonical SMILES is not available. Of canonical SMILES, we prefer ones from
+  # the OpenEye OEToolkits program.
+  canonical_pdbx_smiles = ''
+  fallback_pdbx_smiles = ''
+  descriptor_types = mmcif.get('_pdbx_chem_comp_descriptor.type', [])
+  descriptors = mmcif.get('_pdbx_chem_comp_descriptor.descriptor', [])
+  programs = mmcif.get('_pdbx_chem_comp_descriptor.program', [])
+  for descriptor_type, descriptor, program in zip(
+      descriptor_types, descriptors, programs
+  ):
+    if descriptor_type == 'SMILES_CANONICAL':
+      if (not canonical_pdbx_smiles) or program == 'OpenEye OEToolkits':
+        canonical_pdbx_smiles = descriptor
+    if not fallback_pdbx_smiles and descriptor_type == 'SMILES':
+      fallback_pdbx_smiles = descriptor
+  pdbx_smiles = canonical_pdbx_smiles or fallback_pdbx_smiles
+
+  return ComponentInfo(
+      name=front_or_empty(names),
+      type=type_,
+      pdbx_synonyms=front_or_empty(pdbx_synonyms),
+      formula=front_or_empty(formulas),
+      formula_weight=front_or_empty(formula_weights),
+      mon_nstd_parent_comp_id=mon_nstd_parent_comp_id,
+      mon_nstd_flag=mon_nstd_flag,
+      pdbx_smiles=pdbx_smiles,
+  )
+
+
+@functools.lru_cache(maxsize=128)
+def component_name_to_info(ccd: Ccd, res_name: str) -> ComponentInfo | None:
+  component = ccd.get(res_name)
+  if component is None:
+    return None
+  return mmcif_to_info(component)
+
+
+def type_symbol(ccd: Ccd, res_name: str, atom_name: str) -> str:
+  """Returns the element type for the given component name and atom name.
+
+  Args:
+    ccd: The chemical components dictionary.
+    res_name: The component name, e.g. ARG.
+    atom_name: The atom name, e.g. CB, OXT, or NH1.
+
+  Returns:
+    Element type, e.g. C for (ARG, CB), O for (ARG, OXT), N for (ARG, NH1).
+  """
+  res = ccd.get(res_name)
+  if res is None:
+    return '?'
+  try:
+    return res['_chem_comp_atom.type_symbol'][
+        res['_chem_comp_atom.atom_id'].index(atom_name)
+    ]
+  except (ValueError, IndexError, KeyError):
+    return '?'

flax_model/alphafold3/constants/converters/ccd_pickle_gen.py

@@ -0,0 +1,42 @@
+
+
+"""Reads Chemical Components gz file and generates a CCD pickle file."""
+
+from collections.abc import Sequence
+import gzip
+import pickle
+import sys
+
+from flax_model.alphafold3.cpp import cif_dict
+import tqdm
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) != 3:
+    raise ValueError('Must specify input_file components.cif and output_file')
+
+  _, input_file, output_file = argv
+
+  print(f'Parsing {input_file}', flush=True)
+  if input_file.endswith('.gz'):
+    opener = gzip.open
+  else:
+    opener = open
+
+  with opener(input_file, 'rb') as f:
+    whole_file = f.read()
+  result = {
+      key: {k: tuple(v) for k, v in value.items()}
+      for key, value in tqdm.tqdm(
+          cif_dict.parse_multi_data_cif(whole_file).items(), disable=None
+      )
+  }
+  assert len(result) == whole_file.count(b'data_')
+
+  print(f'Writing {output_file}', flush=True)
+  with open(output_file, 'wb') as f:
+    pickle.dump(result, f, protocol=pickle.HIGHEST_PROTOCOL)
+  print('Done', flush=True)
+
+if __name__ == '__main__':
+  main(sys.argv)

flax_model/alphafold3/constants/converters/chemical_component_sets_gen.py

@@ -0,0 +1,72 @@
+
+
+"""Script for updating chemical_component_sets.py."""
+
+from collections.abc import Mapping, Sequence
+import pathlib
+import pickle
+import re
+import sys
+
+from flax_model.alphafold3.common import resources
+import tqdm
+
+
+_CCD_PICKLE_FILE = resources.filename(
+    'constants/converters/ccd.pickle'
+)
+
+
+def find_ions_and_glycans_in_ccd(
+    ccd: Mapping[str, Mapping[str, Sequence[str]]],
+) -> dict[str, frozenset[str]]:
+  """Finds glycans and ions in all version of CCD."""
+  glycans_linking = []
+  glycans_other = []
+  ions = []
+  for name, comp in tqdm.tqdm(ccd.items(), disable=None):
+    if name == 'UNX':
+      continue  # Skip "unknown atom or ion".
+    comp_type = comp['_chem_comp.type'][0].lower()
+    # Glycans have the type 'saccharide'.
+    if re.findall(r'\bsaccharide\b', comp_type):
+      # Separate out linking glycans from others.
+      if 'linking' in comp_type:
+        glycans_linking.append(name)
+      else:
+        glycans_other.append(name)
+
+    # Ions have the word 'ion' in their name.
+    comp_name = comp['_chem_comp.name'][0].lower()
+    if re.findall(r'\bion\b', comp_name):
+      ions.append(name)
+  result = dict(
+      glycans_linking=frozenset(glycans_linking),
+      glycans_other=frozenset(glycans_other),
+      ions=frozenset(ions),
+  )
+
+  return result
+
+
+def main(argv: Sequence[str]) -> None:
+  if len(argv) != 2:
+    raise ValueError(
+        'Directory to write to must be specified as a command-line arguments.'
+    )
+
+  print(f'Loading {_CCD_PICKLE_FILE}', flush=True)
+  with open(_CCD_PICKLE_FILE, 'rb') as f:
+    ccd: Mapping[str, Mapping[str, Sequence[str]]] = pickle.load(f)
+  output_path = pathlib.Path(argv[1])
+  output_path.parent.mkdir(exist_ok=True)
+  print('Finding ions and glycans', flush=True)
+  result = find_ions_and_glycans_in_ccd(ccd)
+  print(f'writing to {output_path}', flush=True)
+  with output_path.open('wb') as f:
+    pickle.dump(result, f)
+  print('Done', flush=True)
+
+
+if __name__ == '__main__':
+  main(sys.argv)

flax_model/alphafold3/constants/mmcif_names.py

@@ -0,0 +1,208 @@
+
+
+"""Names of things in mmCIF format.
+
+See https://www.iucr.org/__data/iucr/cifdic_html/2/cif_mm.dic/index.html
+"""
+
+from collections.abc import Mapping, Sequence, Set
+from typing import Final
+
+from flax_model.alphafold3.constants import atom_types
+from flax_model.alphafold3.constants import residue_names
+
+
+# The following are all possible values for the "_entity.type".
+# https://mmcif.wwpdb.org/dictionaries/mmcif_pdbx_v50.dic/Items/_entity.type.html
+BRANCHED_CHAIN: Final[str] = 'branched'
+MACROLIDE_CHAIN: Final[str] = 'macrolide'
+NON_POLYMER_CHAIN: Final[str] = 'non-polymer'
+POLYMER_CHAIN: Final[str] = 'polymer'
+WATER: Final[str] = 'water'
+
+CYCLIC_PSEUDO_PEPTIDE_CHAIN: Final[str] = 'cyclic-pseudo-peptide'
+DNA_CHAIN: Final[str] = 'polydeoxyribonucleotide'
+DNA_RNA_HYBRID_CHAIN: Final[str] = (
+    'polydeoxyribonucleotide/polyribonucleotide hybrid'
+)
+OTHER_CHAIN: Final[str] = 'other'
+PEPTIDE_NUCLEIC_ACID_CHAIN: Final[str] = 'peptide nucleic acid'
+POLYPEPTIDE_D_CHAIN: Final[str] = 'polypeptide(D)'
+PROTEIN_CHAIN: Final[str] = 'polypeptide(L)'
+RNA_CHAIN: Final[str] = 'polyribonucleotide'
+
+# Most common _entity_poly.types.
+STANDARD_POLYMER_CHAIN_TYPES: Final[Set[str]] = {
+    PROTEIN_CHAIN,
+    DNA_CHAIN,
+    RNA_CHAIN,
+}
+
+# Possible values for _entity.type other than polymer and water.
+LIGAND_CHAIN_TYPES: Final[Set[str]] = {
+    BRANCHED_CHAIN,
+    MACROLIDE_CHAIN,
+    NON_POLYMER_CHAIN,
+}
+
+# Possible values for _entity.type other than polymer.
+NON_POLYMER_CHAIN_TYPES: Final[Set[str]] = {
+    *LIGAND_CHAIN_TYPES,
+    WATER,
+}
+
+# Peptide possible values for _entity_poly.type.
+PEPTIDE_CHAIN_TYPES: Final[Set[str]] = {
+    CYCLIC_PSEUDO_PEPTIDE_CHAIN,
+    POLYPEPTIDE_D_CHAIN,
+    PROTEIN_CHAIN,
+    PEPTIDE_NUCLEIC_ACID_CHAIN,
+}
+
+
+# Nucleic-acid possible values for _entity_poly.type.
+NUCLEIC_ACID_CHAIN_TYPES: Final[Set[str]] = {
+    RNA_CHAIN,
+    DNA_CHAIN,
+    DNA_RNA_HYBRID_CHAIN,
+}
+
+# All possible values for _entity_poly.type.
+POLYMER_CHAIN_TYPES: Final[Set[str]] = {
+    *NUCLEIC_ACID_CHAIN_TYPES,
+    *PEPTIDE_CHAIN_TYPES,
+    OTHER_CHAIN,
+}
+
+
+TERMINAL_OXYGENS: Final[Mapping[str, str]] = {
+    PROTEIN_CHAIN: 'OXT',
+    DNA_CHAIN: 'OP3',
+    RNA_CHAIN: 'OP3',
+}
+
+
+# For each chain type, which atom should be used to represent each residue.
+RESIDUE_REPRESENTATIVE_ATOMS: Final[Mapping[str, str]] = {
+    PROTEIN_CHAIN: atom_types.CA,
+    DNA_CHAIN: atom_types.C1PRIME,
+    RNA_CHAIN: atom_types.C1PRIME,
+}
+
+# Methods involving crystallization. See the documentation at
+# mmcif.wwpdb.org/dictionaries/mmcif_pdbx_v50.dic/Items/_exptl.method.html
+# for the full list of experimental methods.
+CRYSTALLIZATION_METHODS: Final[Set[str]] = {
+    'X-RAY DIFFRACTION',
+    'NEUTRON DIFFRACTION',
+    'ELECTRON CRYSTALLOGRAPHY',
+    'POWDER CRYSTALLOGRAPHY',
+    'FIBER DIFFRACTION',
+}
+
+# Possible bond types.
+COVALENT_BOND: Final[str] = 'covale'
+HYDROGEN_BOND: Final[str] = 'hydrog'
+METAL_COORDINATION: Final[str] = 'metalc'
+DISULFIDE_BRIDGE: Final[str] = 'disulf'
+
+
+def is_standard_polymer_type(chain_type: str) -> bool:
+  """Returns if chain type is a protein, DNA or RNA chain type.
+
+  Args:
+     chain_type: The type of the chain.
+
+  Returns:
+    A bool for if the chain_type matches protein, DNA, or RNA.
+  """
+  return chain_type in STANDARD_POLYMER_CHAIN_TYPES
+
+
+def guess_polymer_type(chain_residues: Sequence[str]) -> str:
+  """Guess the polymer type (protein/rna/dna/other) based on the residues.
+
+  The polymer type is guessed by first checking for any of the standard
+  protein residues. If one is present then the chain is considered to be a
+  polypeptide. Otherwise we decide by counting residue types and deciding by
+  majority voting (e.g. mostly DNA residues -> DNA). If there is a tie between
+  the counts, the ordering is rna > dna > other.
+
+  Note that we count MSE and UNK as protein residues.
+
+  Args:
+    chain_residues: A sequence of full residue name (1-letter for DNA, 2-letters
+      for RNA, 3 for protein). The _atom_site.label_comp_id column in mmCIF.
+
+  Returns:
+    The most probable chain type as set in the _entity_poly mmCIF table:
+    protein - polypeptide(L), rna - polyribonucleotide,
+    dna - polydeoxyribonucleotide or other.
+  """
+  residue_types = {
+      **{r: RNA_CHAIN for r in residue_names.RNA_TYPES},
+      **{r: DNA_CHAIN for r in residue_names.DNA_TYPES},
+      **{r: PROTEIN_CHAIN for r in residue_names.PROTEIN_TYPES_WITH_UNKNOWN},
+      residue_names.MSE: PROTEIN_CHAIN,
+  }
+
+  counts = {PROTEIN_CHAIN: 0, RNA_CHAIN: 0, DNA_CHAIN: 0, OTHER_CHAIN: 0}
+  for residue in chain_residues:
+    residue_type = residue_types.get(residue, OTHER_CHAIN)
+    # If we ever see a protein residue we'll consider this a polypeptide(L).
+    if residue_type == PROTEIN_CHAIN:
+      return residue_type
+    counts[residue_type] += 1
+
+  # Make sure protein > rna > dna > other if there is a tie.
+  tie_braker = {PROTEIN_CHAIN: 3, RNA_CHAIN: 2, DNA_CHAIN: 1, OTHER_CHAIN: 0}
+
+  def order_fn(item):
+    name, count = item
+    return count, tie_braker[name]
+
+  most_probable_type = max(counts.items(), key=order_fn)[0]
+  return most_probable_type
+
+
+def fix_non_standard_polymer_res(*, res_name: str, chain_type: str) -> str:
+  """Returns the res_name of the closest standard protein/RNA/DNA residue.
+
+  Optimized for the case where a single residue needs to be converted.
+
+  If res_name is already a standard type, it is returned unaltered.
+  If a match cannot be found, returns 'UNK' for protein chains and 'N' for
+    RNA/DNA chains.
+
+  Args:
+     res_name: A residue_name (monomer code from the CCD).
+     chain_type: The type of the chain, must be PROTEIN_CHAIN, RNA_CHAIN or
+       DNA_CHAIN.
+
+  Returns:
+    An element from PROTEIN_TYPES_WITH_UNKNOWN | RNA_TYPES | DNA_TYPES | {'N'}.
+
+  Raises:
+    ValueError: If chain_type not in PEPTIDE_CHAIN_TYPES or
+      {OTHER_CHAIN, RNA_CHAIN, DNA_CHAIN, DNA_RNA_HYBRID_CHAIN}.
+  """
+  # Map to one letter code, then back to common res_names.
+  one_letter_code = residue_names.letters_three_to_one(res_name, default='X')
+
+  if chain_type in PEPTIDE_CHAIN_TYPES or chain_type == OTHER_CHAIN:
+    return residue_names.PROTEIN_COMMON_ONE_TO_THREE.get(one_letter_code, 'UNK')
+  elif chain_type == RNA_CHAIN:
+    # RNA's CCD monomer code is single-letter.
+    return (
+        one_letter_code if one_letter_code in residue_names.RNA_TYPES else 'N'
+    )
+  elif chain_type == DNA_CHAIN:
+    return residue_names.DNA_COMMON_ONE_TO_TWO.get(one_letter_code, 'N')
+  elif chain_type == DNA_RNA_HYBRID_CHAIN:
+    return (
+        res_name
+        if res_name in residue_names.NUCLEIC_TYPES_WITH_UNKNOWN
+        else 'N'
+    )
+  else:
+    raise ValueError(f'Expected a protein/DNA/RNA chain but got {chain_type}')

flax_model/alphafold3/constants/periodic_table.py

@@ -0,0 +1,385 @@
+
+
+"""Periodic table of elements."""
+
+from collections.abc import Mapping, Sequence
+import dataclasses
+from typing import Final
+
+import numpy as np
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Element:
+  name: str
+  number: int
+  symbol: str
+  weight: float
+
+
+# Weights taken from rdkit/Code/GraphMol/atomic_data.cpp for compatibility.
+# pylint: disable=invalid-name
+
+# X is an unknown element that can be present in the CCD,
+# https://www.rcsb.org/ligand/UNX.
+X: Final[Element] = Element(name='Unknown', number=0, symbol='X', weight=0.0)
+H: Final[Element] = Element(name='Hydrogen', number=1, symbol='H', weight=1.008)
+He: Final[Element] = Element(name='Helium', number=2, symbol='He', weight=4.003)
+Li: Final[Element] = Element(
+    name='Lithium', number=3, symbol='Li', weight=6.941
+)
+Be: Final[Element] = Element(
+    name='Beryllium', number=4, symbol='Be', weight=9.012
+)
+B: Final[Element] = Element(name='Boron', number=5, symbol='B', weight=10.812)
+C: Final[Element] = Element(name='Carbon', number=6, symbol='C', weight=12.011)
+N: Final[Element] = Element(
+    name='Nitrogen', number=7, symbol='N', weight=14.007
+)
+O: Final[Element] = Element(name='Oxygen', number=8, symbol='O', weight=15.999)
+F: Final[Element] = Element(
+    name='Fluorine', number=9, symbol='F', weight=18.998
+)
+Ne: Final[Element] = Element(name='Neon', number=10, symbol='Ne', weight=20.18)
+Na: Final[Element] = Element(
+    name='Sodium', number=11, symbol='Na', weight=22.99
+)
+Mg: Final[Element] = Element(
+    name='Magnesium', number=12, symbol='Mg', weight=24.305
+)
+Al: Final[Element] = Element(
+    name='Aluminium', number=13, symbol='Al', weight=26.982
+)
+Si: Final[Element] = Element(
+    name='Silicon', number=14, symbol='Si', weight=28.086
+)
+P: Final[Element] = Element(
+    name='Phosphorus', number=15, symbol='P', weight=30.974
+)
+S: Final[Element] = Element(name='Sulfur', number=16, symbol='S', weight=32.067)
+Cl: Final[Element] = Element(
+    name='Chlorine', number=17, symbol='Cl', weight=35.453
+)
+Ar: Final[Element] = Element(
+    name='Argon', number=18, symbol='Ar', weight=39.948
+)
+K: Final[Element] = Element(
+    name='Potassium', number=19, symbol='K', weight=39.098
+)
+Ca: Final[Element] = Element(
+    name='Calcium', number=20, symbol='Ca', weight=40.078
+)
+Sc: Final[Element] = Element(
+    name='Scandium', number=21, symbol='Sc', weight=44.956
+)
+Ti: Final[Element] = Element(
+    name='Titanium', number=22, symbol='Ti', weight=47.867
+)
+V: Final[Element] = Element(
+    name='Vanadium', number=23, symbol='V', weight=50.942
+)
+Cr: Final[Element] = Element(
+    name='Chromium', number=24, symbol='Cr', weight=51.996
+)
+Mn: Final[Element] = Element(
+    name='Manganese', number=25, symbol='Mn', weight=54.938
+)
+Fe: Final[Element] = Element(name='Iron', number=26, symbol='Fe', weight=55.845)
+Co: Final[Element] = Element(
+    name='Cobalt', number=27, symbol='Co', weight=58.933
+)
+Ni: Final[Element] = Element(
+    name='Nickel', number=28, symbol='Ni', weight=58.693
+)
+Cu: Final[Element] = Element(
+    name='Copper', number=29, symbol='Cu', weight=63.546
+)
+Zn: Final[Element] = Element(name='Zinc', number=30, symbol='Zn', weight=65.39)
+Ga: Final[Element] = Element(
+    name='Gallium', number=31, symbol='Ga', weight=69.723
+)
+Ge: Final[Element] = Element(
+    name='Germanium', number=32, symbol='Ge', weight=72.61
+)
+As: Final[Element] = Element(
+    name='Arsenic', number=33, symbol='As', weight=74.922
+)
+Se: Final[Element] = Element(
+    name='Selenium', number=34, symbol='Se', weight=78.96
+)
+Br: Final[Element] = Element(
+    name='Bromine', number=35, symbol='Br', weight=79.904
+)
+Kr: Final[Element] = Element(
+    name='Krypton', number=36, symbol='Kr', weight=83.8
+)
+Rb: Final[Element] = Element(
+    name='Rubidium', number=37, symbol='Rb', weight=85.468
+)
+Sr: Final[Element] = Element(
+    name='Strontium', number=38, symbol='Sr', weight=87.62
+)
+Y: Final[Element] = Element(
+    name='Yttrium', number=39, symbol='Y', weight=88.906
+)
+Zr: Final[Element] = Element(
+    name='Zirconium', number=40, symbol='Zr', weight=91.224
+)
+Nb: Final[Element] = Element(
+    name='Niobiu', number=41, symbol='Nb', weight=92.906
+)
+Mo: Final[Element] = Element(
+    name='Molybdenum', number=42, symbol='Mo', weight=95.94
+)
+Tc: Final[Element] = Element(
+    name='Technetium', number=43, symbol='Tc', weight=98
+)
+Ru: Final[Element] = Element(
+    name='Ruthenium', number=44, symbol='Ru', weight=101.07
+)
+Rh: Final[Element] = Element(
+    name='Rhodium', number=45, symbol='Rh', weight=102.906
+)
+Pd: Final[Element] = Element(
+    name='Palladium', number=46, symbol='Pd', weight=106.42
+)
+Ag: Final[Element] = Element(
+    name='Silver', number=47, symbol='Ag', weight=107.868
+)
+Cd: Final[Element] = Element(
+    name='Cadmium', number=48, symbol='Cd', weight=112.412
+)
+In: Final[Element] = Element(
+    name='Indium', number=49, symbol='In', weight=114.818
+)
+Sn: Final[Element] = Element(name='Tin', number=50, symbol='Sn', weight=118.711)
+Sb: Final[Element] = Element(
+    name='Antimony', number=51, symbol='Sb', weight=121.76
+)
+Te: Final[Element] = Element(
+    name='Tellurium', number=52, symbol='Te', weight=127.6
+)
+I: Final[Element] = Element(
+    name='Iodine', number=53, symbol='I', weight=126.904
+)
+Xe: Final[Element] = Element(
+    name='Xenon', number=54, symbol='Xe', weight=131.29
+)
+Cs: Final[Element] = Element(
+    name='Caesium', number=55, symbol='Cs', weight=132.905
+)
+Ba: Final[Element] = Element(
+    name='Barium', number=56, symbol='Ba', weight=137.328
+)
+La: Final[Element] = Element(
+    name='Lanthanum', number=57, symbol='La', weight=138.906
+)
+Ce: Final[Element] = Element(
+    name='Cerium', number=58, symbol='Ce', weight=140.116
+)
+Pr: Final[Element] = Element(
+    name='Praseodymium', number=59, symbol='Pr', weight=140.908
+)
+Nd: Final[Element] = Element(
+    name='Neodymium', number=60, symbol='Nd', weight=144.24
+)
+Pm: Final[Element] = Element(
+    name='Promethium', number=61, symbol='Pm', weight=145
+)
+Sm: Final[Element] = Element(
+    name='Samarium', number=62, symbol='Sm', weight=150.36
+)
+Eu: Final[Element] = Element(
+    name='Europium', number=63, symbol='Eu', weight=151.964
+)
+Gd: Final[Element] = Element(
+    name='Gadolinium', number=64, symbol='Gd', weight=157.25
+)
+Tb: Final[Element] = Element(
+    name='Terbium', number=65, symbol='Tb', weight=158.925
+)
+Dy: Final[Element] = Element(
+    name='Dysprosium', number=66, symbol='Dy', weight=162.5
+)
+Ho: Final[Element] = Element(
+    name='Holmium', number=67, symbol='Ho', weight=164.93
+)
+Er: Final[Element] = Element(
+    name='Erbium', number=68, symbol='Er', weight=167.26
+)
+Tm: Final[Element] = Element(
+    name='Thulium', number=69, symbol='Tm', weight=168.934
+)
+Yb: Final[Element] = Element(
+    name='Ytterbium', number=70, symbol='Yb', weight=173.04
+)
+Lu: Final[Element] = Element(
+    name='Lutetium', number=71, symbol='Lu', weight=174.967
+)
+Hf: Final[Element] = Element(
+    name='Hafnium', number=72, symbol='Hf', weight=178.49
+)
+Ta: Final[Element] = Element(
+    name='Tantalum', number=73, symbol='Ta', weight=180.948
+)
+W: Final[Element] = Element(
+    name='Tungsten', number=74, symbol='W', weight=183.84
+)
+Re: Final[Element] = Element(
+    name='Rhenium', number=75, symbol='Re', weight=186.207
+)
+Os: Final[Element] = Element(
+    name='Osmium', number=76, symbol='Os', weight=190.23
+)
+Ir: Final[Element] = Element(
+    name='Iridium', number=77, symbol='Ir', weight=192.217
+)
+Pt: Final[Element] = Element(
+    name='Platinum', number=78, symbol='Pt', weight=195.078
+)
+Au: Final[Element] = Element(
+    name='Gold', number=79, symbol='Au', weight=196.967
+)
+Hg: Final[Element] = Element(
+    name='Mercury', number=80, symbol='Hg', weight=200.59
+)
+Tl: Final[Element] = Element(
+    name='Thallium', number=81, symbol='Tl', weight=204.383
+)
+Pb: Final[Element] = Element(name='Lead', number=82, symbol='Pb', weight=207.2)
+Bi: Final[Element] = Element(
+    name='Bismuth', number=83, symbol='Bi', weight=208.98
+)
+Po: Final[Element] = Element(
+    name='Polonium', number=84, symbol='Po', weight=209
+)
+At: Final[Element] = Element(
+    name='Astatine', number=85, symbol='At', weight=210
+)
+Rn: Final[Element] = Element(name='Radon', number=86, symbol='Rn', weight=222)
+Fr: Final[Element] = Element(
+    name='Francium', number=87, symbol='Fr', weight=223
+)
+Ra: Final[Element] = Element(name='Radium', number=88, symbol='Ra', weight=226)
+Ac: Final[Element] = Element(
+    name='Actinium', number=89, symbol='Ac', weight=227
+)
+Th: Final[Element] = Element(
+    name='Thorium', number=90, symbol='Th', weight=232.038
+)
+Pa: Final[Element] = Element(
+    name='Protactinium', number=91, symbol='Pa', weight=231.036
+)
+U: Final[Element] = Element(
+    name='Uranium', number=92, symbol='U', weight=238.029
+)
+Np: Final[Element] = Element(
+    name='Neptunium', number=93, symbol='Np', weight=237
+)
+Pu: Final[Element] = Element(
+    name='Plutonium', number=94, symbol='Pu', weight=244
+)
+Am: Final[Element] = Element(
+    name='Americium', number=95, symbol='Am', weight=243
+)
+Cm: Final[Element] = Element(name='Curium', number=96, symbol='Cm', weight=247)
+Bk: Final[Element] = Element(
+    name='Berkelium', number=97, symbol='Bk', weight=247
+)
+Cf: Final[Element] = Element(
+    name='Californium', number=98, symbol='Cf', weight=251
+)
+Es: Final[Element] = Element(
+    name='Einsteinium', number=99, symbol='Es', weight=252
+)
+Fm: Final[Element] = Element(
+    name='Fermium', number=100, symbol='Fm', weight=257
+)
+Md: Final[Element] = Element(
+    name='Mendelevium', number=101, symbol='Md', weight=258
+)
+No: Final[Element] = Element(
+    name='Nobelium', number=102, symbol='No', weight=259
+)
+Lr: Final[Element] = Element(
+    name='Lawrencium', number=103, symbol='Lr', weight=262
+)
+Rf: Final[Element] = Element(
+    name='Rutherfordium', number=104, symbol='Rf', weight=267
+)
+Db: Final[Element] = Element(
+    name='Dubnium', number=105, symbol='Db', weight=268
+)
+Sg: Final[Element] = Element(
+    name='Seaborgium', number=106, symbol='Sg', weight=269
+)
+Bh: Final[Element] = Element(
+    name='Bohrium', number=107, symbol='Bh', weight=270
+)
+Hs: Final[Element] = Element(
+    name='Hassium', number=108, symbol='Hs', weight=269
+)
+Mt: Final[Element] = Element(
+    name='Meitnerium', number=109, symbol='Mt', weight=278
+)
+Ds: Final[Element] = Element(
+    name='Darmstadtium', number=110, symbol='Ds', weight=281
+)
+Rg: Final[Element] = Element(
+    name='Roentgenium', number=111, symbol='Rg', weight=281
+)
+Cn: Final[Element] = Element(
+    name='Copernicium', number=112, symbol='Cn', weight=285
+)
+Nh: Final[Element] = Element(
+    name='Nihonium', number=113, symbol='Nh', weight=284
+)
+Fl: Final[Element] = Element(
+    name='Flerovium', number=114, symbol='Fl', weight=289
+)
+Mc: Final[Element] = Element(
+    name='Moscovium', number=115, symbol='Mc', weight=288
+)
+Lv: Final[Element] = Element(
+    name='Livermorium', number=116, symbol='Lv', weight=293
+)
+Ts: Final[Element] = Element(
+    name='Tennessine', number=117, symbol='Ts', weight=292
+)
+Og: Final[Element] = Element(
+    name='Oganesson', number=118, symbol='Og', weight=294
+)
+# pylint: enable=invalid-name
+
+# fmt: off
+# Lanthanides
+_L: Final[Sequence[Element]] = (
+    La, Ce, Pr, Nd, Pm, Sm, Eu, Gd, Tb, Dy, Ho, Er, Tm, Yb, Lu)
+# Actinides
+_A: Final[Sequence[Element]] = (
+    Ac, Th, Pa, U, Np, Pu, Am, Cm, Bk, Cf, Es, Fm, Md, No, Lr)
+
+# pylint: disable=bad-whitespace
+PERIODIC_TABLE: Final[Sequence[Element]] = (
+    X,  # Unknown
+    H,                                                                   He,
+    Li, Be,                                          B,  C,  N,  O,  F,  Ne,
+    Na, Mg,                                          Al, Si, P,  S,  Cl, Ar,
+    K,  Ca,  Sc, Ti, V,  Cr, Mn, Fe, Co, Ni, Cu, Zn, Ga, Ge, As, Se, Br, Kr,
+    Rb, Sr,  Y,  Zr, Nb, Mo, Tc, Ru, Rh, Pd, Ag, Cd, In, Sn, Sb, Te, I,  Xe,
+    Cs, Ba, *_L, Hf, Ta, W,  Re, Os, Ir, Pt, Au, Hg, Tl, Pb, Bi, Po, At, Rn,
+    Fr, Ra, *_A, Rf, Db, Sg, Bh, Hs, Mt, Ds, Rg, Cn, Nh, Fl, Mc, Lv, Ts, Og
+)
+# pylint: enable=bad-whitespace
+# fmt: on
+ATOMIC_SYMBOL: Mapping[int, str] = {e.number: e.symbol for e in PERIODIC_TABLE}
+ATOMIC_NUMBER = {e.symbol: e.number for e in PERIODIC_TABLE}
+# Add Deuterium as previous table contained it.
+ATOMIC_NUMBER['D'] = 1
+
+ATOMIC_NUMBER: Mapping[str, int] = ATOMIC_NUMBER
+ATOMIC_WEIGHT: np.ndarray = np.zeros(len(PERIODIC_TABLE), dtype=np.float64)
+
+for e in PERIODIC_TABLE:
+  ATOMIC_WEIGHT[e.number] = e.weight
+ATOMIC_WEIGHT.setflags(write=False)

flax_model/alphafold3/constants/residue_names.py

@@ -0,0 +1,411 @@
+
+
+"""Constants associated with residue names."""
+
+from collections.abc import Mapping
+import functools
+import sys
+
+# pyformat: disable
+# common_typos_disable
+CCD_NAME_TO_ONE_LETTER: Mapping[str, str] = {
+    '00C': 'C', '01W': 'X', '02K': 'A', '03Y': 'C', '07O': 'C', '08P': 'C',
+    '0A0': 'D', '0A1': 'Y', '0A2': 'K', '0A8': 'C', '0AA': 'V', '0AB': 'V',
+    '0AC': 'G', '0AD': 'G', '0AF': 'W', '0AG': 'L', '0AH': 'S', '0AK': 'D',
+    '0AM': 'A', '0AP': 'C', '0AU': 'U', '0AV': 'A', '0AZ': 'P', '0BN': 'F',
+    '0C': 'C', '0CS': 'A', '0DC': 'C', '0DG': 'G', '0DT': 'T', '0FL': 'A',
+    '0G': 'G', '0NC': 'A', '0SP': 'A', '0U': 'U', '10C': 'C', '125': 'U',
+    '126': 'U', '127': 'U', '128': 'N', '12A': 'A', '143': 'C', '193': 'X',
+    '1AP': 'A', '1MA': 'A', '1MG': 'G', '1PA': 'F', '1PI': 'A', '1PR': 'N',
+    '1SC': 'C', '1TQ': 'W', '1TY': 'Y', '1X6': 'S', '200': 'F', '23F': 'F',
+    '23S': 'X', '26B': 'T', '2AD': 'X', '2AG': 'A', '2AO': 'X', '2AR': 'A',
+    '2AS': 'X', '2AT': 'T', '2AU': 'U', '2BD': 'I', '2BT': 'T', '2BU': 'A',
+    '2CO': 'C', '2DA': 'A', '2DF': 'N', '2DM': 'N', '2DO': 'X', '2DT': 'T',
+    '2EG': 'G', '2FE': 'N', '2FI': 'N', '2FM': 'M', '2GT': 'T', '2HF': 'H',
+    '2LU': 'L', '2MA': 'A', '2MG': 'G', '2ML': 'L', '2MR': 'R', '2MT': 'P',
+    '2MU': 'U', '2NT': 'T', '2OM': 'U', '2OT': 'T', '2PI': 'X', '2PR': 'G',
+    '2SA': 'N', '2SI': 'X', '2ST': 'T', '2TL': 'T', '2TY': 'Y', '2VA': 'V',
+    '2XA': 'C', '32S': 'X', '32T': 'X', '3AH': 'H', '3AR': 'X', '3CF': 'F',
+    '3DA': 'A', '3DR': 'N', '3GA': 'A', '3MD': 'D', '3ME': 'U', '3NF': 'Y',
+    '3QN': 'K', '3TY': 'X', '3XH': 'G', '4AC': 'N', '4BF': 'Y', '4CF': 'F',
+    '4CY': 'M', '4DP': 'W', '4FB': 'P', '4FW': 'W', '4HT': 'W', '4IN': 'W',
+    '4MF': 'N', '4MM': 'X', '4OC': 'C', '4PC': 'C', '4PD': 'C', '4PE': 'C',
+    '4PH': 'F', '4SC': 'C', '4SU': 'U', '4TA': 'N', '4U7': 'A', '56A': 'H',
+    '5AA': 'A', '5AB': 'A', '5AT': 'T', '5BU': 'U', '5CG': 'G', '5CM': 'C',
+    '5CS': 'C', '5FA': 'A', '5FC': 'C', '5FU': 'U', '5HP': 'E', '5HT': 'T',
+    '5HU': 'U', '5IC': 'C', '5IT': 'T', '5IU': 'U', '5MC': 'C', '5MD': 'N',
+    '5MU': 'U', '5NC': 'C', '5PC': 'C', '5PY': 'T', '5SE': 'U', '64T': 'T',
+    '6CL': 'K', '6CT': 'T', '6CW': 'W', '6HA': 'A', '6HC': 'C', '6HG': 'G',
+    '6HN': 'K', '6HT': 'T', '6IA': 'A', '6MA': 'A', '6MC': 'A', '6MI': 'N',
+    '6MT': 'A', '6MZ': 'N', '6OG': 'G', '70U': 'U', '7DA': 'A', '7GU': 'G',
+    '7JA': 'I', '7MG': 'G', '8AN': 'A', '8FG': 'G', '8MG': 'G', '8OG': 'G',
+    '9NE': 'E', '9NF': 'F', '9NR': 'R', '9NV': 'V', 'A': 'A', 'A1P': 'N',
+    'A23': 'A', 'A2L': 'A', 'A2M': 'A', 'A34': 'A', 'A35': 'A', 'A38': 'A',
+    'A39': 'A', 'A3A': 'A', 'A3P': 'A', 'A40': 'A', 'A43': 'A', 'A44': 'A',
+    'A47': 'A', 'A5L': 'A', 'A5M': 'C', 'A5N': 'N', 'A5O': 'A', 'A66': 'X',
+    'AA3': 'A', 'AA4': 'A', 'AAR': 'R', 'AB7': 'X', 'ABA': 'A', 'ABR': 'A',
+    'ABS': 'A', 'ABT': 'N', 'ACB': 'D', 'ACL': 'R', 'AD2': 'A', 'ADD': 'X',
+    'ADX': 'N', 'AEA': 'X', 'AEI': 'D', 'AET': 'A', 'AFA': 'N', 'AFF': 'N',
+    'AFG': 'G', 'AGM': 'R', 'AGT': 'C', 'AHB': 'N', 'AHH': 'X', 'AHO': 'A',
+    'AHP': 'A', 'AHS': 'X', 'AHT': 'X', 'AIB': 'A', 'AKL': 'D', 'AKZ': 'D',
+    'ALA': 'A', 'ALC': 'A', 'ALM': 'A', 'ALN': 'A', 'ALO': 'T', 'ALQ': 'X',
+    'ALS': 'A', 'ALT': 'A', 'ALV': 'A', 'ALY': 'K', 'AN8': 'A', 'AP7': 'A',
+    'APE': 'X', 'APH': 'A', 'API': 'K', 'APK': 'K', 'APM': 'X', 'APP': 'X',
+    'AR2': 'R', 'AR4': 'E', 'AR7': 'R', 'ARG': 'R', 'ARM': 'R', 'ARO': 'R',
+    'ARV': 'X', 'AS': 'A', 'AS2': 'D', 'AS9': 'X', 'ASA': 'D', 'ASB': 'D',
+    'ASI': 'D', 'ASK': 'D', 'ASL': 'D', 'ASM': 'X', 'ASN': 'N', 'ASP': 'D',
+    'ASQ': 'D', 'ASU': 'N', 'ASX': 'B', 'ATD': 'T', 'ATL': 'T', 'ATM': 'T',
+    'AVC': 'A', 'AVN': 'X', 'AYA': 'A', 'AZK': 'K', 'AZS': 'S', 'AZY': 'Y',
+    'B1F': 'F', 'B1P': 'N', 'B2A': 'A', 'B2F': 'F', 'B2I': 'I', 'B2V': 'V',
+    'B3A': 'A', 'B3D': 'D', 'B3E': 'E', 'B3K': 'K', 'B3L': 'X', 'B3M': 'X',
+    'B3Q': 'X', 'B3S': 'S', 'B3T': 'X', 'B3U': 'H', 'B3X': 'N', 'B3Y': 'Y',
+    'BB6': 'C', 'BB7': 'C', 'BB8': 'F', 'BB9': 'C', 'BBC': 'C', 'BCS': 'C',
+    'BE2': 'X', 'BFD': 'D', 'BG1': 'S', 'BGM': 'G', 'BH2': 'D', 'BHD': 'D',
+    'BIF': 'F', 'BIL': 'X', 'BIU': 'I', 'BJH': 'X', 'BLE': 'L', 'BLY': 'K',
+    'BMP': 'N', 'BMT': 'T', 'BNN': 'F', 'BNO': 'X', 'BOE': 'T', 'BOR': 'R',
+    'BPE': 'C', 'BRU': 'U', 'BSE': 'S', 'BT5': 'N', 'BTA': 'L', 'BTC': 'C',
+    'BTR': 'W', 'BUC': 'C', 'BUG': 'V', 'BVP': 'U', 'BZG': 'N', 'C': 'C',
+    'C1X': 'K', 'C25': 'C', 'C2L': 'C', 'C2S': 'C', 'C31': 'C', 'C32': 'C',
+    'C34': 'C', 'C36': 'C', 'C37': 'C', 'C38': 'C', 'C3Y': 'C', 'C42': 'C',
+    'C43': 'C', 'C45': 'C', 'C46': 'C', 'C49': 'C', 'C4R': 'C', 'C4S': 'C',
+    'C5C': 'C', 'C66': 'X', 'C6C': 'C', 'CAF': 'C', 'CAL': 'X', 'CAR': 'C',
+    'CAS': 'C', 'CAV': 'X', 'CAY': 'C', 'CB2': 'C', 'CBR': 'C', 'CBV': 'C',
+    'CCC': 'C', 'CCL': 'K', 'CCS': 'C', 'CDE': 'X', 'CDV': 'X', 'CDW': 'C',
+    'CEA': 'C', 'CFL': 'C', 'CG1': 'G', 'CGA': 'E', 'CGU': 'E', 'CH': 'C',
+    'CHF': 'X', 'CHG': 'X', 'CHP': 'G', 'CHS': 'X', 'CIR': 'R', 'CLE': 'L',
+    'CLG': 'K', 'CLH': 'K', 'CM0': 'N', 'CME': 'C', 'CMH': 'C', 'CML': 'C',
+    'CMR': 'C', 'CMT': 'C', 'CNU': 'U', 'CP1': 'C', 'CPC': 'X', 'CPI': 'X',
+    'CR5': 'G', 'CS0': 'C', 'CS1': 'C', 'CS3': 'C', 'CS4': 'C', 'CS8': 'N',
+    'CSA': 'C', 'CSB': 'C', 'CSD': 'C', 'CSE': 'C', 'CSF': 'C', 'CSI': 'G',
+    'CSJ': 'C', 'CSL': 'C', 'CSO': 'C', 'CSP': 'C', 'CSR': 'C', 'CSS': 'C',
+    'CSU': 'C', 'CSW': 'C', 'CSX': 'C', 'CSZ': 'C', 'CTE': 'W', 'CTG': 'T',
+    'CTH': 'T', 'CUC': 'X', 'CWR': 'S', 'CXM': 'M', 'CY0': 'C', 'CY1': 'C',
+    'CY3': 'C', 'CY4': 'C', 'CYA': 'C', 'CYD': 'C', 'CYF': 'C', 'CYG': 'C',
+    'CYJ': 'X', 'CYM': 'C', 'CYQ': 'C', 'CYR': 'C', 'CYS': 'C', 'CZ2': 'C',
+    'CZZ': 'C', 'D11': 'T', 'D1P': 'N', 'D3': 'N', 'D33': 'N', 'D3P': 'G',
+    'D3T': 'T', 'D4M': 'T', 'D4P': 'X', 'DA': 'A', 'DA2': 'X', 'DAB': 'A',
+    'DAH': 'F', 'DAL': 'A', 'DAR': 'R', 'DAS': 'D', 'DBB': 'T', 'DBM': 'N',
+    'DBS': 'S', 'DBU': 'T', 'DBY': 'Y', 'DBZ': 'A', 'DC': 'C', 'DC2': 'C',
+    'DCG': 'G', 'DCI': 'X', 'DCL': 'X', 'DCT': 'C', 'DCY': 'C', 'DDE': 'H',
+    'DDG': 'G', 'DDN': 'U', 'DDX': 'N', 'DFC': 'C', 'DFG': 'G', 'DFI': 'X',
+    'DFO': 'X', 'DFT': 'N', 'DG': 'G', 'DGH': 'G', 'DGI': 'G', 'DGL': 'E',
+    'DGN': 'Q', 'DHA': 'S', 'DHI': 'H', 'DHL': 'X', 'DHN': 'V', 'DHP': 'X',
+    'DHU': 'U', 'DHV': 'V', 'DI': 'I', 'DIL': 'I', 'DIR': 'R', 'DIV': 'V',
+    'DLE': 'L', 'DLS': 'K', 'DLY': 'K', 'DM0': 'K', 'DMH': 'N', 'DMK': 'D',
+    'DMT': 'X', 'DN': 'N', 'DNE': 'L', 'DNG': 'L', 'DNL': 'K', 'DNM': 'L',
+    'DNP': 'A', 'DNR': 'C', 'DNS': 'K', 'DOA': 'X', 'DOC': 'C', 'DOH': 'D',
+    'DON': 'L', 'DPB': 'T', 'DPH': 'F', 'DPL': 'P', 'DPP': 'A', 'DPQ': 'Y',
+    'DPR': 'P', 'DPY': 'N', 'DRM': 'U', 'DRP': 'N', 'DRT': 'T', 'DRZ': 'N',
+    'DSE': 'S', 'DSG': 'N', 'DSN': 'S', 'DSP': 'D', 'DT': 'T', 'DTH': 'T',
+    'DTR': 'W', 'DTY': 'Y', 'DU': 'U', 'DVA': 'V', 'DXD': 'N', 'DXN': 'N',
+    'DYS': 'C', 'DZM': 'A', 'E': 'A', 'E1X': 'A', 'ECC': 'Q', 'EDA': 'A',
+    'EFC': 'C', 'EHP': 'F', 'EIT': 'T', 'ENP': 'N', 'ESB': 'Y', 'ESC': 'M',
+    'EXB': 'X', 'EXY': 'L', 'EY5': 'N', 'EYS': 'X', 'F2F': 'F', 'FA2': 'A',
+    'FA5': 'N', 'FAG': 'N', 'FAI': 'N', 'FB5': 'A', 'FB6': 'A', 'FCL': 'F',
+    'FFD': 'N', 'FGA': 'E', 'FGL': 'G', 'FGP': 'S', 'FHL': 'X', 'FHO': 'K',
+    'FHU': 'U', 'FLA': 'A', 'FLE': 'L', 'FLT': 'Y', 'FME': 'M', 'FMG': 'G',
+    'FMU': 'N', 'FOE': 'C', 'FOX': 'G', 'FP9': 'P', 'FPA': 'F', 'FRD': 'X',
+    'FT6': 'W', 'FTR': 'W', 'FTY': 'Y', 'FVA': 'V', 'FZN': 'K', 'G': 'G',
+    'G25': 'G', 'G2L': 'G', 'G2S': 'G', 'G31': 'G', 'G32': 'G', 'G33': 'G',
+    'G36': 'G', 'G38': 'G', 'G42': 'G', 'G46': 'G', 'G47': 'G', 'G48': 'G',
+    'G49': 'G', 'G4P': 'N', 'G7M': 'G', 'GAO': 'G', 'GAU': 'E', 'GCK': 'C',
+    'GCM': 'X', 'GDP': 'G', 'GDR': 'G', 'GFL': 'G', 'GGL': 'E', 'GH3': 'G',
+    'GHG': 'Q', 'GHP': 'G', 'GL3': 'G', 'GLH': 'Q', 'GLJ': 'E', 'GLK': 'E',
+    'GLM': 'X', 'GLN': 'Q', 'GLQ': 'E', 'GLU': 'E', 'GLX': 'Z', 'GLY': 'G',
+    'GLZ': 'G', 'GMA': 'E', 'GMS': 'G', 'GMU': 'U', 'GN7': 'G', 'GND': 'X',
+    'GNE': 'N', 'GOM': 'G', 'GPL': 'K', 'GS': 'G', 'GSC': 'G', 'GSR': 'G',
+    'GSS': 'G', 'GSU': 'E', 'GT9': 'C', 'GTP': 'G', 'GVL': 'X', 'H2U': 'U',
+    'H5M': 'P', 'HAC': 'A', 'HAR': 'R', 'HBN': 'H', 'HCS': 'X', 'HDP': 'U',
+    'HEU': 'U', 'HFA': 'X', 'HGL': 'X', 'HHI': 'H', 'HIA': 'H', 'HIC': 'H',
+    'HIP': 'H', 'HIQ': 'H', 'HIS': 'H', 'HL2': 'L', 'HLU': 'L', 'HMR': 'R',
+    'HOL': 'N', 'HPC': 'F', 'HPE': 'F', 'HPH': 'F', 'HPQ': 'F', 'HQA': 'A',
+    'HRG': 'R', 'HRP': 'W', 'HS8': 'H', 'HS9': 'H', 'HSE': 'S', 'HSL': 'S',
+    'HSO': 'H', 'HTI': 'C', 'HTN': 'N', 'HTR': 'W', 'HV5': 'A', 'HVA': 'V',
+    'HY3': 'P', 'HYP': 'P', 'HZP': 'P', 'I': 'I', 'I2M': 'I', 'I58': 'K',
+    'I5C': 'C', 'IAM': 'A', 'IAR': 'R', 'IAS': 'D', 'IC': 'C', 'IEL': 'K',
+    'IG': 'G', 'IGL': 'G', 'IGU': 'G', 'IIL': 'I', 'ILE': 'I', 'ILG': 'E',
+    'ILX': 'I', 'IMC': 'C', 'IML': 'I', 'IOY': 'F', 'IPG': 'G', 'IPN': 'N',
+    'IRN': 'N', 'IT1': 'K', 'IU': 'U', 'IYR': 'Y', 'IYT': 'T', 'IZO': 'M',
+    'JJJ': 'C', 'JJK': 'C', 'JJL': 'C', 'JW5': 'N', 'K1R': 'C', 'KAG': 'G',
+    'KCX': 'K', 'KGC': 'K', 'KNB': 'A', 'KOR': 'M', 'KPI': 'K', 'KST': 'K',
+    'KYQ': 'K', 'L2A': 'X', 'LA2': 'K', 'LAA': 'D', 'LAL': 'A', 'LBY': 'K',
+    'LC': 'C', 'LCA': 'A', 'LCC': 'N', 'LCG': 'G', 'LCH': 'N', 'LCK': 'K',
+    'LCX': 'K', 'LDH': 'K', 'LED': 'L', 'LEF': 'L', 'LEH': 'L', 'LEI': 'V',
+    'LEM': 'L', 'LEN': 'L', 'LET': 'X', 'LEU': 'L', 'LEX': 'L', 'LG': 'G',
+    'LGP': 'G', 'LHC': 'X', 'LHU': 'U', 'LKC': 'N', 'LLP': 'K', 'LLY': 'K',
+    'LME': 'E', 'LMF': 'K', 'LMQ': 'Q', 'LMS': 'N', 'LP6': 'K', 'LPD': 'P',
+    'LPG': 'G', 'LPL': 'X', 'LPS': 'S', 'LSO': 'X', 'LTA': 'X', 'LTR': 'W',
+    'LVG': 'G', 'LVN': 'V', 'LYF': 'K', 'LYK': 'K', 'LYM': 'K', 'LYN': 'K',
+    'LYR': 'K', 'LYS': 'K', 'LYX': 'K', 'LYZ': 'K', 'M0H': 'C', 'M1G': 'G',
+    'M2G': 'G', 'M2L': 'K', 'M2S': 'M', 'M30': 'G', 'M3L': 'K', 'M5M': 'C',
+    'MA': 'A', 'MA6': 'A', 'MA7': 'A', 'MAA': 'A', 'MAD': 'A', 'MAI': 'R',
+    'MBQ': 'Y', 'MBZ': 'N', 'MC1': 'S', 'MCG': 'X', 'MCL': 'K', 'MCS': 'C',
+    'MCY': 'C', 'MD3': 'C', 'MD6': 'G', 'MDH': 'X', 'MDR': 'N', 'MEA': 'F',
+    'MED': 'M', 'MEG': 'E', 'MEN': 'N', 'MEP': 'U', 'MEQ': 'Q', 'MET': 'M',
+    'MEU': 'G', 'MF3': 'X', 'MG1': 'G', 'MGG': 'R', 'MGN': 'Q', 'MGQ': 'A',
+    'MGV': 'G', 'MGY': 'G', 'MHL': 'L', 'MHO': 'M', 'MHS': 'H', 'MIA': 'A',
+    'MIS': 'S', 'MK8': 'L', 'ML3': 'K', 'MLE': 'L', 'MLL': 'L', 'MLY': 'K',
+    'MLZ': 'K', 'MME': 'M', 'MMO': 'R', 'MMT': 'T', 'MND': 'N', 'MNL': 'L',
+    'MNU': 'U', 'MNV': 'V', 'MOD': 'X', 'MP8': 'P', 'MPH': 'X', 'MPJ': 'X',
+    'MPQ': 'G', 'MRG': 'G', 'MSA': 'G', 'MSE': 'M', 'MSL': 'M', 'MSO': 'M',
+    'MSP': 'X', 'MT2': 'M', 'MTR': 'T', 'MTU': 'A', 'MTY': 'Y', 'MVA': 'V',
+    'N': 'N', 'N10': 'S', 'N2C': 'X', 'N5I': 'N', 'N5M': 'C', 'N6G': 'G',
+    'N7P': 'P', 'NA8': 'A', 'NAL': 'A', 'NAM': 'A', 'NB8': 'N', 'NBQ': 'Y',
+    'NC1': 'S', 'NCB': 'A', 'NCX': 'N', 'NCY': 'X', 'NDF': 'F', 'NDN': 'U',
+    'NEM': 'H', 'NEP': 'H', 'NF2': 'N', 'NFA': 'F', 'NHL': 'E', 'NIT': 'X',
+    'NIY': 'Y', 'NLE': 'L', 'NLN': 'L', 'NLO': 'L', 'NLP': 'L', 'NLQ': 'Q',
+    'NMC': 'G', 'NMM': 'R', 'NMS': 'T', 'NMT': 'T', 'NNH': 'R', 'NP3': 'N',
+    'NPH': 'C', 'NPI': 'A', 'NSK': 'X', 'NTY': 'Y', 'NVA': 'V', 'NYM': 'N',
+    'NYS': 'C', 'NZH': 'H', 'O12': 'X', 'O2C': 'N', 'O2G': 'G', 'OAD': 'N',
+    'OAS': 'S', 'OBF': 'X', 'OBS': 'X', 'OCS': 'C', 'OCY': 'C', 'ODP': 'N',
+    'OHI': 'H', 'OHS': 'D', 'OIC': 'X', 'OIP': 'I', 'OLE': 'X', 'OLT': 'T',
+    'OLZ': 'S', 'OMC': 'C', 'OMG': 'G', 'OMT': 'M', 'OMU': 'U', 'ONE': 'U',
+    'ONH': 'A', 'ONL': 'X', 'OPR': 'R', 'ORN': 'A', 'ORQ': 'R', 'OSE': 'S',
+    'OTB': 'X', 'OTH': 'T', 'OTY': 'Y', 'OXX': 'D', 'P': 'G', 'P1L': 'C',
+    'P1P': 'N', 'P2T': 'T', 'P2U': 'U', 'P2Y': 'P', 'P5P': 'A', 'PAQ': 'Y',
+    'PAS': 'D', 'PAT': 'W', 'PAU': 'A', 'PBB': 'C', 'PBF': 'F', 'PBT': 'N',
+    'PCA': 'E', 'PCC': 'P', 'PCE': 'X', 'PCS': 'F', 'PDL': 'X', 'PDU': 'U',
+    'PEC': 'C', 'PF5': 'F', 'PFF': 'F', 'PFX': 'X', 'PG1': 'S', 'PG7': 'G',
+    'PG9': 'G', 'PGL': 'X', 'PGN': 'G', 'PGP': 'G', 'PGY': 'G', 'PHA': 'F',
+    'PHD': 'D', 'PHE': 'F', 'PHI': 'F', 'PHL': 'F', 'PHM': 'F', 'PIV': 'X',
+    'PLE': 'L', 'PM3': 'F', 'PMT': 'C', 'POM': 'P', 'PPN': 'F', 'PPU': 'A',
+    'PPW': 'G', 'PQ1': 'N', 'PR3': 'C', 'PR5': 'A', 'PR9': 'P', 'PRN': 'A',
+    'PRO': 'P', 'PRS': 'P', 'PSA': 'F', 'PSH': 'H', 'PST': 'T', 'PSU': 'U',
+    'PSW': 'C', 'PTA': 'X', 'PTH': 'Y', 'PTM': 'Y', 'PTR': 'Y', 'PU': 'A',
+    'PUY': 'N', 'PVH': 'H', 'PVL': 'X', 'PYA': 'A', 'PYO': 'U', 'PYX': 'C',
+    'PYY': 'N', 'QMM': 'Q', 'QPA': 'C', 'QPH': 'F', 'QUO': 'G', 'R': 'A',
+    'R1A': 'C', 'R4K': 'W', 'RE0': 'W', 'RE3': 'W', 'RIA': 'A', 'RMP': 'A',
+    'RON': 'X', 'RT': 'T', 'RTP': 'N', 'S1H': 'S', 'S2C': 'C', 'S2D': 'A',
+    'S2M': 'T', 'S2P': 'A', 'S4A': 'A', 'S4C': 'C', 'S4G': 'G', 'S4U': 'U',
+    'S6G': 'G', 'SAC': 'S', 'SAH': 'C', 'SAR': 'G', 'SBL': 'S', 'SC': 'C',
+    'SCH': 'C', 'SCS': 'C', 'SCY': 'C', 'SD2': 'X', 'SDG': 'G', 'SDP': 'S',
+    'SEB': 'S', 'SEC': 'A', 'SEG': 'A', 'SEL': 'S', 'SEM': 'S', 'SEN': 'S',
+    'SEP': 'S', 'SER': 'S', 'SET': 'S', 'SGB': 'S', 'SHC': 'C', 'SHP': 'G',
+    'SHR': 'K', 'SIB': 'C', 'SLA': 'P', 'SLR': 'P', 'SLZ': 'K', 'SMC': 'C',
+    'SME': 'M', 'SMF': 'F', 'SMP': 'A', 'SMT': 'T', 'SNC': 'C', 'SNN': 'N',
+    'SOC': 'C', 'SOS': 'N', 'SOY': 'S', 'SPT': 'T', 'SRA': 'A', 'SSU': 'U',
+    'STY': 'Y', 'SUB': 'X', 'SUN': 'S', 'SUR': 'U', 'SVA': 'S', 'SVV': 'S',
+    'SVW': 'S', 'SVX': 'S', 'SVY': 'S', 'SVZ': 'X', 'SYS': 'C', 'T': 'T',
+    'T11': 'F', 'T23': 'T', 'T2S': 'T', 'T2T': 'N', 'T31': 'U', 'T32': 'T',
+    'T36': 'T', 'T37': 'T', 'T38': 'T', 'T39': 'T', 'T3P': 'T', 'T41': 'T',
+    'T48': 'T', 'T49': 'T', 'T4S': 'T', 'T5O': 'U', 'T5S': 'T', 'T66': 'X',
+    'T6A': 'A', 'TA3': 'T', 'TA4': 'X', 'TAF': 'T', 'TAL': 'N', 'TAV': 'D',
+    'TBG': 'V', 'TBM': 'T', 'TC1': 'C', 'TCP': 'T', 'TCQ': 'Y', 'TCR': 'W',
+    'TCY': 'A', 'TDD': 'L', 'TDY': 'T', 'TFE': 'T', 'TFO': 'A', 'TFQ': 'F',
+    'TFT': 'T', 'TGP': 'G', 'TH6': 'T', 'THC': 'T', 'THO': 'X', 'THR': 'T',
+    'THX': 'N', 'THZ': 'R', 'TIH': 'A', 'TLB': 'N', 'TLC': 'T', 'TLN': 'U',
+    'TMB': 'T', 'TMD': 'T', 'TNB': 'C', 'TNR': 'S', 'TOX': 'W', 'TP1': 'T',
+    'TPC': 'C', 'TPG': 'G', 'TPH': 'X', 'TPL': 'W', 'TPO': 'T', 'TPQ': 'Y',
+    'TQI': 'W', 'TQQ': 'W', 'TRF': 'W', 'TRG': 'K', 'TRN': 'W', 'TRO': 'W',
+    'TRP': 'W', 'TRQ': 'W', 'TRW': 'W', 'TRX': 'W', 'TS': 'N', 'TST': 'X',
+    'TT': 'N', 'TTD': 'T', 'TTI': 'U', 'TTM': 'T', 'TTQ': 'W', 'TTS': 'Y',
+    'TY1': 'Y', 'TY2': 'Y', 'TY3': 'Y', 'TY5': 'Y', 'TYB': 'Y', 'TYI': 'Y',
+    'TYJ': 'Y', 'TYN': 'Y', 'TYO': 'Y', 'TYQ': 'Y', 'TYR': 'Y', 'TYS': 'Y',
+    'TYT': 'Y', 'TYU': 'N', 'TYW': 'Y', 'TYX': 'X', 'TYY': 'Y', 'TZB': 'X',
+    'TZO': 'X', 'U': 'U', 'U25': 'U', 'U2L': 'U', 'U2N': 'U', 'U2P': 'U',
+    'U31': 'U', 'U33': 'U', 'U34': 'U', 'U36': 'U', 'U37': 'U', 'U8U': 'U',
+    'UAR': 'U', 'UCL': 'U', 'UD5': 'U', 'UDP': 'N', 'UFP': 'N', 'UFR': 'U',
+    'UFT': 'U', 'UMA': 'A', 'UMP': 'U', 'UMS': 'U', 'UN1': 'X', 'UN2': 'X',
+    'UNK': 'X', 'UR3': 'U', 'URD': 'U', 'US1': 'U', 'US2': 'U', 'US3': 'T',
+    'US5': 'U', 'USM': 'U', 'VAD': 'V', 'VAF': 'V', 'VAL': 'V', 'VB1': 'K',
+    'VDL': 'X', 'VLL': 'X', 'VLM': 'X', 'VMS': 'X', 'VOL': 'X', 'X': 'G',
+    'X2W': 'E', 'X4A': 'N', 'XAD': 'A', 'XAE': 'N', 'XAL': 'A', 'XAR': 'N',
+    'XCL': 'C', 'XCN': 'C', 'XCP': 'X', 'XCR': 'C', 'XCS': 'N', 'XCT': 'C',
+    'XCY': 'C', 'XGA': 'N', 'XGL': 'G', 'XGR': 'G', 'XGU': 'G', 'XPR': 'P',
+    'XSN': 'N', 'XTH': 'T', 'XTL': 'T', 'XTR': 'T', 'XTS': 'G', 'XTY': 'N',
+    'XUA': 'A', 'XUG': 'G', 'XX1': 'K', 'Y': 'A', 'YCM': 'C', 'YG': 'G',
+    'YOF': 'Y', 'YRR': 'N', 'YYG': 'G', 'Z': 'C', 'Z01': 'A', 'ZAD': 'A',
+    'ZAL': 'A', 'ZBC': 'C', 'ZBU': 'U', 'ZCL': 'F', 'ZCY': 'C', 'ZDU': 'U',
+    'ZFB': 'X', 'ZGU': 'G', 'ZHP': 'N', 'ZTH': 'T', 'ZU0': 'T', 'ZZJ': 'A',
+}
+# common_typos_enable
+# pyformat: enable
+
+
+@functools.lru_cache(maxsize=64)
+def letters_three_to_one(restype: str, *, default: str) -> str:
+  """Returns single letter name if one exists otherwise returns default."""
+  return CCD_NAME_TO_ONE_LETTER.get(restype, default)
+
+
+ALA = sys.intern('ALA')
+ARG = sys.intern('ARG')
+ASN = sys.intern('ASN')
+ASP = sys.intern('ASP')
+CYS = sys.intern('CYS')
+GLN = sys.intern('GLN')
+GLU = sys.intern('GLU')
+GLY = sys.intern('GLY')
+HIS = sys.intern('HIS')
+ILE = sys.intern('ILE')
+LEU = sys.intern('LEU')
+LYS = sys.intern('LYS')
+MET = sys.intern('MET')
+PHE = sys.intern('PHE')
+PRO = sys.intern('PRO')
+SER = sys.intern('SER')
+THR = sys.intern('THR')
+TRP = sys.intern('TRP')
+TYR = sys.intern('TYR')
+VAL = sys.intern('VAL')
+UNK = sys.intern('UNK')
+GAP = sys.intern('-')
+
+# Unknown ligand.
+UNL = sys.intern('UNL')
+
+# Non-standard version of MET (with Se instead of S), but often appears in PDB.
+MSE = sys.intern('MSE')
+
+# 20 standard protein amino acids (no unknown).
+PROTEIN_TYPES: tuple[str, ...] = (
+    ALA, ARG, ASN, ASP, CYS, GLN, GLU, GLY, HIS, ILE, LEU, LYS, MET, PHE, PRO,
+    SER, THR, TRP, TYR, VAL,
+)  # pyformat: disable
+
+# 20 standard protein amino acids plus the unknown (UNK) amino acid.
+PROTEIN_TYPES_WITH_UNKNOWN: tuple[str, ...] = PROTEIN_TYPES + (UNK,)
+
+# This is the standard residue order when coding AA type as a number.
+# Reproduce it by taking 3-letter AA codes and sorting them alphabetically.
+# For legacy reasons this only refers to protein residues.
+
+PROTEIN_TYPES_ONE_LETTER: tuple[str, ...] = (
+    'A', 'R', 'N', 'D', 'C', 'Q', 'E', 'G', 'H', 'I', 'L', 'K', 'M', 'F', 'P',
+    'S', 'T', 'W', 'Y', 'V',
+)  # pyformat: disable
+
+PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN: tuple[str, ...] = (
+    PROTEIN_TYPES_ONE_LETTER + ('X',)
+)
+PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP: tuple[str, ...] = (
+    PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN + (GAP,)
+)
+
+PROTEIN_TYPES_ONE_LETTER_TO_INT: Mapping[str, int] = {
+    r: i for i, r in enumerate(PROTEIN_TYPES_ONE_LETTER)
+}
+PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_TO_INT: Mapping[str, int] = {
+    r: i for i, r in enumerate(PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN)
+}
+
+PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP_TO_INT: Mapping[str, int] = {
+    r: i for i, r in enumerate(PROTEIN_TYPES_ONE_LETTER_WITH_UNKNOWN_AND_GAP)
+}
+
+
+PROTEIN_COMMON_ONE_TO_THREE: Mapping[str, str] = {
+    'A': ALA,
+    'R': ARG,
+    'N': ASN,
+    'D': ASP,
+    'C': CYS,
+    'Q': GLN,
+    'E': GLU,
+    'G': GLY,
+    'H': HIS,
+    'I': ILE,
+    'L': LEU,
+    'K': LYS,
+    'M': MET,
+    'F': PHE,
+    'P': PRO,
+    'S': SER,
+    'T': THR,
+    'W': TRP,
+    'Y': TYR,
+    'V': VAL,
+}
+
+PROTEIN_COMMON_THREE_TO_ONE: Mapping[str, str] = {
+    v: k for k, v in PROTEIN_COMMON_ONE_TO_THREE.items()
+}
+
+A = sys.intern('A')
+G = sys.intern('G')
+C = sys.intern('C')
+U = sys.intern('U')
+T = sys.intern('T')
+
+DA = sys.intern('DA')
+DG = sys.intern('DG')
+DC = sys.intern('DC')
+DT = sys.intern('DT')
+
+UNK_NUCLEIC_ONE_LETTER = sys.intern('N')  # Unknown nucleic acid single letter.
+UNK_RNA = sys.intern('N')  # Unknown RNA.
+UNK_DNA = sys.intern('DN')  # Unknown DNA residue (differs from N).
+
+RNA_TYPES: tuple[str, ...] = (A, G, C, U)
+DNA_TYPES: tuple[str, ...] = (DA, DG, DC, DT)
+
+NUCLEIC_TYPES: tuple[str, ...] = RNA_TYPES + DNA_TYPES
+# Without UNK DNA.
+NUCLEIC_TYPES_WITH_UNKNOWN: tuple[str, ...] = NUCLEIC_TYPES + (
+    UNK_NUCLEIC_ONE_LETTER,
+)
+NUCLEIC_TYPES_WITH_2_UNKS: tuple[str, ...] = NUCLEIC_TYPES + (
+    UNK_RNA,
+    UNK_DNA,
+)
+
+RNA_TYPES_ONE_LETTER_WITH_UNKNOWN: tuple[str, ...] = RNA_TYPES + (UNK_RNA,)
+RNA_TYPES_ONE_LETTER_WITH_UNKNOWN_TO_INT: Mapping[str, int] = {
+    r: i for i, r in enumerate(RNA_TYPES_ONE_LETTER_WITH_UNKNOWN)
+}
+
+DNA_TYPES_WITH_UNKNOWN: tuple[str, ...] = DNA_TYPES + (UNK_DNA,)
+DNA_TYPES_ONE_LETTER: tuple[str, ...] = (A, G, C, T)
+DNA_TYPES_ONE_LETTER_WITH_UNKNOWN: tuple[str, ...] = DNA_TYPES_ONE_LETTER + (
+    UNK_NUCLEIC_ONE_LETTER,
+)
+DNA_TYPES_ONE_LETTER_WITH_UNKNOWN_TO_INT: Mapping[str, int] = {
+    r: i for i, r in enumerate(DNA_TYPES_ONE_LETTER_WITH_UNKNOWN)
+}
+DNA_COMMON_ONE_TO_TWO: Mapping[str, str] = {
+    'A': 'DA',
+    'G': 'DG',
+    'C': 'DC',
+    'T': 'DT',
+}
+
+STANDARD_POLYMER_TYPES: tuple[str, ...] = PROTEIN_TYPES + NUCLEIC_TYPES
+POLYMER_TYPES: tuple[str, ...] = PROTEIN_TYPES_WITH_UNKNOWN + NUCLEIC_TYPES
+POLYMER_TYPES_WITH_UNKNOWN: tuple[str, ...] = (
+    PROTEIN_TYPES_WITH_UNKNOWN + NUCLEIC_TYPES_WITH_UNKNOWN
+)
+POLYMER_TYPES_WITH_GAP: tuple[str, ...] = PROTEIN_TYPES + (GAP,) + NUCLEIC_TYPES
+POLYMER_TYPES_WITH_UNKNOWN_AND_GAP: tuple[str, ...] = (
+    PROTEIN_TYPES_WITH_UNKNOWN + (GAP,) + NUCLEIC_TYPES_WITH_UNKNOWN
+)
+POLYMER_TYPES_WITH_ALL_UNKS_AND_GAP: tuple[str, ...] = (
+    PROTEIN_TYPES_WITH_UNKNOWN + (GAP,) + NUCLEIC_TYPES_WITH_2_UNKS
+)
+
+POLYMER_TYPES_ORDER = {restype: i for i, restype in enumerate(POLYMER_TYPES)}
+
+POLYMER_TYPES_ORDER_WITH_UNKNOWN = {
+    restype: i for i, restype in enumerate(POLYMER_TYPES_WITH_UNKNOWN)
+}
+
+POLYMER_TYPES_ORDER_WITH_UNKNOWN_AND_GAP = {
+    restype: i for i, restype in enumerate(POLYMER_TYPES_WITH_UNKNOWN_AND_GAP)
+}
+
+POLYMER_TYPES_ORDER_WITH_ALL_UNKS_AND_GAP = {
+    restype: i for i, restype in enumerate(POLYMER_TYPES_WITH_ALL_UNKS_AND_GAP)
+}
+
+POLYMER_TYPES_NUM = len(POLYMER_TYPES)  # := 29.
+POLYMER_TYPES_NUM_WITH_UNKNOWN = len(POLYMER_TYPES_WITH_UNKNOWN)  # := 30.
+POLYMER_TYPES_NUM_WITH_GAP = len(POLYMER_TYPES_WITH_GAP)  # := 29.
+POLYMER_TYPES_NUM_WITH_UNKNOWN_AND_GAP = len(
+    POLYMER_TYPES_WITH_UNKNOWN_AND_GAP
+)  # := 31.
+POLYMER_TYPES_NUM_ORDER_WITH_ALL_UNKS_AND_GAP = len(
+    POLYMER_TYPES_WITH_ALL_UNKS_AND_GAP
+)  # := 32.
+
+WATER_TYPES: tuple[str, ...] = ('HOH', 'DOD')
+
+UNKNOWN_TYPES: tuple[str, ...] = (UNK, UNK_RNA, UNK_DNA, UNL)

flax_model/alphafold3/constants/side_chains.py

@@ -0,0 +1,103 @@
+
+
+"""Constants associated with side chains."""
+
+from collections.abc import Mapping, Sequence
+import itertools
+
+# Format: The list for each AA type contains chi1, chi2, chi3, chi4 in
+# this order (or a relevant subset from chi1 onwards). ALA and GLY don't have
+# chi angles so their chi angle lists are empty.
+CHI_ANGLES_ATOMS: Mapping[str, Sequence[tuple[str, ...]]] = {
+    'ALA': [],
+    # Chi5 in arginine is always 0 +- 5 degrees, so ignore it.
+    'ARG': [
+        ('N', 'CA', 'CB', 'CG'),
+        ('CA', 'CB', 'CG', 'CD'),
+        ('CB', 'CG', 'CD', 'NE'),
+        ('CG', 'CD', 'NE', 'CZ'),
+    ],
+    'ASN': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'OD1')],
+    'ASP': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'OD1')],
+    'CYS': [('N', 'CA', 'CB', 'SG')],
+    'GLN': [
+        ('N', 'CA', 'CB', 'CG'),
+        ('CA', 'CB', 'CG', 'CD'),
+        ('CB', 'CG', 'CD', 'OE1'),
+    ],
+    'GLU': [
+        ('N', 'CA', 'CB', 'CG'),
+        ('CA', 'CB', 'CG', 'CD'),
+        ('CB', 'CG', 'CD', 'OE1'),
+    ],
+    'GLY': [],
+    'HIS': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'ND1')],
+    'ILE': [('N', 'CA', 'CB', 'CG1'), ('CA', 'CB', 'CG1', 'CD1')],
+    'LEU': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'CD1')],
+    'LYS': [
+        ('N', 'CA', 'CB', 'CG'),
+        ('CA', 'CB', 'CG', 'CD'),
+        ('CB', 'CG', 'CD', 'CE'),
+        ('CG', 'CD', 'CE', 'NZ'),
+    ],
+    'MET': [
+        ('N', 'CA', 'CB', 'CG'),
+        ('CA', 'CB', 'CG', 'SD'),
+        ('CB', 'CG', 'SD', 'CE'),
+    ],
+    'PHE': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'CD1')],
+    'PRO': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'CD')],
+    'SER': [('N', 'CA', 'CB', 'OG')],
+    'THR': [('N', 'CA', 'CB', 'OG1')],
+    'TRP': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'CD1')],
+    'TYR': [('N', 'CA', 'CB', 'CG'), ('CA', 'CB', 'CG', 'CD1')],
+    'VAL': [('N', 'CA', 'CB', 'CG1')],
+}
+
+CHI_GROUPS_FOR_ATOM = {}
+for res_name, chi_angle_atoms_for_res in CHI_ANGLES_ATOMS.items():
+  for chi_group_i, chi_group in enumerate(chi_angle_atoms_for_res):
+    for atom_i, atom in enumerate(chi_group):
+      CHI_GROUPS_FOR_ATOM.setdefault((res_name, atom), []).append(
+          (chi_group_i, atom_i)
+      )
+
+# Mapping from (residue_name, atom_name) pairs to the atom's chi group index
+# and atom index within that group.
+CHI_GROUPS_FOR_ATOM: Mapping[tuple[str, str], Sequence[tuple[int, int]]] = (
+    CHI_GROUPS_FOR_ATOM
+)
+
+MAX_NUM_CHI_ANGLES: int = 4
+ATOMS_PER_CHI_ANGLE: int = 4
+
+# A list of atoms for each AA type that are involved in chi angle calculations.
+CHI_ATOM_SETS: Mapping[str, set[str]] = {
+    residue_name: set(itertools.chain(*atoms))
+    for residue_name, atoms in CHI_ANGLES_ATOMS.items()
+}
+
+# If chi angles given in fixed-length array, this matrix determines how to mask
+# them for each AA type. The order is as per restype_order (see below).
+CHI_ANGLES_MASK: Sequence[Sequence[float]] = (
+    (0.0, 0.0, 0.0, 0.0),  # ALA
+    (1.0, 1.0, 1.0, 1.0),  # ARG
+    (1.0, 1.0, 0.0, 0.0),  # ASN
+    (1.0, 1.0, 0.0, 0.0),  # ASP
+    (1.0, 0.0, 0.0, 0.0),  # CYS
+    (1.0, 1.0, 1.0, 0.0),  # GLN
+    (1.0, 1.0, 1.0, 0.0),  # GLU
+    (0.0, 0.0, 0.0, 0.0),  # GLY
+    (1.0, 1.0, 0.0, 0.0),  # HIS
+    (1.0, 1.0, 0.0, 0.0),  # ILE
+    (1.0, 1.0, 0.0, 0.0),  # LEU
+    (1.0, 1.0, 1.0, 1.0),  # LYS
+    (1.0, 1.0, 1.0, 0.0),  # MET
+    (1.0, 1.0, 0.0, 0.0),  # PHE
+    (1.0, 1.0, 0.0, 0.0),  # PRO
+    (1.0, 0.0, 0.0, 0.0),  # SER
+    (1.0, 0.0, 0.0, 0.0),  # THR
+    (1.0, 1.0, 0.0, 0.0),  # TRP
+    (1.0, 1.0, 0.0, 0.0),  # TYR
+    (1.0, 0.0, 0.0, 0.0),  # VAL
+)

flax_model/alphafold3/cpp.cc

@@ -0,0 +1,45 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// 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
+
+#include "alphafold3/data/cpp/msa_profile_pybind.h"
+#include "alphafold3/model/mkdssp_pybind.h"
+#include "alphafold3/parsers/cpp/cif_dict_pybind.h"
+#include "alphafold3/parsers/cpp/fasta_iterator_pybind.h"
+#include "alphafold3/parsers/cpp/msa_conversion_pybind.h"
+#include "alphafold3/structure/cpp/aggregation_pybind.h"
+#include "alphafold3/structure/cpp/membership_pybind.h"
+#include "alphafold3/structure/cpp/mmcif_atom_site_pybind.h"
+#include "alphafold3/structure/cpp/mmcif_layout_pybind.h"
+#include "alphafold3/structure/cpp/mmcif_struct_conn_pybind.h"
+#include "alphafold3/structure/cpp/mmcif_utils_pybind.h"
+#include "alphafold3/structure/cpp/string_array_pybind.h"
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+namespace {
+
+// Include all modules as submodules to simplify building.
+PYBIND11_MODULE(cpp, m) {
+  RegisterModuleCifDict(m.def_submodule("cif_dict"));
+  RegisterModuleFastaIterator(m.def_submodule("fasta_iterator"));
+  RegisterModuleMsaConversion(m.def_submodule("msa_conversion"));
+  RegisterModuleMmcifLayout(m.def_submodule("mmcif_layout"));
+  RegisterModuleMmcifStructConn(m.def_submodule("mmcif_struct_conn"));
+  RegisterModuleMembership(m.def_submodule("membership"));
+  RegisterModuleMmcifUtils(m.def_submodule("mmcif_utils"));
+  RegisterModuleAggregation(m.def_submodule("aggregation"));
+  RegisterModuleStringArray(m.def_submodule("string_array"));
+  RegisterModuleMmcifAtomSite(m.def_submodule("mmcif_atom_site"));
+  RegisterModuleMkdssp(m.def_submodule("mkdssp"));
+  RegisterModuleMsaProfile(m.def_submodule("msa_profile"));
+}
+
+}  // namespace
+}  // namespace alphafold3

flax_model/alphafold3/jackhmmer_seq_limit.patch

@@ -0,0 +1,32 @@
+--- hmmer-3.4/src/jackhmmer.c
++++ hmmer-3.4/src/jackhmmer.c
+@@ -73,6 +73,7 @@ static ESL_OPTIONS options[] = {
+   { "--noali",      eslARG_NONE,        FALSE, NULL, NULL,      NULL,    NULL,  NULL,            "don't output alignments, so output is smaller",                2 },
+   { "--notextw",    eslARG_NONE,         NULL, NULL, NULL,      NULL,    NULL, "--textw",        "unlimit ASCII text output line width",                         2 },
+   { "--textw",      eslARG_INT,         "120", NULL, "n>=120",  NULL,    NULL, "--notextw",      "set max width of ASCII text output lines",                     2 },
++  { "--seq_limit",  eslARG_INT,          NULL, NULL, NULL,      NULL,    NULL, "--seq_limit",    "if set, truncate all hits after this value is reached",        2 },
+ /* Control of scoring system */
+   { "--popen",      eslARG_REAL,       "0.02", NULL, "0<=x<0.5",NULL,    NULL,  NULL,            "gap open probability",                                         3 },
+   { "--pextend",    eslARG_REAL,        "0.4", NULL, "0<=x<1",  NULL,    NULL,  NULL,            "gap extend probability",                                       3 },
+@@ -298,6 +299,7 @@ output_header(FILE *ofp, ESL_GETOPTS *go
+   if (esl_opt_IsUsed(go, "--noali")      && fprintf(ofp, "# show alignments in output:       no\n")                                                   < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+   if (esl_opt_IsUsed(go, "--notextw")    && fprintf(ofp, "# max ASCII text line length:      unlimited\n")                                            < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+   if (esl_opt_IsUsed(go, "--textw")      && fprintf(ofp, "# max ASCII text line length:      %d\n",             esl_opt_GetInteger(go, "--textw"))    < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
++  if (esl_opt_IsUsed(go, "--seq_limit")  && fprintf(ofp, "# set max sequence hits to return: %d\n",             esl_opt_GetInteger(go, "--seq_limit")) < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+   if (esl_opt_IsUsed(go, "--popen")      && fprintf(ofp, "# gap open probability:            %f\n",             esl_opt_GetReal   (go, "--popen"))    < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+   if (esl_opt_IsUsed(go, "--pextend")    && fprintf(ofp, "# gap extend probability:          %f\n",             esl_opt_GetReal   (go, "--pextend"))  < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+   if (esl_opt_IsUsed(go, "--mx")         && fprintf(ofp, "# subst score matrix (built-in):   %s\n",             esl_opt_GetString (go, "--mx"))       < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+@@ -674,6 +676,13 @@ serial_master(ESL_GETOPTS *go, struct cf
+ 	  /* Print the results. */
+ 	  p7_tophits_SortBySortkey(info->th);
+ 	  p7_tophits_Threshold(info->th, info->pli);
++	  /* Limit the number of hits if specified. */
++	  if (esl_opt_IsOn(go, "--seq_limit"))
++	  {
++	    int seq_limit = esl_opt_GetInteger(go, "--seq_limit");
++	    info->th->N = ESL_MIN(info->th->N, seq_limit);
++	  }
++
+ 	  p7_tophits_CompareRanking(info->th, kh, &nnew_targets);
+ 	  p7_tophits_Targets(ofp, info->th, info->pli, textw); if (fprintf(ofp, "\n\n") < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");
+ 	  p7_tophits_Domains(ofp, info->th, info->pli, textw); if (fprintf(ofp, "\n\n") < 0) ESL_EXCEPTION_SYS(eslEWRITE, "write failed");

flax_model/alphafold3/package_config.py

@@ -0,0 +1,68 @@
+#!/usr/bin/env python3
+"""Package metadata helpers for the split AlphaFold3 module."""
+
+ALPHAFOLD3_PACKAGE_DATA = {
+    "flax_model.alphafold3": [
+        "*.so",
+        "*.pyd",
+        "*.dll",
+        "*.dylib",
+        "README.md",
+        "test_data/**/*",
+        "**/*.pyi",
+    ],
+    "flax_model.alphafold3.constants.converters": [
+        "*.pickle",
+    ],
+}
+
+ALPHAFOLD3_MANIFEST_RULES = [
+    "include flax_model/alphafold3/*.so",
+    "include flax_model/alphafold3/*.pyd",
+    "include flax_model/alphafold3/*.dll",
+    "include flax_model/alphafold3/*.dylib",
+    "recursive-include flax_model/alphafold3/constants/converters *.pickle",
+    "include flax_model/alphafold3/README.md",
+    "recursive-include flax_model/alphafold3/test_data *",
+]
+
+
+def get_package_data():
+    return ALPHAFOLD3_PACKAGE_DATA
+
+
+def get_manifest_rules():
+    return ALPHAFOLD3_MANIFEST_RULES
+
+
+_build_done = False
+
+
+def build_hook(project_root, env, python_executable, subprocess_module):
+    del project_root, env, python_executable, subprocess_module
+    global _build_done
+    if _build_done:
+        return
+
+    from flax_model.alphafold3 import _build as af3_build
+
+    print("[AF3] local build hook triggered")
+    try:
+        print(
+            "[AF3] "
+            f"should_build={af3_build.should_build()} "
+            f"artifacts_exist={af3_build.artifacts_exist()} "
+            f"force_rebuild={af3_build.force_rebuild()}"
+        )
+        af3_build.build_if_needed()
+        print("[AF3] local build hook finished")
+    except af3_build.AF3BuildError as exc:
+        print(f"[AF3] local build skipped or failed: {exc}")
+        if af3_build.is_strict():
+            raise
+
+    _build_done = True
+
+
+def get_build_hook():
+    return build_hook

flax_model/alphafold3/structure/__init__.py

@@ -0,0 +1,38 @@
+
+
+"""Structure module initialization."""
+
+# pylint: disable=g-importing-member
+from flax_model.alphafold3.structure.bioassemblies import BioassemblyData
+from flax_model.alphafold3.structure.bonds import Bonds
+from flax_model.alphafold3.structure.chemical_components import ChemCompEntry
+from flax_model.alphafold3.structure.chemical_components import ChemicalComponentsData
+from flax_model.alphafold3.structure.chemical_components import get_data_for_ccd_components
+from flax_model.alphafold3.structure.chemical_components import populate_missing_ccd_data
+from flax_model.alphafold3.structure.mmcif import BondParsingError
+from flax_model.alphafold3.structure.parsing import BondAtomId
+from flax_model.alphafold3.structure.parsing import from_atom_arrays
+from flax_model.alphafold3.structure.parsing import from_mmcif
+from flax_model.alphafold3.structure.parsing import from_parsed_mmcif
+from flax_model.alphafold3.structure.parsing import from_res_arrays
+from flax_model.alphafold3.structure.parsing import from_sequences_and_bonds
+from flax_model.alphafold3.structure.parsing import ModelID
+from flax_model.alphafold3.structure.parsing import NoAtomsError
+from flax_model.alphafold3.structure.parsing import SequenceFormat
+from flax_model.alphafold3.structure.structure import ARRAY_FIELDS
+from flax_model.alphafold3.structure.structure import AuthorNamingScheme
+from flax_model.alphafold3.structure.structure import Bond
+from flax_model.alphafold3.structure.structure import CascadeDelete
+from flax_model.alphafold3.structure.structure import concat
+from flax_model.alphafold3.structure.structure import enumerate_residues
+from flax_model.alphafold3.structure.structure import fix_non_standard_polymer_residues
+from flax_model.alphafold3.structure.structure import GLOBAL_FIELDS
+from flax_model.alphafold3.structure.structure import make_empty_structure
+from flax_model.alphafold3.structure.structure import MissingAtomError
+from flax_model.alphafold3.structure.structure import MissingAuthorResidueIdError
+from flax_model.alphafold3.structure.structure import multichain_residue_index
+from flax_model.alphafold3.structure.structure import stack
+from flax_model.alphafold3.structure.structure import Structure
+from flax_model.alphafold3.structure.structure_tables import Atoms
+from flax_model.alphafold3.structure.structure_tables import Chains
+from flax_model.alphafold3.structure.structure_tables import Residues

flax_model/alphafold3/structure/bioassemblies.py

@@ -0,0 +1,319 @@
+
+
+"""Utilities for parsing and manipulating bioassembly data."""
+
+from collections.abc import Mapping, Sequence
+import copy
+import dataclasses
+from typing import Self
+
+from flax_model.alphafold3.structure import mmcif
+import numpy as np
+
+
+@dataclasses.dataclass(frozen=True)
+class Operation:
+  """A rigid transformation operation."""
+
+  trans: np.ndarray  # shape: (3,)
+  rot: np.ndarray  # shape: (3, 3)
+
+  def apply_to_coords(self, coords: np.ndarray) -> np.ndarray:
+    """Applies the rotation followed by the translation to `coords`."""
+    return np.dot(coords, self.rot.T) + self.trans[np.newaxis, :]
+
+
+@dataclasses.dataclass(frozen=True)
+class Transform:
+  """A rigid transformation composed of a sequence of `Operation`s."""
+
+  # The sequence of operations that form the transform. These will be applied
+  # right-to-left (last-to-first).
+  operations: Sequence[Operation]
+
+  # The chain IDs that this transform should be applied to. These are
+  # label_asym_ids in the mmCIF spec.
+  chain_ids: Sequence[str]
+
+  # A mapping from chain IDs (of chains that participate in this transform)
+  # to their new values in the bioassembly.
+  chain_id_rename_map: Mapping[str, str]
+
+  def apply_to_coords(self, coords: np.ndarray) -> np.ndarray:
+    """Applies the `operations` in right-to-left order."""
+    for operation in reversed(self.operations):
+      coords = operation.apply_to_coords(coords)
+    return coords
+
+
+def _get_operation(oper_data: Mapping[str, str]) -> Operation:
+  """Parses an `Operation` from a mmCIF _pdbx_struct_oper_list row."""
+  trans = np.zeros((3,), dtype=np.float32)
+  rot = np.zeros((3, 3), dtype=np.float32)
+  for i in range(3):
+    trans[i] = float(oper_data[f'_pdbx_struct_oper_list.vector[{i + 1}]'])
+  for i in range(3):
+    for j in range(3):
+      rot[i][j] = float(
+          oper_data[f'_pdbx_struct_oper_list.matrix[{i + 1}][{j + 1}]']
+      )
+  return Operation(trans=trans, rot=rot)
+
+
+class MissingBioassemblyDataError(Exception):
+  """Raised when bioassembly data is missing from an mmCIF."""
+
+
+class BioassemblyData:
+  """Stores and processes bioassembly data from mmCIF tables."""
+
+  # Not all of these columns are required for internal operations, but all
+  # should be present whenever bioassemblies are defined in an mmCIF to stay
+  # consistent with external mmCIFs.
+  _REQUIRED_COLUMNS = (
+      '_pdbx_struct_assembly.id',
+      '_pdbx_struct_assembly.details',
+      '_pdbx_struct_assembly.method_details',
+      '_pdbx_struct_assembly.oligomeric_details',
+      '_pdbx_struct_assembly.oligomeric_count',
+      '_pdbx_struct_assembly_gen.assembly_id',
+      '_pdbx_struct_assembly_gen.oper_expression',
+      '_pdbx_struct_assembly_gen.asym_id_list',
+      '_pdbx_struct_oper_list.id',
+      '_pdbx_struct_oper_list.type',
+      '_pdbx_struct_oper_list.name',
+      '_pdbx_struct_oper_list.symmetry_operation',
+      '_pdbx_struct_oper_list.matrix[1][1]',
+      '_pdbx_struct_oper_list.matrix[1][2]',
+      '_pdbx_struct_oper_list.matrix[1][3]',
+      '_pdbx_struct_oper_list.vector[1]',
+      '_pdbx_struct_oper_list.matrix[2][1]',
+      '_pdbx_struct_oper_list.matrix[2][2]',
+      '_pdbx_struct_oper_list.matrix[2][3]',
+      '_pdbx_struct_oper_list.vector[2]',
+      '_pdbx_struct_oper_list.matrix[3][1]',
+      '_pdbx_struct_oper_list.matrix[3][2]',
+      '_pdbx_struct_oper_list.matrix[3][3]',
+      '_pdbx_struct_oper_list.vector[3]',
+  )
+
+  def __init__(
+      self,
+      *,
+      pdbx_struct_assembly: Mapping[str, Mapping[str, str]],
+      pdbx_struct_assembly_gen: Mapping[str, Sequence[Mapping[str, str]]],
+      pdbx_struct_oper_list: Mapping[str, Mapping[str, str]],
+      assembly_ids: Sequence[str],
+      oper_ids: Sequence[str],
+  ):
+    for assembly_id in assembly_ids:
+      for table, table_name in (
+          (pdbx_struct_assembly, '_pdbx_struct_assembly'),
+          (pdbx_struct_assembly_gen, '_pdbx_struct_assembly_gen'),
+      ):
+        if assembly_id not in table:
+          raise ValueError(
+              f'Assembly ID "{assembly_id}" missing from {table_name} '
+              f'with keys: {table.keys()}'
+          )
+    for oper_id in oper_ids:
+      if oper_id not in pdbx_struct_oper_list:
+        raise ValueError(
+            f'Oper ID "{oper_id}" missing from _pdbx_struct_oper_list '
+            f'with keys: {pdbx_struct_oper_list.keys()}'
+        )
+
+    self._pdbx_struct_assembly = pdbx_struct_assembly
+    self._pdbx_struct_assembly_gen = pdbx_struct_assembly_gen
+    self._pdbx_struct_oper_list = pdbx_struct_oper_list
+    self._operations = {
+        oper_id: _get_operation(oper_data)
+        for oper_id, oper_data in self._pdbx_struct_oper_list.items()
+    }
+    self._assembly_ids = assembly_ids
+    self._oper_ids = oper_ids
+
+  @classmethod
+  def from_mmcif(cls, cif: mmcif.Mmcif) -> Self:
+    """Constructs an instance of `BioassemblyData` from an `Mmcif` object."""
+    for col in cls._REQUIRED_COLUMNS:
+      if col not in cif:
+        raise MissingBioassemblyDataError(col)
+
+    pdbx_struct_assembly = cif.extract_loop_as_dict(
+        prefix='_pdbx_struct_assembly.', index='_pdbx_struct_assembly.id'
+    )
+    pdbx_struct_oper_list = cif.extract_loop_as_dict(
+        prefix='_pdbx_struct_oper_list.', index='_pdbx_struct_oper_list.id'
+    )
+
+    # _pdbx_struct_assembly_gen is unlike the other two tables because it can
+    # have multiple rows share the same assembly ID. This can happen when an
+    # assembly is constructed by applying different sets of transforms to
+    # different sets of chain IDs. Each of these would have its own row.
+    # Here we group rows by their assembly_id.
+    pdbx_struct_assembly_gen = {}
+    for assembly_id, oper_expression, asym_id_list in zip(
+        cif['_pdbx_struct_assembly_gen.assembly_id'],
+        cif['_pdbx_struct_assembly_gen.oper_expression'],
+        cif['_pdbx_struct_assembly_gen.asym_id_list'],
+    ):
+      pdbx_struct_assembly_gen.setdefault(assembly_id, []).append({
+          '_pdbx_struct_assembly_gen.assembly_id': assembly_id,
+          '_pdbx_struct_assembly_gen.oper_expression': oper_expression,
+          '_pdbx_struct_assembly_gen.asym_id_list': asym_id_list,
+      })
+
+    # We provide these separately to keep track of the original order that they
+    # appear in the mmCIF.
+    assembly_ids = cif['_pdbx_struct_assembly.id']
+    oper_ids = cif['_pdbx_struct_oper_list.id']
+    return cls(
+        pdbx_struct_assembly=pdbx_struct_assembly,
+        pdbx_struct_assembly_gen=pdbx_struct_assembly_gen,
+        pdbx_struct_oper_list=pdbx_struct_oper_list,
+        assembly_ids=assembly_ids,
+        oper_ids=oper_ids,
+    )
+
+  @property
+  def assembly_ids(self) -> Sequence[str]:
+    return self._assembly_ids
+
+  def asym_id_by_assembly_chain_id(self, assembly_id: str) -> Mapping[str, str]:
+    asym_id_by_assembly_chain_id = {}
+    for transform in self.get_transforms(assembly_id):
+      for asym_id, assembly_chain_id in transform.chain_id_rename_map.items():
+        asym_id_by_assembly_chain_id[assembly_chain_id] = asym_id
+    return asym_id_by_assembly_chain_id
+
+  def assembly_chain_ids_by_asym_id(
+      self, assembly_id: str
+  ) -> Mapping[str, set[str]]:
+    assembly_chain_ids_by_asym_id = {}
+    for transform in self.get_transforms(assembly_id):
+      for asym_id, assembly_chain_id in transform.chain_id_rename_map.items():
+        assembly_chain_ids_by_asym_id.setdefault(asym_id, set()).add(
+            assembly_chain_id
+        )
+    return assembly_chain_ids_by_asym_id
+
+  def get_default_assembly_id(self) -> str:
+    """Gets a default assembly ID."""
+    # The first assembly is usually (though not always) the best choice.
+    # If we find a better heuristic for picking bioassemblies then this
+    # method should be updated.
+    return min(self._assembly_ids)
+
+  def get_assembly_info(self, assembly_id: str) -> Mapping[str, str]:
+    return {
+        k.replace('_pdbx_struct_assembly.', ''): v
+        for k, v in self._pdbx_struct_assembly[assembly_id].items()
+    }
+
+  def get_transforms(self, assembly_id: str) -> Sequence[Transform]:
+    """Returns the transforms required to generate the given assembly."""
+    partial_transforms = []
+    all_chain_ids = set()
+    for row in self._pdbx_struct_assembly_gen[assembly_id]:
+      oper_expression = row['_pdbx_struct_assembly_gen.oper_expression']
+      parsed_oper_id_seqs = mmcif.parse_oper_expr(oper_expression)
+      label_asym_ids = row['_pdbx_struct_assembly_gen.asym_id_list'].split(',')
+      all_chain_ids |= set(label_asym_ids)
+      for parsed_oper_id_seq in parsed_oper_id_seqs:
+        partial_transforms.append((parsed_oper_id_seq, label_asym_ids))
+
+    # We start assigning new chain IDs by finding the largest chain ID in
+    # the original structure that is involved in this bioassembly, and then
+    # starting from the next one.
+    max_int_chain_id = max(mmcif.str_id_to_int_id(c) for c in all_chain_ids)
+    next_int_chain_id = max_int_chain_id + 1
+
+    transforms = []
+    has_been_renamed = set()
+    for parsed_oper_id_seq, label_asym_ids in partial_transforms:
+      chain_id_rename_map = {}
+      for label_asym_id in label_asym_ids:
+        if label_asym_id not in has_been_renamed:
+          # The first time we see a label_asym_id we don't need to rename it.
+          # This isn't strictly necessary since we don't provide any
+          # guarantees about chain naming after bioassembly extraction but
+          # can make it a bit easier to inspect and compare structures
+          # pre and post bioassembly extraction.
+          chain_id_rename_map[label_asym_id] = label_asym_id
+          has_been_renamed.add(label_asym_id)
+        else:
+          chain_id_rename_map[label_asym_id] = mmcif.int_id_to_str_id(
+              next_int_chain_id
+          )
+          next_int_chain_id += 1
+      transforms.append(
+          Transform(
+              operations=[
+                  self._operations[oper_id] for oper_id in parsed_oper_id_seq
+              ],
+              chain_ids=label_asym_ids,
+              chain_id_rename_map=chain_id_rename_map,
+          )
+      )
+    return transforms
+
+  def to_mmcif_dict(self) -> Mapping[str, Sequence[str]]:
+    """Returns the bioassembly data as a dict suitable for `mmcif.Mmcif`."""
+    mmcif_dict = {}
+    for assembly_id in self._assembly_ids:
+      for column, val in self._pdbx_struct_assembly[assembly_id].items():
+        mmcif_dict.setdefault(column, []).append(val)
+      for row in self._pdbx_struct_assembly_gen[assembly_id]:
+        for column, val in row.items():
+          mmcif_dict.setdefault(column, []).append(val)
+    for oper_id in self._oper_ids:
+      for column, val in self._pdbx_struct_oper_list[oper_id].items():
+        mmcif_dict.setdefault(column, []).append(val)
+    return mmcif_dict
+
+  def rename_label_asym_ids(
+      self,
+      mapping: Mapping[str, str],
+      present_chains: set[str],
+  ) -> Self:
+    """Returns a new BioassemblyData with renamed label_asym_ids.
+
+    Args:
+      mapping: A mapping from original label_asym_ids to their new values. Any
+        label_asym_ids in this BioassemblyData that are not in this mapping will
+        remain unchanged.
+      present_chains: A set of label_asym_ids that are actually present in the
+        atom site list. All label_asym_ids that are in the BioassemblyData but
+        not in present_chains won't be included in the output BioassemblyData.
+
+    Returns:
+      A new BioassemblyData with renamed label_asym_ids.
+
+    Raises:
+      ValueError: If any two previously distinct chains do not have unique names
+          anymore after the rename.
+    """
+    new_pdbx_struct_assembly_gen = copy.deepcopy(self._pdbx_struct_assembly_gen)
+    for rows in new_pdbx_struct_assembly_gen.values():
+      for row in rows:
+        old_asym_ids = row['_pdbx_struct_assembly_gen.asym_id_list'].split(',')
+        new_asym_ids = [
+            mapping.get(label_asym_id, label_asym_id)
+            for label_asym_id in old_asym_ids
+            if label_asym_id in present_chains
+        ]
+        if len(set(old_asym_ids) & present_chains) != len(set(new_asym_ids)):
+          raise ValueError(
+              'Can not rename chains, the new names are not unique: '
+              f'{sorted(new_asym_ids)}.'
+          )
+        row['_pdbx_struct_assembly_gen.asym_id_list'] = ','.join(new_asym_ids)  # pytype: disable=unsupported-operands
+
+    return BioassemblyData(
+        pdbx_struct_assembly=copy.deepcopy(self._pdbx_struct_assembly),
+        pdbx_struct_assembly_gen=new_pdbx_struct_assembly_gen,
+        pdbx_struct_oper_list=copy.deepcopy(self._pdbx_struct_oper_list),
+        assembly_ids=copy.deepcopy(self._assembly_ids),
+        oper_ids=copy.deepcopy(self._oper_ids),
+    )

flax_model/alphafold3/structure/bonds.py

@@ -0,0 +1,227 @@
+
+
+"""Bond representation for structure module."""
+
+import collections
+from collections.abc import Mapping, Sequence
+import dataclasses
+import typing
+from typing import Self
+
+from flax_model.alphafold3.structure import table
+import numpy as np
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Bonds(table.Table):
+  """Table of atomic bonds."""
+
+  # mmCIF column: _struct_conn.conn_type_id
+  # mmCIF desc: This data item is a pointer to _struct_conn_type.id in the
+  #             STRUCT_CONN_TYPE category.
+  # E.g.: "covale", "disulf", "hydrog", "metalc".
+  type: np.ndarray
+
+  # mmCIF column: _struct_conn.pdbx_role
+  # mmCIF desc: The chemical or structural role of the interaction.
+  # E.g.: "N-Glycosylation", "O-Glycosylation".
+  role: np.ndarray
+
+  # mmCIF columns: _struct_conn.ptnr1_*
+  from_atom_key: np.ndarray
+
+  # mmCIF columns: _struct_conn.ptnr2_*
+  dest_atom_key: np.ndarray
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.empty((0,), dtype=np.int64),
+        from_atom_key=np.empty((0,), dtype=np.int64),
+        dest_atom_key=np.empty((0,), dtype=np.int64),
+        type=np.empty((0,), dtype=object),
+        role=np.empty((0,), dtype=object),
+    )
+
+  def get_atom_indices(
+      self,
+      atom_key: np.ndarray,
+  ) -> tuple[np.ndarray, np.ndarray]:
+    """Returns the indices of the from/dest atoms in the atom_key array."""
+    from_atom_missing = ~np.isin(self.from_atom_key, atom_key)
+    dest_atom_missing = ~np.isin(self.dest_atom_key, atom_key)
+    if np.any(from_atom_missing):
+      raise ValueError(
+          f'No atoms for from_atom_key {self.from_atom_key[from_atom_missing]}'
+      )
+    if np.any(dest_atom_missing):
+      raise ValueError(
+          f'No atoms for dest_atom_key {self.dest_atom_key[dest_atom_missing]}'
+      )
+    sort_indices = np.argsort(atom_key)
+    from_indices_sorted = np.searchsorted(
+        atom_key, self.from_atom_key, sorter=sort_indices
+    )
+    dest_indices_sorted = np.searchsorted(
+        atom_key, self.dest_atom_key, sorter=sort_indices
+    )
+    from_indices = sort_indices[from_indices_sorted]
+    dest_indices = sort_indices[dest_indices_sorted]
+    return from_indices, dest_indices
+
+  def restrict_to_atoms(self, atom_key: np.ndarray) -> Self:
+    if not self.size:  # Early-out for empty table.
+      return self
+    from_atom_mask = np.isin(self.from_atom_key, atom_key)
+    dest_atom_mask = np.isin(self.dest_atom_key, atom_key)
+    mask = np.logical_and(from_atom_mask, dest_atom_mask)
+    return typing.cast(Bonds, self.filter(mask=mask))
+
+  def to_mmcif_dict_from_atom_arrays(
+      self,
+      atom_key: np.ndarray,
+      chain_id: np.ndarray,
+      res_id: np.ndarray,
+      res_name: np.ndarray,
+      atom_name: np.ndarray,
+      auth_asym_id: np.ndarray,
+      auth_seq_id: np.ndarray,
+      insertion_code: np.ndarray,
+  ) -> Mapping[str, Sequence[str] | np.ndarray]:
+    """Returns a dict suitable for building a CifDict, representing bonds.
+
+    Args:
+      atom_key: A (num_atom,) integer array of atom_keys.
+      chain_id: A (num_atom,) array of label_asym_id strings.
+      res_id: A (num_atom,) array of label_seq_id strings.
+      res_name: A (num_atom,) array of label_comp_id strings.
+      atom_name: A (num_atom,) array of label_atom_id strings.
+      auth_asym_id: A (num_atom,) array of auth_asym_id strings.
+      auth_seq_id: A (num_atom,) array of auth_seq_id strings.
+      insertion_code: A (num_atom,) array of insertion code strings.
+    """
+    mmcif_dict = collections.defaultdict(list)
+    ptnr1_indices, ptnr2_indices = self.get_atom_indices(atom_key)
+
+    mmcif_dict['_struct_conn.ptnr1_label_asym_id'] = chain_id[ptnr1_indices]
+    mmcif_dict['_struct_conn.ptnr2_label_asym_id'] = chain_id[ptnr2_indices]
+    mmcif_dict['_struct_conn.ptnr1_label_comp_id'] = res_name[ptnr1_indices]
+    mmcif_dict['_struct_conn.ptnr2_label_comp_id'] = res_name[ptnr2_indices]
+    mmcif_dict['_struct_conn.ptnr1_label_seq_id'] = res_id[ptnr1_indices]
+    mmcif_dict['_struct_conn.ptnr2_label_seq_id'] = res_id[ptnr2_indices]
+    mmcif_dict['_struct_conn.ptnr1_label_atom_id'] = atom_name[ptnr1_indices]
+    mmcif_dict['_struct_conn.ptnr2_label_atom_id'] = atom_name[ptnr2_indices]
+
+    mmcif_dict['_struct_conn.ptnr1_auth_asym_id'] = auth_asym_id[ptnr1_indices]
+    mmcif_dict['_struct_conn.ptnr2_auth_asym_id'] = auth_asym_id[ptnr2_indices]
+    mmcif_dict['_struct_conn.ptnr1_auth_seq_id'] = auth_seq_id[ptnr1_indices]
+    mmcif_dict['_struct_conn.ptnr2_auth_seq_id'] = auth_seq_id[ptnr2_indices]
+    mmcif_dict['_struct_conn.pdbx_ptnr1_PDB_ins_code'] = insertion_code[
+        ptnr1_indices
+    ]
+    mmcif_dict['_struct_conn.pdbx_ptnr2_PDB_ins_code'] = insertion_code[
+        ptnr2_indices
+    ]
+
+    label_alt_id = ['?'] * self.size
+    mmcif_dict['_struct_conn.pdbx_ptnr1_label_alt_id'] = label_alt_id
+    mmcif_dict['_struct_conn.pdbx_ptnr2_label_alt_id'] = label_alt_id
+
+    # We need to set this to make visualisation work in NGL/PyMOL.
+    mmcif_dict['_struct_conn.pdbx_value_order'] = ['?'] * self.size
+
+    # We use a symmetry of 1_555 which is the no-op transformation. Other
+    # values are used when bonds involve atoms that only exist after expanding
+    # the bioassembly, but we don't support this kind of bond at the moment.
+    symmetry = ['1_555'] * self.size
+    mmcif_dict['_struct_conn.ptnr1_symmetry'] = symmetry
+    mmcif_dict['_struct_conn.ptnr2_symmetry'] = symmetry
+    bond_type_counter = collections.Counter()
+    for bond_row in self.iterrows():
+      bond_type = bond_row['type']
+      bond_type_counter[bond_type] += 1
+      mmcif_dict['_struct_conn.id'].append(
+          f'{bond_type}{bond_type_counter[bond_type]}'
+      )
+      mmcif_dict['_struct_conn.pdbx_role'].append(bond_row['role'])
+      mmcif_dict['_struct_conn.conn_type_id'].append(bond_type)
+
+    bond_types = np.unique(self.type)
+    mmcif_dict['_struct_conn_type.id'] = bond_types
+    unknown = ['?'] * len(bond_types)
+    mmcif_dict['_struct_conn_type.criteria'] = unknown
+    mmcif_dict['_struct_conn_type.reference'] = unknown
+
+    return dict(mmcif_dict)
+
+
+def concat_with_atom_keys(
+    bonds_tables: Sequence[Bonds | None],
+    atom_key_arrays: Sequence[np.ndarray],
+) -> tuple[Bonds | None, np.ndarray]:
+  """Concatenates bonds tables and atom keys simultaneously.
+
+  Args:
+    bonds_tables: A sequence of `Bonds` instances to concatenate. If any are
+      None then these are skipped.
+    atom_key_arrays: A sequence of integer `atom_key` arrays, where the n-th
+      bonds_table referrs to the atoms in the n-th atom_key array. These must
+      all be non-None.
+
+  Returns:
+    A pair of (bonds, atom_key) where atom_key is a unique atom_key array with
+    length equal to the sum of the input atom array sizes, and the bonds table
+    contains all the bonds from the individual bonds table inputs.
+  """
+  if not bonds_tables or not atom_key_arrays:
+    if bonds_tables or atom_key_arrays:
+      raise ValueError(
+          'bonds_tables and atom_keys must have same length but got'
+          f' {len(bonds_tables)=} and {len(atom_key_arrays)=}'
+      )
+    return None, np.array([], dtype=np.int64)
+  max_key = -1
+  atom_keys_to_concat = []
+  types_to_concat = []
+  roles_to_concat = []
+  from_atom_keys_to_concat = []
+  dest_atom_keys_to_concat = []
+  for bonds, atom_key in zip(bonds_tables, atom_key_arrays, strict=True):
+    if not atom_key.size:
+      assert bonds is None or bonds.size == 0
+      continue
+    assert np.min(atom_key, initial=0) >= 0  # Should always be non-negative!
+    offset = max_key + 1
+    offset_atom_key = atom_key + offset
+    atom_keys_to_concat.append(offset_atom_key)
+    max_key = np.max(offset_atom_key)
+    if bonds is not None:
+      types_to_concat.append(bonds.type)
+      roles_to_concat.append(bonds.role)
+      from_atom_keys_to_concat.append(bonds.from_atom_key + offset)
+      dest_atom_keys_to_concat.append(bonds.dest_atom_key + offset)
+
+  if atom_keys_to_concat:
+    concatted_atom_keys = np.concatenate(atom_keys_to_concat, axis=0)
+  else:
+    concatted_atom_keys = np.array([], dtype=np.int64)
+
+  if types_to_concat:
+    assert (
+        len(types_to_concat)
+        == len(roles_to_concat)
+        == len(from_atom_keys_to_concat)
+        == len(dest_atom_keys_to_concat)
+    )
+    num_bonds = sum(b.size for b in bonds_tables if b is not None)
+    concatted_bonds = Bonds(
+        key=np.arange(num_bonds, dtype=np.int64),
+        type=np.concatenate(types_to_concat, axis=0),
+        role=np.concatenate(roles_to_concat, axis=0),
+        from_atom_key=np.concatenate(from_atom_keys_to_concat, axis=0),
+        dest_atom_key=np.concatenate(dest_atom_keys_to_concat, axis=0),
+    )
+  else:
+    concatted_bonds = None
+
+  return concatted_bonds, concatted_atom_keys

flax_model/alphafold3/structure/chemical_components.py

@@ -0,0 +1,276 @@
+
+
+"""Utilities for manipulating chemical components data."""
+
+from collections.abc import Iterable, Mapping, Sequence
+import dataclasses
+import functools
+from typing import Self
+
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.structure import mmcif
+import rdkit.Chem as rd_chem
+
+
+@dataclasses.dataclass(frozen=True)
+class ChemCompEntry:
+  """Items of _chem_comp category.
+
+  For the full list of items and their semantics see
+  http://mmcif.rcsb.org/dictionaries/mmcif_pdbx_v50.dic/Categories/chem_comp.html
+  """
+
+  type: str
+  name: str = '?'
+  pdbx_synonyms: str = '?'
+  formula: str = '?'
+  formula_weight: str = '?'
+  mon_nstd_flag: str = '?'
+  pdbx_smiles: str | None = None
+
+  def __post_init__(self):
+    for field, value in vars(self).items():
+      if not value and value is not None:
+        raise ValueError(f"{field} value can't be an empty string.")
+
+  def extends(self, other: Self) -> bool:
+    """Checks whether this ChemCompEntry extends another one."""
+    for field, value in vars(self).items():
+      other_value = getattr(other, field)
+      if _value_is_missing(other_value):
+        continue
+      if value != other_value:
+        return False
+    return True
+
+  @property
+  def rdkit_mol(self) -> rd_chem.Mol:
+    """Returns an RDKit Mol, created via RDKit from entry SMILES string."""
+    if not self.pdbx_smiles:
+      raise ValueError('Cannot construct RDKit Mol with empty pdbx_smiles')
+    return rd_chem.MolFromSmiles(self.pdbx_smiles)
+
+
+_REQUIRED_MMCIF_COLUMNS = ('_chem_comp.id', '_chem_comp.type')
+
+
+class MissingChemicalComponentsDataError(Exception):
+  """Raised when chemical components data is missing from an mmCIF."""
+
+
+@dataclasses.dataclass(frozen=True)
+class ChemicalComponentsData:
+  """Extra information for chemical components occurring in mmCIF.
+
+  Fields:
+    chem_comp: A mapping from _chem_comp.id to associated items in the
+      chem_comp category.
+  """
+
+  chem_comp: Mapping[str, ChemCompEntry]
+
+  @classmethod
+  def from_mmcif(
+      cls, cif: mmcif.Mmcif, fix_mse: bool, fix_unknown_dna: bool
+  ) -> Self:
+    """Constructs an instance of ChemicalComponentsData from an Mmcif object."""
+    for col in _REQUIRED_MMCIF_COLUMNS:
+      if col not in cif:
+        raise MissingChemicalComponentsDataError(col)
+
+    id_ = cif['_chem_comp.id']  # Guaranteed to be present.
+    type_ = cif['_chem_comp.type']  # Guaranteed to be present.
+    name = cif.get('_chem_comp.name', ['?'] * len(id_))
+    synonyms = cif.get('_chem_comp.pdbx_synonyms', ['?'] * len(id_))
+    formula = cif.get('_chem_comp.formula', ['?'] * len(id_))
+    weight = cif.get('_chem_comp.formula_weight', ['?'] * len(id_))
+    mon_nstd_flag = cif.get('_chem_comp.mon_nstd_flag', ['?'] * len(id_))
+    smiles = cif.get('_chem_comp.pdbx_smiles', ['?'] * len(id_))
+    smiles = [None if s == '?' else s for s in smiles]
+
+    chem_comp = {
+        component_name: ChemCompEntry(*entry)
+        for component_name, *entry in zip(
+            id_, type_, name, synonyms, formula, weight, mon_nstd_flag, smiles
+        )
+    }
+
+    if fix_mse and 'MSE' in chem_comp:
+      if 'MET' not in chem_comp:
+        chem_comp['MET'] = ChemCompEntry(
+            type='L-PEPTIDE LINKING',
+            name='METHIONINE',
+            pdbx_synonyms='?',
+            formula='C5 H11 N O2 S',
+            formula_weight='149.211',
+            mon_nstd_flag='y',
+            pdbx_smiles=None,
+        )
+
+    if fix_unknown_dna and 'N' in chem_comp:
+      # Do not delete 'N' as it may be needed for RNA in the system.
+      if 'DN' not in chem_comp:
+        chem_comp['DN'] = ChemCompEntry(
+            type='DNA LINKING',
+            name="UNKNOWN 2'-DEOXYNUCLEOTIDE",
+            pdbx_synonyms='?',
+            formula='C5 H11 O6 P',
+            formula_weight='198.111',
+            mon_nstd_flag='y',
+            pdbx_smiles=None,
+        )
+
+    return ChemicalComponentsData(chem_comp)
+
+  def to_mmcif_dict(self) -> Mapping[str, Sequence[str]]:
+    """Returns chemical components data as a dict suitable for `mmcif.Mmcif`."""
+    mmcif_dict = {}
+
+    mmcif_fields = set()
+    for entry in self.chem_comp.values():
+      for field, value in vars(entry).items():
+        if value:
+          mmcif_fields.add(field)
+    chem_comp_ids = []
+    for component_id in sorted(self.chem_comp):
+      entry = self.chem_comp[component_id]
+      chem_comp_ids.append(component_id)
+      for field in mmcif_fields:
+        mmcif_dict.setdefault(f'_chem_comp.{field}', []).append(
+            getattr(entry, field) or '?'
+        )
+    if chem_comp_ids:
+      mmcif_dict['_chem_comp.id'] = chem_comp_ids
+    return mmcif_dict
+
+
+def _value_is_missing(value: str) -> bool:
+  return not value or value in ('.', '?')
+
+
+def get_data_for_ccd_components(
+    ccd: chemical_components.Ccd,
+    chemical_component_ids: Iterable[str],
+    populate_pdbx_smiles: bool = False,
+) -> ChemicalComponentsData:
+  """Returns `ChemicalComponentsData` for chemical components known by PDB."""
+  chem_comp = {}
+  for chemical_component_id in chemical_component_ids:
+    chem_data = chemical_components.component_name_to_info(
+        ccd=ccd, res_name=chemical_component_id
+    )
+    if not chem_data:
+      continue
+    chem_comp[chemical_component_id] = ChemCompEntry(
+        type=chem_data.type,
+        name=chem_data.name,
+        pdbx_synonyms=chem_data.pdbx_synonyms,
+        formula=chem_data.formula,
+        formula_weight=chem_data.formula_weight,
+        mon_nstd_flag=chem_data.mon_nstd_flag,
+        pdbx_smiles=(
+            chem_data.pdbx_smiles or None if populate_pdbx_smiles else None
+        ),
+    )
+  return ChemicalComponentsData(chem_comp=chem_comp)
+
+
+def populate_missing_ccd_data(
+    ccd: chemical_components.Ccd,
+    chemical_components_data: ChemicalComponentsData,
+    chemical_component_ids: Iterable[str] | None = None,
+    populate_pdbx_smiles: bool = False,
+) -> ChemicalComponentsData:
+  """Populates missing data for the chemical components from CCD.
+
+  Args:
+    ccd: The chemical components database.
+    chemical_components_data: ChemicalComponentsData to populate missing values
+      for. This function doesn't modify the object, extended version is provided
+      as a return value.
+    chemical_component_ids: chemical components to populate missing values for.
+      If not specified, the function will consider all chemical components which
+      are already present in `chemical_components_data`.
+    populate_pdbx_smiles: whether to populate `pdbx_smiles` field using SMILES
+      descriptors from _pdbx_chem_comp_descriptor CCD table. If CCD provides
+      multiple SMILES strings, any of them could be used.
+
+  Returns:
+    New instance of ChemicalComponentsData without missing values for CCD
+    entries.
+  """
+  if chemical_component_ids is None:
+    chemical_component_ids = chemical_components_data.chem_comp.keys()
+
+  ccd_data = get_data_for_ccd_components(
+      ccd, chemical_component_ids, populate_pdbx_smiles
+  )
+  chem_comp = dict(chemical_components_data.chem_comp)
+  for component_id, ccd_entry in ccd_data.chem_comp.items():
+    if component_id not in chem_comp:
+      chem_comp[component_id] = ccd_entry
+    else:
+      already_specified_fields = {
+          field: value
+          for field, value in vars(chem_comp[component_id]).items()
+          if not _value_is_missing(value)
+      }
+      chem_comp[component_id] = ChemCompEntry(
+          **{**vars(ccd_entry), **already_specified_fields}
+      )
+  return ChemicalComponentsData(chem_comp=chem_comp)
+
+
+def get_all_atoms_in_entry(
+    ccd: chemical_components.Ccd, res_name: str
+) -> Mapping[str, Sequence[str]]:
+  """Get all possible atoms and bonds for this residue in a standard order.
+
+  Args:
+    ccd: The chemical components dictionary.
+    res_name: Full CCD name.
+
+  Returns:
+    A dictionary table of the atoms and bonds for this residue in this residue
+    type.
+  """
+  # The CCD version of 'UNK' is weird. It has a CB and a CG atom. We just want
+  # the minimal amino-acid here which is GLY.
+  if res_name == 'UNK':
+    res_name = 'GLY'
+  ccd_data = ccd.get(res_name)
+  if not ccd_data:
+    raise ValueError(f'Unknown residue type {res_name}')
+
+  keys = (
+      '_chem_comp_atom.atom_id',
+      '_chem_comp_atom.type_symbol',
+      '_chem_comp_bond.atom_id_1',
+      '_chem_comp_bond.atom_id_2',
+  )
+
+  # Add terminal hydrogens for protonation of the N-terminal
+  if res_name == 'PRO':
+    res_atoms = {key: [*ccd_data.get(key, [])] for key in keys}
+    res_atoms['_chem_comp_atom.atom_id'].extend(['H2', 'H3'])
+    res_atoms['_chem_comp_atom.type_symbol'].extend(['H', 'H'])
+    res_atoms['_chem_comp_bond.atom_id_1'].extend(['N', 'N'])
+    res_atoms['_chem_comp_bond.atom_id_2'].extend(['H2', 'H3'])
+  elif res_name in residue_names.PROTEIN_TYPES_WITH_UNKNOWN:
+    res_atoms = {key: [*ccd_data.get(key, [])] for key in keys}
+    res_atoms['_chem_comp_atom.atom_id'].append('H3')
+    res_atoms['_chem_comp_atom.type_symbol'].append('H')
+    res_atoms['_chem_comp_bond.atom_id_1'].append('N')
+    res_atoms['_chem_comp_bond.atom_id_2'].append('H3')
+  else:
+    res_atoms = {key: ccd_data.get(key, []) for key in keys}
+
+  return res_atoms
+
+
+@functools.lru_cache(maxsize=128)
+def get_res_atom_names(ccd: chemical_components.Ccd, res_name: str) -> set[str]:
+  """Gets the names of the atoms in a given CCD residue."""
+  atoms = get_all_atoms_in_entry(ccd, res_name)['_chem_comp_atom.atom_id']
+  return set(atoms)

flax_model/alphafold3/structure/cpp/aggregation.pyi

@@ -0,0 +1,5 @@
+
+
+from collections.abc import Sequence
+
+def indices_grouped_by_value(values: Sequence[int]) -> dict[int, list[int]]: ...

flax_model/alphafold3/structure/cpp/aggregation_pybind.cc

@@ -0,0 +1,54 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// 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
+
+#include <cstdint>
+#include <vector>
+
+#include "absl/container/flat_hash_map.h"
+#include "absl/types/span.h"
+#include "pybind11/cast.h"
+#include "pybind11/numpy.h"
+#include "pybind11/pybind11.h"
+#include "pybind11_abseil/absl_casters.h"
+
+namespace {
+
+namespace py = pybind11;
+
+absl::flat_hash_map<int64_t, std::vector<int64_t>> IndicesGroupedByValue(
+    absl::Span<const int64_t> values) {
+  absl::flat_hash_map<int64_t, std::vector<int64_t>> group_indices;
+  for (int64_t i = 0, e = values.size(); i < e; ++i) {
+    group_indices[values[i]].push_back(i);
+  }
+  return group_indices;
+}
+
+constexpr char kIndicesGroupedByValue[] = R"(
+Returns a map from value to a list of indices this value occupies.
+
+E.g. indices_grouped_by_value([1, 1, 2, 3, 3, 1, 1]) returns:
+{1: [0, 1, 5, 6], 2: [2], 3: [3, 4]}
+
+Args:
+  values: a list of values to group.
+)";
+
+}  // namespace
+
+namespace alphafold3 {
+
+void RegisterModuleAggregation(py::module m) {
+  m.def("indices_grouped_by_value", &IndicesGroupedByValue, py::arg("values"),
+        py::doc(kIndicesGroupedByValue + 1),
+        py::call_guard<py::gil_scoped_release>());
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/cpp/aggregation_pybind.h

@@ -0,0 +1,24 @@
+/*
+ * Copyright 2024 DeepMind Technologies Limited
+ *
+ * AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+ * this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+ *
+ * To request access to the AlphaFold 3 model parameters, follow the process set
+ * out at https://github.com/google-deepmind/alphafold3. You may only use these
+ * 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
+ */
+
+#ifndef ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_AGGREGATION_PYBIND_H_
+#define ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_AGGREGATION_PYBIND_H_
+
+#include "pybind11/pybind11.h"
+
+namespace alphafold3 {
+
+void RegisterModuleAggregation(pybind11::module m);
+
+}
+
+#endif  // ALPHAFOLD3_SRC_ALPHAFOLD3_STRUCTURE_PYTHON_AGGREGATION_PYBIND_H_

flax_model/alphafold3/structure/cpp/mmcif_atom_site.pyi

@@ -0,0 +1,15 @@
+
+
+from collections.abc import Callable
+from flax_model.alphafold3.cpp import cif_dict
+
+
+def get_internal_to_author_chain_id_map(
+    mmcif: cif_dict.CifDict
+) -> dict[str,str]: ...
+
+
+def get_or_infer_type_symbol(
+    mmcif: cif_dict.CifDict,
+    atom_id_to_type_symbol: Callable[[str, str], str],
+) -> list[str]: ...

flax_model/alphafold3/structure/cpp/mmcif_struct_conn.pyi

@@ -0,0 +1,5 @@
+
+
+from flax_model.alphafold3.cpp import cif_dict
+
+def get_bond_atom_indices(mmcif_dict: cif_dict.CifDict, model_id: str) -> tuple[list[int],list[int]]: ...

flax_model/alphafold3/structure/cpp/mmcif_struct_conn_lib.cc

@@ -0,0 +1,380 @@
+// Copyright 2024 DeepMind Technologies Limited
+//
+// AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+// this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+//
+// To request access to the AlphaFold 3 model parameters, follow the process set
+// out at https://github.com/google-deepmind/alphafold3. You may only use these
+// 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
+
+#include <cstddef>
+#include <iterator>
+#include <optional>
+#include <string>
+#include <utility>
+#include <vector>
+
+#include "absl/algorithm/container.h"
+#include "absl/container/flat_hash_map.h"
+#include "absl/container/flat_hash_set.h"
+#include "absl/status/status.h"
+#include "absl/status/statusor.h"
+#include "absl/strings/str_cat.h"
+#include "absl/strings/string_view.h"
+#include "absl/types/span.h"
+#include "alphafold3/parsers/cpp/cif_dict_lib.h"
+#include "alphafold3/structure/cpp/mmcif_struct_conn.h"
+
+namespace alphafold3 {
+
+namespace {
+
+struct AtomId {
+  absl::string_view chain_id;
+  absl::string_view res_id_1;
+  absl::string_view res_id_2;
+  absl::string_view atom_name;
+  absl::string_view alt_id;
+
+  friend bool operator==(const AtomId&, const AtomId&) = default;
+  template <typename H>
+  friend H AbslHashValue(H h, const AtomId& m) {
+    return H::combine(std::move(h), m.chain_id, m.res_id_1, m.res_id_2,
+                      m.atom_name, m.alt_id);
+  }
+};
+
+using StringArrayRef = absl::Span<const std::string>;
+using BondIndexByAtom = absl::flat_hash_map<AtomId, std::vector<std::size_t>>;
+using BondAtomIndices = std::vector<std::size_t>;
+
+// Returns whether each container is the same size.
+template <typename C, typename... Cs>
+bool AreSameSize(const C& c, const Cs&... cs) {
+  return ((c.size() == cs.size()) && ...);
+}
+
+struct ColumnSpec {
+  absl::string_view chain_id_col;
+  absl::string_view res_id_1_col;
+  absl::string_view res_id_2_col;
+  absl::string_view atom_name_col;
+  std::optional<absl::string_view> alt_id_col;  // Not used by OpenMM.
+};
+
+class AtomColumns {
+ public:
+  static absl::StatusOr<AtomColumns> Create(const CifDict& mmcif,
+                                            const ColumnSpec& column_spec) {
+    StringArrayRef chain_id = mmcif[column_spec.chain_id_col];
+    StringArrayRef res_id_1 = mmcif[column_spec.res_id_1_col];
+    StringArrayRef res_id_2 = mmcif[column_spec.res_id_2_col];
+    StringArrayRef atom_name = mmcif[column_spec.atom_name_col];
+    if (!AreSameSize(chain_id, res_id_1, res_id_2, atom_name)) {
+      return absl::InvalidArgumentError(absl::StrCat(
+          "Atom columns are not the same size. ",                       //
+          "len(", column_spec.chain_id_col, ")=", chain_id.size(),      //
+          ", len(", column_spec.res_id_1_col, ")=", res_id_1.size(),    //
+          ", len(", column_spec.res_id_2_col, ")=", res_id_2.size(),    //
+          ", len(", column_spec.atom_name_col, ")=", atom_name.size(),  //
+          "."));
+    }
+    if (column_spec.alt_id_col.has_value()) {
+      StringArrayRef alt_id = mmcif[*column_spec.alt_id_col];
+      if (!AreSameSize(alt_id, chain_id)) {
+        return absl::InvalidArgumentError(absl::StrCat(
+            "Atom columns are not the same size. ",                   //
+            "len(", column_spec.chain_id_col, ")=", chain_id.size(),  //
+            ", len(", *column_spec.alt_id_col, ")=", alt_id.size(),   //
+            "."));
+      }
+      return AtomColumns(chain_id, res_id_1, res_id_2, atom_name, alt_id,
+                         column_spec);
+    } else {
+      return AtomColumns(chain_id, res_id_1, res_id_2, atom_name, std::nullopt,
+                         column_spec);
+    }
+  }
+
+  inline std::size_t size() const { return size_; }
+
+  absl::string_view GetNormalizedAltId(const std::size_t index) const {
+    constexpr absl::string_view kFullStop = ".";
+    if (alt_id_.has_value()) {
+      absl::string_view alt_id = (*alt_id_)[index];
+      return alt_id == "?" ? kFullStop : alt_id;
+    } else {
+      return kFullStop;
+    }
+  }
+
+  AtomId GetAtom(const std::size_t index) const {
+    return {.chain_id = chain_id_[index],
+            .res_id_1 = res_id_1_[index],
+            .res_id_2 = res_id_2_[index],
+            .atom_name = atom_name_[index],
+            .alt_id = GetNormalizedAltId(index)};
+  }
+
+  std::string GetAtomString(const std::size_t index) const {
+    std::string alt_id_col;
+    if (column_spec_.alt_id_col.has_value()) {
+      alt_id_col = *column_spec_.alt_id_col;
+    } else {
+      alt_id_col = "default label_alt_id";
+    }
+    return absl::StrCat(
+        column_spec_.chain_id_col, "=", chain_id_[index], ", ",    //
+        column_spec_.res_id_1_col, "=", res_id_1_[index], ", ",    //
+        column_spec_.res_id_2_col, "=", res_id_2_[index], ", ",    //
+        column_spec_.atom_name_col, "=", atom_name_[index], ", ",  //
+        alt_id_col, "=", GetNormalizedAltId(index));               //
+  }
+
+ private:
+  AtomColumns(StringArrayRef chain_id, StringArrayRef res_id_1,
+              StringArrayRef res_id_2, StringArrayRef atom_name,
+              std::optional<StringArrayRef> alt_id,
+              const ColumnSpec& column_spec)
+      : chain_id_(chain_id),
+        res_id_1_(res_id_1),
+        res_id_2_(res_id_2),
+        atom_name_(atom_name),
+        alt_id_(alt_id),
+        column_spec_(column_spec),
+        size_(chain_id.size()) {}
+  StringArrayRef chain_id_;
+  StringArrayRef res_id_1_;
+  StringArrayRef res_id_2_;
+  StringArrayRef atom_name_;
+  std::optional<StringArrayRef> alt_id_;
+  ColumnSpec column_spec_;
+  std::size_t size_;
+};
+
+// Adds the atom index to any rows in the bond table involving that atom.
+absl::Status FillInBondsForAtom(const BondIndexByAtom& bond_index_by_atom,
+                                const AtomId& atom,
+                                const std::size_t atom_index,
+                                BondAtomIndices& bond_atom_indices) {
+  if (auto bond_index_it = bond_index_by_atom.find(atom);
+      bond_index_it != bond_index_by_atom.end()) {
+    for (std::size_t bond_index : bond_index_it->second) {
+      if (bond_index < 0 || bond_index >= bond_atom_indices.size()) {
+        return absl::OutOfRangeError(
+            absl::StrCat("Bond index out of range: ", bond_index));
+      }
+      bond_atom_indices[bond_index] = atom_index;
+    }
+  }
+  return absl::OkStatus();
+}
+
+// Checks that the CifDict has all of the columns in the column spec.
+bool HasAllColumns(const CifDict& mmcif, const ColumnSpec& columns) {
+  return mmcif.Contains(columns.chain_id_col) &&
+         mmcif.Contains(columns.res_id_1_col) &&
+         mmcif.Contains(columns.res_id_2_col) &&
+         mmcif.Contains(columns.atom_name_col) &&
+         (!columns.alt_id_col.has_value() ||
+          mmcif.Contains(*columns.alt_id_col));
+}
+
+// Fully specified ptnr1 atom.
+constexpr ColumnSpec kStructConnPtnr1ColumnsFull{
+    .chain_id_col = "_struct_conn.ptnr1_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr1_auth_seq_id",
+    .res_id_2_col = "_struct_conn.pdbx_ptnr1_PDB_ins_code",
+    .atom_name_col = "_struct_conn.ptnr1_label_atom_id",
+    .alt_id_col = "_struct_conn.pdbx_ptnr1_label_alt_id",
+};
+
+// Fully specified ptnr2 atom.
+constexpr ColumnSpec kStructConnPtnr2ColumnsFull{
+    .chain_id_col = "_struct_conn.ptnr2_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr2_auth_seq_id",
+    .res_id_2_col = "_struct_conn.pdbx_ptnr2_PDB_ins_code",
+    .atom_name_col = "_struct_conn.ptnr2_label_atom_id",
+    .alt_id_col = "_struct_conn.pdbx_ptnr2_label_alt_id",
+};
+
+// Columns used by OpenMM for ptnr1 atoms.
+constexpr ColumnSpec kStructConnPtnr1OpenMM{
+    .chain_id_col = "_struct_conn.ptnr1_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr1_label_seq_id",
+    .res_id_2_col = "_struct_conn.ptnr1_label_comp_id",
+    .atom_name_col = "_struct_conn.ptnr1_label_atom_id",
+    .alt_id_col = std::nullopt,
+};
+
+// Columns used by OpenMM for ptnr2 atoms.
+constexpr ColumnSpec kStructConnPtnr2OpenMM{
+    .chain_id_col = "_struct_conn.ptnr2_label_asym_id",
+    .res_id_1_col = "_struct_conn.ptnr2_label_seq_id",
+    .res_id_2_col = "_struct_conn.ptnr2_label_comp_id",
+    .atom_name_col = "_struct_conn.ptnr2_label_atom_id",
+    .alt_id_col = std::nullopt,
+};
+
+// Fully specified atom sites.
+constexpr ColumnSpec kAtomSiteColumnsFull{
+    .chain_id_col = "_atom_site.label_asym_id",
+    .res_id_1_col = "_atom_site.auth_seq_id",
+    .res_id_2_col = "_atom_site.pdbx_PDB_ins_code",
+    .atom_name_col = "_atom_site.label_atom_id",
+    .alt_id_col = "_atom_site.label_alt_id",
+};
+
+// Atom site columns used to match OpenMM _struct_conn tables.
+constexpr ColumnSpec kAtomSiteColumnsOpenMM{
+    .chain_id_col = "_atom_site.label_asym_id",
+    .res_id_1_col = "_atom_site.label_seq_id",
+    .res_id_2_col = "_atom_site.label_comp_id",
+    .atom_name_col = "_atom_site.label_atom_id",
+    .alt_id_col = "_atom_site.label_alt_id",
+};
+
+}  // namespace
+
+absl::StatusOr<std::pair<BondAtomIndices, BondAtomIndices>> GetBondAtomIndices(
+    const CifDict& mmcif, absl::string_view model_id) {
+  ColumnSpec ptnr1_columns, ptnr2_columns, atom_site_columns;
+
+  if (HasAllColumns(mmcif, kStructConnPtnr1ColumnsFull) &&
+      HasAllColumns(mmcif, kStructConnPtnr2ColumnsFull)) {
+    ptnr1_columns = kStructConnPtnr1ColumnsFull;
+    ptnr2_columns = kStructConnPtnr2ColumnsFull;
+    atom_site_columns = kAtomSiteColumnsFull;
+  } else {
+    ptnr1_columns = kStructConnPtnr1OpenMM;
+    ptnr2_columns = kStructConnPtnr2OpenMM;
+    atom_site_columns = kAtomSiteColumnsOpenMM;
+  }
+
+  absl::StatusOr<AtomColumns> ptnr1_atoms =
+      AtomColumns::Create(mmcif, ptnr1_columns);
+  if (!ptnr1_atoms.ok()) {
+    return ptnr1_atoms.status();
+  }
+  absl::StatusOr<AtomColumns> ptnr2_atoms =
+      AtomColumns::Create(mmcif, ptnr2_columns);
+  if (!ptnr2_atoms.ok()) {
+    return ptnr2_atoms.status();
+  }
+  StringArrayRef struct_conn_id = mmcif["_struct_conn.id"];
+  if (!AreSameSize(struct_conn_id, *ptnr1_atoms, *ptnr2_atoms)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Invalid '_struct_conn.' loop. ",                  //
+        "len(id) = ", struct_conn_id.size(), ", ",         //
+        "len(ptnr1_atoms) = ", ptnr1_atoms->size(), ", ",  //
+        "len(ptnr2_atoms) = ", ptnr2_atoms->size(), "."    //
+        ));
+  }
+
+  absl::StatusOr<AtomColumns> atoms =
+      AtomColumns::Create(mmcif, atom_site_columns);
+  if (!atoms.ok()) {
+    return atoms.status();
+  }
+  StringArrayRef atom_site_id = mmcif["_atom_site.id"];
+  StringArrayRef atom_site_model_id = mmcif["_atom_site.pdbx_PDB_model_num"];
+  if (!AreSameSize(atom_site_id, atom_site_model_id, *atoms)) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Invalid '_atom_site.' loop. ",                                //
+        "len(id)= ", atom_site_id.size(), ", ",                        //
+        "len(pdbx_PDB_model_num)= ", atom_site_model_id.size(), ", ",  //
+        "len(atoms)= ", atoms->size(), "."));                          //
+  }
+
+  // Build maps from atom ID tuples to the rows in _struct_conn where that
+  // atom appears (NB could be multiple).
+  const std::size_t struct_conn_size = struct_conn_id.size();
+  BondIndexByAtom ptnr1_rows_by_atom(struct_conn_size);
+  BondIndexByAtom ptnr2_rows_by_atom(struct_conn_size);
+  for (std::size_t i = 0; i < struct_conn_size; ++i) {
+    ptnr1_rows_by_atom[ptnr1_atoms->GetAtom(i)].push_back(i);
+    ptnr2_rows_by_atom[ptnr2_atoms->GetAtom(i)].push_back(i);
+  }
+
+  // Allocate two output arrays with one element per row in struct_conn, where
+  // each element will be the index of that atom in the atom_site table.
+  // Fill the arrays with atom_site_size, which is an invalid value, so that
+  // we can check at the end that each atom has been found.
+  const std::size_t atom_site_size = atom_site_id.size();
+  BondAtomIndices ptnr1_atom_indices(struct_conn_size, atom_site_size);
+  BondAtomIndices ptnr2_atom_indices(struct_conn_size, atom_site_size);
+
+  bool model_id_ecountered = false;
+  absl::flat_hash_set<absl::string_view> seen_alt_ids;
+  for (std::size_t atom_i = 0; atom_i < atom_site_size; ++atom_i) {
+    if (atom_site_model_id[atom_i] != model_id) {
+      if (!model_id_ecountered) {
+        continue;
+      } else {
+        // Models are contiguous so once we see a different model ID after
+        // encountering our model ID then we can exit early.
+        break;
+      }
+    } else {
+      model_id_ecountered = true;
+    }
+    AtomId atom = atoms->GetAtom(atom_i);
+    seen_alt_ids.insert(atom.alt_id);
+
+    if (auto fill_in_bonds_status1 = FillInBondsForAtom(
+            ptnr1_rows_by_atom, atom, atom_i, ptnr1_atom_indices);
+        !fill_in_bonds_status1.ok()) {
+      return fill_in_bonds_status1;
+    }
+    if (auto fill_in_bonds_status2 = FillInBondsForAtom(
+            ptnr2_rows_by_atom, atom, atom_i, ptnr2_atom_indices);
+        !fill_in_bonds_status2.ok()) {
+      return fill_in_bonds_status2;
+    }
+  }
+  // The seen_alt_ids check is a workaround for a known PDB issue: some mmCIFs
+  // (2evw, 2g0v, 2g0x, 2g0z, 2g10, 2g11, 2g12, 2g14, 2grz, 2ntw as of 2024)
+  // have multiple models and they set different whole-chain altloc in each
+  // model. The bond table however doesn't distinguish between models, so there
+  // are bonds that are valid only for some models. E.g. 2grz has model 1 with
+  // chain A with altloc A, and model 2 with chain A with altloc B. The bonds
+  // table lists a bond for each of these.
+
+  // Check that a ptnr1 atom was found for every bond.
+  if (auto row_it = absl::c_find(ptnr1_atom_indices, atom_site_size);
+      row_it != ptnr1_atom_indices.end()) {
+    if (seen_alt_ids.size() > 1 || seen_alt_ids.contains(".") ||
+        seen_alt_ids.contains("?")) {
+      std::size_t i = std::distance(ptnr1_atom_indices.begin(), row_it);
+      return absl::InvalidArgumentError(
+          absl::StrCat("Error parsing \"", mmcif.GetDataName(), "\". ",
+                       "Cannot find atom for bond ID ", struct_conn_id[i], ": ",
+                       ptnr1_atoms->GetAtomString(i)));
+    }
+  }
+
+  // Check that a ptnr2 atom was found for every bond.
+  if (auto row_it = absl::c_find(ptnr2_atom_indices, atom_site_size);
+      row_it != ptnr2_atom_indices.end()) {
+    if (seen_alt_ids.size() > 1 || seen_alt_ids.contains(".") ||
+        seen_alt_ids.contains("?")) {
+      std::size_t i = std::distance(ptnr2_atom_indices.begin(), row_it);
+      return absl::InvalidArgumentError(
+          absl::StrCat("Error parsing \"", mmcif.GetDataName(), "\". ",
+                       "Cannot find atom for bond ID ", struct_conn_id[i], ": ",
+                       ptnr2_atoms->GetAtomString(i)));
+    }
+  }
+
+  if (!model_id_ecountered) {
+    return absl::InvalidArgumentError(absl::StrCat(
+        "Error parsing \"", mmcif.GetDataName(), "\". model_id \"", model_id,
+        "\" not found in _atom_site.pdbx_PDB_model_num."));
+  }
+
+  return std::make_pair(std::move(ptnr1_atom_indices),
+                        std::move(ptnr2_atom_indices));
+}
+
+}  // namespace alphafold3

flax_model/alphafold3/structure/mmcif.py

@@ -0,0 +1,321 @@
+
+
+"""Low level mmCIF parsing operations and wrappers for nicer C++/Py errors.
+
+Note that the cif_dict.CifDict class has many useful methods to help with data
+extraction which are not shown in this file. You can find them in cif_dict.clif
+together with docstrings. The cif_dict.CifDict class behaves like an immutable
+Python dictionary (some methods are not implemented though).
+"""
+from collections.abc import Callable, Mapping, Sequence
+import functools
+import itertools
+import re
+from typing import ParamSpec, TypeAlias, TypeVar
+
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.cpp import cif_dict
+from flax_model.alphafold3.cpp import mmcif_atom_site
+from flax_model.alphafold3.cpp import mmcif_struct_conn
+from flax_model.alphafold3.cpp import string_array
+import numpy as np
+
+Mmcif = cif_dict.CifDict
+
+
+_P = ParamSpec('_P')
+_T = TypeVar('_T')
+_WappedFn: TypeAlias = Callable[_P, _T]
+
+
+@functools.lru_cache(maxsize=256)
+def int_id_to_str_id(num: int) -> str:
+  """Encodes a number as a string, using reverse spreadsheet style naming.
+
+  Args:
+    num: A positive integer.
+
+  Returns:
+    A string that encodes the positive integer using reverse spreadsheet style,
+    naming e.g. 1 = A, 2 = B, ..., 27 = AA, 28 = BA, 29 = CA, ... This is the
+    usual way to encode chain IDs in mmCIF files.
+  """
+  if num <= 0:
+    raise ValueError(f'Only positive integers allowed, got {num}.')
+
+  num = num - 1  # 1-based indexing.
+  output = []
+  while num >= 0:
+    output.append(chr(num % 26 + ord('A')))
+    num = num // 26 - 1
+  return ''.join(output)
+
+
+@functools.lru_cache(maxsize=256)
+def str_id_to_int_id(str_id: str) -> int:
+  """Encodes an mmCIF-style string chain ID as an integer.
+
+  The integer IDs are one based so this function is the inverse of
+  int_id_to_str_id.
+
+  Args:
+    str_id: A string chain ID consisting only of upper case letters A-Z.
+
+  Returns:
+    An integer that can be used to order mmCIF chain IDs in the standard
+    (reverse spreadsheet style) ordering.
+  """
+  if not re.match('^[A-Z]+$', str_id):
+    raise ValueError(f'String ID must be upper case letters, got {str_id}.')
+
+  offset = ord('A') - 1
+  output = 0
+  for i, c in enumerate(str_id):
+    output += (ord(c) - offset) * int(26**i)
+  return output
+
+
+def from_string(mmcif_string: str | bytes) -> Mmcif:
+  return cif_dict.from_string(mmcif_string)
+
+
+def parse_multi_data_cif(cif_string: str) -> dict[str, Mmcif]:
+  """Parses a CIF string with multiple data records.
+
+  For instance, the CIF string:
+
+  ```
+  data_001
+  _foo bar
+  #
+  data_002
+  _foo baz
+  ```
+
+  is parsed as:
+
+  ```
+  {'001': Mmcif({'_foo': ['bar']}), '002': Mmcif({'_foo': ['baz']})}
+  ```
+
+  Args:
+    cif_string: The multi-data CIF string to be parsed.
+
+  Returns:
+    A dictionary mapping record names to Mmcif objects with data.
+  """
+  return cif_dict.parse_multi_data_cif(cif_string)
+
+
+# def tokenize(mmcif_string: str) -> list[str]:
+#   return cif_dict.tokenize(mmcif_string)
+
+
+# def split_line(line: str) -> list[str]:
+#   return cif_dict.split_line(line)
+
+
+class BondParsingError(Exception):
+  """Exception raised by errors when getting bond atom indices."""
+
+
+def get_bond_atom_indices(
+    mmcif: Mmcif,
+    model_id: str = '1',
+) -> tuple[Sequence[int], Sequence[int]]:
+  """Extracts the indices of the atoms that participate in bonds.
+
+  Args:
+    mmcif: The mmCIF object to process.
+    model_id: The ID of the model that the returned atoms will belong to. This
+      should be a value in the mmCIF's _atom_site.pdbx_PDB_model_num column.
+
+  Returns:
+    Two lists of atom indices, `from_atoms` and `to_atoms`, each one having
+    length num_bonds (as defined by _struct_conn, the bonds table). The bond
+    i, defined by the i'th row in _struct_conn, is a bond from atom at index
+    from_atoms[i], to the atom at index to_atoms[i]. The indices are simple
+    0-based indexes into the columns of the _atom_site table in the input
+    mmCIF, and do not necessarily correspond to the values in _atom_site.id,
+    or any other column.
+
+  Raises:
+    BondParsingError: If any of the required tables or columns are not present
+    in
+      the mmCIF, or if the _struct_conn table refers to atoms that cannot
+      be found in the _atom_site table.
+  """
+  try:
+    return mmcif_struct_conn.get_bond_atom_indices(mmcif, model_id)
+  except ValueError as e:
+    raise BondParsingError(str(e)) from e
+
+
+def get_or_infer_type_symbol(
+    mmcif: Mmcif, ccd: chemical_components.Ccd | None = None
+) -> Sequence[str]:
+  """Returns the type symbol (element) for all of the atoms.
+
+  Args:
+    mmcif: A parsed mmCIF file in the Mmcif format.
+    ccd: The chemical component dictionary. If not provided, defaults to the
+      cached CCD.
+
+  If present, returns the _atom_site.type_symbol. If not, infers it using
+  _atom_site.label_comp_id (residue name), _atom_site.label_atom_id (atom name)
+  and the CCD.
+  """
+  ccd = ccd or chemical_components.Ccd()
+  type_symbol_fn = lambda res_name, atom_name: chemical_components.type_symbol(
+      ccd, res_name, atom_name
+  )
+  return mmcif_atom_site.get_or_infer_type_symbol(mmcif, type_symbol_fn)
+
+
+def get_chain_type_by_entity_id(mmcif: Mmcif) -> Mapping[str, str]:
+  """Returns mapping from entity ID to its type or polymer type if available.
+
+  If the entity is in the _entity_poly table, returns its polymer chain type.
+  If not, returns the type as specified in the _entity table.
+
+  Args:
+    mmcif: CifDict holding the mmCIF.
+  """
+  poly_entity_id = mmcif.get('_entity_poly.entity_id', [])
+  poly_type = mmcif.get('_entity_poly.type', [])
+  poly_type_by_entity_id = dict(zip(poly_entity_id, poly_type, strict=True))
+
+  chain_type_by_entity_id = {}
+  for entity_id, entity_type in zip(
+      mmcif.get('_entity.id', []), mmcif.get('_entity.type', []), strict=True
+  ):
+    chain_type = poly_type_by_entity_id.get(entity_id) or entity_type
+    chain_type_by_entity_id[entity_id] = chain_type
+
+  return chain_type_by_entity_id
+
+
+def get_internal_to_author_chain_id_map(mmcif: Mmcif) -> Mapping[str, str]:
+  """Returns a mapping from internal chain ID to the author chain ID.
+
+  Note that this is not a bijection. One author chain ID can map to multiple
+  internal chain IDs. For example, a protein chain and a ligand bound to it will
+  share the same author chain ID, but they will each have a unique internal
+  chain ID).
+
+  Args:
+    mmcif: CifDict holding the mmCIF.
+  """
+  return mmcif_atom_site.get_internal_to_author_chain_id_map(mmcif)
+
+
+def get_experimental_method(mmcif: Mmcif) -> str | None:
+  field = '_exptl.method'
+  return ','.join(mmcif[field]).lower() if field in mmcif else None
+
+
+def get_release_date(mmcif: Mmcif) -> str | None:
+  """Returns the oldest revision date."""
+  if '_pdbx_audit_revision_history.revision_date' not in mmcif:
+    return None
+
+  # Release dates are ISO-8601, hence sort well.
+  return min(mmcif['_pdbx_audit_revision_history.revision_date'])
+
+
+def get_resolution(mmcif: Mmcif) -> float | None:
+  """Returns the resolution of the structure.
+
+  More than one resolution can be reported in an mmCIF. This function returns
+  the first one (in the order _refine.ls_d_res_high,
+  _em_3d_reconstruction.resolution, _reflns.d_resolution_high) that appears
+  in the mmCIF as is parseable as a float.
+
+  Args:
+    mmcif: An `Mmcif` object.
+
+  Returns:
+    The resolution as reported in the mmCIF.
+  """
+  for res_key in ('_refine.ls_d_res_high',
+                  '_em_3d_reconstruction.resolution',
+                  '_reflns.d_resolution_high'):
+    if res_key in mmcif:
+      try:
+        raw_resolution = mmcif[res_key][0]
+        return float(raw_resolution)
+      except ValueError:
+        continue
+  return None
+
+
+def parse_oper_expr(oper_expression: str) -> list[tuple[str, ...]]:
+  """Determines which transforms to apply based on an MMCIF oper_expression str.
+
+  Args:
+    oper_expression: the field oper_expression from MMCIF format data.
+      Transform ids may be either numbers or single letters. Hyphens are used to
+      denote a numeric range of transforms to apply, and commas are used to
+      delimit a sequence of transforms. Where two sets of parentheses are
+      adjacent without a comma, the two sets of transforms should be combined as
+      a cartesian product, i.e. all possible pairs.
+      example 1,2,3 -> generate 3 copies of each chain by applying 1, 2 or 3.
+      example (1-3) -> generate 3 copies of each chain by applying 1, 2 or 3.
+      example (1-3)(4-6) -> generate 9 copies of each chain by applying one of
+        [(1,4), (1,5), (1,6),
+         (2,4), (2,5), (2,6),
+         (3,4), (3,5), (3,6)]
+      example (P) -> apply transform with id P.
+
+  Raises:
+    ValueError: Failure to parse oper_expression.
+
+  Returns:
+    A list with one element for each chain copy that should be generated.
+    Each element is a list of transform ids to apply.
+  """
+  # Expand ranges, e.g. 1-4 -> 1,2,3,4.
+  def range_expander(match):
+    return ','.join(
+        [str(i) for i in range(int(match.group(1)),
+                               int(match.group(2)) + 1)])
+
+  ranges_expanded = re.sub(r'\b(\d+)-(\d+)', range_expander, oper_expression)
+
+  if re.fullmatch(r'(\w+,)*\w+', ranges_expanded):
+    # No brackets, just a single range, e.g. "1,2,3".
+    return [(t,) for t in ranges_expanded.split(',')]
+  elif re.fullmatch(r'\((\w+,)*\w+\)', ranges_expanded):
+    # Single range in brackets, e.g. "(1,2,3)".
+    return [(t,) for t in ranges_expanded[1:-1].split(',')]
+  elif re.fullmatch(r'\((\w+,)*\w+\)\((\w+,)*\w+\)', ranges_expanded):
+    # Cartesian product of two ranges, e.g. "(1,2,3)(4,5)".
+    part1, part2 = ranges_expanded[1:-1].split(')(')
+    return list(itertools.product(part1.split(','), part2.split(',')))
+  else:
+    raise ValueError(f'Unsupported oper_expression format: {oper_expression}')
+
+
+def format_float_array(
+    values: np.ndarray, num_decimal_places: int) -> Sequence[str]:
+  """Converts 1D array to a list of strings with the given number of decimals.
+
+  This function is faster than converting via Python list comprehension, e.g.:
+  atoms_x = ['%.3f' % x for x in atoms_x]
+
+  Args:
+    values: A numpy array with values to convert. This array is casted to
+      float32 before doing the conversion.
+    num_decimal_places: The number of decimal points to keep, including trailing
+      zeros. E.g. for 1.07 and num_decimal_places=1: 1.1,
+      num_decimal_places=2: 1.07, num_decimal_places=3: 1.070.
+
+  Returns:
+    A list of formatted strings.
+  """
+  if values.ndim != 1:
+    raise ValueError(f'The given array must be 1D, got {values.ndim}D')
+
+  return string_array.format_float_array(
+      values=values.astype(np.float32), num_decimal_places=num_decimal_places
+  )

flax_model/alphafold3/structure/parsing.py

@@ -0,0 +1,1801 @@
+
+
+"""Module for parsing various data sources and producing Structures."""
+
+from collections.abc import Collection, Mapping, MutableMapping, Sequence
+import dataclasses
+import datetime
+import enum
+import functools
+import itertools
+from typing import TypeAlias
+
+from flax_model.alphafold3.constants import chemical_components
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import mmcif_utils
+from flax_model.alphafold3.cpp import string_array
+from flax_model.alphafold3.structure import bioassemblies
+from flax_model.alphafold3.structure import bonds
+from flax_model.alphafold3.structure import chemical_components as struc_chem_comps
+from flax_model.alphafold3.structure import mmcif
+from flax_model.alphafold3.structure import structure
+from flax_model.alphafold3.structure import structure_tables
+import numpy as np
+
+
+ChainIndex: TypeAlias = int
+ResIndex: TypeAlias = int
+AtomName: TypeAlias = str
+BondAtomId: TypeAlias = tuple[ChainIndex, ResIndex, AtomName]
+
+_INSERTION_CODE_REMAP: Mapping[str, str] = {'.': '?'}
+
+
+class NoAtomsError(Exception):
+  """Raise when the mmCIF does not have any atoms."""
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class BondIndices:
+  from_indices: list[int]
+  dest_indices: list[int]
+
+
+@enum.unique
+class ModelID(enum.Enum):
+  """Values for specifying model IDs when parsing."""
+
+  FIRST = 1  # The first model in the file.
+  ALL = 2  # All models in the file.
+
+
+@enum.unique
+class SequenceFormat(enum.Enum):
+  """The possible formats for an input sequence."""
+
+  FASTA = 'fasta'  # One-letter code used in FASTA.
+  CCD_CODES = 'ccd_codes'  # Multiple-letter chemical components dictionary ids.
+  LIGAND_SMILES = 'ligand_smiles'  # SMILES string defining a molecule.
+
+
+def _create_bond_lookup(
+    bonded_atom_pairs: Sequence[tuple[BondAtomId, BondAtomId]],
+) -> Mapping[tuple[ChainIndex, ResIndex], Mapping[AtomName, BondIndices]]:
+  """Creates maps to help find bonds during a loop over residues."""
+  bond_lookup = {}
+  for bond_i, (from_atom_id, dest_atom_id) in enumerate(bonded_atom_pairs):
+    from_chain_i, from_res_i, from_atom_name = from_atom_id
+    dest_chain_i, dest_res_i, dest_atom_name = dest_atom_id
+    bonds_by_from_atom_name = bond_lookup.setdefault(
+        (from_chain_i, from_res_i), {}
+    )
+    bonds_by_dest_atom_name = bond_lookup.setdefault(
+        (dest_chain_i, dest_res_i), {}
+    )
+    bonds_by_from_atom_name.setdefault(
+        from_atom_name, BondIndices(from_indices=[], dest_indices=[])
+    ).from_indices.append(bond_i)
+    bonds_by_dest_atom_name.setdefault(
+        dest_atom_name, BondIndices(from_indices=[], dest_indices=[])
+    ).dest_indices.append(bond_i)
+  return bond_lookup
+
+
+def _get_atom_element(
+    ccd: chemical_components.Ccd, res_name: str, atom_name: str
+) -> str:
+  type_symbol = chemical_components.type_symbol(
+      ccd, res_name=res_name, atom_name=atom_name
+  )
+  return type_symbol or '?'
+
+
+def _get_representative_atom(
+    ccd: chemical_components.Ccd,
+    res_name: str,
+    chain_type: str,
+    sequence_format: SequenceFormat,
+) -> tuple[str, str]:
+  match sequence_format:
+    case SequenceFormat.CCD_CODES:
+      atom_name = _get_first_non_leaving_atom(ccd=ccd, res_name=res_name)
+      atom_element = _get_atom_element(
+          ccd=ccd, res_name=res_name, atom_name=atom_name
+      )
+      return atom_name, atom_element
+    case SequenceFormat.LIGAND_SMILES:
+      return '', '?'
+    case SequenceFormat.FASTA:
+      if chain_type in mmcif_names.PEPTIDE_CHAIN_TYPES:
+        return 'CA', 'C'
+      if chain_type in mmcif_names.NUCLEIC_ACID_CHAIN_TYPES:
+        return "C1'", 'C'
+      else:
+        raise ValueError(chain_type)
+    case _:
+      raise ValueError(sequence_format)
+
+
+@functools.lru_cache(maxsize=128)
+def _get_first_non_leaving_atom(
+    ccd: chemical_components.Ccd, res_name: str
+) -> str:
+  """Returns first definitely non-leaving atom if exists, as a stand-in."""
+  all_atoms = struc_chem_comps.get_all_atoms_in_entry(ccd, res_name=res_name)[
+      '_chem_comp_atom.atom_id'
+  ]
+  representative_atom = all_atoms[0]
+  if representative_atom == 'O1' and len(all_atoms) > 1:
+    representative_atom = all_atoms[1]
+  return representative_atom
+
+
+def _add_ligand_to_chem_comp(
+    chem_comp: MutableMapping[str, struc_chem_comps.ChemCompEntry],
+    ligand_id: str,
+    ligand_smiles: str,
+):
+  """Adds a ligand to chemical components. Raises ValueError on mismatch."""
+  new_entry = struc_chem_comps.ChemCompEntry(
+      type='non-polymer', pdbx_smiles=ligand_smiles
+  )
+
+  existing_entry = chem_comp.get(ligand_id)
+  if existing_entry is None:
+    chem_comp[ligand_id] = new_entry
+  elif existing_entry != new_entry:
+    raise ValueError(
+        f'Mismatching data for ligand {ligand_id}: '
+        f'{new_entry} != {existing_entry}'
+    )
+
+
+def _get_first_model_id(cif: mmcif.Mmcif) -> str:
+  """Returns cheaply the first model ID from the mmCIF."""
+  return cif.get_array(
+      '_atom_site.pdbx_PDB_model_num', dtype=object, gather=slice(1)
+  )[0]
+
+
+def _get_str_model_id(
+    cif: mmcif.Mmcif,
+    model_id: ModelID | int,
+) -> str:
+  """Converts a user-specified model_id argument into a string."""
+  match model_id:
+    case int():
+      str_model_id = str(model_id)
+    case enum.Enum():
+      # We compare the enum's value attribute since regular enum comparison
+      # breaks when adhoc importing.
+      match model_id.value:
+        case ModelID.FIRST.value:
+          try:
+            str_model_id = _get_first_model_id(cif)
+          except IndexError as e:
+            raise NoAtomsError(
+                'The mmCIF does not have any atoms or'
+                ' _atom_site.pdbx_PDB_model_num is missing.'
+            ) from e
+        case ModelID.ALL.value:
+          str_model_id = ''
+        case _:
+          raise ValueError(
+              f'Model ID {model_id} with value {model_id.value} not recognized.'
+          )
+    case _:
+      raise ValueError(
+          f'Model ID {model_id} with type {type(model_id)} not recognized.'
+      )
+  return str_model_id
+
+
+def _parse_bonds(
+    cif: mmcif.Mmcif,
+    atom_key: np.ndarray,
+    model_id: str,
+) -> bonds.Bonds:
+  """Returns the bonds table extracted from the mmCIF.
+
+  Args:
+    cif: The raw mmCIF to extract the bond information from.
+    atom_key: A numpy array defining atom key for each atom in _atom_site. Note
+      that the atom key must be computed before resolving alt-locs since this
+      function operates on the raw mmCIF!
+    model_id: The ID of the model to get bonds for.
+  """
+  if '_struct_conn.id' not in cif:
+    # This is the category key item for the _struct_conn table, therefore
+    # we use it to determine whether to parse bond info.
+    return bonds.Bonds.make_empty()
+  from_atom, dest_atom = mmcif.get_bond_atom_indices(cif, model_id)
+  from_atom = np.array(from_atom, dtype=np.int64)
+  dest_atom = np.array(dest_atom, dtype=np.int64)
+  num_bonds = from_atom.shape[0]
+  bond_key = np.arange(num_bonds, dtype=np.int64)
+  bond_type = cif.get_array('_struct_conn.conn_type_id', dtype=object)
+  if '_struct_conn.pdbx_role' in cif:  # This column isn't always present.
+    bond_role = cif.get_array('_struct_conn.pdbx_role', dtype=object)
+  else:
+    bond_role = np.full((num_bonds,), '?', dtype=object)
+
+  bonds_mask = np.ones((num_bonds,), dtype=bool)
+  # Symmetries other than 1_555 imply the atom is not part of the asymmetric
+  # unit, and therefore this is a bond that only exists in the expanded
+  # bioassembly.
+  # We do not currently support parsing these types of bonds.
+  if '_struct_conn.ptnr1_symmetry' in cif:
+    ptnr1_symmetry = cif.get_array('_struct_conn.ptnr1_symmetry', dtype=object)
+    np.logical_and(bonds_mask, ptnr1_symmetry == '1_555', out=bonds_mask)
+  if '_struct_conn.ptnr2_symmetry' in cif:
+    ptnr2_symmetry = cif.get_array('_struct_conn.ptnr2_symmetry', dtype=object)
+    np.logical_and(bonds_mask, ptnr2_symmetry == '1_555', out=bonds_mask)
+  # Remove bonds that involve atoms that are not part of the structure,
+  # e.g. waters if include_water=False. In a rare case this also removes invalid
+  # bonds that are indicated by a key that is set to _atom_site size.
+  np.logical_and(bonds_mask, np.isin(from_atom, atom_key), out=bonds_mask)
+  np.logical_and(bonds_mask, np.isin(dest_atom, atom_key), out=bonds_mask)
+  return bonds.Bonds(
+      key=bond_key[bonds_mask],
+      type=bond_type[bonds_mask],
+      role=bond_role[bonds_mask],
+      from_atom_key=from_atom[bonds_mask],
+      dest_atom_key=dest_atom[bonds_mask],
+  )
+
+
+@dataclasses.dataclass(frozen=True, slots=True)
+class _MmcifHeader:
+  name: str
+  resolution: float | None
+  release_date: datetime.date | None
+  structure_method: str | None
+  bioassembly_data: bioassemblies.BioassemblyData | None
+  chemical_components_data: struc_chem_comps.ChemicalComponentsData | None
+
+
+def _get_mmcif_header(
+    cif: mmcif.Mmcif,
+    fix_mse: bool,
+    fix_unknown_dna: bool,
+) -> _MmcifHeader:
+  """Extract header fields from an mmCIF object."""
+  entry_id = cif.get('_entry.id')
+  name = entry_id[0] if entry_id else cif.get_data_name()
+  resolution = mmcif.get_resolution(cif)
+
+  release_date = mmcif.get_release_date(cif)
+  if release_date is not None:
+    release_date = datetime.date.fromisoformat(release_date)
+
+  experiments = cif.get('_exptl.method')
+  structure_method = ','.join(experiments) if experiments else None
+
+  try:
+    bioassembly_data = bioassemblies.BioassemblyData.from_mmcif(cif)
+  except bioassemblies.MissingBioassemblyDataError:
+    bioassembly_data = None
+
+  try:
+    chemical_components_data = (
+        struc_chem_comps.ChemicalComponentsData.from_mmcif(
+            cif, fix_mse=fix_mse, fix_unknown_dna=fix_unknown_dna
+        )
+    )
+  except struc_chem_comps.MissingChemicalComponentsDataError:
+    chemical_components_data = None
+
+  return _MmcifHeader(
+      name=name,
+      resolution=resolution,
+      release_date=release_date,
+      structure_method=structure_method,
+      bioassembly_data=bioassembly_data,
+      chemical_components_data=chemical_components_data,
+  )
+
+
+def from_parsed_mmcif(
+    mmcif_object: mmcif.Mmcif,
+    *,
+    name: str | None = None,
+    fix_mse_residues: bool = False,
+    fix_arginines: bool = False,
+    fix_unknown_dna: bool = False,
+    include_water: bool = False,
+    include_other: bool = False,
+    include_bonds: bool = False,
+    model_id: int | ModelID = ModelID.FIRST,
+) -> structure.Structure:
+  """Construct a Structure from a parsed mmCIF object.
+
+  This function is called by `from_mmcif` but can be useful when an mmCIF has
+  already been parsed e.g. to extract extra information from the header before
+  then converting to Structure for further manipulation.
+
+  Args:
+    mmcif_object: A parsed mmcif.Mmcif object.
+    name: Optional name for the structure. If not provided, the name will be
+      taken from the mmCIF data_ field.
+    fix_mse_residues: If True, selenium atom sites (SE) in selenomethionine
+      (MSE) residues will be changed to sulphur atom sites (SD). This is because
+      methionine (MET) residues are often replaced with MSE to aid X-Ray
+      crystallography. If False, the SE MSE atom sites won't be modified.
+    fix_arginines: If True, NH1 and NH2 in arginine will be swapped if needed so
+      that NH1 is always closer to CD than NH2. If False, no atom sites in
+      arginine will be touched. Note that HH11, HH12, HH21, HH22 are fixed too.
+    fix_unknown_dna: If True, residues with name N in DNA chains will have their
+      res_name replaced with DN. Atoms are not changed.
+    include_water: If True, water (HOH) molecules will be parsed. Water
+      molecules may be grouped into chains, where number of residues > 1. Water
+      molecules are usually grouped into chains but do not necessarily all share
+      the same chain ID.
+    include_other: If True, all other atoms that are not included by any of the
+      above parameters will be included. This covers e.g. "polypeptide(D)" and
+      "macrolide" entities, as well as all other non-standard types.
+    include_bonds: If True, bond information will be parsed from the mmCIF and
+      stored in the Structure.
+    model_id: Either the integer model ID to parse, or one of ModelID.FIRST to
+      parse the first model, or ModelID.ALL to parse all models.
+
+  Returns:
+    A Structure representation of the mmCIF object.
+  """
+  str_model_id = _get_str_model_id(cif=mmcif_object, model_id=model_id)
+  header = _get_mmcif_header(
+      mmcif_object, fix_mse=fix_mse_residues, fix_unknown_dna=fix_unknown_dna
+  )
+
+  chains, residues, atoms = get_tables(
+      cif=mmcif_object,
+      fix_mse_residues=fix_mse_residues,
+      fix_arginines=fix_arginines,
+      fix_unknown_dna=fix_unknown_dna,
+      include_water=include_water,
+      include_other=include_other,
+      model_id=str_model_id,
+  )
+
+  if include_bonds and atoms.size > 0:
+    # NB: parsing the atom table before the bonds table allows for a more
+    # informative error message when dealing with bad multi-model mmCIFs.
+    # Also always use a specific model ID, even when parsing all models.
+    if str_model_id == '':  # pylint: disable=g-explicit-bool-comparison
+      bonds_model_id = _get_first_model_id(mmcif_object)
+    else:
+      bonds_model_id = str_model_id
+
+    bonds_table = _parse_bonds(
+        mmcif_object,
+        atom_key=atoms.key,
+        model_id=bonds_model_id,
+    )
+  else:
+    bonds_table = bonds.Bonds.make_empty()
+
+  return structure.Structure(
+      name=name if name is not None else header.name,
+      resolution=header.resolution,
+      release_date=header.release_date,
+      structure_method=header.structure_method,
+      bioassembly_data=header.bioassembly_data,
+      chemical_components_data=header.chemical_components_data,
+      bonds=bonds_table,
+      chains=chains,
+      residues=residues,
+      atoms=atoms,
+  )
+
+
+def from_mmcif(
+    mmcif_string: str | bytes,
+    *,
+    name: str | None = None,
+    fix_mse_residues: bool = False,
+    fix_arginines: bool = False,
+    fix_unknown_dna: bool = False,
+    include_water: bool = False,
+    include_other: bool = False,
+    include_bonds: bool = False,
+    model_id: int | ModelID = ModelID.FIRST,
+) -> structure.Structure:
+  """Construct a Structure from a mmCIF string.
+
+  Args:
+    mmcif_string: The string contents of an mmCIF file.
+    name: Optional name for the structure. If not provided, the name will be
+      taken from the mmCIF data_ field.
+    fix_mse_residues: If True, selenium atom sites (SE) in selenomethionine
+      (MSE) residues will be changed to sulphur atom sites (SD). This is because
+      methionine (MET) residues are often replaced with MSE to aid X-Ray
+      crystallography. If False, the SE MSE atom sites won't be modified.
+    fix_arginines: If True, NH1 and NH2 in arginine will be swapped if needed so
+      that NH1 is always closer to CD than NH2. If False, no atom sites in
+      arginine will be touched. Note that HH11, HH12, HH21, HH22 are fixed too.
+    fix_unknown_dna: If True, residues with name N in DNA chains will have their
+      res_name replaced with DN. Atoms are not changed.
+    include_water: If True, water (HOH) molecules will be parsed. Water
+      molecules may be grouped into chains, where number of residues > 1. Water
+      molecules are usually grouped into chains but do not necessarily all share
+      the same chain ID.
+    include_other: If True, all other atoms that are not included by any of the
+      above parameters will be included. This covers e.g. "polypeptide(D)" and
+      "macrolide" entities, as well as all other non-standard types.
+    include_bonds: If True, bond information will be parsed from the mmCIF and
+      stored in the Structure.
+    model_id: Either the integer model ID to parse, or one of ModelID.FIRST to
+      parse the first model, or ModelID.ALL to parse all models.
+
+  Returns:
+    A Structure representation of the mmCIF string.
+  """
+  mmcif_object = mmcif.from_string(mmcif_string)
+
+  return from_parsed_mmcif(
+      mmcif_object,
+      name=name,
+      fix_mse_residues=fix_mse_residues,
+      fix_arginines=fix_arginines,
+      fix_unknown_dna=fix_unknown_dna,
+      include_water=include_water,
+      include_other=include_other,
+      include_bonds=include_bonds,
+      model_id=model_id,
+  )
+
+
+def from_res_arrays(atom_mask: np.ndarray, **kwargs) -> structure.Structure:
+  """Returns Structure created from from arrays with a residue dimension.
+
+  All unset fields are filled with defaults (e.g. 1.0 for occupancy) or
+  unset/unknown values (e.g. UNK for residue type, or '.' for atom element).
+
+  Args:
+    atom_mask: A array with shape (num_res, num_atom). This is used to decide
+      which atoms in the atom dimension are present in a given residue. Present
+      atoms should have a nonzero value, e.g. 1.0 or True.
+    **kwargs: A mapping from field name to values. For all array-valued fields
+      these arrays must have a dimension of length num_res. Chain and residue
+      fields should have this as their only dimension and atom fields should be
+      shaped (num_res, num_atom). Coordinate fields may also have arbitrary
+      leading dimensions (they must be the same across all coordinate fields).
+      See structure.{CHAIN,RESIDUE,ATOM}_FIELDS for a list of allowed fields.
+  """
+  num_res, num_atom = atom_mask.shape
+  included_indices = np.flatnonzero(atom_mask)
+
+  array_fields = (
+      structure.CHAIN_FIELDS.keys()
+      | structure.RESIDUE_FIELDS.keys()
+      | structure.ATOM_FIELDS.keys()
+  )
+  initializer_kwargs = {}
+  fields = {}
+  for k, val in kwargs.items():
+    if k not in array_fields:
+      # The kwarg key isn't an array field name. Such kwargs are forwarded as-is
+      # to the constructor. They are expected to be global fields (e.g. name).
+      # Other values will raise an error when the constructor is called.
+      if k in structure.TABLE_FIELDS:
+        raise ValueError(f'Table fields must not be set. Got {k}.')
+      initializer_kwargs[k] = val
+      continue
+    elif val is None:
+      raise ValueError(f'{k} must be non-None.')
+
+    if not isinstance(val, np.ndarray):
+      raise TypeError(f'Value for {k} must be a NumPy array. Got {type(val)}.')
+    if k in structure.CHAIN_FIELDS or k in structure.RESIDUE_FIELDS:
+      if val.shape != (num_res,):
+        raise ValueError(
+            f'{k} must have shape ({num_res=},). Got {val.shape=}.'
+        )
+      # Do not reshape the chain/residue arrays, they have the shape we need.
+      fields[k] = val
+    else:
+      assert k in structure.ATOM_FIELDS
+      if val.shape[-2:] != (num_res, num_atom):
+        raise ValueError(
+            f'{k} must have final two dimensions of length '
+            f'{(num_res, num_atom)=}. Got {val.shape=}.'
+        )
+      leading_dims = val.shape[:-2]
+      flat_val = val.reshape(leading_dims + (-1,), order='C')
+      masked_val = flat_val[..., included_indices]
+      fields[k] = masked_val
+
+  # Get chain IDs or assume this is a single-chain structure.
+  chain_id = kwargs.get('chain_id', np.array(['A'] * num_res, dtype=object))
+  # Find chain starts in res-sized arrays, use these to make chain-sized arrays.
+  chain_start = np.concatenate(
+      ([0], np.where(chain_id[1:] != chain_id[:-1])[0] + 1)
+  )
+  if len(set(chain_id)) != len(chain_start):
+    raise ValueError(f'Chain IDs must be contiguous, but got {chain_id}')
+
+  chain_lengths = np.diff(chain_start, append=len(chain_id))
+  chain_key = np.repeat(np.arange(len(chain_start)), chain_lengths)
+
+  chain_entity_id = fields.get('chain_entity_id')
+  if chain_entity_id is not None:
+    entity_id = chain_entity_id[chain_start]
+  else:
+    entity_id = np.array(
+        [str(mmcif.str_id_to_int_id(cid)) for cid in chain_id[chain_start]],
+        dtype=object,
+    )
+  chain_str_empty = np.full((num_res,), '.', dtype=object)
+  chains_table = structure_tables.Chains(
+      key=chain_key[chain_start],
+      id=chain_id[chain_start],
+      type=fields.get('chain_type', chain_str_empty)[chain_start],
+      auth_asym_id=fields.get('chain_auth_asym_id', chain_id)[chain_start],
+      entity_id=entity_id,
+      entity_desc=fields.get('chain_entity_desc', chain_str_empty)[chain_start],
+  )
+
+  # Since all arrays are residue-shaped, we can use them directly.
+  res_key = np.arange(num_res, dtype=np.int64)
+  res_id = fields.get('res_id', res_key + 1).astype(np.int32)
+  residues_table = structure_tables.Residues(
+      key=res_key,
+      chain_key=chain_key,
+      id=res_id,
+      name=fields.get('res_name', np.full(num_res, 'UNK', dtype=object)),
+      auth_seq_id=fields.get(
+          'res_auth_seq_id', np.char.mod('%d', res_id).astype(object)
+      ),
+      insertion_code=fields.get(
+          'res_insertion_code', np.full(num_res, '?', dtype=object)
+      ),
+  )
+
+  # The atom-sized arrays have already been masked and reshaped.
+  num_atoms_per_res = np.sum(atom_mask, axis=1, dtype=np.int32)
+  num_atoms_total = np.sum(num_atoms_per_res, dtype=np.int32)
+  # Structure is immutable, so use the same array multiple times to save RAM.
+  atom_str_empty = np.full(num_atoms_total, '.', dtype=object)
+  atom_float32_zeros = np.zeros(num_atoms_total, dtype=np.float32)
+  atom_float32_ones = np.ones(num_atoms_total, dtype=np.float32)
+  atoms_table = structure_tables.Atoms(
+      key=np.arange(num_atoms_total, dtype=np.int64),
+      chain_key=np.repeat(chain_key, num_atoms_per_res),
+      res_key=np.repeat(res_key, num_atoms_per_res),
+      name=fields.get('atom_name', atom_str_empty),
+      element=fields.get('atom_element', atom_str_empty),
+      x=fields.get('atom_x', atom_float32_zeros),
+      y=fields.get('atom_y', atom_float32_zeros),
+      z=fields.get('atom_z', atom_float32_zeros),
+      b_factor=fields.get('atom_b_factor', atom_float32_zeros),
+      occupancy=fields.get('atom_occupancy', atom_float32_ones),
+  )
+
+  return structure.Structure(
+      chains=chains_table,
+      residues=residues_table,
+      atoms=atoms_table,
+      bonds=structure_tables.Bonds.make_empty(),  # Currently not set.
+      **initializer_kwargs,
+  )
+
+
+def expand_sequence(
+    sequence: str, chain_type: str, sequence_format: SequenceFormat
+) -> Sequence[str]:
+  """Returns full residue names based on a sequence string.
+
+  Args:
+    sequence: A string representing the sequence.
+    chain_type: The chain type of the sequence.
+    sequence_format: The format of the sequence argument.
+  """
+  match sequence_format:
+    case SequenceFormat.FASTA:
+      if not all(c.isalpha() for c in sequence):
+        raise ValueError(f'Sequence "{sequence}" has non-alphabetic characters')
+      match chain_type:
+        case mmcif_names.PROTEIN_CHAIN:
+          res_name_map = residue_names.PROTEIN_COMMON_ONE_TO_THREE
+          default_res_name = residue_names.UNK
+        case mmcif_names.RNA_CHAIN:
+          res_name_map = {r: r for r in residue_names.RNA_TYPES}
+          default_res_name = residue_names.UNK_RNA
+        case mmcif_names.DNA_CHAIN:
+          res_name_map = residue_names.DNA_COMMON_ONE_TO_TWO
+          default_res_name = residue_names.UNK_DNA
+        case _:
+          raise ValueError(f'{chain_type=} not supported for FASTA format.')
+      return [
+          res_name_map.get(one_letter_res, default_res_name)
+          for one_letter_res in sequence
+      ]
+    case SequenceFormat.CCD_CODES:
+      return sequence.strip('()').split(')(')
+    case SequenceFormat.LIGAND_SMILES:
+      ligand_id, _ = sequence.split(':', maxsplit=1)
+      return [ligand_id]
+
+
+def from_sequences_and_bonds(
+    *,
+    sequences: Sequence[str],
+    chain_types: Sequence[str],
+    sequence_formats: Sequence[SequenceFormat],
+    bonded_atom_pairs: Sequence[tuple[BondAtomId, BondAtomId]] | None,
+    ccd: chemical_components.Ccd,
+    chain_ids: Sequence[str] | None = None,
+    name: str = 'from_sequences_and_bonds',
+    bond_type: str | None = None,
+    **constructor_args,
+) -> structure.Structure:
+  """Returns a minimal structure for the input sequences and bonds.
+
+  The returned structure will have at least one atom per residue. If the
+  residue has any bonded atoms, according to `bonded_atom_pairs`, then
+  all (and only) those atoms will be present for that residue. If the residue
+  is not involved in any bond then an arbitrary atom will be created.
+
+  Args:
+    sequences: A sequence of strings, each one representing a single chain.
+    chain_types: The types of each chain, e.g. polypeptide(L). The n-th element
+      describes the n-th sequence in `sequences`.
+    sequence_formats: The format of each sequence. The n-th element describes
+      the n-th sequence in `sequences`.
+    bonded_atom_pairs: A sequence of bonded atom pairs. Each atom is described
+      as a tuple of (chain_index, res_index, atom_name), where the first two
+      values are 0-based indices. The chain_index is the index of the chain in
+      the `sequences` argument, and the res_index is the index of the residue in
+      that sequence. The atom_name is the name of the atom in the residue, e.g.
+      CA. If the atom is not found in the standard atoms for that residue
+      (according to the CCD) then an error is raised.
+    ccd: The chemical components dictionary.
+    chain_ids: A sequence of chain IDs, one for each chain in `sequences`. If
+      not provided, then the chain IDs will be generated automatically based on
+      sequence indices.
+    name: A name for the returned structure.
+    bond_type: This type will be used for all bonds in the structure, where type
+      follows PDB scheme, e.g. unknown (?), hydrog, metalc, covale, disulf.
+    **constructor_args: These arguments are passed directly to the
+      structure.Structure constructor.
+  """
+  chain_id = []
+  chain_type = []
+  chain_res_count = []
+  res_id = []
+  res_name = []
+  res_atom_count = []
+  atom_name = []
+  atom_element = []
+  chem_comp = {}
+
+  num_bonds = len(bonded_atom_pairs or ())
+  from_atom_key = np.full((num_bonds,), -1, dtype=np.int64)
+  dest_atom_key = np.full((num_bonds,), -1, dtype=np.int64)
+
+  # Create map (chain_i, res_i) -> {atom_name -> (from_idxs dest_idxs)}.
+  # This allows quick lookup of whether a residue has any bonded atoms, and
+  # which bonds those atoms participate in.
+  bond_lookup = _create_bond_lookup(bonded_atom_pairs or ())
+
+  current_atom_key = 0
+  for chain_i, (sequence, curr_chain_type, sequence_format) in enumerate(
+      zip(sequences, chain_types, sequence_formats, strict=True)
+  ):
+    if chain_ids is not None:
+      current_chain_id = chain_ids[chain_i]
+    else:
+      current_chain_id = mmcif.int_id_to_str_id(chain_i + 1)
+    num_chain_residues = 0
+    for res_i, full_res_name in enumerate(
+        expand_sequence(sequence, curr_chain_type, sequence_format)
+    ):
+      current_res_id = res_i + 1
+      num_res_atoms = 0
+
+      # Look for bonded atoms in the bond lookup and if any are found, add
+      # their atom keys to the bond atom_key columns.
+      if bond_indices_by_atom_name := bond_lookup.get((chain_i, res_i)):
+        comp_atoms = None
+        if sequence_format != SequenceFormat.LIGAND_SMILES:
+          comp_atoms = set(ccd.get(full_res_name)['_chem_comp_atom.atom_id'])
+        for bond_atom_name, bond_indices in bond_indices_by_atom_name.items():
+          if comp_atoms is not None and bond_atom_name not in comp_atoms:
+            raise ValueError(
+                f'Bonded atom "{bond_atom_name}" was not found in the list of'
+                f' atoms of the chemical component {full_res_name}. Valid atom'
+                f' names for {full_res_name} are: {sorted(comp_atoms)}.'
+                ' This is likely caused by an invalid atom name in the bonded'
+                f' atom (chain_id={current_chain_id}, res_id={current_res_id},'
+                f' atom_name={bond_atom_name}) specified in `bondedAtomPairs`'
+                ' in the input JSON.'
+            )
+          atom_name.append(bond_atom_name)
+          atom_element.append(
+              _get_atom_element(
+                  ccd=ccd, res_name=full_res_name, atom_name=bond_atom_name
+              )
+          )
+          for from_bond_i in bond_indices.from_indices:
+            from_atom_key[from_bond_i] = current_atom_key
+          for dest_bond_i in bond_indices.dest_indices:
+            dest_atom_key[dest_bond_i] = current_atom_key
+          current_atom_key += 1
+          num_res_atoms += 1
+      else:
+        # If this residue has no bonded atoms then we need to add one atom
+        # like in from_sequences.
+        assert num_res_atoms == 0
+        rep_atom_name, rep_atom_element = _get_representative_atom(
+            ccd=ccd,
+            res_name=full_res_name,
+            chain_type=curr_chain_type,
+            sequence_format=sequence_format,
+        )
+        atom_name.append(rep_atom_name)
+        atom_element.append(rep_atom_element)
+        num_res_atoms += 1
+        current_atom_key += 1
+
+      if sequence_format == SequenceFormat.LIGAND_SMILES:
+        # Sequence expect to be in the format <ligand_id>:<ligand_smiles>,
+        # which always corresponds to a single-residue chain.
+        ligand_id, ligand_smiles = sequence.split(':', maxsplit=1)
+        if ccd.get(ligand_id) is not None:
+          raise ValueError(
+              f'Ligand name {ligand_id} is in CCD - it is not supported to give'
+              ' ligands created from SMILES the same name as CCD components.'
+          )
+        # We need to provide additional chemical components metadata for
+        # ligands specified via SMILES strings since they might not be in CCD.
+        _add_ligand_to_chem_comp(chem_comp, ligand_id, ligand_smiles)
+
+      assert num_res_atoms >= 1
+      res_atom_count.append(num_res_atoms)
+      num_chain_residues += 1
+      res_id.append(current_res_id)
+      res_name.append(full_res_name)
+
+    chain_id.append(current_chain_id)
+    chain_type.append(curr_chain_type)
+    chain_res_count.append(num_chain_residues)
+
+  chem_comp_data = struc_chem_comps.ChemicalComponentsData(chem_comp)
+  chem_comp_data = struc_chem_comps.populate_missing_ccd_data(
+      ccd=ccd,
+      chemical_components_data=chem_comp_data,
+      chemical_component_ids=set(res_name),
+  )
+
+  if bonded_atom_pairs is not None:
+    unknown_bond_col = np.full((num_bonds,), '?', dtype=object)
+    if bond_type is None:
+      bond_type_col = unknown_bond_col
+    else:
+      bond_type_col = np.full((num_bonds,), bond_type, dtype=object)
+    bonds_table = bonds.Bonds(
+        key=np.arange(num_bonds, dtype=np.int64),
+        type=bond_type_col,
+        role=unknown_bond_col,
+        from_atom_key=from_atom_key,
+        dest_atom_key=dest_atom_key,
+    )
+  else:
+    bonds_table = structure_tables.Bonds.make_empty()
+
+  chain_key = np.arange(len(sequences), dtype=np.int64)  # 1 chain per sequence.
+  chain_id = np.array(chain_id, dtype=object)
+  chains_table = structure_tables.Chains(
+      key=chain_key,
+      id=chain_id,
+      type=np.array(chain_type, dtype=object),
+      auth_asym_id=chain_id,
+      entity_id=np.char.mod('%d', chain_key + 1).astype(object),
+      entity_desc=np.array(['.'] * len(chain_key), dtype=object),
+  )
+
+  res_key = np.arange(len(res_name), dtype=np.int64)
+  res_chain_key = np.repeat(chain_key, chain_res_count)
+  residues_table = structure_tables.Residues(
+      key=res_key,
+      chain_key=res_chain_key,
+      id=np.array(res_id, dtype=np.int32),
+      name=np.array(res_name, dtype=object),
+      auth_seq_id=np.char.mod('%d', res_id).astype(object),
+      insertion_code=np.full(len(res_name), '?', dtype=object),
+  )
+
+  num_atoms = current_atom_key
+  atom_float32_zeros = np.zeros(num_atoms, dtype=np.float32)
+  atoms_table = structure_tables.Atoms(
+      key=np.arange(num_atoms, dtype=np.int64),
+      chain_key=np.repeat(res_chain_key, res_atom_count),
+      res_key=np.repeat(res_key, res_atom_count),
+      name=np.array(atom_name, dtype=object),
+      element=np.array(atom_element, dtype=object),
+      x=atom_float32_zeros,
+      y=atom_float32_zeros,
+      z=atom_float32_zeros,
+      b_factor=atom_float32_zeros,
+      occupancy=np.ones(num_atoms, np.float32),
+  )
+
+  return structure.Structure(
+      name=name,
+      atoms=atoms_table,
+      residues=residues_table,
+      chains=chains_table,
+      bonds=bonds_table,
+      chemical_components_data=chem_comp_data,
+      **constructor_args,
+  )
+
+
+class _ChainResBuilder:
+  """Class for incrementally building chain and residue tables."""
+
+  def __init__(
+      self,
+      *,
+      chain_key_by_chain_id: Mapping[str, int],
+      entity_id_by_chain_id: Mapping[str, str],
+      chain_type_by_entity_id: Mapping[str, str],
+      entity_desc_by_entity_id: Mapping[str, str],
+      fix_mse_residues: bool,
+      fix_unknown_dna: bool,
+  ):
+    # Len: num_chains.
+    self.chain_key = []
+    self.chain_id = []
+    self.chain_type = []
+    self.chain_auth_asym_id = []
+    self.chain_entity_id = []
+    self.chain_entity_desc = []
+
+    # Len: num_residues.
+    self.res_key = []
+    self.res_chain_key = []
+    self.res_id = []
+    self.res_name = []
+    self.res_auth_seq_id = []
+    self.res_insertion_code = []
+
+    self.chain_key_by_chain_id = chain_key_by_chain_id
+    self.entity_id_by_chain_id = entity_id_by_chain_id
+    self.chain_type_by_entity_id = chain_type_by_entity_id
+    self.entity_desc_by_entity_id = entity_desc_by_entity_id
+    self.key_for_res: dict[tuple[str, str, str, str], int] = {}
+
+    self._fix_mse_residues = fix_mse_residues
+    self._fix_unknown_dna = fix_unknown_dna
+
+  def add_residues(
+      self,
+      *,
+      chain_ids: np.ndarray,
+      chain_auth_asym_ids: np.ndarray,
+      res_ids: np.ndarray,
+      res_names: np.ndarray,
+      res_auth_seq_ids: np.ndarray,
+      res_ins_codes: np.ndarray,
+  ):
+    """Adds a residue (and its chain) to the tables."""
+    # Create chain table data.
+    if chain_ids.size == 0:
+      return
+
+    chain_ids_with_prev = np.concatenate(
+        (([self.chain_id[-1] if self.chain_id else None], chain_ids))
+    )
+    chain_change_mask = chain_ids_with_prev[:-1] != chain_ids_with_prev[1:]
+    chain_change_ids = chain_ids[chain_change_mask]
+    chain_keys = string_array.remap(
+        chain_change_ids, self.chain_key_by_chain_id, inplace=False
+    )
+    self.chain_key.extend(chain_keys)
+    self.chain_id.extend(chain_change_ids)
+    self.chain_auth_asym_id.extend(chain_auth_asym_ids[chain_change_mask])
+    chain_entity_id = string_array.remap(
+        chain_change_ids, self.entity_id_by_chain_id, inplace=False
+    )
+    self.chain_entity_id.extend(chain_entity_id)
+    chain_type = string_array.remap(
+        chain_entity_id, self.chain_type_by_entity_id, inplace=False
+    )
+    self.chain_type.extend(chain_type)
+    chain_entity_desc = string_array.remap(
+        chain_entity_id, self.entity_desc_by_entity_id, inplace=False
+    )
+    self.chain_entity_desc.extend(chain_entity_desc)
+
+    # Create residue table data.
+    num_prev_res = len(self.res_id)
+    res_keys = np.arange(num_prev_res, num_prev_res + len(res_ids))
+    res_iter = zip(
+        chain_ids,
+        res_auth_seq_ids,
+        res_names,
+        res_ins_codes,
+        strict=True,
+    )
+    key_for_res_update = {
+        res_unique_id: res_key
+        for res_key, res_unique_id in enumerate(res_iter, num_prev_res)
+    }
+    self.key_for_res.update(key_for_res_update)
+    self.res_key.extend(res_keys)
+    self.res_chain_key.extend(
+        string_array.remap(chain_ids, self.chain_key_by_chain_id, inplace=False)
+    )
+    self.res_id.extend(res_ids)
+    self.res_name.extend(res_names)
+    self.res_auth_seq_id.extend(res_auth_seq_ids)
+    self.res_insertion_code.extend(res_ins_codes)
+
+  def make_chains_table(self) -> structure_tables.Chains:
+    """Returns the Structure chains table."""
+    chain_key = np.array(self.chain_key, dtype=np.int64)
+    if not np.all(chain_key[:-1] <= chain_key[1:]):
+      # If the order is inconsistent with the atoms table, sort so that it is.
+      order = np.argsort(self.chain_key, kind='stable')
+      return structure_tables.Chains(
+          key=chain_key[order],
+          id=np.array(self.chain_id, dtype=object)[order],
+          type=np.array(self.chain_type, dtype=object)[order],
+          auth_asym_id=np.array(self.chain_auth_asym_id, dtype=object)[order],
+          entity_id=np.array(self.chain_entity_id, dtype=object)[order],
+          entity_desc=np.array(self.chain_entity_desc, dtype=object)[order],
+      )
+    return structure_tables.Chains(
+        key=chain_key,
+        id=np.array(self.chain_id, dtype=object),
+        type=np.array(self.chain_type, dtype=object),
+        auth_asym_id=np.array(self.chain_auth_asym_id, dtype=object),
+        entity_id=np.array(self.chain_entity_id, dtype=object),
+        entity_desc=np.array(self.chain_entity_desc, dtype=object),
+    )
+
+  def make_residues_table(self) -> structure_tables.Residues:
+    """Returns the Structure residues table."""
+    res_name = np.array(self.res_name, dtype=object)
+    res_chain_key = np.array(self.res_chain_key, dtype=np.int64)
+
+    if self._fix_mse_residues:
+      string_array.remap(res_name, mapping={'MSE': 'MET'}, inplace=True)
+
+    if self._fix_unknown_dna:
+      # Remap residues from N -> DN in DNA chains only.
+      dna_chain_mask = (
+          np.array(self.chain_type, dtype=object) == mmcif_names.DNA_CHAIN
+      )
+      dna_chain_key = np.array(self.chain_key, dtype=object)[dna_chain_mask]
+      res_name[(res_name == 'N') & np.isin(res_chain_key, dna_chain_key)] = 'DN'
+
+    if not np.all(res_chain_key[:-1] <= res_chain_key[1:]):
+      # If the order is inconsistent with the atoms table, sort so that it is.
+      order = np.argsort(res_chain_key, kind='stable')
+      return structure_tables.Residues(
+          key=np.array(self.res_key, dtype=np.int64)[order],
+          chain_key=res_chain_key[order],
+          id=np.array(self.res_id, dtype=np.int32)[order],
+          name=res_name[order],
+          auth_seq_id=np.array(self.res_auth_seq_id, dtype=object)[order],
+          insertion_code=np.array(self.res_insertion_code, dtype=object)[order],
+      )
+    return structure_tables.Residues(
+        key=np.array(self.res_key, dtype=np.int64),
+        chain_key=res_chain_key,
+        id=np.array(self.res_id, dtype=np.int32),
+        name=res_name,
+        auth_seq_id=np.array(self.res_auth_seq_id, dtype=object),
+        insertion_code=np.array(self.res_insertion_code, dtype=object),
+    )
+
+
+def _get_string_array_default(cif: mmcif.Mmcif, key: str, default: list[str]):
+  try:
+    return cif.get_array(key, dtype=object)
+  except KeyError:
+    return default
+
+
+def _generate_required_tables_if_missing(
+    cif: mmcif.Mmcif,
+) -> Mapping[str, Sequence[str]]:
+  """Generates all required tables and columns if missing."""
+  update = {}
+
+  atom_site_entities = _get_string_array_default(
+      cif, '_atom_site.label_entity_id', []
+  )
+
+  # OpenMM produces files that don't have any of the tables and also have
+  # _atom_site.label_entity_id set to '?' for all atoms. We infer the entities
+  # based on the _atom_site.label_asym_id column. We start with cheaper O(1)
+  # checks to prevent running the expensive O(n) check on most files.
+  if (
+      len(atom_site_entities) > 0  # pylint: disable=g-explicit-length-test
+      and '_entity.id' not in cif  # Ignore if the _entity table exists.
+      and atom_site_entities[0] == '?'  # Cheap check.
+      and set(atom_site_entities) == {'?'}  # Expensive check.
+  ):
+    label_asym_ids = cif.get_array('_atom_site.label_asym_id', dtype=object)
+    atom_site_entities = [
+        str(mmcif.str_id_to_int_id(cid)) for cid in label_asym_ids
+    ]
+    # Update _atom_site.label_entity_id to be consistent with the new tables.
+    update['_atom_site.label_entity_id'] = atom_site_entities
+
+  # Check table existence by checking the presence of its primary key.
+  if '_struct_asym.id' not in cif:
+    # Infer the _struct_asym table using the _atom_site table.
+    asym_ids = _get_string_array_default(cif, '_atom_site.label_asym_id', [])
+
+    if len(atom_site_entities) == 0 or len(asym_ids) == 0:  # pylint: disable=g-explicit-length-test
+      raise ValueError(
+          'Could not parse an mmCIF with no _struct_asym table and also no '
+          '_atom_site.label_entity_id or _atom_site.label_asym_id columns.'
+      )
+
+    # Deduplicate, but keep the order intact - dict.fromkeys maintains order.
+    entity_id_chain_id_pairs = list(
+        dict.fromkeys(zip(atom_site_entities, asym_ids, strict=True))
+    )
+    update['_struct_asym.entity_id'] = [e for e, _ in entity_id_chain_id_pairs]
+    update['_struct_asym.id'] = [c for _, c in entity_id_chain_id_pairs]
+
+  if '_entity.id' not in cif:
+    # Infer the _entity_poly and _entity tables using the _atom_site table.
+    residues = _get_string_array_default(cif, '_atom_site.label_comp_id', [])
+    group_pdb = _get_string_array_default(cif, '_atom_site.group_PDB', [])
+    if '_atom_site.label_entity_id' in cif:
+      entities = atom_site_entities
+    else:
+      # If _atom_site.label_entity_id not set, use the asym_id -> entity_id map.
+      asym_to_entity = dict(
+          zip(
+              cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True
+          )
+      )
+      entities = string_array.remap(
+          cif.get_array('_atom_site.label_asym_id', dtype=object),
+          mapping=asym_to_entity,
+      )
+
+    entity_ids = []
+    entity_types = []
+    entity_poly_entity_ids = []
+    entity_poly_types = []
+    entity_poly_table_missing = '_entity_poly.entity_id' not in cif
+    for entity_id, group in itertools.groupby(
+        zip(entities, residues, group_pdb, strict=True), key=lambda e: e[0]
+    ):
+      _, entity_residues, entity_group_pdb = zip(*group, strict=True)
+      entity_type = _guess_entity_type(
+          chain_residues=entity_residues, atom_types=entity_group_pdb
+      )
+      entity_ids.append(entity_id)
+      entity_types.append(entity_type)
+
+      if entity_poly_table_missing and entity_type == mmcif_names.POLYMER_CHAIN:
+        polymer_type = mmcif_names.guess_polymer_type(entity_residues)
+        entity_poly_entity_ids.append(entity_id)
+        entity_poly_types.append(polymer_type)
+
+    update['_entity.id'] = entity_ids
+    update['_entity.type'] = entity_types
+    if entity_poly_table_missing:
+      update['_entity_poly.entity_id'] = entity_poly_entity_ids
+      update['_entity_poly.type'] = entity_poly_types
+
+  if '_atom_site.type_symbol' not in cif:
+    update['_atom_site.type_symbol'] = mmcif.get_or_infer_type_symbol(cif)
+
+  return update
+
+
+def _maybe_add_missing_scheme_tables(
+    cif: mmcif.Mmcif,
+    res_starts: Sequence[int],
+    label_asym_ids: np.ndarray,
+    label_seq_ids: np.ndarray,
+    label_comp_ids: np.ndarray,
+    auth_seq_ids: np.ndarray,
+    pdb_ins_codes: np.ndarray,
+) -> Mapping[str, Sequence[str]]:
+  """If missing, infers the scheme tables from the _atom_site table."""
+  update = {}
+
+  required_poly_seq_scheme_cols = (
+      '_pdbx_poly_seq_scheme.asym_id',
+      '_pdbx_poly_seq_scheme.pdb_seq_num',
+      '_pdbx_poly_seq_scheme.pdb_ins_code',
+      '_pdbx_poly_seq_scheme.seq_id',
+      '_pdbx_poly_seq_scheme.mon_id',
+      '_pdbx_poly_seq_scheme.pdb_strand_id',
+  )
+  if not all(col in cif for col in required_poly_seq_scheme_cols):
+    # Create a mask for atoms where each polymer residue start.
+    entity_id_by_chain_id = dict(
+        zip(cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True)
+    )
+    chain_type_by_entity_id = dict(
+        zip(cif['_entity.id'], cif['_entity.type'], strict=True)
+    )
+    # Remap asym ID -> entity ID.
+    label_entity_id = string_array.remap(
+        label_asym_ids, mapping=entity_id_by_chain_id, inplace=False
+    )
+    # Remap entity ID -> chain type.
+    chain_type = string_array.remap(
+        label_entity_id, mapping=chain_type_by_entity_id, inplace=False
+    )
+    res_mask = np.zeros_like(label_seq_ids, dtype=bool)
+    res_mask[res_starts] = True
+    res_mask &= chain_type == mmcif_names.POLYMER_CHAIN
+
+    entity_poly_seq_cols = (
+        '_entity_poly_seq.entity_id',
+        '_entity_poly_seq.num',
+        '_entity_poly_seq.mon_id',
+    )
+    if all(col in cif for col in entity_poly_seq_cols):
+      # Use _entity_poly_seq if available.
+      poly_seq_num = cif.get_array('_entity_poly_seq.num', dtype=object)
+      poly_seq_mon_id = cif.get_array('_entity_poly_seq.mon_id', dtype=object)
+      poly_seq_entity_id = cif.get_array(
+          '_entity_poly_seq.entity_id', dtype=object
+      )
+      # We have to add the entity ID to the residue ID because multiple residues
+      # can share the same ID. This also allows using string_array.remap.
+      label_seq_id_to_auth_seq_id = dict(
+          zip(
+              np.char.add(label_entity_id[res_mask], label_seq_ids[res_mask]),
+              auth_seq_ids[res_mask],
+              strict=True,
+          )
+      )
+      scheme_pdb_seq_num = string_array.remap(
+          np.char.add(poly_seq_entity_id, poly_seq_num),
+          mapping=label_seq_id_to_auth_seq_id,
+          default_value='.',
+      )
+      label_seq_id_to_ins_code = dict(
+          zip(label_seq_ids[res_mask], pdb_ins_codes[res_mask], strict=True)
+      )
+      scheme_pdb_ins_code = string_array.remap(
+          poly_seq_num, mapping=label_seq_id_to_ins_code, default_value='.'
+      )
+
+      # The _entity_poly_seq table is entity-based, while _pdbx_poly_seq_scheme
+      # is chain-based. A single entity could mean multiple chains (asym_ids),
+      # we therefore need to replicate each entity for all of the chains.
+      scheme_asym_id = []
+      select = []
+      indices = np.arange(len(poly_seq_entity_id), dtype=np.int32)
+      for asym_id, entity_id in zip(
+          cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True
+      ):
+        entity_mask = poly_seq_entity_id == entity_id
+        select.extend(indices[entity_mask])
+        scheme_asym_id.extend([asym_id] * sum(entity_mask))
+
+      scheme_pdb_strand_id = string_array.remap(
+          np.array(scheme_asym_id, dtype=object),
+          mapping=mmcif.get_internal_to_author_chain_id_map(cif),
+          inplace=False,
+      )
+
+      update['_pdbx_poly_seq_scheme.asym_id'] = scheme_asym_id
+      update['_pdbx_poly_seq_scheme.pdb_strand_id'] = scheme_pdb_strand_id
+      update['_pdbx_poly_seq_scheme.pdb_seq_num'] = scheme_pdb_seq_num[select]
+      update['_pdbx_poly_seq_scheme.pdb_ins_code'] = scheme_pdb_ins_code[select]
+      update['_pdbx_poly_seq_scheme.seq_id'] = poly_seq_num[select]
+      update['_pdbx_poly_seq_scheme.mon_id'] = poly_seq_mon_id[select]
+    else:
+      # _entity_poly_seq not available, fallback to _atom_site.
+      res_asym_ids = label_asym_ids[res_mask]
+      res_strand_ids = string_array.remap(
+          array=res_asym_ids,
+          mapping=mmcif.get_internal_to_author_chain_id_map(cif),
+          inplace=False,
+      )
+      update['_pdbx_poly_seq_scheme.asym_id'] = res_asym_ids
+      update['_pdbx_poly_seq_scheme.pdb_strand_id'] = res_strand_ids
+      update['_pdbx_poly_seq_scheme.pdb_seq_num'] = auth_seq_ids[res_mask]
+      update['_pdbx_poly_seq_scheme.pdb_ins_code'] = pdb_ins_codes[res_mask]
+      update['_pdbx_poly_seq_scheme.seq_id'] = label_seq_ids[res_mask]
+      update['_pdbx_poly_seq_scheme.mon_id'] = label_comp_ids[res_mask]
+
+  required_nonpoly_scheme_cols = (
+      '_pdbx_nonpoly_scheme.mon_id',
+      '_pdbx_nonpoly_scheme.asym_id',
+      '_pdbx_nonpoly_scheme.pdb_seq_num',
+      '_pdbx_nonpoly_scheme.pdb_ins_code',
+  )
+  required_branch_scheme_cols = (
+      '_pdbx_branch_scheme.mon_id',
+      '_pdbx_branch_scheme.asym_id',
+      '_pdbx_branch_scheme.pdb_seq_num',
+  )
+
+  # Generate _pdbx_nonpoly_scheme only if both tables are missing.
+  if not (
+      all(col in cif for col in required_nonpoly_scheme_cols)
+      or all(col in cif for col in required_branch_scheme_cols)
+  ):
+    # To be strictly semantically correct, multi-residue ligands should be
+    # written in _pdbx_branch_scheme. However, Structure parsing handles
+    # correctly multi-residue ligands in _pdbx_nonpoly_scheme and the tables
+    # constructed here live only while parsing, hence this is unnecessary.
+    entity_id_by_chain_id = dict(
+        zip(cif['_struct_asym.id'], cif['_struct_asym.entity_id'], strict=True)
+    )
+    chain_type_by_entity_id = dict(
+        zip(cif['_entity.id'], cif['_entity.type'], strict=True)
+    )
+    # Remap asym ID -> entity ID.
+    chain_type = string_array.remap(
+        label_asym_ids, mapping=entity_id_by_chain_id, inplace=False
+    )
+    # Remap entity ID -> chain type.
+    string_array.remap(
+        chain_type, mapping=chain_type_by_entity_id, inplace=True
+    )
+    res_mask = np.zeros_like(label_seq_ids, dtype=bool)
+    res_mask[res_starts] = True
+    res_mask &= chain_type != mmcif_names.POLYMER_CHAIN
+
+    if not np.any(res_mask):
+      return update  # Shortcut: no non-polymer residues.
+
+    ins_codes = string_array.remap(
+        pdb_ins_codes[res_mask], mapping={'?': '.'}, inplace=False
+    )
+
+    update['_pdbx_nonpoly_scheme.asym_id'] = label_asym_ids[res_mask]
+    update['_pdbx_nonpoly_scheme.pdb_seq_num'] = auth_seq_ids[res_mask]
+    update['_pdbx_nonpoly_scheme.pdb_ins_code'] = ins_codes
+    update['_pdbx_nonpoly_scheme.mon_id'] = label_comp_ids[res_mask]
+
+  return update
+
+
+def _get_chain_key_by_chain_id(
+    resolved_chain_ids: np.ndarray, struct_asym_chain_ids: np.ndarray
+) -> Mapping[str, int]:
+  """Returns chain key for each chain ID respecting resolved chain ordering."""
+  # Check that all chain IDs found in the (potentially filtered) _atom_site
+  # table are present in the _struct_asym table.
+  unique_resolved_chain_ids = set(resolved_chain_ids)
+  if not unique_resolved_chain_ids.issubset(set(struct_asym_chain_ids)):
+    unique_resolved_chain_ids = sorted(unique_resolved_chain_ids)
+    unique_struct_asym_chain_ids = sorted(set(struct_asym_chain_ids))
+    raise ValueError(
+        'Bad mmCIF: chain IDs in _atom_site.label_asym_id '
+        f'{unique_resolved_chain_ids} is not a subset of chain IDs in '
+        f'_struct_asym.id {unique_struct_asym_chain_ids}.'
+    )
+
+  resolved_mask = string_array.isin(
+      struct_asym_chain_ids, unique_resolved_chain_ids
+  )
+  # For all resolved chains, use the _atom_site order they appear in. E.g.
+  # resolved_chain_ids     = [B A   E D F]
+  # struct_asym_chain_ids  = [A B C D E F]
+  # consistent_chain_order = [B A C E D F]
+  # chain_keys             = [0 1 2 3 4 5]
+  consistent_chain_order = struct_asym_chain_ids.copy()
+  consistent_chain_order[resolved_mask] = resolved_chain_ids
+  return dict(zip(consistent_chain_order, range(len(struct_asym_chain_ids))))
+
+
+def get_tables(
+    cif: mmcif.Mmcif,
+    fix_mse_residues: bool,
+    fix_arginines: bool,
+    fix_unknown_dna: bool,
+    include_water: bool,
+    include_other: bool,
+    model_id: str,
+) -> tuple[
+    structure_tables.Chains, structure_tables.Residues, structure_tables.Atoms
+]:
+  """Returns chain, residue, and atom tables from a parsed mmcif.
+
+  Args:
+    cif: A parsed mmcif.Mmcif.
+    fix_mse_residues: See from_mmcif.
+    fix_arginines: See from_mmcif.
+    fix_unknown_dna: See from_mmcif.
+    include_water: See from_mmcif.
+    include_other: See from_mmcif.
+    model_id: A string defining which model ID to use. If set, only coordinates,
+      b-factors and occupancies for the given model are returned. If empty,
+      coordinates, b-factors and occupanciesall for models are returned with a
+      leading dimension of num_models. Note that the model_id argument in
+      from_mmcif is an integer and has slightly different use (see from_mmcif).
+  """
+  # Add any missing tables and columns we require for parsing.
+  if cif_update := _generate_required_tables_if_missing(cif):
+    cif = cif.copy_and_update(cif_update)
+
+  # Resolve alt-locs, selecting only a single option for each residue. Also
+  # computes the layout, which defines where chain and residue boundaries are.
+  atom_site_all_models, layout = mmcif_utils.filter(
+      cif,
+      include_nucleotides=True,
+      include_ligands=True,
+      include_water=include_water,
+      include_other=include_other,
+      model_id=model_id,
+  )
+  atom_site_first_model = atom_site_all_models[0]
+
+  # Get atom information from the _atom_site table.
+  def _first_model_string_array(col: str) -> np.ndarray:
+    return cif.get_array(col, dtype=object, gather=atom_site_first_model)
+
+  def _requested_models_float_array(col: str) -> np.ndarray:
+    if not model_id:
+      # Return data for all models with a leading dimension of num_models.
+      return cif.get_array(col, dtype=np.float32, gather=atom_site_all_models)
+    else:
+      # Return data only for the single requested model.
+      return cif.get_array(col, dtype=np.float32, gather=atom_site_first_model)
+
+  # These columns are the same for all models, fetch them just for the 1st one.
+  label_comp_ids = _first_model_string_array('_atom_site.label_comp_id')
+  label_asym_ids = _first_model_string_array('_atom_site.label_asym_id')
+  label_seq_ids = _first_model_string_array('_atom_site.label_seq_id')
+  label_atom_ids = _first_model_string_array('_atom_site.label_atom_id')
+  if '_atom_site.auth_seq_id' in cif:
+    auth_seq_ids = _first_model_string_array('_atom_site.auth_seq_id')
+  else:
+    auth_seq_ids = label_seq_ids  # auth_seq_id unset, fallback to label_seq_id.
+  type_symbols = _first_model_string_array('_atom_site.type_symbol')
+  pdbx_pdb_ins_codes = _first_model_string_array('_atom_site.pdbx_PDB_ins_code')
+
+  # These columns are different for all models, fetch them as requested.
+  atom_x = _requested_models_float_array('_atom_site.Cartn_x')
+  atom_y = _requested_models_float_array('_atom_site.Cartn_y')
+  atom_z = _requested_models_float_array('_atom_site.Cartn_z')
+  atom_b_factor = _requested_models_float_array('_atom_site.B_iso_or_equiv')
+  atom_occupancy = _requested_models_float_array('_atom_site.occupancy')
+
+  # Make sure the scheme (residue) tables exist in case they are not present.
+  if cif_update := _maybe_add_missing_scheme_tables(
+      cif,
+      res_starts=layout.residue_starts(),
+      label_asym_ids=label_asym_ids,
+      label_seq_ids=label_seq_ids,
+      label_comp_ids=label_comp_ids,
+      auth_seq_ids=auth_seq_ids,
+      pdb_ins_codes=pdbx_pdb_ins_codes,
+  ):
+    cif = cif.copy_and_update(cif_update)
+
+  # Fix common issues found in mmCIF files, like swapped arginine NH atoms.
+  mmcif_utils.fix_residues(
+      layout,
+      comp_id=label_comp_ids,
+      atom_id=label_atom_ids,
+      atom_x=atom_x[0] if not model_id else atom_x,
+      atom_y=atom_y[0] if not model_id else atom_y,
+      atom_z=atom_z[0] if not model_id else atom_z,
+      fix_arg=fix_arginines,
+  )
+
+  # Get keys for chains in the order they appear in _atom_site while also
+  # dealing with empty chains.
+  resolved_chain_ids = label_asym_ids[layout.chain_starts()]
+  struct_asym_chain_ids = cif.get_array('_struct_asym.id', dtype=object)
+
+  chain_key_by_chain_id = _get_chain_key_by_chain_id(
+      resolved_chain_ids=resolved_chain_ids,
+      struct_asym_chain_ids=struct_asym_chain_ids,
+  )
+  entity_id_by_chain_id = dict(
+      zip(struct_asym_chain_ids, cif['_struct_asym.entity_id'], strict=True)
+  )
+  entity_description = cif.get(
+      '_entity.pdbx_description', ['?'] * len(cif['_entity.id'])
+  )
+  entity_desc_by_entity_id = dict(
+      zip(cif['_entity.id'], entity_description, strict=True)
+  )
+  chain_type_by_entity_id = mmcif.get_chain_type_by_entity_id(cif)
+  auth_asym_id_by_chain_id = mmcif.get_internal_to_author_chain_id_map(cif)
+
+  chain_res_builder = _ChainResBuilder(
+      chain_key_by_chain_id=chain_key_by_chain_id,
+      entity_id_by_chain_id=entity_id_by_chain_id,
+      chain_type_by_entity_id=chain_type_by_entity_id,
+      entity_desc_by_entity_id=entity_desc_by_entity_id,
+      fix_mse_residues=fix_mse_residues,
+      fix_unknown_dna=fix_unknown_dna,
+  )
+
+  # Collect data for polymer chain and residue tables. _pdbx_poly_seq_scheme is
+  # guaranteed to be present thanks to _maybe_add_missing_scheme_tables.
+  def _get_poly_seq_scheme_col(col: str) -> np.ndarray:
+    return cif.get_array(key=f'_pdbx_poly_seq_scheme.{col}', dtype=object)
+
+  poly_seq_asym_ids = _get_poly_seq_scheme_col('asym_id')
+  poly_seq_pdb_seq_nums = _get_poly_seq_scheme_col('pdb_seq_num')
+  poly_seq_seq_ids = _get_poly_seq_scheme_col('seq_id')
+  poly_seq_mon_ids = _get_poly_seq_scheme_col('mon_id')
+  poly_seq_pdb_strand_ids = _get_poly_seq_scheme_col('pdb_strand_id')
+  poly_seq_pdb_ins_codes = _get_poly_seq_scheme_col('pdb_ins_code')
+  string_array.remap(
+      poly_seq_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+
+  # We resolved alt-locs earlier for the atoms table. In cases of heterogeneous
+  # residues (a residue with an alt-loc that is of different residue type), we
+  # need to also do the same resolution in the residues table. Compute a mask
+  # for the residues that were selected in the atoms table.
+  poly_seq_mask = mmcif_utils.selected_polymer_residue_mask(
+      layout=layout,
+      atom_site_label_asym_ids=label_asym_ids[layout.residue_starts()],
+      atom_site_label_seq_ids=label_seq_ids[layout.residue_starts()],
+      atom_site_label_comp_ids=label_comp_ids[layout.residue_starts()],
+      poly_seq_asym_ids=poly_seq_asym_ids,
+      poly_seq_seq_ids=poly_seq_seq_ids,
+      poly_seq_mon_ids=poly_seq_mon_ids,
+  )
+
+  if not include_other and poly_seq_mask:
+    # Mask filtered-out residues so that they are not treated as missing.
+    # Instead, we don't want them included in the chains/residues tables at all.
+    keep_mask = string_array.remap(
+        poly_seq_asym_ids,
+        mapping={cid: True for cid in resolved_chain_ids},
+        default_value=False,
+        inplace=False,
+    ).astype(bool)
+    poly_seq_mask &= keep_mask
+
+  chain_res_builder.add_residues(
+      chain_ids=poly_seq_asym_ids[poly_seq_mask],
+      chain_auth_asym_ids=poly_seq_pdb_strand_ids[poly_seq_mask],
+      res_ids=poly_seq_seq_ids[poly_seq_mask].astype(np.int32),
+      res_names=poly_seq_mon_ids[poly_seq_mask],
+      res_auth_seq_ids=poly_seq_pdb_seq_nums[poly_seq_mask],
+      res_ins_codes=poly_seq_pdb_ins_codes[poly_seq_mask],
+  )
+
+  # Collect data for ligand chain and residue tables. _pdbx_nonpoly_scheme
+  # could be empty/unset if there are only branched ligands.
+  def _get_nonpoly_scheme_col(col: str) -> np.ndarray:
+    key = f'_pdbx_nonpoly_scheme.{col}'
+    if f'_pdbx_nonpoly_scheme.{col}' in cif:
+      return cif.get_array(key=key, dtype=object)
+    else:
+      return np.array([], dtype=object)
+
+  nonpoly_asym_ids = _get_nonpoly_scheme_col('asym_id')
+  nonpoly_auth_seq_ids = _get_nonpoly_scheme_col('pdb_seq_num')
+  nonpoly_pdb_ins_codes = _get_nonpoly_scheme_col('pdb_ins_code')
+  nonpoly_mon_ids = _get_nonpoly_scheme_col('mon_id')
+  nonpoly_auth_asym_id = string_array.remap(
+      nonpoly_asym_ids, mapping=auth_asym_id_by_chain_id, inplace=False
+  )
+
+  def _get_branch_scheme_col(col: str) -> np.ndarray:
+    key = f'_pdbx_branch_scheme.{col}'
+    if f'_pdbx_branch_scheme.{col}' in cif:
+      return cif.get_array(key=key, dtype=object)
+    else:
+      return np.array([], dtype=object)
+
+  branch_asym_ids = _get_branch_scheme_col('asym_id')
+  branch_auth_seq_ids = _get_branch_scheme_col('pdb_seq_num')
+  branch_pdb_ins_codes = _get_branch_scheme_col('pdb_ins_code')
+  branch_mon_ids = _get_branch_scheme_col('mon_id')
+  branch_auth_asym_id = string_array.remap(
+      branch_asym_ids, mapping=auth_asym_id_by_chain_id, inplace=False
+  )
+
+  if branch_asym_ids.size > 0 and branch_pdb_ins_codes.size == 0:
+    branch_pdb_ins_codes = np.array(['.'] * branch_asym_ids.size, dtype=object)
+
+  # Compute the heterogeneous residue masks as above, this time for ligands.
+  nonpoly_mask, branch_mask = mmcif_utils.selected_ligand_residue_mask(
+      layout=layout,
+      atom_site_label_asym_ids=label_asym_ids[layout.residue_starts()],
+      atom_site_label_seq_ids=label_seq_ids[layout.residue_starts()],
+      atom_site_auth_seq_ids=auth_seq_ids[layout.residue_starts()],
+      atom_site_label_comp_ids=label_comp_ids[layout.residue_starts()],
+      atom_site_pdbx_pdb_ins_codes=pdbx_pdb_ins_codes[layout.residue_starts()],
+      nonpoly_asym_ids=nonpoly_asym_ids,
+      nonpoly_auth_seq_ids=nonpoly_auth_seq_ids,
+      nonpoly_pdb_ins_codes=nonpoly_pdb_ins_codes,
+      nonpoly_mon_ids=nonpoly_mon_ids,
+      branch_asym_ids=branch_asym_ids,
+      branch_auth_seq_ids=branch_auth_seq_ids,
+      branch_pdb_ins_codes=branch_pdb_ins_codes,
+      branch_mon_ids=branch_mon_ids,
+  )
+
+  if not include_water:
+    if nonpoly_mask:
+      nonpoly_mask &= (nonpoly_mon_ids != 'HOH') & (nonpoly_mon_ids != 'DOD')
+    if branch_mask:
+      # Fix for bad mmCIFs that have water in the branch scheme table.
+      branch_mask &= (branch_mon_ids != 'HOH') & (branch_mon_ids != 'DOD')
+
+  string_array.remap(
+      pdbx_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+  string_array.remap(
+      nonpoly_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+  string_array.remap(
+      branch_pdb_ins_codes, mapping=_INSERTION_CODE_REMAP, inplace=True
+  )
+
+  def _ligand_residue_ids(chain_ids: np.ndarray) -> np.ndarray:
+    """Computes internal residue ID for ligand residues that don't have it."""
+
+    # E.g. chain_ids=[A, A, A, B, C, C, D, D, D] -> [1, 2, 3, 1, 1, 2, 1, 2, 3].
+    indices = np.arange(chain_ids.size, dtype=np.int32)
+    return (indices + 1) - np.maximum.accumulate(
+        indices * (chain_ids != np.roll(chain_ids, 1))
+    )
+
+  branch_residue_ids = _ligand_residue_ids(branch_asym_ids[branch_mask])
+  nonpoly_residue_ids = _ligand_residue_ids(nonpoly_asym_ids[nonpoly_mask])
+
+  chain_res_builder.add_residues(
+      chain_ids=branch_asym_ids[branch_mask],
+      chain_auth_asym_ids=branch_auth_asym_id[branch_mask],
+      res_ids=branch_residue_ids,
+      res_names=branch_mon_ids[branch_mask],
+      res_auth_seq_ids=branch_auth_seq_ids[branch_mask],
+      res_ins_codes=branch_pdb_ins_codes[branch_mask],
+  )
+
+  chain_res_builder.add_residues(
+      chain_ids=nonpoly_asym_ids[nonpoly_mask],
+      chain_auth_asym_ids=nonpoly_auth_asym_id[nonpoly_mask],
+      res_ids=nonpoly_residue_ids,
+      res_names=nonpoly_mon_ids[nonpoly_mask],
+      res_auth_seq_ids=nonpoly_auth_seq_ids[nonpoly_mask],
+      res_ins_codes=nonpoly_pdb_ins_codes[nonpoly_mask],
+  )
+
+  chains = chain_res_builder.make_chains_table()
+  residues = chain_res_builder.make_residues_table()
+
+  # Construct foreign residue keys for the atoms table.
+  res_ends = np.array(layout.residues(), dtype=np.int32)
+  res_starts = np.array(layout.residue_starts(), dtype=np.int32)
+  res_lengths = res_ends - res_starts
+
+  # Check just for HOH, DOD can be part e.g. of hydroxycysteine.
+  if include_water:
+    res_chain_types = chains.apply_array_to_column(
+        column_name='type', arr=residues.chain_key
+    )
+    water_mask = res_chain_types != mmcif_names.WATER
+    if 'HOH' in set(residues.name[water_mask]):
+      raise ValueError('Bad mmCIF file: non-water entity has water molecules.')
+  else:
+    # Include resolved and unresolved residues.
+    if 'HOH' in set(residues.name) | set(label_comp_ids[res_starts]):
+      raise ValueError('Bad mmCIF file: non-water entity has water molecules.')
+
+  atom_chain_key = string_array.remap(
+      label_asym_ids, mapping=chain_res_builder.chain_key_by_chain_id
+  ).astype(int)
+
+  # If any of the residue lookups failed, the mmCIF is corrupted.
+  try:
+    atom_res_key_per_res = string_array.remap_multiple(
+        (
+            label_asym_ids[res_starts],
+            auth_seq_ids[res_starts],
+            label_comp_ids[res_starts],
+            pdbx_pdb_ins_codes[res_starts],
+        ),
+        mapping=chain_res_builder.key_for_res,
+    )
+  except KeyError as e:
+    raise ValueError(
+        'Lookup for the following atom from the _atom_site table failed: '
+        f'(label_asym_id, auth_seq_id, res_name, ins_code)={e}. This typically '
+        'indicates that the _pdbx_poly_seq_scheme, _pdbx_nonpoly_scheme, or '
+        '_pdbx_branch_scheme tables do not have data for all residues present '
+        'in the _atom_site table. It could also be due to a known issue with '
+        'a small number of multi-model mmCIFs.'
+    ) from e
+
+  # The residue ID will be shared for all atoms within that residue.
+  atom_res_key = np.repeat(atom_res_key_per_res, repeats=res_lengths)
+
+  if fix_mse_residues:
+    met_residues_mask = (residues.name == 'MET')[atom_res_key]
+    unfixed_mse_selenium_mask = met_residues_mask & (label_atom_ids == 'SE')
+    label_atom_ids[unfixed_mse_selenium_mask] = 'SD'
+    type_symbols[unfixed_mse_selenium_mask] = 'S'
+
+  atoms = structure_tables.Atoms(
+      key=atom_site_first_model,
+      chain_key=atom_chain_key,
+      res_key=atom_res_key,
+      name=label_atom_ids,
+      element=type_symbols,
+      x=atom_x,
+      y=atom_y,
+      z=atom_z,
+      b_factor=atom_b_factor,
+      occupancy=atom_occupancy,
+  )
+
+  return chains, residues, atoms
+
+
+def from_atom_arrays(
+    *,
+    res_id: np.ndarray,
+    name: str = 'unset',
+    release_date: datetime.date | None = None,
+    resolution: float | None = None,
+    structure_method: str | None = None,
+    all_residues: Mapping[str, Sequence[tuple[str, int]]] | None = None,
+    bioassembly_data: bioassemblies.BioassemblyData | None = None,
+    chemical_components_data: (
+        struc_chem_comps.ChemicalComponentsData | None
+    ) = None,
+    bond_table: structure_tables.Bonds | None = None,
+    chain_id: np.ndarray | None = None,
+    chain_type: np.ndarray | None = None,
+    res_name: np.ndarray | None = None,
+    atom_key: np.ndarray | None = None,
+    atom_name: np.ndarray | None = None,
+    atom_element: np.ndarray | None = None,
+    atom_x: np.ndarray | None = None,
+    atom_y: np.ndarray | None = None,
+    atom_z: np.ndarray | None = None,
+    atom_b_factor: np.ndarray | None = None,
+    atom_occupancy: np.ndarray | None = None,
+) -> structure.Structure:
+  """Returns a Structure constructed from atom array level data.
+
+  All fields except name and, res_id are optional, all array fields consist of a
+  value for each atom in the structure - so residue and chain values should hold
+  the same value for each atom in the chain or residue. Fields which are not
+  defined are filled with default values.
+
+  Validation is performed by the Structure constructor where possible - but
+  author_naming scheme and all_residues must be checked in this function.
+
+  It is not possible to construct structures with chains that do not contain
+  any resolved residues using this function. If this is necessary, use the
+  structure.Structure constructor directly.
+
+  Args:
+    res_id: Integer array of shape [num_atom]. The unique residue identifier for
+      each residue. mmCIF field - _atom_site.label_seq_id.
+    name: The name of the structure. E.g. a PDB ID.
+    release_date: The release date of the structure as a `datetime.date`.
+    resolution: The resolution of the structure in Angstroms.
+    structure_method: The method used to solve this structure's coordinates.
+    all_residues: An optional mapping from each chain ID (i.e. label_asym_id) to
+      a sequence of (label_comp_id, label_seq_id) tuples, one per residue. This
+      can contain residues that aren't present in the atom arrays. This is
+      common in experimental data where some residues are not resolved but are
+      known to be present.
+    bioassembly_data: An optional instance of bioassembly.BioassemblyData. If
+      present then a new Structure representing a specific bioassembly can be
+      extracted using `Structure.generate_bioassembly(assembly_id)`.
+    chemical_components_data: An optional instance of ChemicalComponentsData.
+      Its content will be used for providing metadata about chemical components
+      in this Structure instance. If not specified information will be retrieved
+      from the standard chemical component dictionary (CCD, for more details see
+      https://www.wwpdb.org/data/ccd).
+    bond_table: A table representing manually-specified bonds. This corresponds
+      to the _struct_conn table in an mmCIF. Atoms are identified by their key,
+      as specified by the atom_key column. If this table is provided then the
+      atom_key column must also be defined.
+    chain_id: String array of shape [num_atom] of unique chain identifiers.
+      mmCIF field - _atom_site.label_asym_id.
+    chain_type: String array of shape [num_atom]. The molecular type of the
+      current chain (e.g. polyribonucleotide). mmCIF field - _entity_poly.type
+      OR _entity.type (for non-polymers).
+    res_name: String array of shape [num_atom].. The name of each residue,
+      typically a 3 letter string for polypeptides or 1-2 letter strings for
+      polynucleotides. mmCIF field - _atom_site.label_comp_id.
+    atom_key: A unique sorted integer array, used only by the bonds table to
+      identify the atoms participating in each bond. If the bonds table is
+      specified then this column must be non-None.
+    atom_name: String array of shape [num_atom]. The name of each atom (e.g CA,
+      O2', etc.). mmCIF field - _atom_site.label_atom_id.
+    atom_element: String array of shape [num_atom]. The element type of each
+      atom (e.g. C, O, N, etc.). mmCIF field - _atom_site.type_symbol.
+    atom_x: Float array of shape [..., num_atom] of atom x coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_y: Float array of shape [..., num_atom] of atom y coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_z: Float array of shape [..., num_atom] of atom z coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_b_factor: Float array of shape [..., num_atom] or [num_atom] of atom
+      b-factors or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+    atom_occupancy: Float array of shape [..., num_atom] or [num_atom] of atom
+      occupancies or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+  """
+
+  atoms, residues, chains = structure_tables.tables_from_atom_arrays(
+      res_id=res_id,
+      all_residues=all_residues,
+      chain_id=chain_id,
+      chain_type=chain_type,
+      res_name=res_name,
+      atom_key=atom_key,
+      atom_name=atom_name,
+      atom_element=atom_element,
+      atom_x=atom_x,
+      atom_y=atom_y,
+      atom_z=atom_z,
+      atom_b_factor=atom_b_factor,
+      atom_occupancy=atom_occupancy,
+  )
+
+  return structure.Structure(
+      name=name,
+      release_date=release_date,
+      resolution=resolution,
+      structure_method=structure_method,
+      bioassembly_data=bioassembly_data,
+      chemical_components_data=chemical_components_data,
+      atoms=atoms,
+      chains=chains,
+      residues=residues,
+      bonds=bond_table or structure_tables.Bonds.make_empty(),
+  )
+
+
+def _guess_entity_type(
+    chain_residues: Collection[str], atom_types: Collection[str]
+) -> str:
+  """Guess the entity type (polymer/non-polymer/water) based on residues/atoms.
+
+  We treat both arguments as unordered collections since we care only whether
+  all elements satisfy come conditions. The chain_residues can be either
+  grouped by residue (length num_res), or it can be raw (length num_atoms).
+  Atom type is unique for each atom in a residue, so don't group atom_types.
+
+  Args:
+    chain_residues: A sequence of full residue name (1-letter for DNA, 2-letters
+      for RNA, 3 for protein). The _atom_site.label_comp_id column in mmCIF.
+    atom_types: Atom type: ATOM or HETATM. The _atom_site.group_PDB column in
+      mmCIF.
+
+  Returns:
+    One of polymer/non-polymer/water based on the following criteria:
+    * If all atoms are HETATMs and all residues are water -> water.
+    * If all atoms are HETATMs and not all residues are water -> non-polymer.
+    * Otherwise -> polymer.
+  """
+  if not chain_residues or not atom_types:
+    raise ValueError(
+        f'chain_residues (len {len(chain_residues)}) and atom_types (len '
+        f'{len(atom_types)}) must be both non-empty. Got: {chain_residues=} '
+        f'and {atom_types=}'
+    )
+
+  if all(a == 'HETATM' for a in atom_types):
+    if all(c in residue_names.WATER_TYPES for c in chain_residues):
+      return mmcif_names.WATER
+    return mmcif_names.NON_POLYMER_CHAIN
+  return mmcif_names.POLYMER_CHAIN

flax_model/alphafold3/structure/sterics.py

@@ -0,0 +1,132 @@
+
+
+"""Functions relating to spatial locations of atoms within a structure."""
+
+from collections.abc import Collection, Sequence
+
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.structure import mmcif
+import numpy as np
+import scipy
+
+
+def _make_atom_has_clash_mask(
+    kd_query_result: np.ndarray,
+    struc: structure.Structure,
+    ignore_chains: Collection[str],
+) -> np.ndarray:
+  """Returns a boolean NumPy array representing whether each atom has a clash.
+
+  Args:
+    kd_query_result: NumPy array containing N-atoms arrays, each array
+      containing indices to atoms that clash with the N'th atom.
+    struc: Structure over which clashes were detected.
+    ignore_chains: Collection of chains that should not be considered clashing.
+      A boolean NumPy array of length N atoms.
+  """
+  atom_is_clashing = np.zeros((struc.num_atoms,), dtype=bool)
+  for atom_index, clashes in enumerate(kd_query_result):
+    chain_i = struc.chain_id[atom_index]
+    if chain_i in ignore_chains:
+      continue
+    islig_i = struc.is_ligand_mask[atom_index]
+    for clashing_atom_index in clashes:
+      chain_c = struc.chain_id[clashing_atom_index]
+      if chain_c in ignore_chains:
+        continue
+      islig_c = struc.is_ligand_mask[clashing_atom_index]
+      if (
+          clashing_atom_index == atom_index
+          or chain_i == chain_c
+          or islig_i != islig_c
+      ):
+        # Ignore clashes within chain or between ligand and polymer.
+        continue
+      atom_is_clashing[atom_index] = True
+  return atom_is_clashing
+
+
+def find_clashing_chains(
+    struc: structure.Structure,
+    clash_thresh_angstrom: float = 1.7,
+    clash_thresh_fraction: float = 0.3,
+) -> Sequence[str]:
+  """Finds chains that clash with others.
+
+  Clashes are defined by polymer backbone atoms and all ligand atoms.
+  Ligand-polymer clashes are not dropped.
+
+  Will not find clashes if all coordinates are 0. Coordinates are all 0s if
+  the structure is generated from sequences only, as done for inference in
+  dendro for example.
+
+  Args:
+    struc: The structure defining the chains and atom positions.
+    clash_thresh_angstrom: Below this distance, atoms are considered clashing.
+    clash_thresh_fraction: Chains with more than this fraction of their atoms
+      considered clashing will be dropped. This value should be in the range (0,
+      1].
+
+  Returns:
+    A sequence of chain ids for chains that clash.
+
+  Raises:
+    ValueError: If `clash_thresh_fraction` is not in range (0,1].
+  """
+  if not 0 < clash_thresh_fraction <= 1:
+    raise ValueError('clash_thresh_fraction must be in range (0,1]')
+
+  struc_backbone = struc.filter_polymers_to_single_atom_per_res()
+  if struc_backbone.num_chains == 0:
+    return []
+
+  # If the coordinates are all 0, do not search for clashes.
+  if not np.any(struc_backbone.coords):
+    return []
+
+  coord_kdtree = scipy.spatial.cKDTree(struc_backbone.coords)
+
+  # For each atom coordinate, find all atoms within the clash thresh radius.
+  clashing_per_atom = coord_kdtree.query_ball_point(
+      struc_backbone.coords, r=clash_thresh_angstrom
+  )
+  chain_ids = struc_backbone.chains
+  if struc_backbone.atom_occupancy is not None:
+    chain_occupancy = np.array([
+        np.mean(struc_backbone.atom_occupancy[start:end])
+        for start, end in struc_backbone.iter_chain_ranges()
+    ])
+  else:
+    chain_occupancy = None
+
+  # Remove chains until no more significant clashing.
+  chains_to_remove = set()
+  for _ in range(len(chain_ids)):
+    # Calculate maximally clashing.
+    atom_has_clash = _make_atom_has_clash_mask(
+        clashing_per_atom, struc_backbone, chains_to_remove
+    )
+    clashes_per_chain = np.array([
+        atom_has_clash[start:end].mean()
+        for start, end in struc_backbone.iter_chain_ranges()
+    ])
+    max_clash = np.max(clashes_per_chain)
+    if max_clash <= clash_thresh_fraction:
+      # None of the remaining chains exceed the clash fraction threshold, so
+      # we can exit.
+      break
+
+    # Greedily remove worst with the lowest occupancy.
+    most_clashes = np.nonzero(clashes_per_chain == max_clash)[0]
+    if chain_occupancy is not None:
+      occupancy_clashing = chain_occupancy[most_clashes]
+      last_lowest_occupancy = (
+          len(occupancy_clashing) - np.argmin(occupancy_clashing[::-1]) - 1
+      )
+      worst_and_last = most_clashes[last_lowest_occupancy]
+    else:
+      worst_and_last = most_clashes[-1]
+
+    chains_to_remove.add(chain_ids[worst_and_last])
+
+  return sorted(chains_to_remove, key=mmcif.str_id_to_int_id)

flax_model/alphafold3/structure/structure_tables.py

@@ -0,0 +1,811 @@
+
+
+"""Table implementations for the Structure class."""
+
+import collections
+from collections.abc import Mapping, Sequence
+import dataclasses
+import functools
+import itertools
+import typing
+from typing import Any, ClassVar, Self
+
+from flax_model.alphafold3.constants import mmcif_names
+from flax_model.alphafold3.constants import residue_names
+from flax_model.alphafold3.cpp import aggregation
+from flax_model.alphafold3.cpp import string_array
+from flax_model.alphafold3.structure import bonds as bonds_module
+from flax_model.alphafold3.structure import mmcif
+from flax_model.alphafold3.structure import table
+import numpy as np
+
+
+Bonds = bonds_module.Bonds
+
+
+def _residue_name_to_record_name(
+    residue_name: np.ndarray,
+    polymer_mask: np.ndarray,
+) -> np.ndarray:
+  """Returns record names (ATOM/HETATM) given residue names and polymer mask."""
+  record_name = np.array(['HETATM'] * len(residue_name), dtype=object)
+  record_name[polymer_mask] = string_array.remap(
+      residue_name[polymer_mask],
+      mapping={r: 'ATOM' for r in residue_names.STANDARD_POLYMER_TYPES},
+      default_value='HETATM',
+  )
+  return record_name
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class AuthorNamingScheme:
+  """A mapping from internal values to author values in a mmCIF.
+
+  Fields:
+    auth_asym_id: A mapping from label_asym_id to auth_asym_id.
+    auth_seq_id: A mapping from label_asym_id to a mapping from
+      label_seq_id to auth_seq_id.
+    insertion_code: A mapping from label_asym_id to a mapping from
+      label_seq_id to insertion codes.
+    entity_id: A mapping from label_asym_id to _entity.id.
+    entity_desc: A mapping from _entity.id to _entity.pdbx_description.
+  """
+
+  auth_asym_id: Mapping[str, str]
+  auth_seq_id: Mapping[str, Mapping[int, str]]
+  insertion_code: Mapping[str, Mapping[int, str | None]]
+  entity_id: Mapping[str, str]
+  entity_desc: Mapping[str, str]
+
+
+def _default(
+    candidate_value: np.ndarray | None, default_value: Sequence[Any], dtype: Any
+) -> np.ndarray:
+  if candidate_value is None:
+    return np.array(default_value, dtype=dtype)
+  return np.array(candidate_value, dtype=dtype)
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Atoms(table.Table):
+  """Table of atoms in a Structure."""
+
+  chain_key: np.ndarray
+  res_key: np.ndarray
+  name: np.ndarray
+  element: np.ndarray
+  x: np.ndarray
+  y: np.ndarray
+  z: np.ndarray
+  b_factor: np.ndarray
+  occupancy: np.ndarray
+  multimodel_cols: ClassVar[tuple[str, ...]] = (
+      'x',
+      'y',
+      'z',
+      'b_factor',
+      'occupancy',
+  )
+
+  def __post_init__(self):
+    # Validates that the atom coordinates, b-factors and occupancies are finite.
+    for column_name in ('x', 'y', 'z', 'b_factor', 'occupancy'):
+      column = self.get_column(column_name)
+      if not np.isfinite(column).all():
+        raise ValueError(
+            f'Column {column_name} must not contain NaN/inf values.'
+        )
+    # super().__post_init__() can't be used as that causes the following error:
+    # TypeError: super(type, obj): obj must be an instance or subtype of type
+    super(Atoms, self).__post_init__()
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.array([], dtype=np.int64),
+        chain_key=np.array([], dtype=np.int64),
+        res_key=np.array([], dtype=np.int64),
+        name=np.array([], dtype=object),
+        element=np.array([], dtype=object),
+        x=np.array([], dtype=np.float32),
+        y=np.array([], dtype=np.float32),
+        z=np.array([], dtype=np.float32),
+        b_factor=np.array([], dtype=np.float32),
+        occupancy=np.array([], dtype=np.float32),
+    )
+
+  @classmethod
+  def from_defaults(
+      cls,
+      *,
+      chain_key: np.ndarray,
+      res_key: np.ndarray,
+      key: np.ndarray | None = None,
+      name: np.ndarray | None = None,
+      element: np.ndarray | None = None,
+      x: np.ndarray | None = None,
+      y: np.ndarray | None = None,
+      z: np.ndarray | None = None,
+      b_factor: np.ndarray | None = None,
+      occupancy: np.ndarray | None = None,
+  ) -> Self:
+    """Create an Atoms table with minimal user inputs."""
+    num_atoms = len(chain_key)
+    if not num_atoms:
+      return cls.make_empty()
+    return Atoms(
+        chain_key=chain_key,
+        res_key=res_key,
+        key=_default(key, np.arange(num_atoms), np.int64),
+        name=_default(name, ['?'] * num_atoms, object),
+        element=_default(element, ['?'] * num_atoms, object),
+        x=_default(x, [0.0] * num_atoms, np.float32),
+        y=_default(y, [0.0] * num_atoms, np.float32),
+        z=_default(z, [0.0] * num_atoms, np.float32),
+        b_factor=_default(b_factor, [0.0] * num_atoms, np.float32),
+        occupancy=_default(occupancy, [1.0] * num_atoms, np.float32),
+    )
+
+  def get_value_by_index(
+      self, column_name: str, index: int
+  ) -> table.TableEntry | np.ndarray:
+    if column_name in self.multimodel_cols:
+      return self.get_column(column_name)[..., index]
+    else:
+      return self.get_column(column_name)[index]
+
+  def copy_and_update_coords(self, coords: np.ndarray) -> Self:
+    """Returns a copy with the x, y and z columns updated."""
+    if coords.shape[-1] != 3:
+      raise ValueError(
+          f'Expecting 3-dimensional coordinates, got {coords.shape}'
+      )
+    return typing.cast(
+        Atoms,
+        self.copy_and_update(
+            x=coords[..., 0], y=coords[..., 1], z=coords[..., 2]
+        ),
+    )
+
+  @property
+  def shape(self) -> tuple[int, ...]:
+    return self.x.shape
+
+  @property
+  def ndim(self) -> int:
+    return len(self.shape)
+
+  @functools.cached_property
+  def num_models(self) -> int:
+    """The number of models of this Structure."""
+    leading_dims = self.shape[:-1]
+    match leading_dims:
+      case ():
+        return 1
+      case (single_leading_dim_size,):
+        return single_leading_dim_size
+      case _:
+        raise ValueError(
+            'num_models not defined for atom tables with more than one '
+            'leading dimension.'
+        )
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Residues(table.Table):
+  """Table of residues in a Structure."""
+
+  chain_key: np.ndarray
+  id: np.ndarray
+  name: np.ndarray
+  auth_seq_id: np.ndarray
+  insertion_code: np.ndarray
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.array([], dtype=np.int64),
+        chain_key=np.array([], dtype=np.int64),
+        id=np.array([], dtype=np.int32),
+        name=np.array([], dtype=object),
+        auth_seq_id=np.array([], dtype=object),
+        insertion_code=np.array([], dtype=object),
+    )
+
+  @classmethod
+  def from_defaults(
+      cls,
+      *,
+      id: np.ndarray,  # pylint:disable=redefined-builtin
+      chain_key: np.ndarray,
+      key: np.ndarray | None = None,
+      name: np.ndarray | None = None,
+      auth_seq_id: np.ndarray | None = None,
+      insertion_code: np.ndarray | None = None,
+  ) -> Self:
+    """Create a Residues table with minimal user inputs."""
+    num_res = len(id)
+    if not num_res:
+      return cls.make_empty()
+    return Residues(
+        key=_default(key, np.arange(num_res), np.int64),
+        id=id,
+        chain_key=chain_key,
+        name=_default(name, ['UNK'] * num_res, object),
+        auth_seq_id=_default(auth_seq_id, id.astype(str), object),
+        insertion_code=_default(insertion_code, ['?'] * num_res, object),
+    )
+
+
+@dataclasses.dataclass(frozen=True, slots=True, kw_only=True)
+class Chains(table.Table):
+  """Table of chains in a Structure."""
+
+  id: np.ndarray
+  type: np.ndarray
+  auth_asym_id: np.ndarray
+  entity_id: np.ndarray
+  entity_desc: np.ndarray
+
+  @classmethod
+  def make_empty(cls) -> Self:
+    return cls(
+        key=np.array([], dtype=np.int64),
+        id=np.array([], dtype=object),
+        type=np.array([], dtype=object),
+        auth_asym_id=np.array([], dtype=object),
+        entity_id=np.array([], dtype=object),
+        entity_desc=np.array([], dtype=object),
+    )
+
+  @classmethod
+  def from_defaults(
+      cls,
+      *,
+      id: np.ndarray,  # pylint:disable=redefined-builtin
+      key: np.ndarray | None = None,
+      type: np.ndarray | None = None,  # pylint:disable=redefined-builtin
+      auth_asym_id: np.ndarray | None = None,
+      entity_id: np.ndarray | None = None,
+      entity_desc: np.ndarray | None = None,
+  ) -> Self:
+    """Create a Chains table with minimal user inputs."""
+    num_chains = len(id)
+    if not num_chains:
+      return cls.make_empty()
+
+    return Chains(
+        key=_default(key, np.arange(num_chains), np.int64),
+        id=id,
+        type=_default(type, [mmcif_names.PROTEIN_CHAIN] * num_chains, object),
+        auth_asym_id=_default(auth_asym_id, id, object),
+        entity_id=_default(
+            entity_id, np.arange(1, num_chains + 1).astype(str), object
+        ),
+        entity_desc=_default(entity_desc, ['.'] * num_chains, object),
+    )
+
+
+def to_mmcif_sequence_and_entity_tables(
+    chains: Chains,
+    residues: Residues,
+    atom_res_key: np.ndarray,
+) -> Mapping[str, Sequence[str]]:
+  """Returns raw sequence and entity mmCIF tables."""
+  raw_mmcif = collections.defaultdict(list)
+  chains_by_entity_id = {}
+  written_entity_poly_seq_ids = set()
+  present_res_keys = set(atom_res_key)
+
+  # Performance optimisation: Find residue indices for each chain in advance, so
+  # that we don't have to do redunant masking work for each chain.
+  res_indices_for_chain = aggregation.indices_grouped_by_value(
+      residues.chain_key
+  )
+
+  for chain in chains.iterrows():
+    # Add all chain information to the _struct_asym table.
+    chain_id = chain['id']  # Saves multiple dict lookups.
+    auth_asym_id = chain['auth_asym_id']
+    entity_id = chain['entity_id']
+    chains_by_entity_id.setdefault(entity_id, []).append(chain)
+    raw_mmcif['_struct_asym.id'].append(chain_id)
+    raw_mmcif['_struct_asym.entity_id'].append(entity_id)
+
+    res_chain_indices = res_indices_for_chain[chain['key']]
+    chain_type = chain['type']
+    is_polymer = chain_type in mmcif_names.POLYMER_CHAIN_TYPES
+    is_water = chain_type == mmcif_names.WATER
+    is_branched = len(res_chain_indices) > 1 and not is_polymer and not is_water
+    write_entity_poly_seq = entity_id not in written_entity_poly_seq_ids
+
+    # Iterate over the individual masked residue table columns, as that doesn't
+    # create a copy (only a view), while residues[res_chain_indices] does.
+    for res_key, res_name, res_id, pdb_seq_num, res_ins_code in zip(
+        residues.key[res_chain_indices],
+        residues.name[res_chain_indices],
+        residues.id[res_chain_indices],
+        residues.auth_seq_id[res_chain_indices],
+        residues.insertion_code[res_chain_indices],
+        strict=True,
+    ):
+      is_missing = res_key not in present_res_keys
+      str_res_id = str(res_id)
+      # While atom_site uses "?" for insertion codes, scheme tables use ".".
+      ins_code = (res_ins_code or '.').replace('?', '.')
+      auth_seq_num = '?' if is_missing else pdb_seq_num
+
+      if is_polymer:
+        raw_mmcif['_pdbx_poly_seq_scheme.asym_id'].append(chain_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.entity_id'].append(entity_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.seq_id'].append(str_res_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.mon_id'].append(res_name)
+        raw_mmcif['_pdbx_poly_seq_scheme.pdb_seq_num'].append(pdb_seq_num)
+        raw_mmcif['_pdbx_poly_seq_scheme.auth_seq_num'].append(auth_seq_num)
+        raw_mmcif['_pdbx_poly_seq_scheme.pdb_strand_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_poly_seq_scheme.pdb_ins_code'].append(ins_code)
+        # Structure doesn't support heterogeneous sequences.
+        raw_mmcif['_pdbx_poly_seq_scheme.hetero'].append('n')
+        if write_entity_poly_seq:
+          raw_mmcif['_entity_poly_seq.entity_id'].append(entity_id)
+          raw_mmcif['_entity_poly_seq.num'].append(str_res_id)
+          raw_mmcif['_entity_poly_seq.mon_id'].append(res_name)
+          # Structure doesn't support heterogeneous sequences.
+          raw_mmcif['_entity_poly_seq.hetero'].append('n')
+          written_entity_poly_seq_ids.add(entity_id)
+      elif is_branched:
+        raw_mmcif['_pdbx_branch_scheme.asym_id'].append(chain_id)
+        raw_mmcif['_pdbx_branch_scheme.entity_id'].append(entity_id)
+        raw_mmcif['_pdbx_branch_scheme.mon_id'].append(res_name)
+        raw_mmcif['_pdbx_branch_scheme.num'].append(str_res_id)
+        raw_mmcif['_pdbx_branch_scheme.pdb_asym_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_branch_scheme.pdb_seq_num'].append(pdb_seq_num)
+        raw_mmcif['_pdbx_branch_scheme.auth_asym_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_branch_scheme.auth_seq_num'].append(auth_seq_num)
+        raw_mmcif['_pdbx_branch_scheme.pdb_ins_code'].append(ins_code)
+        # Structure doesn't support heterogeneous sequences.
+        raw_mmcif['_pdbx_branch_scheme.hetero'].append('n')
+      else:
+        raw_mmcif['_pdbx_nonpoly_scheme.asym_id'].append(chain_id)
+        raw_mmcif['_pdbx_nonpoly_scheme.entity_id'].append(entity_id)
+        raw_mmcif['_pdbx_nonpoly_scheme.mon_id'].append(res_name)
+        raw_mmcif['_pdbx_nonpoly_scheme.pdb_seq_num'].append(pdb_seq_num)
+        raw_mmcif['_pdbx_nonpoly_scheme.auth_seq_num'].append(auth_seq_num)
+        raw_mmcif['_pdbx_nonpoly_scheme.pdb_strand_id'].append(auth_asym_id)
+        raw_mmcif['_pdbx_nonpoly_scheme.pdb_ins_code'].append(ins_code)
+
+  # Add _entity and _entity_poly tables.
+  for entity_id, chains in chains_by_entity_id.items():
+    # chains should always be a non-empty list because of how we constructed
+    # chains_by_entity_id.
+    assert chains
+    # All chains for a given entity should have the same type and sequence
+    # so we can pick the first one without losing information.
+    key_chain = chains[0]
+    raw_mmcif['_entity.id'].append(entity_id)
+    raw_mmcif['_entity.pdbx_description'].append(key_chain['entity_desc'])
+    entity_type = key_chain['type']
+    if entity_type not in mmcif_names.POLYMER_CHAIN_TYPES:
+      raw_mmcif['_entity.type'].append(entity_type)
+    else:
+      raw_mmcif['_entity.type'].append('polymer')
+      raw_mmcif['_entity_poly.entity_id'].append(entity_id)
+      raw_mmcif['_entity_poly.type'].append(entity_type)
+
+      # _entity_poly.pdbx_strand_id is a comma-separated list of
+      # auth_asym_ids that are part of the entity.
+      raw_mmcif['_entity_poly.pdbx_strand_id'].append(
+          ','.join(chain['auth_asym_id'] for chain in chains)
+      )
+  return raw_mmcif
+
+
+def to_mmcif_atom_site_and_bonds_table(
+    *,
+    chains: Chains,
+    residues: Residues,
+    atoms: Atoms,
+    bonds: Bonds,
+    coords_decimal_places: int,
+) -> Mapping[str, Sequence[str]]:
+  """Returns raw _atom_site and _struct_conn mmCIF tables."""
+  raw_mmcif = collections.defaultdict(list)
+  # Use [value] * num wherever possible since it is about 10x faster than list
+  # comprehension in such cases. Also use f-strings instead of str() - faster.
+  total_atoms = atoms.size * atoms.num_models
+  raw_mmcif['_atom_site.id'] = [f'{i}' for i in range(1, total_atoms + 1)]
+  raw_mmcif['_atom_site.label_alt_id'] = ['.'] * total_atoms
+  # Use format_float_array instead of list comprehension for performance.
+  raw_mmcif['_atom_site.Cartn_x'] = mmcif.format_float_array(
+      values=atoms.x.ravel(), num_decimal_places=coords_decimal_places
+  )
+  raw_mmcif['_atom_site.Cartn_y'] = mmcif.format_float_array(
+      values=atoms.y.ravel(), num_decimal_places=coords_decimal_places
+  )
+  raw_mmcif['_atom_site.Cartn_z'] = mmcif.format_float_array(
+      values=atoms.z.ravel(), num_decimal_places=coords_decimal_places
+  )
+
+  # atoms.b_factor or atoms.occupancy can be flat even when the coordinates have
+  # leading dimensions. In this case we tile it to match.
+  if atoms.b_factor.ndim == 1:
+    atom_b_factor = np.tile(atoms.b_factor, atoms.num_models)
+  else:
+    atom_b_factor = atoms.b_factor.ravel()
+  raw_mmcif['_atom_site.B_iso_or_equiv'] = mmcif.format_float_array(
+      values=atom_b_factor, num_decimal_places=2
+  )
+
+  if atoms.occupancy.ndim == 1:
+    atom_occupancy = np.tile(atoms.occupancy, atoms.num_models)
+  else:
+    atom_occupancy = atoms.occupancy.ravel()
+  raw_mmcif['_atom_site.occupancy'] = mmcif.format_float_array(
+      values=atom_occupancy.ravel(), num_decimal_places=2
+  )
+
+  label_atom_id = atoms.name
+  type_symbol = atoms.element
+  label_comp_id = residues.apply_array_to_column('name', atoms.res_key)
+  label_asym_id = chains.apply_array_to_column('id', atoms.chain_key)
+  label_entity_id = chains.apply_array_to_column('entity_id', atoms.chain_key)
+  # Performance optimisation: Do the int->str conversion on num_residue-sized,
+  # array, then select instead of selecting and then converting.
+  label_seq_id = residues.id.astype('str').astype(object)[
+      ..., residues.index_by_key[atoms.res_key]
+  ]
+
+  # _atom_site.label_seq_id is '.' for non-polymers.
+  non_polymer_chain_mask = string_array.isin(
+      chains.type, mmcif_names.POLYMER_CHAIN_TYPES, invert=True
+  )
+  non_polymer_chain_keys = chains.key[non_polymer_chain_mask]
+  non_polymer_atom_mask = np.isin(atoms.chain_key, non_polymer_chain_keys)
+  label_seq_id[non_polymer_atom_mask] = '.'
+
+  auth_asym_id = chains.apply_array_to_column('auth_asym_id', atoms.chain_key)
+  auth_seq_id = residues.apply_array_to_column('auth_seq_id', atoms.res_key)
+  pdbx_pdb_ins_code = residues.apply_array_to_column(
+      'insertion_code', atoms.res_key
+  )
+  string_array.remap(pdbx_pdb_ins_code, mapping={None: '?'}, inplace=True)
+
+  group_pdb = _residue_name_to_record_name(
+      residue_name=label_comp_id, polymer_mask=~non_polymer_atom_mask
+  )
+
+  def tile_for_models(arr: np.ndarray) -> list[str]:
+    if atoms.num_models == 1:
+      return arr.tolist()  # Memory optimisation: np.tile(arr, 1) does a copy.
+    return np.tile(arr, atoms.num_models).tolist()
+
+  raw_mmcif['_atom_site.group_PDB'] = tile_for_models(group_pdb)
+  raw_mmcif['_atom_site.label_atom_id'] = tile_for_models(label_atom_id)
+  raw_mmcif['_atom_site.type_symbol'] = tile_for_models(type_symbol)
+  raw_mmcif['_atom_site.label_comp_id'] = tile_for_models(label_comp_id)
+  raw_mmcif['_atom_site.label_asym_id'] = tile_for_models(label_asym_id)
+  raw_mmcif['_atom_site.label_entity_id'] = tile_for_models(label_entity_id)
+  raw_mmcif['_atom_site.label_seq_id'] = tile_for_models(label_seq_id)
+  raw_mmcif['_atom_site.auth_asym_id'] = tile_for_models(auth_asym_id)
+  raw_mmcif['_atom_site.auth_seq_id'] = tile_for_models(auth_seq_id)
+  raw_mmcif['_atom_site.pdbx_PDB_ins_code'] = tile_for_models(pdbx_pdb_ins_code)
+  model_id = np.array(
+      [str(i + 1) for i in range(atoms.num_models)], dtype=object
+  )
+  raw_mmcif['_atom_site.pdbx_PDB_model_num'] = np.repeat(
+      model_id, [atoms.size] * atoms.num_models
+  ).tolist()
+
+  if bonds.key.size > 0:
+    raw_mmcif.update(
+        bonds.to_mmcif_dict_from_atom_arrays(
+            atom_key=atoms.key,
+            chain_id=label_asym_id,
+            res_id=label_seq_id,
+            res_name=label_comp_id,
+            atom_name=label_atom_id,
+            auth_asym_id=auth_asym_id,
+            auth_seq_id=auth_seq_id,
+            insertion_code=np.array(pdbx_pdb_ins_code),
+        )
+    )
+  return raw_mmcif
+
+
+def _flatten_author_naming_scheme_table(
+    res_table: Mapping[str, Mapping[int, str]],
+    chain_ids: np.ndarray,
+    res_chain_ids: np.ndarray,
+    res_ids: np.ndarray,
+    default_if_missing: str,
+    table_name: str,
+) -> np.ndarray:
+  """Flattens an author naming scheme table consistently with res_ids."""
+  if not set(chain_ids).issubset(res_table):
+    raise ValueError(
+        f'Chain IDs in the chain_id array must be a subset of {table_name} in '
+        'author naming scheme:\n'
+        f'chain_ids: {sorted(chain_ids)}\n'
+        f'{table_name} keys: {sorted(res_table.keys())}'
+    )
+
+  chain_change_mask = res_chain_ids[1:] != res_chain_ids[:-1]
+  res_chain_boundaries = np.concatenate(
+      ([0], np.where(chain_change_mask)[0] + 1, [len(res_chain_ids)])
+  )
+
+  flat_vals = np.empty(len(res_ids), dtype=object)
+  for chain_start, chain_end in itertools.pairwise(res_chain_boundaries):
+    chain_id = res_chain_ids[chain_start]
+    chain_res_ids = res_ids[chain_start:chain_end]
+    chain_mapping = res_table[chain_id]
+    flat_vals[chain_start:chain_end] = [
+        chain_mapping.get(r, default_if_missing) for r in chain_res_ids
+    ]
+
+  return flat_vals
+
+
+def tables_from_atom_arrays(
+    *,
+    res_id: np.ndarray,
+    author_naming_scheme: AuthorNamingScheme | None = None,
+    all_residues: Mapping[str, Sequence[tuple[str, int]]] | None = None,
+    chain_id: np.ndarray | None = None,
+    chain_type: np.ndarray | None = None,
+    res_name: np.ndarray | None = None,
+    atom_key: np.ndarray | None = None,
+    atom_name: np.ndarray | None = None,
+    atom_element: np.ndarray | None = None,
+    atom_x: np.ndarray | None = None,
+    atom_y: np.ndarray | None = None,
+    atom_z: np.ndarray | None = None,
+    atom_b_factor: np.ndarray | None = None,
+    atom_occupancy: np.ndarray | None = None,
+) -> tuple[Atoms, Residues, Chains]:
+  """Returns Structure tables constructed from atom array level data.
+
+  All fields except name and, res_id are optional, all array fields consist of a
+  value for each atom in the structure - so residue and chain values should hold
+  the same value for each atom in the chain or residue. Fields which are not
+  defined are filled with default values.
+
+  Validation is performed by the Structure constructor where possible - but
+  author_naming scheme and all_residues must be checked in this function.
+
+  It is not possible to construct structures with chains that do not contain
+  any resolved residues using this function. If this is necessary, use the
+  structure.Structure constructor directly.
+
+  Args:
+    res_id: Integer array of shape [num_atom]. The unique residue identifier for
+      each residue. mmCIF field - _atom_site.label_seq_id.
+    author_naming_scheme: An optional instance of AuthorNamingScheme to use when
+      converting this structure to mmCIF.
+    all_residues: An optional mapping from each chain ID (i.e. label_asym_id) to
+      a sequence of (label_comp_id, label_seq_id) tuples, one per residue. This
+      can contain residues that aren't present in the atom arrays. This is
+      common in experimental data where some residues are not resolved but are
+      known to be present.
+    chain_id: String array of shape [num_atom] of unique chain identifiers.
+      mmCIF field - _atom_site.label_asym_id.
+    chain_type: String array of shape [num_atom]. The molecular type of the
+      current chain (e.g. polyribonucleotide). mmCIF field - _entity_poly.type
+      OR _entity.type (for non-polymers).
+    res_name: String array of shape [num_atom].. The name of each residue,
+      typically a 3 letter string for polypeptides or 1-2 letter strings for
+      polynucleotides. mmCIF field - _atom_site.label_comp_id.
+    atom_key: A unique sorted integer array, used only by the bonds table to
+      identify the atoms participating in each bond. If the bonds table is
+      specified then this column must be non-None.
+    atom_name: String array of shape [num_atom]. The name of each atom (e.g CA,
+      O2', etc.). mmCIF field - _atom_site.label_atom_id.
+    atom_element: String array of shape [num_atom]. The element type of each
+      atom (e.g. C, O, N, etc.). mmCIF field - _atom_site.type_symbol.
+    atom_x: Float array of shape [..., num_atom] of atom x coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_y: Float array of shape [..., num_atom] of atom y coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_z: Float array of shape [..., num_atom] of atom z coordinates. May have
+      arbitrary leading dimensions, provided that these are consistent across
+      all coordinate fields.
+    atom_b_factor: Float array of shape [..., num_atom] or [num_atom] of atom
+      b-factors or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+    atom_occupancy: Float array of shape [..., num_atom] or [num_atom] of atom
+      occupancies or equivalent. If there are no extra leading dimensions then
+      these values are assumed to apply to all coordinates for a given atom. If
+      there are leading dimensions then these must match those used by the
+      coordinate fields.
+  """
+  num_atoms = len(res_id)
+
+  for arr_name, array, dtype in (
+      ('chain_id', chain_id, object),
+      ('chain_type', chain_type, object),
+      ('res_id', res_id, np.int32),
+      ('res_name', res_name, object),
+      ('atom_key', atom_key, np.int64),
+      ('atom_name', atom_name, object),
+      ('atom_element', atom_element, object),
+  ):
+    if array is not None and array.shape != (num_atoms,):
+      raise ValueError(f'{arr_name} shape {array.shape} != ({num_atoms},)')
+    if array is not None and array.dtype != dtype:
+      raise ValueError(f'{arr_name} dtype {array.dtype} != {dtype}')
+
+  for arr_name, array in (
+      ('atom_x', atom_x),
+      ('atom_y', atom_y),
+      ('atom_z', atom_z),
+      ('atom_b_factor', atom_b_factor),
+      ('atom_occupancy', atom_occupancy),
+  ):
+    if array is not None and array.shape[-1] != num_atoms:
+      raise ValueError(f'{arr_name} last dim {array.shape[-1]} != {num_atoms=}')
+    if (
+        array is not None
+        and array.dtype != np.float32
+        and array.dtype != np.float64
+    ):
+      raise ValueError(
+          f'{arr_name} must be np.float32 or np.float64, got {array.dtype=}'
+      )
+
+  if all_residues is not None and (res_name is None or res_id is None):
+    raise ValueError(
+        'If all_residues != None, res_name and res_id must not be None either.'
+    )
+
+  if num_atoms == 0:
+    return Atoms.make_empty(), Residues.make_empty(), Chains.make_empty()
+
+  if chain_id is None:
+    chain_id = np.full(shape=num_atoms, fill_value='A', dtype=object)
+  if res_name is None:
+    res_name = np.full(shape=num_atoms, fill_value='UNK', dtype=object)
+
+  chain_change_mask = chain_id[1:] != chain_id[:-1]
+  chain_start = np.concatenate(([0], np.where(chain_change_mask)[0] + 1))
+  res_start = np.concatenate(
+      ([0], np.where((res_id[1:] != res_id[:-1]) | chain_change_mask)[0] + 1)
+  )
+
+  if len(set(chain_id)) != len(chain_start):
+    raise ValueError(f'Chain IDs must be contiguous, but got {chain_id}')
+
+  # We do not support chains with unresolved residues-only in this function.
+  chain_ids = chain_id[chain_start]
+  if all_residues and set(all_residues.keys()) != set(chain_ids):
+    raise ValueError(
+        'all_residues must contain the same set of chain IDs as the chain_id '
+        f'array:\nall_residues keys: {sorted(all_residues.keys())}\n'
+        f'chain_ids: {sorted(chain_ids)}.'
+    )
+  # Make sure all_residue ordering is consistent with chain_id.
+  if all_residues and np.any(list(all_residues.keys()) != chain_ids):
+    all_residues = {cid: all_residues[cid] for cid in chain_ids}
+
+  # Create the chains table.
+  num_chains = len(chain_ids)
+  chain_keys = np.arange(num_chains, dtype=np.int64)
+  chain_key_by_chain_id = dict(zip(chain_ids, chain_keys, strict=True))
+
+  if chain_type is not None:
+    chain_types = chain_type[chain_start]
+  else:
+    chain_types = np.full(num_chains, mmcif_names.PROTEIN_CHAIN, dtype=object)
+
+  if author_naming_scheme is not None:
+    auth_asym_id = string_array.remap(
+        chain_ids, author_naming_scheme.auth_asym_id
+    )
+    entity_id = string_array.remap(
+        chain_ids, author_naming_scheme.entity_id, default_value='.'
+    )
+    entity_desc = string_array.remap(
+        entity_id, author_naming_scheme.entity_desc, default_value='.'
+    )
+  else:
+    auth_asym_id = chain_ids
+    entity_id = (chain_keys + 1).astype(str).astype(object)
+    entity_desc = np.full(num_chains, '.', dtype=object)
+
+  chains = Chains(
+      key=chain_keys,
+      id=chain_ids,
+      type=chain_types,
+      auth_asym_id=auth_asym_id,
+      entity_id=entity_id,
+      entity_desc=entity_desc,
+  )
+
+  # Create the residues table.
+  if all_residues is not None:
+    residue_order = []
+    for cid, residues in all_residues.items():
+      residue_order.extend((cid, rname, int(rid)) for (rname, rid) in residues)
+    res_chain_ids, res_names, res_ids = zip(*residue_order)
+    res_chain_ids = np.array(res_chain_ids, dtype=object)
+    res_ids = np.array(res_ids, dtype=np.int32)
+    res_names = np.array(res_names, dtype=object)
+  else:
+    res_chain_ids = chain_id[res_start]
+    res_ids = res_id[res_start]
+    res_names = res_name[res_start]
+    residue_order = list(zip(res_chain_ids, res_names, res_ids))
+
+  if author_naming_scheme is not None and author_naming_scheme.auth_seq_id:
+    auth_seq_id = _flatten_author_naming_scheme_table(
+        author_naming_scheme.auth_seq_id,
+        chain_ids=chain_ids,
+        res_chain_ids=res_chain_ids,
+        res_ids=res_ids,
+        default_if_missing='.',
+        table_name='auth_seq_id',
+    )
+  else:
+    auth_seq_id = res_ids.astype(str).astype(object)
+
+  if author_naming_scheme is not None and author_naming_scheme.insertion_code:
+    insertion_code = _flatten_author_naming_scheme_table(
+        author_naming_scheme.insertion_code,
+        chain_ids=chain_ids,
+        res_chain_ids=res_chain_ids,
+        res_ids=res_ids,
+        default_if_missing='?',
+        table_name='insertion_code',
+    )
+    # Make sure insertion code of None is mapped to '.'.
+    insertion_code = string_array.remap(insertion_code, {None: '?'})
+  else:
+    insertion_code = np.full(shape=len(res_ids), fill_value='?', dtype=object)
+
+  res_key_by_res = {res: i for i, res in enumerate(residue_order)}
+  res_keys = np.arange(len(residue_order), dtype=np.int64)
+  res_chain_keys = string_array.remap(
+      res_chain_ids, chain_key_by_chain_id
+  ).astype(np.int64)
+  residues = Residues(
+      chain_key=res_chain_keys,
+      key=res_keys,
+      id=res_ids,
+      name=res_names,
+      auth_seq_id=auth_seq_id,
+      insertion_code=insertion_code,
+  )
+
+  if atom_key is None:
+    atom_key = np.arange(num_atoms, dtype=np.int64)
+
+  atom_chain_keys = string_array.remap(chain_id, chain_key_by_chain_id).astype(
+      np.int64
+  )
+
+  try:
+    atom_res_keys = [res_key_by_res[r] for r in zip(chain_id, res_name, res_id)]
+  except KeyError as e:
+    missing_chain_id, missing_res_name, missing_res_id = e.args[0]
+    raise ValueError(
+        'Inconsistent res_name, res_id and all_residues. Could not find '
+        f'residue with chain_id={missing_chain_id}, '
+        f'res_name={missing_res_name}, res_id={missing_res_id} in all_residues.'
+    ) from e
+
+  atoms = Atoms(
+      key=atom_key,
+      chain_key=atom_chain_keys,
+      res_key=np.array(atom_res_keys, dtype=np.int64),
+      name=_default(atom_name, ['?'] * num_atoms, object),
+      element=_default(atom_element, ['?'] * num_atoms, object),
+      x=_default(atom_x, [0.0] * num_atoms, np.float32),
+      y=_default(atom_y, [0.0] * num_atoms, np.float32),
+      z=_default(atom_z, [0.0] * num_atoms, np.float32),
+      b_factor=_default(atom_b_factor, [0.0] * num_atoms, np.float32),
+      occupancy=_default(atom_occupancy, [1.0] * num_atoms, np.float32),
+  )
+  return atoms, residues, chains

flax_model/alphafold3/structure/table.py

@@ -0,0 +1,555 @@
+
+
+"""Table module for atom/residue/chain tables in Structure.
+
+Tables are intended to be lightweight collections of columns, loosely based
+on a pandas dataframe, for use in the Structure class.
+"""
+
+import abc
+from collections.abc import Callable, Collection, Iterable, Iterator, Mapping, Sequence
+import dataclasses
+import functools
+import graphlib
+import typing
+from typing import Any, Protocol, Self, TypeAlias, TypeVar, overload
+
+from flax_model.alphafold3.cpp import string_array
+import numpy as np
+
+
+TableEntry: TypeAlias = str | int | float | None
+FilterPredicate: TypeAlias = (
+    TableEntry
+    | Iterable[Any]  # Workaround for b/326384670. Tighten once fixed.
+    | Callable[[Any], bool]  # Workaround for b/326384670. Tighten once fixed.
+    | Callable[[np.ndarray], bool]
+)
+
+
+class RowLookup(Protocol):
+
+  def get_row_by_key(
+      self,
+      key: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> Mapping[str, Any]:
+    ...
+
+
+@dataclasses.dataclass(frozen=True, kw_only=True)
+class Table:
+  """Parent class for structure tables.
+
+  A table is a collection of columns of equal length, where one column is the
+  key. The key uniquely identifies each row in the table.
+
+  A table can refer to other tables by including a foreign key column, whose
+  values are key values from the other table's key column. These column can have
+  arbitrary names and are treated like any other integer-valued column.
+
+  See the `Database` class in this module for utilities for handing sets of
+  tables that are related via foreign keys.
+
+  NB: This does not correspond to an mmCIF table.
+  """
+
+  key: np.ndarray
+
+  def __post_init__(self):
+    for col_name in self.columns:
+      if (col_len := self.get_column(col_name).shape[-1]) != self.size:
+        raise ValueError(
+            f'All columns should have length {self.size} but got "{col_name}"'
+            f' with length {col_len}.'
+        )
+      self.get_column(col_name).flags.writeable = False  # Make col immutable.
+    if self.key.size and self.key.min() < 0:
+      raise ValueError(
+          'Key values must be non-negative. Got negative values:'
+          f' {set(self.key[self.key < 0])}'
+      )
+    self.key.flags.writeable = False  # Make key immutable.
+
+  def __getstate__(self) -> dict[str, Any]:
+    """Returns members with cached properties removed for pickling."""
+    cached_props = {
+        k
+        for k, v in self.__class__.__dict__.items()
+        if isinstance(v, functools.cached_property)
+    }
+    return {k: v for k, v in self.__dict__.items() if k not in cached_props}
+
+  @functools.cached_property
+  def index_by_key(self) -> np.ndarray:
+    """Mapping from key values to their index in the column arrays.
+
+    i.e.: self.key[index_by_key[k]] == k
+    """
+    if not self.key.size:
+      return np.array([], dtype=np.int64)
+    else:
+      index_by_key = np.zeros(np.max(self.key) + 1, dtype=np.int64)
+      index_by_key[self.key] = np.arange(self.size)
+      return index_by_key
+
+  @functools.cached_property
+  def columns(self) -> tuple[str, ...]:
+    """The names of the columns in the table, including the key column."""
+    return tuple(field.name for field in dataclasses.fields(self))
+
+  @functools.cached_property
+  def items(self) -> Mapping[str, np.ndarray]:
+    """Returns the mapping from column names to column values."""
+    return {col: getattr(self, col) for col in self.columns}
+
+  @functools.cached_property
+  def size(self) -> int:
+    """The number of rows in the table."""
+    return self.key.shape[-1]
+
+  def __len__(self) -> int:
+    return self.size
+
+  def get_column(self, column_name: str) -> np.ndarray:
+    """Gets a column by name."""
+    # Performance optimisation: use the cached columns, instead of getattr.
+    return self.items[column_name]
+
+  def apply_array(self, arr: np.ndarray) -> Self:
+    """Returns a sliced table using a key (!= index) array or a boolean mask."""
+    if arr.dtype == bool and np.all(arr):
+      return self  # Shortcut: No-op, so just return.
+
+    return self.copy_and_update(**{
+        column_name: self.apply_array_to_column(column_name, arr)
+        for column_name in self.columns
+    })
+
+  def apply_index(self, index_arr: np.ndarray) -> Self:
+    """Returns a sliced table using an index (!= key) array."""
+    if index_arr.dtype == bool:
+      raise ValueError('The index array must not be a boolean mask.')
+
+    return self.copy_and_update(
+        **{col: self.get_column(col)[..., index_arr] for col in self.columns}
+    )
+
+  def apply_array_to_column(
+      self,
+      column_name: str,
+      arr: np.ndarray,
+  ) -> np.ndarray:
+    """Returns a sliced column array using a key array or a boolean mask."""
+    if arr.dtype == bool:
+      return self.get_column(column_name)[..., arr]
+    else:
+      return self.get_column(column_name)[..., self.index_by_key[arr]]
+
+  def get_value_by_index(self, column_name: str, index: int) -> Any:
+    return self.get_column(column_name)[index]
+
+  def get_value_by_key(
+      self,
+      column_name: str,
+      key: int | np.integer,
+  ) -> TableEntry:
+    """Gets the value of a column at the row with specified key value."""
+    return self.get_value_by_index(column_name, self.index_by_key[key])
+
+  @overload
+  def __getitem__(self, key: str) -> np.ndarray:
+    ...
+
+  @overload
+  def __getitem__(self, key: np.ndarray) -> 'Table':
+    ...
+
+  @overload
+  def __getitem__(self, key: tuple[str, int | np.integer]) -> TableEntry:
+    ...
+
+  @overload
+  def __getitem__(self, key: tuple[str, np.ndarray]) -> np.ndarray:
+    ...
+
+  def __getitem__(self, key):
+    match key:
+      case str():
+        return self.get_column(key)
+      case np.ndarray() as key_arr_or_mask:
+        return self.apply_array(key_arr_or_mask)
+      case str() as col, int() | np.integer() as key_val:
+        return self.get_value_by_key(col, key_val)
+      case str() as col, np.ndarray() as key_arr_or_mask:
+        return self.apply_array_to_column(col, key_arr_or_mask)
+      case _:
+        if isinstance(key, tuple):
+          err_msg = f'{key}, type: tuple({[type(v) for v in key]})'
+        else:
+          err_msg = f'{key}, type: {type(key)}'
+        raise KeyError(err_msg)
+
+  def get_row_by_key(
+      self,
+      key: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> dict[str, Any]:
+    """Gets the row with specified key value."""
+    return self.get_row_by_index(
+        self.index_by_key[key], column_name_map=column_name_map
+    )
+
+  def get_row_by_index(
+      self,
+      index: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> dict[str, Any]:
+    """Gets the row at the specified index."""
+    if column_name_map is not None:
+      return {
+          renamed_col: self.get_value_by_index(col, index)
+          for renamed_col, col in column_name_map.items()
+      }
+    else:
+      return {col: self.get_value_by_index(col, index) for col in self.columns}
+
+  def iterrows(
+      self,
+      *,
+      row_keys: np.ndarray | None = None,
+      column_name_map: Mapping[str, str] | None = None,
+      **table_by_foreign_key_col: RowLookup,
+  ) -> Iterator[Mapping[str, Any]]:
+    """Yields rows from the table.
+
+    This can be used to easily convert a table to a Pandas dataframe:
+
+    ```py
+    df = pd.DataFrame(table.iterrows())
+    ```
+
+    Args:
+      row_keys: An optional array of keys of rows to yield. If None, all rows
+        will be yielded.
+      column_name_map: An optional mapping from desired keys in the row dicts to
+        the names of the columns they correspond to.
+      **table_by_foreign_key_col: An optional mapping from column names in this
+        table, which are expected to be columns of foreign keys, to the table
+        that the foreign keys point into. If provided, then the yielded rows
+        will include data from the foreign tables at the appropriate key.
+    """
+    if row_keys is not None:
+      row_indices = self.index_by_key[row_keys]
+    else:
+      row_indices = range(self.size)
+    for i in row_indices:
+      row = self.get_row_by_index(i, column_name_map=column_name_map)
+      for key_col, table in table_by_foreign_key_col.items():
+        foreign_key = self[key_col][i]
+        foreign_row = table.get_row_by_key(foreign_key)
+        row.update(foreign_row)
+      yield row
+
+  def with_column_names(
+      self, column_name_map: Mapping[str, str]
+  ) -> 'RenamedTableView':
+    """Returns a view of this table with mapped column names."""
+    return RenamedTableView(self, column_name_map=column_name_map)
+
+  def make_filter_mask(
+      self,
+      mask: np.ndarray | None = None,
+      *,
+      apply_per_element: bool = False,
+      **predicate_by_col: FilterPredicate,
+  ) -> np.ndarray | None:
+    """Returns a boolean array of rows to keep, or None if all can be kept.
+
+    Args:
+      mask: See `Table.filter`.
+      apply_per_element: See `Table.filter`.
+      **predicate_by_col: See `Table.filter`.
+
+    Returns:
+      Either a boolean NumPy array of length `(self.size,)` denoting which rows
+      should be kept according to the input mask and predicates, or None. None
+      implies there is no filtering required, and is used where possible
+      instead of an all-True array to save time and space.
+    """
+    if mask is None:
+      if not predicate_by_col:
+        return None
+      else:
+        mask = np.ones((self.size,), dtype=bool)
+    else:
+      if mask.shape != (self.size,):
+        raise ValueError(
+            f'mask must have shape ({self.size},). Got: {mask.shape}.'
+        )
+      if mask.dtype != bool:
+        raise ValueError(f'mask must have dtype bool. Got: {mask.dtype}.')
+
+    for col, predicate in predicate_by_col.items():
+      if self[col].ndim > 1:
+        raise ValueError(
+            f'Cannot filter by column {col} with more than 1 dimension.'
+        )
+
+      callable_predicates = []
+      if not callable(predicate):
+        if isinstance(predicate, Iterable) and not isinstance(predicate, str):
+          target_vals = predicate
+        else:
+          target_vals = [predicate]
+        for target_val in target_vals:
+          callable_predicates.append(lambda x, target=target_val: x == target)
+      else:
+        callable_predicates.append(predicate)
+
+      field_mask = np.zeros_like(mask)
+      for callable_predicate in callable_predicates:
+        if not apply_per_element:
+          callable_predicate = typing.cast(
+              Callable[[np.ndarray], bool], callable_predicate
+          )
+          predicate_result = callable_predicate(self.get_column(col))
+        else:
+          predicate_result = np.array(
+              [callable_predicate(elem) for elem in self.get_column(col)]
+          )
+        np.logical_or(field_mask, predicate_result, out=field_mask)
+      np.logical_and(mask, field_mask, out=mask)  # Update in-place.
+    return mask
+
+  def filter(
+      self,
+      mask: np.ndarray | None = None,
+      *,
+      apply_per_element: bool = False,
+      invert: bool = False,
+      **predicate_by_col: FilterPredicate,
+  ) -> Self:
+    """Filters the table using mask and/or predicates and returns a new table.
+
+    Predicates can be either:
+      1. A constant value, e.g. `'CA'`. In this case then only rows that match
+        this value for the given column are retained.
+      2. A (non-string) iterable e.g. `('A', 'B')`. In this
+        case then rows are retained if they match any of the provided values for
+        the given column.
+      3. A boolean function e.g. `lambda b_fac: b_fac < 100.0`.
+        In this case then only rows that evaluate to `True` are retained. By
+        default this function's parameter is expected to be an array, unless
+        `apply_per_element=True`.
+
+    Args:
+      mask: An optional boolean NumPy array with length equal to the table size.
+        If provided then this will be combined with the other predicates so that
+        a row is included if it is masked-in *and* matches all the predicates.
+      apply_per_element: Whether apply predicates to each element in the column
+        individually, or to pass the whole column array to the predicate.
+      invert: If True then the returned table will contain exactly those rows
+        that would be removed if this was `False`.
+      **predicate_by_col: A mapping from column name to a predicate. Filtered
+        columns must be 1D arrays. If multiple columns are provided as keyword
+        arguments then each predicate is applied and the results are combined
+        using a boolean AND operation, so an atom is only retained if it passes
+        all predicates.
+
+    Returns:
+      A new table with the desired rows retained (or filtered out if
+      `invert=True`).
+
+    Raises:
+      ValueError: If mask is provided and is not a bool array with shape
+        `(num_atoms,)`.
+    """
+    filter_mask = self.make_filter_mask(
+        mask, apply_per_element=apply_per_element, **predicate_by_col
+    )
+    if filter_mask is None:
+      # No mask or predicate was specified, so we can return early.
+      if not invert:
+        return self
+      else:
+        return self[np.array((), dtype=np.int64)]
+    else:
+      return self[~filter_mask if invert else filter_mask]
+
+  def _validate_keys_are_column_names(self, keys: Collection[str]) -> None:
+    """Raises an error if any of the keys are not column names."""
+    if mismatches := set(keys) - set(self.columns):
+      raise ValueError(f'Invalid column names: {sorted(mismatches)}.')
+
+  def copy_and_update(self, **new_column_by_column_name: np.ndarray) -> Self:
+    """Returns a copy of this table with the specified changes applied.
+
+    Args:
+      **new_column_by_column_name: New values for the specified columns.
+
+    Raises:
+      ValueError: If a specified column name is not a column in this table.
+    """
+    self._validate_keys_are_column_names(new_column_by_column_name)
+    return dataclasses.replace(self, **new_column_by_column_name)
+
+  def copy_and_remap(
+      self, **mapping_by_col: Mapping[TableEntry, TableEntry]
+  ) -> Self:
+    """Returns a copy of the table with the specified columns remapped.
+
+    Args:
+      **mapping_by_col: Each kwarg key should be the name of one of this table's
+        columns, and each value should be a mapping. The values in the column
+        will be looked up in the mapping and replaced with the result if one is
+        found.
+
+    Raises:
+      ValueError: If a specified column name is not a column in this table.
+    """
+    self._validate_keys_are_column_names(mapping_by_col)
+    if not self.size:
+      return self
+    remapped_cols = {}
+    for column_name, mapping in mapping_by_col.items():
+      col_arr = self.get_column(column_name)
+      if col_arr.dtype == object:
+        remapped = string_array.remap(col_arr, mapping)
+      else:
+        remapped = np.vectorize(lambda x: mapping.get(x, x))(col_arr)  # pylint: disable=cell-var-from-loop
+      remapped_cols[column_name] = remapped
+    return self.copy_and_update(**remapped_cols)
+
+
+class RenamedTableView:
+  """View of a table with renamed column names."""
+
+  def __init__(self, table: Table, column_name_map: Mapping[str, str]):
+    self._table = table
+    self._column_name_map = column_name_map
+
+  def get_row_by_key(
+      self,
+      key: int,
+      column_name_map: Mapping[str, str] | None = None,
+  ) -> Mapping[str, Any]:
+    del column_name_map
+    return self._table.get_row_by_key(
+        key, column_name_map=self._column_name_map
+    )
+
+
+_DatabaseT = TypeVar('_DatabaseT', bound='Database')
+
+
+class Database(abc.ABC):
+  """Relational database base class."""
+
+  @property
+  @abc.abstractmethod
+  def tables(self) -> Collection[str]:
+    """The names of the tables in this database."""
+
+  @abc.abstractmethod
+  def get_table(self, table_name: str) -> Table:
+    """Gets the table with the given name."""
+
+  @property
+  @abc.abstractmethod
+  def foreign_keys(self) -> Mapping[str, Collection[tuple[str, str]]]:
+    """Describes the relationship between keys in the database.
+
+    Returns:
+      A map from table names to pairs of `(column_name, foreign_table_name)`
+      where `column_name` is a column containing foreign keys in the table named
+      by the key, and the `foreign_table_name` is the name of the table that
+      those foreign keys refer to.
+    """
+
+  @abc.abstractmethod
+  def copy_and_update(
+      self: _DatabaseT,
+      **new_field_by_field_name: ...,
+  ) -> _DatabaseT:
+    """Returns a copy of this database with the specified changes applied."""
+
+
+def table_dependency_order(db: Database) -> Iterable[str]:
+  """Yields the names of the tables in the database in dependency order.
+
+  This order guarantees that a table appears after all other tables that
+  it refers to using foreign keys. Specifically A < B implies that A contains
+  no column that refers to B.key as a foreign key.
+
+  Args:
+    db: The database that defines the table names and foreign keys.
+  """
+  connections: dict[str, set[str]] = {}
+  for table_name in db.tables:
+    connection_set = set()
+    for _, foreign_table in db.foreign_keys.get(table_name, ()):
+      connection_set.add(foreign_table)
+    connections[table_name] = connection_set
+  yield from graphlib.TopologicalSorter(connections).static_order()
+
+
+def concat_databases(dbs: Sequence[_DatabaseT]) -> _DatabaseT:
+  """Concatenates the tables across a sequence of databases.
+
+  Args:
+    dbs: A non-empty sequence of database instances of the same type.
+
+  Returns:
+    A new database containing the concatenated tables from the input databases.
+
+  Raises:
+    ValueError: If `dbs` is empty or `dbs` contains different Database
+      types.
+  """
+  if not dbs:
+    raise ValueError('Need at least one value to concatenate.')
+  distinct_db_types = {type(db) for db in dbs}
+  if len(distinct_db_types) > 1:
+    raise ValueError(
+        f'All `dbs` must be of the same type, got: {distinct_db_types}'
+    )
+
+  first_db, *other_dbs = dbs
+  concatted_tables: dict[str, Table] = {}
+  key_offsets: dict[str, list[int]] = {}
+  for table_name in table_dependency_order(first_db):
+    first_table = first_db.get_table(table_name)
+    columns: dict[str, list[np.ndarray]] = {
+        column_name: [first_table.get_column(column_name)]
+        for column_name in first_table.columns
+    }
+    key_offsets[table_name] = [
+        first_table.key.max() + 1 if first_table.size else 0
+    ]
+
+    for prev_index, db in enumerate(other_dbs):
+      table = db.get_table(table_name)
+      for col_name in table.columns:
+        columns[col_name].append(table.get_column(col_name))
+      key_offset = key_offsets[table_name][prev_index]
+      offset_key = table.key + key_offset
+      columns['key'][-1] = offset_key
+      if table.size:
+        key_offsets[table_name].append(offset_key.max() + 1)
+      else:
+        key_offsets[table_name].append(key_offsets[table_name][prev_index])
+      for fkey_col_name, foreign_table_name in first_db.foreign_keys.get(
+          table_name, []
+      ):
+        fkey_columns = columns[fkey_col_name]
+        fkey_columns[-1] = (
+            fkey_columns[-1] + key_offsets[foreign_table_name][prev_index]
+        )
+
+    concatted_columns = {
+        column_name: np.concatenate(values, axis=-1)
+        for column_name, values in columns.items()
+    }
+    concatted_tables[table_name] = (type(first_table))(**concatted_columns)
+  return first_db.copy_and_update(**concatted_tables)

flax_model/alphafold3/structure/test_utils.py

@@ -0,0 +1,161 @@
+
+
+"""Utilities for structure module testing."""
+
+import dataclasses
+
+from absl.testing import parameterized
+from flax_model.alphafold3 import structure
+from flax_model.alphafold3.common.testing import data
+import numpy as np
+import tree
+
+
+class StructureTestCase(parameterized.TestCase):
+  """Testing utilities for working with structure.Structure."""
+
+  def assertAuthorNamingSchemeEqual(self, ans1, ans2):  # pylint: disable=invalid-name
+    """Walks naming scheme, making sure all elements are equal."""
+    if ans1 is None or ans2 is None:
+      self.assertIsNone(ans1)
+      self.assertIsNone(ans2)
+      return
+    flat_ans1 = dict(tree.flatten_with_path(dataclasses.asdict(ans1)))
+    flat_ans2 = dict(tree.flatten_with_path(dataclasses.asdict(ans2)))
+    for k, v in flat_ans1.items():
+      self.assertEqual(v, flat_ans2[k], msg=str(k))
+    for k, v in flat_ans2.items():
+      self.assertEqual(v, flat_ans1[k], msg=str(k))
+
+  def assertAllResiduesEqual(self, all_res1, all_res2):  # pylint: disable=invalid-name
+    """Walks all residues, making sure alll elements are equal."""
+    if all_res1 is None or all_res2 is None:
+      self.assertIsNone(all_res1)
+      self.assertIsNone(all_res2)
+      return
+    self.assertSameElements(all_res1.keys(), all_res2.keys())
+    for chain_id, chain_res in all_res1.items():
+      self.assertSequenceEqual(chain_res, all_res2[chain_id], msg=chain_id)
+
+  def assertBioassemblyDataEqual(self, data1, data2):  # pylint: disable=invalid-name
+    if data1 is None or data2 is None:
+      self.assertIsNone(data1)
+      self.assertIsNone(data2)
+      return
+    self.assertDictEqual(data1.to_mmcif_dict(), data2.to_mmcif_dict())
+
+  def assertChemicalComponentsDataEqual(  # pylint: disable=invalid-name
+      self,
+      data1,
+      data2,
+      allow_chem_comp_data_extension,
+  ):
+    """Checks whether two ChemicalComponentData objects are considered equal."""
+    if data1 is None or data2 is None:
+      self.assertIsNone(data1)
+      self.assertIsNone(data2)
+      return
+    if (not allow_chem_comp_data_extension) or (
+        data1.chem_comp.keys() ^ data2.chem_comp.keys()
+    ):
+      self.assertDictEqual(data1.chem_comp, data2.chem_comp)
+    else:
+      mismatching_values = []
+      for component_id in data1.chem_comp:
+        found = data1.chem_comp[component_id]
+        expected = data2.chem_comp[component_id]
+        if not found.extends(expected):
+          mismatching_values.append((component_id, expected, found))
+
+      if mismatching_values:
+        mismatch_err_msgs = '\n'.join(
+            f'{component_id}: {expected} or its extension expected,'
+            f' but {found} found.'
+            for component_id, expected, found in mismatching_values
+        )
+        self.fail(
+            f'Mismatching values for `_chem_comp` table: {mismatch_err_msgs}',
+        )
+
+  def assertBondsEqual(self, bonds1, bonds2, atom_key1, atom_key2):  # pylint: disable=invalid-name
+    """Checks whether two Bonds objects are considered equal."""
+    # An empty bonds table is functionally equivalent to an empty bonds table.
+    # NB: this can only ever be None in structure v1.
+    if bonds1 is None or not bonds1.size or bonds2 is None or not bonds2.size:
+      self.assertTrue(bonds1 is None or not bonds1.size, msg=f'{bonds1=}')
+      self.assertTrue(bonds2 is None or not bonds2.size, msg=f'{bonds2=}')
+      return
+
+    ptnr1_indices1, ptnr2_indices1 = bonds1.get_atom_indices(atom_key1)
+    ptnr1_indices2, ptnr2_indices2 = bonds2.get_atom_indices(atom_key2)
+    np.testing.assert_array_equal(ptnr1_indices1, ptnr1_indices2)
+    np.testing.assert_array_equal(ptnr2_indices1, ptnr2_indices2)
+    np.testing.assert_array_equal(bonds1.type, bonds2.type)
+    np.testing.assert_array_equal(bonds1.role, bonds2.role)
+
+  def assertStructuresEqual(  # pylint: disable=invalid-name
+      self,
+      struc1,
+      struc2,
+      *,
+      ignore_fields=None,
+      allow_chem_comp_data_extension=False,
+      atol=0,
+  ):
+    """Checks whether two Structure objects could be considered equal.
+
+    Args:
+      struc1: First Structure object.
+      struc2: Second Structure object.
+      ignore_fields: Fields not taken into account during comparison.
+      allow_chem_comp_data_extension: Whether to allow data of `_chem_comp`
+        table to differ if `struc2` is missing some fields, but `struc1` has
+        specific values for them.
+      atol: Absolute tolerance for floating point comparisons (in
+        np.testing.assert_allclose).
+    """
+    for field in sorted(structure.GLOBAL_FIELDS):
+      if ignore_fields and field in ignore_fields:
+        continue
+      if field == 'author_naming_scheme':
+        self.assertAuthorNamingSchemeEqual(struc1[field], struc2[field])
+      elif field == 'all_residues':
+        self.assertAllResiduesEqual(struc1[field], struc2[field])
+      elif field == 'bioassembly_data':
+        self.assertBioassemblyDataEqual(struc1[field], struc2[field])
+      elif field == 'chemical_components_data':
+        self.assertChemicalComponentsDataEqual(
+            struc1[field], struc2[field], allow_chem_comp_data_extension
+        )
+      elif field == 'bonds':
+        self.assertBondsEqual(
+            struc1.bonds, struc2.bonds, struc1.atom_key, struc2.atom_key
+        )
+      else:
+        self.assertEqual(struc1[field], struc2[field], msg=field)
+
+    # The chain order within a structure is arbitrary so in order to
+    # directly compare arrays we first align struc1 to struc2 and check that
+    # the number of atoms doesn't change.
+    num_atoms = struc1.num_atoms
+    self.assertEqual(struc2.num_atoms, num_atoms)
+    struc1 = struc1.order_and_drop_atoms_to_match(struc2)
+    self.assertEqual(struc1.num_atoms, num_atoms)
+
+    for field in sorted(structure.ARRAY_FIELDS):
+      if field == 'atom_key':
+        # atom_key has no external meaning, so it doesn't matter whether it
+        # differs between two structures.
+        continue
+      if ignore_fields and field in ignore_fields:
+        continue
+      self.assertEqual(struc1[field] is None, struc2[field] is None, msg=field)
+
+      if np.issubdtype(struc1[field].dtype, np.inexact):
+        np.testing.assert_allclose(
+            struc1[field], struc2[field], err_msg=field, atol=atol
+        )
+      else:
+        np.testing.assert_array_equal(
+            struc1[field], struc2[field], err_msg=field
+        )

flax_model/alphafold3/version.py

@@ -0,0 +1,13 @@
+# Copyright 2024 DeepMind Technologies Limited
+#
+# AlphaFold 3 source code is licensed under CC BY-NC-SA 4.0. To view a copy of
+# this license, visit https://creativecommons.org/licenses/by-nc-sa/4.0/
+#
+# To request access to the AlphaFold 3 model parameters, follow the process set
+# out at https://github.com/google-deepmind/alphafold3. You may only use these
+# 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
+
+"""Single source of truth for the AlphaFold version."""
+
+__version__ = '3.0.1'