clustering.py

"""
Utile python file for Clusturing analysis after Normal Mode Sampling    
"""

__author__  = "Pierre-Alexandre HO"
__date__    = "2025-06-25"
__version__ = "1.0"


# standard python modules

import numpy as np
import pandas as pd
import seaborn as sns
import matplotlib.pyplot as plt
import random

from typing import List, Union, Sequence, Any
from scipy.stats import norm
from scipy.cluster.hierarchy import linkage, dendrogram, fcluster


random.seed(10)

def add_coordinates (file:str)-> dict:
    """
    Reads a file containing molecular coordinates of conformers data and constructs a dictionary.
    
    The file is expected to have the following format:
      - The first line with a single number indicates the number of atoms per molecule (mol_len).
      - A line starting with "Converted" or "NMS" signals the start of a new conformer.
      - Subsequent lines with 4 entries contain the atom name and its x, y, z coordinates.
      - When the number of coordinate lines equals mol_len, the conformer is considered complete.
      
    Note: The 1st conformation: 'Original' is not taking account avoiding to appears many times
    
    Parameters
    ----------
    file : str
        Path to the input file containing molecular coordinate data.
    
    Returns
    -------
    conformers_coordinates : dict
        Dictionary where each key is a conformer identifier (e.g., "conf_1") and each value is a list:
            {conf_i: [atom_list, np.array(x_list), np.array(y_list), np.array(z_list)]}
    """
    
    mol_len = 0
    conformers_coordinates = {}
    count = 0
    conf_mum = 0
    with open (file, "r") as source:
        for line in source:
            line = line.strip().split()
            
            if  mol_len == 0 and len(line) == 1:
                mol_len = int(line[0])
            
            elif line[0] in ["Converted", 'NMS']:
                atom_list = []
                x_list = []
                y_list = []
                z_list = []
                count = 0
                conf_mum += 1
            
            elif len(line) == 4 and conf_mum > 0:
                atom_list.append(line[0])
                x_list.append(float(line[1]))
                z_list.append(float(line[3]))
                y_list.append(float(line[2]))
                count += 1
            
            if count == mol_len:
                conformers_coordinates[f"conf_{conf_mum}"] = [
                    atom_list,
                    np.array(x_list),
                    np.array(y_list),
                    np.array(z_list)
                ]
    
    return conformers_coordinates


def calculate_rmsd(coords1:np.ndarray, coords_ref:np.ndarray)-> float:
    """
    Calculate the Root Mean Square Deviation (RMSD) between two sets of atomic coordinates.

    Parameters
    ----------
    coords1 : np.ndarray
        A 2D array of shape (n_atoms, 3) representing the first set of atomic coordinates.
    coords_ref : np.ndarray
        A 2D array of shape (n_atoms, 3) representing the second set of atomic coordinates use as reference.

    Returns
    -------
    float
        The RMSD value between the two sets of coordinates.

    Raises
    ------
    ValueError
        If the input arrays do not have the same shape.
    """
    if coords1.shape != coords_ref.shape:
        raise ValueError("Input coordinate arrays must have the same shape.")

    diff = coords1 - coords_ref
    return np.sqrt(np.mean(np.sum(diff**2, axis=1)))


def compute_rmsd_matrix(structures: List[np.ndarray]) -> np.ndarray:
    """
    Compute a distance matrix using RMSD as the distance metric to compare each pair of structures.

    This function calculates the Root Mean Square Deviation (RMSD) between each pair of structures
    and constructs a symmetric distance matrix where each element (i, j) represents the RMSD
    between the i-th and j-th structures.

    Parameters
    ----------
    structures : List[np.ndarray]
        A list of structures, where each structure is represented as a NumPy array.
        Each array should have the same shape and represent atomic coordinates.

    Returns
    -------
    np.ndarray
        A symmetric distance matrix of shape (n, n), where n is the number of structures.
        Each element (i, j) in the matrix represents the RMSD between the i-th and j-th structures.

    Examples
    --------
    >>> structures = [np.array([[1, 2, 3], [4, 5, 6]]), np.array([[1, 2, 3], [7, 8, 9]])]
    >>> compute_rmsd_matrix(structures)
    array([[ 0.        ,  3.        ],
           [ 3.        ,  0.        ]])
    """
    n = len(structures)
    distance_matrix = np.zeros((n, n))
    for i in range(n):
        for j in range(i + 1, n):
            d = calculate_rmsd(structures[i], structures[j])
            distance_matrix[i, j] = d
            distance_matrix[j, i] = d
    return distance_matrix


