From 06bac09dca74740c57cd3e76b79db9377bcfe507 Mon Sep 17 00:00:00 2001
From: David Kastner <kastner@mit.edu>
Date: Fri, 26 Jul 2024 13:05:55 -0400
Subject: [PATCH] Combine NEBs into a single trajectory

---
 pyqmmm/cli.py             |  6 +++++
 pyqmmm/qm/combine_nebs.py | 46 +++++++++++++++++++++++++++++++++++++++
 2 files changed, 52 insertions(+)
 create mode 100644 pyqmmm/qm/combine_nebs.py

diff --git a/pyqmmm/cli.py b/pyqmmm/cli.py
index 068d316..207eced 100644
--- a/pyqmmm/cli.py
+++ b/pyqmmm/cli.py
@@ -208,6 +208,7 @@ def md(
 @click.option("--orca_scan", "-os", is_flag=True, help="Plots an ORCA scan.")
 @click.option("--orca_neb_restart", "-rneb", is_flag=True, help="Prepare to restart an ORCA NEB.")
 @click.option("--create_neb_mep", "-mep", is_flag=True, help="Creates MEP_trj from out.")
+@click.option("--combine_nebs", "-cneb", is_flag=True, help="Combines NEBs into a single trajectory.")
 @click.help_option('--help', '-h', is_flag=True, help='Exiting pyQMMM.')
 def qm(
     plot_energy,
@@ -219,6 +220,7 @@ def qm(
     orca_scan,
     orca_neb_restart,
     create_neb_mep,
+    combine_nebs,
     ):
     """
     Functions for quantum mechanics (QM) simulations.
@@ -294,6 +296,10 @@ def qm(
         import pyqmmm.qm.create_mep_trj
         pyqmmm.qm.create_mep_trj.create_neb_mep_trj_from_out()
 
+    if combine_nebs:
+        import pyqmmm.qm.combine_nebs
+        pyqmmm.qm.combine_nebs.get_combined_trajectory()
+
 
 if __name__ == "__main__":
     # Run the command-line interface when this script is executed
diff --git a/pyqmmm/qm/combine_nebs.py b/pyqmmm/qm/combine_nebs.py
new file mode 100644
index 0000000..fdd60a0
--- /dev/null
+++ b/pyqmmm/qm/combine_nebs.py
@@ -0,0 +1,46 @@
+import os
+import re
+
+def get_xyz_files():
+    # List all files in the current working directory
+    files = os.listdir('.')
+    # Filter files to get only those with a single number in their name ending with .xyz
+    xyz_files = [f for f in files if re.match(r'^\d+\.xyz$', f)]
+    # Sort files based on the numerical value in their name
+    xyz_files.sort(key=lambda x: int(re.match(r'(\d+)', x).group()))
+    return xyz_files
+
+def read_xyz_frames(file):
+    with open(file, 'r') as f:
+        lines = f.readlines()
+    frames = []
+    i = 0
+    while i < len(lines):
+        num_atoms = int(lines[i].strip())
+        header = lines[i + 1].strip()
+        frame = lines[i:i + num_atoms + 2]
+        frames.append(frame)
+        i += num_atoms + 2
+    return frames
+
+def combine_xyz_files(xyz_files, output_file):
+    combined_frames = []
+    
+    for i, file in enumerate(xyz_files):
+        frames = read_xyz_frames(file)
+        if i < len(xyz_files) - 1:
+            combined_frames.extend(frames[:-1])  # Exclude the last frame of the file
+        else:
+            combined_frames.extend(frames)  # Include all frames of the last file
+    
+    with open(output_file, 'w') as f:
+        for frame in combined_frames:
+            f.writelines(frame)
+
+def get_combined_trajectory():
+    xyz_files = get_xyz_files()
+    if xyz_files:
+        combine_xyz_files(xyz_files, 'combined.xyz')
+        print("Combined file 'combined.xyz' created successfully.")
+    else:
+        print("No suitable .xyz files found.")