def show_sticks(xyz=None, species=None, project_directory=None):
"""
Draws the molecule in a "sticks" style according to the supplied xyz coordinates
Returns whether successful of not
If successful, save an image using draw_3d
"""
xyz = check_xyz_species_for_drawing(xyz, species)
coordinates, _, _, _, _ = get_xyz_matrix(xyz)
s_mol, b_mol = molecules_from_xyz(xyz)
mol = b_mol if b_mol is not None else s_mol
try:
_, rd_mol, _ = rdkit_conf_from_mol(mol, coordinates)
except ValueError:
return False
mb = Chem.MolToMolBlock(rd_mol)
p = p3d.view(width=400, height=400)
p.addModel(mb, 'sdf')
p.setStyle({'stick': {}})
# p.setBackgroundColor('0xeeeeee')
p.zoomTo()
p.show()
draw_3d(xyz=xyz,
species=species,
project_directory=project_directory,
save_only=True)
return True
def draw_3d(xyz=None, species=None, project_directory=None, save_only=False):
"""
Draws the molecule in a "3D-balls" style
If xyz ig given, it will be used, otherwise the function looks for species.final_xyz
Input coordinates are in string format
Saves an image if a species and `project_directory` are provided
If `save_only` is ``True``, then don't plot, only save the image
"""
xyz = check_xyz_species_for_drawing(xyz, species)
_, atoms, x, y, z = get_xyz_matrix(xyz)
atoms = [str(a) for a in atoms]
ase_atoms = list()
for i, atom in enumerate(atoms):
ase_atoms.append(Atom(symbol=atom, position=(x[i], y[i], z[i])))
ase_mol = Atoms(ase_atoms)
if not save_only:
display(view(ase_mol, viewer='x3d'))
if project_directory is not None and species is not None:
folder_name = 'rxns' if species.is_ts else 'Species'
geo_path = os.path.join(project_directory, 'output', folder_name,
species.label, 'geometry')
if not os.path.exists(geo_path):
os.makedirs(geo_path)
ase_write(filename=os.path.join(geo_path, 'geometry.png'),
images=ase_mol,
scale=100)
def test_xyz_format_conversion(self):
"""Test conversions from string to list xyz formats"""
xyz_str0 = """N 2.24690600 -0.00006500 0.11597700
C -1.05654800 1.29155000 -0.02642500
C -1.05661400 -1.29150400 -0.02650600
C -0.30514100 0.00000200 0.00533200
C 1.08358900 -0.00003400 0.06558000
H -0.39168300 2.15448600 -0.00132500
H -1.67242600 1.35091400 -0.93175000
H -1.74185400 1.35367700 0.82742800
H -0.39187100 -2.15447800 0.00045500
H -1.74341400 -1.35278100 0.82619100
H -1.67091600 -1.35164600 -0.93286400
"""
xyz_list, atoms, x, y, z = get_xyz_matrix(xyz_str0)
# test all forms of input for get_xyz_string():
xyz_str1 = get_xyz_string(coords=xyz_list, symbols=atoms)
xyz_str2 = get_xyz_string(coords=xyz_list, numbers=[7, 6, 6, 6, 6, 1, 1, 1, 1, 1, 1])
mol, _ = molecules_from_xyz(xyz_str0)
xyz_str3 = get_xyz_string(coords=xyz_list, mol=mol)
self.assertEqual(xyz_str0, xyz_str1)
self.assertEqual(xyz_str1, xyz_str2)
self.assertEqual(xyz_str2, xyz_str3)
self.assertEqual(atoms, ['N', 'C', 'C', 'C', 'C', 'H', 'H', 'H', 'H', 'H', 'H'])
self.assertEqual(x, [2.246906, -1.056548, -1.056614, -0.305141, 1.083589, -0.391683, -1.672426, -1.741854,
-0.391871, -1.743414, -1.670916])
self.assertEqual(y[1], 1.29155)
self.assertEqual(z[-1], -0.932864)
def test_rdkit_conf_from_mol(self):
"""Test rdkit_conf_from_mol"""
_, b_mol = converter.molecules_from_xyz(self.xyz2)
xyz, _, _, _, _ = converter.get_xyz_matrix(self.xyz2)
conf, rd_mol, indx_map = converter.rdkit_conf_from_mol(mol=b_mol,
coordinates=xyz)
self.assertTrue(conf.Is3D())
self.assertEqual(rd_mol.GetNumAtoms(), 5)
self.assertEqual(indx_map, {0: 0, 1: 1, 2: 2, 3: 3, 4: 4})
def test_get_xyz_matrix(self):
"""Test conversion of xyz string to array format"""
xyz, symbols, x, y, z = converter.get_xyz_matrix(self.xyz2)
self.assertEqual(
xyz,
[[1.1746411, -0.15309781, 0.0], [0.06304988, 0.35149648, 0.0],
[-1.12708952, -0.11333971, 0.0], [-1.93800144, 0.60171738, 0.0],
[-1.29769464, -1.18742971, 0.0]])
self.assertEqual(symbols, ['O', 'N', 'C', 'H', 'H'])
self.assertEqual(
x, [1.1746411, 0.06304988, -1.12708952, -1.93800144, -1.29769464])
self.assertEqual(
y, [-0.15309781, 0.35149648, -0.11333971, 0.60171738, -1.18742971])
self.assertEqual(z, [0.0, 0.0, 0.0, 0.0, 0.0])
def plot_3d_mol_as_scatter(xyz, path=None, plot_h=True, show_plot=True):
"""
Draws the molecule as scattered balls in space according to the supplied xyz coordinates
`xyz` is in string form
`path` is the species output path to save the image
"""
xyz, atoms, x, y, z = get_xyz_matrix(xyz)
x = np.array(x, dtype=float)
y = np.array(y, dtype=float)
z = np.array(z, dtype=float)
colors = []
sizes = []
for i, atom in enumerate(atoms):
size = 500
if atom == 'H':
if plot_h:
colors.append('gray')
size = 250
else:
colors.append('white')
atoms[i] = ''
x[i], y[i], z[i] = 0, 0, 0
elif atom == 'C':
colors.append('k')
elif atom == 'N':
colors.append('b')
elif atom == 'O':
colors.append('r')
elif atom == 'S':
colors.append('orange')
else:
colors.append('g')
sizes.append(size)
fig = plt.figure()
ax = Axes3D(fig)
ax.scatter(xs=x, ys=y, zs=z, s=sizes, c=colors, depthshade=True)
for i, atom in enumerate(atoms):
ax.text(x[i] + 0.01, y[i] + 0.01, z[i] + 0.01, str(atom), size=7)
plt.axis('off')
if show_plot:
plt.show()
if path is not None:
image_path = os.path.join(path, "scattered_balls_structure.png")
plt.savefig(image_path, bbox_inches='tight')
