mini_numerical.py

import numpy as np
import math
import sys
import os
import matplotlib.pyplot as plt
from vmd_movie import generate_vmd_movie
from tqdm import tqdm

# Constantes
K_BOND = 450  # kcal/mol/Ų
I_EQ = 0.9572  # Å, distance d'équilibre
TETA_EQ = 104.5200  # degrés, angle d'équilibre
K_ANGLE = 55  # kcal/mol/rad²
H = 1e-4   # Pas pour le calcul numérique du gradient
STEP = 1e-3  # Pas de mise à jour des coordonnées
THRESHOLD =  0.5  # kcal/mol/A, Valeur par défaut pour le seuil GRMS 
MAX_ITERATIONS = 1000  # Nombre maximal d'itérations


def extract_coord_pdb(pdbfile: str):
    """
    Extraction des coordonnées atomiques à partir d'un fichier PDB.

    Paramètres:
        pdbfile (str): Chemin vers le fichier PDB contenant la molécule.

    Retourne:
        np.ndarray: Tableau contenant les coordonnées des atomes sous forme aplatie.
    """
    coordinates = []
    try:
        with open(pdbfile, "r") as pdb_file:
            for line in pdb_file:
                if not line.startswith("ATOM"):
                    continue
                x = float(line[30:39])
                y = float(line[38:46])
                z = float(line[46:54])
                coordinates.append([x, y, z])

        coordinates_array = np.array(coordinates)
        if coordinates_array.shape != (3, 3):
            raise ValueError("Le fichier PDB ne contient pas exactement 3 atomes.")
        return coordinates_array.flatten()
    
    except Exception as e:
        print(f"Erreur lors de l'extraction des coordonnées : {e}")
        sys.exit(1)


def calc_len_angle(coordinates: np.ndarray):
    """
    Calcul des longueurs de liaison O-H1, O-H2 et de l'angle H1-O-H2.

    Paramètres:
        coordinates (np.ndarray): Tableau des coordonnées atomiques.

    Retourne:
        tuple[float, float, float]: Longueur O-H1, Longueur O-H2, et angle H1-O-H2 (radians).
    """
    xO, yO, zO, xH1, yH1, zH1, xH2, yH2, zH2 = coordinates
    len1 = round(np.linalg.norm([xH1 - xO, yH1 - yO, zH1 - zO]), 4)
    len2 = round(np.linalg.norm([xH2 - xO, yH2 - yO, zH2 - zO]), 4)

    # Calcul du produit scalaire des vecteurs O->H1 et O->H2
    dot_product = (
        (xH1 - xO) * (xH2 - xO) + (yH1 - yO) * (yH2 - yO) + (zH1 - zO) * (zH2 - zO)
    )
    # Calcul de l'angle H1-O-H2
    angle = math.acos(dot_product / (len1 * len2))

    return len1, len2, angle


def calc_pot_energy(coordinates: np.ndarray):
    """
    Calcul de l'énergie totale de la molécule en fonction des coordonnées atomiques.
    
    Paramètres:
        coordinates (np.ndarray): Tableau des coordonnées atomiques (shape=(9,)).
    
    Retourne:
        float: Énergie totale en kcal/mol.
    """
    len1, len2, angle = calc_len_angle(coordinates)
    angle_eq_rad = TETA_EQ * math.pi / 180
    
    e_bond = K_BOND * (len1 - I_EQ) ** 2 + K_BOND * (len2 - I_EQ) ** 2
    e_angle = K_ANGLE * (angle - angle_eq_rad) ** 2
    
    return e_bond + e_angle

