169 lines
6.1 KiB
Python
169 lines
6.1 KiB
Python
# setup_pymol.py - PyMOL 自定义设置脚本
|
||
|
||
from pymol import cmd
|
||
|
||
def setup_rendering(bg_color='white', transparency=0.5, ray_shadow=False):
|
||
"""设置 PyMOL 的背景、透明度和渲染优化"""
|
||
cmd.bg_color(bg_color)
|
||
|
||
if ray_shadow:
|
||
cmd.set('ray_shadows', 'on')
|
||
cmd.set('ray_shadow', 0.2)
|
||
else:
|
||
cmd.set('ray_shadows', 'off')
|
||
cmd.set('ray_shadow', 0)
|
||
|
||
cmd.set('transparency_mode', 2)
|
||
cmd.set('transparency', transparency)
|
||
cmd.set('surface_quality', 1)
|
||
|
||
# 移除了 surface_color_by_atom，因为某些版本不支持
|
||
# 使用 color 命令可以达到类似效果
|
||
|
||
print(f"✓ 渲染设置完成: 背景={bg_color}, 透明度={transparency}, 阴影={'开' if ray_shadow else '关'}")
|
||
|
||
|
||
def clean_structure(remove_water=True, remove_metals=True, remove_ions=True):
|
||
"""
|
||
清理结构中的水分子、金属离子和其他离子
|
||
|
||
参数:
|
||
remove_water (bool): 是否删除水分子，默认 True
|
||
remove_metals (bool): 是否删除金属离子，默认 True
|
||
remove_ions (bool): 是否删除其他常见离子，默认 True
|
||
"""
|
||
removed_items = []
|
||
|
||
# 1. 删除水分子
|
||
if remove_water:
|
||
water_count = cmd.count_atoms('resn HOH+WAT+H2O+DOD+D2O')
|
||
if water_count > 0:
|
||
cmd.remove('resn HOH+WAT+H2O+DOD+D2O')
|
||
removed_items.append(f'水分子 ({water_count} 个原子)')
|
||
|
||
# 2. 删除常见金属离子
|
||
if remove_metals:
|
||
metals = [
|
||
'NA', 'K', 'MG', 'CA', 'ZN', 'FE', 'MN', 'CU', 'NI', 'CO',
|
||
'CD', 'HG', 'PB', 'AG', 'AU', 'PT', 'CR', 'MO', 'W', 'V'
|
||
]
|
||
metal_selection = '+'.join(metals)
|
||
metal_count = cmd.count_atoms(f'resn {metal_selection}')
|
||
if metal_count > 0:
|
||
cmd.remove(f'resn {metal_selection}')
|
||
removed_items.append(f'金属离子 ({metal_count} 个原子)')
|
||
|
||
# 3. 删除其他常见离子（氯离子、硫酸根等）
|
||
if remove_ions:
|
||
ions = ['CL', 'BR', 'I', 'F', 'SO4', 'PO4', 'NO3', 'NH4']
|
||
ion_selection = '+'.join(ions)
|
||
ion_count = cmd.count_atoms(f'resn {ion_selection}')
|
||
if ion_count > 0:
|
||
cmd.remove(f'resn {ion_selection}')
|
||
removed_items.append(f'其他离子 ({ion_count} 个原子)')
|
||
|
||
if removed_items:
|
||
print("✓ 结构清理完成:")
|
||
for item in removed_items:
|
||
print(f" - 已删除 {item}")
|
||
else:
|
||
print("✓ 未发现需要清理的分子")
|
||
|
||
|
||
def create_colored_objects(protein_color='orange', nucleic_color='white', ligand_color='red'):
|
||
"""创建蛋白质、核酸和小分子对象，并应用指定颜色"""
|
||
if len(cmd.get_object_list()) == 0:
|
||
print("⚠ 警告: 未检测到已加载的结构，请先加载 PDB 文件")
|
||
return
|
||
|
||
obj_created = []
|
||
|
||
# 1. 创建蛋白质对象
|
||
if cmd.count_atoms('polymer.protein') > 0:
|
||
cmd.create('Protein_obj', 'polymer.protein')
|
||
cmd.color(protein_color, 'Protein_obj')
|
||
cmd.show('surface', 'Protein_obj')
|
||
obj_created.append(f'Protein_obj ({protein_color}, {cmd.count_atoms("Protein_obj")} 原子)')
|
||
|
||
# 2. 创建核酸对象
|
||
if cmd.count_atoms('polymer.nucleic') > 0:
|
||
cmd.create('Nucleic_obj', 'polymer.nucleic')
|
||
cmd.color(nucleic_color, 'Nucleic_obj')
|
||
cmd.show('surface', 'Nucleic_obj')
|
||
obj_created.append(f'Nucleic_obj ({nucleic_color}, {cmd.count_atoms("Nucleic_obj")} 原子)')
|
||
|
||
# 3. 创建小分子对象（排除水分子和离子）
|
||
ligand_selection = 'organic and not (resn HOH+WAT+H2O+DOD+D2O)'
|
||
if cmd.count_atoms(ligand_selection) > 0:
|
||
cmd.create('Ligand_obj', ligand_selection)
|
||
cmd.color(ligand_color, 'Ligand_obj')
|
||
cmd.show('surface', 'Ligand_obj')
|
||
obj_created.append(f'Ligand_obj ({ligand_color}, {cmd.count_atoms("Ligand_obj")} 原子)')
|
||
|
||
# 缩放到所有对象
|
||
cmd.zoom('all')
|
||
|
||
if obj_created:
|
||
print("✓ 对象创建完成:")
|
||
for obj in obj_created:
|
||
print(f" - {obj}")
|
||
else:
|
||
print("⚠ 未创建任何对象，请检查结构类型")
|
||
|
||
|
||
def quick_setup(bg_color='white', transparency=0.5,
|
||
protein_color='orange', nucleic_color='white', ligand_color='red',
|
||
clean=True):
|
||
"""
|
||
一键完成所有设置：清理结构 + 渲染优化 + 对象创建与着色
|
||
|
||
参数:
|
||
bg_color (str): 背景颜色
|
||
transparency (float): 透明度 (0-1)
|
||
protein_color (str): 蛋白质颜色
|
||
nucleic_color (str): 核酸颜色
|
||
ligand_color (str): 小分子颜色
|
||
clean (bool): 是否清理水分子和离子，默认 True
|
||
|
||
使用示例:
|
||
quick_setup() # 默认设置+清理
|
||
quick_setup(clean=False) # 不清理
|
||
quick_setup('black', 0.3, 'cyan', 'yellow', 'magenta')
|
||
"""
|
||
print("\n" + "="*60)
|
||
print("开始 PyMOL 快速设置...")
|
||
print("="*60)
|
||
|
||
# 1. 清理结构（如果启用）
|
||
if clean:
|
||
clean_structure()
|
||
|
||
# 2. 设置渲染
|
||
setup_rendering(bg_color, transparency)
|
||
|
||
# 3. 创建并着色对象
|
||
create_colored_objects(protein_color, nucleic_color, ligand_color)
|
||
|
||
print("\n" + "="*60)
|
||
print("✓ PyMOL 快速设置完成！")
|
||
print("="*60 + "\n")
|
||
|
||
|
||
# 注册函数到 PyMOL 的全局命名空间
|
||
cmd.extend('setup_rendering', setup_rendering)
|
||
cmd.extend('clean_structure', clean_structure)
|
||
cmd.extend('create_colored_objects', create_colored_objects)
|
||
cmd.extend('quick_setup', quick_setup)
|
||
|
||
print("\n" + "="*60)
|
||
print("PyMOL 自定义函数已加载")
|
||
print("\n可用命令:")
|
||
print(" setup_rendering - 设置背景和渲染")
|
||
print(" clean_structure - 清理水分子和离子")
|
||
print(" create_colored_objects - 创建并着色对象")
|
||
print(" quick_setup - 一键完成所有设置")
|
||
print("\n使用示例:")
|
||
print(" quick_setup() # 默认设置（包含清理）")
|
||
print(" quick_setup(clean=False) # 不清理水和离子")
|
||
print(" clean_structure() # 单独清理")
|
||
print("="*60 + "\n") |