Skip to content

Commit a8a1031

Browse files
bobbypatonbobbypaton
committed
add tensor mode, --dp option, multi-structure CLI & output formatting
Add --tensor/--save for 3D binary occupancy tensor output with PyMOL visualization. Add --dp for configurable decimal places. Change default grid spacing to 0.05 Å. Wire multi-structure XYZ/SDF iteration into main() CLI. Display structure names from multi-structure files. Support RDKit mol objects in filename display.
1 parent dda60be commit a8a1031

3 files changed

Lines changed: 186 additions & 12 deletions

File tree

‎dbstep/Dbstep.py‎

Lines changed: 102 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -32,6 +32,8 @@ def __init__(self, *args, **kwargs):
3232
self.L, self.Bmin, self.Bmax = False, False, False
3333
# Volume Parameters
3434
self.occ_vol, self.bur_vol, self.bur_shell = False, False, False
35+
# Tensor Parameters
36+
self.tensor, self.tensor_grid = False, False
3537

3638
if "options" in kwargs:
3739
self.options = kwargs["options"]
@@ -59,6 +61,17 @@ def __init__(self, *args, **kwargs):
5961
# flag volume if buried shell requested
6062
if options.vshell:
6163
options.volume = True
64+
# tensor mode: force volume (for grid) and sterimol (for alignment)
65+
if options.tensor:
66+
if not options.atom3:
67+
sys.exit("ERROR: --tensor requires --atom3 to fully define the molecular orientation.")
68+
# default to coarser grid if user didn't explicitly set grid spacing
69+
if options.grid == 0.05:
70+
options.grid = 1.0
71+
options._tensor_only_sterimol = not options.sterimol
72+
options._tensor_only_volume = not options.volume
73+
options.volume = True
74+
options.sterimol = True
6275
# sterimol scan requires grid-based measurement for per-radius slicing
6376
if options.sterimol and options.scan and options.measure == "classic":
6477
options.measure = "grid"
@@ -82,9 +95,33 @@ def __init__(self, *args, **kwargs):
8295
r_min, r_max, r_intervals, strip_width = self._parse_scan(options)
8396

