AF3 unify_alignment_db_indices脚本解读

AlphaFold3 unify_alignment_db_indices 脚本在源代码的scripts/alignment_db_scripts文件夹下。该脚本是统一多个对齐数据库索引文件（.index），生成一个总的索引文件 super.index，便于模型后续统一查找对齐数据。即把由多个 create_alignment_db.py 脚本生成的 .index 索引文件合并成一个统一的大索引。

源代码：

import argparse
import json
import os
""" Unifies databases created with create_alignment_db.py """def main(args):super_index = {}for f in os.listdir(args.alignment_db_dir):if not os.path.splitext(f)[-1] == ".index":continuewith open(os.path.join(args.alignment_db_dir, f), "r") as fp:index = json.load(fp)db_name = f"{os.path.splitext(f)[0]}.db"for k in index:super_index[k] = {"db": db_name,"files": index[k],}with open(os.path.join(args.output_dir, "super.index"), "w") as fp:json.dump(super_index, fp)if __name__ == "__main__":parser = argparse.ArgumentParser()parser.add_argument("alignment_db_dir", type=str, help="Path to directory containing alignment_dbs")parser.add_argument("output_dir", type=str, help="Path in which to output super index")args = parser.parse_args()main(args)

代码解读：

核心逻辑：main(args)

def main(args):super_index = {}for f in os.listdir(args.alignment_db_dir):if not os.path.splitext(f)[-1] == ".index":continue

• 遍历指定目录中所有文件

• 筛选出 .index 结尾的索引文件（忽略其他文件，如 .db 数据文件）

加载每个 .index 并合并到 super_index

        with open(os.path.join(args.alignment_db_dir, f), "r") as fp:index = json.load(fp)db_name = f"{os.path.splitext(f)[0]}.db"for k in index:super_index[k] = {"db": db_name,"files": index[k],}

• 读取单个 .index 文件，内容是 {chain_id: {"files": [...]}}

• 从文件名推测该 .index 对应的 .db 数据文件名（例如 align_0.index → align_0.db）

• 构建统一格式：{chain_id: {"db": ..., "files": ...}}

• 并合并到一个大的 super_index 中

写出 super.index

    with open(os.path.join(args.output_dir, "super.index"), "w") as fp:json.dump(super_index, fp)

• 将所有链的统一索引写入新的文件 super.index

• 格式为 JSON，内容类似：

• {"1abc_A": {"db": "align_0.db","files": [["file1.a3m", 0, 1234],["file2.a3m", 1234, 5678]]},...
  }
  

  总结脚本功能

  步骤	内容
  输入	一个目录，包含多个 .index 和 .db 文件（由多个分片 create_alignment_db 脚本生成）
  操作	遍历所有 .index 文件，合并它们
  输出	一个统一的 super.index 文件，记录所有链的对齐数据在哪个 .db 中、偏移和大小
  用途	模型/下游程序可以只加载 super.index 即可定位所有链的对齐数据位置

和create_alignment_db_sharded.py脚本比较