def calc_num_gradient(coordinates: np.ndarray, h: float):
    """
    Calcul du gradient numérique de l'énergie par différences finies.

    Paramètres:
        coordinates (np.ndarray): Tableau des coordonnées atomiques.
        h (float): Pas de dérivation numérique.

    Retourne:
        np.ndarray: Gradient de l'énergie par rapport aux coordonnées.
    """
    if np.linalg.norm(coordinates) < 1e-10:
        return np.zeros_like(coordinates)

    gradient = np.zeros_like(coordinates)
    for i in range(len(coordinates)):
        coords_plus_h = coordinates.copy()
        coords_minus_h = coordinates.copy()
        coords_plus_h[i] += h
        coords_minus_h[i] -= h

        e_plus_h = calc_pot_energy(coords_plus_h)
        e_minus_h = calc_pot_energy(coords_minus_h)

        gradient[i] = (e_plus_h - e_minus_h) / (2 * h)

    return gradient


def update_coordinates(coordinates: np.ndarray, gradient: np.ndarray, step: float):
    """
    Mise à jour des coordonnées atomiques en fonction du gradient calculé.

    Paramètres:
        coordinates (np.ndarray): Tableau des coordonnées atomiques actuelles.
        gradient (np.ndarray): Gradient de l'énergie.
        step (float): Pas de mise à jour.

    Retourne:
        np.ndarray: Nouvelles coordonnées après mise à jour.
    """
    norm_grad = np.linalg.norm(gradient)
    if norm_grad < 1e-10:
        return coordinates  # Évite la division par zéro
    
    # Normalisation du gradient
    direction = -gradient / norm_grad  
    return coordinates + step * direction

def minimize_energy_numerical(pdbfile: str):
    """
    Minimisation de l'énergie de la molécule par descente de gradient avec enregistrement dans un fichier.
    
    Paramètres:
        pdbfile (str): Chemin vers le fichier PDB contenant la molécule.
    
    Retourne:
        dict: Historique des paramètres enregistrés à chaque itération.
    """
    coordinates = extract_coord_pdb(pdbfile)
    history = {"iterations": [], "len1": [], "len2": [], "angle": [], "e_bond": [], "e_angle": [], "e_pot": [], "grms": [], "coordinates": []}
    
    for iteration in range(MAX_ITERATIONS):
        len1, len2, angle = calc_len_angle(coordinates)
        Epot = calc_pot_energy(coordinates)
        e_bond = K_BOND * (len1 - I_EQ) ** 2 + K_BOND * (len2 - I_EQ) ** 2
        angle_eq_rad = TETA_EQ * math.pi / 180
        e_angle = K_ANGLE * (angle - angle_eq_rad) ** 2
        
        gradient = calc_num_gradient(coordinates, H)
        grms = np.linalg.norm(gradient) / np.sqrt(len(coordinates))
        coordinates = update_coordinates(coordinates, gradient, STEP)
        
        history["iterations"].append(iteration)
        history["len1"].append(len1)
        history["len2"].append(len2)
        history["angle"].append(angle * 180 / math.pi)
        history["e_bond"].append(e_bond)
        history["e_angle"].append(e_angle)
        history["e_pot"].append(Epot)
        history["grms"].append(grms)
        history["coordinates"].append(coordinates.tolist())
        
        if grms < THRESHOLD:
            break
    else:
        print("\nNombre maximal d'itérations atteint sans convergence.\n")
    
    print("\n--- Après Minimisation ---\n")
    print(f"Nombre d'itérations : {iteration + 1}")
    print(f"Énergie potentielle : {Epot:.2e} kcal/mol")
    print(f"GRMS : {grms:.5e} kcal/(mol·Å)")
    print(f"Angle après minimisation : {angle * 180 / np.pi:.4f} degrés")
    print(f"Longueur de liaison l1  : {len1:.4f} Å")
    print(f"Longueur de liaison l2  : {len2:.4f} Å")

    plot_energies(history)
    save_history(history)
    write_pdb(coordinates, "minimized.pdb")
    
    return history  

