ChrisHayduk/minAlphaFold2

解决的问题

minAlphaFold2 提供了一个极简且具教学意义的 AlphaFold2 PyTorch 实现。它通过移除生产环境级别的复杂性（如 JAX 优化代码或繁重的架构），并以代码与原始 62 页补充论文之间清晰的 1:1 对应关系取代，解决了理解原始系统的困难。

运作方式

本项目使用纯 PyTorch 原语（nn.Linear、torch.einsum 等）实现 AlphaFold2 架构，无需外部机器学习库。它完全遵循论文规范，包括：

• 架构： 所有 32 个补充算法皆实现为独立模块，从 Evoformer 堆叠到结构模块 (Structure Module)。
• 训练： 采用两阶段训练配方（初始与微调），利用梯度累积在单一 GPU 上模拟大规模 TPU 批次。
• 数据管线： 整合 OpenProteinSet 进行 MSA 与模板处理，包含针对分辨率与序列组成的确定性筛选。
• 优化： 忠实移植迭代约束 Amber 最小化，用于放松预测结构。

适用对象

• 希望学习 AlphaFold2 内部运作机制，但不想钻研生产环境代码库的 AI 研究人员与学生。
• 寻求精简、可训练的参考实现，以便进行分支与调整以进行蛋白质折叠研究的开发者。
• 具备编程背景，希望以易读方式进入 AI 驱动结构生物学领域的生物学家。

特色亮点

• 论文与代码对应： 每个模块皆对应 AF2 补充论文中的编号算法。
• 纯 PyTorch： 无自定义 CUDA 内核或 einops，使代码高度可读且标准化。
• 单 GPU 可训练： 使用梯度检查点与累积技术，让 48 层 Evoformer 块可在消费级或云端 GPU 上运行。
• 端到端管线： 包含数据下载、预处理、训练与结构放松脚本。