Skip to content
Snippets Groups Projects
qc_support.py 4.28 KiB
Newer Older
"""
qc_support.py

Support functions for Quantum Chemistry (QC) analysis and visualization.

This module provides utility functions to support QC analysis and visualization,
including 3D molecular visualization and file conversion.

Meant to be used with Jupyter Notebooks in the computational chemistry course
of the TU Darmstadt.
"""

import py3Dmol as pm
from pathlib import Path
import subprocess as sub
import os

from ipywidgets import interact,fixed,IntSlider
import ipywidgets

def mol_to_xyz(file_path: str) -> None:
    """
    Convert a mol file to an xyz file. A new file with the same name
    but .xyz file ending is created.
    @param file_path: name/path of the mol file that is converted.
    @return: None
    """
    lines = Path(file_path).read_text().split("\n")
    natoms = int(lines[3].split()[0])
    lines = lines[4:4+natoms]
    lines = [line.split()[:4] for line in lines]
    lines = [f"{line[3]} {line[0]} {line[1]} {line[2]}" for line in lines]
    out = [f"{natoms}", f"converted from {file_path}"]
    out += lines
    Path(file_path.replace(".mol", ".xyz")).write_text("\n".join(out))


def show_trajectory(file_path: str) -> pm.view:
    v = pm.view()
    v.addModelsAsFrames(Path(file_path).read_text(), "xyz")
    v.zoomTo()
    v.setStyle({'sphere':{'radius':0.2},'stick':{'radius':0.1}})
    v.animate({"loop" : "forward", "interval" : "2000"})
    return v

def show_mol(file_path: str, numbers: bool=False, labels: list[str]=[]) -> pm.view:
    v = pm.view()
    xyz = Path(file_path).read_text()
    v.addModel(xyz, "xyz")
    v.zoomTo()
    v.setStyle({'sphere':{'radius':0.2},'stick':{'radius':0.1}})
    if numbers:
        num = int(xyz.split()[0])
        for i in range(num):
            v.addLabel(f"{i}",
                       {"backgroundColor" : "white",
                        "backgroundOpacity" : "0.2",
                        "fontColor" : "black"},
                       {"serial" : f"{i}"})
    if len(labels) != 0:
        num = int(xyz.split()[0])
        assert len(labels) == num, "Length of label list does not match number of atoms."
        for i,label in enumerate(labels):
            v.addLabel(f"{label}",
                       {"backgroundColor" : "white",
                        "backgroundOpacity" : "0.2",
                        "fontColor" : "black"},
                       {"serial" : f"{i}"})
    return v

def read_xyz_trajectory_str(file: str) -> list[str]:
    with open(file) as f:
        lines = f.readlines()
    block_size = int(lines[0]) +2
    frames = []
    for i in range(0,len(lines), block_size):
        block = lines[i:i+block_size]
        frames.append("".join(block))
    return frames

def show_vibration(file_path: str, index: int) -> pm.view:
    """
    Returns a 3Dmol view that is an animation of a vibrational mode.
    @param file_path: String. Path to an ORCA output file (.out) or hessian (.hess)
    @param index: Integer. Number of the vibrational mode to be animated
    @return: 3Dmol view. Animation of the given vibration.
    """
    assert ".out" in file_path or ".hess" in file_path, f"""\
The given file is not compatible. Provide either an ORCA output
file (.out) or a hessian (.hess) file.
    """
    proc = sub.run(f"orca_pltvib {file_path} {index}", shell=True, text=True, capture_output=True)
    v = show_trajectory(f"{file_path}.v{index:03d}.xyz")
    os.remove(f"{file_path}.v{index:03d}.xyz")
    return v

class FrameViewer():
    def __init__(self, traj_file: str) -> None:
        self.frames = read_xyz_trajectory_str(traj_file)
        self.v = pm.view()
        self.v.addModel(self.frames[0], "xyz")
        self.v.zoomTo()
        self.v.setStyle({'sphere':{'radius':0.2},'stick':{'radius':0.1}})
        return

    def show_frame(self, index=0):
        """Updates the FrameViewer to show the structure of the
        molecule at the given index."""
        self.v.removeAllModels()
        self.v.addModel(self.frames[index], "xyz")
        self.v.setStyle({'sphere':{'radius':0.2},'stick':{'radius':0.1}})
        return self.v.show()
    
    def frame_slider(self):
        """Returns a Slider with a view on the given optimization
        trajectory. The slider can be used to change the frame."""
        return interact(self.show_frame, 
                        index=ipywidgets.IntSlider(min=0,max=len(self.frames)-1, step=1))