def _load_snapshot_chimera_obj(session, data, model_by_id):
"""Map data from _save_snapshot_chimera_obj back to a ChimeraX object"""
if data is None:
return None
elif data['type'] == 'Atoms':
if 'single_structure' in data:
if data['single_structure'] is None:
return Atoms()
m = model_by_id[data['single_structure']]
atoms = [m.atoms[x] for x in data['indices']]
else:
atoms = []
for s, ind in zip(data['structures'], data['indices']):
atoms.append(model_by_id[s].atoms[ind])
obj = Atoms(atoms)
return obj
elif data['type'] == 'Atom':
return model_by_id[data['structure']].atoms[data['index']]
elif data['type'] == 'Bond':
return model_by_id[data['structure']].bonds[data['index']]
elif data['type'] == 'Pseudobond':
f = model_by_id[data['structure']]._get_features()
obj = f.pseudobonds[data['index']]
return obj
raise TypeError("Don't know how to load snapshot %s" % str(data))
def apply_restraints(trs, rrs, adjust_for_confidence, confidence_type):
template_as = []
restrained_as = []
for tr, rr in zip(trs, rrs):
ta_names = set(tr.atoms.names).intersection(atom_names)
ra_names = set(rr.atoms.names).intersection(atom_names)
common_names = list(ta_names.intersection(ra_names))
template_as.extend([tr.find_atom(name) for name in common_names])
restrained_as.extend([rr.find_atom(name) for name in common_names])
# template_as.append(tr.atoms[numpy.in1d(tr.atoms.names, common_names)])
# restrained_as.append(rr.atoms[numpy.in1d(rr.atoms.names, common_names)])
from chimerax.atomic import Atoms
template_as = Atoms(template_as)
restrained_as = Atoms(restrained_as)
template_coords = template_as.coords
from math import sqrt
for i, ra1 in enumerate(restrained_as):
query_coord = numpy.array([template_coords[i]])
indices = find_close_points(query_coord, template_coords,
distance_cutoff)[1]
indices = indices[indices != i]
for ind in indices:
ra2 = restrained_as[ind]
if ra1.residue == ra2.residue:
continue
if adjust_for_confidence:
if confidence_type == 'plddt':
scores = [
template_as[i].bfactor * confidence_multiplier,
template_as[ind].bfactor * confidence_multiplier
]
elif confidence_type == 'pae':
scores = [
pae_matrix[template_as[i].residue.number - 1,
template_as[ind].residue.number - 1]
]
kappa_adj, tol_adj, falloff_adj = adjust_distance_restraint_terms_by_confidence(
scores, confidence_type)
if kappa_adj == 0:
continue
else:
kappa_adj = tol_adj = 1
falloff_adj = 0
try:
dr = adrm.add_restraint(ra1, ra2)
except ValueError:
continue
dist = distance(query_coord[0], template_coords[ind])
dr.tolerance = tolerance * dist * tol_adj
dr.target = dist
dr.c = max(sqrt(dist) * well_half_width, 0.1)
#dr.effective_spring_constant = spring_constant
dr.kappa = kappa * kappa_adj
from math import log
dr.alpha = -1 - fall_off * log(
(max(dist - 1, 1))) - falloff_adj
dr.enabled = True
def added_atoms(session, atoms, to_atoms):
amap = {(a.residue.chain_id, a.residue.number, a.name): a for a in atoms}
from chimerax.atomic import Atoms
added = Atoms([
a for a in to_atoms
if (a.residue.chain_id, a.residue.number, a.name) not in amap
])
session.selection.clear()
added.selected = True
session.logger.status('Added %d atoms' % len(added), log=True)
return added
def paired_principal_atoms(aligned_residues):
from chimerax.atomic import Atoms
result = [
Atoms(aa)
for aa in zip(*((a1, a2)
for a1, a2 in zip(aligned_residues[0].principal_atoms,
aligned_residues[1].principal_atoms)
if a1 is not None and a2 is not None))
]
if not len(result):
return [Atoms(), Atoms()]
return result
def __init__(self, structure):
self._log = Logger('structuretugger.log' if write_logs else None)
self.structure = structure
self._minimized = False
# OpenMM requires the atoms to be sorted by residue number
self._structure_atoms = sa = structure.atoms
satoms = list(sa)
satoms.sort(key=lambda a: (a.residue.chain_id, a.residue.number))
from chimerax.atomic import Atoms
self.atoms = Atoms(satoms)
# Atom being tugged
self.atom = None
initialize_openmm()
# OpenMM objects
self._topology = None
self._system = None
self._force = None # CustomExternalForce pulling force
self._platform = None
self._simulation = None
self._sim_forces = None # Current forces on atoms
# OpenMM simulation parameters
global openmm_forcefield_parameters
self._forcefields = openmm_forcefield_parameters
self._sim_steps = 50 # Simulation steps between mouse position updates
self._force_constant = 10000
from simtk import unit
#self._temperature = 300*unit.kelvin
self._temperature = 100 * unit.kelvin
#self._constraint_tolerance = 0.00001
#self._time_step = 2.0*unit.femtoseconds
self._integrator_tolerance = 0.001
self._constraint_tolerance = 0.001
self._friction = 1.0 / unit.picoseconds # Coupling to heat bath
self._platform_name = 'CPU'
#self._platform_name = 'OpenCL' # Works on Mac
#self._platform_name = 'CUDA' # This is 3x faster but requires env DYLD_LIBRARY_PATH=/usr/local/cuda/lib Chimera.app/Contents/MacOS/ChimeraX so paths to cuda libraries are found.
self._max_allowable_force = 50000.0 # kJ/mol/nm
# OpenMM particle data
self._particle_number = None # Integer index of tugged atom
self._particle_positions = None # Numpy array, Angstroms
self._particle_force_index = {}
self._particle_masses = None # Original particle masses
self._create_openmm_system()
def indexLabel(session, models=None):
from chimerax.core.objects import Objects
from chimerax.atomic import AtomicStructure, Atoms
from chimerax.label.label3d import label
from SEQCROW.utils import contrast_bw
if models is None:
models = session.models.list(type=AtomicStructure)
elif isinstance(models, AtomicStructure):
models = [models]
for m in models:
for i, atom in enumerate(m.atoms):
l = str(i + 1)
ele_color = atom.color[:-1]
if contrast_bw(ele_color) == "black":
label_color = (0, 0, 0, 255)
else:
label_color = (255, 255, 255, 255)
label(
session,
objects=Objects(atoms=Atoms([atom])),
object_type='atoms',
text=l,
offset=(-0.11 * len(l), -0.2, -0.2),
height=0.4,
on_top=True,
color=label_color,
)
def atoms(self):
'''
Returns an unsorted `Atoms` object encompassing all atoms in chi dihedrals. Read only.
'''
from chimerax.atomic import concatenate, Atoms
dihedrals = concatenate([r.chi_dihedrals for r in self])
return concatenate([Atoms(d.atoms) for d in dihedrals]).unique()
def swap_equivalent_atoms(residue):
rname = residue.name
if rname in rings:
rname = 'aromatic'
equivalent_pairs = equivalent_heavy_atoms[rname]
from chimerax.atomic import Atoms
paired_atoms = []
for a1name, a2name in equivalent_pairs.items():
a1, a2 = [residue.find_atom(name) for name in (a1name, a2name)]
if a1 is None or a2 is None:
# Bail out if any equivalent atom is missing
return False
pair = []
for a in (a1, a2):
group = [a]
for n in a.neighbors:
if n.residue != a.residue:
# Don't flip if any equivalent atom is bonded to something else
return False
if n.element.name == 'H':
group.append(n)
pair.append(Atoms(group))
paired_atoms.append(pair)
for pair in paired_atoms:
if len(pair[0]) != len(pair[1]):
# Something weird going on - probably one of the atoms is missing a hydrogen. Bail out.
return False
for pair in paired_atoms:
pair[0].coords, pair[1].coords = pair[1].coords, pair[0].coords
return True
def correct_o_position(res, psi):
n, ca, c, o = [res.find_atom(name) for name in ['N', 'CA', 'C', 'O']]
from chimerax.geometry import rotation, dihedral
d = dihedral(*[a.coord for a in (n, ca, c, o)])
r = rotation(c.coord - ca.coord, psi + 180 - d, center=c.coord)
from chimerax.atomic import Atoms
Atoms([o]).transform(r)
def break_disulfide(cys1, cys2):
from chimerax.core.errors import UserError
from chimerax.atomic import Atoms
s1 = cys1.find_atom('SG')
s2 = cys2.find_atom('SG')
if s1 is None or s2 is None:
raise UserError(
'Missing SG atom! Are both residues complete cysteines?')
if s2 not in s1.neighbors:
raise UserError('These residues are not disulfide bonded!')
has_hydrogens = ('H' in cys1.atoms.element_names)
b = Atoms((s1, s2)).intra_bonds[0]
b.delete()
if has_hydrogens:
from chimerax.build_structure import modify_atom
modify_atom(s1, s1.element, 2, res_name='CYS', connect_back=False)
modify_atom(s2, s2.element, 2, res_name='CYS', connect_back=False)
def test_save_snapshot_chimera_obj(self):
"""Test save_snapshot of Chimera objects"""
self.assertIsNone(src.io._save_snapshot_chimera_obj(None))
self.assertRaises(TypeError, src.io._save_snapshot_chimera_obj,
'non-chimera object')
session = make_session()
state = src.io._RMFState(session)
state2 = src.io._RMFState(session)
residue = state.new_residue('ALA', 'A', 1)
atom1 = state.new_atom('C', 'C')
residue.add_atom(atom1)
atom2 = state.new_atom('N', 'N')
residue.add_atom(atom2)
st2_residue = state2.new_residue('ALA', 'A', 1)
st2_atom1 = state2.new_atom('C', 'C')
st2_residue.add_atom(st2_atom1)
d = src.io._save_snapshot_chimera_obj(atom1)
self.assertEqual(d['type'], 'Atom')
# single structure
atoms = Atoms((atom1, atom2))
d = src.io._save_snapshot_chimera_obj(atoms)
self.assertEqual(d['type'], 'Atoms')
self.assertIn('single_structure', d)
# list of zero atoms
atoms = Atoms()
d = src.io._save_snapshot_chimera_obj(atoms)
self.assertEqual(d['type'], 'Atoms')
self.assertIn('single_structure', d)
# multiple structures
atoms = Atoms((atom1, st2_atom1))
d = src.io._save_snapshot_chimera_obj(atoms)
self.assertEqual(d['type'], 'Atoms')
self.assertIn('structures', d)
bond = state.new_bond(atom1, atom2)
d = src.io._save_snapshot_chimera_obj(bond)
self.assertEqual(d['type'], 'Bond')
pbond = state._add_pseudobond((atom1, atom2))
d = src.io._save_snapshot_chimera_obj(pbond)
self.assertEqual(d['type'], 'Pseudobond')
def all_connected_selector(session, models, results):
"""select all atoms connected to the current selection"""
# TODO: right mouse mode for this
cur_sel = selected_atoms(session)
bond_sel = selected_bonds(session)
for bond in bond_sel:
cur_sel = cur_sel.merge(Atoms(bond.atoms))
atoms = Atoms()
for atom in cur_sel:
if atom in atoms:
continue
elif atom.structure not in models:
continue
connected_atoms = get_fragment(atom, max_len=len(atom.structure.atoms))
atoms = atoms.merge(connected_atoms)
results.add_atoms(atoms)
def new_feature(self, atoms):
if len(atoms) == 2:
# Make and return a pseudobond if the feature acts on 2 atoms
state = atoms[0].structure
return state._add_pseudobond(atoms)
else:
# Otherwise, return the list of atoms the feature acts on
return Atoms(atoms)
def move_models_and_atoms(tf, models, atoms, move_whole_molecules, base_model):
if move_whole_molecules and atoms is not None and len(atoms) > 0:
models = list(models) + list(atoms.unique_structures)
from chimerax.atomic import Atoms
atoms = Atoms()
if atoms is None:
from chimerax.atomic import Atoms
atoms = Atoms()
global position_history
position_history.record_position(models, atoms, base_model)
for m in models:
# TODO: Handle case where parent of volume has non-identity position. tf is in scene coords.
m.position = tf * m.position
if len(atoms) > 0:
atoms.coords = tf * atoms.coords
position_history.record_position(models, atoms, base_model)
def __init__(self, triangle_range, description, atoms):
# triangle_range is a range object that corresponds to the indices
# of the triangles for that shape in the vertex array
self.triangle_range = triangle_range
self.description = description
from chimerax.atomic import Atoms
if atoms is not None and not isinstance(atoms, Atoms):
atoms = Atoms(atoms)
self.atoms = atoms
def _step_forward(self, *_):
'''
Move one step forward along the spline, and adjust the target
positions accordingly.
'''
isolde = self.isolde
xr = self._extended_residues
xa = self._extended_atoms
if abs(self._current_position_on_spline) >= abs(self._shift_length):
# Our target has reached the end of the spline. Switch to a
# final "polishing" routine
isolde.sim_handler.triggers.remove_handler(self._handler)
self._handler = None
self.finished = True
self.triggers.activate_trigger('register shift finished', self)
#self._handler = self.isolde.triggers.add_handler('completed simulation step', self._final_polish)
self._current_position_on_spline += self._spline_step
sp = self._positions_along_spline
sp += self._spline_step
outside = numpy.logical_or(sp < 0, sp > self._spline_length - 1)
inside = numpy.argwhere(numpy.invert(outside)).ravel()
outside = numpy.argwhere(outside).ravel()
# Release restraints on atoms outside of the spline
out_a = xr[outside].atoms
if len(out_a):
isolde.release_xyz_restraints_on_selected_atoms(sel=out_a)
# Update the restraints on atoms inside of the spline
in_a = xa[inside]
in_pos = sp[inside]
nspl = self.n_spline
caspl = self._ca_spline
cspl = self._c_spline
cbspl = self._cb_spline
splev = interpolate.splev
#restraints = isolde._sim_pos_restr
pr_mgr = self._pr_mgr
k = self.spring_constant
n_targets = numpy.column_stack(splev(in_pos, nspl[0]))
ca_targets = numpy.column_stack(splev(in_pos, caspl[0]))
c_targets = numpy.column_stack(splev(in_pos, cspl[0]))
cb_targets = numpy.column_stack(splev(in_pos, cbspl[0]))
#all_targets = numpy.column_stack((n_targets, ca_targets, c_targets, cb_targets))
from chimerax.atomic import Atoms
for i, targets in enumerate(
(n_targets, ca_targets, c_targets, cb_targets)):
atoms = in_a[:, i]
mask = atoms != None
targets = targets[mask]
atoms = Atoms(atoms[mask])
prs = pr_mgr.get_restraints(atoms)
prs.targets = targets
prs.spring_constants = k
prs.enableds = True
def __init__(self, residues, cartesian, log = None):
# Create interpolation function for each residue.
from .interp_residue import internal_residue_interpolator, cartesian_residue_interpolator
residue_interpolator = cartesian_residue_interpolator if cartesian else internal_residue_interpolator
res_interp = {}
t0 = time()
cartesian_atoms = []
dihedral_atoms = []
nr = len(residues)
for i,r in enumerate(residues):
if log and i%100 == 0:
log.status('Making interpolator for residue %d of %d' % (i, nr))
residue_interpolator(r, cartesian_atoms, dihedral_atoms)
t1 = time()
global rsit
rsit += t1-t0
from chimerax.atomic import Atoms
self.cartesian_atom_indices = Atoms(cartesian_atoms).coord_indices
self.dihedral_atom_indices = Atoms(dihedral_atoms).coord_indices
def split_connected(atoms):
aset = set(atoms)
reached = {} # Map atom to tuple of connected atoms
for a in atoms:
j = set([a])
for b in a.bonds:
a2 = b.other_atom(a)
if a2 in aset and a2 in reached:
j.update(reached[a2])
j = tuple(j)
for a3 in j:
reached[a3] = j
cats = list(set(reached.values()))
cats.sort(key = lambda cat: len(cat))
cats.reverse() # Number largest to smallest
from chimerax.atomic import Atoms
pieces = ([('', Atoms(cats[0]))] if len(cats) == 1
else [('%d' % (i+1,), Atoms(cat)) for i,cat in enumerate(cats)])
return pieces
def _pick_atoms(self, pick):
if pick is None:
return None
a = None
r = None
if hasattr(pick, 'atom'):
a = pick.atom
r = a.residue
elif hasattr(pick, 'bond'):
b = pick.bond
coords = [a.coord for a in b.atoms]
from chimerax.geometry import distance
distances = [distance(c, pick.position) for c in coords]
a = b.atoms[distances.index(min(distances))]
r = a.residue
elif hasattr(pick, 'residue'):
r = pick.residue
# Tug heavy atoms instead of hydrogens
if a:
if a.element.name == 'H':
h_mode = self._tug_mgr.allow_hydrogens
if h_mode == 'no':
self.session.logger.warning(
'Tugging of hydrogens is not enabled. '
'Applying tug to the nearest bonded heavy atom.')
a = a.neighbors[0]
elif h_mode == 'polar':
for n in a.neighbors:
if n.element.name == 'C':
self.session.logger.warning(
'Tugging of non-polar hydrogens is not enabled. '
'Applying tug to the nearest bonded heavy atom.'
)
a = n
break
self._focal_atom = a #a = self._focal_atom = pick
tm = self.tug_mode
if tm == "atom" and a:
from chimerax.atomic import Atoms
pa = Atoms([a])
elif tm == "residue" and r:
pa = r.atoms
if a is not None:
self._focal_atom = a
else:
self._focal_atom = r.principal_atom
else:
pa = None
return pa
def tm_selector(session, models, results):
"""select transition metals using AaronTools' TMETAL dictionary"""
atoms = []
for model in models:
if isinstance(model, AtomicStructure):
for atom in model.atoms:
if atom.element.name in TMETAL:
atoms.append(atom)
#need to add a Collection, not just a list of atoms
results.add_atoms(Atoms(atoms))
def _label_objects(self, r):
# Label CA atom so label does not jump around.
la = [a for a in r.atoms if a.name == self._label_atom_name]
from chimerax.core.objects import Objects
if len(la) == 1:
from chimerax.atomic import Atoms
objects = Objects(atoms=Atoms(la))
otype = 'atoms'
else:
# If no CA atom them label center of residue
objects = Objects(atoms=r.atoms)
otype = 'residues'
return objects, otype
def shortcut_atoms(session):
matoms = []
sel = session.selection
atoms_list = sel.items('atoms')
from chimerax.atomic import concatenate, Atoms
if atoms_list:
atoms = concatenate(atoms_list)
elif sel.empty():
# Nothing selected, so operate on all atoms
from chimerax.atomic import all_atoms
atoms = all_atoms(session)
else:
atoms = Atoms()
return atoms
def call_c_plus_plus(cpp_func, structures, return_collection, *args):
import os
num_cpus = os.cpu_count() if not None else 1
from chimerax.atomic import Atom, Atoms, AtomicStructure
groups = []
for structure in structures:
if not isinstance(structure, AtomicStructure):
continue
cpp_args = args + (num_cpus, return_collection)
grps = cpp_func(structure.cpp_pointer, *cpp_args)
if return_collection:
# accumulate the numpy arrays to later be concatenated and turned into a Collection
groups.append(grps)
else:
groups.extend(grps)
if return_collection:
if groups:
import numpy
groups = Atoms(numpy.concatenate(groups))
else:
groups = Atoms()
return groups
def selected_items(self, itype):
if itype in ('atoms', 'bonds'):
from chimerax.atomic import Atoms
atoms = Atoms(None)
for s in self._selected_shapes:
a = s.atoms
if a is not None:
atoms |= s
if itype == 'bonds':
return atoms.intra_bonds
return atoms
elif itype == 'shapes':
return list(self._selected_shapes)
return []
def save_residue_graphs_to_text(residues, filename):
import numpy
from chimerax.atomic import Atoms
with open(filename, 'wt') as outfile:
for r in residues:
ratoms = r.atoms
labels = ratoms.elements.numbers
edges = numpy.array([ratoms.indices(Atoms(b.atoms)) for b in ratoms.intra_bonds])
outfile.write('graph_{}\n'.format(r.name))
outfile.write('labels\n')
outfile.write(', '.join([str(l) for l in labels])+'\n')
outfile.write('edges\n')
for e in edges:
outfile.write('{},{}\n'.format(e[0],e[1]))
def indexLabel(session, models=None):
from chimerax.core.objects import Objects
from chimerax.atomic import AtomicStructure, Atoms
from chimerax.label.label3d import label
if models is None:
models = session.models.list(type=AtomicStructure)
elif isinstance(models, AtomicStructure):
models = [models]
for m in models:
for i, atom in enumerate(m.atoms):
l = str(i + 1)
label(session, objects=Objects(atoms=Atoms([atom])), object_type='atoms', \
text=l, offset=(-0.11*len(l),-0.2,-0.2), height=0.4, on_top=True)
def get_fragment(start, stop=None, max_len=100000):
"""
see AaronTools.geometry.Geometry.get_fragment
"""
stack = deque([start])
frag = [start]
stop = set([stop])
while len(stack) > 0 and len(frag) <= max_len:
connected = stack.popleft()
connected = set(connected.neighbors) - stop - set(frag)
stack.extend(connected)
frag.extend(connected)
return Atoms(frag)
def apply_strict_ncs(template_atoms, ncs_map):
'''
Copy the transformed coordinates of the template atoms to the given NCS
atoms.
Args:
* template_atoms: a :class:`Atoms` instance
* ncs_map: a list of (:class:`Place`, :class:`Atoms`) pairs. Each
array of atoms must have a 1:1 correspondence to the template_atoms.
Atoms will be sorted according to (chain, residue number, insertion
code, atom name) before applying the NCS operators.
'''
from chimerax.atomic import Atoms
template_atoms = Atoms(
sorted(template_atoms,
key=lambda a: (a.residue.chain_id, a.residue.number, a.residue.
insertion_code, a.name)))
for (place, ncs_atoms) in ncs_map:
ncs_atoms = Atoms(
sorted(ncs_atoms,
key=lambda a: (a.residue.chain_id, a.residue.number, a.
residue.insertion_code, a.name)))
ncs_atoms.coords = place * template_atoms.coords
def make_graph_from_residue(residue, label='element'):
'''
Create a graph representation of a residue's topology, with nodes labelled
by either atomic number or an integer representation of their name.
'''
import numpy
from chimerax.atomic import Atoms
ratoms = residue.atoms
if label=='element':
labels = ratoms.elements.numbers
elif label=='name':
labels = numpy.array([simple_name_to_int(name) for name in ratoms.names])
else:
raise TypeError('"label" argument should be one of "element" or "name"!')
edges = numpy.array([ratoms.indices(Atoms(b.atoms)) for b in ratoms.intra_bonds])
return Graph(labels, edges)
def populate_clash_table_from_geometry(self, model):
'''
If no simulation is running, use ChimeraX's built-in clash toolset to find
and list geometric clashes in decreasing order of overlap severity.
'''
atoms = model.atoms
from chimerax.clashes import clashes
clash_dict = clashes.find_clashes(self.session,
atoms,
inter_model=False)
t = self._geom_clash_table
t.setRowCount(0)
if not clash_dict:
return
from functools import reduce
# Clash dict is bi-directional, so number of clashes is half the total count
clash_count = 1 / 2 * reduce(lambda x, y: x + y,
(len(d.keys())
for d in clash_dict.values()))
t.setRowCount(clash_count)
seen = set()
from chimerax.atomic import Atoms
# Make list of unique clashing atom pairs and their distances
clash_list = []
for a1, clashes in clash_dict.items():
seen.add(a1)
for a2, dist in clashes.items():
if a2 in seen:
continue
clash_list.append((Atoms([a1, a2]), dist))
# Sort clashes in decreasing order of overlap
clash_list = sorted(clash_list, key=lambda x: x[1], reverse=True)
from Qt.QtWidgets import QTableWidgetItem
for i, (catoms, overlap) in enumerate(clash_list):
a1, a2 = catoms
r1, r2 = catoms.residues
data = ("{} {}{}: {}".format(r1.name, r1.chain_id, r1.number,
a1.name),
"{} {}{}: {}".format(r2.name, r2.chain_id, r2.number,
a2.name), "{:0.2f}".format(overlap))
for j, d in enumerate(data):
item = QTableWidgetItem(d)
item.data = catoms
t.setItem(i, j, item)
t.resizeColumnsToContents()
