# Example for: model.pick_atoms()

# This will pick various subsets of atoms in the MODEL and compare them 
# with MODEL2.

env = environ()
log.level(1, 1, 1, 1, 0)

# Set some defaults (the same as in top.ini):
env.selection_mode = 'ATOM' # only the selected atoms, not whole residues
env.selection_from = 'ALL'  # scanning of all atoms, not selected atoms
env.selection_search = 'SEGMENT'  # scan over a segment
env.selection_segment = ('FIRST:', 'LAST:')   # the whole chain as a segment
env.res_types = 'ALL'               # all residue types
env.pick_atoms_set = 1              # put the selected atoms in set 1
env.selection_status = 'INITIALIZE' # select only the selected atoms

# Read the models and the alignment:
mdl  = model(env, file='1fas')
mdl2 = model(env, file='2ctx')
aln = alignment(env, file='toxin.ali', align_codes=('1fas', '2ctx'))
aln.write(file='toxin.pap', alignment_format='PAP')

# Pick and superpose mainchain atoms:
mdl.pick_atoms(aln, atom_types='MNCH')
mdl.superpose(mdl2, aln)

# Pick and superpose sidechain atoms:
mdl.pick_atoms(aln, atom_types='SDCH')
mdl.superpose(mdl2, aln)

# Pick and superpose CA and CB atoms:
mdl.pick_atoms(aln, atom_types='CA CB')
mdl.superpose(mdl2, aln)

# Pick and superpose all atoms:
mdl.pick_atoms(aln, atom_types='ALL')
mdl.superpose(mdl2, aln)

# Pick and superpose CA and CB atoms in one segment only:
mdl.pick_atoms(aln, atom_types='CA CB', selection_segment=('2:', '10:'))
mdl.superpose(mdl2, aln)

# Pick and superpose all atoms within 6 angstroms of the 'CA' atom in residue '10:':
mdl.pick_atoms(aln, atom_types='ALL', sphere_radius=6.0,
               selection_search='SPHERE', sphere_center=('10:', 'CA'))
mdl.superpose(mdl2, aln)

# Pick and superpose all atoms within 6 angstroms of any atom in 
# segment 2: to 10:
mdl.pick_atoms(aln, atom_types='ALL', selection_segment=('2:', '10:'),
               selection_search='SPHERE_SEGMENT', sphere_radius=6.0)
mdl.superpose(mdl2, aln)

# Pick and superpose all atoms in all loops (ie residues within 2 positions 
# of any gap in the alignment):
mdl.pick_atoms(aln, atom_types='ALL', selection_segment=('LOOPS', ''),
               selection_search='SEGMENT', gap_extension=(2, 2))
mdl.superpose(mdl2, aln)

# Pick and superpose all atoms within 6 angstroms of all loops (ie residues
# within 2 positions of any gap in the alignment):
mdl.pick_atoms(aln, atom_types='ALL', selection_segment=('LOOPS', ''),
               selection_search='SPHERE_SEGMENT', sphere_radius=6.0,
               gap_extension=(2, 2))
mdl.superpose(mdl2, aln)