def extract(s, force_recalc, sites, custom_symbol):
# First, we need the molecules
if Molecules.default_name not in s.info or force_recalc:
Molecules.get(s)
mols = s.info[Molecules.default_name]
elems = s.get_chemical_symbols()
if sites is None:
sites = range(len(s))
labels = []
for s_i in sites:
# Find the molecule that it belongs to
s_mol = [m for m in mols if s_i in m.indices][0]
s_elems = np.array(elems, dtype='S2')
if custom_symbol is not None:
s_elems[s_i] = custom_symbol
# Grab the bonds
bonds = s_mol.get_array(Bonds.default_name)
# This is a necessary step since the bonds are not classified
# by original structure index yet
bonds = {a: bonds[i] for i, a in enumerate(s_mol.indices)}
labels.append(recursive_mol_label(s_i, s_mol.indices,
bonds, s_elems))
return labels
def extract(s, force_recalc, save_info, save_asarray):
# First, we need the molecules
if Molecules.default_name not in s.info or force_recalc:
Molecules.get(s)
elems = s.get_chemical_symbols()
mol_sites = []
for mol_i, mol in enumerate(s.info[Molecules.default_name]):
# For each atom we do a depth-first traversal of the network
sites = {}
bonds = mol.get_array(Bonds.default_name)
# This is a necessary step since the bonds are not classified
# by original structure index yet
bonds = {a: bonds[i] for i, a in enumerate(mol.indices)}
for a in mol.indices:
sites[a] = recursive_mol_label(a, mol.indices, bonds, elems)
# Now grab the unique sites and pick the name of the molecule
site_names = sorted(list(set(sites.values())))
site_dict = {'name': site_names[0]}
# Now rename the sites
elem_sites = {}
for a in sites:
s_i = site_names.index(sites[a])
if elems[a] not in elem_sites:
elem_sites[elems[a]] = [s_i]
elif s_i not in elem_sites[elems[a]]:
elem_sites[elems[a]].append(s_i)
sites[a] = '{0}_{1}'.format(elems[a],
elem_sites[elems[a]].index(s_i)+1)
site_dict['sites'] = sites
mol_sites.append(site_dict)
if save_asarray:
arr = [sites[a] for a in mol.indices]
mol.set_array(MoleculeSites.default_name, arr)
if save_info:
s.info[MoleculeSites.default_name] = mol_sites
return mol_sites