From c9c7898d429dfbb65a1dabb26cdad7cb2aad1ada Mon Sep 17 00:00:00 2001
From: ehhartmann <69237857+ehhartmann@users.noreply.github.com>
Date: Wed, 21 Feb 2024 17:44:39 +0100
Subject: [PATCH] Feat: analyse reaction participation (#365)

* trigger PR

* add new config options to schema

* add structure for passing residuetypes to FF class

* fix: typo

* expose reference to residuetypes file in config

* finish radical part and add tests

* use new black version

* extend radical and residuetypes treatment to kimmdy-modify-top

* black

* fix bugs for lowercase atomtypes (gaff)

* implement tool

---------

Co-authored-by: Eric Hartmann <hartmaec@rh05659.villa-bosch.de>
Co-authored-by: KRiedmiller <kai.riedmiller@web.de>
---
 src/kimmdy/analysis.py | 83 +++++++++++++++++++++++++++++++++++++++---
 1 file changed, 78 insertions(+), 5 deletions(-)

diff --git a/src/kimmdy/analysis.py b/src/kimmdy/analysis.py
index 476c1811..cfb248c6 100644
--- a/src/kimmdy/analysis.py
+++ b/src/kimmdy/analysis.py
@@ -18,7 +18,7 @@
 
 from kimmdy.utils import run_shell_cmd
 from kimmdy.parsing import read_json, write_json
-from kimmdy.recipe import RecipeCollection
+from kimmdy.recipe import RecipeCollection, Break, Bind, Place, Relax
 
 
 def get_analysis_dir(dir: Path) -> Path:
@@ -336,9 +336,12 @@ def radical_population(
     atoms.tempfactors = list(occupancy.values())
     pdb_output = f"{analysis_dir}/radical_population_selection.pdb"
     atoms.write(pdb_output)
-    protein = u.select_atoms("protein")
-    pdb_output_protein = f"{analysis_dir}/radical_population.pdb"
-    protein.write(pdb_output_protein)
+    try:
+        protein = u.select_atoms("protein")
+        pdb_output_protein = f"{analysis_dir}/radical_population.pdb"
+        protein.write(pdb_output_protein)
+    except IndexError:
+        print("Problem with writing protein pdb. Continuing.")
 
     if open_plot:
         sp.call(("xdg-open", output_path))
@@ -479,6 +482,59 @@ def time_from_logfile(log_path: Path, sep: int, factor: float = 1.0):
         sp.call(("xdg-open", output_path))
 
 
+def reaction_participation(dir: str, open_plot: bool = False):
+    """Plot runtime of all tasks.
+
+    Parameters
+    ----------
+    dir
+        Directory of KIMMDY run
+    open_plot
+        Open plot in default system viewer.
+    """
+    print(
+        "Count reaction participation\n"
+        f"dir: \t\t\t{dir}\n"
+        f"open_plot: \t\t{open_plot}\n"
+    )
+
+    run_dir = Path(dir).expanduser().resolve()
+    analysis_dir = get_analysis_dir(run_dir)
+
+    reaction_count = {"overall": 0}
+    for recipes in run_dir.glob("*decide_recipe/recipes.csv"):
+        # get picked recipe
+        rc, picked_rp = RecipeCollection.from_csv(recipes)
+        # get involved atoms
+        reaction_atom_ids = set()
+        for step in picked_rp.recipe_steps:
+            if isinstance(step, Break) or isinstance(step, Bind):
+                reaction_atom_ids |= set([step.atom_id_1, step.atom_id_2])
+            elif isinstance(step, Place):
+                reaction_atom_ids |= set([step.id_to_place])
+            else:
+                continue
+        # update count
+        for atom_id in reaction_atom_ids:
+            if atom_id in reaction_count:
+                reaction_count[atom_id] += 1
+            else:
+                reaction_count[atom_id] = 1
+        reaction_count["overall"] += 1
+
+    reaction_count = dict(sorted(reaction_count.items()))
+    sns.barplot(
+        x=list(reaction_count.keys()), y=list(reaction_count.values()), errorbar=None
+    )
+
+    plt.xlabel("Atom identifier")
+    plt.ylabel("Reaction participation count")
+    output_path = str(analysis_dir / "reaction_participation_fingerprint.png")
+    plt.savefig(output_path, dpi=300)
+    if open_plot:
+        sp.call(("xdg-open", output_path))
+
+
 def get_analysis_cmdline_args() -> argparse.Namespace:
     """Parse command line arguments.
 
@@ -621,7 +677,19 @@ def get_analysis_cmdline_args() -> argparse.Namespace:
         action="store_true",
         help="Open plot in default system viewer.",
     )
-
+    parser_reaction_participation = subparsers.add_parser(
+        name="reaction_participation",
+        help="Plot counts of reaction participation per atom id",
+    )
+    parser_reaction_participation.add_argument(
+        "dir", type=str, help="KIMMDY run directory to be analysed."
+    )
+    parser_reaction_participation.add_argument(
+        "--open-plot",
+        "-p",
+        action="store_true",
+        help="Open plot in default system viewer.",
+    )
     return parser.parse_args()
 
 
@@ -646,6 +714,11 @@ def entry_point_analysis():
         plot_rates(args.dir)
     elif args.module == "runtime":
         plot_runtime(args.dir, args.md_tasks, args.datefmt, args.open_plot)
+    elif args.module == "reaction_participation":
+        reaction_participation(
+            args.dir,
+            args.open_plot,
+        )
     else:
         print(
             "No analysis module specified. Use -h for help and a list of available modules."