Pymol Iterate函数工作原理及输出处理与CSV生成问题

一、Iterate 核心工作机制 #

Iterate 是 PyMOL 专门用于遍历原子选择集的工具函数，它会逐个遍历你指定的原子（或对象），将每个原子作为上下文，执行你给定的代码逻辑。可以理解为给每个原子单独“运行一段小代码”，遍历完成后就能批量处理整个选择集的原子数据。

二、用 Iterate 执行自定义函数 #

不管是获取原子名称、计算二面角还是其他自定义逻辑，都可以通过 Iterate 批量执行。关键是把自定义函数注入到 Iterate 的执行上下文里，示例如下：

# 自定义函数1：获取原子基础信息
def get_atom_details(atom):
    # 返回包含残基名、残基ID、原子名、原子ID的元组
    return (atom.resn, atom.resi, atom.name, atom.id)

# 自定义函数2：计算目标二面角（示例以当前原子为中心关联其他三个原子）
def compute_dihedral(atom):
    # 构建包含四个原子的选择表达式（需根据实际需求调整）
    target_sel = f"{atom.id} and (name N or name CA or name C or name O)"
    if cmd.count_atoms(target_sel) == 4:
        return cmd.get_dihedral(target_sel)
    # 不满足条件时返回空值
    return None

# 初始化存储结果的列表
output_list = []

# 用 Iterate 遍历"all"选择集，执行get_atom_details函数
cmd.iterate(
    selection="all",
    expression="output_list.append(get_atom_details(atom))",
    space={"get_atom_details": get_atom_details, "output_list": output_list}
)

# 针对特定选择集（比如丙氨酸的CA原子）执行二面角计算
cmd.iterate(
    selection="resn ALA and name CA",
    expression="output_list.append((atom.id, compute_dihedral(atom)))",
    space={"compute_dihedral": compute_dihedral, "output_list": output_list}
)

解释：

• selection：指定要遍历的原子集，支持PyMOL所有选择表达式
• expression：要执行的代码字符串，这里是把函数返回的元组追加到结果列表
• space：传递外部函数和变量的字典，让Iterate能识别你的自定义逻辑

三、将结果存入列表并生成CSV #

Iterate 每次执行会把函数返回的元组追加到你定义的列表中，直接用Python的csv模块就能生成指定格式的CSV：

import csv

# 假设第一个自定义函数的结果为主，定义CSV列名
csv_columns = ["Residue_Name", "Residue_ID", "Atom_Name", "Atom_ID"]

# 写入CSV文件
with open("pymol_atom_data.csv", "w", newline="", encoding="utf-8") as csv_file:
    writer = csv.writer(csv_file)
    # 写入表头
    writer.writerow(csv_columns)
    # 批量写入所有元组数据
    writer.writerows(output_list)

如果你的列表里混合了不同结构的元组（比如原子信息和二面角数据），建议分开存储后再分别写入CSV，避免格式混乱。

四、为什么必须用字典传递函数？ #

PyMOL 的 iterate 函数执行的代码字符串，是在PyMOL内部的独立命名空间中运行的，而你在脚本里定义的自定义函数、变量默认属于脚本的本地命名空间——这两个空间是完全隔离的。

用space参数传递字典，相当于把脚本里的函数、变量“注入”到Iterate的执行上下文里，让代码字符串能识别到这些外部定义的内容。如果不传递，Iterate执行时会直接报错“名称未定义”，因为它找不到你写的自定义函数。

简单来说：字典是连接脚本本地命名空间和PyMOL内部执行空间的唯一桥梁。

内容的提问来源于stack exchange，提问作者Riccardo Fusco