def multi_line_chart(df: pd.DataFrame, x_axis: str, y_axis: str, color_class: str, title: str, y_label: str, 
    color_label: str) -> None:
    """
    Generates a multi-line chart with a colorbar using Seaborn and Matplotlib.

    Parameters
    ----------
    df : pd.DataFrame
        The dataset containing the data to plot.
    x_axis : str
        Column name for the x-axis.
    y_axis : str
        Column name for the y-axis.
    color_class : str
        Column name used for coloring the lines.
    title : str
        Title of the plot.
    y_label : str
        Label for the y-axis.
    color_label : str
        Label for the colorbar.

    Returns
    -------
    None
        The function directly plots the figure.
    """

    fig, ax = plt.subplots(figsize=(20, 10))
    cmap = sns.color_palette("rocket_r", as_cmap=True)
    sns.lineplot(data=df, x=x_axis, y=y_axis, ax=ax, hue=color_class, palette=cmap)

    # Colorbar
    norm = plt.Normalize(vmin=df[color_class].min(), vmax=df[color_class].max())
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    cbar = fig.colorbar(sm, ax=ax)

    cbar.set_label(color_label)
    ax.set_ylabel(y_label)
    ax.set_title(title)
    ax.get_legend().remove()
    plt.show()
    

def plot_dendrogram(distance_matrix: np.ndarray, link_method: str, title: str, xlabel: str,
                    ylabel: str, labels: list, cut_off: float = None) -> dict:
    """
    Plots a dendrogram based on the provided distance matrix using hierarchical clustering.
    Optionally, cuts the dendrogram at the specified cut-off threshold and identifies the clusters.
    
    Parameters:
    ----------
    distance_matrix : np.ndarray
        A square distance matrix (2D array) containing pairwise distances between points.
    
    link_method : str
        The linkage method to use for hierarchical clustering. Options are: 
        'single', 'complete', 'average', 'ward', etc.
    
    title : str
        The title of the dendrogram plot.
    
    xlabel : str
        The label for the x-axis.
    
    ylabel : str
        The label for the y-axis.
    
    labels : list
        A list of labels corresponding to the data points in the distance matrix.
    

    cut_off : float, optional
        The threshold value for cutting the dendrogram into flat clusters. 
        If None (default), no cut-off is applied.
    
    Returns:
    -------
    np.ndarray
        A structured array with the columns: Temperature, Conformation, and Cluster ID.
    """
    # Perform hierarchical clustering using the specified linkage method
    Z = linkage(distance_matrix, method=link_method)
    
    # Apply the cut-off threshold if specified, otherwise use default hierarchical clusters
    # Default cut-off as 70% of the maximum height of the linkage as it write in the doc
    if cut_off is None:
        cut_off =  0.7 * max(Z[:, 2])  
    clusters = fcluster(Z, t=cut_off, criterion='distance')
    # Group the structures into clusters (mapping Temperature/Conformation to Cluster)
    result = {}
    for i, label in enumerate(labels):
        result[label] = clusters[i]
    
    # Plot the dendrogram
    fig, ax = plt.subplots(figsize=(10, 7))
    dendrogram(Z, labels=labels, color_threshold=cut_off)
    ax.set_title(title)
    ax.set_xlabel(xlabel)
    ax.set_ylabel(ylabel)    
    ax.axhline(cut_off, linestyle="--", color="black", linewidth =0.3)

    return result

