274 lines
9.7 KiB
Python
274 lines
9.7 KiB
Python
#!/usr/bin/env python
|
||
# -*- encoding: utf-8 -*-
|
||
'''
|
||
@file :pymol_color_plugin.py
|
||
@Description: :pymol配色方案
|
||
@Date :2023/8/29 10:52:25
|
||
@Author :hotwa
|
||
@version :1.0
|
||
'''
|
||
|
||
|
||
from pymol import cmd
|
||
from pathlib import Path
|
||
from dataclasses import dataclass
|
||
from json import dumps
|
||
from Bio.Data import IUPACData
|
||
from Bio.SeqUtils import seq3
|
||
from collections import defaultdict
|
||
|
||
STANDARD_AMINO_ACIDS_1 = set(IUPACData.protein_letters)
|
||
STANDARD_AMINO_ACIDS_3 = set([seq3(i).upper() for i in STANDARD_AMINO_ACIDS_1])
|
||
|
||
|
||
def bejson(d:dict(help='need to beauty json')) -> dict:
|
||
return dumps(d,indent=4,ensure_ascii=False)
|
||
|
||
@dataclass
|
||
class pymolbase():
|
||
file: Path
|
||
|
||
def __post_init__(self):
|
||
cmd.reinitialize('everything')
|
||
cmd.load(filename=self.file.as_posix())
|
||
PDBs = cmd.get_names()
|
||
if len(PDBs) == 1:
|
||
PDB = PDBs[0]
|
||
else:
|
||
raise ValueError(f'this pdb have more than one object! PDBs:{PDBs}')
|
||
self.CAindex = cmd.identify(f"{PDB} and name CA") # get CA index
|
||
self.pdbstrList = [cmd.get_pdbstr("%s and id %s" % (PDB, CAid)).splitlines() for CAid in CAindex]
|
||
self.ProtChainResiList = [[PDB, i[0][21], i[0][22:26].strip()] for i in pdbstrList] # get pdb chain line string
|
||
resn_hetatm = self.get_resn('hetatm')
|
||
|
||
def mutate(self, site:str, chain:str, mutation_type: str):
|
||
'''
|
||
mutation_type:ALA,ARG,ASN,ASP,CYS,GLN,GLU,GLY,HIS,ILE,LEU,LYS,MET,PHE,PRO,SER,THR,TRP,TYR,VAL
|
||
'''
|
||
for i, j, k in self.ProtChainResiList:
|
||
if int(k) == int(site) and j == str(chain):
|
||
cmd.wizard("mutagenesis")
|
||
cmd.refresh_wizard()
|
||
cmd.get_wizard().set_mode(mutation_type)
|
||
selection = f"/{i}//{j}/{k}"
|
||
cmd.get_wizard().do_select(selection)
|
||
cmd.get_wizard().apply()
|
||
cmd.set_wizard("done")
|
||
|
||
def add_polar_hydrogen(self, pymol_obj='sele'):
|
||
'''
|
||
add polar hydrogen
|
||
'''
|
||
cmd.h_add(pymol_obj) # add all hydrogens in this molecular
|
||
cmd.remove(f"{pymol_obj} & hydro & not nbr. (don.|acc.)") # remove no-polar-hydrogen
|
||
|
||
def save(self, file_name: Path, selection: str):
|
||
cmd.save(file_name.as_posix(), f"{selection}")
|
||
|
||
# def get_resn(self, name:str = 'all'):
|
||
# resn_list = []
|
||
# if name == 'protein':
|
||
# cmd.iterate(f"all and not hetatm", expression=lambda atom: resn_list.append(atom.resn))
|
||
# return tuple(set(resn_list))
|
||
# else:
|
||
# cmd.iterate(selection=name, expression=lambda atom: resn_list.append(atom.resn))
|
||
# return tuple(set(resn_list))
|
||
|
||
def get_resn(self, name:str = 'all'): # ! not test
|
||
resn_dict = defaultdict(set)
|
||
|
||
def collect_resn(atom):
|
||
category = 'protein' if atom.chain and not atom.hetatm else 'small_molecule'
|
||
resn_dict[category].add(atom.resn)
|
||
|
||
cmd.iterate(name, expression=collect_resn)
|
||
|
||
protein_resns = resn_dict.get('protein', set())
|
||
small_mol_resns = resn_dict.get('small_molecule', set())
|
||
|
||
standard_aa = protein_resns.intersection(STANDARD_AMINO_ACIDS_3)
|
||
non_standard_aa = protein_resns - STANDARD_AMINO_ACIDS_3
|
||
|
||
return {
|
||
'small_molecules': list(small_mol_resns),
|
||
'protein_resns': list(protein_resns),
|
||
'standard_aa': list(standard_aa),
|
||
'non_standard_aa': list(non_standard_aa)
|
||
}
|
||
|
||
class colorp(pymolbase):
|
||
color1 = (('color1', '[186,182,217]'), 'purple')
|
||
color2 = (('color2', '[233,195,153]'), 'yellow')
|
||
color3 = (('color3', '[43,113,216]'), 'blue-N')
|
||
color4 = (('color4', '[206,155,198]'), 'purple')
|
||
color5 = (('color5', '[251,187,62]'), 'orange')
|
||
color6 = (('color6', '[245,157,158]'), 'red')
|
||
color7 = (('color7', '[133,188,135]'), 'green')
|
||
color8 = (('color8', '[30,230,30]'),'green-CL') # Cl卤素配色
|
||
color9 = (('color9', '[141,215,247]'),'blue-C') # C配色
|
||
color10 = (('color10', '[0,132,55]'),'green-F') # F卤素配色
|
||
grey1 = ('grey1','[224,224,224]')
|
||
colors = (color1, color2, color3, color4, color5, color6, color7, color8, color9, color10)
|
||
|
||
def __init__(self, path, id):
|
||
self.id = id.lower()
|
||
cmd.reinitialize()
|
||
p = Path(path)
|
||
if p.is_dir():
|
||
file = p.joinpath(f"{id}.pdb")
|
||
else:
|
||
raise ValueError('path params error')
|
||
cmd.load(file,id)
|
||
|
||
def pretty(self):
|
||
cmd.remove('solvent metals') # 移除金属离子和溶剂
|
||
cmd.remove("resn SO4 or resn PO4 or resn CL")
|
||
cmd.pretty(selection=self.id)
|
||
|
||
@staticmethod
|
||
def defcolor():
|
||
# 定义常用配色
|
||
color_gennerate = map(lambda x:x[0],colorp.colors)
|
||
list(map(lambda x:cmd.set_color(*x),color_gennerate))
|
||
# 定义灰度配色
|
||
cmd.set_color(*colorp.grey1)
|
||
|
||
@staticmethod
|
||
def grey(obj='sele',opt=True):
|
||
colorp.defcolor()
|
||
if opt: method.optimisation() # 优化
|
||
# 对某个对象进行灰化
|
||
cmd.color('grey1',f'{obj}')
|
||
|
||
@staticmethod
|
||
def color_mole(obj='hetatm'):
|
||
# color blue,sele in name c*
|
||
colorp.defcolor()
|
||
cmd.color('color9',f'{obj} in name c*')
|
||
cmd.color('color6',f'{obj} in name o*')
|
||
cmd.color('color5',f'{obj} in name s*')
|
||
cmd.color('color3',f'{obj} in name n*')
|
||
cmd.color('color8',f'{obj} in name cl*')
|
||
cmd.color('color10',f'{obj} in name f*')
|
||
|
||
@staticmethod
|
||
def pretty(obj = None):
|
||
if not obj: obj = cmd.get_names()[0]
|
||
cmd.remove('solvent metals') # 移除金属离子和溶剂
|
||
cmd.pretty(selection=obj)
|
||
|
||
|
||
class font():
|
||
font_size = 28 # 单位就是正常的px。你也可以用负值，则单位是Å
|
||
|
||
class mole():
|
||
stick_radius=0.10
|
||
|
||
class atom():
|
||
size = 0.28
|
||
|
||
class method():
|
||
|
||
@staticmethod
|
||
def optimisation(selection, grey=True):
|
||
colorp.defcolor()
|
||
if grey: cmd.set('cartoon_color','grey1', selection=selection)
|
||
cmd.set('stick_transparency','0.1', selection=selection)
|
||
cmd.set("ray_shadows","off", selection=selection)
|
||
cmd.set('cartoon_highlight_color', selection=selection)
|
||
cmd.set('cartoon_fancy_helices', selection=selection)
|
||
cmd.set('cartoon_transparency','0.5', selection=selection)
|
||
|
||
|
||
@staticmethod
|
||
def remove():
|
||
cmd.remove('resn SO4') # SO4 remove
|
||
cmd.remove('solvent metals') # 移除金属离子和溶剂
|
||
|
||
def get_resn(name:str = 'all'):
|
||
resn_list = []
|
||
cmd.iterate(selection=name, expression=lambda atom: resn_list.append(atom.resn))
|
||
return tuple(set(resn_list))
|
||
|
||
def protgrey(selection='all'):
|
||
method.remove()
|
||
method.optimisation(selection)
|
||
|
||
# 观察小分子
|
||
# protgrey()
|
||
# colorp.color_mole('sele')
|
||
# add label
|
||
# label sele, "your label"
|
||
# black background
|
||
# set bg_rgb, [0, 0, 0]
|
||
# define label
|
||
# cmd.get_position("sele")
|
||
# set label_position, [x,y,z]
|
||
# cmd.h_add(pymol_obj) # add all hydrogens in this molecular
|
||
# cmd.remove(f"{pymol_obj} & hydro & not nbr. (don.|acc.)") # remove no-polar-hydrogen
|
||
# cmd.remove(f"sele & hydro") # 移除所有的H
|
||
# cmd.set('cartoon_color','color8', 'obj')
|
||
# cmd.hide("everything", "chain A and resn UHT and alt B") # 小分子隐藏B构象
|
||
# cmd.label("chain A and resn UHT and alt A", "resn") # 添加标签
|
||
#cmd.select('sele', f'resn {resn} and chain {chain}')
|
||
#cmd.center('sele')
|
||
#cmd.orient('sele')
|
||
#obj_name = 'around_and_self'
|
||
# cmd.create(obj_name, f'(resn {resn} and chain {chain}) around {distance}') # 不扩展残基
|
||
# cmd.create(obj_name, f'byres (resn {resn} and chain {chain}) around {distance}') # 扩展残基
|
||
# cmd.create(obj_name, f'(resn {resn} and chain {chain}) around {distance} or (resn {resn} and chain {chain})') # 选择resn的对象和resn对象周围3A的原子
|
||
#cmd.create(obj_name, f'byres (resn {resn} and chain {chain}) around {distance} or (resn {resn} and chain {chain})') # 选择resn的对象和resn对象周围3A的原子扩展至残基的原子
|
||
#cmd.hide('everything', 'all')
|
||
#cmd.show('lines', obj_name)
|
||
# 选择所有核酸（DNA 或 RNA）
|
||
# select nucleic_acids, polymer.nucleic
|
||
# # 选择所有蛋白质
|
||
# select proteins, polymer.protein
|
||
# # 锁定到链
|
||
# select chainA_protein, polymer.protein and chain A
|
||
## pymol 编辑模式
|
||
#选择mouse，3 botton edit 模式，需要在3 botton view 模式下先选择对象。
|
||
#按住shift + 左键进行旋转
|
||
#按住shift + 中键进行平移
|
||
cmd.extend('colorp',colorp)
|
||
cmd.extend('method',method)
|
||
|
||
if __name__ == '__main__':
|
||
PDB = '4i24'
|
||
cmd.reinitialize('everything')
|
||
cmd.fetch('4i24')
|
||
protgrey('4i24')
|
||
colorp.color_mole(obj='hetatm')
|
||
# 设置动画帧数
|
||
# cmd.mset("1 x60 61 x60") # 60帧展示初始结构，第61帧用于突变，再60帧展示突变后结构
|
||
|
||
# 保存第1帧
|
||
cmd.scene('frame1', 'store')
|
||
cmd.mview('store', state=1, scene='frame1')
|
||
|
||
# 执行残基突变。这里以将位于链A和编号为797的残基突变为ALA为例。
|
||
selection1 = "/4i24//A/797"
|
||
selection2 = "/4i24//A/1C9"
|
||
cmd.select("both_residues", f"{selection1} or {selection2}")
|
||
cmd.orient("both_residues")
|
||
cmd.show('sticks', selection1)
|
||
# 保存第2帧
|
||
cmd.scene('frame60', 'store')
|
||
cmd.mview('store', state=60, scene='frame60')
|
||
# 关闭电影模式
|
||
cmd.mstop() # 停止电影播放
|
||
cmd.mdelete() # 删除所有电影帧
|
||
cmd.wizard("mutagenesis")
|
||
cmd.refresh_wizard()
|
||
cmd.get_wizard().set_mode("GLY")
|
||
cmd.get_wizard().do_select(selection1)
|
||
# 保存第3帧
|
||
cmd.scene('frame120', 'store')
|
||
cmd.mview('store', state=120, scene='frame120')
|
||
cmd.get_wizard().apply()
|
||
cmd.set_wizard()
|
||
# 保存第4帧
|
||
cmd.scene('frame180', 'store')
|
||
cmd.mview('store', state=180, scene='frame180')
|
||
|
||
cmd.mpng('frame_prefix_', 1, 4) |