乐于分享
好东西不私藏

Pymol使用之零依赖的序列设计插件PymolMPNN

Pymol使用之零依赖的序列设计插件PymolMPNN

最近写报告写得头大,就想着看点别的换换脑子,看到“投机解码”的时候突然想到,如果LLM可以这样提高输出效率,那蛋白设计工具能不能也这么搞。于是我就问龙虾(Opus5),ProteinMPNN可以投机解码或者蒸馏吗,让它算得快一点或者权重文件小一点,最好在CPU上能一秒一条。

结果龙虾回答说:MPNN效率本身已经很高了,在CPU上一秒一条的要求原始版本的MPNN(vanilla MPNN)上应该就能达到。(因为我之前和龙虾讨论过Pymol内建MPNN功能)然后龙虾接着说,如果非说MPNN使用有什么不便的话,可能主要是部署,因为它依赖PyTorch,“在Pymol里内嵌一个2GBtorch可能是灾难”。

接下来龙虾又话锋一转——“但是我有个好消息,我在拆解开源Pymol的时候发现numpyPymol的硬依赖,而MPNN只有~1.7M参数,numpy重写整个forward完全可行”,这样下来的话整个插件只需要十几MB,而且零外部依赖。

其实我前几天刷到那个HaloMPNN工作的时候就有构建torch-free版MPNN的想法,这又一次听到龙虾说这事可行,我就让龙虾写了个执行清单,看看是不是真的可行。龙虾写完之后我没当回事就先放在了一边。

又过了两天,我写报告写得有点力竭了,突然想起这个清单。我就想,我写这么累不能让ai闲着,于是就给codex充了20块钱,这样我写报告它写插件。

最后呢,就得到了文末的零外部依赖的Pymol序列设计插件——PymolMPNN。故有此篇。

如果有想先尝试再了解的读者,可以直奔文末,下载体验PymolMPNN。

以下是PymolMPNN开发的全过程:

1.可行性分析

