helper.py

"""Helper functions for the tutorial on coarse-grained molecular optimization"""

import warnings
from pathlib import Path
from subprocess import run, CalledProcessError
import logging
import pandas as pd
import numpy as np
import torch
from scipy.stats import norm
import matplotlib.pyplot as plt
import seaborn as sns
import MDAnalysis as mda
from sklearn.gaussian_process import GaussianProcessRegressor
from sklearn.gaussian_process.kernels import RBF

logging.getLogger("pymbar").setLevel(logging.ERROR)
from alchemlyb.estimators import MBAR
from alchemlyb.parsing.gmx import extract_u_nk

SEED = 18
torch.manual_seed(SEED)
np.random.seed(SEED)


def _run_command(command: str, tries: int = 2):
    """
    Run a shell command with a specified number of tries in case of failure.
    If the command fails after the specified number of tries, the exception is raised.
    :param command: The shell command to execute.
    :param tries: The number of total attempts to run the command.
    """
    for attempt in range(tries):
        try:
            run(command, shell=True, check=True)
            return
        except CalledProcessError as e:
            print("An error occured. Retrying...")
            if attempt == tries - 1:
                raise e


def run_molecule_simulations(
    molecule: str, n_threads: int = 1, delete_on_failure: bool = False
):
    """
    Run all simulations required for thermodynamic integration of a two-bead molecule in water,
    hexane, and a water-hexane mixture. Since the minimized structure files are provided, the
    energy minimization step is skipped.
    :param molecule: A string representing the two bead molecule in the format 'A-B'.
    :param n_threads: Number of threads to use for the simulations (GROMACS mdrun -nt option).
    :param delete_on_failure: Delete the system simulation directory if simulation fails.
    """
    if "-" not in molecule or len(molecule.split("-")) != 2:
        raise ValueError("Molecule must be in the format 'A-B'")
    n_threads = max(1, int(n_threads))
    try:
        molecule_path = Path("simulations") / molecule
        for system in ["water", "hexane", "mixture"]:
            ### Directory setup ###
            system_path = molecule_path / system
            system_path.mkdir(parents=True, exist_ok=True)
            ### Setup the system topology ###
            if not Path(system_path / "system.top").exists():
                topology = Path(f"{system}-lig.top").read_text(encoding="utf-8")
                for placeholder, replacement in zip(["B1", "B2"], molecule.split("-")):
                    topology = topology.replace(placeholder, replacement)
                Path(system_path / "system.top").write_text(topology, encoding="utf-8")
            ### Adjust ligand pulling in equilibration based on the environment. ###
            if system == "mixture":
                _run_command("sed -i 's/^;pull/pull/g' equilibration.mdp")
            else:
                _run_command("sed -i 's/^pull/;pull/g' equilibration.mdp")
            ### Perform system equilibration ###
            if not Path(system_path / "equilibration.gro").exists():
                command = (
                    # Prepare the simulation input file
                    "gmx grompp -f equilibration.mdp "  # Input: Parameters
                    + f"-c {system}-lig.gro "  # Input: Starting structure
                    + f"-p {system_path}/system.top "  # Input: System topology
                    + (
                        "-n mixture-lig.ndx " if system == "mixture" else ""
                    )  # Input: Index file if solvent mixture simulation
                    + f"-o {system_path}/equilibration.tpr "  # Output: Simulation file
                    + f"-po {system_path}/equilibration.out.mdp "  # Output: Full parameter backup
                    + f">> {system_path}/simulation.run.log 2>&1 "  # Redirect output to a log file
                    # Perform the simulation
                    + f"&& gmx mdrun -deffnm {system_path}/equilibration -nt {n_threads} "
                    + f">> {system_path}/simulation.run.log 2>&1"  # Redirect output to a log file
                )
                _run_command(command)
            print(f"Running simulation 1/9 in {system}", end="\r")
            ### Adjust the number of integration steps and pull setup based on the environment. ###
            command = 'sed -i "s/^nsteps.*/nsteps                   = {nsteps}/" lambda-run.mdp'
            if system == "mixture":
                _run_command(command.format(nsteps=30000))
                # Enable ligand pulling in mixture simulations
                _run_command("sed -i 's/^;pull/pull/g' lambda-run.mdp")
            else:
                _run_command(command.format(nsteps=20000))
                # Disable ligand pulling in pure solvent simulations
                _run_command("sed -i 's/^pull/;pull/g' lambda-run.mdp")
            ### Run eight lambda-step simulations ###
            for i in range(8):
                lambda_path = system_path / f"lambda{i}"
                lambda_path.mkdir(exist_ok=True)
                if not Path(lambda_path / "production.gro").exists():
                    _run_command(
                        f'sed -i "s/^init-lambda-state.*/init-lambda-state        = {i}/" '
                        + "lambda-run.mdp"
                    )
                    command = (
                        # Prepare the simulation input file
                        "gmx grompp -f lambda-run.mdp "  # Input: Parameters
                        + f"-c {system_path}/equilibration.gro "  # Input: Starting structure
                        + f"-p {system_path}/system.top "  # Input: System topology
                        + (
                            "-n mixture-lig.ndx " if system == "mixture" else ""
                        )  # Input: Index file if solvent mixture simulation
                        + f"-o {lambda_path}/production.tpr "  # Output: Simulation file
                        + f"-po {lambda_path}/production.out.mdp "  # Output: Full parameter backup
                        + f">> {lambda_path}/simulation.run.log 2>&1 "  # Redirect output to file
                        # Perform the simulation
                        + f"&& gmx mdrun -deffnm {lambda_path}/production -nt {n_threads} "
                        + f">> {lambda_path}/simulation.run.log 2>&1"  # Redirect output to file
                    )
                    _run_command(command)
                print(f"Running simulation {i + 2}/9 in {system}", end="\r")
    except Exception as e:
        if delete_on_failure:
            _run_command(f"rm -r {system_path}")
        print(f"Failed to simulate {molecule}, please retry")
        raise e


