From 5dcba3817b6f1c83bafde1d48124c44ec10f8dda Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Fri, 7 Jun 2024 16:34:20 -0400
Subject: [PATCH 01/23] bond angles dihedrals backbone code

---
 .../bond_angles_dihedrals_tool.py             | 249 ++++++++++++++++++
 mdagent/utils/__init__.py                     |  12 +-
 mdagent/utils/data_handling.py                | 109 ++++++++
 3 files changed, 369 insertions(+), 1 deletion(-)
 create mode 100644 mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
 create mode 100644 mdagent/utils/data_handling.py

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
new file mode 100644
index 00000000..abcb2fbc
--- /dev/null
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -0,0 +1,249 @@
+import mdtraj as md
+from langchain.tools import BaseTool
+
+from mdagent.utils import PathRegistry, load_single_traj
+
+# def load_single_traj(path_registry, traj_file, top_file=None):
+#     if top_file is not None:
+#         traj = md.load(traj_file, top=top_file)
+#     else:
+#         traj = md.load(traj_file)
+#     return traj if traj else None
+
+
+class ComputeAngles(BaseTool):
+    name = "compute_angles"
+    description = """Calculate the bond angles for the given sets of three atoms in
+    each snapshot of a molecular simulation."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, angle_indices, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+
+        if (
+            not angle_indices
+            or not isinstance(angle_indices, list)
+            or not all(len(indices) == 3 for indices in angle_indices)
+        ):
+            return (
+                "Invalid angle_indices. It should be a list of tuples, each "
+                "containing three atom indices."
+            )
+
+        return md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, angle_indices, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputeDihedrals(BaseTool):
+    name = "compute_dihedrals"
+    description = """Calculate the dihedral angles for the given groups of four atoms
+    in each snapshot of a molecular simulation."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, indices, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+
+        if (
+            not indices
+            or not isinstance(indices, list)
+            or not all(
+                isinstance(tup, tuple) and all(isinstance(i, int) for i in tup)
+                for tup in indices
+            )
+        ):
+            return (
+                "Invalid indices. It should be a list of tuples, each containing"
+                "atom indices as integers."
+            )
+
+        # Assuming a generic computation method for demonstration
+        # Replace `md.compute_properties` with the actual computation method you need
+
+        return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, indices, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputePsi(BaseTool):
+    name = "compute_psi"
+    description = """Calculate the psi angles for each snapshot in a molecular dynamics
+      simulation. These angles involve the alpha and carbonyl carbons, allowing rotation
+        that shapes the protein's secondary structure and overall 3D conformation.."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_psi(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputePhi(BaseTool):
+    name = "compute_phi"
+    description = """This class is designed to calculate the phi torsion angles in a
+    molecular dynamics simulation.
+    Involve the nitrogen and alpha carbon, allowing for rotation that contributes to
+    the overall folding pattern
+     More flexible, allowing for a wide range of conformations that define the protein’s
+     three-dimensional structure ."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_phi(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputeChi1(BaseTool):
+    name = "compute_chi1"
+    description = """Calculate the chi1 angles (the first side chain torsion angle
+    formed between four atoms around the CA-CB axis) for each snapshot in a molecular
+      dynamics simulation."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_chi1(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class Compute_Chi2(BaseTool):
+    name = "compute_chi2"
+    description = """Calculate the chi2 angles (the second side chain torsion angle
+      formed between four atoms around the CB-CG axis) for each snapshot in a
+      molecular dynamics simulation."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_chi2(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputeChi3(BaseTool):
+    name = "compute_chi3"
+    description = """Calculate the chi3 angles (the third side chain torsion angle
+    formed between four atoms around the CG-CD axis) for each snapshot in a molecular
+    dynamics simulation. Note: Only the residues ARG, GLN, GLU, LYS, and MET have
+    these atoms."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_chi3(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputeChi4(BaseTool):
+    name = "compute_chi4"
+    description = """Calculate the chi4 angles (the fourth side chain torsion angle
+      formed between four atoms around the CD-CE or CD-NE axis) for each snapshot in
+      a molecular dynamics simulation. Note: Only the residues ARG and LYS have these
+        atoms."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_chi4(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+class ComputeOmega(BaseTool):
+    name = "compute_omega"
+    description = """Calculate the omega angles (a specific type of bond angle) for
+    each snapshot in a molecular dynamics simulation.
+    Role: omega angles are primarily involved in the peptide bond, crucial for
+    determining the planarity
+      and rigidity of the peptide bond. They usually have less variability due to
+      the preference for trans configuration.
+     omega angles are more rigid due to the partial double-bond character of the
+     peptide bond, leading to limited rotational freedom
+     (mostly fixed at 180° or 0°)."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        traj = load_single_traj(self.path_registry, traj_file, top_file)
+        if not traj:
+            return "Trajectory could not be loaded."
+        return md.compute_omega(traj, periodic=True, opt=True)
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
+
+
+# class Ramachandran Plot
diff --git a/mdagent/utils/__init__.py b/mdagent/utils/__init__.py
index dbb247a4..fc3eaf7d 100644
--- a/mdagent/utils/__init__.py
+++ b/mdagent/utils/__init__.py
@@ -1,5 +1,15 @@
+from .data_handling import load_single_traj, load_traj_with_ref, save_plot, save_to_csv
 from .makellm import _make_llm
 from .path_registry import FileType, PathRegistry
 from .set_ckpt import SetCheckpoint
 
-__all__ = ["_make_llm", "PathRegistry", "FileType", "SetCheckpoint"]
+__all__ = [
+    "load_single_traj",
+    "load_traj_with_ref",
+    "save_plot",
+    "save_to_csv",
+    "_make_llm",
+    "FileType",
+    "PathRegistry",
+    "SetCheckpoint",
+]
diff --git a/mdagent/utils/data_handling.py b/mdagent/utils/data_handling.py
new file mode 100644
index 00000000..be111c0e
--- /dev/null
+++ b/mdagent/utils/data_handling.py
@@ -0,0 +1,109 @@
+import os
+import warnings
+
+import matplotlib.pyplot as plt
+import mdtraj as md
+import numpy as np
+
+from .path_registry import FileType, PathRegistry
+
+
+def load_single_traj(path_registry, top_fileid, traj_fileid=None, traj_required=False):
+    """
+    Load a single trajectory file using mdtraj. Check for file IDs in the path registry.
+    Parameters:
+    path_registry (PathRegistry): mapping file IDs to file paths.
+    top_fileid (str): File ID for the topology file.
+    traj_fileid (str, optional): File ID for the trajectory file.
+    traj_required (bool, optional): Whether the traj file is required. Default is False.
+    Returns:
+    mdtraj.Trajectory: Trajectory object.
+    """
+    if not isinstance(path_registry, PathRegistry):
+        raise ValueError("path_registry must be an instance of PathRegistry.")
+    all_fileids = path_registry.list_path_names()
+    if top_fileid not in all_fileids:
+        raise ValueError(f"Topology File ID '{top_fileid}' not found in PathRegistry")
+    top_path = path_registry.get_mapped_path(top_fileid)
+
+    if traj_fileid is None:
+        if not traj_required:
+            warnings.warn(
+                (
+                    "Trajectory File ID is not provided but is not required; "
+                    f"loading MDTrajectory from topology {top_fileid} only."
+                ),
+                UserWarning,
+            )
+            return md.load(top_path)
+        else:
+            raise ValueError("Trajectory File ID is required, and it's not provided.")
+
+    if traj_fileid not in all_fileids:
+        raise ValueError(
+            f"Trajectory File ID '{traj_fileid}' not found in PathRegistry."
+        )
+    traj_path = path_registry.get_mapped_path(traj_fileid)
+    return md.load(traj_path, top=top_path)
+
+
+def load_traj_with_ref(
+    path_registry, top_id, traj_id=None, ref_top_id=None, ref_traj_id=None
+):
+    traj = load_single_traj(path_registry, top_id, traj_id)
+    if ref_top_id is None:
+        ref_traj = traj
+    else:
+        ref_traj = load_single_traj(path_registry, ref_top_id, ref_traj_id)
+    return traj, ref_traj
+
+
+def save_to_csv(
+    path_registry, data_to_save, analysis_name, description=None, header=""
+):
+    """
+    Saves data to a csv file and maps the file ID to the file path in the path registry.
+    Parameters:
+    path_registry (PathRegistry): mapping file IDs to file paths.
+    data_to_save (np.ndarray): Data to save to a csv file.
+    analysis_name (str): Name of the analysis or data. This will be used as the file ID.
+    description (str, optional): Description of the data.
+    header (str, optional): Header for the csv file.
+    Returns:
+    str: File ID for the saved data.
+    """
+    if not isinstance(path_registry, PathRegistry):
+        raise ValueError("path_registry must be an instance of PathRegistry.")
+    if not isinstance(data_to_save, np.ndarray):
+        raise TypeError("data_to_save must be an instance of np.ndarray.")
+
+    base_path = f"{path_registry.ckpt_records}/{analysis_name}"
+    file_path = f"{base_path}.csv"
+    i = 0
+    while os.path.exists(file_path):
+        i += 1
+        file_path = f"{base_path}_{i}.csv"
+    file_id = analysis_name if i == 0 else f"{analysis_name}_{i}"
+    np.savetxt(file_path, data_to_save, delimiter=",", header=header)
+    path_registry.map_path(file_id, file_path, description=description)
+    print(f"Data saved to {file_path}")
+    return file_id
+
+
+def save_plot(path_registry, fig_analysis, description=None):
+    if not isinstance(path_registry, PathRegistry):
+        raise ValueError("path_registry must be an instance of PathRegistry.")
+    if plt.gcf().get_axes() == []:  # if there's no plot
+        raise ValueError("No plot detected. Failed to save.")
+
+    fig_name = path_registry.write_file_name(
+        type=FileType.FIGURE,
+        fig_analysis=fig_analysis,
+        file_format="png",
+    )
+    fig_id = path_registry.get_fileid(file_name=fig_name, type=FileType.FIGURE)
+    fig_path = f"{path_registry.ckpt_figures}/{fig_name}"
+    plt.savefig(fig_path)
+    path_registry.map_path(fig_id, fig_path, description=description)
+    print(f"Plot saved to {fig_path}")
+    return fig_id

From 520e16c001e9c406ce6fc255d09b6d8781ffa9df Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sat, 8 Jun 2024 11:19:38 -0400
Subject: [PATCH 02/23] path registry and code revisions

---
 mdagent/tools/base_tools/__init__.py          | 20 +++++++++++++++++++
 .../base_tools/analysis_tools/__init__.py     | 20 +++++++++++++++++++
 .../bond_angles_dihedrals_tool.py             |  2 +-
 mdagent/tools/maketools.py                    | 18 +++++++++++++++++
 4 files changed, 59 insertions(+), 1 deletion(-)

diff --git a/mdagent/tools/base_tools/__init__.py b/mdagent/tools/base_tools/__init__.py
index 0cf63290..a807acf7 100644
--- a/mdagent/tools/base_tools/__init__.py
+++ b/mdagent/tools/base_tools/__init__.py
@@ -1,3 +1,14 @@
+from .analysis_tools.bond_angles_dihedrals_tool import (
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
+)
 from .analysis_tools.plot_tools import SimulationOutputFigures
 from .analysis_tools.ppi_tools import PPIDistance
 from .analysis_tools.rdf_tool import RDFTool
@@ -22,6 +33,15 @@
 from .util_tools.search_tools import Scholar2ResultLLM
 
 __all__ = [
+    "ComputeAngles",
+    "ComputeChi1",
+    "ComputeChi2",
+    "ComputeChi3",
+    "ComputeChi4",
+    "ComputeDihedrals",
+    "ComputeOmega",
+    "ComputePhi",
+    "ComputePsi",
     "ListRegistryPaths",
     "MapPath2Name",
     "ProteinName2PDBTool",
diff --git a/mdagent/tools/base_tools/analysis_tools/__init__.py b/mdagent/tools/base_tools/analysis_tools/__init__.py
index f2193e08..ac00364f 100644
--- a/mdagent/tools/base_tools/analysis_tools/__init__.py
+++ b/mdagent/tools/base_tools/analysis_tools/__init__.py
@@ -1,3 +1,14 @@
+from .bond_angles_dihedrals_tool import (
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
+)
 from .plot_tools import SimulationOutputFigures
 from .ppi_tools import PPIDistance
 from .rgy import RadiusofGyrationAverage, RadiusofGyrationPerFrame, RadiusofGyrationPlot
@@ -5,6 +16,15 @@
 from .vis_tools import VisFunctions, VisualizeProtein
 
 __all__ = [
+    "ComputeAngles",
+    "ComputeChi1",
+    "ComputeChi2",
+    "ComputeChi3",
+    "ComputeChi4",
+    "ComputeDihedrals",
+    "ComputeOmega",
+    "ComputePhi",
+    "ComputePsi",
     "PPIDistance",
     "RMSDCalculator",
     "RadiusofGyrationPerFrame",
diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index abcb2fbc..c73d7535 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -150,7 +150,7 @@ async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-class Compute_Chi2(BaseTool):
+class ComputeChi2(BaseTool):
     name = "compute_chi2"
     description = """Calculate the chi2 angles (the second side chain torsion angle
       formed between four atoms around the CB-CG axis) for each snapshot in a
diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py
index 34b72404..1728dd0d 100644
--- a/mdagent/tools/maketools.py
+++ b/mdagent/tools/maketools.py
@@ -9,6 +9,15 @@
 
 from .base_tools import (
     CleaningToolFunction,
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
     ListRegistryPaths,
     ModifyBaseSimulationScriptTool,
     PackMolTool,
@@ -45,6 +54,15 @@ def make_all_tools(
 
     # add base tools
     base_tools = [
+        ComputeAngles(path_registry=path_instance),
+        ComputeChi1(path_registry=path_instance),
+        ComputeChi2(path_registry=path_instance),
+        ComputeChi3(path_registry=path_instance),
+        ComputeChi4(path_registry=path_instance),
+        ComputeDihedrals(path_registry=path_instance),
+        ComputeOmega(path_registry=path_instance),
+        ComputePhi(path_registry=path_instance),
+        ComputePsi(path_registry=path_instance),
         Scholar2ResultLLM(llm=llm, path_registry=path_instance),
         CleaningToolFunction(path_registry=path_instance),
         ListRegistryPaths(path_registry=path_instance),

From 3b1539f99eb6c0752df60fb59597edb68746af48 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sat, 8 Jun 2024 12:41:42 -0400
Subject: [PATCH 03/23] revised class descriptions for clarity

---
 .../bond_angles_dihedrals_tool.py             | 55 +++++++++----------
 1 file changed, 25 insertions(+), 30 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index c73d7535..548cb2f2 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -14,7 +14,8 @@
 class ComputeAngles(BaseTool):
     name = "compute_angles"
     description = """Calculate the bond angles for the given sets of three atoms in
-    each snapshot of a molecular simulation."""
+    each snapshot, and provide a list of indices specifying which atoms are involved
+    in each bond angle calculation.."""
 
     path_registry: PathRegistry | None = None
 
@@ -46,7 +47,8 @@ async def _arun(self, traj_file, angle_indices, top_file=None):
 class ComputeDihedrals(BaseTool):
     name = "compute_dihedrals"
     description = """Calculate the dihedral angles for the given groups of four atoms
-    in each snapshot of a molecular simulation."""
+      in each snapshot, and provide a list of dihedral angles along with a list of
+      indices specifying which atoms are involved in each dihedral angle calculation."""
 
     path_registry: PathRegistry | None = None
 
@@ -83,9 +85,9 @@ async def _arun(self, traj_file, indices, top_file=None):
 
 class ComputePsi(BaseTool):
     name = "compute_psi"
-    description = """Calculate the psi angles for each snapshot in a molecular dynamics
-      simulation. These angles involve the alpha and carbonyl carbons, allowing rotation
-        that shapes the protein's secondary structure and overall 3D conformation.."""
+    description = """Calculate the psi angles for each snapshot, providing a list of
+    psi angles for each frame in the trajectory and a list of indices specifying the
+    atoms involved in calculating each psi angle"""
 
     path_registry: PathRegistry | None = None
 
@@ -105,12 +107,8 @@ async def _arun(self, traj_file, top_file=None):
 
 class ComputePhi(BaseTool):
     name = "compute_phi"
-    description = """This class is designed to calculate the phi torsion angles in a
-    molecular dynamics simulation.
-    Involve the nitrogen and alpha carbon, allowing for rotation that contributes to
-    the overall folding pattern
-     More flexible, allowing for a wide range of conformations that define the protein’s
-     three-dimensional structure ."""
+    description = """This class calculates phi torsion angles and provides a list of phi
+      angles and indices specifying which atoms are involved in the calculations"""
 
     path_registry: PathRegistry | None = None
 
@@ -131,8 +129,9 @@ async def _arun(self, traj_file, top_file=None):
 class ComputeChi1(BaseTool):
     name = "compute_chi1"
     description = """Calculate the chi1 angles (the first side chain torsion angle
-    formed between four atoms around the CA-CB axis) for each snapshot in a molecular
-      dynamics simulation."""
+    formed between four atoms around the CA-CB axis) for each snapshot, providing a
+    list of chi1 angles and indices specifying the atoms involved in each chi1 angle
+    calculation."""
 
     path_registry: PathRegistry | None = None
 
@@ -153,8 +152,9 @@ async def _arun(self, traj_file, top_file=None):
 class ComputeChi2(BaseTool):
     name = "compute_chi2"
     description = """Calculate the chi2 angles (the second side chain torsion angle
-      formed between four atoms around the CB-CG axis) for each snapshot in a
-      molecular dynamics simulation."""
+    formed between four atoms around the CB-CG axis) for each snapshot, providing a
+    list of chi2 angles and a list of indices specifying the atoms involved in
+    calculating each chi2 angle."""
 
     path_registry: PathRegistry | None = None
 
@@ -175,9 +175,10 @@ async def _arun(self, traj_file, top_file=None):
 class ComputeChi3(BaseTool):
     name = "compute_chi3"
     description = """Calculate the chi3 angles (the third side chain torsion angle
-    formed between four atoms around the CG-CD axis) for each snapshot in a molecular
-    dynamics simulation. Note: Only the residues ARG, GLN, GLU, LYS, and MET have
-    these atoms."""
+    formed between four atoms around the CG-CD axis) for each snapshot in the
+    trajectory, providing a list of chi3 angles and indices specifying the atoms
+    involved in the calculation of each chi3 angle. Note: Only the residues ARG, GLN,
+      GLU, LYS, and MET have these atoms."""
 
     path_registry: PathRegistry | None = None
 
@@ -198,9 +199,9 @@ async def _arun(self, traj_file, top_file=None):
 class ComputeChi4(BaseTool):
     name = "compute_chi4"
     description = """Calculate the chi4 angles (the fourth side chain torsion angle
-      formed between four atoms around the CD-CE or CD-NE axis) for each snapshot in
-      a molecular dynamics simulation. Note: Only the residues ARG and LYS have these
-        atoms."""
+    formed between four atoms around the CD-CE or CD-NE axis) for each snapshot in the
+    trajectory, providing a list of indices specifying which atoms are involved in the
+    chi4 angle calculations. """
 
     path_registry: PathRegistry | None = None
 
@@ -220,15 +221,9 @@ async def _arun(self, traj_file, top_file=None):
 
 class ComputeOmega(BaseTool):
     name = "compute_omega"
-    description = """Calculate the omega angles (a specific type of bond angle) for
-    each snapshot in a molecular dynamics simulation.
-    Role: omega angles are primarily involved in the peptide bond, crucial for
-    determining the planarity
-      and rigidity of the peptide bond. They usually have less variability due to
-      the preference for trans configuration.
-     omega angles are more rigid due to the partial double-bond character of the
-     peptide bond, leading to limited rotational freedom
-     (mostly fixed at 180° or 0°)."""
+    description = """Calculate the omega angles for each snapshot in the trajectory,
+    providing a list of indices specifying which atoms are involved in the omega angle
+    calculations.."""
 
     path_registry: PathRegistry | None = None
 

From 06d5561169d350214e599a750c0e7edcab5dc567 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sat, 8 Jun 2024 20:10:09 -0400
Subject: [PATCH 04/23] unit test

---
 .../test_bond_angles_dihedrals.py             | 500 ++++++++++++++++++
 1 file changed, 500 insertions(+)
 create mode 100644 tests/test_analysis/test_bond_angles_dihedrals.py

diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
new file mode 100644
index 00000000..3c004daa
--- /dev/null
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -0,0 +1,500 @@
+from unittest.mock import MagicMock, patch
+
+import pytest
+
+from mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool import (
+    ComputeAngles,
+    ComputeChi1,
+    ComputeChi2,
+    ComputeChi3,
+    ComputeChi4,
+    ComputeDihedrals,
+    ComputeOmega,
+    ComputePhi,
+    ComputePsi,
+)
+
+
+@pytest.fixture
+def compute_angles_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeAngles(path_registry)
+
+
+@pytest.fixture
+def compute_dihedrals_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeDihedrals(path_registry)
+
+
+@pytest.fixture
+def compute_phi_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputePhi(path_registry)
+
+
+@pytest.fixture
+def compute_psi_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputePsi(path_registry)
+
+
+@pytest.fixture
+def compute_chi1_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeChi1(path_registry)
+
+
+@pytest.fixture
+def compute_chi2_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeChi2(path_registry)
+
+
+@pytest.fixture
+def compute_chi3_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeChi3(path_registry)
+
+
+@pytest.fixture
+def compute_chi4_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeChi4(path_registry)
+
+
+@pytest.fixture
+def compute_omega_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return ComputeOmega(path_registry)
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_angles")
+def test_run_success_compute_angles(
+    mock_compute_angles, mock_load_single_traj, compute_angles_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_angles
+    expected_angles = [0.1, 0.2, 0.3]
+    mock_compute_angles.return_value = expected_angles
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    angle_indices = [(0, 1, 2), (1, 2, 3)]
+    result = compute_angles_tool._run(traj_file, angle_indices, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_angles_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_angles.assert_called_once_with(
+        mock_traj, angle_indices, periodic=True, opt=True
+    )
+    assert result == expected_angles
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_angles(mock_load_single_traj, compute_angles_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    angle_indices = [(0, 1, 2), (1, 2, 3)]
+    result = compute_angles_tool._run(traj_file, angle_indices, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_angles_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_dihedrals")
+def test_run_success_compute_dihedrals(
+    mock_compute_dihedrals, mock_load_single_traj, compute_dihedrals_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_dihedrals
+    expected_dihedrals = [0.4, 0.5, 0.6]
+    mock_compute_dihedrals.return_value = expected_dihedrals
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    indices = [(0, 1, 2, 3), (1, 2, 3, 4)]
+    result = compute_dihedrals_tool._run(traj_file, indices, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_dihedrals_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_dihedrals.assert_called_once_with(
+        mock_traj, indices, periodic=True, opt=True
+    )
+    assert result == expected_dihedrals
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_dihedrals(mock_load_single_traj, compute_dihedrals_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    indices = [(0, 1, 2, 3), (1, 2, 3, 4)]
+    result = compute_dihedrals_tool._run(traj_file, indices, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_dihedrals_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_phi")
+def test_run_success_compute_phi(
+    mock_compute_phi, mock_load_single_traj, compute_phi_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_phi
+    expected_phi = [0.7, 0.8, 0.9]
+    mock_compute_phi.return_value = expected_phi
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_phi_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_phi_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_phi
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_phi(mock_load_single_traj, compute_phi_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_phi_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_phi_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_psi")
+def test_run_success_compute_psi(
+    mock_compute_psi, mock_load_single_traj, compute_psi_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_psi
+    expected_psi = [1.0, 1.1, 1.2]
+    mock_compute_psi.return_value = expected_psi
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_psi_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_psi_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_psi.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_psi
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_psi(mock_load_single_traj, compute_psi_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_psi_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_psi_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_chi1")
+def test_run_success_compute_chi1(
+    mock_compute_chi1, mock_load_single_traj, compute_chi1_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_chi1
+    expected_chi1 = [1.3, 1.4, 1.5]
+    mock_compute_chi1.return_value = expected_chi1
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi1_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi1_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_chi1.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_chi1
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_chi1(mock_load_single_traj, compute_chi1_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi1_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi1_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_chi2")
+def test_run_success_compute_chi2(
+    mock_compute_chi2, mock_load_single_traj, compute_chi2_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_chi2
+    expected_chi2 = [1.6, 1.7, 1.8]
+    mock_compute_chi2.return_value = expected_chi2
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi2_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi2_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_chi2.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_chi2
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_chi2(mock_load_single_traj, compute_chi2_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi2_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi2_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_chi3")
+def test_run_success_compute_chi3(
+    mock_compute_chi3, mock_load_single_traj, compute_chi3_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_chi3
+    expected_chi3 = [1.9, 2.0, 2.1]
+    mock_compute_chi3.return_value = expected_chi3
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi3_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi3_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_chi3.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_chi3
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_chi3(mock_load_single_traj, compute_chi3_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi3_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi3_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_chi4")
+def test_run_success_compute_chi4(
+    mock_compute_chi4, mock_load_single_traj, compute_chi4_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_chi4
+    expected_chi4 = [2.2, 2.3, 2.4]
+    mock_compute_chi4.return_value = expected_chi4
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi4_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi4_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_chi4.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_chi4
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_chi4(mock_load_single_traj, compute_chi4_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_chi4_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_chi4_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+@patch("mdtraj.compute_omega")
+def test_run_success_compute_omega(
+    mock_compute_omega, mock_load_single_traj, compute_omega_tool
+):
+    # Create a mock trajectory
+    mock_traj = MagicMock()
+    mock_load_single_traj.return_value = mock_traj
+
+    # Define the expected output from compute_omega
+    expected_omega = [2.5, 2.6, 2.7]
+    mock_compute_omega.return_value = expected_omega
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_omega_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_omega_tool.path_registry, traj_file, top_file
+    )
+    mock_compute_omega.assert_called_once_with(mock_traj, periodic=True, opt=True)
+    assert result == expected_omega
+
+
+@patch(
+    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+)
+def test_run_fail_compute_omega(mock_load_single_traj, compute_omega_tool):
+    # Simulate the trajectory loading failure
+    mock_load_single_traj.return_value = None
+
+    # Call the _run method
+    traj_file = "rec0_butane_123456"
+    top_file = "top_sim0_butane_123456"
+    result = compute_omega_tool._run(traj_file, top_file)
+
+    # Assertions
+    mock_load_single_traj.assert_called_once_with(
+        compute_omega_tool.path_registry, traj_file, top_file
+    )
+    assert result == "Trajectory could not be loaded."

From 64e64b99f0d259fa1e03c5e8a7948c68625bdd1f Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sat, 8 Jun 2024 20:36:13 -0400
Subject: [PATCH 05/23] minor changes to the bond_angle file

---
 .../analysis_tools/bond_angles_dihedrals_tool.py           | 7 -------
 1 file changed, 7 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 548cb2f2..4d5324c3 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -3,13 +3,6 @@
 
 from mdagent.utils import PathRegistry, load_single_traj
 
-# def load_single_traj(path_registry, traj_file, top_file=None):
-#     if top_file is not None:
-#         traj = md.load(traj_file, top=top_file)
-#     else:
-#         traj = md.load(traj_file)
-#     return traj if traj else None
-
 
 class ComputeAngles(BaseTool):
     name = "compute_angles"

From a142cf92242684dcf0dd7d3bf99ad773bbd65765 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sat, 8 Jun 2024 20:39:59 -0400
Subject: [PATCH 06/23] minor changes to the bond_angle file

---
 .../base_tools/analysis_tools/bond_angles_dihedrals_tool.py    | 3 ---
 1 file changed, 3 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 4d5324c3..179b8c6c 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -67,9 +67,6 @@ def _run(self, traj_file, indices, top_file=None):
                 "atom indices as integers."
             )
 
-        # Assuming a generic computation method for demonstration
-        # Replace `md.compute_properties` with the actual computation method you need
-
         return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
 
     async def _arun(self, traj_file, indices, top_file=None):

From 5019578af1828062428d54ae5ca42dd4ef921be4 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sun, 9 Jun 2024 11:30:04 -0400
Subject: [PATCH 07/23] added try except to every run function

---
 .../bond_angles_dihedrals_tool.py             | 158 +++++++++++-------
 1 file changed, 97 insertions(+), 61 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 179b8c6c..322b4866 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -17,21 +17,25 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, angle_indices, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-
-        if (
-            not angle_indices
-            or not isinstance(angle_indices, list)
-            or not all(len(indices) == 3 for indices in angle_indices)
-        ):
-            return (
-                "Invalid angle_indices. It should be a list of tuples, each "
-                "containing three atom indices."
-            )
-
-        return md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+
+            if (
+                not angle_indices
+                or not isinstance(angle_indices, list)
+                or not all(len(indices) == 3 for indices in angle_indices)
+            ):
+                return (
+                    "Invalid angle_indices. It should be a list of tuples, each "
+                    "containing three atom indices."
+                )
+
+            return md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, angle_indices, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -50,24 +54,28 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, indices, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-
-        if (
-            not indices
-            or not isinstance(indices, list)
-            or not all(
-                isinstance(tup, tuple) and all(isinstance(i, int) for i in tup)
-                for tup in indices
-            )
-        ):
-            return (
-                "Invalid indices. It should be a list of tuples, each containing"
-                "atom indices as integers."
-            )
-
-        return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+
+            if (
+                not indices
+                or not isinstance(indices, list)
+                or not all(
+                    isinstance(tup, tuple) and all(isinstance(i, int) for i in tup)
+                    for tup in indices
+                )
+            ):
+                return (
+                    "Invalid indices. It should be a list of tuples, each containing"
+                    "atom indices as integers."
+                )
+
+            return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, indices, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -86,10 +94,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_psi(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_psi(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -107,10 +119,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_phi(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_phi(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -130,10 +146,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_chi1(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_chi1(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -153,10 +173,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_chi2(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_chi2(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -177,10 +201,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_chi3(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_chi3(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -200,10 +228,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_chi4(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_chi4(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
@@ -222,10 +254,14 @@ def __init__(self, path_registry: PathRegistry):
         self.path_registry = path_registry
 
     def _run(self, traj_file, top_file=None):
-        traj = load_single_traj(self.path_registry, traj_file, top_file)
-        if not traj:
-            return "Trajectory could not be loaded."
-        return md.compute_omega(traj, periodic=True, opt=True)
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Trajectory could not be loaded."
+            return md.compute_omega(traj, periodic=True, opt=True)
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
 
     async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")

From 327925ddfdcc3207e8d2b60cb9167356fdf38a50 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Thu, 13 Jun 2024 09:33:49 -0400
Subject: [PATCH 08/23] a hopeful small correction to pass the git check

---
 mdagent/tools/maketools.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py
index 4d6799cf..ea4a62a4 100644
--- a/mdagent/tools/maketools.py
+++ b/mdagent/tools/maketools.py
@@ -15,10 +15,10 @@
     ComputeChi3,
     ComputeChi4,
     ComputeDihedrals,
+    ComputeLPRMSD,
     ComputeOmega,
     ComputePhi,
     ComputePsi,
-    ComputeLPRMSD,
     ComputeRMSD,
     ComputeRMSF,
     ContactsTool,

From 847f764fce8bc716b81f249e1120735b0d4d2909 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Thu, 13 Jun 2024 20:49:13 -0400
Subject: [PATCH 09/23]  updated the toools to include Ram plot.

---
 mdagent/tools/base_tools/__init__.py          |   2 +
 .../base_tools/analysis_tools/__init__.py     |   2 +
 .../bond_angles_dihedrals_tool.py             | 115 +++--
 mdagent/tools/maketools.py                    |   2 +
 .../test_bond_angles_dihedrals.py             | 416 ++----------------
 5 files changed, 125 insertions(+), 412 deletions(-)

diff --git a/mdagent/tools/base_tools/__init__.py b/mdagent/tools/base_tools/__init__.py
index 8dbfc292..a2a6c6df 100644
--- a/mdagent/tools/base_tools/__init__.py
+++ b/mdagent/tools/base_tools/__init__.py
@@ -8,6 +8,7 @@
     ComputeOmega,
     ComputePhi,
     ComputePsi,
+    RamachandranPlot,
 )
 from .analysis_tools.distance_tools import ContactsTool, DistanceMatrixTool
 from .analysis_tools.inertia import MomentOfInertia
@@ -66,6 +67,7 @@
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",
+    "RamachandranPlot",
     "RDFTool",
     "RMSDCalculator",
     "Scholar2ResultLLM",
diff --git a/mdagent/tools/base_tools/analysis_tools/__init__.py b/mdagent/tools/base_tools/analysis_tools/__init__.py
index 7bfb3e54..f083cb8f 100644
--- a/mdagent/tools/base_tools/analysis_tools/__init__.py
+++ b/mdagent/tools/base_tools/analysis_tools/__init__.py
@@ -8,6 +8,7 @@
     ComputeOmega,
     ComputePhi,
     ComputePsi,
+    RamachandranPlot,
 )
 from .distance_tools import ContactsTool, DistanceMatrixTool
 from .inertia import MomentOfInertia
@@ -40,6 +41,7 @@
     "RadiusofGyrationAverage",
     "RadiusofGyrationPerFrame",
     "RadiusofGyrationPlot",
+    "RamachandranPlot",
     "RMSDCalculator",
     "SimulationOutputFigures",
     "SolventAccessibleSurfaceArea",
diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 322b4866..3ee48300 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -1,3 +1,4 @@
+import matplotlib.pyplot as plt
 import mdtraj as md
 from langchain.tools import BaseTool
 
@@ -20,7 +21,7 @@ def _run(self, traj_file, angle_indices, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
+                return "Failed.Trajectory could not be loaded."
 
             if (
                 not angle_indices
@@ -28,11 +29,12 @@ def _run(self, traj_file, angle_indices, top_file=None):
                 or not all(len(indices) == 3 for indices in angle_indices)
             ):
                 return (
-                    "Invalid angle_indices. It should be a list of tuples, each "
-                    "containing three atom indices."
+                    "Failed. Invalid angle_indices. It should be a list of tuples, "
+                    "each containing three atom indices."
                 )
 
-            return md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+            result = md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -57,7 +59,7 @@ def _run(self, traj_file, indices, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
+                return " Failed. Trajectory could not be loaded."
 
             if (
                 not indices
@@ -67,12 +69,10 @@ def _run(self, traj_file, indices, top_file=None):
                     for tup in indices
                 )
             ):
-                return (
-                    "Invalid indices. It should be a list of tuples, each containing"
-                    "atom indices as integers."
-                )
+                return "Failed. Invalid indices. It should be a list of tuples."
 
-            return md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+            result = md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -81,11 +81,10 @@ async def _arun(self, traj_file, indices, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-class ComputePsi(BaseTool):
-    name = "compute_psi"
-    description = """Calculate the psi angles for each snapshot, providing a list of
-    psi angles for each frame in the trajectory and a list of indices specifying the
-    atoms involved in calculating each psi angle"""
+class ComputePhi(BaseTool):
+    name = "compute_phi"
+    description = """This class calculates phi torsion angles and provides a list of phi
+      angles and indices specifying which atoms are involved in the calculations"""
 
     path_registry: PathRegistry | None = None
 
@@ -97,8 +96,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_psi(traj, periodic=True, opt=True)
+                return " Failed. Trajectory could not be loaded."
+
+            result = md.compute_phi(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -107,10 +108,11 @@ async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-class ComputePhi(BaseTool):
-    name = "compute_phi"
-    description = """This class calculates phi torsion angles and provides a list of phi
-      angles and indices specifying which atoms are involved in the calculations"""
+class ComputePsi(BaseTool):
+    name = "compute_psi"
+    description = """Calculate the psi angles for each snapshot, providing a list of
+    psi angles for each frame in the trajectory and a list of indices specifying the
+    atoms involved in calculating each psi angle"""
 
     path_registry: PathRegistry | None = None
 
@@ -122,8 +124,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_phi(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_psi(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -149,8 +153,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi1(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi1(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -176,8 +182,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi2(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi2(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -204,8 +212,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi3(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi3(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -231,8 +241,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_chi4(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_chi4(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -257,8 +269,10 @@ def _run(self, traj_file, top_file=None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Trajectory could not be loaded."
-            return md.compute_omega(traj, periodic=True, opt=True)
+                return "Failed. Trajectory could not be loaded."
+
+            result = md.compute_omega(traj, periodic=True, opt=True)
+            return f"Succeeded. {result}"
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -267,4 +281,37 @@ async def _arun(self, traj_file, top_file=None):
         raise NotImplementedError("Async version not implemented")
 
 
-# class Ramachandran Plot
+class RamachandranPlot(BaseTool):
+    name = "ramachandran_plot"
+    description = """Generate a Ramachandran plot for the given trajectory, showing
+    the distribution of phi and psi angles for each frame."""
+
+    path_registry: PathRegistry | None = None
+
+    def __init__(self, path_registry: PathRegistry):
+        super().__init__()
+        self.path_registry = path_registry
+
+    def _run(self, traj_file, top_file=None):
+        try:
+            traj = load_single_traj(self.path_registry, traj_file, top_file)
+            if not traj:
+                return "Failed. Trajectory could not be loaded."
+
+            phi_indices, phi_angles = md.compute_phi(traj, periodic=True, opt=True)
+            psi_indices, psi_angles = md.compute_psi(traj, periodic=True, opt=True)
+
+            plt.figure(figsize=(10, 8))
+            plt.scatter(phi_angles.flatten(), psi_angles.flatten(), s=1, color="blue")
+            plt.xlabel("Phi Angles (radians)")
+            plt.ylabel("Psi Angles (radians)")
+            plt.title("Ramachandran Plot")
+            plt.grid(True)
+            plt.show()
+            return "Succeeded. Ramachandran plot generated."
+
+        except Exception as e:
+            return f"Failed. {type(e).__name__}: {e}"
+
+    async def _arun(self, traj_file, top_file=None):
+        raise NotImplementedError("Async version not implemented")
diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py
index ea4a62a4..df5b597b 100644
--- a/mdagent/tools/maketools.py
+++ b/mdagent/tools/maketools.py
@@ -33,6 +33,7 @@
     RadiusofGyrationAverage,
     RadiusofGyrationPerFrame,
     RadiusofGyrationPlot,
+    RamachandranPlot,
     RDFTool,
     Scholar2ResultLLM,
     SetUpandRunFunction,
@@ -86,6 +87,7 @@ def make_all_tools(
         RadiusofGyrationAverage(path_registry=path_instance),
         RadiusofGyrationPerFrame(path_registry=path_instance),
         RadiusofGyrationPlot(path_registry=path_instance),
+        RamachandranPlot(path_registry=path_instance),
         RDFTool(path_registry=path_instance),
         SetUpandRunFunction(path_registry=path_instance),
         SimulationOutputFigures(path_registry=path_instance),
diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
index 3c004daa..8da00862 100644
--- a/tests/test_analysis/test_bond_angles_dihedrals.py
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -12,9 +12,19 @@
     ComputeOmega,
     ComputePhi,
     ComputePsi,
+    RamachandranPlot,
 )
 
 
+# Fixture to patch 'load_single_traj'
+@pytest.fixture
+def patched_load_single_traj():
+    with patch(
+        "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
+    ) as mock_load_single_traj:
+        yield mock_load_single_traj
+
+
 @pytest.fixture
 def compute_angles_tool(get_registry):
     path_registry = get_registry("raw", True)
@@ -69,16 +79,19 @@ def compute_omega_tool(get_registry):
     return ComputeOmega(path_registry)
 
 
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
+@pytest.fixture
+def ramachandran_plot_tool(get_registry):
+    path_registry = get_registry("raw", True)
+    return RamachandranPlot(path_registry)
+
+
 @patch("mdtraj.compute_angles")
 def test_run_success_compute_angles(
-    mock_compute_angles, mock_load_single_traj, compute_angles_tool
+    mock_compute_angles, patched_load_single_traj, compute_angles_tool
 ):
     # Create a mock trajectory
     mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
+    patched_load_single_traj.return_value = mock_traj
 
     # Define the expected output from compute_angles
     expected_angles = [0.1, 0.2, 0.3]
@@ -91,21 +104,18 @@ def test_run_success_compute_angles(
     result = compute_angles_tool._run(traj_file, angle_indices, top_file)
 
     # Assertions
-    mock_load_single_traj.assert_called_once_with(
+    patched_load_single_traj.assert_called_once_with(
         compute_angles_tool.path_registry, traj_file, top_file
     )
     mock_compute_angles.assert_called_once_with(
         mock_traj, angle_indices, periodic=True, opt=True
     )
-    assert result == expected_angles
+    assert result == f"Succeeded. {expected_angles}"
 
 
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_angles(mock_load_single_traj, compute_angles_tool):
+def test_run_fail_compute_angles(patched_load_single_traj, compute_angles_tool):
     # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
+    patched_load_single_traj.return_value = None
 
     # Call the _run method
     traj_file = "rec0_butane_123456"
@@ -114,387 +124,37 @@ def test_run_fail_compute_angles(mock_load_single_traj, compute_angles_tool):
     result = compute_angles_tool._run(traj_file, angle_indices, top_file)
 
     # Assertions
-    mock_load_single_traj.assert_called_once_with(
+    patched_load_single_traj.assert_called_once_with(
         compute_angles_tool.path_registry, traj_file, top_file
     )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-@patch("mdtraj.compute_dihedrals")
-def test_run_success_compute_dihedrals(
-    mock_compute_dihedrals, mock_load_single_traj, compute_dihedrals_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_dihedrals
-    expected_dihedrals = [0.4, 0.5, 0.6]
-    mock_compute_dihedrals.return_value = expected_dihedrals
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    indices = [(0, 1, 2, 3), (1, 2, 3, 4)]
-    result = compute_dihedrals_tool._run(traj_file, indices, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_dihedrals_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_dihedrals.assert_called_once_with(
-        mock_traj, indices, periodic=True, opt=True
-    )
-    assert result == expected_dihedrals
+    assert result == "Failed. Trajectory could not be loaded."
 
 
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_dihedrals(mock_load_single_traj, compute_dihedrals_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    indices = [(0, 1, 2, 3), (1, 2, 3, 4)]
-    result = compute_dihedrals_tool._run(traj_file, indices, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_dihedrals_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
+# Similar tests for other classes (ComputeChi1, ComputeChi2, etc.)
+# ...
 @patch("mdtraj.compute_phi")
-def test_run_success_compute_phi(
-    mock_compute_phi, mock_load_single_traj, compute_phi_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_phi
-    expected_phi = [0.7, 0.8, 0.9]
-    mock_compute_phi.return_value = expected_phi
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_phi_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_phi_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_phi
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_phi(mock_load_single_traj, compute_phi_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_phi_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_phi_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
 @patch("mdtraj.compute_psi")
-def test_run_success_compute_psi(
-    mock_compute_psi, mock_load_single_traj, compute_psi_tool
+def test_run_success_ramachandran_plot(
+    mock_compute_psi, mock_compute_phi, patched_load_single_traj, ramachandran_plot_tool
 ):
     # Create a mock trajectory
     mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
+    patched_load_single_traj.return_value = mock_traj
 
-    # Define the expected output from compute_psi
-    expected_psi = [1.0, 1.1, 1.2]
+    # Define the expected output from compute_phi and compute_psi
+    expected_phi = ([(0, 1, 2, 3)], [[0.7, 0.8, 0.9]])
+    expected_psi = ([(0, 1, 2, 3)], [[1.0, 1.1, 1.2]])
+    mock_compute_phi.return_value = expected_phi
     mock_compute_psi.return_value = expected_psi
 
     # Call the _run method
     traj_file = "rec0_butane_123456"
     top_file = "top_sim0_butane_123456"
-    result = compute_psi_tool._run(traj_file, top_file)
-
+    result = ramachandran_plot_tool._run(traj_file, top_file)
     # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_psi_tool.path_registry, traj_file, top_file
+    patched_load_single_traj.assert_called_once_with(
+        ramachandran_plot_tool.path_registry, traj_file, top_file
     )
+    mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True)
     mock_compute_psi.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_psi
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_psi(mock_load_single_traj, compute_psi_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_psi_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_psi_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-@patch("mdtraj.compute_chi1")
-def test_run_success_compute_chi1(
-    mock_compute_chi1, mock_load_single_traj, compute_chi1_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_chi1
-    expected_chi1 = [1.3, 1.4, 1.5]
-    mock_compute_chi1.return_value = expected_chi1
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi1_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi1_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_chi1.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_chi1
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_chi1(mock_load_single_traj, compute_chi1_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi1_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi1_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-@patch("mdtraj.compute_chi2")
-def test_run_success_compute_chi2(
-    mock_compute_chi2, mock_load_single_traj, compute_chi2_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_chi2
-    expected_chi2 = [1.6, 1.7, 1.8]
-    mock_compute_chi2.return_value = expected_chi2
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi2_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi2_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_chi2.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_chi2
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_chi2(mock_load_single_traj, compute_chi2_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi2_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi2_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-@patch("mdtraj.compute_chi3")
-def test_run_success_compute_chi3(
-    mock_compute_chi3, mock_load_single_traj, compute_chi3_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_chi3
-    expected_chi3 = [1.9, 2.0, 2.1]
-    mock_compute_chi3.return_value = expected_chi3
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi3_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi3_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_chi3.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_chi3
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_chi3(mock_load_single_traj, compute_chi3_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi3_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi3_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-@patch("mdtraj.compute_chi4")
-def test_run_success_compute_chi4(
-    mock_compute_chi4, mock_load_single_traj, compute_chi4_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_chi4
-    expected_chi4 = [2.2, 2.3, 2.4]
-    mock_compute_chi4.return_value = expected_chi4
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi4_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi4_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_chi4.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_chi4
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_chi4(mock_load_single_traj, compute_chi4_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_chi4_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_chi4_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-@patch("mdtraj.compute_omega")
-def test_run_success_compute_omega(
-    mock_compute_omega, mock_load_single_traj, compute_omega_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    mock_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_omega
-    expected_omega = [2.5, 2.6, 2.7]
-    mock_compute_omega.return_value = expected_omega
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_omega_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_omega_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_omega.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == expected_omega
-
-
-@patch(
-    "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-)
-def test_run_fail_compute_omega(mock_load_single_traj, compute_omega_tool):
-    # Simulate the trajectory loading failure
-    mock_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = compute_omega_tool._run(traj_file, top_file)
-
-    # Assertions
-    mock_load_single_traj.assert_called_once_with(
-        compute_omega_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Trajectory could not be loaded."
+    assert result == "Succeeded. Ramachandran plot generated."

From e732d230b0152c48cd48d670b0357c5998dd4fd4 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Sun, 16 Jun 2024 14:22:07 -0400
Subject: [PATCH 10/23]  added type hint, added graphs, save to path and file,
 added helper function

---
 .../bond_angles_dihedrals_tool.py             | 252 +++++++++++++++---
 .../test_bond_angles_dihedrals.py             |  38 ++-
 2 files changed, 242 insertions(+), 48 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 3ee48300..f13d0fc0 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -1,5 +1,6 @@
 import matplotlib.pyplot as plt
 import mdtraj as md
+import numpy as np
 from langchain.tools import BaseTool
 
 from mdagent.utils import PathRegistry, load_single_traj
@@ -17,11 +18,11 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, angle_indices, top_file=None):
+    def _run(self, traj_file: str, angle_indices: list, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return "Failed.Trajectory could not be loaded."
+                return "Failed. Trajectory could not be loaded."
 
             if (
                 not angle_indices
@@ -33,13 +34,25 @@ def _run(self, traj_file, angle_indices, top_file=None):
                     "each containing three atom indices."
                 )
 
-            result = md.compute_angles(traj, angle_indices, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            angles = md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+
+            # Check if path_registry is not None
+            if self.path_registry is not None:
+                plot_save_path = self.path_registry.get_mapped_path("angles_plot.png")
+                plot_angles(angles, title="Bond Angles", save_path=plot_save_path)
+                return "Succeeded. Bond angles computed, saved to file and plot saved."
+            else:
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, angle_indices, top_file=None):
+    async def _arun(
+        self,
+        traj_file: str,
+        angle_indices: list,
+        top_file: str | None = None,
+    ):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -55,11 +68,11 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, indices, top_file=None):
+    def _run(self, traj_file: str, indices: list, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return " Failed. Trajectory could not be loaded."
+                return "Failed. Trajectory could not be loaded."
 
             if (
                 not indices
@@ -71,13 +84,26 @@ def _run(self, traj_file, indices, top_file=None):
             ):
                 return "Failed. Invalid indices. It should be a list of tuples."
 
-            result = md.compute_dihedrals(traj, indices, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            dihedrals = md.compute_dihedrals(traj, indices, periodic=True, opt=True)
+
+            # Check if path_registry is not None
+            if self.path_registry is not None:
+                plot_save_path = self.path_registry.get_mapped_path(
+                    "dihedrals_plot.png"
+                )
+                plot_angles(
+                    dihedrals, title="Dihedral Angles", save_path=plot_save_path
+                )
+                return (
+                    "Succeeded. Dihedral angles computed, saved to file and plot saved."
+                )
+            else:
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, indices, top_file=None):
+    async def _arun(self, traj_file: str, indices: list, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -92,19 +118,33 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
-                return " Failed. Trajectory could not be loaded."
+                return "Failed. Trajectory could not be loaded."
+
+            indices, angles = md.compute_phi(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("phi_results.npz", indices, angles)
 
-            result = md.compute_phi(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("phi_plot.png")
+                plot_angles(angles, title="Phi Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. Phi angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -120,19 +160,33 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
                 return "Failed. Trajectory could not be loaded."
 
-            result = md.compute_psi(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            indices, angles = md.compute_psi(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("psi_results.npz", indices, angles)
+
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("psi_plot.png")
+                plot_angles(angles, title="Psi Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. Psi angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -149,14 +203,28 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
                 return "Failed. Trajectory could not be loaded."
 
-            result = md.compute_chi1(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            indices, angles = md.compute_chi1(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("chi1_results.npz", indices, angles)
+
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("chi1_plot.png")
+                plot_angles(angles, title="Chi1 Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. chi1 angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
@@ -178,19 +246,33 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
                 return "Failed. Trajectory could not be loaded."
 
-            result = md.compute_chi2(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            indices, angles = md.compute_chi2(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("chi2_results.npz", indices, angles)
+
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("chi2_plot.png")
+                plot_angles(angles, title="Chi2 Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. chi2 angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -208,19 +290,33 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
                 return "Failed. Trajectory could not be loaded."
 
-            result = md.compute_chi3(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            indices, angles = md.compute_chi3(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("chi3_results.npz", indices, angles)
+
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("chi3_plot.png")
+                plot_angles(angles, title="Chi3 Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. chi3 angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -237,19 +333,33 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
                 return "Failed. Trajectory could not be loaded."
 
-            result = md.compute_chi4(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            indices, angles = md.compute_chi4(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("chi4_results.npz", indices, angles)
+
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("chi4_plot.png")
+                plot_angles(angles, title="Chi4 Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. chi4 angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -265,19 +375,33 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
                 return "Failed. Trajectory could not be loaded."
 
-            result = md.compute_omega(traj, periodic=True, opt=True)
-            return f"Succeeded. {result}"
+            indices, angles = md.compute_omega(traj, periodic=True, opt=True)
+
+            # Check if path_registry is initialized
+            if self.path_registry is not None:
+                # Save results to a file
+                save_results_to_file("omega_results.npz", indices, angles)
+
+                # Generate and save a plot
+                plot_save_path = self.path_registry.get_mapped_path("omega_plot.png")
+                plot_angles(angles, title="Omega Angles", save_path=plot_save_path)
+
+                # Return success message
+                return "Succeeded. omega angles computed, saved to file and plot saved."
+            else:
+                # Return failure message if path_registry is not initialized
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
 
 
@@ -292,7 +416,7 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file, top_file=None):
+    def _run(self, traj_file: str, top_file: str | None = None):
         try:
             traj = load_single_traj(self.path_registry, traj_file, top_file)
             if not traj:
@@ -301,17 +425,61 @@ def _run(self, traj_file, top_file=None):
             phi_indices, phi_angles = md.compute_phi(traj, periodic=True, opt=True)
             psi_indices, psi_angles = md.compute_psi(traj, periodic=True, opt=True)
 
+            # Map indices to residues for further analysis or reporting
+            map_indices_to_residues(traj, phi_indices)
+            map_indices_to_residues(traj, psi_indices)
+
+            # can add further analysis or reporting here using phi_residues and
+            # psi_residues
             plt.figure(figsize=(10, 8))
             plt.scatter(phi_angles.flatten(), psi_angles.flatten(), s=1, color="blue")
             plt.xlabel("Phi Angles (radians)")
             plt.ylabel("Psi Angles (radians)")
             plt.title("Ramachandran Plot")
             plt.grid(True)
-            plt.show()
-            return "Succeeded. Ramachandran plot generated."
+
+            # Check if path_registry is not None
+            if self.path_registry is not None:
+                plot_save_path = self.path_registry.get_mapped_path(
+                    "ramachandran_plot.png"
+                )
+                plt.savefig(plot_save_path)
+                return "Succeeded. Ramachandran plot generated and saved to file."
+            else:
+                return "Failed. Path registry is not initialized."
 
         except Exception as e:
             return f"Failed. {type(e).__name__}: {e}"
 
-    async def _arun(self, traj_file, top_file=None):
+    async def _arun(self, traj_file: str, top_file: str | None = None):
         raise NotImplementedError("Async version not implemented")
+
+
+# Helper functions suggested by Jorge
+def map_indices_to_residues(traj, indices):
+    atom_to_residue = {atom.index: atom.residue for atom in traj.topology.atoms}
+    residues_per_angle = [
+        [atom_to_residue[idx] for idx in angle_set] for angle_set in indices
+    ]
+    return residues_per_angle
+
+
+def save_results_to_file(filename, indices, angles):
+    np.savez(filename, indices=indices, angles=angles)
+
+
+def plot_angles(angles, title="Angles", save_path=None):
+    print(f"Save path received: {save_path}")  # Debugging help
+    plt.figure(figsize=(10, 8))
+    for angle_set in angles.T:
+        plt.plot(angle_set, label="Angle")
+    plt.xlabel("Frame")
+    plt.ylabel("Angle (radians)")
+    plt.title(title)
+    plt.legend()
+    plt.grid(True)
+    if save_path:
+        print(f"Calling savefig with path: {save_path}")  # Debugging help
+        plt.savefig(save_path)
+    else:
+        plt.show()
diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
index 8da00862..7f25c647 100644
--- a/tests/test_analysis/test_bond_angles_dihedrals.py
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -1,5 +1,6 @@
 from unittest.mock import MagicMock, patch
 
+import numpy as np
 import pytest
 
 from mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool import (
@@ -86,17 +87,23 @@ def ramachandran_plot_tool(get_registry):
 
 
 @patch("mdtraj.compute_angles")
+@patch("matplotlib.pyplot.savefig")
 def test_run_success_compute_angles(
-    mock_compute_angles, patched_load_single_traj, compute_angles_tool
+    mock_savefig, mock_compute_angles, patched_load_single_traj, compute_angles_tool
 ):
     # Create a mock trajectory
     mock_traj = MagicMock()
     patched_load_single_traj.return_value = mock_traj
 
     # Define the expected output from compute_angles
-    expected_angles = [0.1, 0.2, 0.3]
+    expected_angles = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
     mock_compute_angles.return_value = expected_angles
 
+    # Mock the path registry get_mapped_path method
+    compute_angles_tool.path_registry.get_mapped_path = MagicMock(
+        return_value="angles_plot.png"
+    )
+
     # Call the _run method
     traj_file = "rec0_butane_123456"
     top_file = "top_sim0_butane_123456"
@@ -110,7 +117,11 @@ def test_run_success_compute_angles(
     mock_compute_angles.assert_called_once_with(
         mock_traj, angle_indices, periodic=True, opt=True
     )
-    assert result == f"Succeeded. {expected_angles}"
+    compute_angles_tool.path_registry.get_mapped_path.assert_called_once_with(
+        "angles_plot.png"
+    )
+    mock_savefig.assert_called_once_with("angles_plot.png")
+    assert result == "Succeeded. Bond angles computed, saved to file and plot saved."
 
 
 def test_run_fail_compute_angles(patched_load_single_traj, compute_angles_tool):
@@ -131,11 +142,16 @@ def test_run_fail_compute_angles(patched_load_single_traj, compute_angles_tool):
 
 
 # Similar tests for other classes (ComputeChi1, ComputeChi2, etc.)
-# ...
+
+
 @patch("mdtraj.compute_phi")
 @patch("mdtraj.compute_psi")
 def test_run_success_ramachandran_plot(
-    mock_compute_psi, mock_compute_phi, patched_load_single_traj, ramachandran_plot_tool
+    mock_savefig,
+    mock_compute_psi,
+    mock_compute_phi,
+    patched_load_single_traj,
+    ramachandran_plot_tool,
 ):
     # Create a mock trajectory
     mock_traj = MagicMock()
@@ -147,14 +163,24 @@ def test_run_success_ramachandran_plot(
     mock_compute_phi.return_value = expected_phi
     mock_compute_psi.return_value = expected_psi
 
+    # Mock the path registry get_mapped_path method
+    ramachandran_plot_tool.path_registry.get_mapped_path = MagicMock(
+        return_value="ramachandran_plot.png"
+    )
+
     # Call the _run method
     traj_file = "rec0_butane_123456"
     top_file = "top_sim0_butane_123456"
     result = ramachandran_plot_tool._run(traj_file, top_file)
+
     # Assertions
     patched_load_single_traj.assert_called_once_with(
         ramachandran_plot_tool.path_registry, traj_file, top_file
     )
     mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True)
     mock_compute_psi.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    assert result == "Succeeded. Ramachandran plot generated."
+    ramachandran_plot_tool.path_registry.get_mapped_path.assert_called_once_with(
+        "ramachandran_plot.png"
+    )
+    mock_savefig.assert_called_once_with("ramachandran_plot.png")
+    assert result == "Succeeded. Ramachandran plot generated and saved to file."

From 7ebb3a484cd6c5c1a2e90e4aa9991aa60fe21f7d Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Fri, 21 Jun 2024 09:35:13 -0400
Subject: [PATCH 11/23] updated unit test

---
 tests/test_analysis/test_bond_angles_dihedrals.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
index 7f25c647..c14d3e74 100644
--- a/tests/test_analysis/test_bond_angles_dihedrals.py
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -144,12 +144,13 @@ def test_run_fail_compute_angles(patched_load_single_traj, compute_angles_tool):
 # Similar tests for other classes (ComputeChi1, ComputeChi2, etc.)
 
 
+@patch("matplotlib.pyplot.savefig")
 @patch("mdtraj.compute_phi")
 @patch("mdtraj.compute_psi")
 def test_run_success_ramachandran_plot(
-    mock_savefig,
     mock_compute_psi,
     mock_compute_phi,
+    mock_savefig,
     patched_load_single_traj,
     ramachandran_plot_tool,
 ):

From 244de1cb0dc00e41ce4791642e9d41f355ce7a1e Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Mon, 24 Jun 2024 19:21:23 -0400
Subject: [PATCH 12/23]  did some tweaking to try to get unit test to be
 functional.

---
 .../bond_angles_dihedrals_tool.py             | 19 ++++++++++---------
 1 file changed, 10 insertions(+), 9 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index f13d0fc0..545f2e7b 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -429,21 +429,22 @@ def _run(self, traj_file: str, top_file: str | None = None):
             map_indices_to_residues(traj, phi_indices)
             map_indices_to_residues(traj, psi_indices)
 
-            # can add further analysis or reporting here using phi_residues and
-            # psi_residues
-            plt.figure(figsize=(10, 8))
-            plt.scatter(phi_angles.flatten(), psi_angles.flatten(), s=1, color="blue")
-            plt.xlabel("Phi Angles (radians)")
-            plt.ylabel("Psi Angles (radians)")
-            plt.title("Ramachandran Plot")
-            plt.grid(True)
-
             # Check if path_registry is not None
             if self.path_registry is not None:
                 plot_save_path = self.path_registry.get_mapped_path(
                     "ramachandran_plot.png"
                 )
+                plt.figure(figsize=(10, 8))
+                plt.scatter(
+                    phi_angles.flatten(), psi_angles.flatten(), s=1, color="blue"
+                )
+                plt.xlabel("Phi Angles (radians)")
+                plt.ylabel("Psi Angles (radians)")
+                plt.title("Ramachandran Plot")
+                plt.grid(True)
+                print(f"Saving plot to: {plot_save_path}")
                 plt.savefig(plot_save_path)
+                print(f"Ramachandran plot saved to: {plot_save_path}")
                 return "Succeeded. Ramachandran plot generated and saved to file."
             else:
                 return "Failed. Path registry is not initialized."

From 615b89544426569185e7d72358170be3078da827 Mon Sep 17 00:00:00 2001
From: Brittany Watterson <bwatters@ur.rochester.edu>
Date: Mon, 24 Jun 2024 21:26:06 -0400
Subject: [PATCH 13/23]  waving white flag. Need a second set of eyes to see
 what I am missing on my unit test

---
 tests/test_analysis/test_bond_angles_dihedrals.py | 15 +++++++++------
 1 file changed, 9 insertions(+), 6 deletions(-)

diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
index c14d3e74..9d343c27 100644
--- a/tests/test_analysis/test_bond_angles_dihedrals.py
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -81,7 +81,7 @@ def compute_omega_tool(get_registry):
 
 
 @pytest.fixture
-def ramachandran_plot_tool(get_registry):
+def ramachandran_plot(get_registry):
     path_registry = get_registry("raw", True)
     return RamachandranPlot(path_registry)
 
@@ -152,7 +152,7 @@ def test_run_success_ramachandran_plot(
     mock_compute_phi,
     mock_savefig,
     patched_load_single_traj,
-    ramachandran_plot_tool,
+    ramachandran_plot,
 ):
     # Create a mock trajectory
     mock_traj = MagicMock()
@@ -165,23 +165,26 @@ def test_run_success_ramachandran_plot(
     mock_compute_psi.return_value = expected_psi
 
     # Mock the path registry get_mapped_path method
-    ramachandran_plot_tool.path_registry.get_mapped_path = MagicMock(
+    ramachandran_plot.path_registry.get_mapped_path = MagicMock(
         return_value="ramachandran_plot.png"
     )
 
     # Call the _run method
     traj_file = "rec0_butane_123456"
     top_file = "top_sim0_butane_123456"
-    result = ramachandran_plot_tool._run(traj_file, top_file)
+    result = ramachandran_plot._run(traj_file, top_file)
 
     # Assertions
     patched_load_single_traj.assert_called_once_with(
-        ramachandran_plot_tool.path_registry, traj_file, top_file
+        ramachandran_plot.path_registry, traj_file, top_file
     )
     mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True)
     mock_compute_psi.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    ramachandran_plot_tool.path_registry.get_mapped_path.assert_called_once_with(
+    ramachandran_plot.path_registry.get_mapped_path.assert_called_once_with(
         "ramachandran_plot.png"
     )
+    # Ensure savefig is called
+    print(mock_savefig.call_args_list)
     mock_savefig.assert_called_once_with("ramachandran_plot.png")
+
     assert result == "Succeeded. Ramachandran plot generated and saved to file."

From 16c0be9efa543f609824745ca3bc19097e2aaa48 Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Thu, 23 Jan 2025 20:44:17 -0500
Subject: [PATCH 14/23] ComputeAngle tool implementation

---
 .../bond_angles_dihedrals_tool.py             | 353 ++++++++++++++++--
 1 file changed, 323 insertions(+), 30 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 545f2e7b..82951d20 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -1,16 +1,44 @@
+from typing import Optional
+
 import matplotlib.pyplot as plt
 import mdtraj as md
 import numpy as np
 from langchain.tools import BaseTool
-
-from mdagent.utils import PathRegistry, load_single_traj
+from pydantic import BaseModel, Field
+
+from mdagent.utils import FileType, PathRegistry, load_single_traj
+
+
+class ComputingAnglesSchema(BaseModel):
+    trajectory_fileid: str = Field(
+        description="Trajectory File ID of the simulation to be analyzed"
+    )
+    topology_fileid: str = Field(
+        description=("Topology File ID of the simulation to be analyzed")
+    )
+    analysis: str = Field(
+        "all",
+        description=(
+            "Which analysis to be done. Availables are: "
+            "phi-psi (saves a Ramachandran plot and histograms for the Phi-Psi angles),"
+            "chi1-chi2 (gets the chi1 and chi2 dihedral angles and the chi1-chi2 plot"
+            "is saved. For the plots it only uses sidechains with enough carbons),"
+            "all (makes all of the previous analysis)"
+        ),
+    )
+    selection: Optional[str] = Field(
+        "backbone and sidechain",
+        description=(
+            "Which selection of atoms from the simulation "
+            "to use for the pca analysis"
+        ),
+    )
 
 
 class ComputeAngles(BaseTool):
     name = "compute_angles"
-    description = """Calculate the bond angles for the given sets of three atoms in
-    each snapshot, and provide a list of indices specifying which atoms are involved
-    in each bond angle calculation.."""
+    description = """Analyze dihedral angles from a trajectory file. The tool allows for
+    analysis of the phi-psi angles, chi1-chi2 angles, or both. """
 
     path_registry: PathRegistry | None = None
 
@@ -18,42 +46,307 @@ def __init__(self, path_registry: PathRegistry):
         super().__init__()
         self.path_registry = path_registry
 
-    def _run(self, traj_file: str, angle_indices: list, top_file: str | None = None):
+    def _run(self, input):
+
         try:
-            traj = load_single_traj(self.path_registry, traj_file, top_file)
-            if not traj:
-                return "Failed. Trajectory could not be loaded."
+            input = self.validate_input(**input)
+
+        except ValueError as e:
+            return f"Failed. Error using the PCA Tool: {str(e)}"
+
+        (
+            traj_id,
+            top_id,
+            analysis,
+            selection,
+            error,
+            system_input_message,
+        ) = self.get_values(input)
 
+        if error:
+            return f"Failed. Error with the tool inputs: {error} "
+        if system_input_message == "Tool Messages:":
+            system_input_message = ""
+
+        try:
+            traj = load_single_traj(
+                self.path_registry,
+                top_id,
+                traj_fileid=traj_id,
+                traj_required=True,
+            )
+        except ValueError as e:
             if (
-                not angle_indices
-                or not isinstance(angle_indices, list)
-                or not all(len(indices) == 3 for indices in angle_indices)
+                "The topology and the trajectory files might not\
+                  contain the same atoms"
+                in str(e)
             ):
                 return (
-                    "Failed. Invalid angle_indices. It should be a list of tuples, "
-                    "each containing three atom indices."
+                    "Failed. Error loading trajectory. Make sure the topology file"
+                    " is from the initial positions of the trajectory. Error: {str(e)}"
                 )
+            return f"Failed. Error loading trajectory: {str(e)}"
+        except OSError as e:
+            if (
+                "The topology is loaded by filename extension, \
+                and the detected"
+                in str(e)
+            ):
+                return (
+                    "Failed. Error loading trajectory. Make sure you include the"
+                    "correct file for the topology. Supported extensions are:"
+                    "'.pdb', '.pdb.gz', '.h5', '.lh5', '.prmtop', '.parm7', '.prm7',"
+                    "  '.psf', '.mol2', '.hoomdxml', '.gro', '.arc', '.hdf5' and '.gsd'"
+                )
+            return f"Failed. Error loading trajectory: {str(e)}"
+        except Exception as e:
+            return f"Failed. Error loading trajectory: {str(e)}"
 
-            angles = md.compute_angles(traj, angle_indices, periodic=True, opt=True)
+        return self.analyze_trajectory(traj, analysis, self.path_registry, traj_id)
 
-            # Check if path_registry is not None
-            if self.path_registry is not None:
-                plot_save_path = self.path_registry.get_mapped_path("angles_plot.png")
-                plot_angles(angles, title="Bond Angles", save_path=plot_save_path)
-                return "Succeeded. Bond angles computed, saved to file and plot saved."
-            else:
-                return "Failed. Path registry is not initialized."
+    async def _arun(self, input):
+        raise NotImplementedError("Async version not implemented")
+
+    # Example helper functions (optional). You can instead just keep them as
+    # blocks in the if-statements.
+    def compute_and_plot_phi_psi(self, traj, path_registry, sim_id):
+        """
+        Computes phi-psi angles, saves results to file, and produces Ramachandran plot.
+        """
+        try:
+            # Compute phi and psi angles
+            phi_indices, phi_angles = md.compute_phi(traj)
+            psi_indices, psi_angles = md.compute_psi(traj)
 
+            # Convert angles to degrees
+            phi_angles = phi_angles * (180.0 / np.pi)
+            psi_angles = psi_angles * (180.0 / np.pi)
         except Exception as e:
-            return f"Failed. {type(e).__name__}: {e}"
+            return None, f"Failed. Error computing phi-psi angles: {str(e)}"
+
+        # If path_registry is available, save files and produce plot
+        if path_registry is not None:
+            # Save angle results
+            save_results_to_file("phi_results.npz", phi_indices, phi_angles)
+            save_results_to_file("psi_results.npz", psi_indices, psi_angles)
+
+            # Make Ramachandran plot
+            try:
+                plt.hist2d(
+                    phi_angles.flatten(), psi_angles.flatten(), bins=150, cmap="Blues"
+                )
+                plt.xlabel(r"$\phi$")
+                plt.ylabel(r"$\psi$")
+                plt.colorbar()
+
+                file_name = path_registry.write_file_name(
+                    FileType.FIGURE,
+                    fig_analysis="ramachandran",
+                    file_format="png",
+                    Sim_id=sim_id,
+                )
+                desc = f"Ramachandran plot for the simulation {sim_id}"
+                plot_id = path_registry.get_fileid(file_name, FileType.FIGURE)
+                path = path_registry.ckpt_dir + "/figures/"
+                plt.savefig(path + file_name)
+                path_registry.map_path(plot_id, path + file_name, description=desc)
+                plt.clf()  # Clear the current figure so it does not overlay next plot
+                print("Ramachandran plot saved to file")
+                return plot_id, "Succeeded. Ramachandran plot saved."
+            except Exception as e:
+                return None, f"Failed. Error saving Ramachandran plot: {str(e)}"
+        else:
+            return (
+                None,
+                "Succeeded. Computed phi-psi angles (no path_registry to save).",
+            )
+
+    def compute_and_plot_chi1_chi2(self, traj, path_registry, sim_id):
+        """
+        Computes chi1-chi2 angles, saves results to file, and produces Chi1-Chi2 plot.
+        """
+        try:
+            # Compute chi1 and chi2 angles
+            chi1_indices, chi1_angles = md.compute_chi1(traj)
+            chi2_indices, chi2_angles = md.compute_chi2(traj)
 
-    async def _arun(
-        self,
-        traj_file: str,
-        angle_indices: list,
-        top_file: str | None = None,
-    ):
-        raise NotImplementedError("Async version not implemented")
+            # Convert angles to degrees
+            chi1_angles = chi1_angles * (180.0 / np.pi)
+            chi2_angles = chi2_angles * (180.0 / np.pi)
+        except Exception as e:
+            return None, f"Failed. Error computing chi1-chi2 angles: {str(e)}"
+
+        # If path_registry is available, save files and produce plot
+        if path_registry is not None:
+            # Get the indices of the first side-chain atoms from chi1 and chi2
+            chi1_atoms = [atom_idx[1] for atom_idx in chi1_indices]
+            chi2_atoms = [atom_idx[0] for atom_idx in chi2_indices]
+
+            # Filter chi1 angles to match atoms that appear in chi2
+            chi1_angles_long = np.array(
+                [
+                    chi1_angles[:, i]
+                    for i, chi1_atom in enumerate(chi1_atoms)
+                    if chi1_atom in chi2_atoms
+                ]
+            )
+
+            # Save angle results
+            save_results_to_file("chi1_results.npz", chi1_indices, chi1_angles)
+            save_results_to_file("chi2_results.npz", chi2_indices, chi2_angles)
+
+            # Make Chi1-Chi2 plot
+            try:
+                plt.hist2d(
+                    chi1_angles_long.T.flatten(),
+                    chi2_angles.flatten(),
+                    bins=200,
+                    cmap="Blues",
+                )
+                plt.xlabel(r"$\chi1$")
+                plt.ylabel(r"$\chi2$")
+                plt.title(f"Chi1-Chi2 plot for the simulation {sim_id}")
+                plt.colorbar()
+
+                file_name = path_registry.write_file_name(
+                    FileType.FIGURE,
+                    fig_analysis="chi1-chi2",
+                    file_format="png",
+                    Sim_id=sim_id,
+                )
+                desc = f"Chi1-Chi2 plot for the simulation {sim_id}"
+                chi_plot_id = path_registry.get_fileid(file_name, FileType.FIGURE)
+                path = path_registry.ckpt_dir + "/figures/"
+                plt.savefig(path + file_name)
+                path_registry.map_path(chi_plot_id, path + file_name, description=desc)
+                plt.clf()  # Clear the current figure so it does not overlay next plot
+                print("Chi1-Chi2 plot saved to file")
+                return chi_plot_id, "Succeeded. Chi1-Chi2 plot saved."
+            except Exception as e:
+                return None, f"Failed. Error saving Chi1-Chi2 plot: {str(e)}"
+        else:
+
+            return None, "Succeeded. Computed chi1-chi2 angles."
+
+    def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
+        """
+        Main function to decide which analysis to do:
+        'phi-psi', 'chi1-chi2', or 'all'.
+        """
+        # Store optional references for convenience
+        self_path_registry = path_registry
+        self_sim_id = sim_id
+
+        # ================ PHI-PSI ONLY =================
+        if analysis == "phi-psi":
+            plot_id, message = self.compute_and_plot_phi_psi(
+                traj, self_path_registry, self_sim_id
+            )
+            return message
+
+        # ================ CHI1-CHI2 ONLY ================
+        elif analysis == "chi1-chi2":
+            plot_id, message = self.compute_and_plot_chi1_chi2(
+                traj, self_path_registry, self_sim_id
+            )
+            return message
+
+        # ================ ALL =================
+        elif analysis == "all":
+            # First do phi-psi
+            phi_plot_id, phi_message = self.compute_and_plot_phi_psi(
+                traj, self_path_registry, self_sim_id
+            )
+            if "Failed." in phi_message:
+                return phi_message
+
+            # Then do chi1-chi2
+            chi_plot_id, chi_message = self.compute_and_plot_chi1_chi2(
+                traj, self_path_registry, self_sim_id
+            )
+            if "Failed." in chi_message:
+                return chi_message
+
+            return (
+                "Succeeded. All analyses completed. "
+                f"Ramachandran plot message: {phi_message} "
+                f"Chi1-Chi2 plot message: {chi_message}"
+            )
+
+        else:
+            # Unknown analysis type
+            return f"Failed. Unknown analysis type: {analysis}"
+
+    def validate_input(self, **input):
+        input = input.get("action_input", input)
+        input = input.get("input", input)
+        trajectory_id = input.get("trajectory_fileid", None)
+        topology_id = input.get("topology_fileid", None)
+        analysis = input.get("analysis", "all")
+        selection = input.get("selection", "backbone and sidechain")
+        if not trajectory_id:
+            raise ValueError("Incorrect Inputs: trajectory_fileid is required")
+        if not topology_id:
+            raise ValueError("Incorrect Inputs: topology_fileid is required")
+        # check if trajectory id is valid
+        fileids = self.path_registry.list_path_names()
+        error = ""
+        system_message = "Tool Messages:"
+        if trajectory_id not in fileids:
+            error += " Trajectory File ID not in path registry"
+        if topology_id not in fileids:
+            error += " Topology File ID not in path registry"
+
+        if analysis.lower() not in [
+            "all",
+            "phi-psi",
+            "chi1-chi2",
+        ]:
+            analysis = "all"
+            system_message += (
+                " analysis arg not recognized, using analysis = 'all' as default"
+            )
+
+        if selection not in [
+            "backbone",
+            "name CA",
+            "backbone and name CA",
+            "protein",
+            "backbone and sidechain",
+            "sidechain",
+            "all",
+        ]:
+            selection = "all"  # just alpha carbons
+        # get all the kwargs:
+        keys = input.keys()
+        for key in keys:
+            if key not in [
+                "trajectory_fileid",
+                "pc_percentage",
+                "analysis",
+                "selection",
+            ]:
+                system_message += f"{key} is not part of admitted tool inputs"
+        if error == "":
+            error = None
+        return {
+            "trajectory_fileid": trajectory_id,
+            "topology_fileid": topology_id,
+            "analysis": analysis,
+            "selection": selection,
+            "error": error,
+            "system_message": system_message,
+        }
+
+    def get_values(self, input):
+        traj_id = input.get("trajectory_fileid")
+        top_id = input.get("topology_fileid")
+        analysis = input.get("analysis")
+        sel = input.get("selection")
+        error = input.get("error")
+        syst_mes = input.get("system_message")
+
+        return traj_id, top_id, analysis, sel, error, syst_mes
 
 
 class ComputeDihedrals(BaseTool):

From 54671322dbf4ba3c8d9104f41d3bbd07cdaafad9 Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Fri, 24 Jan 2025 00:30:57 -0500
Subject: [PATCH 15/23] removing unnecesary angle tools

---
 .../base_tools/analysis_tools/__init__.py     | 22 +------------------
 .../bond_angles_dihedrals_tool.py             |  1 +
 mdagent/tools/maketools.py                    | 20 +----------------
 3 files changed, 3 insertions(+), 40 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/__init__.py b/mdagent/tools/base_tools/analysis_tools/__init__.py
index 1d85b38d..5684852a 100644
--- a/mdagent/tools/base_tools/analysis_tools/__init__.py
+++ b/mdagent/tools/base_tools/analysis_tools/__init__.py
@@ -1,15 +1,4 @@
-from .bond_angles_dihedrals_tool import (
-    ComputeAngles,
-    ComputeChi1,
-    ComputeChi2,
-    ComputeChi3,
-    ComputeChi4,
-    ComputeDihedrals,
-    ComputeOmega,
-    ComputePhi,
-    ComputePsi,
-    RamachandranPlot,
-)
+from .bond_angles_dihedrals_tool import ComputeAngles
 from .distance_tools import ContactsTool, DistanceMatrixTool
 from .inertia import MomentOfInertia
 from .pca_tools import PCATool
@@ -22,14 +11,6 @@
 
 __all__ = [
     "ComputeAngles",
-    "ComputeChi1",
-    "ComputeChi2",
-    "ComputeChi3",
-    "ComputeChi4",
-    "ComputeDihedrals",
-    "ComputeOmega",
-    "ComputePhi",
-    "ComputePsi",
     "ComputeLPRMSD",
     "ComputeRMSD",
     "ComputeRMSF",
@@ -39,7 +20,6 @@
     "PCATool",
     "PPIDistance",
     "RadiusofGyrationTool",
-    "RamachandranPlot",
     "RMSDCalculator",
     "SimulationOutputFigures",
     "SolventAccessibleSurfaceArea",
diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 82951d20..3b9faee7 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -41,6 +41,7 @@ class ComputeAngles(BaseTool):
     analysis of the phi-psi angles, chi1-chi2 angles, or both. """
 
     path_registry: PathRegistry | None = None
+    args_schema = ComputingAnglesSchema
 
     def __init__(self, path_registry: PathRegistry):
         super().__init__()
diff --git a/mdagent/tools/maketools.py b/mdagent/tools/maketools.py
index 3cd31d9a..88e63c79 100644
--- a/mdagent/tools/maketools.py
+++ b/mdagent/tools/maketools.py
@@ -13,20 +13,12 @@
 from .base_tools import (
     CleaningToolFunction,
     ComputeAcylindricity,
+    ComputeAngles,
     ComputeAsphericity,
     ComputeDSSP,
     ComputeGyrationTensor,
-    ComputeAngles,
-    ComputeChi1,
-    ComputeChi2,
-    ComputeChi3,
-    ComputeChi4,
-    ComputeDihedrals,
     ComputeLPRMSD,
     ComputeRelativeShapeAntisotropy,
-    ComputeOmega,
-    ComputePhi,
-    ComputePsi,
     ComputeRMSD,
     ComputeRMSF,
     ContactsTool,
@@ -56,7 +48,6 @@
     PPIDistance,
     ProteinName2PDBTool,
     RadiusofGyrationTool,
-    RamachandranPlot,
     RDFTool,
     Scholar2ResultLLM,
     SetUpandRunFunction,
@@ -96,14 +87,6 @@ def make_all_tools(
         ComputeGyrationTensor(path_registry=path_instance),
         ComputeRelativeShapeAntisotropy(path_registry=path_instance),
         ComputeAngles(path_registry=path_instance),
-        ComputeChi1(path_registry=path_instance),
-        ComputeChi2(path_registry=path_instance),
-        ComputeChi3(path_registry=path_instance),
-        ComputeChi4(path_registry=path_instance),
-        ComputeDihedrals(path_registry=path_instance),
-        ComputeOmega(path_registry=path_instance),
-        ComputePhi(path_registry=path_instance),
-        ComputePsi(path_registry=path_instance),
         CleaningToolFunction(path_registry=path_instance),
         ComputeLPRMSD(path_registry=path_instance),
         ComputeRMSD(path_registry=path_instance),
@@ -117,7 +100,6 @@ def make_all_tools(
         PPIDistance(path_registry=path_instance),
         ProteinName2PDBTool(path_registry=path_instance),
         RadiusofGyrationTool(path_registry=path_instance),
-        RamachandranPlot(path_registry=path_instance),
         RDFTool(path_registry=path_instance),
         SetUpandRunFunction(path_registry=path_instance),
         SimulationOutputFigures(path_registry=path_instance),

From 59071e357be54cb321971dab979d9bd57a0ec49e Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Fri, 24 Jan 2025 11:22:01 -0500
Subject: [PATCH 16/23] adding small peptide traj in conftest.py

---
 tests/conftest.py | 64 ++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 63 insertions(+), 1 deletion(-)

diff --git a/tests/conftest.py b/tests/conftest.py
index 38a35ffa..27a775e9 100644
--- a/tests/conftest.py
+++ b/tests/conftest.py
@@ -357,6 +357,58 @@ def butane_trajectory_without_hydrogens(request):
     request.addfinalizer(lambda: safe_remove(pdb_file))
 
 
+@pytest.fixture(scope="module")
+def small_peptide_gag_trajectory(request):
+    """
+    Writes out a small multi-model PDB with two frames of a Gly–Ala–Gly peptide.
+    Yields the filename so that tests can load and analyze phi/psi angles.
+    """
+    pdb_content = """\
+MODEL        1
+ATOM      1  N   GLY A   1       1.000   1.000   1.000  1.00  0.00           N
+ATOM      2  CA  GLY A   1       2.000   1.000   1.000  1.00  0.00           C
+ATOM      3  C   GLY A   1       2.500   2.250   1.000  1.00  0.00           C
+ATOM      4  O   GLY A   1       3.000   3.000   1.000  1.00  0.00           O
+ATOM      5  N   ALA A   2       1.700   2.900   1.000  1.00  0.00           N
+ATOM      6  CA  ALA A   2       2.000   4.300   1.100  1.00  0.00           C
+ATOM      7  CB  ALA A   2       1.000   5.250   1.000  1.00  0.00           C
+ATOM      8  C   ALA A   2       3.400   5.000   1.000  1.00  0.00           C
+ATOM      9  O   ALA A   2       4.100   5.900   1.500  1.00  0.00           O
+ATOM     10  N   GLY A   3       3.900   4.300   0.200  1.00  0.00           N
+ATOM     11  CA  GLY A   3       5.200   4.850   0.000  1.00  0.00           C
+ATOM     12  C   GLY A   3       6.100   3.700   0.000  1.00  0.00           C
+ATOM     13  O   GLY A   3       7.000   3.700   0.700  1.00  0.00           O
+ENDMDL
+MODEL        2
+ATOM      1  N   GLY A   1       1.100   1.000   0.900  1.00  0.00           N
+ATOM      2  CA  GLY A   1       2.050   1.100   1.100  1.00  0.00           C
+ATOM      3  C   GLY A   1       2.550   2.300   1.100  1.00  0.00           C
+ATOM      4  O   GLY A   1       3.050   3.100   1.050  1.00  0.00           O
+ATOM      5  N   ALA A   2       1.650   2.900   1.100  1.00  0.00           N
+ATOM      6  CA  ALA A   2       2.000   4.300   1.100  1.00  0.00           C
+ATOM      7  CB  ALA A   2       1.200   5.350   1.050  1.00  0.00           C
+ATOM      8  C   ALA A   2       3.450   4.900   1.050  1.00  0.00           C
+ATOM      9  O   ALA A   2       4.150   5.800   1.600  1.00  0.00           O
+ATOM     10  N   GLY A   3       3.900   4.200   0.350  1.00  0.00           N
+ATOM     11  CA  GLY A   3       5.300   4.850   0.300  1.00  0.00           C
+ATOM     12  C   GLY A   3       6.200   3.650   0.250  1.00  0.00           C
+ATOM     13  O   GLY A   3       7.100   3.600   0.900  1.00  0.00           O
+ENDMDL
+END
+"""
+
+    # You can name this file however you want; using a random suffix is common.
+    pdb_filename = "SMALL_GAG_trajectory_987654.pdb"
+    with open(pdb_filename, "w") as f:
+        f.write(pdb_content)
+
+    # Yield the filename to the test function(s).
+    yield pdb_filename
+
+    # Cleanup after tests complete
+    request.addfinalizer(lambda: safe_remove(pdb_filename))
+
+
 @pytest.fixture(scope="module")
 def get_registry(
     raw_alanine_pdb_file,
@@ -364,6 +416,7 @@ def get_registry(
     butane_static_trajectory_with_hydrogens,
     butane_dynamic_trajectory_with_hydrogens,
     butane_trajectory_without_hydrogens,
+    small_peptide_gag_trajectory,
     request,
 ):
     created_paths = []  # Keep track of created directories for cleanup
@@ -374,7 +427,12 @@ def get_new_ckpt():
         return base_path, registry
 
     def create(
-        raw_or_clean, with_files, dynamic=False, include_hydrogens=False, map_path=True
+        raw_or_clean,
+        with_files,
+        dynamic=False,
+        include_hydrogens=False,
+        map_path=True,
+        include_peptide_trajectory=False,
     ):
         base_path, registry = get_new_ckpt()
         created_paths.append(base_path)
@@ -402,6 +460,10 @@ def create(
                     traj_file, top_file = butane_static_trajectory_with_hydrogens
             else:
                 traj_file, top_file = butane_trajectory_without_hydrogens
+            if include_peptide_trajectory:
+                pep_traj_file = small_peptide_gag_trajectory
+                files["pep_traj_987654"] = {"name": pep_traj_file, "dir": record_path}
+
             files["rec0_butane_123456"] = {"name": traj_file, "dir": record_path}
             files["top_sim0_butane_123456"] = {"name": top_file, "dir": pdb_path}
             for f in files:

From 15b7ccc81c057834a6b7f5dbc5d6d8ca7c1d6502 Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Fri, 24 Jan 2025 16:13:05 -0500
Subject: [PATCH 17/23] adding initial tests that check correct workflow for
 angle tool

---
 .../bond_angles_dihedrals_tool.py             |   8 +-
 .../test_bond_angles_dihedrals.py             | 236 ++++--------------
 2 files changed, 58 insertions(+), 186 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 3b9faee7..02fefcf6 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -26,6 +26,8 @@ class ComputingAnglesSchema(BaseModel):
             "all (makes all of the previous analysis)"
         ),
     )
+    # This arg is here, but is not used in the code. As of now it will get the analysis
+    # of all the residues in the simulation
     selection: Optional[str] = Field(
         "backbone and sidechain",
         description=(
@@ -36,7 +38,7 @@ class ComputingAnglesSchema(BaseModel):
 
 
 class ComputeAngles(BaseTool):
-    name = "compute_angles"
+    name = "ComputeAngles"
     description = """Analyze dihedral angles from a trajectory file. The tool allows for
     analysis of the phi-psi angles, chi1-chi2 angles, or both. """
 
@@ -53,7 +55,7 @@ def _run(self, input):
             input = self.validate_input(**input)
 
         except ValueError as e:
-            return f"Failed. Error using the PCA Tool: {str(e)}"
+            return f"Failed. Error using the ComputeAngle Tool: {str(e)}"
 
         (
             traj_id,
@@ -323,7 +325,7 @@ def validate_input(self, **input):
         for key in keys:
             if key not in [
                 "trajectory_fileid",
-                "pc_percentage",
+                "topology_fileid",
                 "analysis",
                 "selection",
             ]:
diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
index 9d343c27..0e04c81b 100644
--- a/tests/test_analysis/test_bond_angles_dihedrals.py
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -1,190 +1,60 @@
-from unittest.mock import MagicMock, patch
-
-import numpy as np
-import pytest
+from unittest.mock import patch
 
 from mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool import (
     ComputeAngles,
-    ComputeChi1,
-    ComputeChi2,
-    ComputeChi3,
-    ComputeChi4,
-    ComputeDihedrals,
-    ComputeOmega,
-    ComputePhi,
-    ComputePsi,
-    RamachandranPlot,
 )
 
 
-# Fixture to patch 'load_single_traj'
-@pytest.fixture
-def patched_load_single_traj():
-    with patch(
-        "mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.load_single_traj"
-    ) as mock_load_single_traj:
-        yield mock_load_single_traj
-
-
-@pytest.fixture
-def compute_angles_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeAngles(path_registry)
-
-
-@pytest.fixture
-def compute_dihedrals_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeDihedrals(path_registry)
-
-
-@pytest.fixture
-def compute_phi_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputePhi(path_registry)
-
-
-@pytest.fixture
-def compute_psi_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputePsi(path_registry)
-
-
-@pytest.fixture
-def compute_chi1_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeChi1(path_registry)
-
-
-@pytest.fixture
-def compute_chi2_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeChi2(path_registry)
-
-
-@pytest.fixture
-def compute_chi3_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeChi3(path_registry)
-
-
-@pytest.fixture
-def compute_chi4_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeChi4(path_registry)
-
-
-@pytest.fixture
-def compute_omega_tool(get_registry):
-    path_registry = get_registry("raw", True)
-    return ComputeOmega(path_registry)
-
+def test_compute_angles_tool_bad_inputs(get_registry):
+    reg = get_registry("raw", True, map_path=True, include_peptide_trajectory=True)
+    angles_tool = ComputeAngles(path_registry=reg)
+    bad_input_files = {
+        "trajectory_fileid": "pep_traj_987654_3",
+        "topology_fileid": "pep_traj_987654_3",
+        "analysis": "both",
+    }
+
+    error_catching = angles_tool._run(bad_input_files)
+    assert "Trajectory File ID not in path registry" in error_catching
+    assert "Topology File ID not in path registry" in error_catching
+
+
+# patch and or moch save_results_to_file
+# @patch("mdagent.tools.base_tools.analysis_tools.bond_angles_dihedrals_tool.save_results_to_file")
+def test_compute_angles_ram_values(get_registry):
+    reg = get_registry("raw", True, dynamic=True, include_hydrogens=True)
+    angles_tool = ComputeAngles(path_registry=reg)
+    phi_psi_input_files = {
+        "trajectory_fileid": "pep_traj_987654",
+        "topology_fileid": "pep_traj_987654",
+        "analysis": "phi-psi",
+    }
+    chi_innput_files = {
+        "trajectory_fileid": "pep_traj_987654",
+        "topology_fileid": "pep_traj_987654",
+        "analysis": "chi1-chi2",
+    }
+    # traj = md.load(reg.get_mapped_path("pep_traj_987654"))
 
-@pytest.fixture
-def ramachandran_plot(get_registry):
-    path_registry = get_registry("raw", True)
-    return RamachandranPlot(path_registry)
-
-
-@patch("mdtraj.compute_angles")
-@patch("matplotlib.pyplot.savefig")
-def test_run_success_compute_angles(
-    mock_savefig, mock_compute_angles, patched_load_single_traj, compute_angles_tool
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    patched_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_angles
-    expected_angles = np.array([[0.1, 0.2, 0.3], [0.4, 0.5, 0.6]])
-    mock_compute_angles.return_value = expected_angles
-
-    # Mock the path registry get_mapped_path method
-    compute_angles_tool.path_registry.get_mapped_path = MagicMock(
-        return_value="angles_plot.png"
-    )
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    angle_indices = [(0, 1, 2), (1, 2, 3)]
-    result = compute_angles_tool._run(traj_file, angle_indices, top_file)
-
-    # Assertions
-    patched_load_single_traj.assert_called_once_with(
-        compute_angles_tool.path_registry, traj_file, top_file
-    )
-    mock_compute_angles.assert_called_once_with(
-        mock_traj, angle_indices, periodic=True, opt=True
-    )
-    compute_angles_tool.path_registry.get_mapped_path.assert_called_once_with(
-        "angles_plot.png"
-    )
-    mock_savefig.assert_called_once_with("angles_plot.png")
-    assert result == "Succeeded. Bond angles computed, saved to file and plot saved."
-
-
-def test_run_fail_compute_angles(patched_load_single_traj, compute_angles_tool):
-    # Simulate the trajectory loading failure
-    patched_load_single_traj.return_value = None
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    angle_indices = [(0, 1, 2), (1, 2, 3)]
-    result = compute_angles_tool._run(traj_file, angle_indices, top_file)
-
-    # Assertions
-    patched_load_single_traj.assert_called_once_with(
-        compute_angles_tool.path_registry, traj_file, top_file
-    )
-    assert result == "Failed. Trajectory could not be loaded."
-
-
-# Similar tests for other classes (ComputeChi1, ComputeChi2, etc.)
-
-
-@patch("matplotlib.pyplot.savefig")
-@patch("mdtraj.compute_phi")
-@patch("mdtraj.compute_psi")
-def test_run_success_ramachandran_plot(
-    mock_compute_psi,
-    mock_compute_phi,
-    mock_savefig,
-    patched_load_single_traj,
-    ramachandran_plot,
-):
-    # Create a mock trajectory
-    mock_traj = MagicMock()
-    patched_load_single_traj.return_value = mock_traj
-
-    # Define the expected output from compute_phi and compute_psi
-    expected_phi = ([(0, 1, 2, 3)], [[0.7, 0.8, 0.9]])
-    expected_psi = ([(0, 1, 2, 3)], [[1.0, 1.1, 1.2]])
-    mock_compute_phi.return_value = expected_phi
-    mock_compute_psi.return_value = expected_psi
-
-    # Mock the path registry get_mapped_path method
-    ramachandran_plot.path_registry.get_mapped_path = MagicMock(
-        return_value="ramachandran_plot.png"
-    )
-
-    # Call the _run method
-    traj_file = "rec0_butane_123456"
-    top_file = "top_sim0_butane_123456"
-    result = ramachandran_plot._run(traj_file, top_file)
-
-    # Assertions
-    patched_load_single_traj.assert_called_once_with(
-        ramachandran_plot.path_registry, traj_file, top_file
-    )
-    mock_compute_phi.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    mock_compute_psi.assert_called_once_with(mock_traj, periodic=True, opt=True)
-    ramachandran_plot.path_registry.get_mapped_path.assert_called_once_with(
-        "ramachandran_plot.png"
-    )
-    # Ensure savefig is called
-    print(mock_savefig.call_args_list)
-    mock_savefig.assert_called_once_with("ramachandran_plot.png")
-
-    assert result == "Succeeded. Ramachandran plot generated and saved to file."
+    with patch(
+        "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_and_plot_phi_psi"
+    ) as mock_compute_and_plot_phi_psi:
+        with patch(
+            "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_and_plot_chi1_chi2"
+        ) as mock_compute_and_plot_chi1_chi2:
+            mock_compute_and_plot_phi_psi.return_value = ("mockid", "mockresult")
+            # instance.return_value = ("mockid", "mockresult")
+            angles_tool._run(phi_psi_input_files)
+            # print(result)
+            assert mock_compute_and_plot_phi_psi.called
+            # assert compute_and_plot_chi1_chi2 is not called
+            assert not mock_compute_and_plot_chi1_chi2.called
+
+            # =========================================================================#
+            mock_compute_and_plot_chi1_chi2.return_value = ("mockid", "mockresult")
+            angles_tool._run(chi_innput_files)
+            assert mock_compute_and_plot_chi1_chi2.called
+            # assert compute_and_plot_phi_psi is not called
+            assert (
+                mock_compute_and_plot_phi_psi.assert_called_once
+            )  # already called once

From 52aae411499ef3f35b8e801d8be8ef9d8dc6e64a Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Mon, 27 Jan 2025 15:52:32 -0500
Subject: [PATCH 18/23] redoing chi angle analysis for computeangle tools

---
 .../bond_angles_dihedrals_tool.py             | 218 +++++++++++-------
 1 file changed, 141 insertions(+), 77 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 02fefcf6..2f6655da 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -21,9 +21,9 @@ class ComputingAnglesSchema(BaseModel):
         description=(
             "Which analysis to be done. Availables are: "
             "phi-psi (saves a Ramachandran plot and histograms for the Phi-Psi angles),"
-            "chi1-chi2 (gets the chi1 and chi2 dihedral angles and the chi1-chi2 plot"
-            "is saved. For the plots it only uses sidechains with enough carbons),"
-            "all (makes all of the previous analysis)"
+            "chis (gets the chis 1-4 angles and plots a time evolutiuon plot for all"
+            "residues is saved. For the plots it only uses sidechains with enough "
+            "carbons), all (makes all of the previous analysis)"
         ),
     )
     # This arg is here, but is not used in the code. As of now it will get the analysis
@@ -164,77 +164,141 @@ def compute_and_plot_phi_psi(self, traj, path_registry, sim_id):
                 "Succeeded. Computed phi-psi angles (no path_registry to save).",
             )
 
-    def compute_and_plot_chi1_chi2(self, traj, path_registry, sim_id):
+    def classify_chi(self, ang_deg, res_name=""):
+        """Return an integer code depending on angle range."""
+        # Example classification with made-up intervals:
+        if res_name == "PRO" or res_name == "P":
+            if ang_deg < 0:
+                return 3  # e.g. "p-"
+            else:
+                return 4  # e.g. "p+"
+        # angles for g+
+        if 0 <= ang_deg < 120:
+            return 0  # e.g. "g+"
+        # angles for t
+        elif -120 >= ang_deg or ang_deg > 120:
+            return 1  # e.g. "t"
+        # angles for g-
+        elif -120 <= ang_deg < 0:
+            return 2  # e.g. "g-"
+
+    # function that takes an array and classifies the angles
+    def classify_chi_angles(self, angles, res_name=""):
+        return [self.classify_chi(ang, res_name) for ang in angles]
+
+    def _plot_one_chi_angle(self, ax, angle_array, residue_names, title=None):
         """
-        Computes chi1-chi2 angles, saves results to file, and produces Chi1-Chi2 plot.
+        Classify angles per residue/frame, then do imshow on a given Axes.
+        angle_array: shape (n_frames, n_residues) or (n_residues, n_frames)
+        residue_names: length n_residues
         """
-        try:
-            # Compute chi1 and chi2 angles
-            chi1_indices, chi1_angles = md.compute_chi1(traj)
-            chi2_indices, chi2_angles = md.compute_chi2(traj)
-
-            # Convert angles to degrees
-            chi1_angles = chi1_angles * (180.0 / np.pi)
-            chi2_angles = chi2_angles * (180.0 / np.pi)
-        except Exception as e:
-            return None, f"Failed. Error computing chi1-chi2 angles: {str(e)}"
-
-        # If path_registry is available, save files and produce plot
-        if path_registry is not None:
-            # Get the indices of the first side-chain atoms from chi1 and chi2
-            chi1_atoms = [atom_idx[1] for atom_idx in chi1_indices]
-            chi2_atoms = [atom_idx[0] for atom_idx in chi2_indices]
-
-            # Filter chi1 angles to match atoms that appear in chi2
-            chi1_angles_long = np.array(
-                [
-                    chi1_angles[:, i]
-                    for i, chi1_atom in enumerate(chi1_atoms)
-                    if chi1_atom in chi2_atoms
-                ]
-            )
-
-            # Save angle results
-            save_results_to_file("chi1_results.npz", chi1_indices, chi1_angles)
-            save_results_to_file("chi2_results.npz", chi2_indices, chi2_angles)
-
-            # Make Chi1-Chi2 plot
-            try:
-                plt.hist2d(
-                    chi1_angles_long.T.flatten(),
-                    chi2_angles.flatten(),
-                    bins=200,
-                    cmap="Blues",
-                )
-                plt.xlabel(r"$\chi1$")
-                plt.ylabel(r"$\chi2$")
-                plt.title(f"Chi1-Chi2 plot for the simulation {sim_id}")
-                plt.colorbar()
-
-                file_name = path_registry.write_file_name(
-                    FileType.FIGURE,
-                    fig_analysis="chi1-chi2",
-                    file_format="png",
-                    Sim_id=sim_id,
-                )
-                desc = f"Chi1-Chi2 plot for the simulation {sim_id}"
-                chi_plot_id = path_registry.get_fileid(file_name, FileType.FIGURE)
-                path = path_registry.ckpt_dir + "/figures/"
-                plt.savefig(path + file_name)
-                path_registry.map_path(chi_plot_id, path + file_name, description=desc)
-                plt.clf()  # Clear the current figure so it does not overlay next plot
-                print("Chi1-Chi2 plot saved to file")
-                return chi_plot_id, "Succeeded. Chi1-Chi2 plot saved."
-            except Exception as e:
-                return None, f"Failed. Error saving Chi1-Chi2 plot: {str(e)}"
-        else:
-
-            return None, "Succeeded. Computed chi1-chi2 angles."
+        state_sequence = np.array(
+            [
+                [self.classify_chi_angles(a, str(name)[:3])]
+                for i, (a, name) in enumerate(zip(angle_array.T, residue_names))
+            ]
+        )
+        states_per_res = state_sequence.reshape(
+            state_sequence.shape[0], state_sequence.shape[2]
+        )  # shape = (#res,1, #frames)
+        # -> (#res, #frames)
+
+        n_residues = len(residue_names)
+        unique_states = np.unique(states_per_res)
+        n_states = len(unique_states)
+        cmap = plt.get_cmap("tab20", n_states)
+
+        im = ax.imshow(
+            states_per_res,
+            aspect="auto",
+            interpolation="none",
+            cmap=cmap,
+            origin="upper",
+        )
+
+        ax.set_xlabel("Frame index")
+        ax.set_ylabel("Residue")
+        if title:
+            ax.set_title(title)
+
+        ax.set_yticks(np.arange(n_residues))
+        ax.set_yticklabels([str(r) for r in residue_names], fontsize=8)
+
+        cbar = plt.colorbar(im, ax=ax, ticks=range(n_states), pad=0.01)
+
+        # Example state -> label mapping
+        state_labels_map = {0: "g+", 1: "t", 2: "g-", 3: "Cγ endo", 4: "Cγ exo"}
+        tick_labels = [state_labels_map.get(s, f"State {s}") for s in unique_states]
+        cbar.ax.set_yticklabels(tick_labels, fontsize=8)
+
+    ###################################################
+    # Main function to produce a single figure w/ 4 subplots
+    ###################################################
+    def compute_plot_all_chi_angles(self, traj, sim_id="sim"):
+        """
+        Create one figure with 4 subplots (2x2):
+        - subplot(0,0): χ1
+        - subplot(0,1): χ2
+        - subplot(1,0): χ3
+        - subplot(1,1): χ4
+        """
+        chi1_indices, chi_1_angles = md.compute_chi1(traj)
+        chi2_indices, chi_2_angles = md.compute_chi2(traj)
+        chi3_indices, chi_3_angles = md.compute_chi3(traj)
+        chi4_indices, chi_4_angles = md.compute_chi4(traj)
+
+        chi_1_angles_degrees = np.rad2deg(chi_1_angles)
+        chi_2_angles_degrees = np.rad2deg(chi_2_angles)
+        chi_3_angles_degrees = np.rad2deg(chi_3_angles)
+        chi_4_angles_degrees = np.rad2deg(chi_4_angles)
+        residue_names_1 = [traj.topology.atom(i).residue for i in chi1_indices[:, 1]]
+        residue_names_2 = [traj.topology.atom(i).residue for i in chi2_indices[:, 1]]
+        residue_names_3 = [traj.topology.atom(i).residue for i in chi3_indices[:, 1]]
+        residue_names_4 = [traj.topology.atom(i).residue for i in chi4_indices[:, 1]]
+        fig, axes = plt.subplots(2, 2, figsize=(12, 8))
+
+        # Top-left: χ1
+        self._plot_one_chi_angle(
+            axes[0, 0], chi_1_angles_degrees, residue_names_1, title=r"$\chi$1"
+        )
+
+        # Top-right: χ2
+        self._plot_one_chi_angle(
+            axes[0, 1], chi_2_angles_degrees, residue_names_2, title="$\chi$2"
+        )
+
+        # Bottom-left: χ3
+        self._plot_one_chi_angle(
+            axes[1, 0], chi_3_angles_degrees, residue_names_3, title="$\chi$3"
+        )
+
+        # Bottom-right: χ4
+        self._plot_one_chi_angle(
+            axes[1, 1], chi_4_angles_degrees, residue_names_4, title="$\chi$4"
+        )
+        # add title
+        fig.suptitle("Chi angles per residue for simulation {sim}", fontsize=16)
+        plt.tight_layout()
+        # plt.show()
+        # Save the figure
+        file_name = self.path_registry.write_file_name(
+            FileType.FIGURE,
+            fig_analysis="chi_angles",
+            file_format="png",
+            Sim_id=sim_id,
+        )
+        desc = f"Chi angles plot for the simulation {sim_id}"
+        plot_id = self.path_registry.get_fileid(file_name, FileType.FIGURE)
+        path = self.path_registry.ckpt_dir + "/figures/"
+        plt.savefig(path + file_name)
+        self.path_registry.map_path(plot_id, path + file_name, description=desc)
+        plt.clf()  # Clear the current figure so it does not overlay next plot
+        return plot_id, "Succeeded. Chi angles plot saved."
 
     def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
         """
         Main function to decide which analysis to do:
-        'phi-psi', 'chi1-chi2', or 'all'.
+        'phi-psi', 'chis', or 'all'.
         """
         # Store optional references for convenience
         self_path_registry = path_registry
@@ -242,17 +306,17 @@ def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
 
         # ================ PHI-PSI ONLY =================
         if analysis == "phi-psi":
-            plot_id, message = self.compute_and_plot_phi_psi(
+            ram_plot_id, phi_message = self.compute_and_plot_phi_psi(
                 traj, self_path_registry, self_sim_id
             )
-            return message
+            return f"Ramachandran plot with ID {ram_plot_id}, message: {phi_message} "
 
         # ================ CHI1-CHI2 ONLY ================
-        elif analysis == "chi1-chi2":
-            plot_id, message = self.compute_and_plot_chi1_chi2(
+        elif analysis == "chis":
+            chi_plot_id, chi_message = self.compute_plot_all_chi_angles(
                 traj, self_path_registry, self_sim_id
             )
-            return message
+            return f"Chis plot with ID {chi_plot_id}, message: {chi_message}"
 
         # ================ ALL =================
         elif analysis == "all":
@@ -264,7 +328,7 @@ def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
                 return phi_message
 
             # Then do chi1-chi2
-            chi_plot_id, chi_message = self.compute_and_plot_chi1_chi2(
+            chi_plot_id, chi_message = self.compute_plot_all_chi_angles(
                 traj, self_path_registry, self_sim_id
             )
             if "Failed." in chi_message:
@@ -272,8 +336,8 @@ def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
 
             return (
                 "Succeeded. All analyses completed. "
-                f"Ramachandran plot message: {phi_message} "
-                f"Chi1-Chi2 plot message: {chi_message}"
+                f"Ramachandran plot with ID {ram_plot_id}, message: {phi_message} "
+                f"Chis plot with ID {chi_plot_id}, message: {chi_message}"
             )
 
         else:
@@ -303,7 +367,7 @@ def validate_input(self, **input):
         if analysis.lower() not in [
             "all",
             "phi-psi",
-            "chi1-chi2",
+            "chis",
         ]:
             analysis = "all"
             system_message += (

From f834550d67afb8cf582f083c616b94046e87ee3c Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Mon, 27 Jan 2025 15:54:43 -0500
Subject: [PATCH 19/23] typo

---
 .../base_tools/analysis_tools/bond_angles_dihedrals_tool.py     | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index 2f6655da..d7c0b1da 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -336,7 +336,7 @@ def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
 
             return (
                 "Succeeded. All analyses completed. "
-                f"Ramachandran plot with ID {ram_plot_id}, message: {phi_message} "
+                f"Ramachandran plot with ID {phi_plot_id}, message: {phi_message} "
                 f"Chis plot with ID {chi_plot_id}, message: {chi_message}"
             )
 

From 01ebab57956dc0a8f7d90c4580d4bd282da87910 Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Tue, 28 Jan 2025 10:44:21 -0500
Subject: [PATCH 20/23] fixing input error (Removing path registry as input) in
 the analyze angles function

---
 .../bond_angles_dihedrals_tool.py              | 18 +++++++-----------
 1 file changed, 7 insertions(+), 11 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index d7c0b1da..d5637a39 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -105,7 +105,7 @@ def _run(self, input):
         except Exception as e:
             return f"Failed. Error loading trajectory: {str(e)}"
 
-        return self.analyze_trajectory(traj, analysis, self.path_registry, traj_id)
+        return self.analyze_trajectory(traj, analysis, traj_id)
 
     async def _arun(self, input):
         raise NotImplementedError("Async version not implemented")
@@ -295,41 +295,37 @@ def compute_plot_all_chi_angles(self, traj, sim_id="sim"):
         plt.clf()  # Clear the current figure so it does not overlay next plot
         return plot_id, "Succeeded. Chi angles plot saved."
 
-    def analyze_trajectory(self, traj, analysis, path_registry=None, sim_id="sim"):
+    def analyze_trajectory(self, traj, analysis, sim_id="sim"):
         """
         Main function to decide which analysis to do:
         'phi-psi', 'chis', or 'all'.
         """
         # Store optional references for convenience
-        self_path_registry = path_registry
+
         self_sim_id = sim_id
 
         # ================ PHI-PSI ONLY =================
         if analysis == "phi-psi":
-            ram_plot_id, phi_message = self.compute_and_plot_phi_psi(
-                traj, self_path_registry, self_sim_id
-            )
+            ram_plot_id, phi_message = self.compute_and_plot_phi_psi(traj, self_sim_id)
             return f"Ramachandran plot with ID {ram_plot_id}, message: {phi_message} "
 
         # ================ CHI1-CHI2 ONLY ================
         elif analysis == "chis":
             chi_plot_id, chi_message = self.compute_plot_all_chi_angles(
-                traj, self_path_registry, self_sim_id
+                traj, self_sim_id
             )
             return f"Chis plot with ID {chi_plot_id}, message: {chi_message}"
 
         # ================ ALL =================
         elif analysis == "all":
             # First do phi-psi
-            phi_plot_id, phi_message = self.compute_and_plot_phi_psi(
-                traj, self_path_registry, self_sim_id
-            )
+            phi_plot_id, phi_message = self.compute_and_plot_phi_psi(traj, self_sim_id)
             if "Failed." in phi_message:
                 return phi_message
 
             # Then do chi1-chi2
             chi_plot_id, chi_message = self.compute_plot_all_chi_angles(
-                traj, self_path_registry, self_sim_id
+                traj, self_sim_id
             )
             if "Failed." in chi_message:
                 return chi_message

From 777797d25dd35508d0f545c8e8044c8d9641cee2 Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Tue, 28 Jan 2025 11:41:10 -0500
Subject: [PATCH 21/23] fixing typos and input descriptions

---
 .../bond_angles_dihedrals_tool.py             | 32 ++++++++-----------
 1 file changed, 13 insertions(+), 19 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index d5637a39..e60f05f3 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -40,7 +40,7 @@ class ComputingAnglesSchema(BaseModel):
 class ComputeAngles(BaseTool):
     name = "ComputeAngles"
     description = """Analyze dihedral angles from a trajectory file. The tool allows for
-    analysis of the phi-psi angles, chi1-chi2 angles, or both. """
+    analysis of the phi-psi angles, chis angles, or both. """
 
     path_registry: PathRegistry | None = None
     args_schema = ComputingAnglesSchema
@@ -105,14 +105,14 @@ def _run(self, input):
         except Exception as e:
             return f"Failed. Error loading trajectory: {str(e)}"
 
-        return self.analyze_trajectory(traj, analysis, traj_id)
+        return self.analyze_trajectory(traj, analysis, sim_id=traj_id)
 
     async def _arun(self, input):
         raise NotImplementedError("Async version not implemented")
 
     # Example helper functions (optional). You can instead just keep them as
     # blocks in the if-statements.
-    def compute_and_plot_phi_psi(self, traj, path_registry, sim_id):
+    def compute_and_plot_phi_psi(self, traj, sim_id):
         """
         Computes phi-psi angles, saves results to file, and produces Ramachandran plot.
         """
@@ -128,7 +128,7 @@ def compute_and_plot_phi_psi(self, traj, path_registry, sim_id):
             return None, f"Failed. Error computing phi-psi angles: {str(e)}"
 
         # If path_registry is available, save files and produce plot
-        if path_registry is not None:
+        if self.path_registry is not None:
             # Save angle results
             save_results_to_file("phi_results.npz", phi_indices, phi_angles)
             save_results_to_file("psi_results.npz", psi_indices, psi_angles)
@@ -142,17 +142,17 @@ def compute_and_plot_phi_psi(self, traj, path_registry, sim_id):
                 plt.ylabel(r"$\psi$")
                 plt.colorbar()
 
-                file_name = path_registry.write_file_name(
+                file_name = self.path_registry.write_file_name(
                     FileType.FIGURE,
                     fig_analysis="ramachandran",
                     file_format="png",
                     Sim_id=sim_id,
                 )
                 desc = f"Ramachandran plot for the simulation {sim_id}"
-                plot_id = path_registry.get_fileid(file_name, FileType.FIGURE)
-                path = path_registry.ckpt_dir + "/figures/"
+                plot_id = self.path_registry.get_fileid(file_name, FileType.FIGURE)
+                path = self.path_registry.ckpt_dir + "/figures/"
                 plt.savefig(path + file_name)
-                path_registry.map_path(plot_id, path + file_name, description=desc)
+                self.path_registry.map_path(plot_id, path + file_name, description=desc)
                 plt.clf()  # Clear the current figure so it does not overlay next plot
                 print("Ramachandran plot saved to file")
                 return plot_id, "Succeeded. Ramachandran plot saved."
@@ -277,7 +277,7 @@ def compute_plot_all_chi_angles(self, traj, sim_id="sim"):
             axes[1, 1], chi_4_angles_degrees, residue_names_4, title="$\chi$4"
         )
         # add title
-        fig.suptitle("Chi angles per residue for simulation {sim}", fontsize=16)
+        fig.suptitle(f"Chi angles per residue for simulation {sim_id}", fontsize=16)
         plt.tight_layout()
         # plt.show()
         # Save the figure
@@ -302,31 +302,25 @@ def analyze_trajectory(self, traj, analysis, sim_id="sim"):
         """
         # Store optional references for convenience
 
-        self_sim_id = sim_id
-
         # ================ PHI-PSI ONLY =================
         if analysis == "phi-psi":
-            ram_plot_id, phi_message = self.compute_and_plot_phi_psi(traj, self_sim_id)
+            ram_plot_id, phi_message = self.compute_and_plot_phi_psi(traj, sim_id)
             return f"Ramachandran plot with ID {ram_plot_id}, message: {phi_message} "
 
         # ================ CHI1-CHI2 ONLY ================
         elif analysis == "chis":
-            chi_plot_id, chi_message = self.compute_plot_all_chi_angles(
-                traj, self_sim_id
-            )
+            chi_plot_id, chi_message = self.compute_plot_all_chi_angles(traj, sim_id)
             return f"Chis plot with ID {chi_plot_id}, message: {chi_message}"
 
         # ================ ALL =================
         elif analysis == "all":
             # First do phi-psi
-            phi_plot_id, phi_message = self.compute_and_plot_phi_psi(traj, self_sim_id)
+            phi_plot_id, phi_message = self.compute_and_plot_phi_psi(traj, sim_id)
             if "Failed." in phi_message:
                 return phi_message
 
             # Then do chi1-chi2
-            chi_plot_id, chi_message = self.compute_plot_all_chi_angles(
-                traj, self_sim_id
-            )
+            chi_plot_id, chi_message = self.compute_plot_all_chi_angles(traj, sim_id)
             if "Failed." in chi_message:
                 return chi_message
 

From 1f902cf0c6b741febf695ca549167dc9f0b10fa5 Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Tue, 4 Feb 2025 14:04:18 -0500
Subject: [PATCH 22/23] change name for function in tests

---
 tests/test_analysis/test_bond_angles_dihedrals.py | 10 +++++-----
 1 file changed, 5 insertions(+), 5 deletions(-)

diff --git a/tests/test_analysis/test_bond_angles_dihedrals.py b/tests/test_analysis/test_bond_angles_dihedrals.py
index 0e04c81b..e71cac8e 100644
--- a/tests/test_analysis/test_bond_angles_dihedrals.py
+++ b/tests/test_analysis/test_bond_angles_dihedrals.py
@@ -40,20 +40,20 @@ def test_compute_angles_ram_values(get_registry):
         "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_and_plot_phi_psi"
     ) as mock_compute_and_plot_phi_psi:
         with patch(
-            "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_and_plot_chi1_chi2"
-        ) as mock_compute_and_plot_chi1_chi2:
+            "mdagent.tools.base_tools.analysis_tools.ComputeAngles.compute_plot_all_chi_angles"
+        ) as compute_plot_all_chi_angles:
             mock_compute_and_plot_phi_psi.return_value = ("mockid", "mockresult")
             # instance.return_value = ("mockid", "mockresult")
             angles_tool._run(phi_psi_input_files)
             # print(result)
             assert mock_compute_and_plot_phi_psi.called
             # assert compute_and_plot_chi1_chi2 is not called
-            assert not mock_compute_and_plot_chi1_chi2.called
+            assert not compute_plot_all_chi_angles.called
 
             # =========================================================================#
-            mock_compute_and_plot_chi1_chi2.return_value = ("mockid", "mockresult")
+            compute_plot_all_chi_angles.return_value = ("mockid", "mockresult")
             angles_tool._run(chi_innput_files)
-            assert mock_compute_and_plot_chi1_chi2.called
+            assert compute_plot_all_chi_angles.called
             # assert compute_and_plot_phi_psi is not called
             assert (
                 mock_compute_and_plot_phi_psi.assert_called_once

From 3a3cd723cdae772180557aca720343235610293a Mon Sep 17 00:00:00 2001
From: Jorge <jmedina9@ur.rochester.edu>
Date: Tue, 4 Feb 2025 15:53:08 -0500
Subject: [PATCH 23/23] adding selection usage in angle tools

---
 .../analysis_tools/bond_angles_dihedrals_tool.py    | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
index e60f05f3..2ed2e7d5 100644
--- a/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
+++ b/mdagent/tools/base_tools/analysis_tools/bond_angles_dihedrals_tool.py
@@ -31,8 +31,10 @@ class ComputingAnglesSchema(BaseModel):
     selection: Optional[str] = Field(
         "backbone and sidechain",
         description=(
-            "Which selection of atoms from the simulation "
-            "to use for the pca analysis"
+            "A string specifying which atoms to select from the trajectory, using "
+            "MDTraj’s selection syntax. Common examples include expressions like 'resid "
+            "1 to 10', 'name CA', or 'backbone' to define particular subsets of atoms "
+            "for analysis."
         ),
     )
 
@@ -104,6 +106,13 @@ def _run(self, input):
             return f"Failed. Error loading trajectory: {str(e)}"
         except Exception as e:
             return f"Failed. Error loading trajectory: {str(e)}"
+        # make selection
+        if selection:
+            try:
+                traj = traj.atom_slice(traj.top.select(selection))
+            except Exception as e:
+                # return f"Failed. Error selecting atoms: {str(e)}"
+                print(f"Error selecting atoms: {str(e)}, defaulting to all atoms")
 
         return self.analyze_trajectory(traj, analysis, sim_id=traj_id)