操作最开始我先把龙虾写的清单复制给了codex5.6 Sol Medium),问它思路是否可行,codex表示了认可——“整体方案可行,而且“先导出权重和 golden reference,再逐层 NumPy 对齐,最后封装 PyMOL”的路线是正确的。”,同时提出了一个快速的验证路线:

    1.锁定commit和v_48_020。2.打印权重、模型配置和参数量。3.插桩确认sample()每一步的真实shape。4.分别测feature、encoder、decoder时间。5.导出一个最小1UBQ trace。6.在目标PyMOL里确认Python、NumPy和Qt版本。7.用一个NumPy脚本只复现最终Linear和一个decoder layer,测量小矩阵性能。

    2.确定项目目标

    看到codex分析的头头是道,我就让他按照我之前写的开源pymol安装教程从头构建环境安装pymol用于测试,然后根据机器配置把CPU版本的ProteinMPNN也装好,然后写一个流程规划给我确认。其实开始这里,我对于codex能不能胜任MPNN重写的任务并没有什么把握,所以就先让它配环境装程序试试,如果连程序都装不好,那重写估计也够呛,如果可以,那重写就有些希望。然后codex work26分钟之后说,环境已经配好了,程序也已经装好了(ProtienMPNN源码和Pymol whl文件是我提供的),而且很有主见地分别创建了Pymol调试环境py313forpymol和MPNN的调试环境proteinmpnn-dev避免相互污染,接着给出了整个流程规划:

    Step 0:环境和源码固定

    1.验证:import pymolPyMOL GUI 启动PyQt5 插件加载ffmpeg 可调用numpy 可用2.获取官方 ProteinMPNN 源码并固定 commit。3.确认 v_48_020.pt 是否存在并打印完整参数 shape 清单。4.插桩官方 sample(),确认 decoder 是否已经是位置级增量计算。

    Step 1:权重导出和 golden reference

    产出:mpnn_v48_020.npz模型配置 manifest1UBQ golden reference中等长度单体 golden reference多链复合物 golden reference先只做确定性 argmax,不做随机采样一致性承诺。

    Step2:NumPy 几何特征

    严格按固定 commit 中的 ProteinFeatures 实现,逐项验证:虚拟Cβ距离与RBFKNN/E_idx相对位置编码edge embeddingnormalization每一步都和Torch reference比较,单独定位误差。

    Step 3:NumPy 网络 forward

    实现:LinearLayerNorm官方实际使用的激活函数Encoder message passingDecoder layerlogits先完成朴素等价实现,再根据 profiling 决定是否需要额外优化。官方实现如果已经按单个位置解码,就不再重复设计所谓“增量解码”。

    Step4:端到端验证

    验收分层进行:中间张量误差;logits 数值误差;argmax 序列逐位一致;同一预生成随机数流下的采样一致;性能 benchmark:L=50、100、300。

    Step 5:PyMOL 插件

    一期建议支持:标准 20 种氨基酸;完整 N/CA/C/O backbone;单对象和多链;chain breakfixed positions;random seed;temperature;设计结果输出;每位概率/entropy 写入副本对象的 B-factor一期暂缓:tied positions;PSSM/bias;非标准残基;缺失原子自动修复;空间并行近似解码。

    Step 6:打包与发布

    最终结构:

    proteinmpnn_plugin/├── __init__.py├── mpnn_numpy.py├── features.py├── weights/│   ├── mpnn_v48_020.npz│   └── manifest.json├── ui/└── README.md

    目标是“不依赖 PyTorch”,但仍依赖 PyMOL 自带的 Python 和 NumPy;需要在当前 PyMOL 版本上实际验证,不能提前承诺兼容所有冻结版。

    3. 开整!

    我看了看它写的步骤,也没看出不合理的地方,就让它继续。Codex工作了157之后,检查好了源码和权重。

    还修订了执行流程,然后就找我确认:

    step1前的确认

    我确认之后,codex工作了2218step1  step2 的权重导出和Numpy几何特征实现就都完成了:

    我继续ok之后1758,codex把第三步decoder的Numpy实现也搞定了,还修复了Conda 的 NumPy 2.5.1 + MKL 2026 在 Python 3.13/PyMOL 环境中触发原生异常的问题。

    到这一步其实用Numpy重写ProteinMPNN的目标基本已经实现了。下一步应该是封装插件。codex主动提出:现在的性能目标还没达到,“下一步应先做性能优化,再完善插件功能。”

    咱也不太懂,就只能让它继续。4分钟之后,它把性能提升了大概25%

    并且继续完成了固定残基序列设计的功能:

    看它这么给力,我决定彻底开启自动驾驶,让它一直推进到需要我来验证为止:

    在连续工作了3946之后,codex完成了所有可以自动处理的部分:

    截止到这里,作为Pymol插件的PymolMPNN已经基本成型。但很多时候,真是上手了才能意识到最确切的需求。我安装验证的时候想当然地觉得已经实现的功能是把设计的序列映射到结构,但运行了两遍发现侧链都没变,才意识到插件的功能还停留在设计序列的直接输出。

    序列设计最直观的输出,还是应该直接映射到结构上,于是放下大功告成的高兴,继续和codex对齐:

    codex一下就理解了我的想法,工作了552之后,实现了需求:

    此时的版本号是v1.2,之后发现调用pymol自带的Mutagenesis时会在某些情况会报错,于是修改了几版之后版本号水涨船高涨到了v1.4。

    截至这里,在codex说一不二的执行下,我们就实现了最初的目标,完成了把ProteinMPNN依赖的PyTorch算子用Numpy重写,Pymol插件封装,调用Pymol Mutagenesis重建相应位置的侧链用于检视的完整流程。

    这时候就回到了我最开始看到HaloMPNN工作时候的设想,那能不能把其他的MPNN工作也打包进来作为候选权重,比如SolubleMPNN、ThermoMPNN,LigandMPNN和HaloMPNN等。结果codex研究了一下说,SolubleMPNN是ProteinMPNN自带的权重,我们现在已经支持了,ThermoMPNN是做稳定性预测的,不涉及序列设计,而LigandMPNN因为考虑了配体信息,这个模型的图结构和ProteinMPNN有很大差别,所以不好兼容。HaloMPNN的信息codex一直没查到,不过现在已经有了默认版和Soluble版可选,作为一个初步版本也可以接受。

    既然权重已经能有选项,下面就是完善PymolMPNN的功能,和ProteinMPNN的功能对齐。其实主要就是把omit AA 和bias AA加上,实现对特定残基的排除或加权。再加上默认MPNN(vanilla MPNN)和SolubleMPNN的选项,最终得到了现在的v1.8版本。

    命令格式为:

    mpnn_design selection, output_object, fixed_selection, temperature, seed, pdb_path, fasta_path, omit_aas, bias, bias_by_res, model

    GUI界面为:

    迭代了好几版的GUI

    4. 搞定

    至此,PymolMPNN已经实现了原版主要功能。支持了这么多功能之后,这个插件(包含权重)现在多大呢?答案是11.92MB^_^ 我认为完全符合我最初对轻量的预期。

    当然,如果要进行大规模、批量的序列设计,可能还是需要专门的服务器和原版的ProteinMPNN。但是我觉得PymolMPNN可以作为一种低门槛体验蛋白序列设计的方式(只需要有Pymol,不需要配置任何环境),也可以用来做一些小规模设计,或者在特定位点进行残基突变时提供一个序列层面的参考。

    总之,欢迎大家取用,也欢迎专业的同学来审计codex的代码,有任何使用上的建议、意见和bug也欢迎大家反馈~
    全文完。
    插件下载地址:
    https://github.com/guozheng-git/pymol-scripts/blob/main/PymolMPNN_v1.8.zip
    如果有读者不方便从github下载,也可以在后台发送“PymolMPNN”获取插件压缩包。
    (最近也在了解工作和实习机会,如有合适机会,也欢迎后台联系^_^)