Secure your code as it's written. Use Snyk Code to scan source code in minutes - no build needed - and fix issues immediately.
def _atom_id_matches(atom_id, atom, typemap):
""" Helper func for comparing atomic indices, symbols, neighbors, rings """
atomic_num = atom.element
if atom_id.data == 'atomic_num':
return atomic_num == int(atom_id.children[0])
elif atom_id.data == 'atom_symbol':
if str(atom_id.children[0]) == '*':
return True
elif str(atom_id.children[0]).startswith('_'):
# Store non-element elements in .name
return atom.name == str(atom_id.children[0])
else:
return atomic_num == pt.AtomicNum[str(atom_id.children[0])]
elif atom_id.data == 'has_label':
label = atom_id.children[0][1:] # Strip the % sign from the beginning.
return label in typemap[atom.idx]['whitelist']
elif atom_id.data == 'neighbor_count':
return len(atom.bond_partners) == int(atom_id.children[0])
elif atom_id.data == 'ring_size':
cycle_len = int(atom_id.children[0])
for cycle in typemap[atom.idx]['cycles']:
if len(cycle) == cycle_len:
return True
return False
elif atom_id.data == 'ring_count':
n_cycles = len(typemap[atom.idx]['cycles'])
if n_cycles == int(atom_id.children[0]):
return True
return False
typemap = {atom.idx: {'whitelist': set(), 'blacklist': set(),
'atomtype': None} for atom in structure.atoms}
rules = _load_rules(forcefield, typemap)
# Only consider rules for elements found in topology
subrules = dict()
system_elements = set()
for a in structure.atoms:
# First add non-element types, which are strings, then elements
if a.name.startswith('_'):
if a.name in forcefield.non_element_types:
system_elements.add(a.name)
else:
if 0 < a.atomic_number <= pt.KNOWN_ELEMENTS:
element = pt.Element[a.atomic_number]
system_elements.add(element)
else:
raise FoyerError(
'Parsed atom {} as having neither an element '
'nor non-element type.'.format(a)
)
for key, val in rules.items():
atom = val.nodes[0]['atom']
if len(list(atom.find_data('atom_symbol'))) == 1 and \
not list(atom.find_data('not_expression')):
try:
element = next(atom.find_data('atom_symbol')).children[0]
except IndexError:
try:
def _atom_id_matches(self, atom_id, atom):
atomic_num = atom.element._atomic_number
if atom_id.head == 'atomic_num':
return atomic_num == int(atom_id.tail[0])
elif atom_id.head == 'atom_symbol':
if str(atom_id.tail[0]) == '*':
return True
elif str(atom_id.tail[0]).startswith('_'):
return atom.element.name == str(atom_id.tail[0])
else:
return atomic_num == pt.AtomicNum[str(atom_id.tail[0])]
elif atom_id.head == 'has_label':
label = atom_id.tail[0][1:] # Strip the % sign from the beginning.
return label in (rule_name for rule_name in atom.whitelist)
elif atom_id.head == 'neighbor_count':
return len(atom.bond_partners) == int(atom_id.tail[0])
elif atom_id.head == 'ring_size':
cycle_len = int(atom_id.tail[0])
for cycle in atom.cycles:
if len(cycle) == cycle_len:
return True
return False
elif atom_id.head == 'matches_string':
raise NotImplementedError('matches_string feature is not yet implemented')
top_graph = nx.Graph()
top_graph.add_nodes_from(((a.idx, {'atom': a})
for a in structure.atoms))
top_graph.add_edges_from(((b.atom1.idx, b.atom2.idx)
for b in structure.bonds))
if self._graph_matcher is None:
atom = nx.get_node_attributes(self, name='atom')[0]
if len(list(atom.find_data('atom_symbol'))) == 1 and \
not list(atom.find_data('not_expression')):
try:
element = next(atom.find_data('atom_symbol')).children[0]
except IndexError:
try:
atomic_num = next(atom.find_data('atomic_num')).children[0]
element = pt.Element[int(atomic_num)]
except IndexError:
element = None
else:
element = None
self._graph_matcher = SMARTSMatcher(top_graph, self,
node_match=self._node_match,
element=element,
typemap=typemap)
matched_atoms = set()
for mapping in self._graph_matcher.subgraph_isomorphisms_iter():
mapping = {node_id: atom_id for atom_id, node_id in mapping.items()}
# The first node in the smarts graph always corresponds to the atom
# that we are trying to match.
atom_index = mapping[0]
# Don't yield duplicate matches found via matching the pattern in a
'atomtype': None} for atom in structure.atoms}
rules = _load_rules(forcefield, typemap)
# Only consider rules for elements found in topology
subrules = dict()
system_elements = set()
for a in structure.atoms:
# First add non-element types, which are strings, then elements
if a.name.startswith('_'):
if a.name in forcefield.non_element_types:
system_elements.add(a.name)
else:
if 0 < a.atomic_number <= pt.KNOWN_ELEMENTS:
element = pt.Element[a.atomic_number]
system_elements.add(element)
else:
raise FoyerError(
'Parsed atom {} as having neither an element '
'nor non-element type.'.format(a)
)
for key, val in rules.items():
atom = val.nodes[0]['atom']
if len(list(atom.find_data('atom_symbol'))) == 1 and \
not list(atom.find_data('not_expression')):
try:
element = next(atom.find_data('atom_symbol')).children[0]
except IndexError:
try:
atomic_num = next(atom.find_data('atomic_num')).children[0]