8497
# Build occupancy grid
85-
occ_grid, occ_vol, point_tree, grid_axes = self._build_grid(
98+
occ_grid, occ_vol, point_tree, grid_axes, occ_mask = self._build_grid(
8699
mol, name, options, origin, x_min, x_max, y_min, y_max, z_min, z_max)
87100

101+
# Store tensor and metadata if requested
102+
if options.tensor and occ_mask is not None:
103+
self.tensor = occ_mask.astype(int)
104+
self.tensor_grid = {
105+
"origin": [x_min, y_min, z_min],
106+
"spacing": options.grid,
107+
"shape": occ_mask.shape,
108+
"x_vals": grid_axes[0] if grid_axes else None,
109+
"y_vals": grid_axes[1] if grid_axes else None,
110+
"z_vals": grid_axes[2] if grid_axes else None,
111+
}
112+
if options.save:
113+
save_base = name if isinstance(name, str) else "tensor"
114+
np.save(save_base + "_tensor.npy", self.tensor)
115+
if not options.quiet:
116+
print(" Tensor saved to {}_tensor.npy (shape: {})".format(save_base, self.tensor.shape))
117+
118+
# Restore sterimol/volume flags if they were only set for tensor alignment
119+
if options.tensor:
120+
if getattr(options, '_tensor_only_sterimol', False):
121+
options.sterimol = False
122+
if getattr(options, '_tensor_only_volume', False):
123+
options.volume = False
124+
88125
# Print column headers (once across multi-file runs)
89126
self._print_column_header(options)
90127

@@ -105,7 +142,10 @@ def __init__(self, *args, **kwargs):
105142
if options.sterimol:
106143
cylinders.append(" CYLINDER, 0., 0., 0., 0., 0., {:5.3f}, 0.1, 1.0, 1.0, 1.0, 0., 0.0, 1.0,".format(self.L))
107144
writer.xyz_export(file, mol)
108-
writer.pymol_export(file, mol, spheres, cylinders, options.isoval, options.visv, options.viss)
145+
if options.sterimol or options.volume:
146+
writer.pymol_export(file, mol, spheres, cylinders, options.isoval, options.visv, options.viss)
147+
if options.tensor and self.tensor is not False:
148+
writer.tensor_pymol_export(file, mol, self.tensor, self.tensor_grid)
109149

110150
def _assign_surface(self, mol, file, options, origin):
111151
"""Assign VDW radii or parse density cube, translate molecule, and remove metals.
@@ -190,15 +230,16 @@ def _parse_scan(self, options):
190230
exit()
191231

192232
def _build_grid(self, mol, name, options, origin, x_min, x_max, y_min, y_max, z_min, z_max):
193-
"""Construct occupancy grid. Returns (occ_grid, occ_vol, point_tree, grid_axes)."""
233+
"""Construct occupancy grid. Returns (occ_grid, occ_vol, point_tree, grid_axes, occ_mask)."""
194234
grid_axes = None
195235
occ_grid = occ_vol = point_tree = None
236+
occ_mask = None
196237

197238
if options.surface == "vdw":
198239
# User can override grid dimensions
199240
if options.gridsize:
200241
gs = [float(val) for val in options.gridsize.replace(":", ",").split(",")]
201-
if gs[1] < x_max or gs[0] > x_min or gs[3] < y_max or gs[2] > y_min or gs[5] < z_max or gs[4] > z_min:
242+
if not options.tensor and (gs[1] < x_max or gs[0] > x_min or gs[3] < y_max or gs[2] > y_min or gs[5] < z_max or gs[4] > z_min):
202243
sys.exit("ERROR: Your molecule is larger than the gridsize you selected,\n please try again with a larger gridsize")
203244
x_min, x_max, y_min, y_max, z_min, z_max = gs
204245

@@ -209,7 +250,10 @@ def _build_grid(self, mol, name, options, origin, x_min, x_max, y_min, y_max, z_
209250
if options.volume or options.measure == "grid":
210251
if options.volume and not (options.sterimol and options.measure == "grid"):
211252
# Fast path: skip full grid construction when grid sterimol not needed
212-
occ_grid, occ_vol = sterics.occupied_direct(mol.CARTESIANS, mol.RADII, origin, x_vals, y_vals, z_vals, options)
253+
if options.tensor:
254+
occ_grid, occ_vol, occ_mask = sterics.occupied_direct(mol.CARTESIANS, mol.RADII, origin, x_vals, y_vals, z_vals, options, return_mask=True)
255+
else:
256+
occ_grid, occ_vol = sterics.occupied_direct(mol.CARTESIANS, mol.RADII, origin, x_vals, y_vals, z_vals, options)
213257
point_tree = None
214258
grid_axes = (x_vals, y_vals, z_vals)
215259
else:
@@ -232,7 +276,7 @@ def _build_grid(self, mol, name, options, origin, x_min, x_max, y_min, y_max, z_
232276
if options.volume:
233277
grid, point_tree = sterics.resize_grid(x_max, y_max, z_max, x_min, y_min, z_min, options, mol)
234278

235-
return occ_grid, occ_vol, point_tree, grid_axes
279+
return occ_grid, occ_vol, point_tree, grid_axes, occ_mask
236280

237281
def _print_column_header(self, options):
238282
"""Print column headers once across multi-file runs."""
@@ -265,7 +309,18 @@ def _compute(self, mol, file, options, origin, occ_grid, occ_vol, point_tree, gr
265309
else:
266310
occ_dist2 = None
267311

268-
fname = os.path.basename(file)
312+
if isinstance(file, str):
313+
fname = os.path.basename(file)
314+
else:
315+
try:
316+
fname = file.GetProp("_Name")
317+
except Exception:
318+
fname = "rdkit_mol"
319+
# Use structure name from multi-XYZ file if available
320+
if hasattr(mol, 'structure_name') and mol.structure_name:
321+
# Extract clean name from comment line (first word, strip extension)
322+
sname = mol.structure_name.split()[0]
323+
fname = os.path.splitext(sname)[0] if '.' in sname else sname
269324
fw = dbstep._file_col_width
270325

271326
for rad in np.linspace(r_min, r_max, r_intervals):
@@ -293,21 +348,27 @@ def _compute(self, mol, file, options, origin, occ_grid, occ_vol, point_tree, gr
293348
cylinders.extend(cyl)
294349

295350
# Tabulate result
351+
dp = options.dp
352+
vfmt = "{{:10.{}f}}".format(dp)
353+
rfmt = "{{:6.{}f}}".format(dp)
296354
if options.volume and options.sterimol:
297355
if options.pymol:
298356
spheres.append(" SPHERE, 0.000, 0.000, 0.000, {:5.3f},".format(rad))
299357
if not options.quiet:
300358
atom2_str = ",".join(str(a) for a in options.spec_atom_2)
301-
print(" {:>{fw}} {:>6} {:>6} {:6.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f} {:10.2f}".format(fname, options.spec_atom_1, atom2_str, rad, occ_vol, bur_vol, bur_shell, Bmin, Bmax, L, fw=fw))
359+
fmt = " {:>" + str(fw) + "} {:>6} {:>6} " + rfmt + " " + " ".join([vfmt] * 6)
360+
print(fmt.format(fname, options.spec_atom_1, atom2_str, rad, occ_vol, bur_vol, bur_shell, Bmin, Bmax, L))
302361
elif options.volume:
303362
if options.pymol:
304363
spheres.append(" SPHERE, 0.000, 0.000, 0.000, {:5.3f},".format(rad))
305364
if not options.quiet:
306-
print(" {:>{fw}} {:>6} {:6.2f} {:10.2f} {:10.2f} {:10.2f}".format(fname, options.spec_atom_1, rad, occ_vol, bur_vol, bur_shell, fw=fw))
365+
fmt = " {:>" + str(fw) + "} {:>6} " + rfmt + " " + " ".join([vfmt] * 3)
366+
print(fmt.format(fname, options.spec_atom_1, rad, occ_vol, bur_vol, bur_shell))
307367
elif options.sterimol:
308368
if not options.quiet:
309369
atom2_str = ",".join(str(a) for a in options.spec_atom_2)
310-
print(" {:>{fw}} {:>6} {:>6} {:10.2f} {:10.2f} {:10.2f}".format(fname, options.spec_atom_1, atom2_str, Bmin, Bmax, L, fw=fw))
370+
fmt = " {:>" + str(fw) + "} {:>6} {:>6} " + " ".join([vfmt] * 3)
371+
print(fmt.format(fname, options.spec_atom_1, atom2_str, Bmin, Bmax, L))
311372

312373
# Store results on self
313374
if occ_vol is not None:
@@ -390,7 +451,7 @@ def set_options(kwargs):
390451
var_dict = {
391452
"verbose": ["verbose", False],
392453
"v": ["verbose", False],
393-
"grid": ["grid", 0.1],
454+
"grid": ["grid", 0.05],
394455
"scalevdw": ["SCALE_VDW", 1.0],
395456
"noH": ["noH", False],
396457
"nometals": ["no_metals", False],
@@ -408,6 +469,9 @@ def set_options(kwargs):
408469
"debug": ["debug", False],
409470
"b": ["volume", False],
410471
"volume": ["volume", False],
472+
"tensor": ["tensor", False],
473+
"t": ["tensor", False],
474+
"save": ["save", False],
411475
"vshell": ["vshell", False],
412476
"pymol": ["pymol", False],
413477
"quiet": ["quiet", False],
@@ -416,6 +480,7 @@ def set_options(kwargs):
416480
"gridsize": ["gridsize", False],
417481
"measure": ["measure", "classic"],
418482
"pos": ["pos", False],
483+
"dp": ["dp", 2],
419484
"graph": ["graph", False],
420485
"fg": ["shared_fg", False],
421486
"shared_fg": ["shared_fg", False],
@@ -424,6 +489,7 @@ def set_options(kwargs):
424489
"voltype": ["voltype", "crippen"],
425490
"visv": ["visv", "circle"],
426491
"viss": ["viss", False],
492+
"structure": ["structure", None],
427493
}
428494

429495
for key in var_dict:
@@ -457,6 +523,7 @@ def main():
457523
parser.add_option("--noH", dest="noH", action="store_true", help="Exclude hydrogen atoms from steric measurements", default=False)
458524
parser.add_option("--nometals", dest="no_metals", action="store_true", help="Exclude metal atoms from steric measurements", default=False)
459525
parser.add_option("--norot", dest="norot", action="store_true", help="Do not rotate the molecule (use if structures have been pre-aligned)", default=False)
526+
parser.add_option("--dp", dest="dp", action="store", type="int", help="Number of decimal places for output values (default: 2)", default=2, metavar="dp")
460527
parser.add_option("--pos", dest="pos", action="store_true", help="Measure Sterimol parameters in positive direction (from atom1 toward atom2)", default=False)
461528
parser.add_option("--quiet", dest="quiet", action="store_true", help="Suppress all print output", default=False)
462529
parser.add_option("--radii", dest="radii", action="store", choices=["bondi", "charry-tkatchenko"], help="VDW radii set: bondi or charry-tkatchenko (default: bondi)", default="bondi")
@@ -467,6 +534,8 @@ def main():
467534
parser.add_option("-s", "--sterimol", dest="sterimol", action="store_true", help="Compute Sterimol parameters (L, Bmin, Bmax)", default=False)
468535
parser.add_option("--measure", dest="measure", action="store", choices=["classic", "grid"], help="Sterimol method: classic (Verloop, default) or grid-based", default="classic", metavar="measure")
469536
parser.add_option("--surface", dest="surface", action="store", choices=["vdw", "density"], help="Surface type: Bondi VDW radii or density cube file (default: vdw)", default="vdw", metavar="surface")
537+
parser.add_option("-t", "--tensor", dest="tensor", action="store_true", help="Return 3D binary occupancy tensor (requires --atom1, --atom2, --atom3)", default=False)
538+
parser.add_option("--save", dest="save", action="store_true", help="Save tensor to .npy file (use with --tensor)", default=False)
470539
parser.add_option("-b", "--vbur", dest="volume", action="store_true", help="Calculate buried volume of input molecule", default=False)
471540
parser.add_option("-v", "--verbose", dest="verbose", action="store_true", help="Print verbose output", default=False)
472541
parser.add_option("--viss", dest="viss", action="store_true", help="Visualize Sterimol Bmin and Bmax in PyMOL as circle outlines", default=False)
@@ -480,6 +549,10 @@ def main():
480549
options.SCALE_VDW = 1.17
481550
options.noH = True
482551

552+
# Tensor mode: default to coarser grid spacing if user didn't set --grid
553+
if options.tensor and options.grid == 0.05:
554+
options.grid = 1.0
555+
483556
# Sterimol scan requires grid-based measurement for per-radius slicing
484557
if options.sterimol and options.scan and options.measure == "classic":
485558
options.measure = "grid"
@@ -544,6 +617,24 @@ def main():
544617
vec_df[numeric_cols] = vec_df[numeric_cols].round(2)
545618
vec_df.to_csv(file.split(".")[0] + "_2d_output.csv", index=False)
546619
else:
620+
# Detect multi-structure files
621+
_, ext = os.path.splitext(file)
622+
if ext == ".xyz":
623+
structures = parse_data.get_xyz_structures(file)
624+
if len(structures) > 1:
625+
for idx in range(len(structures)):
626+
options.structure = idx
627+
dbstep(file, options=options)
628+
options.structure = None
629+
continue
630+
elif ext in [".sdf", ".mol"]:
631+
structures = parse_data.get_sdf_structures(file)
632+
if len(structures) > 1:
633+
for idx in range(len(structures)):
634+
options.structure = idx
635+
dbstep(file, options=options)
636+
options.structure = None
637+
continue
547638
dbstep(file, options=options)
548639

549640
if dbstep._column_width and not options.quiet:

‎dbstep/sterics.py‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@ def count_grid_points_in_sphere(x_vals, y_vals, z_vals, origin, R):
2626
return count
2727

2828

29-
def occupied_direct(coords, radii, origin, x_vals, y_vals, z_vals, options):
29+
def occupied_direct(coords, radii, origin, x_vals, y_vals, z_vals, options, return_mask=False):
3030
"""Generate occupied grid points without allocating the full meshgrid.
3131
Per-atom work is vectorized within the atom's bounding box via broadcasting.
3232
The boolean mask provides automatic deduplication where atom spheres overlap."""
@@ -67,6 +67,8 @@ def occupied_direct(coords, radii, origin, x_vals, y_vals, z_vals, options):
6767
occ_vol = n_occ * spacing**3
6868
if options.verbose:
6969
print(" Molecular volume is {:5.4f} Ang^3".format(occ_vol))
70+
if return_mask:
71+
return occ_grid, occ_vol, occ_mask
7072
return occ_grid, occ_vol
7173

7274

‎dbstep/writer.py‎

Lines changed: 81 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -174,6 +174,87 @@ def pymol_export(file, mol, spheres, cylinders, isoval, visv, viss):
174174
log.Writeonlyfile('cmd.set("orthoscopic", "on")')
175175

176176

177+
def tensor_pymol_export(file, mol, tensor, tensor_grid):
178+
"""Outputs a PyMOL script to visualize the 3D occupancy tensor as a wireframe grid with occupied voxels"""
179+
import numpy as np
180+
base, ext = os.path.splitext(file)
181+
182+
log = Logger(base, "py", "tensor")
183+
log.Writeonlyfile("from pymol.cgo import *")
184+
log.Writeonlyfile("from pymol import cmd\n")
185+
186+
x_vals = tensor_grid["x_vals"]
187+
y_vals = tensor_grid["y_vals"]
188+
z_vals = tensor_grid["z_vals"]
189+
spacing = tensor_grid["spacing"]
190+
191+
# Wireframe grid using CGO LINES
192+
log.Writeonlyfile("grid = [")
193+
log.Writeonlyfile(" BEGIN, LINES,")
194+
log.Writeonlyfile(" COLOR, 0.7, 0.7, 0.7,")
195+
196+
# Lines parallel to X axis
197+
for y in y_vals:
198+
for z in z_vals:
199+
log.Writeonlyfile(" VERTEX, {:.4f}, {:.4f}, {:.4f},".format(x_vals[0], y, z))
200+
log.Writeonlyfile(" VERTEX, {:.4f}, {:.4f}, {:.4f},".format(x_vals[-1], y, z))
201+
202+
# Lines parallel to Y axis
203+
for x in x_vals:
204+
for z in z_vals:
205+
log.Writeonlyfile(" VERTEX, {:.4f}, {:.4f}, {:.4f},".format(x, y_vals[0], z))
206+
log.Writeonlyfile(" VERTEX, {:.4f}, {:.4f}, {:.4f},".format(x, y_vals[-1], z))
207+
208+
# Lines parallel to Z axis
209+
for x in x_vals:
210+
for y in y_vals:
211+
log.Writeonlyfile(" VERTEX, {:.4f}, {:.4f}, {:.4f},".format(x, y, z_vals[0]))
212+
log.Writeonlyfile(" VERTEX, {:.4f}, {:.4f}, {:.4f},".format(x, y, z_vals[-1]))
213+
214+
log.Writeonlyfile(" END,")
215+
log.Writeonlyfile("]")
216+
log.Writeonlyfile('cmd.load_cgo(grid, "grid")')
217+
log.Writeonlyfile('cmd.set("cgo_line_width", 1.0, "grid")\n')
218+
219+
# Occupied voxels as spheres
220+
occ_indices = np.argwhere(tensor)
221+
r = spacing / 2.0
222+
if len(occ_indices) > 0:
223+
log.Writeonlyfile("occupied = [")
224+
for idx in occ_indices:
225+
x = x_vals[idx[0]]
226+
y = y_vals[idx[1]]
227+
z = z_vals[idx[2]]
228+
log.Writeonlyfile(" COLOR, 0.2, 0.6, 1.0,")
229+
log.Writeonlyfile(" SPHERE, {:.4f}, {:.4f}, {:.4f}, {:.4f},".format(x, y, z, r))
230+
log.Writeonlyfile("]")
231+
log.Writeonlyfile('cmd.load_cgo(occupied, "occupied")')
232+
log.Writeonlyfile('cmd.set("cgo_transparency", 0.4, "occupied")\n')
233+
234+
# Empty (unoccupied) voxels as spheres
235+
empty_indices = np.argwhere(tensor == 0)
236+
if len(empty_indices) > 0:
237+
log.Writeonlyfile("empty = [")
238+
for idx in empty_indices:
239+
x = x_vals[idx[0]]
240+
y = y_vals[idx[1]]
241+
z = z_vals[idx[2]]
242+
log.Writeonlyfile(" COLOR, 0.9, 0.4, 0.4,")
243+
log.Writeonlyfile(" SPHERE, {:.4f}, {:.4f}, {:.4f}, {:.4f},".format(x, y, z, r))
244+
log.Writeonlyfile("]")
245+
log.Writeonlyfile('cmd.load_cgo(empty, "empty")')
246+
log.Writeonlyfile('cmd.set("cgo_transparency", 0.7, "empty")')
247+
log.Writeonlyfile('cmd.disable("empty")\n')
248+
249+
# Load transformed molecule
250+
full_path = os.path.abspath(file)
251+
name, ext = os.path.splitext(full_path)
252+
log.Writeonlyfile('cmd.load("' + name + '_transform.xyz")')
253+
log.Writeonlyfile('cmd.show_as("spheres", "' + base.split("/")[-1] + '_transform")')
254+
log.Writeonlyfile('cmd.set("sphere_transparency", 0.5)')
255+
log.Writeonlyfile('cmd.set("orthoscopic", "on")')
256+
257+
177258
def xyz_export(file, mol):
178259
"""Write xyz coordinates of molecule to file"""
179260
name, ext = os.path.splitext(file)

0 commit comments

Comments
 (0)