def surfaces_from_nodes(nodes, color, place, instances, session):
from collada.scene import GeometryNode, Node
from chimerax.geometry import Place
from chimerax.core.models import Surface
splist = []
for n in nodes:
if isinstance(n, GeometryNode):
materials = dict((m.symbol, m.target) for m in n.materials)
g = n.geometry
colors = g.sourceById
if g.id in instances:
add_geometry_instance(g.primitives, instances[g.id], place,
color, materials)
else:
instances[g.id] = spieces = geometry_node_surfaces(
g.primitives, place, color, materials, colors, session)
splist.extend(spieces)
elif isinstance(n, Node):
pl = place * Place(n.matrix[:3, :])
spieces = surfaces_from_nodes(n.children, color, pl, instances,
session)
name = n.xmlnode.get('name')
if len(spieces) > 1:
m = Surface(name, session)
m.add(spieces)
splist.append(m)
elif len(spieces) == 1:
s = spieces[0]
s.name = name
splist.append(s)
return splist
def vbur_vis(
session,
geom,
targets,
radii,
scale,
radius,
center,
point_spacing,
intersection_scale,
volume_type,
vbur,
labels,
basis=None,
key_atoms=None,
):
# from cProfile import Profile
#
# profile = Profile()
# profile.enable()
# number of points is based on surface area
n_grid = int(4 * np.pi * radius**2 / point_spacing)
if volume_type == "buried":
model = Surface("%%Vbur for %s" % geom.name, session)
else:
model = Surface("%%Vfree for %s" % geom.name, session)
# verts, norms, and triangles for the drawing
vertices = []
normals = []
triangles = []
if isinstance(center, np.ndarray):
center_coords = center
else:
center_coords = geom.COM(center)
if isinstance(radii, dict):
radii_dict = radii
elif radii.lower() == "umn":
radii_dict = VDW_RADII
elif radii.lower() == "bondi":
radii_dict = BONDI_RADII
else:
raise RuntimeError("received %s for radii, must be umn or bondi" %
radii)
if not hasattr(center, "__iter__"):
center = [center]
if targets is None:
if len(center) == 1:
atoms = [atom for atom in geom.atoms if atom not in center]
else:
atoms = geom.atoms
else:
atoms = geom.find(targets)
atoms_within_radius = []
for atom in atoms:
# determine which atom's radii extend within the sphere
# reduces the number of distances we need to calculate
d = np.linalg.norm(center_coords - atom.coords)
inner_edge = d - scale * radii_dict[atom.element]
if inner_edge < radius:
atoms_within_radius.append(atom)
# sort atoms based on their distance to the center
# this makes is so we usually break out of looping over the atoms faster
atoms_within_radius.sort(
key=lambda a, c=center_coords: np.linalg.norm(a.coords - c))
radius_list = []
for atom in atoms_within_radius:
radius_list.append(scale * radii_dict[atom.element])
coords = geom.coordinates(atoms_within_radius)
atom_dist = distance_matrix(coords, coords)
sphere = fibonacci_sphere(num=n_grid, radius=radius)
# add points right where spheres intersect
# this makes the intersections look less pokey
d_ac = distance_matrix(coords, [center_coords])
center_added_points = []
atom_added_points = [[] for atom in atoms_within_radius]
for i in range(0, len(coords)):
r1 = radius_list[i]
v_n = coords[i] - center_coords
v_n /= np.linalg.norm(v_n)
d = d_ac[i, 0]
# intersection with big sphere
if d + r1 > radius:
theta = np.arccos((r1**2 - radius**2 - d**2) / (-2 * d * radius))
h = radius * np.sin(theta)
b = radius * np.cos(theta)
p = b * v_n
circ = 2 * np.pi * h
n_added = int(np.ceil(intersection_scale * circ / point_spacing))
r_t = perp_vector(v_n)
r0 = np.cross(r_t, v_n)
r0 *= h / np.linalg.norm(r0)
rot_angle = 2 * np.pi / n_added
R = rotation_matrix(rot_angle, v_n)
prev_r = r0
prev_r_list = []
for x in np.linspace(0, 2 * np.pi, num=n_added):
prev_r = np.dot(R, prev_r)
prev_r_list.append(prev_r)
prev_r_list = np.array(prev_r_list)
d_ep = distance_matrix(prev_r_list + p + center_coords, coords)
# if the point is obscured by another atom, don't bother adding it
# this saves time on triangulation
diff_mat = d_ep - radius_list
diff_mat[:, i] = 1
mask = np.invert(np.any(diff_mat < 0, axis=1))
center_added_points.extend(prev_r_list[mask] + p)
atom_added_points[i].extend(prev_r_list[mask] + p + center_coords)
for j in range(0, i):
d = atom_dist[i, j]
r2 = radius_list[j]
if d < r1 + r2 and d > abs(r1 - r2):
v_n = coords[j] - coords[i]
v_n /= np.linalg.norm(v_n)
theta = np.arccos((r2**2 - r1**2 - d**2) / (-2 * r1 * d))
h = r1 * np.sin(theta)
b = r1 * np.cos(theta)
p = b * v_n
circ = 2 * np.pi * h
n_added = int(
np.ceil(intersection_scale * circ / point_spacing))
r_t = perp_vector(v_n)
r0 = np.cross(r_t, v_n)
r0 *= h / np.linalg.norm(r0)
rot_angle = 2 * np.pi / n_added
R = rotation_matrix(rot_angle, v_n)
prev_r = r0
prev_r_list = []
for x in np.linspace(0, 2 * np.pi, num=n_added):
prev_r = np.dot(R, prev_r)
prev_r_list.append(prev_r)
prev_r_list = np.array(prev_r_list)
norm_mask = np.linalg.norm(center_coords -
(prev_r_list + p + coords[i]),
axis=1) < radius
prev_r_list = prev_r_list[norm_mask]
if not len(prev_r_list):
continue
d_ep = distance_matrix(prev_r_list + p + coords[i], coords)
diff_mat = d_ep - radius_list
diff_mat[:, i] = 1
diff_mat[:, j] = 1
mask = np.invert(np.any(diff_mat < 0, axis=1))
if any(mask):
atom_added_points[i].extend(prev_r_list[mask] + p +
coords[i])
atom_added_points[j].extend(prev_r_list[mask] + p +
coords[i])
tol = 0.3 * point_spacing
for i in range(0, len(coords)):
# get a grid of points around each atom
# remove any points that are close to an intersection
# this makes it less likely for the triangulation to choose
# one of these points instead of one that we already added
# then, if we have to remove a triangle involving one of these points later,
# it won't leave a gap
n_atom_grid = int(radius_list[i]**2 * n_grid / radius**2)
atom_sphere = fibonacci_sphere(radius=radius_list[i], num=n_atom_grid)
mask = np.ones(len(atom_sphere), dtype=bool)
n_atom_grid = len(atom_sphere)
if len(atom_added_points[i]) > 0:
remove_ndx = []
dist_mat = distance_matrix(
atom_sphere,
np.array(atom_added_points[i]) - coords[i])
mask *= np.any((dist_mat - tol) < 0)
atom_sphere = atom_sphere[mask]
atom_sphere = np.array(atom_sphere)
n_atom_grid = len(atom_sphere)
atom_sphere = np.concatenate(
(atom_sphere, np.array(atom_added_points[i]) - coords[i]))
if len(atom_sphere) < 4:
continue
atom_hull = ConvexHull(atom_sphere / radius_list[i])
tri = atom_hull.simplices
atom_sphere += coords[i]
remove_v = []
new_ndx = np.zeros(len(atom_sphere), dtype=int)
# remove any points that are covered by another atom
# only loop over n_atom_grid points so we don't remove
# any points right on the intersection b/c of
# numerical issues
del_count = 0
atom_grid_dist = distance_matrix(atom_sphere, coords)
center_atom_grid_dist = distance_matrix(atom_sphere,
[center_coords])[:, 0]
for j in range(0, n_atom_grid):
if center_atom_grid_dist[j] > radius:
remove_v.append(j)
del_count += 1
continue
for k in range(0, len(coords)):
if i == k:
continue
if atom_grid_dist[j, k] < radius_list[k]:
remove_v.append(j)
del_count += 1
break
new_ndx[j] = j - del_count
for j in range(n_atom_grid, len(atom_sphere)):
new_ndx[j] = j - del_count
if len(remove_v) == len(atom_sphere):
continue
atom_sphere = atom_sphere.tolist()
for vi in remove_v[::-1]:
atom_sphere.pop(vi)
tri = tri[np.all(tri != vi, axis=1)]
for j, ti in enumerate(tri):
for k, v in enumerate(ti):
tri[j][k] = new_ndx[v]
atom_sphere = np.array(atom_sphere)
norms = -(atom_sphere - coords[i]) / radius_list[i]
triangles.extend(tri + len(vertices))
vertices.extend(atom_sphere)
normals.extend(norms)
remove_ndx = []
if len(center_added_points) > 0:
dist_mat = distance_matrix(sphere, center_added_points)
for i in range(0, len(sphere)):
for j in range(0, len(center_added_points)):
if dist_mat[i, j] < tol:
remove_ndx.append(i)
break
sphere = sphere.tolist()
for ndx in remove_ndx[::-1]:
sphere.pop(ndx)
sphere = np.array(sphere)
n_sphere = len(sphere)
if center_added_points:
sphere = np.concatenate((sphere, np.array(center_added_points)))
center_hull = ConvexHull(sphere / radius)
sphere += center_coords
tri = center_hull.simplices
remove_v = []
new_ndx = np.zeros(len(sphere), dtype=int)
del_count = 0
center_grid_dist = distance_matrix(sphere, coords)
for i in range(0, n_sphere):
if volume_type == "free":
for j in range(0, len(coords)):
if center_grid_dist[i, j] < radius_list[j]:
remove_v.append(i)
del_count += 1
break
elif volume_type == "buried":
if all(center_grid_dist[i, j] > radius_list[j]
for j in range(0, len(coords))):
remove_v.append(i)
del_count += 1
new_ndx[i] = i - del_count
for i in range(n_sphere, len(sphere)):
new_ndx[i] = i - del_count
sphere = sphere.tolist()
mask = np.ones(len(tri), dtype=bool)
for vi in remove_v[::-1]:
sphere.pop(vi)
mask *= np.invert(np.any(tri == vi, axis=1))
tri = tri[mask]
new_t = tri
if volume_type == "buried":
mask = np.invert(np.all(new_t >= n_sphere, axis=1))
new_t = new_t[mask]
new_t = np.array(new_t)
for i, ti in enumerate(new_t):
new_t[i] = new_ndx[ti]
norms = []
if sphere:
sphere = np.array(sphere)
norms = (sphere - center_coords) / radius
triangles.extend(new_t + len(vertices))
vertices.extend(sphere)
normals.extend(norms)
# the triangles need to be reordered so the points are
# clockwise (or counterclockwise? i don't remember)
vertices = np.array(vertices)
normals = np.array(normals)
triangles = np.array(triangles)
v1 = vertices[triangles[:, 1]] - vertices[triangles[:, 0]]
v2 = vertices[triangles[:, 2]] - vertices[triangles[:, 0]]
c = np.cross(v1, v2)
mask = np.sum(c * normals[triangles[:, 0]], axis=1) < 0
triangles[mask] = triangles[mask, ::-1]
model.set_geometry(vertices, normals, triangles)
# profile.disable()
# from TestManager.stream_holder import StreamHolder
# import pstats
# stream = StreamHolder(session)
# pstats.Stats(profile, stream=stream).strip_dirs().sort_stats(-1).print_stats()
# stream.flush()
if labels != "none":
d = model.new_drawing("octants")
d.display_style = d.Mesh
d.use_lighting = False
d.casts_shadows = False
d.pickable = False
d_theta = np.pi / 180
oct_radius = 1.01 * radius
verts = [np.zeros(3)]
triangles = []
for k, v in enumerate(basis.T):
if labels == "quadrants" and k == 2:
continue
start_ndx = len(verts)
prev_vert = oct_radius * perp_vector(v)
verts.append(prev_vert)
R = rotation_matrix(d_theta, v)
for i in range(1, 360):
triangles.append([0, start_ndx + i - 1, start_ndx + i])
prev_vert = np.dot(R, prev_vert)
verts.append(prev_vert)
triangles.append([0, start_ndx, start_ndx + i - 1])
for v2 in basis.T[:k]:
v3 = np.cross(v, v2)
v3 /= np.linalg.norm(v3)
v3 *= oct_radius
ndx = len(verts)
verts.append(v)
verts.append(v3)
triangles.append([ndx, 0, ndx + 1])
ndx = len(verts)
verts.append(v)
verts.append(-v3)
triangles.append([ndx, 0, ndx + 1])
ndx = len(verts)
verts.append(-v)
verts.append(-v3)
triangles.append([ndx, 0, ndx + 1])
ndx = len(verts)
verts.append(-v)
verts.append(v3)
triangles.append([ndx, 0, ndx + 1])
norms = np.array(verts) / oct_radius
verts = np.array(verts) + center_coords
triangles = np.array(triangles)
d.set_geometry(verts, norms, triangles)
d.set_edge_mask(2 * np.ones(len(triangles), dtype=int))
label_list = []
x = basis.T[0]
x = x / np.linalg.norm(x)
y = basis.T[1]
y = y / np.linalg.norm(y)
z = basis.T[2]
z = z / np.linalg.norm(z)
if labels == "octants":
for i, val in enumerate(vbur):
coordinates = np.zeros(3)
if any(i == n for n in [1, 2, 5, 6]):
coordinates -= x
else:
coordinates += x
if any(i == n for n in [2, 3, 4, 5]):
coordinates -= y
else:
coordinates += y
if i > 3:
coordinates -= z
else:
coordinates += z
if volume_type == "free":
l = "%.1f%%" % (100 / 8 - val)
else:
l = "%.1f%%" % (val)
coordinates /= np.linalg.norm(coordinates)
coordinates *= radius
coordinates += center_coords
label_list.append(VoidLabel(l, coordinates, session.main_view))
else:
quad_vbur = [
vbur[0] + vbur[7],
vbur[1] + vbur[6],
vbur[2] + vbur[5],
vbur[3] + vbur[4],
]
for i, val in enumerate(quad_vbur):
coordinates = np.zeros(3)
if any(i == n for n in [1, 2]):
coordinates -= x
else:
coordinates += x
if any(i == n for n in [2, 3]):
coordinates -= y
else:
coordinates += y
if volume_type == "free":
l = "%.1f%%" % (25 - val)
else:
l = "%.1f%%" % (val)
coordinates /= np.linalg.norm(coordinates)
coordinates *= radius
coordinates += center_coords
label_list.append(VoidLabel(l, coordinates, session.main_view))
label_object = VoidLabels(session)
label_object.add_labels(label_list)
model.add([label_object])
return [model]
return [model]