def save_history(history: dict, filename: str = "history.txt"):
    """
    Sauvegarde de l'historique de minimisation, incluant le GRMS.
    
    Paramètres:
        history (dict): Dictionnaire contenant les données d'historique.
        filename (str): Nom du fichier de sortie.
    """
    save_path = os.path.join(SAVE_DIR, filename)
    os.makedirs(SAVE_DIR, exist_ok=True)
    
    with open(save_path, "w") as f:
        f.write("Iter\tlen1(Å)\tlen2(Å)\tAngle(°)\tE_bond\tE_angle\tE_total\tGRMS\tO_x\tO_y\tO_z\tH1_x\tH1_y\tH1_z\tH2_x\tH2_y\tH2_z\n")
        for i in range(len(history["iterations"])):
            f.write(
                f"{history['iterations'][i]}\t{history['len1'][i]:.4f}\t{history['len2'][i]:.4f}\t"
                f"{history['angle'][i]:.4f}\t{history['e_bond'][i]:.4e}\t{history['e_angle'][i]:.4e}\t"
                f"{history['e_pot'][i]:.4e}\t{history['grms'][i]:.4e}\t"
                f"{history['coordinates'][i][0]:.6f}\t{history['coordinates'][i][1]:.6f}\t{history['coordinates'][i][2]:.6f}\t"
                f"{history['coordinates'][i][3]:.6f}\t{history['coordinates'][i][4]:.6f}\t{history['coordinates'][i][5]:.6f}\t"
                f"{history['coordinates'][i][6]:.6f}\t{history['coordinates'][i][7]:.6f}\t{history['coordinates'][i][8]:.6f}\n"
            )
    
    print("\n--- Sauvegarde des résultats ---")
    print(f"\nHistorique sauvegardé dans : {save_path}")

    
def plot_energies(history: dict):
    """
    Tracé de l'énergie potentielle, de l'énergie de liaison et de l'énergie d'angle
    en fonction des itérations de la descente de gradient.

    Paramètres:
        history (dict): Dictionnaire contenant l'historique des énergies et itérations.
    """
    save_path = os.path.join(SAVE_DIR, "energy_plot.png")
    os.makedirs(SAVE_DIR, exist_ok=True)

    iterations = history["iterations"]
    energies_pot = history["e_pot"]
    energies_bond = history["e_bond"]
    energies_angle = history["e_angle"]

    plt.figure(figsize=(8, 6))
    plt.plot(iterations, energies_pot, label="Énergie potentielle", linestyle="-", linewidth=2)
    plt.plot(iterations, energies_bond, label="Énergie de liaison", linestyle="--", linewidth=2)
    plt.plot(iterations, energies_angle, label="Énergie d'angle", linestyle="-.", linewidth=2)

    plt.xlabel("Itérations")
    plt.ylabel("Énergie (kcal/mol)")
    plt.title("Évolution des énergies au cours de la minimisation")
    plt.legend()
    plt.grid()

    plt.savefig(save_path, dpi=300)
    print(f"Plot des énergies enregistré sous : {save_path}")

def write_pdb(coordinates: np.ndarray, filename=None, file_handle=None):
    """
    Écriture des coordonnées minimisées dans un fichier PDB formaté correctement.

    Paramètres:
        coordinates (np.ndarray): Tableau contenant les coordonnées (x, y, z) des atomes (shape=(9,)).
        filename (str, optionnel): Nom du fichier PDB à sauvegarder si `file_handle` n'est pas utilisé.
        file_handle (file object, optionnel): Fichier ouvert où écrire les données.
    """
    atom_names = ["O", "H1", "H2"]
    res_name = "H2O"
    chain_id = "A"
    res_seq = 1
    occupancy = 1.00
    temp_factor = 0.00
    element_symbols = ["O", "H", "H"]

    if file_handle is None:
        if filename is None:
            raise ValueError("write_pdb() nécessite soit un `filename`, soit `file_handle`.")
        save_path = os.path.join(SAVE_DIR, filename)
        os.makedirs(SAVE_DIR, exist_ok=True)
        file_handle = open(save_path, "w")
        file_handle.write("REMARK    Structure minimisée\n")

    for j in range(3):  
        x, y, z = coordinates[j * 3: (j + 1) * 3]
        file_handle.write(
            f"ATOM  {j+1:5d} {atom_names[j]:<4} {res_name:<4}{chain_id:1}{res_seq:4d}    "
            f"{x:8.3f}{y:8.3f}{z:8.3f}{occupancy:6.2f}{temp_factor:6.2f}          {element_symbols[j]:>2}\n"
        )

    if filename is not None:
        file_handle.write("END\n")
        file_handle.close()
        print(f"Fichier PDB généré : {filename}")


