Skip to content
Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
11 changes: 8 additions & 3 deletions pymol_arrows.py
Original file line number Diff line number Diff line change
Expand Up @@ -6,11 +6,13 @@
import argparse


def write_pymol_arrows(base, atoms, scale, color, radius):
def write_pymol_arrows(base, atoms, scale, color, radius, threshold):
pymol_file = base + '_arrows.pymol'
lines = []
lines.append('run cgo_arrow.py')
arrow_objs = []
t2 = threshold**2
s2 = scale**2
for i, atom in enumerate(atoms):
arrow_obj = base + '_arrow_' + str(i)
arrow_objs.append(arrow_obj)
Expand All @@ -25,7 +27,8 @@ def write_pymol_arrows(base, atoms, scale, color, radius):
if color:
line += ', color={}'.format(color)
line += ', name={}'.format(arrow_obj)
lines.append(line)
if( (dx**2 + dy**2 + dz**2)*s2 > t2 ): # Check threshold
lines.append(line)
arrow_group = base + '_arrows'
line = 'group {}, {}'.format(arrow_group, ' '.join(arrow_objs))
lines.append(line)
Expand Down Expand Up @@ -94,14 +97,16 @@ def parse_args():
help='Output a .pdb file where the b-factor is gradient magnitude')
parser.add_argument('--sum', action='store_true', default=False,
help='Sum gradient components instead of taking magnitude')
parser.add_argument('-t', '--threshold', type=float, default=0,
help="Lower thrashold for arrow length (scaled)")
return parser.parse_args()


if __name__ == '__main__':
args = parse_args()
atoms = read_xyz_file(args.xyz_file)
base_name = args.xyz_file.replace('.xyz', '')
write_pymol_arrows(base_name, atoms, args.scale, args.color, args.radius)
write_pymol_arrows(base_name, atoms, args.scale, args.color, args.radius, args.threshold)
if args.pdb_file:
pdb_file = base_name + '.pdb'
write_pdb_file(pdb_file, atoms, args.sum)
Expand Down