import os
import re
from past.utils import old_div
import numpy as np
import schrodinger.infra.mm as mm
import schrodinger.structure as structure
import schrodinger.utils.cmdline as cmdline
[docs]class CNSData():
[docs] def __init__(self, fname, create_grid_st=True):
self._cell_handle = None
self.grid_st = None
self._parse_cns(fname)
if create_grid_st:
self._prepare_grid_st()
def __del__(self):
if self._cell_handle:
mm.mmct_delete_distance_cell(self._cell_handle)
def _parse_cns(self, fname):
with open(fname) as fh:
lines = fh.readlines()
n = int(lines[1]) + 2
line = lines[n]
self.dims = []
for i in range(0, 72, 8):
self.dims.append(int(line[i:i + 8]))
self.n_grid = np.array([
self.dims[2] - self.dims[1] + 1, self.dims[5] - self.dims[4] + 1,
self.dims[8] - self.dims[7] + 1
])
self.min_dim = np.array([self.dims[1], self.dims[4], self.dims[7]])
self.max_dim = np.array([self.dims[2], self.dims[5], self.dims[8]])
n += 1
line = lines[n]
self.cell_size = []
for i in range(0, 72, 12):
self.cell_size.append(float(line[i:i + 12]))
self.grid_width = np.array([
old_div(self.cell_size[0], self.dims[0]),
old_div(self.cell_size[1], self.dims[3]),
old_div(self.cell_size[2], self.dims[6])
])
n += 2
values = []
for line in lines[n:]:
if line.strip() == "-9999":
break
a = re.findall(r'^\s*[0-9]+$', line)
if len(a) > 0:
continue
i = 0
while (i + 12 < len(line)):
values.append(float(line[i:i + 12]))
i += 12
self.grid_data = np.array(values).reshape(self.n_grid[2],
self.n_grid[1],
self.n_grid[0])
def _prepare_grid_st(self):
grid_data = self.grid_data.transpose().ravel()
self.grid_st = structure.create_new_structure(
self.n_grid[0] * self.n_grid[1] * self.n_grid[2])
grid_location = []
for i in range(self.min_dim[0], self.max_dim[0] + 1):
x = i * self.grid_width[0]
for j in range(self.min_dim[1], self.max_dim[1] + 1):
y = j * self.grid_width[1]
for k in range(self.min_dim[2], self.max_dim[2] + 1):
z = k * self.grid_width[2]
grid_location.append([x, y, z])
self.grid_st.setXYZ(np.array(grid_location))
for a in self.grid_st.atom:
a.name = "GR%d" % a.index
a.atom_type = 61
a.temperature_factor = grid_data[a.index - 1]
self._cell_handle = mm.mmct_create_distance_cell(
self.grid_st.handle, self.grid_width.max())
[docs] def write(self, cns_fname, remark=''):
fh = open(cns_fname, 'w')
s = '\n'
s += '%d\n' % len(remark.splitlines())
if remark:
s += remark
for i in range(len(self.dims)):
s += '%8d' % self.dims[i]
s += '\n'
for i in range(len(self.cell_size)):
s += '%12.5E' % self.cell_size[i]
s += '\n'
s += 'ZYX\n'
shape = self.grid_data.shape
for i in range(shape[0]):
n = 1
s += '%8d\n' % i
for j in range(shape[1]):
for k in range(shape[2]):
s += '%12.5E' % self.grid_data[i, j, k]
if n % 6 == 0:
s += '\n'
n += 1
if ((n - 1) % 6) != 0:
s += '\n'
s += '%8d\n' % -9999
s += '%12.5E%12.5E\n' % (np.mean(self.grid_data), np.std(
self.grid_data))
fh.write(s)
fh.close()
[docs] def get_grid_data(self, x, y, z):
list_handle = mm.mmct_query_distance_cell(self._cell_handle, x, y, z)
list_size = mm.mmlist_get_size(list_handle)
if list_size == 0:
print('cannot find grid point.')
mm.mmlist_delete(list_handle)
return None, None
min_atom = 0
min_distance = 1.0E100
b = np.array([x, y, z])
for i in range(list_size):
iatom = mm.mmlist_get(list_handle, i)
a = self.grid_st.atom[iatom].xyz
dist = np.linalg.norm(a - b)
if dist < min_distance:
min_distance = dist
min_atom = self.grid_st.atom[iatom]
mm.mmlist_delete(list_handle)
return min_atom.xyz, min_atom.temperature_factor
[docs]def write_cns(grid, length, spacing, filename, center=(0., 0., 0.)):
"""
This function writes out a volumetric file (.cns) in CNS/XPLOR format
:param grid: a 3 dimensional `numpy.array` of `float`
:type grid: `numpy.array`
:param length: the length of the box (x,y,z)
:type length: `list` of `float`, or `numpy.array`
:param spacing: the spacing of the grid (x,y,z)
:type spacing: `list` of `float`, or `numpy.array`
:param filename: filename to write the grid to
:type filename: `str`
:param center: geometric center
:type center: iterable of `float`, or 1x3 `numpy.array`
"""
grid = grid.astype(np.float32)
spacing = np.asarray(spacing, dtype=np.float32)
center = np.ceil(np.asarray(center, dtype=np.float32) / spacing)
with open(filename, 'w') as fh:
# Write header
fh.write("\n%8i\nREMARKS Map type: CNS map Regular\n" % 1)
# Write grid dimensions
shape = np.array(grid.shape, np.float32)
origin_by_spacing = center - shape / 2 # origin is at the center
origin_by_spacing = np.ceil(origin_by_spacing).astype(np.int32)
grid_stop = origin_by_spacing + np.array(shape) - 1
for v1, v2, v3 in zip(shape, origin_by_spacing, grid_stop):
fh.write("%8i%8i%8i" % (v1, v2, v3))
fh.write("\n")
# Write box length and angles (degrees).
fh.write("%12.5e%12.5e%12.5e" % (length[0], length[1], length[2]))
fh.write("%12.5e%12.5e%12.5e\n" % (90.0, 90.0, 90.0))
fh.write("ZYX\n")
# Write grid values
for c in range(grid.shape[2]):
fh.write("%8i\n" % c)
count = 0
for b in range(grid.shape[1]):
for a in range(grid.shape[0]):
fh.write("%12.5e" % grid[a][b][c])
count += 1
if (count % 6 == 0):
fh.write("\n")
if (count % 6 != 0):
fh.write("\n")
# write average and std. dev.
fh.write("%8i\n" % (-9999))
fh.write("%12.5e%12.5e\n" % (grid.mean(), grid.std()))
if (__name__ == "__main__"):
usage = '''
$SCHRODINGER/run %prog -wm watermap.mae -ligand ligand.mae -cavity watermap-tpi.cns output.mae
Description:
%prog compute scores from watermap and cavity
Example to retain ligand:
'''
parser = cmdline.SingleDashOptionParser(usage)
options, args = parser.parse_args()
input_cns = args[0]
output_cns = args[1]
output_st = args[2]
if os.path.exists(output_cns):
os.remove(output_cns)
if os.path.exists(output_st):
os.remove(output_st)
cns = CNSData(input_cns)
cns.write(output_cns)
cns.grid_st.write(output_st)