def generate_vmd_script():
    """
    Génère un script VMD pour afficher uniquement la molécule minimisée en CPK,
    et l'enregistre dans le dossier 'results'.
    """
    save_path = os.path.join(SAVE_DIR, "visualize.vmd")
    os.makedirs(SAVE_DIR, exist_ok=True)

    script_content = """\
# Chargement de la molécule après minimisation
mol new minimized.pdb
mol delrep 0 0
mol representation CPK
mol color Name
mol addrep 0

# Ajustement de la vue
display resetview
axes location off
menu graphics on
menu display off
"""

    with open(save_path, "w") as vmd_file:
        vmd_file.write(script_content)

    print(f"Script VMD généré : {save_path}")


def plot_angle(ax, origin, angle, size=50, unit="points", **kwargs):
    """Ajoute une annotation d'angle à un graphique 3D."""
    x, y, z = origin
    angle_text = f"{angle:.1f}°"
    ax.text(x, y, z, angle_text, fontsize=12, ha='center', va='center', **kwargs)

def plot_molecule_3d(pdbfile: str, history: dict):
    """
    Trace la molécule en 3D avant et après minimisation en récupérant les valeurs depuis `history`.

    Paramètres:
        pdbfile (str): Fichier PDB initial.
        history (dict): Dictionnaire contenant les valeurs avant et après minimisation.
    """
    save_path = os.path.join(SAVE_DIR, "3D_plot.png")

    # Récupérer les coordonnées initiales et finales depuis l'historique
    initial_coords = np.array(history["coordinates"][0]).reshape(3, 3)
    final_coords = np.array(history["coordinates"][-1]).reshape(3, 3)

    # Récupération des valeurs AVANT minimisation (dernière itération)
    len1_init = history["len1"][-1]
    len2_init = history["len2"][-1]
    angle_init = history["angle"][-1]

    # Récupération des valeurs APRÈS minimisation (première itération)
    len1_final = history["len1"][0]
    len2_final = history["len2"][0]
    angle_final = history["angle"][0]

    _, axes = plt.subplots(1, 2, subplot_kw={'projection': '3d'}, figsize=(12, 6))

    axes = axes[::-1]  # Inverser pour que le premier soit "Avant" et le second "Après"
    
    for ax, coords_current, title, len1, len2, angle in zip(
        axes, [initial_coords, final_coords], ["Après", "Avant"],  
        [len1_init, len1_final], [len2_init, len2_final], [angle_init, angle_final]
    ):
        ax.set_title(title)
        ax.scatter(coords_current[:, 0], coords_current[:, 1], coords_current[:, 2], c=['red', 'blue', 'blue'], s=200)
        ax.plot(coords_current[:2, 0], coords_current[:2, 1], coords_current[:2, 2], 'k--')
        ax.plot(coords_current[[0, 2], 0], coords_current[[0, 2], 1], coords_current[[0, 2], 2], 'k--')

        ax.text((coords_current[0, 0] + coords_current[1, 0]) / 2, 
                (coords_current[0, 1] + coords_current[1, 1]) / 2, 
                (coords_current[0, 2] + coords_current[1, 2]) / 2, 
                f"{len1:.4f} Å", color='black', fontsize=10)

        ax.text((coords_current[0, 0] + coords_current[2, 0]) / 2, 
                (coords_current[0, 1] + coords_current[2, 1]) / 2, 
                (coords_current[0, 2] + coords_current[2, 2]) / 2, 
                f"{len2:.4f} Å", color='black', fontsize=10)

        plot_angle(ax, coords_current[0], angle, size=50, unit="points", color='black')

        ax.set_xlabel('X')
        ax.set_ylabel('Y')
        ax.set_zlabel('Z')

    plt.savefig(save_path, dpi=300)
    print(f"Graphique enregistré sous : {save_path}")


