def com_distance(st, atoms1, st2=None, atoms2=None):
"""
Return Center of Mass distance for two sets of atoms.
If one set of atoms and one structure is provided,
the distance between the subset and the full structure is returned
If a second stucture (st2) is provided atoms2 is a subset of st2.
:param st:
:type st: schrodinger.structure.Structure
:param atoms1:
:type atoms1: list
:param st2:
:type st2: schrodinger.structure.Structure
:param atoms2:
:type atoms2: list
:return:
"""
if st2 is None:
st2 = st
if atoms2 is None:
atoms2 = [atm.index for atm in st2.atom]
com1 = center_of_mass(st, atoms1)
com2 = center_of_mass(st2, atoms2)
return np.sqrt(np.sum((com1 - com2)**2))
def move_molecule1(atoms, space_group):
print("atoms.get_distance(1,5): ",
atoms.get_distance(1, 5, mic=False, vector=True))
print("atoms.get_distance(1,21): ",
atoms.get_distance(1, 21, mic=False, vector=True))
a, b, c, alpha, beta, gamma = atoms.get_cell_lengths_and_angles()
X = atoms.get_positions()
old_scaled_position = atoms.get_scaled_positions()
print("a, b, c, alpha, beta, gamma: ", a, b, c, alpha, beta, gamma)
#print("X: ", X)
#print("old_scaled_position: ", old_scaled_position)
molecule1 = molecule_lists(atoms)[0]
st = ase_atoms_to_structure(atoms)
x_com1 = center_of_mass(st, molecule1)
X_molecule1 = np.array([X[x - 1] for x in molecule1])
#print("X_molecule1: ", X_molecule1)
F = np.array([[1, 0, 0], [0, 1, 0], [0, 0, 1]])
delta_a, delta_b, delta_c = random_draw(0.1, 0.1), random_draw(
0.1, 0.1), random_draw(0.1, 0.1)
new_F = F + np.array([[delta_a, 0, 0], [0, delta_b, 0], [0, 0, delta_c]])
x_com1_new = np.dot(x_com1, new_F)
diff_com = x_com1_new - x_com1
#print("diff_com: ", diff_com)
new_X = []
for i, row in enumerate(X):
if i + 1 in molecule1:
new_X.append(row + diff_com)
else:
new_X.append(row)
L = atoms.get_cell()
new_L = np.dot(L, new_F)
new_atoms = atoms.copy()
new_atoms.set_positions(new_X)
new_atoms.set_cell(new_L, scale_atoms=False)
scaled_positions = new_atoms.get_scaled_positions()
molecule1_scaled_positions = [scaled_positions[i - 1] for i in molecule1]
final_scaled_positions = get_equivalent_sites(molecule1_scaled_positions,
space_group)
#print("final_scaled_positions: ", final_scaled_positions)
new_atoms.set_scaled_positions(final_scaled_positions)
print("new_atoms.get_distance(1,5): ",
new_atoms.get_distance(1, 5, mic=False, vector=True))
print("new_atoms.get_distance(1,21): ",
new_atoms.get_distance(1, 21, mic=False, vector=True))
return new_atoms
def main():
coms = {}
with open(prot_file) as fp:
for line in tqdm(fp, desc='protein file'):
if line[0] == '#': continue
protein, target, start = line.strip().split()
if protein not in coms:
coms[protein] = {}
protein_folder = os.path.join(data_root,
protein + '/structures/aligned')
if start not in coms[protein]:
file = os.path.join(protein_folder, start + '_prot.mae')
s = list(StructureReader(file))[0]
coms[protein][start] = analyze.center_of_mass(s)
file = os.path.join(protein_folder, start + '_prot.mae')
s = list(StructureReader(file))[0]
coms[protein][start] = analyze.center_of_mass(s)
with open('com_after_aligned.pkl', 'wb') as f:
pickle.dump(coms, f)
def get_com(atoms, molecule_lists):
"""
return a (n_molecules, 3) for center of mass
"""
ret_arr = []
st = ase_atoms_to_structure(atoms)
for molecule_list in molecule_lists:
com = center_of_mass(st, molecule_list)
ret_arr.append(com)
return np.array(ret_arr)
def get_pocket_res(protein, ligand):
"""
Given a co-crystallized protein and ligand, extract residues within specified distance of ligand.
Args:
protein (Biopython Structure object): receptor protein
ligand (RDKit Mol object): co-crystallized ligand
dist (float): distance cutoff for defining binding site
Returns:
key_residues (set of Biopython Residue objects): set of key binding site residues
"""
# get protein coordinates
prot_atoms = protein.getAtomIndices()
prot_coords = protein.getXYZ()
# get ligand coordinates
lig_coords = ligand.getXYZ()
kd_tree = scipy.spatial.KDTree(prot_coords)
key_pts = kd_tree.query_ball_point(lig_coords, r=DIST, p=2.0)
key_pts = set([k for l in key_pts for k in l])
return analyze.center_of_mass(protein,
list(key_pts.intersection(prot_atoms)))
def main():
args = parse_args()
print("args", args)
if args.implementation == TORCHANI:
from torchani_calculator import torchani_calculator
calculator = torchani_calculator(args.network_type, args.numb_networks)
elif args.implementation == AES_ANI:
from ani_ase import ani_ase_calculator
calculator = ani_ase_calculator(args.network_type)
elif args.implementation == KHAN:
from khan_calculator import khan_calculator
calculator = khan_calculator(args.network_type, args.khan_network,
args.numb_networks)
#assert args.cif_file.endswith('.cif') or args.cif_file.endswith('.xyz')
if args.cif_file.endswith('.mae'):
st = next(StructureReader(args.cif_file))
if args.truncate_box:
# reduced size box
a = 12.0
deletions = []
for mol in st.molecule:
rcom = center_of_mass(st, mol.getAtomIndices())
inbox = [abs(x) < a / 2 for x in rcom]
if False in [abs(x) < a / 2 for x in rcom]:
deletions.extend(mol.getAtomIndices())
st.deleteAtoms(deletions)
rd = rawdataset_from_st([st])
ase_molecs = khan_molec_to_ase_atoms(rd.all_Xs)
atoms = ase_molecs[0]
if args.truncate_box:
# hard coded box
#a = 22.6868
vecs = [[a, 0.0, 0.0], [0.0, a, 0.0], [0.0, 0.0, a]]
atoms.set_cell(vecs)
atoms.set_pbc([True, True, True])
print("after set cell", atoms.get_pbc())
if args.cif_file.endswith('.cif'):
atoms = ase_io.read(args.cif_file)
if args.box_length is not None:
atoms.set_cell([
[args.box_length, 0.0, 0.0],
[0.0, args.box_length, 0.0],
[0.0, 0.0, args.box_length],
])
atoms.set_pbc([True, True, True])
if args.supercell is not None:
atoms = build_supercell(atoms, [args.supercell] * 3)
periodic = False not in atoms.get_pbc()
print("cell")
print(atoms.get_cell())
if periodic:
#atoms.wrap()
space_group = spacegroup.get_spacegroup(atoms)
print("Space group of crystal: %s" % space_group)
print("initial unit cell")
print(atoms.cell)
# this shows how to get unique atoms and there equivalents
#scaled_positions = atoms.get_scaled_positions()
#unique_positions = space_group.unique_sites(scaled_positions)
#all_positions, symmetry_map = space_group.equivalent_sites(unique_positions)
#symmetry_groups = defaultdict(list)
#for igroup, position in zip(symmetry_map, all_positions):
# symmetry_groups[igroup].append(position)
#print("unique positions")
#for xyz in unique_positions:
# print(xyz)
#for igroup in sorted(symmetry_groups.keys()):
# print("positions in symmetry group %d" % igroup)
# for xyz in symmetry_groups[igroup]:
# print(xyz)
torsions = ()
if args.torsion is not None:
torsions = [(
float(args.torsion[0]),
int(args.torsion[1]),
int(args.torsion[2]),
int(args.torsion[3]),
)]
if args.optimize:
# cannot optimize cell until we implement a stress tensor
energy = optimize_molecule(
atoms,
calculator,
torsions=torsions,
thresh=0.05,
optimize_cell=args.relax_cell,
)
else:
start_time = time.time()
energy = single_point_energy(atoms, calculator)
print("--- %s seconds ---" % (time.time() - start_time))
if args.dynamics:
traj = molecular_dynamics(atoms,
calculator,
total_time=args.total_time,
time_step=args.time_step)
if args.cif_file.endswith('.mae'):
st = next(StructureReader(args.cif_file))
with StructureWriter("md_path.mae") as writer:
for m in traj:
st0 = st.copy()
st0.setXYZ(m.get_positions())
writer.append(st0)
else:
with open("md_path.xyz", "w") as fout:
xyz_lst = []
for atoms in traj:
xyz_lst.append("%d\n" % len(atoms))
species = atoms.get_chemical_symbols()
carts = atoms.get_positions()
for ch, xyz in zip(species, carts):
xyz_lst.append("%s %.8f %.8f %.8f" %
(ch, xyz[0], xyz[1], xyz[2]))
fout.write("\n".join(xyz_lst))
print("energy of cell (au):", energy)
if periodic:
space_group = spacegroup.get_spacegroup(atoms)
print("Final Space group of crystal: %s" % space_group)
print("energy of cell/volume (au):", energy / atoms.get_volume())
print("final unit cell")
print(atoms.cell)
print("Uncertainty %.4f (kcal/mol)" %
(calculator.uncertainty(atoms) * 627.509))
if periodic:
outfile = os.path.splitext(args.cif_file)[0] + "_optimized.cif"
else:
outfile = os.path.splitext(args.cif_file)[0] + "_optimized.xyz"
ase_io.write(outfile, atoms)
print("final structure written to file: %s" % outfile)