def build_coordinate_matrix(
    data: dict[str, dict[str, tuple[np.ndarray, np.ndarray, np.ndarray]]])->tuple[np.ndarray, list[str], list[str]]:
    """
    Converts a nested dictionary into a numerical matrix where each row represents a conformation, 
    and each column corresponds to atomic coordinates (x, y, z).

    Parameters
    ----------
    data : dict
        Nested dictionary where:
        - Keys (str) represent different sampling conditions (e.g., temperature).
        - Values (dict) contain conformations as keys and a tuple of atomic data as values:
            conformation # : (atom_list, np.array(x), np.array(y), np.array(z)).

    Returns
    -------
    X : np.ndarray
        A 2D array of shape (n_conformations, n_atoms * 3) containing concatenated coordinates.
    temperature : list of str
        List of sampling labels (e.g., temperatures) corresponding to each conformation.
    conf_num : list of str
        List of conformation labels corresponding to each structure.
    """
    conformations = []
    temperature = []
    conf_num = []

    for sampling, conformations_dict in data.items():
        for conf_name, (atom_list, x, y, z) in conformations_dict.items():
            coords = np.concatenate([x, y, z])  # Flatten coordinates into a single vector
            conformations.append(coords)
            temperature.append(int(sampling.replace('K', '')))  # Store the sampling condition (e.g., temperature)
            conf_num.append(int(conf_name.replace('conf_', '')))  # Store conformation identifier

    return np.array(conformations, dtype=float), temperature, conf_num  # Ensure numerical data


def plot_pca(
    df_pca: pd.DataFrame, explained_var: Union[np.ndarray, Sequence[float]], color_class: Union[str, Sequence]
    )->None:
    """
    Plots the PCA projection of conformations.

    Parameters
    ----------
    df_pca : pd.DataFrame
        DataFrame containing PCA projections with columns "PC1" and "PC2".
    explained_var : array-like
        Array or list containing the explained variance ratios for each principal component.
        Values are expected as fractions (e.g., 0.65 for 65%).
    color_class : str or array-like
        Either the name of a column in df_pca to be used for color coding the points, 
        or an array-like object of categorical values corresponding to each point.

    Returns
    -------
    None
        Displays a scatter plot of the PCA results.
    """
    fig, ax = plt.subplots(figsize=(12, 10))
    
    # Create scatter plot using seaborn
    sns.scatterplot(data=df_pca, x="PC1", y="PC2", hue=color_class, ax=ax)
    
    # Set axis labels with explained variance percentages
    ax.set_xlabel(f"PC1 ({explained_var[0]:.1%})")
    ax.set_ylabel(f"PC2 ({explained_var[1]:.1%})")
    
    # Set title and legend
    ax.set_title("PCA Projection of Conformations")
    plt.legend(title=str(color_class))
    
    plt.show()
    
    
def plot_density_pca(df_pca: pd.DataFrame, explained_var: list) -> None:
    """
    Plot the density distribution of the first two principal components from a PCA analysis.

    This function uses seaborn's kdeplot to visualize the density distribution of the first
    two principal components (PC1 and PC2) and annotates the plot with the percentage of
    explained variance for each component.

    Parameters
    ----------
    df_pca : pd.DataFrame
        A DataFrame containing the results of a PCA analysis. It should include columns named
        'PC1' and 'PC2' representing the first and second principal components, respectively.
    explained_var : list
        A list containing the explained variance ratios for the principal components.
        The first two elements should correspond to PC1 and PC2.

    Returns
    -------
    None
        This function does not return any value but displays a plot.
    """
    fig, ax = plt.subplots(figsize=(12, 10))
    sns.kdeplot(data=df_pca, x="PC1", y='PC2', ax=ax, fill=True, cmap="Blues")

    # Set axis labels with explained variance percentages
    ax.set_xlabel(f"PC1 ({explained_var[0]:.1%})")
    ax.set_ylabel(f"PC2 ({explained_var[1]:.1%})")

    # Set title and legend
    ax.set_title("Density Distribution of PCA Components", fontsize=14)
    plt.legend()
    plt.show()
    
    
def calc_logl(x: np.ndarray, mu: float, sd: float) -> float:
    """
    Helper function to calculate log-likelihood of data given a normal distribution.

    Parameters
    ----------
    x : np.ndarray
        Array of data points.
    mu : float
        Mean of the normal distribution.
    sd : float
        Standard deviation of the normal distribution.

    Returns
    -------
    float
        The log-likelihood of the data.
    """
    logl = 0
    for i in x:
        logl += np.log(norm.pdf(i, mu, sd))
    return logl