def calculate_free_energy(
    molecule: str, system: str, print_results: bool = False
) -> tuple[float, float]:
    """
    Calculate solvation free energies for a given molecule and system using the
    MBAR algorithm (https://doi.org/10.1063/1.2978177). This function assumes
    completed simulations.
    :param molecule: A string representing the two bead molecule in the format 'A-B'.
    :param system: One of the three systems: 'water', 'hexane', 'mixture'
    :param print_results: Print free energy and error estimation
    :returns: The solvation free energy and an uncertainty estimate in kcal/mol.
    """
    ### Collect data from simulation output files ###
    path = Path("simulations") / molecule / system
    if not path.exists():
        raise ValueError(f"Simulations for {molecule}/{system} not found")
    xvg_files = [p / "production.xvg" for p in path.iterdir() if p.is_dir()]
    u_nk_list = [extract_u_nk(f, T=300) for f in xvg_files]
    u_nk_combined = pd.concat(u_nk_list)
    ### Perform MBAR calculation ###
    with warnings.catch_warnings():
        warnings.simplefilter("ignore")
        mbar = MBAR().fit(u_nk_combined)
    ### Extract result and convert unit ###
    free_energy = float(mbar.delta_f_.iloc[0, -1]) * 0.5924  # Convert kT to kcal/mol
    d_free_energy = (
        float(mbar.d_delta_f_.iloc[0, -1]) * 0.5924
    )  # Convert kT to kcal/mol
    if print_results:
        print(f"dG_{system} = {free_energy:.3f} ± {d_free_energy:.3f} kcal/mol")
    return free_energy, d_free_energy


def convert_to_sdf(structure_path, trajectory_path, output_path, stride=1):
    """GROMACS .tpr/.xtc -> multi-frame .sdf with explicit bonds."""
    with warnings.catch_warnings():
        warnings.filterwarnings("ignore")
        u = mda.Universe(structure_path, trajectory_path)
    bonds = [(b.atoms[0].index + 1, b.atoms[1].index + 1) for b in u.bonds]
    elem = [
        "O" if r == "W" else "C" for r in u.atoms.resnames
    ]  # hexane=C, water=O (color tags)
    na, nb = len(u.atoms), len(bonds)
    with open(output_path, "w") as fh:
        for ts in u.trajectory[::stride]:
            fh.write("frame\n  MDAnalysis\n\n")
            fh.write("%3d%3d  0  0  0  0  0  0  0  0999 V2000\n" % (na, nb))
            for i, a in enumerate(u.atoms):
                x, y, z = a.position
                fh.write(
                    "%10.4f%10.4f%10.4f %-3s 0  0  0  0  0  0  0  0  0  0  0  0\n"
                    % (x, y, z, elem[i])
                )
            for p, q in bonds:
                fh.write("%3d%3d  1  0\n" % (p, q))
            fh.write("M  END\n$$$$\n")