def save_multipdb(initial_coords: np.ndarray, history: list, filename: str = "multi_minimized.pdb"):
    """
    Sauvegarde de l'évolution de la minimisation dans un fichier PDB unique avec plusieurs modèles,
    en incluant le premier modèle correspondant à l'état initial avant minimisation.

    Paramètres:
        initial_coords (np.ndarray): Coordonnées initiales avant minimisation.
        history (list of np.ndarray): Liste des coordonnées atomiques à chaque itération.
        filename (str): Nom du fichier PDB de sortie.
    """
    save_path = os.path.join(SAVE_DIR, filename)
    os.makedirs(SAVE_DIR, exist_ok=True)

    with open(save_path, "w") as pdb_file:
        pdb_file.write("REMARK    Multi-PDB des étapes de minimisation\n")

        # Ajouter l'état initial avant minimisation
        pdb_file.write("MODEL 000\n")  
        write_pdb(initial_coords, file_handle=pdb_file)  
        pdb_file.write("TER\nENDMDL\n")  

        # Ajouter les étapes de minimisation
        for i, coords in enumerate(history):
            pdb_file.write(f"MODEL {i+1:03d}\n")  
            write_pdb(coords, file_handle=pdb_file)  
            pdb_file.write("TER\nENDMDL\n")  

        pdb_file.write("END\n")

    print(f"Fichier Multi-PDB généré : {save_path}")


# === MAIN ===
if __name__ == "__main__":
    if len(sys.argv) < 2:
        print("Utilisation : python script.py <fichier_pdb> [threshold]")
        sys.exit(1)

    pdbfile = sys.argv[1]

    # Vérifier si un second argument est passé pour threshold
    if len(sys.argv) > 2:
        try:
            THRESHOLD = float(sys.argv[2])  
            print(f"\nSeuil GRMS défini à : {THRESHOLD} kcal/mol/A")
        except ValueError:
            print("Le seuil threshold doit être un nombre valide.")
            sys.exit(1)
    else : 
        print(f"\nSeuil GRMS par défaut : {THRESHOLD} kcal/mol/A")
    
    pdb_filename = os.path.splitext(os.path.basename(pdbfile))[0] 
    SAVE_DIR = f"results_{pdb_filename}"  
    os.makedirs(SAVE_DIR, exist_ok=True) 
    
    coordinates = extract_coord_pdb(pdbfile)
    print("\n--- Avant Minimisation ---")
    print(f"\nCoordonnées extraites :\n{coordinates.reshape(3, 3)}")

    len1, len2, angle = calc_len_angle(coordinates)
    print(f"\n Longueur O-H1 : {len1:.4f} Å")
    print(f" Longueur O-H2 : {len2:.4f} Å")
    print(f" Angle H1-O-H2 : {angle * 180 / math.pi:.4f} degrés")

    e_pot = calc_pot_energy(coordinates)
    print(f" Énergie potentielle : {e_pot:.4e} kcal/mol")
    grad = calc_num_gradient(coordinates, H)
    print(f"\n Gradient numérique : \n{grad.reshape(3, 3)}")

    # === Sauvegarde des résultats  ===
    history = minimize_energy_numerical(pdbfile)  
    generate_vmd_script()
    plot_molecule_3d(pdbfile, history)
    save_multipdb(coordinates, history["coordinates"])
    generate_vmd_movie(SAVE_DIR)
    
    # === End ===
    print("\nMinimisation et génération des fichiers terminées avec succès !")