def find_optimal_k(data: np.ndarray) -> int:
    """
    Provide a numpy array, returns index to serve as cut-off for optimal segmentation.

    This function calculates the log-likelihood for different segmentations of the data
    and returns the index that maximizes the combined log-likelihood of the two segments.

    Parameters
    ----------
    data : np.ndarray
        Array of data points to be segmented.

    Returns
    -------
    int
        The index at which to split the data to maximize the log-likelihood.
    """
    profile_logl = []
    for q in range(1, len(data)):
        n = len(data)
        s1 = data[0:q]
        s2 = data[q:]
        mu1 = s1.mean()
        mu2 = s2.mean()
        sd1 = s1.std()
        sd2 = s2.std()

        # Calculate pooled standard deviation
        sd_pooled = np.sqrt((((q - 1) * (sd1 ** 2) + (n - q - 1) * (sd2 ** 2)) / (n - 2)))

        # Calculate combined log-likelihood for the current segmentation
        combined_logl = calc_logl(s1, mu1, sd_pooled) + calc_logl(s2, mu2, sd_pooled)
        profile_logl.append(combined_logl)

    # Return the index that maximizes the log-likelihood
    return np.argmax(profile_logl)


def plot_energy_rmsd(ener_rmsd: pd.DataFrame) -> None:
    """
    Plot the relationship between RMSD and Energy, colored by Temperature.

    This function creates a scatter plot of RMSD versus Energy, with points colored
    according to the Temperature. It includes a color bar to indicate the temperature scale.

    Parameters
    ----------
    ener_rmsd : pd.DataFrame
        A DataFrame containing the columns 'RMSD', 'Energy', and 'Temperature'.
        'RMSD' and 'Energy' are the coordinates for the scatter plot, and 'Temperature'
        is used to color the points.

    Returns
    -------
    None
        This function does not return any value but displays a plot.
    """
    fig, ax = plt.subplots(figsize=(12, 10))
    cmap = sns.color_palette("rocket_r", as_cmap=True)
    sns.scatterplot(data=ener_rmsd, x="RMSD", y="Energy", hue="Temperature", palette=cmap)

    ax.set_ylabel("Energy (Eh)")
    ax.set_xlabel("RMSD ($\AA$)")

    # Colorbar
    norm = plt.Normalize(vmin=ener_rmsd["Temperature"].min(), vmax=ener_rmsd["Temperature"].max())
    sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
    sm.set_array([])
    cbar = fig.colorbar(sm, ax=ax)
    cbar.set_label("Temperature (°K)")
    ax.get_legend().remove()

    plt.show()
    
    import pandas as pd
from typing import Dict, Any


def parse_cluster(nms_conformations: Dict[str, Dict[str, Any]], cluster_dic: Dict[str, Any]
                  ) -> pd.DataFrame:
    """
    Parse and organize data from two dictionaries into a pandas DataFrame.

    This function iterates over the `nms_conformations` dictionary, checks for specific conditions,
    and collects relevant data into lists. The collected data is then used to create a DataFrame.

    Parameters
    ----------
    nms_conformations : Dict[str, Dict[str, Any]]
        A nested dictionary where the outer key is a temperature (as a string) and the inner dictionary
        contains conformations with their associated values.
    cluster_dic : Dict[str, Any]
        A dictionary where keys are names formed by combining temperature and conformation,
        and values are cluster information.

    Returns
    -------
    pd.DataFrame
        A DataFrame with columns 'Temperature', 'Conformations', 'Clusters', and 'Structures'.
    """
    temperatures, conformations, structures, clusters = list(), list(), list(), list()

    for temp, confs in nms_conformations.items():
        if temp != '0K':
            for conf, values in confs.items():
                name = f'{temp}_{conf}'
                if name in cluster_dic:
                    temperatures.append(temp)
                    conformations.append(conf)
                    structures.append(values)
                    clusters.append(cluster_dic[name])

    return pd.DataFrame({
        'Temperature': temperatures,
        'Conformations': conformations,
        'Clusters': clusters,
        'Structures': structures
    })