def visualize_latent_space(latent_space: torch.Tensor, molecules: list[str]):
    """
    Visualize the latent space representation of a set of molecules
    :param latent_space: The Nx2 dimenesional encoding values of the N molecules
    :param molecules: The list of molecule labels used for annotations of the scatter plot.
    """
    ### Scatter plot of all encoded molecules ###
    fig, ax = plt.subplots(figsize=(8, 8))
    ax.scatter(latent_space[:, 0], latent_space[:, 1], s=15, color="red")
    for n, molecule in enumerate(molecules):
        ax.annotate(str(molecule), (latent_space[n, 0], latent_space[n, 1]), fontsize=6)
    ax.set(xlabel="Encoding dimension 1", ylabel="Encoding dimension 2")
    ### Visualize encoding space densities using KDE plots ###
    ax_right = fig.add_axes([0.901, 0.11, 0.04, ax.get_position().height])
    ax_top = fig.add_axes([0.125, 0.88, ax.get_position().width, 0.04])
    sns.kdeplot(
        y=latent_space[:, 1],
        ax=ax_right,
        color="gray",
        fill=True,
        linewidth=0.2,
        bw_adjust=0.5,
    )
    sns.kdeplot(
        x=latent_space[:, 0],
        ax=ax_top,
        color="gray",
        fill=True,
        linewidth=0.2,
        bw_adjust=0.5,
    )
    ax_right.set(xlabel=None, ylabel=None, ylim=ax.get_ylim())
    ax_right.axis("off")
    ax_top.set(xlabel=None, ylabel=None, xlim=ax.get_xlim())
    ax_top.axis("off")
    plt.show()


def acquisition_function(
    values: np.ndarray,
    uncertainty: np.ndarray,
    best_known_value: float,
    xi: float = 0.0,
) -> np.ndarray:
    """
    The acquisition function implemented here is the expected improvement. See this
    link for an explanation:
    https://ekamperi.github.io/machine%20learning/2021/06/11/acquisition-functions.html#expected-improvement-ei
    :param values: Mean prediction values from the surrogate model.
    :param uncertainty: Predicted standard deviation from the surrogate model.
    :param best_known_value: Best so far observed value.
    :param xi: Shift best known value to achieve more (xi > 0) or less (xi < 0) exploration
    """
    z = values - best_known_value - xi
    return z * norm.cdf(z / uncertainty) + uncertainty * norm.pdf(z / uncertainty)


def argmax_with_excluded_indices(
    values: np.ndarray, excluded_indices: list[int]
) -> int:
    """
    Find the index of the maximum value in an array, excluding specified indices. If
    there are multiple maximum values, one of them is randomly selected. The numpy
    masked array is used to handle the exclusion of indices.
    :param values: A numpy array of values.
    :param excluded_indices: A list of indices to exclude from the search for the maximum.
    :returns: The index of the maximum value, excluding specified indices.
    """
    np.random.seed(SEED)
    if len(values) == 0:
        raise ValueError("The input array is empty.")
    mask = np.isin(np.arange(len(values)), list(excluded_indices))
    values = np.ma.masked_array(values, mask=mask)
    return np.random.choice(np.nonzero(values == values.max())[0])


class SurrogateModel:
    """
    A surrogate model for predicting target values based on a latent space representation.
    This model uses a Gaussian Process Regressor with a radial basis function (RBF) kernel.
    It is designed to fit a set of data points and predict values for the latent space.
    """

    def __init__(self, latent_space: np.ndarray):
        """
        Initialize the surrogate model with a latent space representation. The latent space
        is used for all predictions of the model.
        :param latent_space: A Nx2 numpy array representing the 2D latent space of N molecules.
        """
        kernel = RBF(length_scale=0.5, length_scale_bounds=(0.05, 2))
        self.gaussian_process = GaussianProcessRegressor(
            kernel=kernel, n_restarts_optimizer=9, alpha=0.05, random_state=SEED
        )
        self.latent_space = latent_space
        if not isinstance(self.latent_space, np.ndarray):
            raise ValueError("Latent space must be a numpy array.")
        if self.latent_space.ndim != 2 or self.latent_space.shape[1] != 2:
            raise ValueError("Latent space must be a Nx2 dimensional array.")

    def fit(self, data: dict[int, float]):
        """
        Fit the surrogate model to the provided data. The data should be a dictionary
        where keys are indices of the latent space and values are the target values.
        :param data: A dictionary with keys as indices of the latent space and values as
            target values.
        """
        if not isinstance(data, dict):
            raise ValueError("Data must be a dictionary.")
        if len(data) == 0:
            return
        x = [self.latent_space[i] for i in data.keys()]
        y = list(data.values())
        with warnings.catch_warnings():
            warnings.simplefilter("ignore")
            self.gaussian_process.fit(x, y)

    def predict(self) -> tuple[np.ndarray, np.ndarray]:
        """
        Predict the target values for the latent space using the fitted surrogate model.
        The function does not take arguments, as it uses the latent space provided during
        the initialization of the model for predictions.
        :returns: A tuple containing the predicted values and their standard deviations.
        """
        return self.gaussian_process.predict(self.latent_space, return_std=True)