Coverage for cosmolayer/cosmosac/visualize.py: 88%
195 statements
« prev ^ index » next coverage.py v7.15.4, created at 2026-08-26 00:09 +0000
1"""
2.. module:: cosmolayer.cosmosac.visualize
3 :synopsis: Visualize COSMO-SAC surface segments.
5.. functionauthor:: Charlles Abreu <craabreu@gmail.com>
6"""
8import argparse
9import pathlib
11import cmap
12import networkx as nx
13import numpy as np
14import open3d as o3d
15import periodictable as pt
17from cosmolayer.cosmosac import Component
19RADII_MULTIPLIERS: tuple[float, float, float] = (1.5, 2.5, 4.0)
21ELEMENT_COLORS = { # https://pymolwiki.org/Color_Values
22 "Br": (0.650980392, 0.160784314, 0.160784314),
23 "C": (0.2, 1.0, 0.2),
24 "Cl": (0.121568627, 0.941176471, 0.121568627),
25 "F": (0.701960784, 1.0, 1.0),
26 "H": (0.9, 0.9, 0.9),
27 "I": (0.580392157, 0.0, 0.580392157),
28 "N": (0.2, 0.2, 1.0),
29 "O": (1.0, 0.3, 0.3),
30 "P": (1.0, 0.501960784, 0.0),
31 "Si": (0.941176471, 0.784313725, 0.627450980),
32 "S": (0.9, 0.775, 0.25),
33}
36TOLERANCE: float = 1e-10
37DOT_PRODUCT_TOLERANCE: float = 0.9
38X_AXIS: np.ndarray = np.array([1.0, 0.0, 0.0])
39Y_AXIS: np.ndarray = np.array([0.0, 1.0, 0.0])
40Z_AXIS: np.ndarray = np.array([0.0, 0.0, 1.0])
43def estimate_vdw_radius(element: str) -> float:
44 covalent_radius = pt.elements.symbol(element).covalent_radius
45 if covalent_radius is None:
46 raise ValueError(f"Unknown covalent radius for element {element!r}")
47 return float(covalent_radius) + 0.8 # Å
50def create_atom_spheres(
51 component: Component,
52 radius_scale: float,
53 resolution: int = 40,
54 default_color: tuple[float, float, float] = (0.7, 0.7, 0.7),
55) -> list[o3d.geometry.TriangleMesh]:
56 atom_df = component.atom_data
57 spheres: list[o3d.geometry.TriangleMesh] = []
58 for element, x, y, z in zip(
59 atom_df["element"],
60 atom_df["x"],
61 atom_df["y"],
62 atom_df["z"],
63 strict=True,
64 ):
65 element_name = str(element).strip()
66 radius = estimate_vdw_radius(element_name) * radius_scale
67 sphere = o3d.geometry.TriangleMesh.create_sphere(
68 radius=radius,
69 resolution=resolution,
70 )
71 sphere.compute_vertex_normals()
72 sphere.translate((float(x), float(y), float(z)))
73 rgb = np.array(ELEMENT_COLORS.get(element_name, default_color))
74 sphere.paint_uniform_color(rgb)
75 spheres.append(sphere)
76 return spheres
79def compute_rotation_matrix(
80 original_axis: np.ndarray, target_axis: np.ndarray, normalize: bool = False
81) -> np.ndarray:
82 """Rodrigues' rotation formula between two unit-direction vectors."""
83 if normalize:
84 original_axis = original_axis / np.linalg.norm(original_axis)
85 target_axis = target_axis / np.linalg.norm(target_axis)
86 v = np.cross(original_axis, target_axis)
87 c = original_axis.dot(target_axis)
88 s2 = v.dot(v)
89 if s2 < TOLERANCE:
90 if c > 0:
91 return np.eye(3) # Parallel (c ≈ 1) → identity
92 arbitrary = X_AXIS if abs(original_axis[0]) < DOT_PRODUCT_TOLERANCE else Y_AXIS
93 arbitrary -= original_axis * original_axis.dot(arbitrary)
94 orthogonal = arbitrary / np.linalg.norm(arbitrary)
95 return 2.0 * np.outer(orthogonal, orthogonal) - np.eye(3)
96 kmat = np.array([[0, -v[2], v[1]], [v[2], 0, -v[0]], [-v[1], v[0], 0]])
97 rotation: np.ndarray = np.eye(3) + kmat + ((1 - c) / s2) * kmat @ kmat
98 return rotation
101def create_bond_sticks(
102 component: Component,
103 atom_radius_scale: float,
104 bond_radius: float,
105 resolution: int = 100,
106 default_color: tuple[float, float, float] = (0.7, 0.7, 0.7),
107) -> list[o3d.geometry.TriangleMesh]:
108 atom_df = component.atom_data
109 coords = atom_df[["x", "y", "z"]].values
110 elements = atom_df["element"].values
111 radii = atom_df["element"].apply(estimate_vdw_radius).values * atom_radius_scale
112 bonds = component.bonds
113 cylinders: list[o3d.geometry.TriangleMesh] = []
114 for i, j in bonds:
115 vector = coords[j] - coords[i]
116 length = np.linalg.norm(vector)
117 if length < radii[i] + radii[j]:
118 continue
119 axis = vector / length
120 rotation = compute_rotation_matrix(Z_AXIS, axis)
121 midpoint = (coords[i] + coords[j] + (radii[i] - radii[j]) * axis) / 2
122 for k in (i, j):
123 cylinder = o3d.geometry.TriangleMesh.create_cylinder(
124 radius=bond_radius,
125 height=np.linalg.norm(coords[k] - midpoint),
126 resolution=resolution,
127 )
128 cylinder.rotate(rotation, center=np.zeros(3))
129 cylinder.translate((coords[k] + midpoint) / 2)
130 cylinder.compute_vertex_normals()
131 rgb = np.array(ELEMENT_COLORS.get(elements[k], default_color))
132 cylinder.paint_uniform_color(rgb)
133 cylinders.append(cylinder)
134 return cylinders
137def ball_pivoting_algorithm(
138 points: np.ndarray,
139 normals: np.ndarray,
140 vertex_rgb: np.ndarray,
141 radii_multipliers: tuple[float, float, float],
142) -> tuple[o3d.geometry.TriangleMesh, np.ndarray]:
143 pcd = o3d.geometry.PointCloud(o3d.utility.Vector3dVector(points))
144 pcd.normals = o3d.utility.Vector3dVector(normals)
146 spacing = np.asarray(pcd.compute_nearest_neighbor_distance()).mean().item()
147 radii = o3d.utility.DoubleVector([m * spacing for m in radii_multipliers])
149 mesh_bpa = o3d.geometry.TriangleMesh.create_from_point_cloud_ball_pivoting(
150 pcd, radii
151 )
153 mesh_bpa.remove_degenerate_triangles()
154 mesh_bpa.remove_duplicated_triangles()
155 mesh_bpa.remove_non_manifold_edges()
156 mesh_bpa.remove_unreferenced_vertices()
158 kdtree = o3d.geometry.KDTreeFlann(pcd)
159 indices = np.empty(len(mesh_bpa.vertices), dtype=int)
160 for vi, v in enumerate(mesh_bpa.vertices):
161 _, idx, _ = kdtree.search_knn_vector_3d(v, 1)
162 indices[vi] = int(idx[0])
164 vertex_rgb = vertex_rgb[indices]
165 mesh_bpa.vertex_colors = o3d.utility.Vector3dVector(vertex_rgb)
167 return mesh_bpa, indices
170def find_loops(
171 mesh: o3d.geometry.TriangleMesh, edge_color: str
172) -> list[o3d.geometry.LineSet]:
173 graph = nx.Graph()
174 for triangle in mesh.triangles:
175 _, j, k = map(int, triangle)
176 graph.add_edge(j, k)
177 loops = nx.cycle_basis(graph)
179 vertices = np.asarray(mesh.vertices, dtype=float)
180 linesets: list[o3d.geometry.LineSet] = []
182 for loop in loops:
183 idx = np.asarray(loop + [loop[0]], dtype=int)
184 pts = vertices[idx]
185 lines = np.column_stack(
186 [np.arange(len(idx) - 1), np.arange(1, len(idx))]
187 ).astype(np.int32)
188 lineset = o3d.geometry.LineSet(
189 points=o3d.utility.Vector3dVector(pts),
190 lines=o3d.utility.Vector2iVector(lines),
191 )
192 if edge_color is not None:
193 rgb = np.asarray(cmap.Color(edge_color))[:3]
194 lineset.paint_uniform_color(rgb)
195 linesets.append(lineset)
197 return linesets
200def geodesic_centroid(center: np.ndarray, *vertices: np.ndarray) -> np.ndarray:
201 num_vertices = len(vertices)
202 vectors = [v - center for v in vertices]
203 norms = [np.linalg.norm(v) for v in vectors]
204 radius = sum(norms) / num_vertices
205 mean_vector = sum(vectors) / num_vertices
206 centroid: np.ndarray = center + radius * mean_vector / np.linalg.norm(mean_vector)
207 return centroid
210def surface_tessellation(
211 component: Component,
212 original_charge_densities: bool = False,
213 interpolated_colors: bool = False,
214 colormap: str = "jet",
215) -> o3d.geometry.TriangleMesh:
216 segment_data = component.segment_data
217 atom_data = component.atom_data
218 sigma_grid = component.sigma_grid
219 vmin, vmax = sigma_grid[0], sigma_grid[-1]
220 sigmas = segment_data[
221 "sigma" if original_charge_densities else "sigma_avg"
222 ].values.clip(vmin, vmax)
224 atom_coords = np.stack(
225 [segment_data["atom"].map(atom_data[axis]).values for axis in "xyz"], axis=1
226 )
227 pts = segment_data[["x", "y", "z"]].values
228 displacements = pts - atom_coords
229 normals = displacements / np.linalg.norm(displacements, axis=1, keepdims=True)
231 normalized_sigmas = (sigmas.clip(vmin, vmax) - vmin) / (vmax - vmin)
232 mapper = cmap.Colormap(colormap)
233 vertex_rgb = mapper(normalized_sigmas)[:, :3]
235 mesh_bpa, indices = ball_pivoting_algorithm(
236 pts, normals, vertex_rgb, RADII_MULTIPLIERS
237 )
239 if interpolated_colors:
240 return mesh_bpa
242 vertices = np.asarray(mesh_bpa.vertices, dtype=float)
243 triangles = np.asarray(mesh_bpa.triangles, dtype=int)
244 colors = np.asarray(mesh_bpa.vertex_colors, dtype=float)
245 atoms = segment_data["atom"].values[indices]
247 new_vertices = vertices.tolist()
248 new_colors = colors.tolist()
250 def add_vertex(v: np.ndarray, c: np.ndarray) -> int:
251 idx = len(new_vertices)
252 new_vertices.append(v)
253 new_colors.append(c)
254 return idx
256 midpoint_cache: dict[tuple[int, int], int] = {}
258 def midpoint_vertices(i: int, j: int) -> tuple[int, int]:
259 if (i, j) in midpoint_cache:
260 return midpoint_cache[(i, j)], midpoint_cache[(j, i)]
262 if atoms[i] == atoms[j]:
263 midpoint = geodesic_centroid(
264 atom_coords[atoms[i]], vertices[i], vertices[j]
265 )
266 else:
267 midpoint = (vertices[i] + vertices[j]) / 2
269 mij = midpoint_cache[(i, j)] = add_vertex(midpoint, colors[i])
270 mji = midpoint_cache[(j, i)] = add_vertex(midpoint, colors[j])
271 return mij, mji
273 new_triangles = []
275 for triangle in triangles:
276 i, j, k = map(int, triangle)
277 mij, mji = midpoint_vertices(i, j)
278 mjk, mkj = midpoint_vertices(j, k)
279 mik, mki = midpoint_vertices(i, k)
281 if atoms[i] == atoms[j] == atoms[k]:
282 centroid = geodesic_centroid(
283 atom_coords[atoms[i]], vertices[i], vertices[j], vertices[k]
284 )
285 else:
286 centroid = (vertices[i] + vertices[j] + vertices[k]) / 3
288 mijk = add_vertex(centroid, colors[i])
289 mjki = add_vertex(centroid, colors[j])
290 mkij = add_vertex(centroid, colors[k])
292 new_triangles += [
293 [i, mij, mijk],
294 [i, mik, mijk],
295 [j, mji, mjki],
296 [j, mjk, mjki],
297 [k, mkj, mkij],
298 [k, mki, mkij],
299 ]
301 mesh = o3d.geometry.TriangleMesh(
302 vertices=o3d.utility.Vector3dVector(new_vertices),
303 triangles=o3d.utility.Vector3iVector(new_triangles),
304 )
305 mesh.vertex_colors = o3d.utility.Vector3dVector(new_colors)
306 mesh.compute_vertex_normals()
308 return mesh
311def generate_geometries(
312 component: Component,
313 original_charge_densities: bool = False,
314 use_continuous_colors: bool = False,
315 colormap: str = "jet",
316 segment_edge_color: str | None = None,
317) -> tuple[o3d.geometry.Geometry3D, ...]:
318 """Build Open3D geometries for visualizing a component's COSMO surface.
320 Returns a tuple of Open3D geometries:
322 (1) a tessellated surface mesh colored by screening charge density;
323 (2) optionally, segment-boundary loop line sets when ``segment_edge_color`` is set;
324 (3) atom spheres; and
325 (4) bond sticks.
327 Parameters
328 ----------
329 component : Component
330 The molecular component whose COSMO surface is to be visualized.
331 original_charge_densities : bool, optional
332 If ``True``, color the surface using the original (unsmoothed) segment
333 charge densities instead of the distance-weighted averages. Default is
334 ``False``.
335 use_continuous_colors : bool, optional
336 If ``True``, use interpolated colors across the surface; otherwise,
337 segments are uniformly colored. Default is ``False``.
338 colormap : str, optional
339 Name of the colormap used to map charge density to color (e.g.
340 ``"jet"``, ``"viridis"``). Default is ``"jet"``.
341 segment_edge_color : str or None, optional
342 Color name for the edges between segments (e.g. ``"black"``).
343 If ``None`` or if ``use_continuous_colors`` is ``True``, no edge
344 loops are drawn. Default is ``None``.
346 Returns
347 -------
348 tuple of Geometry3D
349 A sequence of Open3D geometries: mesh, loops (if any), atom spheres,
350 and bond sticks.
352 Examples
353 --------
354 >>> from importlib.resources import files
355 >>> from cosmolayer.cosmosac import Component
356 >>> from cosmolayer.cosmosac.visualize import generate_geometries
357 >>> path = files("cosmolayer.data") / "C=C(N)O.cosmo"
358 >>> component = Component(path.read_text())
359 >>> geometries = generate_geometries(component)
360 >>> len(geometries) >= 1
361 True
362 >>> type(geometries[0]).__name__
363 'TriangleMesh'
364 >>> geometries_loops = generate_geometries(component, segment_edge_color="black")
365 >>> len(geometries_loops) > len(geometries)
366 True
367 """
368 mesh = surface_tessellation(
369 component,
370 original_charge_densities,
371 use_continuous_colors,
372 colormap,
373 )
374 if segment_edge_color is None or use_continuous_colors:
375 loops = []
376 else:
377 loops = find_loops(mesh, segment_edge_color)
378 atom_spheres = create_atom_spheres(component, 0.25)
379 bond_sticks = create_bond_sticks(component, 0.25, 0.1)
380 return (mesh, *loops, *atom_spheres, *bond_sticks)
383def get_parser() -> argparse.ArgumentParser:
384 """Return the argument parser for cosmoviz (used by sphinx-argparse)."""
385 parser = argparse.ArgumentParser(
386 prog="cosmoviz",
387 formatter_class=argparse.RawTextHelpFormatter,
388 description="Visualize COSMO files",
389 )
390 parser.add_argument(
391 "cosmo_file",
392 type=pathlib.Path,
393 help="Path to a COSMO quantum mechanical output file",
394 )
395 parser.add_argument(
396 "--show-original-charge-densities",
397 action="store_true",
398 help="Show original charge densities instead of smoothed ones",
399 )
400 parser.add_argument(
401 "--use-continuous-colors",
402 action="store_true",
403 help="Use continuous colors instead of uniformly colored segments",
404 )
405 parser.add_argument(
406 "--segment-edge-color",
407 type=str,
408 default=None,
409 help="Color of the edges between segments (default: None)",
410 )
411 parser.add_argument(
412 "--colormap",
413 type=str,
414 default="jet",
415 help="Matplotlib colormap name (default: jet)",
416 )
417 return parser
420def main() -> None:
421 args = get_parser().parse_args()
422 component = Component(args.cosmo_file.read_text())
423 geometries = generate_geometries(
424 component,
425 args.show_original_charge_densities,
426 args.use_continuous_colors,
427 args.colormap,
428 args.segment_edge_color,
429 )
430 o3d.visualization.draw_geometries( # ty: ignore[possibly-missing-submodule]
431 geometries,
432 mesh_show_back_face=True,
433 window_name=f"Surface Segments from {args.cosmo_file.name}",
434 )
437if __name__ == "__main__":
438 main()