def compute_medoid(df: pd.DataFrame, method: str = 'min') -> pd.Series:
    """
    Compute the medoid structure based on RMSD and return associated metadata.

    This function calculates the RMSD matrix for a set of structures and determines the medoid
    based on the specified method ('min' or 'max'). The medoid is the structure with the minimum
    or maximum sum of RMSD distances to all other structures.

    Parameters
    ----------
    df : pd.DataFrame
        A DataFrame containing a column 'Structures' with structural data.
    method : str, optional
        The method to determine the medoid. Options are 'min' (default) for the structure with
        the minimum sum of RMSD distances, and 'max' for the structure with the maximum sum.

    Returns
    -------
    pd.Series
        The row of the DataFrame corresponding to the medoid structure.
    """
    # Extract structures from the DataFrame
    structures = [np.array([struct[1], struct[2], struct[3]]) for struct in df['Structures']]

    # Compute the RMSD matrix
    rmsd_matrix = compute_rmsd_matrix(structures)

    # Determine the medoid index based on the specified method
    if method == 'min':
        medoid_index = np.argmin(rmsd_matrix.sum(axis=1))
    elif method == 'max':
        medoid_index = np.argmax(rmsd_matrix.sum(axis=1))
    else:
        raise ValueError("Method must be either 'min' or 'max'.")

    # Return the row of the DataFrame corresponding to the medoid structure
    return df.iloc[medoid_index]


def parse_medoid(cluster_df: pd.DataFrame, method: str = 'min') -> pd.DataFrame:
    """
    Parse a DataFrame containing cluster information and compute the medoid for each unique cluster.

    This function iterates over each unique cluster in the DataFrame, computes the medoid structure
    for each cluster using the specified method, and compiles the results into a new DataFrame.

    Parameters
    ----------
    cluster_df : pd.DataFrame
        A DataFrame containing columns 'Temperature', 'Conformations', 'Clusters', and 'Structures'.
    method : str, optional
        The method to determine the medoid. Options are 'min' (default) for the structure with
        the minimum sum of RMSD distances, and 'max' for the structure with the maximum sum.

    Returns
    -------
    pd.DataFrame
        A DataFrame with columns 'Temperature', 'Conformations', 'Clusters', and 'Structures',
        containing the medoid information for each unique cluster.
    """
    temperatures, conformations, structures, clusters = list(), list(), list(), list()

    for clust in set(cluster_df['Clusters']):
        # Filter the DataFrame for the current cluster
        df_temp = cluster_df[cluster_df['Clusters'] == clust]

        # Compute the medoid for the current cluster
        medoid_info = compute_medoid(df_temp, method=method)

        # Append the medoid information to the lists
        temperatures.append(medoid_info["Temperature"])
        conformations.append(medoid_info["Conformations"])
        structures.append(medoid_info["Structures"])
        clusters.append(medoid_info["Clusters"])

        # Clean up the temporary DataFrame
        del df_temp

    # Create and return a new DataFrame with the medoid information
    return pd.DataFrame({
        'Temperature': temperatures,
        'Conformations': conformations,
        'Clusters': clusters,
        'Structures': structures
    })
    
    
def write_to_xyz(name_file: str, representative_structures: pd.DataFrame) -> None:
    """
    Write representative molecular structures to an XYZ file.

    This function takes a DataFrame of representative structures and writes each structure
    to an XYZ file.

    Parameters
    ----------
    name_file : str
        The name of the file to write the XYZ data to. The file will be saved in the
        'freq_XTB/representative/' directory with the '.xyz' extension.
    representative_structures : pd.DataFrame
        A DataFrame containing molecular structures and associated metadata. The DataFrame
        should have columns 'Structures', 'Temperature', and 'Conformations'.

    Returns
    -------
    None
        This function does not return any value but writes the XYZ data to a file.
    """
    with open(f'freq_XTB/representative/{name_file}.xyz', 'w') as file:
        for row in representative_structures.itertuples(index=True):
            # Write the number of atoms
            file.write(f"{len(row.Structures[0])}\n")

            # Write a comment line with metadata
            file.write(f"NMS round: {row.Index + 1} - temperature = {row.Temperature} - conformation {row.Conformations}\n")

            # Write the atom coordinates
            for atom, x, y, z in zip(*row.Structures):
                file.write(f"{atom} {float(x):.6f} {float(y):.6f} {float(z):.6f}\n")