|
# Thermal ellipsoids in ORTEP-like style |
|
# by Takanori Nakane, 2025 |
|
|
|
from math import sin, cos, pi |
|
from pymol.cgo import * |
|
|
|
import numpy as np |
|
|
|
def normalize(v): |
|
return v / np.linalg.norm(v) |
|
|
|
def cgo_ellipse_line(center, ax1, ax2, ndiv=32): |
|
obj = [] |
|
step = 2 * pi / ndiv |
|
for i in range(ndiv): |
|
p1 = center + ax1 * cos(step * i) + ax2 * sin(step * i) |
|
p2 = center + ax1 * cos(step * (i + 1)) + ax2 * sin(step * (i + 1)) |
|
obj.append(VERTEX) |
|
obj.extend(p1) |
|
obj.append(VERTEX) |
|
obj.extend(p2) |
|
|
|
return obj |
|
|
|
def cgo_ellipse_surface(center, ax1, ax2, ax3, cut_octant=True, ndiv=32): |
|
assert (not cut_octant) or (ndiv % 4 == 0) # necessary to cut out an octant exactly |
|
half_ndiv = ndiv // 2 |
|
|
|
obj = [] |
|
scale = np.array([np.linalg.norm(ax1), np.linalg.norm(ax2), np.linalg.norm(ax3)]) |
|
camera_matrix = np.array(cmd.get_view()[0:9]).reshape((3, 3)) |
|
camera_direction = camera_matrix.dot(np.array([0, 0, 1])) |
|
|
|
ax1_sign = np.sign(np.dot(ax1, camera_direction)) |
|
ax2_sign = np.sign(np.dot(ax2, camera_direction)) |
|
ax3_sign = np.sign(np.dot(ax3, camera_direction)) |
|
|
|
# we must use the same tessellation as cgo_ellipse_disc |
|
step = 2 * pi / ndiv |
|
for i in range(ndiv // 2): |
|
p0 = cos(step * i) * ax1 |
|
q0 = cos(step * (i + 1)) * ax1 |
|
|
|
f = sin(step * i) |
|
g = sin(step * (i + 1)) |
|
|
|
q0_sign = np.sign(np.dot(q0, camera_direction)) |
|
|
|
for j in range(ndiv): |
|
if cut_octant and q0_sign > 0: |
|
if ax2_sign >= 0 and ax3_sign >= 0 and j >= 0 and j < ndiv // 4: continue |
|
if ax2_sign < 0 and ax3_sign >= 0 and j >= ndiv // 4 and j < ndiv // 2: continue |
|
if ax2_sign >= 0 and ax3_sign < 0 and j >= 3 * ndiv // 4 and j < ndiv: continue |
|
if ax2_sign < 0 and ax3_sign < 0 and j >= ndiv // 2 and j < 3 * ndiv // 4: continue |
|
|
|
p1 = ax2 * cos(2 * pi * j / ndiv) + ax3 * sin(2 * pi * j / ndiv) |
|
p2 = ax2 * cos(2 * pi * (j + 1) / ndiv) + ax3 * sin(2 * pi * (j + 1) / ndiv) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normalize(scale * (p0 + f * p1))) |
|
obj.append(VERTEX) |
|
obj.extend(center + p0 + f * p1) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normalize(scale * (q0 + g * p1))) |
|
obj.append(VERTEX) |
|
obj.extend(center + q0 + g * p1) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normalize(scale * (q0 + g * p2))) |
|
obj.append(VERTEX) |
|
obj.extend(center + q0 + g * p2) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normalize(scale * (p0 + f * p1))) |
|
obj.append(VERTEX) |
|
obj.extend(center + p0 + f * p1) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normalize(scale * (q0 + g * p2))) |
|
obj.append(VERTEX) |
|
obj.extend(center + q0 + g * p2) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normalize(scale * (p0 + f * p2))) |
|
obj.append(VERTEX) |
|
obj.extend(center + p0 + f * p2) |
|
|
|
return obj |
|
|
|
def cgo_ellipse_disc(center, ax1, ax2, ndiv=32): |
|
delta=0.0001 # displacement to avoid Z-fighting |
|
normal = np.cross(normalize(ax1), normalize(ax2)) |
|
|
|
obj = [] |
|
step = 2 * pi / ndiv |
|
for i in range(ndiv): |
|
p1 = center + ax1 * cos(step * i) + ax2 * sin(step * i) |
|
p2 = center + ax1 * cos(step * (i + 1)) + ax2 * sin(step * (i + 1)) |
|
|
|
# one face |
|
obj.append(NORMAL) |
|
obj.extend(normal) |
|
obj.append(VERTEX) |
|
obj.extend(delta * normal + p1) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normal) |
|
obj.append(VERTEX) |
|
obj.extend(delta * normal + p2) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(normal) |
|
obj.append(VERTEX) |
|
obj.extend(delta * normal + center) |
|
|
|
# the other side |
|
obj.append(NORMAL) |
|
obj.extend(-normal) |
|
obj.append(VERTEX) |
|
obj.extend(- delta * normal + p2) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(-normal) |
|
obj.append(VERTEX) |
|
obj.extend(- delta * normal + p1) |
|
|
|
obj.append(NORMAL) |
|
obj.extend(-normal) |
|
obj.append(VERTEX) |
|
obj.extend(- delta * normal + center) |
|
|
|
return obj |
|
|
|
# Probability to Radius |
|
# |
|
# THE PROBLEM: |
|
# Find the radius `r` of a sphere such that |
|
# the integrated probability of a 3D standard Gaussian within it is `p`. |
|
# |
|
# SOLUTION: |
|
# We work backwards. We can think of a 3D standard Gaussian (i.e. no covariance) |
|
# as a distribution of THREE 1D Gaussian random variables `x`, `y` and `z`. |
|
# The sum of squares `r^2 = x^2 + y^2 + z^2` obeys the chi-squared distribution |
|
# whose degree of freedom is 3. Therefore, the integrated probability within |
|
# the sphere is the same as the cumulative probability of this chi-squared |
|
# distribution up to `r^2`. |
|
# The cumulative distribution function CDF is `p = P(3 / 2, r * r / 2)`, |
|
# where P is the regularized lower incomplete Gamma function. |
|
# This can be calculated by `scipy.special.gammainc`. |
|
# To get `r` from `p`, we numerically solve this equation. |
|
# |
|
# CODE: |
|
# [float(np.round(fsolve(lambda x:scipy.special.gammainc(1.5, x*x/2)-p, 1.0)[0], 5)) |
|
# for p in np.linspace(0.01, 1.0, 100)] |
|
# |
|
# The result matches the table in PyMOL's layer2/RepEllipsoid.cpp, |
|
# which samples p every 0.02. |
|
# The same constants are also found in the ORTEP source code. |
|
# https://digital.library.unt.edu/ark:/67531/metadc678816/m1/145/ |
|
|
|
# for p = 0.01, 0.02, ..., 1.00 |
|
SPHERE_RADII = [0.33887, 0.42992, 0.49507, 0.54786, 0.59317, 0.63338, 0.66988, |
|
0.70353, 0.73491, 0.76444, 0.79245, 0.81915, 0.84475, 0.86938, 0.89318, 0.91624, |
|
0.93864, 0.96046, 0.98175, 1.00258, 1.02299, 1.04302, 1.06270, 1.08207, 1.10115, |
|
1.11998, 1.13857, 1.15696, 1.17515, 1.19317, 1.21103, 1.22876, 1.24636, 1.26385, |
|
1.28124, 1.29855, 1.31578, 1.33296, 1.35009, 1.36718, 1.38424, 1.40128, 1.41831, |
|
1.43535, 1.45240, 1.46947, 1.48657, 1.50372, 1.52092, 1.53817, 1.55550, 1.57291, |
|
1.59041, 1.60801, 1.62573, 1.64357, 1.66155, 1.67968, 1.69797, 1.71644, 1.73510, |
|
1.75396, 1.77304, 1.79236, 1.81194, 1.83179, 1.85193, 1.87240, 1.89321, 1.91439, |
|
1.93596, 1.95796, 1.98043, 2.00340, 2.02691, 2.05100, 2.07575, 2.10119, 2.12739, |
|
2.15444, 2.18242, 2.21143, 2.24158, 2.27300, 2.30587, 2.34037, 2.37674, 2.41526, |
|
2.45628, 2.50028, 2.54783, 2.59975, 2.65713, 2.72156, 2.79548, 2.88291, 2.99120, |
|
3.13646, 3.36821, 10.42106] |
|
|
|
def thermal_ellipsoid(selection, cgo_name="thermal_ellipsoid", draw_rings=False): |
|
''' |
|
DESCRIPTION |
|
|
|
Draw thermal ellipsoids in a style similar to ORTEP. |
|
|
|
USAGE |
|
|
|
thermal_ellipsoid selection, [cgo_name, [draw_rings]] |
|
|
|
ARGUMENTS |
|
|
|
selection = string: selection of the atoms to draw thermal ellipsoids |
|
|
|
cgo_name = string: name of the resulting CGO object. {defalt: thermal_ellipsoid} |
|
When the object is already present, it will be cleared. |
|
|
|
draw_rings = bool: draw rings instead of solid ellipsoids {default: False} |
|
|
|
SETTINGS |
|
|
|
The following global settings are respected: |
|
|
|
- ellipsoid_quality |
|
|
|
- ellipsoid_probability |
|
|
|
LIMITATIONS |
|
|
|
- Object matrix (TTT matrix) is ignored. |
|
|
|
- You have to rerun the command to update the octant positions to match the camera. |
|
|
|
- Isotropic atoms are ignored. |
|
|
|
EXAMPLES |
|
|
|
- Rings over native ellipsoids |
|
|
|
as ellipsoids |
|
show lines |
|
thermal_ellipsoid all, draw_rings=True |
|
color black, thermal_ellipsoid |
|
set cgo_line_radius, 0.005 |
|
bg_color white |
|
set opaque_background, 1 |
|
ray |
|
|
|
- Classic ORTEP-like style (black and white) |
|
|
|
as lines |
|
thermal_ellipsoid all |
|
set antialias, 2 |
|
set ray_trace_mode, 2 |
|
bg_color white |
|
set opaque_background, 1 |
|
set antialias, 2 |
|
ray |
|
''' |
|
cgo_obj = [] |
|
div = 32 * int(cmd.get("ellipsoid_quality")) |
|
|
|
# In the latest PyMOL, we can use "iterate_to_list" in pymol.editor |
|
stored._thermal_ellipsoid_color_indices = [] |
|
stored._thermal_ellipsoid_probabilities = [] |
|
cmd.iterate(selection, "stored._thermal_ellipsoid_color_indices.append(color)") |
|
cmd.iterate(selection, "stored._thermal_ellipsoid_probabilities.append(s.ellipsoid_probability)") |
|
color_indices = stored._thermal_ellipsoid_color_indices |
|
probabilities = stored._thermal_ellipsoid_probabilities |
|
del stored._thermal_ellipsoid_color_indices |
|
del stored._thermal_ellipsoid_probabilities |
|
for atom, color_idx, prob in zip(cmd.get_model(selection).atom, color_indices, probabilities): |
|
if not hasattr(atom, "u_aniso"): |
|
continue # skip isotropic atoms |
|
|
|
# TODO: handle object's TTT matrix |
|
center = np.array(atom.coord) |
|
|
|
iprob = np.clip(int((prob + 0.01) * 100) - 1, 0, 99) |
|
pradius = SPHERE_RADII[iprob] |
|
matrix = np.array([[atom.u_aniso[0], atom.u_aniso[3], atom.u_aniso[4]], |
|
[atom.u_aniso[3], atom.u_aniso[1], atom.u_aniso[5]], |
|
[atom.u_aniso[4], atom.u_aniso[5], atom.u_aniso[2]]]) |
|
eigenval, eigenvec = np.linalg.eig(matrix) |
|
|
|
if any(eigenval <= 0): # non positive definite |
|
continue |
|
eivenval_sqrt = np.sqrt(eigenval) |
|
|
|
if np.cross(eigenvec[:, 0], eigenvec[:, 1]).dot(eigenvec[:, 2]) < 0: |
|
eigenvec[:, 2] *= -1 # make sure axes are right-handed |
|
|
|
color = cmd.get_color_tuple(color_idx) |
|
if not draw_rings: |
|
cgo_obj.append(COLOR) |
|
cgo_obj.extend(color) |
|
cgo_obj += cgo_ellipsoid(center, eigenvec[:, 0] * eivenval_sqrt[0] * pradius, |
|
eigenvec[:, 1] * eivenval_sqrt[1] * pradius, |
|
eigenvec[:, 2] * eivenval_sqrt[2] * pradius, |
|
ndiv=div) |
|
else: |
|
cgo_obj += cgo_ellipsoid_rings(center, eigenvec[:, 0] * eivenval_sqrt[0] * pradius, |
|
eigenvec[:, 1] * eivenval_sqrt[1] * pradius, |
|
eigenvec[:, 2] * eivenval_sqrt[2] * pradius, |
|
ndiv=div) |
|
|
|
cmd.delete(cgo_name) |
|
cmd.load_cgo(cgo_obj, cgo_name) |
|
|
|
def cgo_ellipsoid(center, ax1, ax2, ax3, ndiv=32, cut_octant=True): |
|
cgo_obj = [BEGIN, TRIANGLES] |
|
|
|
ax1n = ax1 / np.linalg.norm(ax1) |
|
ax2n = ax2 / np.linalg.norm(ax2) |
|
ax3n = ax3 / np.linalg.norm(ax3) |
|
|
|
if cut_octant: |
|
cgo_obj += cgo_ellipse_disc(center, ax1, ax2, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_disc(center, ax2, ax3, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_disc(center, ax3, ax1, ndiv=ndiv) |
|
|
|
cgo_obj += cgo_ellipse_surface(center, ax1, ax2, ax3, ndiv=ndiv) |
|
|
|
cgo_obj.append(END) |
|
return cgo_obj |
|
|
|
def cgo_ellipsoid_rings(center, ax1, ax2, ax3, ndiv=32, draw_center_cross=False): |
|
cgo_obj = [BEGIN, LINES] |
|
|
|
if draw_center_cross: |
|
cgo_obj.append(VERTEX) |
|
cgo_obj.extend(center + ax1) |
|
cgo_obj.append(VERTEX) |
|
cgo_obj.extend(center - ax1) |
|
cgo_obj.append(VERTEX) |
|
cgo_obj.extend(center + ax2) |
|
cgo_obj.append(VERTEX) |
|
cgo_obj.extend(center - ax2) |
|
cgo_obj.append(VERTEX) |
|
cgo_obj.extend(center + ax3) |
|
cgo_obj.append(VERTEX) |
|
cgo_obj.extend(center - ax3) |
|
|
|
for i in range(1): |
|
f = i / (ndiv + 1) |
|
rest = np.sqrt(1.0 - f * f) |
|
cgo_obj += cgo_ellipse_line(center + f * ax1, rest * ax2, rest * ax3, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_line(center - f * ax1, rest * ax2, rest * ax3, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_line(center + f * ax2, rest * ax3, rest * ax1, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_line(center - f * ax2, rest * ax3, rest * ax1, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_line(center + f * ax3, rest * ax1, rest * ax2, ndiv=ndiv) |
|
cgo_obj += cgo_ellipse_line(center - f * ax3, rest * ax1, rest * ax2, ndiv=ndiv) |
|
|
|
cgo_obj.append(END) |
|
return cgo_obj |
|
|
|
cmd.extend("thermal_ellipsoid", thermal_ellipsoid) |
