import numpy as np
import plotly.graph_objects as go
import atom_properties as atomp
from itertools import accumulate
import qc_parser as qcp
import subprocess as sub
import os
class XYZ:
sphere_mode = "ball"
def __init__(self, coords: list[dict[str, str|float]]):
self.xyz = np.array([[coord["x"], coord["y"], coord["z"]] for coord in coords])
self.elements = [coord["element"] for coord in coords]
@classmethod
def from_file(cls, file: str):
"""
Reads in files with the structure 'element x y z' and returns
a numpy array with the stored coordinates.
@param file: String. Path to the file to be read in.
@return: 2D numpy array. Contains the coordinates from the file at the
corresponding indices.
"""
with open(file) as f:
lines = f.readlines()
if len(lines[0].split()) == 1:
lines = lines[2:]
return cls([{"element": line.split()[0],
"x": float(line.split()[1]),
"y": float(line.split()[2]),
"z": float(line.split()[3])} for line in lines])
@classmethod
def from_list(cls, lines: list[list[str|float]]):
return cls([{"element": line[0],
"x": float(line[1]),
"y": float(line[2]),
"z": float(line[3])} for line in lines])
def get_bond_length(self, i: int, j: int) -> float:
"""
Returns the bond length between atoms i and j.
Args:
i (int): Index of the first atom.
j (int): Index of the second atom.
Returns:
float: The bond length between the two atoms.
"""
return float(np.linalg.norm(self.xyz[i] - self.xyz[j]))
def circle_coordinates(self, point: np.ndarray, v: np.ndarray, radius: float, num_points: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Generate cartesian coordinates of a circle in 3D space.
Args:
point (np.ndarray): A point on the circle.
v (np.ndarray): A vector perpendicular to the circle.
radius (float): The radius of the circle.
num_points (int): The number of points to generate on the circle.
Returns:
np.ndarray: A 2D array of shape (num_points, 3) containing the cartesian coordinates of the circle.
"""
# Normalize the vector v
v = v / np.linalg.norm(v)
# Create a vector u that is perpendicular to v
u = np.cross(v, [1, 0, 0]) if np.linalg.norm(np.cross(v, [1, 0, 0])) > 1e-6 else np.cross(v, [0, 1, 0])
u = u / np.linalg.norm(u)
# Create a vector w that is perpendicular to both u and v
w = np.cross(u, v)
# Generate the cartesian coordinates of the circle
theta = np.linspace(0, 2 * np.pi, num_points)
circle_coords = point + radius * (u * np.cos(theta)[:, np.newaxis] + w * np.sin(theta)[:, np.newaxis])
x, y, z = circle_coords[:, 0], circle_coords[:, 1], circle_coords[:, 2]
return x, y, z
def make_fibonacci_sphere(self, center: np.ndarray, radius: float = 0.1, resolution: int = 32) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Return cartesian coordinates of points evenly distributed on the surface of a sphere.
Args:
center (np.ndarray): The coordinates of the center of the sphere.
radius (float, optional): The radius of the sphere. Defaults to 0.1.
resolution (int, optional): The number of points to be generated. Defaults to 32.
Returns:
tuple[np.ndarray, np.ndarray, np.ndarray]: The cartesian coordinates of the points
on the surface of the sphere. The three arrays are the x, y and z coordinates
of the points.
"""
num_points = resolution
indices = np.arange(0, num_points, dtype=float) + 0.5
phi = np.arccos(1 - 2*indices/num_points)
theta = np.pi * (1 + 5**0.5) * indices
x = radius * np.sin(phi) * np.cos(theta) + center[0]
y = radius * np.sin(phi) * np.sin(theta) + center[1]
z = radius * np.cos(phi) + center[2]
return x, y, z
def make_cylinder(self, center_i: int, center_j: int, resolution: int = 32, radius: float = 0.1) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Return cartesian coordinates of points on the edges of a cylinder.
The cylinder is defined by the line between atoms i and j, and the radius is the given value. The number of points is determined by the resolution.
Args:
center_i (int): Index of the first atom that defines the cylinder.
center_j (int): Index of the second atom that defines the cylinder.
resolution (int, optional): The number of points to be generated. Defaults to 32.
radius (float, optional): The radius of the cylinder. Defaults to 0.1.
Returns:
tuple[np.ndarray, np.ndarray, np.ndarray]: The cartesian coordinates of the points on the edges of the cylinder. The three arrays are the x, y and z coordinates of the points.
"""
v = self.xyz[center_j] - self.xyz[center_i]
x1, y1, z1 = self.circle_coordinates(self.xyz[center_i], v, radius=radius, num_points=resolution//4)
x2, y2, z2 = self.circle_coordinates(self.xyz[center_j], v, radius=radius, num_points=resolution//4)
return np.concatenate((x1, x2)), np.concatenate((y1, y2)), np.concatenate((z1, z2))
def make_atom_mesh(self, i: int, resolution: int = 32) -> go.Mesh3d:
"""
Return a Mesh3d object that represents the atom at index i.
The atom is represented as a sphere with a radius depending on its van der Waals radius.
Args:
i (int): Index of the atom to be represented.
resolution (int, optional): The number of points to be generated. Defaults to 32.
Returns:
go.Mesh3d: A Mesh3d object that represents the atom.
"""
if self.sphere_mode == "vdw":
radius = atomp.vdw_radii_dict[self.elements[i]]
elif self.sphere_mode == "ball":
radius = atomp.vdw_radii_dict[self.elements[i]] * 0.2
x, y, z = self.make_fibonacci_sphere(self.xyz[i], radius=radius, resolution=resolution)
atom_mesh = go.Mesh3d(
x=x,
y=y,
z=z,
color=atomp.atom_colors_dict[self.elements[i]],
opacity=1,
alphahull=0,
name=f'{self.elements[i]}{i}', # label is too short to also show coordinates
hoverinfo='name', # Only show the name on hover
)
return atom_mesh
def make_bond_mesh(self, i: int, j: int, resolution: int = 32, radius: float = 0.1) -> go.Mesh3d:
"""
Return a Mesh3d object that represents a bond between atoms i and j.
The bond is represented as a cylinder with a radius depending on the given radius.
Args:
i (int): Index of the first atom in the bond.
j (int): Index of the second atom in the bond.
resolution (int, optional): The number of points to be generated. Defaults to 32.
radius (float, optional): The radius of the cylinder. Defaults to 0.1.
Returns:
go.Mesh3d: A Mesh3d object that represents the bond.
"""
x, y, z = self.make_cylinder(i, j, resolution=resolution, radius=radius)
bond_mesh = go.Mesh3d(
x=x,
y=y,
z=z,
color="#444444",
opacity=1,
alphahull=0,
# name=f'{self.elements[i]}-{self.elements[j]}',
hoverinfo='none', # No hover info at all
)
return bond_mesh
def get_molecular_mesh(self, resolution: int = 64, rel_cutoff: float = 0.5) -> go.Figure:
"""
Returns a list of Mesh3d objects that represent the molecular structure.
Args:
resolution (int, optional): The number of points to be generated. Defaults to 64.
rel_cutoff (float, optional): Bond are only shown if the interatomic distance is smaller than the sum of the atomic vdw radii multiplied with this parameter.
Defaults to 0.5.
Returns:
go.Figure: A list of Mesh3d objects that represent the molecular structure.
"""
mesh_list = []
# Add the bonds
for i in range(len(self.elements)):
for j in range(i+1, len(self.elements)):
bond_length = self.get_bond_length(i, j)
max_bond_length = rel_cutoff * (atomp.vdw_radii_dict[self.elements[i]] + atomp.vdw_radii_dict[self.elements[j]])
if bond_length < max_bond_length:
mesh_list.append(self.make_bond_mesh(i, j, resolution=resolution, radius=0.1))
# Add the atoms
for i in range(len(self.elements)):
mesh_list.append(self.make_atom_mesh(i, resolution=resolution))
return mesh_list
def get_molecular_viewer(self, resolution: int = 64, rel_cutoff: float = 0.5) -> go.Figure:
"""
Returns a plotly figure with the molecular structure.
The molecular structure is represented as a collection of spheres at the atomic positions,
with the bonds represented as cylinders between the atoms.
Args:
resolution (int, optional): The number of points to be generated. Defaults to 64.
rel_cutoff (float, optional): Bond are only shown if the interatomic distance is smaller than the sum of the atomic vdw radii multiplied with this parameter.
Defaults to 0.5.
Returns:
go.Figure: A plotly figure with the molecular structure.
"""
fig = go.Figure()
for mesh in self.get_molecular_mesh(resolution=resolution, rel_cutoff=rel_cutoff):
fig.add_trace(mesh)
# Fix aspect ratio such that the molecule is displayed undistorted
fig.update_layout(scene_aspectmode='data')
# Remove axes and labels
# fig.update_scenes(
# xaxis=dict(showgrid=False, zeroline=False, showticklabels=False, title=""),
# yaxis=dict(showgrid=False, zeroline=False, showticklabels=False, title=""),
# zaxis=dict(showgrid=False, zeroline=False, showticklabels=False, title=""),
# )
return fig
def circle_coordinates(point: np.ndarray, v: np.ndarray, radius: float, num_points: int) -> tuple[np.ndarray, np.ndarray, np.ndarray]:
"""
Generate cartesian coordinates of a circle in 3D space.
Args:
point (np.ndarray): A point on the circle.
v (np.ndarray): A vector perpendicular to the circle.
radius (float): The radius of the circle.
num_points (int): The number of points to generate on the circle.
Returns:
np.ndarray: A 2D array of shape (num_points, 3) containing the cartesian coordinates of the circle.
"""
# Normalize the vector v
v = v / np.linalg.norm(v)
# Create a vector u that is perpendicular to v
u = np.cross(v, [1, 0, 0]) if np.linalg.norm(np.cross(v, [1, 0, 0])) > 1e-6 else np.cross(v, [0, 1, 0])
u = u / np.linalg.norm(u)
# Create a vector w that is perpendicular to both u and v
w = np.cross(u, v)
# Generate the cartesian coordinates of the circle
theta = np.linspace(0, 2 * np.pi, num_points)
circle_coords = point + radius * (u * np.cos(theta)[:, np.newaxis] + w * np.sin(theta)[:, np.newaxis])
x, y, z = circle_coords[:, 0], circle_coords[:, 1], circle_coords[:, 2]
return x, y, z
def make_arrow_mesh(at: np.ndarray, dir: np.ndarray, resolution: int = 16, radius: float = 0.05) -> go.Mesh3d:
tip = at + dir
bottom = at - dir
v = bottom - tip
x1, y1, z1 = circle_coordinates(bottom, v, radius=radius, num_points=resolution//4)
x2, y2, z2 = circle_coordinates(tip, v, radius=radius, num_points=resolution//4)
x, y, z = np.concatenate((x1, x2)), np.concatenate((y1, y2)), np.concatenate((z1, z2))
arrow_mesh = go.Mesh3d(
x=x,
y=y,
z=z,
color="#444444",
opacity=1,
alphahull=0,
# name=f'{self.elements[i]}-{self.elements[j]}',
hoverinfo='none', # No hover info at all
)
return arrow_mesh
class Trajectory:
def __init__(self, coordinates: list[XYZ]):
self.coordinates = coordinates
self.vibration_vectors = None
@classmethod
def from_opt_output(cls, output_file: str):
"""
Creates a Trajectory object from an ORCA output file.
Args:
output_file (str): The path to the ORCA output file.
Returns:
Trajectory: A Trajectory object.
"""
parsed = qcp.read_qc_file(output_file)
return cls([XYZ(frame) for frame in parsed["coordinates"]])
@classmethod
def from_trajectory_file(cls, trajectory_file: str):
"""
Creates a Trajectory object from a trajectory file.
Args:
trajectory_file (str): The path to the trajectory file.
Returns:
Trajectory: A Trajectory object.
"""
return cls([frame for frame in Trajectory.xyz_generator(trajectory_file)])
@classmethod
def from_vibration_output(cls, output_file: str, mode: int):
"""
Creates a Trajectory object from an ORCA output file
containing a frequency calculation at the given mode.
Args:
output_file (str): The path to the ORCA output file.
Returns:
Trajectory: A Trajectory object.
"""
proc = sub.run(f"orca_pltvib {output_file} {mode}", shell=True, text=True, capture_output=True)
trajectory_file = f"{output_file}.v{mode:03d}.xyz"
xyz = Trajectory.from_trajectory_file(trajectory_file)
xyz.vibration_vectors = Trajectory.get_vibration_vectors(trajectory_file)
os.remove(trajectory_file)
return xyz
@staticmethod
def xyz_generator(trajectory_file: str):
"""
Generator that yields XYZ objects from a trajectory file.
Args:
trajectory_file (str): The path to the trajectory file.
Yields:
XYZ: An XYZ object.
"""
with open(trajectory_file) as f:
lines = f.readlines()
block_size = int(lines[0].strip()) + 2
for i in range(0, len(lines), block_size):
yield XYZ.from_list([line.split() for line in lines[i+2:i+block_size]])
@staticmethod
def get_vibration_vectors(trajectory_file: str):
"""
Returns the vibration vectors from a trajectory file. Must be generated
by orca_pltvib, so it contains the vectors in the last three columns.
Args:
trajectory_file (str): The path to the trajectory file.
Returns:
list[list[float]]: A list of vibration vectors.
"""
with open(trajectory_file) as f:
lines = f.readlines()
block_size = int(lines[0].strip())
lines = lines[2:]
vibration_vectors = [[float(element) for element in line.split()[-3:]] for line in lines[:block_size]]
print(vibration_vectors)
return vibration_vectors
def get_vibration_vectors_mesh(self) -> list[go.Mesh3d]:
"""
"""
assert self.vibration_vectors is not None, "Vibration vectors must be generated first"
meshes = []
first_frame_coords = self.coordinates[0].xyz
for i,vector in enumerate(self.vibration_vectors):
meshes.append(make_arrow_mesh(at=first_frame_coords[i], dir=vector))
return meshes
def get_molecular_viewer_animated(self, resolution: int = 64, rel_cutoff: float = 0.5, arrows: bool = False) -> go.Figure:
"""
Returns a plotly figure with the molecular structure of the trajectory. The figure comes with an animation.
Args:
resolution (int, optional): The number of points to be generated. Defaults to 64.
rel_cutoff (float, optional): Bond are only shown if the interatomic distance is smaller than the sum of the atomic vdw radii multiplied with this parameter.
Defaults to 0.5.
Returns:
go.Figure: A plotly figure with the molecular structure.
"""
fig = go.Figure()
# Get and add all the meshes
mesh_elements = [0]
frames = []
for xyz in self.coordinates:
data = []
meshes = xyz.get_molecular_mesh(resolution=resolution, rel_cutoff=rel_cutoff)
mesh_elements.append(len(meshes))
for mesh in meshes:
data.append(mesh) # was in double parenthesis
if arrows and self.vibration_vectors is not None:
data += self.get_vibration_vectors_mesh()
frames.append(go.Frame(data=data))
fig = go.Figure(
data=data, # initial frame before animation
frames=frames
)
fig.update_layout(
updatemenus=[{
"type": "buttons",
"buttons": [{
"label": "Play",
"method": "animate",
"args": [None, {
"frame": {"duration": 200}
}]
}],
}],
scene_aspectmode='data'
)
return fig
def get_molecular_viewer(self, resolution: int = 64, rel_cutoff: float = 0.5) -> go.Figure:
"""
Returns a plotly figure with the molecular structure of the trajectory. The figure comes with a slider that allows to switch between frames.
Args:
resolution (int, optional): The number of points to be generated. Defaults to 64.
rel_cutoff (float, optional): Bond are only shown if the interatomic distance is smaller than the sum of the atomic vdw radii multiplied with this parameter.
Defaults to 0.5.
Returns:
go.Figure: A plotly figure with the molecular structure.
"""
# Get and add all the meshes
mesh_elements = [0]
fig = go.Figure()
for xyz in self.coordinates:
meshes = xyz.get_molecular_mesh(resolution=resolution, rel_cutoff=rel_cutoff)
mesh_elements.append(len(meshes))
for mesh in meshes:
fig.add_trace(mesh)
mesh_elements_cumulative = list(accumulate(mesh_elements))
# Make the visibility list
slider_steps: list[dict] = []
for i in range(len(self.coordinates)):
step = {
"method": "restyle",
"args": [{"visible": [False] * mesh_elements_cumulative[-1]}]
}
step["args"][0]["visible"][mesh_elements_cumulative[i]:mesh_elements_cumulative[i+1]] = [True] * mesh_elements[i+1]
slider_steps.append(step)
# Make the slider
slider = {
"active": 0,
"currentvalue": {"prefix": "Frame: "},
"steps": slider_steps
}
fig.update_layout(
sliders=[slider]
)
fig.update_layout(scene_aspectmode='data')
return fig